using .Networks

using JuMP
using Dates
using DataStructures
using Printf
using Parameters
"""
REF_SCHEDULE_TYPE_IN_TSO : Indicates which schedule to use as reference for pilotables state/levels needed
                            for sequencing constraints and TSO objective function.
"""
@with_kw mutable struct TSOBilevelConfigs
    CONSIDER_DELTAS::Bool = true
    CONSIDER_N_1_CSTRS::Bool = false
    TSO_LIMIT_PENALTY::Float64 = 1e-3
    TSO_LOL_PENALTY::Float64 = 1e5
    TSO_CAPPING_COST::Float64 = 1.
    TSO_PILOTABLE_BOUNDING_COST::Float64 = 1.
    USE_UNITS_PROP_COST_AS_TSO_BOUNDING_COST::Bool = true
    MARKET_LOL_PENALTY::Float64 = 1e5
    MARKET_CAPPING_COST::Float64 = 1.
    out_path::Union{Nothing,String} = nothing
    problem_name::String = "TSOBilevel"
    LINK_SCENARIOS_LIMIT::Bool = true
    LINK_SCENARIOS_PILOTABLE_LEVEL::Bool = false
    LINK_SCENARIOS_PILOTABLE_ON::Bool = false
    LINK_SCENARIOS_PILOTABLE_LEVEL_MARKET::Bool = false
    big_m = 1e6
    REF_SCHEDULE_TYPE_IN_TSO::Union{Market,TSO} = Market();
end

##########################################################
#        upper problem : TSO
##########################################################

@with_kw struct TSOBilevelTSOLimitableModel <: AbstractLimitableModel
    #gen,ts,s
    delta_p = SortedDict{Tuple{String,DateTime,String},VariableRef}();
    #gen,ts,s
    p_injected = SortedDict{Tuple{String,DateTime,String},VariableRef}();
    #gen,ts,s
    p_limit = SortedDict{Tuple{String,DateTime,String},VariableRef}();
    #gen,ts,s
    b_is_limited = SortedDict{Tuple{String,DateTime,String},VariableRef}();
    #gen,ts,s
    p_limit_x_is_limited = SortedDict{Tuple{String,DateTime,String},VariableRef}();
    #gen,ts,s
    p_capping = SortedDict{Tuple{String,DateTime,String},VariableRef}();
    #ts,s
    p_global_capping = SortedDict{Tuple{DateTime,String},VariableRef}();
end

@with_kw struct TSOBilevelTSOPilotableModel <: AbstractPilotableModel
    #gen,ts,s
    p_imposition_min = SortedDict{Tuple{String,DateTime,String},VariableRef}();
    #gen,ts,s
    p_imposition_max = SortedDict{Tuple{String,DateTime,String},VariableRef}();
    #gen,ts,s
    b_start = SortedDict{Tuple{String,DateTime,String},VariableRef}();
    #gen,ts,s
    b_on = SortedDict{Tuple{String,DateTime,String},VariableRef}();
end

@with_kw struct TSOBilevelTSOLoLModel <: AbstractLoLModel
    #bus,ts,s #Loss of Load
    p_loss_of_load = SortedDict{Tuple{String,DateTime,String},VariableRef}();
    #ts,s
    p_global_loss_of_load = SortedDict{Tuple{DateTime,String},VariableRef}();
end

@with_kw mutable struct TSOBilevelTSOObjectiveModel <: AbstractObjectiveModel
    deltas = GenericAffExpr{Float64, VariableRef}()

    limitable_cost = GenericAffExpr{Float64, VariableRef}()
    pilotable_cost = GenericAffExpr{Float64, VariableRef}()

    penalty = GenericAffExpr{Float64, VariableRef}()

    full_obj_1 = GenericAffExpr{Float64, VariableRef}()
    full_obj_2 = GenericAffExpr{Float64, VariableRef}()
end

##########################################################
#        lower problem : Market
##########################################################
@with_kw struct TSOBilevelMarketLimitableModel <: AbstractLimitableModel
    #ts,s #FIXME not sure if this should be a lower or an upper variable
    p_global_injected = SortedDict{Tuple{DateTime,String},VariableRef}();
    #ts,s
    p_global_capping = SortedDict{Tuple{DateTime,String},VariableRef}();
end

"""
# Note :
    deltas are added here to allow automatically computing them and their constraints
    but the constraints are not included in the reformulation (stationarity nor complementarity)
    cause they can be seen as part of the upper
    constraints discussed here are : delta >= +- (p_injected - P_ref)
"""
@with_kw struct TSOBilevelMarketPilotableModel <: AbstractPilotableModel
    #gen,ts,s
    delta_p = SortedDict{Tuple{String,DateTime,String},VariableRef}();
    #gen,ts,s
    p_injected = SortedDict{Tuple{String,DateTime,String},VariableRef}();
end

@with_kw struct TSOBilevelMarketLoLModel <: AbstractLoLModel
    #ts,s
    p_global_loss_of_load = SortedDict{Tuple{DateTime,String},VariableRef}();
end

@with_kw mutable struct TSOBilevelMarketObjectiveModel <: AbstractObjectiveModel
    limitable_cost = GenericAffExpr{Float64, VariableRef}()
    pilotable_cost = GenericAffExpr{Float64, VariableRef}()

    penalty = GenericAffExpr{Float64, VariableRef}()

    full_obj = GenericAffExpr{Float64, VariableRef}()
end

##########################################################
#        TSOBilevel Model
##########################################################

@with_kw struct TSOBilevelTSOModelContainer <: AbstractModelContainer
    model::Model
    limitable_model::TSOBilevelTSOLimitableModel = TSOBilevelTSOLimitableModel()
    pilotable_model::TSOBilevelTSOPilotableModel = TSOBilevelTSOPilotableModel()
    lol_model::TSOBilevelTSOLoLModel = TSOBilevelTSOLoLModel()
    objective_model::TSOBilevelTSOObjectiveModel = TSOBilevelTSOObjectiveModel()
    #branch,ts,s,ptdf_case
    flows::SortedDict{Tuple{String,DateTime,String,String},AffExpr} =
        SortedDict{Tuple{String,DateTime,String,String},AffExpr}()
    rso_constraint::SortedDict{Tuple{String,DateTime,String,String},ConstraintRef} =
        SortedDict{Tuple{String,DateTime,String,String},ConstraintRef}()
end
@with_kw struct TSOBilevelMarketModelContainer <: AbstractModelContainer
    model::Model
    limitable_model::TSOBilevelMarketLimitableModel = TSOBilevelMarketLimitableModel()
    pilotable_model::TSOBilevelMarketPilotableModel = TSOBilevelMarketPilotableModel()
    lol_model::TSOBilevelMarketLoLModel = TSOBilevelMarketLoLModel()
    objective_model::TSOBilevelMarketObjectiveModel = TSOBilevelMarketObjectiveModel()
    #ts,s
    eod_constraint::SortedDict{Tuple{Dates.DateTime,String}, ConstraintRef} =
        SortedDict{Tuple{Dates.DateTime,String}, ConstraintRef}()
end
@with_kw struct TSOBilevelKKTModelContainer <: AbstractModelContainer
    #TODO create adequate struct to link each dual kkt variable, indicator variable, and constraint to each other
    # e.g.
    # eod_model = ts,s -> Struct{dual, indicator, ConstraintRef}
    # pilotable_min = id,ts,s -> Struct{dual, indicator, ConstraintRef} or Struct{dual, Nothing, ConstraintRef}
    # than simply call reformulate_kkt(cstr) when building lower problem's constraints provided that objective is created first
    model::Model
    #ts,s
    eod_duals::SortedDict{Tuple{DateTime,String},VariableRef} =
        SortedDict{Tuple{DateTime,String},VariableRef}()
    capping_duals::SortedDict{Tuple{DateTime,String},VariableRef} =
        SortedDict{Tuple{DateTime,String},VariableRef}()
    capping_indicators::SortedDict{Tuple{DateTime,String},VariableRef} =
        SortedDict{Tuple{DateTime,String},VariableRef}()
    loss_of_load_duals::SortedDict{Tuple{DateTime,String},VariableRef} =
        SortedDict{Tuple{DateTime,String},VariableRef}()
    loss_of_load_indicators::SortedDict{Tuple{DateTime,String},VariableRef} =
        SortedDict{Tuple{DateTime,String},VariableRef}()
    #pilotable_id,ts,s
    firmness_duals::SortedDict{Tuple{String,DateTime,String},VariableRef} =
        SortedDict{Tuple{String,DateTime,String},VariableRef}()
    pmin_duals::SortedDict{Tuple{String,DateTime,String},VariableRef} =
        SortedDict{Tuple{String,DateTime,String},VariableRef}()
    pmin_indicators::SortedDict{Tuple{String,DateTime,String},VariableRef} =
        SortedDict{Tuple{String,DateTime,String},VariableRef}()
    pmax_duals::SortedDict{Tuple{String,DateTime,String},VariableRef} =
        SortedDict{Tuple{String,DateTime,String},VariableRef}()
    pmax_indicators::SortedDict{Tuple{String,DateTime,String},VariableRef} =
        SortedDict{Tuple{String,DateTime,String},VariableRef}()
end
TSOBilevelModel = BilevelModelContainer{TSOBilevelTSOModelContainer, TSOBilevelMarketModelContainer, TSOBilevelKKTModelContainer}
function BilevelModelContainer{TSOBilevelTSOModelContainer,TSOBilevelMarketModelContainer,TSOBilevelKKTModelContainer}()
    bilevel_model = Model()
    upper = TSOBilevelTSOModelContainer(model=bilevel_model)
    lower = TSOBilevelMarketModelContainer(model=bilevel_model)
    kkt_model = TSOBilevelKKTModelContainer(model=bilevel_model)
    return BilevelModelContainer(bilevel_model, upper, lower, kkt_model)
end

"""
Computes the actual capping in the Limitable model

    The global capping of the upper limitable model can be lower than the actual capping
    The actual capping is given by the localised capping in the upper problem or the global capping in the lower problem
"""
function sum_capping(limitable_model::TSOBilevelTSOLimitableModel, ts,s, network::Networks.Network)
    sum_l = 0.
    for gen in Networks.get_generators_of_type(network, Networks.LIMITABLE)
        gen_id = Networks.get_id(gen)
        sum_l += limitable_model.p_capping[gen_id,ts,s]
    end
    return sum_l
end

"""
Computes the actual loss of load (LoL) in the LoL model

    The global LoL of the upper limitable model can be lower than the actual LoL
    The actual LoL is given by the localised LoL in the upper problem or the global LoL in the lower problem
"""
function sum_lol(lol_model::TSOBilevelTSOLoLModel, ts, s, network::Networks.Network)
    sum_l = 0.
    for bus in Networks.get_buses(network)
        bus_id = Networks.get_id(bus)
        sum_l += lol_model.p_loss_of_load[bus_id,ts,s]
    end
    return sum_l
end

function has_positive_slack(model_container::TSOBilevelModel)::Bool
    return has_positive_value(model_container.lower.lol_model.p_global_loss_of_load) #If TSO cut => market did too
end

function has_positive_slack(model_container::TSOBilevelTSOModelContainer)::Bool
    return has_positive_value(model_container.lol_model.p_loss_of_load) #If market cut => TSO localised the LoL
end

function get_upper_obj_expr(bilevel_model::TSOBilevelModel)
    return bilevel_model.upper.objective_model.full_obj_2
end
function get_lower_obj_expr(bilevel_model::TSOBilevelModel)
    return bilevel_model.lower.objective_model.full_obj
end

