############################################################################################
"""
$(TYPEDEF)

Contains Particle container for propagation.

# Fields
$(TYPEDFIELDS)
"""
mutable struct Particles{R,P<:ParticleKernel,B,I<:Integer} <: AbstractParticles
    "Particle trajectories, for a discussion about possible shapes for the trajectories."
    #=
    Field that hold Nparticles Ndata times, CRITERIA:
    1) Must be able to sort values efficiently for a Column Major ~ Base Vector{Vector{T}} not performant
    2) must be able to be extended in Ndata dimension if new data arrives
    -> ElasticMatrix can do both
    + Can easily return a single particle trajectory
    -> Need to use view(elasticmatrix,...)
    + Can resize in Nparticles dimension for tuning
    -> not possible at the moment
    3) Interesting conversations:
    https://discourse.julialang.org/t/preallocating-vector-of-vectors/69304/10."
    https://discourse.julialang.org/t/elasticarrays-arraysofarrays-for-vector-of-arrays/70426/2
    -> a VectorOfVectors of SimilarArrays that is push!-able sounds interesting, but how
    would we depict multivariate trajectories in this case? The Matrix Variant would just store a vector in each field and thats it.
    =#
    val::ElasticMatrix{B,Vector{B}}
    "Saved ancestors of resampling step in pf"
    ancestor::ElasticMatrix{I,Vector{I}}
    "Distributions to propagate particle forward."
    kernel::P
    "Particle weights."
    weights::BaytesCore.ParameterWeights
    "Necessary buffer values for resampling particles."
    buffer::ParticleBuffer{B,I,R}
    "Log likelihood estimate."
    ℓobjective::Accumulator
    function Particles(
        reference::AbstractArray{T},
        prediction::Vector{R},
        kernel::P,
        ancestortype::Type{I},
        Nparticles::Integer,
        Ndata::Integer
    ) where {T,R,P<:ParticleKernel,I<:Integer}
        ## Create val
        val = ElasticMatrix{T}(undef, Nparticles, Ndata)
        ## Create ancestors - switch Nparticles, Ndata for easier access in resampling step
        ancestor = ElasticMatrix{ancestortype}(undef, Nparticles, Ndata)
        ## Create initial weights
        weights = BaytesCore.ParameterWeights(Nparticles)
        ## Create buffer
        buffer = ParticleBuffer(prediction, reference, Nparticles, Ndata, ancestortype)
        ## Return Particles
        return new{R,P,T,I}(val, ancestor, kernel, weights, buffer, Accumulator())
    end
end

function initial!(_rng::Random.AbstractRNG, particles::Particles, tune::ParticleFilterTune)
    return initial!(_rng, particles.kernel, particles.val, tune.iter.current)
end
function transition!(
    _rng::Random.AbstractRNG, particles::Particles, tune::ParticleFilterTune
)
    return transition!(_rng, particles.kernel, particles.val, tune.iter.current)
end

function weight!(dataₜ::D, particles::Particles, tune::ParticleFilterTune) where {D}
    return weight!(
        tune.weighting,
        particles.weights,
        dataₜ,
        particles.kernel,
        particles.val,
        tune.iter.current,
    )
end

#!NOTE: Here we can directly input tune.iter.current-1 for ancestors
function ancestors!(
    _rng::Random.AbstractRNG,
    particles::Particles,
    tune::ParticleFilterTune,
    ancestor=particles.ancestor,
)
    return ancestors!(
        _rng,
        tune.referencing,
        tune.resampling,
        ancestor,
        tune.iter.current - 1,
        tune.chains.Nchains,
        particles.weights.buffer,
    )
end
#!NOTE: This step is AFTER resample!() [ancestors!/get_reference!] during particle propagation hence tune.iter.current is used
function set_reference!(
    particles::Particles, tune::ParticleFilterTune, reference::AbstractArray{P}
) where {P}
    return set_reference!(tune.referencing, particles.val, reference, tune.iter.current)
end

function shuffle_forward!(particles::Particles, tune::ParticleFilterTune)
    return shuffle_forward!(
        particles.val,
        particles.buffer.parameter.val,
        particles.ancestor,
        particles.buffer.parameter.index,
        max(tune.memory), #tune.memory.latent,
        tune.iter.current - 1,
    )
end
function shuffle_backward!(particles::Particles, tune::ParticleFilterTune)
    return shuffle_backward!(
        particles.val, particles.buffer.parameter.val, particles.ancestor, particles.buffer.parameter.index
    )
end

############################################################################################
"""
$(SIGNATURES)
Resize particles struct with new `Nparticles` and `Ndata` size.

# Examples
```julia
```

"""
function resize!(
    particles::Particles,
    kernel::P,
    reference::AbstractArray{T},
    Nparticles::Integer,
    Ndata::Integer,
) where {T,P<:ParticleKernel}
    ## Change val
    particles.val = ElasticMatrix{T}(undef, Nparticles, Ndata)
    ## Change ancestor size
    F = eltype(particles.ancestor)
    particles.ancestor = ElasticMatrix{F}(undef, Nparticles, Ndata)
    ## Update buffer
    update!(particles.buffer, Nparticles, Ndata)
    ## Compute new initial weights
    particles.weights = BaytesCore.ParameterWeights(Nparticles)
    ##
    return nothing
end

############################################################################################
"""
$(SIGNATURES)
Update particles struct with new number of particles. Note that kernel and log likelihood are adjusted in another step.

# Examples
```julia
```

"""
function update!(
    particles::Particles,
    ParticleUpdate::BaytesCore.UpdateFalse,
    DataUpdate::BaytesCore.UpdateFalse,
    kernel::P,
    reference::AbstractArray{T},
    Nparticles::Integer,
    Ndata::Integer,
) where {T,P<:ParticleKernel}
    ## Set loglikelihood counter back to 0
    init!(particles.ℓobjective)
    ## Update kernel
    particles.kernel = kernel
    ## Return
    return nothing
end
function update!(
    particles::Particles,
    ParticleUpdate::BaytesCore.UpdateTrue,
    DataUpdate::D,
    kernel::P,
    reference::AbstractArray{T},
    Nparticles::Integer,
    Ndata::Integer,
) where {D<:BaytesCore.UpdateBool,T,P<:ParticleKernel}
    ## Set loglikelihood counter back to 0
    init!(particles.ℓobjective)
    ## Update kernel
    particles.kernel = kernel
    ## Resize Particles with new Nparticles and Ndata
    resize!(particles, kernel, reference, Nparticles, Ndata)
    ## Return
    return nothing
end
function update!(
    particles::Particles,
    ParticleUpdate::BaytesCore.UpdateFalse,
    DataUpdate::BaytesCore.UpdateTrue,
    kernel::P,
    reference::AbstractArray{T},
    Nparticles::Integer,
    Ndata::Integer,
) where {T,P<:ParticleKernel}
    ## Set loglikelihood counter back to 0
    init!(particles.ℓobjective)
    ## Update kernel
    particles.kernel = kernel
    ## Resize val and ancestor size
    resize!(particles.val, Nparticles, Ndata)
    resize!(particles.ancestor, Nparticles, Ndata)
    ## Resize buffer
    update!(particles.buffer, Nparticles, Ndata)
    ##
    return nothing
end

############################################################################################
"""
$(SIGNATURES)
Resample particle ancestors and shuffle current particles dependency.

# Examples
```julia
```

"""
function resample!(
    _rng::Random.AbstractRNG,
    particles::Particles,
    tune::ParticleFilterTune,
    reference::AbstractArray{P},
) where {P}
    #!NOTE: Resampling step for t-1 done at iteration t, before propagation starts. So ancestors_t-1 for particle_t are stored
    if BaytesCore.issmaller(
        BaytesCore.computeESS(particles.weights), tune.chains.Nchains * tune.chains.threshold
    )
        ## Resample ancestors
        ancestors!(_rng, particles, tune)
        ## Compute new reference ancestor
        get_reference!(_rng, particles, tune, reference)
        ## Set resampled == true in corresponding iteration -> happens at start of iteration, so iter-1 for resampled particles
        particles.buffer.resampled[tune.iter.current - 1] = true
        ## Equal weight normalized weights for next iteration memory
        Base.fill!(particles.weights.ℓweightsₙ, log(1.0 / tune.chains.Nchains))
        ## Shuffle particles according to resampled indices
        #!NOTE: We shuffle particles at iter-1 since we are doing this before transition at iter
        shuffle_forward!(particles, tune)
    else
        ## Set resampled to false
        particles.buffer.resampled[tune.iter.current - 1] = false
        ## Assign default order to ancestors
        @inbounds for Nrow in Base.OneTo(length(particles.buffer.parameter.val)) #size(particles.ancestor[tune.iter.current-1], 1) )
            particles.ancestor[Nrow, tune.iter.current - 1] = Nrow
        end
    end
    return nothing
end

"""
$(SIGNATURES)
Initialize particles.

# Examples
```julia
```

"""
function initial!(
    _rng::AbstractRNG,
    particles::Particles,
    tune::ParticleFilterTune,
    reference::AbstractArray{P},
    objective::Objective,
) where {P}
    ## Start iterating over initial distributions of particles ~ depends on memory of particles, usually 1
    for t in 1:tune.memory.initial #max(1, tune.memory.latent)
        ## Sample from initial distribution
        initial!(_rng, particles, tune)
        ## Update Last particle value based on Particle Filter settings (Marginal/Conditiona/Ancestral)
        set_reference!(particles, tune, reference)
        ## Calculate particle weights and log likelihood
        @inbounds if t > tune.memory.data #maxmemory
            weight!(BaytesCore.grab(objective.data, t, tune.config.data), particles, tune)
            update!(particles.ℓobjective, logmeanexp(particles.weights.ℓweights))
        end
        #!NOTE: No Resample step at init! for particle ancestors ~ first resampling should happen in propagate!() before particle propagation starts
        ## Update current iteration
        update!(tune.iter)
        ## Assign default order to ancestors
        @inbounds for Nrow in Base.OneTo(length(particles.buffer.parameter.val))
            particles.ancestor[Nrow, t] = Nrow
        end
    end
    return nothing
end

############################################################################################
"""
$(SIGNATURES)
Propagate particles forward.

# Examples
```julia
```

"""
function propagate!(
    _rng::AbstractRNG,
    particles::Particles,
    tune::ParticleFilterTune,
    reference::AbstractArray{P},
    objective::Objective,
) where {P}
    #!NOTE: Ndata a bit more flexible than size(reference, 1), as one could only propagate a few iterations even if more samples are known.
    for t in (tune.iter.current):(tune.chains.Ndata)
        ## Resample particle ancestors if resampling criterion fullfiled
        resample!(_rng, particles, tune, reference)
        ## Transition particles
        transition!(_rng, particles, tune)
        ## Update Last particle value based on Particle Filter settings (Marginal/Conditiona/Ancestral)
        set_reference!(particles, tune, reference)
        ## Calculate particle weights and log likelihood
        #!NOTE: cannot do this at the same time if particles are resampled adaptively, as normalized weights will change if not resampled
        @inbounds if t > tune.memory.data #maxmemory
            weight!(
                BaytesCore.grab(objective.data, tune.iter.current, tune.config.data),
                particles,
                tune,
            )
            update!(particles.ℓobjective, logmeanexp(particles.weights.ℓweights))
        end
        ## Update current iteration to Ndata+1 for propagation step
        update!(tune.iter)
    end
    ## Assign default order to ancestors for final index
    @inbounds for Nrow in Base.OneTo(length(particles.buffer.parameter.val))
        particles.ancestor[Nrow, tune.iter.current - 1] = Nrow
    end
    #!NOTE: NO more resampling step at final iteration, as we (1) resample ancestor (2) propagate particles, not (1) propagate particles (2) resample particles
    return nothing
end

############################################################################################
"""
$(SIGNATURES)
Predict new latent and observed data.

# Examples
```julia
```

"""

function predict(
    _rng::Random.AbstractRNG,
    trajectory::AbstractArray{P},
    kernel::ParticleKernel,
    reference::AbstractArray{P},
    iter::Integer,
) where {P}
    ## Predict new state
    stateᵗ⁺¹ = transition(_rng, kernel, trajectory, iter)
    ## Add prediction to proposal
    push!(trajectory, stateᵗ⁺¹)
    ## Predict new data
    dataᵗ⁺¹ = rand(_rng, kernel.evidence(trajectory, iter))
    return (eltype(reference)(stateᵗ⁺¹), dataᵗ⁺¹)
end
function predict(
    _rng,
    particles::Particles,
    tune::ParticleFilterTune,
    reference::AbstractArray{P},
    path::Integer,
) where {P}
    return predict(
        _rng,
        assign!(particles.buffer.proposal, view(particles.val, path, :)),
        particles.kernel,
        reference,
        tune.iter.current,
    )
