text:
module Atmos

export AtmosModel,
    AtmosAcousticLinearModel,
    AtmosAcousticGravityLinearModel,
    HLLCNumericalFlux,
    RoeNumericalFlux,
    RoeNumericalFluxMoist

using UnPack
using CLIMAParameters
using CLIMAParameters.Planet: grav, cp_d, R_v, LH_v0, e_int_v0
using CLIMAParameters.Atmos.SubgridScale: C_smag
using DocStringExtensions
using LinearAlgebra, StaticArrays
using ..ConfigTypes
using ..Orientations
import ..Orientations:
    vertical_unit_vector,
    altitude,
    latitude,
    longitude,
    projection_normal,
    gravitational_potential,
    ∇gravitational_potential,
    projection_tangential

using ..VariableTemplates
using ..Thermodynamics
using ..TemperatureProfiles

using ..TurbulenceClosures
import ..TurbulenceClosures: turbulence_tensors
using ..TurbulenceConvection

import ..Thermodynamics: internal_energy
using ..MPIStateArrays: MPIStateArray
using ..Mesh.Grids:
    VerticalDirection,
    HorizontalDirection,
    min_node_distance,
    EveryDirection,
    Direction

using ..BalanceLaws
using ClimateMachine.Problems

import ..BalanceLaws:
    vars_state,
    flux_first_order!,
    flux_second_order!,
    source!,
    eq_tends,
    flux,
    source,
    wavespeed,
    boundary_conditions,
    boundary_state!,
    compute_gradient_argument!,
    compute_gradient_flux!,
    transform_post_gradient_laplacian!,
    init_state_auxiliary!,
    init_state_prognostic!,
    update_auxiliary_state!,
    indefinite_stack_integral!,
    reverse_indefinite_stack_integral!,
    integral_load_auxiliary_state!,
    integral_set_auxiliary_state!,
    reverse_integral_load_auxiliary_state!,
    reverse_integral_set_auxiliary_state!

import ClimateMachine.DGMethods:
    LocalGeometry, lengthscale, resolutionmetric, DGModel

import ..DGMethods.NumericalFluxes:
    boundary_state!,
    boundary_flux_second_order!,
    normal_boundary_flux_second_order!,
    NumericalFluxFirstOrder,
    NumericalFluxGradient,
    NumericalFluxSecondOrder,
    CentralNumericalFluxHigherOrder,
    CentralNumericalFluxDivergence,
    CentralNumericalFluxFirstOrder,
    numerical_flux_first_order!,
    NumericalFluxFirstOrder
using ..DGMethods.NumericalFluxes:
    RoeNumericalFlux,
    HLLCNumericalFlux,
    RusanovNumericalFlux,
    RoeNumericalFluxMoist

import ..Courant: advective_courant, nondiffusive_courant, diffusive_courant

"""
    AtmosModel <: BalanceLaw

A `BalanceLaw` for atmosphere modeling. Users may over-ride prescribed
default values for each field.

# Usage

    AtmosModel(
        param_set,
        problem,
        orientation,
        ref_state,
        turbulence,
        hyperdiffusion,
        spongelayer,
        moisture,
        precipitation,
        radiation,
        source,
        tracers,
        data_config,
    )

# Fields
$(DocStringExtensions.FIELDS)
"""
struct AtmosModel{FT, PS, PR, O, RS, T, TC, HD, VS, M, P, R, S, TR, LF, DC} <:
       BalanceLaw
    "Parameter Set (type to dispatch on, e.g., planet parameters. See CLIMAParameters.jl package)"
    param_set::PS
    "Problem (initial and boundary conditions)"
    problem::PR
    "An orientation model"
    orientation::O
    "Reference State (For initial conditions, or for linearisation when using implicit solvers)"
    ref_state::RS
    "Turbulence Closure (Equations for dynamics of under-resolved turbulent flows)"
    turbulence::T
    "Turbulence Convection Closure (e.g., EDMF)"
    turbconv::TC
    "Hyperdiffusion Model (Equations for dynamics of high-order spatial wave attenuation)"
    hyperdiffusion::HD
    "Viscous sponge layers"
    viscoussponge::VS
    "Moisture Model (Equations for dynamics of moist variables)"
    moisture::M
    "Precipitation Model (Equations for dynamics of precipitating species)"
    precipitation::P
    "Radiation Model (Equations for radiative fluxes)"
    radiation::R
    "Source Terms (Problem specific source terms)"
    source::S
    "Tracer Terms (Equations for dynamics of active and passive tracers)"
    tracers::TR
    "Large-scale forcing (Forcing information from GCMs, reanalyses, or observations)"
    lsforcing::LF
    "Data Configuration (Helper field for experiment configuration)"
    data_config::DC
end

"""
    AtmosModel{FT}()

Constructor for `AtmosModel` (where `AtmosModel <: BalanceLaw`) for LES
and single stack configurations.
"""
function AtmosModel{FT}(
    ::Union{Type{AtmosLESConfigType}, Type{SingleStackConfigType}},
    param_set::AbstractParameterSet;
    init_state_prognostic::ISP = nothing,
    problem::PR = AtmosProblem(init_state_prognostic = init_state_prognostic),
    orientation::O = FlatOrientation(),
    ref_state::RS = HydrostaticState(DecayingTemperatureProfile{FT}(param_set),),
    turbulence::T = SmagorinskyLilly{FT}(0.21),
    turbconv::TC = NoTurbConv(),
    hyperdiffusion::HD = NoHyperDiffusion(),
    viscoussponge::VS = NoViscousSponge(),
    moisture::M = EquilMoist{FT}(),
    precipitation::P = NoPrecipitation(),
    radiation::R = NoRadiation(),
    source::S = (
        Gravity(),
        Coriolis(),
        GeostrophicForcing(FT, 7.62e-5, 0, 0),
        turbconv_sources(turbconv)...,
    ),
    tracers::TR = NoTracers(),
    lsforcing::LF = NoLSForcing(),
    data_config::DC = nothing,
) where {
    FT <: AbstractFloat,
    ISP,
    PR,
    O,
    RS,
    T,
    TC,
    HD,
    VS,
    M,
    P,
    R,
    S,
    TR,
    LF,
    DC,
}
    @assert !any(isa.(source, Tuple))

    atmos = (
        param_set,
        problem,
        orientation,
        ref_state,
        turbulence,
        turbconv,
        hyperdiffusion,
        viscoussponge,
        moisture,
        precipitation,
        radiation,
        source,
        tracers,
        lsforcing,
        data_config,
    )

    return AtmosModel{FT, typeof.(atmos)...}(atmos...)
end

"""
    AtmosModel{FT}()

Constructor for `AtmosModel` (where `AtmosModel <: BalanceLaw`) for GCM
configurations.
"""
function AtmosModel{FT}(
    ::Type{AtmosGCMConfigType},
    param_set::AbstractParameterSet;
    init_state_prognostic::ISP = nothing,
    problem::PR = AtmosProblem(init_state_prognostic = init_state_prognostic),
    orientation::O = SphericalOrientation(),
    ref_state::RS = HydrostaticState(DecayingTemperatureProfile{FT}(param_set),),
    turbulence::T = SmagorinskyLilly{FT}(C_smag(param_set)),
    turbconv::TC = NoTurbConv(),
    hyperdiffusion::HD = NoHyperDiffusion(),
    viscoussponge::VS = NoViscousSponge(),
    moisture::M = EquilMoist{FT}(),
    precipitation::P = NoPrecipitation(),
    radiation::R = NoRadiation(),
    source::S = (Gravity(), Coriolis(), turbconv_sources(turbconv)...),
    tracers::TR = NoTracers(),
    lsforcing::LF = NoLSForcing(),
    data_config::DC = nothing,
) where {
    FT <: AbstractFloat,
    ISP,
    PR,
    O,
    RS,
    T,
    TC,
    HD,
    VS,
    M,
    P,
    R,
    S,
    TR,
    LF,
    DC,
}

    @assert !any(isa.(source, Tuple))

    atmos = (
        param_set,
        problem,
        orientation,
        ref_state,
        turbulence,
        turbconv,
        hyperdiffusion,
        viscoussponge,
        moisture,
        precipitation,
        radiation,
        source,
        tracers,
        lsforcing,
        data_config,
    )

    return AtmosModel{FT, typeof.(atmos)...}(atmos...)
end

"""
    vars_state(m::AtmosModel, ::Prognostic, FT)

Conserved state variables (prognostic variables).

!!! warning

    The order of the fields for `AtmosModel` needs to match the one for
    `AtmosLinearModel` since a shared state is used
"""
function vars_state(m::AtmosModel, st::Prognostic, FT)
    @vars begin
        # start of inclusion in `AtmosLinearModel`
        ρ::FT
        ρu::SVector{3, FT}
        ρe::FT
        turbulence::vars_state(m.turbulence, st, FT)
        hyperdiffusion::vars_state(m.hyperdiffusion, st, FT)
        moisture::vars_state(m.moisture, st, FT)
        # end of inclusion in `AtmosLinearModel`
        precipitation::vars_state(m.precipitation, st, FT)
        turbconv::vars_state(m.turbconv, st, FT)
        radiation::vars_state(m.radiation, st, FT)
        tracers::vars_state(m.tracers, st, FT)
        lsforcing::vars_state(m.lsforcing, st, FT)
    end
end

"""
    vars_state(m::AtmosModel, ::Gradient, FT)

Pre-transform gradient variables.
"""
function vars_state(m::AtmosModel, st::Gradient, FT)
    @vars begin
        u::SVector{3, FT}
        h_tot::FT
        turbulence::vars_state(m.turbulence, st, FT)
        turbconv::vars_state(m.turbconv, st, FT)
        hyperdiffusion::vars_state(m.hyperdiffusion, st, FT)
        moisture::vars_state(m.moisture, st, FT)
        lsforcing::vars_state(m.lsforcing, st, FT)
        precipitation::vars_state(m.precipitation, st, FT)
        tracers::vars_state(m.tracers, st, FT)
    end
end

"""
    vars_state(m::AtmosModel, ::GradientFlux, FT)

Post-transform gradient variables.
"""
function vars_state(m::AtmosModel, st::GradientFlux, FT)
    @vars begin
        ∇h_tot::SVector{3, FT}
        turbulence::vars_state(m.turbulence, st, FT)
        turbconv::vars_state(m.turbconv, st, FT)
        hyperdiffusion::vars_state(m.hyperdiffusion, st, FT)
        moisture::vars_state(m.moisture, st, FT)
        lsforcing::vars_state(m.lsforcing, st, FT)
        precipitation::vars_state(m.precipitation, st, FT)
        tracers::vars_state(m.tracers, st, FT)
    end
end

"""
    vars_state(m::AtmosModel, ::GradientLaplacian, FT)

Pre-transform hyperdiffusive variables.
"""
function vars_state(m::AtmosModel, st::GradientLaplacian, FT)
    @vars begin
        hyperdiffusion::vars_state(m.hyperdiffusion, st, FT)
    end
end

"""
    vars_state(m::AtmosModel, ::Hyperdiffusive, FT)

Post-transform hyperdiffusive variables.
"""
function vars_state(m::AtmosModel, st::Hyperdiffusive, FT)
    @vars begin
        hyperdiffusion::vars_state(m.hyperdiffusion, st, FT)
    end
end

"""
    vars_state(m::AtmosModel, ::Auxiliary, FT)

Auxiliary variables, such as vertical (stack) integrals, coordinates,
orientation information, reference states, subcomponent auxiliary vars,
debug variables.
"""
function vars_state(m::AtmosModel, st::Auxiliary, FT)
    @vars begin
        ∫dz::vars_state(m, UpwardIntegrals(), FT)
        ∫dnz::vars_state(m, DownwardIntegrals(), FT)
        coord::SVector{3, FT}
        orientation::vars_state(m.orientation, st, FT)
        ref_state::vars_state(m.ref_state, st, FT)
        turbulence::vars_state(m.turbulence, st, FT)
        turbconv::vars_state(m.turbconv, st, FT)
        hyperdiffusion::vars_state(m.hyperdiffusion, st, FT)
        moisture::vars_state(m.moisture, st, FT)
        precipitation::vars_state(m.precipitation, st, FT)
        tracers::vars_state(m.tracers, st, FT)
        radiation::vars_state(m.radiation, st, FT)
        lsforcing::vars_state(m.lsforcing, st, FT)
    end
end

"""
    vars_state(m::AtmosModel, ::UpwardIntegrals, FT)
"""
function vars_state(m::AtmosModel, st::UpwardIntegrals, FT)
    @vars begin
        radiation::vars_state(m.radiation, st, FT)
        turbconv::vars_state(m.turbconv, st, FT)
    end
end

"""
    vars_state(m::AtmosModel, ::DownwardIntegrals, FT)
"""
function vars_state(m::AtmosModel, st::DownwardIntegrals, FT)
    @vars begin
        radiation::vars_state(m.radiation, st, FT)
    end
end

####
#### Forward orientation methods
####
projection_normal(bl, aux, u⃗) =
    projection_normal(bl.orientation, bl.param_set, aux, u⃗)
projection_tangential(bl, aux, u⃗) =
    projection_tangential(bl.orientation, bl.param_set, aux, u⃗)
latitude(bl, aux) = latitude(bl.orientation, aux)
longitude(bl, aux) = longitude(bl.orientation, aux)
altitude(bl, aux) = altitude(bl.orientation, bl.param_set, aux)
vertical_unit_vector(bl, aux) =
    vertical_unit_vector(bl.orientation, bl.param_set, aux)
gravitational_potential(bl, aux) = gravitational_potential(bl.orientation, aux)
∇gravitational_potential(bl, aux) =
    ∇gravitational_potential(bl.orientation, aux)

turbulence_tensors(atmos::AtmosModel, args...) =
    turbulence_tensors(atmos.turbulence, atmos, args...)

export AbstractSource
abstract type AbstractSource end

include("declare_prognostic_vars.jl") # declare prognostic variables
include("multiphysics_types.jl")      # types for multi-physics tendencies
include("tendencies_mass.jl")         # specify mass tendencies
include("tendencies_momentum.jl")     # specify momentum tendencies
include("tendencies_energy.jl")       # specify energy tendencies
include("tendencies_moisture.jl")     # specify moisture tendencies
include("tendencies_precipitation.jl")# specify precipitation tendencies

include("problem.jl")
include("ref_state.jl")
include("moisture.jl")
include("precipitation.jl")
include("thermo_states.jl")
include("radiation.jl")
include("source.jl")
include("tracers.jl")
include("lsforcing.jl")
include("linear.jl")
include("courant.jl")
include("filters.jl")

include("atmos_tendencies.jl")        # specify atmos tendencies
include("get_prognostic_vars.jl")     # get tuple of prognostic variables

