using OrderedCollections

"""
    LazyBridgeOptimizer{OT<:MOI.ModelLike} <: AbstractBridgeOptimizer

The `LazyBridgeOptimizer` combines several bridges, which are added using the [`add_bridge`](@ref) function.
Whenever a constraint is added, it only attempts to bridge it if it is not supported by the internal model (hence its name `Lazy`).
When bridging a constraint, it selects the minimal number of bridges needed.
For instance, a constraint `F`-in-`S` can be bridged into a constraint `F1`-in-`S1` (supported by the internal model) using bridge 1 or
bridged into a constraint `F2`-in-`S2` (unsupported by the internal model) using bridge 2 which can then be
bridged into a constraint `F3`-in-`S3` (supported by the internal model) using bridge 3,
it will choose bridge 1 as it allows to bridge `F`-in-`S` using only one bridge instead of two if it uses bridge 2 and 3.
"""
mutable struct LazyBridgeOptimizer{OT<:MOI.ModelLike} <: AbstractBridgeOptimizer
    # Internal model
    model::OT
    # Bridged variables
    variable_map::Variable.Map
    var_to_name::Dict{MOI.VariableIndex, String}
    name_to_var::Union{Dict{String, MOI.VariableIndex}, Nothing}
    # Bridged constraints
    constraint_map::Constraint.Map
    con_to_name::Dict{MOI.ConstraintIndex, String}
    name_to_con::Union{Dict{String, MOI.ConstraintIndex}, Nothing}
    # Bridged objective
    objective_map::Objective.Map
    # Bellman-Ford
    # List of types of available bridges
    variable_bridge_types::Vector{Any}
    # (S,) -> Number of bridges that need to be used for constrained variables in `S`
    variable_dist::Dict{Tuple{DataType}, Int}
    # (S,) -> Bridge to be used for constrained variables in `S`
    variable_best::Dict{Tuple{DataType}, DataType}
    # List of types of available bridges
    constraint_bridge_types::Vector{Any}
    # (F, S) -> Number of bridges that need to be used for an `F`-in-`S` constraint
    constraint_dist::Dict{Tuple{DataType, DataType}, Int}
    # (F, S) -> Bridge to be used for an `F`-in-`S` constraint
    constraint_best::Dict{Tuple{DataType, DataType}, DataType}
    # List of types of available bridges
    objective_bridge_types::Vector{Any}
    # (F, S) -> Number of bridges that need to be used for an `F`-in-`S` constraint
    objective_dist::Dict{Tuple{DataType}, Int}
    # (F, S) -> Bridge to be used for an `F`-in-`S` constraint
    objective_best::Dict{Tuple{DataType}, DataType}
end
function LazyBridgeOptimizer(model::MOI.ModelLike)
    return LazyBridgeOptimizer{typeof(model)}(
        model,
        Variable.Map(), Dict{MOI.VariableIndex, String}(), nothing,
        Constraint.Map(), Dict{MOI.ConstraintIndex, String}(), nothing,
        Objective.Map(),
        Any[], Dict{Tuple{DataType}, Int}(),
        Dict{Tuple{DataType}, DataType}(),
        Any[], Dict{Tuple{DataType, DataType}, Int}(),
        Dict{Tuple{DataType, DataType}, DataType}(),
        Any[], Dict{Tuple{DataType}, Int}(),
        Dict{Tuple{DataType}, DataType}())
end

function Variable.bridges(bridge::LazyBridgeOptimizer)
    return bridge.variable_map
end
function Constraint.bridges(bridge::LazyBridgeOptimizer)
    return bridge.constraint_map
end
function Objective.bridges(b::LazyBridgeOptimizer)
    return b.objective_map
end

function _dist(b::LazyBridgeOptimizer, S::Type{<:MOI.AbstractSet})
    F = MOIU.variable_function_type(S)
    if MOI.supports_constraint(b.model, F, S)
        return 0
    else
        return min(get(b.variable_dist, (S,), typemax(Int)),
                   get(b.constraint_dist, (F, S), typemax(Int) - 1) + 1)
    end
end

function _dist(b::LazyBridgeOptimizer, F::Type{<:MOI.AbstractFunction}, S::Type{<:MOI.AbstractSet})
    if MOI.supports_constraint(b.model, F, S)
        return 0
    else
        return get(b.constraint_dist, (F, S), typemax(Int))
    end
end

function _dist(b::LazyBridgeOptimizer, F::Type{<:MOI.AbstractScalarFunction})
    if MOI.supports(b.model, MOI.ObjectiveFunction{F}())
        return 0
    else
        return get(b.objective_dist, (F,), typemax(Int))
    end
end

function _added_dist(b::LazyBridgeOptimizer, args...)
    dist = mapreduce(C -> _dist(b, C[1]), +,
                     added_constrained_variable_types(args...), init = 0)
    dist += mapreduce(C -> _dist(b, C[1], C[2]), +,
                      added_constraint_types(args...), init = 0)
    return dist
end
function added_dist(b::LazyBridgeOptimizer,
                    BT::Type{<:Union{Variable.AbstractBridge,
                                     Constraint.AbstractBridge}},
                    args...)
    return _added_dist(b, BT, args...)
end
function added_dist(b::LazyBridgeOptimizer,
                    BT::Type{<:Objective.AbstractBridge},
                    args...)
    F = set_objective_function_type(BT, args...)
    return _added_dist(b, BT, args...) + _dist(b, F)
end

function _supports_added_no_update(b::LazyBridgeOptimizer, args...)
    return all(C -> supports_no_update(b, C[1]), added_constrained_variable_types(args...)) &&
        all(C -> supports_no_update(b, C[1], C[2]), added_constraint_types(args...))
end
function supports_added_no_update(
    b::LazyBridgeOptimizer,
    BT::Type{<:Union{Variable.AbstractBridge,
                     Constraint.AbstractBridge}},
    args...
)
    return _supports_added_no_update(b, BT, args...)
end
function supports_added_no_update(
    b::LazyBridgeOptimizer,
    BT::Type{<:Objective.AbstractBridge},
    args...
)
    F = set_objective_function_type(BT, args...)
    return _supports_added_no_update(b, BT, args...) &&
        supports_no_update(b, F)
end

# Update `b.variable_dist`, `b.constraint_dist` `b.variable_best` and
# `b.constraint_best` for constrained variable types in `variables` and
# constraint types in `constraints`.
function update_dist!(b::LazyBridgeOptimizer, variables, constraints, objectives)
    # Bellman-Ford algorithm
    changed = true # Has b.constraint_dist changed in the last iteration ?
    while changed
        changed = false
        for BT in b.variable_bridge_types
            for (S,) in variables
                if Variable.supports_constrained_variable(BT, S) &&
                    supports_added_no_update(b, BT, S)
                    # Number of bridges needed using BT
                    dist = 1 + added_dist(b, BT, S)
                    # Is it better that what can currently be done ?
                    if dist < _dist(b, S)
                        b.variable_dist[(S,)] = dist
                        b.variable_best[(S,)] = Variable.concrete_bridge_type(BT, S)
                        changed = true
                    end
                end
            end
        end
        for BT in b.constraint_bridge_types
            for (F, S) in constraints
                if MOI.supports_constraint(BT, F, S) &&
                    supports_added_no_update(b, BT, F, S)
                    # Number of bridges needed using BT
                    dist = 1 + added_dist(b, BT, F, S)
                    # Is it better that what can currently be done ?
                    if dist < _dist(b, F, S)
                        b.constraint_dist[(F, S)] = dist
                        b.constraint_best[(F, S)] = Constraint.concrete_bridge_type(BT, F, S)
                        changed = true
                    end
                end
            end
        end
        for BT in b.objective_bridge_types
            for (F,) in objectives
                if Objective.supports_objective_function(BT, F) &&
                    supports_added_no_update(b, BT, F)
                    # Number of bridges needed using BT
                    dist = 1 + added_dist(b, BT, F)
                    # Is it better that what can currently be done ?
                    if dist < _dist(b, F)
                        b.objective_dist[(F,)] = dist
                        b.objective_best[(F,)] = Objective.concrete_bridge_type(BT, F)
                        changed = true
                    end
                end
            end
        end
    end
end

function fill_required!(required_variables::OrderedSet{Tuple{DataType}},
                        required_constraints::OrderedSet{Tuple{DataType, DataType}},
                        required_objectives::OrderedSet{Tuple{DataType}},
                        b::LazyBridgeOptimizer,
                        BT::Type{<:AbstractBridge})
    for C in added_constrained_variable_types(BT)
        # The constrained variables in `C[1]` might be supported while the
        # constraints `F`-in-`C[1]` are not supported. In this case, the second
        # `fill_requires!` would add `F`-in-`C[1]` while we know that:
        # * either `C[1]` is natively supported in which case we don't need to
        #   worry about `F`-in-`C[1]`;
        # * either, `C[1]` was found to be supported via variable bridges in
        #   but in this run of Bellman-Ford, `F`-in-`C[1]` was also in required
        #   so we know that it is not supported and we don't need to add it
        #   again in required as it would arrive to the same conclusion.
        if !supports_no_update(b, C[1])
            fill_required!(required_variables, required_constraints,
                           required_objectives, b, C[1])
            F = MOIU.variable_function_type(C[1])
            fill_required!(required_variables, required_constraints,
                           required_objectives, b, F, C[1])
        end
    end
    for C in added_constraint_types(BT)
        fill_required!(required_variables, required_constraints,
                       required_objectives, b, C[1], C[2])
    end
    if BT <: Objective.AbstractBridge
        fill_required!(required_variables, required_constraints, required_objectives, b,
                       set_objective_function_type(BT))
    end
end

function fill_required!(required_variables::OrderedSet{Tuple{DataType}},
                        required_constraints::OrderedSet{Tuple{DataType, DataType}},
                        required_objectives::OrderedSet{Tuple{DataType}},
                        b::LazyBridgeOptimizer, S::Type{<:MOI.AbstractSet})
    if supports_no_update(b, S)
        return # The constraint is supported
    end
    if (S,) in required_variables
        return # The requirements for these constrained variables have already been added or are being added
    end
    # The constrained variables are not supported yet, add in
    # `required_variables` the required variables types to bridge it.
    push!(required_variables, (S,))
    for BT in b.variable_bridge_types
        if Variable.supports_constrained_variable(BT, S)
            fill_required!(required_variables, required_constraints,
                           required_objectives, b,
                           Variable.concrete_bridge_type(BT, S))
        end
    end
end

function fill_required!(required_variables::OrderedSet{Tuple{DataType}},
                        required_constraints::OrderedSet{Tuple{DataType, DataType}},
                        required_objectives::OrderedSet{Tuple{DataType}},
                        b::LazyBridgeOptimizer, F::Type{<:MOI.AbstractFunction},
                        S::Type{<:MOI.AbstractSet})
    if supports_no_update(b, F, S)
        return # The constraint is supported
    end
    if (F, S) in required_constraints
        return # The requirements for this constraint have already been added or are being added
    end
    # The constraint is not supported yet, add
    # * in `required_variables` the required constrained variables types and
    # * in `required_constraints` the required constraint types
    # to bridge it.
    push!(required_constraints, (F, S))
    for BT in b.constraint_bridge_types
        if MOI.supports_constraint(BT, F, S)
            fill_required!(required_variables, required_constraints,
                           required_objectives, b,
                           Constraint.concrete_bridge_type(BT, F, S))
        end
    end
end

