using DifferentialEquations
using Random
using .Constants

"""
    fhn(du, u, p:FHNParams, t:)

    Implements the spatial dynamics of FitzHugh-Nagumo (fhn). Designed to be 
    within a larger numerical solver of partial differential equations.

    # Arguments
    - `du`: output argument which stores the calculated derivatives 
    - `u`: input vector containing the current state of the system at time t
    - `p`: holds all the fixed parameters of the FHN model
    - `t`: current time 

    # Returns
    - `du`: calculated derivatives put back into the du array
"""
function fhn!(du, u, p::FHNParams, t = 0)
    u_mat = reshape(u[1:p.N^2], p.N, p.N) # activation variable
    v_mat = reshape(u[p.N^2+1:end], p.N, p.N) # deactivation variable

    Δu = reshape(laplacian(u_mat, p.N, p.dx), p.N, p.N)
    Δv = reshape(laplacian(v_mat, p.N, p.dx), p.N, p.N)

    fu = p.Du * Δu .+ u_mat .- u_mat .^ 3 ./ 3 .- v_mat
    fv = p.Dv * Δv .+ p.ϵ * (u_mat .+ p.a .- p.b .* v_mat)

    du .= vcat(vec(fu), vec(fv))
end

"""
    run_simulation(tspan::Tuple{Float64,Float64}, N::Int)

    solving the ODE and modelling it after FHN

    # Arguments
    - `tspan`: tuple of two Float64's representing start and end times for simulation
    - `N`: size of the N×N grid

    # Returns
    - `sol`: solved differential equation (ODE)
"""
function run_simulation(tspan::Tuple{Float64,Float64}, N::Int)
    # Turing-spot parameters
    p = FHNParams(N = N, dx = 1.0, Du = 1e-5, Dv = 1e-3, ϵ = 0.01, a = 0.1, b = 0.5)

    # Initial conditions (random noise)
    Random.seed!(1234)
    # Create two vectors with length N*N with numbers between 0.1 and 0.11
    u0 = vec(0.1 .+ 0.01 .* rand(N, N))
    v0 = vec(0.1 .+ 0.01 .* rand(N, N))

    y0 = vcat(u0, v0)

    prob = ODEProblem(fhn!, y0, tspan, p)
    sol = solve(prob, BS3(), saveat=1.0)  # You can try `Rosenbrock23()` too

    return sol
end