using CairoMakie
using LaTeXStrings: @L_str
using Statistics: Statistics

using ReCo: ReCo

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

function plot_mean_kappa(; rl_dir::String, n_last_episodes::Int64)
    dir_content = readdir(rl_dir; join=true, sort=true)
    n_content = length(dir_content)

    sim_dirs = Vector{String}(undef, n_last_episodes)
    sim_dir_counter = 1

    # Skip first sim_dir for the case that the simulation is still running
    skipped_first_sim_dir = false

    for file_or_dir_ind in n_content:-1:1
        file_or_dir = dir_content[file_or_dir_ind]

        if isdir(file_or_dir)
            if skipped_first_sim_dir
                sim_dirs[sim_dir_counter] = file_or_dir
                sim_dir_counter += 1
                if sim_dir_counter > n_last_episodes
                    break
                end
            else
                skipped_first_sim_dir = true
            end
        end
    end

    if sim_dir_counter < n_last_episodes
        error("The rl_dir does not have n_last_episodes + 1 dirs!")
    end

    sim_consts = ReCo.load_sim_consts(sim_dirs[1])
    half_box_len = sim_consts.half_box_len
    total_n_snapshots = ReCo.BundlesInfo(sim_dirs[1]).total_n_snapshots

    snapshot_κs = zeros(Float64, total_n_snapshots)

    for sim_dir in sim_dirs
        bundles_info = ReCo.BundlesInfo(sim_dir)
        total_n_snapshots = bundles_info.total_n_snapshots

        for snapshot_ind in 1:total_n_snapshots
            bundle, bundle_snapshot = ReCo.get_bundle_to_snapshot(
                bundles_info, snapshot_ind
            )

            cs_view = view(bundle.c, :, bundle_snapshot)

            κ = ReCo.gyration_tensor_eigvals_ratio(cs_view, half_box_len)
            snapshot_κs[snapshot_ind] += κ
        end
    end

    snapshot_κs ./= n_last_episodes

    mean_κ = Statistics.mean(snapshot_κs)

    init_cairomakie!()
    fig = gen_figure(; padding=9)

    if n_last_episodes > 1
        title = "Averaged over last $n_last_episodes episodes"
    else
        title = "Result of only one episode"
    end

    ax = Axis(
        fig[1, 1];
        xlabel="Frame",
        ylabel=L"\kappa",
        limits=(1, total_n_snapshots, 0.0, 1.04),
        title=title,
    )

    lines!(ax, 1:total_n_snapshots, snapshot_κs; label=L"\kappa")

    rounded_mean_κ = round(mean_κ; digits=2)

    lines!(
        ax,
        [1, total_n_snapshots],
        [mean_κ, mean_κ];
        label=L"Mean $\tilde{\kappa} = %$rounded_mean_κ$",
        linestyle=:dash,
        color=:red,
    )

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

    set_gaps!(fig)

    save_fig("mean_kappa.pdf", fig; parent_dir=rl_dir)

    return nothing
end