function fill_required!(required_variables::OrderedSet{Tuple{DataType}},
                        required_constraints::OrderedSet{Tuple{DataType, DataType}},
                        required_objectives::OrderedSet{Tuple{DataType}},
                        b::LazyBridgeOptimizer,
                        F::Type{<:MOI.AbstractScalarFunction})
    if supports_no_update(b, F)
        return # The objective is supported
    end
    if (F,) in required_objectives
        return # The requirements for this objective have already been added or are being added
    end
    # The objective is not supported yet, add
    # * in `required_variables` the required constrained variables types,
    # * in `required_constraints` the required constraint types and
    # * in `required_objectives` the required objective types
    # to bridge it.
    push!(required_objectives, (F,))
    for BT in b.objective_bridge_types
        if Objective.supports_objective_function(BT, F)
            fill_required!(required_variables, required_constraints,
                           required_objectives, b,
                           Objective.concrete_bridge_type(BT, F))
        end
    end
end

function required(b::LazyBridgeOptimizer, types::Tuple)
    required_variables = OrderedSet{Tuple{DataType}}()
    required_constraints = OrderedSet{Tuple{DataType, DataType}}()
    required_objectives = OrderedSet{Tuple{DataType}}()
    fill_required!(required_variables, required_constraints,
                   required_objectives, b, types...)
    return required_variables, required_constraints, required_objectives
end

# Compute dist[(F, S)], dist[(S,)] and best[(F, S)], best[(S,)]
function update!(b::LazyBridgeOptimizer, types::Tuple)
    update_dist!(b, required(b, types)...)
end

# After `add_bridge(b, BT)`, some constrained variables `(S,)` in
# `keys(b.variable_best)` or constraints `(F, S)` in `keys(b.constraint_best)`
# or `(F,)` in `keys(b.objective_best)` may be bridged
# with less bridges than `b.variable_dist[(S,)]`,
# `b.constraint_dist[(F, S)]` or `b.objective_dist[(F,)]` using `BT`.
function _update_key_dists!(b)
    # TODO we should call `fill_required!`.
    update_dist!(b, keys(b.variable_best), keys(b.constraint_best), keys(b.objective_best))
end

"""
    add_bridge(b::LazyBridgeOptimizer, BT::Type{<:Variable.AbstractBridge})

Enable the use of the variable bridges of type `BT` by `b`.
"""
function add_bridge(b::LazyBridgeOptimizer, BT::Type{<:Variable.AbstractBridge})
    push!(b.variable_bridge_types, BT)
    update_dist!(b, keys(b.variable_best), keys(b.constraint_best), keys(b.objective_best))
end

"""
    add_bridge(b::LazyBridgeOptimizer, BT::Type{<:Constraint.AbstractBridge})

Enable the use of the constraint bridges of type `BT` by `b`.
"""
function add_bridge(b::LazyBridgeOptimizer, BT::Type{<:Constraint.AbstractBridge})
    push!(b.constraint_bridge_types, BT)
    update_dist!(b, keys(b.variable_best), keys(b.constraint_best), keys(b.objective_best))
end

"""
    add_bridge(b::LazyBridgeOptimizer, BT::Type{<:Objective.AbstractBridge})

Enable the use of the objective bridges of type `BT` by `b`.
"""
function add_bridge(b::LazyBridgeOptimizer, BT::Type{<:Objective.AbstractBridge})
    push!(b.objective_bridge_types, BT)
    update_dist!(b, keys(b.variable_best), keys(b.constraint_best), keys(b.objective_best))
end

function _remove_bridge(bridge_types::Vector, BT::Type)
    i = findfirst(isequal(BT), bridge_types)
    if i === nothing
        error("Cannot remove bridge `BT` as it was never added or was already",
              " removed.")
    else
        deleteat!(bridge_types, i)
    end
    return
end
function _reset_dist(b::LazyBridgeOptimizer)
    empty!(b.variable_dist)
    empty!(b.variable_best)
    empty!(b.constraint_dist)
    empty!(b.constraint_best)
end

"""
    remove_bridge(b::LazyBridgeOptimizer, BT::Type{<:Variable.AbstractBridge})

Disable the use of the variable bridges of type `BT` by `b`.
"""
function remove_bridge(b::LazyBridgeOptimizer, BT::Type{<:Variable.AbstractBridge})
    _remove_bridge(b.variable_bridge_types, BT)
    _reset_dist(b)
end

"""
    remove_bridge(b::LazyBridgeOptimizer, BT::Type{<:Constraint.AbstractBridge})

Disable the use of the constraint bridges of type `BT` by `b`.
"""
function remove_bridge(b::LazyBridgeOptimizer, BT::Type{<:Constraint.AbstractBridge})
    _remove_bridge(b.constraint_bridge_types, BT)
    _reset_dist(b)
end

# It only bridges when the constraint is not supporting, hence the name "Lazy"
function is_bridged(b::LazyBridgeOptimizer, S::Type{<:MOI.AbstractSet})
    return !MOI.supports_constraint(b.model, MOIU.variable_function_type(S), S)
end
function is_bridged(b::LazyBridgeOptimizer, F::Type{<:MOI.AbstractFunction}, S::Type{<:MOI.AbstractSet})
    return !MOI.supports_constraint(b.model, F, S)
end
function is_bridged(b::LazyBridgeOptimizer, F::Type{<:MOI.AbstractScalarFunction})
    return !MOI.supports(b.model, MOI.ObjectiveFunction{F}())
end
# Same as supports_constraint but do not trigger `update!`. This is
# used inside `update!`.
function supports_no_update(b::LazyBridgeOptimizer, S::Type{<:MOI.AbstractSet})
    F = MOIU.variable_function_type(S)
    return MOI.supports_constraint(b.model, MOIU.variable_function_type(S), S) ||
        (S,) in keys(b.variable_best) || (F, S) in keys(b.constraint_best)
end
function supports_no_update(b::LazyBridgeOptimizer, F::Type{<:MOI.AbstractFunction}, S::Type{<:MOI.AbstractSet})
    return MOI.supports_constraint(b.model, F, S) || (F, S) in keys(b.constraint_best)
end
function supports_no_update(b::LazyBridgeOptimizer, F::Type{<:MOI.AbstractScalarFunction})
    return MOI.supports(b.model, MOI.ObjectiveFunction{F}()) || (F,) in keys(b.objective_best)
end

supports_constraint_bridges(::LazyBridgeOptimizer) = true
function supports_bridging_constrained_variable(
    b::LazyBridgeOptimizer, S::Type{<:MOI.AbstractSet}
)
    update!(b, (S,))
    return (S,) in keys(b.variable_best)
end
function supports_bridging_constraint(
    b::LazyBridgeOptimizer, F::Type{<:MOI.AbstractFunction},
    S::Type{<:MOI.AbstractSet}
)
    update!(b, (F, S))
    return (F, S) in keys(b.constraint_best)
end
function supports_bridging_objective_function(
    b::LazyBridgeOptimizer, F::Type{<:MOI.AbstractScalarFunction}
)
    update!(b, (F,))
    return (F,) in keys(b.objective_best)
end
function bridge_type(b::LazyBridgeOptimizer, S::Type{<:MOI.AbstractSet})
    update!(b, (S,))
    result = get(b.variable_best, (S,), nothing)
    if result === nothing
        throw(MOI.UnsupportedConstraint{MOIU.variable_function_type(S), S}())
    end
    return result
end
function bridge_type(b::LazyBridgeOptimizer, F::Type{<:MOI.AbstractFunction}, S::Type{<:MOI.AbstractSet})
    update!(b, (F, S))
    result = get(b.constraint_best, (F, S), nothing)
    if result === nothing
        throw(MOI.UnsupportedConstraint{F, S}())
    end
    return result
end
function bridge_type(b::LazyBridgeOptimizer, F::Type{<:MOI.AbstractScalarFunction})
    update!(b, (F,))
    result = get(b.objective_best, (F,), nothing)
    if result === nothing
        throw(MOI.UnsupportedAttribute(MOI.ObjectiveFunction{F}()))
    end
    return result
end

_func_name(::Type{Constraint.ScalarFunctionizeBridge}) = "SingleVariable", "constraint"
_func_name(::Type{Constraint.VectorFunctionizeBridge}) = "VectorOfVariables", "constraint"
_func_name(::Type{Objective.FunctionizeBridge}) = "SingleVariable", "objective"
function _functionize_bridge(bridge_types, target_type)
    index = findfirst(bridge_type -> bridge_type <: target_type,
                      bridge_types)
    if index === nothing
        func, name = _func_name(target_type)
        error("Need to apply a `$target_type` to a",
              " `$func` $name because the variable is",
              " bridged but no such $name bridge type was added. Add one",
              " with `add_bridge`.")
    end
    return bridge_types[index]
end
function constraint_vector_functionize_bridge(b::LazyBridgeOptimizer)
    return _functionize_bridge(b.constraint_bridge_types, Constraint.VectorFunctionizeBridge)
end
function constraint_scalar_functionize_bridge(b::LazyBridgeOptimizer)
    return _functionize_bridge(b.constraint_bridge_types, Constraint.ScalarFunctionizeBridge)
end
function objective_functionize_bridge(b::LazyBridgeOptimizer)
    return _functionize_bridge(b.objective_bridge_types, Objective.FunctionizeBridge)
end


struct ODEQuadratureAdjointSensitivityFunction{rateType,uType,uType2,UF,G,JC,GC,DG,TJ,SType,CV} <: SensitivityFunction
  uf::UF
  g::G
  J::TJ
  dg_val::uType
  jac_config::JC
  g_grad_config::GC
  sensealg::QuadratureAdjoint
  f_cache::rateType
  discrete::Bool
  y::uType2
  sol::SType
  dg::DG
  colorvec::CV
end

@noinline function ODEQuadratureAdjointSensitivityFunction(g,u0,p,sensealg,discrete,sol,dg,tspan,colorvec)
  numparams = length(p)
  numindvar = length(u0)
  # if there is an analytical Jacobian provided, we are not going to do automatic `jac*vec`
  f = sol.prob.f
  isautojacvec = DiffEqBase.has_jac(f) ? false : get_jacvec(sensealg)
  J = isautojacvec ? nothing : similar(u0, numindvar, numindvar)

  if !discrete
    if dg != nothing || isautojacvec
      pg = nothing
      pg_config = nothing
    else
      pg = UGradientWrapper(g,tspan[2],p)
      pg_config = build_grad_config(sensealg,pg,u0,p)
    end
  else
    pg = nothing
    pg_config = nothing
  end

  if isautojacvec
    jac_config = nothing
    uf = nothing
  else
    uf = DiffEqDiffTools.UJacobianWrapper(f,tspan[2],p)
    jac_config = build_jac_config(sensealg,uf,u0)
  end

  y = copy(sol(tspan[1])) # TODO: Has to start at interpolation value!
  dg_val = similar(u0, numindvar) # number of funcs size
  f_cache = deepcopy(u0)

  return ODEQuadratureAdjointSensitivityFunction(uf,pg,J,dg_val,
                                       jac_config,pg_config,
                                       sensealg,f_cache,
                                       discrete,y,sol,dg,
                                       colorvec)
end

# u = λ'
function (S::ODEQuadratureAdjointSensitivityFunction)(du,u,p,t)
  @unpack y, sol, J, uf, sensealg, f_cache, jac_config, discrete, dg, dg_val, g, g_grad_config = S
  idx = length(y)
  f = sol.prob.f
  isautojacvec = DiffEqBase.has_jac(f) ? false : get_jacvec(sensealg)
  sol(y,t)
  λ     = u
  dλ    = du

  if !isautojacvec
    if DiffEqBase.has_jac(f)
      f.jac(J,y,p,t) # Calculate the Jacobian into J
    else
      uf.t = t
      jacobian!(J, uf, y, f_cache, sensealg, jac_config)
    end
    mul!(dλ',λ',J)
  else
    _dy, back = Tracker.forward(y) do u
      if DiffEqBase.isinplace(sol.prob)
        out_ = map(zero, u)
        f(out_, u, p, t)
        Tracker.collect(out_)
      else
        vec(f(u, p, t))
      end
    end
    dλ[:] = Tracker.data(back(λ)[1])
  end

  dλ .*= -one(eltype(λ))

  if !discrete
    if dg != nothing
      dg(dg_val,y,p,t)
    else
      g.t = t
      gradient!(dg_val, g, y, sensealg, g_grad_config)
    end
    dλ .+= dg_val
  end

  nothing
