## Expressions ##

### Expression templates ###

"Defines branch indicator as a function of corresponding ne_branch indicator variables."
function expression_branch_indicator(pm::_PM.AbstractPowerModel, br_idx::Int; nw::Int=_PM.nw_id_default)
    if !haskey(_PM.var(pm, nw), :z_branch)
        _PM.var(pm, nw)[:z_branch] = Dict{Int,Any}()
    end

    branch = _PM.ref(pm, nw, :branch, br_idx)
    f_bus = branch["f_bus"]
    t_bus = branch["t_bus"]

    expression_branch_indicator(pm, nw, br_idx, f_bus, t_bus)
end

### Actual expressions ###

function expression_branch_indicator(pm::_PM.AbstractPowerModel, n::Int, br_idx, f_bus, t_bus)
    branch_ne_sum = sum(_PM.var(pm, n, :branch_ne, l) for l in _PM.ref(pm, n, :ne_buspairs, (f_bus,t_bus), "branches"))

    _PM.var(pm, n, :z_branch)[br_idx] = 1 - branch_ne_sum
end

## Constraints ##

### Constraint templates ###

""
function constraint_ohms_yt_from_repl(pm::_PM.AbstractPowerModel, i::Int; nw::Int=_PM.nw_id_default)
    branch = _PM.ref(pm, nw, :branch, i)
    f_bus = branch["f_bus"]
    t_bus = branch["t_bus"]
    f_idx = (i, f_bus, t_bus)
    t_idx = (i, t_bus, f_bus)

    g, b = _PM.calc_branch_y(branch)
    tr, ti = _PM.calc_branch_t(branch)
    g_fr = branch["g_fr"]
    b_fr = branch["b_fr"]
    tm = branch["tap"]

    vad_min = _PM.ref(pm, nw, :off_angmin)
    vad_max = _PM.ref(pm, nw, :off_angmax)

    # track if a certain candidate branch is replacing a line
    replace, ne_br_idx = replace_branch(pm, nw, f_bus, t_bus)
    # If lines is to be repalced use formulations below, else use PowerModels constraint for existing branches
    if replace  == 0
        _PM.constraint_ohms_yt_from(pm, nw, f_bus, t_bus, f_idx, t_idx, g, b, g_fr, b_fr, tr, ti, tm)
    else
        constraint_ohms_yt_from_repl(pm, nw, ne_br_idx, f_bus, t_bus, f_idx, b, vad_min, vad_max)
    end
end

""
function constraint_ohms_yt_to_repl(pm::_PM.AbstractPowerModel, i::Int; nw::Int=_PM.nw_id_default)
    branch = _PM.ref(pm, nw, :branch, i)
    f_bus = branch["f_bus"]
    t_bus = branch["t_bus"]
    f_idx = (i, f_bus, t_bus)
    t_idx = (i, t_bus, f_bus)

    g, b = _PM.calc_branch_y(branch)
    tr, ti = _PM.calc_branch_t(branch)
    g_to = branch["g_to"]
    b_to = branch["b_to"]
    tm = branch["tap"]

    vad_min = _PM.ref(pm, nw, :off_angmin)
    vad_max = _PM.ref(pm, nw, :off_angmax)

    # track if a certain candidate branch is replacing a line
    replace, ne_br_idx = replace_branch(pm, nw, f_bus, t_bus)
    # If lines is to be repalced use formulations below, else use PowerModels constraint for existing branches
    if replace  == 0
        _PM.constraint_ohms_yt_to(pm, nw, f_bus, t_bus, f_idx, t_idx, g, b, g_to, b_to, tr, ti, tm)
    else
        constraint_ohms_yt_to_repl(pm, nw, ne_br_idx, f_bus, t_bus, t_idx, b, vad_min, vad_max)
    end
end

function constraint_voltage_angle_difference_repl(pm::_PM.AbstractPowerModel, i::Int; nw::Int=_PM.nw_id_default)
    branch = _PM.ref(pm, nw, :branch, i)
    f_bus = branch["f_bus"]
    t_bus = branch["t_bus"]
    f_idx = (i, branch["f_bus"], branch["t_bus"])

    vad_min = _PM.ref(pm, nw, :off_angmin)
    vad_max = _PM.ref(pm, nw, :off_angmax)

    # track if a certain candidate branch is replacing a line
    replace, ne_br_idx = replace_branch(pm, nw, f_bus, t_bus)
    # If lines is to be repalced use formulations below, else use PowerModels constraint for existing branches
    if replace  == 0
        _PM.constraint_voltage_angle_difference(pm, nw, f_idx, branch["angmin"], branch["angmax"])
    else
        constraint_voltage_angle_difference_repl(pm, nw, ne_br_idx, f_idx, branch["angmin"], branch["angmax"], vad_min, vad_max)
    end
end

function constraint_thermal_limit_from_repl(pm::_PM.AbstractPowerModel, i::Int; nw::Int=_PM.nw_id_default)
    branch = _PM.ref(pm, nw, :branch, i)
    f_bus = branch["f_bus"]
    t_bus = branch["t_bus"]
    f_idx = (i, f_bus, t_bus)

    if !haskey(branch, "rate_a")
        Memento.error(_LOGGER, "constraint_thermal_limit_from_ne requires a rate_a value on all branches, calc_thermal_limits! can be used to generate reasonable values")
    end

    # track if a certain candidate branch is replacing a line
    replace, ne_br_idx = replace_branch(pm, nw, f_bus, t_bus)
    # If lines is to be repalced use formulations below, else use PowerModels constraint for existing branches
    if replace  == 0
        _PM.constraint_thermal_limit_from(pm, nw, f_idx, branch["rate_a"])
    else
        constraint_thermal_limit_from_repl(pm, nw, ne_br_idx, f_idx, branch["rate_a"])
    end
end

""
function constraint_thermal_limit_to_repl(pm::_PM.AbstractPowerModel, i::Int; nw::Int=_PM.nw_id_default)
    branch = _PM.ref(pm, nw, :branch, i)
    f_bus = branch["f_bus"]
    t_bus = branch["t_bus"]
    t_idx = (i, t_bus, f_bus)

    if !haskey(branch, "rate_a")
        Memento.error(_LOGGER, "constraint_thermal_limit_to_ne requires a rate_a value on all branches, calc_thermal_limits! can be used to generate reasonable values")
    end
    # track if a certain candidate branch is replacing a line
    replace, ne_br_idx = replace_branch(pm, nw, f_bus, t_bus)
    # If lines is to be repalced use formulations below, else use PowerModels constraint for existing branches
    if replace  == 0
        _PM.constraint_thermal_limit_to(pm, nw, t_idx, branch["rate_a"])
    else
        constraint_thermal_limit_to_repl(pm, nw, ne_br_idx, t_idx, branch["rate_a"])
    end
end

#### Constraint templates used in radial networks ####

"States that at most one of the ne_branches sharing the same bus pair must be built."
function constraint_branch_complementarity(pm::_PM.AbstractPowerModel, br_idx::Int; nw::Int=_PM.nw_id_default)
    branch = _PM.ref(pm, nw, :branch, br_idx)
    f_bus = branch["f_bus"]
    t_bus = branch["t_bus"]

    constraint_branch_complementarity(pm, nw, br_idx, f_bus, t_bus)
end

""
function constraint_power_losses_on_off(pm::_PM.AbstractPowerModel, br_idx::Int; nw::Int=_PM.nw_id_default)
    branch = _PM.ref(pm, nw, :branch, br_idx)
    f_bus = branch["f_bus"]
    t_bus = branch["t_bus"]
    f_idx = (br_idx, f_bus, t_bus)
    t_idx = (br_idx, t_bus, f_bus)

    r = branch["br_r"]
    x = branch["br_x"]
    tm = branch["tap"]
    g_sh_fr = branch["g_fr"]
    g_sh_to = branch["g_to"]
    b_sh_fr = branch["b_fr"]
    b_sh_to = branch["b_to"]

    vad_min = _PM.ref(pm, nw, :off_angmin)
    vad_max = _PM.ref(pm, nw, :off_angmax)

    constraint_power_losses_on_off(pm, nw, br_idx, f_bus, t_bus, f_idx, t_idx, r, x, g_sh_fr, g_sh_to, b_sh_fr, b_sh_to, tm, vad_min, vad_max)
end

""
function constraint_power_losses_frb_on_off(pm::_PM.AbstractPowerModel, br_idx::Int; nw::Int=_PM.nw_id_default)
    branch = _PM.ref(pm, nw, :branch, br_idx)
    f_bus = branch["f_bus"]
    t_bus = branch["t_bus"]
    f_idx = (br_idx, f_bus, t_bus)
    t_idx = (br_idx, t_bus, f_bus)

    r = branch["br_r"]
    x = branch["br_x"]
    tm = branch["tap"]
    g_sh_fr = branch["g_fr"]
    g_sh_to = branch["g_to"]
    b_sh_fr = branch["b_fr"]
    b_sh_to = branch["b_to"]

    vad_min = _PM.ref(pm, nw, :off_angmin)
    vad_max = _PM.ref(pm, nw, :off_angmax)

    constraint_power_losses_frb_on_off(pm, nw, br_idx, f_bus, t_bus, f_idx, t_idx, r, x, g_sh_fr, g_sh_to, b_sh_fr, b_sh_to, tm, vad_min, vad_max)
end

""
function constraint_power_losses_oltc_on_off(pm::_PM.AbstractBFModel, br_idx::Int; nw::Int=_PM.nw_id_default)
    branch = _PM.ref(pm, nw, :branch, br_idx)
    f_bus = branch["f_bus"]
    t_bus = branch["t_bus"]
    f_idx = (br_idx, f_bus, t_bus)
    t_idx = (br_idx, t_bus, f_bus)

    r = branch["br_r"]
    x = branch["br_x"]
    g_sh_fr = branch["g_fr"]
    g_sh_to = branch["g_to"]
    b_sh_fr = branch["b_fr"]
    b_sh_to = branch["b_to"]

    vad_min = _PM.ref(pm, nw, :off_angmin)
    vad_max = _PM.ref(pm, nw, :off_angmax)

    constraint_power_losses_oltc_on_off(pm, nw, br_idx, f_bus, t_bus, f_idx, t_idx, r, x, g_sh_fr, g_sh_to, b_sh_fr, b_sh_to, vad_min, vad_max)
end

""
function constraint_ne_power_losses_parallel(pm::_PM.AbstractPowerModel, ne_br_idx::Int; nw::Int=_PM.nw_id_default)
    ne_branch = _PM.ref(pm, nw, :ne_branch, ne_br_idx)
    f_bus = ne_branch["f_bus"]
    t_bus = ne_branch["t_bus"]
    f_idx = (ne_br_idx, f_bus, t_bus)
    t_idx = (ne_br_idx, t_bus, f_bus)
    vad_min = _PM.ref(pm, nw, :off_angmin)
    vad_max = _PM.ref(pm, nw, :off_angmax)

    ne_r       = ne_branch["br_r"]
    ne_x       = ne_branch["br_x"]
    ne_tm      = ne_branch["tap"]
    ne_g_sh_fr = ne_branch["g_fr"]
    ne_g_sh_to = ne_branch["g_to"]
    ne_b_sh_fr = ne_branch["b_fr"]
    ne_b_sh_to = ne_branch["b_to"]

    br_idx  = branch_idx(pm, nw, f_bus, t_bus)
    branch  = _PM.ref(pm, nw, :branch, br_idx)
    r       = branch["br_r"]
    x       = branch["br_x"]
    tm      = branch["tap"]
    g_sh_fr = branch["g_fr"]
    g_sh_to = branch["g_to"]
    b_sh_fr = branch["b_fr"]
    b_sh_to = branch["b_to"]

    if ne_tm != tm
        Memento.error(_LOGGER, "ne_branch $(ne_br_idx) cannot be built in parallel to branch $(br_idx) because has a different tap ratio")
    end

    constraint_ne_power_losses_parallel(pm, nw, br_idx, ne_br_idx, f_bus, t_bus, f_idx, t_idx, r, x, g_sh_fr, g_sh_to, b_sh_fr, b_sh_to, ne_r, ne_x, ne_g_sh_fr, ne_g_sh_to, ne_b_sh_fr, ne_b_sh_to, tm, vad_min, vad_max)
end

""
function constraint_ne_power_losses_frb_parallel(pm::_PM.AbstractPowerModel, ne_br_idx::Int; nw::Int=_PM.nw_id_default)
    ne_branch = _PM.ref(pm, nw, :ne_branch, ne_br_idx)
    f_bus = ne_branch["f_bus"]
    t_bus = ne_branch["t_bus"]
    f_idx = (ne_br_idx, f_bus, t_bus)
    t_idx = (ne_br_idx, t_bus, f_bus)
    vad_min = _PM.ref(pm, nw, :off_angmin)
    vad_max = _PM.ref(pm, nw, :off_angmax)

    ne_r       = ne_branch["br_r"]
    ne_x       = ne_branch["br_x"]
    ne_tm      = ne_branch["tap"]
    ne_g_sh_fr = ne_branch["g_fr"]
    ne_g_sh_to = ne_branch["g_to"]
    ne_b_sh_fr = ne_branch["b_fr"]
    ne_b_sh_to = ne_branch["b_to"]

    br_idx  = branch_idx(pm, nw, f_bus, t_bus)
    branch  = _PM.ref(pm, nw, :branch, br_idx)
    r       = branch["br_r"]
    x       = branch["br_x"]
    tm      = branch["tap"]
    g_sh_fr = branch["g_fr"]
    g_sh_to = branch["g_to"]
    b_sh_fr = branch["b_fr"]
    b_sh_to = branch["b_to"]

    if ne_tm != tm
        Memento.error(_LOGGER, "ne_branch $(ne_br_idx) cannot be built in parallel to branch $(br_idx) because has a different tap ratio")
    end

    constraint_ne_power_losses_frb_parallel(pm, nw, br_idx, ne_br_idx, f_bus, t_bus, f_idx, t_idx, r, x, g_sh_fr, g_sh_to, b_sh_fr, b_sh_to, ne_r, ne_x, ne_g_sh_fr, ne_g_sh_to, ne_b_sh_fr, ne_b_sh_to, tm, vad_min, vad_max)
