text:
module NumericalFluxes

export NumericalFluxGradient,
    NumericalFluxFirstOrder,
    NumericalFluxSecondOrder,
    RusanovNumericalFlux,
    RoeNumericalFlux,
    HLLCNumericalFlux,
    RoeNumericalFluxMoist,
    CentralNumericalFluxGradient,
    CentralNumericalFluxFirstOrder,
    CentralNumericalFluxSecondOrder,
    CentralNumericalFluxDivergence,
    CentralNumericalFluxHigherOrder,
    LMARSNumericalFlux

using StaticArrays, LinearAlgebra
using ClimateMachine.VariableTemplates
using KernelAbstractions.Extras: @unroll
using ...BalanceLaws
import ...BalanceLaws:
    vars_state,
    boundary_state!,
    wavespeed,
    flux_first_order!,
    flux_second_order!,
    compute_gradient_flux!,
    compute_gradient_argument!,
    transform_post_gradient_laplacian!,
    boundary_conditions

"""
    NumericalFluxGradient

Any `P <: NumericalFluxGradient` should define methods for:

    numerical_flux_gradient!(
        gnf::P,
        balance_law::BalanceLaw,
        diffF, n⁻,
        Q⁻, Qstate_gradient_flux⁻, Qaux⁻,
        Q⁺, Qstate_gradient_flux⁺, Qaux⁺,
        t
    )

    numerical_boundary_flux_gradient!(
        gnf::P,
        balance_law::BalanceLaw,
        local_state_gradient_flux,
        n⁻,
        local_transform⁻, local_state_prognostic⁻, local_state_auxiliary⁻,
        local_transform⁺, local_state_prognostic⁺, local_state_auxiliary⁺,
        bctype,
        t
    )

"""
abstract type NumericalFluxGradient end

"""
    CentralNumericalFluxGradient <: NumericalFluxGradient

"""
struct CentralNumericalFluxGradient <: NumericalFluxGradient end

function numerical_flux_gradient!(
    ::CentralNumericalFluxGradient,
    balance_law::BalanceLaw,
    transform_gradient::MMatrix,
    normal_vector::AbstractArray,
    state_gradient⁻::AbstractArray,
    state_prognostic⁻::AbstractArray,
    state_auxiliary⁻::AbstractArray,
    state_gradient⁺::AbstractArray,
    state_prognostic⁺::AbstractArray,
    state_auxiliary⁺::AbstractArray,
    t,
)

    transform_gradient .=
        SVector(normal_vector) .* (state_gradient⁺ .+ state_gradient⁻)' ./ 2
end

function numerical_boundary_flux_gradient!(
    numerical_flux::CentralNumericalFluxGradient,
    bctype,
    balance_law::BalanceLaw,
    transform_gradient::MMatrix,
    normal_vector::SVector,
    state_gradient⁻::Vars{T},
    state_prognostic⁻::Vars{S},
    state_auxiliary⁻::Vars{A},
    state_gradient⁺::Vars{T},
    state_prognostic⁺::Vars{S},
    state_auxiliary⁺::Vars{A},
    t,
    state1⁻::Vars{S},
    aux1⁻::Vars{A},
) where {D, T, S, A}
    boundary_state!(
        numerical_flux,
        bctype,
        balance_law,
        state_prognostic⁺,
        state_auxiliary⁺,
        normal_vector,
        state_prognostic⁻,
        state_auxiliary⁻,
        t,
        state1⁻,
        aux1⁻,
    )

    compute_gradient_argument!(
        balance_law,
        state_gradient⁺,
        state_prognostic⁺,
        state_auxiliary⁺,
        t,
    )
    transform_gradient .= normal_vector .* parent(state_gradient⁺)'
end

"""
    NumericalFluxFirstOrder

Any `N <: NumericalFluxFirstOrder` should define the a method for

    numerical_flux_first_order!(
        numerical_flux::N,
        balance_law::BalanceLaw,
        flux,
        normal_vector⁻,
        Q⁻, Qaux⁻,
        Q⁺, Qaux⁺,
        t
    )

where
- `flux` is the numerical flux array
- `normal_vector⁻` is the unit normal
- `Q⁻`/`Q⁺` are the minus/positive state arrays
- `t` is the time

An optional method can also be defined for

    numerical_boundary_flux_first_order!(
        numerical_flux::N,
        balance_law::BalanceLaw,
        flux,
        normal_vector⁻,
        Q⁻, Qaux⁻,
        Q⁺, Qaux⁺,
        bctype, t
    )

"""
abstract type NumericalFluxFirstOrder end

function numerical_flux_first_order! end

function numerical_boundary_flux_first_order!(
    numerical_flux::NumericalFluxFirstOrder,
    bctype,
    balance_law::BalanceLaw,
    fluxᵀn::Vars{S},
    normal_vector::SVector,
    state_prognostic⁻::Vars{S},
    state_auxiliary⁻::Vars{A},
    state_prognostic⁺::Vars{S},
    state_auxiliary⁺::Vars{A},
    t,
    direction,
    state1⁻::Vars{S},
    aux1⁻::Vars{A},
) where {S, A}

    boundary_state!(
        numerical_flux,
        bctype,
        balance_law,
        state_prognostic⁺,
        state_auxiliary⁺,
        normal_vector,
        state_prognostic⁻,
        state_auxiliary⁻,
        t,
        state1⁻,
        aux1⁻,
    )

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

"""
    RusanovNumericalFlux <: NumericalFluxFirstOrder

The RusanovNumericalFlux (aka local Lax-Friedrichs) numerical flux.

# Usage

    RusanovNumericalFlux()

Requires a `flux_first_order!` and `wavespeed` method for the balance law.
"""
struct RusanovNumericalFlux <: NumericalFluxFirstOrder end

update_penalty!(::RusanovNumericalFlux, ::BalanceLaw, _...) = nothing

function numerical_flux_first_order!(
    numerical_flux::RusanovNumericalFlux,
    balance_law::BalanceLaw,
    fluxᵀn::Vars{S},
    normal_vector::SVector,
    state_prognostic⁻::Vars{S},
    state_auxiliary⁻::Vars{A},
    state_prognostic⁺::Vars{S},
    state_auxiliary⁺::Vars{A},
    t,
    direction,
) where {S, A}

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

    fluxᵀn = parent(fluxᵀn)
    wavespeed⁻ = wavespeed(
        balance_law,
        normal_vector,
        state_prognostic⁻,
        state_auxiliary⁻,
        t,
        direction,
    )
    wavespeed⁺ = wavespeed(
        balance_law,
        normal_vector,
        state_prognostic⁺,
        state_auxiliary⁺,
        t,
        direction,
    )
    max_wavespeed = max.(wavespeed⁻, wavespeed⁺)
    penalty =
        max_wavespeed .* (parent(state_prognostic⁻) - parent(state_prognostic⁺))

    # TODO: should this operate on ΔQ or penalty?
    update_penalty!(
        numerical_flux,
        balance_law,
        normal_vector,
        max_wavespeed,
        Vars{S}(penalty),
        state_prognostic⁻,
        state_auxiliary⁻,
        state_prognostic⁺,
        state_auxiliary⁺,
        t,
    )

    fluxᵀn .+= penalty / 2
end

"""
    CentralNumericalFluxFirstOrder() <: NumericalFluxFirstOrder

The central numerical flux for nondiffusive terms

# Usage

    CentralNumericalFluxFirstOrder()

Requires a `flux_first_order!` method for the balance law.
"""
struct CentralNumericalFluxFirstOrder <: NumericalFluxFirstOrder end

function numerical_flux_first_order!(
    ::CentralNumericalFluxFirstOrder,
    balance_law::BalanceLaw,
    fluxᵀn::Vars{S},
    normal_vector::SVector,
    state_prognostic⁻::Vars{S},
    state_auxiliary⁻::Vars{A},
    state_prognostic⁺::Vars{S},
    state_auxiliary⁺::Vars{A},
    t,
    direction,
) where {S, A}

    FT = eltype(fluxᵀn)
    num_state_prognostic = number_states(balance_law, Prognostic())
    fluxᵀn = parent(fluxᵀn)

    flux⁻ = similar(fluxᵀn, Size(3, num_state_prognostic))
    fill!(flux⁻, -zero(FT))
    flux_first_order!(
        balance_law,
        Grad{S}(flux⁻),
        state_prognostic⁻,
        state_auxiliary⁻,
        t,
        direction,
    )

    flux⁺ = similar(fluxᵀn, Size(3, num_state_prognostic))
    fill!(flux⁺, -zero(FT))
    flux_first_order!(
        balance_law,
        Grad{S}(flux⁺),
        state_prognostic⁺,
        state_auxiliary⁺,
        t,
        direction,
    )

    fluxᵀn .+= (flux⁻ + flux⁺)' * (normal_vector / 2)
end

"""
    RoeNumericalFlux() <: NumericalFluxFirstOrder

A numerical flux based on the approximate Riemann solver of Roe

# Usage

    RoeNumericalFlux()

Requires a custom implementation for the balance law.
"""
struct RoeNumericalFlux <: NumericalFluxFirstOrder end

"""
    HLLCNumericalFlux() <: NumericalFluxFirstOrder

A numerical flux based on the approximate Riemann solver of the
HLLC method. The HLLC flux is a modification of the Harten, Lax, van-Leer
(HLL) flux, where an additional contact property is introduced in order
to restore missing rarefraction waves. The HLLC flux requires
model-specific information, hence it requires a custom implementation
based on the underlying balance law.

# Usage

    HLLCNumericalFlux()

Requires a custom implementation for the balance law.

 - [Toro2013](@cite)
"""
struct HLLCNumericalFlux <: NumericalFluxFirstOrder end

"""
    LMARSNumericalFlux <: NumericalFluxFirstOrder
Low Mach Number Approximate Riemann Solver. Upwind biased
first order flux function.

- [Chen2013](@cite)
"""
struct LMARSNumericalFlux <: NumericalFluxFirstOrder end

"""
    RoeNumericalFluxMoist <: NumericalFluxFirstOrder

A moist implementation of the numerical flux based on the approximate Riemann solver of Roe

Requires a custom implementation for the balance law.
"""
struct RoeNumericalFluxMoist <: NumericalFluxFirstOrder
    " set to true for low Mach number correction"
    LM::Bool
    " set to true for Hartman Hyman correction"
    HH::Bool
    " set to true for LeVeque correction"
    LV::Bool
    " set to true for Positivity preserving LeVeque Correction"
    LVPP::Bool
end

RoeNumericalFluxMoist(;
    LM::Bool = false,
    HH::Bool = false,
    LV::Bool = false,
    LVPP::Bool = false,
) = RoeNumericalFluxMoist(LM, HH, LV, LVPP)

"""
    NumericalFluxSecondOrder

Any `N <: NumericalFluxSecondOrder` should define the a method for

    numerical_flux_second_order!(
        numerical_flux::N,
        balance_law::BalanceLaw,
        flux,
        normal_vector⁻,
        Q⁻, Qstate_gradient_flux⁻, Qaux⁻,
        Q⁺, Qstate_gradient_flux⁺, Qaux⁺,
        t
    )

where
- `flux` is the numerical flux array
- `normal_vector⁻` is the unit normal
- `Q⁻`/`Q⁺` are the minus/positive state arrays
- `Qstate_gradient_flux⁻`/`Qstate_gradient_flux⁺` are the minus/positive diffusive state arrays
- `Qstate_gradient_flux⁻`/`Qstate_gradient_flux⁺` are the minus/positive auxiliary state arrays
- `t` is the time

An optional method can also be defined for

    numerical_boundary_flux_second_order!(
        numerical_flux::N,
        balance_law::BalanceLaw,
        flux,
        normal_vector⁻,
        Q⁻, Qstate_gradient_flux⁻, Qaux⁻,
        Q⁺, Qstate_gradient_flux⁺, Qaux⁺,
        bctype,
        t
    )

"""
abstract type NumericalFluxSecondOrder end

function numerical_flux_second_order! end

function numerical_boundary_flux_second_order! end

"""
    CentralNumericalFluxSecondOrder <: NumericalFluxSecondOrder

The central numerical flux for diffusive terms

# Usage

    CentralNumericalFluxSecondOrder()

Requires a `flux_second_order!` for the balance law.
"""
struct CentralNumericalFluxSecondOrder <: NumericalFluxSecondOrder end

function numerical_flux_second_order!(
    ::CentralNumericalFluxSecondOrder,
    balance_law::BalanceLaw,
    fluxᵀn::Vars{S},
    normal_vector⁻::SVector,
    state_prognostic⁻::Vars{S},
    state_gradient_flux⁻::Vars{D},
    state_hyperdiffusive⁻::Vars{HD},
    state_auxiliary⁻::Vars{A},
    state_prognostic⁺::Vars{S},
    state_gradient_flux⁺::Vars{D},
    state_hyperdiffusive⁺::Vars{HD},
    state_auxiliary⁺::Vars{A},
    t,
) where {S, D, HD, A}

    FT = eltype(fluxᵀn)
    num_state_prognostic = number_states(balance_law, Prognostic())
    fluxᵀn = parent(fluxᵀn)

    flux⁻ = similar(fluxᵀn, Size(3, num_state_prognostic))
    fill!(flux⁻, -zero(FT))
    flux_second_order!(
        balance_law,
        Grad{S}(flux⁻),
        state_prognostic⁻,
        state_gradient_flux⁻,
        state_hyperdiffusive⁻,
        state_auxiliary⁻,
        t,
    )

    flux⁺ = similar(fluxᵀn, Size(3, num_state_prognostic))
    fill!(flux⁺, -zero(FT))
    flux_second_order!(
        balance_law,
        Grad{S}(flux⁺),
        state_prognostic⁺,
        state_gradient_flux⁺,
        state_hyperdiffusive⁺,
        state_auxiliary⁺,
        t,
    )

    fluxᵀn .+= (flux⁻ + flux⁺)' * (normal_vector⁻ / 2)
end

