text:
using FromFile
@from "Core.jl" import CONST_TYPE, Node, copyNode, Options
@from "EquationUtils.jl" import countNodes
@from "CustomSymbolicUtilsSimplification.jl" import custom_simplify
@from "InterfaceSymbolicUtils.jl" import node_to_symbolic_safe, symbolic_to_node
@from "CheckConstraints.jl" import check_constraints
@from "Utils.jl" import isbad, isgood

# Simplify tree
function combineOperators(tree::Node, options::Options)::Node
    # NOTE: (const (+*-) const) already accounted for. Call simplifyTree before.
    # ((const + var) + const) => (const + var)
    # ((const * var) * const) => (const * var)
    # ((const - var) - const) => (const - var)
    # (want to add anything commutative!)
    # TODO - need to combine plus/sub if they are both there.
    if tree.degree == 0
        return tree
    elseif tree.degree == 1
        tree.l = combineOperators(tree.l, options)
    elseif tree.degree == 2
        tree.l = combineOperators(tree.l, options)
        tree.r = combineOperators(tree.r, options)
    end

    top_level_constant = tree.degree == 2 && (tree.l.constant || tree.r.constant)
    if tree.degree == 2 && (options.binops[tree.op] == (*) || options.binops[tree.op] == (+)) && top_level_constant
        op = tree.op
        # Put the constant in r. Need to assume var in left for simplification assumption.
        if tree.l.constant
            tmp = tree.r
            tree.r = tree.l
            tree.l = tmp
        end
        topconstant = tree.r.val
        # Simplify down first
        below = tree.l
        if below.degree == 2 && below.op == op
            if below.l.constant
                tree = below
                tree.l.val = options.binops[op](tree.l.val, topconstant)
            elseif below.r.constant
                tree = below
                tree.r.val = options.binops[op](tree.r.val, topconstant)
            end
        end
    end

    if tree.degree == 2 && options.binops[tree.op] == (-) && top_level_constant
        # Currently just simplifies subtraction. (can't assume both plus and sub are operators)
        # Not commutative, so use different op.
        if tree.l.constant
            if tree.r.degree == 2 && options.binops[tree.r.op] == (-)
                if tree.r.l.constant
                    #(const - (const - var)) => (var - const)
                    l = tree.l
                    r = tree.r
                    simplified_const = -(l.val - r.l.val) #neg(sub(l.val, r.l.val))
                    tree.l = tree.r.r
                    tree.r = l
                    tree.r.val = simplified_const
                elseif tree.r.r.constant
                    #(const - (var - const)) => (const - var)
                    l = tree.l
                    r = tree.r
                    simplified_const = l.val + r.r.val #plus(l.val, r.r.val)
                    tree.r = tree.r.l
                    tree.l.val = simplified_const
                end
            end
        else #tree.r.constant is true
            if tree.l.degree == 2 && options.binops[tree.l.op] == (-)
                if tree.l.l.constant
                    #((const - var) - const) => (const - var)
                    l = tree.l
                    r = tree.r
                    simplified_const = l.l.val - r.val#sub(l.l.val, r.val)
                    tree.r = tree.l.r
                    tree.l = r
                    tree.l.val = simplified_const
                elseif tree.l.r.constant
                    #((var - const) - const) => (var - const)
                    l = tree.l
                    r = tree.r
                    simplified_const = r.val + l.r.val #plus(r.val, l.r.val)
                    tree.l = tree.l.l
                    tree.r.val = simplified_const
                end
            end
        end
    end
    return tree
end

# Simplify tree
function simplifyTree(tree::Node, options::Options)::Node
    if tree.degree == 1
        tree.l = simplifyTree(tree.l, options)
        l = tree.l.val
        if tree.l.degree == 0 && tree.l.constant && isgood(l)
            out = options.unaops[tree.op](l)
            if isbad(out)
                return tree
            end
            return Node(convert(CONST_TYPE, out))
        end
    elseif tree.degree == 2
        tree.l = simplifyTree(tree.l, options)
        tree.r = simplifyTree(tree.r, options)
        constantsBelow = (
             tree.l.degree == 0 && tree.l.constant &&
             tree.r.degree == 0 && tree.r.constant
        )
        if constantsBelow
            # NaN checks:
            l = tree.l.val
            r = tree.r.val
            if isbad(l) || isbad(r)
                return tree
            end

            # Actually compute:
            out = options.binops[tree.op](l, r)
            if isbad(out)
                return tree
            end
            return Node(convert(CONST_TYPE, out))
        end
    end
    return tree
end

# Expensive but powerful simplify using SymbolicUtils
function simplifyWithSymbolicUtils(tree::Node, options::Options, curmaxsize::Int)::Node
    if !(((+) in options.binops) && ((*) in options.binops))
        return tree
    end
    init_node = copyNode(tree)
    init_size = countNodes(tree)
    symbolic_util_form, complete = node_to_symbolic_safe(tree, options, index_functions=true)
    if !complete
        return init_node
    end
    eqn_form, complete2 = custom_simplify(symbolic_util_form, options)
    if !complete2
        return init_node
    end
    final_node = symbolic_to_node(eqn_form, options)
    final_size = countNodes(tree)
    did_simplification_improve = (final_size <= init_size) && (check_constraints(final_node, options, curmaxsize))
    output = did_simplification_improve ? final_node : init_node

    return output
end

function simplifyWithSymbolicUtils(tree::Node, options::Options)::Node
    simplifyWithSymbolicUtils(tree, options, options.maxsize)
end


text:
struct NegativeBinomial <: Distribution{Int} end

"""
    neg_binom(r::Real, p::Real)
Sample an `Int` from a Negative Binomial distribution. Returns the number of
failures before the `r`th success in a sequence of independent Bernoulli trials.
`r` is the number of successes (which may be fractional) and `p` is the
probability of success per trial.
"""
const neg_binom = NegativeBinomial()

function logpdf(::NegativeBinomial, x::Int, r::Real, p::Real)
    Distributions.logpdf(Distributions.NegativeBinomial(r, p), x)
end

function logpdf_grad(::NegativeBinomial, x::Int, r::Real, p::Real)
    r_grad = x == 0 ? log(p) : sum(1/(x+r-i) for i in 1:x) + log(p)
    p_grad = x >= 0 ? r/p - (1/(1-p) * x) : 0.0
    return (nothing, r_grad, p_grad)
end

function random(::NegativeBinomial, r::Real, p::Real)
    rand(Distributions.NegativeBinomial(r, p))
end

is_discrete(::NegativeBinomial) = true

(::NegativeBinomial)(r, p) = random(NegativeBinomial(), r, p)

has_output_grad(::NegativeBinomial) = false
has_argument_grads(::NegativeBinomial) = (true, true)

export neg_binom

ext: jl


text:
using Gnuplot
x = y = -15:0.33:15
fz(x,y) = sin.(sqrt.(x.^2 + y.^2))./sqrt.(x.^2+y.^2)
fxy = [fz(x,y) for x in x, y in y]

@gsp(x, y, fxy, "w pm3d", "set ylabel 'y' textcolor rgb 'red'",
     "set xlabel 'x' textcolor rgb 'blue'", "set key off")

# Draw contours using the same palette as the surface plot
z = -2.5  # z coordinate at which contours are drawn
@gsp :- "set cbrange [$(minimum(fxy)):$(maximum(fxy))]"
@gsp :- "set xyplane 0"
levels = contourlines(x, y, fxy)
for i in 1:length(levels)
    @gsp :- levels[i].data "u 1:2:($z):($(levels[i].z)) w l notit lc pal" :-
end
@gsp

max_stars_repo_name: lazarusA/gnuplot-examples

max_issues_repo_name: lazarusA/gnuplot-examples

max_forks_repo_name: lazarusA/gnuplot-examples


text:
# block_diag.jl
# test

using LinearMapsAA: LinearMapAA, LinearMapAM, LinearMapAO
using SparseArrays: blockdiag, sparse
using Test: @test, @testset

M1 = rand(3,2)
M2 = rand(5,4)
M = [M1 zeros(size(M1,1),size(M2,2));
	zeros(size(M2,1),size(M1,2)) M2]

A1 = LinearMapAA(M1) # WrappedMap
A2 = LinearMapAA(x -> M2*x, y -> M2'y, size(M2), T=eltype(M2)) # FunctionMap
B = block_diag(A1, A2)

@testset "block_diag for AM" begin
	@test Matrix(B) == M
	@test Matrix(B)' == Matrix(B')
end

# test LMAO
@testset "block_diag for AO" begin
	Ao = LinearMapAA(M2 ; odim=(1,5))
	@test block_diag(A1, A1) isa LinearMapAM
	@test block_diag(A1, A1 ; tryop=true) isa LinearMapAO
	@test block_diag(A1, Ao) isa LinearMapAM
	Bo = block_diag(Ao, Ao)
	@test Bo isa LinearMapAO

	Md = blockdiag(sparse(M2), sparse(M2))
	X = rand(4,2)
	Yo = Bo * X
	Yd = Md * vec(X)
	@test vec(Yo) ≈ Yd

	@test Matrix(Bo)' == Matrix(Bo')
end


text:
##########################
### Variance estimator ###
##########################
abstract type VarEstimator{TI,TF} end

# Ref： https://github.com/stan-dev/math/blob/develop/stan/math/prim/mat/fun/welford_var_estimator.hpp
mutable struct WelfordVar{TI,TF} <: VarEstimator{TI,TF}
    n :: TI
    μ :: Vector{TF}
    M :: Vector{TF}
end

function reset!(wv::WelfordVar{TI,TF}) where {TI<:Integer,TF<:Real}
    wv.n = zero(TI)
    wv.μ .= zero(TF)
    wv.M .= zero(TF)
end

function add_sample!(wv::WelfordVar, s::AbstractVector)
    wv.n += 1
    δ = s .- wv.μ
    wv.μ .+= δ ./ wv.n
    wv.M .+= δ .* (s .- wv.μ)
end

# https://github.com/stan-dev/stan/blob/develop/src/stan/mcmc/var_adaptation.hpp
function get_var(wv::VarEstimator{TI,TF})::Vector{TF} where {TI<:Integer,TF<:Real}
    n, M = wv.n, wv.M
    @assert n >= 2 "Cannot get variance with only one sample"
    return (n / ((n + 5) * (n - 1))) .* M .+ 1e-3 * (5 / (n + 5))
end

abstract type CovarEstimator{TI<:Integer,TF<:Real} end

# NOTE: this naive covariance estimator is used only in testing
mutable struct NaiveCovar{TI,TF} <: CovarEstimator{TI,TF}
    n :: TI
    S :: Vector{Vector{TF}}
end

function add_sample!(nc::NaiveCovar, s::AbstractVector)
    nc.n += 1
    push!(nc.S, s)
end

function reset!(nc::NaiveCovar{TI,TF}) where {TI<:Integer,TF<:Real}
    nc.n = zero(TI)
    nc.S = Vector{Vector{TF}}()
end

function get_covar(nc::NaiveCovar{TI,TF})::Matrix{TF} where {TI<:Integer,TF<:Real}
    @assert nc.n >= 2 "Cannot get variance with only one sample"
    return Statistics.cov(nc.S)
end

