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

using ReCo: ReCo

includet("../src/Visualization/common_CairoMakie.jl")

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

const δt = 1e-4
const Dₜ = ReCo.DEFAULT_Dₜ

function max_possible_displacement(T::Float64, v₀::Float64, δt::Float64=δt, Dₜ::Float64=Dₜ)
    return T * v₀ + T / δt * sqrt(2 * Dₜ * δt)
end

function msd_simulation(
    v₀::Float64,
    half_box_len::Float64,
    T::Float64,
    snapshot_at::Float64,
    parent_dir::String;
    comment::String="",
    seed::Int64=42,
)
    Random.seed!(seed)

    dir = ReCo.init_sim(;
        n_particles=1,
        v₀=v₀,
        parent_dir=parent_dir,
        comment=comment,
        half_box_len=half_box_len,
        δt=δt,
    )

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

    return dir
end

function mean_squared_displacement(;
    n_simulations::Int64, v₀s::NTuple{N,Float64}, T::Float64
) where {N}
    n_v₀s = length(v₀s)

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

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

    progress = ProgressMeter.Progress(n_v₀s * n_simulations; dt=3, desc="MSD: ")

    for (v₀_ind, v₀) in enumerate(v₀s)
        half_box_len = max_possible_displacement(T, v₀)

        parent_dir = main_parent_dir * "/$v₀"

        Threads.@threads for sim_ind in 1:n_simulations
            dir = msd_simulation(v₀, half_box_len, T, 0.5, parent_dir; comment="$sim_ind")

            sim_dirs[sim_ind, v₀_ind] = dir

            ProgressMeter.next!(progress; showvalues=[(:v₀, v₀)])
        end
    end

    ts = Float64[]
    ReCo.append_bundle_properties!(
        (ts,), (:t,), sim_dirs[1, 1]; particle_slice=1, snapshot_slice=:, first_bundle=2
    ) # Skip the first bundle to avoid t = 0

    mean_sq_displacements = zeros(Float64, (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) # Skip the first bundle to avoid t = 0
                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::NTuple{N,Float64}
) where {N}
    init_cairomakie!()

    fig = gen_figure()

    ax = Axis(fig[1, 1]; xlabel=L"t", ylabel=L"\mathbf{MSD}(t)", 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_0 = %$v₀" for v₀ in v₀s])

    set_gaps!(fig)

    save_fig("mean_squared_displacement.pdf", fig)

    return nothing
end

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

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

    plot_mean_sq_displacement_with_expectation(ts, mean_sq_displacements, v₀s)

    return nothing
end

function plot_random_walk(; T::Float64, v₀::Float64, seed::Int64)
    half_box_len = max_possible_displacement(T, v₀)
    dir = msd_simulation(v₀, half_box_len, T, 8.0, "random_walk_$(Dates.now())"; seed=seed)

    ts = Float64[]
    cs = SVector{2,Float64}[]
    ReCo.append_bundle_properties!(
        (ts, cs), (:t, :c), dir; particle_slice=1, snapshot_slice=:
    )

    init_cairomakie!()

    fig = gen_figure()

    max_displacement = maximum(ReCo.norm2d.(cs))
    expected_mean_displacement = sqrt(expected_mean_squared_displacement(ts[end], v₀))

    limit = max(max_displacement, expected_mean_displacement)
    limit *= 1.04

    ax = Axis(
        fig[1, 1];
        xlabel=L"x",
        ylabel=L"y",
        aspect=AxisAspect(1),
        limits=(-limit, limit, -limit, limit),
    )

    path_plot = lines!(ax, cs; linewidth=0.5)

    marker_size = 4

    start_point_plot = scatter!(ax, cs[1]; markersize=marker_size, color=:lime)

    end_point_plot = scatter!(ax, cs[end]; markersize=marker_size, color=:red)

    expected_mean_squared_displacement_plot = arc!(
        ax,
        Point2(0.0, 0.0),
        expected_mean_displacement,
        0,
        2 * π;
        color=:orange,
        linewidth=1,
    )

    Legend(
        fig[1, 2],
        [
            start_point_plot,
            end_point_plot,
            path_plot,
            expected_mean_squared_displacement_plot,
        ],
        ["Start point", "End point", "Path", "Expected mean displacement"],
    )

    set_gaps!(fig)

    save_fig("random_walk.pdf", fig)

    return nothing
end

function run_random_walk()
    plot_random_walk(; T=100_000.0, v₀=0.0, seed=12345)

    return nothing
end