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


# 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


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


include("abstract.jl")

"""
Analysis data of determinant quantum Monte Carlo (DQMC) simulation
"""
@with_kw mutable struct DQMCAnalysis
    acc_rate::Float64 = 0.
    prop_local::Int = 0
    acc_local::Int = 0
    acc_rate_global::Float64 = 0.
    prop_global::Int = 0
    acc_global::Int = 0
end

"""
Parameters of determinant quantum Monte Carlo (DQMC)
"""
@with_kw struct DQMCParameters
    global_moves::Bool = false
    global_rate::Int = 5
    thermalization::Int = 100 # number of thermalization sweeps
    sweeps::Int = 100 # number of sweeps (after thermalization)

    all_checks::Bool = true # e.g. check if propagation is stable/instable
    safe_mult::Int = 10

    delta_tau::Float64 = 0.1
    beta::Float64
    slices::Int = beta / delta_tau
    @assert isinteger(beta / delta_tau) string("beta/delta_tau", "
        (= number of imaginary time slices) must be an integer but is",
        beta / delta_tau, ".")

    measure_rate::Int = 10
end

"""
Determinant quantum Monte Carlo (DQMC) simulation
"""
mutable struct DQMC{M<:Model, CB<:Checkerboard, ConfType<:Any,
        Stack<:AbstractDQMCStack} <: MonteCarloFlavor
    model::M
    conf::ConfType
    s::Stack

    p::DQMCParameters
    a::DQMCAnalysis
    thermalization_measurements::Dict{Symbol, AbstractMeasurement}
    measurements::Dict{Symbol, AbstractMeasurement}

    DQMC{M, CB, ConfType, Stack}() where {M<:Model, CB<:Checkerboard,
        ConfType<:Any, Stack<:AbstractDQMCStack} = new()
end

include("stack.jl")
include("slice_matrices.jl")

"""
    DQMC(m::M; kwargs...) where M<:Model

Create a determinant quantum Monte Carlo simulation for model `m` with
keyword parameters `kwargs`.
"""
function DQMC(m::M;
        seed::Int=-1,
        checkerboard::Bool=false,
        thermalization_measurements = Dict{Symbol, AbstractMeasurement}(),
        measurements = :default,
        kwargs...
    ) where M<:Model
    # default params
    # paramskwargs = filter(kw->kw[1] in fieldnames(DQMCParameters), kwargs)
    p = DQMCParameters(; kwargs...)

    geltype = greenseltype(DQMC, m)
    heltype = hoppingeltype(DQMC, m)
    conf = rand(DQMC, m, p.slices)
    mc = DQMC{M, checkerboard ? CheckerboardTrue : CheckerboardFalse,
        typeof(conf), DQMCStack{geltype, heltype}}()
    mc.model = m
    mc.p = p
    mc.s = DQMCStack{geltype, heltype}()

    init!(
        mc, seed = seed, conf = conf,
        thermalization_measurements = thermalization_measurements,
        measurements = measurements
    )
    return mc
end

"""
    DQMC(m::M, params::Dict)
    DQMC(m::M, params::NamedTuple)

Create a determinant quantum Monte Carlo simulation for model `m` with
(keyword) parameters as specified in the dictionary/named tuple `params`.
"""
DQMC(m::Model, params::Dict{Symbol, T}) where T = DQMC(m; params...)
DQMC(m::Model, params::NamedTuple) = DQMC(m; params...)

# convenience
@inline beta(mc::DQMC) = mc.p.beta
@inline nslices(mc::DQMC) = mc.p.slices
@inline model(mc::DQMC) = mc.model
@inline conf(mc::DQMC) = mc.conf
@inline current_slice(mc::DQMC) = mc.s.current_slice

# cosmetics
import Base.summary
import Base.show
Base.summary(mc::DQMC) = "DQMC simulation of $(summary(mc.model))"
function Base.show(io::IO, mc::DQMC)
    print(io, "Determinant quantum Monte Carlo simulation\n")
    print(io, "Model: ", mc.model, "\n")
    print(io, "Beta: ", mc.p.beta, " (T ≈ $(round(1/mc.p.beta, sigdigits=3)))\n")
    N_th_meas = length(mc.thermalization_measurements)
    N_me_meas = length(mc.measurements)
    print(io, "Measurements: ", N_th_meas + N_me_meas, " ($N_th_meas + $N_me_meas)")
end
Base.show(io::IO, m::MIME"text/plain", mc::DQMC) = print(io, mc)

"""
    init!(mc::DQMC[; seed::Real=-1])

Initialize the determinant quantum Monte Carlo simulation `mc`.
If `seed !=- 1` the random generator will be initialized with `Random.seed!(seed)`.
"""
function init!(mc::DQMC;
        seed::Real = -1,
        conf = rand(DQMC,model(mc),nslices(mc)),
        thermalization_measurements = Dict{Symbol, AbstractMeasurement}(),
        measurements = :default
    )
    seed == -1 || Random.seed!(seed)

    mc.conf = conf

    init_hopping_matrices(mc, mc.model)
    initialize_stack(mc)

    mc.a = DQMCAnalysis()

    mc.thermalization_measurements = thermalization_measurements
    if measurements isa Dict{Symbol, AbstractMeasurement}
        mc.measurements = measurements
    elseif measurements == :default
        mc.measurements = default_measurements(mc, mc.model)
    else
        @warn(
            "`measurements` should be of type Dict{Symbol, AbstractMeasurement}, but is " *
            "$(typeof(measurements)). No measurements have been set."
        )
        mc.measurements = Dict{Symbol, AbstractMeasurement}()
    end

    nothing
