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


#=
This file is auto-generated. Do not edit.
=#
"""
    mutable struct BaseMachine <: Machine
        R::Float64
        Xd_p::Float64
        eq_p::Float64
        MVABase::Float64
        ext::Dict{String, Any}
        states::Vector{Symbol}
        n_states::Int64
        internal::InfrastructureSystemsInternal
    end

Parameters of an Automatic Voltage Regulator Type II -  Typical static exciter model

# Arguments
- `R::Float64`: Resistance after EMF in machine per unit, validation range: (0, nothing)
- `Xd_p::Float64`: Reactance after EMF in machine per unit, validation range: (0, nothing)
- `eq_p::Float64`: Fixed EMF behind the impedance, validation range: (0, nothing)
- `MVABase::Float64`: Nominal Capacity in MVA, validation range: (0, nothing)
- `ext::Dict{String, Any}`
- `states::Vector{Symbol}`
- `n_states::Int64`
- `internal::InfrastructureSystemsInternal`: power system internal reference, do not modify
"""
mutable struct BaseMachine <: Machine
    "Resistance after EMF in machine per unit"
    R::Float64
    "Reactance after EMF in machine per unit"
    Xd_p::Float64
    "Fixed EMF behind the impedance"
    eq_p::Float64
    "Nominal Capacity in MVA"
    MVABase::Float64
    ext::Dict{String, Any}
    states::Vector{Symbol}
    n_states::Int64
    "power system internal reference, do not modify"
    internal::InfrastructureSystemsInternal
end

function BaseMachine(R, Xd_p, eq_p, MVABase, ext=Dict{String, Any}(), )
    BaseMachine(R, Xd_p, eq_p, MVABase, ext, Vector{Symbol}(), 0, InfrastructureSystemsInternal(), )
end

function BaseMachine(; R, Xd_p, eq_p, MVABase, ext=Dict{String, Any}(), )
    BaseMachine(R, Xd_p, eq_p, MVABase, ext, )
end

# Constructor for demo purposes; non-functional.
function BaseMachine(::Nothing)
    BaseMachine(;
        R=0,
        Xd_p=0,
        eq_p=0,
        MVABase=0,
        ext=Dict{String, Any}(),
    )
end

"""Get BaseMachine R."""
get_R(value::BaseMachine) = value.R
"""Get BaseMachine Xd_p."""
get_Xd_p(value::BaseMachine) = value.Xd_p
"""Get BaseMachine eq_p."""
get_eq_p(value::BaseMachine) = value.eq_p
"""Get BaseMachine MVABase."""
get_MVABase(value::BaseMachine) = value.MVABase
"""Get BaseMachine ext."""
get_ext(value::BaseMachine) = value.ext
"""Get BaseMachine states."""
get_states(value::BaseMachine) = value.states
"""Get BaseMachine n_states."""
get_n_states(value::BaseMachine) = value.n_states
"""Get BaseMachine internal."""
get_internal(value::BaseMachine) = value.internal


@testset "io.jl" begin
    # Only test threading if 1.3+.
    if VERSION > v"1.2"
        @testset "threaded sampling" begin
            # Test that chains with the same seed will sample identically.
            @testset "rng" begin
                model = gdemo_default

                # multithreaded sampling with PG causes segfaults on Julia 1.5.4
                # https://github.com/TuringLang/Turing.jl/issues/1571
                samplers = @static if VERSION <= v"1.5.3" || VERSION >= v"1.6.0"
                    (
                        HMC(0.1, 7),
                        PG(10),
                        IS(),
                        MH(),
                        Gibbs(PG(3, :s), HMC(0.4, 8, :m)),
                        Gibbs(HMC(0.1, 5, :s), ESS(:m)),
                    )
                else
                    (
                        HMC(0.1, 7),
                        IS(),
                        MH(),
                        Gibbs(HMC(0.1, 5, :s), ESS(:m)),
                    )
                end
                for sampler in samplers
                    Random.seed!(5)
                    chain1 = sample(model, sampler, MCMCThreads(), 1000, 4)

                    Random.seed!(5)
                    chain2 = sample(model, sampler, MCMCThreads(), 1000, 4)

                    @test chain1.value == chain2.value
                end
            end

            # Smoke test for default sample call.
            Random.seed!(100)
            chain = sample(gdemo_default, HMC(0.1, 7), MCMCThreads(), 1000, 4)
            check_gdemo(chain)

            # run sampler: progress logging should be disabled and
            # it should return a Chains object
            sampler = Sampler(HMC(0.1, 7), gdemo_default)
            chains = sample(gdemo_default, sampler, MCMCThreads(), 1000, 4)
            @test chains isa MCMCChains.Chains
        end
    end
    @testset "chain save/resume" begin
        Random.seed!(1234)

        alg1 = HMCDA(1000, 0.65, 0.15)
        alg2 = PG(20)
        alg3 = Gibbs(PG(30, :s), HMCDA(500, 0.65, 0.05, :m))

        chn1 = sample(gdemo_default, alg1, 5000; save_state=true)
        check_gdemo(chn1)

        chn1_resumed = Turing.Inference.resume(chn1, 1000)
        check_gdemo(chn1_resumed)

        chn1_contd = sample(gdemo_default, alg1, 5000; resume_from=chn1)
        check_gdemo(chn1_contd)

        chn1_contd2 = sample(gdemo_default, alg1, 5000; resume_from=chn1, reuse_spl_n=1000)
        check_gdemo(chn1_contd2)

        chn2 = sample(gdemo_default, alg2, 1000; save_state=true)
        check_gdemo(chn2)

        chn2_contd = sample(gdemo_default, alg2, 1000; resume_from=chn2)
        check_gdemo(chn2_contd)

        chn3 = sample(gdemo_default, alg3, 1000; save_state=true)
        check_gdemo(chn3)

        chn3_contd = sample(gdemo_default, alg3, 1000; resume_from=chn3)
        check_gdemo(chn3_contd; atol=0.25)
    end
    @testset "Contexts" begin
        # Test LikelihoodContext
        @model testmodel(x) = begin
            a ~ Beta()
            lp1 = getlogp(_varinfo)
            x[1] ~ Bernoulli(a)
            global loglike = getlogp(_varinfo) - lp1
        end
        model = testmodel([1.0])
        varinfo = Turing.VarInfo(model)
        model(varinfo, Turing.SampleFromPrior(), Turing.LikelihoodContext())
        @test getlogp(varinfo) == loglike

        # Test MiniBatchContext
        @model testmodel(x) = begin
            a ~ Beta()
            x[1] ~ Bernoulli(a)
        end
        model = testmodel([1.0])
        varinfo1 = Turing.VarInfo(model)
        varinfo2 = deepcopy(varinfo1)
        model(varinfo1, Turing.SampleFromPrior(), Turing.LikelihoodContext())
        model(varinfo2, Turing.SampleFromPrior(), Turing.MiniBatchContext(Turing.LikelihoodContext(), 10))
        @test isapprox(getlogp(varinfo2) / getlogp(varinfo1), 10)
    end
    @testset "Prior" begin
        N = 5000

        # Note that all chains contain 3 values per sample: 2 variables + log probability
        Random.seed!(100)
        chains = sample(gdemo_d(), Prior(), N)
        @test chains isa MCMCChains.Chains
        @test size(chains) == (N, 3, 1)
        @test mean(chains, :s) ≈ 3 atol=0.1
        @test mean(chains, :m) ≈ 0 atol=0.1

        Random.seed!(100)
        chains = sample(gdemo_d(), Prior(), MCMCThreads(), N, 4)
        @test chains isa MCMCChains.Chains
        @test size(chains) == (N, 3, 4)
        @test mean(chains, :s) ≈ 3 atol=0.1
        @test mean(chains, :m) ≈ 0 atol=0.1

        Random.seed!(100)
        chains = sample(gdemo_d(), Prior(), N; chain_type = Vector{NamedTuple})
        @test chains isa Vector{<:NamedTuple}
        @test length(chains) == N
        @test all(length(x) == 3 for x in chains)
        @test all(haskey(x, :lp) for x in chains)
        @test mean(x[:s][1] for x in chains) ≈ 3 atol=0.1
        @test mean(x[:m][1] for x in chains) ≈ 0 atol=0.1
    end
end


module UtilsWrapper

using JavaCall