end

""
function constraint_voltage_magnitude_difference_on_off(pm::_PM.AbstractPowerModel, br_idx::Int; nw::Int=_PM.nw_id_default)
    branch = _PM.ref(pm, nw, :branch, br_idx)
    f_bus = branch["f_bus"]
    t_bus = branch["t_bus"]
    f_idx = (br_idx, f_bus, t_bus)
    t_idx = (br_idx, t_bus, f_bus)

    r = branch["br_r"]
    x = branch["br_x"]
    g_sh_fr = branch["g_fr"]
    b_sh_fr = branch["b_fr"]
    tm = branch["tap"]

    constraint_voltage_magnitude_difference_on_off(pm, nw, br_idx, f_bus, t_bus, f_idx, t_idx, r, x, g_sh_fr, b_sh_fr, tm)
end

""
function constraint_voltage_magnitude_difference_frb_on_off(pm::_PM.AbstractPowerModel, br_idx::Int; nw::Int=_PM.nw_id_default)
    branch = _PM.ref(pm, nw, :branch, br_idx)
    f_bus = branch["f_bus"]
    t_bus = branch["t_bus"]
    f_idx = (br_idx, f_bus, t_bus)
    t_idx = (br_idx, t_bus, f_bus)

    r = branch["br_r"]
    x = branch["br_x"]
    g_sh_fr = branch["g_fr"]
    b_sh_fr = branch["b_fr"]
    tm = branch["tap"]

    constraint_voltage_magnitude_difference_frb_on_off(pm, nw, br_idx, f_bus, t_bus, f_idx, t_idx, r, x, g_sh_fr, b_sh_fr, tm)
end

""
function constraint_voltage_magnitude_difference_oltc_on_off(pm::_PM.AbstractBFModel, br_idx::Int; nw::Int=_PM.nw_id_default)
    branch = _PM.ref(pm, nw, :branch, br_idx)
    f_bus = branch["f_bus"]
    t_bus = branch["t_bus"]
    f_idx = (br_idx, f_bus, t_bus)
    t_idx = (br_idx, t_bus, f_bus)

    r = branch["br_r"]
    x = branch["br_x"]
    g_sh_fr = branch["g_fr"]
    b_sh_fr = branch["b_fr"]
    tm_min = branch["tm_min"]
    tm_max = branch["tm_max"]

    constraint_voltage_magnitude_difference_oltc_on_off(pm, nw, br_idx, f_bus, t_bus, f_idx, t_idx, r, x, g_sh_fr, b_sh_fr, tm_min, tm_max)
end

""
function constraint_ne_voltage_magnitude_difference_parallel(pm::_PM.AbstractPowerModel, ne_br_idx::Int; nw::Int =_PM.nw_id_default)
    ne_branch = _PM.ref(pm, nw, :ne_branch, ne_br_idx)
    f_bus = ne_branch["f_bus"]
    t_bus = ne_branch["t_bus"]
    f_idx = (ne_br_idx, f_bus, t_bus)
    t_idx = (ne_br_idx, t_bus, f_bus)

    ne_r       = ne_branch["br_r"]
    ne_x       = ne_branch["br_x"]
    ne_g_sh_fr = ne_branch["g_fr"]
    ne_b_sh_fr = ne_branch["b_fr"]
    ne_tm      = ne_branch["tap"]

    br_idx  = branch_idx(pm, nw, f_bus, t_bus)
    branch  = _PM.ref(pm, nw, :branch, br_idx)
    r       = branch["br_r"]
    x       = branch["br_x"]
    g_sh_fr = branch["g_fr"]
    b_sh_fr = branch["b_fr"]
    tm      = branch["tap"]
    if is_oltc_branch(pm, br_idx, nw = nw)
        Memento.error(_LOGGER, "ne_branch $ne_br_idx cannot be built in parallel to an OLTC (branch $br_idx)")
    end
    if ne_tm != tm
        Memento.error(_LOGGER, "ne_branch $ne_br_idx cannot be built in parallel to branch $br_idx because has a different tap ratio")
    end

    constraint_ne_voltage_magnitude_difference_parallel(pm, nw, br_idx, ne_br_idx, f_bus, t_bus, f_idx, t_idx, r, x, g_sh_fr, b_sh_fr, ne_r, ne_x, ne_g_sh_fr, ne_b_sh_fr, tm)
end

""
function constraint_ne_voltage_magnitude_difference_frb_parallel(pm::_PM.AbstractPowerModel, ne_br_idx::Int; nw::Int =_PM.nw_id_default)
    ne_branch = _PM.ref(pm, nw, :ne_branch, ne_br_idx)
    f_bus = ne_branch["f_bus"]
    t_bus = ne_branch["t_bus"]
    f_idx = (ne_br_idx, f_bus, t_bus)
    t_idx = (ne_br_idx, t_bus, f_bus)

    ne_r       = ne_branch["br_r"]
    ne_x       = ne_branch["br_x"]
    ne_g_sh_fr = ne_branch["g_fr"]
    ne_b_sh_fr = ne_branch["b_fr"]
    ne_tm      = ne_branch["tap"]

    br_idx  = branch_idx(pm, nw, f_bus, t_bus)
    branch  = _PM.ref(pm, nw, :branch, br_idx)
    r       = branch["br_r"]
    x       = branch["br_x"]
    g_sh_fr = branch["g_fr"]
    b_sh_fr = branch["b_fr"]
    tm      = branch["tap"]

    if is_oltc_branch(pm, br_idx, nw = nw)
        Memento.error(_LOGGER, "ne_branch $ne_br_idx cannot be built in parallel to an OLTC (branch $br_idx)")
    end
    if ne_tm != tm
        Memento.error(_LOGGER, "ne_branch $ne_br_idx cannot be built in parallel to branch $br_idx because has a different tap ratio")
    end

    constraint_ne_voltage_magnitude_difference_frb_parallel(pm, nw, br_idx, ne_br_idx, f_bus, t_bus, f_idx, t_idx, r, x, g_sh_fr, b_sh_fr, ne_r, ne_x, ne_g_sh_fr, ne_b_sh_fr, tm)
end

""
function constraint_ne_thermal_limit_from_parallel(pm::_PM.AbstractPowerModel, ne_br_idx::Int; nw::Int=_PM.nw_id_default)
    ne_branch = _PM.ref(pm, nw, :ne_branch, ne_br_idx)
    f_bus     = ne_branch["f_bus"]
    t_bus     = ne_branch["t_bus"]
    f_idx     = (ne_br_idx, f_bus, t_bus)
    if !haskey(ne_branch, "rate_a")
        Memento.error(_LOGGER, "constraint_ne_thermal_limit_from_parallel requires a rate_a value on all ne_branches, calc_thermal_limits! can be used to generate reasonable values")
    end
    ne_rate_a = ne_branch["rate_a"]

    br_idx = branch_idx(pm, nw, f_bus, t_bus)
    branch = _PM.ref(pm, nw, :branch, br_idx)
    rate_a = branch["rate_a"]

    constraint_ne_thermal_limit_from_parallel(pm, nw, br_idx, ne_br_idx, f_idx, rate_a, ne_rate_a)
end

""
function constraint_ne_thermal_limit_to_parallel(pm::_PM.AbstractPowerModel, ne_br_idx::Int; nw::Int=_PM.nw_id_default)
    ne_branch = _PM.ref(pm, nw, :ne_branch, ne_br_idx)
    f_bus     = ne_branch["f_bus"]
    t_bus     = ne_branch["t_bus"]
    t_idx     = (ne_br_idx, t_bus, f_bus)
    if !haskey(ne_branch, "rate_a")
        Memento.error(_LOGGER, "constraint_ne_thermal_limit_to_parallel requires a rate_a value on all ne_branches, calc_thermal_limits! can be used to generate reasonable values")
    end
    ne_rate_a = ne_branch["rate_a"]

    br_idx = branch_idx(pm, nw, f_bus, t_bus)
    branch = _PM.ref(pm, nw, :branch, br_idx)
    rate_a = branch["rate_a"]

    constraint_ne_thermal_limit_to_parallel(pm, nw, br_idx, ne_br_idx, t_idx, rate_a, ne_rate_a)
end

### Actual constraints ###

function constraint_ohms_yt_from_repl(pm::_PM.AbstractDCPModel, n::Int, ne_br_idx, f_bus, t_bus, f_idx, b, vad_min, vad_max)
    p_fr  = _PM.var(pm, n,   :p, f_idx)
    va_fr = _PM.var(pm, n,   :va, f_bus)
    va_to = _PM.var(pm, n,   :va, t_bus)
    z = _PM.var(pm, n, :branch_ne, ne_br_idx)

    JuMP.@constraint(pm.model, p_fr <= -b*(va_fr - va_to + vad_max*z))
    JuMP.@constraint(pm.model, p_fr >= -b*(va_fr - va_to + vad_min*z))
end

"nothing to do, this model is symetric"
function constraint_ohms_yt_to_repl(pm::_PM.AbstractAPLossLessModels, n::Int, ne_br_idx, f_bus, t_bus, t_idx, b, vad_min, vad_max)
end

function constraint_voltage_angle_difference_repl(pm::_PM.AbstractDCPModel, n::Int, ne_br_idx, f_idx, angmin, angmax, vad_min, vad_max)
    i, f_bus, t_bus = f_idx

    va_fr = _PM.var(pm, n, :va, f_bus)
    va_to = _PM.var(pm, n, :va, t_bus)
    z = _PM.var(pm, n, :branch_ne, ne_br_idx)

    JuMP.@constraint(pm.model, va_fr - va_to <= angmax*(1-z) + vad_max*z)
    JuMP.@constraint(pm.model, va_fr - va_to >= angmin*(1-z) + vad_min*z)
end

""
function constraint_thermal_limit_from_repl(pm::_PM.AbstractActivePowerModel, n::Int, ne_br_idx, f_idx, rate_a)
    p_fr = _PM.var(pm, n, :p, f_idx)
    z = _PM.var(pm, n, :branch_ne, ne_br_idx)

    JuMP.@constraint(pm.model, p_fr <=  rate_a*(1-z))
    JuMP.@constraint(pm.model, p_fr >= -rate_a*(1-z))
end

""
function constraint_thermal_limit_to_repl(pm::_PM.AbstractActivePowerModel, n::Int, ne_br_idx, t_idx, rate_a)
    p_to = _PM.var(pm, n, :p, t_idx)
    z = _PM.var(pm, n, :branch_ne, ne_br_idx)

    JuMP.@constraint(pm.model, p_to <=  rate_a*(1-z))
    JuMP.@constraint(pm.model, p_to >= -rate_a*(1-z))
end

#### Actual constraints used in radial networks ####

"States that at most one of the ne_branches sharing the same bus pair must be built."
function constraint_branch_complementarity(pm::_PM.AbstractPowerModel, n::Int, i, f_bus, t_bus)
    JuMP.@constraint(pm.model, sum(_PM.var(pm, n, :branch_ne, l) for l in ne_branch_ids(pm, n, f_bus, t_bus)) <= 1)
end

""
function constraint_voltage_magnitude_difference_on_off(pm::_PM.AbstractBFAModel, n::Int, i, f_bus, t_bus, f_idx, t_idx, r, x, g_sh_fr, b_sh_fr, tm)
    branch = _PM.ref(pm, n, :branch, i)
    fr_bus = _PM.ref(pm, n, :bus, f_bus)
    to_bus = _PM.ref(pm, n, :bus, t_bus)
    M_hi   =  fr_bus["vmax"]^2/tm^2 - to_bus["vmin"]^2
    M_lo   = -fr_bus["vmin"]^2/tm^2 + to_bus["vmax"]^2

    p_fr = _PM.var(pm, n, :p, f_idx)
    q_fr = _PM.var(pm, n, :q, f_idx)
    w_fr = _PM.var(pm, n, :w, f_bus)
    w_to = _PM.var(pm, n, :w, t_bus)
    z    = _PM.var(pm, n, :z_branch, i)

    JuMP.@constraint(pm.model, (w_fr/tm^2) - w_to <= 2*(r*p_fr + x*q_fr) + M_hi*(1-z) )
    JuMP.@constraint(pm.model, (w_fr/tm^2) - w_to >= 2*(r*p_fr + x*q_fr) - M_lo*(1-z) )
end

""
function constraint_voltage_magnitude_difference_frb_on_off(pm::_PM.AbstractBFAModel, n::Int, i, f_bus, t_bus, f_idx, t_idx, r, x, g_sh_fr, b_sh_fr, tm)
    branch = _PM.ref(pm, n, :branch, i)
    fr_bus = _PM.ref(pm, n, :bus, f_bus)
    to_bus = _PM.ref(pm, n, :bus, t_bus)
    M_hi   =  1.0^2/tm^2 - to_bus["vmin"]^2
    M_lo   = -1.0^2/tm^2 + to_bus["vmax"]^2

    p_fr = _PM.var(pm, n, :p, f_idx)
    q_fr = _PM.var(pm, n, :q, f_idx)
    w_to = _PM.var(pm, n, :w, t_bus)
    z    = _PM.var(pm, n, :z_branch, i)
    # w_fr is assumed equal to 1.0

    JuMP.@constraint(pm.model, (1.0/tm^2) - w_to <= 2*(r*p_fr + x*q_fr) + M_hi*(1-z) )
    JuMP.@constraint(pm.model, (1.0/tm^2) - w_to >= 2*(r*p_fr + x*q_fr) - M_lo*(1-z) )