abstract type DivNumericalPenalty end
struct CentralNumericalFluxDivergence <: DivNumericalPenalty end

function numerical_flux_divergence!(
    ::CentralNumericalFluxDivergence,
    balance_law::BalanceLaw,
    div_penalty::Vars{GL},
    normal_vector::SVector,
    grad⁻::Grad{GL},
    grad⁺::Grad{GL},
) where {GL}
    parent(div_penalty) .=
        (parent(grad⁺) .+ parent(grad⁻))' * (normal_vector / 2)
end

function numerical_boundary_flux_divergence!(
    numerical_flux::CentralNumericalFluxDivergence,
    bctype,
    balance_law::BalanceLaw,
    div_penalty::Vars{GL},
    normal_vector::SVector,
    grad⁻::Grad{GL},
    state_auxiliary⁻::Vars{A},
    grad⁺::Grad{GL},
    state_auxiliary⁺::Vars{A},
    t,
) where {GL, A}
    boundary_state!(
        numerical_flux,
        bctype,
        balance_law,
        grad⁺,
        state_auxiliary⁺,
        normal_vector,
        grad⁻,
        state_auxiliary⁻,
        t,
    )
    numerical_flux_divergence!(
        numerical_flux,
        balance_law,
        div_penalty,
        normal_vector,
        grad⁻,
        grad⁺,
    )
end

abstract type GradNumericalFlux end
struct CentralNumericalFluxHigherOrder <: GradNumericalFlux end

function numerical_flux_higher_order!(
    ::CentralNumericalFluxHigherOrder,
    balance_law::BalanceLaw,
    hyperdiff::Vars{HD},
    normal_vector::SVector,
    lap⁻::Vars{GL},
    state_prognostic⁻::Vars{S},
    state_auxiliary⁻::Vars{A},
    lap⁺::Vars{GL},
    state_prognostic⁺::Vars{S},
    state_auxiliary⁺::Vars{A},
    t,
) where {HD, GL, S, A}
    G = normal_vector .* (parent(lap⁺) .- parent(lap⁻))' ./ 2
    transform_post_gradient_laplacian!(
        balance_law,
        hyperdiff,
        Grad{GL}(G),
        state_prognostic⁻,
        state_auxiliary⁻,
        t,
    )
end

function numerical_boundary_flux_higher_order!(
    numerical_flux::CentralNumericalFluxHigherOrder,
    bctype,
    balance_law::BalanceLaw,
    hyperdiff::Vars{HD},
    normal_vector::SVector,
    lap⁻::Vars{GL},
    state_prognostic⁻::Vars{S},
    state_auxiliary⁻::Vars{A},
    lap⁺::Vars{GL},
    state_prognostic⁺::Vars{S},
    state_auxiliary⁺::Vars{A},
    t,
) where {HD, GL, S, A}
    boundary_state!(
        numerical_flux,
        bctype,
        balance_law,
        state_prognostic⁺,
        state_auxiliary⁺,
        lap⁺,
        normal_vector,
        state_prognostic⁻,
        state_auxiliary⁻,
        lap⁻,
        t,
    )
    numerical_flux_higher_order!(
        numerical_flux,
        balance_law,
        hyperdiff,
        normal_vector,
        lap⁻,
        state_prognostic⁻,
        state_auxiliary⁻,
        lap⁺,
        state_prognostic⁺,
        state_auxiliary⁺,
        t,
    )
end

numerical_boundary_flux_second_order!(
    numerical_flux::CentralNumericalFluxSecondOrder,
    bctype,
    balance_law::BalanceLaw,
    fluxᵀn::Vars{S},
    normal_vector::SVector,
    state_prognostic⁻::Vars{S},
    state_gradient_flux⁻::Vars{D},
    state_hyperdiffusive⁻::Vars{HD},
    state_auxiliary⁻::Vars{A},
    state_prognostic⁺::Vars{S},
    state_gradient_flux⁺::Vars{D},
    state_hyperdiffusive⁺::Vars{HD},
    state_auxiliary⁺::Vars{A},
    t,
    state1⁻::Vars{S},
    diff1⁻::Vars{D},
    aux1⁻::Vars{A},
) where {S, D, HD, A} = normal_boundary_flux_second_order!(
    numerical_flux,
    bctype,
    balance_law,
    fluxᵀn,
    normal_vector,
    state_prognostic⁻,
    state_gradient_flux⁻,
    state_hyperdiffusive⁻,
    state_auxiliary⁻,
    state_prognostic⁺,
    state_gradient_flux⁺,
    state_hyperdiffusive⁺,
    state_auxiliary⁺,
    t,
    state1⁻,
    diff1⁻,
    aux1⁻,
)

function normal_boundary_flux_second_order!(
    numerical_flux,
    bctype,
    balance_law::BalanceLaw,
    fluxᵀn::Vars{S},
    normal_vector,
    state_prognostic⁻,
    state_gradient_flux⁻,
    state_hyperdiffusive⁻,
    state_auxiliary⁻,
    state_prognostic⁺,
    state_gradient_flux⁺,
    state_hyperdiffusive⁺,
    state_auxiliary⁺,
    t,
    state1⁻,
    diff1⁻,
    aux1⁻,
) where {S}
    FT = eltype(fluxᵀn)
    num_state_prognostic = number_states(balance_law, Prognostic())
    fluxᵀn = parent(fluxᵀn)

    flux = similar(fluxᵀn, Size(3, num_state_prognostic))
    fill!(flux, -zero(FT))
    boundary_flux_second_order!(
        numerical_flux,
        bctype,
        balance_law,
        Grad{S}(flux),
        state_prognostic⁺,
        state_gradient_flux⁺,
        state_hyperdiffusive⁺,
        state_auxiliary⁺,
        normal_vector,
        state_prognostic⁻,
        state_gradient_flux⁻,
        state_hyperdiffusive⁻,
        state_auxiliary⁻,
        t,
        state1⁻,
        diff1⁻,
        aux1⁻,
    )

    fluxᵀn .+= flux' * normal_vector
end

# This is the function that my be overloaded for flux-based BCs
function boundary_flux_second_order!(
    numerical_flux::NumericalFluxSecondOrder,
    bctype,
    balance_law,
    flux,
    state_prognostic⁺,
    state_gradient_flux⁺,
    state_hyperdiffusive⁺,
    state_auxiliary⁺,
    normal_vector,
    state_prognostic⁻,
    state_gradient_flux⁻,
    state_hyperdiffusive⁻,
    state_auxiliary⁻,
    t,
    state1⁻,
    diff1⁻,
    aux1⁻,
)
    boundary_state!(
        numerical_flux,
        bctype,
        balance_law,
        state_prognostic⁺,
        state_gradient_flux⁺,
        state_hyperdiffusive⁺,
        state_auxiliary⁺,
        normal_vector,
        state_prognostic⁻,
        state_gradient_flux⁻,
        state_hyperdiffusive⁻,
        state_auxiliary⁻,
        t,
        state1⁻,
        diff1⁻,
        aux1⁻,
    )
    flux_second_order!(
        balance_law,
        flux,
        state_prognostic⁺,
        state_gradient_flux⁺,
        state_hyperdiffusive⁺,
        state_auxiliary⁺,
        t,
    )
end

# Array wrappers for numerical Fluxes
function numerical_flux_first_order!(
    numerical_flux_first_order,
    balance_law,
    flux::AbstractArray,
    normal_vector::AbstractArray,
    state_prognostic⁻::AbstractArray,
    state_auxiliary⁻::AbstractArray,
    state_prognostic⁺::AbstractArray,
    state_auxiliary⁺::AbstractArray,
    t,
    face_direction,
)
    FT = eltype(flux)
    numerical_flux_first_order!(
        numerical_flux_first_order,
        balance_law,
        Vars{vars_state(balance_law, Prognostic(), FT)}(flux),
        SVector(normal_vector),
        Vars{vars_state(balance_law, Prognostic(), FT)}(state_prognostic⁻),
        Vars{vars_state(balance_law, Auxiliary(), FT)}(state_auxiliary⁻),
        Vars{vars_state(balance_law, Prognostic(), FT)}(state_prognostic⁺),
        Vars{vars_state(balance_law, Auxiliary(), FT)}(state_auxiliary⁺),
        t,
        face_direction,
    )
end

function numerical_boundary_flux_first_order!(
    numerical_flux_first_order,
    bctag::Int,
    balance_law,
    flux::AbstractArray,
    normal_vector::AbstractArray,
    state_prognostic⁻::AbstractArray,
    state_auxiliary⁻::AbstractArray,
    state_prognostic⁺::AbstractArray,
    state_auxiliary⁺::AbstractArray,
    t,
    face_direction,
    state_prognostic_bottom1::AbstractArray,
    state_auxiliary_bottom1::AbstractArray,
)
    FT = eltype(flux)
    bcs = boundary_conditions(balance_law)
    # TODO: there is probably a better way to unroll this loop
    Base.Cartesian.@nif 7 d -> bctag == d <= length(bcs) d -> begin
        bc = bcs[d]
        numerical_boundary_flux_first_order!(
            numerical_flux_first_order,
            bc,
            balance_law,
            Vars{vars_state(balance_law, Prognostic(), FT)}(flux),
            SVector(normal_vector),
            Vars{vars_state(balance_law, Prognostic(), FT)}(state_prognostic⁻),
            Vars{vars_state(balance_law, Auxiliary(), FT)}(state_auxiliary⁻),
            Vars{vars_state(balance_law, Prognostic(), FT)}(state_prognostic⁺),
            Vars{vars_state(balance_law, Auxiliary(), FT)}(state_auxiliary⁺),
            t,
            face_direction,
            Vars{vars_state(balance_law, Prognostic(), FT)}(
                state_prognostic_bottom1,
            ),
            Vars{vars_state(balance_law, Auxiliary(), FT)}(
                state_auxiliary_bottom1,
            ),
        )
    end d -> throw(BoundsError(bcs, bctag))
end

function numerical_flux_second_order!(
    numerical_flux_second_order,
    balance_law,
    flux::AbstractArray,
    normal_vector::AbstractArray,
    state_prognostic⁻::AbstractArray,
    state_gradient_flux⁻::AbstractArray,
    state_hyperdiffusive⁻::AbstractArray,
    state_auxiliary⁻::AbstractArray,
    state_prognostic⁺::AbstractArray,
    state_gradient_flux⁺::AbstractArray,
    state_hyperdiffusive⁺::AbstractArray,
    state_auxiliary⁺::AbstractArray,
    t,
)
    FT = eltype(flux)
    numerical_flux_second_order!(
        numerical_flux_second_order,
        balance_law,
        Vars{vars_state(balance_law, Prognostic(), FT)}(flux),
        SVector(normal_vector),
        Vars{vars_state(balance_law, Prognostic(), FT)}(state_prognostic⁻),
        Vars{vars_state(balance_law, GradientFlux(), FT)}(state_gradient_flux⁻),
        Vars{vars_state(balance_law, Hyperdiffusive(), FT)}(
            state_hyperdiffusive⁻,
        ),
        Vars{vars_state(balance_law, Auxiliary(), FT)}(state_auxiliary⁻),
        Vars{vars_state(balance_law, Prognostic(), FT)}(state_prognostic⁺),
        Vars{vars_state(balance_law, GradientFlux(), FT)}(state_gradient_flux⁺),
        Vars{vars_state(balance_law, Hyperdiffusive(), FT)}(
            state_hyperdiffusive⁺,
        ),
        Vars{vars_state(balance_law, Auxiliary(), FT)}(state_auxiliary⁺),
        t,
    )
end

function numerical_boundary_flux_second_order!(
    numerical_flux_second_order,
    bctag::Int,
    balance_law,
    flux::AbstractArray,
    normal_vector::AbstractArray,
    state_prognostic⁻::AbstractArray,
    state_gradient_flux⁻::AbstractArray,
    state_hyperdiffusive⁻::AbstractArray,
    state_auxiliary⁻::AbstractArray,
    state_prognostic⁺::AbstractArray,
    state_gradient_flux⁺::AbstractArray,
    state_hyperdiffusive⁺::AbstractArray,
    state_auxiliary⁺::AbstractArray,
    t,
    state_prognostic_bottom1::AbstractArray,
    state_auxiliary_bottom1::AbstractArray,
    state_gradient_flux_bottom1::AbstractArray,
)
    FT = eltype(flux)
    bcs = boundary_conditions(balance_law)
    # TODO: there is probably a better way to unroll this loop
    Base.Cartesian.@nif 7 d -> bctag == d <= length(bcs) d -> begin
        bc = bcs[d]
        numerical_boundary_flux_second_order!(
            numerical_flux_second_order,
            bc,
            balance_law,
            Vars{vars_state(balance_law, Prognostic(), FT)}(flux),
            SVector(normal_vector),
            Vars{vars_state(balance_law, Prognostic(), FT)}(state_prognostic⁻),
            Vars{vars_state(balance_law, GradientFlux(), FT)}(
                state_gradient_flux⁻,
            ),
            Vars{vars_state(balance_law, Hyperdiffusive(), FT)}(
                state_hyperdiffusive⁻,
            ),
            Vars{vars_state(balance_law, Auxiliary(), FT)}(state_auxiliary⁻),
            Vars{vars_state(balance_law, Prognostic(), FT)}(state_prognostic⁺),
            Vars{vars_state(balance_law, GradientFlux(), FT)}(
                state_gradient_flux⁺,
            ),
            Vars{vars_state(balance_law, Hyperdiffusive(), FT)}(
                state_hyperdiffusive⁺,
            ),
            Vars{vars_state(balance_law, Auxiliary(), FT)}(state_auxiliary⁺),
            t,
            Vars{vars_state(balance_law, Prognostic(), FT)}(
                state_prognostic_bottom1,
            ),
            Vars{vars_state(balance_law, GradientFlux(), FT)}(
                state_gradient_flux_bottom1,
            ),
            Vars{vars_state(balance_law, Auxiliary(), FT)}(
                state_auxiliary_bottom1,
            ),
        )
    end d -> throw(BoundsError(bcs, bctag))
