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

# Laplacian with 5-point stencil
function laplacianA(A::Matrix{Float64}, N::Int, dx::Float64)
    dx2 = dx^2
    lap = zeros(N, N)
    lap[2:end-1, 2:end-1] .= (
        A[1:end-2, 2:end-1] .+ A[3:end, 2:end-1]  # up/down
        .+
        A[2:end-1, 1:end-2] .+ A[2:end-1, 3:end]  # left/right
        .-
        4 .* A[2:end-1, 2:end-1]
    ) ./ dx2
    return lap
end

# Gray-Scott model: reaction + diffusion
function grayscott!(du, u, p::GSParams, t=0)
    N = p.N
    U = reshape(u[1:N^2], N, N)
    V = reshape(u[N^2+1:end], N, N)

    ΔU = laplacianA(U, N, p.dx)
    ΔV = laplacianA(V, N, p.dx)

    UVV = U .* V .* V

    dU = p.Du * ΔU .- UVV .+ p.F .* (1 .- U)
    dV = p.Dv * ΔV .+ UVV .- (p.F .+ p.k) .* V

    du .= vcat(vec(dU), vec(dV))
end

# Run simulation
function run_simulationG(tspan::Tuple{Float64,Float64}, N::Int, GSP::GSParams)
    # Replace this in run_simulation
    p = GSP

    # Initial conditions: U = 1, V = 0 everywhere
    U = ones(N, N)
    V = zeros(N, N)

    # Seed a square in the center
    center = N ÷ 2
    radius = 10
    U[center-radius:center+radius, center-radius:center+radius] .= 0.50
    V[center-radius:center+radius, center-radius:center+radius] .= 0.25

    y0 = vcat(vec(U), vec(V))

    prob = ODEProblem(grayscott!, y0, tspan, p)
    sol = solve(prob, BS3(), saveat=10.0)  # or Rosenbrock23()

    return sol
end

function run_simulationG_no_ode(tspan::Tuple{Float64,Float64}, N::Int, GSP::GSParams)
    p = GSP
    dx2 = p.dx^2
    dt = 1.0  # fixed timestep

    # Initialize U and V
    U = ones(N, N)
    V = zeros(N, N)

    # Seed pattern in the center
    center = N ÷ 2
    radius = 10
    U[center-radius:center+radius, center-radius:center+radius] .= 0.50
    V[center-radius:center+radius, center-radius:center+radius] .= 0.25

    # Store history for visualization (optional)
    snapshots = []

    for i in 1:tspan[2]
        print("Running $i/$(tspan[2])\r")

        lap_U = laplacianA(U, N, p.dx)
        lap_V = laplacianA(V, N, p.dx)

        UVV = U .* V .* V
        dU = p.Du * lap_U .- UVV .+ p.F * (1 .- U)
        dV = p.Dv * lap_V .+ UVV .- (p.F + p.k) .* V

        U .+= dU * dt
        V .+= dV * dt

        if i % 100 == 0  # save every 100 steps
            push!(snapshots, (copy(U), copy(V)))
        end
    end
    return snapshots
end