text:
""" Structs and methods for dynamics integrators.
    Standard (non-damped) Landau-Lifshitz dynamics integrates:
        dS/dt = -S × B,
    where `B = -∇ H` is the effective field.
"""

abstract type Integrator end

"""
    HeunP(sys::SpinSystem)

Integrates Landau-Lifshitz spin dynamics using the Heun method, with a final
projection step that exactly constrains |S|=1. The method is locally second
order accurate.
"""
mutable struct HeunP <: Integrator
    sys :: SpinSystem
    _S₁ :: Array{Vec3, 4}
    _S₂ :: Array{Vec3, 4}
    _B  :: Array{Vec3, 4}
    _f₁ :: Array{Vec3, 4}

    function HeunP(sys::SpinSystem)
        return new(
            sys, zero(sys.sites), zero(sys.sites),
            zero(sys.sites), zero(sys.sites)
        )
    end
end

"""
    LangevinHeunP(sys, kT, α)

Implements Langevin dynamics on `sys` targetting a temperature `kT`, with a
damping coefficient `α`. Provided `α` should not be normalized by the spin
magnitude -- this is done internally.

Uses the 2nd-order Heun + projection scheme."
"""
mutable struct LangevinHeunP <: Integrator
    α   :: Float64        # Damping coeff
    D   :: Float64        # Stochastic strength
    sys :: SpinSystem
    _S₁ :: Array{Vec3, 4} # Intermediate integration variable space
    _S₂ :: Array{Vec3, 4}
    _B  :: Array{Vec3, 4}
    _f₁ :: Array{Vec3, 4}
    _r₁ :: Array{Vec3, 4}
    _ξ  :: Array{Vec3, 4}

    function LangevinHeunP(sys::SpinSystem, kT::Float64, α::Float64)
        return new(
            α, α*kT/(1+α*α), sys,
            zero(sys.sites), zero(sys.sites), zero(sys.sites),
            zero(sys.sites), zero(sys.sites), zero(sys.sites)
        )
    end
end

"""
    SphericalMidpoint(sys::SpinSystem; atol=1e-12)

Integrates Landau-Lifshitz spin dynamics using the spherical-midpoint
integrator, which is symplectic and implicit. Each step is converged to absolute
tolerance `atol`.
"""
mutable struct SphericalMidpoint <: Integrator
    sys :: SpinSystem
    _S̄  :: Array{Vec3, 4}
    _Ŝ  :: Array{Vec3, 4}
    _S̄′ :: Array{Vec3, 4}
    _B  :: Array{Vec3, 4}
    atol :: Float64

    function SphericalMidpoint(sys::SpinSystem; atol=1e-12)
        return new(
            sys, zero(sys.sites), zero(sys.sites),
            zero(sys.sites), zero(sys.sites), atol
        )
    end
end

@inline f(S, B) = -S × B
@inline f(S, B, α) = -S × (B + α * (S × B))

"""
    evolve!(integrator, Δt)

Performs a single integrator timestep of size Δt.
"""
function evolve!(integrator::HeunP, Δt::Float64)
    @unpack sys, _S₁, _S₂, _B, _f₁ = integrator
    S = sys.sites

    # Euler step
    field!(_B, S, sys.hamiltonian)
    @. _f₁ = f(S, _B)
    @. _S₁ = S + Δt * _f₁

    # Corrector step
    field!(_B, _S₁, sys.hamiltonian)
    @. _S₂ = normalize(S + 0.5 * Δt * (_f₁ + f(_S₁, _B)))

    # Swap buffers
    sys.sites, integrator._S₂ = integrator._S₂, sys.sites
    nothing
end

function evolve!(integrator::LangevinHeunP, Δt::Float64)
    @unpack α, D, sys, _S₁, _S₂, _B, _f₁, _r₁, _ξ = integrator
    S = sys.sites

    randn!(_ξ)
    _ξ .*= √(2D)
    # Normalize fluctuations by 1/√S
    for b in 1:nbasis(sys)
        selectdim(_ξ, 1, b) .*= 1. / sqrt(sys.sites_info[b].S)
    end

    # Euler step
    field!(_B, S, sys.hamiltonian)
    @. _f₁ = f(S, _B, α)
    @. _r₁ = f(S, _ξ, α)
    @. _S₁ = S + Δt * _f₁ + √Δt * _r₁

    # Corrector step
    field!(_B, _S₁, sys.hamiltonian)
    @. _S₂ = normalize(S + 0.5 * Δt * (_f₁ + f(_S₁, _B, α)) + 0.5 * √Δt * (_r₁ + f(_S₁, _ξ, α)))

    # Swap buffers
    sys.sites, integrator._S₂ = integrator._S₂, sys.sites
    nothing
end

function evolve!(integrator::SphericalMidpoint, Δt::Float64)
    @unpack sys, _S̄, _Ŝ, _S̄′, _B, atol = integrator
    S = sys.sites

    # Initial guess for midpoint
    @. _S̄ = S

    max_iters = 100
    for iter in 1:max_iters
        # Integration step for current best guess of midpoint _S̄. Produces
        # improved midpoint estimator _S̄′.
        @. _Ŝ = normalize(_S̄)
        field!(_B, _Ŝ, sys.hamiltonian)
        @. _S̄′ = S + 0.5 * Δt * f(_Ŝ, _B)

        # Convergence is reached if every element of _S̄ and _S̄′ agree to
        # within tolerance.
        converged = true
        for i = 1:length(S)
            converged = converged && isapprox(_S̄[i], _S̄′[i]; atol)
        end

        # If converged, then we can return
        if converged
            # Normalization here should not be necessary in principle, but it
            # could be useful in practice for finite `atol`.
            @. S = normalize(2*_S̄′ - S)
            return
        end

        @. _S̄ = _S̄′
    end

    error("Spherical midpoint method failed to converge to tolerance $atol after $max_iters iterations.")
end

"""
    A sampler which produces new samples using Langevin Landau-Lifshitz dynamics
"""
mutable struct LangevinSampler <: AbstractSampler
    integrator :: LangevinHeunP
    kT         :: Float64
    Δt         :: Float64
    nsteps     :: Int
end

"""
    LangevinSampler(sys, kT, α, Δt, nsteps)

Creates a `LangevinSampler` which samples the spin system's Hamiltonian using Langevin
 dynamics at a temperature `kT`, damping coefficient `α`, and producing a new sample
 by integrating with `nsteps` timesteps of size `Δt`.
"""
function LangevinSampler(sys::SpinSystem, kT::Float64, α::Float64, Δt::Float64, nsteps::Int)
    integrator = LangevinHeunP(sys, kT, α)
    LangevinSampler(integrator, kT, Δt, nsteps)
end

@inline function set_temp!(sampler::LangevinSampler, kT::Float64)
    sampler.kT = kT
    α = sampler.integrator.α
    sampler.integrator.D = α * kT / (1 + α * α)
    nothing
end

get_temp(sampler::LangevinSampler) = sampler.kT
get_system(sampler::LangevinSampler) = sampler.integrator.system

@inline function sample!(sampler::LangevinSampler)
    for _ in 1:sampler.nsteps
        evolve!(sampler.integrator, sampler.Δt)
    end
end


text:
@testset "nonmaxima supression" begin
    @info "Test: NonmaximaSuppression"

    @testset "Numerical" begin
        img = Gray{N0f8}.(load("algorithms/References/circle.png"))
        nms_ref = Gray{N0f8}.(load("algorithms/References/circle_nms.png"))

        f = NonmaximaSuppression()

        g₁, g₂ = imgradients(img, KernelFactors.scharr)
        mag = hypot.(g₁, g₂)

        nms = zeros(eltype(mag), axes(mag))

        f(nms, mag, g₁, g₂)

        for i in CartesianIndices(nms_ref)
            nms_ref[i] == 0.0 ? (@test nms[i] == 0) : (@test nms[i] == mag[i])
        end
    end
end

ext: jl


text:
using XLSX, CSV, DataFrames

function main()
  inscritos = CSV.read("inscritos.csv", DataFrame)
  n = size(inscritos, 1)
  presenca = [inscritos[!, ["Email Address", "Nome Completo"]] DataFrame(horas = zeros(Int, n))]
  xlsx = XLSX.readxlsx("presença.xlsx")
  pesos = [3, 2, 3, 2, 4, 2, 4, 2, 2, 4]
  perdidos = []
  for i in 1:XLSX.sheetcount(xlsx)
    sh = xlsx[i]
    emails = unique(sort(lowercase.(sh["B"][2:end])))
    I = indexin(emails, presenca[!, "Email Address"])
    M = findall(I .=== nothing)
    append!(perdidos, emails[M])
    I = I[findall(I .!== nothing)]
    presenca[I, "horas"] .+= pesos[i]
  end
  I = findall(presenca.horas .> 0)
  CSV.write("participantes.csv", presenca[I,:])
  df = DataFrame(email = unique(sort(perdidos)))
  CSV.write("faltantes.csv", df)
end

main()

max_stars_repo_name: CiDAMO/scripts-certificados

max_issues_repo_name: CiDAMO/scripts-certificados

max_forks_repo_name: CiDAMO/scripts-certificados


text:
export transient_polynomial, transient_mean, transient_variance, transient_covariance_ellipsoid

@doc raw"""
	transient_pop(MP::MarkovProcess, μ0::Dict, v::APL, d::Int, trange::AbstractVector{<:Real}, solver)

returns SOS program of degree `d` for computing a **lower** bound on ``\mathbb{E}[v(x(T))]``
where ``v`` is a polynomial and ``x(T)`` the state of the Markov process `MP` at time `T = trange[end]`.
μ0 encodes the distribution of the initial state of the process in terms of its moments;
specifically, it maps monomials to the respective moments of the initial distribution.
`trange` is an **ordered** collection of time points used to discretize the time
horizon ``[0,T]``, i.e., `T = trange[end]`. Populating `trange` and increasing `d` has a
tightening effect on the bound furnished by the assembled SOS program.
"""
function transient_pop(MP::MarkovProcess, μ0::Dict, p::APL, d::Int,
                       trange::AbstractVector{<:Real}, solver, P::Partition = trivial_partition(MP.X))
    if trange[1] == 0
        trange = trange[2:end]
    end
    t = MP.iv
    nT = length(trange)
    Δt = vcat(trange[1], [trange[i] - trange[i-1] for i in 2:nT])
    T  = @set(t >= 0 && t <= 1)

    model = SOSModel(solver)
    w = Dict((k, v) => (props(P.graph, v)[:cell] isa Singleton ?
                    @variable(model, [1], Poly(monomials(t, 0:d)))[1] :
                    @variable(model, [1], Poly(monomials(sort(vcat(MP.x, t), rev = true), 0:d)))[1]) for v in vertices(P.graph), k in 1:nT)

    for v in vertices(P.graph)
        add_dynamics_constraints!(model, MP, v, P, props(P.graph, v)[:cell], T, Δt, w, 0)
    end

    for v in vertices(P.graph)
        add_transversality_constraints!(model, MP, props(P.graph, v)[:cell], w[nT,v], p, v)
    end

    for e in edges(P.graph)
        add_coupling_constraints!(model, MP, e, P, T, Δt, w)
    end
	@objective(model, Max, expectation(μ0[first(keys(μ0))] isa Dict ? [subs(w[1, v], t => 0) for v in vertices(P.graph)] : subs(w[1, 1], t => 0), μ0))
    return model, w
end

function transient_pop(MP::MarkovProcess, μ0::Dict, p::Num, d::Int,
					   trange::AbstractVector{<:Real}, solver, P::Partition = trivial_partition(MP.X))
	if μ0[first(keys(μ0))] isa Dict
		μ0 = Dict(v => Dict(polynomialize_expr(mono, MP.poly_vars) => μ0[v][mono] for mono in keys(μ0[v])) for v in keys(μ0))
	else
		μ0 = Dict(polynomialize_expr(mono, MP.poly_vars) => μ0[mono] for mono in keys(μ0))
	end
	return transient_pop(MP, μ0, polynomialize_expr(p, MP.poly_vars), d, trange, solver, P)
end

@doc raw"""
	transient_polynomial(MP::MarkovProcess, μ0::Dict, v::APL, d::Int, trange::AbstractVector{<:Real}, solver)

returns a **lower** bound on ``\mathbb{E}[v(x(T))]`` where ``v`` is a polynomial and ``x(T)`` the state
of the Markov process `MP` at time `T = trange[end]`. `μ0` encodes the
distribution of the initial state of the process in terms of its moments; specifically,
it maps monomials to the respective moments of the initial distribution.
`trange` is an **ordered** collection of time points used to discretize the time
horizon ``[0,T]``, i.e., `T = trange[end]`. Populating `trange` and increasing `d` improves
the computed bound.
"""
function transient_polynomial(MP::MarkovProcess, μ0::Dict, p::APL, d::Int, trange::AbstractVector{<:Real}, solver, P::Partition = trivial_partition(MP.X))
    model, w = transient_pop(MP, μ0, p, d, trange, solver, P)
    optimize!(model)
    return Bound(objective_value(model), model, P, dual_poly(w, MP.iv, trange))
end

function transient_polynomial(MP::MarkovProcess, μ0::Dict, p::Num, d::Int, trange::AbstractVector{<:Real}, solver, P::Partition = trivial_partition(MP.X))
	if μ0[first(keys(μ0))] isa Dict
		μ0 = Dict(v => Dict(polynomialize_expr(mono, MP.poly_vars) => μ0[v][mono] for mono in keys(μ0[v])) for v in keys(μ0))
	else
		μ0 = Dict(polynomialize_expr(mono, MP.poly_vars) => μ0[mono] for mono in keys(μ0))
	end
	return transient_polynomial(MP, μ0, polynomialize_expr(p, MP.poly_vars), d, trange, solver, P)
end

@doc raw"""
	transient_mean(MP::MarkovProcess, μ0::Dict, x::APL, d::Int, trange::AbstractVector{<:Real}, solver)

returns a **lower** and **upper** bound on ``\mathbb{E}[v(x(T))]`` where ``v`` is a polynomial and ``x(T)`` the state
of the Markov process `MP` at time `T = trange[end]`. `μ0` encodes the
distribution of the initial state of the process in terms of its moments; specifically,
it maps monomials to the respective moments of the initial distribution.
trange is an **ordered** collection of time points used to discretize the time
horizon ``[0,T]``, i.e., `T = trange[end]`. Populating `trange` and increasing `d` improves
the computed bounds.
"""
function transient_mean(MP::MarkovProcess, μ0::Dict, p::APL, d::Int, trange::AbstractVector{<:Real}, solver, P::Partition = trivial_partition(MP.X))
    lb = transient_polynomial(MP, μ0, p, d, trange, solver)
    ub = transient_polynomial(MP, μ0, -p, d, trange, solver)
	ub.value *= -1
    return lb, ub
end

function transient_mean(MP::MarkovProcess, μ0::Dict, p::Num, d::Int, trange::AbstractVector{<:Real}, solver, P::Partition = trivial_partition(MP.X))
	if μ0[first(keys(μ0))] isa Dict
		μ0 = Dict(v => Dict(polynomialize_expr(mono, MP.poly_vars) => μ0[v][mono] for mono in keys(μ0[v])) for v in keys(μ0))
	else
		μ0 = Dict(polynomialize_expr(mono, MP.poly_vars) => μ0[mono] for mono in keys(μ0))
	end
	return transient_mean(MP, μ0, polynomialize_expr(p, MP.poly_vars), d, trange, solver, P)