"""
    flux_first_order!(
        m::AtmosModel,
        flux::Grad,
        state::Vars,
        aux::Vars,
        t::Real
    )

Computes and assembles non-diffusive fluxes in the model
equations.
"""
@inline function flux_first_order!(
    m::AtmosModel,
    flux::Grad,
    state::Vars,
    aux::Vars,
    t::Real,
    direction,
)
    ρu_pad = SVector(1, 1, 1)
    ts = recover_thermo_state(m, state, aux)
    tend = Flux{FirstOrder}()
    args = (m, state, aux, t, ts, direction)
    flux.ρ = Σfluxes(eq_tends(Mass(), m, tend), args...)
    flux.ρu = Σfluxes(eq_tends(Momentum(), m, tend), args...) .* ρu_pad
    flux.ρe = Σfluxes(eq_tends(Energy(), m, tend), args...)

    flux_first_order!(m.moisture, m, flux, state, aux, t, ts, direction)
    flux_first_order!(m.precipitation, m, flux, state, aux, t, ts, direction)
    flux_first_order!(m.tracers, m, flux, state, aux, t, ts, direction)
    flux_first_order!(m.turbconv, m, flux, state, aux, t, ts, direction)
end

function compute_gradient_argument!(
    atmos::AtmosModel,
    transform::Vars,
    state::Vars,
    aux::Vars,
    t::Real,
)
    ρinv = 1 / state.ρ
    transform.u = ρinv * state.ρu
    ts = recover_thermo_state(atmos, state, aux)
    e_tot = state.ρe * (1 / state.ρ)
    transform.h_tot = total_specific_enthalpy(ts, e_tot)

    compute_gradient_argument!(atmos.moisture, transform, state, aux, t)
    compute_gradient_argument!(atmos.precipitation, transform, state, aux, t)
    compute_gradient_argument!(atmos.turbulence, transform, state, aux, t)
    compute_gradient_argument!(
        atmos.hyperdiffusion,
        atmos,
        transform,
        state,
        aux,
        t,
    )
    compute_gradient_argument!(atmos.tracers, transform, state, aux, t)
    compute_gradient_argument!(atmos.lsforcing, transform, state, aux, t)
    compute_gradient_argument!(atmos.turbconv, atmos, transform, state, aux, t)
end

function compute_gradient_flux!(
    atmos::AtmosModel,
    diffusive::Vars,
    ∇transform::Grad,
    state::Vars,
    aux::Vars,
    t::Real,
)
    diffusive.∇h_tot = ∇transform.h_tot

    # diffusion terms required for SGS turbulence computations
    compute_gradient_flux!(
        atmos.turbulence,
        atmos.orientation,
        diffusive,
        ∇transform,
        state,
        aux,
        t,
    )
    # diffusivity of moisture components
    compute_gradient_flux!(atmos.moisture, diffusive, ∇transform, state, aux, t)
    compute_gradient_flux!(
        atmos.lsforcing,
        diffusive,
        ∇transform,
        state,
        aux,
        t,
    )
    compute_gradient_flux!(
        atmos.precipitation,
        diffusive,
        ∇transform,
        state,
        aux,
        t,
    )
    compute_gradient_flux!(atmos.tracers, diffusive, ∇transform, state, aux, t)
    compute_gradient_flux!(
        atmos.turbconv,
        atmos,
        diffusive,
        ∇transform,
        state,
        aux,
        t,
    )
end

function transform_post_gradient_laplacian!(
    atmos::AtmosModel,
    hyperdiffusive::Vars,
    hypertransform::Grad,
    state::Vars,
    aux::Vars,
    t::Real,
)
    transform_post_gradient_laplacian!(
        atmos.hyperdiffusion,
        atmos,
        hyperdiffusive,
        hypertransform,
        state,
        aux,
        t,
    )
end

"""
    flux_second_order!(
        atmos::AtmosModel,
        flux::Grad,
        state::Vars,
        diffusive::Vars,
        hyperdiffusive::Vars,
        aux::Vars,
        t::Real
    )
Diffusive fluxes in AtmosModel. Viscosity, diffusivity are calculated
in the turbulence subcomponent and accessed within the diffusive flux
function. Contributions from subcomponents are then assembled (pointwise).
"""
@inline function flux_second_order!(
    atmos::AtmosModel,
    flux::Grad,
    state::Vars,
    diffusive::Vars,
    hyperdiffusive::Vars,
    aux::Vars,
    t::Real,
)
    ρu_pad = SVector(1, 1, 1)
    ts = recover_thermo_state(atmos, state, aux)
    tend = Flux{SecondOrder}()
    args = (atmos, state, aux, t, ts, diffusive, hyperdiffusive)
    flux.ρ = Σfluxes(eq_tends(Mass(), atmos, tend), args...) .* ρu_pad
    flux.ρu = Σfluxes(eq_tends(Momentum(), atmos, tend), args...) .* ρu_pad
    flux.ρe = Σfluxes(eq_tends(Energy(), atmos, tend), args...)

    ν, D_t, τ = turbulence_tensors(atmos, state, diffusive, aux, t)

    flux_second_order!(atmos.moisture, flux, state, diffusive, aux, t, D_t)
    flux_second_order!(atmos.precipitation, flux, args...)
    flux_second_order!(
        atmos.hyperdiffusion,
        flux,
        state,
        diffusive,
        hyperdiffusive,
        aux,
        t,
    )
    flux_second_order!(atmos.tracers, flux, state, diffusive, aux, t, D_t)
    flux_second_order!(atmos.turbconv, atmos, flux, state, diffusive, aux, t)
end

@inline function wavespeed(
    m::AtmosModel,
    nM,
    state::Vars,
    aux::Vars,
    t::Real,
    direction,
)
    ρinv = 1 / state.ρ
    u = ρinv * state.ρu
    uN = abs(dot(nM, u))
    ts = recover_thermo_state(m, state, aux)
    ss = soundspeed_air(ts)

    FT = typeof(state.ρ)
    ws = fill(uN + ss, MVector{number_states(m, Prognostic()), FT})
    vars_ws = Vars{vars_state(m, Prognostic(), FT)}(ws)

    wavespeed_tracers!(m.tracers, vars_ws, nM, state, aux, t)

    return ws
end

function update_auxiliary_state!(
    dg::DGModel,
    m::AtmosModel,
    Q::MPIStateArray,
    t::Real,
    elems::UnitRange,
)
    FT = eltype(Q)
    state_auxiliary = dg.state_auxiliary

    if number_states(m, UpwardIntegrals()) > 0
        indefinite_stack_integral!(dg, m, Q, state_auxiliary, t, elems)
        reverse_indefinite_stack_integral!(dg, m, Q, state_auxiliary, t, elems)
    end

    update_auxiliary_state!(nodal_update_auxiliary_state!, dg, m, Q, t, elems)

    # TODO: Remove this hook. This hook was added for implementing
    # the first draft of EDMF, and should be removed so that we can
    # rely on a single vertical element traversal. This hook allows
    # us to compute globally vertical quantities specific to EDMF
    # until we're able to remove them or somehow incorporate them
    # into a higher level hierarchy.
    update_auxiliary_state!(dg, m.turbconv, m, Q, t, elems)

    return true
end

function nodal_update_auxiliary_state!(
    m::AtmosModel,
    state::Vars,
    aux::Vars,
    t::Real,
)
    atmos_nodal_update_auxiliary_state!(m.moisture, m, state, aux, t)
    atmos_nodal_update_auxiliary_state!(m.precipitation, m, state, aux, t)
    atmos_nodal_update_auxiliary_state!(m.radiation, m, state, aux, t)
    atmos_nodal_update_auxiliary_state!(m.tracers, m, state, aux, t)
    turbulence_nodal_update_auxiliary_state!(m.turbulence, m, state, aux, t)
    turbconv_nodal_update_auxiliary_state!(m.turbconv, m, state, aux, t)
end

function integral_load_auxiliary_state!(
    m::AtmosModel,
    integ::Vars,
    state::Vars,
    aux::Vars,
)
    integral_load_auxiliary_state!(m.radiation, integ, state, aux)
    integral_load_auxiliary_state!(m.turbconv, m, integ, state, aux)
end

function integral_set_auxiliary_state!(m::AtmosModel, aux::Vars, integ::Vars)
    integral_set_auxiliary_state!(m.radiation, aux, integ)
    integral_set_auxiliary_state!(m.turbconv, m, aux, integ)
end

function reverse_integral_load_auxiliary_state!(
    m::AtmosModel,
    integ::Vars,
    state::Vars,
    aux::Vars,
)
    reverse_integral_load_auxiliary_state!(m.radiation, integ, state, aux)
end

function reverse_integral_set_auxiliary_state!(
    m::AtmosModel,
    aux::Vars,
    integ::Vars,
)
    reverse_integral_set_auxiliary_state!(m.radiation, aux, integ)
end

function atmos_nodal_init_state_auxiliary!(
    m::AtmosModel,
    aux::Vars,
    tmp::Vars,
    geom::LocalGeometry,
)
    aux.coord = geom.coord
    init_aux_turbulence!(m.turbulence, m, aux, geom)
    atmos_init_aux!(m.ref_state, m, aux, tmp, geom)
    init_aux_hyperdiffusion!(m.hyperdiffusion, m, aux, geom)
    atmos_init_aux!(m.tracers, m, aux, geom)
    init_aux_turbconv!(m.turbconv, m, aux, geom)
    m.problem.init_state_auxiliary(m.problem, m, aux, geom)
end

"""
    init_state_auxiliary!(
        m::AtmosModel,
        aux::Vars,
        grid,
        direction
    )

Initialise auxiliary variables for each AtmosModel subcomponent.
Store Cartesian coordinate information in `aux.coord`.
"""
function init_state_auxiliary!(
    m::AtmosModel,
    state_auxiliary::MPIStateArray,
    grid,
    direction,
)
    init_aux!(m, m.orientation, state_auxiliary, grid, direction)

    init_state_auxiliary!(
        m,
        (m, aux, tmp, geom) ->
            atmos_init_ref_state_pressure!(m.ref_state, m, aux, geom),
        state_auxiliary,
        grid,
        direction,
    )

    ∇p = ∇reference_pressure(m.ref_state, state_auxiliary, grid)

    init_state_auxiliary!(
        m,
        atmos_nodal_init_state_auxiliary!,
        state_auxiliary,
        grid,
        direction;
        state_temporary = ∇p,
    )
end

precompute(atmos::AtmosModel, args, ::Source) =
    (ts = recover_thermo_state(atmos, args.state, args.aux),)

"""
    source!(
        m::AtmosModel,
        source::Vars,
        state::Vars,
        diffusive::Vars,
        aux::Vars,
        t::Real,
        direction::Direction,
    )

Computes (and assembles) source terms `S(Y)` in:
```
∂Y
-- = - ∇ • F + S(Y)
∂t
```
"""
function source!(
    m::AtmosModel,
    source::Vars,
    state::Vars,
    diffusive::Vars,
    aux::Vars,
    t::Real,
    direction,
)
    ρu_pad = SVector(1, 1, 1)
    tend = Source()

    _args = (
        state = state,
        aux = aux,
        t = t,
        direction = direction,
        diffusive = diffusive,
    )

    args = merge(_args, (precomputed = precompute(m, _args, tend),))

    source.ρ = Σsources(eq_tends(Mass(), m, tend), m, args)
    source.ρu = Σsources(eq_tends(Momentum(), m, tend), m, args) .* ρu_pad
    source.ρe = Σsources(eq_tends(Energy(), m, tend), m, args)
    source!(m.moisture, source, m, args)
    source!(m.precipitation, source, m, args)

    atmos_source!(m.source, m, source, state, diffusive, aux, t, direction)
end

"""
    init_state_prognostic!(
        m::AtmosModel,
        state::Vars,
        aux::Vars,
        localgeo,
        t,
        args...,
    )

Initialise state variables. `args...` provides an option to include
configuration data (current use cases include problem constants,
spline-interpolants).
"""
function init_state_prognostic!(
    m::AtmosModel,
    state::Vars,
    aux::Vars,
    localgeo,
    t,
    args...,
)
    m.problem.init_state_prognostic(
        m.problem,
        m,
        state,
        aux,
        localgeo,
        t,
        args...,
    )
end

