export FreeTheory

using Combinatorics

struct FreeTheory <: AbstractTheory end

struct FreeTerm <: AbstractTerm
    root::Σ
    args::Vector{Union{Variable,AbstractTerm}}
end

term(::FreeTheory, root, args) = FreeTerm(root, args)
Base.convert(::Type{Expr}, t::FreeTerm) = Expr(:call, t.root, convert.(Expr, t.args)...)

theory(::Type{FreeTerm}) = FreeTheory()
priority(::Type{FreeTerm}) = 1000

vars(t::FreeTerm) = mapreduce(vars, ∪, t.args; init=Set{Variable}())

function Base.:(==)(a::FreeTerm, b::FreeTerm)
    a.root == b.root || return false

    for i ∈ eachindex(a.args)
        a.args[i] == b.args[i] || return false
    end

    return true
end

function (a::FreeTerm >ₜ b::FreeTerm)
    a.root == b.root || return a.root >ₑ b.root
    @assert length(a.args) == length(b.args)

    for (x, y) ∈ zip(a.args, b.args)
        x == y || return x >ₜ y
    end

    return false
end

function Base.hash(t::FreeTerm, h::UInt)
    init = hash(t.root, hash(FreeTerm, h))
    foldr(hash, t.args; init=init)
end

Base.map(f, t::FreeTerm) = FreeTerm(t.root, map(f, t.args))

@enum FreeKind VAR NODE ALIEN
struct FreeAux
    kind::FreeKind
    idx::UInt
    args::Vector{FreeAux}
end
struct FreeMatcher <: AbstractMatcher
    m::FreeAux
    syms::Vector{Σ}
    vars::Vector{Variable}
    matchers::Vector{AbstractMatcher}
    ϕ::Vector{Int}
end

function _build_free_context!(t::FreeTerm, syms, vars, aliens)
    args = similar(t.args, FreeAux)
    node_idx = findfirst(==(t.root), syms)
    node_idx === nothing && (push!(syms, t.root); node_idx = length(syms))

    for i ∈ eachindex(t.args)
        arg = t.args[i]
        if isa(arg, Variable)
            idx = findfirst(==(arg), vars)
            idx === nothing && (push!(vars, arg); idx = length(vars))
            args[i] = FreeAux(VAR, idx, [])
        elseif isa(arg, FreeTerm)
            args[i] = _build_free_context!(arg, syms, vars, aliens)
        else
            push!(aliens, arg)
            args[i] = FreeAux(ALIEN, length(aliens), [])
        end
    end

    return FreeAux(NODE, node_idx, args)
end

function _find_permutation(aliens, V)
    # TODO: improve efficiency using monotonicity
    matchers, V_best = many_matchers(aliens, V)
    ϕ = eachindex(aliens)

    for ψ ∈ permutations(eachindex(aliens))
        result, V′ = many_matchers(aliens[ψ], V)
        if length(V′) > length(V_best)
            (matchers, V_best, ϕ) = (result, num_fixed, ψ)
        end
    end

    return (matchers, V_best, ϕ)
end

function matcher(t::FreeTerm, V)
    syms = Σ[]
    vars = Variable[]
    aliens = AbstractTerm[]
    m = _build_free_context!(t, syms, vars, aliens)

    (matchers, V′, ϕ) = _find_permutation(aliens, V ∪ vars)
    return FreeMatcher(m, syms, vars, matchers, ϕ), V′
end

struct FreeSubproblem <: AbstractSubproblem
    subproblems::Vector{AbstractSubproblem}
end

function match!(σ, A::FreeMatcher, t::FreeTerm)
    if isempty(A.matchers)
        match_aux!(A.m, A, σ, nothing, t) === 0 || return nothing
        return EmptySubproblem()
    end

    aliens = similar(A.matchers, AbstractTerm)
    aliens_found = match_aux!(A.m, A, σ, aliens, t)
    aliens_found === length(A.matchers) || return nothing

    subproblems = similar(A.matchers, AbstractSubproblem)
    for i ∈ A.ϕ
        a_match = match!(σ, A.matchers[i], aliens[i])
        a_match === nothing && return nothing
        subproblems[i] = a_match
    end

    FreeSubproblem(subproblems)
end

function match_aux!(m, A, σ, aliens, t)
    if m.kind === VAR
        x = A.vars[m.idx]
        if haskey(σ, x)
            σ[x] == t || return -1
        else
            σ[x] = t
        end
        return 0
    end

    if m.kind === NODE
        isa(t, FreeTerm) || return -1

        t.root === A.syms[m.idx] || return -1
        length(t.args) === length(m.args) || return -1
        res = 0
        for (arg, aux) ∈ zip(t.args, m.args)
            res += match_aux!(aux, A, σ, aliens, arg)
        end
        return res
    else
        aliens[m.idx] = t
        return 1
    end
end

function compile(A::FreeMatcher, V)
    fn_name = gensym(:match!_free)
    V′ = V ∪ A.vars

    aliens = similar(A.matchers, AbstractTerm)
    subproblems = similar(A.matchers, AbstractSubproblem)

    vars_seen = [x ∈ V for x ∈ A.vars]
    recursive_matchers = _compile_free_aux(A.m, :t, aliens, A.syms, A.vars, vars_seen)

    alien_matches = []
    matcher_fns = []
    for i ∈ A.ϕ
        a_match = gensym(:a_match)
        fn, fn_expr = compile(A.matchers[i], V′)
        append!(matcher_fns, fn_expr.args)
        body = quote
            $a_match = $fn(σ, $aliens[$i])
            $a_match === nothing && return nothing
            $subproblems[$i] = $a_match
        end
        @assert body.head === :block
        append!(alien_matches, body.args)
    end

    subs_expr = isempty(subproblems) ? EmptySubproblem() :
                length(subproblems) == 1 ? :($subproblems[1]) : FreeSubproblem(subproblems)

    fn_name, quote
        $(matcher_fns...)
        function $fn_name(σ, t)
            $(recursive_matchers.args...)
            $(alien_matches...)
            $subs_expr
        end
    end
end
function _compile_free_aux(m, t, aliens, syms, vars, vars_seen)
    idx = m.idx

    if m.kind === VAR
        x = vars[idx]
        if vars_seen[idx]
            return quote
                σ[$x] == $t || return nothing
            end
        else
            vars_seen[idx] = true
            return quote
                σ[$x] = $t
            end
        end
    end

    if m.kind === NODE
        t′ = t
        t = gensym(:t)
        recs = [_compile_free_aux(arg, :($t.args[$i]), aliens, syms, vars, vars_seen)
                for (i, arg) ∈ enumerate(m.args)]
        recursive_calls = []
        for rec ∈ recs
            if rec.head === :block
                append!(recursive_calls, rec.args)
            else
                push!(recursive_calls, rec)
            end
        end

        return quote
            $t = $t′
            isa($t, $FreeTerm) || return
            $t.root == $(Meta.quot(syms[idx])) || return
            length($t.args) == $(length(m.args)) || return
            $(recursive_calls...)
        end
    else
        quote
            $aliens[$(m.idx)] = $t
        end
    end
end

Base.iterate(iter::Matches{FreeSubproblem}) = _aiterate(iter.p, iter.s.subproblems...)
Base.iterate(iter::Matches{FreeSubproblem}, st) = _aiterate1(iter.p, (iter.s.subproblems...,), st)

struct FreeRewriter <: AbstractRewriter
    rules::Dict{Σ,Vector{Pair{FreeMatcher,Any}}}
end

rewriter(::FreeTheory) = FreeRewriter(Dict{Σ,Vector{Pair{FreeMatcher,Any}}}())
function Base.push!(rw::FreeRewriter, (p, b)::Pair{FreeTerm})
    haskey(rw.rules, p.root) || (rw.rules[p.root] = Pair{FreeMatcher,Any}[])
    push!(rw.rules[p.root], matcher(p) => b)
    rw
end

function compile(rw::FreeRewriter)
    fn_name = gensym(:rewrite_free)

    matcher_exprs = Expr[]

    rules = Dict()
    for (root, matchers) ∈ rw.rules
        rules[root] = []
        for (p, b) ∈ matchers
            fn, expr = compile(p)
            push!(matcher_exprs, expr)
            push!(rules[root], fn => b)
        end
    end

    index = gensym(:rewrite_free_tree)

    fn_name, quote
        $(matcher_exprs...)
        $index = Dict($((
            :($(Meta.quot(root)) => $(:([$((:($fn => $b) for (fn, b) ∈ rs)...)])))
            for (root, rs) ∈ rules
        )...))
        function $fn_name(t)
            $haskey($rules, t.root) || return

            σ = $Substitution()

            for (match_fn, builder) ∈ $index[t.root]
                m = match_fn(σ, t)
                m === nothing && continue
                next = iterate($Matches(σ, m))
                next === nothing && continue
                σ = next[1]
                return builder(σ)::$AbstractTerm
            end

            return
        end
    end
end


type Objective
    func
    args
    dim
    inherit
    constraint
    history

    function Objective(dim; func= Nullable(), args=Nullable(),
        inherit=obj_default_inherit, constraint=Nullable())
        return new(func, args, dim, obj_default_inherit, constraint, Float64[])
    end
end

function obj_construct_solution(objective::Objective, x; parent=Nullable())
    sol = Solution()
    sol.x = x
    sol.attach = objective.inherit(parent=parent)
    return sol
end

# evaluate the objective function of a solution
function obj_eval(objective, solution)
    solution.value = objective.func(solution, objective.args)
    push!(objective.history, solution.value)
    return solution.value
end

function obj_eval_constraint(objective, solution)
  #Todo