end

"""
$(SIGNATURES)
Assign 'trajectory' elements to 'buffer' up until index 'iter'.

# Examples
```julia
```

"""
function assign!(buffer::AbstractVector{P}, trajectory::AbstractVector{T}) where {P,T}
    ## Resize buffer
    resize!(buffer, length(trajectory))
    ## Add trajectory elements to buffer
    @inbounds @simd for iter in eachindex(trajectory)
        buffer[iter] = trajectory[iter]
    end
    return buffer
end

############################################################################################
#export
export Particles, resample!, initial!, propagate!, predict


export DQNLearner

using Random
using Flux

mutable struct DQNLearner{
    Tq<:AbstractApproximator,
    Tt<:AbstractApproximator,
    Tf,
    S<:Union{Int,Nothing},
    R<:AbstractRNG,
} <: AbstractLearner
    approximator::Tq
    target_approximator::Tt
    loss_func::Tf
    stack_size::S
    γ::Float32
    batch_size::Int
    update_horizon::Int
    min_replay_history::Int
    update_freq::Int
    target_update_freq::Int
    update_step::Int
    rng::R
end

function DQNLearner(;
    approximator::Tq,
    target_approximator::Tt,
    loss_func::Tf,
    stack_size::Union{Int,Nothing} = 4,
    γ::Float32 = 0.99f0,
    batch_size::Int = 32,
    update_horizon::Int = 1,
    min_replay_history::Int = 32,
    update_freq::Int = 1,
    target_update_freq::Int = 100,
    update_step::Int = 0,
    seed = nothing,
) where {Tq,Tt,Tf}
    copyto!(approximator, target_approximator)
    rng = MersenneTwister(seed)
    DQNLearner(
        approximator,
        target_approximator,
        loss_func,
        stack_size,
        γ,
        batch_size,
        update_horizon,
        min_replay_history,
        update_freq,
        target_update_freq,
        update_step,
        rng,
    )
end

"""

!!! note
    The state of the observation is assumed to have been stacked,
    if `!isnothing(stack_size)`.
"""
(learner::DQNLearner)(obs) =
    obs |> get_state |>
    x ->
        send_to_device(device(learner.approximator), x) |> learner.approximator |>
        send_to_host

function RLBase.update!(learner::DQNLearner, batch)
    learner.update_step += 1
    learner.update_step % learner.update_freq == 0 || return nothing

    Q, Qₜ, γ, loss_func, update_horizon, batch_size = learner.approximator,
    learner.target_approximator,
    learner.γ,
    learner.loss_func,
    learner.update_horizon,
    learner.batch_size
    states, rewards, terminals, next_states = map(
        x -> send_to_device(device(Q), x),
        (batch.states, batch.rewards, batch.terminals, batch.next_states),
    )
    actions = CartesianIndex.(batch.actions, 1:batch_size)

    gs = gradient(params(Q)) do
        q = batch_estimate(Q, states)[actions]
        q′ = dropdims(maximum(batch_estimate(Qₜ, next_states); dims = 1), dims = 1)
        G = rewards .+ γ^update_horizon .* (1 .- terminals) .* q′
        loss_func(G, q)
    end

    update!(Q, gs)

    if learner.update_step % learner.target_update_freq == 0
        copyto!(Qₜ, Q)
    end
end

function RLBase.extract_experience(t::AbstractTrajectory, learner::DQNLearner)
    s = learner.stack_size
    h = learner.update_horizon
    n = learner.batch_size
    γ = learner.γ
    valid_ind_range = isnothing(s) ? (1:(length(t)-h)) : (s:(1:(length(t)-h)))
    if length(t) > learner.min_replay_history
        inds = rand(learner.rng, valid_ind_range, n)
        states = consecutive_view(get_trace(t, :state), inds; n_stack = s)
        actions = consecutive_view(get_trace(t, :action), inds)
        next_states = consecutive_view(get_trace(t, :state), inds .+ h; n_stack = s)
        consecutive_rewards = consecutive_view(get_trace(t, :reward), inds; n_horizon = h)
        consecutive_terminals =
            consecutive_view(get_trace(t, :terminal), inds; n_horizon = h)
        rewards, terminals = zeros(Float32, n), fill(false, n)

        # make sure that we only consider experiences in current episode
        for i in 1:n
            m = findfirst(view(consecutive_terminals, :, i))
            if isnothing(m)
                terminals[i] = false
                rewards[i] = discount_rewards_reduced(view(consecutive_rewards, :, i), γ)
            else
                terminals[i] = true
                rewards[i] = discount_rewards_reduced(view(consecutive_rewards, 1:m, i), γ)
            end
        end
        (
            states = states,
            actions = actions,
            rewards = rewards,
            terminals = terminals,
            next_states = next_states,
        )
    else
        nothing
    end
end


@testset "math partitioning" begin
    s = raw"A $B \com{C}$ D"
    parts = FP.md_partition(s)
    mathb = parts[2]
    @test mathb.name == :MATH_INLINE
    parts = FP.math_partition(FP.content(mathb))
    @test parts[1].name == :TEXT
    @test parts[2].name == :LX_COMMAND
    @test parts[3].name == :CU_BRACKETS
    @test FP.content(parts[1]) == "B "
    @test FP.content(parts[3]) == "C"
end


@everywhere using POMDPToolbox
@everywhere using POMDPModels

pomdp = BabyPOMDP()
fwc = FeedWhenCrying()
rnd = solve(RandomSolver(MersenneTwister(7)), pomdp)

q = [] # vector of the simulations to be run
push!(q, Sim(pomdp, fwc, max_steps=32, rng=MersenneTwister(4), metadata=Dict(:policy=>"feed when crying")))
push!(q, Sim(pomdp, rnd, max_steps=32, rng=MersenneTwister(4), metadata=Dict(:policy=>"random")))

# this creates two simulations, one with the feed-when-crying policy and one with a random policy

data = run_parallel(q)

# by default, the dataframe output contains the reward and the contents of `metadata`
@show data
# data = 2×2 DataFrames.DataFrame
# │ Row │ policy             │ reward   │
# ├─────┼────────────────────┼──────────┤
# │ 1   │ "feed when crying" │ -4.5874  │
# │ 2   │ "random"           │ -27.4139 │

# to perform additional analysis on each of the simulations one can define a processing function with the `do` syntax:
data2 = run_parallel(q,progress=false) do sim, hist
    println("finished a simulation - final state was $(last(state_hist(hist)))")
    return [:steps=>n_steps(hist), :reward=>discounted_reward(hist)]
end

@show data2


export det

import ..DTree

using mlpack._Internal.io

import mlpack_jll
const detLibrary = mlpack_jll.libmlpack_julia_det

# Call the C binding of the mlpack det binding.
function det_mlpackMain()
  success = ccall((:det, detLibrary), Bool, ())
  if !success
    # Throw an exception---false means there was a C++ exception.
    throw(ErrorException("mlpack binding error; see output"))
  end
end

" Internal module to hold utility functions. "
module det_internal
  import ..detLibrary

import ...DTree

# Get the value of a model pointer parameter of type DTree.
function IOGetParamDTree(paramName::String)::DTree
  DTree(ccall((:IO_GetParamDTreePtr, detLibrary), Ptr{Nothing}, (Cstring,), paramName))
end

# Set the value of a model pointer parameter of type DTree.
function IOSetParamDTree(paramName::String, model::DTree)
  ccall((:IO_SetParamDTreePtr, detLibrary), Nothing, (Cstring, Ptr{Nothing}), paramName, model.ptr)
end

# Serialize a model to the given stream.
function serializeDTree(stream::IO, model::DTree)
  buf_len = UInt[0]
  buf_ptr = ccall((:SerializeDTreePtr, detLibrary), Ptr{UInt8}, (Ptr{Nothing}, Ptr{UInt}), model.ptr, Base.pointer(buf_len))
  buf = Base.unsafe_wrap(Vector{UInt8}, buf_ptr, buf_len[1]; own=true)
  write(stream, buf)
end
# Deserialize a model from the given stream.
function deserializeDTree(stream::IO)::DTree
  buffer = read(stream)
  DTree(ccall((:DeserializeDTreePtr, detLibrary), Ptr{Nothing}, (Ptr{UInt8}, UInt), Base.pointer(buffer), length(buffer)))
end
end # module

"""
    det(; [folds, input_model, max_leaf_size, min_leaf_size, path_format, skip_pruning, test, training, verbose])

This program performs a number of functions related to Density Estimation Trees.
 The optimal Density Estimation Tree (DET) can be trained on a set of data
(specified by `training`) using cross-validation (with number of folds specified
with the `folds` parameter).  This trained density estimation tree may then be
saved with the `output_model` output parameter.

The variable importances (that is, the feature importance values for each
dimension) may be saved with the `vi` output parameter, and the density
estimates for each training point may be saved with the `training_set_estimates`
output parameter.

Enabling path printing for each node outputs the path from the root node to a
leaf for each entry in the test set, or training set (if a test set is not
provided).  Strings like 'LRLRLR' (indicating that traversal went to the left
child, then the right child, then the left child, and so forth) will be output.
If 'lr-id' or 'id-lr' are given as the `path_format` parameter, then the ID
(tag) of every node along the path will be printed after or before the L or R
character indicating the direction of traversal, respectively.

This program also can provide density estimates for a set of test points,
specified in the `test` parameter.  The density estimation tree used for this
task will be the tree that was trained on the given training points, or a tree
given as the parameter `input_model`.  The density estimates for the test points
may be saved using the `test_set_estimates` output parameter.

# Arguments

 - `folds::Int`: The number of folds of cross-validation to perform for
      the estimation (0 is LOOCV)  Default value `10`.

 - `input_model::DTree`: Trained density estimation tree to load.
 - `max_leaf_size::Int`: The maximum size of a leaf in the unpruned, fully
      grown DET.  Default value `10`.

 - `min_leaf_size::Int`: The minimum size of a leaf in the unpruned, fully
      grown DET.  Default value `5`.

 - `path_format::String`: The format of path printing: 'lr', 'id-lr', or
      'lr-id'.  Default value `lr`.

 - `skip_pruning::Bool`: Whether to bypass the pruning process and output
      the unpruned tree only.  Default value `false`.

 - `test::Array{Float64, 2}`: A set of test points to estimate the density
      of.
 - `training::Array{Float64, 2}`: The data set on which to build a density
      estimation tree.
 - `verbose::Bool`: Display informational messages and the full list of
      parameters and timers at the end of execution.  Default value `false`.

# Return values

 - `output_model::DTree`: Output to save trained density estimation tree
      to.
 - `tag_counters_file::String`: The file to output the number of points
      that went to each leaf.  Default value ``.

 - `tag_file::String`: The file to output the tags (and possibly paths)
      for each sample in the test set.  Default value ``.

 - `test_set_estimates::Array{Float64, 2}`: The output estimates on the
      test set from the final optimally pruned tree.
 - `training_set_estimates::Array{Float64, 2}`: The output density
      estimates on the training set from the final optimally pruned tree.
 - `vi::Array{Float64, 2}`: The output variable importance values for each
      feature.

"""
function det(;
             folds::Union{Int, Missing} = missing,
             input_model::Union{DTree, Missing} = missing,
             max_leaf_size::Union{Int, Missing} = missing,
             min_leaf_size::Union{Int, Missing} = missing,
             path_format::Union{String, Missing} = missing,
             skip_pruning::Union{Bool, Missing} = missing,
             test = missing,
             training = missing,
             verbose::Union{Bool, Missing} = missing,
             points_are_rows::Bool = true)
  # Force the symbols to load.
  ccall((:loadSymbols, detLibrary), Nothing, ());

  IORestoreSettings("Density Estimation With Density Estimation Trees")

  # Process each input argument before calling mlpackMain().
  if !ismissing(folds)
    IOSetParam("folds", convert(Int, folds))
  end
  if !ismissing(input_model)
    det_internal.IOSetParamDTree("input_model", convert(DTree, input_model))
  end
  if !ismissing(max_leaf_size)
    IOSetParam("max_leaf_size", convert(Int, max_leaf_size))
  end
  if !ismissing(min_leaf_size)
    IOSetParam("min_leaf_size", convert(Int, min_leaf_size))
  end
  if !ismissing(path_format)
    IOSetParam("path_format", convert(String, path_format))
  end
  if !ismissing(skip_pruning)
    IOSetParam("skip_pruning", convert(Bool, skip_pruning))
  end
  if !ismissing(test)
    IOSetParamMat("test", test, points_are_rows)
  end
  if !ismissing(training)
    IOSetParamMat("training", training, points_are_rows)
  end
  if verbose !== nothing && verbose === true
    IOEnableVerbose()
  else
    IODisableVerbose()
  end

  IOSetPassed("output_model")
  IOSetPassed("tag_counters_file")
  IOSetPassed("tag_file")
  IOSetPassed("test_set_estimates")
  IOSetPassed("training_set_estimates")
  IOSetPassed("vi")
  # Call the program.
  det_mlpackMain()

  return det_internal.IOGetParamDTree("output_model"),
         Base.unsafe_string(IOGetParamString("tag_counters_file")),
         Base.unsafe_string(IOGetParamString("tag_file")),
         IOGetParamMat("test_set_estimates", points_are_rows),
         IOGetParamMat("training_set_estimates", points_are_rows),
         IOGetParamMat("vi", points_are_rows)
