"""
    TurbulenceClosures

Functions for turbulence, sub-grid scale modelling. These include
viscosity terms, diffusivity and stress tensors.

- [`ConstantViscosity`](@ref)
- [`ViscousSponge`](@ref)
- [`SmagorinskyLilly`](@ref)
- [`Vreman`](@ref)
- [`AnisoMinDiss`](@ref)

"""
module TurbulenceClosures

# ## Turbulence Closures
# In `turbulence.jl` we specify turbulence closures. Currently,
# pointwise models of the eddy viscosity/eddy diffusivity type are
# supported for turbulent shear and tracer diffusivity. Methods currently supported
# are:\
# [`ConstantViscosity`](@ref constant-viscosity)\
# [`ViscousSponge`](@ref viscous-sponge)\
# [`SmagorinskyLilly`](@ref smagorinsky-lilly)\
# [`Vreman`](@ref vreman)\
# [`AnisoMinDiss`](@ref aniso-min-diss)\

#md # !!! note
#md #     Usage: This is a quick-ref guide to using turbulence models as a subcomponent
#md #     of `BalanceLaw` \
#md #     $\nu$ is the kinematic viscosity, $C_smag$ is the Smagorinsky Model coefficient,
#md #     `turbulence=ConstantDynamicViscosity(ρν)`\
#md #     `turbulence=ConstantKinematicViscosity(ν)`\
#md #     `turbulence=ViscousSponge(ν, z_max, z_sponge, α, γ)`\
#md #     `turbulence=SmagorinskyLilly(C_smag)`\
#md #     `turbulence=Vreman(C_smag)`\
#md #     `turbulence=AnisoMinDiss(C_poincare)`

using ClimateMachine
using DocStringExtensions
using LinearAlgebra
using StaticArrays

import CLIMAParameters: AbstractParameterSet
import ClimateMachine.Mesh.Geometry:
    LocalGeometry, resolutionmetric, lengthscale

using ClimateMachine.Orientations
using ClimateMachine.VariableTemplates
using ClimateMachine.BalanceLaws
using CLIMAParameters.Atmos.SubgridScale: inv_Pr_turb

import ClimateMachine.BalanceLaws:
    vars_state,
    flux_first_order!,
    flux_second_order!,
    source!,
    compute_gradient_argument!,
    compute_gradient_flux!,
    transform_post_gradient_laplacian!,
    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!

export TurbulenceClosureModel,
    ConstantViscosity,
    ConstantDynamicViscosity,
    ConstantKinematicViscosity,
    SmagorinskyLilly,
    Vreman,
    AnisoMinDiss,
    HyperDiffusion,
    NoHyperDiffusion,
    DryBiharmonic,
    EquilMoistBiharmonic,
    NoViscousSponge,
    UpperAtmosSponge,
    turbulence_tensors,
    init_aux_turbulence!,
    init_aux_hyperdiffusion!,
    turbulence_nodal_update_auxiliary_state!,
    sponge_viscosity_modifier

# ### Abstract Type
# We define a `TurbulenceClosureModel` abstract type and
# default functions for the generic turbulence closure
# which will be overloaded with model specific functions.

"""
    Abstract type with default do-nothing behaviour for
arbitrary turbulence closure models.
"""
abstract type TurbulenceClosureModel end

vars_state(::TurbulenceClosureModel, ::AbstractStateType, FT) = @vars()

"""
    ConstantViscosity <: TurbulenceClosureModel
Abstract type for constant viscosity models
"""
abstract type ConstantViscosity <: TurbulenceClosureModel end

"""
    Abstract type for Hyperdiffusion models
"""
abstract type HyperDiffusion end

"""
    Abstract type for viscous sponge layers.
Modifier for viscosity computed from existing turbulence closures.
"""
abstract type ViscousSponge end

"""
    init_aux_turbulence!
Initialise auxiliary variables for turbulence models.
Overload for specific turbulence closure type.
"""
function init_aux_turbulence!(
    ::TurbulenceClosureModel,
    ::BalanceLaw,
    aux::Vars,
    geom::LocalGeometry,
) end

"""
    turbulence_nodal_update_auxiliary_state!
Update auxiliary variables for turbulence models.
Overload for specific turbulence closure type.
"""
function turbulence_nodal_update_auxiliary_state!(
    ::TurbulenceClosureModel,
    ::BalanceLaw,
    state::Vars,
    aux::Vars,
    t::Real,
) end

"""
    compute_gradient_argument!
Assign pre-gradient-transform variables specific to turbulence models.
"""
function compute_gradient_argument!(
    ::TurbulenceClosureModel,
    transform::Vars,
    state::Vars,
    aux::Vars,
    t::Real,
) end
"""
    compute_gradient_flux!(::TurbulenceClosureModel, _...)
Post-gradient-transformed variables specific to turbulence models.
"""
function compute_gradient_flux!(
    ::TurbulenceClosureModel,
    ::Orientation,
    diffusive,
    ∇transform,
    state,
    aux,
    t,
) end

# Fallback functions for hyperdiffusion model
vars_state(::HyperDiffusion, ::AbstractStateType, FT) = @vars()

function init_aux_hyperdiffusion!(
    ::HyperDiffusion,
    ::BalanceLaw,
    aux::Vars,
    geom::LocalGeometry,
) end
function hyperdiffusion_nodal_update_auxiliary_state!(
    ::HyperDiffusion,
    ::BalanceLaw,
    state::Vars,
    aux::Vars,
    t::Real,
) end
function compute_gradient_argument!(
    ::HyperDiffusion,
    ::BalanceLaw,
    transform::Vars,
    state::Vars,
    aux::Vars,
    t::Real,
) end
function transform_post_gradient_laplacian!(
    h::HyperDiffusion,
    bl::BalanceLaw,
    hyperdiffusive::Vars,
    gradvars::Grad,
    state::Vars,
    aux::Vars,
    t::Real,
) end
function flux_second_order!(
    h::HyperDiffusion,
    flux::Grad,
    state::Vars,
    diffusive::Vars,
    hyperdiffusive::Vars,
    aux::Vars,
    t::Real,
) end
function compute_gradient_flux!(
    h::HyperDiffusion,
    diffusive::Vars,
    ∇transform::Grad,
    state::Vars,
    aux::Vars,
    t::Real,
) end

function turbulence_tensors end
function hyperviscosity_tensors end

"""
    ν, D_t, τ = turbulence_tensors(
                    ::TurbulenceClosureModel,
                    orientation::Orientation,
                    param_set::AbstractParameterSet,
                    state::Vars,
                    diffusive::Vars,
                    aux::Vars,
                    t::Real
                )
Compute the kinematic viscosity (`ν`), the diffusivity (`D_t`) and SGS momentum flux tensor (`τ`)
for a given turbulence closure. Each closure overloads this method with the appropriate calculations
for the returned quantities.

# Arguments

- `::TurbulenceClosureModel` = Struct identifier for turbulence closure model
- `orientation` = `BalanceLaw.orientation`
- `param_set` parameter set
- `state` = Array of prognostic (state) variables. See `vars_state` in `BalanceLaw`
- `diffusive` = Array of diffusive variables
- `aux` = Array of auxiliary variables
- `t` = time
"""
function turbulence_tensors(
    m::TurbulenceClosureModel,
    bl::BalanceLaw,
    state,
    diffusive,
    aux,
    t,
)
    ν, D_t, τ = turbulence_tensors(
        m,
        bl.orientation,
        bl.param_set,
        state,
        diffusive,
        aux,
        t,
    )
    ν, D_t, τ = sponge_viscosity_modifier(bl, bl.viscoussponge, ν, D_t, τ, aux)
    return (ν, D_t, τ)
end

# We also provide generic math functions for use within the turbulence closures,
# commonly used quantities such as the [principal tensor invariants](@ref tensor-invariants), handling of
# [symmetric tensors](@ref symmetric-tensors) and [tensor norms](@ref tensor-norms)are addressed.

# ### [Pricipal Invariants](@id tensor-invariants)
# ```math
# \textit{I}_{1} = \mathrm{tr(X)} \\
# \textit{I}_{2} = (\mathrm{tr(X)}^2 - \mathrm{tr(X^2)}) / 2 \\
# \textit{I}_{3} = \mathrm{det(X)} \\
# ```
"""
    principal_invariants(X)

Calculates principal invariants of a tensor `X`. Returns 3 element tuple containing the invariants.
"""
function principal_invariants(X)
    first = tr(X)
    second = (first^2 - tr(X .^ 2)) / 2
    third = det(X)
    return (first, second, third)
end

# ### [Symmetrize](@id symmetric-tensors)
# ```math
# \frac{\mathrm{X} + \mathrm{X}^{T}}{2} \\
# ```
"""
    symmetrize(X)

Given a (3,3) second rank tensor X, compute `(X + X')/2`, returning a
symmetric `SHermitianCompact` object.
"""
function symmetrize(X::StaticArray{Tuple{3, 3}})
    SHermitianCompact(SVector(
        X[1, 1],
        (X[2, 1] + X[1, 2]) / 2,
        (X[3, 1] + X[1, 3]) / 2,
        X[2, 2],
        (X[3, 2] + X[2, 3]) / 2,
        X[3, 3],
    ))
end

# ### [2-Norm](@id tensor-norms)
# Given a tensor X, return the tensor dot product
# ```math
# \sum_{i,j} S_{ij}^2
# ```
"""
    norm2(X)
Given a tensor `X`, computes X:X.
"""
function norm2(X::SMatrix{3, 3, FT}) where {FT}
    abs2(X[1, 1]) +
    abs2(X[2, 1]) +
    abs2(X[3, 1]) +
    abs2(X[1, 2]) +
    abs2(X[2, 2]) +
    abs2(X[3, 2]) +
    abs2(X[1, 3]) +
    abs2(X[2, 3]) +
    abs2(X[3, 3])
end
function norm2(X::SHermitianCompact{3, FT, 6}) where {FT}
    abs2(X[1, 1]) +
    2 * abs2(X[2, 1]) +
    2 * abs2(X[3, 1]) +
    abs2(X[2, 2]) +
    2 * abs2(X[3, 2]) +
    abs2(X[3, 3])
end

# ### [Strain-rate Magnitude](@id strain-rate-magnitude)
# By definition, the strain-rate magnitude, as defined in
# standard turbulence modelling is computed such that
# ```math
# |\mathrm{S}| = \sqrt{2 \sum_{i,j} \mathrm{S}_{ij}^2}
# ```
# where
# ```math
# \vec{S}(\vec{u}) = \frac{1}{2}  \left(\nabla\vec{u} +  \left( \nabla\vec{u} \right)^T \right)
# ```
# \mathrm{S} is the rate-of-strain tensor. (Symmetric component of the velocity gradient). Note that the
# skew symmetric component (rate-of-rotation) is not currently computed.
"""
    strain_rate_magnitude(S)
Given the rate-of-strain tensor `S`, computes its magnitude.
"""
function strain_rate_magnitude(S::SHermitianCompact{3, FT, 6}) where {FT}
    return sqrt(2 * norm2(S))
end

"""
    WithDivergence
A divergence type which includes the divergence term in the momentum flux tensor
"""
struct WithDivergence end
export WithDivergence
"""
    WithoutDivergence
A divergence type which does not include the divergence term in the momentum flux tensor
"""
struct WithoutDivergence end
export WithoutDivergence

# ### [Constant Viscosity Model](@id constant-viscosity)
# `ConstantViscosity` requires a user to specify the constant viscosity (dynamic or kinematic)
# and appropriately computes the turbulent stress tensor based on this term. Diffusivity can be
# computed using the turbulent Prandtl number for the appropriate problem regime.
# ```math
# \tau =
#     \begin{cases}
#     - 2 \nu \mathrm{S} & \mathrm{WithoutDivergence},\\
#     - 2 \nu \mathrm{S} + \frac{2}{3} \nu \mathrm{tr(S)} I_3 & \mathrm{WithDivergence}.
#     \end{cases}
# ```

"""
    ConstantDynamicViscosity <: ConstantViscosity

Turbulence with constant dynamic viscosity (`ρν`).
Divergence terms are included in the momentum flux tensor if divergence_type is WithDivergence.

# Fields

$(DocStringExtensions.FIELDS)
"""
struct ConstantDynamicViscosity{FT, DT} <: ConstantViscosity
    "Dynamic Viscosity [kg/m/s]"
    ρν::FT
    divergence_type::DT
    function ConstantDynamicViscosity(
        ρν::FT,
        divergence_type::Union{WithDivergence, WithoutDivergence} = WithoutDivergence(),
    ) where {FT}
        return new{FT, typeof(divergence_type)}(ρν, divergence_type)
    end
end

"""
    ConstantKinematicViscosity <: ConstantViscosity

Turbulence with constant kinematic viscosity (`ν`).
Divergence terms are included in the momentum flux tensor if divergence_type is WithDivergence.

# Fields

$(DocStringExtensions.FIELDS)
"""
struct ConstantKinematicViscosity{FT, DT} <: ConstantViscosity
    "Kinematic Viscosity [m2/s]"
    ν::FT
    divergence_type::DT
    function ConstantKinematicViscosity(
        ν::FT,
        divergence_type::Union{WithDivergence, WithoutDivergence} = WithoutDivergence(),
    ) where {FT}
        return new{FT, typeof(divergence_type)}(ν, divergence_type)
    end
end

vars_state(::ConstantViscosity, ::Gradient, FT) = @vars()
vars_state(::ConstantViscosity, ::GradientFlux, FT) =
    @vars(S::SHermitianCompact{3, FT, 6})

function compute_gradient_flux!(
    ::ConstantViscosity,
    ::Orientation,
    diffusive::Vars,
    ∇transform::Grad,
    state::Vars,
    aux::Vars,
    t::Real,
)

    diffusive.turbulence.S = symmetrize(∇transform.u)
end

compute_stress(div_type::WithoutDivergence, ν, S) = (-2 * ν) * S
compute_stress(div_type::WithDivergence, ν, S) =
    (-2 * ν) * S + (2 * ν / 3) * tr(S) * I