end

function obj_default_inherit(; parent=Nullable())
    return Nullable()
end

function get_history_bestsofar(objective)
    history_bestsofar = Float64[]
    bestsofar = Inf
    for i in 1:length(objective.history)
        if objective.history[i] < bestsofar
            bestsofar = objective.history[i]
        end
        push!(history_bestsofar, bestsofar)
    end
    return history_bestsofar
end

function obj_clean_history(objective::Objective)
    objective.history = []
end


using Test

# the source code to be tested
include("../src/SBRG.jl")

# @testset "SUVector" begin
@testset "constructor" begin
    vs = (
        [1, 2, 5, 5, 7],
        [1, 2, 5, 7, 7],
        [1, 1, 2, 5, 7],
        [1, 5, 2, 5, 7],
        [1, 5, 2, 7],
        [2, 1, 5, 7],
        [1, 2, 7, 5],
    )
    for v in vs
        @test_throws AssertionError SUVector(v)
    end
end
@testset "indexing, push, append" begin
    v = SUVector([1, 5, 10])
    @test size(v) == (3,)
    @test v[2] == 5
    @test v[end] == 10
    @test (v[2] = 4) == 4
    @test_throws AssertionError v[2] = 11
    push!(v, 11)
    append!(v, [14, 18])
    @test v[4] == 11
    @test v[end] == 18
end
@testset "intersect" begin
    v1 = SUVector([1, 2, 3, 4])
    v2 = SUVector([1, 2])
    v3 = SUVector([2, 4])
    v4 = SUVector{Int}()
    args = ((v1, v2), (v1, v3), (v1, v4), (v4, v4))
    outs = (
        SUVector([1, 2]),
        SUVector([2, 4]),
        v4,
        v4,
    )
    for (arg, out) in zip(args, outs)
        @test all(intersect(arg...) .== out)
    end
end
@testset "symdiff" begin
    v1 = SUVector([2, 5, 7, 8, 10, 12, 15])
    v2 = SUVector([5, 8, 11, 12])
    v3 = SUVector{Int}()
    args = ((v1, v2), (v2, v3), (v3, v3))
    outs = (
        SUVector([2, 7, 10, 11, 15]),
        v2,
        v3,
    )
    for (arg, out) in zip(args, outs)
        @test all(symdiff(arg...) .== out)
    end
end
@testset "union" begin
    v1 = SUVector([1, 4, 6, 7])
    v2 = SUVector([2, 5, 6])
    v3 = SUVector([4, 7, 9])
    v4 = SUVector{Int}()
    args = ((v1, v2), (v1, v3), (v1, v4))
    outs = (
        SUVector([1, 2, 4, 5, 6, 7]),
        SUVector([1, 4, 6, 7, 9]),
        v1,
    )
    # check for correctness
    for (arg, out) in zip(args, outs)
        @test all(union(arg...) .== out)
    end
end
# end

@testset "Pauli & Term" begin
    # set up test arguments and outputs
    X1Y2 = Pauli(1, [1, 2], [2])
    X2Y3 = Pauli(1, [2, 3], [3])
    X1X2 = Pauli(0, [1, 2], Int64[])
    Z1Z2 = Pauli(0, Int64[], [1, 2])
    Y1Y2 = Pauli(0, [1, 2], [1, 2])
    T1 = Term(2, X1Y2)
    T2 = Term(3, X2Y3)
    # set up outputs
    X1Y2X2Y3 = Pauli(0, [1, 3], [2, 3])
    X2Y3X1Y2 = Pauli(2, [1, 3], [2, 3])
    T1T2 = Term(6, X1Y2X2Y3)
    # check for correctness
    @test X1Y2 == X1Y2
    @test X1Y2 != X2Y3
    @test X1Y2 * X2Y3 == X1Y2X2Y3
    @test X2Y3 * X1Y2 == X2Y3X1Y2
    @test !commute(X1Y2, X2Y3)
    @test commute(X1X2, Z1Z2)
    @test commute(X1X2, Y1Y2)
    @test commute(Y1Y2, Z1Z2)
    @test_throws AssertionError Pauli(5, Int64[], Int64[])
    @test_throws AssertionError Pauli(2, [2, 1], [1])
    @test all(sites(X1Y2) .== [1, 2])
    @test center(X1Y2) == 1.5
    @test T1*T2 == T1T2
end

@testset "RGHamiltonian" begin
    # some pauli operators on three sites
    X1X2 = Pauli(0, [1,2], Int[]);
    X2X3 = Pauli(0, [2,3], Int[]);
    Z1 = Pauli(0, Int[], [1])
    Z2 = Pauli(0, Int[], [2])
    Z3 = Pauli(0, Int[], [3])
    Os = (X1X2, X2X3, Z1, Z2, Z3)
    # coefficients
    hs = rand(length(Os))
    # hamiltonian
    H = RGHamiltonian([Term(h, O) for (h, O) in zip(hs, Os)])
end


#include("benchmark.jl")

module FlatArraysTests

using Test, FlatArrays, Random
using Base:
    IndexStyle,
    has_offset_axes

for d in 1:3,
    (f, g) in ((firstindex, :firstindex),
               (lastindex,  :lastindex),
               (size,       :size),
               (axes,       :axis),
               (stride,     :stride))
    @eval $(Symbol(g,d))(A) = $f(A, $d)
end

#function memcpy!(A::DenseArray{T}, i::Integer,
#                 B::DenseArray{T}, j::Integer, n::Integer) where {T}
#    @assert firstindex(A) ≤ i ≤ i + n - 1 ≤ lastindex(A)
#    @assert firstindex(B) ≤ j ≤ j + n - 1 ≤ lastindex(B)
#    ccall(:memcpy, Ptr{Cvoid}, (Ptr{Cvoid}, Ptr{Cvoid}, Csize_t),
#          A, B, n*sizeof(T))
#end

function memcpy!(A::DenseArray{T}, B::DenseArray{T}, n::Integer) where {T}
    @assert firstindex(A) == 1 && 0 ≤ n ≤ lastindex(A)
    @assert firstindex(B) == 1 && 0 ≤ n ≤ lastindex(B)
    ccall(:memcpy, Ptr{Cvoid}, (Ptr{Cvoid}, Ptr{Cvoid}, Csize_t),
          A, B, n*sizeof(T))
    return A
end

function samevalues(A::AbstractArray{T}, B::AbstractArray{T},
                    exhaustive::Bool=false) where {T}
    @assert length(A) == length(B)
    j = 0
    for a in A
        j += 1
        B[j] == a || return false
    end
    if exhaustive
        i = 0
        for b in B
            i += 1
            A[i] == b || return false
        end
    end
    return true
end

@testset "FlatArrays" begin
    dims = (2,3,4,5)
    A = randn(dims)
    nelem = length(A)
    nrows = prod(dims[1:2])
    ncols = prod(dims[3:end])
    V = flatten(A, nelem)
    M = flatten(A, nrows, ncols)
    Q = ((firstindex,      1),
         (firstindex1,     1),
         (lastindex,       nelem),
         (lastindex1,      nelem),
         (length,          nelem),
         (ndims,           1),
         (eltype,          eltype(A)),
         (size,           (nelem,)),
         (size1,           nelem),
         (axes,           (1:nelem,)),
         (axis1,           1:nelem),
         (strides,        (1,)),
         (stride1,         1),
         (stride2,         nelem),
         (isflat,          true),
         (has_offset_axes, false),
         (IndexStyle,      IndexLinear()))
    for i in randperm(length(Q))
        (f, r) = Q[i]
        @test f(V) == r
    end
    Q = ((firstindex,      1),
         (firstindex1,     1),
         (firstindex2,     1),
         (lastindex,       nelem),
         (lastindex1,      nrows),
         (lastindex2,      ncols),
         (length,          nelem),
         (ndims,           2),
         (eltype,          eltype(A)),
         (size,           (nrows, ncols)),
         (size1,           nrows),
         (size2,           ncols),
         (axes,           (1:nrows, 1:ncols)),
         (axis1,           1:nrows),
         (axis2,           1:ncols),
         (strides,        (1, nrows)),
         (stride1,         1),
         (stride2,         nrows),
         (stride3,         nelem),
         (isflat,          true),
         (has_offset_axes, false),
         (IndexStyle,      IndexLinear()))
    for i in randperm(length(Q))
        (f, r) = Q[i]
        @test f(M) == r
    end
    @test isflat(A) == true
    @test isflat(A) == true
    V1 = view(A, :,  :, 1:3, 3)  # flat view
    @test isflat(V1) == true
    V2 = view(A, :,  :, 2:3, :) # non-flat view
    @test isflat(V2) == false
    @test reshape(A, nelem) == V
    @test reshape(A, nrows, ncols) == M
    B = similar(A, nrows, ncols)
    @test memcpy!(B, M, nelem) == reshape(A, nrows, ncols)
    @test samevalues(B, M, true)
    @test samevalues(B, V, true)
    @test samevalues(flatten(V1,length(V1)), V1, true)
    @test samevalues(flatten(V2,length(V2)), V2, true)
    @test_throws BoundsError V[-1]
    @test_throws BoundsError V[length(V)+1]
    @test_throws BoundsError M[1,0]
    @test_throws BoundsError M[nrows+1,0] == M[1,1]
end

end # module


if parse(Bool, get(ENV, "CI", "false"))
    @info "CI detected: skipping dataset download"