end

# g is either g(t,u,p) or discrete g(t,u,i)
@noinline function ODEAdjointProblem(sol,sensealg::QuadratureAdjoint,g,
                                     t=nothing,dg=nothing,
                                     callback=CallbackSet())
  f = sol.prob.f
  tspan = (sol.prob.tspan[2],sol.prob.tspan[1])
  discrete = t != nothing

  p = sol.prob.p
  p === DiffEqBase.NullParameters() && error("Your model does not have parameters, and thus it is impossible to calculate the derivative of the solution with respect to the parameters. Your model must have parameters to use parameter sensitivity calculations!")

  u0 = zero(sol.prob.u0)

  len = length(u0)
  λ = similar(u0, len)
  sense = ODEQuadratureAdjointSensitivityFunction(g,u0,
                                        p,sensealg,discrete,
                                        sol,dg,tspan,f.colorvec)

  init_cb = t !== nothing && sol.prob.tspan[2] == t[end]
  cb = generate_callbacks(sense, g, λ, t, callback, init_cb)
  z0 = vec(zero(λ))
  ODEProblem(sense,z0,tspan,p,callback=cb)
end

struct AdjointSensitivityIntegrand{pType,uType,rateType,S,AS,PF,PJC,PJT}
  sol::S
  adj_sol::AS
  p::pType
  y::uType
  λ::uType
  pf::PF
  f_cache::rateType
  pJ::PJT
  paramjac_config::PJC
  sensealg::QuadratureAdjoint
end

function AdjointSensitivityIntegrand(sol,adj_sol,sensealg)
  prob = sol.prob
  @unpack f, p, tspan, u0 = prob
  y = similar(sol.prob.u0)
  λ = similar(adj_sol.prob.u0)
  # we need to alias `y`
  pf = DiffEqDiffTools.ParamJacobianWrapper(f,tspan[1],y)
  f_cache = similar(y)
  isautojacvec = get_jacvec(sensealg)
  pJ = isautojacvec ? nothing : similar(u0,length(u0),length(p))

  if DiffEqBase.has_paramjac(f) || isautojacvec
    paramjac_config = nothing
  else
    paramjac_config = build_param_jac_config(sensealg,pf,y,p)
  end
  AdjointSensitivityIntegrand(sol,adj_sol,p,y,λ,pf,f_cache,pJ,paramjac_config,sensealg)
end

function (S::AdjointSensitivityIntegrand)(out,t)
  @unpack y, λ, pJ, pf, p, f_cache, paramjac_config, sensealg, sol, adj_sol = S
  f = sol.prob.f
  sol(y,t)
  adj_sol(λ,t)
  λ .*= -one(eltype(λ))
  isautojacvec = get_jacvec(sensealg)
  # y is aliased
  pf.t = t

  if !isautojacvec
    if DiffEqBase.has_paramjac(f)
      f.paramjac(pJ,y,p,t) # Calculate the parameter Jacobian into pJ
    else
      jacobian!(pJ, pf, p, f_cache, sensealg, paramjac_config)
    end
    mul!(out',λ',pJ)
  else
    _, back = Tracker.forward(y, p) do u, p
      if DiffEqBase.isinplace(sol.prob)
        out_ = map(zero, u)
        f(out_, u, p, t)
        Tracker.collect(out_)
      else
        vec(f(u, p, t))
      end
    end
    out[:] = vec(Tracker.data(back(λ)[2]))
  end
  out'
end

function (S::AdjointSensitivityIntegrand)(t)
  out = similar(S.p)
  S(out,t)
end

function _adjoint_sensitivities_u0(sol,sensealg::QuadratureAdjoint,args...;kwargs...)
  error("Can't get sensitivities of u0 with quadrature.")
end

function _adjoint_sensitivities(sol,sensealg::QuadratureAdjoint,alg,g,
                                t=nothing,dg=nothing;
                                abstol=1e-6,reltol=1e-3,
                                iabstol=abstol, ireltol=reltol,
                                kwargs...)
  adj_prob = ODEAdjointProblem(sol,sensealg,g,t,dg)
  adj_sol = solve(adj_prob,alg;abstol=abstol,reltol=reltol,
                               save_everystep=true,save_start=true,kwargs...)
  integrand = AdjointSensitivityIntegrand(sol,adj_sol,sensealg)

  if t === nothing
    res,err = quadgk(integrand,sol.prob.tspan[1],sol.prob.tspan[2],
                   atol=iabstol,rtol=ireltol)
  else
    res = zero(integrand.p)'
    for i in 1:length(t)-1
      res .+= quadgk(integrand,t[i],t[i+1],
                     atol=iabstol,rtol=ireltol)[1]
    end
    if t[1] != sol.prob.tspan[1]
      res .+= quadgk(integrand,sol.prob.tspan[1],t[1],
                     atol=iabstol,rtol=ireltol)[1]
    end
  end
  res
end


module Morbit

export morbit_formatter;
export MixedMOP, optimize, AlgorithmConfig, add_objective!, add_vector_objective!;
export ExactConfig, TaylorConfig, RbfConfig, LagrangeConfig;
export save_config, save_database, save_iter_data;
export load_config, load_database, load_iter_data;

using Printf: @sprintf

# steepest descent and directed search
using LinearAlgebra: norm, pinv
import JuMP;
import OSQP;

# LagrangeModels and PS
import NLopt;

using Parameters: @with_kw, @unpack, @pack!
using MathOptInterface;
const MOI = MathOptInterface;

import ThreadSafeDicts: ThreadSafeDict;
using Memoize;

import UUIDs;

import FiniteDiff#erences
const FD = FiniteDiff#erences

import ForwardDiff
const AD = ForwardDiff

import Logging: LogLevel, @logmsg
import Logging

include("custom_logging.jl")

include("shorthands.jl");
include("Interfaces.jl");
include("diff_wrappers.jl");

# implementations (order should not matter)
include("VectorObjectiveFunction.jl");
include("MixedMOP.jl");

include("ResultImplementation.jl")
include("DataBaseImplementation.jl")
include("IterDataImplementation.jl")
include("SurrogatesImplementation.jl");

include("ConfigImplementations.jl")

# utilities
include("adding_objectives.jl");
include("descent.jl")
include("saving.jl")
include("utilities.jl")

function shrink_radius( ac :: AbstractConfig, Δ :: NumOrVec, steplength :: NumOrVec)
    if radius_update_method(ac) == :standard
        return Δ .* _gamma_shrink(ac);
    elseif radius_update_method(ac) == :steplength
        return steplength .* _gamma_shrink(ac);
    end
end

function shrink_radius_much( ac :: AbstractConfig, Δ :: NumOrVec, steplength :: NumOrVec)
    if radius_update_method(ac) == :standard
        return Δ .* _gamma_shrink_much(ac);
    elseif radius_update_method(ac) == :steplength
        return steplength .* _gamma_shrink_much(ac);
    end
end

function grow_radius( ac :: AbstractConfig, Δ :: NumOrVec, steplength :: NumOrVec)
    if radius_update_method(ac) == :standard
        return min( Δᵘ(ac), _gamma_grow(ac) .* Δ )
    elseif radius_update_method(ac) == :steplength
        return min( Δᵘ(ac), ( _gamma_grow .+ steplength ./ Δ ) .* Δ );
    end
end

"""
Perform initialization of the data passed to `optimize` function.
"""
function initialize_data( mop :: AbstractMOP, x0 :: RVec, fx0 :: Vec = Float32[];
    algo_config :: Union{AbstractConfig, Nothing} = nothing,
    populated_db :: Union{AbstractDB, Nothing} = nothing )

    if num_objectives(mop) == 0
        error("`mop` has no objectives!")
    end

    @warn "The evaluation counter of `mop` is reset."
    reset_evals!( mop );
    # for backwards-compatibility with unconstrained problems:
    if num_vars(mop) == 0
        MOI.add_variables(mop, length(x0))
    end

    # initialize first iteration site
    @assert !isempty( x0 ) "Please provide a non-empty feasible starting point `x0`."

    x_scaled = scale( x0, mop );
    tfn = TransformerFn(mop)

    # initalize first objective vector
    if isempty( fx0 )
        # if no starting function value was provided, eval objectives
        fx_sorted = eval_all_objectives( mop, tfn, x_scaled );
    else
        fx_sorted = apply_internal_sorting( fx0, mop );
    end

    # ensure at least half-precision
    F = Base.promote_eltype( x_scaled, fx_sorted, Float32 )
    x = F.(x_scaled)
    fx = F.(fx_sorted)

    if isnothing( algo_config )
        ac = DefaultConfig{F}()
    else
        ac = algo_config
    end

    # initialize database
    if !isnothing(populated_db)
        # has a database been provided? if yes, prepare
        data_base = populated_db;
        transform!( data_base, mop );
    else
        data_base = init_db( use_db(ac) );
        set_transformed!(data_base, true)
    end

    # make sure, x & fx are in database
    start_res = init_res(Result{F}, x, fx )
    ensure_contains_result!(data_base, start_res);

    iter_data = init_iter_data(IterData, x, fx, Δ⁰(ac));
    xᵗ_index!(iter_data, get_id(start_res));

    # initialize surrogate models
    sc = init_surrogates( mop, iter_data, ac );

    return (mop, tfn, ac, iter_data, sc)
end

