using .CuArrays.CUDNN: @check, libcudnn, cudnnStatus_t, libcudnn_handle,
  cudnnDataType, TensorDesc, FilterDesc

using LinearAlgebra

mutable struct DropoutDesc
  ptr::Ptr{Nothing}
  states::CuVector{UInt8}
end

Base.unsafe_convert(::Type{Ptr{Nothing}}, dd::DropoutDesc) = dd.ptr

function DropoutDesc(ρ::Real; seed::Integer=0)
  d = [C_NULL]
  s = Csize_t[0]
  @check ccall((:cudnnCreateDropoutDescriptor,libcudnn), cudnnStatus_t, (Ptr{Ptr{Nothing}},), d)
  @check ccall((:cudnnDropoutGetStatesSize,libcudnn),cudnnStatus_t,(Ptr{Nothing},Ptr{Csize_t}),libcudnn_handle[],s)
  states = CuArray{UInt8}(s[]) # TODO: can we drop this when ρ=0?
  desc = DropoutDesc(d[], states)
  @check ccall((:cudnnSetDropoutDescriptor,libcudnn),cudnnStatus_t,(Ptr{Nothing},Ptr{Nothing},Cfloat,Ptr{Nothing},Csize_t,Culonglong),
    desc,libcudnn_handle[],ρ,states,length(states),seed)
  finalizer(desc) do x
    @check ccall((:cudnnDestroyDropoutDescriptor,libcudnn),cudnnStatus_t,(Ptr{Nothing},),x)
  end
  return desc
end

const RNN_RELU = 0 # Stock RNN with ReLu activation
const RNN_TANH = 1 # Stock RNN with tanh activation
const LSTM = 2     # LSTM with no peephole connections
const GRU = 3      # Using h' = tanh(r * Uh(t-1) + Wx) and h = (1 - z) * h' + z * h(t-1)

const LINEAR_INPUT = 0
const SKIP_INPUT = 1

const UNIDIRECTIONAL = 0
const BIDIRECTIONAL = 1

const RNN_ALGO_STANDARD = 0
const RNN_ALGO_PERSIST_STATIC = 1
const RNN_ALGO_PERSIST_DYNAMIC = 2

# param layout:
# RNN: [weight, bias] × [input, hidden]
# GRU: [weight, bias] × [input, hidden] × [reset, update, newmem]
# LSTM: [weight, bias] × [input, hidden] × [input, forget, newmem, output]

function params(w::CuVector, input, hidden, n = 1)
  slice(offset, shape) = reshape(view(w, offset.+(1:prod(shape))), shape)
  wx = slice(0, (input, hidden*n))
  wh = slice(length(wx), (hidden, hidden*n))
  bias = view(w, length(wx)+length(wh) .+ (1:hidden*n))
  (wx, wh), bias
end

mutable struct RNNDesc{T}
  mode::Int
  input::Int
  hidden::Int
  params::CuVector{T}
  weights::NTuple{2,CuMatrix{T}}
  bias::CuVector{T}
  ptr::Ptr{Nothing}
end

Base.unsafe_convert(::Type{Ptr{Nothing}}, d::RNNDesc) = d.ptr

function rnnParamSize(T, r, input)
  size = Csize_t[0]
  @check ccall((:cudnnGetRNNParamsSize, libcudnn), cudnnStatus_t, (Ptr{Nothing},Ptr{Nothing},Ptr{Nothing},Ptr{Csize_t},Cint),
    libcudnn_handle[], r, TensorDesc(T, (1,input,1)), size, cudnnDataType(T))
  return Int(size[])÷sizeof(T)
end

ngates(mode) = [1, 1, 4, 3][mode+1]
ngates(r::RNNDesc) = ngates(r.mode)

function RNNDesc{T}(mode::Int, input::Int, hidden::Int; layers = 1) where T
  d = [C_NULL]
  @check ccall((:cudnnCreateRNNDescriptor,libcudnn),cudnnStatus_t,(Ptr{Ptr{Nothing}},),d)

  dropoutDesc = DropoutDesc(0)
  inputMode = LINEAR_INPUT
  direction = UNIDIRECTIONAL
  algo = RNN_ALGO_STANDARD
  @check ccall((:cudnnSetRNNDescriptor_v6,libcudnn), cudnnStatus_t, (Ptr{Nothing},Ptr{Nothing},Cint,Cint,Ptr{Nothing},Cint,Cint,Cint,Cint,Cint),
    libcudnn_handle[],d[],hidden,layers,dropoutDesc,inputMode,direction,mode,algo,cudnnDataType(T))

  w = cuzeros(T, rnnParamSize(T, d[], input))
  # TODO: avoid reserve allocation here
  rd = RNNDesc{T}(mode, input, hidden, w, params(w, input, hidden, ngates(mode))..., d[])
  finalizer(rd) do x
    @check ccall((:cudnnDestroyRNNDescriptor,libcudnn),cudnnStatus_t,(Ptr{Nothing},),x)
  end
  return rd
end

function rnnWorkspaceSize(r::RNNDesc, seqlen, xdesc)
  size = Csize_t[0]
  @check ccall((:cudnnGetRNNWorkspaceSize, libcudnn), cudnnStatus_t, (Ptr{Nothing},Ptr{Nothing},Cint,Ptr{Ptr{Nothing}},Ptr{Csize_t}),
    libcudnn_handle[], r, seqlen, xdesc, size)
  return Int(size[])
end

const workspace = [CuVector{UInt8}(1)]

getworkspace(bytes) =
  length(workspace[]) ≥ bytes ?
    workspace[] :
    (workspace[] = CuVector{UInt8}(bytes))

getworkspace(r::RNNDesc, seqlen, xdesc) =
  getworkspace(rnnWorkspaceSize(r, seqlen, xdesc))

function rnnTrainingReserveSize(r::RNNDesc, seqlen, xdesc)
  size = Csize_t[0]
  @check ccall((:cudnnGetRNNTrainingReserveSize,libcudnn), cudnnStatus_t, (Ptr{Nothing}, Ptr{Nothing}, Cint, Ptr{Ptr{Nothing}}, Ptr{Csize_t}),
    libcudnn_handle[], r, seqlen, xdesc, size)
  return Int(size[])
end

function cudnnRNNForward(rnn::RNNDesc{T}, seqlen, xd, x, hd, h, cd, c, wd, w, yd, y, hod, ho, cod, co,
                         workspace, reserve=nothing) where T
  if reserve == nothing
    @check ccall((:cudnnRNNForwardInference, libcudnn), cudnnStatus_t,
                 (Ptr{Nothing}, Ptr{Nothing}, Cint,
                  Ptr{Ptr{Nothing}}, Ptr{T}, Ptr{Nothing}, Ptr{T}, Ptr{Nothing}, Ptr{T},
                  Ptr{Nothing}, Ptr{T}, Ptr{Ptr{Nothing}}, Ptr{T}, Ptr{Nothing}, Ptr{T},
                  Ptr{Nothing}, Ptr{T},
                  Ptr{Nothing}, Csize_t),
                 libcudnn_handle[], rnn, seqlen,
                 xd, x, hd, h, cd, c, wd, w, yd, y, hod, ho, cod, co,
                 workspace, length(workspace))
  else
    @check ccall((:cudnnRNNForwardTraining, libcudnn), cudnnStatus_t,
                 (Ptr{Nothing}, Ptr{Nothing}, Cint,
                  Ptr{Ptr{Nothing}}, Ptr{T}, Ptr{Nothing}, Ptr{T}, Ptr{Nothing}, Ptr{T}, Ptr{Nothing}, Ptr{T}, Ptr{Ptr{Nothing}}, Ptr{T}, Ptr{Nothing}, Ptr{T}, Ptr{Nothing}, Ptr{T},
                  Ptr{Nothing}, Csize_t, Ptr{Nothing}, Csize_t),
                 libcudnn_handle[], rnn, seqlen,
                 xd, x, hd, h, cd, c, wd, w, yd, y, hod, ho, cod, co,
                 workspace, length(workspace), reserve, length(reserve))
  end
end

xDesc(x) = [TensorDesc(eltype(x), (1, size(x, 1), size(x, 2)))]

hDesc(h::Nothing) = C_NULL, C_NULL
hDesc(x::Integer) = (@assert x == 0; hDesc(nothing))
function hDesc(h::CuArray)
  TensorDesc(eltype(h), (size(h, 1), size(h, 2), 1)), h
end

# TODO: can we just manipulate strides here?
# TODO: should use repmat, but this isn't implemented.
hBatch(x::AbstractVector, h::CuVector) = h
hBatch(x::AbstractMatrix, h::CuVector) = h .* cuones(1, size(x, 2))
hBatch(x::AbstractMatrix, h::CuMatrix) = h .* cuones(1, size(h,2) == 1 ? size(x,2) : 1)

function forward(rnn::RNNDesc{T}, x::CuArray{T}, h_::CuArray{T}, c_ = nothing, train = Val{false}) where T
  h = hBatch(x, h_)
  c = c_ == nothing ? nothing : hBatch(x, c_)
  @assert size(x, 1) == rnn.input
  @assert size(h, 1) == rnn.hidden
  @assert size(x, 2) == size(h, 2)
  seqLength = 1
  xdesc = xDesc(x)
  y = x isa AbstractVector ? similar(x, rnn.hidden) : similar(x, rnn.hidden, size(x, 2))
  ho = similar(h)
  ydesc = xDesc(y)
  workspace = getworkspace(rnn, seqLength, xdesc)
  reserve = train == Val{true} ?
    CuVector{UInt8}(rnnTrainingReserveSize(rnn, seqLength, xdesc)) :
    nothing
  co = c == nothing ? c : similar(c)
  cudnnRNNForward(rnn, seqLength,
                  xdesc, x,
                  hDesc(h)...,
                  hDesc(c)...,
                  FilterDesc(T, (1, 1, length(rnn.params))), rnn.params,
                  ydesc, y,
                  hDesc(ho)...,
                  hDesc(co)...,
                  workspace, reserve)
  result = c == nothing ? (y, ho) : (y, ho, co)
  return train == Val{true} ? (reserve, result) : result
end

forwardTrain(rnn::RNNDesc{T}, x::CuArray{T}, h::CuArray{T}, c = nothing) where T =
  forward(rnn, x, h, c, Val{true})

function cudnnRNNBackwardData(rnn::RNNDesc{T}, seqlen, yd, y, dyd, dy, dhod, dho, dcod, dco,
                              wd, w, hd, h, cd, c, dxd, dx, dhd, dh, dcd, dc, ws, rs) where T
  @check ccall((:cudnnRNNBackwardData,libcudnn),cudnnStatus_t,
               (Ptr{Nothing}, Ptr{Nothing}, Cint,
                Ptr{Ptr{Nothing}}, Ptr{T}, Ptr{Ptr{Nothing}}, Ptr{T}, Ptr{Nothing}, Ptr{T},
                Ptr{Nothing}, Ptr{T}, Ptr{Nothing}, Ptr{T}, Ptr{Nothing}, Ptr{T}, Ptr{Nothing},
                Ptr{T}, Ptr{Ptr{Nothing}}, Ptr{T}, Ptr{Nothing}, Ptr{T}, Ptr{Nothing}, Ptr{T},
                Ptr{Nothing}, Csize_t, Ptr{Nothing}, Csize_t),
               libcudnn_handle[], rnn, seqlen, yd, y, dyd, dy, dhod, dho, dcod, dco,
               wd, w, hd, h, cd, c, dxd, dx, dhd, dh, dcd, dc, ws, length(ws), rs, length(rs))
end

function backwardData(rnn::RNNDesc{T}, y, dy_, dho, dco, h, c, reserve) where T
  # Same as above, any more efficient way?
  dy = dy_ isa Integer ? zero(y) : dy_
  yd = xDesc(y)
  dx = y isa AbstractVector ? similar(dy, rnn.input) : similar(dy, rnn.input, size(dy, 2))
  dh = similar(h)
  dc = c == nothing ? nothing : similar(c)
  cudnnRNNBackwardData(rnn, 1,
    yd, y, yd, dy, hDesc(dho)..., hDesc(dco)...,
    FilterDesc(T, (1, 1, length(rnn.params))), rnn.params,
    hDesc(h)..., hDesc(c)..., xDesc(dx), dx, hDesc(dh)..., hDesc(dc)...,
    workspace[], reserve)
  return c == nothing ? (dx, dh) : (dx, dh, dc)