const AbsoluteDate = @jimport org.orekit.time.AbsoluteDate
const AbsolutePVCoordinates = @jimport org.orekit.utils.AbsolutePVCoordinates
const AbstractMultipleShooting = @jimport org.orekit.utils.AbstractMultipleShooting
const AngularCoordinates = @jimport org.orekit.utils.AngularCoordinates
const AngularDerivativesFilter = @jimport org.orekit.utils.AngularDerivativesFilter
const AttitudeProvider = @jimport org.orekit.attitudes.AttitudeProvider
const CartesianDerivativesFilter = @jimport org.orekit.utils.CartesianDerivativesFilter
const Class = @jimport java.lang.Class
const Collection = @jimport java.util.Collection
const Differentiation = @jimport org.orekit.utils.Differentiation
const EOPHistory = @jimport org.orekit.frames.EOPHistory
const ElevationMask = @jimport org.orekit.utils.ElevationMask
const Enum = @jimport java.lang.Enum
const ExtendedPVCoordinatesProvider = @jimport org.orekit.utils.ExtendedPVCoordinatesProvider
const Field = @jimport org.hipparchus.Field
const FieldAbsoluteDate = @jimport org.orekit.time.FieldAbsoluteDate
const FieldAbsolutePVCoordinates = @jimport org.orekit.utils.FieldAbsolutePVCoordinates
const FieldAngularCoordinates = @jimport org.orekit.utils.FieldAngularCoordinates
const FieldPVCoordinates = @jimport org.orekit.utils.FieldPVCoordinates
const FieldPVCoordinatesProvider = @jimport org.orekit.utils.FieldPVCoordinatesProvider
const FieldRotation = @jimport org.hipparchus.geometry.euclidean.threed.FieldRotation
const FieldTimeInterpolable = @jimport org.orekit.time.FieldTimeInterpolable
const FieldTimeSpanMap = @jimport org.orekit.utils.FieldTimeSpanMap
const FieldVector3D = @jimport org.hipparchus.geometry.euclidean.threed.FieldVector3D
const Frame = @jimport org.orekit.frames.Frame
const FundamentalNutationArguments = @jimport org.orekit.data.FundamentalNutationArguments
const GenericTimeStampedCache = @jimport org.orekit.utils.GenericTimeStampedCache
const Gradient = @jimport org.hipparchus.analysis.differentiation.Gradient
const IERSConventions = @jimport org.orekit.utils.IERSConventions
const IERSConventions_NutationCorrectionConverter = @jimport org.orekit.utils.IERSConventions$NutationCorrectionConverter
const ImmutableTimeStampedCache = @jimport org.orekit.utils.ImmutableTimeStampedCache
const InputStream = @jimport java.io.InputStream
const InterpolationTableLoader = @jimport org.orekit.utils.InterpolationTableLoader
const JString = @jimport java.lang.String
const LagrangianPoints = @jimport org.orekit.utils.LagrangianPoints
const List = @jimport java.util.List
const LoveNumbers = @jimport org.orekit.utils.LoveNumbers
const Map = @jimport java.util.Map
const MultipleShooter = @jimport org.orekit.utils.MultipleShooter
const MultipleShooting = @jimport org.orekit.utils.MultipleShooting
const NavigableSet = @jimport java.util.NavigableSet
const Object = @jimport java.lang.Object
const Optional = @jimport java.util.Optional
const OrbitType = @jimport org.orekit.orbits.OrbitType
const OrekitConfiguration = @jimport org.orekit.utils.OrekitConfiguration
const PVCoordinates = @jimport org.orekit.utils.PVCoordinates
const PVCoordinatesProvider = @jimport org.orekit.utils.PVCoordinatesProvider
const ParameterDriver = @jimport org.orekit.utils.ParameterDriver
const ParameterDriversList = @jimport org.orekit.utils.ParameterDriversList
const ParameterDriversList_DelegatingDriver = @jimport org.orekit.utils.ParameterDriversList$DelegatingDriver
const ParameterFunction = @jimport org.orekit.utils.ParameterFunction
const ParameterObserver = @jimport org.orekit.utils.ParameterObserver
const PoissonSeries_CompiledSeries = @jimport org.orekit.data.PoissonSeries$CompiledSeries
const PositionAngle = @jimport org.orekit.orbits.PositionAngle
const RealFieldElement = @jimport org.hipparchus.RealFieldElement
const Rotation = @jimport org.hipparchus.geometry.euclidean.threed.Rotation
const SecularAndHarmonic = @jimport org.orekit.utils.SecularAndHarmonic
const SortedSet = @jimport java.util.SortedSet
const SpacecraftState = @jimport org.orekit.propagation.SpacecraftState
const StateFunction = @jimport org.orekit.utils.StateFunction
const StateJacobian = @jimport org.orekit.utils.StateJacobian
const Stream = @jimport java.util.stream.Stream
const TimeInterpolable = @jimport org.orekit.time.TimeInterpolable
const TimeScalarFunction = @jimport org.orekit.time.TimeScalarFunction
const TimeScale = @jimport org.orekit.time.TimeScale
const TimeScales = @jimport org.orekit.time.TimeScales
const TimeSpanMap = @jimport org.orekit.utils.TimeSpanMap
const TimeSpanMap_Span = @jimport org.orekit.utils.TimeSpanMap$Span
const TimeStamped = @jimport org.orekit.time.TimeStamped
const TimeStampedAngularCoordinates = @jimport org.orekit.utils.TimeStampedAngularCoordinates
const TimeStampedCache = @jimport org.orekit.utils.TimeStampedCache
const TimeStampedFieldAngularCoordinates = @jimport org.orekit.utils.TimeStampedFieldAngularCoordinates
const TimeStampedFieldPVCoordinates = @jimport org.orekit.utils.TimeStampedFieldPVCoordinates
const TimeStampedGenerator = @jimport org.orekit.utils.TimeStampedGenerator
const TimeStampedPVCoordinates = @jimport org.orekit.utils.TimeStampedPVCoordinates
const TimeVectorFunction = @jimport org.orekit.time.TimeVectorFunction
const Vector3D = @jimport org.hipparchus.geometry.euclidean.threed.Vector3D

export AbsolutePVCoordinates
export AbstractMultipleShooting
export AngularCoordinates
export AngularDerivativesFilter
export CartesianDerivativesFilter
export Constants
export Differentiation
export ElevationMask
export ExtendedPVCoordinatesProvider
export FieldAbsolutePVCoordinates
export FieldAngularCoordinates
export FieldPVCoordinates
export FieldPVCoordinatesProvider
export FieldTimeSpanMap
export GenericTimeStampedCache
export IERSConventions
export ImmutableTimeStampedCache
export InterpolationTableLoader
export LagrangianPoints
export LoveNumbers
export MultipleShooter
export MultipleShooting
export OrekitConfiguration
export PVCoordinates
export PVCoordinatesProvider
export ParameterDriver
export ParameterDriversList
export ParameterFunction
export ParameterObserver
export SecularAndHarmonic
export StateFunction
export StateJacobian
export TimeSpanMap
export TimeStampedAngularCoordinates
export TimeStampedCache
export TimeStampedFieldAngularCoordinates
export TimeStampedFieldPVCoordinates
export TimeStampedGenerator
export TimeStampedPVCoordinates
export add
export add_constraint
export add_observer
export add_offset
export add_point
export add_valid_after
export add_valid_before
export apply_to
export approximate_as_polynomial_only
export compare_to
export compute
export create_from_modified_rodrigues
export cross_product
export describe_constable
export differentiate
export empty_cache
export equals
export estimate_rate
export estimate_velocity
export evaluate_tc
export extract_range
export filter
export find_by_name
export fit
export generate
export get
export get_abscissa_grid
export get_acceleration
export get_all
export get_angular_velocity
export get_cache_slots_number
export get_class
export get_date
export get_declaring_class
export get_drivers
export get_earliest
export get_earth_orientation_angle_function
export get_elevation
export get_entries
export get_eop_tidal_correction
export get_filter
export get_fitted_parameters
export get_frame
export get_gast_function
export get_generate_calls
export get_generator
export get_get_neighbors_calls
export get_gmst_function
export get_gmst_rate_function
export get_harmonic_amplitude
export get_identity
export get_imaginary
export get_latest
export get_love_numbers
export get_max_order
export get_max_slots
export get_max_span
export get_max_value
export get_mean_obliquity_function
export get_min_value
export get_modified_rodrigues
export get_momentum
export get_name
export get_nb_params
export get_neighbors
export get_neighbors_size
export get_new_slot_quantum_gap
export get_nominal_tidal_displacement
export get_normalized_value
export get_nutation_arguments
export get_nutation_correction_converter
export get_nutation_function
export get_nutation_reference_epoch
export get_observers
export get_ocean_pole_tide
export get_ordinate_grid
export get_permanent_tide
export get_plus
export get_position
export get_precession_function
export get_pv_coordinates
export get_real
export get_reference_date
export get_reference_value
export get_rotation
export get_rotation_acceleration
export get_rotation_rate
export get_scale
export get_size
export get_slots
export get_slots_evictions
export get_solid_pole_tide
export get_span
export get_tidal_displacement_frequency_correction_diurnal
export get_tidal_displacement_frequency_correction_zonal
export get_tide_frequency_dependence_function
export get_transitions
export get_value
export get_values_samples
export get_velocity
export get_xy_sp_xy2_function
export get_zero
export hash_code
export interpolate
export is_selected
export load_data
export max_value_changed
export mean_derivative
export mean_second_derivative
export mean_value
export min_value_changed
export name
export name_changed
export negate
export normalize
export notify
export notify_all
export ordinal
export osculating_derivative
export osculating_second_derivative
export osculating_value
export reference_date_changed
export reference_value_changed
export remove_observer
export replace_observer
export reset_fitting
export revert
export scale_changed
export selection_changed
export set_cache_slots_number
export set_closed_orbit_constraint
export set_epoch_freedom
export set_max_value
export set_min_value
export set_name
export set_normalized_value
export set_patch_point_component_freedom
export set_reference_date
export set_reference_value
export set_scale
export set_selected
export set_value
export shifted_by
export sort
export still_accepts_data
export subtract_offset
export to_absolute_pv_coordinates
export to_angular_coordinates
export to_derivative_structure_pv
export to_derivative_structure_rotation
export to_derivative_structure_vector
export to_field_pv_coordinates_provider
export to_pv_coordinates
export to_string
export to_taylor_provider
export to_time_stamped_pv_coordinates
export to_univariate_derivative1_rotation
export to_univariate_derivative1_vector
export to_univariate_derivative1pv
export to_univariate_derivative2_vector
export to_univariate_derivative2pv
export value
export value_changed
export value_of
export values
export wait