end

"""
	transient_mean(rn::ReactionSystem, S0::Dict, S, d::Int, trange::AbstractVector{<:Number}, solver,
			scales = Dict(s => 1 for s in species(rn));
			auto_scaling = false)

returns a **lower** and **upper** bound on the mean of the molecular count of species
`S` in reaction network `rn` at time `T = trange[end]`. `S0` refers to the **deterministic** initial state
of the reaction system (including all species!).
`trange` is an **ordered** collection of time points used to discretize the time
horizon ``[0,T]``, i.e., `T = trange[end]`. Populating trange and increasing `d` improves
the computed bounds.

For numerical stability, it is recommended to provide scales of the expected
magnitude of molecular counts for the different species at steady state. If
the system is **closed** it is also possible to enable auto_scaling which will
find the maximum molecular counts for each species under stoichiometry
constraints (via LP).
"""
function transient_mean(rn::ReactionSystem, S0::Dict, S, d::Int, trange::AbstractVector{<:Number}, solver,
		  				scales = Dict(s => 1 for s in species(rn));
						params::Dict = Dict(), auto_scaling = false)
	RP, S0 = reaction_process_setup(rn, S0, scales = scales, auto_scaling = auto_scaling, solver = solver, params = params)
	μ0 = init_moments(RP.JumpProcess.x, S0, d + maximum(maxdegree.(RP.JumpProcess.a)))
 	return transient_mean(RP.JumpProcess, μ0, RP.species_to_state[S], d, trange, solver)
end

@doc raw"""
	transient_variance(MP::MarkovProcess, μ0::Dict, v::APL, d::Int, trange::AbstractVector{<:Real}, solver)

returns an **upper** bound on ``\mathbb{E}[v(x(T))^2] - \mathbb{E}[v(x(T))]^2`` where ``v`` is a polynomial
and ``x(T)`` the state of the Markov process `MP` at time `T = trange[end]`.
"""
function transient_variance(MP::MarkovProcess, μ0::Dict, p::APL, d::Int, trange::AbstractVector{<:Real}, solver, P::Partition = trivial_partition(MP.X))
	if trange[1] == 0
		trange = trange[2:end]
	end
	t = MP.iv
	nT = length(trange)
	Δt = vcat(trange[1], [trange[i] - trange[i-1] for i in 2:nT])
	T  = @set(t >= 0 && t <= 1)

	model = SOSModel(solver)
	w = Dict((k, v) => (props(P.graph, v)[:cell] isa Singleton ?
                    @variable(model, [1], Poly(monomials(t, 0:d)))[1] :
                    @variable(model, [1], Poly(monomials(sort(vcat(MP.x, t), rev = true), 0:d)))[1]) for v in vertices(P.graph), k in 1:nT)
	@variable(model, S[1:2])
	@constraint(model, [1 S[1]; S[1] S[2]] in PSDCone())

	for v in vertices(P.graph)
		add_dynamics_constraints!(model, MP, v, P, props(P.graph, v)[:cell], T, Δt, w, 0)
	end

	for v in vertices(P.graph)
		add_transversality_constraints!(model, MP, props(P.graph, v)[:cell], w[nT, v], -(p^2 + 2*S[1]*p), v)
	end

	for e in edges(P.graph)
		add_coupling_constraints!(model, MP, e, P, T, Δt, w)
	end

	@objective(model, Max, expectation(μ0[first(keys(μ0))] isa Dict ? [subs(w[1, v], t => 0) for v in vertices(P.graph)] : subs(w[1, 1], t => 0), μ0) - S[2])
	optimize!(model)
	return Bound(-objective_value(model), model, P, dual_poly(w, t, trange))
end

function transient_variance(MP::MarkovProcess, μ0::Dict, p::Num, d::Int, trange::AbstractVector{<:Real}, solver, P::Partition = trivial_partition(MP.X))
	if μ0[first(keys(μ0))] isa Dict
		μ0 = Dict(v => Dict(polynomialize_expr(mono, MP.poly_vars) => μ0[v][mono] for mono in keys(μ0[v])) for v in keys(μ0))
	else
		μ0 = Dict(polynomialize_expr(mono, MP.poly_vars) => μ0[mono] for mono in keys(μ0))
	end
	return transient_variance(MP, μ0, polynomialize_expr(p, MP.poly_vars), d, trange, solver, P)
end

"""
	transient_variance(rn::ReactionSystem, S0::Dict, S, d::Int, trange::AbstractVector{<:Real}, solver,
			    scales = Dict(s => 1 for s in species(rn));
				auto_scaling = false)

returns an **upper** bound on the variance of species `S` in the reaction network `rn`
at time `T = trange[end]`. `S0` refers to the **deterministic** initial state
of the reaction system (including all species!).
`trange` is an **ordered** collection of time points used to discretize the time
horizon ``[0,T]``, i.e., `T = trange[end]`. Populating `trange` and increasing `d` improves
the computed bounds.

For numerical stability, it is recommended to provide scales of the expected
magnitude of molecular counts for the different species at steady state. If
the system is **closed** it is also possible to enable `auto_scaling` which will
find the maximum molecular counts for each species under stoichiometry
constraints (via LP).
"""
function transient_variance(rn::ReactionSystem, S0::Dict, S, d::Int, trange::AbstractVector{<:Real}, solver,
						    scales = Dict(s => 1 for s in species(rn));
							auto_scaling = false, params::Dict = Dict())
	RP, S0 = reaction_process_setup(rn, S0, scales = scales, auto_scaling = auto_scaling, solver = solver, params = params)
	μ0 = init_moments(RP.JumpProcess.x, S0, d + maximum(maxdegree.(RP.JumpProcess.a)))
 	return transient_variance(RP.JumpProcess, μ0, RP.species_to_state[S], d, trange, solver)
end

@doc raw"""
	transient_covariance_ellipsoid(MP::MarkovProcess, μ0::Dict, v::Vector{APL}, d::Int, trange::AbstractVector{<:Real}, solver)

returns an **upper** bound on the volume of the covariance ellipsoid ``\text{det}(\mathbb{E}[v(x(T))v(x(T))^\top] - \mathbb{E}[v(x(T))] \mathbb{E}[v(x(T))]^\top)``,
where ``v`` is a polynomial and ``x(T)`` the state of the Markov process `MP` at time `T = trange[end]`.
"""
function transient_covariance_ellipsoid(MP::MarkovProcess, μ0::Dict, p::Vector{<:APL}, d::Int, trange::AbstractVector{<:Real}, solver, P::Partition = trivial_partition(MP.X))
	if trange[1] == 0
		trange = trange[2:end]
	end
	t = MP.iv
	nT = length(trange)
	Δt = vcat(trange[1], [trange[i] - trange[i-1] for i in 2:nT])
	T  = @set(t >= 0 && t <= 1)
	n = length(p)

	model = SOSModel(solver)
	w = Dict((k, v) => (props(P.graph, v)[:cell] isa Singleton ?
                    @variable(model, [1], Poly(monomials(t, 0:d)))[1] :
                    @variable(model, [1], Poly(monomials(sort(vcat(MP.x, t), rev = true), 0:d)))[1]) for v in vertices(P.graph), k in 1:nT)
	@variable(model, S[1:n+1,1:n+1], PSD)
 	@variable(model, U[1:2n,1:2n], PSD)
 	@variable(model, r[1:n])
    @variable(model, q[1:n])

	@constraint(model, [1 S[1]; S[1] S[2]] in PSDCone())
	@constraint(model, S[1:n,1:n] .== U[1:n, 1:n])
    @constraint(model, [i in 1:n, j in 1:i], -2*U[n+i,j] - (i == j ? U[n+i,n+i] - r[i] : 0) == 0)
    @constraint(model, [i in 1:n], [-1, q[i], r[i]] in MOI.DualExponentialCone())

	for v in vertices(P.graph)
		add_dynamics_constraints!(model, MP, v, P, props(P.graph, v)[:cell], T, Δt, w, 0)
	end

	for v in vertices(P.graph)
		add_transversality_constraints!(model, MP, props(P.graph, v)[:cell], w[nT, v], -(p'*S[1:n, 1:n]*p + 2*S[end,1:n]'*p), v)
	end

	for e in edges(P.graph)
		add_coupling_constraints!(model, MP, e, P, T, Δt, w)
	end

	@objective(model, Max, expectation(μ0[first(keys(μ0))] isa Dict ? [subs(w[1, v], t => 0) for v in vertices(P.graph)] : subs(w[1, 1], t => 0), μ0) - S[n+1,n+1] - sum(q))
	optimize!(model)
	return Bound(exp(-objective_value(model)), model, P, dual_poly(w, t, trange))
end

function transient_covariance_ellipsoid(MP::MarkovProcess, μ0::Dict, v::Vector{Num}, d::Int, trange::AbstractVector{<:Real}, solver, P::Partition = trivial_partition(MP.X))
	if μ0[first(keys(μ0))] isa Dict
		μ0 = Dict(v => Dict(polynomialize_expr(mono, MP.poly_vars) => μ0[v][mono] for mono in keys(μ0[v])) for v in keys(μ0))
	else
		μ0 = Dict(polynomialize_expr(mono, MP.poly_vars) => μ0[mono] for mono in keys(μ0))
	end
	return transient_covariance_ellipsoid(MP, μ0, polynomialize_expr(v, MP.poly_vars), d, trange, solver, P)
end

"""
	transient_covariance_ellipsoid(rn::ReactionSystem, S0::Dict, S::AbstractVector, d::Int, trange::AbstractVector{<:Real}, solver,
							scales = Dict(s => 1 for s in species(rn));
							auto_scaling = false)

returns an **upper** bound on the volume of the covariance ellipsoid associated with any
collection of chemical species in the reaction network `rn` at time `T = trange[end]`.
S0 refers to the **deterministic** initial state of the reaction system
(including all species!). `trange` is an **ordered** collection of time points used to discretize the time
horizon ``[0,T]``, i.e., `T = trange[end]`. Populating `trange` and increasing `d` improves
the computed bounds.

For numerical stability, it is recommended to provide scales of the expected
magnitude of molecular counts for the different species at steady state. If
the system is **closed** it is also possible to enable auto_scaling which will
find the maximum molecular counts for each species under stoichiometry
constraints (via LP).
"""
function transient_covariance_ellipsoid(rn::ReactionSystem, S0::Dict, S::AbstractVector, d::Int, trange::AbstractVector{<:Real}, solver,
										scales = Dict(s => 1 for s in species(rn));
										auto_scaling = false, params::Dict = Dict())
	RP, x0 = reaction_process_setup(rn, S0, scales = scales, auto_scaling = auto_scaling, solver = solver, params = params)
	μ0 = init_moments(RP.JumpProcess.x, x0, d + maximum(maxdegree.(RP.JumpProcess.a)))
 	return transient_covariance_ellipsoid(RP.JumpProcess, μ0, [RP.species_to_state[x] for x in S], d, trange, solver)
end

# distributed
function transient_indicator(MP::MarkovProcess, μ0::Dict, v_target::Int, d::Int,
							trange::AbstractVector{<:Real}, solver, P::Partition = trivial_partition(MP.X); sense = 1)
	if trange[1] == 0
		trange = trange[2:end]
	end
	t = MP.iv
	nT = length(trange)
	Δt = vcat(trange[1], [trange[i] - trange[i-1] for i in 2:nT])
	T  = @set(t >= 0 && t <= 1)

	model = SOSModel(solver)
	w = Dict((k, v) => (props(P.graph, v)[:cell] isa Singleton ?
						@variable(model, [1], Poly(monomials(t, 0:d)))[1] :
						@variable(model, [1], Poly(monomials(sort(vcat(MP.x, t), rev = true), 0:d)))[1]) for v in vertices(P.graph), k in 1:nT)

	for v in vertices(P.graph)
		add_dynamics_constraints!(model, MP, v, P, props(P.graph, v)[:cell], T, Δt, w, 0)
	end

	for v in vertices(P.graph)
		flag = (v_target == v)
		add_transversality_constraints!(model, MP, props(P.graph, v)[:cell], w[nT,v], sense*flag, v)
	end

	for e in edges(P.graph)
		add_coupling_constraints!(model, MP, e, P, T, Δt, w)
	end
	@objective(model, Max, expectation(μ0[first(keys(μ0))] isa Dict ? [subs(w[1, v], t => 0) for v in vertices(P.graph)] : subs(w[1, 1], t => 0), μ0))

	optimize!(model)
	return Bound(objective_value(model), model, P, dual_poly(w, t, trange))
end
# add transient probability mass

function approximate_transient_measure(MP::MarkovProcess, μ0::Dict, p::APL, d::Int,
									   trange::AbstractVector{<:Real}, solver, P::Partition = trivial_partition(MP.X))
	if trange[1] == 0
		trange = trange[2:end]
	end
	t = MP.iv
	nT = length(trange)
	Δt = vcat(trange[1], [trange[i] - trange[i-1] for i in 2:nT])
	T  = @set(t >= 0 && t <= 1)

	model = SOSModel(solver)
	w = Dict((k, v) => (props(P.graph, v)[:cell] isa Singleton ?
						@variable(model, [1], Poly(monomials(t, 0:d)))[1] :
						@variable(model, [1], Poly(monomials(sort(vcat(MP.x, t), rev = true), 0:d)))[1]) for v in vertices(P.graph), k in 1:nT)

	for v in vertices(P.graph)
		add_dynamics_constraints!(model, MP, v, P, props(P.graph, v)[:cell], T, Δt, w, 0)
	end

	cons = Dict()
	for v in vertices(P.graph)
		cons[v] = add_transversality_constraints!(model, MP, props(P.graph, v)[:cell], w[nT,v], p, v)
	end

	for e in edges(P.graph)
		add_coupling_constraints!(model, MP, e, P, T, Δt, w)
	end
	@objective(model, Max, expectation(μ0[first(keys(μ0))] isa Dict ? [subs(w[1, v], t => 0) for v in vertices(P.graph)] : subs(w[1, 1], t => 0), μ0))

	optimize!(model)
	dist = []
    for v in vertices(P.graph)
        if props(P.graph, v)[:cell] isa Singleton
            push!(dist, dual(cons[v]).a[end])
        elseif props(P.graph, v)[:cell] isa Vector{BasicSemialgebraicSet}
            push!(dist, sum(dual.(cons[v])).a[end])
        elseif props(P.graph, v)[:cell] isa BasicSemialgebraicSet
            push!(dist, dual(cons[v]).a[end])
        end
    end
	return Bound(objective_value(model), model, P, dual_poly(w, t, trange)), dist