end

function numerical_boundary_flux_gradient!(
    numerical_flux_gradient,
    bctag::Int,
    balance_law,
    flux::AbstractArray,
    normal_vector::AbstractArray,
    state_gradient⁻::AbstractArray,
    state_prognostic⁻::AbstractArray,
    state_auxiliary⁻::AbstractArray,
    state_gradient⁺::AbstractArray,
    state_prognostic⁺::AbstractArray,
    state_auxiliary⁺::AbstractArray,
    t,
    state_prognostic_bottom1::AbstractArray,
    state_auxiliary_bottom1::AbstractArray,
)
    FT = eltype(flux)
    bcs = boundary_conditions(balance_law)
    # TODO: there is probably a better way to unroll this loop
    Base.Cartesian.@nif 7 d -> bctag == d <= length(bcs) d -> begin
        bc = bcs[d]
        # Computes G* incorporating boundary conditions
        numerical_boundary_flux_gradient!(
            numerical_flux_gradient,
            bc,
            balance_law,
            flux,
            SVector(normal_vector),
            Vars{vars_state(balance_law, Gradient(), FT)}(state_gradient⁻),
            Vars{vars_state(balance_law, Prognostic(), FT)}(state_prognostic⁻),
            Vars{vars_state(balance_law, Auxiliary(), FT)}(state_auxiliary⁻),
            Vars{vars_state(balance_law, Gradient(), FT)}(state_gradient⁺),
            Vars{vars_state(balance_law, Prognostic(), FT)}(state_prognostic⁺),
            Vars{vars_state(balance_law, Auxiliary(), FT)}(state_auxiliary⁺),
            t,
            Vars{vars_state(balance_law, Prognostic(), FT)}(
                state_prognostic_bottom1,
            ),
            Vars{vars_state(balance_law, Auxiliary(), FT)}(
                state_auxiliary_bottom1,
            ),
        )
    end d -> throw(BoundsError(bcs, bctag))
end

end


text:
using Flux
println("NNPDE_tests")
using DiffEqFlux
println("Starting Soon!")
using ModelingToolkit
using DiffEqBase
using Test, NeuralPDE
println("Starting Soon!")
using GalacticOptim
using Optim
using Quadrature,Cubature, Cuba
using QuasiMonteCarlo
using SciMLBase
using OrdinaryDiffEq

using Random
Random.seed!(100)

cb = function (p,l)
    println("Current loss is: $l")
    return false
end

## Example 1, 1D ode
println("Example 1, 1D ode")
@parameters θ
@variables u(..)
Dθ = Differential(θ)

# 1D ODE
eq = Dθ(u(θ)) ~ θ^3 + 2*θ + (θ^2)*((1+3*(θ^2))/(1+θ+(θ^3))) - u(θ)*(θ + ((1+3*(θ^2))/(1+θ+θ^3)))

# Initial and boundary conditions
bcs = [u(0.) ~ 1.0]

# Space and time domains
domains = [θ ∈ IntervalDomain(0.0,1.0)]
# Discretization
dt = 0.1f0
# Neural network
chain = FastChain(FastDense(1,12,Flux.σ),FastDense(12,1))
initθ = DiffEqFlux.initial_params(chain)

strategy = NeuralPDE.GridTraining(0.1)
discretization = NeuralPDE.PhysicsInformedNN(chain,
                                             strategy;
                                             init_params = nothing,
                                             phi = nothing,
                                             derivative = nothing,
                                             )

pde_system = PDESystem(eq,bcs,domains,[θ],[u])
prob = NeuralPDE.discretize(pde_system,discretization)
sym_prob = NeuralPDE.symbolic_discretize(pde_system,discretization)

res = GalacticOptim.solve(prob, ADAM(0.1); cb = cb, maxiters=1000)
prob2 = remake(prob,u0=res.minimizer)
res = GalacticOptim.solve(prob2, ADAM(0.001); cb = cb, maxiters=1000)
phi = discretization.phi

analytic_sol_func(t) = exp(-(t^2)/2)/(1+t+t^3) + t^2
ts = [domain.domain.lower:dt/10:domain.domain.upper for domain in domains][1]
u_real  = [analytic_sol_func(t) for t in ts]
u_predict  = [first(phi(t,res.minimizer)) for t in ts]

@test u_predict ≈ u_real atol = 0.2

# using Plots
# t_plot = collect(ts)
# plot(t_plot ,u_real)
# plot!(t_plot ,u_predict)

## Example 2, 2D Poisson equation
function test_2d_poisson_equation(chain_, strategy_)
    println("Example 2, 2D Poisson equation, chain: $(typeof(chain_)), strategy: $strategy_")
    @parameters x y
    @variables u(..)
    Dxx = Differential(x)^2
    Dyy = Differential(y)^2

    # 2D PDE
    eq  = Dxx(u(x,y)) + Dyy(u(x,y)) ~ -sin(pi*x)*sin(pi*y)

    # Initial and boundary conditions
    bcs = [u(0,y) ~ 0.f0, u(1,y) ~ -sin(pi*1)*sin(pi*y),
           u(x,0) ~ 0.f0, u(x,1) ~ -sin(pi*x)*sin(pi*1)]
    # Space and time domains
    domains = [x ∈ IntervalDomain(0.0,1.0),
               y ∈ IntervalDomain(0.0,1.0)]

    # chain_ = FastChain(FastDense(2,12,Flux.σ),FastDense(12,12,Flux.σ),FastDense(12,1))
    discretization = NeuralPDE.PhysicsInformedNN(chain_,
                                                 strategy_)

    pde_system = PDESystem(eq,bcs,domains,[x,y],[u])
    prob = NeuralPDE.discretize(pde_system,discretization)
    sym_prob = NeuralPDE.symbolic_discretize(pde_system,discretization)
    res = GalacticOptim.solve(prob, ADAM(0.1); cb = cb, maxiters=500)
    phi = discretization.phi

    dx = 0.1
    xs,ys = [domain.domain.lower:dx/10:domain.domain.upper for domain in domains]
    analytic_sol_func(x,y) = (sin(pi*x)*sin(pi*y))/(2pi^2)

    u_predict = reshape([first(phi([x,y],res.minimizer)) for x in xs for y in ys],(length(xs),length(ys)))
    u_real = reshape([analytic_sol_func(x,y) for x in xs for y in ys], (length(xs),length(ys)))
    diff_u = abs.(u_predict .- u_real)

    @test u_predict ≈ u_real atol = 3.0

    # p1 = plot(xs, ys, u_real, linetype=:contourf,title = "analytic");
    # p2 = plot(xs, ys, u_predict, linetype=:contourf,title = "predict");
    # p3 = plot(xs, ys, diff_u,linetype=:contourf,title = "error");
    # plot(p1,p2,p3)
end

# Discretization
dx = 0.1
grid_strategy = NeuralPDE.GridTraining(dx)
fastchain = FastChain(FastDense(2,12,Flux.σ),FastDense(12,12,Flux.σ),FastDense(12,1))
fluxchain = Chain(Dense(2,12,Flux.σ),Dense(12,12,Flux.σ),Dense(12,1))
chains = [fluxchain, fastchain]
for chain in chains
    test_2d_poisson_equation(chain, grid_strategy)
end

stochastic_strategy = NeuralPDE.StochasticTraining(100) #points
quadrature_strategy = NeuralPDE.QuadratureTraining(quadrature_alg=CubatureJLh(),
                                                    reltol=1e-3,abstol=1e-3,
                                                    maxiters =50, batch=100)
quasirandom_strategy = NeuralPDE.QuasiRandomTraining(100; #points
                                                     sampling_alg = UniformSample(),
                                                     minibatch = 100)

strategies = [stochastic_strategy, quadrature_strategy,quasirandom_strategy]
for strategy_ in strategies
    chain_ = FastChain(FastDense(2,12,Flux.σ),FastDense(12,12,Flux.σ),FastDense(12,1))
    test_2d_poisson_equation(chain_, strategy_)
end

function run_2d_poisson_equation(strategy_)
    @parameters x y
    @variables u(..)
    Dxx = Differential(x)^2
    Dyy = Differential(y)^2

    # 2D PDE
    eq  = Dxx(u(x,y)) + Dyy(u(x,y)) ~ -sin(pi*x)*sin(pi*y)

    # Initial and boundary conditions
    bcs = [u(0,y) ~ 0.f0, u(1,y) ~ -sin(pi*1)*sin(pi*y),
           u(x,0) ~ 0.f0, u(x,1) ~ -sin(pi*x)*sin(pi*1)]
    # Space and time domains
    domains = [x ∈ IntervalDomain(0.0,1.0),
               y ∈ IntervalDomain(0.0,1.0)]
    # Discretization
    dx = 0.1

    chain = FastChain(FastDense(2,4,Flux.σ),FastDense(4,1))

    discretization = NeuralPDE.PhysicsInformedNN(chain,
                                                 strategy_)

    pde_system = PDESystem(eq,bcs,domains,[x,y],[u])
    prob = NeuralPDE.discretize(pde_system,discretization)

    res = GalacticOptim.solve(prob, ADAM(0.01); cb = cb,  maxiters=5)
end

algs = [CubatureJLh(), CubatureJLp()]
for alg in algs
    strategy_ =  NeuralPDE.QuadratureTraining(quadrature_alg = alg,reltol=1e-4,abstol=1e-3,maxiters=30, batch=10)
    run_2d_poisson_equation(strategy_)
end

## Example 3, high-order ode
println("Example 3, high-order ode")
@parameters x
@variables u(..)
Dxxx = Differential(x)^3
Dx = Differential(x)

# ODE
eq = Dxxx(u(x)) ~ cos(pi*x)

# Initial and boundary conditions
bcs = [u(0.) ~ 0.0,
       u(1.) ~ cos(pi),
       Dx(u(1.)) ~ 1.0]

# Space and time domains
domains = [x ∈ IntervalDomain(0.0,1.0)]

# Discretization
dx = 0.05

# Neural network
chain = FastChain(FastDense(1,8,Flux.σ),FastDense(8,1))
quasirandom_strategy = NeuralPDE.QuasiRandomTraining(100; #points
                                                     sampling_alg = UniformSample(),
                                                     minibatch = 100)
discretization = NeuralPDE.PhysicsInformedNN([chain],quasirandom_strategy)
pde_system = PDESystem(eq,bcs,domains,[x],[u])
prob = NeuralPDE.discretize(pde_system,discretization)
sym_prob = NeuralPDE.symbolic_discretize(pde_system,discretization)

res = GalacticOptim.solve(prob, ADAM(0.1); cb = cb, maxiters=300)
prob = remake(prob,u0=res.minimizer)
res = GalacticOptim.solve(prob, ADAM(0.01); cb = cb, maxiters=800)
prob = remake(prob,u0=res.minimizer)
res = GalacticOptim.solve(prob, ADAM(0.001); cb = cb, maxiters=600)
phi = discretization.phi[1]

analytic_sol_func(x) = (π*x*(-x+(π^2)*(2*x-3)+1)-sin(π*x))/(π^3)

xs = [domain.domain.lower:dx/10:domain.domain.upper for domain in domains][1]
u_real  = [analytic_sol_func(x) for x in xs]
u_predict  = [first(phi(x,res.minimizer)) for x in xs]

@test u_predict ≈ u_real atol = 1.0

# x_plot = collect(xs)
# plot(x_plot ,u_real)
# plot!(x_plot ,u_predict)

## Example 4, system of pde
println("Example 4, system of pde")
@parameters x, y
@variables u1(..), u2(..)
Dx = Differential(x)
Dy = Differential(y)

# System of pde
eqs = [Dx(u1(x,y)) + 4*Dy(u2(x,y)) ~ 0,
      Dx(u2(x,y)) + 9*Dy(u1(x,y)) ~ 0,
      3*u1(x,0) ~ 2*u2(x,0)]

# Initial and boundary conditions
bcs = [u1(x,0) ~ 2x, u2(x,0) ~ 3x]

# Space and time domains
domains = [x ∈ IntervalDomain(0.0,1.0), y ∈ IntervalDomain(0.0,1.0)]

# Neural network
chain1 = FastChain(FastDense(2,10,Flux.σ),FastDense(10,1))
chain2 = FastChain(FastDense(2,10,Flux.σ),FastDense(10,1))

strategy = NeuralPDE.GridTraining(0.1)
discretization = NeuralPDE.PhysicsInformedNN([chain1,chain2],strategy)
pde_system = PDESystem(eqs,bcs,domains,[x,y],[u1,u2])
prob = NeuralPDE.discretize(pde_system,discretization)
sym_prob = NeuralPDE.symbolic_discretize(pde_system,discretization)

res = GalacticOptim.solve(prob,ADAM(0.05); cb = cb, maxiters=2000)
phi = discretization.phi

analytic_sol_func(x,y) =[1/3*(6x - y), 1/2*(6x - y)]
xs,ys = [domain.domain.lower:0.01:domain.domain.upper for domain in domains]
u_real  = [[analytic_sol_func(x,y)[i] for x in xs  for y in ys] for i in 1:2]

initθ = discretization.init_params
acum =  [0;accumulate(+, length.(initθ))]
sep = [acum[i]+1 : acum[i+1] for i in 1:length(acum)-1]
minimizers = [res.minimizer[s] for s in sep]
u_predict  = [[phi[i]([x,y],minimizers[i])[1] for x in xs  for y in ys] for i in 1:2]

@test u_predict ≈ u_real atol = 3.0

# p1 =plot(xs, ys, u_predict, st=:surface);
# p2 = plot(xs, ys, u_real, st=:surface);
# plot(p1,p2)

## Example 5, 2d wave equation, neumann boundary condition
println("Example 5, 2d wave equation, neumann boundary condition")
#here we use low level api for build solution
@parameters x, t
@variables u(..)
Dxx = Differential(x)^2
Dtt = Differential(t)^2
Dt = Differential(t)

#2D PDE
C=1
eq  = Dtt(u(x,t)) ~ C^2*Dxx(u(x,t))

