#TODO: Document interface to support static size
# Define maybestaticsize, maybestaticlength, and maybestaticfirstindex

@inline static(::Val{N}) where {N} = StaticInt{N}()
@inline static(::Nothing) = nothing

@inline maybestaticfirst(a) = static_first(a)
@inline maybestaticlast(a) = static_last(a)
@inline maybestaticlength(a) = static_length(a)
@inline maybestaticlength(a::UnitRange{T}) where {T} = last(a) - first(a) + oneunit(T)

@inline maybestaticrange(r::Base.OneTo{T}) where {T} =
  ArrayInterface.OptionallyStaticUnitRange(StaticInt{1}(), last(r))
@inline maybestaticrange(r::UnitRange) = r
@inline maybestaticrange(r) = maybestaticfirst(r):maybestaticlast(r)

@inline maybestaticsize(::NTuple{N}, ::Val{1}) where {N} = StaticInt{N}() # should we assert that i == 1?
@inline maybestaticsize(
  ::LinearAlgebra.Adjoint{T,V},
  ::Val{1},
) where {T,V<:AbstractVector{T}} = One()
@inline maybestaticsize(
  ::LinearAlgebra.Transpose{T,V},
  ::Val{1},
) where {T,V<:AbstractVector{T}} = One()
@inline maybestaticsize(A, ::Val{N}) where {N} = ArrayInterface.size(A)[N]

# These have versions that may allow for more optimizations, so we override base methods with a single `StaticInt` argument.
for (f, ff) ∈ [
  (:(Base.:+), :vadd_fast),
  (:(Base.:-), :vsub_fast),
  (:(Base.:*), :vmul_fast),
  (:(Base.:+), :vadd_nsw),
  (:(Base.:-), :vsub_nsw),
  (:(Base.:*), :vmul_nsw),
  (:(Base.:+), :vadd_nuw),
  (:(Base.:-), :vsub_nuw),
  (:(Base.:*), :vmul_nuw),
  (:(Base.:+), :vadd_nw),
  (:(Base.:-), :vsub_nw),
  (:(Base.:*), :vmul_nw),
  (:(Base.:<<), :vshl),
  (:(Base.:÷), :vdiv),
  (:(Base.:%), :vrem),
  (:(Base.:>>>), :vashr),
]
  @eval begin
    # @inline $f(::StaticInt{M}, ::StaticInt{N}) where {M, N} = StaticInt{$f(M, N)}()
    # If `M` and `N` are known at compile time, there's no need to add nsw/nuw flags.
    @inline $ff(::StaticInt{M}, ::StaticInt{N}) where {M,N} =
      $f(StaticInt{M}(), StaticInt{N}())
    # @inline $f(::StaticInt{M}, x) where {M} = $ff(M, x)
    # @inline $f(x, ::StaticInt{M}) where {M} = $ff(x, M)
    @inline $ff(::StaticInt{M}, x::T) where {M,T<:IntegerTypesHW} = $ff(M % T, x)
    @inline $ff(x::T, ::StaticInt{M}) where {M,T<:IntegerTypesHW} = $ff(x, M % T)
    @inline $ff(::StaticInt{M}, x::T) where {M,T} = $ff(T(M), x)
    @inline $ff(x::T, ::StaticInt{M}) where {M,T} = $ff(x, T(M))
  end
end
for f ∈ [:vadd_fast, :vsub_fast, :vmul_fast]
  @eval begin
    @inline $f(::StaticInt{M}, n::T) where {M,T<:Number} = $f(T(M), n)
    @inline $f(m::T, ::StaticInt{N}) where {N,T<:Number} = $f(m, T(N))
  end
end
for f ∈ [:vsub, :vsub_fast, :vsub_nsw, :vsub_nuw, :vsub_nw]
  @eval begin
    @inline $f(::Zero, m::Number) = -m
    @inline $f(::Zero, m::IntegerTypesHW) = -m
    @inline $f(m::Number, ::Zero) = m
    @inline $f(m::IntegerTypesHW, ::Zero) = m
    @inline $f(::Zero, ::Zero) = Zero()
    @inline $f(::Zero, ::StaticInt{N}) where {N} = -StaticInt{N}()
    @inline $f(::StaticInt{N}, ::Zero) where {N} = StaticInt{N}()
  end
end
for f ∈ [:vadd, :vadd_fast, :vadd_nsw, :vadd_nuw, :vadd_nw]
  @eval begin
    @inline $f(::StaticInt{N}, ::Zero) where {N} = StaticInt{N}()
    @inline $f(::Zero, ::StaticInt{N}) where {N} = StaticInt{N}()
    @inline $f(::Zero, ::Zero) = Zero()
    @inline $f(a::Number, ::Zero) = a
    @inline $f(a::IntegerTypesHW, ::Zero) = a
    @inline $f(::Zero, a::Number) = a
    @inline $f(::Zero, a::IntegerTypesHW) = a
  end
end

@inline vmul_fast(::StaticInt{N}, ::Zero) where {N} = Zero()
@inline vmul_fast(::Zero, ::StaticInt{N}) where {N} = Zero()
@inline vmul_fast(::Zero, ::Zero) = Zero()
@inline vmul_fast(::StaticInt{N}, ::One) where {N} = StaticInt{N}()
@inline vmul_fast(::One, ::StaticInt{N}) where {N} = StaticInt{N}()
@inline vmul_fast(::One, ::One) = One()
@inline vmul_fast(a::Number, ::One) = a
@inline vmul_fast(a::MM, ::One) = a
@inline vmul_fast(a::IntegerTypesHW, ::One) = a
@inline vmul_fast(::One, a::Number) = a
@inline vmul_fast(::One, a::MM) = a
@inline vmul_fast(::One, a::IntegerTypesHW) = a
@inline vmul_fast(::Zero, ::One) = Zero()
@inline vmul_fast(::One, ::Zero) = Zero()

for T ∈ [:VecUnroll, :AbstractMask, :MM]
  @eval begin
    @inline Base.:(+)(x::$T, ::Zero) = x
    @inline Base.:(+)(::Zero, x::$T) = x
    @inline Base.:(-)(x::$T, ::Zero) = x
    @inline Base.:(*)(x::$T, ::One) = x
    @inline Base.:(*)(::One, x::$T) = x
    @inline Base.:(*)(::$T, ::Zero) = Zero()
    @inline Base.:(*)(::Zero, ::$T) = Zero()
  end
end
@inline Base.:(+)(m::AbstractMask{W}, ::StaticInt{N}) where {N,W} =
  m + vbroadcast(Val{W}(), N)
@inline Base.:(+)(::StaticInt{N}, m::AbstractMask{W}) where {N,W} =
  vbroadcast(Val{W}(), N) + m
# @inline Base.:(*)(::StaticInt{N}, m::Mask{W}) where {N,W} = vbroadcast(Val{W}(), N) * m
@inline vadd_fast(x::VecUnroll, ::Zero) = x
@inline vadd_fast(::Zero, x::VecUnroll) = x
@inline vsub_fast(x::VecUnroll, ::Zero) = x
@inline vmul_fast(x::VecUnroll, ::One) = x
@inline vmul_fast(::One, x::VecUnroll) = x
@inline vmul_fast(::VecUnroll, ::Zero) = Zero()
@inline vmul_fast(::Zero, ::VecUnroll) = Zero()

for V ∈ [:AbstractSIMD, :MM]
  @eval begin
    @inline Base.FastMath.mul_fast(::Zero, x::$V) = Zero()
    @inline Base.FastMath.mul_fast(::One, x::$V) = x
    @inline Base.FastMath.mul_fast(x::$V, ::Zero) = Zero()
    @inline Base.FastMath.mul_fast(x::$V, ::One) = x

    @inline Base.FastMath.add_fast(::Zero, x::$V) = x
    @inline Base.FastMath.add_fast(x::$V, ::Zero) = x

    @inline Base.FastMath.sub_fast(::Zero, x::$V) = -x
    @inline Base.FastMath.sub_fast(x::$V, ::Zero) = x
  end
end


##############################################################################
##
## Iterate on terms
##
##############################################################################

eachterm(x::AbstractTerm) = (x,)
eachterm(x::NTuple{N, AbstractTerm}) where {N} = x
TermOrTerms = Union{AbstractTerm, NTuple{N, AbstractTerm} where N}
hasintercept(t::TermOrTerms) =
    InterceptTerm{true}() ∈ terms(t) ||
    ConstantTerm(1) ∈ terms(t)
omitsintercept(f::FormulaTerm) = omitsintercept(f.rhs)
omitsintercept(t::TermOrTerms) =
    InterceptTerm{false}() ∈ terms(t) ||
    ConstantTerm(0) ∈ terms(t) ||
    ConstantTerm(-1) ∈ terms(t)
##############################################################################
##
## Parse IV
##
##############################################################################

function parse_iv(@nospecialize(f::FormulaTerm))
	for term in eachterm(f.rhs)
		if term isa FormulaTerm
			formula_endo = FormulaTerm(ConstantTerm(0), tuple(ConstantTerm(0), eachterm(term.lhs)...))
			formula_iv = FormulaTerm(ConstantTerm(0), tuple(ConstantTerm(0), eachterm(term.rhs)...))
            exos = Tuple((term for term in eachterm(f.rhs) if !isa(term, FormulaTerm)))
            return FormulaTerm(f.lhs, exos), formula_endo, formula_iv
		end
	end
	return f, nothing, nothing
end

##############################################################################
##
## Parse FixedEffect
##
##############################################################################
struct FixedEffectTerm <: AbstractTerm
    x::Symbol
end
StatsModels.termvars(t::FixedEffectTerm) = [t.x]
fe(x::Term) = FixedEffectTerm(Symbol(x))

has_fe(::FixedEffectTerm) = true
has_fe(::FunctionTerm{typeof(fe)}) = true
has_fe(t::InteractionTerm) = any(has_fe(x) for x in t.terms)
has_fe(::AbstractTerm) = false
has_fe(t::FormulaTerm) = any(has_fe(x) for x in eachterm(t.rhs))

fesymbol(t::FixedEffectTerm) = t.x
fesymbol(t::FunctionTerm{typeof(fe)}) = Symbol(t.args_parsed[1])

function parse_fixedeffect(df::AbstractDataFrame, @nospecialize(formula::FormulaTerm))
    fes = FixedEffect[]
    ids = Symbol[]
    for term in eachterm(formula.rhs)
        result = parse_fixedeffect(df, term)
        if result != nothing
            push!(fes, result[1])
            push!(ids, result[2])
        end
    end
    if !isempty(fes)
        if any(fe.interaction isa Ones for fe in fes)
            formula = FormulaTerm(formula.lhs, tuple(InterceptTerm{false}(), (term for term in eachterm(formula.rhs) if (term != ConstantTerm(1)) & (term != InterceptTerm{true}()) & !has_fe(term))...))
        else
            formula = FormulaTerm(formula.lhs, Tuple(term for term in eachterm(formula.rhs) if !has_fe(term)))
        end
    end
    return fes, ids, formula
end

# Constructors from dataframe + Term
function parse_fixedeffect(df::AbstractDataFrame, t::AbstractTerm)
    if has_fe(t)
        st = fesymbol(t)
        return FixedEffect(df[!, st]), Symbol(:fe_, st)
    end
end

# Constructors from dataframe + InteractionTerm
function parse_fixedeffect(df::AbstractDataFrame, t::InteractionTerm)
    fes = (x for x in t.terms if has_fe(x))
    interactions = (x for x in t.terms if !has_fe(x))
    if !isempty(fes)
        # x1&x2 from (x1&x2)*id
        fe_names = [fesymbol(x) for x in fes]
        fe = FixedEffect(group((df[!, fe_name] for fe_name in fe_names)...); interaction = _multiply(df, Symbol.(interactions)))
        interactions = setdiff(Symbol.(terms(t)), fe_names)
        s = vcat(["fe_" * string(fe_name) for fe_name in fe_names], string.(interactions))
        return fe, Symbol(reduce((x1, x2) -> x1*"&"*x2, s))
    end
end

function _multiply(df, ss::Vector)
    if isempty(ss)
        out = Ones(size(df, 1))
    else
        out = ones(size(df, 1))
        for j in eachindex(ss)
            _multiply!(out, df[!, ss[j]])
        end
    end
    return out
end

function _multiply!(out, v)
    if v isa CategoricalVector
        throw("Fixed Effects cannot be interacted with Categorical Vector. Use fe(x)&fe(y)")
    end
    for i in eachindex(out)
        if v[i] === missing
            # may be missing when I remove singletons
            out[i] = 0.0
        else
            out[i] = out[i] * v[i]
        end
    end
end


using BenchmarkTools

include("helpers.jl")

const SUITE = BenchmarkGroup()

kwargs = Dict(:verbose=>false)

BenchmarkTools.DEFAULT_PARAMETERS.seconds = 60

n = 300
p = 3
covars, counts = bdata(n, p, 10^2)

for d in 10 .^ (1:4)
  SUITE[d] = BenchmarkGroup()
  covars, counts = bdata(n, p, d)
  for f in (DMR, HDMR)
    SUITE[d][string(f)] = BenchmarkGroup()
    for par = ("parallel", "serial")
      parallel = (par == "parallel")
      SUITE[d][string(f)][par] = BenchmarkGroup()
      if parallel
        clustertypes = ["local", "remote"]
      else
        clustertypes = ["local"]
      end
      for clus in clustertypes
        local_cluster = (clus == "local")
        # SUITE[d][f][par][clus] = @benchmarkable $(f)($covars, $counts; parallel=$parallel, local_cluster=$local_cluster, $kwargs...)
        SUITE[d][string(f)][par][clus] = @benchmarkable fit($f, $covars, $counts; parallel=$parallel, local_cluster=$local_cluster, $kwargs...)
      end
    end
  end
end


push!(LOAD_PATH, "/Users/czhang/Desktop/Projects/dw_/julialib/")
import DimmWitted
DimmWitted.set_libpath("/Users/czhang/Desktop/Projects/dw_/libdw_julia")

