const HUnitLowerTriangular = UnitLowerTriangular{<:Any,<:HMatrix}
const HUpperTriangular     = UpperTriangular{<:Any,<:HMatrix}

function ldiv!(L::HUnitLowerTriangular, B::AbstractMatrix)
    H = parent(L)
    if isleaf(H)
        d = data(H)
        ldiv!(UnitLowerTriangular(d), B) # B <-- L\B
    else
        @assert !hasdata(H) "only leaves are allowed to have data when using `ldiv`!"
        shift = pivot(H) .- 1
        chdH  = children(H)
        m, n   = size(chdH)
        @assert m === n
        for i = 1:m
            irows  = colrange(chdH[i,i]) .- shift[2]
            bi     = view(B, irows, :)
            for j = 1:(i - 1)# j<i
                jrows  = colrange(chdH[i,j]) .- shift[2]
                bj     = view(B, jrows, :)
                _mul131!(bi, chdH[i,j], bj, -1)
            end
            # recursion stage
            ldiv!(UnitLowerTriangular(chdH[i,i]), bi)
        end
    end
    return B
end

function ldiv!(L::HUnitLowerTriangular, R::RkMatrix)
    ldiv!(L, R.A) # change R.A in-place
    return R
end

function ldiv!(L::HUnitLowerTriangular, X::HMatrix, compressor)
    H = parent(L)
    @assert isclean(H)
    if isleaf(X)
        d = data(X)
        @timeit_debug "dense ldiv!" ldiv!(L, d)
    elseif isleaf(H) # X not a leaf, but L is a leaf. This should not happen.
        error()
    else
        chdH = children(H)
        chdX = children(X)
        m, n  = size(chdH)
        @assert m == n
        for k in 1:size(chdX, 2)
            for i = 1:m
                for j = 1:(i - 1)# j<i
                    @timeit_debug "hmul!" begin
                        hmul!(chdX[i,k], chdH[i,j], chdX[j,k], -1, 1, compressor)
                    end
                end
                ldiv!(UnitLowerTriangular(chdH[i,i]), chdX[i,k], compressor)
            end
        end
    end
    return X
end

function ldiv!(U::HUpperTriangular, B::AbstractMatrix)
    H = parent(U)
    if isleaf(H)
        d = data(H)
        ldiv!(UpperTriangular(d), B) # B <-- L\B
    else
        @assert !hasdata(H) "only leaves are allowed to have data when using `ldiv`!"
        shift = pivot(H) .- 1
        chdH  = children(H)
        m, n   = size(chdH)
        @assert m === n
        for i = m:-1:1
            irows  = colrange(chdH[i,i]) .- shift[2]
            bi     = view(B, irows, :)
            for j = i+1:n # j>i
                jrows  = colrange(chdH[i,j]) .- shift[2]
                bj     = view(B, jrows, :)
                _mul131!(bi, chdH[i,j], bj, -1)
            end
            # recursion stage
            ldiv!(UpperTriangular(chdH[i,i]), bi)
        end
    end
    return B
end

# 1.3
function rdiv!(B::StridedMatrix, U::HUpperTriangular)
    H = parent(U)
    if isleaf(H)
        d = data(H)
        rdiv!(B, UpperTriangular(d)) # b <-- b/L
    else
        @assert !hasdata(H) "only leaves are allowed to have data when using `rdiv`!"
        shift = pivot(H) .- 1 |> reverse
        chdH  = children(H)
        m, n   = size(chdH)
        @assert m === n
        for i = 1:m
            icols  = rowrange(chdH[i,i]) .- shift[1]
            bi     = view(B, :, icols)
            for j = 1:(i - 1)
                jcols  = rowrange(chdH[j,i]) .- shift[1]
                bj     = view(B, :, jcols)
                _mul113!(bi, bj, chdH[j,i], -1)
            end
            # recursion stage
            rdiv!(bi, UpperTriangular(chdH[i,i]))
        end
    end
    return B
end

# 2.3
function rdiv!(R::RkMatrix, U::HUpperTriangular)
    Bt = rdiv!(Matrix(R.Bt), U)
    adjoint!(R.B, Bt)
    return R
end

# 3.3
function rdiv!(X::AbstractHMatrix, U::HUpperTriangular, compressor)
    H = parent(U)
    if isleaf(X)
        d = data(X)
        @timeit_debug "dense rdiv!" begin
            rdiv!(d, U) # b <-- b/L
        end
    elseif isleaf(H)
        error()
    else
        @assert !hasdata(H) # only leaves are allowed to have data for the inversion
        chdX = children(X)
        chdH = children(H)
        m, n = size(chdH)
        for k in 1:size(chdX, 1)
            for i = 1:m
                for j = 1:(i - 1)
                    @timeit_debug "hmul!" begin
                        hmul!(chdX[k,i], chdX[k,j], chdH[j,i], -1, 1, compressor)
                    end
                end
                rdiv!(chdX[k,i], UpperTriangular(chdH[i,i]), compressor)
            end
        end
    end
    return X
end


# flips bits between ith and jth bit in n
flipbits(n, i, j) = n ⊻ (2^(j - i + 1) - 1) << (i - 1)

# checks if given pattern has same number of home and away games
isPattern(nslots, pat) =
    count_ones(pat) == count_zeros(pat) - sizeof(pat) * 8 + nslots

ishard(c::T) where {T<:Constraint} = c.pen == 0

@inline teams(c::Constraint) = c.teams
@inline teams(c::CAcon) = IdxSet(c.teams1.binrepr | c.teams2.binrepr)
@inline teams(c::GAcon) = unique(Iterators.flatten(c.meetings))

@inline patconteam(c::CAcon) = only(c.teams1) # get team for a patetrn con
@inline patconteam(c::BRcon) = only(c.teams)

@inline patsetconteams(c::CAcon) = c.teams1
@inline patsetconteams(c::BRcon) = c.teams

@inline slots(c::Constraint) = c.slots

@inline constraints(i::Instance) =
    vcat(i.constrCA, i.constrBR, i.constrGA, i.constrFA, i.constrSE)

# base case
@inline is_pattern_constraint(rr2::RR2, con::Constraint) = false

@inline is_pattern_constraint(rr2::RR2, con::BRcon) =
    is_pattern_constraint(con::BRcon)
@inline is_pattern_constraint(con::BRcon) = length(con.teams) == 1
@inline is_pattern_constraint(rr2::RR2, con::CAcon) =
    length(con.teams1) == 1 && length(con.teams2) == rr2.nteams

@inline is_pattern_set_constraint(rr2::RR2, con::Constraint) = false
@inline is_pattern_set_constraint(rr2::RR2, con::CAcon) =
    length(con.teams1) > 1 && length(con.teams2) == rr2.nteams
@inline is_pattern_set_constraint(rr2::RR2, con::BRcon) = length(con.teams) > 1

function patternset(sched::Matrix{Int})
    return [
        sum(2^(s - 1) for s in 1:size(sched, 2) if sign(sched[i, s]) == 1) for
        i = 1:size(sched, 1)
    ]
end

"""
    is_schedule(rr2::RR2, S::Array{Int, 2})

Checks if a schedule satisfies all base constraints.

# Arguments

- `rr2::RR2`: Format data of the problem instance
- `S::Array{Int,2}`: The potential schedule to be verified. Assuming teams `i`
    and `j` meet in slot `s` we have `S[i,s] == j`, if team `i` plays at home
    and `S[i,s] == -j` if team `i` plays away.

# Notes

This function is purely intended for debugging purposes.
"""
function is_schedule(rr2::RR2, S::Array{Int,2})
    @assert size(S, 1) == rr2.nteams && size(S, 2) == rr2.nslots

    for s = 1:rr2.nslots
        # exactly rr2.nteams÷2 teams play at home in each slot
        if count(x -> sign(x) == 1, S[:, s]) != rr2.nteams ÷ 2
            return false
        end
        # each team participates in each slot exactly once
        if sort(abs.(S[:, s])) != 1:rr2.nteams
            return false
        end
    end

    # each team plays against each other team once at home and once away
    for r = 1:rr2.nteams
        if sort(S[r, :]) !=
           vcat(-rr2.nteams:-r-1, -r+1:-1, 1:r-1, r+1:rr2.nteams)
            return false
        end
    end
    # check if first half of tournament is a single round robin
    if rr2.isphased
        for i = 1:rr2.nteams
            if sort(abs.(S[i, 1:rr2.nslots÷2])) != vcat(1:i-1, i+1:rr2.nteams)
                return false
            end
        end
    end

    return true
end


### A Hamiltonian is composed of blocks (kpoints), which have a list of RealFourierOperator
# corresponding to each term
# This is the "high-level" interface, provided for convenience

abstract type HamiltonianBlock end

# Generic HamiltonianBlock
struct GenericHamiltonianBlock <: HamiltonianBlock
    basis::PlaneWaveBasis
    kpoint::Kpoint

    # The operators are vectors of RealFourierOperator,
    # not typed because of type invariance issues.
    operators::Vector  # the original list of RealFourierOperator
                       # (as many as there are terms), kept for easier exploration
    optimized_operators::Vector  # Optimized list of RealFourierOperator, for application
end

# More optimized HamiltonianBlock for the important case of a DFT Hamiltonian
struct DftHamiltonianBlock <: HamiltonianBlock
    basis::PlaneWaveBasis
    kpoint::Kpoint
    operators::Vector

    # Individual operators for easy access
    fourier_op::FourierMultiplication
    local_op::RealSpaceMultiplication
    nonlocal_op::Union{Nothing,NonlocalOperator}
    divAgrad_op::Union{Nothing,DivAgradOperator}

    scratch  # Pre-allocated scratch arrays for fast application
end

function HamiltonianBlock(basis, kpoint, operators, scratch=ham_allocate_scratch_(basis))
    optimized_operators = optimize_operators_(operators)
    fourier_ops  = filter(o -> o isa FourierMultiplication,   optimized_operators)
    real_ops     = filter(o -> o isa RealSpaceMultiplication, optimized_operators)
    nonlocal_ops = filter(o -> o isa NonlocalOperator,        optimized_operators)
    divAgrid_ops = filter(o -> o isa DivAgradOperator,        optimized_operators)

    is_dft_ham = (   length(fourier_ops) == 1 && length(real_ops) == 1
                  && length(nonlocal_ops) < 2 && length(divAgrid_ops) < 2)
    if is_dft_ham
        nonlocal_op = isempty(nonlocal_ops) ? nothing : only(nonlocal_ops)
        divAgrid_op = isempty(divAgrid_ops) ? nothing : only(divAgrid_ops)
        DftHamiltonianBlock(basis, kpoint, operators,
                            only(fourier_ops), only(real_ops),
                            nonlocal_op, divAgrid_op, scratch)
    else
        GenericHamiltonianBlock(basis, kpoint, operators, optimized_operators)
    end
end
function ham_allocate_scratch_(basis::PlaneWaveBasis{T}) where {T}
    (ψ_reals=[zeros(complex(T), basis.fft_size...) for _ = 1:Threads.nthreads()], )
end

Base.:*(H::HamiltonianBlock, ψ) = mul!(similar(ψ), H, ψ)
Base.eltype(block::HamiltonianBlock) = complex(eltype(block.basis))
Base.size(block::HamiltonianBlock, i::Integer) = i < 3 ? size(block)[i] : 1
function Base.size(block::HamiltonianBlock)
    n_G = length(G_vectors(block.basis, block.kpoint))
    (n_G, n_G)
end

import Base: Matrix, Array
Array(block::HamiltonianBlock)  = Matrix(block)
Matrix(block::HamiltonianBlock) = sum(Matrix, block.operators)
Matrix(block::GenericHamiltonianBlock) = sum(Matrix, block.optimized_operators)

struct Hamiltonian
    basis::PlaneWaveBasis
    blocks::Vector{HamiltonianBlock}
end

function LinearAlgebra.mul!(Hψ, H::Hamiltonian, ψ)
    for ik = 1:length(H.basis.kpoints)
        mul!(Hψ[ik], H.blocks[ik], ψ[ik])
    end
    Hψ
end
# need `deepcopy` here to copy the elements of the array of arrays ψ (not just pointers)
Base.:*(H::Hamiltonian, ψ) = mul!(deepcopy(ψ), H, ψ)

# Loop through bands, IFFT to get ψ in real space, loop through terms, FFT and accumulate into Hψ
# For the common DftHamiltonianBlock there is an optimized version below
@views @timing "Hamiltonian multiplication" function LinearAlgebra.mul!(Hψ::AbstractArray,
                                                                        H::GenericHamiltonianBlock,
                                                                        ψ::AbstractArray)
    T = eltype(H.basis)
    n_bands = size(ψ, 2)
    Hψ_fourier = similar(Hψ[:, 1])
    ψ_real  = zeros(complex(T), H.basis.fft_size...)
    Hψ_real = zeros(complex(T), H.basis.fft_size...)

    # take ψi, IFFT it to ψ_real, apply each term to Hψ_fourier and Hψ_real, and add it to Hψ
    for iband = 1:n_bands
        Hψ_real .= 0
        Hψ_fourier .= 0
        G_to_r!(ψ_real, H.basis, H.kpoint, ψ[:, iband])
        for op in H.optimized_operators
            @timing "$(nameof(typeof(op)))" begin
                apply!((fourier=Hψ_fourier, real=Hψ_real),
                       op,
                       (fourier=ψ[:, iband], real=ψ_real))
            end
        end
        Hψ[:, iband] .= Hψ_fourier
        r_to_G!(Hψ_fourier, H.basis, H.kpoint, Hψ_real)
        Hψ[:, iband] .+= Hψ_fourier
    end

    Hψ