end

""
function constraint_voltage_magnitude_difference_oltc_on_off(pm::_PM.AbstractBFAModel, n::Int, i, f_bus, t_bus, f_idx, t_idx, r, x, g_sh_fr, b_sh_fr, tm_min, tm_max)
    branch = _PM.ref(pm, n, :branch, i)
    fr_bus = _PM.ref(pm, n, :bus, f_bus)
    to_bus = _PM.ref(pm, n, :bus, t_bus)
    M_hi   =  1.0^2/tm_min^2 - to_bus["vmin"]^2
    M_lo   = -1.0^2/tm_max^2 + to_bus["vmax"]^2

    p_fr = _PM.var(pm, n, :p, f_idx)
    q_fr = _PM.var(pm, n, :q, f_idx)
    ttmi = _PM.var(pm, n, :ttmi, i)
    w_to = _PM.var(pm, n, :w, t_bus)
    z    = _PM.var(pm, n, :z_branch, i)
    # w_fr is assumed equal to 1.0 to preserve the linearity of the model

    JuMP.@constraint(pm.model, 1.0*ttmi - w_to <= 2*(r*p_fr + x*q_fr) + M_hi*(1-z) )
    JuMP.@constraint(pm.model, 1.0*ttmi - w_to >= 2*(r*p_fr + x*q_fr) - M_lo*(1-z) )
end

""
function constraint_ne_voltage_magnitude_difference_parallel(pm::_PM.AbstractBFAModel, n::Int, br_idx_e, br_idx_c, f_bus, t_bus, f_idx_c, t_idx_c, r_e, x_e, g_sh_fr_e, b_sh_fr_e, r_c, x_c, g_sh_fr_c, b_sh_fr_c, tm)
    # Suffixes: _e: existing branch; _c: candidate branch; _p: parallel equivalent
    r_p  = (r_e*(r_c^2+x_c^2)+r_c*(r_e^2+x_e^2)) / ((r_e+r_c)^2+(x_e+x_c)^2)
    x_p  = (x_e*(r_c^2+x_c^2)+x_c*(r_e^2+x_e^2)) / ((r_e+r_c)^2+(x_e+x_c)^2)

    constraint_ne_voltage_magnitude_difference(pm::_PM.AbstractBFAModel, n::Int, br_idx_c, f_bus, t_bus, f_idx_c, t_idx_c, r_p, x_p, 0, 0, tm)
end

""
function constraint_ne_voltage_magnitude_difference_frb_parallel(pm::_PM.AbstractBFAModel, n::Int, br_idx_e, br_idx_c, f_bus, t_bus, f_idx_c, t_idx_c, r_e, x_e, g_sh_fr_e, b_sh_fr_e, r_c, x_c, g_sh_fr_c, b_sh_fr_c, tm)
    # Suffixes: _e: existing branch; _c: candidate branch; _p: parallel equivalent
    r_p  = (r_e*(r_c^2+x_c^2)+r_c*(r_e^2+x_e^2)) / ((r_e+r_c)^2+(x_e+x_c)^2)
    x_p  = (x_e*(r_c^2+x_c^2)+x_c*(r_e^2+x_e^2)) / ((r_e+r_c)^2+(x_e+x_c)^2)

    constraint_ne_voltage_magnitude_difference_frb(pm::_PM.AbstractBFAModel, n::Int, br_idx_c, f_bus, t_bus, f_idx_c, t_idx_c, r_p, x_p, 0, 0, tm)
end

## Auxiliary functions ##

function replace_branch(pm, nw, f_bus, t_bus)
    replace = 0
    ne_br_idx = 0
    for (br, ne_branch) in _PM.ref(pm, nw, :ne_branch)
        if ((ne_branch["f_bus"] == f_bus && ne_branch["t_bus"] == t_bus) || (ne_branch["f_bus"] == t_bus && ne_branch["t_bus"] == f_bus)) && ne_branch["replace"] == true
            replace = 1
            ne_br_idx = br
        end
    end
    return replace, ne_br_idx
end

"Returns the index of `branch` connecting `f_bus` to `t_bus`, if such a `branch` exists; 0 otherwise"
function branch_idx(pm::_PM.AbstractPowerModel, nw::Int, f_bus, t_bus)
    buspairs = _PM.ref(pm, nw, :buspairs)
    buspair = get(buspairs, (f_bus,t_bus), Dict("branch"=>0))
    return buspair["branch"]
end

"Returns a list of indices of `ne_branch`es relative to `branch` `br_idx`"
function ne_branch_ids(pm::_PM.AbstractPowerModel, br_idx::Int; nw::Int=_PM.nw_id_default)
    branch = _PM.ref(pm, nw, :branch, br_idx)
    f_bus = branch["f_bus"]
    t_bus = branch["t_bus"]
    ne_branch_ids(pm, nw, f_bus, t_bus)
end

"Returns a list of indices of `ne_branch`es connecting `f_bus` to `t_bus`"
function ne_branch_ids(pm::_PM.AbstractPowerModel, nw::Int, f_bus, t_bus)
    ne_buspairs = _PM.ref(pm, nw, :ne_buspairs)
    ne_buspair = get(ne_buspairs, (f_bus,t_bus), Dict("branches"=>Int[]))
    return ne_buspair["branches"]
end

"Returns whether a `ne_branch` is intended to replace the existing branch or to be added in parallel"
function ne_branch_replace(pm::_PM.AbstractPowerModel, ne_br_idx::Int; nw::Int=_PM.nw_id_default)
    ne_branch = _PM.ref(pm, nw, :ne_branch, ne_br_idx)
    if !haskey(ne_branch, "replace")
        Memento.error(_LOGGER, "a `replace` value is required on all `ne_branch`es")
    end
    return ne_branch["replace"] == 1
end


notary = Notary()
cypher = CypherSuite(notary)

# Somewhere far far away

mixkey = newsigner(notary)

MIX_PORT = 3000
MIX_ID = id(mixkey)

mix = Mix(MIX_PORT, MIX_ID, cypher)
@async SynchronicBallot.serve(mix, mixkey)

sleep(1.)

# Using the mixer

guardiankey = newsigner(notary)
transactions = Trigger(TransactionVector(), t -> println(typeof(t.document)))
tlog = BraidChain(transactions, Guardian(id(guardiankey)), notary)

registrator = newsigner(notary)
reg = Registrator(id(registrator))
push!(tlog, Transaction(reg, guardiankey))

braider = newsigner(notary)

push!(tlog, Transaction(Braider(id(braider)), guardiankey))

gk = GateKeeper(3001, id(braider), cypher, mix)
sb = BraiderConfig(UInt8(3), UInt8(64), gk)
bo = BraidOfficer(sb, tlog)
botask = @async BraidChains.serve!(bo, braider)
sleep(1.)

### Registration of members ###

member1 = KeyChain(notary)
member2 = KeyChain(notary)
member3 = KeyChain(notary)

# on the server
push!(tlog, Transaction(Member(id(member1)), registrator))
push!(tlog, Transaction(Member(id(member2)), registrator))
push!(tlog, Transaction(Member(id(member3)), registrator))

update!(member1, tlog)
update!(member2, tlog)
update!(member3, tlog)

@show istaskstarted(botask)

### First Braid

@sync begin
    @async braid!(member1, bo)
    @async braid!(member2, bo)
    @async braid!(member3, bo)
end

### After that I shall repeat the procedure

sleep(1.)

update!(member1, tlog)
update!(member2, tlog)
update!(member3, tlog)

@test id(member1[end]) in pseudonyms(tlog)
@test id(member2[end]) in pseudonyms(tlog)
@test id(member3[end]) in pseudonyms(tlog)

### Second Braid

@sync begin
    @async braid!(member1, bo)
    @async braid!(member2, bo)
    @async braid!(member3, bo)
end

sleep(1.)

update!(member1, tlog)
update!(member2, tlog)
update!(member3, tlog)

@test id(member1[end]) in pseudonyms(tlog)
@test id(member2[end]) in pseudonyms(tlog)
@test id(member3[end]) in pseudonyms(tlog)


# This file is WIP while the interface for templates is finalized
@testset "Manual Operations Template" begin
    template = ProblemTemplate(CopperPlatePowerModel)
    set_device_model!(template, PowerLoad, StaticPowerLoad)
    set_device_model!(template, ThermalStandard, ThermalStandardUnitCommitment)
    set_device_model!(template, Line, StaticBranchUnbounded)
    @test !isempty(template.devices)
    @test !isempty(template.branches)
    @test isempty(template.services)
end

@testset "Operations Template Overwrite" begin
    template = ProblemTemplate(CopperPlatePowerModel)
    set_device_model!(template, PowerLoad, StaticPowerLoad)
    set_device_model!(template, ThermalStandard, ThermalStandardUnitCommitment)
    @test_logs (:warn, "Overwriting ThermalStandard existing model") set_device_model!(
        template,
        DeviceModel(ThermalStandard, ThermalBasicUnitCommitment),
    )
    @test PSI.get_formulation(template.devices[:ThermalStandard]) ==
          ThermalBasicUnitCommitment
end

@testset "Provided Templates Tests" begin
    uc_template = template_unit_commitment()
    @test !isempty(uc_template.devices)
    @test PSI.get_formulation(uc_template.devices[:ThermalStandard]) ==
          ThermalBasicUnitCommitment
    uc_template = template_unit_commitment(network=DCPPowerModel)
    @test get_network_formulation(uc_template) == DCPPowerModel
    @test !isempty(uc_template.branches)
    @test !isempty(uc_template.services)

    ed_template = template_economic_dispatch()
    @test !isempty(ed_template.devices)
    @test PSI.get_formulation(ed_template.devices[:ThermalStandard]) == ThermalBasicDispatch
    ed_template = template_economic_dispatch(network=ACPPowerModel)
    @test get_network_formulation(ed_template) == ACPPowerModel
    @test !isempty(ed_template.branches)
    @test !isempty(ed_template.services)
end


module Stheno

    using Reexport

    using AbstractGPs
    using BlockArrays
    using ChainRulesCore
    using FillArrays
    @reexport using KernelFunctions
    using LinearAlgebra
    using Random
    using Zygote
    using ZygoteRules

    import Base.Broadcast: broadcasted

    using AbstractGPs: AbstractGP, FiniteGP, GP
    import AbstractGPs:
        mean,
        cov,
        var,
        mean_and_cov,
        mean_and_var,
        rand,
        logpdf,
        elbo,
        dtc,
        posterior,
        approx_posterior,
        marginals

    using MacroTools: @capture, combinedef, postwalk, splitdef

    const AV{T} = AbstractVector{T}

    # Various bits of utility that aren't inherently GP-related. Often very type-piratic.
    include(joinpath("util", "zygote_rules.jl"))
    include(joinpath("util", "covariance_matrices.jl"))
    include(joinpath("util", "block_arrays", "dense.jl"))
    include(joinpath("util", "block_arrays", "diagonal.jl"))
    include(joinpath("util", "abstract_data_set.jl"))
    include(joinpath("util", "proper_type_piracy.jl"))

    # Supertype for GPs.
    include("abstract_gp.jl")

    # Atomic GP objects.
    include(joinpath("gp", "gp.jl"))

    # Composite GPs, constructed via affine transformation of CompositeGPs and GPs.
    include(joinpath("composite", "composite_gp.jl"))
    include(joinpath("composite", "cross.jl"))
    include(joinpath("composite", "product.jl"))
    include(joinpath("composite", "addition.jl"))
    include(joinpath("composite", "compose.jl"))
    # include(joinpath("composite", "gradient.jl"))
    # include(joinpath("composite", "integrate.jl"))

    # Gaussian Process Probabilistic Programme object which implements the AbstractGPs API.
    include("gaussian_process_probabilistic_programme.jl")

    # Sparse GP hack to make pseudo-point approximations play nicely with Turing.jl.
    include("sparse_finite_gp.jl")

    include("deprecate.jl")

    export wrap, BlockData, GPC, GPPPInput, @gppp
    export ∘, select, stretch, periodic, shift
    export cov_diag, mean_and_cov_diag
end # module


using SeisRequests
using CSV
using DataFrames
using Dates
import DataStructures: OrderedDict

const MBool = Union{Missing,Bool}
const MDateTime = Union{Missing,DateTime}
const MFloat = Union{Missing,Float64}
const MInt = Union{Missing,Int}
const MString = Union{Missing,String}

"Available processing filters and allowable types for them"
const IRISTimeSeries_PROCESSING_FIELDS = Dict(
    :taper => Real,
    :taper_type => AbstractString,
    :envelope => Bool,
    :lpfilter => Real,
    :hpfilter => Real,
    :bpfilter => Tuple{Number,Number},
    :demean => Bool,
    :scale => Union{Real,AbstractString},
    :divscale => Real,
    :correct => Bool,
    :freqlimits => Tuple{Number,Number,Number,Number},
    :autolimits => Union{AbstractString,Tuple{Number,Number}},
    :units => AbstractString,
    :diff => Bool,
    :int => Bool,
    :decimate => Real)

const IRISTimeSeries_OUTPUT_TYPES = ("ascii", "ascii2", "ascii", "geocsv", "geocsv.tspair",
"geocsv.slist", "audio", "miniseed", "plot", "saca", "sacbb", "sacbl")

struct IRISMassRequest
  # Obspy Domain & Restriction class
  network::String
  station::String
  location::String
  channel::String
  starttime::DateTime
  endtime::DateTime
  duration::MFloat
  output::String
  process::OrderedDict{Symbol,Any}
  minlatitude::MFloat
  maxlatitude::MFloat
  minlongitude::MFloat
  maxlongitude::MFloat
  longitude::MFloat
  latitude::MFloat
  minradius::MFloat
  maxradius::MFloat
  startbefore::MDateTime
  startafter::MDateTime
  endbefore::MDateTime
  endafter::MDateTime
  reject_channels_with_gaps::MBool
  minimum_length::MFloat
  minimum_interstation_distance_in_m::MInt