######################################
# The following piece of code creates a
# synthetic data set:
#    - Data type is Cdouble
#    - Modle type is Array{Cdouble}
#
nexp = 100000
nfeat = 100
examples = Array(Cdouble, nexp, nfeat+1)
for row = 1:nexp
	for col = 1:nfeat
		examples[row, col] = 1
	end
	if rand() > 0.8
		examples[row, nfeat+1] = 0
	else
		examples[row, nfeat+1] = 1
	end
end
model = Cdouble[0 for i = 1:nfeat]

######################################
# Define the loss function and gradient
# function for logistic regression
#
function loss(row::Array{Cdouble,1}, model::Array{Cdouble,1})
	@inbounds begin
		const label = row[length(row)]
		const nfeat = length(model)
		d = 0.0
		for i = 1:nfeat
			d = d + row[i]*model[i]
		end
	end
	return (-label * d + log(exp(d) + 1.0))
end

function grad(col::Array{Cdouble,1}, _colid::Cint, rows::Array{Array{Cdouble, 1}}, model::Array{Cdouble,1})
	colid = _colid + 1
	nfeat = length(model)
	nrows = length(rows)
	if colid > nfeat	# if it is the last column, then it is the labels, do nothing
		return 1.0
	end

	sum_term = 0.0
	pat_term = 0.0
	for ir = 1:length(rows)
		label = rows[ir][nfeat+1]
		d = 0.0
		for i = 1:nfeat
			d = d + rows[ir][i]*model[i]
		end
		sum_term = sum_term + label*rows[ir][colid]
		pat_term = pat_term + rows[ir][colid]*1.0/(1.0+exp(-d))
	end

	model[colid] = model[colid] - 0.00001* (-sum_term + pat_term)

	return 1.0
end

function avg(p_models::Ptr{Array{Cdouble,1}}, nrepl::Cint, irepl::Cint)
	models = pointer_to_array(p_models, convert(Int64, nrepl+1))
	julia_irepl = irepl + 1
	const nfeat = length(models[julia_irepl])
	for j = 1:nfeat
		d = 0.0
		for i = 1:nrepl
			d = d + models[i][j]
		end
		models[julia_irepl][j] = d/nrepl
	end
	return 1.0
end

######################################
# Create a DimmWitted object using data
# and model. You do not need to specify
# the type, they are infer'ed by the
# open() function, which is parametric.
#
dw = DimmWitted.open(examples, model,
                DimmWitted.MR_PERMACHINE,
                DimmWitted.DR_SHARDING,
                DimmWitted.AC_C2R)

######################################
# Register functions.
#
handle_loss = DimmWitted.register_row(dw, loss)
handle_grad = DimmWitted.register_c2r(dw, grad)
DimmWitted.register_model_avg(dw, handle_loss, avg, true)
DimmWitted.register_model_avg(dw, handle_grad, avg, true)

######################################
# Run 10 epoches.
#
for iepoch = 1:10
	rs = DimmWitted.exec(dw, handle_loss)
	println("LOSS: ", rs/nexp)
	println(sum(model))
	rs = DimmWitted.exec(dw, handle_grad)
end


module MNIST

using GZip, Colors

const Gray = Colors.Gray{Colors.N0f8}

const dir = joinpath(@__DIR__, "../../deps/mnist")

function load()
  mkpath(dir)
  cd(dir) do
    for file in ["train-images-idx3-ubyte",
                 "train-labels-idx1-ubyte",
                 "t10k-images-idx3-ubyte",
                 "t10k-labels-idx1-ubyte"]
      isfile(file) && continue
      info("Downloading MNIST dataset")
      download("https://cache.julialang.org/http://yann.lecun.com/exdb/mnist/$file.gz", "$file.gz")
      open(file, "w") do io
        write(io, GZip.open(read, "$file.gz"))
      end
    end
  end
end

const IMAGEOFFSET = 16
const LABELOFFSET = 8

const NROWS = 28
const NCOLS = 28

const TRAINIMAGES = joinpath(dir, "train-images-idx3-ubyte")
const TRAINLABELS = joinpath(dir, "train-labels-idx1-ubyte")
const TESTIMAGES = joinpath(dir, "t10k-images-idx3-ubyte")
const TESTLABELS = joinpath(dir, "t10k-labels-idx1-ubyte")

function imageheader(io::IO)
  magic_number = bswap(read(io, UInt32))
  total_items = bswap(read(io, UInt32))
  nrows = bswap(read(io, UInt32))
  ncols = bswap(read(io, UInt32))
  return magic_number, Int(total_items), Int(nrows), Int(ncols)
end

function labelheader(io::IO)
  magic_number = bswap(read(io, UInt32))
  total_items = bswap(read(io, UInt32))
  return magic_number, Int(total_items)
end

function rawimage(io::IO)
  img = Array{Gray}(NCOLS, NROWS)
  for i in 1:NCOLS, j in 1:NROWS
    img[i, j] = reinterpret(Colors.N0f8, read(io, UInt8))
  end
  return img
end

function rawimage(io::IO, index::Integer)
  seek(io, IMAGEOFFSET + NROWS * NCOLS * (index - 1))
  return rawimage(io)
end

rawlabel(io::IO) = Int(read(io, UInt8))

function rawlabel(io::IO, index::Integer)
  seek(io, LABELOFFSET + (index - 1))
  return rawlabel(io)
end

getfeatures(io::IO, index::Integer) = vec(getimage(io, index))

"""
    images()
    images(:test)

Load the MNIST images.

Each image is a 28×28 array of `Gray` colour values (see Colors.jl).

Returns the 60,000 training images by default; pass `:test` to retreive the
10,000 test images.
"""
function images(set = :train)
  load()
  io = IOBuffer(read(set == :train ? TRAINIMAGES : TESTIMAGES))
  _, N, nrows, ncols = imageheader(io)
  [rawimage(io) for _ in 1:N]
end

"""
    labels()
    labels(:test)

Load the labels corresponding to each of the images returned from `images()`.
Each label is a number from 0-9.

Returns the 60,000 training labels by default; pass `:test` to retreive the
10,000 test labels.
"""
function labels(set = :train)
  load()
  io = IOBuffer(read(set == :train ? TRAINLABELS : TESTLABELS))
  _, N = labelheader(io)
  [rawlabel(io) for _ = 1:N]
end

end # module


include("battlesnake.jl")
include("algorithms.jl")
include("minimax.jl")
include("agent.jl")
include("maxn.jl")
println("Import all done")

# Turn 7 https://play.battlesnake.com/g/73639a35-ade8-42ee-9584-72f2a44266ef/
snake1 = Snake(
    [
        Position(5,6),
        Position(5,5),
        Position(5,4),
        Position(4,4)
    ],
    false,
    100
    )
snake2 = Snake(
    [
        Position(6,5),
        Position(6,4),
        Position(6,3),
        Position(7,3),
    ],
    false,
    100
    )
snake3 = Snake(
    [
        Position(1,1),
        Position(1,2),
        Position(1,3),
        Position(1,4),
    ],
    false,
    100
    )

food = Food([Position(6,6)])

# Weird head2head
# Looks like "working as intended" since we don't move simultaneously in minimax
# Why is this no problem in python though?
# Would't the opponent win in the next step?
# Turn 11 https://play.battlesnake.com/g/1a6794ac-9da9-4e32-a74b-b246797dbb16/
snake1 = Snake(
    [
        Position(7,4),
        Position(8,4),
        Position(8,3),
        Position(7,3)
    ],
    false,
    49
    )
snake2 = Snake(
    [
        Position(6,5),
        Position(6,6),
        Position(6,7),
        Position(5,7),
        Position(4,7),
    ],
    false,
    100
    )
snake4 = Snake(
        [
            Position(9,5),
            Position(9,6)
        ],
        false,
        100
    )
food = Food([Position(11,5), Position(9,2)])
### Play best move
height, width = 11, 11
snakes = [snake1, snake2, snake3, snake4]
board = Boardstate(snakes, food, height, width)
print_board(board)
move(board, 1, 0.4)
# println(a_star_search(Position(1,1), Position(5,5), board))
# flood = floodfill([snake1, snake2], width, height, 1)
# println(flood)
# for x in 1:10
#     for i in 1:length(snakes)
#         a = iterative_deepening_maxn(board, 400.0, height, width, 3)
#         move_snake!(snakes[3], food, a)
#         # winner = check_winner(snakes, 1, height, width)
#         # println("Winnner is $winner")
#     end
#     print_board(board)
# end
# print("-----------------")
# print_board(board)
# a = iterative_deepening(board, 400.0, height, width, 2)
# println(a)
# move_snake!(snake2, food, a)
# winner = check_winner([snake1, snake2], 2, height, width)
# println("Winnner is $winner")
# # print("-----------------")
# print_board(board)


function FieldStopOnEvent()
    return FieldStopOnEvent(())
end

function event_occurred(obj::FieldStopOnEvent, arg0::FieldSpacecraftState, arg1::FieldEventDetector, arg2::jboolean)
    return jcall(obj, "eventOccurred", Action, (FieldSpacecraftState, FieldEventDetector, jboolean), arg0, arg1, arg2)
end


# Reporting.jl

module Reporting

using Dates

export aging

import ..AppliAR: retrieve_unpaid_invoices, retrieve_paid_invoices

import ..AppliAR: id, name, date, meta, header, body, students, price_per_student, vat_perc

struct Aging
    id_inv::String
    csm::String
    inv_date::Date
    amount::Float64
    days::Day
end

aging(path_file_unpaid, path_file_paid) = begin
    unpaid_invoices = retrieve_unpaid_invoices(;path=path_file_unpaid)
    paid_invoices = retrieve_paid_invoices(;path=path_file_paid)

    paid = [id(x) for x in paid_invoices]
    unpaid = filter(x -> id(x) ∉ paid, unpaid_invoices)

    list = []
    for invoice in unpaid
        id_inv = id(invoice)
        csm = name(header(invoice))
        inv_date = Date(date(meta(invoice)))
        b = body(invoice)
        amount = (price_per_student(b) * length(students(b))) * (1 + vat_perc(b))
        days = Dates.today() - inv_date
        aging_item = Aging(id_inv, csm, inv_date, amount, days)
        push!(list, aging_item)
    end
    return list
end # aging

end # module


module CodedComputing

using LinearAlgebra, SparseArrays, HDF5, DataFrames
using Distributions
using DataStructures

include("Linalg.jl")
include("concurrentcircularbuffer.jl")
include("mul.jl")
include("partitioning.jl")
include("Datasets.jl")
include("profiling.jl")
include("eventdriven.jl")
include("loadbalancing.jl")
include("treegradient.jl")

end


using .CuYao
using .CuYao: CUDA
using .CUDA: CuArray, @linearidx, GPUArrays
using LinearAlgebra

export togpu

CUDA.ones(::Type{Tropical{T}}, dims...) where T = fill!(CuArray{Tropical{T}}(undef, dims...), one(Tropical{T}))
CUDA.zeros(::Type{Tropical{T}}, dims...) where T = fill!(CuArray{Tropical{T}}(undef, dims...), zero(Tropical{T}))
CUDA.ones(::Type{CountingTropical{T}}, dims...) where T = fill!(CuArray{CountingTropical{T}}(undef, dims...), one(CountingTropical{T}))
CUDA.zeros(::Type{CountingTropical{T}}, dims...) where T = fill!(CuArray{CountingTropical{T}}(undef, dims...), zero(CountingTropical{T}))

function _init_reg(::Type{T}, L::Int, usecuda::Val{:true}) where T
    ArrayReg(CUDA.ones(T, 1<<L))
end

function _init_reg(::Type{T}, lt::MaskedSquareLattice, ::Val{:true}) where T
    nbit = size(lt, 2) + 2
    state = CUDA.ones(T, 1<<nbit)
    ArrayReg(state)
end

function togpu(tn::TensorNetwork)
    TensorNetwork(togpu.(tn.tensors); metas=tn.metas)
end

function togpu(t::LabeledTensor)
    LabeledTensor(CuArray(t.array), t.labels)
end


################################################################################
# ptree.jl
#
# Code implementing a persistent hash array mapped trie tyoe on which the
# persistent collection library of Air is built. This type is not intended for
# use outside of the library.
#
# @author Noah C. Benson
#
# MIT License
# Copyright (c) 2020 Noah C. Benson

# ==============================================================================
# Private
#
# This section contains private constants and function definitions intendef for
# use only in this library or occasionally within the library, but not outside
# of the current module.

# # Constants ==================================================================
#   This is the type of the unsigned integer used for the leaves.
const HASH_T              = typejoin(typeof(hash(nothing)),
                                     typeof(objectid(nothing)))
const HASH_ZERO           = HASH_T(0x00)
const HASH_ONE            = HASH_T(0x01)
const HASH_MAX            = ~HASH_ZERO
# Number of bits in the total addressable hash-space of a PTree.
const HASH_BITCOUNT       = sizeof(HASH_T) * 8
const HASH_TOPBIT         = HASH_T(0x01) << (HASH_BITCOUNT - 1)
#   We use a constant shift of 5 throughout except at the root node (which is
#   Due to the implemeentation details, the twig shifts *MUST* be >= 3 and all
#   shifts must be <= 7.
const PTREE_NODE_SHIFT    = 6
const PTREE_TWIG_SHIFT    = 6
const PTREE_ROOT_SHIFT    = (let tmp = rem(HASH_BITCOUNT - PTREE_TWIG_SHIFT,
                                           PTREE_NODE_SHIFT)
                                 (tmp == 0 ? PTREE_NODE_SHIFT : tmp)
                             end)
const PTREE_ROOT_FIRSTBIT = HASH_BITCOUNT - PTREE_ROOT_SHIFT
const PTREE_ROOT_NCHILD   = (1 << PTREE_ROOT_SHIFT)
const PTREE_NODE_NCHILD   = (1 << PTREE_NODE_SHIFT)
const PTREE_TWIG_NCHILD   = (1 << PTREE_TWIG_SHIFT)
const PTREE_BITS_BITCOUNT = max(PTREE_ROOT_NCHILD,
                                PTREE_NODE_NCHILD,
                                PTREE_TWIG_NCHILD)
const PTREE_BITS_T        = (if     (PTREE_BITS_BITCOUNT <= 8)  UInt8
                             elseif (PTREE_BITS_BITCOUNT <= 16) UInt16
                             elseif (PTREE_BITS_BITCOUNT <= 32) UInt32
                             elseif (PTREE_BITS_BITCOUNT <= 64) UInt64
                             else                               UInt128
                             end)