end

# Fast version, specialized on DFT models. Minimizes the number of FFTs and allocations
@views @timing "DftHamiltonian multiplication" function LinearAlgebra.mul!(Hψ::AbstractArray,
                                                                           H::DftHamiltonianBlock,
                                                                           ψ::AbstractArray)
    n_bands = size(ψ, 2)
    have_divAgrad = !isnothing(H.divAgrad_op)

    # Notice that we use unnormalized plans for extra speed
    potential = H.local_op.potential / prod(H.basis.fft_size)

    # parallelize the loop by breaking it into nthreads() chunks. The ψ_reals are chunk-local
    chunk_length = cld(n_bands, Threads.nthreads())
    @sync for (ichunk, chunk) in enumerate(Iterators.partition(1:n_bands, chunk_length))
        Threads.@spawn for iband in chunk # spawn a task per chunk
            to = TimerOutput()  # Thread-local timer output
            ψ_real = H.scratch.ψ_reals[ichunk]

            @timeit to "local+kinetic" begin
                G_to_r!(ψ_real, H.basis, H.kpoint, ψ[:, iband]; normalize=false)
                ψ_real .*= potential
                r_to_G!(Hψ[:, iband], H.basis, H.kpoint, ψ_real; normalize=false)  # overwrites ψ_real
                Hψ[:, iband] .+= H.fourier_op.multiplier .* ψ[:, iband]
            end

            if have_divAgrad
                @timeit to "divAgrad" begin
                    apply!((fourier=Hψ[:, iband], real=nothing),
                           H.divAgrad_op,
                           (fourier=ψ[:, iband], real=nothing),
                           ψ_real)  # ψ_real used as scratch
                end
            end

            if Threads.threadid() == 1
                merge!(DFTK.timer, to; tree_point=[t.name for t in DFTK.timer.timer_stack])
            end
       end
    end

    # Apply the nonlocal operator
    if !isnothing(H.nonlocal_op)
        @timing "nonlocal" begin
            apply!((fourier=Hψ, real=nothing), H.nonlocal_op, (fourier=ψ, real=nothing))
        end
    end

    Hψ
end

# Get energies and Hamiltonian
# kwargs is additional info that might be useful for the energy terms to precompute
# (eg the density ρ)
@timing function energy_hamiltonian(basis::PlaneWaveBasis, ψ, occupation; kwargs...)
    # it: index into terms, ik: index into kpoints
    @timing "ene_ops" ene_ops_arr = [ene_ops(term, basis, ψ, occupation; kwargs...)
                                     for term in basis.terms]
    energies  = [eh.E for eh in ene_ops_arr]
    operators = [eh.ops for eh in ene_ops_arr]         # operators[it][ik]

    # flatten the inner arrays in case a term returns more than one operator
    function flatten(arr)
        ret = []
        for a in arr
            if a isa RealFourierOperator
                push!(ret, a)
            else
                push!(ret, a...)
            end
        end
        ret
    end
    hks_per_k   = [flatten([blocks[ik] for blocks in operators])
                   for ik = 1:length(basis.kpoints)]      # hks_per_k[ik][it]

    H = Hamiltonian(basis, [HamiltonianBlock(basis, kpt, hks)
                            for (hks, kpt) in zip(hks_per_k, basis.kpoints)])
    E = Energies(basis.model.term_types, energies)
    (E=E, H=H)
end
function Hamiltonian(basis::PlaneWaveBasis; ψ=nothing, occupation=nothing, kwargs...)
    _, H = energy_hamiltonian(basis, ψ, occupation; kwargs...)
    H
end

"""
Get the total local potential of the given Hamiltonian, in real space
in the spin components.
"""
function total_local_potential(ham::Hamiltonian)
    n_spin = ham.basis.model.n_spin_components
    pots = map(1:n_spin) do σ
        # Get the potential from the first Hamiltonian block of this spin component
        # (works since all local potentials are the same)
        i_σ = first(krange_spin(ham.basis, σ))
        total_local_potential(ham.blocks[i_σ])
    end
    cat(pots..., dims=4)
end
total_local_potential(Hk::DftHamiltonianBlock) = Hk.local_op.potential
function total_local_potential(Hk::GenericHamiltonianBlock)
    only(o for o in Hk.optimized_operators if o isa RealSpaceMultiplication).potential
end

"""
Returns a new Hamiltonian with local potential replaced by the given one
"""
function hamiltonian_with_total_potential(ham::Hamiltonian, V)
    @assert size(V, 4) == ham.basis.model.n_spin_components
    newblocks = [hamiltonian_with_total_potential(Hk, V[:, :, :, Hk.kpoint.spin])
                 for Hk in ham.blocks]
    Hamiltonian(ham.basis, newblocks)
end
function hamiltonian_with_total_potential(Hk::HamiltonianBlock, V)
    operators = [op for op in Hk.operators if !(op isa RealSpaceMultiplication)]
    push!(operators, RealSpaceMultiplication(Hk.basis, Hk.kpoint, V))
    HamiltonianBlock(Hk.basis, Hk.kpoint, operators, Hk.scratch)
end


using Pkg
pkg"activate Burner"
# using Revise
using Burner
using Images
using VideoIO
using BSON

width = 512
height = 288

demodisp = Burner.Display(ones(3, height, width), 0.0015)
video = VideoIO.openvideo("oblivion.mp4")
imageshow = Float64.(Images.channelview(imresize(load("descargar-images-full-hd-1080p.jpg"), (height, width))))

Burner.burnVideo(demodisp, video, 1.0/30.0)
imagepresented = colorview(RGB, Burner.present(demodisp, imageshow))
imagepresented
bson("results.bson", Dict(:display => demodisp, :demoimage => imagepresented))


using POMDPModels
using POMDPs

let
    problem = BabyPOMDP(-5, -10, 0.1, 0.8, 0.1, 0.9)

    n = 100000
    results = Vector{Float64}(n)

    i = 1
    rng = MersenneTwister(i)
    init_state = false
    od = observation(problem, init_state, false, init_state)
    obs = rand(rng, od)
    sim = RolloutSimulator(rng=rng, initialstate=init_state, eps=0.0001)
    policy = FeedWhenCrying()
    results[i] = simulate(sim, problem, policy, updater(policy), PreviousObservation(obs))

    rngs = [MersenneTwister(i) for i in 1:n]

    @time for i in 1:n
        rng = rngs[i]
        init_state = false
        od = observation(problem, init_state, false, init_state)
        obs = rand(rng, od)
        sim = RolloutSimulator(rng=rng, initialstate=init_state, eps=0.0001)
        policy = FeedWhenCrying()
        results[i] = simulate(sim, problem, policy, updater(policy), PreviousObservation(obs))
    end
end


#=

Burning phase for tribe.

Ignacio Quintero Mächler

t(-_-t)

April 27 2017

=#

