module Visualization

include("utils/constants.jl")
include("utils/templates.jl")
include("gray_scott_solver.jl")
include("fhn_solver.jl")

using Observables, Makie, GLMakie
using .Constants
using .Templates
using .GrayScottSolver: step_gray_scott!
using .FHNSolver: step_fhn!

const STEP_FUNCTIONS = Dict(
    :gray_scott => GrayScottSolver.step_gray_scott!,
    :fhn => FHNSolver.step_fhn!
)

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

"""
    reset(U, V, heat_obs)

    Resets heatmap to original state by replacing each cell.
    Currently only places a square in the center

    # Arguments
    - `U`: Matrix filled with ones
    - `V`: Empty matrix filled with zeros
    - `heat_obs`: Heatmap observable

    # Returns
    - ``: resetted heatmap observable
"""
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

function param_box!(grid, row, labeltxt, observable::Observable; col=1)
    label_cell = 2 * col - 1
    textbox_cell = 2 * col
    Label(grid[row, label_cell], labeltxt)
    box = Textbox(grid[row, textbox_cell], 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
    on(observable) do val
        box.displayed_string[] = string(val)
    end
end

function multi_step!(state, n_steps, heat_obs::Observable, params_obs::Observable; step_method=step_gray_scott!, dx=1)
    for _ in 1:n_steps
        heat_obs[] = step_method(state[1], state[2], params_obs; dx=1)
    end
end

function build_ui(U, V, param_obs_map::NamedTuple, params_obs, heat_obs)

    reset!(U, V, heat_obs)
    fig = Figure(size=(1300, 950))

    gh = GridLayout(fig[1, 1])
    dropdown = Menu(fig, options=collect(zip(["Gray-Scott", "FHN"], [:gray_scott, :fhn])))
    gh[1, 1] = dropdown
    ax = Axis(gh[2, 1])

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

    run_label = Observable{String}("Run")
    stepsize = Observable(30)
    spoint = select_point(ax.scene)
    step_method = Observable{Function}(step_gray_scott!)

    # 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:50, format="{}x", width=750, startvalue=stepsize[], tellwidth=false))

    speed_slider = slidergrid.sliders[1].value

    gh[1, 2] = templategrid = GridLayout(ax.scene, tellwidth=false)
    templategrid[1, 1] = Label(fig, "Templates:")
    btn_waves = Button(templategrid[1, 2], width=100, label="Wave Pattern")
    btn_cow = Button(templategrid[1, 3], width=100, label="Cow Spots")
    btn_cheetah = Button(templategrid[1, 4], width=100, label="Cheetah Spots")

    # place all the parameter boxes
    gh[2, 2] = textboxgrid = GridLayout(ax.scene, tellwidth=false)

    # Create and assign column header labels
    textboxgrid[1, 1] = Label(fig, "GrayScott:", halign=:center)
    textboxgrid[1, 3] = Label(fig, "FHN:", halign=:center)

    # GrayScott column (col 1)
    param_box!(textboxgrid, 2, "Du", param_obs_map.Du, col=1)
    param_box!(textboxgrid, 3, "Dv", param_obs_map.Dv, col=1)
    param_box!(textboxgrid, 4, "Feed", param_obs_map.F, col=1)
    param_box!(textboxgrid, 5, "Kill", param_obs_map.k, col=1)

    # FHN column (col 2)
    param_box!(textboxgrid, 2, "Du", param_obs_map.Du, col=2)
    param_box!(textboxgrid, 3, "Dv", param_obs_map.Dv, col=2)
    param_box!(textboxgrid, 4, "ϵ", param_obs_map.ϵ, col=2)
    param_box!(textboxgrid, 5, "a", param_obs_map.a, col=2)
    param_box!(textboxgrid, 6, "b", param_obs_map.b, col=2)

    rowsize!(gh, 1, Relative(0.2))  # small row for the menu
    rowsize!(gh, 2, Relative(0.8))
    for c in 1:4
        colsize!(textboxgrid, c, Relative(0.1))
    end

    # ======== Events ========
    # Timer and state for animation
    running = Observable(false)

    on(running) do r
        run_label[] = r ? "Pause" : "Run"
    end
    on(speed_slider) do s
        try
            stepsize[] = s
            println("Changed stepsize to $s")
        catch
            @warn "Invalid input for $s"
        end
    end
    # Control Listeners
    on(btn_step.clicks) do _
        multi_step!((U, V), stepsize[] * 5, heat_obs, params_obs; step_method=step_method[])
    end

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

    on(btn_start.clicks) do _
        @async while running[]
            multi_step!((U, V), stepsize[], heat_obs, params_obs; step_method=step_method[])
            sleep(0.0015)  # ~20 FPS
        end
    end

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

    end

    # Template Control
    on(btn_waves.clicks) do _
        # add column to center of matrix
        U, V = Templates.blocks_ic(params_obs[].N)

        hm.colormap[] = :viridis

        # change params
        param_obs_map.Du[] = 0.0008
        param_obs_map.Dv[] = 0.1
        param_obs_map.ϵ[] = 0.01
        param_obs_map.a[] = 0.5
        param_obs_map.b[] = 0.15
        heat_obs[] = copy(U)
    end

    # Template Control
    on(btn_cow.clicks) do _
        # fill matrix with random noise
        U = 0.05 .* (2 .* rand(params_obs[].N, params_obs[].N) .- 1)  # noise in [-0.05, 0.05]
        V = 0.05 .* (2 .* rand(params_obs[].N, params_obs[].N) .- 1)

        hm.colormap[] = :grayC

        # change params
        param_obs_map.Du[] = 1.0
        param_obs_map.Dv[] = 10.0
        param_obs_map.ϵ[] = 0.01
        param_obs_map.a[] = -0.1
        param_obs_map.b[] = 1.2
        heat_obs[] = copy(U)
    end

    on(btn_cheetah.clicks) do _
        # fill matrix with random noise
        U = 0.05 .* (2 .* rand(params_obs[].N, params_obs[].N) .- 1)  # noise in [-0.05, 0.05]
        V = 0.05 .* (2 .* rand(params_obs[].N, params_obs[].N) .- 1)

        hm.colormap[] = cgrad(:lajolla, rev=true)

        # change params
        param_obs_map.Du[] = 0.182
        param_obs_map.Dv[] = 0.5
        param_obs_map.ϵ[] = 0.05
        param_obs_map.a[] = 0.2
        param_obs_map.b[] = 2.0
        heat_obs[] = copy(U)
    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

    on(dropdown.selection) do sel
        if sel isa Tuple
            label, selected_model = sel
        else
            selected_model = sel
            label = string(sel)
        end

        @info "Selected model: $label → $selected_model"

        if haskey(STEP_FUNCTIONS, selected_model)
            step_method[] = STEP_FUNCTIONS[selected_model]
        else
            @warn "Unknown model selected: $selected_model"
        end
    end

    return fig
end

export build_ui
end