const BITS_ZERO           = PTREE_BITS_T(0x0)
const BITS_ONE            = PTREE_BITS_T(0x1)
const PTREE_NODE_BITS     = HASH_BITCOUNT - PTREE_ROOT_SHIFT - PTREE_TWIG_SHIFT
const PTREE_NODE_LEVELS   = div(PTREE_NODE_BITS, PTREE_NODE_SHIFT)
const PTREE_LEVELS        = PTREE_NODE_LEVELS + 2
const PTREE_ROOT_DEPTH    = 0
const PTREE_TWIG_DEPTH    = PTREE_ROOT_DEPTH + PTREE_NODE_LEVELS + 1
const PTREE_LEAF_DEPTH    = PTREE_TWIG_DEPTH + 1
# Functions for making masks
"""
    lowmask(bitno, T)
    lowmask(bitno)

Yields a mask of (unsigned integer) type T with all bits above the given bit
number set to false and all bits below that number set to true. The bit itself
is set to false. Bits are indexed starting at 0.

The default value of T is the PTree hash type (HASH_T).

lowmask(bitno) is equal to ~highmask(bitno).
"""
lowmask(bitno::K, ::Type{T}) where {T<:Unsigned,K<:Integer} =
    ((T(0x01) << bitno) - 0x01)
lowmask(bitno::K) where {K <: Integer} = lowmask(bitno, HASH_T)
"""
    highmask(bitno, T)
    highmask(bitno)

Yields a mask of (unsigned integer) type T with all bits above the given bit
number set to true and all bits below that number set to false. The bit itself
is set to true. Bits are indexed starting at 0.

The default value of T is the PTree hash type (HASH_T).

highmask(bitno) is equal to ~lowmask(bitno).
"""
highmask(bitno::K, ::Type{T}) where {T<:Unsigned,K<:Integer} =
    ~lowmask(bitno, T)
highmask(bitno::K) where {K <: Integer} = highmask(bitno, HASH_T)
# The mask of the bits that represent the depth of the node-id in a PTree.
const PTREE_DEPTH_MASK = lowmask(PTREE_TWIG_SHIFT)
# The pair type for the dict interface.
const HASHPAIR_T{T} = Pair{HASH_T, T} where {T}

# # Functions
"""
    ptree_depth(nodeaddr)

Yields the depth of the node with the given node address. This depth is in the
theoretical complete tree, not in the reified tree represented with memory.
"""
ptree_depth(nodeid::HASH_T) = nodeid & PTREE_DEPTH_MASK
"""
    depth_to_bitshift(depth)

Yields the tuple (B0,S) of the first bit index and the shift for the given
depth.
"""
depth_to_bitshift(depth::Integer) = depth_to_bitshift(Int(depth))
depth_to_bitshift(depth::Int) = begin
    if depth == PTREE_TWIG_DEPTH
        return (0, PTREE_TWIG_SHIFT)
    elseif depth == 0
        return (PTREE_ROOT_FIRSTBIT, PTREE_ROOT_SHIFT)
    else
        return (PTREE_ROOT_FIRSTBIT - depth*PTREE_NODE_SHIFT,
                PTREE_NODE_SHIFT)
    end
end
"""
    ptree_bitshift(nodeid)

Yields the bitshift for the given ptree node id.
"""
ptree_bitshift(nodeid::HASH_T) = depth_to_bitshift(ptree_depth(nodeid))
"""
    ptree_shift(nodeid)

Yields the shift for the given ptree node id.
"""
ptree_shift(nodeid::HASH_T) = ptree_bitshift(nodeid)[2]
"""
    ptree_firstbit(nodeid)

Yields the first-bit for the given ptree node id.
"""
ptree_firstbit(nodeid::HASH_T) = ptree_bitshift(nodeid)[1]
"""
    ptree_minleaf(nodeid)

Yields the minimum child leaf index associated with the given nodeid.
"""
ptree_minleaf(nodeid::HASH_T) = nodeid & ~PTREE_DEPTH_MASK
"""
    ptree_maxleaf(nodeid)

Yields the maximum child leaf index assiciated with the given nodeid.
"""
ptree_maxleaf(nodeid::HASH_T) = begin
    (b0,s) = ptree_bitshift(nodeid)
    mask = (HASH_ONE << (b0 + s)) - HASH_ONE
    return nodeid | mask
end
"""
    ptree_minmaxleaf(nodeid)

Yields the (min, max) child leaf index assiciated with the given nodeid.
"""
ptree_minmaxleaf(nodeid::HASH_T) = begin
    mn = ptree_minleaf(nodeid)
    (bit0,shift) = ptree_bitshift(nodeid)
    mask = (HASH_ONE << (bit0+shift)) - HASH_ONE
    return (nodeid & ~mask, nodeid | mask)
end
"""
    ptree_id(minleaf, depth)

Yields the node-id for the node whose minimum leaf and depth are given.
"""
ptree_id(minleaf::Integer,depth::Integer) =
    ptree_id(HASH_T(minleaf), HASH_T(depth))
ptree_id(minleaf::HASH_T, depth::HASH_T) = minleaf | depth
"""
    ptree_parentid(nodeid0)

Yields the node-id of the parent of the given node. Note that node 0 (the tree's
theoretical root) has no parent. If given a node id of 0, this function will
return an arbitrary large number.
"""
ptree_parentid(nodeid0::HASH_T) = begin
    d = ptree_depth(nodeid0) - 0x1
    node = nodeid0 & ~((-HASH_T(1)) >> (d*PTREE_NODE_SHIFT))
    return node | d
end
"""
    ptree_isbeneath(nodeid, leafid)

Yields true if the given leafid can be found beneath the given node-id.
"""
ptree_isbeneath(nodeid::HASH_T, leafid::HASH_T) = begin
    (mn,mx) = ptree_minmaxleaf(nodeid)
    return leafid <= mx && leafid >= mn
end
"""
    ptree_highbitdiff(id1, id2)

Yields the highest bit that is different between id1 and id2.
"""
ptree_highbitdiff(id1::HASH_T, id2::HASH_T) =
    HASH_BITCOUNT - leading_zeros(xor(id1, id2)) - 1

# ==============================================================================
# Structures / Types
#
# Definitions of structures and types that should come before the various
# methods and constructors that operate on them.
# #PTree =======================================================================
"""
    PTree{T}

PTree is a persistent tree/array hybrid that maps arbitrary unsigned integers
(type HASH_T) to values. It can efficiently be used for arrays or for
hash-maps, and supports efficient lookup, association, and dissociation.

All fields of a PTree should be considered strictly private, as modification
of the fields may result in a kernel crash. The properties of a PTree are both
immutable and safe to inpect.
"""
struct PTree{T} <: AbstractDict{HASH_T, T}
    id::HASH_T
    bits::PTREE_BITS_T
    numel::Int
    cells::Union{Nothing, Vector{T}, Vector{PTree{T}}}
end
const PTREE_NODE_CELL_T{T} = Vector{PTree{T}} where {T}
const PTREE_TWIG_CELL_T{T} = Vector{T} where {T}
const PTREE_EMPTY_CELL_T = Nothing
const PTREE_CELLS_T{T} = Union{PTREE_EMPTY_CELL_T,
                               PTREE_NODE_CELL_T{T},
                               PTREE_TWIG_CELL_T{T}} where {T}
# We will occasionally want to do something across all the children of a node;
# we can exploit the bits data to speed this up in sparse cases. The code for
# doing that is in these two macros.
macro _ptree_forcells(tree, var::Symbol, expr::Expr)
    b = gensym("bits")
    return quote
        let $b = $(tree).bits, $var
            while true
                $var = trailing_zeros($b)
                if $var == PTREE_BITS_BITCOUNT
                    break
                else
                    $expr
                end
                $b &= ~($(HASH_T(0x01)) << $var)
            end
        end
    end |> esc
end
macro _ptree_forcells_from(tree, var::Symbol, from, expr::Expr)
    b = gensym("bits")
    f = gensym("from")
    return quote
        let $b = $(tree).bits & highmask($from), $var
            while true
                $var = trailing_zeros($b)
                if $var == PTREE_BITS_BITCOUNT
                    break
                else
                    $expr
                end
                b &= ~($(HASH_T(0x01)) << $var)
            end
        end
    end |> esc
end
"""
    ptree_cellindex(ptree, leafid)
    ptree_cellindex(nodeid, bits, leafid)
Yields a tuple (present, bitindex, cellindex) in which [1] present is a boolean
indicating whether a ptree containing the given leafid or the leaf itself is a
child of the given ptree; [2] bitindex is the index into ptree's bits integer
for the given leaf, and cellindex is the inndex innto the ptree's cell array
where that child is or would be found. If the leafid is outside of the given
ptree (i.e., it cannot exist beneath this ptree) then the cellindex returned
returned is 0, but the bitindex will still match the appropriate shift for the
ptree's depth.
"""
ptree_cellindex(id::HASH_T, bits::PTREE_BITS_T, leafid::HASH_T) = begin
    # Check that the leaf is below this leaf.
    (bit0, shift) = ptree_bitshift(id)
    # Grab the index out of the leaf id.
    idx = (leafid >> bit0) & lowmask(shift)
    ptree_isbeneath(id, leafid) || return (false, idx, 0)
    # See if the bit is set?
    flag = HASH_ONE << idx
    chno = 1 + count_ones(bits & lowmask(idx))
    return (bits & flag == flag, idx, chno)
end
ptree_cellindex(u::PTree{T}, leafid::HASH_T) where {T} =
    ptree_cellindex(getfield(u, :id), getfield(u, :bits), leafid)
"""
    ptree_cellindex!(ptree, leafid)
    ptree_cellindex!(nodeid, bits, leafid)
Yields the index into the ptree's cell vector of the child containing the leaf
with the given id. If the ptree does not contain the given leafid because the
appropriate bit is not set then 0 is returned. However, unlike the function
ptree_cellindex(), this function does not check whether or not the given leafid
is in the set of possible children of the tree.
"""
ptree_cellindex!(id::HASH_T, bits::PTREE_BITS_T, leafid::HASH_T) = begin
    # Check that the leaf is below this leaf.
    (bit0, shift) = ptree_bitshift(id)
    # Grab the index out of the leaf id.
    idx = (leafid >> bit0) & lowmask(shift)
    # See if the bit is set?
    flag = HASH_ONE << idx
    chno = 1 + count_ones(bits & lowmask(idx))
    return (bits & flag == flag, idx, chno)
end
ptree_cellindex!(ptree::PTree{T}, leafid::HASH_T) where {T} =
    ptree_cellindex!(getfield(ptree, :id), getfield(ptree, :bits), leafid)
"""
    ptree_cellkey(ptree, childidx)
Yields the leafid (a HASH_T value) of the key that goes with the particular
child index that is given. This only works correctly for twig nodes.
"""
ptree_cellkey(id::HASH_T, k::HASH_T) = begin
    mn = ptree_minleaf(id)
    return mn | k
end
ptree_cellkey(id::HASH_T, k::II) where {II<:Integer} =
    ptree_cellkey(id, (HASH_T(k)))
ptree_cellkey(u::PTree{T}, k::II) where {T,II<:Integer} =
    ptree_cellkey(getfield(u, :id), HASH_T(k))
ptree_cellkey(u::PTree{T}, k::HASH_T) where {T} =
    ptree_cellkey(getfield(u, :id), k)

# # Constructors ===============================================================
#   First, constructors for making empty Trees.
PTree{T}() where {T} = PTree{T}(0x0, 0x0, 0, nothing)
#   Now for copying trees of the same type.
PTree{T}(u::PTree{T}) where {T}   = u
PTree{T}(u::PTree{U}) where {T,U} = begin
    (id, bits, n, cells) = (getfield(u, :id),    getfield(u, :bits),
                            getfield(u, :numel), getfield(u, :cells))
    d = ptree_depth(id)
    if d == PTREE_TWIG_DEPTH
        cells::Vector{U}
        return PTree{T}(id, bits, n, convert(Vector{T}, cells))
    else
        cells::Vector{PTree{U}}
        return PTree{T}(id, bits, n, PTree{T}[PTree{T}(c) for c in cells])
    end
end
PTree{T}(d::AbstractDict{HASH_T,S}) where {T,S} = begin
    u = PTree{T}()
    for (k,v) in d
        u = setindex(u, v, k)
    end
    return u
end
#   We want to be able to construct the twig for a pair:
PTree{T}(kv::Pair{HASH_T,V}) where {T,V} = begin
    id = ptree_id(kv[1] & ~PTREE_DEPTH_MASK, PTREE_TWIG_DEPTH)
    bits = PTREE_BITS_T(0x1) << (kv[1] & PTREE_DEPTH_MASK)
    return PTree{T}(id, bits, 1, T[kv[2]])