"""
    burn_compete(...)

Burning & adaptive phase for MCMC.
"""
function burn_tribe(total_llf     ::Function,
                    ω1upd_llr     ::Function,
                    ω0upd_llr     ::Function,
                    λ1upd_llr     ::Function,
                    λ0upd_llr     ::Function,
                    Xupd_llr      ::Function,
                    Rupd_llr      ::Function,
                    ωxupd_llr     ::Function,
                    σ²upd_llr     ::Function,
                    bgiid         ::Function,
                    bgiid_br      ::Function,
                    mhr_upd_Xbr   ::Function,
                    mhr_upd_Xtrio ::Function,
                    mhr_upd_Ybr   ::Function,
                    mhr_upd_Ytrio ::Function,
                    mhr_upd_Ystem ::Function,
                    mhr_upd_XYbr  ::Function,
                    mhr_upd_XYtrio::Function,
                    Xc      ::Array{Float64,2},
                    Yc      ::Array{Int64,3},
                    δXc     ::Array{Float64,3},
                    δYc     ::Array{Float64,3},
                    LApc    ::Array{Float64,2},
                    LAnc    ::Array{Float64,2},
                    LDc     ::Array{Float64,3},
                    σ²c     ::Float64,
                    ωxc     ::Float64,
                    ω1c     ::Float64,
                    ω0c     ::Float64,
                    λ1c     ::Float64,
                    λ0c     ::Float64,
                    Xnc1    ::Array{Int64,1},
                    Xnc2    ::Array{Int64,1},
                    brl     ::Array{Float64,1},
                    wcol    ::Array{Array{Int64,1},1},
                    bridx_a ::Array{Array{UnitRange{Int64},1},1},
                    brs     ::Array{Int64,3},
                    stemevc ::Array{Array{Float64,1},1},
                    trios   ::Array{Array{Int64,1},1},
                    wXp     ::Array{Int64,1},
                    λprior  ::Float64,
                    ωxprior ::NTuple{2,Float64},
                    ω1prior ::NTuple{2,Float64},
                    ω0prior ::NTuple{2,Float64},
                    σ²prior ::Float64,
                    np      ::Int64,
                    parvec  ::Array{Int64,1},
                    nburn   ::Int64,
                    obj_ar  ::Float64 = 0.234,
                    tune_int::Int64   = 100)

  m, ntip, narea = size(Yc)

  nedge = size(brs, 1)

  # likelihood and prior
  if ωxc >= 0.0
    llc = total_llf(Xc, Yc, LApc, LDc, ωxc, ω1c, ω0c, λ1c, λ0c,
                    stemevc, brs, σ²c)
  else
    llc = total_llf(Xc, Yc, LAnc, LDc, ωxc, ω1c, ω0c, λ1c, λ0c,
                    stemevc, brs, σ²c)
  end

  prc = allλpr(   λ1c, λ0c, λprior)           +
        logdexp(  σ²c, σ²prior)               +
        logdnorm(ωxc, ωxprior[1], ωxprior[2]) +
        logdnorm(ω1c, ω1prior[1], ω1prior[2]) +
        logdnorm(ω0c, ω0prior[1], ω0prior[2])

  # make scaling function
  scalef = makescalef(obj_ar)

  # rest of tuning parameters
  ptn = fill(0.1, np)

  # initialize acceptance log
  ltn = zeros(Int64, np)
  lup = zeros(Int64, np)
  lac = zeros(Int64, np)

  # row i proposals for X
  xpi  = fill(NaN, ntip)             # proposal x slice
  δxi  = fill(NaN, ntip, ntip)       # lineage pairwise differences
  lapi = fill(NaN, ntip)             # lineage average for ωx > 0
  lani = fill(NaN, ntip)             # lineage average for ωx < 0
  ldi  = fill(NaN, ntip, narea)      # lineage difference

  rj = wcol[1][2]

  # progress bar
  p = Progress(nburn, dt=5, desc="burn...", barlen=20, color=:green)

  # print number of parameters
  printstyled(
    "\nωx updates per iter = ", lastindex(filter(x -> x == 2,parvec)),
    "\nω1 updates per iter = ", lastindex(filter(x -> x == 3,parvec)),
    "\nω0 updates per iter = ", lastindex(filter(x -> x == 4,parvec)),
    "\nσ² updates per iter = ", lastindex(filter(x -> x == 1,parvec)),
    "\nλ1 updates per iter = ", lastindex(filter(x -> x == 5,parvec)),
    "\nλ0 updates per iter = ", lastindex(filter(x -> x == 6,parvec)),
    "\nθ  updates per iter = ", length(parvec), "\n", color = :green)

  Xcidx = CartesianIndices(Xc)

  #start burnin
  for it = Base.OneTo(nburn)

    # Update vector
    shuffle!(parvec)

    for up = parvec

      # update X
      if up > 6

        # X updates
        @inbounds begin

          upx = wXp[up - 6]

          # if root
          if upx == 1

            # allocate
            @simd for i = Base.OneTo(ntip)
              xpi[i] =  Xc[1,i]
            end

            # update xi
            addupt!(xpi, ptn, 1, up)

            xpi[rj] = xpi[1]::Float64

            llr = Rupd_llr(xpi, Xc, σ²c)::Float64

            if -randexp() < llr
              llc     += llr::Float64
              Xc[1,:]  = xpi::Array{Float64,1}
              lac[up] += 1   # log acceptance
            end

          else

            xi, xj = Xcidx[upx].I

            # allocate
            for j = Base.OneTo(ntip)
              xpi[j]  = Xc[xi,j]
              lapi[j] = LApc[xi,j]
              lani[j] = LAnc[xi,j]
              @simd for k = Base.OneTo(narea)
                ldi[j,k] = LDc[xi,j,k]
              end
              @simd for i = Base.OneTo(ntip)
                δxi[i,j] = δXc[i,j,xi]
              end
            end

            # update xi
            addupt!(xpi, ptn, xj, up)

            if in(upx, Xnc1)        # if an internal node
              xpi[Xcidx[Xnc2[findfirst(isequal(upx),Xnc1)]][2]] = xpi[xj]
            end

            # calculate new averages
            Xupd_linavg!(δxi, lapi, lani, ldi, wcol,
                         xpi, xi, xj, Yc, δYc, narea)

            if ωxc >= 0.0
              llr = Xupd_llr(xi, xpi, Xc, lapi, ldi, LApc, LDc, Yc,
                             ωxc, ω1c, ω0c, λ1c, λ0c, σ²c)::Float64
            else
              llr = Xupd_llr(xi, xpi, Xc, lani, ldi, LAnc, LDc, Yc,
                             ωxc, ω1c, ω0c, λ1c, λ0c, σ²c)::Float64
            end

            if -randexp() < llr
              llc        += llr::Float64
              Xc[xi,:]    = xpi::Array{Float64,1}
              δXc[:,:,xi] = δxi::Array{Float64,2}
              LApc[xi,:]  = lapi::Array{Float64,1}
              LAnc[xi,:]  = lani::Array{Float64,1}
              LDc[xi,:,:] = ldi::Array{Float64,2}
              lac[up]    += 1   # log acceptance
            end
          end
        end

      # update λ1
      elseif up == 5

        # update λ
        λ1p = mulupt(λ1c, ptn[5])::Float64

        llr = λ1upd_llr(Yc, λ1c, λ1p, λ0c, ω1c, LDc, stemevc, brs)

        prr = llrdexp_x(λ1p, λ1c, λprior)

        if -randexp() < (llr + prr + log(λ1p/λ1c))
          llc    += llr::Float64
          prc    += prr::Float64
          λ1c     = λ1p::Float64
          lac[5] += 1
        end

        # update internal node
        if rand() < 0.4
          λϕ1, λϕ0 = λϕprop()
          llc = mhr_upd_Ytrio(rand(trios), Xc, Yc,
                      λ1c, λ0c, ωxc, ω1c, ω0c, σ²c, λϕ1, λϕ0, llc, prc,
                      LApc, LAnc, LDc, δXc, δYc, brs, stemevc)
        end

      # update λ0
      elseif up == 6

        # update λ
        λ0p = mulupt(λ0c, ptn[6])::Float64

        llr = λ0upd_llr(Yc, λ1c, λ0c, λ0p, ω0c, LDc, stemevc, brs)::Float64

        prr = llrdexp_x(λ0p, λ0c, λprior)

        if -randexp() < (llr + prr + log(λ0p/λ0c))
          llc     += llr::Float64
          prc     += prr::Float64
          λ0c      = λ0p::Float64
          lac[6] += 1
        end

        # update internal node
        if rand() < 0.4
          λϕ1, λϕ0 = λϕprop()
          llc = mhr_upd_Ytrio(rand(trios), Xc, Yc,
                    λ1c, λ0c, ωxc, ω1c, ω0c, σ²c, λϕ1, λϕ0, llc, prc,
                    LApc, LAnc, LDc, δXc, δYc, brs, stemevc)
        end

      # if σ² is updated
      elseif up == 1

        σ²p = mulupt(σ²c, ptn[1])::Float64

        #likelihood ratio
        if ωxc >= 0.0
          llr = σ²upd_llr(Xc, LApc, ωxc, σ²c, σ²p)::Float64
        else
          llr = σ²upd_llr(Xc, LAnc, ωxc, σ²c, σ²p)::Float64
        end

        # prior ratio
        prr = llrdexp_x(σ²p, σ²c, σ²prior)

        if -randexp() < (llr + prr + log(σ²p/σ²c))
          llc += llr::Float64
          prc += prr::Float64
          σ²c  = σ²p::Float64
          lac[1] += 1
        end

      #update ωx
      elseif up == 2

        ωxp = addupt(ωxc, ptn[2])::Float64

        #likelihood ratio
        llr = ωxupd_llr(Xc, LApc, LAnc, ωxc, ωxp, σ²c)

        # prior ratio
        prr = llrdnorm_x(ωxp, ωxc, ωxprior[1], ωxprior[2])

        if -randexp() < (llr + prr)
          llc += llr::Float64
          prc += prr::Float64
          ωxc  = ωxp::Float64
          lac[2] += 1
        end

      #update ω1
      elseif up == 3

        ω1p = addupt(ω1c, ptn[3])::Float64

        # proposal likelihood and prior
        llr = ω1upd_llr(Yc, λ1c, ω1c, ω1p, LDc)

        # prior ratio
        prr = llrdnorm_x(ω1p, ω1c, ω1prior[1], ω1prior[2])

        if -randexp() < (llr + prr)
          llc += llr::Float64
          prc += prr::Float64
          ω1c  = ω1p::Float64
          lac[3] += 1
        end

        # update internal node
        if rand() < 0.4
          λϕ1, λϕ0 = λϕprop()
          llc = mhr_upd_Ytrio(rand(trios), Xc, Yc,
                    λ1c, λ0c, ωxc, ω1c, ω0c, σ²c, λϕ1, λϕ0, llc, prc,
                    LApc, LAnc, LDc, δXc, δYc, brs, stemevc)
        end

      # update ω0
      else

        ω0p = addupt(ω0c, ptn[4])

        # proposal likelihood and prior
        llr = ω0upd_llr(Yc, λ0c, ω0c, ω0p, LDc)::Float64

        # prior ratio
        prr = llrdnorm_x(ω0p, ω0c, ω0prior[1], ω0prior[2])

        if -randexp() < (llr + prr)
          llc += llr
          prc += prr
          ω0c  = ω0p
          lac[4] += 1
        end

        # update internal node
        if rand() < 0.4
          λϕ1, λϕ0 = λϕprop()
          llc = mhr_upd_Ytrio(rand(trios), Xc, Yc,
                    λ1c, λ0c, ωxc, ω1c, ω0c, σ²c, λϕ1, λϕ0, llc, prc,
                    LApc, LAnc, LDc, δXc, δYc, brs, stemevc)
        end

      end

      ## make DA updates with some Pr
      # make Y branch update
      if rand() < 2e-3
        λϕ1, λϕ0 = λϕprop()
        llc = mhr_upd_Ybr(rand(Base.OneTo(nedge)),
                        Xc, Yc, λ1c, λ0c, ωxc, ω1c, ω0c, σ²c, λϕ1, λϕ0,
                        llc, prc, LApc, LAnc, LDc, δXc, δYc, brs, stemevc)
      end

      # make X branch update
      if rand() < 2e-3
        σ²ϕ = σ²ϕprop()
        llc = mhr_upd_Xbr(rand(Base.OneTo(nedge-1)),
                          Xc, Yc, λ1c, λ0c, ωxc, ω1c, ω0c, σ²c, σ²ϕ, llc,
                          LApc, LAnc, LDc, δXc, δYc, brs, stemevc)
      end

      # make X trio update
      if rand() < 2e-3
        σ²ϕ = σ²ϕprop()
        llc = mhr_upd_Xtrio(rand(trios),
                            Xc, Yc, λ1c, λ0c, ωxc, ω1c, ω0c, σ²c, σ²ϕ, llc,
                            LApc, LAnc, LDc, δXc, δYc, brs, stemevc)
      end

      # make joint X & Y branch update
      if rand() < 2e-3
        λϕ1, λϕ0 = λϕprop()
        σ²ϕ = σ²ϕprop()
        llc = mhr_upd_XYbr(rand(Base.OneTo(nedge-1)),
                           Xc, Yc, λ1c, λ0c, ωxc, ω1c, ω0c, σ²c, σ²ϕ, λϕ1, λϕ0,
                           llc, prc, LApc, LAnc, LDc, δXc, δYc, brs, stemevc)
      end

      # make joint X & Y trio update
      if rand() < 2e-3
        λϕ1, λϕ0 = λϕprop()
        σ²ϕ = σ²ϕprop()
        llc      = mhr_upd_XYtrio(rand(trios),
                           Xc, Yc, λ1c, λ0c, ωxc, ω1c, ω0c, σ²c, σ²ϕ, λϕ1, λϕ0,
                           llc, prc, LApc, LAnc, LDc, δXc, δYc, brs, stemevc)
      end

      # update stem branch
      if rand() < 2e-3
        λϕ1, λϕ0 = λϕprop()
        llc = mhr_upd_Ystem(λ1c, λ0c, λϕ1, λϕ0, llc, stemevc, brs)
      end

      # log number of updates
      ltn[up] += 1
      lup[up] += 1

      if (in(tune_int,ltn))
        wts = findall(isequal(tune_int),ltn)      # which to scale
        for j = wts
          ar     = lac[j]/lup[j]
          ptn[j] = scalef(ptn[j],ar)
          ltn[j] = 0
        end
      end
    end

    next!(p)
  end

  return llc, prc, Xc, Yc, LApc, LAnc, LDc, δXc, δYc, stemevc,
         brs, σ²c, ωxc, ω1c, ω0c, λ1c, λ0c, ptn
end


 #returns the specification symbol in multicomponent
#returns :singlecomponent if there are no specifications
#throws error if more than one specification is found

const AMOUNT_CONST =Dict{Int,Any}(
     0000 => OneMol()
    ,0001 => MaterialAmount{MOLAR}()
    ,0002 => MaterialAmount{MASS}()
    ,0099 => SingleComponent()
    ,1100 => MaterialCompounds{MOLAR,FRACTION}()
    ,1200 => MaterialCompounds{MASS,FRACTION}()
    ,2100 => MaterialCompounds{MOLAR,TOTAL_AMOUNT}()
    ,2200 => MaterialCompounds{MASS,TOTAL_AMOUNT}()
    ,3100 => HumiditySpec{WetBulbTemperature}()
    ,3200 => HumiditySpec{HumidityRatio}()
    ,3300 => HumiditySpec{MolarHumidity}()
    ,3400 => HumiditySpec{MassHumidity}()
    ,3500 => HumiditySpec{RelativeHumidity}()
    ,3600 => HumiditySpec{HumidityDewPoint}()
)

#the following functions are essencial in the ordering
#do not change the relative order!!

#equilibria properties are always first
_stateorder(x::PhaseFractions) = 1
_stateorder(x::TwoPhaseEquilibrium) = 2
_stateorder(x::VaporFraction) = 3

#pressure is always first in the scalar properties
_stateorder(x::Pressure) = 4
_stateorder(x::Helmholtz) = 5
_stateorder(x::Gibbs) = 6
_stateorder(x::InternalEnergy) = 7
_stateorder(x::VolumeAmount) = 8
_stateorder(x::Enthalpy) = 9
_stateorder(x::Entropy) = 10
#temperature always for last of the scalar properties
_stateorder(x::Temperature) = 11

#From here, it doesnt matter the order, as those are treated specially

_stateorder(x::HumiditySpec) = 1000

_stateorder(x::MaterialCompounds) = 10002

_stateorder(x::MaterialAmount ) = 10001
_stateorder(x::PhaseTag) = 10000
_stateorder(x::Options) = 10004
_stateorder(x::MolecularWeight) = 10005
_stateorder(x::VariableSpec) = 10006

#using specs here allows to delete all reduce_check_mass(::Type)
function _static_specs(x::Type{ThermodynamicState{S,C}}) where {S,C}
    @nospecialize S,C
    tvec = S.parameters
    res1 = [tvec[i].parameters[1]() for i in 1:length(tvec)]
    if C === Tuple{}
        return res1
    else
        tvec2 = C.parameters
        res2 = [tvec2.parameters[i]() for i in 1:length(tvec)]
        append!(res1,res2)
        return res1
    end
end

@generated function static_amount_type(x::T) where T<: ThermodynamicState
    sps = _static_specs(x)
    val = sum(_reduce_check_mass(i) for i in sps)
    modval = mod(val,MATERIAL_SINGLE_MAX)
    if val in (2100,2200)
        _1 = val
        _2 = val
    else
        _1 = max(val-modval,99)
        _2 = modval
    end
    res =  (AMOUNT_CONST[_1],AMOUNT_CONST[_2])

    return :($res)
end

function _state_from_type(x::Type{T}) where T<: ThermodynamicState
    @nospecialize T
    state_specs =  _static_specs(T)
    has_compound = false
    for spec in state_specs
        if spec isa MaterialCompounds
            has_compound = true
            comp_spec = spec
            break
        end
    end

    #double iteration, used in moist air to create an invalid
    #state where compound and humidity spec are present
    #in this case, the resulting type is MultiPT instead of HumPT
    for spec in state_specs
        if (spec isa HumiditySpec) & !has_compound
            has_compound = true
            comp_spec = spec
            break
        end
    end

    if !has_compound
        comp_spec = SingleComponent()
    end
    sorted_specs = sort(state_specs,by= _stateorder)
    res = (sorted_specs[1],sorted_specs[2],comp_spec)
    return res
end

__unwrap(x::Type{T}) where T = x.parameters[1]
__unwrap(x) = x
function _dynamic_get_spec_idx(@nospecialize(x::Type{ThermodynamicState{S,C}}),@nospecialize(sp::Type{S2}))::Int where {S,C,S2}
    state_specs =  _static_specs(x)
    S3 = __unwrap(sp)
    for (i,spec) in pairs(state_specs)
        if typeof(spec) <: S3
            return i
        end
    end
    return 0
end

@generated function _static_get_spec_idx(@nospecialize(x::ThermodynamicState{S,C}),@nospecialize(sp::Type{S2}))::Int where {S,C,S2}
    res = _dynamic_get_spec_idx(x,sp)
    return :($res)
end
"""
    state_type(x::ThermodynamicState)

Returns a sorted tuple of thermodynamic specifications of the input state,
 the tuple has 3 elements, corresponding to two general specifications and a last amount specification,
  indicating if the state is single or multicomponent.

"""
@generated function state_type(x::ThermodynamicState)
    res = _state_from_type(x)
    return :($res)
end