end

"""
    run!(mc::DQMC[; verbose::Bool=true, sweeps::Int, thermalization::Int])

Runs the given Monte Carlo simulation `mc`.
Progress will be printed to `stdout` if `verbose=true` (default).
"""
function run!(mc::DQMC; verbose::Bool=true, sweeps::Int=mc.p.sweeps,
        thermalization=mc.p.thermalization)

    do_th_measurements = !isempty(mc.thermalization_measurements)
    do_me_measurements = !isempty(mc.measurements)
    !do_me_measurements && @warn(
        "There are no measurements set up for this simulation!"
    )
    total_sweeps = sweeps + thermalization

    start_time = now()
    verbose && println("Started: ", Dates.format(start_time, "d.u yyyy HH:MM"))

    # fresh stack
    verbose && println("Preparing Green's function stack")
    initialize_stack(mc) # redundant ?!
    build_stack(mc)
    propagate(mc)

    _time = time()
    verbose && println("\n\nThermalization stage - ", thermalization)
    do_th_measurements && prepare!(mc.thermalization_measurements, mc, mc.model)
    for i in 1:total_sweeps
        verbose && (i == thermalization + 1) &&println("\n\nMeasurement stage - ", sweeps)
        for u in 1:2 * nslices(mc)
            update(mc, i)

            # For optimal performance whatever is most likely to fail should be
            # checked first.
            if current_slice(mc) == nslices(mc) && i <= thermalization && mc.s.direction == -1 &&
                    iszero(mod(i, mc.p.measure_rate)) && do_th_measurements
                measure!(mc.thermalization_measurements, mc, mc.model, i)
            end
            if (i == thermalization+1)
                do_th_measurements && finish!(mc.thermalization_measurements, mc, mc.model)
                do_me_measurements && prepare!(mc.measurements, mc, mc.model)
            end
            if current_slice(mc) == nslices(mc) && mc.s.direction == -1 && i > thermalization &&
                    iszero(mod(i, mc.p.measure_rate)) && do_me_measurements
                measure!(mc.measurements, mc, mc.model, i)
            end

        end

        if mod(i, 10) == 0
            mc.a.acc_rate = mc.a.acc_rate / (10 * 2 * nslices(mc))
            mc.a.acc_rate_global = mc.a.acc_rate_global / (10 / mc.p.global_rate)
            sweep_dur = (time() - _time)/10
            if verbose
                println("\t", i)
                @printf("\t\tsweep dur: %.3fs\n", sweep_dur)
                @printf("\t\tacc rate (local) : %.1f%%\n", mc.a.acc_rate*100)
                if mc.p.global_moves
                  @printf("\t\tacc rate (global): %.1f%%\n", mc.a.acc_rate_global*100)
                  @printf("\t\tacc rate (global, overall): %.1f%%\n",
                    mc.a.acc_global/mc.a.prop_global*100)
                end
            end

            mc.a.acc_rate = 0.0
            mc.a.acc_rate_global = 0.0
            flush(stdout)
            _time = time()
        end
    end
    do_me_measurements && finish!(mc.measurements, mc, mc.model)

    mc.a.acc_rate = mc.a.acc_local / mc.a.prop_local
    mc.a.acc_rate_global = mc.a.acc_global / mc.a.prop_global

    end_time = now()
    verbose && println("Ended: ", Dates.format(end_time, "d.u yyyy HH:MM"))
    verbose && @printf("Duration: %.2f minutes", (end_time - start_time).value/1000. /60.)

    nothing
end

"""
    update(mc::DQMC, i::Int)

Propagates the Green's function and performs local and global updates at
current imaginary time slice.
"""
function update(mc::DQMC, i::Int)
    propagate(mc)

    # global move
    # if mc.p.global_moves && (current_slice(mc) == mc.p.slices &&
    #        mc.s.direction == -1 && iszero(mod(i, mc.p.global_rate)))
    #     mc.a.prop_global += 1
    #     b = global_move(mc, mc.model, mc.conf) # not yet in DQMC_optional, i.e. unsupported
    #     mc.a.acc_global += b
    # end

    # local moves
    sweep_spatial(mc)

    nothing
end

