mutable struct Prop3DAcoTTIDenQ_DEO2_FDTD
    p::Ptr{Cvoid}
end

function Prop3DAcoTTIDenQ_DEO2_FDTD(;
        nz=0,
        ny=0,
        nx=0,
        nsponge=60,
        nbz=512,
        nby=8,
        nbx=8,
        dz=5.0,
        dy=8.0,
        dx=5.0,
        dt=0.001,
        nthreads=Sys.CPU_THREADS,
        freesurface=true,
        freqQ=5.0,
        qMin=0.1,
        qInterior=100.0)
    nz,ny,nx,nsponge,nbz,nby,nbx,nthreads = map(x->round(Int,x), (nz,ny,nx,nsponge,nbz,nby,nbx,nthreads))
    dz,dy,dx,dt = map(x->Float32(x), (dz,dy,dx,dt))

    @assert nx > 0
    @assert ny > 0
    @assert nz > 0
    @assert nsponge > 0
    @assert nthreads > 0
    @assert nbx > 0
    @assert nby > 0
    @assert nbz > 0

    fs = freesurface ? 1 : 0

    p = ccall((:Prop3DAcoTTIDenQ_DEO2_FDTD_alloc, libprop3DAcoTTIDenQ_DEO2_FDTD),
        Ptr{Cvoid},
        (Clong, Clong,    Clong, Clong, Clong, Clong,   Cfloat, Cfloat, Cfloat, Cfloat, Clong, Clong, Clong),
         fs,    nthreads, nx,    ny,    nz,    nsponge, dx,     dy,     dz,     dt,     nbx,   nby,   nbz)

    ccall((:Prop3DAcoTTIDenQ_DEO2_FDTD_SetupDtOmegaInvQ, libprop3DAcoTTIDenQ_DEO2_FDTD),
        Cvoid,
        (Ptr{Cvoid}, Cfloat, Cfloat, Cfloat),
         p,          freqQ,  qMin,   qInterior)

    Prop3DAcoTTIDenQ_DEO2_FDTD(p)
end

free(prop::Prop3DAcoTTIDenQ_DEO2_FDTD) = ccall((:Prop3DAcoTTIDenQ_DEO2_FDTD_free, libprop3DAcoTTIDenQ_DEO2_FDTD), Cvoid, (Ptr{Cvoid},), prop.p)

function size(prop::Prop3DAcoTTIDenQ_DEO2_FDTD)
    nz = ccall((:Prop3DAcoTTIDenQ_DEO2_FDTD_getNz, libprop3DAcoTTIDenQ_DEO2_FDTD), (Clong), (Ptr{Cvoid},), prop.p)
    ny = ccall((:Prop3DAcoTTIDenQ_DEO2_FDTD_getNy, libprop3DAcoTTIDenQ_DEO2_FDTD), (Clong), (Ptr{Cvoid},), prop.p)
    nx = ccall((:Prop3DAcoTTIDenQ_DEO2_FDTD_getNx, libprop3DAcoTTIDenQ_DEO2_FDTD), (Clong), (Ptr{Cvoid},), prop.p)
    (nz,ny,nx)
end

for _f in (:V, :Eps, :Eta, :B, :F, :PSpace, :MSpace, :PCur, :POld, :MCur, :MOld, :SinTheta, :CosTheta, :SinPhi, :CosPhi, :DtOmegaInvQ)
        symf = "Prop3DAcoTTIDenQ_DEO2_FDTD_get" * string(_f)
    @eval $(_f)(prop::Prop3DAcoTTIDenQ_DEO2_FDTD) = unsafe_wrap(Array, ccall(($symf, libprop3DAcoTTIDenQ_DEO2_FDTD), Ptr{Float32}, (Ptr{Cvoid},), prop.p), size(prop), own=false)
end

propagateforward!(prop::Prop3DAcoTTIDenQ_DEO2_FDTD) = ccall((:Prop3DAcoTTIDenQ_DEO2_FDTD_TimeStep, libprop3DAcoTTIDenQ_DEO2_FDTD), Cvoid, (Ptr{Cvoid},), prop.p)
propagateadjoint!(prop::Prop3DAcoTTIDenQ_DEO2_FDTD) = propagateforward!(prop) # self-adjoint

scale_spatial_derivatives!(prop::Prop3DAcoTTIDenQ_DEO2_FDTD) =
    ccall((:Prop3DAcoTTIDenQ_DEO2_FDTD_ScaleSpatialDerivatives, libprop3DAcoTTIDenQ_DEO2_FDTD), Cvoid, (Ptr{Cvoid},), prop.p)

abstract type Prop3DAcoTTIDenQ_DEO2_FDTD_Model end
struct Prop3DAcoTTIDenQ_DEO2_FDTD_Model_VEA <: Prop3DAcoTTIDenQ_DEO2_FDTD_Model end
struct Prop3DAcoTTIDenQ_DEO2_FDTD_Model_V <: Prop3DAcoTTIDenQ_DEO2_FDTD_Model end

# v,ϵ and η in model-space
function forwardBornInjection!(prop::Prop3DAcoTTIDenQ_DEO2_FDTD, modeltype::Prop3DAcoTTIDenQ_DEO2_FDTD_Model_VEA, dmodel, wavefields)
    ccall((:Prop3DAcoTTIDenQ_DEO2_FDTD_ForwardBornInjection_VEA, libprop3DAcoTTIDenQ_DEO2_FDTD), Cvoid,
        (Ptr{Cvoid}, Ptr{Cfloat}, Ptr{Cfloat}, Ptr{Cfloat}, Ptr{Cfloat},        Ptr{Cfloat},        Ptr{Cfloat},          Ptr{Cfloat}),
         prop.p,     dmodel["v"], dmodel["ϵ"], dmodel["η"], wavefields["pold"], wavefields["mold"], wavefields["pspace"], wavefields["mspace"])
end

function adjointBornAccumulation!(prop::Prop3DAcoTTIDenQ_DEO2_FDTD, modeltype::Prop3DAcoTTIDenQ_DEO2_FDTD_Model_VEA,
        imagingcondition::ImagingConditionStandard, dmodel, wavefields)
    ccall((:Prop3DAcoTTIDenQ_DEO2_FDTD_AdjointBornAccumulation_VEA, libprop3DAcoTTIDenQ_DEO2_FDTD), Cvoid,
        (Ptr{Cvoid}, Ptr{Cfloat}, Ptr{Cfloat}, Ptr{Cfloat}, Ptr{Cfloat},        Ptr{Cfloat},        Ptr{Cfloat},          Ptr{Cfloat}),
         prop.p,     dmodel["v"], dmodel["ϵ"], dmodel["η"], wavefields["pold"], wavefields["mold"], wavefields["pspace"], wavefields["mspace"])
end

# v in model-space
function forwardBornInjection!(prop::Prop3DAcoTTIDenQ_DEO2_FDTD, modeltype::Prop3DAcoTTIDenQ_DEO2_FDTD_Model_V, dmodel, wavefields)
    ccall((:Prop3DAcoTTIDenQ_DEO2_FDTD_ForwardBornInjection_V, libprop3DAcoTTIDenQ_DEO2_FDTD), Cvoid,
        (Ptr{Cvoid}, Ptr{Cfloat}, Ptr{Cfloat},          Ptr{Cfloat}),
         prop.p,     dmodel["v"], wavefields["pspace"], wavefields["mspace"])
end

function adjointBornAccumulation!(prop::Prop3DAcoTTIDenQ_DEO2_FDTD, modeltype::Prop3DAcoTTIDenQ_DEO2_FDTD_Model_V,
        imagingcondition::ImagingConditionStandard, dmodel, wavefields)
    ccall((:Prop3DAcoTTIDenQ_DEO2_FDTD_AdjointBornAccumulation_V, libprop3DAcoTTIDenQ_DEO2_FDTD), Cvoid,
        (Ptr{Cvoid}, Ptr{Cfloat}, Ptr{Cfloat},          Ptr{Cfloat}),
         prop.p,     dmodel["v"], wavefields["pspace"], wavefields["mspace"])
end

function adjointBornAccumulation!(prop::Prop3DAcoTTIDenQ_DEO2_FDTD, modeltype::Prop3DAcoTTIDenQ_DEO2_FDTD_Model_V,
        imagingcondition::ImagingConditionWaveFieldSeparationFWI, dmodel, wavefields)
    ccall((:Prop3DAcoTTIDenQ_DEO2_FDTD_AdjointBornAccumulation_wavefieldsep_V, libprop3DAcoTTIDenQ_DEO2_FDTD), Cvoid,
        (Ptr{Cvoid}, Ptr{Cfloat}, Ptr{Cfloat},          Ptr{Cfloat},          Clong),
         prop.p,     dmodel["v"], wavefields["pspace"], wavefields["mspace"], 1)
end

function adjointBornAccumulation!(prop::Prop3DAcoTTIDenQ_DEO2_FDTD, modeltype::Prop3DAcoTTIDenQ_DEO2_FDTD_Model_V,
        imagingcondition::ImagingConditionWaveFieldSeparationRTM, dmodel, wavefields)
    ccall((:Prop3DAcoTTIDenQ_DEO2_FDTD_AdjointBornAccumulation_wavefieldsep_V, libprop3DAcoTTIDenQ_DEO2_FDTD), Cvoid,
        (Ptr{Cvoid}, Ptr{Cfloat}, Ptr{Cfloat},          Ptr{Cfloat},          Clong),
         prop.p,     dmodel["v"], wavefields["pspace"], wavefields["mspace"], 0)
end

function show(io::IO, prop::Prop3DAcoTTIDenQ_DEO2_FDTD)
    nx = ccall((:Prop3DAcoTTIDenQ_DEO2_FDTD_getNx, libprop3DAcoTTIDenQ_DEO2_FDTD), (Clong), (Ptr{Cvoid},), prop.p)
    ny = ccall((:Prop3DAcoTTIDenQ_DEO2_FDTD_getNy, libprop3DAcoTTIDenQ_DEO2_FDTD), (Clong), (Ptr{Cvoid},), prop.p)
    nz = ccall((:Prop3DAcoTTIDenQ_DEO2_FDTD_getNz, libprop3DAcoTTIDenQ_DEO2_FDTD), (Clong), (Ptr{Cvoid},), prop.p)
    write(io, "Prop3DAcoTTIDenQ_DEO2_FDTD -- nx,ny,nz; $nx,$ny,$nz")
end

function tilesize(::Type{Prop3DAcoTTIDenQ_DEO2_FDTD}; nz, ny, nx, rngy=4:2:64, rngx=4:2:64, N=10, nthreads=Sys.CPU_THREADS-1)
    speeds = Matrix{Float64}(undef, length(rngy), length(rngx))
    for (iby,by) in enumerate(rngy), (ibx,bx) in enumerate(rngx)
        @info "by=$by, bx=$bx"
        prop = Prop3DAcoTTIDenQ_DEO2_FDTD(nz=nz, ny=ny, nx=nx, nbz=nz, nby=by, nbx=bx, nthreads=nthreads)

        v = V(prop)
        ϵ = Eps(prop)
        η = Eta(prop)
        sinθ = SinTheta(prop)
        cosθ = CosTheta(prop)
        sinϕ = SinPhi(prop)
        cosϕ = CosPhi(prop)
        b = B(prop)
        f = F(prop)
        pcur = PCur(prop)
        pold = POld(prop)
        mcur = MCur(prop)
        mold = MOld(prop)

        v .= 1500
        ϵ .= 1
        η .= 1
        sinθ .= 0
        cosθ .= 1
        sinϕ .= 0
        cosϕ .= 1
        b .= 1
        f .= 1
        rand!(pcur)
        rand!(pold)
        rand!(mcur)
        rand!(mold)

        t = 0.0
        set_zero_subnormals(true)
        for i = 1:N
            t += @elapsed propagateforward!(prop)
        end
        set_zero_subnormals(false)
        speeds[iby,ibx] = N*prod(size(prop))/1000/1000/t
        free(prop)
    end
    rngx,rngy,speeds
end


module midiModule

using MIDI
function getMidiDatas(filePath)

#midiファイル読み込み
println("koko")
midi = MIDI.readMIDIFile(filePath)
#1トラック目取得
notes = getnotes(midi.tracks[1], midi.tpq)
#ピッチ、発音位置、ベロシティー取得
pitchSequence = []
positionSequence = []
velocitySequence = []
unit = 120
beforeNotePosition = 0
beforeNoteDuration = 0
cnt = 0
for note in notes
    cnt += 1
    push!(pitchSequence, Int(note.pitch))
    push!(positionSequence, Int(note.position))
    # ベロシティーは、休符のところは0、音符が伸びて鳴っている箇所はそのノートが持つベロシティーを、duration=120ごとに1要素として配列で取得
    # 休符がある場合
    if note.position - beforeNotePosition - beforeNoteDuration + 1 > 0 && cnt != 1
        #その幅だけ休符を追加
        append!(velocitySequence, repeat([0], div(note.position - beforeNotePosition - beforeNoteDuration + 1, unit)))
    end
    append!(velocitySequence, repeat([Int(note.velocity)], div(Int(note.duration + 1), unit)))
    beforeNotePosition = Int(note.position)
    beforeNoteDuration = Int(note.duration)