# Ref: https://github.com/stan-dev/math/blob/develop/stan/math/prim/mat/fun/welford_covar_estimator.hpp
mutable struct WelfordCovar{TI<:Integer,TF<:Real} <: CovarEstimator{TI,TF}
    n :: TI
    μ :: Vector{TF}
    M :: Matrix{TF}
end

function reset!(wc::WelfordCovar{TI,TF}) where {TI<:Integer,TF<:Real}
    wc.n = zero(TI)
    wc.μ .= zero(TF)
    wc.M .= zero(TF)
end

function add_sample!(wc::WelfordCovar, s::AbstractVector)
    wc.n += 1
    δ = s .- wc.μ
    wc.μ .+= δ ./ wc.n
    wc.M .+= (s .- wc.μ) * δ'
end

# Ref: https://github.com/stan-dev/stan/blob/develop/src/stan/mcmc/covar_adaptation.hpp
function get_covar(wc::WelfordCovar{TI,TF})::Matrix{TF} where {TI<:Integer,TF<:Real}
    n, M = wc.n, wc.M
    @assert n >= 2 "Cannot get variance with only one sample"
    return (n / ((n + 5) * (n - 1))) .* M + 1e-3 * (5 / (n + 5)) * LinearAlgebra.I
end

################
### Adapters ###
################

abstract type PreConditioner <: AbstractAdapter end

struct UnitPreConditioner <: PreConditioner end

function Base.string(::UnitPreConditioner)
    return string([1.0])
end

struct DiagPreConditioner{TI<:Integer,TF<:Real} <: PreConditioner
    ve  :: VarEstimator{TI,TF}
    std :: Vector{TF}
end

function DiagPreConditioner(d::Integer)
    ve = WelfordVar(0, zeros(d), zeros(d))
    return DiagPreConditioner(ve, Vector(ones(d)))
end

function Base.string(dpc::DiagPreConditioner)
    return string(dpc.std)
end

function adapt!(dpc::DiagPreConditioner, θ, is_addsample::Bool, is_updatestd::Bool)
    if is_addsample
        add_sample!(dpc.ve, θ)
    end
    if is_updatestd
        var = get_var(dpc.ve)
        dpc.std .= sqrt.(var)
        reset!(dpc.ve)
        return true
    end
    return false
end

struct DensePreConditioner{TI<:Integer,TF<:Real} <: PreConditioner
    ce    :: CovarEstimator{TI,TF}
    covar :: Matrix{TF}
end

function DensePreConditioner(d::Integer)
    ce = WelfordCovar(0, zeros(d), zeros(d,d))
    # TODO: take use of the line below when we have an interface to set which pre-conditioner to use
    # ce = NaiveCovar(0, Vector{Vector{Float64}}())
    return DensePreConditioner(ce, LinearAlgebra.diagm(0 => ones(d)))
end

function Base.string(dpc::DensePreConditioner)
    return string(LinearAlgebra.diag(dpc.covar))
end

function adapt!(dpc::DensePreConditioner, θ, is_addsample::Bool, is_updatestd::Bool)
    if is_addsample
        add_sample!(dpc.ce, θ)
    end
    if is_updatestd
        covar = get_covar(dpc.ce)
        dpc.covar .= covar
        reset!(dpc.ce)
        return true
    end
    return false
end


text:
using CompScienceMeshes
import PlotlyJS

Faces = meshsphere(1.0, 0.55)

faces = barycentric_refinement(Faces)
PlotlyJS.plot(patch(faces.mesh))
# error("top")
NV = numvertices(Faces)

# Nodes = skeleton(Faces, 0)
# node_ctrs = [vertices(Nodes)[node][1] for node in cells(Nodes)]
# # C = [c[i] for c in node_ctrs, i in 1:3]
# Nodes.faces = Nodes.faces[CompScienceMeshes.sort_sfc(node_ctrs)]
#
# fine = Mesh(verts, fcs)
# D = connectivity(Nodes, fine)
# rows, vals = rowvals(D), nonzeros(D)
# sorted_fcs = Vector{celltype(mesh)}()
# for (i,Node) in enumerate(cells(Nodes))
#     for k in nzrange(D,i)
#         j = rows[k]
#         push!(sorted_fcs, fcs[j])
#     end
# end

vertices(Faces)[1:NV] == vertices(faces)[1:NV]

Nodes = skeleton(Faces,0)

D = connectivity(Nodes, faces)

import PlotlyJS
numcells(faces)
nf = 224
t1 = patch(Mesh(faces.mesh.vertices, faces.mesh.faces[1:nf]), rand(numcells(faces)))
PlotlyJS.plot(t1)

ctrs = [cartesian(center(chart(Nodes,c))) for c in cells(Nodes)]
Ctrs = [ctr[i] for ctr in ctrs, i in 1:3]
nc = div(length(ctrs),3)
t2 = PlotlyJS.scatter3d(x=Ctrs[1:nc,1], y=Ctrs[1:nc,2], z=Ctrs[1:nc,3])

Edges = skeleton(Faces,1)
ctrs = [cartesian(center(chart(Edges,c))) for c in cells(Edges)]
Ctrs = [ctr[i] for ctr in ctrs, i in 1:3]
nc = div(length(ctrs),3)
t3 = PlotlyJS.scatter3d(x=Ctrs[1:nc,1], y=Ctrs[1:nc,2], z=Ctrs[1:nc,3])

PlotlyJS.plot([t1,t2,t3])

aux = skeleton(faces,2)

t4 = patch(Mesh(aux.vertices, aux.faces[1:nf]), rand(nf))

PlotlyJS.plot(t4)


text:
# -----------------------------------------------------------
#
# Constructing the matrix elements of the derivative
# operator in the 1-D DG Basis
#
# -----------------------------------------------------------

# Efficiency criticality: LOW
# Computations only performed once

# Accuracy criticality: HIGH
# Critical for accurate PDE evolution

# -----------------------------------------------------
# 1-D Derivatives
# -----------------------------------------------------

function dLegendreP(k,x)
    k<=K_max || throw(DomainError())
    return array2poly(symbolic_diff(leg_coeffs[k+1]),x)
end

function dh(k,mode,x)
    mode<=k || throw(DomainError())
    return array2poly(symbolic_diff((@view dg_coeffs[k][:,mode])),x)
end

function dleg(mode::Int, x::T) where T <: Real
    return sqrt(2.0)*dLegendreP(mode-1, 2*x-1) *2
end

function dbasis(level::Int, cell::Int, mode::Int, x::T) where T <: Real
    return dleg(mode, (1<<level)*x - (cell-1)) * (2.0)^(level/2) * (1<<level)
end

function dbasis(level::Int, cell::Int, mode::Int)
    return x->dbasis(level, cell, mode, x)
end

# -----------------------------------------------------
# Shifted and scaled derivatives: v
# -----------------------------------------------------

function dv(k::Int, level::Int, cell::Int, mode::Int, x::Real)
    if level==0 # At the base level, we are not discontinuous, and we simply
                # use Legendre polynomials up to degree k-1 as a basis
        return 2*dLegendreP(mode-1,2*x-1)*sqrt(2.0)
    else
        return (1<<(level))*dh(k, mode, (1<<level)*x - (2*cell-1)) * (2.0)^(level/2)
        # Otherwise we use an appropriately shifted, scaled, and normalized
        # DG function
    end
end

function dv(k::Int, level::Int, cell::Int, mode::Int)
    return (xs::Real -> dv(k,level,cell,mode,xs))
end

# -----------------------------------------------------
# Numerical Inner Product
# -----------------------------------------------------

# TODO: There is a way to do this all symbolically, with no use for numerics
function inner_product1D(f::Function, g::Function, lvl::Int, cell::Int;
                                rtol = REL_TOL, atol = ABS_TOL, maxevals=MAX_EVALS)
    xmin = (cell-1)/(1<<max(0, lvl-1))
    xmax = (cell)/(1<<max(0, lvl-1))
    h = (x-> f(x)*g(x))
    (val, err) = hquadrature(h, xmin, xmax; rtol=rtol, atol=atol, maxevals=maxevals)
    return val
end

# -----------------------------------------------------
# Taking Derivative of Array Representing Polynomial
# -----------------------------------------------------

function symbolic_diff(v::AbstractArray{T}) where T <: Real
    n=length(v)
    k=div(n,2)
    ans = zeros(T,n)
    for i in 1:n
        if i<k
            ans[i] = i*v[i+1]
        elseif i > k && i<2k
            ans[i] = (i-k) * v[i+1]
        else
            ans[i]=0
        end
    end
    return ans
end

# -----------------------------------------------------
# < f_1 | D | f_2 > matrix elements
# -----------------------------------------------------

# Matrix element of legendre basis on [-1, 1]
function legendreDlegendre(mode1::Int, mode2::Int)
    return inner_product(leg_coeffs[mode1], symbolic_diff(leg_coeffs[mode2]))
end

# Matrix element of position basis elements
function legvDv(level, cell1, mode1, cell2, mode2;
                rtol = REL_TOL, atol=ABS_TOL, maxevals=MAX_EVALS)
    if cell1 == cell2
        return (1<<(level+1))*legendreDlegendre(mode1, mode2)
    end
    return 0.0
end

# Matrix element of DG basis on [-1, 1]
function hDh(k::Int, mode1::Int, mode2::Int)
    return inner_product((@view dg_coeffs[k][:,mode1]),
                         symbolic_diff((@view dg_coeffs[k][:,mode2])))
end

# Matrix element of hierarchical basis elements
function vDv(k::Int, lvl1::Int, cell1::Int, mode1::Int, lvl2::Int, cell2::Int, mode2::Int;
                    rtol = REL_TOL, atol = ABS_TOL, maxevals=MAX_EVALS)
    if lvl1 == lvl2
        if lvl1 == 0
            return 2*legendreDlegendre(mode1, mode2)
        end
        if cell1 == cell2
            return hDh(k, mode1, mode2)*(1<<max(0, lvl1))
        end
        return 0.0
    end
    if lvl1 < lvl2
        if lvl1 == 0
            return inner_product1D(v(k,0,1,mode1), dv(k,lvl2,cell2,mode2), lvl2, cell2;
                                    rtol = rtol, atol = atol, maxevals=maxevals)
        end
        if (1<<(lvl2-lvl1))*cell1 >= cell2 && (1<<(lvl2-lvl1))*(cell1-1) < cell2
            return inner_product1D(v(k,lvl1,cell1,mode1), dv(k,lvl2,cell2,mode2), lvl2, cell2;
                                    rtol = rtol, atol = atol, maxevals=maxevals)
        end
        return 0.0
    end
    return 0.0

end

# -----------------------------------------------------
# < f_1 | D | f_2 > for a specific f_2 over all f_1
# -----------------------------------------------------

function diff_basis_DG(k::Int, level::Int, cell::Int, mode::Int)
    dcoeffs = Array{Float64}(undef, level+1, k)
    p = cell
    for l in level:-1:0
        for f_n in 1:k
            dcoeffs[l+1,f_n] = vDv(k, l, p, f_n, level, cell, mode)
        end
        p = ceil(Int, p/2)
    end
    return dcoeffs
end


