text:
"""
    One-moment bulk microphysics scheme, which includes:

  - terminal velocity of precipitation
  - condensation and evaporation of cloud liquid water and
    deposition and sublimation of cloud ice (relaxation to equilibrium)
  - autoconversion of cloud liquid water into rain and of cloud ice into snow
  - accretion due to collisions between categories of condensed species
  - evaporation and sublimation of hydrometeors
  - melting of snow into rain
"""
module Microphysics

using SpecialFunctions

using ClimateMachine.Thermodynamics

using CLIMAParameters
using CLIMAParameters.Planet: ρ_cloud_liq, R_v, grav, T_freeze
using CLIMAParameters.Atmos.Microphysics

const APS = AbstractParameterSet
const ACloudPS = AbstractCloudParameterSet
const APrecipPS = AbstractPrecipParameterSet
const ALPS = AbstractLiquidParameterSet
const AIPS = AbstractIceParameterSet
const ARPS = AbstractRainParameterSet
const ASPS = AbstractSnowParameterSet

export τ_relax

export terminal_velocity

export conv_q_vap_to_q_liq_ice

export conv_q_liq_to_q_rai
export conv_q_ice_to_q_sno

export accretion
export accretion_rain_sink
export accretion_snow_rain

export evaporation_sublimation
export snow_melt

"""
    v0_rai(param_set, ρ)

 - `param_set` - abstract set with earth parameters
 - `ρ` air density

Returns the proportionality coefficient in terminal velocity(r/r0).
"""
function v0_rai(param_set::APS, rain_param_set::ARPS, ρ::FT) where {FT <: Real}

    _ρ_cloud_liq::FT = ρ_cloud_liq(param_set)
    _C_drag::FT = Microphysics.C_drag(param_set)
    _grav::FT = grav(param_set)
    _r0::FT = r0(rain_param_set)

    return sqrt(FT(8 / 3) / _C_drag * (_ρ_cloud_liq / ρ - FT(1)) * _grav * _r0)
end

"""
    n0_sno(snow_param_set, q_sno, ρ)

 - `snow_param_set` - abstract set with snow parameters
 - `q_sno` -  snow specific humidity
 - `ρ` - air density

Returns the intercept parameter of the assumed Marshall-Palmer distribution of
snow particles.
"""
function n0_sno(snow_param_set::ASPS, q_sno::FT, ρ::FT) where {FT <: Real}

    _ν_sno::FT = ν_sno(snow_param_set)
    _μ_sno::FT = μ_sno(snow_param_set)

    return _μ_sno * (ρ * q_sno)^_ν_sno
end

"""
    unpack_params(param_set, microphysics_param_set, ρ, q_)

 - `param_set` - abstract set with earth parameters
 - `microphysics_param_set` - abstract set with microphysics parameters
 - `q_` - specific humidity
 - `ρ` - air density

Utility function that unpacks microphysics parameters.
"""
function unpack_params(
    param_set::APS,
    ice_param_set::AIPS,
    ρ::FT,
    q_ice::FT,
) where {FT <: Real}
    #TODO - make ρ and q_ice optional
    _n0::FT = n0(ice_param_set)
    _r0::FT = r0(ice_param_set)

    _m0::FT = m0(param_set, ice_param_set)
    _me::FT = me(ice_param_set)

    _χm::FT = χm(ice_param_set)
    _Δm::FT = Δm(ice_param_set)

    return (_n0, _r0, _m0, _me, _χm, _Δm)
end
function unpack_params(
    param_set::APS,
    rain_param_set::ARPS,
    ρ::FT,
    q_rai::FT,
) where {FT <: Real}
    #TODO - make q_rai optional
    _n0::FT = n0(rain_param_set)
    _r0::FT = r0(rain_param_set)

    _m0::FT = m0(param_set, rain_param_set)
    _me::FT = me(rain_param_set)
    _a0::FT = a0(rain_param_set)
    _ae::FT = ae(rain_param_set)
    _v0::FT = v0_rai(param_set, rain_param_set, ρ)
    _ve::FT = ve(rain_param_set)

    _χm::FT = χm(rain_param_set)
    _Δm::FT = Δm(rain_param_set)
    _χa::FT = χa(rain_param_set)
    _Δa::FT = Δa(rain_param_set)
    _χv::FT = χv(rain_param_set)
    _Δv::FT = Δv(rain_param_set)

    return (
        _n0,
        _r0,
        _m0,
        _me,
        _χm,
        _Δm,
        _a0,
        _ae,
        _χa,
        _Δa,
        _v0,
        _ve,
        _χv,
        _Δv,
    )
end
function unpack_params(
    param_set::APS,
    snow_param_set::ASPS,
    ρ::FT,
    q_sno::FT,
) where {FT <: Real}

    _n0::FT = n0_sno(snow_param_set, q_sno, ρ)
    _r0::FT = r0(snow_param_set)

    _m0::FT = m0(snow_param_set)
    _me::FT = me(snow_param_set)
    _a0::FT = a0(snow_param_set)
    _ae::FT = ae(snow_param_set)
    _v0::FT = v0(snow_param_set)
    _ve::FT = ve(snow_param_set)

    _χm::FT = χm(snow_param_set)
    _Δm::FT = Δm(snow_param_set)
    _χa::FT = χa(snow_param_set)
    _Δa::FT = Δa(snow_param_set)
    _χv::FT = χv(snow_param_set)
    _Δv::FT = Δv(snow_param_set)

    return (
        _n0,
        _r0,
        _m0,
        _me,
        _χm,
        _Δm,
        _a0,
        _ae,
        _χa,
        _Δa,
        _v0,
        _ve,
        _χv,
        _Δv,
    )
end

"""
    lambda(q, ρ, n0, m0, me, r0, χm, Δm)

 - `q` - specific humidity of rain, ice or snow
 - `ρ` - air density
 - `n0` - size distribution parameter
 - `m0`, `me`, `χm`, `Δm`, `r0` - mass(radius) parameters

Returns the rate parameter of the assumed size distribution of
particles (rain drops, ice crystals, snow crystals).
"""
function lambda(
    q::FT,
    ρ::FT,
    n0::FT,
    m0::FT,
    me::FT,
    r0::FT,
    χm::FT,
    Δm::FT,
) where {FT <: Real}

    λ::FT = FT(0)

    if q > FT(0)
        λ =
            (
                χm * m0 * n0 * gamma(me + Δm + FT(1)) / ρ / q / r0^(me + Δm)
            )^FT(1 / (me + Δm + 1))
    end
    return λ
end

"""
    τ_relax(liquid_param_set)
    τ_relax(ice_param_set)

 - `liquid_param_set` - abstract set with cloud liquid water parameters
 - `ice_param_set` - abstract set with cloud ice parameters

Returns the relaxation timescale for condensation and evaporation of
cloud liquid water or the relaxation timescale for sublimation and
deposition of cloud ice.
"""
function τ_relax(liquid_param_set::ALPS)

    _τ_relax = τ_cond_evap(liquid_param_set)
    return _τ_relax
end
function τ_relax(ice_param_set::AIPS)

    _τ_relax = τ_sub_dep(ice_param_set)
    return _τ_relax
end

"""
    G_func(param_set, T, Liquid())
    G_func(param_set, T, Ice())

 - `param_set` - abstract set with earth parameters
 - `T` - air temperature
 - `Liquid()`, `Ice()` - liquid or ice phase to dispatch over.

Utility function combining thermal conductivity and vapor diffusivity effects.
"""
function G_func(param_set::APS, T::FT, ::Liquid) where {FT <: Real}

    _K_therm::FT = K_therm(param_set)
    _R_v::FT = R_v(param_set)
    _D_vapor::FT = D_vapor(param_set)

    L = latent_heat_vapor(param_set, T)
    p_vs = saturation_vapor_pressure(param_set, T, Liquid())

    return FT(1) / (
        L / _K_therm / T * (L / _R_v / T - FT(1)) + _R_v * T / _D_vapor / p_vs
    )
end
function G_func(param_set::APS, T::FT, ::Ice) where {FT <: Real}

    _K_therm::FT = K_therm(param_set)
    _R_v::FT = R_v(param_set)
    _D_vapor::FT = D_vapor(param_set)

    L = latent_heat_sublim(param_set, T)
    p_vs = saturation_vapor_pressure(param_set, T, Ice())

    return FT(1) / (
        L / _K_therm / T * (L / _R_v / T - FT(1)) + _R_v * T / _D_vapor / p_vs
    )
end

"""
    terminal_velocity(param_set, precip_param_set, ρ, q_)

 - `param_set` - abstract set with earth parameters
 - `precip_param_set` - abstract set with rain or snow parameters
 - `ρ` - air density
 - `q_` - rain or snow specific humidity

Returns the mass weighted average terminal velocity assuming
a Marshall-Palmer (1948) distribution of rain drops and snow crystals.
"""
function terminal_velocity(
    param_set::APS,
    precip_param_set::APrecipPS,
    ρ::FT,
    q_::FT,
) where {FT <: Real}
    fall_w = FT(0)
    if q_ > FT(0)

        (_n0, _r0, _m0, _me, _χm, _Δm, _a0, _ae, _χa, _Δa, _v0, _ve, _χv, _Δv) =
            unpack_params(param_set, precip_param_set, ρ, q_)

        _λ::FT = lambda(q_, ρ, _n0, _m0, _me, _r0, _χm, _Δm)

        fall_w =
            _χv *
            _v0 *
            (_λ * _r0)^(-_ve - _Δv) *
            gamma(_me + _ve + _Δm + _Δv + FT(1)) / gamma(_me + _Δm + FT(1))
    end

    return fall_w
end

"""
    conv_q_vap_to_q_liq_ice(liquid_param_set::ALPS, q_sat, q)
    conv_q_vap_to_q_liq_ice(ice_param_set::AIPS, q_sat, q)

 - `liquid_param_set` - abstract set with cloud water parameters
 - `ice_param_set` - abstract set with cloud ice parameters
 - `q_sat` - PhasePartition at equilibrium
 - `q` - current PhasePartition

Returns the cloud water tendency due to condensation and evaporation
or cloud ice tendency due to sublimation and vapor deposition.
The tendency is obtained assuming a relaxation to equilibrium with
a constant timescale.
"""
function conv_q_vap_to_q_liq_ice(
    liquid_param_set::ALPS,
    q_sat::PhasePartition{FT},
    q::PhasePartition{FT},
) where {FT <: Real}

    _τ_cond_evap::FT = τ_relax(liquid_param_set)

    return (q_sat.liq - q.liq) / _τ_cond_evap
end
function conv_q_vap_to_q_liq_ice(
    ice_param_set::AIPS,
    q_sat::PhasePartition{FT},
    q::PhasePartition{FT},
) where {FT <: Real}

    _τ_sub_dep::FT = τ_relax(ice_param_set)

    return (q_sat.ice - q.ice) / _τ_sub_dep
end

"""
    conv_q_liq_to_q_rai(rain_param_set, q_liq)

 - `rain_param_set` - abstract set with rain microphysics parameters
 - `q_liq` - liquid water specific humidity

Returns the q_rai tendency due to collisions between cloud droplets
(autoconversion), parametrized following Kessler (1995).
"""
function conv_q_liq_to_q_rai(rain_param_set::ARPS, q_liq::FT) where {FT <: Real}

    _τ_acnv::FT = τ_acnv(rain_param_set)
    _q_liq_threshold::FT = q_liq_threshold(rain_param_set)

    return max(0, q_liq - _q_liq_threshold) / _τ_acnv
end

