mutable struct RandWalker
  i::Int
  x::Float64
  y::Float64
  ymin::Float64
  reporter
  function RandWalker(reporter = nothing)
    new(0, 0.0, 0.0, Inf, reporter)
  end
end

function step!(rw::RandWalker)
  rw.i += 1
  rw.x += randn()
  rw.y += (randn() - 0.05) # Slight downward trend to make more fitness improvement happen...
  if rw.y < rw.ymin
    println("New min found: ($(rw.i),$(rw.y))")
    rw.ymin = rw.y
    if rw.reporter != nothing
      rw.reporter(rw.i, rw.ymin)
    end
  end
end

function rand_walk!(rw::RandWalker)
  while true
    print("."); flush(stdout);
    step!(rw)
    sleep(0.4 * rand())
  end
end

using Gadfly

function fitness_plot(iterations, fitnesses)
  plot(x = iterations, y = fitnesses,
    Geom.point, Geom.line,
    Guide.xlabel("Iteration"), Guide.ylabel("Fitness"),
    Stat.step,
    #Scale.y_log10,
  ) |> drawing(10inch, 6inch)
end

function main(window)
  push!(window.assets, "widgets")

  # If we don't already have a fitness history we create one and start a
  # random walker that updates it.
  if !haskey(task_local_storage(), :fhistory)
    println("Starting random walker...")
    fitness_history = Input( (Int, Float64)[] )

    function report_new_min(i::Int, newmin::Float64)
      push!(fitness_history, push!(value(fitness_history), (i, newmin)))
    end

    # Create random walker, start task for walking
    rw = RandWalker(report_new_min)
    t = @async rand_walk!(rw)

    task_local_storage(:fhistory, fitness_history)
    task_local_storage(:rwtask, t)
  end

  local content

  lift(task_local_storage(:fhistory)) do hist
    println("Changed! len = ", length(hist))

    if length(hist) >= 1
      its  = map(first, hist)
      fits = map(t -> t[2], hist)
      content = vbox(
        hbox("Best fitness:       ", hskip(1em), @sprintf("%.3f", fits[end]) |> emph),
        hbox("Found at iteration: ", hskip(1em), string(its[end]) |> emph),
        fitness_plot(its, fits),
      ) |> packacross(center)
    else
      content = hbox("No min found yet...")
    end

    task_started = task_running = false
    if haskey(task_local_storage(), :rwtask)
      t = task_local_storage(:rwtask)
      task_running = !istaskdone(t)
      #task_started = istaskstarted(t) # Only in 0.4
    end

    vbox(
      h1("Fitness progress") |> emph,
      content,
      #hbox("Task started: ", string(task_started)), # Only in 0.4
      #hbox("Task running: ", hskip(1em), string(task_running)),
    ) |> packacross(center)
  end
end


#####################
# Credal circuits
#####################
abstract type CredalΔNode{O} <: DecoratorΔNode{O} end
abstract type CredalLeafNode{O} <: CredalΔNode{O} end
abstract type CredalInnerNode{O} <: CredalΔNode{O} end

mutable struct CredalLiteral{O} <: CredalLeafNode{O}
    origin::O
    data
    bit::Bool
    CredalLiteral(n) = new{node_type(n)}(n, nothing, false)
end

mutable struct Credal⋀{O} <: CredalInnerNode{O}
    origin::O
    children::Vector{<:CredalΔNode{<:O}}
    data
    bit::Bool
    Credal⋀(n, children) = begin
        new{node_type(n)}(n, convert(Vector{CredalΔNode{node_type(n)}},children), nothing, false)
    end
end

mutable struct Credal⋁{O} <: CredalInnerNode{O}
    origin::O
    children::Vector{<:CredalΔNode{<:O}}
    log_thetas::Vector{Float64}
    log_thetas_u::Vector{Float64}
    data
    bit::Bool
    Credal⋁(n, children) = new{node_type(n)}(n, convert(Vector{CredalΔNode{node_type(n)}},children), some_vector(Float64, length(children)), some_vector(Float64, length(children)), nothing, false)
end

const CredalΔ{O} = AbstractVector{<:CredalΔNode{<:O}}

Base.eltype(::Type{CredalΔ{O}}) where {O} = CredalΔNode{<:O}

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

import LogicCircuits.GateType # make available for extension
import LogicCircuits.node_type

@inline GateType(::Type{<:CredalLiteral}) = LiteralGate()
@inline GateType(::Type{<:Credal⋀}) = ⋀Gate()
@inline GateType(::Type{<:Credal⋁}) = ⋁Gate()

@inline node_type(::CredalΔNode) = CredalΔNode

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

const CredalCache = Dict{ΔNode, CredalΔNode}

function CredalΔ2(circuit::Δ)::CredalΔ
    node2dag(CredalΔ2(circuit[end]))
end

function CredalΔ2(circuit::ΔNode)::CredalΔNode
    f_con(n) = error("Cannot construct a probabilistic circuit from constant leafs: first smooth and remove unsatisfiable branches.")
    f_lit(n) = CredalLiteral(n)
    f_a(n, cn) = Credal⋀(n, cn)
    f_o(n, cn) = Credal⋁(n, cn)
    foldup_aggregate(circuit, f_con, f_lit, f_a, f_o, CredalΔNode{node_type(circuit)})
end

function CredalΔ(circuit::Δ, cache::CredalCache = CredalCache())

    sizehint!(cache, length(circuit)*4÷3)

    pc_node(::LiteralGate, n::ΔNode) = CredalLiteral(n)
    pc_node(::ConstantGate, n::ΔNode) = error("Cannot construct a probabilistic circuit from constant leafs: first smooth and remove unsatisfiable branches.")

    pc_node(::⋀Gate, n::ΔNode) = begin
        children = map(c -> cache[c], n.children)
        Credal⋀(n, children)
    end

    pc_node(::⋁Gate, n::ΔNode) = begin
        children = map(c -> cache[c], n.children)
        Credal⋁(n, children)
    end

    map(circuit) do node
        pcn = pc_node(GateType(node), node)
        cache[node] = pcn
        pcn
    end
end

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

import LogicCircuits: literal, children # make available for extension

@inline literal(n::CredalLiteral)::Lit  = literal(n.origin)
@inline children(n::CredalInnerNode) = n.children

num_parameters(n::Credal⋁) = num_children(n)
num_parameters(c::CredalΔ) = sum(n -> num_parameters(n), ⋁_nodes(c))

"Return the first origin that is a probabilistic circuit node"
prob_origin(n::DecoratorΔNode)::CredalΔNode = origin(n, CredalΔNode)

"Return the first origin that is a probabilistic circuit"
prob_origin(c::DecoratorΔ)::CredalΔ = origin(c, CredalΔNode)

function estimate_credal_parameters2(pc::CredalΔ, data::XData{Bool}, s_idm::Float64; pseudocount::Float64)
    Logical.pass_up_down2(pc, data)
    w = (data isa PlainXData) ? nothing : weights(data)
    estimate_credal_parameters_cached2(pc, w, s_idm; pseudocount=pseudocount)
end

function estimate_credal_parameters_cached2(pc::CredalΔ, w, s_idm::Float64; pseudocount::Float64)
    flow(n) = Float64(sum(sum(n.data)))
    children_flows(n) = sum.(map(c -> c.data[1] .& n.data[1], children(n)))

    if issomething(w)
        flow_w(n) = sum(Float64.(n.data[1]) .* w)
        children_flows_w(n) = sum.(map(c -> Float64.(c.data[1] .& n.data[1]) .* w, children(n)))
        flow = flow_w
        children_flows = children_flows_w
    end

    estimate_credal_parameters_node2(n::CredalΔNode, s_idm::Float64) = ()
    function estimate_credal_parameters_node2(n::Credal⋁, s_idm::Float64)
        if num_children(n) == 1
            n.log_thetas_u .= 0.0
            n.log_thetas .= 0.0
        else
            smoothed_flow = flow(n) + pseudocount
            uniform_pseudocount = pseudocount / num_children(n)
            n.log_thetas .= log.((children_flows(n) .+ uniform_pseudocount)  ./ (smoothed_flow .+ s_idm))
            n.log_thetas_u .= log.(((children_flows(n) .+ uniform_pseudocount) .+ s_idm) ./ (smoothed_flow .+ s_idm))
            # @assert isapprox(sum(exp.(n.log_thetas)), 1.0, atol=1e-6) "Parameters do not sum to one locally"
            # normalize away any leftover error
            # n.log_thetas .- logsumexp(n.log_thetas)
        end
    end

    foreach(estimate_credal_parameters_node2, pc, s_idm)
end

function log_likelihood_per_instance2(pc::CredalΔ, data::XData{Bool})
    Logical.pass_up_down2(pc, data)
    log_likelihood_per_instance_cached(pc, data)
end

function log_likelihood_per_instance_cached(pc::CredalΔ, data::XData{Bool})
    log_likelihoods = zeros(num_examples(data))
    indices = some_vector(Bool, num_examples(data))::BitVector
    for n in pc
         if n isa Credal⋁ && num_children(n) != 1 # other nodes have no effect on likelihood
            foreach(n.children, n.log_thetas) do c, log_theta
                indices = n.data[1] .& c.data[1]
                view(log_likelihoods, indices::BitVector) .+=  log_theta # see MixedProductKernelBenchmark.jl
            end
         end
    end
    log_likelihoods
end

import LogicCircuits: conjoin_like, disjoin_like, literal_like, copy_node, normalize, replace_node # make available for extension

"Conjoin nodes in the same way as the example"
@inline function conjoin_like(example::CredalΔNode, arguments::Vector)
    if isempty(arguments)
        # @assert false "Credalabilistic circuit does not have anonymous true node"
        nothing
    elseif example isa Credal⋀ && children(example) == arguments
        example
    else
        n = conjoin_like(origin(example), origin.(arguments))
        Credal⋀(n, arguments)
    end
end

"Disjoin nodes in the same way as the example"
@inline function disjoin_like(example::CredalΔNode, arguments::Vector)
    if isempty(arguments)
        # @assert false "Probabilistic circuit does not have false node"
        nothing
    elseif example isa Credal⋁ && children(example) == arguments
        example
    else
        n = disjoin_like(origin(example), origin.(arguments))
        # normalize parameters
        thetas = zeros(Float64, length(arguments))
        flag = falses(length(arguments))
        for (i, c) in enumerate(arguments)
            ind = findfirst(x -> x == c, children(example))
            if issomething(ind)
                thetas[i] = exp(example.log_thetas[ind])
                flag[i] = true
            end
        end
        if all(flag)
            thetas = thetas / sum(thetas)
        end
        p = Credal⋁(n, arguments)
        p.log_thetas .= log.(thetas)
        p
    end
end

"Construct a new literal node like the given node's type"
@inline literal_like(::CredalΔNode, lit::Lit) = CredalLiteral(lit)

@inline copy_node(n::Credal⋁, cns) = begin
    orig = copy_node(origin(n), origin.(cns))
    p = Credal⋁(orig, cns)
    p.log_thetas .= copy(n.log_thetas)
    p
end

@inline copy_node(n::Credal⋀, cns) = begin
    orig = copy_node(origin(n), origin.(cns))
    Credal⋀(orig, cns)
end

import LogicCircuits.normalize

@inline normalize(n::Credal⋁, old_n::Credal⋁, kept::Union{Vector{Bool}, BitArray}) = begin
     thetas = exp.(old_n.log_thetas[kept])
     n.log_thetas .= log.(thetas / sum(thetas))
end

function estimate_credal_parameters(pc::CredalΔ, data::XBatches{Bool}, s_idm::Float64; pseudocount::Float64)
    estimate_credal_parameters(AggregateFlowΔ(pc, aggr_weight_type(data)), data, s_idm; pseudocount=pseudocount)
end

function estimate_credal_parameters(afc::AggregateFlowΔ, data::XBatches{Bool}, s_idm::Float64; pseudocount::Float64)
    @assert feature_type(data) == Bool "Can only learn probabilistic circuits on Bool data"
    @assert (afc[end].origin isa CredalΔNode) "AggregateFlowΔ must originate in a CredalΔ"
    collect_aggr_flows(afc, data)
    estimate_credal_parameters_cached(afc, s_idm; pseudocount=pseudocount)
    afc
end

function estimate_credal_parameters(fc::FlowΔ, data::XBatches{Bool}, s_idm::Float64; pseudocount::Float64)
    @assert feature_type(data) == Bool "Can only learn probabilistic circuits on Bool data"
    @assert (prob_origin(afc[end]) isa CredalΔNode) "FlowΔ must originate in a CredalΔ"
    collect_aggr_flows(fc, data)
    estimate_credal_parameters_cached(origin(fc), s_idm; pseudocount=pseudocount)
end

 # turns aggregate statistics into theta parameters
function estimate_credal_parameters_cached(afc::AggregateFlowΔ, s_idm::Float64; pseudocount::Float64)
    foreach(n -> estimate_credal_parameters_node(n, s_idm; pseudocount=pseudocount), afc)
end

estimate_credal_parameters_node(::AggregateFlowΔNode, s_idm::Float64; pseudocount::Float64) = () # do nothing
function estimate_credal_parameters_node(n::AggregateFlow⋁, s_idm::Float64; pseudocount)
    origin = n.origin::Credal⋁
    if num_children(n) == 1
        origin.log_thetas_u .= 0.0
        origin.log_thetas .= 0.0
    else
        smoothed_aggr_flow = (n.aggr_flow + pseudocount)
        uniform_pseudocount = pseudocount / num_children(n)
        origin.log_thetas .= log.( (n.aggr_flow_children .+ uniform_pseudocount) ./ (smoothed_aggr_flow .+ s_idm ))
        origin.log_thetas_u .= log.( ((n.aggr_flow_children .+ uniform_pseudocount) .+ s_idm) ./ (smoothed_aggr_flow .+ s_idm ))
        # @assert isapprox(sum(exp.(origin.log_thetas)), 1.0, atol=1e-6) "Parameters do not sum to one locally: $(exp.(origin.log_thetas)), estimated from $(n.aggr_flow) and $(n.aggr_flow_children). Did you actually compute the aggregate flows?"
        #normalize away any leftover error ()
        # origin.log_thetas .- logsumexp(origin.log_thetas)
    end