##########################################################
#        upper problem : TSO functions : TSOBilevelTSOModelContainer
##########################################################
function create_tso_vars!( model_container::TSOBilevelTSOModelContainer,
                            network::Networks.Network,
                            target_timepoints::Vector{Dates.DateTime},
                            scenarios::Vector{String},
                            preceding_market_schedule::Schedule
                            )
    pilotables_list_l = Networks.get_generators_of_type(network, Networks.PILOTABLE)
    limitables_list_l = Networks.get_generators_of_type(network, Networks.LIMITABLE)
    buses_list_l = Networks.get_buses(network)
    add_pilotables_vars!(model_container,
                        pilotables_list_l, target_timepoints, scenarios,
                        imposition_vars=true, commitment_vars=true,
                        prefix="tso")
    add_limitables_vars!(model_container,
                        target_timepoints, scenarios,
                        limitables_list_l,
                        preceding_market_schedule,
                        delta_vars=true,
                        injection_vars=true, limit_vars=true, local_capping_vars=true, global_capping_vars=true,
                        prefix="tso"
                        )
    add_lol_vars!(model_container,
                target_timepoints, scenarios, buses_list_l,
                global_lol_vars=true, local_lol_vars=true,
                prefix="tso")
end

function add_tso_constraints!(bimodel_container::TSOBilevelModel,
                            target_timepoints, scenarios, network,
                            firmness, reference_schedule, generators_initial_state,
                            uncertainties_at_ech::UncertaintiesAtEch,
                            configs::TSOBilevelConfigs)
    model = bimodel_container.model
    tso_model_container::TSOBilevelTSOModelContainer = bimodel_container.upper
    market_model_container::TSOBilevelMarketModelContainer = bimodel_container.lower

    pilotables_list_l = Networks.get_generators_of_type(network, Networks.PILOTABLE)
    unit_commitment_constraints!(model,
                    tso_model_container.pilotable_model, pilotables_list_l,  target_timepoints, scenarios,
                    firmness, reference_schedule, generators_initial_state,
                    always_link_scenarios=configs.LINK_SCENARIOS_PILOTABLE_ON,
                    prefix="tso")
    power_imposition_constraints!(model, tso_model_container.pilotable_model,
                    pilotables_list_l, target_timepoints, scenarios,
                    firmness, reference_schedule;
                    always_link_scenarios=configs.LINK_SCENARIOS_PILOTABLE_LEVEL,
                    prefix="tso")

    limitables_list_l = Networks.get_generators_of_type(network, Networks.LIMITABLE)
    limitables_ids_l = map(lim_gen->Networks.get_id(lim_gen), limitables_list_l)
    limitable_model = tso_model_container.limitable_model
    limitable_power_constraints!(model,
                        limitable_model, limitables_list_l, target_timepoints, scenarios,
                        firmness, uncertainties_at_ech,
                        always_link_scenarios=configs.LINK_SCENARIOS_LIMIT,
                        prefix="tso")
    local_capping_constraints!(model,
                            limitable_model, limitables_list_l, target_timepoints, scenarios,
                            uncertainties_at_ech, prefix="tso")
    global_capping_constraints!(model,
                            limitable_model, limitables_list_l,
                            target_timepoints, scenarios,
                            uncertainties_at_ech,
                            prefix="tso")
    distribution_constraint!(model,
                            get_global_capping(market_model_container.limitable_model),
                            get_local_capping(tso_model_container.limitable_model),
                            limitables_ids_l, target_timepoints, scenarios,
                            cstr_prefix_name="tso_distribute_capping")
    distribution_constraint!(model,
                            market_model_container.limitable_model.p_global_injected, #TODO rename to avoid onfusion between localised injections and global ones
                            get_p_injected(tso_model_container.limitable_model),
                            limitables_ids_l, target_timepoints, scenarios,
                            cstr_prefix_name="tso_distribute_enr_injections")

    buses_list_l = Networks.get_buses(network)
    buses_ids_l = map(bus->Networks.get_id(bus), buses_list_l)
    lol_model = tso_model_container.lol_model
    local_lol_constraints!(model,
                        lol_model,
                        buses_list_l, target_timepoints, scenarios,
                        uncertainties_at_ech,
                        prefix="tso")
    global_lol_constraints!(model,
                        lol_model, buses_list_l, target_timepoints, scenarios,
                        uncertainties_at_ech,
                        prefix="tso")
    distribution_constraint!(model,
                            get_global_lol(market_model_container.lol_model),
                            get_local_lol(tso_model_container.lol_model),
                            buses_ids_l, target_timepoints, scenarios,
                            cstr_prefix_name="tso_distribute_lol")

    add_rso_flows_exprs(tso_model_container.flows,
                market_model_container.pilotable_model,
                tso_model_container.limitable_model,
                tso_model_container.lol_model,
                all_combinations(network, target_timepoints, scenarios),
                uncertainties_at_ech, network)
    combinations = (configs.CONSIDER_N_1_CSTRS) ? all_combinations(network, target_timepoints, scenarios) :
                                                    all_n_combinations(network, target_timepoints, scenarios)
    rso_constraints!(bimodel_container.model,
                    tso_model_container.rso_constraint,
                    tso_model_container.flows,
                    combinations,
                    network,
                    prefix="tso")

    return bimodel_container
end

function create_tso_objectives!(bimodel_container::TSOBilevelModel,
                                target_timepoints, scenarios, network,
                                preceding_market_schedule::Schedule,
                                capping_cost, loss_of_load_cost,
                                limit_penalty,
                                pilotable_bounding_cost,
                                use_prop_cost_for_bounding::Bool)
    tso_model_container = bimodel_container.upper
    market_model_container = bimodel_container.lower
    objective_model = tso_model_container.objective_model

    # cost for deviating from market schedule
    for (_, var_delta) in market_model_container.pilotable_model.delta_p
        objective_model.deltas += var_delta
    end
    for (_, var_delta) in tso_model_container.limitable_model.delta_p
        objective_model.deltas += var_delta
    end

    # objective_model.pilotable_cost
    add_tsobilevel_impositions_cost!(tso_model_container,
                                    target_timepoints, scenarios, network,
                                    preceding_market_schedule,
                                    pilotable_bounding_cost,
                                    use_prop_cost_for_bounding)

    # limitable_cost : capping (fr. ecretement)
    objective_model.limitable_cost += coeffxsum(tso_model_container.limitable_model.p_global_capping, capping_cost)

    # cost for cutting consumption (lol) and avoid limiting for no reason
    objective_model.penalty += coeffxsum(tso_model_container.lol_model.p_global_loss_of_load, loss_of_load_cost)
    objective_model.penalty += coeffxsum(tso_model_container.limitable_model.b_is_limited, limit_penalty)

    # cost for cutting load/consumption by the market
    objective_model.penalty += coeffxsum(market_model_container.lol_model.p_global_loss_of_load, loss_of_load_cost)

    #Objective 1 : add penalty ?
    objective_model.full_obj_1 = objective_model.deltas + objective_model.penalty
    #Objective 2 : add penalty ?
    objective_model.full_obj_2 = ( objective_model.pilotable_cost +
                                objective_model.limitable_cost +
                                objective_model.penalty )

    return tso_model_container
end

function add_tsobilevel_impositions_cost!(model_container::TSOBilevelTSOModelContainer,
                                        target_timepoints, scenarios, network,
                                        preceding_market_schedule,
                                        pilotable_bounding_cost,
                                        use_prop_cost_for_bounding::Bool)
    tso_pilotable_model = model_container.pilotable_model
    objective_expr = model_container.objective_model.pilotable_cost

    for gen in Networks.get_generators_of_type(network, Networks.PILOTABLE)
        gen_id = Networks.get_id(gen)
        if use_prop_cost_for_bounding
            cost = Networks.get_prop_cost(gen)
        else
            cost = pilotable_bounding_cost
        end

        for ts in target_timepoints
            for s in scenarios
                commitment = get_commitment_value(preceding_market_schedule, gen_id, ts, s)
                if  !Networks.needs_commitment(gen) || (!ismissing(commitment) && (commitment==ON))
                    add_tsobilevel_started_impositions_cost!(objective_expr, gen,
                                                            tso_pilotable_model.p_imposition_min[gen_id, ts, s],
                                                            tso_pilotable_model.p_imposition_max[gen_id, ts, s],
                                                            cost)
                else
                    add_tsobilevel_non_started_impositions_cost!(objective_expr, model_container,
                                                                 gen, ts, s, cost)
                end
            end
        end
    end

    return model_container
end
function add_tsobilevel_started_impositions_cost!(objective_expr::AffExpr,
                                                gen::Generator,
                                                pmin_var, pmax_var,
                                                pilotable_bounding_cost
                                                )
    p_max = Networks.get_p_max(gen)
    p_min = Networks.get_p_min(gen)

    add_to_expression!(objective_expr, pilotable_bounding_cost * (p_max - pmax_var))
    add_to_expression!(objective_expr, pilotable_bounding_cost * (pmin_var - p_min))

    return objective_expr
end
function add_tsobilevel_non_started_impositions_cost!(objective_expr::AffExpr,
                                                      tso_model_container::TSOBilevelTSOModelContainer,
                                                    gen, ts, s,
                                                    pilotable_bounding_cost
                                                    )
    #need a second expression for units that do not need a commitment (p_min=0 and no b_on)
    @assert Networks.needs_commitment(gen)

    tso_pilotable_model = tso_model_container.pilotable_model
    gen_id = Networks.get_id(gen)

    p_max = Networks.get_p_max(gen)
    # p_min = Networks.get_p_min(gen)

    #need to cost reducing pmax otherwise TSO may always limit pmax when starting a unit
    pmax_var = tso_pilotable_model.p_imposition_max[gen_id, ts, s]
    b_on_var = tso_pilotable_model.b_on[gen_id, ts, s]
    add_to_expression!(objective_expr, pilotable_bounding_cost * (p_max*b_on_var - pmax_var) )

    pmin_var = tso_pilotable_model.p_imposition_min[gen_id, ts, s]
    add_to_expression!(objective_expr, pilotable_bounding_cost * pmin_var)
    # add_to_expression!(objective_expr, pilotable_bounding_cost * p_min * b_on_var)

    return objective_expr
end

##########################################################
#        lower problem : Market functtons  : TSOBilevelMarketModelContainer
##########################################################

function create_market_vars!(model_container::TSOBilevelMarketModelContainer,
                            network::Networks.Network,
                            target_timepoints::Vector{Dates.DateTime},
                            scenarios::Vector{String},
                            preceding_market_schedule::Schedule
                            )

    pilotables_list_l = Networks.get_generators_of_type(network, Networks.PILOTABLE)
    add_pilotables_vars!(model_container,
                        pilotables_list_l, target_timepoints, scenarios,
                        preceding_market_schedule,
                        delta_vars=true,
                        injection_vars=true,
                        prefix="market_"
                        )
    add_limitables_vars!(model_container,
                        target_timepoints, scenarios,
                        global_capping_vars=true,
                        global_injection_vars=true,
                        prefix="market_"
                        )
    add_lol_vars!(model_container,
                target_timepoints, scenarios,
                global_lol_vars=true,
                prefix="market_")
end

function add_firmness_duals!(kkt_model_container::TSOBilevelKKTModelContainer,
                            pilotables_list, target_timepoints, scenarios,
                            firmness,
                            link_scenarios_pilotable_level_market;
                            prefix::String="")
    for gen_pilotable in pilotables_list
        gen_id = Networks.get_id(gen_pilotable)
        for ts in target_timepoints
            decision_firmness_l = get_power_level_firmness(firmness, gen_id, ts)
            for s in scenarios[2:end]
                if requires_linking(decision_firmness_l, link_scenarios_pilotable_level_market)
                    #create duals relative to firmness constraints of injected pilotable power in market
                    name = @sprintf("c_firmness[%s,%s,%s]",gen_id,ts,s)
                    add_dual!(kkt_model_container.model, kkt_model_container.firmness_duals,
                                (gen_id,ts,s), name, false)
                end
            end
        end
    end
    return kkt_model_container
end

function add_eod_constraints_duals!(kkt_model_container::TSOBilevelKKTModelContainer,
                            target_timepoints, scenarios)
    for ts in target_timepoints
        for s in scenarios
            #create duals relative to EOD constraint
            name = @sprintf("c_eod[%s,%s]",ts,s)
            add_dual!(kkt_model_container.model, kkt_model_container.eod_duals, (ts,s), name, false)
        end
    end
end