"""
    get_spec(val::AbstractSpec,st::ThermodynamicState)
    get_spec(val::Type{AbstractSpec},st::ThermodynamicState)

Returns the first specification of the type `val` present in the state `st`. returns `nothing` if not found.

## Example

```julia-repl
julia> a1 = state(n = [1.1,2.2,3.3],T=273.15,mol_h = 101.342)
ThermodynamicState with 3 properties:
  Molar amounts : [1.1, 2.2, 3.3][mol]
  Temperature : 273.15[K]
  Molar enthalpy : 101.342[J mol^-1]

julia> get_spec(Enthalpy,a1)
  spec(mol_h = 101.342[J mol^-1])

julia> get_spec(VolumeAmount,a1) == nothing
  true
```

"""
function get_spec end
get_spec(val::T,st::ThermodynamicState) where T<:AbstractSpec = get_spec(T,st)

function get_spec(val::Type{T},st::ThermodynamicState) where T<: AbstractSpec
    i = _static_get_spec_idx(st,T)
    i === 0 && return nothing
    return st.specs[i]
    #return nothing
end

function throw_get_spec(val,st::ThermodynamicState)::Spec
    res = get_spec(val,st)
    res === nothing && throw(error("$st does not have $val"))
    return res
end

"""
    has_spec(val::AbstractSpec,st::ThermodynamicState)::Bool
    has_spec(val::Type{AbstractSpec},st::ThermodynamicState)::Bool

checks if the specification of the type `val` present in the state `st`.

## Example

```julia-repl
julia> a1 = state(n = [1.1,2.2,3.3],T=273.15,mol_h = 101.342)
ThermodynamicState with 3 properties:
  Molar amounts : [1.1, 2.2, 3.3][mol]
  Temperature : 273.15[K]
  Molar enthalpy : 101.342[J mol^-1]

julia> has_spec(Enthalpy,a1)
  true

julia> has_spec(VolumeAmount,a1) == nothing
  false
```

"""
function has_spec end
function has_spec(val::Type{T},st::ThermodynamicState) where T<: AbstractSpec
    static_has_spec(st,val)
end

"""
    amount_type(st::ThermodynamicState)

given a state, returns a tuple with information about how is the amount of matter specified.

The first element of the tuple is related to the number of components present in the state. if it has only one component. it will return `SingleComponent`. if a multicomponent specification (`MaterialCompounds`)
is defined, it will return that type instead.

The second element is related to the type of material specified (molar or mass). if there isn't any specification, it will assume one mol, via the type `OneMol`. It will return a `MaterialAmount`
if it is specified.

## Examples

```julia-repl
julia> tp = state(t=1,p=2)
ThermodynamicState with 2 properties:
  Temperature : 1[K]
  Pressure : 2[Pa]

julia> amount_type(tp)
(SingleComponent(), OneMol())

julia> tpn = state(t=1,p=2,mass=3)
ThermodynamicState with 3 properties:
  Temperature : 1[K]
  Pressure : 2[Pa]
  Mass : 3[kg]

julia> amount_type(tpn)
(SingleComponent(), MaterialAmount{MASS}())

julia> tpxm = state(t=1,p=2,mass=3,xn = [0.1,0.9])
ThermodynamicState with 4 properties:
  Temperature : 1[K]
  Pressure : 2[Pa]
  Mass : 3[kg]
  Molar fraction : [0.1, 0.9]

julia> amount_type(tpxm)
(MaterialCompounds{MOLAR, FRACTION}(), MaterialAmount{MASS}())

julia> tpm = state(t=1,p=2,m = [1,2])
ThermodynamicState with 3 properties:
  Temperature : 1[K]
  Pressure : 2[Pa]
  Mass amounts : [1, 2][kg]

julia> amount_type(tpm)
(MaterialCompounds{MASS, TOTAL_AMOUNT}(), MaterialCompounds{MASS, TOTAL_AMOUNT}())
```
"""
function amount_type(st::T) where T<:ThermodynamicState
    return static_amount_type(st)
end

@generated function static_has_spec(st::ThermodynamicState,val::Type{T}) where T<:AbstractSpec
    idx =_static_get_spec_idx(st,T)
    res = !iszero(idx)
    return :($res)
end

reduce_numtype(::T) where T <: Spec = reduce_numtype(T)
reduce_numtype(::Type{Spec{V,T}}) where {V,T<:Unitful.Quantity{W}} where W = W
reduce_numtype(::Type{Spec{V,T}}) where {V,T<:AbstractVector{<:Unitful.Quantity{W}}} where W = W
reduce_numtype(::Type{Spec{V,T}}) where {V,T<:AbstractVector{<:W}} where W = W
reduce_numtype(::Type{Spec{V,T}}) where {V,T} = Bool
reduce_numtype(::Type{Spec{V,T}}) where {V,T<:Number} = T
reduce_numtype(::Type{T}) where T<:Number = T
reduce_numtype(A,B) = promote_type(reduce_numtype(A),reduce_numtype(B))
function numtype(st::T) where T <: ThermodynamicState
    res = foldl(reduce_numtype,st.specs)
    return res
end
_one(st::T) where T <: ThermodynamicState = one(numtype(st))


#=
@testset "test NGramHMM..." begin
        @time d = ZhuXian()
        @time m = KongYiji.NGramHMM(2, d, KongYiji.dir("usrdict1"), 0.)
        @time m2 = KongYiji.HMM(d, KongYiji.dir("usrdict1"), 0.)
        @test m.hpr == m2.hpr
        @test m.h2h == m2.h2h
        @test m.h2v == m2.h2v
        @test m.INF == m2.INF

        xs = "田不易独自一人站在树林里的僻静处，负手而立。"
        r = m(xs)
        r2 = m2(xs)
        println(r)
        println(r2)
end
  =#

@testset "Cross validating NGramHMM..." begin
        @time ctb = ChTreebank()

        nb = 10
        batches = KongYiji.foldbatch(ctb, nb)
        tbs2 = KongYiji.HmmScoreTable[]
        tbs3 = KongYiji.HmmScoreTable[]
        @showprogress 1 "Cross Validating ooHMM..." for (tr, te) in batches
                m3 = KongYiji.NGramHMM(3, tr, KongYiji.dir("usrdict1"), 0.)
                m2 = KongYiji.HMM(tr, KongYiji.dir("usrdict1"), 0.)
                x = KongYiji.rawsents(te)
                z = KongYiji.wordsents(te)
                @time y2 = m2(x)
                @time y3 = m3(x)
                t2 = KongYiji.HmmScoreTable(z, y2)
                t3 = KongYiji.HmmScoreTable(z, y3)
                @show t2
                @show t3
                push!(tbs2, t2)
                push!(tbs3, t3)
                #println(tbs[end])
        end
        println(sum(tbs2))
        println(sum(tbs3))
end

#=
@testset "Compare NGramHMM with HMM..." begin
        @time d = ZhuXian()
        @time m = KongYiji.NGramHMM(2, d, KongYiji.dir("usrdict1"), 0.)
        @time m2 = KongYiji.HMM(d, KongYiji.dir("usrdict1"), 0.)
        for sent in rawsents(d)
                y = m(sent)
                y2 = m2(sent)
                @test y == y2
                if y != y2; println(sent); end
        end
end

  =#


struct Summary <: Output
  filename::String
  last_infections::Vector{Float32}
  num_infections::Vector{UInt32}
  Summary(filename::AbstractString, num_trajectories::Integer) = new(filename, fill(NaN32, num_trajectories), Vector{UInt32}(undef, num_trajectories))
end

function pushtrajectory!(s::Summary, trajectory_id::Integer, writelock::Base.AbstractLock, state::MocosSim.SimState, ::MocosSim.SimParams, ::DetectionCallback)
  try
    lock(writelock)
    s.last_infections[trajectory_id] = maximum(MocosSim.time, state.forest.inedges)
    s.num_infections[trajectory_id] = count(MocosSim.istransmission, state.forest.inedges)
  finally
    unlock(writelock)
  end
  nothing
end

function aftertrajectories(s::Summary, ::MocosSim.SimParams)
  file = jldopen(s.filename, "w")
  try
    file["num_infections"] = s.num_infections
    file["last_infections"] = s.last_infections
  finally
    close(file)
  end
  nothing
end


using StreamReader
using FactCheck

facts("All") do
    @fact DEFAULT_PART_SIZE --> 1024
    @fact calculate(10) --> (1, 10, 10)
    @fact calculate(10, 3) --> (4, 3, 1)
    @fact calculate(1024) --> (1, 1024, 1024)
    @fact calculate(2048) --> (2, 1024, 1024)
    @fact calculate(2024) --> (2, 1024, 1000)

    pit = PartsIterator(10)
    @fact pit.size --> 10
    @fact pit.part_size --> 10
    @fact pit.left --> 0
    @fact pit.length --> 1
    @fact pit.part --> 0
    @fact pit.current_size --> 0
    @fact pit.loaded --> 0

    pit = PartsIterator(10, 3)

    @fact pit.size --> 10
    @fact pit.part_size --> 3
    @fact pit.left --> 0
    @fact pit.length --> 4
    @fact pit.part --> 0
    @fact pit.current_size --> 0
    @fact pit.loaded --> 0

    s = start(pit)
    @fact s --> 3
    @fact pit.part_size --> 3
    @fact pit.left --> 10
    @fact pit.length --> 4
    @fact pit.part --> 0
    @fact pit.current_size --> 0
    @fact pit.loaded --> 0

    @fact done(pit, s) --> false
    @fact pit.part_size --> 3
    @fact pit.left --> 10
    @fact pit.length --> 4
    @fact pit.part --> 0
    @fact pit.current_size --> 0
    @fact pit.loaded --> 0

    d, s = next(pit, s)
    @fact d --> 3
    @fact s --> 3
    @fact pit.part_size --> 3
    @fact pit.left --> 7
    @fact pit.length --> 4
    @fact pit.part --> 1
    @fact pit.current_size --> 3
    @fact pit.loaded --> 3

    @fact done(pit, s) --> false

    d, s = next(pit, s)
    @fact d --> 3
    @fact s --> 3
    @fact pit.part_size --> 3
    @fact pit.left --> 4
    @fact pit.length --> 4
    @fact pit.part --> 2
    @fact pit.current_size --> 3
    @fact pit.loaded --> 6

    @fact done(pit, s) --> false

    d, s = next(pit, s)
    @fact d --> 3
    @fact s --> 1
    @fact pit.part_size --> 3
    @fact pit.left --> 1
    @fact pit.length --> 4
    @fact pit.part --> 3
    @fact pit.current_size --> 3
    @fact pit.loaded --> 9

    @fact done(pit, s) --> false

    d, s = next(pit, s)
    @fact d --> 1
    @fact s --> 0
    @fact pit.part_size --> 3
    @fact pit.left --> 0
    @fact pit.length --> 4
    @fact pit.part --> 4
    @fact pit.current_size --> 1
    @fact pit.loaded --> 10

    @fact done(pit, s) --> true
    @fact [(1, 3, 3, 7, 30.0), (2, 3, 6, 4, 60.0), (3, 3, 9, 1, 90.0), (4, 1, 10, 0, 100.0)] --> [(pit.part, d, pit.loaded, pit.left, ceil(loadedpct(pit), 2)) for d in pit]

    #### READER ITERATOR TEST #####

    io = IOBuffer()
    s = write(io, "God is Love!!")
    seekstart(io)
    ir = IOStringReaderIterator(io, s, 4)
    edata = String["God ", "is L", "ove!", "!"]
    data = collect(ir)
    @fact edata --> data

    #### READER LIST ITERATOR TEST #####

    data = "God is Love!!"
    io = IOBuffer()
    s = write(io, data)
    seekstart(io)
    num_parts, parts_size, last_part_size = calculate(s, 4)
    parts_size_arr = Integer[parts_size for i in 1:num_parts-1]
    push!(parts_size_arr, last_part_size)

    edata = {{"Go", "d "}, {"is", " L"}, {"ov", "e!"}, {"!"}}

    rli = ReaderListIterator((mri) -> IOStringReaderIterator(io, parts_size_arr[mri.num], 2), num_parts)
    for (num, sri) in rli
        for d in sri
            @fact d --> edata[num][sri.num] "Mismatch edata[$num][$(sri.num)]"
        end
    end
end


const checker_width_dialog = Dialog()
const checker_per_width_dialog = Dialog()
const checker_per_height_dialog = Dialog()

function newboard(designations)
    @label desig
    dialog = Dialog()
    examples = isempty(designations) ? "" : "(already existing designations: $designations) "
    designation = request("Give a designation for this new board:$examples", dialog)
    if designation ∈ designations
        @warn "designation $designation is already taken, try again…"
        @goto desig
    end
    isempty(designation) && @goto desig
    @label width
    _checker_width_cm = requestᵐ("What is the width (same as height) of the checkers in cm?", checker_width_dialog)
    if !all(isnumeric, filter(!isequal('.'), _checker_width_cm))
        @warn "width $_checker_width_cm is not a number, try again…"
        @goto width
    end
    checker_width_cm = parse(Float64, _checker_width_cm)
    if checker_width_cm < 0
        @warn "width $_checker_width_cm has to be larger than zero, try again…"
        @goto width
    end
    @label perwidth
    _checker_per_width = requestᵐ("How many checkers are there across the width of the board?", checker_per_width_dialog)
    if !all(isnumeric, _checker_per_width)
        @warn "$_checker_per_width is not an integer, try again…"
        @goto perwidth
    end
    checker_per_width = parse(Int, _checker_per_width)
    if checker_per_width < 1
        @warn "there has to be more than one checkers, try again…"
        @goto perwidth
    end
    @label perheight
    _checker_per_height = requestᵐ("How many checkers are there across the height of the board?", checker_per_height_dialog)
    if !all(isnumeric, _checker_per_height)
        @warn "$_checker_per_height is not an integer, try again…"
        @goto perheight
    end
    checker_per_height = parse(Int, _checker_per_height)
    if checker_per_height < 1
        @warn "there has to be more than one checkers, try again…"
        @goto perheight
    end
    @label describe
    dialog = Dialog()
    board_description = request("""Describe the board to facilitate recognizing it in the future (e.g. "a small board, black tape framming it on the long side, cardboard on the short"):""", dialog)
    if isempty(board_description)
        @warn "you must give some minimal description. Think about the future generations!"
        @goto describe
    end
    try
        Board(designation, checker_width_cm, (checker_per_width, checker_per_height), board_description)
    catch ex
        @warn "something was wrong" ex
        @goto desig
    end