end

# compute log likelihood
function compute_log_likelihood(pc::CredalΔ, data::XBatches{Bool})
    compute_log_likelihood(AggregateFlowΔ(pc, aggr_weight_type(data)))
end

# compute log likelihood, reusing AggregateFlowΔ but ignoring its current aggregate values
function compute_log_likelihood(afc::AggregateFlowΔ, data::XBatches{Bool})
    @assert feature_type(data) == Bool "Can only test probabilistic circuits on Bool data"
    collect_aggr_flows(afc, data)
    ll = log_likelihood(afc)
    (afc, ll)
end

# return likelihoods given current aggregate flows.
function log_likelihood(afc::AggregateFlowΔ)
    sum(n -> log_likelihood(n), afc)
end

log_likelihood(::AggregateFlowΔNode) = 0.0
log_likelihood(n::AggregateFlow⋁) = sum(n.origin.log_thetas .* n.aggr_flow_children)

"""
Calculates log likelihood for a batch of fully observed samples.
(Also retures the generated FlowΔ)
"""
function log_likelihood_upper_per_instance(pc::CredalΔ, batch::PlainXData{Bool})
    fc = FlowΔ(pc, num_examples(batch), Bool)
    (fc, log_likelihood_upper_per_instance(fc, batch))
end

function log_likelihood_lower_per_instance(pc::CredalΔ, batch::PlainXData{Bool})
    fc = FlowΔ(pc, num_examples(batch), Bool)
    (fc, log_likelihood_lower_per_instance(fc, batch))
end
"""
Calculate log likelihood per instance for batches of samples.
"""
function log_likelihood_per_instance(pc::CredalΔ, batches::XBatches{Bool})::Vector{Float64}
    mapreduce(b -> log_likelihood_per_instance(pc, b)[2], vcat, batches)
end

function log_prob_upper(pc::CredalΔ, batch::PlainXData{Bool})
    log_likelihood_upper_per_instance(pc, batch)[2]
end

function log_prob_lower(pc::CredalΔ, batch::PlainXData{Bool})
    log_likelihood_lower_per_instance(pc, batch)[2]
end

function log_marginal_upper(pc::CredalΔ, batch::PlainXData{Int8})
    marginal_log_likelihood_upper_per_instance(pc, batch)[2]
end

function log_marginal_lower(pc::CredalΔ, batch::PlainXData{Int8})
    marginal_log_likelihood_lower_per_instance(pc, batch)[2]
end

"""
Calculate log likelihood for a batch of fully observed samples.
(This is for when you already have a FlowΔ)
"""
function log_likelihood_upper_per_instance(fc::FlowΔ, batch::PlainXData{Bool})
    @assert (prob_origin(fc[end]) isa CredalΔNode) "FlowΔ must originate in a CredalΔ"
    pass_up_down(fc, batch)
    log_likelihoods = zeros(num_examples(batch))
    indices = some_vector(Bool, flow_length(fc))::BitVector
    for n in fc
         if n isa DownFlow⋁ && num_children(n) != 1 # other nodes have no effect on likelihood
            origin = prob_origin(n)::Credal⋁
            foreach(n.children, origin.log_thetas) do c, log_theta
                #  be careful here to allow for the Boolean multiplication to be done using & before switching to float arithmetic, or risk losing a lot of runtime!
                # log_likelihoods .+= prod_fast(downflow(n), pr_factors(c)) .* log_theta
                assign_prod(indices, downflow(n), pr_factors(c))
                view(log_likelihoods, indices::BitVector) .+=  log_theta # see MixedProductKernelBenchmark.jl
                # TODO put the lines above in Utils in order to ensure we have specialized types
            end
         end
    end
    log_likelihoods
end

function log_likelihood_lower_per_instance(fc::FlowΔ, batch::PlainXData{Bool})
    @assert (prob_origin(fc[end]) isa CredalΔNode) "FlowΔ must originate in a CredalΔ"
    pass_up_down(fc, batch)
    log_likelihoods = zeros(num_examples(batch))
    indices = some_vector(Bool, flow_length(fc))::BitVector
    for n in fc
         if n isa DownFlow⋁ && num_children(n) != 1 # other nodes have no effect on likelihood
            origin = prob_origin(n)::Credal⋁
            foreach(n.children, origin.log_thetas) do c, log_theta
                #  be careful here to allow for the Boolean multiplication to be done using & before switching to float arithmetic, or risk losing a lot of runtime!
                # log_likelihoods .+= prod_fast(downflow(n), pr_factors(c)) .* log_theta
                assign_prod(indices, downflow(n), pr_factors(c))
                view(log_likelihoods, indices::BitVector) .+=  log_theta # see MixedProductKernelBenchmark.jl
                # TODO put the lines above in Utils in order to ensure we have specialized types
            end
         end
    end
    log_likelihoods
end

"""
Calculate the upper log likelihood for a batch of samples with partial evidence P(e).
(Also returns the generated FlowΔ)

To indicate a variable is not observed, pass -1 for that variable.
"""
function marginal_log_likelihood_upper_per_instance(pc::CredalΔ, batch::PlainXData{Int8})
    opts = (flow_opts★..., el_type=Float64, compact⋀=false, compact⋁=false)
    fc = UpFlowΔ(pc, num_examples(batch), Float64, opts)
    (fc, marginal_log_likelihood_upper_per_instance(fc, batch))
end

"""
Calculate the lower log likelihood for a batch of samples with partial evidence P(e).
(Also returns the generated FlowΔ)

To indicate a variable is not observed, pass -1 for that variable.
"""
function marginal_log_likelihood_lower_per_instance(pc::CredalΔ, batch::PlainXData{Int8})
    opts = (flow_opts★..., el_type=Float64, compact⋀=false, compact⋁=false)
    fc = UpFlowΔ(pc, num_examples(batch), Float64, opts)
    (fc, marginal_log_likelihood_lower_per_instance(fc, batch))
end

"""
Calculate the upper log likelihood for a batch of samples with partial evidence P(e).
(If you already have a FlowΔ)

To indicate a variable is not observed, pass -1 for that variable.
"""
function marginal_log_likelihood_upper_per_instance(fc::UpFlowΔ, batch::PlainXData{Int8})
    @assert (prob_origin(fc[end]) isa CredalΔNode) "FlowΔ must originate in a CredalΔ"
    credal_marginal_upper_pass_up(fc, batch)
    pr(fc[end])
end

"""
Calculate the lower log likelihood for a batch of samples with partial evidence P(e).
(If you already have a FlowΔ)

To indicate a variable is not observed, pass -1 for that variable.
"""
function marginal_log_likelihood_lower_per_instance(fc::UpFlowΔ, batch::PlainXData{Int8})
    @assert (prob_origin(fc[end]) isa CredalΔNode) "FlowΔ must originate in a CredalΔ"
    credal_marginal_lower_pass_up(fc, batch)
    pr(fc[end])
end

"""
Calculate conditional inference upper bound
"""

function conditional_upper(pc::CredalΔ, batch::PlainXData{Int8})
    conditional_upper_per_instance(pc, batch)[2]
end

"""
Calculate conditional inference lower bound
"""

function conditional_lower(pc::CredalΔ, batch::PlainXData{Int8}, mu::Array{Float64,1})
    conditional_lower_per_instance(pc, batch, mu)[2]
end

"""
Calculate conditional inference upper bound per instance
"""
function conditional_upper_per_instance(pc::CredalΔ, batch::PlainXData{Int8})
    opts = (compact⋀=false, compact⋁=false)
    #Initialize the flows as a tuple instead of a single value
    fc = UpFlowΔ(pc, num_examples(batch), Tuple{Float64,Float64,Float64,Float64},opts);
    (fc, conditional_upper_instance(fc, batch))
end

"""
Calculate conditional inference upper bound per instance
"""
function conditional_lower_per_instance(pc::CredalΔ, batch::PlainXData{Int8}, mu::Array{Float64,1})
    opts = (compact⋀=false, compact⋁=false)
    #Initialize the flows as a tuple instead of a single value
    fc = UpFlowΔ(pc, num_examples(batch), Tuple{Float64,Float64,Float64,Float64},opts);
    (fc, conditional_lower_instance(fc, batch, mu))
end

"""
Calculate conditional inference upper bound per instance

"""
function conditional_upper_instance(fc::UpFlowΔ, batch::PlainXData{Int8})
    # @assert (prob_origin(fc[end]) isa ProbΔNode) "FlowΔ must originate in a ProbΔ"
    conditional_upper_pass_up(fc, batch)
    pr(fc[end])
end

"""
Calculate conditional inference lowe bound per instance

"""
function conditional_lower_instance(fc::UpFlowΔ, batch::PlainXData{Int8}, mu::Array{Float64,1})
    # @assert (prob_origin(fc[end]) isa ProbΔNode) "FlowΔ must originate in a ProbΔ"
    conditional_lower_pass_up(fc, batch, mu)
    pr(fc[end])
end
"""
# function check_parameter_integrity(circuit::CredalΔ)
#     for node in filter(n -> GateType(n) isa Credal⋁, circuit)
#         @assert all(θ -> !isnan(θ), node.log_thetas) "There is a NaN in one of the log_thetas"
#     end
#     true
# end

# ##################
# # Sampling from a psdd
# ##################

# """
# Sample from a PSDD without any evidence
# """
# function sample(circuit::CredalΔ)::AbstractVector{Bool}
#     inst = Dict{Var,Int64}()
#     simulate(circuit[end], inst)
#     len = length(keys(inst))
#     ans = Vector{Bool}()
#     for i = 1:len
#         push!(ans, inst[i])
#     end
#     ans
# end

# # Uniformly sample based on the probability of the items
# # and return the selected index
# function sample(probs::AbstractVector{<:Number})::Int32
#     z = sum(probs)
#     q = rand() * z
#     cur = 0.0
#     for i = 1:length(probs)
#         cur += probs[i]
#         if q <= cur
#             return i
#         end
#     end
#     return length(probs)
# end

# function simulate(node::CredalLiteral, inst::Dict{Var,Int64})
#     if positive(node)
#         inst[variable(node.origin)] = 1
#     else
#         inst[variable(node.origin)] = 0
#     end
# end

# function simulate(node::Credal⋁, inst::Dict{Var,Int64})
#     idx = sample(exp.(node.log_thetas))
#     simulate(node.children[idx], inst)
# end
# function simulate(node::Credal⋀, inst::Dict{Var,Int64})
#     for child in node.children
#         simulate(child, inst)
#     end
# end

# """
# Sampling with Evidence from a psdd.
# Internally would call marginal pass up on a newly generated flow circuit.
# """
# function sample(circuit::CredalΔ, evidence::PlainXData{Int8})::AbstractVector{Bool}
#     opts= (compact⋀=false, compact⋁=false)
#     flow_circuit = UpFlowΔ(circuit, 1, Float64, opts)
#     marginal_pass_up(flow_circuit, evidence)
#     sample(flow_circuit)
# end

# """
# Sampling with Evidence from a psdd.
# Assuming already marginal pass up has been done on the flow circuit.
# """
# function sample(circuit::UpFlowΔ)::AbstractVector{Bool}
#     inst = Dict{Var,Int64}()
#     simulate2(circuit[end], inst)
#     len = length(keys(inst))
#     ans = Vector{Bool}()
#     for i = 1:len
#         push!(ans, inst[i])
#     end
#     ans
# end

# function simulate2(node::UpFlowLiteral, inst::Dict{Var,Int64})
#     if positive(node)
#         #TODO I don't think we need these 'grand_origin' parts below
#         inst[variable(grand_origin(node))] = 1
#     else
#         inst[variable(grand_origin(node))] = 0
#     end
# end

# function simulate2(node::UpFlow⋁, inst::Dict{Var,Int64})
#     prs = [ pr(ch)[1] for ch in children(node) ]
#     idx = sample(exp.(node.origin.log_thetas .+ prs))
#     simulate2(children(node)[idx], inst)
# end

# function simulate2(node::UpFlow⋀, inst::Dict{Var,Int64})
#     for child in children(node)
#         simulate2(child, inst)
#     end
# end

# ##################
# # Most Probable Explanation MPE of a psdd
# #   aka MAP
# ##################

# @inline function MAP(circuit::CredalΔ, evidence::PlainXData{Int8})::Matrix{Bool}
#     MPE(circuit, evidence)
# end

# function MPE(circuit::CredalΔ, evidence::PlainXData{Int8})::Matrix{Bool}
#     # Computing Marginal Likelihood for each node
#     fc, lls = marginal_log_likelihood_per_instance(circuit, evidence)

#     ans = Matrix{Bool}(zeros(size(evidence.x)))
#     active_samples = Array{Bool}(ones( num_examples(evidence) ))

#     mpe_simulate(fc[end], active_samples, ans)
#     ans
# end

# """
# active_samples: bool vector indicating which samples are active for this node during mpe
# result: Matrix (num_samples, num_variables) indicating the final result of mpe
# """
# function mpe_simulate(node::UpFlowLiteral, active_samples::Vector{Bool}, result::Matrix{Bool})
#     if positive(node)
#         result[active_samples, variable(node)] .= 1
#     else
#         result[active_samples, variable(node)] .= 0
#     end
# end
# function mpe_simulate(node::UpFlow⋁, active_samples::Vector{Bool}, result::Matrix{Bool})
#     prs = zeros( length(node.children), size(active_samples)[1] )
#     @simd  for i=1:length(node.children)
#         prs[i,:] .= pr(node.children[i]) .+ (node.origin.log_thetas[i])
#     end

