using Pkg
Pkg.activate(".")
Pkg.instantiate()

# Load packages on master process
using Pipe
using CSV
using DataFrames
using Agents
using Random
using StatsBase
using Distributions
using CViM

# runs the sensitivity analysis for the parameters of interests, transforms and collects datas
function run_sens(seeds::Vector{UInt32}, param::Symbol, param_range::Vector{Float64}; shock::String = "Missing")
    # collect agent variables
    adata = [:value, :type, :status, :ON_interbank, :flow, 
        :Term_interbank, :loans, :loans_demand, :GDP, 
        :wages, :consumption, :income, :bills, :taxes]

    for x in param_range
        # Setup model properties
        properties = (
            param => x,
            shock = shock,
        )

        println("Running parameter scans for $(param) at $(x)...")
        
        models = [CViM.init_model(; seed, properties...) for seed in seeds]

        # run parallel models
        adf, _ =  ensemblerun!(models, dummystep, CViM.model_step!, 800;
            adata, parallel = true, showprogress = true)

        println("Collecting data for $(param) at $(x)...")

        # Aggregate agent data
        adf = @pipe adf |>
            groupby(_, [:step, :id, :status]) |>
            combine(_, adata[1:3] .=> unique, adata[4:end] .=> mean; renamecols = false)
        adf[!, param] .= x

        # Write data to disk
        println("Saving to disk for $(param) at $(x)...")
        datapath = mkpath("data/sensitivity_analysis/$(param)/$(x)")
        filepath = "$datapath/df.csv"
        isfile(filepath) && rm(filepath)
        CSV.write(filepath, adf; append = true, header = !isfile(filepath))
        empty!(adf)
    end
    return nothing
end

function run(seeds::Vector{UInt32})
    run_sens(seeds, :r, [0.9, 1.1, 1.3])
    run_sens(seeds, :γ, [0.0, 0.1, 0.3])
    run_sens(seeds, :pref, [0.01, 0.05, 0.09])
    run_sens(seeds, :g, [100.0, 200.0, 300.0])
    printstyled("Paramascan and data collection finished."; color = :blue)
    return nothing
end

begin
    Random.seed!(95129)
    seeds = rand(UInt32, 1)
    run(seeds)
end