text:
"""
    Matern Kernel
"""
struct MaternKernel{T<:Real,KT<:KernelType} <: Kernel{T,KT}
    fields::KernelFields{T,KT}
    ν::T
    function MaternKernel{T,KT}(θ::Vector{T},ν::T,variance::T=one(T);dim::Integer=0) where {T<:Real,KT<:KernelType}
        @assert ν > 0 "ν should be bigger than 0!"
        if KT == ARDKernel
            if length(θ)==1 && dim ==0
                error("You defined an ARD Matern kernel without precising the number of dimensions or giving a vector for the lengthscale                   Please set dim in your kernel initialization")
            elseif dim!=0 && (length(θ)!=dim && length(θ)!=1)
                @warn "You did not use the same dimension for your params and dim, using the first value of params for all dimensions"
                θ = ones(dim,T=T)*θ[1]
            elseif length(θ)==1 && dim!=0
                θ = ones(dim)*θ[1]
            end
            dim = length(θ)
            return new{T,ARDKernel}(KernelFields{T,ARDKernel}(
                                        "General Matern",
                                        HyperParameter{T}(variance,interval(OpenBound{T}(zero(T)),NullBound{T}()),fixed=false),
                                        HyperParameters{T}(θ,[interval(OpenBound(zero(T)),NullBound{T}()) for _ in 1:dim]),dim,
                                        WeightedEuclidean(one(T)./(θ.^2))),ν)
        else
            return new{T,IsoKernel}(KernelFields{T,IsoKernel}(
                                        "General Matern",
                                        HyperParameter{T}(variance,interval(OpenBound{T}(zero(T)),NullBound{T}()),fixed=false),
                                        HyperParameters{T}(θ,[interval(OpenBound{T}(zero(T)),NullBound{T}())]),1,
                                        Euclidean()),ν)
        end
    end
end

#TODO Implement specific kernels for matern kernel
function MaternKernel(θ::T1=1.0,ν::T2=2.5;variance::T3=one(T1),dim::Integer=0,ARD::Bool=false) where {T1<:Real,T2<:Real,T3<:Real}
    maxtype = floattype(T1,T2,T3)
    if ARD
        # if ν≈0.5
            # ExponentialKernel{floattype(T1,T2),ARDKernel}([θ],ν,variance=variance,dim=dim)
        # elseif ν≈1.5
            # Matern3_2Kernel{floattype(T1,T2),ARDKernel}([θ],ν,variance=variance,dim=dim)
        # elseif ν≈2.5
            # Matern5_2Kernel{floattype(T1,T2),ARDKernel}([θ],ν,variance=variance,dim=dim)
        # else
            MaternKernel{maxtype,ARDKernel}([θ],maxtype(ν),maxtype(variance),dim=dim)
        # end
    else
        # if ν≈0.5
            # ExponentialKernel{floattype(T1,T2),IsoKernel}([θ],ν,variance=variance)
            # elseif ν≈1.5
            # Matern3_2Kernel{floattype(T1,T2),IsoKernel}([θ],ν,variance=variance)
        # elseif ν≈2.5
            # Matern5_2Kernel{floattype(T1,T2),ARDKernel}([θ],ν,variance=variance)
        # else
            MaternKernel{maxtype,IsoKernel}([θ],maxtype(ν),maxtype(variance))
        # end
    end
 end

function MaternKernel(θ::Array{T1,1},ν::T2=2.5;variance::T3=one(T1),dim::Integer=0) where {T1<:Real,T2<:Real,T3<:Real}
    maxtype = floattype(T1,T2,T3)
    MaternKernel{maxtype,ARDKernel}(θ,maxtype(ν),maxtype(variance),dim=dim)
end

@inline function maternkernel(z::T, l::T, ν::T) where {T<:Real}
    v = sqrt(2.0*ν)*z/l
    v = v < eps(T) ? eps(T) : v
    return 2.0 * (v/2.0)^ν * besselk(ν,v) / gamma(ν)
end

@inline function maternkernel(z::T, ν::T) where {T<:Real}
    v = sqrt(2.0*ν)*z
    v = v < eps(T) ? eps(T) : v
    return 2.0 * (v/2.0)^ν * besselk(ν,v) / gamma(ν)
end

function kappa(k::MaternKernel{T,IsoKernel}) where {T<:Real,KT}
    return z->maternkernel(z,getlengthscales(k),k.ν)
end

function kappa(k::MaternKernel{T,ARDKernel}) where {T<:Real,KT}
    return z->maternkernel(z,k.ν)
end

function updateweights!(k::MaternKernel{T,KT},w::Vector{T}) where {T,KT}
    k.fields.metric.weights .= 1.0./(w.^2)
end