"""
    conv_q_ice_to_q_sno(param_set, ice_param_set, q, ρ, T)

 - `param_set` - abstract set with earth parameters
 - `ice_param_set` - abstract set with ice microphysics parameters
 - `q` - phase partition
 - `ρ` - air density
 - `T` - air temperature

Returns the q_sno tendency due to autoconversion from ice.
Parameterized following Harrington et al. (1996) and Kaul et al. (2015).
"""
function conv_q_ice_to_q_sno(
    param_set::APS,
    ice_param_set::AIPS,
    q::PhasePartition{FT},
    ρ::FT,
    T::FT,
) where {FT <: Real}
    acnv_rate = FT(0)
    _S::FT = supersaturation(param_set, q, ρ, T, Ice())

    if (q.ice > FT(0) && _S > FT(0))

        _G::FT = G_func(param_set, T, Ice())

        _r_ice_snow::FT = r_ice_snow(ice_param_set)

        (_n0, _r0, _m0, _me, _χm, _Δm) =
            unpack_params(param_set, ice_param_set, ρ, q.ice)

        _λ::FT = lambda(q.ice, ρ, _n0, _m0, _me, _r0, _χm, _Δm)

        acnv_rate =
            4 * FT(π) * _S * _G * _n0 / ρ *
            exp(-_λ * _r_ice_snow) *
            (
                _r_ice_snow^FT(2) / (_me + _Δm) +
                (_r_ice_snow * _λ + FT(1)) / _λ^FT(2)
            )
    end
    return acnv_rate
end

"""
    accretion(param_set, cloud_param_set, precip_param_set, q_clo, q_pre, ρ)

 - `param_set` - abstract set with earth parameters
 - `cloud_param_set` - abstract set with cloud water or cloud ice parameters
 - `precip_param_set` - abstract set with rain or snow parameters
 - `q_clo` - cloud water or cloud ice specific humidity
 - `q_pre` - rain water or snow specific humidity
 - `ρ` - rain water or snow specific humidity

Returns the sink of cloud water (liquid or ice) due to collisions
with precipitating water (rain or snow).
"""
function accretion(
    param_set::APS,
    cloud_param_set::ACloudPS,
    precip_param_set::APrecipPS,
    q_clo::FT,
    q_pre::FT,
    ρ::FT,
) where {FT <: Real}

    accr_rate = FT(0)
    if (q_clo > FT(0) && q_pre > FT(0))

        (_n0, _r0, _m0, _me, _χm, _Δm, _a0, _ae, _χa, _Δa, _v0, _ve, _χv, _Δv) =
            unpack_params(param_set, precip_param_set, ρ, q_pre)

        _λ::FT = lambda(q_pre, ρ, _n0, _m0, _me, _r0, _χm, _Δm)
        _E::FT = E(cloud_param_set, precip_param_set)

        accr_rate =
            q_clo * _E * _n0 * _a0 * _v0 * _χa * _χv / _λ *
            gamma(_ae + _ve + _Δa + _Δv + FT(1)) /
            (_λ * _r0)^(_ae + _ve + _Δa + _Δv)
    end
    return accr_rate
end

"""
    accretion_rain_sink(param_set, ice_param_set, rain_param_set,
                        q_ice, q_rai, ρ)

 - `param_set` - abstract set with earth parameters
 - `ice_param_set` - abstract set with cloud ice parameters
 - `rain_param_set` - abstract set with rain parameters
 - `q_ice` - cloud ice specific humidity
 - `q_rai` - rain water specific humidity
 - `ρ` - air density

Returns the sink of rain water (partial source of snow) due to collisions
with cloud ice.
"""
function accretion_rain_sink(
    param_set::APS,
    ice_param_set::AIPS,
    rain_param_set::ARPS,
    q_ice::FT,
    q_rai::FT,
    ρ::FT,
) where {FT <: Real}

    accr_rate = FT(0)
    if (q_ice > FT(0) && q_rai > FT(0))

        (_n0_ice, _r0_ice, _m0_ice, _me_ice, _χm_ice, _Δm_ice) =
            unpack_params(param_set, ice_param_set, ρ, q_ice)

        (
            _n0_rai,
            _r0_rai,
            _m0_rai,
            _me_rai,
            _χm_rai,
            _Δm_rai,
            _a0_rai,
            _ae_rai,
            _χa_rai,
            _Δa_rai,
            _v0_rai,
            _ve_rai,
            _χv_rai,
            _Δv_rai,
        ) = unpack_params(param_set, rain_param_set, ρ, q_rai)

        _E::FT = E(ice_param_set, rain_param_set)

        _λ_rai::FT = lambda(
            q_rai,
            ρ,
            _n0_rai,
            _m0_rai,
            _me_rai,
            _r0_rai,
            _χm_rai,
            _Δm_rai,
        )
        _λ_ice::FT = lambda(
            q_ice,
            ρ,
            _n0_ice,
            _m0_ice,
            _me_ice,
            _r0_ice,
            _χm_ice,
            _Δm_ice,
        )

        accr_rate =
            _E / ρ *
            _n0_rai *
            _n0_ice *
            _m0_rai *
            _a0_rai *
            _v0_rai *
            _χm_rai *
            _χa_rai *
            _χv_rai / _λ_ice / _λ_rai * gamma(
                _me_rai +
                _ae_rai +
                _ve_rai +
                _Δm_rai +
                _Δa_rai +
                _Δv_rai +
                FT(1),
            ) /
            (
                _r0_rai * _λ_rai
            )^(_me_rai + _ae_rai + _ve_rai + _Δm_rai + _Δa_rai + _Δv_rai)
    end
    return accr_rate
end

"""
    accretion_snow_rain(param_set, i_param_set, j_param_set, q_i, q_j, ρ)

 - `i` - snow for temperatures below freezing
         or rain for temperatures above freezing
 - `j` - rain for temperatures below freezing
         or rain for temperatures above freezing
 - `param_set` - abstract set with earth parameters
 - `i_param_set`, `j_param_set` - abstract set with snow or rain
    microphysics parameters
 - `q_` - specific humidity of snow or rain
 - `ρ` - air density

Returns the accretion rate between rain and snow.
Collisions between rain and snow result in
snow at temperatures below freezing and in rain at temperatures above freezing.
"""
function accretion_snow_rain(
    param_set::APS,
    i_param_set::APrecipPS,
    j_param_set::APrecipPS,
    q_i::FT,
    q_j::FT,
    ρ::FT,
) where {FT <: Real}

    accr_rate = FT(0)
    if (q_i > FT(0) && q_j > FT(0))

        (
            _n0_i,
            _r0_i,
            _m0_i,
            _me_i,
            _χm_i,
            _Δm_i,
            _a0_i,
            _ae_i,
            _χa_i,
            _Δa_i,
            _v0_i,
            _ve_i,
            _χv_i,
            _Δv_i,
        ) = unpack_params(param_set, i_param_set, ρ, q_i)
        (
            _n0_j,
            _r0_j,
            _m0_j,
            _me_j,
            _χm_j,
            _Δm_j,
            _a0_j,
            _ae_j,
            _χa_j,
            _Δa_j,
            _v0_j,
            _ve_j,
            _χv_j,
            _Δv_j,
        ) = unpack_params(param_set, j_param_set, ρ, q_j)

        _E_ij::FT = E(i_param_set, j_param_set)

        _λ_i::FT = lambda(q_i, ρ, _n0_i, _m0_i, _me_i, _r0_i, _χm_i, _Δm_i)
        _λ_j::FT = lambda(q_j, ρ, _n0_j, _m0_j, _me_j, _r0_j, _χm_j, _Δm_j)

        _v_ti = terminal_velocity(param_set, i_param_set, ρ, q_i)
        _v_tj = terminal_velocity(param_set, j_param_set, ρ, q_j)

        accr_rate =
            FT(π) / ρ *
            _n0_i *
            _n0_j *
            _m0_j *
            _χm_j *
            _E_ij *
            abs(_v_ti - _v_tj) / _r0_j^(_me_j + _Δm_j) * (
                FT(2) * gamma(_me_j + _Δm_j + FT(1)) / _λ_i^FT(3) /
                _λ_j^(_me_j + _Δm_j + FT(1)) +
                FT(2) * gamma(_me_j + _Δm_j + FT(2)) / _λ_i^FT(2) /
                _λ_j^(_me_j + _Δm_j + FT(2)) +
                gamma(_me_j + _Δm_j + FT(3)) / _λ_i /
                _λ_j^(_me_j + _Δm_j + FT(3))
            )
    end
    return accr_rate
end

"""
    evaporation_sublimation(param_set, rain_param_set, q, q_rai, ρ, T)
    evaporation_sublimation(param_set, snow_param_set, q, q_sno, ρ, T)

 - `param_set` - abstract set with earth parameters
 - `rain_param_set` - abstract set with rain microphysics parameters
 - `snow_param_set` - abstract set with snow microphysics parameters
 - `q` - phase partition
 - `q_rai` - rain specific humidity
 - `q_sno` - snow specific humidity
 - `ρ` - air density
 - `T` - air temperature

Returns the tendency due to rain evaporation or snow sublimation.
"""
function evaporation_sublimation(
    param_set::APS,
    rain_param_set::ARPS,
    q::PhasePartition{FT},
    q_rai::FT,
    ρ::FT,
    T::FT,
) where {FT <: Real}
    evap_subl_rate = FT(0)
    _S::FT = supersaturation(param_set, q, ρ, T, Liquid())

    if (q_rai > FT(0) && _S < FT(0))

        _a_vent::FT = a_vent(rain_param_set)
        _b_vent::FT = b_vent(rain_param_set)
        _ν_air::FT = ν_air(param_set)
        _D_vapor::FT = D_vapor(param_set)

        _G::FT = G_func(param_set, T, Liquid())

        (_n0, _r0, _m0, _me, _χm, _Δm, _a0, _ae, _χa, _Δa, _v0, _ve, _χv, _Δv) =
            unpack_params(param_set, rain_param_set, ρ, q_rai)

        _λ::FT = lambda(q_rai, ρ, _n0, _m0, _me, _r0, _χm, _Δm)

        evap_subl_rate =
            4 * FT(π) * _n0 / ρ * _S * _G / _λ^FT(2) * (
                _a_vent +
                _b_vent * (_ν_air / _D_vapor)^FT(1 / 3) /
                (_r0 * _λ)^((_ve + _Δv) / FT(2)) *
                (FT(2) * _v0 * _χv / _ν_air / _λ)^FT(1 / 2) *
                gamma((_ve + _Δv + FT(5)) / FT(2))
            )
    end
    return evap_subl_rate
end
function evaporation_sublimation(
    param_set::APS,
    snow_param_set::ASPS,
    q::PhasePartition{FT},
    q_sno::FT,
    ρ::FT,
    T::FT,
) where {FT <: Real}
    evap_subl_rate = FT(0)
    if q_sno > FT(0)

        _a_vent::FT = a_vent(snow_param_set)
        _b_vent::FT = b_vent(snow_param_set)
        _ν_air::FT = ν_air(param_set)
        _D_vapor::FT = D_vapor(param_set)

        _S::FT = supersaturation(param_set, q, ρ, T, Ice())
        _G::FT = G_func(param_set, T, Ice())

        (_n0, _r0, _m0, _me, _χm, _Δm, _a0, _ae, _χa, _Δa, _v0, _ve, _χv, _Δv) =
            unpack_params(param_set, snow_param_set, ρ, q_sno)
        _λ::FT = lambda(q_sno, ρ, _n0, _m0, _me, _r0, _χm, _Δm)

        evap_subl_rate =
            4 * FT(π) * _n0 / ρ * _S * _G / _λ^FT(2) * (
                _a_vent +
                _b_vent * (_ν_air / _D_vapor)^FT(1 / 3) /
                (_r0 * _λ)^((_ve + _Δv) / FT(2)) *
                (FT(2) * _v0 * _χv / _ν_air / _λ)^FT(1 / 2) *
                gamma((_ve + _Δv + FT(5)) / FT(2))
            )
    end
    return evap_subl_rate
end