function iterate!( iter_data :: AbstractIterData, mop :: AbstractMOP, tfn :: TransformerFn,
    sc :: SurrogateContainer, algo_config :: AbstractConfig )

    # config params
    ν_succ = _nu_success( algo_config );
    ν_acc = _nu_accept( algo_config );
    μ = _mu( algo_config );
    _beta = _beta( algo_config );
    _eps_crit = _eps_crit( algo_config );
    _gamma_crit = _gamma_crit( algo_config );
    count_nonlin = count_nonlinear_iterations( algo_config );

    # read iter data to handy variables
    x = xᵗ(iter_data);
    fx = fxᵗ(iter_data);
    Δ = Δᵗ(iter_data);

    # check some stopping conditions
    if num_iterations(iter_data) >= max_iter(algo_config)
        @logmsg loglevel1 "Stopping. Maximum number of iterations reached."
        return true;
    end

    if !_budget_okay(mop, algo_config)
        @logmsg loglevel1 "Stopping. Computational budget is exhausted."
        return true; # TODO stop codes
    end

    if Δ_abs_test( Δ, algo_config )
        return true;
    end

    # set iteration counter
    if it_stat(iter_data) != MODELIMPROVING || count_nonlin
        inc_iterations!(iter_data)
        set_model_improvements!(iter_data, 0);
    else
        inc_model_improvements!(iter_data)
    end

    @logmsg loglevel1 """\n
        |--------------------------------------------
        |Iteration $(num_iterations(iter_data)).$(num_model_improvements(iter_data))
        |--------------------------------------------
        |  Current trust region radius is $(Δ).
        |  Current number of function evals is $(num_evals(mop)).
        |  Iterate is $(_prettify(unscale(x, mop)))
        |  Values are $(_prettify(reverse_internal_sorting(fx,mop)))
        |--------------------------------------------
    """

    # update surrogate models
    if num_iterations(iter_data) > 1
        if it_stat == MODELIMPROVING
            improve_surrogates!( sc, mop, iter_data, algo_config; ensure_fully_linear = false );
        else
            update_surrogates!( sc, mop, iter_data, algo_config; ensure_fully_linear = false );
        end
    end

    # calculate descent step and criticality value
    ω, x₊, mx₊, steplength = compute_descent_step( algo_config, mop, iter_data, sc )
    @logmsg loglevel1 "Criticality is ω = $(ω)."

    # stop at the end of the this loop run?
    if ω_Δ_rel_test(ω, Δ, algo_config) || ω_abs_test( ω, algo_config )
        return true;
    end

    # Criticallity test
    _fully_linear = fully_linear(sc)
    num_critical_loops = 0;

    if ω <= _eps_crit && (!_fully_linear || all(Δ .> mu * ω))
        @logmsg loglevel1 "Entered Criticallity Test."
        if !_fully_linear
            @logmsg loglevel1 "Ensuring all models to be fully linear."
            update_surrogates!( sc, mop, iter_data, algo_config; ensure_fully_linear = true );

            ω, x₊, mx₊, steplength = compute_descent_step(algo_config,mop,iter_data,sc);
            if !fully_linear(sc)
                @logmsg loglevel1 "Could not make all models fully linear. Trying one last descent step."
                @goto MAIN;
            end
        end

        while all(Δᵗ(iter_data) .> mu * ω)
            @logmsg loglevel1 "Criticality loop $(num_critical_loops + 1)."
            if num_critical_loops >= max_critical_loops(algo_config)
                @logmsg loglevel1 "Maximum number ($(max_critical_loops(algo_config))) of critical loops reached. Exiting..."
                @goto EXIT_MAIN
            end
            if !_budget_okay(mop, algo_config)
                @logmsg loglevel1 "Computational budget exhausted. Exiting…"
                @goto EXIT_MAIN
            end

            # shrink radius
            Δᵗ!( iter_data, Δᵗ(iter_data) .* _gamma_crit );
            # make model linear
            update_surrogates!( sc, mop, iter_data, algo_config; ensure_fully_linear = true );
            # (re)calculate criticality
            # TODO make backtracking optional and don't do here
            ω, data = get_criticality

            if Δ_abs_test( Δᵗ(iter_data) , algo_config ) ||
                ω_Δ_rel_test(ω, Δᵗ(iter_data), algo_config) || ω_abs_test( ω, algo_config )

                @goto EXIT_MAIN
            end

            if !fully_linear(sc)
                @logmsg loglevel1 "Could not make all models fully linear. Trying one last descent step."
                @goto MAIN;
            end

            num_critical_loops += 1;
        end
        "Exiting after $(num_critical_loops) loops with ω = $(ω) and Δ = $(Δᵗ(iter_data))."
        @goto MAIN
        @label EXIT_MAIN
        return true;
    end# Crit test if

    @label MAIN # re-entry point after criticality test

    mx = eval_models(sc, x);
    fx₊ = eval_all_objectives(mop, tfn, x₊);

    if strict_acceptance_test( algo_config )
        ρ = minimum( (fx .- fx₊) ./ (mx .- mx₊) )
    else
        ρ = (maximum(fx) - maximum( fx₊ ))/(maximum(mx)-maximum(mx₊))
    end

    @logmsg loglevel2 """\n
    Attempting descent of length $steplength.
    | f(x)  | $(_prettify(fx))
    | f(x₊) | $(_prettify(fx₊))
    | m(x)  | $(_prettify(mx))
    | m(x₊) | $(_prettify(mx₊))
    The error betwenn f(x) and m(x) is $(sum(abs.(fx .- mx))).
    $(strict_acceptance_test(algo_config) ? "All" : "One") of the components must decrease.
    Thus, ρ is $ρ.
    """

    # update trust region radius
    ρ = isnan(ρ) ? -Inf : ρ;
    old_Δ = copy(Δᵗ(iter_data));  # if it was changed in criticality test
    if ρ >= ν_succ
        if Δ < _beta * ω
            Δᵗ!(iter_data, grow_radius(algo_config, Δ, steplength) );
        end
        it_stat!(iter_data, SUCCESSFULL)
    else
        if fully_linear(sc)
            if ρ < ν_acc
                Δᵗ!(iter_data, shrink_radius_much(algo_config, Δ, steplength) );
                it_stat!(iter_data, INACCEPTABLE)
            else
                Δᵗ!(iter_data, shrink_radius(algo_config, Δ, steplength) );
                it_stat!(iter_data, ACCEPTABLE)
            end
        else
            it_stat!(iter_data, MODELIMPROVING)
        end
    end

    # accept x?
    old_x_index = xᵗ_index(iter_data); # for stamp!ing
    if it_stat(iter_data) == SUCCESSFULL || it_stat(iter_data) == ACCEPTABLE
        trial_point_index = set_next_iterate!(iter_data, x₊, fx₊);
    else
        trial_point_index = keep_current_iterate!(iter_data, x₊, fx₊);
    end
    @logmsg loglevel1 """\n
        The trial point was $( (it_stat(iter_data) == SUCCESSFULL) || (it_stat(iter_data) == ACCEPTABLE ) ? "" : "not ")accepted.
        The iteration is $(it_stat(iter_data)).
        Moreover, the radius was updated as below:
        old radius : $(old_Δ)
        new radius : $(Δᵗ(iter_data)) ($(round(Δᵗ(iter_data)/old_Δ * 100; digits=1)) %)
    """

    stamp!(db(iter_data), Dict(
            "iter_status" => it_stat(iter_data),
            "xᵗ_index" => old_x_index,
            "xᵗ₊_index" => trial_point_index,
            "ρ" => ρ,
            "Δ" => Δᵗ(iter_data),
            "ω" => ω,
            "num_critical_loops" => num_critical_loops,
            "stepsize" => steplength,
            "model_meta" => [ deepcopy(sw.meta) for sw ∈ sc.surrogates ]
        )
    );

    if x_tol_rel_test( x, x₊, algo_config  ) ||
        x_tol_abs_test( x, x₊, algo_config ) ||
        f_tol_rel_test( fx, fx₊, algo_config  ) ||
        f_tol_abs_test( fx, fx₊, algo_config )
        return true;
    end
    return false
end

function get_return_values( iter_data :: AbstractIterData, mop :: AbstractMOP)
    ret_x = unscale(xᵗ(iter_data),mop);
    ret_fx = reverse_internal_sorting(fxᵗ(iter_data),mop);
    # TODO make tests for this: are returen values the same before and after `finalize_iter_data!` ?
    return (ret_x, ret_fx)
end

function finalize_iter_data!( iter_data :: AbstractIterData, mop :: AbstractMOP )
    if !isnothing( db(iter_data) )
        untransform!(db(iter_data), mop)
    end
    nothing
end

############################################
function optimize( mop :: AbstractMOP, x0 :: Vec, fx0 :: Vec = Float16[];
    algo_config :: Union{Nothing, AbstractConfig} = nothing,
    populated_db :: Union{AbstractDB, Nothing} = nothing # TODO make passing of AbstractIterData possible
    )

    mop, tfn, ac, iter_data, sc = initialize_data( mop, x0, fx0; algo_config, populated_db )

    @logmsg loglevel1 _stop_info_str( ac, mop )

    @logmsg loglevel1 "Entering main optimization loop."
    while true
        abbort = iterate!(iter_data, mop, tfn, sc, ac)
        abbort && break;
    end# while

    ret_x, ret_fx = get_return_values( iter_data, mop )
    # unscale sites and re-sort values to return to user
    finalize_iter_data!(iter_data, mop)

    # @assert all((ret_x, ret_fx) .== get_return_values( iter_data, mop ))

    @logmsg loglevel1 _fin_info_str(iter_data, mop)

    return ret_x, ret_fx, iter_data
end# function optimize

end#module


function Base.sort!(dt::DataTable; cols=Any[], alg=nothing,
                    lt=isless, by=identity, rev=false, order=Forward)
    if !(isa(by, Function) || eltype(by) <: Function)
        msg = "'by' must be a Function or a vector of Functions. Perhaps you wanted 'cols'."
        throw(ArgumentError(msg))
    end
    ord = ordering(dt, cols, lt, by, rev, order)
    _alg = Sort.defalg(dt, ord; alg=alg, cols=cols)
    sort!(dt, _alg, ord)
end

function Base.sort!(dt::DataTable, a::Base.Sort.Algorithm, o::Base.Sort.Ordering)
    p = sortperm(dt, a, o)
    pp = similar(p)
    c = columns(dt)

    for (i,col) in enumerate(c)
        # Check if this column has been sorted already
        if any(j -> c[j]===col, 1:i-1)
            continue
        end

        copy!(pp,p)
        Base.permute!!(col, pp)
    end
    dt
end


using Test
using MPI
using CLIMA
using CLIMA.Mesh.Topologies
using CLIMA.Mesh.Grids
using CLIMA.Mesh.Grids: VerticalDirection, HorizontalDirection, EveryDirection
using CLIMA.VTK
using Logging
using Printf
using LinearAlgebra
using CLIMA.DGmethods: DGModel, init_ode_state, LocalGeometry, courant
using CLIMA.DGmethods.NumericalFluxes: Rusanov, CentralNumericalFluxGradient,
                                       CentralNumericalFluxDiffusive
using CLIMA.Courant
using CLIMA.PlanetParameters: kappa_d
using CLIMA.Atmos: AtmosModel,
                   AtmosAcousticLinearModel, RemainderModel,
                   FlatOrientation,
                   NoReferenceState, ReferenceState,
                   DryModel, NoRadiation, PeriodicBC, NoPrecipitation,
                   Gravity, HydrostaticState, IsothermalProfile,
                   ConstantViscosityWithDivergence, vars_state, soundspeed
using CLIMA.Atmos
using CLIMA.ODESolvers
const ArrayType = CLIMA.array_type()

using CLIMA.MoistThermodynamics: air_density, total_energy, internal_energy,
                                 soundspeed_air

using CLIMA.VariableTemplates: Vars
using StaticArrays

function initialcondition!(bl, state, aux, coords, t)
    FT = eltype(state)

    p∞::FT = 10 ^ 5
    T∞::FT = 300
    translation_speed::FT = 150
    translation_angle::FT = pi / 4
    α = translation_angle
    u∞ = SVector(translation_speed * coords[1], translation_speed * coords[1], 0)

    u = u∞
    T = T∞
    # adiabatic/isentropic relation
    p = p∞ * (T / T∞) ^ (FT(1) / kappa_d)
    ρ = air_density(T, p)

    state.ρ = ρ
    state.ρu = ρ * u
    e_kin = u' * u / 2
    state.ρe = ρ * total_energy(e_kin, FT(0), T)

    nothing
end