else
    data_dir = datadep"EMNIST"

    n_features = (28, 28)
    n_targets = 1

    for (name, n_obs) in [(:balanced,  (112800, 18800)),
                          (:byclass,  (697932, 116323)),
                          (:bymerge,  (697932, 116323)),
                          (:digits,  (240000, 40000)),
                          (:letters,  (124800, 20800)),
                          (:mnist,  (60000, 10000))]

        d = EMNIST(name)

        @test d.name == name
        @test d.split == :train
        @test extrema(d.features) == (0, 1)
        @test convert2image(d, 1) isa AbstractMatrix{<:Gray}
        @test convert2image(d, 1:2) isa AbstractArray{<:Gray, 3}

        test_supervised_array_dataset(d;
            n_features, n_targets, n_obs=n_obs[1],
            Tx=Float32, Ty=Int,
            conv2img=true)

        d = EMNIST(name, split=:test, Tx=UInt8)

        @test d.name == name
        @test d.split == :test
        @test extrema(d.features) == (0, 255)
        @test convert2image(d, 1) isa AbstractMatrix{<:Gray}

        test_supervised_array_dataset(d;
            n_features, n_targets, n_obs=n_obs[2],
            Tx=UInt8, Ty=Int,
            conv2img=true)
    end
end


function init_pseudo_observable_mappings!(m::Model990)

    pseudo_names = [:y_t, :y_f_t, :NaturalRate, :π_t, :OutputGap, :ExAnteRealRate, :LongRunInflation,
                    :MarginalCost, :Wages, :FlexibleWages, :Hours, :FlexibleHours, :z_t, :Z_t, :ztil_t,
                    :NominalFFR, :RealFFR, :ExpectedNominalNaturalRate]

    # Create PseudoObservable objects
    pseudo = OrderedDict{Symbol,PseudoObservable}()
    for k in pseudo_names
        pseudo[k] = PseudoObservable(k)
    end

    # Fill in names and reverse transforms
    pseudo[:y_t].name = "Output Growth"
    pseudo[:y_t].longname = "Output Growth Per Capita"

    pseudo[:y_f_t].name = "Flexible Output Growth"
    pseudo[:y_f_t].longname = "Output that would obtain in a flexible-price economy."

    pseudo[:NaturalRate].name = "Real Natural Rate"
    pseudo[:NaturalRate].longname = "The real interest rate that would prevail in a flexible-price economy."
    pseudo[:NaturalRate].rev_transform = quartertoannual

    pseudo[:π_t].name = "Inflation"
    pseudo[:π_t].longname = "Inflation"
    pseudo[:π_t].rev_transform = quartertoannual

    pseudo[:OutputGap].name = "Output Gap"
    pseudo[:OutputGap].longname = "Output Gap"

    pseudo[:ExAnteRealRate].name = "Ex Ante Real Rate"
    pseudo[:ExAnteRealRate].longname = "Ex Ante Real Rate"
    pseudo[:ExAnteRealRate].rev_transform = quartertoannual

    pseudo[:LongRunInflation].name = "Long Run Inflation"
    pseudo[:LongRunInflation].longname = "Long Run Inflation"
    pseudo[:LongRunInflation].rev_transform = quartertoannual

    pseudo[:MarginalCost].name = "Marginal Cost"
    pseudo[:MarginalCost].longname = "Marginal Cost"

    pseudo[:Wages].name = "Wages"
    pseudo[:Wages].longname = "Wages"

    pseudo[:FlexibleWages].name = "Flexible Wages"
    pseudo[:FlexibleWages].longname = "Wages that would obtain in a flexible-wage economy"

    pseudo[:Hours].name = "Hours"
    pseudo[:Hours].longname = "Hours"

    pseudo[:FlexibleHours].name     = "Flexible Hours"
    pseudo[:FlexibleHours].longname = "Flexible Hours"

    pseudo[:z_t].name     = "z_t"
    pseudo[:z_t].longname = "z_t"

    pseudo[:Z_t].name     = "Z_t"
    pseudo[:Z_t].longname = "Z_t"

    pseudo[:z_t].name     = "ztil_t"
    pseudo[:z_t].longname = "ztil_t"

    pseudo[:NominalFFR].name     = "Nominal FFR"
    pseudo[:NominalFFR].longname = "Nominal FFR at an annual rate"
    pseudo[:NominalFFR].rev_transform = quartertoannual

    pseudo[:RealFFR].name     = "Real FFR"
    pseudo[:RealFFR].longname = "Real FFR at an annual rate"
    pseudo[:RealFFR].rev_transform = quartertoannual

    pseudo[:ExpectedNominalNaturalRate].name  = "Expected Nominal Natural Rate"
    pseudo[:ExpectedNominalNaturalRate].longname = "Natural Rate + Expected Inflation"
    pseudo[:ExpectedNominalNaturalRate].rev_transform = quartertoannual

    # Add to model object
    m.pseudo_observable_mappings = pseudo
end


using Base: promote_eltype

function _map(f, X)
    if isempty(X)
        return isa(f, Type) ? similar(X, f) : similar(X)
    end
    st = start(X)
    x1, st = next(X, st)
    first = f(x1)
    dest = similar(X, typeof(first))
    dest[1] = first
    return map_to!(f, 2, st, dest, X)
end

function _map(f, X1, X2)
    shp = promote_shape(size(X1), size(X2))
    if prod(shp) == 0
        return similar(X1, promote_type(eltype(X1), eltype(X2)), shp)
    end
    first = f(X1[1], X2[1])
    dest = similar(X1, typeof(first), shp)
    dest[1] = first
    return map_to!(f, 2, dest, X1, X2)
end

function _map(f, Xs...)
    shape = mapreduce(size, promote_shape, Xs)
    if prod(shape) == 0
        return similar(Xs[1], promote_eltype(Xs...), shape)
    end
    first = f(ith_all(1, Xs)...)
    dest = similar(Xs[1], typeof(first), shape)
    dest[1] = first
    return map_to_n!(f, 2, dest, Xs)
end

function map_to!{T,F}(f::F, offs, st, dest::NullableArray{T}, X)
    # map to dest array, checking the type of each result. if a result does not
    # match, widen the result type and re-dispatch.
    i = offs
    while !done(X, st)
        @inbounds Xi, st = next(X, st)
        el = f(Xi)
        S = typeof(el)
        if S === T || S <: T
            @inbounds dest[i] = el::T
            i += 1
        else
            R = typejoin(T, S)
            new = similar(dest, R)
            copy!(new,1, dest,1, i-1)
            @inbounds new[i] = el
            return map_to!(f, i+1, st, new, X)
        end
    end
    return dest
end

function map_to!{T,F}(f::F, offs, dest::NullableArray{T}, X1::NullableArray, X2::NullableArray)
    for i = offs:length(X1)
        @inbounds X1i, X2i = X1[i], X2[i]
        el = f(X1i, X2i)
        S = typeof(el)
        if (S !== T) && !(S <: T)
            R = typejoin(T, S)
            new = similar(dest, R)
            copy!(new,1, dest,1, i-1)
            @inbounds new[i] = el
            return map_to!(f, i+1, new, X1, X2)
        end
        @inbounds dest[i] = el::T
    end
    return dest
end

function map_to_n!{T,F}(f::F, offs, dest::NullableArray{T}, Xs)
    for i = offs:length(Xs[1])
        el = f(ith_all(i, Xs)...)
        S = typeof(el)
        if (S !== T) && !(S <: T)
            R = typejoin(T, S)
            new = similar(dest, R)
            copy!(new,1, dest,1, i-1)
            @inbounds new[i] = el
            return map_to_n!(f, i+1, new, Xs)
        end
        @inbounds dest[i] = el::T
    end
    return dest
end


include("./shared.jl")

## 3D tests
include("../scripts/diff_3D/diff_3D_nonlin_perf.jl")
@reset_parallel_stencil()
indsx = Int.(ceil.(LinRange(1, length(xc), 12)))
indsy = Int.(ceil.(LinRange(1, length(yc), 12)))
indsz = Int.(ceil.(LinRange(1, length(zc), 12)))
d3d   = Dict(:X=> xc[indsx], :H=>H[indsx,indsy,indsz])

@testset "Reference-tests diffusion 3D nonlin perf (noMPI)" begin
    @test_reference "reftest-files/test_diff_3D_nonlin_perf.bson" d3d by=comp
end

@reset_parallel_stencil()


# Specify model parameters and/or initial values to optimize
function get_search_index()::Tuple{Array{Int64,1},Array{Int64,1}}
    # parameters
    search_idx_params::Vector{Int} = [
        C.kassaikk
        C.kassaikknfkb
        C.kassanfkbikk
        C.kdisaikk
        C.kdisaikknfkb
        C.kdisanfkbikk
        C.kassa
        C.kdisa
        C.kdegbounda
        C.kdegboundaIKK
        C.kdegfreea
        C.kdegfreeaIKK
        C.kshutboundikbain
        C.kshutboundikbaout
        C.kshutfreeikbain
        C.kshutfreeikbaout
        C.k0mrnaikba
        C.kdegmrnaikba
        C.khillprodmrnaikba
        C.kpikba
        C.kprodmrnaikba

        C.kassbikk
        C.kassbikknfkb
        C.kassbnfkbikk
        C.kdisbikk
        C.kdisbikknfkb
        C.kdisbnfkbikk
        C.kassb
        C.kdegboundb
        C.kdegboundbIKK
        C.kdegfreeb
        C.kdegfreebIKK
        C.kdisb
        C.kshutboundikbbin
        C.kshutboundikbbout
        C.kshutfreeikbbin
        C.kshutfreeikbbout
        C.k0mrnaikbb
        C.kdegmrnaikbb
        # C.khillprodmrnaikbb
        C.kpikbb
        # C.kprodmrnaikbb
        C.kasseikk
        C.kasseikknfkb
        C.kassenfkbikk
        C.kdiseikk
        C.kdiseikknfkb
        C.kdisenfkbikk

        C.kasse
        C.kdegbounde
        C.kdegboundeIKK
        C.kdegfreee
        C.kdegfreeeIKK
        C.kdise
        C.kshutboundikbein
        C.kshutboundikbeout
        C.kshutfreeikbein
        C.kshutfreeikbeout
        C.k0mrnaikbe
        C.kdegmrnaikbe
        C.khillprodmrnaikbe
        C.kpikbe
        C.kprodmrnaikbe

        C.kshutnfkbin
        C.kshutnfkbout

        C.k0mrnaa20
        C.kdegmrnaa20
        C.kprodmrnaa20
        C.kpa20
        C.kdega20
        C.khillprodmrnaa20
    ]

    # initial values
    search_idx_initials::Vector{Int} = [
        # V.(variableName)
    ]

    return search_idx_params, search_idx_initials