end
PTree(kv::Pair{HASH_T,V}) where {V} = PTree{V}(kv)
#   We also construct from arrays; they are considered to be 0-indexed when used
#   with trees.
PTree{T}(u::AbstractArray{S,N}) where {T,S,N} = begin
    depth = PTREE_TWIG_DEPTH
    height = 1
    n = length(u)
    (n == 0) && return PTree{T}()
    (n == 1) && return PTree{T}(ptree_id(0x0, depth), 0x1, 1, T[u[1]])
    # Start by making all the twigs
    maxidx = fld(n + PTREE_TWIG_NCHILD - 1, PTREE_TWIG_NCHILD)
    bits = ~PTREE_BITS_T(0x0)
    nodes = Vector{PTree{T}}(undef, maxidx)
    for ii in 1:maxidx
        kk0 = (ii - 1)*PTREE_TWIG_NCHILD
        nn = min(PTREE_TWIG_NCHILD, n - kk0)
        nid = ptree_id(PTREE_TWIG_NCHILD * (ii-1), depth)
        ch = Vector{T}(undef, PTREE_TWIG_NCHILD)
        if nn == PTREE_TWIG_NCHILD
            for kk in 1:PTREE_TWIG_NCHILD
                @inbounds ch[kk] = u[kk + kk0]
            end
            nodes[ii] = PTree{T}(nid, bits, Int(PTREE_TWIG_NCHILD), ch)
        else
            for kk in 1:nn
                @inbounds ch[kk] = u[kk + kk0]
            end
            nodes[ii] = PTree{T}(nid, lowmask(nn, PTREE_BITS_T), Int(nn), ch)
        end
    end
    # Then iterate over the nodes doing the same thing.
    while true
        # Even if we are given an array with the max number of elements, we must
        # eventually have a list of only one node; that's the node we return.
        (maxidx == 1) && return nodes[1]
        # Increment
        depth -= 1
        height += 1
        n = length(nodes)
        oldnodes = nodes
        maxidx = div(n + PTREE_NODE_NCHILD - 1, PTREE_NODE_NCHILD)
        nodes = Vector{PTree{T}}(undef, maxidx)
        hnp = PTREE_NODE_NCHILD^(height-1) # number per node at the prev level
        hn  = Int(PTREE_NODE_NCHILD*hnp) # number per node at this level
        bits = ~PTREE_BITS_T(0x0)
        for ii in 1:maxidx
            kk0 = (ii - 1)*PTREE_NODE_NCHILD
            nn = min(PTREE_NODE_NCHILD, n - kk0)
            nid = ptree_id(hn * (ii-1), depth)
            ch = Vector{PTree{T}}(undef, nn)
            if nn == PTREE_NODE_NCHILD
                # split this out for faster looping
                for kk in 1:PTREE_NODE_NCHILD
                    @inbounds ch[kk] = (@inbounds oldnodes[kk + kk0])
                end
                nodes[ii] = PTree{T}(nid, bits, hn - hnp + getfield(ch[end], :numel), ch)
            else
                for kk in 1:nn
                    @inbounds ch[kk] = (@inbounds oldnodes[kk + kk0])
                end
                nodes[ii] = PTree{T}(nid, (PTREE_BITS_T(0x1) << nn) - 0x1,
                                     (nn - 1)*hnp + getfield(ch[end], :numel), ch)
            end
        end
    end
end
# Finally, we want to duplicate all of these constructors using the convert
Base.convert(::Type{PTree{T}}, u::PTree{T}) where {T} = u
Base.convert(::Type{PTree{T}}, u::PTree{U}) where {T,U} = PTree{T}(u)
Base.convert(::Type{PTree{T}}, u::AbstractDict{HASH_T,U}) where {T,U} =
    PTree{T}(u)
Base.convert(::Type{PTree{T}}, u::AbstractArray{1,U}) where {T,U} =
    PTree{T}(u)

# ==============================================================================
# Methods
#
# Definitions of methods of the above type.
Base.empty(u::PTree{T}) where {T} = PTree{T}(r)
Base.empty(u::PTree{T}, S::Type) where {T} = PTree{S}()
Base.isempty(u::PTree{T}) where {T} = (getfield(u, :numel) == 0)
Base.length(u::PTree{T}) where {T} = getfield(u, :numel)
Base.isequal(t::PTree{T}, s::PTree{S}) where {T,S} = begin
    bits = getfield(t, :bits)
    (bits == getfiield(s, :bits)) || return false
    (getfield(t, :id) == getfield(s, :id)) || return false
    tcells = getfield(t, :cells)
    scells = getfield(s, :cells)
    for k in 1:count_ones(bits)
        t = @inbounds tcells[k]
        s = @inbounds scells[k]
        isequal(t, s) || return false
    end
    return true
end
Base.copy(t::PTree{T}) where {T} = t
Base.propertynames(u::PTree{T}, private::Bool=false) where {T} =
    (:leafcount,
     :children,
     :minleaf,
     :maxleaf,
     :depth,
     :address)
Base.getproperty(u::PTree{T}, symbol::Symbol) where {T} = begin
    if symbol == :leafcount
        return getfield(u, :numel)
    elseif symbol == :children
        id = getfield(u, :id)
        (getfield(u, :numel) == 0) && return ()
        if ptree_depth(id) == PTREE_TWIG_DEPTH
            return tuple(u...)
        else
            cells = getfield(u, :cells)::Vector{PTree{T}}
            return tuple(cells...)
        end
    elseif symbol == :minleaf
        return ptree_minleaf(getfield(u, :id))
    elseif symbol == :maxleaf
        return ptree_maxleaf(getfield(u, :id))
    elseif symbol == :depth
        return ptree_depth(getfield(u, :id))
    elseif symbol == :address
        return getfield(u, :id)
    else
        throw(ArgumentError("Type PTree has no property $symbol"))
    end
end
getpair(u::PTree{T}, k::HASH_T) where {T} = begin
    # Start by making sure that the address spaces are compatible:
    # we want to make sure (once) that k is beneath u.
    id = getfield(u, :id)
    ptree_isbeneath(id, k) || return missing
    # Also, if we are empty, we need to return right away.
    (getfield(u, :numel) == 0) && return missing
    # Okay, let's try to find the child.
    while true
        # Is it in this Tree?
        (inq, bitidx, idx) = ptree_cellindex!(u, k)
        inq || return missing
        # Are we a twig?
        #if isa(getfield(u, :cells), Vector{T})
        if ptree_depth(id) == PTREE_TWIG_DEPTH
            qq = getfield(u, :cells)::Vector{T}
            return k => (@inbounds qq[idx])
        else
            q = getfield(u, :cells)::Vector{PTree{T}}
            u = (@inbounds q[idx])
        end
        id = getfield(u, :id)
    end
end
Base.get(u::PTree{T}, k::HASH_T, df) where {T} = begin
    # Start by making sure that the address spaces are compatible:
    # we want to make sure (once) that k is beneath u.
    id = getfield(u, :id)
    ptree_isbeneath(id, k) || return df
    # Also, if we are empty, we need to return right away.
    (getfield(u, :numel) == 0) && return df
    # Okay, let's try to find the child.
    while true
        # Is it in this Tree?
        (inq, bitidx, idx) = ptree_cellindex(u, k)
        inq || return df
        # Are we a twig?
        #if isa(getfield(u, :cells), Vector{T})
        if ptree_depth(id) == PTREE_TWIG_DEPTH
            qq = getfield(u, :cells)::Vector{T}
            return (@inbounds qq[idx])
        else
            q = getfield(u, :cells)::Vector{PTree{T}}
            u = (@inbounds q[idx])
        end
        id = getfield(u, :id)
    end
end
Base.in(kv::Pair{HASH_T,T}, u::PTree{T}, f::F) where {T,F<:Function} = begin
    id = getfield(u, :id)
    ptree_isbeneath(id, k) || return df
    # Also, if we are empty, we need to return right away.
    (getfield(u, :numel) == 0) && return df
    # Okay, let's try to find the child.
    while true
        # Is it in this Tree?
        (inq, bitidx, idx) = ptree_cellindex!(u, k)
        inq || return df
        # Are we a twig?
        #if isa(getfield(u, :cells), Vector{T})
        if ptree_depth(id) == PTREE_TWIG_DEPTH
            qq = getfield(u, :cells)::Vector{T}
            return f(kv[2], (@inbounds qq[idx])) == true
        else
            q = getfield(u, :cells)::Vector{PTree{T}}
            u = (@inbounds q[idx])
        end
        id = getfield(u, :id)
    end
end
Base.in(kv::Pair{HASH_T,T}, u::PTree{T}) where {T} = in(kv, u, (===))

Base.iterate(u::PTree{T}) where {T} = begin
    (getfield(u, :numel) == 0) && return nothing
    id = getfield(u, :id)
    d = ptree_depth(id)
    while d < PTREE_TWIG_DEPTH
        cells = getfield(u, :cells)::Vector{PTree{T}}
        u = cells[1]
        id = getfield(u, :id)
        d = ptree_depth(id)
    end
    leaves = getfield(u, :cells)::Vector{T}
    bits = getfield(u, :bits)
    bitno = trailing_zeros(bits)
    k = ptree_cellkey(id, bitno)
    v = leaves[1] # @inbounds leaves[1]
    return (Pair{HASH_T,T}(k,v), k)
end
iterkeys(u::PTree{T}) where {T} = begin
    (getfield(u, :numel) == 0) && return nothing
    id = getfield(u, :id)
    d = ptree_depth(id)
    while d < PTREE_TWIG_DEPTH
        cells = getfield(u, :cells)::Vector{PTree{T}}
        u = cells[1]
        id = getfield(u, :id)
        d = ptree_depth(id)
    end
    leaves = getfield(u, :cells)::Vector{T}
    bits = getfield(u, :bits)
    bitno = trailing_zeros(bits)
    k = ptree_cellkey(id, bitno)
    v = leaves[1] # @inounds leaves[1]
    return (k, k)
end
itervals(u::PTree{T}) where {T} = begin
    (getfield(u, :numel) == 0) && return nothing
    id = getfield(u, :id)
    d = ptree_depth(id)
    while d < PTREE_TWIG_DEPTH
        cells = getfield(u, :cells)::Vector{PTree{T}}
        u = cells[1]
        id = getfield(u, :id)
        d = ptree_depth(id)
    end
    leaves = getfield(u, :cells)::Vector{T}
    bits = getfield(u, :bits)
    bitno = trailing_zeros(bits)
    k = ptree_cellkey(id, bitno)
    v = leaves[1] # @inounds leaves[1]
    return (v, k)
end

macro _ptree_iterate_gencode(rtype::Symbol)
    cids = [gensym("cid") for _ in 1:PTREE_LEVELS-1]
    cels = [gensym("cel") for _ in 1:PTREE_LEVELS-1]
    bits = [gensym("bit") for _ in 1:PTREE_LEVELS-1]
    if rtype == :keys
        rexpr = :(k)
        fn = :iterkeys
    elseif rtype == :vals
        rexpr = :(v)
        fn = :itervals
    else
        rexpr = :(Pair{HASH_T,T}(k,v))
        fn = :(Base.iterate)
    end
    twigexpr = quote
        # We must be a twig at this point.
        leaves = getfield(u, :cells)::Vector{T}
        id = getfield(u, :id)
        bits = getfield(u, :bits)
        (inq, bitno, idx) = ptree_cellindex(id, bits, k0)
        mask = (BITS_ONE << (bitno + 1)) - BITS_ONE
        nextbitno = trailing_zeros(bits & ~mask)
        if nextbitno < PTREE_BITS_BITCOUNT
            k = ptree_cellkey(id, nextbitno)
            idx = count_ones(bits & mask) + 1
            v = (@inbounds leaves[idx])
            return ($rexpr, k)
        end
        # If we reach here, we did not find a next leaf.
    end
    # At the lowest level, the try and twig expressions are the same.
    tryexpr = twigexpr
    for ii in PTREE_LEVELS-1:-1:1
        cel = cels[ii]
        cid = cids[ii]
        bit = bits[ii]
        tryexpr = quote
            if d == PTREE_TWIG_DEPTH
                $twigexpr
            else
                $cel = getfield(u, :cells)::Vector{PTree{T}}
                $bit = getfield(u, :bits)
                (inq, bitno, $cid) = ptree_cellindex(id, $bit, k0)
                u = $cel[$cid]
                id = getfield(u, :id)
                d = ptree_depth(id)
                $tryexpr
                # If we reach this point, we haven't found a next pair yet.
                n = count_ones($bit)
                if $cid < n
                    u = $cel[$cid + 1]
                    break
                end
            end
        end
    end
    # Put these together into the functions.
    return quote
        $fn(u::PTree{T}, k0::HASH_T) where {T} = begin
            id = getfield(u, :id)
            d = ptree_depth(id)
            while true
                $tryexpr
                # If we reach the end of the loop, there's nothing past k0.
                return nothing
            end
            # Upon finding a cell that hasn't been iterated, the code breaks
            # from the abbove loop and brings us here.
            id = getfield(u, :id)
            d = ptree_depth(id)
            while d < PTREE_TWIG_DEPTH
                cells = getfield(u, :cells)::Vector{PTree{T}}
                u = (@inbounds cells[1])
                id = getfield(u, :id)
                d = ptree_depth(id)
            end
            leaves = getfield(u, :cells)::Vector{T}
            bits = getfield(u, :bits)
            nextbitno = trailing_zeros(bits)
            k = ptree_cellkey(id, nextbitno)
            v = (@inbounds leaves[1])
            return ($rexpr, k)
        end
    end |> esc
end
# Run the macro to generate the iteration functions:
(@_ptree_iterate_gencode keys)
(@_ptree_iterate_gencode vals)
(@_ptree_iterate_gencode pairs)

setindex(u::PTree{T}, v::V, k::HASH_T) where {T,V} = begin
    # First of all, if this node is empty, we just return a new node
    (getfield(u, :numel) == 0) && return PTree{T}(k => v)
    # Next, if k is a child of this node, we return a new ancestor node that
    # contains us both and is the highest possible ancestor that contains us
    # both.
    id = getfield(u, :id)
    if !ptree_isbeneath(id, k)
        # Make the new twig that will hold k => v.
        kv = PTree{T}(k => v)
        # What's the lowest bit at which the mismatch in address starts?
        topmiss = ptree_highbitdiff(id, k)
        if topmiss <= HASH_BITCOUNT - PTREE_ROOT_SHIFT
            nodelev = div(topmiss - PTREE_TWIG_SHIFT, PTREE_NODE_SHIFT)
            bit0 = PTREE_TWIG_SHIFT + nodelev * PTREE_NODE_SHIFT
            shift = PTREE_NODE_SHIFT
            newdepth = PTREE_LEVELS - 2 - nodelev
        else
            newdepth = 0
            bit0 = HASH_BITCOUNT - PTREE_ROOT_SHIFT
            shift = PTREE_ROOT_SHIFT
        end
        newminleaf = id & ~((HASH_ONE << bit0) - HASH_ONE)
        newid = ptree_id(newminleaf, newdepth)
        newmask = (HASH_ONE << shift) - HASH_ONE
        ii_k = newmask & (id >> bit0)
        ii_u = newmask & (k >> bit0)
        newbits = (BITS_ONE << ii_k) | (BITS_ONE << ii_u)
        newchs = ii_u < ii_k ? PTree{T}[kv, u] : PTree{T}[u, kv]
        return PTree{T}(newid, newbits, 1 + getfield(u, :numel), newchs)
    end
    # Otherwise, this key goes beneath this node, so find out where.
    bits = getfield(u, :bits)
    numel = getfield(u, :numel)
    (bit0, shift) = ptree_bitshift(id)
    # Grab the index out of the leaf id.
    idx = (k >> bit0) & lowmask(shift)
    # We'll want to see if this bit is set and then insert/edit at this index:
    flag = BITS_ONE << idx
    idx = 1 + count_ones(bits & lowmask(idx))
    # However, what we do next depends on whether we are in a twig node or not.
    if ptree_depth(id) == PTREE_TWIG_DEPTH
        cells = getfield(u, :cells)::Vector{T}
        if bits & flag == flag
            # We are replacing an item
            return PTree{T}(id, bits, numel, setindex(cells, v, idx))
        else
            # If we get here, we are adding a new node, which means we are
            # inserting a leaf.
            return PTree{T}(id, bits | flag, numel + 1, insert(cells, idx, v))
        end
    else
        cells = getfield(u, :cells)::Vector{PTree{T}}
        if bits & flag == flag
            # We pass control along then return a copy of ourselves on the
            # way back up the control stack.
            oldc = cells[idx]
            oldn = getfield(oldc, :numel)
            newc = setindex(cells[idx], v, k)
            (newc === cells[idx]) && return u
            newn = getfield(newc, :numel)
            newcells = setindex(cells, newc, idx)
            return PTree{T}(id, bits, numel - oldn + newn, newcells)
        else
            # We add a twig beneath this node.
            newcells = insert(cells, idx, PTree{T}(k => v))
            return PTree{T}(id, bits | flag, numel + 1, newcells)
        end
    end
