using CairoMakie
using LaTeXStrings: @L_str
using StaticArrays: SVector
using JLD2: JLD2
using Dates: Dates
using CSV: CSV
using DataFrames: DataFrames
using Random: Random

using ReCo: ReCo

include("../../visualization/common_CairoMakie.jl")

function radial_distribution_simulation(;
    n_particles::Int64, v₀s::NTuple{N,Float64}, T::Float64, packing_fraction::Float64
) where {N}
    Random.seed!(42)

    n_v₀s = length(v₀s)

    sim_dirs = Vector{String}(undef, n_v₀s)

    parent_dir = "radial_distribution_$(Dates.now())"

    Threads.@threads for v₀_ind in 1:n_v₀s
        v₀ = v₀s[v₀_ind]

        sim_dir = ReCo.init_sim(;
            n_particles=n_particles,
            v₀=v₀,
            packing_fraction=packing_fraction,
            parent_dir=parent_dir,
            comment="$v₀",
        )

        ReCo.run_sim(sim_dir; duration=T, seed=v₀_ind)

        sim_dirs[v₀_ind] = sim_dir
    end

    return sim_dirs
end

function circular_shell_volume(lower_radius, Δradius)
    return π * (2 * lower_radius * Δradius + Δradius^2)
end

function radial_distribution(;
    sim_dirs::Vector{String}, n_radii::Int64, n_last_snapshots::Int64, n_particles::Int64
)
    sim_consts = ReCo.load_sim_consts(sim_dirs[1])

    particle_radius = sim_consts.particle_radius

    min_lower_radius = 0.0
    max_lower_radius = 5 * (2 * particle_radius)
    Δradius = (max_lower_radius - min_lower_radius) / n_radii

    lower_radii = LinRange(min_lower_radius, max_lower_radius, n_radii)

    box_volume = (2 * sim_consts.half_box_len)^2
    volume_per_particle = box_volume / n_particles

    n_sim_dirs = length(sim_dirs)

    gs = Vector{Vector{Float64}}(undef, n_sim_dirs)

    Threads.@threads for sim_dir_ind in 1:n_sim_dirs
        sim_dir = sim_dirs[sim_dir_ind]

        cs = Matrix{SVector{2,Float64}}(undef, (n_particles, n_last_snapshots))

        bundle_paths = ReCo.sorted_bundle_paths(sim_dir; rev=true)

        snapshot_counter = 0
        break_bundle_path_loop = false

        for bundle_path in bundle_paths
            bundle::ReCo.Bundle = JLD2.load_object(bundle_path)

            for snapshot_ind in (bundle.n_snapshots):-1:1
                snapshot_counter += 1

                @simd for particle_ind in 1:n_particles
                    cs[particle_ind, snapshot_counter] = bundle.c[
                        particle_ind, snapshot_ind
                    ]
                end

                if snapshot_counter == n_last_snapshots
                    break_bundle_path_loop = true
                    break
                end
            end

            if break_bundle_path_loop
                break
            end
        end

        if snapshot_counter != n_last_snapshots
            error("snapshot_counter != n_last_snapshots")
        end

        g = zeros(Float64, n_radii)

        for snapshot_ind in 1:n_last_snapshots
            for p1_ind in 1:n_particles
                for p2_ind in (p1_ind + 1):n_particles
                    c1 = cs[p1_ind, snapshot_ind]
                    c2 = cs[p2_ind, snapshot_ind]

                    r⃗₁₂ = c2 - c1
                    r⃗₁₂ = ReCo.restrict_coordinates(r⃗₁₂, sim_consts.half_box_len)

                    distance = ReCo.norm2d(r⃗₁₂)

                    lower_radius_ind = ceil(Int64, distance / Δradius)

                    if lower_radius_ind <= n_radii
                        g[lower_radius_ind] += 2
                    end
                end
            end
        end

        for (lower_radius_ind, lower_radius) in enumerate(lower_radii)
            g[lower_radius_ind] *=
                volume_per_particle / (
                    n_last_snapshots *
                    n_particles *
                    circular_shell_volume(lower_radius, Δradius)
                )
        end

        gs[sim_dir_ind] = g
    end

    return (lower_radii, gs, particle_radius)
end

function plot_radial_distributions(;
    v₀s::NTuple{N,Float64},
    lower_radii::AbstractVector{Float64},
    gs::Vector{Vector{Float64}},
    particle_radius::Float64,
    include_comparison::Bool,
    filename_addition::String="",
) where {N}
    println("Plotting the radial distributions")

    init_cairomakie!()

    fig = gen_figure()

    min_lower_radius = 0.0
    max_lower_radius = maximum(lower_radii)

    min_g = 0.0
    max_g = maximum(maximum.(gs))

    ax = Axis(
        fig[1, 1];
        xticks=0:(2 * particle_radius):floor(Int64, max_lower_radius),
        yticks=0:ceil(Int64, max_g),
        xlabel=L"r / d",
        ylabel=L"g(r)",
        limits=(min_lower_radius - 0.03, max_lower_radius + 0.03, min_g, max_g * 1.03),
    )

    lines!(
        ax,
        SVector(min_lower_radius, max_lower_radius),
        SVector(1.0, 1.0);
        linestyle=:dash,
        color=:red,
        linewidth=1,
    )

    for (g_ind, g) in enumerate(gs)
        v₀ = round(Int64, v₀s[g_ind])
        scatterlines!(
            ax,
            lower_radii,
            g;
            markersize=3,
            linewidth=1,
            label=L"Simulation with $v_0 = %$(v₀)$",
        )
    end

    if include_comparison
        comparison_curve = CSV.read(
            "analysis/radial_distribution_function/g_of_r_d_ratio_with_0_37_packing_fraction.csv",
            DataFrames.DataFrame;
            header=3,
        )
        lines!(
            ax,
            comparison_curve.r_d_ratio,
            comparison_curve.g_of_r_d_ratio;
            label=L"Reference with $v_0 = 0$",
            color=:orange,
        )
    end

    axislegend(ax; position=:rt, padding=3, rowgap=-3)

    save_fig("radial_distribution$filename_addition.pdf", fig)

    return nothing
end

function run_radial_distribution_analysis()
    v₀s = (0.0, 40.0, 80.0)

    n_particles = 1000

    sim_dirs = radial_distribution_simulation(;
        n_particles=n_particles, v₀s=v₀s, T=100.0, packing_fraction=0.37
    )

    lower_radii, gs, particle_radius = radial_distribution(;
        sim_dirs=sim_dirs, n_radii=58, n_last_snapshots=200, n_particles=n_particles
    )

    plot_radial_distributions(;
        v₀s=v₀s[1:1], lower_radii, gs=gs[1:1], particle_radius, include_comparison=true
    )

    plot_radial_distributions(;
        v₀s=v₀s[1:end],
        lower_radii,
        gs=gs[1:end],
        particle_radius,
        include_comparison=false,
        filename_addition="_all_vs",
    )

    return nothing
end