end


function worker_size()
    ws = VectorizationBase.nextpow2(num_threads())
    IfElse.ifelse(Static.lt(ws,StaticInt{8}()), StaticInt{8}(), ws)
end
worker_type() = VectorizationBase.mask_type(worker_size())
worker_pointer_type() = Ptr{worker_type()}

const WORKERS = Ref(zero(UInt128)) # 0 = unavailable, 1 = available
const STATES = UInt8[]

worker_pointer() = Base.unsafe_convert(worker_pointer_type(), pointer_from_objref(WORKERS))

function reserved(id)
    p = Base.unsafe_convert(worker_pointer_type(), STATES)
    __vload(p, VectorizationBase.lazymul(cache_linesize(), id), False(), register_size())
end
function reserve_threads!(id, reserve)
    p = Base.unsafe_convert(worker_pointer_type(), STATES)
    __vstore!(p, reserve % worker_type(), VectorizationBase.lazymul(cache_linesize(), id), False(), False(), False(), register_size())
    nothing
end
function free_threads!(freed_threads)
    ThreadingUtilities._atomic_or!(worker_pointer(), freed_threads)
    nothing
end
function free_local_threads!()
    tid = Base.Threads.threadid()
    tmask = one(worker_type()) << (tid - one(tid))
    r = reserved(tid) | tmask
    reserve_threads!(tid, zero(worker_type()))
    free_threads!(r)
end

@inline function _request_threads(id::UInt32, num_requested::UInt32)
    reserved_threads = reserved(id)
    reserved_count = count_ones(reserved_threads)%UInt32
    no_threads = zero(worker_type())
    # reserved_count ≥ num_requested && return reserved_threads, no_threads
    reserved_count%Int32 ≥ num_requested%Int32 && return UnsignedIteratorEarlyStop(reserved_threads, num_requested), no_threads
    # to get more, we xchng, setting all to `0`
    # then see which we need, and free those we aren't using.
    wp = worker_pointer()
    _all_threads = all_threads = ThreadingUtilities._atomic_xchg!(wp, no_threads)
    additional_threads = count_ones(all_threads)
    # num_requested === StaticInt{-1}() && return reserved_threads, all_threads
    num_requested === StaticInt{-1}() && return UnsignedIteratorEarlyStop(reserved_threads | all_threads), all_threads
    excess = additional_threads + reserved_count - num_requested
    # signed(excess) ≤ 0 && return reserved_threads, all_threads
    signed(excess) ≤ 0 && return UnsignedIteratorEarlyStop(reserved_threads | all_threads), all_threads
    # we need to return the `excess` to the pool.
    lz = leading_zeros(all_threads) % UInt32
    # i = 8
    while true
        # start by trying to trim off excess from lz
        lz += excess%UInt32
        m = (one(worker_type()) << (UInt32(worker_size()) - lz)) - one(worker_type())
        masked = (all_threads & m) ⊻ all_threads
        excess -= count_ones(masked)
        all_threads &= (~masked)
        # @show bitstring(masked), count_ones(masked), bitstring(unused_threads), excess, lz, bitstring(all_threads)
        excess == 0 && break
        # i -= 1
        # @assert i > 0
    end
    ThreadingUtilities._atomic_store!(wp, _all_threads & (~all_threads))
    return UnsignedIteratorEarlyStop(reserved_threads | all_threads, num_requested), all_threads
end
@inline function request_threads(id, num_requested)
    _request_threads(id % UInt32, num_requested % UInt32)
end
reserved_threads(id) = UnsignedIteratorEarlyStop(reserved(id))
reserved_threads(id, count) = UnsignedIteratorEarlyStop(reserved(id), count % UInt32)


export sachdr, writesac, writesacpz

# ============================================================================
# Utility functions not for export

# Bytes 305:308 as a littleendian Int32 should read 0x06 0x00 0x00 0x00; compare each end to 0x0a to allow older SAC versions (if version in same place?)
function should_bswap(io::IO)
  skip(io, 304)
  u = read(io, UInt8)
  skip(io, 2)
  v = read(io, UInt8)
  q::Bool = (
    # Least significant byte in u
    if 0x00 < u < 0x0a && v == 0x00
      false
    # Most significant byte in u
    elseif u == 0x00 && 0x00 < v < 0x0a
      true
    else
      error("Invalid SAC file.")
    end
    )
  return q
end

function write_sac_file(fname::String, x::AbstractArray, tdata::Array{Float32,1}, xy::Bool)
  open(fname, "w") do io
    write(io, BUF.sac_fv)
    write(io, BUF.sac_iv)
    write(io, BUF.sac_cv)
    if eltype(x) == Float32
      write(io, x)
    elseif eltype(x) <: Complex
      write(io, Float32.(real.(x)))
    else
      write(io, Float32.(x))
    end
    if xy
      write(io, tdata)
    end
  end
  return nothing
end

function fill_sac(si::Int64, nx::Int32, ts::Int64, id::Array{String,1})
  @assert nx ≤ typemax(Int32)
  tt = [Base.parse(Int32, k) for k in split(string(u2d(ts*μs)), r"[\.\:T\-]")]
  L = length(tt)
  (L < 7) && append!(tt, zeros(Int32, 7-L))

  # Ints
  y = tt[1]
  j = Int32(md2j(y, tt[2], tt[3]))
  BUF.sac_iv[1] = y
  BUF.sac_iv[2] = j
  for i in 3:6
    BUF.sac_iv[i] = tt[i+1]
  end
  BUF.sac_iv[10] = Int32(nx)

  # Floats
  BUF.sac_fv[6] = rem(ts, 1000)*1.0f-3
  BUF.sac_fv[7] = BUF.sac_fv[6] + BUF.sac_fv[1]*nx

  # Filename
  y_s = lpad(y, 4, '0')
  j_s = lpad(j, 3, '0')
  h_s = lpad(tt[4], 2, '0')
  m_s = lpad(tt[5], 2, '0')
  s_s = lpad(tt[6], 2, '0')
  ms_s = lpad(tt[7], 3, '0')
  fname = join([y_s, j_s, h_s, m_s, s_s, ms_s, id[1], id[2], id[3], id[4], "R.SAC"], '.')
  return fname
end

function read_sac_stream(f::IO, fv::Array{Float32,1}, iv::Array{Int32,1}, cv::Array{UInt8,1}, full::Bool, swap::Bool)
  read!(f, fv)
  read!(f, iv)
  read!(f, cv)
  if swap == true
    fv .= bswap.(fv)
    iv .= bswap.(iv)
  end
  nx = getindex(iv, 10)
  x = Array{Float32,1}(undef, nx)
  read!(f, x)
  if swap == true
    x .= bswap.(x)
  end

  # floats
  gain = Float64(getindex(fv, 4))
  if gain == Float64(sac_nul_f)
    gain = one(Float64)
  end
  loc = GeoLoc()
  j = 0
  lf = (:lat, :lon, :el, :az, :inc)
  for k in (32,33,34,58,59)
    j += 1
    val = getindex(fv, k)
    if val != sac_nul_f
      setfield!(loc, lf[j], Float64(val))
    end
  end

  # ints
  ts = mktime(getindex(iv, 1),
              getindex(iv, 2),
              getindex(iv, 3),
              getindex(iv, 4),
              getindex(iv, 5),
              getindex(iv, 6)*Int32(1000))
  b = getindex(fv, 6)
  if b != sac_nul_f
    ts += round(Int64, b*1.0f3)
  end

  # chars
  id = zeros(UInt8, 15)
  i = 1
  while i < 193
    c = getindex(cv, i)
    if c == sac_nul_start && i < 188
      val = getindex(cv, i+1:i+5)
      if val == sac_nul_Int8
        cv[i:i+5] .= 0x20
      end
    elseif c == 0x00
      setindex!(cv, i, 0x20)
    end
    # fill ID
    if i == 1
      i = fill_id!(id, cv, i, 8, 4, 8)
    elseif i == 17
      i = fill_id!(id, cv, i, 24, 10, 11)
    elseif i == 161
      i = fill_id!(id, cv, i, 168, 13, 15)
    elseif i == 169
      i = fill_id!(id, cv, i, 176, 1, 2)
    else
      i = i+1
    end
  end
  deleteat!(id, id.==0x00)

  # Create a seischannel
  C = SeisChannel()
  setfield!(C, :id, String(id))
  setfield!(C, :fs, Float64(1.0f0/getindex(fv,1)))
  setfield!(C, :gain, Float64(gain))
  setfield!(C, :loc, loc)
  t = Array{Int64,2}(undef, 2, 2)
  setindex!(t, one(Int64), 1)
  setindex!(t, Int64(nx), 2)
  setindex!(t, ts, 3)
  setindex!(t, zero(Int64), 4)
  setfield!(C, :t, t)
  setfield!(C, :x, x)

  # Create dictionary if full headers are desired
  if full == true
    D = getfield(C, :misc)

    fk = getindex(sac_keys, 1)
    ik = getindex(sac_keys, 2)
    ck = getindex(sac_keys, 3)

    # Parse floats
    z = zero(Int16)
    o = one(Int16)
    k = ""
    m = z
    while m < Int16(70)
      m += o
      if fv[m] != sac_nul_f
        k = getindex(fk, m)
        D[k] = getindex(fv, m)
      end
    end

    # Parse ints
    m = z
    while m < Int16(40)
      m += o
      if iv[m] != sac_nul_i
        k = getindex(ik, m)
        D[k] = getindex(iv, m)
      end
    end

    m = z
    j = o
    while j < Int16(24)
      n = Int16(j == 2 ? 16 : 8)
      p = m + n
      ii = m
      while ii < p
        ii += o
        if getindex(cv,ii) != 0x20
          k = getindex(ck, j)
          D[k] = String(getindex(cv, m+o:p))
          break
        end
      end
      m += n
      j += o
    end
  end

  return C
end

function read_sac_file!(S::SeisData, fname::String, fv::Array{Float32,1}, iv::Array{Int32,1}, cv::Array{UInt8,1}, full::Bool)
  f = open(fname, "r")
  q = should_bswap(f)
  seekstart(f)
  C = read_sac_stream(f, fv, iv, cv, full, q)
  setfield!(C, :src, fname)
  note!(C, string("+src: ", fname))
  close(f)
  push!(S,C)
  return nothing
end

# ============================================================================
# NOTE: Leave keyword arguments, even if some aren't type-stable!
# Use of "optional" variables instead is a 5x **slowdown**

"""
    sachdr(f)

Print formatted SAC headers from file `f` to stdout. Does not accept wildcard
file strings.
"""
function sachdr(fname::String)
  S = read_data("sac", fname, full=true)
  for i = 1:S.n
    D = getindex(getfield(S, :misc), i)
    src = getindex(getfield(S, :src), i)
    printstyled(string(src, "\n"), color=:light_green, bold=true)
    for k in sort(collect(keys(D)))
      println(stdout, uppercase(k), ": ", string(D[k]))
    end
  end
  return nothing
end

function reset_sacbuf()
  checkbuf_strict!(BUF.sac_fv, 70)
  checkbuf_strict!(BUF.sac_iv, 40)
  checkbuf_strict!(BUF.sac_cv, 192)

  fill!(BUF.sac_fv, sac_nul_f)
  fill!(BUF.sac_iv, sac_nul_i)
  for i in 1:24
    BUF.sac_cv[1+8*(i-1):8*i] = sac_nul_c
  end
  BUF.sac_cv[17:24] .= 0x20
  BUF.sac_iv[7] = Int32(6)
end

function fill_sac_id(id_str::String)
  id = split_id(id_str)

  # Chars, all segments
  ci = [169, 1, 25, 161]
  for j = 1:4
    sj = ci[j]
    if isempty(id[j])
      BUF.sac_cv[sj:sj+7] .= sac_nul_c
    else
      s = codeunits(id[j])
      L = min(length(s), 8)
      copyto!(BUF.sac_cv, sj, s, 1, L)
      if L < 8
        BUF.sac_cv[sj+L:sj+7] .= 0x20
      end
    end
  end
  return id