"""
    sweep_spatial(mc::DQMC)

Performs a sweep of local moves along spatial dimension at current
imaginary time slice.
"""
function sweep_spatial(mc::DQMC)
    m = model(mc)
    N = nsites(m)

    @inbounds for i in 1:N
        detratio, ΔE_boson, Δ = propose_local(mc, m, i, current_slice(mc), conf(mc))
        mc.a.prop_local += 1

        if abs(imag(detratio)) > 1e-6
            println("Did you expect a sign problem? imag. detratio: ",
                abs(imag(detratio)))
            @printf "%.10e" abs(imag(detratio))
        end
        p = real(exp(- ΔE_boson) * detratio)

        # Metropolis
        if p > 1 || rand() < p
            accept_local!(mc, m, i, current_slice(mc), conf(mc), Δ, detratio,
                ΔE_boson)
            mc.a.acc_rate += 1/N
            mc.a.acc_local += 1
        end
    end
    nothing
end

"""
    greens(mc::DQMC)

Obtain the current equal-time Green's function.

Internally, `mc.s.greens` is an effective Green's function. This method transforms
this effective one to the actual Green's function by multiplying hopping matrix
exponentials from left and right.
"""
function greens(mc::DQMC_CBFalse)
    eThalfminus = mc.s.hopping_matrix_exp
    eThalfplus = mc.s.hopping_matrix_exp_inv

    greens = copy(mc.s.greens)
    greens .= greens * eThalfminus
    greens .= eThalfplus * greens
    return greens
end
function greens(mc::DQMC_CBTrue)
    chkr_hop_half_minus = mc.s.chkr_hop_half
    chkr_hop_half_plus = mc.s.chkr_hop_half_inv

    greens = copy(mc.s.greens)

    @inbounds @views begin
        for i in reverse(1:mc.s.n_groups)
          greens .= greens * chkr_hop_half_minus[i]
        end
        for i in reverse(1:mc.s.n_groups)
          greens .= chkr_hop_half_plus[i] * greens
        end
    end
    return greens
end

include("DQMC_mandatory.jl")
include("DQMC_optional.jl")
include("measurements.jl")


#Period types
value{P<:Period}(x::P) = x.value

# The default constructors for Periods work well in almost all cases
# P(x) = new((convert(Int64,x))
# The following definitions are for Period-specific safety
for p in (:Year,:Month,:Week,:Day,:Hour,:Minute,:Second,:Millisecond)
    # This ensures that we can't convert between Periods
    @eval $p(x::Period) = throw(ArgumentError("Can't convert $(typeof(x)) to $($p)"))
    # Unless the Periods are the same type
    @eval $p(x::$p) = x
    # Convenience method for show()
    @eval _units(x::$p) = " " * lowercase(string($p)) * (abs(value(x)) == 1 ? "" : "s")
    # periodisless
    @eval periodisless(x::$p,y::$p) = value(x) < value(y)
    # String parsing (mainly for IO code)
    @eval $p(x::String) = $p(parseint(Int64,x))
    # Period accessors
    @eval $p(x::TimeType) = $p($(symbol(lowercase(string(p))))(x))
end
# Now we're safe to define Period-Number conversions
# Anything an Int64 can convert to, a Period can convert to
Base.convert{T<:Number}(::Type{T},x::Period) = convert(T,value(x))
Base.convert{T<:Period}(::Type{T},x::Real) = T(int64(x))

#Print/show/traits
Base.string{P<:Period}(x::P) = string(value(x),_units(x))
Base.show(io::IO,x::Period) = print(io,string(x))
Base.zero{P<:Period}(::Union(Type{P},P)) = P(0)
Base.one{P<:Period}(::Union(Type{P},P)) = P(1)
Base.typemin{P<:Period}(::Type{P}) = P(typemin(Int64))
Base.typemax{P<:Period}(::Type{P}) = P(typemax(Int64))

# Default values (as used by TimeTypes)
default{T<:DatePeriod}(p::Union(T,Type{T})) = one(p)
default{T<:TimePeriod}(p::Union(T,Type{T})) = zero(p)

(-){P<:Period}(x::P) = P(-value(x))
Base.isless{P<:Period}(x::P,y::P) = isless(value(x),value(y))
Base.isless{R<:Real}(x::Period,y::R) = isless(int64(y),value(x))
Base.isless{R<:Real}(y::R,x::Period) = isless(int64(y),value(x))
=={P<:Period}(x::P,y::P) = ===(value(x),value(y))
=={R<:Real}(x::Period,y::R) = ==(int64(y),value(x))
=={R<:Real}(y::R,x::Period) = ==(int64(y),value(x))

Base.isless(x::Period,y::Period) = throw(ArgumentError("Can't compare $(typeof(x)) and $(typeof(y))"))
==(x::Period,y::Period) = throw(ArgumentError("Can't compare $(typeof(x)) and $(typeof(y))"))

#Period Arithmetic:
import Base.div, Base.mod, Base.gcd, Base.lcm
let vec_ops = [:.+,:.-,:.*,:.%,:div]
    for op in [:+,:-,:*,:%,:mod,:gcd,:lcm,vec_ops]
        @eval begin
        #Period-Period
        ($op){P<:Period}(x::P,y::P) = P(($op)(value(x),value(y)))
        #Period-Real
        ($op){P<:Period}(x::P,y::Real) = P(($op)(value(x),convert(Int64,y)))
        ($op){P<:Period}(x::Real,y::P) = P(($op)(convert(Int64,x),value(y)))
        end
        #Vectorized
        if op in vec_ops
            @eval begin
            ($op){P<:Period}(x::AbstractArray{P}, y::P) = reshape([$op(i,y) for i in x], size(x))
            ($op){P<:Period}(y::P, x::AbstractArray{P}) = reshape([$op(i,y) for i in x], size(x))
            end
        end
    end