end

backwardData(rnn, y, dy, dho, hx, reserve) =
  backwardData(rnn, y, dy, dho, nothing, hx, nothing, reserve)

function cudnnRNNBackwardWeights(rnn::RNNDesc{T}, seqlen, xd, x, hd, h, yd, y, dwd, dw,
                                 workspace, reserve) where T
  @check ccall((:cudnnRNNBackwardWeights,libcudnn), cudnnStatus_t,
               (Ptr{Nothing}, Ptr{Nothing}, Cint,  # handle, rnnDesc, seqLength
                Ptr{Ptr{Nothing}}, Ptr{T}, #x
                Ptr{Nothing}, Ptr{T}, #hx
                Ptr{Ptr{Nothing}}, Ptr{T}, #y
                Ptr{Nothing}, Csize_t, #ws
                Ptr{Nothing}, Ptr{T}, #dw
                Ptr{Nothing}, Csize_t), #rs
               libcudnn_handle[], rnn, seqlen, xd, x, hd, h, yd, y,
               workspace, length(workspace), dwd, dw, reserve, length(reserve))
end

function backwardWeights(rnn::RNNDesc{T}, x, h, y, reserve) where T
  dw = zero(rnn.params)
  cudnnRNNBackwardWeights(rnn, 1,
    xDesc(x), x, hDesc(h)..., xDesc(y), y,
    FilterDesc(T, (1, 1, length(dw))), dw,
    workspace[], reserve)
  return params(dw, rnn.input, rnn.hidden, ngates(rnn))
end

# Interface

import ..Flux: Flux, relu
import ..Tracker: TrackedArray
using .CuArrays.CUDAnative
using .CuArrays: @cuindex, cudims

function LinearAlgebra.copy_transpose!(dst::CuArray, src::CuArray)
  function kernel(dst, src)
    I = @cuindex dst
    dst[I...] = src[reverse(I)...]
    return
  end
  blk, thr = cudims(dst)
  @cuda blocks=blk threads=thr kernel(dst, src)
  return dst
end

CuParam{T,N} = Union{CuArray{T,N},TrackedArray{T,N,CuArray{T,N}}}
CuRNN{T} = Flux.RNNCell{<:Union{typeof(tanh),typeof(relu)},<:CuParam{T,2},<:CuParam{T,1}}
CuGRU{T} = Flux.GRUCell{<:CuParam{T,2},<:CuParam{T,1}}
CuLSTM{T} = Flux.LSTMCell{<:CuParam{T,2},<:CuParam{T,1}}
CuRNNs{T} = Union{CuRNN{T},CuGRU{T},CuLSTM{T}}

function copyparams!(m::CuRNNs, d::RNNDesc)
  Wi, Wh = d.weights
  copy_transpose!(Wi, Flux.data(m.Wi))
  copy_transpose!(Wh, Flux.data(m.Wh))
  copy_transpose!(d.bias, Flux.data(m.b))
  return
end

function RNNDesc(m::CuRNNs{T}) where T
  h, i = length(m.h), size(m.Wi, 2)
  mode = m isa CuRNN ?
    (m.σ == tanh ? RNN_TANH : RNN_RELU) :
    m isa CuGRU ? GRU : LSTM
  r = RNNDesc{T}(mode, i, h)
  return r
end

const descs = WeakKeyDict()

function desc(rnn)
  d = haskey(descs, rnn) ? descs[rnn] : (descs[rnn] = RNNDesc(rnn))
  copyparams!(rnn, d)
  return d
end

import Flux.Tracker
import Flux.Tracker: data, istracked, track, unbroadcast, @grad, nobacksies

istrain(m::CuRNNs, args...) = any(x -> x isa TrackedArray, (m.Wi, m.Wh, m.b, args...))

function (m::CuRNN{T})(h::CuParam{T}, x::CuParam{T}) where T <: Union{Float32,Float64}
  result = istrain(m, h, x) ?
    track(m, x, h, m.Wi, m.Wh, m.b) :
    forward(desc(m), x, h)
  return result[2], result[1]
end

function (m::CuGRU{T})(h::CuParam{T}, x::CuParam{T}) where T <: Union{Float32,Float64}
  result = istrain(m, h, x) ?
    track(m, x, h, m.Wi, m.Wh, m.b) :
    forward(desc(m), x, h)
  return result[2], result[1]
end

function (m::CuLSTM{T})(h::NTuple{2,CuParam{T}}, x::CuParam{T}) where T <: Union{Float32,Float64}
  result = istrain(m, h, x) ?
    track(m, x, h[1], h[2], m.Wi, m.Wh, m.b) :
    forward(desc(m), x, h[1], h[2])
  return (result[2], result[3]), result[1]
end

(m::CuRNN{T})(h::CuParam{T}, x) where T <: Union{Float32,Float64} = m(h, CuArray{T}(x))
(m::CuGRU{T})(h::CuParam{T}, x) where T <: Union{Float32,Float64} = m(h, CuArray{T}(x))
(m::CuLSTM{T})(h::NTuple{2,CuParam{T}}, x) where T <: Union{Float32,Float64} = m(h, CuArray{T}(x))

@grad function (m::Union{CuRNN,CuGRU})(x, h, Wi, Wh, b)
  reserve, result = forwardTrain(desc(m), data(x), data(h))
  result, function (Δ)
    y, ho = result
    dy, dho = Δ
    h_ = hBatch(x, data(h))
    dx, dh = backwardData(descs[m], y, dy, dho, h_, reserve)
    (dWi, dWh), db = backwardWeights(descs[m], data(x), h_, y, reserve)
    nobacksies(:RNN, (dx, unbroadcast(h, dh), transpose(dWi), transpose(dWh), db))
  end
end

@grad function (m::CuLSTM)(x, h, c, Wi, Wh, b)
  reserve, result = forwardTrain(desc(m), data.((x, h, c))...)
  result, function (Δ)
    y, ho = result
    dy, dho, dco = Δ
    h_ = hBatch(x, data(h))
    c_ = hBatch(x, data(c))
    dx, dh, dc = backwardData(descs[m], y, dy, dho, dco, h_, c_, reserve)
    (dWi, dWh), db = backwardWeights(descs[m], data(x), h_, y, reserve)
    nobacksies(:RNN,
      (dx, unbroadcast(h, dh), unbroadcast(c, dc),
       transpose(dWi), transpose(dWh), db))
  end
end


# # Battle Royale
# ![](battle.gif)
#
# This example illustrates how to leverage higher dimensions of a `GridSpace` to identify
# the distance from neighbors not just spatially, but also categorically. We'll also use
# the [`walk!`](@ref) function extensively.

# The `Models` module includes this example as [`Models.battle`](@ref).

# ## Rules of Engagement
# Agents wander around the map looking for opponents. When a grid space is occupied by
# two or more agents there will be a battle. With experience gained from the fight, the
# victor searches for more opponents to crush and losers scurry away defeated or possibly
# even die. This process repeats until there is a single, definitive winner.
#
# For this battle ground to exist, the following rules must be followed:
# - Agents have an experience level, starting at level 1 up to a maximum of 10.
# - Agents will search for the nearest worthy opponent (one with equal or ±1 experience level) and move towards them to attack, so long as something more important doesn't happen, which could be
#   - A tougher opponent (with experience level +2 or higher) is nearby: run!
#   - There are no worthy opponents available, but there are weak ones (with experience level -2 or lower): chase them down.
#   - Capture and taunt a weaker opponent, then kill them.
#   - Notice a tough opponent is occupied, sneak up and kill them.
#   - There is no-one worthy to fight, but also no-one left to taunt. All bets are off: THERE CAN BE ONLY ONE.

# Battles are won by weighted chance - a higher level gives an agent a larger chance of
# winning, but does not guarantee it. When a victor is chosen
#   - The difference in experience between opponents is swapped.
#   - If an agents experience reaches 0, they die.
#
# Captured opponents will be killed once taunted. The captor will gain half of their
# experience. If an opportunist manages to take the captor by surprise, they can gain
# up to half of the captor's experience. This means a level 1 agent may eliminate a
# level 10 captor and jump straight to level 6.
#
# Once all rules of engagement have been exhausted, the final showdown begins. Opponents
# fight their closest adversary regardless of experience level. Winner takes all.

# ## Model Setup

cd(@__DIR__) #src
using Random # hide
using Agents
using Plots
gr() # hide

mutable struct Fighter <: AbstractAgent
    id::Int
    pos::Dims{3}
    has_prisoner::Bool
    capture_time::Int
    shape::Symbol # For plotting
end

# As you can see, the properties of out agent are very simple and contain only two
# parameters that are needed to store context from one time step to the next. All
# other properties needed are stored in the space. `pos` is three-dimensional, two
# for the actual space agents move within, and a third categorical dimension representing
# their level. `shape` is used solely for plotting (well, used once just for convenience).

# Now let's set up the battle field:
function battle(; fighters = 50)
    model = ABM(
        Fighter,
        GridSpace((100, 100, 10); periodic = false);
        scheduler = random_activation,
    )

    n = 0
    while n != fighters
        pos = (rand(1:100, 2)..., 1) # Start at level 1
        if isempty(pos, model)
            add_agent!(pos, model, false, 0, :diamond)
            n += 1
        end
    end

    return model
end

Random.seed!(6547) # hide
model = battle()

# 50 opponents positioned randomly on a 100x100 grid, with no escape
# (`periodic = false`). To leverage categorical dimensions fully, non-periodic chebyshev
# space is necessary.

# ## Game Dynamics
#
# To implement the rules of engagement, only an `agent_step!` function is required,
# along with a few helper functions.

space(agent) = agent.pos[1:2]
level(agent) = agent.pos[3]

# `space` allows us to invoke a number of helpful utilities provided by Agents.jl
# but only operate on our spatial dimensions, `level` is a wrapper to access the agent's
# experience easily.

# Nearest agents that satisfy our search criteria can be identified via Euclidean
# distance solely on the spatial dimensions of our `GridSpace`.

function closest_target(agent::Fighter, ids::Vector{Int}, model::ABM)
    if length(ids) == 1
        closest = ids[1]
    else
        close_id = argmin(map(id -> edistance(space(agent), space(model[id]), model), ids))
        closest = ids[close_id]
    end
    return model[closest]
end

# Since our battles are only between opponents with equal, or as much as one level apart,
# the odds can be set explicitly. Stronger opponents have twice the capacity of winning a
# match.

function battle!(one::Fighter, two::Fighter, model)
    if level(one) == level(two)
        ## Odds are equivalent
        one_winner = rand() < 0.5
    elseif level(one) > level(two)
        ## Odds are in favor of one
        one_winner = 2 * rand() > rand()
    else
        ## Odds are in favor of two
        one_winner = rand() > 2 * rand()
    end

    one_winner ? (up = one; down = two) : (up = two; down = one)

    new_lvl_up = min(level(up) + 1, 10)
    new_pos_up = clamp.(rand(-1:1, 2) .+ space(up), [1, 1], size(model.space)[1:2])
    move_agent!(up, (new_pos_up..., new_lvl_up), model)
    new_lvl_down = level(down) - 1
    if new_lvl_down == 0
        kill_agent!(down, model)
    else
        move_agent!(down, (space(down)..., new_lvl_down), model)
    end
end

# If an agent has a prisoner, it will taunt it for a time, then kill it, so long as
# an opportunist doesn't sneak up on them first!
# Here we use the tuple constructor with [`nearby_ids`](@ref) to look for agents at
# the same position as the captor `(0, 0)`, and *any* level `(..., 10)`. We could also
# use the range constructor in this instance
# `nearby_ids(agent, model, [(1, 0:0), (2, 0:0)])`, meaning which is more performant but not as
# readable.

