text:
# Variables and constraints related to flexible loads

## Variables

function variable_flexible_demand(pm::_PM.AbstractPowerModel; kwargs...)
    variable_total_flex_demand(pm; kwargs...)
    variable_demand_reduction(pm; kwargs...)
    variable_demand_shifting_upwards(pm; kwargs...)
    variable_demand_shifting_downwards(pm; kwargs...)
    variable_demand_curtailment(pm; kwargs...)
    variable_flexible_demand_indicator(pm; kwargs..., relax=true)
    variable_flexible_demand_investment(pm; kwargs...)
end

"Variable: whether flexible demand is enabled at a flex load point"
function variable_flexible_demand_indicator(pm::_PM.AbstractPowerModel; nw::Int=_PM.nw_id_default, relax::Bool=false, report::Bool=true)
    first_n = first_id(pm, nw, :hour, :scenario)
    if nw == first_n
        if !relax
            z = _PM.var(pm, nw)[:z_flex] = JuMP.@variable(pm.model,
                [i in _PM.ids(pm, nw, :flex_load)], base_name="$(nw)_z_flex",
                binary = true,
                start = 0
            )
        else
            z = _PM.var(pm, nw)[:z_flex] = JuMP.@variable(pm.model,
                [i in _PM.ids(pm, nw, :flex_load)], base_name="$(nw)_z_flex",
                lower_bound = 0,
                upper_bound = 1,
                start = 0
            )
        end
    else
        z = _PM.var(pm, nw)[:z_flex] = _PM.var(pm, first_n)[:z_flex]
    end
    if report
        _PM.sol_component_value(pm, nw, :load, :flex, _PM.ids(pm, nw, :flex_load), z)
        _IM.sol_component_fixed(pm, _PM.pm_it_sym, nw, :load, :flex, _PM.ids(pm, nw, :fixed_load), 0.0)
    end
end

"Variable: investment decision to enable flexible demand at a flex load point"
function variable_flexible_demand_investment(pm::_PM.AbstractPowerModel; nw::Int=_PM.nw_id_default, relax::Bool=false, report::Bool=true)
    first_n = first_id(pm, nw, :hour, :scenario)
    if nw == first_n
        if !relax
            investment = _PM.var(pm, nw)[:z_flex_investment] = JuMP.@variable(pm.model,
                [i in _PM.ids(pm, nw, :flex_load)], base_name="$(nw)_z_flex_investment",
                binary = true,
                start = 0
            )
        else
            investment = _PM.var(pm, nw)[:z_flex_investment] = JuMP.@variable(pm.model,
                [i in _PM.ids(pm, nw, :flex_load)], base_name="$(nw)_z_flex_investment",
                lower_bound = 0,
                upper_bound = 1,
                start = 0
            )
        end
    else
        investment = _PM.var(pm, nw)[:z_flex_investment] = _PM.var(pm, first_n)[:z_flex_investment]
    end
    if report
        _PM.sol_component_value(pm, nw, :load, :investment, _PM.ids(pm, nw, :flex_load), investment)
        _IM.sol_component_fixed(pm, _PM.pm_it_sym, nw, :load, :investment, _PM.ids(pm, nw, :fixed_load), 0.0)
    end
end

function variable_total_flex_demand(pm::_PM.AbstractPowerModel; kwargs...)
    variable_total_flex_demand_active(pm; kwargs...)
    variable_total_flex_demand_reactive(pm; kwargs...)
end

"Variable for the actual (flexible) real load demand at each load point and each time step"
function variable_total_flex_demand_active(pm::_PM.AbstractPowerModel; nw::Int=_PM.nw_id_default, bounded::Bool=true, report::Bool=true)
    pflex = _PM.var(pm, nw)[:pflex] = JuMP.@variable(pm.model,
        [i in _PM.ids(pm, nw, :load)], base_name="$(nw)_pflex",
        lower_bound = 0,
        upper_bound = _PM.ref(pm, nw, :load, i, "pd") * (1 + get(_PM.ref(pm, nw, :load, i), "pshift_up_rel_max", 0.0)), # Not strictly necessary: redundant due to other bounds
        start = _PM.comp_start_value(_PM.ref(pm, nw, :load, i), "pd")
    )
    report && _PM.sol_component_value(pm, nw, :load, :pflex, _PM.ids(pm, nw, :load), pflex)
end

function variable_total_flex_demand_reactive(pm::_PM.AbstractActivePowerModel; nw::Int=_PM.nw_id_default, bounded::Bool=true, report::Bool=true)
end

"Variable for the actual (flexible) reactive load demand at each load point and each time step"
function variable_total_flex_demand_reactive(pm::_PM.AbstractPowerModel; nw::Int=_PM.nw_id_default, bounded::Bool=true, report::Bool=true)
    qflex = _PM.var(pm, nw)[:qflex] = JuMP.@variable(pm.model,
        [i in _PM.ids(pm, nw, :load)], base_name="$(nw)_qflex",
        start = _PM.comp_start_value(_PM.ref(pm, nw, :load, i), "qd")
    )
    report && _PM.sol_component_value(pm, nw, :load, :qflex, _PM.ids(pm, nw, :load), qflex)
end

"Variable for load curtailment (i.e. involuntary demand reduction) at each load point and each time step"
function variable_demand_curtailment(pm::_PM.AbstractPowerModel; nw::Int=_PM.nw_id_default, bounded::Bool=true, report::Bool=true)
    pcurt = _PM.var(pm, nw)[:pcurt] = JuMP.@variable(pm.model,
        [i in _PM.ids(pm, nw, :load)], base_name="$(nw)_pcurt",
        lower_bound = 0,
        upper_bound = _PM.ref(pm, nw, :load, i, "pd"),
        start = 0
    )
    report && _PM.sol_component_value(pm, nw, :load, :pcurt, _PM.ids(pm, nw, :load), pcurt)
end

"Variable for the power not consumed (voluntary load reduction) at each flex load point and each time step"
function variable_demand_reduction(pm::_PM.AbstractPowerModel; nw::Int=_PM.nw_id_default, bounded::Bool=true, report::Bool=true)
    # This is bounded for each time step by a fixed share (0 ≤ pred_rel_max ≤ 1) of the
    # reference load demand pd for that time step. (Thus, while pred_rel_max is a scalar
    # input parameter, the variable bounds become a time series.)
    pred = _PM.var(pm, nw)[:pred] = JuMP.@variable(pm.model,
        [i in _PM.ids(pm, nw, :flex_load)], base_name="$(nw)_pred",
        lower_bound = 0,
        upper_bound = _PM.ref(pm, nw, :load, i, "pd") * _PM.ref(pm, nw, :flex_load, i, "pred_rel_max"),
        start = 0
    )
    if report
        _PM.sol_component_value(pm, nw, :load, :pred, _PM.ids(pm, nw, :flex_load), pred)
        _IM.sol_component_fixed(pm, _PM.pm_it_sym, nw, :load, :pred, _PM.ids(pm, nw, :fixed_load), 0.0)
    end
end

"Variable for keeping track of the energy not consumed (i.e. the accumulated voluntary load reduction) over the operational planning horizon at each flex load point"
function variable_energy_not_consumed(pm::_PM.AbstractPowerModel; nw::Int=_PM.nw_id_default, bounded::Bool=true, report::Bool=true)
    first_nw = first_id(pm, nw, :hour)
    ered = _PM.var(pm, nw)[:ered] = JuMP.@variable(pm.model,
        [i in _PM.ids(pm, nw, :flex_load)], base_name="$(nw)_ered",
        lower_bound = 0,
        upper_bound = _PM.ref(pm, nw, :flex_load, i, "ered_rel_max") * _PM.ref(pm, first_nw, :flex_load, i, "ed"),
        start = 0
    )
    if report
        _PM.sol_component_value(pm, nw, :load, :ered, _PM.ids(pm, nw, :flex_load), ered)
        _IM.sol_component_fixed(pm, _PM.pm_it_sym, nw, :load, :ered, _PM.ids(pm, nw, :fixed_load), 0.0)
    end
end

"Variable for the upward demand shifting at each flex load point and each time step"
function variable_demand_shifting_upwards(pm::_PM.AbstractPowerModel; nw::Int=_PM.nw_id_default, bounded::Bool=true, report::Bool=true)
    # This is bounded for each time step by a fixed share (0 ≤ pshift_up_rel_max ≤ 1) of the
    # reference load demand pd for that time step. (Thus, while pshift_up_rel_max is a
    # scalar input parameter, the variable bounds become a time series.)
    pshift_up = _PM.var(pm, nw)[:pshift_up] = JuMP.@variable(pm.model,
        [i in _PM.ids(pm, nw, :flex_load)], base_name="$(nw)_pshift_up",
        lower_bound = 0,
        upper_bound = _PM.ref(pm, nw, :load, i, "pd") * _PM.ref(pm, nw, :flex_load, i, "pshift_up_rel_max"),
        start = 0
    )
    if report
        _PM.sol_component_value(pm, nw, :load, :pshift_up, _PM.ids(pm, nw, :flex_load), pshift_up)
        _IM.sol_component_fixed(pm, _PM.pm_it_sym, nw, :load, :pshift_up, _PM.ids(pm, nw, :fixed_load), 0.0)
    end
end

"Variable for keeping track of the accumulated upward demand shifting over the operational planning horizon at each flex_load point"
function variable_total_demand_shifting_upwards(pm::_PM.AbstractPowerModel; nw::Int=_PM.nw_id_default, bounded::Bool=true, report::Bool=true)
    first_nw = first_id(pm, nw, :hour)
    eshift_up = _PM.var(pm, nw)[:eshift_up] = JuMP.@variable(pm.model,
        [i in _PM.ids(pm, nw, :flex_load)], base_name="$(nw)_eshift_up",
        lower_bound = 0,
        # The accumulated load shifted up should equal the accumulated load shifted down, so this constraint is probably redundant
        upper_bound = _PM.ref(pm, nw, :flex_load, i, "eshift_rel_max") * _PM.ref(pm, first_nw, :flex_load, i, "ed"),
        start = 0
    )
    if report
        _PM.sol_component_value(pm, nw, :load, :eshift_up, _PM.ids(pm, nw, :flex_load), eshift_up)
        _IM.sol_component_fixed(pm, _PM.pm_it_sym, nw, :load, :eshift_up, _PM.ids(pm, nw, :fixed_load), 0.0)
    end
end

"Variable for the downward demand shifting at each flex load point and each time step"
function variable_demand_shifting_downwards(pm::_PM.AbstractPowerModel; nw::Int=_PM.nw_id_default, bounded::Bool=true, report::Bool=true)
    # This is bounded for each time step by a fixed share (0 ≤ pshift_down_rel_max ≤ 1) of
    # the reference load demand pd for that time step. (Thus, while pshift_down_rel_max is a
    # scalar input parameter, the variable bounds become a time series.)
    pshift_down = _PM.var(pm, nw)[:pshift_down] = JuMP.@variable(pm.model,
        [i in _PM.ids(pm, nw, :flex_load)], base_name="$(nw)_pshift_down",
        lower_bound = 0,
        upper_bound = _PM.ref(pm, nw, :load, i, "pd") * _PM.ref(pm, nw, :flex_load, i, "pshift_down_rel_max"),
        start = 0
    )
    if report
        _PM.sol_component_value(pm, nw, :load, :pshift_down, _PM.ids(pm, nw, :flex_load), pshift_down)
        _IM.sol_component_fixed(pm, _PM.pm_it_sym, nw, :load, :pshift_down, _PM.ids(pm, nw, :fixed_load), 0.0)
    end
end

"Variable for keeping track of the accumulated upward demand shifting over the operational planning horizon at each flex load point"
function variable_total_demand_shifting_downwards(pm::_PM.AbstractPowerModel; nw::Int=_PM.nw_id_default, bounded::Bool=true, report::Bool=true)
    first_nw = first_id(pm, nw, :hour)
    eshift_down = _PM.var(pm, nw)[:eshift_down] = JuMP.@variable(pm.model,
        [i in _PM.ids(pm, nw, :flex_load)], base_name="$(nw)_eshift_down",
        lower_bound = 0,
        upper_bound = _PM.ref(pm, nw, :flex_load, i, "eshift_rel_max") * _PM.ref(pm, first_nw, :flex_load, i, "ed"),
        start = 0
    )
    if report
        _PM.sol_component_value(pm, nw, :load, :eshift_down, _PM.ids(pm, nw, :flex_load), eshift_down)
        _IM.sol_component_fixed(pm, _PM.pm_it_sym, nw, :load, :eshift_down, _PM.ids(pm, nw, :fixed_load), 0.0)
    end
end

## Constraint templates

function constraint_flexible_demand_activation(pm::_PM.AbstractPowerModel, i::Int, prev_nws::Vector{Int}, nw::Int)
    investment_horizon = [nw]
    lifetime = _PM.ref(pm, nw, :load, i, "lifetime")
    for n in Iterators.reverse(prev_nws[max(end-lifetime+2,1):end])
        i in _PM.ids(pm, n, :load) ? push!(investment_horizon, n) : break
    end
    constraint_flexible_demand_activation(pm, nw, i, investment_horizon)
end

function constraint_flex_bounds_ne(pm::_PM.AbstractPowerModel, i::Int; nw::Int=_PM.nw_id_default)
    load = _PM.ref(pm, nw, :load, i)
    constraint_flex_bounds_ne(pm, nw, i, load["pd"], load["pshift_up_rel_max"], load["pshift_down_rel_max"], load["pred_rel_max"])
end

function constraint_total_flexible_demand(pm::_PM.AbstractPowerModel, i::Int; nw::Int=_PM.nw_id_default)
    load     = _PM.ref(pm, nw, :load, i)
    pd       = load["pd"]
    pf_angle = get(load, "pf_angle", 0.0) # Power factor angle, in radians
    constraint_total_flexible_demand(pm, nw, i, pd, pf_angle)
end

function constraint_total_fixed_demand(pm::_PM.AbstractPowerModel, i::Int; nw::Int=_PM.nw_id_default)
    load     = _PM.ref(pm, nw, :load, i)
    pd       = load["pd"]
    pf_angle = get(load, "pf_angle", 0.0) # Power factor angle, in radians
    constraint_total_fixed_demand(pm, nw, i, pd, pf_angle)