end

# intfuncs
Base.gcdx{T<:Period}(a::T,b::T) = ((g,x,y)=gcdx(value(a),value(b)); return T(g),x,y)
Base.abs{T<:Period}(a::T) = T(abs(value(a)))

periodisless(::Period,::Year)        = true
periodisless(::Period,::Month)       = true
periodisless(::Year,::Month)         = false
periodisless(::Period,::Week)        = true
periodisless(::Year,::Week)          = false
periodisless(::Month,::Week)         = false
periodisless(::Period,::Day)         = true
periodisless(::Year,::Day)           = false
periodisless(::Month,::Day)          = false
periodisless(::Week,::Day)           = false
periodisless(::Period,::Hour)        = false
periodisless(::Minute,::Hour)        = true
periodisless(::Second,::Hour)        = true
periodisless(::Millisecond,::Hour)   = true
periodisless(::Period,::Minute)      = false
periodisless(::Second,::Minute)      = true
periodisless(::Millisecond,::Minute) = true
periodisless(::Period,::Second)      = false
periodisless(::Millisecond,::Second) = true
periodisless(::Period,::Millisecond) = false

# Stores multiple periods in greatest to least order by type, not values
type CompoundPeriod
    periods::Array{Period,1}
end
function Base.string(x::CompoundPeriod)
    s = ""
    for p in x.periods
        s *= ", " * string(p)
    end
    return s[3:end]
end
Base.show(io::IO,x::CompoundPeriod) = print(io,string(x))
# E.g. Year(1) + Day(1)
(+)(x::Period,y::Period) = CompoundPeriod(sort!(Period[x,y],rev=true,lt=periodisless))
(+)(x::CompoundPeriod,y::Period) = (sort!(push!(x.periods,y) ,rev=true,lt=periodisless); return x)
# E.g. Year(1) - Month(1)
(-)(x::Period,y::Period) = CompoundPeriod(sort!(Period[x,-y],rev=true,lt=periodisless))
(-)(x::CompoundPeriod,y::Period) = (sort!(push!(x.periods,-y),rev=true,lt=periodisless); return x)

# Capture TimeType+-Period methods
(+)(a::TimeType,b::Period,c::Period) = (+)(a,b+c)
(-)(a::TimeType,b::Period,c::Period) = (-)(a,b-c)
(+)(a::TimeType,b::Period,c::Period,d::Period...) = (+)((+)(a,b+c),d...)
(-)(a::TimeType,b::Period,c::Period,d::Period...) = (-)((-)(a,b-c),d...)

function (+)(x::TimeType,y::CompoundPeriod)
    for p in y.periods
        x += p
    end
    return x
end
(+)(x::CompoundPeriod,y::TimeType) = y + x

# Convert fixed value Periods to # of milliseconds
typealias FixedPeriod Union(Week,Day,Hour,Minute,Second,Millisecond)

toms(c::Millisecond) = value(c)
toms(c::Second)      = 1000*value(c)
toms(c::Minute)      = 60000*value(c)
toms(c::Hour)        = 3600000*value(c)
toms(c::Day)         = 86400000*value(c)
toms(c::Week)        = 604800000*value(c)
toms(c::Month)       = 86400000.0*30.436875*value(c)
toms(c::Year)        = 86400000.0*365.2425*value(c)

Millisecond{T<:FixedPeriod}(c::T) = Millisecond(toms(c))

days(c::Millisecond) = div(value(c),86400000)
days(c::Second)      = div(value(c),86400)
days(c::Minute)      = div(value(c),1440)
days(c::Hour)        = div(value(c),24)
days(c::Day)         = value(c)
days(c::Week)        = 7*value(c)
days(c::Year)        = 365.2425*value(c)
days(c::Month)       = 30.436875*value(c)

Day{T<:FixedPeriod}(c::T) = Day(days(c))


##  Copyright 2022 Myrl Marmarelis
##
##  Licensed under the Apache License, Version 2.0 (the "License");
##  you may not use this file except in compliance with the License.
##  You may obtain a copy of the License at
##
##    http://www.apache.org/licenses/LICENSE-2.0
##
##  Unless required by applicable law or agreed to in writing, software
##  distributed under the License is distributed on an "AS IS" BASIS,
##  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
##  See the License for the specific language governing permissions and
##  limitations under the License.

using ForwardDiff

abstract type Constraint{ T <: Real } end # do I really need the `T` for the supertype?

draw_instance(::Constraint{T}, ::Predictor{In,POut,T}, ::PredictionParam{T},
    ::Vector{Instance{In,Out,T}}) where {In,POut,Out,T} = #::ConstraintInstance{PIn,Out,T}
  error("unimplemented")

