"""
Initialise model and agent types, compute initial values and distribute them among the model's sectors, 
set initial market interactions and agents' personal values, leverage and liquidity measures.
"""

"""
    init_model(; seed::UInt32 = UInt32(95129), properties...) → model

Initialise the model.    
"""
function init_model(; seed::UInt32 = UInt32(95129), shock = "Missing",
        properties...)
    if shock in ("Missing", "Corridor", "Corridor-unique", "Expansionary")
        model = ABM(Union{Government, CentralBank, Firm, Household, Bank};
            properties = Parameters(; shock, properties...), 
            scheduler = Schedulers.Randomly(),
            rng = Xoshiro(seed),
            warn = false # turns off Agents.jl warning of Union types
        )

        init_agents!(model)
        distribute_SS_values(model)
        real_sector_interactions!(model)
        credit_sector_interactions!(model)
        init_liq_values(model)
        init_lev(model)
    else
        error("You provided a shock named $(shock) that is not yet implemented. 
        Check for typos or add the shock!")
    end
    return model
end

"""
    init_agents!(model::ABM) → model

Initialise and add agent types to the model.    
"""
function init_agents!(model::ABM)
    # initialise Households
    for id = 1:model.n_hh
        a = Household(
            id = id,
            value = init_personal_values(id, model.n_hh, 0, cumsum(model.value_dist)),
        )
        add_agent!(a, model)
    end

    # initialise Firms
    for id = (model.n_hh + 1):(model.n_hh + model.n_f)
        a = Firm(
            id = id,
            value = init_personal_values(id, model.n_f, model.n_hh, cumsum(model.value_dist)),
        )
        add_agent!(a, model)
    end

    # initialise commercial Banks 
    for id = (model.n_hh + model.n_f + 1):(model.n_hh + model.n_f + model.n_bj)
        a = Bank(
            id = id,
            type = :commercial,
            value = init_personal_values(id, model.n_bj, (model.n_hh + model.n_f), cumsum(model.value_dist)),
        )
        add_agent!(a, model)
    end

    # initialise business Banks 
    for id = (model.n_hh + model.n_f + model.n_bj + 1):(model.n_hh + model.n_f + model.n_bj + model.n_bk)
        a = Bank(
            id = id,
            type = :business, 
            value = init_personal_values(id, model.n_bk, (model.n_hh + model.n_f + model.n_bj), cumsum(model.value_dist)),
        )     
        add_agent!(a, model)
    end

    # initialise Government
    a = Government(
        id = model.n_hh + model.n_f + model.n_bj + model.n_bk + 1
    )
    add_agent!(a, model)

    # initialise CentralBank
    a = CentralBank(
        id = model.n_hh + model.n_f + model.n_bj + model.n_bk + 2
    )
    add_agent!(a, model)

    return model
end

"""
    SS_initial(model) → all_variables

Compute the SS values of the model. Please note that the notation used in this function roughly matches the one used in the paper as 
in the code we differentiate between business and commercial banks more explicitly. We thus denote with `j` variables
related to commercial banks and with `k` those related to business ones. The function returns `all_variables` which is an
NTuple that links variable names to their corrisponding initial values. More, it updates the SS-given value
of `model.α2`, trhows a warning if variables are negative and performs the initial Stock-Flow consistency checks.
"""
function SS_initial(model)
    # exogenous variables and parameters
    icb = model.ib # => Pcb = 0.0
    ilh = 0.0335
    ilf = 0.0325
    id = 0.0155
    levf = 0.5
    levh = 0.5
    Inv = 303.8
    # eqs
    Pcb = 0.0
    Y = (model.g - Pcb)/((model.τ/ (1+model.ρ)) - model.r * model.ib)
    W = Y /(1 + model.ρ)
    T = model.τ * W
    B = Y * model.r
    Lf = levf * W
    C = Y - model.g # - Inv
    Yd = C
    Df = Lf - Inv
    Lh = levh * C
    α2_SS  = ((C * (1 - model.α1 - model.α3 * levh)) / ((B - Df + Lf)))
    NWh = (C * (1 - model.α1 - model.α3 * levh)) / (α2_SS)
    Dh = NWh + Lh
    Hj = (model.μ + model.v) * Dh
    Hk = (model.μ + model.v) * Df
    H = Hj + Hk
    Bj = max(0.0, Dh - Lh - model.μ * Dh)
    Bk = max(0.0, Df - Lf - model.μ * Df)
    Aj = if Bj == 0.0
            Hj + Lh - Dh
        else
            model.v * Dh
        end
    Ak = if Bk == 0.0 
            Hk + Lf - Df
        else
            model.v * Df
        end
    A = Aj + Ak
    Bcb = B - Bj - Bk
    Pbj = ilh * Lh - id * Dh + model.ib * Bj + icb * Hj - icb * Aj
    Pbk = ilf * Lf - id * Df + model.ib * Bk + icb * Hk - icb * Ak
    Pf = C + model.g - W + id * Df - ilf * Lf 
    # sectoral net worths: balance sheet matrix
    GD = B
    # NWh = Dh - Lh
    NWcb = Bcb - H + A
    NWbj = Lh + Hj - Dh - Aj + Bj
    NWbk = Lf + Hk - Df - Ak + Bk
    NWf = Inv + Df - Lf

    # collect all variables
    all_variables = (Inv, Y, W, T, B, Lf, C, 
                    Yd, Df, NWh, Lh, Dh, Hj, Hk, 
                    H, Bj, Bk, Ak, Aj, A, Bcb, 
                    Pbj, Pbk, Pf, ilh, ilf, id, NWbj, NWbk)

    # checks for negative initial values
    any(<(0.0), all_variables) && @warn "Some initial values are negative, check the parameters and/or the equations again!"

    # checks for networth balance and hidden equation at SS
    SS_SFC_checks(GD, NWh, NWcb, NWbj, NWbk, NWf, Inv)
    
    # update SS-given parameters
    model.α2 = α2_SS

    return all_variables
end

"""
    SS_SFC_checks(GD, NWh, NWcb, NWbj, NWbk, NWf, Inv; tol::Float64 = 1e-06) → nothing

Perform SFC checks for initial values calculation at the Steady State.
"""
function SS_SFC_checks(GD, NWh, NWcb, NWbj, NWbk, NWf, Inv; tol::Float64 = 1e-06)
    if (GD - (NWh + NWcb + NWbj + NWbk + NWf)) - Inv  > tol
        @warn "Initial values calculation does not respect stock-flow consistency - Net Worth!
            Check equations and parameters and try again."
    end
    if abs(NWcb) > tol
        @warn "Initial values calculation does not respect stock-flow consistency - Hidden Equation!
            Check equations and parameters and try again."
    end
    return nothing
end

"""
    distribute_SS_values(model) → model

Distribute SS values to aggregate sectors (Government and Central Bank) and to heterogenous sectors homogeneously.
"""
function distribute_SS_values(model)
    # take variables names and corresponding initial values
    (Inv, Y, W, T, B, Lf, C,
    Yd, Df, NWh, Lh, Dh, Hj, Hk,
    H, Bj, Bk, Ak, Aj, A, Bcb,
    Pbj, Pbk, Pf, ilh, ilf, id, NWbj, NWbk) = SS_initial(model)

    for a in allagents(model)
        if isa(a, Government)
            a.taxes = T
            a.bills = B
            a.bills_prev = a.bills
        elseif isa(a, CentralBank)
            a.bills = Bcb
            a.bills_prev = a.bills
            a.hpm = H
            a.hpm_prev = a.hpm
            a.advances = A
            a.advances_prev = a.advances
        elseif isa(a, Household)
            a.consumption = C / model.n_hh
            a.income = Yd / model.n_hh
            a.wages = W / model.n_hh
            a.taxes = T / model.n_hh
            a.loans = Lh / model.n_hh
            a.loans_prev = a.loans
            a.loans_demand = a.loans
            a.deposits = Dh / model.n_hh
            a.deposits_prev = a.deposits_prev
            a.networth = NWh / model.n_hh
            a.networth_prev = a.networth
        elseif isa(a, Firm)
            a.consumption = C / model.n_f
            a.inventories = Inv / model.n_f
            a.inventories_prev = a.inventories
            a.GDP = Y / model.n_f
            a.wages = W / model.n_f
            a.wages_prev = a.wages
            a.deposits = Df / model.n_f
            a.deposits_prev = a.deposits
            a.loans = Lf / model.n_f
            a.loans_prev = a.loans
            a.loans_demand = a.loans
            a.profits = Pf / model.n_f
        elseif isa(a, Bank) && a.type == :commercial
            a.deposits = Dh / model.n_bj
            a.deposits_prev = a.deposits
            a.advances = Aj / model.n_bj
            a.advances_prev = a.advances
            a.bills = Bj / model.n_bj
            a.bills_prev = a.bills
            a.hpm = Hj / model.n_bj
            a.hpm_prev = a.hpm
            a.loans = Lh / model.n_bj
            a.loans_prev = a.loans
            a.funding_costs = model.icb
            a.ilh_rate = ilh
            a.id_rate = id
            a.profits = Pbj / model.n_bj
            a.networth = NWbj / model.n_bj
        elseif isa(a, Bank) && a.type == :business
            a.deposits = Df / model.n_bk
            a.deposits_prev = a.deposits
            a.advances = Ak / model.n_bk
            a.advances_prev = a.advances
            a.bills = Bk / model.n_bk
            a.bills_prev = a.bills
            a.hpm = Hk / model.n_bk
            a.hpm_prev = a.hpm
            a.loans = Lf / model.n_bk
            a.loans_prev = a.loans
            a.funding_costs = model.icb
            a.ilf_rate = ilf
            a.id_rate = id
            a.profits = Pbk / model.n_bk
            a.networth = NWbk / model.n_bk
        end
    end
    return model
end

"""
    real_sector_interactions!(model) → model

Match households and firms for consumption and wages exchanges in the real sector.    
"""
function real_sector_interactions!(model)
    HPF = round(Int, model.n_hh/model.n_f)

    for i in 1:model.n_f
        for id in (HPF*(i-1)+1):(HPF*i)
            model[id].belongToFirm = model.n_hh + i
            push!(model[model.n_hh + i].customers, id)
        end
    end
    return model
end

"""
    credit_sector_interactions!(model) → model

Match households and firms to banks in the credit market.    
"""
function credit_sector_interactions!(model)
    APBk = round(Int, model.n_f/model.n_bk)
    for i in 1:model.n_bk
        for id in (APBk*(i-1) + model.n_hh + 1):(APBk*i + model.n_hh)
            model[id].belongToBank = model.n_hh + model.n_f + model.n_bj +  i
            push!(model[model.n_hh + model.n_f +  model.n_bj + i].firms_customers, id)
        end
    end

    APBj = round(Int, model.n_hh/model.n_bj)
    for i in 1:model.n_bj
        for id in (APBj*(i-1)+1):(APBj*i)
            model[id].belongToBank = model.n_hh + model.n_f + i
            push!(model[model.n_hh + model.n_f + i].hh_customers, id)
        end
    end
    return model
end

"""
    init_personal_values(id, n, x, value_dist; VALUES = (:C, :O, :SE, :ST)) → nothing

Initialise agents' value types based on their `id`, total number of agents of the same type `n`,
total number of previously initialised agents of different type `x`, and values distribution `value_dist` which 
depends on the `scenario` implemented.
"""
function init_personal_values(id, n, x, value_dist; VALUES = (:C, :O, :SE, :ST))
    for i in 1:lastindex(value_dist)
        if id ≤ (value_dist[i] * n) + x || id ≈ (value_dist[i] * n) + x
            return VALUES[i]
        else
            continue
        end
    end
    return nothing
end

"""
    init_liq_values(model) → model

Initialise banks' liquidity preferences depending on value types.
"""
function init_liq_values(model)
    for id in ids_by_type(Bank, model)
        if model[id].value == :C || model[id].value == :SE
            if model[id].type == :business
                model[id].liq_values = rand(model.rng, Uniform(0.1, 0.5))
            else
                model[id].liq_values = rand(model.rng, Uniform(0.6, 1.0))
            end
        elseif model[id].value == :O || model[id].value == :ST
            if model[id].type == :business
                model[id].liq_values = rand(model.rng, Uniform(0.6, 1.0))
            else
                model[id].liq_values = rand(model.rng, Uniform(0.1, 0.5))
            end
        end
        model[id].liq_pref = model[id].liq_values
    end
    return model
end

"""
    init_lev(model) → model

Initialise households' and firms' levarage measure depending on value types.
"""
function init_lev(model)
    for id in ids_by_type(Household, model)
        if model[id].value == :C || model[id].value == :ST
            model[id].lev = rand(model.rng, Uniform(0.1, 0.5))
        else
            model[id].lev = rand(model.rng, Uniform(0.6, 1.0))
        end
    end
    for id in ids_by_type(Firm, model)
        if model[id].value == :C || model[id].value == :SE
            model[id].lev = rand(model.rng, Uniform(0.1, 0.5))
        else
            model[id].lev = rand(model.rng, Uniform(0.6, 1.0))
        end
    end
    return model
end