text:
@inline function putindex(ids::NTuple{N,T},val_axis::Val{axis},i::T) where {T,N,axis}
    @inbounds A=ntuple(i->ids[i],Val(axis-1))
    @inbounds B=ntuple(i->ids[i+axis-1],Val(N-axis+1))
    (A...,i,B...)
end

@inline function extrudeindex(ids::NTuple{N},val_axis::Val{axis}) where {N,axis}
    @inbounds A=ntuple(i->ids[i],Val(axis-1))
    @inbounds B=ntuple(i->ids[i+axis],Val(N-axis))
    (A...,B...)
end

@inline function purge(x::R) where {N,T,R<:AbstractArray{T,N}}
    x
end
@inline function purge(x::R) where {T,R<:AbstractArray{T,0}}
    x[]
end
@inline function integrate(x::V, y::M, val_axis::Val{axis}) where {axis,Tm, Nm, M <: AbstractArray{Tm,Nm},Tv,Nv,V <: AbstractArray{Tv,Nv}}
    sy=size(y)
    N=length(x)
    @boundscheck begin
        @assert sy[axis]==N "Integration axis over `y` is incompatible with `x`. Make sure their length match!"
        @assert maximum(size(x))==N "`x` is not vector-like."
    end
    out_type = typeof(oneunit(Tm) * oneunit(Tv))
    res=zeros(out_type,extrudeindex(sy,val_axis))
    @inline idx(I,j) = CartesianIndex(putindex(Tuple(I),val_axis,j))
    N <= 1 && return purge(res)
    @simd for ids in CartesianIndices(res)
        @inbounds res[ids]+=(x[2]-x[1])*y[idx(ids,1)]
    end
    for i in 2:(N-1)
        @simd for ids in CartesianIndices(res)
            @inbounds res[ids]+=(x[i+1]-x[i-1])*y[idx(ids,i)]
        end
    end
    @simd for ids in CartesianIndices(res)
        @inbounds res[ids]+=(x[end]-x[end-1])*y[idx(ids,N)]
    end
    @simd for ids in eachindex(res)
        @inbounds res[ids]/=2
    end
    purge(res)
end

@inline function integrate(x::Colon, y,a)
    y
end
function integrate(x::AbstractArray{S}, y::AT) where {S,T,AT<:AbstractArray{T,0}}
    return zero(promote_type(T,S))
end

@inline @generated function integrate(x::NTuple{N,Union{Colon,AbstractArray}}, y::M) where {Tm, N, M <: AbstractArray{Tm,N}}
    axis=ntuple(i->Val(i),Val(N))
    I=:(y)
    for i in N:-1:1
        I=:(integrate(x[$i],$I,$(axis[i])))
    end
    I
end

@inline function integrate(x::V, y::M) where {Tv, Nv, V <: AbstractArray{Tv,Nv} ,Tm, Nm, M <: AbstractArray{Tm,Nm}}
    integrate(x,y,Val(Nm))
end

"""

    @trapz range variable expression
Calculates integral of [expression] over [variable] in [range]

#### Example
julia> @trapz 0:0.01:1 x x*x

0.33335

"""
macro trapz(range,var,expr)
    quote
        let
        local r=$(esc(range))
        N=length(r)
        @assert N>=2 "null integration range"
        @inline f($(esc(var)))=$(esc(expr))
        @inbounds begin
            local t = (r[2].-r[1]).*f(r[1])
            @simd for i in 2:(N-1)
                t=t .+ f(r[i]).*(r[i+1].-r[i-1])
            end
            t = t .+ (r[end].-r[end-1]).*f(r[end])
        end
        t./2
        end
    end
end


text:
export bsizeof, bdistance, bit_length, onehot, log2i, hypercubic, log2dim1, indices_with
# NOTE: all binary specified operations begin with b

"""
    bsizeof(::Type)

Returns the size of given type in number of binary digits.
"""
bsizeof(::Type{T}) where T = sizeof(T) << 3
bsizeof(::T) where T = bsizeof(T)

"""
    bdistance(i::Integer, j::Integer) -> Int

Return number of different bits.
"""
bdistance(i::Ti, j::Ti) where Ti<:Integer = count_ones(i ⊻ j)

"""
    onehot([T=Float64], nbits, x::Integer; nbatch::Int])

Create an onehot vector in type `Vector{T}` or a batch of onehot vector in type `Matrix{T}`,
where index `x + 1` is one.
"""
function onehot(::Type{T}, nbits::Int, x::Integer) where T
    v = zeros(T, 1 << nbits)
    v[x + 1] = 1
    return v
end

onehot(nbits::Int, x::Integer) = onehot(Float64, nbits, x)

function onehot(::Type{T}, nbits::Int, x::Integer, nbatch::Int) where T
    v = zeros(T, 1 << nbits, nbatch)
    v[x + 1, :] .= 1
    return v
end

onehot(nbits::Int, x::Integer, nbatch::Int) = onehot(Float64, nbits, x, nbatch)

"""
    unsafe_sub(a::UnitRange, b::NTuple{N}) -> NTuple{N}

Returns result in type `Tuple` of `a .- b`. This will not check the length of `a` and `b`, use
at your own risk.
"""
@generated function unsafe_sub(a::UnitRange{T}, b::NTuple{N, T}) where {N, T}
    ex = Expr(:tuple)
    for k in 1:N
        push!(ex.args, :(a.start + $(k-1) - b[$k]))
    end
    return ex
end

"""
    unsafe_sub(a::UnitRange{T}, b::Vector{T}) where T

Returns `a .- b`, fallback version when b is a `Vector`.
"""
unsafe_sub(a::UnitRange{T}, b::Vector{T}) where T = a .- b

"""
    bit_length(x::Integer) -> Int

Return the number of bits required to represent input integer x.
"""
bit_length(x::Integer)  =  sizeof(x)*8 - leading_zeros(x)

"""
    log2i(x::Integer) -> Integer

Return log2(x), this integer version of `log2` is fast but only valid for number equal to 2^n.
"""
function log2i end

for N in [8, 16, 32, 64, 128]
    T = Symbol(:Int, N)
    UT = Symbol(:UInt, N)
    @eval begin
        log2i(x::$T) = !signbit(x) ? ($(N - 1) - leading_zeros(x)) : throw(ErrorException("nonnegative expected ($x)"))
        log2i(x::$UT) = $(N - 1) - leading_zeros(x)
    end
end

"""
    log2dim1(X)

Returns the `log2` of the first dimension's size.
"""
log2dim1(X) = log2i(size(X, 1))

"""
    hypercubic(A::Array) -> Array

get the hypercubic representation for an array.
"""
hypercubic(A::Array) = reshape(A, fill(2, size(A) |> prod |> log2i)...)

"""
    indices_with(n::Int, locs::Vector{Int}, vals::Vector{Int}) -> Vector{Int}

Return indices with specific positions `locs` with value `vals` in a hilbert space of `n` qubits.
"""
function indices_with(n::Int, locs::Vector{Int}, vals::Vector{Int})
    mask = bmask(locs)
    onemask = bmask(locs[vals.!=0])
    return filter(x->ismatch(x, mask, onemask), basis(n))
end

ext: jl


text:
# This module provides an interface to the Moving AI MAPF benchmark suite found
# at https://www.movingai.com/benchmarks/index.html
export BenchmarkInterface

module BenchmarkInterface

using GraphUtils
using TOML
using ..CRCBS

struct ScenarioAgentSpec
    start::Tuple{Int,Int}
    goal::Tuple{Int,Int}
end
struct MAPFScenario
    mapfile::String
    buckets::Vector{Vector{ScenarioAgentSpec}}
end

"""
    parse_map_file

Parses a .map file (see citation below) into an indicator grid. Each cell of the map is
encoded by one of the following characters:
    . - passable terrain
    G - passable terrain
    @ - out of bounds
    O - out of bounds
    T - trees (unpassable)
    S - swamp (passable from regular terrain)
    W - water (traversable, but not passable from terrain)

Returns an array of integers encoded as `IMPASSABLE=>1, FREE=>0` (we treat only
'G', 'S', and '.' as free).

@article{sturtevant2012benchmarks,
  title={Benchmarks for Grid-Based Pathfinding},
  author={Sturtevant, N.},
  journal={Transactions on Computational Intelligence and AI in Games},
  volume={4},
  number={2},
  pages={144 -- 148},
  year={2012},
  url = {http://web.cs.du.edu/~sturtevant/papers/benchmarks.pdf},
}
"""
function parse_map_file(filename,encoding=Dict('@'=>1,'O'=>1,'T'=>1,'W'=>1,'.'=>0,'G'=>0,'S'=>0))
    @assert splitext(filename)[end] == ".map"
    grid = zeros(Int,1,1)
    open(filename,"r") do io
        line = readline(io)
        if findfirst("type",line) == nothing
            throw(ErrorException("First line of map file should give type!"))
        end
        map_type = split(line)[end]
        if map_type != "octile"
            throw(ErrorException("Parser cannot handle non-octile maps!"))
        end
        line = readline(io)
        if findfirst("height",line) == nothing
            throw(ErrorException("Second line of map file should give height!"))
        end
        height = parse(Int,split(line)[end])
        line = readline(io)
        if findfirst("width",line) == nothing
            throw(ErrorException("Third line of map file should give width!"))
        end
        width = parse(Int,split(line)[end])
        @assert readline(io) == "map"
        grid = zeros(Int,height,width)
        for i in 1:height
            line = readline(io)
            for (j,c) in enumerate(line)
                if haskey(encoding,c)
                    grid[i,j] = encoding[c]
                else
                    throw(ErrorException("unrecognized character type $c"))
                end
            end
        end
    end
    return grid
end

"""
    parse_mapf_scenario(filename,map_path="")

Parses a .scen file into a set of 'buckets', where each bucket contains a list
of (start location,goal location) pairs. Each bucket can be used to instantiate
MAPF instances some (or all) of these pairs. The benchmarking approach proposed
on the benchmark website (https://www.movingai.com/benchmarks/index.html) is to
start with a 2-agent MAPF for each bucket, and increase the number of agents
until solver time out.
"""
function parse_mapf_scenario(filename,map_path="")
    @assert splitext(filename)[end] == ".scen"
    buckets = Dict{Int,Vector{ScenarioAgentSpec}}()
    mapfile = ""
    open(filename,"r") do io
        for (i,line) in enumerate(eachline(io))
            if i == 1
                if findfirst("version", line) != nothing
                    continue
                end
            end
            split_line = split(line)
            bucket = parse(Int,split_line[1])
            mapfile = split_line[2]
            map_width = parse(Int,split_line[3])
            map_height = parse(Int,split_line[4])
            start_x = parse(Int,split_line[5]) + 1
            start_y = parse(Int,split_line[6]) + 1
            goal_x = parse(Int,split_line[7]) + 1
            goal_y = parse(Int,split_line[8]) + 1
            optimal_length = parse(Float64,split_line[9])
            if !haskey(buckets,bucket)
                buckets[bucket] = valtype(buckets)()
            end
            push!(buckets[bucket],ScenarioAgentSpec(
                (start_x,start_y),
                (goal_x,goal_y)
            ))
        end
    end
    if isfile(map_path)
        mapfile = map_path
    else
        mapfile = joinpath(map_path,mapfile)
    end
    @assert isfile(mapfile)
    return MAPFScenario(
        mapfile,
        map(k->buckets[k],sort(collect(keys(buckets))))
    )
end

function generate_cbs_env(indicator_grid)
    CBSEnv.LowLevelEnv(graph=construct_factory_env_from_indicator_grid(indicator_grid))
end
function generate_state_map(env::GridFactoryEnvironment)
    (x,y) -> CBSEnv.State(env.vtx_map[x,y],0)
end

function construct_base_mapf(scen::MAPFScenario,env_generator=generate_cbs_env)
    grid = parse_map_file(scen.mapfile)
    env = env_generator(grid)
    return MAPF(env,Vector{state_type(env)}(),Vector{state_type(env)}())
end

function gen_mapf_problem_from_scenario(mapf::M,
        scen::MAPFScenario,
        bucket_id::Int,
        n_agents::Int,
        state_map=generate_state_map(mapf.env.graph),
        ) where {M<:MAPF}
    @assert(length(scen.buckets) >= bucket_id,
        "bucket_id $bucket_id out of range!")
    bucket = scen.buckets[bucket_id]
    @assert(length(bucket) >= n_agents,
        string("bucket only has $(length(bucket)) agents.",
        " Cannot instantiate a problem with $n_agents agents."))
    starts = Vector{state_type(mapf)}()
    goals = Vector{state_type(mapf)}()
    for i in 1:n_agents
        push!(starts, state_map(bucket[i].start...))
        push!(goals, state_map(bucket[i].goal...))
    end
    M(mapf.env,starts,goals)
end

"""
    MovingAIBenchmarkFile

Points to a TOML-formatted file that contains the following elements:
    scenario = "/path/to/scenario/file.scen"
    map_file = "/path/to/map/file.map"
    bucket_idx = 1 # an integer
    n_agents = 2 # an integer

Extension is .bm
"""
struct MovingAIBenchmarkFile
    path::String
end

function pad_file_number(id,pad=4)
    # hacky 0 padding to avoid dependency on Sprintf
    padding = prod(map(i->"0",1:(pad-length(string(id)))))
    return "$(padding)$(id)"
end

