infection_sim.jl

using Random
using Printf
import GLMakie

# -----------------------------
# Infection simulation
# -----------------------------
const EMPTY = 0
const SUSC  = 1
const INF   = 2

@inline function neighbor_offsets(allow_diagonal_infections::Bool)
    neighbors4 = ((-1,0), (1,0), (0,-1), (0,1))
    neighbors8 = ((-1,-1), (-1,0), (-1,1),
                  ( 0,-1),         ( 0,1),
                  ( 1,-1), ( 1,0), ( 1,1))
    return allow_diagonal_infections ? neighbors8 : neighbors4
end

"Return an iterator over valid neighbor indices (ii, jj) of (i, j)."
@inline function neighbors(i::Int, j::Int, n::Int, allow_diagonal_infections::Bool)
    neigh = neighbor_offsets(allow_diagonal_infections)
    return ((i+di, j+dj) for (di, dj) in neigh
            if 1 ≤ i+di ≤ n && 1 ≤ j+dj ≤ n)
end

function init_grid(n, density)
    grid = fill(EMPTY, n, n)
    occ  = rand(n, n) .< density # fill with 1
    grid[occ] .= SUSC
    return grid
end

function seed_infection!(grid)
    healthy_inds = findall(==(SUSC), grid)
    isempty(healthy_inds) && return grid
    grid[rand(healthy_inds)] = INF
    return grid
end

"try to advance the infection front: look at every cell that is infected 
 and if it is next to susceptible cell, try to infect it with p_infect"
function step(grid::Array{Int,2}, p_infect;
              allow_diagonal_infections::Bool = false)

    n = size(grid, 1)
    new_grid = copy(grid)

    for i in 1:n, j in 1:n
        if grid[i,j] == INF # For each infected cell, try to infect its neighbors
            for (ii, jj) in neighbors(i, j, n, allow_diagonal_infections)
                if grid[ii,jj] == SUSC && rand() < p_infect
                    new_grid[ii,jj] = INF
                end
            end
        end
    end

    return new_grid
end

"Return true if *any* infected cell has a susceptible neighbor."
function can_still_spread(grid::Array{Int,2};
                          allow_diagonal_infections::Bool = false)

    n = size(grid, 1)

    for i in 1:n, j in 1:n
        if grid[i,j] == INF
            for (ii, jj) in neighbors(i, j, n, allow_diagonal_infections)
                if grid[ii,jj] == SUSC
                    return true
                end
            end
        end
    end

    return false
end


"Simulate starting from a given initial grid (no infection yet),
 allowing retries until spread is impossible or max_steps is reached."
function simulate_from_grid(initial_grid,
                            p_infect;
                            max_steps,
                            allow_diagonal_infections = false)

    grid0 = copy(initial_grid)
    seed_infection!(grid0)

    states = [grid0]
    infected_counts = [count(==(INF), grid0)]

    for _ in 1:max_steps
        new_grid = step(states[end], p_infect; allow_diagonal_infections = allow_diagonal_infections)

        push!(states, new_grid)
        push!(infected_counts, count(==(INF), new_grid))

        if !can_still_spread(new_grid;
                             allow_diagonal_infections = allow_diagonal_infections)
            break
        end
    end

    return states, infected_counts
end

# -----------------------------
# UI + animation
# -----------------------------
function infection_ui(; n = 60, max_steps = 300)
    density0  = 0.6
    p_infect0 = 0.3

    initial_grid = GLMakie.Observable(init_grid(n, density0))
    data         = GLMakie.Observable(initial_grid[])

    infected_x = GLMakie.Observable(Float64[])
    infected_y = GLMakie.Observable(Float64[])

    fig = GLMakie.Figure(size = (1100, 700))

    # --- top row: grid + time series ---
    heat_ax = GLMakie.Axis(fig[1, 1];
        aspect = GLMakie.DataAspect(),
        xticksvisible = false,
        yticksvisible = false,
        xgridvisible  = false,
        ygridvisible  = false,
        xlabelvisible = false,
        ylabelvisible = false,
        title = "Infection Spread"
    )
    GLMakie.heatmap!(heat_ax, data; colormap = [:white,:green,:red], colorrange = (0,2))

    ts_ax = GLMakie.Axis(fig[1, 2];
        title  = "Infected count over time",
        xlabel = "Step",
        ylabel = "# infected",
        limits = ((0.0, max_steps |> float), (0.0, n^2 |> float)),
    )
    GLMakie.lines!(ts_ax, infected_x, infected_y)

    # --- second row: controls ---
    controls = GLMakie.GridLayout()
    fig[2, 1:2] = controls

    # Area density
    GLMakie.Label(controls[1, 1], "Population density:")
    density_slider = GLMakie.Slider(controls[1, 2],
        range = 0.0:0.01:1.0, startvalue = density0)
    GLMakie.Label(controls[1, 3],
        GLMakie.lift(x -> @sprintf("%.2f", x), density_slider.value))

    # Infection probability
    GLMakie.Label(controls[2, 1], "Infection probability:")
    p_slider = GLMakie.Slider(controls[2, 2],
        range = 0.0:0.01:1.0, startvalue = p_infect0)
    GLMakie.Label(controls[2, 3],
        GLMakie.lift(x -> @sprintf("%.2f", x), p_slider.value))

    # Buttons + diagonal checkbox
    restart_button = GLMakie.Button(controls[3, 1], label = "Restart grid")
    run_button     = GLMakie.Button(controls[3, 2], label = "Run simulation")

    diag_checkbox  = GLMakie.Checkbox(controls[3, 3]; checked = false)
    GLMakie.Label(controls[3, 4], "Allow diagonal infection")

    GLMakie.display(fig)

    # --- Restart grid ---
    GLMakie.on(restart_button.clicks) do _
        density = density_slider.value[]
        initial_grid[] = init_grid(n, density)
        data[] = initial_grid[]
        infected_x[] = Float64[]
        infected_y[] = Float64[]
        heat_ax.title[] = "Infection Spread (reset)"
    end

    # --- Run simulation ---
    GLMakie.on(run_button.clicks) do _
        p_inf      = p_slider.value[]
        allow_diag = diag_checkbox.checked[]
        grid0      = initial_grid[]

        states, counts = simulate_from_grid(
            grid0, p_inf;
            max_steps = max_steps,
            allow_diagonal_infections = allow_diag,
        )

        @async begin
            infected_x[] = Float64[]
            infected_y[] = Float64[]
            heat_ax.title[] = "Infection Spread"

            for (t, g) in enumerate(states)
                data[] = g
                infected_x[] = collect(0:(t-1)) |> x -> Float64.(x)
                infected_y[] = Float64.(counts[1:t])
                sleep(0.03)
            end
        end
    end

    return fig
end

infection_ui()