"""
    snow_melt(param_set, snow_param_set, q_sno, ρ, T)

 - `param_set` - abstract set with earth parameters
 - `snow_param_set` - abstract set with snow microphysics parameters
 - `q_sno` - snow water specific humidity
 - `ρ` - air density
 - `T` - air temperature

Returns the tendency due to snow melt.
"""
function snow_melt(
    param_set::APS,
    snow_param_set::ASPS,
    q_sno::FT,
    ρ::FT,
    T::FT,
) where {FT <: Real}

    snow_melt_rate = FT(0)
    _T_freeze = T_freeze(param_set)

    if (q_sno > FT(0) && T > _T_freeze)

        _a_vent::FT = a_vent(snow_param_set)
        _b_vent::FT = b_vent(snow_param_set)
        _ν_air::FT = ν_air(param_set)
        _D_vapor::FT = D_vapor(param_set)
        _K_therm::FT = K_therm(param_set)

        L = latent_heat_fusion(param_set, T)

        (_n0, _r0, _m0, _me, _χm, _Δm, _a0, _ae, _χa, _Δa, _v0, _ve, _χv, _Δv) =
            unpack_params(param_set, snow_param_set, ρ, q_sno)
        _λ::FT = lambda(q_sno, ρ, _n0, _m0, _me, _r0, _χm, _Δm)

        snow_melt_rate =
            4 * FT(π) * _n0 / ρ * _K_therm / L * (T - _T_freeze) / _λ^FT(2) * (
                _a_vent +
                _b_vent * (_ν_air / _D_vapor)^FT(1 / 3) /
                (_r0 * _λ)^((_ve + _Δv) / FT(2)) *
                (FT(2) * _v0 * _χv / _ν_air / _λ)^FT(1 / 2) *
                gamma((_ve + _Δv + FT(5)) / FT(2))
            )
    end
    return snow_melt_rate
end

end #module Microphysics.jl

max_stars_repo_name: mwarusz/CLIMA

max_issues_repo_name: mwarusz/CLIMA

max_forks_repo_name: mwarusz/CLIMA


text:
using POMDPGym, POMDPs, Test, Distributions, Random

mdp = ContinuumWorldMDP()
x_dist = MvNormal([0,0], [1,1])

rand(x_dist)
adv_mdp = AdditiveAdversarialMDP(mdp, x_dist)

@test adv_mdp.mdp == mdp
@test adv_mdp.x_distribution isa MvNormal

s = rand(initialstate(adv_mdp))
x = rand(adv_mdp.x_distribution)
a = -x

sp,r = gen(adv_mdp, s, a, x)
sp_2,r_2 = gen(adv_mdp.mdp, s, [0,0])

@test all(sp .≈ sp_2)
@test r ≈ r_2

gen(adv_mdp, s, [0.,0.])

ext: jl

max_forks_repo_name: ancorso/POMDPGym


text:
# functions related to Poisson distribution

import .RFunctions:
    poispdf,
    poislogpdf,
    poiscdf,
    poisccdf,
    poislogcdf,
    poislogccdf,
    poisinvcdf,
    poisinvccdf,
    poisinvlogcdf,
    poisinvlogccdf

# generic versions
poispdf(λ::Real, x::Real) = exp(poislogpdf(λ, x))

poislogpdf(λ::T, x::T) where {T <: Real} = xlogy(x, λ) - λ - lgamma(x + 1)

poislogpdf(λ::Number, x::Number) = poislogpdf(promote(float(λ), x)...)

#=
function poislogpdf(λ::Union{Float32,Float64}, x::Union{Float64,Float32,Integer})
    if iszero(λ)
        iszero(x) ? zero(λ) : oftype(λ, -Inf)
    elseif iszero(x)
        -λ
    else
        -lstirling_asym(x + 1)
=#


text:
# This file is a part of Julia. License is MIT: https://julialang.org/license

module ModInts
export ModInt

import Base: +, -, *, /, inv

struct ModInt{n} <: Integer
    k::Int
    ModInt{n}(k) where {n} = new(mod(k,n))
end

Base.show(io::IO, k::ModInt{n}) where {n} =
    print(io, get(io, :compact, false) ? k.k : "$(k.k) mod $n")

(+)(a::ModInt{n}, b::ModInt{n}) where {n} = ModInt{n}(a.k+b.k)
(-)(a::ModInt{n}, b::ModInt{n}) where {n} = ModInt{n}(a.k-b.k)
(*)(a::ModInt{n}, b::ModInt{n}) where {n} = ModInt{n}(a.k*b.k)
(-)(a::ModInt{n}) where {n} = ModInt{n}(-a.k)

inv(a::ModInt{n}) where {n} = ModInt{n}(invmod(a.k, n))
(/)(a::ModInt{n}, b::ModInt{n}) where {n} = a*inv(b) # broaden for non-coprime?

Base.promote_rule(::Type{ModInt{n}}, ::Type{Int}) where {n} = ModInt{n}
Base.convert(::Type{ModInt{n}}, i::Int) where {n} = ModInt{n}(i)

end # module

max_stars_repo_name: caodi/julia

max_issues_repo_name: caodi/julia

max_forks_repo_name: caodi/julia


text:
mutable struct Change

    var
    added::Set{Int}
    removed::Set{Int}

    fixed::Bool
    cardAdded::Int
    cardRemoved::Int
end

function Change(var)

    return Change(var,
                    Set{Int}(),
                    Set{Int}(),
                    false,
                    0,
                    0)
end


text:
using CompScienceMeshes
using BEAST

using Base.Test
using StaticArrays

T = Float64
P = SVector{3,T}

tol = eps(T) * 10^3

function shapevals(ϕ, ts)

    numpoints = length(ts)
    numshapes = numfunctions(ϕ)

    y = Array{P}(numshapes, numpoints)
    for i in 1 : numpoints
        t = ts[i]
        u = ϕ(t)
        for j in 1 : numshapes
            y[j,i] = u[j][1]
        end
    end

    return y
end

const third = one(T)/3
m = meshrectangle(1.0,1.0,0.5);
rwg = raviartthomas(m)
ref = refspace(rwg)

# vrts = vertices(m, first(cells(m)))
# ptch = simplex(vrts[1], vrts[2], vrts[3])
ptch = chart(m, first(cells(m)))

ctrd = neighborhood(ptch, [1,1]/3)
vals = shapevals(ref, [ctrd])

# test edge detection
edges = skeleton(m, 1)
vp = cellpairs(m,edges)

# test function evaluation
v = [
    point(0.0, 0.0, 0.0),
    point(2.0, 0.0, 0.0),
    point(0.0, 3.0, 0.0)]

cell = simplex(v[1], v[2], v[3])

mp = neighborhood(cell,[third, third]);
ϕ = BEAST.RTRefSpace{Float64}()
fx = ϕ(mp)[2][1]

A  = volume(cell)
r = cartesian(mp)
gx = (r - v[2]) / 2A
@test norm(fx - gx) < tol

@test norm((r - v[1])/2A -  ϕ(mp)[1][1]) < tol
@test norm((r - v[2])/2A -  ϕ(mp)[2][1]) < tol
@test norm((r - v[3])/2A -  ϕ(mp)[3][1]) < tol

# repeat the test but now on a random cell
randpoint() = point(2*rand()-1, 2*rand()-1, 2*rand()-1)
v = [randpoint(), randpoint(), randpoint()]
cell = simplex(v[1], v[2], v[3])
mp = neighborhood(cell,[third, third]);

fx = ϕ(mp)[2][1]
#fx = evalfun(rs, mp)

A = volume(cell)
r = cartesian(mp)
gx = (r-v[2]) / 2A
@test norm(fx - gx) < tol

@test norm((r - v[1])/2A -  ϕ(mp)[1][1]) < tol
@test norm((r - v[2])/2A -  ϕ(mp)[2][1]) < tol
@test norm((r - v[3])/2A -  ϕ(mp)[3][1]) < tol


text:
using Statistics, BSON
using Flux, Flux.Optimise
using MLDatasets: CIFAR10
using Images.ImageCore
using Flux: onehotbatch, onecold
using Base.Iterators: partition
using CUDA

# Data
x, y = CIFAR10.traindata(Float32)

# Turn into binary problem (cats/dogs)
class_names = CIFAR10.classnames()
n_classes = length(class_names)
c_cat = findall(class_names.=="cat") .- 1
c_dog = findall(class_names.=="dog") .- 1
cats_and_dogs = findall(y .∈ Ref(vcat(c_cat, c_dog)))

x = x[:,:,:,cats_and_dogs]
y = y[cats_and_dogs]
y = ifelse.(y.==c_cat,0,1)

n = length(y)
batch_size = Int(n/10)
train_set = 1:(n-batch_size)
val_set = (n-batch_size+1):n
train_x = x[:,:,:,train_set]
train_y = y[train_set]
val_x = x[:,:,:,val_set] |> gpu
val_y = y[val_set] |> gpu
data = Flux.DataLoader((data=train_x, label=train_y), batchsize=batch_size) |> gpu

# CNN
nn = Chain(
    Conv((5,5), 3=>16, relu),
    MaxPool((2,2)),
    Conv((5,5), 16=>8, relu),
    MaxPool((2,2)),
    x -> reshape(x, :, size(x, 4)),
    Dense(200, 120),
    Dense(120, 84),
    Dense(84, 1)
) |> gpu

using Flux: Momentum
using Flux.Losses: logitbinarycrossentropy
loss(x, y) = logitbinarycrossentropy(vec(nn(x)), y)
opt = Momentum(0.01)
accuracy(x, y) = mean(vec(round.(Flux.σ.(nn(x)))) .== y)
avg_loss(data) = mean(map(d -> loss(d[1],d[2]), data))
epochs = 200

if retrain

    # Training
    for epoch = 1:epochs
        for d in data
            gs = gradient(Flux.params(nn)) do
            l = loss(d...)
            end
            update!(opt, Flux.params(nn), gs)
        end
        @info "Epoch " * string(epoch)
        println("Training loss:")
        @show avg_loss(data)
        println("Training accuracy:")
        @show accuracy(train_x,train_y)
        println("Test accuracy:")
        @show accuracy(val_x, val_y)
        if epoch % 10 == 0
            BSON.@save "data/CIFAR10_nn_" * string(epoch) * ".bson" nn
        end
    end
end


text:
file_name_default = joinpath(dirname(@__FILE__), "type_db_default.tql")
file_name_schema = joinpath(dirname(@__FILE__), "type_db_schema.tql")
file_name = joinpath(dirname(@__FILE__), "type_db.tql")

if !ispath(file_name)
    # customize the runscript to the environment
    text = read(open(file_name_default), String)
    text = replace(text, "\$file_path"=>file_name_schema)
    write(file_name, text)
    # run the customized script
    run(Cmd(`typedb console --script=$file_name`))
    # remove the line of the creation of the database
    lines = readlines(file_name)
    text = join(deleteat!(lines, 1),"\n")
    write(file_name, text)
end
# run the schema script
run(Cmd(`typedb console --script=$file_name`))

max_stars_repo_name: Humans-of-Julia/TypeDBORM

max_issues_repo_name: Humans-of-Julia/TypeDBORM

max_forks_repo_name: Humans-of-Julia/TypeDBORM


text:
# MacroTools.prewalk(MacroTools.rmlines, MacroTools.gensym_ids(expr2)) |> MacroTools.flatten

coinflip_1(; y) = coinflip_1(y)