apply_operator(::Constraint{T}, ::Predictor{In,POut,T}, ::PredictionParam{T},
    ::ConstraintInstance{In,COut,T}) where {In,POut,COut,T} =
  error("unimplemented")

constraint_instance_type(::Constraint{T}, ::PredictionParam{T}, ::Predictor{In,POut,T}
    ) where {In,POut,T} =
  error("unimplemented")

struct ConservationConstraint{ T } <: Constraint{T}
  input_indices :: Matrix{Int} # batched as columns so that gradient only has to be taken once
  coefficients :: Matrix{T}
  positive :: Vector{Bool} # per column
  absolute :: Vector{Bool}
  mimicked_indices :: Vector{Int} # indices to grab from a random instance; recycled
end

constraint_target(::LomaxParam, (alpha, beta)::AbstractVector) =
  (alpha - 1) / beta # mode of the rate parameter

@generated function constraint_instance_type(::ConservationConstraint{T},
    ::Predictor{In,POut,T}, ::PredictionParam{T}) where {In,POut,T} # make this a macro? (harder to know types of arguments)
  ConstraintInstance{In,0,T} # do I need to quote this even though it's just a type value?
end

function draw_instance(constraint::ConservationConstraint{T}, predictor::Predictor{In,POut,T},
    parametrization::PredictionParam{T}, instances::Vector{Instance{In,Out,T}}
    )::ConstraintInstance{In,0,T} where {In,POut,Out,T}
  grabbed_instance = sample(instances)
  input = @MVector rand(T, In)
  for i in constraint.mimicked_indices # just overwrite previously randomized
    input[i] = grabbed_instance.input[i]
  end
  # I guess I could always subclass the ConstraintInstance?
  # generic structure seems to work fine for now.
  ConstraintInstance(
    convert(SVector, input),
    SVector{0,T}() )
end

function apply_operator(constraint::ConservationConstraint{T}, predictor::Predictor{In,POut,T},
    parametrization::PredictionParam{T}, instance::ConstraintInstance{In,0,T})::T where {In,POut,T}
  n_indices, n_laws = size(constraint.input_indices)
  @assert (n_indices, n_laws) == size(constraint.coefficients)
  n_targets = n_instance_outputs(predictor, parametrization)
  sum(1:n_targets) do target_index
    # recent versions of Zygote can ALMOST differentiate their own gradients
    derivatives = ForwardDiff.gradient(instance.input) do input # input -> begin ... end
      map_prediction(params -> constraint_target(parametrization, params),
        predictor, parametrization, input, target_index)
    end
    #derivatives, = grad(central_fdm(diff_order, 1), differentiable, instance.input)
    sum(1:n_laws) do law
      operator = sum(1:n_indices) do index
        input = constraint.input_indices[index, law]
        coefficient = constraint.coefficients[index, law]
        metric = derivatives[input]
        if constraint.absolute[law]
          metric = abs(metric)
        end
        coefficient * metric
      end
      if constraint.positive[law]
        min(operator, 0) ^ 2
      else
        operator ^ 2 # penalize nonzero-ness
      end
    end
  end
end