#     max_child_ids = [a[1] for a in argmax(prs, dims = 1) ]
#     @simd for i=1:length(node.children)
#         ids = Vector{Bool}( active_samples .* (max_child_ids .== i)[1,:] )  # Only active for this child if it was the max for that sample
#         mpe_simulate(node.children[i], ids, result)
#     end
# end
# function mpe_simulate(node::UpFlow⋀, active_samples::Vector{Bool}, result::Matrix{Bool})
#     for child in node.children
#         mpe_simulate(child, active_samples, result)
#     end
# end


## wrappers for multi-gpu functionality in cusolverMg

# auxiliary functionality

# NOTE: in the cublasMg preview, which also relies on this functionality, a separate library
#       called 'cudalibmg' is introduced. factor this out when we actually ship that.

mutable struct MatrixDescriptor
    desc::cudaLibMgMatrixDesc_t

    function MatrixDescriptor(a, grid; rowblocks = size(a, 1), colblocks = size(a, 2), elta=eltype(a) )
        desc = Ref{cudaLibMgMatrixDesc_t}()
        cusolverMgCreateMatrixDesc(desc, size(a, 1), size(a, 2), rowblocks, colblocks, elta, grid)
        return new(desc[])
    end
end

Base.unsafe_convert(::Type{cudaLibMgMatrixDesc_t}, obj::MatrixDescriptor) = obj.desc

mutable struct DeviceGrid
    desc::cudaLibMgGrid_t

    function DeviceGrid(num_row_devs, num_col_devs, deviceIds, mapping)
        @assert num_row_devs == 1 "Only 1-D column block cyclic is supported, so numRowDevices must be equal to 1."
        desc = Ref{cudaLibMgGrid_t}()
        cusolverMgCreateDeviceGrid(desc, num_row_devs, num_col_devs, deviceIds, mapping)
        return new(desc[])
    end
end

Base.unsafe_convert(::Type{cudaLibMgGrid_t}, obj::DeviceGrid) = obj.desc

function allocateBuffers(n_row_devs, n_col_devs, mat::Matrix)
    mat_row_block_size = div(size(mat, 1), n_row_devs)
    mat_col_block_size = div(size(mat, 2), n_col_devs)
    mat_buffers  = Vector{CuMatrix{eltype(mat)}}(undef, ndevices())
    mat_numRows  = Vector{Int64}(undef, ndevices())
    mat_numCols  = Vector{Int64}(undef, ndevices())
    typesize = sizeof(eltype(mat))
    ldas = Vector{Int64}(undef, ndevices())
    current_dev = device()
    mat_cpu_bufs = Vector{Matrix{eltype(mat)}}(undef, ndevices())
    for (di, dev) in enumerate(devices())
        ldas[di]    = mat_col_block_size
        dev_row     = mod(di - 1, n_row_devs) + 1
        dev_col     = div(di - 1, n_row_devs) + 1

        mat_row_inds     = ((dev_row-1)*mat_row_block_size+1):min(dev_row*mat_row_block_size, size(mat, 1))
        mat_col_inds     = ((dev_col-1)*mat_col_block_size+1):min(dev_col*mat_col_block_size, size(mat, 2))
        mat_cpu_bufs[di] = Array(mat[mat_row_inds, mat_col_inds])
    end
    for (di, dev) in enumerate(devices())
        device!(dev)
        mat_gpu_buf = CuMatrix{eltype(mat)}(undef, size(mat_cpu_bufs[di]))
        copyto!(mat_gpu_buf, mat_cpu_bufs[di])
        mat_buffers[di] = mat_gpu_buf
    end
    for (di, dev) in enumerate(devices())
        device!(dev)
        synchronize()
    end
    device!(current_dev)
    return mat_buffers
end

function returnBuffers(n_row_devs, n_col_devs, row_block_size, col_block_size, dDs, D)
    row_block_size = div(size(D, 1), n_row_devs)
    col_block_size = div(size(D, 2), n_col_devs)
    numRows  = [row_block_size for dev in 1:ndevices()]
    numCols  = [col_block_size for dev in 1:ndevices()]
    typesize = sizeof(eltype(D))
    current_dev = device()
    cpu_bufs = Vector{Matrix{eltype(D)}}(undef, ndevices())
    for (di, dev) in enumerate(devices())
        device!(dev)
        dev_row = mod(di - 1, n_row_devs) + 1
        dev_col = div(di - 1, n_row_devs) + 1
        row_inds = ((dev_row-1)*row_block_size+1):min(dev_row*row_block_size, size(D, 1))
        col_inds = ((dev_col-1)*col_block_size+1):min(dev_col*col_block_size, size(D, 2))
        cpu_bufs[di] = Matrix{eltype(D)}(undef, length(row_inds), length(col_inds))
        copyto!(cpu_bufs[di], dDs[di])
    end
    for (di, dev) in enumerate(devices())
        device!(dev)
        synchronize()
        dev_row = mod(di - 1, n_row_devs) + 1
        dev_col = div(di - 1, n_row_devs) + 1
        row_inds = ((dev_row-1)*row_block_size+1):min(dev_row*row_block_size, size(D, 1))
        col_inds = ((dev_col-1)*col_block_size+1):min(dev_col*col_block_size, size(D, 2))
        D[row_inds, col_inds] = cpu_bufs[di]
    end
    device!(current_dev)
    return D
end

## wrappers

function mg_syevd!(jobz::Char, uplo::Char, A; dev_rows=1, dev_cols=ndevices()) # one host-side array A
    dev = device()
    grid = DeviceGrid(1, ndevices(), devices(), CUDALIBMG_GRID_MAPPING_COL_MAJOR)
    if uplo != 'L'
        throw(ArgumentError("only lower fill mode (uplo='L') supported"))
    end
    m, n    = size(A)
    N       = div(size(A, 2), ndevices()) # dimension of the sub-matrix
    lwork         = Ref{Int64}(0)
    workspace     = Vector{CuArray}(undef, ndevices())
    W             = Vector{real(eltype(A))}(undef, n)
    desc = MatrixDescriptor(A, grid; colblocks=N)
    A_arr     = allocateBuffers(dev_rows, dev_cols, A)
    IA            = 1 # for now
    JA            = 1
    GC.@preserve A_arr begin
        cusolverMgSyevd_bufferSize(mg_handle(), jobz, uplo, n, pointer.(A_arr), IA, JA, desc, W, real(eltype(A)), eltype(A), lwork)
    end
    for (di, dev) in enumerate(devices())
        device!(dev)
        workspace[di] = CUDA.zeros(eltype(A), lwork[])
        synchronize()
    end
    device!(dev)
    info = Ref{Cint}(C_NULL)
    GC.@preserve A_arr workspace begin
        cusolverMgSyevd(mg_handle(), jobz, uplo, n, pointer.(A_arr), IA, JA, desc, W, real(eltype(A)), eltype(A), pointer.(workspace), lwork[], info)
    end
    chkargsok(BlasInt(info[]))
    A = returnBuffers(dev_rows, dev_cols, div(size(A, 1), dev_rows), div(size(A, 2), dev_cols), A_arr, A)
    if jobz == 'N'
        return W
    elseif jobz == 'V'
        return W, A
    end
end

function mg_potrf!(uplo::Char, A; dev_rows=1, dev_cols=ndevices()) # one host-side array A
    dev = device()
    grid = DeviceGrid(1, ndevices(), devices(), CUDALIBMG_GRID_MAPPING_COL_MAJOR)
    if uplo != 'L'
        throw(ArgumentError("only lower fill mode (uplo='L') supported"))
    end
    m, n    = size(A)
    N       = div(size(A, 2), ndevices()) # dimension of the sub-matrix
    lwork         = Ref{Int64}(0)
    workspace     = Vector{CuArray}(undef, ndevices())
    desc = MatrixDescriptor(A, grid; colblocks=N)
    A_arr     = allocateBuffers(dev_rows, dev_cols, A)
    IA      = 1 # for now
    JA      = 1
    GC.@preserve A_arr begin
        cusolverMgPotrf_bufferSize(mg_handle(), uplo, n, pointer.(A_arr), IA, JA, desc, eltype(A), lwork)
    end
    for (di, dev) in enumerate(devices())
        device!(dev)
        workspace[di]     = CUDA.zeros(eltype(A), lwork[])
        synchronize()
    end
    device!(dev)
    info = Ref{Cint}(C_NULL)
    GC.@preserve A_arr workspace begin
        cusolverMgPotrf(mg_handle(), uplo, n, pointer.(A_arr), IA, JA, desc, eltype(A), pointer.(workspace), lwork[], info)
    end
    chkargsok(BlasInt(info[]))

    returnBuffers(dev_rows, dev_cols, div(size(A, 1), dev_rows), div(size(A, 2), dev_cols), A_arr, A)
end

function mg_potri!(uplo::Char, A; dev_rows=1, dev_cols=ndevices()) # one host-side array A
    dev = device()
    grid = DeviceGrid(1, ndevices(), devices(), CUDALIBMG_GRID_MAPPING_COL_MAJOR)
    if uplo != 'L'
        throw(ArgumentError("only lower fill mode (uplo='L') supported"))
    end
    m, n    = size(A)
    N       = div(size(A, 2), ndevices()) # dimension of the sub-matrix
    lwork         = Ref{Int64}(0)
    workspace     = Vector{CuArray}(undef, ndevices())
    desc = MatrixDescriptor(A, grid; colblocks=N)
    A_arr     = allocateBuffers(dev_rows, dev_cols, A)
    IA      = 1 # for now
    JA      = 1
    GC.@preserve A_arr begin
        cusolverMgPotri_bufferSize(mg_handle(), uplo, n, pointer.(A_arr), IA, JA, desc, eltype(A), lwork)
    end
    for (di, dev) in enumerate(devices())
        device!(dev)
        workspace[di]     = CUDA.zeros(eltype(A), lwork[])
        synchronize()
    end
    device!(dev)
    info = Ref{Cint}(C_NULL)
    GC.@preserve A_arr workspace begin
        cusolverMgPotri(mg_handle(), uplo, n, pointer.(A_arr), IA, JA, desc, eltype(A), pointer.(workspace), lwork[], info)
    end
    chkargsok(BlasInt(info[]))

    returnBuffers(dev_rows, dev_cols, div(size(A, 1), dev_rows), div(size(A, 2), dev_cols), A_arr, A)
end

function mg_potrs!(uplo::Char, A, B; dev_rows=1, dev_cols=ndevices()) # one host-side array A
    dev = device()
    grid = DeviceGrid(1, ndevices(), devices(), CUDALIBMG_GRID_MAPPING_COL_MAJOR)
    if uplo != 'L'
        throw(ArgumentError("only lower fill mode (uplo='L') supported"))
    end
    ma, na   = size(A)
    mb, nb   = size(B)
    NA       = div(size(A, 2), ndevices()) # dimension of the sub-matrix
    NB       = div(size(B, 2), ndevices()) # dimension of the sub-matrix
    lwork         = Ref{Int64}(0)
    workspace     = Vector{CuArray}(undef, ndevices())
    descA = MatrixDescriptor(A, grid; colblocks=NA)
    descB = MatrixDescriptor(A, grid; colblocks=NB)
    A_arr     = allocateBuffers(dev_rows, dev_cols, A)
    B_arr     = allocateBuffers(dev_rows, dev_cols, B)
    IA      = 1 # for now
    JA      = 1
    IB      = 1 # for now
    JB      = 1
    GC.@preserve A_arr B_arr begin
        cusolverMgPotrs_bufferSize(mg_handle(), uplo, na, nb, pointer.(A_arr), IA, JA, descA, pointer.(B_arr), IB, JB, descB, eltype(A), lwork)
    end
    for (di, dev) in enumerate(devices())
        device!(dev)
        workspace[di]     = CUDA.zeros(eltype(A), lwork[])
        synchronize()
    end
    device!(dev)
    info = Ref{Cint}(C_NULL)
    GC.@preserve A_arr B_arr workspace begin
        cusolverMgPotrs(mg_handle(), uplo, na, nb, pointer.(A_arr), IA, JA, descA, pointer.(B_arr), IB, JB, descB, eltype(A), pointer.(workspace), lwork[], info)
    end
    chkargsok(BlasInt(info[]))

    returnBuffers(dev_rows, dev_cols, div(size(B, 1), dev_rows), div(size(B, 2), dev_cols), B_arr, B)
end

function mg_getrf!(A; dev_rows=1, dev_cols=ndevices()) # one host-side array A
    dev = device()
    grid = DeviceGrid(1, ndevices(), devices(), CUDALIBMG_GRID_MAPPING_COL_MAJOR)
    m, n    = size(A)
    N       = div(size(A, 2), ndevices()) # dimension of the sub-matrix
    lwork         = Ref{Int64}(0)
    ipivs         = Vector{CuVector{Cint}}(undef, ndevices())
    workspace     = Vector{CuArray}(undef, ndevices())
    desc = MatrixDescriptor(A, grid; colblocks=N)
    A_arr     = allocateBuffers(dev_rows, dev_cols, A)
    IA      = 1 # for now
    JA      = 1
    for (di, dev) in enumerate(devices())
        device!(dev)
        ipivs[di]     = CUDA.zeros(Cint, N)
        synchronize()
    end
    device!(dev)
    GC.@preserve A_arr ipivs begin
        cusolverMgGetrf_bufferSize(mg_handle(), m, n, pointer.(A_arr), IA, JA, desc, pointer.(ipivs), eltype(A), lwork)
    end
    device_synchronize()
    for (di, dev) in enumerate(devices())
        device!(dev)
        workspace[di]     = CUDA.zeros(eltype(A), lwork[])
        synchronize()
    end
    device!(dev)
    info = Ref{Cint}(C_NULL)
    GC.@preserve A_arr ipivs workspace begin
        cusolverMgGetrf(mg_handle(), m, n, pointer.(A_arr), IA, JA, desc, pointer.(ipivs), eltype(A), pointer.(workspace), lwork[], info)
    end
    device_synchronize()
    chkargsok(BlasInt(info[]))

    A = returnBuffers(dev_rows, dev_cols, div(size(A, 1), dev_rows), div(size(A, 2), dev_cols), A_arr, A)
    ipiv = Vector{Int}(undef, n)
    for (di, dev) in enumerate(devices())
        device!(dev)
        ipiv[((di-1)*N + 1):min((di*N), n)] = collect(ipivs[di])
    end
    device!(dev)
    return A, ipiv