function generate_problem_files_from_moving_ai(
        scenario_paths,
        base_map_path,
        problem_dir
        ;
        verbose=true,
        pad=6,
        recursive=false
        )
    problem_id = 1
    mkpath(problem_dir)
    for scen_file in scenario_paths
        scenario = parse_mapf_scenario(scen_file,base_map_path)
        for (bucket_idx,bucket) in enumerate(scenario.buckets)
            for n_agents in 2:length(bucket)
                config_dict = Dict(
                    "scenario" => scen_file,
                    "map_file" => scenario.mapfile,
                    "bucket_idx" => bucket_idx,
                    "n_agents" => n_agents,
                )
                problem_filename = joinpath(
                    problem_dir,"problem_$(pad_file_number(problem_id,pad)).bm")
                if isfile(problem_filename)
                    if verbose
                        println("File ",problem_filename," already exists.",
                            " Skipping...")
                    end
                else
                    open(problem_filename,"w") do io
                        TOML.print(io,config_dict)
                    end
                end
                problem_id += 1
            end
        end
    end
    return true
end

"""
    ProblemLoader

Can be queried via `get_problem(iterator,problem_filename)` to return MAPF
instances.
"""
struct ProblemLoader{M}
    scenarios::Dict{String,MAPFScenario}
    base_mapfs::Dict{String,M}
    ProblemLoader() = new{MAPF}(
        Dict{String,MAPFScenario}(),
        Dict{String,MAPF}()
    )
end
function get_scenario!(loader::ProblemLoader,problem_config::Dict)
    scen_file   = problem_config["scenario"]
    map_file    = problem_config["map_file"]
    if !haskey(loader.scenarios,scen_file)
        loader.scenarios[scen_file] = parse_mapf_scenario(scen_file, map_file)
    end
    return loader.scenarios[scen_file]
end
function get_base_mapf!(loader::ProblemLoader,problem_config::Dict)
    scen_file   = problem_config["scenario"]
    map_file    = problem_config["map_file"]
    scenario = get_scenario!(loader,problem_config)
    if !haskey(loader.base_mapfs,map_file)
        loader.base_mapfs[map_file] = construct_base_mapf(scenario)
    end
    return loader.base_mapfs[map_file]
end

struct LoaderState
    probfile_idx::Int
end

function init_mapf_loader(problem_dir)
    loader = ProblemLoader()
    for problem_file in readdir(problem_dir;join=true)
        problem_config = TOML.parsefile(problem_file)
        get_scenario!(loader,problem_config)
        get_base_mapf!(loader,problem_config)
    end
    return loader
end

function CRCBS.load_problem(loader::ProblemLoader,probfile)
    problem_config = TOML.parsefile(probfile)
    scenario = get_scenario!(loader,problem_config)
    base_mapf = get_base_mapf!(loader,problem_config)
    bucket_idx  = problem_config["bucket_idx"]
    n_agents    = problem_config["n_agents"]
    mapf = gen_mapf_problem_from_scenario(
        base_mapf, scenario, bucket_idx, n_agents)
    return mapf
end

"""
    is_same_scenario_and_bucket(config1,config2)

Checks that two problem_configs are on the same scenario and bucket.
"""
function is_same_scenario_and_bucket(config1,config2)
    for k in ["scenario","bucket_idx"]
        if !(haskey(config1,k) && haskey(config2,k))
            return false
        end
        if !(config1[k] == config2[k])
            return false
        end
    end
    return true
end

"""
    profile_with_skipping!(config,loader)

Follow the Moving AI benchmarking paradigm, where problems are skipped if the
solver has already failed on the same problem instance with fewer agents.
"""
function profile_with_skipping!(config,loader)
    for solver_config in config.solver_configs
        mkpath(solver_config.results_path)
    end
    previous_config = Dict{String,Any}()
    for problem_file in readdir(config.problem_dir;join=true)
        problem_name = splitext(splitdir(problem_file)[end])[1]
        problem_config = TOML.parsefile(problem_file)
        # Reset all solvers if we have moved to a new scenario or bucket
        if !is_same_scenario_and_bucket(problem_config,previous_config)
            previous_config = problem_config
            for solver_config in config.solver_configs
                reset_solver!(solver_config.solver)
            end
        end
        scenario = get_scenario!(loader,problem_config)
        base_mapf = get_base_mapf!(loader,problem_config)
        mapf = load_problem(loader,problem_file)
        for solver_config in config.solver_configs
            solver = solver_config.solver
            outfile = joinpath(solver_config.results_path,string(problem_name,".results"))
            if isfile(outfile)
                @log_info(-1,verbosity(solver),"Results already stored for ",problem_name,
                    " at ", outfile)
                continue
            end
            if failed_status(solver)
                @log_info(1,verbosity(solver),"Skipping problem ",problem_name)
                continue # Keep skipping until new bucket
            end
            solution, timer_results = profile_solver!(solver,mapf)
            results_dict = compile_results(
                solver,
                config.feats,
                mapf,
                solution,
                timer_results
                )
            results_dict["problem_file"] = problem_file
            open(outfile,"w") do io
                TOML.print(io,results_dict)
            end
        end
    end
end

end # moduler


text:
# julia 1.3 crashes with some precompile calls

if VERSION >= v"1.6"
    for T in (Float16, Float32, Float64, Int16, Int32, Int64, UInt8, UInt16, UInt32)
        DS = NCDataset{Nothing}
        precompile(defVar, (DS, String, T))
        for N in (1, 2, 3, 4, 5)
            A = NCDatasets.Attributes{DS}
            Var = NCDatasets.Variable{T,N,DS}
            St = Dict{Symbol,Any}
            CF = NCDatasets.CFVariable{T,N,Var,A,St}
            precompile(Var, (DS, Cint, NTuple{N,Cint}, A))
            precompile(CF, (Var, A, St))
            precompile(dimsize, (CF,))
            precompile(getindex, (CF, Colon))
            precompile(getindex, (CF, Int))
            precompile(getindex, (CF, ntuple(Int, N)...))
            precompile(Array, (CF,))
            precompile(Array, (Var,))
            precompile(show, (IO, Var,))
            precompile(show, (IO, CF,))
            precompile(show, (IO, MIME"text/plain", Var,))
            precompile(show, (IO, MIME"text/plain", CF,))
        end
    end

    precompile(NCDataset{Nothing}, (String,))
    precompile(keys, (NCDataset{Nothing},))
    precompile(show, (IO, NCDataset{Nothing},))
    precompile(getindex, (NCDataset{Nothing}, String))
    precompile(getindex, (NCDataset{Nothing}, Symbol))
end


text:
using LibGit2

url = "https://github.com/hildebrandmw/X-Mem/"
branch = "master"
localdir = joinpath(@__DIR__, "X-Mem")

# Cleanup leftovers
ispath(localdir) && rm(localdir; force = true, recursive = true)

LibGit2.clone(url, localdir; branch = branch)

nprocs = parse(Int, read(`nproc`, String))
cd(localdir)
run(`./build-linux.sh x64_avx $nprocs`)


text:
struct DeviceEnergyBalanceConstraintSpec
    constraint_name::Symbol
    energy_variable::Symbol
    initial_condition::Type{<:InitialConditionType}
    pin_variable_names::Vector{Symbol}
    pout_variable_names::Vector{Symbol}
    parameter_name::Union{Nothing, String}
    forecast_label::Union{Nothing, String}
    multiplier_func::Union{Nothing, Function}
    constraint_func::Function
end

function DeviceEnergyBalanceConstraintSpec(;
    constraint_name::Symbol,
    energy_variable::Symbol,
    initial_condition::Type{<:InitialConditionType},
    constraint_func::Function,
    pin_variable_names::Vector{Symbol} = Vector{Symbol}(),
    pout_variable_names::Vector{Symbol} = Vector{Symbol}(),
    parameter_name::Union{Nothing, String} = nothing,
    forecast_label::Union{Nothing, String} = nothing,
    multiplier_func::Union{Nothing, Function} = nothing,
)
    return DeviceEnergyBalanceConstraintSpec(
        constraint_name,
        energy_variable,
        initial_condition,
        pin_variable_names,
        pout_variable_names,
        parameter_name,
        forecast_label,
        multiplier_func,
        constraint_func,
    )
end

function add_constraints!(
    optimization_container::OptimizationContainer,
    ::Type{T},
    ::Type{U},
    devices::IS.FlattenIteratorWrapper{V},
    model::DeviceModel{V, W},
    ::Type{X},
    feedforward::Union{Nothing, AbstractAffectFeedForward},
) where {
    T <: EnergyBalanceConstraint,
    U <: VariableType,
    V <: PSY.Device,
    W <: AbstractDeviceFormulation,
    X <: PM.AbstractPowerModel,
}
    use_parameters = model_has_parameters(optimization_container)
    use_forecasts = model_uses_forecasts(optimization_container)
    # @assert !(use_parameters && !use_forecasts)
    spec = DeviceEnergyBalanceConstraintSpec(
        T,
        U,
        V,
        W,
        X,
        feedforward,
        use_parameters,
        use_forecasts,
    )
    energy_balance_constraints!(optimization_container, devices, model, feedforward, spec)
end

function energy_balance_constraints!(
    optimization_container::OptimizationContainer,
    devices::IS.FlattenIteratorWrapper{T},
    model::DeviceModel{T, U},
    feedforward::Union{Nothing, AbstractAffectFeedForward},
    spec::DeviceEnergyBalanceConstraintSpec,
) where {T <: PSY.Device, U <: AbstractDeviceFormulation}
    initial_conditions =
        get_initial_conditions(optimization_container, spec.initial_condition, T)
    _apply_energy_balance_constraint_spec!(
        optimization_container,
        spec,
        devices,
        initial_conditions,
        model,
        feedforward,
    )
end

struct DeviceEnergyBalanceConstraintSpecInternal
    constraint_infos::Vector{<:EnergyBalanceConstraintInfo}
    constraint_name::Symbol
    energy_variable::Symbol
    pin_variable_names::Vector{Symbol}
    pout_variable_names::Vector{Symbol}
    param_reference::Union{Nothing, UpdateRef}
end

function _apply_energy_balance_constraint_spec!(
    optimization_container,
    spec,
    devices::IS.FlattenIteratorWrapper{T},
    initial_conditions::Vector{InitialCondition},
    model,
    ff_affected_variables,
) where {T <: PSY.Device}
    constraint_infos = Vector{EnergyBalanceConstraintInfo}(undef, length(devices))
    for (ix, ic) in enumerate(initial_conditions)
        device = ic.device
        dev_name = PSY.get_name(device)
        if !isnothing(spec.forecast_label)
            ts_vector = get_time_series(optimization_container, device, spec.forecast_label)
            multiplier = spec.multiplier_func(device)
            constraint_info = EnergyBalanceConstraintInfo(
                dev_name,
                get_efficiency(device, spec.initial_condition),
                ic,
                multiplier,
                ts_vector,
            )
        else
            constraint_info = EnergyBalanceConstraintInfo(;
                component_name = dev_name,
                efficiency_data = get_efficiency(device, spec.initial_condition),
                ic_energy = ic,
            )
        end
        constraint_infos[ix] = constraint_info
    end

    spec.constraint_func(
        optimization_container,
        DeviceEnergyBalanceConstraintSpecInternal(
            constraint_infos,
            spec.constraint_name,
            spec.energy_variable,
            spec.pin_variable_names,
            spec.pout_variable_names,
            spec.parameter_name === nothing ? nothing :
            UpdateRef{T}(spec.parameter_name, spec.forecast_label),
        ),
    )
    return
end

@doc raw"""
Constructs multi-timestep constraint from initial condition, efficiency data, and variable tuple
# Constraints
If t = 1:
``` varenergy[name, 1] == initial_conditions[ix].value + varin[name, 1]*eff_in*fraction_of_hour - varout[name, 1]*fraction_of_hour/eff_out ```
If t > 1:
``` varenergy[name, t] == varenergy[name, t-1] + varin[name, t]*eff_in*fraction_of_hour - varout[name, t]*fraction_of_hour/eff_out ```
# LaTeX
`` x^{energy}_1 == x^{energy}_{init} + frhr \eta^{in} x^{in}_1 - \frac{frhr}{\eta^{out}} x^{out}_1, \text{ for } t = 1 ``
`` x^{energy}_t == x^{energy}_{t-1} + frhr \eta^{in} x^{in}_t - \frac{frhr}{\eta^{out}} x^{out}_t, \forall t \geq 2 ``
# Arguments
* optimization_container::OptimizationContainer : the optimization_container model built in PowerSimulations
* inputs::Vector{DeviceEnergyBalanceConstraintSpecInternal} : stores constraint information
"""
function energy_balance!(
    optimization_container::OptimizationContainer,
    inputs::DeviceEnergyBalanceConstraintSpecInternal,
)
    time_steps = model_time_steps(optimization_container)
    resolution = model_resolution(optimization_container)
    fraction_of_hour = Dates.value(Dates.Minute(resolution)) / MINUTES_IN_HOUR
    names = [get_component_name(x) for x in inputs.constraint_infos]

    varenergy = get_variable(optimization_container, inputs.energy_variable)
    varin = [get_variable(optimization_container, var) for var in inputs.pin_variable_names]
    varout =
        [get_variable(optimization_container, var) for var in inputs.pout_variable_names]

    constraint = add_cons_container!(
        optimization_container,
        inputs.constraint_name,
        names,
        time_steps,
    )

    for info in inputs.constraint_infos
        eff_in = info.efficiency_data.in
        eff_out = info.efficiency_data.out
        name = get_component_name(info)
        # Create the PGAE outside of the constraint definition
        !isnothing(info.timeseries) ? ts_value = info.timeseries[1] * info.multiplier :
        ts_value = 0.0
        expr = JuMP.AffExpr(0.0)
        for var in varin
            JuMP.add_to_expression!(expr, var[name, 1], eff_in * fraction_of_hour)
        end
        for var in varout
            JuMP.add_to_expression!(expr, var[name, 1], -1.0 * fraction_of_hour / eff_out)
        end
        constraint[name, 1] = JuMP.@constraint(
            optimization_container.JuMPmodel,
            varenergy[name, 1] == info.ic_energy.value + expr + ts_value
        )
    end

    for t in time_steps[2:end], info in inputs.constraint_infos
        eff_in = info.efficiency_data.in
        eff_out = info.efficiency_data.out
        name = get_component_name(info)
        !isnothing(info.timeseries) ? ts_value = info.timeseries[t] * info.multiplier :
        ts_value = 0.0
        expr = JuMP.AffExpr(0.0, varenergy[name, t - 1] => 1.0)
        for var in varin
            JuMP.add_to_expression!(expr, var[name, t], eff_in * fraction_of_hour)
        end
        for var in varout
            JuMP.add_to_expression!(expr, var[name, t], -1.0 * fraction_of_hour / eff_out)
        end
        constraint[name, t] = JuMP.@constraint(
            optimization_container.JuMPmodel,
            varenergy[name, t] == expr + ts_value
        )
    end

    return