end
delete(u::PTree{T}, k::HASH_T) where {T} = begin
    id = getfield(u, :id)
    bits = getfield(u, :bits)
    (inq, bitidx, idx) = ptree_cellindex(id, bits, k)
    # If the node lies outside of this node or the bit isn't set, there's
    # nothing to do.
    inq || return u
    # There's something to delete; start by determining if we are or are not a
    # twig node:
    numel = getfield(u, :numel)
    # Check if we're a twig or something else.
    if ptree_depth(id) == PTREE_TWIG_DEPTH
        # If we are deleting the only element, we return a new empty node.
        (numel == 1) && return PTree{T}()
        cells = getfield(u, :cells)::Vector{T}
        newcells = delete(cells, idx)
        return PTree{T}(id, bits & ~(BITS_ONE << bitidx), numel-1, newcells)
    else
        cells = getfield(u, :cells)::Vector{PTree{T}}
        oldc = cells[idx]
        oldn = getfield(oldc, :numel)
        newc = delete(oldc, k)
        newn = getfield(newc, :numel)
        numel += newn - oldn
        (numel == 0) && return newc
        if newn == 0
            newcells = delete(cells, idx)
            return PTree{T}(id, bits & ~(BITS_ONE << bitidx), numel, newcells)
        else
            newcells = setindex(cells, newc, idx)
            return PTree{T}(id, bits, numel, newcells)
        end
    end
end


using Merlin

function maxrows(m::Matrix)
  _, inds = findmax(m, 1)
  map!(i -> ind2sub(size(m), i)[1], inds)
  vec(inds)
end

function train(path)
  # data
  worddict = read_wordlist("$(path)/words.lst")
  chardict, tagdict = Dict(), Dict()
  traindata = read_conll("$(path)/wsj_00-18.conll", true, worddict, chardict, tagdict)
  println("#word: $(length(worddict)), #char: $(length(chardict)), #tag: $(length(tagdict))")
  traindata = traindata[1:5000]
  testdata = read_conll("$(path)/wsj_22-24.conll", false, worddict, chardict, tagdict)

  # model
  model = POSModel(path)
  opt = SGD(0.0075)

  for iter = 1:10
    println("iter: $(iter)")
    golds, preds = Int[], Int[]
    opt.learnrate = 0.0075 / iter
    loss = 0.0

    for i in randperm(length(traindata))
      tokens = traindata[i]
      append!(golds, map(t -> t.tagid, tokens))

      # forward & prediction
      out = forward(model, tokens)
      maxidx = maxrows(out.value)
      append!(preds, maxidx)

      # loss function
      px = zeros(Float32, size(out.value,1), length(tokens))
      for j = 1:length(tokens)
        px[tokens[j].tagid, j] = 1.0
      end
      out = CrossEntropy(px)(out)
      loss += sum(out.value)

      # backward & update
      backward!(out)
      update!(model, opt)
    end
    println("loss: $(loss)")
    acc = accuracy(golds, preds)
    println("train acc: $(acc)")
    decode(model, testdata)
    println("")
  end
  println("finish")
end

function decode(m::POSModel, data::Vector{Vector{Token}})
  golds, preds = Int[], Int[]
  for i = 1:length(data)
    tokens = data[i]
    append!(golds, map(t -> t.tagid, tokens))
    out = forward(m, tokens)
    append!(preds, maxrows(out.value))
  end
  acc = accuracy(golds, preds)
  println("test acc: $(acc)")
end


module Dolfin

using PyCall

dolfin = PyCall.pyimport_conda("dolfin", "dolfin")
operator = PyCall.pyimport_conda("operator", "operator")

dolfin_methods =
["ALE", "AdaptiveLinearVariationalSolver", "AdaptiveNonlinearVariationalSolver",
 "And", "Argument", "Arguments", "Assembler", "AssemblerBase", "AutoSubDomain",
 "BDF1", "BackwardEuler", "BasisFunction", "BlockMatrix", "BlockVector",
 "BoolParameter", "BoostGraphOrdering", "BoundaryMesh", "BoundingBoxTree",
 "BoundingBoxTree3D", "BoxMesh", "BrokenElement", "ButcherMultiStageScheme",
 "CN2", "CRITICAL", "Cell", "CellFunction", "CellFunctionBool",
"CellFunctionDouble", "CellFunctionInt", "CellFunctionSizet", "CellNormal",
"CellSize", "CellType", "CellVolume", "Circumradius", "Coefficient",
"Coefficients", "CompiledSubDomain", "Constant", "CoordinateMatrix",
"CrankNicolson", "DBG", "DEBUG", "DOLFIN_EPS", "DOLFIN_EPS_LARGE",
"DOLFIN_LINELENGTH", "DOLFIN_PI", "DOLFIN_SQRT_EPS", "DOLFIN_TERM_WIDTH",
"DOLFIN_VERSION_GIT", "DOLFIN_VERSION_MAJOR", "DOLFIN_VERSION_MICRO",
"DOLFIN_VERSION_MINOR", "DOLFIN_VERSION_RELEASE", "DOLFIN_VERSION_STRING",
"DefaultFactory", "DirichletBC", "Dn", "DofMap", "Domain", "DomainBoundary",
"DoubleArray", "DoubleParameter", "Dx", "DynamicMeshEditor", "ERK", "ERK1",
"ERK4", "ERROR", "ESDIRK3", "ESDIRK4", "Edge", "EdgeFunction",
"EdgeFunctionBool", "EdgeFunctionDouble", "EdgeFunctionInt",
"EdgeFunctionSizet", "EigenFactory", "EigenKrylovSolver",
"EigenLUSolver", "EigenMatrix", "EigenVector", "EnrichedElement",
"Equation", "ErrorControl", "Event", "ExplicitEuler",
"ExplicitMidPoint", "Expression", "Extrapolation", "Face",
"FaceFunction", "FaceFunctionBool", "FaceFunctionDouble",
"FaceFunctionInt", "FaceFunctionSizet", "Facet", "FacetArea",
"FacetCell", "FacetElement", "FacetFunction", "FacetFunctionBool",
"FacetFunctionDouble", "FacetFunctionInt", "FacetFunctionSizet",
"FacetNormal", "File", "FiniteElement", "FiniteElementBase", "Form",
"ForwardEuler", "Function", "FunctionAXPY", "FunctionAssigner",
"FunctionSpace", "FunctionSpaceBase", "FunctionSpaceFromCpp", "GRL1",
"GRL2", "GenericAdaptiveVariationalSolver", "GenericBoundingBoxTree",
"GenericDofMap", "GenericFile", "GenericFunction", "GenericLUSolver",
"GenericLinearAlgebraFactory", "GenericLinearOperator",
"GenericLinearSolver", "GenericMatrix", "GenericPreconditioner",
"GenericSparsityPattern", "GenericTensor", "GenericVector",
"GlobalParameters", "Graph", "GraphBuilder", "H1", "H2", "HCurl",
"HDF5Attribute", "HDF5File", "HDiv", "HierarchicalDirichletBC",
"HierarchicalErrorControl", "HierarchicalForm",
"HierarchicalFunction", "HierarchicalFunctionSpace",
"HierarchicalLinearVariationalProblem", "HierarchicalMesh",
"HierarchicalMeshFunctionBool", "HierarchicalMeshFunctionDouble",
"HierarchicalMeshFunctionInt", "HierarchicalMeshFunctionSizet",
"HierarchicalMeshFunctionUInt",
"HierarchicalNonlinearVariationalProblem", "INFO", "Identity",
"ImplicitEuler", "Index", "IndexSet", "IntArray",
"IntParameter", "Integral", "InteriorElement", "IntervalMesh",
"KrylovSolver", "L2", "LUSolver", "Lagrange",
"LagrangeInterpolator", "Legendre", "LinearAlgebraObject",
"LinearOperator", "LinearSolver", "LinearVariationalProblem",
"LinearVariationalSolver", "LocalAssembler", "LocalMeshData",
"LocalSolver", "MPI", "MPICH_IGNORE_CXX_SEEK", "MPIInfo",
"MPI_Comm", "Matrix", "Max", "MaxCellEdgeLength",
"MaxFacetEdgeLength", "Measure", "Mesh", "MeshColoring",
"MeshConnectivity", "MeshCoordinates", "MeshData",
"MeshDisplacement", "MeshDomains", "MeshEditor",
"MeshEntity", "MeshFunction", "MeshFunctionBool",
"MeshFunctionDouble", "MeshFunctionInt",
"MeshFunctionSizet", "MeshGeometry", "MeshHierarchy",
"MeshPartitioning", "MeshPointIntersection",
"MeshQuality", "MeshRenumbering", "MeshTopology",
"MeshTransformation", "MeshValueCollection",
"MeshValueCollectionBool", "MeshValueCollectionDouble",
"MeshValueCollectionInt", "MeshValueCollectionSizet", "Min",
"MinCellEdgeLength", "MinFacetEdgeLength", "MixedElement", "MixedFunctionSpace",
"MultiMesh", "MultiMeshAssembler", "MultiMeshDirichletBC", "MultiMeshDofMap",
"MultiMeshForm", "MultiMeshFunction", "MultiMeshFunctionSpace",
"MultiMeshSubSpace", "MultiStageScheme", "NewtonSolver", "NonlinearProblem",
"NonlinearVariationalProblem", "NonlinearVariationalSolver", "Not",
"OpenMpAssembler", "OptimisationProblem", "Or", "OuterProductCell",
"OuterProductElement", "OuterProductVectorElement", "PETScBaseMatrix",
"PETScFactory", "PETScKrylovSolver", "PETScLUSolver", "PETScLinearOperator",
"PETScMatrix", "PETScObject", "PETScOptions", "PETScPreconditioner",
"PETScSNESSolver", "PETScUserPreconditioner", "PETScVector", "PROGRESS",
"ParameterValue", "Parameters", "PeriodicBoundaryComputation",
"PermutationSymbol", "Point", "PointIntegralSolver", "PointSource",
"ProductCell", "ProductDomain", "Progress", "RK4", "RKSolver", "RL1", "RL2",
"RectangleMesh", "RestrictedElement", "Restriction", "RushLarsenScheme",
"SCOTCH", "SLEPcEigenSolver", "STLFactory", "STLMatrix", "Scalar",
"SparsityPattern", "SparsityPatternBuilder", "SpatialCoordinate",
"SpecialFacetFunction", "StringParameter", "SubDomain", "SubMesh", "SubSpace",
"SubSystemsManager", "SubsetIterator", "SystemAssembler", "TRACE", "Table",
"TableEntry", "TensorConstant", "TensorElement", "TensorFunctionSpace",
"TensorLayout", "TensorProductElement", "TestFunction", "TestFunctions",
"Time", "TimeSeries", "Timer", "TimingClear_clear", "TimingClear_keep",
"TimingType_system", "TimingType_user", "TimingType_wall", "TraceElement",
"TrialFunction", "TrialFunctions", "UFLException", "UIntArray", "UnitCubeMesh",
"UnitHexMesh", "UnitIntervalMesh", "UnitQuadMesh", "UnitSquareMesh",
"UnitTetrahedronMesh", "UnitTriangleMesh", "VTKPlotter", "Variable", "Vector",
"VectorConstant", "VectorElement", "VectorFunctionSpace", "VectorSpaceBasis",
"Vertex", "VertexFunction", "VertexFunctionBool", "VertexFunctionDouble",
"VertexFunctionInt", "VertexFunctionSizet", "WARNING", "XDMFFile",
"acos", "action",
"adapt", "adapt_markers", "add_logfile", "adjoint", "as_backend_type",
"as_cell", "as_domain", "as_matrix", "as_tensor", "as_ufl", "as_vector",
"asin", "assemble", "assemble_system", "assign", "atan", "atan_2", "avg",
 "bessel_I", "bessel_J", "bessel_K", "bessel_Y", "between",
"cell_avg", "cells", "cofac", "common", "compile_expressions",
"compile_extension_module", "compilemodules", "compute_edge_map",
"compute_vertex_map", "conditional", "cos", "cosh", "cpp",
"create_ufc_function_spaces", "cross", "curl", "dC", "dI", "dO", "dP", "dS",
"dS_h", "dS_v", "dc", "debug", "deprecation", "derivative", "det", "dev",
"diag", "diag_vector", "diff", "div", "dof_to_vertex_map", "dolfin",
"dolfin_error", "dolfin_version", "dorfler_mark", "dot", "down_cast",
"ds", "ds_b", "ds_t", "ds_tb", "ds_v", "dump_timings_to_xml", "dx", "e",
"edges", "elem_div", "elem_mult", "elem_op", "elem_pow", "energy_norm",
"entities", "eq", "erf", "error", "errornorm", "exp", "exterior_derivative",
"f", "faces", "facet", "facet_avg", "facets", "fem", "fem_solve",
"functional", "functions", "ge", "generate_error_control",
"generate_error_control_forms", "get_global_parameters", "get_handler",
"get_log_level", "get_logger", "get_tensor_type", "git_commit_hash", "grad",
"gt", "has_cholmod", "has_hdf5", "has_krylov_solver_method",
"has_krylov_solver_preconditioner", "has_linear_algebra_backend",
"has_lu_solver_method", "has_mpi", "has_openmp", "has_parmetis", "has_petsc",
"has_petsc4py", "has_scotch", "has_slepc", "has_slepc4py", "has_tao",
"has_type", "has_umfpack", "has_zlib", "hexahedron", "homogenize", "i",
"importhandler", "in_nullspace", "increase_order", "indices", "info",
"info_blue", "info_green", "info_red", "info_stream", "info_underline",
"init", "init_parent_edge_indices", "inner", "integral_types", "interactive",
"interpolate", "intersect", "interval", "inv", "io", "ipow", "j", "jit", "jump",
"k", "krylov_solver_methods", "krylov_solver_preconditioners", "l", "la",
"la_index_dtype", "la_solve", "le", "lhs", "linear_algebra_backends",
"linear_solver_methods", "list_krylov_solver_methods",
"list_krylov_solver_preconditioners", "list_linear_algebra_backends",
"list_linear_solver_methods", "list_lu_solver_methods", "list_timings", "ln",
"log", "lt", "lt_coordinate", "lu_solver_methods", "mark", "max_value",
"memory_usage", "mesh", "min_value", "monitor_memory_usage", "mpi_comm_self",
"mpi_comm_world", "multistage", "nabla_div", "nabla_grad", "ne", "near",
"norm", "normalize", "not_working_in_parallel", "old_init", "outer", "p",
"parameters", "perp", "plot", "product", "project", "q", "quadrilateral",
"r", "rand", "rank", "refine", "register_element", "register_integral_type",
"relabel", "replace", "replace_integral_domains", "residual", "rhs", "rot",
"s", "seed", "sensitivity_rhs", "set_log_active", "set_log_level",
"set_output_stream", "shape", "show_elements", "sign", "sin", "sinh",
"sizeof_la_index", "skew", "solve", "split", "sqr", "sqrt",
"supported_elements", "supported_elements_for_plotting", "sym", "system",
"tan", "tanh", "tear", "tetrahedron", "tic", "time", "timing", "timings",
"toc", "tr", "transpose", "triangle", "unit_matrices", "unit_matrix",
"unit_vector", "unit_vectors", "variable", "vertex", "vertex_to_dof_map",
"vertices", "warning", "zero"]