function computeIndPointsJmm(k::MaternKernel{T,KT},X::Matrix{T},iPoint::Integer,K::Symmetric{T,Matrix{T}}) where {T,KT}
    l = (getlengthscales(k)); v = getvariance(k); ν = k.ν; C = 2^(1.0-ν)/gamma(ν)
    P = sqrt(2.0*ν).*pairwise(getmetric(k),X,dims=1)
    P .= ifelse.(P.<eps(T),eps(T),P)
    if KT == IsoKernel
        return -(2.0*C*v*ν) .* (X[iPoint,:]'.-X)./(l^2) .* (P[iPoint,:]./l).^(ν-1.0) .* besselk.(ν-1.0,P[iPoint,:]./l)
    elseif KT == ARDKernel
        return -(2.0*C*v*ν) .* (X[iPoint,:]'.-X)./(l.^2)' .* (P[iPoint,:]).^(ν-1.0) .* besselk.(ν-1.0,P[iPoint,:])
    end
end

function computeIndPointsJnm(k::MaternKernel{T,KT},X::Matrix{T},x::Vector{T},iPoint::Integer,K::Matrix{T}) where {T,KT}
    l = (getlengthscales(k)); v = getvariance(k); ν = k.ν; C = 2^(1.0-ν)/gamma(ν)
    P = sqrt(2.0*ν).*pairwise(getmetric(k),x[:,:],X,dims=1)[:]
    P .= ifelse.(P.<eps(T),eps(T),P)
    if KT == IsoKernel
        return -(2.0*C*v*ν).* (x'.-X)./(l^2) .* (P./l).^(ν-1.0) .* besselk.(ν-1.0,P./l)
    else
        return -(2.0*C*v*ν).* (x'.-X)./(l.^2)' .* (P).^(ν-1.0) .* besselk.(ν-1.0,P)
    end
end

################# Matrix derivatives for the Matern kernel###################################
"""Return the derivatives of Knn for the Iso MaternKernel"""
function kernelderivativematrix(X::Array{T},kernel::MaternKernel{T,IsoKernel}) where {T<:Real}
    v = getvariance(kernel); l = getlengthscales(kernel); ν=kernel.ν; C = (2^(1.0-ν))/gamma(ν)
    P = (sqrt(2.0*ν)/l).*pairwise(Euclidean(),X,dims=1);
    P .= ifelse.(P.<eps(T),eps(T),P)
    return Symmetric(lmul!(C*v/l, P.^(ν+1.0) .* besselk.(ν-1.0,P)))
end

"""Return the derivatives of Knn for the Iso MaternKernel with Knn precomputed"""
function kernelderivativematrix_K(X::Array{T,N},K::Symmetric{T,Array{T,N}},kernel::MaternKernel{T,IsoKernel}) where {T,N}
    v = getvariance(kernel); l = getlengthscales(kernel); ν=kernel.ν;  C = (2^(1.0-ν))/gamma(ν)
    P = (sqrt(2.0*ν)/l).*pairwise(Euclidean(),X,dims=1);
    P .= ifelse.(P.<eps(T),eps(T),P)
    return Symmetric(lmul!(C*v/l,P.^(ν+1.0) .* besselk.(ν-1.0,P)))
end

"Return the derivatives of Knn for the ARD MaternKernel"
function kernelderivativematrix(X::Array{T,N},kernel::MaternKernel{T,ARDKernel}) where {T,N}
    v = getvariance(kernel); ls = getlengthscales(kernel); ν = kernel.ν; C = (2^(1.0-ν))/gamma(ν)
    P = (sqrt(2.0*ν)).*pairwise(getmetric(kernel),X,dims=1) #d/ρ
    P .= ifelse.(P.<eps(T),eps(T),P)
    Pi = [pairwise(SqEuclidean(),X[:,i]',dims=2) for i in 1:length(ls)] # (x_i-x_i')
    return Symmetric.(map((pi,l)->lmul!(2.0*C*v*ν/(l^3),pi .* P.^(ν-1.0) .* besselk.(ν-1.0,P)),Pi,ls))
end

"""Return the derivatives of Knn for the ARD MaternKernel with Knn precomputed"""
function kernelderivativematrix_K(X::Array{T,N},K::Symmetric{T,Array{T,N}},kernel::MaternKernel{T,ARDKernel}) where {T,N}
    v = getvariance(kernel); ls = getlengthscales(kernel); ν = kernel.ν; C = (2^(1.0-ν))/gamma(ν)
    P = (sqrt(2.0*ν)).*pairwise(getmetric(kernel),X,dims=1) #d/ρ
    P .= ifelse.(P.<eps(T),eps(T),P)
    Pi = [pairwise(SqEuclidean(),X[:,i]',dims=2) for i in 1:length(ls)] # (x_i-x_i')
    return Symmetric.(map((pi,l)->lmul!(2.0*C*v*ν/(l^3),pi .* P.^(ν-1.0) .* besselk.(ν-1.0,P)),Pi,ls))
end

########## DERIVATIVE MATRICES FOR TWO MATRICES #######
"""Return the derivatives of Knm for the Iso MaternKernel"""
function kernelderivativematrix(X::Array{T,N},Y::Array{T,N},kernel::MaternKernel{T,IsoKernel}) where {T,N}
    v = getvariance(kernel); l = getlengthscales(kernel); ν=kernel.ν;  C = 2^(1.0-ν)/gamma(ν)
    P = (sqrt(2.0*ν)/l)*pairwise(Euclidean(),X,Y,dims=1);
    P .= ifelse.(P.<eps(T),eps(T),P)
    return lmul!(C*v/l,P.^(ν+1.0).*besselk.(ν-1.0,P))
end

"""Return the derivatives of Knn for the Iso MaternKernel with Knm precomputed"""
function kernelderivativematrix_K(X::Array{T,N},Y::Array{T,N},K::Array{T,N},kernel::MaternKernel{T,IsoKernel}) where {T,N}
    v = getvariance(kernel); l = getlengthscales(kernel); ν=kernel.ν; C = 2^(1.0-ν)/gamma(ν)
    P = (sqrt(2.0*ν)/l)*pairwise(Euclidean(),X,Y,dims=1);
    P .= ifelse.(P.<eps(T),eps(T),P)
    return lmul!(C*v/l,P.^(ν+1.0).*besselk.(ν-1.0,P))
end

"""Return the derivatives of Knm for the ARD MaternKernel"""
function kernelderivativematrix(X::Array{T,N},Y::Array{T,N},kernel::MaternKernel{T,ARDKernel}) where {T,N}
    v = getvariance(kernel); ls = getlengthscales(kernel); ν = kernel.ν; C = 2^(1.0-ν)/gamma(ν)
    P = (sqrt(2.0*ν)).*pairwise(getmetric(kernel),X,Y,dims=1)
    P .= ifelse.(P.<eps(T),eps(T),P)
    Pi = [pairwise(SqEuclidean(),X[:,i]',Y[:,i]',dims=2) for i in 1:length(ls)] # (x_i-x_i')
    return map((pi,l)->lmul!(2.0*C*v*ν/(l^3),pi .* P.^(ν-1.0).*besselk.(ν-1.0,P)),Pi,ls)
end

"""Return the derivatives of Knn for the ARD MaternKernel with Knm precomputed"""
function kernelderivativematrix_K(X::Array{T,N},Y::Array{T,N},K::Array{T,N},kernel::MaternKernel{T,ARDKernel}) where {T,N}
    v = getvariance(kernel); ls = getlengthscales(kernel); ν = kernel.ν; C = 2^(1.0-ν)/gamma(ν)
    P = (sqrt(2.0*ν)).*pairwise(getmetric(kernel),X,Y,dims=1)
    P .= ifelse.(P.<eps(T),eps(T),P)
    Pi = [pairwise(SqEuclidean(),X[:,i]',Y[:,i]',dims=2) for i in 1:length(ls)] # (x_i-x_i')
    return map((pi,l)->lmul!(2.0*C*v*ν/(l^3),pi .* P.^(ν-1.0).*besselk.(ν-1.0,P)),Pi,ls)
end

############ DIAGONAL DERIVATIVES ###################

"""Return the derivate of the diagonal covariance matrix for the Iso MaternKernel"""
function kernelderivativediagmatrix(X::Array{T,N},kernel::MaternKernel{T,IsoKernel}) where {T,N}
    n = size(X,1); P = zeros(T,n)
    return P
end

"""Return the derivate of the diagonal covariance matrix for the Iso MaternKernel"""
function kernelderivativediagmatrix(X::Array{T,N},kernel::MaternKernel{T,ARDKernel}) where {T,N}
    n = size(X,1); P = zeros(T,n);
    return [P for _ in 1:kernel.fields.Ndim]
end


text:
using LinearAlgebra, StaticArrays, StableRNGs, Vecchia

# Data size:
const sz = 724

# Covariance function (simple SVector format for pts):
kfn(x,y,p) = p[1]*exp(-norm(x-y)/p[2])*(1.0+norm(x-y)/p[2])

# Covariance function (scalar format for pts to enable SIMD assembly):
function kfn_scalar(x1, x2, y1, y2, p)
  nrm = sqrt((x1-y1)^2 + (x2-y2)^2)
  p[1]*exp(-nrm/p[2])*(1+nrm/p[2])
end

# Choose some random locations to make measurements.
const rng = StableRNG(12345)
const pts = [SVector{2, Float64}(randn(rng, 2)) for _ in 1:sz]
const pts_1d  = range(0.0, 1.0, length=32)
const reg_pts = vec(map(x->SVector{2,Float64}(x...),
                        Iterators.product(pts_1d, pts_1d)))

# Pick true parameters, get the true covariance matrix, and simulate with
# Cholesky factor.
tru_parm = [3.0, 0.5]
truK = Symmetric([kfn(x,y,tru_parm) for x in pts, y in pts])
const sim = cholesky(truK).L*randn(StableRNG(123), sz)

# Create a VecchiaConfig object. You are very much welcome to create your own
# struct that is <:Vecchia.VecchiaConfig to do something more thoughtful for
# your specific problem. The kdtree method included here is very general-use.
# This object uses "chunks" of size 64, and the 3 nearest "chunk" neighbors.
const vecc     = Vecchia.kdtreeconfig(sim, pts, 64, 3, kfn)
const nys_vecc = Vecchia.nystrom_kdtreeconfig(sim, pts, 64, 3, kfn, 72)

# A "Scalarized" VecchiaConfig object, which now means that the covariance
# matrices will be assembled using SIMD, which can provide a serious speedup
# that is completely independent of the speedup gained by threading.
const vecc_s   = Vecchia.scalarize(vecc, kfn_scalar)

ext: jl


text:
using Distributions
using Random
using StaticArrays
using Test
using DocStringExtensions
using LinearAlgebra

using CLIMA
using CLIMA.Atmos
using CLIMA.DGmethods.NumericalFluxes
using CLIMA.GenericCallbacks
using CLIMA.ODESolvers
using CLIMA.Mesh.Filters
using CLIMA.MoistThermodynamics
using CLIMA.PlanetParameters
using CLIMA.VariableTemplates

import CLIMA.DGmethods: vars_state, vars_aux,
                        vars_integrals, vars_reverse_integrals,
                        indefinite_stack_integral!,
                        reverse_indefinite_stack_integral!,
                        integral_load_aux!, integral_set_aux!,
                        reverse_integral_load_aux!,
                        reverse_integral_set_aux!

import CLIMA.DGmethods: boundary_state!
import CLIMA.Atmos: atmos_boundary_state!, atmos_boundary_flux_diffusive!, flux_diffusive!
import CLIMA.DGmethods.NumericalFluxes: boundary_flux_diffusive!

# -------------------- Radiation Model -------------------------- #
vars_state(::RadiationModel, FT) = @vars()
vars_aux(::RadiationModel, FT) = @vars()
vars_integrals(::RadiationModel, FT) = @vars()
vars_reverse_integrals(::RadiationModel, FT) = @vars()

function atmos_nodal_update_aux!(::RadiationModel, ::AtmosModel, state::Vars, aux::Vars, t::Real) end
function preodefun!(::RadiationModel, aux::Vars, state::Vars, t::Real) end
function integral_load_aux!(::RadiationModel, integ::Vars, state::Vars, aux::Vars) end
function integral_set_aux!(::RadiationModel, aux::Vars, integ::Vars) end
function reverse_integral_load_aux!(::RadiationModel, integ::Vars, state::Vars, aux::Vars) end
function reverse_integral_set_aux!(::RadiationModel, aux::Vars, integ::Vars) end
function flux_radiation!(::RadiationModel, flux::Grad, state::Vars, aux::Vars, t::Real) end

# ---------------------------- Begin Boundary Conditions ----------------- #
"""
  DYCOMS_BC <: BoundaryCondition
  Prescribes boundary conditions for Dynamics of Marine Stratocumulus Case
#Fields
$(DocStringExtensions.FIELDS)
"""
struct DYCOMS_BC{FT} <: BoundaryCondition
  "Drag coefficient"
  C_drag::FT
  "Latent Heat Flux"
  LHF::FT
  "Sensible Heat Flux"
  SHF::FT
end

"""
    atmos_boundary_state!(nf::Union{NumericalFluxNonDiffusive, NumericalFluxGradient},
                          bc::DYCOMS_BC, args...)

For the non-diffussive and gradient terms we just use the `NoFluxBC`
"""
atmos_boundary_state!(nf::Union{NumericalFluxNonDiffusive, NumericalFluxGradient},
                      bc::DYCOMS_BC,
                      args...) = atmos_boundary_state!(nf, NoFluxBC(), args...)

"""
    atmos_boundary_flux_diffusive!(nf::NumericalFluxDiffusive,
                                   bc::DYCOMS_BC, atmos::AtmosModel,
                                   F,
                                   state⁺, diff⁺, aux⁺, n⁻,
                                   state⁻, diff⁻, aux⁻,
                                   bctype, t,
                                   state1⁻, diff1⁻, aux1⁻)

When `bctype == 1` the `NoFluxBC` otherwise the specialized DYCOMS BC is used
"""
function atmos_boundary_flux_diffusive!(nf::CentralNumericalFluxDiffusive,
                                        bc::DYCOMS_BC,
                                        atmos::AtmosModel, F,
                                        state⁺, diff⁺, aux⁺,
                                        n⁻,
                                        state⁻, diff⁻, aux⁻,
                                        bctype, t,
                                        state1⁻, diff1⁻, aux1⁻)
  if bctype != 1
    atmos_boundary_flux_diffusive!(nf, NoFluxBC(), atmos, F,
                                   state⁺, diff⁺, aux⁺, n⁻,
                                   state⁻, diff⁻, aux⁻,
                                   bctype, t,
                                   state1⁻, diff1⁻, aux1⁻)
  else
    # Start with the noflux BC and then build custom flux from there
    atmos_boundary_state!(nf, NoFluxBC(), atmos,
                          state⁺, diff⁺, aux⁺, n⁻,
                          state⁻, diff⁻, aux⁻,
                          bctype, t)

    # ------------------------------------------------------------------------
    # (<var>_FN) First node values (First interior node from bottom wall)
    # ------------------------------------------------------------------------
    u_FN = state1⁻.ρu / state1⁻.ρ
    windspeed_FN = norm(u_FN)

    # ----------------------------------------------------------
    # Extract components of diffusive momentum flux (minus-side)
    # ----------------------------------------------------------
    _, τ⁻ = turbulence_tensors(atmos.turbulence, state⁻, diff⁻, aux⁻, t)

    # ----------------------------------------------------------
    # Boundary momentum fluxes
    # ----------------------------------------------------------
    # Case specific for flat bottom topography, normal vector is n⃗ = k⃗ = [0, 0, 1]ᵀ
    # A more general implementation requires (n⃗ ⋅ ∇A) to be defined where A is
    # replaced by the appropriate flux terms
    C_drag = bc.C_drag
    @inbounds begin
      τ13⁺ = - C_drag * windspeed_FN * u_FN[1]
      τ23⁺ = - C_drag * windspeed_FN * u_FN[2]
      τ21⁺ = τ⁻[2,1]
    end

    # Assign diffusive momentum and moisture fluxes
    # (i.e. ρ𝛕 terms)
    FT = eltype(state⁺)
    τ⁺ = SHermitianCompact{3, FT, 6}(SVector(0   ,
                                             τ21⁺, τ13⁺,
                                             0   , τ23⁺, 0))

    # ----------------------------------------------------------
    # Boundary moisture fluxes
    # ----------------------------------------------------------
    # really ∇q_tot is being used to store d_q_tot
    d_q_tot⁺  = SVector(0, 0, bc.LHF/(LH_v0))

    # ----------------------------------------------------------
    # Boundary energy fluxes
    # ----------------------------------------------------------
    # Assign diffusive enthalpy flux (i.e. ρ(J+D) terms)
    d_h_tot⁺ = SVector(0, 0, bc.LHF + bc.SHF)

    # Set the flux using the now defined plus-side data
    flux_diffusive!(atmos, F, state⁺, τ⁺, d_h_tot⁺)
    flux_diffusive!(atmos.moisture, F, state⁺, d_q_tot⁺)
  end
end
# ------------------------ End Boundary Condition --------------------- #

# ------------------------ Begin Radiation Model ---------------------- #
"""
  DYCOMSRadiation <: RadiationModel

Stevens et. al (2005) approximation of longwave radiative fluxes in DYCOMS.
Analytical description as a function of the liquid water path and inversion height zᵢ

* Stevens, B. et. al. (2005) "Evaluation of Large-Eddy Simulations via Observations of Nocturnal Marine Stratocumulus". Mon. Wea. Rev., 133, 1443–1462, https://doi.org/10.1175/MWR2930.1
"""
struct DYCOMSRadiation{FT} <: RadiationModel
  "mass absorption coefficient `[m^2/kg]`"
  κ::FT
  "Troposphere cooling parameter `[m^(-4/3)]`"
  α_z::FT
  "Inversion height `[m]`"
  z_i::FT
  "Density"
  ρ_i::FT
  "Large scale divergence `[s^(-1)]`"
  D_subsidence::FT
  "Radiative flux parameter `[W/m^2]`"
  F_0::FT
  "Radiative flux parameter `[W/m^2]`"
  F_1::FT
end

vars_aux(m::DYCOMSRadiation, FT) = @vars(Rad_flux::FT)

vars_integrals(m::DYCOMSRadiation, FT) = @vars(attenuation_coeff::FT)
function integral_load_aux!(m::DYCOMSRadiation, integrand::Vars, state::Vars, aux::Vars)
  FT = eltype(state)
  integrand.radiation.attenuation_coeff = state.ρ * m.κ * aux.moisture.q_liq
end
function integral_set_aux!(m::DYCOMSRadiation, aux::Vars, integrand::Vars)
  integrand = integrand.radiation.attenuation_coeff
  aux.∫dz.radiation.attenuation_coeff = integrand
end

vars_reverse_integrals(m::DYCOMSRadiation, FT) = @vars(attenuation_coeff::FT)
function reverse_integral_load_aux!(m::DYCOMSRadiation, integrand::Vars, state::Vars, aux::Vars)
  FT = eltype(state)
  integrand.radiation.attenuation_coeff = state.ρ * m.κ * aux.moisture.q_liq
end
function reverse_integral_set_aux!(m::DYCOMSRadiation, aux::Vars, integrand::Vars)
  aux.∫dnz.radiation.attenuation_coeff = integrand.radiation.attenuation_coeff
end

function flux_radiation!(m::DYCOMSRadiation, atmos::AtmosModel, flux::Grad, state::Vars,
                         aux::Vars, t::Real)
  FT = eltype(flux)
  z = altitude(atmos.orientation, aux)
  Δz_i = max(z - m.z_i, -zero(FT))
  # Constants
  upward_flux_from_cloud  = m.F_0 * exp(-aux.∫dnz.radiation.attenuation_coeff)
  upward_flux_from_sfc = m.F_1 * exp(-aux.∫dz.radiation.attenuation_coeff)
  free_troposphere_flux = m.ρ_i * FT(cp_d) * m.D_subsidence * m.α_z * cbrt(Δz_i) * (Δz_i/4 + m.z_i)
  F_rad = upward_flux_from_sfc + upward_flux_from_cloud + free_troposphere_flux
  ẑ = vertical_unit_vector(atmos.orientation, aux)
  flux.ρe += F_rad * ẑ
end
function preodefun!(m::DYCOMSRadiation, aux::Vars, state::Vars, t::Real)
end
# -------------------------- End Radiation Model ------------------------ #

"""
  Initial Condition for DYCOMS_RF01 LES
@article{doi:10.1175/MWR2930.1,
author = {Stevens, Bjorn and Moeng, Chin-Hoh and Ackerman,
          Andrew S. and Bretherton, Christopher S. and Chlond,
          Andreas and de Roode, Stephan and Edwards, James and Golaz,
          Jean-Christophe and Jiang, Hongli and Khairoutdinov,
          Marat and Kirkpatrick, Michael P. and Lewellen, David C. and Lock, Adrian and
          Maeller, Frank and Stevens, David E. and Whelan, Eoin and Zhu, Ping},
title = {Evaluation of Large-Eddy Simulations via Observations of Nocturnal Marine Stratocumulus},
journal = {Monthly Weather Review},
volume = {133},
number = {6},
pages = {1443-1462},
year = {2005},
doi = {10.1175/MWR2930.1},
URL = {https://doi.org/10.1175/MWR2930.1},
eprint = {https://doi.org/10.1175/MWR2930.1}
}
"""
function init_dycoms!(bl, state, aux, (x,y,z), t)
    FT = eltype(state)

    z = FT(z)

    # These constants are those used by Stevens et al. (2005)
    qref       = FT(9.0e-3)
    q_pt_sfc   = PhasePartition(qref)
    Rm_sfc     = FT(gas_constant_air(q_pt_sfc))
    T_sfc      = FT(290.4)
    P_sfc      = FT(MSLP)

    # Specify moisture profiles
    q_liq      = FT(0)
    q_ice      = FT(0)
    zb         = FT(600)         # initial cloud bottom
    zi         = FT(840)         # initial cloud top

    if z <= zi
        θ_liq  = FT(289.0)
        q_tot  = qref
    else
        θ_liq  = FT(297.0) + (z - zi)^(FT(1/3))
        q_tot  = FT(1.5e-3)
    end

    ugeo = FT(7)
    vgeo = FT(-5.5)
    u, v, w = ugeo, vgeo, FT(0)

    # Perturb initial state to break symmetry and trigger turbulent convection
    r1 = FT(rand(Uniform(-0.002, 0.002)))
    r2 = FT(rand(Uniform(-0.00001, 0.00001)))
    r3 = FT(rand(Uniform(-0.001, 0.001)))
    r4 = FT(rand(Uniform(-0.001, 0.001)))
    if z <= 400.0
        θ_liq += r1 * θ_liq
        q_tot += r2 * q_tot
        u     += r3 * u
        v     += r4 * v
    end

    # Pressure
    H     = Rm_sfc * T_sfc / grav
    p     = P_sfc * exp(-z / H)

    # Density, Temperature
    ts    = LiquidIcePotTempSHumEquil_given_pressure(θ_liq, p, q_tot)
    ρ     = air_density(ts)

    e_kin = FT(1/2) * FT((u^2 + v^2 + w^2))
    e_pot = grav * z
    E     = ρ * total_energy(e_kin, e_pot, ts)

    state.ρ               = ρ
    state.ρu              = SVector(ρ*u, ρ*v, ρ*w)
    state.ρe              = E
    state.moisture.ρq_tot = ρ * q_tot

    return nothing
end

function config_dycoms(FT, N, resolution, xmax, ymax, zmax)
    # Reference state
    T_min   = FT(289)
    T_s     = FT(290.4)
    Γ_lapse = FT(grav/cp_d)
    T       = LinearTemperatureProfile(T_min, T_s, Γ_lapse)
    rel_hum = FT(0)
    ref_state = HydrostaticState(T, rel_hum)

    # Radiation model
    κ             = FT(85)
    α_z           = FT(1)
    z_i           = FT(840)
    ρ_i           = FT(1.13)
    D_subsidence  = FT(0) # 0 for stable testing, 3.75e-6 in practice
    F_0           = FT(70)
    F_1           = FT(22)
    radiation = DYCOMSRadiation{FT}(κ, α_z, z_i, ρ_i, D_subsidence, F_0, F_1)

    # Sources
    f_coriolis    = FT(1.03e-4)
    u_geostrophic = FT(7.0)
    v_geostrophic = FT(-5.5)
    w_ref         = FT(0)
    u_relaxation  = SVector(u_geostrophic, v_geostrophic, w_ref)
    # Sponge
    c_sponge = 1
    # Rayleigh damping
    zsponge = FT(1500.0)
    rayleigh_sponge = RayleighSponge{FT}(zmax, zsponge, c_sponge, u_relaxation, 2)
    # Geostrophic forcing
    geostrophic_forcing = GeostrophicForcing{FT}(f_coriolis, u_geostrophic, v_geostrophic)

    # Boundary conditions
    # SGS Filter constants
    C_smag = FT(0.21) # 0.21 for stable testing, 0.18 in practice
    C_drag = FT(0.0011)
    LHF    = FT(115)
    SHF    = FT(15)
    bc = DYCOMS_BC{FT}(C_drag, LHF, SHF)
    ics = init_dycoms!
    source = (Gravity(),
              rayleigh_sponge,
              Subsidence{FT}(D_subsidence),
              geostrophic_forcing)

    model = AtmosModel{FT}(AtmosLESConfiguration;
                           ref_state=ref_state,
                          turbulence=SmagorinskyLilly{FT}(C_smag),
                            moisture=EquilMoist(5),
                           radiation=radiation,
                              source=source,
                   boundarycondition=bc,
                          init_state=ics)

    config = CLIMA.Atmos_LES_Configuration("DYCOMS", N, resolution, xmax, ymax, zmax,
                                           init_dycoms!,
                                           solver_type=CLIMA.ExplicitSolverType(solver_method=LSRK144NiegemannDiehlBusch),
                                           model=model)

    return config
end

function main()
    CLIMA.init()

    FT = Float64

    # DG polynomial order
    N = 4

    # Domain resolution and size
    Δh = FT(40)
    Δv = FT(20)
    resolution = (Δh, Δh, Δv)

    xmax = 1000
    ymax = 1000
    zmax = 2500

    t0 = FT(0)
    timeend = FT(100)

    driver_config = config_dycoms(FT, N, resolution, xmax, ymax, zmax)
    solver_config = CLIMA.setup_solver(t0, timeend, driver_config, forcecpu=true)

    cbtmarfilter = GenericCallbacks.EveryXSimulationSteps(2) do (init=false)
        Filters.apply!(solver_config.Q, 6, solver_config.dg.grid, TMARFilter())
        nothing
    end

    result = CLIMA.invoke!(solver_config;
                          user_callbacks=(cbtmarfilter,),
                          check_euclidean_distance=true)
end

main()

max_stars_repo_name: favba/CLIMA

max_issues_repo_name: favba/CLIMA

max_forks_repo_name: favba/CLIMA


text:
include("../../module/CustomStencil.jl")
include("../../misc/misc.jl")
include("../../misc/cpu_stencils.jl")
## Star Stencil definition with radius = 4

##

dense_coefs = ones(9,9,9)

for i = parse(Int, ARGS[1])
    global dense_coefs
    for x = (i-1):-1:0
        for y = (i-1):-1:0
            for z = (i-1):-1:0
                dense_coefs[1+x,1+y,1+z] = 0
                dense_coefs[1+x,1+y,9-z] = 0
                dense_coefs[1+x,9-y,1+z] = 0
                dense_coefs[1+x,9-y,9-z] = 0
                dense_coefs[9-x,1+y,1+z] = 0
                dense_coefs[9-x,1+y,9-z] = 0
                dense_coefs[9-x,9-y,1+z] = 0
                dense_coefs[9-x,9-y,9-z] = 0
            end
        end
    end
    @show sum(dense_coefs .==1)/729
    st_def = CreateStencilDefinition(dense_coefs)
    bdim = 32
    if length(ARGS) == 2
        bdim = 16
    end
    st_inst1 = NewStencilInstance(st_def, bdim=bdim, m_step=false)

    radius = 4
    nx = 8
    ny = 8
    nz = 8
    dx = 1<<(nx)
    dy = 1<<(ny)
    dz = 1<<(nz)

    data = CreateData(dx,dy,dz)

    t_steps = 16
    gpu_out = ApplyFFTstencil(st_inst1, data, t_steps)
    gpu_out = ApplyStencil(st_inst1, data, t_steps)
    CUDA.cuProfilerStart()
    NVTX.@range "standard" begin
        gpu_out = ApplyStencil(st_inst1, data, t_steps) end
    ApplyFFTstencil(st_inst1, data, t_steps,1)
    NVTX.@range "r1" begin
        gpu_out = ApplyFFTstencil(st_inst1, data, t_steps,1) end

    ApplyFFTstencil(st_inst1, data, t_steps,16)
    NVTX.@range "r2" begin
        gpu_out = ApplyFFTstencil(st_inst1, data, t_steps,4) end

    ApplyFFTstencil(st_inst1, data, 32,8)
    NVTX.@range "r3" begin
        gpu_out = ApplyFFTstencil(st_inst1, data, 32,8) end

    ApplyFFTstencil(st_inst1, data, 48,12)
    NVTX.@range "r4" begin
        gpu_out = ApplyFFTstencil(st_inst1, data, 48,12) end
    exit()

end


text:
@testset "Function quantile" begin
    @test quantile([0.5377;    1.8339;
                   -2.2588;    0.8622;    0.3188;
                      -1.3077;   -0.4336;    0.3426;
                          3.5784;    2.7694],
                          collect(1:3)./(3+1.0);
                          sorted=false, alpha = 0.5, beta = 0.5) == [-0.4336; 0.44015; 1.8339]

    @test isapprox(quantile([0.5377;    1.8339;
             -2.2588;    0.8622;    0.3188;
                -1.3077;   -0.4336;    0.3426;
                    3.5784;    2.7694], collect(1:11)./(11+1.0); sorted=false, alpha = 0.5, beta = 0.5),
[-1.9417666666666666;
     -1.162016666666667;
     -0.4336 ;
      0.19339999999999974;
      0.33466666666666667;
      0.44015   ;
      0.6458666666666668 ;
      1.0241499999999994;
      1.8339    ;
      2.6134833333333343 ;
      3.308733333333333], atol = 1e-10)
end

@testset "Center and standard deviations of a sample for diagonal component I" begin

    Nx = 1
    Ne = 10
    p = 3
    γ = 3.0

    X = reshape([-1.5438
       -1.5518
        0.8671
       -0.1454
       -0.3862
        1.3162
       -0.7965
        0.1354
        0.4178
        0.8199],1,Ne)

    S = RadialMapComponent(Nx, p);
    sort!(X; dims = 2)
    center_std_diag!(S, X, γ)

    @test norm(S.ξ[1]-[-1.4192500000000001, -0.45458333333333345, -0.0050000000000000044, 0.48481666666666645, 0.8592333333333333])<1e-10
    @test norm(S.σ[1]-[2.894, 2.121375, 1.4090999999999998, 1.29635, 1.1232500000000005])<1e-10

end

@testset "Center and standard deviations of a sample for diagonal component II" begin

    Nx = 1
    Ne = 10
    p = 6
    γ = 1.5

    X   = reshape([-1.5438
                   -1.5518
                    0.8671
                   -0.1454
                   -0.3862
                    1.3162
                   -0.7965
                    0.1354
                    0.4178
                    0.8199],1,Ne)

    S = RadialMapComponent(Nx, p);
    sort!(X; dims = 2)
    center_std_diag!(S, X, γ)

    @test norm(S.ξ[1]-[-1.5469111111111111, -1.0040833333333332, -0.45458333333333345, -0.15877777777777774, 0.15108888888888883, 0.48481666666666645, 0.8330111111111111, 1.0417500000000002])<1e-10
    @test norm(S.σ[1]-[0.8142416666666669, 0.8192458333333333, 0.6339791666666666, 0.45425416666666674, 0.4826958333333331, 0.5114416666666668, 0.4177000000000003, 0.31310833333333354])<1e-10

end

@testset "Center and standard deviations of a sample for diagonal component III" begin
    Nx = 2
    Ne = 10
    p = 1
    γ = 3.0
    λ = 0.0

    X = [-0.8544   -0.7938;0.3573    0.5410; 2.7485   -0.5591; -1.5130    1.9766; 0.4340    0.5447; -0.2298   -0.1379;-0.8271   0.6199;-0.8320   -0.0056;0.4979    1.1072;2.3156   -0.1856]'

    S = RadialMapComponent(Nx, p)
    sort!(X, dims = 2)

    center_std_diag!(S, X, γ)

    @test norm(S.ξ[2] -  [ -0.1856,0.2677, 0.6199])<1e-10
    @test norm(S.σ[2] -  [1.3598999999999999, 1.20825, 1.0566])<1e-10
end

@testset "Center and standard deviations of a sample for off-diagonal component" begin

    Nx = 2
    Ne = 10
    p = 2
    γ = 3.0
    λ = 0.0

    X = [-0.8544   -0.7938;0.3573    0.5410; 2.7485   -0.5591; -1.5130    1.9766; 0.4340    0.5447; -0.2298   -0.1379;-0.8271   0.6199;-0.8320   -0.0056;0.4979    1.1072;2.3156   -0.1856]'

    S = RadialMapComponent(Nx, p)
    sort!(X; dims = 2)
    center_std_diag!(S, X, γ)
    center_std_off!(S, X, γ)

    @test norm(S.ξ[1] - [-0.8279166666666666, 0.44464999999999993])<1e-10
    @test norm(S.ξ[2] - [-0.37235, -0.07175, 0.54285, 0.86355])<1e-10

    @test norm(S.σ[1] - [3.8177, 3.8177])<1e-10
    @test norm(S.σ[2] - [0.9018000000000002, 1.3728, 1.4029500000000001, 0.9621])<1e-10

end

@testset "Center and standard deviations of a sample" begin

    Nx = 2
    Ne = 10
    p = 2
    γ = 3.0
    λ = 0.0

    X = Matrix([-0.8544   -0.7938;0.3573    0.5410; 2.7485   -0.5591; -1.5130    1.9766; 0.4340    0.5447; -0.2298   -0.1379;-0.8271   0.6199;-0.8320   -0.0056;0.4979    1.1072;2.3156   -0.1856]')

    S = RadialMap(Nx, p; γ = γ)

    center_std!(S, X)

    @test norm(X - [-0.8544   -0.7938;0.3573    0.5410; 2.7485   -0.5591; -1.5130    1.9766; 0.4340    0.5447; -0.2298   -0.1379;-0.8271   0.6199;-0.8320   -0.0056;0.4979    1.1072;2.3156   -0.1856]')<1e-10

    @test norm(S.C[1].ξ[1] - [-0.8432, -0.52845, 0.39565, 1.40675])<1e-10
    @test norm(S.C[2].ξ[1] - [-0.8279166666666666, 0.44464999999999993])<1e-10
    @test norm(S.C[2].ξ[2] - [-0.37235, -0.07175, 0.54285, 0.86355])<1e-10

    @test norm(S.C[1].σ[1] -  [0.9442499999999999, 1.858275, 2.9028, 3.0332999999999997])<1e-10
    @test norm(S.C[2].σ[1] - [3.8177, 3.8177])<1e-10
    @test norm(S.C[2].σ[2] - [0.9018000000000002, 1.3728, 1.4029500000000001, 0.9621])<1e-10
end

## Perform the same validation with a Sparse structure

@testset "SparseRadialMapComponent: Center and standard deviations of a sample for diagonal component I" begin

    Nx = 1
    Ne = 10
    p = 3
    γ = 3.0

    X = reshape([-1.5438
                 -1.5518
                  0.8671
                 -0.1454
                 -0.3862
                  1.3162
                 -0.7965
                  0.1354
                  0.4178
                  0.8199],1,Ne)

    S = SparseRadialMapComponent(Nx, p);
    sort!(X; dims = 2)
    center_std_diag!(S, X, γ)

    @test norm(S.ξ[1]-[-1.4192500000000001, -0.45458333333333345, -0.0050000000000000044, 0.48481666666666645, 0.8592333333333333])<1e-10
    @test norm(S.σ[1]-[2.894, 2.121375, 1.4090999999999998, 1.29635, 1.1232500000000005])<1e-10

end

@testset "SparseRadialMapComponent Center and standard deviations of a sample for diagonal component II" begin

    Nx = 1
    Ne = 10
    p = 6
    γ = 1.5

    X  =  reshape([-1.5438
                   -1.5518
                    0.8671
                   -0.1454
                   -0.3862
                    1.3162
                   -0.7965
                    0.1354
                    0.4178
                    0.8199],1,Ne)

    S = SparseRadialMapComponent(Nx, p);
    sort!(X; dims = 2)
    center_std_diag!(S, X, γ)

    @test norm(S.ξ[1]-[-1.5469111111111111, -1.0040833333333332, -0.45458333333333345, -0.15877777777777774, 0.15108888888888883, 0.48481666666666645, 0.8330111111111111, 1.0417500000000002])<1e-10
    @test norm(S.σ[1]-[0.8142416666666669, 0.8192458333333333, 0.6339791666666666, 0.45425416666666674, 0.4826958333333331, 0.5114416666666668, 0.4177000000000003, 0.31310833333333354])<1e-10

end

@testset "SparseRadialMapComponent Center and standard deviations of a sample for diagonal component III" begin
    Nx = 2
    Ne = 10
    p = 1
    γ = 3.0
    λ = 0.0

    X = Matrix([-0.8544   -0.7938;0.3573    0.5410; 2.7485   -0.5591; -1.5130    1.9766; 0.4340    0.5447; -0.2298   -0.1379;-0.8271   0.6199;-0.8320   -0.0056;0.4979    1.1072;2.3156   -0.1856]')

    S = SparseRadialMapComponent(Nx, p)
    sort!(X; dims = 2)

    center_std_diag!(S, X, γ)

    @test norm(S.ξ[2] -  [ -0.1856,0.2677, 0.6199])<1e-10
    @test norm(S.σ[2] -  [1.3598999999999999, 1.20825, 1.0566])<1e-10
end

@testset "SparseRadialMapComponent Center and standard deviations of a sample for off-diagonal component" begin

    Nx = 2
    Ne = 10
    p = 2
    γ = 3.0
    λ = 0.0

    X = Matrix([-0.8544   -0.7938;0.3573    0.5410; 2.7485   -0.5591; -1.5130    1.9766; 0.4340    0.5447; -0.2298   -0.1379;-0.8271   0.6199;-0.8320   -0.0056;0.4979    1.1072;2.3156   -0.1856]')

    S = SparseRadialMapComponent(Nx, p)
    sort!(X; dims = 2)
    center_std_diag!(S, X, γ)
    center_std_off!(S, X, γ)

    @test norm(S.ξ[1] - [-0.8279166666666666, 0.44464999999999993])<1e-10
    @test norm(S.ξ[2] - [-0.37235, -0.07175, 0.54285, 0.86355])<1e-10

    @test norm(S.σ[1] - [3.8177, 3.8177])<1e-10
    @test norm(S.σ[2] - [0.9018000000000002, 1.3728, 1.4029500000000001, 0.9621])<1e-10

end

@testset "Center and standard deviations of a sample" begin

    Nx = 2
    Ne = 10
    p = 2
    γ = 3.0
    λ = 0.0

    X = Matrix([-0.8544   -0.7938;0.3573    0.5410; 2.7485   -0.5591; -1.5130    1.9766; 0.4340    0.5447; -0.2298   -0.1379;-0.8271   0.6199;-0.8320   -0.0056;0.4979    1.1072;2.3156   -0.1856]')

    S = SparseRadialMap(Nx, [p; p]; γ = γ, λ = λ)

    center_std!(S, X)

    @test norm(X - [-0.8544   -0.7938;0.3573    0.5410; 2.7485   -0.5591; -1.5130    1.9766; 0.4340    0.5447; -0.2298   -0.1379;-0.8271   0.6199;-0.8320   -0.0056;0.4979    1.1072;2.3156   -0.1856]')<1e-10

    @test norm(S.C[1].ξ[1] - [-0.8432, -0.52845, 0.39565, 1.40675])<1e-10
    @test norm(S.C[2].ξ[1] - [-0.8279166666666666, 0.44464999999999993])<1e-10
    @test norm(S.C[2].ξ[2] - [-0.37235, -0.07175, 0.54285, 0.86355])<1e-10

    @test norm(S.C[1].σ[1] -  [0.9442499999999999, 1.858275, 2.9028, 3.0332999999999997])<1e-10
    @test norm(S.C[2].σ[1] - [3.8177, 3.8177])<1e-10
    @test norm(S.C[2].σ[2] - [0.9018000000000002, 1.3728, 1.4029500000000001, 0.9621])<1e-10
end


text:
function adhoc_partition(xs, n)
    @check firstindex(xs) == 1
    m = cld(length(xs), n)
    return (view(xs, i*n+1:min((i+1)*n, length(xs))) for i in 0:m-1)
end

if VERSION >= v"1.4"
    const _partition = Iterators.partition
else
    const _partition = adhoc_partition
end

@inline _asval(x::Val) = x
@inline _asval(x) = Val(x)

const SIMDValFlag = Union{Val{true}, Val{false}, Val{:ivdep}}
const SIMDFlag = Union{Bool, Symbol, SIMDValFlag}

__verify_simd_flag(simd::SIMDValFlag, _) = simd
__verify_simd_flag(_, simd) =
    throw(ArgumentError("Invalid `simd` option: $(repr(simd))"))

verify_simd_flag(simd) = __verify_simd_flag(_asval(simd), simd)

# TODO: handle Union
function elsizeof(::Type{T}) where {T}
    if Base.allocatedinline(Some{T})
        return sizeof(Some{T})
    else
        return sizeof(Ptr{Any})
    end
end

function _median(order, (a, b, c)::NTuple{3,Any})
    # Sort `(a, b, c)`:
    if Base.lt(order, b, a)
        a, b = b, a
    end
    if Base.lt(order, c, a)
        a, c = c, a
    end
    if Base.lt(order, c, b)
        b, c = c, b
    end
    return b
end

_median(order, (a, b, c, d, e, f, g, h, i)::NTuple{9,Any}) = _median(
    order,
    (_median(order, (a, b, c)), _median(order, (d, e, f)), _median(order, (g, h, i))),
)

function maptasks(f, xs)
    tasks = Task[]
    @sync for x in xs
        push!(tasks, @spawn f(x))
    end
    return map(fetch, tasks)
end

function exclusive_cumsum!(xs, acc = zero(eltype(xs)))
    @inbounds for i in eachindex(xs)
        xs[i], x = acc, xs[i]
        acc += x
    end
    return acc
end

struct Unroll{N,A}
    xs::A
end

Unroll{N}(xs::A) where {N,A} = Unroll{N,A}(xs)

@inline _foldlargs(op, acc) = acc
@inline _foldlargs(op, acc, x, xs...) = _foldlargs(op, op(acc, x), xs...)

@inline function _foldl(op::F, acc, itr::Unroll{N}) where {N,F}
    i = firstindex(itr.xs)
    n = lastindex(itr.xs) - N
    while i <= n
        acc = let i = i
            _foldlargs(acc, ntuple(identity, Val{N}())...) do acc, k
                op(acc, @inbounds itr.xs[i + (k - 1)])
            end
        end
        i += N
    end
    while i <= lastindex(itr.xs)
        acc = op(acc, @inbounds itr.xs[i])
        i += 1
    end
    return acc
end

function define_docstrings()
    docstrings = [:ThreadsX => joinpath(dirname(@__DIR__), "README.md")]
    docsdir = joinpath(@__DIR__, "docs")
    for filename in readdir(docsdir)
        stem, ext = splitext(filename)
        ext == ".md" || continue
        name = Symbol(stem)
        name in names(ThreadsX, all=true) || continue
        push!(docstrings, name => joinpath(docsdir, filename))
    end
    for (name, path) in docstrings
        include_dependency(path)
        doc = read(path, String)
        doc = replace(doc, r"^```julia"m => "```jldoctest $name")
        doc = replace(doc, "<kbd>TAB</kbd>" => "_TAB_")
        @eval ThreadsX $Base.@doc $doc $name
    end
end


text:
struct Interpolated{X, Y, I}
    x0::X
    xf::X
    y0::Y
    yf::Y
    interpolator::I
    clamp::Bool

    function Interpolated(x0::X, xf::X, y0::Y, yf::Y, f=identity; min_num_derivs=Val(0), clamp::Bool=true) where {X, Y}
        interpolator = make_interpolator(f, min_num_derivs)
        new{X, Y, typeof(interpolator)}(x0, xf, y0, yf, interpolator, clamp)
    end
end

function Interpolated(x0, xf, y0, yf, f=identity; min_num_derivs=Val(0))
    Interpolated(promote(x0, xf)..., promote(y0, yf)..., f, min_num_derivs=min_num_derivs)
end

# Call overload with derivatives
function (trajectory::Interpolated)(x, ::Val{num_derivs}) where num_derivs
    x0, xf = trajectory.x0, trajectory.xf
    if !trajectory.clamp && (x < x0 || x > xf)
        throw_trajectory_domain_error(x, x0, xf)
    end

    # 1. Compute θ(x) = clamp((x - x0) / (xf - x0), 0, 1) and dθ/dx.
    Δx = xf - x0
    θ = (x - x0) / Δx
    if θ <= zero(θ)
        θ = zero(θ)
        dθdx = zero(θ)
    elseif θ >= one(θ)
        θ = one(θ)
        dθdx = zero(θ)
    else
        dθdx = one(x - x0) / Δx
    end

    # 2. Warp θ using interpolator to obtain α(θ(x)) and its partial derivatives w.r.t θ.
    α, α_θ_derivs = let α_vals = trajectory.interpolator(θ, Val(num_derivs))
        α_vals[1], Base.tail(α_vals)
    end

    # 3. Compute derivatives of α w.r.t x using higher-order chain rule,
    # where higher order derivatives of θ w.r.t. x are zero.
    α_x_derivs = let α_θ_derivs=α_θ_derivs, dθdx=dθdx
        ntuple(i -> α_θ_derivs[i] * dθdx^i, Val(num_derivs))
    end

    # 3. Interpolate between y0 and yf using α.
    y0, yf = trajectory.y0, trajectory.yf
    y, dydα = interpolate(y0, yf, α)

    # 4. Compute derivatives of y w.r.t. x using higher-order chain rule,
    # where higher order derivatives of y w.r.t. α are zero.
    y_x_derivs = ntuple(i -> dydα * α_x_derivs[i], Val(num_derivs))

    y, y_x_derivs...
end

# Single-argument call overload just returns the value:
(trajectory::Interpolated)(x) = trajectory(x, Val(0))[1]

# Interpolation
function interpolate(y0, yf, α)
    Δy = yf - y0
    y = y0 + α * Δy
    dydα = Δy
    y, dydα
end

function interpolate(y0::Rotation{3}, yf::Rotation{3}, α)
    interpolate(Quat(y0), Quat(yf), α)
end

function interpolate(y0::Quat, yf::Quat, α)
    Δy = AngleAxis(y0 \ yf)
    angle = rotation_angle(Δy)
    axis = rotation_axis(Δy)
    y = y0 * Quat(AngleAxis(α * angle, axis...))
    dydα = axis * angle # Lie derivative
    y, dydα
end

# `make_interpolator` overloads
# function returned by `make_interpolator` has signature (θ, n) -> (f(θ), (f′(θ), ...,f^{(n)}(θ)))
function make_interpolator(f, ::Val)
    function (θ, ::Val{num_derivs}) where {num_derivs}
        # Default to calling the `nthderiv` function
        f(θ), ntuple(n -> nthderiv(f, θ, n), Val(num_derivs))...
    end
end

function make_interpolator(p::Polynomial, ::Val{min_num_derivs}) where min_num_derivs
    PolyDerivEvaluator(p, Val(min_num_derivs))
end

# `nthderiv` overloads
nthderiv(::typeof(identity), θ, n::Int) = n === 1 ? one(θ) : zero(θ)


text:
include("bkfiltertest.jl")

x1,p1 = augment(s,1e-10*eye(2))
@test x1 == [1.0,0.0,0.0,0.0]
@test p1 == [4*eye(2,2) zeros(2,2);
             zeros(2,2) 1e-10*eye(2)]

function fa(x::Vector)
    x1 = zeros(2)
    x1[1] = x[1] + 0.01*x[2] + x[3]
    x1[2] = x[2] - 0.01*x[1] + x[4]
    x1
end

fam = AugmentedUnscentedModel(fa,1e-10*eye(2))

function ha(x)
    x1 = zeros(1)
    x1[1] = x[1] + x[3]
end

zam = AugmentedUnscentedObservationModel(ha,[0.1]')

kaf = AugmentedUnscentedKalmanFilter(s,fam,zam)

kaf1 = predict(kaf)
kaf2 = update(kaf1,y)

@test_approx_eq kaf2.x.x kf2.x.x
@test_approx_eq kaf2.x.p kf2.x.p


text:
# This file is part of the IntervalArithmetic.jl package; MIT licensed

## floor, ceil, trunc, sign, roundTiesToEven, roundTiesToAway
for f in (:floor, :ceil, :trunc)
    @eval ($f)(a::Interval) = unsafe_interval(($f)(a.lo), ($f)(a.hi))
    @eval ($f)(x::EmptyInterval) = x
end

const RoundTiesToEven = RoundNearest
const RoundTiesToAway = RoundNearestTiesAway

"""
    round(a::Interval[, RoundingMode])

Returns the interval with rounded to an interger limits.

For compliance with the IEEE Std 1788-2015, "roundTiesToEven" corresponds
to `round(a)` or `round(a, RoundNearest)`, and "roundTiesToAway"
to `round(a, RoundNearestTiesAway)`.
"""
round(a::ValidInterval) = round(a, RoundNearest)
round(a::ValidInterval, ::RoundingMode{:ToZero}) = trunc(a)
round(a::ValidInterval, ::RoundingMode{:Up}) = ceil(a)
round(a::ValidInterval, ::RoundingMode{:Down}) = floor(a)

round(x::EmptyInterval, ::RoundingMode) = x

round(a::Interval, ::RoundingMode{:Nearest}) =
    unsafe_interval(round(a.lo), round(a.hi))

round(a::Interval, ::RoundingMode{:NearestTiesAway}) =
    unsafe_interval(
        round(a.lo, RoundNearestTiesAway),
        round(a.hi, RoundNearestTiesAway)
    )


text:
# Portions translated from SLICOT-Reference distribution
# Copyright (c) 2002-2020 NICONET e.V.
function run_mb02cd(datfile, io=stdout)
    ZERO = 0.0e0
    NIN = 5
    NOUT = 6
    KMAX = 20
    NMAX = 20
    LDG = 2*KMAX
    LDL = NMAX*KMAX
    LDR = NMAX*KMAX
    LDT = KMAX
    LDWORK = ( NMAX - 1 )*KMAX
    LCS = 3*LDWORK
    T = Array{Float64,2}(undef, LDT,NMAX*KMAX)
    G = Array{Float64,2}(undef, LDG,NMAX*KMAX)
    R = Array{Float64,2}(undef, LDR,NMAX*KMAX)
    L = Array{Float64,2}(undef, LDL,NMAX*KMAX)
    CS = Array{Float64,1}(undef, LCS)
    DWORK = Array{Float64,1}(undef, LDWORK)
f = open(datfile,"r")
    readline(f)
    vs = split(readline(f))
    N = parse(BlasInt, vs[1])
    K = parse(BlasInt, vs[2])
    JOB =  vs[3][1]
    TYPET = 'R'
    M = N*K
if ( N<=0 || N>NMAX )
else
if ( K<=0 || K>KMAX )
else

    vs = String[]
    _isz,_jsz = (K,M)
    while length(vs) < _isz*_jsz
        append!(vs, replace.(split(readline(f)),'D'=>'E'))
    end
    for i in 1:_isz
       _i0 = (i-1)*_jsz
       T[i,1:_jsz] .= parsex.(Float64, vs[_i0+1:_i0+_jsz])
    end
# interp call 1

    INFO = SLICOT.mb02cd!(JOB, TYPET, K, N, T, G, R, L, CS, LCS, LDWORK)
    @test INFO == 0
    INFO == 0 || return

if ( INFO!=0 )
else
if ( LSAME( JOB, 'G' ) || LSAME( JOB, 'A' ) || LSAME( JOB, 'L' ) || LSAME( JOB, 'R' ) )
# interp call 2

    #FOREIGN.dlaset!( 'Full', K, K, ZERO, ZERO, G(K+1,1), LDG )
    G[K+1:2*K,1:K] .= ZERO

# interp output 1
    println(io, "G:")
    _nc = M
    _nr = 2*K
    show(io, "text/plain", G[1:_nr,1:_nc])
    println(io)

end # if
if ( LSAME( JOB, 'L' ) || LSAME( JOB, 'A' ) )
# interp output 2
    println(io, "L:")
    _nc = M
    _nr = M
    show(io, "text/plain", L[1:_nr,1:_nc])
    println(io)

end # if
if ( LSAME( JOB, 'R' ) || LSAME( JOB, 'A' ) || LSAME( JOB, 'O' ) )
# interp output 3
    println(io, "R:")
    _nc = M
    _nr = M
    show(io, "text/plain", R[1:_nr,1:_nc])
    println(io)

end # if
end # if
end # if
end # if
    close(f)
end # run_mb02cd()


text:
using Test
using TestSetExtensions

include("../src/Groebner.jl")

using .Groebner
using .Groebner.AbstractAlgebra

# Taken from JuMP/test/solvers.jl
function try_import(name::Symbol)
    try
        @eval import $name
        return true
    catch e
        return false
    end
end

⊂(xs, ys) = all(in(ys), xs)
# set equality
≂(xs, ys) = ⊂(xs, ys) && ⊂(ys, xs)

@testset "All tests" begin

    @includetests ["commons", "io", "io_consistency",
                    "univariate_aa"]

    @includetests ["crt_reconstruction"] # "mod_reconstruction"]#,

    @includetests ["core_f4_reduce", "core_f4_stress",
                    "core_f4", "rational_f4", "groebner_certify"]

    @includetests ["core_isgroebner", "core_isgroebner_stress",
                    "isgroebner_certify"]

    @includetests ["core_normalform", "core_normalform_stress"]

    @includetests ["hard_problems_f4", "large_problems_f4"]

    @includetests ["probabilistic_linalg"]

    if try_import(:DynamicPolynomials)
        @includetests ["io_dynamic", "core_dynamic"]
    end

    @includetests ["onthefly_order_change"]

    @includetests ["core_fglm", "core_kbase"]

    @includetests ["benchmark_systems"]

end


text:
function MOI.eval_objective(nlp::NLPData{T}, variables::Vector{T}) where T
    trajectory!(
        nlp.trajopt.states,
        nlp.trajopt.actions,
        variables,
        nlp.indices.states,
        nlp.indices.actions)
    cost(
        nlp.trajopt.objective,
        nlp.trajopt.states,
        nlp.trajopt.actions,
        nlp.trajopt.parameters)
end

function MOI.eval_objective_gradient(nlp::NLPData{T}, gradient, variables) where T
    fill!(gradient, 0.0)
    trajectory!(
        nlp.trajopt.states,
        nlp.trajopt.actions,
        variables,
        nlp.indices.states,
        nlp.indices.actions)
    gradient!(gradient,
        nlp.indices.state_action,
        nlp.trajopt.objective,
        nlp.trajopt.states,
        nlp.trajopt.actions,
        nlp.trajopt.parameters)
    return
end

function MOI.eval_constraint(nlp::NLPData{T}, violations, variables) where T
    fill!(violations, 0.0)
    trajectory!(
        nlp.trajopt.states,
        nlp.trajopt.actions,
        variables,
        nlp.indices.states,
        nlp.indices.actions)
    constraints!(
        violations,
        nlp.indices.dynamics_constraints,
        nlp.trajopt.dynamics,
        nlp.trajopt.states,
        nlp.trajopt.actions,
        nlp.trajopt.parameters)
    !isempty(nlp.indices.stage_constraints) && constraints!(violations, nlp.indices.stage_constraints, nlp.trajopt.constraints, nlp.trajopt.states, nlp.trajopt.actions, nlp.trajopt.parameters)
    nlp.general_constraint.num_constraint != 0 && constraints!(violations, nlp.indices.general_constraint, nlp.general_constraint, variables, nlp.parameters)
    return
end

function MOI.eval_constraint_jacobian(nlp::NLPData{T}, jacobian, variables) where T
    fill!(jacobian, 0.0)
    trajectory!(
        nlp.trajopt.states,
        nlp.trajopt.actions,
        variables,
        nlp.indices.states,
        nlp.indices.actions)
    jacobian!(
        jacobian,
        nlp.indices.dynamics_jacobians,
        nlp.trajopt.dynamics,
        nlp.trajopt.states,
        nlp.trajopt.actions,
        nlp.trajopt.parameters)
    !isempty(nlp.indices.stage_jacobians) && jacobian!(jacobian, nlp.indices.stage_jacobians, nlp.trajopt.constraints, nlp.trajopt.states, nlp.trajopt.actions, nlp.trajopt.parameters)
    nlp.general_constraint.num_constraint != 0 && jacobian!(jacobian, nlp.indices.general_jacobian, nlp.general_constraint, variables, nlp.parameters)
    return
end

function MOI.eval_hessian_lagrangian(nlp::MOI.AbstractNLPEvaluator, hessian, variables, scaling, duals)
    fill!(hessian, 0.0)
    trajectory!(
        nlp.trajopt.states,
        nlp.trajopt.actions,
        variables,
        nlp.indices.states,
        nlp.indices.actions)
    duals!(
        nlp.trajopt.duals_dynamics,
        nlp.trajopt.duals_constraints,
        nlp.duals_general,
        duals,
        nlp.indices.dynamics_constraints,
        nlp.indices.stage_constraints,
        nlp.indices.general_constraint)
    hessian!(
        hessian,
        nlp.indices.objective_hessians,
        nlp.trajopt.objective,
        nlp.trajopt.states,
        nlp.trajopt.actions,
        nlp.trajopt.parameters,
        scaling)
    hessian_lagrangian!(
        hessian,
        nlp.indices.dynamics_hessians,
        nlp.trajopt.dynamics,
        nlp.trajopt.states,
        nlp.trajopt.actions,
        nlp.trajopt.parameters,
        nlp.trajopt.duals_dynamics)
    hessian_lagrangian!(
        hessian,
        nlp.indices.stage_hessians,
        nlp.trajopt.constraints,
        nlp.trajopt.states,
        nlp.trajopt.actions,
        nlp.trajopt.parameters,
        nlp.trajopt.duals_constraints)
    hessian_lagrangian!(
        hessian,
        nlp.indices.general_hessian,
        nlp.general_constraint,
        variables,
        nlp.parameters,
        nlp.duals_general)
    return
end

MOI.features_available(nlp::MOI.AbstractNLPEvaluator) = nlp.hessian_lagrangian ? [:Grad, :Jac, :Hess] : [:Grad, :Jac]
MOI.initialize(nlp::MOI.AbstractNLPEvaluator, features) = nothing
MOI.jacobian_structure(nlp::MOI.AbstractNLPEvaluator) = nlp.jacobian_sparsity
MOI.hessian_lagrangian_structure(nlp::MOI.AbstractNLPEvaluator) = nlp.hessian_lagrangian_sparsity


text: include("Planet.jl")

max_stars_repo_name: trontrytel/CLIMA

max_issues_repo_name: trontrytel/CLIMA

max_forks_repo_name: trontrytel/CLIMA


text:
module AutomotiveDrivingModels

using Printf
using LinearAlgebra
using Parameters
using StaticArrays
using Distributions
using Reexport
using Random
using DataFrames
using Tricks: static_hasmethod

include(joinpath(@__DIR__, "Vec", "Vec.jl"))
@reexport using .Vec

# Roadways

export CurvePt,
       Curve,
       CurveIndex,
       CurveProjection,
       CURVEINDEX_START,
       get_lerp_time,
       index_closest_to_point,
       get_curve_index,
       curveindex_end

include("roadways/curves.jl")

export
    LaneTag,
    LaneBoundary,
    Lane,
    LaneConnection,
    SpeedLimit,
    RoadSegment,
    Roadway,
    RoadProjection,
    RoadIndex,

    NULL_BOUNDARY,
    NULL_LANETAG,
    NULL_ROADINDEX,

    DEFAULT_SPEED_LIMIT,
    DEFAULT_LANE_WIDTH,

    is_in_exits,
    is_in_entrances,
    is_at_curve_end,
    is_between_segments_lo,
    is_between_segments_hi,
    is_between_segments,
    has_segment,
    has_lanetag,
    has_next,
    has_prev,
    next_lane,
    prev_lane,
    next_lane_point,
    prev_lane_point,
    connect!,
    move_along,
    n_lanes_left,
    leftlane,
    n_lanes_right,
    rightlane,
    lanes,
    lanetags,
    read

include("roadways/roadways.jl")

export
    Frenet,
    NULL_FRENET

include("roadways/frenet.jl")

export
        gen_straight_curve,
        gen_straight_segment,
        gen_straight_roadway,
        gen_stadium_roadway,
        gen_bezier_curve

include("roadways/roadway_generation.jl")

## agent definitions

export
    AbstractAgentDefinition,
    AgentClass,
    VehicleDef,
    BicycleModel,
    length,
    width,
    class

include("agent-definitions/agent_definitions.jl")

export
    Entity,
    Scene,
    EntityScene,
    capacity,
    id2index,
    get_by_id,
    get_first_available_id,
    posf,
    posg,
    vel,
    velf,
    velg,
    VehicleState,
    CustomVehicleState,
    get_center,
    get_footpoint,
    get_front,
    get_rear,
    get_lane

include("states/entities.jl")
include("states/scenes.jl")
include("states/interface.jl")
include("states/vehicle_state.jl")

## Collision Checkers

export
    ConvexPolygon,
    CPAMemory,
    CollisionCheckResult,
    to_oriented_bounding_box!,
    get_oriented_bounding_box,
    is_colliding,
    is_potentially_colliding,
    get_collision_time,
    get_first_collision,
    is_collision_free,
    get_distance,
    get_edge,
    collision_checker,
    polygon

include("collision-checkers/minkowski.jl")
include("collision-checkers/parallel_axis.jl")

## Feature Extraction

export AbstractFeature,
       EntityFeature,
       SceneFeature,
       TemporalFeature,
       extract_features,
       extract_feature,
       featuretype,

       # provided feature functions
       posgx,
       posgy,
       posgθ,
       posfs,
       posft,
       posfϕ,
       vel,
       velfs,
       velft,
       velgx,
       velgy,
       time_to_crossing_right,
       time_to_crossing_left,
       estimated_time_to_lane_crossing,
       iswaiting,
       iscolliding,
       distance_to,
       acc,
       accfs,
       accft,
       jerk,
       jerkft,
       turn_rate_g,
       turn_rate_f,
       isbraking,
       isaccelerating

include("feature-extraction/features.jl")

export lane_width,
       markerdist_left,
       markerdist_right,
       road_edge_dist_left,
       road_edge_dist_right,
       lane_offset_left,
       lane_offset_right,
       has_lane_left,
       has_lane_right,
       lane_curvature

include("feature-extraction/lane_features.jl")

export
    VehicleTargetPoint,
    VehicleTargetPointFront,
    VehicleTargetPointCenter,
    VehicleTargetPointRear,
    targetpoint_delta,
    find_neighbor,
    NeighborLongitudinalResult,
    FrenetRelativePosition,
    get_frenet_relative_position,
    dist_to_front_neighbor,
    front_neighbor_speed,
    time_to_collision

include("feature-extraction/neighbors_features.jl")

export
    LidarSensor,
    nbeams,
    observe!,
    RoadlineLidarSensor,
    nlanes,
    LanePortion,
    RoadwayLidarCulling,
    ensure_leaf_in_rlc!,
    get_lane_portions

include("feature-extraction/lidar_sensor.jl")

## Actions

export
    propagate,
    EntityAction,
    LaneFollowingAccel,
    AccelTurnrate,
    AccelDesang,
    LatLonAccel,
    AccelSteeringAngle,
    PedestrianLatLonAccel

include("actions/interface.jl")
include("actions/lane_following_accel.jl")
include("actions/accel_turnrate.jl")
include("actions/accel_desang.jl")
include("actions/lat_lon_accel.jl")
include("actions/accel_steering.jl")
include("actions/pedestrian_lat_lon_accel.jl")

export
    DriverModel,
    StaticDriver,
    action_type,
    set_desired_speed!,
    observe!,
    reset_hidden_state!,
    reset_hidden_states!

include("behaviors/interface.jl")

export
    LaneFollowingDriver,
    StaticLaneFollowingDriver,
    PrincetonDriver,
    IntelligentDriverModel,
    ProportionalSpeedTracker,
    track_longitudinal!,
    LateralDriverModel,
    ProportionalLaneTracker,
    LatLonSeparableDriver,
    Tim2DDriver,
    SidewalkPedestrianModel,
    LaneChangeChoice,
    LaneChangeModel,
    get_lane_offset,
    DIR_RIGHT,
    DIR_MIDDLE,
    DIR_LEFT,
    MOBIL,
    TimLaneChanger

include("behaviors/lane_following_drivers.jl")
include("behaviors/princeton_driver.jl")
include("behaviors/speed_trackers.jl")
include("behaviors/intelligent_driver_model.jl")
include("behaviors/lateral_driver_models.jl")
include("behaviors/lane_change_models.jl")
include("behaviors/MOBIL.jl")
include("behaviors/tim_lane_changer.jl")
include("behaviors/lat_lon_separable_driver.jl")
include("behaviors/tim_2d_driver.jl")
include("behaviors/sidewalk_pedestrian_model.jl")

export
    simulate,
    simulate!,
    run_callback,
    CollisionCallback,
    observe_from_history!,
    simulate_from_history!,
    simulate_from_history,
    maximum_entities

include("simulation/simulation.jl")
include("simulation/callbacks.jl")
include("simulation/simulation_from_history.jl")

end # AutomotiveDrivingModels


text:
using Distributed
using Dates
using Statistics
using JLD

num_of_cores_for_cal = 5
if length(workers()) < num_of_cores_for_cal
    addprocs(num_of_cores_for_cal - length(workers()) + 1)
end

@everywhere push!(LOAD_PATH,".")
@everywhere using MBL_ED

print("开始计算：")
ts = now()
println(ts)

the_ws = [i for i in 0.5:0.5:2.0]
append!(the_ws, [i for i in 2.25:0.25:5.0])
append!(the_ws, [5.5, 6.0, 7.0, 8.0, 9.0, 10.0])

r = give_empty_dataframe_for_format()
## 计算部分
# for n in [10]
for n in [10, 12]
    append!(r, primitive_run(tol_n = n, up = div(n, 2), W = the_ws, samples = 10))
end
# k = r[1]
# for i in 2:length(r)
#     append!(k, r[i])
# end
# save("/data/aybai/Dr._Li/fin_MBL_CUDA.jld", "result", k);

print("计算完成：")
tf = now()
println(tf)
println(Dates.value(tf - ts)/1000, "seconds used")

JLD.save("MBL_ED_test.jld", "result", r);


text:
module Losers
using Statistics

include("./dice.jl")
include("./hausdorff.jl")

export
    dice,
    hausdorff

end


text:
# This file includes testset for StaticSystems

@testset "StaticSystems" begin
    @info "Running StaticSystemTestSet ..."

    # StaticSystem construction
    @def_static_system struct Mysystem{IP, OP, RO} <: AbstractStaticSystem
        input::IP = Inport(2)
        output::OP = Outport(3)
        readout::RO = (u, t) -> [u[1] + u[2], u[1] - u[2], u[1] * u[2]]
    end
    # ofunc(u, t) = [u[1] + u[2], u[1] - u[2], u[1] * u[2]]
    ss = Mysystem()
    @test isimmutable(ss)
    @test length(ss.input) == 2
    @test length(ss.output) == 3
    @test typeof(ss.input) == Inport{Inpin{Float64}}
    @test typeof(ss.output) == Outport{Outpin{Float64}}
    @test typeof(ss.trigger) == Inpin{Float64}
    @test typeof(ss.handshake) == Outpin{Bool}

    # ofunc2(u, t) = nothing
    ss = Mysystem(readout=nothing, input=nothing, output=nothing)  # Input or output may be nothing
    @test ss.input === nothing
    @test ss.output === nothing

    # Driving of StaticSystem
    ofunc3(u, t) = u[1] + u[2]
    ss = Mysystem(readout = ofunc3, input=Inport(2), output=Outport(1))
    iport, oport, ipin, opin = Inport(1), Outport(2), Inpin{Bool}(), Outpin()
    connect!(oport, ss.input)
    connect!(ss.output, iport)
    connect!(opin, ss.trigger)
    connect!(ss.handshake, ipin)
    tsk = launch(ss)
    tsk2 = @async while true
        all(take!(iport) .=== NaN) && break
    end
    for t in 1. : 10.
        put!(opin, t)
        put!(oport, [t, t])
        take!(ipin)
    end
    @test content(iport[1].link.buffer) == collect(1 : 10) * 2
    @test !istaskdone(tsk)
    put!(opin, NaN)
    sleep(0.1)
    @test istaskdone(tsk)
    put!(ss.output, [NaN])
    sleep(0.1)
    @test istaskdone(tsk2)

    # Adder tests
    ss = Adder()
    @test isimmutable(ss)
    @test ss.signs == (+, +)
    @test length(ss.output) == 1

    ss = Adder(signs=(+, +, -))
    @test ss.signs == (+, +, -)
    oport, iport, opin, ipin = Outport(3), Inport(1), Outpin(), Inpin{Bool}()
    connect!(opin, ss.trigger)
    connect!(oport, ss.input)
    connect!(ss.handshake, ipin)
    connect!(ss.output, iport)
    tsk = launch(ss)
    tsk2 = @async while true
        all(take!(iport) .=== NaN) && break
    end
    put!(opin, 1.)
    put!(oport, [1, 3, 5])
    take!(ipin)
    @test read(iport[1].link.buffer) == 1 + 3 - 5
    put!(opin, NaN)
    sleep(0.1)
    @test istaskdone(tsk)
    put!(ss.output, [NaN])
    sleep(0.1)
    @test istaskdone(tsk2)

    # Multiplier tests
    ss = Multiplier(ops=(*, *))
    @test isimmutable(ss)
    @test ss.ops == (*, *)
    @test length(ss.output) == 1

    ss = Multiplier(ops=(*, *, /,*))
    @test ss.ops == (*, *, /, *)
    oport, iport, opin, ipin = Outport(4), Inport(1), Outpin(), Inpin{Bool}()
    connect!(opin, ss.trigger)
    connect!(oport, ss.input)
    connect!(ss.handshake, ipin)
    connect!(ss.output, iport)
    tsk = launch(ss)
    tsk2 = @async while true
        all(take!(iport) .=== NaN) && break
    end
    put!(opin, 1.)
    put!(oport, [1, 3, 5, 6])
    take!(ipin)
    @test read(iport[1].link.buffer) == 1 * 3 / 5 * 6
    put!(opin, NaN)
    sleep(0.1)
    @test istaskdone(tsk)
    put!(ss.output, [NaN])
    sleep(0.1)
    @test istaskdone(tsk2)

    # Gain tests
    ss = Gain(input=Inport(3))
    @test isimmutable(ss)
    @test ss.gain == 1.
    @test length(ss.output) == 3
    K = rand(3, 3)
    ss = Gain(input=Inport(3), gain=K)
    oport, iport, opin, ipin = Outport(3), Inport(3), Outpin(), Inpin{Bool}()
    connect!(oport, ss.input)
    connect!(opin, ss.trigger)
    connect!(ss.handshake, ipin)
    connect!(ss.output, iport)
    tsk = launch(ss)
    tsk2 = @async while true
        all(take!(iport) .=== NaN) && break
    end
    u = rand(3)
    put!(opin, 1.)
    put!(oport, u)
    take!(ipin)
    @test [read(pin.link.buffer) for pin in iport] == K * u
    put!(opin, NaN)
    sleep(0.1)
    @test istaskdone(tsk)
    put!(ss.output, NaN * ones(3))
    sleep(0.1)
    @test istaskdone(tsk2)

    # Terminator tests
    ss = Terminator(input=Inport(3))
    @test isimmutable(ss)
    @test ss.readout === nothing
    @test ss.output === nothing
    @test typeof(ss.trigger) == Inpin{Float64}
    @test typeof(ss.handshake) == Outpin{Bool}
    oport, opin, ipin = Outport(3), Outpin(), Inpin{Bool}()
    connect!(oport, ss.input)
    connect!(opin, ss.trigger)
    connect!(ss.handshake, ipin)
    tsk = launch(ss)
    put!(opin, 1.)
    put!(oport, [1., 2., 3.])
    take!(ipin)
    @test [read(pin.link.buffer) for pin in ss.input] == [1., 2., 3.]
    put!(opin, NaN)
    sleep(0.1)
    @test istaskdone(tsk)

    # Memory tests
    ss = Memory(delay=1., numtaps=10, initial=zeros(3))
    @test isimmutable(ss)
    @test size(ss.databuf) == (3, 10)
    @test size(ss.timebuf) == (10,)
    @test mode(ss.databuf) == Cyclic
    @test mode(ss.timebuf) == Cyclic
    @test typeof(ss.trigger) == Inpin{Float64}
    @test typeof(ss.handshake) == Outpin{Bool}
    @test typeof(ss.input) == Inport{Inpin{Float64}}
    @test typeof(ss.output) == Outport{Outpin{Float64}}
    @test outbuf(ss.databuf) == zeros(3, 10)
    oport, iport, opin, ipin = Outport(3), Inport(3), Outpin(), Inpin{Bool}()
    connect!(oport, ss.input)
    connect!(opin, ss.trigger)
    connect!(ss.handshake, ipin)
    connect!(ss.output, iport)
    tsk = launch(ss)
    tsk2 = @async while true
        all(take!(iport) .=== NaN) && break
    end
    put!(opin, 1.)
    put!(oport, [10, 20, 30])
    take!(ipin)
    @test [read(pin.link.buffer) for pin in iport] == zeros(3)
    @test ss.databuf[:, 1] == [10, 20, 30]
    put!(opin, NaN)
    sleep(0.1)
    @test istaskdone(tsk)
    put!(ss.output, NaN * ones(3))
    sleep(0.1)
    @test istaskdone(tsk2)

    # Coupler test
    conmat =  [-1 1; 1 -1]
    cplmat = [1 0 0; 0 0 0; 0 0 0]
    ss = Coupler(conmat=conmat, cplmat=cplmat)
    @test isimmutable(ss)
    @test typeof(ss.trigger) == Inpin{Float64}
    @test typeof(ss.handshake) == Outpin{Bool}
    @test typeof(ss.input) == Inport{Inpin{Float64}}
    @test typeof(ss.output) == Outport{Outpin{Float64}}
    @test length(ss.input) == 6
    @test length(ss.output) == 6
    oport, iport, opin, ipin = Outport(6), Inport(6), Outpin(), Inpin{Bool}()
    connect!(oport, ss.input)
    connect!(opin, ss.trigger)
    connect!(ss.handshake, ipin)
    connect!(ss.output, iport)
    tsk = launch(ss)
    tsk2 = @async while true
        all(take!(iport) .=== NaN) && break
    end
    put!(opin, 1.)
    u = rand(6)
    put!(oport, u)
    take!(ipin)
    @test [read(pin.link.buffer) for pin in iport] == kron(conmat, cplmat) * u
    put!(opin, NaN)
    sleep(0.1)
    @test istaskdone(tsk)
    put!(ss.output, NaN * ones(6))
    sleep(0.1)
    @test istaskdone(tsk2)

    # Test defining new statik systems
    # The type must be a subtype of AbstractStaticSystem
    @test_throws Exception @eval @def_static_system struct MyStaticSystem{RO, OP}
        reaout::RO = (u, t) -> u
        output::OP = Outport()
    end

    # The type must be a subtype of AbstractStaticSystem
    @test_throws Exception @eval @def_static_system struct MyStaticSystem{RO, OP} <: MyDummyAbstractStaticSystem
        reaout::RO = (u, t) -> u
        output::OP = Outport()
    end

    @info "Done StaticSystemTestSet ..."
end # testset


text:
abstract type AbstractIntegratorSPARK{dType, tType} <: DeterministicIntegrator{dType, tType} end
abstract type AbstractIntegratorHSPARK{dType, tType} <: AbstractIntegratorSPARK{dType, tType} end
abstract type AbstractIntegratorVSPARK{dType, tType} <: AbstractIntegratorSPARK{dType, tType} end


text:
using ProgressMeter
using JLD2
using ArgParse
using Random
using Statistics
using LinearAlgebra

using Resampling
using Resampling: TabularRL
using Resampling.Environments

import StatsBase

function measure_variance_sample_martha_time_norm(iterations, batch_size, ER, WER;
                                                  rng=Random.GLOBAL_RNG,
                                                  chain_size=10,
                                                  random_value=false,
                                                  value_function=nothing,
                                                  norm_l=2)

    sum_is = IWER.total(WER.sumtree)
    avg_is = sum_is/size(WER)[1]
    if value_function == nothing
        value_function = TabularRL.StateValueFunction(chain_size)
        if random_value
            value_function.values = randn(rng, chain_size)
        else
            fill!(value_function.values, 0.5)
        end
    end

    update_error = Dict{String, Array{Float64}}()

    update_error["IS"] = zeros(iterations)
    update_error["IR"] = zeros(iterations)
    update_error["VTrace_1.0"] = zeros(iterations)
    update_error["BC-IR"] = zeros(iterations)
    update_error["WIS-Batch"] = zeros(iterations)
    update_error["WIS-Buffer"] = zeros(iterations)
    update_error["WIS-Buffersum"] = zeros(iterations)

    # sample
    for i in 1:iterations
        s_er = sample(ER, batch_size; rng=rng)
        s_wer = sample(WER, batch_size; rng=rng)
        batch_is_sum = sum(s_er[:,9])

        # Note: Experience order
        # (state_t[1], state_tp1[1], action, reward, gamma, terminal,
        #   mu[action + 1], pi[action + 1], pi[action + 1]/mu[action + 1])

        args_er = (s_er[:, 9], s_er[:,1], s_er[:,2], s_er[:,4], s_er[:,5], s_er[:,6])
        args_er = (s_er[:, 9], s_er[:,1], s_er[:,2], s_er[:,4], s_er[:,5], s_er[:,6])
        args_wer = (ones(batch_size), s_wer[:,1], s_wer[:,2], s_wer[:,4], s_wer[:,5], s_wer[:,6])

        update_error["IS"][i] = norm(mean(tdupdate(value_function, args_er...)), norm_l)
        update_error["VTrace_1.0"][i] = norm(mean(tdupdate_vtrace(1.0, value_function, args_er...)), norm_l)
        update_error["IR"][i] = norm(mean(tdupdate(value_function, args_wer...)), norm_l)
        update_error["BC-IR"][i] = norm(mean((1.0/avg_is) .* tdupdate(value_function, args_wer...)), norm_l)
        update_error["WIS-Buffersum"][i] = norm(sum((1.0/sum_is) .* tdupdate(value_function, args_er...)), norm_l)
        update_error["WIS-Buffer"][i] = norm(mean((size(ER)[1]/sum_is) .* tdupdate(value_function, args_er...)), norm_l)
        update_error["WIS-Batch"][i] = norm(sum((1.0/(sum(s_er[:,end]))) .* tdupdate(value_function, args_er...)), norm_l)
    end

    update_variance = Dict{String, Float64}()

    for key in keys(update_error)
        update_variance[key] = var(update_error[key])
    end

    return update_variance

end

function add_res_to_dict!(num_runs, r, res_dict, res)
    if length(keys(res_dict)) == 0
            # res_dict = typeof(res)()
        for k in keys(res)
            res_dict[k] = zeros(num_runs)
            res_dict[k][r] = res[k]
        end
    else
        for k in keys(res)
            res_dict[k][r] = res[k]
        end
    end
end

function create_ret_dict(num_runs, res_dict)

    ret_dict = Dict{String, Array{Float64,1}}()
    for k in keys(res_dict)
        ret_dict[k] = [mean(res_dict[k]), std(res_dict[k])/sqrt(num_runs)]
    end
    return ret_dict
end

function measure_variance_over_time_norm(num_runs, buffer_size, iterations, batch_size, mu, pi; measure_steps=[1000], gamma=0.9, chain_size=10, verbose=false, kwargs...)

    func = measure_variance_sample_martha_time_norm
    truth = MarkovChain.DynamicProgramming(MarkovChain.environment(10), pi, gamma)
    res_dict = Dict{Int64, Dict{String, Array{Float64, 1}}}()
    for step in measure_steps
        res_dict[step] = Dict{String, Array{Float64, 1}}()
    end

    @showprogress 0.1 "Runs: " for r in 1:num_runs
    # for r in 1:num_runs
        rng=Random.MersenneTwister(r)
        ER, WER = get_markov_chain_experience(buffer_size, mu, pi; rng=rng)
        sum_is = IWER.total(WER.sumtree)
        algo = BatchTD(1.0, buffer_size, 1.0/sum_is)
        value_function = TabularRL.StateValueFunction(chain_size)
        fill!(value_function.values, 0.0)

        for iter in 0:maximum(measure_steps)
            if iter in measure_steps
                res = func(iterations, batch_size, ER, WER;
                           rng=rng, value_function=value_function,
                           chain_size=chain_size, kwargs...)
                add_res_to_dict!(num_runs, r, res_dict[iter], res)
            end
            s_opt = IWER.getrow(ER.buffer, :)
            TabularRL.update!(value_function, algo,
                              s_opt[:, 9], s_opt[:, 1],
                              s_opt[:, 2], s_opt[:, 4],
                              s_opt[:, 5], s_opt[:, 6])
            if verbose
                println(iter, " ", rmse(value_function, pi, truth)[1])
            end
        end
    end

    ret_dict = Dict{Int64, Any}()
    for step in measure_steps
        ret_dict[step] = create_ret_dict(num_runs, res_dict[step])
    end

    time_series_dict = Dict{String, Array{Float64, 1}}()
    time_series_stderr_dict = Dict{String, Array{Float64, 1}}()
    for k in keys(ret_dict[measure_steps[1]])
        time_series_dict[k] = [ret_dict[s][k][1] for s in measure_steps]
        time_series_stderr_dict[k] = [ret_dict[s][k][2] for s in measure_steps]
    end

    return ret_dict, time_series_dict, time_series_stderr_dict
end


text:
export check_dir
function check_dir(dir::String)
    if !isdir(dir)
        mkpath(dir)
    end
end

export rm_dir
function rm_dir(dir::String)
    if !isdir(dir) return end
    dir = abspath(dir)
    a = readdir(dir)
    rm_files = [".DS_Store",]
    for irf in rm_files
        if irf in a
            rm(joinpath(dir, irf))
            filter!(x->x != irf, a)
        end
    end
    if isempty(a)
        rm(dir)
        rm_dir(dirname(dir))
    end
end

export add_lockfile
function add_lockfile(filename::String, sleep_time::Float64 = 16.0)
    count = 0
    # sleep_time = 16.9
    p = dirname(filename)
    # @show filename, p
    isdir(p) || mkpath(p)
    while true
        if isfile(filename)
            print(string("\r $count exist! wait $sleep_time s. ", filename, ))
            sleep(sleep_time)
            count += 1
        else
            # run(`touch $filename`)
            touch(filename)
            break
        end
    end
end

export rm_lockfile
function rm_lockfile(filename::String, rm_path = true)
    if isfile(filename)
        # run(`rm $filename`)
        try
            rm(filename)
        catch
            @warn "deleted? $filename"
        end
    end
    if rm_path
        rm_dir(dirname(filename))
    end
end