include("absolute_pv_coordinates.jl")
include("abstract_multiple_shooting.jl")
include("angular_coordinates.jl")
include("angular_derivatives_filter.jl")
include("cartesian_derivatives_filter.jl")
include("constants.jl")
include("differentiation.jl")
include("elevation_mask.jl")
include("extended_pv_coordinates_provider.jl")
include("field_absolute_pv_coordinates.jl")
include("field_angular_coordinates.jl")
include("field_pv_coordinates.jl")
include("field_pv_coordinates_provider.jl")
include("field_time_span_map.jl")
include("generic_time_stamped_cache.jl")
include("iers_conventions.jl")
include("immutable_time_stamped_cache.jl")
include("interpolation_table_loader.jl")
include("lagrangian_points.jl")
include("love_numbers.jl")
include("multiple_shooter.jl")
include("multiple_shooting.jl")
include("orekit_configuration.jl")
include("parameter_driver.jl")
include("parameter_drivers_list.jl")
include("parameter_function.jl")
include("parameter_observer.jl")
include("pv_coordinates.jl")
include("pv_coordinates_provider.jl")
include("secular_and_harmonic.jl")
include("state_function.jl")
include("state_jacobian.jl")
include("time_span_map.jl")
include("time_stamped_angular_coordinates.jl")
include("time_stamped_cache.jl")
include("time_stamped_field_angular_coordinates.jl")
include("time_stamped_field_pv_coordinates.jl")
include("time_stamped_generator.jl")
include("time_stamped_pv_coordinates.jl")

end


export
    InstantAlgorithm,
    infer

"""
    InstantAlgorithm

    Algorithm that runs once on an `InstantNetwork`.
"""
abstract type InstantAlgorithm <: Algorithm end

# FIXME:
# placeholder_beliefs is intended to provide information about the distribution of placeholders to an algorithm.
# The current interface allows marginal distributions over placeholders to be passed, which is fine for an
# algorithm like BP, but if an algorithm needs joint information it has to get it in a different way.
# In particular, during particle filtering, the importance sampling algorithm needs joint samples of the previous state.
# Fortunately, it can get it out of the :particles state in the runtime, but this is clunky and not in line with the
# intentions of the general infer method.
# We should introduce sfuncs that define distributions over dictionaries, and provide operators to produce joint samples
# or compute marginals over individual variables, as they are able to support them. placeholder_beliefs should be
# such an sfunc.
"""
    infer(algorithm::InstantAlgorithm, runtime::InstantRuntime,
        evidence::Dict{Symbol, Score},
        interventions::Dict{Symbol, Dist},
        placeholder_beliefs::Dict{Symbol, Dist})

Run the inference algorithm.

Stores the results in `runtime`. The format of these results is up to
`algorithm`, but they should be usable by queries with this algorithm.

# Arguments
- `algorithm`: The instant algorithm to run.
- `runtime`: The runtime in which to run the algorithm.
- `evidence`: The supplied evidence, which defaults to `Dict()`.
- `interventions`: The supplied interventions, which defaults to `Dict()`.
- `placeholder_beliefs`: The beliefs associated with the placeholders in the
network, which default to `Dict()`. Instant algorithms might require that a belief be
supplied for every placeholder in `network`.
"""


lattice = load_lattice(joinpath(@__DIR__, "CrFe2Ga_qe.in"))
atoms   = load_atoms(  joinpath(@__DIR__, "CrFe2Ga_qe.in"))
atoms   = attach_pseudos(atoms, Cr="hgh/pbe/cr-q14.hgh",
                                Fe="hgh/pbe/fe-q16.hgh",
                                Ga="hgh/pbe/ga-q3.hgh")
magnetic_moments = build_magnetic_moments(atoms; Ga=0.0, Fe=5.0, Cr=5.0)

cases = Dict{String, Any}(
    "adaptive" => (damping=DFTK.AdaptiveDamping(), ),  # Parameter-free adaptive damping
)
for α in 5:5:20
    cases["standard_$(@sprintf "%03d" α)"] = (damping=DFTK.FixedDamping(α / 100), )
end
for i in [1, 3]
    cases["adaptive_0$i"] = (damping=DFTK.AdaptiveDamping(i / 10), )
end

config[@__DIR__] = (
    directory        =  @__DIR__,
    lattice          = lattice,
    atoms            = atoms,
    magnetic_moments = magnetic_moments,
    #
    kwargs_model     = (smearing=Smearing.Gaussian(), temperature=0.01),
    Ecut             = 45,
    kgrid            = [13, 13, 13],
    supersampling    = 2.0,
    potential        = nothing,
    kwargs_scf       = (acceleration=DFTK.AndersonAcceleration(;m=0), ),
    mixing           = KerkerMixing(),
    cases            = cases,
)


# TODO: this is really a module for face-centered diamond-cubic
#       and over time it should become much more general of course

module Si

using JuLIP, JuLIP.Potentials, QuadGK, ForwardDiff, LinearAlgebra

import JuLIPMaterials

using JuLIPMaterials.CLE: elastic_moduli, voigt_moduli,
         fourth_order_basis, sextic_roots, A_coefficients, D_coefficients,
         little_a

const CLE = JuLIPMaterials.CLE
const FCC = JuLIPMaterials.FCC
# CauchyBorn = JuLIPMaterials.CauchyBorn