function captor_behavior!(agent, model)
    close_ids = collect(nearby_ids(agent, model, (0, 0, 10)))
    if length(close_ids) == 1
        ## Taunt prisoner or kill it
        prisoner = model[close_ids[1]]
        if prisoner.capture_time > 10
            agent.shape = :square
            gain = ceil(Int, level(prisoner) / 2)
            new_lvl = min(level(agent) + gain, 10)
            kill_agent!(prisoner, model)
            agent.has_prisoner = false
            move_agent!(agent, (space(agent)..., new_lvl), model)
        end
    else
        ## Someone is here to kill the captor. Could be more than one opponent
        prisoner = [model[id] for id in close_ids if model[id].capture_time > 0][1]
        exploiter = rand([
            model[id]
            for
            id in close_ids if
            model[id].capture_time == 0 && model[id].has_prisoner == false
        ])
        exploiter.shape = :square
        gain = ceil(Int, level(agent) / 2)
        new_lvl = min(level(agent) + rand(1:gain), 10)
        kill_agent!(agent, model)
        move_agent!(exploiter, (space(exploiter)..., new_lvl), model)
        ## Prisoner runs away in the commotion
        prisoner.shape = :utriangle
        prisoner.capture_time = 0
        walk!(prisoner, (rand(-1:1, 2)..., 0), model)
    end
end

# When there are only few fighters standing, the stakes are higher. Prior experience is
# paramount since there is no gain, and fights are to the death.

function endgame!(agent, model)
    origin = space(agent)
    end_ids = collect(Iterators.filter(
        id -> model[id].shape == :circle && id != agent.id,
        allids(model),
    ))
    agent.shape = :circle
    if !isempty(end_ids)
        opponent = closest_target(agent, end_ids, model)
        target = space(opponent)
        if origin == target
            ## Battle
            agent.shape = :square
            opponent.shape = :square
            showdown!(agent, opponent, model)
        else
            walk!(agent, (sign.(target .- origin)..., 0), model)
        end
    end
end

function showdown!(one::Fighter, two::Fighter, model)
    if level(one) == level(two)
        ## Odds are equivalent
        one_winner = rand() < 0.5
    elseif level(one) > level(two)
        ## Odds are in favor of one
        one_winner = level(one) - level(two) * rand() > rand()
    else
        ## Odds are in favor of two
        one_winner = rand() > level(two) - level(one) * rand()
    end

    one_winner ? kill_agent!(two, model) : kill_agent!(one, model)
end

# The rest of our interactions flow down a hierarchy, so we'll place them directly in the
# `agent_step!` function. We use the tuple search for `occupied_ids` here, as we did with
# `close_ids` above. The rest of the searches however use the range search to provide a
# more precise criteria.
#
# The easiest context to explore is `worthy_ids`: all we want to do is find an agent with
# a similar level. If we used the tuple search here, we would have to search
# `(100, 100, 1)` - even though we are not at all interested in the spatial location of
# the neighbors at this time. `(3, -1:1)` is therefore more accurate representation.
#
# A more complex example is that of `strong_ids`. We are looking for agents with a level
# 2-4 points higher withing a distance of `(5, 5)`. The range search becomes a little
# verbose, but precise. An equivalent tuple search is not completely possible however.
# The closest solution is `(5, 5, 4)`, which also looks for *weaker* opponents and must
# be filtered to the correct neighbor set after the fact. In this instance the range
# search has significant performance gains.

function agent_step!(agent, model)
    if agent.capture_time > 0
        ## Captured agents are powerless, but we need to keep track of how long
        ## they have been in this state
        agent.capture_time += 1
    elseif agent.has_prisoner
        captor_behavior!(agent, model)
    else
        origin = space(agent)
        ## Find agents that have captives, they are not focused
        occupied_ids = collect(Iterators.filter(
            id -> model[id].has_prisoner,
            nearby_ids(agent, model, (7, 7, 10)),
        ))
        if !isempty(occupied_ids)
            ## Sneak up behind them
            target = space(closest_target(agent, occupied_ids, model))
            agent.shape = :pentagon
            walk!(agent, (sign.(target .- origin)..., 0), model)
        else
            ## Opponents that are greatly higher in rank that the current agent
            strong_ids = collect(nearby_ids(agent, model, [(1, -5:5), (2, -5:5), (3, 2:4)]))
            if !isempty(strong_ids)
                ## Run away from nearest
                target = space(closest_target(agent, strong_ids, model))
                agent.shape = :utriangle
                walk!(agent, (sign.(origin .- target)..., 0), model)
            else
                ## There are no distractions. Search for the closest worthy opponent
                worthy_ids = collect(nearby_ids(agent, model, [(3, -1:1)]))
                if !isempty(worthy_ids)
                    opponent = closest_target(agent, worthy_ids, model)
                    target = space(opponent)
                    if origin == target
                        ## Battle
                        agent.shape = :square
                        opponent.shape = :square
                        battle!(agent, opponent, model)
                    else
                        ## Move towards worthy opponent
                        agent.shape = :diamond
                        walk!(agent, (sign.(target .- origin)..., 0), model)
                    end
                else
                    ## Find any weak targets in the vicinity
                    weak_ids = collect(nearby_ids(
                        agent,
                        model,
                        [(1, -10:10), (2, -10:10), (3, -4:-2)],
                    ))
                    if !isempty(weak_ids)
                        prisoner = closest_target(agent, weak_ids, model)
                        target = space(prisoner)
                        if origin == target
                            ## Capture and taunt target
                            agent.has_prisoner = true
                            agent.shape = :vline
                            prisoner.capture_time += 1
                            prisoner.shape = :hline
                        else
                            ## Chase down nearest (can move 2 steps at a time!)
                            agent.shape = :star4
                            walk!(agent, (2 .* sign.(target .- origin)..., 0), model)
                        end
                    else
                        ## Abandon honour. This is the end
                        endgame!(agent, model)
                    end
                end
            end
        end
    end

    return nothing
end

# ## Let the Battle Begin
# Plotting is relatively straightforward. [`plotabm`](@ref) cannot be used explicitly (yet)
# since it expects our categorical dimension is actually a third spatial one.
# We start with some custom legends to easier understand the dynamics.

clr(agent) = cgrad(:tab10)[level(agent)]
lvl = plot(
    (1:10)',
    label = ["Level 1" "Level 2" "Level 3" "Level 4" "Level 5" "Level 6" "Level 7" "Level 8" "Level 9" "Level 10"],
    palette = :tab10,
    legend = :top,
    framestyle = :none,
)
state = scatter(
    (1:8)',
    xlims = (-2, -1),
    label = ["Battle" "Run" "Showdown" "Sneak" "Duel" "Captor" "Prisoner" "Chase"],
    markershape = [:square :utriangle :circle :pentagon :diamond :vline :hline :star4],
    color = :black,
    legend = :top,
    framestyle = :none,
)
anim = @animate for i in 0:225
    posn = [space(model[id]) for id in by_id(model)]
    cm = [clr(model[id]) for id in by_id(model)]
    shp = [model[id].shape for id in by_id(model)]
    ll = @layout [a{0.7w} [b; c]]
    battle = scatter(
        posn,
        legend = :none,
        color = cm,
        markersize = 7,
        markershape = shp,
        xlims = (-2, 103),
        ylims = (-2, 103),
        showaxis = false,
        minorgrid = true,
        ticks = (0:10:100, []),
    )
    plot(battle, lvl, state, layout = ll, size = (850, 600))
    step!(model, agent_step!, 1)
end
gif(anim, "battle.gif", fps = 10)

# Some interesting behaviour emerges: sometimes you see a group of diamonds chasing one triangle.
# What ends up happening here is usually a close pair that wishes to fight gets caught
# out by the weaker one of the two running away from an even stronger opponent. Problem
# is that this stronger opponent is chasing the stronger of the pair, but since the
# weakest of the pair is still closer to the newcomer, there is a stalemate.
# This is usually resolved by hitting a boundary or other opponents.


###### This file is part of the MendelImpute.jl package.
###### It contains relevant code to decompress a vector of `HaplotypeMosaicPair`
###### to a genotype matrix

"""
    convert_compressed(t<:Real, phaseinfo::String, reffile::String)

Converts `phaseinfo` into a phased genotype matrix of type `t` using the full
reference haplotype panel `H`

# Inputs
- `t`: Type of matrix. If `bool`, genotypes are converted to a `BitMatrix`
- `phaseinfo`: Vector of `HaplotypeMosaicPair`s stored in `.jlso` format
- `reffile`: The complete (uncompressed) haplotype reference file

# Output
+ `X1`: allele 1 of the phased genotype. Each column is a sample. `X = X1 + X2`.
+ `X2`: allele 2 of the phased genotype. Each column is a sample. `X = X1 + X2`.
+ `phase`: the original data structure after phasing and imputation.
+ `sampleID`: The ID's of each imputed person.
+ `H`: the complete reference haplotype panel. Columns of `H` are haplotypes.
"""
function convert_compressed(
    t::Type{T},
    phaseinfo::AbstractString,
    reffile::AbstractString
    ) where T <: Real
    endswith(phaseinfo, ".jlso") || error("phaseinfo does not end with '.jlso'")
    H = convert_ht(Bool, reffile, trans=true, msg="importing reference data...")

    phase_and_sampleID = JLSO.load(phaseinfo)
    phase = phase_and_sampleID[:ph]
    sampleID = phase_and_sampleID[:sampleID]

    X1, X2 = convert_compressed(t, phase, H)
    return X1, X2, phase, sampleID, H
end

"""
    convert_compressed(t<:Real, phaseinfo::Vector{HaplotypeMosaicPair}, H::AbstractMatrix)

Columns of `H` are haplotypes.
"""
function convert_compressed(
    t::Type{T},
    phaseinfo::Vector{HaplotypeMosaicPair},
    H::AbstractMatrix
    ) where T <: Real

    people = length(phaseinfo)
    snps = phaseinfo[1].strand1.length
    M = (t == Bool ? BitArray{2} : Matrix{t})

    X1 = M(undef, snps, people)
    X2 = M(undef, snps, people)
    impute!(X1, X2, H, phaseinfo)

    return X1, X2
end


function parser(m::JuMP.Model)
    map::Map = Map()

    variable_list = JuMP.all_variables(m)
    @assert all(v -> JuMP.has_lower_bound(v) && JuMP.has_upper_bound(v), variable_list) "All variables must be bounded"

    for var in variable_list
        lower_bound = JuMP.lower_bound(var)
        upper_bound = JuMP.upper_bound(var)

        scale = upper_bound - lower_bound + 1

        num_created_city::Int64 = Int(ceil(log2(scale)))

        print(num_created_city)
    end

end


mutable struct OnlineSVGP{
    T<:Real,
    TLikelihood<:AbstractLikelihood,
    TInference<:AbstractInference,
    TData<:AbstractDataContainer,
    N,
} <: AbstractGP{T,TLikelihood,TInference,N}
    data::TData
    f::NTuple{N,OnlineVarLatent{T}}
    likelihood::TLikelihood
    inference::TInference
    verbose::Int64
    atfrequency::Int64
    trained::Bool
end