end

function get_search_region()::Matrix{Float64}
    p::Vector{Float64} = param_values()
    u0::Vector{Float64} = initial_values()

    search_idx::Tuple{Array{Int64,1},Array{Int64,1}} = get_search_index()
    search_param::Vector{Float64} = initialize_search_param(search_idx, p, u0)

    search_rgn::Matrix{Float64} = zeros(2, length(p) + length(u0))

    # Default: 0.1 ~ 10x
    for (i, j) in enumerate(search_idx[1])
        search_rgn[1,j] = search_param[i] * 0.1  # lower bound
        search_rgn[2,j] = search_param[i] * 10.0  # upper bound
    end

    # Default: 0.5 ~ 2x
    for (i, j) in enumerate(search_idx[2])
        search_rgn[1,j + length(p)] = search_param[i + length(search_idx[1])] * 0.5  # lower bound
        search_rgn[2,j + length(p)] = search_param[i + length(search_idx[1])] * 2.0  # upper bound
    end

    # search_rgn[:, C.param_name] = [lower_bound, upper_bound]
    # search_rgn[:, V.var_name+length(p)] = [lower_bound, upper_bound]
    search_rgn[:, C.khillprodmrnaikba] = [1.0, 2.0]
    search_rgn[:, C.khillprodmrnaikbe] = [1.0, 2.0]
    search_rgn[:, C.khillprodmrnaa20] = [1.0, 2.0]

    search_rgn = convert_scale!(search_rgn, search_idx)

    return search_rgn
end

function update_param(indiv::Vector{Float64})::Tuple{Array{Float64,1},Array{Float64,1}}
    p::Vector{Float64} = param_values()
    u0::Vector{Float64} = initial_values()

    search_idx::Tuple{Array{Int64,1},Array{Int64,1}} = get_search_index()

    for (i, j) in enumerate(search_idx[1])
        @inbounds p[j] = indiv[i]
    end
    for (i, j) in enumerate(search_idx[2])
        @inbounds u0[j] = indiv[i + length(search_idx[1])]
    end

    return p, u0
end

function decode_gene2val(indiv_gene)::Vector{Float64}
    search_rgn::Matrix{Float64} = get_search_region()
    indiv::Vector{Float64} = zeros(length(indiv_gene))

    for i in eachindex(indiv_gene)
        indiv[i] = 10^(
            indiv_gene[i] * (
                search_rgn[2,i] - search_rgn[1,i]
            ) + search_rgn[1,i]
        )
    end

    return round.(indiv, sigdigits=7)
end

function encode_val2gene(indiv::Vector{Float64})
    search_rgn::Matrix{Float64} = get_search_region()
    indiv_gene::Vector{Float64} = zeros(length(indiv))

    for i in eachindex(indiv)
        indiv_gene[i] = (
            log10(indiv[i]) - search_rgn[1,i]
        ) / (
            search_rgn[2,i] - search_rgn[1,i]
        )
    end

    return indiv_gene
end

function encode_bestIndivVal2randGene(
        gene_idx::Int64,
        best_indiv::Vector{Float64},
        p0_bounds::Vector{Float64})::Float64
    search_rgn::Matrix{Float64} = get_search_region()
    rand_gene::Float64 = (
        log10(
            best_indiv[gene_idx] * 10^(
                rand() * log10(p0_bounds[2] / p0_bounds[1]) + log10(p0_bounds[1])
            )
        ) - search_rgn[1,gene_idx]
    ) / (
        search_rgn[2,gene_idx] - search_rgn[1,gene_idx]
    )
    return rand_gene
end

function initialize_search_param(
        search_idx::Tuple{Array{Int64,1},Array{Int64,1}},
        p::Vector{Float64},
        u0::Vector{Float64})::Vector{Float64}
    duplicate::Vector{String} = []
    if length(search_idx[1]) != length(unique(search_idx[1]))
        for idx in findall(
            [count(x -> x == i, search_idx[1]) for i in unique(search_idx[1])] .!= 1
        )
            push!(duplicate, C.NAMES[search_idx[1][idx]])
        end
        error(
            "Duplicate parameters (C.): $duplicate"
        )
    elseif length(search_idx[2]) != length(unique(search_idx[2]))
        for idx in findall(
            [count(x -> x == i, search_idx[2]) for i in unique(search_idx[2])] .!= 1
        )
            push!(duplicate, V.NAMES[search_idx[2][idx]])
        end
        error(
            "Duplicate initial conditions (V.): $duplicate"
        )
    end
    search_param = zeros(
        length(search_idx[1]) + length(search_idx[2])
    )
    for (i, j) in enumerate(search_idx[1])
        @inbounds search_param[i] = p[j]
    end
    for (i, j) in enumerate(search_idx[2])
        @inbounds search_param[i + length(search_idx[1])] = u0[j]
    end

    if any(x -> x == 0.0, search_param)
        msg::String = "search_param must not contain zero."
        for idx in search_idx[1]
            if p[idx] == 0.0
                error(
                    @sprintf(
                        "`C.%s` in search_idx_params: ", C.NAMES[idx]
                    ) * msg
                )
            end
        end
        for idx in search_idx[2]
            if u0[idx] == 0.0
                error(
                    @sprintf(
                        "`V.%s` in search_idx_initials: ", V.NAMES[idx]
                    ) * msg
                )
            end
        end
    end

    return search_param
end

function convert_scale!(
        search_rgn::Matrix{Float64},
        search_idx::Tuple{Array{Int64,1},Array{Int64,1}})::Matrix{Float64}
    for i = 1:size(search_rgn, 2)
        if minimum(search_rgn[:,i]) < 0.0
            msg = "search_rgn[lower_bound,upper_bound] must be positive.\n"
            if i <= C.NUM
                error(@sprintf("`C.%s` ", C.NAMES[i]) * msg)
            else
                error(@sprintf("`V.%s` ", V.NAMES[i - C.NUM]) * msg)
            end
        elseif minimum(search_rgn[:,i]) == 0.0 && maximum(search_rgn[:,i]) != 0.0
            msg = "lower_bound must be larger than 0.\n"
            if i <= C.NUM
                error(@sprintf("`C.%s` ", C.NAMES[i]) * msg)
            else
                error(@sprintf("`V.%s` ", V.NAMES[i - C.NUM]) * msg)
            end
        elseif search_rgn[2,i] - search_rgn[1,i] < 0.0
            msg = "lower_bound must be smaller than upper_bound.\n"
            if i <= C.NUM
                error(@sprintf("`C.%s` ", C.NAMES[i]) * msg)
            else
                error(@sprintf("`V.%s` ", V.NAMES[i - C.NUM]) * msg)
            end
        end
    end

    nonzero_idx::Vector{Int} = []
    for i = 1:size(search_rgn, 2)
        if search_rgn[:,i] != [0.0,0.0]
            push!(nonzero_idx, i)
        end
    end
    difference::Vector{Int} = collect(
        symdiff(
            Set(nonzero_idx),
            Set(append!(search_idx[1], C.NUM .+ search_idx[2]))
        )
    )
    if length(difference) > 0
        for idx in difference
            if j <= C.NUM
                println(@sprintf("`C.%s`", C.NAMES[Int(idx)]))
            else
                println(@sprintf("`V.%s`", V.NAMES[Int(idx) - C.NUM]))
            end
        end
        error(
            "Set these search_params in both search_idx and search_rgn."
        )
    end

    search_rgn = search_rgn[:,nonzero_idx]

    return log10.(search_rgn)
end


using Test

@testset "masked_relu.jl" begin

    @testset "MaskedReLU" begin
        mask = [-1, 1, 0, 0, 1, 1]
        @testset "Initialize without tightening algorithm" begin
            p = MaskedReLU(mask)
            @test p.tightening_algorithm === nothing
        end

        @testset "Initialize with tightening algorithm" begin
            p = MaskedReLU(mask, MIPVerify.mip)
            @test !(p.tightening_algorithm === nothing)
            @test p.tightening_algorithm == MIPVerify.mip
        end

        @testset "Base.show" begin
            p = MaskedReLU(mask)
            io = IOBuffer()
            Base.show(io, p)
            @test String(take!(io)) ==
                  "MaskedReLU with expected input size (6,). (1 zeroed, 3 as-is, 2 rectified)."
        end
    end

end


mutable struct SimulateContext{T <: Trace, P <: Parameters} <: ExecutionContext
    tr::T
    score::Float64
    visited::Visitor
    params::P
    SimulateContext() = new{HierarchicalTrace, EmptyParameters}(Trace(), 0.0, Visitor(), Parameters())
    SimulateContext(params::P) where P = new{HierarchicalTrace, P}(Trace(), 0.0, Visitor(), params)