end
return pitchSequence, positionSequence, velocitySequence
end

end


using Pkg, InteractiveUtils
ENV["PYTHON"] = ""
Pkg.build("Conda")
Pkg.build("PyCall")
using Conda
Conda.add("pytorch"; channel = "pytorch")
pkg"precompile"


#!/usr/bin/env julia

using ArgParse
using JSON
using Word2Vec

s = ArgParseSettings()
@add_arg_table s begin
    "--vwify", "-v"
    help = "convert JSON-GAM input to vw compatible format"
    action = :store_true
    "--importance", "-i"
    help = "give errors a weight proportional to their mapping quality"
    action = :store_true
    "--phred", "-p"
    help = "regressor already outputs phred format, so don't convert"
    action = :store_true
    "--min-mq", "-q"
    help = "take the minimum of the previous and new mapping quality"
    action = :store_true
    "--mq-model", "-m"
    arg_type = String
    help = "set the MQ using the given model"
    default = ""
    "--dna-embedding", "-e"
    arg_type = String
    help = "use the DNA sequence embedding to generate sequence features"
    default = ""
    "--vw-opts", "-o"
    arg_type = String
    help = "additional options to pass to vw"
    default = ""
    "--kmer-size", "-k"
    arg_type = Int
    help = "kmer size to break input DNA sequences into"
    default = 5
    "--kmer-stride", "-s"
    arg_type = Int
    help = "bases between each successive kmer"
    default = 1
    "--cigar", "-c"
    help = "write the cigar base-by-base in a namespace"
    action = :store_true
    "--sequence", "-S"
    help = "write the read sequence base-by-base"
    action = :store_true
end

parsed_args = parse_args(ARGS, s)

embedding_file = parsed_args["dna-embedding"]
embedding = (embedding_file != "" ? wordvectors(embedding_file) : nothing)
kmer_size = parsed_args["kmer-size"]
kmer_stride = parsed_args["kmer-stride"]
write_cigar = parsed_args["cigar"]
write_sequence = parsed_args["sequence"]

global MQ_MAX = 60

function float2phred(prob)
    return min(-10 * log10(prob), MQ_MAX)
end

function phred2float(qual)
    return 10 ^ (qual * -.1)
end

function prob2mq(prob)
    v = 1-prob
    if v == 1
        return 0
    elseif v > 0
        return min(MQ_MAX, Int32(round(float2phred(v))))
    else
        return MQ_MAX
    end
end

function applymodel(alns, model, params)
    (path,io) = mktemp()
    for aln in alns
        println(io, vwify(aln, kmer_size, kmer_stride))
    end
    close(io)
    lines = split(chomp(readstring(pipeline(path, `vw -t -p /dev/stdout -i $model --quiet $(split(params, " "))`))), "\n")
    rm(path)
    assert(length(lines) == length(alns))
    for (line,aln) in zip(lines, alns)
        (result,name) = split(line, " ")
        assert(name == aln["name"])
        result = parsed_args["phred"] ? Int(round(float(result))) : prob2mq(float(result))
        mapping_quality = haskey(aln, "mapping_quality") ? aln["mapping_quality"] : 0
        aln["mapping_quality"] = (parsed_args["min-mq"] ? min(mapping_quality, result) : result)
    end
    return alns
end

function kmers_of(k::Int, j::Int, s::String)
    [s[i:(i+k-1)] for i in range(1,j,Int(floor((length(s)-k)/j))+1)]
end

function seq_embedding(s::String, k::Int, j::Int)
    function in_vocab(kmer::String)
        haskey(embedding.vocab_hash, kmer)
    end
    kmers = filter(in_vocab, kmers_of(k, j, s))
    reduce(+, [get_vector(embedding, kmer) for kmer in kmers])/length(kmers)
end

function vwify(aln, kmer_size::Int, kmer_stride::Int, mqimportance=false)
    name = haskey(aln, "name") ? aln["name"] : ""
    correct = haskey(aln, "correct") ? aln["correct"] : 0
    correct = correct > 0.5 ? 1 : 0
    if haskey(aln, "correct") aln["correct"] = correct end
    score = haskey(aln, "score") ? aln["score"] : 0
    identity = haskey(aln, "identity") ? aln["identity"] : 0
    mapqual = haskey(aln, "mapping_quality") ? aln["mapping_quality"] : 0
    path = haskey(aln, "path") && haskey(aln["path"], "mapping") ? aln["path"]["mapping"] : []
    i = 1
    cigar = []
    if write_cigar
        for mapping in path
            for elem in mapping["edit"]
                #println(elem)
                fromlen = haskey(elem, "from_length") ? elem["from_length"] : 0
                tolen = haskey(elem, "to_length") ? elem["to_length"] : 0
                seq = haskey(elem, "sequence") ? elem["sequence"] : ""
                if fromlen == tolen
                    if seq != ""
                        for j in range(i, tolen) push!(cigar, string(j, "S")) end
                    else
                        for j in range(i, tolen) push!(cigar, string(j, "M")) end
                    end
                elseif fromlen > tolen
                    for j in range(i, fromlen) push!(cigar, string(j, "D")) end
                elseif tolen > fromlen
                    for j in range(i, tolen) push!(cigar, string(j, "I")) end
                end
                i += tolen
            end
        end
    end
    sequence = write_sequence ? string("|bases ", join([string(x, y) for (x,y) in zip(1:length(aln["sequence"]),aln["sequence"])], " ")) : ""
    seqvec = (embedding == nothing ? "" : string("|dnavec ", join(["d$i:$(Float16(j))" for (i,j) in enumerate(seq_embedding(aln["sequence"], kmer_size, kmer_stride))], " ")))
    cigar = write_cigar ? string("|cigar ", join(cigar, " ")) : ""
    secondaries = haskey(aln, "secondary_score") ? aln["secondary_score"] : []
    secondaries = join(["s$i:$score" for (i,score) in zip(1:length(secondaries),secondaries)], " ")
    importance = (mqimportance && correct == 0) ? "$(mapqual+2) " : ((mqimportance && correct == 1 && mapqual < 60) ? "$(60-mapqual+1) " : " ")
    string("$correct $importance '$name |alignment identity:$identity score:$score mapqual:$mapqual |scores s0:$score $secondaries $seqvec $sequence $cigar")
end

alns = []
for line in eachline(STDIN)
    aln = JSON.parse(line)
    if parsed_args["vwify"]
        println(vwify(aln, kmer_size, kmer_stride, parsed_args["importance"]))
    elseif parsed_args["mq-model"] != ""
        # we should run the model on the input and set a new mapping quality
        push!(alns, aln)
        if length(alns) >= 1000
            assert(length(alns) > 0)
            for aln in applymodel(alns, parsed_args["mq-model"], parsed_args["vw-opts"])
                println(JSON.json(aln))
            end
            alns = []
        end
    else
        # default is to do nothing to input
        println(JSON.json(aln))
    end
end

# clean up the buffer
if parsed_args["mq-model"] != "" && length(alns) > 0
    for aln in applymodel(alns, parsed_args["mq-model"], parsed_args["vw-opts"])
        println(JSON.json(aln))
    end
    alns = []
end


"""
    DeterministicEquivalent

Deterministic equivalent memory structure. Stochastic program is stored as one large optimization problem. Supported by any standard `AbstractOptimizer`.

"""
struct DeterministicEquivalent{N, M, S <: NTuple{M, Scenarios}} <: AbstractStochasticStructure{N}
    decisions::Decisions
    decision_variables::NTuple{M, Vector{DecisionRef}}
    scenarios::S
    sub_objectives::NTuple{M, Vector{MOI.AbstractScalarFunction}}
    model::JuMP.Model

    function DeterministicEquivalent(scenarios::NTuple{M, Scenarios}) where M
        N = M + 1
        decisions = Decisions()
        decision_variables = ntuple(Val(M)) do i
            Vector{DecisionRef}()
        end
        sub_objectives = ntuple(Val(M)) do i
            Vector{MOI.AbstractScalarFunction}()
        end
        S = typeof(scenarios)
        return new{N,M,S}(decisions, decision_variables, scenarios, sub_objectives, Model())
    end
end

function StochasticStructure(scenario_types::ScenarioTypes{M}, ::Deterministic) where M
    scenarios = ntuple(Val(M)) do i
        Vector{scenario_types[i]}()
    end
    return DeterministicEquivalent(scenarios)
end

function StochasticStructure(scenarios::NTuple{M, Vector{<:AbstractScenario}}, ::Deterministic) where M
    return DeterministicEquivalent(scenarios)
end

# Base overloads #
# ========================== #
function Base.print(io::IO, structure::DeterministicEquivalent)
    print(io, "Deterministic equivalent problem\n")
    print(io, structure.model)
end
# ========================== #

# MOI #
# ========================== #
function MOI.get(structure::DeterministicEquivalent, attr::MOI.AbstractModelAttribute)
    return MOI.get(backend(structure.model), attr)
end
function MOI.get(structure::DeterministicEquivalent, attr::MOI.AbstractVariableAttribute, index::MOI.VariableIndex)
    return MOI.get(backend(structure.model), attr, index)
end
function MOI.get(structure::DeterministicEquivalent, attr::Type{MOI.VariableIndex}, name::String)
    return MOI.get(backend(structure.model), attr, name)
end
function MOI.get(structure::DeterministicEquivalent, attr::MOI.AbstractConstraintAttribute, cindex::MOI.ConstraintIndex)
    return MOI.get(backend(structure.model), attr, cindex)
end

function MOI.set(structure::DeterministicEquivalent, attr::MOI.Silent, flag)
    MOI.set(backend(structure.model), attr, flag)
    return nothing
end
function MOI.set(structure::DeterministicEquivalent, attr::MOI.AbstractModelAttribute, value)
    MOI.set(backend(structure.model), attr, value)
    return nothing
end
function MOI.set(structure::DeterministicEquivalent, attr::MOI.AbstractVariableAttribute,
                 index::MOI.VariableIndex, value)
    MOI.set(backend(structure.model), attr, index, value)
    return nothing
end
function MOI.set(structure::DeterministicEquivalent, attr::MOI.AbstractConstraintAttribute,
                 cindex::MOI.ConstraintIndex, value)
    MOI.set(backend(structure.model), attr, cindex, value)
    return nothing
end

function MOI.is_valid(structure::DeterministicEquivalent, index::MOI.VariableIndex)
    return MOI.is_valid(backend(structure.model), index)
end

function MOI.add_constraint(structure::DeterministicEquivalent, f::MOI.AbstractFunction, s::MOI.AbstractSet)
    return MOI.add_constraint(backend(structure.model), f, s)
end

function MOI.delete(structure::DeterministicEquivalent, index::MOI.Index)
    # TODO: more to do if index is decision
    MOI.delete(backend(structure.model), index)
    return nothing
end

# Getters #
# ========================== #
function structure_name(structure::DeterministicEquivalent)
    return "Deterministic equivalent"
end
function decision(structure::DeterministicEquivalent, index::MOI.VariableIndex)
    return decision(structure.decisions, index)
end
function decisions(structure::DeterministicEquivalent)
    return structure.decisions
end
function all_decisions(structure::DeterministicEquivalent)
    return structure.decisions.undecided
end
function num_decisions(structure::DeterministicEquivalent{N}, s::Integer = 2) where N
    1 <= s <= N || error("Stage $s not in range 1 to $N.")
    s == N && error("The final stage does not have decisions.")
    return length(structure.decision_variables[s])
end
function scenario(structure::DeterministicEquivalent{N}, i::Integer, s::Integer = 2) where N
    1 <= s <= N || error("Stage $s not in range 1 to $N.")
    s == 1 && error("The first stage does not have scenarios.")
    return structure.scenarios[s-1][i]
end
function scenarios(structure::DeterministicEquivalent{N}, s::Integer = 2) where N
    1 <= s <= N || error("Stage $s not in range 1 to $N.")
    s == 1 && error("The first stage does not have scenarios.")
    return structure.scenarios[s-1]
end
function subproblem(structure::DeterministicEquivalent, i::Integer, s::Integer = 2)
    return subproblem(scenarioproblems(structure, s), i)
end
function subproblems(structure::DeterministicEquivalent, s::Integer = 2)
    return subproblems(scenarioproblems(structure, s))
end
function num_subproblems(structure::DeterministicEquivalent, s::Integer = 2)
    return 0
end
function deferred(structure::DeterministicEquivalent)
    return num_variables(structure.model) == 0
end
# ========================== #

# Setters
# ========================== #
function update_decisions!(structure::DeterministicEquivalent, change::DecisionModification)
    update_decisions!(structure.model, change)
end