end

const calibration_type_menu = RadioMenu(["Stationary", "Moving"])
const calibration_comment_dialog = Dialog()

function newcalibration(boards)
    board = if isempty(boards)
        @info "registering a new board"
        newboard(String[])
    else
        designations = getfield.(boards, :designation)
        options = designations
        pushfirst!(options, "Register a new board")
        board_menu = RadioMenu(options, min(2, length(options)))
        i = requestᵐ("Which board was used?", board_menu)
        i == 1 ? newboard(designations) : boards[i - 1]
    end

    @label newcalibl

    i = requestᵐ("Which type of calibration is it?", calibration_type_menu)
    intrinsic = if i == 1
        missing
    else
        println("Registrating the intrinsic calibration POI (waving the checkerboard around)")
        newinterval(true)
    end
    println("Registrating the extrinsic calibration POI (the checkerboard on the ground)")
    extrinsic = newinterval(false)

    comment = requestᵐ("Any comments about this calibration?", calibration_comment_dialog)
    try
        Calibration(intrinsic, extrinsic, board, comment)
    catch ex
        @warn "something was wrong" ex
        @goto newcalibl
    end
end


export
        LaneChangeAction,
        LaneChangeChoice,
        LaneChangeModel,
        get_lane_offset,

        DIR_RIGHT,
        DIR_MIDDLE,
        DIR_LEFT,

        TimLaneChanger

abstract type LaneChangeAction <: DriveAction end

"""
    LaneChangeChoice
A choice of whether to change lanes, and what direction to do it in
"""
const DIR_RIGHT = -1
const DIR_MIDDLE =  0
const DIR_LEFT =  1
immutable LaneChangeChoice <: LaneChangeAction
    dir::Int # -1, 0, 1
end
Base.show(io::IO, a::LaneChangeChoice) = @printf(io, "LaneChangeChoice(%d)", dir)
Base.length(::Type{LaneChangeChoice}) = 1
Base.convert(::Type{LaneChangeChoice}, v::Vector{Float64}) = LaneChangeChoice(convert(Int, v[1]))
function Base.copy!(v::Vector{Float64}, a::LaneChangeChoice)
    v[1] = a.dir
    v
end
function get_lane_offset(a::LaneChangeChoice, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0)
    if a.dir == DIR_MIDDLE
        rec[vehicle_index, pastframe].state.posF.t
    elseif a.dir == DIR_LEFT
        convert(Float64, get(LANEOFFSETLEFT, rec, roadway, vehicle_index, pastframe))
    else
        @assert(a.dir == DIR_RIGHT)
        convert(Float64, get(LANEOFFSETRIGHT, rec, roadway, vehicle_index, pastframe))
    end
end

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

abstract type LaneChangeModel{LaneChangeAction} end
get_name(::LaneChangeModel) = "???"
set_desired_speed!(::LaneChangeModel, v_des::Float64) = model # # do nothing by default
reset_hidden_state!(model::LaneChangeModel) = model # do nothing by default
observe!(model::LaneChangeModel, scene::Scene, roadway::Roadway, egoid::Int) = model  # do nothing by default
Base.rand(model::LaneChangeModel) = error("rand not implemented for model $model")

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

type TimLaneChanger <: LaneChangeModel{LaneChangeAction}
    dir::Int
    rec::SceneRecord

    v_des::Float64
    threshold_fore::Float64
    threshold_lane_change_gap_fore::Float64
    threshold_lane_change_gap_rear::Float64

    function TimLaneChanger(
        action_context::IntegratedContinuous;
        v_des::Float64=29.0,
        rec::SceneRecord=SceneRecord(2,action_context.Δt),
        threshold_fore::Float64 = 50.0,
        threshold_lane_change_gap_fore::Float64 = 10.0,
        threshold_lane_change_gap_rear::Float64 = 10.0,
        dir::Int=DIR_MIDDLE,
        )

        retval = new()
        retval.dir = dir
        retval.rec = rec
        retval.v_des = v_des
        retval.threshold_fore = threshold_fore
        retval.threshold_lane_change_gap_fore = threshold_lane_change_gap_fore
        retval.threshold_lane_change_gap_rear = threshold_lane_change_gap_rear
        retval
    end
end
get_name(::TimLaneChanger) = "TimLaneChanger"
function set_desired_speed!(model::TimLaneChanger, v_des::Float64)
    model.v_des = v_des
    model
end
function observe!(model::TimLaneChanger, scene::Scene, roadway::Roadway, egoid::Int)
    #debug = open("debug.log", "a"); println(debug, "observe!(model::TimLaneChanger, scene::Scene, roadway::Roadway, egoid::Int)");  close(debug)

    rec = model.rec
    update!(rec, scene)
    vehicle_index = get_index_of_first_vehicle_with_id(rec, egoid)

    veh_ego = scene[vehicle_index]
    v = veh_ego.state.v

    left_lane_exists = convert(Float64, get(N_LANE_LEFT, rec, roadway, vehicle_index)) > 0
    right_lane_exists = convert(Float64, get(N_LANE_RIGHT, rec, roadway, vehicle_index)) > 0
    fore_M = get_neighbor_fore_along_lane(scene, vehicle_index, roadway, VehicleTargetPointFront(), VehicleTargetPointRear(), VehicleTargetPointFront(), max_distance_fore=model.threshold_fore)
    fore_L = get_neighbor_fore_along_left_lane(scene, vehicle_index, roadway, VehicleTargetPointRear(), VehicleTargetPointRear(), VehicleTargetPointFront())
    fore_R = get_neighbor_fore_along_right_lane(scene, vehicle_index, roadway, VehicleTargetPointRear(), VehicleTargetPointRear(), VehicleTargetPointFront())
    rear_L = get_neighbor_rear_along_left_lane(scene, vehicle_index, roadway, VehicleTargetPointFront(), VehicleTargetPointFront(), VehicleTargetPointRear())
    rear_R = get_neighbor_rear_along_right_lane(scene, vehicle_index, roadway, VehicleTargetPointFront(), VehicleTargetPointFront(), VehicleTargetPointRear())

    model.dir = DIR_MIDDLE
    if fore_M.Δs < model.threshold_fore # there is a lead vehicle
        veh_M = scene[fore_M.ind]
        speed_M = veh_M.state.v
        if speed_M ≤ min(model.v_des, v) # they are driving slower than we want

            speed_ahead = speed_M

            # consider changing to a different lane
            if right_lane_exists &&
               fore_R.Δs > model.threshold_lane_change_gap_rear && # there is space rear
               rear_R.Δs > model.threshold_lane_change_gap_fore && # there is space fore
               (rear_R.ind == 0 || scene[rear_R.ind].state.v ≤ v) && # we are faster than any follower
               (fore_R.ind == 0 || scene[fore_R.ind].state.v > speed_ahead) # we are faster than any leader

                speed_ahead = fore_R.ind != 0 ? scene[fore_R.ind].state.v : Inf
                model.dir = DIR_RIGHT
            end
            if left_lane_exists &&
               fore_L.Δs > model.threshold_lane_change_gap_rear && # there is space rear
               rear_L.Δs > model.threshold_lane_change_gap_fore && # there is space fore
               (rear_L.ind == 0 || scene[rear_L.ind].state.v ≤ v) && # we are faster than any follower
               (fore_L.ind == 0 || scene[fore_L.ind].state.v > speed_ahead) # we are faster than any leader

                speed_ahead = fore_L.ind != 0 ? scene[fore_L.ind].state.v : Inf
                model.dir = DIR_LEFT
            end
        end
    end

    model
end
Base.rand(model::TimLaneChanger) = LaneChangeChoice(model.dir)


#=
This file is auto-generated. Do not edit.
=#

#! format: off

"""
    mutable struct EX4VSA <: AVR
        Iflim::Float64
        d::Float64
        f::Float64
        Spar::Float64
        K1::Float64
        K2::Float64
        Oel_lim::NamedTuple{(:min, :max), Tuple{Float64, Float64}}
        G::Float64
        Ta::Float64
        Tb::Float64
        Te::Float64
        E_lim::NamedTuple{(:min, :max), Tuple{Float64, Float64}}
        V_ref::Float64
        ext::Dict{String, Any}
        states::Vector{Symbol}
        n_states::Int
        internal::InfrastructureSystemsInternal
    end

IEEE Excitation System for Voltage Security Assesment

# Arguments
- `Iflim::Float64`: OEL Field current limit, validation range: `(0, nothing)`, action if invalid: `warn`
- `d::Float64`: OEL parameter d, validation range: `(0, nothing)`, action if invalid: `warn`
- `f::Float64`: OEL parameter f, validation range: `(0, nothing)`, action if invalid: `warn`
- `Spar::Float64`: OEL parameter Spar, validation range: `(0, nothing)`, action if invalid: `warn`
- `K1::Float64`: OEL delay time constant, validation range: `(0, nothing)`, action if invalid: `warn`
- `K2::Float64`: OEL parameter K2, validation range: `(0, nothing)`, action if invalid: `warn`
- `Oel_lim::NamedTuple{(:min, :max), Tuple{Float64, Float64}}`: Oel integrator limits (Oel_min, Oel_max)
- `G::Float64`: AVR Exciter Gain, validation range: `(0, nothing)`, action if invalid: `warn`
- `Ta::Float64`: Numerator lead-lag (lag) time constant in s, validation range: `(0, nothing)`, action if invalid: `warn`
- `Tb::Float64`: Denominator lead-lag (lag) time constant in s, validation range: `(0, nothing)`, action if invalid: `warn`
- `Te::Float64`: Exciter Time Constant in s, validation range: `(0, nothing)`, action if invalid: `warn`
- `E_lim::NamedTuple{(:min, :max), Tuple{Float64, Float64}}`: Voltage regulator limits (regulator output) (E_min, E_max)
- `V_ref::Float64`: Reference Voltage Set-point, validation range: `(0, nothing)`
- `ext::Dict{String, Any}`
- `states::Vector{Symbol}`: The states are:
	Vll: Lead-lag internal state,
	Vex: Exciter Output,
	oel: OEL integrator state
- `n_states::Int`: The EX4VSA has 3 states
- `internal::InfrastructureSystemsInternal`: power system internal reference, do not modify
"""
mutable struct EX4VSA <: AVR
    "OEL Field current limit"
    Iflim::Float64
    "OEL parameter d"
    d::Float64
    "OEL parameter f"
    f::Float64
    "OEL parameter Spar"
    Spar::Float64
    "OEL delay time constant"
    K1::Float64
    "OEL parameter K2"
    K2::Float64
    "Oel integrator limits (Oel_min, Oel_max)"
    Oel_lim::NamedTuple{(:min, :max), Tuple{Float64, Float64}}
    "AVR Exciter Gain"
    G::Float64
    "Numerator lead-lag (lag) time constant in s"
    Ta::Float64
    "Denominator lead-lag (lag) time constant in s"
    Tb::Float64
    "Exciter Time Constant in s"
    Te::Float64
    "Voltage regulator limits (regulator output) (E_min, E_max)"
    E_lim::NamedTuple{(:min, :max), Tuple{Float64, Float64}}
    "Reference Voltage Set-point"
    V_ref::Float64
    ext::Dict{String, Any}
    "The states are:
	Vll: Lead-lag internal state,
	Vex: Exciter Output,
	oel: OEL integrator state"
    states::Vector{Symbol}
    "The EX4VSA has 3 states"
    n_states::Int
    "power system internal reference, do not modify"
    internal::InfrastructureSystemsInternal
end

function EX4VSA(Iflim, d, f, Spar, K1, K2, Oel_lim, G, Ta, Tb, Te, E_lim, V_ref=1.0, ext=Dict{String, Any}(), )
    EX4VSA(Iflim, d, f, Spar, K1, K2, Oel_lim, G, Ta, Tb, Te, E_lim, V_ref, ext, [:Vll, :Vex, :oel], 3, InfrastructureSystemsInternal(), )
end

function EX4VSA(; Iflim, d, f, Spar, K1, K2, Oel_lim, G, Ta, Tb, Te, E_lim, V_ref=1.0, ext=Dict{String, Any}(), states=[:Vll, :Vex, :oel], n_states=3, internal=InfrastructureSystemsInternal(), )
    EX4VSA(Iflim, d, f, Spar, K1, K2, Oel_lim, G, Ta, Tb, Te, E_lim, V_ref, ext, states, n_states, internal, )
end