"""

"""
  function IRISMassRequest(network, station, location, channel, starttime, endtime, duration,
                           output, process, minlatitude, maxlatitude, minlongitude, maxlongitude,
                            longitude, latitude, minradius, maxradius, startbefore, startafter,
                            endbefore, endafter, reject_channels_with_gaps, minimum_length,
                            minimum_interstation_distance_in_m)
        location == "  " && (location = "--")
        !ismissing(starttime) && !ismissing(endtime) && starttime > endtime &&
            throw(ArgumentError("`starttime` must be before endtime"))
        coalesce(duration, 1) > 0 || throw(ArgumentError("`duration` must be positive"))
        0 <= get(process, :taper, 0.1) <= 1 || throw(ArgumentError("`taper` length must be between 0 and 1"))
        # Check types for processing
        for (f, T) in IRISTimeSeries_PROCESSING_FIELDS
            if haskey(process, f)
                try
                    process[f] = convert(T, process[f])
                catch err
                    err isa MethodError || err isa InexactError &&
                        throw(ArgumentError("field `$f` must be of type $T"))
                    rethrow(err)
                end
            end
        end
        if haskey(process, :taper_type)
            haskey(process, :taper) || throw(ArgumentError("`taper_type` must be specficied with taper"))
            taper_type = uppercase(process[:taper_type])
            taper_type in ("HANNING", "HAMMING", "COSINE") ||
            throw(ArgumentError("taper_type must be one of \"HANNING\", \"HAMMING\" or \"COSINE\""))
        end
        haskey(process, :scale) && haskey(process, :divscale) &&
            throw(ArgumentError("Cannot specify both `scale` and `divscale`"))
        if haskey(process, :scale)
            if process[:scale] isa AbstractString
                process[:scale] = uppercase(process[:scale])
                process[:scale] == "AUTO" || throw(ArgumentError("`scale` must be a float or \"AUTO\""))
            end
        end
        if haskey(process, :units)
            !haskey(process, :correct) || !process[:correct] &&
                throw(ArgumentError("`units` can only be used when correct is true"))
            process[:units] in ("DIS", "VEL", "ACC", "DEF") ||
                throw(ArgumentError("`units` must be one of \"DIS\", \"VEL\", \"ACC\" or \"DEF\""))
        end
        haskey(process, :autolimits) && haskey(process, :freqlimits) &&
            throw(ArgumentError("Cannot specify both `autolimits` and `freqlimits`"))

        # Disallowed values
        output in IRISTimeSeries_OUTPUT_TYPES ||
            throw(ArgumentError("`output` must be one of $(IRISTimeSeries_OUTPUT_TYPES)"))

        # MassRequest values

        new(network, station, location, channel, starttime, endtime, duration,
            output, process, minlatitude, maxlatitude, minlongitude, maxlongitude,
            longitude, latitude, minradius, maxradius, startbefore, startafter,
            endbefore, endafter, reject_channels_with_gaps, minimum_length,
            minimum_interstation_distance_in_m)
    end
end

function IRISMassRequest(;
                        network=missing, station=missing, location=missing, channel=missing,
                        starttime=missing, endtime=missing, duration=missing,
                        output=missing, minlatitude=missing, maxlatitude=missing,
                        minlongitude=missing, maxlongitude=missing, longitude=missing,
                        latitude=missing, minradius=missing, maxradius=missing,
                        startbefore=missing, startafter=missing, endbefore=missing,
                        endafter=missing, reject_channels_with_gaps=missing,
                        minimum_length=missing, minimum_interstation_distance_in_m=missing,
                        kwargs...)
    any(ismissing, (network, station, location, channel, starttime, output)) &&
        throw(ArgumentError("network, station, location, channel, starttime " *
                            "and output must all be specified"))
    all(ismissing, (endtime, duration)) &&
        throw(ArgumentError("One of `endtime` or `duration` must be specified"))
    process = OrderedDict{Symbol,Any}()
    for (k, v) in kwargs
        k in keys(IRISTimeSeries_PROCESSING_FIELDS) ||
            throw(ArgumentError("Field `$k` is not a valid IRISTimeSeries processing field"))
        process[k] = v
    end
    IRISMassRequest(network, station, location, channel, starttime, endtime, duration,
                    lowercase(output), process, minlatitude, maxlatitude, minlongitude,
                    maxlongitude, longitude, latitude, minradius, maxradius, startbefore,
                    startafter, endbefore, endafter, reject_channels_with_gaps, minimum_length,
                   minimum_interstation_distance_in_m)
end

"""
  station_request(station,network,channel)

"""
function station_availability(mass_request::IRISMassRequest;
                         level::String="channel", format::String="text")

  # if ismissing(starttime)
  #   starttime = DateTime(1900,1,1)
  # end
  # if ismissing(endtime)
  #   endtime = DateTime(2599,12,31,23,59,59)
  # end

  request = FDSNStation(station = mass_request.station, network = mass_request.network,
                        channel = mass_request.channel,
                        level = level, format = format,
                        starttime = mass_request.starttime,
                        endtime = mass_request.endtime,
                        startbefore = mass_request.startbefore,
                        startafter = mass_request.startafter,
                        endbefore = mass_request.endbefore,
                        endafter = mass_request.endafter,
                        location = mass_request.location,
                        minlatitude = mass_request.minlatitude,
                        maxlatitude = mass_request.maxlatitude,
                        minlongitude = mass_request.minlongitude,
                        maxlongitude = mass_request.maxlongitude,
                        longitude = mass_request.longitude,
                        latitude = mass_request.latitude,
                        minradius = mass_request.minradius,
                        maxradius = mass_request.maxradius)
  receive = get_request(request)
  text = String(receive.body)
  text = strip(text, '#')
  df = CSV.read(IOBuffer(text),delim='|', typemap=Dict(MDateTime=>DateTime,
                MFloat=>Float64,MString=>String))
  newnames =  map(lowercase, String.(names(df)))
  newnames = strip.(newnames)
  newnames = Symbol.(newnames)
  names!(df, newnames)

  return df
end

"""
  mass_download(mass_request, mseed_storage, stationxml_storage)

Launch a data download.

INPUTS TO DO:
  - define domain structure for lat/lon filtering
  - define restrictions structure for time/network/station/channel filtering
"""
function mass_download(mass_request::IRISMassRequest,
                  mseed_storage::String,
                  stationxml_storage::String;
                  download_chunk_size_in_mb::Int=20,
                  threads_per_client::Int=3)
  # function will mirror the obspy mass_downloader.download() function
  # Order of operations TO-DO:
  # 1. Get station availability with station_availability() DONE
  # 2. Setup directory structure based on mseed_storage path DONE
  # 3. Setup directory stucture based on stationxml_storage path DONE
  # 4. Download stationxml for each station - should include response DONE
  # 5. Chunk download into requests based on length of time per mseed file DONE
  #    This should be done in parallel in future using pmap or threads
  # 6. Send HTTP requests to IRIS using IRISTimeSeries DONE
  # 7. Save request to miniseed DONE

  df = station_availability(mass_request)
  # update latest date to today
  latest = DateTime(year(now()),month(now()), day(now()))
  df[:endtime][df[:endtime] .> latest,:] .= latest

  # create directory structure for mseed
  n = size(df,1)
  for ii in 1:n
    current = joinpath(mseed_storage,df[:network][ii],df[:station][ii],
                       df[:channel][ii])
    if isdir(current) == false
      mkpath(current)
    end
  end

  # make directory for stationxml_storage
  if isdir(stationxml_storage) == false
    mkpath(stationxml_storage)
  end

  # download stationxml for each station
  download_stationxml(stationxml_storage,df,starttime=mass_request.starttime,
                      endtime=mass_request.endtime,
                      startbefore=mass_request.startbefore,
                      startafter=mass_request.startafter,
                      endbefore=mass_request.endbefore,
                      endafter=mass_request.endafter,
                      location=mass_request.location)

  # download each network/channel/station
  for ii = 1:n
    NSLC = df[ii,:]
    s, e = NSLC[[:starttime,:endtime]]
    if s < mass_request.starttime; s = mass_request.starttime end
    if e > mass_request.endtime; e = mass_request.endtime end
    station, network, location, channel = NSLC[[:station, :network,
                                                :location, :channel]]
    if ismissing(location); location = "--" end

    # check begin/end dates
    if s != DateTime(year(s),month(s),day(s))
      s = DateTime(year(s),month(s),day(s)) + Dates.Day(1)
    end

    if e != DateTime(year(e),month(e),day(e))
      e = DateTime(year(e),month(e),day(e))
    end

    # start/end date ranges
    start_range = s:Dates.Day(1):e-Dates.Day(1)
    end_range = s+Dates.Day(1) - Dates.Millisecond(1/NSLC[:samplerate][1]
                                                  * 1000) :Dates.Day(1):e

    # download from starttime to endtime in chunks
    for jj=1:length(start_range)
      # get filepath
      filename = get_mseed_name(channel, location, start_range[jj],
                                end_range[jj])
      filepath = joinpath(mseed_storage, network, station, channel, filename)
      download_mseed(station, network, channel,start_range[jj],end_range[jj],
                     filepath)
    end
  end

end

"""
  download_stationxml(args)

Downloads StationXML instrument response.
TODO - verify stationXML are formatted correctly
     - will need to load each file with LightXML
     - this should probably be a separate function
     - make this function threaded with using `threads_per_client`
"""
function download_stationxml(stationxml_storage::String, df::DataFrame;
                             starttime::MDateTime=missing,
                             endtime::MDateTime=missing,
                             startbefore::MDateTime=missing,
                             startafter::MDateTime=missing,
                             endbefore::MDateTime=missing,
                             endafter::MDateTime=missing,
                             location::MString=missing)

  # get all unique network/station combinations
  station_df = unique(df[[:network, :station]])
  n = size(station_df,1)
  for ii in 1:n
    net, sta = station_df[:network][ii], station_df[:station][ii]
    request = FDSNStation(station=sta, network=net,
                        level="response", format="xml", starttime=starttime,
                        endtime=endtime,startbefore=startbefore,
                        startafter=startafter,endbefore=endbefore,
                        endafter=endafter,location=location)
    receive = get_request(request)
    text = String(receive.body)
    stationxml_filename = join([net,sta,"xml"],'.')
    stationxml_path = joinpath(stationxml_storage,stationxml_filename)
    open(stationxml_path,"w") do f
      write(f,text)
    end
    print("Downloaded stationXML for $net.$sta $ii out of $n")
  end
end

"""
   download_mseed(station, network, channel, starttime, endtime, filename)

Download mseed from IRIS.
TO-DO: - try/catch with bad requests
       - print statement with completed requests
       - allow for chunked requests with `download_chunk_size_in_mb`
"""
function download_mseed(station::String, network::String, channel::String,
                      starttime::DateTime,endtime::DateTime,
                      filepath::String)

  try
      request = get_request(IRISTimeSeries(network=network,
                                           station=station,
                                           channel=channel,
                                           location=location,
                                           starttime=starttime,
                                           endtime=endtime,
                                           output="miniseed"))
      text = request.body
      # write to mseed to filepath
      open(filepath, "w") do file
        write(file,text)
      end
      println("Downloaded $network.$station.$location.$channel.$starttime.$endtime")
  catch
      println("Unable to download $network.$station.$location.$channel.$starttime.$endtime")
  end
end

"""
  get_mseed_name(channel, location starttime, endtime)

"""
function get_mseed_name(channel, location, starttime, endtime)
  fmt_string = "%Y%m%dT%H:%M:%S"
  s = Dates.format(starttime, fmt_string)
  e = Dates.format(endtime, fmt_string)
  if location == "--"; location = "" end
  return "$channel.$location.$(starttime)Z.$(endtime)Z.mseed"
  end

network = "CI"
station = "MWC"
channel = "BHZ"
location = "--"
starttime = DateTime(2019, 2, 1, 0, 0, 0)
endtime = now()
req = IRISMassRequest(network=network, station=station, channel=channel,
                      location=location, starttime=starttime,
                      endtime=endtime, output="miniseed")
mseed_storage = "/home/timclements/TEST/MSEED"
stationxml_storage = "/home/timclements/TEST/XML"
mass_download(req,mseed_storage, stationxml_storage)


using FluxRM
using FluxUtils

using FluxRM
const flux = Flux()

const kvs = FluxRM.KVS(flux)
FluxRM.transaction(kvs) do txn
    FluxRM.mkdir!(txn, "run_snap")
end

const __job_counter = Ref{Int64}(0)
function job_counter()
    counter = __job_counter[]
    __job_counter[] = counter + 1
    counter
end

function md_job(params)
    counter = job_counter()

    path = "run_snap.job_$counter"

    FluxRM.transaction(kvs) do txn
        FluxRM.mkdir!(txn, path)
        put_raw!(txn, path * ".input", Dict(:data => params))
    end

    script = """
    const kvs = FluxRM.KVS(Flux())
    input = FluxRM.lookup(kvs, $path * ".input")
    params = input[:data]

    out = run_md(params)

    FluxRM.transaction(kvs) do txn
        put_raw!(txn, $path * ".output", Dict(:data => out))
    end
    """

    return FluxUtils.Job(".", "utils.jl", script)
end

tasks = []
for (i, β) in enumerate(β_samples)
    md_params = MDParams(
        vcat([βi for βi in β]...),
        params, starting_file, save_dir*"$(i)/", Tend,
        dt, i, Temp, dT)

    job = md_job(md_params)
    push!(tasks, run(job))
end

# forall(tasks) do task
#     # wait(task)
#     if success(task)
#         out = fetch(tasks)
#         #  something productive here
#     else
#         @warn "task failed" task
#     end
# end