"""
    OnlineSVGP(args...; kwargs...)

Online Sparse Variational Gaussian Process

## Arguments
- `kernel::Kernel` : Covariance function, can be any kernel from KernelFunctions.jl
- `likelihood` : Likelihood of the model. For compatibilities, see [`Likelihood Types`](@ref likelihood_user)
- `inference` : Inference for the model, see the [`Compatibility Table`](@ref compat_table))
- `Zalg` : Algorithm selecting how inducing points are selected

## Keywords arguments
- `verbose` : How much does the model print (0:nothing, 1:very basic, 2:medium, 3:everything)
- `optimiser` : Optimiser used for the kernel parameters. Should be an Optimiser object from the [Flux.jl](https://github.com/FluxML/Flux.jl) library, see list here [Optimisers](https://fluxml.ai/Flux.jl/stable/training/optimisers/) and on [this list](https://github.com/theogf/AugmentedGaussianProcesses.jl/tree/master/src/inference/optimisers.jl). Default is `ADAM(0.001)`
- `atfrequency::Int=1` : Choose how many variational parameters iterations are between hyperparameters optimization
- `mean=ZeroMean()` : PriorMean object, check the documentation on it [`MeanPrior`](@ref meanprior)
- `Zoptimiser` : Optimiser for inducing points locations
- `T::DataType=Float64` : Hint for what the type of the data is going to be.
"""
function OnlineSVGP(
    kernel::Kernel,
    likelihood::AbstractLikelihood,
    inference::AbstractInference,
    Zalg::InducingPoints.OnIPSA=OIPS(0.9);
    verbose::Integer=0,
    optimiser=ADAM(0.01),
    atfrequency::Integer=1,
    mean::Union{<:Real,AbstractVector{<:Real},PriorMean}=ZeroMean(),
    Zoptimiser=nothing,
    T::DataType=Float64,
)
    data = OnlineDataContainer()
    inference isa AnalyticVI || error("The inference object should be of type `AnalyticVI`")
    implemented(likelihood, inference) ||
        error("The $likelihood is not compatible or implemented with the $inference")

    if isa(optimiser, Bool)
        optimiser = optimiser ? ADAM(0.01) : nothing
    end

    if typeof(mean) <: Real
        mean = ConstantMean(mean)
    elseif typeof(mean) <: AbstractVector{<:Real}
        mean = EmpiricalMean(mean)
    end

    nLatent = num_latent(likelihood)
    latentf = ntuple(_ -> OnlineVarLatent(T, 0, 0, [], Zalg, kernel, mean, optimiser, Zoptimiser), nLatent)
    inference = tuple_inference(inference, nLatent, 0, 0, 0, [], [])
    inference.nIter = 1
    return OnlineSVGP{T,typeof(likelihood),typeof(inference),typeof(data),nLatent}(
        data, latentf, likelihood, inference, verbose, atfrequency, false
    )
end

function Base.show(io::IO, model::OnlineSVGP) where {T}
    return print(
        io,
        "Online Variational Gaussian Process with a $(likelihood(model)) infered by $(inference(model)) ",
    )
end

@traitimpl IsSparse{OnlineSVGP}

Zviews(m::OnlineSVGP) = Zview.(m.f)
objective(m::OnlineSVGP) = ELBO(m)
MBIndices(m::OnlineSVGP) = 1:nSamples(m)
xview(m::OnlineSVGP) = input(m)
yview(m::OnlineSVGP) = output(m)
nFeatures(m::OnlineSVGP) = collect(dim.(getf(m)))

## Accessors to InducingPoints methods
# function InducingPoints.init(Z::OptimIP, m::OnlineSVGP, gp::OnlineVarLatent)
#     return InducingPoints.init(Z.Z, m, gp)
# end
# InducingPoints.init(Z::OIPS, m, gp) = InducingPoints.init(Z, input(m), kernel(gp))

# function InducingPoints.add_point!(Z::OptimIP, m::OnlineSVGP, gp::OnlineVarLatent)
#     return InducingPoints.add_point!(Z.Z, m, gp)
# end
# function InducingPoints.add_point!(Z::OIPS, m::OnlineSVGP, gp::OnlineVarLatent)
#     return InducingPoints.add_point!(Z, input(m), kernel(gp))
# end

# function InducingPoints.remove_point!(Z::OptimIP, m::OnlineSVGP, gp::OnlineVarLatent)
#     return InducingPoints.remove_point!(Z.Z, m, gp)
# end
# function InducingPoints.remove_point!(Z::OIPS, m::OnlineSVGP, gp::OnlineVarLatent)
#     return InducingPoints.remove_point!(Z, pr_cov(gp), kernel(gp))
# end


function GeographicZoneDetector(arg0::BodyShape, arg1::SphericalPolygonsSet, arg2::jdouble)
    return GeographicZoneDetector((BodyShape, SphericalPolygonsSet, jdouble), arg0, arg1, arg2)
end

function GeographicZoneDetector(arg0::jdouble, arg1::jdouble, arg2::BodyShape, arg3::SphericalPolygonsSet, arg4::jdouble)
    return GeographicZoneDetector((jdouble, jdouble, BodyShape, SphericalPolygonsSet, jdouble), arg0, arg1, arg2, arg3, arg4)
end

function g(obj::GeographicZoneDetector, arg0::SpacecraftState)
    return jcall(obj, "g", jdouble, (SpacecraftState,), arg0)
end

function get_body(obj::GeographicZoneDetector)
    return jcall(obj, "getBody", BodyShape, ())
end

function get_margin(obj::GeographicZoneDetector)
    return jcall(obj, "getMargin", jdouble, ())
end

function get_zone(obj::GeographicZoneDetector)
    return jcall(obj, "getZone", SphericalPolygonsSet, ())
end

function with_margin(obj::GeographicZoneDetector, arg0::jdouble)
    return jcall(obj, "withMargin", GeographicZoneDetector, (jdouble,), arg0)
end


using FileIO

include("tiffxml.jl")
include("Segmentation3D.jl")

data_dir = "/datahub/rawdata/tandeng/mRNA_imaging/mRNA_confocal_hamamatsu-60X-TIRF/20200316_result"
#cell_name = filter(x->occursin(r"^s",x), readdir(data_dir) )
cell_name = filter(x->occursin(r"^s[0-9][0-9]?-c[0-9][0-9]?.tiff$",x), readdir(data_dir) )

for i in cell_name
	print("processing $i")
	@time cell_raw = load(File(format"TIFF", "$data_dir/$i"))
	@time nucleus_all, threshold_all = extract3dnucleus(cell_raw)
    save(File(format"TIFF", "$data_dir/$i-clear.ome.tiff"), N0f16.(nucleus_all))
	save("$data_dir/$i.jld", "threshold", threshold_all)
	x, y, z_all = size(nucleus_all)
	@time embedxml(x, y, 20, z_all÷20, "$data_dir/$i-clear.ome.tiff")
end


"""
Construct model for HydroGen with FixedOutput Formulation
"""
function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ArgumentConstructStage,
    model::DeviceModel{H, FixedOutput},
    ::Type{S},
) where {H <: PSY.HydroGen, S <: PM.AbstractPowerModel}
    devices = get_available_components(H, sys)

    add_parameters!(container, ActivePowerTimeSeriesParameter, devices, model)
    add_parameters!(container, ReactivePowerTimeSeriesParameter, devices, model)

    # Expression
    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerTimeSeriesParameter,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ReactivePowerBalance,
        ReactivePowerTimeSeriesParameter,
        devices,
        model,
        S,
    )
    return
end

function construct_device!(
    ::OptimizationContainer,
    ::PSY.System,
    ::ModelConstructStage,
    ::DeviceModel{H, FixedOutput},
    ::Type{S},
) where {H <: PSY.HydroGen, S <: PM.AbstractPowerModel}
    # FixedOutput doesn't add any constraints to the model. This function covers
    # AbstractPowerModel and AbtractActivePowerModel
    return
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ArgumentConstructStage,
    model::DeviceModel{H, FixedOutput},
    ::Type{S},
) where {H <: PSY.HydroGen, S <: PM.AbstractActivePowerModel}
    devices = get_available_components(H, sys)

    add_parameters!(container, ActivePowerTimeSeriesParameter, devices, model)

    # Expression
    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerTimeSeriesParameter,
        devices,
        model,
        S,
    )
    return
end

"""
Construct model for HydroGen with RunOfRiver Dispatch Formulation
"""
function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ArgumentConstructStage,
    model::DeviceModel{H, D},
    ::Type{S},
) where {
    H <: PSY.HydroGen,
    D <: AbstractHydroDispatchFormulation,
    S <: PM.AbstractPowerModel,
}
    devices = get_available_components(H, sys)

    add_variables!(container, ActivePowerVariable, devices, D())
    add_variables!(container, ReactivePowerVariable, devices, D())
    add_variables!(container, EnergyOutput, devices, D())
    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ReactivePowerBalance,
        ReactivePowerVariable,
        devices,
        model,
        S,
    )

    add_parameters!(container, ActivePowerTimeSeriesParameter, devices, model)

    add_expressions!(container, ProductionCostExpression, devices, model)

    add_to_expression!(
        container,
        ActivePowerRangeExpressionLB,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ActivePowerRangeExpressionUB,
        ActivePowerVariable,
        devices,
        model,
        S,
    )

    add_feedforward_arguments!(container, model, devices)
    return
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ModelConstructStage,
    model::DeviceModel{H, D},
    ::Type{S},
) where {
    H <: PSY.HydroGen,
    D <: AbstractHydroDispatchFormulation,
    S <: PM.AbstractPowerModel,
}
    devices = get_available_components(H, sys)

    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerRangeExpressionLB,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerRangeExpressionUB,
        devices,
        model,
        S,
    )

    add_constraints!(
        container,
        ReactivePowerVariableLimitsConstraint,
        ReactivePowerVariable,
        devices,
        model,
        S,
    )
    add_feedforward_constraints!(container, model, devices)

    objective_function!(container, devices, model, S)
    add_constraint_dual!(container, sys, model)

    return
end

"""
Construct model for HydroGen with RunOfRiver Dispatch Formulation
with only Active Power.
"""
function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ArgumentConstructStage,
    model::DeviceModel{H, D},
    ::Type{S},
) where {
    H <: PSY.HydroGen,
    D <: AbstractHydroDispatchFormulation,
    S <: PM.AbstractActivePowerModel,
}
    devices = get_available_components(H, sys)

    add_variables!(container, ActivePowerVariable, devices, D())
    add_variables!(container, EnergyOutput, devices, D())
    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerVariable,
        devices,
        model,
        S,
    )

    add_parameters!(container, ActivePowerTimeSeriesParameter, devices, model)
    add_to_expression!(
        container,
        ActivePowerRangeExpressionLB,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ActivePowerRangeExpressionUB,
        ActivePowerVariable,
        devices,
        model,
        S,
    )

    add_expressions!(container, ProductionCostExpression, devices, model)

    add_feedforward_arguments!(container, model, devices)
    return
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ModelConstructStage,
    model::DeviceModel{H, D},
    ::Type{S},
) where {
    H <: PSY.HydroGen,
    D <: AbstractHydroDispatchFormulation,
    S <: PM.AbstractActivePowerModel,
}
    devices = get_available_components(H, sys)

    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerRangeExpressionLB,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerRangeExpressionUB,
        devices,
        model,
        S,
    )

    add_feedforward_constraints!(container, model, devices)

    objective_function!(container, devices, model, S)

    add_constraint_dual!(container, sys, model)
    return
end

"""
Construct model for HydroGen with ReservoirBudget Dispatch Formulation
"""
function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ArgumentConstructStage,
    model::DeviceModel{H, HydroDispatchReservoirBudget},
    ::Type{S},
) where {H <: PSY.HydroEnergyReservoir, S <: PM.AbstractPowerModel}
    devices = get_available_components(H, sys)

    add_variables!(container, ActivePowerVariable, devices, HydroDispatchReservoirBudget())
    add_variables!(
        container,
        ReactivePowerVariable,
        devices,
        HydroDispatchReservoirBudget(),
    )
    add_variables!(container, EnergyOutput, devices, HydroDispatchReservoirBudget())
    add_parameters!(container, EnergyBudgetTimeSeriesParameter, devices, model)

    add_expressions!(container, ProductionCostExpression, devices, model)

    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ReactivePowerBalance,
        ReactivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ActivePowerRangeExpressionLB,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ActivePowerRangeExpressionUB,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_feedforward_arguments!(container, model, devices)
    return
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ModelConstructStage,
    model::DeviceModel{H, HydroDispatchReservoirBudget},
    ::Type{S},
) where {H <: PSY.HydroEnergyReservoir, S <: PM.AbstractPowerModel}
    devices = get_available_components(H, sys)

    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerRangeExpressionLB,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerRangeExpressionUB,
        devices,
        model,
        S,
    )

    add_constraints!(
        container,
        ReactivePowerVariableLimitsConstraint,
        ReactivePowerVariable,
        devices,
        model,
        S,
    )

    # Energy Budget Constraint
    add_constraints!(container, EnergyBudgetConstraint, devices, model, S)

    add_feedforward_constraints!(container, model, devices)

    objective_function!(container, devices, model, S)

    add_constraint_dual!(container, sys, model)
    return
end