"""
`si110_plane(s::AbstractString) -> at::ASEAtoms, b, xcore, a `

Generates a unit cell for an FCC crystal with orthogonal cell vectors chosen
such that the F1 direction is the burgers vector and the F3 direction the normal
to the standard edge dislocation:
   b = F1 ~ a1;  ν = F3 ~ a2-a3
The third cell vector: F2 ~ a * e1. Here, ~ means they are rotated from the
a1, a2, a3 directions. This cell contains two atoms.

Returns
* `at`: Atoms object
* `b` : Burgers vector
* `xcore` : a core-offset (to add to any lattice position)
* `a` : lattice parameter
"""
function si110_plane(s::Symbol;
                     a = cell(bulk(s, cubic=true))[1,1] )
   # TODO: can si110_plane be combined with FCC.fcc_110_plane?
   @assert s == :Si
   # for Si, default a = 5.43
   # TODO: ensure only that s is an FCC species
   # ====================================================
   # construct the cell matrix
   F = a * JMat( [ sqrt(2)/2  0    0;
                         0    1    0;
                         0    0    sqrt(2)/2 ] )
   X = a * [ JVec([         0.0,   0.0,          0.0 ]),
             JVec([ 0.5/sqrt(2),   0.5,  0.5/sqrt(2) ]),
             JVec([         0.0, -0.25,  0.5/sqrt(2) ]),
             JVec([ 0.5/sqrt(2),  0.25,          0.0 ]) ]
   # construct ASEAtoms
   at = Atoms(s, X)
   set_cell!(at, F')
   # compute a burgers vector in these coordinates
   b = a * sqrt(2)/2 * JVec([1.0,0.0,0.0])
   # compute a core-offset (to add to any lattice position)
   xcore = [-.7, 1.0, 0] * a / 5.43
   # return the information
   return at, b, xcore, a
end

"a fully equilibrated SW potential"
function sw_eq()
    T(σ, at) = tr(stress(StillingerWeber(σ=σ), at))
    at = bulk(:Si, pbc=true)
    r0 = 2.09474
    r1 = r0 - 0.1
    s0, s1 = T(r0, at), T(r1, at)
    while (abs(s1) > 1e-8) && abs(r0 - r1) > 1e-8
        rnew = (r0 * s1 - r1 * s0) / (s1 - s0)
        r0, r1 = r1, rnew
        s0, s1 = s1, T(rnew, at)
    end
    return StillingerWeber(σ=r1)
end

function si110_cluster(sym, L; kwargs...)
    # @assert isodd(L)   # TODO: why?
    atu, b, _, a = si110_plane(sym; kwargs...)
    at = atu * (L, L, 1)
    set_pbc!(at, (false, false, true))
    b = b[1]
    X = positions(at)
    # This choice picks the lower left and upper right atom (not site) positions
    #   TODO I don't like the [-1, 0.1, 0] at all
    xcore = (1/2)*(X[length(X)-2]+X[3]) + [-1, 0.1, 0]
    return at, b, xcore
end

# """
# a function that identifies multi-lattice structure in 2 layers of bulk-Si
# (yes - very restrictive but will do for now!)
# """
# function si_multilattice(at)
#     TOL = 0.2
#     J0 = Int[]
#     J1 = Int[]
#     Jdel = Int[]
#     for (i, j, r, R, _) in sites(at, rnn("Si")+0.1)
#         foundneig = false
#         for (jj, RR) in zip(j, R)
#             if (RR[1] == 0.0) && (abs(RR[2] - 1.3575) < TOL)
#                 # neighbour above >> make (i, jj) a site
#                 push!(J0, i)
#                 push!(J1, jj)
#                 foundneig = true
#                 break
#             elseif (RR[1] == 0.0) && (abs(RR[2] + 1.3575) < TOL)
#                 # neighbour below >> (jj, i) is a site that will be pushed when i ↔ jj
#                 foundneig = true
#                 break
#             end
#         end
#         if !foundneig
#             # i has no neighbour above or below >> probably we just get rid of it
#             push!(Jdel, i)
#         end
#     end
#     return J0, J1, Jdel
# end
#
#
#
# function symml_displacement!(at, u)
#     I0, I1, Idel = si_multilattice(at)
#     @assert isempty(Idel)  # if Idel is not empty then (for now) we don't know what to do
#     X = positions(at)
#     W = CauchyBorn.WcbQuad(calculator(at))   # TODO: generalise this to general calculators
#     F0 = cell(W.at)'
#     p0 = W(F0)
#     # transformation matrices
#     Tp = [0 1/√2  -1/√2; 1 0 0; 0 1/√2 1/√2]
#     Tm = diagm([1,1,-1]) * Tp
#
#     for (i0, i1) in zip(I0, I1)   # each pair (i0, i1) corresponds to a ML lattice site
#         x0, x1 = X[i0], X[i1]
#         x1[3] > x0[3] ? T = Tp : T = Tm
#         x̄ = 0.5 * (x0 + x1)   # centre of mass of the bond
#         U, ∇U = u(x̄)          # displacement and displacement gradient
#         F = T' * (I + ∇U) * T
#         q = T * (W(F * F0) - p0)    # construct the shift corresponding to F = Id + ∇U
#         X[i0], X[i1] = x0 + U - 0.5 * q, x1 + U + 0.5 * q
#     end
#     set_positions!(at, X)
#     return at
# end
#
# function ml_displacement!(at, u)
#     I0, I1, Idel = si_multilattice(at)
#     @assert isempty(Idel)  # if Idel is not empty then (for now) we don't know what to do
#     X = positions(at)
#     W = CauchyBorn.WcbQuad(calculator(at))
#
#     # transformation matrices
#     Tp = [0 1/√2  -1/√2; 1 0 0; 0 1/√2 1/√2]
#     Tm = diagm([1,1,-1]) * Tp
#
#     F0 = cell(W.at)'  # get reference information
#     p0 = W(F0)
#
#     for (i0, i1) in zip(I0, I1)   # each pair (i0, i1) corresponds to a ML lattice site
#         x0, x1 = X[i0], X[i1]
#         x1[3] > x0[3] ? T = Tp : T = Tm
#         U, ∇U = u(x0)            # displacement and displacement gradient
#         F = T' * (I + ∇U) * T
#         q = T * (W(F * F0) - p0)   # construct the shift corresponding to F = Id + ∇U
#         X[i0], X[i1] = x0 + U, x1 + U + q
#     end
#     set_positions!(at, X)
#     return at
# end

"""
`edge110` : produces a high-quality CLE solution for an edge dislocation
in bulk silicon, 110 orientation.

## Keyword Arguments

* a : lattice parameter, default value is the default lattice parameter of the species, other allows values are `:equilibrate`
"""
function edge110(sym::Symbol, R::Real;
                  truncate=true, cle=:anisotropic, ν=0.25,
                  calc=sw_eq(), symmetrize = nothing,  # switch back to sym = true
                  TOL=1e-4, zDir=1,
                  eos_correction = true,
                  a = cell(bulk(sym, cubic=true))[1,1])

   @assert sym == :Si

   # compute the lattice parameter
   if a == :equilibrate
       atu = bulk(sym, cubic=true, pbc=true)
       set_calculator!(atu, calc)
       variablecell!(atu)
       minimise!(atu)
       a = cell(atu)[1,1]
   end

   # setup undeformed geometry
   at, b, x0 = si110_cluster(sym, ceil(Int, R/rnn(:Si)+3); a = a)
   set_calculator!(at, calc)

   # if cle == :anisotropic
   #    W = CauchyBorn.WcbQuad(calc)
   #    C = elastic_moduli(W)
   #    Cv = round.(voigt_moduli(C), 8)
   #    U = CLE.EdgeCubic(b, Cv, a, x0 = x0)
   # elseif cle == :isotropic
   #    @assert sym == nothing
   #    U = xx -> ([CLE.u_edge_isotropic(xx[1]-x0[1], xx[2]-x0[2], b, ν)...; 0.0], nothing)
   # else
   #    error("unknown `cle` option")
   # end

  @assert symmetrize == nothing
  U = xx -> ([CLE.u_edge_isotropic(xx[1]-x0[1], xx[2]-x0[2], b, ν)...; 0.0], nothing)

   if symmetrize == true
      @assert cle == :anisotropic
      symml_displacement!(at, U)
   elseif symmetrize == false
      @assert cle == :anisotropic
      ml_displacement!(at, U)
   elseif symmetrize == nothing
      X = positions(at)
      for (i,x) in enumerate(X)
         X[i] += U(x)[1]
      end
      set_positions!(at, X)
   else
      error("edge110: unknown parameters symmetrize = $(symmetrize)")
   end

   if truncate
      r = [ norm(x[1:2] - x0[1:2]) for x in positions(at) ]
      Ikeep = findall( r .<= R )
      X = positions(at)[Ikeep]
      C = cell(at)
      at = Atoms(:Si, X)
      set_cell!(at, C)
   end

   return at, x0
end

end


#****************************************************************************
#                     Molecular Dynamics Potentials (MDP)
#                            CESMIX-MIT Project
#
#  Contributing authors: Ngoc-Cuong Nguyen (cuongng@mit.edu, exapde@gmail.com)
#****************************************************************************

function readLAMMPS(filename,datamode)

if datamode==0 # binary
    tmp = reinterpret(Float64,read(filename));
else # text
    tmp = Main.DelimitedFiles.readdlm(filename);
end

config = CONFIGStruct();

config.dim = 3;
config.ncq = 0;
config.ncv = 0;
config.nce = 1;
config.ncf = 3;
config.ncx = 3;

k = 1;
sz1, sz2 = size(tmp);
while (k < sz1)
    n = tmp[k,1]; # number of atoms
    if k ==1
        config.natom = reshape([n],1,1);
    else
        config.natom = [config.natom n];
    end
    b = transpose(reshape(tmp[(k+1):(k + n),:], n, 7));
    if k==1
        config.t = reshape(b[1,:],1,n);
        config.x = b[2:4,:];
        config.f = b[5:7,:];
    else
        config.t = [config.t reshape(b[1,:],1,n)];
        config.x = [config.x b[2:4,:]];
        config.f = [config.f b[5:7,:]];
    end
    b = tmp[(k+n+1):(k+n+2),:];
    if k == 1
        config.e = reshape([b[1,1]],1,1);
        config.a = reshape([b[2,1] 0 0],3,1);
        config.b = reshape([0 b[2,2] 0],3,1);
        config.c = reshape([0 0 b[2,3]],3,1);
    else
        config.e = [config.e reshape([b[1,1]],1,1)];
        config.a = [config.a [b[2,1] 0 0]'];
        config.b = [config.b [0 b[2,2] 0]'];
        config.c = [config.c [0 0 b[2,3]]'];
    end
    k = k + n + 3;
end

config.q = reshape([],0,2);
config.v = reshape([],0,2);

return config

end


# define the generative model for the multiple pedestrian crosswalk environment

### REWARD MODEL ##################################################################################

function POMDPs.reward(pomdp::OCPOMDP, s::OCState, a::OCAction, sp::OCState)
    r = 0.
    ego = sp[findfirst(EGO_ID, sp)]
    if is_crash(sp)
        r += pomdp.collision_cost
    end
    if ego.state.posG.x >= pomdp.ego_goal
        r += pomdp.goal_reward
    elseif a.acc > 0.
        r += pomdp.action_cost
    else
        r += pomdp.action_cost
    end
    return r
end

### TERMINAL STATES ###############################################################################

function POMDPs.isterminal(pomdp::OCPOMDP, s::OCState)
    ego = s[findfirst(EGO_ID, s)]
    return is_crash(s) || ego.state.posG.x >= pomdp.ego_goal
end

### TRANSITION MODEL ##############################################################################

function POMDPs.gen(::DDNNode{:sp}, pomdp::OCPOMDP, s::OCState, a::OCAction, rng::AbstractRNG)
    actions = Array{Any}(undef, length(s))
    pomdp.models[1].a = a
    sp = deepcopy(s) #XXX bad
    # max_id = 0
    # for veh in sp
    #     if veh.id > max_id
    #         max_id = veh.id
    #     end
    # end
    # if rand(rng) < pomdp.p_birth && max_id < pomdp.max_peds+1 && !pomdp.no_ped
    # if rand(rng) < 0.01 && max_id < pomdp.max_peds+1 && !pomdp.no_ped
    #     # println("Spawning new pedestrians")
    #     new_ped = initial_pedestrian(pomdp, sp, rng)
    #     pomdp.models[new_ped.id] = ConstantPedestrian(dt = pomdp.ΔT)#TODO parameterized
    #     push!(sp, new_ped)
    #     actions = Array{Any}(undef, length(sp))
    # end
    get_actions!(actions, sp, pomdp.env.roadway, pomdp.models)
    tick!(sp, pomdp.env.roadway, actions, pomdp.ΔT)
    clean_scene!(pomdp.env, sp)
    return sp
end

# function clean_scene!(env::CrosswalkEnv, scene::Scene)
#     for veh in scene
#         if veh.id == 1
#             continue
#         end
#         y = veh.state.posG.y
#         if y >= env.params.crosswalk_length/4
#             deleteat!(scene, findfirst(veh.id, scene))
#         end
#     end
# end

function AutomotiveDrivingModels.propagate(veh::Vehicle, action::OCAction, roadway::Roadway, Δt::Float64)
    x_ = veh.state.posG.x + veh.state.v*Δt + action.acc*Δt^2/2
    if x_ <= veh.state.posG.x
        x_ = veh.state.posG.x
    end
    v_ = veh.state.v + action.acc*Δt
    v_ = clamp(v_, 0., 8.) # parameterized
    return VehicleState(VecSE2(x_, veh.state.posG.y, veh.state.posG.θ), roadway, v_)
end

### INITIAL STATES ################################################################################

function POMDPs.initialstate(pomdp::OCPOMDP, rng::AbstractRNG)
    pomdp.no_ped = rand(rng) < pomdp.no_ped_prob
    scene = Scene()
    if !pomdp.no_ped
        for i=1:pomdp.max_peds
            if rand(rng) < pomdp.p_birth # pedestrian appear
                new_ped = initial_pedestrian(pomdp, scene, rng, true)
                pomdp.models[new_ped.id] = ConstantPedestrian(dt = pomdp.ΔT)
                push!(scene, new_ped)
            end
        end
    end
    ego = initial_ego(pomdp, rng)
    push!(scene, ego)
    return scene
end

"""
    initial_ego(pomdp::OCPOMDP, rng::AbstractRNG)
generate an initial state for the ego car, returns a Vehicle object
"""
function initial_ego(pomdp::OCPOMDP, rng::AbstractRNG)
    env = pomdp.env
    x0 = pomdp.ego_start
    y0 = 0. # parameterize
    v0 = rand(rng, Distributions.Uniform(6., pomdp.env.params.speed_limit)) #TODO parameterize
    return Vehicle(VehicleState(VecSE2(x0, y0, 0.0), env.roadway.segments[1].lanes[1], env.roadway, v0),
                   VehicleDef(), 1)
end

"""
    initial_pedestrian(pomdp::OCPOMDP, scene::Scene, rng::AbstractRNG)
Create a new pedestrian entity at its initial state
"""
function initial_pedestrian(pomdp::OCPOMDP, scene::Scene, rng::AbstractRNG, first_scene::Bool = false)
    env = pomdp.env
    # lateral position
    d_lat = Distributions.Uniform(env.params.roadway_length - env.params.crosswalk_width + 2, env.params.roadway_length + env.params.crosswalk_width-1)
    x0 = 0.5*rand(rng, d_lat)

    # vertical position
    y0 = pomdp.ped_start
    if first_scene
        y0 = rand(rng, pomdp.ped_start:pomdp.ped_goal)
    end

    # orientation
    θ = π/2

    #velocity
    v0 = rand(rng, Distributions.Uniform(0., env.params.ped_max_speed))
    cw_roadway = Roadway([RoadSegment(2, [env.crosswalk])]);
    ped_initialstate = VehicleState(VecSE2(x0, y0, θ), env.crosswalk, cw_roadway, v0);

    # new id, increment last id
    max_id = 0
    for veh in scene
        if veh.id > max_id
            max_id = veh.id
        end
    end
    id = max_id + 1
    if max_id == 0
        id = 2
    end

    return Vehicle(ped_initialstate, PEDESTRIAN_DEF, id)
end

function POMDPs.initialstate_distribution(pomdp::OCPOMDP)
    return GenerativeDist(pomdp)
end

### Observations ##################################################################################

# uncomment for LIDAR observation
# function POMDPs.generate_o(pomdp::OCPOMDP, s::OCState, a::OCAction, sp::OCState, rng::AbstractRNG)
#     observe!(pomdp.sensor, sp, pomdp.env.roadway, EGO_ID)
#     return pomdp.sensor
# end
#
# function POMDPs.generate_o(pomdp::OCPOMDP, s::OCState, rng::AbstractRNG)
#     observe!(pomdp.sensor, s, pomdp.env.roadway, EGO_ID)
#     return pomdp.sensor
# end
#
# function POMDPs.convert_o(::Type{Vector{Float64}}, o::OCObs, pomdp::OCPOMDP)
#     return [o.ranges/pomdp.sensor.max_range, o.range_rates/pomdp.env.params.speed_limit]
# end

# function POMDPs.convert_o(::Type{Vector{Float64}}, o::OCObs, pomdp::OCPOMDP)
#     return [o.ranges/pomdp.sensor.max_range, o.range_rates/pomdp.env.params.speed_limit]
# end

# uncomment for vector representation
function POMDPs.gen(::DDNNode{:o}, pomdp::OCPOMDP, s::OCState, a::OCAction, sp::OCState, rng::AbstractRNG)
    pos_noise = pomdp.pos_obs_noise
    vel_noise = pomdp.vel_obs_noise
    o = zeros(2 + 2*pomdp.max_peds)
    ego = sp[findfirst(EGO_ID, sp)].state
    o[1] = ego.posG.x
    o[2] = ego.v
    ped_off = get_off_the_grid(pomdp)
    for i=2:pomdp.max_peds+1
        o[2*i - 1] = ped_off.posG.y
        o[2*i] = 0.
    end
    for veh in sp
        if veh.id == EGO_ID
            continue
        end
        @assert veh.id <= pomdp.max_peds+1
        ped = veh.state
        if is_observable_fixed(ped, ego, pomdp.env)
            o[2*veh.id - 1] = ped.posG.y + pos_noise*randn(rng)
            o[2*veh.id] = ped.v + vel_noise*randn(rng)
        end
    end
    # rescale
    o[1] /= pomdp.ego_goal
    o[2] /= pomdp.env.params.speed_limit
    for i=2:pomdp.max_peds+1
        o[2*i - 1] /= pomdp.ped_goal
        o[2*i] /= 2. # XXX parameterized
    end
    return o
end

function POMDPs.initialobs(pomdp::OCPOMDP, s::OCState, rng::AbstractRNG)
    return gen(DDNNode(:o), pomdp, s, OCAction(0.), s, rng::AbstractRNG)
end

# function POMDPs.convert_o(::Type{Vector{Float64}}, o::OCObs, pomdp::OCPOMDP)
#     o[1] /= pomdp.ego_goal
#     o[2] /= pomdp.env.params.speed_limit
#     for i=2:pomdp.max_peds+1
#         o[2*i - 1] /= pomdp.ped_goal
#         o[2*i] /= 2. # XXX parameterized
#     end
#     return o
# end

### All together ##################################################################################

function POMDPs.gen(::DDNOut{(:sp,:o,:r)}, pomdp::OCPOMDP, s::OCState, a::OCAction, rng::AbstractRNG)
    sp = gen(DDNOut(:sp), pomdp, s, a, rng)
    o = gen(DDNNode(:o), pomdp, s, a, sp, rng)
    r = reward(pomdp, s, a, sp)
    return sp, o, r
end

#### Helpers

"""
    get_off_the_grid(pomdp::OCPOMDP)
return the off the grid pedestrian state
"""
function get_off_the_grid(pomdp::OCPOMDP)
    env = pomdp.env
    x_ped = 0.5*(env.params.roadway_length - env.params.crosswalk_width + 1)
    cl = pomdp.env.params.crosswalk_length
    pos = VecSE2(x_ped, -cl/2 - 1, pi/2)
    return VehicleState(pos, pomdp.env.roadway, Inf)
end


const use_return  = haskey(ENV, "USE_RETURN" ) ? parse(Bool, ENV["USE_RETURN"] ) : false
const USE_GPU     = haskey(ENV, "USE_GPU"    ) ? parse(Bool, ENV["USE_GPU"]    ) : false
const do_viz      = haskey(ENV, "DO_VIZ"     ) ? parse(Bool, ENV["DO_VIZ"]     ) : true
const do_save     = haskey(ENV, "DO_SAVE"    ) ? parse(Bool, ENV["DO_SAVE"]    ) : false
const do_save_viz = haskey(ENV, "DO_SAVE_VIZ") ? parse(Bool, ENV["DO_SAVE_VIZ"]) : false
const nx          = haskey(ENV, "NX"         ) ? parse(Int , ENV["NX"]         ) : 64 - 1
const ny          = haskey(ENV, "NY"         ) ? parse(Int , ENV["NY"]         ) : 64 - 1
const nz          = haskey(ENV, "NZ"         ) ? parse(Int , ENV["NZ"]         ) : 64 - 1
###
using ParallelStencil
using ParallelStencil.FiniteDifferences3D
@static if USE_GPU
    @init_parallel_stencil(CUDA, Float64, 3)
else
    @init_parallel_stencil(Threads, Float64, 3)
end
using ImplicitGlobalGrid, Plots, Printf, Statistics, LinearAlgebra, MAT
import MPI
# Global reductions
mean_g(A)    = (mean_l = mean(A); MPI.Allreduce(mean_l, MPI.SUM, MPI.COMM_WORLD)/MPI.Comm_size(MPI.COMM_WORLD))
norm_g(A)    = (sum2_l = sum(A.^2); sqrt(MPI.Allreduce(sum2_l, MPI.SUM, MPI.COMM_WORLD)))
maximum_g(A) = (max_l  = maximum(A); MPI.Allreduce(max_l,  MPI.MAX, MPI.COMM_WORLD))
minimum_g(A) = (min_l  = minimum(A); MPI.Allreduce(min_l,  MPI.MIN, MPI.COMM_WORLD))
# CPU functions
@views av_xza(A) = (A[1:end-1,:,1:end-1] .+ A[1:end-1,:,2:end] .+ A[2:end,:,1:end-1] .+ A[2:end,:,2:end]).*0.25
@views av_zi(A)  = (A[2:end-1,2:end-1,2:end-2] .+ A[2:end-1,2:end-1,3:end-1]).*0.5
@views av_za(A)  = (A[:,:,1:end-1] .+ A[:,:,2:end]).*0.5
@views inn(A)    =  A[2:end-1,2:end-1,2:end-1]

@parallel function smooth!(A2::Data.Array, A::Data.Array, fact::Data.Number)
    @inn(A2)   = @inn(A) + 1.0/6.1/fact*(@d2_xi(A) + @d2_yi(A) + @d2_zi(A))
    return
end

@parallel function compute_maxloc!(Musτ2::Data.Array, Musτ::Data.Array)
    @inn(Musτ2) = @maxloc(Musτ)
    return
end

macro Mu_eff() esc(:(1.0/(1.0/@all(Musτ) + 1.0/(G*dt)))) end
@parallel function compute_iter_params!(dτ_Rho::Data.Array, Gdτ::Data.Array, Musτ::Data.Array, Vpdτ::Data.Number, G::Data.Number, dt::Data.Number, Re::Data.Number, r::Data.Number, max_lxyz::Data.Number)
    @all(dτ_Rho) = Vpdτ*max_lxyz/Re/@Mu_eff()
    @all(Gdτ)    = Vpdτ^2/@all(dτ_Rho)/(r+2)
    return
end

@parallel function assign_τ!(τxx::Data.Array, τyy::Data.Array, τzz::Data.Array, τxy::Data.Array,  τxz::Data.Array,  τyz::Data.Array, τxx_o::Data.Array, τyy_o::Data.Array, τzz_o::Data.Array, τxy_o::Data.Array, τxz_o::Data.Array, τyz_o::Data.Array)
    @all(τxx_o) = @all(τxx)
    @all(τyy_o) = @all(τyy)
    @all(τzz_o) = @all(τzz)
    @all(τxy_o) = @all(τxy)
    @all(τxz_o) = @all(τxz)
    @all(τyz_o) = @all(τyz)
    return
end

@parallel function compute_P!(∇V::Data.Array, Pt::Data.Array, Vx::Data.Array, Vy::Data.Array, Vz::Data.Array, Gdτ::Data.Array, r::Data.Number, dx::Data.Number, dy::Data.Number, dz::Data.Number)
    @all(∇V)  = @d_xa(Vx)/dx + @d_ya(Vy)/dy + @d_za(Vz)/dz
    @all(Pt)  = @all(Pt) - r*@all(Gdτ)*@all(∇V)
    return
end

macro inn_yz_Gr() esc(:( @inn_yz(Gdτ)/(G*dt) )) end
macro inn_xz_Gr() esc(:( @inn_xz(Gdτ)/(G*dt) )) end
macro inn_xy_Gr() esc(:( @inn_xy(Gdτ)/(G*dt) )) end
macro av_xyi_Gr() esc(:( @av_xyi(Gdτ)/(G*dt) )) end
macro av_xzi_Gr() esc(:( @av_xzi(Gdτ)/(G*dt) )) end
macro av_yzi_Gr() esc(:( @av_yzi(Gdτ)/(G*dt) )) end
@parallel function compute_τ!(τxx::Data.Array, τyy::Data.Array, τzz::Data.Array, τxy::Data.Array, τxz::Data.Array, τyz::Data.Array, τxx_o::Data.Array, τyy_o::Data.Array, τzz_o::Data.Array, τxy_o::Data.Array, τxz_o::Data.Array, τyz_o::Data.Array, Vx::Data.Array, Vy::Data.Array, Vz::Data.Array, Mus::Data.Array, Gdτ::Data.Array, G::Data.Number, dt::Data.Number, dx::Data.Number, dy::Data.Number, dz::Data.Number)
    @all(τxx) = (@all(τxx) + @all(τxx_o)*@inn_yz_Gr() + 2.0*@inn_yz(Gdτ)*(@d_xi(Vx)/dx))/(1.0 + @inn_yz(Gdτ)/@inn_yz(Mus) + @inn_yz_Gr())
    @all(τyy) = (@all(τyy) + @all(τyy_o)*@inn_xz_Gr() + 2.0*@inn_xz(Gdτ)*(@d_yi(Vy)/dy))/(1.0 + @inn_xz(Gdτ)/@inn_xz(Mus) + @inn_xz_Gr())
    @all(τzz) = (@all(τzz) + @all(τzz_o)*@inn_xy_Gr() + 2.0*@inn_xy(Gdτ)*(@d_zi(Vz)/dz))/(1.0 + @inn_xy(Gdτ)/@inn_xy(Mus) + @inn_xy_Gr())
    @all(τxy) = (@all(τxy) + @all(τxy_o)*@av_xyi_Gr() + 2.0*@av_xyi(Gdτ)*(0.5*(@d_yi(Vx)/dy + @d_xi(Vy)/dx))) /(1.0 + @av_xyi(Gdτ)/@av_xyi(Mus) + @av_xyi_Gr())
    @all(τxz) = (@all(τxz) + @all(τxz_o)*@av_xzi_Gr() + 2.0*@av_xzi(Gdτ)*(0.5*(@d_zi(Vx)/dz + @d_xi(Vz)/dx))) /(1.0 + @av_xzi(Gdτ)/@av_xzi(Mus) + @av_xzi_Gr())
    @all(τyz) = (@all(τyz) + @all(τyz_o)*@av_yzi_Gr() + 2.0*@av_yzi(Gdτ)*(0.5*(@d_zi(Vy)/dz + @d_yi(Vz)/dy))) /(1.0 + @av_yzi(Gdτ)/@av_yzi(Mus) + @av_yzi_Gr())
    return
end

@parallel function compute_dV!(dVx::Data.Array, dVy::Data.Array, dVz::Data.Array, Pt::Data.Array, τxx::Data.Array, τyy::Data.Array, τzz::Data.Array, τxy::Data.Array, τxz::Data.Array, τyz::Data.Array, dτ_Rho::Data.Array, dx::Data.Number, dy::Data.Number, dz::Data.Number)
    @all(dVx) = (@d_xa(τxx)/dx + @d_ya(τxy)/dy + @d_za(τxz)/dz - @d_xi(Pt)/dx)*@av_xi(dτ_Rho)
    @all(dVy) = (@d_ya(τyy)/dy + @d_xa(τxy)/dx + @d_za(τyz)/dz - @d_yi(Pt)/dy)*@av_yi(dτ_Rho)
    @all(dVz) = (@d_za(τzz)/dz + @d_xa(τxz)/dx + @d_ya(τyz)/dy - @d_zi(Pt)/dz)*@av_zi(dτ_Rho)
    return
end

@parallel function compute_V!(Vx::Data.Array, Vy::Data.Array, Vz::Data.Array, dVx::Data.Array, dVy::Data.Array, dVz::Data.Array)
    @inn(Vx) = @inn(Vx) + @all(dVx)
    @inn(Vy) = @inn(Vy) + @all(dVy)
    @inn(Vz) = @inn(Vz) + @all(dVz)
    return
end

@parallel function compute_Res!(Rx::Data.Array, Ry::Data.Array, Rz::Data.Array, Pt::Data.Array, τxx::Data.Array, τyy::Data.Array, τzz::Data.Array, τxy::Data.Array, τxz::Data.Array, τyz::Data.Array, dx::Data.Number, dy::Data.Number, dz::Data.Number)
    @all(Rx)    = @d_xa(τxx)/dx + @d_ya(τxy)/dy + @d_za(τxz)/dz - @d_xi(Pt)/dx
    @all(Ry)    = @d_ya(τyy)/dy + @d_xa(τxy)/dx + @d_za(τyz)/dz - @d_yi(Pt)/dy
    @all(Rz)    = @d_za(τzz)/dz + @d_xa(τxz)/dx + @d_ya(τyz)/dy - @d_zi(Pt)/dz
    return
end

@parallel_indices (iy,iz) function bc_x!(A::Data.Array)
    A[  1, iy,  iz] = A[    2,   iy,   iz]
    A[end, iy,  iz] = A[end-1,   iy,   iz]
    return
end

@parallel_indices (ix,iz) function bc_y!(A::Data.Array)
    A[ ix,  1,  iz] = A[   ix,    2,   iz]
    A[ ix,end,  iz] = A[   ix,end-1,   iz]
    return
end

@parallel_indices (ix,iy) function bc_z!(A::Data.Array)
    A[ ix,  iy,  1] = A[   ix,   iy,    2]
    A[ ix,  iy,end] = A[   ix,   iy,end-1]
    return
end

@views function Stokes3D_()
    # Physics
    lx, ly, lz = 10.0, 10.0, 10.0  # domain extends
    ξ          = 1.0               # Maxwell relaxation time
    μs0        = 1.0               # matrix viscosity
    μsi        = 1e-3              # inclusion viscosity
    G          = 1.0               # elastic shear modulus
    εbg        = 1.0               # background strain-rate
    dt         = μs0/(G*ξ)
    # Numerics
    nt         = 5                 # number of time steps
    iterMax    = 2e5               # maximum number of pseudo-transient iterations
    nout       = 200               # error checking frequency
    Re         = 6.0π              # Reynolds number
    r          = 1.0               # Bulk to shear elastic modulus ratio
    CFL        = 0.8/sqrt(3)       # CFL number
    ε          = 1e-8              # nonlinear absolute tolerence
    # nx, ny, nz = 127, 127, 127     # numerical grid resolution; should be a mulitple of 32-1 for optimal GPU perf
    b_width    = (8, 4, 4)         # boundary width for comm/comp overlap
    # Derived numerics
    me, dims   = init_global_grid(nx, ny, nz) # MPI initialisation
    @static if USE_GPU select_device() end    # select one GPU per MPI local rank (if >1 GPU per node)
    dx, dy, dz = lx/nx_g(), ly/ny_g(), lz/nz_g()      # cell sizes
    max_lxyz   = max(lx,ly,lz)
    Vpdτ       = min(dx,dy,dz)*CFL
    xc, yc, zc = LinRange(dx/2, lx-dx/2, nx), LinRange(dy/2, ly-dy/2, ny), LinRange(dz/2, lz-dz/2, nz)
    # Array allocations
    Pt         = @zeros(nx  ,ny  ,nz  )
    RhoG       = @zeros(nx  ,ny  ,nz  )
    dτ_Rho     = @zeros(nx  ,ny  ,nz  )
    Gdτ        = @zeros(nx  ,ny  ,nz  )
    ∇V         = @zeros(nx  ,ny  ,nz  )
    Vy         = @zeros(nx  ,ny+1,nz  )
    τxx        = @zeros(nx  ,ny-2,nz-2)
    τyy        = @zeros(nx-2,ny  ,nz-2)
    τzz        = @zeros(nx-2,ny-2,nz  )
    τxy        = @zeros(nx-1,ny-1,nz-2)
    τxz        = @zeros(nx-1,ny-2,nz-1)
    τyz        = @zeros(nx-2,ny-1,nz-1)
    τxx_o      = @zeros(nx  ,ny-2,nz-2)
    τyy_o      = @zeros(nx-2,ny  ,nz-2)
    τzz_o      = @zeros(nx-2,ny-2,nz  )
    τxy_o      = @zeros(nx-1,ny-1,nz-2)
    τxz_o      = @zeros(nx-1,ny-2,nz-1)
    τyz_o      = @zeros(nx-2,ny-1,nz-1)
    dVx        = @zeros(nx-1,ny-2,nz-2)
    dVy        = @zeros(nx-2,ny-1,nz-2)
    dVz        = @zeros(nx-2,ny-2,nz-1)
    Rx         = @zeros(nx-1,ny-2,nz-2)
    Ry         = @zeros(nx-2,ny-1,nz-2)
    Rz         = @zeros(nx-2,ny-2,nz-1)
    Mus2       = @zeros(nx  ,ny  ,nz  )
    Musτ       = @zeros(nx  ,ny  ,nz  )
    # Initial conditions
    Rad2       =  zeros(nx  ,ny  ,nz  )
    Vx         =  zeros(nx+1,ny  ,nz  )
    Vz         =  zeros(nx  ,ny  ,nz+1)
    Rad2      .= [(x_g(ix,dx,Rad2) +0.5*dx -0.5*lx)^2 + (y_g(iy,dy,Rad2) +0.5*dy -0.5*ly)^2 + (z_g(iz,dz,Rad2) +0.5*dz -0.5*lz)^2 for ix=1:size(Rad2,1), iy=1:size(Rad2,2), iz=1:size(Rad2,3)]
    Vx         = Data.Array( -εbg.*[(x_g(ix,dx,Vx) +0.5*dx -0.5*lx) for ix=1:size(Vx,1), iy=1:size(Vx,2), iz=1:size(Vx,3)] )
    Vz         = Data.Array(  εbg.*[(z_g(iz,dz,Vz) +0.5*dz -0.5*lz) for ix=1:size(Vz,1), iy=1:size(Vz,2), iz=1:size(Vz,3)] )
    Mus        = μs0*ones(nx,ny,nz)
    Mus[Rad2.<1.0] .= μsi
    Mus        = Data.Array( Mus )
    Mus2      .= Mus
    for ism=1:10#15
        @hide_communication b_width begin # communication/computation overlap
            @parallel smooth!(Mus2, Mus, 1.0)
            Mus, Mus2 = Mus2, Mus
            update_halo!(Mus)
        end
    end
    Musτ      .= Mus
    @hide_communication b_width begin # communication/computation overlap
        @parallel compute_maxloc!(Musτ, Mus)
        @parallel (1:size(Musτ,2), 1:size(Musτ,3)) bc_x!(Musτ)
        @parallel (1:size(Musτ,1), 1:size(Musτ,3)) bc_y!(Musτ)
        @parallel (1:size(Musτ,1), 1:size(Musτ,2)) bc_z!(Musτ)
        update_halo!(Musτ)
    end
    len_Rx_g   = ((nx-2-1)*dims[1]+2)*((ny-2-2)*dims[2]+2)*((nz-2-2)*dims[3]+2)
    len_Ry_g   = ((nx-2-2)*dims[1]+2)*((ny-2-1)*dims[2]+2)*((nz-2-2)*dims[3]+2)
    len_Rz_g   = ((nx-2-2)*dims[1]+2)*((ny-2-2)*dims[2]+2)*((nz-2-1)*dims[3]+2)
    len_∇V_g   = ((nx-2  )*dims[1]+2)*((ny-2  )*dims[2]+2)*((nz-2  )*dims[3]+2)
    # Preparation of visualisation
    if do_viz || do_save_viz
        if (me==0) ENV["GKSwstype"]="nul"; if do_viz !ispath("../../figures") && mkdir("../../figures") end; end
        nx_v, ny_v, nz_v = (nx-2)*dims[1], (ny-2)*dims[2], (nz-2)*dims[3]
        if (nx_v*ny_v*nz_v*sizeof(Data.Number) > 0.8*Sys.free_memory()) error("Not enough memory for visualization.") end
        Pt_v   = zeros(nx_v, ny_v, nz_v) # global array for visu
        Vz_v   = zeros(nx_v, ny_v, nz_v)
        Rz_v   = zeros(nx_v, ny_v, nz_v)
        Mus_v  = zeros(nx_v, ny_v, nz_v)
        τxz_v  = zeros(nx_v, ny_v, nz_v)
        Pt_inn = zeros(nx-2, ny-2, nz-2) # no halo local array for visu
        Vz_inn = zeros(nx-2, ny-2, nz-2)
        Rz_inn = zeros(nx-2, ny-2, nz-2)
        Mus_inn= zeros(nx-2, ny-2, nz-2)
        τxz_inn= zeros(nx-2, ny-2, nz-2)
        y_sl2, y_sl = Int(ceil((ny_g()-2)/2)), Int(ceil(ny_g()/2))
        Xi_g, Zi_g  = dx+dx/2:dx:(lx-dx-dx/2), dz+dz/2:dz:(lz-dz-dz/2) # inner points only
    end
    # Time loop
    @parallel compute_iter_params!(dτ_Rho, Gdτ, Musτ, Vpdτ, G, dt, Re, r, max_lxyz)
    t=0.0; ittot=0; evo_t=[]; evo_τzz=[]
    for it = 1:nt
        err=2*ε; iter=0; err_evo1=[]; err_evo2=[]
        @parallel assign_τ!(τxx, τyy, τzz, τxy, τxz, τyz, τxx_o, τyy_o, τzz_o, τxy_o, τxz_o, τyz_o)
        # Pseudo-transient iteration
        while err > ε && iter <= iterMax
            if (it==1 && iter==11)  tic()  end
            @parallel compute_P!(∇V, Pt, Vx, Vy, Vz, Gdτ, r, dx, dy, dz)
            @parallel compute_τ!(τxx, τyy, τzz, τxy, τxz, τyz, τxx_o, τyy_o, τzz_o, τxy_o, τxz_o, τyz_o, Vx, Vy, Vz, Mus, Gdτ, G, dt, dx, dy, dz)
            @parallel compute_dV!(dVx, dVy, dVz, Pt, τxx, τyy, τzz, τxy, τxz, τyz, dτ_Rho, dx, dy, dz)
            @hide_communication b_width begin # communication/computation overlap
                @parallel compute_V!(Vx, Vy, Vz, dVx, dVy, dVz)
                @parallel (1:size(Vy,2), 1:size(Vy,3)) bc_x!(Vy)
                @parallel (1:size(Vz,2), 1:size(Vz,3)) bc_x!(Vz)
                @parallel (1:size(Vx,1), 1:size(Vx,3)) bc_y!(Vx)
                @parallel (1:size(Vz,1), 1:size(Vz,3)) bc_y!(Vz)
                @parallel (1:size(Vx,1), 1:size(Vx,2)) bc_z!(Vx)
                @parallel (1:size(Vy,1), 1:size(Vy,2)) bc_z!(Vy)
                update_halo!(Vx, Vy, Vz)
            end
            iter += 1
            if iter % nout == 0
                @parallel compute_Res!(Rx, Ry, Rz, Pt, τxx, τyy, τzz, τxy, τxz, τyz, dx, dy, dz)
                Vmin, Vmax = minimum_g(Vx), maximum_g(Vx)
                Pmin, Pmax = minimum_g(Pt), maximum_g(Pt)
                norm_Rx    = norm_g(Rx)/(Pmax-Pmin)*lx/sqrt(len_Rx_g)
                norm_Ry    = norm_g(Ry)/(Pmax-Pmin)*lx/sqrt(len_Ry_g)
                norm_Rz    = norm_g(Rz)/(Pmax-Pmin)*lx/sqrt(len_Rz_g)
                norm_∇V    = norm_g(∇V)/(Vmax-Vmin)*lx/sqrt(len_∇V_g)
                # norm_Rx = norm_g(Rx)/len_Rx_g; norm_Ry = norm_g(Ry)/len_Ry_g; norm_Rz = norm_g(Rz)/len_Rz_g; norm_∇V = norm_g(∇V)/len_∇V_g
                err = maximum([norm_Rx, norm_Ry, norm_Rz, norm_∇V])
                if isnan(err) error("NaN") end
                push!(err_evo1,maximum([norm_Rx, norm_Ry, norm_Rz, norm_∇V])); push!(err_evo2,iter)
                if (me==0) @printf("Step = %d, iter = %d, err = %1.3e [norm_Rx=%1.3e, norm_Ry=%1.3e, norm_Rz=%1.3e, norm_∇V=%1.3e] \n", it, iter, err, norm_Rx, norm_Ry, norm_Rz, norm_∇V) end
            end
        end
        ittot += iter; t += dt
        push!(evo_t, t); push!(evo_τzz, maximum(τzz))
    end
    # Performance
    wtime    = toc()
    A_eff    = (10*2 +1)/1e9*nx*ny*nz*sizeof(Data.Number) # Effective main memory access per iteration [GB] (Lower bound of required memory access: Te has to be read and written: 2 whole-array memaccess; Ci has to be read: : 1 whole-array memaccess)
    wtime_it = wtime/(ittot-10)                           # Execution time per iteration [s]
    T_eff    = A_eff/wtime_it                             # Effective memory throughput [GB/s]
    if (me==0) @printf("Total iters = %d (%d steps), time = %1.3e sec (@ T_eff = %1.2f GB/s) \n", ittot, nt, wtime, round(T_eff, sigdigits=2)) end
    # Visualisation
    if do_viz || do_save_viz
        Pt_inn .= inn(Pt);   gather!(Pt_inn, Pt_v)
        Vz_inn .= av_zi(Vz); gather!(Vz_inn, Vz_v)
        Rz_inn .= av_za(Rz); gather!(Rz_inn, Rz_v)
        Mus_inn.= inn(Mus);  gather!(Mus_inn, Mus_v)
        τxz_inn.= av_xza(τxz); gather!(τxz_inn, τxz_v)
        if me==0 && do_viz
            p1 = heatmap(Xi_g, Zi_g, Pt_v[:,y_sl,:]', aspect_ratio=1, xlims=(Xi_g[1],Xi_g[end]), zlims=(Zi_g[1],Zi_g[end]), c=:viridis, title="Pressure")
            p2 = heatmap(Xi_g, Zi_g, Vz_v[:,y_sl,:]', aspect_ratio=1, xlims=(Xi_g[1],Xi_g[end]), zlims=(Zi_g[1],Zi_g[end]), c=:viridis, title="Vz")
            p4 = heatmap(Xi_g, Zi_g, log10.(abs.(Rz_v[:,y_sl2,:]')), aspect_ratio=1,  xlims=(Xi_g[1],Xi_g[end]), zlims=(Zi_g[1],Zi_g[end]), c=:viridis, title="log10(Rz)")
            #p5 = plot(err_evo2,err_evo1, legend=false, xlabel="# iterations", ylabel="log10(error)", linewidth=2, markershape=:circle, markersize=3, labels="max(error)", yaxis=:log10)
            p3 = plot(evo_t, evo_τzz, legend=false, xlabel="time", ylabel="max(τzz)", linewidth=0, markershape=:circle, framestyle=:box, markersize=3)
               #plot!(evo_t, 2.0.*εbg.*μs0.*(1.0.-exp.(.-evo_t.*G./μs0)), linewidth=2.0) # analytical solution
            plot(p1, p2, p4, p3)
            savefig("../../figures/Stokes_3D_ve_$(nx_g()).png")
        end
    end
    if me==0 && do_save
        !ispath("../../output") && mkdir("../../output")
        open("../../output/out_Stokes3D_ve.txt","a") do io
            println(io, "$(nx_g()) $(ny_g()) $(nz_g()) $(ittot) $(nt)")
        end
    end
    if me==0 && do_save_viz
        !ispath("../../out_visu") && mkdir("../../out_visu")
        matwrite("../../out_visu/Stokes_3D_ve.mat", Dict("Pt_3D"=> Pt_v, "Mus_3D"=> Mus_v, "Txz_3D"=> τxz_v, "Vz_3D"=> Vz_v, "dx_3D"=> dx, "dy_3D"=> dy, "dz_3D"=> dz); compress = true)
    end
    finalize_global_grid()
    return xc, yc, zc, Pt
end

if use_return
    xc, yc, zc, P = Stokes3D_();
else
    Stokes3D = begin Stokes3D_(); return; end
end


"""
    process_pokes(filepath::String)

Given a path of flipping data it preprocess the main variables such as Trial and Block. To deal with the different formats of raw data
collect over the year it checks the presence of a set of preexisting column
"""

function process_pokes(filepath::String)
    df = CSV.read(filepath, DataFrame)
    if !iscolumn(df,:Poke)
        DataFrames.rename!(df, propertynames(df)[1] => :Poke) #change poke counter name
    end
    if iscolumn(df,:delta)
        DataFrames.rename!(df, :delta => :Delta)
    end
    df[!,:Poke] = df[!,:Poke].+1
    start_time = df[1,:PokeIn]
    df[!,:PokeIn] = df[!,:PokeIn] .- start_time
    df[!,:PokeOut] = df[!,:PokeOut] .- start_time
    df[!,:PokeDur] = df[!,:PokeOut] - df[!,:PokeIn]
    if !iscolumn(df,:Wall)
        df[:,:Wall] = zeros(size(df,1))
    end
    booleans=[:Reward,:Side,:SideHigh,:Stim,:Wall]#columns to convert to Bool
    for x in booleans
        if df[1,x] isa AbstractString
                df[!,x] = parse.(Bool,df[!,x])
        elseif df[1,x] isa Real
            df[!,x] = Bool.(df[:,x])
        end
    end
    df[!,:Side] = [a ? "L" : "R" for a in df[!,:Side]]
    if iscolumn(df,:ProbVec0)
        integers=[:Protocollo,:ProbVec0,:ProbVec1,:GamVec0,:GamVec1,:Delta]; #columns to convert to Int64
        for x in integers
            df[!,x] = Int64.(df[!,x])
        end
        df[!,:Protocol] = [r.Protocollo == 0 ? string(r.ProbVec0,"/",r.GamVec0) : string(r.ProbVec1,"/",r.GamVec1) for r in eachrow(df)]
        for x in[:ProbVec0,:ProbVec1,:GamVec0,:GamVec1,:Protocollo]
            DataFrames.select!(df,DataFrames.Not(x))
        end
        if !iscolumn(df,:StimFreq)
            df[!,:StimFreq] .= 10000
        end
        df[!,:StimFreq] = [a == 10000 ? 25 : a  for a in df[!,:StimFreq]]
        df[!,:Box] .= "Box0"
    elseif iscolumn(df,:Prwd)
        df[!,:Protocol] = string.(df[!,:Prwd],'/',df[!,:Ptrs])
        df[!,:Box] = "Box".*string.(df[:,:Box])
    end
    df[!,:StimDay] .= length(findall(df[:,:Stim])) == 0 ? false : true
    df[!,:Leave] = lead(check_changes(df[:,:Side]), default = false)
    df[!,:Bout] = count_different(df.Reward .| check_changes(df.Side))
    df[!,:Streak] = count_different(check_changes(df[:,:Side]))
    df[!,:ReverseStreak] = reverse(df[:,:Streak])
    df[!,:PokeInStreak] .= 0
    df[!,:PokeHierarchy] .= 0.0
    df[!,:PokeInStreak] = Vector{Union{Float64,Missing}}(undef,size(df,1))
    df[!,:PreInterpoke] = Vector{Union{Float64,Missing}}(undef,size(df,1))
    df[!,:PostInterpoke] = Vector{Union{Float64,Missing}}(undef,size(df,1))
    combine(groupby(df,:Streak)) do dd
        dd[:,:PokeInStreak] = count_different(check_changes(dd[!,:Poke]))
        dd[:,:PreInterpoke] =  dd[!,:PokeIn] .-lag(dd[!,:PokeOut],default = missing)
        dd[:,:PostInterpoke] = lead(dd[!,:PokeIn],default = missing).- dd[!,:PokeOut]
        dd[:,:PokeHierarchy] = get_hierarchy(dd[!,:Reward])
    end
    df[!,:Block] = count_different(check_changes(df[!,:Wall]))
    df[!,:StreakInBlock] .= 0
    combine(groupby(df,:Block)) do dd
        dd[:,:StreakInBlock] = count_different(check_changes(dd[!,:Side]))
    end
    df[!,:SideHigh] = [x ? "L" : "R" for x in df[!,:SideHigh]]
    df[!,:Correct] = df[!,:Side] .== df[!,:SideHigh]
    return df[:, [:Poke,:Side,:SideHigh, :Correct,
    :Reward, :PokeHierarchy, :PokeInStreak,
    :Leave, :Bout, :Streak, :Stim, :StimFreq, :Wall, :Block, :StreakInBlock,
    :Delta, :Protocol, :Box, :StimDay,
    :PokeIn, :PokeOut, :PokeDur, :PreInterpoke, :PostInterpoke,:ReverseStreak]]
end