function add_min_global_capping_duals!(kkt_model_container::TSOBilevelKKTModelContainer,
                                    target_timepoints, scenarios)
    for ts in target_timepoints
        for s in scenarios
            #create duals and indicators relative to TSO min capping constraint
            name = @sprintf("c_min_e[%s,%s]",ts,s)
            add_dual_and_indicator!(kkt_model_container.model,
                                    kkt_model_container.capping_duals, kkt_model_container.capping_indicators, (ts,s),
                                    name, true)
        end
    end
end

function add_min_global_lol_duals!(market_model_container::TSOBilevelMarketModelContainer,
                                kkt_model_container::TSOBilevelKKTModelContainer,
                                target_timepoints, scenarios)
    for ts in target_timepoints
        for s in scenarios
            #create duals and indicators relative to TSO min LoL constraint
            name = @sprintf("c_min_lol[%s,%s]",ts,s)
            add_dual_and_indicator!(kkt_model_container.model,
                                    kkt_model_container.loss_of_load_duals, kkt_model_container.loss_of_load_indicators, (ts,s),
                                    name, true)
        end
    end
    return market_model_container
end
function add_impositions_duals!(kkt_model_container::TSOBilevelKKTModelContainer,
                            pilotables_list, target_timepoints, scenarios)
    for ts in target_timepoints
        for s in scenarios
            for gen in pilotables_list
                gen_id = Networks.get_id(gen)
                add_pilotable_pmin_duals!(kkt_model_container,
                                        gen_id, ts, s)
                add_pilotable_pmax_duals!(kkt_model_container,
                                        gen_id, ts, s)
            end
        end
    end

    return kkt_model_container
end
function add_pilotable_pmin_duals!(kkt_model_container::TSOBilevelKKTModelContainer,
                                gen_id, ts, s)
    #create duals and indicators relative to tso pmin constraint
    name = @sprintf("c_tso_pmin[%s,%s,%s]",gen_id,ts,s)
    add_dual_and_indicator!(kkt_model_container.model,
                            kkt_model_container.pmin_duals, kkt_model_container.pmin_indicators, (gen_id,ts,s),
                            name, true)

    return kkt_model_container
end
function add_pilotable_pmax_duals!(kkt_model_container::TSOBilevelKKTModelContainer,
                                gen_id, ts, s)
    #create duals and indicators relative to tso pmax constraint
    name = @sprintf("c_tso_pmax[%s,%s,%s]",gen_id,ts,s)
    add_dual_and_indicator!(kkt_model_container.model,
                            kkt_model_container.pmax_duals, kkt_model_container.pmax_indicators, (gen_id,ts,s),
                            name, true)

    return kkt_model_container
end

function add_market_constraints!(bimodel_container::TSOBilevelModel,
                            target_timepoints, scenarios, network,
                            firmness, #reference_schedule, generators_initial_state,
                            uncertainties_at_ech::UncertaintiesAtEch,
                            configs::TSOBilevelConfigs)
    model = bimodel_container.model
    tso_model_container::TSOBilevelTSOModelContainer = bimodel_container.upper
    market_model_container::TSOBilevelMarketModelContainer = bimodel_container.lower
    kkt_model_container::TSOBilevelKKTModelContainer = bimodel_container.kkt_model

    pilotables_list_l = Networks.get_generators_of_type(network, Networks.PILOTABLE)
    pilotable_power_constraints!(model,
                                market_model_container.pilotable_model,
                                pilotables_list_l, target_timepoints, scenarios,
                                firmness,
                                missing,
                                ignore_sequencing_cstrs=true, #Sequencing constraints are imposed in TSO upper problem
                                always_link_scenarios=configs.LINK_SCENARIOS_PILOTABLE_LEVEL_MARKET,
                                prefix="market_")
    add_firmness_duals!(kkt_model_container,
                        pilotables_list_l, target_timepoints, scenarios, firmness,
                        configs.LINK_SCENARIOS_PILOTABLE_LEVEL_MARKET,
                        prefix="market_")
    respect_impositions_constraints!(model,
                                market_model_container.pilotable_model, #Injections are in the lower problem
                                tso_model_container.pilotable_model, #Impositions are in the upper problem
                                pilotables_list_l,  target_timepoints, scenarios,
                                prefix="market_")
    add_impositions_duals!(kkt_model_container,
                        pilotables_list_l, target_timepoints, scenarios)

    limitables_list_l = Networks.get_generators_of_type(network, Networks.LIMITABLE)
    global_capping_constraints!(model,
                            market_model_container.limitable_model,
                            limitables_list_l, target_timepoints, scenarios,
                            uncertainties_at_ech,
                            min_cap=get_global_capping(tso_model_container.limitable_model)
                            )
    add_min_global_capping_duals!(kkt_model_container,
                            target_timepoints, scenarios)
    global_injected_constraints!(model,
                            market_model_container.limitable_model,
                            limitables_list_l, target_timepoints, scenarios,
                            uncertainties_at_ech, prefix="market_")

    buses_list_l = Networks.get_buses(network)
    lol_model = market_model_container.lol_model
    global_lol_constraints!(model,
                        lol_model, buses_list_l, target_timepoints, scenarios,
                        uncertainties_at_ech,
                        min_lol=get_global_lol(tso_model_container.lol_model),
                        prefix="market")
    add_min_global_lol_duals!(market_model_container, kkt_model_container,
                        target_timepoints, scenarios)

    eod_constraints!(market_model_container.model, market_model_container.eod_constraint,
                    market_model_container.pilotable_model,
                    tso_model_container.limitable_model,
                    market_model_container.lol_model,
                    target_timepoints, scenarios,
                    uncertainties_at_ech, network
                    )
    add_eod_constraints_duals!(kkt_model_container, target_timepoints, scenarios)

    return bimodel_container
end

function create_market_objectives!(model_container::TSOBilevelMarketModelContainer,
                                network,
                                capping_cost, loss_of_load_cost)
    objective_model = model_container.objective_model

    # model_container.pilotable_cost
    add_prop_cost!(model_container.objective_model.pilotable_cost,
                            model_container.pilotable_model.p_injected, network)

    # limitable_cost : capping (fr. ecretement)
    objective_model.limitable_cost += coeffxsum(model_container.limitable_model.p_global_capping, capping_cost)

    # cost for cutting load/consumption
    objective_model.penalty += coeffxsum(model_container.lol_model.p_global_loss_of_load, loss_of_load_cost)

    objective_model.full_obj = ( objective_model.pilotable_cost +
                                objective_model.limitable_cost +
                                objective_model.penalty )

    return model_container
end

##########################################################
#        kkt reformulation : TSOBilevelKKTModelContainer
##########################################################

function add_kkt_stationarity_constraints!(kkt_model::TSOBilevelKKTModelContainer,
                                            target_timepoints, scenarios, network,
                                            capping_cost, loss_of_load_cost)
    #FIXME can be generic by iterating on lower variables to construct each stationarity constraint
    # iterate on the objective and lower constraints to extract their coefficients, but need to link each cnstraint to its dual var
    add_capping_stationarity_constraints!(kkt_model, target_timepoints, scenarios, capping_cost)
    add_loss_of_load_stationarity_constraints!(kkt_model, target_timepoints, scenarios, loss_of_load_cost)
    add_pilotable_stationarity_constraints!(kkt_model, target_timepoints, scenarios, network)
end

function add_loss_of_load_stationarity_constraints!(kkt_model::TSOBilevelKKTModelContainer,
                                                target_timepoints, scenarios, loss_of_load_cost)
    for ts in target_timepoints
        for s in scenarios
            # @assert ( capping_cost ≈ coefficient(market_model.objective_model.full_obj,
            #                                                 market_model.lol_model.p_global_loss_of_load[ts,s]) )
            name = @sprintf("c_stationarity_lol[%s,%s]",ts,s)
            @constraint(kkt_model.model,
                        loss_of_load_cost + kkt_model.eod_duals[ts,s] - kkt_model.loss_of_load_duals[ts,s] == 0,
                        base_name=name)
        end
    end
end

function add_capping_stationarity_constraints!(kkt_model::TSOBilevelKKTModelContainer,
                                                target_timepoints, scenarios, capping_cost)
    for ts in target_timepoints
        for s in scenarios
            # @assert ( capping_cost ≈ coefficient(market_model.objective_model.full_obj,
            #                                                 market_model.limitable_model.p_global_capping[ts,s]) )
            name = @sprintf("c_stationarity_e[%s,%s]",ts,s)
            @constraint(kkt_model.model,
                        capping_cost - kkt_model.eod_duals[ts,s] - kkt_model.capping_duals[ts,s] == 0,
                        base_name=name)
        end
    end
end

function firmness_duals_sationarity_expr(kkt_model,
                                        gen_id, ts, scenario,
                                        scenarios)::AffExpr
    if length(scenarios) == 1
        result_l = 0.
    elseif scenario == scenarios[1]
        sum_l = sum( get(kkt_model.firmness_duals, (gen_id,ts,s), 0.) for s in scenarios[2:end])
        result_l = -sum_l
    else
        result_l = get(kkt_model.firmness_duals, (gen_id,ts,scenario), 0.)
    end
    return result_l
end
function add_pilotable_stationarity_constraints!(kkt_model::TSOBilevelKKTModelContainer,
                                                target_timepoints, scenarios, network)
    for pilotable_gen in Networks.get_generators_of_type(network, Networks.PILOTABLE)
        gen_id = Networks.get_id(pilotable_gen)
        gen_prop_cost = Networks.get_prop_cost(pilotable_gen)
        for ts in target_timepoints
            for s in scenarios
                # @assert ( pilotable_bounding_cost ≈ coefficient(market_model.objective_model.full_obj,
                #                                             market_model.pilotable_model.p_injected[gen_id,ts,s]) )
                name = @sprintf("c_stationarity_pilotable_p[%s,%s,%s]",gen_id,ts,s)
                firmness_sationarity_expr = firmness_duals_sationarity_expr(kkt_model,
                                                                                gen_id, ts, s,
                                                                                scenarios)
                @constraint(kkt_model.model,
                            0 == gen_prop_cost
                                    + kkt_model.eod_duals[ts,s]
                                    - kkt_model.pmin_duals[gen_id,ts,s]
                                    + kkt_model.pmax_duals[gen_id,ts,s]
                                    + firmness_sationarity_expr,
                            base_name=name)
            end
        end
    end
end

function add_kkt_complementarity_constraints!(model_container::TSOBilevelModel,
                                            big_m, target_timepoints, scenarios, network)
    #FIXME can be done iteratively and generically if we loop on constraint expressions and know their corresponding kkt vars
    add_emin_complementarity_constraints!(model_container, big_m, target_timepoints, scenarios)
    add_lolmin_complementarity_constraints!(model_container, big_m, target_timepoints, scenarios)
    add_pmin_complementarity_constraints!(model_container, big_m, target_timepoints, scenarios, network)
    add_pmax_complementarity_constraints!(model_container, big_m, target_timepoints, scenarios, network)
end

"""
    complementarity constraint linked to @ref PSCOPF.global_capping_constraints!
"""
function add_emin_complementarity_constraints!(model_container::TSOBilevelModel,
                                            big_m, target_timepoints, scenarios)
    tso_limitable_model = model_container.upper.limitable_model
    market_limitable_model = model_container.lower.limitable_model
    kkt_model = model_container.kkt_model

    for ts in target_timepoints
        for s in scenarios
            kkt_var = kkt_model.capping_duals[ts,s]
            cstr_expr = market_limitable_model.p_global_capping[ts,s] - tso_limitable_model.p_global_capping[ts,s]
            b_indicator = kkt_model.capping_indicators[ts,s]
            ub_cstr = compute_ub(cstr_expr, big_m)
            formulate_complementarity_constraints!(kkt_model.model, kkt_var, cstr_expr, b_indicator, big_m, ub_cstr)
        end
    end
    return model_container
end

