using CairoMakie
using LaTeXStrings: @L_str

using ReCo: ReCo

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

function run_reward_discount_analysis()
    γs = 0.0:0.2:1.0
    n_γs = length(γs)
    γ_rewards = Vector{Vector{Float64}}(undef, n_γs)

    Threads.@threads for γ_ind in 1:n_γs
        γ = γs[γ_ind]
        env_helper = ReCo.run_rl(;
            EnvType=ReCo.OriginEnv,
            n_episodes=3,
            episode_duration=15.0,
            n_particles=150,
            update_actions_at=0.08,
            ϵ_stable=0.00001,
            process_dir="reward_discount_analysis/$γ_ind",
            reward_discount=γ,
            show_simulation_progress=false,
        )

        rewards = env_helper.shared.hook.rewards
        γ_rewards[γ_ind] = rewards
    end

    init_cairomakie!()

    fig = gen_figure()

    ax = Axis(fig[1, 1]; xlabel="Episode", ylabel="Reward")

    rewards_plots = []
    for rewards in γ_rewards
        rewards_plot = lines!(ax, 1:length(rewards), rewards)
        push!(rewards_plots, rewards_plot)
    end

    Legend(fig[1, 2], rewards_plots, [L"\gamma = %$γ" for γ in γs])

    set_gaps!(fig)

    save_fig("reward_discount_analysis.pdf", fig)

    return nothing
end