# Constructor for demo purposes; non-functional.
function EX4VSA(::Nothing)
    EX4VSA(;
        Iflim=0,
        d=0,
        f=0,
        Spar=0,
        K1=0,
        K2=0,
        Oel_lim=(min=0.0, max=0.0),
        G=0,
        Ta=0,
        Tb=0,
        Te=0,
        E_lim=(min=0.0, max=0.0),
        V_ref=0,
        ext=Dict{String, Any}(),
    )
end

"""Get [`EX4VSA`](@ref) `Iflim`."""
get_Iflim(value::EX4VSA) = value.Iflim
"""Get [`EX4VSA`](@ref) `d`."""
get_d(value::EX4VSA) = value.d
"""Get [`EX4VSA`](@ref) `f`."""
get_f(value::EX4VSA) = value.f
"""Get [`EX4VSA`](@ref) `Spar`."""
get_Spar(value::EX4VSA) = value.Spar
"""Get [`EX4VSA`](@ref) `K1`."""
get_K1(value::EX4VSA) = value.K1
"""Get [`EX4VSA`](@ref) `K2`."""
get_K2(value::EX4VSA) = value.K2
"""Get [`EX4VSA`](@ref) `Oel_lim`."""
get_Oel_lim(value::EX4VSA) = value.Oel_lim
"""Get [`EX4VSA`](@ref) `G`."""
get_G(value::EX4VSA) = value.G
"""Get [`EX4VSA`](@ref) `Ta`."""
get_Ta(value::EX4VSA) = value.Ta
"""Get [`EX4VSA`](@ref) `Tb`."""
get_Tb(value::EX4VSA) = value.Tb
"""Get [`EX4VSA`](@ref) `Te`."""
get_Te(value::EX4VSA) = value.Te
"""Get [`EX4VSA`](@ref) `E_lim`."""
get_E_lim(value::EX4VSA) = value.E_lim
"""Get [`EX4VSA`](@ref) `V_ref`."""
get_V_ref(value::EX4VSA) = value.V_ref
"""Get [`EX4VSA`](@ref) `ext`."""
get_ext(value::EX4VSA) = value.ext
"""Get [`EX4VSA`](@ref) `states`."""
get_states(value::EX4VSA) = value.states
"""Get [`EX4VSA`](@ref) `n_states`."""
get_n_states(value::EX4VSA) = value.n_states
"""Get [`EX4VSA`](@ref) `internal`."""
get_internal(value::EX4VSA) = value.internal

"""Set [`EX4VSA`](@ref) `Iflim`."""
set_Iflim!(value::EX4VSA, val) = value.Iflim = val
"""Set [`EX4VSA`](@ref) `d`."""
set_d!(value::EX4VSA, val) = value.d = val
"""Set [`EX4VSA`](@ref) `f`."""
set_f!(value::EX4VSA, val) = value.f = val
"""Set [`EX4VSA`](@ref) `Spar`."""
set_Spar!(value::EX4VSA, val) = value.Spar = val
"""Set [`EX4VSA`](@ref) `K1`."""
set_K1!(value::EX4VSA, val) = value.K1 = val
"""Set [`EX4VSA`](@ref) `K2`."""
set_K2!(value::EX4VSA, val) = value.K2 = val
"""Set [`EX4VSA`](@ref) `Oel_lim`."""
set_Oel_lim!(value::EX4VSA, val) = value.Oel_lim = val
"""Set [`EX4VSA`](@ref) `G`."""
set_G!(value::EX4VSA, val) = value.G = val
"""Set [`EX4VSA`](@ref) `Ta`."""
set_Ta!(value::EX4VSA, val) = value.Ta = val
"""Set [`EX4VSA`](@ref) `Tb`."""
set_Tb!(value::EX4VSA, val) = value.Tb = val
"""Set [`EX4VSA`](@ref) `Te`."""
set_Te!(value::EX4VSA, val) = value.Te = val
"""Set [`EX4VSA`](@ref) `E_lim`."""
set_E_lim!(value::EX4VSA, val) = value.E_lim = val
"""Set [`EX4VSA`](@ref) `V_ref`."""
set_V_ref!(value::EX4VSA, val) = value.V_ref = val
"""Set [`EX4VSA`](@ref) `ext`."""
set_ext!(value::EX4VSA, val) = value.ext = val


export gradient!

function gradient!(top::Var; debug=false)
    sorted = topsort(top)
    isgraddefined(top) || initgrad!(top)
    for v in sorted
        f = v.creator
        f isa SymbolFunc && continue
        debug && println(typeof(f))
        gys = [output.value.grad for output in f.outputs]
        gxs = backward(v.creator, gys...)
        gxs isa Tuple || (gxs = (gxs,))
        for (x, gx) in zip(f.args, gxs)
            debug && println("\t", typeof(gx), ":", size(gx))
            isgraddefined(x) && !isnothing(gx) ? x.grad = x.grad + gx : x.grad = gx
        end
    end
    return sorted[end]
end


module PyFenicsAD

import ChainRulesCore
using PyCall

const fecr = PyNULL()
const to_numpy = PyNULL()
const from_numpy = PyNULL()
const evaluate_primal = PyNULL()
const evaluate_pullback = PyNULL()
const evaluate_pushforward = PyNULL()

function __init__()
    copy!(fecr, pyimport("fecr"))
    copy!(to_numpy, fecr.to_numpy)
    copy!(from_numpy, fecr.from_numpy)
    copy!(evaluate_primal, fecr.evaluate_primal)
    copy!(evaluate_pullback, fecr.evaluate_pullback)
    copy!(evaluate_pushforward, fecr.evaluate_pushforward)
end

export fecr, to_numpy, from_numpy, evaluate_primal, evaluate_pullback, evaluate_pushforward

"""
A convenience macro that creates a julia function and registers forward or reverse rules for the provided FEniCS/Firedrake `fem_function`.
"""
macro register_fem_function(julia_function_name, fem_templates, fem_function)
    name_str = "$julia_function_name"
    quote
        function $(esc(julia_function_name))(inputs...)
            return evaluate_primal($(esc(fem_function)), $(esc(fem_templates)), inputs...)[1]
        end

        function ChainRulesCore.rrule(::typeof($(esc(julia_function_name))), inputs...)
            pyout = pycall(
                evaluate_primal,
                PyObject,
                $(esc(fem_function)),
                $(esc(fem_templates)),
                inputs...,
            )
            numpy_output, fem_output, fem_inputs, tape =
                [get(pyout, PyObject, i) for i = 0:3]

            # Only single-output functions are supported for now
            function fem_pullback(g)
                vjp_out = evaluate_pullback(fem_output, fem_inputs, tape, g)
                return (ChainRulesCore.NO_FIELDS, vjp_out...)
            end
            return get(pyout, 0), fem_pullback
        end

        function ChainRulesCore.frule(
            (_, Δinputs),
            ::typeof($(esc(julia_function_name))),
            inputs...,
        )
            pyout = pycall(
                evaluate_primal,
                PyObject,
                $(esc(fem_function)),
                $(esc(fem_templates)),
                inputs...,
            )
            numpy_output, fem_output, fem_inputs, tape =
                [get(pyout, PyObject, i) for i = 0:3]

            ∂numpy_output = evaluate_pushforward(fem_output, fem_inputs, tape, Δinputs)
            return get(pyout, 0), ∂numpy_output
        end

        println("Registered FEniCS/Firedrake function with name ", $name_str, "!")
    end
end

export @register_fem_function

end # module


#jl #! format: off
# # Sequential Simulations with [PowerSimulations.jl](https://github.com/NREL-SIIP/PowerSimulations.jl)

# **Originally Contributed by**: Clayton Barrows

# ## Introduction

# PowerSimulations.jl supports simulations that consist of sequential optimization problems
# where results from previous problems inform subsequent problems in a variety of ways. This
# example demonstrates some of these capabilities to represent electricity market clearing.
# This example is intended to be an extension of the
# [OperationsProblem example.](https://nbviewer.jupyter.org/github/NREL-SIIP/SIIPExamples.jl/blob/master/notebook/3_PowerSimulations_examples/01_operations_problems.ipynb)

# ## Dependencies
using SIIPExamples

# ### Modeling Packages
using PowerSystems
using PowerSimulations
const PSI = PowerSimulations
using PowerSystemCaseBuilder
using Dates

# ### Optimization packages
using HiGHS #solver

# ### Optimizer
# It's most convenient to define an optimizer instance upfront and pass it into the
# `DecisionModel` constructor. For this example, we can use the free HiGHS solver with a
# relatively relaxed MIP gap (`ratioGap`) setting to improve speed.

solver = optimizer_with_attributes(HiGHS.Optimizer, "mip_rel_gap" => 0.5)

# ### Hourly day-ahead system
# First, we'll create a `System` with hourly data to represent day-ahead forecasted wind,
# solar, and load profiles:
sys_DA = build_system(PSITestSystems, "modified_RTS_GMLC_DA_sys")

# ### 5-Minute system
# The RTS data also includes 5-minute resolution time series data. So, we can create another
# `System` to represent 15 minute ahead forecasted data for a "real-time" market:
sys_RT = build_system(PSITestSystems, "modified_RTS_GMLC_RT_sys")

# ## `ProblemTemplate`s define stages
# Sequential simulations in PowerSimulations are created by defining `OperationsProblems`
# that represent stages, and how information flows between executions of a stage and
# between different stages.
#
# Let's start by defining a two stage simulation that might look like a typical day-Ahead
# and real-time electricity market clearing process.

# ### Day-ahead unit commitment stage
# First, we can define a unit commitment template for the day ahead problem. We can use the
# included UC template, but in this example, we'll replace the `ThermalBasicUnitCommitment`
# with the slightly more complex `ThermalStandardUnitCommitment` for the thermal generators.
template_uc = template_unit_commitment()
set_device_model!(template_uc, ThermalStandard, ThermalStandardUnitCommitment)

# ### Define the reference model for the real-time economic dispatch
# In addition to the manual specification process demonstrated in the OperationsProblem
# example, PSI also provides pre-specified templates for some standard problems:
template_ed = template_economic_dispatch(
    network = NetworkModel(StandardPTDFModel, PTDF = PTDF(sys_DA), use_slacks = true),#NetworkModel(CopperPlatePowerModel, use_slacks = true),
)

# ### Define the `SimulationModels`
# `DecisionModel`s define the problems that are executed in the simulation.
# The actual problem will change as the stage gets updated to represent
# different time periods, but the formulations applied to the components is constant within
# a stage. In this case, we want to define two stages with the `ProblemTemplate`s
# and the `System`s that we've already created.
models = SimulationModels(
    decision_models = [
        DecisionModel(template_uc, sys_DA, optimizer = solver, name = "UC"),
        DecisionModel(template_ed, sys_RT, optimizer = solver, name = "ED"),
    ],
)

# ### `SimulationSequence`
# Similar to an `ProblemTemplate`, the `SimulationSequence` provides a template of
# how to execute a sequential set of operations problems.

#nb # print_struct(SimulationSequence)

# Let's review some of the `SimulationSequence` arguments.

# ### Chronologies
# In PowerSimulations, chronologies define where information is flowing. There are two types
# of chronologies.
#  - inter-stage chronologies: Define how information flows between stages. e.g. day-ahead
# solutions are used to inform economic dispatch problems
#  - intra-stage chronologies: Define how information flows between multiple executions of a
# single stage. e.g. the dispatch setpoints of the first period of an economic dispatch problem
# are constrained by the ramping limits from setpoints in the final period of the previous problem.

# ### `FeedForward`
# The definition of exactly what information is passed using the defined chronologies is
# accomplished with `FeedForward`. Specifically, `FeedForward` is used
# to define what to do with information being passed with an inter-stage chronology. Let's
# define a `FeedForward` that affects the semi-continuous range constraints of thermal generators
# in the economic dispatch problems based on the value of the unit-commitment variables.

feedforward = Dict(
    "ED" => [
        SemiContinuousFeedforward(
            component_type = ThermalStandard,
            source = OnVariable,
            affected_values = [ActivePowerVariable],
        ),
    ],
)

# ### Sequencing
# The stage problem length, look-ahead, and other details surrounding the temporal Sequencing
# of stages are controlled using the structure of the time series data in the `System`s.
# So, to define a typical day-ahead - real-time sequence:
#  - Day ahead problems should represent 48 hours, advancing 24 hours after each execution (24-hour look-ahead)
#  - Real time problems should represent 1 hour (12 5-minute periods), advancing 15 min after each execution (15 min look-ahead)
# We can adjust the time series data to reflect this structure in each `System`:
# - `transform_single_time_series!(sys_DA, 48, Hour(1))`
# - `transform_single_time_series!(sys_RT, 12, Minute(15))`
#
# Now we can put it all together to define a `SimulationSequence`

DA_RT_sequence = SimulationSequence(
    models = models,
    ini_cond_chronology = InterProblemChronology(),
    feedforwards = feedforward,
)

# ## `Simulation`
# Now, we can build and execute a simulation using the `SimulationSequence` and `Stage`s
# that we've defined.
sim = Simulation(
    name = "rts-test",
    steps = 2,
    models = models,
    sequence = DA_RT_sequence,
    simulation_folder = mktempdir(".", cleanup = true),
)

# ### Build simulation
build!(sim)

# ### Execute simulation
# the following command returns the status of the simulation (0: is proper execution) and
# stores the results in a set of HDF5 files on disk.
execute!(sim, enable_progress_bar = false)