function __init__()
  for meth in dolfin_methods
    ex = string(meth, "=", "dolfin[\"", meth, "\"]")
    prog = parse(ex)
    eval(prog)
    eval(Expr(:export, parse(meth)))
  end
end

import Base:+, -,*,/, ==, <<
import PyCall:isequal
+(f::PyObject, g::PyObject) = operator[:add](f,g)
+(f::PyObject, g) = operator[:add](f,Constant(g))

-(f::PyObject, g::PyObject) = operator[:sub](f,g)
-(f::PyObject, g) = operator[:sub](f,Constant(g))

*(f::PyObject, g::PyObject) = operator[:mul](f,g)
*(f::PyObject, g) = operator[:mul](f,Constant(g))

/(f::PyObject, g::PyObject) = operator[:div](f,g)
/(f::PyObject, g) = operator[:div](f,Constant(g))

==(f::PyObject, g::PyObject) = operator[:eq](f,g)
-(f::PyObject) = operator[:neg](f)
<<(f::PyObject, g::PyObject) = operator[:lshift](f, g)

end


real_rep(a::AbstractArray{DualT}) where {TAG, T, DualT<:ForwardDiff.Dual{TAG, T}} = reinterpret(reshape, T, a)
_view1(B::AbstractMatrix) = @view(B[1,:])
_view1(B::AbstractArray{<:Any,3}) = @view(B[1,:,:])

# multiplication of dual vector/matrix by standard matrix from the left
function _matmul!(_C::AbstractVecOrMat{DualT}, A::AbstractMatrix, _B::AbstractVecOrMat{DualT},
                  α, β=Zero(), nthread::Nothing=nothing, MKN=nothing, contig_axis=nothing) where {DualT <: ForwardDiff.Dual}
  B = real_rep(_B)
  C = real_rep(_C)

  @tturbo for n ∈ indices((C, B), 3), m ∈ indices((C, A), (2, 1)), l in indices((C, B), 1)
    Cₗₘₙ = zero(eltype(C))
    for k ∈ indices((A, B), 2)
      Cₗₘₙ += A[m, k] * B[l, k, n]
    end
    C[l, m, n] = α * Cₗₘₙ + β * C[l, m, n]
  end

  _C
end

# multiplication of dual matrix by standard vector/matrix from the right
@inline function _matmul!(_C::AbstractVecOrMat{DualT}, _A::AbstractMatrix{DualT}, B::AbstractVecOrMat,
                          α=One(), β=Zero(), nthread::Nothing=nothing, MKN=nothing) where {TAG, T, DualT <: ForwardDiff.Dual{TAG, T}}
  if Bool(ArrayInterface.is_dense(_C)) && Bool(ArrayInterface.is_column_major(_C)) &&
    Bool(ArrayInterface.is_dense(_A)) && Bool(ArrayInterface.is_column_major(_A))
    # we can avoid the reshape and call the standard method
    A = reinterpret(T, _A)
    C = reinterpret(T, _C)
    _matmul!(C, A, B, α, β, nthread, nothing)
  else
    # we cannot use the standard method directly
    A = real_rep(_A)
    C = real_rep(_C)

    @tturbo for n ∈ indices((C, B), (3, 2)), m ∈ indices((C, A), 2), l in indices((C, A), 1)
      Cₗₘₙ = zero(eltype(C))
      for k ∈ indices((A, B), (3, 1))
        Cₗₘₙ += A[l, m, k] * B[k, n]
      end
      C[l, m, n] = α * Cₗₘₙ + β * C[l, m, n]
    end
  end

  _C
end

@inline function _matmul!(_C::AbstractVecOrMat{DualT}, _A::AbstractMatrix{DualT}, _B::AbstractVecOrMat{DualT},
                          α=One(), β=Zero(), nthread::Nothing=nothing, MKN=nothing, contig=nothing) where {TAG, T, P, DualT <: ForwardDiff.Dual{TAG, T, P}}
  A = real_rep(_A)
  C = real_rep(_C)
  B = real_rep(_B)
  if Bool(ArrayInterface.is_dense(_C)) && Bool(ArrayInterface.is_column_major(_C)) &&
    Bool(ArrayInterface.is_dense(_A)) && Bool(ArrayInterface.is_column_major(_A))
    # we can avoid the reshape and call the standard method
    Ar = reinterpret(T, _A)
    Cr = reinterpret(T, _C)
    _matmul!(Cr, Ar, _view1(B), α, β, nthread, nothing)
  else
    # we cannot use the standard method directly
    @tturbo for n ∈ indices((C, B), 3), m ∈ indices((C, A), 2), l in indices((C, A), 1)
      Cₗₘₙ = zero(eltype(C))
      for k ∈ indices((A, B), (3, 2))
        Cₗₘₙ += A[l, m, k] * B[1, k, n]
      end
      C[l, m, n] = α * Cₗₘₙ + β * C[l, m, n]
    end
  end
  Pstatic = static(P)
  @tturbo for n ∈ indices((B,C),3), m ∈ indices((A,C),2), p ∈ 1:Pstatic
    Cₚₘₙ = zero(eltype(C))
    for k ∈ indices((A,B),(3,2))
      Cₚₘₙ += A[1,m,k] * B[p+1,k,n]
    end
    C[p+1,m,n] = C[p+1,m,n] + α*Cₚₘₙ
  end
  _C
end

# multiplication of dual vector/matrix by standard matrix from the left
function _matmul_serial!(_C::AbstractVecOrMat{DualT}, A::AbstractMatrix, _B::AbstractVecOrMat{DualT},
                         α, β, MKN) where {DualT <: ForwardDiff.Dual}
  B = real_rep(_B)
  C = real_rep(_C)

  @turbo for n ∈ indices((C, B), 3), m ∈ indices((C, A), (2, 1)), l in indices((C, B), 1)
    Cₗₘₙ = zero(eltype(C))
    for k ∈ indices((A, B), 2)
      Cₗₘₙ += A[m, k] * B[l, k, n]
    end
    C[l, m, n] = α * Cₗₘₙ + β * C[l, m, n]
  end

  _C
end

# multiplication of dual matrix by standard vector/matrix from the right
@inline function _matmul_serial!(_C::AbstractVecOrMat{DualT}, _A::AbstractMatrix{DualT}, B::AbstractVecOrMat,
                                 α, β, MKN) where {TAG, T, DualT <: ForwardDiff.Dual{TAG, T}}
  if Bool(ArrayInterface.is_dense(_C)) && Bool(ArrayInterface.is_column_major(_C)) &&
    Bool(ArrayInterface.is_dense(_A)) && Bool(ArrayInterface.is_column_major(_A))
    # we can avoid the reshape and call the standard method
    A = reinterpret(T, _A)
    C = reinterpret(T, _C)
    _matmul_serial!(C, A, B, α, β, nothing)
  else
    # we cannot use the standard method directly
    A = real_rep(_A)
    C = real_rep(_C)

    @turbo for n ∈ indices((C, B), (3, 2)), m ∈ indices((C, A), 2), l in indices((C, A), 1)
      Cₗₘₙ = zero(eltype(C))
      for k ∈ indices((A, B), (3, 1))
        Cₗₘₙ += A[l, m, k] * B[k, n]
      end
      C[l, m, n] = α * Cₗₘₙ + β * C[l, m, n]
    end
  end

  _C
end

@inline function _matmul_serial!(_C::AbstractVecOrMat{DualT}, _A::AbstractMatrix{DualT}, _B::AbstractVecOrMat{DualT},
                                 α, β, MKN) where {TAG, T, P, DualT <: ForwardDiff.Dual{TAG, T, P}}
  A = real_rep(_A)
  C = real_rep(_C)
  B = real_rep(_B)
  if Bool(ArrayInterface.is_dense(_C)) && Bool(ArrayInterface.is_column_major(_C)) &&
    Bool(ArrayInterface.is_dense(_A)) && Bool(ArrayInterface.is_column_major(_A))
    # we can avoid the reshape and call the standard method
    Ar = reinterpret(T, _A)
    Cr = reinterpret(T, _C)
    _matmul_serial!(Cr, Ar, _view1(B), α, β, nothing)
  else
    # we cannot use the standard method directly
    @turbo for n ∈ indices((C, B), 3), m ∈ indices((C, A), 2), l in indices((C, A), 1)
      Cₗₘₙ = zero(eltype(C))
      for k ∈ indices((A, B), (3, 2))
        Cₗₘₙ += A[l, m, k] * B[1, k, n]
      end
      C[l, m, n] = α * Cₗₘₙ + β * C[l, m, n]
    end
  end
  Pstatic = static(P)
  @turbo for n ∈ indices((B,C),3), m ∈ indices((A,C),2), p ∈ 1:Pstatic
    Cₚₘₙ = zero(eltype(C))
    for k ∈ indices((A,B),(3,2))
      Cₚₘₙ += A[1,m,k] * B[p+1,k,n]
    end
    C[p+1,m,n] = C[p+1,m,n] + α*Cₚₘₙ
  end
  _C
end


# Space-walking recursive functions for random search and tree-parzen estimation.

walkable(x::Delayed.AbstractParam) = true
walkable(x) = false

# Entrypoint (only for TPE, not for random search)
"""
$(TYPEDSIGNATURES)

Resolves a posterior inference space by depth-first iteration, replacing prior random-
variables with new posterior distributions that make use of observations.
"""
function node(
    space::Types.SPACE_TYPE, trials::Vector{Trials.Trial}, config::Config
)::Tuple{Trials.ValsDict, Any}

    params = Dict{Symbol, Types.AbstractDelayed}(
        item.label => item.obj
            for item in Graph.dfs(space)
                if walkable(item)
    )

    vals = Trials.ValsDict()

    hyperparams = node(space, vals, params, trials, config)

    return vals, hyperparams
end

function obs_memo(
    item::Delayed.AbstractDistDelayed, params::Dict{Symbol, Types.AbstractDelayed}
)::Symbol
    for (k, v) in params
        v == item && return k
    end

    throw(DomainError("item was not a value in params"))
end

# Node resolvers. For each node type there is first a random-search resolver and then a TPE
# resolver. They are distinguished by the TPE methods requiring extra items.
"""
$(TYPEDSIGNATURES)

Resolves random search AbstractParam nodes and places parameter results in the vals dictionary.
"""
function node(item::Delayed.AbstractParam, vals::Trials.ValsDict)::Union{IndexObjects.IndexInt, Real}
    if haskey(vals, item.label)
        throw(KeyError("Key $(item) already present in $(vals[item.label])"))
    end

    obj = node(item.obj, vals)
    vals[item.label] = IndexObjects.getval(obj)

    return obj
end
"""
$(TYPEDSIGNATURES)

Resolves random search AbstractParam nodes and places parameter results in the vals dictionary.
"""
function node(
    item::Delayed.AbstractParam, vals::Trials.ValsDict, params::Dict{Symbol, Types.AbstractDelayed},
    trials::Vector{Trials.Trial}, config::Config
)::Union{IndexObjects.IndexInt, Real}
    if haskey(vals, item.label)
        throw(KeyError("Key $(item.label) already present $(vals[item.label])"))
    end

    obj = node(item.obj, vals, params, trials, config)
    vals[item.label] = IndexObjects.getval(obj)

    return obj
end

function node(item::Delayed.BinaryOperator, vals::Trials.ValsDict)::Real

    left = node(item.left, vals)
    right = node(item.right, vals)

    return item.operator(left, right)