end
Simulate() = SimulateContext()

# ------------ Convenience ------------ #

function simulate(fn::Function, args...)
    ctx = SimulateContext()
    ret = ctx(fn, args...)
    return ret, HierarchicalCallSite(ctx.tr, ctx.score, fn, args, ret)
end

function simulate(params, fn::Function, args...)
    ctx = SimulateContext(params)
    ret = ctx(fn, args...)
    return ret, HierarchicalCallSite(ctx.tr, ctx.score, fn, args, ret)
end

function simulate(fn::typeof(rand), d::Distribution{T}) where T
    ctx = SimulateContext()
    addr = gensym()
    ret = ctx(rand, addr, d)
    return ret, get_top(ctx.tr, addr)
end

function simulate(params, fn::typeof(rand), d::Distribution{T}) where T
    ctx = SimulateContext(params)
    addr = gensym()
    ret = ctx(rand, addr, d)
    return ret, get_top(ctx.tr, addr)
end

function simulate(c::typeof(plate), fn::Function, args::Vector)
    ctx = SimulateContext()
    addr = gensym()
    ret = ctx(plate, addr, fn, args)
    return ret, get_sub(ctx.tr, addr)
end

function simulate(params, c::typeof(plate), fn::Function, args::Vector)
    addr = gensym()
    v_ps = parameters(addr => params)
    ctx = SimulateContext(v_ps)
    ret = ctx(plate, addr, fn, args)
    return ret, get_sub(ctx.tr, addr)
end

function simulate(fn::typeof(plate), d::Distribution{T}, len::Int) where T
    ctx = SimulateContext()
    addr = gensym()
    ret = ctx(plate, addr, d, len)
    return ret, get_sub(ctx.tr, addr)
end

function simulate(c::typeof(markov), fn::Function, len::Int, args...)
    ctx = SimulateContext()
    addr = gensym()
    ret = ctx(markov, addr, fn, len, args...)
    return ret, get_sub(ctx.tr, addr)
end

function simulate(params, c::typeof(markov), fn::Function, len::Int, args...)
    addr = gensym()
    v_ps = parameters(addr => params)
    ctx = SimulateContext(v_ps)
    ret = ctx(markov, addr, fn, len, args...)
    return ret, get_sub(ctx.tr, addr)
end

# ------------ includes ------------ #

include("hierarchical/simulate.jl")
include("plate/simulate.jl")
include("markov/simulate.jl")
include("conditional/simulate.jl")
include("factor/simulate.jl")

# ------------ Documentation ------------ #

@doc(
"""
```julia
mutable struct SimulateContext{T <: Trace} <: ExecutionContext
    tr::T
    visited::Visitor
    params::LearnableParameters
    SimulateContext(params) where T <: Trace = new{T}(Trace(), Visitor(), params)
end
```

`SimulateContext` is used to simulate traces without recording likelihood weights. `SimulateContext` can be instantiated with custom `LearnableParameters` instances, which is useful when used for gradient-based learning.

Inner constructors:
```julia
SimulateContext(params) = new{HierarchicalTrace}(Trace(), Visitor(), params)
```
""", SimulateContext)

@doc(
"""
```julia
ret, cl = simulate(fn::Function, args...; params = LearnableParameters())
ret, cl = simulate(fn::typeof(rand), d::Distribution{T}; params = LearnableParameters()) where T
ret, v_cl = simulate(c::typeof(plate), fn::Function, args::Vector; params = LearnableParameters()) where T
ret, v_cl = simulate(fn::typeof(plate), d::Distribution{T}, len::Int; params = LearnableParameters()) where T
ret, v_cl = simulate(c::typeof(markov), fn::Function, len::Int, args...; params = LearnableParameters())
```

`simulate` function provides an API to the `SimulateContext` execution context. You can use this function on any of the matching signatures above - it will return the return value `ret`, and a `RecordSite` instance specialized to the call. `simulate` is used to express unconstrained generation of a probabilistic program trace, without likelihood weight recording.
""", simulate)


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


mutable struct SimulationParameters
  endtime::Float64
  filenamestub::String
  time::Float64
  iteration::Int
end

SimulationParameters(endtime, filenamestub) = SimulationParameters(endtime, filenamestub, 0.0, 0)

@concrete struct Simulation{F<:AbstractField}
  plasma::Plasma
  field::F
  timeintegrator
  diagnosticdumpevery::Int64
  endtime::Float64
  filenamestub::String
  time::Ref{Float64}
  iteration::Ref{Int}
  diagnosticdumpcounter::Ref{Int}
end

function Simulation(plasma::Plasma, field::AbstractField, timeintegrator::AbstractTimeIntegrator;
  diagnosticdumpevery=1, endtime=10.0, filenamestub=nothing)
  isnothing(filenamestub) && (filenamestub = "$(typeof(eltype(plasma[1])))-$(typeof(field))-$(timeintegrator)-")
  return Simulation(plasma, field, timeintegrator, diagnosticdumpevery, endtime, filenamestub, Ref(0.0), Ref(0), Ref(-1))
end
filenamestub(sim::Simulation) = sim.filenamestub
filename(filepathstub::String, i::Int) = filepathstub * lpad("$i", 6, "0") * ".jld2"
filename(sim::Simulation) = filename(sim.filenamestub, diagnosticdumpcounter(sim))
timestep(sim::Simulation) = timestep(sim.timeintegrator)
iteration(sim::Simulation) = sim.iteration[]
diagnosticdumpcounter(sim::Simulation) = sim.diagnosticdumpcounter[]
time(sim::Simulation) = sim.time[]
iterate!(sim::Simulation) = (sim.iteration[] += 1)
function save(sim::Simulation)
  sim.diagnosticdumpcounter[] += 1
  @save "$(filename(sim))" sim
end
function load(filepathstub::String, diagnosticdumpcount::Int)
  @load "$(filename(filepathstub, diagnosticdumpcount))" sim
  return sim
end

function init!(sim::Simulation)
  sim.diagnosticdumpcounter[] = -1
  zero!(sim.field)
  deposit!(sim.field, sim.plasma)
  solve!(sim.field)
  return sim
end
function run!(sim::Simulation)
  save(sim) # save initial conditions as file 0
  while time(sim) <= sim.endtime
    runtimestep!(sim)
  end
end

function runtimestep!(sim::Simulation)
  integrate!(sim.plasma, sim.field, sim.timeintegrator)
  sim.iteration[] += 1
  sim.time[] += timestep(sim.timeintegrator)
  if mod(iteration(sim), sim.diagnosticdumpevery) == 0
    save(sim)
  end
  return sim
end


function CalTrack(Depth,cal,bit;DepthFrom=minimum(Depth),DepthTo=maximum(Depth),Dtick=false)

    calax=gca()
    calax.plot(cal,Depth, color="red",linewidth=1, linestyle="--")
    calax.plot(bit,Depth, color="black",linewidth=0.5)
    calax.set_ylim([DepthTo,DepthFrom])
    calax.set_xlim([5,17])
    calax.set_xlabel("Diameter Inches")
    calax.set_yticks(range(DepthFrom, stop=DepthTo, length=51),minor=true)
    calax.set_yticks(range(DepthFrom, stop=DepthTo, length=11))
    calax.xaxis.tick_top()
    calax.xaxis.set_label_position("top")
    calax.tick_params("both",labelsize=8)
    setp(calax.get_yticklabels(),visible=Dtick)
    calax.fill_betweenx(Depth,cal,bit,where=(cal .< bit),color="yellow")
    calax.fill_betweenx(Depth,cal,bit,where=(cal .> bit),color=(0.2,0.4,0.8))
end


using Dates, .CentralBanks

CentralBanks.ECB_FxRates("//datos02/9763-AnalisisyPlanificacion_Financiera/D. Datos/7. Tipos", Dates.Date(2005))

CentralBanks.downloadchapter(4,"c:\\Users\\Luis\\OneDrive\\Documents\\JuliaProjects\\CentralBanks")
CentralBanks.downloadchapter(2,"c:\\Users\\Luis\\OneDrive\\Documents\\JuliaProjects\\CentralBanks")
CentralBanks.downloadchapter(19,"c:\\Users\\Luis\\OneDrive\\Documents\\JuliaProjects\\CentralBanks")

CentralBanks.downloadInterestRates("c:\\Users\\Luis\\OneDrive\\Documents\\JuliaProjects\\CentralBanks")
dfs = CentralBanks.CSVtoDataFrame("c:\\Users\\Luis\\OneDrive\\Documents\\JuliaProjects\\CentralBanks\\tipos",Date(2005))
sort!(dfs["tInteres"], 1,rev = true) # Va a fallar porque no está el using DataFrame
CentralBanks.DataFrametoCSV(dfs["tInteres"], "c:\\Users\\Luis\\OneDrive\\Documents\\JuliaProjects\\CentralBanks\\tipos","tinteres.csv")

CentralBanks.combinaciones(ans)


module Yeppp

# @unix_only    const libyeppp = "libyeppp"
# @windows_only const libyeppp = "yeppp"

using Libdl # BinDeps injects a dependency on Libdl

if isfile(joinpath(dirname(@__FILE__),"..","deps","deps.jl"))
    include("../deps/deps.jl")
else
    error("Yeppp not properly installed. Please run Pkg.build(\"Yeppp\")")
end

function __init__()
    status = ccall( (:yepLibrary_Init, libyeppp), Cint, ())
    status != 0 && error("yepLibrary_Init: error: ", status)
end