# function get_params!(server)
# end

# function update_params!(server, gradient)
# end

# for iter in interations
#     params = get_params!(server)
#     gradient = Zygote.get_gradients(run_md, params)
#     update_params!(server, gradient)
# end


################################################################################
#
#   HELPER FUNCTIONS FOR CONSTRUCTION OF LATTICES / NEIGHBORS
#
#   FILE CONTAINS
#   - bond concanation
#
################################################################################

################################################################################
#
#   CONCANATION OF BONDS
#
################################################################################

# N=0
function concanateBonds(
            b1 :: B,
            b2 :: B,
            labelfunction :: Function = (bo1,bo2) -> getDefaultLabel(L)
        ) :: B where {L,B<:AbstractBond{L,0}}

    # assert that the bonds can indeed by concanated
    @assert to(b1) == from(b2) "Bonds cannot be concanated as "*string(from(b1))*"-"*string(to(b1))*" (-) "*string(from(b2))*"-"*string(to(b2))*" not possible"

    # create a new bond
    return newBond(
        B,
        from(b1) :: Int64,
        to(b2) :: Int64,
        labelfunction(b1,b2) :: L,
        NTuple{0,Int64}()
    )
end
# N=1
function concanateBonds(
            b1 :: B,
            b2 :: B,
            labelfunction :: Function = (bo1,bo2) -> getDefaultLabel(L)
        ) :: B where {L,B<:AbstractBond{L,1}}

    # assert that the bonds can indeed by concanated
    @assert to(b1) == from(b2) "Bonds cannot be concanated as "*string(from(b1))*"-"*string(to(b1))*" (-) "*string(from(b2))*"-"*string(to(b2))*" not possible"

    # create a new bond
    return newBond(
        B,
        from(b1) :: Int64,
        to(b2) :: Int64,
        labelfunction(b1,b2) :: L,
        (
            wrap(b1)[1] + wrap(b2)[1],
        )
    )
end
# N=2
function concanateBonds(
            b1 :: B,
            b2 :: B,
            labelfunction :: Function = (bo1,bo2) -> getDefaultLabel(L)
        ) :: B where {L,B<:AbstractBond{L,2}}

    # assert that the bonds can indeed by concanated
    @assert to(b1) == from(b2) "Bonds cannot be concanated as "*string(from(b1))*"-"*string(to(b1))*" (-) "*string(from(b2))*"-"*string(to(b2))*" not possible"

    # create a new bond
    return newBond(
        B,
        from(b1) :: Int64,
        to(b2) :: Int64,
        labelfunction(b1,b2) :: L,
        (
            wrap(b1)[1] + wrap(b2)[1],
            wrap(b1)[2] + wrap(b2)[2]
        )
    )
end
# N=3
function concanateBonds(
            b1 :: B,
            b2 :: B,
            labelfunction :: Function = (bo1,bo2) -> getDefaultLabel(L)
        ) :: B where {L,B<:AbstractBond{L,3}}

    # assert that the bonds can indeed by concanated
    @assert to(b1) == from(b2) "Bonds cannot be concanated as "*string(from(b1))*"-"*string(to(b1))*" (-) "*string(from(b2))*"-"*string(to(b2))*" not possible"

    # create a new bond
    return newBond(
        B,
        from(b1) :: Int64,
        to(b2) :: Int64,
        labelfunction(b1,b2) :: L,
        (
            wrap(b1)[1] + wrap(b2)[1],
            wrap(b1)[2] + wrap(b2)[2],
            wrap(b1)[3] + wrap(b2)[3]
        )
    )
end

# export the concanation function
export concanateBonds


using Test
using Printf: @printf
using LinearAlgebra: tr
using QuantumLattices.Essentials.Frameworks
using QuantumLattices.Essentials.Terms: Parameters
using QuantumLattices.Interfaces: id, register!
using QuantumLattices.Prerequisites: Float
import QuantumLattices.Interfaces: prepare!, run!, update!
import QuantumLattices.Essentials.Frameworks: dependences

struct FEngine <: Engine end
Base.repr(::FEngine) = "FEngine"
Base.show(io::IO, ::FEngine) = @printf io "FEngine"
struct FApp <: App end

@testset "Assignment" begin
    @test FApp() == FApp()
    @test isequal(FApp(), FApp())
    @test update!(FApp()) == FApp()

    assign = Assignment(:FApp, FApp(), (t=1.0, U=8.0), dependences=(:FApp1, :FApp2))
    @test deepcopy(assign) == assign
    @test isequal(deepcopy(assign), assign)
    @test assign|>valtype == assign|>typeof|>valtype == Any
    @test assign|>id == assign|>typeof|>id == :FApp
    update!(assign, t=2.0)
    @test assign.parameters == (t=2.0, U=8.0)
end

@testset "Algorithm" begin
    @test FEngine() == FEngine()
    @test isequal(FEngine(), FEngine())
    @test update!(FEngine()) == FEngine()

    alg = Algorithm("Alg", FEngine(), parameters=(t=1.0, U=8.0))
    @test string(alg) == repr(alg) == "Alg_FEngine_1.0_8.0"
    @test repr(alg, (:U,)) == "Alg_FEngine_1.0"
    register!(alg, :FApp1, FApp(), parameters=(U=5.0,))
    register!(alg, :FApp2, FApp(), parameters=(U=6.0,), dependences=(:FApp1,))
    @test get(alg, Val(:FApp1)) == get(alg, :FApp1) == Assignment(:FApp1, FApp(), (t=1.0, U=5.0), virgin=false)
    @test get(alg, Val(:FApp2)) == get(alg, :FApp2) == Assignment(:FApp2, FApp(), (t=1.0, U=6.0), dependences=(:FApp1,), virgin=false)
    @test prepare!(alg, get(alg, :FApp1)) === nothing
    @test run!(alg, get(alg, :FApp1)) === nothing
    @test dependences(alg, get(alg, :FApp1)) == ()
    @test dependences(alg, get(alg, :FApp2)) == (:FApp1,)
    @test dependences(alg, get(alg, :FApp2), (:FApp1,)) == ()
end

mutable struct VCA <: Engine
    t::Float
    U::Float
end
function update!(vca::VCA; kwargs...)
    vca.t = get(kwargs, :t, vca.t)
    vca.U = get(kwargs, :U, vca.U)
    return vca
end
Parameters(vca::VCA) = Parameters{(:t, :U)}(vca.t, vca.U)
gf(alg::Algorithm{VCA}) = get(alg, Val(:_VCAGF_)).data

mutable struct GF <: App
    dim::Int
    count::Int
end

mutable struct DOS <: App
    mu::Float
end
update!(eb::DOS; kwargs...) = (eb.mu = get(kwargs, :mu, eb.mu); eb)

dependences(alg::Algorithm{VCA}, assign::Assignment{GF}, ::Tuple{}=()) = assign.dependences
dependences(alg::Algorithm{VCA}, assign::Assignment, ::Tuple{}=()) = (:_VCAGF_, assign.dependences...)
prepare!(alg::Algorithm{VCA}, assign::Assignment{GF}) = assign.virgin && (assign.data = Matrix{valtype(assign)|>eltype}(undef, assign.app.dim, assign.app.dim))
run!(alg::Algorithm{VCA}, assign::Assignment{GF}) = (assign.app.count += 1; assign.data[:, :] .= alg.engine.t+alg.engine.U)

function run!(alg::Algorithm{VCA}, assign::Assignment{DOS})
    rundependences!(alg, assign)
    assign.data = tr(gf(alg))
end

@testset "Framework" begin
    engine = VCA(1.0, 8.0)
    gf = Assignment(:_VCAGF_, GF(4, 0), Parameters(engine), data=Matrix{Complex{Float}}(undef, 0, 0))
    vca = Algorithm("Test", engine, assignments=(gf,))
    register!(vca, :DOS, DOS(-3.5), parameters=(U=7.0,), map=params::Parameters->Parameters{(:t, :U, :mu)}(params.t, params.U, -params.U/2))
    dos = get(vca, :DOS)
    @test dos.data == 32.0
    @test dos.app.mu == -3.5
    update!(dos, U=6.0)
    run!(vca, :DOS, false)
    @test dos.data == 28.0
    @test dos.app.mu == -3.0
end


module EnergyModels

# using Requires
using CSV
using Destruct
using JuMP
using AxisArrays
using LightGraphs
using MetaGraphs
using DataFrames
using Logging
using SparseArrays

import PowerModels
const PM = PowerModels

import MathOptInterface
const MOI = MathOptInterface

using Base.Iterators: flatten

# using Memento

# # Create our module level logger (this will get precompiled)
# const logger = getlogger(@__MODULE__)

# # Register the module level logger at runtime so that folks can access the logger via `get_logger(MyModule)`
# # NOTE: If this line is not included then the precompiled `MyModule.LOGGER` won't be registered at runtime.
# function __init__()
#     Memento.register(logger)
# end

function __init__()
    # @require TimeSeries="9e3dc215-6440-5c97-bce1-76c03772f85e" nothing
    # @require PowerSystems="bcd98974-b02a-5e2f-9ee0-a103f5c450dd" include("data/powersystems.jl")
    # @require PyCall="438e738f-606a-5dbb-bf0a-cddfbfd45ab0" include("data/pypsanetwork.jl")
end

include("abstracttypes.jl")
include("core.jl")

include("formulation.jl")
include("compat.jl")
include("macros.jl")
include("modelview.jl")
include("wrappedarray.jl")
include("registry.jl")
include("graph.jl")
include("containerviews.jl")

include("components/components.jl")
include("data/data.jl")

export add!, addto!, build!, optimize!, store_results!, set_snapshots!, devices,
    subnetworks, buses, push!, axis, graph, determine_subnetworks!, jumpmodel,
    get, getjump, getvalue, getdual, getparam

# Components and Devices
export EnergyModel, SubNetwork, Device, Bus, Line, Transformer, Link,
    Load, Generator, StorageUnit, Store, Branch, PassiveBranch, ActiveBranch,
    OnePort

# TODO Generate a precompile file with SnoopCompiler

end # module


using Flux: BilinearUpsample
using Test

@testset "BilinearUpsample" begin
  @test size(BilinearUpsample((2, 2))(rand(2, 2, 1, 1))) == (4, 4, 1, 1)
  @test size(BilinearUpsample((3, 3))(rand(2, 2, 1, 1))) == (6, 6, 1, 1)
  @test size(BilinearUpsample((2, 2))(rand(2, 2, 10, 10))) == (4, 4, 10, 10)
  @test size(BilinearUpsample((3, 3))(rand(2, 2, 10, 10))) == (6, 6, 10, 10)

  @test_throws BoundsError BilinearUpsample((2, 2))(rand(2, 2))
end


"""
# Description

A logical node that represents a group-by operation. It has two fields:

1. `source`: A node describing a relation that will be divided into groups.
2. `expressions`: A tuple of `Expression` objects that define the groups.
"""
struct GroupBy <: LogicalNode
    source::Any
    expressions::Tuple
end

function Base.print(io::IO, node::GroupBy)
    println(io, "GroupBy LogicalNode")
    println(io, type(node.source))
    for e in node.expressions
        println(io, e)
    end
end


module MzXML

using Base64
using LightXML, Unitful, ProgressMeter
using MzCore
using MzCore.IntervalSets, MzCore.AxisArrays
using CodecZlib

"""
    MSscan(polarity::Char, msLevel::Int, retentionTime::typeof(1.0u"s"),
           lowMz::Float64, highMz::Float64, basePeakMz::Float64,
           totIonCurrent::Float64, mz::Vector, I::Vector, children::Vector{MSscan})

Represent mass spectrometry data at a single time point, scanning over all `m/z`.
`polarity` is `+` or `-` for positive-mode or negative-mode, respectively.
`msLevel` is 1 for MS, 2 for MS², and so on.
Retention time is represented using the Unitful package.

`mz` and `I` should be vectors of the same length, with an ion count of `I[i]` at `mz[i]`.
`children` refers to MS² or higher for particular ions in this scan.
"""
struct MSscan{T<:AbstractFloat,TI<:Real}
    polarity::Char
    msLevel::Int
    retentionTime::typeof(1.0u"s")
    lowMz::Float64
    highMz::Float64
    basePeakMz::Float64
    totIonCurrent::Float64
    mz::Vector{T}
    I::Vector{TI}
    children::Vector{MSscan{T,TI}}
end

MzCore.intensitytype(::Type{MSscan{T,TI}}) where {T<:AbstractFloat,TI<:Real} = TI
MzCore.mztype(::Type{MSscan{T,TI}}) where {T<:AbstractFloat,TI<:Real} = T

const empty32 = MSscan{Float32,Float32}[]
const empty64 = MSscan{Float64,Float64}[]

function Base.show(io::IO, scan::MSscan)
    if scan.msLevel > 1
        print(io, "  "^(scan.msLevel-2), "└─basePeak ", scan.basePeakMz, ": ")
    end
    print(io, scan.retentionTime, ": ")
    mzI = [mz=>I for (mz, I) in zip(scan.mz, scan.I)]
    println(io, mzI)
end

"""
    scanpos = MzXML.index(filename)

Return the vector of file positions corresponding to each scan.
"""
function index(filename)
    scanpositions = open(filename) do file
        # Find the indexOffset element
        seekend(file)
        p = position(file)
        l = 200
        if p < l
            error("Too short to be mzXML")
        end
        skip(file, -l)
        str = loadrest(file)
        m = match(r"<indexOffset>([0-9].*)</indexOffset>", str)
        if m == nothing
            error("Cannot find indexOffset element")
        end
        length(m.captures) == 1 || error("indexOffset should contain a single number")
        indexpos = parse(Int, m.captures[1])
        length_tail = l - m.offset
        # Read the index
        seek(file, indexpos)
        str = loadrest(file)
        xindex = parse_string(str[1:end-length_tail-1])
        xroot = root(xindex)
        scanpositions = Int[]
        for c in child_elements(xroot)
            name(c) == "offset" || error("index could not be parsed")
            push!(scanpositions, parse(Int, content(c)))
        end
        scanpositions
    end
    scanpositions