end

function write_sac_channel(S::GphysData, i::Int64, v::Int64, xy::Bool, fn::String)
  id = fill_sac_id(S.id[i])
  fs = S.fs[i]
  t = S.t[i]
  tdata = Array{Float32,1}(undef, 0)

  # Floats, all segments
  BUF.sac_fv[1] = (fs == 0.0 ? 0.0f0 : Float32(1.0/fs))
  BUF.sac_fv[4] = Float32(S.gain[i])
  for i in (32, 33, 34, 58, 59)
    BUF.sac_fv[i] = 0.0f0
  end
  if !isempty(S.loc[i])
    loc = S.loc[i]
    if typeof(loc) == GeoLoc
      BUF.sac_fv[32] = Float32(getfield(loc, :lat))
      BUF.sac_fv[33] = Float32(getfield(loc, :lon))
      BUF.sac_fv[34] = Float32(getfield(loc, :el))
      BUF.sac_fv[58] = Float32(getfield(loc, :az))
      BUF.sac_fv[59] = Float32(getfield(loc, :inc))
    end
  end

  if (xy == true || fs == 0.0)
    si = 1
    nx = Int32(length(S.x[i]))
    @assert nx ≤ typemax(Int32)
    ts = t[1,2]
    tdata = Float32.(μs*(t_expand(t, fs) .- ts))
    if fn == ""
      fname = fill_sac(si, nx, ts, id)
    else
      fill_sac(si, nx, ts, id)
      fname = fn
    end
    BUF.sac_fv[7] += sum(t[2:end,2])*μs
    BUF.sac_iv[16] = Int32(4)
    BUF.sac_iv[36] = zero(Int32)
    write_sac_file(fname, S.x[i], tdata, true)
    v > 0  && println(stdout, now(), ": Wrote SAC xy file ", fname, " from channel ", i)
  else
    W = t_win(S.t[i], fs)
    inds = x_inds(S.t[i])
    BUF.sac_iv[16] = one(Int32)
    BUF.sac_iv[36] = one(Int32)
    for j in 1:size(inds,1)
      si = inds[j,1]
      ei = inds[j,2]
      ts = W[j,1]
      nx = Int32(ei-si+1)
      if fn == ""
        fname = fill_sac(si, nx, ts, id)
      else
        fill_sac(si, nx, ts, id)
        fname = fn
      end
      vx = view(S.x[i], si:ei)
      write_sac_file(fname, vx, tdata, false)
      v > 0  && println(stdout, now(), ": Wrote SAC ts file ", fname, " from channel ", i, " segment ", j)
    end
  end

end

"""
    writesac(S::Union{SeisData}[; ts=false, v=0])

Write all data in SeisData structure `S` to auto-generated SAC files.

Keywords:
* `fname=FF` uses filename FF, rather than creating file names automatically.
(Only works with GphysChannel objects)
* `xy=true` writes generic x-y data with time as the independent variable.
"""
function writesac(S::GphysData; fn::String="", xy::Bool=false, v::Int64=KW.v)
  reset_sacbuf()
  for i = 1:S.n
    write_sac_channel(S, i, v, xy, fn)
  end
  return nothing
end

function writesac(S::GphysChannel;
  fname::String="",
  xy::Bool=false,
  v::Int64=KW.v)

  fstr = String(
    if fname == ""
      fname
    else
      if endswith(lowercase(fname), ".sac")
        fname
      else
        fname * ".sac"
      end
    end
    )
  writesac(SeisData(S), fn=fstr, xy=xy, v=v)
  return nothing
end

function add_pzchan!(S::GphysData, D::Dict{String, Any}, file::String)
  id  = D["NETWORK   (KNETWK)"] * "." *
        D["STATION    (KSTNM)"] * "." *
        D["LOCATION   (KHOLE)"] * "." *
        D["CHANNEL   (KCMPNM)"]
  i = findid(id, S)
  loc   = GeoLoc( lat = parse(Float64, D["LATITUDE"]),
                  lon = parse(Float64, D["LONGITUDE"]),
                  el  = parse(Float64, D["ELEVATION"]),
                  dep = parse(Float64, D["DEPTH"]),
                  az  = parse(Float64, D["AZIMUTH"]),
                  inc = parse(Float64, D["DIP"])-90.0
                )
  fs    = parse(Float64, D["SAMPLE RATE"])

  # gain, units; note, not "INSTGAIN", that's just a scalar multipler
  gu    = split(D["SENSITIVITY"], limit=2, keepempty=false)
  gain  = parse(Float64, gu[1])
  units = lowercase(String(gu[2]))
  if startswith(units, "(")
    units = units[2:end]
  end
  if endswith(units, ")")
    units = units[1:end-1]
  end
  units = fix_units(units2ucum(units))

  #= fix for poorly-documented fundamental shortcoming:
    "INPUT UNIT"         => "M"
    I have no idea why SACPZ uses displacement PZ =#
  u_in = fix_units(units2ucum(D["INPUT UNIT"]))
  Z = get(D, "Z", ComplexF32[])
  if u_in != units
    if u_in == "m" && units == "m/s"
      deleteat!(Z, 1)
    elseif u_in == "m" && units == "m/s2"
      deleteat!(Z, 1:2)
    end
  end

  resp = PZResp(parse(Float32, D["A0"]),
                0.0f0,
                get(D, "P", ComplexF32[]),
                Z
                )

  if i == 0
    # resp
    C = SeisChannel()
    setfield!(C, :id, id)
    setfield!(C, :name, D["DESCRIPTION"])
    setfield!(C, :loc, loc)
    setfield!(C, :fs, fs)
    setfield!(C, :gain, gain)
    setfield!(C, :resp, resp)
    setfield!(C, :units, units)
    setfield!(C, :src, file)
    setfield!(C, :misc, D)
    note!(C, "+src : read_sacpz " * file * ")")
    push!(S, C)
  else
    ts = Dates.DateTime(get(D, "START", "1970-01-01T00:00:00")).instant.periods.value*1000 - dtconst
    te = Dates.DateTime(get(D, "END", "2599-12-31T23:59:59")).instant.periods.value*1000 - dtconst
    t0 = isempty(S.t[i]) ? ts : S.t[i][1,2]
    if ts ≤ t0 ≤ te
      if S.fs[i] == 0.0
        S.fs[i] = fs
      end

      if isempty(S.units[i])
        S.units[i] = units
      end

      if S.gain[i] == 1.0
        S.gain[i] = gain
      end

      if typeof(S.resp[i]) == GenResp || isempty(S.resp[i])
        S.resp[i] = resp
      end

      if isempty(S.name[i])
        S.name[i] = D["DESCRIPTION"]
      end

      if isempty(S.loc[i])
        S.loc[i]  = loc
      end

      S.misc[i] = merge(D, S.misc[i])
    end
  end
  return nothing
end

@doc """
    read_sacpz!(S::GphysData, pzfile::String)

Read sacpz file `pzfile` into SeisIO struct `S`.

If an ID in the pz file matches channel `i` at times in `S.t[i]`:
* Fields :fs, :gain, :loc, :name, :resp, :units are overwritten if empty/unset
* Information from the pz file is merged into :misc if the corresponding keys
aren't in use.
""" read_sacpz!
function read_sacpz!(S::GphysData, file::String)
  io = open(file, "r")
  read_state = 0x00
  D = Dict{String, Any}()
  kv = Array{String, 1}(undef, 2)

  # Do this for each channel
  while true

    # EOF
    if eof(io)
      add_pzchan!(S, D, file)
      break
    end

    line = readline(io)

    # Header section
    if startswith(line, "*")
      if endswith(strip(line), "**")
        read_state += 0x01
        if read_state == 0x03
          add_pzchan!(S, D, file)
          read_state = 0x01
          D = Dict{String, Any}()
        end
      else
        kv .= strip.(split(line[2:end], ":", limit=2, keepempty=false))
        D[kv[1]] = kv[2]
      end

    # Zeros section
    elseif startswith(line, "ZEROS")
      N = parse(Int64, split(line, limit=2, keepempty=false)[2])
      D["Z"] = Array{Complex{Float32},1}(undef, N)
      for i = 1:N
        try
          mark(io)
          zc = split(readline(io), limit=2, keepempty=false)
          D["Z"][i] = complex(parse(Float32, zc[1]), parse(Float32, zc[2]))
        catch
          D["Z"][i:N] .= zero(ComplexF32)
          reset(io)
        end
      end

    # Poles section
    elseif startswith(line, "POLES")
      N = parse(Int64, split(line, limit=2, keepempty=false)[2])
      D["P"] = Array{Complex{Float32},1}(undef, N)
      for i = 1:N
        pc = split(readline(io), limit=2, keepempty=false)
        D["P"][i] = complex(parse(Float32, pc[1]), parse(Float32, pc[2]))
      end

    # Constant section
    elseif startswith(line, "CONSTANT")
      D["CONSTANT"] = String(split(line, limit=2, keepempty=false)[2])
    end
  end
  close(io)
  return S
end

@doc (@doc read_sacpz)
function read_sacpz(file::String)
  S = SeisData()
  read_sacpz!(S, file)
  return S
end

@doc """
    writesacpz(S::GphysData, pzfile::String)

Write fields from SeisIO struct `S` into sacpz file `pzfile`. Uses information
from fields :fs, :gain, :loc, :misc, :name, :resp, :units.
""" writesacpz
function writesacpz(S::GphysData, file::String)
  io = open(file, "w")
  for i in 1:S.n
    id = split(S.id[i], ".")
    created   = get(S.misc[i], "CREATED", string(u2d(time())))
    ts_str    = isempty(S.t[i]) ? "1970-01-01T00:00:00" : string(u2d(S.t[i][1,2]*1.0e-6))
    t_start   = get(S.misc[i], "START", ts_str)
    t_end     = get(S.misc[i], "END", "2599-12-31T23:59:59")
    unit_in   = get(S.misc[i], "INPUT UNIT", "?")
    unit_out  = get(S.misc[i], "OUTPUT UNIT", "?")

    Y = typeof(S.resp[i])
    if Y == GenResp
      a0 = 1.0
      P = S.resp[i][:,1]
      Z = deepcopy(S.resp[i][:,2])
    elseif Y in (PZResp, PZResp64)
      a0 = getfield(S.resp[i], :a0)
      P = getfield(S.resp[i], :p)
      Z = deepcopy(getfield(S.resp[i], :z))
    elseif Y == MultiStageResp
      j = 0
      for k = 1:length(S.resp[i].stage)
        stg = S.resp[i].stage[k]
        if typeof(stg) in (PZResp, PZResp64)
          j = k
          break
        end
      end
      if j == 0
        @warn(string("Skipped channel ", i, " (", id, "): incompatible response Type"))
        continue
      else
        a0 = getfield(S.resp[i].stage[j], :a0)
        P = getfield(S.resp[i].stage[j], :p)
        Z = deepcopy(getfield(S.resp[i].stage[j], :z))
      end
    end

    write(io, 0x2a)
    write(io, 0x20)
    write(io, fill!(zeros(UInt8, 34), 0x2a))
    write(io, 0x0a)

    write(io, "* NETWORK   (KNETWK): ", id[1], 0x0a)
    write(io, "* STATION    (KSTNM): ", id[2], 0x0a)
    write(io, "* LOCATION   (KHOLE): ", isempty(id[3]) ? "  " : id[3], 0x0a)
    write(io, "* CHANNEL   (KCMPNM): ", id[4], 0x0a)

    write(io, "* CREATED           : ", created, 0x0a)
    write(io, "* START             : ", t_start, 0x0a)
    write(io, "* END               : ", t_end, 0x0a)
    write(io, "* DESCRIPTION       : ", S.name[i], 0x0a)
    write(io, "* LATITUDE          : ", @sprintf("%0.6f", S.loc[i].lat), 0x0a)
    write(io, "* LONGITUDE         : ", @sprintf("%0.6f", S.loc[i].lon), 0x0a)
    write(io, "* ELEVATION         : ", string(S.loc[i].el), 0x0a)
    write(io, "* DEPTH             : ", string(S.loc[i].dep), 0x0a)
    write(io, "* DIP               : ", string(S.loc[i].inc+90.0), 0x0a)
    write(io, "* AZIMUTH           : ", string(S.loc[i].az), 0x0a)
    write(io, "* SAMPLE RATE       : ", string(S.fs[i]), 0x0a)

    for j in ("INPUT UNIT", "OUTPUT UNIT", "INSTTYPE", "INSTGAIN", "COMMENT")
      write(io, 0x2a, 0x20)
      write(io, rpad(j, 18))
      write(io, 0x3a, 0x20)
      v = get(S.misc[i], j, "")
      write(io, v)
      if j == "INSTGAIN" && v == ""
        write(io, "1.0E+00 (", S.units[i], ")")
      end
      write(io, 0x0a)
    end

    write(io, "* SENSITIVITY       : ", @sprintf("%12.6e", S.gain[i]), 0x20, 0x28, uppercase(S.units[i]), 0x29, 0x0a)
    NZ = length(Z)
    NP = length(P)
    write(io, "* A0                : ", @sprintf("%12.6e", a0), 0x0a)
    CONST = get(S.misc[i], "CONSTANT", string(a0*S.gain[i]))

    write(io, 0x2a)
    write(io, 0x20)
    write(io, fill!(zeros(UInt8, 34), 0x2a))
    write(io, 0x0a)

    # fix for units_in always being m
    if S.units[i] == "m/s"
      NZ += 1
      pushfirst!(Z, zero(ComplexF32))
    elseif S.units[i] == "m/s2"
      NZ += 2
      prepend!(Z, zeros(ComplexF32, 2))
    end
    write(io, "ZEROS\t", string(NZ), 0x0a)
    for i = 1:NZ
      write(io, 0x09, @sprintf("%+12.6e", real(Z[i])), 0x09, @sprintf("%+12.6e", imag(Z[i])), 0x09, 0x0a)
    end

    write(io, "POLES\t", string(NP), 0x0a)
    for i = 1:NP
      write(io, 0x09, @sprintf("%+12.6e", real(P[i])), 0x09, @sprintf("%+12.6e", imag(P[i])), 0x09, 0x0a)
    end

    write(io, "CONSTANT\t", CONST, 0x0a)
    write(io, 0x0a, 0x0a)
  end
  close(io)
  return nothing