text:
"""
    flux_first_order!(
        bl::BalanceLaw,
        flux::Grad,
        state::Vars,
        aux::Vars,
        t::Real
    )

Computes (and assembles) flux terms `F¹(Y)` in:

```
∂Y
-- + ∇ • F¹(Y) + ∇ • F²(Y,G) = S(Y, G),     G = ∇Y
∂t
```

Computes and assembles non-diffusive
fluxes in the model equations.

For this fallback to work, several methods must be defined:
 - [`prognostic_vars`](@ref)
 - [`eq_tends`](@ref)
 - [`get_prog_state`](@ref)
optionally,
 - [`precompute`](@ref)
and individual [`flux`](@ref) kernels that
are defined for each type that `eq_tends` returns.
"""
@inline function flux_first_order!(
    bl::BalanceLaw,
    flux,
    state,
    aux,
    t,
    direction,
)

    tend = Flux{FirstOrder}()
    _args = (; state, aux, t, direction)
    args = merge(_args, (precomputed = precompute(bl, _args, tend),))

    map(prognostic_vars(bl)) do prog
        var, name = get_prog_state(flux, prog)
        val = Σfluxes(prog, eq_tends(prog, bl, tend), bl, args)
        setproperty!(var, name, val)
    end
    nothing
end

"""
    flux_second_order!(
        bl::BalanceLaw,
        flux::Grad,
        state::Vars,
        diffusive::Vars,
        hyperdiffusive::Vars,
        aux::Vars,
        t::Real
    )

Computes (and assembles) flux terms `F²(Y, G)` in:

```
∂Y
-- + ∇ • F¹(Y) + ∇ • F²(Y,G) = S(Y, G),     G = ∇Y
∂t
```

Diffusive fluxes in BalanceLaw. Viscosity, diffusivity are calculated
in the turbulence subcomponent and accessed within the diffusive flux
function. Contributions from subcomponents are then assembled (pointwise).

For this fallback to work, several methods must be defined:
 - [`prognostic_vars`](@ref)
 - [`eq_tends`](@ref)
 - [`get_prog_state`](@ref)
optionally,
 - [`precompute`](@ref)
and individual [`flux`](@ref) kernels that
are defined for each type that `eq_tends` returns.
"""
@inline function flux_second_order!(
    bl::BalanceLaw,
    flux,
    state,
    diffusive,
    hyperdiffusive,
    aux,
    t,
)
    tend = Flux{SecondOrder}()
    _args = (; state, aux, t, diffusive, hyperdiffusive)
    args = merge(_args, (precomputed = precompute(bl, _args, tend),))

    map(prognostic_vars(bl)) do prog
        var, name = get_prog_state(flux, prog)
        val = Σfluxes(prog, eq_tends(prog, bl, tend), bl, args)
        setproperty!(var, name, val)
    end
    nothing
end

"""
    source!(
        bl::BalanceLaw,
        source::Vars,
        state::Vars,
        diffusive::Vars,
        aux::Vars,
        t::Real,
        direction::Direction,
    )

Computes (and assembles) source terms `S(Y)` in:

```
∂Y
-- + ∇ • F¹(Y) + ∇ • F²(Y,G) = S(Y, G),     G = ∇Y
∂t
```

For this fallback to work, several methods must be defined:
 - [`prognostic_vars`](@ref)
 - [`eq_tends`](@ref)
 - [`get_prog_state`](@ref)
optionally,
 - [`precompute`](@ref)
and individual [`source`](@ref) kernels that
are defined for each type that `eq_tends` returns.
"""
function source!(bl::BalanceLaw, source, state, diffusive, aux, t, direction)
    tend = Source()
    _args = (; state, aux, t, direction, diffusive)
    args = merge(_args, (precomputed = precompute(bl, _args, tend),))

    map(prognostic_vars(bl)) do prog
        var, name = get_prog_state(source, prog)
        val = Σsources(prog, eq_tends(prog, bl, tend), bl, args)
        setproperty!(var, name, val)
    end
    nothing
end


text:
export salt_pepper, salt_pepper_chn

salt_or_pepper(salt_prob, salt, pepper) = rand() < salt_prob ? salt : pepper

 # function for a raw value
function f_sp(prob, salt_prob, salt, pepper)
    return x -> rand() < prob ? salt_or_pepper(salt_prob, salt, pepper) : x
end

@inline function comb_sp(x, n)
    return n
end

function salt_pepper(X, prob=0.1; salt_prob=0.5, salt=1, pepper=0)
    return salt_pepper!(copy(X), prob, salt_prob=salt_prob, salt=salt, pepper=pepper)
end

 # for arbitrary array, RGB and Gray images
function salt_pepper!(X, prob=0.1; salt_prob=0.5, salt=1, pepper=0)

    # type conversion of salt & pepper makes it faster
    a = X[1]
    salt = convert(eltype(a), salt)
    pepper = convert(eltype(a), pepper)

    return apply_noise!(comb_sp, f_sp(prob, salt_prob, salt, pepper), X, false)
end

"""
    salt_pepper(X; salt_prob=0.5, salt=1.0, pepper=0.0[, prob=0.1])

Returns array `X` affected by salt and pepper noise.
`X` can be an array or an RGB or Gray image
`prob` is a optional argument for the probability that a pixel will be affected by the noise.
`salt_prob` is a keyword argument representing the probability for salt noise.
The probability for pepper noise is therefore 1-`salt_prob`.
`salt` is a keyword argument for specifying the value of salt noise.
`pepper` is a keyword argument for specifying the value of pepper noise.
"""
salt_pepper

function salt_pepper_chn(X, prob=0.1; salt_prob=0.5, salt=1.0, pepper=0.0)
    return salt_pepper_chn!(copy(X), prob, salt_prob=salt_prob, salt=salt, pepper=pepper)
end

function salt_pepper_chn!(X, prob=0.1; salt_prob=0.5, salt=1.0, pepper=0.0)
    return apply_noise_chn!(comb_sp, f_sp(prob, salt_prob, salt, pepper), X, false)
end

"""
    salt_pepper_chn(X; salt_prob=0.5, salt=1.0, pepper=0.0[, prob=0.1])

Returns a RGB Image `X` affected by salt and pepper noise.
When a salt or pepper occurs, it is applied to all channels of the RGB making a real salt
and pepper on the whole image.
`prob` is a optional argument for the probability that a pixel will be affected by the noise.
`salt_prob` is a keyword argument representing the probability for salt noise.
The probability for pepper noise is therefore 1-`salt_prob`.
`salt` is a keyword argument for specifying the value of salt noise.
`pepper` is a keyword argument for specifying the value of pepper noise.
"""
salt_pepper_chn

max_issues_repo_name: roflmaostc/Noise

max_forks_repo_name: roflmaostc/Noise


text:
using SOAPlite
using Test, LinearAlgebra, BenchmarkTools
using JuLIP, JuLIP.Testing

@testset "SOAPlite.jl" begin
    include("test.jl")
end


text:
using Color, FixedPointNumbers, Compat
using Base.Test

const red = color("red")
const red24 = RGB24(0x00ff0000)
const red32 = ARGB32(0xffff0000)
for T in (Float64, Float32, Ufixed8)
    c = RGB(one(T), zero(T), zero(T))
    @test eltype(c) == T
    c64 = convert(RGB{Float64}, c)
    @test typeof(c64) == RGB{Float64}
    @test c64 == red
    cr = convert(RGB{T}, red)
    @test cr == c
end
@test RGB(1,0,0) == red
@test RGB(convert(UInt8, 1),0,0) == red
@test RGB(convert(UInt8, 1),convert(UInt8, 0),convert(UInt8, 0)) == red
@test convert(RGB, red24) == red

for Cto in Color.CVparametric
    for Cfrom in Color.CVparametric
        for Tto in (Float32, Float64)
            for Tfrom in (Float32, Float64)
                c = convert(Cfrom{Tfrom}, red)
                @test typeof(c) == Cfrom{Tfrom}
                c1 = convert(Cto, c)
                @test eltype(c1) == Tfrom
                c2 = convert(Cto{Tto}, c)
                @test typeof(c2) == Cto{Tto}
            end
        end
    end
end
for Cto in Color.CVparametric
    @test typeof(convert(Cto, red24)) == Cto{Float64}
    @test typeof(convert(Cto{Float32}, red24)) == Cto{Float32}
end

# Test conversion from Ufixed types
for Cto in Color.CVfloatingpoint
    for Cfrom in Color.CVfractional
        for Tto in (Float32, Float64)
            for Tfrom in (Ufixed8, Ufixed10, Ufixed12, Ufixed14, Ufixed16)
                c = convert(Cfrom{Tfrom}, red)
                @test typeof(c) == Cfrom{Tfrom}
                c1 = convert(Cto, c)
                @test eltype(c1) == Float64
                c2 = convert(Cto{Tto}, c)
                @test typeof(c2) == Cto{Tto}
            end
        end
    end
end

# Test conversion to Ufixed types
for Cto in Color.CVfractional
    for Cfrom in Color.CVfloatingpoint
        for Tto in (Ufixed8, Ufixed10, Ufixed12, Ufixed14, Ufixed16)
            for Tfrom in (Float32, Float64)
                c = convert(Cfrom{Tfrom}, red)
                @test typeof(c) == Cfrom{Tfrom}
                c2 = convert(Cto{Tto}, c)
                @test typeof(c2) == Cto{Tto}
            end
        end
    end
end

ac = rgba(red)

@test convert(RGB, ac) == RGB(1,0,0)
@test convert(RGB{Ufixed8}, ac) == RGB{Ufixed8}(1,0,0)
@test convert(RGBA{Ufixed8}, ac) == RGBA{Ufixed8}(1,0,0,1)
@test convert(HSVA, ac) == HSVA{Float64}(convert(HSV, red), 1.0)
@test convert(HSVA{Float32}, ac) == HSVA{Float32}(convert(HSV{Float32}, red), 1.0f0)
@test convert(RGBA, red) == ac

@test convert(ARGB32, ac) == ARGB32(0xffff0000)
@test convert(Uint32, convert(ARGB32, ac)) == 0xffff0000
@test convert(RGB24, RGB(0xffuf8,0x00uf8,0x00uf8)) == RGB24(0x00ff0000)
@test convert(Uint32, convert(RGB24, RGB(0xffuf8,0x00uf8,0x00uf8))) == 0x00ff0000
redhsv = convert(HSV, red)
@test convert(RGB24, redhsv) == RGB24(0x00ff0000)

@test convert(RGB{Ufixed8}, red24) == RGB{Ufixed8}(1,0,0)
@test convert(RGBA{Ufixed8}, red32) == RGBA{Ufixed8}(1,0,0,1)
@test convert(HSVA{Float64}, red32) == HSVA{Float64}(360, 1, 1, 1)

@test_throws MethodError AlphaColorValue(RGB(1,0,0), 0xffuf8)

# Test vector space operations
import Base.full
full(T::ColorValue) = map(x->getfield(T, x), fieldnames(T)) #Allow test to access numeric elements
# Generated from:
#=
julia> for t in subtypes(ColorValue)
          isleaftype(t) || isleaftype(t{Float64}) || continue

          try
              println("@test_approx_eq_eps ", t(0.125, 0.5, 0), "+", t(0.2, 0.7, 0.4), " ", t(0.125, 0.5, 0) + t(0.2, 0.7, 0.4), " 91eps()")
              println("@test_approx_eq_eps 3", t(0.125, 0.5, 0.03), " ", 3t(0.125, 0.5, 0.03), " 91eps()\n")
          catch    continue
          end
       end