end

function constraint_red_state(pm::_PM.AbstractPowerModel, i::Int; nw::Int=_PM.nw_id_default)
    if haskey(_PM.ref(pm, nw), :time_elapsed)
        time_elapsed = _PM.ref(pm, nw, :time_elapsed)
    else
        Memento.warn(_LOGGER, "network data should specify time_elapsed, using 1.0 as a default")
        time_elapsed = 1.0
    end
    constraint_red_state_initial(pm, nw, i, time_elapsed)
end

function constraint_red_state(pm::_PM.AbstractPowerModel, i::Int, nw_1::Int, nw_2::Int)
    if haskey(_PM.ref(pm, nw_2), :time_elapsed)
        time_elapsed = _PM.ref(pm, nw_2, :time_elapsed)
    else
        Memento.warn(_LOGGER, "network $(nw_2) should specify time_elapsed, using 1.0 as a default")
        time_elapsed = 1.0
    end
    constraint_red_state(pm, nw_1, nw_2, i, time_elapsed)
end

function constraint_shift_duration(pm::_PM.AbstractPowerModel, i::Int, first_hour_nw::Int, nw::Int)
    constraint_shift_duration_up(pm, i, first_hour_nw, nw)
    constraint_shift_duration_down(pm, i, first_hour_nw, nw)
end

function constraint_shift_duration_up(pm::_PM.AbstractPowerModel, i::Int, first_hour_nw::Int, nw::Int)
    load         = _PM.ref(pm, nw, :load, i)
    start_period = max(nw-load["tshift_up"], first_hour_nw)
    constraint_shift_duration_up(pm, nw, i, load["pd"], load["pshift_up_rel_max"], start_period)
end

function constraint_shift_duration_down(pm::_PM.AbstractPowerModel, i::Int, first_hour_nw::Int, nw::Int)
    load         = _PM.ref(pm, nw, :load, i)
    start_period = max(nw-load["tshift_down"], first_hour_nw)
    constraint_shift_duration_down(pm, nw, i, load["pd"], load["pshift_down_rel_max"], start_period)
end

function constraint_shift_up_state(pm::_PM.AbstractPowerModel, i::Int; nw::Int=_PM.nw_id_default)
    if haskey(_PM.ref(pm, nw), :time_elapsed)
        time_elapsed = _PM.ref(pm, nw, :time_elapsed)
    else
        Memento.warn(_LOGGER, "network data should specify time_elapsed, using 1.0 as a default")
        time_elapsed = 1.0
    end
    constraint_shift_up_state_initial(pm, nw, i, time_elapsed)
end

function constraint_shift_up_state(pm::_PM.AbstractPowerModel, i::Int, nw_1::Int, nw_2::Int)
    if haskey(_PM.ref(pm, nw_2), :time_elapsed)
        time_elapsed = _PM.ref(pm, nw_2, :time_elapsed)
    else
        Memento.warn(_LOGGER, "network $(nw_2) should specify time_elapsed, using 1.0 as a default")
        time_elapsed = 1.0
    end
    constraint_shift_up_state(pm, nw_1, nw_2, i, time_elapsed)
end

function constraint_shift_down_state(pm::_PM.AbstractPowerModel, i::Int; nw::Int=_PM.nw_id_default)
    if haskey(_PM.ref(pm, nw), :time_elapsed)
        time_elapsed = _PM.ref(pm, nw, :time_elapsed)
    else
        Memento.warn(_LOGGER, "network data should specify time_elapsed, using 1.0 as a default")
        time_elapsed = 1.0
    end
    constraint_shift_down_state_initial(pm, nw, i, time_elapsed)
end

function constraint_shift_down_state(pm::_PM.AbstractPowerModel, i::Int, nw_1::Int, nw_2::Int)
    if haskey(_PM.ref(pm, nw_2), :time_elapsed)
        time_elapsed = _PM.ref(pm, nw_2, :time_elapsed)
    else
        Memento.warn(_LOGGER, "network $(nw_2) should specify time_elapsed, using 1.0 as a default")
        time_elapsed = 1.0
    end
    constraint_shift_down_state(pm, nw_1, nw_2, i, time_elapsed)
end

function constraint_shift_state_final(pm::_PM.AbstractPowerModel, i::Int; nw::Int=_PM.nw_id_default)
    constraint_shift_state_final(pm, nw, i)
end

# This way of enforcing a balance between power shifted upward and power shifted downward:
# - does not use `eshift_up` and `eshift_down` variables;
# - is alternative to `constraint_shift_up_state`, `constraint_shift_down_state`, and
#   `constraint_shift_state_final`.
# It must be called only on last hour nws.
function constraint_shift_balance_periodic(pm::_PM.AbstractPowerModel, i::Int, period::Int; nw::Int=_PM.nw_id_default)
    timeseries_nw_ids = similar_ids(pm, nw, hour = 1:dim_length(pm,:hour))
    time_elapsed = Int(_PM.ref(pm, nw, :time_elapsed))
    if period % time_elapsed ≠ 0
        Memento.error(_LOGGER, "\"period\" must be a multiple of \"time_elapsed\".")
    end
    for horizon in Iterators.partition(timeseries_nw_ids, period÷time_elapsed)
        constraint_shift_balance_periodic(pm, horizon, i)
    end
end

## Constraint implementations

function constraint_flexible_demand_activation(pm::_PM.AbstractPowerModel, n::Int, i::Int, horizon::Vector{Int})
    indicator = _PM.var(pm, n, :z_flex, i)
    investments = _PM.var.(Ref(pm), horizon, :z_flex_investment, i)

    # Activate the flexibility depending on the investment decisions in the load's horizon.
    JuMP.@constraint(pm.model, indicator == sum(investments))
end

function constraint_flex_bounds_ne(pm::_PM.AbstractPowerModel, n::Int, i::Int, pd, pshift_up_rel_max, pshift_down_rel_max, pred_rel_max)
    pshift_up   = _PM.var(pm, n, :pshift_up, i)
    pshift_down = _PM.var(pm, n, :pshift_down, i)
    pred        = _PM.var(pm, n, :pred, i)
    z           = _PM.var(pm, n, :z_flex, i)

    # Bounds on the demand flexibility decision variables (demand shifting and voluntary load reduction)
    JuMP.@constraint(pm.model, pshift_up   <=   pshift_up_rel_max * pd * z)
    JuMP.@constraint(pm.model, pshift_down <= pshift_down_rel_max * pd * z)
    JuMP.@constraint(pm.model, pred        <=        pred_rel_max * pd * z)
end

function constraint_total_flexible_demand(pm::_PM.AbstractPowerModel, n::Int, i, pd, pf_angle)
    pflex       = _PM.var(pm, n, :pflex, i)
    qflex       = _PM.var(pm, n, :qflex, i)
    pcurt       = _PM.var(pm, n, :pcurt, i)
    pred        = _PM.var(pm, n, :pred, i)
    pshift_up   = _PM.var(pm, n, :pshift_up, i)
    pshift_down = _PM.var(pm, n, :pshift_down, i)

    # Active power demand is the reference demand `pd` plus the contributions from all the demand flexibility decision variables
    JuMP.@constraint(pm.model, pflex == pd - pcurt - pred + pshift_up - pshift_down)

    # Reactive power demand is given by the active power demand and the power factor angle of the load
    JuMP.@constraint(pm.model, qflex == tan(pf_angle) * pflex)
end

function constraint_total_flexible_demand(pm::_PM.AbstractActivePowerModel, n::Int, i, pd, pf_angle)
    pflex       = _PM.var(pm, n, :pflex, i)
    pcurt       = _PM.var(pm, n, :pcurt, i)
    pred        = _PM.var(pm, n, :pred, i)
    pshift_up   = _PM.var(pm, n, :pshift_up, i)
    pshift_down = _PM.var(pm, n, :pshift_down, i)

    # Active power demand is the reference demand `pd` plus the contributions from all the demand flexibility decision variables
    JuMP.@constraint(pm.model, pflex == pd - pcurt - pred + pshift_up - pshift_down)
end

function constraint_total_fixed_demand(pm::_PM.AbstractPowerModel, n::Int, i, pd, pf_angle)
    pflex = _PM.var(pm, n, :pflex, i)
    qflex = _PM.var(pm, n, :qflex, i)
    pcurt = _PM.var(pm, n, :pcurt, i)

    # Active power demand is the difference between reference demand `pd` and involuntary curtailment
    JuMP.@constraint(pm.model, pflex == pd - pcurt)

    # Reactive power demand is given by the active power demand and the power factor angle of the load
    JuMP.@constraint(pm.model, qflex == tan(pf_angle) * pflex)
end

function constraint_total_fixed_demand(pm::_PM.AbstractActivePowerModel, n::Int, i, pd, pf_angle)
    pflex = _PM.var(pm, n, :pflex, i)
    pcurt = _PM.var(pm, n, :pcurt, i)

    # Active power demand is the difference between reference demand `pd` and involuntary curtailment
    JuMP.@constraint(pm.model, pflex == pd - pcurt)
end

function constraint_red_state_initial(pm::_PM.AbstractPowerModel, n::Int, i::Int, time_elapsed)
    pred = _PM.var(pm, n, :pred, i)
    ered = _PM.var(pm, n, :ered, i)

    # Initialization of not consumed energy variable (accumulated voluntary load reduction)
    JuMP.@constraint(pm.model, ered == time_elapsed * pred)
end

function constraint_red_state(pm::_PM.AbstractPowerModel, n_1::Int, n_2::Int, i::Int, time_elapsed)
    pred   = _PM.var(pm, n_2, :pred, i)
    ered_2 = _PM.var(pm, n_2, :ered, i)
    ered_1 = _PM.var(pm, n_1, :ered, i)

    # Accumulation of not consumed energy (accumulation of voluntary load reduction for each time step)
    JuMP.@constraint(pm.model, ered_2 - ered_1 == time_elapsed * pred)
end

function constraint_shift_duration_up(pm::_PM.AbstractPowerModel, n::Int, i::Int, pd, pshift_up_rel_max, start_period::Int)
    # Apply an upper bound to the demand shifted upward during the recovery period
    JuMP.@constraint(pm.model, sum(_PM.var(pm, t, :pshift_up, i) for t in start_period:n) <= pshift_up_rel_max * pd)
end

function constraint_shift_duration_down(pm::_PM.AbstractPowerModel, n::Int, i::Int, pd, pshift_down_rel_max, start_period::Int)
    # Apply an upper bound to the demand shifted downward during the recovery period
    JuMP.@constraint(pm.model, sum(_PM.var(pm, t, :pshift_down, i) for t in start_period:n) <= pshift_down_rel_max * pd)
end

function constraint_shift_up_state_initial(pm::_PM.AbstractPowerModel, n::Int, i::Int, time_elapsed)
    pshift_up = _PM.var(pm, n, :pshift_up, i)
    eshift_up = _PM.var(pm, n, :eshift_up, i)

    # Initialization of accumulated upward demand shifting variable
    JuMP.@constraint(pm.model, eshift_up == time_elapsed * pshift_up)
end

function constraint_shift_up_state(pm::_PM.AbstractPowerModel, n_1::Int, n_2::Int, i::Int, time_elapsed)
    pshift_up   = _PM.var(pm, n_2, :pshift_up, i)
    eshift_up_2 = _PM.var(pm, n_2, :eshift_up, i)
    eshift_up_1 = _PM.var(pm, n_1, :eshift_up, i)

    # Accumulation of upward demand shifting for each time step
    JuMP.@constraint(pm.model, eshift_up_2 - eshift_up_1 == time_elapsed * pshift_up)
end

function constraint_shift_down_state_initial(pm::_PM.AbstractPowerModel, n::Int, i::Int, time_elapsed)
    pshift_down = _PM.var(pm, n, :pshift_down, i)
    eshift_down = _PM.var(pm, n, :eshift_down, i)

    # Initialization of accumulated downward demand shifting variable
    JuMP.@constraint(pm.model, eshift_down == time_elapsed * pshift_down)
end

function constraint_shift_down_state(pm::_PM.AbstractPowerModel, n_1::Int, n_2::Int, i::Int, time_elapsed)
    pshift_down   = _PM.var(pm, n_2, :pshift_down, i)
    eshift_down_2 = _PM.var(pm, n_2, :eshift_down, i)
    eshift_down_1 = _PM.var(pm, n_1, :eshift_down, i)

    # Accumulation of downward demand shifting for each time step
    JuMP.@constraint(pm.model, eshift_down_2 - eshift_down_1 == time_elapsed * pshift_down)
end

function constraint_shift_state_final(pm::_PM.AbstractPowerModel, n::Int, i::Int)
    eshift_up   = _PM.var(pm, n, :eshift_up, i)
    eshift_down = _PM.var(pm, n, :eshift_down, i)

    # The accumulated upward demand shifting over the operational planning horizon should equal the accumulated downward
    # demand shifting (since this is demand shifted and not reduced or curtailed)
    JuMP.@constraint(pm.model, eshift_up == eshift_down)
end

function constraint_shift_balance_periodic(pm::_PM.AbstractPowerModel, horizon::AbstractVector{Int}, i::Int)
    pshift_up   = _PM.var.(Ref(pm), horizon, :pshift_up, i)
    pshift_down = _PM.var.(Ref(pm), horizon, :pshift_down, i)

    JuMP.@constraint(pm.model, sum(pshift_up) == sum(pshift_down))
end


text:
module Impute

using DataFrames
using Statistics

import DataFrames: DataFrameRow
import Base.Iterators

export impute, impute!, chain, chain!, drop, drop!, interp, interp!, ImputeError

const Dataset = Union{AbstractArray, DataFrame}

"""
    ImputeError{T} <: Exception

Is thrown by `impute` methods when the limit of imputable values has been exceeded.

# Fields
* msg::T - the message to print.
"""
struct ImputeError{T} <: Exception
    msg::T
end

Base.showerror(io::IO, err::ImputeError) = println(io, "ImputeError: $(err.msg)")

include("context.jl")
include("imputors.jl")

const global imputation_methods = Dict{Symbol, Type}(
    :drop => Drop,
    :interp => Interpolate,
    :fill => Fill,
    :locf => LOCF,
    :nocb => NOCB,
)