end

"""
    scans, info = MzXML.load(filename; productlevels=0, timeinterval=0.0u"s" .. Inf*u"s", timeshift=0.0u"s")

Load a `*.mzXML` file. `scans` is a vector of scan structures, each at a particular scan time.

`productlevels` controls the number of product-ion levels that will be loaded in combined MS/MS² or similar runs.
The "base" MS level (the first encountered in the file) will always be loaded--`productlevels` controls the
number of additional levels.
By default no product-ion data will be loaded (`productlevels=0`), but you can request it by increasing `productlevels` to 1 or higher.

`scans` will include only those scans occuring within `timeinterval`, with the default to include all times.
The `..` is from the `IntervalSets` package.
"""
function load(filename; productlevels=0, kwargs...)
    xdoc = parse_file(filename)
    xroot = root(xdoc)
    if name(xroot) != "mzXML"
        error("Not an mzXML file")
    end

    # Find the msRun node
    msRun = find_element(xroot, "msRun")

    props = Dict{Symbol,Any}()
    props[:startTime] = parse_time(attribute(msRun, "startTime"))
    props[:endTime] = parse_time(attribute(msRun, "endTime"))
    el = find_element(msRun, "msInstrument")
    el = find_element(el, "msModel")
    props[:msModel] = attribute(el, "value")
    props[:productlevels] = productlevels

    scans = load_scans(msRun, productlevels; kwargs...)
    LightXML.free(xdoc)
    return scans, props
end

function load_scans(elm, productlevels=0; kwargs...)
    # First we discover the data type
    local T
    local TI
    for c in child_elements(elm)
        n = name(c)
        if n != "scan"
            continue
        end
        peak = find_element(c, "peaks")
        TI, T, nochildren = precisiondict[attribute(peak, "precision")]
        break
    end
    # Now load the data
    scans = MSscan{T,TI}[]
    load_scans!(scans, elm, productlevels; kwargs...)
end

@noinline function load_scans!(scans, elm, productlevels; kwargs...)
    prog = ProgressUnknown("Number of scans read:")
    for c in child_elements(elm)
        n = name(c)
        if n != "scan"
            continue
        end
        scan = load_scan(c, productlevels; kwargs...)::Union{Nothing,eltype(scans)}
        if scan !== nothing
            push!(scans, scan)
            next!(prog)
        end
    end
    finish!(prog)
    scans
end

polaritydict = Dict("+" => '+', "-" => '-')
precisiondict = Dict("32" => (Float32, Float32, empty32), "64" => (Float64, Float64, empty64))

function load_scan(elm, productlevels; timeinterval=0.0u"s" .. Inf*u"s", timeshift=0.0u"s")
    polarity = polaritydict[attribute(elm, "polarity")]
    retentionTime = parse_time(attribute(elm, "retentionTime")) + timeshift
    retentionTime ∈ timeinterval || return nothing
    lMza, hMza = attribute(elm, "lowMz"), attribute(elm, "highMz")
    if lMza !== nothing && hMza !== nothing
        lowMz, highMz = parse(Float64, lMza), parse(Float64, hMza)
    else
        lowMz = highMz = NaN
    end
    msLevela = attribute(elm, "msLevel")
    msLevel = msLevela === nothing ? 1 : parse(Int, msLevela)
    basePeakMz = parse(Float64, attribute(elm, "basePeakMz"))
    totIonCurrent = parse(Float64, attribute(elm, "totIonCurrent"))
    npeaks = parse(Int, attribute(elm, "peaksCount"))
    peak = find_element(elm, "peaks")
    data = base64decode(content(peak))
    compression_type = attribute(peak, "compressionType")
    if compression_type == "zlib"
        data = transcode(ZlibDecompressor, data)
    elseif compression_type != "none"
        error("compressionType $compression_type not supported")
    end
    TI, T, nochildren = precisiondict[attribute(peak, "precision")]
    A = reinterpret(TI, data)
    bo = attribute(peak, "byteOrder")
    if bo == "network"
        ntoh!(A)
    else
        error("Don't know what to do with byteOrder $bo")
    end
    po = attribute(peak, "pairOrder")
    if po === nothing
        po = attribute(peak, "contentType")
    end
    po == "m/z-int" || error("Don't know what to do with pairOrder/contentType $po")
    I = A[2:2:end]
    mz = reinterpret(T, A[1:2:end])
    if !issorted(mz)   # it's unclear whether the standard requires these to be sorted, so better to check
        perm = sortperm(mz)
        I, mz = I[perm], mz[perm]
    end
    children = productlevels > 0 ? load_scans!(MSscan{T,TI}[], elm, productlevels-1) : nochildren
    MSscan{T,TI}(polarity, msLevel, retentionTime, lowMz, highMz, basePeakMz, totIonCurrent, mz, I, children)
end

function loadrest(file::IOStream)
    nb = filesize(file) - position(file)
    b = read(file, nb)
    return String(b)
end

function parse_time(tstr)
    if startswith(tstr, "PT") && endswith(tstr, "S")
        return parse(Float64, tstr[3:end-1])*u"s"
    end
    error("Time string $tstr not recognized")
end

function ntoh!(A)
    for i = 1:length(A)
        A[i] = ntoh(A[i])
    end
end

Base.pairs(scan::MSscan) = zip(scan.mz, scan.I)

function val2index(ax::Axis, val)
    f, s = first(ax.val), step(ax.val)
    r = (val - f) / s
    idx = floor(Int, r)
    return idx + 1, r - idx
end

"""
    copyto!(dest::AxisArray, src::Vector{<:MSscan})

Convert the "sparse-matrix" representation `src` of mass spectrometry
data into a dense 2d matrix `dest`. `dest` must have one axis named
`:mz` and the other `:time`; the `:time` axis must have physical units of
time.

Any previous contents of `dest` are erased. The values of `src` are
accumulated into `dest` using bilinear interpolation.
"""
function Base.copyto!(dest::AxisArray, src::Vector{<:MSscan})
    axmz, axt = AxisArrays.axes(dest, Axis{:mz}), AxisArrays.axes(dest, Axis{:time})
    fill!(dest, 0)
    for scan in src
        j, ft = val2index(axt, scan.retentionTime)
        firstindex(axt.val) - 1 <= j <= lastindex(axt.val) || continue
        for (mz, I) in pairs(scan)
            i, fmz = val2index(axmz, mz)
            # Accumulate into `dest`, using bilinear interpolation to
            # avoid aliasing
            if firstindex(axt.val) <= j <= lastindex(axt.val)
                if firstindex(axmz.val) <= i <= lastindex(axmz.val)
                    dest[axmz(i),axt(j)] += (1-fmz)*(1-ft)*I
                end
                if firstindex(axmz.val) <= i+1 <= lastindex(axmz.val)
                    dest[axmz(i+1),axt(j)] += fmz*(1-ft)*I
                end
            end
            if firstindex(axt.val) <= j+1 <= lastindex(axt.val)
                if firstindex(axmz.val) <= i <= lastindex(axmz.val)
                    dest[axmz(i),axt(j+1)] += (1-fmz)*ft*I
                end
                if firstindex(axmz.val) <= i+1 <= lastindex(axmz.val)
                    dest[axmz(i+1),axt(j+1)] += fmz*ft*I
                end
            end
        end
    end
    return dest
end

# An internal type for nested `extrema` calls (yielding the min of all mins, max of all maxs)
struct ExPair{T}
    min::T
    max::T
end
Base.min(a::ExPair, b::ExPair) = a.min <= b.min ? a : b
Base.max(a::ExPair, b::ExPair) = a.max >= b.max ? a : b

function MzCore.limits(scans::Vector{<:MSscan})
    function extrema_mz(scan)
        isfinite(scan.lowMz) && isfinite(scan.highMz) && return ExPair(scan.lowMz, scan.highMz)
        return ExPair(first(scan.mz), last(scan.mz))  # safe because guaranteed to be sorted
    end
    tmin, tmax   = extrema(scan->scan.retentionTime, scans)
    mzmin, mzmax = extrema(extrema_mz, scans)
    return Axis{:mz}(mzmin.min .. mzmax.max), Axis{:time}(tmin..tmax)
end

end


#### Obstacle generation ####

function empty_obstacles!(env::OccludedEnv)
    env.obstacles = Vector{ConvexPolygon}[]
end

function add_obstacle!(env::OccludedEnv, pos::VecSE2, width::Float64, height::Float64)
    obstacle = get_oriented_bounding_box(pos, width, height)
    push!(env.obstacles, obstacle)
end

function left_obstacle_space(env::OccludedEnv; y_min=env.params.y_min/4)
    x_min = env.params.x_min
    x_max = env.params.inter_x - env.params.nlanes*env.params.lane_width - maximum(env.params.crosswalk_width/2)
    y_max = env.params.inter_y-env.params.nlanes_main*env.params.lane_width
    return x_min, x_max, y_min, y_max
end

function get_width(poly::ConvexPolygon)
    width_seg = get_edge(poly, 1)
    width = dist(width_seg.B, width_seg.A)
    return width
end

function get_height(poly::ConvexPolygon)
    height_seg = get_edge(poly, 2)
    height = dist(height_seg.B, height_seg.A)
    return height
end

function right_obstacle_space(env::UrbanEnv; delta::Float64=1.0, y_min=env.params.y_min/4)
    x_min = env.params.inter_x + env.params.nlanes*env.params.lane_width + maximum(env.params.crosswalk_width/2)
    x_max = env.params.x_max
    y_max = env.params.inter_y-env.params.nlanes_main*env.params.lane_width
    return x_min, x_max, y_min, y_max
end

function upper_obstacle_space(env::UrbanEnv; max_width = 2.0)
    # x_min = env.params.x_min
    x_min = minimum([pos.x for pos in env.params.crosswalk_pos]) + maximum(env.params.crosswalk_width/2)
    # x_max = env.params.crosswalk_pos.x - env.params.crosswalk_width/2
    x_max = maximum([pos.x for pos in env.params.crosswalk_pos]) - maximum(env.params.crosswalk_width/2)
    y_min = env.params.inter_y + env.params.nlanes_main*env.params.lane_width
    y_max = y_min + max_width
    return x_min, x_max, y_min, y_max
end

function max_left_obstacle(env::UrbanEnv)
    x_min, x_max, y_min, y_max = left_obstacle_space(env)
    obstacle = get_oriented_bounding_box(VecSE2(0.5*(x_max + x_min), 0.5*(y_max + y_min)), x_max - x_min, y_max-y_min)
    return obstacle
end
function max_right_obstacle(env::UrbanEnv)
    x_min, x_max, y_min, y_max = right_obstacle_space(env)
    obstacle = get_oriented_bounding_box(VecSE2(0.5*(x_max + x_min), 0.5*(y_max + y_min)), x_max - x_min, y_max-y_min)
    return obstacle
end
function max_upper_obstacle(env::UrbanEnv)
    x_min, x_max, y_min, y_max = upper_obstacle_space(env)
    obstacle = get_oriented_bounding_box(VecSE2(0.5*(x_max + x_min), 0.5*(y_max + y_min)), x_max - x_min, y_max-y_min)
    return obstacle
end

struct LeftObsDist
    pos::VecSE2
    widths::Vector{Float64}
end
function LeftObsDist(env::UrbanEnv)
    xmin, xmax, ymin, ymax = left_obstacle_space(env)
    x_pos = 0.5*xmin + 0.5*xmax
    obs_pos = VecSE2(x_pos,0.5*ymax+0.5*ymin)
    max_width = xmax - xmin
    widths = [max_width, max_width/2, max_width/4, 3*max_width/4]
    return LeftObsDist(obs_pos, widths)
end
struct RightObsDist
    pos::VecSE2
    widths::Vector{Float64}
end
function RightObsDist(env::UrbanEnv)
    xmin, xmax, ymin, ymax = right_obstacle_space(env)
    x_pos = 0.5*xmin + 0.5*xmax
    obs_pos = VecSE2(x_pos,0.5*ymax+0.5*ymin)
    max_width = xmax - xmin
    widths = [max_width, max_width/2, max_width/4, 3*max_width/4]
    return RightObsDist(obs_pos, widths)
end
struct UpperObsDist
    pos::VecSE2
    widths::Vector{Float64}
end
function UpperObsDist(env::UrbanEnv)
    xmin, xmax, ymin, ymax = upper_obstacle_space(env)
    x_pos = 0.5*xmin + 0.5*xmax
    obs_pos = VecSE2(x_pos,0.5*ymax+0.5*ymin)
    max_width = xmax - xmin
    widths = [max_width, max_width/2, max_width/4, 3*max_width/4]
    return UpperObsDist(obs_pos, widths)
end

"""
structure with obstacle distribution parameters
"""
struct ObstacleDistribution
    upper_obs_pres_prob::Float64
    left_obs_pres_prob::Float64
    right_obs_pres_prob::Float64
    upper_obs_dist::UpperObsDist
    left_obs_dist::LeftObsDist
    right_obs_dist::RightObsDist

    # obs_widths::Vector{Float64}
    # obs_upper::Vector{Float64}
    # obs_left::Vector{Float64}
    # obs_right::Vector{Float64}
end

function ObstacleDistribution(env::UrbanEnv ;
                              upper_obs_pres_prob::Float64 = 0.75,
                              left_obs_pres_prob::Float64 = 0.75,
                              right_obs_pres_prob::Float64 = 0.75)
    left_obs_dist = LeftObsDist(env)
    right_obs_dist = RightObsDist(env)
    upper_obs_dist = UpperObsDist(env)
    return ObstacleDistribution(upper_obs_pres_prob, left_obs_pres_prob, right_obs_pres_prob,
                                upper_obs_dist, left_obs_dist, right_obs_dist)