end
"""
    max_entropy_measure(MP::MarkovProcess, μ0::Dict, d::Int, trange::AbstractVector{<:Real}, solver, P::Partition;
                        side_infos::BasicSemialgebraicSet)

returns approximate values for the transient measure which maximizes the entropy on the partition `P` at time t.
"""
function max_entropy_measure(MP::MarkovProcess, μ0::Dict, d::Int, trange::AbstractVector{<:Real}, solver, P::Partition)
	if trange[1] == 0
		trange = trange[2:end]
	end
	t = MP.iv
	nT = length(trange)
	Δt = vcat(trange[1], [trange[i] - trange[i-1] for i in 2:nT])
	T  = @set(t >= 0 && t <= 1)

	model = SOSModel(solver)
	w = Dict((k, v) => (props(P.graph, v)[:cell] isa Singleton ?
						@variable(model, [1], Poly(monomials(t, 0:d)))[1] :
						@variable(model, [1], Poly(monomials(sort(vcat(MP.x, t), rev = true), 0:d)))[1]) for v in vertices(P.graph), k in 1:nT)

	for v in vertices(P.graph)
		add_dynamics_constraints!(model, MP, v, P, props(P.graph, v)[:cell], T, Δt, w, 0)
	end

	cons = Dict()
	@variable(model, u[vertices(P.graph)])
	@variable(model, q[vertices(P.graph)])
	for v in vertices(P.graph)
		cons[v] = add_transversality_constraints!(model, MP, props(P.graph, v)[:cell], w[nT,v], -q[v], v)
	end

	for e in edges(P.graph)
		add_coupling_constraints!(model, MP, e, P, T, Δt, w)
	end

	@constraint(model, [v in vertices(P.graph)], [-1, q[v], u[v]] in MOI.DualExponentialCone())

	@objective(model, Max, expectation(μ0[first(keys(μ0))] isa Dict ? [subs(w[1, v], t => 0) for v in vertices(P.graph)] : subs(w[1, 1], t => 0), μ0)
						   - sum(u))

	optimize!(model)
	dist = []
    for v in vertices(P.graph)
        if props(P.graph, v)[:cell] isa Singleton
            push!(dist, dual(cons[v]).a[end])
        elseif props(P.graph, v)[:cell] isa Vector{BasicSemialgebraicSet}
            push!(dist, sum(dual.(cons[v])).a[end])
        elseif props(P.graph, v)[:cell] isa BasicSemialgebraicSet
            push!(dist, dual(cons[v]).a[end])
        end
    end
	return Bound(objective_value(model), model, P, dual_poly(w, t, trange)), dist
end


text:
"""Filtration of an abstract complex

We call this sequence of complexes the **filtration** of `f` and
think of it as a construction by adding chunks of simplices at a time `t::FI`.
∅ = K0 ⊆ K1 ⊆ . . . ⊆ Kn = K.
"""
mutable struct Filtration{C<:AbstractComplex, FI}
    # underlying abstract cell complex
    complex::C
    # total order of simplices as array of (dim, simplex id, filtation value)
    total::Vector{Tuple{Int,Int,FI}}
    divisions::Number
end
Base.show(io::IO, flt::Filtration{C,FI}) where {C <: AbstractComplex, FI} = print(io, "Filtration($(complex(flt)), $FI)")
Base.valtype(flt::Filtration{C,FI}) where {C <: AbstractComplex, FI} = FI

Base.complex(flt::Filtration) = flt.complex
order(flt::Filtration) = flt.total
Base.minimum(flt::Filtration) = order(flt)[1][3]
Base.maximum(flt::Filtration) = order(flt)[end][3]
#
# Constructors
#
Filtration(::Type{C}, ::Type{FI}) where {C <: AbstractComplex, FI} =
    Filtration(C(), Vector{Tuple{Int,Int,FI}}(), Inf)
Base.similar(flt::Filtration{C,FI}) where {C <: AbstractComplex, FI} = Filtration(C, FI)

"""Construct filtration from a cell complex using the order of their appearence in the complex"""
function filtration(cplx::AbstractComplex)
    idx = Vector{Tuple{Int,Int,Int}}()
    i = 1
    for d in 0:dim(cplx)
        for c in cells(cplx, d)
            push!(idx, (dim(c), c[:index], i))
            i += 1
        end
    end
    return Filtration(cplx, idx, Inf)
end

"""Construct filtration from a cell complex and a complex weight function"""
function filtration(cplx::C, w::Dict{Int,Vector{FI}}; divisions::Number=Inf) where {C<:AbstractComplex, FI}
    idx = Vector{Tuple{Int,Int,FI}}()
    for d in 0:dim(cplx)
        for c in cells(cplx, d)
            ci = c[:index]
            push!(idx, (d, ci, w[d][ci]))
        end
    end
    sort!(idx, by=x->(x[3], x[1])) # sort by dimension & filtration value
    return Filtration(cplx, idx, divisions)
end

function Base.push!(flt::Filtration{C,FI}, cl::AbstractCell, v::FI; recursive=false) where {C<:AbstractComplex, FI}
    cplx = complex(flt)
    @assert isa(cl, celltype(cplx)) "Complex $(cplx) does not accept $(typeof(cl))"
    cls = push!(cplx, cl, recursive=recursive)
    ord = order(flt)
    idx = length(ord) == 0 ? 1 : findlast(e->e[3]<=v, ord)
    for c in sort!(cls, by=s->dim(s))
        if idx == length(ord)
            push!(ord, (dim(c), c[:index], v))
        else
            insert!(ord, idx, (dim(c), c[:index], v))
        end
        idx += 1
    end
    return flt
end

"""Generate a combined boundary matrix from the filtration `flt` for the persistent homology calculations."""
function boundary_matrix(flt::Filtration; reduced=false)
    ridx = reduced ? 1 : 0
    # initialize boundary matrix
    cplx = complex(flt)
    bm = map(i->BitSet(), 1:sum(size(cplx))+ridx)
    # fill boundary matrix
    ord = order(flt)
    for (i, (d, ci, fv)) in enumerate(ord)
        if d > 0
            splx = cplx[ci, d]
            for face in faces(splx)
                fi = cplx[face, d-1]
                push!(bm[i+ridx], findfirst(e->e[1] == d-1 && e[2] == fi, ord)+ridx)
            end
        elseif reduced
            push!(bm[i+ridx], 1)
        end
    end
    return bm
end

function SparseArrays.sparse(∂::Vector{BitSet})
    m = length(∂)
    ret = spzeros(Int, m, m)
    for i in 1:m
        bm = ∂[i]
        for (l, j) in enumerate(bm)
            ret[j,i] = j # (-1)^((l-1)%2) # coefs require exact order of faces in provides simplex
        end
    end
    return ret
end

#
# I/O
#

function Base.write(io::IO, flt::Filtration)
    cplx = complex(flt)
    for (d, ci, fv) in order(flt)
        for k in cplx[ci,d][:values]
            write(io, "$k,")
        end
        write(io, "$fv\n")
    end
end

function Base.read(io::IO, ::Type{Filtration{C,FI}}) where {C <: AbstractComplex, FI}
    flt = Filtration(C,FI)
    ET = eltype(celltype(complex(flt)))
    while !eof(io)
        l = readline(io)
        vals = split(l, ',')
        svals = map(v->parse(ET, v), vals[1:end-1])
        fval = parse(FI, vals[end])
        push!(flt, Simplex(svals), fval)
    end
    return flt
end

"Write a combined boundary matrix to a text file"
function writeboundarymatrix(io::IO, bm::Vector, zeroindex = true)
    for smplx in bm
        if length(smplx) == 0
            write(io, "0")
        else
            write(io, "$(length(smplx)-1)")
        end
        for i in smplx
            write(io, " $(zeroindex ? i-1 : i)")
        end
        write(io, 0x0A)
    end
end

#
# Iterators
#

# Filtration simplicial complex iterator
Base.length(flt::Filtration) = isinf(flt.divisions) ? length(unique(e->e[3], order(flt))) : flt.divisions
Base.eltype(flt::Filtration{C,FI}) where {C <: AbstractComplex, FI} = C

"""Loop through the filtration `flt` producing growing simplicial complexes on every iteration"""
function Base.iterate(flt::Filtration, state=(1, -Inf, 0))
    ord = order(flt)
    idx = state[1]
    idx > length(ord) && return nothing # done
    if state[2] == -Inf # calculate initial state
        fval = ord[idx][3]
        incr = (maximum(flt)-minimum(flt)) / flt.divisions
    else
        fval = state[2]
        incr = state[3]
    end
    splxs = Tuple{Int,Int}[] #simplex dim & index
    while idx <= length(ord) && (fval+incr) >= ord[idx][3]
        push!(splxs, ord[idx][1:2])
        idx += 1
    end
    nextfval = fval+incr
    if idx <= length(ord) && isinf(flt.divisions)
        nextfval = ord[idx][3]
    end
    return (fval, splxs), (idx, nextfval, incr)
end

# Filtration simplex iterator
Base.length(splxs::Simplices{<:Filtration}) = length(splxs.itr)
Base.eltype(splxs::Simplices{<:Filtration}) = celltype(splxs.itr)

function Base.iterate(splxs::Simplices{<:Filtration},state=nothing)
    # call underlying iterator
    if state === nothing
        res = iterate(splxs.itr)
    else
        res = iterate(splxs.itr, state)
    end
    # final state
    res == nothing && return nothing
    # get complex
    cplx = complex(splxs.itr)
    ss = [cplx[i, d] for (d, i) in res[1][2]]
    return (res[1][1], ss), res[2] # state of filtration iterator
end


text:
export postproc1d, postproc2d

function postproc1d!(fwdsol::FWDSol)
    @unpack gb, pp, pm, btilde, fsl = fwdsol
    qtot1d, qd1d = postproc1d(gb, pp, pm,
                              btilde, fsl)
    fwdsol.misc[:qtot1d] = qtot1d
    fwdsol.misc[:qd1d] = qd1d
    nothing
end
"1D post-processing"
function postproc1d(gb, pp, pm,
                    btilde1d=gb.bdot-gb.dhdt,
                    fsl=gb.fsl
                    )
    qtot1d = flux_n(gb, btilde1d)
    qd1d = flux_deform(qtot1d, on_nodes(fsl), pp.r)
    return qtot1d, qd1d
end

"2D post-processing"
function postproc2d()
end

"""
2D errors of a field vs a measurement field of gl (:h, :iv)

See also iv2d_err.
"""
function error2d(h_or_iv::Matrix, gl, field, mask=nothing; fillval=NaN)
    obs_ = getfield(gl, field).vals
    samesize = size(obs_.v)==size(h_or_iv)
    obs  = samesize ? obs_ :
        Interp.interpolate((obs_.x, obs_.y), obs_.v, Interp.Gridded(Interp.Linear()) )[gl.dem.x, gl.dem.y]
    out = obs.v .- h_or_iv
    if mask==nothing
        return out
    else
        out[.!mask] = fillval
        return out
    end
end

"""
    error_h(hs2d::Matrix, gl)

Error along the radar trajectories: model - measurement

Returns the trajectories with the errors on them:

err, bed_model, surf, radar_bed, thick_model

See also compare_to_radar, `thick_err`.
"""
function error_h(hs2d::Matrix, gl, rlines_to_drop=[], in_glacier=in_glacier(gl))
    tr = gl.h.vals
    hi = Interp.interpolate((gl.dem.x, gl.dem.y), hs2d, Interp.Gridded(Interp.Linear()) )
    si = Interp.interpolate((gl.dem.x, gl.dem.y), gl.dem.v, Interp.Gridded(Interp.Linear()) )
    err = deepcopy(gl.h.vals)
    bed = deepcopy(gl.h.vals)
    radar_bed = deepcopy(gl.h.vals)
    surf = deepcopy(gl.h.vals)
    thick = deepcopy(gl.h.vals)
    for i in eachindex(tr.v)
        x,y = tr.x[i], tr.y[i]
        if in_glacier[x,y]==1 # only inside glacier
            err_ = hi[x,y] - tr.v[i]
            err.v[i] = err_
            err.err[i] = 0 # hi[tr.x[i],tr.y[i]] # if the interpolation values are also desired
            bed.v[i] = si[x,y] - hi[x,y]
            bed.err[i] = 0
            surf.v[i] = si[x,y]
            surf.err[i] = 0
            radar_bed.v[i] = si[x,y] - tr.v[i]
            radar_bed.err[i] = 0
            thick.v[i] = hi[x,y]
            thick.err[i] = 0
        else
            # otherwise drop the point
            err.v[i] = NaN
            err.err[i] = NaN
            bed.v[i] = NaN
            bed.err[i] = NaN
            surf.v[i] = NaN
            surf.err[i] = NaN
            radar_bed.v[i] = NaN
            radar_bed.err[i] = NaN
            thick.v[i] = NaN
            thick.err[i] = NaN
        end
    end
    return err, bed, surf, radar_bed, thick
end


text:
module RulesBIFMZprevTest

using Test
using ReactiveMP
using Random
import ReactiveMP: @test_rules

@testset "rules:BIFM:zprev" begin

     @testset "Belief Propagation: (m_out::MultivariateNormalDistributionsFamily, m_in::MultivariateNormalDistributionsFamily, m_znext::MultivariateNormalDistributionsFamily, meta::BIFMMeta)" begin

        meta = BIFMMeta([2.0 0; 0 1], # A
                        [3.0 0; 0 2], # B
                        [4.0 0; 0 3], # C
                        [5.0 0; 0 4], # H
                        [6.0 0; 0 5], # BHBt
                        [1.0, 2],     # ξz
                        [7.0 0; 0 6], # Λz
                        [2.0, 3],     # ξztilde
                        [8.0 0; 0 7], # Λztilde
                        [3.0, 4],     # μu
                        [9.0 0; 0 8]  # Σu
                        )

        @test_rules [ with_float_conversions = true, float32_atol = 1e-2, bigfloat_atol = 1e-8 ] BIFM(:zprev, Marginalisation) [
            (input = (m_out = MvNormalMeanPrecision([1,2], [2 0; 0 1]),     m_in = MvNormalMeanPrecision([1,2], [1 0; 0 2]), m_znext = MvNormalMeanPrecision([1,2], [1 0; 0 2]), meta = meta),
                output = MvNormalWeightedMeanPrecision([-0.60402684563758, -1.4782608695652169], [0.44295302013420557 0.0; 0.0 0.4782608695652171])),
            (input = (m_out = MvNormalMeanPrecision([3,4], [6 0; 0 1]),     m_in = MvNormalMeanPrecision([1,6], [2 0; 0 2]), m_znext = MvNormalMeanPrecision([8,2], [2 0; 0 2]), meta = meta),
                output = MvNormalWeightedMeanPrecision([-0.9321266968325688, -5.043478260869566],[0.886877828054339 0.0; 0.0 0.4782608695652171])),
            (input = (m_out = MvNormalMeanPrecision([5,6], [2 0; 0 5]),     m_in = MvNormalMeanPrecision([6,2], [1 0; 0 1]), m_znext = MvNormalMeanPrecision([1,9], [1 0; 0 1]), meta = meta),
                output = MvNormalWeightedMeanPrecision([-3.7114093959731633, -0.45945945945943834], [0.44295302013420557 0.0; 0.0 0.24864864864865122]))
        ]

    end

end

end


text:
# AUV msg slam client model

using Distributed
# addprocs(8)

using Pkg
Pkg.activate(@__DIR__)
pkg"instantiate"
pkg"precompile"