end

function mg_getrs!(trans, A, ipiv, B; dev_rows=1, dev_cols=ndevices()) # one host-side array A
    dev = device()
    grid = DeviceGrid(1, ndevices(), devices(), CUDALIBMG_GRID_MAPPING_COL_MAJOR)
    ma, na  = size(A)
    mb, nb  = size(B)
    NA      = div(size(A, 2), ndevices()) # dimension of the sub-matrix
    NB      = div(size(B, 2), ndevices()) # dimension of the sub-matrix
    lwork         = Ref{Int64}(0)
    ipivs         = Vector{CuVector{Cint}}(undef, ndevices())
    workspace     = Vector{CuArray}(undef, ndevices())
    descA = MatrixDescriptor(A, grid; colblocks=NA)
    descB = MatrixDescriptor(A, grid; colblocks=NB)
    A_arr     = allocateBuffers(dev_rows, dev_cols, A)
    B_arr     = allocateBuffers(dev_rows, dev_cols, B)
    IA      = 1 # for now
    JA      = 1
    IB      = 1 # for now
    JB      = 1
    for (di, dev) in enumerate(devices())
        device!(dev)
        local_ipiv    = Cint.(ipiv[(di-1)*NA+1:min(di*NA,length(ipiv))])
        ipivs[di]     = CuArray(local_ipiv)
        synchronize()
    end
    device!(dev)
    GC.@preserve A_arr B_arr ipivs begin
        cusolverMgGetrs_bufferSize(mg_handle(), trans, na, nb, pointer.(A_arr), IA, JA, descA, pointer.(ipivs), pointer.(B_arr), IB, JB, descB, eltype(A), lwork)
    end
    for (di, dev) in enumerate(devices())
        device!(dev)
        workspace[di]     = CUDA.zeros(eltype(A), lwork[])
        synchronize()
    end
    device!(dev)
    info = Ref{Cint}(C_NULL)
    GC.@preserve A_arr B_arr ipivs workspace begin
        cusolverMgGetrs(mg_handle(), trans, na, nb, pointer.(A_arr), IA, JA, descA, pointer.(ipivs), pointer.(B_arr), IB, JB, descB, eltype(A), pointer.(workspace), lwork[], info)
    end
    chkargsok(BlasInt(info[]))

    returnBuffers(dev_rows, dev_cols, div(size(B, 1), dev_rows), div(size(B, 2), dev_cols), B_arr, B)
end


module BART

  using Distributed
  using Distributions
  using LinearAlgebra
  using MCMCChains
  using StatsBase
  using SpecialFunctions

  import DecisionTree

  include("trees.jl")
  include("models.jl")
  include("proposals.jl")
  include("predict.jl")
  include("fit.jl")

  export
    BartModel,
    Chains,
    Hypers,
    Opts,
    TrainData

  export
    depth,
    fit,
    ptd,
    predict,
    update

end


export DenseConvDims

"""
    DenseConvDims

Concrete subclass of `ConvDims` for a normal, dense, conv2d/conv3d.
"""
struct DenseConvDims{N,K,C_in,C_out,S,P,D,F} <: ConvDims{N,S,P,D,F}
    I::NTuple{N,Int}
end

# Getters for the fields
input_size(c::DenseConvDims) = c.I
kernel_size(c::DenseConvDims{N,K,C_in,C_out}) where {N,K,C_in,C_out} = K
channels_in(c::DenseConvDims{N,K,C_in,C_out}) where {N,K,C_in,C_out} = C_in::Int
channels_out(c::DenseConvDims{N,K,C_in,C_out}) where {N,K,C_in,C_out} = C_out::Int

# Convenience wrapper to create DenseConvDims objects
function DenseConvDims(x_size::NTuple{M}, w_size::NTuple{M};
                       stride=1, padding=0, dilation=1, flipkernel::Bool=false) where M
    # Do common parameter validation
    stride, padding, dilation = check_spdf(x_size, w_size, stride, padding, dilation)

    # Ensure channels are equal
    if x_size[end-1] != w_size[end-1]
        xs = x_size[end-1]
        ws = w_size[end-1]
        throw(DimensionMismatch("Input channels must match! ($xs vs. $ws)"))
    end

    # The type parameters are what
    return DenseConvDims{
        M - 2,
        w_size[1:end-2],
        x_size[end-1],
        w_size[end],
        stride,
        padding,
        dilation,
        flipkernel
    }(
        # Input spatial size
        x_size[1:end-2],
    )
end

# Auto-extract sizes and sub out to big brother above
function DenseConvDims(x::AbstractArray, w::AbstractArray; kwargs...)
    if ndims(x) != ndims(w)
        throw(DimensionMismatch("Rank of x and w must match! ($(ndims(x)) vs. $(ndims(w)))"))
    end
    return DenseConvDims(size(x), size(w); kwargs...)
end

# Useful for constructing a new DenseConvDims that has only a few elements different
# from the original progenitor object that it inherits shapes from.
function DenseConvDims(c::ConvDims; N=spatial_dims(c), I=input_size(c), K=kernel_size(c),
                       C_in=channels_in(c), C_out=channels_out(c), S=stride(c),
                       P=padding(c), D=dilation(c), F=flipkernel(c))
    return DenseConvDims{N, K, C_in, C_out, S, P, D, F}(I)
end

function check_dims(x::NTuple{M}, w::NTuple{M}, y::NTuple{M}, cdims::DenseConvDims) where {M}
    # First, check that channel counts are all correct:
    @assert x[M-1] == channels_in(cdims) DimensionMismatch("Data input channel count ($(x[M-1]) vs. $(channels_in(cdims)))")
    @assert y[M-1] == channels_out(cdims) DimensionMismatch("Data output channel count ($(y[M-1]) vs. $(channels_out(cdims)))")
    @assert w[M-1] == channels_in(cdims) DimensionMismatch("Kernel input channel count ($(w[M-1]) vs. $(channels_in(cdims)))")
    @assert w[M] == channels_out(cdims) DimensionMismatch("Kernel output channel count ($(w[M]) vs. $(channels_out(cdims)))")

    # Next, check that the spatial dimensions match up
    @assert x[1:M-2] == input_size(cdims) DimensionMismatch("Data input spatial size ($(x[1:M-2]) vs. $(input_size(cdims)))")
    @assert y[1:M-2] == output_size(cdims) DimensionMismatch("Data output spatial size ($(y[1:M-2]) vs. $(output_size(cdims)))")
    @assert w[1:M-2] == kernel_size(cdims) DimensionMismatch("Kernel spatial size ($(w[1:M-2]) vs. $(kernel_size(cdims)))")

    # Finally, check that the batch size matches
    @assert x[M] == y[M] DimensionMismatch("Batch size ($(x[M]) vs. $(y[M]))")
end


# --------------------------------------------------------------------------
# ACE.jl: Julia implementation of the Atomic Cluster Expansion
# Copyright (c) 2019 Christoph Ortner <christophortner0@gmail.com>
# All rights reserved.
# --------------------------------------------------------------------------

module Bonds

include("cylindrical_coords.jl")
include("fourier.jl")

include("bondalist.jl")

include("bond_basis.jl")

# include("bond_eval.jl")

include("utils.jl")

end


using CLIMAParameters.Planet: R_d, cv_d, T_0, e_int_v0, e_int_i0

"""
    linearized_air_pressure(ρ, ρe_tot, ρe_pot, ρq_tot=0, ρq_liq=0, ρq_ice=0)

The air pressure, linearized around a dry rest state, from the equation of state
(ideal gas law) where:

 - `ρ` (moist-)air density
 - `ρe_tot` total energy density
 - `ρe_pot` potential energy density
and, optionally,
 - `ρq_tot` total water density
 - `ρq_liq` liquid water density
 - `ρq_ice` ice density
"""
function linearized_air_pressure(
    param_set::AbstractParameterSet,
    ρ::FT,
    ρe_tot::FT,
    ρe_pot::FT,
    ρq_tot::FT = FT(0),
    ρq_liq::FT = FT(0),
    ρq_ice::FT = FT(0),
) where {FT <: Real, PS}
    _R_d::FT = R_d(param_set)
    _cv_d::FT = cv_d(param_set)
    _T_0::FT = T_0(param_set)
    _e_int_v0::FT = e_int_v0(param_set)
    _e_int_i0::FT = e_int_i0(param_set)
    return ρ * _R_d * _T_0 +
           _R_d / _cv_d * (
        ρe_tot - ρe_pot - (ρq_tot - ρq_liq) * _e_int_v0 +
        ρq_ice * (_e_int_i0 + _e_int_v0)
    )
end

@inline function linearized_pressure(
    ::DryModel,
    param_set::AbstractParameterSet,
    orientation::Orientation,
    state::Vars,
    aux::Vars,
)
    ρe_pot = state.ρ * gravitational_potential(orientation, aux)
    return linearized_air_pressure(param_set, state.ρ, state.ρe, ρe_pot)
end
@inline function linearized_pressure(
    ::EquilMoist,
    param_set::AbstractParameterSet,
    orientation::Orientation,
    state::Vars,
    aux::Vars,
)
    ρe_pot = state.ρ * gravitational_potential(orientation, aux)
    linearized_air_pressure(
        param_set,
        state.ρ,
        state.ρe,
        ρe_pot,
        state.moisture.ρq_tot,
    )
end
@inline function linearized_pressure(
    ::NonEquilMoist,
    param_set::AbstractParameterSet,
    orientation::Orientation,
    state::Vars,
    aux::Vars,
)
    ρe_pot = state.ρ * gravitational_potential(orientation, aux)
    linearized_air_pressure(
        param_set,
        state.ρ,
        state.ρe,
        ρe_pot,
        state.moisture.ρq_tot,
        state.moisture.ρq_liq,
        state.moisture.ρq_ice,
    )
end

abstract type AtmosLinearModel <: BalanceLaw end

function vars_state(lm::AtmosLinearModel, st::Prognostic, FT)
    @vars begin
        ρ::FT
        ρu::SVector{3, FT}
        ρe::FT
        turbulence::vars_state(lm.atmos.turbulence, st, FT)
        hyperdiffusion::vars_state(lm.atmos.hyperdiffusion, st, FT)
        moisture::vars_state(lm.atmos.moisture, st, FT)
    end
end
vars_state(lm::AtmosLinearModel, ::AbstractStateType, FT) = @vars()
vars_state(lm::AtmosLinearModel, st::Auxiliary, FT) =
    vars_state(lm.atmos, st, FT)

function update_auxiliary_state!(
    dg::DGModel,
    lm::AtmosLinearModel,
    Q::MPIStateArray,
    t::Real,
    elems::UnitRange,
)
    return false
end
function flux_second_order!(
    lm::AtmosLinearModel,
    flux::Grad,
    state::Vars,
    diffusive::Vars,
    hyperdiffusive::Vars,
    aux::Vars,
    t::Real,
)
    nothing
end
integral_load_auxiliary_state!(
    lm::AtmosLinearModel,
    integ::Vars,
    state::Vars,
    aux::Vars,
) = nothing
integral_set_auxiliary_state!(lm::AtmosLinearModel, aux::Vars, integ::Vars) =
    nothing
reverse_integral_load_auxiliary_state!(
    lm::AtmosLinearModel,
    integ::Vars,
    state::Vars,
    aux::Vars,
) = nothing
reverse_integral_set_auxiliary_state!(
    lm::AtmosLinearModel,
    aux::Vars,
    integ::Vars,
) = nothing
flux_second_order!(
    lm::AtmosLinearModel,
    flux::Grad,
    state::Vars,
    diffusive::Vars,
    aux::Vars,
    t::Real,
) = nothing
function wavespeed(
    lm::AtmosLinearModel,
    nM,
    state::Vars,
    aux::Vars,
    t::Real,
    direction,
)
    ref = aux.ref_state
    return soundspeed_air(lm.atmos.param_set, ref.T)
end

function boundary_state!(
    nf::NumericalFluxFirstOrder,
    atmoslm::AtmosLinearModel,
    args...,
)
    atmos_boundary_state!(nf, AtmosBC(), atmoslm, args...)
end
function boundary_state!(
    nf::NumericalFluxSecondOrder,
    atmoslm::AtmosLinearModel,
    args...,
)
    nothing
end
init_state_auxiliary!(
    lm::AtmosLinearModel,
    aux::MPIStateArray,
    grid,
    direction,
) = nothing
init_state_prognostic!(
    lm::AtmosLinearModel,
    state::Vars,
    aux::Vars,
    coords,
    t,
) = nothing

struct AtmosAcousticLinearModel{M} <: AtmosLinearModel
    atmos::M
    function AtmosAcousticLinearModel(atmos::M) where {M}
        if atmos.ref_state === NoReferenceState()
            error("AtmosAcousticLinearModel needs a model with a reference state")
        end
        new{M}(atmos)
    end
end

function flux_first_order!(
    lm::AtmosAcousticLinearModel,
    flux::Grad,
    state::Vars,
    aux::Vars,
    t::Real,
    direction,
)
    FT = eltype(state)
    ref = aux.ref_state
    e_pot = gravitational_potential(lm.atmos.orientation, aux)

    flux.ρ = state.ρu
    pL = linearized_pressure(
        lm.atmos.moisture,
        lm.atmos.param_set,
        lm.atmos.orientation,
        state,
        aux,
    )
    flux.ρu += pL * I
    flux.ρe = ((ref.ρe + ref.p) / ref.ρ - e_pot) * state.ρu
    nothing