end

"""
sample obstacle configuration from obs_dist and add them to the environment
"""
function sample_obstacles!(env::UrbanEnv, obs_dist::ObstacleDistribution, rng::AbstractRNG)
    empty_obstacles!(env)
    # sample upper_obstacle
    upper_xmin, upper_xmax, upper_ymin, upper_ymax = upper_obstacle_space(env)
    if rand(rng) < obs_dist.upper_obs_pres_prob
        upper_pos = obs_dist.upper_obs_dist.pos
        upper_width = clamp(rand(rng, obs_dist.upper_obs_dist.widths), 0., 2*min(upper_pos.x - upper_xmin, upper_xmax - upper_pos.x))
        # upper_width = rand(rng, obs_dist.upper_obs_dist.widths)
        add_obstacle!(env, upper_pos , upper_width , upper_ymax - upper_ymin)
    end

    # sample left_obstacle
    left_xmin, left_xmax, left_ymin, left_ymax = left_obstacle_space(env)
    if rand(rng) < obs_dist.left_obs_pres_prob
        left_pos = obs_dist.left_obs_dist.pos
        left_width = clamp(rand(rng, obs_dist.left_obs_dist.widths), 0., 2*min(left_pos.x - left_xmin, left_xmax - left_pos.x))
        add_obstacle!(env, left_pos, left_width , left_ymax - left_ymin)
    end

    # sample right_obstacle
    right_xmin, right_xmax, right_ymin, right_ymax = right_obstacle_space(env)
    if rand(rng) < obs_dist.right_obs_pres_prob
        right_pos = obs_dist.right_obs_dist.pos
        # right_width = rand(rng, obs_dist.right_obs_dist.widths)
        right_width = clamp(rand(rng, obs_dist.right_obs_dist.widths), 0., 2*min(right_pos.x - right_xmin, right_xmax - right_pos.x))
        add_obstacle!(env, right_pos , right_width , right_ymax - right_ymin)
    end
end

"""
sample one obstacle from obs_dist and add it to the environment
"""
function sample_obstacle!(env::UrbanEnv, obs_dist::ObstacleDistribution, rng::AbstractRNG)
    # choose from distribution
    dist_ind = rand(1:3)
    presence = rand(rng)
    if dist_ind == 1 && presence < obs_dist.left_obs_pres_prob
        left_xmin, left_xmax, left_ymin, left_ymax = left_obstacle_space(env)
        left_pos = obs_dist.left_obs_dist.pos
        left_width = clamp(rand(rng, obs_dist.left_obs_dist.widths), 0., 2*min(left_pos.x - left_xmin, left_xmax - left_pos.x))
        add_obstacle!(env, left_pos, left_width , left_ymax - left_ymin)
    elseif dist_ind == 2  && presence < obs_dist.upper_obs_pres_prob
        upper_xmin, upper_xmax, upper_ymin, upper_ymax = upper_obstacle_space(env)
        upper_pos = obs_dist.upper_obs_dist.pos
        upper_width = clamp(rand(rng, obs_dist.upper_obs_dist.widths), 0., 2*min(upper_pos.x - upper_xmin, upper_xmax - upper_pos.x))
        # upper_width = rand(rng, obs_dist.upper_obs_dist.widths)
        add_obstacle!(env, upper_pos , upper_width , upper_ymax - upper_ymin)
    elseif dist_ind == 3 && presence < obs_dist.right_obs_pres_prob
        right_xmin, right_xmax, right_ymin, right_ymax = right_obstacle_space(env)
        right_pos = obs_dist.right_obs_dist.pos
        # right_width = rand(rng, obs_dist.right_obs_dist.widths)
        right_width = clamp(rand(rng, obs_dist.right_obs_dist.widths), 0., 2*min(right_pos.x - right_xmin, right_xmax - right_pos.x))
        add_obstacle!(env, right_pos , right_width , right_ymax - right_ymin)
    end
end


abstract type Model end

struct Gaussian <: Model end

struct Student <: Model end

struct Logistic <: Model end

struct Pois <: Model end

struct NegBin <: Model end


"""
    Ensembles

This module provides the main function [`run_ensemble`](@ref Ensembles.run_ensemble).
This serves to run multiple trajectories for a given simulation type, sampling
from an initial distribution.
"""
module Ensembles

using SciMLBase: SciMLBase, EnsembleProblem, solve
using RecursiveArrayTools: ArrayPartition
using TypedTables: TypedTables

using NQCBase: NQCBase
using NQCDynamics:
    AbstractSimulation,
    Simulation,
    RingPolymerSimulation,
    DynamicsUtils,
    DynamicsMethods,
    DynamicsOutputs
using RingPolymerArrays: RingPolymerArray

export run_ensemble

include("selections.jl")
include("outputs.jl")
include("reductions.jl")

"""
    run_ensemble(sim::AbstractSimulation, tspan, distribution;
        selection::Union{Nothing,AbstractVector}=nothing,
        output=(sol,i)->(sol,false),
        reduction=:append
        reduction::Symbol=:append,
        ensemble_algorithm=SciMLBase.EnsembleThreads(),
        algorithm=DynamicsMethods.select_algorithm(sim),
        saveat=[],
        trajectories=1,
        kwargs...
        )

Run multiple trajectories for timespan `tspan` sampling from `distribution`.
The DifferentialEquations ensemble interface is used which allows us to specify
functions to modify the output and how it is reduced across trajectories.

# Keywords

* `selection` should be an `AbstractVector` containing the indices to sample from the `distribution`. By default, `nothing` leads to random sampling.
* `output` can be a function that transforms the DiffEq solution to an output, or a tuple of output quantities as for `run_trajectory`.
* `reduction` defines how the data is reduced across trajectories. Options are `:append`, `:mean` or `:sum`.
* `ensemble_algorithm` is the algorithm from DifferentialEquations which determines which form of parallelism is used.
* `algorithm` is the algorithm used to integrate the equations of motion.
* `saveat` mirrors the DiffEq keyword and labels the time points to save the output.
* `trajectories` is the number of trajectories to perform.

This function wraps the EnsembleProblem from DifferentialEquations and passes the `kwargs`
to the `solve` function.
"""
function run_ensemble(
    sim::AbstractSimulation, tspan, distribution;
    selection::Union{Nothing,AbstractVector}=nothing,
    output=SciMLBase.DEFAULT_OUTPUT_FUNC,
    reduction::Symbol=:append,
    ensemble_algorithm=SciMLBase.EnsembleThreads(),
    algorithm=DynamicsMethods.select_algorithm(sim),
    saveat=[],
    trajectories=1,
    u_init=nothing,
    kwargs...
)

    kwargs = NQCBase.austrip_kwargs(;kwargs...)
    trajectories = convert(Int, trajectories)
    saveat = austrip.(saveat)
    tspan = austrip.(tspan)
    (output isa Symbol) && (output = (output,))
    reduction = Reduction(reduction)

    ensemble_problem = EnsembleProblem(sim, tspan, distribution, selection, output, reduction, u_init, saveat, trajectories, kwargs)

    @info "Performing $trajectories trajectories."
    stats = @timed solve(ensemble_problem, algorithm, ensemble_algorithm; saveat, trajectories, kwargs...)
    @info "Finished after $(stats.time) seconds."

    return extract_output(stats, output, reduction, trajectories)
end

function SciMLBase.EnsembleProblem(sim::AbstractSimulation, tspan, distribution, selection, output, reduction, u_init, saveat, trajectories, kwargs)

    prob_func = Selection(distribution, selection)

    problem = create_problem(sim, distribution, tspan)
    output_func = Output(output, sim)
    if output_func isa DynamicsOutputs.EnsembleSaver
        u_init = get_u_init(reduction, saveat, kwargs, tspan, problem.u0, output_func)
    end

    return EnsembleProblem(problem; prob_func, output_func, reduction, u_init)
end

Output(output::Tuple, sim) = DynamicsOutputs.EnsembleSaver(output, sim)
Output(output, _) = output

function create_problem(sim::AbstractSimulation, distribution, tspan)
    u0 = sample_distribution(sim, distribution)
    return DynamicsMethods.create_problem(u0, tspan, sim)
end

function extract_output(stats, output, reduction, trajectories)
    sol = stats.value

    if output isa Tuple
        process_output(reduction, sol, output, trajectories)
    else
        return sol.u
    end
end

function process_output(::AppendReduction, sol, output, trajectories)
    return [generate_table(output, s) for s in sol.u]
end

function process_output(::SumReduction, sol, output, trajectories)
    for i in eachindex(sol.u)
        sol.u[i][1] /= trajectories
    end
    return generate_table(output, sol.u)
end

function process_output(::MeanReduction, sol, output, trajectories)
    return generate_table(output, sol.u)
end

function generate_table(output, u)
    output = (:t, output...)
    out = [(;zip(output, quantity)...) for quantity in u]
    return TypedTables.Table(out)
end

end # module


@testset "Element" begin
    # Need some Nodes
    node1 = Node{2}(1,   Vec{2}((2.3, 4.5)), [1,7,9], 3)
    node3 = Node{2}(3,   Vec{2}((2.3, 4.5)), [1,7,9], 3)
    node13 = Node{2}(13, Vec{2}((2.3, 4.5)), [1,7,9], 3)
    node18 = Node{2}(18, Vec{2}((2.3, 4.5)), [1,7,9], 3)

    # Test Full Construction
    element = Element(11, [node1, node3, node13, node18], [12,17,81], 3)

    @test element.id == 11
    @test element.nodes[2].id == 3
    @test element.dofs == [12,17,81]

    # Test DofObject Interface
    @test dofs(element) == [12,17,81]
    @test processor_id(element) == 3
end


# This file is a part of Julia. License is MIT: https://julialang.org/license

## general machinery for irrational mathematical constants

"""
    AbstractIrrational <: Real

Number type representing an exact irrational value.
"""
abstract type AbstractIrrational <: Real end

"""
    Irrational{sym} <: AbstractIrrational

Number type representing an exact irrational value denoted by the
symbol `sym`.
"""
struct Irrational{sym} <: AbstractIrrational end

show(io::IO, x::Irrational{sym}) where {sym} = print(io, sym)

function show(io::IO, ::MIME"text/plain", x::Irrational{sym}) where {sym}
    if get(io, :compact, false)
        print(io, sym)
    else
        print(io, sym, " = ", string(float(x))[1:15], "...")
    end
end

promote_rule(::Type{<:AbstractIrrational}, ::Type{Float16}) = Float16
promote_rule(::Type{<:AbstractIrrational}, ::Type{Float32}) = Float32
promote_rule(::Type{<:AbstractIrrational}, ::Type{<:AbstractIrrational}) = Float64
promote_rule(::Type{<:AbstractIrrational}, ::Type{T}) where {T<:Real} = promote_type(Float64, T)
promote_rule(::Type{S}, ::Type{T}) where {S<:AbstractIrrational,T<:Number} = promote_type(promote_type(S, real(T)), T)

AbstractFloat(x::AbstractIrrational) = Float64(x)
Float16(x::AbstractIrrational) = Float16(Float32(x))
Complex{T}(x::AbstractIrrational) where {T<:Real} = Complex{T}(T(x))

@pure function Rational{T}(x::AbstractIrrational) where T<:Integer
    o = precision(BigFloat)
    p = 256
    while true
        setprecision(BigFloat, p)
        bx = BigFloat(x)
        r = rationalize(T, bx, tol=0)
        if abs(BigFloat(r) - bx) > eps(bx)
            setprecision(BigFloat, o)
            return r
        end
        p += 32
    end
end
(::Type{Rational{BigInt}})(x::AbstractIrrational) = throw(ArgumentError("Cannot convert an AbstractIrrational to a Rational{BigInt}: use rationalize(Rational{BigInt}, x) instead"))

@pure function (t::Type{T})(x::AbstractIrrational, r::RoundingMode) where T<:Union{Float32,Float64}
    setprecision(BigFloat, 256) do
        T(BigFloat(x), r)
    end
end

float(::Type{<:AbstractIrrational}) = Float64

==(::Irrational{s}, ::Irrational{s}) where {s} = true
==(::AbstractIrrational, ::AbstractIrrational) = false

<(::Irrational{s}, ::Irrational{s}) where {s} = false
function <(x::AbstractIrrational, y::AbstractIrrational)
    Float64(x) != Float64(y) || throw(MethodError(<, (x, y)))
    return Float64(x) < Float64(y)
end

<=(::Irrational{s}, ::Irrational{s}) where {s} = true
<=(x::AbstractIrrational, y::AbstractIrrational) = x==y || x<y

# Irrationals, by definition, can't have a finite representation equal them exactly
==(x::AbstractIrrational, y::Real) = false
==(x::Real, y::AbstractIrrational) = false

# Irrational vs AbstractFloat
<(x::AbstractIrrational, y::Float64) = Float64(x,RoundUp) <= y
<(x::Float64, y::AbstractIrrational) = x <= Float64(y,RoundDown)
<(x::AbstractIrrational, y::Float32) = Float32(x,RoundUp) <= y
<(x::Float32, y::AbstractIrrational) = x <= Float32(y,RoundDown)
<(x::AbstractIrrational, y::Float16) = Float32(x,RoundUp) <= y
<(x::Float16, y::AbstractIrrational) = x <= Float32(y,RoundDown)
<(x::AbstractIrrational, y::BigFloat) = setprecision(precision(y)+32) do
    big(x) < y
end
<(x::BigFloat, y::AbstractIrrational) = setprecision(precision(x)+32) do
    x < big(y)
end

<=(x::AbstractIrrational, y::AbstractFloat) = x < y
<=(x::AbstractFloat, y::AbstractIrrational) = x < y