"""
    impute!(data::Dataset, method::Symbol=:interp, args...; limit::Float64=0.1)

Looks up the `Imputor` type for the `method`, creates it and calls
`impute!(imputor::Imputor, data::Dataset, limit::Float64)` with it.

# Arguments
* `data::Dataset`: the datset containing missing elements we should impute.
* `method::Symbol`: the imputation method to use
    (options: [`:drop`, `:fill`, `:interp`, `:locf`, `:nocb`])
* `args::Any...`: any arguments you should pass to the `Imputor` constructor.
* `limit::Float64`: missing data ratio limit/threshold (default: 0.1)
"""
function impute!(data::Dataset, method::Symbol, args...; limit::Float64=0.1)
    imputor_type = imputation_methods[method]
    imputor = length(args) > 0 ? imputor_type(args...) : imputor_type()
    return impute!(imputor, data, limit)
end

"""
    impute!(data::Dataset, missing::Function, method::Symbol=:interp, args...; limit::Float64=0.1)

Creates the appropriate `Imputor` type and `Context` (using `missing` function) in order to call
`impute!(imputor::Imputor, ctx::Context, data::Dataset)` with them.

# Arguments
* `data::Dataset`: the datset containing missing elements we should impute.
* `missing::Function`: the missing data function to use
* `method::Symbol`: the imputation method to use
    (options: [`:drop`, `:fill`, `:interp`, `:locf`, `:nocb`])
* `args::Any...`: any arguments you should pass to the `Imputor` constructor.
* `limit::Float64`: missing data ratio limit/threshold (default: 0.1)
"""
function impute!(data::Dataset, missing::Function, method::Symbol, args...; limit::Float64=0.1)
    imputor_type = imputation_methods[method]
    imputor = length(args) > 0 ? imputor_type(args...) : imputor_type()
    ctx = Context(*(size(data)...), 0, limit, missing)
    return impute!(imputor, ctx, data)
end

"""
    impute(data::Dataset, args...; kwargs...)

Copies the `data` before calling `impute!(new_data, args...; kwargs...)`
"""
function impute(data::Dataset, args...; kwargs...)
    return impute!(deepcopy(data), args...; kwargs...)
end

"""
    chain!(data::Dataset, missing::Function, imputors::Imputor...; kwargs...)

Creates a `Chain` with `imputors` and calls `impute!(imputor, missing, data; kwargs...)`
"""
function chain!(data::Dataset, missing::Function, imputors::Imputor...; kwargs...)
    imputor = Chain(imputors...)
    return impute!(imputor, missing, data; kwargs...)
end

"""
    chain!(data::Dataset, imputors::Imputor...; kwargs...)

Creates a `Chain` with `imputors` and calls `impute!(imputor, data; kwargs...)`
"""
function chain!(data::Dataset, imputors::Imputor...; kwargs...)
    imputor = Chain(imputors...)
    return impute!(imputor, data; kwargs...)
end

"""
    chain(data::Dataset, args...; kwargs...)

Copies the `data` before calling `chain!(data, args...; kwargs...)`
"""
function chain(data::Dataset, args...; kwargs...)
    result = deepcopy(data)
    return chain!(data, args...; kwargs...)
end

"""
    drop!(data::Dataset; limit=1.0)

Utility method for `impute!(data, :drop; limit=limit)`
"""
drop!(data::Dataset; limit=1.0) = impute!(data, :drop; limit=limit)

"""
    drop(data::Dataset; limit=1.0)

Utility method for `impute(data, :drop; limit=limit)`
"""
Iterators.drop(data::Dataset; limit=1.0) = impute(data, :drop; limit=limit)

"""
    interp!(data::Dataset; limit=1.0)

Utility method for `impute!(data, :interp; limit=limit)`
"""
interp!(data::Dataset; limit=1.0) = impute!(data, :interp; limit=limit)

"""
    interp(data::Dataset; limit=1.0)

Utility method for `impute(data, :interp; limit=limit)`
"""
interp(data::Dataset; limit=1.0) = impute(data, :interp; limit=limit)

end  # module

ext: jl


text:
## matrix class stuff
## Work with Array{Sym}, not python array objects, as possible

## The basic idea is: we use pytype mapping to make sympy matrices into Array{Sym}.
## this is achieved with the convert(Array{Sym}, o::PyObject) method below.

## Array{Sym} objects are converted into Python objects via
PyCall.PyObject(a::AbstractArray{Sym}) = pycall(sympy[:Matrix], PyObject, PyCall.array2py(a))

## Matrix methods and objects

## For calling methods we have  call_matrix_meth(M, :meth, ...) for M.meth(...)
## This helps, grabbing the M.meth part
function getindex(s::AbstractArray{Sym}, i::Symbol)
    PyObject(s)[i]
end

# call a matrix method M[:det](). Matrix or vector arguments are converted via symmatrix
# though this may need to be done specially for some arguments that are passed in as collections
global call_matrix_meth(object, meth, args...; kwargs...) = begin
    o = PyObject(object)
    o[meth](args...; kwargs...)
end

## support for convert(Array{Sym}, o::PyObject)
function matrix_size(x::PyObject)
    a = x[:shape]            # tuple
    if isa(a, PyObject)
        a =  (a[:__getitem__](0), a[:__getitem__](1))
    end
    while a[end] == 1
        a = a[1:end-1]
    end
    a
end

get_matrix_index(s::PyObject, i::Integer...) = get(s, Sym, ntuple(k -> i[k]-1, length(i)))
function get_matrix_index(s::PyObject, i::Integer)
    ind = ind2sub(matrix_size(s), i-1) # 0-base PyThon, 1 base Julia
    if length(ind) == 1
        get(s, Sym, ind[1])
    else
        get(s, Sym, ind)
    end
end

function convert(::Type{Array{Sym}}, o::PyObject)
    sz = matrix_size(o)
    ndims = length(sz)
    if ndims == 0
        Sym(o)
    elseif ndims == 1
        Sym[get_matrix_index(o,i) for i in 1:sz[1]]
    elseif ndims == 2
        Sym[get_matrix_index(o,i,j) for i in 1:sz[1], j in 1:sz[2]]
    else
        ## need something else for arrays... XXX -- can't linear index a
        b = Sym[get_matrix_index(o,i) for i in 1:prod(sz)]
        reshape(b, sz)
    end
end

## covert back to Array{Sym}. Could just use broadcast (subs.(...)) here
## once v0.4 support is dropped.
subs(ex::AbstractArray{Sym}, args...; kwargs...) = map(u -> subs(u, args...; kwargs...), ex)

## Methods

## these are sympy methods that need exporting
sympy_matrix_methods = (:jordan_cell,
                        )

for meth in sympy_matrix_methods
    meth_name=string(meth)
    @eval begin
        @doc """
`$($meth_name)`: a SymPy function.
The SymPy documentation can be found through: http://docs.sympy.org/latest/search.html?q=$($meth_name)

Specific docs may also be found at [SymPy Docs for matrices](http://docs.sympy.org/latest/modules/matrices/matrices.html#module-sympy.matrices.matrices)
""" ->
         ($meth)(args...; kwargs...) = sympy_meth(Symbol($meth_name), args...; kwargs...)
    end
    eval(Expr(:export, meth))
end

VERSION >= v"0.7.0-" && import Base.adjoint

## These are matrix methods that need exporting
matrix_methods = (:LDLsolve,
                  :LDLdecomposition, :LDLdecompositionFF,
                  :LUdecomposition_Simple,
                  :LUdecomposition,
                  :LUsolve,
                  :QRdecomposition, :QRsolve,
                  :adjoint, :adjugate,
                  :cholesky, :cholesky_solve, :conjugate,
                  :diagonal_solve, :diagonalize, :dual,
                  :expand,
                  :integrate,
                  :is_symmetric,
                  :inverse_ADJ, :inverse_GE, :inverse_LU,
                  :jordan_form,
                  :limit,
                  :lower_triangular_solve,
                  :minorEntry, :minorMatrix,
                  :normalized, :nullspace,
                  :permuteBkwd, :permuteFwd,
                  :print_nonzero,
                  :rref,
                  :singular_values,
                  :upper_triangular_solve,
                  :vech
)

for meth in matrix_methods
    meth_name = string(meth)
    @eval begin
      @doc """
`$($meth_name)`: a SymPy function.
The SymPy documentation can be found through: http://docs.sympy.org/latest/search.html?q=$($meth_name)

Specific docs may also be found at [SymPy Docs for matrices](http://docs.sympy.org/latest/modules/matrices/matrices.html#module-sympy.matrices.matrices)
""" ->
        ($meth)(ex::Matrix{Sym}, args...; kwargs...) =  call_matrix_meth(ex, Symbol($meth_name), args...;kwargs...)
    end
    eval(Expr(:export, meth))
end

cofactor(A::Matrix{Sym}, i, j) = call_matrix_meth(A, :cofactor, i-1, j-1)
jacobian(X::AbstractArray{Sym}, Y::AbstractArray{Sym}) = call_matrix_meth(X, :jacobian, Y)

export cofactor

## These are SymPy properties that we want to call as methods, exported
matrix_properties = (:H, :C,
                    :is_lower, :is_lower_hessenberg, :is_square, :is_upper,
                    :is_upper_hessenberg, :is_zero
                     )

for meth in matrix_properties
     meth_name = string(meth)
    @eval begin
      @doc """
`$($meth_name)`: a SymPy function.
The SymPy documentation can be found through: http://docs.sympy.org/latest/search.html?q=$($meth_name)
""" ->
        ($meth)(ex::Matrix{Sym}, args...; kwargs...) =  ex[Symbol($meth_name)]
    end
    eval(Expr(:export, meth))
end

## These are special cased
norm(a::AbstractVector{Sym}, args...; kwargs...) = call_matrix_meth(a, :norm, args...; kwargs...)
norm(a::AbstractMatrix{Sym}, args...; kwargs...) = call_matrix_meth(a, :norm, args...; kwargs...)
chol(a::Matrix{Sym}) = cholesky(a)
exp(a::Matrix{Sym}) = call_matrix_meth(a, :exp)
conj(a::Sym) = conjugate(a)
eigvals(a::Matrix{Sym}) = collect(keys(call_matrix_meth(a, :eigenvals))) # a[:eigevnals]() has multiplicity
function eigvecs(a::Matrix{Sym})
    ds =  call_matrix_meth(a, :eigenvects)
    out = Any[]
    for d in ds
        append!(out, d[3])
    end
    out
end
rref(a::Matrix{T}) where {T <: Integer} = N(rref(convert(Matrix{Sym}, a)))
rref(a::Matrix{Rational{T}}) where {T <: Integer} = N(rref(convert(Matrix{Sym}, a)))

"""
Return orthogonal basis from a set of vectors

Example:
```
L = [Sym[1,2,3], Sym[2,5,9], Sym[1,4,2]] # need Sym vectors.
GramSchmidt(L, true)
```
"""
GramSchmidt(vecs::Vector{Vector{T}}, args...; kwargs...) where {T} = sympy_meth(:GramSchmidt, map(u->convert(Vector{Sym},u),vecs), args...; kwargs...)
## :cross?
## hessian, wronskian, GramSchmidt

"""

Find Hessian of a symbolic expression

Example:
```
x, y = symbols("x y")
f = x^2 - 2x*y
hessian(f, [x,y])

u = SymFunction("u")(x,y)
hessian(u, [x,y])
```
"""
hessian(u::SymbolicObject, vars::AbstractArray{Sym}, args...) = sympy_meth(:hessian, u, vars, args...)
hessian(u::SymbolicObject) = hessian(u, free_symbols(u))

"""
Wronskian of functions

Example
```
u = SymFunction("u")(x)
v = SymFunction("v")(x)
wronskian([u,v], x)  # determinant of [u v; u' v']
```

"""
wronskian(fs::Vector{T}, x::Sym, args...; kwargs...) where {T <: SymbolicObject} = sympy_meth(:wronskian, fs, x, args...; kwargs...)

export GramSchmidt, hessian, wronskian


text:
using JuMP
using ConditionalJuMP
using Gurobi

struct State{T}
    q::T
    v::T
end

JuMP.getvalue(s::State) = State(getvalue(s.q), getvalue(s.v))

const k = 1000
const Δt = 0.1
const u_max = 10

function update(x::State, u)
    contact_force = ifelse_conditional(@?(x.q <= 0),
        -k * x.q,
        zero(x.q))
    acceleration = u + contact_force
    x_next = State(x.q + Δt * x.v + 0.5 * Δt^2 * acceleration,
                   x.v + Δt * acceleration)
end

function simulate(x0, us)
    xs = [update(x0, us[1])]
    for i in 2:length(us)
        push!(xs, update(xs[end], us[i]))
    end
    xs
end

function run_mpc(x0::State, N=10)
    m = Model(solver=GurobiSolver(OutputFlag=0))
    @variable m -u_max <= u[1:N] <= u_max
    xs = [State(@variable(m, basename="q", lowerbound=-10, upperbound=10),
                @variable(m, basename="v", lowerbound=-10, upperbound=10)) for i in 1:N]
    @constraint m xs[1].q == x0.q
    @constraint m xs[1].v == x0.v
    for i in 2:N
        xi = update(xs[i-1], u[i-1])
        @constraint m xi.q == xs[i].q
        @constraint m xi.v == xs[i].v
    end
    x_f = xs[end]
    @objective(m, Min, 10 * ((x_f.q - 1.0)^2 + x_f.v^2) + 0.01 * sum(u.^2))
    warmstart!(m)
    solve(m)
    getvalue.(xs), getvalue.(u)
end


