text:
#-------------------------------------------------------------------------------
#general wrapper for interpolating profiles in log pressure coordinates

export AtmosphericProfile

struct AtmosphericProfile{Q}
    ϕ::LinearInterpolator{Q,NoBoundaries}
end

function Base.show(io::IO, ::AtmosphericProfile{Q}) where {Q}
    print(io, "AtmosphericProfile{$Q}")
end

Base.copy(x::AtmosphericProfile) = AtmosphericProfile(copy(x.ϕ))

function AtmosphericProfile(P::AbstractVector{<:Real}, y::AbstractVector{Q}) where {Q<:Real}
    @assert length(P) == length(y) "cannot form AtmosphericProfile with unequal numbers of points"
    #pressure in ascending order
    idx = sortperm(P)
    P = collect(Q, P[idx])
    y = y[idx]
    ϕ = LinearInterpolator(log.(P), y, NoBoundaries())
    AtmosphericProfile(ϕ)
end

(x::AtmosphericProfile)(P) = x.ϕ(log(P))

#-------------------------------------------------------------------------------
#constructing generalized hydrostatic pressure profiles and inverting for z

export scaleheight, hydrostatic, altitude
export Hydrostatic

"""
    scaleheight(g, μ, T)

Evaluate the [atmospheric scale height](https://en.wikipedia.org/wiki/Scale_height),

``\\frac{RT}{μg}``

where ``R`` is the [universial gas constant](https://en.wikipedia.org/wiki/Gas_constant).

# Arguments
* `g`: gravitational acceleration [m/s``^s``]
* `μ`: mean molar mass [kg/mole]
* `T`: temperature [K]
"""
scaleheight(g, μ, T) = 𝐑*T/(μ*g)

#parameterized hydrostatic relation in log coordinates
function dlnPdz(z, lnP, param::Tuple)
    #unpack parameters
    Pₛ, g, fT, fμ = param
    #evaluate temperature and mean molar mass [kg/mole]
    P = exp(lnP)
    P < 𝐏ₘᵢₙ && return zero(lnP)
    P = min(P, Pₛ) #don't allow tiny unphysical dips below Pₛ
    T = fT(P)
    μ = fμ(T, P)
    #evaluate derivative
    -μ*g/(𝐑*T)
end

"""
    hydrostatic(z, Pₛ, g, fT, fμ)

Compute the hydrostatic pressure [Pa] at a specific altitude using arbitrary atmospheric profiles of temperature and mean molar mass. This function integrates the hydrostatic relation,

``\\frac{dP}{dz} = \\frac{\\mu g}{R T}

from the surface to a height of ``z``, where ``R`` is the [universial gas constant](https://en.wikipedia.org/wiki/Gas_constant).

# Arguments

* `z`: altitude [m] to compute pressure at
* `Pₛ`: surface pressure [Pa]
* `g`: gravitational acceleration [m/s``^2``]
* `fT`: temperature [K] as a function of pressure, `fT(P)`
* `fμ`: mean molar mass [kg/mole] as a function of pressure and temperature `fμ(T,P)`
"""
function hydrostatic(z, Pₛ, g, fT::T, fμ::U) where {T,U}
    @assert z >= 0 "cannot compute pressure at negative altitude $z m"
    @assert Pₛ > 𝐏ₘᵢₙ "pressure cannot be less than $𝐏ₘᵢₙ Pa"
    #integration parameters
    param = (Pₛ, g, fT, fμ)
    #integrate in log coordinates and return
    exp(radau(dlnPdz, log(Pₛ), zero(z), z, param))
end

"""
    altitude(P, Pₛ, g, fT, fμ)

Compute the altitude [m] at which a specific hydrostatic pressure occurs using arbitrary atmospheric profiles of temperature and mean molar mass. This function applies a root finder to the [`hydrostatic`](@ref) function.

# Arguments

* `P`: pressure [Pa] to compute altitude at
* `Pₛ`: surface pressure [Pa]
* `g`: gravitational acceleration [m/s``^2``]
* `fT`: temperature [K] as a function of pressure, `fT(P)`
* `fμ`: mean molar mass [kg/mole] as a function of pressure and temperature `fμ(T,P)`
"""
function altitude(P, Pₛ, g, fT::T, fμ::U) where {T,U}
    @assert P < Pₛ "surface pressure must be greater than pressure aloft"
    #pressure decreases monotonically, find altitudes bracketing Pₜ
    z₁ = zero(P)
    z₂ = 1e2*one(P)
    P₁ = Pₛ
    P₂ = hydrostatic(z₂, Pₛ, g, fT, fμ)
    while P₂ > P
        z₁ = z₂
        z₂ *= 2
        P₁ = P₂
        P₂ = hydrostatic(z₂, Pₛ, g, fT, fμ)
    end
    #find precise altitude where P = Pₜ
    fₕ(z,::Any) = log(hydrostatic(z, Pₛ, g, fT, fμ)) - log(P)
    regulafalsi(fₕ, z₁, z₂)
end

"""
[Function-like type](https://docs.julialang.org/en/v1/manual/methods/#Function-like-objects) for initializing and evaluating a hydrostatic pressure profile with arbitrary temperature and mean molar mass profiles. A `Hydrostatic` object maps altitude to pressure. Internally, a pressure vs altitude profile is generated and used for interpolation.

# Constructor

    Hydrostatic(Pₛ, Pₜ, g, fT, fμ, N=250)

* `Pₛ`: surface pressure [Pa]
* `Pₜ`: top of profile pressure [Pa]
* `g`: gravitational acceleration [m/s``^2``]
* `fT`: temperature [K] as a function of presssure, `fT(P)`
* `fμ`: mean molar mass [kg/mole] as a function of temperature and pressure, `fμ(T,P)`
* `N`: optional, number of interpolation nodes

For a constant molar mass or temperature, you can use [anonymous functions](https://docs.julialang.org/en/v1/manual/functions/#man-anonymous-functions) directly. For example, to construct a hydrostatic pressure profile for a crude Earth-like atmosphere:

```@example
#moist adiabatic temperature profile
M = MoistAdiabat(288, 1e5, 1040, 1996, 0.029, 0.018, 2.3e6, psatH2O, Ptropo=1e4);
#hydrostatic pressure profile with constant mean molar mass
H = Hydrostatic(1e5, 1, 9.8, M, (T,P)->0.029);
#evaluate pressures at a few different altitudes
H.([0, 1e3, 1e4])
```
"""
struct Hydrostatic{T}
    ϕ::LinearInterpolator{T,WeakBoundaries}
    zₜ::T
end

function Hydrostatic(Pₛ, Pₜ, g, fT::T, fμ::U, N::Int=100) where {T,U}
    Pₛ, Pₜ = promote(Pₛ, Pₜ)
    #find the altitude corresponding to Pₜ
    zₜ = altitude(Pₜ, Pₛ, g, fT, fμ)
    #interpolation knots and output array
    z = logrange(zero(zₜ), zₜ, N)
    lnP = zeros(typeof(Pₛ), N)
    #integration parameters
    param = (Pₛ, g, fT, fμ)
    #integrate to get a full pressure profile
    radau!(lnP, z, dlnPdz, log(Pₛ), zero(zₜ), zₜ, param)
    #construct and return
    Hydrostatic(LinearInterpolator(z, lnP, WeakBoundaries()), zₜ)
end

(H::Hydrostatic)(z) = exp(H.ϕ(z))

"""
    altitude(H::Hydrostatic, P)

Compute the altitude at which a specific pressure occurs in a [`Hydrostatic`](@ref) pressure profile. A root finder is applied to the object.
"""
function altitude(H::Hydrostatic, P::Real)::Float64
    regulafalsi((z,p) -> log(H(z)) - log(P), zero(H.zₜ), H.zₜ)
end

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

#general function for adiabat with one condensible in bulk non-condensible
function dTdP(P, T, cₚₙ, cₚᵥ, μₙ, μᵥ, L, psat::F) where {F}
    #molar mixing ratio of condensible
    α = psat(T)/P
    #specific gas constants
    Rₙ = 𝐑/μₙ
    Rᵥ = 𝐑/μᵥ
    #numerator
    𝑵 = 1.0 + α*L/(Rₙ*T)
    #denominator
    𝑫 = 1.0 + α*(cₚᵥ/cₚₙ + (L/(T*Rᵥ) - 1.0)*L/(cₚₙ*T))
    #final expression
    (T/P)*(Rₙ/cₚₙ)*(𝑵/𝑫)
end

#same function in ω coordinates
function dTdω(ω, T, cₚₙ, cₚᵥ, μₙ, μᵥ, L, psat::F) where {F}
    P = ω2P(ω)
    -dωfac(P)*dTdP(P, T, cₚₙ, cₚᵥ, μₙ, μᵥ, L, psat::F)
end

#slurp up the parameters for integration
dTdω(ω, T, param::Tuple) = dTdω(ω, T, param...)

#-------------------------------------------------------------------------------
#exported wrappers of dTdω

export lapserate, lapse!

#single-condensible moist lapse rate
function lapserate(T, P, cₚₙ, cₚᵥ, μₙ, μᵥ, L, psat::F) where {F}
    dTdP(P, T, cₚₙ, cₚᵥ, μₙ, μᵥ, L, psat::F)
end

#dry lapse rate
function lapserate(T, P, cₚ, μ)
    dTdP(P, T, cₚ, 1.0, μ, 1.0, 0.0, T->0.0)
end

function lapse!(T, P, cₚ, μ)
    @assert length(P) == length(T)
    idx = sortperm(P, rev=true) #pressure sorting in descending order
    for n ∈ 1:length(idx)-1
        i, j = idx[n], idx[n+1]
        #expected lapse rate
        Γₑ = lapserate(T[i], P[i], cₚ, μ)
        #lapse rate of profile
        Γₚ = (T[j] - T[i])/(P[j] - P[i])
        #heat the upper point if needed
        if Γₚ > Γₑ
            T[j] = T[i] + Γₑ*(P[j] - P[i])
        end
    end
end

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

abstract type AbstractAdiabat end

export MoistAdiabat, DryAdiabat
export tropopause

function checkadiabat(Tₛ, Pₛ, Pₜ, Tstrat, Ptropo, smooth)
    @assert Pₛ > Pₜ "Pₛ must be greater than Pₜ"
    @assert Pₜ > 0 "Pₜ must be greater than 0"
    @assert Tstrat >= 0 "stratosphere temperature cannot be negative"
    @assert Ptropo >= 0 "tropopause pressure cannot be negative"
    @assert smooth >= 0 "smoothing distance cannot be negative"
    if Tstrat > 0
        @assert Tstrat < Tₛ "Tstrat cannot be greater than Tₛ"
    end
    if (Tstrat != 0) & (Ptropo != 0)
        throw("Cannot have nonzero Tstrat and Ptropo, must use one or the other")
    end
end

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

"""
[Function-like type](https://docs.julialang.org/en/v1/manual/methods/#Function-like-objects) for initializing and evaluating a dry adiabatic temperature profile. Optionally, a uniform upper atmospheric temperature can be set below a specified temperature or pressure.

# Constructor

    DryAdiabat(Tₛ, Pₛ, cₚ, μ; Tstrat=0.0, Ptropo=0.0, Pₜ=$𝐏ₘᵢₙ)

* `Tₛ`: surface temperature [K]
* `Pₛ`: surface pressure [K]
* `cₚ`: specific heat of the atmosphere [J/kg/K]
* `μ`: molar mass of the atmosphere [kg/mole]

If `Tstrat` is greater than zero, the temperature profile will never drop below that temperature. If `Ptropo` is greater than zero, the temperature profile at pressures lower than `Ptropo` will be equal to the temperature at exactly `Ptropo`. `Tstrat` and `Ptropo` cannot both greater than zero.

`Pₜ` defines the highest pressure [Pa] in the temperature profile. If should generally be small but cannot be zero.

# Example

Once constructed, use a `DryAdiabat` like a function to compute temperature at a given pressure.

```@example
Tₛ = 288; #surface temperature [K]
Pₛ = 1e5; #surface pressure [Pa]
cₚ = 1040; #specific heat of air [J/kg/K]
μ = 0.029; #molar mass of air [kg/mole]

#construct the dry adiabat with an upper atmosphere temperature of 190 K
D = DryAdiabat(Tₛ, Pₛ, cₚ, μ, Tstrat=190);

#temperatures at 40-10 kPa
D.([4e4, 3e4, 2e4, 1e4])
```
"""
struct DryAdiabat{U} <: AbstractAdiabat
    Tₛ::U
    Pₛ::U
    Pₜ::U
    cₚ::U
    μ::U
    Tstrat::U
    Ptropo::U
    #smoothing quantities
    smooth::U #smoothing distance
    T₂::U #temperature at beginning of smoothing
    h₂::U #hermite thing at T₂
end

function Base.show(io::IO, Γ::DryAdiabat{T}) where {T}
    print(io, "DryAdiabat{$T}:\n")
    print(io, "  Tₛ     = $(Γ.Tₛ) K\n")
    print(io, "  Pₛ     = $(Γ.Pₛ) Pa\n")
    print(io, "  cₚ     = $(Γ.cₚ) J/kg/K\n")
    print(io, "  μ      = $(Γ.μ) kg/mole\n")
    if (Γ.Tstrat != 0) & (Γ.Ptropo != 0)
        Tstrat = round(Γ.Tstrat, sigdigits=6)
        print(io, "  Tstrat = $Tstrat K\n")
        Ptropo = round(Γ.Ptropo, sigdigits=6)
        print(io, "  Ptropo = $Ptropo Pa\n")
        print(io, "  smoothing interval of $(Γ.smooth) Pa")
    end
end

function DryAdiabat(Tₛ, Pₛ, cₚ, μ;
                    Tstrat=0.0,
                    Ptropo=0.0,
                    smooth=1e2,
                    Pₜ=𝐏ₘᵢₙ)
    checkadiabat(Tₛ, Pₛ, Pₜ, Tstrat, Ptropo, smooth)
    #fill in Tstrat or Ptropo if needed
    if Tstrat != 0
        Ptropo = regulafalsi((P,_) -> temperature(P, Tₛ, Pₛ, cₚ, μ) - Tstrat, Pₛ, Pₜ)
    elseif Ptropo != 0
        Tstrat = temperature(Ptropo, Tₛ, Pₛ, cₚ, μ)
    end
    #get smoothing connection ready if Ptropo is nonzero
    h₂ = zero(Tₛ)
    T₂ = zero(Tₛ)
    if Ptropo != 0
        P₂ = Ptropo + smooth
        T₂ = temperature(P₂, Tₛ, Pₛ, cₚ, μ)
        T₂′ = lapserate(T₂, P₂, cₚ, μ)
        h₂ = smooth*T₂′
    end
    DryAdiabat(promote(Tₛ, Pₛ, Pₜ, cₚ, μ, Tstrat, Ptropo, smooth, T₂, h₂)...)