end
source!(::AtmosAcousticLinearModel, _...) = nothing

struct AtmosAcousticGravityLinearModel{M} <: AtmosLinearModel
    atmos::M
    function AtmosAcousticGravityLinearModel(atmos::M) where {M}
        if atmos.ref_state === NoReferenceState()
            error("AtmosAcousticGravityLinearModel needs a model with a reference state")
        end
        new{M}(atmos)
    end
end
function flux_first_order!(
    lm::AtmosAcousticGravityLinearModel,
    flux::Grad,
    state::Vars,
    aux::Vars,
    t::Real,
    direction,
)
    FT = eltype(state)
    ref = aux.ref_state
    e_pot = gravitational_potential(lm.atmos.orientation, aux)

    flux.ρ = state.ρu
    pL = linearized_pressure(
        lm.atmos.moisture,
        lm.atmos.param_set,
        lm.atmos.orientation,
        state,
        aux,
    )
    flux.ρu += pL * I
    flux.ρe = ((ref.ρe + ref.p) / ref.ρ) * state.ρu
    nothing
end
function source!(
    lm::AtmosAcousticGravityLinearModel,
    source::Vars,
    state::Vars,
    diffusive::Vars,
    aux::Vars,
    t::Real,
    ::NTuple{1, Dir},
) where {Dir <: Direction}
    if Dir === VerticalDirection || Dir === EveryDirection
        ∇Φ = ∇gravitational_potential(lm.atmos.orientation, aux)
        source.ρu -= state.ρ * ∇Φ
    end
    nothing
end

function numerical_flux_first_order!(
    numerical_flux::RoeNumericalFlux,
    balance_law::AtmosLinearModel,
    fluxᵀn::Vars{S},
    normal_vector::SVector,
    state_prognostic⁻::Vars{S},
    state_auxiliary⁻::Vars{A},
    state_prognostic⁺::Vars{S},
    state_auxiliary⁺::Vars{A},
    t,
    direction,
) where {S, A}
    @assert balance_law.atmos.moisture isa DryModel

    numerical_flux_first_order!(
        CentralNumericalFluxFirstOrder(),
        balance_law,
        fluxᵀn,
        normal_vector,
        state_prognostic⁻,
        state_auxiliary⁻,
        state_prognostic⁺,
        state_auxiliary⁺,
        t,
        direction,
    )

    atmos = balance_law.atmos
    param_set = atmos.param_set

    ρu⁻ = state_prognostic⁻.ρu

    ref_ρ⁻ = state_auxiliary⁻.ref_state.ρ
    ref_ρe⁻ = state_auxiliary⁻.ref_state.ρe
    ref_T⁻ = state_auxiliary⁻.ref_state.T
    ref_p⁻ = state_auxiliary⁻.ref_state.p
    ref_h⁻ = (ref_ρe⁻ + ref_p⁻) / ref_ρ⁻
    ref_c⁻ = soundspeed_air(param_set, ref_T⁻)

    pL⁻ = linearized_pressure(
        atmos.moisture,
        param_set,
        atmos.orientation,
        state_prognostic⁻,
        state_auxiliary⁻,
    )

    ρu⁺ = state_prognostic⁺.ρu

    ref_ρ⁺ = state_auxiliary⁺.ref_state.ρ
    ref_ρe⁺ = state_auxiliary⁺.ref_state.ρe
    ref_T⁺ = state_auxiliary⁺.ref_state.T
    ref_p⁺ = state_auxiliary⁺.ref_state.p
    ref_h⁺ = (ref_ρe⁺ + ref_p⁺) / ref_ρ⁺
    ref_c⁺ = soundspeed_air(param_set, ref_T⁺)

    pL⁺ = linearized_pressure(
        atmos.moisture,
        param_set,
        atmos.orientation,
        state_prognostic⁺,
        state_auxiliary⁺,
    )

    # not sure if arithmetic averages are a good idea here
    h̃ = (ref_h⁻ + ref_h⁺) / 2
    c̃ = (ref_c⁻ + ref_c⁺) / 2

    ΔpL = pL⁺ - pL⁻
    Δρuᵀn = (ρu⁺ - ρu⁻)' * normal_vector

    fluxᵀn.ρ -= ΔpL / 2c̃
    fluxᵀn.ρu -= c̃ * Δρuᵀn * normal_vector / 2
    fluxᵀn.ρe -= h̃ * ΔpL / 2c̃
end

function numerical_flux_first_order!(
    ::HLLCNumericalFlux,
    balance_law::AtmosLinearModel,
    fluxᵀn::Vars{S},
    normal_vector::SVector,
    state_prognostic⁻::Vars{S},
    state_auxiliary⁻::Vars{A},
    state_prognostic⁺::Vars{S},
    state_auxiliary⁺::Vars{A},
    t,
    direction,
) where {S, A}

    # There is no intermediate speed for the AtmosLinearModel.
    # As a result, HLLC simplifies to Rusanov.
    numerical_flux_first_order!(
        RusanovNumericalFlux(),
        balance_law,
        fluxᵀn,
        normal_vector,
        state_prognostic⁻,
        state_auxiliary⁻,
        state_prognostic⁺,
        state_auxiliary⁺,
        t,
        direction,
    )
end


"""
    Code for LineList types used by CCfs
Author: Eric Ford
Created: August 2020
"""

"""A struct implementing a line list should be a subtype of AbstractLineList. """
abstract type AbstractLineList end

""" A basic line list for passing to compute CCFs.
Contains (views into) arrays specifying the minimum and maximum wavelength range and weight for each line. """
struct BasicLineList{T<:Real, AA<:AbstractArray{T,1} } <: AbstractLineList
    λ::AA
    weight::AA
end

""" `BasicLineList( λ, weight )` """
function BasicLineList{T}(λ::AA, w::AA) where { T<:Real, AA<:AbstractArray{T,1} }
    @assert length(λ) == length(w)
    @assert length(λ) >= 1
    @assert 0.0 .<= w .<= 1.0
    BasicLineList{eltype(w),typeof(w)}(λ,w)
end

function BasicLineList( df::DataFrame )
    @assert hasproperty(df,:lambda)
    @assert hasproperty(df,:weight)
    BasicLineList{eltype(df[:,:lambda]),typeof(df[:,:weight])}(df[:,:lambda],df[:,:weight])
end

function EmptyBasicLineList()
    return BasicLineList{Float64,Array{Float64,1}}(zeros(0),zeros(0))
end

""" A basic line list for passing to compute CCFs.
Contains (views into) arrays specifying the minimum and maximum wavelength range and weight for each line. """
struct BasicLineList2D{T1<:Real, AA1<:AbstractArray{T1,1}, T2<:Integer,  AA2<:AbstractArray{T2,1} } <: AbstractLineList
    λ::AA1
    weight::AA1
    order::AA2
end

""" `BasicLineList( λ, weight )` """
function BasicLineList2D{T1,T2}(λ::AA1, w::AA1, o::AA2) where {T1<:Real, AA1<:AbstractArray{T1,1}, T2<:Integer,  AA2<:AbstractArray{T2,1} }
    @assert length(λ) == length(w) == length(o)
    @assert length(λ) >= 1
    @assert all(0.0 .<= w .<= 1.0)
    @assert all(1 .<= o .<= 200)           # arbitrary limit for maximum number of orders
    BasicLineList2D{eltype(w),typeof(w),eltype{o},typeof(o)}(λ,w,o)
end

function BasicLineList2D( df::DataFrame )
    @assert hasproperty(df,:lambda)
    @assert hasproperty(df,:weight)
    @assert hasproperty(df,:order)
    @assert all(0.0 .<= df[!,:weight] .<= 1.0)
    @assert all(1 .<= df[!,:order] .<= 200)           # arbitrary limit for maximum number of orders
    BasicLineList2D{eltype(df[:,:lambda]),typeof(df[:,:weight]),eltype(df[:,:order]),typeof(df[:,:order])}(df[!,:lambda],df[!,:weight],df[!,:order])
end

function EmptyBasicLineList2D()
    return BasicLineList2D{Float64,Array{Float64,1},Int32,Array{Int32,1}}(zeros(0),zeros(0), zeros(Int32,0) )
end

#=  Not fully implemented/tested yet
""" A line list for passing to compute CCFs with variable line widths.
Contains (views into) arrays specifying the minimum and maximum wavelength range and weight for each line. """
struct VarWidthLineList{T<:Real, AA<:AbstractArray{T,1} } <: AbstractLineList
    λ_lo::AA
    λ_hi::AA
    weight::AA
end

""" VarWidthLineList( λ_lo, λ_hi, weight ) """
function VarWidthLineList{T}(lo::AA, hi::AA, w::AA) where { T<:Real, AA<:AbstractArray{T,1} }
    @assert length(lo) == length(hi) == length(w)
    @assert length(lo) >= 1
    @assert 0.0 .<= λ_hi.-λ_lo .<= 2.0
    @assert 0.0 .<= w .<= 1.0
    VarWidthLineList{eltype(w),typeof(w)}(lo,hi,w)
end
=#

import Base.length
""" Return length of line linst. """
length(ll::AbstractLineList) = length(ll.weight)


module LiquidsDynamics

### Imports

# Load submodules code
include("Projections.jl")
include("LiquidsStructureGrids.jl")

# Load the actual modules
using Reexport
using Parameters
using LinearAlgebra: I

@reexport using .Projections
@reexport using .LiquidsStructureGrids

### Exports
export asymptotics, asymptotics!, dynamics, dynamics!, initialize_asymptotics,
       initialize_dynamics

### Implementation
abstract type SCGLE end
abstract type Dynamics    <: SCGLE end
abstract type Asymptotics <: SCGLE end

include("utils.jl")
include("dynamics/dynamics.jl")
include("dynamics/init.jl")
include("dynamics/shorttimes.jl")
include("dynamics/intermediatetimes.jl")
include("dynamics/utils.jl")
include("asymptotics/asymptotics.jl")
include("asymptotics/init.jl")
include("asymptotics/utils.jl")

end # module


###
### iterate
###
checkindexlo(r, i::AbstractVector) = checkindexlo(r, minimum(i))
checkindexlo(r, i) = firstindex(r) <= i
checkindexlo(r, i::CartesianIndex{1}) = firstindex(r) <= first(i.I)

checkindexhi(r, i::AbstractVector) = checkindexhi(r, maximum(i))
checkindexhi(r, i) = lastindex(r) >= i
checkindexhi(r, i::CartesianIndex{1}) = firstindex(r) <= first(i.I)

###
### Generic array traits
###
# TODO should this be in Static.jl
static_isempty(x::OrdinalRange) = _static_isempty(static_first(x), static_step(x), static_last(x))
function static_isempty(x)
    len = static_length(x)
    return Static.eq(zero(len), len)
end
function _static_isempty(start::F, step::S, stop::L) where {F,S,L}
    return Static.ne(start, stop) & Static.ne(Static.gt(step, zero(step)), Static.gt(stop, start))
end

# FIXME these absolutely needs to go in ArrayInterface
function ArrayInterface.known_length(::Type{T}) where {T<:AbstractRange}
    if parent_type(T) <: T
        return nothing
    else
        return known_length(parent_type(T))
    end
end

#Base.rem(::Static.StaticFloat64{X}, ::Static.StaticFloat64{Y}) where {X,Y} = static(rem(X, Y))
static_div(x::X, y::Y) where {X,Y} = _div(is_static(X) & is_static(Y), x, y)
_div(::True, x, y) = static(div(known(x), known(y)))
_div(::False, x, y) = div(dynamic(x), dynamic(y))

static_rem(x::X, y::Y) where {X,Y} = _rem(is_static(X) & is_static(Y), x, y)
_rem(::True, x, y) = static(rem(known(x), known(y)))
_rem(::False, x, y) = rem(dynamic(x), dynamic(y))

_sub1(x::T) where {T} = x - oneunit(T)
_add1(x::T) where {T} = x + oneunit(T)
_int(idx) = round(Integer, idx, RoundToZero)::Int
_int(idx::Integer) = Int(idx)::Int
_int(idx::StaticInt{N}) where {N} = idx
_int(idx::TwicePrecision{T}) where {T} = round(Integer, T(idx), RoundToZero)

_drop_unit(x::X) where {X} = div(x, oneunit(x))
_drop_unit(x::Real) = x

### empty
_empty_ur(::Type{T}) where {T} = one(T):zero(T)

_empty(x::X, y::Y) where {X,Y} = Vector{Int}()
@inline function _empty(x::X, y::Y) where {X<:AbstractRange,Y<:AbstractRange}
    if (known_step(X) === nothing) | (known_step(Y) === nothing)
        return 1:1:0
    else
        if known_first(x) === one(eltype(X))  && known_first(y) === one(eltype(Y))
            if known_last(x) isa Nothing || known_last(y) isa Nothing
                return static(1):0
            else
                return static(1):static(0)
            end
        else
            return 1:0
        end
    end
end

const Equal{T} = Union{Fix2{typeof(==),T},Fix2{typeof(isequal),T}}
const NotEqual{T} = Fix2{typeof(!=),T}
const NotIn{T} = (typeof(!in(Any)).name.wrapper){Fix2{typeof(in),T}}
const In{T} = Fix2{typeof(in),T}

_maybe_static(::True, x::Int) = static(x)
_maybe_static(::True, x::StaticInt) = x
_maybe_static(::False, x::Int) = x
_maybe_static(::False, x::StaticInt) = dynamic(x)

ifelseop(::True, t, f, args...) = t(args...)
ifelseop(::False, t, f, args...) = f(args...)
@inline function ifelseop(b::Bool, t, f, args...)
    if b
        return dynamic(t(args...))
    else
        return dynamic(f(args...))
    end
