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

# Start workers
using Distributed
addprocs()

# Set up package environment on workers
@everywhere begin
    using Pkg
    Pkg.activate(".")
end

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

# Load packages on workers
@everywhere begin
    using Agents
    using StatsBase
    using Distributions
    using CViM
end

# runs the model, transforms and saves data
function run_model(number_of_runs::Int = 50)
    scenarios = (
        (scenario = "Baseline", value_dist = (0.25, 0.25, 0.25, 0.25)),
        (scenario = "C-skewed",  value_dist = (0.7, 0.1, 0.1, 0.1)),
        (scenario = "O-skewed", value_dist = (0.1, 0.7, 0.1, 0.1)),
        (scenario = "SE-skewed", value_dist = (0.1, 0.1, 0.7, 0.1)),
        (scenario = "ST-skewed", value_dist = (0.1, 0.1, 0.1, 0.7)),
    )

    shocks = ("Missing", "Corridor", "Corridor-unique", "Expansionary")

    # collect agent variables
    adata = [:value, :type, :status, :lev, :liq_pref, :ON_interbank, :flow,
        :Term_interbank, :loans, :loans_demand, :consumption, :GDP, 
        :lending_facility, :deposit_facility, :ib_on_demand, :ib_term_demand,
        :networth, :bills]
        
    # collect model variables
    mdata = [:n_hh, :n_f, :gy, :ion, :iterm, :icbl, :icbd, :icb]

    for (i, scenario) in enumerate(scenarios)
        seeds = rand(UInt32, number_of_runs)

        for shock in shocks
            properties = (scenario...,
                shock = shock)

            println("Creating $number_of_runs seeded $(properties.shock)-shock and $(scenario.scenario)-scenario models and running...")

            models = [CViM.init_model(; seed, properties...) for seed in seeds]
            
            adf, mdf, _ =  ensemblerun!(models, dummystep, CViM.model_step!, 800;
                adata, mdata, parallel = true, showprogress = true)
                
            println("Collecting data for $(properties.shock)-shock and $(scenario.scenario)-scenario...")

            # Aggregate model data over replicates
            mdf = @pipe mdf |>
                groupby(_, :step) |>
                combine(_, mdata[1:2] .=> unique, mdata[3:end] .=> mean; renamecols = false)
            mdf[!, :shock] = fill(properties.shock, nrow(mdf))
            mdf[!, :scenario] = fill(properties.scenario, nrow(mdf))

            # Aggregate agent data over replicates
            adf = @pipe adf |>
                groupby(_, [:step, :id, :status]) |>
                combine(_, adata[1:3] .=> unique, adata[4:end] .=> mean; renamecols = false)
            adf[!, :shock] = fill(properties.shock, nrow(adf))
            adf[!, :scenario] = fill(properties.scenario, nrow(adf))

            # Write data to disk
            println("Saving to disk for $(properties.shock)-shock and $(scenario.scenario)-scenario...")
            datapath = mkpath("data/shock=$(properties.shock)/$(scenario.scenario)")
            filepath = "$datapath/adf.csv"
            isfile(filepath) && rm(filepath)
            CSV.write(filepath, adf)
            filepath = "$datapath/mdf.csv"
            isfile(filepath) && rm(filepath)
            CSV.write(filepath, mdf)
            println("Finished for $(properties.shock) shock and $(scenario.scenario) scenario.")
        end
    end

    printstyled("Simulations finished and data saved!"; color = :blue)
    return nothing
end

Random.seed!(95129)
run_model()