end

@doc raw"""
Constructs multi-timestep constraint from initial condition, efficiency data, and variable tuple
# Constraints
If t = 1:
``` varenergy[name, 1] == initial_conditions[ix].value + varin[name, 1]*eff_in*fraction_of_hour - varout[name, 1]*fraction_of_hour/eff_out ```
If t > 1:
``` varenergy[name, t] == varenergy[name, t-1] + varin[name, t]*eff_in*fraction_of_hour - varout[name, t]*fraction_of_hour/eff_out ```
# LaTeX
`` x^{energy}_1 == x^{energy}_{init} + frhr \eta^{in} x^{in}_1 - \frac{frhr}{\eta^{out}} x^{out}_1, \text{ for } t = 1 ``
`` x^{energy}_t == x^{energy}_{t-1} + frhr \eta^{in} x^{in}_t - \frac{frhr}{\eta^{out}} x^{out}_t, \forall t \geq 2 ``
# Arguments
* optimization_container::OptimizationContainer : the optimization_container model built in PowerSimulations
* inputs::Vector{DeviceEnergyBalanceConstraintSpecInternal} : stores constraint information
"""
function energy_balance_param!(
    optimization_container::OptimizationContainer,
    inputs::DeviceEnergyBalanceConstraintSpecInternal,
)
    time_steps = model_time_steps(optimization_container)
    resolution = model_resolution(optimization_container)
    fraction_of_hour = Dates.value(Dates.Minute(resolution)) / MINUTES_IN_HOUR
    names = [get_component_name(x) for x in inputs.constraint_infos]
    has_parameter_data = !isnothing(inputs.param_reference)

    varenergy = get_variable(optimization_container, inputs.energy_variable)
    varin = [get_variable(optimization_container, var) for var in inputs.pin_variable_names]
    varout =
        [get_variable(optimization_container, var) for var in inputs.pout_variable_names]

    if has_parameter_data
        container = add_param_container!(
            optimization_container,
            inputs.param_reference,
            names,
            time_steps,
        )
        multiplier = get_multiplier_array(container)
        param = get_parameter_array(container)
    end

    constraint = add_cons_container!(
        optimization_container,
        inputs.constraint_name,
        names,
        time_steps,
    )

    for info in inputs.constraint_infos
        eff_in = info.efficiency_data.in
        eff_out = info.efficiency_data.out
        name = get_component_name(info)
        # Create the PGAE outside of the constraint definition
        expr = zero(PGAE)
        for var in varin
            JuMP.add_to_expression!(expr, var[name, 1], eff_in * fraction_of_hour)
        end
        for var in varout
            JuMP.add_to_expression!(expr, var[name, 1], -1 * fraction_of_hour / eff_out)
        end
        if has_parameter_data
            multiplier[name, 1] = info.multiplier
            param[name, 1] =
                PJ.add_parameter(optimization_container.JuMPmodel, info.timeseries[1])
            JuMP.add_to_expression!(expr, param[name, 1], multiplier[name, 1])
        end
        constraint[name, 1] = JuMP.@constraint(
            optimization_container.JuMPmodel,
            varenergy[name, 1] == info.ic_energy.value + expr
        )
    end

    for t in time_steps[2:end], info in inputs.constraint_infos
        eff_in = info.efficiency_data.in
        eff_out = info.efficiency_data.out
        name = get_component_name(info)
        expr = expr = zero(PGAE)
        JuMP.add_to_expression!(expr, varenergy[name, t - 1])
        for var in varin
            JuMP.add_to_expression!(expr, var[name, t], eff_in * fraction_of_hour)
        end
        for var in varout
            JuMP.add_to_expression!(expr, var[name, t], -1 * fraction_of_hour / eff_out)
        end
        if has_parameter_data
            multiplier[name, t] = info.multiplier
            param[name, t] =
                PJ.add_parameter(optimization_container.JuMPmodel, info.timeseries[t])
            JuMP.add_to_expression!(expr, param[name, t], multiplier[name, t])
        end
        constraint[name, t] =
            JuMP.@constraint(optimization_container.JuMPmodel, varenergy[name, t] == expr)
    end

    return
end

########## Old constraint function that will be deleted after storage refactor is complete

@doc raw"""
Constructs multi-timestep constraint from initial condition, efficiency data, and variable tuple

# Constraints

If t = 1:

``` varenergy[name, 1] == initial_conditions[ix].value + varin[name, 1]*eff_in*fraction_of_hour - varout[name, 1]*fraction_of_hour/eff_out ```

If t > 1:

``` varenergy[name, t] == varenergy[name, t-1] + varin[name, t]*eff_in*fraction_of_hour - varout[name, t]*fraction_of_hour/eff_out ```

# LaTeX

`` x^{energy}_1 == x^{energy}_{init} + frhr \eta^{in} x^{in}_1 - \frac{frhr}{\eta^{out}} x^{out}_1, \text{ for } t = 1 ``

`` x^{energy}_t == x^{energy}_{t-1} + frhr \eta^{in} x^{in}_t - \frac{frhr}{\eta^{out}} x^{out}_t, \forall t \geq 2 ``

# Arguments
* optimization_container::OptimizationContainer : the optimization_container model built in PowerSimulations
* initial_conditions::Vector{InitialCondition} : for time zero 'varenergy'
* efficiency_data::Tuple{Vector{String}, Vector{InOut}} :: charging/discharging efficiencies
* cons_name::Symbol : name of the constraint
* var_names::Tuple{Symbol, Symbol, Symbol} : the names of the variables
- : var_names[1] : varin
- : var_names[2] : varout
- : var_names[3] : varenergy

"""
function energy_balance(
    optimization_container::OptimizationContainer,
    initial_conditions::Vector{InitialCondition},
    efficiency_data::Tuple{Vector{String}, Vector{InOut}},
    cons_name::Symbol,
    var_names::Tuple{Symbol, Symbol, Symbol},
)
    parameters = model_has_parameters(optimization_container)
    time_steps = model_time_steps(optimization_container)
    resolution = model_resolution(optimization_container)
    fraction_of_hour = Dates.value(Dates.Minute(resolution)) / MINUTES_IN_HOUR
    name_index = efficiency_data[1]

    varin = get_variable(optimization_container, var_names[1])
    varout = get_variable(optimization_container, var_names[2])
    varenergy = get_variable(optimization_container, var_names[3])

    constraint =
        add_cons_container!(optimization_container, cons_name, name_index, time_steps)

    for (ix, name) in enumerate(name_index)
        eff_in = efficiency_data[2][ix].in
        eff_out = efficiency_data[2][ix].out
        # Create the PGAE outside of the constraint definition
        balance =
            initial_conditions[ix].value + varin[name, 1] * eff_in * fraction_of_hour -
            (varout[name, 1]) * fraction_of_hour / eff_out
        constraint[name, 1] = JuMP.@constraint(
            optimization_container.JuMPmodel,
            varenergy[name, 1] == balance
        )
    end

    for t in time_steps[2:end], (ix, name) in enumerate(name_index)
        eff_in = efficiency_data[2][ix].in
        eff_out = efficiency_data[2][ix].out

        constraint[name, t] = JuMP.@constraint(
            optimization_container.JuMPmodel,
            varenergy[name, t] ==
            varenergy[name, t - 1] + varin[name, t] * eff_in * fraction_of_hour -
            (varout[name, t]) * fraction_of_hour / eff_out
        )
    end

    return
end

@doc raw"""
Constructs multi-timestep constraint from initial condition, efficiency data, and variable tuple
# Constraints
If t = 1:
`` varenergy[name, 1] == initial_conditions[ix].value + (paraminflow[name, t] - varspill[name, 1] - varout[name, 1])*fraction_of_hour ``
If t > 1:
`` varenergy[name, t] == varenergy[name, t-1] + (paraminflow[name, t] - varspill[name, t] - varout[name, t])*fraction_of_hour ``
`` varenergy[name, end] >= paramenergytarget[name, end]
# LaTeX
`` x^{energy}_1 == x^{energy}_{init} + frhr  (x^{in}_1 - x^{spillage}_1 -  x^{out}_1), \text{ for } t = 1 ``
`` x^{energy}_t == x^{energy}_{t-1} + frhr (x^{in}_t - x^{spillage}_t - x^{out}_t), \forall t \geq 2 ``
`` x^{energy}_t >= x^{energy}_{target} \text{ for } t = end ``
# Arguments
* optimization_container::OptimizationContainer : the optimization_container model built in PowerSimulations
* initial_conditions::Vector{InitialCondition} : for time zero 'varenergy'
* time_series_data::Tuple{Vector{DeviceTimeSeriesConstraintInfo}, Vector{DeviceTimeSeriesConstraintInfo}} : forecast information
- : time_series_data[1] : Inflow energy forecast information
- : time_series_data[2] : Target reservoir storage forecast information
* cons_names::Tuple{Symbol, Symbol} : name of the constraints
- : cons_names[1] : energy balance constraint name
- : cons_names[2] : energy target constraint name
* var_names::Tuple{Symbol, Symbol, Symbol} : the names of the variables
- : var_names[1] : varspill
- : var_names[2] : varout
- : var_names[3] : varenergy
* param_reference::UpdateRef : UpdateRef to access the inflow parameter
"""
function energy_balance_hydro_param!(
    optimization_container::OptimizationContainer,
    initial_conditions::Vector{InitialCondition},
    time_series_data::Tuple{
        Vector{DeviceTimeSeriesConstraintInfo},
        Vector{DeviceTimeSeriesConstraintInfo},
    },
    cons_names::Tuple{Symbol, Symbol},
    var_names::Tuple{Symbol, Symbol, Symbol},
    param_references::Tuple{UpdateRef, UpdateRef},
)
    time_steps = model_time_steps(optimization_container)
    resolution = model_resolution(optimization_container)
    fraction_of_hour = Dates.value(Dates.Second(resolution)) / SECONDS_IN_HOUR

    inflow_data = time_series_data[1]
    target_data = time_series_data[2]

    name_index = [get_component_name(d) for d in inflow_data]

    varspill = get_variable(optimization_container, var_names[1])
    varout = get_variable(optimization_container, var_names[2])
    varenergy = get_variable(optimization_container, var_names[3])

    balance_cons_name = cons_names[1]
    target_cons_name = cons_names[2]
    balance_param_reference = param_references[1]
    target_param_reference = param_references[2]

    container_inflow = add_param_container!(
        optimization_container,
        balance_param_reference,
        name_index,
        time_steps,
    )
    param_inflow = get_parameter_array(container_inflow)
    multiplier_inflow = get_multiplier_array(container_inflow)
    container_target = add_param_container!(
        optimization_container,
        target_param_reference,
        name_index,
        time_steps,
    )
    param_target = get_parameter_array(container_target)
    multiplier_target = get_multiplier_array(container_target)

    balance_constraint = add_cons_container!(
        optimization_container,
        balance_cons_name,
        name_index,
        time_steps,
    )
    target_constraint =
        add_cons_container!(optimization_container, target_cons_name, name_index, 1)

    for (ix, d) in enumerate(inflow_data)
        name = get_component_name(d)
        multiplier_inflow[name, 1] = d.multiplier
        param_inflow[name, 1] =
            add_parameter(optimization_container.JuMPmodel, d.timeseries[1])
        exp =
            initial_conditions[ix].value +
            (
                multiplier_inflow[name, 1] * param_inflow[name, 1] - varspill[name, 1] -
                varout[name, 1]
            ) * fraction_of_hour
        balance_constraint[name, 1] =
            JuMP.@constraint(optimization_container.JuMPmodel, varenergy[name, 1] == exp)

        for t in time_steps[2:end]
            multiplier_inflow[name, t] = d.multiplier
            param_inflow[name, t] =
                add_parameter(optimization_container.JuMPmodel, d.timeseries[t])
            exp =
                varenergy[name, t - 1] +
                (
                    multiplier_inflow[name, t] * param_inflow[name, t] - varspill[name, t] -
                    varout[name, t]
                ) * fraction_of_hour
            balance_constraint[name, t] = JuMP.@constraint(
                optimization_container.JuMPmodel,
                varenergy[name, t] == exp
            )
        end
    end

    for (ix, d) in enumerate(target_data)
        name = get_component_name(d)
        for t in time_steps
            param_target[name, t] =
                add_parameter(optimization_container.JuMPmodel, d.timeseries[t])
            multiplier_target[name, t] = d.multiplier
        end
        target_constraint[name, 1] = JuMP.@constraint(
            optimization_container.JuMPmodel,
            varenergy[name, time_steps[end]] >=
            d.multiplier * param_target[name, time_steps[end]]
        )
    end

    return
end