@everywhere begin
    using Pkg
    Pkg.activate(@__DIR__)
end

using Caesar, RoME, DistributedFactorGraphs
@everywhere using Caesar, RoME, DistributedFactorGraphs
using LCMCore, BotCoreLCMTypes
using Dates
using DataStructures
using Logging
using JSON2
using DSP
# using JLD2, FileIO

# mute LCM broadcasting beyond this computer
run(`$(ENV["HOME"])/mutelcm.sh`)

include(joinpath(@__DIR__, "CommonUtils.jl"))

global ODOSCALE = [1.0; 1.0; 1.0]

# bring required utilities and handlers into context
include(joinpath(@__DIR__, "MsgHandlers.jl"))
include(joinpath(@__DIR__, "SandsharkUtils.jl"))

function main(;parsed_args=parse_commandline(),
               lcm=LCM(),
               logSpeed::Float64=parsed_args["speed"],
               iters::Int=parsed_args["iters"],
               dbg::Bool=parsed_args["dbg"],
               savePlotting::Bool=parsed_args["savePlotting"]  )
  #

  # data containers
  dashboard = Dict{Symbol,Any}()
  lblDict = Dict{DateTime,Vector{Float64}}()
  magDict = Dict{DateTime,Float64}()

  # fg object and initialization
  fg = LightDFG{SolverParams}(params=SolverParams())
  # fg = initfg()
  initializeAUV_noprior(fg, dashboard, stride_range=parsed_args["stride_range"], stride_solve=parsed_args["fixedlag"])

  # scale the odometry noise
  dashboard[:odoCov] .*= parsed_args["kappa_odo"]
  dashboard[:savePlotting] = savePlotting

  # store args, dead reckon tether, and lbl values in files.
  Base.mkpath(getLogPath(fg))
  # write the arguments to file
  argsfile = open(joinLogPath(fg,"args.txt"),"w")
  println(argsfile, ARGS)
  println(argsfile, parsed_args)
  close(argsfile)
  argsfile = open(joinLogPath(fg,"args.json"),"w")
  JSON2.write(argsfile, parsed_args)
  close(argsfile)
  DRTLog = open(joinLogPath(fg,"DRT.csv"),"w")
  LBLLog = open(joinLogPath(fg,"LBL.csv"),"w")
  Odolog = open(joinLogPath(fg,"ODO.csv"),"w")
  dirodolog = open(joinLogPath(fg,"DIRODO.csv"),"w")
  rawodolog = open(joinLogPath(fg,"RAWODO.csv"),"w")
  allthtlog = open(joinLogPath(fg,"ALLTHT.csv"),"w")
  posetiminglog = open(joinLogPath(fg,"timing_pose.csv"),"w")
  solvetiminglog = open(joinLogPath(fg,"timing_solve.csv"),"w")

  # prepare the solver in the background
  defaultFixedLagOnTree!(fg, parsed_args["fixedlag"], limitfixeddown=parsed_args["limitfixeddown"])
  ST = manageSolveTree!(fg, dashboard, dbg=dbg, timinglog=solvetiminglog, limitfixeddown=parsed_args["limitfixeddown"])

  dshfile = open(joinLogPath(fg,"dashboard_start.json"),"w")
  JSON2.write(dshfile, dashboard)
  close(dshfile)

  # middleware handlers
  # start with LBL and magnetometer
  subscribe(lcm, "AUV_LBL_INTERPOLATED", (c,d)->lbl_hdlr(c,d,fg,dashboard,lblDict,LBLLog), pose_t)
  subscribe(lcm, "AUV_MAGNETOMETER",     (c,d)->mag_hdlr(c,d,fg,dashboard,magDict), pose_t)

  @info "get a starting position and mag orientation"
  while length(lblDict) == 0 || length(magDict) == 0
    handle(lcm)
  end
  lblKeys = keys(lblDict) |> collect |> sort
  magKeys = keys(magDict) |> collect |> sort
  initPose = [lblDict[lblKeys[1]];magDict[magKeys[1]]]
  dashboard[:odoTime] = lblKeys[1]
  @show startT = lblKeys[1]

  setTimestamp!(fg, :x0, lblKeys[1])
  # setTimestamp!(getVariable(fg, :x0), lblKeys[1])

  # add starting prior
  addFactor!(fg, [:x0;], PriorPose2(MvNormal(initPose,Matrix(Diagonal([0.5; 0.5; 1.0].^2)))), autoinit=false)
  setSolvable!(fg, :x0, 1)
  doautoinit!(fg,:x0)

  @info "Start with the real-time tracking aspect..."
  odoChannel = parsed_args["odoGyroBias"] ? "AUV_ODOMETRY_GYROBIAS" : "AUV_ODOMETRY"
  subscribe(lcm, odoChannel,             (c,d)->pose_hdlr(c,d,fg,dashboard, posetiminglog, DRTLog, Odolog, dirodolog, rawodolog, allthtlog), pose_t)
  # subscribe(lcm, "AUV_ODOMETRY_GYROBIAS", (c,d)->pose_hdlr(c,d,fg,dashboard), pose_t)
  subscribe(lcm, "AUV_RANGE_CORRL",      (c,d)->range_hdlr(c,d,fg,dashboard), raw_t)
  # subscribe(lcm, "AUV_BEARING_CORRL",callback, raw_t)

  # repeat single native odometry
  # subscribe(lcm, "AUV_ODOMETRY",         (c,d)->odo_ensure_hdlr(c,d,fg,dashboard,dirodolog), pose_t)

  recordWTDSH = [true;]
  WTDSH = Vector{Tuple{DateTime, Symbol, Int, HandlerStateMachine, SolverStateMachine, Millisecond, Int}}()
  dashboard[:taskSignals] = @async begin
    # mkpath(getLogPath(fg))
    open(joinLogPath(fg, "dash.csv"), "w") do io
      while recordWTDSH[1]
        line = [now() getLastPoses(fg,filterLabel=r"x\d",number=1)[1] dashboard[:poseStride] Int(isready(dashboard[:solveSettings].canTakePoses)) dashboard[:solveSettings].solveInProgress dashboard[:realTimeSlack].value length(dashboard[:solveSettings].poseSolveToken.data)]
        # push!(WTDSH, (line...))
        writedlm(io, line, ',')
        sleep(0.005)
      end
    end
  end
  # record trees for later reconstruction into an animation
  dashboard[:taskRecTrees] = @async begin
    Base.mkpath(joinLogPath(fg, "trees"))
    treepath = joinLogPath(fg, "bt.pdf")
    counter = 0
    tcfid = open(joinLogPath(fg, "trees/timestamps.log"),"w")
    while 0 < parsed_args["recordTrees"]
      if ispath(treepath)
        counter += 1
        println(tcfid, "$(dashboard[:lastMsgTime]), $(dashboard[:lastPose]), $counter")
        Base.cp(treepath, joinLogPath(fg, "trees/bt_$(counter).pdf") )
      end
      sleep(1.0/parsed_args["recordTrees"])
    end
    close(tcfid)
  end

  # run handler
  # dashboard[:solveSettings].canTakePoses = HSMHandling
  holdTimeRTS = now()
  offsetT = now() - startT
  dashboard[:doDelay] = lcm isa LCMLog
  for i in 1:iters
      handle(lcm)
      # slow down in case of using an LCMLog object
      if dashboard[:doDelay]
        # might be delaying the solver -- ie real time slack grows above zero
        holdTimeRTS = now()
        while 1 < length(dashboard[:solveSettings].poseSolveToken.data)  # !HSMCanContinue(dashboard)
            # while dashboard[:canTakePoses] == HSMBlocking && dashboard[:solveInProgress] == SSMSolving
          flush(DRTLog)
          @info "delay for solver, tokens=$(dashboard[:solveSettings].poseSolveToken.data), RTS=$(dashboard[:realTimeSlack])"
          # , canTakePoses=$(dashboard[:solveSettings].canTakePoses)
          # map(x->flush(x.stream), [DRTLog;LBLLog;Odolog;dirodolog;rawodolog;allthtlog;posetiminglog,solvetiminglog])
          sleep(0.2)
        end
        # real time slack update -- this is the compute overrun
        dashboard[:realTimeSlack] += now()-holdTimeRTS
        # playback speed
        dt = (dashboard[:lastMsgTime]-startT)
        nMsgT = startT + typeof(dt)(round(Int, dt.value/logSpeed)) + dashboard[:realTimeSlack]
        while now() < (nMsgT + offsetT)
          flush(DRTLog)
          sleep(0.01)
        end
      end
  end

  # close out solver and other loops
  parsed_args["recordTrees"] = -1.0
  recordWTDSH[1] = false
  dashboard[:solveSettings].loopSolver = false

  # close UDP listening on LCM other file handles
  close(lcm)
  close(DRTLog)
  close(LBLLog)
  close(Odolog)
  close(dirodolog)
  close(rawodolog)
  close(allthtlog)
  close(posetiminglog)
  close(solvetiminglog)

  # return last factor graph as built and solved
  return fg, dashboard, WTDSH, ST
end

# user intentions
parsed_args=parse_commandline()

@show ODOSCALE[1] *= parsed_args["scaleOdoX"]

# ## Uncomment for different defaults in Juno
# parsed_args["iters"] = 15000
# parsed_args["kappa_odo"] = 0.1
# parsed_args["speed"] = 1.0
# parsed_args["recordTrees"] = false

# maybe warmup
# Precompile large swath of the solver functions
if parsed_args["warmupjit"]
  warmUpSolverJIT()
  warmUpSolverJIT() # twice to ensure larger footprint for multiprocess
end

## Do the actual run -- on live UDP traffic
# fg = main()

## Do run on LCM data from file
# lcmlogfile = joinpath(ENV["HOME"],"data","sandshark","lcmlog","lcm-sandshark-med.log")
lcmlogfile = joinpath(ENV["HOME"],"data","sandshark","lcmlog","sandshark-long.lcmlog")
fg, dashboard, wtdsh, ST = main(parsed_args=parsed_args, lcm=LCMLog(lcmlogfile) )

# @save joinLogPath(fg, "dash.jld2") wtdsh
# fid = open(joinLogPath(fg, "dash.json"), "w").
# JSON2.write(fid, wtdsh)
# close(fid)

## moved to CommonUtils.jl
# In case any variables had not been solved yet
ensureAllInitialized!(fg)
saveDFG(fg, joinpath(getLogPath(fg),"fg_final") )

# draw plots in getLogPath(fg)
if parsed_args["genResults"]
  include(joinpath(@__DIR__, "GenerateResults.jl"))
end

## BATCH SOLVE

dontMarginalizeVariablesAll!(fg)
setSolvable!(getVariable(fg, :drt_ref), 0)
getSolverParams(fg).maxincidence = 1000
tree = solveTree!(fg)

saveDFG(fg, joinpath(getLogPath(fg),"fg_batchsolve") )

if parsed_args["genResults"]
  plb = plotSandsharkFromDFG(fg, drawTriads=false)
  plb |> PDF(joinpath(getLogPath(fg),"traj_batch.pdf"))
end

@show getLogPath(fg)

#


text:
@testset "TotalVariation" begin
    p, q = synthetic_gradient(SynthSphere(50))
    z = zeros(Float64, size(p))
    Z = ShapeFromShading.TotalVariation(z=z)(p, q)
    @test Z[38,38] ≈ -1.1185373026725053e90
    @test Z[1,1] ≈ -9.432843507802872e89
    @test Z[76,76] ≈ 3.637059020826754e90
end


text:
function JSONRPC.parse_params(::Type{Val{Symbol("\$/cancelRequest")}}, params)
    return CancelParams(params)
end

function process(r::JSONRPC.Request{Val{Symbol("\$/cancelRequest")},CancelParams}, server)
end

function JSONRPC.parse_params(::Type{Val{Symbol("\$/setTraceNotification")}}, params)
    return params
end

function process(r::JSONRPC.Request{Val{Symbol("\$/setTraceNotification")},Dict{String,Any}}, server)
end

function JSONRPC.parse_params(::Type{Val{Symbol("julia/lint-package")}}, params)
    return
end

function process(r::JSONRPC.Request{Val{Symbol("julia/lint-package")},Nothing}, server)
    for (uri, f) in server.documents
        @info basename(uri._uri), " ", f.code.index
    end
end

function JSONRPC.parse_params(::Type{Val{Symbol("julia/toggle-lint")}}, params)
    return TextDocumentIdentifier(params["textDocument"])
end

function process(r::JSONRPC.Request{Val{Symbol("julia/toggle-lint")},TextDocumentIdentifier}, server)
    doc = server.documents[URI2(r.uri)]
    doc._runlinter = !doc._runlinter
end

function JSONRPC.parse_params(::Type{Val{Symbol("julia/reload-modules")}}, params)
end

function process(r::JSONRPC.Request{Val{Symbol("julia/reload-modules")},Nothing}, server)
end

function JSONRPC.parse_params(::Type{Val{Symbol("julia/toggleFileLint")}}, params)
    return params
end

function process(r::JSONRPC.Request{Val{Symbol("julia/toggleFileLint")}}, server)
    path = get(r.params, "path", "")
    uri = get(r.params, "external", "")
    if isdir(uri2filepath(path))
        for doc in values(server.documents)
            uri2 = doc._uri
            server.debug_mode && @info "LINT: ignoring $path"
            if startswith(uri2, uri)
                toggle_file_lint(doc, server)
            end
        end
    else
        if uri in map(i->i._uri, values(server.documents))
            server.debug_mode && @info "LINT: ignoring $path"
            doc = server.documents[URI2(uri)]
            toggle_file_lint(doc, server)
        end
    end
end

function JSONRPC.parse_params(::Type{Val{Symbol("julia/toggle-log")}}, params)
end

function process(r::JSONRPC.Request{Val{Symbol("julia/toggle-log")},Nothing}, server)
    server.debug_mode = !server.debug_mode
end

function JSONRPC.parse_params(::Type{Val{Symbol("julia/getCurrentBlockOffsetRange")}}, params)
    return TextDocumentPositionParams(params)
end

function process(r::JSONRPC.Request{Val{Symbol("julia/getCurrentBlockOffsetRange")}}, server)
    if !haskey(server.documents, URI2(r.params.textDocument.uri))
        send(JSONRPC.Response(r.id, CancelParams(r.id)), server)
        return
    end
    tdpp = r.params
    doc = server.documents[URI2(tdpp.textDocument.uri)]
    offset = get_offset(doc, tdpp.position.line + 1, tdpp.position.character)
    stack, offsets = StaticLint.get_stack(doc.code.cst, offset)
    if length(stack) > 1 && stack[1] isa CSTParser.EXPR{CSTParser.FileH}
        if stack[2] isa  CSTParser.EXPR{CSTParser.ModuleH} && length(stack) > 3
            p1, p2, p3 = (offsets[4] + 1, offsets[4] + stack[4].span, offsets[4] + stack[4].fullspan)
        else
            p1, p2, p3 = (offsets[2] + 1, offsets[2] + stack[2].span, offsets[2] + stack[2].fullspan)
        end
    else
        p1 = p2 = p3 = length(doc._content)
    end

    response = JSONRPC.Response(r.id, (length(doc._content, 1, max(1, p1)), length(doc._content, 1, p2), length(doc._content, 1, p3)))

    send(response, server)