=#
@test_approx_eq_eps DIN99o{Float64}(0.125,0.5,0.0)+DIN99o{Float64}(0.2,0.7,0.4) DIN99o{Float64}(0.3249177178200238,1.1851041291240045,0.39613501908850973) 91eps()
@test_approx_eq_eps 3DIN99o{Float64}(0.125,0.5,0.03) DIN99o{Float64}(0.3748457441768566,1.4684725814274917,0.08810835488564604) 91eps()

@test_approx_eq_eps DIN99{Float64}(0.125,0.5,0.0)+DIN99{Float64}(0.2,0.7,0.4) DIN99{Float64}(0.3247634199662825,1.1845967897238356,0.3960040943105819) 91eps()
@test_approx_eq_eps 3DIN99{Float64}(0.125,0.5,0.03) DIN99{Float64}(0.37455660462520735,1.467408770792347,0.08804452624754316) 91eps()

@test_approx_eq_eps HSL{Float64}(0.125,0.5,0.0)+HSL{Float64}(0.2,0.7,0.4) HSL{Float64}(0.2000361824889771,0.7000000631323785,0.279999991890043) 91eps()
@test_approx_eq_eps 3HSL{Float64}(0.125,0.5,0.03) HSL{Float64}(0.17819347263565605,0.3941314428230693,0.029133626691685854) 91eps()

@test_approx_eq_eps HSV{Float64}(0.125,0.5,0.0)+HSV{Float64}(0.2,0.7,0.4) HSV{Float64}(0.20002334068804029,0.7000000071033119,0.3999999666630183) 91eps()
@test_approx_eq_eps 3HSV{Float64}(0.125,0.5,0.03) HSV{Float64}(0.15547682101166416,0.42727995497106286,0.07819689447804021) 91eps()

@test_approx_eq_eps LCHab{Float64}(0.125,0.5,0.0)+LCHab{Float64}(0.2,0.7,0.4) LCHab{Float64}(0.3249999999999993,1.1999928922680212,0.23333346495867205) 91eps()
@test_approx_eq_eps 3LCHab{Float64}(0.125,0.5,0.03) LCHab{Float64}(0.375,1.4999999999999816,0.03000000000047518) 91eps()

@test_approx_eq_eps LCHuv{Float64}(0.125,0.5,0.0)+LCHuv{Float64}(0.2,0.7,0.4) LCHuv{Float64}(0.32500000000000007,1.200147006104681,0.2329439337672688) 91eps()
@test_approx_eq_eps 3LCHuv{Float64}(0.125,0.5,0.03) LCHuv{Float64}(0.375,1.5000000000000007,0.03000000000000911) 91eps()

@test_approx_eq_eps LMS{Float64}(0.125,0.5,0.0)+LMS{Float64}(0.2,0.7,0.4) LMS{Float64}(0.32500000000000007,1.2000000000000002,0.4000000000000001) 91eps()
@test_approx_eq_eps 3LMS{Float64}(0.125,0.5,0.03) LMS{Float64}(0.37499999999999994,1.4999999999999998,0.09000000000000001) 91eps()

@test_approx_eq_eps Lab{Float64}(0.125,0.5,0.0)+Lab{Float64}(0.2,0.7,0.4) Lab{Float64}(0.3249999999999993,1.2000000000000066,0.40000000000000036) 91eps()
@test_approx_eq_eps 3Lab{Float64}(0.125,0.5,0.03) Lab{Float64}(0.375,1.4999999999999876,0.09000000000001229) 91eps()

@test_approx_eq_eps Luv{Float64}(0.125,0.5,0.0)+Luv{Float64}(0.2,0.7,0.4) Luv{Float64}(0.32500000000000007,1.2112171964760525,0.3551312140957904) 91eps()
@test_approx_eq_eps 3Luv{Float64}(0.125,0.5,0.03) Luv{Float64}(0.375,1.4999999999999998,0.09000000000000033) 91eps()

@test_approx_eq_eps XYZ{Float64}(0.125,0.5,0.0)+XYZ{Float64}(0.2,0.7,0.4) XYZ{Float64}(0.325,1.2,0.4) 91eps()
@test_approx_eq_eps 3XYZ{Float64}(0.125,0.5,0.03) XYZ{Float64}(0.375,1.5,0.09) 91eps()

@test_approx_eq_eps xyY{Float64}(0.125,0.5,0.0)+xyY{Float64}(0.2,0.7,0.4) xyY{Float64}(0.2,0.7,0.4) 91eps()
@test_approx_eq_eps 3xyY{Float64}(0.125,0.5,0.03) xyY{Float64}(0.125,0.5,0.09) 91eps()

#59
@test Color.xyz_to_uv(XYZ{Float64}(0.0, 0.0, 0.0)) === (0.0, 0.0)
@test Color.xyz_to_uv(XYZ{Float64}(0.0, 1.0, 0.0)) === (0.0, 0.6)
@test Color.xyz_to_uv(XYZ{Float64}(0.0, 1.0, 1.0)) === (0.0, 0.5)
@test Color.xyz_to_uv(XYZ{Float64}(1.0, 0.0, 1.0)) === (1.0, 0.0)
@test Color.xyz_to_uv(XYZ{Float64}(1.0, 0.0, 0.0)) === (4.0, 0.0)

@test 1.0LCHuv(0.0, 0.0, 0.0) === LCHuv(0.0, 0.0, 0.0)


text:
#
println("Dd) Test of the PhotoRecombination module with ASF from an internally generated initial- and final-state multiplet.")
#
setDefaults("print summary: open", "zzz-PhotoRecombination.sum")
setDefaults("method: continuum, asymptotic Coulomb")  ## setDefaults("method: continuum, Galerkin")
setDefaults("method: normalization, pure Coulomb")   ## setDefaults("method: normalization, pure Coulomb")    setDefaults("method: normalization, pure sine")

wa = Atomic.Computation("xx",  Nuclear.Model(12.); grid=JAC.Radial.Grid("grid: by given parameters"; rnt = 2.0e-6, h = 5.0e-2, hp = 2.0e-2, NoPoints = 600),
                        initialConfigs=[Configuration("1s^2 2s^2 2p^5"), Configuration("1s^2 2s 2p^6") ],
                        finalConfigs  =[Configuration("1s^2 2s^2 2p^6")],
                        process = JAC.Rec,
                        processSettings=PhotoRecombination.Settings([E1, M1], [JAC.UseCoulomb, JAC.UseBabushkin], [10., 30., 50.], [0.],
                                        false, true, true, true, true, Tuple{Int64,Int64}[(1,1)]) )

wb = perform(wa)
setDefaults("print summary: close", "")


text:
"""
    LNAInput

This is a structure which holds the inputs needed for computing the Linear Noise Approximation (LNA). This structure will hold the stochastic system as provided by the user; uniquely defined through kinetic parameters, the rates of the system and the stoichiometry matrix.

# Arguments
- `params::AbstractArray{Float64,1}`: The rate parameters of the stochastic model.
- `S::AbstractArray{Float64,2}`: the stochiometry matrix of the system. Size: number of reactions x number of species.
- `reaction_rate_function::Function,`: This is a function f(x, parameters) which should return an array of the reaction rates of the system, i.e. S*f would describe the ODE representation of the system.
- `volume::Float64`: The volume of the reactants of the system.
"""
struct LNAInput
    params::AbstractArray{Float64,1}
    S::AbstractArray{Float64,2}
    reaction_rate_function::Function
    volume::Float64
end

"""
    LNA

This is a structure which will hold the LNA: the mean of the trajectories and the covariance between the species.

# Arguments
- `traj_means`: A (number of species) x (number of time points) array which holds the mean trajectory for each species on each row of the array.
- `traj_covars`: An array which holds the covariance matrix of the species at each time point.
- `time_points`: The timepoints the system was solved for.
"""
struct LNA
    traj_means::AbstractArray{Float64,2}
    traj_covars::AbstractArray{AbstractArray{Float64,2},1}
    time_points::AbstractArray{Float64,1}
end