function add_scenario!(structure::DeterministicEquivalent, scenario::AbstractScenario, stage::Integer = 2)
    push!(scenarios(structure, stage), scenario)
    return nothing
end
function add_worker_scenario!(structure::DeterministicEquivalent, scenario::AbstractScenario, w::Integer, stage::Integer = 2)
    add_scenario!(structure, scenario, stage)
    return nothing
end
function add_scenario!(scenariogenerator::Function, structure::DeterministicEquivalent, stage::Integer = 2)
    add_scenario!(structure, scenariogenerator(), stage)
    return nothing
end
function add_worker_scenario!(scenariogenerator::Function, structure::DeterministicEquivalent, w::Integer, stage::Integer = 2)
    add_scenario!(scenariogenerator, structure, stage)
    return nothing
end
function add_scenarios!(structure::DeterministicEquivalent, _scenarios::Vector{<:AbstractScenario}, stage::Integer = 2)
    append!(scenarios(structure, stage), _scenarios)
    return nothing
end
function add_worker_scenarios!(structure::DeterministicEquivalent, scenarios::Vector{<:AbstractScenario}, w::Integer, stage::Integer = 2)
    add_scenarios!(structure, scenarios, stasge)
    return nothing
end
function add_scenarios!(scenariogenerator::Function, structure::DeterministicEquivalent, n::Integer, stage::Integer = 2)
    for i = 1:n
        add_scenario!(structure, stage) do
            return scenariogenerator()
        end
    end
    return nothing
end
function add_worker_scenarios!(scenariogenerator::Function, structure::DeterministicEquivalent, n::Integer, w::Integer, stage::Integer = 2)
    add_scenarios!(scenariogenerator, structure, n, stage)
    return nothing
end
function sample!(structure::DeterministicEquivalent, sampler::AbstractSampler, n::Integer, stage::Integer = 2)
    sample!(scenarios(structure, stage), sampler, n)
    return nothing
end
# ========================== #


### TODO: Rethink the CUDNN interface

using CUDNN: cudnnConvolutionDescriptor_t, cudnnCreateConvolutionDescriptor, cudnnSetConvolutionNdDescriptor_v3, cudnnDestroyConvolutionDescriptor, cudnnHandle, cptr, TD, FD, cudnnDataType
using CUDNN: CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_PRECOMP_GEMM, CUDNN_CONVOLUTION_BWD_FILTER_ALGO_0, CUDNN_CONVOLUTION_BWD_DATA_ALGO_0
using CUDNN: cudnnGetConvolutionNdForwardOutputDim, cudnnGetConvolutionForwardWorkspaceSize
using CUDNN: cudnnConvolutionBackwardFilter_v3, cudnnConvolutionBackwardData_v3

function cudnnConvolutionForward_v4(src, filter, dest;
                                    handle=cudnnHandle, alpha=1.0, beta=0.0,
                                    algorithm=CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_PRECOMP_GEMM,
                                    workSpace=C_NULL, workSpaceSizeInBytes=0,
                                    padding=0, stride=1, upscale=1, mode=CUDNN_CONVOLUTION)
    @assert eltype(filter) == eltype(src)
    cd = cudnnConvolutionDescriptor_v4(ndims(src), eltype(src), padding, stride, upscale, mode)
    osize = cudnnGetConvolutionNdForwardOutputDim(src,filter;convDesc=cd)
    (dest == nothing) && (dest = CudaArray(eltype(src), osize))
    @assert osize == size(dest)
    @assert eltype(dest) == eltype(src)
    wsize = cudnnGetConvolutionForwardWorkspaceSize(src, filter, dest; algorithm=algorithm)
    if ((wsize > 0) && (workSpace == C_NULL || workSpaceSizeInBytes < wsize))
        workSpaceSizeInBytes = wsize
        ws = CudaArray(Int8, workSpaceSizeInBytes)
    else
        ws = workSpace
    end
    cudnnConvolutionForward(handle,
                            cptr(alpha,src),TD(src),src,
                            FD(filter),filter,
                            cd,algorithm,ws,workSpaceSizeInBytes,
                            cptr(beta,dest),TD(dest),dest)
    free(cd)
    ws === workSpace || free(ws)
    return dest
end

# I am guessing if y=w*x+b going forward, the arguments below
# correspond to src=x, diff=dy, grad=dw.
function cudnnConvolutionBackwardFilter_v4(src, diff, grad;
                                           handle=cudnnHandle, alpha=1.0, beta=0.0,
                                           algo=CUDNN_CONVOLUTION_BWD_FILTER_ALGO_0,
                                           workSpace=C_NULL, workSpaceSizeInBytes=0,
                                           padding=0, stride=1, upscale=1, mode=CUDNN_CONVOLUTION)
    cd = cudnnConvolutionDescriptor_v4(ndims(src), eltype(src), padding, stride, upscale, mode)
    cudnnConvolutionBackwardFilter_v3(handle,
                                      cptr(alpha,src),TD(src),src,
                                      TD(diff),diff,cd,
                                      algo, workSpace, workSpaceSizeInBytes,
                                      cptr(beta,grad),FD(grad),grad)
    free(cd)
    return grad
end

# I am guessing if y=w*x+b going forward, the arguments below
# correspond to filter=w, diff=dy, grad=dx.
function cudnnConvolutionBackwardData_v4(filter::AbstractCudaArray, diff::AbstractCudaArray, grad::AbstractCudaArray;
                                         handle=cudnnHandle, alpha=1.0, beta=0.0,
                                         algo=CUDNN_CONVOLUTION_BWD_DATA_ALGO_0,
                                         workSpace=C_NULL, workSpaceSizeInBytes=0,
                                         padding=0, stride=1, upscale=1, mode=CUDNN_CONVOLUTION)
    cd = cudnnConvolutionDescriptor_v4(ndims(filter), eltype(filter), padding, stride, upscale, mode)
    cudnnConvolutionBackwardData_v3(handle,cptr(alpha,diff),
                                    FD(filter),filter,
                                    TD(diff),diff,cd,
                                    algo, workSpace, workSpaceSizeInBytes,
                                    cptr(beta,grad),TD(grad),grad)
    free(cd)
    return grad
end

# Decided not to keep the cd in the Pool structure: not serializable.
function cudnnConvolutionDescriptor_v4(xdims, xtype, padding, stride, upscale, mode)
    cd = cudnnConvolutionDescriptor_t[0]
    cudnnCreateConvolutionDescriptor(cd)
    nd = xdims-2
    cudnnSetConvolutionNdDescriptor_v3(cd[1],nd,cdsize(padding,nd),cdsize(stride,nd),cdsize(upscale,nd),mode,cudnnDataType(xtype))
    return cd[1]
end

function cdsize(w, nd)
    isa(w,Integer) ? Cint[fill(w,nd)...] :
    length(w)!=nd ? error("Dimension mismatch") :
    Cint[reverse(w)...]
end

CUDArt.free(cd::cudnnConvolutionDescriptor_t)=cudnnDestroyConvolutionDescriptor(cd)

### TODO: THIS INTERFACE SHOULD GO INTO CUDNN:

using CUDNN: cudnnPoolingDescriptor_t, cudnnCreatePoolingDescriptor, cudnnSetPoolingNdDescriptor, cudnnDestroyPoolingDescriptor, cudnnHandle, cptr, TD, cudnnPoolingMode_t, cudnnGetPoolingNdDescriptor

function cudnnPoolingForward_v4(src, dest;
                                window=2, padding=0, stride=window, mode=CUDNN_POOLING_MAX,
                                handle=cudnnHandle, alpha=1.0, beta=0.0)
    pd = cudnnPoolingDescriptor_v4(ndims(src)-2, window, padding, stride, mode)
    cudnnPoolingForward(handle, pd,
                        cptr(alpha,src), TD(src), src,
                        cptr(beta,dest), TD(dest), dest)
    free(pd)
    return dest
end

function cudnnPoolingBackward_v4(src, srcDiff, dest, destDiff;
                                 window=2, padding=0, stride=window, mode=CUDNN_POOLING_MAX,
                                 handle=cudnnHandle, alpha=1.0, beta=0.0)
    pd = cudnnPoolingDescriptor_v4(ndims(src)-2, window, padding, stride, mode)
    cudnnPoolingBackward(handle, pd,
                         cptr(alpha,src), TD(src), src,
                         TD(srcDiff), srcDiff,
                         TD(dest), dest,
                         cptr(beta,destDiff), TD(destDiff), destDiff)
    free(pd)
    return destDiff
end

# Decided not to keep the pd in the Pool structure: not serializable.
function cudnnPoolingDescriptor_v4(nd, window, padding, stride, mode)
    pd = cudnnPoolingDescriptor_t[0]
    cudnnCreatePoolingDescriptor(pd)
    cudnnSetPoolingNdDescriptor(pd[1],mode,nd,pdsize(window,nd),pdsize(padding,nd),pdsize(stride,nd))
    return pd[1]
end

function cudnnGetPoolingNdDescriptor_v4(pd::cudnnPoolingDescriptor_t, nd)
    m = cudnnPoolingMode_t[0]
    n = Cint[0]
    s = Array(Cint, nd)
    p = Array(Cint, nd)
    t = Array(Cint, nd)
    cudnnGetPoolingNdDescriptor(pd, nd, m, n, s, p, t)
    inttuple(x)=tuple(Int[x...]...)
    (m[1], n[1], inttuple(s), inttuple(p), inttuple(t))
end

CUDArt.free(pd::cudnnPoolingDescriptor_t)=cudnnDestroyPoolingDescriptor(pd)

pdsize(w, nd)=Cint[reverse(psize(w,nd))...]

psize(w, nd)=(isa(w,Integer)  ? fill(w,nd) : length(w) != nd ? error("Dimension mismatch") : w)


"""
# KM3NeT Data Definitions v1.3.1-8-gae7dfb2
https://git.km3net.de/common/km3net-dataformat
"""

module Reconstruction
  const JPP_RECONSTRUCTION_TYPE = 4000
  const JMUONBEGIN = 0
  const JMUONPREFIT = 1
  const JMUONSIMPLEX = 2
  const JMUONGANDALF = 3
  const JMUONENERGY = 4
  const JMUONSTART = 5
  const JLINEFIT = 6
  const JMUONEND = 99
  const JSHOWERBEGIN = 100
  const JSHOWERPREFIT = 101
  const JSHOWERPOSITIONFIT = 102
  const JSHOWERCOMPLETEFIT = 103
  const JSHOWER_BJORKEN_Y = 104
  const JSHOWERENERGYPREFIT = 105
  const JSHOWERPOINTSIMPLEX = 106
  const JSHOWERDIRECTIONPREFIT = 107
  const JSHOWEREND = 199
  const DUSJ_RECONSTRUCTION_TYPE = 200
  const DUSJSHOWERBEGIN = 200
  const DUSJSHOWERPREFIT = 201
  const DUSJSHOWERPOSITIONFIT = 202
  const DUSJSHOWERCOMPLETEFIT = 203
  const DUSJSHOWEREND = 299
  const AANET_RECONSTRUCTION_TYPE = 101
  const AASHOWERBEGIN = 300
  const AASHOWERFITPREFIT = 302
  const AASHOWERFITPOSITIONFIT = 303
  const AASHOWERFITDIRECTIONENERGYFIT = 304
  const AASHOWEREND = 399
  const JUSERBEGIN = 1000
  const JMUONVETO = 1001
  const JMUONPATH = 1003
  const JMCEVT = 1004
  const JUSEREND = 1099
  const RECTYPE_UNKNOWN = -1
  const RECSTAGE_UNKNOWN = -1
end


using Pkg
Pkg.activate(normpath(joinpath(@__DIR__, "..")))

using ATools
using NReco

using ArgParse
using Configurations
using DataFrames
using Glob
using HDF5
using Logging
using Printf

logger = SimpleLogger(stdout, Logging.Warn)
old_logger = global_logger(logger)