end

function JSONRPC.parse_params(::Type{Val{Symbol("julia/activateenvironment")}}, params)
end

function process(r::JSONRPC.Request{Val{Symbol("julia/activateenvironment")}}, server)
end


text:
function mk_method(
    trait     :: Any,
    interface :: MethodInterface,
    syms      :: Syms, # all typeclass parameters
    fn_deps   :: Exprs # all functional dependencies
) where {Syms <: AbstractArray{Symbol}, Exprs <: AbstractArray{Pair{Symbol, Expr}}}
    argsty     = interface.argsty
    retty      = interface.retty
    f          = interface.name
    infer      = interface.infer
    fresh      = interface.fresh

    # these typeclass parameters can be inferred directly by arguments
    infer_by_dispatch = intersect!(infer, syms)

    # these typeclass parameters must be able to inferred by functional dependencies,
    # otherwise raise error.
    infer_by_fn_deps  = setdiff!(Set(syms), infer_by_dispatch)
    help_infer = Expr[]
    for (sym, fn_dep) in fn_deps
        if sym in infer_by_fn_deps
            push!(help_infer, fn_dep)
            delete!(infer_by_fn_deps, sym)
        end
    end

    # remainder of typeclass parameters that cannot get inferred
    if !isempty(infer_by_fn_deps)
        cannot_infer_anyway = join(map(string, collect(infer_by_fn_deps)), ", ")
        error(
            "Cannot infer type param(s) ($cannot_infer_anyway) for method $f, trait $trait; "*
            "Please add functional dependencies."
        )
    end

    @when let :($tp{$(argtys...)}) = argsty,
              (tp === Tuple).?

        basename = gensym(string(f))
        argnames = [Symbol(basename, i) for i in 1:length(argtys)]
        annos    = [:($argname :: $argty) for (argname, argty) in zip(argnames, argtys)]
        @q begin
            function $f($(annos...)) where {$(infer_by_dispatch...), $(fresh...)}
                $(help_infer...)
                ($trait($(syms...)).$f($(argnames...)))::$retty
            end
        end
    @otherwise
        error("Malformed method $f($argsty) for trait $trait.")
    end
end

function mk_default_impl(tparams::TypeParams, default_method :: DefaultMethod)
    mk_name = Symbol("mk.", default_method.name)
    mkfn   = @q begin
        $mk_name($(tparams.as_arguments...)) where {$(tparams.as_where...)} = $(default_method.impl)
    end
    (default_method.name, mk_name, mkfn)
end

mk_trait(syms::Syms) where Syms <: AbstractArray{Symbol} =
    @match syms begin
    [] => TraitParamNil
    [hd, tl...] => :($TraitParamCons{$hd, $(mk_trait(tl))})
    end

"""
get_instance_type
"""
function instance end
function check_inheritance end
function implement end
function default_method_maker(trait, n)
    @assert n isa Val
    f(::Val{n}) where n = n
    error("No default method $(f(n)) for $trait.")
end

function trait(line::LineNumberNode, @nospecialize(sig), @nospecialize(block))
    fn_deps = Vector{Pair{Symbol, Expr}}()
    sups, sig = extract_trait_mk(sig)

    @when :($hd where {$(args...)}) = sig begin
        sig = hd
        if !isempty(args)
            foreach(args) do arg
                @when :($(f :: Symbol) = $_) = arg begin
                    push!(fn_deps, f => arg)
                @otherwise
                    error("Malform functional dependency $arg for trait $sig.")
                end
            end
        end
    end

    @when let :($trait_ty{$(tvars...)}) = sig,
              Expr(:block, stmts...)    = block

        tsyms         = map(extract_tvars, tvars)
        tparams       = TypeParams(tsyms)
        methods       = map(extract_method, stmts)
        methods       = filter(x -> !(x isa LineNumberNode), methods)
        interfaces    = filter(x -> x isa MethodInterface, methods)
        default_impls = filter(x -> x isa DefaultMethod, methods)

        cg_interfaces   = [mk_method(trait_ty, method_interface, tsyms, fn_deps)
                           for method_interface in interfaces]

        instance_name   = Symbol(string(trait_ty), "#instance")
        interface_names = map(interfaces) do x; x.name end
        method_types    = map(x -> Symbol(x, "#t"), interface_names)
        cg_fields       = [:($a :: $b) for (a, b) in zip(interface_names, method_types)]

        (make_default_makers,  ask_default_makers) =
            let triples = map(x -> mk_default_impl(tparams, x), default_impls)
                map(x->x[3], triples),
                map(triples) do (default_name, default_maker_name, _)
                    @q begin
                        $line
                        $CanonicalTraits.default_method_maker(
                            ::Type{$trait_ty},
                            ::Val{$(QuoteNode(default_name))}
                        ) = $default_maker_name
                        $line
                    end
                end

            end
        @q begin
            $line
            abstract type $sig <: $Trait{$(mk_trait(tsyms))} end
            struct $instance_name{$(method_types...)}
                $(cg_fields...)
            end
           $CanonicalTraits.check_inheritance(::Type{$trait_ty}, $(tsyms...)) = begin
               $(sups...)
               nothing
            end
            $CanonicalTraits.instance(::Type{$trait_ty}) = $instance_name
            $(cg_interfaces...)
            $(make_default_makers...)
            $(ask_default_makers...)
            $line
        end
    @otherwise
        error("Malformed trait $sig definition.")
    end
end

macro trait(@nospecialize(sig), block=Expr(:block))
    trait(__source__, sig, block) |> esc
end

(trait::Type{<:Trait})(types...) = begin
    CanonicalTraits.check_inheritance(trait, types...)
    types = join(map(string, types), ", ")
    error("Not implemented trait $trait for ($types).")
end


text:
"""

 █████╗  ██████╗██████╗ ██╗
██╔══██╗██╔════╝██╔══██╗██║
███████║██║     ██████╔╝██║
██╔══██║██║     ██╔══██╗██║
██║  ██║╚██████╗██║  ██║███████╗
╚═╝  ╚═╝ ╚═════╝╚═╝  ╚═╝╚══════╝

File:       PlanarQuadMPC.jl
Author:     Gabriel Barsi Haberfeld, 2020. gbh2@illinois.edu
Function:   This program implemetns an MPC controller for trajectory tracking of
            a planar quadrotor.

Instructions:   Run this file in juno with Julia 1.5.1 or later.
Requirements:   JuMP, Ipopt, Plots, LinearAlgebra, BenchmarkTools.

"""

using JuMP, Ipopt
using Plots, LinearAlgebra
gr()
m = 0.483
g = 9.81
I = 0.01532

function dynamics(x = 0.0 .* zeros(6), u = 0.0 .* zeros(2), dt = 1.0)
    #x = [px pz θ vx vz θ̇]
    #     1  2  3 4  5  6
    px = x[1]
    pz = x[2]
    θ = x[3]
    vx = x[4]
    vz = x[5]
    θ̇ = x[6]
    x0 = copy(x)
    x[1] = vx * cos(θ) - vz * sin(θ)
    x[2] = vx * sin(θ) + vz * cos(θ)
    x[3] = θ̇
    x[4] = vz * θ̇ - g * sin(θ)
    x[5] = -vx * θ̇ - g * cos(θ) + (u[1] + u[2]) / m
    x[6] = 0.0 + (u[1] - u[2]) / I
    x = x * dt + x0
end

function angelaDynamics(x = 0.0 .* zeros(6), u = 0.0 .* zeros(2), dt = 1.0)
    #x = [px pz θ vx vz θ̇]
    #     1  2  3 4  5  6
    px = x[1]
    pz = x[2]
    θ = x[3]
    vx = x[4]
    vz = x[5]
    θ̇ = x[6]
    x0 = copy(x)
    x[1] = vx
    x[2] = vz
    x[3] = θ̇
    x[4] = -1 / m * sin(θ)
    x[5] = 1 / m * cos(θ) * u[1] - g
    x[6] = 0.0 + u[2] / I
    x = x * dt + x0
end

function test_dynamics()
    dt = 0.1
    t = Array(0:0.1:1)
    N = length(t)
    x = 0.0 .* zeros(6)
    x[3] = 0.0
    u = 0.0 .* zeros(2)
    xv = 0.0 .* zeros(6, N)
    for i = 1:N
        u[1] = g
        u[2] = 0.1
        x = dynamics(x, u, dt)
        @show u, x
    end
end

function SimpleMPC(
    x0,
    xref = 0.0 .* zeros(6),
    θlim = pi / 4,
    θ̇lim = pi / 3,
    vxlim = 2,
    vzlim = 1,
    dt = 0.1,
    N = 30,
)
    MPC = Model(optimizer_with_attributes(
        Ipopt.Optimizer,
        "max_iter" => 35,
        "print_level" => 0,
    ))
    @variable(MPC, px[i = 0:N])
    @variable(MPC, pz[i = 0:N])
    @variable(MPC, -θlim <= θ[i = 0:N] <= θlim)
    @variable(MPC, -vxlim <= vx[i = 0:N] <= vxlim)
    @variable(MPC, -vzlim <= vz[i = 0:N] <= vxlim)
    @variable(MPC, -θ̇lim <= θ̇[i = 0:N] <= θ̇lim)
    @variable(MPC, 0 <= u1[i = 0:N])
    @variable(MPC, 0 <= u2[i = 0:N])
    @variable(MPC, β >= 0)

    @constraint(MPC, px[0] == x0[1])
    @constraint(MPC, pz[0] == x0[2])
    @constraint(MPC, θ[0] == x0[3])
    @constraint(MPC, vx[0] == x0[4])
    @constraint(MPC, vz[0] == x0[5])
    @constraint(MPC, θ̇[0] == x0[6])

    #x = [px pz θ vx vz θ̇]
    #     1  2  3 4  5  6
    for k = 0:N-1
        @NLconstraint(MPC, px[k+1] == px[k] + (vx[k] * cos(θ[k]) - vz[k] * sin(θ[k])) * dt)
        @NLconstraint(MPC, pz[k+1] == pz[k] + (vx[k] * sin(θ[k]) + vz[k] * cos(θ[k])) * dt)
        @constraint(MPC, θ[k+1] == θ[k] + θ̇[k] * dt)
        @NLconstraint(MPC, vx[k+1] == vx[k] + (vz[k] * θ̇[k] - g * sin(θ[k])) * dt)
        @NLconstraint(
            MPC,
            vz[k+1] == vz[k] + (-vx[k] * θ̇[k] - g * cos(θ[k]) + (u1[k] + u2[k]) / m) * dt
        )
        @constraint(MPC, θ̇[k+1] == θ̇[k] + (u1[k] - u2[k]) / I * dt)
        #@constraint(MPC, px[k] >= β)
    end

    @objective(
        MPC,
        Min,
        0 + sum(pz[i]^2 + px[i]^2 + 1e-6 * u1[i]^2 + 1e-6 * u2[i]^2 for i = 0:N)
    )
    optimize!(MPC)
    return value.(u1), value.(u2), value.(px), value.(pz)
end

function runSimpleMPC()
    x = [2.0 2.0 0.0 0.0 0.0 0.0]
    dt = 0.1
    N = 30
    xv = zeros(6, N)
    tv = Array(0:dt:dt*(N-1))
    for t = 1:N
        u1, u2, px, pz = SimpleMPC(x, 0.0 * zeros(6), pi / 4, pi / 3, 2, 1, dt)
        u = [u1[0] u2[0]]
        x = dynamics(x, u, dt)
        xv[:, t] = x
        @show u, x
    end
    plot(xv[1, :], xv[2, :], label = "Position", lw = 3)
end

function upperTrack(x)
    return sin.(x ./ 2) .+ 2.0
end

function lowerTrack(x)
    return sin.(x ./ 2) .- 0.0
end

function plotTrack(xi = 0.0, xf = 10.0)
    xv = Array(xi:0.1:xf)
    plot(xv, upperTrack(xv), color = :black)
    plot!(xv, lowerTrack(xv), color = :black)
end

function infinteTrackMPC(
    x0,
    xref = 0.0 .* zeros(6),
    θlim = pi / 4,
    θ̇lim = pi / 3,
    vxlim = 2,
    vzlim = 1,
    dt = 0.1,
    N = 10,
)
    MPC = Model(optimizer_with_attributes(
        Ipopt.Optimizer,
        "max_iter" => 55,
        "print_level" => 0,
    ))
    #State variables
    @variable(MPC, px[i = 0:N])
    @variable(MPC, pz[i = 0:N])
    @variable(MPC, -θlim <= θ[i = 0:N] <= θlim)
    @variable(MPC, -vxlim <= vx[i = 0:N] <= vxlim)
    @variable(MPC, -vzlim <= vz[i = 0:N] <= vxlim)
    @variable(MPC, -θ̇lim <= θ̇[i = 0:N] <= θ̇lim)
    @variable(MPC, 0 <= u1[i = 0:N])
    @variable(MPC, 0 <= u2[i = 0:N])
    @variable(MPC, β >= 0) #optional slack variable

    #State initial constraints
    @constraint(MPC, px[0] == x0[1])
    @constraint(MPC, pz[0] == x0[2])
    @constraint(MPC, θ[0] == x0[3])
    @constraint(MPC, vx[0] == x0[4])
    @constraint(MPC, vz[0] == x0[5])
    @constraint(MPC, θ̇[0] == x0[6])

    #x = [px pz θ vx vz θ̇]
    #     1  2  3 4  5  6
    for k = 0:N-1
        #Dynamics constraints
        @NLconstraint(MPC, px[k+1] == px[k] + (vx[k] * cos(θ[k]) - vz[k] * sin(θ[k])) * dt)
        @NLconstraint(MPC, pz[k+1] == pz[k] + (vx[k] * sin(θ[k]) + vz[k] * cos(θ[k])) * dt)
        @constraint(MPC, θ[k+1] == θ[k] + θ̇[k] * dt)
        @NLconstraint(MPC, vx[k+1] == vx[k] + (vz[k] * θ̇[k] - g * sin(θ[k])) * dt)
        @NLconstraint(
            MPC,
            vz[k+1] == vz[k] + (-vx[k] * θ̇[k] - g * cos(θ[k]) + (u1[k] + u2[k]) / m) * dt
        )
        @constraint(MPC, θ̇[k+1] == θ̇[k] + (u1[k] - u2[k]) / I * dt)

        #Track Constraints
        @NLconstraint(MPC, pz[k] <= sin(px[k] / 2) + 2.0) #upper track
        @NLconstraint(MPC, pz[k] >= sin(px[k] / 2) - 0.0) #lower track
        @constraint(MPC, (u1[k+1] - u1[k])^2 <= 0.01)
        @constraint(MPC, (u2[k+1] - u2[k])^2 <= 0.01)
    end

    @NLobjective(MPC, Max, sum(vx[i] * cos(θ[i] - vz[i] * sin(θ[i])) for i = 0:N))
    optimize!(MPC)
    return value.(u1), value.(u2), value.(px), value.(pz), MPC
end