text:
module CalibrateParameters
	include("calibration_utils.jl")
	using JLD2, FileIO, ImpvolEquilibrium, Base.Test

	function calibrate_parameters!(parameters, fname="../../data/impvol_data.jld2")
		data = load(fname)

		_, N, J, T = size(data["beta"])
		parameters[:N], parameters[:J], parameters[:T] = N, J, T
		parameters[:beta_j] = mean(data["beta"],(1,2,4))

		parameters[:gamma_jk] = compute_gamma(parameters, data)

		parameters[:S_nt] = zeros(1,N,1,T)

		parameters[:d] = expenditure_shares(parameters, data)

		parameters[:kappa_mnjt] = trade_costs(parameters)

		final_expenditure_shares = calculate_expenditure_shares(parameters, data)
		parameters[:final_expenditure_shares] = final_expenditure_shares

		# broad country weights for final expenditure
		country_weights = sum(data["va"], (1,3,4))
		country_weights = country_weights ./ sum(country_weights, 2)

		calculate_p_and_nu!(parameters, data, final_expenditure_shares, country_weights)

		parameters[:w_njt] = calculate_nominal_wages(parameters, data)

		parameters[:B_j] = calculate_B(parameters)

		parameters[:xi] = calculate_xi(parameters)

		parameters[:A] = calculate_A(parameters, data)

		# total world expenditure in the data - needed to get reasonable starting values
		parameters[:nominal_world_expenditure] = sum(data["va"] ./ parameters[:beta_j], (1,2,3))
		# deflate trade imbalance to 1972 dollars
		deflator = CES_price_index(parameters[:nu_njt][:,end:end,:,:], parameters[:p_sectoral][:,end:end,:,:], parameters[:sigma])
		info(deflator[:])

		parameters[:S_nt_data] = (data["trade_balance"] .- mean(data["trade_balance"],2)) ./ deflator

		# global, all-time average of sector final expenditure shares
		importance_weight = mean(parameters[:nu_njt], (1, 2, 4))
		# special-case CES, when nu does not have direct meaning
		if abs(parameters[:sigma]-1)>0.1
			share = data["va"] ./ sum(data["va"], 3)
			importance_weight = mean(share, (1, 2, 4))
		end
		parameters[:importance_weight] = importance_weight
		decompose_shocks!(parameters, importance_weight)
		draw_productivity_shocks!(parameters)
	end

	function compute_gamma(parameters, data)
		N, J, T = parameters[:N], parameters[:J], parameters[:T]

		beta = parameters[:beta_j]
		io_values = data["io_values"]
		total_output = data["total_output"]
		output_shares = data["output_shares"]
		intermediate_input_shares = data["intermediate_input_shares"]

		# Summing sectors
		# Agriculture & mining
		io_values[1,:,:,:]  = sum(io_values[1:2,:,:,:],1)
		io_values         = io_values[setdiff(1:end,2),:,:,:]
		io_values[:,1,:,:]  = sum(io_values[:,1:2,:,:],2)
		io_values         = io_values[:,setdiff(1:end,2),:,:]
		total_output[:,1,:,:] = sum(total_output[:,1:2,:,:],2)
		total_output      = total_output[:,setdiff(1:end,2),:,:]

		# Services
		io_values[18,:,:,:]  = sum(io_values[18:end,:,:,:],1)
		io_values          = io_values[setdiff(1:end,19:end),:,:,:]
		io_values[:,18,:,:]  = sum(io_values[:,18:end,:,:],2)
		io_values          = io_values[:,setdiff(1:end,19:end),:,:]
		total_output[:,18,:,:] = sum(total_output[:,18:end,:,:],2)
		total_output       = total_output[:,setdiff(1:end,19:end),:,:]

		# Split rows
		dupl_idx = [1, 2, 2, 3, 3, 3, 4, 5, 5, 6, 7, 8, 9, 10, 11, 12, 13, 13, 13, 14, 15, 16, 17, 18]
		io_values_dupl = io_values[dupl_idx,:,:,:]

		output_shares_full = ones(size(io_values_dupl))
		split_idx = [2, 3, 4, 5, 6, 8, 9, 17, 18, 19]
		output_shares_full[split_idx,:,:,:] = permutedims(repeat(output_shares, outer = [1,1,size(io_values_dupl,2),1]),(2,3,1,4))

		io_values_new = io_values_dupl .* output_shares_full

		# Split columns
		io_values_dupl = io_values_new[:,dupl_idx,:,:]

		intermediate_input_shares_full = ones(size(io_values_dupl))
		intermediate_input_shares_full[:,split_idx,:,:] = repeat(intermediate_input_shares, outer = [size(io_values_dupl,1),1,1,1])

		io_values_new = io_values_dupl .* intermediate_input_shares_full

		total_output = total_output[:,dupl_idx,:,:]
		output_shares_full = ones(size(total_output))
		output_shares_full[:,split_idx,:,:] = output_shares
		total_output = total_output .* output_shares_full

		# Correct order of sectors
		io_values_new[18,:,:,:], io_values_new[19,:,:,:], io_values_new[20,:,:,:] = io_values_new[20,:,:,:], io_values_new[18,:,:,:], io_values_new[19,:,:,:]
		io_values_new[:,18,:,:], io_values_new[:,19,:,:], io_values_new[:,20,:,:] = io_values_new[:,20,:,:], io_values_new[:,18,:,:], io_values_new[:,19,:,:]
		total_output[:,18,:,:], total_output[:,19,:,:], total_output[:,20,:,:] = total_output[:,20,:,:], total_output[:,18,:,:], total_output[:,19,:,:]

		# Compute gamma
		gamma = io_values_new ./ repeat(total_output, outer = [size(io_values_new,1),1,1,1])
		gamma = mean(gamma,4)
		gamma = gamma .* permutedims(1-beta,(1,3,2,4)) ./ sum(gamma,1)
		return gamma = squeeze(gamma,(3,4))
	end

	function compute_alpha(parameters, data)
		N, J, T = parameters[:N], parameters[:J], parameters[:T]

		va = data["va"]
		beta = parameters[:beta_j]
		gamma = parameters[:gamma_jk]
		weights = parameters[:bp_weights]

		alpha = zeros(J,T)

		for t in 1:T
			va_t = transpose(sum(va[1,:,:,t],1))
			alpha[:,t] = (eye(J) - gamma) * diagm(1 ./ beta[:],0) * va_t / sum(va_t)
		end

		# Replace negative elements with 0
		alpha = (alpha + abs.(alpha)) / 2

		# Smooth the series
		alpha_c, alpha_t = DetrendUtilities.detrend(alpha, weights)

		# Normalization
		alpha = alpha_t ./ sum(alpha_t,1)

		return alpha = permutedims(cat(ndims(alpha) + 2,alpha),(3,4,1,2))
	end

	function trade_costs(parameters)
		N, J, T = parameters[:N], parameters[:J], parameters[:T]

		d = parameters[:d]
		theta = parameters[:theta]
		n_zero = parameters[:numerical_zero]

		# smooth kappa over time so that it does not introduce shocks
		weights = parameters[:bp_weights]
		_, share_trend = DetrendUtilities.detrend(d, weights)
		# DEBUG: undo all this
		share_trend = d

		kappa = zeros(size(d))
		for j in 1:(J-1)
			for t = 1:T
				kappa[:,:,j,t] = ((share_trend[:,:,j,t] .* transpose(share_trend[:,:,j,t])) ./ (diag(share_trend[:,:,j,t]) * transpose(diag(share_trend[:,:,j,t])))).^(1 / (2 * theta))
			end
		end

		kappa[kappa .< n_zero] = n_zero
		kappa[:,:,end,:] = repeat(eye(N), outer = [1,1,1,T]) # Services

		kappa = min.(kappa,1)
		# smooth kappa over time so that it does not introduce shocks
		_, kappa_trend = DetrendUtilities.detrend(kappa, weights)
		# DEBUG: return kappa, not kappa_trend
		return kappa
	end

	function expenditure_shares(parameters, data)
		N, J, T = parameters[:N], parameters[:J], parameters[:T]

		import_shares = data["import_shares"]
		n_zero = parameters[:numerical_zero]

		d = import_shares

		within_import = d ./ sum(d, 2)
		domestic_per_import = 1 ./ sum(d, 2) - 1
		domestic_per_import[domestic_per_import .< n_zero] = n_zero
		d = within_import ./ (1 + domestic_per_import)

		for n in 1:N
			d[n,n,:,:] = ones(J,T) - squeeze(sum(d[n,:,:,:],1),1)
		end

		d[d .< n_zero] = n_zero

		return d ./ sum(d, 2)
	end

	function calculate_xi(parameters)
		theta = parameters[:theta]
		eta = parameters[:eta]

		return gamma((theta + 1 - eta)/theta)
	end

	function calculate_B(parameters)
		beta = parameters[:beta_j]
		gamma = parameters[:gamma_jk]

		gamma = permutedims(cat(ndims(gamma) + 2,gamma),[1,3,2,4])
		return B = (beta .^ -beta) .* prod(gamma .^ -gamma, 1)
	end

	function calculate_nominal_wages(parameters, data)
		nulla = parameters[:numerical_zero]
		weights = parameters[:bp_weights]
		value_added_shares = data["va"] ./ sum(data["va"], 3)
		V_c, V_t = DetrendUtilities.detrend(value_added_shares, weights)

		if parameters[:one_over_rho]>0.0
			trend = 0.5*(V_t - parameters[:one_over_rho])
			labor_share = trend .+ (trend .^2 .+ parameters[:one_over_rho]*value_added_shares) .^ 0.5
			wage_ratio = value_added_shares ./ labor_share
			info("Unweighted wage ratio should be 1: ", mean(wage_ratio))
		else
			# if no labor adjustment, the ratio of value added = the ratio of wages
			wage_ratio = value_added_shares ./ V_t
		end
		nominal_GDP = sum(data["va"], 3)
		return nominal_GDP .* wage_ratio
	end

	function calculate_p_and_nu!(parameters, data, final_expenditure_shares, country_weights)
		N = parameters[:N]
		J = parameters[:J]
		T = parameters[:T]

		p_sectoral_data = data["p_sectoral_data"]
		d = parameters[:d]
		kappa = parameters[:kappa_mnjt]
		theta = parameters[:theta]
		sigma = parameters[:sigma]
		nulla = parameters[:numerical_zero]

		p_sectoral_base = p_sectoral_data ./ p_sectoral_data[:,:,:,1]

		# step 1: calculate nu and price index for base country
		# NB: US as the last country in the matrix
		p_sectoral_US = p_sectoral_base[:,end:end,:,:]
		nu_US = final_expenditure_shares[:,end:end,:,:] .* p_sectoral_US .^ (sigma-1)
		nu_US = nu_US ./ sum(nu_US, 3)
		P_US = CES_price_index(nu_US, p_sectoral_US, sigma)
		@test p_sectoral_US[1,1,:,1] ≈ ones(J) atol=1e-9
		@test P_US[1,1,1,1] ≈ 1.0 atol=1e-9

		# step 2: calculate sectoral prices from market shares relative to US
		# US is assumed to be chosen as a base country (US = end), else pwt should be used to do the conversion
		# normalization: p_sectoral[1,end,:,1] = 1.0
		p_sectoral = array_transpose(exp.( mean(1 / theta * log.(d ./ permutedims(cat(ndims(d),d[end,:,:,:]),[4,1,2,3])) - log.(kappa ./ permutedims(cat(ndims(kappa),kappa[end,:,:,:]),[4,1,2,3])), 2) + repeat(permutedims(cat(ndims(p_sectoral_base),log.(p_sectoral_base[:,end,:,:])), [1,4,2,3]), outer = [size(d,1),1,1,1]) ))
		@test any(isnan, p_sectoral[:,:,1:end-1,:]) == false
		# step 3: calculate tradable nu and infer nontradable nu
		nu = final_expenditure_shares .* (p_sectoral ./ (data["pwt"] .* P_US)) .^ (sigma-1)
		@test any(isnan, nu[:,:,1:end-1,:]) == false

		nontradable_nu = 1 .- sum(final_expenditure_shares[:,:,1:end-1,:], 3)
		# Replace negative elements with second smallest positive
		nu[:,:,end:end,:] .= DetrendUtilities.winsorize(nontradable_nu, 1)
		nu .= nu ./ sum(nu, 3)

		# demand shifter only varies across sectors and over time, not across countries
		parameters[:nu_njt] = sum(country_weights .* nu, (1,2))

		# enforce comformity of model with data
		final_expenditure_shares = parameters[:nu_njt] .* (p_sectoral ./ (data["pwt"] .* P_US)) .^ (1-sigma)
		nontradable_nu = 1 .- sum(final_expenditure_shares[:,:,1:end-1,:], 3)
		# Replace negative elements with second smallest positive
		final_expenditure_shares[:,:,end:end,:] .= DetrendUtilities.winsorize(nontradable_nu, 1)
		final_expenditure_shares .= final_expenditure_shares ./ sum(final_expenditure_shares, 3)

		# step 4: calculate nontradable prices
		# NB: DO NOT recalibrate tradable prices, expenditure_shares are very noisy for small sectors
		if abs(sigma-1)>0.01
			p_sectoral[:,:,end:end,:] = data["pwt"] .* P_US .* (parameters[:nu_njt][:,:,end:end,:] ./ final_expenditure_shares[:,:,end:end,:]) .^ (1/(sigma-1))
		else
			p_sectoral[:,:,end:end,:] = (data["pwt"] .* P_US ./ (prod(p_sectoral[:,:,1:end-1,:] .^ parameters[:nu_njt][:,:,1:end-1,:], 3))) .^ (1 ./ parameters[:nu_njt][:,:,end:end,:])
		end
		parameters[:p_sectoral] = p_sectoral
		parameters[:final_expenditure_shares] = final_expenditure_shares
	end

	function calculate_A(parameters, data)
		N, J, T = parameters[:N], parameters[:J], parameters[:T]

		p_njt = parameters[:p_sectoral]
		p_mjt = array_transpose(p_njt)
		beta_j = parameters[:beta_j]
		gamma = parameters[:gamma_jk]
		kappa_mnjt = parameters[:kappa_mnjt]
		w_njt = parameters[:w_njt]
		B = parameters[:B_j]
		d_mnjt = parameters[:d]
		xi = parameters[:xi]
		theta = parameters[:theta]

		z = zeros(1, N, J, T)

		# use eq 15 in algorithm.pdf
		rho_mnjt = kappa_mnjt .* p_mjt .* d_mnjt .^ (-1/theta)
		rho_njt = exp.(mean(log.(rho_mnjt),1))
		# nontradable input price equals output price
		rho_njt[1,:,end,:] = p_njt[1,:,end,:]
		input_price_index = exp.(rotate_sectors(gamma', log.(p_njt)))

		A_njt = xi .* B ./ rho_njt .* w_njt .^ beta_j .* input_price_index
		return A_njt
	end

	function calculate_expenditure_shares(parameters, data)
		N, J, T = parameters[:N], parameters[:J], parameters[:T]

		nulla = parameters[:numerical_zero]
		weights = parameters[:bp_weights]

		beta = parameters[:beta_j]
		gamma = parameters[:gamma_jk]
		d = parameters[:d]
		# service import shares are NaN
		for t=1:T
			d[:,:,J,t] = eye(N)
		end
		va = data["va"]

		#beta = squeeze(beta,(1,2,4))
		revenue = va ./ beta
		expenditure = zeros(revenue)
		for j=1:J
			for t=1:T
				expenditure[1,:,j,t]  = revenue[1,:,j,t]' * inv(d[:,:,j,t])
			end
		end
		intermediate = rotate_sectors(gamma, revenue)
		final_expenditure = expenditure - intermediate

		# Smooth the series
		_, nu_guess = DetrendUtilities.detrend(final_expenditure, weights)

		# Replace negative elements with smallest positive
		nu_guess .= DetrendUtilities.winsorize(nu_guess ./ sum(nu_guess, 3), 0)
		# Smooth the series
		nu_c, nu_t = DetrendUtilities.detrend(nu_guess, weights)

		# Normalization
		return nu_t ./ sum(nu_t, 3)
end

	function estimate_AR1(data)
		# data is M,N,J,T
		_, N, J, T = size(data)
		current = data[:,:,:,2:T]
		lag = data[:,:,:,1:T-1]

		constant = zeros(1, N, J, 1)
		rho = zeros(1, N, J, 1)
		sigma = zeros(1, N, J, 1)

		# estimate a separate AR(1) for each series
		for n=1:N
			for j=1:J
				y = current[1,n,j,:]
				X = cat(2, ones(T-1), lag[1,n,j,:])

				constant[1,n,j,1], rho[1,n,j,1] = X \ y
				sigma[1,n,j,1] = std(y - X * [constant[1,n,j,1], rho[1,n,j,1]])
			end
		end

		return (constant, rho, sigma)
	end

	function draw_productivity_shocks!(parameters)
		S, T = parameters[:S], parameters[:T]

		parameters[:global_sectoral_shock_njs] = draw_random_realizations(parameters[:global_sectoral_shock], S)
		parameters[:country_shock_njs] = draw_random_realizations(parameters[:country_shock], S)
		parameters[:idiosyncratic_shock_njs] = draw_random_realizations(parameters[:idiosyncratic_shock], S)

		parameters[:A_njs] = map(t ->
			exp.(ImpvolEquilibrium.non_random_variable(parameters[:productivity_trend], t)
				.+ parameters[:global_sectoral_shock_njs][t]
				.+ parameters[:country_shock_njs][t]
				.+ parameters[:idiosyncratic_shock_njs][t]),
				 1:T)
	end

	function draw_random_realizations(data, S)
		# data is M,N,J,T
		_, N, J, T = size(data)
		constant, rho, sigma = estimate_AR1(data)

		draws = Array{Array{Float64, 4}}(T)
		draws[1] = ImpvolEquilibrium.non_random_variable(data, 1)
		for t=2:T
			innovation = sigma .* randn(1,N,J,S - 1)
			random_realization = ImpvolEquilibrium.non_random_variable(data, t)
			past_productivity = ImpvolEquilibrium.non_random_variable(data, t-1)
			# reversion towards mean
			draws[t] = cat(4, random_realization, constant .* (1-rho) .+ past_productivity .* rho .+ innovation)
		end
		return draws
	end

	function decompose_shocks!(parameters, sectoral_weights)
		# M,N,J,T
		# Smooth the series
		weights = parameters[:bp_weights]
		detrended_log_productivity, parameters[:productivity_trend] = DetrendUtilities.detrend(log.(parameters[:A]), weights)

		global_sectoral_shock = mean(detrended_log_productivity, 2)
		# weighted by sector importance, see https://github.com/ceumicrodata/impvol/commit/91d92905678df96d7068b8dd729e6f6d7cf470d8
		country_shock = sum(sectoral_weights .* (detrended_log_productivity .- global_sectoral_shock), 3) ./ sum(sectoral_weights, 3)
		idiosyncratic_shock = detrended_log_productivity .- global_sectoral_shock .- country_shock

		parameters[:global_sectoral_shock] = global_sectoral_shock
		parameters[:country_shock] = country_shock
		parameters[:idiosyncratic_shock] = idiosyncratic_shock
	end
	function jld_saver(data, file_name="results.jld2")
		jldopen(file_name, "w") do file
	   		file["results"] = data
		end
	end
	function jld_loader(file_name="results.jld2")
		file = jldopen(file_name, "r")
		return file["results"]
	end
end


text:
using OrderedCollections
using Oceananigans.OutputReaders: FieldTimeSeries

#####
##### The LESbrary (so to speak)
#####

include("./legacy_data_time_serieses.jl")

function get_times(path)
   file = jldopen(path)
   file_iterations = parse.(Int, keys(file["timeseries/t"]))
   file_times = [file["timeseries/t/$i"] for i in file_iterations]
   return file_times
end

# https://engaging-web.mit.edu/~alir/lesbrary/6DaySuite/
function SyntheticObservationsBatch(suite; first_iteration = 1, stride = 1, last_iteration = nothing, normalize = ZScore, Nz)

   observations = Vector{SyntheticObservations}()

   for case in values(suite)
      path = case.filename
      times = get_times(path)
      last_iteration = isnothing(last_iteration) ? length(times) : last_iteration
      times = times[first_iteration:stride:last_iteration]
      field_names = case.fields

      field_time_serieses = legacy_data_field_time_serieses(path, field_names, times)

      observation = SyntheticObservations(path; field_names, normalize, times, field_time_serieses, grid_size = (1, 1, Nz))
      push!(observations, observation)
   end

   return observations
end

# https://engaging-web.mit.edu/~alir/lesbrary/2DaySuite/
function TwoDaySuite(directory; first_iteration = 13, stride = 1, last_iteration = nothing, normalize = ZScore, Nz = 128)
   directory = joinpath(directory, "2DaySuite")
   suite = OrderedDict(
      "2d_free_convection" => (
         filename = joinpath(directory, "free_convection/instantaneous_statistics.jld2"),
         fields = (:b, :e)),
      "2d_strong_wind" => (
         filename = joinpath(directory, "strong_wind/instantaneous_statistics.jld2"),
         fields = (:b, :u, :v, :e)),
      "2d_strong_wind_no_rotation" => (
         filename = joinpath(directory, "strong_wind_no_rotation/instantaneous_statistics.jld2"),
         fields = (:b, :u, :e)),
      "2d_strong_wind_weak_cooling" => (
         filename = joinpath(directory, "strong_wind_weak_cooling/instantaneous_statistics.jld2"),
         fields = (:b, :u, :v, :e)),
      "2d_weak_wind_strong_cooling" => (
         filename = joinpath(directory, "weak_wind_strong_cooling/instantaneous_statistics.jld2"),
         fields = (:b, :u, :v, :e)),
   )

   return SyntheticObservationsBatch(suite; first_iteration, stride, last_iteration, normalize, Nz)
end

# https://engaging-web.mit.edu/~alir/lesbrary/4DaySuite/
function FourDaySuite(directory; first_iteration = 13, stride = 1, last_iteration = nothing, normalize = ZScore, Nz = 128)
   directory = joinpath(directory, "4DaySuite")
   suite = OrderedDict(
      "4d_free_convection" => (
         filename = joinpath(directory, "free_convection/instantaneous_statistics.jld2"),
         fields = (:b, :e)),
      "4d_strong_wind" => (
         filename = joinpath(directory, "strong_wind/instantaneous_statistics.jld2"),
         fields = (:b, :u, :v, :e)),
      "4d_strong_wind_no_rotation" => (
         filename = joinpath(directory, "strong_wind_no_rotation/instantaneous_statistics.jld2"),
         fields = (:b, :u, :e)),
      "4d_strong_wind_weak_cooling" => (
         filename = joinpath(directory, "strong_wind_weak_cooling/instantaneous_statistics.jld2"),
         fields = (:b, :u, :v, :e)),
      "4d_weak_wind_strong_cooling" => (
         filename = joinpath(directory, "weak_wind_strong_cooling/instantaneous_statistics.jld2"),
         fields = (:b, :u, :v, :e)),
   )

   return SyntheticObservationsBatch(suite; first_iteration, stride, last_iteration, normalize, Nz)
end

function SixDaySuite(directory; first_iteration = 13, stride = 1, last_iteration = nothing, normalize = ZScore, Nz = 128)
   directory = joinpath(directory, "6DaySuite")
   suite = OrderedDict(
      "6d_free_convection" => (
         filename = joinpath(directory, "free_convection/instantaneous_statistics.jld2"),
         fields = (:b, :e)),
      "6d_strong_wind" => (
         filename = joinpath(directory, "strong_wind/instantaneous_statistics.jld2"),
         fields = (:b, :u, :v, :e)),
      "6d_strong_wind_no_rotation" => (
         filename = joinpath(directory, "strong_wind_no_rotation/instantaneous_statistics.jld2"),
         fields = (:b, :u, :e)),
      "6d_strong_wind_weak_cooling" => (
         filename = joinpath(directory, "strong_wind_weak_cooling/instantaneous_statistics.jld2"),
         fields = (:b, :u, :v, :e)),
      "6d_weak_wind_strong_cooling" => (
         filename = joinpath(directory, "weak_wind_strong_cooling/instantaneous_statistics.jld2"),
         fields = (:b, :u, :v, :e)),
   )

   return SyntheticObservationsBatch(suite; first_iteration, stride, last_iteration, normalize, Nz)
end

function GeneralStrat(directory)
   directory = joinpath(directory, "6DaySuite")
   suite = OrderedDict(
      "general_strat_4" => (
         filename = joinpath(directory, "general_strat_4/instantaneous_statistics.jld2"),
         fields = (:b,),
         first = 37, # cut out the first 6 hours
         last = 288),
      "general_strat_8" => (
         filename = joinpath(directory, "general_strat_8/instantaneous_statistics.jld2"),
         fields = (:b,),
         first = 13,
         last = 648),
      "general_strat_16" => (
         filename = joinpath(directory, "general_strat_16/instantaneous_statistics.jld2"),
         fields = (:b,),
         first = 13,
         last = nothing),
      "general_strat_32" => (
         filename = joinpath(directory, "general_strat_32/instantaneous_statistics.jld2"),
         fields = (:b,),
         first = 13,
         last = nothing),
   )
end


text:
using Random
using Statistics
using Test

# Include testing functions
include("test_dimensions.jl")
include("test_symbolic_functions.jl")
include("test_demo_large.jl")

_testset_output_dim("small")
_testset_symbolic_functions("small", ["simple", "none"], ["abs", "square"])

@testset "Broken inputs" begin
    @test_throws ArgumentError tvdiff([0, 1, 2], 1, 0.1; scale="small", precond="bad_input")
    @test_throws ArgumentError tvdiff(
        [0, 1, 2], 1, 0.1; scale="small", diff_kernel="bad_input"
    )
    @test_throws MethodError tvdiff([0, 1, 2], 1, 0.1; scale="sm4ll")
    @test_throws DimensionMismatch tvdiff(
        [0, 1, 2], 1, 0.1; u_0=[1, 2, 3, 4, 5], scale="small"
    )
end


text:
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

module TestUtil

using MXNet
using Test

function test_getdocdefine()
  @info("Util::_getdocdefine")
  @test occursin("Defined in", mx._getdocdefine("sgd_update"))
end  # function test_getdocdefine

function test_firstarg()
  @info("Util::_firstarg")
  @test mx._firstarg(:(f(x, y))) == :x
  @test mx._firstarg(:(f(x::mx.NDArray, y))) == :x
  @test mx._firstarg(:(f(x::mx.NDArray, y::mx.NDArray))) == :x
  @test mx._firstarg(:(f(x::Int, y::mx.NDArray))) == :x
  @test mx._firstarg(:(f(x::Int, y::mx.NDArray; other = 42))) == :x
  @test mx._firstarg(:(f(x::mx.NDArray{T}, y) where {T})) == :x
  @test mx._firstarg(:(f(x::mx.NDArray{T,N}, y) where {T,N})) == :x
  @test mx._firstarg(:(f(x::mx.NDArray{T,N} where {T,N}, y))) == :x
  @test mx._firstarg(:(broadcast_(::typeof(asin), x::mx.NDArray))) == :x
  @test mx._firstarg(:(broadcast_(::typeof(asin), x::mx.NDArray, y::mx.NDArray))) == :x
end  # function test_firstarg

@testset "Util Test" begin
  test_firstarg()
  test_getdocdefine()
end  # @testset "Util"

end  # module TestUtil


text:
module test_bukdu_resources

using Test
using Bukdu # ApplicationController Conn Router Utils routes resources
import Bukdu.Actions: index, show, new, edit, create, delete, update

struct ArticleController <: ApplicationController
    conn::Conn
end

index(::ArticleController)   = (:index,)
show(c::ArticleController)   = (:show, c.params.id)
new(::ArticleController)     = (:new,)
edit(c::ArticleController)   = (:edit, c.params.id)
create(::ArticleController)  = (:create,)
delete(c::ArticleController) = (:delete, c.params.id)
update(c::ArticleController) = (:update, c.params.id)

routes() do
    resources("/articles", ArticleController)
end

@test Router.call(get, "/articles").got   == (:index,)
@test Router.call(get, "/articles/2").got == (:show, "2")

@test Utils.read_stdout(CLI.routes) == """
GET     /articles           ArticleController  index
GET     /articles/new       ArticleController  new
GET     /articles/:id/edit  ArticleController  edit
GET     /articles/:id       ArticleController  show
POST    /articles           ArticleController  create
DELETE  /articles/:id       ArticleController  delete
PATCH   /articles/:id       ArticleController  update
PUT     /articles/:id       ArticleController  update"""

Router.call(delete, "/articles/1")
Router.call(patch, "/articles/1")
Router.call(put, "/articles/1")

Routing.empty!()

routes() do
    resources("/articles", ArticleController, only=[index, show])
end
@test Utils.read_stdout(CLI.routes) == """
GET  /articles      ArticleController  index
GET  /articles/:id  ArticleController  show"""
Routing.empty!()

routes() do
    resources("/articles", ArticleController, except=[new, edit, create, delete, update])
end
@test Utils.read_stdout(CLI.routes) == """
GET  /articles      ArticleController  index
GET  /articles/:id  ArticleController  show"""
Routing.empty!()

end # module test_bukdu_resources

max_stars_repo_name: wookay/TestBukdu

max_issues_repo_name: wookay/TestBukdu

max_forks_repo_name: wookay/TestBukdu


text:
module DiskArrayTools
import DiskArrays: AbstractDiskArray, eachchunk, haschunks, Chunked,
estimate_chunksize, GridChunks, findints, readblock!, writeblock!
using Interpolations
using IterTools: imap
using Base.Iterators: product
using OffsetArrays: OffsetArray

#Define a new chunk grid type for irregular, but still block-shaped grid
struct IrregularGridChunks{N} <: AbstractArray{CartesianIndices{N,NTuple{N,UnitRange{Int}}},N}
    cs::NTuple{N,Vector{UnitRange{Int}}}
    s::NTuple{N,Int}
end
Base.size(cs::IrregularGridChunks) = cs.s
Base.size(cs::IrregularGridChunks,i) = cs.s[i]
function IrregularGridChunks(cs::Vector{UnitRange{Int}}...)
    s = length.(cs)
    IrregularGridChunks((cs...,),s)
end
Base.eltype(::IrregularGridChunks{N}) where N = CartesianIndices{N,NTuple{N,UnitRange{Int}}}
function Base.show(io::IO,g::IrregularGridChunks)
    stot = last.(g.cs)
    griddes = join(string.(g.s), "x")
    sizedes = join(string.(stot), "x")
    print(io,"Irregular chunk grid of size $griddes over DiskArray of size $stot")
end
function Base.getindex(cs::IrregularGridChunks{N},i::Vararg{Int, N}) where N
    rout = getindex.(cs.cs,i)
    CartesianIndices(rout)
end

export DiskArrayStack, diskstack, ConcatDiskArray, CFDiskArray
struct DiskArrayStack{T,N,M,NO}<:AbstractDiskArray{T,N}
    arrays::Array{M,NO}
end
function diskstack(a::Array{M,N}) where {N,M<:AbstractArray{T,NO}} where {T,NO}
    length(unique(size.(a))) == 1 || error("All arrays in the stack must have the same size")
    DiskArrayStack{T,N+NO,M,N}(a)
end
Base.size(r::DiskArrayStack) = (size(r.arrays[1])...,size(r.arrays)...)
haschunks(a::DiskArrayStack) = haschunks(a.arrays[1])
iscompressed(a::DiskArrayStack) = any(iscompressed,a.arrays)
function eachchunk(a::DiskArrayStack{<:Any,<:Any,<:Any,NO}) where NO
    iterold = eachchunk(a.arrays[1])
    cs = (iterold.chunksize...,ntuple(one,NO)...)
    co = (iterold.offset...,ntuple(zero,NO)...)
    GridChunks(a,cs,offset=co)
end

function readblock!(a::DiskArrayStack{<:Any,N,<:Any,NO},aout,i::AbstractVector...) where {N,NO}

  innerinds = i[1:(N-NO)]

  outerinds = i[(N-NO+1):N]
  innercolon = map(_->(:), innerinds)
  iiter = CartesianIndices(outerinds)
  inum  = CartesianIndices(size(iiter))
  foreach(zip(iiter,inum)) do (iouter,iret)
    arnow = a.arrays[iouter]
    if isa(arnow, AbstractDiskArray)
      readblock!(a.arrays[iouter],view(aout,innercolon...,iret.I...),innerinds...)
    else
      aout[innercolon...,iret.I...] = arnow[innerinds...]
    end
  end
  nothing
end

function writeblock!(a::DiskArrayStack{<:Any,N,<:Any,NO},v,i::AbstractVector...) where {N,NO}
  innerinds = i[1:(N-NO)]

  outerinds = i[(N-NO+1):N]
  innercolon = map(_->(:), innerinds)
  iiter = CartesianIndices(outerinds)
  inum  = CartesianIndices(size(iiter))
  foreach(zip(iiter,inum)) do (iouter,iret)
    arnow = a.arrays[iouter]
    if isa(arnow, AbstractDiskArray)
      writeblock!(a.arrays[iouter],view(v,innercolon...,iret.I...),innerinds...)
    else
      arnow[innerinds...] = v[innercolon...,iret.I...]
    end
  end
  nothing
end

function Base.view(a::DiskArrayStack{<:Any,N,<:Any,NO},i...) where {N,NO}
    iinner = i[1:(N-NO)]
    ashort = map(view(a.arrays,i[(N-NO+1):N]...)) do ai
        view(ai,iinner...)
    end
    if ndims(ashort)==0
      ashort[]
    else
      diskstack(ashort)
    end
end

abstract type ResampledDiskArray{T,N} <: AbstractDiskArray{T,N} end

function readblock!(a::ResampledDiskArray, aout, i::AbstractUnitRange...)
  parentranges = map((ir,ip)->ip[ir],i,a.newinds)
  rr = get_readinds(a,parentranges)
  atemp = a.a[rr...]
  atemp2 = OffsetArray(atemp,map(r->(first(r)-1),rr))
  resample_disk(a,aout,atemp2,parentranges)
end

struct InterpolatedDiskArray{T,N,A<:AbstractArray{T,N},I,O,BC,CS<:Union{Nothing,GridChunks{N}}} <: ResampledDiskArray{T,N}
    a::A
    newinds::I
    meth::O
    bc::BC
    chunksize::CS
end
fixin(i,s) = max(1,min(s,i))
function round_readinds(s,r,an)
  mi,ma = extrema(r)
  if mi == ma
    m = fixin(round(Int,mi),s)
    m:m
  else
    a = if an
      (round(Int,mi,RoundNearestTiesUp), -round(Int,-ma,RoundNearestTiesUp))
    else
      (floor(Int,mi),ceil(Int,ma))
    end
    fixin(a[1],s):fixin(a[2],s)
  end
end

function get_readinds(a::InterpolatedDiskArray,r)
  allnearest = all(i->isa(i,Union{Nothing,BSpline{Constant}}),a.meth)
  round_readinds.(size(a.a),r,allnearest)
end

allmeths(order::Tuple,newinds) = map(order,newinds) do o,ni
    ni === nothing ? NoInterp() : BSpline(o)
end
function allmeths(order, newinds)
    allorders = map(newinds) do _
        order
    end
    allmeths(allorders,newinds)
end
function InterpolatedDiskArray(a::AbstractArray,chunksize,newinds...; order=Linear(), bc=Flat())
    eltype(a) <: Union{Missing,Real,Complex} || error("Can only interpolate real or complex values")
    s = size(a)
    ni2 = map((s,i)->i===nothing ? (1:s) : i,s,newinds)
    me = allmeths(order,newinds)
    InterpolatedDiskArray(a,ni2,me,bc,chunksize)
end
Base.size(a::InterpolatedDiskArray) = map(length,a.newinds)
haschunks(a::InterpolatedDiskArray{<:Any,<:Any,<:Any,<:Any,<:Any,<:Any,<:GridChunks}) = true
haschunks(a::InterpolatedDiskArray{<:Any,<:Any,<:Any,<:Any,<:Any,<:Any,Nothing}) = false
eachchunk(a::InterpolatedDiskArray{<:Any,<:Any,<:Any,<:Any,<:Any,<:Any,<:GridChunks}) = a.chunksize

function resample_disk(a::InterpolatedDiskArray,aout,atemp,parentranges)
  meth = map((s,m)->s==1 ? NoInterp() : m,size(atemp),a.meth)
  if all(isequal(NoInterp()),meth)
      aout .= atemp.parent
  else
      fremap(aout,atemp,parentranges,meth,bc=a.bc)
  end
end

function fremap(xout,xin,newinds,ipmeth;bc = Flat())
  interp = extrapolate(interpolate(xin, ipmeth), bc)
  for (i,ic) in enumerate(Iterators.product(newinds...))
    xout[i]=interp(ic...)
  end
end

#Use of Sentinel missing value
struct CFDiskArray{T,N,ST,P<:AbstractArray{ST,N}} <: AbstractDiskArray{Union{T,Missing},N}
    a::P
    mv::ST
    add_offset::T
    scale_factor::T
end
function CFDiskArray(a::AbstractArray{T}, attr::Dict) where T
  mv = get(attr,"missing_value",get(attr,"_FillValue",typemax(T)))
  offs,sc = if haskey(attr,"add_offset") || haskey(attr,"scale_factor")
    offs = get(attr,"add_offset",zero(Float16))
    sc = get(attr,"scale_factor",one(Float16))
    promote(offs,sc)
  else
    zero(T), one(T)
  end
  CFDiskArray(a, T(mv), offs, sc)
end

Base.size(a::CFDiskArray, args...) = size(a.a, args...)
haschunks(a::CFDiskArray) = haschunks(a.a)
eachchunk(a::CFDiskArray) = eachchunk(a.a)
iscompressed(a::CFDiskArray) = iscompressed(a.a)

function readblock!(a::CFDiskArray, aout, r::AbstractVector...)
    mv = a.mv
    sc,offs = a.scale_factor, a.add_offset
    broadcast!(j->scaleoffs(j,mv,sc,offs), aout, view(a.a,r...))
    nothing
end
scaleoffs(x,mv,sc,offs) = x==mv ? missing : x*sc+offs
function writeblock!(a::CFDiskArray, v, r::AbstractVector...)
    mv = a.mv
    sc,offs = a.scale_factor, a.add_offset
    broadcast!(j->scaleoffsinv(j,mv,sc,offs), view(a.a,r...), v)
    nothing
end
scaleoffsinv(x,mv::Integer,sc,offs) = ismissing(x) ? mv : round(typeof(mv),(x-offs)/sc)
scaleoffsinv(x,mv,sc,offs) = ismissing(x) ? mv : ((x-offs)/sc)

struct ConcatDiskArray{T,N,P} <: AbstractDiskArray{T,N}
    parents::P
    startinds::NTuple{N,Vector{Int}}
    size::NTuple{N,Int}
end

function ConcatDiskArray(arrays::AbstractArray)
    #First do some consistency checks
    arrinfo = map(arrays) do a
        eltype(a), size(a)
    end
    et = first(arrinfo)[1]
    nd = length(first(arrinfo)[2])
    sa = size.(arrays)
    all(i->i[1]==et,arrinfo) || error("Arrays don't have the same element type")
    all(i->length(i[2])==nd, arrinfo) || error("Arrays don't have the same dimensions")
    if ndims(arrays)<nd
        arrays = reshape(arrays,size(arrays)...,ones(Int,nd-ndims(arrays))...)
    end
    function othersize(x,id)
        (x[1:id-1]..., x[id+1:end]...)
    end
    si = map(1:nd) do id
        otherdims = ((1:id-1)...,(id+1:nd)...)
        a = size.(arrays)
        a = reduce(a, dims=id, init=ntuple(zero,nd)) do i,j
                if all(iszero,i)
                    j
                elseif othersize(i,id)==othersize(j,id)
                    j
                else
                    error("Dimensions don't match")
                end
        end
        I = fill!(Array{Union{Int,Colon},1}(undef,nd),1)
        I[id] = Colon()
        ar = sa[I...]
        ari = map(i->i[id],ar)
        sl = sum(ari)
        r = cumsum(ari)
        pop!(pushfirst!(r,0))
        r.+1, sl
    end
    ConcatDiskArray{et,nd,typeof(arrays)}(arrays, (getindex.(si,1)...,),(getindex.(si,2)...,))
end

Base.size(a::ConcatDiskArray) = a.size
haschunks(::ConcatDiskArray) = Chunked()
function readblock!(a::ConcatDiskArray, aout, inds::AbstractUnitRange...)
    #Find affected blocks and indices in blocks
    blockinds  = map(inds, a.startinds, size(a.parents)) do i, si, s
        bi1 = max(searchsortedlast(si,first(i)),1)
        bi2 = min(searchsortedfirst(si,last(i)+1)-1,s)
        bi1:bi2
    end
    map(CartesianIndices(blockinds)) do cI
        myar = a.parents[cI]
        mysize = size(myar)
        array_range = map(cI.I, a.startinds, mysize, inds) do ii, si, ms, indstoread
            max(first(indstoread)-si[ii]+1,1) : min(last(indstoread)-si[ii]+1, ms)
        end
        outer_range = map(cI.I, a.startinds, array_range, inds) do ii, si, ar, indstoread
            (first(ar)+si[ii]-first(indstoread)):(last(ar)+si[ii]-first(indstoread))
        end
        aout[outer_range...] .= a.parents[cI][array_range...]
    end
end
function writeblock!(a::ConcatDiskArray, aout, inds::AbstractUnitRange...)
  error("No method yet for writing into a ConcatDiskArray")
end

function eachchunk_fallback(aconc::ConcatDiskArray)
  nested = imap(zip(aconc.parents, Iterators.product(aconc.startinds...))) do (par, si)
      imap(eachchunk(par)) do cI
          cI .+ CartesianIndex(si.-1)
      end
  end
  Iterators.flatten(nested)
end

#This function will work for DiskArrayStack as well, we only need a fallback_chunksize function
function eachchunk(aconc::ConcatDiskArray)
  fbchunks = eachchunk_fallback(aconc)
  nd = ndims(first(fbchunks))
  dout = ntuple(_->Dict{NTuple{nd-1,UnitRange{Int}},Set{UnitRange{Int}}}(),nd)
  foreach(fbchunks) do c
    map(enumerate(c.indices)) do (ind,r)
        other = (c.indices[1:ind-1]...,c.indices[ind+1:end]...)
        s = get!(dout[ind],other) do
            Set{UnitRange{Int}}()
        end
        push!(s,c.indices[ind])
    end
  end
  #First check if all
  uniqueranges = unique.(values.(dout)) # Here we give up
  if any(length.(uniqueranges) .> 1)
    return fbchunks
  end
  simplifyaxes = map(uniqueranges) do ur
    allr = first(ur)
    allr = sort!(collect(allr),by=first)
    l = length.(allr)
    if length(l)==1
        allr,(cs = l[1], offs = 0, l = l[1])
    elseif all(isequal(l[2]),l[2:end-1]) && l[end]<=l[2]
        allr,(cs = l[2], offs = l[2]-l[1], l = allr[end][end])
    else
        allr,nothing
    end
  end
  if all(!isnothing,getindex.(simplifyaxes,2))
    grid = getindex.(simplifyaxes,2)
    cs,offs,l = getproperty.(grid,:cs),getproperty.(grid,:offs),getproperty.(grid,:l)
    return GridChunks(l,cs, offset = offs)
  else
    return IrregularGridChunks(getindex.(simplifyaxes,1)...)
  end
end
end # module


text:
using NCDatasets

"""
   ProfileNative

Container for a multivariate profile read from a GDAC Argo file.

- 1D arrays: lon,lat,date,ymd,hms,pnum_txt,direc,DATA_MODE,isBAD
- 2D arrays: T,S,pressure,depth,T_ERR,SERR
"""
struct ProfileNative
    lon
    lat
    date
    ymd
    hms
    T
    S
    pressure
    depth
    T_ERR
    S_ERR
    pnum_txt
    direc
    DATA_MODE
    isBAD
end

"""
   ProfileNative

Container for a multivariate profile in MITprof format.

- 2D arrays: T,S,Testim,Sestim,Tweight,Sweight,Ttest,Stest,T_ERR,SERR
"""
struct ProfileStandard
    T
    Testim
    Tweight
    Ttest
    T_ERR
    S
    Sestim
    Sweight
    Stest
    S_ERR
end

ProfileStandard(nz::Int) = ProfileStandard(
    Array{Union{Float64,Missing},1}(undef,nz),
    Array{Union{Float64,Missing},1}(undef,nz),
    Array{Union{Float64,Missing},1}(undef,nz),
    Array{Union{Float64,Missing},1}(undef,nz),
    Array{Union{Float64,Missing},1}(undef,nz),
    Array{Union{Float64,Missing},1}(undef,nz),
    Array{Union{Float64,Missing},1}(undef,nz),
    Array{Union{Float64,Missing},1}(undef,nz),
    Array{Union{Float64,Missing},1}(undef,nz),
    Array{Union{Float64,Missing},1}(undef,nz)
)

"""
   MITprofStandard

Container for a MITprof format file data.

- filename : file name
- 1D arrays: lon,lat,date,depth,ID
- 2D arrays: T,S,Te,Se,Tw,Sw
"""
struct MITprofStandard
    filename
    lon
    lat
    date
    depth
    ID
    #
    T
    Te
    Tw
    S
    Se
    Sw
end

"""
   MITprofStandard

Create a `MITprofStandard` view of an MITprof file.

```
fil="nc/5903955_MITprof.nc"
mp=MITprof.MITprofStandard(fil)
```
"""
function MITprofStandard(fil::String)
    ds=Dataset(fil)
    haskey(ds,"prof_descr") ? ID=ds["prof_descr"] : ID=split(basename(fil),"_")[1]
    MITprofStandard(fil,
        ds["prof_lon"],ds["prof_lat"],ds["prof_date"],ds["prof_depth"],
        fill(ID,size(ds["prof_lon"])),
        ds["prof_T"],ds["prof_Testim"],ds["prof_Tweight"],
        ds["prof_S"],ds["prof_Sestim"],ds["prof_Sweight"]
        )
end

function Base.show(io::IO, mp::MITprofStandard) where {T}
    printstyled(io, "File name is ",color=:normal)
    printstyled(io, "$(basename(mp.filename)) \n",color=:blue)
    printstyled(io, "List of variables : \n",color=:normal)
    #
    printstyled(io, "  lon   is ",color=:normal)
    printstyled(io, "$(size(mp.lon)) \n",color=:blue)
    printstyled(io, "  lat   is ",color=:normal)
    printstyled(io, "$(size(mp.lat)) \n",color=:blue)
    printstyled(io, "  date  is ",color=:normal)
    printstyled(io, "$(size(mp.date)) \n",color=:blue)
    printstyled(io, "  depth is ",color=:normal)
    printstyled(io, "$(size(mp.depth)) \n",color=:blue)
    printstyled(io, "  ID    is ",color=:normal)
    printstyled(io, "$(size(mp.ID)) \n",color=:blue)
    #
    printstyled(io, "  T     is ",color=:normal)
    printstyled(io, "$(size(mp.T)) \n",color=:blue)
    printstyled(io, "  Te    is ",color=:normal)
    printstyled(io, "$(size(mp.Te)) \n",color=:blue)
    printstyled(io, "  Tw    is ",color=:normal)
    printstyled(io, "$(size(mp.Tw)) \n",color=:blue)
    printstyled(io, "  S     is ",color=:normal)
    printstyled(io, "$(size(mp.S)) \n",color=:blue)
    printstyled(io, "  Se    is ",color=:normal)
    printstyled(io, "$(size(mp.Se)) \n",color=:blue)
    printstyled(io, "  Sw    is ",color=:normal)
    printstyled(io, "$(size(mp.Sw))",color=:blue)
    #
    return
end


text:
# ---
# title: 335. Self Crossing
# id: problem335
# author: Tian Jun
# date: 2020-10-31
# difficulty: Hard
# categories: Math
# link: <https://leetcode.com/problems/self-crossing/description/>
# hidden: true
# ---
#
# You are given an array _x_ of `n` positive numbers. You start at point `(0,0)`
# and moves `x[0]` metres to the north, then `x[1]` metres to the west, `x[2]`
# metres to the south, `x[3]` metres to the east and so on. In other words,
# after each move your direction changes counter-clockwise.
#
# Write a one-pass algorithm with `O(1)` extra space to determine, if your path
# crosses itself, or not.
#
#
#
# **Example 1:**
#
#
#
#     **┌───┐
#     │   │
#     └───┼──>
#         │
#
#     Input: **[2,1,1,2]
#     Output: true
#
#
# **Example 2:**
#
#
#
#     **┌──────┐
#     │      │
#     │
#     │
#     └────────────>
#
#     Input:** [1,2,3,4]
#     Output: false
#
#
# **Example 3:**
#
#
#
#     **┌───┐
#     │   │
#     └───┼>
#
#     Input:** [1,1,1,1]
#     Output: true
#
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


text:
function encode(s)
    (length(s) == 0) && (return s)

    count = 0
    for char in s
        (char == s[1]) ? (count += 1) : break
    end
    count_str = (count == 1) ? "" : dec(count)
    return string(count_str, s[1], encode(s[count + 1 : end]))
end

function decode(s)
    function expand(count_str, char)
        count = (count_str == "") ? 1 : parse(Int, count_str)
        return char ^ count
    end

    result = ""
    return string((expand(rm.captures...) for rm in eachmatch(r"(\d*)(\D)", s))...)
end


text:
n = 10 # rand(1:10)
A = matrixdepot("erdrey", n)
B = matrixdepot("erdrey", n, ceil(Int, n*log(n)/2))
@test nnz(A) == nnz(B)
@test issymmetric(A)

C = matrixdepot("erdrey", 10, 3)
@test nnz(C) == 6
println("'erdrey' passed test...")


text:
function AbstractPlotting.plot!(
        lscene::LScene, P::AbstractPlotting.PlotFunc,
        attributes::AbstractPlotting.Attributes, args...;
        kw_attributes...)

    plot = AbstractPlotting.plot!(lscene.scene, P, attributes, args...; kw_attributes...)[end]

    plot
end

protrusionnode(ls::LScene) = ls.layoutobservables.protrusions
reportedsizenode(ls::LScene) = ls.layoutobservables.reportedsize

function LScene(parent::Scene; bbox = nothing, scenekw = NamedTuple(), kwargs...)

    default_attrs = default_attributes(LScene, parent).attributes
    theme_attrs = subtheme(parent, :LScene)
    attrs = merge!(merge!(Attributes(kwargs), theme_attrs), default_attrs)

    layoutobservables = LayoutObservables(LScene, attrs.width, attrs.height, attrs.tellwidth, attrs.tellheight,
        attrs.halign, attrs.valign, attrs.alignmode; suggestedbbox = bbox)

    scene = Scene(parent, lift(IRect2D_rounded, layoutobservables.computedbbox); scenekw...)

    LScene(scene, attrs, layoutobservables)
end

max_stars_repo_name: jkrumbiegel/MakieLayout


text:
# MIT license
# Copyright (c) Microsoft Corporation. All rights reserved.
# See LICENSE in the project root for full license information.

@testset "ParaxialLens" begin
    @testset "Virtual point" begin
        lens = ParaxialLensRect(10.0,100.0,100.0,[0.0,0.0,1.0],[0.0,0.0,0.0])
        displaypoint = [0.0,0.0,-8.0]
        r1 = Ray(displaypoint,[1.0,0.0,1.0])
        intsct = OpticSim.surfaceintersection(lens,r1)
        intsctpt = point(intsct)
        # compute the refracted ray and project this backwards to its intersection with the optical axis. This should be the same position as returned by virtualpoint()
        refrac,_,_ = OpticSim.processintersection(OpticSim.interface(lens),OpticSim.point(intsct),OpticSim.normal(lens),OpticalRay(r1,1.0,.55),OpticSim.TEMP_REF,OpticSim.PRESSURE_REF,false)
        # compute intersection of ray with optical axis. this should match the position of the virtual point
        slope = refrac[1]/refrac[3]
        virtptfromslope = [0.0,0.0,-intsctpt[1]/slope]
        # insert test virtptfromslope == point(virtualpoint(lens,displaypoint))
        @test isapprox(virtptfromslope,point(virtualpoint(lens,displaypoint)))
    end

    function rayintersection(lens,incomingray)
        intsct = OpticSim.surfaceintersection(lens,incomingray)
        refrac,_,_ = OpticSim.processintersection(OpticSim.interface(lens),OpticSim.point(intsct),OpticSim.normal(lens),OpticalRay(incomingray,1.0,.55),OpticSim.TEMP_REF,OpticSim.PRESSURE_REF,false)
        return refrac
    end

    @testset "Refracted rays" begin
         #test all 4 combinations: n⋅r > 0, n⋅r < 0, focal length > 0, focal length < 0
        #n⋅r > 0 focal length > 0
        lens = ParaxialLensRect(1.0,100.0,100.0,[0.0,0.0,1.0],[0.0,0.0,0.0])
        r = Ray([1.0,0.0,-1.0],[0.0,0.0,1.0])
        refrac = rayintersection(lens,r)
        @test isapprox([-sqrt(2)/2,0.0,sqrt(2)/2],refrac)

        #n⋅r < 0 focal length > 0
        r = Ray([1.0,0.0,1.0],[0.0,0.0,-1.0])
        refrac = rayintersection(lens,r)
        @test isapprox([-sqrt(2)/2,0.0,-sqrt(2)/2],refrac)

        #n⋅r > 0 focal length < 0
        lens = ParaxialLensRect(-1.0,100.0,100.0,[0.0,0.0,1.0],[0.0,0.0,0.0])
        r = Ray([1.0,0.0,-1.0],[0.0,0.0,1.0])
        refrac = rayintersection(lens,r)
        @test isapprox([sqrt(2)/2,0.0,sqrt(2)/2],refrac)

        #n⋅r < 0 focal length < 0
        r = Ray([1.0,0.0,1.0],[0.0,0.0,-1.0])
        refrac = rayintersection(lens,r)
        @test isapprox([sqrt(2)/2,0.0,-sqrt(2)/2],refrac)
    end

    @testset "Reversibility of rays for paraxial lenses" begin
        lens = ParaxialLensRect(10.0,100.0,100.0,[0.0,0.0,1.0],[0.0,0.0,0.0])
        displaypoint = [0.0,0.0,-8.0]
        direction = [0.0,0.0,1.0]
        r1 = Ray(displaypoint,direction)
        intsct = OpticSim.surfaceintersection(lens,r1)
        intsctpt = point(intsct)
        # compute the refracted ray and project this backwards to its intersection with the optical axis. This should be the same position as returned by virtualpoint()
        refrac1,_,_ = OpticSim.processintersection(OpticSim.interface(lens),OpticSim.point(intsct),OpticSim.normal(lens),OpticalRay(r1,1.0,.55),OpticSim.TEMP_REF,OpticSim.PRESSURE_REF,false)
        r2 = Ray(-displaypoint,-direction)

        intsct = OpticSim.surfaceintersection(lens,r2)
        intsctpt = point(intsct)
        # compute the refracted ray and project this backwards to its intersection with the optical axis. This should be the same position as returned by virtualpoint()
        refrac2,_,_ = OpticSim.processintersection(OpticSim.interface(lens),OpticSim.point(intsct),OpticSim.normal(lens),OpticalRay(r2,1.0,.55),OpticSim.TEMP_REF,OpticSim.PRESSURE_REF,false)

        @test isapprox(-refrac2,refrac1)
    end
end


text:
using EffectiveWaves

## Choose two species randomly (uniformly) distributed in the medium
# Usage Specie(ρ = density, r = radius, c = wavespeed, volfrac = volume fraction)
species = [
    Specie(ρ=WaterDistilled.ρ,r=30.e-6, c=WaterDistilled.c, volfrac=0.1),
    Specie(ρ=Inf, r=100.0e-6, c=2.0, volfrac=0.2)
]
# background medium
background = Glycerol # for other materials, see materials.jl

# angular frequencies
ωs = LinRange(0.01,1.0,60)*30.0e6

# using only one effective wavenumber per angular frequencies. This wavenumber is calculated from an asymptotic formula for low particle volume fraction
k_effs = wavenumber_low_volfrac(ωs, background, species)

speeds = ωs./real(k_effs)
attenuations = imag(k_effs)

using Plots
pyplot()

p1 = plot(ωs./real(background.c), speeds,  ylabel="wave speed (m/s)", xlabel="k (1/m)");
p2 = plot(ωs./real(background.c), attenuations, ylabel="attenuation (1/m)", xlabel="k (1/m)");
plot(p1,p2)

## An example where we vary the species

height=450
# unicodeplot() # alternative plot backend
pyplot(linewidth=3, size=(2*height,height))
Plots.scalefontsizes(1.5)

# for fixed total volfrac fraction
background = Medium(ρ=1.0,c = 1.0)
ωs = LinRange(0.01,1.0,60)
volfrac = 0.25
r1 = 0.5
r2 = 1.5
N=5
vols = LinRange(0.0,1.0,N)*volfrac

kTs_arr = [
  begin
    sp1 = Specie(0.0, r1; volfrac=vols[i])
    sp2 = Specie(Inf, r1; volfrac=volfrac-vols[i])
    [ wavenumber_low_volfrac(ω, background, [sp1,sp2]) for ω in ωs]
  end
for i = 1:N];

speeds = [ ωs./real(kTs) for kTs in kTs_arr]
attenuations = imag(kTs_arr)

labs = reshape( map(v -> "void vol = $(Int(round(100*v)))%",vols),1, length(vols));
p1 = plot(ωs./real(background.c), speeds, labels=labs, ylabel="wave speed (m/s)" ,xlabel="k (1/m)");
p2 = plot(ωs./real(background.c), attenuations, labels=labs, xlabel="k (1/m)", ylabel="attenuation (1/m)");

plot(p1,p2)

Plots.scalefontsizes(1./1.5)


text:
import FireflyModule.Firefly

function move_firefly(f1::Firefly, f2::Firefly, r:Int)
    """Moving f1 to f2 less than r swap operations.
    """
    number_of_swaps = Int(ceil(rand() * (r-2) ))
    d, d_info = hamming_distance(f1, f2)

    while number_of_swaps > 0
        d, d_info = hamming_distance(f1, f2)
        random_index = rand([i for i in 1:length(d_info) if d_info[i]])
        value_to_copy = f2.path[random_index]
        index_to_move = findfirst(x -> x==value_to_copy, f1.path)

        if number_of_swaps == 1 && f1.path[index_to_move] == f2.path[random_index] && f1.path[random_index] == f2.path[index_to_move]
            break
        end

        f1.path[random_index], f1.path[index_to_move] = f1.path[index_to_move], f1.path[random_index]
        # println("swapped! ", f1.path[random_index], " -- ", f1.path[index_to_move])
        if f1.path[index_to_move] == f2.path[index_to_move] number_of_swaps -= 1 end
        number_of_swaps -= 1
    end
end

max_stars_repo_name: zahidkizmaz/firefly-algorithm-tsp

max_issues_repo_name: zahidkizmaz/firefly-algorithm-tsp

max_forks_repo_name: zahidkizmaz/firefly-algorithm-tsp


text:
"""
    SeisAnimation(path, d; <keyword arguments>)

make animation of a 3D cube

# Arguments
- `path::String`: path to save the generated animation.
- `d::Array{Real,3}`: 3D data to plot.

# Keyword arguments

- `wbox=6`: width of figure in inches.
- `hbox=6`: height of figure in inches.
- `dpi=100`: dots-per-inch of figure.

- `cmap="gray"`: colormap
- `pclip=98`: percentile for determining clip.
- `vmin="NULL"`: minimum value used in colormapping data.
- `vmax="NULL"`: maximum value used in colormapping data.
- `aspect="auto"`: image aspect ratio.
- `interpolation="Hanning"`: interpolation method for colormapping data.
- `title=" "`: title of plot.
- `titlesize=16`: size of title.
- `xlabel=" "`: label on x-axis.
- `ylabel=" "`: label on x-axis.
- `labelsize=14`: size of labels on axis.

- `ox=0`: first point of x-axis.
- `dx=1`: increment of x-axis.
- `oy=0`: first point of y-axis.
- `dy=1`: increment of y-axis.

- `xticks="NULL"`: ticks on x-axis.
- `yticks="NULL"`: ticks on y-axis.
- `ticksize=11`: size of labels on ticks.

- `jump=1`: number of slices jumped when making animation.
- `interal=200`: delay for each frames, default 200 ms.

# Example
```julia
julia> path = joinpath(homedir(), "Desktop/animation")
julia> SeisAnimation(path, randn(100,100,10); wbox=6, hbox=4, cmap="gray", vmax=1, vmin=-1,
                     ox=10, dx=10, oy=2, dy=2, title="wavefield", xlabel="X (m)", ylabel="Z (m)", jump=1, interval=400)
```
"""
function SeisAnimation(pathout::Ts, d::Array{Tv,3}; wbox=6, hbox=6, dpi=100,
                       cmap="gray", vmin="NULL", vmax="NULL", pclip=98,
                       ox=0, dx=1, oy=0, dy=1, aspect="auto", interpolation="Hanning",
                       title=" ", titlesize=16, xlabel=" ", ylabel=" ", labelsize=14,
                       xticks="NULL", yticks="NULL", ticksize=11, jump=1,
                       interval=200) where {Ts<:String, Tv<:AbstractFloat}

    # dimensions of data
    (n1, n2, n3) = size(d)

    # open canvas
    fig = pl.figure(figsize=(wbox,hbox), dpi=dpi, facecolor="w", edgecolor="k")

    # generate figure for each frontal slice
    ims = PyCall.PyObject[]
    for i = 1 : jump : n3

        s = d[:,:,i]

        # dynamic range
        if (vmin=="NULL" || vmax=="NULL")
           if (pclip<=100)
    	        a = -quantile(abs.(s[:]), (pclip/100))
    	     else
    	        a = -quantile(abs.(s[:]), 1)*pclip/100
    	     end
    	     b = -a
        else
    	     a = vmin
    	     b = vmax
        end

        # generate figure
        image = pl.imshow(s, cmap=cmap, vmin=a, vmax=b,
                          extent=[ox, ox+(size(s,2)-1)*dx, oy+(size(s,1)-1)*dy, oy],
                          aspect=aspect, interpolation=interpolation)

        # tile and ticks
        pl.title(title, fontsize=titlesize)
        pl.xlabel(xlabel, fontsize=labelsize)
        pl.ylabel(ylabel, fontsize=labelsize)

        xticks == "NULL" ? nothing : pl.xticks(xticks)
        yticks == "NULL" ? nothing : pl.yticks(yticks)
        ax = pl.gca()
        pl.setp(ax.get_xticklabels(), fontsize=ticksize)
        pl.setp(ax.get_yticklabels(), fontsize=ticksize)
        tight_layout()

        # save as one element of the array
        push!(ims, PyCall.PyObject[image])
    end

    # make a animation from vector of image
    ani = anim.ArtistAnimation(fig, ims, interval=interval, blit=true)

    # save the result as mp4 file
    pathout = join([pathout ".mp4"])
    ani.save(pathout)

    return nothing
end


text: Dict{String, Any}("albums" => Any[Dict{String, Any}("name" => Dict{String, Any}("value" => "Born to Run", "type" => "string"), "songs" => Any[Dict{String, Any}("name" => Dict{String, Any}("value" => "Jungleland", "type" => "string")), Dict{String, Any}("name" => Dict{String, Any}("value" => "Meeting Across the River", "type" => "string"))]), Dict{String, Any}("name" => Dict{String, Any}("value" => "Born in the USA", "type" => "string"), "songs" => Any[Dict{String, Any}("name" => Dict{String, Any}("value" => "Glory Days", "type" => "string")), Dict{String, Any}("name" => Dict{String, Any}("value" => "Dancing in the Dark", "type" => "string"))])])


text:
###
### tbegin/tend
###

struct TBeginEnd{F}
    "maps `(firstindex(::Tensor, d), lastindex(::Tensor, d))` to an index"
    f::F
    ast
end

for (tend, _tend, l) in [(:tbegin, :_tbegin, :f), (:tend, :_tend, :l)]
    @eval begin
        $_tend(f, l) = $l
        const $tend = TBeginEnd($_tend, $(Meta.quot(tend)))
    end
end

Base.show(io::IO, tend::TBeginEnd) = print(io, tend.ast)

macro _tbeginend(ex)
    @assert Meta.isexpr(ex, :call)
    showexpr = Expr(
        :call,
        ex.args[1],
        [
            if Meta.isexpr(i, :call) && Meta.isexpr(i.args[1], :.)
                Expr(:$, Expr(:., i.args[1].args[1], Meta.quot(:ast)))
            else
                Expr(:$, i)
            end
            for i in ex.args[2:end]
        ]...,
    )
    :(TBeginEnd((f, l) -> $ex, $(Meta.quot(showexpr))))
end

for f in [:(:), :+, :-, :*, :min, :max, :%, :mod, :mod1, :div, :fld]
    @eval begin
        Base.$f(x::TBeginEnd, y::Int) = @_tbeginend $f(x.f(f, l), y)
        Base.$f(x::Int, y::TBeginEnd) = @_tbeginend $f(x, y.f(f, l))
        Base.$f(x::TBeginEnd, y::TBeginEnd) = @_tbeginend $f(x.f(f, l), y.f(f, l))
    end
end

Base.:(:)(x::TBeginEnd, s::Int, y::Int) = @_tbeginend x.f(f, l):s:y
Base.:(:)(x::Int, s::Int, y::TBeginEnd) = @_tbeginend x:s:y.f(f, l)
Base.:(:)(x::TBeginEnd, s::Int, y::TBeginEnd) = @_tbeginend x.f(f, l):s:y.f(f, l)

Base.to_index(i::TBeginEnd) = i
expanded_dims(::Type{<:TBeginEnd}) = 1

expand_tbeginend(t, dim, x::TBeginEnd) = x.f(firstindex(t, dim), lastindex(t, dim))
expand_tbeginend(t, _, x) = x

function expand_tbeginends(t, inds::Tuple)
    ntuple(i -> expand_tbeginend(t, i, inds[i]), length(inds))
end
function expand_tbeginends(t, ti::TIndex{N,ipos}) where {N,ipos}
    TIndex{N,ipos}(expand_tbeginends(t, ti.I))
end


text:
using Einsum

function get_Kσs(sp::StiffnessTopOptProblem{xdim,TT}, u_dofs, cellvalues) where {xdim,TT}
    E = getE(sp)
    ν = getν(sp)
    dh = sp.ch.dh
    # usually ndof_pc = xdim * n_basefuncs
    # ? number of nodes per cell == n_basefuncs per cell
    ndof_pc = ndofs_per_cell(dh)
    n_basefuncs = getnbasefunctions(cellvalues)
    global_dofs = zeros(Int, ndof_pc)
    Kσs = [zeros(TT, ndof_pc, ndof_pc) for i in 1:getncells(dh.grid)]
    Kσ_e = zeros(TT, ndof_pc, ndof_pc)
    # block-diagonal - block σ_e = σ_ij, i,j in xdim
    # ! shouldn't this be xdim*xdim by xdim*xdim?
    # ? ψ_e = zeros(TT, xdim*ndof_pc, xdim*ndof_pc)
    ψ_e = zeros(TT, xdim * xdim, xdim * xdim)
    # ? G = zeros(TT, xdim*ndof_pc, ndof_pc)
    G = zeros(TT, xdim * xdim, xdim * n_basefuncs)
    δ = Matrix(TT(1.0)I, xdim, xdim)
    ϵ = zeros(TT, xdim, xdim)
    σ = zeros(TT, xdim, xdim)
    # u_i,j: partial derivative
    u_p = zeros(TT, xdim, xdim)
    for (cellidx, cell) in enumerate(CellIterator(dh))
        Kσ_e .= 0
        reinit!(cellvalues, cell)
        # get cell's dof's global dof indices, i.e. CC_a^e
        celldofs!(global_dofs, dh, cellidx)
        for q_point in 1:getnquadpoints(cellvalues)
            dΩ = getdetJdV(cellvalues, q_point)
            for d in 1:xdim
                ψ_e[((d - 1) * xdim + 1):(d * xdim), ((d - 1) * xdim + 1):(d * xdim)] .= 0
            end
            for a in 1:n_basefuncs
                ∇ϕ = shape_gradient(cellvalues, q_point, a)
                _u = @view u_dofs[(@view global_dofs[xdim * (a - 1) .+ (1:xdim)])]
                # u_i,j, i for spatial xdim, j for partial derivative
                @einsum u_p[i, j] = _u[i] * ∇ϕ[j]
                # effect of the quadratic term in the strain formula have on the stress field is ignored
                @einsum ϵ[i, j] = 1 / 2 * (u_p[i, j] + u_p[j, i])
                # isotropic solid
                @einsum σ[i, j] =
                    E * ν / (1 - ν^2) * δ[i, j] * ϵ[k, k] + E * ν * (1 + ν) * ϵ[i, j]
                for d in 1:xdim
                    # block diagonal
                    ψ_e[
                        ((d - 1) * xdim .+ 1):(d * xdim), ((d - 1) * xdim .+ 1):(d * xdim)
                    ] .+= σ
                    G[(xdim * (d - 1) + 1):(xdim * d), (a - 1) * xdim + d] .= ∇ϕ
                end
            end
            Kσ_e .+= G' * ψ_e * G * dΩ
        end
        Kσs[cellidx] .= Kσ_e
    end

    return Kσs
end

function buckling(problem::StiffnessTopOptProblem{xdim,T}, ginfo, einfo) where {xdim,T}
    dh = problem.ch.dh

    u = ginfo.K \ ginfo.f
    Kσs = get_Kσs(problem, u, einfo.cellvalues)
    Kσ = deepcopy(ginfo.K)

    if Kσ isa Symmetric
        Kσ.data.nzval .= 0
        assembler = Ferrite.AssemblerSparsityPattern(Kσ.data, T[], Int[], Int[])
    else
        Kσ.nzval .= 0
        assembler = Ferrite.AssemblerSparsityPattern(Kσ, T[], Int[], Int[])
    end

    # * assemble global geometric stiffness matrix
    global_dofs = zeros(Int, ndofs_per_cell(dh))
    Kσ_e = zeros(T, size(Kσs[1]))
    celliteratortype = CellIterator{typeof(dh).parameters...}
    _celliterator::celliteratortype = CellIterator(dh)
    TK = eltype(Kσs)
    for (i, cell) in enumerate(_celliterator)
        celldofs!(global_dofs, dh, i)
        if TK <: Symmetric
            Ferrite.assemble!(assembler, global_dofs, Kσs[i].data)
        else
            Ferrite.assemble!(assembler, global_dofs, Kσs[i])
        end
    end

    return ginfo.K, Kσ
end