"""
    complementarity constraint linked to @ref PSCOPF.global_lol_constraints!
"""
function add_lolmin_complementarity_constraints!(model_container::TSOBilevelModel,
                                                big_m, target_timepoints, scenarios)
    tso_lol_model = model_container.upper.lol_model
    market_lol_model = model_container.lower.lol_model
    kkt_model = model_container.kkt_model

    for ts in target_timepoints
        for s in scenarios
            kkt_var = kkt_model.loss_of_load_duals[ts,s]
            cstr_expr = market_lol_model.p_global_loss_of_load[ts,s] - tso_lol_model.p_global_loss_of_load[ts,s]
            b_indicator = kkt_model.loss_of_load_indicators[ts,s]
            ub_cstr = compute_ub(cstr_expr, big_m)
            formulate_complementarity_constraints!(kkt_model.model, kkt_var, cstr_expr, b_indicator, big_m, ub_cstr)
        end
    end
    return model_container
end

"""
    complementarity constraint linked to the p_min part of @ref PSCOPF.respect_impositions_constraints!
"""
function add_pmin_complementarity_constraints!(model_container::TSOBilevelModel,
                                                big_m, target_timepoints, scenarios, network)
    tso_pilotable_model = model_container.upper.pilotable_model
    market_pilotable_model = model_container.lower.pilotable_model
    kkt_model = model_container.kkt_model

    for ts in target_timepoints
        for s in scenarios
            for pilotable_gen in Networks.get_generators_of_type(network, Networks.PILOTABLE)
                gen_id = Networks.get_id(pilotable_gen)

                kkt_var = kkt_model.pmin_duals[gen_id,ts,s]
                cstr_expr = market_pilotable_model.p_injected[gen_id,ts,s] - tso_pilotable_model.p_imposition_min[gen_id,ts,s]
                b_indicator = kkt_model.pmin_indicators[gen_id,ts,s]
                ub_cstr = compute_ub(cstr_expr, big_m) #or get_p_max(pilotable_gen)
                formulate_complementarity_constraints!(kkt_model.model, kkt_var, cstr_expr, b_indicator, big_m, ub_cstr)
            end
        end
    end
    return model_container
end

"""
    complementarity constraint linked to the p_max part of @ref PSCOPF.respect_impositions_constraints!
"""
function add_pmax_complementarity_constraints!(model_container::TSOBilevelModel,
                                                big_m, target_timepoints, scenarios, network)
    tso_pilotable_model = model_container.upper.pilotable_model
    market_pilotable_model = model_container.lower.pilotable_model
    kkt_model = model_container.kkt_model

    for ts in target_timepoints
        for s in scenarios
            for pilotable_gen in Networks.get_generators_of_type(network, Networks.PILOTABLE)
                gen_id = Networks.get_id(pilotable_gen)

                kkt_var = kkt_model.pmax_duals[gen_id,ts,s]
                cstr_expr = tso_pilotable_model.p_imposition_max[gen_id,ts,s] - market_pilotable_model.p_injected[gen_id,ts,s]
                b_indicator = kkt_model.pmax_indicators[gen_id,ts,s]
                ub_cstr = compute_ub(cstr_expr, big_m) #or get_p_max(pilotable_gen)
                formulate_complementarity_constraints!(kkt_model.model, kkt_var, cstr_expr, b_indicator, big_m, ub_cstr)
            end
        end
    end
    return model_container
end

#############################
# Utils
#############################
function coeffxsum(vars_dict::AbstractDict{T,V}, coeff::Float64
                )::GenericAffExpr{Float64, VariableRef} where T <: Tuple where V <: AbstractVariableRef
    terms = [var_l=>coeff for (_, var_l) in vars_dict]
    result = GenericAffExpr{Float64, VariableRef}(0., terms)
    return result
end

##########################################################
#        TSOBilevel
##########################################################
function tso_bilevel(network::Networks.Network,
                    target_timepoints::Vector{Dates.DateTime},
                    generators_initial_state::SortedDict{String,GeneratorState},
                    scenarios::Vector{String},
                    uncertainties_at_ech::UncertaintiesAtEch,
                    firmness::Firmness,
                    preceding_market_schedule::Schedule,
                    preceding_tso_schedule::Schedule,
                    # gratis_starts::Set{Tuple{String,Dates.DateTime}},
                    configs::TSOBilevelConfigs
                    )

    @assert(configs.big_m >= configs.MARKET_LOL_PENALTY)
    @assert(configs.big_m >= configs.MARKET_CAPPING_COST)
    @assert(all( configs.big_m >= Networks.get_prop_cost(gen)
                for gen in Networks.get_generators_of_type(network, Networks.PILOTABLE) ))

    if is_market(configs.REF_SCHEDULE_TYPE_IN_TSO)
        reference_schedule = preceding_market_schedule
    elseif is_tso(configs.REF_SCHEDULE_TYPE_IN_TSO)
        reference_schedule = preceding_tso_schedule
    else
        throw( error("Invalid REF_SCHEDULE_TYPE_IN_TSO config.") )
    end

    bimodel_container_l = TSOBilevelModel()

    create_tso_vars!(bimodel_container_l.upper,
                    network, target_timepoints, scenarios,
                    preceding_market_schedule)
    create_market_vars!(bimodel_container_l.lower,
                        network, target_timepoints, scenarios,
                        preceding_market_schedule)

    #this is the expression no objective is added to the jump model
    create_market_objectives!(bimodel_container_l.lower, network,
                            configs.MARKET_CAPPING_COST, configs.MARKET_LOL_PENALTY)

    #constraints may use upper and lower vars at the same time
    add_tso_constraints!(bimodel_container_l, target_timepoints, scenarios, network,
                        firmness, reference_schedule, generators_initial_state,
                        uncertainties_at_ech, configs)
    add_market_constraints!(bimodel_container_l, target_timepoints, scenarios, network, firmness, uncertainties_at_ech, configs)
    #kkt stationarity
    add_kkt_stationarity_constraints!(bimodel_container_l.kkt_model,
                                    target_timepoints, scenarios, network,
                                    configs.MARKET_CAPPING_COST, configs.MARKET_LOL_PENALTY)
    #kkt complementarity
    add_kkt_complementarity_constraints!(bimodel_container_l, configs.big_m, target_timepoints, scenarios, network)

    create_tso_objectives!(bimodel_container_l,
                        target_timepoints, scenarios, network,
                        reference_schedule, #reference to see which units are currently on
                        configs.TSO_CAPPING_COST, configs.TSO_LOL_PENALTY,
                        configs.TSO_LIMIT_PENALTY,
                        configs.TSO_PILOTABLE_BOUNDING_COST, configs.USE_UNITS_PROP_COST_AS_TSO_BOUNDING_COST)

    launch_solve!(bimodel_container_l, configs)

    return bimodel_container_l
end

function launch_solve!(bimodel_container::TSOBilevelModel, configs::TSOBilevelConfigs)
    # first_iter = true
    # while(!isempty(to_add) || first_iter)
    #     first_iter = false
    #     solve!(bimodel_container, configs.problem_name, configs.out_path)

    #     to_add = verify_rso()
    #     add_constraints(to_add)
    # end

    #the upper problem's objective model holds the two objective expressions and the deltas expression
    if configs.CONSIDER_DELTAS
        solve_2steps_deltas!(bimodel_container.upper, configs)
    else
        obj = bimodel_container.upper.objective_model.full_obj_2
        @objective(get_model(bimodel_container), Min, obj)
        solve!(bimodel_container, configs.problem_name, configs.out_path)
    end
    @info("Lower Objective Value : $(value(bimodel_container.lower.objective_model.full_obj))")
end


module BenefitCap
#
# Apply the benefit cap as described in CPAG 19/20
# ch 52. Somewhat impressionistic and could do
# with another go.
#
using ScottishTaxBenefitModel

using .ModelHousehold:
    BenefitUnit,
    Household,
    Person

using .STBParameters:
    BenefitCapSys

using .Results:
    BenefitUnitResult

using .Definitions

using .Intermediate:
    MTIntermediate

using .STBIncomes

export apply_benefit_cap!

"""
Apply a benefit cap to a benefit unit.
"""
function apply_benefit_cap!(
    benefit_unit_result :: BenefitUnitResult,
    region           :: Standard_Region,
    benefit_unit     :: BenefitUnit,
    intermed         :: MTIntermediate,
    caps             :: BenefitCapSys,
    route            :: LegacyOrUC )
    # println( "apply_benefit_cap entered; route = $route")
    bu = benefit_unit # shortcut
    bur = benefit_unit_result # shortcut

    if route == legacy_bens
        for pid in bu.adults
            if bur.pers[pid].income[WORKING_TAX_CREDIT] > 0
                # cpag 21/2 p 1182
                # println("wtc bailing out")
                return
            end
        end
    else
        gross_earnings = 0.0
        for pid in bu.adults
            gross_earnings += isum( bur.pers[pid].income, [WAGES,SELF_EMPLOYMENT_INCOME] )
        end
        if gross_earnings >= caps.uc_incomes_limit
            # println("gross earn; bailing out ")
            return
        end
    end

    if intermed.someone_pension_age ||
        intermed.someone_is_carer ||
        (intermed.num_severely_disabled_adults > 0) ||
        caps.abolished
        # println("pension age bailing out")
        return
    end

    cap = intermed.num_people == 1 ?
        caps.outside_london_single :
        caps.outside_london_couple
    if region == London
        cap = intermed.num_people == 1 ?
            caps.inside_london_single :
            caps.inside_london_couple
    end
    # println("got cap as $cap")
    totbens = 0.0
    included = UC_CAP_BENEFITS
    target_ben = UNIVERSAL_CREDIT
    min_amount = bur.uc.childcare_costs
    if route == legacy_bens
        included = LEGACY_CAP_BENEFITS
        target_ben = HOUSING_BENEFIT
        min_amount = 0.5
    end
    recip_pers :: BigInt = -1
    recip_ben = 0.0

    for pid in bu.adults
        totbens += isum( bur.pers[pid].income, included )
        if bur.pers[pid].income[target_ben] > 0
            recip_pers = pid
            recip_ben = bur.pers[pid].income[target_ben]
        end
    end
    if recip_ben == 0.0
        # println("uc/hb0; returning ")
        return
    end
    excess = totbens - cap
    # println("totbens=$totbens cap=$cap excess=$excess")
    if excess > min_amount
        rd = max( min_amount, recip_ben - excess )
        bur.bencap.reduction = recip_ben - rd
        bur.pers[recip_pers].income[target_ben] = rd
    end
    bur.bencap.cap_benefits = totbens
    bur.bencap.cap = cap
end # cap_benefits

end # module