function runTrackMPC()
    default(dpi = 300)
    default(thickness_scaling = 2)
    default(size = [1200, 800])
    x = [0.0 0.5 0.0 0.0 0.0 0.0]
    dt = 0.1
    N = 200
    xv = zeros(6, N)
    tv = Array(0:dt:dt*(N-1))
    for t = 1:N
        u1, u2, px, pz = infinteTrackMPC(x, 0.0 * zeros(6), pi / 4, pi / 3, 2, 1, dt)
        u = [u1[0] u2[0]]
        x = dynamics(x, u, dt)
        xv[:, t] = x
        @show u, x, t
    end
    plotTrack(min(xv[1, :]...), max(xv[1, :]...))
    plot!(xv[1, :], xv[2, :], label = "Position", lw = 3)
end

function ObstacleMPC(
    x0,
    xref = 0.0 .* zeros(6),
    Obstacles = nothing,
    θlim = pi / 4,
    θ̇lim = pi / 3,
    vxlim = 2,
    vzlim = 1,
    dt = 0.1,
    N = 30,
)
    MPC = Model(optimizer_with_attributes(
        Ipopt.Optimizer,
        "max_iter" => 600,
        "print_level" => 0,
    ))
    #State variables
    @variable(MPC, px[i = 0:N])
    @variable(MPC, pz[i = 0:N])
    @variable(MPC, -θlim <= θ[i = 0:N] <= θlim)
    @variable(MPC, -vxlim <= vx[i = 0:N] <= vxlim)
    @variable(MPC, -vzlim <= vz[i = 0:N] <= vxlim)
    @variable(MPC, -θ̇lim <= θ̇[i = 0:N] <= θ̇lim)
    @variable(MPC, 0 <= u1[i = 0:N])
    @variable(MPC, 0 <= u2[i = 0:N])
    @variable(MPC, β >= 0) #optional slack variable

    #State initial constraints
    @constraint(MPC, px[0] == x0[1])
    @constraint(MPC, pz[0] == x0[2])
    @constraint(MPC, θ[0] == x0[3])
    @constraint(MPC, vx[0] == x0[4])
    @constraint(MPC, vz[0] == x0[5])
    @constraint(MPC, θ̇[0] == x0[6])

    #x = [px pz θ vx vz θ̇]
    #     1  2  3 4  5  6
    for k = 0:N-1
        #Dynamics constraints
        @NLconstraint(MPC, px[k+1] == px[k] + (vx[k] * cos(θ[k]) - vz[k] * sin(θ[k])) * dt)
        @NLconstraint(MPC, pz[k+1] == pz[k] + (vx[k] * sin(θ[k]) + vz[k] * cos(θ[k])) * dt)
        @constraint(MPC, θ[k+1] == θ[k] + θ̇[k] * dt)
        @NLconstraint(MPC, vx[k+1] == vx[k] + (vz[k] * θ̇[k] - g * sin(θ[k])) * dt)
        @NLconstraint(
            MPC,
            vz[k+1] == vz[k] + (-vx[k] * θ̇[k] - g * cos(θ[k]) + (u1[k] + u2[k]) / m) * dt
        )
        @constraint(MPC, θ̇[k+1] == θ̇[k] + (u1[k] - u2[k]) / I * dt)
        #@constraint(MPC, (u1[k+1] - u1[k])^2 <= 0.1)
        #@constraint(MPC, (u2[k+1] - u2[k])^2 <= 0.1)
        for i = 1:size(Obstacles, 1)
            @constraint(
                MPC,
                (px[k] - Obstacles[i, 1])^2 + (pz[k] - Obstacles[i, 2])^2 >=
                (Obstacles[i, 3])^2
            )
        end
    end

    @objective(
        MPC,
        Min,
        0 + sum(
            (pz[i] - xref[2])^2 + (px[i] - xref[1])^2 + 1e-12 * u1[i]^2 + 1e-12 * u2[i]^2
            for i = 0:N
        )
    )
    optimize!(MPC)
    return value.(u1), value.(u2), value.(px), value.(pz)
end

function runObstacleMPC()
    x = [0.0 0.0 0.0 0.0 0.0 0.0]
    xref = [3.0 5.0 0.0 0.0 0.0 0.0]
    dt = 0.1
    N = 40
    xv = zeros(6, N)
    tv = Array(0:dt:dt*(N-1))
    xCenters = [4.0; 1.0]
    yCenters = [3.0; 2.5]
    Diameters = [1.0; 1.0]
    Obstacles = [xCenters yCenters Diameters]
    for t = 1:N
        u1, u2, px, pz = ObstacleMPC(x, xref, Obstacles, pi / 4, pi / 3, 2, 1, dt)
        u = [u1[0] u2[0]]
        x = dynamics(x, u, dt)
        xv[:, t] = x
        @show u, x
    end
    default(dpi = 300)
    default(thickness_scaling = 2)
    default(size = [1200, 800])
    h = plot(xv[1, :], xv[2, :], label = "Position", lw = 3)
    h = plotObstacles(Obstacles)
    display(h)
end

function plotObstacles(Obstacles)
    nObstacles = size(Obstacles, 1)
    h = nothing
    for i = 1:nObstacles
        h = plotCircle(Obstacles[i, 1], Obstacles[i, 2], Obstacles[i, 3])
    end
    return h
end

function circleShape(h, k, r)
    θ = LinRange(0, 2 * pi, 500)
    h .+ r * sin.(θ), k .+ r * cos.(θ)
end

function plotCircle(x, y, r)
    h = plot!(
        circleShape(x, y, r),
        seriesType = [:shape],
        lw = 0.5,
        c = :blue,
        linecolor = :black,
        legend = false,
        fillalpha = 0.2,
        aspect_ratio = 1,
    )
    return h
end

max_stars_repo_name: gbarsih/planar-quad-mpc

max_issues_repo_name: gbarsih/planar-quad-mpc

max_forks_repo_name: gbarsih/planar-quad-mpc


text:
include("abstract_trajectory.jl")
include("trajectory.jl")
include("trajectory_extension.jl")
include("reservoir_trajectory.jl")


text:
mutable struct Cell{Data, N, T, L}
    boundary::HyperRectangle{N, T}
    data::Data
    divisions::SVector{N, T}
    children::Union{TwosArray{N, Cell{Data, N, T, L}, L}, Nothing}
    parent::Union{Cell{Data, N, T, L}, Nothing}
end

function Cell(origin::SVector{N, T}, widths::SVector{N, T}, data::Data=nothing) where {Data, N, T}
    Cell(HyperRectangle(origin, widths), data)
end

@generated function Cell(boundary::HyperRectangle{N, T}, data::Data=nothing) where {Data, N, T}
    L = 2^N
    return quote
        T2 = Base.promote_op(/, T, Int)
        Cell{Data, N, T2, $L}(boundary,
             data,
             boundary.origin + boundary.widths / 2,
             nothing,
             nothing)
    end
end

@inline isleaf(cell::Cell) = cell.children === nothing
@inline children(cell::Cell) = cell.children
@inline parent(cell::Cell) = cell.parent
@inline center(cell::Cell) = center(cell.boundary)
@inline vertices(cell::Cell) = vertices(cell.boundary)

@inline size(::Type{C}) where {C <: Cell} = ()
@inline size(cell::Cell) = size(typeof(cell))

show(io::IO, cell::Cell) = print(io, "Cell: $(cell.boundary)")

@inline getindex(cell::Cell) = cell
@inline getindex(cell::Cell, ::CartesianIndex{0}) = cell
@inline getindex(cell::Cell, I) = getindex(cell.children, I)
@inline getindex(cell::Cell, I...) = getindex(cell.children, I...)

function child_boundary(cell::Cell, indices)
    half_widths = cell.boundary.widths ./ 2
    HyperRectangle(
        cell.boundary.origin + (SVector(indices) - 1) .* half_widths,
        half_widths)
end

@generated function map_children(f::Function, cell::Cell{Data, N, T, L}) where {Data, N, T, L}
    Expr(:call, :TwosArray, Expr(:tuple,
        [:(f(cell, $(I.I))) for I in CartesianIndices(ntuple(_ -> 2, Val(N)))]...))
end

child_indices(cell::Cell{Data, N, T, L}) where {Data, N, T, L} = child_indices(Val(N))

@generated function child_indices(::Val{N}) where N
    Expr(:call, :TwosArray, Expr(:tuple,
        [I.I for I in CartesianIndices(ntuple(_ -> 2, Val(N)))]...))
end

function split!(cell::Cell{Data, N}) where {Data, N}
    split!(cell, (c, I) -> cell.data)
end

@generated function split!(cell::Cell{Data, N}, child_data::AbstractArray) where {Data, N}
    split!_impl(cell, child_data, Val(N))
end

split!(cell::Cell, child_data_function::Function) =
    split!(cell, map_children(child_data_function, cell))

function split!_impl(::Type{C}, child_data, ::Val{N}) where {C <: Cell, N}
    child_exprs = [:(Cell(child_boundary(cell, $(I.I)),
                          child_data[$i])) for (i, I) in enumerate(CartesianIndices(ntuple(_ -> 2, Val(N))))]
    quote
        @assert isleaf(cell)
        cell.children = $(Expr(:call, :TwosArray, Expr(:tuple, child_exprs...)))
        for child in cell.children
            child.parent = cell
        end
        cell
    end
end

@generated function findleaf(cell::Cell{Data, N}, point::AbstractVector) where {Data, N}
    quote
        while true
            if isleaf(cell)
                return cell
            end
            length(point) == $N || throw(DimensionMismatch("expected a point of length $N"))
            @inbounds cell = $(Expr(:ref, :cell, [:(ifelse(point[$i] >= cell.divisions[$i], 2, 1)) for i in 1:N]...))
        end
    end
end

function allcells(cell::Cell)
    Channel() do c
        queue = [cell]
        while !isempty(queue)
            current = pop!(queue)
            put!(c, current)
            if !isleaf(current)
                append!(queue, children(current))
            end
        end
    end
end

function allleaves(cell::Cell)
    Channel() do c
        for child in allcells(cell)
            if isleaf(child)
                put!(c, child)
            end
        end
    end
end

function allparents(cell::Cell)
    Channel() do c
        queue = [cell]
        while !isempty(queue)
            current = pop!(queue)
            p = parent(current)
            if ! (p === nothing)
                put!(c, p)
                push!(queue, p)
            end
        end
    end
end


text:
using Stheno: MeanFunction, Kernel, AV, pairwise, ew, pw, BlockData, blocks
using Stheno: block_diagonal, AbstractGP
using FiniteDifferences: j′vp
import FiniteDifferences: to_vec

const _rtol = 1e-10
const _atol = 1e-10

_to_psd(A::Matrix{<:Real}) = A * A' + I
_to_psd(a::Vector{<:Real}) = exp.(a) .+ 1
_to_psd(σ::Real) = exp(σ) + 1
_to_psd(As::Vector{<:Matrix{<:Real}}) = block_diagonal(_to_psd.(As))

Base.length(::Nothing) = 0

function print_adjoints(adjoint_ad, adjoint_fd, rtol, atol)
    @show typeof(adjoint_ad), typeof(adjoint_fd)
    adjoint_ad, adjoint_fd = to_vec(adjoint_ad)[1], to_vec(adjoint_fd)[1]
    println("atol is $atol, rtol is $rtol")
    println("ad, fd, abs, rel")
    abs_err = abs.(adjoint_ad .- adjoint_fd)
    rel_err = abs_err ./ adjoint_ad
    display([adjoint_ad adjoint_fd abs_err rel_err])
    println()
end

# AbstractArrays.
function to_vec(x::ColVecs{<:Real})
    x_vec, back = to_vec(x.X)
    return x_vec, x_vec -> ColVecs(back(x_vec))
end

function to_vec(X::BlockArray)
    X_Array = Array(X)
    x, X_Array_from_vec = to_vec(X_Array)
    function BlockArray_from_vec(x::Vector)
        X_Array = X_Array_from_vec(x)
        return BlockArray(X_Array, axes(X))
    end
    return x, BlockArray_from_vec
end

function to_vec(x::BlockData)
    x_vecs, x_backs = zip(map(to_vec, blocks(x))...)
    sz = cumsum([map(length, x_vecs)...])
    return vcat(x_vecs...), function(v)
        return BlockData([x_backs[n](v[sz[n]-length(blocks(x)[n])+1:sz[n]])
            for n in 1:length(blocks(x))])
    end
end

function to_vec(X::T) where T<:Union{Adjoint,Transpose}
    U = T.name.wrapper
    return vec(Matrix(X)), x_vec->U(permutedims(reshape(x_vec, size(X))))
end

Base.zero(d::Dict) = Dict([(key, zero(val)) for (key, val) in d])
Base.zero(x::Array) = zero.(x)

# My version of isapprox
function fd_isapprox(x_ad::Nothing, x_fd, rtol, atol)
    return fd_isapprox(x_fd, zero(x_fd), rtol, atol)
end
function fd_isapprox(x_ad::AbstractArray, x_fd::AbstractArray, rtol, atol)
    return all(fd_isapprox.(x_ad, x_fd, rtol, atol))
end
function fd_isapprox(x_ad::Real, x_fd::Real, rtol, atol)
    return isapprox(x_ad, x_fd; rtol=rtol, atol=atol)
end
function fd_isapprox(x_ad::NamedTuple, x_fd, rtol, atol)
    f = (x_ad, x_fd)->fd_isapprox(x_ad, x_fd, rtol, atol)
    return all([f(getfield(x_ad, key), getfield(x_fd, key)) for key in keys(x_ad)])
end
function fd_isapprox(x_ad::Tuple, x_fd::Tuple, rtol, atol)
    return all(map((x, x′)->fd_isapprox(x, x′, rtol, atol), x_ad, x_fd))
end
function fd_isapprox(x_ad::Dict, x_fd::Dict, rtol, atol)
    return all([fd_isapprox(get(()->nothing, x_ad, key), x_fd[key], rtol, atol) for
        key in keys(x_fd)])
end

function adjoint_test(
    f, ȳ, x...;
    rtol=_rtol,
    atol=_atol,
    fdm=FiniteDifferences.Central(5, 1),
    print_results=false,
)
    # Compute forwards-pass and j′vp.
    y, back = Zygote.pullback(f, x...)
    @timeit to "adj_ad" adj_ad = back(ȳ)
    @timeit to "adj_fd" adj_fd = j′vp(fdm, f, ȳ, x...)

    # Check that forwards-pass agrees with plain forwards-pass.
    @test y ≈ f(x...)

    # Check that ad and fd adjoints (approximately) agree.
    print_results && print_adjoints(adj_ad, adj_fd, rtol, atol)
    @test fd_isapprox(adj_ad, adj_fd, rtol, atol)
end

"""
    mean_function_tests(m::MeanFunction, X::AbstractVector)

Test _very_ basic consistency properties of the mean function `m`.
"""
function mean_function_tests(m::MeanFunction, x::AbstractVector)
    @test ew(m, x) isa AbstractVector
    @test length(ew(m, x)) == length(x)
end