function makenema(args)

	lorf    = args["loralgo"]
	detconf = args["detconf"]
	dr      = args["dir"]
	outd    = args["odir"]
	outf    = args["ofile"]
	file_i  = args["filei"]
	file_l  = args["filel"]

	lor_algo = NReco.lor_kmeans
	if lorf == "lor_qmax"
		lor_algo = NReco.lor_maxq
	end

	if isdir(outd) == false
		mkdir(outd)
	end
	output = string(outd,"/", outf)

	## Patch here so we can read actually select in a directory which
	## is already partially processed. Needs to be improved in general.
	files  = sort(glob("*.h5", dr), by=x->parse(Int64, split(x, "-")[end][1:end-3]))

	if detconf != "default"
		dconf = from_toml(NReco.DetConf, detconf)
	else
		dconf = NReco.DetConf()
	end

	println("makenema configuration")
	println("detector configuration", dconf)
	println(" lor_algo  = $lorf")

	println("number of files in data dir = $(length(files))")
	println("reading = $(file_l - file_i + 1) files")
	println("output file  = $output")

	h5open(output, "w") do h5out
		pars_grp    = create_group(h5out, "selected_events")
		pars_dtype  = ATools.generate_hdf5_datatype(ATools.EventParameters)
		conf_grp    = create_group(h5out, "configuration")
		conf_dtype  = ATools.generate_hdf5_datatype(ATools.SimConfiguration)

		nevt, evts  = NReco.nemareco(files, dconf, file_i, file_l, lor_algo)

		rmin, rmax  = extrema(getfield.(evts, :r1))
		# Maybe combine the DetConf and SimConfiguration types?
		config_pars = ATools.SimConfiguration(SiPMPDE=dconf.pde,
						ElecSTD=dconf.sigma_tof, SiPMThr=dconf.ecut,
						SumQmin=dconf.qmin, SumQmax=dconf.qmax,
						Ntof=dconf.ntof, NEvent=nevt, Rmin=rmin,
						Rmax=rmax)#TODO: Calibration option implementation

		## Save to file.
		pars_dspace = dataspace(evts)
		pars_dset   = create_dataset(pars_grp, "EventParameters", pars_dtype, pars_dspace)
		write_dataset(pars_dset, pars_dtype, evts)
		conf_dspace = dataspace([config_pars])
		conf_dset   = create_dataset(conf_grp, "RunConfiguration", conf_dtype, conf_dspace)
		write_dataset(conf_dset, conf_dtype, [config_pars])#Must be a better way
	end
end

function parse_commandline()
    s = ArgParseSettings()

    @add_arg_table s begin
        "--dir", "-d"
            help = "directory with nema simulations"
            arg_type = String
		"--odir", "-o"
            help = "output directory"
            arg_type = String
		"--ofile", "-x"
            help = "output file"
            arg_type = String
            default = "evtdf.h5"
		"--filei", "-i"
	        help = "number of initial file in glob list"
	        default  = 1
			arg_type = Int
		"--filel", "-l"
		    help = "number of last file in glob list"
		    default  = 1
			arg_type = Int
		"--loralgo", "-g"
			help = "algorithm to use for LOR reconstruction "
			arg_type = String
			default = "lor_kmeans"
		"--detconf", "-c"
			help = "Detector configuration"
			arg_type = String
			default = "default"
    end

    return parse_args(s)
end

function main()
	parsed_args = parse_commandline()
	println("Running makenema with arguments", parsed_args)
	makenema(parsed_args)
end

@time main()


using Pkg; Pkg.activate(first(splitdir(@__DIR__)))

using CUDA
using DLPack
using PyCall

@pyimport hoomd
@pyimport hoomd.md as md
@pyimport hoomd.dump as dump
@pyimport hoomd.group as group
@pyimport hoomd.benchmark as benchmark

@pyimport numpy as np

N = 10000
mode_string = "gpu"

if mode_string ∉ ("cpu", "gpu")
    throw("""Execution mode argument must be either "cpu" or "gpu".""")
end

# %%
hoomd.context.initialize("--mode=" * mode_string)
rcut = 3.0
sqrt_N = int(sqrt(N))

system = hoomd.init.create_lattice(
    unitcell = hoomd.lattice.sq(a = 2.0),
    n = [sqrt_N, sqrt_N]
)
nlist = md.nlist.cell(check_period = 1)
# basic LJ forces from Hoomd
lj = md.pair.lj(rcut, nlist)
lj.pair_coeff.set('A', 'A', epsilon=1.0, sigma=1.0)
md.integrate.mode_standard(dt=0.005)
md.integrate.langevin(group = group.all(), kT=1.0, seed=42)
# equilibrate for 4k steps first
hoomd.run(4000)

# %%
import pysages

pysages.set_backend("hoomd")
simulation = hoomd.context.current
simulation_view = pysages.wrap(simulation)

# %%
ξ1 = pysages.cvs.distance(np.array([0, 1]))
ξ = pysages.collective_variable(ξ1)

grid = pysages.Grid(
    lower = (0,),
    upper = (2.0,),
    shape = (1024,),
    periodicity = (True,)
)

sampling_method = pysages.methods.abf(simulation_view, grid, ξ, N = 100)

# %%
sampler = pysages.bind(simulation, sampling_method)

# %%
hoomd.run(5000)

# %%
benchmark_results = hoomd.benchmark.series(
    warmup = 6000,
    repeat = 5,
    steps = 50000,
    limit_hours = 2
)


struct LCMControlReceiver
    result::DynamicsResult{Float64, Float64}
    tprev::Base.RefValue{Float64}
    τprev::Vector{Float64}
    robot_info::HumanoidRobotInfo{Float64}
    lcm::LCM
    robot_state_channel::String
    robot_state_msg::robot_state_t
    atlas_command_msg::atlas_command_t
    encodebuffer::FastWriteBuffer
    decodebuffer::FastReadBuffer
    new_command::Base.RefValue{Bool}
end

# TODO: reduce cost of encoding

function LCMControlReceiver(robot_info::HumanoidRobotInfo;
        robot_state_channel::String = "EST_ROBOT_STATE",
        robot_command_channel::String = "ROBOT_COMMAND")
    mechanism = robot_info.mechanism
    result = DynamicsResult(mechanism)
    tprev = Ref(0.)
    τprev = zeros(num_velocities(mechanism))
    lcm = LCM()
    robot_state_msg = robot_state_t()
    atlas_command_msg = atlas_command_t()
    encodebuffer = FastWriteBuffer()
    decodebuffer = FastReadBuffer()
    receiver = LCMControlReceiver(result, tprev, τprev, robot_info, lcm, robot_state_channel,
        robot_state_msg, atlas_command_msg, encodebuffer, decodebuffer, Ref(false))

    for joint in tree_joints(mechanism)
        num_velocities(joint) == 1 || continue
        push!(robot_state_msg.joint_name, joint.name)
    end
    robot_state_msg.num_joints = length(robot_state_msg.joint_name)
    resize!(robot_state_msg)

    let receiver = receiver
        sub = subscribe(lcm, robot_command_channel, (channel, data) -> handle_robot_command_msg(receiver, data))
        set_queue_capacity(sub, 1)
    end

    receiver
end

function initialize!(receiver::LCMControlReceiver)
    receiver.tprev[] = 0
    receiver.τprev .= 0
    nothing
end

function handle_robot_command_msg(receiver::LCMControlReceiver, data::AbstractVector{UInt8})
    setdata!(receiver.decodebuffer, data)
    msg = receiver.atlas_command_msg
    decode!(msg, receiver.decodebuffer)
    receiver.new_command[] = true
end

function publish_robot_state(receiver::LCMControlReceiver, t::Number, state::MechanismState)
    set!(receiver.robot_state_msg, receiver.result, receiver.robot_info, receiver.τprev, t, state)
    encode(receiver.encodebuffer, receiver.robot_state_msg)
    bytes = take!(receiver.encodebuffer)
    publish(receiver.lcm, receiver.robot_state_channel, bytes)
end

struct NoCommandError <: Exception end
Base.showerror(io::IO, e::NoCommandError) = print(io, "Didn't receive a command message.")

function (receiver::LCMControlReceiver)(τ::AbstractVector, t::Number, state::MechanismState)
    # send state info on first tick to get things started
    if t == 0 # TODO: consider creating a separate flag in the controller for doing this
        publish_robot_state(receiver, t, state)
        sleep(0.5); publish_robot_state(receiver, t, state) # send a second time to work around a bug in the controller
    end

    # process command
    LCMCore.handle(receiver.lcm, Second(1))
    receiver.new_command[] || throw(NoCommandError())
    set!(τ, receiver.tprev[] - t, state, receiver.τprev, receiver.atlas_command_msg, receiver.robot_info)
    receiver.new_command[] = false
    receiver.tprev[] = t
    receiver.τprev .= τ

    # send state info
    publish_robot_state(receiver, t, state)
    τ
end

function RigidBodySim.PeriodicController(Δt::Number, receiver::LCMControlReceiver) # TODO: switch order?
    initialize = let receiver = receiver
        (c, t, u, integrator) -> initialize!(receiver)
    end
    PeriodicController(zero(receiver.τprev), Δt, receiver; initialize = initialize)
end


abstract type AbstractReservesFormulation <: AbstractServiceFormulation end
struct RangeReserve <: AbstractReservesFormulation end
struct StepwiseCostReserve <: AbstractReservesFormulation end
############################### Reserve Variables` #########################################
"""
This function add the variables for reserves to the model
"""
function AddVariableSpec(
    ::Type{ActiveServiceVariable},
    ::PSIContainer,
    service::T,
) where {T <: PSY.Reserve}
    return AddVariableSpec(;
        variable_name = make_variable_name(PSY.get_name(service), T),
        binary = false,
        lb_value_func = x -> 0,
        devices_filter_func = x -> PSY.get_available(x),
    )
end

function AddVariableSpec(
    ::Type{T},
    ::Type{U},
    ::PSIContainer,
) where {T <: ServiceRequirementVariable, U <: PSY.ReserveDemandCurve}
    return AddVariableSpec(;
        variable_name = make_variable_name(T, U),
        binary = false,
        lb_value_func = x -> 0.0,
    )
end

################################## Reserve Requirement Constraint ##########################
function service_requirement_constraint!(
    psi_container::PSIContainer,
    service::SR,
    ::ServiceModel{SR, T},
) where {SR <: PSY.Reserve, T <: AbstractReservesFormulation}
    parameters = model_has_parameters(psi_container)
    use_forecast_data = model_uses_forecasts(psi_container)
    initial_time = model_initial_time(psi_container)
    @debug initial_time
    time_steps = model_time_steps(psi_container)
    name = PSY.get_name(service)
    constraint = get_constraint(psi_container, make_constraint_name(REQUIREMENT, SR))
    reserve_variable = get_variable(psi_container, name, SR)
    use_slacks = get_services_slack_variables(psi_container.settings)

    ts_vector = get_time_series(psi_container, service, "requirement")

    use_slacks && (slack_vars = reserve_slacks(psi_container, name))

    requirement = PSY.get_requirement(service)
    if parameters
        param = get_parameter_array(
            psi_container,
            UpdateRef{SR}(SERVICE_REQUIREMENT, "requirement"),
        )
        for t in time_steps
            param[name, t] = PJ.add_parameter(psi_container.JuMPmodel, ts_vector[t])
            if use_slacks
                resource_expression = sum(reserve_variable[:, t]) + slack_vars[t]
            else
                resource_expression = sum(reserve_variable[:, t])
            end
            constraint[name, t] = JuMP.@constraint(
                psi_container.JuMPmodel,
                resource_expression >= param[name, t] * requirement
            )
        end
    else
        for t in time_steps
            constraint[name, t] = JuMP.@constraint(
                psi_container.JuMPmodel,
                sum(reserve_variable[:, t]) >= ts_vector[t] * requirement
            )
        end
    end
    return
end

function cost_function!(
    psi_container::PSIContainer,
    service::SR,
    ::ServiceModel{SR, T},
) where {SR <: PSY.Reserve, T <: AbstractReservesFormulation}
    reserve = get_variable(psi_container, PSY.get_name(service), SR)
    for r in reserve
        JuMP.add_to_expression!(psi_container.cost_function, r, 1.0)
    end
    return
end

function service_requirement_constraint!(
    psi_container::PSIContainer,
    service::SR,
    ::ServiceModel{SR, StepwiseCostReserve},
) where {SR <: PSY.Reserve}
    initial_time = model_initial_time(psi_container)
    @debug initial_time
    time_steps = model_time_steps(psi_container)
    name = PSY.get_name(service)
    constraint = get_constraint(psi_container, make_constraint_name(REQUIREMENT, SR))
    reserve_variable = get_variable(psi_container, name, SR)
    requirement_variable = get_variable(psi_container, SERVICE_REQUIREMENT, SR)

    for t in time_steps
        constraint[name, t] = JuMP.@constraint(
            psi_container.JuMPmodel,
            sum(reserve_variable[:, t]) >= requirement_variable[name, t]
        )
    end

    return
end

function _convert_to_variablecost(val::TS.TimeArray)
    cost_col = Vector{Symbol}()
    load_col = Vector{Symbol}()
    variable_costs = Vector{Array{NTuple{2, Float64}}}()
    col_names = TS.colnames(val)
    for c in col_names
        if occursin("cost_bp", String(c))
            push!(cost_col, c)
        elseif occursin("load_bp", String(c))
            push!(load_col, c)
        end
    end
    for row in DataFrames.eachrow(DataFrames.DataFrame(val))
        push!(variable_costs, [Tuple(row[[c, l]]) for (c, l) in zip(cost_col, load_col)])
    end
    return variable_costs
end