# ## Results
# To access the results, we need to load the simulation result metadata and then make
# requests to the specific data of interest. This allows you to efficiently access the
# results of interest without overloading resources.
results = SimulationResults(sim);
uc_results = get_decision_problem_results(results, "UC"); # UC stage result metadata
ed_results = get_decision_problem_results(results, "ED"); # ED stage result metadata

# We can read all the result variables
read_variables(uc_results)

# or all the parameters
read_parameters(uc_results)

# We can just list the variable names contained in `uc_results`:
list_variable_names(uc_results)

# and a number of parameters (this pattern also works for aux_variables, expressions, and duals)
list_parameter_names(uc_results)

# Now we can read the specific results of interest for a specific problem, time window (optional),
# and set of variables, duals, or parameters (optional)

Dict([
    v => read_variable(uc_results, v) for v in [
        "ActivePowerVariable__RenewableDispatch",
        "ActivePowerVariable__HydroDispatch",
        "StopVariable__ThermalStandard",
    ]
])

# Or if we want the result of just one variable, parameter, or dual (must be defined in the
# problem definition), we can use:

read_parameter(
    ed_results,
    "ActivePowerTimeSeriesParameter__RenewableFix",
    initial_time = DateTime("2020-01-01T06:00:00"),
    count = 5,
)

# * note that this returns the results of each execution step in a separate dataframe *
# If you want the realized results (without lookahead periods), you can call `read_realized_*`:

read_realized_variables(
    uc_results,
    ["ActivePowerVariable__ThermalStandard", "ActivePowerVariable__RenewableDispatch"],
)

# ## Plotting
# Take a look at the [plotting examples.](https://nbviewer.jupyter.org/github/NREL-SIIP/SIIPExamples.jl/blob/master/notebook/3_PowerSimulations_examples/04_bar_stack_plots.ipynb)


# Copyright © nhdalyMadeThis, LLC
# Released under MIT License

module HueClone

using Colors
using Rematch # For matching text-based user input

export play_blink, play_juno

const DEFAULT_ROWS, DEFAULT_COLS = 6,4

function create_colors_grid(rows, cols)
    corners = rand(RGB, 4)
    topleft,topright,bottomleft,bottomright = corners

    leftedge = range(topleft, stop=bottomleft, length=rows)
    rightedge = range(topright, stop=bottomright, length=rows)

    tiles = Array{RGB,2}(undef, (rows,cols));
    for r in 1:rows
        tiles[r,:] = range(leftedge[r], stop=rightedge[r], length=cols)
    end

    tiles
end

macro swap!(a,b)
    quote
        tmp = $(esc(a))
        $(esc(a)) = $(esc(b))
        $(esc(b)) = tmp
    end
end

function user_attempt_swap!(a, b)
    a,b = Tuple(a), Tuple(b)
    (r,c) = size(tiles)
    fixed_tile_positions = [(1,1), (r,1), (1,c), (r,c)]
    if a == b || a in fixed_tile_positions || b in fixed_tile_positions
        return
    end
    @swap!(tiles[a...], tiles[b...])
end

function new_game(rows, cols)
    global goal = create_colors_grid(rows, cols)
    global tiles = deepcopy(goal)
    for _ in 1:100
        user_attempt_swap!(rand(CartesianIndices(tiles)),
                           rand(CartesianIndices(tiles)))
    end
end

function play_game(io_device; size=(DEFAULT_ROWS, DEFAULT_COLS))
    new_game(size...)
    while tiles != goal
        print_board(io_device, tiles)
        a,b = collect_userinput(io_device)
        user_attempt_swap!(a, b)
    end
    print_board(io_device, tiles)
    win_game(io_device)
end

# ------------------------------------
# ------ Juno-REPL based -------------
# ------------------------------------

struct JunoDisplay end

"""
    play_juno()

Play the game in Atom, with the Juno REPL and Juno Plots display
"""
play_juno(;size=(DEFAULT_ROWS, DEFAULT_COLS)) = play_game(JunoDisplay(), size=size)

print_board(::JunoDisplay, tiles) = display(tiles)

function collect_userinput(::JunoDisplay)
    while true
        try
            println("Enter tile coordinates to swap (col,row):")
            # Read a and b from user input
            print("a: ")
            @match Expr(Tuple, a) = Meta.parse(readline())
            print("b: ")
            @match Expr(Tuple, b) = Meta.parse(readline())
            return reverse(a), reverse(b)
        catch e
            e isa InterruptException && rethrow()
            @warn("Must enter valid coordinate tuple (x,y)")
            @error(e)
        end
    end
end
function win_game(::JunoDisplay)
    println("Hooray you win! :)")
end

# ------------------------------------
# ------ Blink based -----------------
# ------------------------------------

using Blink

"""
    play_blink()

Play the game using a Blink html window.
"""
function play_blink(;size=(DEFAULT_ROWS, DEFAULT_COLS))
    w = create_board(size...)
    handle(w, "click") do pos
        handle_click(pos)
    end

    play_game(w; size=size)
end

function create_board(rows, cols)
    sq_size = 50
    w = Blink.Window(Dict(:width => sq_size*cols + 4,
                          :height => sq_size*rows + 40),
                     async=false)
    body!(w, """

    """, async=false)
    return w
end
function print_board(w::Blink.Window, tiles)
    content!(w, ".grid",
            join([
                """

                    $(join([
                        """

"""
                        for r in 1:size(tiles)[1]
                        ], "")
                    )

                """
                for c in 1:size(tiles)[2]
                ], "")
        ; fade=false, async=false)
end

clicksChannel = Channel(2)
function handle_click(pos)
    global clicksChannel
    put!(clicksChannel, Tuple(pos))
end
function collect_userinput(w::Blink.Window)
    return (take!(clicksChannel), take!(clicksChannel))
end

function win_game(w::Blink.Window)
    content!(w, ".messages", "YOU WIN! 🤣")
end

end


#!/usr/bin/env julia

using Luxor

using Test

using Colors

using Random

Random.seed!(42)

demomovie = Movie(400, 400, "test", 0:359)

function backdrop(scene, framenumber)
    background("black")
end

function frame(scene, framenumber)
    sethue(Colors.HSV(framenumber, 1, 1))
    p = scene.easingfunction(framenumber, 0, 1, scene.framerange.stop)
    circle(polar(100, -pi/2 - p * 2pi), 80, :fill)
    text(string("frame $framenumber of $(length(scene.movie.movieframerange))"), Point(O.x, O.y-190))
end

mktempdir() do tmpdir
    @test animate(demomovie, [Scene(demomovie, backdrop), Scene(demomovie, frame, 0:359, easingfunction=easeinoutquad)], creategif=false) == true
end

demo = Movie(400, 100, "test")

backdrop(scene, framenumber) =  background("black")

function frame1(scene, framenumber)
    sethue("purple")
    eased_n = scene.easingfunction(framenumber, 0, 1, scene.framerange.stop)
    circle(Point(-180 + (360 * eased_n), -20), 10, :fill)
end

function frame2(scene, framenumber)
    sethue("green")
    eased_n = scene.easingfunction(framenumber, 0, 1, scene.framerange.stop)
    circle(Point(-180 + (360 * eased_n), 0), 10, :fill)
end

function frame3(scene, framenumber)
    sethue("red")
    eased_n = scene.easingfunction(framenumber, 0, 1, scene.framerange.stop)
    circle(Point(-180 + (360 * eased_n), 20), 10, :fill)
end

function frame4(scene, framenumber)
    sethue("red")
    eased_n = scene.easingfunction(framenumber, 0, 1, scene.framerange.stop)
    circle(Point(-180 + (360 * eased_n), 20), 10, :fill)
    sethue("white")
    fontsize(20)
    text(string(scene.opts))
end

mktempdir() do tmpdir
    @test animate(demo, [
        Scene(demo, backdrop, 0:200),
        Scene(demo, frame1,   0:200, easingfunction=easeinsine),
        Scene(demo, frame2,   0:200, easingfunction=easeinoutcubic),
        Scene(demo, frame3,   0:200, easingfunction=easeinoutquint),
        Scene(demo, frame4,   0:200, easingfunction=easeinoutquint, optarg=42),
        ],
        creategif=false) == true
end

println("...finished animation tests")


using .JSort
using PyCall
import PyPlot; const plt = PyPlot;

function localalign(detector::ChannelDetector,
                    region; nbins=5000, emin=-150, emax=15000,
                    kwargs...)
    full, edges = histogram(detector, nbins=nbins)
    chns = goodchannels(detector)

    signal = pyimport("scipy.signal")
    smoother(x) = signal.savgol_filter(x, 51, 4)

    eclip(x) = clip(x, emin, emax)

    kw = get(kwargs, :a1, Dict())
    if get(kw, :plot, false)
        plothist(full, edges)
    end
    delta = get(kw, :delta, [3500, 3500])
    searchregion = eclip(region[1]-delta[1]):eclip(region[2]+delta[2])

    println("Aligning")

    coeff, _ = alignspectra(full, edges, lowregion=searchregion, lowsmoother=smoother,
                            lownumregions=get(kw, :nregions, 2),
                            lowregionwidth=get(kw, :swidth, 200),
                            lowsearchwidth=get(kw, :lwidth, 60),
                            plot=get(kw, :plot, false), order=1)
    coefficients = repeat([0.0 1.0], outer=(33, 1))
    coefficients[chns, :] .= hcat(coeff...)'
    calib = GammaCalibrator(coefficients)
    aligned = calib(detector, inplace=false)

    if get(kw, :plot, false)
        plot(aligned)
    end

    alignedhist = histogram(aligned, edges)

    kw = get(kwargs, :a2, Dict())
    delta = get(kw, :delta, [3500, 3500])
    searchregion = eclip(region[1]-delta[1]):eclip(region[2]+delta[2])

    coeff, _ = alignspectra(alignedhist, edges, lowregion=searchregion, lowsmoother=smoother,
                            lownumregions=get(kw, :nregions, 3),
                            lowregionwidth=get(kw, :swidth, 200),
                            lowsearchwidth=get(kw, :lwidth, -10),
                            plot=get(kw, :plot, false), order=1)

    coefficients = repeat([0.0 1.0], outer=(33, 1))
    coefficients[chns, :] .= hcat(coeff...)'
    combine!(calib, GammaCalibrator(coefficients))
    aligned = calib(detector, inplace=false)

    if get(kw, :plot, false)
      plot(aligned)
    end

    #S = sum(aligned)
    #hist = histogram(S, edges)
    #plot(S)
    # Energy - index is wrong
    #peak = findpeaks(hist[region], numpeaks=1)[1]
    #peak += first(region)
    #@show peak
    #@show peaks = map(x -> x(peak), invcalib.poly)
    #ax.scatter(peaks[chns], 1.5:length(chns)+0.5, marker="x", color="r", s=2)

    return calib
end

function regionalign(d::ChannelDetector, regions; nbins=1000,
                     emin=100, emax=2500, kwargs...)
    strips, edges = JSort.histogram(d, nbins=nbins)
    shift, gain = first(edges), step(edges)
    indextoe(x) = shift + gain*x
    etoindex(x) = round(Int, -shift/gain + 1/gain*x)
    #regions = [map(etoindex, region) for region in regions]
    eclip(x) = clip(x, emin, emax)
    delta = 400

    peaks = []
    for region in regions
        low, high = region[1], region[2]
        # Align all strips to the reference strip
        lowregion = eclip(low-delta):eclip(high+delta)
        @show lowregion
        println("Aligning")
        al, calib, _ = align(d, lowregion=lowregion, lownumregions=2,
                             lowregionwidth=200,
                             lowsearchwidth=80, order=1, kwargs...);
        println("Done aliging")
        println("Region peaking")
        # Find the peak(s) in the aligned region
        peak = regionpeaks(al, [region])
        println("Done region peaking")
        # Backshift the peaks to the original coordinate system
        peak = backshift(peak, goodchannels(d), invert(calib))
        push!(peaks, peak)
    end
    peaks
end

function regionpeaks(d::ChannelDetector, regions; nbins=5000)
    strips, edges = JSort.histogram(d, nbins=nbins)
    shift, gain = first(edges), step(edges)
    indextoe(x) = shift + gain*x
    etoindex(x) = round(Int, -shift/gain + 1/gain*x)
    regions = [map(etoindex, region) for region in regions]
    println("Regioning for $regions")
    peaks = regionpeaks(strips, regions)
    [map(indextoe, peak) for peak in peaks]
end

function regionpeaks(strips, regions)
    peaks = Vector{Int}[]
    for strip in strips
        peaks_ = regionpeaks(strip, regions)
        push!(peaks, peaks_)
    end
    peaks
end

function regionpeaks(strip::AbstractVector{<:Number}, regions)::Vector{Int}
    peaks = Int[]
    for region in regions
        peak = regionpeak(strip, region[1], region[2])
        push!(peaks, peak)
    end
    return peaks
end

function regionpeak(strip::AbstractVector{<:Number}, lower::Int, upper::Int)::Int
    region = strip[lower:upper]
    peak = findpeaks(region, numpeaks=1)[1]
    return peak + lower - 2 #BUG -2 is indicative of a offset bug
end

function backshift(points, channels, calib)
    shiftedpoints = similar(points)
    for i in 1:length(points)
        shiftedpoints[i] = calib[channels[i]](points[i])
    end
    shiftedpoints
end