# Initial and boundary conditions
bcs = [u(0,t) ~ 0.,# for all t > 0
       u(1,t) ~ 0.,# for all t > 0
       u(x,0) ~ x*(1. - x), #for all 0 < x < 1
       Dt(u(x,0)) ~ 0. ] #for all  0 < x < 1]

# Space and time domains
domains = [x ∈ IntervalDomain(0.0,1.0),
           t ∈ IntervalDomain(0.0,1.0)]

# Neural network
chain = FastChain(FastDense(2,16,Flux.σ),FastDense(16,16,Flux.σ),FastDense(16,1))

phi = NeuralPDE.get_phi(chain)
derivative = NeuralPDE.get_numeric_derivative()
initθ = DiffEqFlux.initial_params(chain)

indvars = [x,t]
depvars = [u]
dim = length(domains)
quadrature_strategy = NeuralPDE.QuadratureTraining(quadrature_alg=CubatureJLh(),reltol= 1e-4,abstol= 1e-3,maxiters=10, batch=10)

_pde_loss_function = NeuralPDE.build_loss_function(eq,indvars,depvars,phi, derivative,chain,initθ,quadrature_strategy)
_pde_loss_function(rand(2,10), initθ)

bc_indvars = NeuralPDE.get_argument(bcs,indvars,depvars)
_bc_loss_functions = [NeuralPDE.build_loss_function(bc,indvars,depvars, phi, derivative,chain,initθ,quadrature_strategy,
                                              bc_indvars = bc_indvar) for (bc,bc_indvar) in zip(bcs,bc_indvars)]
map(loss_f -> loss_f(rand(1,10), initθ),_bc_loss_functions)

dx = 0.1
train_sets = NeuralPDE.generate_training_sets(domains,dx,[eq],bcs,indvars,depvars)
pde_train_set,bcs_train_set = train_sets
pde_bounds, bcs_bounds = NeuralPDE.get_bounds(domains,bcs,indvars,depvars,quadrature_strategy)

pde_loss_function = NeuralPDE.get_loss_function([_pde_loss_function],
                                                pde_bounds,
                                                quadrature_strategy;
                                                τ = 1/100)

pde_loss_function(initθ)
quadrature_strategy = NeuralPDE.QuadratureTraining(quadrature_alg=CubatureJLh(),reltol= 1e-2,abstol= 1e-1,maxiters=5, batch=100)
bc_loss_function = NeuralPDE.get_loss_function(_bc_loss_functions,
                                               bcs_bounds,
                                               quadrature_strategy;
                                               τ = 1/40)
bc_loss_function(initθ)

function loss_function_(θ,p)
    return pde_loss_function(θ) + bc_loss_function(θ)
end

f_ = OptimizationFunction(loss_function_, GalacticOptim.AutoZygote())
prob = GalacticOptim.OptimizationProblem(f_, initθ)

cb_ = function (p,l)
    println("Current losses are: ", pde_loss_function(p), " , ",  bc_loss_function(p))
    return false
end

res = GalacticOptim.solve(prob,Optim.BFGS(); cb = cb_, maxiters=400)

xs,ts = [domain.domain.lower:dx:domain.domain.upper for domain in domains]
analytic_sol_func(x,t) =  sum([(8/(k^3*pi^3)) * sin(k*pi*x)*cos(C*k*pi*t) for k in 1:2:50000])

u_predict = reshape([first(phi([x,t],res.minimizer)) for x in xs for t in ts],(length(xs),length(ts)))
u_real = reshape([analytic_sol_func(x,t) for x in xs for t in ts], (length(xs),length(ts)))

@test u_predict ≈ u_real atol = 1.0

# diff_u = abs.(u_predict .- u_real)
# p1 = plot(xs, ts, u_real, linetype=:contourf,title = "analytic");
# p2 =plot(xs, ts, u_predict, linetype=:contourf,title = "predict");
# p3 = plot(xs, ts, diff_u,linetype=:contourf,title = "error");
# plot(p1,p2,p3)

## Example 6, pde with mixed derivative
println("Example 6, pde with mixed derivative")
@parameters x y
@variables u(..)
Dxx = Differential(x)^2
Dyy = Differential(y)^2
Dx = Differential(x)
Dy = Differential(y)

eq = Dxx(u(x,y)) + Dx(Dy(u(x,y))) - 2*Dyy(u(x,y)) ~  -1.

# Initial and boundary conditions
bcs = [u(x,0) ~ x,
       Dy(u(x,0)) ~ x,
       u(x,0) ~ Dy(u(x,0))]

# Space and time domains
domains = [x ∈ IntervalDomain(0.0,1.0), y ∈ IntervalDomain(0.0,1.0)]

# Discretization
dx = 0.1
strategy = NeuralPDE.GridTraining(dx)
# Neural network
chain = FastChain(FastDense(2,12,Flux.σ),FastDense(12,12,Flux.σ),FastDense(12,1))
discretization = NeuralPDE.PhysicsInformedNN(chain, strategy)
pde_system = PDESystem(eq,bcs,domains,[x,y],[u])
prob = NeuralPDE.discretize(pde_system,discretization)

res = GalacticOptim.solve(prob,ADAM(0.1); cb = cb, maxiters=1000)
phi = discretization.phi

analytic_sol_func(x,y) = x + x*y +y^2/2
xs,ys = [domain.domain.lower:dx:domain.domain.upper for domain in domains]

u_predict = reshape([first(phi([x,y],res.minimizer)) for x in xs for y in ys],(length(xs),length(ys)))
u_real = reshape([analytic_sol_func(x,y) for x in xs for y in ys], (length(xs),length(ys)))
diff_u = abs.(u_predict .- u_real)

@test u_predict ≈ u_real atol = 1.0

# p1 = plot(xs, ys, u_real, linetype=:contourf,title = "analytic");
# p2 = plot(xs, ys, u_predict, linetype=:contourf,title = "predict");
# p3 = plot(xs, ys, diff_u,linetype=:contourf,title = "error");
# plot(p1,p2,p3)

## Example 7, Fokker-Planck equation
println("Example 7, Fokker-Planck equation")
# the example took from this article https://arxiv.org/abs/1910.10503
@parameters x
@variables p(..)
Dx = Differential(x)
Dxx = Differential(x)^2

#2D PDE
α = 0.3
β = 0.5
_σ = 0.5
# Discretization
dx = 0.05
# here we use normalization condition: dx*p(x) ~ 1, in order to get non-zero solution.
#(α - 3*β*x^2)*p(x) + (α*x - β*x^3)*Dx(p(x)) ~ (_σ^2/2)*Dxx(p(x))
eq  = Dx((α*x - β*x^3)*p(x)) ~ (_σ^2/2)*Dxx(p(x))+dx*p(x) - 1.

# Initial and boundary conditions
bcs = [p(-2.2) ~ 0. ,p(2.2) ~ 0. , p(-2.2) ~ p(2.2)]

# Space and time domains
domains = [x ∈ IntervalDomain(-2.2,2.2)]

# Neural network
inn = 18
chain = FastChain(FastDense(1,inn,Flux.σ),
                  FastDense(inn,inn,Flux.σ),
                  FastDense(inn,1))

discretization = NeuralPDE.PhysicsInformedNN(chain,
                                             NeuralPDE.GridTraining(dx))

pde_system = PDESystem(eq,bcs,domains,[x],[p])
prob = NeuralPDE.discretize(pde_system,discretization)

res = GalacticOptim.solve(prob,ADAM(0.1); cb = cb, maxiters=1500)
prob = remake(prob,u0=res.minimizer)
res = GalacticOptim.solve(prob,Optim.BFGS(initial_stepnorm=0.01);allow_f_increases=true, cb = cb, maxiters=500)

dx = dx/5
discretization = remake(discretization; strategy = NeuralPDE.GridTraining(dx), init_params =res.minimizer)
prob = NeuralPDE.discretize(pde_system,discretization)

res = GalacticOptim.solve(prob,ADAM(0.01); cb = cb, maxiters=1000)
prob = remake(prob,u0=res.minimizer)
res = GalacticOptim.solve(prob,Optim.BFGS(initial_stepnorm=0.01);allow_f_increases=true, cb = cb, maxiters=500)

dx= dx/5
discretization = remake(discretization; strategy = NeuralPDE.GridTraining(dx), init_params =res.minimizer)
prob = NeuralPDE.discretize(pde_system,discretization)

res = GalacticOptim.solve(prob,Optim.BFGS(initial_stepnorm=0.01);allow_f_increases=true, cb = cb, maxiters=500)
prob = remake(prob,u0=res.minimizer)
res = GalacticOptim.solve(prob,ADAM(0.001); cb = cb, maxiters=1000)

phi = discretization.phi
analytic_sol_func(x) = 28*exp((1/(2*_σ^2))*(2*α*x^2 - β*x^4))
xs = [domain.domain.lower:dx:domain.domain.upper for domain in domains][1]
u_real  = [analytic_sol_func(x) for x in xs]
u_predict  = [first(phi(x,res.minimizer)) for x in xs]

@test u_predict ≈ u_real atol = 100.0

# plot(xs ,u_real, label = "analytic")
# plot!(xs ,u_predict, label = "predict")

## Example 8, Lorenz System (Parameter Estimation)
println("Example 8, Lorenz System")

Random.seed!(1234)
@parameters t ,σ_ ,β, ρ
@variables x(..), y(..), z(..)
Dt = Differential(t)
eqs = [Dt(x(t)) ~ σ_*(y(t) - x(t)),
       Dt(y(t)) ~ x(t)*(ρ - z(t)) - y(t),
       Dt(z(t)) ~ x(t)*y(t) - β*z(t)]

bcs = [x(0) ~ 1.0, y(0) ~ 0.0, z(0) ~ 0.0]
domains = [t ∈ IntervalDomain(0.0,1.0)]
dt = 0.05f0

input_ = length(domains)
n = 8

chain1 = FastChain(FastDense(input_,n,Flux.σ),FastDense(n,n,Flux.σ),FastDense(n,1))
chain2 = FastChain(FastDense(input_,n,Flux.σ),FastDense(n,n,Flux.σ),FastDense(n,1))
chain3 = FastChain(FastDense(input_,n,Flux.σ),FastDense(n,n,Flux.σ),FastDense(n,1))

#Generate Data
function lorenz!(du,u,p,t)
 du[1] = 10.0*(u[2]-u[1])
 du[2] = u[1]*(28.0-u[3]) - u[2]
 du[3] = u[1]*u[2] - (8/3)*u[3]
end

u0 = [1.0;0.0;0.0]
tspan = (0.0,1.0)
prob = ODEProblem(lorenz!,u0,tspan)
sol = solve(prob, Tsit5(), dt=0.1)
ts = [domain.domain.lower:dt/5:domain.domain.upper for domain in domains][1]

function getData(sol)
    data = []
    us = hcat(sol(ts).u...)
    ts_ = hcat(sol(ts).t...)
    return [us,ts_]
end

data = getData(sol)

#Additional Loss Function
initθs = DiffEqFlux.initial_params.([chain1,chain2,chain3])
acum =  [0;accumulate(+, length.(initθs))]
sep = [acum[i]+1 : acum[i+1] for i in 1:length(acum)-1]
(u_ , t_) = data
len = length(data[2])

function additional_loss(phi, θ , p)
    return sum(sum(abs2, phi[i](t_ , θ[sep[i]]) .- u_[[i], :])/len for i in 1:1:3)
end

discretization = NeuralPDE.PhysicsInformedNN([chain1 , chain2, chain3],NeuralPDE.GridTraining(dt), param_estim=true, additional_loss=additional_loss)
pde_system = PDESystem(eqs,bcs,domains,[t],[x, y, z],[σ_, ρ, β],defaults=Dict([p .=> 1f0 for p in [σ_, ρ, β]]))
prob = NeuralPDE.discretize(pde_system,discretization)
sym_prob = NeuralPDE.symbolic_discretize(pde_system,discretization)

res = GalacticOptim.solve(prob, Optim.BFGS(); cb = cb, maxiters=6000)
p_ = res.minimizer[end-2:end]
@test sum(abs2, p_[1] - 10.00) < 0.1
@test sum(abs2, p_[2] - 28.00) < 0.1
@test sum(abs2, p_[3] - (8/3)) < 0.1
#Plotting the system
# initθ = discretization.init_params
# acum =  [0;accumulate(+, length.(initθ))]
# sep = [acum[i]+1 : acum[i+1] for i in 1:length(acum)-1]
# minimizers = [res.minimizer[s] for s in sep]
# ts = [domain.domain.lower:dt/10:domain.domain.upper for domain in domains][1]
# u_predict  = [[discretization.phi[i]([t],minimizers[i])[1] for t in ts] for i in 1:3]
# plot(sol)
# plot!(ts, u_predict, label = ["x(t)" "y(t)" "z(t)"])

ext: jl


text:
using LayeredLayouts
using Documenter

makedocs(;
    modules=[LayeredLayouts],
    authors="Lyndon White <lyndon.white@invenialabs.co.uk> and contributors",
    repo="https://github.com/oxinabox/LayeredLayouts.jl/blob/{commit}{path}#L{line}",
    sitename="LayeredLayouts.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://oxinabox.github.io/LayeredLayouts.jl",
        assets=String[],
    ),
    pages=[
        "Home" => "index.md",
    ],
)

deploydocs(;
    repo="github.com/oxinabox/LayeredLayouts.jl",
)


text:
function test_daxpy()
    n = 1000000
    count = 1000
    a::Float64 = 0.5
    x = Array{Float64, 1}(n)
    y = Array{Float64, 1}(n)

    for i in 1:n
        x[i] = i-1
        y[i] = 1.0
    end

    @timev for j in 1:count
        @inbounds for i in 1:n
            y[i] = a*x[i] + y[i]
        end
    end

    println(y[n])
end

#@code_llvm
test_daxpy()


text:
abstract type AbstractCovarianceKernel end

