module Visualization

using GLMakie
"""
    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)

    # Initialize heatmap with first time step
    u0 = reshape(sol[1][1:N^2], N, N)
    heat_obs = Observable(u0)
    hmap = heatmap!(ax, heat_obs, colormap=:RdGy)

    # 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

end