end
function node(
    item::Delayed.BinaryOperator, vals::Trials.ValsDict, params::Dict{Symbol, Types.AbstractDelayed},
    trials::Vector{Trials.Trial}, config::Config
)::Float64

    left = node(item.left, vals, params, trials, config)
    right = node(item.right, vals, params, trials, config)

    return item.operator(left, right)
end

function node(item::Delayed.UnaryOperator, vals::Trials.ValsDict)::Real

    operand = node(item.operand, vals)

    return item.operator(operand)
end
function node(
    item::Delayed.UnaryOperator, vals::Trials.ValsDict, params::Dict{Symbol, Types.AbstractDelayed},
    trials::Vector{Trials.Trial}, config::Config
)::Float64

    operand = node(item.operand, vals, params, trials, config)

    return item.operator(operand)
end

function node(item::Delayed.CategoricalIndex, vals::Trials.ValsDict)::IndexObjects.IndexInt
    return IndexObjects.IndexInt(first(Delayed.categoricalindex(item.probabilities, 1).v))
end
function node(
    item::Delayed.CategoricalIndex, vals::Trials.ValsDict,
    params::Dict{Symbol, Types.AbstractDelayed}, trials::Vector{Trials.Trial},
    config::Config
)::IndexObjects.IndexInt

    nid = obs_memo(item, params)

    return Resolve.posterior(item, item.probabilities, nid, trials, config)
end

function node(item::Dict{Symbol, T} where T, vals::Trials.ValsDict)::Dict{Symbol, Any}

    result = Dict{Symbol, Any}()
    # Cannot use a dictionary comprehension because the values will be evaluated twice.
    for (k, v) in item
        result[k] = node(v, vals)
    end

    return result
end
function node(
    item::Dict{Symbol, T} where T, vals::Trials.ValsDict,
    params::Dict{Symbol, Types.AbstractDelayed}, trials::Vector{Trials.Trial},
    config::Config
)::Dict{Symbol, Any}

    result = Dict{Symbol, Any}()
    # Cannot use a dictionary comprehension because the values will be evaluated twice.
    for (k, v) in item
        result[k] = node(v, vals, params, trials, config)
    end

    return result
end

function node(item::Delayed.Float, vals::Trials.ValsDict)::Float64

    return float(node(item.arg, vals))
end
function node(
    item::Delayed.Float, vals::Trials.ValsDict, params::Dict{Symbol, Types.AbstractDelayed},
    trials::Vector{Trials.Trial}, config::Config
)::Float64

    return float(node(item.arg, vals, params, trials, config))
end

function node(item::Delayed.LogNormal, vals::Trials.ValsDict)::Float64
    mu = node(item.mu, vals)
    sigma = node(item.sigma, vals)

    return Delayed.lognormal(mu, sigma)
end
function node(
    item::Delayed.LogNormal, vals::Trials.ValsDict,
    params::Dict{Symbol, Types.AbstractDelayed}, trials::Vector{Trials.Trial},
    config::Config
)::Real

    mu = node(item.mu, vals, params, trials, config)
    sigma = node(item.sigma, vals, params, trials, config)
    nid = obs_memo(item, params)

    return Resolve.posterior(item, mu, sigma, nid, trials, config)
end

function node(item::Delayed.LogNormalQuantDist, vals::Trials.ValsDict)::Float64
    mu = node(item.mu, vals)
    sigma = node(item.sigma, vals)
    q = node(item.q, vals)

    return Delayed.lognormalquant(item, mu, sigma, q)
end
function node(
    item::Delayed.LogNormalQuantDist, vals::Trials.ValsDict,
    params::Dict{Symbol, Types.AbstractDelayed}, trials::Vector{Trials.Trial},
    config::Config
)::Real

    mu = node(item.mu, vals, params, trials, config)
    sigma = node(item.sigma, vals, params, trials, config)
    q = node(item.q, vals, params, trials, config)
    nid = obs_memo(item, params)

    return Resolve.posterior(item, mu, sigma, q, nid, trials, config)
end

function node(item::Delayed.Normal, vals::Trials.ValsDict)::Float64
    mu = node(item.mu, vals)
    sigma = node(item.sigma, vals)

    return Delayed.normal(mu, sigma)
end
function node(
    item::Delayed.Normal, vals::Trials.ValsDict, params::Dict{Symbol, Types.AbstractDelayed},
    trials::Vector{Trials.Trial}, config::Config
)::Real

    mu = node(item.mu, vals, params, trials, config)
    sigma = node(item.sigma, vals, params, trials, config)
    nid = obs_memo(item, params)

    return Resolve.posterior(item, mu, sigma, nid, trials, config)
end

function node(item::Delayed.QuantNormal, vals::Trials.ValsDict)::Float64
    mu = node(item.mu, vals)
    sigma = node(item.sigma, vals)
    q = node(item.q, vals)

    return Delayed.quantnormal(mu, sigma, q)
end
function node(
    item::Delayed.QuantNormal, vals::Trials.ValsDict, params::Dict{Symbol, Types.AbstractDelayed},
    trials::Vector{Trials.Trial}, config::Config
)::Real

    mu = node(item.mu, vals, params, trials, config)
    sigma = node(item.sigma, vals, params, trials, config)
    q = node(item.q, vals, params, trials, config)
    nid = obs_memo(item, params)

    return Resolve.posterior(item, mu, sigma, q, nid, trials, config)
end

function node(item::Delayed.RandIndex, vals::Trials.ValsDict)::IndexObjects.IndexInt

    upper = node(item.upper, vals)
    if upper < 1
        throw(ArgumentError("upper will be used as index so must be greater than 0"))
    end

    return IndexObjects.IndexInt(rand(1:upper))
end
function node(
    item::Delayed.RandIndex, vals::Trials.ValsDict, params::Dict{Symbol, Types.AbstractDelayed},
    trials::Vector{Trials.Trial}, config::Config
)::IndexObjects.IndexInt

    upper = node(item.upper, vals, params, trials, config)
    if upper < 1
        throw(ArgumentError("upper will be used as index so must be greater than 0"))
    end
    nid = obs_memo(item, params)

    return Resolve.posterior(item, upper, nid, trials, config)
end

function node(item::Delayed.AbstractSwitch, vals::Trials.ValsDict)

    # Randomly generate a number to indicate which index of the options will be returned
    choice = node(item.choice, vals)

    return node(item.options[choice.v], vals)
end
function node(
    item::Delayed.AbstractSwitch, vals::Trials.ValsDict, params::Dict{Symbol, Types.AbstractDelayed},
    trials::Vector{Trials.Trial}, config::Config
)
    choice = node(item.choice, vals, params, trials, config)

    return node(item.options[choice.v], vals, params, trials, config)
end

function node(item::Delayed.Uniform, vals::Trials.ValsDict)::Float64
    low = node(item.low, vals)
    high = node(item.high, vals)

    return Delayed.uniform(low, high)
end
function node(
    item::Delayed.Uniform, vals::Trials.ValsDict, params::Dict{Symbol, Types.AbstractDelayed},
    trials::Vector{Trials.Trial}, config::Config
)::Real

    low = node(item.low, vals, params, trials, config)
    high = node(item.high, vals, params, trials, config)
    nid = obs_memo(item, params)

    return Resolve.posterior(item, low, high, nid, trials, config)
end

function node(item::Delayed.QuantUniform, vals::Trials.ValsDict)::Float64
    low = node(item.low, vals)
    high = node(item.high, vals)
    q = node(item.q, vals)

    return Delayed.quantuniform(low, high, q)
end
function node(
    item::Delayed.QuantUniform, vals::Trials.ValsDict, params::Dict{Symbol, Types.AbstractDelayed},
    trials::Vector{Trials.Trial}, config::Config
)::Real

    low = node(item.low, vals, params, trials, config)
    high = node(item.high, vals, params, trials, config)
    q = node(item.q, vals, params, trials, config)
    nid = obs_memo(item, params)

    return Resolve.posterior(item, low, high, q, nid, trials, config)
end

function node(item::Delayed.LogUniform, vals::Trials.ValsDict)::Float64

    low = node(item.low, vals)
    high = node(item.high, vals)

    return Delayed.loguniform(low, high)
end

function node(
    item::Delayed.LogUniform, vals::Trials.ValsDict,
    params::Dict{Symbol, Types.AbstractDelayed}, trials::Vector{Trials.Trial},
    config::Config
)::Real

    low = node(item.low, vals, params, trials, config)
    high = node(item.high, vals, params, trials, config)
    nid = obs_memo(item, params)

    return Resolve.posterior(item, low, high, nid, trials, config)
end

function node(item::Delayed.LogUniformQuantDist, vals::Trials.ValsDict)::Float64
    low = node(item.low, vals)
    high = node(item.high, vals)
    q = node(item.q, vals)

    return Delayed.loguniformquant(item, low, high, q)
end
function node(
    item::Delayed.LogUniformQuantDist, vals::Trials.ValsDict,
    params::Dict{Symbol, Types.AbstractDelayed}, trials::Vector{Trials.Trial},
    config::Config
)::Real

    low = node(item.low, vals, params, trials, config)
    high = node(item.high, vals, params, trials, config)
    q = node(item.q, vals, params, trials, config)
    nid = obs_memo(item, params)

    return Resolve.posterior(item, low, high, q, nid, trials, config)
end

function node(items::Vector, vals::Trials.ValsDict)::Vector{<: Any}
    return [
        node(item, vals::Trials.ValsDict)
            for item in items
    ]
end
function node(
    items::Vector, vals::Trials.ValsDict, params::Dict{Symbol, Types.AbstractDelayed},
    trials::Vector{Trials.Trial}, config::Config
)::Vector{<: Any}

    return [
        node(item, vals, params, trials, config)
        for item in items
    ]
end

function node(items::Tuple, vals::Trials.ValsDict)::Tuple
    return tuple((
        node(item, vals::Trials.ValsDict)
            for item in items
    )...)
end
function node(
    items::Tuple, vals::Trials.ValsDict, params::Dict{Symbol, Types.AbstractDelayed},
    trials::Vector{Trials.Trial}, config::Config
)::Tuple

    return tuple((
        node(item, vals, params, trials, config)
            for item in items
    )...)
end

function node(item::T, vals::Trials.ValsDict)::T where T <: Union{Real, Symbol, String}
    return item
end
function node(
    item::T, vals::Trials.ValsDict, params::Dict{Symbol, Types.AbstractDelayed},
    trials::Vector{Trials.Trial}, config::Config
)::T where T <: Union{Real, Symbol, String}

    return item
end


using .Networks

using Dates
using JuMP
using Printf

@with_kw mutable struct EnergyMarket <: AbstractMarket
    configs = EnergyMarketConfigs()
end

function run(runnable::Union{EnergyMarket,BalanceMarket},
            ech::Dates.DateTime, firmness, TS::Vector{Dates.DateTime},
            context::AbstractContext)
    problem_name_l = @sprintf("%s_%s", typeof(runnable), ech)

    gratis_starts = Set{Tuple{String,Dates.DateTime}}()
    if runnable.configs.CONSIDER_GRATIS_STARTS
        gratis_starts = init_gratis_start(context, runnable.configs.REF_SCHEDULE_TYPE)
    end
    @debug("gratis_starts : ", gratis_starts)

    runnable.configs.out_path = context.out_dir
    runnable.configs.problem_name = problem_name_l

    return energy_market(get_network(context),
                        TS,
                        get_generators_initial_state(context),
                        get_scenarios(context, ech),
                        get_uncertainties(context, ech),
                        firmness,
                        get_market_schedule(context),
                        get_tso_schedule(context),
                        get_tso_actions(context),
                        gratis_starts,
                        runnable.configs
                        )
end

function update_market_schedule!(context::AbstractContext, ech,
                                result::EnergyMarketModel,
                                firmness,
                                runnable::Union{EnergyMarket,BalanceMarket})
    market_schedule = get_market_schedule(context)
    market_schedule.decider_type = DeciderType(runnable)
    market_schedule.decision_time = ech

    # Pilotables levels
    for ((gen_id, ts, s), p_injected_var) in result.pilotable_model.p_injected
        if get_power_level_firmness(firmness, gen_id, ts) == FREE
            set_prod_value!(market_schedule, gen_id, ts, s, value(p_injected_var))
        elseif get_power_level_firmness(firmness, gen_id, ts) == TO_DECIDE
            set_prod_definitive_value!(market_schedule, gen_id, ts, value(p_injected_var))
        elseif get_power_level_firmness(firmness, gen_id, ts) == DECIDED
            @assert( ismissing(get_prod_value(market_schedule, gen_id, ts))
                    || (value(p_injected_var) ≈ get_prod_value(market_schedule, gen_id, ts)) )
            set_prod_definitive_value!(market_schedule, gen_id, ts, value(p_injected_var))
        end
    end

    # Commitment
    for ((gen_id, ts, s), b_on_var) in result.pilotable_model.b_on
        gen_state_value = parse(GeneratorState, value(b_on_var))
        if get_commitment_firmness(firmness, gen_id, ts) == FREE
            set_commitment_value!(market_schedule, gen_id, ts, s, gen_state_value)
        elseif get_commitment_firmness(firmness, gen_id, ts) in [TO_DECIDE, DECIDED]
            set_commitment_definitive_value!(market_schedule, gen_id, ts, gen_state_value)
        end
    end

    # Capping
    update_schedule_capping!(market_schedule, context, ech, result.limitable_model, runnable.configs.CONSIDER_TSOACTIONS_LIMITATIONS)

    # Limitables levels : needs to be after capping update
    uncertainties_l = get_uncertainties(context, ech)
    for gen in Networks.get_generators_of_type(get_network(context), Networks.LIMITABLE)
        gen_id = Networks.get_id(gen)
        for ts in get_target_timepoints(context)
            for s in get_scenarios(context)
                capped_l = safeget_capping(market_schedule, gen_id, ts, s)
                injected_l = get_uncertainties(uncertainties_l,gen_id,ts,s) - capped_l
                set_prod_value!(market_schedule, gen_id, ts, s, injected_l)
            end
        end
    end

    # loss_of_load (load-shedding)
    update_schedule_loss_of_load!(market_schedule, context, ech, result.lol_model)

    return market_schedule
end