function release()
    status = ccall( (:yepLibrary_Release, libyeppp), Cint, ())
    status != 0 && error("yepLibrary_Release: error: ", status)
end

_iscontiguous(A::DenseArray) = true
# Fallback
_iscontiguous(A::StridedArray) = false
# Cases we can verify
_iscontiguous(A::SubArray) = Base.iscontiguous(A)
_iscontiguous(A::StridedVector) = stride(A,1) == 1

macro yepppfunsAA_A(fname, libname, BT)
    errorname = libname * ": error: "
    quote
        global $(fname)
        function $(fname)(res::Union{DenseArray{$(BT)},StridedArray{$(BT)}}, x::Union{DenseArray{$(BT)},StridedArray{$(BT)}}, y::Union{DenseArray{$(BT)},StridedArray{$(BT)}})
            n = length(x)
            @assert(n == length(y) == length(res))
            @assert !Base.has_offset_axes(res, x, y)
            @assert _iscontiguous(res) && _iscontiguous(x) && _iscontiguous(y) "could not verify that arguments are contigous is memory"

            status = ccall( ($(libname), libyeppp), Cint, (Ptr{$(BT)}, Ptr{$(BT)}, Ptr{$(BT)}, Culong), x, y, res, n)
            status != 0 && error($(errorname), status)
            res
        end
    end
end

macro yepppfunsA_A(fname, libname, BT)
    errorname = libname * ": error: "
    quote
        global $(fname)
        function $(fname)(res::Union{DenseArray{$(BT)},StridedArray{$(BT)}}, x::Union{DenseArray{$(BT)},StridedArray{$(BT)}})
            n = length(x)
            @assert(n == length(res))
            @assert !Base.has_offset_axes(res, x)
            @assert _iscontiguous(res) && _iscontiguous(x) "could not verify that arguments are contigous is memory"

            status = ccall( ($(libname), libyeppp), Cint, (Ptr{$(BT)}, Ptr{$(BT)}, Culong), x, res, n)
            status != 0 && error($(errorname), status)
            res
        end
    end
end

macro yepppfunsA_S(fname, libname, BT)
    errorname = libname * ": error: "
    quote
        global $(fname)
        function $(fname)(x::Union{DenseArray{$(BT)},StridedArray{$(BT)}})
            n = length(x)
            @assert !Base.has_offset_axes(x)
            @assert _iscontiguous(x) "could not verify that arguments are contigous is memory"
            res = Ref{$BT}()
            status = ccall( ($(libname), libyeppp), Cint, (Ptr{$(BT)}, Ptr{$(BT)}, Culong), x, res, n)
            status != 0 && error($(errorname), status)
            res[]
        end
    end
end

## == Float64 == ##
@yepppfunsAA_A add! "yepCore_Add_V64fV64f_V64f" Float64
@yepppfunsAA_A multiply! "yepCore_Multiply_V64fV64f_V64f" Float64
@yepppfunsAA_A subtract! "yepCore_Subtract_V64fV64f_V64f" Float64
@yepppfunsAA_A max! "yepCore_Max_V64fV64f_V64f" Float64
@yepppfunsAA_A min! "yepCore_Min_V64fV64f_V64f" Float64

@yepppfunsA_A log! "yepMath_Log_V64f_V64f" Float64
@yepppfunsA_A exp! "yepMath_Exp_V64f_V64f" Float64
@yepppfunsA_A sin! "yepMath_Sin_V64f_V64f" Float64
@yepppfunsA_A cos! "yepMath_Cos_V64f_V64f" Float64
@yepppfunsA_A tan! "yepMath_Tan_V64f_V64f" Float64

@yepppfunsA_S sum "yepCore_Sum_V64f_S64f" Float64
@yepppfunsA_S sumabs "yepCore_SumAbs_V64f_S64f" Float64
@yepppfunsA_S sumabs2 "yepCore_SumSquares_V64f_S64f" Float64

## == Float32 == ##
@yepppfunsAA_A add! "yepCore_Add_V32fV32f_V32f" Float32
@yepppfunsAA_A multiply! "yepCore_Multiply_V32fV32f_V32f" Float32
@yepppfunsAA_A subtract! "yepCore_Subtract_V32fV32f_V32f" Float32
@yepppfunsAA_A max! "yepCore_Max_V32fV32f_V32f" Float32
@yepppfunsAA_A min! "yepCore_Min_V32fV32f_V32f" Float32

@yepppfunsA_S sum "yepCore_Sum_V32f_S32f" Float32
@yepppfunsA_S sumabs "yepCore_SumAbs_V32f_S32f" Float32
@yepppfunsA_S sumabs2 "yepCore_SumSquares_V32f_S32f" Float32

"""
    dot(x::Union{DenseVector{Float64},StridedVector{Float64}}, y::Union{DenseVector{Float64},StridedVector{Float64}})

Compute the dot product of x and y.
"""
function dot(x::Union{DenseVector{Float64},StridedVector{Float64}}, y::Union{DenseVector{Float64},StridedVector{Float64}})
    n = length(x)
    @assert(n == length(y))
    @assert !Base.has_offset_axes(x, y)
    @assert _iscontiguous(x) && _iscontiguous(y)

    dotproduct = Ref{Float64}()
    status = ccall( (:yepCore_DotProduct_V64fV64f_S64f, libyeppp), Cint, (Ptr{Float64}, Ptr{Float64}, Ptr{Float64}, Culong), x, y, dotproduct, n)
    status != 0 && error("yepCore_DotProduct_V64fV64f_S64f: error: ", status)
    dotproduct[]
end

"""
    dot(x::Union{DenseVector{Float32},StridedVector{Float32}}, y::Union{DenseVector{Float32},StridedVector{Float32}})

Compute the dot product of x and y.
"""
function dot(x::Union{DenseVector{Float32},StridedVector{Float32}}, y::Union{DenseVector{Float32},StridedVector{Float32}})
    n = length(x)
    @assert(n == length(y))
    @assert !Base.has_offset_axes(x, y)
    @assert _iscontiguous(x) && _iscontiguous(y)

    dotproduct = Ref{Float32}()
    status = ccall( (:yepCore_DotProduct_V32fV32f_S32f, libyeppp), Cint, (Ptr{Float32}, Ptr{Float32}, Ptr{Float32}, Culong), x, y, dotproduct, n)
    status != 0 && error("yepCore_DotProduct_V32fV32f_S32f: error: ", status)
    dotproduct[]
end

"""
    max(x, y)

Compares the vectors `x` and `y` and return the element wise maxima.
"""
max(x, y) = max!(similar(x), x, y)

"""
    max(x, y)

Compares the vectors `x` and `y` and return the element wise minima.
"""
min(x, y) = min!(similar(x), x, y)

function evalpoly!(res::Union{DenseArray{Float64},StridedVector{Float64}}, coef::Union{DenseArray{Float64},StridedVector{Float64}}, x::Union{DenseArray{Float64},StridedVector{Float64}})
    n = length(coef)
    arraysize = length(x)
    @assert(arraysize == length(res))
    @assert !Base.has_offset_axes(res, coef, x)
    @assert _iscontiguous(res) && _iscontiguous(coef) && _iscontiguous(x)

    status = ccall( (:yepMath_EvaluatePolynomial_V64fV64f_V64f, libyeppp), Cint, (Ptr{Float64}, Ptr{Float64}, Ptr{Float64}, Culong, Culong), coef, x, res, n, arraysize)
    status != 0 && error("yepMath_EvaluatePolynomial_V64fV64f_V64f: error: ", status)
    res
end

"""
    evalpoly(coeff, x)

Evaluates polynomial with double precision (64-bit) floating-point coefficients `coeff` on an array of double precision (64-bit) floating-point elements `x`.
"""
evalpoly(coef, x) = evalpoly!(similar(x), coef, x)

"""
    add(x, y)

Perform element wise addition of the two array `x` and `y`.
"""
add(x, y) = add!(similar(x), x, y)

"""
    subtract(x, y)

Perform element wise subtraction of the two array `x` and `y`.
"""
subtract(x, y) = subtract!(similar(x), x, y)

"""
    multiply(x, y)

Perform element wise multiplication of the two array `x` and `y`.
"""
multiply(x, y) = multiply!(similar(x), x, y)

"""
    log(x)

Returns the element wise natural logarithm of `x`.
"""
log(x) = log!(similar(x), x)

"""
    log!(x)

Computes the element wise natural logarithm of `x` inplace.
"""
log!(x) = log!(x, x)

"""
    exp(x)

Returns the element wise exponential of `x`.
"""
exp(x) = exp!(similar(x), x)

"""
    exp(x)

Computes the element wise exponential of `x` inplace.
"""
exp!(x) = exp!(x, x)

"""
    sin(x)

Returns element wise `sin` of `x`.
"""
sin(x) = sin!(similar(x), x)

"""
    sin(x)

Computes element wise `sin` of `x` inplace.
"""
sin!(x) = sin!(x, x)

"""
    cos(x)

Returns element wise `cos` of `x`.
"""
cos(x) = cos!(similar(x), x)

"""
    cos(x)

Returns element wise `cos` of `x`.
"""
cos!(x) = cos!(x, x)

"""
    tan(x)

Returns element wise `tan` of `x`.
"""
tan(x) = tan!(similar(x), x)

"""
    tan(x)

Returns element wise `tan` of `x`.
"""
tan!(x) = tan!(x, x)
end # module


using DSGE, ModelConstructors, Dates, Optim, Test, Nullables, HDF5
path = dirname(@__FILE__)
m = AnSchorfheide()
m <= Setting(:date_zlb_end, get_setting(m, :date_zlb_start))
m <= Setting(:hessian_path, "")