end

#direct calculation of raw temperature profile
temperature(P, Tₛ, Pₛ, cₚ, μ) = Tₛ*(P/Pₛ)^(𝐑/(μ*cₚ))

temperature(Γ::DryAdiabat, P) = temperature(P, Γ.Tₛ, Γ.Pₛ, Γ.cₚ, Γ.μ)

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

"""
[Function-like type](https://docs.julialang.org/en/v1/manual/methods/#Function-like-objects) for initializing and evaluating a moist adiabatic temperature profile. Optional uniform upper atmospheric temperature below a specified temperature or pressure.

# Constructor

    MoistAdiabat(Tₛ, Pₛ, cₚₙ, cₚᵥ, μₙ, μᵥ, L, psat; Tstrat=0, Ptropo=0, N=1000, Pₜ=$𝐏ₘᵢₙ)

* `Tₛ`: surface temperature [K]
* `Pₛ`: surface pressure [K]
* `cₚₙ`: specific heat of the non-condensible atmospheric component (air) [J/kg/K]
* `cₚᵥ`: specific heat of the condensible atmospheric component [J/kg/K]
* `μₙ`: molar mass of the non-condensible atmospheric component (air) [kg/mole]
* `μᵥ`: molar mass of the condensible atmospheric component [kg/mole]
* `L`: condsible component's latent heat of vaporization [J/kg]
* `psat`: function defining the saturation vapor pressure for a given temperature, `psat(T)`

If `Tstrat` is greater than zero, the temperature profile will never drop below that temperature. If `Ptropo` is greater than zero, the temperature profile at pressures lower than `Ptropo` will be equal to the temperature at exactly `Ptropo`. `Tstrat` and `Ptropo` cannot both greater than zero.

`Pₜ` defines the highest pressure [Pa] in the temperature profile. If should generally be small but cannot be zero.

The profile is evaluated along a number of pressure values in the atmosphere set by `N`. Those points are then used to construct a cubic spline interpolator for efficient and accurate temperature calculation. Experience indicates that 1000 points is very accurate and also fast.

# Example

Once constructed, use a `MoistAdiabat` like a function to compute temperature at a given pressure.

```@example
Tₛ = 288; #surface temperature [K]
Pₛ = 1e5; #surface pressure [Pa]
cₚₙ = 1040; #specific heat of air [J/kg/K]
cₚᵥ = 1996; #specific heat of H2O [J/kg/K]
μₙ = 0.029; #molar mass of air [kg/mole]
μᵥ = 0.018; #molar mass of H2O [kg/mole]
L = 2.3e6; #H2O latent heat of vaporization [J/kg]

#a saturation vapor pressure function for H2O is built in
psat = psatH2O;

#construct the moist adiabat with a tropopause pressure of 1e4 Pa
M = MoistAdiabat(Tₛ, Pₛ, cₚₙ, cₚᵥ, μₙ, μᵥ, L, psat, Ptropo=1e4);

#temperatures at 30-5 kPa
M.([3e4, 2e4, 1e4, 5e3])
```
"""
struct MoistAdiabat{U} <: AbstractAdiabat
    ϕ::LinearInterpolator{U,WeakBoundaries}
    Pₛ::U
    Pₜ::U
    Tstrat::U
    Ptropo::U
    #smoothing quantities
    smooth::U #smoothing distance
    T₂::U #temperature at beginning of smoothing
    h₂::U #hermite thing at T₂
end

function Base.show(io::IO, Γ::MoistAdiabat{T}) where {T}
    print(io, "MoistAdiabat{$T}:\n")
    print(io, "  Tₛ     = $(Γ(Γ.Pₛ)) K\n")
    print(io, "  Pₛ     = $(Γ.Pₛ) Pa\n")
    if (Γ.Tstrat != 0) & (Γ.Ptropo != 0)
        Tstrat = round(Γ.Tstrat, sigdigits=6)
        print(io, "  Tstrat = $Tstrat K\n")
        Ptropo = round(Γ.Ptropo, sigdigits=6)
        print(io, "  Ptropo = $Ptropo Pa\n")
        print(io, "  smoothing interval of $(Γ.smooth) Pa")
    end
end

function Base.copy(Γ::MoistAdiabat)
    MoistAdiabat(copy(Γ.ϕ), Γ.Pₛ, Γ.Pₜ, Γ.Tstrat, Γ.Ptropo, Γ.smooth, Γ.T₂, Γ.h₂)
end

function MoistAdiabat(Tₛ, Pₛ, cₚₙ, cₚᵥ, μₙ, μᵥ, L, psat::F;
                      Tstrat=0.0,
                      Ptropo=0.0,
                      smooth=1e2,
                      N::Int=100,
                      Pₜ=𝐏ₘᵢₙ) where {F}
    #basic checks
    checkadiabat(Tₛ, Pₛ, Pₜ, Tstrat, Ptropo, smooth)
    #type uniformity
    Tₛ, Pₛ, Pₜ, Tstrat, Ptropo, smooth = promote(Tₛ, Pₛ, Pₜ, Tstrat, Ptropo, smooth)
    #interpolation knots and output vector
    ω₁, ω₂ = P2ω(Pₛ, Pₜ)
    ω = logrange(ω₁, ω₂, N)
    T = zeros(typeof(Tₛ), N)
    #pack the parameters
    param = (cₚₙ, cₚᵥ, μₙ, μᵥ, L, psat)
    #integrate with in-place dense output
    radau!(T, ω, dTdω, Tₛ, ω[1], ω[end], param)
    #interpolator ω pressure coordinates
    ϕ = LinearInterpolator(ω, T, WeakBoundaries())
    #fill in Tstrat or Ptropo if needed
    if Tstrat != 0
        Ptropo = regulafalsi((P,_) -> temperature(ϕ, P) - Tstrat, Pₛ, Pₜ)
    elseif Ptropo != 0
        Tstrat = temperature(ϕ, Ptropo)
    end
    #get smoothing connection ready if Ptropo is nonzero
    h₂ = zero(Tₛ)
    T₂ = zero(Tₛ)
    if Ptropo != 0
        P₂ = Ptropo + smooth
        T₂ = temperature(ϕ, P₂)
        T₂′ = lapserate(T₂, P₂, cₚₙ, cₚᵥ, μₙ, μᵥ, L, psat)
        h₂ = smooth*T₂′
    end
    MoistAdiabat(ϕ, Pₛ, Pₜ, Tstrat, Ptropo, smooth, T₂, h₂)
end

#direct calculation of raw temperature profile
temperature(ϕ::LinearInterpolator, P) = ϕ(P2ω(P))

temperature(Γ::MoistAdiabat, P) = temperature(Γ.ϕ, P)

#------------------------------------
#general operations with an adiabat

#find the pressure corresponding to a temperature, ignoring Tstrat & Ptropo
function pressure(Γ::AbstractAdiabat, T)
    #surface temperature
    Tₛ = temperature(Γ, Γ.Pₛ)
    #TOA temperature (at very tiny pressure)
    Tₜ = temperature(Γ, Γ.Pₜ)
    #T profile should always decrease with P, so demand bracketing
    @assert Tₛ >= T >= Tₜ "temperature $T K out of adiabat range [$(Tₛ),$(Tₜ)] K"
    #find the root with false position
    regulafalsi((P,::Any) -> temperature(Γ, P) - T, Γ.Pₛ, Γ.Pₜ)
end

function (Γ::AbstractAdiabat)(P)
    #return tropopause temperature if P is below Ptropo
    P < Γ.Ptropo && return Γ.Tstrat*one(P)
    #check if P is inside the smoothing region
    if (Γ.Ptropo != 0) & (Γ.smooth != 0)
        P₁ = Γ.Ptropo
        if Γ.Ptropo < P < P₁ + Γ.smooth
            ψ = (P - Γ.Ptropo)/Γ.smooth
            T₁ = Γ.Tstrat
            T₂ = Γ.T₂
            h₂ = Γ.h₂
            #smooth cubic connection
            return ψ^3*(2*T₁ - 2*T₂ + h₂) + ψ^2*(-3*T₁ + 3*T₂ - h₂) + T₁
        end
    end
    #raw temperature from the profile, unadjusted
    T = temperature(Γ, P)
    #return stratosphere temperature if T is below Tstrat
    T < Γ.Tstrat && return Γ.Tstrat
    #ensure positive
    @assert T > 0 "non-positive temperature ($T K) encountered in adiabat at $P Pa"
    return T
end

"""
    tropopause(Γ::AbstractAdiabat)

Compute the temperature [K] and pressure [Pa] at which the tropopause occurs in an adiabatic temperature profile. This function can be called on a `DryAdiabat` or a `MoistAdiabat` if it was constructed with nonzero `Tstrat` or `Ptropo`. Returns the tuple `(T,P)`.
"""
function tropopause(Γ::AbstractAdiabat)
    Γ.Ptropo != 0 && Γ.Tstrat != 0 && return (Γ.Tstrat,Γ.Ptropo)
    error("no stratosphere temperature or pressure has been defined (Tstrat/Ptropo)")
end

#-------------------------------------------------------------------------------
export psatH2O, tsatCO2, ozonelayer

"""
    psatH2O(T)

Compute the saturation partial pressure of water vapor at a given temperature using expressions from

* [Murphy, D. M. & Koop, T. Review of the vapour pressures of ice and supercooled water for atmospheric applications. Q. J. R. Meteorol. Soc. 131, 1539–1565 (2005).](https://rmets.onlinelibrary.wiley.com/doi/10.1256/qj.04.94)

This function uses equation 10 in the paper above when ``T >= 273.15`` K and equation 7 otherwise.
"""
function psatH2O(T)
    a = log(T)
    b = 1/T
    if T >= 273.15
        #equation 10
        c = 53.878 - 1331.22*b - 9.44523*a + 0.014025*T
        d = c*tanh(0.0415*(T - 218.8))
        P = exp(54.842763 - 6763.22*b - 4.21*a + 3.67e-4*T + d)
    else
        #equation 7
        P = exp(9.550426 - 5723.265*b + 3.53068*a - 0.00728332*T)
    end
    return P
end

"""
    tsatCO2(P)

Compute the saturation pressure of carbon dioxide at a certain pressure using equation 19 from

* [Fanale, F. P., Salvail, J. R., Bruce Banerdt, W. & Steven Saunders, R. Mars: The regolith-atmosphere-cap system and climate change. Icarus 50, 381–407 (1982)](https://doi.org/10.1016/0019-1035(82)90131-2)

The equation is inverted to express temperature as a function of pressure.
"""
function tsatCO2(P)
    @assert P <= 518000.0 "Pressure cannot be above 518000 Pa for CO2 saturation temperature"
    A = 1.2264e12 #[Pa]
    B = -3167.8 #[K]
    B/log(P/A)
end

"""
    ozonelayer(P, Cmax=8e-6)

Approximate the molar concentration of ozone in Earth's ozone layer using an 8 ppm peak at 1600 Pa which falls to zero at 100 Pa and 25500 Pa. Peak concentration is defined by `Cmax`. This approximation is discussed in

* Jacob, D. Introduction to Atmospheric Chemistry. (Princeton University Press, 1999).

"""
function ozonelayer(P, Cmax=8e-6)
    P = log(P)
    P₁ = 10.146433731146518 #ln(25500)
    P₂ = 7.3777589082278725 #ln(1600)
    P₃ = 4.605170185988092  #ln(100)
    if P₂ <= P <= P₁
        return Cmax*(P₁ - P)/(P₁ - P₂)
    elseif P₃ <= P <= P₂
        return Cmax*(P - P₃)/(P₂ - P₃)
    end
    return 0.0*P
end

#-------------------------------------------------------------------------------
#miscellaneous stuff

export condensibleprofile, haircut!, rayleighCO2

"""
    condensibleprofile(Γ::AbstractAdiabat, fPₛ)

Create a function defining concentration vs pressure for a condensible with uniform upper-atmosphere (stratosphere) concentration. The new concentration profile is created with reference to an existing adiabatic profile ([`DryAdiabat`](@ref) or [`MoistAdiabat`](@ref)), which must have `Ptropo != 0` or `Tstrato != 0`. Lower atmospheric concentration is determined by the temperature dependent partial pressure function `fPₛ(T)`. The concentration is `P/(fPₛ + P)`, where `P` is the dry/non-condensible pressure.

"""
function condensibleprofile(Γ::AbstractAdiabat, fPₛ::F)::Function where {F}
    #insist on an isothermal stratosphere
    @assert ((Γ.Ptropo != 0) | (Γ.Tstrat != 0)) "adiabat must have isothermal stratosphere"
    #compute tropopause pressure and temperature
    Tₜ, Pₜ = tropopause(Γ)
    #compute saturation partial pressure at tropopause
    Pₛₜ = fPₛ(Tₜ)
    #create concentration function and return it
    let (Pₜ, Pₛₜ) = (Pₜ, Pₛₜ)
        function (T, P)
            if P >= Pₜ
                Pₛ = fPₛ(T)
                C = Pₛ/(Pₛ + P)
            else
                C = Pₛₜ/(Pₜ + Pₛₜ)
            end
            return C
        end
    end
end

#this is already exported in gases.jl
"""
    reconcentrate(G::Gas, Γ::AbstractAdiabat, fPₛ)

Create a new concentration function for a [`Gas`](@ref) and use it to [`reconcentrate`](@the gas). This does not automatically copy gas data.
"""
function reconcentrate(g::Gas, Γ::AbstractAdiabat, fPₛ)::Gas
    #assign the gas a new concentration profile
    reconcentrate(g, condensibleprofile(Γ, fPₛ))
end

export haircut!
"""
    haircut!(T, P, fTₛ)

Put a temperature floor on a temperature profile using the saturation temperature function `fTₛ(P)`.
"""
function haircut!(T, P, fTₛ::F)::Nothing where {F}
    @assert length(T) == length(P)
    for i ∈ eachindex(T)
        Tₛ = fTₛ(P[i])
        if Tₛ > T[i]
            T[i] = Tₛ
        end
    end
    nothing
end

function rayleighCO2(ν, Pₛ, g, θ)
    #convert to wavelength in micrometers
    λ = ν2λ(ν)*1e6
    #formula 2.32 from
    #  Hansen and Travis, Light Scattering in Planetary Atmospheres (1974)
    τ₀ = 1.527*(1/λ^4)*(1 + 0.013*(1/λ^2))
    #conversion from Venus conditions
    τ₀ *= 8.7/(93*𝐀)
    #final optical depth
    τ = (Pₛ/g)*τ₀

    γ = 0.75
    μ = cos(θ)
    β = 1 - exp(-τ/μ)
    f = γ*τ
    R⁻ = ((0.5 - γ*μ)*β + f)/(1 + f)
    R⁺ = f/(1 + f)
    R = 1 - (1 - R⁺)*(1 - R⁻)/((1 - R⁻))
end


text:
# Autogenerated wrapper script for Libcerf_jll for i686-linux-gnu
export libcerf