"""
    cross_kernel_tests(k::Kernel, x0::AV, x1::AV, x2::AV)

Tests that any cross kernel `k` should be able to pass. Requires that
`length(x0) == length(x1)` and `length(x0) ≠ length(x2)`.
"""
function cross_kernel_tests(k::Kernel, x0::AV, x1::AV, x2::AV; atol=1e-9)
    @assert length(x0) == length(x1)
    @assert length(x0) ≠ length(x2)

    # Check that elementwise basically works.
    @test ew(k, x0, x1) isa AbstractVector
    @test length(ew(k, x0, x1)) == length(x0)

    # Check that pairwise basically works.
    @test pw(k, x0, x2) isa AbstractMatrix
    @test size(pw(k, x0, x2)) == (length(x0), length(x2))

    # Check that elementwise is consistent with pairwise.
    @test ew(k, x0, x1) ≈ diag(pw(k, x0, x1)) atol=atol
end

"""
    kernel_tests(k::Kernel, X0::AbstractVector, X1::AbstractVector, X2::AbstractVector)

Tests that any kernel `k` should be able to pass. Requires that `length(X0) == length(X1)`
and `length(X0) ≠ length(X2)`.
"""
function kernel_tests(k::Kernel, x0::AV, x1::AV, x2::AV; atol=1e-9)
    @assert length(x0) == length(x1)
    @assert length(x0) ≠ length(x2)

    # Check that all of the binary methods work as expected.
    cross_kernel_tests(k, x0, x1, x2)

    # Check additional binary elementwise properties for kernels.
    @test ew(k, x0, x1) ≈ ew(k, x1, x0)
    @test pw(k, x0, x2) ≈ pw(k, x2, x0)' atol=atol

    # Check that unary elementwise basically works.
    @test ew(k, x0) isa AbstractVector
    @test length(ew(k, x0)) == length(x0)

    # Check that unary pairwise basically works.
    @test pw(k, x0) isa AbstractMatrix
    @test size(pw(k, x0)) == (length(x0), length(x0))
    @test pw(k, x0) ≈ pw(k, x0)' atol=atol

    # Check that unary elementwise is consistent with unary pairwise.
    @test ew(k, x0) ≈ diag(pw(k, x0)) atol=atol

    # Check that unary pairwise produces a positive definite matrix (approximately).
    @test all(eigvals(Matrix(pw(k, x0))) .> -atol)

    # Check that unary elementwise / pairwise are consistent with the binary versions.
    @test ew(k, x0) ≈ ew(k, x0, x0) atol=atol
    @test pw(k, x0) ≈ pw(k, x0, x0) atol=atol
end

"""
    differentiable_mean_function_tests(m::MeanFunction, ȳ::AV, x::AV)

Ensure that the gradient w.r.t. the inputs of `MeanFunction` `m` are approximately correct.
"""
function differentiable_mean_function_tests(
    m::MeanFunction,
    ȳ::AbstractVector{<:Real},
    x::AbstractVector{<:Real};
    rtol=_rtol,
    atol=_atol,
)
    # Run forward tests.
    mean_function_tests(m, x)

    # Check adjoint.
    @assert length(ȳ) == length(x)
    adjoint_test(x->ew(m, x), ȳ, x; rtol=rtol, atol=atol)
end
function differentiable_mean_function_tests(
    m::MeanFunction,
    ȳ::AbstractVector{<:Real},
    x::ColVecs{<:Real};
    rtol=_rtol,
    atol=_atol,
)
    # Run forward tests.
    mean_function_tests(m, x)

    @assert length(ȳ) == length(x)
    adjoint_test(X->ew(m, ColVecs(X)), ȳ, x.X; rtol=rtol, atol=atol)
end
function differentiable_mean_function_tests(
    rng::AbstractRNG,
    m::MeanFunction,
    x::AV;
    rtol=_rtol,
    atol=_atol,
)
    return differentiable_mean_function_tests(
        m,
        randn(rng, length(x)),
        x;
        rtol=rtol,
        atol=atol,
    )
end

"""
    differentiable_cross_kernel_tests(
        k::Kernel,
        ȳ::AbstractVector{<:Real},
        Ȳ::AbstractMatrix{<:Real},
        x0::AbstractVector,
        x1::AbstractVector,
        x2::AbstractVector,
    )

Ensure that the adjoint w.r.t. the inputs of a `Kernel` which is supposed to be
differentiable are approximately correct.
"""
function differentiable_cross_kernel_tests(
    k::Kernel,
    ȳ::AbstractVector{<:Real},
    Ȳ::AbstractMatrix{<:Real},
    x0::AbstractVector,
    x1::AbstractVector,
    x2::AbstractVector;
    rtol=_rtol,
    atol=_atol,
)
    # Run forwards-pass cross kernel tests.
    cross_kernel_tests(k, x0, x1, x2)

    # Ensure that the inputs are as required.
    @assert length(ȳ) == length(x0)
    @assert length(ȳ) == length(x1)
    @assert size(Ȳ) == (length(x0), length(x2))

    # Binary elementwise.
    adjoint_test((x, x′)->ew(k, x, x′), ȳ, x0, x1; rtol=rtol, atol=atol)

    # Binary pairwise.
    adjoint_test((x, x′)->pw(k, x, x′), Ȳ, x0, x2; rtol=rtol, atol=atol)
end
function differentiable_cross_kernel_tests(
    rng::AbstractRNG,
    k::Kernel,
    x0::AV,
    x1::AV,
    x2::AV;
    rtol=_rtol,
    atol=_atol,
)
    ȳ, Ȳ = randn(rng, length(x0)), randn(rng, length(x0), length(x2))
    return differentiable_cross_kernel_tests(k, ȳ, Ȳ, x0, x1, x2; rtol=rtol, atol=atol)
end

"""
    differentiable_kernel_tests(
        k::Kernel,
        ȳ::AbstractVector{<:Real},
        Ȳ::AbstractMatrix{<:Real},
        Ȳ_sq::AbstractMatrix{<:Real},
        x0::AbstractVector,
        x1::AbstractVector,
        x2::AbstractVector,
    )

A superset of the tests provided by `differentiable_cross_kernel_tests` designed to test
kernels (which provide unary, in addition to binary, methods for `elementwise` and
`pairwise`.)
"""
function differentiable_kernel_tests(
    k::Kernel,
    ȳ::AbstractVector{<:Real},
    Ȳ::AbstractMatrix{<:Real},
    Ȳ_sq::AbstractMatrix{<:Real},
    x0::AbstractVector,
    x1::AbstractVector,
    x2::AbstractVector;
    rtol=_rtol,
    atol=_atol,
)
    # Run the forwards-pass kernel tests.
    kernel_tests(k, x0, x1, x2)

    # Ensure that the inputs are as required.
    @assert length(ȳ) == length(x0)
    @assert length(ȳ) == length(x1)
    @assert size(Ȳ) == (length(x0), length(x2))
    @assert size(Ȳ_sq, 1) == size(Ȳ_sq, 2)
    @assert size(Ȳ_sq, 1) == length(x0)

    # Run the Kernel tests.
    differentiable_cross_kernel_tests(k, ȳ, Ȳ, x0, x1, x2; rtol=rtol, atol=atol)

    # Unary elementwise tests.
    adjoint_test(x->ew(k, x), ȳ, x0; rtol=rtol, atol=atol)

    # Unary pairwise test.
    adjoint_test(x->pw(k, x), Ȳ_sq, x0; rtol=rtol, atol=atol)
end
function differentiable_kernel_tests(
    rng::AbstractRNG,
    k::Kernel,
    x0::AV,
    x1::AV,
    x2::AV;
    rtol=_rtol,
    atol=_atol,
)
    N, N′ = length(x0), length(x2)
    ȳ, Ȳ, Ȳ_sq = randn(rng, N), randn(rng, N, N′), randn(rng, N, N)
    return differentiable_kernel_tests(k, ȳ, Ȳ, Ȳ_sq, x0, x1, x2; rtol=rtol, atol=atol)
end

"""
    abstractgp_interface_tests(
        f::AbstractGP, f′::AbstractGP, x0::AV, x1::AV, x2::AV, x3::AV;
        atol=1e-9, rtol=1e-9,
    )

Check that the `AbstractGP` interface is at least implemented for `f` and is
self-consistent. `x0` and `x1` must be valid inputs for `f`. `x2` and `x3` must be a valid
input for `f′`.
"""
function abstractgp_interface_tests(
    f::AbstractGP, f′::AbstractGP, x0::AV, x1::AV, x2::AV, x3::AV;
    atol=1e-9, rtol=1e-9,
)
    m = mean_vector(f, x0)
    @test m isa AbstractVector{<:Real}
    @test length(m) == length(x0)

    @assert length(x0) ≠ length(x1)
    @assert length(x0) ≠ length(x2)
    @assert length(x0) == length(x3)

    # Check that binary cov conforms to the API
    K_ff′_x0_x2 = cov(f, f′, x0, x2)
    @test K_ff′_x0_x2 isa AbstractMatrix{<:Real}
    @test size(K_ff′_x0_x2) == (length(x0), length(x2))
    @test K_ff′_x0_x2 ≈ cov(f′, f, x2, x0)'

    # Check that unary cov is consistent with binary cov and conforms to the API
    K_x0 = cov(f, x0)
    @test K_x0 isa AbstractMatrix{<:Real}
    @test size(K_x0) == (length(x0), length(x0))
    @test K_x0 ≈ cov(f, f, x0, x0) atol=atol rtol=rtol
    @test minimum(eigvals(K_x0)) > -atol
    @test K_x0 ≈ K_x0' atol=atol rtol=rtol

    # Check that single-process binary cov is consistent with binary-process binary-cov
    K_x0_x1 = cov(f, x0, x1)
    @test K_x0_x1 isa AbstractMatrix{<:Real}
    @test size(K_x0_x1) == (length(x0), length(x1))
    @test K_x0_x1 ≈ cov(f, f, x0, x1)

    # Check that binary cov_diag conforms to the API and is consistent with binary cov
    K_x0_x3_diag = cov_diag(f, f′, x0, x3)
    @test K_x0_x3_diag isa AbstractVector{<:Real}
    @test length(K_x0_x3_diag) == length(x0)
    @test K_x0_x3_diag ≈ diag(cov(f, f′, x0, x3)) atol=atol rtol=rtol
    @test K_x0_x3_diag ≈ cov_diag(f′, f, x3, x0) atol=atol rtol=rtol

    # Check that unary-binary cov_diag is consistent.
    K_x0_x0_diag = cov_diag(f, x0, x0)
    @test K_x0_x0_diag isa AbstractVector{<:Real}
    @test length(K_x0_x0_diag) == length(x0)
    @test K_x0_x0_diag ≈ diag(cov(f, x0, x0)) atol=atol rtol=rtol

    # Check that unary cov_diag conforms to the API and is consistent with unary cov
    K_x0_diag = cov_diag(f, x0)
    @test K_x0_diag isa AbstractVector{<:Real}
    @test length(K_x0_diag) == length(x0)
    @test K_x0_diag ≈ diag(cov(f, x0)) atol=atol rtol=rtol
end


text:
using LinearAlgebra
using DifferentialEquations
using ForwardDiff
using Optim

export RaySolver

"""
$(TYPEDEF)
A pure Julia implementation of a ray/Gaussian beam propapagation model. The model supports complex
environments, but retains differentiability.
"""
Base.@kwdef struct RaySolver{T1,T2} <: PropagationModel{T1}
  env::T1
  nbeams::Int = 0
  minangle::Float64 = -80°
  maxangle::Float64 = +80°
  ds::Float64 = 1.0
  atol::Float64 = 1e-4
  rugosity::Float64 = 1.5
  athreshold::Float64 = 1e-5
  solver::T2 = missing
  solvertol::Float64 = 1e-4
  function RaySolver(env, nbeams, minangle, maxangle, ds, atol, rugosity, athreshold, solver, solvertol)
    nbeams < 0 && (nbeams = 0)
    -π/2 ≤ minangle ≤ π/2 || throw(ArgumentError("minangle should be between -π/2 and π/2"))
    -π/2 ≤ maxangle ≤ π/2 || throw(ArgumentError("maxangle should be between -π/2 and π/2"))
    minangle < maxangle || throw(ArgumentError("maxangle should be more than minangle"))
    if solver === missing
      if ssp(env) isa IsoSSP || (ssp(env) isa SampledSSP && length(ssp(env).z) ≤ 2)
        solver = Tsit5()
      else
        solver = Rodas5()
      end
    end
    new{typeof(env),typeof(solver)}(check(RaySolver, env), nbeams, minangle, maxangle, ds, atol, rugosity, athreshold, solver, solvertol)
  end
end

"""
    RaySolver(env; nbeams, minangle, maxangle, ds, atol, rugosity, athreshold, solver, solvertol)

Create a RaySolver propagation model.
"""
RaySolver(env; kwargs...) = RaySolver(; env=env, kwargs...)

### interface functions

function check(::Type{RaySolver}, env::Union{<:UnderwaterEnvironment,Missing})
  env
end

function arrivals(model::RaySolver, tx1::AcousticSource, rx1::AcousticReceiver)
  check2d([tx1], [rx1])
  a = [RayArrival(r.time, r.phasor, r.surface, r.bottom, r.launchangle, r.arrivalangle) for r ∈ eigenrays(model, tx1, rx1; ds=0.0)]
  sort(a; by = a1 -> a1.time)
end

function eigenrays(model::RaySolver, tx1::AcousticSource, rx1::AcousticReceiver; ds=model.ds)
  check2d([tx1], [rx1])
  p2 = location(rx1)
  nbeams = model.nbeams
  if nbeams == 0
    p1 = location(tx1)
    R = abs(p2[1] - p1[1])
    h = maxdepth(bathymetry(model.env))
    nbeams = ceil(Int, 16 * (model.maxangle - model.minangle) / atan(h, R))
  end
  θ = range(model.minangle, model.maxangle; length=nbeams)
  n = length(θ)
  rmax = ForwardDiff.value(p2[1])
  err = fill(NaN64, n)
  Threads.@threads for i ∈ 1:n
    p3 = traceray(model, tx1, θ[i], rmax).raypath[end]
    if isapprox(p3[1], p2[1]; atol=model.atol) && isapprox(p3[2], p2[2]; atol=model.atol)
      err[i] = ForwardDiff.value(p3[3] - p2[3])
    end
  end
  T = promote_type(envrealtype(model.env), eltype(location(tx1)), typeof(nominalfrequency(tx1)), eltype(location(rx1)))
  erays = Array{RayArrival{T,T},1}(undef, 0)
  for i ∈ 1:n
    if isapprox(err[i], 0.0; atol=model.atol)
      eray = traceray(model, tx1, convert(T, θ[i]), rmax, ds)
      push!(erays, eray)
    elseif i > 1 && !isnan(err[i-1]) && !isnan(err[i]) && sign(err[i-1]) * sign(err[i]) == -1
      a, b = ordered(θ[i-1], θ[i])
      soln = optimize(ϕ -> abs2(traceray(model, tx1, ϕ, rmax).raypath[end][3] - p2[3]), a, b; abs_tol=1e-6)
      if Optim.converged(soln) && soln.minimum ≤ 1e-2
        eray = traceray(model, tx1, convert(T, soln.minimizer), rmax, ds)
        push!(erays, eray)
      end
    elseif i > 2 && isnearzero(err[i-2], err[i-1], err[i])
      a, b = ordered(θ[i-2], θ[i])
      soln = optimize(ϕ -> abs2(traceray(model, tx1, ϕ, rmax).raypath[end][3] - p2[3]), a, b; abs_tol=1e-6)
      if Optim.converged(soln) && soln.minimum ≤ 1e-2
        eray = traceray(model, tx1, convert(T, soln.minimizer), rmax, ds)
        push!(erays, eray)
      end
    end
  end
  erays