end


Base.promote_rule(::Type{VaxInt16}, ::Type{T}) where T <: Union{Int8,Int16} = Int16
Base.promote_rule(::Type{VaxInt16}, ::Type{T}) where T <: Union{Int32,Int64,Int128} = T
Base.promote_rule(::Type{VaxInt16}, ::Type{T}) where T <: IEEEFloat = T

Base.promote_rule(::Type{VaxInt32}, ::Type{T}) where T <: Union{Int8,Int16,Int32,VaxInt16} = Int32
Base.promote_rule(::Type{VaxInt32}, ::Type{T}) where T <: Union{Int64,Int128,IEEEFloat} = T
Base.promote_rule(::Type{VaxInt32}, ::Type{T}) where T <: IEEEFloat = T

Base.promote_rule(::Type{VaxFloatF}, ::Type{T}) where T <: Union{Int8,Int16,Int32,Int64,Int128} = Float32
Base.promote_rule(::Type{VaxFloatF}, ::Type{T}) where T <: Union{Float16,Float32} = Float32
Base.promote_rule(::Type{VaxFloatF}, ::Type{T}) where T <: VaxInt = Float32
Base.promote_rule(::Type{VaxFloatF}, ::Type{Float64}) = Float64

Base.promote_rule(::Type{VaxFloatD}, ::Type{T}) where T <: Union{Int8,Int16,Int32,Int64,Int128} = Float64
Base.promote_rule(::Type{VaxFloatD}, ::Type{T}) where T <: Union{VaxFloatF,VaxFloatG} = Float64
Base.promote_rule(::Type{VaxFloatD}, ::Type{T}) where T <: VaxInt = Float64
Base.promote_rule(::Type{VaxFloatD}, ::Type{T}) where T <: IEEEFloat = Float64

Base.promote_rule(::Type{VaxFloatG}, ::Type{T}) where T <: Union{Int8,Int16,Int32,Int64,Int128} = Float64
Base.promote_rule(::Type{VaxFloatG}, ::Type{T}) where T <: Union{VaxFloatF,VaxFloatG} = Float64
Base.promote_rule(::Type{VaxFloatG}, ::Type{T}) where T <: VaxInt = Float64
Base.promote_rule(::Type{VaxFloatG}, ::Type{T}) where T <: IEEEFloat = Float64

Base.promote_rule(::Type{BigFloat}, ::Type{<:VaxFloat}) = BigFloat
Base.promote_rule(::Type{BigFloat}, ::Type{<:VaxInt}) = BigFloat

Base.promote_rule(::Type{BigInt}, ::Type{<:VaxFloat}) = BigFloat
Base.promote_rule(::Type{BigInt}, ::Type{<:VaxInt}) = BigFloat

Base.promote_type(::Type{VaxInt16}, ::Type{VaxInt16}) = Int16
Base.promote_type(::Type{VaxInt32}, ::Type{VaxInt32}) = Int32

Base.promote_type(::Type{VaxFloatF}, ::Type{VaxFloatF}) = Float32
Base.promote_type(::Type{VaxFloatD}, ::Type{VaxFloatD}) = Float64
Base.promote_type(::Type{VaxFloatG}, ::Type{VaxFloatG}) = Float64


### SELECTED DATA FILE INDEXES REQUIRED FOR TGFb MODEL ONLY

# INDEX VALUES CHANGE WITH UPDATES TO MODEL.NET

initial_cond_update_array = [];
rate_constant_update_array = [];
control_parameter_update_array = [];

# initial conditions
#push!(initial_cond_update_array,[tgfb12_initial_cond_index,0.0]) # TGFB12 initial conditions

# This gets hacky, we are using the TGFB infastructure to run the inhibitors etc.

# rates
# TGFB dosing
push!(rate_constant_update_array,[yreg_translation_rate_index,10.0]) # addition_TGFb12: [] =([])=> TGFB12
push!(rate_constant_update_array,[yreg_deg_rate_index,yreg_ox_deg_rate]) # addition_TGFb12: [] =([])=> TGFB12

# Run snail slug instead
#push!(rate_constant_update_array,[152,10.0]) # induction_gene_SNAIL_SLUG: gene_SNAIL_SLUG = mRNA_SNAIL_SLUG
#push!(rate_constant_update_array,[ss_virus_forward_rate_index,4.0])
#push!(rate_constant_update_array,[ss_virus_reverse_rate_index,0.1])

# give MAPK inhibitor

#push!(rate_constant_update_array,[U0126_foward_rate_index,4.0])
#push!(rate_constant_update_array,[U0126_reverse_rate_index,0.1])

# ECADHERINchange this to transcription rate
# push!(rate_constant_update_array,[ecad_trans_rate_index,ecad_trans_rate*5.0])

# Block TGFB3
#push!(rate_constant_update_array,[TGFB3_Ab_foward_rate_index,4.0])
#push!(rate_constant_update_array,[TGFB3_Ab_reverse_rate_index,0.1])

# TGFB3 dosing

#push!(rate_constant_update_array,[TGFB3_foward_rate_index,4.0])
#push!(rate_constant_update_array,[TGFB3_reverse_rate_index,0.1])

# DN LEF1 temoprary
#lef1_induciton_rate_index = 141
#lef1_induciton_rate = 0.01

# DN SMAD temoprary
#SMAD_induciton_rate_index = 134
#SMAD_induciton_rate = 0.01

# push modifications to RAM
#UpdateArray(data_dictionary,"INITIAL_CONDITION_ARRAY",initial_cond_update_array)
#UpdateArray(data_dictionary,"RATE_CONSTANT_ARRAY",rate_constant_update_array)

# mRNA degradation test
#tubulin_induction_index = 145
#push!(rate_constant_update_array,[tubulin_induction_index,0.0])


### Latent models ####
include("prior.jl")
include("posterior.jl")

## Exact Gaussian Process
struct LatentGP{T,Tpr<:GPPrior,Tpo<:Posterior{T},O} <: AbstractLatent{T,Tpr,Tpo}
    prior::Tpr
    post::Tpo
    opt::O
end

function LatentGP(T::DataType, dim::Int, kernel::Kernel, mean::PriorMean, opt)
    return LatentGP(
        GPPrior(deepcopy(kernel), deepcopy(mean), cholesky(Matrix{T}(I, dim, dim))),
        Posterior(dim, zeros(T, dim), cholesky(Matrix{T}(I(dim)))),
        deepcopy(opt),
    )
end

@traitimpl IsFull{LatentGP}

## AbstractVarLatent

abstract type AbstractVarLatent{T,Tpr,Tpo} <: AbstractLatent{T,Tpr,Tpo} end

## Variational Gaussian Process
mutable struct VarLatent{T,Tpr<:GPPrior,Tpo<:VarPosterior{T},O} <:
               AbstractVarLatent{T,Tpr,Tpo}
    prior::Tpr
    post::Tpo
    opt::O
end

function VarLatent(T::DataType, dim::Int, kernel::Kernel, mean::PriorMean, opt)
    return VarLatent(
        GPPrior(deepcopy(kernel), deepcopy(mean), cholesky(Matrix{T}(I, dim, dim))),
        VarPosterior{T}(dim),
        deepcopy(opt),
    )
end

@traitimpl IsFull{VarLatent}

## Sparse Variational Gaussian Process

mutable struct SparseVarLatent{
    T,Tpr<:GPPrior,Tpo<:VarPosterior{T},Topt,TZ<:AbstractVector,TZopt
} <: AbstractVarLatent{T,Tpr,Tpo}
    prior::Tpr
    post::Tpo
    Z::TZ
    Knm::Matrix{T}
    κ::Matrix{T}
    K̃::Vector{T}
    opt::Topt
    Zopt::TZopt
end

function SparseVarLatent(
    T::DataType,
    dim::Int,
    S::Int,
    Z::AbstractVector,
    kernel::Kernel,
    mean::PriorMean,
    opt=nothing,
    Zopt=nothing
)
    return SparseVarLatent(
        GPPrior(deepcopy(kernel), deepcopy(mean), cholesky(Matrix{T}(I(dim)))),
        VarPosterior{T}(dim),
        deepcopy(Z),
        Matrix{T}(undef, S, dim),
        Matrix{T}(undef, S, dim),
        Vector{T}(undef, S),
        deepcopy(opt),
        deepcopy(Zopt),
    )
end

@traitimpl IsSparse{SparseVarLatent}

## Monte-Carlo Gaussian Process

struct SampledLatent{T,Tpr<:GPPrior,Tpo<:SampledPosterior{T}} <: AbstractLatent{T,Tpr,Tpo}
    prior::Tpr
    post::Tpo
end

function SampledLatent(T::DataType, dim::Int, kernel::Kernel, mean::PriorMean)
    return SampledLatent(
        GPPrior(deepcopy(kernel), deepcopy(mean), cholesky(Matrix{T}(I, dim, dim))),
        SampledPosterior(dim, zeros(T, dim), Symmetric(Matrix{T}(I(dim)))),
    )
end

@traitimpl IsFull{SampledLatent}

## Online Sparse Variational Process

mutable struct OnlineVarLatent{T,Tpr<:GPPrior,Tpo<:AbstractVarPosterior{T},Topt,TZ<:AbstractVector,
    TZalg<:InducingPoints.OnIPSA,TZopt} <:
               AbstractVarLatent{T,Tpo,Tpr}
    prior::Tpr
    post::Tpo
    Z::TZ
    Zalg::TZalg
    Knm::Matrix{T}
    κ::Matrix{T}
    K̃::Vector{T}
    Zupdated::Bool
    opt::Topt
    Zopt::TZopt
    Zₐ::AbstractVector
    Kab::Matrix{T}
    κₐ::Matrix{T}
    K̃ₐ::Matrix{T}
    invDₐ::Symmetric{T,Matrix{T}}
    prev𝓛ₐ::T
    prevη₁::Vector{T}
end

function OnlineVarLatent(
    T::DataType,
    dim::Int,
    nSamplesUsed::Int,
    Z::AbstractVector,
    Zalg::InducingPoints.OnIPSA,
    kernel::Kernel,
    mean::PriorMean,
    opt=nothing,
    Zopt=nothing
)
    return OnlineVarLatent(
        GPPrior(deepcopy(kernel), deepcopy(mean), cholesky(Matrix{T}(I, dim, dim))),
        OnlineVarPosterior{T}(dim),
        Z,
        Zalg,
        Matrix{T}(undef, nSamplesUsed, dim),
        Matrix{T}(undef, nSamplesUsed, dim),
        Vector{T}(undef, nSamplesUsed),
        false,
        deepcopy(opt),
        deepcopy(Zopt),
        deepcopy(Z),
        Matrix{T}(I, dim, dim),
        Matrix{T}(I, dim, dim),
        Matrix{T}(I, dim, dim),
        Symmetric(Matrix{T}(I, dim, dim)),
        zero(T),
        Vector{T}(undef, dim),
    )
end

@traitimpl IsSparse{OnlineVarLatent}

## Variational Student-T Process

mutable struct TVarLatent{T<:Real,Tpr<:TPrior,Tpo<:VarPosterior{T},O} <:
               AbstractVarLatent{T,Tpr,Tpo}
    prior::Tpr
    post::Tpo
    opt::O
