module SnapshotPlot

export plot_snapshot

using CairoMakie
using ColorSchemes: ColorSchemes
using LaTeXStrings: @L_str

using ..ReCo: ReCo

include("common.jl")
include("common_CairoMakie.jl")

const DEFAULT_SHOW_CENTER_OF_MASS = false
const DEFAULT_SHOW_KAPPA = true

function get_wanted_snapshot_number(total_n_snapshots::Int64)
    print("There are $total_n_snapshots snapshots. Enter the wanted snapshot number: ")
    answer = readline()

    snapshot = parse(Int64, answer)

    return snapshot
end

"""
    plot_snapshot(sim_dir::String; <keyword arguments>)

Plot one snapshot of a simulation.

The function will ask for the number of the snapshot to plot out of the total number of snapshots. The output is `sim_dir/graphics/N.pdf` with `N` as the number of the chosen snapshot.

# Arguments
- `sim_dir::String`: Simulation directory.
- `show_center_of_mass::Bool=$DEFAULT_SHOW_CENTER_OF_MASS`: Show the center of mass as a transparent black circle.
- `show_κ::Bool=$DEFAULT_SHOW_KAPPA`: Show κ as the ratio of the eigenvalues of the gyration tensor.
"""
function plot_snapshot(
    sim_dir::String;
    show_center_of_mass::Bool=DEFAULT_SHOW_CENTER_OF_MASS,
    show_κ::Bool=DEFAULT_SHOW_KAPPA,
)
    bundles_info = ReCo.BundlesInfo(sim_dir)
    wanted_snapshot_out_of_total = get_wanted_snapshot_number(
        bundles_info.total_n_snapshots
    )

    bundle, bundle_snapshot = ReCo.get_bundle_to_snapshot(
        bundles_info, wanted_snapshot_out_of_total
    )

    sim_consts = ReCo.load_sim_consts(sim_dir)

    println("Initializing CairoMakie...")

    init_cairomakie!()
    fig = gen_figure()

    cs_view = view(bundle.c, :, bundle_snapshot)
    center_of_mass = ReCo.center_of_mass(cs_view, sim_consts.half_box_len)

    if show_κ
        eigvals_ratio = ReCo.gyration_tensor_eigvals_ratio(
            cs_view, sim_consts.half_box_len, center_of_mass
        )
        κ = round(eigvals_ratio; digits=2)
        title = L"\kappa = %$κ"

        ax, color_scheme = gen_axis_and_colorbar(fig, sim_consts; axis_title=title)
    else
        ax, color_scheme = gen_axis_and_colorbar(fig, sim_consts)
    end

    particle_xs = Vector{Float64}(undef, sim_consts.n_particles)
    particle_ys = Vector{Float64}(undef, sim_consts.n_particles)
    particle_colors = Vector{ColorSchemes.ColorTypes.RGB}(undef, sim_consts.n_particles)

    @simd for particle_ind in 1:(sim_consts.n_particles)
        c = bundle.c[particle_ind, bundle_snapshot]
        φ = bundle.φ[particle_ind, bundle_snapshot]

        color = angle_color(φ, color_scheme)

        particle_xs[particle_ind] = c[1]
        particle_ys[particle_ind] = c[2]
        particle_colors[particle_ind] = color
    end

    scatter!(
        ax,
        particle_xs,
        particle_ys;
        markersize=2 * sim_consts.particle_radius,
        markerspace=SceneSpace,
        color=particle_colors,
    )

    if show_center_of_mass
        scatter!(
            ax,
            Point(center_of_mass);
            markersize=6 * sim_consts.particle_radius,
            markerspace=SceneSpace,
            color=ColorSchemes.ColorTypes.RGBA(0.0, 0.0, 0.0, 0.6),
        )
    end

    set_gaps!(fig)

    save_fig("$wanted_snapshot_out_of_total.pdf", fig; parent_dir="$sim_dir/graphics")

    return nothing
end

end # module