using CairoMakie
using LaTeXStrings: @L_str
using Dates: Dates
using Random: Random
using StaticArrays: SVector
using JLD2: JLD2
using CellListMap: CellListMap

using ReCo: ReCo

# IMPORTANT: Disable the periodic boundary conditions
# The arguments types have to match for the function to be overwritten!
ReCo.update_verlet_lists!(::Any, ::CellListMap.CellList) = nothing
ReCo.gen_cell_list(::Vector{SVector{2,Float64}}, ::CellListMap.Box) = nothing
ReCo.gen_cell_list_box(::Float64, ::Float64) = nothing
ReCo.push_to_verlet_list!(::Any, ::Any, ::Any) = nothing
ReCo.restrict_coordinate(value::Float64, ::Float64) = value
ReCo.restrict_coordinates(v::SVector{2,Float64}, ::Float64) = v
ReCo.restrict_coordinates!(::Particle, ::Float64) = nothing
ReCo.minimum_image_coordinate(value::Float64, ::Float64) = value
ReCo.minimum_image(v::SVector{2,Float64}, ::Float64) = v

function mean_squared_displacement(;
    n_simulations::Int64, v₀s::AbstractVector{Float64}, T::Float64
)
    Random.seed!(42)

    n_v₀s = length(v₀s)

    δt = ReCo.DEFAULT_δt
    Dₜ = ReCo.DEFAULT_Dₜ

    main_parent_dir = "mean_squared_displacement_$(Dates.now())"
    mkpath(main_parent_dir)

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

    for (v₀_ind, v₀) in enumerate(v₀s)
        max_possible_displacement = T * v₀ + T / δt * sqrt(2 * Dₜ * δt)

        parent_dir = main_parent_dir * "/$v₀"

        Threads.@threads for sim_ind in 1:n_simulations
            dir = ReCo.init_sim(;
                n_particles=1,
                v₀=v₀,
                parent_dir=parent_dir,
                comment="$sim_ind",
                half_box_len=max_possible_displacement,
            )

            sim_dirs[sim_ind, v₀_ind] = dir

            ReCo.run_sim(
                dir;
                duration=T,
                seed=rand(1:typemax(Int64)),
                snapshot_at=0.01,
                n_bundle_snapshots=200,
            )
        end
    end

    ts = Float64[]

    bundle_paths = ReCo.sorted_bundle_paths(sim_dirs[1, 1])

    for i in 2:length(bundle_paths) # Skip the first bundle to avoid t = 0
        bundle::ReCo.Bundle = JLD2.load_object(bundle_paths[i])

        append!(ts, bundle.t)
    end

    mean_sq_displacements = zeros((length(ts), n_v₀s))

    @simd for v₀_ind in 1:n_v₀s
        for sim_ind in 1:n_simulations
            sim_dir = sim_dirs[sim_ind, v₀_ind]

            bundle_paths = ReCo.sorted_bundle_paths(sim_dir)
            snapshot_ind = 1

            for i in 2:length(bundle_paths)
                bundle::ReCo.Bundle = JLD2.load_object(bundle_paths[i])

                for c in bundle.c
                    sq_displacement = ReCo.sq_norm2d(c)
                    mean_sq_displacements[snapshot_ind, v₀_ind] += sq_displacement

                    snapshot_ind += 1
                end
            end
        end
    end

    mean_sq_displacements ./= n_simulations

    return (ts, mean_sq_displacements)
end

function expected_mean_squared_displacement(t::Float64, v₀::Float64)
    Dₜ = ReCo.DEFAULT_Dₜ
    particle_radius = ReCo.DEFAULT_PARTICLE_RADIUS

    Dᵣ = 3 * Dₜ / ((2 * particle_radius)^2)

    return (4 * Dₜ + 2 * v₀^2 / Dᵣ) * t + 2 * v₀^2 * (exp(-Dᵣ * t) - 1) / (Dᵣ^2)
end

function plot_mean_sq_displacement_with_expectation(
    ts::Vector{Float64},
    mean_sq_displacements::Matrix{Float64},
    v₀s::AbstractVector{Float64},
)
    CairoMakie.activate!()
    set_theme!()

    text_width_in_pt = 405

    fig = Figure(;
        resolution=(text_width_in_pt, 0.55 * text_width_in_pt),
        fontsize=10,
        figure_padding=1,
    )

    ax = Axis(fig[1, 1]; xlabel=L"t", ylabel=L"\mathbf{MSD}", xscale=log10, yscale=log10)

    t_linrange = LinRange(ts[1], ts[end], 1000)

    v₀_scatter_plots = []

    for (v₀_ind, v₀) in enumerate(v₀s)
        scatter_plot = scatter!(
            ax, ts, view(mean_sq_displacements, :, v₀_ind); markersize=4
        )

        push!(v₀_scatter_plots, scatter_plot)

        expected_mean_sq_displacements = expected_mean_squared_displacement.(t_linrange, v₀)

        lines!(ax, t_linrange, expected_mean_sq_displacements)
    end

    Legend(fig[1, 2], v₀_scatter_plots, [L"v₀ = %$v₀" for v₀ in v₀s])

    colgap!(fig.layout, 5)
    rowgap!(fig.layout, 5)

    parent_dir = "exports/graphics/"
    mkpath(parent_dir)

    save("$parent_dir/mean_squared_displacement.pdf", fig; pt_per_unit=1)

    return nothing
end

function run_analysis()
    v₀s = SVector(0.0, 20.0, 40.0, 60.0, 80.0)

    ts, mean_sq_displacements = mean_squared_displacement(;
        n_simulations=3 * Threads.nthreads(), v₀s=v₀s, T=10.0
    )

    plot_mean_sq_displacement_with_expectation(ts, mean_sq_displacements, v₀s)

    return nothing
end

# run_analysis()