using DifferentialEquations
using Plots

function iz!(du, u, p, t)
    (C, a, b, k, vᵣ, vₜ), I = p
    v, i, _ = u
    
    du[1] = C \ (k * (v - vᵣ) * (v - vₜ) - i + I(t))
    du[2] = a * (b * (v - vᵣ) - i)
    du[3] = 0
end

params = (
    C = 100,
    a = 0.03,
    b = -2.0,
    k = 0.7,
    vᵣ = -60,
    vₜ = -40
)

condition(u, t, integrator) = u[1] - 35

function affect!(integrator)
    integrator.u[1] = -50
    integrator.u[2] += 100
    integrator.u[3] = integrator.t
end

cb = ContinuousCallback(condition, affect!)

I(t) = ifelse(50 ≤ t ≤ 350, 100, 0)
tspan = (0.0, 400.0)
u0 = [-60.0, 0.0, 0.0]
alg = Tsit5() # Vern7()

prob = ODEProblem(iz!, u0, tspan, (params, I))
@time sol = solve(prob, alg; callback=cb)
spiketimings = unique(sol[3, :])[2:end]

plot(sol;
    vars=(0, 1),
    plotdensity=20000,
    xlabel="Time (ms)",
    ylabel="Membrane potential (mV)",
    legend=false
)

plot!(t -> I(t) / 20 - 70, range(tspan...; length=1000))

scatter!(_ -> 40, spiketimings; shape=:dtriangle)

function firerate(current)
    prob = ODEProblem(iz!, [-60.0, 0.0, 0.0], (0.0, 100000.0), (params, _ -> current))
    sol = solve(prob, Tsit5(); callback=cb)
    ts = unique(sol[3, :])[2:end]
    length(ts) ≥ 2 ? 1000 / (ts[end] - ts[end-1]) : 0.0
end

currents = range(50, 55; length=200)
@time frequencies = firerate.(currents)

scatter(currents, frequencies;
    xlabel="Input current (pA)",
    ylabel="Firing rate (Hz)",
    xrange=extrema(currents),
    yrange=extrema(frequencies),
    markersize=2,
    legend=false
)