using GLMakie, Observables
include("../src/constants.jl")
using .Constants
include("../src/visualization.jl")
using .Visualization

# # Gray Scott Model Parameters
# Parameters and initial conditions
N = 128
dx = 1.0
# diffusion rates of substance 'u' and 'v'
Du, Dv = Observable(0.16), Observable(0.08)
# feed rate of 'u' and kill rate of 'v'
F, k = Observable(0.060), Observable(0.062)
dt = 1.0
params_obs = Observable(GSParams(N, dx, Du[], Dv[], F[], k[]))

# # GUI Parameters
run_label = Observable{String}("Run")
stepsize = Observable{Int}(30)



function update_params!(params_obs::Observable, u, v, feed, kill)
    old = params_obs[]
    params_obs[] = GSParams(old.N, old.dx, u, v, feed, kill)
end

# Whenever a value gets changed via the textbox, update params object to reflect changes in simulation
lift(Du, Dv, F, k) do u, v, F, k
    update_params!(params_obs, u, v, F, k)
end
U = ones(N, N)
V = zeros(N, N)
center = N ÷ 2
radius = 10
# set a cube in the center with starting concentrations for 'u' and 'v'
U[center-radius:center+radius, center-radius:center+radius] .= 0.50
V[center-radius:center+radius, center-radius:center+radius] .= 0.25

# Observable holding current U for heatmap
heat_obs = Observable(U)


function laplacian5(f)
    h2 = dx^2
    left = f[2:end-1, 1:end-2]
    right = f[2:end-1, 3:end]
    down = f[3:end, 2:end-1]
    up = f[1:end-2, 2:end-1]
    center = f[2:end-1, 2:end-1]
    return (left .+ right .+ down .+ up .- 4 .* center) ./ h2
end

function step!(U, V, params_obs::Observable)
    lap_u = laplacian5(U)
    lap_v = laplacian5(V)
    diff_u = params_obs[].Du
    diff_v = params_obs[].Dv
    feed_u = params_obs[].F
    kill_v = params_obs[].k
    u = U[2:end-1, 2:end-1]
    v = V[2:end-1, 2:end-1]

    uvv = u .* v .* v
    u_new = u .+ (diff_u .* lap_u .- uvv .+ feed_u .* (1 .- u))
    v_new = v .+ (diff_v .* lap_v .+ uvv .- (feed_u .+ kill_v) .* v)

    # Update with new values
    U[2:end-1, 2:end-1] .= u_new
    V[2:end-1, 2:end-1] .= v_new

    # Periodic boundary conditions
    U[1, :] .= U[end-1, :]
    U[end, :] .= U[2, :]
    U[:, 1] .= U[:, end-1]
    U[:, end] .= U[:, 2]

    V[1, :] .= V[end-1, :]
    V[end, :] .= V[2, :]
    V[:, 1] .= V[:, end-1]
    V[:, end] .= V[:, 2]

    # Update heatmap observable
    return U
end
function multi_step!(state, n_steps)
    for _ in 1:n_steps
        heat_obs[] = step!(state[1], state[2], params_obs)
    end
end

# Build GUI
fig = Figure(size=(600, 600))


gh = GridLayout(fig[1, 1])
ax = Axis(gh[1, 1])

hm = Makie.heatmap!(ax, heat_obs, colormap=:viridis)
Makie.deactivate_interaction!(ax, :rectanglezoom)
ax.aspect = DataAspect()

spoint = select_point(ax.scene)

function coord_to_index(x, y, N)
    ix = clamp(round(Int, x), 1, N)
    iy = clamp(round(Int, y), 1, N)
    return ix, iy
end
r = 5

on(spoint) do pt
    if pt === nothing
        return
    end
    x, y = pt
    i, j = coord_to_index(x, y, N)

    # get corners of square that will get filled with concentration
    imin = max(i - r, 1)
    imax = min(i + r, N)
    jmin = max(j - r, 1)
    jmax = min(j + r, N)

    # set disbalanced concentration of U and V
    U[imin:imax, jmin:jmax] .= 0.5
    V[imin:imax, jmin:jmax] .= 0.25

    heat_obs[] = copy(U)
end

# # Controls
fig[2, 1] = buttongrid = GridLayout(ax.scene, tellwidth=false)
btn_step = Button(buttongrid[1, 1], width=50, label="Step")
btn_start = Button(buttongrid[1, 2], width=50, label=run_label)
btn_reset = Button(buttongrid[1, 3], width=50, label="Reset")
slidergrid = SliderGrid(fig[3, 1], (label="Speed", range=1:1:100, format="{}x", width=350, startvalue=stepsize[]))

speed_slider = slidergrid.sliders[1].value
on(speed_slider) do s
    try
        stepsize[] = s
        println("Changed stepsize to $s")
    catch
        @warn "Invalid input for $s"
    end
end

gh[1, 2] = textboxgrid = GridLayout(ax.scene, tellwidth=false)
rowsize!(gh, 1, Relative(1.0))
function param_box!(row, labeltxt, observable)
    Label(textboxgrid[row, 1], labeltxt)
    box = Textbox(textboxgrid[row, 2], validator=Float64, width=50, placeholder=labeltxt, stored_string="$(observable[])")
    on(box.stored_string) do s
        try
            observable[] = parse(Float64, s)
            println("changed $labeltxt to $s")
        catch
            @warn "Invalid input for $labeltxt: $s"
        end
    end
end
param_box!(1, "Du", Du)
param_box!(2, "Dv", Dv)
param_box!(3, "Feed", F)
param_box!(4, "kill", k)


# Timer and state for animation
running = Observable(false)
function reset!(U, V, heat_obs)
    U .= 1.0
    V .= 0.0
    center = size(U, 1) ÷ 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
    heat_obs[] = copy(U)
end

on(running) do r
    run_label[] = r ? "Pause" : "Run"
end


# Button Listeners
on(btn_step.clicks) do _
    multi_step!((U, V), stepsize[])
end

on(btn_start.clicks) do _
    running[] = !running[]
end

on(btn_start.clicks) do _
    @async while running[]
        multi_step!((U, V), stepsize[])
        sleep(0.05)  # ~20 FPS
    end
end

on(btn_reset.clicks) do _
    running[] = false
    reset!(U, V, heat_obs)

end
display(fig)