function update!(mc::MarkovChain{T},
                 mcmc::MCMC{T},
                 Σ::Array{Float64, 2},
                 σ::Float64;
                 condition_on_network::Bool=false,
                 event_batches::Int64=1,
                 transmission_rate_cache::Union{Nothing, TransmissionRateCache}=nothing) where T <: EpidemicModel
  # Initialize
  new_tr_cache = transmission_rate_cache
  new_events = mc.events[end]
  new_events_array = new_events[_state_progressions[T][2:end]]
  new_params = mc.risk_parameters[end]
  new_network = mc.transmission_network[end]
  new_lposterior = mc.log_posterior[end]
  # Randomize event time augmentation
  event_indices = findall(new_events_array[:] .> -Inf)
  aug_order = sample(event_indices, length(event_indices), replace=false)
  @debug "$(length(aug_order)) events to augment"
  if event_batches < 0
    @error "Cannot have negative amount of event batches"
  end
  if event_batches > length(aug_order)
    @warn "More event batches than there are events to augment ($(event_batches) > $(length(aug_order))), setting to maximum ($(length(aug_order)))"
    event_batches = length(aug_order)
  end
  batch_size = fld(length(aug_order), event_batches)
  @debug "Performing data augmentation in batches of $batch_size events at a time"
  for i = 1:(1 + event_batches)
    if i == 1
      # Propose new risk parameters
      proposed_events = new_events
      proposed_events_array = new_events_array
      proposed_params = generate(RiskParameters{T}, new_params, Σ)
      @debug "Generating new TransmissionRateCache for TN-ILM parameter proposal"
      proposed_tr_cache = TransmissionRateCache(size(new_events_array, 1))
    else
      proposed_events_array = copy(new_events_array)
      proposed_events = Events{T}(proposed_events_array)
      for j = (batch_size*(i-2) + 1):min(batch_size*(i-1),length(aug_order))
        id, state_index = Tuple(CartesianIndices((new_events.individuals,
                                            length(_state_progressions[T][2:end])))[aug_order[j]])
        new_state = _state_progressions[T][state_index+1]
        time = new_events[new_state][id]

        # Conditioning on network means that only event times valid under current network will be proposed.
        # This is useful for models which may have additional contributions to the posterior based on network,
        # and require more modest proposals (e.g. phylodynamic models).
        if condition_on_network
          proposed_event = generate(Event,
                                    Event{T}(time, id, new_state),
                                    σ,
                                    mcmc.event_extents,
                                    mcmc.event_observations,
                                    proposed_events,
                                    new_network)
        else
          proposed_event = generate(Event,
                                    Event{T}(time, id, new_state),
                                    σ,
                                    mcmc.event_extents,
                                    mcmc.event_observations,
                                    proposed_events)
        end
        proposed_events_array = reshape([proposed_events_array[1:(aug_order[j]-1)]
                                         _time(proposed_event)
                                         proposed_events_array[(aug_order[j]+1):end]],
                                        size(new_events_array))
        proposed_events = Events{T}(proposed_events_array)
      end
      proposed_params = new_params
      proposed_tr_cache = new_tr_cache
    end
    proposed_lprior = logpriors(proposed_params, mcmc.risk_priors)
    # Based on the logprior and competiting MCMC iteration, this loglikelihood is required for acceptance
    # Calculating this in advance allows us to cut loglikelihood calculation short if it goes below threshold
    ll_acceptance_threshold = log(rand()) + new_lposterior - proposed_lprior
    if ll_acceptance_threshold < Inf
      proposed_llikelihood = loglikelihood(proposed_params,
                                           mcmc.risk_functions,
                                           proposed_events,
                                           mcmc.population,
                                           mcmc.starting_states,
                                           transmission_rates_output = false,
                                           transmissions_output = false,
                                           early_decision_value = ll_acceptance_threshold,
                                           transmission_rate_cache = proposed_tr_cache)[1]
      proposed_lposterior = proposed_lprior + proposed_llikelihood
    else
      proposed_llikelihood = -Inf
      proposed_lposterior = -Inf
    end
    if proposed_llikelihood >= ll_acceptance_threshold
      @debug "MCMC proposal accepted (acceptance probability = $(round(min(1.0, exp(proposed_lposterior - new_lposterior)), digits=3)))"
      new_params = proposed_params
      new_tr_cache = proposed_tr_cache
      new_events = proposed_events
      new_events_array = proposed_events_array
      new_lposterior = proposed_lposterior
    else
      @debug "next!: MCMC proposal rejected (acceptance probability = $(round(min(1.0, exp(proposed_lposterior - new_lposterior)), digits=3)))"
    end
  end
  mc.iterations += 1
  tnr, tx = loglikelihood(new_params,
                          mcmc.risk_functions,
                          new_events,
                          mcmc.population,
                          mcmc.starting_states,
                          loglikelihood_output = false,
                          transmission_rate_cache = new_tr_cache)[[2; 3]]
  new_network = generate(TransmissionNetwork, tnr, mcmc, tx)
  push!(mc.events, new_events)
  push!(mc.transmission_network, new_network)
  push!(mc.risk_parameters, new_params)
  push!(mc.log_posterior, new_lposterior)
  return mc
end

function update!(mcmc::MCMC{T},
                 Σ::Array{Float64, 2},
                 σ::Float64) where T <: EpidemicModel
  @simd for mc in mcmc.markov_chains
    next!(mc, mcmc, Σ, σ)
  end
  return mcmc
end


# operation_type: Forward Modeling, RTM, FWI, TWRI
function H5Modeling(;
  # models
  h5vel,
  h5density=nothing,
  h5qualityFactor=nothing,
  nb::Real=40,  # number of absorbing boundaries points on each side
  h5epsilon=nothing,
  h5delta=nothing,
  h5tetha=nothing,
  h5phi=nothing,
  model_xkey::String="CDP_X",
  model_ykey::String="CDP_Y",

  # geometry
  h5geom=nothing,
  geom_src_pkey::String="SP",
  geom_src_xkey::String="SRCX",
  geom_src_ykey::String="SRCY",
  geom_src_zkey::String="SES",
  geom_rec_xkey::String="GRPX",
  geom_rec_ykey::String="GRPY",
  geom_rec_zkey::String="RGE",

  # judi
  opt::JUDI.JUDIOptions,

  # common settings
  computation_type::String,
  # LSRTM
  lsrtm_niter::Int=nothing,
  lsrtm_batchsize::Int=nothing,
  # FWI
  fwi_niter::Int=nothing,
  fwi_batchsize::Int=nothing,
  fwi_vmin::Number,
  fwi_vmax::Number,
  # TWRI
  twri_opt::TWRIOptions=nothing
  )

  # pylogging = pyimport("logging")
  h5gt = pyimport("h5gtpy._h5gt")
  h5geo = pyimport("h5geopy._h5geo")

  @info "H5Modeling started\n"

  if isnothing(h5vel)
    @error "No velocity model\n"
    return
  end

  if isnothing(opt)
    @error "No JUDI Options for modeling\n"
    return
  end

  if !isnothing(opt.dt_comp) && opt.dt_comp <= 0
    @error "dt_comp can't be less or equal to zero\n"
    return
  end

  # prepare physical parameters
  phpvel, model_orientation = H5ReadPhysicalParameter(
    h5vel, phptype=VELOCITY, xkey=model_xkey, ykey=model_ykey)
  if isnothing(phpvel)
    @error "Unable to read PhysicalParameter: VELOCITY\n"
    return
  end

  # declare PHP vars to be able use thm as arguments (or unknown var error will appear)
  phpdensity = nothing
  phpqualityFactor = nothing
  phpepsilon = nothing
  phpdelta = nothing
  phptheta = nothing
  phpphi = nothing

  if !isnothing(h5density)
    phpdensity, model_orientation = H5ReadPhysicalParameter(
      h5density, phptype=DENSITY, xkey=model_xkey, ykey=model_ykey)
    if isnothing(phpdensity)
      @error "Unable to read PhysicalParameter: DENSITY\n"
      return
    end
  end

  if !isnothing(h5qualityFactor)
    phpqualityFactor, model_orientation = H5ReadPhysicalParameter(
      h5qualityFactor, phptype=QUALITYFACTOR, xkey=model_xkey, ykey=model_ykey)
    if isnothing(phpqualityFactor)
      @error "Unable to read PhysicalParameter: QUALITYFACTOR\n"
      return
    end
  end

  if !isnothing(h5epsilon)
    phpepsilon, model_orientation = H5ReadPhysicalParameter(
      h5epsilon, phptype=EPSILON, xkey=model_xkey, ykey=model_ykey)
    if isnothing(phpepsilon)
      @error "Unable to read PhysicalParameter: EPSILON\n"
      return
    end
  end

  if !isnothing(h5delta)
    phpdelta, model_orientation = H5ReadPhysicalParameter(
      h5delta, phptype=DELTA, xkey=model_xkey, ykey=model_ykey)
    if isnothing(phpdelta)
      @error "Unable to read PhysicalParameter: DELTA\n"
      return
    end
  end

  if !isnothing(h5tetha)
    phptheta, model_orientation = H5ReadPhysicalParameter(
      h5tetha, phptype=THETA, xkey=model_xkey, ykey=model_ykey)
    if isnothing(phptheta)
      @error "Unable to read PhysicalParameter: THETA\n"
      return
    end
  end

  if !isnothing(h5phi)
    phpphi, model_orientation = H5ReadPhysicalParameter(
      h5phi, phptype=PHI, xkey=model_xkey, ykey=model_ykey)
    if isnothing(phpphi)
      @error "Unable to read PhysicalParameter: PHI\n"
      return
    end
  end

  model = JUDI.Model(phpvel.n, phpvel.d, phpvel.o, phpvel.data,
          rho=isnothing(phpdensity) ? nothing : phpdensity.data,
          qp=isnothing(phpqualityFactor) ? nothing : phpqualityFactor.data,
          epsilon=isnothing(phpepsilon) ? nothing : phpepsilon.data,
          delta=isnothing(phpdelta) ? nothing : phpdelta.data,
          theta=isnothing(phptheta) ? nothing : phptheta.data,
          phi=isnothing(phpphi) ? nothing : phpphi.data,
          nb=nb)

  # Geometry
  if isnothing(h5geom)
    @error "Geometry H5Seis is NULL\n"
    return
  end

  model_origin_x = model.o[1]
  model_origin_y = 0
  if length(model.o) == 3
    model_origin_y = model.o[2]
  end

  con = H5SeisCon(seis=h5geom, pkey=geom_src_pkey)
  recGeometry = H5GeometryOOC(
    h5geo=h5geo,
    container=con,
    key="receiver",
    xkey=geom_rec_xkey,
    ykey=geom_rec_ykey,
    zkey=geom_rec_zkey,
    do_coord_transform=true,
    model_origin_x=model_origin_x,
    model_origin_y=model_origin_y,
    model_orientation=model_orientation)
  srcGeometry = H5GeometryOOC(
    h5geo=h5geo,
    container=con,
    key="source",
    xkey=geom_src_xkey,
    ykey=geom_src_ykey,
    zkey=geom_src_zkey,
    do_coord_transform=true,
    model_origin_x=model_origin_x,
    model_origin_y=model_origin_y,
    model_orientation=model_orientation)

  # recGeometry = Geometry(
  #   container=con,
  #   xkey=geom_rec_xkey,
  #   ykey=geom_rec_ykey,
  #   zkey=geom_rec_zkey,
  #   do_coord_transform=true,
  #   model_origin_x=model_origin_x,
  #   model_origin_y=model_origin_y,
  #   model_orientation=model_orientation)
  # srcGeometry = Geometry(
  #   container=con,
  #   xkey=geom_src_xkey,
  #   ykey=geom_src_ykey,
  #   zkey=geom_src_zkey,
  #   do_coord_transform=true,
  #   model_origin_x=model_origin_x,
  #   model_origin_y=model_origin_y,
  #   model_orientation=model_orientation)

  if isnothing(srcGeometry) || isnothing(srcGeometry)
    @error "Unable to prepare Source and Receiver JUDI Geometry objects from H5Seis\n"
    return
  end

  # setup wavelet
  dt = Float32(abs(h5geom.getSampRate("ms")))
  nt = Integer(h5geom.getNSamp())
  t = Float32((nt-1)*dt)
  wavelet = ricker_wavelet(t, dt, opt.f0)
  q = judiVector(srcGeometry, wavelet)

  # JUDI (Devito) doesn't create a dir to save data
  if opt.save_data_to_disk
    mkpath(opt.file_path)
  end

  # simply to print the value
  ntComp = get_computational_nt(srcGeometry, recGeometry, model, dt=opt.dt_comp)
  dtComp = calculate_dt(model)

  # print modeling info
  @info "computation_type: $computation_type\n"
  @info "ntComp: $ntComp\n"
  @info "dtComp: $dtComp\n"
  @info "model settings:\n"
  @info "model.o: $(model.o)\n"
  @info "model.n: $(model.n)\n"
  @info "model.d: $(model.d)\n"
  @info "model_xkey: $model_xkey\n"
  @info "model_ykey: $model_ykey\n"
  @info "model_origin_x: $model_origin_x\n"
  @info "model_origin_y: $model_origin_y\n"
  @info "model_orientation: $model_orientation\n"
  @info "geometry settings:\n"
  @info "geom_src_xkey: $geom_src_xkey\n"
  @info "geom_src_ykey: $geom_src_ykey\n"
  @info "geom_src_zkey: $geom_src_zkey\n"
  @info "geom_rec_xkey: $geom_rec_xkey\n"
  @info "geom_rec_ykey: $geom_rec_ykey\n"
  @info "geom_rec_zkey: $geom_rec_zkey\n"
  @info "dt: $dt\n"
  @info "t: $t\n"
  @info "src_frq (kHz): $(opt.f0)"
  @info "nsrc: $(length(con))\n"
  @info "pkey: $(con.pkey)\n"
  @info "pkeyvals: $(con.pkeyvals)\n"
  @info "options: $opt\n"

  if computation_type == "Forward Modeling"
    H5ForwardModeling(model=model, opt=opt, q=q, recGeometry=recGeometry)
  elseif computation_type == "RTM"
    dobs = judiVector(recGeometry, con)
    mkpath(opt.file_path)
    H5RTM(model=model, opt=opt, q=q, dobs=dobs)
  elseif computation_type == "LSRTM"
    @error "LSRTM not ready yet\n"
  elseif computation_type == "FWI"
    dobs = judiVector(recGeometry, con)
    mkpath(opt.file_path)
    H5FWI(model=model, opt=opt, q=q, dobs=dobs,
          niterations=fwi_niter, batchsize=fwi_batchsize,
          vmin=fwi_vmin, vmax=fwi_vmax)
  elseif computation_type == "TWRI"
    @error "TWRI not ready yet\n"
  end