"""
Construct model for HydroGen with ReservoirBudget Dispatch Formulation
with only Active Power.
"""
function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ArgumentConstructStage,
    model::DeviceModel{H, HydroDispatchReservoirBudget},
    ::Type{S},
) where {H <: PSY.HydroEnergyReservoir, S <: PM.AbstractActivePowerModel}
    devices = get_available_components(H, sys)

    add_variables!(container, ActivePowerVariable, devices, HydroDispatchReservoirBudget())
    add_variables!(container, EnergyOutput, devices, HydroDispatchReservoirBudget())
    add_parameters!(container, EnergyBudgetTimeSeriesParameter, devices, model)

    add_expressions!(container, ProductionCostExpression, devices, model)

    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerVariable,
        devices,
        model,
        S,
    )

    add_to_expression!(
        container,
        ActivePowerRangeExpressionLB,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ActivePowerRangeExpressionUB,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_feedforward_arguments!(container, model, devices)
    return
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ModelConstructStage,
    model::DeviceModel{H, HydroDispatchReservoirBudget},
    ::Type{S},
) where {H <: PSY.HydroEnergyReservoir, S <: PM.AbstractActivePowerModel}
    devices = get_available_components(H, sys)

    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerRangeExpressionLB,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerRangeExpressionUB,
        devices,
        model,
        S,
    )

    # Energy Budget Constraint
    add_constraints!(container, EnergyBudgetConstraint, devices, model, S)

    add_feedforward_constraints!(container, model, devices)

    objective_function!(container, devices, model, S)
    add_constraint_dual!(container, sys, model)

    return
end

"""
Construct model for HydroGen with ReservoirStorage Dispatch Formulation
"""
function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ArgumentConstructStage,
    model::DeviceModel{H, HydroDispatchReservoirStorage},
    ::Type{S},
) where {H <: PSY.HydroEnergyReservoir, S <: PM.AbstractPowerModel}
    devices = get_available_components(H, sys)

    add_variables!(container, ActivePowerVariable, devices, HydroDispatchReservoirStorage())
    add_variables!(
        container,
        ReactivePowerVariable,
        devices,
        HydroDispatchReservoirStorage(),
    )
    add_variables!(container, EnergyVariable, devices, HydroDispatchReservoirStorage())
    add_variables!(
        container,
        WaterSpillageVariable,
        devices,
        HydroDispatchReservoirStorage(),
    )
    add_variables!(
        container,
        EnergyShortageVariable,
        devices,
        HydroDispatchReservoirStorage(),
    )
    add_variables!(
        container,
        EnergySurplusVariable,
        devices,
        HydroDispatchReservoirStorage(),
    )
    add_variables!(container, EnergyOutput, devices, HydroDispatchReservoirStorage())
    add_parameters!(container, EnergyTargetTimeSeriesParameter, devices, model)
    add_parameters!(container, InflowTimeSeriesParameter, devices, model)

    add_expressions!(container, ProductionCostExpression, devices, model)

    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ReactivePowerBalance,
        ReactivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ActivePowerRangeExpressionLB,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ActivePowerRangeExpressionUB,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_feedforward_arguments!(container, model, devices)
    return
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ModelConstructStage,
    model::DeviceModel{H, HydroDispatchReservoirStorage},
    ::Type{S},
) where {H <: PSY.HydroEnergyReservoir, S <: PM.AbstractPowerModel}
    devices = get_available_components(H, sys)

    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerRangeExpressionLB,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerRangeExpressionUB,
        devices,
        model,
        S,
    )

    add_constraints!(
        container,
        ReactivePowerVariableLimitsConstraint,
        ReactivePowerVariable,
        devices,
        model,
        S,
    )

    add_initial_condition!(
        container,
        devices,
        HydroDispatchReservoirStorage(),
        InitialEnergyLevel(),
    )
    # Energy Balance Constraint
    add_constraints!(container, EnergyBalanceConstraint, devices, model, S)
    add_constraints!(container, EnergyTargetConstraint, devices, model, S)
    add_feedforward_constraints!(container, model, devices)

    objective_function!(container, devices, model, S)

    add_constraint_dual!(container, sys, model)
    return
end

"""
Construct model for HydroGen with ReservoirStorage Dispatch Formulation
with only Active Power
"""
function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ArgumentConstructStage,
    model::DeviceModel{H, HydroDispatchReservoirStorage},
    ::Type{S},
) where {H <: PSY.HydroEnergyReservoir, S <: PM.AbstractActivePowerModel}
    devices = get_available_components(H, sys)

    add_variables!(container, ActivePowerVariable, devices, HydroDispatchReservoirStorage())
    add_variables!(container, EnergyVariable, devices, HydroDispatchReservoirStorage())
    add_variables!(
        container,
        WaterSpillageVariable,
        devices,
        HydroDispatchReservoirStorage(),
    )
    add_variables!(
        container,
        EnergyShortageVariable,
        devices,
        HydroDispatchReservoirStorage(),
    )
    add_variables!(
        container,
        EnergySurplusVariable,
        devices,
        HydroDispatchReservoirStorage(),
    )
    add_variables!(container, EnergyOutput, devices, HydroDispatchReservoirStorage())
    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerVariable,
        devices,
        model,
        S,
    )

    add_parameters!(container, EnergyTargetTimeSeriesParameter, devices, model)
    add_parameters!(container, InflowTimeSeriesParameter, devices, model)

    add_expressions!(container, ProductionCostExpression, devices, model)

    add_to_expression!(
        container,
        ActivePowerRangeExpressionLB,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ActivePowerRangeExpressionUB,
        ActivePowerVariable,
        devices,
        model,
        S,
    )

    add_feedforward_arguments!(container, model, devices)

    return
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ModelConstructStage,
    model::DeviceModel{H, HydroDispatchReservoirStorage},
    ::Type{S},
) where {H <: PSY.HydroEnergyReservoir, S <: PM.AbstractActivePowerModel}
    devices = get_available_components(H, sys)

    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerRangeExpressionLB,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerRangeExpressionUB,
        devices,
        model,
        S,
    )

    add_initial_condition!(
        container,
        devices,
        HydroDispatchReservoirStorage(),
        InitialEnergyLevel(),
    )
    # Energy Balance Constraint
    add_constraints!(container, EnergyBalanceConstraint, devices, model, S)
    add_constraints!(container, EnergyTargetConstraint, devices, model, S)
    add_feedforward_constraints!(container, model, devices)

    objective_function!(container, devices, model, S)
    add_constraint_dual!(container, sys, model)

    return
end

"""
Construct model for HydroGen with RunOfRiver Commitment Formulation
"""
function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ArgumentConstructStage,
    model::DeviceModel{H, D},
    ::Type{S},
) where {H <: PSY.HydroGen, D <: HydroCommitmentRunOfRiver, S <: PM.AbstractPowerModel}
    devices = get_available_components(H, sys)

    add_variables!(container, ActivePowerVariable, devices, D())
    add_variables!(container, ReactivePowerVariable, devices, D())
    add_variables!(container, OnVariable, devices, D())
    add_variables!(container, EnergyOutput, devices, D())
    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerVariable,
        devices,
        model,
        S,
    )

    add_to_expression!(
        container,
        ReactivePowerBalance,
        ReactivePowerVariable,
        devices,
        model,
        S,
    )

    add_parameters!(container, ActivePowerTimeSeriesParameter, devices, model)

    add_expressions!(container, ProductionCostExpression, devices, model)

    add_to_expression!(
        container,
        ActivePowerRangeExpressionLB,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ActivePowerRangeExpressionUB,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_feedforward_arguments!(container, model, devices)
    return
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ModelConstructStage,
    model::DeviceModel{H, D},
    ::Type{S},
) where {H <: PSY.HydroGen, D <: HydroCommitmentRunOfRiver, S <: PM.AbstractPowerModel}
    devices = get_available_components(H, sys)

    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerRangeExpressionLB,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerRangeExpressionUB,
        devices,
        model,
        S,
    )

    add_constraints!(
        container,
        ReactivePowerVariableLimitsConstraint,
        ReactivePowerVariable,
        devices,
        model,
        S,
    )

    add_feedforward_constraints!(container, model, devices)

    objective_function!(container, devices, model, S)

    add_constraint_dual!(container, sys, model)
    return
end

"""
Construct model for HydroGen with RunOfRiver Commitment Formulation
with only Active Power.
"""
function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ArgumentConstructStage,
    model::DeviceModel{H, D},
    ::Type{S},
) where {
    H <: PSY.HydroGen,
    D <: HydroCommitmentRunOfRiver,
    S <: PM.AbstractActivePowerModel,
}
    devices = get_available_components(H, sys)

    add_variables!(container, ActivePowerVariable, devices, D())
    add_variables!(container, OnVariable, devices, D())
    add_variables!(container, EnergyOutput, devices, D())
    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerVariable,
        devices,
        model,
        S,
    )

    add_expressions!(container, ProductionCostExpression, devices, model)

    add_parameters!(container, ActivePowerTimeSeriesParameter, devices, model)
    add_to_expression!(
        container,
        ActivePowerRangeExpressionLB,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ActivePowerRangeExpressionUB,
        ActivePowerVariable,
        devices,
        model,
        S,
    )

    add_feedforward_arguments!(container, model, devices)
    return
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ModelConstructStage,
    model::DeviceModel{H, D},
    ::Type{S},
) where {
    H <: PSY.HydroGen,
    D <: HydroCommitmentRunOfRiver,
    S <: PM.AbstractActivePowerModel,
}
    devices = get_available_components(H, sys)

    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerRangeExpressionLB,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerRangeExpressionUB,
        devices,
        model,
        S,
    )

    add_feedforward_constraints!(container, model, devices)

    objective_function!(container, devices, model, S)

    add_constraint_dual!(container, sys, model)
    return
end

"""
Construct model for HydroGen with ReservoirBudget Commitment Formulation
"""
function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ArgumentConstructStage,
    model::DeviceModel{H, D},
    ::Type{S},
) where {H <: PSY.HydroGen, D <: HydroCommitmentReservoirBudget, S <: PM.AbstractPowerModel}
    devices = get_available_components(H, sys)

    add_variables!(container, ActivePowerVariable, devices, D())
    add_variables!(container, ReactivePowerVariable, devices, D())
    add_variables!(container, OnVariable, devices, D())
    add_variables!(container, EnergyOutput, devices, D())
    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ActivePowerBalance,
        ReactivePowerVariable,
        devices,
        model,
        S,
    )

    add_parameters!(container, EnergyBudgetTimeSeriesParameter, devices, model)

    add_expressions!(container, ProductionCostExpression, devices, model)

    add_to_expression!(
        container,
        ActivePowerRangeExpressionLB,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ActivePowerRangeExpressionUB,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_feedforward_arguments!(container, model, devices)
    return
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ModelConstructStage,
    model::DeviceModel{H, D},
    ::Type{S},
) where {H <: PSY.HydroGen, D <: HydroCommitmentReservoirBudget, S <: PM.AbstractPowerModel}
    devices = get_available_components(H, sys)

    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerRangeExpressionLB,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerRangeExpressionUB,
        devices,
        model,
        S,
    )

    add_constraints!(
        container,
        ReactivePowerVariableLimitsConstraint,
        ReactivePowerVariable,
        devices,
        model,
        S,
    )
    # Energy Budget Constraint
    add_constraints!(container, EnergyBudgetConstraint, devices, model, S)

    add_feedforward_constraints!(container, model, devices)

    objective_function!(container, devices, model, S)
    add_constraint_dual!(container, sys, model)

    return
end

"""
Construct model for HydroGen with ReservoirBudget Commitment Formulation
with only Active Power.
"""
function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ArgumentConstructStage,
    model::DeviceModel{H, D},
    ::Type{S},
) where {
    H <: PSY.HydroGen,
    D <: HydroCommitmentReservoirBudget,
    S <: PM.AbstractActivePowerModel,
}
    devices = get_available_components(H, sys)

    add_variables!(container, ActivePowerVariable, devices, D())
    add_variables!(container, OnVariable, devices, D())
    add_variables!(container, EnergyOutput, devices, D())
    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerVariable,
        devices,
        model,
        S,
    )

    add_parameters!(container, EnergyBudgetTimeSeriesParameter, devices, model)

    add_expressions!(container, ProductionCostExpression, devices, model)
    add_to_expression!(
        container,
        ActivePowerRangeExpressionLB,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ActivePowerRangeExpressionUB,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_feedforward_arguments!(container, model, devices)

    return
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ModelConstructStage,
    model::DeviceModel{H, D},
    ::Type{S},
) where {
    H <: PSY.HydroGen,
    D <: HydroCommitmentReservoirBudget,
    S <: PM.AbstractActivePowerModel,
}
    devices = get_available_components(H, sys)

    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerRangeExpressionLB,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerRangeExpressionUB,
        devices,
        model,
        S,
    )

    # Energy Budget Constraint
    add_constraints!(container, EnergyBudgetConstraint, devices, model, S)

    add_feedforward_constraints!(container, model, devices)

    objective_function!(container, devices, model, S)

    add_constraint_dual!(container, sys, model)
    return