"""
    cnse_boundary_state!(::NumericalFluxFirstOrder, ::Impenetrable{FreeSlip}, ::CNSE3D)
apply free slip boundary condition for velocity
sets reflective ghost point
"""
@inline function cnse_boundary_state!(
    ::NumericalFluxFirstOrder,
    ::Impenetrable{FreeSlip},
    ::CNSE3D,
    ::TurbulenceClosure,
    state⁺,
    aux⁺,
    n⁻,
    state⁻,
    aux⁻,
    t,
    args...,
)
    state⁺.ρ = state⁻.ρ

    ρu⁻ = state⁻.ρu
    state⁺.ρu = ρu⁻ - 2 * n⁻ ⋅ ρu⁻ .* SVector(n⁻)

    return nothing
end

"""
    cnse_boundary_state!(::NumericalFluxGradient, ::Impenetrable{FreeSlip}, ::CNSE3D)
apply free slip boundary condition for velocity
sets non-reflective ghost point
"""
function cnse_boundary_state!(
    ::NumericalFluxGradient,
    ::Impenetrable{FreeSlip},
    ::CNSE3D,
    ::ConstantViscosity,
    state⁺,
    aux⁺,
    n⁻,
    state⁻,
    aux⁻,
    t,
    args...,
)
    state⁺.ρ = state⁻.ρ

    ρu⁻ = state⁻.ρu
    state⁺.ρu = ρu⁻ - n⁻ ⋅ ρu⁻ .* SVector(n⁻)

    return nothing
end

"""
    shallow_normal_boundary_flux_second_order!(::NumericalFluxSecondOrder, ::Impenetrable{FreeSlip}, ::CNSE3D)
apply free slip boundary condition for velocity
apply zero numerical flux in the normal direction
"""
function cnse_boundary_state!(
    ::NumericalFluxSecondOrder,
    ::Impenetrable{FreeSlip},
    ::CNSE3D,
    ::ConstantViscosity,
    state⁺,
    gradflux⁺,
    aux⁺,
    n⁻,
    state⁻,
    gradflux⁻,
    aux⁻,
    t,
    args...,
)
    state⁺.ρu = state⁻.ρu
    gradflux⁺.ν∇u = n⁻ * (@SVector [-0, -0, -0])'

    return nothing
end

"""
    cnse_boundary_state!(::NumericalFluxFirstOrder, ::Impenetrable{NoSlip}, ::CNSE3D)
apply no slip boundary condition for velocity
sets reflective ghost point
"""
@inline function cnse_boundary_state!(
    ::NumericalFluxFirstOrder,
    ::Impenetrable{NoSlip},
    ::CNSE3D,
    ::TurbulenceClosure,
    state⁺,
    aux⁺,
    n⁻,
    state⁻,
    aux⁻,
    t,
    args...,
)
    state⁺.ρ = state⁻.ρ
    state⁺.ρu = -state⁻.ρu

    return nothing
end

"""
    cnse_boundary_state!(::NumericalFluxGradient, ::Impenetrable{NoSlip}, ::CNSE3D)
apply no slip boundary condition for velocity
set numerical flux to zero for U
"""
@inline function cnse_boundary_state!(
    ::NumericalFluxGradient,
    ::Impenetrable{NoSlip},
    ::CNSE3D,
    ::ConstantViscosity,
    state⁺,
    aux⁺,
    n⁻,
    state⁻,
    aux⁻,
    t,
    args...,
)
    FT = eltype(state⁺)
    state⁺.ρu = @SVector zeros(FT, 3)

    return nothing
end

"""
    cnse_boundary_state!(::NumericalFluxSecondOrder, ::Impenetrable{NoSlip}, ::CNSE3D)
apply no slip boundary condition for velocity
sets ghost point to have no numerical flux on the boundary for U
"""
@inline function cnse_boundary_state!(
    ::NumericalFluxSecondOrder,
    ::Impenetrable{NoSlip},
    ::CNSE3D,
    ::ConstantViscosity,
    state⁺,
    gradflux⁺,
    aux⁺,
    n⁻,
    state⁻,
    gradflux⁻,
    aux⁻,
    t,
    args...,
)
    state⁺.ρu = -state⁻.ρu
    gradflux⁺.ν∇u = gradflux⁻.ν∇u

    return nothing
end

"""
    cnse_boundary_state!(::Union{NumericalFluxFirstOrder, NumericalFluxGradient}, ::Penetrable{FreeSlip}, ::CNSE3D)
no mass boundary condition for penetrable
"""
cnse_boundary_state!(
    ::Union{NumericalFluxFirstOrder, NumericalFluxGradient},
    ::Penetrable{FreeSlip},
    ::CNSE3D,
    ::ConstantViscosity,
    _...,
) = nothing

"""
    cnse_boundary_state!(::NumericalFluxSecondOrder, ::Penetrable{FreeSlip}, ::CNSE3D)
apply free slip boundary condition for velocity
apply zero numerical flux in the normal direction
"""
function cnse_boundary_state!(
    ::NumericalFluxSecondOrder,
    ::Penetrable{FreeSlip},
    ::CNSE3D,
    ::ConstantViscosity,
    state⁺,
    gradflux⁺,
    aux⁺,
    n⁻,
    state⁻,
    gradflux⁻,
    aux⁻,
    t,
    args...,
)
    state⁺.ρu = state⁻.ρu
    gradflux⁺.ν∇u = n⁻ * (@SVector [-0, -0, -0])'

    return nothing
end

"""
    cnse_boundary_state!(::Union{NumericalFluxFirstOrder, NumericalFluxGradient}, ::Impenetrable{MomentumFlux}, ::HBModel)
apply kinematic stress boundary condition for velocity
applies free slip conditions for first-order and gradient fluxes
"""
function cnse_boundary_state!(
    nf::Union{NumericalFluxFirstOrder, NumericalFluxGradient},
    ::Impenetrable{<:MomentumFlux},
    model::CNSE3D,
    turb::TurbulenceClosure,
    args...,
)
    return cnse_boundary_state!(
        nf,
        Impenetrable(FreeSlip()),
        model,
        turb,
        args...,
    )
end

"""
    cnse_boundary_state!(::NumericalFluxSecondOrder, ::Impenetrable{MomentumFlux}, ::HBModel)
apply kinematic stress boundary condition for velocity
sets ghost point to have specified flux on the boundary for ν∇u
"""
@inline function cnse_boundary_state!(
    ::NumericalFluxSecondOrder,
    bc::Impenetrable{<:MomentumFlux},
    model::CNSE3D,
    ::ConstantViscosity,
    state⁺,
    gradflux⁺,
    aux⁺,
    n⁻,
    state⁻,
    gradflux⁻,
    aux⁻,
    t,
    args...,
)
    state⁺.ρu = state⁻.ρu
    gradflux⁺.ν∇u = n⁻ * bc.drag(state⁻, aux⁻, t)'

    return nothing
end

"""
    cnse_boundary_state!(::Union{NumericalFluxFirstOrder, NumericalFluxGradient}, ::Penetrable{MomentumFlux}, ::HBModel)
apply kinematic stress boundary condition for velocity
applies free slip conditions for first-order and gradient fluxes
"""
function cnse_boundary_state!(
    nf::Union{NumericalFluxFirstOrder, NumericalFluxGradient},
    ::Penetrable{<:MomentumFlux},
    model::CNSE3D,
    turb::TurbulenceClosure,
    args...,
)
    return cnse_boundary_state!(
        nf,
        Penetrable(FreeSlip()),
        model,
        turb,
        args...,
    )
end

"""
    cnse_boundary_state!(::NumericalFluxSecondOrder, ::Penetrable{MomentumFlux}, ::HBModel)
apply kinematic stress boundary condition for velocity
sets ghost point to have specified flux on the boundary for ν∇u
"""
@inline function cnse_boundary_state!(
    ::NumericalFluxSecondOrder,
    bc::Penetrable{<:MomentumFlux},
    shallow::CNSE3D,
    ::ConstantViscosity,
    state⁺,
    gradflux⁺,
    aux⁺,
    n⁻,
    state⁻,
    gradflux⁻,
    aux⁻,
    t,
    args...,
)
    state⁺.ρu = state⁻.ρu
    gradflux⁺.ν∇u = n⁻ * bc.drag(state⁻, aux⁻, t)'

    return nothing
end


#TODO have to make manager into a specific thing with components to optimize...
@inline @inbounds function set_optim_values!(m::AbstractLedger, vals)
    # if length(@entities_in(m[Optimize])) <= 11
    #     for i = 1:8
    #         m[Spin].data[i] = Spin(ϕ=vals[i])
    #     end
    #     for i = 1:2
    #         m[L].data[i] = L(ϕ=vals[i+8])
    #     end
    #     if length(@entities_in(m[Optimize])) == 11
    #         m[Pb].data[1] = Pb(vals[end])
    #     end
    # else
        for (i, e) in enumerate(@entities_in(m[Optimize]))

            if e ∈ m[Spin]
                m[e] = Spin(ϕ=vals[i])
            elseif e ∈ m[L]
                m[e] = L(ϕ=vals[i])
            else
                m[e] = Pb(vals[i])
            end
        end
    # end
end
@export function optim_value(m, e)
    if e ∈ m[Spin]
        return m[Spin][e].ϕ
    elseif e ∈ m[L]
        return m[L][e].ϕ
    else
        return m[Pb][e].pb
    end
end
@export optim_values(m) = map(x->optim_value(m, x), @entities_in(m[Optimize]))

@export function obj(x, m)
    set_optim_values!(m, x)
    # set_angles!(m, x)
    m[Etot][1].e = 0.0
    update(stage(m, :E_all), m)
    return m[Etot][1].e
end

function BlackBoxOptim.bboptimize(m::AbstractLedger, cs...)
    empty!(m[Optimize])
    comps_to_optimize = map(x->m[x], cs)
    for c in comps_to_optimize
        for e in @entities_in(c)
            m[e] = Optimize()
        end
    end
    vals = optim_values(m)
    orig_angles = copy(vals)
    # res = optimize(x->obj(x, m), vals, ParticleSwarm(fill(-π, length(vals)), fill(1.0π, length(vals)), 140), Optim.Options(iterations=80))
    obj2 = x -> begin
        set_optim_values!(m, x)

        res = optimize(x->obj(x, m), optim_values(m, cs), ConjugateGradient())
        return res.minimum
    end
    res = bboptimize(x->obj2(x),
                     SearchRange = (0.0, 2.0π),
                     NumDimensions=length(vals),
                     MaxTime=0.1, Method=:generating_set_search,TraceMode=:silent)
    return res
end

function Optim.optimize(m::AbstractLedger, cs...)
    empty!(m[Optimize])
    comps_to_optimize = map(x->m[x], cs)
    for c in comps_to_optimize
        for e in @entities_in(c)
            m[e] = Optimize()
        end
    end
    vals = optim_values(m)
    orig_angles = copy(vals)
    res = optimize(x->obj(x, m), orig_angles, method=ConjugateGradient(),g_tol=1e-5)
    return res
end
Optim.optimize(m::AbstractLedger) = optimize(m, Spin, L, Pb)

@export function optimize!(m::AbstractLedger, cs...)
    res = optimize(m, cs...)
    set_optim_values!(m, res.minimizer)
    return m
end

"Takes the states and spaces the angles such that they are uniformly distributed from the starting point to the end point."
function space_out_states!(ledgers)
    n = length(ledgers[1][Optimize])
    total_rotations = zeros(n)
    nl = length(ledgers)
    rotation_directions = Vector{Int}[]
    for i=2:nl
        rot = all_angles(ledgers[i]) .- all_angles(ledgers[i-1])
        total_rotations .+= abs.(rot)
        push!(rotation_directions, sign.(rot))
    end
    rotation_per_ledger = total_rotations./ (nl - 1)
    for i=2:length(ledgers)-1
        rot_directions = rotation_directions[i-1]
        set_optim_values!(ledgers[i], all_angles(ledgers[i-1]) .+ rot_directions .* rotation_per_ledger)
    end
end

"Takes a set of states and performs nudged elastic band optimization on them, returning the optimized images."
function neb(ledgers, c=0.1)
    opt_ledgers = deepcopy.(ledgers)

    n = length(ledgers[1][Optimize])
    #10 is because 10 angles
    nebrange = i -> n*(i-1)+1:n*i
    start_ϕs = zeros(n*(length(ledgers)-2))
    for (i, l) in enumerate(ledgers[2:end-1])
        all_angles!(view(start_ϕs, nebrange(i)), l)
    end

    angle_vec = zeros(n)

    res = optimize(start_ϕs, ConjugateGradient()) do x
        tot = 0.0
        for il in 1:div(length(x), n)
            l = opt_ledgers[il+1]
            tot += obj(view(x, nebrange(il)), l)
        end
        #elastic part

        tot += c*sum((all_angles!(angle_vec, ledgers[1]) .- view(x, 1:n)) .^ 2)
        for i in 2:div(length(x),n)
            tot += c * sum((view(x, nebrange(i)) .- view(x, nebrange(i-1))) .^ 2)
        end
        tot += c*sum((all_angles!(angle_vec, ledgers[end]) .- view(x, length(x)-(n-1):length(x))).^2)
        return tot
    end
    for i = 2:length(opt_ledgers)-1
        ix = i-1
        set_optim_values!(opt_ledgers[i], view(res.minimizer, nebrange(ix)))
    end
    return opt_ledgers
end


#!/usr/bin/env julia

import Celeste.DeterministicVI
import Celeste.DeterministicVI: ElboIntermediateVariables, elbo_likelihood

include(joinpath(Pkg.dir("Celeste"), "test", "SampleData.jl"))

import SampleData

function benchmark_elbo_likelihood()
    ea, vp, catalog = SampleData.gen_sample_star_dataset()

    # Warm up---this compiles the code
    DeterministicVI.elbo_likelihood(ea, vp)

    # clear allocations in case julia is running with --track-allocations=user
    Profile.clear_malloc_data()

    if isempty(ARGS)
        @time DeterministicVI.elbo_likelihood(ea, vp)
    elseif ARGS[1] == "--profile"
        Profile.init(delay=0.0001)
        @profile DeterministicVI.elbo_likelihood(ea, vp)
        Profile.print(format=:flat, sortedby=:count)
    end
end

benchmark_elbo_likelihood()