JLLWrappers.@generate_wrapper_header("Libcerf")
JLLWrappers.@declare_library_product(libcerf, "libcerf.so.1")
function __init__()
    JLLWrappers.@generate_init_header()
    JLLWrappers.@init_library_product(
        libcerf,
        "lib/libcerf.so",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()

ext: jl


text:
abstract type AbstractPriorExample end

function getproperty(prior_ex::AbstractPriorExample, name::Symbol)
    if name == :distribution
       _distribution(prior_ex)
    end
end

@recipe function f(prior_ex::AbstractPriorExample)
    dbn = prior_ex.distribution
    if isa(dbn, MixtureModel)
        components --> false
    end
   xlab --> L"\mu"
   label --> ""
   title --> _pretty_name(prior_ex)
   dbn
end

"""
    AshSpiky, AshNearnormal, AshFlattop, AshSkew, AshBignormal, AshBimodal

Empirical Bayes priors that are used in the simulations of:

> Stephens, M., 2017. False discovery rates: a new deal. Biostatistics, 18(2), pp.275-294.
"""
AshSpiky, AshNearNormal, AshFlatTop, AshSkew, AshBignormal, AshBimodal

abstract type AshPrior <: AbstractPriorExample end

struct AshSpiky <: AshPrior end
_distribution(::AshSpiky) =  MixtureModel([Normal(0,0.25); Normal(0,0.5); Normal(0,1); Normal(0,2)] ,
                                          [0.4; 0.2; 0.2; 0.2 ])
_pretty_name(::AshSpiky) = "Spiky"

struct AshNearNormal <: AshPrior end
_distribution(::AshNearNormal) = MixtureModel([Normal(0,1); Normal(0,2)], [1/3; 2/3])
_pretty_name(::AshNearNormal) = "Near Normal"

struct AshFlatTop <: AshPrior end
_distribution(::AshFlatTop) = MixtureModel([Normal(x,0.5) for x=-1.5:0.5:1.5])
_pretty_name(::AshFlatTop) = "Flat Top"

struct AshSkew <: AshPrior end
_distribution(::AshSkew) =  MixtureModel([Normal(-2,2), Normal(-1,1.5), Normal(0,1), Normal(1,1)],
                                            [1/4; 1/4; 1/3; 1/6])
_pretty_name(::AshSkew) = "Skew"

struct AshBigNormal <: AshPrior end
_distribution(::AshBigNormal) =  Normal(0,4)
_pretty_name(::AshBigNormal) = "Skew"

struct AshBimodal <: AshPrior end
_distribution(::AshBimodal) =  MixtureModel( [Normal(-2,1); Normal(2,1)],
                            [0.5;0.5] )
_pretty_name(::AshBimodal) = "AshBimodal"

"""
    IWUnimod, IWBimod

Empirical Bayes priors that are used in the simulations of:

> Ignatiadis, N. and Wager, S., 2019. Bias-aware confidence intervals for empirical Bayes analysis. arXiv preprint arXiv:1902.02774.
"""
IWUnimod, IWBimod

abstract type IWPrior <: AbstractPriorExample end

struct IWUnimod <: IWPrior end
_distribution(::IWUnimod) =  MixtureModel([ Normal(-0.2,.2), Normal(0,0.9)],[0.7, 0.3])
_pretty_name(::IWUnimod) = "Unimodal"

struct IWBimod <: IWPrior end
_distribution(::IWBimod) =  MixtureModel([Normal(-1.5,.2), Normal(1.5, .2)])
_pretty_name(::IWBimod) = "Bimodal"

"""
    EfronUnifNormal

Empirical Bayes priors that are used in the simulations of:

> Efron, B., 2016. Empirical Bayes deconvolution estimates. Biometrika, 103(1), pp.1-20.
"""
struct EfronUnifNormal <: AbstractPriorExample end
_distribution(::EfronUnifNormal) =  MixtureModel([Uniform(-3.0,3.0), Normal(0.0, 0.5)], [1/8, 7/8])
_pretty_name(::EfronUnifNormal) = "Uniform / Normal-Mixture"


text:
using DSP, DelimitedFiles, LsqFit

function DT_model(t, β::Array{Float64,1}, ρ::Float64, ndet::Float64, nmed::Float64)
	#optimzed semi-infinite model

    n::Float64 = nmed/ndet
    μa::Float64 = β[1]
    μsp::Float64 = β[2]
    D::Float64 = 1/3μsp
	ν::Float64 = 29.9792345/nmed

    Rt = Array{Float64}(undef, length(t))

	if n > 1.0
		A = 504.332889 - 2641.00214n + 5923.699064n^2 - 7376.355814n^3 +
		 5507.53041n^4 - 2463.357945n^5 + 610.956547n^6 - 64.8047n^7
	elseif n < 1.0
		A = 3.084635 - 6.531194n + 8.357854n^2 - 5.082751n^3
	else
		A = 1.0
	end

	zs::Float64 = 1/μsp
	ze::Float64 = 2A*D

	z3m::Float64 = - zs
    z4m::Float64 = 2ze +zs

	Threads.@threads for n in eachindex(t)

   		Rt[n] = -exp(-(ρ^2/(4D*ν*t[n])) - μa*ν*t[n])
   		Rt[n] = Rt[n]/(2*(4π*D*ν)^(3/2)*t[n]^(5/2))

   		Rt[n] = Rt[n]*(z3m*exp(-(z3m^2/(4D*ν*t[n]))) - z4m*exp(-(z4m^2/(4D*ν*t[n]))))

	end

   replace!(Rt, NaN => 0)
   return Rt./maximum(Rt)

end

function get_fit_window(counts, risefactor, tailfactor)

    maxval = maximum(counts)
    ind1::Int = findfirst(x -> x > maxval*risefactor, counts)
    ind2::Int =  findlast(x -> x >maxval*tailfactor, counts)

    return ind1, ind2
end

function pads(u::Array{Float64,1})
    outsize = 2 .* length(u) .- 1
    nffts = nextfastfft(outsize)
    out = zeros(eltype(u), nffts)
end

function get_fftplan(u::Array{Float64,1}, out)
    upad = zeros(eltype(u), length(out))
    copyto!(upad,u)
    pl = plan_rfft(upad)
    upadRfft = pl * upad
    return pl, upadRfft
end

function conv_DT(t, β::Array{Float64,1}, data::input_data, out, pl,upadRfft;
    vpad = zeros(eltype(data.DTOF), length(out)),
    RtDT = Array{Float64}(undef, length(data.IRF)),
    vpadRfft = zeros(eltype(upadRfft), length(upadRfft))
    )

        RtDT = DT_model(data.t, β, data.ρ, data.nmed, data.ndet)
        copyto!(vpad, RtDT)

        vpadRfft = pl * vpad
        vpadRfft .*= upadRfft
        out = irfft(vpadRfft, length(out))

        out = out./maximum(out)

	    tidx = findfirst(x -> x == t[1], data.t)

	    out = out[tidx:tidx+length(t)-1]
	return log.(out)
end

function getfit(input_data, model_params)
    ## get FFT padding and plans
    out::Array{eltype(input_data.IRF)} = pads(input_data.IRF)
    pl, upadRfft = get_fftplan(input_data.IRF, out)

    #get fit windows
    ind1, ind2 = get_fit_window(input_data.DTOF, model_params.risefactor, model_params.tailfactor)

    curve_fit((t, β) -> conv_DT(t, β, input_data, out, pl, upadRfft), input_data.t[ind1:ind2], log.(input_data.DTOF[ind1:ind2]),
    model_params.initparams, lower=model_params.lb, upper=model_params.ub)

end

#= 12/3/2020
@benchmark getfit(data,modelp)
BenchmarkTools.Trial:
  memory estimate:  30.68 MiB
  allocs estimate:  9561
  --------------
  minimum time:     15.015 ms (0.00% GC)
  median time:      16.626 ms (5.95% GC)
  mean time:        18.319 ms (3.91% GC)
  maximum time:     64.462 ms (1.59% GC)
  --------------
  samples:          273
  evals/sample:     1
=#

max_stars_repo_name: sammzer/LightPropagation

max_issues_repo_name: sammzer/LightPropagation

max_forks_repo_name: sammzer/LightPropagation


text:
#=

Two-dimensional ergodic Hidden Markov Model

HMM2d([])

=#

struct HMM2d
end


text:
"""
Based on
https://hyperopt.github.io/hyperopt/
"""
module TestQuantNormal

using Test
using TreeParzen

# Simple quantnormal
# Called HP.QuantNormal multiple number of times to check if the distribution of values
# matched the expected given input parameters and also a test that the values appear at
# correct intervals
d = []
for i in 1:10000
    a = HP.QuantNormal(:x, 0.0, 1.0, 2.0)
    b = TreeParzen.Resolve.node(a, TreeParzen.Trials.ValsDict())
    push!(d, b)
end
@test all(in.(d, [[-4.0, -2.0, 0, 2.0, 4.0]]))

end # module TestQuantNormal
true


text:
libraries(::WidgetTheme) = [style_css()]

Base.@deprecate_binding backend Widgets.backends

const themes = OrderedDict{Symbol, WidgetTheme}(:nativehtml => NativeHTML())

registertheme!(key::Union{Symbol, AbstractString}, val::WidgetTheme) = setindex!(themes, val, Symbol(key))

"""
`settheme!(s::Union{Symbol, AbstractString})`

Set theme of Interact globally. See `availablethemes()` to know what themes are currently available.
"""
function settheme!(s::Union{Symbol, AbstractString})
    theme = get(themes, Symbol(s)) do
        error("Theme $s is not available.")
    end
    settheme!(theme)
end

settheme!(b::WidgetTheme) = isa(Widgets.get_backend(), WidgetTheme) && Widgets.set_backend!(b)
gettheme() = isa(Widgets.get_backend(), WidgetTheme) ? Widgets.get_backend() : nothing
availablethemes() = sort(collect(keys(themes)))
resettheme!() = isa(Widgets.get_backend(), WidgetTheme) && Widgets.reset_backend!()


text:
module CytoscapeJS

using Artifacts

js = joinpath(artifact"cytoscape.min.js", "cytoscape.min.js")

end


text:
## module 可以视为一个namespace的全局作用域

module A
    a = 1
end;

module B
    module C
        c = 2
    end
    println(C.c)
    import ..A
    println(A.a)
end;

module D
    #b = a # 没有变量a
end;

module E
    import ..A
    #A.a = 2 # 不嫩个其他module全局变量赋值
end;

x = 0;
for i = 1:10
    x = i + 1
end
# for-loop中局部x不改变外部变量
println("local x = ",x)

x = 0;
for i = 1:10
    global x
    x = i + 1
end
# global关键字能引用外部全局x
# 但尽量不使用全局变量
println("global x = ",x)

x, y = 1, 2
function foo()
    x = 10 # 声明了x，x为局部变量
    return x + y # 没有声明y，y来自全局
end
println("foo() = ", foo())

# 嵌套函数会改变其父作用域的局部变量
x, y = 1, 2;
function baz()
    x = 2 # introduces a new local
    function bar()
        x = 10       # modifies the parent's x
        return x + y # y is global
    end
    return bar() + x # 12 + 10 (x is modified in call of bar())
end;
# 之所以这样设计是为了允许构建闭包
# 内部函数从包含它的作用域中继承的变量有时被称为被捕获变量
# 被捕获变量可能带来性能问题

println("baz() = ", baz())
println("global (x, y) = ", (x, y))

max_stars_repo_name: SenZhangAI/example

max_issues_repo_name: SenZhangAI/example

max_forks_repo_name: SenZhangAI/example


text:
using Test

include("armstrong-numbers.jl")

@testset "armstrong numbers" begin
    @test  isarmstrong(0)
    @test  isarmstrong(5)
    @test !isarmstrong(10)
    @test !isarmstrong(100)
    @test  isarmstrong(153)
    @test  isarmstrong(9474)
    @test !isarmstrong(9475)
    @test  isarmstrong(9926315)
    @test !isarmstrong(9926314)
end


text:
module GBIF

using HTTP
using JSON
using Dates
using Requires

const gbifurl = "http://api.gbif.org/v1/"
const gbifenums = Dict(
  "basisOfRecord" => [
    "FOSSIL_SPECIMEN",
    "HUMAN_OBSERVATION",
    "LITERATURE",
    "LIVING_SPECIMEN",
    "MACHINE_OBSERVATION",
    "MATERIAL_SAMPLE",
    "OBSERVATION",
    "PRESERVED_SPECIMEN",
    "UNKNOWN"
  ],
  "continent" => [
    "AFRICA",
    "ANTARCTICA",
    "ASIA",
    "EUROPE",
    "NORTH_AMERICA",
    "OCEANIA",
    "SOUTH_AMERICA"
  ],
  "establishmentMeans" => [
    "INTRODUCED",
    "INVASIVE",
    "MANAGED",
    "NATIVE",
    "NATURALISED",
    "UNCERTAIN"
  ],
  "issue" => [
    "BASIS_OF_RECORD_INVALID",
    "CONTINENT_COUNTRY_MISMATCH",
    "CONTINENT_DERIVED_FROM_COORDINATES",
    "CONTINENT_INVALID",
    "COORDINATE_INVALID",
    "COORDINATE_OUT_OF_RANGE",
    "COORDINATE_PRECISION_INVALID",
    "COORDINATE_REPROJECTED",
    "COORDINATE_REPROJECTION_FAILED",
    "COORDINATE_REPROJECTION_SUSPICIOUS",
    "COORDINATE_ROUNDED",
    "COORDINATE_UNCERTAINTY_METERS_INVALID",
    "COUNTRY_COORDINATE_MISMATCH",
    "COUNTRY_DERIVED_FROM_COORDINATES",
    "COUNTRY_INVALID",
    "COUNTRY_MISMATCH",
    "DEPTH_MIN_MAX_SWAPPED",
    "DEPTH_NON_NUMERIC",
    "DEPTH_NOT_METRIC",
    "DEPTH_UNLIKELY",
    "ELEVATION_MIN_MAX_SWAPPED",
    "ELEVATION_NON_NUMERIC",
    "ELEVATION_NOT_METRIC",
    "ELEVATION_UNLIKELY",
    "GEODETIC_DATUM_ASSUMED_WGS84",
    "GEODETIC_DATUM_INVALID",
    "IDENTIFIED_DATE_INVALID",
    "IDENTIFIED_DATE_UNLIKELY",
    "INDIVIDUAL_COUNT_INVALID",
    "INTERPRETATION_ERROR",
    "MODIFIED_DATE_INVALID",
    "MODIFIED_DATE_UNLIKELY",
    "MULTIMEDIA_DATE_INVALID",
    "MULTIMEDIA_URI_INVALID",
    "PRESUMED_NEGATED_LATITUDE",
    "PRESUMED_NEGATED_LONGITUDE",
    "PRESUMED_SWAPPED_COORDINATE",
    "RECORDED_DATE_INVALID",
    "RECORDED_DATE_MISMATCH",
    "RECORDED_DATE_UNLIKELY",
    "REFERENCES_URI_INVALID",
    "TAXON_MATCH_FUZZY",
    "TAXON_MATCH_HIGHERRANK",
    "TAXON_MATCH_NONE",
    "TYPE_STATUS_INVALID",
    "ZERO_COORDINATE"
  ]
 )

#=
HACK this is required because some GBIF strings fail to parse, and I do not know
why.
=#
import Base: convert
function convert(::Type{AbstractString}, t::T) where {T <: Nothing}
  return "<nothing>"
end

# Load the main functions

include("query.jl")

include("types/GBIFTaxon.jl")
export GBIFTaxon

include("types/GBIFRecords.jl")
export GBIFRecord, GBIFRecords

include("types/iterators.jl")
include("types/show.jl")

include("taxon.jl")
export taxon

include("occurrence.jl")
include("paging.jl")
export occurrence, occurrences
export occurrences!

# Extends with DataFrames functionalities
function __init__()
  @require DataFrames="a93c6f00-e57d-5684-b7b6-d8193f3e46c0" include("requires/dataframes.jl")
end

end # module


text:
@add_problem constraints_benchmark function LT_constraints(
    handle_problem!,
    args...,
)
    p = satisfy(; numeric_type = T)

    x = [Variable() for _ in 1:1000]
    for (i, xi) in enumerate(x)
        push!(p.constraints, xi <= 1.0 * i)
    end
    handle_problem!(p)
    return nothing
end

@add_problem constraints_benchmark function LT_constraint(
    handle_problem!,
    args...,
)
    p = satisfy(; numeric_type = T)

    x = Variable(1000)
    push!(p.constraints, x <= collect(1.0:1000.0))
    handle_problem!(p)
    return nothing
end

@add_problem constraints_benchmark function GT_constraints(
    handle_problem!,
    args...,
)
    p = satisfy(; numeric_type = T)

    x = [Variable() for _ in 1:1000]
    for (i, xi) in enumerate(x)
        push!(p.constraints, xi >= 1.0 * i)
    end
    handle_problem!(p)
    return nothing
end

@add_problem constraints_benchmark function GT_constraint(
    handle_problem!,
    args...,
)
    p = satisfy(; numeric_type = T)

    x = Variable(1000)
    push!(p.constraints, x >= collect(1.0:1000.0))
    handle_problem!(p)
    return nothing
end

@add_problem constraints_benchmark function equality_constraints(
    handle_problem!,
    args...,
)
    p = satisfy(; numeric_type = T)

    x = [Variable() for _ in 1:1000]
    for (i, xi) in enumerate(x)
        push!(p.constraints, xi == 1.0 * i)
    end
    handle_problem!(p)
    return nothing
end

@add_problem constraints_benchmark function equality_constraint(
    handle_problem!,
    args...,
)
    p = satisfy(; numeric_type = T)

    x = Variable(1000)
    push!(p.constraints, x == collect(1.0:1000.0))
    handle_problem!(p)
    return nothing
end

@add_problem constraints_benchmark function sdp_constraint(
    handle_problem!,
    args...,
)
    p = satisfy(; numeric_type = T)

    x = Variable(44, 44) # 990 vectorized entries
    push!(p.constraints, x ⪰ 0)
    handle_problem!(p)
    return nothing
end

@add_problem constraints_benchmark function sdp_constraints(
    handle_problem!,
    args...,
)
    p = satisfy(; numeric_type = T)

    x = [Variable(4, 4) for _ in 1:100] # 1000 total vectorized entries
    for v in x
        push!(p.constraints, v ⪰ 0)
    end
    handle_problem!(p)
    return nothing
end


text:
using PrettyPlots, Plots
gr()
#pyplot()
#pgfplots()
# trying to match the driver steering angle
a1=-0.0206672256
a2=-0.0208976376

b1=-0.0064290878
b2=-0.0066664742

c1=0.0162199327
c2=0.0173577308

#L=n.numStatePoints;

L=linspace(0,6,60)
s0_1=c1;s0=c2;
s_k=zeros(L);
for k in 1:length(L)
  s_k[k]=s0+k*(s0-s0_1)
end

plot(L,s0*ones(L),label="s0")
plot!(L,s0_1*ones(L),label="s0_1")
plot!(L,s_k,label="s_k")

driver_obj=integrate!(mdl,n,r.x[:,6];D=sa_param,(:variable=>:control),(:integrand=>:squared),(:integrandAlgebra=>:subtract));


text:
@testset "cdp.jl" begin

    @testset "Deterministic optimal growth" begin
        alpha = 0.65
        beta = 0.95
        f(s, x) = log(x)
        g(s, x, e) = s^alpha - x
        shocks = [1.]
        weights = [1.]
        x_lb(s) = 0
        x_ub(s) = s

        # Analytical solution
        ab = alpha * beta
        c1 = (log(1 - ab) + log(ab) * ab / (1 - ab)) / (1 - beta)
        c2 = alpha / (1 - ab)
        v_star(s) = c1 + c2 * log(s)
        x_star(s) = (1 - ab) * s^alpha

        # Analytical state path
        s_init = 0.1
        ts_length = 25
        s_path_star = Array{Float64}(undef, ts_length)
        s_path_star[1] = s_init
        shock = 1.  # Arbitrary
        for t in 1:ts_length-1
            s_path_star[t+1] = g(s_path_star[t], x_star(s_path_star[t]), shock)
        end

        # Bases
        bases = Basis[]

        # Chebyshev
        n = 30
        s_min, s_max = 0.1, 2.
        push!(bases, Basis(ChebParams(n, s_min, s_max)))

        # Spline
        k = 3
        m = 101
        breaks = m - (k-1)
        s_min, s_max = 0.1, 2.
        push!(bases, Basis(SplineParams(breaks, s_min, s_max, k)))

        methods = [PFI, VFI]

        for basis in bases
            cdp = ContinuousDP(f, g, beta, shocks, weights, x_lb, x_ub, basis)

            for method in methods
                # solve
                tol = sqrt(eps())
                max_iter = 500
                res = @inferred(solve(cdp, method, tol=tol, max_iter=max_iter))

                rtol = 1e-5
                @test isapprox(res.V, v_star.(cdp.interp.S); rtol=rtol)
                @test isapprox(res.X, x_star.(cdp.interp.S); rtol=rtol)

                # set_eval_nodes!
                grid_size = 200
                eval_nodes = collect(range(s_min, stop=s_max, length=grid_size))
                set_eval_nodes!(res, eval_nodes);

                # simulate
                s_path = @inferred(simulate(res, s_init, ts_length))
                atol = 1e-5
                @test isapprox(s_path[end], s_path_star[end]; atol=atol)
            end
        end

        @testset "Test warning" begin
            cdp = ContinuousDP(
                f, g, beta, shocks, weights, x_lb, x_ub, bases[1]
            )
            for max_iter in [0, 1]
                @test_logs (:warn, r".*max_iter.*")
                           solve(cdp, max_iter=max_iter)
           end
        end
    end

end


text:
using Turing, ReverseDiff, Memoization
using DifferentialEquations
using Plots
using Statistics
using StatsPlots
Turing.setadbackend(:reversediff)
Turing.setrdcache(true)
using Distributions
using LinearAlgebra

# here we are simulating the signal of two coupled oscillators
# with amplitudes A1 and A2 of kBT/k and kBT/(k+c) where c is the coupling constant
# D = kBT/gamma
# in terms of the standard parameters of a OU process:
# sigma = sqrt(2*D)
# Theta = D/A

function corr_osc(c)
    # function that returns two time series that are correlated through a coupling coefficient c
    μ = 0.0 # mean is zero
    σ = sqrt(2) # D=1
    Θ1 = 1.0
    Θ2 = 1.0+abs(c)

    W1 = OrnsteinUhlenbeckProcess(Θ1,μ,σ,0.0,1.0)
    W2 = OrnsteinUhlenbeckProcess(Θ2,μ,σ,0.0,1.0)
    prob1 = NoiseProblem(W1,(0.0,100.0))
    prob2 = NoiseProblem(W2,(0.0,100.0))
    sol1 = solve(prob1;dt=0.1)
    sol2 = solve(prob2;dt=0.1)

    # creating the two correlated
    x1 = (sol1.u .+ sol2.u)/2
    if c>0
        x2 = (sol1.u .- sol2.u)/2
    else
        x2 = (sol2.u .- sol1.u)/2
    end
    return x1,x2
end

# Ornstein-Uhlenbeck process
@model ou(rn,T,delta_t) = begin
    ampl ~ Uniform(0.0,5.0)
    b ~ beta(5.0,1.0)

    rn[1] ~ Normal(0,sqrt(ampl))

    for i=2:T
        rn[i] ~ Normal(rn[i-1]*b,sqrt(ampl*(1-b^2)))
    end
end

# Ornstein-Uhlenbeck process of two coupled oscillators
@model ou_corr(rn1,rn2,T,delta_t) = begin
    ampl1 ~ Uniform(0.0,5.0)
    ampl2 ~ Uniform(0.0,5.0)
    d ~ Uniform(0.0,5.0)
    b1 = exp(-delta_t*d/ampl1)
    b2 = exp(-delta_t*d/ampl2)

    rn1[1] ~ Normal(0,sqrt(ampl1))
    rn2[1] ~ Normal(0,sqrt(ampl2))

    for i=2:T
        rn1[i] ~ Normal(rn1[i-1]*b1,sqrt(ampl1*(1-b1^2)))
        rn2[i] ~ Normal(rn2[i-1]*b2,sqrt(ampl2*(1-b2^2)))
    end
end

# Ornstein-Uhlenbeck process with added Gaussian noise
@model oupn(rn,T,delta_t,::Type{R}=Vector{Float64}) where {R} = begin
    ampl ~ Uniform(0.0,5.0)
    b ~ Beta(5.0,1.0)
    noise_ampl ~ Uniform(0.0,1)

    b = exp(-delta_t/tau)
    r = R(undef, T)

    r[1] ~ Normal(0,sqrt(ampl))

    for i=2:T
        r[i] ~ Normal(r[i-1]*b,sqrt(ampl*(1-b^2)))
    end
    rn ~ MvNormal(r,sqrt(noise_ampl))
end

x1, x2 = corr_osc(-0.8)

p1 = Plots.plot(x1)
p1 = Plots.plot!(x2)

pearson = Statistics.cor(x1,x2)
println(pearson)

# lets see whether we can estimate c from the data
y1 = x1 .+ x2
y2 = x1 .- x2
chn = sample(ou_corr(y1,y2,length(y1),0.1), NUTS(0.65), 10000)

print(describe(chn))
p2 = plot(chn)

ampl1 = Array(chn[:ampl1])
ampl2 = Array(chn[:ampl2])

if ampl1>ampl2
    c = (ampl1 .- ampl2)./ampl2
else
    c = (ampl1 .- ampl2)./ampl1
end

println("A1,A2: ",mean(ampl1),",",mean(ampl2))
println("c estimate: ",mean(c),"std: ",std(c))


text:
using HypothesisTests, Test

@testset "White tests" begin

X = [ 4.18      0.77      1.00;
     -3.41      0.73      1.00;
      5.75      0.81      1.00;
      0.05      0.80      1.00;
     -2.18      0.82      1.00;
      3.88      0.81      1.00;
      0.73      0.77      1.00;
      5.70      0.77      1.00;
     -0.69      0.83      1.00;
     -8.14      0.87      1.00;
      5.37      0.99      1.00;
      1.87      0.95      1.00;
      5.76      0.80      1.00;
     -0.79      0.89      1.00;
    -13.23      1.21      1.00;
      3.97      1.26      1.00;
      5.20      0.81      1.00;
      3.16      0.61      1.00;
      6.41      0.53      1.00;
      1.72      0.64      1.00;
      2.20      0.75      1.00;
      1.06      0.95      1.00;
      9.54      0.96      1.00;
     -4.75      1.31      1.00;
     -5.05      1.04      1.00;
      0.48      1.07      1.00;
      3.41      1.21      1.00;
     -2.21      1.08      1.00;
      0.21      1.15      1.00;
     -2.37      1.35      1.00;
     -1.55      1.24      1.00;
     -6.91      1.28      1.00;
     -7.62      1.24      1.00;
      4.81      1.21      1.00;
      3.51      1.07      1.00;
     -3.68      0.87      1.00;
     -3.42      0.80      1.00;
     -6.03      0.92      1.00;
     -1.99      0.98      1.00;
      3.20      1.13      1.00;
     -3.88      1.06      1.00;
     -3.35      0.96      1.00;
     -3.10      1.05      1.00;
     11.14      0.76      1.00;
      1.17      0.51      1.00;
     11.27      0.59      1.00;
      4.56      0.63      1.00;
      0.78      0.67      1.00;
      3.50      0.69      1.00;
      2.40      0.62      1.00;
      2.84      0.63      1.00;
      6.71      0.71      1.00;
      0.63      0.69      1.00;
      3.11      0.67      1.00;
     -3.90      0.74      1.00;
     -0.41      0.76      1.00;
      0.85      0.76      1.00;
     -3.56      0.76      1.00;
      2.26      0.70      1.00;
     -1.78      0.73      1.00;
     -2.06      0.76      1.00;
     -4.62      0.71      1.00;
      0.61      0.73      1.00;
     -0.56      0.81      1.00;
     -6.01      0.78      1.00;
      1.59      0.75      1.00;
     -2.88      0.82      1.00;
     10.44      0.83      1.00;
     -0.82      0.86      1.00;
     -1.01      1.00      1.00;
     -1.80      0.73      1.00;
      1.73      0.64      1.00;
      7.92      0.65      1.00;
      1.11      0.58      1.00;
     -0.79      0.62      1.00;
     -0.94      0.72      1.00;
      4.92      0.66      1.00;
      1.16      0.55      1.00;
     -0.65      0.62      1.00;
     -1.03      0.55      1.00;
     -4.58      0.60      1.00;
      3.79      0.65      1.00;
      6.31      0.61      1.00;
      3.66      0.65      1.00;
      0.42      0.56      1.00;
      6.72      0.53      1.00;
      4.79      0.60      1.00;
     -1.31      0.52      1.00;
      4.59      0.49      1.00;
      0.90      0.52      1.00;
     -6.49      0.52      1.00;
      6.16      0.46      1.00;
     -8.35      0.45      1.00;
      4.47      0.46      1.00;
      1.12      0.39      1.00;
     -3.13      0.49      1.00;
     12.43      0.42      1.00;
      4.36      0.43      1.00;
      1.90      0.47      1.00;
     -2.14      0.44      1.00 ]

e = [ 5.19346;
      2.22296;
      3.82622;
      2.93541;
      0.70453;
      1.98955;
      0.56655;
      1.54306;
      0.57363;
     -0.13346;
      1.68565;
      9.24558;
      7.39317;
      2.45484;
     -4.83724;
     -0.74742;
     -0.22633;
      0.51537;
      2.82822;
      9.35943;
      5.38263;
      6.59691;
      1.79080;
      0.54506;
      4.06661;
     -4.18808;
      1.69414;
      5.40963;
      2.92762;
     -2.67351;
     -2.74613;
     -3.52752;
     -4.27775;
      2.88245;
     -8.12566;
      0.08760;
      2.20495;
     -0.45509;
     -0.06959;
      1.51330;
      1.37878;
     -2.33046;
      0.35730;
    -11.41189;
     -0.48083;
      1.08917;
      2.52333;
      1.73218;
     15.60426;
     -0.15961;
     -0.22174;
     -1.19174;
     13.92022;
      1.16044;
      0.08087;
     -5.81547;
     -5.14199;
     -5.69918;
     -0.36585;
     -3.70294;
      2.83687;
     -1.37631;
     -2.55403;
     -4.10796;
      1.32092;
     -1.64056;
     -4.09963;
     -4.85728;
     -1.45496;
     -3.12612;
     -4.36648;
     -3.23379;
      5.22893;
      5.06730;
     -0.63041;
     -1.70377;
     -5.71335;
     -2.76831;
      2.20443;
     -0.10174;
     -2.36727;
     -5.49859;
     -4.78093;
     -3.36665;
      5.12027;
     -3.03180;
     -2.52035;
      3.10913;
     -3.51388;
     -0.70389;
     -1.01963;
     -8.78989;
      2.95170;
     -3.95464;
     -4.01258;
     -1.92315;
     -4.94209;
      5.29138;
      1.28435;
      0.28832 ]

    atol = 1e-4
    w_test = WhiteTest(X, e, type = :linear)
    @test pvalue(w_test) ≈ 0.1287 atol=atol

    w_test = WhiteTest(X, e, type = :linear_and_squares)
    @test pvalue(w_test) ≈ 0.2774 atol=atol

    w_test = WhiteTest(X, e)
    w_pval = pvalue(w_test)
    @test w_pval ≈ 0.3458 atol=atol

    refstr = """
    White's (or Breusch-Pagan's) test for heteroskedasticity
    --------------------------------------------------------
    Population details:
        parameter of interest:   T*R2
        value under h_0:         0
        point estimate:          $(w_test.lm)

    Test summary:
        outcome with 95% confidence: fail to reject h_0
        one-sided p-value:           $(round(w_pval,digits=4))

    Details:
        T*R^2 statistic:        $(w_test.lm)
        degrees of freedom:     $(w_test.dof)
        type:                   $(w_test.type)
    """

    buffer = IOBuffer()
    show(buffer, w_test)
    str = String(take!(buffer))
    @test str == refstr

    @test_throws DimensionMismatch WhiteTest([rand(3) ones(3)], rand(4))

    @test_throws ArgumentError WhiteTest(ones(4,1), rand(4))

    @test_throws ArgumentError WhiteTest([zeros(4) rand(4)], rand(4))

    @test_throws ArgumentError WhiteTest(rand(4,2), rand(4))

    bp_test = BreuschPaganTest(X, e)
    @test pvalue(bp_test) ≈ 0.1287 atol=atol
    show(IOBuffer(), bp_test)
end


text:
# # Wasserstein non-negative matrix factorisation
#
# In this example, we implement Wasserstein non-negative matrix factorisation (W-NMF)
# following the paper [^RCP16] by Rolet et al.
#
# [^RCP16]: Rolet, Antoine, Marco Cuturi, and Gabriel Peyré. ["Fast dictionary learning with a smoothed Wasserstein loss."](https://marcocuturi.net/Papers/rolet16fast.pdf) Artificial Intelligence and Statistics. PMLR, 2016.
#
# ## Introduction
#
# Matrix decomposition is a classical problem in machine learning. Given a $m \times n$ matrix $X$ representing $n$ repeated $m$-dimensional observations, one may seek matrices $D, \Lambda$ of appropriate dimensions such that
# ```math
# X \approx D \Lambda.
# ```
# For a target rank $k < \min \{ m, n \}$, i.e. $D \in \mathbb{R}^{m \times k}, \Lambda \in \mathbb{R}^{k \times n}$, this problem can be thought of as seeking a low-dimensional linear representation $D \Lambda$ of the potentially high-dimensional dataset $X$.
#
# Lee and Seung [^LS99] observed that the data matrix $X$ is non-negative in many practical applications, and that naturally one may want the factor matrices $D, \Lambda$ to be also non-negative.
#
# [^LS99]: Lee, Daniel D., and H. Sebastian Seung. "Learning the parts of objects by non-negative matrix factorization." Nature 401.6755 (1999): 788-791.
#
# For a given $m \times n$ matrix $X \ge 0$, finding the factor matrices $D \in \mathbb{R}^{m \times k}, \Lambda \in \mathbb{R}^{k \times n}$ such that $X \approx D \Lambda$ with $D \ge 0, \Lambda \ge 0$ is known as the rank-$k$ non-negative matrix factorization (NMF) problem. Typically, such an approximate representation is sought by solving a minimisation problem
# ```math
# \min_{D \in \mathbb{R}^{m \times k}, \Lambda \in \mathbb{R}^{k \times n}, D \ge 0, \Lambda \ge 0} \Phi(X, D \Lambda),
# ```
# where $(X, Y) \mapsto \Psi(X, Y)$ is a loss function defined on matrices. Commonly used choices of $\Phi$ include the squared Frobenius loss $\Phi(X, Y) = \| X - Y \|_F^2$ and Kullback-Leibler divergence $\Phi(X, Y) = \sum_{ij} X_{ij} \log(X_{ij}/Y_{ij})$.
#
# Both these loss functions (and many other common choices) decompose elementwise in their arguments, that is, they can be written as $\Phi(X, Y) = \sum_{ij} f(X_{ij}, Y_{ij})$ for some function $f$ acting on scalars.
#
# Rolet et al. note that pointwise loss functions cannot account for spatial correlations in datasets with underlying geometry, and propose to use entropy-regularised optimal transport as a loss function that is sensitive to the spatial arrangement of the data.
# They argue that for datasets such as images, optimal transport is a more natural choice of loss function, and that it achieves superior performance.
#
# In particular, suppose that the columns of $X$ encode image data, and that $C \in \mathbb{R}^{m \times m}$ encodes the squared Euclidean distances on the imaging domain. The problem that Rolet et al. pose is
# ```math
# \min_{D \in \mathbb{R}^{m \times k}_{\ge 0}, \Lambda \in \mathbb{R}^{k \times n}_{\ge 0}} \sum_{i = 1}^{n} \operatorname{OT}_{\varepsilon}(X_i, (D\Lambda)_i) + \rho_1 E(D) + \rho_2 E(\Lambda),
# ```
# where $\operatorname{OT}_{\varepsilon}(\alpha, \beta)$ is the entropy-regularised optimal transport loss between two probability distributions $\alpha, \beta$ for a cost $C$, and $E$ is an entropy barrier function (a smooth approximation to the non-negativity constraint),
# ```math
# E(A) = \sum_{ij} (A_{ij} \log(A_{ij}) - 1).
# ```
# The parameters $\rho_1, \rho_2$ control how "sharp" the non-negativity constraints are. As $\rho_1, \rho_2 \to 0$, the smoothed constraint approaches the hard non-negativity constraint. Finally, $\varepsilon$ controls the regularisation level for the optimal transport loss.
#
# This example shows how this method can be implemented using functions from the `Dual` submodule of OptimalTransport.jl.
#
# ## Set up prerequisites
#
# Load packages that will be required later.
#
using OptimalTransport
import OptimalTransport.Dual: Dual
using MLDatasets: MLDatasets
using StatsBase
using Plots;
default(; palette=:Set1_3)
using LogExpFunctions
using NNlib: NNlib
using LinearAlgebra
using Distances
using Base.Iterators
using NMF
using Optim
#
# Define `simplex_norm!`, which normalises `x` to sum to `1` along `dims`.
#
function simplex_norm!(x; dims=1)
    return x .= x ./ sum(x; dims=dims)
end
#
# ## Implementation
#
# We now implement Wasserstein-NMF. Rolet et al. split the original non-convex problem into a pair of convex problems, one for $D$ and one for $\Lambda$.
# ```math
# \begin{aligned}
# &\min_{D \in \mathbb{R}^{m \times k}} \sum_{i = 1}^{n} \operatorname{OT}_{\varepsilon}(X_i, (D\Lambda)_i) + \rho_2 E(D), \\
# &\min_{\Lambda \in \mathbb{R}^{k \times n}} \sum_{i = 1}^{n} \operatorname{OT}_{\varepsilon}(X_i, (D\Lambda)_i) + \rho_1 E(\Lambda).
# \end{aligned}
# ```
# For each of these problems, the dual problem can be derived (see Section 3.3 of Rolet et al. for details on how this is done). These turn out to be
# ```math
# \begin{aligned}
# &\min_{G \in \mathbb{R}^{m \times n}} \sum_{i = 1}^n \operatorname{OT}_{\varepsilon}^*(X_i, G_i) + \rho_2 \sum_{i = 1}^{n} E^*(-(G \Lambda^\top)_i / \rho_2) \\
# &\min_{G \in \mathbb{R}^{m \times n}} \sum_{i = 1}^n \operatorname{OT}_{\varepsilon}^*(X_i, G_i) + \rho_1 \sum_{i = 1}^{n} E^*(-(D^\top G)_i / \rho_1).
# \end{aligned}
# ```
# The semi-dual of entropy-regularised optimal transport loss, $\operatorname{OT}_{\varepsilon}$, is implemented as `Dual.ot_entropic_semidual` and its gradient can be computed by `Dual.ot_entropic_semidual_grad`. $E^*$ turns out to be `logsumexp`, for which we implement a wrapper function:
#
function E_star(x; dims=1)
    return logsumexp(x; dims=dims)
end;
#
# The gradient of `logsumexp` is `softmax`, so we define also its gradient:
#
function E_star_grad(x; dims=1)
    return NNlib.softmax(x; dims=1)
end;
#
# Thus, for each problem we may define the dual objective and its gradient. We note that `ot_entropic_semidual(_grad)` automatically broadcasts along columns of its input. There is therefore no need to make multiple calls to the function, thus allowing for more efficient evaluation.
#
function dual_obj_weights(X, K, ε, D, G, ρ1)
    return sum(Dual.ot_entropic_semidual(X, G, ε, K)) + ρ1 * sum(E_star(-D' * G / ρ1))
end
function dual_obj_weights_grad!(∇, X, K, ε, D, G, ρ1)
    return ∇ .= Dual.ot_entropic_semidual_grad(X, G, ε, K) - D * E_star_grad(-D' * G / ρ1)
end
function dual_obj_dict(X, K, ε, Λ, G, ρ2)
    return sum(Dual.ot_entropic_semidual(X, G, ε, K)) + ρ2 * sum(E_star(-G * Λ' / ρ2))
end
function dual_obj_dict_grad!(∇, X, K, ε, Λ, G, ρ2)
    return ∇ .= Dual.ot_entropic_semidual_grad(X, G, ε, K) - E_star_grad(-G * Λ' / ρ2) * Λ
end;
#
# The only remaining part of Wasserstein-NMF to implement is the conversion at optimality from the dual variable $G$ to the primal variables $D, \Lambda$. From the results of Theorems 3 and 4 in Rolet et al., we have for $\Lambda$:
# ```math
# \Lambda_i = \operatorname{softmax}((-D^\top G)_i / \rho_1),
# ```
# and we have for $D$:
# ```math
# D_i = \operatorname{softmax}((-G \Lambda^\top)_i / \rho_2).
# ```
#
function getprimal_weights(D, G, ρ1)
    return NNlib.softmax(-D' * G / ρ1; dims=1)
end
function getprimal_dict(Λ, G, ρ2)
    return NNlib.softmax(-G * Λ' / ρ2; dims=1)
end;
#
# We can now implement functions `solve_weights` and `solve_dict` that solve the respective dual problems for the next iterates of `Λ` and `D`.
#
function solve_weights(X, K, ε, D, ρ1; alg, options)
    opt = optimize(
        g -> dual_obj_weights(X, K, ε, D, g, ρ1),
        (∇, g) -> dual_obj_weights_grad!(∇, X, K, ε, D, g, ρ1),
        zero.(X),
        alg,
        options,
    )
    return getprimal_weights(D, Optim.minimizer(opt), ρ1)
end
function solve_dict(X, K, ε, Λ, ρ2; alg, options)
    opt = optimize(
        g -> dual_obj_dict(X, K, ε, Λ, g, ρ2),
        (∇, g) -> dual_obj_dict_grad!(∇, X, K, ε, Λ, g, ρ2),
        zero.(X),
        alg,
        options,
    )
    return getprimal_dict(Λ, Optim.minimizer(opt), ρ2)
end;

#
# ## Example: noisy univariate Gaussians
#
# We set up each observation as a mixture of 3 Gaussians with means sampled from `N(6, σ), N(0, σ), N(-6, σ)` respectively, and mixture weights sampled uniformly from `[0, 1]`. The resulting mixture model is normalised to sum to 1 on the discrete domain `coord`.
#
f(x, μ, σ) = exp.(-(x .- μ) .^ 2)
coord = range(-12, 12; length=100)
N = 100
σ = 1
X = hcat(
    [
        rand() * f(coord, σ * randn() + 6, 1) +
        rand() * f(coord, σ * randn(), 1) +
        rand() * f(coord, σ * randn() - 6, 1) for _ in 1:N
    ]...,
)
X = simplex_norm!(X);
#
# We visualise the  observations.
#
plot(coord, X; alpha=0.1, color=:blue, title="Input data", legend=nothing)
#
# We can apply NMF with a squared Frobenius loss using the NMF.jl package. We seek `k = 3` components.  This performs poorly, since the pointwise nature of the loss function cannot handle the translational noise in the data.
#
k = 3
result = nnmf(X, k; alg=:multmse)
plot(coord, result.W; title="NMF with Frobenius loss", palette=:Set1_3)
#
# We can now set up a cost matrix corresponding to the domain `coord`.
#
C = pairwise(SqEuclidean(), coord)
C = C / mean(C);
# Specify parameters
ε = 0.025
ρ1, ρ2 = (5e-2, 5e-2);
# Compute Gibbs kernel
K = exp.(-C / ε);
#
# Now we use a random initialisation, where columns of `D` and `Λ` are normalised to sum to 1.
#
D = rand(size(X, 1), k) # dictionary
simplex_norm!(D; dims=1) # norm columnwise
Λ = rand(k, size(X, 2)) # weights
simplex_norm!(Λ; dims=1); # norm rowwise
#
# We now run 10 iterations of Wasserstein-NMF.
#
n_iter = 10
for iter in 1:n_iter
    @info "Wasserstein-NMF: iteration $iter"
    D .= solve_dict(
        X,
        K,
        ε,
        Λ,
        ρ2;
        alg=LBFGS(),
        options=Optim.Options(;
            iterations=250, g_tol=1e-4, show_trace=false, show_every=10
        ),
    )
    Λ .= solve_weights(
        X,
        K,
        ε,
        D,
        ρ1;
        alg=LBFGS(),
        options=Optim.Options(;
            iterations=250, g_tol=1e-4, show_trace=false, show_every=10
        ),
    )
end
#
# We observe that Wasserstein-NMF learns atoms (columns of $D$) corresponding to the three Gaussians used to generate the input data.
#
plot(coord, D; title="NMF with Wasserstein loss", palette=:Set1_3)

# ## Example: image data (MNIST)
# Here we will download MNIST dataset using MLDatasets.jl and downscale each image to 14x14 to allow for faster runtime (since we are running on CPU).
#
sizex, sizey = 28, 28
factor = 2 # downscale factor
Σ = hcat([sum(I(sizex)[:, i:(i + factor - 1)]; dims=2) for i in 1:factor:sizex]...)
sizex, sizey = sizex ÷ factor, sizey ÷ factor
N = 256
x, y = MLDatasets.MNIST.traindata(Float64, sample(1:60_000, N; replace=false))
x = permutedims(x, (2, 1, 3))
x = cat([Σ' * x[:, :, i] * Σ for i in 1:N]...; dims=3)
X = simplex_norm!(reshape(x, (sizex * sizey, :)));
#
# The columns of `X` now correspond to images "flattened" as vectors. We can preview a few images.
#
M = 25
plot(
    [
        heatmap(
            reshape(X[:, i], sizex, sizey);
            legend=:none,
            axis=nothing,
            showaxis=false,
            aspect_ratio=:equal,
            c=:Blues,
            yflip=true,
        ) for i in 1:M
    ]...;
    layout=(5, M ÷ 5),
    plot_title="Input images",
)

# Now we can set up `coord`, cost matrix `C`, and specify parameters.
#
coord = reshape(collect(product(1:sizex, 1:sizey)), :)
C = pairwise(SqEuclidean(), coord)
C = C / mean(C);
ε = 0.0025
ρ1, ρ2 = (5e-3, 5e-3);
# We compute the Gibbs kernel from `C`:
K = exp.(-C / ε);
#
# Let us aim to learn `k = 25` atoms.
k = 25;
#
# Initialise again randomly
#
D = rand(size(X, 1), k) # dictionary
simplex_norm!(D; dims=1) # norm columnwise
Λ = rand(k, size(X, 2)) # weights
simplex_norm!(Λ; dims=1); # norm rowwise
#
# We now run 15 iterations of Wasserstein-NMF.
#
n_iter = 15
for iter in 1:n_iter
    @info "Wasserstein-NMF: iteration $iter"
    D .= solve_dict(
        X,
        K,
        ε,
        Λ,
        ρ2;
        alg=LBFGS(),
        options=Optim.Options(;
            iterations=250, g_tol=1e-4, show_trace=false, show_every=10
        ),
    )
    Λ .= solve_weights(
        X,
        K,
        ε,
        D,
        ρ1;
        alg=LBFGS(),
        options=Optim.Options(;
            iterations=250, g_tol=1e-4, show_trace=false, show_every=10
        ),
    )
end
#
# We can inspect the atoms learned (columns of `D`):
#

plot(
    [
        heatmap(
            reshape(D[:, i], sizex, sizey);
            legend=:none,
            axis=nothing,
            showaxis=false,
            aspect_ratio=:equal,
            c=:Blues,
            yflip=true,
        ) for i in 1:k
    ]...;
    layout=(5, k ÷ 5),
    plot_title="Learned atoms",
)

#
# Finally, we can look at the images constructed by the low-rank approximation `DΛ` and compare to the original images that we previewed earlier.
#
X_hat = D * Λ
plot(
    [
        heatmap(
            reshape(X_hat[:, i], sizex, sizey);
            legend=:none,
            axis=nothing,
            showaxis=false,
            aspect_ratio=:equal,
            c=:Blues,
            yflip=true,
        ) for i in 1:M
    ]...;
    layout=(5, M ÷ 5),
    plot_title="Low rank approximation",
)


text:
using Test
using SOFA

@testset "SOFA" begin
    testdir = joinpath(dirname(@__DIR__), "test")
    @testset "sofa" begin
        include(joinpath(testdir, "test_sofa.jl"))
    end
end


text:
using Test
using PETSc, MPI, LinearAlgebra, SparseArrays

MPI.Initialized() || MPI.Init()

@testset "SNES" begin

  comm    = MPI.COMM_WORLD
  mpirank = MPI.Comm_rank(comm)
  mpisize = MPI.Comm_size(comm)

  for petsclib in PETSc.petsclibs
      PETSc.initialize(petsclib)
      PetscScalar = PETSc.scalartype(petsclib)
      PetscInt    = PETSc.inttype(petsclib)

      # This shows some examples of how SNES can be used.
      # Note that you will receive pointers to PETSc vectors within your user routines.
      #       That is important for parallel simulations, where global residual
      #       and solution vectors are passed to the user-routines, but we work with
      #       local vectors

      function fn!(cfx, cx, user_ctx)
        # We could do Global->Local here on cfx/cx, provided a pointer to the local
        #  vector is available in user_ctx
        x_in  = PETSc.unsafe_localarray(PetscScalar, cx;  write=false)   # read array
        fx_in = PETSc.unsafe_localarray(PetscScalar, cfx; read=false)    # write array

        fx_in[1] = x_in[1]^2       +  x_in[1]*x_in[2] - 3
        fx_in[2] = x_in[1]*x_in[2] +  x_in[2]^2 - 6

        Base.finalize(fx_in)
        Base.finalize(x_in)
      end

      function update_jac!(cx, J1, args...)
          x_in  = PETSc.unsafe_localarray(PetscScalar, cx;  write=false)
          J1[1,1] = 2x_in[1] + x_in[2]
          J1[1,2] = x_in[1]
          J1[2,1] = x_in[2]
          J1[2,2] = x_in[1] + 2x_in[2]

          Base.finalize(x_in)
          PETSc.assemble(J1)
      end

      # structure with which we can pass data to the user-routines above
      mutable struct Data
          vec
          julia
      end

      julia_vec = 0;  # we want pointers to local vectors

      # NOTE: for some reason, snes_monitor and ksp_monitor is not working if we specify that here.
      # To be sorted out why that is (converged_reason does work)
      S = PETSc.SNES{PetscScalar}(petsclib,comm;
                      ksp_rtol=1e-4,
                      pc_type="none",
                      ksp_monitor=true,
                      snes_monitor=true,
                      snes_converged_reason=false,
                      ksp_converged_reason=false)

      data        = Data(PetscScalar.([100;2]), 1)
      S.user_ctx  = data;      # we can pack anything we need in this struct

      PJ = PETSc.MatSeqDense(zeros(PetscScalar,(2,2)))
      PETSc.setfunction!(S, fn!, PETSc.VecSeq( PetscScalar.(zeros(2))))
      PETSc.setjacobian!(S, update_jac!, PJ, PJ)

      x  = PETSc.VecSeq(PetscScalar.([2.0, 3.0]));
      b  = PETSc.VecSeq(PetscScalar.([0.0, 0.0]));
      PETSc.solve!(x, S, b)

      sol = PETSc.unsafe_localarray(PetscScalar, x.ptr; read=false)
      @test sol ≈ [1.0,2.0] rtol=1e-4

      # cleanup
      PETSc.destroy(x);
      PETSc.destroy(b);
      PETSc.destroy(PJ);
      #PETSc.destroy(S);

      PETSc.finalize(petsclib)
    end
end


text:
module JoinTreeInference

export Node
export Potential
export parse_net
export create_moral_graph
export triangulate_graph
export Sepset
export create_junction_tree

using DataStructures

net_pattern = r"net[\s\n]*{(\n|.)*}"
node_pattern = r"node\s+([_\w-]+)[\n\s]*{([^}]*)}"
potential_pattern = r"potential\s*\(\s*([_\w-]+)\s*(\|\s*([_\w-]+\s*)+)?\)[\n\s]*{([^}]*)}"
statement_pattern = r"([_\w-]+)[]\s\n]*=[\n\s]*([^;]+);"
word_pattern = r"(\w+)"

type Node
    name::String
    states::Array{String}
end

type Potential
    node::String
    other_nodes::Array{String, 1}
    data::Array{Float64, 1}
end

type Sepset
    first::Int64
    second::Int64
    nodes::Set{String}
    mass::Int64
    cost::Int64
end

function parse_node(nodematch::RegexMatch)
    node_name = nodematch.captures[1]
    node_values = Array(String, 0)
    state_match = match(statement_pattern, nodematch[2])
    if state_match[1] == "states"
        for v in eachmatch(word_pattern, state_match[2])
            push!(node_values, v.match)
        end
    end
    return Node(node_name, node_values)
end

function parse_data_statement(value)
    value = replace(value, r"[\n\r]", "")
    value = replace(value, r"\)\s+\(", ")(")
    value = replace(value, r"[)(]", " ")
    data = Array(Float64, 0)
    output = ""
    for char in value
        if char == ' '
            if length(output) > 0
                push!(data, float(output))
            end
            output = ""
        else
            output = string(output, char)
        end
    end
    return data
end

function parse_potential(potentialmatch::RegexMatch)
    node = potentialmatch[1]
    other_nodes = potentialmatch[2]
    others = Array(String, 0)
    if other_nodes != nothing
        for n in eachmatch(r"([_\w-]+)", other_nodes)
            push!(others, n.match)
        end
    end
    body = potentialmatch[4]
    values = match(statement_pattern, body)
    if values[1] == "data"
        data = parse_data_statement(values[2])
    end
    return Potential(node, others, data)
end

function parse_net(net_fname)
    f = open(net_fname)
    net_str = readstring(f)
    close(f)

    net_str = lowercase(replace(net_str, r"%.*", ""))
    node_list = Array(Node, 0)
    for nodematch in eachmatch(node_pattern, net_str)
        push!(node_list, parse_node(nodematch))
    end

    potential_list = Array(Potential, 0)
    for potential_match in eachmatch(potential_pattern, net_str)
        push!(potential_list, parse_potential(potential_match))
    end
    return node_list, potential_list
end

function create_moral_graph(node_list::Array{Node,1},
                potential_list::Array{Potential,1})
    a_list = Dict{String, Set{String}}()
    for n in node_list
        a_list[n.name] = Set{String}()
    end

    function add_edge(i, j)
        push!(a_list[i], j)
        push!(a_list[j], i)
    end

    for p in potential_list
        for o in p.other_nodes
            add_edge(p.node, o)
        end
        for i in 1:length(p.other_nodes)
            for j in (i+1):length(p.other_nodes)
                add_edge(p.other_nodes[i], p.other_nodes[j])
            end
        end
    end
    return a_list
end

function get_next_node(g::Dict{String,Set{String}},
    node_weights::Dict{String,Int})

    best_f = 0
    best_w = 0
    best_k = 0
    function update(f, w, k)
        best_f = f
        best_w = w
        best_k = k
    end

    nodes =  collect(keys(g))
    first = true
    for (k, node) in enumerate(nodes)
        nei = collect(g[node])
        f = 0
        w = node_weights[node]
        for i in 1:length(nei)
            w += node_weights[nei[i]]
            for j in (i+1):length(nei)
                if !(nei[i] in g[nei[j]])
                    f += 1
                end
            end
        end
        if first
            first = false
            update(f, w, k)
        else
            if (f < best_f) || (f == best_f && w < best_w)
                update(f, w, k)
            end
        end
    end
    return nodes[best_k]
end

function triangulate_graph(g::Dict{String,Set{String}},
    node_list::Array{Node, 1})
    g1 = deepcopy(g)
    g2 = deepcopy(g)
    function add_edge(i, j)
        for g in (g1, g2)
            push!(g[i], j)
            push!(g[j], i)
        end
    end

    node_weights = Dict{String, Int}()
    for node in node_list
        node_weights[node.name] = length(node.states)
    end

    clusters = Array{Set{String}, 1}()
    function check_add(g::Dict{String, Set{String}}, node::String)
        cluster = copy(g[node])
        push!(cluster, node)
        subsumed = false
        for c in clusters
            if issubset(cluster, c)
                subsumed = true
                break
            end
        end
        if !subsumed
            push!(clusters, cluster)
        end
    end

    while !isempty(keys(g1))
        next_node = get_next_node(g1, node_weights)
        check_add(g1, next_node)
        nei = collect(g1[next_node])
        for i in 1:length(nei)
            for j in (i+1):length(nei)
                add_edge(nei[i], nei[j])
            end
            delete!(g1[nei[i]], next_node)
        end
        delete!(g1, next_node)
    end
    return g2, clusters
end

function triangulate_graph(node_list::Array{Node,1},
                potential_list::Array{Potential,1})
    mg = create_moral_graph(node_list, potential_list)
    return triangulate_graph(mg, node_list)
end

function create_sepsets(clusters, node_list)
    node_weights = Dict{String, Int64}()
    for node in node_list
        node_weights[node.name] = length(node.states)
    end
    n = length(clusters)
    weights = Array{Int64, 1}()
    for i in 1:n
        w = 1
        for v in clusters[i]
            w *= node_weights[v]
        end
        push!(weights, w)
    end

    sepsets = Array{Sepset, 1}()
    for i in 1:n
        for j in (i+1):n
            # creat a new sepset
            nodes = intersect(clusters[i], clusters[j])
            mass = length(nodes)
            cost = weights[i] + weights[j]
            push!(sepsets, Sepset(i, j, nodes, mass, cost))
        end
    end

    sepset_comp(x, y) = (x.mass > y.mass) || ((x.mass == y.mass) && (x.cost < y.cost))
    return sort(sepsets, lt = sepset_comp)
end

function create_junction_tree(clusters::Array{Set{String}, 1}, sepsets::Array{Sepset, 1})
    n = length(clusters)
    output_tree = Dict{Int, Set{Int}}()
    for i in 1:n
        output_tree[i] = Set{Int}()
    end

    tree = IntDisjointSets(n)
    num_edges = 0
    for sepset in sepsets
        if num_edges == n - 1
            break
        end
        if ! in_same_set(tree, sepset.first, sepset.second)
            union!(tree, sepset.first, sepset.second)
            push!(output_tree[sepset.first], sepset.second)
            push!(output_tree[sepset.second], sepset.first)
            num_edges += 1
        end
    end
    return output_tree
end

function create_junction_tree(clusters::Array{Set{String}, 1}, node_list::Array{Node, 1})
    sepsets = create_sepsets(clusters, node_list)
    return create_junction_tree(clusters, sepsets)
end

end

max_stars_repo_name: bhrzslm/uncertainty-reasoning

max_issues_repo_name: bhrzslm/uncertainty-reasoning

max_forks_repo_name: bhrzslm/uncertainty-reasoning


text:
"""
    Constructing all Bayesian Weights from the set of parameters and information
    precision.

    Longer than it should be for here, but does the job!
"""
function CstrBPar(βᵢⱼ, βᵢ, β, par::Pars)::BPars
    @unpack_Pars par
    ## Economy
    σ²      = 1/(σ²ₐ^-1 + β + β*σ²ᵤ^-1)

    # sliced by signal
    w_z      = β*σ²ᵤ^-1*σ²
    w_zε     = √β*σ²ᵤ^-1*σ²
    w_x      = β*σ²
    w_xε     = √β*σ²

    # sliced by fund/noise
    w_p      = σ²ₐ^-1*σ²
    w_a      = w_z + w_x
    w_ε      = w_zε + w_xε

    ## Market
    σ²ᵢ      = 1/(σ²ₐ^-1 + βᵢ + βᵢ*σ²ᵤ^-1)

    # sliced by signal
    wᵢ_z      = βᵢ*σ²ᵤ^-1*σ²ᵢ
    wᵢ_zε     = √βᵢ*σ²ᵤ^-1*σ²ᵢ
    wᵢ_x      = βᵢ*σ²ᵢ
    wᵢ_xε     = √βᵢ*σ²ᵢ

    # sliced by fund/noise
    wᵢ_p      = σ²ₐ^-1*σ²ᵢ
    wᵢ_a      = wᵢ_z + wᵢ_x
    wᵢ_ε      = wᵢ_zε + wᵢ_xε

    ## Objective, no private signal!
    σ²ᵢ_obj   = 1/(σ²ₐ^-1 + βᵢ*σ²ᵤ^-1)

    # sliced by signal
    wᵢ_obj_z  = βᵢ*σ²ᵤ^-1*σ²ᵢ_obj
    wᵢ_obj_zε = √βᵢ*σ²ᵤ^-1*σ²ᵢ_obj

    # sliced by fund/noise
    wᵢ_obj_p  = σ²ₐ^-1*σ²ᵢ_obj
    wᵢ_obj_a  = wᵢ_obj_z
    wᵢ_obj_ε  = wᵢ_obj_zε

    # Individual
    σ²ᵢⱼ      = 1/(σ²ₐ^-1 + βᵢⱼ + βᵢ*σ²ᵤ^-1)
    # sliced by signal
    wᵢⱼ_z     = βᵢ*σ²ᵤ^-1*σ²ᵢⱼ
    wᵢⱼ_zε    = √βᵢ*σ²ᵤ^-1*σ²ᵢⱼ
    wᵢⱼ_x     = βᵢⱼ*σ²ᵢⱼ
    wᵢⱼ_xε    = √βᵢⱼ*σ²ᵢⱼ

    # sliced by fund/noise
    wᵢⱼ_p     = σ²ₐ^-1*σ²ᵢⱼ
    wᵢⱼ_a     = wᵢⱼ_z  + wᵢⱼ_x
    wᵢⱼ_ε     = wᵢⱼ_zε + wᵢⱼ_xε

    Ωᵢⱼ_x     = wᵢ_x/wᵢⱼ_x
    Ωᵢⱼ_z     = (wᵢ_z - wᵢⱼ_z)/wᵢⱼ_x
    Ωᵢⱼ_p     = (wᵢ_p - wᵢⱼ_p)/wᵢⱼ_x
    Ωᵢⱼ_ε     = (wᵢ_z - wᵢⱼ_z)/wᵢⱼ_xε

    # Perturbed Variances
    σ²ₐi_tilde   = (1/σ²ₐ + βᵢⱼ*(1 - Ωᵢⱼ_z)^2/(1 + σ²ᵤ*Ωᵢⱼ_z^2*βᵢⱼ/β))^-1
    σ²_i_ε_tilde = (1/σ²ᵤ + (Ωᵢⱼ_z^2*βᵢⱼ/β)/(1 + βᵢⱼ*σ²ₐ*(1 - Ωᵢⱼ_z)^2))^-1

    BPars(βᵢ       = βᵢ,
          βᵢⱼ      = βᵢⱼ,
          β        = β,

          ## Economy
          σ²      = σ²,
          w_z     = w_z,
          w_zε    = w_zε,
          w_x     = w_x,
          w_xε    = w_xε,
          w_p     = w_p,
          w_a     = w_a,
          w_ε     = w_ε,

          ### Objective
          #σ²_obj   = σ²_obj,
          #w_obj_z  = w_obj_z,
          #w_obj_zε = w_obj_zε,

          ## Market
          σ²ᵢ      = σ²ᵢ,
          wᵢ_z     = wᵢ_z,
          wᵢ_zε    = wᵢ_zε,
          wᵢ_x     = wᵢ_x,
          wᵢ_xε    = wᵢ_xε,
          wᵢ_p     = wᵢ_p,
          wᵢ_a     = wᵢ_a,
          wᵢ_ε     = wᵢ_ε,

          ## Objective
          σ²ᵢ_obj   = σ²ᵢ_obj,
          wᵢ_obj_z  = wᵢ_obj_z,
          wᵢ_obj_zε = wᵢ_obj_zε,

          wᵢ_obj_p  = wᵢ_obj_p,
          wᵢ_obj_a  = wᵢ_obj_a,
          wᵢ_obj_ε  = wᵢ_obj_ε,

          # Individual
          σ²ᵢⱼ      = σ²ᵢⱼ,
          wᵢⱼ_z     = wᵢⱼ_z,
          wᵢⱼ_zε    = wᵢⱼ_zε,
          wᵢⱼ_x     = wᵢⱼ_x,
          wᵢⱼ_xε    = wᵢⱼ_xε,
          wᵢⱼ_p     = wᵢⱼ_p,
          wᵢⱼ_a     = wᵢⱼ_a,
          wᵢⱼ_ε     = wᵢⱼ_ε,

          Ωᵢⱼ_x     = Ωᵢⱼ_x,
          Ωᵢⱼ_z     = Ωᵢⱼ_z,
          Ωᵢⱼ_p     = Ωᵢⱼ_p,
          Ωᵢⱼ_ε     = Ωᵢⱼ_ε,

          σ²ₐi_tilde   = σ²ₐi_tilde,
          σ²_i_ε_tilde = σ²_i_ε_tilde)
end

# Helper function
CstrBPar(β, par::Pars) = CstrBPar(β, β, β, par)

"""
    Information Acquisition Costs
"""
function infocost(βᵢⱼ,par)
    par.ca*βᵢⱼ^par.cb
end
"""
    Marginal Cost of Information Acquisition
"""
function mc(βᵢⱼ,par)
    par.ca*par.cb*βᵢⱼ^(par.cb-1)
end

# Helper function
mc(βᵢⱼ; par::Pars) = mc(βᵢⱼ,par)

"""
    I only cover the case with an log-linear payoff function. But other cases are also possible!
"""
payoff(a) = exp(a)

"""
    The price function.
"""
function price(z, par, bpar)
    @unpack_Pars par
    @unpack_BPars bpar
    exp(wᵢ_p*aₜ + wᵢ_x*z + wᵢ_z*(z - εₜ/√β) + 1/2*σ²ᵢ)
end

"""
    findbeta(par::Pars; βₑ = 1.)

Finds the symmetric choice of β for which mb(β,β) = mc(β).

# Arguments:
- `par::Pars`: set of parameters
- `βₑ`: keyword argument for market beta for partial equilibrium analysis
"""
function findbeta(par::Pars; βₑ = 1.)

    f = β -> mb(β^2, par) - mc(β^2, par) # Small trick to keep beta positive, but use unbounded rootfinding.
    find_zero(f, 1., atol = 1e-8)^2
end

"""
    Compute the marginal benefit of information acquisition with Gaussian
    Quadrature.
"""
function mb(βᵢⱼ, βᵢ, β, par::Pars; bpar::BPars = CstrBPar(βᵢⱼ, βᵢ, β, par))
    n              = 5
    points, weight = qnwnorm([n,n],[par.aₜ, par.εₜ],[par.σ²ₐ, par.σ²ᵤ])
    a              = points[:,1]
    ϵ              = points[:,2]
    z              = a + ϵ/√βᵢ
    probdiff       = TradeProbDiff.(a, z, par = par, bpar = bpar)
    prof           = profit.(a, z, par = par, bpar = bpar)
    bounds         = 2

    return sum(bounds.*probdiff.*weight.*prof)
end

mb(β, par::Pars) = mb(β, β, β, par)

"""
    Substract the price from the payoff
"""
profit(θ, z; par::Pars, bpar::BPars) = payoff(θ) - price(z, par, bpar)

"""
    Derivative of probability of trading with respect to βᵢ.

    This part depends on the payoff function. Here I provide the analytical expression,
    but in principle it could also be computed numerically. Details can be found in my JMP.
"""
function TradeProbDiff(a, z; par::Pars, bpar::BPars)
    term      = TradeProbTermGrowth(a, z, par, bpar)
    term_diff = TradeProbTermDiffGrowth(a, z, par, bpar)
    term_diff*pdf(Normal(), term)
end

function TradeProbTermGrowth(a, z, par::Pars, bpar::BPars)
    @unpack_Pars par
    @unpack_BPars bpar
    -√βᵢⱼ*(Ωᵢⱼ_p*aₜ + Ωᵢⱼ_x*z + Ωᵢⱼ_z*(z - εₜ/√βᵢ) + 1/(2*wᵢⱼ_x)*(σ²ᵢ - σ²ᵢⱼ) - a)
end

function TradeProbTermDiffGrowth(a, z, par::Pars, bpar::BPars)
    @unpack_Pars par
    @unpack_BPars bpar
    postᵢ          = wᵢ_p*aₜ + wᵢ_x*z + wᵢ_z*(z - εₜ/√βᵢ) + 1/2*σ²ᵢ
    wi_wijx_diff   = -(1/√βᵢⱼ - 1/(2*βᵢⱼ^(3/2))*σ²ᵢⱼ^-1)*postᵢ
    wij_wijx_diff  = -(σ²ₐ^-1*aₜ + βᵢ*σ²ᵤ^-1*(z - εₜ/√βᵢ) + 1/2)/(2*βᵢⱼ^(3/2))

    return term_diff      = 1/(2*√βᵢⱼ)*a + wi_wijx_diff + wij_wijx_diff
end


text:
struct IPv4HeaderRaw
    version_ihl::UInt8
    dscp_ecn::UInt8
    total_length::UInt16
    id::UInt16
    flags_fragmentoffset::UInt16
    ttl::UInt8
    protocol::UInt8
    header_checksum::UInt16
    src_ip::UInt32
    dst_ip::UInt32
end

const IPPROTOCOL_IGMP = UInt8(0x02)
const IPPROTOCOL_TCP = UInt8(0x06)
const IPPROTOCOL_UDP = UInt8(0x11)

struct IPv4Header
    header_length::UInt8
    dscp::UInt8
    total_length::UInt16
    id::UInt16
    flags::UInt8
    fragment_offset::UInt16
    ttl::UInt8
    protocol::UInt8
    src_ip::IPv4
    dst_ip::IPv4
end

struct IPv4Packet
    header::IPv4Header
    payload::UnsafeArray{UInt8, 1}
end

function decode_ipv4(data::DenseVector{UInt8})
    p = Base.unsafe_convert(Ptr{UInt8}, data)
    GC.@preserve data begin
        rh = unsafe_load(convert(Ptr{IPv4HeaderRaw}, p))
    end
    h = IPv4Header(
        (rh.version_ihl & 0x0f) * 4,
        rh.dscp_ecn >> 2,
        ntoh(rh.total_length),
        ntoh(rh.id),
        ntoh(rh.flags_fragmentoffset) >> 13,
        (ntoh(rh.flags_fragmentoffset) & 0x1fff) * 8,
        rh.ttl,
        rh.protocol,
        IPv4(ntoh(rh.src_ip)),
        IPv4(ntoh(rh.dst_ip)))
    IPv4Packet(
        h,
        UnsafeArray{UInt8, 1}(
            p + h.header_length,
            (Int(h.total_length - h.header_length),)))
end

max_stars_repo_name: ancapdev/NetworkProtocols

max_issues_repo_name: ancapdev/NetworkProtocols


text:
# This file was generated by the Julia Swagger Code Generator
# Do not modify this file directly. Modify the swagger specification instead.

mutable struct DiskEncryptionSettings <: SwaggerModel
    diskEncryptionKey::Any # spec type: Union{ Nothing, KeyVaultSecretReference } # spec name: diskEncryptionKey
    keyEncryptionKey::Any # spec type: Union{ Nothing, KeyVaultKeyReference } # spec name: keyEncryptionKey
    enabled::Any # spec type: Union{ Nothing, Bool } # spec name: enabled

    function DiskEncryptionSettings(;diskEncryptionKey=nothing, keyEncryptionKey=nothing, enabled=nothing)
        o = new()
        validate_property(DiskEncryptionSettings, Symbol("diskEncryptionKey"), diskEncryptionKey)
        setfield!(o, Symbol("diskEncryptionKey"), diskEncryptionKey)
        validate_property(DiskEncryptionSettings, Symbol("keyEncryptionKey"), keyEncryptionKey)
        setfield!(o, Symbol("keyEncryptionKey"), keyEncryptionKey)
        validate_property(DiskEncryptionSettings, Symbol("enabled"), enabled)
        setfield!(o, Symbol("enabled"), enabled)
        o
    end
end # type DiskEncryptionSettings

const _property_map_DiskEncryptionSettings = Dict{Symbol,Symbol}(Symbol("diskEncryptionKey")=>Symbol("diskEncryptionKey"), Symbol("keyEncryptionKey")=>Symbol("keyEncryptionKey"), Symbol("enabled")=>Symbol("enabled"))
const _property_types_DiskEncryptionSettings = Dict{Symbol,String}(Symbol("diskEncryptionKey")=>"KeyVaultSecretReference", Symbol("keyEncryptionKey")=>"KeyVaultKeyReference", Symbol("enabled")=>"Bool")
Base.propertynames(::Type{ DiskEncryptionSettings }) = collect(keys(_property_map_DiskEncryptionSettings))
Swagger.property_type(::Type{ DiskEncryptionSettings }, name::Symbol) = Union{Nothing,eval(Base.Meta.parse(_property_types_DiskEncryptionSettings[name]))}
Swagger.field_name(::Type{ DiskEncryptionSettings }, property_name::Symbol) =  _property_map_DiskEncryptionSettings[property_name]

function check_required(o::DiskEncryptionSettings)
    true
end

function validate_property(::Type{ DiskEncryptionSettings }, name::Symbol, val)
end


text:
include("dependencies_for_runtests.jl")

using Oceananigans.Models.HydrostaticFreeSurfaceModels: VectorInvariant, PrescribedVelocityFields, PrescribedField
using Oceananigans.Models.HydrostaticFreeSurfaceModels: ExplicitFreeSurface, ImplicitFreeSurface
using Oceananigans.Coriolis: VectorInvariantEnergyConserving, VectorInvariantEnstrophyConserving
using Oceananigans.TurbulenceClosures: VerticallyImplicitTimeDiscretization, ExplicitTimeDiscretization
using Oceananigans.TurbulenceClosures: CATKEVerticalDiffusivity, HorizontallyCurvilinearAnisotropicBiharmonicDiffusivity

function time_step_hydrostatic_model_works(grid;
                                           coriolis = nothing,
                                           free_surface = ExplicitFreeSurface(),
                                           momentum_advection = nothing,
                                           closure = nothing,
                                           velocities = nothing)

    tracers = [:T, :S]
    closure isa CATKEVerticalDiffusivity && push!(tracers, :e)

    model = HydrostaticFreeSurfaceModel(grid = grid,
                                        momentum_advection = momentum_advection,
                                        free_surface = free_surface,
                                        coriolis = coriolis,
                                        tracers = tracers,
                                        velocities = velocities,
                                        closure = closure)

    simulation = Simulation(model, Δt=1.0, stop_iteration=1)

    run!(simulation)

    return model.clock.iteration == 1
end

function hydrostatic_free_surface_model_tracers_and_forcings_work(arch)
    grid = RectilinearGrid(arch, size=(1, 1, 1), extent=(2π, 2π, 2π))
    model = HydrostaticFreeSurfaceModel(grid=grid, tracers=(:T, :S, :c, :d))

    @test model.tracers.T isa Field
    @test model.tracers.S isa Field
    @test model.tracers.c isa Field
    @test model.tracers.d isa Field

    @test haskey(model.forcing, :u)
    @test haskey(model.forcing, :v)
    @test haskey(model.forcing, :η)
    @test haskey(model.forcing, :T)
    @test haskey(model.forcing, :S)
    @test haskey(model.forcing, :c)
    @test haskey(model.forcing, :d)

    simulation = Simulation(model, Δt=1.0, stop_iteration=1)
    run!(simulation)

    @test model.clock.iteration == 1

    return nothing
end

topo_1d = (Flat, Flat, Bounded)

topos_2d = ((Periodic, Flat, Bounded),
            (Flat, Bounded,  Bounded),
            (Bounded, Flat,  Bounded))

topos_3d = ((Periodic, Periodic, Bounded),
            (Periodic, Bounded,  Bounded),
            (Bounded,  Bounded,  Bounded))

@testset "Hydrostatic free surface Models" begin
    @info "Testing hydrostatic free surface models..."

    @testset "$topo_1d model construction" begin
        @info "  Testing $topo_1d model construction..."
        for arch in archs, FT in [Float64] #float_types
            grid = RectilinearGrid(arch, FT, topology=topo_1d, size=(1), extent=(1))
            model = HydrostaticFreeSurfaceModel(grid=grid)
            @test model isa HydrostaticFreeSurfaceModel
        end
    end

    for topo in topos_2d
        @testset "$topo model construction" begin
            @info "  Testing $topo model construction..."
            for arch in archs, FT in float_types
                grid = RectilinearGrid(arch, FT, topology=topo, size=(1, 1), extent=(1, 2))
                model = HydrostaticFreeSurfaceModel(grid=grid)
                @test model isa HydrostaticFreeSurfaceModel
            end
        end
    end

    for topo in topos_3d
        @testset "$topo model construction" begin
            @info "  Testing $topo model construction..."
            for arch in archs, FT in float_types
                grid = RectilinearGrid(arch, FT, topology=topo, size=(1, 1, 1), extent=(1, 2, 3))
                model = HydrostaticFreeSurfaceModel(grid=grid)
                @test model isa HydrostaticFreeSurfaceModel
            end
        end
    end

    @testset "Halo size check in model constructor" begin
        for topo in topos_3d
            grid = RectilinearGrid(topology=topo, size=(1, 1, 1), extent=(1, 2, 3))
            hcabd_closure = HorizontallyCurvilinearAnisotropicBiharmonicDiffusivity()

            @test_throws ArgumentError HydrostaticFreeSurfaceModel(grid=grid, tracer_advection=CenteredFourthOrder())
            @test_throws ArgumentError HydrostaticFreeSurfaceModel(grid=grid, tracer_advection=UpwindBiasedThirdOrder())
            @test_throws ArgumentError HydrostaticFreeSurfaceModel(grid=grid, tracer_advection=UpwindBiasedFifthOrder())
            @test_throws ArgumentError HydrostaticFreeSurfaceModel(grid=grid, momentum_advection=UpwindBiasedFifthOrder())
            @test_throws ArgumentError HydrostaticFreeSurfaceModel(grid=grid, closure=hcabd_closure)

            # Big enough
            bigger_grid = RectilinearGrid(topology=topo, size=(1, 1, 1), extent=(1, 2, 3), halo=(3, 3, 3))

            model = HydrostaticFreeSurfaceModel(grid=bigger_grid, closure=hcabd_closure)
            @test model isa HydrostaticFreeSurfaceModel

            model = HydrostaticFreeSurfaceModel(grid=bigger_grid, momentum_advection=UpwindBiasedFifthOrder())
            @test model isa HydrostaticFreeSurfaceModel

            model = HydrostaticFreeSurfaceModel(grid=bigger_grid, closure=hcabd_closure)
            @test model isa HydrostaticFreeSurfaceModel

            model = HydrostaticFreeSurfaceModel(grid=bigger_grid, tracer_advection=UpwindBiasedFifthOrder())
            @test model isa HydrostaticFreeSurfaceModel
        end
    end

    @testset "Setting HydrostaticFreeSurfaceModel fields" begin
        @info "  Testing setting hydrostatic free surface model fields..."
        for arch in archs, FT in float_types
            N = (4, 4, 1)
            L = (2π, 3π, 5π)

            grid = RectilinearGrid(arch, FT, size=N, extent=L)
            model = HydrostaticFreeSurfaceModel(grid=grid)

            x, y, z = nodes((Face, Center, Center), model.grid, reshape=true)

            u₀(x, y, z) = x * y^2
            u_answer = @. x * y^2

            η₀ = rand(size(grid)...)
            η_answer = deepcopy(η₀)

            set!(model, u=u₀, η=η₀)

            u, v, w = model.velocities
            η = model.free_surface.η

            @test all(Array(interior(u)) .≈ u_answer)
            @test all(Array(interior(η)) .≈ η_answer)
        end
    end

    for arch in archs
        for topo in topos_3d
            grid = RectilinearGrid(arch, size=(1, 1, 1), extent=(1, 1, 1), topology=topo)

            @testset "Time-stepping Rectilinear HydrostaticFreeSurfaceModels [$arch, $topo]" begin
                @info "  Testing time-stepping Rectilinear HydrostaticFreeSurfaceModels [$arch, $topo]..."
                @test time_step_hydrostatic_model_works(grid)
            end
        end

        z_face_generator(; Nz=1, p=1, H=1) = k -> -H + (k / (Nz+1))^p # returns a generating function

        rectilinear_grid = RectilinearGrid(arch, size=(1, 1, 1), extent=(1, 1, 1), halo=(3, 3, 3))

        lat_lon_sector_grid = LatitudeLongitudeGrid(arch, size=(1, 1, 1), longitude=(0, 60), latitude=(15, 75), z=(-1, 0), precompute_metrics=true)
        lat_lon_strip_grid  = LatitudeLongitudeGrid(arch, size=(1, 1, 1), longitude=(-180, 180), latitude=(15, 75), z=(-1, 0), precompute_metrics=true)

        vertically_stretched_grid = RectilinearGrid(arch, size=(1, 1, 1), x=(0, 1), y=(0, 1), z=z_face_generator(), halo=(3, 3, 3))
        lat_lon_sector_grid_stretched = LatitudeLongitudeGrid(arch, size=(1, 1, 1), longitude=(0, 60), latitude=(15, 75), z=z_face_generator(), precompute_metrics=true)
        lat_lon_strip_grid_stretched  = LatitudeLongitudeGrid(arch, size=(1, 1, 1), longitude=(-180, 180), latitude=(15, 75), z=z_face_generator(), precompute_metrics=true)

        grids = (rectilinear_grid, lat_lon_sector_grid, lat_lon_strip_grid, lat_lon_sector_grid_stretched, lat_lon_strip_grid_stretched, vertically_stretched_grid)
        free_surfaces = (ExplicitFreeSurface(), ImplicitFreeSurface(), ImplicitFreeSurface(solver_method=:MatrixIterativeSolver))

        for grid in grids
            for free_surface in free_surfaces
                topo = topology(grid)
                grid_type = typeof(grid).name.wrapper
                free_surface_type = typeof(free_surface).name.wrapper
                test_label = "[$arch, $grid_type, $topo, $free_surface_type]"
                @testset "Time-stepping HydrostaticFreeSurfaceModels with different grids $test_label" begin
                    @info "  Testing time-stepping HydrostaticFreeSurfaceModels with different grids $test_label..."
                    @test time_step_hydrostatic_model_works(grid, free_surface=free_surface)
                end
            end
        end

        for coriolis in (nothing, FPlane(f=1), BetaPlane(f₀=1, β=0.1))
            @testset "Time-stepping HydrostaticFreeSurfaceModels [$arch, $(typeof(coriolis))]" begin
                @info "  Testing time-stepping HydrostaticFreeSurfaceModels [$arch, $(typeof(coriolis))]..."
                @test time_step_hydrostatic_model_works(rectilinear_grid, coriolis=coriolis)
            end
        end

        for coriolis in (nothing,
                         HydrostaticSphericalCoriolis(scheme=VectorInvariantEnergyConserving()),
                         HydrostaticSphericalCoriolis(scheme=VectorInvariantEnstrophyConserving()))

            @testset "Time-stepping HydrostaticFreeSurfaceModels [$arch, $(typeof(coriolis))]" begin
                @test time_step_hydrostatic_model_works(lat_lon_sector_grid, coriolis=coriolis)
                @test time_step_hydrostatic_model_works(lat_lon_strip_grid, coriolis=coriolis)
            end
        end

        for momentum_advection in (VectorInvariant(), CenteredSecondOrder(), WENO5())
            @testset "Time-stepping HydrostaticFreeSurfaceModels [$arch, $(typeof(momentum_advection))]" begin
                @info "  Testing time-stepping HydrostaticFreeSurfaceModels [$arch, $(typeof(momentum_advection))]..."
                @test time_step_hydrostatic_model_works(rectilinear_grid, momentum_advection=momentum_advection)
            end
        end

        momentum_advection = VectorInvariant()
        @testset "Time-stepping HydrostaticFreeSurfaceModels [$arch, $(typeof(momentum_advection))]" begin
            @info "  Testing time-stepping HydrostaticFreeSurfaceModels [$arch, $(typeof(momentum_advection))]..."
            @test time_step_hydrostatic_model_works(lat_lon_sector_grid, momentum_advection=momentum_advection)
        end

        for closure in (IsotropicDiffusivity(),
                        HorizontallyCurvilinearAnisotropicDiffusivity(),
                        HorizontallyCurvilinearAnisotropicDiffusivity(time_discretization=VerticallyImplicitTimeDiscretization()),
                        CATKEVerticalDiffusivity(),
                        CATKEVerticalDiffusivity(time_discretization=ExplicitTimeDiscretization()))

            @testset "Time-stepping Curvilinear HydrostaticFreeSurfaceModels [$arch, $(typeof(closure).name.wrapper)]" begin
                @info "  Testing time-stepping Curvilinear HydrostaticFreeSurfaceModels [$arch, $(typeof(closure).name.wrapper)]..."
                @test_skip time_step_hydrostatic_model_works(arch, vertically_stretched_grid, closure=closure)
                @test time_step_hydrostatic_model_works(lat_lon_sector_grid, closure=closure)
                @test time_step_hydrostatic_model_works(lat_lon_strip_grid, closure=closure)
            end
        end

        closure = IsotropicDiffusivity()
        @testset "Time-stepping Rectilinear HydrostaticFreeSurfaceModels [$arch, $(typeof(closure).name.wrapper)]" begin
            @info "  Testing time-stepping Rectilinear HydrostaticFreeSurfaceModels [$arch, $(typeof(closure).name.wrapper)]..."
            @test time_step_hydrostatic_model_works(rectilinear_grid, closure=closure)
        end

        @testset "PrescribedVelocityFields [$arch]" begin
            @info "  Testing PrescribedVelocityFields [$arch]..."

            grid = RectilinearGrid(arch, size=1, x = (0, 1), topology = (Periodic, Flat, Flat))

            u₀, v₀ = 0.1, 0.2

            U = Field{Face, Center, Center}(grid)
            V = Field{Center, Face, Center}(grid)

            CUDA.@allowscalar begin
                parent(U)[2, 1, 1] = u₀
                parent(V)[2, 1, 1] = v₀
            end

            u = PrescribedField(Face, Center, Center, U, grid)
            v = PrescribedField(Face, Center, Center, V, grid)

            CUDA.@allowscalar begin
                @test parent(u)[1, 1, 1] == u₀
                @test parent(u)[3, 1, 1] == u₀
                @test parent(v)[1, 1, 1] == v₀
                @test parent(v)[3, 1, 1] == v₀
            end
        end

        @testset "Time-stepping HydrostaticFreeSurfaceModels with PrescribedVelocityFields [$arch]" begin
            @info "  Testing time-stepping HydrostaticFreeSurfaceModels with PrescribedVelocityFields [$arch]..."

            # Non-parameterized functions
            u(x, y, z, t) = 1
            v(x, y, z, t) = exp(z)
            w(x, y, z, t) = sin(z)
            velocities = PrescribedVelocityFields(u=u, v=v, w=w)

            @test time_step_hydrostatic_model_works(rectilinear_grid, momentum_advection  = nothing, velocities = velocities)
            @test time_step_hydrostatic_model_works(lat_lon_sector_grid, momentum_advection = nothing, velocities = velocities)

            parameters = (U=1, m=0.1, W=0.001)
            u(x, y, z, t, p) = p.U
            v(x, y, z, t, p) = exp(p.m * z)
            w(x, y, z, t, p) = p.W * sin(z)

            velocities = PrescribedVelocityFields(u=u, v=v, w=w, parameters=parameters)

            @test time_step_hydrostatic_model_works(rectilinear_grid, momentum_advection  = nothing, velocities = velocities)
            @test time_step_hydrostatic_model_works(lat_lon_sector_grid, momentum_advection = nothing, velocities = velocities)
        end

        @testset "HydrostaticFreeSurfaceModel with tracers and forcings [$arch]" begin
            @info "  Testing HydrostaticFreeSurfaceModel with tracers and forcings [$arch]..."
            hydrostatic_free_surface_model_tracers_and_forcings_work(arch)
        end
    end
end