check_covpars(varpar, lenpar) = varpar > 0 && lenpar > 0 ||
    throw(DomainError("Covariance parameters must be positive"))

"""
    ExpCov{T}
        σ²::T
        ϕ::T

The exponential covariance function with variance σ² and correlation distance
paramter ϕ, so that covariance at distance ``d`` is given by

```math
\\Cov(d) = \\sigma^2 \\exp\\left(-\\frac{d}{\\phi})
```

Note that this is equivalent to a Matérn covariance with smoothness parameter ν
= 1/2
"""
struct ExpCov{T} <: AbstractCovarianceKernel
    σ²::T
    ϕ::T

    function ExpCov(σ²::T, ϕ::T) where T
        check_covpars(σ², ϕ)
        new{T}(σ², ϕ)
    end
end

function ExpCov(σ², ϕ)
    σ², ϕ = promote(σ², ϕ)
    ExpCov(σ², ϕ)
end

function (K::ExpCov)(d)
    K.σ² * exp(-d / K.ϕ)
end

"""
    Matérn32Cov{T}
        σ²::T
        ϕ::T

The Matérn covariance function with fixed smoothness ν = 3/2. Covariance at distance ``d`` is given by

```math
\\Cov(d) = \\sigma^2 \\left(1 + \\frac{\\sqrt{3} d}{\\phi}\\right)
     \\exp\\left(-\\frac{\\sqrt{3} d}{\\phi}\\right)
```
"""
struct Matérn32Cov{T} <: AbstractCovarianceKernel
    σ²::T
    ϕ::T

    function Matérn32Cov(σ²::T, ϕ::T) where T
        check_covpars(σ², ϕ)
        new{T}(σ², ϕ)
    end
end

function Matérn32Cov(σ², ϕ)
    σ², ϕ = promote(σ², ϕ)
    Matérn32Cov(σ², ϕ)
end

# Add a non-Unicode name
const Matern32Cov = Matérn32Cov

function (K::Matérn32Cov)(d)
    K.σ² * (1 + √3 * d / K.ϕ) * exp(-√3 * d / K.ϕ)
end

function cov(K::AbstractCovarianceKernel, Ω::AbstractFisheryDomain)
    Σ = zeros(length(Ω), length(Ω))
    @inbounds for jdx in eachindex(Ω), idx in eachindex(Ω)
        idx > jdx && continue
        Σ[idx, jdx] = K(Ω.distances[idx, jdx])
    end
    PDMat(Symmetric(Σ, :U))
end

function cov(K::AbstractCovarianceKernel, L::Vector{T}) where {T<:Union{Tuple,Real}}
    n = length(L)
    Σ = zeros(n, n)
    @inbounds for jdx in eachindex(L), idx in 1:jdx
        idx > jdx && continue
        Σ[idx, jdx] = K(hypot((L[idx] .- L[jdx])...))
    end
    PDMat(Symmetric(Σ, :U))
end

"""
    AR1{T}
        σ²::T
        ρ::T

Kernel of a discrete time AR(1) process with correlation parameter ρ.
"""
struct AR1{T} <: AbstractCovarianceKernel
    σ²::T
    ρ::T

    AR1(σ²::T, ρ::T) where T = σ² > 0 && ρ > 0 && new{T}(σ², ρ)
end

AR1(σ², ρ) = AR1(promote(σ², ρ)...)

function (ar1::AR1)(d::Integer)
    ar1.σ² / (1 - ar1.ρ^2) * ar1.ρ ^ abs(d)
end

function cov(ar1::AR1, d::Integer)
    c = ar1.(0:d - 1)
    Σ = zeros(d, d)
    for jdx in 1:d, idx in jdx:d
        Σ[idx, jdx] = c[idx - jdx + 1]
    end
    PDMat(Symmetric(Σ, :L))
end


text:
# argument types too specific?
function sumprod(A::Vector{Float64}, B::Vector{Float64})
    return sum(A .* B)
end

# what options are available?
sumprod_1(A::Vector{Float64}, B::Vector{Float64}) = sum(A .* B)
sumprod_2(A::Vector{Number}, B::Vector{Number}) = sum(A .* B)
sumprod_3(A::Vector{T}, B::Vector{T}) where T <: Number = sum(A .* B)
sumprod_4(A::Vector{S}, B::Vector{T}) where {S <: Number, T <: Number} = sum(A .* B)
sumprod_5(A::Array{S,N}, B::Array{T,N}) where {N, S <: Number, T <: Number} = sum(A .* B)
sumprod_6(A::AbstractArray{S,N}, B::AbstractArray{T,N}) where {N, S <: Number, T <: Number} = sum(A .* B)
sumprod_7(A, B) = sum(A .* B)

# let's build a test harness to testing various scenarios
function test_harness(f, scenario, args...)
    try
        f(args...)
        println(f, " #$(scenario) success")
    catch ex
        if ex isa MethodError
            println(f, " #$(scenario) failure (method not selected)")
        else
            println(f, " #$(scenario) failure (unknown error $ex)")
        end
    end
end

# Test a specifc sumprod function against several scenarios
function test_sumprod(f)
    test_harness(f, 1, [1.0,2.0], [3.0, 4.0]);
    test_harness(f, 2, [1,2], [3,4]);
    test_harness(f, 3, [1,2], [3.0,4.0]);
    test_harness(f, 4, rand(2,2), rand(2,2));
    test_harness(f, 5, Number[1,2.0], Number[3.0, 4]);
end

test_sumprod(sumprod_1)
#=
julia> test_sumprod(sumprod_1)
sumprod_1 #1 success
sumprod_1 #2 failure (method not selected)
sumprod_1 #3 failure (method not selected)
sumprod_1 #4 failure (method not selected)
sumprod_1 #5 failure (method not selected)
=#

test_sumprod(sumprod_2)
#=
julia> test_sumprod(sumprod_2)
sumprod_2 #1 failure (method not selected)
sumprod_2 #2 failure (method not selected)
sumprod_2 #3 failure (method not selected)
sumprod_2 #4 failure (method not selected)
sumprod_2 #5 success
=#

Vector{Float64} <: Vector{Number}
#=
julia> Vector{Float64} <: Vector{Number}
false
=#

test_sumprod(sumprod_3)
#=
julia> test_sumprod(sumprod_3)
sumprod_3 #1 success
sumprod_3 #2 success
sumprod_3 #3 failure (method not selected)
sumprod_3 #4 failure (method not selected)
sumprod_3 #5 success
=#

test_sumprod(sumprod_4)
#=
julia> test_sumprod(sumprod_4)
sumprod_4 #1 success
sumprod_4 #2 success
sumprod_4 #3 success
sumprod_4 #4 failure (method not selected)
sumprod_4 #5 success
=#

test_sumprod(sumprod_5)
#=
julia> test_sumprod(sumprod_5)
sumprod_5 #1 success
sumprod_5 #2 success
sumprod_5 #3 success
sumprod_5 #4 success
sumprod_5 #5 success
=#

test_sumprod(sumprod_6)
#=
julia> test_sumprod(sumprod_6)
sumprod_6 #1 success
sumprod_6 #2 success
sumprod_6 #3 success
sumprod_6 #4 success
sumprod_6 #5 success
=#

using SparseArrays
A = sparse([1,10,100], [1,10,100], [1,2,3])
B = sparse([1,10,100], [1,10,100], [4,5,6])
sumprod_6(A, B)

test_sumprod(sumprod_7)
#=
julia> test_sumprod(sumprod_7)
sumprod_7 #1 success
sumprod_7 #2 success
sumprod_7 #3 success
sumprod_7 #4 success
sumprod_7 #5 success
=#

#--- any performance impact???
using BenchmarkTools

A = rand(10_000);
B = rand(10_000);

@btime sumprod_1($A, $B);
@btime sumprod_5($A, $B);
@btime sumprod_6($A, $B);
@btime sumprod_7($A, $B);

#=
julia> @btime sumprod_1($A, $B);
  11.769 μs (2 allocations: 78.20 KiB)

julia> @btime sumprod_5($A, $B);
  11.792 μs (2 allocations: 78.20 KiB)

julia> @btime sumprod_6($A, $B);
  11.944 μs (2 allocations: 78.20 KiB)

julia> @btime sumprod_7($A, $B);
  11.626 μs (2 allocations: 78.20 KiB)
=#


text:
using Pkg
Pkg.activate(".")
using ConceptnetNumberbatch
using Serialization

# Load a serialized version of ConceptNetEnglish
fid = open("./_conceptnet_/numberbatch-en-17.06.txt.bin")
cptnet = deserialize(fid)
close(fid)

## get similar words from NN q-gram model
dictionary = collect(keys(cptnet))
@time modelstuff = ConceptnetNumberbatch.build_nn_model(dictionary, ngram_size=2)
open("./_conceptnet_/model.bin", "w") do fid
    serialize(fid, modelstuff)
end
#modelstuff = open(deserialize, "./_conceptnet_/model.bin", "r")
targets = ["phrasis", "appled", "fdellity", "moanster"]
found = ConceptnetNumberbatch.get_similar_words(targets[1], modelstuff...)
for target in targets
    @time found = ConceptnetNumberbatch.get_similar_words(target, modelstuff...)
    println("$target => $(found[1])")
end


text:
immutable HyperbolicIntercept{T} <: Model
    β₀::T
    a::T
    α::T
end

function Base.show(io::IO, ::Type{HyperbolicIntercept})
    return print(io, "Intercept Hyperbolic")
end

function Base.show(io::IO, m::HyperbolicIntercept)
    @printf(
        io,
        "Intercept Hyperbolic (β₀ = %.3f, a = %.3f, α = %.3f)",
        m.β₀,
        m.a,
        m.α,
    )
end

function initial_parameters(::Type{HyperbolicIntercept})
    return Float64[0.0, log(1.0), log(1.0)]
end

canonical_parameters(m::HyperbolicIntercept) = (m.β₀, m.a, m.α)

function HyperbolicIntercept(parameters::Vector)
    β₀ = parameters[1]
    a = exp(parameters[2])
    α = exp(parameters[3])
    return HyperbolicIntercept(β₀, a, α)
end

@inline function predict(
    m::HyperbolicIntercept,
    x2::Real,
    t2::Real,
    x1::Real,
    t1::Real,
)
    β₀, a, α = m.β₀, m.a, m.α
    u1 = x1 * (1 / (1 + α * t1))
    u2 = x2 * (1 / (1 + α * t2))
    z = β₀ + a * (u2 - u1)
    p = invlogit(z)
    return p
end

@inline function gradient_component(
    m::HyperbolicIntercept,
    x2::Real,
    t2::Real,
    x1::Real,
    t1::Real,
)
    β₀, a, α = m.β₀, m.a, m.α
    d_t2 = 1 / (1 + α * t2)
    d_t1 = 1 / (1 + α * t1)
    gr = (
        1.0,
        (x2 * d_t2 - x1 * d_t1) * a,
        a * (x1 * t1 * d_t1^2 - x2 * t2 * d_t2^2) * α,
    )
    return gr
end


text:
using CENG0038-Droplet-Distribution-research-project
using Test

@testset "CENG0038-Droplet-Distribution-research-project.jl" begin
    # Write your tests here.
end


text:
using MotifSequenceGenerator

export random_notes_sequence

function notes_limits(notes::Notes)
    start = minimum(n.position for n in notes)
    fine  = maximum(n.position + n.duration for n in notes)
    return (start, fine)
end

"""
    random_notes_sequence(motifs::Vector{Notes{N}}, q, δq = 0; weights)
Create a random sequence from a pool of notes (`motifs`) such that
it has total length `ℓ` exactly `q - δq ≤ ℓ ≤ q + δq`.
Notice that `q` is measured in **ticks**. Optionally pass keyword `weights`
to sample different motifs with different weights (either as frequencies or as
probabilities).

Return the result as a single `Notes` container, and also return the sequence
of motifs used.

This function uses [`random_sequence`](@ref) from the module
[`MotifSequenceGenerator`](@ref), adapted to the [`Notes`](@ref) struct.
"""
function random_notes_sequence(motifs::Vector{Notes{N}}, q, δq = 0;
                               weights = ones(length(motifs))) where N

    tpq = motifs[1].tpq
    any(x -> x != tpq, notes.tpq for notes in motifs) && throw(ArgumentError(
    "The pool of motifs does no share the same `tpq`."))

    res, seq = random_sequence(motifs, q, notes_limits, translate, δq; weights=weights)
    ret = N[]
    for notes in res
        append!(ret, notes.notes)
    end
    return Notes(ret, motifs[1].tpq), seq
end

# function minimum_subdivision(motifs::Vector{M}) where {M<:Notes}
#     Int(minimum(minimum(n.duration for n in notes) for notes in motifs))
# end
# q % minimum_subdivision(motifs) == 0 || throw(ArgumentError("
# Given window `q` is not divisible by the minimum subdivision present in `motifs`."))


text:
module Insolation

using Dates

using CLIMAParameters: AbstractParameterSet
using CLIMAParameters
using CLIMAParameters.Planet
const APS = AbstractParameterSet
Base.broadcastable(param_set::APS) = Ref(param_set)

include("ZenithAngleCalc.jl")
include("InsolationCalc.jl")

end # module


text:
module ParametricPolyhedra

export ParametricTriangle, edges

type ParametricTriangle{T}
    # edge lengths
    a::Nullable{T}
    b::Nullable{T}
    c::Nullable{T}
    # angles (radians)
    alpha::Nullable{T}
    beta::Nullable{T}
    gamma::Nullable{T}
end

# functions to allow mixed constructors
makenullable(a::Nullable) = a
makenullable(a) = Nullable(a)

# mixed case
function ParametricTriangle(a, b, c,
                            alpha, beta, gamma)
    ParametricTriangle(makenullable(a),
                       makenullable(b),
                       makenullable(c),
                       makenullable(alpha),
                       makenullable(beta),
                       makenullable(gamma))
end