function cost_function!(
    psi_container::PSIContainer,
    service::SR,
    ::Type{StepwiseCostReserve},
) where {SR <: PSY.Reserve}
    use_forecast_data = model_uses_forecasts(psi_container)
    initial_time = model_initial_time(psi_container)
    @debug initial_time
    time_steps = model_time_steps(psi_container)

    function pwl_reserve_cost(
        psi_container::PSIContainer,
        variable::JV,
        cost_component::Vector{NTuple{2, Float64}},
    ) where {JV <: JuMP.AbstractVariableRef}
        return _pwlgencost_sos!(psi_container, variable, cost_component)
    end

    if use_forecast_data
        ts_vector = _convert_to_variablecost(PSY.get_data(PSY.get_forecast(
            PSY.PiecewiseFunction,
            service,
            initial_time,
            "variable",
            length(time_steps),
        )))

    else
        ts_vector =
            repeat(PSY.get_variable(PSY.get_operation_cost(service)), time_steps[end])
    end

    resolution = model_resolution(psi_container)
    dt = Dates.value(Dates.Second(resolution)) / SECONDS_IN_HOUR
    variable = get_variable(psi_container, SERVICE_REQUIREMENT, SR)
    gen_cost = JuMP.GenericAffExpr{Float64, _variable_type(psi_container)}()
    time_steps = model_time_steps(psi_container)
    name = PSY.get_name(service)
    container = add_var_container!(
        psi_container,
        make_variable_name("$(name)_pwl_cost_vars", SR),
        [name],
        time_steps,
        1:length(ts_vector[1]);
        sparse = true,
    )
    for (t, var) in enumerate(variable[name, :])
        c, pwlvars = pwl_reserve_cost(psi_container, var, ts_vector[t])
        for (ix, v) in enumerate(pwlvars)
            container[(name, t, ix)] = v
        end
        JuMP.add_to_expression!(gen_cost, c)
    end

    cost_expression = gen_cost * dt
    T_ce = typeof(cost_expression)
    T_cf = typeof(psi_container.cost_function)
    if T_cf <: JuMP.GenericAffExpr && T_ce <: JuMP.GenericQuadExpr
        psi_container.cost_function += cost_expression
    else
        JuMP.add_to_expression!(psi_container.cost_function, cost_expression)
    end
    return
end

function modify_device_model!(
    devices_template::Dict{Symbol, DeviceModel},
    service_model::ServiceModel{<:PSY.Reserve, <:AbstractReservesFormulation},
    contributing_devices::Vector{<:PSY.Device},
)
    device_types = unique(typeof.(contributing_devices))
    for dt in device_types
        for (device_model_name, device_model) in devices_template
            # add message here when it exists
            device_model.device_type != dt && continue
            service_model in device_model.services && continue
            push!(device_model.services, service_model)
        end
    end

    return
end

function include_service!(
    constraint_info::T,
    services,
    ::ServiceModel{SR, <:AbstractReservesFormulation},
) where {
    T <: Union{AbstractRangeConstraintInfo, AbstractRampConstraintInfo},
    SR <: PSY.Reserve{PSY.ReserveUp},
}
    for (ix, service) in enumerate(services)
        push!(
            constraint_info.additional_terms_ub,
            make_constraint_name(PSY.get_name(service), SR),
        )
    end
    return
end

function include_service!(
    constraint_info::T,
    services,
    ::ServiceModel{SR, <:AbstractReservesFormulation},
) where {
    T <: Union{AbstractRangeConstraintInfo, AbstractRampConstraintInfo},
    SR <: PSY.Reserve{PSY.ReserveDown},
}
    for (ix, service) in enumerate(services)
        push!(
            constraint_info.additional_terms_lb,
            make_constraint_name(PSY.get_name(service), SR),
        )
    end
    return
end

function add_device_services!(
    constraint_info::T,
    device::D,
    model::DeviceModel,
) where {
    T <: Union{AbstractRangeConstraintInfo, AbstractRampConstraintInfo},
    D <: PSY.Device,
}
    for service_model in get_services(model)
        if PSY.has_service(device, service_model.service_type)
            services =
                (s for s in PSY.get_services(device) if isa(s, service_model.service_type))
            @assert !isempty(services)
            include_service!(constraint_info, services, service_model)
        end
    end
    return
end

function add_device_services!(
    constraint_data_in::AbstractRangeConstraintInfo,
    constraint_data_out::AbstractRangeConstraintInfo,
    device::D,
    model::DeviceModel{D, <:AbstractStorageFormulation},
) where {D <: PSY.Storage}
    for service_model in get_services(model)
        if PSY.has_service(device, service_model.service_type)
            services =
                (s for s in PSY.get_services(device) if isa(s, service_model.service_type))
            @assert !isempty(services)
            if service_model.service_type <: PSY.Reserve{PSY.ReserveDown}
                for service in services
                    push!(
                        constraint_data_in.additional_terms_ub,
                        make_constraint_name(
                            PSY.get_name(service),
                            service_model.service_type,
                        ),
                    )
                end
            elseif service_model.service_type <: PSY.Reserve{PSY.ReserveUp}
                for service in services
                    push!(
                        constraint_data_out.additional_terms_ub,
                        make_constraint_name(
                            PSY.get_name(service),
                            service_model.service_type,
                        ),
                    )
                end
            end
        end
    end
    return
end


using AbstractNumbers, SpecialFunctions
using Test, Aqua

Aqua.test_ambiguities([AbstractNumbers, Base, Core])
Aqua.test_unbound_args(AbstractNumbers)
Aqua.test_undefined_exports(AbstractNumbers)
Aqua.test_project_extras(AbstractNumbers)
Aqua.test_stale_deps(AbstractNumbers)
Aqua.test_deps_compat(AbstractNumbers)
Aqua.test_project_toml_formatting(AbstractNumbers)

struct MyNumber{T} <: AbstractNumbers.AbstractNumber{T}
    number::T
end
const SF = SpecialFunctions

AbstractNumbers.number(x::MyNumber) = x.number
AbstractNumbers.basetype(::Type{<: MyNumber}) = MyNumber

x = MyNumber(1)
@test convert(Number, x) === x

all_funcs = (
    :~, :conj, :abs, :sin, :cos, :tan, :sinh, :cosh, :tanh, :asin, :acos, :atan,
    :asinh, :acosh, :atanh, :sec, :csc, :cot, :asec, :acsc, :acot, :sech, :csch,
    :coth, :asech, :acsch, :acoth, :sinc, :cosc, :cosd, :cotd, :cscd, :secd,
    :sind, :tand, :acosd, :acotd, :acscd, :asecd, :asind, :atand, :rad2deg,
    :deg2rad, :log, :log2, :log10, :log1p, :exponent, :exp, :exp2, :expm1,
    :cbrt, :sqrt, :ceil, :floor, :trunc, :round, :significand,
    :frexp, :ldexp, :modf, :real, :imag, :!, :identity,
    :zero, :one, :<<, :>>, :abs2, :sign, :sinpi, :cospi, :exp10,
    :iseven, :ispow2, :isfinite, :isinf, :isodd, :isinteger, :isreal,
    :isnan, :isempty, :iszero, :transpose, :copysign, :flipsign, :signbit,
    :+, :-, :*, :/, :\, :^, :(==), :(!=), :<, :(<=), :>, :(>=), :min, :max,
    :div, :fld, :rem, :mod, :mod1, :cmp, :&, :|, :xor,
    :clamp
)

special_funcs = (
    :airyai,
    :airyaiprime,
    :airybi,
    :airybiprime,
    :airyaix,
    :airyaiprimex,
    :airybix,
    :airybiprimex,
    :besselh,
    :besselhx,
    :besseli,
    :besselix,
    :besselj,
    :besselj0,
    :besselj1,
    :besseljx,
    :besselk,
    :besselkx,
    :bessely,
    :bessely0,
    :bessely1,
    :besselyx,
    :dawson,
    :erf,
    :erfc,
    :erfcinv,
    :erfcx,
    :erfi,
    :erfinv,
    :eta,
    :digamma,
    :invdigamma,
    :polygamma,
    :trigamma,
    :hankelh1,
    :hankelh1x,
    :hankelh2,
    :hankelh2x,
    :zeta
)

nancompare(a, b) = a == b
function nancompare(a::Number, b::Number)
    isnan(a) && isnan(b) && return true
    a == b
end

function myrand(::Type{T}) where T <: Integer
    x = rand(T(-10):T(10))
    x == 0 ? T(1) : x
end
myrand(::Type{T}) where T <: Unsigned = rand(T(1):T(20))
myrand(::Type{T}) where T = rand(T)
@testset "AbstractNumbers" begin
    for (mod, funcs) in ((Base, all_funcs),), f in funcs
        println(f)
        func = getfield(mod, f)
        @testset "testing $f" begin
            for i = 1:4
                for T in (Float32, Float64, ComplexF64, Int, UInt)
                    args = ntuple(x-> myrand(T), i)
                    if i == 3 && func === cmp
                        continue
                    end
                    if T <: Complex
                        (func in (
                            cotd, cosd, cscd, secd, sind, tand, acosd, acotd, acscd,
                            asecd, asind, atand, rem, modf, (<), (>), (<=), (>=),
                            min, max, cmp, &, |, xor, clamp, div, fld
                        )) && continue
                    end
                    if T <: AbstractFloat
                        (func in (
                            &, |, xor
                        )) && continue
                    end
                    # i have no idea what these functions are - seem to throw exceptions when not careful with input values
                    # also, they're not usable since they're deprecated, even if you get them from specialfunctions
                    if applicable(func, args...)
                        res = try
                            a = func(args...)
                            b = func(MyNumber.(args)...)
                            @testset "$T $args" begin
                                @test nancompare(a, b)
                            end
                        catch e
                            isa(e, DomainError) || rethrow(e)
                        end
                    end
                end
            end
        end
    end

    # isapprox
    for T in (Float32, Float64, ComplexF64, Int, UInt)
        T = Float64
        args = ntuple(x-> myrand(T), 2)
        a = ≈(args...)
        b = ≈(MyNumber.(args)...)
        a == b
    end

    @testset "bessel functions" begin
        bessel_funcs = [(SF.bessely0, SF.bessely1, SF.bessely), (SF.besselj0, SF.besselj1, SF.besselj)]
        @testset "$z, $o" for (z, o, f) in bessel_funcs
            @test z(Float32(2.0)) ≈ z(Float64(2.0))
            @test o(Float32(2.0)) ≈ o(Float64(2.0))
            @test z(MyNumber(2)) ≈ z(MyNumber(2.0))
            @test o(MyNumber(2)) ≈ o(MyNumber(2.0))
            @test z(MyNumber(2.0 + im)) ≈ f(MyNumber(0), MyNumber(2.0 + im))
            @test o(MyNumber(2.0 + im)) ≈ f(MyNumber(1), MyNumber(2.0 + im))
        end
        @testset "besselj error throwing" begin
            @test_throws MethodError SF.besselj(MyNumber(1.2), MyNumber(big(1.0)))
            @test_throws MethodError SF.besselj(MyNumber(1), MyNumber(complex(big(1.0))))
            @test_throws MethodError SF.besseljx(MyNumber(1), MyNumber(big(1.0)))
            @test_throws MethodError SF.besseljx(MyNumber(1), MyNumber(complex(big(1.0))))
        end
    end

end

nothing


import Base:in, +, -, *, ==, convert, isempty, isequal, length, push!, sizeof, append!

export BeamData

const beamfields = [:name, :bp_len,:bp_step, :f, :c, :θ]

"""
   BeamData
A structure for Beam power of ambient noise data.
## Fields: BeamData
| **Field** | **Description** |
|:--------------------|:--------------------|
| :name       | name of input array |
| :misc         | misc dictionary. |
| :notes        | notes |
| :bp_len         | time window length to compute beam power |
| :bp_step       | time window step to compute beam power  |
| :f    | examined frequency [1/s] |
| :c    | examined phase velocity [m/s] |
| :θ    | examined aximuth of slowness vector [rad] |
| :t    | time stamp of each time window |
| :bp   | beam power b(f, t, θ, c)|
| :tstack  | time stamp of stacked data. |
| :tstack_len       | time window length of stacked beam power. |
| :bpstack      | stacked beam power. |
"""
mutable struct BeamData

  name     ::String                          # name of input file
  misc     ::Dict{String,Any}                # misc
  notes    ::Array{String,1}                 # notes
  bp_len   ::Int                             # beam power time window length [s]
  bp_step  ::Int                             # beam power time window step   [s]
  f::Array{Float64,1}                         # examined frequency [1/s]
  c::Array{Float64,1}                        # examined wave velocity
  θ::Array{Float64,1}                        # examined azimuth
  t::Array{Float64,1}                         # time
  bp::Array{<:Union{Complex{Float32},Complex{Float64}},3}  # beam power data
  tstack::Float64                             #stacked time
  tstack_len::Int                             #stacked time length
  bpstack::Array{Complex{Float32},2}          #stacked beam power

  function BeamData(
      name     ::String,
      misc     ::Dict{String,Any},
      notes    ::Array{String,1},
      bp_len   ::Int,
      bp_step  ::Int,
      f        ::Array{Float64,1},
      c        ::Array{Float64,1},
      θ        ::Array{Float64,1},
      t        ::Array{Float64,1},
      bp       ::Array{<:Union{Complex{Float32},Complex{Float64}},3},
      tstack   ::Float64,
      tstack_len::Int,
      bpstack  ::Array{Complex{Float32},2}
      )

      return new(name, misc, notes, bp_len, bp_step, f, c, θ, t, bp, tstack, tstack_len, bpstack)
    end