roe_average(ρ⁻, ρ⁺, var⁻, var⁺) =
    (sqrt(ρ⁻) * var⁻ + sqrt(ρ⁺) * var⁺) / (sqrt(ρ⁻) + sqrt(ρ⁺))

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

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

    FT = eltype(fluxᵀn)
    param_set = balance_law.param_set
    _cv_d::FT = cv_d(param_set)
    _T_0::FT = T_0(param_set)

    Φ = gravitational_potential(balance_law, state_auxiliary⁻)

    ρ⁻ = state_prognostic⁻.ρ
    ρu⁻ = state_prognostic⁻.ρu
    ρe⁻ = state_prognostic⁻.ρe
    ts⁻ = recover_thermo_state(
        balance_law,
        balance_law.moisture,
        state_prognostic⁻,
        state_auxiliary⁻,
    )

    u⁻ = ρu⁻ / ρ⁻
    uᵀn⁻ = u⁻' * normal_vector
    e⁻ = ρe⁻ / ρ⁻
    h⁻ = total_specific_enthalpy(ts⁻, e⁻)
    p⁻ = air_pressure(ts⁻)
    c⁻ = soundspeed_air(ts⁻)

    ρ⁺ = state_prognostic⁺.ρ
    ρu⁺ = state_prognostic⁺.ρu
    ρe⁺ = state_prognostic⁺.ρe

    # TODO: state_auxiliary⁺ is not up-to-date
    # with state_prognostic⁺ on the boundaries
    ts⁺ = recover_thermo_state(
        balance_law,
        balance_law.moisture,
        state_prognostic⁺,
        state_auxiliary⁺,
    )

    u⁺ = ρu⁺ / ρ⁺
    uᵀn⁺ = u⁺' * normal_vector
    e⁺ = ρe⁺ / ρ⁺
    h⁺ = total_specific_enthalpy(ts⁺, e⁺)
    p⁺ = air_pressure(ts⁺)
    c⁺ = soundspeed_air(ts⁺)

    ρ̃ = sqrt(ρ⁻ * ρ⁺)
    ũ = roe_average(ρ⁻, ρ⁺, u⁻, u⁺)
    h̃ = roe_average(ρ⁻, ρ⁺, h⁻, h⁺)
    c̃ = sqrt(roe_average(ρ⁻, ρ⁺, c⁻^2, c⁺^2))

    ũᵀn = ũ' * normal_vector

    Δρ = ρ⁺ - ρ⁻
    Δp = p⁺ - p⁻
    Δu = u⁺ - u⁻
    Δuᵀn = Δu' * normal_vector

    w1 = abs(ũᵀn - c̃) * (Δp - ρ̃ * c̃ * Δuᵀn) / (2 * c̃^2)
    w2 = abs(ũᵀn + c̃) * (Δp + ρ̃ * c̃ * Δuᵀn) / (2 * c̃^2)
    w3 = abs(ũᵀn) * (Δρ - Δp / c̃^2)
    w4 = abs(ũᵀn) * ρ̃

    fluxᵀn.ρ -= (w1 + w2 + w3) / 2
    fluxᵀn.ρu -=
        (
            w1 * (ũ - c̃ * normal_vector) +
            w2 * (ũ + c̃ * normal_vector) +
            w3 * ũ +
            w4 * (Δu - Δuᵀn * normal_vector)
        ) / 2
    fluxᵀn.ρe -=
        (
            w1 * (h̃ - c̃ * ũᵀn) +
            w2 * (h̃ + c̃ * ũᵀn) +
            w3 * (ũ' * ũ / 2 + Φ - _T_0 * _cv_d) +
            w4 * (ũ' * Δu - ũᵀn * Δuᵀn)
        ) / 2

    if !(balance_law.tracers isa NoTracers)
        ρχ⁻ = state_prognostic⁻.tracers.ρχ
        χ⁻ = ρχ⁻ / ρ⁻

        ρχ⁺ = state_prognostic⁺.tracers.ρχ
        χ⁺ = ρχ⁺ / ρ⁺

        χ̃ = roe_average(ρ⁻, ρ⁺, χ⁻, χ⁺)
        Δρχ = ρχ⁺ - ρχ⁻

        wt = abs(ũᵀn) * (Δρχ - χ̃ * Δp / c̃^2)

        fluxᵀn.tracers.ρχ -= ((w1 + w2) * χ̃ + wt) / 2
    end
end

"""
    NumericalFluxFirstOrder()
        ::HLLCNumericalFlux,
        balance_law::AtmosModel,
        fluxᵀn,
        normal_vector,
        state_prognostic⁻,
        state_auxiliary⁻,
        state_prognostic⁺,
        state_auxiliary⁺,
        t,
        direction,
    )

An implementation of the numerical flux based on the HLLC method for
the AtmosModel. For more information on this particular implementation,
see Chapter 10.4 in the provided reference below.

## References

 - [Toro2013](@cite)

"""
function numerical_flux_first_order!(
    ::HLLCNumericalFlux,
    balance_law::AtmosModel,
    fluxᵀn::Vars{S},
    normal_vector::SVector,
    state_prognostic⁻::Vars{S},
    state_auxiliary⁻::Vars{A},
    state_prognostic⁺::Vars{S},
    state_auxiliary⁺::Vars{A},
    t,
    direction,
) where {S, A}
    FT = eltype(fluxᵀn)
    num_state_prognostic = number_states(balance_law, Prognostic())
    param_set = balance_law.param_set

    # Extract the first-order fluxes from the AtmosModel (underlying BalanceLaw)
    # and compute normals on the positive + and negative - sides of the
    # interior facets
    flux⁻ = similar(parent(fluxᵀn), Size(3, num_state_prognostic))
    fill!(flux⁻, -zero(FT))
    flux_first_order!(
        balance_law,
        Grad{S}(flux⁻),
        state_prognostic⁻,
        state_auxiliary⁻,
        t,
        direction,
    )
    fluxᵀn⁻ = flux⁻' * normal_vector

    flux⁺ = similar(flux⁻)
    fill!(flux⁺, -zero(FT))
    flux_first_order!(
        balance_law,
        Grad{S}(flux⁺),
        state_prognostic⁺,
        state_auxiliary⁺,
        t,
        direction,
    )
    fluxᵀn⁺ = flux⁺' * normal_vector

    # Extract relevant fields and thermodynamic variables defined on
    # the positive + and negative - sides of the interior facets
    ρ⁻ = state_prognostic⁻.ρ
    ρu⁻ = state_prognostic⁻.ρu
    ρe⁻ = state_prognostic⁻.ρe
    ts⁻ = recover_thermo_state(
        balance_law,
        balance_law.moisture,
        state_prognostic⁻,
        state_auxiliary⁻,
    )

    u⁻ = ρu⁻ / ρ⁻
    c⁻ = soundspeed_air(ts⁻)

    uᵀn⁻ = u⁻' * normal_vector
    e⁻ = ρe⁻ / ρ⁻
    p⁻ = air_pressure(ts⁻)

    ρ⁺ = state_prognostic⁺.ρ
    ρu⁺ = state_prognostic⁺.ρu
    ρe⁺ = state_prognostic⁺.ρe
    ts⁺ = recover_thermo_state(
        balance_law,
        balance_law.moisture,
        state_prognostic⁺,
        state_auxiliary⁺,
    )

    u⁺ = ρu⁺ / ρ⁺
    uᵀn⁺ = u⁺' * normal_vector
    e⁺ = ρe⁺ / ρ⁺
    p⁺ = air_pressure(ts⁺)
    c⁺ = soundspeed_air(ts⁺)

    # Wave speeds estimates S⁻ and S⁺
    S⁻ = min(uᵀn⁻ - c⁻, uᵀn⁺ - c⁺)
    S⁺ = max(uᵀn⁻ + c⁻, uᵀn⁺ + c⁺)

    # Compute the middle wave speed S⁰ in the contact/star region
    S⁰ =
        (p⁺ - p⁻ + ρ⁻ * uᵀn⁻ * (S⁻ - uᵀn⁻) - ρ⁺ * uᵀn⁺ * (S⁺ - uᵀn⁺)) /
        (ρ⁻ * (S⁻ - uᵀn⁻) - ρ⁺ * (S⁺ - uᵀn⁺))

    p⁰ =
        (
            p⁺ +
            p⁻ +
            ρ⁻ * (S⁻ - uᵀn⁻) * (S⁰ - uᵀn⁻) +
            ρ⁺ * (S⁺ - uᵀn⁺) * (S⁰ - uᵀn⁺)
        ) / 2

    # Compute p * D = p * (0, n₁, n₂, n₃, S⁰)
    pD = @MVector zeros(FT, num_state_prognostic)
    if balance_law.ref_state isa HydrostaticState
        # pressure should be continuous but it doesn't hurt to average
        ref_p⁻ = state_auxiliary⁻.ref_state.p
        ref_p⁺ = state_auxiliary⁺.ref_state.p
        ref_p⁰ = (ref_p⁻ + ref_p⁺) / 2

        momentum_p = p⁰ - ref_p⁰
    else
        momentum_p = p⁰
    end

    pD[2] = momentum_p * normal_vector[1]
    pD[3] = momentum_p * normal_vector[2]
    pD[4] = momentum_p * normal_vector[3]
    pD[5] = p⁰ * S⁰

    # Computes both +/- sides of intermediate flux term flux⁰
    flux⁰⁻ =
        (S⁰ * (S⁻ * parent(state_prognostic⁻) - fluxᵀn⁻) + S⁻ * pD) / (S⁻ - S⁰)
    flux⁰⁺ =
        (S⁰ * (S⁺ * parent(state_prognostic⁺) - fluxᵀn⁺) + S⁺ * pD) / (S⁺ - S⁰)

    if 0 <= S⁻
        parent(fluxᵀn) .= fluxᵀn⁻
    elseif S⁻ < 0 <= S⁰
        parent(fluxᵀn) .= flux⁰⁻
    elseif S⁰ < 0 <= S⁺
        parent(fluxᵀn) .= flux⁰⁺
    else # 0 > S⁺
        parent(fluxᵀn) .= fluxᵀn⁺
    end
end

function numerical_flux_first_order!(
    numerical_flux::RoeNumericalFluxMoist,
    balance_law::AtmosModel,
    fluxᵀn::Vars{S},
    normal_vector::SVector,
    state_prognostic⁻::Vars{S},
    state_auxiliary⁻::Vars{A},
    state_prognostic⁺::Vars{S},
    state_auxiliary⁺::Vars{A},
    t,
    direction,
) where {S, A}
    balance_law.moisture isa EquilMoist ||
        error("Must use a EquilMoist model for RoeNumericalFluxMoist")
    numerical_flux_first_order!(
        CentralNumericalFluxFirstOrder(),
        balance_law,
        fluxᵀn,
        normal_vector,
        state_prognostic⁻,
        state_auxiliary⁻,
        state_prognostic⁺,
        state_auxiliary⁺,
        t,
        direction,
    )

    FT = eltype(fluxᵀn)
    param_set = balance_law.param_set
    _cv_d::FT = cv_d(param_set)
    _T_0::FT = T_0(param_set)
    γ::FT = cp_d(param_set) / cv_d(param_set)
    _e_int_v0::FT = e_int_v0(param_set)
    Φ = gravitational_potential(balance_law, state_auxiliary⁻)

    ρ⁻ = state_prognostic⁻.ρ
    ρu⁻ = state_prognostic⁻.ρu
    ρe⁻ = state_prognostic⁻.ρe
    ρq_tot⁻ = state_prognostic⁻.moisture.ρq_tot

    u⁻ = ρu⁻ / ρ⁻
    e⁻ = ρe⁻ / ρ⁻
    ts⁻ = recover_thermo_state(balance_law, state_prognostic⁻, state_auxiliary⁻)
    h⁻ = total_specific_enthalpy(ts⁻, e⁻)
    qt⁻ = ρq_tot⁻ / ρ⁻
    c⁻ = soundspeed_air(ts⁻)

    ρ⁺ = state_prognostic⁺.ρ
    ρu⁺ = state_prognostic⁺.ρu
    ρe⁺ = state_prognostic⁺.ρe
    ρq_tot⁺ = state_prognostic⁺.moisture.ρq_tot

    u⁺ = ρu⁺ / ρ⁺
    e⁺ = ρe⁺ / ρ⁺
    ts⁺ = recover_thermo_state(balance_law, state_prognostic⁺, state_auxiliary⁺)
    h⁺ = total_specific_enthalpy(ts⁺, e⁺)
    qt⁺ = ρq_tot⁺ / ρ⁺
    c⁺ = soundspeed_air(ts⁺)
    ũ = roe_average(ρ⁻, ρ⁺, u⁻, u⁺)
    e_tot = roe_average(ρ⁻, ρ⁺, e⁻, e⁺)
    h̃ = roe_average(ρ⁻, ρ⁺, h⁻, h⁺)
    qt = roe_average(ρ⁻, ρ⁺, qt⁻, qt⁺)
    ρ = sqrt(ρ⁻ * ρ⁺)
    e_int⁻ = internal_energy(ts⁻)
    e_int⁺ = internal_energy(ts⁺)
    e_int = roe_average(ρ⁻, ρ⁺, e_int⁻, e_int⁺)
    ts = PhaseEquil(
        param_set,
        e_int,
        ρ,
        qt,
        balance_law.moisture.maxiter,
        balance_law.moisture.tolerance,
    )
    c̃ = sqrt((γ - 1) * (h̃ - (ũ[1]^2 + ũ[2]^2 + ũ[3]^2) / 2))
    (R_m, _cp_m, _cv_m, gamma) = gas_constants(ts)
    # chosen by fair dice roll
    # guaranteed to be random
    ω = FT(π) / 3
    δ = FT(π) / 5
    random_unit_vector = SVector(sin(ω) * cos(δ), cos(ω) * cos(δ), sin(δ))
    # tangent space basis
    τ1 = random_unit_vector × normal_vector
    τ2 = τ1 × normal_vector
    ũᵀn⁻ = u⁻' * normal_vector
    ũᵀn⁺ = u⁺' * normal_vector
    ũᵀn = ũ' * normal_vector
    ũc̃⁻ = ũ + c̃ * normal_vector
    ũc̃⁺ = ũ - c̃ * normal_vector
    e_kin_pot = h̃ - _e_int_v0 * qt - _cp_m * c̃^2 / R_m
    if (numerical_flux.LM == true)
        Mach⁺ = sqrt(u⁺' * u⁺) / c⁺
        Mach⁻ = sqrt(u⁻' * u⁻) / c⁻
        Mach = (Mach⁺ + Mach⁻) / 2
        c̃_LM = c̃ * min(Mach * sqrt(4 + (1 - Mach^2)^2) / (1 + Mach^2), 1)
    else
        c̃_LM = c̃
    end
    #Standard Roe
    Λ = SDiagonal(
        abs(ũᵀn - c̃_LM),
        abs(ũᵀn),
        abs(ũᵀn),
        abs(ũᵀn),
        abs(ũᵀn + c̃_LM),
        abs(ũᵀn),
    )
    #Harten Hyman
    if (numerical_flux.HH == true)
        Λ = SDiagonal(
            max(
                abs(ũᵀn - c̃_LM),
                max(
                    0,
                    ũᵀn - c̃_LM - (u⁻' * normal_vector - c⁻),
                    u⁺' * normal_vector - c⁺ - (ũᵀn - c̃_LM),
                ),
            ),
            max(
                abs(ũᵀn),
                max(
                    0,
                    ũᵀn - (u⁻' * normal_vector),
                    u⁺' * normal_vector - (ũᵀn),
                ),
            ),
            max(
                abs(ũᵀn),
                max(
                    0,
                    ũᵀn - (u⁻' * normal_vector),
                    u⁺' * normal_vector - (ũᵀn),
                ),
            ),
            max(
                abs(ũᵀn),
                max(
                    0,
                    ũᵀn - (u⁻' * normal_vector),
                    u⁺' * normal_vector - (ũᵀn),
                ),
            ),
            max(
                abs(ũᵀn + c̃_LM),
                max(
                    0,
                    ũᵀn + c̃_LM - (u⁻' * normal_vector + c⁻),
                    u⁺' * normal_vector + c⁺ - (ũᵀn + c̃_LM),
                ),
            ),
            max(
                abs(ũᵀn),
                max(
                    0,
                    ũᵀn - (u⁻' * normal_vector),
                    u⁺' * normal_vector - (ũᵀn),
                ),
            ),
        )
    end
    if (numerical_flux.LV == true)
        #Pseudo LeVeque Fix
        δ_L_1 = max(0, ũᵀn - ũᵀn⁻)
        δ_L_2 = max(0, ũᵀn - c̃_LM - (ũᵀn⁻ - c⁻))
        δ_L_3 = max(0, ũᵀn + c̃_LM - (ũᵀn⁻ + c⁻))
        δ_R_1 = max(0, ũᵀn⁺ - ũᵀn)
        δ_R_2 = max(0, ũᵀn⁺ - c⁺ - (ũᵀn - c̃_LM))
        δ_R_3 = max(0, ũᵀn⁺ + c⁺ - (ũᵀn + c̃_LM))
        if (ũᵀn < δ_L_1 && ũᵀn > -δ_R_1)
            qa1 = ((δ_L_1 - δ_R_1) * ũᵀn + 2 * δ_L_1 * δ_R_1) / (δ_L_1 + δ_R_1)
        else
            qa1 = abs(ũᵀn)
        end
        if (ũᵀn - c̃ < δ_L_2 && ũᵀn - c̃_LM > -δ_R_2)
            qa2 =
                ((δ_L_2 - δ_R_2) * (ũᵀn - c̃_LM) + 2 * δ_L_2 * δ_R_2) /
                (δ_L_2 + δ_R_2)
        else
            qa2 = abs(ũᵀn - c̃_LM)
        end
        if (ũᵀn + c̃_LM < δ_L_3 && ũᵀn + c̃ > -δ_R_3)
            qa3 =
                ((δ_L_3 - δ_R_3) * (ũᵀn + c̃_LM) + 2 * δ_R_3 * δ_R_3) /
                (δ_L_3 + δ_R_3)
        else
            qa3 = abs(ũᵀn + c̃_LM)
        end
        Λ = SDiagonal(qa2, qa1, qa1, qa1, qa3, qa1)
    end
    if (numerical_flux.LVPP == true)
        #PosPreserving with LeVeque
        b_L = min(ũᵀn - c̃_LM, ũᵀn⁻ - c⁻)
        b_R = max(ũᵀn + c̃_LM, ũᵀn⁺ + c⁺)
        b⁻ = min(0, b_L)
        b⁺ = max(0, b_R)
        δ_L_1 = max(0, ũᵀn - b⁻)
        δ_L_2 = max(0, ũᵀn - c̃_LM - b⁻)
        δ_L_3 = max(0, ũᵀn + c̃_LM - b⁻)
        δ_R_1 = max(0, b⁺ - ũᵀn)
        δ_R_2 = max(0, b⁺ - (ũᵀn - c̃_LM))
        δ_R_3 = max(0, b⁺ - (ũᵀn + c̃_LM))
        if (ũᵀn < δ_L_1 && ũᵀn > -δ_R_1)
            qa1 = ((δ_L_1 - δ_R_1) * ũᵀn + 2 * δ_L_1 * δ_R_1) / (δ_L_1 + δ_R_1)
        else
            qa1 = abs(ũᵀn)
        end
        if (ũᵀn - c̃_LM < δ_L_2 && ũᵀn - c̃_LM > -δ_R_2)
            qa2 =
                ((δ_L_2 - δ_R_2) * (ũᵀn - c̃) + 2 * δ_L_2 * δ_R_2) /
                (δ_L_2 + δ_R_2)
        else
            qa2 = abs(ũᵀn - c̃_LM)
        end
        if (ũᵀn + c̃_LM < δ_L_3 && ũᵀn + c̃_LM > -δ_R_3)
            qa3 =
                ((δ_L_3 - δ_R_3) * (ũᵀn + c̃_LM) + 2 * δ_R_3 * δ_R_3) /
                (δ_L_3 + δ_R_3)
        else
            qa3 = abs(ũᵀn + c̃_LM)
        end
        Λ = SDiagonal(qa2, qa1, qa1, qa1, qa3, qa1)
    end

    M = hcat(
        SVector(1, ũc̃⁺[1], ũc̃⁺[2], ũc̃⁺[3], h̃ - c̃ * ũᵀn, qt),
        SVector(0, τ1[1], τ1[2], τ1[3], τ1' * ũ, 0),
        SVector(0, τ2[1], τ2[2], τ2[3], τ2' * ũ, 0),
        SVector(1, ũ[1], ũ[2], ũ[3], ũ' * ũ / 2 + Φ - _T_0 * _cv_m, 0),
        SVector(1, ũc̃⁻[1], ũc̃⁻[2], ũc̃⁻[3], h̃ + c̃ * ũᵀn, qt),
        SVector(0, 0, 0, 0, _e_int_v0, 1),
    )
    Δρ = ρ⁺ - ρ⁻
    Δρu = ρu⁺ - ρu⁻
    Δρe = ρe⁺ - ρe⁻
    Δρq_tot = ρq_tot⁺ - ρq_tot⁻
    Δstate = SVector(Δρ, Δρu[1], Δρu[2], Δρu[3], Δρe, Δρq_tot)
    parent(fluxᵀn) .-= M * Λ * (M \ Δstate) / 2
end

end # module


text:
#
# Structure to contain the trajectory data
#
struct MDTraj

  n :: Int64
  nsteps :: Int64
  atoms :: Vector{Atoms}
  potential :: Vector{Float64}
  kinetic :: Vector{Float64}
  total :: Vector{Float64}
  time :: Vector{Float64}
  nenc :: Vector{Float64} # Number of encounters

  # Results along the trajectory
  U :: Vector{Float64} # Susceptible
  S :: Vector{Float64} # Sick
  D :: Vector{Float64} # Dead
  I :: Vector{Float64} # Immune

end

# Initialize from number of atoms and number of steps to be saved
function MDTraj(n :: Int64, nsave :: Int64)
  nsteps = nsave
  traj = MDTraj(
                n, # number of particles
                nsteps, # nsteps
                Vector{Atoms}(undef,nsteps), # atoms
                zeros(nsteps), # potential
                zeros(nsteps), # kinetic
                zeros(nsteps), # total
                zeros(nsteps), # time
                zeros(nsteps), # nenc
                zeros(nsteps), # U
                zeros(nsteps), # S
                zeros(nsteps), # D
                zeros(nsteps), # I
                )
  for i in 1:nsteps
    traj.atoms[i] = Atoms(n)
  end
  return traj
end

ext: jl


text:
using TuringModels, StatsFuns

Turing.setadbackend(:reversediff);

d = CSV.read(joinpath(@__DIR__, "..", "..", "data", "chimpanzees.csv"), delim=';');
size(d) # Should be 504x8

# pulled_left, actors, condition, prosoc_left
@model m10_4(y, actors, x₁, x₂) = begin
    # Number of unique actors in the data set
    N_actor = length(unique(actors))

    # Set an TArray for the priors/param
    α ~ filldist(Normal(0, 10), N_actor)
    βp ~ Normal(0, 10)
    βpC ~ Normal(0, 10)

    logits = α[actors] .+ (βp .+ βpC * x₁) .* x₂
    y .~ BinomialLogit.(1, logits)
end;

chns = sample(m10_4(d[:,:pulled_left], d[:,:actor],d[:,:condition],
  d[:,:prosoc_left]), Turing.NUTS(0.65), 1000);

# Rethinking/CmdStan results

m_10_04s_result = "
Iterations = 1:1000
Thinning interval = 1
Chains = 1,2,3,4
Samples per chain = 1000

Empirical Posterior Estimates:
        Mean        SD       Naive SE       MCSE      ESS
a.1 -0.74503184 0.26613979 0.0042080396 0.0060183398 1000
a.2 10.77955494 5.32538998 0.0842018089 0.1269148045 1000
a.3 -1.04982353 0.28535997 0.0045119373 0.0049074219 1000
a.4 -1.04898135 0.28129307 0.0044476339 0.0056325117 1000
a.5 -0.74390933 0.26949936 0.0042611590 0.0052178124 1000
a.6  0.21599365 0.26307574 0.0041595927 0.0045153523 1000
a.7  1.81090866 0.39318577 0.0062168129 0.0071483527 1000
bp  0.83979926 0.26284676 0.0041559722 0.0059795826 1000
bpC -0.12913322 0.29935741 0.0047332562 0.0049519863 1000
";

# Describe the draws

describe(chns)

# End of 10/m10.04t.jl


text:
using Documenter, PlutoUI

# i can't convice Documenter.jl to stop printing `PlutoUI.BuiltinsNotebook.Slider`, etc, so here is a nasty override:
Base.show(io::IO, m::Module) =
    if parentmodule(m) === PlutoUI
        write(io, "PlutoUI")
    else
        write(io, nameof(m))
    end

makedocs(;
    sitename = "PlutoUI",
    modules = [PlutoUI]
)


text:
using InteractionNetworkModels, Plots, Distributions
using StructuredDistances
using StatsPlots, Plots.Measures, StatsBase

E = [[1,2,1,2],
    [1,2,1],
    [3,4,3],
    [3,4],
    [1,2],
    [1,2,1],
    [1,2,3],
    [4,5],
    [7,8]]
d = MatchingDistance(FastLCS(101))
K_inner, K_outer = (DimensionRange(1,10), DimensionRange(1,25))
model = SIM(
    E, 3.5,
    d,
    1:10,
    K_inner, K_outer)

mcmc_sampler = SimMcmcInsertDelete(
    ν_ed=1, ν_td=1, β=0.7,
    len_dist=TrGeometric(0.8, 1, model.K_inner.u),
    burn_in=2000, lag=75, init=InitRandIns(10)
)

mcmc_sampler_prop = SimMcmcInsertDeleteProportional(
    ν_ed=1, ν_td=1, β=0.7,
    len_dist=TrGeometric(0.8, 1, model.K_inner.u),
    burn_in=2000, lag=75, init=InitRandIns(10)
)

@time test = mcmc_sampler(
    model, desired_samples=5000, burn_in=0, lag=1
)
plot(test)
summaryplot(test)
@time test = mcmc_sampler_prop(
    model, desired_samples=5000, burn_in=0, lag=1
)
plot(test)
summaryplot(test)

test.sample
@time test = mcmc_sampler(
    model, desired_samples=500
)
plot(test)
summaryplot(test)

@time mcmc_out = mcmc_sampler(
    model,
    desired_samples=50,
    lag=500,
    burn_in=10000
    )

plot(mcmc_out)
summaryplot(mcmc_out)

data = mcmc_out.sample
E_prior = SIM(E, 0.1, model.dist, model.V, model.K_inner, model.K_outer)
γ_prior = Uniform(0.5,7.0)

posterior = SimPosterior(data, E_prior, γ_prior)

# Construct posterior sampler
posterior_sampler = SimIexInsertDelete(
    mcmc_sampler,
    len_dist=TrGeometric(0.9,K_inner.l,K_inner.u),
    ν_ed=1, ν_td=1,
    β=0.7, ε=0.3
)

posterior_sampler_prop = SimIexInsertDeleteProportional(
    mcmc_sampler_prop,
    len_dist=TrGeometric(0.9,K_inner.l,K_inner.u),
    ν_ed=1, ν_td=1,
    β=0.7, ε=0.3
)

# Mode Conditional
E_init = sample_frechet_mean(posterior.data, posterior.dist)[1]
d(E_init, E)

E_init = [[10,10]]
@time posterior_out = posterior_sampler(
    posterior,
    # 4.9,
    desired_samples=100, lag=1, burn_in=0,
    S_init=E_init, γ_init=2.8,
    aux_init_at_prev=true,
);
plot(
    posterior_out, E, size=(600,400),
    xlabel=["" "Sample Index"],
    ylabel=["Distance to True Mode" "γ"],
    yguidefontsize=[9 9],
    xguidefontsize=9
)
savefig("C:/users/boltg/OneDrive - Lancaster University/PhD/Outputs/Posters/Data on Lake 2022/img/posterior_sample.pdf")

model.mode
posterior_out.S_sample
plot(posterior_out, E)
plot(length.(posterior_out.S_sample))
findall(diff(length.(posterior_out.S_sample)).>0)

plot(length.(posterior_out.S_sample))
print_map_est(posterior_out)

plot(posterior_out, E)

E_init[1] = [1,2,1,4,4,4,4,4,4,4,4,4,4,1,2]
@time posterior_out = posterior_sampler(
    posterior,
    # 4.9,
    desired_samples=1000, lag=1, burn_in=0,
    S_init=E_init, γ_init=2.8,
    aux_init_at_prev=true,
);
plot(posterior_out, E)
posterior_out.S_sample
@time posterior_out_prop = posterior_sampler_prop(
    posterior,
    # 4.9,
    desired_samples=1000, lag=1, burn_in=0,
    S_init=E_init, γ_init=2.8,
    aux_init_at_prev=true,
);
plot(posterior_out_prop, E)
posterior_out_prop.S_sample

posterior_out.S_sample
# Dispersion conditional
E_fix = Eᵐ
posterior_out = posterior_sampler(
    target,
    E_fix,
    desired_samples=1000, lag=1, burn_in=0,
    init=4.6
)
plot(posterior_out)


text:
module Adapt

using LinearAlgebra

adapt_(T, x) = x

adapt(T, x) = adapt_(T, x)

# Base integrations

adapt(T, x::Adjoint) = Adjoint(adapt(T, parent(x)))

end # module


text:
"""
    function configurator_(kernel)

Function to configurate the number of blocks and threads to be launched at runtime.
"""
function configurator_(kernel,N)
    config = launch_configuration(kernel.fun)

    threads = Base.min(max(N,1), config.threads)
    blocks = Base.min(cld(max(N,1), threads),config.blocks)

    return (max(1,round(Int,threads)), max(1,round(Int,blocks)))
end

max_stars_repo_name: gatocor/embryogenesisJulia

max_issues_repo_name: gatocor/embryogenesisJulia

max_forks_repo_name: gatocor/embryogenesisJulia


text:
export
    radrec,
    rav2xf,
    raxisa,
    reccyl,
    recgeo,
    reclat,
    recpgr,
    recsph,
    recrad,
    removc!,
    removc,
    removd!,
    removd,
    removi!,
    removi,
    reordc,
    reordd,
    reordi,
    reordl,
    rotate,
    rotmat,
    rotvec,
    rpd,
    rquad

"""
    radrec(range, ra, dec)

Convert from range, right ascension, and declination to rectangular coordinates.

### Arguments ###

- `range`: Distance of a point from the origin
- `ra`: Right ascension of point in radians
- `dec`: Declination of point in radians

### Output ###

Returns the rectangular coordinates of the point.

### References ###

- [NAIF Documentation](https://naif.jpl.nasa.gov/pub/naif/toolkit_docs/C/cspice/radrec_c.html)
"""
function radrec(range, ra, dec)
    rectan = Array{SpiceDouble}(undef, 3)
    ccall((:radrec_c, libcspice), Cvoid,
          (SpiceDouble, SpiceDouble, SpiceDouble, Ref{SpiceDouble}),
          range, ra, dec, rectan)
    rectan
end

"""
    rav2xf(rot, av)

Determine a state transformation matrix from a rotation matrix and the angular
velocity of the rotation.

### Arguments ###

- `rot`: Rotation matrix
- `av`: Angular velocity vector

### Output ###

Returns state transformation matrix associated with `rot` and `av`.

### References ###

- [NAIF Documentation](https://naif.jpl.nasa.gov/pub/naif/toolkit_docs/C/cspice/rav2xf_c.html)
"""
function rav2xf(rot, av)
    @checkdims 3 3 rot
    @checkdims 3 av
    xform = Array{SpiceDouble}(undef, 6, 6)
    ccall((:rav2xf_c, libcspice), Cvoid,
          (Ref{SpiceDouble}, Ref{SpiceDouble}, Ref{SpiceDouble}),
          permutedims(rot), av, xform)
    permutedims(xform)
end

"""
    raxisa(matrix)

Compute the axis of the rotation given by an input matrix and the angle of the rotation
about that axis.

### Arguments ###

- `matrix`: A 3x3 rotation matrix

### Output ###

- `axis`: Axis of the rotation
- `angle`: Angle through which the rotation is performed

### References ###

- [NAIF Documentation](https://naif.jpl.nasa.gov/pub/naif/toolkit_docs/C/cspice/raxisa_c.html)
"""
function raxisa(matrix)
    @checkdims 3 3 matrix
    axis = Array{SpiceDouble}(undef, 3)
    angle = Ref{SpiceDouble}()
    ccall((:raxisa_c, libcspice), Cvoid,
          (Ref{SpiceDouble}, Ref{SpiceDouble}, Ref{SpiceDouble}),
          permutedims(matrix), axis, angle)
    handleerror()
    axis, angle[]
end

"""
    reccyl(rectan)

Convert from rectangular to cylindrical coordinates.

### Arguments ###

- `rectan`: Rectangular coordinates of a point

### Output ###

- `r`: Distance of the point from the Z axis
- `lon`: Angle (radians) of the point from the XZ plane
- `z`: Height of the point above the XY plane

### References ###

- [NAIF Documentation](https://naif.jpl.nasa.gov/pub/naif/toolkit_docs/C/cspice/reccyl_c.html)
"""
function reccyl(rectan)
    @checkdims 3 rectan
    r = Ref{SpiceDouble}()
    lon = Ref{SpiceDouble}()
    z = Ref{SpiceDouble}()
    ccall((:reccyl_c, libcspice), Cvoid,
          (Ref{SpiceDouble}, Ref{SpiceDouble}, Ref{SpiceDouble}, Ref{SpiceDouble}),
          rectan, r, lon, z)
    r[], lon[], z[]
end

"""
    recgeo(rectan, re, f)

Convert from rectangular coordinates to geodetic coordinates.

### Arguments ###

- `rectan`: Rectangular coordinates of a point
- `re`: Equatorial radius of the reference spheroid
- `f`: Flattening coefficient

### Output ###

- `lon`: Geodetic longitude of the point (radians)
- `lat`: Geodetic latitude  of the point (radians)
- `alt`: Altitude of the point above reference spheroid

### References ###

- [NAIF Documentation](https://naif.jpl.nasa.gov/pub/naif/toolkit_docs/C/cspice/recgeo_c.html)
"""
function recgeo(rectan, re, f)
    @checkdims 3 rectan
    lon = Ref{SpiceDouble}()
    lat = Ref{SpiceDouble}()
    alt = Ref{SpiceDouble}()
    ccall((:recgeo_c, libcspice), Cvoid,
          (Ref{SpiceDouble}, SpiceDouble, SpiceDouble,
           Ref{SpiceDouble}, Ref{SpiceDouble}, Ref{SpiceDouble}),
          rectan, re, f, lon, lat, alt)
    handleerror()
    lon[], lat[], alt[]
end

"""
    reclat(rectan)

Convert from rectangular coordinates to latitudinal coordinates.

### Arguments ###

- `rectan`: Rectangular coordinates of a point

### Output ###

Returns a tuple consisting of:

- `rad`: Distance of the point from the origin
- `lon`: Planetographic longitude of the point (radians)
- `lat`: Planetographic latitude of the point (radians)

### References ###

- [NAIF Documentation](https://naif.jpl.nasa.gov/pub/naif/toolkit_docs/C/cspice/reclat_c.html)
"""
function reclat(rectan)
    @checkdims 3 rectan
    lon = Ref{SpiceDouble}()
    lat = Ref{SpiceDouble}()
    rad = Ref{SpiceDouble}()
    ccall((:reclat_c, libcspice), Cvoid,
          (Ref{SpiceDouble}, Ref{SpiceDouble}, Ref{SpiceDouble}, Ref{SpiceDouble}),
          collect(rectan), rad, lon, lat)
    rad[], lon[], lat[]
end

"""
    recpgr(body, rectan, re, f)

Convert rectangular coordinates to planetographic coordinates.

### Arguments ###

- `body`: Body with which coordinate system is associated
- `rectan`: Rectangular coordinates of a point
- `re`: Equatorial radius of the reference spheroid
- `f`: flattening coefficient

### Output ###

- `lon`: Planetographic longitude of the point (radians).
- `lat`: Planetographic latitude of the point (radians).
- `alt`: Altitude of the point above reference spheroid.

### References ###

- [NAIF Documentation](https://naif.jpl.nasa.gov/pub/naif/toolkit_docs/C/cspice/recpgr_c.html)
"""
function recpgr(body, rectan, re, f)
    @checkdims 3 rectan
    lon = Ref{SpiceDouble}()
    lat = Ref{SpiceDouble}()
    alt = Ref{SpiceDouble}()
    ccall((:recpgr_c, libcspice), Cvoid,
          (Cstring, Ref{SpiceDouble}, SpiceDouble, SpiceDouble,
           Ref{SpiceDouble}, Ref{SpiceDouble}, Ref{SpiceDouble}),
          body, rectan, re, f, lon, lat, alt)
    handleerror()
    lon[], lat[], alt[]
end

"""
    recrad(rectan)

Convert rectangular coordinates to range, right ascension, and
declination.

### Arguments ###

- `rectan`: Rectangular coordinates of a point

### Output ###

Return the tuple `(range, ra, dec)`.

- `range`: Distance of the point from the origin
- `ra`: Right ascension in radians
- `dec`: Declination in radians

### References ###

- [NAIF Documentation](https://naif.jpl.nasa.gov/pub/naif/toolkit_docs/C/cspice/recrad_c.html)
"""
function recrad(rectan)
    @checkdims 3 rectan
    range = Ref{SpiceDouble}()
    ra = Ref{SpiceDouble}()
    dec = Ref{SpiceDouble}()
    ccall((:recrad_c, libcspice), Cvoid,
          (Ref{SpiceDouble}, Ref{SpiceDouble}, Ref{SpiceDouble}, Ref{SpiceDouble}),
          rectan, range, ra, dec)
    range[], ra[], dec[]
end

"""
    recsph(rectan)

Convert from rectangular coordinates to spherical coordinates.

### Arguments ###

- `rectan`: Rectangular coordinates of a point

### Output ###

- `r     `: Distance of the point from the origin
- `colat `: Angle of the point from the Z-axis in radian
- `lon   `: Longitude of the point in radians

### References ###

- [NAIF Documentation](https://naif.jpl.nasa.gov/pub/naif/toolkit_docs/C/cspice/recsph_c.html)
"""
function recsph(rectan)
    @checkdims 3 rectan
    r = Ref{SpiceDouble}()
    colat = Ref{SpiceDouble}()
    lon = Ref{SpiceDouble}()
    ccall((:recsph_c, libcspice), Cvoid,
          (Ref{SpiceDouble}, Ref{SpiceDouble}, Ref{SpiceDouble}, Ref{SpiceDouble}),
          rectan, r, colat, lon)
    r[], colat[], lon[]
end

"""
    removc!(set, item)

Remove an item from a character set.

### Arguments ###

- `set`: A set
- `item`: Item to be removed

### Output ###

Returns the updated set.

### References ###

- [NAIF Documentation](https://naif.jpl.nasa.gov/pub/naif/toolkit_docs/C/cspice/removc_c.html)
"""
function removc!(set, item)
    data = set.data
    GC.@preserve data begin
        ccall((:removc_c, libcspice), Cvoid, (Cstring, Ref{Cell{SpiceChar}}), item, set.cell)
    end
    set
end

@deprecate removc removc!

"""
    removd!(set, item)

Remove an item from a double set.

### Arguments ###

- `set`: A set
- `item`: Item to be removed

### Output ###

Returns the updated set.

### References ###

- [NAIF Documentation](https://naif.jpl.nasa.gov/pub/naif/toolkit_docs/C/cspice/removd_c.html)
"""
function removd!(set, item)
    data = set.data
    GC.@preserve data begin
        ccall((:removd_c, libcspice), Cvoid, (SpiceDouble, Ref{Cell{SpiceDouble}}), item, set.cell)
    end
    set
end

@deprecate removd removd!

"""
    removi!(set, item)

Remove an item from a character set.

### Arguments ###

- `set`: A set
- `item`: Item to be removed

### Output ###

Returns the updated set.

### References ###

- [NAIF Documentation](https://naif.jpl.nasa.gov/pub/naif/toolkit_docs/C/cspice/removi_c.html)
"""
function removi!(set, item)
    data = set.data
    GC.@preserve data begin
        ccall((:removi_c, libcspice), Cvoid, (SpiceInt, Ref{Cell{SpiceInt}}), item, set.cell)
    end
    set
end

@deprecate removi removi!

function _reordc(iorder, array)
    iorder_ = SpiceInt.(iorder .- 1)
    array_, m, n = chararray(array)
    ccall((:reordc_c, libcspice), Cvoid,
          (Ref{SpiceInt}, SpiceInt, SpiceInt, Ref{SpiceChar}),
          iorder_, m, n, array_)
    chararray_to_string(array_, m)
end

@deprecate reordc(iorder, array) array[iorder]

"""
    reordc(iorder, array)

!!! warning "Deprecated"
    Use `array[iorder]` instead.
"""
reordc

function _reordd(iorder, array)
    iorder_ = SpiceInt.(iorder .- 1)
    n = length(iorder)
    ccall((:reordd_c, libcspice), Cvoid,
          (Ref{SpiceInt}, SpiceInt, Ref{SpiceDouble}),
          iorder_, n, array)
    array
end

@deprecate reordd(iorder, array) array[iorder]

"""
    reordd(iorder, array)

!!! warning "Deprecated"
    Use `array[iorder]` instead.
"""
reordd

function _reordi(iorder, array)
    iorder_ = SpiceInt.(iorder .- 1)
    array_ = SpiceInt.(array)
    n = length(iorder)
    ccall((:reordi_c, libcspice), Cvoid,
          (Ref{SpiceInt}, SpiceInt, Ref{SpiceInt}),
          iorder_, n, array_)
    Int.(array_)
end

@deprecate reordi(iorder, array) array[iorder]

"""
    reordi(iorder, array)

!!! warning "Deprecated"
    Use `array[iorder]` instead.
"""
reordi

function _reordl(iorder, array)
    iorder_ = SpiceInt.(iorder .- 1)
    array_ = SpiceBoolean.(array)
    n = length(iorder)
    ccall((:reordl_c, libcspice), Cvoid,
          (Ref{SpiceInt}, SpiceInt, Ref{SpiceBoolean}),
          iorder_, n, array_)
    Bool.(array_)
end

@deprecate reordl(iorder, array) array[iorder]

"""
    reordl(iorder, array)

!!! warning "Deprecated"
    Use `array[iorder]` instead.
"""
reordl

function _repmc(input, marker, value)
    lenout = length(input) - length(marker) + length(value) + 1
    out = Array{SpiceChar}(undef, lenout)
    ccall((:repmc_c, libcspice), Cvoid,
          (Cstring, Cstring, Cstring, SpiceInt, Ref{SpiceChar}),
          input, marker, value, lenout, out)
    handleerror()
    chararray_to_string(out)
end

@deprecate repmc(input, marker, value) replace(input, marker=>value)

"""
    repmc(input, marker, value)

!!! warning "Deprecated"
    Use `replace(input, marker=>value)` instead.
"""
repmc

@deprecate repmct replace

"""
    repmct

!!! warning "Deprecated"
    Use `replace` instead.
"""
repmct

@deprecate repmd replace

"""
    repmd

!!! warning "Deprecated"
    Use `replace` instead.
"""
repmd

@deprecate repmf replace

"""
    repmf

!!! warning "Deprecated"
    Use `replace` instead.
"""
repmf

@deprecate repmi replace

"""
    repmi

!!! warning "Deprecated"
    Use `replace` instead.
"""
repmi

@deprecate repmot replace

"""
    repmot

!!! warning "Deprecated"
    Use `replace` instead.
"""
repmot

"""
    rotate(angle, iaxis)

Calculate the 3x3 rotation matrix generated by a rotation
of a specified angle about a specified axis. This rotation
is thought of as rotating the coordinate system.

### Arguments ###

- `angle`: Angle of rotation (radians)
- `iaxis`: Axis of rotation (X=1, Y=2, Z=3)

### Output ###

Returns rotation matrix associated with `angle` and `iaxis`.

### References ###

- [NAIF Documentation](https://naif.jpl.nasa.gov/pub/naif/toolkit_docs/C/cspice/rotate_c.html)
"""
function rotate(angle, iaxis)
    r = Matrix{SpiceDouble}(undef, 3, 3)
    ccall((:rotate_c, libcspice), Cvoid, (SpiceDouble, SpiceInt, Ref{SpiceDouble}), angle, iaxis, r)
    permutedims(r)
end

"""
    rotmat(m1, angle, iaxis)

Applies a rotation of `angle` radians about axis `iaxis` to a matrix `m1`. This rotation is thought
of as rotating the coordinate system.

### Arguments ###

- `m1`: Matrix to be rotated
- `angle`: Angle of rotation (radians)
- `iaxis`: Axis of rotation (X=1, Y=2, Z=3)

### Output ###

Returns the resulting rotated matrix.

### References ###

- [NAIF Documentation](https://naif.jpl.nasa.gov/pub/naif/toolkit_docs/C/cspice/rotmat_c.html)
"""
function rotmat(m1, angle, iaxis)
    @checkdims 3 3 m1
    mout = Array{SpiceDouble}(undef, 3, 3)
    ccall((:rotmat_c, libcspice), Cvoid,
          (Ref{SpiceDouble}, SpiceDouble, SpiceInt, Ref{SpiceDouble}),
          permutedims(m1), angle, iaxis, mout)
    permutedims(mout)
end

"""
    rotvec(v1, angle, iaxis)

Transform a vector to a new coordinate system rotated by `angle` radians about axis `iaxis`.
This transformation rotates `v1` by `-angle` radians about the specified axis.

### Arguments ###

- `v1`: Vector whose coordinate system is to be rotated
- `angle`: Angle of rotation in radians
- `iaxis`: Axis of rotation (X=1, Y=2, Z=3)

### Output ###

Returns the resulting vector expressed in the new coordinate system.

### References ###

- [NAIF Documentation](https://naif.jpl.nasa.gov/pub/naif/toolkit_docs/C/cspice/rotvec_c.html)
"""
function rotvec(v1, angle, iaxis)
    @checkdims 3 v1
    vout = Array{SpiceDouble}(undef, 3)
    ccall((:rotvec_c, libcspice), Cvoid,
          (Ref{SpiceDouble}, SpiceDouble, SpiceInt, Ref{SpiceDouble}),
          v1, angle, iaxis, vout)
    vout
end

function _rpd()
    ccall((:rpd_c, libcspice), SpiceDouble, ())
end

@deprecate rpd() deg2rad(1.0)

"""
    rpd()

!!! warning "Deprecated"
    Use `deg2rad` instead.
"""
rpd

"""
    rquad(a, b, c)

Find the roots of a quadratic equation.

### Arguments ###

- `a`: Coefficient of quadratic term
- `b`: Coefficient of linear term
- `c`: Constant

### Output ###

- `root1`: Root built from positive discriminant term
- `root2`: Root built from negative discriminant term

### References ###

- [NAIF Documentation](https://naif.jpl.nasa.gov/pub/naif/toolkit_docs/C/cspice/rquad_c.html)
"""
function rquad(a, b, c)
    root1 = Array{SpiceDouble}(undef, 2)
    root2 = Array{SpiceDouble}(undef, 2)
    ccall((:rquad_c, libcspice), Cvoid,
          (SpiceDouble, SpiceDouble, SpiceDouble, Ref{SpiceDouble}, Ref{SpiceDouble}),
          a, b, c, root1, root2)
    handleerror()
    root1, root2
end


text:
# function Base.cconvert(::Type{Vector{CFacet{T}}}, triangles::Vector{NGonFace{N, Cint}}) where {N, T}
#     polygons = map(1:length(triangles)) do i
#         CPolygon(Ptr{Cint}(pointer(triangles, i)), Cint(N))
#     end
#     facets = map(1:length(polygons)) do i
#         CPolygon(Ptr{CPolygon}(pointer(polygons, i)), Cint(1), Ptr{T}(C_NULL), Cint(0))
#     end
#     return facets
# end

struct JLFacet{T, A <: AbstractVector{Cint}}
    polygons::Vector{A}
    holelist::Vector{T}
end
JLFacet(x::AbstractVector{Cint}) = JLFacet([x])
JLFacet(x::Vector{<: AbstractVector{Cint}}) = JLFacet{Float64, eltype(typeof(x))}(x, Float64[])

function JLPolygon(x::CPolygon)
    JLPolygon(unsafe_wrap(Array, x.vertexlist, x.numberofvertices, own = true))
end

function JLFacet(f::CFacet{T}) where T
    c_polys = unsafe_wrap(Array, f.polygonlist, f.numberofpolygons, own = true)
    holes = unsafe_wrap(Array, f.holelist, f.numberofholes, own = true)
    JLFacet{T}(JLPolygon.(c_polys), holes)
end

function convert_poly(x::Vector{T}) where T <: Union{NgonFace{N, Cint}, NTuple{N, Cint}} where N
    return map(1:length(x)) do i
        CPolygon(pointer(x, i), N)
    end
end

function convert_poly(x::Vector{Vector{Cint}})
    return map(1:length(x)) do i
        CPolygon(Base.unsafe_convert(Ptr{Cint}, x[i]), length(x[i]))
    end
end

struct Region{T}
    pos::Point{3, T}
    attribute::T
    maxvol::T
end

struct FacetConstraint{IT, T}
    facet_marker::IT
    max_area_bound::T
end
struct SegmentationConstraint{IT, T}
    index::IT # index into pointlist
    max_length_bound::T
end

inttype(::Type{Float64}) = Int64
inttype(::Type{Float32}) = Int32

struct JLTetGenIO{T, NSimplex, NAttributes, NMTr, IT, FT}
    points::Vector{Point{3, T}}
    pointattributes::Vector{SVector{NAttributes, T}}
    pointmtrs::Vector{SVector{NMTr, T}}
    pointmarkers::Vector{Cint}

    tetrahedra::Vector{SimplexFace{NSimplex, Cint}}
    tetrahedronattributes::Vector{T}
    tetrahedronvolumes::Vector{T}
    neighbors::Vector{Cint}

    facets::FT
    facetmarkers::Vector{Cint}

    holes::Vector{Point{3, T}}

    regions::Vector{Region{T}}

    facetconstraints::Vector{FacetConstraint{IT, T}}

    segmentconstraints::Vector{SegmentationConstraint{IT, T}}

    trifaces::Vector{TriangleFace{Cint}}
    trifacemarkers::Vector{Cint}

    edges::Vector{LineFace{Cint}}
    edgemarkers::Vector{Cint}

    function JLTetGenIO(
            points::Vector{Point{3, T}},
            pointattributes::Vector{SVector{NAttributes, T}},
            pointmtrs::Vector{SVector{NMTr, T}},
            pointmarkers::Vector{Cint},
            tetrahedrons::Vector{SimplexFace{NSimplex, Cint}},
            tetrahedronattributes::Vector{T},
            tetrahedronvolumes::Vector{T},
            neighbors::Vector{Cint},
            facets::FT,
            facetmarkers::Vector{Cint},
            holes::Vector{Point{3, T}},
            regions::Vector{Region{T}},
            facetconstraints::Vector{FacetConstraint{IT, T}},
            segmentconstraints::Vector{SegmentationConstraint{IT, T}},
            trifaces::Vector{TriangleFace{Cint}},
            trifacemarkers::Vector{Cint},
            edges::Vector{LineFace{Cint}},
            edgemarkers::Vector{Cint},
        ) where {T, NSimplex, NAttributes, NMTr, IT, FT}

        new{T, NSimplex, NAttributes, NMTr, IT, FT}(
            points,
            pointattributes,
            pointmtrs,
            pointmarkers,
            tetrahedrons,
            tetrahedronattributes,
            tetrahedronvolumes,
            neighbors,
            facets,
            facetmarkers,
            holes,
            regions,
            facetconstraints,
            segmentconstraints,
            trifaces,
            trifacemarkers,
            edges,
            edgemarkers,
        )
    end

end

function JLTetGenIO(
        points::Vector{Point{3, T}};
        pointattributes = SVector{0, T}[],
        pointmtrs = SVector{0, T}[],
        pointmarkers = Cint[],
        tetrahedrons = SimplexFace{4, Cint}[],
        tetrahedronattributes = T[],
        tetrahedronvolumes = T[],
        neighbors = Cint[],
        facets = JLFacet{T, Vector{Cint}}[],
        facetmarkers = Cint[],
        holes = Point{3, T}[],
        regions = Region{T}[],
        facetconstraints = FacetConstraint{inttype(T), T}[],
        segmentconstraints = SegmentationConstraint{inttype(T), T}[],
        trifaces = TriangleFace{Cint}[],
        trifacemarkers = Cint[],
        edges = LineFace{Cint}[],
        edgemarkers = Cint[]

    ) where T
    JLTetGenIO(
        points,
        pointattributes,
        pointmtrs,
        pointmarkers,
        tetrahedrons,
        tetrahedronattributes,
        tetrahedronvolumes,
        neighbors,
        facets,
        facetmarkers,
        holes,
        regions,
        facetconstraints,
        segmentconstraints,
        trifaces,
        trifacemarkers,
        edges,
        edgemarkers,
    )

end

function number_of_elements_field(io, basename)
    nfield = if basename == "tetrahedra"
        Symbol("numberof" * basename)
    else
        Symbol("numberof" * basename * "s")
    end
    n = if nfield in fieldnames(typeof(io))
        return nfield
    elseif startswith(basename, "point")
        return :numberofpoints
    elseif startswith(basename, "facet")
        return :numberoffacets
    elseif startswith(basename, "triface")
        return :numberoftrifaces
    elseif startswith(basename, "edge")
        return :numberofedges
    else
        error("No numberof for $nfield")
    end
end

function get_array(io, field::Symbol, Typ)
    basename = if field == :tetrahedra
        string(field)
    else
        string(field)[1:end-1] # cut of s
    end
    ptr = if field == :tetrahedra
        io.tetrahedronlist
    else
        getfield(io, Symbol(basename * "list"))
    end
    ptr == C_NULL && return Typ[]
    n = getfield(io, number_of_elements_field(io, basename))
    # own true, since we disable freeing in the c-part
    return unsafe_wrap(Array, Base.unsafe_convert(Ptr{Typ}, ptr), n, own = true)
end

function Base.convert(
        IOT::Type{JLTetGenIO{T, NSimplex, NAttributes, NMTr, IT, A}}, io::CPPTetGenIO{T}
    ) where {T, NSimplex, NAttributes, NMTr, IT, A}
    JLTetGenIO(ntuple(fieldcount(IOT)) do i
        fname, ftype = fieldname(IOT, i), fieldtype(IOT, i)
        get_array(io, fname, eltype(ftype))
    end...)
end

function Base.convert(::Type{JLTetGenIO}, io::CPPTetGenIO{T}) where T
    NSimplex = Int(io.numberofcorners)
    NAttributes = Int(io.numberofpointattributes)
    NMTr = Int(io.numberofpointmtrs)
    IT = inttype(T) # TODO these are indices,
    convert(JLTetGenIO{T, NSimplex, NAttributes, NMTr, IT, Vector{CFacet{T}}}, io)
end

function Base.cconvert(::Type{Ptr{CFacet{T}}}, facets::Vector{JLFacet{T, A}}) where {T, A}
    c_polys = map(facets) do facet
        convert_poly(facet.polygons)
    end
    c_facets = map(c_polys) do cpoly
        ptr = Base.unsafe_convert(Ptr{CPolygon}, cpoly)
        CFacet{T}(ptr, length(cpoly), C_NULL, 0)
    end
    return (facets, c_polys, c_facets)
end

function Base.cconvert(::Type{Ptr{CFacet{T}}}, facets::Vector{NgonFace{N, Cint}}) where {N, T}
    nfacets = length(facets)
    c_polys = map(1:nfacets) do i
        CPolygon(Ptr{Cint}(pointer(facets, i)), N)
    end
    c_facets = map(1:nfacets) do i
        ptr = pointer(c_polys, i)
        CFacet{T}(ptr, 1, C_NULL, 0)
    end
    return (facets, c_polys, c_facets)
end

function unsafe_array_convert(
        ::Type{Ptr{CFacet{T}}},
        facets::Tuple{Vector{F}, Vector{Vector{CPolygon}}, Vector{CFacet{T}}}
    ) where {F, T}
    return Base.unsafe_convert(Ptr{CFacet{T}}, facets[3])
end

function unsafe_array_convert(
        ::Type{Ptr{CFacet{T}}},
        facets::Tuple{Vector{F}, Vector{CPolygon}, Vector{CFacet{T}}}
    ) where {F, T}
    return Base.unsafe_convert(Ptr{CFacet{T}}, facets[3])
end

function unsafe_array_convert(P::Type{Ptr{T}}, x::Vector{T}) where T
    Base.unsafe_convert(P, x)
end

function unsafe_array_convert(P::Type{Ptr{T1}}, x::Vector{T2}) where {T1, T2}
    Ptr{T1}(Base.unsafe_convert(Ptr{T2}, x))
end

function Base.cconvert(CIO::Type{CPPTetGenIO{T}}, obj::JLTetGenIO{T, NSimplex, NAttributes, NMTr, IT, A}) where {T, NSimplex, NAttributes, NMTr, IT, A}
    cfnames = fieldnames(CIO)
    dict = Dict{Symbol, Any}(
        :numberofpointattributes => NAttributes,
        :numberofpointmtrs => NMTr,
        :numberofcorners => NSimplex,
        :firstnumber => Cint(1),
        :mesh_dim => Cint(3),
    )
    gc_tack_cconvert = []
    for field in fieldnames(typeof(obj))
        basename = if field == :tetrahedra
            string(field)
        else
            string(field)[1:end-1] # cut of s
        end
        listname = Symbol(replace(basename * "list", "tetrahedra" => "tetrahedron"))
        FT = fieldtype(CIO, listname)
        array = getfield(obj, field)
        converted = Base.cconvert(FT, array)
        push!(gc_tack_cconvert, converted)
        if isempty(array)
            # make sure we forward NULL for empty
            dict[listname] = Base.unsafe_convert(FT, C_NULL)
        else
            dict[listname] = unsafe_array_convert(FT, converted)
        end
        nfield = if basename == "tetrahedra"
            Symbol("numberof" * basename)
        else
            Symbol("numberof" * basename * "s")
        end
        if nfield in cfnames
            dict[nfield] = length(array)
        end
    end
    fields_tuple = getindex.((dict,), cfnames)
    return (CPPTetGenIO{T}(fields_tuple...), gc_tack_cconvert)
end

function Base.unsafe_convert(::Type{CPPTetGenIO{T}}, x::Tuple{CPPTetGenIO{T}, Vector{Any}}) where T
    return x[1]
end

function tetrahedralize(input::JLTetGenIO{Float64}, command::String)
    rc=Cint[0]
    cres = ccall((:tetrahedralize2_f64, libtet), CPPTetGenIO{Float64}, (CPPTetGenIO{Float64}, Cstring, Ptr{Cint}), input, command,rc)
    if rc[1]!=0
        throw(TetGenError(rc[1]))
    end
    return convert(JLTetGenIO, cres)
end


text:
function save_noise!(W::BoxWedgeTail)
  if W.t != W.curt
    push!(W.W,copy(W.curW))
    push!(W.t,copy(W.curt))
    push!(W.A,copy(W.curA))
    if W.Z != nothing
      push!(W.Z,copy(W.curZ))
    end
  end
  return nothing
end

@inline function accept_step!(W::BoxWedgeTail,dt,u,p,setup_next=true)

  W.curt += W.dt
  W.iter += 1

  if isinplace(W)
    @.. W.curW += W.dW
    if W.Z != nothing
      @.. W.curZ += W.dZ
    end
  else
    W.curW += W.dW
    if W.Z != nothing
      W.curZ += W.dZ
    end
  end

  if W.save_everystep
    push!(W.W,copy(W.curW))
    push!(W.t,copy(W.curt))
    push!(W.A,copy(W.curA))
    if W.Z != nothing
      push!(W.Z,copy(W.curZ))
    end
  end

  W.dt = dt #dtpropose
  # Setup next step
  if setup_next
    setup_next_step!(W::BoxWedgeTail,u,p)
  end
  return nothing
end

@inline function setup_next_step!(W::BoxWedgeTail,u,p)
  calculate_step!(W,W.dt,u,p)
  return nothing
end

@inline function calculate_step!(W::BoxWedgeTail,dt,u,p)
  if isinplace(W)
    sample_distribution(W.dW,W,dt,u,p,W.curt,W.rng)
    if W.Z != nothing
      W.dist(W.dZ,W,dt,u,p,W.curt,W.rng)
    end
  else
    W.dW, W.curA = sample_distribution(W,dt,u,p,W.curt,W.rng)
    if W.Z != nothing
      W.dZ = W.dist(W,dt,u,p,W.curt,W.rng)
    end
  end
  W.dt = dt
  return nothing
end

@inline function reject_step!(W::BoxWedgeTail,dtnew,u,p)
  error("BoxWedgeTail cannot be used with adaptivity rejections")
end

@inline function interpolate!(W::BoxWedgeTail,u,p,t)
  if sign(W.dt)*t > sign(W.dt)*W.t[end] # Steps past W
    dt = t - W.t[end]
    if isinplace(W)
      sample_distribution(W.dW,W,dt,u,p,t,W.rng)
      W.curW .+= W.dW
      if W.Z != nothing
        W.dist(W.dZ,W,dt,u,p,t,W.rng)
        W.curZ .+= W.dZ
      end
    else
      W.dW, W.curA = sample_distribution(W,dt,u,p,t,W.rng)
      W.curW += W.dW
      if W.Z != nothing
        W.dZ = W.dist(W,dt,u,p,t,W.rng)
        W.curZ += W.dZ
      end
    end
    out1 = copy(W.curW)
    if W.save_everystep
      push!(W.t,t)
      push!(W.W,out1)
    end
    if W.Z != nothing
      out2= copy(W.curZ)
      if W.save_everystep
        push!(W.Z,out2)
      end
    else
      out2 = nothing
    end
    return out1,out2
  else # Bridge
    i = searchsortedfirst(W.t,t)
    if t == W.t[i]
      if isinplace(W)
        W.curW .= W.W[i]
      else
        W.curW = W.W[i]
      end
      if W.Z != nothing
        if isinplace(W)
          W.curZ .= W.Z[i]
        else
          W.curZ = W.Z[i]
        end
        return copy(W.curW),copy(W.curZ)
      else
        return copy(W.curW),nothing
      end
    else
      W0,Wh = W.W[i-1],W.W[i]
      if W.Z != nothing
        Z0,Zh = W.Z[i-1],W.Z[i]
      end
      h = W.t[i]-W.t[i-1]
      q = (t-W.t[i-1])/h
      if isinplace(W)
        new_curW = similar(W.dW)
        W.bridge(new_curW,W,W0,Wh,q,h,u,p,t,W.rng)
        if iscontinuous(W)
          @. new_curW += (1-q)*W0
        else
          @. new_curW += W0
        end
        if W.Z != nothing
          new_curZ = similar(W.dZ)
          W.bridge(new_curZ,W,Z0,Zh,q,h,u,p,t,W.rng)
          if iscontinuous(W)
            @. new_curZ += (1-q)*Z0
          else
            @. new_curZ += Z0
          end
        else
          new_curZ = nothing
        end
      else
        new_curW = W.bridge(W,W0,Wh,q,h,u,p,t,W.rng)
        if iscontinuous(W)
          # This should actually be based on the function for computing the mean
          # flow of the noise process, but for now we'll just handle Wiener and
          # Poisson
          new_curW += (1-q)*W0
        else
          new_curW += W0
        end
        if W.Z != nothing
          new_curZ = W.bridge(W,Z0,Zh,q,h,u,p,t,W.rng)
          if iscontinuous(W)
            new_curZ += (1-q)*Z0
          else
            new_curZ += Z0
          end
        else
          new_curZ = nothing
        end
      end
      W.curW = new_curW
      if W.save_everystep
        insert!(W.W,i,new_curW)
        insert!(W.t,i,t)
      end
      if W.Z != nothing
        W.curZ = new_curZ
        if W.save_everystep
          insert!(W.Z,i,new_curZ)
        end
      end
      return new_curW,new_curZ
    end
  end
end

@inline function interpolate!(out1,out2,W::BoxWedgeTail,u,p,t)
  if sign(W.dt)*t > sign(W.dt)*W.t[end] # Steps past W
    dt = t - W.t[end]
    sample_distribution(W.dW,W,dt,u,p,t,W.rng)
    out1 .+= W.dW
    if W.Z != nothing
      W.dist(W.dZ,W,dt,u,p,t,W.rng)
      out2 .+= W.dZ
    end
    if W.save_everystep
      push!(W.t,t)
      push!(W.W,copy(out1))
      if W.Z != nothing
        push!(W.Z,copy(out2))
      end
    end
  else # Bridge
    i = searchsortedfirst(W.t,t)
    if t == W.t[i]
      out1 .= W.W[i]
      if W.Z != nothing
        out2 .= W.Z[i]
      end
    else
      W0,Wh = W.W[i-1],W.W[i]
      if W.Z != nothing
        Z0,Zh = W.Z[i-1],W.Z[i]
      end
      h = W.t[i]-W.t[i-1]
      q = (t-W.t[i-1])/h
      W.bridge(out1,W,W0,Wh,q,h,u,p,t,W.rng)
      out1 .+= (1-q)*W0
      if W.Z != nothing
        W.bridge(out2,W,Z0,Zh,q,h,u,p,t,W.rng)
        out2 .+= (1-q)*Z0
      end
      W.curW .= out1
      if W.save_everystep
        insert!(W.W,i,copy(out1))
        insert!(W.t,i,t)
      end
      if W.Z != nothing
        W.curZ .= out2
        if W.save_everystep
          insert!(W.Z,i,copy(out2))
        end
      end
    end
  end
end

# joint probability densitity function
function joint_density_function(r,a,rtol)
  integral, err = QuadGK.quadgk(x ->r/π*(x/sinh(x))*exp(-r^2*x/(2*tanh(x)))*cos(a*x), 0, Inf, rtol=rtol)
  return integral
end

# boxes
abstract type AbstractBoxGeneration end

mutable struct BoxGeneration1{boxesType, pType, offType, distType} <: AbstractBoxGeneration
  boxes::boxesType
  probability::pType
  offset::offType
  dist::distType
end

struct BoxGeneration2{boxesType, pType, offType, distType} <: AbstractBoxGeneration
  boxes::boxesType
  probability::pType
  offset::offType
  dist::distType
end

struct BoxGeneration3{boxesType, pType, offType, distType} <: AbstractBoxGeneration
  boxes::boxesType
  probability::pType
  offset::offType
  dist::distType
end

function generate_boxes1(densf, Δr, Δa, Δz, rM, aM, offset=nothing, scale=1)
  boxes = Array{typeof(Δr),1}[]
  probability = Vector{typeof(Δr)}(undef, 0)
  rs = collect(zero(rM):Δr:rM)
  as = collect(zero(aM):Δa:aM)

  if offset===nothing
    offset = zeros(typeof(Δr), Int(rM/Δr), Int(aM/Δa))
  end

  for (i,r) in enumerate(rs[1:end-1])
    for (j,a) in enumerate(as[1:end-1])
      # height of the box
      z = Δz
      if offset !== nothing
        indx1 = (i-1)*scale+one(scale)
        indx2 = (j-1)*scale+one(scale)
        z += offset[indx1,indx2]
      end

      counter = 0

      # check for complete inclusion of box in volume
      while (z <= densf(r,a) && z <= densf(r+Δr,a) && z <= densf(r,a+Δa) && z <= densf(r+Δr,a+Δa))
        counter += one(counter)
        z += Δz
      end
      z -= Δz
      # store position of top corner of box and width
      if counter != 0
        push!(boxes, [r, a, Δr, Δa])
        push!(probability, counter*2*Δr*Δa*Δz)
      end

      if offset !== nothing
        offset[indx1:indx1+(scale-one(scale)),indx2:indx2+(scale-one(scale))] .= z
      end

    end
  end
  return boxes, probability, offset
end

function generate_boxes2(densf, Δrmin, Δamin, Δzmin, Δrmax, Δamax, Δzmax, rM, aM, offset=nothing)
  boxes = Array{typeof(Δrmin),1}[]
  probability = Vector{typeof(Δrmin)}(undef, 0)
  # start with largest possible size, then subsequently decrease size and fill remaining space
  Δr = Δrmax
  Δa = Δamax
  Δz = Δzmax

  if offset===nothing
    offset = zeros(typeof(Δrmin), Int(rM/Δrmin), Int(aM/Δamin))
  end

  while (Δr >= Δrmin && Δa >= Δamin && Δz >= Δzmin)
    scale = Int(Δr/Δrmin)
    boxestmp, probabilitytmp = generate_boxes1(densf, Δr, Δa, Δz, rM, aM, offset, scale)
    boxes = cat(boxes, boxestmp, dims=1)
    probability = cat(probability, probabilitytmp, dims=1)
    Δr = Δr/2
    Δa = Δa/2
    Δz = Δz/2
  end
  # merge boxes
  return boxes, probability, offset
end

function generate_boxes3(densf, Δr, Δa, Δz, rM, aM, offset=nothing, scale=1)
  boxes = Array{typeof(Δr),1}[]
  probability = Vector{typeof(Δr)}(undef, 0)
  rs = collect(zero(rM):Δr:rM)
  as = collect(zero(aM):Δa:aM)

  if offset===nothing
    offset = zeros(typeof(Δr), Int(rM/Δr), Int(aM/Δa))
  end

  for (i,r) in enumerate(rs[1:end-1])
    for (j,a) in enumerate(as[1:end-1])
      # height of the box
      z = Δz
      if offset !== nothing
        indx1 = (i-1)*scale+one(scale)
        indx2 = (j-1)*scale+one(scale)
        z += offset[indx1,indx2]
      end
      # check for complete inclusion of box in volume
      while (z <= densf(r,a) && z <= densf(r+Δr,a) && z <= densf(r,a+Δa) && z <= densf(r+Δr,a+Δa))
        # store position of top corner of box and width
        push!(boxes, [r, a, Δr, Δa])
        # probability = 2*Δr*Δa*Δz for all small boxes, push!(probability, 1)
        push!(probability, 2*Δr*Δa*Δz)

        z += Δz
      end
      z -= Δz
      if offset !== nothing
        offset[indx1:indx1+(scale-one(scale)),indx2:indx2+(scale-one(scale))] .= z
      end
    end
  end
  return boxes, probability, offset
end

# entropy of partition
function entropy(probs)
  TotalVolume = sum(probs)
  qi = probs./TotalVolume
  return - sum(qi.*log2.(qi))
end

# sampling from boxes
function sample_box(W::BoxWedgeTail, Boxes::AbstractBoxGeneration)
  indx = rand(W.rng, Boxes.dist)
  # boxes store r, a, Δr, Δa
  ri, ai, Δr, Δa = Boxes.boxes[indx]
  DU = Distributions.Product(Distributions.Uniform.([ri,ai],[ri+Δr,ai+Δa]))

  r, a = rand(W.rng, DU)
  return r,a
end

# Wedges
struct Wedges{boxesType, pType, offType, distType} <: AbstractBoxGeneration
  boxes::boxesType
  probability::pType
  fvalues::offType
  dist::distType
end

# linear approximation for squeezing method
function linear_interpolation_wedges(fij, fij2, fij3, fij4, r, a, ri, ai, Δr, Δa)
  t = (r - ri)/Δr
  u = (a - ai)/Δa
  return (one(t)-t)*(one(u)-u)*fij + t*(one(u)-u)*fij2 + (one(t)-t)*u*fij3 + t*u*fij4
end

function contrained_optimization_problem(densf, fij, fij2, fij3, fij4, ri, ai, Δr, Δa)
  function difference(x)
    densf(x[1], x[2]) - linear_interpolation_wedges(fij, fij2, fij3, fij4, x[1], x[2], ri, ai, Δr, Δa)
  end
  ϵijmax = Optim.optimize(x->-difference(x), [ri, ai], [ri+Δr, ai+Δa], [ri+Δr/2, ai+Δa/2], Optim.Fminbox(Optim.NelderMead()))
  ϵijmin = Optim.optimize(x-> difference(x), [ri, ai], [ri+Δr, ai+Δa], [ri+Δr/2, ai+Δa/2], Optim.Fminbox(Optim.NelderMead()))
  return ϵijmin.minimum, ϵijmax.minimum
end

function generate_wedges(densf, Δr, Δa, Δz, rM, aM, offset, sqeezing)
  boxes = Array{typeof(Δr),1}[]
  probability = Vector{typeof(Δr)}(undef, 0)
  f = Array{typeof(Δr),1}[]
  rs = collect(zero(rM):Δr:rM)
  as = collect(zero(aM):Δa:aM)
  for (i,r) in enumerate(rs[1:end-1])
    for (j,a) in enumerate(as[1:end-1])
      # base height hij of the box enclosing the wedge
      hij = offset[i,j]
      # max height of the box f̃_ij
      fij = densf(r,a) # f_{ij}
      fij2 = densf(r+Δr,a) # f_{i+1,j}
      fij3 = densf(r,a+Δa) # f_{i,j+1}
      fij4 = densf(r+Δr,a+Δa) # f_{i+1,j+1}

      f̃ij = max(fij,fij2,fij3,fij4)

      # store position of top corner of box and width
      if sqeezing
        ϵijmin, ϵijmax = contrained_optimization_problem(densf, fij, fij2, fij3, fij4, r, a, Δr, Δa)
        push!(boxes, [f̃ij, hij, abs(ϵijmin), abs(ϵijmax), r, a])
        # store PDF values
        push!(f, [fij, fij2, fij3, fij4])
      else
        push!(boxes, [f̃ij, hij, zero(hij), zero(hij), r, a])
      end

      # probability = 2*Δr*Δa*Δz
      push!(probability, 2*Δr*Δa*(f̃ij-hij))
    end
  end
  return boxes, probability, f
end

# sampling from wedges
function sample_wedge(W::BoxWedgeTail, wedges::Wedges)
  indx = rand(W.rng, wedges.dist)
  # wedges store f̃ij, hij, ϵijmin, ϵijmax, r, a, Δr
  f̃ij, hij, ϵijmin, ϵijmax, ri, ai = wedges.boxes[indx]
  DU = Distributions.Product(Distributions.Uniform.([ri,ai,hij],[ri+W.Δr,ai+W.Δa,f̃ij]))
  if W.sqeezing
    fij, fij2, fij3, fij4 = wedges.fvalues[indx]
    while true
      r, a, z = rand(W.rng, DU)
      flin = linear_interpolation_wedges(fij, fij2, fij3, fij4, r, a, ri, ai, W.Δr, W.Δa)
      if z > flin + ϵijmax
        continue
      elseif z < flin - ϵijmin
        break
      end

      if z <= W.jpdf(r,a)
        break
      end
    end
  else
    r, a, z = rand(W.rng, DU)
    while z > W.jpdf(r,a)
      r, a, z = rand(W.rng, DU)
    end
  end
  return r, a
end

# Tail approximation
abstract type AbstractTail end

struct Tail1{pType, distType, pdfType, cType} <: AbstractTail
  p::pType
  dist::distType
  candpdf::pdfType
  cvalue::cType # success probability 1/c.

  function Tail1(rM, aM)
    p = convert(typeof(rM),0.0002982405821953734)
    candpdf = (r,a) -> exp(-r^2/2)
    dist1 = Distributions.truncated(Distributions.Normal(zero(rM), one(rM)), rM, 12*one(rM))
    dist2 = Distributions.Uniform(zero(aM), aM)

    dist = Distributions.Product([dist1,dist2])

    c = convert(typeof(rM),2.3)

    new{typeof(p),typeof(dist),typeof(candpdf),typeof(c)}(p,dist,candpdf,c)
  end
end

struct Tail2{pType, distType, pdfType, cType, boundsType} <: AbstractTail
  p::pType
  dist::distType
  candpdf::pdfType
  cvalue::cType # success probability 1/c.
  alower::boundsType
  aupper::boundsType

  function Tail2(rM, aM)
    p = convert(typeof(rM),3.648002197730662e-5)
    candpdf = (r,a) -> 3//2*exp(-r^2/2-3*a^2/(2*r^2))
    # store only distribution of r, when sampling re-create distribution for a|r
    dist = Distributions.truncated(Distributions.Normal(zero(rM), one(rM)), rM, 8*one(rM))

    c = convert(typeof(rM),2.0)

    alower = aM
    aupper = 2*aM

    new{typeof(p),typeof(dist),typeof(candpdf),typeof(c),typeof(alower)}(p,dist,candpdf,c,alower,aupper)
  end
end

struct Tail3{pType, distType, pdfType, cType} <: AbstractTail
  p::pType
  dist::distType
  candpdf::pdfType
  cvalue::cType # success probability 1/c.

  function Tail3(rM, aM)
    p = convert(typeof(rM),0.0018026889941638695)
    candpdf = (r,a) -> exp(convert(typeof(rM),-pi/2)*a)
    dist1 = Distributions.Uniform(2*one(rM), rM)
    dist2 = Distributions.truncated(Distributions.Exponential(convert(typeof(rM),2/pi)), aM, 8*one(aM))

    dist = Distributions.Product([dist1,dist2])

    c = convert(typeof(rM),2.6)

    new{typeof(p),typeof(dist),typeof(candpdf),typeof(c)}(p,dist,candpdf,c)
  end
end

struct Tail4{pType, distType, pdfType, cType} <: AbstractTail
  p::pType
  dist::distType
  candpdf::pdfType
  cvalue::cType # success probability 1/c.

  function Tail4(rM, aM)
    p = convert(typeof(rM),1.6145323459108908e-6)
    candpdf = (r,a) -> 15*r*exp(convert(typeof(rM),-pi)*a)
    dist1 = Distributions.TriangularDist(zero(rM), one(rM)/2, one(rM)/2)
    dist2 = Distributions.truncated(Distributions.Exponential(convert(typeof(rM),1/pi)), aM, 6*one(aM))

    dist = Distributions.Product([dist1,dist2])

    c = convert(typeof(rM),2.6)

    new{typeof(p),typeof(dist),typeof(candpdf),typeof(c)}(p,dist,candpdf,c)
  end
end

struct Tail5{pType, distType, pdfType, cType} <: AbstractTail
  p::pType
  dist::distType
  candpdf::pdfType
  cvalue::cType # success probability 1/c.

  function Tail5(rM, aM)
    p = convert(typeof(rM),1.9637270929071978e-5)
    candpdf = (r,a) -> 15*r*exp(convert(typeof(rM),-2.8)*a)
    dist1 = Distributions.TriangularDist(one(rM)/2, one(rM), one(rM))
    dist2 = Distributions.truncated(Distributions.Exponential(convert(typeof(rM),1/2.8)), aM, 6*one(aM))

    dist = Distributions.Product([dist1,dist2])

    c = convert(typeof(rM),2.8)

    new{typeof(p),typeof(dist),typeof(candpdf),typeof(c)}(p,dist,candpdf,c)
  end
end

struct Tail6{pType, distType, pdfType, cType} <: AbstractTail
  p::pType
  dist::distType
  candpdf::pdfType
  cvalue::cType # success probability 1/c.

  function Tail6(rM, aM)
    p = convert(typeof(rM),0.00012022274714710923)
    candpdf = (r,a) -> 25*r*exp(convert(typeof(rM),-2.6)*a)
    dist1 = Distributions.TriangularDist(one(rM), 3/2*one(rM), 3/2*one(rM))
    dist2 = Distributions.truncated(Distributions.Exponential(convert(typeof(rM),1/2.6)), aM, 6*one(aM))

    dist = Distributions.Product([dist1,dist2])

    c = convert(typeof(rM),3.0)

    new{typeof(p),typeof(dist),typeof(candpdf),typeof(c)}(p,dist,candpdf,c)
  end
end

struct Tail7{pType, distType, pdfType, cType} <: AbstractTail
  p::pType
  dist::distType
  candpdf::pdfType
  cvalue::cType # success probability 1/c.

  function Tail7(rM, aM)
    p = convert(typeof(rM),0.00038595770876898326)
    candpdf = (r,a) -> 25*r*exp(convert(typeof(rM),-2.4)*a)
    dist1 = Distributions.TriangularDist(3/2*one(rM), 2*one(rM), 2*one(rM))
    dist2 = Distributions.truncated(Distributions.Exponential(convert(typeof(rM),1/2.4)), aM, 6*one(aM))

    dist = Distributions.Product([dist1,dist2])

    c = convert(typeof(rM),3.2)

    new{typeof(p),typeof(dist),typeof(candpdf),typeof(c)}(p,dist,candpdf,c)
  end
end

struct Tail8{pType, distType, pdfType, cType} <: AbstractTail
  p::pType
  dist::distType
  candpdf::pdfType
  cvalue::cType # success probability 1/c.

  function Tail8(rM, aM)
    p = convert(typeof(rM),6.566973361922312e-6)
    candpdf = (r,a) -> 40*r*exp(convert(typeof(rM),-2.4)*a)
    dist1 = Distributions.TriangularDist(one(rM), 2*one(rM), 2*one(rM))
    dist2 = Distributions.truncated(Distributions.Exponential(convert(typeof(rM),1/2.4)), 6*one(aM), 8*one(aM))

    dist = Distributions.Product([dist1,dist2])

    c = convert(typeof(rM),4.2)

    new{typeof(p),typeof(dist),typeof(candpdf),typeof(c)}(p,dist,candpdf,c)
  end
end

struct Tail9{pType, distType, pdfType, cType} <: AbstractTail
  p::pType
  dist::distType
  candpdf::pdfType
  cvalue::cType # success probability 1/c.

  function Tail9(rM, aM)
    p = convert(typeof(rM),4.149989421681627e-6)
    candpdf = (r,a) -> 7//10*exp(convert(typeof(rM),-pi/2)*a)
    dist1 = Distributions.Uniform(2*one(rM), 5*one(rM))
    dist2 = Distributions.truncated(Distributions.Exponential(convert(typeof(rM),2/pi)), 8*one(aM), 10*one(aM))

    dist = Distributions.Product([dist1,dist2])

    c = convert(typeof(rM),2.4)

    new{typeof(p),typeof(dist),typeof(candpdf),typeof(c)}(p,dist,candpdf,c)
  end
end

struct TailApproxs{T1<:Tail1,T2<:Tail2,T3<:Tail3,T4<:Tail4,T5<:Tail5,T6<:Tail6,
    T7<:Tail7,T8<:Tail8,T9<:Tail9,distType}
  tail1::T1
  tail2::T2
  tail3::T3
  tail4::T4
  tail5::T5
  tail6::T6
  tail7::T7
  tail8::T8
  tail9::T9
  dist::distType

  function TailApproxs(rM, aM)
    tail1 = Tail1(rM, aM)
    tail2 = Tail2(rM, aM)
    tail3 = Tail3(rM, aM)
    tail4 = Tail4(rM, aM)
    tail5 = Tail5(rM, aM)
    tail6 = Tail6(rM, aM)
    tail7 = Tail7(rM, aM)
    tail8 = Tail8(rM, aM)
    tail9 = Tail9(rM, aM)

    probability = [tail1.p,tail2.p,tail3.p,tail4.p,tail5.p,tail6.p,tail7.p,tail8.p,tail9.p]

    dist = Distributions.Categorical(probability/sum(probability))

    new{typeof(tail1),typeof(tail2),typeof(tail3),typeof(tail4),typeof(tail5),
      typeof(tail6),typeof(tail7),typeof(tail8),typeof(tail9),typeof(dist)}(tail1,tail2,
      tail3,tail4,tail5,tail6,tail7,tail8,tail9,dist)
  end
end

# sampling from tail approximations by rejection sampling
function sample_tail(W::BoxWedgeTail, tails::TailApproxs)
  indx = rand(W.rng, tails.dist)
  # decide from which tail region to sample
  if indx == 1
    T = tails.tail1
  elseif indx == 2
    T = tails.tail2
  elseif indx == 3
    T = tails.tail3
  elseif indx == 4
    T = tails.tail4
  elseif indx == 5
    T = tails.tail5
  elseif indx == 6
    T = tails.tail6
  elseif indx == 7
    T = tails.tail7
  elseif indx == 8
    T = tails.tail8
  else
    T = tails.tail9
  end

  r, a = sample_tail(W.rng, W.jpdf, T)
end

function sample_tail(rng, densf, T::Tail2)

  local r,a

  while true
    # draw uniform
    U = rand(rng)
    # draw a candidate X∼g from the candidate density
    # draw first the r value
    r = rand(rng, T.dist)
    # generate a from the conditional distribution
    dista = Distributions.truncated(Distributions.Normal(zero(r), sqrt(r/3)), T.alower, T.aupper)
    a = rand(rng, dista)
    if U <= densf(r,a)/(T.cvalue*T.candpdf(r,a))
      break
    end
  end

  return r,a
end

function sample_tail(rng, densf, T)
  local r, a
  while true
    # draw uniform
    U = rand(rng)
    # draw a candidate X∼g from the candidate density
    r,a = rand(rng, T.dist)
    if U <= densf(r,a)/(T.cvalue*T.candpdf(r,a))
      break
    end
  end

  return r,a
end

function sample_distribution(dW,W::BoxWedgeTail,dt,u,p,t,rng)
  # decide if sample should come from boxes, wedges, or tails
  indx = rand(rng, W.distBWT)

  # draw (r,a) pair
  if indx == 1
    r, a = DiffEqNoiseProcess.sample_box(W, W.boxes)
  elseif indx == 2
    r, a = DiffEqNoiseProcess.sample_wedge(W, W.wedges)
  else
    r, a = DiffEqNoiseProcess.sample_tail(W, W.tails)
  end

  # get dWs from r and scale by sqrt(dt)
  θ = rand(rng, W.distΠ)
  sqabsdt = @fastmath sqrt(abs(dt))
  @. dW = sqabsdt*r*[cos(θ), sin(θ)]

  # random sign for A and scale by dt
  W.curA = rand(rng, (-1, 1))*a*dt

  return nothing
end

function sample_distribution(W::BoxWedgeTail,dt,u,p,t,rng)
  # decide if sample should come from boxes, wedges, or tails
  indx = rand(rng, W.distBWT)

  # draw (r,a) pair
  if indx == 1
    r, a = DiffEqNoiseProcess.sample_box(W, W.boxes)
  elseif indx == 2
    r, a = DiffEqNoiseProcess.sample_wedge(W, W.wedges)
  else
    r, a = DiffEqNoiseProcess.sample_tail(W, W.tails)
  end

  # get dWs from r and scale by sqrt(dt)
  θ = rand(rng, W.distΠ)
  sqabsdt = @fastmath sqrt(abs(dt))
  dW = sqabsdt*r*[cos(θ), sin(θ)]

  # random sign for A and scale by dt
  A = rand(rng, (-1, 1))*a*dt

  return dW, A
end


text:
module Unification

using InteractiveUtils
using Logging

include("bindings.jl")
include("unify.jl")

end


text:
# Autogenerated wrapper script for Libgpg_error_jll for i686-w64-mingw32
export libgpg_error

JLLWrappers.@generate_wrapper_header("Libgpg_error")
JLLWrappers.@declare_library_product(libgpg_error, "libgpg-error-0.dll")
function __init__()
    JLLWrappers.@generate_init_header()
    JLLWrappers.@init_library_product(
        libgpg_error,
        "bin\\libgpg-error-0.dll",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()


text:
export i_ascending!

"""
	i_ascending!(xs!, inds!, arr)

Find the ascending sequence in `arr` and store the results into `xs!`, indices are stored in `inds!`.
This function can be used to get the maximum value and maximum indices.
"""
@i function i_ascending!(xs!::AbstractVector{T}, inds!, arr::AbstractArray{T}) where T
	@invcheckoff if (length(arr) > 0, ~)
		y ← zero(T)
		y += arr[1]
		xs![end+1] ↔ y
		anc ← 1
		inds![end+1] ↔ anc
		@inbounds for i = 2:length(arr)
			if (arr[i] > xs![end], i==inds![end])
				ind ← i
				x ← zero(T)
				x += arr[i]
				xs![end+1] ↔ x
				inds![end+1] ↔ ind
			end
		end
	end
end


text:
"""
$(TYPEDEF)

Users can create subtypes of `Controller` to implement controllers.
If such a controller object is passed to `simulate!`, at each time step, `controller.control!(mechanism, controller, timestep)` will be called.
This way, the user has access to all attributes of the `mechanism`, the `controller`, and the time step. Note that the control function must be called `control!`.
"""
abstract type Controller end


text:
@diffusion_process FitzHughNagumoAux{K,R} begin
    :dimensions
    process --> 2
    wiener --> 1

    :parameters
    (ϵ, s, γ, β, σ) --> K

    :auxiliary_info
    t0 --> Float64
    T --> Float64
    vT --> R

    :additional
    constdiff --> true
    linear --> true
end

B(t, P::FitzHughNagumoAux) = @SMatrix [
    1/P.ϵ-3*P.vT[1]^2/P.ϵ   -1/P.ϵ;
    P.γ                 -1.0
]

β(t, P::FitzHughNagumoAux) = @SVector [
    P.s/P.ϵ+2*P.vT[1]^3/P.ϵ,
    P.β
]

σ(t, P::FitzHughNagumoAux) = @SVector [0.0, P.σ]


text:
module Neuron_Test

using FactCheck, JNeuron

myimport=input(string(dirname(Base.source_path()),"/../examples/data/cell2.asc"));
blank_neuron=instantiate(myimport);
set_nsegs!(blank_neuron);

facts() do

    @fact typeof(blank_neuron) --> JNeuron.Neuron_0
    @fact length(blank_neuron.nodes) --> 925
    @fact Int64[length(blank_neuron.internal_nodes[i]) for i=1:4] --> [1,179,193,353]

    #Soma
    @fact blank_neuron.nodes[924].parent --> 0
    @fact blank_neuron.nodes[924].children --> [1,217,229,317,321,397,417,453,473,923,925]
    @fact blank_neuron.secs[198].pnode --> 923:925
end

myneuron1=add(blank_neuron,HH());
myneuron2=add(blank_neuron,(HH(),Passive()));
myneuron3=add(blank_neuron,(HH(),Passive()),(HH(),Passive()),Passive(),Passive());

facts() do
    @fact typeof(myneuron1) --> JNeuron.Neuron_1
    @fact typeof(myneuron2) --> JNeuron.Neuron_2
    @fact typeof(myneuron3) --> JNeuron.Neuron_3

    @fact size(myneuron1.soma.p) --> (16,1)
    @fact size(myneuron1.axon.p) --> (16,179)
    @fact size(myneuron1.dendrite.p) --> (16,193)
    @fact size(myneuron1.apical.p) --> (16,353)

    @fact size(myneuron2.soma.p) --> (17,1)
    @fact size(myneuron2.axon.p) --> (17,179)
    @fact size(myneuron2.dendrite.p) --> (17,193)
    @fact size(myneuron2.apical.p) --> (17,353)

    @fact size(myneuron3.soma.p) --> (17,1)
    @fact size(myneuron3.axon.p) --> (17,179)
    @fact size(myneuron3.dendrite.p) --> (3,193)
    @fact size(myneuron3.apical.p) --> (3,353)
end

mynetwork=Network((myneuron1,myneuron2,myneuron3),100.0);
run!(mynetwork,true);
run!(mynetwork,false);

facts() do

    @fact mynetwork.neur.N_1[1].v[1] --> roughly(-64.9604, atol=.01)
    @fact mynetwork.neur.N_2[1].v[1] --> roughly(-67.5302, atol=.01)
    @fact mynetwork.neur.N_3[1].v[1] --> greater_than(-70.0)
    @fact mynetwork.neur.N_3[1].v[1] --> less_than(-65.0)
end

end

max_stars_repo_name: paulmthompson/JNeuron

max_issues_repo_name: paulmthompson/JNeuron

max_forks_repo_name: paulmthompson/JNeuron