using DynamicalSystems

function duffing_drift(u0 = [0.1, 0.25]; ω = 1.0, β = 0.2, ε0 = 0.4, α=0.00045)
    return CoupledODEs(duffing_drift_rule, u0, [ω, β, ε0, α])
end

@inbounds function duffing_drift_rule(x, p, t)
    ω, β, ε0, α = p
    dx1 = x[2]
    dx2 = (ε0+α*t)*cos(ω*t) + x[1] - x[1]^3 - 2β * x[2]
    return SVector(dx1, dx2)
end

function EAPD(ds,init_states::Matrix,T;Ttr,ϵ=sqrt(2)*1e-10)
    N = size(init_states)[1]
    ρ = Float64[]
    times = Int64[]

    #save drifting rate value
    α = current_parameters(ds)[4]

    #add test particle to every ensemble member
    init_states_plus_pert = StateSpaceSet(vcat(init_states,init_states))

    #create a pds for the ensemble
    #pds is a ParallelDynamicalSystem
    pds = ParallelDynamicalSystem(ds,init_states_plus_pert)

    set_parameter!.(pds.systems,4,0.0) #set to non-drifting for initial ensemble
    #setting parameter for all systems in one go doesn't work?

    #step system pds to reach attractor(non-drifting)
    #system starts to drift at t0=0.0
    for _ in 1:Ttr
        step!(pds)
    end
    
    #add perturbation to test particles
    for i in 1:N 
        state_i = current_state(pds.systems[i])
        perturbed_state_i = state_i .+ perturbation(ds,ϵ)
        set_state!(pds.systems[N+i],perturbed_state_i)
    end

    set_parameter!.(pds.systems,4,α) #set to drifting for initial ensemble

    #set back time to t0 = 0
    reinit!(pds,current_states(pds))

    #calculate EAPD at each time step
    for _ in 1:T
        ρ_t = EAPD(pds)
        push!(ρ,ρ_t)
        push!(times,current_time(pds))
        step!(pds)
    end

    return ρ,times

end

function EAPD(pds)

    states = current_states(pds)                        
    N = Int(length(states)/2)

    #calculate distance averages
    dists = [norm(states[i] - states[N+i]) for i in 1:N]
    ρ = mean(log.(dists))
    return ρ 

end

function perturbation(ds,ϵ)
    D, T = dimension(ds), eltype(ds)
    Q0 = randn(SVector{D, T})
    Q0 = ϵ * Q0 / norm(Q0)  
end

duffing = duffing_drift()
duffing_map = StroboscopicMap(duffing,2π)

init_states = randn(1000,2)
@time ρ,times = EAPD(duffing_map,init_states,100;Ttr=20)

pl = plot(times,ρ,ylabel=L"\rho(t)",lc=:gray10,lw=2,xlabel=L"t",legend=false,guidefontsize=28)