function Base.isnull(p::ParametricTriangle)
    isnull(p.a) || isnull(p.b) || isnull(p.c) ||
    isnull(p.alpha) || isnull(p.beta) || isnull(p.gamma)
end

"""
Test if a ParametricTriangle has a valid configuration.
"""
function Base.isvalid(p::ParametricTriangle)
    # underdetermined case
    isnull(p) && return false
    # otherwise check constraints since all values exist
    a = get(p.a)
    b = get(p.b)
    c = get(p.c)
    alpha = get(p.alpha)
    beta = get(p.beta)
    gamma = get(p.gamma)
    return a*cos(beta) + b*cos(alpha) - c == 0 &&
           b*sin(alpha) - a*sin(beta) == 0 &&
           alpha + beta + gamma - pi == 0
end

# version with isapprox for floats
function Base.isvalid{T<:AbstractFloat}(p::ParametricTriangle{T};
                                        rtol=sqrt(eps(T)),
                                        atol=zero(T))
    # underdetermined case
    isnull(p) && return false
    # otherwise check constraints since all values exist
    a = p.a.value
    b = p.b.value
    c = p.c.value
    alpha = p.alpha.value
    beta = p.beta.value
    gamma = p.gamma.value
    return isapprox(a*cos(beta) + b*cos(alpha) - c,0,
                    rtol=rtol,atol=atol) &&
           isapprox(b*sin(alpha) - a*sin(beta),0,
                    rtol=rtol,atol=atol) &&
           isapprox(alpha + beta + gamma - pi,0,
                    rtol=rtol,atol=atol)
end

"""
Given three angles, compute the edge lengths. The circumcircle radius
default is 1.0, but can be specified as a fourth argument.
"""
function edges{T}(alpha::T,beta::T,gamma::T,c=one(T))
    (c*sin(alpha),c*sin(beta),c*sin(gamma))
end

"""
Given angle `a` and circumcircle diameter, `c`,
compute the opposite edge length in the triangle.
"""
function _edge(a,c)
    c*sin(a)
end

"""
Given an underdetermined ParametricTriangle, compute the missing values
and return a new ParametricTriangle
"""
function Base.fill(p::ParametricTriangle)
    # all angles must be specified
    if isnull(p.alpha) || isnull(p.alpha) || isnull(p.gamma)
        error("Cannot fill in values for this triangle. All angles must be specified")
    end
    alpha = get(p.alpha)
    beta = get(p.beta)
    gamma = get(p.gamma)
    # no edges given, use circumcircle=1
    if isnull(p.a) && isnull(p.b) && isnull(p.c)
        e = edges(alpha,beta,gamma)
        return ParametricTriangle(e[1],e[2],e[3],p.alpha,p.beta,p.gamma)
    else
        # find the circumcircle
        D = !isnull(p.a) ? get(p.a)/sin(alpha) :
            !isnull(p.b) ? get(p.b)/sin(beta) :
            get(p.c)/sin(gamma) # one must be specified because of prior check
        # we only need to figure one side that is specified
        # so we can (re)compute the other two
        if !isnull(p.a)
            return ParametricTriangle(p.a, _edge(beta,D), _edge(gamma,D),
                                      p.alpha, p.beta, p.gamma)
        elseif !isnull(p.b)
            return ParametricTriangle(_edge(alpha,D), p.b, _edge(gamma,D),
                                      p.alpha, p.beta, p.gamma)
        elseif !isnull(p.c)
            return ParametricTriangle(_edge(alpha,D), _edge(beta,D), p.c,
                                      p.alpha, p.beta, p.gamma)
        end
    end
    # otherwise we broke it ... TODO remove once we start using this in loops
    error("filling triangle failed! What did you do???")
end

end # module


text:
#

__precompile__()

module LibArchive

using Compat

if isdefined(Base, :unsafe_read)
    # I forgot to add `unsafe_read` to `Compat` ;-P
    import Base: unsafe_read, unsafe_write
else
    const unsafe_read = read
    const unsafe_write = write
end

export ArchiveRetry, ArchiveFailed, ArchiveFatal

const depfile = joinpath(dirname(@__FILE__), "..", "deps", "deps.jl")
if isfile(depfile)
    include(depfile)
else
    error("LibArchive not properly installed. Please run Pkg.build(\"LibArchive\")")
end

include("constants.jl")

###
# Version

"""
libarchive version number
"""
function version()
    vernum = ccall((:archive_version_number, libarchive), Cint, ())
    vernum, patch = divrem(vernum, Cint(1000))
    major, minor = divrem(vernum, Cint(1000))
    VersionNumber(major, minor, patch)
end

@compat abstract type Archive <: IO end

function archive_guard(func::Function, archive::Archive)
    try
        res = func(archive)
        # Only explicitly close if the function didn't throw an error
        # otherwise, only call free to avoid overriding user exceptions
        close(archive)
        return res
    finally
        free(archive)
    end
end

include("error.jl")
include("callback.jl")
include("archive_utils.jl")
include("entry.jl")
include("reader.jl")
include("writer.jl")
include("format.jl")

# /*
#  * Set archive_match object that will be used in archive_read_disk to
#  * know whether an entry should be skipped. The callback function
#  * _excluded_func will be invoked when an entry is skipped by the result
#  * of archive_match.
#  */
# int	archive_read_disk_set_matching(struct archive *,
# 		    struct archive *_matching, void (*_excluded_func)
# 		    (struct archive *, void *, struct archive_entry *),
# 		    void *_client_data);
# int	archive_read_disk_set_metadata_filter_callback(struct archive *,
# 		    int (*_metadata_filter_func)(struct archive *, void *,
# 		    	struct archive_entry *), void *_client_data);

# /*
#  * ARCHIVE_MATCH API
#  */
# struct archive *archive_match_new(void);
# int	archive_match_free(struct archive *);

# /*
#  * Test if archive_entry is excluded.
#  * This is a convenience function. This is the same as calling all
#  * archive_match_path_excluded, archive_match_time_excluded
#  * and archive_match_owner_excluded.
#  */
# int	archive_match_excluded(struct archive *,
# 		    struct archive_entry *);

# /*
#  * Test if pathname is excluded. The conditions are set by following functions.
#  */
# int	archive_match_path_excluded(struct archive *,
# 		    struct archive_entry *);
# /* Add exclusion pathname pattern. */
# int	archive_match_exclude_pattern(struct archive *, const char *);
# /* Add exclusion pathname pattern from file. */
# int	archive_match_exclude_pattern_from_file(struct archive *,
# 		    const char *, int _nullSeparator);
# /* Add inclusion pathname pattern. */
# int	archive_match_include_pattern(struct archive *, const char *);
# /* Add inclusion pathname pattern from file. */
# int	archive_match_include_pattern_from_file(struct archive *,
# 		    const char *, int _nullSeparator);
# /*
#  * How to get statistic information for inclusion patterns.
#  */
# /* Return the amount number of unmatched inclusion patterns. */
# int	archive_match_path_unmatched_inclusions(struct archive *);
# /* Return the pattern of unmatched inclusion with ARCHIVE_OK.
#  * Return ARCHIVE_EOF if there is no inclusion pattern. */
# int	archive_match_path_unmatched_inclusions_next(
# 		    struct archive *, const char **);

# /*
#  * Test if a file is excluded by its time stamp.
#  * The conditions are set by following functions.
#  */
# int	archive_match_time_excluded(struct archive *,
# 		    struct archive_entry *);

# /* Set inclusion time. */
# int	archive_match_include_time(struct archive *, int _flag,
# 		    time_t _sec, long _nsec);
# /* Set inclusion time by a date string. */
# int	archive_match_include_date(struct archive *, int _flag,
# 		    const char *_datestr);
# /* Set inclusion time by a particluar file. */
# int	archive_match_include_file_time(struct archive *,
# 		    int _flag, const char *_pathname);
# /* Add exclusion entry. */
# int	archive_match_exclude_entry(struct archive *,
# 		    int _flag, struct archive_entry *);

# /*
#  * Test if a file is excluded by its uid ,gid, uname or gname.
#  * The conditions are set by following functions.
#  */
# int	archive_match_owner_excluded(struct archive *,
# 		    struct archive_entry *);
# /* Add inclusion uid, gid, uname and gname. */
# int	archive_match_include_uid(struct archive *, int64_t);
# int	archive_match_include_gid(struct archive *, int64_t);
# int	archive_match_include_uname(struct archive *, const char *);
# int	archive_match_include_gname(struct archive *, const char *);

# /*
#  * ARCHIVE_READ_DISK API
#  *
#  * This is still evolving and somewhat experimental.
#  */
# struct archive *archive_read_disk_new(void);
# /* The names for symlink modes here correspond to an old BSD
#  * command-line argument convention: -L, -P, -H */
# /* Follow all symlinks. */
# int archive_read_disk_set_symlink_logical(struct archive *);
# /* Follow no symlinks. */
# int archive_read_disk_set_symlink_physical(struct archive *);
# /* Follow symlink initially, then not. */
# int archive_read_disk_set_symlink_hybrid(struct archive *);
# /* TODO: Handle Linux stat32/stat64 ugliness. <sigh> */
# int archive_read_disk_entry_from_file(struct archive *,
#     struct archive_entry *, int /* fd */, const struct stat *);
# /* Look up gname for gid or uname for uid. */
# /* Default implementations are very, very stupid. */
# const char *archive_read_disk_gname(struct archive *, int64_t);
# const char *archive_read_disk_uname(struct archive *, int64_t);
# /* "Standard" implementation uses getpwuid_r, getgrgid_r and caches the
#  * results for performance. */
# int	archive_read_disk_set_standard_lookup(struct archive *);
# /* You can install your own lookups if you like. */
# int	archive_read_disk_set_gname_lookup(struct archive *,
#     void * /* private_data */,
#     const char *(* /* lookup_fn */)(void *, int64_t),
#     void (* /* cleanup_fn */)(void *));
# int	archive_read_disk_set_uname_lookup(struct archive *,
#     void * /* private_data */,
#     const char *(* /* lookup_fn */)(void *, int64_t),
#     void (* /* cleanup_fn */)(void *));
# /* Start traversal. */
# int	archive_read_disk_open(struct archive *, const char *);
# /*
#  * Request that current entry be visited.  If you invoke it on every
#  * directory, you'll get a physical traversal.  This is ignored if the
#  * current entry isn't a directory or a link to a directory.  So, if
#  * you invoke this on every returned path, you'll get a full logical
#  * traversal.
#  */
# int	archive_read_disk_descend(struct archive *);
# int	archive_read_disk_can_descend(struct archive *);
# int	archive_read_disk_current_filesystem(struct archive *);
# int	archive_read_disk_current_filesystem_is_synthetic(struct archive *);
# int	archive_read_disk_current_filesystem_is_remote(struct archive *);
# /* Request that the access time of the entry visited by travesal be restored. */
# int  archive_read_disk_set_atime_restored(struct archive *);

# /*
#  * Set behavior. The "flags" argument selects optional behavior.
#  */
# int  archive_read_disk_set_behavior(struct archive *,
# 		    int flags);

# /*-
#  * ARCHIVE_WRITE_DISK API
#  *
#  * To create objects on disk:
#  *   1) Ask archive_write_disk_new for a new archive_write_disk object.
#  *   2) Set any global properties.  In particular, you probably
#  *      want to set the options.
#  *   3) For each entry:
#  *      - construct an appropriate struct archive_entry structure
#  *      - archive_write_header to create the file/dir/etc on disk
#  *      - archive_write_data to write the entry data
#  *   4) archive_write_free to cleanup the writer and release resources
#  *
#  * In particular, you can use this in conjunction with archive_read()
#  * to pull entries out of an archive and create them on disk.
#  */
# struct archive	*archive_write_disk_new(void);
# /* This file will not be overwritten. */
# int archive_write_disk_set_skip_file(struct archive *,
#     int64_t, int64_t);
# /* Set flags to control how the next item gets created.
#  * This accepts a bitmask of ARCHIVE_EXTRACT_XXX flags defined above. */
# int		 archive_write_disk_set_options(struct archive *,
# 		     int flags);
# /*
#  * The lookup functions are given uname/uid (or gname/gid) pairs and
#  * return a uid (gid) suitable for this system.  These are used for
#  * restoring ownership and for setting ACLs.  The default functions
#  * are naive, they just return the uid/gid.  These are small, so reasonable
#  * for applications that don't need to preserve ownership; they
#  * are probably also appropriate for applications that are doing
#  * same-system backup and restore.
#  */
# /*
#  * The "standard" lookup functions use common system calls to lookup
#  * the uname/gname, falling back to the uid/gid if the names can't be
#  * found.  They cache lookups and are reasonably fast, but can be very
#  * large, so they are not used unless you ask for them.  In
#  * particular, these match the specifications of POSIX "pax" and old
#  * POSIX "tar".
#  */
# int	 archive_write_disk_set_standard_lookup(struct archive *);
# /*
#  * If neither the default (naive) nor the standard (big) functions suit
#  * your needs, you can write your own and register them.  Be sure to
#  * include a cleanup function if you have allocated private data.
#  */
# int archive_write_disk_set_group_lookup(struct archive *,
#     void * /* private_data */,
#     int64_t (*)(void *, const char *, int64_t),
#     void (* /* cleanup */)(void *));
# int archive_write_disk_set_user_lookup(struct archive *,
#     void * /* private_data */,
#     int64_t (*)(void *, const char *, int64_t),
#     void (* /* cleanup */)(void *));
# int64_t archive_write_disk_gid(struct archive *, const char *, int64_t);
# int64_t archive_write_disk_uid(struct archive *, const char *, int64_t);

end


text:
#=

JuliaBolt is ported from Neobolt: Neo4j Bolt connector for Python
https://github.com/neo4j-drivers/neobolt

Original notices:

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

=#

using Sockets

#####################################
#           Message Frame           #
#####################################

const empty_view = UInt8[]