function coinflip_1(y)
    function inner_function_2(
        vi_3::DynamicPPL.VarInfo,
        sampler_4::DynamicPPL.AbstractSampler,
        ctx_5::DynamicPPL.AbstractContext,
        model_6
    )

        local y
        temp_var_7 = model_6.args.y
        varT_8 = typeof(temp_var_7)

        if temp_var_7 isa DynamicPPL.FloatOrArrayType
            y = DynamicPPL.get_matching_type(sampler_4, vi_3, temp_var_7)
        elseif DynamicPPL.hasmissing(varT_8)
            y = (DynamicPPL.get_matching_type(sampler_4, vi_3, varT_8))(temp_var_7)
        else
            y = temp_var_7
        end

        vi_3.logp = 0
        temp_right_9 = Beta(1, 1)
        DynamicPPL.assert_dist(temp_right_9, msg="Right-hand side of a ~ must be subtype of Distribution or a vector of Distributions on line 449.")

        preprocessed_10 = begin
            sym_11 = Val(:p)
            if !(DynamicPPL.inparams(sym_11, Val{(:y,)}())) || DynamicPPL.inparams(sym_11, DynamicPPL.getmissing(model_6))
                (DynamicPPL.VarName{:p}(""), ())
            else
                if DynamicPPL.inparams(sym_11, Val{(:y,)}())
                    lhs_12 = p
                    if lhs_12 === missing
                        (DynamicPPL.VarName{:p}(""), ())
                    else
                        lhs_12
                    end
                else
                    throw("This point should not be reached. Please report this error.")
                end
            end
        end

        if preprocessed_10 isa Tuple
            (vn_13, inds_14) = preprocessed_10
            out_15 = DynamicPPL.tilde(ctx_5, sampler_4, temp_right_9, vn_13, inds_14, vi_3)
            p = out_15[1]
            vi_3.logp += out_15[2]
        else
            vi_3.logp += DynamicPPL.tilde(ctx_5, sampler_4, temp_right_9, preprocessed_10, vi_3)
        end

        N = length(y)

        for n = 1:N
            temp_right_16 = Bernoulli(p)
            DynamicPPL.assert_dist(temp_right_16, msg="Right-hand side of a ~ must be subtype of Distribution or a vector of Distributions on line 449.")
            preprocessed_17 = begin
                sym_18 = Val(:y)
                if !(DynamicPPL.inparams(sym_18, Val{(:y,)}())) || DynamicPPL.inparams(sym_18, DynamicPPL.getmissing(model_6))
                    (DynamicPPL.VarName{:y}(foldl(*, ("[" * join([if n === (:)
                                                                  "Colon()"
                                                                  else
                                                                  string(n)
                                                                  end], ",") * "]",), init="")), ((n,),))
                else
                    if DynamicPPL.inparams(sym_18, Val{(:y,)}())
                        lhs_19 = y[n]
                        if lhs_19 === missing
                            (DynamicPPL.VarName{:y}(foldl(*, ("[" * join([if n === (:)
                                                                          "Colon()"
                                                                          else
                                                                          string(n)
                                                                          end], ",") * "]",), init="")),
                             ((n,),))
                        else
                            lhs_19
                        end
                    else
                        throw("This point should not be reached. Please report this error.")
                    end
                end
            end

            if preprocessed_17 isa Tuple
                (vn_20, inds_21) = preprocessed_17
                out_22 = DynamicPPL.tilde(ctx_5, sampler_4, temp_right_16, vn_20, inds_21, vi_3)
                y[n] = out_22[1]
                vi_3.logp += out_22[2]
            else
                vi_3.logp += DynamicPPL.tilde(ctx_5, sampler_4, temp_right_16, preprocessed_17, vi_3)
            end
        end
    end

    return DynamicPPL.Model(inner_function_2,
                            DynamicPPL.namedtuple(NamedTuple{(:y,), Tuple{DynamicPPL.get_type(y)}}, (y,)),
                            DynamicPPL.ModelGen{(:y,)}(coinflip_1, NamedTuple()))
end

coinflip = DynamicPPL.ModelGen{(:y,)}(coinflip_1, NamedTuple())


text:
using PyPlot
using SpecialFunctions
using MittagLefflerFunctions
using OffsetArrays

function g(y)
    x = (1+y) / (1-y)
    return ( erfcx(x) - one(y) ) / x
end

function Eneg(a, x)
    y = (x-1) / (x+1)
    return 1 + x * chebyshev_sum(a, y)
end

nmax = 20
a = OffsetArray{Float64}(undef, 0:nmax)
M = 2nmax
chebyshev_coefs!(a, g, M)

figure(1)
y = range(-1, 1, length=201)
plot(y, g.(y))
grid(true)

figure(2)
approx_g = Float64[ chebyshev_sum(a, yval) for yval in y ]
plot(y, approx_g - g.(y))
grid(true)

figure(3)
x = range(0, 5, length=201)
E = Float64[ Eneg(a, xval) for xval in x ]
plot(x, E - erfcx.(x))
grid(true)


text:
"""
    BG96
Data from [Bollerslev and Ghysels (JBES 1996)](https://doi.org/10.2307/1392425).
"""
const BG96 = readdlm(joinpath(dirname(pathof(ARCHModels)), "data", "bollerslev_ghysels.txt"), skipstart=1)[:, 1];

"""
    UnivariateVolatilitySpec{T} <: VolatilitySpec{T} end

Abstract supertype that univariate volatility specifications inherit from.
"""
abstract type UnivariateVolatilitySpec{T} <: VolatilitySpec{T} end

"""
    StandardizedDistribution{T} <: Distributions.Distribution{Univariate, Continuous}

Abstract supertype that standardized distributions inherit from.
"""
abstract type StandardizedDistribution{T} <: Distribution{Univariate, Continuous} end

"""
    UnivariateARCHModel{T<:AbstractFloat,
              		    VS<:UnivariateVolatilitySpec,
              			SD<:StandardizedDistribution{T},
              			MS<:MeanSpec{T}
              			} <: ARCHModel
"""
mutable struct UnivariateARCHModel{T<:AbstractFloat,
                 				   VS<:UnivariateVolatilitySpec,
                 		  	  	   SD<:StandardizedDistribution{T},
                 				   MS<:MeanSpec{T}
                 				   } <: ARCHModel
    spec::VS
    data::Vector{T}
    dist::SD
    meanspec::MS
	fitted::Bool
    function UnivariateARCHModel{T, VS, SD, MS}(spec, data, dist, meanspec, fitted) where {T, VS, SD, MS}
        new(spec, data, dist, meanspec, fitted)
    end
end

mutable struct UnivariateSubsetARCHModel{T<:AbstractFloat,
                 				   VS<:UnivariateVolatilitySpec,
                 		  	  	   SD<:StandardizedDistribution{T},
                 				   MS<:MeanSpec{T},
								   N
                 				   } <: ARCHModel
	spec::VS
    data::Vector{T}
    dist::SD
    meanspec::MS
	fitted::Bool
	subset::NTuple{N, Int}
    function UnivariateSubsetARCHModel{T, VS, SD, MS, N}(spec, data, dist, meanspec, fitted, subset) where {T, VS, SD, MS, N}
        new(spec, data, dist, meanspec, fitted, subset)
    end
end

"""
    UnivariateARCHModel(spec::UnivariateVolatilitySpec, data::Vector; dist=StdNormal(),
	          			meanspec=NoIntercept(), fitted=false
              			)

Create a UnivariateARCHModel.

# Example:
```jldoctest
julia> UnivariateARCHModel(GARCH{1, 1}([1., .9, .05]), randn(10))

GARCH{1, 1} model with Gaussian errors, T=10.

─────────────────────────────────────────
                             ω   β₁    α₁
─────────────────────────────────────────
Volatility parameters:     1.0  0.9  0.05
─────────────────────────────────────────
```
"""
function UnivariateARCHModel(spec::VS,
          		 			 data::Vector{T};
          					 dist::SD=StdNormal{T}(),
          				 	 meanspec::MS=NoIntercept{T}(),
		  			 		 fitted::Bool=false
          					 ) where {T<:AbstractFloat,
                    			 	  VS<:UnivariateVolatilitySpec,
                   					  SD<:StandardizedDistribution,
                   					  MS<:MeanSpec
                   			 		  }
    UnivariateARCHModel{T, VS, SD, MS}(spec, data, dist, meanspec, fitted)
end

function UnivariateSubsetARCHModel(spec::VS,
          		 			 data::Vector{T};
          					 dist::SD=StdNormal{T}(),
          				 	 meanspec::MS=NoIntercept{T}(),
		  			 		 fitted::Bool=false,
							 subset::NTuple{N, Int}
							 ) where {T<:AbstractFloat,
                    			 	  VS<:UnivariateVolatilitySpec,
                   					  SD<:StandardizedDistribution,
                   					  MS<:MeanSpec,
									  N
                   			 		  }
    UnivariateSubsetARCHModel{T, VS, SD, MS, N}(spec, data, dist, meanspec, fitted, subset)
end
loglikelihood(am::UnivariateARCHModel) = loglik(typeof(am.spec), typeof(am.dist),
                                      am.meanspec, am.data,
                                      vcat(am.spec.coefs, am.dist.coefs,
                                           am.meanspec.coefs
                                           )
                                      )

loglikelihood(am::UnivariateSubsetARCHModel) = loglik(typeof(am.spec), typeof(am.dist),
                                      am.meanspec, am.data,
                                      vcat(am.spec.coefs, am.dist.coefs,
                                           am.meanspec.coefs
                                           ),
									  subsetmask(typeof(am.spec), am.subset)
									  )

dof(am::UnivariateARCHModel) = nparams(typeof(am.spec)) + nparams(typeof(am.dist)) + nparams(typeof(am.meanspec))
dof(am::UnivariateSubsetARCHModel) = nparams(typeof(am.spec), am.subset) + nparams(typeof(am.dist)) + nparams(typeof(am.meanspec))
coef(am::UnivariateARCHModel)=vcat(am.spec.coefs, am.dist.coefs, am.meanspec.coefs)
coefnames(am::UnivariateARCHModel) = vcat(coefnames(typeof(am.spec)),
                                coefnames(typeof(am.dist)),
                                coefnames(am.meanspec)
                                )

# documented in general
function simulate(spec::UnivariateVolatilitySpec{T2}, nobs; warmup=100, dist::StandardizedDistribution{T2}=StdNormal{T2}(),
                  meanspec::MeanSpec{T2}=NoIntercept{T2}(),
				  rng=GLOBAL_RNG
                  ) where {T2<:AbstractFloat}
    data = zeros(T2, nobs)
    _simulate!(data,  spec; warmup=warmup, dist=dist, meanspec=meanspec, rng=rng)
    UnivariateARCHModel(spec, data; dist=dist, meanspec=meanspec, fitted=false)
end

function _simulate!(data::Vector{T2}, spec::UnivariateVolatilitySpec{T2};
                  warmup=100,
                  dist::StandardizedDistribution{T2}=StdNormal{T2}(),
                  meanspec::MeanSpec{T2}=NoIntercept{T2}(),
				  rng=GLOBAL_RNG
                  ) where {T2<:AbstractFloat}
	@assert warmup>=0
	append!(data, zeros(T2, warmup))
    T = length(data)
	r1 = presample(typeof(spec))
	r2 = presample(meanspec)
	r = max(r1, r2)
	r = max(r, 1) # make sure this works for, e.g., ARCH{0}; CircularBuffer requires at least a length of 1
    ht = CircularBuffer{T2}(r)
    lht = CircularBuffer{T2}(r)
    zt = CircularBuffer{T2}(r)
	at = CircularBuffer{T2}(r)
    @inbounds begin
        h0 = uncond(typeof(spec), spec.coefs)
		m0 = uncond(meanspec)
        h0 > 0 || error("Model is nonstationary.")
        for t = 1:T
			if t>r2
				themean = mean(at, ht, lht, data, meanspec, meanspec.coefs, t)
			else
				themean = m0
			end
			if t>r1
                update!(ht, lht, zt, at, typeof(spec), spec.coefs)
            else
				push!(ht, h0)
                push!(lht, log(h0))
            end
			push!(zt, rand(rng, dist))
			push!(at, sqrt(ht[end])*zt[end])
			data[t] = themean + at[end]
        end
    end
    deleteat!(data, 1:warmup)
end

@inline function splitcoefs(coefs, VS, SD, meanspec)
    ng = nparams(VS)
    nd = nparams(SD)
    nm = nparams(typeof(meanspec))
    length(coefs) == ng+nd+nm || throw(NumParamError(ng+nd+nm, length(coefs)))
    garchcoefs = coefs[1:ng]
    distcoefs = coefs[ng+1:ng+nd]
    meancoefs = coefs[ng+nd+1:ng+nd+nm]
    return garchcoefs, distcoefs, meancoefs
end
"""
    volatilities(am::UnivariateARCHModel)
Return the conditional volatilities.
"""
function volatilities(am::UnivariateARCHModel{T, VS, SD}) where {T, VS, SD}
	ht = Vector{T}(undef, 0)
	lht = Vector{T}(undef, 0)
	zt = Vector{T}(undef, 0)
	at = Vector{T}(undef, 0)
	loglik!(ht, lht, zt, at, VS, SD, am.meanspec, am.data, vcat(am.spec.coefs, am.dist.coefs, am.meanspec.coefs))
	return sqrt.(ht)