end

BeamData(;
          name     ::String                    = "",
          misc     ::Dict{String,Any}          = Dict{String,Any}(),
          notes    ::Array{String,1}           = Array{String,1}(undef, 0),
          bp_len   ::Int                       = zero(Int),
          bp_step  ::Int                       = zero(Int),
          f        ::Array{Float64,1}          = Array{Float64,1}(undef, 0),
          c        ::Array{Float64,1}          = Array{Float64,1}(undef, 0),
          θ        ::Array{Float64,1}          = Array{Float64,1}(undef, 0),
          t        ::Array{Float64,1}          = Array{Float64,1}(undef, 0),
          bp       ::Array{<:Union{Complex{Float32},Complex{Float64}},3} =
                     Array{Complex{Float32},3}(undef, 0, 0, 0),
          tstack    ::Float64                   = zero(Float64),
          tstack_len::Int                       = zero(Int),
          bpstack  ::Array{Complex{Float32},2}  = Array{ComplexF32,2}(undef, 0, 2)

          ) = BeamData(name, misc, notes, bp_len, bp_step, f, c, θ, t, bp, tstack, tstack_len, bpstack)

in(s::String, BP::BeamData) = BP.name==s

isempty(B::BeamData) = min(B.name=="",
                           B.bp_len == zero(typeof(B.bp_len)),
                           B.bp_step == zero(typeof(B.bp_step)),
                           B.tstack == zero(typeof(B.tstack)),
                           B.tstack_len == zero(typeof(B.tstack_len)),
                           isempty(B.f),isempty(B.c),
                           isempty(B.θ),isempty(B.t),
                           isempty(B.bp),isempty(B.bpstack),
                           isempty(B.misc),isempty(B.notes))

isequal(B::BeamData, U::BeamData) = minimum([hash(getfield(F,i))==hash(getfield(U,i)) for i in beamfields]::Array{Bool,1})
==(B::BeamData, U::BeamData) = isequal(B,U)::Bool

append!(B::BeamData, U::BeamData)  = (if B == U;
  setfield!(B, :t, vcat(getfield(B,:t), getfield(U,:t)));
  setfield!(B, :bp, hcat(getfield(B,:bp), getfield(U,:bp)))
  elseif isempty(B);
  [setfield!(B, i, getfield(U,i)) for i in beamfields];
  [setfield!(B, i, getfield(U,i)) for i in [:t,:bp]];
  end;
  return B)

+(B::BeamData, U::BeamData) = (T = deepcopy(B); return append!(T, U))

function sizeof(B::BeamData)
s = sum([sizeof(getfield(B,f)) for f in beamfields])
if !isempty(B.notes)
  s += sum([sizeof(i) for i in B.notes])
end
if !isempty(B.misc)
  s += sum([sizeof(i) for i in values(B.misc)])
end
return s
end


macro rand(exp)
    rand_macro(exp)
end

function rand_macro(ex)
    whereparams = []
    ex isa Expr && ex.head ∈ (:(=), :function, :->) ||
        throw(ArgumentError("@rand requires an expression defining `rand`"))
    sig = ex.args[1]
    body = ex.args[2]

    if sig.head == :where
        append!(whereparams, sig.args[2:end])
        sig = sig.args[1]
    end

    if ex.head == :function && sig.head == :tuple # anonymous function
        sig = Expr(:call, :rand, sig.args...)
    end
    if ex.head == :->
        # TODO: check that only one argument is passed
        sig = Expr(:call, :rand, sig)
    end

    sig.head == :call &&
        sig.args[1] == :rand || throw(ArgumentError(
            "@rand requires a function expression defining `rand`"))

    argname = sig.args[2].args[1] # x
    namefull = sig.args[2] # x::X
    @assert namefull.head == :(::) # TODO: throw exception

    # sub-samplers; second argument true forces Val(Inf) for the sampler
    sps = Pair{<:Any,Bool}[]

    rng = gensym()
    body = samplerize!(sps, body, argname, rng)
    istrivial = isempty(sps)

    exsig = Expr(:call,
                 :(Random.rand),
                 :($(esc(rng))::AbstractRNG),
                 esc(as_sampler(namefull, istrivial)))

    if !isempty(whereparams)
        exsig = Expr(:where, exsig, map(esc, whereparams)...)
    end

    ex = Expr(:function, exsig, esc(body))

    sp = if istrivial
        # we explicitly define Sampler even in the trivial case to handle
        # redefinitions, where the old rand/sampler pair (for SamplerSimple)
        # is overwritten by a new one (for SamplerTrivial)
        quote
            Random.Sampler(::Type{RNG}, n::Repetition) where {RNG<:AbstractRNG} =
                SamplerTrivial($(esc(argname)))
        end
    else
        quote
            Random.Sampler(::Type{RNG}, n::Repetition) where {RNG<:AbstractRNG} =
                SamplerSimple($(esc(argname)), tuple(SP))
        end
    end
    @assert sp.args[2].args[1].head == :where
    append!(sp.args[2].args[1].args, map(esc, whereparams))

    # insert x::X in the argument list, between RNG and n::Repetition
    insert!(sp.args[2].args[1].args[1].args, 3, esc(namefull))

    # insert inner samplers
    if !istrivial
        SP = [Expr(:call, :Sampler, :RNG, esc(x), many ? Val{Inf}() : :n) for (x, many) in sps]
        @assert :SP == pop!(sp.args[2].args[2].args[2].args[3].args)
        append!(sp.args[2].args[2].args[2].args[3].args, SP)
    end

    quote
        $ex
        $(sp.args[2]) # unwrap the quote/block around the definition
    end
end

function as_sampler(ex, istrivial)
    t = istrivial ? :(RandomExtensions.SamplerTrivial) : :(RandomExtensions.SamplerSimple)
    Expr(:(::),
         ex.args[1],
         Expr(:curly, t,
              Expr(:(<:), ex.args[2])))
end

function samplerize!(sps, ex, name, rng)
    if ex == name
        # not within a rand() call
        return Expr(:ref, name) # name -> name[]
    end
    ex isa Expr || return ex
    if ex.head == :call && ex.args[1] == :rand
        # we assume that if rand has more than one arg, we want
        # a Val(Inf) sampler (e.g. rand(1:9, 2, 3)
        push!(sps, ex.args[2] => length(ex.args) > 2)
        i = length(sps)
        Expr(:call, :rand, rng, :($name.data[$i]), ex.args[3:end]...)
    else
        Expr(ex.head, map(e -> samplerize!(sps, e, name, rng), ex.args)...)
    end
end


export RemoteEstObj,
    LocalEstObj,
    EstimationWrapper,
    update_reo!,
    reset_reo!,
    theta0, theta1, hess, invhess, grad,
    getworkers,
    parallel_simloglik!,
    update!,
    start_up_workers

"workers() but excluding master"
getworkers() = filter(i -> i != 1, workers())

function start_up_workers(ENV::Base.EnvDict; nprocs = Sys.CPU_THREADS)
    # send primitives to workers
    oldworkers = getworkers()
    println_time_flush("removing workers $oldworkers")
    rmprocs(oldworkers)
    flush(stdout)
    if "SLURM_JOBID" in keys(ENV)
        num_cpus_to_request = parse(Int, ENV["SLURM_TASKS_PER_NODE"])
        println("requesting $(num_cpus_to_request) cpus from slurm.")
        flush(stdout)
        pids = addprocs_slurm(num_cpus_to_request)
    else
        cputhrds = Sys.CPU_THREADS
        cputhrds < nprocs && @warn "using nprocs = $cputhrds < $nprocs specified"
        pids = addprocs(min(nprocs, cputhrds))
    end
    println_time_flush("Workers added: $pids")
    return pids
end

# ----------------------------------

@GenGlobal g_RemoteEstObj

# ----------------------------------

abstract type AbstractEstimationObjects end

struct RemoteEstObj{D<:DataSetofSets} <: AbstractEstimationObjects
    data::D
    llvec::SharedVector{Float64}
    scoremat::SharedMatrix{Float64}
    sim::SimulationDrawsMatrix{Float64,Matrix{Float64}}

    function RemoteEstObj(data::D,llvec,scoremat,sim) where {D}
        k, n = size(scoremat)
        n == num_i(data) == length(llvec) == num_i(sim) || throw(DimensionMismatch("num_i"))
        k == _nparm(data) || throw(DimensionMismatch("k"))
        _nparm(drill(data)) == _nparm(sim) || throw(DimensionMismatch("sim k"))
        return new{D}(data, llvec, scoremat, sim)
    end
end

struct LocalEstObj{D<:DataSetofSets} <: AbstractEstimationObjects
    data::D
    theta0::Vector{Float64}
    theta1::Vector{Float64}
    grad::Vector{Float64}
    hess::Matrix{Float64}
    invhess::Matrix{Float64}

    function LocalEstObj(data::D,theta0,theta1,grad,hess,invhess) where {D}
        k = length(theta0)
        k == length(theta1) == length(grad) || throw(DimensionMismatch())
        k == checksquare(hess) == checksquare(invhess) || throw(DimensionMismatch())
        k == _nparm(data) || throw(DimensionMismatch())

        return new{D}(data,theta0,theta1,grad,hess,invhess)
    end
end

data(x::AbstractEstimationObjects) = x.data

llvec(x::RemoteEstObj) = x.llvec
scoremat(x::RemoteEstObj) = x.scoremat
sim(x::RemoteEstObj) = x.sim
size(x::RemoteEstObj) = size(scoremat(x))
num_i(x::RemoteEstObj) = length(llvec(eo))
_nparm(x::RemoteEstObj) = size(scoremat(x),1)
LL(x::RemoteEstObj) = sum(llvec(x))
negLL(x::RemoteEstObj) = -LL(x)
loglikelihood(x::RemoteEstObj) = LL(x)

theta0(x::LocalEstObj) = x.theta0
theta1(x::LocalEstObj) = x.theta1
grad(  x::LocalEstObj) = x.grad
hess(  x::LocalEstObj) = x.hess
invhess(x::LocalEstObj) = x.invhess
num_i( x::LocalEstObj) = num_i(data(x))
_nparm(x::LocalEstObj) = length(theta0(x))

function reset!(x::LocalEstObj)
    fill!(grad(x), 0)
    fill!(hess(x), 0)
    fill!(invhess(x),0)
    fill!(theta1(x),0)
end

function update!(data::DataSetofSets, theta::Vector)
    _nparm(data) == length(theta) || throw(DimensionMismatch())
    thetas = split_thetas(data, theta)
    for (d, θ) in zip(data, thetas)
        update!(d, θ)
    end
end

@noinline function update_reo!(reo::RemoteEstObj, theta::Vector)
    update!(data(reo), theta)
    return nothing
end

update_reo!(theta::Vector) = update_reo!(get_g_RemoteEstObj(), theta)

@noinline function reset_reo!(reo::RemoteEstObj)
    data_drill = drill(data(reo))
    ddm = _model(data_drill)
    vf = value_function(ddm)
    fill!(vf, 0)
end

reset_reo!() = reset_reo!(get_g_RemoteEstObj())

function invhess!(x::LocalEstObj)
    invhess(x) .= inv(hess(x))
    return invhess(x)
end

function RemoteEstObj(leo::LocalEstObj, M, pids = workers())
    n = num_i(leo)
    k = _nparm(leo)
    llvec    = SharedVector{Float64}(n;    pids=pids)
    scoremat = SharedMatrix{Float64}(k, n; pids=pids)
    sim      = SimulationDraws(M, data(leo))
    return RemoteEstObj(data(leo), llvec, scoremat, sim)
end

function LocalEstObj(data, theta)
    k = length(theta)
    k == _nparm(data) || throw(DimensionMismatch("length(theta) = $(length(theta)) but _nparm(data) = $(_nparm(data)) θ = $theta"))

    theta0 = Vector{Float64}(undef, k)
    theta1 = similar(theta0)
    grad   = similar(theta0)
    hess   = Matrix{Float64}(undef, k, k)
    invhess = similar(hess)
    return LocalEstObj(data, theta0, theta1, grad, hess, invhess)
end

struct EstimationWrapper{D} <: AbstractEstimationObjects
    leo::LocalEstObj{D}
    reo::RemoteEstObj{D}
    function EstimationWrapper(leo::LocalEstObj{D},reo::RemoteEstObj{D}) where {D}
        @assert data(leo) === data(reo)
        return new{D}(leo,reo)
    end
end

LocalEstObj(ew::EstimationWrapper) = ew.leo
RemoteEstObj(ew::EstimationWrapper) = ew.reo
data(ew::EstimationWrapper) = data(RemoteEstObj(ew))
OneTo(ew::EstimationWrapper) = OneTo(num_i(LocalEstObj(ew)))
_nparm(ew::EstimationWrapper) = _nparm(LocalEstObj(ew))

