module CViM

using Agents
using Pipe
using Random
using StatsBase
using Distributions
using Combinatorics

include("model/params.jl")
include("model/structs.jl")
include("model/init.jl")
include("model/utils.jl")
include("model/SFC/gov.jl")
include("model/SFC/cb.jl")
include("model/SFC/hh.jl")
include("model/SFC/firms.jl")
include("model/SFC/banks.jl")

"""
    model_step!(model) → model

Define what happens in the model at each step.
"""
function model_step!(model)
    model.step += 1

    #begin: apply shocks
    if model.shock == "Corridor" && iszero(model.step % model.sas)
        model.icbd += 0.005 # 50 basis points
        model.icbl += 0.005
        model.icb = (model.icbl + model.icbd) / 2.0
    elseif model.shock == "Corridor-unique" && model.step == model.sas
        model.icbd += 0.005 # 50 basis points
        model.icbl += 0.005
        model.icb = (model.icbl + model.icbd) / 2.0
    elseif model.shock == "Expansionary" && model.step == model.sas
        model.icbd -= 0.005
        model.icbl -= 0.005 
        model.icb = (model.icbl + model.icbd) / 2.0
    end
    #end: apply shocks

    for id in ids_by_type(Bank, model)
        CViM.prev_vars!(model[id])
        CViM.update_liq_preferences!(model[id], model)
        CViM.credit_rates!(model[id], model.χ1, model.χ2, model.χ3)
        CViM.reset_vars!(model[id])
    end
    CViM.update_model_vars!(model)

    CViM.update_change_rates!(model)
    for id in ids_by_type(Household, model)
        CViM.prev_vars!(model[id])
        CViM.repayment!(model[id], model)
        CViM.interests_payments!(model[id], model)
        CViM.consumption!(model[id], model)
        CViM.update_lev!(model[id], model.change_rates_hh, model.pref)
    end

    CViM.hhs_matching!(model)
    for id in ids_by_type(Household, model)
        CViM.loans_demand!(model[id])
        CViM.loans!(model[id], model)
    end   

    for id in ids_by_type(Firm, model)
        CViM.prev_vars!(model[id])
        CViM.repayment!(model[id], model)
        CViM.networth!(model[id])
        CViM.interests_payments!(model[id], model)
        CViM.update_lev!(model[id], model.change_rates_firms, model.pref)
    end
    
    CViM.firms_matching!(model)
    for id in ids_by_type(Firm, model)
        CViM.loans_demand!(model[id])
        CViM.loans!(model[id], model)
        CViM.deposits!(model[id], model)
        CViM.consumption!(model[id], model)
        CViM.inventories!(model[id])
        CViM.production!(model[id], model.g, model.n_f)
        CViM.wages!(model[id], model)
        CViM.profits!(model[id])
        CViM.current_balance!(model[id])
        CViM.SFC!(model[id], model)
    end
    CViM.GDP_growth_rate(model)

    for id in ids_by_type(Bank, model)
        CViM.interests_payments!(model[id], model)
        CViM.profits!(model[id])
        CViM.current_balance!(model[id])
        CViM.update_status!(model[id]) # updates IB status
    end
    
    banks_profits = sum(a.profits for a in allagents(model) if a isa Bank) 
    firms_profits = sum(a.profits for a in allagents(model) if a isa Firm)
    profits = (banks_profits + firms_profits)/model.n_hh
    for id in ids_by_type(Household, model)
        CViM.taxes!(model[id], model.τ)
        CViM.income!(model[id], profits)
        CViM.SFC!(model[id], model)
    end

    for id in ids_by_type(CentralBank, model)
        CViM.prev_vars!(model[id])
        CViM.profits!(model[id], model)
        CViM.current_balance!(model[id], model)
    end

    for id in ids_by_type(Government, model)
        CViM.SFC!(model[id], model)
    end

    # begin: Interbank Market
    CViM.ib_matching!(model)
    for id in ids_by_type(Bank, model)
        CViM.update_ib_demand_supply!(model[id])
        CViM.ib_on!(model[id], model)
        CViM.ib_term!(model[id], model)
        CViM.lending_facility!(model[id])
        CViM.deposit_facility!(model[id])
        CViM.funding_costs!(model[id], model.icb, model.ion, model.iterm, model.icbl)
        CViM.SFC!(model[id], model)
    end
    CViM.ib_rates!(model)
    # end: Interbank Market
    
    for id in ids_by_type(CentralBank, model)
        CViM.SFC!(model[id], model)
    end

    CViM.SFC_checks!(model) # check for Stock-Flow consistency 
    return model
end

"""
    GDP_growth_rate(model) → model.gy

Updates growth rate of national GDP.
"""
function GDP_growth_rate(model)
    model.gy = sum(a.GDP - a.GDP_prev for a in allagents(model) if a isa Firm) / sum(a.GDP_prev for a in allagents(model) if a isa Firm)
    return model.gy
end

