using Pkg
Pkg.activate("src/plots")
Pkg.instantiate()

# Load packages & code libraries
using DataFrames
using CSV
using Pipe
using Statistics
using CairoMakie
using QuantEcon
using LaTeXStrings

include("lib.jl")

function figures_a(df::DataFrame, param::Symbol, i::Int)
    vars = [:ON_interbank, :Term_interbank]

    # take only bank agents
    df1 = @pipe df |> dropmissing(_, vars) |> filter(:status => x -> x != "neutral", _) |> groupby(_, [:step, param]) |> 
        combine(_, vars .=> mean, renamecols = false)

    p = generic_plot([df1], [:ON_interbank, :Term_interbank], ["Overnight loans", "Term loans"], param)
    save("FigB$(i)a.pdf", p)
end

function figures_b(df::DataFrame, m::DataFrame, param::Symbol, i::Int)
    df_hh = @pipe df |> filter(:id => x -> x >= 1 && x <= mean(m[!, :n_hh]), _) |>
        groupby(_, [:step, param]) |> 
        combine(_, [:loans] .=> mean, renamecols = false)
   
    df_firms = @pipe df |>  filter(:id => x -> x > mean(m[!, :n_hh]) && x <= mean(m[!, :n_hh]) + mean(m[!, :n_f]), _) |>
        groupby(_, [:step, param]) |> 
        combine(_, [:loans] .=> mean, renamecols = false)

    p = generic_plot([df_firms, df_hh], [:loans], ["Firms loans", "Households loans"], param)
    save("FigB$(i)b.pdf", p)
end

function load_data()
    # parameter names as strings
    params_strings = ["r", "γ", "pref", "g"]
    # parameter ranges
    params_range = ["0.9", "1.1", "1.3", 
        "0.0", "0.1", "0.3", 
        "0.01", "0.05", "0.09", 
        "100.0", "200.0", "300.0"]

    df = DataFrame()

    for param in params_strings
        if occursin(param, params_strings[1])
            range_indices = 1:3
        elseif occursin(param, params_strings[2])
            range_indices = 4:6
        elseif occursin(param, params_strings[3])
            range_indices = 7:9
        elseif occursin(param, params_strings[4])
            range_indices = 10:length(params_range)
        end
        
        for val in params_range[range_indices]
            append!(df, CSV.File("data/sensitivity_analysis/$(param)/$(val)/df.csv"); cols = :union, promote = true)
        end
    end
    
    # take model variables from Baseline scenario
    mdf = DataFrame()
    append!(mdf, CSV.File("data/shock=Missing/Baseline/mdf.csv"))

    return df, mdf
end

function create_sens_plots(df, mdf)
    # parameter names as symbols
    param_symbols = [:r, :γ, :pref, :g]

    cd(mkpath("img/pdf/")) do
        for (index, param) in pairs(param_symbols)
            figures_a(filter(param => x -> !ismissing(x), df), param, index)
            figures_b(filter(param => x -> !ismissing(x), df), mdf, param, index)
        end
    end
    printstyled("Sensitivity plots generated."; color = :blue)
end

df, mdf = load_data()
create_sens_plots(df, mdf)