check_theta(ew::EstimationWrapper, theta) = check_theta(data(ew),theta)

function check_theta(data::DataSetofSets, theta)
    thetasvw = split_thetas(data, theta)
    for (dat,θ) in zip(data,thetasvw)
        update!(dat, θ)
    end
end

function update!(ew::EstimationWrapper, theta, dograd)
    l = LocalEstObj(ew)
    r = RemoteEstObj(ew)
    theta1(l) .= theta
    if dograd
        score = scoremat(r)
        h = hess(l)
        g = grad(l)
        mul!(h, score, score')
        sum!(reshape(g, :, 1), score)
        # println("gradient = $g")
        g .*= -1
    end
    nll = negLL(r)
    if dograd
        countplus!(-nll,theta,g)
    else
        countplus!(-nll,theta)
    end
    flush(stdout)
    return nll
end

@noinline function simloglik!(i, theta, dograd, reo::RemoteEstObj; kwargs...)

    check_finite(theta)

    gradi = view(scoremat(reo), :, i)
    grptup = getindex.(data(reo), i)
    simi = view(sim(reo), i)
    θρ = theta_ρ(data(reo), theta)

    fill!(_qm(sim(reo)), 0)
    fill!(gradi, 0)

    # do updates
    @assert isfinite(θρ)
    update!(simi, θρ)

    thetasvw = split_thetas(data(reo), theta)
    idxs = theta_indexes(data(reo))

    llv = llvec(reo)
    llv[i] = simloglik!(gradi, grptup, thetasvw, idxs, simi, dograd; kwargs...)

end

simloglik!(i, theta, dograd; kwargs...) =
    simloglik!(i, theta, dograd, get_g_RemoteEstObj(); kwargs...)

function serial_simloglik!(ew, theta, dograd; n::Integer=num_i(LocalEstObj(ew)), kwargs...)
    check_theta(ew,theta)
    reo = RemoteEstObj(ew)
    update_reo!(reo, theta)
    map(i -> simloglik!(i, theta, dograd, reo; kwargs...), OneTo(n))
    return update!(ew, theta, dograd)
end

function parallel_simloglik!(ew, theta, dograd; kwargs...)
    check_finite(theta)
    check_theta(ew,theta)
    wp = CachingPool(workers())
    @eval @everywhere update_reo!($theta)
    let x=theta, dograd=dograd, kwargs=kwargs
        pmap(i -> simloglik!(i, x, dograd; kwargs...), wp, OneTo(ew))
    end
    return update!(ew, theta, dograd)
end

function test_parallel_simloglik!(ew, theta, dograd; kwargs...)
    @everywhere reset_reo!()
    parallel_simloglik!(ew, theta, dograd; kwargs...)
end

function test_serial_simloglik!(ew, theta, dograd; kwargs...)
    reset_reo!(RemoteEstObj(ew))
    serial_simloglik!(ew, theta, dograd; kwargs...)
end


using .JSort

using .JSort
import Base: read, read!, fill!, write, close

mutable struct QuickReader <: AbstractReader
    parameters::Parameters
    input::CBuffer
    labr_size::Int
    initialtime::Int
    eγ::Vector{Vector{Float32}}
    source::String
end

function QuickReader(parameters::Parameters)
    j(x) = joinpath(parameters.savepath, x)
    source = joinpath(parameters.readpath, "raw.bin")
    if ~isfile(source)
        error("The file $source does not exist")
    end
    buffer = CBuffer(source)
    egs = [Float32[] for i in 1:33]

    QuickReader(parameters, buffer, 0, -1, egs, source)
end

read(reader::QuickReader, gate::Gate) = read(reader, gate=gate)
read(reader::QuickReader, gate::Gate2D) = read(reader, pgate=gate)

function read(reader::QuickReader; gate::Gate=Gate(), pgate=nothing)
    close(reader.input)
    reader.input = CBuffer(reader.source)

    qkinzpoly = reader.parameters.qkinz
    timecalib = reader.parameters.calibrator[:t]
    timecorrection = reader.parameters.calibrator[:tc]
    gammacalib = reader.parameters.calibrator[:γ]
    ecalib = reader.parameters.calibrator[:e]
    Δecalib = reader.parameters.calibrator[:Δe]
    prompt = reader.parameters.gates[:prompt]
    background = reader.parameters.gates[:background]
    event = Event()
    read_events = front = back = 0
    # TODO Same type as calibrators
    # TODO Buffer2D are not overwritten
    e = Δe = eΔe = Float32(0.0)
    eₓᵗ = eᵧ = Float32(0.0)
    t = tc = Float32(0.0)
    reader.eγ = [Float32[] for i in 1:33]
    for eγ in reader.eγ
        sizehint!(eγ, 2^13)
    end

    while isgood(reader)
        read_events += 1
        fill!(event, reader)
        front, back = event.front, event.back
        e = ecalib[front, back](event.e)
        Δe = Δecalib[front, back](event.de)
        if !isnothing(pgate)
            !pgate(e, Δe) && continue
        end
        # Theoretical excitation energy from kinematics
        eΔe = e + Δe
        # This is a `constant` poly. No calibration needed
        eₓᵗ = qkinzpoly[front](eΔe)
        !gate(eₓᵗ) && continue

        @views for labr in event.labr[1:reader.labr_size]
            labr.channel ∈ reader.parameters.badchannels && continue
            # Gamma detectors
            eᵧ = gammacalib[labr.channel](labr.adc)

            # Time
            if reader.initialtime == -1
                reader.initialtime = event.labr[1].tdc
            end
            t = timecalib[labr.channel](labr.tdc)
            !prompt(t) && continue
            #tc = timecorrection[labr.channel](t, eΔe)
            # Also correct using eᵧ?
            push!(reader.eγ[labr.channel], eᵧ)
        end
    end
end

function close(reader::QuickReader)
    close(reader.input)
end

function ChannelDetector(reader::QuickReader)
    data = GammaDetector[GammaDetector(reader.eγ[chn], chn) for chn in 1:33]
    ChannelDetector(data)
end

plot(r::QuickReader; kwargs...) = plot(ChannelDetector(r); kwargs...)


# TODO: FIXME for general case
# wrong for transposed matrices, and certain views/SubArrays.
unitstride(op::Operation, s) = first(op.ref.ref) === s

function cost(op::Operation, unrolled::Symbol, Wshift::Int, size_T::Int = op.elementbytes)
    isconstant(op) && return 0.0, 0, 1
    # Wshift == dependson(op, unrolled) ? Wshift : 0
    # c = first(cost(instruction(op), Wshift, size_T))::Int
    instr = instruction(op)
    opisunrolled = dependson(op, unrolled)
    srt, sl, srp = opisunrolled ? vector_cost(instr, Wshift, size_T) : scalar_cost(instr)
    if accesses_memory(op)
        # either vbroadcast/reductionstore, vmov(a/u)pd, or gather/scatter
        # @show instr, unrolled, loopdependencies(op), unitstride(op, unrolled)
        if opisunrolled
            if !unitstride(op, unrolled)# || !isdense(op) # need gather/scatter
                r = (1 << Wshift)
                srt *= r
                sl *= r
            # else # vmov(a/u)pd
            end
        elseif instr === :setindex! # broadcast or reductionstore; if store we want to penalize reduction
            srt *= 2
            sl *= 2
        end
    end
    srt, sl, srp
end

    # Base._return_type()

function biggest_type_size(ls::LoopSet)
    maximum(elsize, operations(ls))
end
function VectorizationBase.pick_vector_width(ls::LoopSet, u::Symbol)
    VectorizationBase.pick_vector_width(length(ls, u), biggest_type_size(ls))
end
function VectorizationBase.pick_vector_width_shift(ls::LoopSet, u::Symbol)
    VectorizationBase.pick_vector_width_shift(length(ls, u), biggest_type_size(ls))
end
function hasintersection(a, b)
    for aᵢ ∈ a, bᵢ ∈ b
        aᵢ === bᵢ && return true
    end
    false
end

# evaluates cost of evaluating loop in given order
# heuristically, could simplify analysis by just unrolling outer loop?
function evaluate_cost_unroll(
    ls::LoopSet, order::Vector{Symbol}, max_cost = typemax(Float64), vectorized::Symbol = first(order)
)
    # included_vars = Set{UInt}()
    included_vars = fill(false, length(operations(ls)))
    nested_loop_syms = Symbol[]#Set{Symbol}()
    total_cost = 0.0
    iter = 1.0
    # Need to check if fusion is possible
    size_T = biggest_type_size(ls)
    W, Wshift = VectorizationBase.pick_vector_width_shift(length(ls, vectorized), size_T)::Tuple{Int,Int}
    for itersym ∈ order
        # Add to set of defined symbles
        push!(nested_loop_syms, itersym)
        liter = Float64(length(ls, itersym))
        if itersym === vectorized
            liter /= W
        end
        iter *= liter
        # check which vars we can define at this level of loop nest
        for (id,op) ∈ enumerate(operations(ls))
            # won't define if already defined...
            # id = identifier(op)
            included_vars[id] && continue
            # it must also be a subset of defined symbols
            loopdependencies(op) ⊆ nested_loop_syms || continue
            # hasintersection(reduceddependencies(op), nested_loop_syms) && return Inf
            rd = reduceddependencies(op)
            hasintersection(rd, nested_loop_syms[1:end-length(rd)]) && return Inf
            included_vars[id] = true

            total_cost += iter * first(cost(op, vectorized, Wshift, size_T))
            total_cost > max_cost && return total_cost # abort if more expensive; we only want to know the cheapest
        end
    end
    total_cost
end

# only covers vectorized ops; everything else considered lifted?
function depchain_cost!(
    skip::Vector{Bool}, op::Operation, vectorized::Symbol, Wshift::Int, size_T::Int, rt::Float64 = 0.0, sl::Int = 0
)
    skip[identifier(op)] = true
    # depth first search
    for opp ∈ parents(op)
        skip[identifier(opp)] && continue
        rt, sl = depchain_cost!(skip, opp, vectorized, Wshift, size_T, rt, sl)
    end
    # Basically assuming memory and compute don't conflict, but everything else does
    # Ie, ignoring the fact that integer and floating point operations likely don't either
    if iscompute(op)
        rtᵢ, slᵢ = cost(op, vectorized, Wshift, size_T)
        rt += rtᵢ; sl += slᵢ
    end
    rt, sl
end
function parentsnotreduction(op::Operation)
    for opp ∈ parents(op)
        isreduction(opp) && return false
    end
    return true
end
function unroll_no_reductions(ls, order, vectorized, Wshift, size_T)
    innermost = last(order)
    compute_rt = 0.0
    load_rt = 0.0
    # latency not a concern, because no depchains
    for op ∈ operations(ls)
        dependson(op, innermost) || continue
        if iscompute(op)
            compute_rt += first(cost(op, vectorized, Wshift, size_T))
        elseif isload(op)
            load_rt += first(cost(op, vectorized, Wshift, size_T))
        end
    end
    # heuristic guess
    # @show compute_rt, load_rt
    min(4, round(Int, (compute_rt + load_rt + 1) / compute_rt))
end
function determine_unroll_factor(
    ls::LoopSet, order::Vector{Symbol}, unrolled::Symbol, vectorized::Symbol = first(order)
)
    size_T = biggest_type_size(ls)
    W, Wshift = VectorizationBase.pick_vector_width_shift(length(ls, vectorized), size_T)::Tuple{Int,Int}

    # The strategy is to use an unroll factor of 1, unless there appears to be loop carried dependencies (ie, num_reductions > 0)
    # The assumption here is that unrolling provides no real benefit, unless it is needed to enable OOO execution by breaking up these dependency chains
    num_reductions = 0#sum(isreduction, operations(ls))
    for op ∈ operations(ls)
        if isreduction(op) & iscompute(op) && parentsnotreduction(op)
            num_reductions += 1
        end
    end
    if iszero(num_reductions)
        # if only 1 loop, no need to unroll
        # if more than 1 loop, there is some cost. Picking 2 here as a heuristic.
        return length(order) == 1 ? 1 : unroll_no_reductions(ls, order, vectorized, Wshift, size_T)
    end
    # So if num_reductions > 0, we set the unroll factor to be high enough so that the CPU can be kept busy
    # if there are, U = max(1, round(Int, max(latency) * throughput / num_reductions)) = max(1, round(Int, latency / (recip_throughput * num_reductions)))
    # We also make sure register pressure is not too high.
    latency = 0
    compute_recip_throughput = 0.0
    visited_nodes = fill(false, length(operations(ls)))
    load_recip_throughput = 0.0
    store_recip_throughput = 0.0
    for op ∈ operations(ls)
        dependson(op, unrolled) || continue
        if isreduction(op)
            rt, sl = depchain_cost!(visited_nodes, op, vectorized, Wshift, size_T)
            latency = max(sl, latency)
            compute_recip_throughput += rt
        elseif isload(op)
            load_recip_throughput += first(cost(op, vectorized, Wshift, size_T))
        elseif isstore(op)
            store_recip_throughput += first(cost(op, vectorized, Wshift, size_T))
        end
    end
    recip_throughput = max(
        compute_recip_throughput,
        load_recip_throughput,
        store_recip_throughput
    )
    max(1, round(Int, latency / (recip_throughput * num_reductions) ) )
end

function tile_cost(X, U, T)
    X[1] + X[4] + X[2] / T + X[3] / U
end
function solve_tilesize(X, R)
    @inbounds any(iszero, (R[1],R[2],R[3])) && return -1,-1,Inf #solve_smalltilesize(X, R, Umax, Tmax)
    # @inbounds any(iszero, (R[1],R[2],R[3])) && return -1,-1,Inf #solve_smalltilesize(X, R, Umax, Tmax)
    # We use a lagrange multiplier to find floating point values for U and T
    # first solving for U via quadratic formula
    # X is vector of costs, and R is of register pressures
    RR = REGISTER_COUNT - R[3] - R[4]
    a = (R[1])^2*X[2] - (R[2])^2*R[1]*X[3]/RR
    b = 2*R[1]*R[2]*X[3]
    c = -RR*R[1]*X[3]
    Ufloat = (sqrt(b^2 - 4a*c) - b) / (2a)
    Tfloat = (RR - Ufloat*R[2])/(Ufloat*R[1])
    # @show Ufloat, Tfloat
    (isfinite(Tfloat) && isfinite(Ufloat)) || return -1,-1,Inf
    Ulow = max(1, floor(Int, Ufloat)) # must be at least 1
    Tlow = max(1, floor(Int, Tfloat)) # must be at least 1
    Uhigh = Ulow + 1 #ceil(Int, Ufloat)
    Thigh = Tlow + 1 #ceil(Int, Tfloat)

    RR = REGISTER_COUNT - R[3] - R[4]
    U, T = Ulow, Tlow
    tcost = tile_cost(X, Ulow, Tlow)
    # @show Ulow*Thigh*R[1] + Ulow*R[2]
    if RR ≥ Ulow*Thigh*R[1] + Ulow*R[2]
        tcost_temp = tile_cost(X, Ulow, Thigh)
        # @show tcost_temp, tcost
        if tcost_temp < tcost
            tcost = tcost_temp
            U, T = Ulow, Thigh
        end
    end
    # The RR + 1 is a hack to get it to favor Uhigh in more scenarios
    Tl = Tlow
    while RR < Uhigh*Tl*R[1] + Uhigh*R[2]
        Tl -= 1
    end
    tcost_temp = tile_cost(X, Uhigh, Tl)
    if tcost_temp < tcost
        tcost = tcost_temp
        U, T = Uhigh, Tl
    end
    if RR > Uhigh*Thigh*R[1] + Uhigh*R[2]
        throw("Something went wrong when solving for Tfloat and Ufloat.")
    end
    U, T, tcost
end
function solve_tilesize_constU(X, R, U)
    floor(Int, (REGISTER_COUNT - R[3] - R[4] - U*R[2]) / (U * R[1]))
end
function solve_tilesize_constT(X, R, T)
    floor(Int, (REGISTER_COUNT - R[3] - R[4]) / (T * R[1] + R[2]))
end
function solve_tilesize_constT(ls, T)
    R = @view ls.reg_pres[:,1]
    floor(Int, (REGISTER_COUNT - R[3] - R[4]) / (T * R[1] + R[2]))
end
# Tiling here is about alleviating register pressure for the UxT
function solve_tilesize(X, R, Umax, Tmax)
    first(R) == 0 && return -1,-1,Inf #solve_smalltilesize(X, R, Umax, Tmax)
    U, T, cost = solve_tilesize(X, R)
    # T -= T & 1
    # U = min(U, T)
    U_too_large = U > Umax
    T_too_large = T > Tmax
    if U_too_large
        if T_too_large
            U = Umax
            T = Tmax
        else # U too large, resolve T
            U = Umax
            T = solve_tilesize_constU(X, R, U)
        end
    elseif T_too_large
        T = Tmax
        U = solve_tilesize_constT(X, R, T)
    end
    U, T, cost
end
function solve_tilesize(
    ls::LoopSet, unrolled::Symbol, tiled::Symbol,
    cost_vec::AbstractVector{Float64} = @view(ls.cost_vec[:,1]),
    reg_pressure::AbstractVector{Int} = @view(ls.reg_pres[:,1])
)
    maxT = 4#8
    maxU = 4#8
    if isstaticloop(ls, tiled)
        maxT = min(2maxT, looprangehint(ls, tiled))
    end
    if isstaticloop(ls, unrolled)
        maxU = min(2maxU, looprangehint(ls, unrolled))
    end
    solve_tilesize(cost_vec, reg_pressure, maxU, maxT)
end

function set_upstream_family!(adal::Vector{T}, op::Operation, val::T) where {T}
    adal[identifier(op)] = val
    for opp ∈ parents(op)
        set_upstream_family!(adal, opp, val)
    end
end

# Just tile outer two loops?
# But optimal order within tile must still be determined
# as well as size of the tiles.
function evaluate_cost_tile(
    ls::LoopSet, order::Vector{Symbol}, vectorized::Symbol
)
    N = length(order)
    @assert N ≥ 2 "Cannot tile merely $N loops!"
    tiled = order[1]
    unrolled = order[2]
    ops = operations(ls)
    nops = length(ops)
    included_vars = fill(false, nops)
    unrolledtiled = fill(false, 2, nops)
    descendentsininnerloop = fill(false, nops)
    innerloop = last(order)
    iters = fill(-99.9, nops)
    nested_loop_syms = Symbol[]# Set{Symbol}()
    iter = 1.0
    # Need to check if fusion is possible
    size_T = biggest_type_size(ls)
    W, Wshift = VectorizationBase.pick_vector_width_shift(length(ls, vectorized), size_T)::Tuple{Int,Int}
    # costs =
    # cost_mat[1] / ( unrolled * tiled)
    # cost_mat[2] / ( tiled)
    # cost_mat[3] / ( unrolled)
    # cost_mat[4]
    # @show order
    cost_vec = cost_vec_buf(ls)
    reg_pressure = reg_pres_buf(ls)
    # @inbounds reg_pressure[2] = 1
    # @inbounds reg_pressure[3] = 1
    for n ∈ 1:N
        itersym = order[n]
        # Add to set of defined symbles
        push!(nested_loop_syms, itersym)
        if n == 1
            iter = length(ls, itersym) * length(ls, order[2]) / W
        elseif n > 2
            iter *= Float64(length(ls, itersym))
        end
        # check which vars we can define at this level of loop nest
        for (id, op) ∈ enumerate(ops)
            # isconstant(op) && continue
            # @assert id == identifier(op)+1 # testing, for now
            # won't define if already defined...
            included_vars[id] && continue
            # it must also be a subset of defined symbols
            loopdependencies(op) ⊆ nested_loop_syms || continue
            # # @show nested_loop_syms
            # # @show reduceddependencies(op)
            rd = reduceddependencies(op)
            hasintersection(rd, nested_loop_syms[1:end-length(rd)]) && return 0,0,Inf
            included_vars[id] = true
            unrolledtiled[1,id] = unrolled ∈ loopdependencies(op)
            unrolledtiled[2,id] = tiled ∈ loopdependencies(op)
            iters[id] = iter
            innerloop ∈ loopdependencies(op) && set_upstream_family!(descendentsininnerloop, op, true)
        end
    end
    for (id, op) ∈ enumerate(ops)
        iters[id] == -99.9 && continue
        descendentsininnerloop[id] || continue
        isunrolled = unrolledtiled[1,id]
        istiled = unrolledtiled[2,id]
        rt, lat, rp = cost(op, vectorized, Wshift, size_T)
            # @show instruction(op), rt, lat, rp, iter
        rt *= iters[id]
            # @show isunrolled, istiled
        if isunrolled && istiled # no cost decrease; cost must be repeated
            cost_vec[1] += rt
            reg_pressure[1] += rp
        elseif isunrolled # cost decreased by tiling
            cost_vec[2] += rt
            reg_pressure[2] += rp
        elseif istiled # cost decreased by unrolling
            cost_vec[3] += rt
            reg_pressure[3] += rp
        else# neither unrolled or tiled
            cost_vec[4] += rt
            reg_pressure[4] += rp
        end
    end
    # @show order, vectorized cost_vec reg_pressure
    # @show solve_tilesize(ls, unrolled, tiled, cost_vec, reg_pressure)
    solve_tilesize(ls, unrolled, tiled, cost_vec, reg_pressure)
end

struct LoopOrders
    syms::Vector{Symbol}
    buff::Vector{Symbol}
end
function LoopOrders(ls::LoopSet)
    syms = [s for s ∈ keys(ls.loops)]
    LoopOrders(syms, similar(syms))
end
function Base.iterate(lo::LoopOrders)
    lo.syms, zeros(Int, length(lo.syms))# - 1)
end

function swap!(x, i, j)
    xᵢ, xⱼ = x[i], x[j]
    x[j], x[i] = xᵢ, xⱼ
end
function advance_state!(state)
    N = length(state)
    for n ∈ 1:N
        sₙ = state[n]
        if sₙ == N - n
            if n == N
                return false
            else
                state[n] = 0
            end
        else
            state[n] = sₙ + 1
            break
        end
    end
    true
end
# I doubt this is the most efficient algorithm, but it's the simplest thing
# that I could come up with.
function Base.iterate(lo::LoopOrders, state)
    advance_state!(state) || return nothing
    # # @show state
    syms = copyto!(lo.buff, lo.syms)
    for i ∈ eachindex(state)
        sᵢ = state[i]
        sᵢ == 0 || swap!(syms, i, i + sᵢ)
    end
    syms, state
end
function choose_unroll_order(ls::LoopSet, lowest_cost::Float64 = Inf)
    lo = LoopOrders(ls)
    best_order = lo.syms
    new_order, state = iterate(lo) # right now, new_order === best_order
    best_vec = first(new_order)
    while true
        for new_vec ∈ new_order
            cost_temp = evaluate_cost_unroll(ls, new_order, lowest_cost, new_vec)
            if cost_temp < lowest_cost
                lowest_cost = cost_temp
                best_order = new_order
                best_vec = new_vec
            end
        end
        iter = iterate(lo, state)
        iter === nothing && return best_order, best_vec, lowest_cost
        new_order, state = iter
    end
end
function choose_tile(ls::LoopSet)
    lo = LoopOrders(ls)
    best_order = copyto!(ls.loop_order.bestorder, lo.syms)
    best_vec = first(best_order) # filler
    new_order, state = iterate(lo) # right now, new_order === best_order
    U, T, lowest_cost = 0, 0, Inf
    while true
        for new_vec ∈ @view(new_order[2:end]) # view to skip first
            U_temp, T_temp, cost_temp = evaluate_cost_tile(ls, new_order, new_vec)
            if cost_temp < lowest_cost
                lowest_cost = cost_temp
                U, T = U_temp, T_temp
                best_vec = new_vec
                copyto!(best_order, new_order)
                save_tilecost!(ls)
            end
        end
        iter = iterate(lo, state)
        iter === nothing && return best_order, best_vec, U, T, lowest_cost
        new_order, state = iter
    end
end
# Last in order is the inner most loop
function choose_order(ls::LoopSet)
    if num_loops(ls) > 1
        torder, tvec, tU, tT, tc = choose_tile(ls)
    else
        tc = Inf
    end
    uorder, uvec, uc = choose_unroll_order(ls, tc)
    if num_loops(ls) > 1 && tc ≤ uc
        return torder, tvec, min(tU, tT), tT
        # return torder, tvec, 4, 4#5, 5
    else
        return uorder, uvec, determine_unroll_factor(ls, uorder, first(uorder), uvec), -1
    end
end


# UnitCommitment.jl: Optimization Package for Security-Constrained Unit Commitment
# Copyright (C) 2020, UChicago Argonne, LLC. All rights reserved.
# Released under the modified BSD license. See COPYING.md for more details.

using UnitCommitment

basedir = @__DIR__

@testset "read_egret_solution" begin
    solution = UnitCommitment.read_egret_solution(
        "$basedir/../fixtures/egret_output.json.gz",
    )
    for attr in ["Is on", "Production (MW)", "Production cost (\$)"]
        @test attr in keys(solution)
        @test "115_STEAM_1" in keys(solution[attr])
        @test length(solution[attr]["115_STEAM_1"]) == 48
    end
    @test solution["Production cost (\$)"]["315_CT_6"][15:20] ==
          [0.0, 0.0, 884.44, 1470.71, 1470.71, 884.44]
    @test solution["Startup cost (\$)"]["315_CT_6"][15:20] ==
          [0.0, 0.0, 5665.23, 0.0, 0.0, 0.0]
    @test length(keys(solution["Is on"])) == 154
end