@doc raw"""
Constructs multi-timestep constraint from initial condition, efficiency data, and variable tuple
# Constraints
If t = 1:
`` varenergy[name, 1] == initial_conditions[ix].value + (paraminflow[name, t] - varspill[name, 1] - varout[name, 1])*fraction_of_hour ``
If t > 1:
`` varenergy[name, t] == varenergy[name, t-1] + (paraminflow[name, t] - varspill[name, t] - varout[name, t])*fraction_of_hour ``
`` varenergy[name, end] >= paramenergytarget[name, end]
# LaTeX
`` x^{energy}_1 == x^{energy}_{init} + frhr  (x^{in}_1 - x^{spillage}_1 -  x^{out}_1), \text{ for } t = 1 ``
`` x^{energy}_t == x^{energy}_{t-1} + frhr (x^{in}_t - x^{spillage}_t - x^{out}_t), \forall t \geq 2 ``
`` x^{energy}_t >= x^{energy}_{target} \text{ for } t = end ``
# Arguments
* optimization_container::OptimizationContainer : the optimization_container model built in PowerSimulations
* initial_conditions::Vector{InitialCondition} : for time zero 'varenergy'
* time_series_data::Tuple{Vector{DeviceTimeSeriesConstraintInfo}, Vector{DeviceTimeSeriesConstraintInfo}} : forecast information
- : time_series_data[1] : Inflow energy forecast information
- : time_series_data[2] : Target reservoir storage forecast information
* cons_names::Tuple{Symbol, Symbol} : name of the constraints
- : cons_names[1] : energy balance constraint name
- : cons_names[2] : energy target constraint name
* var_names::Tuple{Symbol, Symbol, Symbol} : the names of the variables
- : var_names[1] : varspill
- : var_names[2] : varout
- : var_names[3] : varenergy
"""
function energy_balance_hydro!(
    optimization_container::OptimizationContainer,
    initial_conditions::Vector{InitialCondition},
    time_series_data::Tuple{
        Vector{DeviceTimeSeriesConstraintInfo},
        Vector{DeviceTimeSeriesConstraintInfo},
    },
    cons_names::Tuple{Symbol, Symbol},
    var_names::Tuple{Symbol, Symbol, Symbol},
)
    time_steps = model_time_steps(optimization_container)
    resolution = model_resolution(optimization_container)
    fraction_of_hour = Dates.value(Dates.Minute(resolution)) / MINUTES_IN_HOUR

    inflow_data = time_series_data[1]
    target_data = time_series_data[2]

    name_index = [get_component_name(d) for d in inflow_data]

    varspill = get_variable(optimization_container, var_names[1])
    varout = get_variable(optimization_container, var_names[2])
    varenergy = get_variable(optimization_container, var_names[3])

    balance_cons_name = cons_names[1]
    target_cons_name = cons_names[2]

    balance_constraint = add_cons_container!(
        optimization_container,
        balance_cons_name,
        name_index,
        time_steps,
    )
    target_constraint =
        add_cons_container!(optimization_container, target_cons_name, name_index, 1)

    for (ix, d) in enumerate(inflow_data)
        name = get_component_name(d)
        balance_constraint[name, 1] = JuMP.@constraint(
            optimization_container.JuMPmodel,
            varenergy[name, 1] ==
            initial_conditions[ix].value +
            (d.multiplier * d.timeseries[1] - varspill[name, 1] - varout[name, 1]) *
            fraction_of_hour
        )

        for t in time_steps[2:end]
            balance_constraint[name, t] = JuMP.@constraint(
                optimization_container.JuMPmodel,
                varenergy[name, t] ==
                varenergy[name, t - 1] +
                (d.multiplier * d.timeseries[t] - varspill[name, t] - varout[name, t]) *
                fraction_of_hour
            )
        end
    end

    for (ix, d) in enumerate(target_data)
        name = get_component_name(d)
        target_constraint[name, 1] = JuMP.@constraint(
            optimization_container.JuMPmodel,
            varenergy[name, time_steps[end]] >=
            d.multiplier * d.timeseries[time_steps[end]]
        )
    end

    return
end

@doc raw"""
Constructs multi-timestep constraint from initial condition, efficiency data, and variable tuple for pumped hydro
# Constraints
If t = 1:
``` varenergy_up[name, 1] == initial_conditions[ix].value + (param_inflow[name, t] + varin[name, 1] - varspill[name, 1] - varout[name, 1])*fraction_of_hour ```
If t > 1:
``` varenergy_up[name, t] == varenergy_up[name, t-1] + (param_inflow[name, t] + varin[name, t] - varspill[name, t] - varout[name, t])*fraction_of_hour ```
# LaTeX
`` x^{energy}_1 == x^{energy}_{init} + frhr  (x^{in}_1 + x^{in}_1 - x^{spillage}_1 -  x^{out}_1), \text{ for } t = 1 ``
`` x^{energy}_t == x^{energy}_{t-1} + frhr (x^{in}_t + x^{in}_t - x^{spillage}_t - x^{out}_t), \forall t \geq 2 ``
# Arguments
* optimization_container::OptimizationContainer : the optimization_container model built in PowerSimulations
* initial_conditions::Vector{InitialCondition} : for time zero 'varenergy_up'
* inflow_data::Vector{DeviceTimeSeriesConstraintInfo} :: Inflow energy forecast information
* cons_name::Symbol : name of the constraint
* var_names::Tuple{Symbol, Symbol, Symbol, Symbol} : the names of the variables
- : var_names[1] : varspill
- : var_names[2] : varout
- : var_names[3] : varenergy_up
- : var_names[4] : varin
* param_reference::UpdateRef : UpdateRef to access the inflow parameter
"""
function energy_balance_hydro_param!(
    optimization_container::OptimizationContainer,
    initial_conditions::Vector{InitialCondition},
    ts_data::Tuple{
        Vector{DeviceTimeSeriesConstraintInfo},
        Vector{DeviceTimeSeriesConstraintInfo},
    },
    cons_name::Tuple{Symbol, Symbol},
    var_names::Tuple{Symbol, Symbol, Symbol, Symbol, Symbol},
    param_reference::Tuple{UpdateRef, UpdateRef},
)
    time_steps = model_time_steps(optimization_container)
    resolution = model_resolution(optimization_container)
    fraction_of_hour = Dates.value(Dates.Second(resolution)) / SECONDS_IN_HOUR
    inflow_data = ts_data[1]
    outflow_data = ts_data[2]
    inflow_name_index = [get_component_name(d) for d in inflow_data]
    outflow_name_index = [get_component_name(d) for d in outflow_data]

    varspill = get_variable(optimization_container, var_names[1])
    varout = get_variable(optimization_container, var_names[2])
    varenergy_up = get_variable(optimization_container, var_names[3])
    varin = get_variable(optimization_container, var_names[4])
    varenergy_down = get_variable(optimization_container, var_names[5])

    container_inflow = add_param_container!(
        optimization_container,
        param_reference[1],
        inflow_name_index,
        time_steps,
    )
    param_inflow = get_parameter_array(container_inflow)
    multiplier_inflow = get_multiplier_array(container_inflow)
    container_outflow = add_param_container!(
        optimization_container,
        param_reference[2],
        outflow_name_index,
        time_steps,
    )
    param_outflow = get_parameter_array(container_outflow)
    multiplier_outflow = get_multiplier_array(container_outflow)
    constraint_up = add_cons_container!(
        optimization_container,
        cons_name[1],
        inflow_name_index,
        time_steps,
    )
    constraint_down = add_cons_container!(
        optimization_container,
        cons_name[2],
        outflow_name_index,
        time_steps,
    )

    for (ix, d) in enumerate(inflow_data)
        name = get_component_name(d)
        pump_eff = 1.0 # TODO: get pump efficiency PSY.get_pump_efficiency(d)
        multiplier_inflow[name, 1] = d.multiplier
        param_inflow[name, 1] =
            add_parameter(optimization_container.JuMPmodel, d.timeseries[1])
        exp =
            initial_conditions[ix].value +
            (
                multiplier_inflow[name, 1] * param_inflow[name, 1] +
                varin[name, 1] * pump_eff - varspill[name, 1] - varout[name, 1]
            ) * fraction_of_hour
        constraint_up[name, 1] =
            JuMP.@constraint(optimization_container.JuMPmodel, varenergy_up[name, 1] == exp)

        for t in time_steps[2:end]
            multiplier_inflow[name, t] = d.multiplier
            param_inflow[name, t] =
                add_parameter(optimization_container.JuMPmodel, d.timeseries[t])
            exp =
                varenergy_up[name, t - 1] +
                (
                    multiplier_inflow[name, t] * param_inflow[name, t] +
                    varin[name, t] * pump_eff - varspill[name, t] - varout[name, t]
                ) * fraction_of_hour
            constraint_up[name, t] = JuMP.@constraint(
                optimization_container.JuMPmodel,
                varenergy_up[name, t] == exp
            )
        end
    end

    for (ix, d) in enumerate(outflow_data)
        name = get_component_name(d)
        pump_eff = 1.0 # TODO: get pump efficiency PSY.get_pump_efficiency(d)
        multiplier_outflow[name, 1] = d.multiplier
        param_outflow[name, 1] =
            add_parameter(optimization_container.JuMPmodel, d.timeseries[1])
        exp =
            initial_conditions[ix].value +
            (
                varspill[name, 1] + varout[name, 1] -
                multiplier_outflow[name, 1] * param_outflow[name, 1] -
                varin[name, 1] * pump_eff
            ) * fraction_of_hour
        constraint_down[name, 1] = JuMP.@constraint(
            optimization_container.JuMPmodel,
            varenergy_down[name, 1] == exp
        )

        for t in time_steps[2:end]
            multiplier_outflow[name, t] = d.multiplier
            param_outflow[name, t] =
                add_parameter(optimization_container.JuMPmodel, d.timeseries[t])
            exp =
                varenergy_down[name, t - 1] +
                (
                    varspill[name, t] + varout[name, t] -
                    multiplier_outflow[name, t] * param_outflow[name, t] -
                    varin[name, t] * pump_eff
                ) * fraction_of_hour
            constraint_down[name, t] = JuMP.@constraint(
                optimization_container.JuMPmodel,
                varenergy_down[name, t] == exp
            )
        end
    end

    return
end

@doc raw"""
Constructs multi-timestep constraint from initial condition, efficiency data, and variable tuple for pumped hydro
# Constraints
If t = 1:
``` varenergy[name, 1] == initial_conditions[ix].value + (paraminflow[name, t] + varin[name, 1] - varspill[name, 1] - varout[name, 1])*fraction_of_hour ```
If t > 1:
``` varenergy[name, t] == varenergy[name, t-1] + (paraminflow[name, t] + varin[name, t] - varspill[name, t] - varout[name, t])*fraction_of_hour ```
# LaTeX
`` x^{energy}_1 == x^{energy}_{init} + frhr  (x^{in}_1 + x^{in}_1 - x^{spillage}_1 -  x^{out}_1), \text{ for } t = 1 ``
`` x^{energy}_t == x^{energy}_{t-1} + frhr (x^{in}_t + x^{in}_t - x^{spillage}_t - x^{out}_t), \forall t \geq 2 ``
# Arguments
* optimization_container::OptimizationContainer : the optimization_container model built in PowerSimulations
* initial_conditions::Vector{InitialCondition} : for time zero 'varenergy'
* inflow_data::TVector{DeviceTimeSeriesConstraintInfo} :: Inflow energy forecast information
* cons_name::Symbol : name of the constraint
* var_names::Tuple{Symbol, Symbol, Symbol, Symbol} : the names of the variables
- : var_names[1] : varspill
- : var_names[2] : varout
- : var_names[3] : varenergy
- : var_names[4] : varin
"""
function energy_balance_hydro!(
    optimization_container::OptimizationContainer,
    initial_conditions::Vector{InitialCondition},
    ts_data::Tuple{
        Vector{DeviceTimeSeriesConstraintInfo},
        Vector{DeviceTimeSeriesConstraintInfo},
    },
    cons_name::Tuple{Symbol, Symbol},
    var_names::Tuple{Symbol, Symbol, Symbol, Symbol, Symbol},
)
    time_steps = model_time_steps(optimization_container)
    resolution = model_resolution(optimization_container)
    fraction_of_hour = Dates.value(Dates.Minute(resolution)) / MINUTES_IN_HOUR
    inflow_data = ts_data[1]
    outflow_data = ts_data[2]
    inflow_name_index = [get_component_name(d) for d in inflow_data]
    outflow_name_index = [get_component_name(d) for d in outflow_data]

    varspill = get_variable(optimization_container, var_names[1])
    varout = get_variable(optimization_container, var_names[2])
    varenergy_up = get_variable(optimization_container, var_names[3])
    varin = get_variable(optimization_container, var_names[4])
    varenergy_down = get_variable(optimization_container, var_names[5])

    constraint_up = add_cons_container!(
        optimization_container,
        cons_name[1],
        inflow_name_index,
        time_steps,
    )
    constraint_down = add_cons_container!(
        optimization_container,
        cons_name[2],
        outflow_name_index,
        time_steps,
    )

    for (ix, d) in enumerate(inflow_data)
        name = get_component_name(d)
        pump_eff = 1.0 # TODO: get pump efficiency PSY.get_pump_efficiency(d)
        constraint_up[name, 1] = JuMP.@constraint(
            optimization_container.JuMPmodel,
            varenergy_up[name, 1] ==
            initial_conditions[ix].value +
            (
                d.multiplier * d.timeseries[1] + varin[name, 1] * pump_eff -
                varspill[name, 1] - varout[name, 1]
            ) * fraction_of_hour
        )

        for t in time_steps[2:end]
            constraint_up[name, t] = JuMP.@constraint(
                optimization_container.JuMPmodel,
                varenergy_up[name, t] ==
                varenergy_up[name, t - 1] +
                (
                    d.multiplier * d.timeseries[t] + varin[name, t] * pump_eff -
                    varspill[name, t] - varout[name, t]
                ) * fraction_of_hour
            )
        end
    end

    for (ix, d) in enumerate(outflow_data)
        name = get_component_name(d)
        pump_eff = 1.0 # TODO: get pump efficiency PSY.get_pump_efficiency(d)
        constraint_down[name, 1] = JuMP.@constraint(
            optimization_container.JuMPmodel,
            varenergy_down[name, 1] ==
            initial_conditions[ix].value +
            (
                varspill[name, 1] + varout[name, 1] - d.multiplier * d.timeseries[1] -
                varin[name, 1] * pump_eff
            ) * fraction_of_hour
        )

        for t in time_steps[2:end]
            constraint_down[name, t] = JuMP.@constraint(
                optimization_container.JuMPmodel,
                varenergy_down[name, t] ==
                varenergy_down[name, t - 1] +
                (
                    varspill[name, t] + varout[name, t] - d.multiplier * d.timeseries[t] - varin[name, t] * pump_eff
                ) * fraction_of_hour
            )
        end
    end
    return