function turbulence_tensors(
    m::ConstantDynamicViscosity,
    orientation::Orientation,
    param_set::AbstractParameterSet,
    state::Vars,
    diffusive::Vars,
    aux::Vars,
    t::Real,
)

    FT = eltype(state)
    _inv_Pr_turb::FT = inv_Pr_turb(param_set)
    S = diffusive.turbulence.S
    ν = m.ρν / state.ρ
    D_t = ν * _inv_Pr_turb
    τ = compute_stress(m.divergence_type, ν, S)
    return ν, D_t, τ
end

function turbulence_tensors(
    m::ConstantKinematicViscosity,
    orientation::Orientation,
    param_set::AbstractParameterSet,
    state::Vars,
    diffusive::Vars,
    aux::Vars,
    t::Real,
)

    FT = eltype(state)
    _inv_Pr_turb::FT = inv_Pr_turb(param_set)
    S = diffusive.turbulence.S
    ν = m.ν
    D_t = ν * _inv_Pr_turb
    τ = compute_stress(m.divergence_type, ν, S)
    return ν, D_t, τ
end

# ### [Smagorinsky-Lilly](@id smagorinsky-lilly)
# The Smagorinsky turbulence model, with Lilly's correction to
# stratified atmospheric flows, is included in ClimateMachine.
# The input parameter to this model is the Smagorinsky coefficient.
# For atmospheric flows, the coefficient `C_smag` typically takes values between
# 0.15 and 0.23. Flow dependent `C_smag` are currently not supported (e.g. Germano's
# extension). The Smagorinsky-Lilly model does not contain explicit filtered terms.
# #### Equations
# ```math
# \nu = (C_{s} \mathrm{f}_{b} \Delta)^2 \sqrt{|\mathrm{S}|}
# ```
# with the stratification correction term
# ```math
# f_{b} =
#    \begin{cases}
#    1 & \mathrm{Ri} \leq 0 ,\\
#    \max(0, 1 - \mathrm{Ri} / \mathrm{Pr}_{t})^{1/4} & \mathrm{Ri} > 0 .
#    \end{cases}
# ```
# ```math
# \mathrm{Ri} =  \frac{N^2}{{|S|}^2}
# ```
# ```math
# N = \left( \frac{g}{\theta_v} \frac{\partial \theta_v}{\partial z}\right)^{1/2}
# ```
# Here, $\mathrm{Ri}$ and $\mathrm{Pr}_{t}$ are the Richardson and
# turbulent Prandtl numbers respectively.  $\Delta$ is the mixing length in the
# relevant coordinate direction. We use the DG metric terms to determine the
# local effective resolution (see `src/Mesh/Geometry.jl`), and modify the vertical lengthscale by the
# stratification correction factor $\mathrm{f}_{b}$ so that $\Delta_{vert} = \Delta z f_b$.

"""
    SmagorinskyLilly <: TurbulenceClosureModel

# Fields

$(DocStringExtensions.FIELDS)

# Smagorinsky Model Reference

See [Smagorinsky1963](@cite)

# Lilly Model Reference

See [Lilly1962](@cite)

# Brunt-Väisälä Frequency Reference

See [Durran1982](@cite)

"""
struct SmagorinskyLilly{FT} <: TurbulenceClosureModel
    "Smagorinsky Coefficient [dimensionless]"
    C_smag::FT
end

vars_state(::SmagorinskyLilly, ::Auxiliary, FT) = @vars(Δ::FT)
vars_state(::SmagorinskyLilly, ::Gradient, FT) = @vars(θ_v::FT)
vars_state(::SmagorinskyLilly, ::GradientFlux, FT) =
    @vars(S::SHermitianCompact{3, FT, 6}, N²::FT)

function init_aux_turbulence!(
    ::SmagorinskyLilly,
    ::BalanceLaw,
    aux::Vars,
    geom::LocalGeometry,
)
    aux.turbulence.Δ = lengthscale(geom)
end

function compute_gradient_argument!(
    m::SmagorinskyLilly,
    transform::Vars,
    state::Vars,
    aux::Vars,
    t::Real,
)
    transform.turbulence.θ_v = aux.moisture.θ_v
end

function compute_gradient_flux!(
    ::SmagorinskyLilly,
    orientation::Orientation,
    diffusive::Vars,
    ∇transform::Grad,
    state::Vars,
    aux::Vars,
    t::Real,
)

    diffusive.turbulence.S = symmetrize(∇transform.u)
    ∇Φ = ∇gravitational_potential(orientation, aux)
    diffusive.turbulence.N² =
        dot(∇transform.turbulence.θ_v, ∇Φ) / aux.moisture.θ_v
end

function turbulence_tensors(
    m::SmagorinskyLilly,
    orientation::Orientation,
    param_set::AbstractParameterSet,
    state::Vars,
    diffusive::Vars,
    aux::Vars,
    t::Real,
)

    FT = eltype(state)
    _inv_Pr_turb::FT = inv_Pr_turb(param_set)
    S = diffusive.turbulence.S
    normS = strain_rate_magnitude(S)
    k̂ = vertical_unit_vector(orientation, param_set, aux)

    # squared buoyancy correction
    Richardson = diffusive.turbulence.N² / (normS^2 + eps(normS))
    f_b² = sqrt(clamp(FT(1) - Richardson * _inv_Pr_turb, FT(0), FT(1)))
    ν₀ = normS * (m.C_smag * aux.turbulence.Δ)^2 + FT(1e-5)
    ν = SVector{3, FT}(ν₀, ν₀, ν₀)
    ν_v = k̂ .* dot(ν, k̂)
    ν_h = ν₀ .- ν_v
    ν = SDiagonal(ν_h + ν_v .* f_b²)
    D_t = diag(ν) * _inv_Pr_turb
    τ = -2 * ν * S
    return ν, D_t, τ
end

# ### [Vreman Model](@id vreman)
# Vreman's turbulence model for anisotropic flows, which provides a
# less dissipative solution (specifically in the near-wall and transitional regions)
# than the Smagorinsky-Lilly method. This model
# relies of first derivatives of the velocity vector (i.e., the gradient tensor).
# By design, the Vreman model handles transitional as well as fully turbulent flows adequately.
# The input parameter to this model is the Smagorinsky coefficient - the coefficient is modified
# within the model functions to account for differences in model construction.
# #### Equations
# ```math
# \nu_{t} = 2.5 C_{s}^2 \sqrt{\frac{B_{\beta}}{u_{i,j}u_{i,j}}},
# ```
# where ($i,j, m = (1,2,3)$)
# ```math
# \begin{align}
# B_{\beta} &= \beta_{11}\beta_{22} + \beta_{11}\beta_{33} + \beta_{22}\beta_{33} - (\beta_{13}^2 + \beta_{12}^2 + \beta_{23}^2) \\
# \beta_{ij} &= \Delta_{m}^2 u_{i, m} u_{j, m} \\
# u_{i,j} &= \frac{\partial u_{i}}{\partial x_{j}}.
# \end{align}
# ```

"""
    Vreman{FT} <: TurbulenceClosureModel

Filter width Δ is the local grid resolution calculated from the mesh metric tensor. A Smagorinsky coefficient
is specified and used to compute the equivalent Vreman coefficient.

1) ν_e = √(Bᵦ/(αᵢⱼαᵢⱼ)) where αᵢⱼ = ∂uⱼ∂uᵢ with uᵢ the resolved scale velocity component.
2) βij = Δ²αₘᵢαₘⱼ
3) Bᵦ = β₁₁β₂₂ + β₂₂β₃₃ + β₁₁β₃₃ - β₁₂² - β₁₃² - β₂₃²
βᵢⱼ is symmetric, positive-definite.
If Δᵢ = Δ, then β = Δ²αᵀα

# Fields

$(DocStringExtensions.FIELDS)

# Reference

 - [Vreman2004](@cite)
"""
struct Vreman{FT} <: TurbulenceClosureModel
    "Smagorinsky Coefficient [dimensionless]"
    C_smag::FT
end
vars_state(::Vreman, ::Auxiliary, FT) = @vars(Δ::FT)
vars_state(::Vreman, ::Gradient, FT) = @vars(θ_v::FT)
vars_state(::Vreman, ::GradientFlux, FT) =
    @vars(∇u::SMatrix{3, 3, FT, 9}, N²::FT)

function init_aux_turbulence!(
    ::Vreman,
    ::BalanceLaw,
    aux::Vars,
    geom::LocalGeometry,
)
    aux.turbulence.Δ = lengthscale(geom)
end
function compute_gradient_argument!(
    m::Vreman,
    transform::Vars,
    state::Vars,
    aux::Vars,
    t::Real,
)
    transform.turbulence.θ_v = aux.moisture.θ_v
end
function compute_gradient_flux!(
    ::Vreman,
    orientation::Orientation,
    diffusive::Vars,
    ∇transform::Grad,
    state::Vars,
    aux::Vars,
    t::Real,
)
    diffusive.turbulence.∇u = ∇transform.u
    ∇Φ = ∇gravitational_potential(orientation, aux)
    diffusive.turbulence.N² =
        dot(∇transform.turbulence.θ_v, ∇Φ) / aux.moisture.θ_v
end