end

"""
    predict(am::UnivariateARCHModel, what=:volatility, horizon=1; level=0.01)
Form a `horizon`-step ahead prediction from `am`. `what` controls which object is predicted.
The choices are `:volatility` (the default), `:variance`, `:return`, and `:VaR`. The VaR
level can be controlled with the keyword argument `level`.

Not all prediction targets / volatility specifications support multi-step predictions.
"""
function predict(am::UnivariateARCHModel{T, VS, SD}, what=:volatility, horizon=1; level=0.01) where {T, VS, SD, MS}
	ht = volatilities(am).^2
	lht = log.(ht)
	zt = residuals(am)
	at = residuals(am, standardized=false)
	themean = T(0)
	if horizon > 1
		if what == :VaR
			error("Predicting VaR more than one period ahead is not implemented. Consider predicting one period ahead and scaling by `sqrt(horizon)`.")
		elseif what == :volatility
			error("Predicting volatility more than one period ahead is not implemented.")
		elseif what == :variance && !(VS <: TGARCH)
			error("Predicting variance more than one period ahead is not implemented for $(modname(VS)).")
		end
	end
    data = copy(am.data)
	for current_horizon = (1 : horizon)
		t = length(data) + current_horizon
		if what == :return || what == :VaR
			themean = mean(at, ht, lht, am.data, am.meanspec, am.meanspec.coefs, t)
		end
		update!(ht, lht, zt, at, VS, am.spec.coefs, current_horizon)
		push!(zt, 0.)
		push!(at, 0.)
        push!(data, themean)
	end
	if what == :return
		return themean
	elseif what == :volatility
		return sqrt(ht[end])
	elseif what == :variance
		return ht[end]
	elseif what == :VaR
		return -themean - sqrt(ht[end]) * quantile(am.dist, level)
	else error("Prediction target $what unknown.")
	end
end

"""
    means(am::UnivariateARCHModel)
Return the conditional means of the model.
"""
function means(am::UnivariateARCHModel)
	return am.data-residuals(am; standardized=false)
end

"""
    residuals(am::UnivariateARCHModel; standardized=true)
Return the residuals of the model. Pass `standardized=false` for the non-devolatized residuals.
"""
function residuals(am::UnivariateARCHModel{T, VS, SD}; standardized=true) where {T, VS, SD}
		ht = Vector{T}(undef, 0)
		lht = Vector{T}(undef, 0)
		zt = Vector{T}(undef, 0)
		at = Vector{T}(undef, 0)
		loglik!(ht, lht, zt, at, VS, SD, am.meanspec, am.data, vcat(am.spec.coefs, am.dist.coefs, am.meanspec.coefs))
	return standardized ? zt : at
end

"""
    VaRs(am::UnivariateARCHModel, level=0.01)
Return the in-sample Value at Risk implied by `am`.
"""
function VaRs(am::UnivariateARCHModel, level=0.01)
    return -means(am) .- volatilities(am) .* quantile(am.dist, level)
end

#this works on CircularBuffers. The idea is that ht/lht/zt need to be allocated
#inside of this function, when the type that Optim it with is known (because
#it calls it with dual numbers for autodiff to work). It works with arrays, too,
#but grows them by length(data); hence it should be called with an empty one-
#dimensional array of the right type.
@inline function loglik!(ht::AbstractVector{T2}, lht::AbstractVector{T2},
                         zt::AbstractVector{T2}, at::AbstractVector{T2}, vs::Type{VS}, ::Type{SD}, meanspec::MS,
                         data::Vector{T1}, coefs::AbstractVector{T3}, subsetmask=trues(nparams(vs))
                         ) where {VS<:UnivariateVolatilitySpec, SD<:StandardizedDistribution,
                                  MS<:MeanSpec, T1<:AbstractFloat, T2, T3
                                  }
    garchcoefs, distcoefs, meancoefs = splitcoefs(coefs, VS, SD, meanspec)
	lowergarch, uppergarch = constraints(VS, T1)
	lowerdist, upperdist = constraints(SD, T1)
    lowermean, uppermean = constraints(MS, T1)
    all(lowerdist.<distcoefs.<upperdist) && all(lowermean.<meancoefs.<uppermean) && all(lowergarch[subsetmask].<garchcoefs[subsetmask].<uppergarch[subsetmask]) || return T2(-Inf)
	garchcoefs .*= subsetmask
    T = length(data)
	r1 = presample(VS)
	r2 = presample(meanspec)
    r = max(r1, r2)
    T > r || error("Sample too small.")
	ki = kernelinvariants(SD, distcoefs)
    @inbounds begin
        h0 = var(data) # could be moved outside
		m0 = mean(data)
        #h0 = uncond(VS, garchcoefs)
        #h0 > 0 || return T2(NaN)
        LL = zero(T2)
        for t = 1:T
			if t>r2
				themean = mean(at, ht, lht, data, meanspec, meancoefs, t)
			else
				themean = m0
			end
			if t > r1
                update!(ht, lht, zt, at, VS, garchcoefs)
            else
				push!(ht, h0)
                push!(lht, log(h0))
            end
            ht[end] < 0 && return T2(NaN)
			push!(at, data[t]-themean)
			push!(zt, at[end]/sqrt(ht[end]))
			LL += -lht[end]/2 + logkernel(SD, zt[end], distcoefs, ki...)

        end#for
    end#inbounds
    LL += T*logconst(SD, distcoefs)
end#function

function loglik(spec::Type{VS}, dist::Type{SD}, meanspec::MS,
                   data::Vector{<:AbstractFloat}, coefs::AbstractVector{T2}, subsetmask=trues(nparams(spec))
                   ) where {VS<:UnivariateVolatilitySpec, SD<:StandardizedDistribution,
                            MS<:MeanSpec, T2
                            }
    r = max(presample(VS), presample(meanspec))
	r = max(r, 1) # make sure this works for, e.g., ARCH{0}; CircularBuffer requires at least a length of 1
    ht = CircularBuffer{T2}(r)
    lht = CircularBuffer{T2}(r)
    zt = CircularBuffer{T2}(r)
	at = CircularBuffer{T2}(r)
    loglik!(ht, lht, zt, at, spec, dist, meanspec, data, coefs, subsetmask)

end

function logliks(spec, dist, meanspec, data, coefs::Vector{T}) where {T}
    garchcoefs, distcoefs, meancoefs = splitcoefs(coefs, spec, dist, meanspec)
    ht = T[]
    lht = T[]
    zt = T[]
	at = T[]
    loglik!(ht, lht, zt, at, spec, dist, meanspec, data, coefs)
    LLs = -lht./2 .+ logkernel.(dist, zt, Ref{Vector{T}}(distcoefs), kernelinvariants(dist, distcoefs)...) .+ logconst(dist, distcoefs)
end

function informationmatrix(am::UnivariateARCHModel; expected::Bool=true)
	expected && error("expected informationmatrix is not implemented for UnivariateARCHModel. Use expected=false.")
	g = x -> sum(logliks(typeof(am.spec), typeof(am.dist), am.meanspec, am.data, x))
	H = ForwardDiff.hessian(g, vcat(am.spec.coefs, am.dist.coefs, am.meanspec.coefs))
	J = -H
end

function scores(am::UnivariateARCHModel)
	f = x -> logliks(typeof(am.spec), typeof(am.dist), am.meanspec, am.data, x)
	S = ForwardDiff.jacobian(f, vcat(am.spec.coefs, am.dist.coefs, am.meanspec.coefs))
end

function _fit!(garchcoefs::Vector{T}, distcoefs::Vector{T},
              meancoefs::Vector{T}, ::Type{VS}, ::Type{SD}, meanspec::MS,
              data::Vector{T}; algorithm=BFGS(), autodiff=:forward, kwargs...
              ) where {VS<:UnivariateVolatilitySpec, SD<:StandardizedDistribution,
                       MS<:MeanSpec, T<:AbstractFloat
                       }
    obj = x -> -loglik(VS, SD, meanspec, data, x)
    coefs = vcat(garchcoefs, distcoefs, meancoefs)
    res = optimize(obj, coefs, algorithm; autodiff=autodiff, kwargs...)
    coefs .= Optim.minimizer(res)
    ng = nparams(VS)
    ns = nparams(SD)
    nm = nparams(typeof(meanspec))
    garchcoefs .= coefs[1:ng]
    distcoefs .= coefs[ng+1:ng+ns]
    meancoefs .= coefs[ng+ns+1:ng+ns+nm]
	meanspec.coefs .= meancoefs
    return nothing
end

"""
    fit(VS::Type{<:UnivariateVolatilitySpec}, data; dist=StdNormal, meanspec=Intercept,
        algorithm=BFGS(), autodiff=:forward, kwargs...)

Fit the ARCH model specified by `VS` to `data`. `data` can be a vector or a
GLM.LinearModel (or GLM.TableRegressionModel).

# Keyword arguments:
- `dist=StdNormal`: the error distribution.
- `meanspec=Intercept`: the mean specification, either as a type or instance of that type.
- `algorithm=BFGS(), autodiff=:forward, kwargs...`: passed on to the optimizer.

# Example: EGARCH{1, 1, 1} model without intercept, Student's t errors.
```jldoctest
julia> fit(EGARCH{1, 1, 1}, BG96; meanspec=NoIntercept, dist=StdT)

EGARCH{1, 1, 1} model with Student's t errors, T=1974.

Volatility parameters:
──────────────────────────────────────────────
      Estimate  Std.Error    z value  Pr(>|z|)
──────────────────────────────────────────────
ω   -0.0162014  0.0186806  -0.867286    0.3858
γ₁  -0.0378454  0.018024   -2.09972     0.0358
β₁   0.977687   0.012558   77.8538      <1e-99
α₁   0.255804   0.0625497   4.08961     <1e-04
──────────────────────────────────────────────

Distribution parameters:
─────────────────────────────────────────
   Estimate  Std.Error  z value  Pr(>|z|)
─────────────────────────────────────────
ν   4.12423    0.40059  10.2954    <1e-24
─────────────────────────────────────────
```
"""
function fit(::Type{VS}, data::Vector{T}; dist::Type{SD}=StdNormal{T},
             meanspec::Union{MS, Type{MS}}=Intercept{T}(T[0]), algorithm=BFGS(),
             autodiff=:forward, kwargs...
             ) where {VS<:UnivariateVolatilitySpec, SD<:StandardizedDistribution,
                      MS<:MeanSpec, T<:AbstractFloat
                      }
	#can't use dispatch for this b/c meanspec is a kwarg
	meanspec isa Type ? ms = meanspec(zeros(T, nparams(meanspec))) : ms = deepcopy(meanspec)
    coefs = startingvals(VS, data)
    distcoefs = startingvals(SD, data)
    meancoefs = startingvals(ms, data)
	_fit!(coefs, distcoefs, meancoefs, VS, SD, ms, data; algorithm=algorithm, autodiff=autodiff, kwargs...)
	return UnivariateARCHModel(VS(coefs), data; dist=SD(distcoefs), meanspec=ms, fitted=true)
end

function fitsubset(::Type{VS}, data::Vector{T}, maxlags::Int, subset::Tuple; dist::Type{SD}=StdNormal{T},
             meanspec::Union{MS, Type{MS}}=Intercept{T}(T[0]), algorithm=BFGS(),
             autodiff=:forward, kwargs...
             ) where {VS<:UnivariateVolatilitySpec, SD<:StandardizedDistribution,
                      MS<:MeanSpec, T<:AbstractFloat
                      }
	#can't use dispatch for this b/c meanspec is a kwarg
	meanspec isa Type ? ms = meanspec(zeros(T, nparams(meanspec))) : ms = deepcopy(meanspec)
	VS_large = VS{ntuple(i->maxlags, length(subset))...}
	ng = nparams(VS_large)
	ns = nparams(SD)
	nm = nparams(typeof(ms))
	mask = subsetmask(VS_large, subset)
	garchcoefs = startingvals(VS_large, data, subset)
	distcoefs = startingvals(SD, data)
    meancoefs = startingvals(ms, data)

	obj = x -> -loglik(VS_large, SD, ms, data, x, mask)
    coefs = vcat(garchcoefs, distcoefs, meancoefs)
    res = optimize(obj, coefs, algorithm; autodiff=autodiff, kwargs...)
    coefs .= Optim.minimizer(res)
    garchcoefs .= coefs[1:ng]
	distcoefs .= coefs[ng+1:ng+ns]
    meancoefs .= coefs[ng+ns+1:ng+ns+nm]
	ms.coefs .= meancoefs
    return UnivariateSubsetARCHModel(VS_large(garchcoefs), data; dist=SD(distcoefs), meanspec=ms, fitted=true, subset=subset)