mutable struct MessageFrame
   data_view::SubArray{UInt8}
   current_pane::Integer # index of current pane in panes
   current_offset::Integer # index within the current pane
   panes::Union{Array{Tuple{Integer,Integer}}, Nothing}

   MessageFrame(data_view) = new(data_view, -1, -1, nothing)
   MessageFrame(data_view, panes) = new(data_view, 0, 0, panes)
end

# Message Frame -------------------- #

function close(m::MessageFrame)
    # what would setting data_view to nothing accomplish other than putting the object in an unusable and invalid state?
    # perhaps freeing up memory for the garbage collector?
    return
end

# Message Frame -------------------- #

function next_pane(m::MessageFrame)
    m.current_pane += 1
    if m.current_pane < length(m.panes)
        m.current_offset = 0
    else
        m.current_pane = -1
        m.current_offset = -1
    end
end

# Message Frame -------------------- #

function read_int(m::MessageFrame)
    if m.current_pane == -1
        return -1
    end
    (p, q) = m.panes[m.current_pane + 1] #index
    size = q - p
    value = m.data_view[p + m.current_offset + 1] #index
    m.current_offset += 1
    if m.current_offset == size
        next_pane(m)
    end
    return value
end

# Message Frame -------------------- #

function read(m::MessageFrame, n::Integer)
    if n == 0 || m.current_pane == -1
        return empty_view
    end
    value::Union{AbstractArray, Nothing} = nothing
    is_view = false
    offset = 0

    to_read = n
    while to_read > 0 && m.current_pane >= 0
        (p, q) = m.panes[m.current_pane + 1] #index
        size = q - p
        remaining = size - m.current_offset
        _start = p + m.current_offset
        if to_read <= remaining
            _end = _start + to_read
            if to_read < remaining
                m.current_offset += to_read
            else
                next_pane(m)
            end
        else
            _end = q
            next_pane(m)
        end

        read = _end - _start

        if value !== nothing
            if is_view
                new_value = zeros(UInt8, n)
                new_value[1:offset] = value[1:offset] #index
                value = new_value
                is_view = false
            end
            value[(offset+1):(offset+read)] = m.data_view[(_start+1):_end] #index
        else
            value = Base.view(m.data_view, (_start+1):_end) #index
            is_view = true
        end
        offset += read
        to_read -= read
    end
    return Base.view(value, :)
end

# Message Frame ------------------- #
# # # # # # # # # # # # # # # # # # #

#####################################
#        Chunked Input Buffer       #
#####################################

# TODO: Fix indexing to 1 based

mutable struct ChunkedInputBuffer
    data::Array{UInt8}
    data_view::SubArray{UInt8}
    extent::Integer
    origin::Integer
    limit::Integer
    frame::Union{MessageFrame, Nothing}

    function ChunkedInputBuffer(;capacity::Integer=524288)
        data = zeros(UInt8, capacity)
        return new(data, Base.view(data, :), 0, 0, -1, nothing)
    end
end

# Chunked Input Buffer ------------ #

function capacity(c::ChunkedInputBuffer)
    return length(c.data_view)
end

# Chunked Input Buffer ------------ #

function Base.view(c::ChunkedInputBuffer)
    return Base.view(c.data_view, 1:c.extent) #index
end

# Chunked Input Buffer ------------ #

function load(c::ChunkedInputBuffer, b::AbstractArray{UInt8})
    n = length(b)
    new_extent = c.extent + n
    overflow = new_extent - length(c.data)
    if overflow > 0 # if new data does not fit in data array
        if recycle(c)
            return load(c, b)
        end
        # c.data_view = nothing
        new_extent = c.extent + n

        if new_extent > length(c.data)
            resize!(c.data, new_extent)
        end

        c.data[(c.extent+1):new_extent] = b #index
        c.data_view = Base.view(c.data, :)
    else
        c.data_view[(c.extent + 1):new_extent] = b #index
    end
    c.extent = new_extent
end

function read_some(s::IO, limit=0)
    b = []
    try
        while true
            t = Base.read(s, UInt8)
            #println("> ", Char(t))
            append!(b, t)
            if t == 0x00 || (limit > 0 && length(b) == limit)
                break
            end
        end
    catch EOFError
        println(">EOFError")
    end
    return b
end

function read_some_into(s::IO, b::AbstractArray, limit=length(b))
    try
        i = 1
        while true
            if i > length(b) || (limit != length(b) && i > limit)
                return i - 1
            end
            t = Base.read(s, UInt8)
            #println("> ", Char(t))
            b[i] = t
            i += 1
            if t == 0x00
                return i - 1
            end
        end
    catch EOFError
        println(">EOFError")
    end
end

function receive(c::ChunkedInputBuffer, s::IO, n::Integer)
    try
        new_extent = c.extent + n
        overflow = new_extent - length(c.data)
        if overflow > 0
            if recycle(c)
                return receive(c, s, n)
            end
            data = read_some(s, n)
            data_size = length(data)
            new_extent = c.extent + data_size

            if new_extent > length(c.data)
                resize!(c.data, new_extent)
            end

            c.data[(c.extent+1):new_extent] = data #index
            c.data_view = Base.view(c.data, :)
        else
            data_size = read_some_into(s, Base.view(c.data_view, (c.extent+1):new_extent)) #index
            new_extent = c.extent + data_size
        end
        c.extent = new_extent
        return data_size
    catch e
        println("receive error")
        return -1
    end
end

# Chunked Input Buffer ------------ #

function receive_message(c::ChunkedInputBuffer, s::IO, n::Integer)
    while !frame_message(c)
        received = receive(c, s, n)
        if received <= 0
            return received
        end
    end
    return 1
end

# Chunked Input Buffer ------------ #

function recycle(c::ChunkedInputBuffer)
    origin = c.origin
    if origin == 0
        return false
    end
    available = c.extent - origin
    ### here ###
    c.data[1:available] = c.data[(origin+1):c.extent] #index
    c.extent = available
    c.origin = 0
    return true
end

# Chunked Input Buffer ------------ #

function frame_message(c::ChunkedInputBuffer)
    if c.frame !== nothing
        discard_message(c)
    end
    panes = Tuple{Integer, Integer}[]
    p = origin = c.origin
    extent = c.extent
    while p < extent
        available = extent - p
        if available < 2
            break
        end

        # manually unpack
        # chunk_size, = struct_unpack(">H", c.view[p:(p+2)])
        _chunk_size = reinterpret(UInt16, c.data_view[(p+1):(p+2)]) #index
        chunk_size = ntoh(_chunk_size[1])

        p += 2
        if chunk_size == 0
            c.limit = p
            c.frame = MessageFrame(Base.view(c.data_view, (origin+1):c.limit), panes) #index
            return true
        end
        q = p + chunk_size
        push!(panes, (p - origin, q - origin))
        p = q
    end
    return false
end

# Chunked Input Buffer ------------ #

function discard_message(c::ChunkedInputBuffer)
    if c.frame !== nothing
        close(c.frame)
        c.origin = c.limit
        c.limit = -1
        c.frame = nothing
    end
end

# Chunked Input Buffer ------------ #

function Base.show(io::IO, m::ChunkedInputBuffer)
  Base.show(io, m.data)
end

# Chunked Input Buffer ------------ #
# # # # # # # # # # # # # # # # # # #

#####################################
#       Chunked Output Buffer       #
#####################################

mutable struct ChunkedOutputBuffer
    max_chunk_size::Integer
    header::Integer
    _start::Integer
    _end::Integer
    data::Array{UInt8}

    ChunkedOutputBuffer(;capacity::Integer=1048576, max_chunk_size::Integer=16384) = new(max_chunk_size, 0, 2, 2, zeros(UInt8, capacity))
end

# Chunked Output Buffer ----------- #

function clear(c::ChunkedOutputBuffer)
    c.header = 0
    c._start = 2
    c._end = 2
    c.data[1:2] = [0x00, 0x00] #index
end

# Chunked Output Buffer ----------- #

function Base.write(c::ChunkedOutputBuffer, b::AbstractArray{UInt8})
    to_write = length(b)
    max_chunk_size = c.max_chunk_size
    pos = 0
    while to_write > 0
        chunk_size = c._end - c._start
        remaining = max_chunk_size - chunk_size
        if remaining == 0 || remaining < to_write <= max_chunk_size
            chunk(c)
        else
            wrote = min(to_write, remaining)
            new_end = c._end + wrote

            if new_end > length(c.data)
                resize!(c.data, new_end)
            end

            c.data[(c._end+1):new_end] = b[(pos+1):(pos+wrote)] #index
            c._end = new_end
            pos += wrote

            new_chunk_size = UInt16(c._end - c._start)
            # c.data[c.header:(c.header + 2)] = struct_pack(">H", new_chunk_size)
            _new_chunk_size = hton(new_chunk_size)
            c.data[(c.header+1):(c.header + 2)] = reinterpret(UInt8, [_new_chunk_size]) # index

            to_write -= wrote
        end
    end
end

# Chunked Output Buffer ----------- #

function chunk(c::ChunkedOutputBuffer)
    c.header = c._end
    c._start = c.header + 2
    c._end = c._start

    if c._start > length(c.data)
        resize!(c.data, c._start)
    end
    c.data[(c.header+1):c._start] = [0x00,0x00] #index
end

# Chunked Output Buffer ----------- #

function Base.view(c::ChunkedOutputBuffer)
    _end = c._end
    chunk_size = _end - c._start
    if chunk_size == 0
        return Base.view(c.data, 1:c.header) #index
    else
        return Base.view(c.data, 1:_end) #index
    end
end

# Chunked Output Buffer ----------- #
# # # # # # # # # # # # # # # # # # #

max_stars_repo_name: virtualgraham/JuliaBolt

max_issues_repo_name: virtualgraham/JuliaBolt

max_forks_repo_name: virtualgraham/JuliaBolt


text:
using Lora

println("    Testing GenericModel constructors...")

println("      Initialization without input arguments...")

m = GenericModel()

θ = BasicContUnvParameter(:θ, 1, signature=:low)
x = Data(:x, 2)
λ = Hyperparameter(:λ, 3)

add_vertex!(m, θ)
add_vertex!(m, x)
add_vertex!(m, λ)

add_edge!(m, x, θ)
add_edge!(m, λ, θ)

println("      Initialization via matrix-based specification of dependencies...")

m = GenericModel([θ, x, λ], [x θ; λ θ])

# println("    Testing plotting of GenericModel...")

# g = convert(GenericGraph, m)
# plot(g)


text:
module SDWBA

using QuadGK
using LinearAlgebra
using SpecialFunctions
using DelimitedFiles

export Scatterer,
	rescale,
	resize,
	interpolate,
	rotate,
	length,
	form_function,
	backscatter_xsection,
	target_strength,
	tilt_spectrum,
	freq_spectrum,
	from_csv,
	to_csv,
	Models

import Base: copy, length

include("Scatterer.jl")
include("Models.jl")

end # module


text:
# ------------------------------------------------------------------------------
# ASM preconditioner
# ------------------------------------------------------------------------------

"""
```julia
ASM(operator)
```

ASM diagonal preconditioner for finite element operators

# Arguments:
- `operator`: finite element operator to precondition

# Returns:
- ASM preconditioner object

# Example:
```jldoctest
# setup
mesh = Mesh2D(1.0, 1.0);
mass = GalleryOperator("mass", 4, 4, mesh);

# preconditioner
asm = ASM(mass);

# verify
println(asm)
println(asm.operator)

# output
asm preconditioner
finite element operator:
2d mesh:
    dx: 1.0
    dy: 1.0

2 inputs:
operator field:
  tensor product basis:
    numbernodes1d: 4
    numberquadraturepoints1d: 4
    numbercomponents: 1
    dimension: 2
  evaluation mode:
    interpolation
operator field:
  tensor product basis:
    numbernodes1d: 4
    numberquadraturepoints1d: 4
    numbercomponents: 1
    dimension: 2
  evaluation mode:
    quadratureweights

1 output:
operator field:
  tensor product basis:
    numbernodes1d: 4
    numberquadraturepoints1d: 4
    numbercomponents: 1
    dimension: 2
  evaluation mode:
    interpolation
```
"""
mutable struct ASM <: AbstractPreconditioner
    # data never changed
    operator::Operator

    # data empty until assembled
    subdomaininverse::AbstractArray{Float64, 2}

    Qs::AbstractArray{Float64, 2}

    extendednodecoordinatedifferences::AbstractArray{Float64}

    # inner constructor
    ASM(operator::Operator) = new(operator)
end

# printing
# COV_EXCL_START
Base.show(io::IO, preconditioner::ASM) = print(io, "asm preconditioner")
# COV_EXCL_STOP

# ------------------------------------------------------------------------------
# data for computing symbols
# ------------------------------------------------------------------------------

"""
```julia
getsubdomaininverse(preconditioner)
```

Compute or retrieve the inverse of the symbol matrix diagonal for a ASM
    preconditioner

# Returns:
- Symbol matrix diagonal inverse for the operator

# Example:
```jldoctest
# setup
mesh = Mesh1D(1.0);
diffusion = GalleryOperator("diffusion", 3, 3, mesh);

# preconditioner
asm = ASM(diffusion)

# note: either syntax works
diagonalinverse = LFAToolkit.getsubdomaininverse(asm);
diagonalinverse = asm.subdomaininverse;

# verify
@assert diagonalinverse ≈ [6/7 0; 0 3/4]

# output

```
"""
function getsubdomaininverse(preconditioner::ASM)
    # assemble if needed
    if !isdefined(preconditioner, :subdomaininverse)
        # WARNING: hard-coded to 1d
        elementmatrix = preconditioner.operator.elementmatrix
        Nq, _ = size(elementmatrix)
        p = Nq-1
        Nqe = Nq+2
        pe = Nqe-1
        overlappingop = zeros(Nqe,Nqe)
        overlappingop[2:Nqe-1,2:Nqe-1] = elementmatrix
        overlappingop[1:2,1:2] += elementmatrix[p:Nq,p:Nq]
        overlappingop[pe:Nqe,pe:Nqe] += elementmatrix[1:2,1:2]

        subdomaininverse = pinv(overlappingop)

        preconditioner.subdomaininverse = subdomaininverse
    end

    # return
    return getfield(preconditioner, :subdomaininverse)