end


text:
##
N = 1000
A = rand(N,N)

function fct(B, x, y)
    for i in eachindex(x)
        for j in eachindex(y)
            B[i,j] = A[x[i],y[j]]
        end
    end
end

U,S,V = svd(A)

S = [ i < 15 ? 10.0^(-i) : 0.0 for i = 1:N ]

A = U*diagm(S)*V'

U, V = aca_compression(fct, 1:N, 1:N, T=Float64)

@test U*V ≈ A atol = 1e-14

##
Ns = 200
Nt = 100
A = rand(Nt,Ns)

function fct(B, x, y)
    for i in eachindex(x)
        for j in eachindex(y)
            B[i,j] = A[x[i],y[j]]
        end
    end
end

U,S,V = svd(A)

S .= 0
S[1:15] = [10.0^(-i) for i = 1:15 ]

A = U*diagm(S)*V'

U, V = aca_compression(fct, 1:Nt, 1:Ns, T=Float64)

@test U*V ≈ A atol = 1e-14

##
Ns = 100
Nt = 200
A = rand(Nt,Ns)

function fct(B, x, y)
    for i in eachindex(x)
        for j in eachindex(y)
            B[i,j] = A[x[i],y[j]]
        end
    end
end

U,S,V = svd(A)

S .= 0
S[1:15] = [10.0^(-i) for i = 1:15 ]

A = U*diagm(S)*V'

U, V = aca_compression(fct, 1:Nt, 1:Ns, T=Float64)

@test U*V ≈ A atol = 1e-14

##
a = [1.0 0.0 0.0 0.0 0.0]
b = [7.0 3.0 1.0 9.0 1.0]
A = a'*b

function fct(B, x, y)
    for i in eachindex(x)
        for j in eachindex(y)
            B[i,j] = A[x[i],y[j]]
        end
    end
end

U, V = aca_compression(fct, 1:5, 1:5, T=Float64, svdrecompress=false)

@test U*V == A

##
a = [1.0 -2.0 6.0 4.0 5.0]
b = [7.0 3.0 1.0 9.0 1.0]
A = a'*b

function fct(B, x, y)
    for i in eachindex(x)
        for j in eachindex(y)
            B[i,j] = A[x[i],y[j]]
        end
    end
end

U, V = aca_compression(fct, 1:5, 1:5, T=Float64, svdrecompress=false)

@test U*V == A

##
a1 = [1.0 -2.0 6.0 4.0 5.0]
b1 = [7.0 3.0 1.0 9.0 1.0]
a2 = [9.0 3.0 -6.0 2.0 3.0]
b2 = [5.0 4.0 3.0 1.0 -4.0]
#a3 = [10.0 2.0 -6.0 2.0 3.0]
#b3 = [-11.0 4.0 3.0 1.0 -4.0]
A = a1'*b1 + a2'*b2 #+ a3'*b3

function fct(B, x, y)
    for i in eachindex(x)
        for j in eachindex(y)
            B[i,j] = A[x[i],y[j]]
        end
    end
end

U, V = aca_compression(fct, 1:5, 1:5, T=Float64)

@test U*V ≈ A atol = 1e-13

##
Ns = 100
Nt = 200
A = rand(Nt,Ns)

B = zeros(2*Nt, 2*Ns)

U,S,V = svd(A)

S .= 0
S[1:15] = [10.0^(-i) for i = 1:15 ]

A = U*diagm(S)*V'

rowindices = [2*Nt - 2*i + 1 for i=1:Nt]
colindices = [2*Ns - 2*i + 1 for i=1:Ns]

for i = 1:Ns
    for j=1:Nt
        B[end-2*j+1,end-2*i+1] = A[j,i]
    end
end

function fct(C, x, y)
    for i in eachindex(x)
        for j in eachindex(y)
            C[i,j] = B[x[i],y[j]]
        end
    end
end

U, V = aca_compression(fct, rowindices, colindices, T=Float64)

@test U*V ≈ A atol = 1e-14

##
A = rand(2,1)

function fct(B, x, y)
    for i in eachindex(x)
        for j in eachindex(y)
            B[i,j] = A[x[i],y[j]]
        end
    end
end

U, V = aca_compression(fct, 1:2, 1:1, T=Float64)
@test U*V ≈ A atol = 1e-14

##
A = rand(1,1)

function fct(B, x, y)
    for i in eachindex(x)
        for j in eachindex(y)
            B[i,j] = A[x[i],y[j]]
        end
    end
end

U, V = aca_compression(fct, 1:1, 1:1, T=Float64)
@test U*V ≈ A atol = 1e-14

max_stars_repo_name: sbadrian/FastBEAST

max_issues_repo_name: JoshuaTetzner/FastBEAST

max_forks_repo_name: JoshuaTetzner/FastBEAST


text:
abstract type AbstractImplementation end
struct MatrixImplementation <: AbstractImplementation end
struct ConvolutionImplementation <: AbstractImplementation end

function optflow(first_img::AbstractArray{T, 2}, second_img::AbstractArray{T,2},  algorithm::Farneback{}) where T <: Gray
    displacement = Array{SVector{2, Float64}, 2}(undef,size(first_img))
    for i in eachindex(displacement)
            displacement[i] = SVector{2, Float64}(0.0, 0.0)
    end
    optflow!(first_img, second_img, displacement,  algorithm)
end