# Dates, indices, number of periods for each regime
@testset "Test field access functions for AbstractDSGEModel objects" begin
    @test DSGE.get_parameters(m) == m.parameters
    @test DSGE.get_rng(m) == m.rng
    @test DSGE.get_settings(m) == m.settings
    @test DSGE.get_observables(m) == m.observables
    @test DSGE.get_pseudo_observables(m) == m.pseudo_observables
    @test DSGE.get_exogenous_shocks(m) == m.exogenous_shocks
end

@testset "Test setting access functions for AbstractDSGEModel objects" begin
    n_anticipated_shocks(m) == get_setting(m, :n_anticipated_shocks)
    n_anticipated_shocks_padding(m) == get_setting(m, :n_anticipated_shocks_padding)

    @test DSGE.date_presample_start(m) == get_setting(m, :date_presample_start)
    @test DSGE.date_mainsample_start(m) == get_setting(m, :date_mainsample_start)
    @test DSGE.date_zlb_start(m) == get_setting(m, :date_zlb_start)
    @test DSGE.date_zlb_end(m) == get_setting(m, :date_zlb_end)

    @test DSGE.date_presample_end(m) == Dates.lastdayofquarter(get_setting(m, :date_mainsample_start) - Dates.Month(3))
    @test DSGE.date_prezlb_end(m) == Dates.lastdayofquarter(get_setting(m, :date_zlb_start) - Dates.Month(3))
    @test DSGE.date_mainsample_end(m) == Dates.lastdayofquarter(get_setting(m, :date_forecast_start) - Dates.Month(3))
    @test DSGE.date_conditional_end(m) == get_setting(m, :date_conditional_end)

    @test DSGE.index_presample_start(m) == 1
    @test DSGE.index_mainsample_start(m) == subtract_quarters(DSGE.date_mainsample_start(m), DSGE.date_presample_start(m)) + 1
    @test DSGE.index_zlb_start(m) == subtract_quarters(DSGE.date_zlb_start(m), DSGE.date_presample_start(m)) + 1
    @test DSGE.index_forecast_start(m) == subtract_quarters(DSGE.date_forecast_start(m), DSGE.date_presample_start(m)) + 1

    @test DSGE.index_shockdec_start(m) == subtract_quarters(DSGE.date_shockdec_start(m), DSGE.date_mainsample_start(m)) + 1

    @test DSGE.index_shockdec_end(m) == subtract_quarters(DSGE.date_shockdec_end(m), DSGE.date_mainsample_start(m)) + 1

    @test n_presample_periods(m)   == subtract_quarters(DSGE.date_mainsample_start(m), DSGE.date_presample_start(m))
    @test n_prezlb_periods(m)      == subtract_quarters(DSGE.date_zlb_start(m), DSGE.date_mainsample_start(m))
    @test n_zlb_periods(m)         == subtract_quarters(DSGE.date_forecast_start(m), DSGE.date_zlb_start(m))
    @test n_mainsample_periods(m)  == subtract_quarters(DSGE.date_forecast_start(m), DSGE.date_mainsample_start(m))
    @test n_conditional_periods(m) == subtract_quarters(DSGE.date_conditional_end(m), DSGE.date_mainsample_end(m))

    @test inds_presample_periods(m)  == collect(DSGE.index_presample_start(m):(DSGE.index_mainsample_start(m)-1))
    @test inds_prezlb_periods(m)     == collect(DSGE.index_mainsample_start(m):(DSGE.index_zlb_start(m)-1))
    @test inds_zlb_periods(m)        == collect(DSGE.index_zlb_start(m):(DSGE.index_forecast_start(m)-1))
    @test inds_mainsample_periods(m) == collect(DSGE.index_mainsample_start(m):(DSGE.index_forecast_start(m)-1))

    @test cond_vintage(m)    == get_setting(m, :cond_vintage)
    @test cond_id(m)         == get_setting(m, :cond_id)
    @test cond_full_names(m) == get_setting(m, :cond_full_names)
    @test cond_semi_names(m) == get_setting(m, :cond_semi_names)
    @test use_population_forecast(m) == get_setting(m, :use_population_forecast)
    @test DSGE.hpfilter_population(m)     == get_setting(m, :hpfilter_population)

    # Interface for general computation settings
    @test DSGE.use_parallel_workers(m)    == get_setting(m, :use_parallel_workers)

    # Interface for estimation settings
    @test DSGE.reoptimize(m)          == get_setting(m, :reoptimize)
    @test DSGE.calculate_hessian(m) == get_setting(m, :calculate_hessian)
    @test DSGE.hessian_path(m)      == get_setting(m, :hessian_path)
    @test DSGE.n_hessian_test_params(m) == get_setting(m, :n_hessian_test_params)

    # Interface for Metropolis-Hastings settings
    @test DSGE.n_mh_blocks(m)       ==  get_setting(m, :n_mh_blocks)
    @test DSGE.n_mh_param_blocks(m) ==  get_setting(m, :n_mh_param_blocks)
    @test DSGE.n_mh_simulations(m)  ==  get_setting(m, :n_mh_simulations)
    @test DSGE.n_mh_burn(m)         ==  get_setting(m, :n_mh_burn)
    @test DSGE.mh_thin(m)           ==  get_setting(m, :mh_thin)

    # Interface for forecast settings
    @test DSGE.date_forecast_start(m)   == get_setting(m, :date_forecast_start)
    @test DSGE.forecast_block_size(m)   == get_setting(m, :forecast_block_size)
    @test DSGE.forecast_start_block(m)  == get_setting(m, :forecast_start_block)
    @test DSGE.forecast_input_file_overrides(m) == get_setting(m, :forecast_input_file_overrides)
    @test isnull(DSGE.forecast_uncertainty_override(m)) && isnull(get_setting(m, :forecast_uncertainty_override))
    @test DSGE.forecast_smoother(m)     == get_setting(m, :forecast_smoother)
    @test DSGE.forecast_tdist_df_val(m) == get_setting(m, :forecast_tdist_df_val)
    @test DSGE.forecast_tdist_shocks(m) == get_setting(m, :forecast_tdist_shocks)
    @test DSGE.forecast_zlb_value(m)    == get_setting(m, :forecast_zlb_value)
    @test DSGE.impulse_response_horizons(m) == get_setting(m, :impulse_response_horizons)
    @test DSGE.n_shockdec_periods(m)    == DSGE.index_shockdec_end(m) - DSGE.index_shockdec_start(m) + 1

    # Interface for alternative policy settings
    @test DSGE.alternative_policy(m) == get_setting(m, :alternative_policy)
end

@testset "Test other auxiliary setting functions for AbstractDSGEModel objects" begin
    m <= Setting(:population_mnemonic, :test)
    @test get(DSGE.parse_population_mnemonic(m)[1]) == :test
    m <= Setting(:population_mnemonic, Nullable())
    println(DSGE.parse_population_mnemonic(m))
    @test DSGE.parse_population_mnemonic(m)[1].hasvalue == false
    m1 = Model1002("ss10")
    @test @test_matrix_approx_eq inds_states_no_ant(m) collect(1:8)
    @test @test_matrix_approx_eq inds_states_no_ant(m1) vcat(collect(1:62), collect(69:84))
    @test @test_matrix_approx_eq inds_shocks_no_ant(m) collect(1:3)
    @test @test_matrix_approx_eq inds_shocks_no_ant(m1) collect(1:18)
    @test @test_matrix_approx_eq inds_obs_no_ant(m) collect(1:3)
    @test @test_matrix_approx_eq inds_obs_no_ant(m1) collect(1:13)

    m <= Setting(:date_forecast_start, DSGE.quartertodate("2019-Q3"))
    m <= Setting(:forecast_horizons, 10)
    exp_date_forecast_end = DSGE.lastdayofquarter(DSGE.quartertodate("2019-Q3") + Dates.Month(3 * 9))
    @test forecast_horizons(m) == 10
    m <= Setting(:date_conditional_end, DSGE.date_mainsample_end(m) + Dates.Month(3))
    @test forecast_horizons(m; cond_type = :full) == 9
    @test forecast_horizons(m; cond_type = :semi) == 9
    @test DSGE.date_forecast_end(m) == exp_date_forecast_end
    @test DSGE.date_shockdec_start(m) == date_mainsample_start(m)
    m <= Setting(:shockdec_startdate, nothing)
    @test DSGE.date_shockdec_start(m) == DSGE.date_mainsample_start(m)
    m <= Setting(:shockdec_startdate, Nullable{Date}(DSGE.quartertodate("2019-Q3")))
    @test DSGE.date_shockdec_start(m) == DSGE.quartertodate("2019-Q3")
    m <= Setting(:shockdec_enddate, Nullable{Date}())
    @test DSGE.date_shockdec_end(m) == DSGE.date_forecast_end(m)
    m <= Setting(:shockdec_enddate, Nullable{Date}(DSGE.quartertodate("2019-Q3")))
    @test DSGE.date_shockdec_end(m) == DSGE.quartertodate("2019-Q3")

    @test @test_matrix_approx_eq load_parameters_from_file(m, "$path/reference/load_params.h5") h5read("$path/reference/load_params.h5", "params")
    @test_throws ErrorException load_parameters_from_file(m, "$path/reference/smc.h5")
    @test_throws ErrorException load_parameters_from_file(m, "$path/reference/smc.jld2")
    @test_throws AssertionError load_parameters_from_file(m1, "$path/reference/load_params.h5")
    @test_throws AssertionError load_parameters_from_file(m, "$path/reference/load_params_float32.h5")

    m.parameters[1].value = 1.
    specify_mode!(m, "reference/load_params.h5")
    out_params = map(x -> x.value, m.parameters)
    @test @test_matrix_approx_eq out_params load_parameters_from_file(m, "$path/reference/load_params.h5")

    specify_hessian!(m, "reference/hessian.h5")
    @test get_setting(m, :calculate_hessian) == false
    @test get_setting(m, :hessian_path) == joinpath(dirname("$path"), "test", "reference", "hessian.h5")

    para = deepcopy(map(x -> x.value, m.parameters))
    m.parameters[1].value = 1.
    DSGE.update!(m, para)
    @test m.parameters[1].value != 1.
    @test @test_matrix_approx_eq map(x->x.value,m.parameters) para