"""
    firms_matching!(model) → model
Updates firms' matching in the credit market.
"""
function firms_matching!(model)
    for id in ids_by_type(Firm, model)
        #Select potential partners
        potential_partners = filter(i -> model[i] isa Bank && i != model[id].belongToBank && model[i].type == :business, collect(allids(model)))[1:model.χ]
        #Select new partner with the best interest rate among potential partners
        new_partner = rand(model.rng, filter(i -> i in potential_partners && model[i].ilf_rate == minimum(model[a].ilf_rate for a in potential_partners), potential_partners))
        #Select interest rate of the new potential partner
        inew = model[new_partner].ilf_rate
        #Pick up old partner
        old_partner = model[id].belongToBank
        #PICK UP THE INTEREST OF THE OLD PARTNER
        iold = model[old_partner].ilf_rate
        #COMPARE OLD AND NEW INTERESTS
        if rand(model.rng) < (1 - exp(model.λ * (inew - iold)/inew))
            #THEN SWITCH TO A NEW PARTNER
            deleteat!(model[old_partner].firms_customers, findall(x -> x == id, model[old_partner].firms_customers))
            model[id].belongToBank = new_partner
            push!(model[new_partner].firms_customers, id)
        end
    end
    return model
end

"""
    hhs_matching!(model) → model
Updates households' matching in the credit market.
"""
function hhs_matching!(model)
    for id in ids_by_type(Household, model)
        #Select potential partners
        potential_partners = filter(i -> model[i] isa Bank && i != model[id].belongToBank && model[i].type == :commercial, collect(allids(model)))[1:model.χ]
        #Select new partner with the best interest rate among potential partners
        new_partner = rand(model.rng, filter(i -> i in potential_partners && model[i].ilh_rate == minimum(model[a].ilh_rate for a in potential_partners), potential_partners))
        #Select interest rate of the new potential partner
        inew = model[new_partner].ilh_rate
        #Pick up old partner
        old_partner = model[id].belongToBank
        #PICK UP THE INTEREST OF THE OLD PARTNER
        iold = model[old_partner].ilh_rate
        #COMPARE OLD AND NEW INTERESTS
        if rand(model.rng) < (1 - exp(model.λ * (inew - iold)/inew))
            #THEN SWITCH TO A NEW PARTNER
            deleteat!(model[old_partner].hh_customers, findall(x -> x == id, model[old_partner].hh_customers))
            model[id].belongToBank = new_partner
            push!(model[new_partner].hh_customers, id)
        end
    end
    return model
end

"""
    ib_matching!(model) → model

Update borrowing banks' matching.
"""
function ib_matching!(model)
    for id in ids_by_type(Bank, model)
        # end function prematurely if there are no lender banks available
        isempty([a.id for a in allagents(model) if a isa Bank && a.status == :lender]) && return

        if model[id].status == :borrower
            #Select potential partners with the closest liquidity preference
            potential_partners = filter(i -> model[i] isa Bank && model[i].status == :lender && abs(model[i].liq_pref - model[id].liq_pref) <= 1e-01, collect(allids(model)))
            if !isempty(potential_partners) #&& rand(model.rng, Bool)
                #Select new partner
                new_partner = rand(model.rng, filter(i -> i in potential_partners, potential_partners))
                model[id].belongToBank = new_partner
                push!(model[new_partner].ib_customers, id)
            end
        end
    end
    return model
end

"""
    ib_rates!(model) → model

Update interbank rates on overnight and term segments based on disequilibrium dynamics between demand and supply.
The function also checks that interest rates fall within the central bank's corridor, otherwise a warning is issued.
"""
function ib_rates!(model)
    if length([a.id for a in allagents(model) if a isa Bank && a.status == :borrower && !ismissing(a.belongToBank)]) > 0 && 
        length([a.id for a in allagents(model) if a isa Bank && a.status == :lender && !isempty(a.ib_customers)]) > 0
 
        ON = sum(a.ib_on_demand for a in allagents(model) if a isa Bank && a.status == :borrower) - 
            sum(a.ib_on_supply for a in allagents(model) if a isa Bank && a.status == :lender)
        Term = sum(a.ib_term_demand for a in allagents(model) if a isa Bank && a.status == :borrower) - 
            sum(a.ib_term_supply for a in allagents(model) if a isa Bank && a.status == :lender)

        model.ion = model.icbd + ((model.icbl - model.icbd)/(1 + exp(-model.σib * ON)))
        model.iterm = model.icbd + ((model.icbl - model.icbd)/(1 + exp(-model.σib * Term)))
    end

    # check corridor
    if model.ion > model.icbl || model.ion < model.icbd
        @warn "Interbank ON rate outside the central bank's corridor!"
    elseif model.iterm > model.icbl || model.iterm < model.icbd
        @warn "Interbank Term rate outside the central bank's corridor!"
    end
    return model.ion, model.iterm
end

"""
    update_model_vars!(model) → model

Update model variables.
"""
function update_model_vars!(model)
    model.ion_prev = model.ion
    model.iterm_prev = model.iterm
    return model
end

"""
    update_change_rates!(model) → model

Update credit rates changes for households and firms. Used to update real sector agents' desired leverage.
"""
function update_change_rates!(model)
    model.change_rates_hh = (mean(a.ilh_rate for a in allagents(model) if a isa Bank && a.type == :commercial) - 
        mean(a.ilh_rate_prev for a in allagents(model) if a isa Bank && a.type == :commercial))
    model.change_rates_firms = (mean(a.ilf_rate for a in allagents(model) if a isa Bank && a.type == :business) - 
        mean(a.ilf_rate_prev for a in allagents(model) if a isa Bank && a.type == :business))
    return model
end

end # module CViM