end

function TVarLatent(T::DataType, ν::Real, dim::Int, kernel::Kernel, mean::PriorMean, opt)
    return TVarLatent(
        TPrior(
            deepcopy(kernel),
            deepcopy(mean),
            cholesky(Matrix{T}(I, dim, dim)),
            ν,
            rand(T),
            rand(T),
        ),
        VarPosterior{T}(dim),
        deepcopy(opt),
    )
end

@traitimpl IsFull{TVarLatent}

### Functions

prior(gp::AbstractLatent) = gp.prior
kernel(gp::AbstractLatent) = kernel(prior(gp))
setkernel!(gp::AbstractLatent, kernel::Kernel) = setkernel!(prior(gp), kernel)
pr_mean(gp::AbstractLatent) = mean(prior(gp))
pr_mean(gp::AbstractLatent, X::AbstractVector) = mean(prior(gp), X)
setpr_mean!(gp::AbstractLatent, μ₀::PriorMean) = setmean!(prior(gp), μ₀)
pr_cov(gp::AbstractLatent) = cov(prior(gp))
pr_cov(gp::TVarLatent) = prior(gp).χ * cov(prior(gp))
pr_cov!(gp::AbstractLatent, K::Cholesky) = gp.prior.K = K

posterior(gp::AbstractLatent) = gp.post
Distributions.dim(gp::AbstractLatent) = dim(posterior(gp))
Distributions.mean(gp::AbstractLatent) = mean(posterior(gp))
Distributions.cov(gp::AbstractLatent) = cov(posterior(gp))
Distributions.var(gp::AbstractLatent) = var(posterior(gp))
nat1(gp::AbstractVarLatent) = nat1(posterior(gp))
nat2(gp::AbstractVarLatent) = nat2(posterior(gp))

mean_f(model::AbstractGP) = mean_f.(model.f)

@traitfn mean_f(gp::T) where {T <: AbstractLatent; IsFull{T}} = mean_f(mean(gp))
@traitfn mean_f(gp::T) where {T <: AbstractLatent; !IsFull{T}} = mean_f(mean(gp), gp.κ)

mean_f(μ::AbstractVector) = μ
mean_f(μ::AbstractVector, κ::AbstractMatrix) = κ * μ

var_f(model::AbstractGP) = var_f.(model.f)

@traitfn var_f(gp::T) where {T <: AbstractLatent; IsFull{T}} = var_f(cov(gp))
@traitfn var_f(gp::T) where {T <: AbstractLatent; !IsFull{T}} = var_f(cov(gp), gp.κ, gp.K̃)

var_f(Σ::AbstractMatrix) = diag(Σ)
var_f(Σ::AbstractMatrix, κ::AbstractMatrix, K̃::AbstractVector) = diag_ABt(κ * Σ, κ) + K̃

Zview(gp::SparseVarLatent) = gp.Z
Zview(gp::OnlineVarLatent) = gp.Z

setZ!(gp::AbstractLatent, Z::AbstractVector) = gp.Z = Z

opt(gp::AbstractLatent) = gp.opt
Zopt(::AbstractLatent) = nothing
Zopt(gp::SparseVarLatent) = gp.Zopt
Zopt(gp::OnlineVarLatent) = gp.Zopt

@traitfn function compute_K!(
    gp::TGP, X::AbstractVector, jitt::Real
) where {TGP <: AbstractLatent; IsFull{TGP}}
    return pr_cov!(gp, cholesky(kernelmatrix(kernel(gp), X) + jitt * I))
end

@traitfn function compute_K!(gp::T, jitt::Real) where {T <: AbstractLatent; !IsFull{T}}
    return pr_cov!(gp, cholesky(kernelmatrix(kernel(gp), gp.Z) + jitt * I))
end

function compute_κ!(gp::SparseVarLatent, X::AbstractVector, jitt::Real)
    gp.Knm = kernelmatrix(kernel(gp), X, gp.Z)
    gp.κ = copy(gp.Knm / pr_cov(gp))
    gp.K̃ = kernelmatrix_diag(kernel(gp), X) .+ jitt - diag_ABt(gp.κ, gp.Knm)
    return all(gp.K̃ .> 0) || error("K̃ has negative values")
end

function compute_κ!(gp::OnlineVarLatent, X::AbstractVector, jitt::Real)
    # Covariance with the model at t-1
    gp.Kab = kernelmatrix(kernel(gp), gp.Zₐ, gp.Z)
    gp.κₐ = gp.Kab / pr_cov(gp)
    Kₐ = Symmetric(kernelmatrix(kernel(gp), gp.Zₐ) + jitt * I)
    gp.K̃ₐ = Kₐ - gp.κₐ * transpose(gp.Kab)

    # Covariance with a new batch
    gp.Knm = kernelmatrix(kernel(gp), X, gp.Z)
    gp.κ = gp.Knm / pr_cov(gp)
    gp.K̃ = kernelmatrix_diag(kernel(gp), X) .+ jitt - diag_ABt(gp.κ, gp.Knm)
    all(gp.K̃ .> 0) || error("K̃ has negative values")
end


module Models

export IncompressibleModel, NonDimensionalModel, Clock, tick!, state

using Adapt

using Oceananigans.Architectures
using Oceananigans.Fields
using Oceananigans.Coriolis
using Oceananigans.Buoyancy
using Oceananigans.TurbulenceClosures
using Oceananigans.BoundaryConditions
using Oceananigans.Solvers
using Oceananigans.Forcing
using Oceananigans.Utils

"""
    AbstractModel

Abstract supertype for models.
"""
abstract type AbstractModel end

"""
    state(model)

Returns a `NamedTuple` with fields `velocities, tracers, diffusivities, tendencies`
corresponding to `NamedTuple`s of `OffsetArray`s that reference each of the field's data.
"""
@inline state(model) = (   velocities = datatuple(model.velocities),
                              tracers = datatuple(model.tracers),
                        diffusivities = datatuple(model.diffusivities))

include("clock.jl")
include("incompressible_model.jl")
include("non_dimensional_model.jl")
include("show_models.jl")

end


include("TwoPartCost.jl")
include("ThreePartCost.jl")
include("TechHydro.jl")
include("TechRenewable.jl")
include("TechThermal.jl")
include("Area.jl")
include("LoadZone.jl")
include("Bus.jl")
include("Arc.jl")
include("Line.jl")
include("MonitoredLine.jl")
include("PhaseShiftingTransformer.jl")
include("TapTransformer.jl")
include("Transformer2W.jl")
include("HVDCLine.jl")
include("VSCDCLine.jl")
include("InterruptibleLoad.jl")
include("FixedAdmittance.jl")
include("PowerLoad.jl")
include("HydroEnergyReservoir.jl")
include("HydroDispatch.jl")
include("RenewableDispatch.jl")
include("RenewableFix.jl")
include("ThermalStandard.jl")
include("GenericBattery.jl")
include("StaticReserve.jl")
include("VariableReserve.jl")
include("Transfer.jl")
include("AVRFixed.jl")
include("AVRSimple.jl")
include("AVRTypeI.jl")
include("AVRTypeII.jl")
include("BaseMachine.jl")
include("AndersonFouadMachine.jl")
include("FullMachine.jl")
include("MarconatoMachine.jl")
include("OneDOneQMachine.jl")
include("SimpleAFMachine.jl")
include("SimpleFullMachine.jl")
include("SimpleMarconatoMachine.jl")
include("PSSFixed.jl")
include("PSSSimple.jl")
include("SingleMass.jl")
include("FiveMassShaft.jl")
include("TGFixed.jl")
include("TGTypeI.jl")
include("TGTypeII.jl")
include("AvgCnvFixedDC.jl")
include("FixedDCSource.jl")
include("LCLFilter.jl")
include("LCFilter.jl")
include("PLL.jl")
include("VirtualInertia.jl")
include("ReactivePowerDroop.jl")
include("CombinedVIwithVZ.jl")
include("Source.jl")

export get_Ae
export get_Be
export get_D
export get_D_12
export get_D_23
export get_D_34
export get_D_45
export get_D_ex
export get_D_hp
export get_D_ip
export get_D_lp
export get_Emf
export get_H
export get_H_ex
export get_H_hp
export get_H_ip
export get_H_lp
export get_K0
export get_K_ex
export get_K_hp
export get_K_ip
export get_K_lp
export get_K_p
export get_K_ω
export get_Ka
export get_Ke
export get_Kf
export get_Kv
export get_L_1d
export get_L_1q
export get_L_ad
export get_L_aq
export get_L_d
export get_L_f1d
export get_L_ff
export get_L_q
export get_MVABase
export get_P_max
export get_P_min
export get_R
export get_R_1d
export get_R_1q
export get_R_f
export get_T1
export get_T2
export get_T3
export get_T4
export get_T5
export get_T_AA
export get_Ta
export get_Tc
export get_Td0_p
export get_Td0_pp
export get_Te
export get_Tf
export get_Tq0_p
export get_Tq0_pp
export get_Tr
export get_Ts
export get_V_I
export get_V_R
export get_V_pss
export get_Vr_max
export get_Vr_min
export get_X_th
export get_Xd
export get_Xd_p
export get_Xd_pp
export get_Xq
export get_Xq_p
export get_Xq_pp
export get_Y
export get_activepower
export get_activepower_flow
export get_activepowerlimits
export get_activepowerlimits_from
export get_activepowerlimits_to
export get_angle
export get_anglelimits
export get_arc
export get_area
export get_available
export get_b
export get_basevoltage
export get_bus
export get_bustype
export get_capacity
export get_cf
export get_efficiency
export get_energy
export get_eq_p
export get_ext
export get_fixed
export get_flowlimits
export get_from
export get_fuel
export get_inflow
export get_initial_storage
export get_inputactivepowerlimits
export get_internal
export get_inv_d_fluxlink
export get_inv_q_fluxlink
export get_inverter_firing_angle
export get_inverter_taplimits
export get_inverter_xrc
export get_kad
export get_kd
export get_kffi
export get_kffv
export get_ki_pll
export get_kic
export get_kiv
export get_kp_pll
export get_kpc
export get_kpv
export get_kq
export get_kω
export get_lf
export get_lg
export get_load_zone
export get_loss
export get_lv
export get_maxactivepower
export get_maxreactivepower
export get_model
export get_n_states
export get_number
export get_op_cost
export get_outputactivepowerlimits
export get_powerfactor
export get_primaryshunt
export get_primemover
export get_r
export get_ramplimits
export get_rate
export get_rating
export get_reactivepower
export get_reactivepower_flow
export get_reactivepowerlimits
export get_reactivepowerlimits_from
export get_reactivepowerlimits_to
export get_rectifier_firing_angle
export get_rectifier_taplimits
export get_rectifier_xrc
export get_requirement
export get_rf
export get_rg
export get_rv
export get_s_rated
export get_services
export get_shutdn
export get_startup
export get_states
export get_storage_capacity
export get_tap
export get_tech
export get_timeframe
export get_timelimits
export get_to
export get_v_rated
export get_variable
export get_voltage
export get_voltagelimits
export get_x
export get_α
export get_γd
export get_γq
export get_τ_max
export get_τ_min
export get_ω_lp
export get_ωad
export get_ωb
export get_ωf


using Crazyflie
using ProgressMeter

# temporary
using PyCall
using DelimitedFiles