end

return_nothing(args...) = nothing

return_static_one(args...) = static(1)

ifadd1(x, val) = ifelseop(x, _add1, identity, val)
ifsub1(x, val) = ifelseop(x, _sub1, identity, val)

_static_length(_, r) = static_length(r)
_static_length(_, _, r) = static_length(r)


using MurmurHash3: mmhash128_a

function Base.hash(x::ShortString, h::UInt)
    h += Base.memhash_seed
    last(mmhash128_a(sizeof(x), bswap(x.size_content), h%UInt32)) + h
end


#######################################################################################################################################################################################################
#
# Changes to this type
# General
#     2021-Nov-29: add abstract photosynthesis system type
#     2022-Jan-14: rename to photosynthesis model
#     2022-Jan-14: add hierachy description
#     2022-Jan-25: fix documentation
#
#######################################################################################################################################################################################################
"""

$(TYPEDEF)

Hierachy of `AbstractPhotosynthesisModel`:
- [`C3CytochromeModel`](@ref)
- [`C3VJPModel`](@ref)
- [`C4VJPModel`](@ref)
"""
abstract type AbstractPhotosynthesisModel{FT<:AbstractFloat} end

#######################################################################################################################################################################################################
#
# Changes to this struct
# General
#     2022-Jan-18: add C3CytochromeModel structure for C₃ photosynthesis system
#     2022-Jan-25: fix documentation
#     2022-Feb-07: add more fields to use with Photosynthesis v0.3.1
#     2022-Feb-07: remove j_p680 and j_p700 series variables
#     2022-Feb-11: split COLIMIT to COLIMIT_CJ and COLIMIT_IP (minor breaking)
#     2022-Mar-01: add two more fields: TD_ΗC and TD_ΗL
#     2022-Mar-01: move η_c and η_l from reaction center to photosynthesis model
#     2022-Mar-01: move k_q from prognostic to dignostic section
#     2022-Mar-09: add v_cmax25_ww to use with StomataModels.jl
#
#######################################################################################################################################################################################################
"""

$(TYPEDEF)

Structure that stores C3 Cytochrome photosynthesis system information

# Fields

$(TYPEDFIELDS)

"""
mutable struct C3CytochromeModel{FT<:AbstractFloat} <: AbstractPhotosynthesisModel{FT}
    # parameters that do not change with time
    "[`AbstractColimit`](@ref) type colimitation method for Ac and Aj => Ai"
    COLIMIT_CJ::AbstractColimit{FT}
    "[`AbstractColimit`](@ref) type colimitation method for Ai and Ap => Ag"
    COLIMIT_IP::AbstractColimit{FT}
    "[`AbstractColimit`](@ref) type colimitation method for J"
    COLIMIT_J::AbstractColimit{FT}
    "Coefficient 4.0/4.5 for NADPH/ATP requirement stochiometry, respectively"
    EFF_1::FT
    "Coefficient 8.0/10.5 for NADPH/ATP requirement stochiometry, respectively"
    EFF_2::FT
    "[`AbstractTemperatureDependency`](@ref) type Kc temperature dependency"
    TD_KC::AbstractTemperatureDependency{FT}
    "[`AbstractTemperatureDependency`](@ref) type Ko temperature dependency"
    TD_KO::AbstractTemperatureDependency{FT}
    "[`AbstractTemperatureDependency`](@ref) type Kq temperature dependency"
    TD_KQ::AbstractTemperatureDependency{FT}
    "[`AbstractTemperatureDependency`](@ref) type respiration temperature dependency"
    TD_R::AbstractTemperatureDependency{FT}
    "[`AbstractTemperatureDependency`](@ref) type Vcmax temperature dependency"
    TD_VCMAX::AbstractTemperatureDependency{FT}
    "[`AbstractTemperatureDependency`](@ref) type Γ* temperature dependency"
    TD_Γ::AbstractTemperatureDependency{FT}
    "[`AbstractTemperatureDependency`](@ref) type Η_C temperature dependency"
    TD_ΗC::AbstractTemperatureDependency{FT}
    "[`AbstractTemperatureDependency`](@ref) type Η_L temperature dependency"
    TD_ΗL::AbstractTemperatureDependency{FT}

    # prognostic variables that change with time
    "Total concentration of Cytochrome b₆f `[μmol m⁻²]`"
    b₆f::FT
    "Respiration rate at 298.15 K `[μmol m⁻² s⁻¹]`"
    r_d25::FT
    "Maximal carboxylation rate at 298.15 K `[μmol m⁻² s⁻¹]`"
    v_cmax25::FT
    "Well watered maximal carboxylation rate at 298.15 K `[μmol m⁻² s⁻¹]`"
    v_cmax25_ww::FT

    # dignostic variables that change with time
    "RubisCO limited photosynthetic rate `[μmol m⁻² s⁻¹]`"
    a_c::FT
    "Gross photosynthetic rate `[μmol m⁻² s⁻¹]`"
    a_gross::FT
    "Light limited photosynthetic rate `[μmol m⁻² s⁻¹]`"
    a_j::FT
    "Net photosynthetic rate `[μmol m⁻² s⁻¹]`"
    a_net::FT
    "Product limited photosynthetic rate `[μmol m⁻² s⁻¹]`"
    a_p::FT
    "Electron to CO₂ coefficient"
    e_to_c::FT
    "Potential Electron Transport Rate `[μmol e⁻ m⁻² s⁻¹]`"
    j_pot::FT
    "PSI electron transport rate after colimitation"
    j_psi::FT
    "RubisCO coefficient Kc `[Pa]`"
    k_c::FT
    "Michaelis-Menten's coefficient `[Pa]`"
    k_m::FT
    "RubisCO coefficient Ko `[Pa]`"
    k_o::FT
    "Maximal turnover rate of Cytochrome b₆f `[e⁻ s⁻¹]`"
    k_q::FT
    "Respiration rate at leaf temperature `[μmol m⁻² s⁻¹]`"
    r_d::FT
    "Maximal carboxylation rate at leaf temperature `[μmol m⁻² s⁻¹]`"
    v_cmax::FT
    "Maximal Cytochrome b₆f activity `[μmol e⁻ m⁻² s⁻¹]`"
    v_qmax::FT
    "ratio between J_P700 and J_P680"
    η::FT
    "Coupling efficiency of cyclic electron flow `[mol ATP mol⁻¹ e⁻]`"
    η_c::FT
    "Coupling efficiency of linear electron flow `[mol ATP mol⁻¹ e⁻]`"
    η_l::FT
    "CO₂ compensation point with the absence of Rd `[Pa]`"
    γ_star::FT
end

#######################################################################################################################################################################################################
#
# Changes to this constructor
# General
#     2021-Nov-18: add constructor
#     2022-Jan-25: add documentation
#     2022-Feb-07: add more fields into the constructors
#     2022-Feb-11: split COLIMIT to COLIMIT_CJ and COLIMIT_IP (minor breaking)
#     2022-Feb-11: add colimit option in constructor to enable quick deployment of quadratic colimitation
#     2022-Mar-01: add two more fields: TD_ΗC and TD_ΗL
#     2022-Mar-01: move η_c and η_l from reaction center to photosynthesis model
#     2022-Mar-01: move k_q from prognostic to dignostic section
#     2022-Mar-09: add v_cmax25_ww to use with StomataModels.jl
#
#######################################################################################################################################################################################################
"""

    C3CytochromeModel{FT}(; v_cmax25::Number = 50, r_d25::Number = 0.75, colimit::Bool = false) where {FT<:AbstractFloat}

Constructor for `C3CytochromeModel`, given
- `v_cmax25` Maximal carboxylation rate at 298.15 K
- `r_d25` Respiration rate at 298.15 K
- `colimit` If true, use quadratic colimitations for a_c, a_j, and a_p

---
# Examples
```julia
cy = C3CytochromeModel{Float64}();
cy = C3CytochromeModel{Float64}(v_cmax25 = 30, r_d25 = 1, colimit = true);
```
"""
C3CytochromeModel{FT}(; v_cmax25::Number = 50, r_d25::Number = 0.75, colimit::Bool = false) where {FT<:AbstractFloat} = (
    if colimit
        _colim_cj = QuadraticColimit{FT}(0.98);
        _colim_ip = QuadraticColimit{FT}(0.95);
    else
        _colim_cj = MinimumColimit{FT}();
        _colim_ip = MinimumColimit{FT}();
    end;

    return C3CytochromeModel{FT}(
                _colim_cj,              # COLIMIT_CJ
                _colim_ip,              # COLIMIT_IP
                SerialColimit{FT}(),    # COLIMIT_J
                4,                      # EFF_1
                8,                      # EFF_2
                KcTDCLM(FT),            # TD_KC
                KoTDCLM(FT),            # TD_KO
                KqTDJohnson(FT),        # TD_KQ
                RespirationTDCLM(FT),   # TD_R
                VcmaxTDCLM(FT),         # TD_VCMAX
                ΓStarTDCLM(FT),         # TD_Γ
                ΗCTDJohnson(FT),        # TD_ΗC
                ΗLTDJohnson(FT),        # TD_ΗL
                350 / 300,              # b₆f
                r_d25,                  # r_d25
                v_cmax25,               # v_cmax25,
                v_cmax25,               # v_cmax25_ww
                0,                      # a_c
                0,                      # a_gross
                0,                      # a_j
                -r_d25,                 # a_net
                0,                      # a_p
                0,                      # e_to_c
                0,                      # j_pot
                0,                      # j_psi
                0,                      # k_c
                0,                      # k_m
                0,                      # k_o
                0,                      # k_q
                r_d25,                  # r_d
                v_cmax25,               # v_cmax
                0,                      # v_qmax
                0,                      # η
                0,                      # η_c
                0,                      # η_l
                0)                      # γ_star
);

#######################################################################################################################################################################################################
#
# Changes to this struct
# General
#     2021-Nov-11: add C3VJPModel structure for classic C₃ photosynthesis system
#     2022-Jan-14: add temperature dependency into the structure
#     2022-Jan-14: rename to photosynthesis model
#     2022-Jan-14: add colimitation and e_to_c
#     2022-Jan-25: fix documentation
#     2022-Feb-11: split COLIMIT to COLIMIT_CJ, COLIMIT_IP, and COLIMIT_J (minor breaking)
#     2022-Mar-09: add v_cmax25_ww to use with StomataModels.jl
#
#######################################################################################################################################################################################################
"""

$(TYPEDEF)

Structure that stores C3 photosynthesis system information

# Fields

$(TYPEDFIELDS)

"""
mutable struct C3VJPModel{FT<:AbstractFloat} <: AbstractPhotosynthesisModel{FT}
    # parameters that do not change with time
    "[`AbstractColimit`](@ref) type colimitation method for Ac and Aj => Ai"
    COLIMIT_CJ::AbstractColimit{FT}
    "[`AbstractColimit`](@ref) type colimitation method for Ai and Ap => Ag"
    COLIMIT_IP::AbstractColimit{FT}
    "[`AbstractColimit`](@ref) type colimitation method for J"
    COLIMIT_J::AbstractColimit{FT}
    "Coefficient 4.0/4.5 for NADPH/ATP requirement stochiometry, respectively"
    EFF_1::FT
    "Coefficient 8.0/10.5 for NADPH/ATP requirement stochiometry, respectively"
    EFF_2::FT
    "[`AbstractTemperatureDependency`](@ref) type Jmax temperature dependency"
    TD_JMAX::AbstractTemperatureDependency{FT}
    "[`AbstractTemperatureDependency`](@ref) type Kc temperature dependency"
    TD_KC::AbstractTemperatureDependency{FT}
    "[`AbstractTemperatureDependency`](@ref) type Ko temperature dependency"
    TD_KO::AbstractTemperatureDependency{FT}
    "[`AbstractTemperatureDependency`](@ref) type respiration temperature dependency"
    TD_R::AbstractTemperatureDependency{FT}
    "[`AbstractTemperatureDependency`](@ref) type Vcmax temperature dependency"
    TD_VCMAX::AbstractTemperatureDependency{FT}
    "[`AbstractTemperatureDependency`](@ref) type Γ* temperature dependency"
    TD_Γ::AbstractTemperatureDependency{FT}

    # prognostic variables that change with time
    "Maximal electron transport rate at 298.15 K `[μmol m⁻² s⁻¹]`"
    j_max25::FT
    "Respiration rate at 298.15 K `[μmol m⁻² s⁻¹]`"
    r_d25::FT
    "Maximal carboxylation rate at 298.15 K `[μmol m⁻² s⁻¹]`"
    v_cmax25::FT
    "Well watered maximal carboxylation rate at 298.15 K `[μmol m⁻² s⁻¹]`"
    v_cmax25_ww::FT

    # dignostic variables that change with time
    "RubisCO limited photosynthetic rate `[μmol m⁻² s⁻¹]`"
    a_c::FT
    "Gross photosynthetic rate `[μmol m⁻² s⁻¹]`"
    a_gross::FT
    "Light limited photosynthetic rate `[μmol m⁻² s⁻¹]`"
    a_j::FT
    "Net photosynthetic rate `[μmol m⁻² s⁻¹]`"
    a_net::FT
    "Product limited photosynthetic rate `[μmol m⁻² s⁻¹]`"
    a_p::FT
    "Electron to CO₂ coefficient"
    e_to_c::FT
    "Electron transport `[μmol m⁻² s⁻¹]`"
    j::FT
    "Maximal electron transport rate at leaf temperature `[μmol m⁻² s⁻¹]`"
    j_max::FT
    "Potential Electron Transport Rate `[μmol m⁻² s⁻¹]`"
    j_pot::FT
    "RubisCO coefficient Kc `[Pa]`"
    k_c::FT
    "Michaelis-Menten's coefficient `[Pa]`"
    k_m::FT
    "RubisCO coefficient Ko `[Pa]`"
    k_o::FT
    "Respiration rate at leaf temperature `[μmol m⁻² s⁻¹]`"
    r_d::FT
    "Maximal carboxylation rate at leaf temperature `[μmol m⁻² s⁻¹]`"
    v_cmax::FT
    "CO₂ compensation point with the absence of Rd `[Pa]`"
    γ_star::FT