function tvdrk3!(system::LegSystem,times::CVTimer,n::Int64,splits::Vector{Int64},terms::Vector{Int64},uprox::Float64)
    # allocation for invariant/boundary value RK3 loop
    dW = zeros(2);
    W10 = zeros(2,length(system.branches.ID));
    W20 = zeros(2,length(system.branches.ID));
    W11 = zeros(2,length(system.branches.ID));
    W21 = zeros(2,length(system.branches.ID));
    W12 = zeros(2,length(system.branches.ID));
    W22 = zeros(2,length(system.branches.ID));
    lf = zeros(length(system.branches.ID));
    lb = zeros(length(system.branches.ID));
    ret1 = zeros(2);
    ret2 = zeros(2);
    ret3 = zeros(1);
    ret4 = zeros(1);

    # allocation for interior value RK3 loop
    dU = zeros((length(system.solverparams.colsint))*2);
    Q1 = zeros(length(system.solverparams.colsint),length(system.branches.ID));
    A1 = zeros(length(system.solverparams.colsint),length(system.branches.ID));
    Q2 = zeros(length(system.solverparams.colsint),length(system.branches.ID));
    A2 = zeros(length(system.solverparams.colsint),length(system.branches.ID));
    Acat = zeros(length(system.solverparams.colsint)+2);
    Qcat = zeros(length(system.solverparams.colsint)+2);

    # allocation for WENO3 spatial derivative subroutine
    ISA = zeros(3);
    ISQ = zeros(3);
    PA = Vector{Function}(length(system.solverparams.colsint));
    PQ = Vector{Function}(length(system.solverparams.colsint));
    Aint = Vector{Float64}(length(system.solverparams.colsint));
    Qint = Vector{Float64}(length(system.solverparams.colsint));
    wA = Vector{Float64}(3);
    wQ = Vector{Float64}(3);
    c = Vector{Float64}(3);
    a = Vector{Float64}(2);
    p = Vector{Float64}(2);
    x = Vector{Float64}(2);
    Al = Vector{Float64}(2);
    Ar = Vector{Float64}(2);
    Ql = Vector{Float64}(2);
    Qr = Vector{Float64}(2);
    Fl = Vector{Float64}(2);
    Fr = Vector{Float64}(2);
    Hl = Vector{Float64}(2);
    Hr = Vector{Float64}(2);
    Pi = Vector{Float64}(3);
    σ = Vector{Float64}(2);
    eigs = Vector{Float64}(2);
    abso = Vector{Float64}(2);
    J = zeros(2,2);
    Pret = zeros(1);
    s = zeros(1);

    # first RK step
    # println("First RK:")
    tic();
    for i = 1:length(system.branches.ID)
        # println("First step, artery $i, time step $n.")
        # scalings
        # println("Scalings:")
        As = ((system.solverparams.rho.*system.branches.c0[i][end].^2)./system.branches.beta[i][end]).^2;
        Qs = As.*system.branches.c0[i][end];
        xs = system.branches.lengthinmm[i].*LegHemodynamics.mmTom .- system.branches.k[i];
        ts = As.*xs./Qs;
        Acat .= system.branches.A[i][:,n+1]/As;
        Qcat .= system.branches.Q[i][:,n+1]/Qs;
        # println("Initial area: $(system.branches.A[i][:,n+1]./As)")
        # println("Initial flowrate: $(system.branches.Q[i][:,n+1]./Qs)")
        # invariant step
        # println("Invariant calculation:")
        invariants!(ret1,ret2,ret3,ret4,Acat,Qcat,
            system.branches.c0[i][end],system.branches.beta[i][end],system.solverparams.rho,
            system.solverparams.mu,system.solverparams.diffusioncoeff,Qs,As)
        W10[:,i] .= ret1;
        W20[:,i] .= ret2;
        lf[i] = ret3[1];
        lb[i] = ret4[1];
        # println("Forward invariant at time step $n, artery $i: $(W1n[i]) m/s")
        # println("Invariant update:")
        invariantodes!(dW,W10[:,i],W20[:,i],lf[i],lb[i],Qcat,Acat,
            system.solverparams.rho,system.solverparams.mu,system.solverparams.diffusioncoeff,
            system.branches.k[i]./xs,ts,xs,Qs,As);
        # println("Invariant assignment:")
        if i == 1 # aortic root
            system.branches.W1[i] = (W10[end,i] .+ system.solverparams.h./ts.*dW[1]).*Qs./As;
            system.branches.W2root = (W20[1,i] .+ system.solverparams.h./ts.*dW[2]).*Qs./As;
        elseif isempty(system.branches.children[i]) # terminals
            system.branches.W1end[i] = (W10[end,i] .+ system.solverparams.h./ts.*dW[1]).*Qs./As;
            system.branches.W2[i] = (W20[1,i] .+ system.solverparams.h./ts.*dW[2]).*Qs./As;
        else # all others
            system.branches.W1[i] = (W10[end,i] .+ system.solverparams.h./ts.*dW[1]).*Qs./As;
            system.branches.W2[i] = (W20[1,i] .+ system.solverparams.h./ts.*dW[2]).*Qs./As;
        end
        # primitive step
        # println("WENO3:")
        weno3!(dU,system.branches.k[i]./xs,Acat,
            Qcat,system.branches.beta[i][end],
            system.solverparams.rho,system.solverparams.mu,
            system.solverparams.diffusioncoeff,ts,xs,Qs,As,ISA,ISQ,PA,PQ,wA,wQ,
            c,a,p,x,Al,Ar,Ql,Qr,Hl,Hr,Fl,Fr,Pi,σ,eigs,abso,J,Aint,Qint,Pret,s)
        # println("Intermediate step:")
        A1[:,i] .= (system.branches.A[i][2:end-1,n+1]./As .+
            system.solverparams.h./ts.*dU[system.solverparams.acols[1:end-2]]).*As;
        Q1[:,i] .= (system.branches.Q[i][2:end-1,n+1]./Qs .+
            system.solverparams.h./ts.*dU[system.solverparams.qcols[1:end.-2].-2]).*Qs;
    end
    times.tfd += toq();

    # first set of boundary values
    # println("Split solve:")
    tic();
    for i = 1:length(splits)
        # println("Solving interior boundary, parent ID: $(splits[i])")
        children = system.branches.children[splits[i]];
        Q = zeros(length(children)+1);
        A = zeros(length(children)+1);
        Q[1] = system.branches.Q[splits[i]][system.solverparams.JL,n+1];
        A[1] = system.branches.A[splits[i]][system.solverparams.JL,n+1];
        for j = 1:length(children)
            Q[j+1] = system.branches.Q[children[j]][1,n+1];
            A[j+1] = system.branches.A[children[j]][1,n+1];
        end
        LegHemodynamics.interiorbcs!(system,n,splits[i],Q,A,children);
    end
    times.ts += toq();
    # println("0D-1D coupling:")
    tic();
    LegHemodynamics.coupling!(system,n,system.solverparams.h,terms,uprox);
    times.tc += toq();

    # second RK step
    # println("Second RK:")
    tic()
    for i = 1:length(system.branches.ID)
        # println("Second step, artery $i, time step $n.")
        Acat[1] = system.branches.A[i][1,n+2];
        Acat[2:end-1] .= A1[:,i];
        Acat[end] = system.branches.A[i][end,n+2];
        Qcat[1] = system.branches.Q[i][1,n+2];
        Qcat[2:end-1] .= Q1[:,i];
        Qcat[end] = system.branches.Q[i][end,n+2];
        # scalings
        As = ((system.solverparams.rho.*system.branches.c0[i][end].^2)./system.branches.beta[i][end]).^2;
        Qs = As.*system.branches.c0[i][end];
        xs = system.branches.lengthinmm[i].*LegHemodynamics.mmTom .- system.branches.k[i];
        ts = As.*xs./Qs;
        Acat /= As;
        Qcat /= Qs;
        # println("Area after first step: $(Acat./As)")
        # println("Flowrate after first step: $(Qcat./Qs)")
        # invariant step
        invariants!(ret1,ret2,ret3,ret4,Acat,Qcat,
            system.branches.c0[i][end],system.branches.beta[i][end],system.solverparams.rho,
            system.solverparams.mu,system.solverparams.diffusioncoeff,Qs,As)
        W11[:,i] .= ret1;
        W21[:,i] .= ret2;
        lf[i] = ret3[1];
        lb[i] = ret4[1];
        invariantodes!(dW,W11[:,i],W21[:,i],lf[i],lb[i],Qcat,Acat,
            system.solverparams.rho,system.solverparams.mu,system.solverparams.diffusioncoeff,
            system.branches.k[i]./xs,ts,xs,Qs,As);
        if i == 1 # aortic root
            system.branches.W1[i] = (0.75.*W10[end,i] .+ 0.25.*W11[end,i] .+
                0.25.*system.solverparams.h./ts.*dW[1]).*Qs./As;
            system.branches.W2root = (0.75.*W20[1,i] .+ 0.25.*W21[1,i] .+
                0.25.*system.solverparams.h./ts.*dW[2]).*Qs./As;
        elseif isempty(system.branches.children[i]) # terminals
            system.branches.W1end[i] = (0.75.*W10[end,i] .+ 0.25.*W11[end,i] .+
                0.25.*system.solverparams.h./ts.*dW[1]).*Qs./As;
            system.branches.W2[i] = (0.75.*W20[1,i] .+ 0.25.*W21[1,i] .+
                0.25.*system.solverparams.h./ts.*dW[2]).*Qs./As;
        else # all others
            system.branches.W1[i] = (0.75.*W10[end,i] .+ 0.25.*W11[end,i] .+
                0.25.*system.solverparams.h./ts.*dW[1]).*Qs./As;
            system.branches.W2[i] = (0.75.*W20[1,i] .+ 0.25.*W21[1,i] .+
                0.25.*system.solverparams.h./ts.*dW[2]).*Qs./As;
        end
        # primitive step
        weno3!(dU,system.branches.k[i]./xs,Acat,
            Qcat,system.branches.beta[i][end],
            system.solverparams.rho,system.solverparams.mu,
            system.solverparams.diffusioncoeff,ts,xs,Qs,As,ISA,ISQ,PA,PQ,wA,wQ,
            c,a,p,x,Al,Ar,Ql,Qr,Hl,Hr,Fl,Fr,Pi,σ,eigs,abso,J,Aint,Qint,Pret,s)
        A2[:,i] .= (0.75.*system.branches.A[i][2:end-1,n+1]./As .+ 0.25.*A1[:,i]./As .+
            0.25.*system.solverparams.h./ts.*dU[system.solverparams.acols[1:end-2]]).*As;
        Q2[:,i] .= (0.75.*system.branches.Q[i][2:end-1,n+1]./Qs .+ 0.25.*Q1[:,i]./Qs .+
            0.25.*system.solverparams.h./ts.*dU[system.solverparams.qcols[1:end.-2].-2]).*Qs;
    end
    times.tfd += toq();

    # second set of boundary values
    tic();
    for i = 1:length(splits)
        children = system.branches.children[splits[i]];
        Q = zeros(length(children)+1);
        A = zeros(length(children)+1);
        Q[1] = system.branches.Q[splits[i]][system.solverparams.JL,n+2];
        A[1] = system.branches.A[splits[i]][system.solverparams.JL,n+2];
        for j = 1:length(children)
            Q[j+1] = system.branches.Q[children[j]][1,n+2];
            A[j+1] = system.branches.A[children[j]][1,n+2];
        end
        LegHemodynamics.interiorbcs!(system,n,splits[i],Q,A,children);
    end
    times.ts += toq();
    tic()
    LegHemodynamics.coupling!(system,n,0.5*system.solverparams.h,terms,uprox);
    times.tc += toq();

    # third RK step
    # println("Third RK:")
    tic();
    for i = 1:length(system.branches.ID)
        # println("Third step, artery $i, time step $n.")
        Acat[1] = system.branches.A[i][1,n+2];
        Acat[2:end-1] .= A2[:,i];
        Acat[end] = system.branches.A[i][end,n+2];
        Qcat[1] = system.branches.Q[i][1,n+2];
        Qcat[2:end-1] .= Q2[:,i];
        Qcat[end] = system.branches.Q[i][end,n+2];
        # scalings
        As = ((system.solverparams.rho.*system.branches.c0[i][end].^2)./system.branches.beta[i][end]).^2;
        Qs = As.*system.branches.c0[i][end];
        xs = system.branches.lengthinmm[i].*LegHemodynamics.mmTom .- system.branches.k[i];
        ts = As.*xs./Qs;
        Acat /= As;
        Qcat /= Qs;
        # println("Area after second step: $(Acat./As)")
        # println("Flowrate after second step: $(Qcat./Qs)")
        # invariant step
        invariants!(ret1,ret2,ret3,ret4,Acat,Qcat,
            system.branches.c0[i][end],system.branches.beta[i][end],system.solverparams.rho,
            system.solverparams.mu,system.solverparams.diffusioncoeff,Qs,As)
        W12[:,i] .= ret1;
        W22[:,i] .= ret2;
        lf[i] = ret3[1];
        lb[i] = ret4[1];
        invariantodes!(dW,W12[:,i],W22[:,i],lf[i],lb[i],Qcat,Acat,
            system.solverparams.rho,system.solverparams.mu,system.solverparams.diffusioncoeff,
            system.branches.k[i]./xs,ts,xs,Qs,As);
        if i == 1 # aortic root
            system.branches.W1[i] .= (1/3.*W10[end,i] .+ 2/3.*W12[end,i] .+
                2/3.*system.solverparams.h./ts.*dW[1]).*Qs./As;
            system.branches.W2root = (1/3.*W20[1,i] .+ 2/3.*W22[1,i] .+
                2/3.*system.solverparams.h./ts.*dW[2]).*Qs./As;
        elseif isempty(system.branches.children[i]) # terminals
            system.branches.W1end[i] .= (1/3.*W10[end,i] .+ 2/3.*W12[end,i] .+
                2/3.*system.solverparams.h./ts.*dW[1]).*Qs./As;
            system.branches.W2[i] .= (1/3.*W20[1,i] .+ 2/3.*W22[1,i] .+
                2/3.*system.solverparams.h./ts.*dW[2]).*Qs./As;
        else # all others
            system.branches.W1[i] .= (1/3.*W10[end,i] .+ 2/3.*W12[end,i] .+
                2/3.*system.solverparams.h./ts.*dW[1]).*Qs./As;
            system.branches.W2[i] .= (1/3.*W20[1,i] .+ 2/3.*W22[1,i] .+
                2/3.*system.solverparams.h./ts.*dW[2]).*Qs./As;
        end
        # primitive step
        weno3!(dU,system.branches.k[i]./xs,Acat,
            Qcat,system.branches.beta[i][end],
            system.solverparams.rho,system.solverparams.mu,
            system.solverparams.diffusioncoeff,ts,xs,Qs,As,ISA,ISQ,PA,PQ,wA,wQ,
            c,a,p,x,Al,Ar,Ql,Qr,Hl,Hr,Fl,Fr,Pi,σ,eigs,abso,J,Aint,Qint,Pret,s)
        system.branches.A[i][2:end-1,n+2] .= (1/3.*system.branches.A[i][2:end-1,n+1]./As .+ 2/3.*A2[:,i]./As .+
            2/3.*system.solverparams.h./ts.*dU[system.solverparams.acols[1:end-2]]).*As;
        system.branches.Q[i][2:end-1,n+2] .= (1/3.*system.branches.Q[i][2:end-1,n+1]./Qs .+ 2/3.*Q2[:,i]./Qs .+
            2/3.*system.solverparams.h./ts.*dU[system.solverparams.qcols[1:end.-2].-2]).*Qs;
    end
    times.tfd += toq();

    # third set of boundary values
    # println("Third interior boundaries:")
    tic();
    for i = 1:length(splits)
        # println("Interior boundary allocation:")
        children = system.branches.children[splits[i]];
        Q = zeros(length(children)+1);
        A = zeros(length(children)+1);
        Q[1] = system.branches.Q[splits[i]][system.solverparams.JL,n+2];
        A[1] = system.branches.A[splits[i]][system.solverparams.JL,n+2];
        for j = 1:length(children)
            Q[j+1] = system.branches.Q[children[j]][1,n+2];
            A[j+1] = system.branches.A[children[j]][1,n+2];
        end
        # println("Interior boundary solution:")
        LegHemodynamics.interiorbcs!(system,n,splits[i],Q,A,children);
    end
    times.ts += toq();
    # println("Third coupling:")
    tic();
    LegHemodynamics.coupling!(system,n,system.solverparams.h,terms,uprox);
    times.tc += toq();

end