text:
abstract type AbstractResistor end

"""
    r = Resistor(R0=1e3, α=0.0, T₀=293.15)

A resistor with resistance varying linearly with temperature:

```math
R = R₀ (1 + α (T - T₀))
```

**IMPORTANT** In the constructor, coefficient `α` is a *fraction*, not percentage.
 linear.

# Examples

```jldoctest
julia> R = Resistor(1e3, 0.01)
THW.Resistor(1000.0, 0.0001, 293.15)

julia> resistance(R)
1000.0

julia> resistance(R, 120+273.15)
1010.0

julia> R()
1000.0

julia> R(120+273.15)
1010.0

julia> temperature(R)
293.15

julia> temperature(R, 1010)
120.00000000000009

```
"""
struct Resistor <: AbstractResistor
    "Resistance at reference temperature (Ω)"
    R₀::Float64
    "Linear resistance coefficient"
    α::Float64
    "Reference temperature (K)"
    T₀::Float64
    Resistor(R₀=1e3, α=0.0, T₀=293.15) = new(R₀, α, T₀)
end
Base.broadcastable(R::Resistor) = Ref(R)

"""
    R = Thermistor(R₀, B, T₀)

Models a NTC thermistor. In the constructor, `R0` is the resistance of the thermistor at the reference temperature  `T₀`. For convenience, the temperature should be expressed in K. `B` is a coefficient that characterizes the temperature dependence of the thermistor. In the model used her, the resistance varies with temperature according to the following equation:

```math
R = R₀ \\exp B \\left(\\frac{1}{T₀} - \\frac{1}{T}\\right)
```

where
 * R₀ is the reference resistance in Ω
 * T₀ is the reference temperature in K, usually the adopted value is   293.15 K
 * T is the temperature in K, where the value of resistance should be calculated
 * B is an empirical coefficient in K that characterizes the temperature dependency of the resistance

    When using the struct `Thermistor`, beware: in the constructor, the unit of T₀ is °C NOT K.

# Examples
```jldoctest
julia> R = Thermistor(5e3, 3200, 293.15) # Create an object `Thermistor`
Thermistor(5000.0, 3200.0, 293.15)

julia> resistance(R) # Resistance at the reference temperature
5000.0

julia> resistance(R, 298.15) # Resistance at 25°C
4163.587774917559

julia> R() # The same as resistance(R)
5000.0

julia> R(298.15) # The same as resistance(R, 298.15)
4163.587774917559

julia> temperature(R)
293.15

julia> temperature(R, 4163.588)
298.15

```

"""
struct Thermistor <: AbstractResistor
    "Reference resistance in Ω at temperature `T₀`"
    R₀::Float64
    "Thermistor's B coefficient in K "
    B::Float64
    "Reference temperature in K"
    T₀::Float64
    Thermistor(R₀=5e3, B=0.0, T₀=25.0) = new(R₀, B, T₀)
end
Base.broadcastable(R::Thermistor) = Ref(R)

"""
    resistance(th::AbstractResistor)
    resistance(th::AbstractResistor, Tc)

Calculates the resistance for `AbstractResistor`objects` at temperature `Tc`.
If argumento `Tc` is not provided, the function returns the reference resistance `R0`.

See [`Thermistor`](@ref) to see more details of the  type  `Thermistor`. For more details on
type `Resistor`, see [`Thermistor`](@ref).

"""
resistance(r::AbstractResistor) = r.R₀
resistance(r::Resistor, T) = r.R₀ * (1.0 + r.α * (T - r.T₀))
resistance(th::Thermistor, T) = th.R₀ * exp( -th.B * (1/th.T₀ - 1/T ) )

temperature(r::AbstractResistor) = r.T₀
temperature(r::Resistor, R) = 1/r.α * (R/r.R₀ - 1.0) + r.T₀
temperature(th::Thermistor, R) = 1/( 1/th.T₀ + 1/th.B * log(R/th.R₀) )
temperature(th::Thermistor) = th.T₀

(th::Thermistor)(T) = resistance(th, T)
(th::Thermistor)() = th.R₀
(r::Resistor)(T) = resistance(r, T)
(r::Resistor)() = resistance(r)


text:
"""
$(TYPEDEF)

Dynamic point in 2D space with velocity components: `x, y, dx/dt, dy/dt`

"""
mutable struct DynPoint2 <: IncrementalInference.InferenceVariable
  ut::Int64 # microsecond time
  dims::Int
  manifolds::Tuple{Symbol, Symbol, Symbol, Symbol}
  DynPoint2(;ut::Int64=-9999999999) = new(ut, 4,(:Euclid,:Euclid,:Euclid,:Euclid,))
end
getDimension(::DynPoint2) = 4
# getManifolds(::DynPoint2) = (:Euclid,:Euclid,:Euclid,:Euclid)

ext: jl


text:
const ALCONSTRAINT_PARAMS = Set((
    :use_conic_cost,
    :penalty_initial,
    :penalty_scaling,
    :penalty_max,
    :dual_max
))

Base.@kwdef mutable struct ConstraintOptions{T}
    solveropts::SolverOptions{T} = SolverOptions{T}()
    use_conic_cost::Bool = false
    penalty_initial::T = 1.0
    penalty_scaling::T = 10.0
    penalty_max::T = 1e8
    dual_max::T = 1e8
    usedefault::Dict{Symbol,Bool} = Dict(Pair.(
        ALCONSTRAINT_PARAMS,
        trues(length(ALCONSTRAINT_PARAMS))
    ))
end

function setparams!(conopts::ConstraintOptions; kwargs...)
    for (key,val) in pairs(kwargs)
        if key in ALCONSTRAINT_PARAMS
            setfield!(conopts, key, val)
            conopts.usedefault[key] = false
        end
    end
end

@generated function reset!(conopts::ConstraintOptions)
    # Use generated expression to avoid allocations since the the fields have different types
    optchecks = map(collect(ALCONSTRAINT_PARAMS)) do rawparam
        param = Expr(:quote,rawparam)
        :(conopts.usedefault[$param] && setfield!(conopts, $param, getfield(opts, $param)))
    end
    quote
        opts = conopts.solveropts
        $(Expr(:block, optchecks...))
        return nothing
    end
end

"""
    ALConstraint

A constraint on the optimization problem, which only applies to a single
knot point. Stores all of the data needed to evaluate the constraint,
as well as the Augmented Lagrangian penalty cost and its derivatives.
The default Augmented Lagrangian cost has the following form:

```math
\\frac{1}{2} \\left( || \\Pi_K(\\lambda - I_\\mu c(x) ) ||^2 - || \\lambda ||^2 \\right)
```

for a constraint in cone ``K``. To use the traditional form for equality and
inequality constraints, set `opts.use_conic_cost = false` (default).

This type also provides methods for applying the dual and penalty updates.

## Constructor
A new `ALConstraint` is created using

    ALConstraint{T}(n, m, con, inds; [sig, diffmethod, kwargs...])

with arguments:
- `n` the size of the state vector
- `m` the size of the control vector
- `con` a TrajectoryOptimization.StageConstraint
- `inds` the knot point indices at which the constraint is applied
- `kwargs` keyword arguments passed to the constructor of the [`ConstraintOptions`](@ref).

## Methods
The following methods can be used with an `ALConstraint` object:
- [`evaluate_constraint!(alcon, Z)`](@ref)
- [`constraint_jacobian!(alcon, Z)`](@ref)
- [`alcost(alcon, Z)`](@ref)
- [`algrad!(alcon, Z)`](@ref)
- [`alhess!(alcon, Z)`](@ref)
- [`dualupdate!(alcon)`](@ref)
- [`penaltyupdate!(alcon)`](@ref)
- [`normviolation!(alcon)`](@ref)
- [`max_violation(alcon)`](@ref)
- [`max_penalty(alcon)`](@ref)
- [`reset_duals!(alcon)`](@ref)
- [`reset_penalties!(alcon)`](@ref)
"""

struct ALConstraint{T, C<:TO.StageConstraint, R<:SampledTrajectory}
    nx::Vector{Int}  # state dimension     # TODO: remove these?
    nu::Vector{Int}  # control dimension
    con::C
    sig::FunctionSignature
    diffmethod::DiffMethod
    inds::Vector{Int}              # knot point indices for constraint
    vals::Vector{Vector{T}}        # constraint values
    jac::Vector{Matrix{T}}         # constraint Jacobian
    jac_scaled::Vector{Matrix{T}}  # penalty-scaled constraint Jacobian
    λ::Vector{Vector{T}}           # dual variables
    μ::Vector{Vector{T}}           # penalties
    μinv::Vector{Vector{T}}        # inverted penalties
    λbar::Vector{Vector{T}}        # approximate dual variable
    λproj::Vector{Vector{T}}       # projected dual variables
    λscaled::Vector{Vector{T}}     # scaled projected dual variables
    viol::Vector{Vector{T}}        # constraint violations
    c_max::Vector{T}

    ∇proj::Vector{Matrix{T}}   # Jacobian of projection
    ∇proj_scaled::Vector{Matrix{T}}
    ∇²proj::Vector{Matrix{T}}  # Second-order derivative of projection
    cost::Vector{T}            # (N,) vector of costs (aliased to the one in ALObjective)
    grad::Vector{Vector{T}}    # gradient of Augmented Lagrangian
    hess::Vector{Matrix{T}}    # Hessian of Augmented Lagrangian
    tmp_jac::Matrix{T}

    # Hack to avoid allocations
    # Store references to the trajectory and cost expansion
    # The trajectory is stored in a vector so it can be mutated
    # NOTE: tried to use Refs but it ended up allocating
    Z::Vector{R}
    E::CostExpansion{T}

    opts::ConstraintOptions{T}
    function ALConstraint{T}(Z::R, con::TO.StageConstraint,
                             inds::AbstractVector{<:Integer},
                             costs::Vector{T},
                             E=CostExpansion{T}(RD.dims(Z)[1:2]...);
			                 sig::FunctionSignature=StaticReturn(),
                             diffmethod::DiffMethod=UserDefined(),
                             kwargs...
    ) where {T,R<:SampledTrajectory}
        opts = ConstraintOptions{T}(;kwargs...)

        nx,nu = RD.dims(Z)
        p = RD.output_dim(con)
        P = length(inds)
        w = RD.input_dim(con)
        nz = nx + nu

        vals = [zeros(T, p) for k in inds]
        jac = [zeros(T, p, w) for k in inds]
        jac_scaled = [zeros(T, p, w) for k in inds]
        λ = [zeros(T, p) for k in inds]
        μ = [fill(opts.penalty_initial, p) for k in inds]
        μinv = [inv.(μi) for μi in μ]
        λbar = [zeros(T, p) for k in inds]
        λproj = [zeros(T, p) for k in inds]
        λscaled = [zeros(T, p) for k in inds]
        viol = [zeros(T, p) for k in inds]
        c_max = zeros(T, length(Z))

        ∇proj = [zeros(T, p, p) for k in inds]
        ∇proj_scaled = [zeros(T, p, p) for k in inds]
        ∇²proj = [zeros(T, p, p) for k in inds]

        grad = [zeros(T, nz[k]) for k in inds]
        hess = [zeros(T, nz[k], nz[k]) for k in inds]

        tmp_jac = zeros(T, p, w)

        new{T, typeof(con), R}(
            nx[inds], nu[inds], con, sig, diffmethod, inds, vals, jac, jac_scaled, λ, μ, μinv, λbar,
            λproj, λscaled, viol, c_max, ∇proj, ∇proj_scaled, ∇²proj, costs, grad, hess, tmp_jac,
            [Z], E, opts
        )
    end
end

settraj!(alcon::ALConstraint, Z::SampledTrajectory) = alcon.Z[1] = Z
setparams!(alcon::ALConstraint; kwargs...) = setparams!(alcon.opts; kwargs...)
resetparams!(alcon::ALConstraint) = reset!(alcon.opts)
RD.vectype(alcon::ALConstraint) = RD.vectype(eltype(alcon.Z[1]))

"""
    getinputinds(alcon)

Get the indices of the input to the constraint, determined by the output of
`RobotDynamics.functioninputs`. Returns `1:n+m` for generic stage constraints,
`1:n` for state constraints, and `n+1:n+m` for control constraints.
"""
function getinputinds(alcon::ALConstraint, i::Integer)
    n,m = alcon.nx[i], alcon.nu[i]
    inputtype = RD.functioninputs(alcon.con)
    if inputtype == RD.StateOnly()
        inds = 1:n
    elseif inputtype == RD.ControlOnly()
        inds = n .+ (1:m)
    else
        inds = 1:n+m
    end
end

"""
    getgrad(alcon, i)

Get the view of the entire cost gradient corresponding to the inputs for the constraint.
"""
function getgrad(alcon::ALConstraint, i::Integer)
    inds = getinputinds(alcon, i)
    return view(alcon.grad[i], inds)
end

"""
    gethess(alcon, i)

Get the view of the entire cost Hessian corresponding to the inputs for the constraint.
"""
function gethess(alcon::ALConstraint, i::Integer)
    inds = getinputinds(alcon,i)
    return view(alcon.hess[i], inds, inds)