end


@doc raw"""
    AbstractDecoratorManifold{𝔽} <: AbstractManifold{𝔽}

Declare a manifold to be an abstract decorator.
A manifold which is a subtype of is a __decorated manifold__, i.e. has

* certain additional properties or
* delegates certain properties to other manifolds.

Most prominently, a manifold might be an embedded manifold, i.e. points on a manifold ``\mathcal M``
are represented by (some, maybe not all) points on another manifold ``\mathcal N``.
Depending on the type of embedding, several functions are dedicated to the embedding.
For example if the embedding is isometric, then the [`inner`](@ref) does not have to be
implemented for ``\mathcal M`` but can be automatically implemented by deligation to ``\mathcal N``.

This is modelled by the `AbstractDecoratorManifold` and traits. These are mapped to functions,
which determine the types of transparencies.
"""
abstract type AbstractDecoratorManifold{𝔽} <: AbstractManifold{𝔽} end

"""
    AbstractTrait

An abstract trait type to build a sequence of traits
"""
abstract type AbstractTrait end

"""
    EmptyTrait <: AbstractTrait

A Trait indicating that no feature is present.
"""
struct EmptyTrait <: AbstractTrait end

"""
    IsExplicitDecorator <: AbstractTrait

Specify that a certain type should dispatch per default to its [`decorated_manifold`](@ref).

!!! note
    Any decorator _behind_ this decorator might not have any effect, since the function
    dispatch is moved to its field at this point. Therefore this decorator should always be
    _last_ in the [`TraitList`](@ref).
"""
struct IsExplicitDecorator <: AbstractTrait end

"""
    TraitList <: AbstractTrait

Combine two traits into a combined trait.  Note that this introduces a preceedence.
the first of the traits takes preceedence if a trait is implemented for both functions.

# Constructor

    TraitList(head::AbstractTrait, tail::AbstractTrait)
"""
struct TraitList{T1<:AbstractTrait,T2<:AbstractTrait} <: AbstractTrait
    head::T1
    tail::T2
end

function Base.show(io::IO, t::TraitList)
    return print(io, "TraitList(", t.head, ", ", t.tail, ")")
end

"""
    active_traits(f, args...)

Return the list of traits applicable to the given call of function `f``. This function should be
overloaded for specific function calls.
"""
@inline active_traits(f, args...) = EmptyTrait()

"""
    merge_traits(t1, t2, trest...)

Merge two traits into a nested list of traits. Note that this takes trait preceedence into account,
i.e. `t1` takes preceedence over `t2` is any operations.
It always returns either ab [`EmptyTrait`](@ref) or a [`TraitList`](@ref).

This means that for
* one argument it just returns the trait itself if it is list-like, or wraps the trait in a
    single-element list otherwise,
* two arguments that are list-like, it merges them,
* two arguments of which only the first one is list-like and the second one is not,
    it appends the second argument to the list,
* two arguments of which only the second one is list-like, it prepends the first one to
    the list,
* two arguments of which none is list-like, it creates a two-element list.
* more than two arguments it recursively performs a left-assiciative recursive reduction
    on arguments, that is for example `merge_traits(t1, t2, t3)` is equivalent to
    `merge_traits(merge_traits(t1, t2), t3)`
"""
merge_traits()

@inline merge_traits() = EmptyTrait()
@inline merge_traits(t::EmptyTrait) = t
@inline merge_traits(t::TraitList) = t
@inline merge_traits(t::AbstractTrait) = TraitList(t, EmptyTrait())
@inline merge_traits(t1::EmptyTrait, ::EmptyTrait) = t1
@inline merge_traits(::EmptyTrait, t2::AbstractTrait) = merge_traits(t2)
@inline merge_traits(t1::AbstractTrait, t2::EmptyTrait) = TraitList(t1, t2)
@inline merge_traits(t1::AbstractTrait, t2::TraitList) = TraitList(t1, t2)
@inline merge_traits(::EmptyTrait, t2::TraitList) = t2
@inline function merge_traits(t1::AbstractTrait, t2::AbstractTrait)
    return TraitList(t1, TraitList(t2, EmptyTrait()))
end
@inline merge_traits(t1::TraitList, ::EmptyTrait) = t1
@inline function merge_traits(t1::TraitList, t2::AbstractTrait)
    return TraitList(t1.head, merge_traits(t1.tail, t2))
end
@inline function merge_traits(t1::TraitList, t2::TraitList)
    return TraitList(t1.head, merge_traits(t1.tail, t2))
end
@inline function merge_traits(
    t1::AbstractTrait,
    t2::AbstractTrait,
    t3::AbstractTrait,
    trest::AbstractTrait...,
)
    return merge_traits(merge_traits(t1, t2), t3, trest...)
end

"""
    parent_trait(t::AbstractTrait)

Return the parent trait for trait `t`, that is the more general trait whose behaviour it
inherits as a fallback.
"""
@inline parent_trait(::AbstractTrait) = EmptyTrait()

@inline function trait(f::TF, args...) where {TF}
    bt = active_traits(f, args...)
    return expand_trait(bt)
end

"""
    expand_trait(t::AbstractTrait)

Expand given trait into an ordered [`TraitList`](@ref) list of traits with their parent
traits obtained using [`parent_trait`](@ref).
"""
expand_trait(::AbstractTrait)

@inline expand_trait(e::EmptyTrait) = e
@inline expand_trait(t::AbstractTrait) = merge_traits(t, expand_trait(parent_trait(t)))
@inline function expand_trait(t::TraitList)
    et1 = expand_trait(t.head)
    et2 = expand_trait(t.tail)
    return merge_traits(et1, et2)
end

"""
    next_trait(t::AbstractTrait)

Return the next trait to consider, which by default is no following trait (i.e. [`EmptyTrait`](@ref)).

Expecially for a a [`TraitList`](@ref) this function returns the (remaining) tail of
the remaining traits.
"""
next_trait(::AbstractTrait) = EmptyTrait()

@inline next_trait(t::TraitList) = t.tail

#! format: off
# turn formatting for for the following functions
# due to the if with returns inside (formatter puts a return upfront the if)
function _split_signature(sig::Expr)
    if sig.head == :where
        where_exprs = sig.args[2:end]
        call_expr = sig.args[1]
    elseif sig.head == :call
        where_exprs = []
        call_expr = sig
    else
        error("Incorrect syntax in $sig. Expected a :where or :call expression.")
    end
    fname = call_expr.args[1]
    if isa(call_expr.args[2], Expr) && call_expr.args[2].head == :parameters
        # we have keyword arguments
        callargs = call_expr.args[3:end]
        kwargs_list = call_expr.args[2].args
    else
        callargs = call_expr.args[2:end]
        kwargs_list = []
    end
    argnames = map(callargs) do arg
        if isa(arg, Expr) && arg.head === :kw # default val present
            arg = arg.args[1]
        end
        if isa(arg, Expr) && arg.head === :(::) # typed
            return arg.args[1]
        end
        return arg
    end
    argtypes = map(callargs) do arg
        if isa(arg, Expr) && arg.head === :kw # default val present
            arg = arg.args[1]
        end
        if isa(arg, Expr)
            return arg.args[2]
        else
            return Any
        end
    end

    kwargs_call = map(kwargs_list) do kwarg
        if kwarg.head === :...
            return kwarg
        else
            if isa(kwarg.args[1], Symbol)
                kwargname = kwarg.args[1]
            else
                kwargname = kwarg.args[1].args[1]
            end
            return :($kwargname = $kwargname)
        end
    end

    return (;
        fname = fname,
        where_exprs = where_exprs,
        callargs = callargs,
        kwargs_list = kwargs_list,
        argnames = argnames,
        argtypes = argtypes,
        kwargs_call = kwargs_call,
    )
end
#! format: on

macro invoke_maker(argnum, type, sig)
    parts = ManifoldsBase._split_signature(sig)
    kwargs_list = parts[:kwargs_list]
    callargs = parts[:callargs]
    fname = parts[:fname]
    where_exprs = parts[:where_exprs]
    argnames = parts[:argnames]
    argtypes = parts[:argtypes]
    kwargs_call = parts[:kwargs_call]

    return esc(
        quote
            function ($fname)($(callargs...); $(kwargs_list...)) where {$(where_exprs...)}
                return invoke(
                    $fname,
                    Tuple{
                        $(argtypes[1:(argnum - 1)]...),
                        $type,
                        $(argtypes[(argnum + 1):end]...),
                    },
                    $(argnames...);
                    $(kwargs_call...),
                )
            end
        end,
    )
end

macro trait_function(sig, opts = :())
    parts = ManifoldsBase._split_signature(sig)
    kwargs_list = parts[:kwargs_list]
    callargs = parts[:callargs]
    fname = parts[:fname]
    where_exprs = parts[:where_exprs]
    argnames = parts[:argnames]
    argtypes = parts[:argtypes]
    kwargs_call = parts[:kwargs_call]

    argnametype_exprs = [:(typeof($(argname))) for argname in argnames]

    block = quote
        @inline function ($fname)($(callargs...); $(kwargs_list...)) where {$(where_exprs...)}
            return ($fname)(trait($fname, $(argnames...)), $(argnames...); $(kwargs_call...))
        end
        @inline function ($fname)(
            t::TraitList,
            $(callargs...);
            $(kwargs_list...),
        ) where {$(where_exprs...)}
            return ($fname)(next_trait(t), $(argnames...); $(kwargs_call...))
        end
        @inline function ($fname)(
            t::TraitList{IsExplicitDecorator},
            $(callargs...);
            $(kwargs_list...),
        ) where {$(where_exprs...)}
            arg1 = decorated_manifold($(argnames[1]))
            argt1 = typeof(arg1)
            return invoke(
                $fname,
                Tuple{argt1,$(argnametype_exprs[2:end]...)},
                arg1,
                $(argnames[2:end]...);
                $(kwargs_call...),
            )
        end
    end
    if !(:no_empty in opts.args)
        block = quote
            $block
            @inline function ($fname)(
                ::EmptyTrait,
                $(callargs...);
                $(kwargs_list...),
            ) where {$(where_exprs...)}
                return invoke(
                    $fname,
                    Tuple{supertype($(argtypes[1])),$(argnametype_exprs[2:end]...)},
                    $(argnames...);
                    $(kwargs_call...),
                )
            end
        end
    end
    return esc(block)
end


__precompile__()
#module SeisXcorrelation

using SeisIO, SeisNoise, Dates, FFTW, JLD2, Distributed, Sockets, ORCA

# read and write using jld2
include("io.jl")
# generate, sort, and classify station pairs
include("pairing.jl")
# miscellaneous utilities such as distance computation, windowing, and normalization
include("utils.jl")
# splitting utilities for high-order cross-correlation coda extraction
include("partition.jl")
# FFT wrapper for high-order cross-correlation
include("fft.jl")
# cross-correlation wrapper for high-order cross-correlation
include("correlate.jl")

export seisxcorrelation, seisxcorrelation_highorder