figure8 = [
    [1.050000, 0.000000, -0.000000, 0.000000, -0.000000, 0.830443, -0.276140, -0.384219, 0.180493, -0.000000, 0.000000, -0.000000, 0.000000, -1.356107, 0.688430, 0.587426, -0.329106, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000],  # noqa
    [0.710000, 0.396058, 0.918033, 0.128965, -0.773546, 0.339704, 0.034310, -0.026417, -0.030049, -0.445604, -0.684403, 0.888433, 1.493630, -1.361618, -0.139316, 0.158875, 0.095799, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000],  # noqa
    [0.620000, 0.922409, 0.405715, -0.582968, -0.092188, -0.114670, 0.101046, 0.075834, -0.037926, -0.291165, 0.967514, 0.421451, -1.086348, 0.545211, 0.030109, -0.050046, -0.068177, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000],  # noqa
    [0.700000, 0.923174, -0.431533, -0.682975, 0.177173, 0.319468, -0.043852, -0.111269, 0.023166, 0.289869, 0.724722, -0.512011, -0.209623, -0.218710, 0.108797, 0.128756, -0.055461, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000],  # noqa
    [0.560000, 0.405364, -0.834716, 0.158939, 0.288175, -0.373738, -0.054995, 0.036090, 0.078627, 0.450742, -0.385534, -0.954089, 0.128288, 0.442620, 0.055630, -0.060142, -0.076163, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000],  # noqa
    [0.560000, 0.001062, -0.646270, -0.012560, -0.324065, 0.125327, 0.119738, 0.034567, -0.063130, 0.001593, -1.031457, 0.015159, 0.820816, -0.152665, -0.130729, -0.045679, 0.080444, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000],  # noqa
    [0.700000, -0.402804, -0.820508, -0.132914, 0.236278, 0.235164, -0.053551, -0.088687, 0.031253, -0.449354, -0.411507, 0.902946, 0.185335, -0.239125, -0.041696, 0.016857, 0.016709, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000],  # noqa
    [0.620000, -0.921641, -0.464596, 0.661875, 0.286582, -0.228921, -0.051987, 0.004669, 0.038463, -0.292459, 0.777682, 0.565788, -0.432472, -0.060568, -0.082048, -0.009439, 0.041158, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000],  # noqa
    [0.710000, -0.923935, 0.447832, 0.627381, -0.259808, -0.042325, -0.032258, 0.001420, 0.005294, 0.288570, 0.873350, -0.515586, -0.730207, -0.026023, 0.288755, 0.215678, -0.148061, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000],  # noqa
    [1.053185, -0.398611, 0.850510, -0.144007, -0.485368, -0.079781, 0.176330, 0.234482, -0.153567, 0.447039, -0.532729, -0.855023, 0.878509, 0.775168, -0.391051, -0.713519, 0.391628, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000],  # noqa
]

heart = [
[0.608998,0,0,0,0,8.42094,-20.1848,17.4023,-5.13733,0,0,0,0,-5.25294,19.5805,-26.0513,12.0229,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],
[0.779122,0.195632,0.637789,0.0169511,-0.469071,1.13898,-5.83124,8.3787,-3.63444,-0.0377946,-0.0794284,-0.0646026,0.0244159,-3.78373,9.99798,-9.63602,3.27284,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],
[0.757723,0.467339,-0.036739,-0.110484,0.0639573,0.108074,-1.39392,2.11362,-0.935232,-0.236257,-0.347344,-0.032209,-0.0141377,-1.22068,4.0166,-4.3725,1.62188,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],
[0.957797,0.357255,-0.274841,-0.111489,-0.0115323,-3.42962,8.38992,-7.1466,2.09672,-0.518182,-0.328702,0.0913417,0.0150219,0.733687,-1.05708,0.479298,-0.0521867,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],
[0.95006,-0.108941,-0.502428,-0.0278039,-0.0132172,-2.03991,5.51081,-5.00167,1.53592,-0.639182,0.0542293,0.051927,0.014188,3.04728,-7.84289,6.93303,-2.09398,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],
[0.998431,-0.624157,-0.461845,-0.0138318,-0.00967601,-2.04884,5.5933,-4.99039,1.48711,-0.481169,0.114548,0.029441,0.000554473,3.31264,-7.6003,6.14741,-1.71994,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],
[0.812602,-1.06789,-0.279874,-0.0292519,-0.0503488,-3.7103,10.0379,-9.64748,3.24897,-0.197235,0.269061,-0.00363612,0.0151035,0.916496,-3.79814,4.50747,-1.72174,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],
[0.419508,-1.42038,-0.468108,0.0856745,0.221818,3.92471,-0.277235,-23.883,26.7762,-0.0240585,0.0921,-0.0459625,0.0111955,-3.86908,19.5572,-35.3542,22.4094,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],
[0.64054,-1.5363,-0.00325354,0.514808,0.0803983,-2.16099,9.25223,-14.0528,6.98868,0.000127304,0.0443542,0.0111982,0.0261656,1.57639,-5.36123,6.53792,-2.78901,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],
[0.893144,-1.33358,0.558221,0.00558749,-0.182813,1.17326,-3.86353,4.22095,-1.49454,0.05546,0.116935,0.0458194,0.000332267,2.56277,-6.40182,5.67683,-1.7497,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],
[1.0137,-0.944994,0.264934,0.0564772,0.0488614,3.09438,-7.21848,5.83434,-1.62251,0.277414,0.275656,0.00454676,0.0105108,1.06699,-2.79915,2.41189,-0.698664,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],
[0.941904,-0.480716,0.458841,0.0299512,0.00400315,2.13691,-5.30337,4.62042,-1.38829,0.551501,0.16792,-0.0455683,-0.000158787,-0.00812595,-0.31599,0.457477,-0.175203,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],
[0.902676,0.0448938,0.563892,0.0356819,-0.0204628,3.81762,-12.3689,12.5856,-4.21939,0.632664,-0.0105303,-0.0563382,-0.0149921,-4.47262,11.1707,-9.8176,2.99179,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],
[0.696637,0.437821,0.0271313,-0.0955769,0.0683252,-3.88281,13.7328,-16.7036,6.91373,0.441299,-0.197977,0.0795611,-0.0410307,-17.0929,56.5153,-65.8397,26.5171,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],
[0.790226,0.413439,0.0024842,-0.0197612,-0.0573184,-9.94595,26.5283,-25.7537,8.8682,0.160997,-0.308277,0.0562815,0.00236697,1.27222,-2.73425,2.12724,-0.57832,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],
[0.418106,0.106269,-0.544119,0.234554,1.02684,3.93774,-3.17014,-30.7186,40.0474,0.0139615,-0.0875489,0.0363838,0.0507818,5.39905,-25.4338,43.0935,-25.8168,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]
]

py"""
import time
from cflib.crazyflie.mem import MemoryElement
from cflib.crazyflie.mem import Poly4D

class Uploader:
    def __init__(self):
        self._is_done = False

    def upload(self, trajectory_mem):
        print('Uploading data')
        trajectory_mem.write_data(self._upload_done)

        while not self._is_done:
            time.sleep(0.2)

    def _upload_done(self, mem, addr):
        print('Data uploaded')
        self._is_done = True

def upload_trajectory(cf, trajectory_id, trajectory):
    trajectory_mem = cf.mem.get_mems(MemoryElement.TYPE_TRAJ)[0]

    total_duration = 0
    for row in trajectory:
        duration = row[0]
        x = Poly4D.Poly(row[1:9])
        y = Poly4D.Poly(row[9:17])
        z = Poly4D.Poly(row[17:25])
        yaw = Poly4D.Poly(row[25:33])
        trajectory_mem.poly4Ds.append(Poly4D(duration, x, y, z, yaw))
        total_duration += duration

    Uploader().upload(trajectory_mem)
    cf.high_level_commander.define_trajectory(trajectory_id, 0,
                                              len(trajectory_mem.poly4Ds))
    return total_duration
"""

function reset_estimator(cf)
    cf.param.set_value("kalman.resetEstimation", "1")
    sleep(0.1)
    cf.param.set_value("kalman.resetEstimation", "0")
    lconfig = logger.LogConfig(name="Kalman Variance", period_in_ms=500)
    lconfig.add_variable("kalman.varPX", "float")
    lconfig.add_variable("kalman.varPY", "float")
    lconfig.add_variable("kalman.varPZ", "float")
    varx = ones(10)*1000
    vary = ones(10)*1000
    varz = ones(10)*1000
    record(cf, lconfig) do logs
        for entry in logs
            data = entry[2]
            push!(varx, data["kalman.varPX"])
            push!(vary, data["kalman.varPY"])
            push!(varz, data["kalman.varPZ"])
            popfirst!(varx); popfirst!(vary); popfirst!(varz)
            threshx = (maximum(varx) - minimum(varx)) < 0.001
            threshy = (maximum(vary) - minimum(vary)) < 0.001
            threshz = (maximum(varz) - minimum(varz)) < 0.001
            if threshx && threshy && threshz
                break
            end
        end
    end
end

function enable_commander(cf)
    cf.param.set_value("commander.enHighLevel", "1")
    cf.high_level_commander
end

# because commander.land is not reliable
function land(commander, duration)
    print("Landing\n")
    for i = 1:10
        commander.land(0.0, duration)
        sleep(0.1)
    end
    sleep(duration)
    commander.stop()
end

# because commander.takeoff is not reliable
function takeoff(commander, height, duration)
    print("Taking Off\n")
    for i = 1:5
        commander.takeoff(height, duration)
        sleep(0.1)
    end
    sleep(duration)
end

# because commander.start_trajectory is not reliable
function start_trajectory(commander, id, rel)
    print("Starting Trajectory\n")
    for i = 1:5
        commander.start_trajectory(id, relative=rel)
        sleep(0.1)
    end
    sleep(1.0)
end

function traj(uri="radio://0/80/2M")
    play(uri) do cf
        commander = enable_commander(cf)
        duration = py"upload_trajectory($cf, $1, $heart)"
        reset_estimator(cf)
        takeoff(commander, 1.0, 2.0) # height, duration
        start_trajectory(commander, 1, true)
        sleep(duration)
        land(commander, 5.0) # duration
    end
    return nothing
end


#///////////////////////////////////////
#// File Name: crowd_nav_controller_test.jl
#// Author: Haruki Nishimura (hnishimura@stanford.edu)
#// Date Created: 2021/04/08
#// Description: Test code for src/crowd_nav_controller.jl
#///////////////////////////////////////

using LinearAlgebra
using PyCall

@testset "Crowd Nav Controller Test" begin
    Cep = Matrix(0.5I, 2, 2);                                                           # quadratic position cost matrix
    Cu = Matrix(0.2I, 2, 2);                                                            # quadratic control cost matrix
    β_pos = 0.1;                                                                        # relative weight between instant and terminal pos cost
    α_col = 100.0;                                                                      # scale parameter for exponential collision cost
    β_col = 0.1;                                                                        # relative weight between instant and terminal pos cost
    λ_col = 0.2;                                                                        # bandwidth parameter for exponential collision cost
    σ_risk = 0.0;                                                                       # risk-sensitiveness parameter

    cost_param = CostParameter(Cep, Cu, β_pos, α_col, β_col, λ_col, σ_risk);
    sim_param = SimulationParameter(0.01, 0.4, 12, 10, cost_param);

    tcalc = 0.1;
    dtr = 0.4;
    target_speed = 0.7;
    model_dir = normpath(joinpath(@__DIR__,
                                  "../CrowdNav/crowd_nav/data/output_om_sarl_radius_0.4"))
    cnt_param = CrowdNavControlParameter(model_dir,
                                         "env.config", "policy.config", "sarl",
                                         [0.0, 4.0], tcalc, dtr, target_speed);

    crowd_nav_controller = CrowdNavController(sim_param, cnt_param);
    @test !crowd_nav_controller.rl_robot.visible
    @test crowd_nav_controller.rl_robot.policy.time_step == 0.4;
    @test crowd_nav_controller.rl_robot.v_pref == 0.7;
    @test crowd_nav_controller.rl_robot.time_step == 0.4;
    @test crowd_nav_controller.rl_robot.gx == 0.0;
    @test crowd_nav_controller.rl_robot.gy == 4.0;

    e_init = RobotState([0.0, -4.0, 0.0, 0.0]);
    u_vel = get_action!(crowd_nav_controller, e_init,
                        Dict{String, Vector{Float64}}());
    @test u_vel ≈ [0.0, 8.0]./norm([0.0, 8.0])*crowd_nav_controller.rl_robot.v_pref

    e_init = RobotState([0.0, 0.0, 0.0, 0.0]);
    ado_state_dict = Dict{String, Vector{Float64}}();
    ado_state_dict["PEDESTRIAN/1"] = [5.0, 4.0, -1.0, -3.0];
    ado_state_dict["PEDESTRIAN/2"] = [-3.0, 0.0, 1.0, 1.0];

    u_vel = get_action!(crowd_nav_controller, e_init, ado_state_dict)

    ~, u_acc = crowdnav_control_update!(crowd_nav_controller, e_init, ado_state_dict)
    #@test all(u_vel.*0.4 .≈ get_velocity(e_init).*0.4 .+ 0.5*u_acc.*(0.4)^2)
    @test all(u_vel .≈ get_velocity(e_init) .+ u_acc.*0.4)
end


@testset "DC Power Flow Models Monitored Line Flow Constraints and Static Unbounded" begin
    system = PSB.build_system(PSITestSystems, "c_sys5_ml")
    limits = PSY.get_flow_limits(PSY.get_component(MonitoredLine, system, "1"))
    for model in [DCPPowerModel, StandardPTDFModel]
        template = get_thermal_dispatch_template_network(model)
        op_problem_m = OperationsProblem(
            template,
            system;
            optimizer = OSQP_optimizer,
            PTDF = PSY.PTDF(system),
        )
        @test build!(op_problem_m; output_dir = mktempdir(cleanup = true)) ==
              PSI.BuildStatus.BUILT

        @test check_variable_bounded(op_problem_m, :Fp__MonitoredLine)
        @test check_variable_unbounded(op_problem_m, :Fp__Line)

        @test solve!(op_problem_m) == RunStatus.SUCCESSFUL
        @test check_flow_variable_values(
            op_problem_m,
            :Fp__MonitoredLine,
            "1",
            limits.from_to,
        )
    end