end

nothing


#=
Base.@kwdef mutable struct SIPOptimizer <: MOI.AbstractOptimizer
    objective_sense::MOI.OptimizationSense = MOI.FEASIBLE_SENSE
    time_limit::Float64 = 1000.0
    is_silent::Bool = false
    x::Vector{Float64} = Float64[]
    obj_bound::Float64 = 0.0
    obj_value::Float64 = 0.0
    termination_status::MOI.TerminationStatusCode = MOI.OPTIMIZE_NOT_CALLED
    result_status_code::MOI.ResultStatusCode = MOI.UNKNOWN_RESULT_STATUS
    run_time::Float64 = 0.0
end

function MOI.empty!(m::SIPOptimizer)
end
function MOI.is_empty(m::Optimizer)
end

function MOI.set(m::SIPOptimizer, ::MOI.Silent, value)
    m.is_silent = value
    return nothing
end

function MOI.set(m::SIPOptimizer, ::MOI.TimeLimitSec, value::Nothing)
    m.time_limit = Inf
    return nothing
end
function MOI.set(m::SIPOptimizer, ::MOI.TimeLimitSec, value::Float64)
    m.time_limit = value
    return nothing
end

MOI.set(m::SIPOptimizer, p::MOI.RawParameter, value) = error("The SIPOptimizer does not have any raw parameters.")

function MOI.get(m::SIPOptimizer, ::MOI.ListOfVariableIndices)
end
function MOI.get(m::SIPOptimizer, ::MOI.NumberOfVariables)
end

function MOI.get(m::SIPOptimizer, ::MOI.ObjectiveValue)

end
function MOI.get(m::SIPOptimizer, ::MOI.ObjectiveBound)

end

MOI.get(m::SIPOptimizer, p::MOI.RawParameter) = error("The SIPOptimizer does not have any raw parameters.")

function MOI.get(m::SIPOptimizer, ::MOI.RelativeGap)
    abs(objective_bound() - objective_value())/objective_value()
end
MOI.get(m::SIPOptimizer, ::MOI.SolverName) = "EAGO: SIPOptimizer"
MOI.get(m::SIPOptimizer, ::MOI.TerminationStatus) = m._termination_status_code
MOI.get(m::SIPOptimizer, ::MOI.PrimalStatus) = m.result_status_code
MOI.get(m::SIPOptimizer, ::MOI.SolveTime) = m.run_time
MOI.get(m::SIPOptimizer, ::MOI.ResultCount) = (m.result_status_code === MOI.FEASIBLE_POINT) ? 1 : 0
MOI.get(m::SIPOptimizer, ::MOI.TimeLimitSec) = m.time_limit

MOI.supports(::SIPOptimizer, ::MOI.TimeLimitSec) = true
MOI.supports(::SIPOptimizer, ::MOI.ObjectiveSense) = true
function MOI.set(m::SIPOptimizer, ::MOI.ObjectiveSense, sense::MOI.OptimizationSense)
    m.objective_sense = sense
    return nothing
end

function MOI.optimize!(m::SIPOptimizer)
end
=#


import Celeste: BoundingBox
import Celeste.DECALSIO: DECALSDataSet, _get_overlapping_ccds

@testset "decalsio" begin

    # ensure files downloaded
    DATADIR = joinpath(Pkg.dir("Celeste"), "test", "data", "decam")
    wd = pwd()
    cd(DATADIR)
    make_output = readstring(`make survey-ccds-dr5-0.0-0.5-0.0-0.5.fits`)
    if !startswith(make_output, "make: Nothing to be done for")
        print(make_output)  # only print output if something was done
    end
    cd(wd)

    @testset "overlapping ccds" begin
        dataset = DECALSDataSet(DATADIR,
                                "survey-ccds-dr5-0.0-0.5-0.0-0.5.fits")
        box = BoundingBox(0.0, 0.5, 0.0, 0.5)
        idx = _get_overlapping_ccds(dataset, box)
        @test length(idx) == 134  # all CCDs in FITS file should overlap.

        small_box = BoundingBox(0.2, 0.4, 0.2, 0.4)
        idx2 = _get_overlapping_ccds(dataset, small_box)
        @test length(idx2) == 82  # determined from running code!
                                  # (But should be less than 134)
    end
end


using DataFrames, JSON
using StringDistances, Statistics, DataStructures, Random, Pipe
import ParquetFiles
import Base: isless

include("utils.jl")
include("matchtypes.jl")
include("constructors.jl")
include("steps/init.jl")
include("steps/extend.jl")
include("steps/match.jl")
include("steps/integrate.jl")
include("steps/common.jl")

include("../test/testutils.jl")

## GEN

esp_dic = Dict("fing"=> generate_dfpair(3,1,0))
relp = RelPair([[1 1; 1 2;2 2], [2 2; 2 1]], ("fing", "fing"))
crs = CorefSystem(esp_dic, [relp])

shortside = 10
longside = 10
max_depth = 3
preplen = 3

crr = CorefResolver(crs, shortside, longside, max_depth, preplen);

steps = [
    InitParams("fing"),
    ExtendParams(1, 1),
    MatchParams(2),
    IntegrateParams(),
]

step!(crr, steps[1])

crr.leftspace_main
crr.space_pairs[1]
crr.space_pairs[2]

for erspace in crr |> main |> spaces
    (erspace.width == 1) |> println
    (erspace.height == 3) |> println
    (erspace.space[1:3,1] == [1, 2, 3]) |> println
end

step!(crr, steps[2])

(crr |> main).erspace1.width == 2
(crr |> main).erspace1.space[1:3, 1:2] == [1 1; 1 2; 2 2]

(crr |> main).erspace1.space[1:7,:]

step!(crr, steps[3])

crr.leftspace_main
crr.space_pairs[1]
crr.space_pairs[2]

(crr |> main).erspace1.space
(crr |> main).erspace2.space

step!(crr, steps[4])

crr.crspace.latent_match_vars

### FILE:
crsysarr = [CorefSystem("./test/test_data/test_sys_$i") for i in 0:2]

crs = crsysarr[1]

shortside = 600
longside = 200
max_depth = 3
preplen = 5

crr = CorefResolver(crs, shortside, longside, max_depth, preplen);

steps = [
    InitParams("e"),
    MatchParams(8),
    IntegrateParams(),
]

for params in steps
    step!(crr, params)
end

step!(crr, steps[end])

(crr |> main).erspace1.width#space[1:8,:]

res = crr |> result

tres0 = load_result("./test/test_data/test_sys_0")

res["e"] == tres0["e"]


# is caching useful?  only in extreme cases, apparently
# (but we may need to define equality on states!)

function slow(n)
#    matcher = Repeat(Repeat(Equal("a"), 0, n), n, 0)
#    matcher = Seq(Repeat(Equal("a"), 0, ), Repeat(Equal("a"), 0, n))
#    for greedy in (true, false)
    for greedy in (true,)  # false just eats memory
        println("greedy $greedy")
        matcher = Repeat(Equal("a"), 0, n; greedy=greedy)
        for i in 1:n
            matcher = Seq(Repeat(Equal("a"), 0, n; greedy=greedy), matcher)
        end
        source = repeat("a", n)
        for config in (NoCache, Cache)
            println("$(config)")
            all1 = make_all(config)
            @time collect(all1(source, matcher))
            @time n = length(collect(all1(source, matcher)))
            println("n results: $n")
            debug, all2 = make(Debug, source, matcher; delegate=config)
            collect(all2)
            println("max depth: $(debug.max_depth)")
            println("max iter: $(debug.max_iter)")
            println("n calls: $(debug.n_calls)")
        end
    end
end
slow(3)
#slow(7)  # not for travis!

println("slow ok")


function test_func(x)
    function max_missing(x)
        """Finds the maximum value of a array where missing are present."""
            y = replace(copy(x), missing => 0)
            return(maximum(y))
        end
    max = max_missing(x)
    return x.+ max
end

@testset "cross_apply" begin
    for i in 1:10
        x = rand(10:30)
        y = rand(10:30)
        z = rand(10:30)
        rad = rand(20:100.0, x, y, z)
        @test size(cross_apply(test_func, rad)) == size(rad)
    end
end

@testset "long_apply" begin
    for i in 1:10
        x = rand(10:30)
        y = rand(10:30)
        z = rand(10:30)
        rad = rand(20:100.0, x, y, z)
        mask = rand(0:1, x, y)
        @test size(long_apply(test_func, rad, mask)) == size(rad)
    end
end


abstract type AngularQuadrature end
abstract type UnstructuredMesh_2D end
abstract type LinearUnstructuredMesh_2D <: UnstructuredMesh_2D end
abstract type QuadraticUnstructuredMesh_2D <: UnstructuredMesh_2D end


using jcontrol
using jcontrol.visuals
import jcontrol: Track

plt = draw(:arena)

track = Track(;start_waypoint=4, keep_n_waypoints=-1)
draw!(track)

display(plt)