"""
    seisxcorrelation(data::Dict, tstamp::String, InputDict::Dict)

Compute cross-correlation save data in jld2 file with CorrData format.

# Arguments
- `data`::Dict    : Dictionary containing input data in the form: tstamp/stn => SeisChannel
- `tstamp::String,`    : Time stamp to read from JLD2 file.
- `InputDict::Dict`    : Dictionary containing IO, FFT, xcorr, and stacking parameters

# Output
- `tserrorList`::Array{String,1}    : Array containing the name of failed cross-correlation timestamp/stationpairs
- `basefoname.tstamp.jld2`    : contains CorrData structure with a hierarchical structure (CC function, metadata)

"""
function seisxcorrelation(data::Dict, tstamp::String, InputDict::Dict;verbose=false)
    # IO parameters
    basefoname = InputDict["basefoname"]
    time_unit  = InputDict["timeunit"]
    # FFT parameters
    freqmin    = Float64(InputDict["freqmin"])
    freqmax    = Float64(InputDict["freqmax"])
    fs         = Float64(InputDict["fs"])
    cc_len     = InputDict["cc_len"]
    cc_step    = InputDict["cc_step"]
    to_whiten  = InputDict["to_whiten"]
    time_norm  = InputDict["time_norm"]
    half_win   = InputDict["half_win"]
    water_level   = InputDict["water_level"]
    # correlation parameters
    corrtype   = InputDict["corrtype"]   # "xcorr," "acorr," or "xchancorr" (or any combination thereof)
    corrmethod = InputDict["corrmethod"] # "cross-correlation", "deconv", or "coherence"
    maxtimelag = InputDict["maxtimelag"]
    # stacking parameters
    stack      = InputDict["allstack"]
    stacktype  = InputDict["stacktype"]
    # dictionary to cache FFTs
    FFTDict = Dict{String, FFTData}()
    # list of stations that failed for this timestep
    tserrorList = []

    # split dataset names (keys of data) by "/" to get station list
    # assume form "$tstamp/$station"
    dsk = collect(keys(data))
    stlist = sort([string(split.(dsk[i], "/")[2]) for i=1:length(dsk)])

    # create output file for each time stamp, fill relevant info
    # outFile = jldopen("$basefoname.$tstamp.jld2", "a+")
    outFile = jldopen("$basefoname.$tstamp.jld2", true, true, true, IOStream)
    outFile["info/stationlist"] = stlist
    outFile["info/timeunit"] = time_unit

    println("$tstamp: Computing cross-correlations")
    stniter = 0 # counter to prevent computing duplicate xcorrs with reversed order
    # iterate over station list
    for stn1 in stlist
        stniter+=1

        # don't attempt FFT if this failed already
        if stn1 in tserrorList continue end

        # read station SeisChannels into SeisData before FFT
        S1 = try
            SeisData(data["$tstamp/$stn1"])
        catch;
            push!(tserrorList, "$tstamp/$stn1")
            filter!(a->a≠stn1, stlist)
            println("$tstamp: $stn1 encountered an error on read. Skipping.")
            continue
        end

        if S1.n > 1 @warn "SeisData contains multiple channels. Operating only on the first." end
        delete!(S1[1].misc, "kurtosis")
        delete!(S1[1].misc, "eqtimewindow")

        # make sure the data is the proper length to avoid dimension mismatch
        npts1 = Int(time_unit * S1[1].fs)
        if (length(S1[1].x) > npts1)
            S1[1].x=S1[1].x[1:npts1];
            idxtemp=findall(x->x<=npts1,S1[1].t[:,1])
            S1[1].t=S1[1].t[idxtemp,:];
        end
        # do not attempt fft if data was not available
        try
            if S1[1].misc["dlerror"] == 1
                push!(tserrorList, "$tstamp/$stn1")
                filter!(a->a≠stn1, stlist)
                println("$tstamp: $stn1 encountered an error: dlerror==1. Skipping.")
                continue
            elseif  sum(iszero.(S1[1].x)) > 0.5*length(S1[1].x) # skip if > 50% zeros
                push!(tserrorList, "$tstamp/$stn1")
                filter!(a->a≠stn1, stlist)
                println("$tstamp: $stn1 has >50% zeros. Skipping.")
                continue
            elseif  size(S1[1].t,1)>2
                push!(tserrorList, "$tstamp/$stn2")
                filter!(a->a≠stn2, stlist)
                println("$tstamp: $stn2 has $(size(S1[1].t,1)-2) gaps. Skipping.")
                continue
            end
        catch;
            # assume key "dlerror" does not exist (not downloaded via SeisDownload)
        end

        FFT1 = try
            # try to read FFT from cached FFTs
            if stn1 in keys(FFTDict)
                FFTDict[stn1]
            else
                if verbose; println("fft1 -ing");end
                #the following are new steps to do FFT
                process_raw!(S1,fs)
                R1=RawData(S1,cc_len,cc_step)
                R1.freqmin=freqmin;
                R1.freqmax=freqmax;
                R1.whitened=to_whiten;
                R1.time_norm=time_norm;

                detrend!(R1)
                taper!(R1)

                if time_norm == "one-bit" || time_norm == "phase"
                    if time_norm == "one-bit"
                        onebit!(R1);
                    end
                end
                FFT1 = compute_fft(R1)
                if to_whiten
                    whiten!(FFT1,freqmin,freqmax)
                end
                FFTDict[stn1] = FFT1
                FFT1
            end
        catch y
            push!(tserrorList, "$tstamp/$stn1")
            filter!(a->a≠stn1, stlist)
            println(y)
            exit()
            println("$tstamp: $stn1 encountered an error on FFT. Skipping.")
            continue
        end

        # iterate over station list again
        for stn2 in stlist[stniter:end]
            if stn2 in tserrorList continue end # don't attempt FFT if this failed already

            # see if this is an auto-, cross-, or xchan-correlation
            ct = get_corrtype([stn1, stn2])

            # autocorrelation
            if (ct=="acorr") && ("acorr" in corrtype)
                # set the stn2 FFT to the already computed FFT for stn1
                FFT2 = FFT1

            # cross- or cross-channel correlation
            elseif ct in corrtype
                # read station SeisChannels into SeisData before FFT
                S2 = try
                    SeisData(data["$tstamp/$stn2"])
                catch;
                    push!(tserrorList, "$tstamp/$stn2")
                    filter!(a->a≠stn2, stlist)
                    println("$tstamp: $stn2 encountered an error on read. Skipping.")
                    continue
                end

                if S1.n > 1 @warn "SeisData contains multiple channels. Operating only on the first." end
                delete!(S2[1].misc, "kurtosis")
                delete!(S2[1].misc, "eqtimewindow")
                # make sure the data is the proper length to avoid dimension mismatch
                npts2 = Int(time_unit * S2[1].fs)
                if (length(S2[1].x) > npts2)
                    S2[1].x=S2[1].x[1:npts2];
                    idxtemp=findall(x->x<=npts2,S2[1].t[:,1])
                    S2[1].t=S2[1].t[idxtemp,:];
                end
                # do not attempt fft if data was not available, >50% zeros, or has gaps.
                try
                    if S2[1].misc["dlerror"] == 1
                        push!(tserrorList, "$tstamp/$stn2")
                        filter!(a->a≠stn2, stlist)
                        println("$tstamp: $stn2 encountered an error: dlerror==1. Skipping.")
                        continue
                    elseif  sum(iszero.(S2[1].x)) > 0.5*length(S2[1].x) # skip if > 50% zeros
                        push!(tserrorList, "$tstamp/$stn2")
                        filter!(a->a≠stn2, stlist)
                        println("$tstamp: $stn2 has >50% zeros. Skipping.")
                        continue
                    elseif  size(S2[1].t,1)>2
                        push!(tserrorList, "$tstamp/$stn2")
                        filter!(a->a≠stn2, stlist)
                        println("$tstamp: $stn2 has $(size(S2[1].t,1)-2) gaps. Skipping.")
                        continue
                    end
                catch;
                    # assume that the key dlerror does not exist (data not downloaded via SeisDownload)
                end
                #
                # # make sure the data is the proper length to avoid dimension mismatch
                # npts2 = Int(time_unit * S2[1].fs)
                # if (length(S2[1].x) > npts2) S2[1].x=S2[1].x[1:npts2]; S2[1].t[end,1]=npts2 end

                # try to read FFT from cached FFTs, compute if not found
                FFT2 = try
                    if stn2 in keys(FFTDict)
                        FFTDict[stn2]
                    else
                        if verbose; println("fft2 -ing");end
                        #the following are new steps to do FFT
                        process_raw!(S2,fs)
                        R2=RawData(S2,cc_len,cc_step)
                        R2.freqmin=freqmin;
                        R2.freqmax=freqmax;
                        R2.whitened=to_whiten;
                        R2.time_norm=time_norm;

                        detrend!(R2)
                        taper!(R2)

                        if time_norm == "one-bit" || time_norm == "phase"
                            if time_norm == "one-bit"
                                onebit!(R2);
                            end
                        end
                        FFT2 = compute_fft(R2)
                        if to_whiten
                            whiten!(FFT2,freqmin,freqmax)
                        end
                        # FFT2 = compute_fft(S2, freqmin, freqmax, fs, Int(cc_step), Int(cc_len), to_whiten=to_whiten, time_norm=time_norm)

                        # store FFT2 in dict after smoothing
                        FFTDict[stn2] = FFT2
                        FFT2
                    end
                catch y
                    push!(tserrorList, "$tstamp/$stn2")
                    filter!(a->a≠stn2, stlist)
                    println("$tstamp: $stn2 encountered an error on FFT. Skipping.")
                    continue
                end
            else
                println("Skipping cross-correlation of $stn1 and $stn2.")
                continue
            end

            # TODO: find a way to store FFT2 for deconvolution without storing two sets of FFTs (smooth and unsmoothed)
            # deconvolution divides the cross spectrum by the squared power spectrum of the second signal
            # but each signal can be both source and receiver
            # smooth FFT2 if it hasn't been smoothed already
            if corrmethod == "coherence"
                if verbose; println("coherence -ing");end
                xcorr = compute_cc(coherence(FFT1, half_win,water_level), coherence(FFT2, half_win,water_level), maxtimelag, corr_type=corrmethod)
                # coherence!(FFT1, half_win,water_level)
                # coherence!(FFT2, half_win,water_level)
            elseif corrmethod == "deconv"
                if verbose; println("deconv -ing");end
                # deconvolution!(FFT2, half_win,water_level)
                xcorr = compute_cc(deconvolution(FFT1, half_win,water_level),FFT2,  maxtimelag, corr_type=corrmethod)
            else
                # compute correlation using SeisNoise.jl -- returns type CorrData
                xcorr = compute_cc(FFT1, FFT2, maxtimelag, corr_type=corrmethod)
            end

            varname = "$tstamp/$stn1.$stn2"
            try
                # compute distance between stations
                xcorr.misc["dist"] = dist(FFT1.loc, FFT2.loc)
                # save location of each station
                xcorr.misc["location"] = Dict(stn1=>FFT1.loc, stn2=>FFT2.loc)

                # stack over DL_time_unit
                if stack==true
                    if stacktype=="mean"
                        stack!(xcorr, allstack=true) #default is mean stacking/linear stacking.
                    else
                        stack!(xcorr, allstack=true,stacktype=stacktype)
                    end
                end

                outFile[varname] = xcorr
            catch
                println("$stn1 and $stn2 have no overlap at $tstamp.")
                push!(tserrorList, "$tstamp/$stn1")
            end
        end

        # release memory held by the FFT and time series for this station
        delete!(FFTDict, stn1)
        delete!(data, "$tstamp/$stn1")
    end
    outFile["info/errors"] = tserrorList
    close(outFile)
end