end

"""
Construct model for HydroGen with ReservoirStorage Commitment Formulation
"""
function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ArgumentConstructStage,
    model::DeviceModel{H, HydroCommitmentReservoirStorage},
    ::Type{S},
) where {H <: PSY.HydroEnergyReservoir, S <: PM.AbstractPowerModel}
    devices = get_available_components(H, sys)

    add_variables!(
        container,
        ActivePowerVariable,
        devices,
        HydroCommitmentReservoirStorage(),
    )
    add_variables!(
        container,
        ReactivePowerVariable,
        devices,
        HydroCommitmentReservoirStorage(),
    )
    add_variables!(container, OnVariable, devices, HydroCommitmentReservoirStorage())
    add_variables!(container, EnergyVariable, devices, HydroCommitmentReservoirStorage())
    add_variables!(
        container,
        WaterSpillageVariable,
        devices,
        HydroCommitmentReservoirStorage(),
    )
    add_variables!(
        container,
        EnergyShortageVariable,
        devices,
        HydroCommitmentReservoirStorage(),
    )
    add_variables!(
        container,
        EnergySurplusVariable,
        devices,
        HydroCommitmentReservoirStorage(),
    )
    add_variables!(container, EnergyOutput, devices, HydroCommitmentReservoirStorage())
    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ActivePowerBalance,
        ReactivePowerVariable,
        devices,
        model,
        S,
    )

    add_parameters!(container, EnergyTargetTimeSeriesParameter, devices, model)
    add_parameters!(container, InflowTimeSeriesParameter, devices, model)

    add_expressions!(container, ProductionCostExpression, devices, model)

    add_to_expression!(
        container,
        ActivePowerRangeExpressionLB,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ActivePowerRangeExpressionUB,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_feedforward_arguments!(container, model, devices)
    return
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ModelConstructStage,
    model::DeviceModel{H, HydroCommitmentReservoirStorage},
    ::Type{S},
) where {H <: PSY.HydroEnergyReservoir, S <: PM.AbstractPowerModel}
    devices = get_available_components(H, sys)

    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerRangeExpressionLB,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerRangeExpressionUB,
        devices,
        model,
        S,
    )

    add_constraints!(
        container,
        ReactivePowerVariableLimitsConstraint,
        ReactivePowerVariable,
        devices,
        model,
        S,
    )

    add_initial_condition!(
        container,
        devices,
        HydroCommitmentReservoirStorage(),
        InitialEnergyLevel(),
    )
    # Energy Balance Constraint
    add_constraints!(container, EnergyBalanceConstraint, devices, model, S)
    add_constraints!(container, EnergyTargetConstraint, devices, model, S)
    add_feedforward_constraints!(container, model, devices)

    objective_function!(container, devices, model, S)
    add_constraint_dual!(container, sys, model)

    return
end

"""
Construct model for HydroGen with ReservoirStorage Dispatch Formulation
with only Active Power
"""
function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ArgumentConstructStage,
    model::DeviceModel{H, HydroCommitmentReservoirStorage},
    ::Type{S},
) where {H <: PSY.HydroEnergyReservoir, S <: PM.AbstractActivePowerModel}
    devices = get_available_components(H, sys)

    add_variables!(
        container,
        ActivePowerVariable,
        devices,
        HydroCommitmentReservoirStorage(),
    )
    add_variables!(container, OnVariable, devices, HydroCommitmentReservoirStorage())
    add_variables!(container, EnergyVariable, devices, HydroCommitmentReservoirStorage())
    add_variables!(
        container,
        WaterSpillageVariable,
        devices,
        HydroCommitmentReservoirStorage(),
    )
    add_variables!(
        container,
        EnergyShortageVariable,
        devices,
        HydroCommitmentReservoirStorage(),
    )
    add_variables!(
        container,
        EnergySurplusVariable,
        devices,
        HydroCommitmentReservoirStorage(),
    )

    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_variables!(container, EnergyOutput, devices, HydroCommitmentReservoirStorage())
    add_parameters!(container, EnergyTargetTimeSeriesParameter, devices, model)
    add_parameters!(container, InflowTimeSeriesParameter, devices, model)

    add_expressions!(container, ProductionCostExpression, devices, model)

    add_to_expression!(
        container,
        ActivePowerRangeExpressionLB,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ActivePowerRangeExpressionUB,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_feedforward_arguments!(container, model, devices)
    return
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ModelConstructStage,
    model::DeviceModel{H, HydroCommitmentReservoirStorage},
    ::Type{S},
) where {H <: PSY.HydroEnergyReservoir, S <: PM.AbstractActivePowerModel}
    devices = get_available_components(H, sys)

    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerRangeExpressionLB,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerRangeExpressionUB,
        devices,
        model,
        S,
    )

    add_initial_condition!(
        container,
        devices,
        HydroCommitmentReservoirStorage(),
        InitialEnergyLevel(),
    )

    # Energy Balance Constraint
    add_constraints!(container, EnergyBalanceConstraint, devices, model, S)
    add_constraints!(container, EnergyTargetConstraint, devices, model, S)
    add_feedforward_constraints!(container, model, devices)

    objective_function!(container, devices, model, S)
    add_constraint_dual!(container, sys, model)

    return
end

"""
Construct model for HydroPumpedStorage with PumpedStorage Dispatch Formulation
with only Active Power
"""
function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ArgumentConstructStage,
    model::DeviceModel{H, HydroDispatchPumpedStorage},
    ::Type{S},
) where {H <: PSY.HydroPumpedStorage, S <: PM.AbstractActivePowerModel}
    devices = get_available_components(H, sys)

    add_variables!(container, ActivePowerInVariable, devices, HydroDispatchPumpedStorage())
    add_variables!(container, ActivePowerOutVariable, devices, HydroDispatchPumpedStorage())
    add_variables!(container, EnergyVariableUp, devices, HydroDispatchPumpedStorage())
    add_variables!(container, EnergyVariableDown, devices, HydroDispatchPumpedStorage())
    add_variables!(container, WaterSpillageVariable, devices, HydroDispatchPumpedStorage())
    add_variables!(container, EnergyOutput, devices, HydroDispatchPumpedStorage())
    if get_attribute(model, "reservation")
        add_variables!(
            container,
            ReservationVariable,
            devices,
            HydroDispatchPumpedStorage(),
        )
    end

    add_parameters!(container, InflowTimeSeriesParameter, devices, model)
    add_parameters!(container, OutflowTimeSeriesParameter, devices, model)

    add_expressions!(container, ProductionCostExpression, devices, model)

    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerInVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerOutVariable,
        devices,
        model,
        S,
    )

    add_expressions!(container, ReserveRangeExpressionLB, devices, model)
    add_expressions!(container, ReserveRangeExpressionUB, devices, model)

    add_feedforward_arguments!(container, model, devices)
    return
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ModelConstructStage,
    model::DeviceModel{H, HydroDispatchPumpedStorage},
    ::Type{S},
) where {H <: PSY.HydroPumpedStorage, S <: PM.AbstractActivePowerModel}
    devices = get_available_components(H, sys)

    add_constraints!(
        container,
        OutputActivePowerVariableLimitsConstraint,
        ActivePowerOutVariable,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        InputActivePowerVariableLimitsConstraint,
        ActivePowerInVariable,
        devices,
        model,
        S,
    )

    add_initial_condition!(
        container,
        devices,
        HydroDispatchPumpedStorage(),
        InitialEnergyLevelUp(),
    )
    add_initial_condition!(
        container,
        devices,
        HydroDispatchPumpedStorage(),
        InitialEnergyLevelDown(),
    )

    # Energy Balanace limits
    add_constraints!(container, EnergyCapacityUpConstraint, devices, model, S)
    add_constraints!(container, EnergyCapacityDownConstraint, devices, model, S)
    add_feedforward_constraints!(container, model, devices)

    objective_function!(container, devices, model, S)

    add_constraint_dual!(container, sys, model)
    return
end


#  Copyright 2016, Iain Dunning, Joey Huchette, Miles Lubin, and contributors
#  This Source Code Form is subject to the terms of the Mozilla Public
#  License, v. 2.0. If a copy of the MPL was not distributed with this
#  file, You can obtain one at http://mozilla.org/MPL/2.0/.

function writeMPS(m::Model, fname::AbstractString)
    f = open(fname, "w")

    write(f,"NAME   JuMPModel\n")

    numRows = length(m.linconstr)

    # Objective and constraint names
    gc_enable(false)
    write(f,"ROWS\n")
    write(f," N  OBJ\n")
    hasrange = false
    for c in 1:numRows
        rowsense = sense(m.linconstr[c])
        if rowsense == :(<=)
            senseChar = 'L'
        elseif rowsense == :(==)
            senseChar = 'E'
        elseif rowsense == :(>=)
            senseChar = 'G'
        else
            hasrange = true
            senseChar = 'E'
        end
        @printf(f," %c  CON%d\n",senseChar,c)
    end
    gc_enable(true)

    objlincoef, rowlb, rowub = prepProblemBounds(m)
    if m.objSense == :Max
        println("Warning, MPS does not support maximization sense. Flipping objective coefficients.")
        objlincoef = -objlincoef
    end

    A = prepConstrMatrix(m)
    colptr = A.colptr
    rowval = A.rowval
    nzval = A.nzval

    # Output each column
    gc_enable(false)
    inintegergroup = false
    write(f,"COLUMNS\n")
    for col in 1:m.numCols
        t = m.colCat[col]
        (t == :SemiCont || t == :SemiInt) && error("The MPS file writer does not currently support semicontinuous or semi-integer variables")
        if (t == :Bin || t == :Int) && !inintegergroup
            @printf(f,"    MARKER    'MARKER'                 'INTORG'\n")
            inintegergroup = true
        elseif (t == :Cont || t == :Fixed) && inintegergroup
            @printf(f,"    MARKER    'MARKER'                 'INTEND'\n")
            inintegergroup = false
        end
        for ind in colptr[col]:(colptr[col+1]-1)
            @printf(f,"    VAR%d  CON%d  ",col,rowval[ind])
            print_shortest(f,nzval[ind])
            println(f)
        end
        @printf(f,"    VAR%d  OBJ  ",col)
        print_shortest(f,objlincoef[col])
        println(f)
    end
    if inintegergroup
        @printf(f,"    MARKER    'MARKER'                 'INTEND'\n")
    end
    gc_enable(true)

    # RHSs
    gc_enable(false)
    write(f,"RHS\n")
    for c in 1:numRows
        rowsense = sense(m.linconstr[c])
        if rowsense != :range
            @printf(f,"    rhs    CON%d    ",c)
            print_shortest(f,rhs(m.linconstr[c]))
        else
            @printf(f,"    rhs    CON%d    ",c)
            print_shortest(f,m.linconstr[c].lb)
        end
        println(f)
    end
    gc_enable(true)

    # RANGES
    if hasrange
        gc_enable(false)
        write(f,"RANGES\n")
        for c in 1:numRows
            rowsense = sense(m.linconstr[c])
            if rowsense == :range
                @printf(f,"    rhs    CON%d    ",c)
                print_shortest(f,m.linconstr[c].ub-m.linconstr[c].lb)
                println(f)
            end
        end
    end

    # BOUNDS
    gc_enable(false)
    write(f,"BOUNDS\n")
    for col in 1:m.numCols
        if m.colLower[col] == 0
            if m.colUpper[col] != Inf
                # Default lower 0, and an upper
                @printf(f,"  UP BOUND VAR%d ", col)
                print_shortest(f, m.colUpper[col])
                println(f)
            else
                # Default bounds. Explicitly state for solvers like Gurobi. See issue #792
                @printf(f,"  PL BOUND VAR%d", col)
                println(f)
            end
        elseif m.colLower[col] == -Inf && m.colUpper[col] == +Inf
            # Free
            @printf(f, "  FR BOUND VAR%d\n", col)
        elseif m.colLower[col] != -Inf && m.colUpper[col] == +Inf
            # No upper, but a lower
            @printf(f, "  PL BOUND VAR%d\n  LO BOUND VAR%d ",col,col)
            print_shortest(f,m.colLower[col])
            println(f)
        elseif m.colLower[col] == -Inf && m.colUpper[col] != +Inf
            # No lower, but a upper
            @printf(f,"  MI BOUND VAR%d\n  UP BOUND VAR%d ",col,col)
            print_shortest(f,m.colUpper[col])
            println(f)
        else
            # Lower and upper
            @printf(f, "  LO BOUND VAR%d ",col)
            print_shortest(f,m.colLower[col])
            println(f)
            @printf(f, "  UP BOUND VAR%d ",col)
            print_shortest(f,m.colUpper[col])
            println(f)
        end
    end
    gc_enable(true)

    # Quadratic objective
    gc_enable(false)
    if length(m.obj.qvars1) != 0
        write(f,"QMATRIX\n")
        qv1 = m.obj.qvars1
        qv2 = m.obj.qvars2
        qc  = m.obj.qcoeffs
        for ind = 1:length(qv1)
            if qv1[ind].col == qv2[ind].col
                # Diagonal element
                @printf(f,"  VAR%d VAR%d  ", qv1[ind].col,qv2[ind].col)
                print_shortest(f,2qc[ind])
                println(f)
            else
                # Off diagonal, and we're gonna assume no duplicates
                @printf(f, "  VAR%d VAR%d ", qv1[ind].col,qv2[ind].col)
                print_shortest(f, qc[ind])
                println(f)
                #@printf(f, "  VAR%d VAR%d ", qv2[ind].col,qv1[ind].col)
                #print_shortest(f, qc[ind])
                #println(f)
            end
        end
    end

    write(f,"ENDATA\n")
    close(f)
    gc_enable(true)
