using DifferentialEquations
using Plots
function f(du,u,p,t)
    du[1] = -u[1]
end
u0 = [10.0]
const V = 1
prob = ODEProblem(f,u0,(0.0,10.0))

condition(u,t,integrator) = t==4
affect!(integrator) = integrator.u[1] += 10
cb = DiscreteCallback(condition,affect!)

sol = solve(prob,Tsit5(),callback=cb,tstops=[4.0])
plot(sol)


# control problem
function f(du,u,p,t)
    du[1] = -0.5*u[1] + p
    du[2] = -0.5*u[2]
end

v

u0 = [10.0;10.0]
p = 0.0
prob = ODEProblem(f,u0,(0.0,10.0),p)

const tstop = [5.;8.]
sol = solve(prob,Tsit5(),callback = cbs, tstops=tstop)
plot(sol)

#### vector callback

function f(du,u,p,t)
  du[1] = u[2]
  du[2] = -p
  du[3] = u[4]
  du[4] = 0.0
end

function condition(out,u,t,integrator) # Event when event_f(u,t) == 0
  out[1] = u[1]
  out[2] = (u[3] - 10.0)u[3]
end

function affect!(integrator, idx)
  if idx == 1
    integrator.u[2] = -0.9integrator.u[2]
  elseif idx == 2
    integrator.u[4] = -0.9integrator.u[4]
  end
end

cb = VectorContinuousCallback(condition,affect!,2)


u0 = [50.0,0.0,0.0,2.0]
tspan = (0.0,15.0)
p = 9.8
prob = ODEProblem(f,u0,tspan,p)
sol = solve(prob,Tsit5(),callback=cb,dt=1e-3,adaptive=false)
plot(sol, vars = (3,1))



#### my testing
function dynamics!(du,u,p,t)
  α, β, γ, δ = p
  du[1] = (α - β * u[2]) * u[1]
  du[2] = (-γ + δ * u[1]) * u[2]
end

function condition(out, u , t, integrator)
  out[1] = 20-u[1]
  out[2] = 20-u[2]
end

function affect!(integrator, idx)
  if idx == 1
    integrator.u[1] = 0
  end
  if idx == 2
    integrator.u[2] = 0
  end 
end

cb = VectorContinuousCallback(condition,affect!,2)

u0 = [100, 100]
p = [0.55, 0.028, 0.84, 0.026]
tspan = (0.0,20.0)
prob = ODEProblem{true}(dynamics!,u0,tspan,p)
sol = solve(prob,Tsit5(), callback = cb)
plot(sol)
@show sol[end]
print(sol)
minimum(sol)









# vector callback with just parameters
function dynamics!(du,u,p,t)
  α, β, γ, δ = p
  du[1] = (α - β * u[2]) * u[1]
  du[2] = (-γ + δ * u[1]) * u[2]
end


const stop1 = [1.]
const stop2 = [3.]

function condition(u,t,integrator)
  t in stop1
end
function condition2(u,t,integrator)
  t in stop2
end
function affect!(integrator)
  integrator.p[4] += 4  
end
function affect2!(integrator)
  integrator.p[4] -= 4  
end

save_positions = (true,true)
cb = DiscreteCallback(condition, affect!, save_positions=save_positions)
save_positions = (false,true)
cb2 = DiscreteCallback(condition2, affect2!, save_positions=save_positions)

cbs = CallbackSet(cb,cb2)
# set constant stop values
const stop1 = [1.0]
const stop2 = [3.0]

function time_condition1(out,u,t,integrator)
  t in stop1
end
function time_condition2(out,u,t,integrator)
  t in stop2
end

function time_affect1!(integrator, event_idx)
    integrator.p[4] += 4  
end
function time_affect2!(integrator)
    integrator.p[4] -= 4  
end

save_positions = (true,true)
raise_cb = DiscreteCallback(time_condition1, 
                            time_affect1!, 
                            save_positions=save_positions)

save_positions = (false,true)
lower_cb = DiscreteCallback(time_condition2, 
                            time_affect2!, 
                            save_positions=save_positions)

function pop_condition(out,u,t,integrator)
  out[1] = 20-u[1]
  out[2] = 20-u[2]
end

function pop_condition_test(out,u,t,integrator)
  for i in 1:2
    out[i] = 1e-2 - u[i]
  end
end 
  
  out[1] = 20-u[1]
  out[2] = 20-u[2]
end
function pop_affect!(integrator, event_idx)
  if event_idx == 1
    integrator.u[1] = 0
  end
  if event_idx == 2
    integrator.u[2] = 0
  end
end
function pop_affect_test!(integrator, event_idx)
  for i in 1:2
    if event_idx == i
      integrator.u[i] = 0
    end 
  end
end 

cont_cb1 = VectorContinuousCallback(pop_condition_test,pop_affect_test!,2,save_positions=(true,true))

cb_set = CallbackSet(cb, cb2, cont_cb)


u0 = [100, 100]
p = [0.55, 0.028, 0.84, 0.026]
tspan = (0.0,20.0)
const tstops = [1.;3.]
prob = ODEProblem{true}(dynamics!,u0,tspan,p)
sol = solve(prob,Tsit5(), callback = cont_cb1)
plot(sol, legend = false)

@show sol[end]
print(sol)
minimum(sol)


## calback set with both parameter change and zero condition 

function dynamics!(du,u,p,t)
  α, β, γ, δ = p
  du[1] = (α - β * u[2]) * u[1]
  du[2] = (-γ + δ * u[1]) * u[2]
end

function u1_zero_condition!(u,t,integrator) # Event when event_f(u,t) == 0
  0.5-u[1]
end
function u2_zero_condition!(u,t,integrator) # Event when event_f(u,t) == 0
  0.5-u[2]
end

function u1_zero_affect!(integrator)
  integrator.u[1] = 0
end
function u2_zero_affect!(integrator)
  integrator.u[2] = 0
end

# set constant stop values
const stop1 = [5.0]
const stop2 = [8.0]

function inc_mort_cond(u,t,integrator)
  t in stop1
end

# decrease mortality (stop pulse)
function dec_mort_cond(u,t,integrator)
  t in stop2
end

# apply an affect when condition is reached
function affect_inc!(integrator)
  integrator.p[4] += 10
end

function affect_dec!(integrator)
  integrator.p[4] -= 10
end

save_positions = (true, true)
cb_inc = ContinuousCallback(inc_mort_cond, affect_inc!, save_positions=save_positions)
save_positions = (false, true)
cb_dec = ContinuousCallback(dec_mort_cond, affect_inc!, save_positions=save_positions)


#u1_zero_event = ContinuousCallback(u1_zero_condition!,u1_zero_affect!)
#u2_zero_event = ContinuousCallback(u2_zero_condition!,u2_zero_affect!)

cb_set = CallbackSet(cb_inc, cb_dec)


u0 = [100, 100]
p = [0.55, 0.028, 0.84, 0.026]
tspan = (0.0,20.0)
prob = ODEProblem{true}(dynamics!,u0,tspan,p)
const tstop = [5.;8.]
sol = solve(prob,Tsit5(), callback = cb_set, tstops=tstop)
plot(sol)
@show sol[end]
print(sol)
minimum(sol)