end

function getextendednodecoordinatedifferences(preconditioner::ASM)
    # assemble if needed
    if !isdefined(preconditioner, :extendednodecoordinatedifferences)
        # setup for computation
        inputcoordinates = preconditioner.operator.inputcoordinates
        outputcoordinates = preconditioner.operator.outputcoordinates
        dimension = preconditioner.operator.inputs[1].basis.dimension
        lengths = [
            max(inputcoordinates[:, d]...) - min(inputcoordinates[:, d]...) for
            d = 1:dimension
        ]

        # fill matrix
        numberrows, numbercolumns = size(preconditioner.operator.elementmatrix)
        Nqe = numberrows+2

        extendedinputcoordinates = zeros(Nqe, dimension)
        extendedoutputcoordinates = zeros(Nqe, dimension)

        extendedinputcoordinates[2:numberrows+1,:] = inputcoordinates[:,:]
        extendedoutputcoordinates[2:numberrows+1,:] = outputcoordinates[:,:]

        extendedinputcoordinates[1,:] = (
            inputcoordinates[1,:] - (
                inputcoordinates[numberrows,:] - inputcoordinates[numberrows-1,:]
            )
        )
        extendedoutputcoordinates[1,:] = (
            outputcoordinates[1,:] - (
                outputcoordinates[numberrows,:] - outputcoordinates[numberrows-1,:]
            )
        )

        extendedinputcoordinates[Nqe,:] = (
            inputcoordinates[numberrows,:] + (
                inputcoordinates[2,:] - inputcoordinates[1,:]
            )
        )
        extendedoutputcoordinates[Nqe,:] = (
            outputcoordinates[numberrows,:] + (
                outputcoordinates[2,:] - outputcoordinates[1,:]
            )
        )

        extendednodecoordinatedifferences = zeros(Nqe, Nqe, dimension)

        for i = 1:Nqe, j = 1:Nqe, k = 1:dimension
            extendednodecoordinatedifferences[i, j, k] =
                (extendedinputcoordinates[j, k] - extendedoutputcoordinates[i, k]) / lengths[k]
        end

        preconditioner.extendednodecoordinatedifferences = extendednodecoordinatedifferences
    end

    # return
    return getfield(preconditioner, :extendednodecoordinatedifferences)
end

function getQs(preconditioner::ASM)
    # assemble if needed
    if !isdefined(preconditioner, :Qs)
        # WARNING: hard-coded to 1d
        elementmatrix = preconditioner.operator.elementmatrix
        Nq, _ = size(elementmatrix)
        p = Nq-1
        Nqe = Nq+2
        pe = Nqe-1

        Qs = zeros(Nqe,Nq)
        Qs[1,Nq-1] = 1.0
        Qs[Nqe,2] = 1.0
        identityMat = Matrix{Float64}(I, Nq, Nq)
        Qs[2:Nqe-1,1:Nq] = identityMat[:,:]

        preconditioner.Qs = Qs
    end

    # return
    return getfield(preconditioner, :Qs)
end

# ------------------------------------------------------------------------------
# get/set property
# ------------------------------------------------------------------------------

function Base.getproperty(preconditioner::ASM, f::Symbol)
    if f == :subdomaininverse
        return getsubdomaininverse(preconditioner)
    elseif f == :Qs
        return getQs(preconditioner)
    elseif f == :extendednodecoordinatedifferences
        return getextendednodecoordinatedifferences(preconditioner)
    else
        return getfield(preconditioner, f)
    end
end

function Base.setproperty!(preconditioner::ASM, f::Symbol, value)
    if f == :operator
        throw(ReadOnlyMemoryError()) # COV_EXCL_LINE
    else
        return setfield!(preconditioner, f, value)
    end
end

# ------------------------------------------------------------------------------
# compute symbols
# ------------------------------------------------------------------------------

"""
```julia
computesymbols(preconditioner, ω, θ)
```

Compute or retrieve the symbol matrix for a ASM preconditioned operator

# Arguments:
- `preconditioner`: ASM preconditioner to compute symbol matrix for
- `ω`:              Nothing
- `θ`:              Fourier mode frequency array (one frequency per dimension)

# Returns:
- Symbol matrix for the ASM preconditioned operator

# Example:
```jldoctest
using LinearAlgebra

for dimension in 1:3
    # setup
    mesh = []
    if dimension == 1
        mesh = Mesh1D(1.0);
    elseif dimension == 2
        mesh = Mesh2D(1.0, 1.0);
    elseif dimension == 3
        mesh = Mesh3D(1.0, 1.0, 1.0);
    end
    diffusion = GalleryOperator("diffusion", 3, 3, mesh);

    # preconditioner
    asm = ASM(diffusion);

    # compute symbols
    A = computesymbols(asm, [], π*ones(dimension));

    # verify
    using LinearAlgebra;
    eigenvalues = real(eigvals(A));
    if dimension == 1
        @assert max(eigenvalues...) ≈ 1/7
    elseif dimension == 2
        @assert min(eigenvalues...) ≈ -1/14
    elseif dimension == 3
        @assert min(eigenvalues...) ≈ -0.33928571428571486
    end
end

# output

```
"""
function computesymbols(preconditioner::ASM, ω::Array, θ::Array)

    rowmodemap = preconditioner.operator.rowmodemap
    columnmodemap = preconditioner.operator.columnmodemap
    dimension = preconditioner.operator.dimension
    elementmatrix = preconditioner.operator.elementmatrix
    Minv = preconditioner.subdomaininverse
    Qs = preconditioner.Qs
    extendednodecoordinatedifferences = preconditioner.extendednodecoordinatedifferences
    numberrows, numbercolumns = size(Minv)
    Nq = numberrows-2

    symbolmatrixnodes = zeros(ComplexF64, numberrows, numbercolumns)
    for i = 1:numberrows, j = 1:numbercolumns
        symbolmatrixnodes[i, j] =
            Minv[i, j] *
            ℯ^(im * sum([θ[k] * extendednodecoordinatedifferences[i, j, k] for k = 1:dimension]))
    end
    symbolmatrixmodes = rowmodemap * Qs' * symbolmatrixnodes * Qs * columnmodemap

    A = computesymbols(preconditioner.operator, θ)
    return I - symbolmatrixmodes * A
end

# ------------------------------------------------------------------------------


text:
using OnlineAI

export
    SpikeTrains

# plot the nodes and connections of an ANN from OnlineAI
# node colors correspond to activation level
@recipe function plot(net::NeuralNet)
    :legend     --> :none
    :colorbar   --> :none
    :xlims      --> (1, length(net.layers)+1.3)
    :ylims      --> (-1, 1)
    :size       --> (800,800)
    :linealpha  --> 0.5
    :linecolor  --> :black
    :markersize --> 15

    # these will be per-series
    linetype = []
    marker_z = []
    xlist = []
    ylist = []

    # get the node data
    ns = vcat(net.layers[1].nin, [l.nout for l in net.layers])
    ys = map(n->n>1 ? linspace(0.9,-0.9,n) : zeros(1), ns)
    xs = [i*ones(n) for (i,n) in enumerate(ns)]

    # add the edges
    for i in 2:length(ns)
        x1, x2 = xs[i-1:i]
        y1, y2 = ys[i-1:i]
        pts = Tuple{Float64,Float64}[]
        for i1 in 1:length(x1)
            for i2 in 1:length(x2)
                push!(pts, (x1[i1], y1[i1]))
                push!(pts, (x2[i2], y2[i2]))
            end
        end

        push!(linetype, :path)
        push!(marker_z, nothing)
        push!(xlist, Float64[pt[1] for pt in pts])
        push!(ylist, Float64[pt[2] for pt in pts])
    end

    # add the nodes
    for i in 2:length(ns)
        push!(linetype, :scatter)
        push!(marker_z, net.layers[i-1].a)
        push!(xlist, xs[i])
        push!(ylist, ys[i])
    end

    # set the linetype and marker_z args using row-vectors
    :linetype --> reshape(linetype, 1, length(linetype)), :force
    :marker_z --> reshape(marker_z, 1, length(marker_z)), :force

    # return the args
    xlist, ylist
end

# -------------------------------------------------------

# TODO: should this object should hold the data, and make a recipe on it??

type SpikeTrains
    n::Int
    plt
end

function SpikeTrains(n::Integer; kw...)
    plt = plot(n; leg=false, yticks=nothing, kw...)
    SpikeTrains(n, plt)
end

const _halfheight = 0.4
function Base.push!(spiketrains::SpikeTrains, idx::Integer, t::Real)
    append!(spiketrains.plt, idx, Float64[NaN, t, t],
            Float64[NaN, idx + _halfheight, idx - _halfheight])
    spiketrains
end
function Base.push!{T<:Real}(spiketrains::SpikeTrains, ts::AbstractVector{T})
    for i in 1:length(ts)
        push!(spiketrains, i, ts[i])
    end
    spiketrains
end


text:
@testset "'EB' ................................. " begin

  ss = SampleSpace(5:20, 60)
  prior = p -> Distributions.pdf(Distributions.Beta(2 , 5), p)
  params = EB(ss, .2, .3, prior, 16*10^6, 50*10^3)
  design = optimaldesign(10, params, solver)
  @test power(design, .2) <= .05

  params = EB(ss, .2, .3, prior, 16*10^6, 50*10^3, a = 21, b = 9)
  design = optimaldesign(10, params, solver)
  @test power(design, .2) <= .05

  params = EB(ss, .2, .3, prior, 16*10^6, 50*10^3, a = Inf, b = Inf, targetpower = 0.8)
  design = optimaldesign(10, params, solver)
  @test power(design, .2) <= .05

end

max_stars_repo_name: kkmann/binaryTwoStageDesigns

max_issues_repo_name: kkmann/binaryTwoStageDesigns

max_forks_repo_name: kkmann/binaryTwoStageDesigns


text:
"""
    incrementalProblem(c, α, x, X, pd)

Solves incremental problem with specified costs `c`, parameter ``α \\in [0, 1]``,
first stage solutions `x` and a list of constraints `X` defining a set of feasible solutions.
It is subproblem of of [`evaluationProblem`](@ref) and [`adversarialProblem`](@ref).

Check section 4 _Solving the problems by MIP formulations_ of [publication](https://arxiv.org/abs/1811.06719)
for more information about this algorithm.

# Arguments
- `c`: is a vector of a nonnegative nominal second stage costs.
- `α`: fixed number belonging to ``[0, 1]``
- `x`: first stage solution.
- `X`: is a set of feasible solutions represented as a list functions, each of which
    accepts a list of JuMP variables as an argument and returns a JuMP linear constraint.
- `pd`: an instance of [`ProblemDescriptor`](@ref)
"""
function incrementalProblem(c, α, x, X, pd)
    n = size(c, 1)

    model = Model(solver=CplexSolver(CPXPARAM_ScreenOutput = getProperty("incrementalProblem.cplexLogging")))
    if ndims(c) == 1
        @variable(model, y[1:n], Bin)
        if hasEqualCardinalityProperty(pd)
            @variable(model, z[1:n], Bin)
        end
    elseif ndims(c) == 2
        m = size(c, 2)
        @variable(model, y[1:n, 1:m], Bin)
        if hasEqualCardinalityProperty(pd)
            @variable(model, z[1:n, 1:m], Bin)
        end
    else
        throw(ArgumentError("Unsupported number of dimensions ($(ndims(c)))"))
    end

    @objective(model, Min, vecdot(c, y))

    if hasEqualCardinalityProperty(pd)
        k = ceil(n * (1 - α))
        @constraint(model, sum(z) >= k)
        @constraint(model, z .<= y)
        @constraint(model, z .<= x)
    else
        @constraint(model, vecdot(x, 1 - y) <= α * sum(x))
    end

    # Constraints defining set of feasible solutions
    for constraint in X
        JuMP.addconstraint(model, constraint(y))
    end

    status = solve(model)
    return (getvalue(y), getobjectivevalue(model)) # InexactError is thrown here while solving computeAdversarialLowerBound -> evaluationProblem
end


text:
#################################################################
### This test was lifted from the vingette for the R gmm package.
#################################################################
## library(gmm)

## g1 <- function(tet,x) {
##     m1 <- (tet[1]-x)
##     m2 <- (tet[2]^2 - (x - tet[1])^2)
##     m3 <- x^3-tet[1]*(tet[1]^2+3*tet[2]^2)
##     f <- cbind(m1,m2,m3)
##     return(f)
## }

## Dg <- function(tet,x)
## {
## G <- matrix(c(1, 2*(-tet[1]+mean(x)),
##               -3*tet[1]^2-3*tet[2]^2, 0,
##               2*tet[2], -6*tet[1]*tet[2]),
##             nrow=3,ncol=2)
## return(G)
## }

## x <- as.numeric(as.vector(read.csv2('rand_norm.csv',
##                           header=FALSE)[["V1"]]))
## res <-gmm(g1, x, c(mu = 1.0, sig = 1.0), grad = Dg)

using Distributions

srand(42)
x = rand(Normal(4, 2), 1000)

function g(θ, x)
    m1 = θ[1] - x
    m2 = θ[2]^2 - (x - θ[1]).^2
    m3 = x.^3 - θ[1].*(θ[1]^2 + 3*θ[2]^2)
    return [m1 m2 m3]
end

step_1     = gmm(g, [1.0, 1.0], eye(3), data=x)
step_2_hac = gmm(g, coef(step_1),
                 optimal_W(step_1, QuadraticSpectralKernel(0.91469)), data=x)
step_2_iid = gmm(g, coef(step_1), optimal_W(step_1, HC0()),  data=x)

step_qs_mgr = gmm(g, [1.,1.], eye(3), data = x, mgr = TwoStepGMM(QuadraticSpectralKernel(0.91469)))

step_iid_mgr = gmm(g, [1.,1.], eye(3), data = x, mgr = TwoStepGMM(HC0()))