end

###############################################################################
# LP File Writer
# We use the formatting defined at:
#   http://lpsolve.sourceforge.net/5.0/CPLEX-format.htm
function writeLP(m::Model, fname::AbstractString)

    f = open(fname, "w")

    if length(m.obj.qvars1) != 0
        error("LP writer does not support quadratic objectives.\n")
    end

    # Objective
    if m.objSense == :Max
        write(f,"Maximize\n")
    else
        write(f,"Minimize\n")
    end
    objaff::AffExpr = m.obj.aff
    write(f, " obj: ")
    nnz = length(objaff.coeffs)
    for ind in 1:(nnz-1)
        if ind == 1
            print_shortest(f, objaff.coeffs[ind])
        else
            print_shortest(f, abs(objaff.coeffs[ind]))
        end
        @printf(f, " VAR%d %s ", objaff.vars[ind].col, (objaff.coeffs[ind+1] < 0)? "-" : "+")
    end
    if nnz >= 1
        if nnz == 1
            print_shortest(f, objaff.coeffs[nnz])
        else
            print_shortest(f, abs(objaff.coeffs[nnz]))
        end
        @printf(f, " VAR%d\n", objaff.vars[nnz].col)
    end

    # Constraints
    function writeconstrterms(c::LinearConstraint)
        nnz = length(c.terms.coeffs)
        for ind in 1:(nnz-1)
            if ind == 1
                print_shortest(f, c.terms.coeffs[ind])
            else
                print_shortest(f, abs(c.terms.coeffs[ind]))
            end
            @printf(f, " VAR%d %s ", c.terms.vars[ind].col, (c.terms.coeffs[ind+1] < 0)? "-" : "+")
        end
        if nnz >= 1
            if nnz == 1
                print_shortest(f, c.terms.coeffs[nnz])
            else
                print_shortest(f, abs(c.terms.coeffs[nnz]))
            end
            @printf(f, " VAR%d", c.terms.vars[nnz].col)
        end
    end
    write(f,"Subject To\n")
    constrcount = 1
    for i in 1:length(m.linconstr)
        @printf(f, " c%d: ", constrcount)

        c::LinearConstraint = m.linconstr[i]
        rowsense = sense(c)
        if rowsense != :range
            writeconstrterms(c)
            if rowsense == :(==)
                @printf(f, " = ")
                print_shortest(f, rhs(c))
                println(f)
            elseif rowsense == :<=
                @printf(f, " <= ")
                print_shortest(f, rhs(c))
                println(f)
            else
                @assert rowsense == :>=
                @printf(f, " >= ")
                print_shortest(f, rhs(c))
                println(f)
            end
            constrcount += 1
        else
            writeconstrterms(c)
            @printf(f, " >= ")
            print_shortest(f, c.lb)
            println(f)
            @printf(f, " c%d: ", constrcount+1)
            writeconstrterms(c)
            @printf(f, " <= ")
            print_shortest(f, c.ub)
            println(f)
            constrcount += 2
        end
    end

    # Bounds
    write(f,"Bounds\n")
    for i in 1:m.numCols
        if m.colLower[i] == -Inf
            # No low bound
            if m.colUpper[i] == +Inf
                # Free
                @printf(f, " VAR%d free\n", i)
            else
                # x <= finite
                @printf(f, " -inf <= VAR%d <= ", i)
                print_shortest(f, m.colUpper[i])
                println(f)
            end
        else
            # Low bound exists
            if m.colUpper[i] == +Inf
                # x >= finite
                @printf(f, " ")
                print_shortest(f, m.colLower[i])
                @printf(f," <= VAR%d <= +inf\n", i)
            else
                # finite <= x <= finite
                @printf(f, " ")
                print_shortest(f, m.colLower[i])
                @printf(f, " <= VAR%d <= ", i)
                print_shortest(f, m.colUpper[i])
                println(f)
            end
        end
    end

    # Integer - don't handle binaries specially
    write(f,"General\n")
    for i in 1:m.numCols
        t = m.colCat[i]
        (t == :SemiCont || t == :SemiInt) && error("The LP file writer does not currently support semicontinuous or semi-integer variables")
        if t == :Bin || t == :Int
            @printf(f, " VAR%d\n", i)
        end
    end

    # Done
    write(f,"End\n")
    close(f)
end


# fitnessPath: Synapse id, file or folder (should point to MyProject/Metadata).
#              If it's a folder, all .csv files in the folder will be checked.
function getsamplefitness(syn, fitnessPath)
    fitness = getsamplemetadata(syn, fitnessPath) # hijack metadata loading which is very similar
    rename!(fitness, :MetadataID, :FitnessID)
end

function appendfitness!(metadata::DataFrame, fitnessTable::DataFrame; allowInexactMatches=true)
    ids = metadata[:SampleID]
    @assert length(unique(ids)) == length(ids)
    ids = dropna(fitnessTable[:SampleID])
    @assert length(unique(ids)) == length(ids)
    @assert eltype(metadata[:Passage])<:Integer
    @assert eltype(fitnessTable[:Passage])<:Integer

    fitnessTable = copy(fitnessTable) # create copy since we will modify the table
    fitness = hcat(fitnessTable[:A],fitnessTable[:B],fitnessTable[:C])  # create Nx3 DataArray
    fitness = map(i->mean(dropna(fitness[i,:][:])), 1:size(fitness,1)); # take mean of each row after dropping NAs
    fitnessTable[:Fitness] = fitness
    fitnessTable = fitnessTable[.~isnan.(fitnessTable[:Fitness]), :] # drop measurements where 3 fitness values for a sample are NAs

    # create new columns
    metadata[:Fitness] = 0.0
    metadata[:FitnessID] = ""

    # find exact SampleID matches between metadata and fitnessTable
    matches = indexin(metadata[:SampleID], map(string,fitnessTable[:SampleID])) # converting to string converts NA->"NA" which can be handled by indexin
    hasMatches = matches.!=0
    matches = matches[hasMatches]

    metadata[hasMatches, :Fitness]   = fitnessTable[matches, :Fitness]
    metadata[hasMatches, :FitnessID] = fitnessTable[matches, :FitnessID]

    if allowInexactMatches
        fitnessTable = fitnessTable[setdiff(1:end,matches),:] # remove those fitness entries that have been used
        used = falses(size(fitnessTable,1)) # but from now on, keep track of which ones we have used rather than removing rows in the loop

        for i in find(ismissing.(metadata[:Fitness])) # for those samples we don't have a fitness value for
            if ismissing(metadata[i,:Reference]) || ismissing(metadata[i,:Mutagen]) || ismissing(metadata[i,:Dose]) || ismissing(metadata[i,:Passage])
                continue # skip samples with incomplete metadata
            end

            mask = (metadata[i,:Reference] .== fitnessTable[:Virus]) .&
                   (metadata[i,:Mutagen]   .== fitnessTable[:Mutagen]) .&
                   (metadata[i,:Dose]      .== fitnessTable[:Dose]) .&
                   (metadata[i,:Passage]   .== fitnessTable[:Passage]) .&
                   .~used

            ind = findfirst(mask) # if there are multiple matches, use the first one and leave others to match other samples
            ind==0 && continue # no match

            metadata[i,:Fitness]   = fitnessTable[ind,:Fitness]
            metadata[i,:FitnessID] = fitnessTable[ind,:FitnessID]
            used[ind] = true
        end

        # fitnessTable = fitnessTable[.~used,:]
    end

    metadata
end


# *********************************************************************************
# REopt, Copyright (c) 2019-2020, Alliance for Sustainable Energy, LLC.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without modification,
# are permitted provided that the following conditions are met:
#
# Redistributions of source code must retain the above copyright notice, this list
# of conditions and the following disclaimer.
#
# Redistributions in binary form must reproduce the above copyright notice, this
# list of conditions and the following disclaimer in the documentation and/or other
# materials provided with the distribution.
#
# Neither the name of the copyright holder nor the names of its contributors may be
# used to endorse or promote products derived from this software without specific
# prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
# IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
# INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
# OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED
# OF THE POSSIBILITY OF SUCH DAMAGE.
# *********************************************************************************

"""
    AbsorptionChiller

struct with inner constructor:
```julia
function AbsorptionChiller(;
    min_ton::Real = 0.0,
    max_ton::Real = 0.0,
    chiller_cop::Real,
    chiller_elec_cop::Real = 14.1,
    installed_cost_per_ton::Real,
    om_cost_per_ton::Real,
    macrs_option_years::Real = 0,
    macrs_bonus_pct::Real = 0
)
```
"""
struct AbsorptionChiller <: AbstractThermalTech
    min_ton::Real
    max_ton::Real
    chiller_cop::Real
    chiller_elec_cop::Real
    installed_cost_us_dollars_per_ton::Real
    om_cost_us_dollars_per_ton::Real
    macrs_option_years::Real
    macrs_bonus_pct::Real
    min_kw::Real
    max_kw::Real
    installed_cost_per_kw::Real
    om_cost_per_kw::Real
    function AbsorptionChiller(;
        min_ton::Real = 0.0,
        max_ton::Real = 0.0,
        chiller_cop::Real,
        chiller_elec_cop::Real = 14.1,
        installed_cost_per_ton::Real,
        om_cost_per_ton::Real,
        macrs_option_years::Real = 0,
        macrs_bonus_pct::Real = 0,
        )

        min_kw = min_ton * KWH_THERMAL_PER_TONHOUR
        max_kw = max_ton * KWH_THERMAL_PER_TONHOUR
        installed_cost_per_kw = installed_cost_per_ton / KWH_THERMAL_PER_TONHOUR
        om_cost_per_kw = om_cost_per_ton / KWH_THERMAL_PER_TONHOUR

        new(
            min_ton,
            max_ton,
            chiller_cop,
            chiller_elec_cop,
            installed_cost_per_ton,
            om_cost_per_ton,
            macrs_option_years,
            macrs_bonus_pct,
            min_kw,
            max_kw,
            installed_cost_per_kw,
            om_cost_per_kw
        )
    end
end


#####################
# Downward Flow circuits
#####################

"Representation of downward flows with `in_progress` Bool for effective book keeping"