end

#######################################################################################################################################################################################################
#
# Changes to this constructor
# General
#     2021-Nov-11: add constructor
#     2022-Jan-14: update constructor to match structure
#     2022-Jan-25: fix documentation
#     2022-Feb-11: split COLIMIT to COLIMIT_CJ, COLIMIT_IP, and COLIMIT_J (minor breaking)
#     2022-Feb-11: add colimit option in constructor to enable quick deployment of quadratic colimitation
#     2022-Mar-09: add v_cmax25_ww to use with StomataModels.jl
#
#######################################################################################################################################################################################################
"""

    C3VJPModel{FT}(; v_cmax25::Number = 50, j_max25::Number = 83.5, r_d25::Number = 0.75, colimit::Bool = false) where {FT<:AbstractFloat}

Constructor for [`C3VJPModel`](@ref), given
- `v_cmax25` Maximal carboxylation rate at 298.15 K
- `j_max25` Maximal electron transport rate at 298.15 K
- `r_d25` Respiration rate at 298.15 K
- `colimit` If true, use quadratic colimitations for j and a_c, a_j, and a_p

---
# Examples
```julia
c3 = C3VJPModel{Float64}();
c3 = C3VJPModel{Float64}(v_cmax25 = 30, j_max25 = 50, r_d25 = 1, colimit = true);
```
"""
C3VJPModel{FT}(; v_cmax25::Number = 50, j_max25::Number = 83.5, r_d25::Number = 0.75, colimit::Bool = false) where {FT<:AbstractFloat} = (
    if colimit
        _colim_cj = QuadraticColimit{FT}(0.98);
        _colim_ip = QuadraticColimit{FT}(0.95);
        _colim_j  = QuadraticColimit{FT}(0.7);
    else
        _colim_cj = MinimumColimit{FT}();
        _colim_ip = MinimumColimit{FT}();
        _colim_j  = MinimumColimit{FT}();
    end;

    return C3VJPModel{FT}(
                _colim_cj,              # COLIMIT_CJ
                _colim_ip,              # COLIMIT_IP
                _colim_j,               # COLIMIT_J
                4,                      # EFF_1
                8,                      # EFF_2
                JmaxTDCLM(FT),          # TD_JMAX
                KcTDCLM(FT),            # TD_KC
                KoTDCLM(FT),            # TD_KO
                RespirationTDCLM(FT),   # TD_R
                VcmaxTDCLM(FT),         # TD_VCMAX
                ΓStarTDCLM(FT),         # TD_Γ
                j_max25,                # j_max25
                r_d25,                  # r_d25
                v_cmax25,               # v_cmax25
                v_cmax25,               # v_cmax25_ww
                0,                      # a_c
                0,                      # a_gross
                0,                      # a_j
                -r_d25,                 # a_net
                0,                      # a_p
                0,                      # e_to_c
                0,                      # j
                j_max25,                # j_max
                0,                      # j_pot
                0,                      # k_c
                0,                      # k_m
                0,                      # k_o
                r_d25,                  # r_d
                v_cmax25,               # v_cmax
                0)                      # γ_star
);

#######################################################################################################################################################################################################
#
# Changes to this struct
# General
#     2022-Jan-14: add C4VJPModel structure for classic C₄ photosynthesis system
#     2022-Jan-25: fix documentation
#     2022-Feb-11: remove j from the struct
#     2022-Feb-11: split COLIMIT to COLIMIT_CJ and COLIMIT_IP (minor breaking)
#     2022-Mar-09: add v_cmax25_ww to use with StomataModels.jl
# To do
#     TODO: add Jmax to C4VJPModel and thus JMAX TD in Photosynthesis.jl (not necessary)
#
#######################################################################################################################################################################################################
"""

$(TYPEDEF)

Structure that stores C4 photosynthesis system information

# Fields

$(TYPEDFIELDS)

"""
mutable struct C4VJPModel{FT<:AbstractFloat} <: AbstractPhotosynthesisModel{FT}
    # parameters that do not change with time
    "[`AbstractColimit`](@ref) type colimitation method for Ac and Aj => Ai"
    COLIMIT_CJ::AbstractColimit{FT}
    "[`AbstractColimit`](@ref) type colimitation method for Ai and Ap => Ag"
    COLIMIT_IP::AbstractColimit{FT}
    "[`AbstractTemperatureDependency`](@ref) type Kpep temperature dependency"
    TD_KPEP::AbstractTemperatureDependency{FT}
    "[`AbstractTemperatureDependency`](@ref) type  respiration temperature dependency"
    TD_R::AbstractTemperatureDependency{FT}
    "[`AbstractTemperatureDependency`](@ref) type Vcmax temperature dependency"
    TD_VCMAX::AbstractTemperatureDependency{FT}
    "[`AbstractTemperatureDependency`](@ref) type Vpmax temperature dependency"
    TD_VPMAX::AbstractTemperatureDependency{FT}

    # prognostic variables that change with time
    "Respiration rate at 298.15 K `[μmol m⁻² s⁻¹]`"
    r_d25::FT
    "Maximal carboxylation rate at 298.15 K `[μmol m⁻² s⁻¹]`"
    v_cmax25::FT
    "Well watered maximal carboxylation rate at 298.15 K `[μmol m⁻² s⁻¹]`"
    v_cmax25_ww::FT
    "Maximal PEP carboxylation rate at 298.15 K `[μmol m⁻² s⁻¹]`"
    v_pmax25::FT

    # dignostic variables that change with time
    "RubisCO limited photosynthetic rate `[μmol m⁻² s⁻¹]`"
    a_c::FT
    "Gross photosynthetic rate `[μmol m⁻² s⁻¹]`"
    a_gross::FT
    "Light limited photosynthetic rate `[μmol m⁻² s⁻¹]`"
    a_j::FT
    "Net photosynthetic rate `[μmol m⁻² s⁻¹]`"
    a_net::FT
    "Product limited photosynthetic rate `[μmol m⁻² s⁻¹]`"
    a_p::FT
    "Electron to CO₂ coefficient"
    e_to_c::FT
    "Potential Electron Transport Rate `[μmol m⁻² s⁻¹]`"
    j_pot::FT
    "PEP coefficient Kpep `[Pa]`"
    k_pep::FT
    "Respiration rate at leaf temperature `[μmol m⁻² s⁻¹]`"
    r_d::FT
    "Maximal carboxylation rate at leaf temperature `[μmol m⁻² s⁻¹]`"
    v_cmax::FT
    "Maximal PEP carboxylation rate at leaf temperature `[μmol m⁻² s⁻¹]`"
    v_pmax::FT
end

#######################################################################################################################################################################################################
#
# Changes to this constructor
# General
#     2021-Nov-18: add constructor
#     2022-Jan-25: fix documentation
#     2022-Feb-11: remove j from the struct
#     2022-Feb-11: default e_to_c set to 1/6
#     2022-Feb-11: split COLIMIT to COLIMIT_CJ and COLIMIT_IP (minor breaking)
#     2022-Feb-11: add colimit option in constructor to enable quick deployment of quadratic colimitation
#     2022-Mar-09: add v_cmax25_ww to use with StomataModels.jl
#
#######################################################################################################################################################################################################
"""

    C4VJPModel{FT}(; v_cmax25::Number = 50, v_pmax25::Number = 50, r_d25::Number = 0.75, colimit::Bool = false) where {FT<:AbstractFloat}

Constructor for [`C4VJPModel`](@ref), given
- `v_cmax25` Maximal carboxylation rate at 298.15 K
- `v_pmax25` Maximal PEP carboxylation rate at 298.15 K
- `r_d25` Respiration rate at 298.15 K
- `colimit` If true, use quadratic colimitations for a_c, a_j, and a_p

---
# Examples
```julia
c4 = C4VJPModel{Float64}();
c4 = C4VJPModel{Float64}(v_cmax25 = 30, v_pmax25 = 40, r_d25 = 1, colimit = true);
```
"""
C4VJPModel{FT}(; v_cmax25::Number = 50, v_pmax25::Number = 50, r_d25::Number = 0.75, colimit::Bool = false) where {FT<:AbstractFloat} = (
    if colimit
        _colim_cj = QuadraticColimit{FT}(0.8);
        _colim_ip = QuadraticColimit{FT}(0.95);
    else
        _colim_cj = MinimumColimit{FT}();
        _colim_ip = MinimumColimit{FT}();
    end;

    return C4VJPModel{FT}(
                _colim_cj,              # COLIMIT_CJ
                _colim_ip,              # COLIMIT_IP
                KpepTDCLM(FT),          # TD_KPEP
                RespirationTDCLM(FT),   # TD_R
                VcmaxTDCLM(FT),         # TD_VCMAX
                VpmaxTDBoyd(FT),        # TD_VPMAX
                r_d25,                  # r_d25
                v_cmax25,               # v_cmax25
                v_cmax25,               # v_cmax25_ww
                v_pmax25,               # v_pmax25
                0,                      # a_c
                0,                      # a_gross
                0,                      # a_j
                -r_d25,                 # a_net
                0,                      # a_p
                1/6,                    # e_to_c
                0,                      # j_pot
                0,                      # k_pep
                r_d25,                  # r_d
                v_cmax25,               # v_cmax
                v_pmax25)               # v_pmax
);

#######################################################################################################################################################################################################
#
# Changes to this type
# General
#     2022-Jan-14: add abstract mode type
#     2022-Jan-14: add hierachy description
#     2022-Jan-25: fix documentation
#
#######################################################################################################################################################################################################
"""

$(TYPEDEF)

Hierachy of AbstractSoilVC:
- [`GCO₂Mode`](@ref)
- [`PCO₂Mode`](@ref)
"""
abstract type AbstractPhotosynthesisMode end

#######################################################################################################################################################################################################
#
# Changes to this struct
# General
#     2022-Jan-14: add CO₂ conductance mode
#     2022-Jan-25: fix documentation
#
#######################################################################################################################################################################################################
"""

$(TYPEDEF)

An empty structure to signal the function to calculate photosynthetic rates based on leaf diffusive conductance to CO₂.

---
# Examples
```julia
mode = GCO₂Mode();
```
"""
struct GCO₂Mode <: AbstractPhotosynthesisMode end

#######################################################################################################################################################################################################
#
# Changes to this struct
# General
#     2022-Jan-14: add CO₂ partial pressure mode
#     2022-Jan-25: fix documentation
#
#######################################################################################################################################################################################################
"""

$(TYPEDEF)

An empty structure to signal the function to calculate photosynthetic rates based on CO₂ partial pressure.

---
# Examples
```julia
mode = PCO₂Mode();
```
"""
struct PCO₂Mode <: AbstractPhotosynthesisMode end


function add(obj::RealMatrix, arg0::RealMatrix)
    return jcall(obj, "add", RealMatrix, (RealMatrix,), arg0)
end

function add_to_entry(obj::RealMatrix, arg0::jint, arg1::jint, arg2::jdouble)
    return jcall(obj, "addToEntry", void, (jint, jint, jdouble), arg0, arg1, arg2)
end

function copy(obj::RealMatrix)
    return jcall(obj, "copy", RealMatrix, ())
end

function copy_sub_matrix(obj::RealMatrix, arg0::Vector{jint}, arg1::Vector{jint}, arg2::Vector{Vector{jdouble}})
    return jcall(obj, "copySubMatrix", void, (Vector{jint}, Vector{jint}, Vector{Vector{jdouble}}), arg0, arg1, arg2)
end

function copy_sub_matrix(obj::RealMatrix, arg0::jint, arg1::jint, arg2::jint, arg3::jint, arg4::Vector{Vector{jdouble}})
    return jcall(obj, "copySubMatrix", void, (jint, jint, jint, jint, Vector{Vector{jdouble}}), arg0, arg1, arg2, arg3, arg4)
end

function create_matrix(obj::RealMatrix, arg0::jint, arg1::jint)
    return jcall(obj, "createMatrix", RealMatrix, (jint, jint), arg0, arg1)
end

function get_column(obj::RealMatrix, arg0::jint)
    return jcall(obj, "getColumn", Vector{jdouble}, (jint,), arg0)
end

function get_column_matrix(obj::RealMatrix, arg0::jint)
    return jcall(obj, "getColumnMatrix", RealMatrix, (jint,), arg0)
end

function get_column_vector(obj::RealMatrix, arg0::jint)
    return jcall(obj, "getColumnVector", RealVector, (jint,), arg0)
end

function get_data(obj::RealMatrix)
    return jcall(obj, "getData", Vector{Vector{jdouble}}, ())
end

function get_entry(obj::RealMatrix, arg0::jint, arg1::jint)
    return jcall(obj, "getEntry", jdouble, (jint, jint), arg0, arg1)
end

function get_frobenius_norm(obj::RealMatrix)
    return jcall(obj, "getFrobeniusNorm", jdouble, ())
end

function get_row(obj::RealMatrix, arg0::jint)
    return jcall(obj, "getRow", Vector{jdouble}, (jint,), arg0)
end

function get_row_matrix(obj::RealMatrix, arg0::jint)
    return jcall(obj, "getRowMatrix", RealMatrix, (jint,), arg0)
end

function get_row_vector(obj::RealMatrix, arg0::jint)
    return jcall(obj, "getRowVector", RealVector, (jint,), arg0)
end

function get_sub_matrix(obj::RealMatrix, arg0::Vector{jint}, arg1::Vector{jint})
    return jcall(obj, "getSubMatrix", RealMatrix, (Vector{jint}, Vector{jint}), arg0, arg1)
end