end

"""
    evaluate_constraint!(alcon, Z)

Evaluate the constraint at all time steps, storing the result in the [`ALConstraint`](@ref).
"""
function TO.evaluate_constraints!(alcon::ALConstraint)
    TO.evaluate_constraints!(alcon.sig, alcon.con, alcon.vals, alcon.Z[1], alcon.inds)
end

"""
    constraint_jacobians!(alcon, Z)

Evaluate the constraint Jacobian at all time steps, storing the result in the
[`ALConstraint`](@ref).
"""
function TO.constraint_jacobians!(alcon::ALConstraint)
    Z = alcon.Z[1]
    sig = function_signature(alcon)
    diff = alcon.diffmethod
    for (i,k) in enumerate(alcon.inds)
        RD.jacobian!(sig, diff, alcon.con, alcon.jac[i], alcon.vals[i], Z[k])
    end
end

@doc raw"""
    alcost(alcon)

Calculates the additional cost added by the augmented Lagrangian:

```math
\sum_{i=1}^{P} \frac{1}{2 \mu_i} || \Pi_K(\lambda_i - \mu_i c(x_k)) ||^2 - || \lambda_i ||^2
```

where ``k`` is the ``i``th knot point of ``P`` to which the constraint applies, and ``K`` is the
cone for the constraint.

Assumes that the constraints have already been evaluated via [`evaluate_constraint!`](@ref).
"""
function alcost(alcon::ALConstraint{T}) where T
    use_conic = alcon.opts.use_conic_cost
    cone = TO.sense(alcon.con)
    for (i,k) in enumerate(alcon.inds)
        if use_conic
            # Use generic conic cost
            J = alcost(alcon, i)
        else
            # Special-case on the cone
            J = alcost(cone, alcon, i)
        end
        # Add to the vector of AL penalty costs stored in the ALObjective
        # Hack to avoid allocation
        alcon.cost[k] += J
    end
    return nothing
end

"""
    algrad!(alcon, Z)

Evaluate the gradient of the augmented Lagrangian penalty cost for all time
steps.

Assumes that that the constraint and constraint Jacobians have already been
evaluated via [`evaluate_constraint!`](@ref) and [`constraint_jacobian!`](@ref),
and that [`alcost`](@ref) has already been called, which evaluates the
constraint values and approximate, projected, and scaled projected multipliers.

The gradient for the generic conic cost is of the following form:

```math
    \\nabla c(x)^T \\nabla \\Pi_K(\\bar{\\lambda})^T I_{\\mu}^{-1} \\Pi_K(\\bar{\\lambda})
```

where ``\\bar{\\lambda} = \\lambda - I_{\\mu} c(x)`` are the approximate dual
variables.
"""
function algrad!(alcon::ALConstraint{T}) where T
    use_conic = alcon.opts.use_conic_cost
    cone = TO.sense(alcon.con)
    for i in eachindex(alcon.inds)
        if use_conic
            # Use generic conic cost
            algrad!(alcon, i)
        else
            # Special-case on the cone
            algrad!(cone, alcon, i)
        end
    end
end

"""
    alhess!(alcon, Z)

Evaluate the Gauss-Newton Hessian of the augmented Lagrangian penalty cost for all time
steps.

Assumes that [`alcost`](@ref) and [`algrad`](@ref) have already been called, which evaluate
the constraint values, constraint Jacobians, and approximate, projected, and scaled
projected multipliers.

The Gauss-Newton Hessian for the generic conic cost is of the following form:

```math
    \\nabla c(x)^T \\nabla \\Pi_K(\\bar{\\lambda})^T I_{\\mu}^{-1} \\nabla \\Pi_K(\\bar{\\lambda}) \\nabla c(x)
    +  \\nabla c(x)^T \\nabla^2 \\Pi_K(\\bar{\\lambda}, I_{\\mu}^{-1} \\Pi_K(\\bar{\\lambda})) \\nabla c(x)
```

where ``\\bar{\\lambda} = \\lambda - I_{\\mu} c(x)`` are the approximate dual
variables.
"""
function alhess!(alcon::ALConstraint{T}) where T
    # Assumes Jacobians have already been computed
    use_conic = alcon.opts.use_conic_cost
    cone = TO.sense(alcon.con)
    for i in eachindex(alcon.inds)
        if use_conic
            # Use generic conic cost
            alhess!(alcon, i)
        else
            # Special-case on the cone
            alhess!(cone, alcon, i)
        end
    end
end

"""
    add_alcost_expansion!(alcon, E)

Add the precomputed gradient and Hessian of the AL penalty cost to the
cost expansion stored in `E`. Assumes [`alcost(alcon)`](@ref), [`algrad!(alcon)`](@ref),
and [`alhess!(alcon)`](@ref) have already been called to evaluate these terms
about the current trajectory.
"""
function add_alcost_expansion!(alcon::ALConstraint)
    E = alcon.E  # this is aliased to the one in the iLQR solver
    for (i,k) in enumerate(alcon.inds)
        E[k].grad .+= alcon.grad[i]
        E[k].hess .+= alcon.hess[i]
    end
end

##############################
# Equality Constraints
##############################
function alcost(::TO.Equality, alcon::ALConstraint, i::Integer)
    λ, μ, c = alcon.λ[i], alcon.μ[i], alcon.vals[i]
    Iμ = Diagonal(μ)
    return λ'c + 0.5 * dot(c,Iμ,c)
end