end

@testset "DC Power Flow Models Monitored Line Flow Constraints and Static with Bounds" begin
    system = PSB.build_system(PSITestSystems, "c_sys5_ml")
    set_rate!(PSY.get_component(Line, system, "2"), 1.5)
    for model in [DCPPowerModel, StandardPTDFModel]
        template = get_thermal_dispatch_template_network(model)
        set_device_model!(template, DeviceModel(Line, StaticBranch))
        set_device_model!(template, DeviceModel(MonitoredLine, StaticBranchUnbounded))
        op_problem_m = OperationsProblem(
            template,
            system;
            optimizer = OSQP_optimizer,
            PTDF = PSY.PTDF(system),
        )
        @test build!(op_problem_m; output_dir = mktempdir(cleanup = true)) ==
              PSI.BuildStatus.BUILT

        @test check_variable_unbounded(op_problem_m, :Fp__MonitoredLine)
        # Broken
        # @test check_variable_bounded(op_problem_m, :Fp__Line)

        @test solve!(op_problem_m) == RunStatus.SUCCESSFUL
        @test check_flow_variable_values(op_problem_m, :Fp__Line, "2", 1.5)
    end
end

@testset "AC Power Flow Monitored Line Flow Constraints" begin
    system = PSB.build_system(PSITestSystems, "c_sys5_ml")
    limits = PSY.get_flow_limits(PSY.get_component(MonitoredLine, system, "1"))
    template = get_thermal_dispatch_template_network(ACPPowerModel)
    op_problem_m = OperationsProblem(template, system; optimizer = ipopt_optimizer)
    @test build!(op_problem_m; output_dir = mktempdir(cleanup = true)) ==
          PSI.BuildStatus.BUILT

    @test check_variable_bounded(op_problem_m, :FpFT__MonitoredLine)
    @test check_variable_unbounded(op_problem_m, :FqFT__MonitoredLine)

    @test solve!(op_problem_m) == RunStatus.SUCCESSFUL
    @test check_flow_variable_values(
        op_problem_m,
        :FpFT__MonitoredLine,
        :FqFT__MonitoredLine,
        "1",
        0.0,
        limits.from_to,
    )
end

@testset "DC Power Flow Models for HVDCLine with with Line Flow Constraints, TapTransformer & Transformer2W Unbounded" begin
    ratelimit_constraint_names = [
        :RateLimit_ub__Transformer2W,
        :RateLimit_lb__Transformer2W,
        :RateLimit_ub__TapTransformer,
        :RateLimit_lb__TapTransformer,
    ]

    system = PSB.build_system(PSITestSystems, "c_sys14_dc")
    hvdc_line = PSY.get_component(HVDCLine, system, "DCLine3")
    limits_from = PSY.get_active_power_limits_from(hvdc_line)
    limits_to = PSY.get_active_power_limits_to(hvdc_line)
    limits_min = min(limits_from.min, limits_to.min)
    limits_max = min(limits_from.max, limits_to.max)

    tap_transformer = PSY.get_component(TapTransformer, system, "Trans3")
    rate_limit = PSY.get_rate(tap_transformer)

    transformer = PSY.get_component(Transformer2W, system, "Trans4")
    rate_limit2w = PSY.get_rate(tap_transformer)

    for model in [DCPPowerModel, StandardPTDFModel],
        hvdc_model in [HVDCDispatch, HVDCLossless]

        template = get_template_dispatch_with_network(model)
        set_device_model!(template, HVDCLine, hvdc_model)
        op_problem_m = OperationsProblem(
            template,
            system;
            optimizer = OSQP_optimizer,
            PTDF = PSY.PTDF(system),
        )
        @test build!(op_problem_m; output_dir = mktempdir(cleanup = true)) ==
              PSI.BuildStatus.BUILT

        @test check_variable_bounded(op_problem_m, :Fp__HVDCLine)
        @test check_variable_unbounded(op_problem_m, :Fp__TapTransformer)
        @test check_variable_unbounded(op_problem_m, :Fp__Transformer2W)

        psi_constraint_test(op_problem_m, ratelimit_constraint_names)

        @test solve!(op_problem_m) == RunStatus.SUCCESSFUL

        @test check_flow_variable_values(
            op_problem_m,
            :Fp__HVDCLine,
            "DCLine3",
            limits_min,
            limits_max,
        )
        @test check_flow_variable_values(
            op_problem_m,
            :Fp__TapTransformer,
            "Trans3",
            -rate_limit,
            rate_limit,
        )
        @test check_flow_variable_values(
            op_problem_m,
            :Fp__Transformer2W,
            "Trans4",
            -rate_limit2w,
            rate_limit2w,
        )
    end
end

@testset "DC Power Flow Models for Unbounded HVDCLine, and StaticBranchBounds for TapTransformer & Transformer2W" begin
    system = PSB.build_system(PSITestSystems, "c_sys14_dc")
    hvdc_line = PSY.get_component(HVDCLine, system, "DCLine3")
    limits_from = PSY.get_active_power_limits_from(hvdc_line)
    limits_to = PSY.get_active_power_limits_to(hvdc_line)
    limits_min = min(limits_from.min, limits_to.min)
    limits_max = min(limits_from.max, limits_to.max)

    tap_transformer = PSY.get_component(TapTransformer, system, "Trans3")
    rate_limit = PSY.get_rate(tap_transformer)

    transformer = PSY.get_component(Transformer2W, system, "Trans4")
    rate_limit2w = PSY.get_rate(tap_transformer)

    for model in [DCPPowerModel, StandardPTDFModel]
        template = get_template_dispatch_with_network(model)
        set_device_model!(template, DeviceModel(HVDCLine, HVDCUnbounded))
        set_device_model!(template, DeviceModel(TapTransformer, StaticBranchBounds))
        set_device_model!(template, DeviceModel(Transformer2W, StaticBranchBounds))
        op_problem_m = OperationsProblem(
            template,
            system;
            optimizer = OSQP_optimizer,
            PTDF = PSY.PTDF(system),
        )
        @test build!(op_problem_m; output_dir = mktempdir(cleanup = true)) ==
              PSI.BuildStatus.BUILT

        if model == DCPPowerModel
            # TODO: Currently Broken, remove variable bounds in HVDCUnbounded
            # @test check_variable_unbounded(op_problem_m, :Fp__HVDCLine)

            @test check_variable_bounded(op_problem_m, :Fp__TapTransformer)
            @test check_variable_bounded(op_problem_m, :Fp__TapTransformer)
        end

        @test solve!(op_problem_m) == RunStatus.SUCCESSFUL

        @test check_flow_variable_values(
            op_problem_m,
            :Fp__HVDCLine,
            "DCLine3",
            limits_min,
            limits_max,
        )
        @test check_flow_variable_values(
            op_problem_m,
            :Fp__TapTransformer,
            "Trans3",
            -rate_limit,
            rate_limit,
        )
        @test check_flow_variable_values(
            op_problem_m,
            :Fp__Transformer2W,
            "Trans4",
            -rate_limit2w,
            rate_limit2w,
        )
    end
end

@testset "AC Power Flow Models for HVDCLine Flow Constraints and TapTransformer & Transformer2W Unbounded" begin
    ratelimit_constraint_names = [
        :RateLimitFT__Transformer2W,
        :RateLimitTF__Transformer2W,
        :RateLimitFT__TapTransformer,
        :RateLimitTF__TapTransformer,
        :RateLimitFT__HVDCLine,
        :RateLimitTF__HVDCLine,
    ]

    system = PSB.build_system(PSITestSystems, "c_sys14_dc")

    hvdc_line = PSY.get_component(HVDCLine, system, "DCLine3")
    limits_from = PSY.get_active_power_limits_from(hvdc_line)
    limits_to = PSY.get_active_power_limits_to(hvdc_line)
    limits_min = min(limits_from.min, limits_to.min)
    limits_max = min(limits_from.max, limits_to.max)

    tap_transformer = PSY.get_component(TapTransformer, system, "Trans3")
    rate_limit = PSY.get_rate(tap_transformer)

    transformer = PSY.get_component(Transformer2W, system, "Trans4")
    rate_limit2w = PSY.get_rate(tap_transformer)

    template = get_template_dispatch_with_network(ACPPowerModel)
    set_device_model!(template, DeviceModel(HVDCLine, HVDCDispatch))
    op_problem_m = OperationsProblem(template, system; optimizer = ipopt_optimizer)
    @test build!(op_problem_m; output_dir = mktempdir(cleanup = true)) ==
          PSI.BuildStatus.BUILT

    check_variable_bounded(op_problem_m, :FqTF__HVDCLine)
    check_variable_bounded(op_problem_m, :Fp__HVDCLine)
    @test check_variable_bounded(op_problem_m, :FpFT__TapTransformer)
    @test check_variable_unbounded(op_problem_m, :FqFT__TapTransformer)
    @test check_variable_bounded(op_problem_m, :FpTF__Transformer2W)
    @test check_variable_unbounded(op_problem_m, :FqTF__Transformer2W)

    psi_constraint_test(op_problem_m, ratelimit_constraint_names)

    @test solve!(op_problem_m) == RunStatus.SUCCESSFUL

    @test check_flow_variable_values(
        op_problem_m,
        :Fp__HVDCLine,
        :FqTF__HVDCLine,
        "DCLine3",
        limits_max,
    )
    @test check_flow_variable_values(
        op_problem_m,
        :FpFT__TapTransformer,
        :FqFT__TapTransformer,
        "Trans3",
        rate_limit,
    )
    @test check_flow_variable_values(
        op_problem_m,
        :FpTF__Transformer2W,
        :FqTF__Transformer2W,
        "Trans4",
        rate_limit2w,
    )
end

#= This Test is Broken to do a missing implementation of add_variable
@testset "DC Power Flow Models for HVDCLine Dispatch and TapTransformer & Transformer2W Unbounded" begin
    ratelimit_constraint_names = [
        :RateLimitFT__Transformer2W,
        :RateLimitTF__Transformer2W,
        :RateLimitFT__TapTransformer,
        :RateLimitTF__TapTransformer,
        :RateLimitFT__HVDCLine,
        :RateLimitTF__HVDCLine,
    ]

    system = PSB.build_system(PSITestSystems, "c_sys14_dc")

    hvdc_line = PSY.get_component(HVDCLine, system, "DCLine3")
    limits_from = PSY.get_active_power_limits_from(hvdc_line)
    limits_to = PSY.get_active_power_limits_to(hvdc_line)
    limits_min = min(limits_from.min, limits_to.min)
    limits_max = min(limits_from.max, limits_to.max)

    tap_transformer = PSY.get_component(TapTransformer, system, "Trans3")
    rate_limit = PSY.get_rate(tap_transformer)

    transformer = PSY.get_component(Transformer2W, system, "Trans4")
    rate_limit2w = PSY.get_rate(tap_transformer)

    template = get_template_dispatch_with_network(StandardPTDFModel)
    set_device_model!(template, DeviceModel(HVDCLine, HVDCDispatch))
    op_problem_m = OperationsProblem(template, system; PTDF = PSY.PTDF(system), optimizer = ipopt_optimizer)
    @test build!(op_problem_m; output_dir = mktempdir(cleanup = true)) ==
          PSI.BuildStatus.BUILT

    check_variable_bounded(op_problem_m, :FqTF__HVDCLine)
    check_variable_bounded(op_problem_m, :FpTF__HVDCLine)
    @test check_variable_bounded(op_problem_m, :FpFT__TapTransformer)
    @test check_variable_unbounded(op_problem_m, :FqFT__TapTransformer)
    @test check_variable_bounded(op_problem_m, :FpTF__Transformer2W)
    @test check_variable_unbounded(op_problem_m, :FqTF__Transformer2W)

    psi_constraint_test(op_problem_m, ratelimit_constraint_names)

    @test solve!(op_problem_m) == RunStatus.SUCCESSFUL

    @test check_flow_variable_values(
        op_problem_m,
        :FpTF__HVDCLine,
        :FqTF__HVDCLine,
        "DCLine3",
        limits_max,
    )
    @test check_flow_variable_values(
        op_problem_m,
        :FpFT__TapTransformer,
        :FqFT__TapTransformer,
        "Trans3",
        rate_limit,
    )
    @test check_flow_variable_values(
        op_problem_m,
        :FpTF__Transformer2W,
        :FqTF__Transformer2W,
        "Trans4",
        rate_limit2w,
    )
end
=#