function turbulence_tensors(
    m::Vreman,
    orientation::Orientation,
    param_set::AbstractParameterSet,
    state::Vars,
    diffusive::Vars,
    aux::Vars,
    t::Real,
)
    FT = eltype(state)
    _inv_Pr_turb::FT = inv_Pr_turb(param_set)
    α = diffusive.turbulence.∇u
    S = symmetrize(α)
    k̂ = vertical_unit_vector(orientation, param_set, aux)

    normS = strain_rate_magnitude(S)
    Richardson = diffusive.turbulence.N² / (normS^2 + eps(normS))
    f_b² = sqrt(clamp(1 - Richardson * _inv_Pr_turb, 0, 1))

    β = (aux.turbulence.Δ)^2 * (α' * α)
    Bβ = principal_invariants(β)[2]

    ν₀ = m.C_smag^2 * FT(2.5) * sqrt(abs(Bβ / (norm2(α) + eps(FT))))

    ν = SVector{3, FT}(ν₀, ν₀, ν₀)
    ν_v = k̂ .* dot(ν, k̂)
    ν_h = ν₀ .- ν_v
    ν = SDiagonal(ν_h + ν_v .* f_b²)
    D_t = diag(ν) * _inv_Pr_turb
    τ = -2 * ν * S
    return ν, D_t, τ
end

# ### [Anisotropic Minimum Dissipation](@id aniso-min-diss)
# This method is based Vreugdenhil and Taylor's minimum-dissipation eddy-viscosity model.
# The principles of the Rayleigh quotient minimizer are applied to the energy dissipation terms in the
# conservation equations, resulting in a maximum dissipation bound, and a model for
# eddy viscosity and eddy diffusivity.
# ```math
# \nu_e = (\mathrm{C}\delta)^2  \mathrm{max}\left[0, - \frac{\hat{\partial}_k \hat{u}_{i} \hat{\partial}_k \hat{u}_{j} \mathrm{\hat{S}}_{ij}}{\hat{\partial}_p \hat{u}_{q}} \right]
# ```
"""
    AnisoMinDiss{FT} <: TurbulenceClosureModel

Filter width Δ is the local grid resolution calculated from the mesh metric tensor. A Poincare coefficient
is specified and used to compute the equivalent AnisoMinDiss coefficient (computed as the solution to the
eigenvalue problem for the Laplacian operator).

# Fields
$(DocStringExtensions.FIELDS)

# Reference

See [Vreugdenhil2018](@cite)

"""
struct AnisoMinDiss{FT} <: TurbulenceClosureModel
    C_poincare::FT
end
vars_state(::AnisoMinDiss, ::Auxiliary, FT) = @vars(Δ::FT)
vars_state(::AnisoMinDiss, ::Gradient, FT) = @vars(θ_v::FT)
vars_state(::AnisoMinDiss, ::GradientFlux, FT) =
    @vars(∇u::SMatrix{3, 3, FT, 9}, N²::FT)
function init_aux_turbulence!(
    ::AnisoMinDiss,
    ::BalanceLaw,
    aux::Vars,
    geom::LocalGeometry,
)
    aux.turbulence.Δ = lengthscale(geom)
end
function compute_gradient_argument!(
    m::AnisoMinDiss,
    transform::Vars,
    state::Vars,
    aux::Vars,
    t::Real,
)
    transform.turbulence.θ_v = aux.moisture.θ_v
end
function compute_gradient_flux!(
    ::AnisoMinDiss,
    orientation::Orientation,
    diffusive::Vars,
    ∇transform::Grad,
    state::Vars,
    aux::Vars,
    t::Real,
)
    ∇Φ = ∇gravitational_potential(orientation, aux)
    diffusive.turbulence.∇u = ∇transform.u
    diffusive.turbulence.N² =
        dot(∇transform.turbulence.θ_v, ∇Φ) / aux.moisture.θ_v
end
function turbulence_tensors(
    m::AnisoMinDiss,
    orientation::Orientation,
    param_set::AbstractParameterSet,
    state::Vars,
    diffusive::Vars,
    aux::Vars,
    t::Real,
)
    FT = eltype(state)
    k̂ = vertical_unit_vector(orientation, param_set, aux)
    _inv_Pr_turb::FT = inv_Pr_turb(param_set)

    ∇u = diffusive.turbulence.∇u
    S = symmetrize(∇u)
    normS = strain_rate_magnitude(S)

    δ = aux.turbulence.Δ
    Richardson = diffusive.turbulence.N² / (normS^2 + eps(normS))
    f_b² = sqrt(clamp(1 - Richardson * _inv_Pr_turb, 0, 1))

    δ_vec = SVector(δ, δ, δ)
    δ_m = δ_vec ./ transpose(δ_vec)
    ∇û = ∇u .* δ_m
    Ŝ = symmetrize(∇û)
    ν₀ =
        (m.C_poincare .* δ_vec) .^ 2 * max(
            FT(1e-5),
            -dot(transpose(∇û) * (∇û), Ŝ) / (dot(∇û, ∇û) .+ eps(normS)),
        )

    ν_v = k̂ .* dot(ν₀, k̂)
    ν_h = ν₀ .- ν_v
    ν = SDiagonal(ν_h + ν_v .* f_b²)
    D_t = diag(ν) * _inv_Pr_turb
    τ = -2 * ν * S
    return ν, D_t, τ
end

"""
  NoHyperDiffusion <: HyperDiffusion
Defines a default hyperdiffusion model with zero hyperdiffusive fluxes.
"""
struct NoHyperDiffusion <: HyperDiffusion end

hyperviscosity_tensors(m::HyperDiffusion, bl::BalanceLaw, args...) =
    hyperviscosity_tensors(m, bl.orientation, bl.param_set, args...)

"""
  EquilMoistBiharmonic{FT} <: HyperDiffusion

Assumes equilibrium thermodynamics in compressible flow.
Horizontal hyperdiffusion methods for application in GCM and LES settings
Timescales are prescribed by the user while the diffusion coefficient is
computed as a function of the grid lengthscale.

# Fields
$(DocStringExtensions.FIELDS)
"""
struct EquilMoistBiharmonic{FT} <: HyperDiffusion
    τ_timescale::FT
    τ_timescale_q_tot::FT
end

EquilMoistBiharmonic(τ_timescale::FT) where {FT} =
    EquilMoistBiharmonic(τ_timescale, τ_timescale)

vars_state(::EquilMoistBiharmonic, ::Auxiliary, FT) = @vars(Δ::FT)
vars_state(::EquilMoistBiharmonic, ::Gradient, FT) =
    @vars(u_h::SVector{3, FT}, h_tot::FT, q_tot::FT)
vars_state(::EquilMoistBiharmonic, ::GradientLaplacian, FT) =
    @vars(u_h::SVector{3, FT}, h_tot::FT, q_tot::FT)
vars_state(::EquilMoistBiharmonic, ::Hyperdiffusive, FT) = @vars(
    ν∇³u_h::SMatrix{3, 3, FT, 9},
    ν∇³h_tot::SVector{3, FT},
    ν∇³q_tot::SVector{3, FT}
)

function init_aux_hyperdiffusion!(
    ::EquilMoistBiharmonic,
    ::BalanceLaw,
    aux::Vars,
    geom::LocalGeometry,
)
    aux.hyperdiffusion.Δ = lengthscale(geom)
end

function compute_gradient_argument!(
    h::EquilMoistBiharmonic,
    bl::BalanceLaw,
    transform::Vars,
    state::Vars,
    aux::Vars,
    t::Real,
)
    ρinv = 1 / state.ρ
    u = state.ρu * ρinv
    k̂ = vertical_unit_vector(bl, aux)
    u_h = (SDiagonal(1, 1, 1) - k̂ * k̂') * u
    transform.hyperdiffusion.u_h = u_h
    transform.hyperdiffusion.h_tot = transform.h_tot
    transform.hyperdiffusion.q_tot = state.moisture.ρq_tot * ρinv
end

function transform_post_gradient_laplacian!(
    h::EquilMoistBiharmonic,
    bl::BalanceLaw,
    hyperdiffusive::Vars,
    hypertransform::Grad,
    state::Vars,
    aux::Vars,
    t::Real,
)
    _inv_Pr_turb = eltype(state)(inv_Pr_turb(bl.param_set))
    ∇Δu_h = hypertransform.hyperdiffusion.u_h
    ∇Δh_tot = hypertransform.hyperdiffusion.h_tot
    ∇Δq_tot = hypertransform.hyperdiffusion.q_tot
    # Unpack
    τ_timescale = h.τ_timescale
    τ_timescale_q_tot = h.τ_timescale_q_tot
    # Compute hyperviscosity coefficient
    ν₄ = (aux.hyperdiffusion.Δ / 2)^4 / 2 / τ_timescale
    ν₄_q_tot = (aux.hyperdiffusion.Δ / 2)^4 / 2 / τ_timescale_q_tot
    hyperdiffusive.hyperdiffusion.ν∇³u_h = ν₄ * ∇Δu_h
    hyperdiffusive.hyperdiffusion.ν∇³h_tot = ν₄ * ∇Δh_tot
    hyperdiffusive.hyperdiffusion.ν∇³q_tot = ν₄_q_tot * ∇Δq_tot
end

function flux_second_order!(
    h::EquilMoistBiharmonic,
    flux::Grad,
    state::Vars,
    diffusive::Vars,
    hyperdiffusive::Vars,
    aux::Vars,
    t::Real,
)
    flux.ρu += state.ρ * hyperdiffusive.hyperdiffusion.ν∇³u_h
    flux.ρe += hyperdiffusive.hyperdiffusion.ν∇³u_h * state.ρu
    flux.ρe += hyperdiffusive.hyperdiffusion.ν∇³h_tot * state.ρ
    flux.moisture.ρq_tot += hyperdiffusive.hyperdiffusion.ν∇³q_tot * state.ρ
end

"""
  DryBiharmonic{FT} <: HyperDiffusion

Assumes dry compressible flow.
Horizontal hyperdiffusion methods for application in GCM and LES settings
Timescales are prescribed by the user while the diffusion coefficient is
computed as a function of the grid lengthscale.

# Fields
$(DocStringExtensions.FIELDS)
"""
struct DryBiharmonic{FT} <: HyperDiffusion
    τ_timescale::FT
end
vars_state(::DryBiharmonic, ::Auxiliary, FT) = @vars(Δ::FT)
vars_state(::DryBiharmonic, ::Gradient, FT) =
    @vars(u_h::SVector{3, FT}, h_tot::FT)
vars_state(::DryBiharmonic, ::GradientLaplacian, FT) =
    @vars(u_h::SVector{3, FT}, h_tot::FT)
vars_state(::DryBiharmonic, ::Hyperdiffusive, FT) =
    @vars(ν∇³u_h::SMatrix{3, 3, FT, 9}, ν∇³h_tot::SVector{3, FT})

function init_aux_hyperdiffusion!(
    ::DryBiharmonic,
    ::BalanceLaw,
    aux::Vars,
    geom::LocalGeometry,
)
    aux.hyperdiffusion.Δ = lengthscale(geom)
end

function compute_gradient_argument!(
    h::DryBiharmonic,
    bl::BalanceLaw,
    transform::Vars,
    state::Vars,
    aux::Vars,
    t::Real,
)
    ρinv = 1 / state.ρ
    u = state.ρu * ρinv
    k̂ = vertical_unit_vector(bl, aux)
    u_h = (SDiagonal(1, 1, 1) - k̂ * k̂') * u
    transform.hyperdiffusion.u_h = u_h
    transform.hyperdiffusion.h_tot = transform.h_tot
end

function transform_post_gradient_laplacian!(
    h::DryBiharmonic,
    bl::BalanceLaw,
    hyperdiffusive::Vars,
    hypertransform::Grad,
    state::Vars,
    aux::Vars,
    t::Real,
)
    _inv_Pr_turb = eltype(state)(inv_Pr_turb(bl.param_set))
    ∇Δu_h = hypertransform.hyperdiffusion.u_h
    ∇Δh_tot = hypertransform.hyperdiffusion.h_tot
    # Unpack
    τ_timescale = h.τ_timescale
    # Compute hyperviscosity coefficient
    ν₄ = (aux.hyperdiffusion.Δ / 2)^4 / 2 / τ_timescale
    hyperdiffusive.hyperdiffusion.ν∇³u_h = ν₄ * ∇Δu_h
    hyperdiffusive.hyperdiffusion.ν∇³h_tot = ν₄ * ∇Δh_tot
end

function flux_second_order!(
    h::DryBiharmonic,
    flux::Grad,
    state::Vars,
    diffusive::Vars,
    hyperdiffusive::Vars,
    aux::Vars,
    t::Real,
)
    flux.ρu += state.ρ * hyperdiffusive.hyperdiffusion.ν∇³u_h
    flux.ρe += hyperdiffusive.hyperdiffusion.ν∇³u_h * state.ρu
    flux.ρe += hyperdiffusive.hyperdiffusion.ν∇³h_tot * state.ρ
end

# ### [Viscous Sponge](@id viscous-sponge)
# `ViscousSponge` requires a user to specify a constant viscosity (kinematic),
# a sponge start height, the domain height, a sponge strength, and a sponge
# exponent.
# Given viscosity, diffusivity and stresses from arbitrary turbulence models,
# the viscous sponge enhances diffusive terms within a user-specified layer,
# typically used at the top of the domain to absorb waves. A smooth onset is
# ensured through a weight function that increases weight height from the sponge
# onset height.
# ```
"""
    NoViscousSponge
No modifiers applied to viscosity/diffusivity in sponge layer
# Fields
#
$(DocStringExtensions.FIELDS)
"""
struct NoViscousSponge <: ViscousSponge end
function sponge_viscosity_modifier(
    bl::BalanceLaw,
    m::NoViscousSponge,
    ν,
    D_t,
    τ,
    aux,
)
    return (ν, D_t, τ)
end

"""
    Upper domain viscous relaxation
Applies modifier to viscosity and diffusivity terms
in a user-specified upper domain sponge region
# Fields
#
$(DocStringExtensions.FIELDS)
"""
struct UpperAtmosSponge{FT} <: ViscousSponge
    "Maximum domain altitude (m)"
    z_max::FT
    "Altitude at with sponge starts (m)"
    z_sponge::FT
    "Sponge Strength 0 ⩽ α_max ⩽ 1"
    α_max::FT
    "Sponge exponent"
    γ::FT
end

function sponge_viscosity_modifier(
    bl::BalanceLaw,
    m::UpperAtmosSponge,
    ν,
    D_t,
    τ,
    aux::Vars,
)
    z = altitude(bl.orientation, bl.param_set, aux)
    if z >= m.sponge
        r = (z - m.z_sponge) / (m.z_max - m.z_sponge)
        β_sponge = m.α_max * sinpi(r / 2)^m.γ
        ν += β_sponge * ν
        D_t += β_sponge * D_t
        τ += β_sponge * τ
    end
    return (ν, D_t, τ)
end

end #module TurbulenceClosures.jl


abstract type Spectrum end
abstract type Multispectral <: Spectrum end # Multispectral types
abstract type CMatch <: Multispectral end #Color Matching Functions (x_bar, y_bar, z_bar)
abstract type ConeFund <: Multispectral end # Cone fundamentals (l, m, s)
abstract type COpp <: Multispectral end # Opponent color channels
# abstract type MSpec <: Multispectral end # Generalized multispectral type, maybe implemented later

"""
    LSpec, Luminance Spectrum

Fields:

- `λ::Vector{Real}`, vector containing wavelengths
- `l::Vector{Real}`, vector containing luminance values
"""
struct LSpec <: Spectrum
    λ::Vector{Real}  #wavelength vector
    l::Vector{Real}  #reflectance vector
    function LSpec(λ, l)
        new(λ, l)
    end
end

"""
    RSpec, Reflectance Spectrum

Fields:

- `λ::Vector{Real}`, vector containing wavelengths
- `r::Vector{Real}, vector containing reflectanceS values
"""
struct RSpec <: Spectrum
    λ::Vector{Real}  #wavelength vector
    r::Vector{Real}  #reflectance vector
    function RSpec(λ, r)
        new(λ, r)
    end
end

"""
    TSpec, Transmittance spectrum

Fields:

- `λ::Vector{Real}`, vector containing wavelengths
- `t::Vector{Real}`, vector containing transmittance values
- `x::Real`, unit thickness value
"""
struct TSpec <: Spectrum
    λ::Vector{Real}  #start wavelength
    t::Vector{Real}  #transmittance vector at unit thickness
    x::Real          #thickness in terms of unit thickness
    function TSpec(λ, t, x)
        new(λ, t, x)
    end
end

#=  Multispectral types
    These include color matching functions and cone fundamentals defined by the CIE
=#

"""
    CIE31

Type of CIE 1931 2° observer color matching function.
"""
struct CIE31 <: CMatch
    λ::Vector{Real} #wavelength vector
    x::Vector{Real} #x_bar cone matching function
    y::Vector{Real} #y_bar cone matching function
    z::Vector{Real} #z_bar cone matching function
    function CIE31(λ, l, m, s)
        new(λ, l, m, s)
    end
end

"""
    LMS31

Type of CIE 1931 2° observer cone fundamentals.
"""
struct LMS31 <: ConeFund
    λ::Vector{Real}
    l::Vector{Real}
    m::Vector{Real}
    s::Vector{Real}
    function LMS31(λ, l, m, s)
        new(λ, l, m, s)
    end
end

"""
    CIE31_J

Type of CIE 1931 2° observer color matching function with Judd corrections.
"""
struct CIE31_J <: CMatch
    λ::Vector{Real} #wavelength vector
    x::Vector{Real} #x_bar cone matching function
    y::Vector{Real} #y_bar cone matching function
    z::Vector{Real} #z_bar cone matching function
    function CIE31_J(λ, l, m, s)
        new(λ, l, m, s)
    end
end

"""
    CIE31_JV

Type of CIE 1931 2° observer color matching function with corrections by Judd and Vos.
"""
struct CIE31_JV <: CMatch
    λ::Vector{Real} #wavelength vector
    x::Vector{Real} #x_bar cone matching function
    y::Vector{Real} #y_bar cone matching function
    z::Vector{Real} #z_bar cone matching function
    function CIE31_JV(λ, l, m, s)
        new(λ, l, m, s)
    end
end

"""
    CIE31_J

Type of CIE 1964 10° observer color matching function.
"""
struct CIE64 <: CMatch
    λ::Vector{Real} #wavelength vector
    x::Vector{Real} #x_bar cone matching function
    y::Vector{Real} #y_bar cone matching function
    z::Vector{Real} #z_bar cone matching function
    function CIE64(λ, l, m, s)
        new(λ, l, m, s)
    end
end

"""
    LMS64

Type of CIE 1964 10° observer cone fundamentals.
"""
struct LMS64 <: ConeFund
    λ::Vector{Real}
    l::Vector{Real}
    m::Vector{Real}
    s::Vector{Real}
    function LMS64(λ, l, m, s)
        new(λ, l, m, s)
    end
end

"""
    CIE12_2

Type of CIE 2012 2° observer color matching function. Current CIE 2° standard observer.
"""
struct CIE12_2 <: CMatch
    λ::Vector{Real} # wavelength vector
    x::Vector{Real} #x_bar cone matching function
    y::Vector{Real} #y_bar cone matching function
    z::Vector{Real} #z_bar cone matching function
    function CIE12_2(λ, l, m, s)
        new(λ, l, m, s)
    end
end

"""
    CIE12_10

Type of CIE 2012 10° observer color matching function. Current CIE 10° standard observer.
"""
struct CIE12_10 <: CMatch
    λ::Vector{Real} # wavelength vector
    x::Vector{Real} #x_bar cone matching function
    y::Vector{Real} #y_bar cone matching function
    z::Vector{Real} #z_bar cone matching function
    function CIE12_10(λ, l, m, s)
        new(λ, l, m, s)
    end
end

"""
    LMS06_2

Type of CIE 2006 2° observer cone fundamentals.
"""
struct LMS06_2 <: ConeFund
    λ::Vector{Real} #wavelength vector
    l::Vector{Real} #l cone response function
    m::Vector{Real} #m cone response function
    s::Vector{Real} #s cone response function
    function LMS06_2(λ, l, m, s)
        new(λ, l, m, s)
    end
end

"""
    LMS06_10

Type of CIE 2006 10° observer cone fundamentals.
"""
struct LMS06_10 <: ConeFund
    λ::Vector{Real} #wavelength vector
    l::Vector{Real} #l cone response function
    m::Vector{Real} #m cone response function
    s::Vector{Real} #s cone response function
    function LMS06_10(λ, l, m, s)
        new(λ, l, m, s)
    end
end

"""
    VYBRG

Type of Ingling-Tsou color opponent space functions with `Vλ` (lightness), `yb` (yellow-blue), `rg` (red-green) color channels.
"""
struct VYBRG <: COpp
    λ::Vector{Real}
    Vλ::Vector{Real}
    yb::Vector{Real}
    rg::Vector{Real}
    function VYBRG(λ, Vλ, yb, rg)
        new(λ, Vλ, yb, rg)
    end
end


include("../imports.jl")

@testset ExtendedTestSet "`ImageClassification`" begin
    @testset ExtendedTestSet "Core interface" begin
        @testset ExtendedTestSet "`encodeinput`" begin
            method = ImageClassification(10, (32, 32))
            image = rand(RGB, 64, 96)

            xtrain = encodeinput(method, Training(), image)
            @test size(xtrain) == (32, 32, 3)
            @test eltype(xtrain) == Float32

            xinference = encodeinput(method, Inference(), image)
            @test size(xinference) == (32, 48, 3)
            @test eltype(xinference) == Float32
        end

        @testset ExtendedTestSet "`encodetarget`" begin
            method = ImageClassification(2, (32, 32))
            category = 1
            y = encodetarget(method, Training(), category)
            @test y ≈ [1, 0]
            encodetarget!(y, method, Training(), 2)
            @test y ≈ [0, 1]
        end

        @testset ExtendedTestSet "`encode`" begin
            method = ImageClassification(10, (32, 32))
            image = rand(RGB, 64, 96)
            category = 1
            @test_nowarn encode(method, Training(), (image, category))
        end
    end
end


module YaoTensorNetwork

using BitBasis
using Yao, LuxurySparse
using DelimitedFiles
using Random
using Requires
using YaoBlocks.Optimise

const control_seperated = Ref(false)

export seperatecontrol
"""
    seperatecontrol([val::Bool])

Global switch, seperate control into rank tree tensors if true.
set value if `val` provided.
"""
function seperatecontrol(val::Bool)
    control_seperated[] = val
end

seperatecontrol() = control_seperated[]

include("utils.jl")
include("EinGraph.jl")
include("GraphBuilder.jl")
include("convert.jl")
include("tobasic_ext.jl")

@init @require OMEinsum = "ebe7aa44-baf0-506c-a96f-8464559b3922" include("contract.jl")

end # module


__precompile__()
module FragmentGrammars

export Analysis, BaseDistribution, BaseRule, FragmentRule, AbstractRule, Fragment, Pointer, FragmentGrammar
export sample, add_obs!, rm_obs!
export run_chartparser, sample_tree
export category_type, terminal_type, category_rule_type, terminal_rule_type, startstate, startsymbols, score_type, state_type, completions, prob
export iterate, eltype, IteratorSize
# export show

import Base: iterate, eltype, IteratorSize#, length, show

include("GeneralizedChartparsing\\src\\GeneralizedChartparsing.jl")
using .GeneralizedChartparsing
using .GeneralizedChartparsing.Trees
using .GeneralizedChartparsing: sample_tree, Constituent
import .GeneralizedChartparsing: lhs, rhs, category_type, terminal_type, category_rule_type, terminal_rule_type, startstate, add_rule!, startsymbols, score_type, state_type, completions, prob
import .GeneralizedChartparsing: run_chartparser, categorical_sample

include("CompoundDists.jl");
using .CompoundDists
import .CompoundDists: sample, add_obs!, rm_obs!, logscore

# TODO: ALL OF THIS BADLY NEEDS TO BE REFACTORED.
# TODO: So much code can be turned into broadcasts
# TODO: Coins are flipped for preterminals??? Sheesh. Fix that.

###########################
# Helper structs and functions #
###########################
struct Fragment
    tree :: Tree
    variables :: Vector{Tree}
    leaves :: Vector{Tree}
end

struct Pointer
    fragment :: Fragment
    children :: Dict{Tree, Pointer}
end

function clone(ptr::Pointer)
    function clone(ptr::Pointer, tree::Tree, variables::Vector{Tree}, leaves::Vector{Tree}, children::Dict{Tree, Pointer})
        newtree = TreeNode(tree.data)
        if !isempty(tree.children)
            for child in tree.children
                insert_child!(newtree, clone(ptr, child, variables, leaves, children))
            end
        else
            if tree in ptr.fragment.variables
                if get(ptr.children, tree, nothing) != nothing
                    push!(children, newtree => ptr.children[tree])
                end
                push!(variables, newtree)
                push!(leaves, newtree)
            elseif tree in ptr.fragment.leaves
                push!(leaves, newtree)
            end
        end
        return newtree
    end
    variables = Tree[]
    leaves = Tree[]
    children = Dict{Tree, Pointer}()
    return Pointer(Fragment(clone(ptr, ptr.fragment.tree, variables, leaves, children), variables, leaves), children)
end

clone(frag::Fragment) = clone(Pointer(frag, Dict{Tree,Pointer}())).fragment

struct Analysis{C, CR} # Is this struct even needed? It's basically a Tuple of what comes out of FG sample
    pointer :: Pointer
    dm_obs :: Vector{Tuple{C,CR}}
    bb_obs :: Vector{Pair{Tuple{C, CR, C},Bool}}
    crp_obs :: Vector{Fragment}
end

abstract type AbstractRule{C1,C2} end

struct BaseRule{C1,C2} <: AbstractRule{C1,C2}
    # rule :: ContextFreeRule{C1,C2}
    lhs :: C1
    rhs :: Tuple{Vararg{C2}}
end
BaseRule(rule) = BaseRule(#= rule, =# lhs(rule), rhs(rule))
# BaseRule(lhs::C1, rhs::Tuple{Vararg{C2}}) where {C1,C2} = BaseRule(ContextFreeRule{C1,C2}(lhs, rhs))

struct FragmentRule{C1,C2} <: AbstractRule{C1,C2}
    fragment :: Fragment
    # rule :: ContextFreeRule{C1,C2}
    lhs :: C1
    rhs :: Tuple{Vararg{C2}}
end
FragmentRule(rule, fragment) = FragmentRule(fragment, #= rule, =#lhs(rule), rhs(rule))
FragmentRule(fragment) = FragmentRule(fragment, fragment.tree.data, (getfield.(fragment.leaves, :data)...,))

mutable struct ApproxRule{C1,C2} <: AbstractRule{C1,C2}
    lhs :: C1
    rhs :: Tuple{Vararg{C2}}
    rules :: Vector{AbstractRule{C1,C2}}
    probs :: Vector{LogProb}
    prob :: LogProb
end

approxprob(ar::ApproxRule) = ar.prob
approxprob(br::BaseRule) = one(LogProb)

# Unzip rule-logprob pairs into two lists, cast the logprob list into floats, and call categorical_sample
sample(r::ApproxRule) = categorical_sample(r.rules, r.probs)
# sample(r::ApproxRule) = categorical_sample((x->(x[1], float.(x[2])))(collect(zip(r.rules...)))...)
categorical_sample(tokens::Vector{X}, weights::Vector{LogProb}) where X = categorical_sample(tokens, float.(weights))

lhs(r::AbstractRule) = r.lhs
rhs(r::AbstractRule) = r.rhs
(r::BaseRule)(cat) = rhs(r)
(r::BaseRule)() = rhs(r)
(r::FragmentRule)(cat) = rhs(r)
(r::FragmentRule)() = rhs(r)
(r::ApproxRule)(cat) = rhs(r)
(r::ApproxRule)() = rhs(r)

show(io::IO, r::AbstractRule{C1,C2}) where {C1,C2} =
print("AbstractRule($(lhs(r)),$(rhs(r)))")
show(io::IO, r::BaseRule{C1,C2}) where {C1,C2} =
println("BaseRule($(lhs(r)),$(rhs(r)))")
show(io::IO, r::FragmentRule{C1,C2}) where {C1,C2} =
print("FragmentRule($(lhs(r)),$(rhs(r)))")
show(io::IO, r::ApproxRule{C1,C2}) where {C1,C2} =
print("ApproxRule($(lhs(r)),$(rhs(r)))")

eltype(::Type{Pointer}) = Pointer
IteratorSize(::Type{Pointer}) = Base.SizeUnknown()
function iterate(frag_pointer::Pointer, state = [frag_pointer])
    if isempty(state)
        nothing
    else
        state[1], prepend!(state[2:end], collect(values(state[1].children)))
    end
end

eltype(::Type{State{C,CR}}) where {C,CR} = State{C,CR}
IteratorSize(::Type{State{C,CR}}) where {C,CR} = Base.SizeUnknown()
function iterate(fsa_state::State, state = [fsa_state])
    if isempty(state)
        nothing
    else
        state[1], prepend!(state[2:end], collect(values(state[1].trans)))
    end
end

get_idx(A::AbstractVector{T}, i::T) where T = (
    for (j,k) in enumerate(A)
        if i == k
            return j
        end
    end; error("element $i not found in $A")
)

# Basically copied from add_rule! in GeneralizedChartparsing/src/Grammars.jl
function add_rule!(state::State{C,AbstractRule{C,C}}, rule::AbstractRule{C,C}, head, cats) where {C} #, CR<:ApproxRule{C,C}}
    s = state
    for c in cats
        if is_possible_transition(s, c)
            s = transition(s, c)
        else
            s.isfinal = false
            s = s.trans[c] = State(C, AbstractRule{C,C})
        end
    end
    # All unique lhs-rhs pairs are collapsed to include base rules and fragments, and transitions between states determine RHS, therefore we only need to check LHS here.
    aridx = findall((comp -> comp[1] == head), s.comp)
    # @show aridx, head, cats, rule
    if isempty(aridx) # if no approx rule yet
        ar = ApproxRule(head, cats, AbstractRule{C,C}[rule], LogProb[zero(LogProb)], zero(LogProb))
        push!(s.comp, (head, ar))
    else
        ar = s.comp[aridx[1]][2] # there can be only one! [2] because [1] is lhs and [2] is approxrule
        if !(rule in ar.rules) push!(ar.rules, rule); push!(ar.probs, zero(LogProb)) end
    end
    # push!(s.comp, (head, rule))
end

function rm_rule!(state::State{C,AbstractRule{C,C}}, rule::AbstractRule{C,C}, head, cats) where {C} #, CR<:ApproxRule{C,C}}
    s = state
    for c in cats
        if is_possible_transition(s, c)
            s = transition(s, c)
        else
            s.isfinal = false
            s = s.trans[c] = State(C, CR)
        end
    end
    aridx = findall((comp -> comp[1] == head), s.comp)
    # @show aridx, head, cats, rule
    # if isempty(aridx) # if no approx rule yet
        # Nothing needs to happen here, right?
    # else
    if !isempty(aridx) # Shouldn't this never be the case?
        ar = s.comp[aridx[1]][2] # there can be only one! [2] because [1] is lhs and [2] is approxrule
        # @show ar.rules
        ridx = findall(r -> r===rule, ar.rules)
        if !isempty(ridx)
            ridx = ridx[1] # there can be only one!
            deleteat!(ar.rules, ridx)
            deleteat!(ar.probs, ridx)
            # filter!(r -> r[1]≠rule, ar.rules) # Remove "rule" from the approx rule
            if isempty(ar.rules) && isempty(ar.probs)
                deleteat!(s.comp, aridx) # Remove empty approx rule from completions
            end
        end
    end
end

################################
# Fragment Grammar definitions #
################################
"""
    FragmentGrammar(categories, startcategories, category_rules, terminals, terminal_rules, startstate, terminal_dict, CRP, DM, BB)
"""
mutable struct FragmentGrammar{C, CR, T, TR}
    categories :: Vector{C}
    startcategories :: Vector{C}
    category_rules :: Vector{CR}
    terminals :: Vector{T}
    terminal_rules :: Vector{TR}
    preterminals :: Vector{C}
    startstate :: State{C, AbstractRule{C,C}}
    terminal_dict :: Dict{T, Vector{Tuple{C, TR}}}
    CRP :: Dict{C, ChineseRest{Fragment}}
    DM :: Dict{C, DirCat{CR, Float64}}
    BB :: Dict{Tuple{C, CR, C}, BetaBern{Bool, Int}}
end

show(io::IO, fg::FragmentGrammar{C, CR, T, TR}) where {C, CR, T, TR} =
    print("FragmentGrammar{$C, $(CR), $T, $(TR)}()")

show(io::IO, crp::ChineseRest) =
    print("ChineseRest($(crp.a), $(crp.b), $(crp.num_tables), $(crp.num_customers), $(crp.tables))")

category_type(fg::FragmentGrammar{C, CR, T, TR}) where {C, CR, T, TR} = C
terminal_type(fg::FragmentGrammar{C, CR, T, TR}) where {C, CR, T, TR} = T
category_rule_type(fg::FragmentGrammar{C, CR, T, TR}) where {C, CR, T, TR} = AbstractRule{C,C}
terminal_rule_type(fg::FragmentGrammar{C, CR, T, TR}) where {C, CR, T, TR} = AbstractRule{C,T}
categories(fg::FragmentGrammar) = fg.categories
terminals(fg::FragmentGrammar) = fg.terminals
preterminals(fg::FragmentGrammar) = fg.preterminals
startstate(fg::FragmentGrammar) = fg.startstate
startsymbols(fg::FragmentGrammar) = fg.startcategories
score_type(fg::FragmentGrammar) = LogProb
state_type(fg::FragmentGrammar{C}) where C = State{C,AbstractRule{C,C}} # typeof(startstate(fg))
completions(fg::FragmentGrammar{C}, state::State{C,AbstractRule{C,C}}) where C =
    ((cat, r, prob(fg, cat, r)) for (cat, r) in completions(state))
completions(fg::FragmentGrammar{C, CR, T, TR}, t::T) where {C, T, CR, TR} =
    ((cat, rule, prob(fg, cat, rule)) for (cat, rule) in fg.terminal_dict[t])
prob(fg::FragmentGrammar{C}, cat::C, rule::BaseRule{C,C}) where C = one(score_type(fg))
prob(fg::FragmentGrammar{C}, cat::C, rule::R) where {C, R<:ApproxRule{C,C}} = rule.prob

function update_approx_probs!(rule::ApproxRule{C,C}, fg::FragmentGrammar{C}) where C
    # Should really use broadcasting or something here...
    cat = lhs(rule)
    rule.prob = zero(LogProb)
    for (i,r) in enumerate(rule.rules)
        if r isa BaseRule
            new_p = logscore(fg.DM[cat], r) * new_table_logscore(fg.CRP[cat])
            rule.prob += new_p
            rule.probs[i] = new_p
        elseif r isa FragmentRule
            new_p = logscore(fg.CRP[cat], r.fragment, true)
            rule.prob += new_p
            rule.probs[i] = new_p
        else error("ApproxRule subrule is of type $(typeof(r)), should be BaseRule or FragmentRule") end
    end
end
update_approx_probs!(completion::Tuple{C,ApproxRule{C,C}}, fg::FragmentGrammar{C}) where C = update_approx_probs!(completion[2], fg)
update_approx_probs!(fg::FragmentGrammar) = for state in startstate(fg) update_approx_probs!.(state.comp, Ref(fg)) end

logscore(fg::FragmentGrammar) = reduce(*, logscore.(values(fg.DM))) * reduce(*, logscore.(values(fg.CRP))) * reduce(*, logscore.(values(fg.BB)))

"""
    FragmentGrammar(categories, startcategories, category_rules, terminals, terminal_rules[, a::Float64, b::Float64])
"""
function FragmentGrammar(cats::Vector{C}, starts::Vector{C}, cat_rules::Vector{CR}, terms::Vector{T}, term_rules::Vector{TR}, a=0.01, b=0.2, dm_pseudo=1.0, bb_alpha=1, bb_beta=1) where {C, CR<:BaseRule{C,C}, T, TR}
    # cat_rules = Vector{CR}(cat_rules)
    startstate = State(C, AbstractRule{C,C})# ApproxRule{C,C})
    for r in cat_rules
        add_rule!(startstate, r, lhs(r), rhs(r))
        # rm_rule!(startstate, r, r.lhs, r.rhs)
    end

    CRP = Dict{C,ChineseRest{Fragment}}()
    DM = Dict{C, DirCat{CR, Float64}}(cat => DirCat{CR,Float64}(Dict(x => dm_pseudo for x in CR[r for (i, r) in enumerate(cat_rules) if lhs(r) == cat]), dm_pseudo) for (j, cat) in enumerate(cats))
    BB = Dict{Tuple{C, CR, C}, BetaBern{Bool, Int}}((lhs(r), r, rhs) => BetaBern(bb_alpha, bb_beta) for r in cat_rules for rhs in rhs(r))
    preterminals = C[]
    terminal_dict = Dict{T, Vector{Tuple{C, TR}}}()
    for r in term_rules
        push!(preterminals, lhs(r))
        t = rhs(r)[1]
        if haskey(terminal_dict, t)
            push!(terminal_dict[t], (lhs(r), r))
        else
            terminal_dict[t] = [(lhs(r), r)]
        end
    end

    fg = FragmentGrammar{C, CR, T, TR}(cats, starts, cat_rules, terms, term_rules, preterminals, startstate, terminal_dict, CRP, DM, BB)

    fg.CRP = Dict{C,ChineseRest{Fragment}}(cat => ChineseRest(a, b, BaseDistribution(fg, cat)) for (i, cat) in enumerate(cats))

    return fg
end

function sample(fg :: FragmentGrammar{C,CR,T,TR}, cat :: C) where {C, CR, T, TR}
    crp_sample = sample(fg.CRP[cat])
    fragment, dm_counts, bb_counts, crp_counts = crp_sample isa Tuple ? crp_sample : (crp_sample, Tuple{C,CR}[], Pair{Tuple{C,CR,C}, Bool}[], Fragment[])
    children = Dict{Tree, Pointer}()
    for variable in fragment.variables
        ptr, dm_counts_child, bb_counts_child, crp_counts_child = sample(fg, variable.data)
        append!(dm_counts, dm_counts_child)
        append!(bb_counts, bb_counts_child)
        append!(crp_counts, crp_counts_child)
        push!(children, variable => ptr)
    end
    return Pointer(fragment, children), dm_counts, bb_counts, crp_counts
end

struct BaseDistribution{C} <: Distribution{Fragment}
    fg :: FragmentGrammar{C}
    category :: C
end

show(io::IO, bd::BaseDistribution{C}) where C = print("BaseDistribution{$C}()")

function sample(basedist :: BaseDistribution)
    C, CR, T, TR = category_type(basedist.fg), category_rule_type(basedist.fg), terminal_type(basedist.fg), terminal_rule_type(basedist.fg)

    variables = Tree{C}[]
    leaves = Tree{C}[]
    tree = TreeNode(basedist.category)

    dm_sample = sample(basedist.fg.DM[basedist.category])
    dm_counts = Tuple{C,CR}[(basedist.category, dm_sample)]
    # push!(dm_counts, (basedist.category, dm_sample))
    bb_counts = Pair{Tuple{C, CR, C}, Bool}[]
    crp_counts = Fragment[]

    r = dm_sample
    children = r(basedist.category) # or r.rhs or rhs(r)

    for child in children
        if child in basedist.fg.preterminals
            # Make a leaf node (non-variable)
            leaf_tree = TreeNode(child)
            insert_child!(tree, leaf_tree)
            push!(leaves, leaf_tree)
        else
            bbidx = (basedist.category, dm_sample, child)
            extend = sample(basedist.fg.BB[bbidx])
            push!(bb_counts, bbidx => extend)
            if extend # if we extend the fragment
                # Get (recursive) Pointer, DM counts, BB counts, and CRP counts (fragments)
                ptr, dm_counts_child, bb_counts_child, crp_counts_child = sample(basedist.fg, child)
                fragchild = clone(ptr.fragment) # ptr.fragment # We only need the fragment for the base distribution
                # Merge BB counts
                append!(bb_counts, bb_counts_child)
                # Merge DM counts
                append!(dm_counts, dm_counts_child)
                # Merge CRP counts (recursively counts only if we're *extending* a non-terminal with another fragment)
                # If it's a new fragment, it may contain a fragment in crp_counts_child. If it's already stored, then it contains nothing.
                append!(crp_counts, crp_counts_child)
                # Add to tree
                insert_child!(tree, fragchild.tree) # deepcopy(fragchild.tree))   # Better than modifying the underlying tree if taken from CRP preexisting fragment
                # Merge variables (non-terminal leaves) from recursive calls to FG
                append!(variables, fragchild.variables)
                append!(leaves, fragchild.leaves)
            else
                # Make a leaf node (variable)
                variable_tree = TreeNode(child)
                push!(variables, variable_tree)
                push!(leaves, variable_tree)
                insert_child!(tree, variable_tree)
            end
        end
    end

    frag = Fragment(tree, variables, leaves)
    push!(crp_counts, frag)

    return frag, dm_counts, bb_counts, crp_counts
end

# TODO: Write tests for add/rm_obs!
function add_obs!(fg :: FragmentGrammar, analysis :: Analysis)
    # Add DM counts
    for dm_obs in analysis.dm_obs
        add_obs!(fg.DM[dm_obs[1]], dm_obs[2])
    end
    # Add BB counts
    for bb_obs in analysis.bb_obs
        add_obs!(fg.BB[bb_obs[1]], bb_obs[2])
    end
    # Add fragments to CRP
    for frag in analysis.crp_obs
        add_obs!(fg.CRP[frag.tree.data], frag)
        fr = FragmentRule(frag)
        add_rule!(startstate(fg), fr, lhs(fr), rhs(fr))
    end
    for frag_ptr in analysis.pointer
        add_obs!(fg.CRP[frag_ptr.fragment.tree.data], frag_ptr.fragment)
        # Add completions to finite state machine (for approx. PCFG)
        fr = FragmentRule(frag_ptr.fragment)
        add_rule!(startstate(fg), fr, lhs(fr), rhs(fr))
    end
end

function rm_obs!(fg :: FragmentGrammar, analysis :: Analysis)
    # rm DM counts
    for dm_obs in analysis.dm_obs
        rm_obs!(fg.DM[dm_obs[1]], dm_obs[2])
    end
    # rm BB counts
    for bb_obs in analysis.bb_obs
        rm_obs!(fg.BB[bb_obs[1]], bb_obs[2])
    end
    # rm fragments from CRP
    for frag in analysis.crp_obs
        rm_obs!(fg.CRP[frag.tree.data], frag)
        fr = FragmentRule(frag)
        rm_rule!(startstate(fg), fr, lhs(fr), rhs(fr))
    end
    for frag_ptr in analysis.pointer
        rm_obs!(fg.CRP[frag_ptr.fragment.tree.data], frag_ptr.fragment)
        # rm completions to finite state machine (for approx. PCFG)
        fr = FragmentRule(frag_ptr.fragment)
        rm_rule!(startstate(fg), fr, lhs(fr), rhs(fr))
    end
end

# TODO: Refactor this method so it isn't spaghetti
function sample(parsetree :: Tree{Tuple{Cons,AbstractRule{C,C}}}, fg::FragmentGrammar{C,CR,T,TR}) where {Cons<:Constituent, C, CR, T, TR}
    cons, rule = parsetree.data

    local tree, frag
    variables = Tree{C}[]
    leaves = Tree{C}[]
    children = Dict{Tree, Pointer}()

    dm_obs = Tuple{C,CR}[]
    bb_obs = Pair{Tuple{C, CR, C}, Bool}[]
    crp_obs = Fragment[]

    if rule isa BaseRule{C,C} # If terminal rule, just return nothing and the caller will make a leaf from the preterminal.
        return nothing, dm_obs, bb_obs, crp_obs

    elseif rule isa ApproxRule{C,C} # If it's ApproxRule, we care about it, i.e. it's not a preterminal. We may still sample a BaseRule from this.
        sampled_rule = sample(rule)

        if sampled_rule isa BaseRule{C,C}
            # TODO: -Make a fragment of current contiguous tree fragment and bookkeep (crp_obs vector or something)
            #       -Flip a coin to decide if this fragment's tree will continue being the current contiguous tree fragment
            push!(dm_obs, (lhs(rule), sampled_rule))
            tree = TreeNode(lhs(rule))

            for child in parsetree.children
                child_cat = lhs(child.data[2])
                ptr_child, dm_obs_child, bb_obs_child, crp_obs_child = sample(child, fg)

                if ptr_child != nothing
                    append!(dm_obs, dm_obs_child)
                    append!(bb_obs, bb_obs_child)
                    append!(crp_obs, crp_obs_child)

                    bbidx = (lhs(rule), sampled_rule, child_cat)
                    extend = sample(fg.BB[bbidx])
                    push!(bb_obs, bbidx => extend)

                    if extend
                        ptr_child_clone = clone(ptr_child) # Keeps references to the same other Pointer objects from variable nodes, but creates a clone of the underlying Fragment tree and variables/leaves Vectors.
                        frag_child = ptr_child.fragment
                        tree_child = frag_child.tree
                        insert_child!(tree, tree_child)
                        append!(variables, frag_child.variables)
                        append!(leaves, frag_child.leaves)
                        merge!(children, ptr_child_clone.children)
                    else
                        variable_tree = TreeNode(child_cat)
                        insert_child!(tree, variable_tree)
                        push!(variables, variable_tree)
                        push!(leaves, variable_tree)
                        push!(children, variable_tree => ptr_child)
                    end
                else
                    # TODO: We're at a terminal rule child and are making a new fragment, so add preterminal tree node as leaf of fragment.
                    tree_child = TreeNode(child_cat)
                    insert_child!(tree, tree_child)
                    push!(leaves, tree_child)
                end
            end

            frag = Fragment(tree, variables, leaves)
            push!(crp_obs, frag)
            return Pointer(frag, children), dm_obs, bb_obs, crp_obs

        elseif sampled_rule isa FragmentRule{C,C}
            push!(crp_obs, sampled_rule.fragment)
            for (i, child) in enumerate(parsetree.children)
                leaf_node = sampled_rule.fragment.leaves[i]
                if leaf_node in sampled_rule.fragment.variables # For each of these recurse. If return != nothing, do the logical stuff, etc.
                    ptr_child, dm_obs_child, bb_obs_child, crp_obs_child = sample(child, fg)
                    append!(dm_obs, dm_obs_child)
                    append!(bb_obs, bb_obs_child)
                    append!(crp_obs, crp_obs_child)
                    push!(children, leaf_node => ptr_child)
                end
            end
            return Pointer(sampled_rule.fragment, children), dm_obs, bb_obs, crp_obs
        end
    end
end

approx_rule(r::Tuple{Cons,AbstractRule}) where Cons <: Constituent = r[2]

function logscore(approx_tree::Tree{Tuple{Cons,AbstractRule{C,C}}}, fg::FragmentGrammar{C,CR,T,TR}) where {Cons<:Constituent, C, CR, T, TR}
    reduce(*, approxprob.(approx_rule.(data.(approx_tree))))
end

run_chartparser(input::Tree, fg::FragmentGrammar) = run_chartparser(leaf_data(input), fg, boolean_dependency_dict(input))

function boolean_dependency_dict(tree::Tree{C}) where C
    dependency_dict = Dict{Tuple{C,Int,Int},Bool}()
    function helper(tree, dependency_dict::Dict{Tuple{C,Int,Int}, Bool}, dot::Ref{Int})
        start = dot[]
        if isempty(tree.children)
            dot[] += 1
        else
            for child in tree.children
                helper(child, dependency_dict, dot)
            end
        end
        push!(dependency_dict, (tree.data, start, dot[])=>true)
    end
    helper(tree, dependency_dict, Ref(1))
    # @show dependency_dict
    return dependency_dict
end

function run_mcmc!(fg::FragmentGrammar, inputs::Vector{Tree}, sweeps=1)
    update_approx_probs!(fg)
    analyses = init_mcmc(fg, inputs)
    add_obs!.([fg], getindex.(values(analyses),1))
    update_approx_probs!(fg)

    # local curr_prob::LogProb

    for sweep in 1:sweeps
        for input_tree in keys(analyses)
            curr_prob = logscore(fg)
            curr_analysis, curr_approx_tree = analyses[input_tree]
            rm_obs!(fg, curr_analysis)

            update_approx_probs!(fg) # F-f(i), needed for logscore(::Tree{Tuple{Cons,AbstractRule{C,C}}}, fg::FragmentGrammar) as well.

            proposal_approx_tree = sample_tree(run_chartparser(input_tree, fg))
            proposal_analysis = Analysis(sample(proposal_approx_tree, fg)...)
            forward_trans_prob = logscore(proposal_approx_tree, fg)
            reverse_trans_prob = logscore(curr_approx_tree, fg)

            add_obs!(fg, proposal_analysis) # No need to update approx probs, not going to use them here.

            new_prob = logscore(fg)

            accept_prob = min(one(LogProb), new_prob/curr_prob * reverse_trans_prob/forward_trans_prob)
            if flip(float(accept_prob)) == 0 # If we reject.
                @show "Reject", accept_prob
                rm_obs!(fg, proposal_analysis)
                add_obs!(fg, curr_analysis)
                @show curr_prob.log
            else # If we accept.
                @show "Accept", accept_prob
                analyses[input_tree] = (proposal_analysis, proposal_approx_tree)
                @show new_prob.log
            end
            update_approx_probs!(fg)
        end
    end
    return analyses
end

function init_mcmc(fg::FragmentGrammar, inputs::Vector{Tree})
    analyses = Dict{Tree,Tuple{Analysis,Tree}}()
    for input_tree in inputs
        analysis = Analysis(initialize_fragments(input_tree, fg)...)
        approx_tree = sample_tree(run_chartparser(input_tree, fg)) # This should always be BaseRules
        analyses[input_tree] = (analysis, approx_tree)
    end
    return analyses
end

function initialize_fragments(tree::Tree, fg::FragmentGrammar{C,CR}) where {C,CR}
    variables = Tree{C}[]
    leaves = Tree{C}[]
    children = Dict{Tree, Pointer}()

    dm_obs = Tuple{C,CR}[]
    bb_obs = Pair{Tuple{C, CR, C}, Bool}[]
    crp_obs = Fragment[]

    l = tree.data

    working_tree = TreeNode(l)

    # FIND BASERULE USED
    st = startstate(fg)
    for child in tree.children
        # I can assume that the FG only has base rules in its states right now.
        st = transition(st, child.data)
        insert_child!(working_tree, TreeNode(child.data))
    end
    ar = completions(st)[1][2] # Completion is Tuple{C, CR}
    # ar had sure as hell better be a BaseRule at this point.
    r = ar.rules[1] # There should seriously only be one.
    push!(dm_obs, (lhs(r), r))

    # RECURSE FOR FRAGMENTS, CHECK IF TERMINAL, ETC.
    for (i, child) in enumerate(working_tree.children)
        variable_tree = child
        push!(leaves, variable_tree)

        if !(child.data in preterminals(fg))
            push!(variables, variable_tree)
            bbidx = (lhs(r), r, child.data)
            push!(bb_obs, bbidx => false)

            ptr_child, dm_obs_child, bb_obs_child, crp_obs_child = initialize_fragments(tree.children[i], fg)
            append!(dm_obs, dm_obs_child)
            append!(bb_obs, bb_obs_child)
            append!(crp_obs, crp_obs_child)

            push!(children, variable_tree => ptr_child)
        end
    end

    frag = Fragment(working_tree, variables, leaves)
    push!(crp_obs, frag)

    return Pointer(frag, children), dm_obs, bb_obs, crp_obs
end

end


module DiffEqGPU

using GPUifyLoops, CuArrays, CUDAnative, DiffEqBase, LinearAlgebra, Distributed
using CUDAdrv: CuPtr, CU_NULL, Mem, CuDefaultStream
using CuArrays: CUBLAS

function gpu_kernel(f,du,u,p,t)
    @loop for i in (1:size(u,2); (blockIdx().x-1) * blockDim().x + threadIdx().x)
        @views @inbounds f(du[:,i],u[:,i],p[:,i],t)
        nothing
    end
    nothing
end

function jac_kernel(f,J,u,p,t)
    @loop for i in (0:(size(u,2)-1); (blockIdx().x-1) * blockDim().x + threadIdx().x - 1)
        section = 1 + (i*size(u,1)) : ((i+1)*size(u,1))
        @views @inbounds f(J[section,section],u[:,i+1],p[:,i+1],t)
        nothing
    end
    nothing
end

function discrete_condition_kernel(condition,cur,u,t,p)
    @loop for i in (1:size(u,2); (blockIdx().x-1) * blockDim().x + threadIdx().x)
        @views @inbounds cur[i] = condition(u[:,i],t,FakeIntegrator(u[:,i],t,p[:,i]))
        nothing
    end
    nothing
end

function discrete_affect!_kernel(affect!,cur,u,t,p)
    @loop for i in (1:size(u,2); (blockIdx().x-1) * blockDim().x + threadIdx().x)
        @views @inbounds cur[i] && affect!(FakeIntegrator(u[:,i],t,p[:,i]))
        nothing
    end
    nothing
end

function continuous_condition_kernel(condition,out,u,t,p)
    @loop for i in (1:size(u,2); (blockIdx().x-1) * blockDim().x + threadIdx().x)
        @views @inbounds out[i] = condition(u[:,i],t,FakeIntegrator(u[:,i],t,p[:,i]))
        nothing
    end
    nothing
end

function continuous_affect!_kernel(affect!,event_idx,u,t,p)
    @loop for i in ((event_idx,); (blockIdx().x-1) * blockDim().x + threadIdx().x)
        @views @inbounds affect!(FakeIntegrator(u[:,i],t,p[:,i]))
        nothing
    end
    nothing
end

function GPUifyLoops.launch_config(::Union{typeof(gpu_kernel),
                                           typeof(jac_kernel),
                                           typeof(discrete_condition_kernel),
                                           typeof(discrete_affect!_kernel),
                                           typeof(continuous_condition_kernel),
                                           typeof(continuous_affect!_kernel)},
                                           maxthreads,context,g,f,du,u,args...;
                                           kwargs...)
    t = min(maxthreads,size(u,2))
    blocks = ceil(Int,size(u,2)/t)
    (threads=t,blocks=blocks)
end

function W_kernel(jac, W, u, p, gamma, t)
    len = size(u,1)
    @loop for i in (1:size(u,2); (blockIdx().x-1) * blockDim().x + threadIdx().x)
        _W = @inbounds @view(W[:, :, i])
        @views @inbounds jac(_W,u[:,i],p[:,i],t)
        @inbounds for i in eachindex(_W)
            _W[i] = gamma*_W[i]
        end
        _one = one(eltype(_W))
        @inbounds for i in 1:len
            _W[i, i] = _W[i, i] - _one
        end
        nothing
    end
    return nothing
end

function Wt_kernel(jac, W, u, p, gamma, t)
    len = size(u,1)
    @loop for i in (1:size(u,2); (blockIdx().x-1) * blockDim().x + threadIdx().x)
        _W = @inbounds @view(W[:, :, i])
        @views @inbounds jac(_W,u[:,i],p[:,i],t)
        @inbounds for i in 1:len
            _W[i, i] = -inv(gamma) + _W[i, i]
        end
        nothing
    end
    return nothing
end

function GPUifyLoops.launch_config(::Union{typeof(W_kernel),typeof(Wt_kernel)},
                                           maxthreads,context,g,jac,W,u,args...;
                                           kwargs...)
    t = min(maxthreads,size(u,2))
    blocks = ceil(Int,size(u,2)/t)
    (threads=t,blocks=blocks)
end

function cuda_lufact!(W)
    CuArrays.CUBLAS.getrf_strided_batched!(W, false)
    return nothing
end

function cpu_lufact!(W)
    len = size(W, 1)
    for i in 1:size(W,3)
        _W = @inbounds @view(W[:, :, i])
        generic_lufact!(_W, len)
    end
    return nothing
end

struct FakeIntegrator{uType,tType,P}
    u::uType
    t::tType
    p::P
end

abstract type EnsembleArrayAlgorithm <: DiffEqBase.EnsembleAlgorithm end
struct EnsembleCPUArray <: EnsembleArrayAlgorithm end
struct EnsembleGPUArray <: EnsembleArrayAlgorithm end

function DiffEqBase.__solve(ensembleprob::DiffEqBase.AbstractEnsembleProblem,
                 alg::Union{DiffEqBase.DEAlgorithm,Nothing},
                 ensemblealg::EnsembleArrayAlgorithm;
                 trajectories, batch_size = trajectories, kwargs...)

    num_batches = trajectories ÷ batch_size
    num_batches * batch_size != trajectories && (num_batches += 1)

    if num_batches == 1 && ensembleprob.reduction === DiffEqBase.DEFAULT_REDUCTION
       time = @elapsed sol = batch_solve(ensembleprob,alg,ensemblealg,1:trajectories;kwargs...)
       return DiffEqBase.EnsembleSolution(sol,time,true)
    end

    converged::Bool = false
    u = ensembleprob.u_init === nothing ? similar(batch_solve(ensembleprob,alg,ensemblealg,1:batch_size;kwargs...), 0) : ensembleprob.u_init

    if nprocs() == 1
        # While pmap works, this makes much better error messages.
        time = @elapsed begin
            sols = map(1:num_batches) do i
                if i == num_batches
                  I = (batch_size*(i-1)+1):trajectories
                else
                  I = (batch_size*(i-1)+1):batch_size*i
                end
                batch_data = batch_solve(ensembleprob,alg,ensemblealg,I;kwargs...)
                if ensembleprob.reduction !== DiffEqBase.DEFAULT_REDUCTION
                  u, _ = ensembleprob.reduction(u,batch_data,I)
                  return u
                else
                  batch_data
                end
            end
        end
    else
        time = @elapsed begin
            sols = pmap(1:num_batches) do i
                if i == num_batches
                  I = (batch_size*(i-1)+1):trajectories
                else
                  I = (batch_size*(i-1)+1):batch_size*i
                end
                x = batch_solve(ensembleprob,alg,ensemblealg,I;kwargs...)
                yield()
                if ensembleprob.reduction !== DiffEqBase.DEFAULT_REDUCTION
                  u, _ = ensembleprob.reduction(u,x,I)
                else
                  x
                end

            end
        end
    end

    if ensembleprob.reduction === DiffEqBase.DEFAULT_REDUCTION
        DiffEqBase.EnsembleSolution(hcat(sols...),time,true)
    else
        DiffEqBase.EnsembleSolution(sols[end], time, true)
    end
end

function batch_solve(ensembleprob,alg,ensemblealg,I;kwargs...)
    probs = [ensembleprob.prob_func(deepcopy(ensembleprob.prob),i,1) for i in I]
    @assert all(p->p.tspan == probs[1].tspan,probs)
    @assert !isempty(I)
    #@assert all(p->p.f === probs[1].f,probs)

    len = length(probs[1].u0)
    if ensemblealg isa EnsembleGPUArray
        # it's 1:length(I) since probs is generated above with for i in I
        u0 = CuArray(hcat([probs[i].u0 for i in 1:length(I)]...))
        p  = CuArray(hcat([probs[i].p  for i in 1:length(I)]...))
    elseif ensemblealg isa EnsembleCPUArray
        u0 = hcat([probs[i].u0 for i in 1:length(I)]...)
        p  = hcat([probs[i].p  for i in 1:length(I)]...)
    end

    if DiffEqBase.has_jac(probs[1].f)
        jac_prototype = ensemblealg isa EnsembleGPUArray ?
                        cu(zeros(Float32,len,len,length(I))) :
                        zeros(Float32,len,len,length(I))
        if probs[1].f.colorvec !== nothing
            colorvec = repeat(probs[1].f.colorvec,length(I))
        else
            colorvec = repeat(1:length(probs[1].u0),length(I))
        end
    else
        jac_prototype = nothing
        colorvec = nothing
    end

    _callback = generate_callback(probs[1],length(I),ensemblealg)
    prob = generate_problem(probs[1],u0,p,jac_prototype,colorvec)

    internalnorm(u::CuArray,t) = sqrt(maximum(sum(abs2, u, dims=1)/size(u, 1)))
    internalnorm(u::Union{AbstractFloat,Complex},t) = abs(u)

    sol  = solve(prob,alg; callback = _callback,merge_callbacks = false,
                 #internalnorm=internalnorm,
                 kwargs...)

    us = Array.(sol.u)
    solus = [[@view(us[i][:,j]) for i in 1:length(us)] for j in 1:length(probs)]
    [ensembleprob.output_func(DiffEqBase.build_solution(probs[i],alg,sol.t,solus[i],destats=sol.destats,retcode=sol.retcode),i)[1] for i in 1:length(probs)]
end

function generate_problem(prob::ODEProblem,u0,p,jac_prototype,colorvec)
    _f = let f=prob.f.f
        function (du,u,p,t)
            version = u isa CuArray ? CUDA() : CPU()
            @launch version gpu_kernel(f,du,u,p,t)
        end
    end

    if DiffEqBase.has_jac(prob.f)
        _Wfact! = let jac=prob.f.jac
            function (W,u,p,gamma,t)
                iscuda = u isa CuArray
                version = iscuda ? CUDA() : CPU()
                @launch version W_kernel(jac, W, u, p, gamma, t)
                iscuda ? cuda_lufact!(W) : cpu_lufact!(W)
            end
        end
        _Wfact!_t = let jac=prob.f.jac
            function (W,u,p,gamma,t)
                iscuda = u isa CuArray
                version = iscuda ? CUDA() : CPU()
                @launch version Wt_kernel(jac, W, u, p, gamma, t)
                iscuda ? cuda_lufact!(W) : cpu_lufact!(W)
            end
        end
    else
        _Wfact! = nothing
        _Wfact!_t = nothing
    end

    if DiffEqBase.has_tgrad(prob.f)
        _tgrad = let tgrad=prob.f.tgrad
            function (J,u,p,t)
                version = u isa CuArray ? CUDA() : CPU()
                @launch version gpu_kernel(tgrad,J,u,p,t)
            end
        end
    else
        _tgrad = nothing
    end

    f_func = ODEFunction(_f,Wfact = _Wfact!,
                        Wfact_t = _Wfact!_t,
                        #colorvec=colorvec,
                        jac_prototype = jac_prototype,
                        tgrad=_tgrad)
    prob = ODEProblem(f_func,u0,prob.tspan,p;
                      prob.kwargs...)
end

function generate_problem(prob::SDEProblem,u0,p,jac_prototype,colorvec)
    _f = let f=prob.f.f
        function (du,u,p,t)
            version = u isa CuArray ? CUDA() : CPU()
            @launch version gpu_kernel(f,du,u,p,t)
        end
    end

    _g = let f=prob.f.g
        function (du,u,p,t)
            version = u isa CuArray ? CUDA() : CPU()
            @launch version gpu_kernel(f,du,u,p,t)
        end
    end

    if DiffEqBase.has_jac(prob.f)
        _Wfact! = let jac=prob.f.jac
            function (W,u,p,gamma,t)
                iscuda = u isa CuArray
                version = iscuda ? CUDA() : CPU()
                @launch version W_kernel(jac, W, u, p, gamma, t)
                iscuda ? cuda_lufact!(W) : cpu_lufact!(W)
            end
        end
        _Wfact!_t = let jac=prob.f.jac
            function (W,u,p,gamma,t)
                iscuda = u isa CuArray
                version = iscuda ? CUDA() : CPU()
                @launch version Wt_kernel(jac, W, u, p, gamma, t)
                iscuda ? cuda_lufact!(W) : cpu_lufact!(W)
            end
        end
    else
        _Wfact! = nothing
        _Wfact!_t = nothing
    end

    if DiffEqBase.has_tgrad(prob.f)
        _tgrad = let tgrad=prob.f.tgrad
            function (J,u,p,t)
                version = u isa CuArray ? CUDA() : CPU()
                @launch version gpu_kernel(tgrad,J,u,p,t)
            end
        end
    else
        _tgrad = nothing
    end

    f_func = SDEFunction(_f,_g,Wfact = _Wfact!,
                        Wfact_t = _Wfact!_t,
                        #colorvec=colorvec,
                        jac_prototype = jac_prototype,
                        tgrad=_tgrad)
    prob = SDEProblem(f_func,_g,u0,prob.tspan,p;
                      prob.kwargs...)
end

function generate_callback(prob,I,ensemblealg)
    if :callback ∉ keys(prob.kwargs)
        _callback = nothing
    elseif prob.kwargs[:callback] isa DiscreteCallback
        if ensemblealg isa EnsembleGPUArray
            cur = CuArray([false for i in 1:I])
        else
            cur = [false for i in 1:I]
        end
        _condition = prob.kwargs[:callback].condition
        _affect!   = prob.kwargs[:callback].affect!

        condition = function (u,t,integrator)
            version = u isa CuArray ? CUDA() : CPU()
            @launch version discrete_condition_kernel(_condition,cur,u,t,integrator.p)
            any(cur)
        end

        affect! = function (integrator)
            version = integrator.u isa CuArray ? CUDA() : CPU()
            @launch version discrete_affect!_kernel(_affect!,cur,integrator.u,integrator.t,integrator.p)
        end

        _callback = DiscreteCallback(condition,affect!,save_positions=prob.kwargs[:callback].save_positions)
    elseif prob.kwargs[:callback] isa ContinuousCallback
        _condition   = prob.kwargs[:callback].condition
        _affect!     = prob.kwargs[:callback].affect!
        _affect_neg! = prob.kwargs[:callback].affect_neg!

        condition = function (out,u,t,integrator)
            version = u isa CuArray ? CUDA() : CPU()
            @launch version continuous_condition_kernel(_condition,out,u,t,integrator.p)
            nothing
        end

        affect! = function (integrator,event_idx)
            version = integrator.u isa CuArray ? CUDA() : CPU()
            @launch version continuous_affect!_kernel(_affect!,event_idx,integrator.u,integrator.t,integrator.p)
        end

        affect_neg! = function (integrator,event_idx)
            version = integrator.u isa CuArray ? CUDA() : CPU()
            @launch version continuous_affect!_kernel(_affect_neg!,event_idx,integrator.u,integrator.t,integrator.p)
        end

        _callback = VectorContinuousCallback(condition,affect!,affect_neg!,I,save_positions=prob.kwargs[:callback].save_positions)
    end
    _callback
end

### GPU Factorization

struct LinSolveGPUSplitFactorize
    len::Int
    nfacts::Int
end
LinSolveGPUSplitFactorize() = LinSolveGPUSplitFactorize(0, 0)

function (p::LinSolveGPUSplitFactorize)(x,A,b,update_matrix=false;kwargs...)
    version = b isa CuArray ? CUDA() : CPU()
    copyto!(x, b)
    @launch version ldiv!_kernel(A,x,p.len,p.nfacts)
    return nothing
end

function (p::LinSolveGPUSplitFactorize)(::Type{Val{:init}},f,u0_prototype)
    LinSolveGPUSplitFactorize(size(u0_prototype)...)
end

function ldiv!_kernel(W,u,len,nfacts)
    @loop for i in (1:nfacts; (blockIdx().x-1) * blockDim().x + threadIdx().x)
        section = 1 + ((i-1)*len) : (i*len)
        _W = @inbounds @view(W[:, :, i])
        _u = @inbounds @view u[section]
        naivesolve!(_W, _u, len)
        nothing
    end
    return nothing
end

function GPUifyLoops.launch_config(::typeof(ldiv!_kernel),
                                           maxthreads,context,g,W,x,len,nfacts,
                                           args...;
                                           kwargs...)
    t = min(maxthreads,nfacts)
    blocks = ceil(Int,nfacts/t)
    (threads=t,blocks=blocks)
end

function __printjac(A, ii)
    @cuprintf "[%d, %d]\n" ii.offset1 ii.stride2
    @cuprintf "%d %d %d\n%d %d %d\n%d %d %d\n" ii[1, 1] ii[1, 2] ii[1, 3] ii[2, 1] ii[2, 2] ii[2, 3] ii[3, 1] ii[3, 2] ii[3, 3]
    @cuprintf "%2.2f %2.2f %2.2f\n%2.2f %2.2f %2.2f\n%2.2f %2.2f %2.2f" A[ii[1, 1]] A[ii[1, 2]] A[ii[1, 3]] A[ii[2, 1]] A[ii[2, 2]] A[ii[2, 3]] A[ii[3, 1]] A[ii[3, 2]] A[ii[3, 3]]
end

function generic_lufact!(A::AbstractMatrix{T}, minmn) where {T}
    m = n = minmn
    #@cuprintf "\n\nbefore lufact!\n"
    #__printjac(A, ii)
    #@cuprintf "\n"
    @inbounds for k = 1:minmn
        #@cuprintf "inner factorization loop\n"
        # Scale first column
        Akkinv = inv(A[k,k])
        for i = k+1:m
            #@cuprintf "L\n"
            A[i,k] *= Akkinv
        end
        # Update the rest
        for j = k+1:n, i = k+1:m
            #@cuprintf "U\n"
            A[i,j] -= A[i,k]*A[k,j]
        end
    end
    #@cuprintf "after lufact!"
    #__printjac(A, ii)
    #@cuprintf "\n\n\n"
    return nothing
end

struct MyL{T} # UnitLowerTriangular
    data::T
end
struct MyU{T} # UpperTriangular
    data::T
end

function naivesub!(A::MyU, b::AbstractVector, n)
    x = b
    @inbounds for j in n:-1:1
        xj = x[j] = A.data[j,j] \ b[j]
        for i in j-1:-1:1 # counterintuitively 1:j-1 performs slightly better
            b[i] -= A.data[i,j] * xj
        end
    end
    return nothing
end
function naivesub!(A::MyL, b::AbstractVector, n)
    x = b
    @inbounds for j in 1:n
        xj = x[j]
        for i in j+1:n
            b[i] -= A.data[i,j] * xj
        end
    end
    return nothing
end

function naivesolve!(A::AbstractMatrix, x::AbstractVector, n)
    naivesub!(MyL(A), x, n)
    naivesub!(MyU(A), x, n)
    return nothing
end

export EnsembleCPUArray, EnsembleGPUArray, LinSolveGPUSplitFactorize

end # module


check_int(i::Int, key_path) = i
function check_int(o, key_path)
    try
        convert(Int, o)
    catch
        throw(AssertionError("$o must be an integer [$key_path]"))
    end
end

check_float(f::Float64, key_path) = f
function check_float(o, key_path)
    try
        convert(Float64, o)
    catch
        throw(AssertionError("$o must be a float [$key_path]"))
    end
end

check_boolean(b::Bool, key_path) = b
function check_boolean(o, key_path)
    try
        convert(Bool, o)
    catch
        throw(AssertionError("$o must be a boolean [$key_path]"))
    end
end

function check_array(val::AbstractArray, values::AbstractArray, key_path)
    if !issubset(val, values)
        throw(AssertionError("$val not in $values [$key_path]"))
    end
end
check_array(val, values::AbstractArray, key_path) =
    throw(AssertionError("$val must be an array (subset of $values) [$key_path]"))

function get_values(def, classes)
    if haskey(def, "values")
        return def["values"]
    elseif haskey(def, "class")
        return keys_or_vals(classes[def["class"]])
    else
        throw(AssertionError("array type definitions must specify `values` or `class`"))
    end
end

"""
Checks if an item meets the requirements of the field's definition.
"""
function check_item(val, def, key_path; keys = nothing, classes = nothing)
    # check type of value
    dtype = def["type"]
    if dtype == "int"
        val = check_int(val, key_path)
    elseif dtype == "float"
        val = check_float(val, key_path)
    elseif dtype == "boolean"
        val = check_boolean(val, key_path)
    elseif dtype == "enum"
        values = get_values(def, classes)
        @assert val in values "$val not in $values [$key_path]"
    elseif dtype == "array"
        values = get_values(def, classes)
        check_array(val, values, key_path)
    elseif dtype == "keys"
        check_array(val, keys, key_path)
    end

    # check range
    if haskey(def, "min")
        @assert val >= def["min"] "$val must be greater than $(def["min"]) [$key_path]"
    end
    if haskey(def, "max")
        @assert val <= def["max"] "$val must be less than $(def["max"]) [$key_path]"
    end

    return val
end


## Hashing

equality_data(g::AbstractFPGroupElement) = (normalform!(g); word(g))

bitget(h::UInt, n::Int) = Bool((h & (1 << n)) >> n)
bitclear(h::UInt, n::Int) = h & ~(1 << n)
bitset(h::UInt, n::Int) = h | (1 << n)
bitset(h::UInt, v::Bool, n::Int) = v ? bitset(h, n) : bitclear(h, n)

# We store hash of a word in field `savedhash` to use it as cheap proxy to
# determine non-equal words. Additionally bits of `savehash` store boolean
# information as follows
# * `savedhash & 1` (the first bit): is word in normal form?
# * `savedhash & 2` (the second bit): is the hash valid?
const __BITFLAGS_MASK = ~(~(UInt(0)) << 2)

isnormalform(g::AbstractFPGroupElement) = bitget(g.savedhash, 0)
_isvalidhash(g::AbstractFPGroupElement) = bitget(g.savedhash, 1)

_setnormalform(h::UInt, v::Bool) = bitset(h, v, 0)
_setvalidhash(h::UInt, v::Bool) = bitset(h, v, 1)

_setnormalform!(g::AbstractFPGroupElement, v::Bool) = g.savedhash = _setnormalform(g.savedhash, v)
_setvalidhash!(g::AbstractFPGroupElement, v::Bool) = g.savedhash = _setvalidhash(g.savedhash, v)

# To update hash use this internal method, possibly only after computing the
# normal form of `g`:
function _update_savedhash!(g::AbstractFPGroupElement, data)
    h = hash(data, hash(parent(g)))
    h = (h << count_ones(__BITFLAGS_MASK)) | (__BITFLAGS_MASK & g.savedhash)
    g.savedhash = _setvalidhash(h, true)
    return g
end

function Base.hash(g::AbstractFPGroupElement, h::UInt)
    _isvalidhash(g) || _update_savedhash!(g, equality_data(g))
    return hash(g.savedhash >> count_ones(__BITFLAGS_MASK), h)
end

function Base.copyto!(res::AbstractFPGroupElement, g::AbstractFPGroupElement)
    @assert parent(res) === parent(g)
    resize!(word(res), length(word(g)))
    copyto!(word(res), word(g))
    res.savedhash = g.savedhash
    return res
end


# Bar geometry summarizes data as vertical bars.
immutable BarGeometry <: Gadfly.GeometryElement
    # How bars should be positioned if they are grouped by color.
    # Valid options are:
    #   :stack -> place bars on top of each other (default)
    #   :dodge -> place bar next to each other
    position::Symbol

    # :vertical or :horizontal
    orientation::Symbol

    default_statistic::Gadfly.StatisticElement

    function BarGeometry(default_statistic=Gadfly.Stat.identity();
                         position::Symbol=:stack, orientation::Symbol=:vertical)
        new(position, orientation, default_statistic)
    end
end

const bar = BarGeometry

function histogram(; position=:stack, bincount=nothing,
                   minbincount=3, maxbincount=150,
                   orientation::Symbol=:vertical,
                   density::Bool=false)
    BarGeometry(Gadfly.Stat.histogram(bincount=bincount,
                                      minbincount=minbincount,
                                      maxbincount=maxbincount,
                                      position=position,
                                      orientation=orientation,
                                      density=density),
                position=position,
                orientation=orientation)
end

function element_aesthetics(::BarGeometry)
    [:x, :xmin, :xmax, :y, :ymin, :ymax, :color]
end

function default_statistic(geom::BarGeometry)
    return geom.default_statistic
end

# Render a single color bar chart
function render_colorless_bar(geom::BarGeometry,
                              theme::Gadfly.Theme,
                              aes::Gadfly.Aesthetics,
                              orientation::Symbol)
    if orientation == :horizontal
        ctx = compose!(
            context(),
            rectangle([0.0],
                      [Measure(cy=ymin) + theme.bar_spacing/2 for ymin in aes.ymin],
                      aes.x,
                      [Measure(cy=(ymax - ymin)) - theme.bar_spacing
                       for (ymin, ymax) in zip(aes.ymin, aes.ymax)]),
            svgclass("geometry"))
    else
        ctx = compose!(
            context(),
            rectangle([Measure(cx=xmin) + theme.bar_spacing/2 for xmin in aes.xmin],
                      [0.0],
                      [Measure(cx=(xmax - xmin)) - theme.bar_spacing
                       for (xmin, xmax) in zip(aes.xmin, aes.xmax)],
                      aes.y),
            svgclass("geometry"))
    end

    compose!(ctx, fill(theme.default_color))
    if isa(theme.bar_highlight, Function)
        compose!(ctx, stroke(theme.bar_highlight(theme.default_color)))
    elseif isa(theme.bar_highlight, ColorValue)
        compose!(ctx, stroke(theme.bar_highlight))
    else
        compose!(ctx, stroke(nothing))
    end
    return ctx
end

# Render a bar chart grouped by discrete colors and stacked.
function render_colorful_stacked_bar(geom::BarGeometry,
                                     theme::Gadfly.Theme,
                                     aes::Gadfly.Aesthetics,
                                     orientation::Symbol)
    # preserve factor orders of pooled data arrays
    if isa(aes.color, PooledDataArray)
        idxs = sortperm(aes.color.refs, rev=true)
    else
        idxs = 1:length(aes.color)
    end

    ctx = context()
    if orientation == :horizontal
        stack_height_dict = Dict()
        for y in aes.ymin
            stack_height_dict[y] = zero(eltype(aes.x))
        end
        stack_height = zeros(eltype(aes.x), length(idxs))

        for (i, j) in enumerate(idxs)
            stack_height[i] = stack_height_dict[aes.ymin[j]]
            stack_height_dict[aes.ymin[j]] += aes.x[j]
        end

        compose!(
            ctx,
            rectangle(
                stack_height,
                [aes.ymin[i]*cy - theme.bar_spacing/2 for i in idxs],
                [aes.x[i] for i in idxs],
                [(aes.ymax[i] - aes.ymin[i])*cy + theme.bar_spacing for i in idxs]))
    elseif orientation == :vertical
        stack_height_dict = Dict()
        for x in aes.xmin
            stack_height_dict[x] = zero(eltype(aes.y))
        end
        stack_height = zeros(eltype(aes.y), length(idxs))

        for (i, j) in enumerate(idxs)
            stack_height[i] = stack_height_dict[aes.xmin[j]]
            stack_height_dict[aes.xmin[j]] += aes.y[j]
        end

        compose!(
            ctx,
            rectangle(
                [aes.xmin[i]*cx + theme.bar_spacing/2 for i in idxs],
                stack_height,
                [(aes.xmax[i] - aes.xmin[i])*cx - theme.bar_spacing for i in idxs],
                [aes.y[i] for i in idxs]))
    else
        error("Orientation must be :horizontal or :vertical")
    end

    cs = [aes.color[i] for i in idxs]
    compose!(ctx, fill(cs), svgclass("geometry"))
    if isa(theme.bar_highlight, Function)
        compose!(ctx, stroke([theme.bar_highlight(c) for c in cs]))
    elseif isa(theme.bar_highlight, ColorValue)
        compose!(ctx, stroke(theme.bar_highlight))
    else
        compose!(ctx, stroke(nothing))
    end
    return ctx
end

# Render a bar chart grouped by discrete colors and stuck next to each other.
function render_colorful_dodged_bar(geom::BarGeometry,
                                    theme::Gadfly.Theme,
                                    aes::Gadfly.Aesthetics,
                                    orientation::Symbol)
    # preserve factor orders of pooled data arrays
    if isa(aes.color, PooledDataArray)
        idxs = sortperm(aes.color.refs, rev=true)
    else
        idxs = 1:length(aes.color)
    end

    ctx = context()
    if orientation == :horizontal
        dodge_count = DefaultDict(() -> 0)
        for i in idxs
            dodge_count[aes.ymin[i]] += 1
        end

        dodge_height = Dict()
        dodge_pos_dict = DefaultDict(() -> 0cy)
        for i in idxs
            dodge_height[aes.ymin[i]] =
                ((aes.ymax[i] - aes.ymin[i]) / dodge_count[aes.ymin[i]]) * cy
            dodge_pos_dict[aes.ymin[i]] = aes.ymin[i]*cy
        end

        dodge_pos = Array(Measure, length(idxs))
        for (i, j) in enumerate(idxs)
            dodge_pos[i] = dodge_pos_dict[aes.ymin[j]] - theme.bar_spacing/2
            dodge_pos_dict[aes.ymin[j]] += dodge_height[aes.ymin[j]]
        end

        compose!(
            ctx,
            rectangle(
                [0.0],
                dodge_pos,
                [aes.x[i] for i in idxs],
                [((aes.ymax[i] - aes.ymin[i])*cy + theme.bar_spacing) / dodge_count[aes.ymin[i]]
                 for i in idxs]))
    elseif orientation == :vertical
        dodge_count = DefaultDict(() -> 0)
        for i in idxs
            dodge_count[aes.xmin[i]] += 1
        end

        dodge_width = Dict()
        dodge_pos_dict = DefaultDict(() -> 0cx)
        for i in idxs
            dodge_width[aes.xmin[i]] =
                ((aes.xmax[i] - aes.xmin[i]) / dodge_count[aes.xmin[i]]) * cx
            dodge_pos_dict[aes.xmin[i]] = aes.xmin[i]*cx
        end

        dodge_pos = Array(Measure, length(idxs))
        for (i, j) in enumerate(idxs)
            dodge_pos[i] = dodge_pos_dict[aes.xmin[j]] + theme.bar_spacing/2
            dodge_pos_dict[aes.xmin[j]] += dodge_width[aes.xmin[j]]
        end

        compose!(
            ctx,
            rectangle(
                dodge_pos,
                [0.0],
                [((aes.xmax[i] - aes.xmin[i])*cx - theme.bar_spacing) / dodge_count[aes.xmin[i]]
                 for i in idxs],
                [aes.y[i] for i in idxs]))
    else
        error("Orientation must be :horizontal or :vertical")
    end

    cs = [aes.color[i] for i in idxs]
    compose!(ctx, fill(cs), svgclass("geometry"))
    if isa(theme.bar_highlight, Function)
        compose!(ctx, stroke([theme.bar_highlight(c) for c in cs]))
    elseif isa(theme.bar_highlight, ColorValue)
        compose!(ctx, stroke(theme.bar_highlight))
    else
        compose!(ctx, stroke(nothing))
    end
    return ctx
end

# Render bar geometry.
#
# Args:
#   geom: bar geometry
#   theme: the plot's theme
#   aes: some aesthetics
#
# Returns
#   A compose form.
#
function render(geom::BarGeometry, theme::Gadfly.Theme, aes::Gadfly.Aesthetics)
    if geom.orientation == :horizontal
        if (is(aes.ymin, nothing) || is(aes.ymax, nothing)) && is(aes.y, nothing)
            error("Geom.bar required \"y\" to be bound or both \"y_min\" and \"y_max\".")
        end
        if (aes.y != nothing && length(aes.y) != length(aes.x)) ||
           (aes.ymin != nothing && (length(aes.ymin) != length(aes.x) || length(aes.ymax) != length(aes.x)))
            error("Geom.bar requires x and y to be of equal length.")
        end

        var = :y
        minvar = :ymin
        maxvar = :ymax
    else
        if (is(aes.xmin, nothing) || is(aes.xmax, nothing)) && is(aes.x, nothing)
            error("Geom.bar required \"x\" to be bound or both \"x_min\" and \"x_max\".")
        end
        if (aes.x != nothing && length(aes.x) != length(aes.y)) &&
           (aes.xmin != nothing && (length(aes.xmin) != length(aes.y) || length(aes.xmax) != length(aes.y)))
            error("Geom.bar requires x and y to be of equal length.")
        end
        var = :x
        minvar = :xmin
        maxvar = :xmax
    end

    if getfield(aes, minvar) === nothing
        aes2 = Gadfly.Aesthetics()
        values = getfield(aes, var)
        minvalue, maxvalue = minimum(values), maximum(values)
        T = typeof((maxvalue - minvalue) / 1.0)

        span = zero(T)
        unique_count = length(Set(values))
        if unique_count > 1
            span = (maximum(values) - minimum(values)) / convert(Float64, (unique_count - 1))
        end

        if span == zero(T)
            span = one(T)
        end

        T = promote_type(eltype(values), typeof(span/2.0))
        setfield!(aes2, minvar, Array(T, length(values)))
        setfield!(aes2, maxvar, Array(T, length(values)))

        for (i, x) in enumerate(values)
            getfield(aes2, minvar)[i] = x - span/2.0
            getfield(aes2, maxvar)[i] = x + span/2.0
        end
        aes = inherit(aes, aes2)
    end

    if aes.color === nothing
        ctx = render_colorless_bar(geom, theme, aes, geom.orientation)
    elseif geom.position == :stack
        ctx = render_colorful_stacked_bar(geom, theme, aes, geom.orientation)
    elseif geom.position == :dodge
        ctx = render_colorful_dodged_bar(geom, theme, aes, geom.orientation)
    else
        error("$(geom.position) is not a valid position for the bar geometry.")
    end

    return compose!(
        context(),
        ctx,
        linewidth(theme.highlight_width),
        svgattribute("shape-rendering", "crispEdges"))
end


#=

Idea is: we have a model formula and a design dict of variables with reference values.
all teh terms that are not present in the design get "typified".  we re-write the model
formula replacing all the terms with their "typified" version where available.

then we call modelcols using this new formula + the reference grid

=#

struct TypicalTerm{T,V} <: AbstractTerm
    term::T
    values::V
    function TypicalTerm(t::T, v::V) where {T<:AbstractTerm,V}
        n = length(v)
        m = width(t)
        if n != m
            throw(ArgumentError("mismatch between number of typical values ($n) " *
                                "and number of columns generated by term $t ($m)"))
        end
        return new{T,V}(t, v)
    end
end

function StatsModels.modelcols(t::TypicalTerm, d::NamedTuple)
    cols = ones(length(first(d)), width(t))
    for (idx, v) in enumerate(t.values)
        cols[:, idx] .*= v
    end
    return cols
end

StatsBase.coefnames(t::TypicalTerm) = coefnames(t.term)
# regular show
Base.show(io::IO, t::TypicalTerm) = show(io, t.term)
# long show
Base.show(io::IO, m::MIME"text/plain", t::TypicalTerm) = show(io, m, t.term)

# statsmodels glue code:
StatsModels.width(t::TypicalTerm) = StatsModels.width(t.term)
# don't generate schema entries for terms which are already typified
StatsModels.needs_schema(::TypicalTerm) = false
StatsModels.termsyms(t::TypicalTerm) = StatsModels.termsyms(t.term)

function get_matrix_term(x)
    x = StatsModels.collect_matrix_terms(x)
    x = x isa MatrixTerm ? x : first(x)
    x isa MatrixTerm || throw(ArgumentError("couldn't extract matrix term from $x"))
    if first(x.terms) isa MatrixTerm
        x = only(x.terms)
    end
    return x
end

function typify(refgrid, model_formula::FormulaTerm,
                model_matrix::AbstractMatrix; typical=mean)
    effects_terms = Term.(Tables.columnnames(refgrid))

    # creates a MatrixTerm (and should work for things like MixedModels) which
    # should correspond to the model_matrix
    matrix_term = get_matrix_term(model_formula.rhs)
    typical_terms = Dict()

    urterms = terms(matrix_term)
    nonfunc_terms = [tt for tt in matrix_term.terms if !isa(tt, FunctionTerm)]
    # only include generating terms that exist outside of a FunctionTerm
    filter!(urterms) do tt
        return any(x -> _termsyms(tt) <= _termsyms(x), nonfunc_terms)
    end

    # we need to include FunctionTerms separately so that we can grab the transformed columns
    func_terms = [filter(tt -> isa(tt, FunctionTerm), matrix_term.terms)...]
    for term in [urterms; func_terms]
        if !any(et -> _symequal(et, term), effects_terms)
            typical_terms[term] = typicalterm(term, matrix_term, model_matrix; typical=typical)
        end
    end

    return _replace(matrix_term, typical_terms)
end

_replace(matrix_term::MatrixTerm, typicals::Dict) = MatrixTerm(_replace.(matrix_term.terms, Ref(typicals)))
_replace(term::AbstractTerm, typicals::Dict) = haskey(typicals, term) ? typicals[term] : term
_replace(term::InteractionTerm, typicals::Dict) = InteractionTerm(_replace.(term.terms, Ref(typicals)))

_termsyms(t) = StatsModels.termsyms(t)
_termsyms(t::FunctionTerm{Fo, Fa, Names}) where {Fo, Fa, Names} = Names
_symequal(t1::AbstractTerm, t2::AbstractTerm) = issetequal(_termsyms(t1), _termsyms(t2))

_trmequal(t1::AbstractTerm, t2::AbstractTerm) = _symequal(t1, t2)
_trmequal(t1::AbstractTerm, t2::FunctionTerm) = false
_trmequal(t1::FunctionTerm, t2::AbstractTerm) = false
function _trmequal(t1::FunctionTerm, t2::FunctionTerm)
    return t1.exorig == t2.exorig &&
        _symequal(t1, t2) &&
        t1.forig == t2.forig
end

function typicalterm(term::AbstractTerm, context::MatrixTerm, model_matrix; typical=mean)
    i = findfirst(t -> _trmequal(t, term), context.terms)
    i === nothing && throw(ArgumentError("Can't determine columns corresponding to '$term' in matrix term $context"))
    cols = (i == 1 ? 0 : sum(width, context.terms[1:(i - 1)])) .+ (1:width(term))
    vals = map(typical, eachcol(view(model_matrix, :, cols)))
    all(v -> length(v) == 1, vals) || throw(ArgumentError("Typical function should return a scalar."))
    return TypicalTerm(term, vals)
end