function get_sub_matrix(obj::RealMatrix, arg0::jint, arg1::jint, arg2::jint, arg3::jint)
    return jcall(obj, "getSubMatrix", RealMatrix, (jint, jint, jint, jint), arg0, arg1, arg2, arg3)
end

function get_trace(obj::RealMatrix)
    return jcall(obj, "getTrace", jdouble, ())
end

function multiply(obj::RealMatrix, arg0::RealMatrix)
    return jcall(obj, "multiply", RealMatrix, (RealMatrix,), arg0)
end

function multiply_entry(obj::RealMatrix, arg0::jint, arg1::jint, arg2::jdouble)
    return jcall(obj, "multiplyEntry", void, (jint, jint, jdouble), arg0, arg1, arg2)
end

function operate(obj::RealMatrix, arg0::RealVector)
    return jcall(obj, "operate", RealVector, (RealVector,), arg0)
end

function operate(obj::RealMatrix, arg0::Vector{jdouble})
    return jcall(obj, "operate", Vector{jdouble}, (Vector{jdouble},), arg0)
end

function power(obj::RealMatrix, arg0::jint)
    return jcall(obj, "power", RealMatrix, (jint,), arg0)
end

function pre_multiply(obj::RealMatrix, arg0::RealMatrix)
    return jcall(obj, "preMultiply", RealMatrix, (RealMatrix,), arg0)
end

function pre_multiply(obj::RealMatrix, arg0::RealVector)
    return jcall(obj, "preMultiply", RealVector, (RealVector,), arg0)
end

function pre_multiply(obj::RealMatrix, arg0::Vector{jdouble})
    return jcall(obj, "preMultiply", Vector{jdouble}, (Vector{jdouble},), arg0)
end

function scalar_add(obj::RealMatrix, arg0::jdouble)
    return jcall(obj, "scalarAdd", RealMatrix, (jdouble,), arg0)
end

function scalar_multiply(obj::RealMatrix, arg0::jdouble)
    return jcall(obj, "scalarMultiply", RealMatrix, (jdouble,), arg0)
end

function set_column(obj::RealMatrix, arg0::jint, arg1::Vector{jdouble})
    return jcall(obj, "setColumn", void, (jint, Vector{jdouble}), arg0, arg1)
end

function set_column_matrix(obj::RealMatrix, arg0::jint, arg1::RealMatrix)
    return jcall(obj, "setColumnMatrix", void, (jint, RealMatrix), arg0, arg1)
end

function set_column_vector(obj::RealMatrix, arg0::jint, arg1::RealVector)
    return jcall(obj, "setColumnVector", void, (jint, RealVector), arg0, arg1)
end

function set_entry(obj::RealMatrix, arg0::jint, arg1::jint, arg2::jdouble)
    return jcall(obj, "setEntry", void, (jint, jint, jdouble), arg0, arg1, arg2)
end

function set_row(obj::RealMatrix, arg0::jint, arg1::Vector{jdouble})
    return jcall(obj, "setRow", void, (jint, Vector{jdouble}), arg0, arg1)
end

function set_row_matrix(obj::RealMatrix, arg0::jint, arg1::RealMatrix)
    return jcall(obj, "setRowMatrix", void, (jint, RealMatrix), arg0, arg1)
end

function set_row_vector(obj::RealMatrix, arg0::jint, arg1::RealVector)
    return jcall(obj, "setRowVector", void, (jint, RealVector), arg0, arg1)
end

function set_sub_matrix(obj::RealMatrix, arg0::Vector{Vector{jdouble}}, arg1::jint, arg2::jint)
    return jcall(obj, "setSubMatrix", void, (Vector{Vector{jdouble}}, jint, jint), arg0, arg1, arg2)
end

function subtract(obj::RealMatrix, arg0::RealMatrix)
    return jcall(obj, "subtract", RealMatrix, (RealMatrix,), arg0)
end

function transpose(obj::RealMatrix)
    return jcall(obj, "transpose", RealMatrix, ())
end

function walk_in_column_order(obj::RealMatrix, arg0::RealMatrixChangingVisitor)
    return jcall(obj, "walkInColumnOrder", jdouble, (RealMatrixChangingVisitor,), arg0)
end

function walk_in_column_order(obj::RealMatrix, arg0::RealMatrixChangingVisitor, arg1::jint, arg2::jint, arg3::jint, arg4::jint)
    return jcall(obj, "walkInColumnOrder", jdouble, (RealMatrixChangingVisitor, jint, jint, jint, jint), arg0, arg1, arg2, arg3, arg4)
end

function walk_in_column_order(obj::RealMatrix, arg0::RealMatrixPreservingVisitor)
    return jcall(obj, "walkInColumnOrder", jdouble, (RealMatrixPreservingVisitor,), arg0)
end

function walk_in_column_order(obj::RealMatrix, arg0::RealMatrixPreservingVisitor, arg1::jint, arg2::jint, arg3::jint, arg4::jint)
    return jcall(obj, "walkInColumnOrder", jdouble, (RealMatrixPreservingVisitor, jint, jint, jint, jint), arg0, arg1, arg2, arg3, arg4)
end

function walk_in_optimized_order(obj::RealMatrix, arg0::RealMatrixChangingVisitor)
    return jcall(obj, "walkInOptimizedOrder", jdouble, (RealMatrixChangingVisitor,), arg0)
end

function walk_in_optimized_order(obj::RealMatrix, arg0::RealMatrixChangingVisitor, arg1::jint, arg2::jint, arg3::jint, arg4::jint)
    return jcall(obj, "walkInOptimizedOrder", jdouble, (RealMatrixChangingVisitor, jint, jint, jint, jint), arg0, arg1, arg2, arg3, arg4)
end

function walk_in_optimized_order(obj::RealMatrix, arg0::RealMatrixPreservingVisitor)
    return jcall(obj, "walkInOptimizedOrder", jdouble, (RealMatrixPreservingVisitor,), arg0)
end

function walk_in_optimized_order(obj::RealMatrix, arg0::RealMatrixPreservingVisitor, arg1::jint, arg2::jint, arg3::jint, arg4::jint)
    return jcall(obj, "walkInOptimizedOrder", jdouble, (RealMatrixPreservingVisitor, jint, jint, jint, jint), arg0, arg1, arg2, arg3, arg4)
end

function walk_in_row_order(obj::RealMatrix, arg0::RealMatrixChangingVisitor)
    return jcall(obj, "walkInRowOrder", jdouble, (RealMatrixChangingVisitor,), arg0)
end

function walk_in_row_order(obj::RealMatrix, arg0::RealMatrixChangingVisitor, arg1::jint, arg2::jint, arg3::jint, arg4::jint)
    return jcall(obj, "walkInRowOrder", jdouble, (RealMatrixChangingVisitor, jint, jint, jint, jint), arg0, arg1, arg2, arg3, arg4)
end

function walk_in_row_order(obj::RealMatrix, arg0::RealMatrixPreservingVisitor)
    return jcall(obj, "walkInRowOrder", jdouble, (RealMatrixPreservingVisitor,), arg0)
end

function walk_in_row_order(obj::RealMatrix, arg0::RealMatrixPreservingVisitor, arg1::jint, arg2::jint, arg3::jint, arg4::jint)
    return jcall(obj, "walkInRowOrder", jdouble, (RealMatrixPreservingVisitor, jint, jint, jint, jint), arg0, arg1, arg2, arg3, arg4)
end


module MERAKit

using TensorKit
using TensorKitManifolds
using KrylovKit
using OptimKit
using TupleTools
using NamedTupleTools
using Printf
using LinearAlgebra
using Logging

export GenericMERA
export Layer
export SimpleLayer
export TernaryMERA
export BinaryMERA
export ModifiedBinaryMERA
export TernaryLayer
export BinaryLayer
export ModifiedBinaryLayer
export ModifiedBinaryOp
export ascend
export ascended_operator
export scale_invariant_operator_sum
export descend
export fixedpoint_densitymatrix
export environment
export remove_symmetry
export layertype
export baselayertype
export operatortype
export getlayer
export num_translayers
export outputspace
export inputspace
export internalspace
export causal_cone_width
export scalefactor
export densitymatrix
export densitymatrices
export densitymatrix_entropies
export random_MERA
export randomlayer
export scalingdimensions
export replace_layer
export release_transitionlayer
export expand_bonddim
export expand_internal_bonddim
export expect
export minimize_expectation
export pseudoserialize
export depseudoserialize
export projectisometric
export projectisometric!
export reset_storage
export tensorwise_sum
export tensorwise_scale
export gradient
export retract
export transport!
export inner

include("tensortools.jl")
include("layer.jl")
include("meracache.jl")
include("genericmera.jl")
include("simplelayer.jl")
include("ternarylayer.jl")
include("binarylayer.jl")
include("modifiedbinaryop.jl")
include("modifiedbinarylayer.jl")

# properties of instances
layertype(m::Union{GenericMERA, MERACache, Layer}) = layertype(typeof(m))
baselayertype(m::Union{GenericMERA, MERACache, Layer}) = baselayertype(typeof(m))
operatortype(m::Union{GenericMERA, MERACache, Layer}) = operatortype(typeof(m))
scalefactor(m::Union{GenericMERA, MERACache, Layer}) = scalefactor(typeof(m))
causal_cone_width(m::Union{GenericMERA, MERACache, Layer}) = causal_cone_width(typeof(m))
Base.eltype(m::Union{GenericMERA, MERACache, Layer}) = eltype(typeof(m))

end  # module


using ThreeBodyTB

function example_run_dft()

    # In this example, we fit a simple model. A robust TB model would require much more data.

    #First, it is necessary to run the DFT calculations, for example, we could define
    #a crystal as

    A = [5.33955  5.33955  0.0;5.33955  0.0      5.33955;0.0      5.33955  5.33955] * 0.529177
    pos = [0 0 0; 0.25 0.25 0.25]
    names = ["Mg", "S"]

    c = makecrys(A, pos, names)

    # To run the dft calculation, first we have to tell the code where QE (quantum espresso) is.

    ThreeBodyTB.set_bin_dirs(qe="/home/kfg/codes/qe-6.6/bin/") # you must give the code your own QE !!!!

    # You may also want to include an MPI command like

    ThreeBodyTB.set_bin_dirs(mpi="mpirun -np ")

    # Lets create an output directory

    mkpath("dft_out")

    dft = missing

    try

        # Then we could run a job like

        dft = ThreeBodyTB.DFT.runSCF(c, directory="dft_out", tmpdir="dft_out")

        # if you add a variable nprocs = 8, it will run in parallel on 8 processors. This will take a minute or two on modern computers.

        # This returns a dft output object, which you could also load later with :

        dft2 = ThreeBodyTB.QE.loadXML("dft_out/qe.save/")

    catch

        # if you don't have QE installed, this won't work obviously. However, I already ran this calculation for you, we can load the results

        dft = ThreeBodyTB.QE.loadXML(ThreeBodyTB.TESTDIR * "/data_forces/znse.in_vnscf_vol_2/qe.save")

    end

    # To calculate the tight-binding k-space object, we can use
    tbck=missing

    try

        workd = "dft_out/"
        tbc, tbck, proj_warn = ThreeBodyTB.AtomicProj.projwfc_workf(dft, directory=workd, writefile="projham_K.xml", only_kspace = true)
        #this will take a minute or two as well, and can be run in parallel

    catch

        # using precalculated data

        workd = ThreeBodyTB.TESTDIR * "/data_forces/znse.in_vnscf_vol_2/"
        tbc, tbck, proj_warn = ThreeBodyTB.AtomicProj.projwfc_workf(dft, directory=workd, writefile="projham_K.xml", only_kspace = true)

    end

    #remove self-consistent terms from tbck, so we get the corrent answer back when we solve self-consistently
    tbck_scf = ThreeBodyTB.SCF.remove_scf_from_tbc(tbck)

    return dft
end

function do_fitting(dft_ref)

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

    # Now that we have seen how to prepare one dft calculation, lets load all the data we need
    # which is precalculated.

    TESTDIR = ThreeBodyTB.TESTDIR

    ft = open("$TESTDIR/data_forces/fil_MgS_znse", "r"); #the list of files we need.
    dirst = readlines(ft);
    close(ft);

    tbc_list = []
    dft_list = []

    #here we load the data
    for t in dirst
        tfull = "$TESTDIR/"*t
        try #try to load from xml files

            dft = ThreeBodyTB.QE.loadXML(tfull*"/qe.save")

            tbck = ThreeBodyTB.TB.read_tb_crys_kspace("projham_K.xml.gz", directory=tfull)

            tbck_scf = ThreeBodyTB.SCF.remove_scf_from_tbc(tbck)

            push!(dft_list, dft)
            push!(tbc_list, tbck_scf)

        catch
            println("missing data $tfull")
        end

    end

    # We only want to fit the Mg-S interactions, not the Mg-Mg or S-S. So we will give those to the fitting algorithm as the starting data, which is fixed.

    ThreeBodyTB.ManageDatabase.clear_database() #make sure nothing is pre-cached
    ThreeBodyTB.ManageDatabase.prepare_database(["Mg"])
    ThreeBodyTB.ManageDatabase.prepare_database(["S"])
    database_start = deepcopy(ThreeBodyTB.ManageDatabase.database_cached)
    ThreeBodyTB.ManageDatabase.clear_database() #make sure nothing is pre-cached

    #now we do the new fitting.

    database_new = ThreeBodyTB.FitTB.do_fitting_recursive(tbc_list, dft_list = dft_list, starting_database = database_start)

    #now we run an example calculation with the new database

    energy_tot, tbc, conv_flag = scf_energy(dft_ref.crys, database=database_new)

    println("We can compare our model energy $energy_tot with the DFT atomization energy ", dft_ref.atomize_energy)

    #println("note that this model is badly overfit and needs more data")

end

dft = example_run_dft()
do_fitting(dft)