abstract type DownFlowΔNode{O,F} <: DecoratorΔNode{UpFlowΔNode{O,F}} end
abstract type DownFlowInnerNode{O,F} <: DownFlowΔNode{O,F} end

abstract type DownFlow⋀{O,F} <: DownFlowInnerNode{O,F} end
abstract type DownFlow⋁{O,F} <: DownFlowInnerNode{O,F} end

# use this version of DownFlow⋀ when ⋀ gates are unique to their parent ⋁ to save space (compact=true)
mutable struct DownFlow{F}
    downflow::F
    in_progress::Bool
end

struct DownFlowLeaf{O,F} <: DownFlowΔNode{O,F}
    origin::UpFlowLeafNode{O,F}
    downflow::DownFlow{F}
end

struct DownFlow⋀Compact{O,F} <: DownFlow⋀{O,F}
    origin::UpFlow⋀{O,F}
    children::Vector{<:DownFlowΔNode{<:O,F}}
    cached_downflow_sinks::Vector{DownFlow{F}}
end

# use this version of DownFlow⋀ when ⋀ gates can have multiple parents to save computation (compact=false)
struct DownFlow⋀Cached{O,F} <: DownFlow⋀{O,F}
    origin::UpFlow⋀{<:O,<:F}
    children::Vector{<:DownFlowΔNode{<:O,F}}
    downflow::DownFlow{F}
end

# this version of DownFlow⋁ is efficient for arity-1 gates
struct DownFlow⋁Compact{O,F} <: DownFlow⋁{O,F}
    origin::UpFlow⋁{<:O,<:F}
    child::DownFlowΔNode{<:O,F}
    cached_downflow_sinks::Vector{DownFlow{F}}
end

struct DownFlow⋁Cached{O,F} <: DownFlow⋁{O,F}
    origin::UpFlow⋁{<:O,<:F}
    children::Vector{<:DownFlowΔNode{<:O,F}}
    downflow::DownFlow{F}
end

const DownFlowΔ{O,F} = AbstractVector{<:DownFlowΔNode{O,F}}

const FlowΔNode{O,F} = DownFlowΔNode{O,F}
const FlowΔ{O,F} = AbstractVector{<:FlowΔNode{O,F}}

#####################
# traits
#####################

@inline GateType(::Type{<:DownFlowLeaf{O}}) where O = GateType(O)
@inline GateType(::Type{<:DownFlow⋀}) = ⋀Gate()
@inline GateType(::Type{<:DownFlow⋁}) = ⋁Gate()

const HasDownFlow = Union{DownFlow⋁Cached, DownFlow⋀Cached, DownFlowLeaf}

#####################
# constructors and conversions
#####################

"""Construct a downward flow circuit from a given upward flow circuit"""
function DownFlowΔ(circuit::UpFlowΔ{O,F}, opts = flow_opts★)::DownFlowΔ{O,F} where {O,F}

    m = flow_length(circuit)
    fmem  = () -> some_vector(eltype(F), m)

    cache = Dict{ΔNode, DownFlowΔNode}()
    sizehint!(cache, length(circuit)*4÷3)

    flow_node(::LeafGate, n::ΔNode) = DownFlowLeaf{O,F}(n, DownFlow(fmem(),false))

    flow_node(::⋀Gate, n::ΔNode) = begin
        n_children = map(c -> cache[c], children(n))
        @assert length(n_children)>=2
        sinks = downflow_sinks(n_children)
        if opts.compact⋀ && length(sinks) <= opts.max_factors
            DownFlow⋀Compact{O,F}(n, n_children, sinks)
        else
            DownFlow⋀Cached{O,F}(n, n_children, DownFlow(fmem(),false))
        end
    end

    flow_node(::⋁Gate, n::ΔNode) = begin
        n_children = map(c -> cache[c], children(n))
        @assert length(n_children)>=1
        sinks = downflow_sinks(n_children)
        if opts.compact⋁ && length(n_children)==1 && length(sinks) <= opts.max_factors
            DownFlow⋁Compact{O,F}(n, n_children[1], sinks)
        else
            DownFlow⋁Cached{O,F}(n, n_children, DownFlow(fmem(),false))
        end
    end

    map(circuit) do node
        fnode = flow_node(GateType(node), node)
        cache[node] = fnode
        fnode
    end
end

"""Construct a up and down flow circuit from a given other circuit"""
function FlowΔ(circuit::Δ, m::Int, ::Type{El}, opts = flow_opts★) where El
    up_flow_circuit = UpFlowΔ(circuit, m, El, opts)
    DownFlowΔ(up_flow_circuit, opts)
end

#####################
# methods
#####################

@inline literal(n::DownFlowLeaf)::Lit  = literal(n.origin)
@inline constant(n::DownFlowLeaf)::Bool = constant(n.origin)
@inline children(n::DownFlow⋁Compact) = [n.child]
@inline children(n::Union{DownFlow⋀,DownFlow⋁Cached}) = n.children

@inline downflow(n::HasDownFlow) = n.downflow.downflow

@inline downflow_sinks(n::DownFlowInnerNode) = n.cached_downflow_sinks
@inline downflow_sinks(n::HasDownFlow) = [n.downflow]
# @inline downflow_sinks(n::DownFlowLeaf{O,F}) where {O,F} = Vector{DownFlow{F}}()
@inline function downflow_sinks(ns::Vector{<:DownFlowΔNode{O,F}})::Vector{DownFlow{F}} where {O,F}
    flatmap(c -> downflow_sinks(c)::Vector{DownFlow{F}}, ns, DownFlow{F}[])
end

pr(n::DownFlowΔNode) = pr(origin(n))
pr_factors(n::DownFlowΔNode) = pr_factors(origin(n))

flow_length(circuit::DownFlowΔ) = length(downflow_sinks(circuit[end])[1].downflow)
origin_flow_length(circuit::DownFlowΔ) = flow_length(origin(circuit[end]))

#####################

# change the size of the flow vectors throughout this circuit (e.g., change minibatch size)
function resize_flows(circuit::DownFlowΔ{F}, size::Int) where {F}
    # check if resize is needed
    if size != flow_length(circuit)
        for n in circuit
            resize_path_flow(n,size)
        end
    end
end

@inline resize_path_flow(n::DownFlowΔNode, size) = () #do nothing
@inline resize_path_flow(n::HasDownFlow, size) = resize!(n.downflow.downflow, size)

#####################

function pass_down(circuit::DownFlowΔ{O,F}) where {F,O}
    resize_flows(circuit, flow_length(origin(circuit)))
    for n in circuit
        reset_downflow_in_progress(n)
    end
    for downflow in downflow_sinks(circuit[end])
        # initialize root flows to 1
        downflow.downflow .= one(eltype(F))
    end
    for n in Iterators.reverse(circuit)
        pass_down_node(n)
    end
end

reset_downflow_in_progress(n::DownFlowΔNode) = () # do nothing
reset_downflow_in_progress(n::HasDownFlow) = (n.downflow.in_progress = false)

pass_down_node(n::DownFlowΔNode) = () # do nothing
pass_down_node(n::DownFlowLeaf) = () # too costly? @assert pr(n) ≈ downflow(n) "Postitive leaf flows should be equal to their probabilities: $(pr(n)) ≆ $(downflow(n))"

function pass_down_node(n::DownFlow⋀Cached)
    for c in n.children
        for sink in downflow_sinks(c)
            if !sink.in_progress
                assign(sink.downflow, downflow(n))
                sink.in_progress = true
            else
                accumulate_val(sink.downflow, downflow(n))
            end
        end
    end
end

function pass_down_node(n::DownFlow⋁Cached)
    if num_children(n) == 1
        for sink in downflow_sinks(n.children[1])
            if !sink.in_progress
                assign(sink.downflow, downflow(n))
                sink.in_progress = true
            else
                accumulate_val(sink.downflow, downflow(n))
            end
        end
    else
        for c in n.children
            pr_fs = pr_factors(c)
            for sink in downflow_sinks(c)
                if !sink.in_progress
                    assign_prod_normalized(sink.downflow, pr(n), downflow(n), pr_fs)
                    sink.in_progress = true
                else
                    accumulate_prod_normalized(sink.downflow, pr(n), downflow(n), pr_fs)
                end
            end
        end
    end
end

#####################

function pass_up_down(circuit::DownFlowΔ{O,F}, data::XData{E}) where {E <: eltype(F)} where {F,O}
    pass_up(origin(circuit), data)
    pass_down(circuit)
end


# ------------------------------------------------------------------
# Licensed under the ISC License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    CompositeTask(tasks, [adjacency])

A learning task composed of multiple sub-`tasks`.
The tasks are assumed to be independent, unless
an `adjacency` matrix is provided with the list
of dependencies.

## Examples

A composite task composed of three independent sub-tasks
`t₁`, `t₂` and `t₃`:

```julia
julia> CompositeTask([t₁,t₂,t₃])
```

A composite sequential task `t₁ → t₂ → t₃`:

```julia
julia> CompositeTask([t₁,t₂,t₃], [0 1 0; 0 0 1; 0 0 0])
```

### Notes

Only directed acyclic graphs (DAGs) are permitted.
"""
struct CompositeTask <: AbstractLearningTask
  tasks::Vector{AbstractLearningTask}
  adjacency::Matrix{Int}

  function CompositeTask(tasks, adjacency)
    @assert all(length(tasks) .== size(adjacency)) "invalid adjacency for tasks"
    @assert isacyclic(adjacency) "dependencies must be acyclic"
    new(tasks, adjacency)
  end
end

CompositeTask(tasks::AbstractVector) =
  CompositeTask(tasks, zeros(Int, length(tasks), length(tasks)))

iscomposite(task::CompositeTask) = true

"""
    isacyclic(adjacency)

Check whether or not the `adjacency` matrix is acyclic
using Kahn's topological sort.
"""
function isacyclic(adjacency::AbstractMatrix{Int})
  # copy input to avoid side effects
  A = copy(adjacency)

  # find root nodes
  roots = Vector{Int}()
  for j in 1:size(A, 2)
    if all(A[:,j] .== 0)
      push!(roots, j)
    end
  end

  # Kahn's algorithm
  sorted = Vector{Int}()
  while !isempty(roots)
    i = pop!(roots)
    push!(sorted, i)

    # for all edges i → j
    for j in findall(A[i,:] .== 1)
      # remove edge i → j
      A[i,j] = 0

      # if j has no other incoming edge
      if all(A[:,j] .== 0)
        push!(sorted, j)
      end
    end
  end

  any(A .== 1) ? false : true
end


promote_rule(::Type{FastRational{T,IsReduced}}, ::Type{T}) where {T<:BitInteger} = FastRational{T,IsReduced}
promote_rule(::Type{FastRational{T,MayReduce}}, ::Type{T}) where {T<:BitInteger} = FastRational{T,IsReduced}
promote_rule(::Type{FastRational{T,IsReduced}}, ::Type{Rational{T}}) where {T<:BitInteger} = FastRational{T,IsReduced}
promote_rule(::Type{FastRational{T,MayReduce}}, ::Type{Rational{T}}) where {T<:BitInteger} = FastRational{T,IsReduced}

promote_rule(::Type{FastRational{T,IsReduced}}, ::Type{A}) where {T,A} = FastRational{T,IsReduced}
promote_rule(::Type{FastRational{T,MayReduce}}, ::Type{A}) where {T,A} = FastRational{T,IsReduced}

convert(::Type{FastRational{T,IsReduced}}, x::Rational{T}) where {T} = FastRational{T,IsReduced}(x.num, x.den)
convert(::Type{FastRational{T,MayReduce}}, x::Rational{T}) where {T} = FastRational{T,IsReduced}(x.num, x.den)
convert(::Type{Rational{T}}, x::FastRational{T,MayReduce}) where {T} = Rational{T}(canonical(x.num, x.den))
convert(::Type{Rational{T}}, x::FastRational{T,IsReduced}) where {T} = Rational{T}(x.num, x.den)

convert(::Type{FastRational{T,IsReduced}}, x::T) where {T} = FastRational{T,IsReduced}(x, one(T))
convert(::Type{FastRational{T,MayReduce}}, x::T) where {T} = FastRational{T,IsReduced}(x, one(T))