"""
    compute_LNA(input::LNAInput,
        x0::Tuple{AbstractArray{Float64,1},AbstractArray{Float64,2}},
        Tspan::Tuple{Float64,Float64},
        saveat::Float64,
        solver::OrdinaryDiffEq.OrdinaryDiffEqAlgorithm=RK4();
        kwargs...)

The function computes the linear noise approximation to system through construction of two ODEs: one describing the trajectories of the mean of the LNA and the other describing the change the covariance between the variables. These outputs are held in a LNA structure.

# Arguments
- `x0::Tuple{AbstractArray{Float64,2},AbstractArray{Float64,2}}`: The initial conditions of the system. In the form of (the initial conditions of the species, the initial covariance matrix of the system).
- `Tspan::Tuple{Float64,Float64}`: The start and end times of the simulation.
- `saveat::Float64`: The number of time points the use wishes to solve the system for.
- `solver::DEAlgorithm`: The ODE solver the user wishes to use, for example `DifferentialEquations.RK4()`.

#Returns
- LNA
"""
function compute_LNA(input::LNAInput,
    x0::Tuple{AbstractArray{Float64,1},AbstractArray{Float64,2}},
    Tspan::Tuple{Float64,Float64},
    saveat::Float64,
    solver::OrdinaryDiffEq.OrdinaryDiffEqAlgorithm=RK4();
    kwargs...)

    if input.volume <= 0.0
        throw(ArgumentError("To use LNA please provide a positive volume of the
            reactants"))
    end

    no_of_species, no_of_reactions = size(input.S)

    ###  the below ODE solves the ODE of the mean and the ODE of the covariance together.
    ## this function is the mean variance decomposition of the LNA. The derviation of which can be
    ## found in "Approximation  and  inference  methods  forstochastic  biochemical  kinetics—a  tutorial  review" (2017) by
    ## David  Schnoerr et al.
    ## x is two matrices concatenated on top of each other: a diagonal matrix with the means on the diagonal follwed by the covariance matrix
    function Mean_ODE(dx, x, pars, t)
        mean_vec = diag(x)
        covar_mx = x[no_of_species+1:no_of_species*2,:]
        reaction_rates = input.reaction_rate_function(mean_vec, pars)
        reaction_rates_jac = ForwardDiff.jacobian(y -> input.reaction_rate_function(y, pars), mean_vec)
        A = input.S * reaction_rates_jac
        dx[1:no_of_species, :] .= diagm(0 => input.S*reaction_rates)
        dx[no_of_species + 1:end, :] .= A*covar_mx + covar_mx*A' + input.S * diagm(0 => reaction_rates) * input.S' ./ sqrt(input.volume)
    end

    prob = ODEProblem(Mean_ODE, vcat(diagm(0=>x0[1]), x0[2]), Tspan, input.params)
    mean_and_var = solve(prob, solver, saveat=saveat; kwargs...)

    mean_traj = Array{Float64,2}(undef, no_of_species, length(mean_and_var.t))
    covar_traj = Array{Array{Float64,2},1}(undef, length(mean_and_var.t))

    for j in 1:length(mean_and_var.t)
        mean_traj[:,j] = diag(mean_and_var[j][1:no_of_species,1:no_of_species])
        covar_traj[j] = mean_and_var[j][no_of_species+1:end, 1:no_of_species]
    end

    return LNA(mean_traj, covar_traj, mean_and_var.t)
end

"""
    sample_LNA_trajectories(lna::LNA, n_samples::Int64)

A function which samples from the LNA to output sampled trajectories. The LNA gives the mean of the tracjectories and the covariance between them; hence a single trajectory can be sampled from a Multivariate Normal distribution. The user can also sample more than one trajectory; which are then averaged.

# Arguments
- `lna::LNA`: LNA stucture.
- `n_samples::Int64`: The number of sampled tracjectories to be sampled and then averaged.

#Returns
-  A (number of species) x (number of time points) array which holds the averaged trajectory for each species on each row of the array.
"""
function sample_LNA_trajectories(lna::LNA, n_samples::Int64)
    trajectories = [rand(MvNormal(lna.traj_means[:,j], lna.traj_covars[j]), n_samples) for j in 1:length(lna.time_points)]
    return hcat([mean(traj,dims=2) for traj in trajectories]...)
end

"""
    get_LNA_trajectories(input::LNAInput, n_samples::Int64,
        x0::Tuple{AbstractArray{Float64,1},AbstractArray{Float64,2}},
        Tspan::Tuple{Float64,Float64},
        saveat::Float64,
        solver::OrdinaryDiffEq.OrdinaryDiffEqAlgorithm=RK4();
        kwargs...)

A function which computes the LNA and then samples from the it to output sampled trajectories. The user can also sample more than one trajectory; which are then averaged.

# Arguments
- `input::LNAInput`: LNAInput stucture.
- `n_samples::Int64`: The number of sampled tracjectories to be sampled and then averaged.
- `x0::Tuple{AbstractArray{Float64,2},AbstractArray{Float64,2}}`: The initial conditions of the system. In the form of (the initial conditions of the species, the initial covariance matrix of the system).
- `Tspan::Tuple{Float64,Float64}`: The start and end times of the simulation.
- `saveat::Float64`: The number of time points the use wishes to solve the system for.
- `solver::DEAlgorithm`: The ODE solver the user wishes to use, for example DifferentialEquations.RK4() .

#Returns
-  A (number of species) x (number of time points) array which holds the averaged trajectory for each species on each row of the array.
"""
function get_LNA_trajectories(input::LNAInput, n_samples::Int64,
    x0::Tuple{AbstractArray{Float64,1},AbstractArray{Float64,2}},
    Tspan::Tuple{Float64,Float64},
    saveat::Float64,
    solver::OrdinaryDiffEq.OrdinaryDiffEqAlgorithm=RK4();
    kwargs...)

    lna = compute_LNA(input, x0, Tspan, saveat, solver; kwargs...)
    return sample_LNA_trajectories(lna, n_samples)
end

    #this function is the mean variance decomposition of the LNAS
    ##LNA_Mean_Var takes the inputs and then performs the simulating of the mean (ODE solution) and the covariance matrix which changes over time.
    ##x0 is the initial conditions of the DE and the covariance matrix
    ##S is the stochiometry matrix
    ##reaction_rate_function is a function of the rates - MUST be in the same order as the stochiometry matrix is laid out
    ##volume is the volume of reactants
    #
    #
    #
    #the maths for this can be found in UserManual_v28102013.pdf in the LNA folder in Gaussian Procceses Drop box, p2 equation (7) - the ODE of the covariance. A volume term has been added which can be seen in other sources - the particular source was linked as the layout is the similar to the one implemented.


text:
struct ClosedProblem{C}  # TODO: it's not clear if this specialisation will help or whether Any will do instead
    options::Options
    top_level::Problem
    flat::FlatProblem
    mfuncs::MonitorFunctions
    covering::C
end

DEFAULTCOVER(prob) = Problem("covering", nothing)
init!(::Nothing, closed) = nothing

function ClosedProblem(problem::Problem, options::Options, cover=DEFAULTCOVER)
    top_level = Problem("cont")
    push!(top_level, problem)
    mfuncs_problem = monitor_functions()
    mfuncs = mfuncs_problem.owner
    push!(top_level, mfuncs_problem)
    covering_problem = cover(top_level)
    covering = covering_problem.owner
    push!(top_level, covering_problem)
    flat = flatten(top_level)
    return closed = ClosedProblem(options, top_level, flat, mfuncs, covering)
end

function init!(closed::ClosedProblem)
    signal!(Signal(:pre_init), closed)
    init!(closed.mfuncs, closed)
    init!(closed.covering, closed)
    return signal!(Signal(:post_init), closed)
end

function run!(closed::ClosedProblem)
    signal!(Signal(:pre_run), closed)
    run!(closed.covering, closed)
    return signal!(Signal(:post_run), closed)
end

function signal!(signal::Signal, closed::ClosedProblem, args...)
    return signal!(closed.flat, signal, closed, args...)
end

function signal!(closed::ClosedProblem, signal::Signal, args...)
    return signal!(closed.flat, signal, args...)
end

get_options(problem::ClosedProblem) = problem.options
get_problem(problem::ClosedProblem) = problem.top_level
get_flatproblem(problem::ClosedProblem) = problem.flat
get_mfuncs(problem::ClosedProblem) = problem.mfuncs
get_covering(problem::ClosedProblem) = problem.covering

function get_initial_state(T::Type{<:Number}, closed::ClosedProblem)
    u = Vector{Vector{T}}()
    t = Vector{Vector{T}}()
    for var in closed.flat.var
        if var.initial_dim == 0
            push!(u, T[])
        elseif var.initial_u === nothing
            push!(u, zeros(T, var.initial_dim))
        elseif length(var.initial_u) == var.initial_dim
            push!(u, convert(Vector{T}, var.initial_u))
        else
            throw(ErrorException("Initial data for variable $(var) does not have the correct number of dimensions"))
        end
        if var.initial_dim == 0
            push!(t, T[])
        elseif var.initial_t === nothing
            push!(t, zeros(T, var.initial_dim))
        elseif length(var.initial_t) == var.initial_dim
            push!(t, convert(Vector{T}, var.initial_t))
        else
            throw(ErrorException("Initial tangent for variable $(var) does not have the correct number of dimensions"))
        end
    end
    d = Vector{Any}()
    for data in closed.flat.data
        push!(d, data.initial_data)
    end
    # Allow problems to alter the initial values/data
    signal!(Signal(:initial_state), closed, u, t, d)
    # Condense
    uu = reduce(vcat, u)
    tt = reduce(vcat, t)
    va = ViewAxis(length.(u))
    return (va=va, u=uu, t=tt, data=d)
end


text:
using IndirectTrajOpt
using IndirectCoStateInit
using Random
Random.seed!(1)

# Initialize BVP function
ps = initCR3BPIndirectParams("Low Thrust CR3BP")
tspan = (0.0, 8.6404*24*3600/ps.crp.TU)
bvpFunc(y0) = cr3bpOptIntegrate(y0, tspan, ps, copyParams = true, termCallbacks = true)

# Set initial and final conditions
ics = [-0.0194885115, -0.0160334798, 0.0,
       8.9188819237, -4.081936888, 0.0, 1.0]
fcs = [0.8233851820, 0.0, -0.02225563,
       0.0, 0.134184103, 0.0, 0.0]

# Initialize FFS Initializer
csInitializer = FSSCoStateInitializer(bvpFunc, ics, fcs; numParticles = 10000);

# Initialize co-states
initialize!(csInitializer)


text:
@testset "expected_loglik" begin
    # Test that the various methods of computing expectations return the same
    # result.
    rng = MersenneTwister(123456)
    q_f = Normal.(zeros(10), ones(10))

    likelihoods_to_test = [
        ExponentialLikelihood(),
        GammaLikelihood(),
        PoissonLikelihood(),
        GaussianLikelihood(),
    ]

    @testset "testing all analytic implementations" begin
        # Test that we're not missing any analytic implementation in `likelihoods_to_test`!
        implementation_types = [
            (; quadrature=m.sig.types[2], lik=m.sig.types[5]) for
            m in methods(ApproximateGPs.expected_loglik)
        ]
        analytic_likelihoods = [
            m.lik for m in implementation_types if
            m.quadrature == ApproximateGPs.Analytic && m.lik != Any
        ]
        for lik_type in analytic_likelihoods
            @test any(lik isa lik_type for lik in likelihoods_to_test)
        end
    end

    @testset "$(nameof(typeof(lik)))" for lik in likelihoods_to_test
        methods = [GaussHermite(100), MonteCarlo(1e7)]
        def = ApproximateGPs._default_quadrature(lik)
        if def isa Analytic
            push!(methods, def)
        end
        y = rand.(rng, lik.(zeros(10)))

        results = map(m -> ApproximateGPs.expected_loglik(m, y, q_f, lik), methods)
        @test all(x -> isapprox(x, results[end]; atol=1e-6, rtol=1e-3), results)
    end

    @test ApproximateGPs.expected_loglik(
        MonteCarlo(), zeros(10), q_f, GaussianLikelihood()
    ) isa Real
    @test ApproximateGPs.expected_loglik(
        GaussHermite(), zeros(10), q_f, GaussianLikelihood()
    ) isa Real
    @test ApproximateGPs._default_quadrature(θ -> Normal(0, θ)) isa GaussHermite
end


text:
##################################################
##
## Pattern matching. Follows the easier part of the masters
## thesis referenced here
## https://github.com/HPAC/matchpy/blob/master/matchpy/
##
## Part of the code in `pm.jl` is derived from the MIT licensed code
## of https://github.com/HPAC/matchpy/blob/master/matchpy/
##
##
## SymEngine things
BasicType = SymEngine.BasicType

head(ex::Basic) = SymEngine.get_symengine_class(ex)
args(ex::Basic)::Vector{Basic} = SymEngine.get_args(ex)

## call a SymEngine function from a symbol
call_fun(fn::Symbol, args) = call_fun(Val{fn}, fn, args)
call_fun(::Type{Val{:Add}}, fn, args) = isempty(args) ? Basic(0) : _call_fun(fn, args)
call_fun(::Type{Val{:Mul}}, fn, args) = isempty(args) ? Basic(1) : _call_fun(fn, args)
call_fun(::Type{Val{T}}, fn, args) where {T} = _call_fun(fn, args)
_call_fun(fn, args) = eval(Expr(:call, SymEngine.map_fn(fn, SymEngine.fn_map),args...))

#using Multisets
_setdiff(m1::Multiset, m2::Multiset) = m1 - intersect(m1, m2)

### Utilities ##################################################

## return vectors of vals, counts (sorted) as a tuple
tally(ps::Vector) = tally(Multiset(ps))
function tally(ps::Multiset{T}) where {T}
    d = Dict{T, Int}()
    for k in collect(ps)
        haskey(d, k) ? (d[k]+=1) : (d[k] = 1)
    end
    ks = collect(keys(d))
    vs = collect(values(d))

    sigma = sortperm(vs, rev=true)
    (ks[sigma], vs[sigma])
end

# so we can call repeat(a,m) to get what we want for a vector or not
_repeat(a::Basic;inner=times) = repeat([a], inner=inner)

### Define wildcard patterns ##############################################

# this allocates, and there isn't much to do about it
# x_: regular variable: 1 match (dot variable)
# x__: match 1 or more terms (plus variable) (in * or +)
# x___: match 0, 1 or more terms (star variable)
_check_dashes(u, pattern, ::Type{Val{false}}) =  endswith(u, pattern) ##startswith(u, pattern) || endswith(u, pattern)
function _check_dashes(u, pattern, ::Type{Val{true}})
    _check_dashes(u, pattern, Val{false}) && !_check_dashes(u, pattern*"_", Val{false})
end

iswildcard(x,exact=false) = issymbolic(x) && _check_dashes(SymEngine.toString(x), "_", Val{exact})
ispluswildcard(x, exact=false) = issymbolic(x) && _check_dashes(SymEngine.toString(x), "__", Val{exact})
isstarwildcard(x, exact=false) = issymbolic(x) && _check_dashes(SymEngine.toString(x), "___", Val{exact})

onlywildcard(x) = iswildcard(x,true)
onlypluswildcard(x) = ispluswildcard(x, true)
onlystarwildcard(x) = isstarwildcard(x, false)  ## 3 or more is star wildcard

isconstantpattern(p::Basic) = length(filter(iswildcard, free_symbols(p))) == 0
isconstantpattern(fn::Symbol) = true

## we use \phi_ for wildcards we don't care about "___" in match py
const ignore_wildcard_re = r"^ϕ"
function ignore_wildcard(x)
    issymbolic(x) || return false
    s = SymEngine.toString(x)
    _check_dashes(s, "_") || return false
    return ismatch(ignore_wildcard, s)
end

# ### Constraints XXX not implemented

# struct Constraint{S,T}
# op::S
# lhs::T
# rhs::T
# end
# function Base.show(io::IO, p::Constraint)
#     println(io, p.lhs, " ", p.op, " ", p.rhs)
# end

# ## Evaluate a constraint returning true or false using values in x, a dictionary
# (p::Constraint)(x)::Bool = p.op(p.lhs(x...), p.rhs(x...))

# # the main constraints
# Gt(a,b)  = Constraint(>,  Basic(a),Basic(b))
# Ge(a,b)  = Constraint(>=, Basic(a),Basic(b))
# Eq(a,b)  = Constraint(==, Basic(a),Basic(b))
# Neq(a,b) = Constraint(!=, Basic(a),Basic(b))
# Le(a,b)  = Constraint(<=, Basic(a),Basic(b))
# Lt(a,b)  = Constraint(<,  Basic(a),Basic(b))

## some predicates on values
import Base: isinteger, isnumber, iseven, isodd

issymbolic(::BasicType{Val{:Symbol}}) = true
issymbolic(::BasicType) = false
issymbolic(x::Basic) = issymbolic(BasicType(x))

isnumber(::T) where {T <: SymEngine.BasicNumber} = true
isnumber(::BasicType) = false
isnumber(x::Basic) = isnumber(BasicType(x))

iseven(x::Basic) = isnumber(x) && iseven(N(x))
isodd(x::Basic) = isnumber(x) && isodd(N(x))

ispositive(x::Union{BasicType{Val{:Integer}}, BasicType{Val{:Rational}}}) = x > 0
ispositive(::BasicType) = false
ispositive(x::Basic) = ispositive(BasicType(x))

isinteger(x::BasicType{Val{:Integer}}) = true
isinteger(::BasicType) = false
isinteger(x::Basic) = isinteger(BasicType(x))

## Functions

arity(fn::Type{Val{:Symbol}}) = 0
arity(fn::T) where {T <: SymEngine.BasicNumber} = 0
arity(fn::Type{Val{:Pow}}) = 2
arity(fn::Type{Val{:Add}}) = 99
arity(fn::Type{Val{:Mul}}) = 99
arity(fn::Type{Val{T}}) where {T} = 1
arity(fn::Symbol) = arity(Val{fn})

iscommutative(::Type{Val{:Add}}) = true
iscommutative(::Type{Val{:Mul}}) = true
iscommutative(::Type{Val{T}}) where {T} = false
iscommutative(fn::Symbol) = iscommutative(Val{fn})

isassociative(::Type{Val{:Add}}) = true
isassociative(::Type{Val{:Mul}}) = true
isassociative(::Type{Val{T}}) where {T} = false
isassociative(fn::Symbol) = isassociative(Val{fn})

## Dictionary things
function compatible_substitutions(d::Dict, d1::Dict)
    for k in intersect(keys(d), keys(d1))
        d[k] == d1[k] || return false
    end
    return true
end

# filter for compatability then merge. Returns new Set, does not mutate
function checked_merge(theta::Set, sigma::Dict)
    if length(theta) >= 1
        cmpts = filter(d -> compatible_substitutions(d, sigma), theta)
        map(u->merge(sigma, u), cmpts)
    else
        Set((sigma, ))
    end
end

# what is this???
function checked_merge(theta::Set, thetap::Set)
    if length(theta) == 0
        return thetap
    end
    out = Set()
    for m in thetap
        for m1 in theta
            if compatible_substitutions(m, m1)
                push!(out, merge(m, m1))
            end
        end
    end
    out

end

## Implement matchpy things
# out type
U = Union{Basic, Vector{Basic}, Tuple{Symbol, Vector{Basic}}}
#aDict(args...) = Dict{Basic, U}(args...)
aDict(args...) = Dict{Basic, Any}(args...)
empty_set() = Set(Dict{Basic, Any}[])

const blank_dict = aDict()
const ↯ = (false, blank_dict)
const NO_MATCH = Set(Dict{Basic,Any}[])

## return (Bool, Dict{Basic,Basic})
## Syntactic -- not associative *or* commutative
## must match exactly
## returns a tuple `(success, sigmap)`
## A match, sigmap, extends sigma and satisfies p(sigmap...) = s
function syntactic_match(s, p, sigma=aDict())  # subject, pattern, sigma::Dict

    isconstantpattern(p) && return (p==s, blank_dict)

    if onlywildcard(p)
        sigmap = aDict(p=>s)
        val = compatible_substitutions(sigma, sigmap)
        return (val ? (val, merge(sigma, sigmap)) :  (val, sigma))
    end

    head(s) != head(p) && return (false,↯)

    s_args, p_args = args(s), args(p)
    length(s_args) != length(p_args) && return (false, ↯)

    for i in eachindex(s_args)
        success, sigmap = syntactic_match(s_args[i], p_args[i], sigma)

        !success && return  (false, ↯)
        !compatible_substitutions(sigmap, sigma) && return (false, ↯)

        merge!(sigma, sigmap)
    end
    return (true, sigma)
end

#
## ss an array of subjects
## p a pattern
## fa a symbol representing an *associative function*, :Add, :Mul. Use :nothing for nothing
## Theta a set of valid (partial) substitutions
## returns a sset
"""
    match_one_to_one

Match a subject with one term

Length `ss` must be 1 to match, unless`p` is a star or plus wildcard

* `ss` is a collection of subjects
* `p` a pattern
* `theta` a set of partial substitutions.
* `comm` if true, the ss values may be permuted to match
* `assoc` if true, wildcards are used like plusvariables,
* `fn` a function symbol of an associative function (e.g, `:Add` and `:Mul`.)

Returns a set extending theta.
"""
function match_one_to_one(ss, p, theta::Set=Set((aDict(),)),
                          comm=false,  fn=:nothing, assoc=fn==:nothing?false:isassociative(fn))

    n = length(ss)
    # p in F0
    if isconstantpattern(p)

        n == 1 && p == ss[1] && return checked_merge(theta, aDict())

    elseif onlywildcard(p) && !assoc        # regular variable pattern

        sigmap = Dict(p=>ss[1])
        n == 1 && return checked_merge(theta, sigmap)

    elseif iswildcard(p)        # sequence variable pattern?

        if assoc  && fn != :nothing ## associative function then we can combine
            sigmap = Dict(p=>call_fun(fn, ss))
        elseif assoc || !onlywildcard(p)
            sigmap = Dict(p=>ss)           # matches the rest e.g. a + x_ ~ a + b + c = a + (b+c), so x_ => b+c
        end

        (isstarwildcard(p) || n >= 1) && return checked_merge(theta, sigmap)

    elseif n == 1

        h = head(p)
        if h == head(ss[1])
            qas, pas = args(ss[1]), args(p)

            # really here we want to pass along only :Add and :Mul
            fap = isassociative(h) ? h : :nothing
            out = Set()
            for sigma in theta
                out = union(out, Set(match_sequence(qas, pas, sigma, Val{iscommutative(fap)}, fap)))
            end
            return out
        end
    end

    return NO_MATCH

end
##################################################

## If p is a pattern and s a subject, then a match is a
## dictionary sigma whose keys are wildcards such that
## p(sigma...) = s
##
## This function performs a partial substitution by sigma
function _psubs(patterns, sigma::Pair{Basic, Basic})
    subs.(patterns, sigma)
end

# [x_,x__,y_] subs {x_=>1, x__=>[2,3,4], y_=>[5,6]} --> [1, 2,3,4,5,6] flattened one level
function _psubs(patterns, sigma::Pair{Basic, Vector{T}}) where {T}
    k,v = sigma
    out = Basic[]
    for s in patterns
        if k in free_symbols(s)
            append!(out, [subs(s, k=>vi) for vi in v])
        else
            push!(out, s)
        end
    end
    out
end

function psubs(patterns, sigma::Dict)
    out = patterns
    for (k,v) in sigma
        out = _psubs(out, k=>v)
    end
    out
end

## reduce after partial substitutions. Return multiset
function reduce_subjects_patterns(subjects, patterns, sigma)
    mss = Multiset(collect(subjects))  # collect so tuples work
    ps = psubs(patterns, sigma)
    mps = Multiset(collect(ps))
    mss - mps, mps - mss
end

## is there an issue with patterns.(sigma) ~ subjects
function so_far_so_good(subjects, patterns, sigma)::Bool
    lhs, rhs = collect.(reduce_subjects_patterns(subjects, patterns, sigma))
    ## we are bad if there are non constants on right (as constants are matched up)
    nconsts = length(filter(isconstantpattern, rhs))
    nconsts > 0 && return false
    length(rhs) - nconsts == 0 && length(lhs) > 0 && return false
    ## we are bad if rhs has wildcards or star
    return true
end
##

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

## Helper functions
function remove_constant_patterns(subjects, patterns, sigma=Dict())
    lhs, rhs = reduce_subjects_patterns(subjects, patterns, sigma)
    collect(_setdiff(lhs, rhs)), collect(_setdiff(rhs, lhs))
end

# check that matched variables in sigma are goo
function check_matched_variables(subjects, patterns, sigma::Dict)
    if so_far_so_good(subjects, patterns, sigma)
        (sigma,)
    else
        ()
    end
end

### helpers for non-variable pattern check

function _ms(u)
    p,s,sigma = u
    hp, hs = head(p), head(s)
    hp != hs && return ()

    comm = Val{iscommutative(hp)}
    fn, assoc = isassociative(hp) ? (hp, true) : (:nothing, false)

    match_sequence( args(s), args(p), sigma,
                       comm, fn, assoc)
end

function _check_pp(p, ss, sigma=Dict())
    o = generator_chain((),
                        () -> (s for s in ss),
                        (s) -> _ms((p,s,sigma))
                        )

    o
end

## take expressions like `sin(x_)` and check
function check_non_variable_patterns(ss, ps, sigma)

    ss, ps = reduce_subjects_patterns(ss, ps, sigma)
    ps = collect(Iterators.filter(!iswildcard, collect(ps))) # compound expressions

    isempty(ps) && return (sigma,)

    theta = Set()
    op = generator_chain((sigma,),
                            ((sigma) -> _check_pp(ps[i], collect(ss), sigma) for i in eachindex(ps))...,
                            (sigma) -> sigma in theta ? () : (push!(theta, sigma); (sigma,))
                            )
    op
end

### --------------------------------------------------

## Check regular variables helper
function assign_wildcard(p, k, sks,svs, sigma)
    out = Set()
    for (i, (val, nk)) in enumerate(zip(sks, svs))
        nk < k && continue
        if haskey(sigma, p)
            sigma[p] == val || continue
        end
        svscopy = copy(svs)
        svscopy[i] -= k
        push!(out, (sks, svscopy, merge(sigma, Dict(p=>val))))
    end
    out
end

function check_regular_variables(subjects, patterns, sigma::Dict)

    ## we reduce subjects, patterns, then use our algorithm
    ss, ps = remove_constant_patterns(subjects, patterns, sigma)

    sks, svs =  tally(ss)
    pks, pvs = tally(filter(onlywildcard, ps))

    o = generator_chain((),
                        () -> ((sks, svs, sigma),),
                        (u-> assign_wildcard(p, v, u...) for (p,v) in zip(pks, pvs))...,
                        (u) -> ((u[3],))
                           )
    o
end

### --------------------------------------------------

### --------------------------------------------------

## Check for sequence variables

## return generator
function _check_sv_init(ss, ps, sigma, fn, assoc, allvars, svals, N)
    ss, ps = remove_constant_patterns(ss, ps, sigma)

    plusvs = Basic[]
    starvs = Basic[]
    for p in ps
        if onlypluswildcard(p) || (assoc && onlywildcard(p))
            push!(plusvs, p)
        elseif onlystarwildcard(p)
            push!(starvs, p)
        end
    end

    nplus = length(plusvs)
    if nplus > length(ss)
        return empty_set() # too many plus variables
    end

    nstars = length(starvs)
    _svals, scnts = tally(ss)
    plusvals, pluscnts = tally(plusvs)
    starvals, starcnts = tally(starvs)

    theta = empty_set()

    empty!(allvars)
    append!(allvars,vcat(plusvals, starvals))
    empty!(svals)
    append!(svals, _svals)
    empty!(N)
    push!(N, length(plusvals))

    ## nothing to do
    if isempty(allvars) # nothing to do
        if length(ss) > 0
            return empty_set() # can't eat up ss values
        else
            return Set((sigma,))
        end
    end

    if length(plusvals) > length(ss)
        return Set()
    end
    if length(ss) == 0
        sigmap = Dict(k=>(fn==:nothing ? Basic[] : call_fun(fn, Basic[])) for k in starvals)
        out = empty_set()
        if compatible_substitutions(sigma, sigmap)
            push!(out, merge(sigma, sigmap))
        end
        return out
    end

    cfs = vcat(pluscnts, starcnts)

    gens = (solve_linear_diophantine(cnt, cfs) for cnt in scnts)
    itr = gen_product(gens...)

    itr
end

function _allocate_sols(sols, sigma, fn, allvars, svals, N)

    d = aDict()
    star_ok = true

    for (i,var) in enumerate(allvars)
        this_many = [a[i] for a in sols]
        if i <= N[1] && sum(this_many) == 0
            star_ok=false
            break
        end
        val = Basic[]
        for k in eachindex(this_many)
            nvals = _repeat(svals[k], inner=this_many[k])
            append!(val, nvals)
        end

        # wrap in associative function
        d[var] = isassociative(fn) ? call_fun(fn, val) : val

    end

    !star_ok && return ()
    if compatible_substitutions(sigma, d)
        sigmap = merge(sigma, d)
        return (sigmap,)
    end
    return ()
end

### When matching [a,a,a,a,b,b,c] ~ [x__, x__, y___] we need to
## solve equations 4 = 2x + y where 4 is number of a's, x >=1 and y >= 0
## this has solutions (2,0), (1,2), and (0,4)
## 2 = 2x + y: (1,0), (0,2)
## 1 = 2x + y: (0,1)
## we then combine (2,0),(1,0), and (0,1) -> x__ -> [a,a,b], y__ -> [c]
## (2,0),(0,2),(0,1) -> x__ -> [a,a], y___ -> [b,b,c] ...
## sigma -> generator of sigmas...
##
## The code copies some of that in matchpy for solving linear diophantine equations
function check_sequence_variables(ss, ps, sigma,
                                  fn=:nothing, assoc=false)

    allvars= Basic[]  # not so pretty as we want to share values across chain
    svals = Basic[]
    N = Int[]

    o = generator_chain((),
                        () -> _check_sv_init(ss, ps, sigma, fn, assoc, allvars, svals, N),
                        (sols) -> _allocate_sols(sols, sigma, fn, allvars, svals, N))

    o
end

### --------------------------------------------------

## Assemble pieces for matching ss ~ ps with commutative or associative

abstract type AbstractMatchObject end
function Base.show(io::IO, o::AbstractMatchObject)
    if isempty(o)
        println(io, "Empty match object")
    else
        println(io, "Non-empty, iterable match object. Use `first` to see one, `collect` for all.")
    end
end

## iterator interface for MatchObjects
## These just make the interace nicer
struct MatchObject <: AbstractMatchObject
itr
end
Base.iteratorsize(itertype::MatchObject)= Base.SizeUnknown()
Base.start(o::MatchObject) = start(o.itr)

Base.done(o::MatchObject, st) = done(o.itr, st)
Base.next(o::MatchObject, st) = next(o.itr, st)

###  --------------------------------------------------

# allocate fixed sum values
function allocate_ks_to_vars(ks, ss, ps, sigma, fn, assoc)
    i,j=0,0 #1,1
    thetap = Set((deepcopy(sigma),))

    for pl in ps
        lsub = 1
        if ispluswildcard(pl) || (iswildcard(pl) && assoc)
            lsub = lsub + ks[1+j]
            if onlystarwildcard(pl)
                lsub -= 1
            end
            j += 1
        end
        ssp = ss[(i+1):(i+lsub)]
        fn, assoc = isassociative(fn) ? (fn, true) : (:nothing, assoc)

        thetap = match_one_to_one(ssp, pl, thetap, Val{false}, fn, assoc)
        if isempty(thetap)
            break
        end
        i += lsub
    end

    thetap
end

## non-commutative matching
function match_sequence(ss, ps, sigma, commutative::Type{Val{false}},
                        fn=:nothing, assoc=fn==:nothing?false:isassociative(fn))

    n, m = length(ss), length(ps)

    ## number of star variables (zero, onem or more). Here
    ## m - nstar number of patterns that must match, if more than n too many
    nstar, nplus = 0, 0
    for p in ps
        onlystarwildcard(p) && (nstar += 1)
        (onlypluswildcard(p) || (assoc && onlywildcard(p))) && (nplus += 1)
    end
    m - nstar > n && (return empty_set())  # XXX the plus vars must match
    nfree = n - m + nstar
    nseq = nstar + nplus        #  number of sequence variables

    if nseq == 0 && nfree > 0
        return empty_set()   # generator of substitutions
    end

    itr = generator_chain((),
                             () -> fixed_sum(nseq, nfree),
                             (ks) -> allocate_ks_to_vars(ks, ss, ps, sigma, fn, assoc))

    MatchObject(itr)
end

###  --------------------------------------------------

# commutative matching
function match_sequence(ss, ps, sigma::Dict,
                        commutative::Type{Val{true}},
                        fn=:nothing,  assoc=fn==:missing?false:isassociative(fn))

    ss, ps = remove_constant_patterns(ss, ps)
    for p in ps
        isconstantpattern(p) && return NO_MATCH
    end
    length(ps) == 0 && length(ss) > 0 && return NO_MATCH

    theta = Set((sigma,))

    if length(ps) == 0
        theta = checked_merge(theta, Set((aDict(),)))
        return theta
    end

    itr = generator_chain((),
                          () -> theta,
                          sigma -> check_matched_variables(ss, ps, sigma),                                                           sigma -> check_non_variable_patterns(ss, ps, sigma),
                          sigma -> check_matched_variables(ss, ps, sigma),
                          sigma -> assoc ? ((sigma,)) : check_regular_variables(ss, ps, sigma),
                          sigma -> check_sequence_variables(ss, ps, sigma, fn, assoc)
    )

    MatchObject(itr)
end

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

## Matching API
export matches

"""

    match(p::Basic, s::Basic, sigma)

Match subject `s` against pattern `p`.

This is syntactic matching.

An expression can be broken down into an abstract syntax tree.

For example
* `sin(x^2+1)`  has head `:Sin` and arguments `x^2+1`
* `x^2 + 1` has head `:Add` and arguments `x^2` and `1`. These arguments are ordered,
though how is not described here.
* `x^2` has head `:Pow` and arguments `x` and `2`
* `1` is a numeric value with no head and no arguments.

The expression `sin(x_)` would match `sin(x^2 + 1)` by matching `x_=>x^2 + 1`. Whereas
`x_^2` would match `x^2` with `x_=>x`. But `2^x_` would not match `x^2` as 2 does not match `x`.
The expression `x_^2 + 1` **might** match `x^2 + 1` -- but it *might* not. Why? `:Add` is commutative, but syntactic matching does not take that into account. Use `matches` for such considerations.

In a picture, the tree would look like
```
       sin(x^2 + 1)
             |
       (:Sin, x^2 + 1)
                  |
           (:Add, x^2, 1)
                  /    |
          (:Pow, x, 2)
                /    |
```
The expression `sin(x_)` matches ``(:Sin, x^2+1)`, so `x_ => x^2 + 1`

The expression `sin(x_ + 1)` *might* match `x_ => x^2` but might not. It depends on ordering, as this matching does not take into account associativity (how `1+x_` would match `1 + a + b` through `1 + (a+b)`) or commutivity.

Returns `(success, sigmap)`. If `!success`, `sigmap` may contain nonsense.
"""
function Base.match(p::Basic, s::Basic, sigma=aDict())
    syntactic_match(s, p, sigma)
end

"""
    matches(patterns, subjects, sigma=Dict(); commutative=false, associative=false)

Match a list of pattern terms agains a list of subject terms.

If `commutative=true` all permutations of subjects are considered.

If `associative=true` wildcards are treated as plus wildcards. (so `x_ + y_ ~ 1 + 2 + 3`)

Patterns have "wildcards" defined through a naming convention:

* variables with one trailing slash, e.g. `x_`, are regular wildcards. They will match one term in the syntax tree unless there is an associative element, in which case they are treated as "plus" wildcards.

* variables with two trailing slashes, e.g. `x__`, are plus wildcards. They will match one or more terms.

* variables with three trailing slashes, e.g., `x___`, are star wildcards. They will match zero, one or more terms.

A value `sigma` may be specified ahead of time. This means that any
match will agree with `sigma` and satisfy `subs.(patterns, sigma...) ~
subjects`. The default is no specification

The return value is a generator for all valid substitutions:

* use `isempty` to check if there are any matches extending sigma.

* use `first` to get a match, when present. This mutates the generator.

* use `collect` to get all matches. This mutates the generator.

Note: Operations which return a value (or more) mutate the generator. One
must use `deepcopy` to make a copy should such behaviour be unwelcome.

Examples:

```julia

using SymEngine, SymEngineExtras
@vars a b c x_ x__ x___ y_ y__ y___

o = matches([x_, y__], [a,b,c]) # associative=false,commutative=false are defaults
first(o)                        # `x_=>a, y__=>[b,c]`
o = matches([x_, y__], [a,b,c], associative=true)
collect(o)                      # `x_=>a, y__=>[b,c]` AND `x_=>[a,b], y__=>[c]`
o = matches([x_, y__], [a,b,c], commutative=true)
collect(o)                      # `x_=>a, y__=>[b,c]` AND `x_=>b, y__=>[a,c]` AND `x_=>c, y__=>[a,b]`
                                # That is, `x_` can match ₃C₁ diferent terms
o = matches([x_, y__], [a,b,c], commutative=true, associative=true)
collect(o)                      # 6 matches, as x_ can match ₃C₁ + ₃C₂ terms now

o = matches([x__, y__], [a,b,c])
collect(o)                      # 2 matches, x__=>[a],  x__=>[a,b] (with y__ absorbing the rest)
o = matches([x__, y__], [a,b,c], commutative=true)
collect(o)                      # 6 matches as x__ can match ₃C₁ + ₃C₂ terms

o = matches([x___, y_], [a,b,c])
collect(o)                      # 1 match, as `y_=>[c]`
o = matches([x___, y__], [a,b,c])
collect(o)                      # 3 matches: ₃C₁ matches, as `x___` matches  `x___=>[], x___=>[a]`, and `x___=>[a,b]`
o = matches([x___, y__], [a,b,c], commutative=true)
collect(o)                      # 7 = ₃C₀ + ₃C₁ + ₃C₂ matches for `x___`, as y__ must have 1 match
```
"""
function matches(ps, ss, sigma=aDict(); commutative=false, associative=false)
    out = match_sequence(ss, ps, sigma, Val{commutative}, :nothing, associative)
end

"""
    matches(pattern, subject, sigma=Dict())

Matches subject against the pattern.

"""
function matches(pattern::Basic, subject::Basic, sigma=aDict())
    # depends on head
    sh, ph = head.((subject,pattern))

    if sh == ph && arity(sh) != 0

        match_sequence(args(subject), args(pattern), sigma,
                          Val{iscommutative(sh)},
                          sh, isassociative(sh))
    else
        ## when is it okay that p ~ s but of different heads?
        ## e.g x_ + y___ ~ a with {x_=>a, y___=>[]} --> {x=>a, y=>0}
        if isassociative(ph) && arity(sh) > 0 # Add or Mul
            match_sequence([subject], args(pattern), aDict(),
                           Val{iscommutative(ph)},
                           ph, isassociative(ph))
        else
            matches([pattern], [subject], sigma)
        end
    end
end

## Add in guards?

## Case - not needed
# cases(ss, p) = filter(s -> !isempty(matches(x_^2, s)), [a,a^2,b^2])

## Replace all

mutable struct SyntaxTreeNode
fn::Symbol
ex::Basic
children::Vector
end

function expression_tree(ex::Basic)
    fn = head(ex)

    if fn in [:Symbol, :Integer]
        return SyntaxTreeNode(fn, ex,  SyntaxTreeNode[]) # use no children for terminal
    else
        exs = args(ex)
        children = expression_tree.(args(ex)) #[expression_tree(u) for u in exs]
        SyntaxTreeNode(fn, ex, children)
    end
end

function node_to_ex(n::SyntaxTreeNode)
    if isempty(n.children)
        n.ex
    else
        call_fun(n.fn, node_to_ex.(n.children))
    end
end

## rewrite :Sin -> :Cos, no arity check... so :Add -> :Sin is an issue

## replace function heads
function _replace(n::SyntaxTreeNode, f::Symbol, g::Symbol, traverse=true)
    if n.fn == f
        n.fn = g
    end
    if traverse
        !isempty(n.children) && _replace.(n.children, f, g, traverse)
    end
end

## replace function head with functino call
function _replace(n::SyntaxTreeNode, f::Symbol, g, traverse=true)
    if traverse && !isempty(n.children)
        _replace.(n.children, f, g, traverse)
    end

    if n.fn == f
        xs = node_to_ex.(n.children)
        n.ex = g(xs...)
        n.fn = :Symbol
        empty!(n.children)
    end
end

## replace call f on expression. If true, replace with g(args)
function _replace(n::SyntaxTreeNode, f, g, traverse=true)
    if traverse && !isempty(n.children)
        _replace.(n.children, f, g, traverse)
    end

    val = f(n.ex)
    if val
        n.ex = g(n.ex)
        n.fn = :Symbol
        empty!(n.children)
    end
end

## replace expressions
function _replace(n::SyntaxTreeNode, pattern::Basic, rpattern::Basic, traverse=true)

    if traverse && !isempty(n.children)
        _replace.(n.children, pattern, rpattern, traverse)
    end

    ex = node_to_ex(n)
    m = matches(pattern, ex)
    val = isempty(m)

    if !val
        sigma = first(m)
        n.ex = subs(ex, subs(pattern, sigma...), subs(rpattern, sigma...))
        # now trick node
        n.fn = :Symbol
        empty!(n.children)
    end

end

## replace
import Base: replace
"""

    replace(ex::Basic, pattern, rpattern, traverse=true)

Replace parts o expression `ex` using a pattern and a replacement
pattern. If `traverse=true`, then descend syntax tree for
replacements. Otherwise, only the top-level is considered.

The replacement dependson the type of pattern/rppattern:

* `pattern::Basic, rpattern::Basic`

    The pattern with wildcards is matched agains the (sub) expressions of `ex`. When there is a match, the *first* substitution found is used to substitute in the (sub) expression using this match.

* `pattern::Symbol, rpattern::Symbol`

    The symbols represent function heads. For example, `sin(x)` has head `:Sin`. The syntax tree is looked at and each occurrence of thefunction head `f` is replaced by the function head `g`.

* `pattern::Symbol, rpattern::function`

    The symbol represents a function head. For example, `sin(x)` has head `:Sin`. The syntax tree is looked at and each occurrence of the function head `f` is replaced by the function `g(args...)`.

* `pattern::function, rpattern::function`

        Each subexpression `ex` is tested by the `pattern` function. If true, the expression is replaced by `g(ex)`.

"""
function replace(ex::Basic, pattern, rpattern, traverse=true)
    n = expression_tree(ex)
    _replace(n, pattern, rpattern, traverse)
    node_to_ex(n)
end

replace(ex::Basic, pr::Pair, traverse=true) = replace(ex, pr..., traverse)


text:
# Julia wrapper for header: gdal_vrt.h
# Automatically generated using Clang.jl

"""
    VRTCreate(int nXSize,
              int nYSize) -> VRTDatasetH
"""
function vrtcreate(arg1, arg2)
    aftercare(ccall((:VRTCreate, libgdal), VRTDatasetH, (Cint, Cint), arg1, arg2))
end

"""
    VRTFlushCache(VRTDatasetH hDataset) -> void
"""
function vrtflushcache(arg1)
    aftercare(ccall((:VRTFlushCache, libgdal), Cvoid, (VRTDatasetH,), arg1))
end

"""
    VRTSerializeToXML(VRTDatasetH hDataset,
                      const char * pszVRTPath) -> CPLXMLNode *
"""
function vrtserializetoxml(arg1, arg2)
    aftercare(ccall((:VRTSerializeToXML, libgdal), Ptr{CPLXMLNode}, (VRTDatasetH, Cstring), arg1, arg2))
end

"""
    VRTAddBand(VRTDatasetH hDataset,
               GDALDataType eType,
               char ** papszOptions) -> int
"""
function vrtaddband(arg1, arg2, arg3)
    aftercare(ccall((:VRTAddBand, libgdal), Cint, (VRTDatasetH, GDALDataType, Ptr{Cstring}), arg1, arg2, arg3))
end

"""
    VRTAddSource(VRTSourcedRasterBandH hVRTBand,
                 VRTSourceH hNewSource) -> CPLErr
"""
function vrtaddsource(arg1, arg2)
    aftercare(ccall((:VRTAddSource, libgdal), CPLErr, (VRTSourcedRasterBandH, VRTSourceH), arg1, arg2))
end

"""
    VRTAddSimpleSource(VRTSourcedRasterBandH hVRTBand,
                       GDALRasterBandH hSrcBand,
                       int nSrcXOff,
                       int nSrcYOff,
                       int nSrcXSize,
                       int nSrcYSize,
                       int nDstXOff,
                       int nDstYOff,
                       int nDstXSize,
                       int nDstYSize,
                       const char * pszResampling,
                       double dfNoDataValue) -> CPLErr
"""
function vrtaddsimplesource(arg1, arg2, arg3, arg4, arg5, arg6, arg7, arg8, arg9, arg10, arg11, arg12)
    aftercare(ccall((:VRTAddSimpleSource, libgdal), CPLErr, (VRTSourcedRasterBandH, GDALRasterBandH, Cint, Cint, Cint, Cint, Cint, Cint, Cint, Cint, Cstring, Cdouble), arg1, arg2, arg3, arg4, arg5, arg6, arg7, arg8, arg9, arg10, arg11, arg12))
end

"""
    VRTAddComplexSource(VRTSourcedRasterBandH hVRTBand,
                        GDALRasterBandH hSrcBand,
                        int nSrcXOff,
                        int nSrcYOff,
                        int nSrcXSize,
                        int nSrcYSize,
                        int nDstXOff,
                        int nDstYOff,
                        int nDstXSize,
                        int nDstYSize,
                        double dfScaleOff,
                        double dfScaleRatio,
                        double dfNoDataValue) -> CPLErr
"""
function vrtaddcomplexsource(arg1, arg2, arg3, arg4, arg5, arg6, arg7, arg8, arg9, arg10, arg11, arg12, arg13)
    aftercare(ccall((:VRTAddComplexSource, libgdal), CPLErr, (VRTSourcedRasterBandH, GDALRasterBandH, Cint, Cint, Cint, Cint, Cint, Cint, Cint, Cint, Cdouble, Cdouble, Cdouble), arg1, arg2, arg3, arg4, arg5, arg6, arg7, arg8, arg9, arg10, arg11, arg12, arg13))
end

"""
    VRTAddFuncSource(VRTSourcedRasterBandH hVRTBand,
                     VRTImageReadFunc pfnReadFunc,
                     void * pCBData,
                     double dfNoDataValue) -> CPLErr
"""
function vrtaddfuncsource(arg1, arg2, arg3, arg4)
    aftercare(ccall((:VRTAddFuncSource, libgdal), CPLErr, (VRTSourcedRasterBandH, VRTImageReadFunc, Ptr{Cvoid}, Cdouble), arg1, arg2, arg3, arg4))
end