end

function fit!(am::UnivariateARCHModel; algorithm=BFGS(), autodiff=:forward, kwargs...)
    am.spec.coefs.=startingvals(typeof(am.spec), am.data)
    am.dist.coefs.=startingvals(typeof(am.dist), am.data)
    am.meanspec.coefs.=startingvals(am.meanspec, am.data)
	_fit!(am.spec.coefs, am.dist.coefs, am.meanspec.coefs, typeof(am.spec),
         typeof(am.dist), am.meanspec, am.data; algorithm=algorithm,
         autodiff=autodiff, kwargs...
         )
	am.fitted=true
    am
end

function fit(am::UnivariateARCHModel; algorithm=BFGS(), autodiff=:forward, kwargs...)
    am2=deepcopy(am)
    fit!(am2; algorithm=algorithm, autodiff=autodiff, kwargs...)
    return am2
end

function fit(vs::Type{VS}, lm::TableRegressionModel{<:LinearModel}; kwargs...) where VS<:UnivariateVolatilitySpec
	fit(vs, response(lm.model); meanspec=Regression(modelmatrix(lm.model); coefnames=coefnames(lm)), kwargs...)
end

function fit(vs::Type{VS}, lm::LinearModel; kwargs...) where VS<:UnivariateVolatilitySpec
	fit(vs, response(lm); meanspec=Regression(modelmatrix(lm)), kwargs...)
end

"""
    selectmodel(::Type{VS}, data; kwargs...) -> UnivariateARCHModel

Fit the volatility specification `VS` with varying lag lengths and return that which
minimizes the [BIC](https://en.wikipedia.org/wiki/Bayesian_information_criterion).

# Keyword arguments:
- `dist=StdNormal`: the error distribution.
- `meanspec=Intercept`: the mean specification, either as a type or instance of that type.
- `minlags=1`: minimum lag length to try in each parameter of `VS`.
- `maxlags=3`: maximum lag length to try in each parameter of `VS`.
- `criterion=bic`: function that takes a `UnivariateARCHModel` and returns the criterion to minimize.
- `show_trace=false`: print `criterion` to screen for each estimated model.
- `algorithm=BFGS(), autodiff=:forward, kwargs...`: passed on to the optimizer.

# Example
```jldoctest
julia> selectmodel(EGARCH, BG96)

EGARCH{1, 1, 2} model with Gaussian errors, T=1974.

Mean equation parameters:
───────────────────────────────────────────────
      Estimate   Std.Error    z value  Pr(>|z|)
───────────────────────────────────────────────
μ  -0.00900018  0.00943948  -0.953461    0.3404
───────────────────────────────────────────────

Volatility parameters:
──────────────────────────────────────────────
      Estimate  Std.Error    z value  Pr(>|z|)
──────────────────────────────────────────────
ω   -0.0544398  0.0592073  -0.919478    0.3578
γ₁  -0.0243368  0.0270414  -0.899985    0.3681
β₁   0.960301   0.0388183  24.7384      <1e-99
α₁   0.405788   0.067466    6.0147      <1e-08
α₂  -0.207357   0.114161   -1.81636     0.0693
──────────────────────────────────────────────
```
"""
function selectmodel(::Type{VS}, data::Vector{T};
                     dist::Type{SD}=StdNormal{T}, meanspec::Union{MS, Type{MS}}=Intercept{T},
                     maxlags::Integer=3, minlags::Integer=1, criterion=bic, show_trace=false, algorithm=BFGS(),
                     autodiff=:forward, kwargs...
                     ) where {VS<:UnivariateVolatilitySpec, T<:AbstractFloat,
                              SD<:StandardizedDistribution, MS<:MeanSpec
                              }
	@assert maxlags >= minlags >= 0

	#threading sometimes segfaults in tests locally. possibly https://github.com/JuliaLang/julia/issues/29934
	mylock=Threads.ReentrantLock()
    ndims = max(my_unwrap_unionall(VS)-1, 0) # e.g., two (p and q) for GARCH{p, q, T}
	ndims2 = max(my_unwrap_unionall(MS)-1, 0 )# e.g., two (p and q) for ARMA{p, q, T}
    res = Array{UnivariateSubsetARCHModel, ndims+ndims2}(undef, ntuple(i->maxlags - minlags + 1, ndims+ndims2))
    Threads.@threads for ind in collect(CartesianIndices(size(res)))
		tup = (ind.I[1:ndims] .+ minlags .-1)
		MSi = (ndims2==0 ? deepcopy(meanspec) : meanspec{ind.I[ndims+1:end] .+ minlags .- 1...})
		res[ind] = fitsubset(VS, data, maxlags, tup; dist=dist, meanspec=MSi,
                       algorithm=algorithm, autodiff=autodiff, kwargs...)
        if show_trace
            lock(mylock)
			VSi = VS{tup...}
            Core.print(modname(VSi))
			ndims2>0 && Core.print("-", modname(MSi))
			Core.println(" model has ",
                              uppercase(split("$criterion", ".")[end]), " ",
                              criterion(res[ind]), "."
                              )
            unlock(mylock)
        end
    end
    crits = criterion.(res)
    _, ind = findmin(crits)
	return fit(VS{res[ind].subset...}, data; dist=dist, meanspec=res[ind].meanspec, algorithm=algorithm, autodiff=autodiff, kwargs...)
end

function coeftable(am::UnivariateARCHModel)
    cc = coef(am)
    se = stderror(am)
    zz = cc ./ se
    CoefTable(hcat(cc, se, zz, 2.0 * normccdf.(abs.(zz))),
              ["Estimate", "Std.Error", "z value", "Pr(>|z|)"],
              coefnames(am), 4)
end

function show(io::IO, am::UnivariateARCHModel)
	if isfitted(am)
		cc = coef(am)
	    se = stderror(am)
	    ccg, ccd, ccm = splitcoefs(cc, typeof(am.spec),
	                               typeof(am.dist), am.meanspec
	                               )
	    seg, sed, sem = splitcoefs(se, typeof(am.spec),
	                               typeof(am.dist), am.meanspec
	                               )
	    zzg = ccg ./ seg
	    zzd = ccd ./ sed
	    zzm = ccm ./ sem
	    println(io, "\n", modname(typeof(am.spec)), " model with ",
	            distname(typeof(am.dist)), " errors, T=", nobs(am), ".\n")

	    length(sem) > 0 && println(io, "Mean equation parameters:", "\n",
	                               CoefTable(hcat(ccm, sem, zzm, 2.0 * normccdf.(abs.(zzm))),
	                                         ["Estimate", "Std.Error", "z value", "Pr(>|z|)"],
	                                         coefnames(am.meanspec), 4
	                                         )
	                              )
	    println(io, "\nVolatility parameters:", "\n",
	            CoefTable(hcat(ccg, seg, zzg, 2.0 * normccdf.(abs.(zzg))),
	                      ["Estimate", "Std.Error", "z value", "Pr(>|z|)"],
	                      coefnames(typeof(am.spec)), 4
	                      )
	            )
	    length(sed) > 0 && println(io, "\nDistribution parameters:", "\n",
	                               CoefTable(hcat(ccd, sed, zzd, 2.0 * normccdf.(abs.(zzd))),
	                                         ["Estimate", "Std.Error", "z value", "Pr(>|z|)"],
	                                         coefnames(typeof(am.dist)), 4
	                                         )
	                              )

   else
	   println(io, "\n", modname(typeof(am.spec)), " model with ",
			   distname(typeof(am.dist)), " errors, T=", nobs(am), ".\n\n")
	   length(am.meanspec.coefs) > 0 && println(io, CoefTable(am.meanspec.coefs, coefnames(am.meanspec), ["Mean equation parameters:"]))
	   println(io, CoefTable(am.spec.coefs, coefnames(typeof(am.spec)), ["Volatility parameters:   "]))
	   length(am.dist.coefs) > 0 && println(io, CoefTable(am.dist.coefs, coefnames(typeof(am.dist)), ["Distribution parameters: "]))
   end
end

function modname(::Type{S}) where S<:Union{UnivariateVolatilitySpec, MeanSpec}
    s = "$(S)"
	lastcomma = findlast(isequal(','), s)
    lastcomma == nothing || (s = s[1:lastcomma-1] * '}')
	firstdot = findfirst(isequal('.'), s)
	firstdot == nothing || (s = s[firstdot+1:end])
	s
end


text:
import JuLIP.Potentials: SitePotential
import ACE: ACEConfig

# TODO: nasty hack - must fix this!!!

cutoff(model::ACE.LinearACEModel) = cutoff(model.basis)

cutoff(basis::ACE.SymmetricBasis) = cutoff(basis.pibasis)

cutoff(basis::ACE.PIBasis) = cutoff(basis.basis1p)

cutoff(B1p::ACE.Product1pBasis) = cutoff(B1p.bases[2])

cutoff(Rn::ACE.Rn1pBasis) = cutoff(Rn.R)

# already defined in pair potentials
# cutoff(Rn::ACE.OrthPolys.TransformedPolys) = Rn.ru

"""
`struct ACESitePotential` : wraps an ACE model into a JuLIP calculator
"""
struct ACESitePotential{ENV, TM} <: SitePotential
   models::Dict{AtomicNumber, TM}
end

struct ACESitePotentialBasis{ENV, TM} <: JuLIP.MLIPs.IPBasis
   models::Dict{AtomicNumber, TM}  # model = basis
   inds::Dict{AtomicNumber, UnitRange{Int}}
end

ACESiteCalc{ENV} = Union{ACESitePotential{ENV}, ACESitePotentialBasis{ENV}}

Base.length(ipbasis::ACESitePotentialBasis) = sum(length, values(ipbasis.models))

cutoff(V::ACESiteCalc) =  maximum(cutoff, values(V.models))

function ACESitePotential(models::Dict{Symbol, TM},
                          ENV = ACEConfig) where {TM}
   mods = Dict([ AtomicNumber(sym) => mo for (sym, mo) in models]...)
   return ACESitePotential(mods, ENV)
end

ACESitePotential(models::Dict{AtomicNumber, TM}, ENV = ACEConfig) where {TM} =
     ACESitePotential{ENV, TM}(models)

# ----------------------------- FIO
import ACEbase: write_dict, read_dict
import Base: ==

write_dict(V::ACESitePotential) =
      Dict( "__id__" => "ACEatoms_ACESitePotential",
            "models" => Dict([ string(Int(z)) => write_dict(m)  for (z, m) in V.models ]...))

function read_dict(::Val{:ACEatoms_ACESitePotential}, D::Dict)
   models = Dict([ AtomicNumber(parse(Int, z)) => read_dict(m)
                   for (z, m) in D["models"] ]...)
   return ACESitePotential(models)
end

==(V1::ACESitePotential, V2::ACESitePotential) =
      (V1.models == V2.models)

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

function _get_basisinds(V::ACESitePotential)
   inds = Dict{AtomicNumber, UnitRange{Int}}()
   i0 = 0
   for (z, mo) in V.models
      len = length(mo.basis)
      inds[z] = (i0+1):(i0+len)   # to generalize for general models
      i0 += len
   end
   return inds
end

function basis(V::ACESitePotential{ENV}) where ENV
   models = Dict( [sym => model.basis for (sym, model) in V.models]... )
   inds = _get_basisinds(V)
   return ACESitePotentialBasis{ENV, Base.valtype(models)}(models, inds)