let
    # boiler plate MPI stuff
    CLIMA.init()
    mpicomm = MPI.COMM_WORLD

    # Mesh generation parameters
    N = 4
    Nq = N+1
    Neh = 10
    Nev = 4

    @testset "$(@__FILE__) DGModel matrix" begin
        for FT in (Float64, Float32)
            for dim = (2, 3)
                if dim == 2
                    brickrange = (range(FT(0); length=Neh+1, stop=1),
                                  range(FT(1); length=Nev+1, stop=2))
                elseif dim == 3
                    brickrange = (range(FT(0); length=Neh+1, stop=1),
                                  range(FT(0); length=Neh+1, stop=1),
                                  range(FT(1); length=Nev+1, stop=2))
                end

                topl = StackedBrickTopology(mpicomm, brickrange)

                function warpfun(ξ1, ξ2, ξ3)
                    FT = eltype(ξ1)

                    ξ1 ≥ FT(1//2) && (ξ1 = FT(1//2) + 2(ξ1 - FT(1//2)))
                    if dim == 2
                        ξ2 ≥ FT(3//2) && (ξ2 = FT(3//2) + 2(ξ2 - FT(3//2)))
                    elseif dim == 3
                        ξ2 ≥ FT(1//2) && (ξ2 = FT(1//2) + 2(ξ2 - FT(1//2)))
                        ξ3 ≥ FT(3//2) && (ξ3 = FT(3//2) + 2(ξ3 - FT(3//2)))
                    end
                    return (ξ1, ξ2, ξ3)
                end

                grid = DiscontinuousSpectralElementGrid(topl,
                                                        FloatType = FT,
                                                        DeviceArray = ArrayType,
                                                        polynomialorder = N)
                                                #meshwarp = warpfun)

                model = AtmosModel{FT}(AtmosLESConfiguration;
                                       ref_state=NoReferenceState(),
                                       turbulence=ConstantViscosityWithDivergence(FT(2)),
                                       moisture=DryModel(),
                                       source=Gravity(),
                                       boundarycondition=PeriodicBC(),
                                       init_state=initialcondition!)

                dg = DGModel(model, grid, Rusanov(), CentralNumericalFluxDiffusive(),
                             CentralNumericalFluxGradient())

                Δt = FT(1e-11)#FT(1//2)
                #Δt = FT(0.00001) #FT(1//10)

                Q = init_ode_state(dg, FT(0))
                solver = LSRK54CarpenterKennedy(dg, Q; dt = Δt, t0 = 0)
                solve!(Q, solver; timeend=Δt)
                Δx = min_node_distance(grid, EveryDirection())
                Δx_v = min_node_distance(grid, VerticalDirection())
                Δx_h = min_node_distance(grid, HorizontalDirection())

                T∞ = FT(300)
                translation_speed = FT(212.13203435596427)
                diff_speed_h = 734.8469228349534
                diff_speed_v = 0
                c   = Δt*(translation_speed + soundspeed_air(T∞))/Δx
                c_h = Δt*(translation_speed + soundspeed_air(T∞))/Δx_h
                c_v = Δt*(soundspeed_air(T∞))/Δx_v
                d_h = Δt*diff_speed_h/Δx_h^2
                d_v = Δt*diff_speed_v/Δx_v^2

                if (FT==Float64 && dim == 2)
                    @test abs(courant(nondiffusive_courant, dg, model, Q, Δt, HorizontalDirection()) - c_h) <= 1e-15
                    @test abs(courant(diffusive_courant,    dg, model, Q, Δt, HorizontalDirection()) - d_h) <= 1e-4
                    @test abs(courant(nondiffusive_courant, dg, model, Q, Δt, VerticalDirection()) - c_v) <= 1e-16
                    @test abs(courant(diffusive_courant,    dg, model, Q, Δt, VerticalDirection()) - d_v) <= 1e-8
                elseif (FT==Float64 && dim == 3)
                    @test abs(courant(nondiffusive_courant, dg, model, Q, Δt, HorizontalDirection()) - c_h) <= 1e-11
                    @test abs(courant(diffusive_courant,    dg, model, Q, Δt, HorizontalDirection()) - d_h) <= 1e-4
                    @test abs(courant(nondiffusive_courant, dg, model, Q, Δt, VerticalDirection()) - c_v) <= 1e-11
                    @test abs(courant(diffusive_courant,    dg, model, Q, Δt, VerticalDirection()) - d_v) <= 1e-8
                elseif (dim == 2)
                    @test abs(courant(nondiffusive_courant, dg, model, Q, Δt, HorizontalDirection()) - c_h) <= 1e-11
                    @test abs(courant(diffusive_courant,    dg, model, Q, Δt, HorizontalDirection()) - d_h) <= 1e-4
                    @test abs(courant(nondiffusive_courant, dg, model, Q, Δt, VerticalDirection()) - c_v) <= 1e-16
                    @test abs(courant(diffusive_courant,    dg, model, Q, Δt, VerticalDirection()) - d_v) <= 1e-8
                else
                    @test abs(courant(nondiffusive_courant, dg, model, Q, Δt, HorizontalDirection()) - c_h) <= 1e-11
                    @test abs(courant(diffusive_courant,    dg, model, Q, Δt, HorizontalDirection()) - d_h) <= 1e-4
                    @test abs(courant(nondiffusive_courant, dg, model, Q, Δt, VerticalDirection()) - c_v) <= 1e-11
                    @test abs(courant(diffusive_courant,    dg, model, Q, Δt, VerticalDirection()) - d_v) <= 2e-8
                end
            end
        end
    end
end

nothing


# Mean-field variational update rules
@naiveVariationalRule(:node_type     => GeneralisedFilterX,
                      :outbound_type => Message{GaussianMeanPrecision},
                      :inbound_types => (Nothing, ProbabilityDistribution, ProbabilityDistribution, ProbabilityDistribution, ProbabilityDistribution, ProbabilityDistribution),
                      :name          => VariationalGFXOutNPPPPP)

@naiveVariationalRule(:node_type     => GeneralisedFilterX,
                      :outbound_type => Message{GaussianWeightedMeanPrecision},
                      :inbound_types => (ProbabilityDistribution, Nothing, ProbabilityDistribution, ProbabilityDistribution, ProbabilityDistribution, ProbabilityDistribution),
                      :name          => VariationalGFXIn1PNPPPP)

@naiveVariationalRule(:node_type     => GeneralisedFilterX,
                      :outbound_type => Message{GaussianWeightedMeanPrecision},
                      :inbound_types => (ProbabilityDistribution, ProbabilityDistribution, Nothing, ProbabilityDistribution, ProbabilityDistribution, ProbabilityDistribution),
                      :name          => VariationalGFXIn2PPNPPP)

@naiveVariationalRule(:node_type     => GeneralisedFilterX,
                      :outbound_type => Message{GaussianWeightedMeanPrecision},
                      :inbound_types => (ProbabilityDistribution, ProbabilityDistribution, ProbabilityDistribution, Nothing, ProbabilityDistribution, ProbabilityDistribution),
                      :name          => VariationalGFXIn3PPPNPP)

@naiveVariationalRule(:node_type     => GeneralisedFilterX,
                      :outbound_type => Message{GaussianWeightedMeanPrecision},
                      :inbound_types => (ProbabilityDistribution, ProbabilityDistribution, ProbabilityDistribution, ProbabilityDistribution, Nothing, ProbabilityDistribution),
                      :name          => VariationalGFXIn4PPPPNP)

@naiveVariationalRule(:node_type     => GeneralisedFilterX,
                      :outbound_type => Message{Gamma},
                      :inbound_types => (ProbabilityDistribution, ProbabilityDistribution, ProbabilityDistribution, ProbabilityDistribution, ProbabilityDistribution, Nothing),
                      :name          => VariationalGFXIn5PPPPPN)

# Structured updates
#TODO


module Model

using Compat
import ForwardDiff
using StaticArrays

# parameter types
export Image, SkyPatch, PsfComponent,
       GalaxyComponent, GalaxyPrototype,
       PriorParams, CanonicalParams, BrightnessParams, StarPosParams,
       GalaxyPosParams, GalaxyShapeParams,
       PsfParams, RawPSF, CatalogEntry,
       ParamSet, SkyIntensity

# functions
export align

# constants
export NUM_COLOR_COMPONENTS, NUM_SOURCE_TYPES, NUM_BANDS,
       BAND_LETTERS, galaxy_prototypes, prior,
       shape_standard_alignment,
       brightness_standard_alignment,
       gal_shape_alignment,
       ids_names,
       ids, star_ids, gal_ids, gal_shape_ids, psf_ids, bids

import Base: convert, length, +
import Distributions
import FITSIO, WCS
import WCS.WCSTransform
import ..Log
import ..Celeste: Const, @aliasscope, @unroll_loop
import ..SensitiveFloats
import ..SensitiveFloats: SensitiveFloat
import ..BivariateNormals
using ..BivariateNormals: BivariateNormalDerivatives, BvnComponent,
                          GalaxySigmaDerivs, get_bvn_cov, eval_bvn_pdf!,
                          get_bvn_derivs!, transform_bvn_derivs!,
                          transform_bvn_ux_derivs!

const cfgdir = joinpath(Pkg.dir("Celeste"), "cfg")

const BAND_LETTERS = "ugriz"

# The number of bands (colors).
const NUM_BANDS = length(BAND_LETTERS)

include("model/light_source_model.jl")
include("model/psf_model.jl")
include("model/image_model.jl")
include("model/param_set.jl")
include("model/imaged_sources.jl")
include("model/wcs_utils.jl")
include("model/fsm_util.jl")
include("model/log_prob.jl")

end  # module


## ExoplanetsSysSim/src/transit_observations_unused.jl
## (c) 2015 Eric B. Ford

abstract EphemerisAbstract                     # Abstract, so we can generalize to allow for TTVs.  # TODO OPT: Make sure not a big speed penalty
immutable EphemerisLinear <: EphemerisAbstract  # Currently unused
  P::Float64
  t0::Float64
end

immutable TransitShape                       # Currently Unused
  depth::Float64        # Fractional
  duration::Float64     # days               #  QUERY: What definition of duration here?  Since we're going to use expressions from Price & Rogers for measurement uncertainties, let's treat this as the Full-width, half-max-duration  (until further notice)
  b::Float64            # in stellar radii
end

immutable TransitParameter{EphemT}           # Currently Unusued
  ephem::EphemT
  shape::TransitShape
end

if false  # Warning only works if TransitPlanetObs is mutable.  For testing purposes only
function set!(obs::TransitPlanetObs, P::Float64, t0::Float64, d::Float64, D::Float64)
  obs.period = P
  obs.t0 = t0
  obs.depth = d
  obs.duration = D
end
end


@testset "Test Compiler" begin
    @testset "Integer Arithmetic" begin
        for (input, expected_constants, expected_instructions) in [
            (
                "1; 2",
                [1, 2],
                [
                    m.make(m.OpConstant, 0),
                    m.make(m.OpPop),
                    m.make(m.OpConstant, 1),
                    m.make(m.OpPop),
                ],
            ),
            (
                "1 + 2",
                [1, 2],
                [
                    m.make(m.OpConstant, 0),
                    m.make(m.OpConstant, 1),
                    m.make(m.OpAdd),
                    m.make(m.OpPop),
                ],
            ),
            (
                "1 - 2",
                [1, 2],
                [
                    m.make(m.OpConstant, 0),
                    m.make(m.OpConstant, 1),
                    m.make(m.OpSub),
                    m.make(m.OpPop),
                ],
            ),
            (
                "1 * 2",
                [1, 2],
                [
                    m.make(m.OpConstant, 0),
                    m.make(m.OpConstant, 1),
                    m.make(m.OpMul),
                    m.make(m.OpPop),
                ],
            ),
            (
                "1 / 2",
                [1, 2],
                [
                    m.make(m.OpConstant, 0),
                    m.make(m.OpConstant, 1),
                    m.make(m.OpDiv),
                    m.make(m.OpPop),
                ],
            ),
            ("-1", [1], [m.make(m.OpConstant, 0), m.make(m.OpMinus), m.make(m.OpPop)]),
        ]
            run_compiler_tests(input, expected_constants, expected_instructions)
        end
    end

    @testset "Boolean Expressions" begin
        for (input, expected_constants, expected_instructions) in [
            ("true", [], [m.make(m.OpTrue), m.make(m.OpPop)]),
            ("false", [], [m.make(m.OpFalse), m.make(m.OpPop)]),
            (
                "1 == 2",
                [1, 2],
                [
                    m.make(m.OpConstant, 0),
                    m.make(m.OpConstant, 1),
                    m.make(m.OpEqual),
                    m.make(m.OpPop),
                ],
            ),
            (
                "1 != 2",
                [1, 2],
                [
                    m.make(m.OpConstant, 0),
                    m.make(m.OpConstant, 1),
                    m.make(m.OpNotEqual),
                    m.make(m.OpPop),
                ],
            ),
            (
                "1 < 2",
                [1, 2],
                [
                    m.make(m.OpConstant, 0),
                    m.make(m.OpConstant, 1),
                    m.make(m.OpLessThan),
                    m.make(m.OpPop),
                ],
            ),
            (
                "1 > 2",
                [1, 2],
                [
                    m.make(m.OpConstant, 0),
                    m.make(m.OpConstant, 1),
                    m.make(m.OpGreaterThan),
                    m.make(m.OpPop),
                ],
            ),
            (
                "true == false",
                [],
                [m.make(m.OpTrue), m.make(m.OpFalse), m.make(m.OpEqual), m.make(m.OpPop)],
            ),
            (
                "true != false",
                [],
                [
                    m.make(m.OpTrue),
                    m.make(m.OpFalse),
                    m.make(m.OpNotEqual),
                    m.make(m.OpPop),
                ],
            ),
            ("!true", [], [m.make(m.OpTrue), m.make(m.OpBang), m.make(m.OpPop)]),
        ]
            run_compiler_tests(input, expected_constants, expected_instructions)
        end
    end

    @testset "Conditionals" begin
        for (input, expected_constants, expected_instructions) in [
            (
                "if (true) {};",
                [],
                [
                    m.make(m.OpTrue),
                    m.make(m.OpJumpNotTruthy, 8),
                    m.make(m.OpNull),
                    m.make(m.OpJump, 9),
                    m.make(m.OpNull),
                    m.make(m.OpPop),
                ],
            ),
            (
                "if (true) { 10 }; 3333",
                [10, 3333],
                [
                    m.make(m.OpTrue),
                    m.make(m.OpJumpNotTruthy, 10),
                    m.make(m.OpConstant, 0),
                    m.make(m.OpJump, 11),
                    m.make(m.OpNull),
                    m.make(m.OpPop),
                    m.make(m.OpConstant, 1),
                    m.make(m.OpPop),
                ],
            ),
            (
                "if (true) { 10 } else { 20 }; 3333",
                [10, 20, 3333],
                [
                    m.make(m.OpTrue),
                    m.make(m.OpJumpNotTruthy, 10),
                    m.make(m.OpConstant, 0),
                    m.make(m.OpJump, 13),
                    m.make(m.OpConstant, 1),
                    m.make(m.OpPop),
                    m.make(m.OpConstant, 2),
                    m.make(m.OpPop),
                ],
            ),
        ]
            run_compiler_tests(input, expected_constants, expected_instructions)
        end
    end

    @testset "Global Let Statements" begin
        for (input, expected_constants, expected_instructions) in [
            (
                "let one = 1;\nlet two = 2;\n",
                [1, 2],
                [
                    m.make(m.OpConstant, 0),
                    m.make(m.OpSetGlobal, 0),
                    m.make(m.OpConstant, 1),
                    m.make(m.OpSetGlobal, 1),
                ],
            ),
            (
                "let one = 1;\none;\n",
                [1],
                [
                    m.make(m.OpConstant, 0),
                    m.make(m.OpSetGlobal, 0),
                    m.make(m.OpGetGlobal, 0),
                    m.make(m.OpPop),
                ],
            ),
            (
                "let one = 1;\nlet two = one;\ntwo;\n",
                [1],
                [
                    m.make(m.OpConstant, 0),
                    m.make(m.OpSetGlobal, 0),
                    m.make(m.OpGetGlobal, 0),
                    m.make(m.OpSetGlobal, 1),
                    m.make(m.OpGetGlobal, 1),
                    m.make(m.OpPop),
                ],
            ),
        ]
            run_compiler_tests(input, expected_constants, expected_instructions)
        end
    end

    @testset "Global Let Statements" begin
        for (input, expected_constants, expected_instructions) in [
            (
                "let num = 55;\nfn() { num }",
                [55, vcat(m.make(m.OpGetGlobal, 0), m.make(m.OpReturnValue))],
                [
                    m.make(m.OpConstant, 0),
                    m.make(m.OpSetGlobal, 0),
                    m.make(m.OpClosure, 1, 0),
                    m.make(m.OpPop),
                ],
            ),
            (
                "fn() {\nlet num = 55;\nnum\n}",
                [
                    55,
                    vcat(
                        m.make(m.OpConstant, 0),
                        m.make(m.OpSetLocal, 0),
                        m.make(m.OpGetLocal, 0),
                        m.make(m.OpReturnValue),
                    ),
                ],
                [m.make(m.OpClosure, 1, 0), m.make(m.OpPop)],
            ),
            (
                """
                fn() {
                  let a = 55;
                  let b = 77;
                  a + b
                }
                """,
                [
                    55,
                    77,
                    vcat(
                        m.make(m.OpConstant, 0),
                        m.make(m.OpSetLocal, 0),
                        m.make(m.OpConstant, 1),
                        m.make(m.OpSetLocal, 1),
                        m.make(m.OpGetLocal, 0),
                        m.make(m.OpGetLocal, 1),
                        m.make(m.OpAdd),
                        m.make(m.OpReturnValue),
                    ),
                ],
                [m.make(m.OpClosure, 2, 0), m.make(m.OpPop)],
            ),
        ]
            run_compiler_tests(input, expected_constants, expected_instructions)
        end
    end

    @testset "String Expressions" begin
        for (input, expected_constants, expected_instructions) in [
            ("\"monkey\"", ["monkey"], [m.make(m.OpConstant, 0), m.make(m.OpPop)]),
            (
                "\"mon\" + \"key\"",
                ["mon", "key"],
                [
                    m.make(m.OpConstant, 0),
                    m.make(m.OpConstant, 1),
                    m.make(m.OpAdd),
                    m.make(m.OpPop),
                ],
            ),
        ]
            run_compiler_tests(input, expected_constants, expected_instructions)
        end
    end

    @testset "Array Literals" begin
        for (input, expected_constants, expected_instructions) in [
            ("[]", [], [m.make(m.OpArray, 0), m.make(m.OpPop)]),
            (
                "[1, 2, 3]",
                [1, 2, 3],
                [
                    m.make(m.OpConstant, 0),
                    m.make(m.OpConstant, 1),
                    m.make(m.OpConstant, 2),
                    m.make(m.OpArray, 3),
                    m.make(m.OpPop),
                ],
            ),
            (
                "[1 + 2, 3 - 4, 5 * 6]",
                [1, 2, 3, 4, 5, 6],
                [
                    m.make(m.OpConstant, 0),
                    m.make(m.OpConstant, 1),
                    m.make(m.OpAdd),
                    m.make(m.OpConstant, 2),
                    m.make(m.OpConstant, 3),
                    m.make(m.OpSub),
                    m.make(m.OpConstant, 4),
                    m.make(m.OpConstant, 5),
                    m.make(m.OpMul),
                    m.make(m.OpArray, 3),
                    m.make(m.OpPop),
                ],
            ),
        ]
            run_compiler_tests(input, expected_constants, expected_instructions)
        end
    end

    @testset "Hash Literals" begin
        for (input, expected_constants, expected_instructions) in [
            ("{}", [], [m.make(m.OpHash, 0), m.make(m.OpPop)]),
            (
                "{1: 2, 3: 4, 5: 6}",
                [1, 2, 3, 4, 5, 6],
                [
                    m.make(m.OpConstant, 0),
                    m.make(m.OpConstant, 1),
                    m.make(m.OpConstant, 2),
                    m.make(m.OpConstant, 3),
                    m.make(m.OpConstant, 4),
                    m.make(m.OpConstant, 5),
                    m.make(m.OpHash, 6),
                    m.make(m.OpPop),
                ],
            ),
            (
                "{1: 2 + 3, 4: 5 * 6}",
                [1, 2, 3, 4, 5, 6],
                [
                    m.make(m.OpConstant, 0),
                    m.make(m.OpConstant, 1),
                    m.make(m.OpConstant, 2),
                    m.make(m.OpAdd),
                    m.make(m.OpConstant, 3),
                    m.make(m.OpConstant, 4),
                    m.make(m.OpConstant, 5),
                    m.make(m.OpMul),
                    m.make(m.OpHash, 4),
                    m.make(m.OpPop),
                ],
            ),
        ]
            run_compiler_tests(input, expected_constants, expected_instructions)
        end
    end

    @testset "Index Expressions" begin
        for (input, expected_constants, expected_instructions) in [
            (
                "[1, 2, 3][1 + 1]",
                [1, 2, 3, 1, 1],
                [
                    m.make(m.OpConstant, 0),
                    m.make(m.OpConstant, 1),
                    m.make(m.OpConstant, 2),
                    m.make(m.OpArray, 3),
                    m.make(m.OpConstant, 3),
                    m.make(m.OpConstant, 4),
                    m.make(m.OpAdd),
                    m.make(m.OpIndex),
                    m.make(m.OpPop),
                ],
            ),
            (
                "{1: 2}[2 - 1]",
                [1, 2, 2, 1],
                [
                    m.make(m.OpConstant, 0),
                    m.make(m.OpConstant, 1),
                    m.make(m.OpHash, 2),
                    m.make(m.OpConstant, 2),
                    m.make(m.OpConstant, 3),
                    m.make(m.OpSub),
                    m.make(m.OpIndex),
                    m.make(m.OpPop),
                ],
            ),
        ]
            run_compiler_tests(input, expected_constants, expected_instructions)
        end
    end

    @testset "Functions" begin
        for (input, expected_constants, expected_instructions) in [
            (
                "fn() { return 5 + 10 }",
                [
                    5,
                    10,
                    vcat(
                        m.make(m.OpConstant, 0),
                        m.make(m.OpConstant, 1),
                        m.make(m.OpAdd),
                        m.make(m.OpReturnValue),
                    ),
                ],
                [m.make(m.OpClosure, 2, 0), m.make(m.OpPop)],
            ),
            (
                "fn() { 5 + 10 }",
                [
                    5,
                    10,
                    vcat(
                        m.make(m.OpConstant, 0),
                        m.make(m.OpConstant, 1),
                        m.make(m.OpAdd),
                        m.make(m.OpReturnValue),
                    ),
                ],
                [m.make(m.OpClosure, 2, 0), m.make(m.OpPop)],
            ),
            (
                "fn() { 1; 2 }",
                [
                    1,
                    2,
                    vcat(
                        m.make(m.OpConstant, 0),
                        m.make(m.OpPop),
                        m.make(m.OpConstant, 1),
                        m.make(m.OpReturnValue),
                    ),
                ],
                [m.make(m.OpClosure, 2, 0), m.make(m.OpPop)],
            ),
            (
                "fn() { }",
                [m.make(m.OpReturn)],
                [m.make(m.OpClosure, 0, 0), m.make(m.OpPop)],
            ),
        ]
            run_compiler_tests(input, expected_constants, expected_instructions)
        end
    end

    @testset "Compiler Scopes" begin
        c = m.Compiler()
        g = c.symbol_table
        @test length(c.scopes) == 1

        m.emit!(c, m.OpMul)
        m.enter_scope!(c)
        @test length(c.scopes) == 2

        m.emit!(c, m.OpSub)
        @test length(c) == 1

        last = m.last_instruction(c)
        @test last.op == m.OpSub
        @test c.symbol_table.outer == g

        m.leave_scope!(c)
        @test length(c.scopes) == 1
        @test c.symbol_table == g
        @test isnothing(c.symbol_table.outer)

        m.emit!(c, m.OpAdd)
        @test length(c) == 2

        last = m.last_instruction(c)
        @test last.op == m.OpAdd

        previous = m.prev_instruction(c)
        @test previous.op == m.OpMul
    end

    @testset "Function Calls" begin
        for (input, expected_constants, expected_instructions) in [
            (
                "fn() { 24 }()",
                [24, vcat(m.make(m.OpConstant, 0), m.make(m.OpReturnValue))],
                [m.make(m.OpClosure, 1, 0), m.make(m.OpCall, 0), m.make(m.OpPop)],
            ),
            (
                "let noArg = fn() { 24 }; noArg()",
                [24, vcat(m.make(m.OpConstant, 0), m.make(m.OpReturnValue))],
                [
                    m.make(m.OpClosure, 1, 0),
                    m.make(m.OpSetGlobal, 0),
                    m.make(m.OpGetGlobal, 0),
                    m.make(m.OpCall, 0),
                    m.make(m.OpPop),
                ],
            ),
            (
                "let oneArg = fn(a) { }; oneArg(24)",
                [m.make(m.OpReturn), 24],
                [
                    m.make(m.OpClosure, 0, 0),
                    m.make(m.OpSetGlobal, 0),
                    m.make(m.OpGetGlobal, 0),
                    m.make(m.OpConstant, 1),
                    m.make(m.OpCall, 1),
                    m.make(m.OpPop),
                ],
            ),
            (
                "let manyArg = fn(a, b, c) { }; manyArg(24, 25, 26)",
                [m.make(m.OpReturn), 24, 25, 26],
                [
                    m.make(m.OpClosure, 0, 0),
                    m.make(m.OpSetGlobal, 0),
                    m.make(m.OpGetGlobal, 0),
                    m.make(m.OpConstant, 1),
                    m.make(m.OpConstant, 2),
                    m.make(m.OpConstant, 3),
                    m.make(m.OpCall, 3),
                    m.make(m.OpPop),
                ],
            ),
            (
                "let oneArg = fn(a) { a }; oneArg(24)",
                [vcat(m.make(m.OpGetLocal, 0), m.make(m.OpReturnValue)), 24],
                [
                    m.make(m.OpClosure, 0, 0),
                    m.make(m.OpSetGlobal, 0),
                    m.make(m.OpGetGlobal, 0),
                    m.make(m.OpConstant, 1),
                    m.make(m.OpCall, 1),
                    m.make(m.OpPop),
                ],
            ),
            (
                "let manyArg = fn(a, b, c) { a; b; c; }; manyArg(24, 25, 26)",
                [
                    vcat(
                        m.make(m.OpGetLocal, 0),
                        m.make(m.OpPop),
                        m.make(m.OpGetLocal, 1),
                        m.make(m.OpPop),
                        m.make(m.OpGetLocal, 2),
                        m.make(m.OpReturnValue),
                    ),
                    24,
                    25,
                    26,
                ],
                [
                    m.make(m.OpClosure, 0, 0),
                    m.make(m.OpSetGlobal, 0),
                    m.make(m.OpGetGlobal, 0),
                    m.make(m.OpConstant, 1),
                    m.make(m.OpConstant, 2),
                    m.make(m.OpConstant, 3),
                    m.make(m.OpCall, 3),
                    m.make(m.OpPop),
                ],
            ),
        ]
            run_compiler_tests(input, expected_constants, expected_instructions)
        end
    end

    @testset "Builtin Functions" begin
        for (input, expected_constants, expected_instructions) in [
            (
                "len([]); push([], 1)",
                [1],
                [
                    m.make(m.OpGetBuiltin, 0),
                    m.make(m.OpArray, 0),
                    m.make(m.OpCall, 1),
                    m.make(m.OpPop),
                    m.make(m.OpGetBuiltin, 4),
                    m.make(m.OpArray, 0),
                    m.make(m.OpConstant, 0),
                    m.make(m.OpCall, 2),
                    m.make(m.OpPop),
                ],
            ),
            (
                "fn() { len([]) }",
                [
                    vcat(
                        m.make(m.OpGetBuiltin, 0),
                        m.make(m.OpArray, 0),
                        m.make(m.OpCall, 1),
                        m.make(m.OpReturnValue),
                    ),
                ],
                [m.make(m.OpClosure, 0, 0), m.make(m.OpPop)],
            ),
        ]
            run_compiler_tests(input, expected_constants, expected_instructions)
        end
    end

    @testset "Closures" begin
        for (input, expected_constants, expected_instructions) in [
            (
                """
                fn(a) {
                    fn(b) {
                        a + b
                    }
                }
                """,
                [
                    vcat(
                        m.make(m.OpGetFree, 0),
                        m.make(m.OpGetLocal, 0),
                        m.make(m.OpAdd),
                        m.make(m.OpReturnValue),
                    ),
                    vcat(
                        m.make(m.OpGetLocal, 0),
                        m.make(m.OpClosure, 0, 1),
                        m.make(m.OpReturnValue),
                    ),
                ],
                [m.make(m.OpClosure, 1, 0), m.make(m.OpPop)],
            ),
            (
                """
                fn(a) {
                    fn(b) {
                        fn(c) {
                            a + b + c
                        }
                    }
                }
                """,
                [
                    vcat(
                        m.make(m.OpGetFree, 0),
                        m.make(m.OpGetFree, 1),
                        m.make(m.OpAdd),
                        m.make(m.OpGetLocal, 0),
                        m.make(m.OpAdd),
                        m.make(m.OpReturnValue),
                    ),
                    vcat(
                        m.make(m.OpGetFree, 0),
                        m.make(m.OpGetLocal, 0),
                        m.make(m.OpClosure, 0, 2),
                        m.make(m.OpReturnValue),
                    ),
                    vcat(
                        m.make(m.OpGetLocal, 0),
                        m.make(m.OpClosure, 1, 1),
                        m.make(m.OpReturnValue),
                    ),
                ],
                [m.make(m.OpClosure, 2, 0), m.make(m.OpPop)],
            ),
            (
                """
                let global = 55;
                fn() {
                    let a = 66;
                    fn() {
                        a = a + 2;
                        let b = 77;
                        fn() {
                            global = global + 4;
                            let c = 88;
                            global + a + b + c
                        }
                    }
                }
                """,
                [
                    55,
                    66,
                    2,
                    77,
                    4,
                    88,
                    vcat(
                        m.make(m.OpGetGlobal, 0),
                        m.make(m.OpConstant, 4),
                        m.make(m.OpAdd),
                        m.make(m.OpSetGlobal, 0),
                        m.make(m.OpConstant, 5),
                        m.make(m.OpSetLocal, 0),
                        m.make(m.OpGetGlobal, 0),
                        m.make(m.OpGetFree, 0),
                        m.make(m.OpAdd),
                        m.make(m.OpGetFree, 1),
                        m.make(m.OpAdd),
                        m.make(m.OpGetLocal, 0),
                        m.make(m.OpAdd),
                        m.make(m.OpReturnValue),
                    ),
                    vcat(
                        m.make(m.OpGetFree, 0),
                        m.make(m.OpConstant, 2),
                        m.make(m.OpAdd),
                        m.make(m.OpSetFree, 0),
                        m.make(m.OpConstant, 3),
                        m.make(m.OpSetLocal, 0),
                        m.make(m.OpGetFree, 0),
                        m.make(m.OpGetLocal, 0),
                        m.make(m.OpClosure, 6, 2),
                        m.make(m.OpReturnValue),
                    ),
                    vcat(
                        m.make(m.OpConstant, 1),
                        m.make(m.OpSetLocal, 0),
                        m.make(m.OpGetLocal, 0),
                        m.make(m.OpClosure, 7, 1),
                        m.make(m.OpReturnValue),
                    ),
                ],
                [
                    m.make(m.OpConstant, 0),
                    m.make(m.OpSetGlobal, 0),
                    m.make(m.OpClosure, 8, 0),
                    m.make(m.OpPop),
                ],
            ),
            (
                """
                let i = 2;
                while (i > 0) {
                    i = i - 1;
                    let j = 3;
                    let f = fn() {
                        i + j;
                    }
                    while (j > 0) {
                        j = j - 1;
                        puts(f() + j);
                    }
                }
                """,
                [
                    2,
                    0,
                    1,
                    3,
                    vcat(
                        m.make(m.OpGetGlobal, 0),
                        m.make(m.OpGetFree, 0),
                        m.make(m.OpAdd),
                        m.make(m.OpReturnValue),
                    ),
                    0,
                    1,
                    vcat(
                        m.make(m.OpGetOuter, 1, 1, 0),
                        m.make(m.OpConstant, 6),
                        m.make(m.OpSub),
                        m.make(m.OpSetOuter, 1, 1, 0),
                        m.make(m.OpGetBuiltin, 5),
                        m.make(m.OpGetOuter, 1, 1, 1),
                        m.make(m.OpCall, 0),
                        m.make(m.OpGetOuter, 1, 1, 0),
                        m.make(m.OpAdd),
                        m.make(m.OpCall, 1),
                        m.make(m.OpPop),
                        m.make(m.OpContinue),
                    ),
                    vcat(
                        m.make(m.OpGetGlobal, 0),
                        m.make(m.OpConstant, 2),
                        m.make(m.OpSub),
                        m.make(m.OpSetGlobal, 0),
                        m.make(m.OpConstant, 3),
                        m.make(m.OpSetLocal, 0),
                        m.make(m.OpGetLocal, 0),
                        m.make(m.OpClosure, 4, 1),
                        m.make(m.OpSetLocal, 1),
                        m.make(m.OpGetLocal, 0),
                        m.make(m.OpConstant, 5),
                        m.make(m.OpGreaterThan),
                        m.make(m.OpJumpNotTruthy, 44),
                        m.make(m.OpClosure, 7, 0),
                        m.make(m.OpCall, 0),
                        m.make(m.OpJumpNotTruthy, 44),
                        m.make(m.OpJump, 23),
                        m.make(m.OpNull),
                        m.make(m.OpPop),
                        m.make(m.OpContinue),
                    ),
                ],
                [
                    m.make(m.OpConstant, 0),
                    m.make(m.OpSetGlobal, 0),
                    m.make(m.OpGetGlobal, 0),
                    m.make(m.OpConstant, 1),
                    m.make(m.OpGreaterThan),
                    m.make(m.OpJumpNotTruthy, 28),
                    m.make(m.OpClosure, 8, 0),
                    m.make(m.OpCall, 0),
                    m.make(m.OpJumpNotTruthy, 28),
                    m.make(m.OpJump, 6),
                    m.make(m.OpNull),
                    m.make(m.OpPop),
                ],
            ),
            (
                """
                while (true) {
                    break;
                    continue;
                }
                """,
                [vcat(m.make(m.OpBreak), m.make(m.OpContinue), m.make(m.OpContinue))],
                [
                    m.make(m.OpTrue),
                    m.make(m.OpJumpNotTruthy, 16),
                    m.make(m.OpClosure, 0, 0),
                    m.make(m.OpCall, 0),
                    m.make(m.OpJumpNotTruthy, 16),
                    m.make(m.OpJump, 0),
                    m.make(m.OpNull),
                    m.make(m.OpPop),
                ],
            ),
        ]
            run_compiler_tests(input, expected_constants, expected_instructions)
        end
    end

    @testset "Recursive Functions" begin
        for (input, expected_constants, expected_instructions) in [
            (
                """
                let countDown = fn(x) { countDown(x - 1); };
                countDown(1);
                """,
                [
                    1,
                    vcat(
                        m.make(m.OpCurrentClosure),
                        m.make(m.OpGetLocal, 0),
                        m.make(m.OpConstant, 0),
                        m.make(m.OpSub),
                        m.make(m.OpCall, 1),
                        m.make(m.OpReturnValue),
                    ),
                    1,
                ],
                [
                    m.make(m.OpClosure, 1, 0),
                    m.make(m.OpSetGlobal, 0),
                    m.make(m.OpGetGlobal, 0),
                    m.make(m.OpConstant, 2),
                    m.make(m.OpCall, 1),
                    m.make(m.OpPop),
                ],
            ),
            (
                """
                let wrapper = fn() {
                    let countDown = fn(x) { countDown(x - 1); };
                    countDown(1);
                }
                wrapper();
                """,
                [
                    1,
                    vcat(
                        m.make(m.OpCurrentClosure),
                        m.make(m.OpGetLocal, 0),
                        m.make(m.OpConstant, 0),
                        m.make(m.OpSub),
                        m.make(m.OpCall, 1),
                        m.make(m.OpReturnValue),
                    ),
                    1,
                    vcat(
                        m.make(m.OpClosure, 1, 0),
                        m.make(m.OpSetLocal, 0),
                        m.make(m.OpGetLocal, 0),
                        m.make(m.OpConstant, 2),
                        m.make(m.OpCall, 1),
                        m.make(m.OpReturnValue),
                    ),
                ],
                [
                    m.make(m.OpClosure, 3, 0),
                    m.make(m.OpSetGlobal, 0),
                    m.make(m.OpGetGlobal, 0),
                    m.make(m.OpCall, 0),
                    m.make(m.OpPop),
                ],
            ),
        ]
            run_compiler_tests(input, expected_constants, expected_instructions)
        end
    end

    @testset "Error handling" begin
        @test_throws ErrorException("unknown operator: &") begin
            c = m.Compiler()
            m.compile!(
                c,
                m.InfixExpression(
                    m.Token(m.PLUS, "+"),
                    m.IntegerLiteral(m.Token(m.INT, "2"), 2),
                    "&",
                    m.IntegerLiteral(m.Token(m.INT, "2"), 2),
                ),
            )
        end

        @test_throws ErrorException("unknown operator: +") begin
            c = m.Compiler()
            m.compile!(
                c,
                m.PrefixExpression(
                    m.Token(m.PLUS, "+"),
                    "+",
                    m.IntegerLiteral(m.Token(m.INT, "2"), 2),
                ),
            )
        end
    end
end


#
println("Ai) Test of the jj-LS transformation of levels from a given multiplet.")
#
wa = Atomic.Computation("jj-LS level transformation",  Nuclear.Model(26.);
                        configs=[Configuration("[Ne] 3s^2 3p^4 3d")],
                        asfSettings=AsfSettings(true, false, "meanDFS", "hydrogenic", "xxx", [1],     0, 1.0e-6, JAC.Subshell[],
                                                true, false, NoneQed(), "yyy", ManyElectron.LSjjSettings(true),
                                                false, [1,2,3,4], false, JAC.LevelSymmetry[] ) )

wb = perform(wa)

##                        configs=[Configuration("[Ne] 3s^2 3p^4"), Configuration("[Ne] 3s^2 3d^4")],  , Configuration("1s^2 2s 2p^6")


using DiffOpt
using Test
using OSQP
using Ipopt
using Clp
using Random
using SCS
using LinearAlgebra
using DelimitedFiles
using GLPK

using SparseArrays: spzeros

import MathOptInterface
const MOI = MathOptInterface
const MOIU = MathOptInterface.Utilities
const MOIT = MathOptInterface.Test

import MatrixOptInterface
const MatOI = MatrixOptInterface

const ATOL = 1e-4
const RTOL = 1e-4

@testset "Examples" begin
    include(joinpath(@__DIR__, "../examples/solve-LP.jl"))
    include(joinpath(@__DIR__, "../examples/solve-QP.jl"))
end

@testset "Generate random problems" begin
    include("gen_random_problem.jl")
end

@testset "MOI_wrapper" begin
    include("MOI_wrapper.jl")
end

@testset "JuMP wrapper" begin
    include("jump.jl")
end

@testset "Solver Interface" begin
    include("solver_interface.jl")
end

@testset "Singular error with deleted variables" begin
    include("singular_exception.jl")
end


## All this code breaks with JumP 0.19.

function solve(model::PowerOperationModel{T, F}; Solver = nothing) where {T <: AbstractOperationsModel, F <: Array{<:Function}}

    psmodel = model.model # JuMP model
    model_type = JuMP.ProblemTraits(psmodel)

    #==
    function ProblemTraits(m::Model; relaxation=false)
    int = !relaxation & any(c-> !(c == :Cont || c == :Fixed), m.colCat)
    qp = !isempty(m.obj.qvars1)
    qc = !isempty(m.quadconstr)
    nlp = m.nlpdata !== nothing
    soc = !isempty(m.socconstr)
    # will need to change this when we add support for arbitrary variable cones
    sdp = !isempty(m.sdpconstr) || !isempty(m.varCones)
    sos = !isempty(m.sosconstr)
    ProblemTraits(int, !(qp|qc|nlp|soc|sdp|sos), qp, qc, nlp, soc, sdp, sos, soc|sdp)
    end
    =#

    if Solver == nothing

        if (model_type.qc|model_type.qp|model_type.nlp)

            if model_type.int
                @error "The model is a Mixed Integer Non-Linear Problem, please define an appropiate solver manually using the Solver= argument"
            end

            JuMP.setsolver(psmodel, Ipopt.IpoptSolver())
            @warn "The model contains non-linear elements (QP, QC, NLP), by default the solver is Ipopt Solver"

        elseif model_type.lin & model_type.int & !(model_type.qc|model_type.qp|model_type.nlp)

            JuMP.setsolver(psmodel, Cbc.CbcSolver(logLevel = 1))
            @warn "The model is Mixed Integer Linear, the default is Cbc Solver"

        elseif model_type.lin & !(model_type.qc|model_type.qp|model_type.nlp)

            JuMP.setsolver(psmodel, Clp.ClpSolver(SolveType = 5, LogLevel = 4))

            @warn "The model is linear, by default the solver is Clp Solver"

        else

            @error "The PS Model is not a standard LP, MILP, QP, QC or NLP, please define an appropiate solver manually using the Solver argument"

        end

    else
        JuMP.setsolver(psmodel, Solver)
        @warn "A solver has been defined manually, this might break the results output function"

    end

    JuMP.build(psmodel)
    status = solve(psmodel)

    if status != :Optimal
        println(status)
        @error "Problem has no solution."
    end

end


"""
    function get_Pg_thermal_constprice(args..., kwargs...)

Identify generators which cost function is a constant price, locate their buses and returns
a matrix with their Name, Power output and Bus number (sorted by Bus). Skips Slack
generators, Generators which Active Power is binding and Generators with non zero slope.
Note: For a quadratic cost function of the form: αPg²+ βPg + γ, a constant price generator
will have the parameter α == 0.

# Arguments
- `sys::System`:                            Power system in p.u. (from PowerSystems.jl)
- `res::PowerSimulations.
    OperationsProblemResults`:              Results of the solved OPF for the system
                                             (from PowerSimulations.jl)

# Keywords
- `dual_gen_tol::Float64 = 1e-1`:           Tolerance to identify any binding generators
"""
function get_Pg_thermal_constprice(
    sys::System,
    res::PowerSimulations.OperationsProblemResults;
    dual_gen_tol::Float64 = 1e-1
    )
    # Define elements
    gen_loc = 0
    gen_constprice = 0 #number of generators with constant price (counter starter)
    gens_thermal_constprice_Pg = Array{Float64}(undef,0)
    gens_thermal_constprice_names = Array{String}(undef,0)
    gens_thermal_constprice_busnumbers = Array{Int64}(undef,0)

    # Get all thermal Generators
    gens_thermal = get_components(ThermalStandard, sys)
    all_PGenThermal = get_variables(res)[:P__ThermalStandard] #Optimised PGen

    # Get dual variables from the OPF results
    dualPgub = get_duals(res)[:P_ub__ThermalStandard__RangeConstraint]
    dualPglb = get_duals(res)[:P_lb__ThermalStandard__RangeConstraint]

    # Identify Generators and save their information
    for gen_thermal in gens_thermal
        gen_loc = gen_loc +1
        if get_bustype(gen_thermal.bus) ≠ BusTypes.REF  #skips slack
            bindub = !isapprox(dualPgub[1,gen_loc], 0.0; atol = dual_gen_tol) #is binding?
            bindlb = !isapprox(dualPglb[1,gen_loc], 0.0; atol = dual_gen_tol) #is binding?
            if !bindub && !bindlb #skips binding Pg gens Dual = 0
                α, β = get_cost(get_variable(get_operation_cost(gen_thermal)))
                γ = get_fixed(get_operation_cost(gen_thermal))
                if α == 0 && (β ≠ 0 || γ ≠ 0)#skips non zero slope
                    gen_constprice = gen_constprice + 1
                    gen_thermal_constprice_name = gen_thermal.name
                    resize!(gens_thermal_constprice_Pg,gen_constprice)
                    resize!(gens_thermal_constprice_names,gen_constprice)
                    resize!(gens_thermal_constprice_busnumbers,gen_constprice)
                    gens_thermal_constprice_Pg[gen_constprice] =
                        all_PGenThermal[1, gen_thermal_constprice_name]
                    gens_thermal_constprice_names[gen_constprice] =
                        gen_thermal_constprice_name
                    gens_thermal_constprice_busnumbers[gen_constprice] =
                        gen_thermal.bus.number
                end
            end
        end
    end

    # Matrix of generators with their information sorted
    gens_thermal_constprice_Matrix = Array{Any}(undef,(gen_constprice, 3))
    gens_thermal_constprice_Matrix[:, 1] = gens_thermal_constprice_names
    gens_thermal_constprice_Matrix[:, 2] = gens_thermal_constprice_Pg
    gens_thermal_constprice_Matrix[:, 3] = gens_thermal_constprice_busnumbers
    # Sort by bus to match with PTDF structure
    gens_thermal_constprice_Matrix =
        gens_thermal_constprice_Matrix[sortperm(gens_thermal_constprice_Matrix[:, 3]), :]

    return gens_thermal_constprice_Matrix
end


##############################################################################
#
# RegularisationFunction.jl
#
# Part of cvortex.jl
# Get vortex regularisation methods.
#
# Copyright 2019 HJA Bird
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to
# deal in the Software without restriction, including without limitation the
# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
# sell copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
# IN THE SOFTWARE.
##############################################################################

"""
A 3D vortex regularisation function.

The method by which the singular nature of a vortex particle is managed.
Generally, this structure is best obtained via *_regularisation(). For example
singular_regularisation(), planetary_regularisation(),
gaussian_regularisation() and winckelmans_regulariation().

Exposure of this allows the use of user kernels in the CPU multithreaded
implementations in cvortex. If you're interested in this, looking at the
cvortex library and the source of cvortex.jl is suggested to understand
the required functions signitures.

# Examples
```
using CVortex
my_regularisation_kernel = singular_regularisation()
my_regularisation_kernel = winckelmans_regularisation()
my_regularisation_kernel = planetary_regularisation()
my_regularisation_kernel = gaussian_regularisation()
```
"""
struct RegularisationFunction
	g_3D :: Ptr{Cvoid}			# Actually float(*g_3D)(float rho)
	g_2D :: Ptr{Cvoid}			# Actually float(*g_2D)(float rho)
	zeta_3D :: Ptr{Cvoid}		# Actually float(*zeta_3D)(float rho)
	combined_3D :: Ptr{Cvoid}	# Actually void(*combined_3D)(float rho, float* g, float* zeta)
	eta_3D :: Ptr{Cvoid}		# Actually float(*eta_3D)(float rho)
	eta_2D :: Ptr{Cvoid}		# Actually float(*eta_2D)(float rho)
	cl_kernel_name_ext :: NTuple{32, Cchar}	# Char[32]
end

#= Functions to to get VortexFunc structures =#
"""
Implements singular vortex particles.
returns RegularisationFunction struct.
"""
function singular_regularisation()
	ret = ccall(("cvtx_VortFunc_singular", libcvortex),
		RegularisationFunction, ())
	return ret;
end
"""
Implements Winckelmans' high order algebraic regularisation.
returns RegularisationFunction struct.
"""
function winckelmans_regularisation()
	ret = ccall(("cvtx_VortFunc_winckelmans", libcvortex),
		RegularisationFunction, ())
	return ret;
end
"""
Implements Planetary regularisation.
returns RegularisationFunction struct.
"""
function planetary_regularisation()
	ret = ccall(("cvtx_VortFunc_planetary", libcvortex),
		RegularisationFunction, ())
	return ret;
end
"""
Implements Gaussian regularisation.
returns RegularisationFunction struct.
"""
function gaussian_regularisation()
	ret = ccall(("cvtx_VortFunc_gaussian", libcvortex),
		RegularisationFunction, ())
	return ret;
end