module Visualization
include("gray_scott_solver.jl")
using GLMakie, Observables, Makie
using Random
using .GrayScottSolver

"""
    step_through_solution(sol::SolutionType, N::Int)

    Function for visualization for the output of run_simulation

    # Arguments
    - sol: computed differential equation by run_simulation
    - N: size of the N×N grid

    # Returns
    - Displays created figure
"""
function step_through_solution(sol, N::Int)
    fig = Figure(resolution=(600, 600))
    ax = Axis(fig[1, 1])
    slider = Slider(fig[2, 1], range=1:length(sol), startvalue=1)

    textbox = Textbox(fig[1, 2], placeholder="Feed Rate", validator=Float64)
    F = Observable(0.060)
    # Initialize heatmap with first time step
    u0 = reshape(sol[1][1:N^2], N, N)
    heat_obs = Observable(u0)
    heatmap!(ax, heat_obs, colormap=:magma)

    on(textbox.stored_string) do s
        F[] = parse(Float64, s)
    end

    # Update heatmap on slider movement
    on(slider.value) do i
        u = reshape(sol[i][1:N^2], N, N)
        heat_obs[] = u
    end

    display(fig)
end

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

function reset!(U, V, heat_obs)
    U .= 1.0
    V .= 0.0
    # Default starting point for Run
    default_start!(U, V)
    heat_obs[] = copy(U)
end

function default_start!(U, V)
    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
end

function starting_points!(U, V)
    N = size(U, 1)
    radius = 10
    jitter = 5
    margin = radius + 2
    rels = [
        (1 / 4, 1 / 4), (3 / 4, 1 / 4), (1 / 2, 1 / 2),
        (1 / 4, 3 / 4), (3 / 4, 3 / 4), (1 / 2, 1 / 8)
    ]
    for (px, py) in rels
        cx = clamp(round(Int, px * N) + rand(-jitter:jitter), margin, N - margin)
        cy = clamp(round(Int, py * N) + rand(-jitter:jitter), margin, N - margin)
        for x in (cx-radius):(cx+radius), y in (cy-radius):(cy+radius)
            if 1 ≤ x ≤ N && 1 ≤ y ≤ N && hypot(x - cx, y - cy) ≤ radius
                U[x, y] = 0.50
                V[x, y] = 0.25
            end
        end
    end
end

function param_box!(grid, row, labeltxt, observable::Observable)
    Label(grid[row, 1], labeltxt)
    box = Textbox(grid[row, 2], validator=Float64, width=50,
        placeholder=labeltxt, stored_string="$(observable[])")
    on(box.stored_string) do s
        try
            observable[] = parse(Float64, s)
            box.displayed_string[] = s
        catch
            @warn "Invalid input for $labeltxt: $s"
        end
    end
    on(observable) do val
        box.displayed_string[] = string(val)
    end
    return box
end

function build_ui(U, V, Du, Dv, F, k, params_obs, heat_obs)
    reset!(U, V, heat_obs)
    fig = Figure(size=(800, 800))
    gh = GridLayout(fig[1, 1])
    ax = Axis(gh[1, 1])
    heatmap!(ax, heat_obs, colormap=:viridis)
    deactivate_interaction!(ax, :rectanglezoom)
    ax.aspect = DataAspect()

    run_label = Observable("Run")
    stepsize = Observable(30)
    spoint = select_point(ax.scene)

    # # Controls
    fig[2, 1] = buttongrid = GridLayout(ax.scene, tellwidth=false)
    btn_step = Button(buttongrid[1, 1], width=60, label="Step")
    btn_start = Button(buttongrid[1, 2], width=60, label=run_label)
    btn_reset = Button(buttongrid[1, 3], width=60, label="Reset")
    btn_Jaguar = Button(buttongrid[1, 4], width=60, label="Jaguar")

    slidergrid = SliderGrid(fig[3, 1], (label="Speed", range=1:100,
        format="{}x", width=350,
        startvalue=stepsize[]))
    speed_slider = slidergrid.sliders[1].value

    gh[1, 2] = textboxgrid = GridLayout(ax.scene, tellwidth=false)
    rowsize!(gh, 1, Relative(1.0))
    param_box!(textboxgrid, 1, "Du", Du)
    param_box!(textboxgrid, 2, "Dv", Dv)
    param_box!(textboxgrid, 3, "Feed", F)
    param_box!(textboxgrid, 4, "kill", k)


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

    on(speed_slider) do s
        stepsize[] = s
        println("Changed stepsize to $s")
    end

    on(btn_step.clicks) do _
        multi_step!((U, V), stepsize[], heat_obs, params_obs)
    end

    on(btn_start.clicks) do _
        running[] = !running[]
        if running[]
            @async while running[]
                multi_step!((U, V), stepsize[], heat_obs, params_obs)
                sleep(0.0015)
            end
        end
    end

    on(btn_Jaguar.clicks) do _
        Du[] = 0.142
        Dv[] = 0.078
        F[] = 0.0617
        k[] = 0.062
        U .= 1.0
        V .= 0.0
        starting_points!(U, V)
        heat_obs[] = copy(U)
    end

    on(btn_reset.clicks) do _
        running[] = false
        reset!(U, V, heat_obs)
    end
    on(spoint) do pt
        r = 5
        if pt === nothing
            return
        end
        x, y = pt
        i, j = coord_to_index(x, y, params_obs[].N)

        # get corners of square that will get filled with concentration
        imin = max(i - r, 1)
        imax = min(i + r, params_obs[].N)
        jmin = max(j - r, 1)
        jmax = min(j + r, params_obs[].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
    return fig
end

export step_through_solution, build_ui
end