# Irrational vs Rational
@pure function rationalize(::Type{T}, x::AbstractIrrational; tol::Real=0) where T
    return rationalize(T, big(x), tol=tol)
end
@pure function lessrational(rx::Rational{<:Integer}, x::AbstractIrrational)
    # an @pure version of `<` for determining if the rationalization of
    # an irrational number required rounding up or down
    return rx < big(x)
end
function <(x::AbstractIrrational, y::Rational{T}) where T
    T <: Unsigned && x < 0.0 && return true
    rx = rationalize(T, x)
    if lessrational(rx, x)
        return rx < y
    else
        return rx <= y
    end
end
function <(x::Rational{T}, y::AbstractIrrational) where T
    T <: Unsigned && y < 0.0 && return false
    ry = rationalize(T, y)
    if lessrational(ry, y)
        return x <= ry
    else
        return x < ry
    end
end
<(x::AbstractIrrational, y::Rational{BigInt}) = big(x) < y
<(x::Rational{BigInt}, y::AbstractIrrational) = x < big(y)

<=(x::AbstractIrrational, y::Rational) = x < y
<=(x::Rational, y::AbstractIrrational) = x < y

isfinite(::AbstractIrrational) = true
isinteger(::AbstractIrrational) = false
iszero(::AbstractIrrational) = false
isone(::AbstractIrrational) = false

hash(x::Irrational, h::UInt) = 3*objectid(x) - h

widen(::Type{T}) where {T<:Irrational} = T

-(x::AbstractIrrational) = -Float64(x)
for op in Symbol[:+, :-, :*, :/, :^]
    @eval $op(x::AbstractIrrational, y::AbstractIrrational) = $op(Float64(x),Float64(y))
end
*(x::Bool, y::AbstractIrrational) = ifelse(x, Float64(y), 0.0)

round(x::Irrational, r::RoundingMode) = round(float(x), r)

"""
	@irrational sym val def
	@irrational(sym, val, def)

Define a new `Irrational` value, `sym`, with pre-computed `Float64` value `val`,
and arbitrary-precision definition in terms of `BigFloat`s given be the expression `def`.
"""
macro irrational(sym, val, def)
    esym = esc(sym)
    qsym = esc(Expr(:quote, sym))
    bigconvert = isa(def,Symbol) ? quote
        function Base.BigFloat(::Irrational{$qsym}, r::MPFR.MPFRRoundingMode=MPFR.ROUNDING_MODE[]; precision=precision(BigFloat))
            c = BigFloat(;precision=precision)
            ccall(($(string("mpfr_const_", def)), :libmpfr),
                  Cint, (Ref{BigFloat}, MPFR.MPFRRoundingMode), c, r)
            return c
        end
    end : quote
        function Base.BigFloat(::Irrational{$qsym}; precision=precision(BigFloat))
            setprecision(BigFloat, precision) do
                $(esc(def))
            end
        end
    end
    quote
        const $esym = Irrational{$qsym}()
        $bigconvert
        Base.Float64(::Irrational{$qsym}) = $val
        Base.Float32(::Irrational{$qsym}) = $(Float32(val))
        @assert isa(big($esym), BigFloat)
        @assert Float64($esym) == Float64(big($esym))
        @assert Float32($esym) == Float32(big($esym))
    end
end

big(x::AbstractIrrational) = BigFloat(x)
big(::Type{<:AbstractIrrational}) = BigFloat

# align along = for nice Array printing
function alignment(io::IO, x::AbstractIrrational)
    m = match(r"^(.*?)(=.*)$", sprint(show, x, context=io, sizehint=0))
    m === nothing ? (length(sprint(show, x, context=io, sizehint=0)), 0) :
    (length(m.captures[1]), length(m.captures[2]))
end

# inv
inv(x::AbstractIrrational) = 1/x


Base.setindex!(v::Array{Complex{ArbFloat{P}},1}, c::ArbComplex{P}, idx::Int64) where {P} =
         setindex!(v, Complex{ArbFloat{P}}(ArbFloat{P}(real(c)), ArbFloat{P}(imag(c))), idx)

Base.setindex!(v::Array{Complex{ArbReal{P}},1}, c::ArbComplex{P}, idx::Int64) where {P} =
    setindex!(v, Complex{ArbReal{P}}(real(c), imag(c)), idx)

Base.convert(::ArbFloatMatrix{P}, x::Array{ArbFloat{P}, 2}) where {P} = ArbFloatMatrix(x)
Base.convert(::ArbRealMatrix{P}, x::Array{ArbReal{P}, 2})  where {P} = ArbRealMatrix(x)
Base.convert(::ArbComplexMatrix{P}, x::Array{ArbComplex{P}, 2}) where {P} = ArbComplexMatrix(x)

svdvals(x::Array{ArbComplex{N},2}) where {N} = throw(ErrorException("not supported"))
svdvals(x::Array{ArbComplex{N},M}) where {N,M} = throw(ErrorException("not supported"))
svdvals(x::Array{Complex{ArbFloat{N}},2}) where {N} = throw(ErrorException("not supported"))
svdvals(x::Array{Complex{ArbReal{N}},2}) where {N} = throw(ErrorException("not supported"))
svdvals(x::Array{ArbComplex,2}) = throw(ErrorException("not supported"))
svdvals(x::Array{ArbComplex,M}) where {M} = throw(ErrorException("not supported"))
svdvals(x::Array{Complex{ArbFloat},2}) = throw(ErrorException("not supported"))
svdvals(x::Array{Complex{ArbReal},2}) = throw(ErrorException("not supported"))

eigvals(x::Array{ArbReal{N},2}) where {N} = ArbReal{N}.(eigvals(ArbFloat{N}.(midpoint.(x))))
eigvals(x::Array{ArbReal,2})  = ArbReal{workingprecision{ArbReal}}.(eigvals(ArbFloat{workingprecision(ArbFloat)}.(midpoint.(x))))

eigvals(x::Array{ArbComplex{N},2}) where {N} = throw(ErrorException("not supported"))
eigvals(x::Array{Complex{ArbFloat{N}},2}) where {N} = throw(ErrorException("not supported"))
eigvals(x::Array{Complex{ArbReal{N}},2}) where {N} = throw(ErrorException("not supported"))

eigvecs(x::Array{ArbFloat{N},2}) where {N} = throw(ErrorException("not supported"))
eigvecs(x::Array{ArbReal{N},2}) where {N} = throw(ErrorException("not supported"))
eigvecs(x::Array{ArbComplex{N},2}) where {N} = throw(ErrorException("not supported"))
eigvecs(x::Array{Complex{ArbFloat{N}},2}) where {N} = throw(ErrorException("not supported"))
eigvecs(x::Array{Complex{ArbReal{N}},2}) where {N} = throw(ErrorException("not supported"))
eigvecs(x::Array{ArbFloat{N},M}) where {N,M} = throw(ErrorException("not supported"))
eigvecs(x::Array{ArbReal{N},M}) where {N,M} = throw(ErrorException("not supported"))
eigvecs(x::Array{ArbComplex{N},M}) where {N,M} = throw(ErrorException("not supported"))

eigvecs(x::Array{ArbFloat,2}) = throw(ErrorException("not supported"))
eigvecs(x::Array{ArbReal,2}) = throw(ErrorException("not supported"))
eigvecs(x::Array{ArbComplex,2}) = throw(ErrorException("not supported"))
eigvecs(x::Array{Complex{ArbFloat},2}) = throw(ErrorException("not supported"))
eigvecs(x::Array{Complex{ArbReal},2}) = throw(ErrorException("not supported"))
eigvecs(x::Array{ArbFloat,M}) where {M} = throw(ErrorException("not supported"))
eigvecs(x::Array{ArbReal,M}) where {M} = throw(ErrorException("not supported"))
eigvecs(x::Array{ArbComplex,M}) where {M} = throw(ErrorException("not supported"))

eigen(x::Array{ArbFloat{N},2}) where {N} = throw(ErrorException("not supported"))
eigen(x::Array{ArbReal{N},2}) where {N} = throw(ErrorException("not supported"))
eigen(x::Array{ArbComplex{N},2}) where {N} = throw(ErrorException("not supported"))
eigen(x::Array{Complex{ArbFloat{N}},2}) where {N} = throw(ErrorException("not supported"))
eigen(x::Array{Complex{ArbReal{N}},2}) where {N} = throw(ErrorException("not supported"))
eigen(x::Array{ArbFloat{N},M}) where {N,M} = throw(ErrorException("not supported"))
eigen(x::Array{ArbReal{N},M}) where {N,M} = throw(ErrorException("not supported"))
eigen(x::Array{ArbComplex{N},M}) where {N,M} = throw(ErrorException("not supported"))

eigen(x::Array{ArbFloat,2}) = throw(ErrorException("not supported"))
eigen(x::Array{ArbReal,2}) = throw(ErrorException("not supported"))
eigen(x::Array{ArbComplex,2}) = throw(ErrorException("not supported"))
eigen(x::Array{Complex{ArbFloat},2}) = throw(ErrorException("not supported"))
eigen(x::Array{Complex{ArbReal},2}) = throw(ErrorException("not supported"))
eigen(x::Array{ArbFloat,M}) where {M} = throw(ErrorException("not supported"))
eigen(x::Array{ArbReal,M}) where {M} = throw(ErrorException("not supported"))
eigen(x::Array{ArbComplex,M}) where {M} = throw(ErrorException("not supported"))


@everywhere push!(LOAD_PATH, "../")
@everywhere push!(LOAD_PATH, "./")

include("sawyerfunctions.jl")

####################################### Experiment Parameters / HyperParameters

model_file = "$(modeldir)/sawyer.xml"

# simulator parameters
T = 250
numT = 120
skip = 1
niter = 100

# solver parameters
gamma = 0.995
gae = 0.98

fullFIM = true

norm_step_size = 0.05

cg_iter= 12
cg_reg = 1e-6
cg_tol = 1e-10

#### deterministic randomization
srand(12345) # random seed, set globally

dtype = Float64 # ALWAYS DO FLOAT64

#         qpos qvel
const ns = 13 + 13 + 4 + 3 #3 + 3
my_mjsys       = mjw.load_model(model_file, skip, "normal", ns)

if METHOD == :NPG ######################################### NPG startup
   #my_policy      = Policy.GLP{dtype}(my_mjsys.ns, my_mjsys.nu) # inputs: n, m
   my_policy      = Policy.NN{dtype}(my_mjsys.ns, my_mjsys.nu, 32, 2) # inputs: n, m
   ls = Policy.getls(my_policy)
   ls .= -1.0

   pgmodel        = PolicyGradModel(copy(my_policy.theta),
                                    my_mjsys,
                                    my_policy,
                                    T, numT, 10, niter)

   # NPG baseline aka value function approximation
   baseline    = Baseline.Quadratic{Float64}(1e-5, my_mjsys.ns, T, numT)

   pg_specs   = NPGStrategy{dtype}(pgmodel,
                                     baseline,
                                     fullFIM,
                                     norm_step_size,
                                     gamma, gae,
                                     cg_iter, cg_reg, cg_tol)

elseif METHOD == :MPPI ######################################### MPPI startup

   rand(220)
   mpcT = 100

   mppi = MPPIModel(eye(my_mjsys.nu)*0.3, # sigma
                    0.2,                  # lambda
                    my_mjsys,
                    50, # horizon for mpc; set to T for non-MPC mode
                    T,
                    numT)

   mppi_specs = MPPIStrategy(mppi)
else
   error("WHAT YOU DOING")
end


"""Implements the CreditMetrics model for portfolio credit risk"""
module CreditMetrics

using StatsFuns

include("threshold_links.jl")
include("thresholds.jl")

end # module


__precompile__(true)
module Coord

using Reexport
@reexport using GridInterpolations
@reexport using ODE
@reexport using JLD

include("tables.jl")
include("state.jl")
include("dynamics.jl")
include("uav.jl")
include("encounter.jl")
include("strategies.jl")
include("environment.jl")
include("build_simulation.jl")
include("run_simulation.jl")

end # module


# routines related to the PETSc backend

export initialize_backend, petsc_data

"""
    Stores key PETSc information with the `grid` structure
"""
struct petsc_data
    petsclib
    da
end
petsc_data() = petsc_data(nothing,nothing)

"""
    check_backend(backend{BackendPETSc}; Scalar=Float64, dim=1)

checks PETSc backend for the `Scalar`, which loads PETSc (and, if requested, MPI)
"""
function check_backend(b::Backend{BackendPETSc}; dim=1, Scalar=Float64)

    if !isdefined(Main, :PETSc)
        error("PETSc is not loaded; ensure it is loaded first with: using PETSc")
    end
    if (b.mpi==true && !isdefined(Main, :MPI))
        error("MPI is not loaded; ensure it is loaded first with: using MPI")
    end
    @eval using PETSc

    # get the PETSc lib with our chosen `PetscScalar` type
    petsclib = PETSc.getlib(; PetscScalar = Scalar)

    # Initialize PETSc
    PETSc.initialize(petsclib)

    return petsclib
end

# Internal function that translates the CompGrids topology to PETSc nomenclature
function bcs_translate(bound)
    bound_petsc = PETSc.LibPETSc.DM_BOUNDARY_NONE
    if bound == Ghosted
        bound_petsc = PETSc.LibPETSc.DM_BOUNDARY_GHOSTED
    elseif bound == Periodic
        bound_petsc = PETSc.DM_BOUNDARY_PERIODIC
    end
    return bound_petsc
end

function globalfromlocalsize(size, localsize, opts, stencilwidth, topology, b::Backend{BackendPETSc, FT}) where FT

    if !b.mpi & !isempty(localsize)
        size = localsize   # not using MPI
    elseif b.mpi & !isempty(localsize)
        error("The PETSc backend currently doesn't allow specifying the local grid size")
    else
        size=size
    end

    return size
end