function optflow!(first_img::AbstractArray{T, 2}, second_img::AbstractArray{T,2}, displacement::Array{SVector{2, Float64}, 2}, algorithm::Farneback{}) where T <: Gray

    # Replace NaN with zero in both images.
    first_img = map(x -> isnan(x) ? zero(x) : x, first_img)
    second_img = map(x -> isnan(x) ? zero(x) : x, second_img)

    start_rows = first(first(axes(first_img)))
    start_cols = first(last(axes(first_img)))
    end_rows = last(first(axes(first_img)))
    end_cols = last(last(axes(first_img)))

    A1, B1, _ = polynomial_expansion(ConvolutionImplementation(), first_img, algorithm.expansion_window, algorithm.σ_expansion_window)
    A2, B2, _ = polynomial_expansion(ConvolutionImplementation(), second_img, algorithm.expansion_window, algorithm.σ_expansion_window)

    AtA = zero(similar(A1))
    AtB = zero(similar(B1))

    AtA_smoothed = zero(similar(A1))
    AtB_smoothed = zero(similar(B1))

    neighbourhood = 2*algorithm.estimation_window + 1
    σ, w =  (algorithm.σ_estimation_window,algorithm.σ_estimation_window), (neighbourhood,neighbourhood)
    kern = Kernel.gaussian(σ ,w)

    Aₗ = zeros(Float64,2,2)
    AtAₗ = zeros(Float64,2,2)
    AtBₗ = zeros(Float64,2,1)
    Adₗ = zeros(Float64,2,1)
    bₗ = zeros(Float64,2,1)
    dbₗ = zeros(Float64,2,1)
    dₗ = zeros(Float64,2,1)
    len =  algorithm.expansion_window
    for k = 1:algorithm.iterations
        # We don't utilise the raw polynomial expansion coefficients at the
        # border of the image; hence, we add and subtract `len`.
        for j = start_cols + len:end_cols - len
            for i = start_rows + len:end_rows - len
                d = displacement[i,j]

                mod_i = round(Int, i + d[2])
                mod_j = round(Int, j + d[1])
                mod_i = clamp(mod_i, 1, first(size(first_img)))
                mod_j = clamp(mod_j, 1, last(size(first_img)))

                # (A1[i,j,:,:] + A2[i,j,:,:]) / 2
                for p = 1:2
                    for q = 1:2
                        setindex!(Aₗ, A1[i,j,p,q] / 2 + A2[mod_i,mod_j,p,q] / 2, p, q)
                    end
                end
                # A*d
                setindex!(dₗ, d[1], 1)
                setindex!(dₗ, d[2], 2)
                mul!(Adₗ, Aₗ, dₗ)
                # (B2[i,j,:] - B1[i,j,:]) / 2
                for p = 1:2
                    setindex!(bₗ, (B2[mod_i,mod_j,p] - B1[i,j,p]) / 2 , p)
                end
                # dB = A*d - ((B2 - B1)/2)
                broadcast!(-, dbₗ, Adₗ, bₗ)
                # A'*A
                mul!(AtAₗ,Aₗ',Aₗ)
                # A'*dB
                mul!(AtBₗ,Aₗ',dbₗ)
                # Store result
                for p = 1:2
                    for q = 1:2
                        setindex!(AtA, AtAₗ[p,q], i, j, p, q)
                    end
                    setindex!(AtB, AtBₗ[p], i, j, p)
                end
            end
        end

        for p = 1:2
            for q = 1:2
                AtAview = @view AtA[:,:,p,q]
                AtAview_smoothed = @view AtA_smoothed[:,:,p,q]
                imfilter!(AtAview_smoothed, AtAview, kern, Pad(:replicate))
            end
            AtBview = @view AtB[:,:,p]
            AtBview_smoothed = @view AtB_smoothed[:,:,p]
            imfilter!(AtBview_smoothed, AtBview, kern, Pad(:replicate))
        end

        for j = start_cols:end_cols
            for i = start_rows:end_rows
                P = @view AtA_smoothed[i,j,:,:]
                Q = @view  AtB_smoothed[i,j,:]
                displacement[i,j] = SVector{2, Float64}(pinv2x2(P) * Q)
            end
        end
    end
    # Convert (x,y) coordinate convention to (row, col) convention.
    map!(x -> SVector(last(x), first(x)), displacement, displacement)
    return displacement
end

"""
```
A, B, C = polynomial_expansion(implementation, img, neighbourhood, σ)
```

Returns the polynomial coefficients of the approximation of the neighbourhood of a
point in a 2D function with a polynomial.
The expansion is: f(x) = x'Ax + B'x + C

# Options

Various options for the fields of this type are described in more detail
below.

## Choices for `implementation`

Selects implementation (`::MatrixImplementation` vs `::ConvolutionImplementation`).
The matrix implementation serves as a useful reference against which other
faster implementations can be validated.

## Choices for `img`

Grayscale (Float) image whose polynomial expansion is to be found.

## Choices for `window_size`

Determines the size of the pixel neighbourhood used to find polynomial expansion
for each pixel; larger values mean that the image will be approximated with
smoother surfaces, yielding more robust algorithm and more blurred motion field.
The total size equals `2*window_size + 1`.

## Choices for `σ`

Standard deviation of the Gaussian that is used to smooth the image for the purpose
of approximating it with a polynomial expansion.

## References

Farnebäck, G.: Polynomial Expansion for Orientation and Motion Estimation. PhD thesis, Linköping University,
Sweden, SE-581 83 Linköping, Sweden (2002) Dissertation No 790, ISBN 91-7373-475-6.

"""
function polynomial_expansion(implementation::MatrixImplementation, img::AbstractArray{T, 2}, window_size::Int, σ::Real) where T <: Gray
    len = window_size
    neighbourhood = 2*window_size + 1
    a = ImageFiltering.Kernel.gaussian((σ,σ), (neighbourhood,neighbourhood))[-len:len, -len:len]

    padded_img = parent(padarray(img, Fill(0, (len, len))))

    B0 = zeros(neighbourhood*neighbourhood, 6)
    for x = -len:len
        for y = -len:len
            B0[neighbourhood*(x + len) + (y + len) + 1, :] = [1, x, y, x*x, y*y, x*y]
        end
    end

    B = SMatrix{neighbourhood*neighbourhood, 6}(B0)
    Wa = SDiagonal{neighbourhood^2,Float64}( @view a[:])
    BtWa = B'*Wa

    # If we assume that the confidence is constant we can move several
    # matrix multiplications out of the main loop.
    Wc = SDiagonal(ones(1,neighbourhood^2)...)
    BtWaWc = BtWa*Wc
    G = BtWaWc*B
    G_inv = SMatrix{6,6,Float64}(pinv(convert(Array,G)))
    G_invBtWaWc = G_inv*BtWaWc
    poly_A = zeros(size(img)[1], size(img)[2], 2, 2)
    poly_B = zeros(size(img)[1], size(img)[2], 2)
    poly_C = zeros(size(img))
    f = @MVector zeros(neighbourhood^2)
    for i = (len + 1):(size(padded_img)[1] - len)
        for j = (len + 1):(size(padded_img)[2] - len)
            sub_img = @view padded_img[(i - len):(i + len), (j - len):(j + len)]
            f .= vec(sub_img)
            r = G_invBtWaWc*f

            poly_A[i-len, j-len, 1, 1] =  r[4]
            poly_A[i-len, j-len, 2, 1] =  r[6]/2
            poly_A[i-len, j-len, 1, 2] =  r[6]/2
            poly_A[i-len, j-len, 2, 2] =  r[5]

            poly_B[i-len, j-len, 1] = r[2]
            poly_B[i-len, j-len, 2] = r[3]

            poly_C[i-len, j-len] = r[1]
        end
    end
    return (poly_A, poly_B, poly_C)
end

function polynomial_expansion(imp::ConvolutionImplementation, img::AbstractArray{T, 2}, window_size::Int, σ::Real) where T <: Gray
    #len = floor(Int, neighbourhood/2)
    len = window_size
    neighbourhood = 2*window_size + 1

    # Define basis functions.
    P₂ = construct_basis(len, neighbourhood)

    # Construct Gaussian filter.
    g = collect(KernelFactors.gaussian(σ,neighbourhood))

    # Construct inverse metric using Gaussian "applicability".
    G⁻¹ = construct_inverse_metric(P₂, g)

    unity, x, x², y, xy, y² = calculate_correlations(img, len, g)

    # In-place version of the following code (thus reducing allocations):
    # unity_tmp = G⁻¹[1,1] * unity + G⁻¹[1,4] * x² + G⁻¹[1,5] * y²
    # x = G⁻¹[2,2] * x
    # y = G⁻¹[3,3] * y
    # x² = G⁻¹[4,4] * x² + G⁻¹[4,1] * unity
    # y² = G⁻¹[5,5] * y² + G⁻¹[5,1] * unity
    # xy = G⁻¹[6,6] * xy
    # unity = unity_tmp
    unity_tmp = similar(unity)
    broadcast!(+,unity_tmp, lmul!(G⁻¹[1,1], unity), lmul!(G⁻¹[1,4], x²), lmul!(G⁻¹[1,5] , y²))
    x =  lmul!(G⁻¹[2,2], x)
    y =  lmul!(G⁻¹[3,3], y)
    x² = broadcast!(+,x², lmul!(G⁻¹[4,4]/G⁻¹[1,4], x²), lmul!(G⁻¹[4,1]/G⁻¹[1,1], unity))
    y² = broadcast!(+,y², lmul!(G⁻¹[5,5]/G⁻¹[1,5], y² ), lmul!(G⁻¹[5,1]/G⁻¹[4,1], unity))
    xy = lmul!(G⁻¹[6,6], xy)
    unity = unity_tmp

    # # The result of each multiplication allocated memory which is undesirable.
    # # I've left this version because I am trying to track down an occasional
    # # NaN bug and I want to be sure it is not due to the `in-place` version of
    # # this code.
    # unity_tmp = G⁻¹[1,1] * unity + G⁻¹[1,4] * x² + G⁻¹[1,5] * y²
    # x = G⁻¹[2,2] * x
    # y = G⁻¹[3,3] * y
    # x² = G⁻¹[4,4] * x² + G⁻¹[4,1] * unity
    # y² = G⁻¹[5,5] * y² + G⁻¹[5,1] * unity
    # xy = G⁻¹[6,6] * xy
    # unity = unity_tmp

    poly_A = zeros(Float64,first(size(img)), last(size(img)), 2, 2)
    poly_B = zeros(Float64,first(size(img)), last(size(img)), 2)
    poly_C = unity
    for j =  1:last(size(img))
        for i = 1:first(size(img))
            @inbounds setindex!(poly_A, x²[i,j], i, j, 1, 1)
            @inbounds setindex!(poly_A, xy[i,j]/2, i, j, 1, 2)
            @inbounds setindex!(poly_A, poly_A[i,j,1,2], i, j, 2, 1)
            @inbounds setindex!(poly_A, y²[i,j], i, j, 2, 2)
            @inbounds setindex!(poly_B, x[i,j], i, j, 1)
            @inbounds setindex!(poly_B, y[i,j], i, j, 2)
        end
    end
    return (poly_A, poly_B, poly_C)
end

function construct_basis(len::Int, neighbourhood::Int)
    P₂ = zeros(neighbourhood*neighbourhood, 6)
    for x = -len:len
        for y = -len:len
            P₂[neighbourhood*(x + len) + (y + len) + 1, :] = [1, x, y, x*x, y*y, x*y]
        end
    end
    return P₂
end

function construct_inverse_metric(P₂::AbstractArray, g::AbstractArray)
    applicability = g*g'
    Wa = Matrix(Diagonal(@view applicability[:]))
    P₂_Wa_P₂ = P₂'*Wa*P₂
    G⁻¹ = pinv(P₂_Wa_P₂)
end

function calculate_correlations(img, len, g)
    # Define basis functions.
    l = collect(-len:len)
    lg = l.*g
    l²g = l.*lg

    # Correlations along the first dimensions.
    unity_first_dimension = imfilter(img, ( ImageFiltering.ReshapedOneD{2,0}(centered(g)),), Fill(0))
    y_first_dimension = imfilter(img, ( ImageFiltering.ReshapedOneD{2,0}(centered(lg)),),  Fill(0))
    y²_first_dimension = imfilter(img, ( ImageFiltering.ReshapedOneD{2,0}(centered(l²g)),), Fill(0))

    # Correlations along the second dimensions.
    # Figure 4.1 (Follow correlation structure) 1, y, y², x , xy, x²
    unity = imfilter(unity_first_dimension, ( ImageFiltering.ReshapedOneD{2,1}(centered(g)),), Fill(0))
    x =  imfilter(unity_first_dimension, ( ImageFiltering.ReshapedOneD{2,1}(centered(lg)),), Fill(0))
    x² =  imfilter(unity_first_dimension, ( ImageFiltering.ReshapedOneD{2,1}(centered(l²g)),), Fill(0))

    y =  imfilter(y_first_dimension, ( ImageFiltering.ReshapedOneD{2,1}(centered(g)),), Fill(0))
    xy = imfilter(y_first_dimension, ( ImageFiltering.ReshapedOneD{2,1}(centered(lg)),), Fill(0))

    y² = imfilter(y²_first_dimension, ( ImageFiltering.ReshapedOneD{2,1}(centered(g)),), Fill(0))

    return unity, x, x², y, xy, y²
end


text:
include("../src/Util_multi.jl")
using MultipleParticle

function run_smc(N, lbs, ubs, f, ϵ₀, Δϵ_tol, α, Δϵ_window, output;
                 save_interval=0, target=0.234 , parallel=false, N_sample=100, γ=1., nu=1.)
    mkdir(output)
    lap_start = time()
    lap_end = 0.
    stepsize_flag = 0
    particles = generate_init_particles(N, lbs, ubs, f)

    ϵ_old = ϵ₀
    ϵ_new = 0.0
    cnt = 0
    terminate = false
    ESS=N
    Δϵ_mean = zeros(Δϵ_window)
    logfilename = joinpath(output, "log.csv")

    while terminate==false
        lap_end = time()
        lap = lap_end - lap_start
        lap_start = lap_end
        open(logfilename, "a") do fp
            writecsv(fp, Array([cnt, ϵ_old, ESS, lap])')
        end
        cnt += 1

        ϵ_new = calcPA(particles, α)
        Δϵ = ϵ_old - ϵ_new
        pop!(Δϵ_mean);
        Δϵ_mean = vcat(Δϵ, Δϵ_mean)
        if mean(Δϵ_mean) < Δϵ_tol && cnt > (Δϵ_window+100)
            terminate=true
        end

        # resampling
        candidates = deepcopy(particles)
        while ESS < N
            candidates = vcat(candidates, resample(particles, N-ESS))
            ESS = length(candidates)
        end

        F_old = get_obj_func_vals(candidates)
        candidates = reset_weights(candidates)
        accept_flags = zeros(length(candidates))
        if parallel == true
            candidates, accept_flags = MCstep_all_parallel(candidates, ϵ_old)
        else
            candidates, accept_flags = MCstep_all(candidates, ϵ_old)
        end

        particles = update_weights(candidates, F_old, ϵ_old, ϵ_new)
        particles = drop_zero_weights(particles)
        particles = normalize_weights(particles)

        acceptance = length(particles)/length(candidates)
        MH_acceptance = sum(accept_flags)/length(accept_flags)
        open(joinpath(output, "acceptance.csv"), "a") do fp
            writecsv(fp, Array([MH_acceptance acceptance]))
        end

        if MH_acceptance < target
            stepsize_flag = 1
        end
        if stepsize_flag == 0
            particles, γ, nu = update_stepsize_inRKHS_all(particles, γ, nu, target, MH_acceptance,
                                                          parallel=parallel, N=N_sample, adaptive=false)
        else
            particles, γ, nu = update_stepsize_inRKHS_all(particles, γ, nu, target, MH_acceptance,
                                                          parallel=parallel, N=N_sample, adaptive=true)
        end

        if length(particles) == 0
            println("particles were vanished at ϵ=$ϵ_new")
            break
        end

        if save_interval==0
            nothing
        elseif mod(cnt, save_interval) == 0 && terminate == false
            est_params = (lbs.*(ubs./lbs).^(hcat(get_values(particles)...)./10))'
            open(joinpath(output, @sprintf("est_param_%05d.csv", cnt)), "w") do fp
                writecsv(fp, est_params)
            end
        elseif terminate == true
            ESS = length(particles)
            lap_end = time()
            lap = lap_end - lap_start
            open(logfilename, "a") do fp
                writecsv(fp, Array([cnt, ϵ_old, Int(floor(ESS)), lap])')
            end

            est_params = (lbs.*(ubs./lbs).^(hcat(get_values(particles)...)./10))'
            open(joinpath(output, @sprintf("est_param_%05d.csv", cnt)), "w") do fp
                writecsv(fp, est_params)
            end
            break
        end

        ϵ_old = ϵ_new
        ESS = length(particles)
    end
end


text:
"""
    MPS

A finite size matrix product state type.
Keeps track of the orthogonality center.
"""
mutable struct MPS <: AbstractMPS
  length::Int
  data::Vector{ITensor}
  llim::Int
  rlim::Int
  function MPS(N::Int,
               A::Vector{<:ITensor},
               llim::Int = 0,
               rlim::Int = N+1)
    new(N, A, llim, rlim)
  end
end

"""
    MPS()

Construct an empty MPS with zero sites.
"""
MPS() = MPS(0, Vector{ITensor}(), 0, 0)

"""
    MPS(N::Int)

Construct an MPS with N sites with default constructed
ITensors.
"""
MPS(N::Int) = MPS(N, Vector{ITensor}(undef, N))

"""
    MPS(::Type{T<:Number}, sites)

Construct an MPS from a collection of indices
with element type T.
"""
function MPS(::Type{T}, sites) where {T<:Number}
  N = length(sites)
  v = Vector{ITensor}(undef, N)
  l = [Index(1, "Link,l=$ii") for ii=1:N-1]
  for ii in eachindex(sites)
    s = sites[ii]
    if ii == 1
      v[ii] = ITensor(T,l[ii], s)
    elseif ii == N
      v[ii] = ITensor(T,l[ii-1], s)
    else
      v[ii] = ITensor(T,l[ii-1],s,l[ii])
    end
  end
  return MPS(N, v)
end

"""
    MPS(sites)

Construct an MPS from a collection of indices
with element type Float64.
"""
MPS(sites) = MPS(Float64, sites)

function randomU(s1,s2)
  if !hasqns(s1) && !hasqns(s2)
    mdim = dim(s1)*dim(s2)
    RM = randn(mdim,mdim)
    Q,_ = NDTensors.qr_positive(RM)
    G = itensor(Q,dag(s1),dag(s2),s1',s2')
  else
    M = randomITensor(QN(),s1',s2',dag(s1),dag(s2))
    U,S,V = svd(M,(s1',s2'))
    u = commonind(U,S)
    v = commonind(S,V)
    G = (U*delta(dag(u),v))*V
  end
  return G
end

function randomizeMPS!(M::MPS, sites, linkdim=1)
  N = length(sites)
  c = div(N,2)
  max_pass = 100
  for pass=1:max_pass,half=1:2
    if half==1
      (db,brange) = (+1, 1:1:N-1)
    else
      (db,brange) = (-1, N:-1:2)
    end
    for b=brange
      s1 = sites[b]
      s2 = sites[b+db]
      G = randomU(s1,s2)
      T = noprime(G*M[b]*M[b+db])
      rinds = uniqueinds(M[b],M[b+db])
      U,S,V = svd(T,rinds;maxdim=linkdim)
      M[b] = U
      M[b+db] = S*V
      M[b+db] /= norm(M[b+db])
    end
    if half==2 && dim(commonind(M[c],M[c+1])) >= linkdim
      break
    end
  end
  setleftlim!(M, 0)
  setrightlim!(M, 2)
  if dim(commonind(M[c],M[c+1])) < linkdim
    error("MPS center bond dim less than requested")
  end
end

function randomCircuitMPS(sites,linkdim::Int;kwargs...)::MPS
  N = length(sites)
  M = MPS(N)

  l = Vector{Index}(undef,N)

  d = dim(sites[N])
  chi = min(linkdim,d)
  l[N-1] = Index(chi,"Link,l=$(N-1)")
  O = NDTensors.random_orthog(chi,d)
  M[N] = itensor(O,l[N-1],sites[N])

  for j=N-1:-1:2
    chi *= dim(sites[j])
    chi = min(linkdim,chi)
    l[j-1] = Index(chi,"Link,l=$(j-1)")
    O = NDTensors.random_orthog(chi,dim(sites[j])*dim(l[j]))
    T = reshape(O,(chi,dim(sites[j]),dim(l[j])))
    M[j] = itensor(T,l[j-1],sites[j],l[j])
  end

  O = NDTensors.random_orthog(1,dim(sites[1])*dim(l[1]))
  l0 = Index(1,"Link,l=0")
  T = reshape(O,(1,dim(sites[1]),dim(l[1])))
  M[1] = itensor(T,l0,sites[1],l[1])
  M[1] *= setelt(l0=>1)

  M.llim = 0
  M.rlim = 2

  return M
end

"""
    randomMPS(::Type{T<:Number}, sites; linkdim=1)

Construct a random MPS with link dimension `linkdim` of
type `T`.
"""
function randomMPS(::Type{T}, sites, linkdim::Int=1) where {T<:Number}
  if hasqns(sites[1])
    error("initial state required to use randomMPS with QNs")
  end

  # For non-QN-conserving MPS, instantiate
  # the random MPS directly as a circuit:
  return randomCircuitMPS(sites,linkdim)
end

"""
    randomMPS(sites; linkdim=1)

Construct a random MPS with link dimension `linkdim` of
type `Float64`.
"""
randomMPS(sites, linkdim::Int=1) = randomMPS(Float64, sites, linkdim)

"""
    randomMPS(sites,state; linkdim=1)

Construct a real, random MPS with link dimension `linkdim`,
made by randomizing an initial product state specified by
`state`.
"""
function randomMPS(sites,state,linkdim::Int=1)::MPS
  M = productMPS(sites,state)
  if linkdim > 1
    randomizeMPS!(M,sites,linkdim)
  end
  return M
end

"""
    productMPS(::Type{T<:Number}, ivals::Vector{<:IndexVal})

Construct a product state MPS with element type `T` and
nonzero values determined from the input IndexVals.
"""
function productMPS(::Type{T},
                    ivals::Vector{<:IndexVal}) where {T<:Number}
  N = length(ivals)
  M = MPS(N)
  if hasqns(ind(ivals[1]))
    links = [Index(QN()=>1;tags="Link,l=$n") for n=1:N]
  else
    links = [Index(1,"Link,l=$n") for n=1:N]
  end
  M[1] = ITensor(T,ind(ivals[1]),links[1])
  M[1][ivals[1],links[1](1)] = 1.0
  for n=2:N-1
    s = ind(ivals[n])
    M[n] = ITensor(T,dag(links[n-1]),s,links[n])
    M[n][links[n-1](1),ivals[n],links[n](1)] = 1.0
  end
  M[N] = ITensor(T,dag(links[N-1]),ind(ivals[N]))
  M[N][links[N-1](1),ivals[N]] = 1.0
  return M
end

"""
    productMPS(ivals::Vector{<:IndexVal})

Construct a product state MPS with element type `Float64` and
nonzero values determined from the input IndexVals.
"""
productMPS(ivals::Vector{<:IndexVal}) = productMPS(Float64,
                                                   ivals::Vector{<:IndexVal})

function productMPS(::Type{T},
                    sites,
                    states) where {T<:Number}
  if length(sites) != length(states)
    throw(DimensionMismatch("Number of sites and and initial states don't match"))
  end
  ivals = [state(sites[n],states[n]) for n=1:length(sites)]
  return productMPS(T, ivals)
end

productMPS(sites, states) = productMPS(Float64, sites, states)

function linkind(M::MPS, j::Int)
  N = length(M)
  j ≥ length(M) && error("No link index to the right of site $j (length of MPS is $N)")
  li = commonind(M[j],M[j+1])
  if isnothing(li)
    error("linkind: no MPS link index at link $j")
  end
  return li
end

function siteind(M::MPS, j::Int)
  N = length(M)
  if j == 1
    si = uniqueind(M[j],M[j+1])
  elseif j == N
    si = uniqueind(M[j],M[j-1])
  else
    si = uniqueind(M[j],M[j-1],M[j+1])
  end
  return si
end

function siteinds(M::MPS)
  return [siteind(M,j) for j in 1:length(M)]
end

function replacesiteinds!(M::MPS, sites)
  for j in eachindex(M)
    sj = siteind(M,j)
    replaceind!(M[j], sj, sites[j])
  end
  return M
end

"""
    dot(psi::MPS, phi::MPS; make_inds_match = true)
    inner(psi::MPS, phi::MPS; make_inds_match = true)

Compute <psi|phi>.

If `make_inds_match = true`, the function attempts to make
the site indices match before contracting (so for example, the
inputs can have different site indices, as long as they
have the same dimensions or QN blocks).
"""
function LinearAlgebra.dot(M1::MPS, M2::MPS; make_inds_match::Bool = true)::Number
  N = length(M1)
  if length(M2) != N
    throw(DimensionMismatch("inner: mismatched lengths $N and $(length(M2))"))
  end
  M1dag = dag(M1)
  if make_inds_match
    simlinkinds!(M1dag)
    replacesiteinds!(M1dag, siteinds(M2))
  end
  O = M1dag[1]*M2[1]
  for j in eachindex(M1)[2:end]
    O = (O*M1dag[j])*M2[j]
  end
  return O[]
end

inner(M1::MPS, M2::MPS; kwargs...) = dot(M1, M2; kwargs...)

"""
    replacebond!(M::MPS, b::Int, phi::ITensor; kwargs...)

Factorize the ITensor `phi` and replace the ITensors
`b` and `b+1` of MPS `M` with the factors. Choose
the orthogonality with `ortho="left"/"right"`.
"""
function replacebond!(M::MPS,
                      b::Int,
                      phi::ITensor;
                      kwargs...)
  ortho::String = get(kwargs, :ortho, "left")
  which_decomp::Union{String, Nothing} = get(kwargs, :which_decomp, nothing)
  normalize::Bool = get(kwargs, :normalize, false)

  # Deprecated keywords
  if haskey(kwargs, :dir)
    error("""dir keyword in replacebond! has been replaced by ortho.
          Note that the options are now the same as factorize, so use `left` instead of `fromleft` and `right` instead of `fromright`.""")
  end

  L,R,spec = factorize(phi,inds(M[b]); which_decomp = which_decomp,
                                       tags = tags(linkind(M,b)),
                                       kwargs...)
  M[b]   = L
  M[b+1] = R
  if ortho == "left"
    leftlim(M) == b-1 && setleftlim!(M, leftlim(M)+1)
    rightlim(M) == b+1 && setrightlim!(M, rightlim(M)+1)
    normalize && (M[b+1] ./= norm(M[b+1]))
  elseif ortho == "right"
    leftlim(M) == b && setleftlim!(M, leftlim(M)-1)
    rightlim(M) == b+2 && setrightlim!(M, rightlim(M)-1)
    normalize && (M[b] ./= norm(M[b]))
  else
    error("In replacebond!, got ortho = $ortho, only currently supports `left` and `right`.")
  end
  return spec
end

"""
    sample!(m::MPS)

Given a normalized MPS m, returns a `Vector{Int}`
of `length(m)` corresponding to one sample
of the probability distribution defined by
squaring the components of the tensor
that the MPS represents. If the MPS does
not have an orthogonality center,
orthogonalize!(m,1) will be called before
computing the sample.
"""
function sample!(m::MPS)
  orthogonalize!(m, 1)
  return sample(m)
end

"""
    sample(m::MPS)

Given a normalized MPS m with `orthocenter(m)==1`,
returns a `Vector{Int}` of `length(m)`
corresponding to one sample of the
probability distribution defined by
squaring the components of the tensor
that the MPS represents
"""
function sample(m::MPS)
  N = length(m)

  if orthocenter(m) != 1
    error("sample: MPS m must have orthocenter(m)==1")
  end
  if abs(1.0-norm(m[1])) > 1E-8
    error("sample: MPS is not normalized, norm=$(norm(m[1]))")
  end

  result = zeros(Int,N)
  A = m[1]

  for j=1:N
    s = siteind(m,j)
    d = dim(s)
    # Compute the probability of each state
    # one-by-one and stop when the random
    # number r is below the total prob so far
    pdisc = 0.0
    r = rand()
    # Will need n,An, and pn below
    n = 1
    An = ITensor()
    pn = 0.0
    while n <= d
      projn = ITensor(s)
      projn[s[n]] = 1.0
      An = A*projn
      pn = real(scalar(dag(An)*An))
      pdisc += pn
      (r < pdisc) && break
      n += 1
    end

    result[j] = n

    if j < N
      A = m[j+1]*An
      A *= (1.0/sqrt(pn))
    end
  end
  return result
end


text:
using FuzzySets, Test

empty = Interval() # empty interval
@test isnan(empty.left)
@test isnan(empty.right)

a = Interval(2, 4)
@test mid(a) == 3
@test rad(a) == 1

b = Interval(3, 5)
@test b == Interval(3, 5)

@test a + b == Interval(5, 9)
@test a - b == Interval(-3, 1)
@test a * b == Interval(6, 20)
@test a / b == Interval(2/5, 4/3)

c = 1 .. 2
@test c * (c - c) == Interval(-2, 2)
@test c * c - c * c == Interval(-3, 3)

a⃗ = [a, b, c]
for (idx, x) in enumerate(zip(a, b))
    if idx == 1
        @test x == (2, 3)
    elseif idx == 2
        @test x == (4, 5)
    end
end

# Fuzzy distances
a = Interval(1.5, 2.5)
b = Interval(1.5, 2.5)
@test a - b == Interval(-1.0, 1.0)
@test (a - b) ^ 2 == Interval(0, 1.0)

a = Interval(0.5, 1.5)
b = Interval(0.6, 1.6)
@test a - b == Interval(-1.1, 0.9)
@test (a - b) ^ 2 ≈ Interval(0, 1.21)

a = Interval(0, 2.21)
b = Interval(0, 1.48)
@test a - b == Interval(-1.48, 2.21)
@test (a - b) ^ 2 == Interval(0, 4.8841)

@test 1 / Interval(0.4, 1.6) == Interval(0.625, 2.5)
@test 1 / Interval(-0.4, 1.6) == Interval(-2.5, 0.625)
@test 1 / Interval(0, 25) == Interval(0, 0.04)
@test 1 / (1 + Interval(0.4, 4)) ≈ Interval(0.2, 0.7142857142857143)

a = Interval(0.3, 0.6)
b = Interval(0.4, 0.7)
@test a - b == Interval(0.3-0.7, 0.6-0.4)
@test (a - b)^2 == Interval(0, max((0.3-0.7)^2, (0.6-0.4)^2))


text:
struct BinCovFactor{T<:Unsigned} <: AbstractFactor{T}
    name::String
    levels::AbstractVector{<:AbstractString}
    bins::AbstractVector{<:Real}
    covariate::AbstractCovariate{<:AbstractFloat}
end

Base.length(var::BinCovFactor{T}) where {T<:Unsigned} = length(var.covariate)

function BinCovFactor(name::String, bins::AbstractVector{T}, covariate::AbstractCovariate{S}) where {T<:Real} where {S<:AbstractFloat}
    bins = issorted(bins) ? bins : sort(bins)
    levelcount = length(bins) - 1
    levels = [@sprintf("[%G,%G%s", bins[i], bins[i + 1], i == levelcount ? "]" : ")") for i in 1:levelcount]
    if levelcount <= typemax(UInt8)
        BinCovFactor{UInt8}(name, levels, bins, covariate)
    elseif levelcount <= typemax(UInt16)
        BinCovFactor{UInt16}(name, levels, bins, covariate)
    else
        BinCovFactor{UInt32}(name, levels, bins, covariate)
    end
end

function factor(covariate::AbstractCovariate{S}, bins::AbstractVector{T}) where {T<:Real} where {S<:AbstractFloat}
    BinCovFactor(getname(covariate), bins, covariate)
end

function factor(covariate::AbstractCovariate{S}) where {S<:AbstractFloat}
    fromobs = 1
    toobs = length(covariate)
    slicelength = verifyslicelength(fromobs, toobs, SLICELENGTH)
    slices = slice(covariate, fromobs, toobs, slicelength)
    vset = fold(Set{S}(), slices) do acc, slice
        for i in 1:length(slice)
            v = slice[i]
            if !isnan(v)
                push!(acc, v)
            end
        end
        acc
    end
    uniq = collect(vset)
    sort!(uniq)
    m = length(uniq)
    bins = Vector{S}(undef, m + 1)
    bins[1] = -Inf32
    for i in 2:(m + 1)
        if i == (m + 1)
            bins[i] = Inf32
        else
            bins[i] = uniq[i - 1] + 0.5 * (uniq[i] - uniq[i - 1])
        end
    end
    BinCovFactor(getname(covariate), bins, covariate)
end

function factor(covariate::GroupStatsCovariate{N, U, Int64, T}) where {T<:AbstractFloat} where {N} where {U}
    if covariate.transform == identity
        bins = unique(collect(values(covariate.groupstats.stats)))
        BinCovFactor(getname(covariate), bins, covariate)
    else
        fromobs = 1
        toobs = length(covariate)
        slicelength = verifyslicelength(fromobs, toobs, SLICELENGTH)
        slices = slice(covariate, fromobs, toobs, slicelength)
        vset = fold(Set{S}(), slices) do acc, slice
            for i in 1:length(slice)
                v = slice[i]
                if !isnan(v)
                    push!(acc, v)
                end
            end
            acc
        end
        bins = collect(vset)
        sort!(bins)
        BinCovFactor(getname(covariate), bins, covariate)
    end
end

function slice(factor::BinCovFactor{T}, fromobs::Integer, toobs::Integer, slicelength::Integer) where {T<:Unsigned}
    bins = factor.bins
    binlenm1 = convert(T, length(bins) - 1)
    z = bins[length(bins)]
    f = (x ->
            begin
                if isnan(x)
                    return zero(T)
                else
                    i = searchsortedlast(bins, x)
                    return convert(T, i)
                end
            end
        )
    slicelength = verifyslicelength(fromobs, toobs, slicelength)
    slices = slice(factor.covariate, fromobs, toobs, slicelength)
    mapslice(f, slices, slicelength, T)
end

function Base.map(factor::BinCovFactor{T}, dataframe::AbstractDataFrame; permute::Bool = false) where {T<:Unsigned}
    BinCovFactor(getname(factor), factor.bins, map(factor.covariate, dataframe))
end

function isordinal(factor::BinCovFactor{T}) where {T<:Unsigned}
    true
end

max_stars_repo_name: Statfactory/JuML

max_issues_repo_name: Statfactory/JuML

max_forks_repo_name: Statfactory/JuML


text:
@test cleanzipcode("01234") == "01234"
@test cleanzipcode("1234") == "01234"
@test cleanzipcode(" 1234 ") == "01234"
@test cleanzipcode(" 1234-9999") == "01234"
@test ismissing(cleanzipcode("a"))
@test cleanzipcode("a", enforcestring = true) == "NA"

@test cleanzipcode(1234) == "01234"
@test cleanzipcode(0123456789, returnNA=false) == "123456789"

zipvec1 = ["1234", 1234, "1234-1234"]
@test cleanzipcode(zipvec1, enforcestring=true) == ["01234", "01234", "01234"]

cleaner1 = removewhitespace()
cleaner2 = removezipsuffix()
cleaner3 = padleftzeros()

@test occursin(cleaner1.pattern, " 01234 ")
@test cleaner1.cleanfun(" 01234 ") == "01234"
@test occursin(cleaner2.pattern, "01234-9999")
@test cleaner2.cleanfun("01234-9999") == "01234"
@test occursin(cleaner3.pattern, "1234")
@test cleaner3.cleanfun("1234") == "01234"


text:
module AtomIteratorTest

using Test
using ProximalBase
using CoordinateDescent
using Random

Random.seed!(1)

@testset "Ordered" begin
    p = 5
    x = SparseIterate(p)
    x[2] = 1.
    x[1] = 2.

    it = CoordinateDescent.OrderedIterator(x)

    fullPass   = collect(1:5)
    sparsePass = [2, 1]

    @test collect(it) == fullPass           # this should be a full pass

    CoordinateDescent.reset!(it, true)
    @test collect(it) == fullPass           # this should be a full pass

    CoordinateDescent.reset!(it, false)
    @test collect(it) == sparsePass         # pass over non-zeros

    # SymmetricSparseIterate
    x = SymmetricSparseIterate(3)
    x[2] = 1.
    x[1] = 2.

    it = CoordinateDescent.OrderedIterator(x)

    fullPass   = collect(1:6)
    sparsePass = [2, 1]

    @test collect(it) == fullPass           # this should be a full pass

    CoordinateDescent.reset!(it, true)
    @test collect(it) == fullPass           # this should be a full pass

    CoordinateDescent.reset!(it, false)
    @test collect(it) == sparsePass         # pass over non-zeros
end

@testset "Random" begin
    p = 50
    s = 10

    x = SparseIterate(p)
    for i=1:s
      x[rand(1:p)] = randn()
    end

    it = CoordinateDescent.RandomIterator(x)

    @test collect(it) == collect(1:ProximalBase.numCoordinates(it.iterate))   # this should be a full pass over 1:p

    CoordinateDescent.reset!(it, true)
    @test collect(it) == it.order                       # this should be a full pass over 1:p in a random order

    CoordinateDescent.reset!(it, false)
    @test collect(it) == [x.nzval2ind[it.order[i]] for i=1:nnz(x)] # this should be a sparse pass

    # SymmetricSparseIterate
    x = SymmetricSparseIterate(10)
    for i=1:s
      x[rand(1:55)] = randn()
    end

    it = CoordinateDescent.RandomIterator(x)

    @test collect(it) == collect(1:ProximalBase.numCoordinates(it.iterate))   # this should be a full pass over 1:p

    CoordinateDescent.reset!(it, true)
    @test collect(it) == it.order                       # this should be a full pass over 1:p in a random order

    CoordinateDescent.reset!(it, false)
    @test collect(it) == [x.nzval2ind[it.order[i]] for i=1:nnz(x)] # this should be a sparse pass

end

end


text:
const NATIVE_LABEL_ENCODING = MLLabelUtils.LabelEnc.NativeLabels(["inlier", "outlier"])
const LABEL_ENCODING = MLLabelUtils.LabelEnc.NativeLabels([:inlier,:outlier])
const LEARNING_LABEL_ENCODING = MLLabelUtils.LabelEnc.NativeLabels([:Lin,:Lout])
const SplitType = Union{Val{:train}, Val{:test}, Val{:query}}

function load_data(file_path; header=false, native_label_encoding=NATIVE_LABEL_ENCODING)
    raw_data, _ = header ? readdlm(file_path, ',', header=header) : (readdlm(file_path, ','), nothing)
    data = copy(transpose(float.(raw_data[:, 1:end-1])))
    labels = convert_labels_from_raw(raw_data[:,end])
    return data, labels
end

function convert_labels_from_raw(raw_labels, native_label_encoding=NATIVE_LABEL_ENCODING)
    return MLLabelUtils.convertlabel(LABEL_ENCODING, raw_labels, native_label_encoding)
end

function convert_labels_to_learning(labels, label_encoding=LABEL_ENCODING)
    return MLLabelUtils.convertlabel(LEARNING_LABEL_ENCODING, labels, label_encoding)
end

abstract type SplitStrategy end

struct FullSplitStrat <: SplitStrategy end
struct UnlabeledSplitStrat <: SplitStrategy end
struct UnlabeledAndLabeledInlierSplitStrat <: SplitStrategy end
struct LabeledSplitStrat <: SplitStrategy end
struct LabeledInlierSplitStrat <: SplitStrategy end
struct LabeledOutlierSplitStrat <: SplitStrategy end

struct DataSplits
    train::BitArray
    test::BitArray
    train_strat::SplitStrategy
    test_strat::SplitStrategy
    query_strat::SplitStrategy
    DataSplits(train::BitArray) = DataSplits(train, FullSplitStrat())
    DataSplits(train::BitArray, strat::SplitStrategy) = DataSplits(train, copy(train), strat)
    DataSplits(train::BitArray, train_strat::SplitStrategy, test_strat::SplitStrategy) = new(train, copy(train), train_strat, test_strat, train_strat)
    DataSplits(train::BitArray, train_strat::SplitStrategy, test_strat::SplitStrategy, query_strat::SplitStrategy) = new(train, copy(train), train_strat, test_strat, query_strat)
    DataSplits(train::BitArray, test::BitArray, strat::SplitStrategy) = new(train, test, strat, strat, strat)
    DataSplits(train::BitArray, test::BitArray, train_strat::SplitStrategy, test_strat::SplitStrategy) = new(train, test, train_strat, test_strat, train_strat)
end

get_mask(ds, pools, ::Val{:train}) = calc_mask(ds.train_strat, ds.train, pools)
get_mask(ds, pools, ::Val{:test}) = calc_mask(ds.test_strat, ds.test, pools)
get_mask(ds, pools, ::Val{:query}) = calc_mask(ds.query_strat, ds.train, pools)

function get_local_idx(global_idx::Int, ds, pools, split::T) where T <: SplitType
    mask = get_mask(ds, pools, split)
    mask[global_idx] || error("$(get_val_type(split)) split does not contain observation with global idx $(global_idx).")
    local_idx = count(mask[1:global_idx])
    return local_idx
end

"""
    filter_query_id(query_ids::Vector{Int}, split_strategy, query_pool_labels, ::Val{:train})

    Returns the query ids that also are in the train split.
"""
function filter_query_id(query_ids::Vector{Int}, split_strategy, query_pool_labels, ::Val{:train})
    query_id_mask = calc_mask(split_strategy.train_strat, split_strategy.train[query_ids], query_pool_labels)
    return query_ids[query_id_mask]
end

get_train(ds::DataSplits, data::Array{T, 2}, pools::Vector{Symbol}) where T <: Real = select_subset(ds.train_strat, ds.train, data, pools)
get_test(ds::DataSplits, data::Array{T, 2}, pools::Vector{Symbol}) where T <: Real = select_subset(ds.test_strat, ds.test, data, pools)
get_query(ds::DataSplits, data::Array{T, 2}, pools::Vector{Symbol}) where T <: Real = select_subset(ds.query_strat, ds.train, data, pools)

function select_subset(strat::SplitStrategy, init_mask::BitArray, data::Array{T, 2}, pools::Vector{Symbol}) where T <: Real
    mask = calc_mask(strat, init_mask, pools)
    return (data[:, mask], pools[mask], findall(mask))
end

calc_mask(strat::FullSplitStrat, init_mask, pools) = init_mask
calc_mask(strat::UnlabeledSplitStrat, init_mask, pools) = init_mask .& (pools .== :U)
calc_mask(strat::UnlabeledAndLabeledInlierSplitStrat, init_mask, pools) = init_mask .& ((pools .== :U) .| (pools .== :Lin))
calc_mask(strat::LabeledSplitStrat, init_mask, pools) = init_mask .& ((pools .== :Lin) .| (pools .== :Lout))
calc_mask(strat::LabeledInlierSplitStrat, init_mask, pools) = init_mask .& (pools .== :Lin)
calc_mask(strat::LabeledOutlierSplitStrat, init_mask, pools) = init_mask .& (pools .== :Lout)

function get_initial_pools(data, labels, data_splits, initial_pool_strategy; n=20, p=0.1, x=10)
    if initial_pool_strategy ∉ ["Pu", "Pp", "Pn", "Pnin", "Pa"]
        throw(ArgumentError("Unknown initial pools strategy '$(initial_pool_strategy)'."))
    end
    l = fill(:U, size(data, 2))
    if initial_pool_strategy == "Pa"
        label_candidates = findall(data_splits.train .& (labels .== :inlier))
        n = min(size(data, 1) + x, length(label_candidates))
        l[label_candidates[1:n]] .= convert_labels_to_learning(labels[label_candidates[1:n]])
    elseif initial_pool_strategy ∈ ["Pp", "Pn"]
        p = initial_pool_strategy == "Pp" ? p : min(1.0, n / sum(data_splits.train))
        label_candidates = findall(data_splits.train)
        (label_indices, _), _ = MLDataUtils.stratifiedobs((label_candidates, labels[data_splits.train]), p=p)
        l[label_indices] .= convert_labels_to_learning(labels[label_indices])
    elseif initial_pool_strategy == "Pnin"
        label_candidates = findall(data_splits.train .& (labels .== :inlier))
        label_indices = Random.shuffle(label_candidates)[1:min(length(label_candidates), n)]
        l[label_indices] .= convert_labels_to_learning(labels[label_indices])
    end
    return l
end

function get_splits_and_init_pools(data, labels, split_strategy, initial_pool_strategy; holdout_p=0.2, kwargs...)
    train = trues(size(data, 2))
    if split_strategy == "Sf"
        data_splits = DataSplits(train, FullSplitStrat())
    elseif split_strategy == "Sh"
        (train_indices, train_labels), (test_indices, test_labels) = MLDataUtils.stratifiedobs((collect(1:length(labels)), labels), p=1-holdout_p)
        train[test_indices] .= false
        data_splits = DataSplits(train, .~train, FullSplitStrat())
    elseif split_strategy == "Si"
        data_splits = DataSplits(train, LabeledInlierSplitStrat(), FullSplitStrat(), FullSplitStrat())
    elseif split_strategy == "Sl"
        data_splits = DataSplits(train, LabeledSplitStrat(), FullSplitStrat(), LabeledSplitStrat())
    elseif split_strategy == "Slf"
        data_splits = DataSplits(train, LabeledSplitStrat(), FullSplitStrat(), FullSplitStrat())
    else
        throw(ArgumentError("Unknown split strategy '$(split_strategy)'."))
    end
    initial_pools = get_initial_pools(data, labels, data_splits, initial_pool_strategy; kwargs...)
    return data_splits, initial_pools
end


text:
"""
    get_beta -> Vector{Real}

returns the dispersion of the mode at a given wavelength up to
the given order for the structure's modes
"""
function get_beta(mode::Mode, source)
    effectiveindex = mode.effectiveindex
    ω = frequency(effectiveindex)
    length(ω) == 1 && error("the given refractive index is non-dispersive")
    β₀ = ω/c0 .* get_attr(effectiveindex)
    β = OpticalAttr(ω, β₀, "β")
    return β
end

function get_beta(mode::Mode, source, numorders::Int)
    effectiveindex = mode.effectiveindex
    ω = frequency(effectiveindex)
    length(ω) == 1 && error("the given refractive index is non-dispersive")
    β₀ = ω/c0 .* get_attr(effectiveindex)
    β = OpticalAttr(frequency(effectiveindex), β₀, "β")
    ω_query = frequency(source)
    β_atquery = Any[]
    for order=0:numorders
        push!(β_atquery, der(β, ω_query; order=order))
    end
    return β_atquery
end

function get_beta(mode::ToyMode, source, numorders::Int)
    β = mode.beta
    (numorders+1) > length(β) && error("model does not have sufficient orders of β")
    return β[1:(numorders+1)]
end

"""
    get_stationarydispersion(mode, source)

    Returns the stationary dispersion of the mode at the frequency defined by
    source. This dispersion is centered around ω=0 and has no constant or
    linear component - it is the dispersion in the moving frame.
"""
function get_stationarydispersion(mode, source)
    const source_ω = frequency(source)
    const beta = get_beta(mode, source)
    const beta_linear = get_beta(mode, source, 1)
    return ω->(beta(ω) - beta_linear[1] - beta_linear[2] * (ω - source_ω))
end

function get_dispersionfrombeta(beta_coeff)
    return Poly(beta_coeff, :ω)
end

"""
    get_groupindex -> Real

returns the group index of the mode at a given wavelength
"""
function get_groupindex(mode, source)
    β = get_beta(mode, source, 1)
    n_group = β[2] * c0
    n_group
end

"""
    get_FSR -> Real

returns the free spectral range (FSR) of the resonator's modes
"""
function get_FSR(resonator, mode::Mode, source)
    groupindex = get_groupindex(mode, source)
    FSR = abs(c0/(groupindex*pathlength(resonator)))
    FSR
end


text:
using Turing, Random, Test
import Turing.Inference
import AdvancedMH

dir = splitdir(splitdir(pathof(Turing))[1])[1]
include(dir*"/test/test_utils/AllUtils.jl")

@testset "emcee.jl" begin
    @testset "gdemo" begin
        Random.seed!(9876)

        n_samples = 1000
        n_walkers = 250

        spl = Emcee(n_walkers, 2.0)
        chain = sample(gdemo_default, spl, n_samples)

        check_gdemo(chain)
    end
end