function backshift(points::AbstractVector, coefficients::AbstractVector)
    shiftedpoints = similar(points)
    for i in 1:length(points)
        shift, gain = coefficients[i]
        poly(x) = x/gain - shift/gain
        shiftedpoints[i] = poly(points[i])
    end
    shiftedpoints
end

function backshift(point::Number, coefficients::AbstractVector)
    shiftedpoints = []
    for i in 1:length(points)
        shift, gain = coefficients[i]
        poly(x) = x/gain - shift/gain
        push!(shiftedpoints, poly(point))
    end
    shiftedpoints
end

function clip(x, lower, upper)
    if x < lower
        lower
    elseif x > upper
        upper
    else
        x
    end
end

function plothist(strips, edges; points=nothing)
    mat = cat(strips'..., dims=1)
    fig, ax = plt.subplots()
    ax.pcolormesh(edges, 1:length(strips)+1, mat .|> log10)
    if !isnothing(points)
        for points_ in points
            for (i, point) in enumerate(points_)
                ax.scatter(point, i+0.5, marker="x", color="r", s=3)
            end
        end
    end
    ax
end

function plothist(det::ChannelDetector; nbins=1000, points=nothing)
    det_strips, det_edges = JSort.histogram(det, nbins=nbins)
    plothist(det_strips, det_edges, points=points)
end


#=
A nieve metric for values of any type: simple subtraction of their values.

Nieve for characters,
since it doesn't represent the knowledge that humans have re classes of characters.
=#

function pointDifference(
        targetValue::T,
        corpusValue::T
        ) where{T}
    return corpusValue - targetValue
end


## Solve MergingMDP for different human action spaces

using Revise
using Random
using LinearAlgebra
using POMDPs
using POMDPModelTools
using POMDPPolicies
using POMDPSimulators
using IterTools
using AutomotiveDrivingModels
using AutomotivePOMDPs: polygon, overlap
using Parameters
using GridInterpolations
using StaticArrays
using AutoUrban
using AutoViz
using Printf
using Reel
using DiscreteValueIteration
using SparseArrays
using ProgressMeter
using JLD2
using FileIO

includet("environment.jl")
includet("mdp_definition.jl")
includet("helpers.jl")

mdp = MergingMDP()

@show n_states(mdp)

assumptions = [:wide, :cv, :ca, :cd]
action_spaces = [[-2.0, 0.0, 2.0], [0.0], [2.0], [-2.0]]
# solver = SparseValueIterationSolver(max_iterations=300, verbose=true)

# policies = ValueIterationPolicy[]

# for (i, human_actions) in enumerate(action_spaces)
#     mdp.human_actions = human_actions
#     push!(policies, solve(solver, mdp))
# end

# # mdp.human_actions = [-2.0]
# # policies[end] = solve(solver, mdp)

# @save "/scratch/boutonm/policies.jld2" policies

@load "policies.jld2" policies

## Display policies
using Plots
using Interact
using Blink

function policy_plot(mdp::MergingMDP, policy::ValueIterationPolicy, ve::Float64=0.0, vo::Float64 = 0.0, assumption::Symbol=:wide)
    xs = mdp.ego_grid[1:end]
    ys = mdp.human_grid
    ae = mdp.acceleration_grid[3]
    z = map(((y,x),) -> value(policy, Float64[x, ve, ae, y, vo]), Iterators.product(ys, xs))
    za = map(((y,x),) -> action(policy, Float64[x, ve, ae, y, vo]), Iterators.product(ys, xs))
    p = [heatmap(xs, ys, z, title="Pr(success) ($assumption)",
                     aspect_ratio=1, color=:plasma, clim=(0,1), size=(1400, 600), colorbar=false),
        heatmap(xs, ys, za, title="Action Map",  xlabel="Ego longitudinal position",
                     aspect_ratio=1,color=:viridis, clim=(1,7), size=(1400, 600), colorbar=false)
        ]
    scatter!(p[1], [20.0], [20.0], axis=:none, legend=:none, color=:red)
    scatter!(p[2], [20.0], [20.0], axis=:none, legend=:none, color=:red)
    return p
end

w1 = Window()
ui = @manipulate for ve = mdp.velocity_grid, vo = mdp.velocity_grid
    xs = mdp.ego_grid[1:end]
    ys = mdp.human_grid
    ae = mdp.acceleration_grid[3]
    plots = []
    for (i, policy) in enumerate(policies)
        push!.(Ref(plots), policy_plot(mdp, policy, ve, vo, assumptions[i]))
    end
    plot!(plots[1],  ylabel="Other longitudinal position")
    plot!(plots[end], colorbar=true, zticks=String["brake", "-0.2", "-0.1", "0.0", "0.1", "0.2", "release"])
    plot!(plots[end-1], colorbar=true)
    plot(plot([plots[i] for i=1:2:7]..., layout=(1, 4)),
         plot([plots[i] for i=2:2:8]..., layout=(1, 4)), layout=(2,1))
end
body!(w1, ui)

## Visualize policy running

policy = policies[3]
hr = HistoryRecorder(rng = rng, max_steps=200)
rng = MersenneTwister(1)
s0 = [mdp.ego_grid[1], mdp.velocity_grid[5], mdp.acceleration_grid[3], mdp.human_grid[1], mdp.velocity_grid[3]]
hist = simulate(hr, mdp, policy, s0)

w3 = Window()
ui = @manipulate for step in 1:n_steps(hist)
    s = state_hist(hist)[step]
    a = action_hist(hist)[step]
    v = MergingViz(mdp=mdp, s=s, a=a, action_values = actionvalues(policy, s))
    AutoViz.render([v], cam=StaticCamera(VecE2(0.0, 0.0), 20.0))
end
body!(w3, ui)


"""
This function creates the model for a full themal dispatch formulation depending on combination of devices, device_formulation and system_formulation
"""
function construct_device!(canonical::Canonical, sys::PSY.System,
                           model::DeviceModel{T, D},
                           ::Type{S};
                           kwargs...) where {T<:PSY.ThermalGen,
                                             D<:AbstractThermalFormulation,
                                             S<:PM.AbstractPowerModel}

    devices = PSY.get_components(T, sys)

    if validate_available_devices(devices, T)
        return
    end

    #Variables
    activepower_variables!(canonical, devices)

    reactivepower_variables!(canonical, devices)

    commitment_variables!(canonical, devices)

    #Initial Conditions

    initial_conditions!(canonical, devices, D)

    #Constraints
    activepower_constraints!(canonical, devices, D, S)

    reactivepower_constraints!(canonical, devices, D, S)

    commitment_constraints!(canonical, devices, D, S)

    ramp_constraints!(canonical, devices, D, S)

    time_constraints!(canonical, devices, D, S)

    feedforward!(canonical, T, model.feedforward)

    #Cost Function
    cost_function(canonical, devices, D, S)

    return

end

"""
This function creates the model for a full themal dispatch formulation depending on combination of devices, device_formulation and system_formulation
"""
function construct_device!(canonical::Canonical, sys::PSY.System,
                           model::DeviceModel{T, D},
                           ::Type{S};
                           kwargs...) where {T<:PSY.ThermalGen,
                                             D<:AbstractThermalFormulation,
                                             S<:PM.AbstractActivePowerModel}

    devices = PSY.get_components(T, sys)

    if validate_available_devices(devices, T)
        return
    end

    #Variables
    activepower_variables!(canonical, devices)

    commitment_variables!(canonical, devices)

    #Initial Conditions

    initial_conditions!(canonical, devices, D)

    #Constraints
    activepower_constraints!(canonical, devices, D, S)

    commitment_constraints!(canonical, devices, D, S)

    ramp_constraints!(canonical, devices, D, S)

    time_constraints!(canonical, devices, D, S)

    feedforward!(canonical, T, model.feedforward)

    #Cost Function
    cost_function(canonical, devices, D, S)

    return

end

"""
This function creates the model for a full themal dispatch formulation depending on combination of devices, device_formulation and system_formulation
"""
function construct_device!(canonical::Canonical, sys::PSY.System,
                           model::DeviceModel{T, ThermalBasicUnitCommitment},
                           ::Type{S};
                           kwargs...) where {T<:PSY.ThermalGen,
                                             S<:PM.AbstractPowerModel}

    devices = PSY.get_components(T, sys)

    if validate_available_devices(devices, T)
        return
    end

    #Variables
    activepower_variables!(canonical, devices)

    reactivepower_variables!(canonical, devices)

    commitment_variables!(canonical, devices)

    #Initial Conditions

    initial_conditions!(canonical, devices, model.formulation)

    #Constraints
    activepower_constraints!(canonical, devices, model.formulation, S)

    reactivepower_constraints!(canonical, devices, model.formulation, S)

    commitment_constraints!(canonical, devices, model.formulation, S)

    feedforward!(canonical, T, model.feedforward)

    #Cost Function
    cost_function(canonical, devices, model.formulation, S)

    return

end

"""
This function creates the model for a full themal dispatch formulation depending on combination of devices, device_formulation and system_formulation
"""
function construct_device!(canonical::Canonical, sys::PSY.System,
                           model::DeviceModel{T, ThermalBasicUnitCommitment},
                           ::Type{S};
                           kwargs...) where {T<:PSY.ThermalGen,
                                             S<:PM.AbstractActivePowerModel}

    devices = PSY.get_components(T, sys)

    if validate_available_devices(devices, T)
        return
    end

    #Variables
    activepower_variables!(canonical, devices)

    commitment_variables!(canonical, devices)

    #Initial Conditions

    initial_conditions!(canonical, devices, model.formulation)

    #Constraints
    activepower_constraints!(canonical, devices, model.formulation, S)

    commitment_constraints!(canonical, devices, model.formulation, S)

    feedforward!(canonical, T, model.feedforward)

    #Cost Function
    cost_function(canonical, devices, model.formulation, S)

    return

end

"""
This function creates the model for a full themal dispatch formulation depending on combination of devices, device_formulation and system_formulation
"""
function construct_device!(canonical::Canonical, sys::PSY.System,
                           model::DeviceModel{T, ThermalRampLimited},
                           ::Type{S};
                           kwargs...) where {T<:PSY.ThermalGen,
                                             S<:PM.AbstractPowerModel}

    devices = PSY.get_components(T, sys)

    if validate_available_devices(devices, T)
        return
    end

    #Variables
    activepower_variables!(canonical, devices)

    reactivepower_variables!(canonical, devices)

    #Initial Conditions

    initial_conditions!(canonical, devices, model.formulation)

    #Constraints
    if !(isa(model.feedforward, SemiContinuousFF))
        activepower_constraints!(canonical, devices, ThermalRampLimited, S)
    end

    reactivepower_constraints!(canonical, devices, model.formulation, S)

    ramp_constraints!(canonical, devices, model.formulation, S)

    feedforward!(canonical, T, model.feedforward)

    #Cost Function
    cost_function(canonical, devices, model.formulation, S)

    return

end

"""
This function creates the model for a full themal dispatch formulation depending on combination of devices, device_formulation and system_formulation
"""
function construct_device!(canonical::Canonical, sys::PSY.System,
                           model::DeviceModel{T, ThermalRampLimited},
                           ::Type{S};
                           kwargs...) where {T<:PSY.ThermalGen,
                                             S<:PM.AbstractActivePowerModel}

    devices = PSY.get_components(T, sys)

    if validate_available_devices(devices, T)
        return
    end

    #Variables
    activepower_variables!(canonical, devices)

    #Initial Conditions

    initial_conditions!(canonical, devices, model.formulation)

    #Constraints
    if !(isa(model.feedforward, SemiContinuousFF))
        activepower_constraints!(canonical, devices, ThermalRampLimited, S)
    end

    ramp_constraints!(canonical, devices, model.formulation, S)

    feedforward!(canonical, T, model.feedforward)

    #Cost Function
    cost_function(canonical, devices, model.formulation, S)

    return

end

function construct_device!(canonical::Canonical, sys::PSY.System,
                           model::DeviceModel{T, D},
                           ::Type{S};
                           kwargs...) where {T<:PSY.ThermalGen,
                                             D<:AbstractThermalDispatchFormulation,
                                             S<:PM.AbstractPowerModel}

    devices = PSY.get_components(T, sys)

    if validate_available_devices(devices, T)
        return
    end

    #Variables
    activepower_variables!(canonical, devices)

    reactivepower_variables!(canonical, devices)

    #Initial Conditions

    #Constraints
    if !(isa(model.feedforward, SemiContinuousFF))
        activepower_constraints!(canonical, devices, D, S)
    end

    reactivepower_constraints!(canonical, devices, D, S)

    feedforward!(canonical, T, model.feedforward)

    #Cost Function
    cost_function(canonical, devices, D, S)

    return

end

function construct_device!(canonical::Canonical, sys::PSY.System,
                           model::DeviceModel{T, D},
                           ::Type{S};
                           kwargs...) where {T<:PSY.ThermalGen,
                                             D<:AbstractThermalDispatchFormulation,
                                             S<:PM.AbstractActivePowerModel}

    devices = PSY.get_components(T, sys)

    if validate_available_devices(devices, T)
        return
    end

    #Variables
    activepower_variables!(canonical, devices)

    #Initial Conditions

    #Constraints
    # Slighly hacky for now
    if !(isa(model.feedforward, SemiContinuousFF))
        activepower_constraints!(canonical, devices, D, S)
    end

    feedforward!(canonical, T, model.feedforward)

    #Cost Function
    cost_function(canonical, devices, D, S)

    return

end