function algrad!(::TO.Equality, alcon::ALConstraint, i::Integer)
    λbar = alcon.λbar[i]
    λ, μ, c = alcon.λ[i], alcon.μ[i], alcon.vals[i]
    ∇c = alcon.jac[i]
    grad = getgrad(alcon, i)
    # grad = alcon.grad[i]

    λbar .= λ .+ μ .* c
    matmul!(grad, ∇c', λbar)
    return nothing
end

function alhess!(::TO.Equality, alcon::ALConstraint, i::Integer)
    ∇c = alcon.jac[i]
    hess = gethess(alcon, i)
    # hess = alcon.hess[i]
    tmp = alcon.tmp_jac
    # Iμ = Diagonal(alcon.μ[i])
    # matmul!(tmp, Iμ, ∇c)
    μ = alcon.μ[i]
    for i = 1:size(∇c,1), j = 1:size(∇c,2)
        tmp[i,j] = μ[i] * ∇c[i,j]
    end
    matmul!(hess, ∇c', tmp)
    return nothing
end

##############################
# Inequality Constraints
##############################
function alcost(::TO.Inequality, alcon::ALConstraint, i::Integer)
    λ, μ, c = alcon.λ[i], alcon.μ[i], alcon.vals[i]
    a = alcon.λbar[i]
    for i = 1:length(a)
        isactive = (c[i] >= 0) | (λ[i] > 0)
        a[i] = isactive * μ[i]
    end
    Iμ = Diagonal(a)
    return λ'c + 0.5 * dot(c,Iμ,c)
end

function algrad!(::TO.Inequality, alcon::ALConstraint, i::Integer)
    ∇c, λbar = alcon.jac[i], alcon.λbar[i]
    λ, μ, c = alcon.λ[i], alcon.μ[i], alcon.vals[i]
    grad = getgrad(alcon, i)
    # grad = alcon.grad[i]
    a = alcon.λbar[i]
    for i = 1:length(a)
        isactive = (c[i] >= 0) | (λ[i] > 0)
        a[i] = isactive * μ[i]
    end
    λbar .= λ .+ a .* c
    matmul!(grad, ∇c', λbar)
    return nothing
end

function alhess!(::TO.Inequality, alcon::ALConstraint, i::Integer)
    ∇c = alcon.jac[i]
    c = alcon.vals[i]
    λ, μ = alcon.λ[i], alcon.μ[i]
    hess = gethess(alcon, i)
    # hess = alcon.hess[i]
    tmp = alcon.tmp_jac
    a = alcon.λbar[i]
    for i = 1:length(a)
        isactive = (c[i] >= 0) | (λ[i] > 0)
        a[i] = isactive * μ[i]
    end
    # Iμ = Diagonal(a)
    for i = 1:size(∇c,1), j = 1:size(∇c,2)
        tmp[i,j] = a[i] * ∇c[i,j]
    end

    # matmul!(tmp, Iμ, ∇c)
    matmul!(hess, ∇c', tmp)
    return nothing
end

@inline alcost(::TO.ConstraintSense, alcon::ALConstraint, i::Integer) = alcost(alcon, i)
@inline algrad!(::TO.ConstraintSense, alcon::ALConstraint, i::Integer) = algrad!(alcon, i)
@inline alhess!(::TO.ConstraintSense, alcon::ALConstraint, i::Integer) = alhess!(alcon, i)

##############################
# Generic Cones
##############################

function alcost(alcon::ALConstraint, i::Integer)
    dualcone = TO.dualcone(TO.sense(alcon.con))

    λ, λbar, λp, λs = alcon.λ[i], alcon.λbar[i], alcon.λproj[i], alcon.λscaled[i]
    μ, μinv, c = alcon.μ[i], alcon.μinv[i], alcon.vals[i]

    # Approximate dual
    λbar .= λ .- μ .* c

    # Projected approximate dual
    TO.projection!(dualcone, λp, λbar)

    # Scaled dual
    μinv .= inv.(μ)
    λs .= μinv .* λp

    # Cost
    Iμ = Diagonal(μinv)
    return 0.5 * (λp'λs - λ'Iμ*λ)
end

function algrad!(alcon::ALConstraint, i::Integer)
    dualcone = TO.dualcone(TO.sense(alcon.con))

    # Assume λbar and λp have already been calculated
    λbar, λp, λs = alcon.λbar[i], alcon.λproj[i], alcon.λscaled[i]
    μ, c, ∇c, Iμ∇c = alcon.μ[i], alcon.vals[i], alcon.jac[i], alcon.jac_scaled[i]
    ∇proj = alcon.∇proj[i]
    grad = getgrad(alcon, i)
    # grad = alcon.grad[i]
    tmp = alcon.tmp_jac

    # Scale the Jacobian
    p,nm = size(∇c)
    for i = 1:p
        for j = 1:nm
            Iμ∇c[i,j] = ∇c[i,j] * μ[i]
        end
    end

    # grad = -∇c'∇proj'Iμ*λp
    TO.∇projection!(dualcone, ∇proj, λbar)
    matmul!(tmp, ∇proj, Iμ∇c)  # derivative of λp wrt x
    tmp .*= -1
    matmul!(grad, tmp', λs)
    return nothing
end

function alhess!(alcon::ALConstraint, i::Integer)
    dualcone = TO.dualcone(TO.sense(alcon.con))

    λbar, λs = alcon.λbar[i], alcon.λscaled[i]
    μ, c, ∇c = alcon.μ[i], alcon.vals[i], alcon.jac_scaled[i]
    ∇proj, ∇²proj = alcon.∇proj[i], alcon.∇²proj[i]
    Iμ∇proj = alcon.∇proj_scaled[i]
    hess = gethess(alcon, i)
    # hess = alcon.hess[i]
    tmp = alcon.tmp_jac

    # Assume 𝝯proj is already computed
    # TODO: reuse this from before
    matmul!(tmp, ∇proj, ∇c)  # derivative of λp wrt x
    tmp .*= -1

    # Scale projection Jacobian
    p = length(λbar)
    μinv = alcon.μinv[i]
    for i = 1:p, j = 1:p
        Iμ∇proj[i,j] = ∇proj[i,j] * μinv[i]
    end

    # Calculate second-order projection term
    TO.∇²projection!(dualcone, ∇²proj, λbar, λs)
    matmul!(∇²proj, ∇proj', Iμ∇proj, 1.0, 1.0)

    # hess = ∇c'Iμ*(∇²proj(λs) + ∇proj'Iμ\∇proj)*Iμ*∇c
    matmul!(tmp, ∇²proj, ∇c)
    matmul!(hess, ∇c', tmp)
    return nothing
end

##############################
# Dual and Penalty Updates
##############################

"""
    dualupdate!(alcon)

Update the dual variables for all time times for a single constraint.
For now, always performs the update.
"""
function dualupdate!(alcon::ALConstraint)
    cone = TO.sense(alcon.con)
    use_conic = alcon.opts.use_conic_cost
    λ_max = alcon.opts.dual_max
    for i in eachindex(alcon.inds)
        λ, μ, c = alcon.λ[i], alcon.μ[i], alcon.vals[i]
        if use_conic
            dualupdate!(alcon, i)
        else
            # Special-case to the cone
            dualupdate!(cone, alcon, i)
        end
        # Saturate dual variables
        clamp!(λ, -λ_max, λ_max)
    end
end

function dualupdate!(::TO.Equality, alcon::ALConstraint, i::Integer)
    λbar, λ, μ, c = alcon.λbar[i], alcon.λ[i], alcon.μ[i], alcon.vals[i]
    λbar .= λ .+ μ .* c
    λ .= λbar
    return nothing
end

function dualupdate!(::TO.Inequality, alcon::ALConstraint, i::Integer)
    λbar, λ, μ, c = alcon.λbar[i], alcon.λ[i], alcon.μ[i], alcon.vals[i]
    λbar .= λ .+ μ .* c
    λ .= max.(0, λbar)
    return nothing
end

@inline dualupdate!(::TO.SecondOrderCone, alcon::ALConstraint, i::Integer) =
    dualupdate!(alcon, i)

function dualupdate!(alcon::ALConstraint, i::Integer)
    # TODO: just copy the projected duals already stored
    dualcone = TO.dualcone(TO.sense(alcon.con))
    λbar, λ, μ, c = alcon.λbar[i], alcon.λ[i], alcon.μ[i], alcon.vals[i]
    λbar .= λ .- μ .* c
    TO.projection!(dualcone, λ, λbar)
    return nothing
end

"""
    penaltyupate!(alcon)

Update the penalty terms by the geometric factor `opts.penalty_scaling`.
Always updates the penalties and updates all penalties equally, thresholding at the maximum
specified by `opts.penalty_max`.
"""
function penaltyupdate!(alcon::ALConstraint)
    μ = alcon.μ
    ϕ = alcon.opts.penalty_scaling
    μ_max = alcon.opts.penalty_max
    for i = 1:length(alcon.inds)
        μ[i] .*= ϕ
        clamp!(alcon.μ[i], 0, μ_max)
        alcon.μinv[i] .= inv.(μ[i])
    end
end

##############################
# Max Violation and Penalty
##############################

"""
    normviolation!(alcon, p)

Evaluate the `p`-norm of the constraint violations. The violation is
defined to be

```math
\\Pi_K(c(x)) - c(x)
```
These values for each time step are stored in `alcon.viol`.
"""
function normviolation!(alcon::ALConstraint, p=2, c_max=alcon.c_max)
    cone = TO.sense(alcon.con)
    c_max === alcon.c_max && (c_max .= 0)
    for (i,k) in enumerate(alcon.inds)
        TO.projection!(cone, alcon.viol[i], alcon.vals[i])
        alcon.viol[i] .-= alcon.vals[i]
        v = norm(alcon.viol[i], p)
        if isinf(p)
            c_max[k] = max(c_max[k], v)
        else
            c_max[k] = norm(SA[c_max[k], v], p)
        end
    end
    # return norm(alcon.c_max, p)
    return nothing
end
max_violation!(alcon::ALConstraint, c_max=alcon.c_max) = normviolation!(alcon, Inf, c_max)

"""
    max_penalty(alcon)

Find the maximum penalty value over all time steps.
"""
function max_penalty(alcon::ALConstraint{T}) where T
    μ_max = zero(T)
    for i = 1:length(alcon.inds)
        μ_max = max(maximum(alcon.μ[i]), μ_max)
    end
    return μ_max
end

##############################
# Reset functions
##############################
"""
    reset_duals!(alcon)

Reset the dual variables to zero.
"""
function reset_duals!(alcon::ALConstraint)
    for i = 1:length(alcon.inds)
        alcon.λ[i] .= 0
    end
end

"""
    reset_penalties!(alcon)

Reset tne penalties to the initial penalty specified by `opts.penalty_initial`.
"""
function reset_penalties!(alcon::ALConstraint)
    μ_initial = alcon.opts.penalty_initial
    for i = 1:length(alcon.inds)
        alcon.μ[i] .= μ_initial
    end
end


text:
#!/usr/bin/env julia

using SphericalFourierBesselDecompositions
SFB = SphericalFourierBesselDecompositions

using Test
using LinearAlgebra

@testset "Window Chains" begin
    # Note: a lot of the differences here come down to the choice of nside and
    # details of how healpy is used.

    @testset "Wmix" begin
        rmin = 500.0
        rmax = 1000.0
        amodes = SFB.AnlmModes(3, 5, rmin, rmax)
        cmodes = SFB.ClnnModes(amodes, Δnmax=1)
        wmodes = SFB.ConfigurationSpaceModes(rmin, rmax, 1000, amodes.nside)
        win = SFB.make_window(wmodes, :radial, :ang_quarter, :rotate, :separable)

        cache2 = SFB.WindowChains.WindowChainsCacheFullWmix(win, wmodes, amodes)
        cache3 = SFB.WindowChains.WindowChainsCacheFullWmixOntheflyWmix(win, wmodes, amodes)
        cache4 = SFB.WindowChains.WindowChainsCacheSeparableWmix(win, wmodes, amodes)
        cache5 = SFB.WindowChains.WindowChainsCacheSeparableWmixOntheflyWlmlm(win, wmodes, amodes)

        ## test cache2 consistency
        #@show real.(cache2.wmix') real.(cache2.wmix)
        #@show imag.(cache2.wmix') imag.(cache2.wmix)
        @test cache2.wmix' ≈ cache2.wmix
        #@show real.(cache2.wmix_negm') real.(cache2.wmix_negm)
        #@show imag.(cache2.wmix_negm') imag.(cache2.wmix_negm)
        nlmsize = SFB.getnlmsize(amodes)
        for i=1:nlmsize, i′=1:nlmsize
            n1, l1, m1 = SFB.getnlm(amodes, i)
            n2, l2, m2 = SFB.getnlm(amodes, i′)
            @test cache2.wmix_negm[i′,i] ≈ (-1)^(m1+m2) * cache2.wmix_negm[i,i′]
        end

        get_win(nlm1, nlm2; verbose=false) = begin
            n1, l1, m1 = nlm1
            n2, l2, m2 = nlm2

            nl1 = SFB.getidx(cache2.amodes, n1, l1, 0)
            nl2 = SFB.getidx(cache2.amodes, n2, l2, 0)
            w2 = SFB.WindowChains.get_wmix(cache2.wmix, cache2.wmix_negm,
                                           nl1, m1, nl2, m2)
            verbose && @show cache2.wmix[nl1+abs(m1),nl2+abs(m2)]
            verbose && @show cache2.wmix'[nl1+abs(m1),nl2+abs(m2)]
            verbose && @show conj(cache2.wmix[nl2+abs(m2),nl1+abs(m1)])
            verbose && @show cache2.wmix_negm[nl1+abs(m1),nl2+abs(m2)]
            verbose && @show cache2.wmix_negm'[nl1+abs(m1),nl2+abs(m2)]
            verbose && @show conj(cache2.wmix_negm[nl2+abs(m2),nl1+abs(m1)])

            wlmlm4 = SFB.WindowChains.get_wlmlm(cache4, l1, m1, l2, m2)
            inlnl4 = cache4.Ilnln[l1+1,n1,l2+1,n2]
            w4 = inlnl4 * wlmlm4

            wlmlm5 = SFB.WindowChains.window_wmix(l1, m1, l2, m2,
                                                  cache5.Wlm, cache5.LMcache)
            inlnl5 = cache5.Ilnln[l1+1,n1,l2+1,n2]
            w5 = inlnl5 * wlmlm5
            return w2, w4, w5
        end

        test_mode(nlm1, nlm2; verbose=true) = begin
            w2a, w4a, w5a = get_win(nlm1, nlm2; verbose=verbose)
            verbose && @show nlm1,nlm2
            verbose && @show w2a w4a w5a
            @test w2a ≈ w4a ≈ w5a

            w2b, w4b, w5b = get_win(nlm2, nlm1; verbose=verbose)
            verbose && @show nlm2,nlm1
            verbose && @show w2b w4b w5b
            @test w2b ≈ w4b ≈ w5b

            nlm1[3] = -nlm1[3]
            nlm2[3] = -nlm2[3]
            w2c, w4c, w5c = get_win(nlm1, nlm2; verbose=verbose)
            verbose && @show nlm1,nlm2,nlm1[3]+nlm2[3]
            verbose && @show w2c w4c w5c
            @test w2c ≈ w4c ≈ w5c

            @test w2b ≈ conj(w2a)
            @test w2c ≈ (-1)^(nlm1[3] + nlm2[3]) * conj(w2a)
            @test w4b ≈ conj(w4a)
            @test w4c ≈ (-1)^(nlm1[3] + nlm2[3]) * conj(w4a)
            @test w5b ≈ conj(w5a)
            @test w5c ≈ (-1)^(nlm1[3] + nlm2[3]) * conj(w5a)
        end

        test_mode([1, 1, -1], [1, 0, 0])
        test_mode([1, 1, 1], [1, 1, 1])
        test_mode([1, 1, -1], [1, 1, 1])

        nlmsize = SFB.getnlmsize(amodes)
        long_tests && for j=1:nlmsize, i=1:nlmsize
            n1, l1, m1 = SFB.getnlm(amodes, i)
            n2, l2, m2 = SFB.getnlm(amodes, j)
            test_mode([n1, l1, m1], [n2, l2, m2]; verbose=false)
            test_mode([n1, l1, -m1], [n2, l2, m2]; verbose=false)
            test_mode([n1, l1, m1], [n2, l2, -m2]; verbose=false)
            test_mode([n1, l1, -m1], [n2, l2, -m2]; verbose=false)
        end
    end

    @testset "k = 1" begin
        k = 1
        rmin = 500.0
        rmax = 1000.0
        amodes = SFB.AnlmModes(3, 5, rmin, rmax)
        cmodes = SFB.ClnnModes(amodes, Δnmax=1)
        wmodes = SFB.ConfigurationSpaceModes(rmin, rmax, 1000, amodes.nside)
        win = SFB.make_window(wmodes, :radial, :ang_quarter, :rotate, :separable)

        wlnn = SFB.win_lnn(win, wmodes, cmodes)
        @show wlnn

        cache1 = SFB.WindowChains.WindowChainsCacheWignerChain(win, wmodes, amodes)
        cache2 = SFB.WindowChains.WindowChainsCacheFullWmix(win, wmodes, amodes)
        cache3 = SFB.WindowChains.WindowChainsCacheFullWmixOntheflyWmix(win, wmodes, amodes)
        cache4 = SFB.WindowChains.WindowChainsCacheSeparableWmix(win, wmodes, amodes)
        cache5 = SFB.WindowChains.WindowChainsCacheSeparableWmixOntheflyWlmlm(win, wmodes, amodes)
        cache = SFB.WindowChainsCache(win, wmodes, amodes)

        # single modes
        ell = [1]
        n1 = [1]
        n2 = [1]
        wk0 = (2*ell[1] + 1) * wlnn[SFB.getidx(cmodes, ell[1], n1[1], n2[1])]
        wk1 = SFB.window_chain(ell, n1, n2, cache1)
        wk2 = SFB.window_chain(ell, n1, n2, cache2)
        wk3 = SFB.window_chain(ell, n1, n2, cache3)
        wk4 = SFB.window_chain(ell, n1, n2, cache4)
        wk5 = SFB.window_chain(ell, n1, n2, cache5)
        wk = SFB.window_chain(ell, n1, n2, cache)
        @show wk wk0 wk1 wk2 wk3 wk4 wk5
        @test wk ≈ wk0
        @test_skip wk ≈ wk1
        @test wk ≈ wk2
        @test wk ≈ wk3
        @test wk ≈ wk4
        @test wk ≈ wk5

        # all modes
        wk_lnni = SFB.window_chain(k, win, wmodes, cmodes)
        @show wk_lnni
        @test all(isreal.(wk_lnni))
        wk_lnni = real.(wk_lnni)

        @show size(wk_lnni) size(wlnn)
        long_tests && for i=1:SFB.getlnnsize(cmodes)
            l, n1, n2 = SFB.getlnn(cmodes, i)
            wk = wk_lnni[i] / (2*l+1)
            @show l,n1,n2,wk_lnni[i],wk,wlnn[i]
            @test wk ≈ wlnn[i]
        end
    end

    @testset "k = 2" begin
        k = 2
        rmin = 500.0
        rmax = 1000.0
        amodes = SFB.AnlmModes(2, 2, rmin, rmax)
        cmodes = SFB.ClnnModes(amodes, Δnmax=1)
        wmodes = SFB.ConfigurationSpaceModes(rmin, rmax, 1000, amodes.nside)
        win = SFB.make_window(wmodes, :radial, :ang_quarter, :rotate, :separable)

        # comparison
        @time M = SFB.power_win_mix(win, wmodes, cmodes)
        @time M′ = SFB.power_win_mix(win, wmodes, cmodes, interchange_NN′=true)

        cache1 = SFB.WindowChains.WindowChainsCacheWignerChain(win, wmodes, amodes)
        cache2 = SFB.WindowChains.WindowChainsCacheFullWmix(win, wmodes, amodes)
        cache3 = SFB.WindowChains.WindowChainsCacheFullWmixOntheflyWmix(win, wmodes, amodes)
        cache4 = SFB.WindowChains.WindowChainsCacheSeparableWmix(win, wmodes, amodes)
        cache5 = SFB.WindowChains.WindowChainsCacheSeparableWmixOntheflyWlmlm(win, wmodes, amodes)
        cache = SFB.WindowChainsCache(win, wmodes, amodes)

        # single modes
        test_single(ell, n1, n2) = begin
            @show ell,n1,n2
            lnnA = SFB.getidx(cmodes,ell[2],n2[2],n1[2])
            lnnB = SFB.getidx(cmodes,ell[1],n1[1],n2[1])
            wk0 = (2*ell[2] + 1) * M′[lnnA, lnnB]
            wk1 = SFB.window_chain(ell, n1, n2, cache1)
            wk2 = SFB.window_chain(ell, n1, n2, cache2)
            wk3 = SFB.window_chain(ell, n1, n2, cache3)
            wk4 = SFB.window_chain(ell, n1, n2, cache4)
            wk5 = SFB.window_chain(ell, n1, n2, cache5)
            wk = SFB.window_chain(ell, n1, n2, cache)
            @show wk wk0 wk1 wk2 wk3 wk4 wk5
            @test wk ≈ wk0
            @test_skip wk ≈ wk1
            @test wk ≈ wk2
            @test wk ≈ wk3
            @test wk ≈ wk4
            @test wk ≈ wk5
        end
        test_single([1,1], [1,1], [1,1])
        test_single([2,0], [1,1], [2,1])
        test_single([0,0], [1,1], [1,2])
        test_single([0,0], [1,1], [2,2])
        test_single([0,1], [1,1], [1,1])

        # all modes
        @time wk_lnni = SFB.window_chain(k, win, wmodes, cmodes)
        @time wk_lnni = SFB.window_chain(k, win, wmodes, cmodes)
        @test all(isreal.(wk_lnni))

        @show size(wk_lnni) size(M)
        @test size(wk_lnni) == size(M)
        long_tests && for i2=1:SFB.getlnnsize(cmodes), i1=1:SFB.getlnnsize(cmodes)
            l1, n1, n1′ = SFB.getlnn(cmodes, i1)
            l2, n2, n2′ = SFB.getlnn(cmodes, i2)
            i2′ = SFB.getidx(cmodes, l2, n2′, n2)
            wk = wk_lnni[i1,i2′] / (2*l2+1)
            @show l1,n1,n1′,l2,n2,n2′,wk,M′[i2,i1]
            @test wk ≈ M′[i2,i1]
        end
    end

    @testset "Wk method agreement" begin
        rmin = 500.0
        rmax = 1000.0
        amodes = SFB.AnlmModes(8, 4, rmin, rmax)
        cmodes = SFB.ClnnModes(amodes, Δnmax=0)
        wmodes = SFB.ConfigurationSpaceModes(rmin, rmax, 1000, amodes.nside)
        win = SFB.make_window(wmodes, :radial, :ang_quarter, :rotate, :separable)

        cache1 = SFB.WindowChains.WindowChainsCacheWignerChain(win, wmodes, amodes)
        cache2 = SFB.WindowChains.WindowChainsCacheFullWmix(win, wmodes, amodes)
        cache3 = SFB.WindowChains.WindowChainsCacheFullWmixOntheflyWmix(win, wmodes, amodes)
        cache4 = SFB.WindowChains.WindowChainsCacheSeparableWmix(win, wmodes, amodes)
        cache5 = SFB.WindowChains.WindowChainsCacheSeparableWmixOntheflyWlmlm(win, wmodes, amodes)
        cache = SFB.WindowChainsCache(win, wmodes, amodes)

        test_individual(ell, n1, n2) = begin
            @time wk1 = SFB.window_chain(ell, n1, n2, cache1)
            @time wk2 = SFB.window_chain(ell, n1, n2, cache2)
            @time wk3 = SFB.window_chain(ell, n1, n2, cache3)
            @time wk4 = SFB.window_chain(ell, n1, n2, cache4)
            @time wk5 = SFB.window_chain(ell, n1, n2, cache5)
            @time wk = SFB.window_chain(ell, n1, n2, cache)
            @show ell n1 n2
            @show wk1 wk2 wk3 wk4 wk5 wk
            @test_skip wk1 ≈ wk
            @test wk2 ≈ wk
            @test wk3 ≈ wk
            @test wk4 ≈ wk
            @test wk5 ≈ wk
            return wk
        end

        println("Testing individual modes...")
        long_tests && for l1=0:amodes.lmax
            test_individual([l1], [1], [1])
            for l2=0:amodes.lmax
                test_individual([l1, l2], [1, 1], [1, 1])
                for l3=0:amodes.lmax
                    test_individual([l1, l2, l3], [1, 1, 1], [1, 1, 1])
                end
            end
        end
        #test_individual([0, 0, 0, 1], [1, 1, 1, 1], [1, 1, 1, 1])
        #test_individual([1, 3, 4, 5], [1, 1, 1, 1], [1, 1, 1, 1])
        #test_individual([2, 3, 4, 5], [1, 1, 1, 1], [1, 1, 1, 1])
        #test_individual([3, 3, 3, 3], [1, 1, 2, 2], [2, 2, 1, 1])

        # check n ↔ n′ symmetry doesn't exist
        ell = [2, 1, 0]
        n1 = [4, 2, 2]
        n2 = [2, 3, 3]
        test_individual(ell, n1, n2)
        test_individual(ell, n2, n1)

        # check antisymmetry doesn't exist
        ell = [2, 1, 0, 1]
        n1 = [4, 2, 2, 2]
        n2 = [2, 3, 3, 2]
        test_individual(ell, n1, n2)
        test_individual(ell, n2, n1)
        ell[1:2] .= ell[2:-1:1]
        n1[1:2] .= n1[2:-1:1]
        n2[1:2] .= n2[2:-1:1]
        test_individual(ell, n1, n2)

        # check anticyclic symmetry
        #ell = [2, 1, 0, 1]
        #n1 = [2, 3, 3, 2]
        #n2 = [4, 2, 2, 2]
        for i=1:10
            k = 4
            p1 = k:-1:1
            ell = rand(0:amodes.lmax, k)
            n1 = rand(1:amodes.nmax, k)
            n2 = rand(1:amodes.nmax, k)
            wk = test_individual(ell, n1, n2)
            wkp1 = test_individual(ell[p1], n2[p1], n1[p1])
            @show wk wkp1
            @test wk ≈ wkp1
        end

        return

        # These checks can fail:
        # check k=3 n,N ↔ n′,N′ symmetry, Nah, not a symmetry
        for i=1:10
            k = 3
            ell = rand(0:amodes.lmax, k)
            n1 = rand(1:amodes.nmax, k)
            n2 = rand(1:amodes.nmax, k)
            wk = test_individual(ell, n1, n2)
            n1[2:3] .= n1[3:-1:2]
            n2[2:3] .= n2[3:-1:2]
            wkp1 = test_individual(ell, n2, n1)
            @show wk wkp1
            @test wk ≈ wkp1
        end
    end

    @testset "Symmetricizing f=$features" for features in [(), (:separable,)]
        rmin = 500.0
        rmax = 1000.0
        amodes = SFB.AnlmModes(8, 4, rmin, rmax)
        cmodes = SFB.ClnnModes(amodes, Δnmax=0)
        wmodes = SFB.ConfigurationSpaceModes(rmin, rmax, 1000, amodes.nside)
        win = SFB.make_window(wmodes, :radial, :ang_quarter, :rotate, features...)
        cache = SFB.WindowChainsCache(win, wmodes, amodes)

        ell = [1, 2, 3, 4]
        n1 = [1, 2, 3, 4]
        n2 = [5, 6, 7, 8]
        symmetries = [1=>0, 2=>1, 3=>2]
        wkfull = SFB.window_chain(ell, n1, n2, cache, symmetries)
        @show wkfull
        @test wkfull ≈ -1.2707026010569427e-8  # only useful for detecting changes
    end

    @testset "Wk single pixel mask" begin
        rmin = 500.0
        rmax = 1000.0
        nside = 16
        amodes = SFB.AnlmModes(8, 4, rmin, rmax; nside=nside)
        cmodes = SFB.ClnnModes(amodes, Δnmax=0)
        wmodes = SFB.ConfigurationSpaceModes(rmin, rmax, 1000, amodes.nside)
        mask = fill(0.0, SFB.hp.nside2npix(amodes.nside))
        mask[805+1] = 1.0
        phi = fill(1.0, wmodes.nr)
        win = SFB.SeparableArray(phi, mask, name1=:phi, name2=:mask)

        cache1 = SFB.WindowChains.WindowChainsCacheWignerChain(win, wmodes, amodes)
        cache2 = SFB.WindowChains.WindowChainsCacheFullWmix(win, wmodes, amodes)
        cache3 = SFB.WindowChains.WindowChainsCacheFullWmixOntheflyWmix(win, wmodes, amodes)
        cache4 = SFB.WindowChains.WindowChainsCacheSeparableWmix(win, wmodes, amodes)
        cache5 = SFB.WindowChains.WindowChainsCacheSeparableWmixOntheflyWlmlm(win, wmodes, amodes)
        cache = SFB.WindowChainsCache(win, wmodes, amodes)

        test_individual(ell, n1, n2) = begin
            @time wk1 = SFB.window_chain(ell, n1, n2, cache1)
            @time wk2 = SFB.window_chain(ell, n1, n2, cache2)
            @time wk3 = SFB.window_chain(ell, n1, n2, cache3)
            @time wk4 = SFB.window_chain(ell, n1, n2, cache4)
            @time wk5 = SFB.window_chain(ell, n1, n2, cache5)
            @time wk = SFB.window_chain(ell, n1, n2, cache)
            @show ell n1 n2
            @show wk1 wk2 wk3 wk4 wk5 wk
            @test_skip wk1 ≈ wk
            @test wk2 ≈ wk
            @test wk3 ≈ wk
            @test wk4 ≈ wk
            @test wk5 ≈ wk
            return wk
        end

        println("Testing individual modes...")
        for l1=0:amodes.lmax
            test_individual([l1], [1], [1])
            long_tests && for l2=0:amodes.lmax
                test_individual([l1, l2], [1, 1], [1, 1])
                for l3=0:amodes.lmax
                    test_individual([l1, l2, l3], [1, 1, 1], [1, 1, 1])
                end
            end
        end
    end

    @testset "Other tests" begin
        rmin = 0.0
        rmax = 2000.0
        amodes = SFB.AnlmModes(0.03, rmin, rmax)
        cmodes = SFB.ClnnModes(amodes, Δnmax=0)
        wmodes = SFB.ConfigurationSpaceModes(rmin, rmax, 1000, amodes.nside)
        @show amodes.lmax_n

        phi = fill(1.0, wmodes.nr)
        mask = fill(1.0, SFB.hp.nside2npix(amodes.nside))
        win = SFB.SeparableArray(phi, mask, name1=:phi, name2=:mask)

        wlm, LMcache = SFB.WindowChains.calc_Wlm(mask, amodes.lmax, amodes.nside)
        lrange = 48:48
        long_tests && for l1=lrange, l2=lrange, m1=-l1:l1, m2=-l2:l2
            w = SFB.WindowChains.window_wmix(l1, m1, l2, m2, wlm, LMcache)
            if l1 == l2 && m1 == m2
                @test w ≈ 1
                if !(w ≈ 1)
                    @error "Wlmlm unexpected result" l1,l2 m1,m2 w
                end
            else
                @test w ≈ 0  atol=1e-10
                if !(abs(w) <= 1e-10)
                    @error "Wlmlm unexpected result" l1,l2 m1,m2 w
                end
            end
        end

        #return

        # Note: the following tests can only be done reasonably fast with
        # WignerFamilies instead of WignerSymbols in the window chain.

        cache = SFB.WindowChainsCache(win, wmodes, amodes)
        @show typeof(cache)
        for l=0:amodes.lmax
            nmax = amodes.nmax_l[l+1]
            @test cache.Ilnln[l+1,1:nmax,l+1,1:nmax] ≈ I  atol=1e-4
            Ik = SFB.WindowChains.calc_kprod(cache.Ilnln, [l,l] .+ 1, [1,1], [1,1])
            @test Ik ≈ 1  atol=1e-4
        end

        ell = [48, 48]
        n1 = [1, 1]
        n2 = [1, 1]
        wk = SFB.window_chain(ell, n1, n2, cache)
        @show wk wk/(2*ell[1]+1)
        @test wk/(2*ell[1]+1) ≈ 1

        long_tests && for m1=-ell[1]:ell[1], m2=-ell[2]:ell[2]
            w = SFB.WindowChains.get_wlmlm(cache, ell[end], m2, ell[1], m1)
            w2 = SFB.WindowChains.window_wmix(ell[end], m2, ell[1], m1, wlm, LMcache)
            if m1 != m2 || ell[1] != ell[2]
                @test w ≈ 0  atol=1e-10
                @test w2 ≈ 0  atol=1e-10
            else
                @test w ≈ 1
                @test w2 ≈ 1
                if !(w ≈ 1)
                    @error "Wlmlm unexpected result" ell m1,m2 w w2
                end
            end
        end

        ell = [48, 0]
        n1 = [1, 1]
        n2 = [2, 2]
        wk = SFB.window_chain(ell, n1, n2, cache)
        @show wk
    end
end

# vim: set sw=4 et sts=4 :


text:
"""
Macro if used in a `@resumable function` that returns the `expr` otherwise throws an error.
"""
macro yield(expr=nothing)
  error("@yield macro outside a @resumable function!")
end

"""
Macro if used in a `@resumable function` that creates a not saved variable otherwise throws an error.
"""
macro nosave(expr=nothing)
  error("@nosave macro outside a @resumable function!")
end

"""
Macro that transforms a function definition in a finite-statemachine:

- Defines a new `mutable struct` that implements the iterator interface and is used to store the internal state.
- Makes this new type callable having following characteristics:
  - implementents the statements from the initial function definition but;
  - returns at a `@yield` statement and;
  - continues after the `@yield` statement when called again.
- Defines a constructor function that respects the calling conventions of the initial function definition and returns an object of the new type.
"""
macro resumable(expr::Expr)
  expr.head !== :function && error("Expression is not a function definition!")
  func_def = splitdef(expr)
  rtype = :rtype in keys(func_def) ? func_def[:rtype] : Any
  args, kwargs, arg_dict = get_args(func_def)
  params = ((get_param_name(param) for param in func_def[:whereparams])...,)
  ui8 = BoxedUInt8(zero(UInt8))
  postwalk(x->hasreturnvalue(x), func_def[:body])
  func_def[:body] = postwalk(x->transform_for(x, ui8), func_def[:body])
  slots = get_slots(copy(func_def), arg_dict, __module__)
  type_name = gensym()
  constr_def = copy(func_def)
  if isempty(params)
    struct_name = :($type_name <: ResumableFunctions.FiniteStateMachineIterator{$rtype})
    constr_def[:name] = :($type_name)
  else
    struct_name = :($type_name{$(func_def[:whereparams]...)} <: ResumableFunctions.FiniteStateMachineIterator{$rtype})
    constr_def[:name] = :($type_name{$(params...)})
  end
  constr_def[:args] = tuple()
  constr_def[:kwargs] = tuple()
  constr_def[:rtype] = nothing
  constr_def[:body] = quote
    fsmi = new()
    fsmi._state = 0x00
    fsmi
  end
  constr_expr = combinedef(constr_def) |> flatten
  type_expr = :(
    mutable struct $struct_name
      _state :: UInt8
      $((:($slotname :: $slottype) for (slotname, slottype) in slots)...)
      $(constr_expr)
    end
  )
  @debug type_expr|>MacroTools.striplines
  call_def = copy(func_def)
  if isempty(params)
    call_def[:rtype] = nothing
    call_def[:body] = quote
      fsmi = $type_name()
      $((arg !== Symbol("_") ? :(fsmi.$arg = $arg) : nothing for arg in args)...)
      $((:(fsmi.$arg = $arg) for arg in kwargs)...)
      fsmi
    end
  else
    call_def[:rtype] = nothing
    call_def[:body] = quote
      fsmi = $type_name{$(params...)}()
      $((arg !== Symbol("_") ? :(fsmi.$arg = $arg) : nothing for arg in args)...)
      $((:(fsmi.$arg = $arg) for arg in kwargs)...)
      fsmi
    end
  end
  call_expr = combinedef(call_def) |> flatten
  @debug call_expr|>MacroTools.striplines
  if isempty(params)
    func_def[:name] = :((_fsmi::$type_name))
  else
    func_def[:name] = :((_fsmi::$type_name{$(params...)}))
  end
  func_def[:rtype] = nothing
  func_def[:body] = postwalk(x->transform_slots(x, keys(slots)), func_def[:body])
  func_def[:body] = postwalk(transform_arg, func_def[:body])
  func_def[:body] = postwalk(transform_exc, func_def[:body]) |> flatten
  ui8 = BoxedUInt8(zero(UInt8))
  func_def[:body] = postwalk(x->transform_try(x, ui8), func_def[:body])
  ui8 = BoxedUInt8(zero(UInt8))
  func_def[:body] = postwalk(x->transform_yield(x, ui8), func_def[:body])
  func_def[:body] = postwalk(x->transform_nosave(x, Set{Symbol}()), func_def[:body])
  func_def[:body] = quote
    _fsmi._state === 0x00 && @goto $(Symbol("_STATE_0"))
    $((:(_fsmi._state === $i && @goto $(Symbol("_STATE_",:($i)))) for i in 0x01:ui8.n)...)
    error("@resumable function has stopped!")
    @label $(Symbol("_STATE_0"))
    _fsmi._state = 0xff
    _arg isa Exception && throw(_arg)
    $(func_def[:body])
  end
  func_def[:args] = [Expr(:kw, :(_arg::Any), nothing)]
  func_def[:kwargs] = []
  func_expr = combinedef(func_def) |> flatten
  @debug func_expr|>MacroTools.striplines
  esc(:($type_expr; $func_expr; $call_expr))
end


text:
# percentages
data = [30 40 30; 50 25 25; 30 30 40; 10 50 40]
# cumulative sum so that columns increase
data_cs = cumsum(data; dims=2)
bar(data_cs; stacked=true, fills=["midnightblue", "lightseagreen", "lightsalmon"])


text:
# This file was generated, do not modify it. # hide
using LinearAlgebra

k = 4
n = 7

W = [I; randn(n - k, k)]

max_stars_repo_name: cscherrer/InformativePrior

max_issues_repo_name: cscherrer/InformativePrior

max_forks_repo_name: cscherrer/InformativePrior


text:
export residuals, χ²

@recipe function f(m::AbstractModel, dvis::EHTObservation{T,A}; datamarker=:circle, datacolor=:grey) where {T,A<:EHTVisibilityDatum}
    xguide --> "uv-distance (λ)"
    yguide --> "V (Jy)"
    markershape --> :diamond

    u = getdata(dvis, :u)
    v = getdata(dvis, :v)
    uvdist = hypot.(u,v)
    vis = visibility.(dvis.data)
    error = getdata(dvis, :error)
    vre = real.(vis)
    vim = imag.(vis)
    #add data errorbars
    @series begin
        seriestype := :scatter
        markershape := datamarker
        markercolor := datacolor
        alpha := 0.5
        yerr := error
        linecolor := nothing
        label := "Data"
        uvdist, vre
    end

    @series begin
        seriestype := :scatter
        markershape := datamarker
        markercolor := :white
        markeralpha := 0.01
        markerstrokecolor := :black
        markerstrokealpha := 1.0
        linecolor :=nothing
        label := nothing
        yerr := error
        uvdist, vim
    end
    seriestype-->:scatter
    vmod = visibilities(m, u, v)
    labels --> "Model"
    uvdist, hcat(real.(vmod), imag.(vmod))
end

@recipe function f(dvis::EHTObservation{T,A};) where {T,A<:EHTVisibilityDatum}
    xguide --> "uv-distance (λ)"
    yguide --> "V (Jy)"
    markershape --> :circle

    u = getdata(dvis, :u)
    v = getdata(dvis, :v)
    uvdist = hypot.(u,v)
    vis = visibility.(dvis.data)
    error = getdata(dvis, :error)
    vre = real.(vis)
    vim = imag.(vis)
    #add data errorbars
    @series begin
        seriestype := :scatter
        alpha := 0.5
        yerr := error
        linecolor := nothing
        label := "Real"
        uvdist, vre
    end

    @series begin
        seriestype := :scatter
        markeralpha := 0.1
        markerstrokecolor := :black
        markerstrokealpha := 1.0
        linecolor :=nothing
        label := nothing
        yerr := error
        label := "Imag"
        uvdist, vim
    end
end

@recipe function f(dvis::EHTObservation{T,A};) where {T,A<:EHTVisibilityAmplitudeDatum}
    xguide --> "uv-distance (λ)"
    yguide --> "|V| (Jy)"
    markershape --> :diamond

    u = getdata(dvis, :u)
    v = getdata(dvis, :v)
    uvdist = hypot.(u,v)
    amp = amplitude.(dvis.data)
    error = getdata(dvis, :error)
    #add data errorbars
    seriestype --> :scatter
    alpha --> 0.5
    yerr := error
    linecolor --> nothing
    label := "Data"
    uvdist, amp
end

@recipe function f(acc::ArrayConfiguration)
    xguide --> "u (λ)"
    yguide --> "v (λ)"
    markershape --> :circle

    u, v = getuv(acc)
    #add data errorbars
    seriestype --> :scatter
    linecolor --> nothing
    aspect_ratio --> :equal
    label := "Data"
    title --> "Frequency: $(acc.frequency/1e9) GHz"
    vcat(u,-u), vcat(v,-v)
end

@recipe function f(m::AbstractModel, dvis::EHTObservation{T,A}; datamarker=:circle, datacolor=:grey) where {T,A<:EHTVisibilityAmplitudeDatum}
    xguide --> "uv-distance (λ)"
    yguide --> "V (Jy)"
    markershape --> :diamond

    u = getdata(dvis, :u)
    v = getdata(dvis, :v)
    uvdist = hypot.(u,v)
    amp = amplitude.(dvis.data)
    error = getdata(dvis, :error)
    #add data errorbars
    @series begin
        seriestype := :scatter
        markershape := datamarker
        markercolor := datacolor
        alpha := 0.5
        yerr := error
        linecolor := nothing
        label := "Data"
        uvdist, amp
    end

    seriestype-->:scatter
    amod = amplitudes(m, u, v)
    labels --> "Model"
    uvdist, amod
end

export uvarea
function uvarea(d::EHTClosurePhaseDatum)
    u1, v1, u2, v2, u3, v3 = uvpositions(d)
    a = hypot(u1-u2, v1-v2)
    b = hypot(u2-u3, v2-v3)
    c = hypot(u3-u1, v3-v1)
    heron(a,b,c)
end

function heron(a,b,c)
    s = 0.5*(a+b+c)
    return sqrt(s*(s-a)*(s-b)*(s-c))
end

function uvarea(d::EHTLogClosureAmplitudeDatum)
    u1, v1, u2, v2, u3, v3, u4, v4 = uvpositions(d)
    a = hypot(u1-u2, v1-v2)
    b = hypot(u2-u3, v2-v3)
    c = hypot(u3-u4, v3-v4)
    d = hypot(u4-u1, v4-v1)
    h = hypot(u1-u3, v1-v3)
    return heron(a,b,h)+heron(c,d,h)
end

@recipe function f(dlca::EHTObservation{T,A}) where {T,A<:EHTLogClosureAmplitudeDatum}
    xguide --> "√(quadrangle area) (λ)"
    yguide --> "Log Clos. Amp."
    markershape --> :diamond
    area = sqrt.(uvarea.(dlca.data))
    phase = getdata(dlca, :amp)
    error = getdata(dlca, :error)
    #add data errorbars
    seriestype --> :scatter
    alpha := 0.5
    yerr := error
    label --> "Data"
    linecolor --> nothing
    area, phase
end

@recipe function f(m::AbstractModel, dlca::EHTObservation{T,A}; datamarker=:circle, datacolor=:grey) where {T,A<:EHTLogClosureAmplitudeDatum}
    xguide --> "√(quadrangle area) (λ)"
    yguide --> "Log Clos. Amp."
    markershape --> :diamond
    u1 = getdata(dlca, :u1)
    v1 = getdata(dlca, :v1)
    u2 = getdata(dlca, :u2)
    v2 = getdata(dlca, :v2)
    u3 = getdata(dlca, :u3)
    v3 = getdata(dlca, :v3)
    u4 = getdata(dlca, :u4)
    v4 = getdata(dlca, :v4)
    area = sqrt.(uvarea.(dlca.data))
    phase = getdata(dlca, :amp)
    error = getdata(dlca, :error)
    #add data errorbars
    @series begin
        seriestype := :scatter
        markershape := datamarker
        markercolor := datacolor
        alpha := 0.5
        yerr := error
        linecolor := nothing
        label := "Data"
        area, phase
    end

    seriestype-->:scatter
    amod = logclosure_amplitudes(m, u1, v1, u2, v2, u3, v3, u4, v4)
    labels --> "Model"
    area,amod
end

@recipe function f(dcp::EHTObservation{T,A}) where {T, A<:EHTClosurePhaseDatum}
    xguide --> "√(triangle area) (λ)"
    yguide --> "Phase (rad)"
    markershape --> :circle
    u1 = getdata(dcp, :u1)
    v1 = getdata(dcp, :v1)
    u2 = getdata(dcp, :u2)
    v2 = getdata(dcp, :v2)
    u3 = getdata(dcp, :u3)
    v3 = getdata(dcp, :v3)
    area = sqrt.(uvarea.(dcp.data))
    phase = getdata(dcp, :phase)
    error = getdata(dcp, :error)
    seriestype := :scatter
    alpha --> 0.5
    yerr := error
    linecolor --> nothing
    label := "Data"
    area, phase
end

@recipe function f(m::AbstractModel, dcp::EHTObservation{T,A}; datamarker=:circle, datacolor=:grey) where {T,A<:EHTClosurePhaseDatum}
    xguide --> "√(triangle area) (λ)"
    yguide --> "Phase (rad)"
    markershape --> :diamond
    u1 = getdata(dcp, :u1)
    v1 = getdata(dcp, :v1)
    u2 = getdata(dcp, :u2)
    v2 = getdata(dcp, :v2)
    u3 = getdata(dcp, :u3)
    v3 = getdata(dcp, :v3)
    area = sqrt.(uvarea.(dcp.data))
    phase = getdata(dcp, :phase)
    error = getdata(dcp, :error)
    #add data errorbars
    @series begin
        seriestype := :scatter
        markershape := datamarker
        markercolor := datacolor
        alpha := 0.5
        yerr := error
        linecolor := nothing
        label := "Data"
        area, phase
    end

    seriestype-->:scatter
    amod = closure_phases(m, u1, v1, u2, v2, u3, v3)
    labels --> "Model"
    area,amod
end

@userplot Residual

@recipe function f(h::Residual)
    if length(h.args) != 2 || !(typeof(h.args[1]) <: AbstractModel) ||
        !(typeof(h.args[2]) <: EHTObservation)
        error("Residual should be given a model and data product.  Got: $(typeof(h.args))")
    end
    m, damp = h.args

    uvdist, res = residuals(m, damp)
    chi2 = sum(abs2, res)
    xguide --> "uv-distance (λ)"
    yguide --> "Normalized Residual"
    markershape --> :circle
    linecolor --> nothing
    legend --> false

    title --> @sprintf "<χ²> = %.2f" chi2/length(damp)
    uvdist, res
end

function χ²(m, data)
    return sum(abs2, last(residuals(m, data)))
end

function residuals(m, damp::EHTObservation{T, A}) where {T, A<:EHTVisibilityAmplitudeDatum}
    amp = getdata(damp, :amp)
    u = getdata(damp, :u)
    v = getdata(damp, :v)

    mamp = amplitudes(m, u, v)
    res = (amp - mamp)./getdata(damp, :error)
    return hypot.(u, v), res
end

function residuals(m, dcp::EHTObservation{T, A}) where {T, A<:EHTClosurePhaseDatum}
    u1 = getdata(dcp, :u1)
    v1 = getdata(dcp, :v1)
    u2 = getdata(dcp, :u2)
    v2 = getdata(dcp, :v2)
    u3 = getdata(dcp, :u3)
    v3 = getdata(dcp, :v3)
    area = sqrt.(uvarea.(dcp.data))
    phase = getdata(dcp, :phase)
    error = getdata(dcp, :error)

    mphase = closure_phases(m, u1, v1, u2, v2, u3, v3)
    res = zeros(length(phase))
    for i in eachindex(res)
        s,c  = sincos(phase[i] - mphase[i])
        dθ = atan(s,c)
        res[i] = dθ/error[i]
    end
return area, res
end

function residuals(m, dlca::EHTObservation{T, A}) where {T, A<:EHTLogClosureAmplitudeDatum}
    u1 = getdata(dlca, :u1)
    v1 = getdata(dlca, :v1)
    u2 = getdata(dlca, :u2)
    v2 = getdata(dlca, :v2)
    u3 = getdata(dlca, :u3)
    v3 = getdata(dlca, :v3)
    u4 = getdata(dlca, :u4)
    v4 = getdata(dlca, :v4)
    area = sqrt.(uvarea.(dlca.data))
    phase = getdata(dlca, :amp)
    error = getdata(dlca, :error)

    mphase = logclosure_amplitudes(m, u1, v1, u2, v2, u3, v3, u4, v4)
    res = (phase- mphase)./error
    return area, res
end


text:
# Use baremodule to shave off a few KB from the serialized `.ji` file
baremodule Arb_jll
using Base
using Base: UUID
import JLLWrappers

JLLWrappers.@generate_main_file_header("Arb")
JLLWrappers.@generate_main_file("Arb", UUID("d9960996-1013-53c9-9ba4-74a4155039c3"))
end  # module Arb_jll


text:
import Pkg
# This file will import all necessary Packages for cipher-tools.

# Make sure user really wants to do this
println("This Programs installs packages required for Cipher-Tools. Press enter to begin the process.")
readline()
println("This could take some time...\n\n")

# Plotting tools
# Pkg.add("Plots")

# Dependencies for Plots
# Pkg.add("GR")
# Pkg.add("PGFPlotsX")
# Pkg.add("PlotlyJS"); Pkg.add("ORCA")
# Pkg.add("PyPlot")

Pkg.add("JSON")

max_stars_repo_name: Squalm/Cipher-Tools

max_issues_repo_name: Squalm/Cipher-Tools

max_forks_repo_name: Squalm/Cipher-Tools


text:
using Test

for testscen in 1:2
    valdir, scenario, use_permafrost, use_seaice = get_scenario(testscen)
    println(scenario)

    m = test_page_model()
    include("../src/components/RCPSSPScenario.jl")
    include("../src/components/CO2emissions.jl")

    scenario = addrcpsspscenario(m, scenario)
    co2emit = add_comp!(m, co2emissions)

    co2emit[:er_CO2emissionsgrowth] = scenario[:er_CO2emissionsgrowth]
    set_param!(m, :co2emissions, :e0_baselineCO2emissions, readpagedata(m, "data/e0_baselineCO2emissions.csv"))

    ##running Model
    run(m)

    emissions = m[:co2emissions,  :e_regionalCO2emissions]

    # Recorded data
    emissions_compare = readpagedata(m, "test/validationdata/$valdir/e_regionalCO2emissions.csv")

    @test emissions ≈ emissions_compare rtol = 1e-3
end


text:
"""
    constructtestfields_2layer(gr)

Constructs flow fields for a 2-layer problem with parameters such that
    q1 = Δψ1 + 25*(ψ2-ψ1),
    q2 = Δψ2 + 25/4*(ψ1-ψ2).
"""
function constructtestfields_2layer(gr)
  x, y = gridpoints(gr)
  k₀, l₀ = 2π/gr.Lx, 2π/gr.Ly # fundamental wavenumbers

  # a set of streafunctions ψ1 and ψ2, ...
  ψ1 = @. 1e-3 * ( 1/4*cos(2k₀*x)*cos(5l₀*y) + 1/3*cos(3k₀*x)*cos(3l₀*y) )
  ψ2 = @. 1e-3 * (     cos(3k₀*x)*cos(4l₀*y) + 1/2*cos(4k₀*x)*cos(2l₀*y) )

  # ... their corresponding PVs q1, q2,
  q1 = @. 1e-3 * ( 1/6 *(   75cos(4k₀*x)*cos(2l₀*y) +   2cos(3k₀*x)*( -43cos(3l₀*y) + 75cos(4l₀*y) ) - 81cos(2k₀*x)*cos(5l₀*y) ) )
  q2 = @. 1e-3 * ( 1/48*( -630cos(4k₀*x)*cos(2l₀*y) + 100cos(3k₀*x)*(    cos(3l₀*y) - 15cos(4l₀*y) ) + 75cos(2k₀*x)*cos(5l₀*y) ) )

  # ... and various derived fields, e.g., ∂ψ1/∂x,
  ψ1x = @. 1e-3 * (     -k₀/2*sin(2k₀*x)*cos(5l₀*y) -       k₀*sin(3k₀*x)*cos(3l₀*y) )
  ψ2x = @. 1e-3 * (      -3k₀*sin(3k₀*x)*cos(4l₀*y) -      2k₀*sin(4k₀*x)*cos(2l₀*y) )
  Δψ2 = @. 1e-3 * ( -25*k₀*l₀*cos(3k₀*x)*cos(4l₀*y) - 10*k₀*l₀*cos(4k₀*x)*cos(2l₀*y) )

  q1x = @. 1e-3 * ( 1/6 *( -4k₀*75sin(4k₀*x)*cos(2l₀*y) - 3k₀*2sin(3k₀*x)*( -43cos(3l₀*y) + 75cos(4l₀*y) ) + 2k₀*81sin(2k₀*x)*cos(5l₀*y) ) )
  q2x = @. 1e-3 * ( 1/48*( 4k₀*630sin(4k₀*x)*cos(2l₀*y) - 3k₀*100sin(3k₀*x)*(    cos(3l₀*y) - 15cos(4l₀*y) ) - 2k₀*75sin(2k₀*x)*cos(5l₀*y) ) )

  Δq1 = @. 1e-3 * (k₀*l₀)*( 1/6 *( -20* 75cos(4k₀*x)*cos(2l₀*y) +   2cos(3k₀*x)*( +18*43cos(3l₀*y) - 25*75cos(4l₀*y) ) +29*81cos(2k₀*x)*cos(5l₀*y) ) )
  Δq2 = @. 1e-3 * (k₀*l₀)*( 1/48*( +20*630cos(4k₀*x)*cos(2l₀*y) + 100cos(3k₀*x)*(    -18cos(3l₀*y) + 25*15cos(4l₀*y) ) -29*75cos(2k₀*x)*cos(5l₀*y) ) )

  return ψ1, ψ2, q1, q2, ψ1x, ψ2x, q1x, q2x, Δψ2, Δq1, Δq2
end

"""
    constructtestfields_3layer(gr)

Constructs flow fields for a 3-layer problem with parameters such that
    q1 = Δψ1 + 20ψ2 - 20ψ1,
    q2 = Δψ2 + 20ψ1 - 44ψ2 + 24ψ3,
    q2 = Δψ2 + 12ψ2 - 12ψ3.
"""
function constructtestfields_3layer(gr)
  x, y = gridpoints(gr)
  k₀, l₀ = 2π/gr.Lx, 2π/gr.Ly # fundamental wavenumbers

  # a set of streafunctions ψ1, ψ2, ψ3, ...
  ψ1 = @. 1e-3 * ( 1/4*cos(2k₀*x)*cos(5l₀*y) + 1/3*cos(3k₀*x)*cos(3l₀*y) )
  ψ2 = @. 1e-3 * (     cos(3k₀*x)*cos(4l₀*y) + 1/2*cos(4k₀*x)*cos(2l₀*y) )
  ψ3 = @. 1e-3 * (     cos(1k₀*x)*cos(3l₀*y) + 1/2*cos(2k₀*x)*cos(2l₀*y) )

  Δψ1 = @. -1e-3 * ( 1/4*((2k₀)^2+(5l₀)^2)*cos(2k₀*x)*cos(5l₀*y) + 1/3*((3k₀)^2+(3l₀)^2)*cos(3k₀*x)*cos(3l₀*y) )
  Δψ2 = @. -1e-3 * (     ((3k₀)^2+(4l₀)^2)*cos(3k₀*x)*cos(4l₀*y) + 1/2*((4k₀)^2+(2l₀)^2)*cos(4k₀*x)*cos(2l₀*y) )
  Δψ3 = @. -1e-3 * (     ((1k₀)^2+(3l₀)^2)*cos(1k₀*x)*cos(3l₀*y) + 1/2*((2k₀)^2+(2l₀)^2)*cos(2k₀*x)*cos(2l₀*y) )

  # ... their corresponding PVs q1, q2, q3,
  q1 = @. Δψ1 + 20ψ2 - 20ψ1
  q2 = @. Δψ2 + 20ψ1 - 44ψ2 + 24ψ3
  q3 = @. Δψ3        + 12ψ2 - 12ψ3

  return ψ1, ψ2, ψ3, q1, q2, q3
end

"""
    test_pvtofromstreamfunction_2layer(dev)

Tests the pvfromstreamfunction function that gives qh from ψh. To do so, it
constructs a 2-layer problem with parameters such that
    q1 = Δψ1 + 25*(ψ2-ψ1), q2 = Δψ2 + 25/4*(ψ1-ψ2).
Then given ψ1 and ψ2 it tests if pvfromstreamfunction gives the expected
q1 and q2. Similarly, that streamfunctionfrompv gives ψ1 and ψ2 from q1 and q2.
"""
function test_pvtofromstreamfunction_2layer(dev::Device=CPU())
  n, L = 128, 2π
  gr = TwoDGrid(dev, n, L)

  nlayers = 2      # these choice of parameters give the
  f₀, g = 1, 1     # desired PV-streamfunction relations
  H = [0.2, 0.8]   # q1 = Δψ1 + 25*(ψ2-ψ1), and
  ρ = [4.0, 5.0]   # q2 = Δψ2 + 25/4*(ψ1-ψ2).

  prob = MultiLayerQG.Problem(nlayers, dev; nx=n, Lx=L, f₀=f₀, g=g, H=H, ρ=ρ)
  sol, cl, pr, vs, gr = prob.sol, prob.clock, prob.params, prob.vars, prob.grid

  ψ1, ψ2, q1, q2, ψ1x, ψ2x, q1x, q2x, Δψ2, Δq1, Δq2 = constructtestfields_2layer(gr)

  vs.ψh[:, :, 1] .= rfft(ψ1)
  vs.ψh[:, :, 2] .= rfft(ψ2)

  MultiLayerQG.pvfromstreamfunction!(vs.qh, vs.ψh, pr, gr)
  MultiLayerQG.invtransform!(vs.q, vs.qh, pr)

  vs.qh[:, :, 1] .= rfft(q1)
  vs.qh[:, :, 2] .= rfft(q2)

  MultiLayerQG.streamfunctionfrompv!(vs.ψh, vs.qh, pr, gr)
  MultiLayerQG.invtransform!(vs.ψ, vs.ψh, pr)

  return isapprox(q1, vs.q[:, :, 1], rtol=rtol_multilayerqg) && isapprox(q2, vs.q[:, :, 2], rtol=rtol_multilayerqg) && isapprox(ψ1, vs.ψ[:, :, 1], rtol=rtol_multilayerqg) && isapprox(ψ2, vs.ψ[:, :, 2], rtol=rtol_multilayerqg)
end

"""
    test_pvtofromstreamfunction_3layer()

Tests the pvfromstreamfunction function that gives qh from ψh. To do so, it
constructs a 3-layer problem with parameters such that
  q1 = Δψ1 + 20ψ2 - 20ψ1, q2 = Δψ2 + 20ψ1 - 44ψ2 + 24ψ3, q3 = Δψ3 + 12ψ2 - 12ψ3.
Then given ψ1, ψ2, and ψ2 it tests if pvfromstreamfunction gives the expected
q1, q2, and q3. Similarly, that streamfunctionfrompv gives ψ1, ψ2, and ψ2 from
q1, q2, and q3.
"""
function test_pvtofromstreamfunction_3layer(dev::Device=CPU())
  n, L = 128, 2π
  gr = TwoDGrid(dev, n, L)

  nlayers = 3            # these choice of parameters give the
  f₀, g = 1, 1           # desired PV-streamfunction relations
  H = [0.25, 0.25, 0.5]  # q1 = Δψ1 + 20ψ2 - 20ψ1,
  ρ = [4.0, 5.0, 6.0]    # q2 = Δψ2 + 20ψ1 - 44ψ2 + 24ψ3,
                         # q3 = Δψ3        + 12ψ2 - 12ψ3.

  prob = MultiLayerQG.Problem(nlayers, dev; nx=n, Lx=L, f₀=f₀, g=g, H=H, ρ=ρ)
  sol, cl, pr, vs, gr = prob.sol, prob.clock, prob.params, prob.vars, prob.grid

  ψ1, ψ2, ψ3, q1, q2, q3 = constructtestfields_3layer(gr)

  vs.ψh[:, :, 1] .= rfft(ψ1)
  vs.ψh[:, :, 2] .= rfft(ψ2)
  vs.ψh[:, :, 3] .= rfft(ψ3)

  MultiLayerQG.pvfromstreamfunction!(vs.qh, vs.ψh, pr, gr)
  MultiLayerQG.invtransform!(vs.q, vs.qh, pr)

  vs.qh[:, :, 1] .= rfft(q1)
  vs.qh[:, :, 2] .= rfft(q2)
  vs.qh[:, :, 3] .= rfft(q3)

  MultiLayerQG.streamfunctionfrompv!(vs.ψh, vs.qh, pr, gr)
  MultiLayerQG.invtransform!(vs.ψ, vs.ψh, pr)

  return isapprox(q1, vs.q[:, :, 1], rtol=rtol_multilayerqg) && isapprox(q2, vs.q[:, :, 2], rtol=rtol_multilayerqg) && isapprox(q3, vs.q[:, :, 3], rtol=rtol_multilayerqg) && isapprox(ψ1, vs.ψ[:, :, 1], rtol=rtol_multilayerqg) && isapprox(ψ2, vs.ψ[:, :, 2], rtol=rtol_multilayerqg) && isapprox(ψ3, vs.ψ[:, :, 3], rtol=rtol_multilayerqg)
end

"""
    test_mqg_nonlinearadvection(dt, stepper, dev; kwargs...)

Tests the advection term by timestepping a test problem with timestep dt and
timestepper identified by the string stepper. The test 2-layer problem is
derived by picking a solution q1f, q2f (with streamfunctions ψ1f, ψ2f) for which
the advection terms J(ψn, qn) are non-zero. Next, a forcing Ff is derived such
that a solution to the problem forced by this Ff is then qf.
(This solution may not be realized, at least at long times, if it is unstable.)
"""
function test_mqg_nonlinearadvection(dt, stepper, dev::Device=CPU(); n=128, L=2π, nlayers=2, μ=0.0, ν=0.0, nν=1)

  A = ArrayType(dev)

  tf = 0.5
  nt = round(Int, tf/dt)

  nx, ny = 64, 66
  Lx, Ly = 2π, 2π
  gr = TwoDGrid(dev, nx, Lx, ny, Ly)

  x, y = gridpoints(gr)
  k₀, l₀ = 2π/gr.Lx, 2π/gr.Ly # fundamental wavenumbers

  nlayers = 2       # these choice of parameters give the
  f₀, g = 1, 1      # desired PV-streamfunction relations
  H = [0.2, 0.8]    # q1 = Δψ1 + 25*(ψ2-ψ1), and
  ρ = [4.0, 5.0]    # q2 = Δψ2 + 25/4*(ψ1-ψ2).

  β = 0.35

  U1, U2 = 0.1, 0.05
  u1 = @. 0.5sech(gr.y/(Ly/15))^2; u1 = A(reshape(u1, (1, gr.ny)))
  u2 = @. 0.02cos(3l₀*gr.y);       u2 = A(reshape(u2, (1, gr.ny)))
  uyy1 = real.(ifft( -gr.l.^2 .* fft(u1) ))
  uyy2 = real.(ifft( -gr.l.^2 .* fft(u2) ))

  U = zeros(ny, nlayers)
  CUDA.@allowscalar U[:, 1] = u1 .+ U1
  CUDA.@allowscalar U[:, 2] = u2 .+ U2

  μ, ν, nν = 0.1, 0.05, 1

  η0, σx, σy = 1.0, Lx/25, Ly/20
   η = @. η0*exp( -(x+Lx/8)^2/(2σx^2) - (y-Ly/8)^2/(2σy^2) )
  ηx = @. -(x + Lx/8)/(σx^2) * η

  ψ1, ψ2, q1, q2, ψ1x, ψ2x, q1x, q2x, Δψ2, Δq1, Δq2 = constructtestfields_2layer(gr)

  Ff1 = FourierFlows.jacobian(ψ1, q1, gr)     + @. (β - uyy1 -   25*(U2+u2-U1-u1) )*ψ1x + (U1+u1)*q1x - ν*Δq1
  Ff2 = FourierFlows.jacobian(ψ2, q2 + η, gr) + @. (β - uyy2 - 25/4*(U1+u1-U2-u2) )*ψ2x + (U2+u2)*(q2x + ηx) + μ*Δψ2 - ν*Δq2

  T = eltype(gr)

  Ff = zeros(dev, T, (gr.nx, gr.ny, nlayers))
  @views Ff[:, :, 1] = Ff1
  @views Ff[:, :, 2] = Ff2

  Ffh = zeros(dev, Complex{T}, (gr.nkr, gr.nl, nlayers))
  @views Ffh[:, :, 1] = rfft(Ff1)
  @views Ffh[:, :, 2] = rfft(Ff2)

  function calcFq!(Fqh, sol, t, cl, v, p, g)
    Fqh .= Ffh
    nothing
  end

  prob = MultiLayerQG.Problem(nlayers, dev; nx=nx, ny=ny, Lx=Lx, Ly=Ly,
   f₀=f₀, g=g, H=H, ρ=ρ, U=U, eta=η, β=β, μ=μ, ν=ν, nν=nν, calcFq=calcFq!, stepper=stepper, dt=dt)
  sol, cl, pr, vs, gr = prob.sol, prob.clock, prob.params, prob.vars, prob.grid

  qf = zeros(dev, T, (gr.nx, gr.ny, nlayers))
  @views qf[:, :, 1] = q1
  @views qf[:, :, 2] = q2

  ψf = zeros(dev, T, (gr.nx, gr.ny, nlayers))
  @views ψf[:, :, 1] = ψ1
  @views ψf[:, :, 2] = ψ2

  MultiLayerQG.set_q!(prob, qf)

  stepforward!(prob, nt)
  MultiLayerQG.updatevars!(prob)

  return isapprox(vs.q, qf, rtol=rtol_multilayerqg) && isapprox(vs.ψ, ψf, rtol=rtol_multilayerqg)
end

"""
    test_mqg_linearadvection(dt, stepper, dev; kwargs...)

Tests the advection term of the linearized equations by timestepping a test
problem with timestep dt and timestepper identified by the string stepper.
The test 2-layer problem is derived by picking a solution q1f, q2f (with
streamfunctions ψ1f, ψ2f) for which the advection terms J(ψn, qn) are non-zero.
Next, a forcing Ff is derived such that a solution to the problem forced by this
Ff is then qf. (This solution may not be realized, at least at long times, if it
is unstable.)
"""
function test_mqg_linearadvection(dt, stepper, dev::Device=CPU(); n=128, L=2π, nlayers=2, μ=0.0, ν=0.0, nν=1)

  A = ArrayType(dev)

  tf = 0.5
  nt = round(Int, tf/dt)

  nx, ny = 64, 66
  Lx, Ly = 2π, 2π
  gr = TwoDGrid(dev, nx, Lx, ny, Ly)

  x, y = gridpoints(gr)
  k₀, l₀ = 2π/gr.Lx, 2π/gr.Ly # fundamental wavenumbers

  nlayers = 2       # these choice of parameters give the
  f₀, g = 1, 1      # desired PV-streamfunction relations
  H = [0.2, 0.8]    # q1 = Δψ1 + 25*(ψ2-ψ1), and
  ρ = [4.0, 5.0]    # q2 = Δψ2 + 25/4*(ψ1-ψ2).

  β = 0.35

  U1, U2 = 0.1, 0.05
  u1 = @. 0.5sech(gr.y/(Ly/15))^2; u1 = A(reshape(u1, (1, gr.ny)))
  u2 = @. 0.02cos(3l₀*gr.y);       u2 = A(reshape(u2, (1, gr.ny)))
  uyy1 = real.(ifft( -gr.l.^2 .* fft(u1) ))
  uyy2 = real.(ifft( -gr.l.^2 .* fft(u2) ))

  U = zeros(ny, nlayers)
  CUDA.@allowscalar U[:, 1] = u1 .+ U1
  CUDA.@allowscalar U[:, 2] = u2 .+ U2

  μ, ν, nν = 0.1, 0.05, 1

  η0, σx, σy = 1.0, Lx/25, Ly/20
  η = @. η0*exp( -(x + Lx/8)^2/(2σx^2) -(y-Ly/8)^2/(2σy^2) )
  ηx = @. -(x + Lx/8)/(σx^2) * η

  ψ1, ψ2, q1, q2, ψ1x, ψ2x, q1x, q2x, Δψ2, Δq1, Δq2 = constructtestfields_2layer(gr)

  Ff1 = (β .- uyy1 .- 25*(U2.+u2.-U1.-u1) ).*ψ1x + (U1.+u1).*q1x - ν*Δq1
  Ff2 = FourierFlows.jacobian(ψ2, η, gr) + (β .- uyy2 .- 25/4*(U1.+u1.-U2.-u2) ).*ψ2x + (U2.+u2).*(q2x + ηx) + μ*Δψ2 - ν*Δq2

  T = eltype(gr)

  Ff = zeros(dev, T, (gr.nx, gr.ny, nlayers))
  @views Ff[:, :, 1] = Ff1
  @views Ff[:, :, 2] = Ff2

  Ffh = zeros(dev, Complex{T}, (gr.nkr, gr.nl, nlayers))
  @views Ffh[:, :, 1] = rfft(Ff1)
  @views Ffh[:, :, 2] = rfft(Ff2)

  function calcFq!(Fqh, sol, t, cl, v, p, g)
    Fqh .= Ffh
    nothing
  end

  prob = MultiLayerQG.Problem(nlayers, dev; nx=nx, ny=ny, Lx=Lx, Ly=Ly, f₀=f₀, g=g, H=H, ρ=ρ, U=U, eta=η, β=β, μ=μ, ν=ν, nν=nν, calcFq=calcFq!, stepper=stepper, dt=dt, linear=true)
  sol, cl, pr, vs, gr = prob.sol, prob.clock, prob.params, prob.vars, prob.grid

  qf = zeros(dev, T, (gr.nx, gr.ny, nlayers))
  @views qf[:, :, 1] = q1
  @views qf[:, :, 2] = q2

  ψf = zeros(dev, T, (gr.nx, gr.ny, nlayers))
  @views ψf[:, :, 1] = ψ1
  @views ψf[:, :, 2] = ψ2

  MultiLayerQG.set_q!(prob, qf)

  stepforward!(prob, nt)
  MultiLayerQG.updatevars!(prob)

  return isapprox(vs.q, qf, rtol=rtol_multilayerqg) && isapprox(vs.ψ, ψf, rtol=rtol_multilayerqg)
end

"""
    test_mqg_energies(dev; kwargs...)

Tests the kinetic (KE) and potential (PE) energies function by constructing a
2-layer problem and initializing it with a flow field whose KE and PE are known.
"""
function test_mqg_energies(dev::Device=CPU(); dt=0.001, stepper="ForwardEuler", n=128, L=2π, nlayers=2, μ=0.0, ν=0.0, nν=1)
  nx, ny = 64, 66
  Lx, Ly = 2π, 2π
  gr = TwoDGrid(dev, nx, Lx, ny, Ly)

  x, y = gridpoints(gr)
  k₀, l₀ = 2π/gr.Lx, 2π/gr.Ly # fundamental wavenumbers

  nlayers = 2       # these choice of parameters give the
  f₀, g = 1, 1      # desired PV-streamfunction relations
  H = [0.2, 0.8]    # q1 = Δψ1 + 25*(ψ2-ψ1), and
  ρ = [4.0, 5.0]    # q2 = Δψ2 + 25/4*(ψ1-ψ2).

  prob = MultiLayerQG.Problem(nlayers, dev; nx=nx, ny=ny, Lx=Lx, Ly=Ly, f₀=f₀, g=g, H=H, ρ=ρ)
  sol, cl, pr, vs, gr = prob.sol, prob.clock, prob.params, prob.vars, prob.grid

  ψ1, ψ2, q1, q2, ψ1x, ψ2x, q1x, q2x, Δψ2, Δq1, Δq2 = constructtestfields_2layer(gr)

  qf = zeros(gr.nx, gr.ny, nlayers)
  CUDA.@allowscalar @views qf[:, :, 1] .= q1
  CUDA.@allowscalar @views qf[:, :, 2] .= q2

  MultiLayerQG.set_q!(prob, qf)

  KE, PE = MultiLayerQG.energies(prob)

  return isapprox(KE[1], 61/640*1e-6, rtol=rtol_multilayerqg) && isapprox(KE[2], 3*1e-6, rtol=rtol_multilayerqg) && isapprox(PE[1], 1025/1152*1e-6, rtol=rtol_multilayerqg) && MultiLayerQG.addforcing!(prob.timestepper.RHS₁, sol, cl.t, cl, vs, pr, gr)==nothing
end

function test_mqg_energysinglelayer(dev::Device=CPU(); dt=0.001, stepper="ForwardEuler", nlayers=1, μ=0.0, ν=0.0, nν=1)
  nx, Lx  = 64, 2π
  ny, Ly  = 64, 3π
  gr = TwoDGrid(dev, nx, Lx, ny, Ly)

  x, y = gridpoints(gr)
  k₀, l₀ = 2π/gr.Lx, 2π/gr.Ly # fundamental wavenumbers

  energy_calc = 29/9

  ψ0 = @. sin(2k₀*x)*cos(2l₀*y) + 2sin(k₀*x)*cos(3l₀*y)
  q0 = @. -((2k₀)^2+(2l₀)^2)*sin(2k₀*x)*cos(2l₀*y) - (k₀^2+(3l₀)^2)*2sin(k₀*x)*cos(3l₀*y)

  prob = MultiLayerQG.Problem(nlayers, dev; nx=nx, Lx=Lx, ny=ny, Ly=Ly, stepper=stepper, U=zeros(ny))

  MultiLayerQG.set_q!(prob, reshape(q0, (nx, ny, nlayers)))

  energyq0 = MultiLayerQG.energies(prob)

  return isapprox(energyq0, energy_calc, rtol=rtol_multilayerqg)
end

"""
    test_mqg_fluxes(dt, stepper; kwargs...)

Tests the lateral and vertical eddy fluxes by constructing a 2-layer problem and
initializing it with a flow field whose fluxes are known.
"""
function test_mqg_fluxes(dev::Device=CPU(); dt=0.001, stepper="ForwardEuler", n=128, L=2π, nlayers=2, μ=0.0, ν=0.0, nν=1)
  nx, ny = 128, 126
  Lx, Ly = 2π, 2π
  gr = TwoDGrid(dev, nx, Lx, ny, Ly)

  x, y = gridpoints(gr)
  k₀, l₀ = 2π/gr.Lx, 2π/gr.Ly # fundamental wavenumbers

  nlayers = 2       # these choice of parameters give the
  f₀, g = 1, 1      # desired PV-streamfunction relations
  H = [0.2, 0.8]    # q1 = Δψ1 + 25*(ψ2-ψ1), and
  ρ = [4.0, 5.0]    # q2 = Δψ2 + 25/4*(ψ1-ψ2).
  U = zeros(ny, nlayers)
  U[:, 1] = @. sech(gr.y / 0.2)^2

  prob = MultiLayerQG.Problem(nlayers, dev; nx=nx, ny=ny, Lx=Lx, Ly=Ly, f₀=f₀, g=g, H=H, ρ=ρ, U=U)
  sol, cl, pr, vs, gr = prob.sol, prob.clock, prob.params, prob.vars, prob.grid

  ψ1 = @. cos(k₀*x) * cos(l₀*y) + sin(k₀*x)
  ψ2 = @. cos(k₀*x + π/10) * cos(l₀*y)
  ψ = zeros(gr.nx, gr.ny, nlayers)
  CUDA.@allowscalar @views ψ[:, :, 1] .= ψ1
  CUDA.@allowscalar @views ψ[:, :, 2] .= ψ2
  MultiLayerQG.set_ψ!(prob, ψ)
  lateralfluxes, verticalfluxes = MultiLayerQG.fluxes(prob)

  return CUDA.@allowscalar isapprox(lateralfluxes[1], 0.00626267, rtol=1e-6) && CUDA.@allowscalar isapprox(lateralfluxes[2], 0, atol=1e-12) && CUDA.@allowscalar isapprox(verticalfluxes[1], -0.196539, rtol=1e-6)
end

"""
    test_mqg_fluxessinglelayer(dt, stepper; kwargs...)

Tests the lateral eddy fluxes by constructing a 1-layer problem and initializing
it with a flow field whose fluxes are known.
"""
function test_mqg_fluxessinglelayer(dev::Device=CPU(); dt=0.001, stepper="ForwardEuler", n=128, L=2π, μ=0.0, ν=0.0, nν=1)
  nlayers = 1

  nx, ny = 128, 126
  Lx, Ly = 2π, 2π
  gr = TwoDGrid(dev, nx, Lx, ny, Ly)

  x, y = gridpoints(gr)
  k₀, l₀ = 2π/gr.Lx, 2π/gr.Ly # fundamental wavenumbers

  U = zeros(ny, nlayers)
  U = @. sech(gr.y / 0.2)^2

  prob = MultiLayerQG.Problem(nlayers, dev; nx=nx, ny=ny, Lx=Lx, Ly=Ly, U=U)
  sol, cl, pr, vs, gr = prob.sol, prob.clock, prob.params, prob.vars, prob.grid

  ψ = @. cos(k₀*x) * cos(l₀*y) + sin(k₀*x)
  MultiLayerQG.set_ψ!(prob, ψ)
  lateralfluxes = MultiLayerQG.fluxes(prob)

  return CUDA.@allowscalar isapprox(lateralfluxes[1], 0.0313134, atol=1e-7)
end

"""
    test_setqsetψ(dt, stepper; kwargs...)

Tests the set_q!() and set_ψ!() functions that initialize sol with a flow with
given `q` or `ψ` respectively.
"""
function test_mqg_setqsetψ(dev::Device=CPU(); dt=0.001, stepper="ForwardEuler", n=64, L=2π, nlayers=2, μ=0.0, ν=0.0, nν=1)
  nx, ny = 32, 34
  L = 2π
  gr = TwoDGrid(dev, nx, L, ny, L)

  x, y = gridpoints(gr)
  k₀, l₀ = 2π/gr.Lx, 2π/gr.Ly # fundamental wavenumbers

  nlayers = 2       # these choice of parameters give the
  f₀, g = 1, 1      # desired PV-streamfunction relations
  H = [0.2, 0.8]    # q1 = Δψ1 + 25*(ψ2-ψ1), and
  ρ = [4.0, 5.0]    # q2 = Δψ2 + 25/4*(ψ1-ψ2).

  prob = MultiLayerQG.Problem(nlayers, dev; nx=nx, ny=ny, Lx=L, f₀=f₀, g=g, H=H, ρ=ρ)
  sol, cl, pr, vs, gr = prob.sol, prob.clock, prob.params, prob.vars, prob.grid

  T = eltype(gr)

  f1 = @. 2cos(k₀*x)*cos(l₀*y)
  f2 = @.  cos(k₀*x+π/10)*cos(2l₀*y)
  f = zeros(dev, T, (gr.nx, gr.ny, nlayers))
  f[:, :, 1] .= f1
  f[:, :, 2] .= f2

  ψtest = zeros(dev, T, size(f))
  MultiLayerQG.set_ψ!(prob, f)
  @. vs.qh = sol
  MultiLayerQG.streamfunctionfrompv!(vs.ψh, vs.qh, pr, gr)
  MultiLayerQG.invtransform!(ψtest, vs.ψh, pr)

  qtest = zeros(dev, T, size(f))
  MultiLayerQG.set_q!(prob, f)
  @. vs.qh = sol
  MultiLayerQG.invtransform!(qtest, vs.qh, pr)

  return isapprox(ψtest, f, rtol=rtol_multilayerqg) && isapprox(qtest, f, rtol=rtol_multilayerqg)
end

"""
    test_paramsconstructor(; kwargs...)

Tests that `Params` constructor works with both mean flow `U` being a floats
(i.e., constant `U` in each layer) or vectors (i.e., `U(y)` in each layer).
"""
function test_mqg_paramsconstructor(dev::Device=CPU(); dt=0.001, stepper="ForwardEuler")
  nx, ny = 32, 34
  L = 2π
  gr = TwoDGrid(dev, nx, L, ny, L)

  nlayers = 2       # these choice of parameters give the
  f₀, g = 1, 1      # desired PV-streamfunction relations
  H = [0.2, 0.8]    # q1 = Δψ1 + 25*(ψ2-ψ1), and
  ρ = [4.0, 5.0]    # q2 = Δψ2 + 25/4*(ψ1-ψ2).

  U1, U2 = 0.1, 0.05

  T = eltype(gr)

  Uvectors = zeros(dev, T, (ny, nlayers))
  Uvectors[:, 1] .= U1
  Uvectors[:, 2] .= U2

  Ufloats = zeros(dev, T, nlayers)
  CUDA.@allowscalar Ufloats[1] = U1
  CUDA.@allowscalar Ufloats[2] = U2

  probUvectors = MultiLayerQG.Problem(nlayers, dev; nx=nx, ny=ny, Lx=L, f₀=f₀, g=g, H=H, ρ=ρ, U=Uvectors)
  probUfloats = MultiLayerQG.Problem(nlayers, dev; nx=nx, ny=ny, Lx=L, f₀=f₀, g=g, H=H, ρ=ρ, U=Ufloats)

  return isapprox(probUfloats.params.U, probUvectors.params.U, rtol=rtol_multilayerqg)
end

function test_mqg_problemtype(dev, T)
  prob1 = MultiLayerQG.Problem(1, dev; T=T)
  prob2 = MultiLayerQG.Problem(2, dev; T=T)

  A = ArrayType(dev)

  return (typeof(prob1.sol)<:A{Complex{T}, 3} && typeof(prob1.grid.Lx)==T && typeof(prob1.vars.u)<:A{T, 3}) && (typeof(prob2.sol)<:A{Complex{T}, 3} && typeof(prob2.grid.Lx)==T && typeof(prob2.vars.u)<:A{T, 3})
end

"""
    test_mqg_rossbywave(; kwargs...)

Evolves a Rossby wave on a beta plane with an imposed zonal flow U and compares
with the analytic solution.
"""
function test_mqg_rossbywave(stepper, dt, nsteps, dev::Device=CPU())
  nlayers = 1
       nx = 64
       Lx = 2π
        β = 2
        U = 0.5

  prob = MultiLayerQG.Problem(nlayers, dev; nx=nx, Lx=Lx, U=U, β=β, stepper=stepper, dt=dt)
  sol, cl, v, p, g = prob.sol, prob.clock, prob.vars, prob.params, prob.grid

  x, y = gridpoints(g)

  # the Rossby wave initial condition
   ampl = 1e-2
  kwave, lwave = 3 * 2π/g.Lx, 2 * 2π/g.Ly
      ω = kwave * ( U - p.β / (kwave^2 + lwave^2) ) # Doppler-shifted Rossby frequency
     q0 = @. ampl * cos(kwave * x) * cos(lwave * y)
     ψ0 = @. - q0 / (kwave^2 + lwave^2)

  MultiLayerQG.set_q!(prob, q0)

  stepforward!(prob, nsteps)
  dealias!(sol, g)
  MultiLayerQG.updatevars!(prob)

  q_theory = @. ampl * cos(kwave * (x - ω / kwave * cl.t)) * cos(lwave * y)

  return isapprox(q_theory, v.q, rtol=g.nx*g.ny*nsteps*1e-12)
end

function test_numberoflayers(dev::Device=CPU())
  prob_nlayers1 = MultiLayerQG.Problem(1, dev)
  prob_nlayers2 = MultiLayerQG.Problem(2, dev)

  return MultiLayerQG.numberoflayers(prob_nlayers1)==1 && MultiLayerQG.numberoflayers(prob_nlayers2)==2
end

function test_mqg_stochasticforcedproblemconstructor(dev::Device=CPU())

  function calcFq!(Fqh, sol, t, clock, vars, params, grid)
    Fqh .= Ffh
    return nothing
  end

  prob = MultiLayerQG.Problem(2, dev; calcFq=calcFq!, stochastic=true)

  return typeof(prob.vars.prevsol)==typeof(prob.sol)
end