end

function rays(model::RaySolver, tx1::AcousticSource, θ::Real, rmax)
  check2d([tx1], [])
  traceray(model, tx1, θ, rmax, model.ds)
end

function transfercoef(model::RaySolver, tx1::AcousticSource, rx::AcousticReceiverGrid2D{T1}; mode=:coherent) where {T1}
  check2d([tx1], rx)
  mode === :coherent || mode === :incoherent || throw(ArgumentError("Unknown mode :" * string(mode)))
  # implementation primarily based on ideas from COA (Computational Ocean Acoustics, 2nd ed., ch. 3)
  f = nominalfrequency(tx1)
  T = promote_type(envrealtype(model.env), eltype(location(tx1)), typeof(f), T1, ComplexF64)
  tc = zeros(T, size(rx,1), size(rx,2), Threads.nthreads())
  rmax = maximum(rx.xrange) + 0.1
  nbeams = model.nbeams
  if nbeams == 0
    p1 = location(tx1)
    R = abs(rmax - p1[1])
    h = maxdepth(bathymetry(model.env))
    nbeams = clamp(ceil(Int, 20 * (model.maxangle - model.minangle) / atan(h, R)), 100, 1000)
  end
  θ = range(model.minangle, model.maxangle; length=nbeams)
  δθ = Float64(θ.step)
  c₀ = soundspeed(ssp(model.env), location(tx1)...)
  ω = 2π * f
  G = 1 / (2π)^(1/4)
  Threads.@threads for θ1 ∈ θ
    β = (mode === :incoherent ? 2 : 1) * cos(θ1)/c₀
    traceray(model, tx1, θ1, rmax, 1.0; cb = (s1, u1, s2, u2, A₀, D₀, t₀, cₛ, kmah1, kmah2) -> begin
      r1, z1, ξ1, ζ1, t1, _, q1 = u1
      r2, z2, ξ2, ζ2, t2, _, q2 = u2
      ndx = findall(r -> r1 ≤ r < r2, rx.xrange)
      rz1 = (r1, z1)
      vlen = norm((r2-r1, z2-z1))
      rvlen = norm((ξ1, ζ1))
      tᵥ = (ξ1, ζ1) ./ rvlen
      nᵥ = (ζ1, -ξ1) ./ rvlen
      D = D₀ + (s1 + s2) / 2
      γ = fast_absorption(f, D, salinity(model.env))
      Wmax = max(q1, q2) * δθ
      ndx2 = findall(z -> min(z1, z2) - 4*Wmax ≤ z ≤ max(z1, z2) + 4*Wmax, rx.zrange)
      for j ∈ ndx
        for i ∈ ndx2
          rz = (rx.xrange[j], rx.zrange[i])
          v = rz .- rz1
          s = dot(v, tᵥ)
          n = dot(v, nᵥ)
          α = s / vlen
          t = t₀ + t1 + (t2 - t1) * α
          q = q1 + (q2 - q1) * α
          A = A₀ * γ * G * √abs(β * cₛ / (rz[1] * q)) # COA (3.76)
          kmah = round(Int, kmah1 + (kmah2 - kmah1) * α)
          A *= cis(-π/2 * kmah)                       # COA section 3.4.1 (KMAH correction)
          W = abs(q * δθ)                             # COA (3.74)
          A *= exp(-(n / W)^2)
          tc[j, i, Threads.threadid()] += mode === :coherent ? conj(A) * cis(-ω * t) : Complex(abs2(A), 0.0)
        end
      end
    end)
  end
  rv = dropdims(sum(tc; dims=3); dims=3)
  mode === :incoherent && (rv = sqrt.(rv))
  rv
end

### helper functions

ordered(a, b) = a < b ? (a, b) : (b, a)

function isnearzero(a, b, c)
  (isnan(a) || isnan(b) || isnan(c)) && return false
  sign(a) != sign(b) && return false
  sign(a) != sign(c) && return false
  abs(a) < abs(b) && return false
  abs(c) < abs(b) && return false
  return true
end

function check2d(tx, rx)
  # TODO: remove restrictions on coordinates
  all(location(tx1)[1] == 0.0 for tx1 ∈ tx) || throw(ArgumentError("RaySolver requires transmitters at (0, 0, z)"))
  all(location(tx1)[2] == 0.0 for tx1 ∈ tx) || throw(ArgumentError("RaySolver requires transmitters in the x-z plane"))
  all(location(rx1)[1] >= 0.0 for rx1 ∈ rx) || throw(ArgumentError("RaySolver requires receivers to be in the +x halfspace"))
  all(location(rx1)[2] == 0.0 for rx1 ∈ rx) || throw(ArgumentError("RaySolver requires receivers in the x-z plane"))
end

function rayeqns!(du, u, params, s)
  # implementation based on COA (3.161-164, 3.58-63)
  # assumes range-independent soundspeed
  r, z, ξ, ζ, t, p, q = u
  c, ∂c, ∂²c = params
  cᵥ = c(z)
  cᵥ² = cᵥ * cᵥ
  c̄ = ∂²c(z) * ξ * ξ
  du[1] = cᵥ * ξ
  du[2] = cᵥ * ζ
  du[3] = 0
  du[4] = -∂c(z) / cᵥ²
  du[5] = 1 / cᵥ
  du[6] = -c̄ * q
  du[7] = cᵥ * p
end

function checkray!(out, u, s, integrator, a::Altimetry, b::Bathymetry, rmax)
  out[1] = altitude(a, u[1], 0.0) - u[2]   # surface reflection
  out[2] = u[2] + depth(b, u[1], 0.0)      # bottom reflection
  out[3] = rmax - u[1]                     # maximum range
  out[4] = u[3]                            # ray turned back
end

function traceray1(T, model, r0, z0, θ, rmax, ds, p0, q0)
  a = altimetry(model.env)
  b = bathymetry(model.env)
  c = z -> soundspeed(ssp(model.env), 0.0, 0.0, z)
  ∂c = z -> ForwardDiff.derivative(c, z)
  ∂²c = z -> ForwardDiff.derivative(∂c, z)
  cᵥ = c(z0)
  u0 = [convert(T, r0), z0, cos(θ)/cᵥ, sin(θ)/cᵥ, zero(T), p0, q0]
  prob = ODEProblem{true}(rayeqns!, u0, (zero(T), one(T) * model.rugosity * (rmax-r0)/cos(θ)), (c, ∂c, ∂²c))
  cb = VectorContinuousCallback(
    (out, u, s, i) -> checkray!(out, u, s, i, a, b, rmax),
    (i, ndx) -> terminate!(i), 4; rootfind=true)
  if ds ≤ 0
    soln = solve(prob, model.solver; abstol=model.solvertol, save_everystep=false, callback=cb)
  else
    soln = solve(prob, model.solver; abstol=model.solvertol, saveat=ds, callback=cb)
  end
  s2 = soln[end]
  soln.t[end], s2[1], s2[2], atan(s2[4], s2[3]), s2[5], s2[6], s2[7], soln.u, soln.t
end

function traceray(model::RaySolver, tx1::AcousticSource, θ::Real, rmax, ds=0.0; cb=nothing)
  θ₀ = θ
  f = nominalfrequency(tx1)
  zmin = -maxdepth(bathymetry(model.env))
  ϵ = one(zmin) * model.solvertol
  p = location(tx1)
  c₀ = soundspeed(ssp(model.env), p...)
  T = promote_type(envrealtype(model.env), eltype(p), typeof(f), typeof(θ), typeof(rmax))
  raypath = Array{NTuple{3,T}}(undef, 0)
  A = one(Complex{T})   # phasor
  s = 0                 # surface bounces
  b = 0                 # bottom bounces
  t = zero(T)           # time along ray
  D = zero(T)           # distance along ray
  q = zero(T)           # spreading factor
  qp = one(T) / c₀      # spreading rate
  while true
    dD, r, z, θ, dt, qp, q, u, svec = traceray1(T, model, p[1], p[3], θ, rmax, ds, qp, q)
    oq = u[1][7]
    kmah = 0
    kmahhist = zeros(length(u))
    for i ∈ 1:length(u)
      sign(oq) * sign(u[i][7]) == -1 && (kmah += 1)
      kmahhist[i] = kmah
      u[i][7] != 0.0 && (oq = u[i][7])
      push!(raypath, (u[i][1], 0.0, u[i][2]))
    end
    if cb !== nothing
      for i ∈ 2:length(u)
        cₛ = soundspeed(ssp(model.env), (u[i-1][1]+u[i][1])/2, 0.0, (u[i-1][2]+u[i][2])/2)
        cb(svec[i-1], u[i-1], svec[i], u[i], A, D, t, cₛ, kmahhist[i-1], kmahhist[i])
      end
    end
    t += dt
    D += dD
    A *= cis(-π/2 * kmah)                 # COA section 3.4.1 (KMAH correction)
    zmin = -depth(bathymetry(model.env), r, 0.0)
    zmax = altitude(altimetry(model.env), r, 0.0)
    p = (r, 0.0, clamp(z, zmin+ϵ, zmax-ϵ))
    if r ≥ rmax - 1e-3
      break
    elseif isapprox(z, zmax; atol=1e-3)   # hit the surface
      s += 1
      A *= reflectioncoef(seasurface(model.env), f, π/2 - θ)
      α = atan(ForwardDiff.derivative(x -> altitude(altimetry(model.env), x, 0.0), r))
      θ = -θ + 2α
    elseif isapprox(z, zmin; atol=1e-3)   # hit the bottom
      b += 1
      A *= reflectioncoef(seabed(model.env), f, π/2 + θ)
      α = atan(ForwardDiff.derivative(x -> depth(bathymetry(model.env), x, 0.0), r))
      θ = -θ - 2α
    else
      break
    end
    if abs(θ) ≥ π/2
      break
    end
    if abs(A)/q < model.athreshold
      break
    end
  end
  cₛ = soundspeed(ssp(model.env), p...)
  A *= √abs(cₛ * cos(θ₀) / (p[1] * c₀ * q))        # COA (3.65)
  A *= fast_absorption(f, D, salinity(model.env))
  RayArrival{T,T}(t, conj(A), s, b, θ₀, -θ, raypath)    # conj(A) needed to match with Bellhop
end


text:
const naivebayesstanmodel = "
// supervised naive Bayes

data {
  // training data
  int<lower=1> K;               // num topics
  int<lower=1> V;               // num words
  int<lower=0> M;               // num docs
  int<lower=0> N;               // total word instances
  int<lower=1,upper=K> z[M];    // topic for doc m
  int<lower=1,upper=V> w[N];    // word n
  int<lower=1,upper=M> doc[N];  // doc ID for word n
  // hyperparameters
  vector<lower=0>[K] alpha;     // topic prior
  vector<lower=0>[V] beta;      // word prior
}
parameters {
  simplex[K] theta;   // topic prevalence
  simplex[V] phi[K];  // word dist for topic k
}
model {
  // priors
  theta ~ dirichlet(alpha);
  for (k in 1:K)
    phi[k] ~ dirichlet(beta);
  // likelihood, including latent category
  for (m in 1:M)
    z[m] ~ categorical(theta);
  for (n in 1:N)
    w[n] ~ categorical(phi[z[doc[n]]]);
}
"


text:
using Printf: @printf
@time let sm = 0
          for p in PrimesPaged()
              p > 2000000 && break
              sm += p
          end; @printf("%d\n", sm) end

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode


text:
using DelayEmbeddings, DynamicalSystemsBase, PyPlot, Test
pygui(true)

function ar_process(u0::T, α::T, p::T, N::Int) where {T<:Real}
    x = zeros(T, N+10)
    x[1] = u0
    for i = 2:N+10
        x[i] = α*x[i-1] + p*randn()
    end
    return x[11:end]
end

trajectories = Dict(
    "White noise" => Dataset(randn(5000)),
    "AR(1)" => Dataset(ar_process(.2,.9,.1,5000)),
    "Hénon (chaotic)" => trajectory(Systems.henon(a=1.4, b=0.3), 199, Ttr=1000),
    "Hénon (periodic)" => trajectory(Systems.henon(a=1.054, b=0.3), 199, Ttr=1000),
    "Lorenz" => trajectory(Systems.lorenz(), 1000.0; Ttr=1000, Δt = 0.05),
    "Lorenz96" => trajectory(Systems.lorenz96(5; F = 8.0), 1000.0; Ttr=1000, Δt = 0.05),
    "Roessler" => trajectory(Systems.roessler(), 1000.0; Ttr=1000, Δt = 0.05),
    "Sine wave" => Dataset(map(x->[sin.(x) cos.(x)], StepRangeLen(0.0,0.2,200)))
)
dict_keys = ["White noise","AR(1)","Hénon (chaotic)","Hénon (periodic)","Lorenz","Lorenz96","Roessler","Sine wave"]

# set noise level
σ = 0.01
# set maximum encountered dimension
dmax = 15
γs = 0:dmax-1
dims = γs .+ 1

slope_thres = .1

for names in dict_keys
    tr = trajectories[names]
    tr = standardize(tr)
    s = tr[:, 1].+ σ.*randn(length(tr))
    τ = estimate_delay(s, "mi_min")

    # do a plot of the differently found methods

    fs = (afnn, fnn, f1nn, ifnn)
    methods = ["afnn", "fnn", "f1nn", "ifnn"]
    fig, axs  = subplots(2, 1; sharex = true)
    for (i, f) in enumerate(fs)
        x = f(s, τ, γs)

        if i == 1
            x .= 1 .- x
            E2 = stochastic_indicator(s, τ, γs)
        end

        axs[1].plot(dims, x, label = string(f); color = "C$(i)")
        axs[2].plot(dims[1:end-1], diff(x); color = "C$(i)")
        # lrs, tans = linear_regions(dims, x, tol = 0.5)
        # axs[1].scatter(dims[lrs], x[lrs], color = "C$(i)")
        z = slope_thres # threshold of slopes
        axs[2].axhspan(-z, z; color = "0.5", alpha = 0.2, label = "±$(z)")

        # Estimations of condition for convergence:
        # earliest point where x is both less than a thresohold, 0.05, and local slope is
        # between ± 0.05
        r = 0.05 # threshold of convergence near 0
        d_max = dmax
        Y, _, rat =  optimal_traditional_de(s, "mi_min", methods[i]; dmax = d_max, slope_thres = slope_thres)
        println("$(methods[i]) embedding dimension for $names: $(size(Y,2)) ")
        axs[1].scatter(dims[size(Y,2)], x[size(Y,2)]; color = "C$(i)", edgecolors  = "k")
    end
    axs[1].grid(); axs[2].grid()
    axs[1].legend(loc = "upper right", ncol = 2)
    axs[1].set_ylabel("estimator")
    axs[2].set_xlabel("dimension")
    axs[2].set_ylabel("local slope")
    axs[1].set_title(names)
end