end

function environment(V::ACESiteCalc{<: ACEConfig},
                     at::Atoms, nlist, i::Integer) where {TM}
   Js, Rs, Zs = JuLIP.Potentials.neigsz(nlist, at, i)
   z0 = at.Z[i]
   return environment(V, Rs, Zs, z0), Js
end

atomstate(z::AtomicNumber) = AtomState{Float64}( (mu0 = z,) )
atomstate(z::AtomicNumber, z0::AtomicNumber, rr::SVector{3, T}) where {T} =
      AtomState{T}( (mu = z, mu0 = z0, rr = rr) )

environment(V::ACESiteCalc{<: ACEConfig},
            Rs::AbstractVector, Zs::AbstractVector, z0::AtomicNumber) =
      ACEConfig( atomstate.(Zs, z0, Rs) )

evaluate!(tmp, V::ACESitePotential, Rs, Zs, z0) =
      evaluate(V.models[z0], environment(V, Rs, Zs, z0)).val

# JuLIP.alloc_temp(V::ACESitePotentialBasis, N::Integer) =
#       (JuLIP.Potentials.alloc_temp_site(N)..., )

function evaluate!(B, tmp, V::ACESitePotentialBasis, Rs, Zs, z0)
   # fill!(B, 0)
   Bview = (@view B[V.inds[z0]])
   evaluate!(Bview, V.models[z0], environment(V, Rs, Zs, z0))
   return B
end

import ACEbase
function ACEbase.evaluate_d(V::ACESiteCalc, Rs::AbstractVector{JVec{T}}, Zs, z0) where {T}
   env = environment(V, Rs, Zs, z0)
   dV = zeros(JVec{T}, length(Rs))
   ACE.grad_config!(dV, V.models[z0], env)
   return dV
end

function evaluate_d!(dV, _tmpd, V::ACESitePotential, Rs, Zs, z0)
   env = environment(V, Rs, Zs, z0)
   return ACE.grad_config!(dV, V.models[z0], env)
end

function evaluate_d!(dB, _tmpd, V::ACESitePotentialBasis, Rs, Zs, z0)
   fill!(dB, zero(eltype(dB)))
   dBview = (@view dB[V.inds[z0], 1:length(Rs)])
   env = environment(V, Rs, Zs, z0)
   evaluate_d!(dBview, V.models[z0], env)
   return dBview
end

ACE.nparams(V::ACESitePotential) = sum(ACE.nparams, values(V.models))

function ACE.params(V::ACESitePotential)
   inds = _get_basisinds(V)
   c = zeros(ACE.nparams(V))
   for (z, mo) in V.models
      c[inds[z]] .= ACE.params(mo)
   end
   return c
end

function ACE.set_params!(V::ACESitePotential, c::AbstractVector)
   inds = _get_basisinds(V)
   for (z, mo) in V.models
      ACE.set_params!(V.models[z], c[inds[z]])
   end
   return V
end


text:
using FeedbackParticleFilters, Statistics, StatsBase

println("Testing UnweightedParticleEnsemble.jl:")

@testset "UnweightedParticleEnsemble.jl" begin

    print("  inner constructor for UnweightedParticleEnsemble")
    N    = 6
    d    = 2
    x_pf = [-2.422480820086937 -0.592332203303167 -2.017301296096984 -1.5151245392598531 0.02565906919199346  0.15161614796874012;
            -1.4986273707574955 0.3549802657063476 0.3625127047102   -0.166475807288935  0.05739520762738308 -0.18730971131161062]
    ens  = UnweightedParticleEnsemble(x_pf)
    print("-")
    @test ens.positions == x_pf
    println("DONE")

    print("  method dim")
    print("-")
    @test dim(ens) == d * N
    println("DONE")

    print("  method particle_dim")
    print("-")
    @test particle_dim(ens) == d
    println("DONE")

    print("  method no_of_particles")
    print("-")
    @test no_of_particles(ens) == N
    println("DONE")

    print("  method eff_no_of_particles")
    print("-")
    @test eff_no_of_particles(ens) == N
    println("DONE")

    print("  method get_pos")
    for i in 1:N
        print("-")
        @test get_pos(ens, i) == x_pf[:, i]
    end
    println("DONE")

    print("  method mean")
    print("-")
    @test size(mean(ens)) == (d, 1)
    print("-")
    @test mean(ens)[:,1] == [-1.0616606069310346; -0.17958745188568506]
    println("DONE")

    print("  method cov")
    print("-")
    @test size(cov(ens)) == (d, d)
    print("-")
    @test cov(ens) == [0.9742140049604432 0.29502867377601394; 0.29502867377601394 0.39598231024579167]
    println("DONE")

    print("  method var")
    print("-")
    @test size(var(ens)) == (d, 1)
    print("-")
    @test var(ens)[:,1] == [0.9742140049604432; 0.39598231024579167]
    println("DONE")

    print("  constructor for identical particles")
    vec  = [1., 2., 3., 4.]
    N    = 5
    ens  = UnweightedParticleEnsemble(vec, N)
    print("-")
    @test no_of_particles(ens) == N
    print("-")
    @test particle_dim(ens) == length(vec)
    for i in 1:N
        print("-")
        @test get_pos(ens, i) == vec
    end
    println("DONE")

end; #UnweightedParticleEnsemble.jl

max_stars_repo_name: simsurace/FeedbackParticleFilters

max_issues_repo_name: simsurace/FeedbackParticleFilters

max_forks_repo_name: simsurace/FeedbackParticleFilters


text:
cd(@__DIR__)
using Pkg
Pkg.activate(".")
##
using StochasticPrograms
using GLPK
using Pipe
using Random
using Plots
using JuMP
using DataFrames
using CSV
using Cbc

include("samplers.jl")
include("plot_utils.jl")
##
function plot_results_r(sp, pv, w, d; s=1, stage_1=[:gci, :gco], stage_2=[:gci2, :gco2], debug = false)
    plt_sto = plot()
    plt_invest = plot()
    plt = plot() # create separate plot for storage state of charge
    plot!(plt_invest, pv .* value(sp[1, :u_pv]), label="pv")
    plot!(plt_invest, w .* value(sp[1, :u_wind]), label="wind")
    plot!(plt_invest, d, label="demand")
    stor_flow_i = value.(sp[1, :sto_in]).data
    stor_flow_o = value.(sp[1, :sto_out]).data
    t_xi = scenarios(sp)[s].t_xi
    stor_charge = [-sum(stor_flow_i[1:t])+sum(stor_flow_o[1:t]) for t in value.(sp[1, :sto_in]).axes[1]] .+ 0.5*value(sp[1, :u_storage])
    plot!(plt_sto, value.(sp[1, :sto_in]).axes, stor_charge, label="global storage charge")
    if debug
        #print("Maximum value of storage flow = "*string(maximum(value.(sp[1, :sto_in]).data))*"\n")
    end
    for var in stage_1
        plot!(plt, value.(sp[1, var]).axes, value.(sp[1, var]).data, label=string(var))
        if debug
            print("Maximum value of "*string(var)*" = "*string(maximum(value.(sp[1, var]).data))*"\n")
        end
    end
    #=for var in stage_2
        if debug
            print("Maximum value of "*string(var)*" = "*string(maximum(value.(sp[2, var], s).data))*"\n")
        end
        plot!(plt, value.(sp[2, var], s).axes, value.(sp[2, var], s).data, label=string(var)*string(s), linestyle=:dash, linewidth=2)
    end=#
    for var in [:gci2, :gco2]
        plot!(plt, (t_xi+1):(t_xi+3), value.(sp[2, var], s), label=string(var)*string(s), linestyle=:dash, linewidth=2)
    end
    stor_flow_i2 = value.(sp[2, :sto_in2],s)
    stor_flow_o2 = value.(sp[2, :sto_out2],s)
    stor_charge2 = [-sum(stor_flow_i2[1:t])+sum(stor_flow_o2[1:t]) for t in 1:(recovery_time+1)] .+ stor_charge[t_xi-1]
    plot!(plt_sto, (t_xi-1):(t_xi+recovery_time), vcat(stor_charge[t_xi-1],stor_charge2), label=string("stochastic storage charge")*string(s), linestyle=:dash, linewidth=2)
    if debug
        #print("Maximum value of sto2 = "*string(maximum(value.(sp[2, :sto2], s).data))*"\n")
    end
    display(plot(plt_invest, plt, plt_sto, layout = (3,1)))
    #print(s)
end

##
# Define the time interval
t_s = 1
t_f = 365*24÷4
offset = 0
timesteps = t_s:t_f

F_range = (100., 200.)
# Define weather and demand data
pv = CSV.read("../basic_example.csv", DataFrame)[timesteps.+offset, 3]
wind = CSV.read("../basic_example.csv", DataFrame)[timesteps.+offset, 4]
demand = CSV.read("../basic_example.csv", DataFrame)[timesteps.+offset, 2] ./80

##
# Note on signs: s_xi<0 means that energy is requested, s_xi>0 means an additional consumption
energy_model_r = @stochastic_model begin
    @stage 1 begin
        @parameters begin
            # Expected lifetime of components, years
            asset_lifetime = 10.
            # Costs in Euro/1kWp
            c_pv = 100.
            c_wind = 1000.
            c_storage = 100.
            c_sto_op = 0.00001
            c_i = .03
            c_o = .01
            # Euro
            inv_budget = 500000000.
            # Shiftable demand, kW
            flexible_demand = 0.
        end
        time_scale = asset_lifetime *365*24/length(timesteps)
        # Component units to be invested in, kWp
        @decision(model, u_pv >= 0)
        @decision(model, u_wind >= 0)
        @decision(model, u_storage >= 0)
        @constraint(model, c_pv*u_pv+c_wind*u_wind+c_storage*u_storage<=inv_budget)
        # Grid connection
        @decision(model, gci[t in timesteps] >= 0)
        @decision(model, gco[t in timesteps] >= 0)
        # Flexible demand
        @decision(model, fl_dem[t in timesteps] >= 0)
        @constraint(model, sum(fl_dem) == flexible_demand)
        # Storage model
        @decision(model, sto_in[t in timesteps] >= 0) #into the bus from storage
        @decision(model, sto_out[t in timesteps] >= 0)
        @constraint(model, [t in timesteps], -0.5*u_storage <= -sum(sto_in[t_s:t])+sum(sto_out[t_s:t]))
        @constraint(model, [t in timesteps], -sum(sto_in[t_s:t])+sum(sto_out[t_s:t]) <= 0.5*u_storage)
        @constraint(model, -sum(sto_in)+sum(sto_out) == 0)
        # Energy balance
        @constraint(model, [t in timesteps],
        gci[t]-gco[t]+u_pv*pv[t]+u_wind*wind[t]-demand[t]-fl_dem[t]+sto_in[t]-sto_out[t]==0)
        # Investment costs
        @objective(model, Min, (u_pv*c_pv+u_wind*c_wind+u_storage*c_storage)/time_scale + c_i*sum(gci) - c_o*sum(gco)+
        c_sto_op*sum(sto_in)+c_sto_op*sum(sto_out))
    end
    @stage 2 begin
        @parameters begin
            c_flex = .5
            #F = 10.
            recovery_time = 3
            c_sto_op = 0.00001
            c_i = .03
            c_o = .01
        end
        @uncertain t_xi s_xi F_xi from F_scenario #t_xi the time of flexibility demand, s_xi - sign (±1 or 0)
        @known(model, u_pv)
        @known(model, u_wind)
        @known(model, u_storage)
        @known(model, gci)
        @known(model, gco)
        @known(model, sto_in)
        @known(model, sto_out)
        @known(model, fl_dem)
        # Post event components
        @recourse(model, gci2[t in 1:recovery_time] >= 0)
        @recourse(model, gco2[t in 1:recovery_time] >= 0)
        @recourse(model, sto_in2[t in 1:(recovery_time+1)] >= 0)
        @recourse(model, sto_out2[t in 1:(recovery_time+1)] >= 0)

        # Storage
        @constraint(model, [t in 1:(recovery_time+1)], 0.5*u_storage - sum(sto_in[t_s:(t_xi-1)])+sum(sto_out[t_s:(t_xi-1)]) <= -sum(sto_in2[t_s:t])+sum(sto_out2[t_s:t]))
        @constraint(model, [t in 1:(recovery_time+1)], -sum(sto_in2[t_s:t])+sum(sto_out2[t_s:t]) <= u_storage - sum(sto_in[t_s:(t_xi-1)])+sum(sto_out[t_s:(t_xi-1)]))

        @constraint(model, -sum(sto_in2)+sum(sto_out2) == -sum(sto_in[t_xi:(t_xi+recovery_time)])+sum(sto_out[t_xi:(t_xi+recovery_time)]))

        #Flexible demand
        @recourse(model, fl_dem2[t in timesteps]>=0)
        @constraint(model, [t in t_s:(t_xi-1)], fl_dem[t] == fl_dem2[t])
        @constraint(model, sum(fl_dem2) == sum(fl_dem))

        # Event energy balance
        # The storage and other fast acting components use the recourse variables here.
        # They provide the balance. Grid connection is not allowed, as we are suporting the grid here.
        @constraint(model, gci[t_xi]-gco[t_xi]+u_pv*pv[t_xi]+u_wind*wind[t_xi]-demand[t_xi]-fl_dem2[t_xi] +sto_in2[1]-sto_out2[1]+F_xi*s_xi==0)
        # Post event energy balance
        @constraint(model, [t in (t_xi+1):(t_xi+recovery_time)],
        gci2[t-t_xi]-gco2[t-t_xi]+u_pv*pv[t]+u_wind*wind[t]-demand[t]-fl_dem2[t]+sto_in2[t-t_xi+1]-sto_out2[t-t_xi+1]==0)
        @objective(model, Min, c_i*(sum(gci2)-sum(gci[t_xi:(t_xi+recovery_time)]))-c_o*(sum(gco2)-sum(gco[t_xi:(t_xi+recovery_time)]))+
            c_sto_op*(sum(sto_in2)+sum(sto_out2)- sum(sto_in[(t_xi-1):(t_xi+recovery_time)])- sum(sto_out[(t_xi-1):(t_xi+recovery_time)])))
    end