"""

    seisxcorrelation_highorder(tstamp::String, corrstationlist::Array{String,2}, data::JLD2.JLDFile, InputDict::Dict)

Compute C2 or C3 cross-correlation and save data in jld2 file with CorrData format.

# Arguments
- `tstamp::String,`    : Time stamp to read from JLD2 file.
- `corrstationlist::Array{String,1},`    : List of cross-correlation names, e.g. BP.CCRB..BP1.BP.EADB..BP1
- `data::JLD2.JLDFile`    : JLD2 file object to read data from.
- `InputDict::Dict`    : Dictionary containing IO, FFT, xcorr, and stacking parameters

# Output
- `basefoname.tstamp.jld2`    : contains SeisData structure with a hierarchical structure (CC function, metadata)

"""
function seisxcorrelation_highorder(data::Dict, tstamp::String, corrstationlist::Array{String,2}, allowable_pairs::Array{String,1}, allowable_vs::Array{String,1}, InputDict::Dict)
    # IO parameters
    basefoname = InputDict["basefoname"]
    # Partitioning parameters
    start_lag  = InputDict["start_lag"]
    win_len    = InputDict["window_len"]
    # FFT parameters
    freqmin    = InputDict["freqmin"]
    freqmax    = InputDict["freqmax"]
    fs         = InputDict["fs"]
    cc_len     = InputDict["cc_len"]
    cc_step    = InputDict["cc_step"]
    # correlation parameters
    corrtype   = InputDict["corrtype"]
    maxtimelag = InputDict["maxtimelag"]
    # stacking parameters
    stack      = InputDict["allstack"]

    pairiter = 1 # counter to prevent computing duplicate xcorrs with reversed order

    # dictionary to cache FFTs
    FFTDict = Dict{String, Array{FFTData,1}}()

    # list of stations that failed for this timestep
    tserrorList = []

    xcorrlist = deepcopy(corrstationlist)

    # convert samples given in [s] to the number of samples
    if typeof(start_lag) == Float64
        start_lag = convert(Int64, start_lag*fs)
    end
    if typeof(win_len) == Float64
        win_len  = convert(Int64, win_len*fs)
    end

    outFile = jldopen("$basefoname.$tstamp.jld2", "a+")
    outFile["info/corrstationlist"] = xcorrlist

    println("$tstamp: Computing cross-correlations")
    for p1=1:length(xcorrlist[1,:])
        # first station pair name and its reverse
        pair1 = xcorrlist[:, p1]
        p1name = pair1[1]*"."*pair1[2]
        p1namerev = pair1[2]*"."*pair1[1]

        for p2=pairiter:length(xcorrlist[1,:])
            # second station pair name
            pair2 = xcorrlist[:, p2]
            p2name = pair2[1]*"."*pair2[2]

            # read data only if we have 3 unique stations
            virt_src = intersect(pair1, pair2)

            if length(virt_src) != 1 continue end

            # set stations used as receivers for c3
            stn1 = setdiff(pair1, virt_src)[1]
            stn2 = setdiff(pair2, virt_src)[1]

            # keep string instead of array for ease of use
            virt_src = virt_src[1]

            # read data only if we have 3 unique station pairs
            if get_corrtype([stn1, stn2]) != "xcorr" continue end
            # read data only if the virtual source is approved
            if virt_src ∉ allowable_vs continue end

            # read data only if the station pair is approved
            if ("$stn1.$stn2" ∉ allowable_pairs) && ("$stn2.$stn1" ∉ allowable_pairs) continue end

            # load FFT1 from dict if it exists
            if "$(virt_src)/$stn1" in keys(FFTDict)
                FFT1_neg, FFT1_pos = FFTDict["$(virt_src)/$stn1"]
            else
                # load data and (optionally) reverse it to get virt_src-stn config
                # copy struct to prevent in-place reversal from messing up future xcorrs
                xcorr1 = copy(data["$tstamp/$p1name"])

                if size(xcorr1.corr, 2)>1 stack!(xcorr1) end

                # reverse xcorr if we loaded stn1-virt_src to get virt_src-stn1
                if virt_src==pair1[2]
                    xcorr1.corr=reverse(xcorr1.corr, dims=1)
                    vs_loc = xcorr1.misc["location"][pair1[2]]
                    stn1_loc = xcorr1.misc["location"][pair1[1]]
                else
                    vs_loc = xcorr1.misc["location"][pair1[1]]
                    stn1_loc = xcorr1.misc["location"][pair1[2]]
                end

                # partition xcorr into positive and negative windows
                xcorr1.corr = partition(xcorr1.corr, start_lag, win_len)

                # compute FFT
                FFT1_neg, FFT1_pos = compute_fft_c3(xcorr1, freqmin, freqmax, fs, cc_step, cc_len)

                FFTDict["$virt_src/$stn1"] = [FFT1_neg, FFT1_pos]
            end

            # load FFT2 from dict if it exists
            if "$virt_src/$stn2" in keys(FFTDict)
                FFT2_neg, FFT2_pos = FFTDict["$virt_src/$stn2"]
            else
                # load data and (optionally) reverse it to get virt_src-stn config
                # copy struct to prevent in-place reversal from messing up future xcorrs
                xcorr2 = copy(data["$tstamp/$p2name"])

                if size(xcorr2.corr, 2)>1 stack!(xcorr2) end

                # reverse xcorr if we loaded stn1-virt_src to get virt_src-stn1
                if virt_src==pair2[2]
                    xcorr2.corr=reverse(xcorr2.corr, dims=1)
                    stn2_loc = xcorr2.misc["location"][pair2[1]]
                else
                    stn2_loc = xcorr2.misc["location"][pair2[2]]
                end

                # partition xcorr into positive and negative windows
                xcorr2.corr = partition(xcorr2.corr, start_lag, win_len)

                # compute FFT
                FFT2_neg, FFT2_pos = compute_fft_c3(xcorr2, freqmin, freqmax, fs, cc_step, cc_len)

                FFTDict["$virt_src/$stn2"] = [FFT2_neg, FFT2_pos]
            end

            #compute xcorr - 1st half of dim=2 is neg lag, 2nd half is pos lag
            xcorr_c3_neg = compute_cc_c3(FFT1_neg, FFT2_neg, maxtimelag)
            xcorr_c3_neg.name = "$stn1.$stn2.$(virt_src)_neg"
            xcorr_c3_pos = compute_cc_c3(FFT1_pos, FFT2_pos, maxtimelag)
            xcorr_c3_pos.name = "$stn1.$stn2.$(virt_src)_pos"

            # stacking
            if stack
                stack!(xcorr_c3_neg, allstack=true)
                stack!(xcorr_c3_pos, allstack=true)
            end

            # save cross-correlation
            varname = "$tstamp/$stn1.$stn2/$virt_src"
            try
                outFile[varname] = [xcorr_c3_neg, xcorr_c3_pos]
            catch
                println("$varname encountered an error on saving to JLD2.")
                push!(tserrorList, varname)
                continue
            end
        end

        # release memory held by the FFT for this station pair in FFTDict
        delete!(FFTDict, p1name)
        delete!(FFTDict, p1namerev)
        # release memory held by the input data
        delete!(data, "$tstamp/$p1name")
        delete!(data, "$tstamp/$p1namerev")

        pairiter += 1
    end
    outFile["erorrs"] = tserrorList
    close(outFile)
end

#end # module


__precompile__(true)

# main module for `LazySets.jl`
module LazySets

using GLPKMathProgInterface, LinearAlgebra, MathProgBase, Reexport, Requires,
      SparseArrays
using LinearAlgebra: checksquare
import LinearAlgebra: norm, ×, normalize, normalize!
import Random
using Random: AbstractRNG, GLOBAL_RNG, SamplerType, shuffle
import InteractiveUtils: subtypes

export Arrays
export ×, normalize

# ==========
# Assertions
# ==========

include("Assertions/Assertions.jl")
@reexport using .Assertions
using .Assertions: @assert
import .Assertions: activate_assertions, deactivate_assertions
# activate assertions by default
activate_assertions(LazySets)

# =====================
# Numeric approximation
# =====================
include("Utils/comparisons.jl")

# =======================
# Arrays auxiliary module
# =======================
include("Arrays/Arrays.jl")
using .Arrays
import .Arrays: distance

# ===================
# Auxiliary functions
# ===================
include("Utils/helper_functions.jl")
include("Utils/macros.jl")
include("Utils/iterators.jl")

# ==================
# Abstract set types
# ==================
include("Interfaces/LazySet.jl")
include("Interfaces/AbstractPolyhedron.jl")
include("Sets/HalfSpace.jl")  # must come before AbstractPolyhedron_functions
include("Interfaces/AbstractPolyhedron_functions.jl")
include("Interfaces/AbstractPolytope.jl")
include("Interfaces/AbstractCentrallySymmetric.jl")
include("Interfaces/AbstractCentrallySymmetricPolytope.jl")
include("Interfaces/AbstractZonotope.jl")
include("Interfaces/AbstractHyperrectangle.jl")
include("Interfaces/AbstractPolygon.jl")
include("Interfaces/AbstractSingleton.jl")
include("Interfaces/AbstractHPolygon.jl")
include("Interfaces/AbstractAffineMap.jl")

# =============================
# Types representing basic sets
# =============================
include("Sets/Ball1.jl")
include("Sets/Ball2.jl")
include("Sets/BallInf.jl")
include("Sets/Ballp.jl")
include("Sets/Ellipsoid.jl")
include("Sets/EmptySet.jl")
include("Sets/HPolygon.jl")
include("Sets/HPolygonOpt.jl")
include("Sets/HPolytope.jl")
include("Sets/HPolyhedron.jl")
include("Sets/Hyperplane.jl")
include("Sets/Hyperrectangle.jl")
include("Sets/Interval.jl")
include("Sets/Line2D.jl")
include("Sets/Line.jl")
include("Sets/LineSegment.jl")
include("Sets/Singleton.jl")
include("Sets/Universe.jl")
include("Sets/VPolygon.jl")
include("Sets/VPolytope.jl")
include("Sets/ZeroSet.jl")
include("Sets/Zonotope.jl")
include("Sets/HParallelotope.jl")

# ==================================
# Types representing non-convex sets
# ==================================
include("Sets/PolynomialZonotope.jl")

# =================================
# Types representing set operations
# =================================
include("LazyOperations/Bloating.jl")
include("LazyOperations/CartesianProduct.jl")
include("LazyOperations/CartesianProductArray.jl")
include("LazyOperations/Complement.jl")
include("LazyOperations/ConvexHull.jl")
include("LazyOperations/ConvexHullArray.jl")
include("LazyOperations/ExponentialMap.jl")
include("LazyOperations/Intersection.jl")
include("LazyOperations/IntersectionArray.jl")
include("LazyOperations/LinearMap.jl")
include("LazyOperations/AffineMap.jl")  # must come after LinearMap
include("LazyOperations/MinkowskiSum.jl")
include("LazyOperations/MinkowskiSumArray.jl")
include("LazyOperations/CachedMinkowskiSumArray.jl")
include("LazyOperations/ResetMap.jl")
include("LazyOperations/SymmetricIntervalHull.jl")
include("LazyOperations/Translation.jl")
include("LazyOperations/UnionSet.jl")
include("LazyOperations/UnionSetArray.jl")
include("LazyOperations/Rectification.jl")  # must come after UnionSet

# =======
# Aliases
# =======
include("Interfaces/aliases.jl")

# =============================
# Conversions between set types
# =============================
include("convert.jl")

# ===========================
# Concrete operations on sets
# ===========================
include("ConcreteOperations/convex_hull.jl")
include("ConcreteOperations/difference.jl")
include("ConcreteOperations/intersection.jl")
include("ConcreteOperations/isdisjoint.jl")
include("ConcreteOperations/issubset.jl")
include("ConcreteOperations/minkowski_difference.jl")
include("ConcreteOperations/minkowski_sum.jl")
include("Utils/samples.jl")

# =====================
# Approximations module
# =====================
include("Approximations/Approximations.jl")
# We export all symbols from Approximations.
# Note that the LazySets module is not supposed to depend on Approximations.
# It can, however, happen that we forget to add the `using` statements.
@reexport using .Approximations

# ==================================
# Plotting (requires Approximations)
# ==================================
include("Plotting/plot_recipes.jl")
include("Plotting/mesh.jl")

# ==========================
# Parallel-algorithms module
# ==========================
include("Parallel/Parallel.jl")

# ===================================================
# Load external packages on-demand (using 'Requires')
# ===================================================
include("init.jl")

# ================================================
# Convenience functions to (de)activate assertions
# ================================================
function activate_assertions()
    for m in [LazySets, Arrays, Approximations, Parallel]
        Assertions.activate_assertions(m)
    end
end
function deactivate_assertions()
    for m in [LazySets, Arrays, Approximations, Parallel]
        Assertions.deactivate_assertions(m)
    end
end

end # module