end
##
recovery_time = 3
f_samp = F_sampler_r(timesteps, F_range, recovery_time)
@timed sp0 = instantiate(energy_model_r, [F_scenario(5, 1, 10.)], c_pv = 700., c_wind = 2000., c_i = 0.3, c_o = 0.05, c_storage = 600., flexible_demand = 0., recovery_time = recovery_time, optimizer = LShaped.Optimizer)
set_optimizer(sp0, LShaped.Optimizer)
set_optimizer_attribute(sp0, MasterOptimizer(), GLPK.Optimizer)
set_optimizer_attribute(sp0, SubProblemOptimizer(), GLPK.Optimizer)
##
@timed optimize!(sp0)

od0 = optimal_decision(sp0)
objective_value(sp0)

plot_results_r(sp0, pv, wind, demand, s = 1, stage_1 = [:gci, :gco, :fl_dem], stage_2 = [:gci2, :gco2, :fl_dem2],)

##
@timed sp0 = instantiate(energy_model_r, f_samp, 100, c_pv = 700., c_wind = 2000., c_i = 0.3, c_o = 0.05, c_storage = 600., flexible_demand = 0., recovery_time = recovery_time, optimizer = GLPK.Optimizer)

@timed optimize!(sp0)

od = optimal_decision(sp)
objective_value(sp)

##
#sp = instantiate(energy_model, [simple_scenario(1,0)], F=100., optimizer = Cbc.Optimizer)

##
for s in 1:5
    plot_results_r(sp, pv, wind, demand, s = s, stage_1 = [:gci, :gco, :fl_dem], stage_2 = [:gci2, :gco2, :fl_dem2],)
end
##
# Main result
println("Termination status: $(termination_status(sp))")
println("Objective value: $(objective_value(sp))")
println("Optimal decision: $(optimal_decision(sp))")

##
# Investment result
u_pv = sp[1, :u_pv]
u_wind = sp[1, :u_wind]
u_storage = sp[1, :u_storage]
# First stage
println("value(u_pv) = $(value(u_pv))")
println("value(u_wind) = $(value(u_wind))")
println("value(u_storage) = $(value(u_storage))")

# Second stage
for s in 1:3
    #println("Objective value in scenario $s: $(objective_value(sp, s))")
    println("Optimal recourse in scenario $s: $(optimal_recourse_decision(sp, s))")
    plot_results(sp, pv, wind, demand, s = s)
end

##

scen = scenarios(sp)
print(scen)
##
function test_decision(p, decision, timesteps)

    infeasible_count = 0
    costs = []

    for sign in [-1,1]
        for t in timesteps
            scenario = F_scenario(t,sign,150.)
            try evaluate_decision(p, decision, scenario);
            catch e;
                infeasible_count += 1;
                append!(costs, Inf)
                #append!(scenario_results, [Inf, scenario.t_xi, scenario.s_xi, scen.F_xi]);
            end
            append!(costs, evaluate_decision(p, decision, scenario))
            #append!(scenario_results, [evaluate_decision(p, decision, scenario), scenario.t_xi, scenario.s_xi, scenario.F_xi]);
        end
    end
    N = length(timesteps)
    print(infeasible_count/N/2.)

    scenario_results = hcat(costs, vcat(timesteps, timesteps), -1*vcat(ones(N), ones(N)))
    return scenario_results
end
##

scen_results = test_decision(sp, od, timesteps)

scen_results0 = test_decision(sp0, od0, timesteps)

#av_cost = sum(scen_results[i, 1] for )

print(scen_bad0)

EVPI(sp)

VRP(sp)

VSS(sp)

set_optimizer(sp, SAA.Optimizer)

evaluate_decision(sp, od, simple_scenario(1,1))

t_xi = scenarios(sp)[3].t_xi

max_stars_repo_path: src/prototype recovery_time.jl

max_issues_repo_path: src/prototype recovery_time.jl

max_forks_repo_path: src/prototype recovery_time.jl


text: include("model/RBC.jl")


text:
import Pkg
# Pkg.add("Conda")
# Pkg.add("PyPlot")
# Pkg.add("ProgressMeter")
# Pkg.add("Gurobi")
# Pkg.add("DataFrames")
# Pkg.add("JuMP")
# Pkg.add("Combinatorics")
# Pkg.add("LightGraphs")
# Pkg.add("DelimitedFiles")
# Pkg.add("MathProgBase")
# Pkg.add("GraphPlot")
# Pkg.add("Compose")
# Pkg.add("Cairo")
# Pkg.add("Fontconfig")
# Pkg.add("Distributions")
# Pkg.add("HTTP")
# Pkg.add("Sockets")
# Pkg.add("JSON2")
# Pkg.add("JSON")
# using Conda; Conda.add("pyqt")

using HTTP, Sockets, JSON2

include("./util.jl")
include("./parsers.jl")
include("./Algorithms/TEAVAR.jl")

mutable struct Request
    topology::String
    demand::String
    path::String
    beta::String
    cutoff::String
    k::String
    downscale_demand::String
end

mutable struct Response
    var::Float64
    cvar::Float64
    allocation::Array
    num_nodes::Int64
    capacity::Array
    failure_probabilities::Array
    flows::Array
    demand::Array
    T::Array
    Tf::Array
    links::Array
    scenarios::Array
    probabilities::Array
    X::Array
end

ROUTER = HTTP.Router()

function JSONHandler(req::HTTP.Request)
    if req.method == "OPTIONS"
        res = HTTP.Response(200)
        res.headers = [Pair("Access-Control-Allow-Origin", "*"),
                       Pair("Vary", "Origin"),
                       Pair("Vary", "Access-Control-Request-Method"),
                       Pair("Vary", "Access-Control-Request-Headers"),
                       Pair("Access-Control-Allow-Headers", "Content-Type, Origin, Accept, token"),
                       Pair("Access-Control-Allow-Methods", "GET, POST,OPTIONS")]
        return res
        response_body = HTTP.handle(ROUTER, req)
    else
        response_body = HTTP.handle(ROUTER, req)
    end
    res = HTTP.Response(200, JSON2.write(response_body))
    res.headers = [Pair("Access-Control-Allow-Origin", "*")]
    return res
end

function teavar(req::HTTP.Request)
    json = JSON2.read(IOBuffer(HTTP.payload(req)), Request)
    println(json)

    topology = json.topology
    demand_num = parse(Int64, json.demand)
    beta = parse(Float64, json.beta)
    cutoff = parse(Float64, json.cutoff)
    k = parse(Int64, json.k)
    downscale_demand = parse(Int64, json.downscale_demand)
    links, capacity, link_probs, nodes = readTopology(topology, downscale=1)

    if topology == "B4"
        weibull_probs = [0.00100046, 0.000419485, 0.000612366, 0.00237276, 0.00612306, 0.000365313, 0.00247675, 0.00166383, 0.000588749, 0.00174909, 0.000456167, 0.000834483, 0.00142059, 0.00243267, 0.00488089, 0.000947656, 4.65804e-5, 0.00280921, 0.00112321, 0.000168123, 0.000485535, 0.000119134, 0.000138946, 9.39895e-5, 0.000366816, 0.000420305, 0.000272837, 0.00109307, 0.00200512, 0.000152399, 0.000880158, 0.00019616, 0.00175013, 0.00140933, 0.00150634, 0.000185742, 0.000741036, 0.00261969]
    else
        weibull_probs = weibullProbs(length(links), shape=.8, scale=.001)
    end
    scenarios, scenario_probs = subScenarios(weibull_probs, cutoff, first=true, last=false)
    # w_scenarios, w_probs = subScenarios(weibull_probs, cutoff, first=true, last=false)
    demand, flows = readDemand("$(topology)/demand", length(nodes), demand_num, scale=1.0, downscale=downscale_demand)

    T, Tf = [], []
    try
        T, Tf, k = parsePaths("$(topology)/paths/$(json.path)", links, flows)
    catch
        if json.path == "ED"
            T, Tf, k, g = getTunnels(nodes, links, capacity, flows, 30, edge_disjoint=true)
        else
            T, Tf, k, g = getTunnels(nodes, links, capacity, flows, k)
        end
    end
    # a = parseYatesSplittingRatios("$(topology)/paths/$(algorithm)", k, flows, zeroindex=zeroindex)
    # a = parseYatesAllocations("$(topology)/paths/$(algorithm)", k, demand, flows, zeroindex=zeroindex)

    env = Gurobi.Env()
    var, cvar, a = TEAVAR(env, links, capacity, flows, demand, beta, k, T, Tf, scenarios, scenario_probs)
    println(var)
    println(cvar)
    println(a)

    nscenarios = length(scenarios)
    ntunnels = length(T)
    X  = ones(nscenarios,ntunnels)
    for s in 1:nscenarios
        for t in 1:ntunnels
            if size(T[t],1) == 0
                X[s,t] = 0
            else
                for e in 1:length(links)
                    if scenarios[s][e] == 0
                        back_edge = findfirst(x -> x == (links[e][2],links[e][1]), links)
                        if in(e, T[t]) || in(back_edge, T[t])
                            X[s,t] = 0
                        end
                    end
                end
            end
        end
    end
    http_res = Response(var, cvar, [a[i, :] for i in 1:size(a, 1)], length(nodes), capacity, weibull_probs, flows, demand, T, Tf, links, scenarios, scenario_probs, [X[i, :] for i in 1:size(X, 1)])
    return http_res
end

function main()
    println("Listening on localhost:8080....")
    HTTP.@register(ROUTER, "POST", "/api/teavar", teavar)
    # HTTP.serve(JSONHandler, Sockets.localhost, 8080)
    HTTP.serve(JSONHandler, "128.30.92.156", 8080)
end

main()


text:
@testset "Base Extensions" begin

bins = 0:0.1:1.0
@test nextpow2(bins) == 0:0.1:1.6

@test allequal([false])   # assume equality for length 1 arrays
@test allequal([0.00, 0, 0im])# check for equivilancy across types

end