text:
import Base: +
import LinearAlgebra: tr
const BlockIndices = Vector{<:AbstractVector{Int}}

struct BlockDiagPDMat{T, V<:BlockIndices, PD<:AbstractPDMat{T}} <: AbstractPDMat{T}
    blockindices::V
    blockPD::Vector{PD}
end
function BlockDiagPDMat(blockindices::BlockIndices)
    blockPD = [alloc_cK(length(ind)) for ind in blockindices]
    return BlockDiagPDMat(blockindices, blockPD)
end
size(a::BlockDiagPDMat) = (sum(length,a.blockindices), sum(length,a.blockindices))
size(a::BlockDiagPDMat, d::Int) = sum(length,a.blockindices)
function \(a::BlockDiagPDMat, x::AbstractVector)
    out = similar(x)
    for (pd,ind) in zip(a.blockPD,a.blockindices)
        out[ind] = pd \ x[ind]
    end
    return out
end
function \(a::BlockDiagPDMat, x::AbstractMatrix)
    out = similar(x)
    for (pd,ind) in zip(a.blockPD,a.blockindices)
        out[ind,:] = pd \ x[ind,:]
    end
    return out
end
\(a::PDMat, x::UniformScaling) = a.chol \ Diagonal{Float64}(x, size(a,1))
\(a::PDMat, x::Diagonal) = a.chol \ x
function \(a::BlockDiagPDMat, x::Diagonal)
    out = zeros(size(x)) # unfortunate! should maybe be block diagonal matrix.
    for (pd,ind) in zip(a.blockPD,a.blockindices)
        out[ind,ind] = pd \ Diagonal(x.diag[ind])
    end
    return out
end
logdet(a::BlockDiagPDMat) = sum(logdet, a.blockPD)
tr(pd::AbstractPDMat) = tr(mat(pd))
tr(a::BlockDiagPDMat) = sum(tr, a.blockPD)
"""
    trinv(pd::AbstractPDMat)

Trace of the inverse of a positive definite matrix.
"""
function trinv(pd::AbstractPDMat)
    return tr(pd \ I)
    # Lk = inv(cholfactors(pd))
    # return dot(Lk, Lk)
end
"""
    trinv(a::BlockDiagPDMat)

Trace of the inverse of a block diagonal positive definite matrix.

This is obtained as the sum of the traces of the inverse of each block.
"""
trinv(a::BlockDiagPDMat) = sum(trinv, a.blockPD)
function add!(a::AbstractMatrix, b::BlockDiagPDMat)
    for (pd,ind) in zip(a.blockPD,a.blockindices)
        block = @view(a[ind,ind])
        block[:,:] += mat(pd)
    end
    return a
end
+(a::AbstractMatrix, b::BlockDiagPDMat) = add!(copy(a), b)
function Base.Matrix(a::BlockDiagPDMat)
    K = zeros(size(a))
    for (pd,ind) in zip(a.blockPD,a.blockindices)
        K[ind, ind] = mat(pd)
    end
    return K # + Diagonal(repeat([1e-10], size(K, 1)))
end
mat(a::BlockDiagPDMat) = Matrix(a)

"""
    Positive Definite Matrix for Full Scale Approximation.
"""
mutable struct FullScalePDMat{T,M<:AbstractMatrix,PD<:AbstractPDMat{T},M2<:AbstractMatrix{T},PD2<:BlockDiagPDMat} <: SparsePDMat{T}
    inducing::M
    ΣQR_PD::PD
    Kuu::PD
    Kuf::M2
    Λ::PD2
end
size(a::FullScalePDMat) = (size(a.Kuf,2), size(a.Kuf,2))
size(a::FullScalePDMat, d::Int) = size(a.Kuf,2)
"""
    We have
        Σ ≈ Kuf' Kuu⁻¹ Kuf + Λ
    where Λ is a diagonal matrix (here given by σ²I + diag(Kff - Qff))
    The Woodbury matrix identity gives
        (A+UCV)⁻¹ = A⁻¹ - A⁻¹ U(C⁻¹ + V A⁻¹ U)⁻¹ V A⁻¹
    which we use here with
        A ← Λ
        U = Kuf'
        V = Kuf
        C = Kuu⁻¹
    which gives
        Σ⁻¹ = Λ⁻¹ - Λ⁻¹ Kuf'(Kuu + Kuf Λ⁻¹ Kuf')⁻¹ Kuf Λ⁻¹
            = Λ⁻¹ - Λ⁻¹ Kuf'(        ΣQR       )⁻¹ Kuf Λ⁻¹
"""
function \(a::FullScalePDMat, x)
    Lk = whiten(a.ΣQR_PD, a.Kuf)
    return a.Λ \ (x .- Lk' * (Lk * (a.Λ \ x)))
end

"""
    trinvAB(A::FullScalePDMat, B::Diagonal)

    Computes tr(A⁻¹ B) efficiently under the FSA approximation:
        Σ ≈ Kuf' Kuu⁻¹ Kuf + Λ

    Derivation:
        tr(Σ⁻¹ B) = tr[ (Λ⁻¹ - Λ⁻² Kuf'(Kuu + Kuf Λ⁻¹ Kuf')⁻¹ Kuf) B ]
                  = tr(Λ⁻¹ B) - tr( Λ⁻¹ Kuf' ΣQR⁻¹ Kuf Λ⁻¹ B )
                  = tr(Λ⁻¹ B) - tr( Λ⁻¹ Kuf' ΣQR^(-1/2) ΣQR^(-1/2) Kuf Λ⁻¹ B )
                                                        ╰────────────────╯
                                                            ≡ L'
                  = tr(Λ⁻¹ B) - tr(L*L'*B)
                  = tr(Λ⁻¹ B) - dot(L, B*L)

    See also: [`\`](@ref).
"""
function trinvAB(A::FullScalePDMat, B::Diagonal)
    Λ = A.Λ
    ΣPDinvKfu = convert(Matrix{Float64}, whiten(A.ΣQR_PD, A.Kuf)')
    L = Λ \ ΣPDinvKfu # is there a better way?
    return tr(Λ \ B) - dot(L, B * L)
end

"""
    The matrix determinant lemma states that
        logdet(A+UWV') = logdet(W⁻¹ + V'A⁻¹U) + logdet(W) + logdet(A)
    So for
        Σ ≈ Kuf' Kuu⁻¹ Kuf + Λ
        logdet(Σ) = logdet(Kuu + Kuf Λ⁻¹ Kuf')       + logdet(Kuu⁻¹) + logdet(Λ)
                  = logdet(        ΣQR             ) - logdet(Kuu)   + logdet(Λ)
"""
logdet(a::FullScalePDMat) = logdet(a.ΣQR_PD) - logdet(a.Kuu) + logdet(a.Λ)#sum(log.(a.Λ))
function Base.Matrix(a::FullScalePDMat)
    Lk = whiten(a.Kuu, a.Kuf)
    Σ = Lk'Lk + mat(a.Λ)
    nobs = size(Σ,1)
    return Σ
end

function wrap_cK(cK::FullScalePDMat, inducing, ΣQR_PD, Kuu, Kuf, Λ::BlockDiagPDMat)
    FullScalePDMat(inducing, ΣQR_PD, Kuu, Kuf, Λ)
end
"""
    tr(a::FullScalePDMat)

    Trace of the FSA approximation to the covariance matrix:

    tr(Σ) = tr(Kuf' Kuu⁻¹ Kuf + Λ)
          = tr(Kuf' Kuu⁻¹ Kuf) + tr(Λ)
          = tr(Kuf' Kuu^{-1/2} Kuu^{-1/2} Kuf) + tr(Λ)
                              ╰──────────────╯
                                 ≡  Lk
          = dot(Lk, Lk) + sum(diag(Λ))
"""
function LinearAlgebra.tr(a::FullScalePDMat)
    Lk = whiten(a.Kuu, a.Kuf)
    return tr(a.Λ) + dot(Lk, Lk)
end

"""
    Fully Independent Training Conditional (FSA) covariance strategy.
"""
struct FullScaleApproxStrat{M<:AbstractMatrix, V<:BlockIndices} <: SparseStrategy
    inducing::M
    blockindices::V
end
SubsetOfRegsStrategy(fsa::FullScaleApproxStrat) = SubsetOfRegsStrategy(fsa.inducing)
DeterminTrainCondStrat(fsa::FullScaleApproxStrat) = DeterminTrainCondStrat(fsa.inducing)

function alloc_cK(covstrat::FullScaleApproxStrat, nobs)
    # The objects that need to be allocated are very similar
    # to SoR, so we'll use that as a starting point:
    SoR = SubsetOfRegsStrategy(covstrat)
    cK_SoR = alloc_cK(SoR, nobs)
    # Additionally, we need to store the diagonal corrections.
    Λ = BlockDiagPDMat(covstrat.blockindices)

    cK_fsa = FullScalePDMat(
        covstrat.inducing,
        cK_SoR.ΣQR_PD,
        cK_SoR.Kuu,
        cK_SoR.Kuf,
        Λ)
    return cK_fsa
end
function update_cK!(cK::FullScalePDMat, X::AbstractMatrix, kernel::Kernel,
                    logNoise::Real, kerneldata::KernelData, covstrat::FullScaleApproxStrat)
    inducing = covstrat.inducing
    blockindices = covstrat.blockindices

    Kuu = cK.Kuu
    Kuubuffer = mat(Kuu)
    cov!(Kuubuffer, kernel, inducing)
    Kuubuffer, chol = make_posdef!(Kuubuffer, cholfactors(cK.Kuu))
    Kuu_PD = wrap_cK(cK.Kuu, Kuubuffer, chol)
    Kuf = cov!(cK.Kuf, kernel, inducing, X)
    Kfu = Kuf'

    dim, nobs = size(X)
    Λ = cK.Λ
    @assert Λ.blockindices == blockindices
    for (pd,ind) in zip(Λ.blockPD,Λ.blockindices)
        Xblock = @view(X[:,ind])
        Kblock = cov(kernel, Xblock)
        Qchol = whiten(Kuu_PD, Kuf[:,ind])

        blockbuffer = mat(pd)
        # TODO: could use memory more efficiently
        blockbuffer[:,:] = Kblock - Qchol'Qchol + exp(2*logNoise)*I
        blockbuffer, chol = make_posdef!(blockbuffer, cholfactors(pd))
    end

    ΣQR = Kuf * (Λ \ Kfu) + Kuu
    LinearAlgebra.copytri!(ΣQR, 'U')

    ΣQR, chol = make_posdef!(ΣQR, cholfactors(cK.ΣQR_PD))
    ΣQR_PD = wrap_cK(cK.ΣQR_PD, ΣQR, chol)
    return wrap_cK(cK, inducing, ΣQR_PD, Kuu_PD, Kuf, Λ)
end

#==========================================
  Log-likelihood gradients
===========================================#
function init_precompute(covstrat::FullScaleApproxStrat, X, y, kernel)
    # here we can re-use the Subset of Regressors pre-computations
    SoR = SubsetOfRegsStrategy(covstrat)
    return init_precompute(SoR, X, y, kernel)
end

"""
dmll_kern!(dmll::AbstractVector, kernel::Kernel, X::AbstractMatrix, cK::AbstractPDMat, kerneldata::KernelData,
                    alpha::AbstractVector, Kuu, Kuf, Kuu⁻¹Kuf, Kuu⁻¹KufΣ⁻¹y, Σ⁻¹Kfu, ∂Kuu, ∂Kuf,
                    covstrat::FullScaleApproxStrat)
Derivative of the log likelihood under the Fully Independent Training Conditional (fsa) approximation.

Helpful reference: Vanhatalo, Jarno, and Aki Vehtari.
                   "Sparse log Gaussian processes via MCMC for spatial epidemiology."
                   In Gaussian processes in practice, pp. 73-89. 2007.

Generally, for a multivariate normal with zero mean
    ∂logp(Y|θ) = 1/2 y' Σ⁻¹ ∂Σ Σ⁻¹ y - 1/2 tr(Σ⁻¹ ∂Σ)
                    ╰───────────────╯     ╰──────────╯
                           `V`                 `T`

where Σ = Kff + σ²I.

Notation: `f` is the observations, `u` is the inducing points.
          ∂X stands for ∂X/∂θ, where θ is the kernel hyperparameters.

In the case of the FSA approximation, we have
    Σ = Λ + Qff
    The second component gives
    ∂(Qff) = ∂(Kfu Kuu⁻¹ Kuf)
    which is used by the gradient function for the subset of regressors approximation,
    and so I don't repeat it here.

    for ∂Λ we have (with `i` indexing each block)
    Λi = Ki - Qi + σ²I
       = Ki - Kui' Kuu⁻¹ Kui + σ²
    ∂Λi = ∂Ki - Kui' ∂(Kuu⁻¹) Kui - 2 ∂Kui' Kuu⁻¹ Kui
        = ∂Ki + Kui' Kuu⁻¹ ∂(Kuu) Kuu⁻¹ Kui - 2 ∂Kui' Kuu⁻¹ Kui
"""
function dmll_kern!(dmll::AbstractVector, kernel::Kernel, X::AbstractMatrix,
                    cK::AbstractPDMat, kerneldata::SparseKernelData,
                    alpha::AbstractVector, Kuu, Kuf, Kuu⁻¹Kuf, Kuu⁻¹KufΣ⁻¹y, Σ⁻¹Kfu, ∂Kuu, ∂Kuf,
                    covstrat::FullScaleApproxStrat)
    # first compute the SoR component
    SoR = SubsetOfRegsStrategy(covstrat)
    dmll_kern!(dmll, kernel, X, cK, kerneldata, alpha,
               Kuu, Kuf, Kuu⁻¹Kuf, Kuu⁻¹KufΣ⁻¹y, Σ⁻¹Kfu, ∂Kuu, ∂Kuf,
               SoR)
    nparams = num_params(kernel)
    dim, nobs = size(X)
    inducing = covstrat.inducing
    Λ = cK.Λ
    for iparam in 1:nparams
        # TODO: the grad_slice! calls here are redundant with the ones in the dmll_kern! call above
        grad_slice!(∂Kuu, kernel, inducing, inducing, kerneldata.Kuu, iparam)
        grad_slice!(∂Kuf, kernel, inducing, X       , kerneldata.Kux1, iparam)
        V, T = 0.0, 0.0
        # ∂Λdense = zeros(nobs, nobs)
        for (pd,ind) in zip(Λ.blockPD,Λ.blockindices)
            Kuu⁻¹Kui = Kuu⁻¹Kuf[:,ind]
            ∂Ki = similar(mat(pd))
            Kui, ∂Kui = Kuf[:,ind], ∂Kuf[:,ind]
            Xi = X[:,ind]
            grad_slice!(∂Ki, kernel, Xi, Xi, EmptyData(), iparam)
            ∂Λi = (∂Ki
                  + Kuu⁻¹Kui' * ∂Kuu * Kuu⁻¹Kui
                  - 2 * ∂Kui' * Kuu⁻¹Kui)
            V += dot(alpha[ind], ∂Λi*alpha[ind])

            # see trinvAB to understand next 3 lines
            ΣPDinvKiu = convert(Matrix{Float64}, whiten(cK.ΣQR_PD, Kui)')
            L = pd \ ΣPDinvKiu
            T += trinvAB(pd, ∂Λi) - dot(L, ∂Λi*L)
            # ∂Λdense[ind,ind] = ∂Λi
        end
        # # FOR DEBUG ONLY
        # Talt = tr(cK \ ∂Λdense) # inefficient
        # Valt = dot(alpha, ∂Λdense * alpha)
        # @show T, Talt
        # @show V, Valt
        # # END DEBUG
        dmll[iparam] += (V-T)/2
    end
    return dmll
end
function dmll_kern!(dmll::AbstractVector, gp::GPBase, precomp::SoRPrecompute, covstrat::FullScaleApproxStrat)
    return dmll_kern!(dmll, gp.kernel, gp.x, gp.cK, gp.data, gp.alpha,
                      gp.cK.Kuu, gp.cK.Kuf,
                      precomp.Kuu⁻¹Kuf, precomp.Kuu⁻¹KufΣ⁻¹y, precomp.Σ⁻¹Kfu,
                      precomp.∂Kuu, precomp.∂Kuf,
                      covstrat)
end

"""
    dmll_noise(gp::GPE, precomp::SoRPrecompute, covstrat::FullScaleApproxStrat)

∂logp(Y|θ) = 1/2 y' Σ⁻¹ ∂Σ Σ⁻¹ y - 1/2 tr(Σ⁻¹ ∂Σ)

∂Σ = I for derivative wrt σ², so
∂logp(Y|θ) = 1/2 y' Σ⁻¹ Σ⁻¹ y - 1/2 tr(Σ⁻¹)
            = 1/2[ dot(α,α) - tr(Σ⁻¹) ]

We have:
    Σ⁻¹ = Λ⁻¹ - Λ⁻² Kuf'(        ΣQR       )⁻¹ Kuf
Use the identity tr(A'A) = dot(A,A) to get:
    tr(Σ⁻¹) = tr(Λ⁻¹) - dot(Lk, Lk) .
where
    Lk ≡ ΣQR^(-1/2) Kuf Λ⁻¹
"""
function dmll_noise(gp::GPE, precomp::SoRPrecompute, covstrat::FullScaleApproxStrat)
    cK = gp.cK
    Λ = cK.Λ
    Lk = Λ \ whiten(cK.ΣQR_PD, cK.Kuf)'
    # # DEBUG
    # noiseT = sum(1 ./ Λ) - dot(Lk, Lk)
    # nobs = gp.nobs
    # Talt = tr(cK \ Matrix(1.0*I, nobs, nobs))
    # @show noiseT, Talt # should be same
    # # END DEBUG
    # @show tr(inv(Matrix(cK)))
    # @show trinv(Λ) - dot(Lk, Lk)
    return exp(2*gp.logNoise) * ( # Jacobian
        dot(gp.alpha, gp.alpha)
        - trinv(Λ) # sum(1 ./ Λ)
        + dot(Lk, Lk)
        )
end

function get_alpha_u(Ktrain::FullScalePDMat, xtrain::AbstractMatrix, ytrain::AbstractVector, meanf::Mean)
    ΣQR_PD = Ktrain.ΣQR_PD
    Kuf = Ktrain.Kuf
    meantrain = mean(meanf, xtrain)
    Λ = Ktrain.Λ
    alpha_u = ΣQR_PD \ (Kuf * (Λ \ (ytrain-meantrain)) )
    return alpha_u
end

"""
predictMVN(xpred::AbstractMatrix, xtrain::AbstractMatrix, ytrain::AbstractVector,
                    kernel::Kernel, meanf::Mean, logNoise::Real,
                    alpha::AbstractVector,
                    covstrat::FullScaleApproxStrat, Ktrain::FullScalePDMat)
    See Quiñonero-Candela and Rasmussen 2005, equations 24b.
    Some derivations can be found below that are not spelled out in the paper.

    Notation: Qab = Kau Kuu⁻¹ Kub
              ΣQR = Kuu + σ⁻² Kuf Kuf'

              x: prediction (test) locations
              f: training (observed) locations
              u: inducing point locations

    We have
        Σ ≈ Kuf' Kuu⁻¹ Kuf + Λ
    By Woodbury
        Σ⁻¹ = Λ⁻¹ - Λ⁻² Kuf'(Kuu + Kuf Λ⁻¹ Kuf')⁻¹ Kuf
            = Λ⁻¹ - Λ⁻² Kuf'(        ΣQR       )⁻¹ Kuf

    The predictive mean can be derived (assuming zero mean function for simplicity)
    μ = (Qxf+Λxf) (Qff + Λff)⁻¹ y
      = Qxf (Qff + Λff)⁻¹ y   + Λxf (Qff + Λff)⁻¹ y
        ╰─────────────────╯
            same as FITC
      = Kxu ΣQR⁻¹ Kuf Λff⁻¹ y + Λxf (Qff + Λff)⁻¹ y
           ╰────────────────╯       ╰─────────────╯
              ≡ alpha_u                ≡ alpha

    Similarly for the posterior predictive covariance:
    Σ = Σxx - (Qxf+Λxf) (Qff + Λff)⁻¹ (Qxf+Λxf)'
      = Σxx - Kxu ΣQR⁻¹ Kuf Λ⁻¹ Qxf'                # substituting result from μ
      = Σxx - Kxu ΣQR⁻¹  Kuf Λ⁻¹ Kfu Kuu⁻¹ Kux      # definition of Qxf
      = Σxx - Kxu ΣQR⁻¹ (ΣQR - Kuu) Kuu⁻¹ Kux       # using definition of ΣQR
      = Σxx - Kxu Kuu⁻¹ Kux + Kxu ΣQR⁻¹ Kux         # expanding
      = Σxx - Qxx           + Kxu ΣQR⁻¹ Kux         # definition of Qxx
"""
function predictMVN(xpred::AbstractMatrix, blockindpred::BlockIndices,
                    xtrain::AbstractMatrix, ytrain::AbstractVector,
                    kernel::Kernel, meanf::Mean,
                    alpha::AbstractVector,
                    covstrat::FullScaleApproxStrat, Ktrain::FullScalePDMat)
    blockindtrain = covstrat.blockindices
    ΣQR_PD = Ktrain.ΣQR_PD
    inducing = covstrat.inducing

    Kux = cov(kernel, inducing, xpred)

    nx = size(xpred,2)
    nf = size(xtrain,2)
    # TODO: Λxf is fairly sparse: it should be possible to avoid its construction
    Λxf = zeros(nx, nf)
    for (predblock, trainblock) in zip(blockindpred, blockindtrain)
        Xpredblock, Xtrainblock = xpred[:,predblock], xtrain[:,trainblock]
        sparseblockdata = SparseKernelData(kernel, inducing, Xpredblock, Xtrainblock)
        denseblockdata = KernelData(kernel, Xpredblock, Xtrainblock)
        Kxf_block = cov(kernel, Xpredblock, Xtrainblock, denseblockdata)
        Qxf_block = getQab(Ktrain, kernel, Xpredblock, Xtrainblock, sparseblockdata)
        Λxf[predblock,trainblock] = Kxf_block - Qxf_block
    end

    meanx = mean(meanf, xpred)
    alpha_u = get_alpha_u(Ktrain, xtrain, ytrain, meanf)
    mupred = meanx + (Kux' * alpha_u) + Λxf * alpha

    sparsedata = SparseKernelData(kernel, inducing, xpred, xtrain)
    Qxf = getQab(Ktrain, kernel, xpred, xtrain, sparsedata)
    ΛplusQxf = Qxf + Λxf
    Σxx = cov(kernel, xpred, xpred)
    Σ_FSA = Σxx - ΛplusQxf * (Ktrain \ ΛplusQxf')
    return mupred, Σ_FSA
end

"""
    predict_f(gp::GPBase, X::Matrix{Float64}[; full_cov::Bool = false])

Return posterior mean and variance of the Gaussian Process `gp` at specfic points which are
given as columns of matrix `X`. If `full_cov` is `true`, the full covariance matrix is
returned instead of only variances.
"""
function predict_f(gp::GPBase, x::AbstractMatrix, blockindpred::BlockIndices; full_cov::Bool=false)
    size(x,1) == gp.dim || throw(ArgumentError("Gaussian Process object and input observations do not have consistent dimensions"))
    if full_cov
        return predict_full(gp, x, blockindpred)
    else
        ## Calculate prediction for each point independently
        μ = Array{eltype(x)}(undef, size(x,2))
        σ2 = similar(μ)
        iblock = Vector{Int}(undef, size(x,2))
        for (j,predblock) in enumerate(blockindpred)
            iblock[predblock] = j
        end
        numblocks = length(blockindpred)
        for k in 1:size(x,2)
            subblockind = [j==iblock[k] ? [1] : [] for j in 1:numblocks]
            m, sig = predict_full(gp, x[:,k:k], subblockind)
            μ[k] = m[1]
            σ2[k] = max(diag(sig)[1], 0.0)
        end
        return μ, σ2
    end
end
predict_full(gp::GPE, xpred::AbstractMatrix, blockindpred::BlockIndices) = predictMVN(xpred, blockindpred, gp.x, gp.y, gp.kernel, gp.mean, gp.alpha, gp.covstrat, gp.cK)

function FSA(x::AbstractMatrix, inducing::AbstractMatrix, blockindices::BlockIndices,
             y::AbstractVector, mean::Mean, kernel::Kernel, logNoise::Real)
    covstrat = FullScaleApproxStrat(inducing, blockindices)
    GPE(x, y, mean, kernel, logNoise, covstrat)
end


text:
#################################################################################
#Type Alias for long type signatures
const MinMax = NamedTuple{(:min, :max), NTuple{2, Float64}}
const NamedMinMax = Tuple{String, MinMax}
const UpDown = NamedTuple{(:up, :down), NTuple{2, Float64}}
const InOut = NamedTuple{(:in, :out), NTuple{2, Float64}}

# Type Alias From other Packages
const PM = PowerModels
const PSY = PowerSystems
const PSI = PowerSimulations
const IS = InfrastructureSystems
const MOI = MathOptInterface
const MOIU = MathOptInterface.Utilities
const PJ = ParameterJuMP
const MOPFM = MOI.FileFormats.Model
const TS = TimeSeries

const BUILD_SIMULATION_TIMER = TimerOutputs.TimerOutput()
const RUN_SIMULATION_TIMER = TimerOutputs.TimerOutput()

#Type Alias for JuMP and PJ containers
const JuMPExpressionMatrix = Matrix{<:JuMP.AbstractJuMPScalar}
const PGAE{V} =
    PJ.ParametrizedGenericAffExpr{Float64, V} where {V <: JuMP.AbstractVariableRef}
const GAE{V} = JuMP.GenericAffExpr{Float64, V} where {V <: JuMP.AbstractVariableRef}
const JuMPAffineExpressionArray = Matrix{GAE{V}} where {V <: JuMP.AbstractVariableRef}
const JuMPAffineExpressionVector = Vector{GAE{V}} where {V <: JuMP.AbstractVariableRef}
const JuMPConstraintArray = JuMP.Containers.DenseAxisArray{JuMP.ConstraintRef}
const JuMPParamArray = JuMP.Containers.DenseAxisArray{PJ.ParameterRef}
const DenseAxisArrayContainer = Dict{Symbol, JuMP.Containers.DenseAxisArray}

@enum BUILD_STATUS begin
    BUILT = 1
    IN_PROGRESS = -1
    EMPTY = 0
end

# Settings constants
const UNSET_HORIZON = 0
const UNSET_INI_TIME = Dates.DateTime(0)

# Tolerance of comparisons
# MIP gap tolerances in most solvers are set to 1e-4
const ABSOLUTE_TOLERANCE = 1.0e-3
const SLACK_COST = 1e6
const COST_EPSILON = 1e-3
const MISSING_INITIAL_CONDITIONS_TIME_COUNT = 999.0

const OPERATIONS_ACCEPTED_KWARGS = [
    :horizon,
    :initial_time,
    :use_forecast_data,
    :PTDF,
    :use_parameters,
    :optimizer,
    :warm_start,
    :slack_variables,
    :system_to_file,
    :constraint_duals,
]

const OPERATIONS_SOLVE_KWARGS = [:optimizer, :save_path]

const STAGE_ACCEPTED_KWARGS =
    [:PTDF, :warm_start, :slack_variables, :constraint_duals, :system_to_file]

const PSI_NAME_DELIMITER = "__"

# The constants below are strings instead of enums because there is a requirement that users
# should be able to define their own without changing PowerSimulations.

# Variables / Parameters
const ACTIVE_POWER = "P"
const ENERGY = "E"
const ENERGY_BUDGET = "energy_budget"
const FLOW_ACTIVE_POWER = "Fp"
const ON = "On"
const REACTIVE_POWER = "Q"
const ACTIVE_POWER_IN = "Pin"
const ACTIVE_POWER_OUT = "Pout"
const RESERVE = "R"
const SERVICE_REQUIREMENT = "service_requirement"
const START = "Start"
const STOP = "Stop"
const THETA = "theta"
const VM = "Vm"
const INFLOW = "In"
const SPILLAGE = "Sp"
const SLACK_UP = "γ⁺"
const SLACK_DN = "γ⁻"

# Constraints
const ACTIVE = "active"
const ACTIVE_RANGE = "activerange"
const ACTIVE_RANGE_LB = "activerange_lb"
const ACTIVE_RANGE_UB = "activerange_ub"
const COMMITMENT = "commitment"
const DURATION = "duration"
const DURATION_DOWN = "duration_dn"
const DURATION_UP = "duration_up"
const ENERGY_CAPACITY = "energy_capacity"
const ENERGY_LIMIT = "energy_limit"
const FEEDFORWARD = "FF"
const FEEDFORWARD_BIN = "FF_bin"
const FLOW_LIMIT = "FlowLimit"
const FLOW_LIMIT_FROM_TO = "FlowLimitFT"
const FLOW_LIMIT_TO_FROM = "FlowLimitTF"
const FLOW_REACTIVE_POWER_FROM_TO = "FqFT"
const FLOW_REACTIVE_POWER_TO_FROM = "FqTF"
const FLOW_ACTIVE_POWER_FROM_TO = "FpFT"
const FLOW_ACTIVE_POWER_TO_FROM = "FpTF"
const FLOW_ACTIVE_POWER = "Fp"
const FLOW_REACTIVE_POWER = "Fq"
const INPUT_POWER_RANGE = "inputpower_range"
const OUTPUT_POWER_RANGE = "outputpower_range"
const RAMP = "ramp"
const RAMP_DOWN = "ramp_dn"
const RAMP_UP = "ramp_up"
const RATE_LIMIT = "RateLimit"
const RATE_LIMIT_FT = "RateLimitFT"
const RATE_LIMIT_TF = "RateLimitTF"
const REACTIVE = "reactive"
const REACTIVE_RANGE = "reactiverange"
const REQUIREMENT = "requirement"
const INFLOW_RANGE = "inflowrange"

ext: jl


text:
export Sparsity, NoSparsity, VariableSparsity, MonomialSparsity, SignSymmetry
abstract type Sparsity end
struct NoSparsity <: Sparsity end

include("xor_space.jl")
include("sign.jl")
include("variable_sparsity.jl")
include("monomial_sparsity.jl")

"""
    struct SparsePreorder{S <: Sparsity, C <: AbstractPreorderCertificate} <: AbstractPreorderCertificate
        sparsity::S
        certificate::C
    end

Same certificate as `C` except that the Sum-of-Squares polynomials `σ_i`
are modelled as a sum of Sum-of-Squares polynomials with smaller basis
using the sparsity reduction `sparsity`.
"""
struct SparsePreorder{S <: Sparsity, C <: AbstractPreorderCertificate} <: AbstractPreorderCertificate
    sparsity::S
    certificate::C
end

get(certificate::SparsePreorder, attr::Cone) = get(certificate.certificate, attr)

struct SparseDomain{S, P, B}
    domain::S
    processed::P
    bases::Vector{Vector{B}}
end

function get(certificate::SparsePreorder, attr::PreprocessedDomain, domain::BasicSemialgebraicSet, p)
    basis, preorder_bases = sparsity(p, domain, certificate.sparsity, certificate.certificate)
    return SparseDomain(domain, get(certificate.certificate, attr, domain, p), preorder_bases)
end

function get(certificate::SparsePreorder, attr::PreorderIndices, domain::SparseDomain)
    return get(certificate.certificate, attr, domain.processed)
end

function get(::SparsePreorder, ::MultiplierBasis, index::PreorderIndex, domain::SparseDomain)
    return domain.bases[index.value]
end
function get(::Type{SparsePreorder{S, C}}, attr::MultiplierBasisType) where {S, C}
    return Vector{get(C, attr)}
end

function get(certificate::SparsePreorder, attr::Generator, index::PreorderIndex, domain::SparseDomain)
    return get(certificate.certificate, attr, index, domain.processed)
end

get(certificate::SparsePreorder, attr::IdealCertificate) = SparseIdeal(certificate.sparsity, get(certificate.certificate, attr))
get(::Type{<:SparsePreorder{S, C}}, attr::IdealCertificate) where {S, C} = SparseIdeal{S, get(C, attr)}

SumOfSquares.matrix_cone_type(::Type{SparsePreorder{S, C}}) where {S, C} = SumOfSquares.matrix_cone_type(C)

"""
    struct SparseIdeal{S <: Sparsity, C <: AbstractIdealCertificate} <: AbstractIdealCertificate
        sparsity::S
        certificate::C
    end

Same certificate as `C` except that the Sum-of-Squares polynomial `σ`
is modelled as a sum of Sum-of-Squares polynomials with smaller basis
using the sparsity reduction `sparsity`.
"""
struct SparseIdeal{S <: Sparsity, C <: AbstractIdealCertificate} <: AbstractIdealCertificate
    sparsity::S
    certificate::C
end

function SparseIdeal(sp::VariableSparsity, cone, basis, maxdegree::Nothing, newton_polytope)
    error("`maxdegree` cannot be `nothing` when `sparsity` is `VariableSparsity`.")
end
function SparseIdeal(sp::VariableSparsity, cone, basis, maxdegree::Integer, newton_polytope)
    return SparseIdeal(sp, MaxDegree(cone, basis, maxdegree))
end
function SparseIdeal(sp::Union{MonomialSparsity, SignSymmetry}, cone, basis, maxdegree, newton_polytope)
    return SparseIdeal(sp, Newton(cone, basis, newton_polytope))
end

function sparsity(poly::MP.AbstractPolynomial, ::VariableSparsity, certificate::MaxDegree)
    H, cliques = chordal_csp_graph(poly, FullSpace())
    return map(cliques) do clique
        return maxdegree_gram_basis(certificate.basis, clique, certificate.maxdegree)
    end
end
function sparsity(monos, sp::Union{SignSymmetry, MonomialSparsity}, gram_basis::MB.MonomialBasis)
    return MB.MonomialBasis.(sparsity(monos, sp, gram_basis.monomials))
end
function sparsity(poly::MP.AbstractPolynomial, sp::Union{SignSymmetry, MonomialSparsity}, certificate::AbstractIdealCertificate)
    return sparsity(monomials(poly), sp, get(certificate, GramBasis(), poly))
end
function get(certificate::SparseIdeal, ::GramBasis, poly)
    return sparsity(poly, certificate.sparsity, certificate.certificate)
end
function get(::Type{SparseIdeal{S, C}}, attr::GramBasisType) where {S, C}
    return Vector{<:get(C, attr)}
end
function get(certificate::SparseIdeal, attr::ReducedPolynomial, poly, domain)
    return get(certificate.certificate, attr, poly, domain)
end
function get(certificate::SparseIdeal, attr::Cone)
    return get(certificate.certificate, attr)
end
SumOfSquares.matrix_cone_type(::Type{SparseIdeal{S, C}}) where {S, C} = SumOfSquares.matrix_cone_type(C)

zero_basis(certificate::SparseIdeal) = zero_basis(certificate.certificate)
zero_basis_type(::Type{SparseIdeal{S, C}}) where {S, C} = zero_basis_type(C)


text:
function tdc(T::SymTridiagonal)

    # Tridiagonal divide-and-conquer using arrowhead matrices
    # COMPUTES: eigenvectors U and eigenvalues E of a SymTridiagonal matrix T,
    # T = U*diagm(E)*U'
    # RETURNS: U, E

    n=length(T.dv)
    # set the tolerances for eig
    tols=[1e2,1e2,1e2,1e2,1e2]
    if n==1
        U=[1.0]
        E=T.dv[1]
    elseif n==2
        U,E,rest=eig(SymArrow([T.dv[1]],[T.ev[1]],T.dv[2],2),tols)
    else
        k=round(Int,floor(n/2))
        T1=SymTridiagonal(T.dv[1:k],T.ev[1:k-1])
        T2=SymTridiagonal(T.dv[k+2:end],T.ev[k+2:end])
        U1,E1=tdc(T1)
        U2,E2=tdc(T2)
        # form arrowhead
        D=[E1;E2]
        z=[U1'[:,end]*T.ev[k]; U2'[:,1]*T.ev[k+1]]
        a=T.dv[k+1]
        A=SymArrow(D,z,a,k+1)
        U,E,rest=eig(A,tols)
        U[1:k,:]=U1*U[1:k,:]
        U[k+2:end,:]=U2*U[k+2:end,:]
    end
    return U,E
end

# roots
function rootsah(pol::Union{Poly{Float32},Poly{Float64},Poly{Int32},Poly{Int64}}, D::Vector{Float64})

    # COMPUTES: the roots of polynomials with all distinct real roots.
    # The computation is forward stable. The program uses SymArrow (arrowhead) companion matrix and
    # corresponding eig routine
    # D are barycentric coordinates - elements od D must interpolate the roots of P,
    # for example
    #              D=roots(polyder(pol))
    # RETURNS: roots E

    T = Float64

    # tols = [1e2,1e2,1e2,1e2,1e2] or similar is the vector of tolerances for eig
    tols=[1e2,1e2,1e2,1e2,1e2]

    Dm=map(T,D)
    Dm=sort(Dm,rev=true)
    D=sort(D,rev=true)

    p=map(T,[pol[i] for i=0:1:length(pol)-1])
    p=p[end:-1:1]

    # Compute z of the arrowhead
    # First we compute values s=p(D) we use Horner sheme with Double
    n=length(p)-1
    pD=map(Double,p)
    DD=map(Double,Dm)
    oneD=Double(one(T))
    s=Array(Double{T},n-1)
    for i=1:n-1
        s[i]=oneD
    end
    s=pD[1]*s
    for i=2:n+1
        r=s.*DD
        s=r+pD[i]
    end

    # Compute t's
    t=Array(Double{T},n-1)
    for j=1:n-1
        h=oneD
        for i=1:j-1
            g=DD[j]-DD[i]
            h=h*g
        end
        for i=j+1:n-1
            g=DD[j]-DD[i]
            h=h*g
        end
        t[j]=h
    end
    # Compute alphaD
    alphaD=pD[2]/pD[1]
    for i=1:n-1
        alphaD+=DD[i]
    end
    # alphaD=-alphaD
    alphaD=Double(-alphaD.hi,-alphaD.lo)
    #  Compute z
    zD=Array(Double{T},n-1)
    for i=1:n-1
        #         zD[i]=sqrt((-s[i])/(t[i]*pD[1]))
        zD[i]=sqrt(Double(-s[i].hi,-s[i].lo)/(t[i]*pD[1]))
    end

    z=Array(Float64,n-1)
    for i=1:n-1
        z[i]=Float64(zD[i].lo+zD[i].hi)
    end
    alpha=Float64(alphaD.hi+alphaD.lo)
    A=SymArrow(D,z,alpha,n)

    E=zeros(Float64,n)
    Qout=zeros(Int64,n)
    for k=1:n
        E[k], Qout[k] = eig( A,zD,alphaD,k,tols )
    end

    E, Qout
end

#------------------------
#--------  Functions

#-------- Inverses

function inv{T}(A::SymArrow{T},zD::Vector{Double{T}}, alphaD::Double{T},i::Int64,tols::Vector{Float64})

    # COMPUTES: inverse of a SymArrow matrix A, inv(A-A.D[i]*I) which is again SymArrow
    # uses higher precision to compute top of the arrow element accurately, if
    # needed.
    # tols=[tolb,tolz] are tolerances, usually [1e3, 10*n]
    # [0.0,0.0] forces DoubleDouble, [1e50,1e50] would never use it
    # RETURNS:  SymArrow(D,z,b,i), Kb, Kz, Qout
    # Kb - condition Kb, Kz - condition Kz, Qout = 1 / 0 - double was / was not used

    n=length(A.D)
    D=Array(T,n)
    z=Array(T,n)
    wz=1/A.z[i]
    shift=A.D[i]

    for k=1:i-1
        D[k]=one(T)/(A.D[k]-shift)
        z[k]=-A.z[k]*D[k]*wz
    end
    for k=i+1:n
        D[k-1]=one(T)/(A.D[k]-shift)
        z[k-1]=-A.z[k]*D[k-1]*wz
    end

    # D=[1./(A.D[1:i-1]-shift),1./(A.D[i+1:end]-shift)]
    # wz=1/A.z[i]
    # z=-[A.z[1:i-1], A.z[i+1:end]].*D1*wz1
    a=A.a-A.D[i]

    # compute the sum in a plain loop
    P=zero(T)
    Q=zero(T)
    Qout=0

    nn=n-1
    for k=1:i-1
        D[k]>0.0 ? P+=A.z[k]^2*D[k] : Q+=A.z[k]^2*D[k]
    end
    for k=i:nn
        D[k]>0.0 ? P+=A.z[k+1]^2*D[k] : Q+=A.z[k+1]^2*D[k]
    end
    a<0 ? P=P-a : Q=Q-a

    Kb=(P-Q)/abs(P+Q)
    Kz=maxabs(A.z)*abs(wz)

    if Kb<tols[1] ||  Kz<tols[2]
        b=(P+Q)*wz*wz
    elseif Kb==Inf
       b=0.0
    else  # recompute in Double
        Qout=1
        AD=map(Double{T},A.D)
        shiftd=AD[i]
        #shiftd=map(Double,A.D[i])
        Dd=[Double{T}[AD[k]-shiftd for k=1:i-1];
            Double{T}[AD[k]-shiftd for
                            k=i+1:length(A.D)]]
        wzd=zD[i] # Double(A.z[i])
        ad=alphaD-shiftd

        Pd,Qd=map(Double{T},(0.0,0.0))

        for k=1:i-1
            Dd[k].hi>0.0 ? Pd+=zD[k]^2/Dd[k] :
            Qd+=zD[k]^2/Dd[k]
        end

        for k=i:nn
            Dd[k].hi>0.0 ? Pd+=zD[k+1]^2/Dd[k] :
            Qd+=zD[k+1]^2/Dd[k]
        end

        ad.hi<0 ?   Pd=Pd-ad : Qd=Qd-ad

        bd=(Pd+Qd)/(wzd*wzd)
        b=Float64(bd.hi+bd.lo)
    end

    if i<A.i
        # SymArrow([D[1:A.i-2],zero(T),D[A.i-1:end]],[z[1:A.i-2],wz,z[A.i-1:end]],b,i),Kb,Kz,Qout
        # This reduces memory allocation (4 vectors) and is 30-40% faster
        for k=n-1:-1:A.i-1
            D[k+1]=D[k]
            z[k+1]=z[k]
        end
        D[A.i-1]=zero(T)
        z[A.i-1]=wz
        SymArrow(D,z,b,i),Kb,Kz,Qout
    else
        # SymArrow([D[1:A.i-1],zero(T),D[A.i:end]],[z[1:A.i-1],wz,z[A.i:end]],b,i+1),Kb,Kz,Qout
        for k=n-1:-1:A.i
            D[k+1]=D[k]
            z[k+1]=z[k]
        end
        D[A.i]=zero(T)
        z[A.i]=wz
        SymArrow(D,z,b,i+1),Kb,Kz,Qout
    end

end # inv

function inv{T}(A::SymArrow{T}, zD::Vector{Double{T}}, alphaD::Double{T}, shift::Float64, tolr::Float64)

    # COMPUTES: inverse of the shifted SymArrow A, inv(A-shift*I) which is SymDPR1
    # uses DoubleDouble to compute rho accurately, if needed.
    # tolr is tolerance, usually 1e3,  0.0 forces Double, 1e50 would never use it
    # RETURNS: SymDPR1(D,u,rho), Krho, Qout
    # Krho - condition Krho, Qout = 1 / 0 - Double was / was not used

    n=length(A.D)
    D=Array(T,n+1)
    u=Array(T,n+1)

    # Ds=A.D-shift

    # D=[1./Ds[1:A.i-1];0.0;1./Ds[A.i:end]];
    # u=[A.z[1:A.i-1]./Ds[1:A.i-1];-1.0;A.z[A.i:end]./Ds[A.i:end]];

    for k=1:A.i-1
        D[k]=one(T)/(A.D[k]-shift)
        u[k]=A.z[k]*D[k]
    end
    D[A.i]=zero(T)
    u[A.i]=-one(T)
    for k=A.i:n
        D[k+1]=one(T)/(A.D[k]-shift)
        u[k+1]=A.z[k]*D[k+1]
    end

    # compute rho and Krho
    #--- compute the sum in a plain loop

    a=A.a-shift
    P=zero(T)
    Q=zero(T)
    Qout=0

    for k=1:A.i-1
        D[k]>0.0 ? P=P+A.z[k]^2*D[k] : Q=Q+A.z[k]^2*D[k]
    end
    for k=A.i:n
        D[k+1]>0.0 ? P=P+A.z[k]^2*D[k+1] : Q=Q+A.z[k]^2*D[k+1]
    end

    a<0 ? P=P-a : Q=Q-a

    # Condition of rho
    Krho=(P-Q)/abs(P+Q)

    if Krho<tolr
        rho=-1.0/(P+Q)

    else  # recompute in Double
        Qout=1
        Pd,Qd=map(Double{T},(0.0,0.0))
        AD=map(Double{T},A.D)
        shiftd=map(Double{T},shift)
        ad=alphaD-shiftd
        #   Dd=[Double(A.D[k])-shiftd for k=1:length(A.D) ]
        #   Dd=map(Double,A.D)-shiftd
        for k=1:A.i-1
            D[k]>0.0 ? Pd=Pd+zD[k]^2/(AD[k]-shiftd) : Qd=Qd+zD[k]^2/(AD[k]-shiftd)
        end
        for k=A.i:n
            D[k+1]>0.0 ? Pd=Pd+zD[k]^2/(AD[k]-shiftd) : Qd=Qd+zD[k]^2/(A.D[k]-shiftd)
        end

        # for k=1:length(A.D)
        #    Dd[k].hi>0.0 ? Pd=Pd+Double(A.z[k])^2/Dd[k] : Qd=Qd+Double(A.z[k])^2/Dd[k]
        # end
        ad.hi+ad.lo<0 ? Pd=Pd-ad : Qd=Qd-ad
        r=map(Double{T},1.0)/(Pd+Qd)
        rho=Float64(-(r.hi+r.lo))
    end

    # returns the following
    SymDPR1(D,u,rho), Krho, Qout

end # inv

function inv{T}(A::SymArrow{T}, zD::Vector{Double{T}}, alphaD::Double{T}, shift::Double{Float64})

    # COMPUTES: inverse of the shifted SymArrow A, inv(A-shift*I), which is a SymDPR1
    # here shift is Double so it uses Double to compute everything
    # RETURNS: SymDPR1(D1,u1,rho), Qout
    # Qout = 1 on exit meaning Double was used

    n=length(A.D)
    D=Array(Double,n+1)
    u=Array(Double,n+1)
    shiftd=map(Double{T},shift)

    # D[1:n]=map(Double,A.D)-shift
    for k=1:n
        D[k]=map(Double{T},A.D[k])-shiftd
    end

    u[1:n]=zD
    a=alphaD-shiftd
    i=A.i

    oned=Double(one(T))
    zerod=Double(zero(T))

    for k=1:i-1
        u[k]=u[k]/D[k]
        D[k]=oned/D[k]
    end

    for k=i:n
        u[k+1]=u[k]/D[k]
        D[k+1]=oned/D[k+1]
    end

    D[i]=zerod
    u[i]=-oned

    # compute rho and Krho
    #--- compute the sum in a plain loop

    P,Q=zerod,zerod
    Qout=1

    for k=1:i-1
        # D[k].hi+D[k].lo > 0.0  # .lo cannot change sign, I think
        # D[k].hi > 0.0 ? P=P+Double(A.z[k])^2*D[k] : Q=Q+Double(A.z[k])^2*D[k]
        D[k].hi > 0.0 ? P=P+zD[k]*u[k] : Q=Q+zD[k]*u[k]
    end
    for k=i:n
        D[k+1].hi >0.0 ? P=P+zD[k]*u[k+1] : Q=Q+zD[k]*u[k+1]
    end

    a.hi+a.lo<0.0  ? P=P-a : Q=Q-a
    r=oned/(P+Q)

    rho=Float64(-(r.hi+r.lo))

    # returns the following
    D1=Array(T,n+1)
    u1=Array(T,n+1)
    for k=1:n+1
        D1[k]=Float64(D[k].hi+D[k].lo)
    end
    for k=1:n+1
        u1[k]=Float64(u[k].hi+u[k].lo)
    end

    SymDPR1(D1,u1,rho), Qout

end # inv

function  eig{T}( A::SymArrow{T},zD::Vector{Double{T}},alphaD::Double{T},k::Int64,tols::Vector{Float64})

    # COMPUTES: k-th eigenvalue of an ordered irreducible SymArrow
    # A = [diag (D) z; z' alpha]
    # Specially designed to be used in the polynomial rootfinder rootsah !!!!
    # tols=[tolb,tolz,tolnu,tolrho,tollambda] = [1e3,10.0*n,1e3,1e3,1e3]
    # RETURNS: lambda
    # lambda - k-th eigenvalue
    # Kb, Kz, Knu, Krho - condition numbers
    # Qout = 1 / 0 - Double was / was not used

    # If Qout>0, quad precision was used
    # i was the shift index

    # Set the dimension
    n = length(A.D) + 1

    # Set all conditions initially to zero
    Kb,Kz,Knu,Krho=0.0,0.0,0.0,0.0
    Qout=0
    v=zeros(T,n)
    # Kz is former kappa_nu

    # Determine the shift sigma, the shift index i, and whether lambda
    # is on the left or the right side of the nearest pole
    # Exterior eigenvalues (k = 1 or k = n):

    if k == 1
        sigma,i,side = A.D[1],1,'R'
    elseif k==n
        sigma,i,side = A.D[n-1],n-1,'L'
    else
        # Interior eigenvalues (k in (2,...n-1) ):
        # We need to compute this in Double
        Dd=map(Double,A.D)
        # Dtemp = Dd-Dd[k]
        # atemp = alphaD-Dd[k]
        # middle = Dtemp[k-1]/2.0
        # Fmiddle = (atemp-middle)-sum(A.z.^2./(Dtemp-middle))
        middle=(Dd[k-1]-Dd[k])/Double(2.0)
        Fmiddle=Double(zero(T))
        for l=1:n-1
            Fmiddle+=zD[l]^2/((Dd[l]-Dd[k])-middle)
        end
        Fmiddle=((alphaD-Dd[k])-middle)-Fmiddle
        sigma,i,side = Fmiddle.hi+Fmiddle.lo < 0.0 ? (A.D[k],k,'R') : (A.D[k-1],k-1,'L')
    end

    # Compute the inverse of the shifted matrix, A_i^(-1), Kb and Kz

    Ainv,Kb,Kz,Qout = inv(A,zD,alphaD,i,tols[1:2])

    # Compute the eigenvalue of the inverse shifted matrix

    nu = bisect( Ainv,side )

    #  nu=fastzero([invD1; 0; invD2], [w1;wz;w2], b, side); # implement later
    #  [nu-nu1]/nu, [nu-nueva]/nu, pause  # and compare

    if abs(nu)==Inf
        # this is nonstandard
        # Deflation in aheig (nu=Inf)
        v[i]=1.0
        lambda=sigma
    else
        # standard case, full computation
        # nu1 is the F- or 1-norm of the inverse of the shifted matrix
        # nu10=maximum([sum(abs(Ainv.z))+abs(Ainv.a), maximum(abs(Ainv.D)+abs(Ainv.z))])
        nu1=0.0
        for l=1:n-1
            nu1=max(nu1,abs(Ainv.D[l])+abs(Ainv.z[l]))
        end
        nu1=max(sumabs(Ainv.z)+abs(Ainv.a), nu1)

        Knu=nu1/abs(nu)
        if Knu<tols[3]
            # Accuracy is fine, compute the eigenvector
            mu = 1.0/nu
            # v=[ A.z./((A.D-sigma)-mu);-1.0]
            for l=1:n-1
            	v[l] = A.z[l]/((A.D[l]-sigma)-mu)
            end
            v[n]=-1.0
            lambda, v = mu+sigma, v/norm(v)
        else
            # Remedies according to Remark 3 - we shift between original
            # eigenvalues and compute DPR1 matrix
            # 1/nu1+sigma, 1/nu+sigma
            # println(k, " Remedy 3 ")
            nu = side=='R' ? abs(nu) : -abs(nu)
            nu1=-sign(nu)*nu1
            sigma1=(nu1+nu)/(2.0*nu*nu1)+sigma

            if findfirst(A.D-sigma1,0.0)>0 # we came back with a pole
                # recompute sigmav more accurately according with dekker
                sigmav=(Double(nu1)+Double(nu))/(Double(2.0)*Double(nu)*Double(nu1))+Double(sigma)
                # Compute the inverse of the shifted arrowhead (DPR1)
                Ainv,Qout1=inv(A,sigmav) # Ainv is Float64
                nu1=bisect(Ainv,side)
                mu1 = 1.0/nu1
                D0=map(Double,A.D)-sigmav
                D1=[D0[l].hi+D0[l].lo for l=1:length(D0)]
                if findfirst(D1-mu1,0.0)>0
                    ind=find(D1-mu1==0.0)
                    v=zeros(n)
                    v[ind]=1.0
                else
                    v=[ A.z./(D1-mu1);-1.0]
                end
                # Shift the eigenvalue back in Double
                lam = Double(1.0)/Double(nu1)+sigmav
                # Return this
                lambda,v = lam.hi+lam.lo, v/norm(v)
            else
                # Compute the inverse of the shifted arrowhead (DPR1)
                Ainv, Krho,Qout1=Arrowhead.inv(A,sigma1,tols[4]) # Ainv is Float64
                # Compute the eigenvalue by bisect for DPR1
	            # Note: instead of bisection could use dlaed4 (need a wrapper) but
	            # it is not faster. There norm(u)==1
                nu1= bisect(Ainv,side)
                mu1=1.0/nu1
                # standard v
                v=[ A.z./((A.D-sigma1)-mu1);-1.0]
                # Return this - shift the eigenvalue back and normalize the vector
                lambda, v = mu1+sigma1, v/norm(v)
            end
            Qout==1 && (Qout=Qout+2)
        end

        # Remedy according to Remark 1 - we recompute the the eigenvalue
        # near zero from the inverse of the original matrix (a DPR1 matrix).
        if (abs(A.D[i])+abs(1.0/nu))/abs(lambda)>tols[5]

            if k==1 && A.D[1]<0.0 || k==n && A.D[n-1]>0.0 || i<n-1 && side=='L' && sign(A.D[i])+sign(A.D[i+1])==0 ||
                side=='R' && sign(A.D[i])+sign(A.D[i-1])==0
                # println(k, " Remedy 1 ")
                # Compute the inverse of the original arrowhead (DPR1)
                Ainv,Krho,Qout1 = inv(A,0.0,tols[4]) # Ainv is Float64
                Qout==1 && (Qout=Qout+4)
                if abs(Ainv.r)==Inf
                    lambda=0.0
                else
                    # Here we do not need bisection. We compute the Rayleigh
                    # quotient by using already computed vectors which is
                    # componentwise accurate
                    nu1=sum(v.^2.*Ainv.D)+Ainv.r*sum(v.*Ainv.u)^2;
                    lambda=1.0/nu1
                end
            end

        end
    end

    # Return this
    lambda, Qout

end # eig (k)


text:
# ---
# title: 1240. Tiling a Rectangle with the Fewest Squares
# id: problem1240
# author: Tian Jun
# date: 2020-10-31
# difficulty: Hard
# categories: Dynamic Programming, Backtracking
# link: <https://leetcode.com/problems/tiling-a-rectangle-with-the-fewest-squares/description/>
# hidden: true
# ---
#
# Given a rectangle of size `n` x `m`, find the minimum number of integer-sided
# squares that tile the rectangle.
#
#
#
# **Example 1:**
#
# ![](https://assets.leetcode.com/uploads/2019/10/17/sample_11_1592.png)
#
#
#
#     Input: n = 2, m = 3
#     Output: 3
#     Explanation: 3 squares are necessary to cover the rectangle.
#     2 (squares of 1x1)
#     1 (square of 2x2)
#
# **Example 2:**
#
# ![](https://assets.leetcode.com/uploads/2019/10/17/sample_22_1592.png)
#
#
#
#     Input: n = 5, m = 8
#     Output: 5
#
#
# **Example 3:**
#
# ![](https://assets.leetcode.com/uploads/2019/10/17/sample_33_1592.png)
#
#
#
#     Input: n = 11, m = 13
#     Output: 6
#
#
#
#
# **Constraints:**
#
#   * `1 <= n <= 13`
#   * `1 <= m <= 13`
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


text:
function randindex(
    obs::IndexObjects.IndexVector, upper::Int, sample_size::Int, config::Config
)::Tuple{IndexObjects.IndexVector, Vector{Float64}}

    weights = ForgettingWeights.forgetting_weights(
        length(obs.v), config.linear_forgetting
    )
    counts = Bincounts.bincount(obs.v, weights, upper)
    # -- add in some prior pseudocounts
    pseudocounts = counts .+ config.prior_weight
    probabilities = pseudocounts / sum(pseudocounts)
    post = Delayed.categoricalindex(probabilities, sample_size)
    return post, probabilities
end


text:
# # Heights_1 problem

# We estimate simple linear regression model with a half-T prior.

using DynamicHMCModels

ProjDir = @__DIR__
cd(ProjDir)

# Import the dataset.

data = DataFrame(CSV.read(joinpath("..", "..", "data", "Howell1.csv"), delim=';'));

# Use only adults and standardize

df = filter(row -> row[:age] >= 18, data);

# Half-T for `σ`, see below.

Base.@kwdef mutable struct Heights_1{Ty <: AbstractVector, Tν <: Real}
    "Observations."
    y::Ty
    "Degrees of freedom for prior on sigma."
    v::Tν
end;

# Write a function to return properly dimensioned transformation.

function make_transformation(model::Heights_1)
    as((σ = asℝ₊, μ  = as(Real, 100, 250)), )
end

model = Heights_1(;y = df[:, :height], v=1.0)

# Then make the type callable with the parameters *as a single argument*.

function (model::Heights_1)(θ)
    @unpack y, v = model   # extract the data
    @unpack μ, σ = θ
    loglikelihood(Normal(μ, σ), y) + logpdf(TDist(v), σ)
end;

# Wrap the problem with a transformation, then use Flux for the gradient.

P = TransformedLogDensity(make_transformation(model), model)
∇P = ADgradient(:ForwardDiff, P);

# Tune and sample.

results = mcmc_with_warmup(Random.GLOBAL_RNG, ∇P, 1000)
posterior = P.transformation.(results.chain)

DynamicHMC.Diagnostics.EBFMI(results.tree_statistics) |> display
println()

DynamicHMC.Diagnostics.summarize_tree_statistics(results.tree_statistics) |> display
println()

a3d = Array{Float64, 3}(undef, 1000, 2, 1);
for j in 1:1
  for i in 1:1000
    a3d[i, 1, j] = values(posterior[i].μ)
    a3d[i, 2, j] = values(posterior[i].σ)
  end
end

pnames = ["μ", "σ"]
sections =   Dict(
  :parameters => pnames,
)
chns = create_mcmcchains(a3d, pnames, sections, start=1);

# Stan.jl results

cmdstan_result = "
Iterations = 1:1000
Thinning interval = 1
Chains = 1,2,3,4
Samples per chain = 1000

Empirical Posterior Estimates:
          Mean        SD       Naive SE      MCSE      ESS
sigma   7.7641872 0.29928194 0.004732063 0.0055677898 1000
   mu 154.6055177 0.41989355 0.006639100 0.0085038356 1000

Quantiles:
         2.5%      25.0%       50.0%      75.0%       97.5%
sigma   7.21853   7.5560625   7.751355   7.9566775   8.410391
   mu 153.77992 154.3157500 154.602000 154.8820000 155.431000
";

describe(chns)

# end of m4.1d.jl


text:
module CubeIO
using ..Cubes
using ..DAT
using ..Proc

import NetCDF.ncread, NetCDF.ncclose
import Base.Iterators
import ..Cubes: saveCube, check_overwrite, getsavefolder, cubechunks
import ..Cubes.ESDLZarr: ZArrayCube

function toPointAxis(aout,ain,loninds,latinds)
  iout = 1
  for (ilon,ilat) in zip(loninds,latinds)
    aout[iout]=ain[ilon,ilat]
    iout+=1
  end
end

"""
    extractLonLats(c::AbstractCubeData,pl::Matrix)

Extracts a list of longitude/latitude coordinates from a data cube. The coordinates
are specified through the matrix `pl` where `size(pl)==(N,2)` and N is the number
of extracted coordinates. Returns a data cube without `LonAxis` and `LatAxis` but with a
`SpatialPointAxis` containing the input locations.
"""
function extractLonLats(c::AbstractCubeData,pl::Matrix;kwargs...)
  size(pl,2)==2 || error("Coordinate list must have exactly 2 columns")
  axlist=caxes(c)
  ilon=findAxis(LonAxis,axlist)
  ilat=findAxis(LatAxis,axlist)
  ilon>0 || error("Input cube must contain a LonAxis")
  ilat>0 || error("input cube must contain a LatAxis")
  lonax=axlist[ilon]
  latax=axlist[ilat]
  pointax = SpatialPointAxis([(pl[i,1],pl[i,2]) for i in 1:size(pl,1)])
  loninds = map(ll->axVal2Index(lonax,ll[1]),pointax.values)
  latinds = map(ll->axVal2Index(latax,ll[2]),pointax.values)
  incubes=InDims("Lon","Lat")
  outcubes=OutDims(pointax)
  y=mapCube(toPointAxis,c,loninds,latinds;indims=incubes,outdims=outcubes,max_cache=1e8,kwargs...)
end
export extractLonLats

using NetCDF
import Base.Iterators: product
function writefun(xout,xin::AbstractArray{Union{Missing,T}},a,nd,cont_loop,filename;kwargs...) where T

  x = map(ix->ismissing(ix) ? convert(T,-9999.0) : ix,xin)

  count_vec = fill(-1,nd)
  start_vec = fill(1,nd)

  used_syms = Symbol[]
  for (k,v) in cont_loop
    count_vec[k] = 1
    start_vec[k] = a[Symbol(v)][1]
    push!(used_syms,Symbol(v))
  end

  splinds = Iterators.filter(i->!in(i,used_syms),keys(a))
  vn = join(string.([a[s][2] for s in splinds]),"_")
  isempty(vn) && (vn="layer")
  ncwrite(x,filename,vn,start=start_vec,count=count_vec)
end

"""
    exportcube(r::AbstractCubeData,filename::String)

Saves a cube object to a portable NetCDF file in `filename`.

When saving, every RangeAxis will be converted to an axis in the NetCDF cube,
while every categorical axis will be represented by a different variable
inside the resulting file. Dimensions will be ordered according to the
`priorities` keyword argument, which defaults to `Dict("LON"=>1,"LAT"=>2,"TIME"=>3)`,
which means that the file will be stored with longitudes varying fastest.
"""
function exportcube(r::AbstractCubeData,filename::String;priorities = Dict("LON"=>1,"LAT"=>2,"TIME"=>3), proj=epsg4326)

  ax = caxes(r)
  ax_cont = collect(filter(i->isa(i,RangeAxis),ax))
  ax_cat  = filter(i->!isa(i,RangeAxis),ax)
  prir = map(i->get(priorities,uppercase(axname(i)),10),ax_cont)
  ax_cont=ax_cont[sortperm(prir)]
  dims = map(NcDim,ax_cont)
  isempty(ax_cat) && (ax_cat=[VariableAxis(["layer"])])
  it = map(i->i.values,ax_cat)
  elt = Base.nonmissingtype(eltype(r))
  vars = NcVar[NcVar(join(collect(string.(a)),"_"),dims,t=elt,atts=Dict("missing_value"=>convert(elt,-9999.0), "grid_mapping" => proj["grid_mapping_name"])) for a in product(it...)]
  file = NetCDF.create(filename,vars)
  for d in dims
    ncwrite(d.vals,filename,d.name)
  end
  nccreate(filename, proj["grid_mapping_name"])
  ncputatt(filename, proj["grid_mapping_name"], proj)
  dl = map(i->i.dimlen,dims) |> cumprod
  isplit = findfirst(i->i>5e7,dl)
  isplit isa Nothing && (isplit=length(dl))
  incubes = InDims(ax_cont[1:(isplit-1)]...)
  cont_loop = Dict(ii=>axname(ax_cont[ii]) for ii in isplit:length(ax_cont))
  mapCube(writefun,r,length(ax_cont),cont_loop,filename,indims=incubes,include_loopvars=true,ispar=false,max_cache=5e8,
  nthreads=[1])
  NetCDF.close(file)
  nothing
end

import Zarr: NoCompressor

function saveCube(c::AbstractCubeData, name::AbstractString; overwrite = false, chunksize = cubechunks(c), compressor=NoCompressor(), max_cache=1e8)
  allax = caxes(c)
  firstaxes = findall(i->i>1,chunksize)
  lastaxes = setdiff(1:length(allax),firstaxes)
  allax = allax[[firstaxes;lastaxes]]
  dl = cumprod(length.(caxes(c)))
  isplit = findfirst(i->i>max_cache/sizeof(eltype(c)),dl)
  isplit isa Nothing && (isplit=length(dl)+1)
  forcesingle = (isplit+1)<length(firstaxes)
  axn = axname.(allax[1:isplit-1])
  indims = InDims(axn...)
  path = getsavefolder(name)
  check_overwrite(path,overwrite)
  outdims = OutDims(axn..., retcubetype=ZArrayCube,chunksize=chunksize[1:length(axn)], compressor=compressor, path = path)
  if forcesingle
    nprocs()>1 && println("Forcing single core processing because of bad chunk size")
    o = mapCube(copyto!,c,indims=indims, outdims=outdims,ispar=false,max_cache=max_cache)
  else
    o = mapCube(copyto!,c,indims=indims, outdims=outdims,max_cache=max_cache)
  end
end

export exportcube

global const projection = Dict(
"grid_mapping_name" => "transverse_mercator",
"longitude_of_central_meridian" => -9. ,
"false_easting" => 500000. ,
"false_northing" => 0. ,
"latitude_of_projection_origin" => 0. ,
"scale_factor_at_central_meridian" => 0.9996,
"long_name" => "CRS definition",
"longitude_of_prime_meridian" => 0.,
"semi_major_axis" => 6378137.,
"inverse_flattening" => 298.257223563 ,
"spatial_ref" => "PROJCS[\"WGS 84 / UTM zone 29N\",GEOGCS[\"WGS 84\",
DATUM[\"WGS_1984\",SPHEROID[\"WGS 84\",6378137,298.257223563,
AUTHORITY[\"EPSG\",\"7030\"]],AUTHORITY[\"EPSG\",\"6326\"]],
PRIMEM[\"Greenwich\",0,AUTHORITY[\"EPSG\",\"8901\"]],
UNIT[\"degree\",0.0174532925199433,AUTHORITY[\"EPSG\",\"9122\"]],
AUTHORITY[\"EPSG\",\"4326\"]],
PROJECTION[\"Transverse_Mercator\"],
PARAMETER[\"latitude_of_origin\",0],
PARAMETER[\"central_meridian\",-9],
PARAMETER[\"scale_factor\",0.9996],
PARAMETER[\"false_easting\",500000],
PARAMETER[\"false_northing\",0],
UNIT[\"metre\",1,AUTHORITY[\"EPSG\",\"9001\"]],
AXIS[\"Easting\",EAST],
AXIS[\"Northing\",NORTH],
AUTHORITY[\"EPSG\",\"32629\"]]",
"GeoTransform" => "722576.2320803495 20 0 4115483.464603715 0 -20 ")

global const epsg4326=Dict(
  "grid_mapping_name" => "latitude_longitude",
  "long_name" => "CRS definition",
  "longitude_of_prime_meridian" => 0.,
  "semi_major_axis" => 6378137.,
  "inverse_flattening" => 298.257223563,
  "spatial_ref" => "GEOGCS[\"WGS 84\",DATUM[\"WGS_1984\",SPHEROID[\"WGS 84\",6378137,298.257223563,AUTHORITY[\"EPSG\",\"7030\"]],AUTHORITY[\"EPSG\",\"6326\"]],PRIMEM[\"Greenwich\",0],UNIT[\"degree\",0.0174532925199433],AUTHORITY[\"EPSG\",\"4326\"]]",
  "GeoTransform" => "-98.68712696894481 0.0001796630568239077 0 20.69179551612753 0 -0.0001796630568239077 "
)

end


text:
"""
special Wrapper, which is completely peeled of again

Comparing to Identity, Identity{T} results in Identity{T}, while NoEffect{T} results in plain T.
"""
struct NoEffect{T}
  value::T
end

"""
only for internal purposes, captures the still unevaluated part of an Eff
"""
struct Continuation{Fs}
  functions::Fs
  Continuation(functions...) = new{typeof(functions)}(functions)
end

"""
central data structure which can capture Effects in a way that they can interact, while each
is handled independently on its own
"""
struct Eff{Effectful, Fs}
  effectful::Effectful
  cont::Continuation{Fs}

  function Eff(effectful::T, cont::Continuation{Fs}) where {T, Fs}
    if !isempty(cont) && effectful isa NoEffect
      # Evaluate NoEffect directly to make sure, we don't have a chain of NoEffect function accumulating
      # e.g. with ContextManager this could lead to things not being evaluated, while the syntax suggest
      # everything is evaluated, and hence the ContextManager may finalize resource despite they are still used.
      cont(effectful.value)

    else
      # also run this if isempty(cont) to stop infinite recursion
      # (which happens otherwise because empty cont results returns noeffect for convenience)
      new{T, Fs}(effectful, cont)
    end
  end
end
Eff(effectful) = Eff(effectful, Continuation())

function Base.show(io::IO, eff::Eff)
  print(io, "Eff(effectful=$(eff.effectful), length(cont)=$(length(eff.cont.functions)))")
end

"""
mark a value as no effect, but plain value
"""
noeffect(value) = Eff(NoEffect(value))  # everything reduces to the Identity Monad
noeffect(eff::Eff) = eff  # if we find a Eff effect, we just directly use it (in analogy to `effect`)

"""
mark a value as an effect
"""
effect(value) = Eff(value)
effect(eff::Eff) = eff  # if we find a Eff effect, we just directly use it

# Functionalities for Continuation
# --------------------------------

function (c::Continuation)(value)
  if isempty(c)
    noeffect(value)
  else
    first_func = c.functions[1]
    rest = c.functions[2:end]
    eff = first_func(value)
    Eff(eff.effectful, Continuation(eff.cont.functions..., rest...))
  end
end
Base.isempty(c::Continuation) = Base.isempty(c.functions)
Base.map(f, c::Continuation) = Continuation(c.functions..., noeffect ∘ f)

# Functionalities for Eff
# -----------------------

TypeClasses.pure(::Type{<:Eff}, a) = noeffect(a)
TypeClasses.map(f, eff::Eff) = TypeClasses.flatmap(noeffect ∘ f, eff)
TypeClasses.flatmap(f, eff::Eff) = Eff(eff.effectful, Continuation(eff.cont.functions..., f))

max_stars_repo_name: schlichtanders/ExtensibleEffects

max_issues_repo_name: schlichtanders/ExtensibleEffects

max_forks_repo_name: schlichtanders/ExtensibleEffects


text:
# Data streams

abstract SampleStream

# auxiliary functions
gets(x::DenseVector, i::Integer) = x[i]
gets(x::DenseVector, i::UnitRange) = view(x, i)
gets(x::DenseVector, i::AbstractVector) = x[i]

gets(x::DenseMatrix, i::Integer) = view(x, :, i)
gets(x::DenseMatrix, i::UnitRange) = view(x, :, i)
gets(x::DenseMatrix, i::AbstractVector) = x[i]

## GenericSampleSeq

type GenericSampleSeq{XS, YS, Ord} <: SampleStream
    xs::XS
    ys::YS
    ord::Ord

    function GenericSampleSeq(xs::XS, ys::YS, ord::Ord)
        nx = size(xs, ndims(xs))
        ny = size(ys, ndims(ys))
        nx == ny ||
            throw(DimensionMismatch("xs and ys must have the same number of samples."))
        new(xs, ys, ord)
    end
end

start(str::GenericSampleSeq) = start(str.ord)
done(str::GenericSampleSeq, s) = done(str.ord, s)

function next(str::GenericSampleSeq, s)
    i, s = next(str.ord, s)
    x = gets(str.xs, i)
    y = gets(str.ys, i)
    return (x, y), s
end

# constructors

function sample_seq(X::DenseVecOrMat, Y::DenseVecOrMat, ord)
    GenericSampleSeq{typeof(X), typeof(Y), typeof(ord)}(X, Y, ord)
end

function sample_seq(X::DenseVecOrMat, Y::DenseVecOrMat)
    sample_seq(X,Y,1:size(X, ndims(X)))
end

# mini-batch sequences

function batches(n::Int, bsize::Int)
    m = ceil(Int, n / bsize)
    bs = Array(UnitRange{Int}, m)
    last = 0
    for i = 1:m-1
        bs[i] = last+1:last+bsize
        last += bsize
    end
    bs[m] = last+1:n
    return bs
end

minibatch_seq(X::DenseVecOrMat, Y::DenseVecOrMat, bsize::Int) =
    (n = size(X, ndims(X)); sample_seq(X, Y, batches(n, bsize)))

minibatch_seq(X::DenseVecOrMat, Y::DenseVecOrMat, bsize::Int, ord) =
    (n = size(X, ndims(X)); sample_seq(X, Y, batches(n, bsize)[ord]))


text:
module ArenaSim

import HTTP
import JSON
using Parameters
using Distributions: Categorical, Bernoulli, Beta, cdf
using Random: rand!, shuffle

export deckinfo, SimParameters, AccountState, SimOutput, PityTimerSamples

include("constants.jl")
include("stochastic_samples.jl")
include("cycle.jl")
include("types.jl")
include("card_db.jl")
include("deckreaders.jl")
include("collection_representation.jl")
include("utils.jl")
include("collection_evaluation.jl")
include("open_pack.jl")
include("process_icr.jl")
include("openvault.jl")
include("convert_wildcards.jl")
include("duplicate_prevention.jl")
include("starterdecks.jl")
include("simulation.jl")
include("shufflesim.jl")
include("QC_winrate.jl")
include("negbin.jl")
include("icrgen.jl")

end # module

max_stars_repo_name: asieev/ArenaSim

max_issues_repo_name: asieev/ArenaSim

max_forks_repo_name: asieev/ArenaSim


text:
@testset "Ladderise" begin
	tree = Tree(parse_newick("(((A,B),C),(D,(E,F)));"))
	ladderise!(tree)
	@test newick_string(tree) == "((C,(A,B)),(D,(E,F)));"
	ladderise!(tree, :left)
	@test newick_string(tree) == "(((A,B),C),((E,F),D));"
end


text:
# Test the calling interface for parsers

using Compat
using HDF5,JLD,KUparser,KUnet,Base.Test
infoln(x)=(println();info(x))
include("pxyequal.jl")

infoln("Loading data")
@date @load "conll07.tst.jld4"
@show nbatch = 4
@show nbeam = 4
@show ndeps = length(deprel)
@show ft = Flist.acl11eager
@show pt = ArcEager13
@show wt = wtype(corpus)
@show pxytuple = @compat Tuple{Vector{pt}, Vector{Matrix{wt}}, Vector{Matrix{wt}}}
p0 = pt(1,ndeps)
s1 = corpus[1]

infoln("Let's train a reasonable model")
h0 = 20000
net = [Mmul(h0), Bias(), Relu(),
       Mmul(p0.nmove), Bias(), Logp(),
       LogpLoss()]
@date (p,x,y)=oparse(pt, corpus, ndeps, ft)
@date train(net, hcat(x...), hcat(y...))
@date c = corpus[1:16]

infoln("Testing oparser:")
@date o0 = oparse(pt, c, ndeps); @test isa(o0, Vector{pt})
@date o1 = oparse(pt, c, ndeps; usepmap=true); @test isequal(o1,o0)
@date o2 = oparse(pt, c, ndeps, ft); @test isa(o2, pxytuple); @test isequal(o2[1],o0)
@date o3 = oparse(pt, c, ndeps, ft; usepmap=true); @test pxyequal(o3,o2)

infoln("Testing gparser:")
@date g0 = gparse(pt, c, ndeps, ft, net); @test isa(g0, Vector{pt})
@date g1 = gparse(pt, c, ndeps, ft, net; nbatch=nbatch); @test isequal(g1,g0)
@date g2 = gparse(pt, c, ndeps, ft, net; usepmap=true); @test isequal(g2,g0)
@date g3 = gparse(pt, c, ndeps, ft, net; usepmap=true,nbatch=nbatch); @test isequal(g3,g0)
@date g4 = gparse(pt, c, ndeps, ft, net; returnxy=true); @test isa(g4, pxytuple); @test isequal(g4[1],g0)
@date g5 = gparse(pt, c, ndeps, ft, net; returnxy=true,nbatch=nbatch); @test pxyequal(g5,g4)
@date g6 = gparse(pt, c, ndeps, ft, net; returnxy=true,usepmap=true); @test pxyequal(g6,g4)
@date g7 = gparse(pt, c, ndeps, ft, net; returnxy=true,usepmap=true,nbatch=nbatch); @test pxyequal(g7,g4)

infoln("Testing bparser:")
@date bx = bparse(pt, c, ndeps, ft, net, 1); @test isequal(bx, g0)
@date b0 = bparse(pt, c, ndeps, ft, net, nbeam); @test isa(b0, Vector{pt})
@date b1 = bparse(pt, c, ndeps, ft, net, nbeam; nbatch=nbatch); @test isequal(b1,b0)
@date b2 = bparse(pt, c, ndeps, ft, net, nbeam; usepmap=true); @test isequal(b2,b0)
@date b3 = bparse(pt, c, ndeps, ft, net, nbeam; usepmap=true,nbatch=nbatch); @test isequal(b3,b0)
@date b4 = bparse(pt, c, ndeps, ft, net, nbeam; returnxy=true); @test isa(b4, pxytuple)
@date b5 = bparse(pt, c, ndeps, ft, net, nbeam; returnxy=true,nbatch=nbatch); @test pxyequal(b5,b4)
@date b6 = bparse(pt, c, ndeps, ft, net, nbeam; returnxy=true,usepmap=true); @test pxyequal(b6,b4)
@date b7 = bparse(pt, c, ndeps, ft, net, nbeam; returnxy=true,usepmap=true,nbatch=nbatch); @test pxyequal(b7,b4)

:ok


text:
###
### **arb_hypgeom.h** -- hypergeometric functions of real variables
###

### Generalized hypergeometric function
arbcall"void arb_hypgeom_pfq(arb_t res, arb_srcptr a, slong p, arb_srcptr b, slong q, const arb_t z, int regularized, slong prec)"

### Confluent hypergeometric functions
arbcall"void arb_hypgeom_0f1(arb_t res, const arb_t a, const arb_t z, int regularized, slong prec)"
arbcall"void arb_hypgeom_m(arb_t res, const arb_t a, const arb_t b, const arb_t z, int regularized, slong prec)"
arbcall"void arb_hypgeom_1f1(arb_t res, const arb_t a, const arb_t b, const arb_t z, int regularized, slong prec)"
arbcall"void arb_hypgeom_u(arb_t res, const arb_t a, const arb_t b, const arb_t z, slong prec)"

### Gauss hypergeometric function
arbcall"void arb_hypgeom_2f1(arb_t res, const arb_t a, const arb_t b, const arb_t c, const arb_t z, int regularized, slong prec)"

### Error functions and Fresnel integrals
arbcall"void arb_hypgeom_erf(arb_t res, const arb_t z, slong prec)"
arbcall"void _arb_hypgeom_erf_series(arb_ptr res, arb_srcptr z, slong zlen, slong len, slong prec)"
arbcall"void arb_hypgeom_erf_series(arb_poly_t res, const arb_poly_t z, slong len, slong prec)"
arbcall"void arb_hypgeom_erfc(arb_t res, const arb_t z, slong prec)"
arbcall"void _arb_hypgeom_erfc_series(arb_ptr res, arb_srcptr z, slong zlen, slong len, slong prec)"
arbcall"void arb_hypgeom_erfc_series(arb_poly_t res, const arb_poly_t z, slong len, slong prec)"
arbcall"void arb_hypgeom_erfi(arb_t res, const arb_t z, slong prec)"
arbcall"void _arb_hypgeom_erfi_series(arb_ptr res, arb_srcptr z, slong zlen, slong len, slong prec)"
arbcall"void arb_hypgeom_erfi_series(arb_poly_t res, const arb_poly_t z, slong len, slong prec)"
arbcall"void arb_hypgeom_fresnel(arb_t res1, arb_t res2, const arb_t z, int normalized, slong prec)"
arbcall"void _arb_hypgeom_fresnel_series(arb_ptr res1, arb_ptr res2, arb_srcptr z, slong zlen, int normalized, slong len, slong prec)"
arbcall"void arb_hypgeom_fresnel_series(arb_poly_t res1, arb_poly_t res2, const arb_poly_t z, int normalized, slong len, slong prec)"

### Incomplete gamma and beta functions
arbcall"void arb_hypgeom_gamma_upper(arb_t res, const arb_t s, const arb_t z, int regularized, slong prec)"
arbcall"void _arb_hypgeom_gamma_upper_series(arb_ptr res, const arb_t s, arb_srcptr z, slong zlen, int regularized, slong n, slong prec)"
arbcall"void arb_hypgeom_gamma_upper_series(arb_poly_t res, const arb_t s, const arb_poly_t z, int regularized, slong n, slong prec)"
arbcall"void arb_hypgeom_gamma_lower(arb_t res, const arb_t s, const arb_t z, int regularized, slong prec)"
arbcall"void _arb_hypgeom_gamma_lower_series(arb_ptr res, const arb_t s, arb_srcptr z, slong zlen, int regularized, slong n, slong prec)"
arbcall"void arb_hypgeom_gamma_lower_series(arb_poly_t res, const arb_t s, const arb_poly_t z, int regularized, slong n, slong prec)"
arbcall"void arb_hypgeom_beta_lower(arb_t res, const arb_t a, const arb_t b, const arb_t z, int regularized, slong prec)"
arbcall"void _arb_hypgeom_beta_lower_series(arb_ptr res, const arb_t a, const arb_t b, arb_srcptr z, slong zlen, int regularized, slong n, slong prec)"
arbcall"void arb_hypgeom_beta_lower_series(arb_poly_t res, const arb_t a, const arb_t b, const arb_poly_t z, int regularized, slong n, slong prec)"

### Exponential and trigonometric integrals
arbcall"void arb_hypgeom_expint(arb_t res, const arb_t s, const arb_t z, slong prec)"
arbcall"void arb_hypgeom_ei(arb_t res, const arb_t z, slong prec)"
arbcall"void _arb_hypgeom_ei_series(arb_ptr res, arb_srcptr z, slong zlen, slong len, slong prec)"
arbcall"void arb_hypgeom_ei_series(arb_poly_t res, const arb_poly_t z, slong len, slong prec)"
arbcall"void arb_hypgeom_si(arb_t res, const arb_t z, slong prec)"
arbcall"void _arb_hypgeom_si_series(arb_ptr res, arb_srcptr z, slong zlen, slong len, slong prec)"
arbcall"void arb_hypgeom_si_series(arb_poly_t res, const arb_poly_t z, slong len, slong prec)"
arbcall"void arb_hypgeom_ci(arb_t res, const arb_t z, slong prec)"
arbcall"void _arb_hypgeom_ci_series(arb_ptr res, arb_srcptr z, slong zlen, slong len, slong prec)"
arbcall"void arb_hypgeom_ci_series(arb_poly_t res, const arb_poly_t z, slong len, slong prec)"
arbcall"void arb_hypgeom_shi(arb_t res, const arb_t z, slong prec)"
arbcall"void _arb_hypgeom_shi_series(arb_ptr res, arb_srcptr z, slong zlen, slong len, slong prec)"
arbcall"void arb_hypgeom_shi_series(arb_poly_t res, const arb_poly_t z, slong len, slong prec)"
arbcall"void arb_hypgeom_chi(arb_t res, const arb_t z, slong prec)"
arbcall"void _arb_hypgeom_chi_series(arb_ptr res, arb_srcptr z, slong zlen, slong len, slong prec)"
arbcall"void arb_hypgeom_chi_series(arb_poly_t res, const arb_poly_t z, slong len, slong prec)"
arbcall"void arb_hypgeom_li(arb_t res, const arb_t z, int offset, slong prec)"
arbcall"void _arb_hypgeom_li_series(arb_ptr res, arb_srcptr z, slong zlen, int offset, slong len, slong prec)"
arbcall"void arb_hypgeom_li_series(arb_poly_t res, const arb_poly_t z, int offset, slong len, slong prec)"

### Bessel functions
arbcall"void arb_hypgeom_bessel_j(arb_t res, const arb_t nu, const arb_t z, slong prec)"
arbcall"void arb_hypgeom_bessel_y(arb_t res, const arb_t nu, const arb_t z, slong prec)"
arbcall"void arb_hypgeom_bessel_jy(arb_t res1, arb_t res2, const arb_t nu, const arb_t z, slong prec)"
arbcall"void arb_hypgeom_bessel_i(arb_t res, const arb_t nu, const arb_t z, slong prec)"
arbcall"void arb_hypgeom_bessel_i_scaled(arb_t res, const arb_t nu, const arb_t z, slong prec)"
arbcall"void arb_hypgeom_bessel_k(arb_t res, const arb_t nu, const arb_t z, slong prec)"
arbcall"void arb_hypgeom_bessel_k_scaled(arb_t res, const arb_t nu, const arb_t z, slong prec)"

### Airy functions
arbcall"void arb_hypgeom_airy(arb_t ai, arb_t ai_prime, arb_t bi, arb_t bi_prime, const arb_t z, slong prec)"
arbcall"void arb_hypgeom_airy_jet(arb_ptr ai, arb_ptr bi, const arb_t z, slong len, slong prec)"
arbcall"void _arb_hypgeom_airy_series(arb_ptr ai, arb_ptr ai_prime, arb_ptr bi, arb_ptr bi_prime, arb_srcptr z, slong zlen, slong len, slong prec)"
arbcall"void arb_hypgeom_airy_series(arb_poly_t ai, arb_poly_t ai_prime, arb_poly_t bi, arb_poly_t bi_prime, const arb_poly_t z, slong len, slong prec)"
#ni arbcall"void arb_hypgeom_airy_zero(arb_t a, arb_t a_prime, arb_t b, arb_t b_prime, const fmpz_t n, slong prec)"

### Coulomb wave functions
arbcall"void arb_hypgeom_coulomb(arb_t F, arb_t G, const arb_t l, const arb_t eta, const arb_t z, slong prec)"
arbcall"void arb_hypgeom_coulomb_jet(arb_ptr F, arb_ptr G, const arb_t l, const arb_t eta, const arb_t z, slong len, slong prec)"
arbcall"void _arb_hypgeom_coulomb_series(arb_ptr F, arb_ptr G, const arb_t l, const arb_t eta, arb_srcptr z, slong zlen, slong len, slong prec)"
arbcall"void arb_hypgeom_coulomb_series(arb_poly_t F, arb_poly_t G, const arb_t l, const arb_t eta, const arb_poly_t z, slong len, slong prec)"

### Orthogonal polynomials and functions
arbcall"void arb_hypgeom_chebyshev_t(arb_t res, const arb_t nu, const arb_t z, slong prec)"
arbcall"void arb_hypgeom_chebyshev_u(arb_t res, const arb_t nu, const arb_t z, slong prec)"
arbcall"void arb_hypgeom_jacobi_p(arb_t res, const arb_t n, const arb_t a, const arb_t b, const arb_t z, slong prec)"
arbcall"void arb_hypgeom_gegenbauer_c(arb_t res, const arb_t n, const arb_t m, const arb_t z, slong prec)"
arbcall"void arb_hypgeom_laguerre_l(arb_t res, const arb_t n, const arb_t m, const arb_t z, slong prec)"
arbcall"void arb_hypgeom_hermite_h(arb_t res, const arb_t nu, const arb_t z, slong prec)"
arbcall"void arb_hypgeom_legendre_p(arb_t res, const arb_t n, const arb_t m, const arb_t z, int type, slong prec)"
arbcall"void arb_hypgeom_legendre_q(arb_t res, const arb_t n, const arb_t m, const arb_t z, int type, slong prec)"
arbcall"void arb_hypgeom_legendre_p_ui_deriv_bound(mag_t dp, mag_t dp2, ulong n, const arb_t x, const arb_t x2sub1)"
arbcall"void arb_hypgeom_legendre_p_ui_zero(arb_t res, arb_t res_prime, ulong n, const arb_t x, slong K, slong prec)"
arbcall"void arb_hypgeom_legendre_p_ui_one(arb_t res, arb_t res_prime, ulong n, const arb_t x, slong K, slong prec)"
arbcall"void arb_hypgeom_legendre_p_ui_asymp(arb_t res, arb_t res_prime, ulong n, const arb_t x, slong K, slong prec)"
arbcall"void arb_hypgeom_legendre_p_rec(arb_t res, arb_t res_prime, ulong n, const arb_t x, slong prec)"
arbcall"void arb_hypgeom_legendre_p_ui(arb_t res, arb_t res_prime, ulong n, const arb_t x, slong prec)"
arbcall"void arb_hypgeom_legendre_p_ui_root(arb_t res, arb_t weight, ulong n, ulong k, slong prec)"

### Dilogarithm
arbcall"void arb_hypgeom_dilog(arb_t res, const arb_t z, slong prec)"

### Hypergeometric sequences
arbcall"void arb_hypgeom_central_bin_ui(arb_t res, ulong n, slong prec)"


text:
# requires python3 and rospy to be installed

# ENV["PYTHON"] = "/usr/bin/python3"
# Pkg.build("PyCall")
# using PyCall
# using RobotOS

using Colors
using Caesar
using BSON

using Test
using Pkg

# import Caesar._PCL: FieldMapper, createMapping, PointCloud, PointField, PCLPointCloud2, Header, asType, _PCL_POINTFIELD_FORMAT, FieldMapping, MsgFieldMap, FieldMatches

##
@testset "test Caesar._PCL.PCLPointCloud2 to Caesar._PCL.PointCloud converter." begin
##

# testdatafile = joinpath( dirname(dirname(Pkg.pathof(Caesar))), "test", "testdata", "_PCLPointCloud2.bson")
# load presaved test data to test the coverter
# BSON.@load testdatafile PointCloudRef PointCloudTest

## build PCLPointCloud2 to be converted

datafile = joinpath( dirname(dirname(Pkg.pathof(Caesar))), "test", "testdata", "_PCLPointCloud2_15776.dat")
fid = open(datafile,"r")
data = read(fid)
close(fid)

##

pc2 = Caesar._PCL.PCLPointCloud2(;
  header = Caesar._PCL.Header(;
    seq = 92147,
    stamp = 0x0005b24647c4af10,
    frame_id = "velodyne"
  ),
  height = 1,
  width = 493,
  fields = [
    Caesar._PCL.PointField("x", 0x00000000, Caesar._PCL._PCL_FLOAT32, 0x00000001),
    Caesar._PCL.PointField("y", 0x00000004, Caesar._PCL._PCL_FLOAT32, 0x00000001),
    Caesar._PCL.PointField("z", 0x00000008, Caesar._PCL._PCL_FLOAT32, 0x00000001),
    Caesar._PCL.PointField("intensity", 0x00000010, Caesar._PCL._PCL_FLOAT32, 0x00000001)
  ],
  point_step = 32,
  row_step = 15776,
  is_dense = 1,
  data
)

##

show(pc2)

# test major unpacking / type conversion needed by ROS.sensor_msgs.pointcloud2
pc = Caesar._PCL.PointCloud(pc2)

show(pc)

@test pc.height == 1
@test pc.width == 493
@test pc.is_dense

@test length(pc.points) == 493

##

# truncated copy of reference data on good to 1e-3
@test isapprox( [-0,0,0],              pc.points[1].data[1:3], atol=5e-3)
@test isapprox( [-1.56486,1.09851,0],  pc.points[2].data[1:3], atol=5e-3)
@test isapprox( [-793.383,556.945,0],  pc.points[3].data[1:3], atol=5e-3)
@test isapprox( [-0,0,0],              pc.points[4].data[1:3], atol=5e-3)
@test isapprox( [-1.56148,1.10331,0],  pc.points[5].data[1:3], atol=5e-3)
@test isapprox( [-788.548,557.169,0],  pc.points[6].data[1:3], atol=5e-3)
@test isapprox( [-790.109,558.273,0],  pc.points[7].data[1:3], atol=5e-3)
@test isapprox( [-791.671,559.376,0],  pc.points[8].data[1:3], atol=5e-3)
@test isapprox( [-793.232,560.479,0],  pc.points[9].data[1:3], atol=5e-3)
@test isapprox( [-794.794,561.583,0],  pc.points[10].data[1:3], atol=5e-3)

# for (i,pt) in enumerate(pc.points)
#   @test isapprox( PointCloudRef[1][i,1:3], pt.data[1:3]; atol=1e-6)
# end

##
end

#


text:
import HTTP2
using HTTP2.Frame
using Test

# A server example
HTTP2.serve(8000, b"
Hello, world!
")


text:
using PlutoStaticHTMLDemo
using Documenter
using PlutoStaticHTML

DocMeta.setdocmeta!(PlutoStaticHTMLDemo, :DocTestSetup, :(using PlutoStaticHTMLDemo); recursive=true)

const DOCS_SRC_PATH = joinpath(@__DIR__, "src")
output_format = PlutoStaticHTML.documenter_output
build_opts = PlutoStaticHTML.BuildOptions(
    DOCS_SRC_PATH; previous_dir=DOCS_SRC_PATH, output_format=output_format
)
PlutoStaticHTML.build_notebooks(build_opts)

makedocs(;
    modules=[PlutoStaticHTMLDemo],
    authors="Seth Axen <seth.axen@gmail.com> and contributors",
    repo="https://github.com/sethaxen/PlutoStaticHTMLDemo.jl/blob/{commit}{path}#{line}",
    sitename="PlutoStaticHTMLDemo.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://sethaxen.github.io/PlutoStaticHTMLDemo.jl",
        assets=String[],
    ),
    pages=[
        "Home" => "index.md",
        "Plotting in Pluto" => "notebook.md",
    ],
)

deploydocs(;
    repo="github.com/sethaxen/PlutoStaticHTMLDemo.jl",
    devbranch="main",
)


text:
export CartPoleEnv

struct CartPoleEnvParams{T}
    gravity::T
    masscart::T
    masspole::T
    totalmass::T
    halflength::T
    polemasslength::T
    forcemag::T
    dt::T
    thetathreshold::T
    xthreshold::T
    max_steps::Int
end

Base.show(io::IO, params::CartPoleEnvParams) = print(
    io,
    join(["$p=$(getfield(params, p))" for p in fieldnames(CartPoleEnvParams)], ","),
)

mutable struct CartPoleEnv{T,R<:AbstractRNG} <: AbstractEnv
    params::CartPoleEnvParams{T}
    state::Array{T,1}
    action::Int
    done::Bool
    t::Int
    rng::R
end

"""
    CartPoleEnv(;kwargs...)

# Keyword arguments
- `T = Float64`
- `gravity = T(9.8)`
- `masscart = T(1.0)`
- `masspole = T(0.1)`
- `halflength = T(0.5)`
- `forcemag = T(10.0)`
- `max_steps = 200`
- 'dt = 0.02'
- `rng = Random.GLOBAL_RNG`
"""
function CartPoleEnv(;
    T = Float64,
    gravity = 9.8,
    masscart = 1.0,
    masspole = 0.1,
    halflength = 0.5,
    forcemag = 10.0,
    max_steps = 200,
    dt = 0.02,
    rng = Random.GLOBAL_RNG,
)
    params = CartPoleEnvParams{T}(
        gravity,
        masscart,
        masspole,
        masscart + masspole,
        halflength,
        masspole * halflength,
        forcemag,
        dt,
        2 * 12 * π / 360,
        2.4,
        max_steps,
    )
    high = cp = CartPoleEnv(params, zeros(T, 4), 2, false, 0, rng)
    reset!(cp)
    cp
end

CartPoleEnv{T}(; kwargs...) where {T} = CartPoleEnv(; T = T, kwargs...)

function RLBase.reset!(env::CartPoleEnv{T}) where {T<:Number}
    env.state[:] = T(0.1) * rand(env.rng, T, 4) .- T(0.05)
    env.t = 0
    env.action = 2
    env.done = false
    nothing
end

RLBase.action_space(env::CartPoleEnv) = Base.OneTo(2)

RLBase.state_space(env::CartPoleEnv{T}) where {T} = Space(
    ClosedInterval{T}[
        (-2*env.params.xthreshold)..(2*env.params.xthreshold),
        -1e38..1e38,
        (-2*env.params.thetathreshold)..(2*env.params.thetathreshold),
        -1e38..1e38,
    ],
)

RLBase.reward(env::CartPoleEnv{T}) where {T} = env.done ? zero(T) : one(T)
RLBase.is_terminated(env::CartPoleEnv) = env.done
RLBase.state(env::CartPoleEnv) = env.state

function (env::CartPoleEnv)(a)
    @assert a in (1, 2)
    env.action = a
    env.t += 1
    force = a == 2 ? env.params.forcemag : -env.params.forcemag
    x, xdot, theta, thetadot = env.state
    costheta = cos(theta)
    sintheta = sin(theta)
    tmp = (force + env.params.polemasslength * thetadot^2 * sintheta) / env.params.totalmass
    thetaacc =
        (env.params.gravity * sintheta - costheta * tmp) / (
            env.params.halflength *
            (4 / 3 - env.params.masspole * costheta^2 / env.params.totalmass)
        )
    xacc = tmp - env.params.polemasslength * thetaacc * costheta / env.params.totalmass
    env.state[1] += env.params.dt * xdot
    env.state[2] += env.params.dt * xacc
    env.state[3] += env.params.dt * thetadot
    env.state[4] += env.params.dt * thetaacc
    env.done =
        abs(env.state[1]) > env.params.xthreshold ||
        abs(env.state[3]) > env.params.thetathreshold ||
        env.t > env.params.max_steps
    nothing
end

Random.seed!(env::CartPoleEnv, seed) = Random.seed!(env.rng, seed)


text:
################################################################################
## Linear Regression
##   y ~ N(b0 + b1 * x, s2)
##   b0, b1 ~ N(0, 1000)
##   s2 ~ invgamma(0.001, 0.001)
################################################################################

data = Dict{Symbol,Any}(
    :x => [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
    :y => [1, 3, 3, 3, 5, 5, 7, 7, 7, 9],
)
data[:xmat] = [ones(10) data[:x]]
model = Model(
    y = Stochastic(1, (μ, s2) -> MultivariateNormal(μ, s2), false),
    μ = Logical(1, (β, xmat) -> xmat * β, false),
    β = Stochastic(1, () -> Normal(0, sqrt(1000)), true),
    s2 = Stochastic(() -> InverseGamma(0.001, 0.001), true),
)

inits = [
    Dict{Symbol,Any}(:y => data[:y], :β => rand(Normal(0, 1), 2), :s2 => rand(Gamma(1, 1))) for i = 1:2
]

@testset "Slice" begin
    Random.seed!(123)
    samplers = [Slice([:β, :s2], 1.0)]
    setsamplers!(model, samplers)

    sim = mcmc(
        model,
        data,
        inits,
        100000,
        burnin = 50000,
        thin = 10,
        chains = 2,
        verbose=false
    )

    r_m = summarystats(sim).value[2:3, 1, 1]
    r_sd = summarystats(sim).value[2:3, 2, 1]
    @test r_m[1]-r_sd[1] <= 0.60 <= r_m[1]+r_sd[1]
    @test r_m[2]-r_sd[2] <= 0.80 <= r_m[2]+r_sd[2]
end

@testset "RWM" begin
    Random.seed!(123)
    samplers = [RWM([:β, :s2], 1.0)]
    setsamplers!(model, samplers)

    sim = mcmc(
        model,
        data,
        inits,
        100000,
        burnin = 50000,
        thin = 10,
        chains = 2,
        verbose=false
    )

    r_m = summarystats(sim).value[2:3, 1, 1]
    r_sd = summarystats(sim).value[2:3, 2, 1]
    @test r_m[1]-r_sd[1] <= 0.60 <= r_m[1]+r_sd[1]
    @test r_m[2]-r_sd[2] <= 0.80 <= r_m[2]+r_sd[2]
end

@testset "NUTS_Classic" begin
    Random.seed!(123)
    samplers = [NUTS([:β, :s2], variant=:classic)]
    setsamplers!(model, samplers)

    sim = mcmc(
        model,
        data,
        inits,
        100000,
        burnin = 50000,
        thin = 10,
        chains = 2,
        verbose=false
    )

    r_m = summarystats(sim).value[2:3, 1, 1]
    r_sd = summarystats(sim).value[2:3, 2, 1]
    @test r_m[1]-r_sd[1] <= 0.60 <= r_m[1]+r_sd[1]
    @test r_m[2]-r_sd[2] <= 0.80 <= r_m[2]+r_sd[2]
end

@testset "NUTS_Riemann" begin
    Random.seed!(123)
    samplers = [NUTS([:β, :s2], variant=:riemann)]
    setsamplers!(model, samplers)

    sim = mcmc(
        model,
        data,
        inits,
        100000,
        burnin = 50000,
        thin = 10,
        chains = 2,
        verbose=false
    )

    r_m = summarystats(sim).value[2:3, 1, 1]
    r_sd = summarystats(sim).value[2:3, 2, 1]
    @test r_m[1]-r_sd[1] <= 0.60 <= r_m[1]+r_sd[1]
    @test r_m[2]-r_sd[2] <= 0.80 <= r_m[2]+r_sd[2]
end

@testset "HMC" begin
    Random.seed!(123)
    samplers = [HMC([:β, :s2], 0.007, 8)]
    setsamplers!(model, samplers)

    sim = mcmc(
        model,
        data,
        inits,
        100000,
        burnin = 50000,
        thin = 10,
        chains = 2,
        verbose=false
    )

    r_m = summarystats(sim).value[2:3, 1, 1]
    r_sd = summarystats(sim).value[2:3, 2, 1]
    @test r_m[1]-r_sd[1] <= 0.60 <= r_m[1]+r_sd[1]
    @test r_m[2]-r_sd[2] <= 0.80 <= r_m[2]+r_sd[2]
end

max_stars_repo_name: erathorn/MCPhylo

max_issues_repo_name: erathorn/MCPhylo

max_forks_repo_name: erathorn/MCPhylo


text:
# SAM Header
# ==========

# Implicit type restriction for predefiend tags:
#   * HD: Dict{String,String}
#   * SQ: Vector{Dict{String,Any}}
#   * RG: Vector{Dict{String,Any}}
#   * PG: Vector{Dict{String,Any}}
#   * CO: Vector{String}

type SAMHeader <: Associative{String,Any}
    data::Dict{String,Any}
end

function SAMHeader()
    return SAMHeader(Dict{String,Any}())
end

function Base.length(header::SAMHeader)
    return length(header.data)
end

function Base.keys(header::SAMHeader)
    return keys(header.data)
end

function Base.values(header::SAMHeader)
    return values(header.data)
end

function Base.getindex(header::SAMHeader, key::AbstractString)
    return getindex(header.data, key)
end

function Base.setindex!(header::SAMHeader, val, key::AbstractString)
    if length(key) != 2
        error("SAM header tag must be of length 2")
    end
    # TODO: check value type and validity?
    setindex!(header.data, val, key)
    return header
end

function Base.delete!(header::SAMHeader, key::AbstractString)
    return delete!(header.data, key)
end

function Base.start(header::SAMHeader)
    return start(header.data)
end

function Base.done(header::SAMHeader, s)
    return done(header.data, s)
end

function Base.next(header::SAMHeader, s)
    return next(header.data, s)
end


text:
@userplot Plot_Solution
@recipe function f(ps::Plot_Solution; derivative=0, true_solution=nothing, stdfac=1.96)
    sol = ps.args[1]

    stds = hcat([sqrt.(diag(M)) for M in sol.x.Σ]...)'
    d = sol.solver.constants.d

    # Plot the results
    means = hcat(sol.x.μ...)'

    # return p1
    ribbon --> stdfac*stds[:, 1:d]
    # marker --> :x
    title --> "\$f(t)\$"
    label --> "Filter Estimate"
    xguide --> "t"

    # TODO: Add true solution back in
    # if !(isnothing(true_solution))
    #     plot!(p1, range(sol.t[1], sol.t[end], length=1000), true_solution, label="True solution")
    # end

    return sol.t, means[:, 1:d]
end

@userplot Plot_Stepsizes
@recipe function f(ps::Plot_Stepsizes; rejections=true)
    sol = ps.args[1]

    stepsizes = sol.t[2:end-1] - sol.t[1:end-2]
    if all([h ≈ stepsizes[1] for h in stepsizes])
        stepsizes .= stepsizes[1]
    end

    yscale --> :log10,
    # marker --> :x
    legend --> :bottomleft
    yguide --> "h"

    @series begin
        label --> "h"
        return sol.t[2:end-1], stepsizes
    end

    # TODO Add rejections back in
    rejected = filter(p->!p.accept, sol.proposals)
    if rejections && length(rejected) > 0
        @series begin
            t = [p.t for p in rejected]
            dt = [p.dt for p in rejected]
            marker := :x
            seriescolor := :black
            label := "rejected"
            seriestype := :scatter
            return t, dt
        end
    end
end

@userplot Plot_Analytic_Solution
@recipe function f(pas::Plot_Analytic_Solution)
    sol, analytic = pas.args[1], pas.args[2]
    ts = sol.t
    highres_ts = range(ts[1], ts[end], length=1000)
    label --> "True solution"
    return highres_ts, hcat(analytic.(highres_ts)...)'
end

@userplot Plot_Errors
@recipe function f(pe::Plot_Errors; c=1.96)
    sol, analytic = pe.args[1], pe.args[2]
    ts = sol.t
    d = sol.solver.constants.d

    f_est = sol.u.μ
    if d == 1
        f_est = hcat(f_est...)'
    end
    true_sol = analytic.(ts)
    diffs = true_sol .- f_est

    @series begin
        yscale --> :log10
        # marker --> :x
        label --> "Global Error"
        yguide --> "function error"
        return ts[2:end], norm.(diffs, 2)[2:end]
    end

    stds = [sqrt.(diag(cov)) for cov in sol.u.Σ]
    @series begin
        label --> "$c * std"
        linestyle := :dash
        return ts[2:end], norm.(c.*stds[2:end], 2)
    end
end

@userplot Plot_Calibration
@recipe function f(pc::Plot_Calibration; interval=0.95)
    sol, analytic = pc.args[1], pc.args[2]
    ts = sol.t
    d = sol.solver.constants.d

    accepted_proposals = [p for p in sol.proposals if p.accept]
    filter_estimates = [p.filter_estimate for p in accepted_proposals]
    f_est = sol.u.μ[2:end]
    f_covs = sol.u.Σ[2:end]
    if d == 1
        f_est = stack(f_est)
    end
    true_sol = analytic.(ts[2:end])
    diffs = true_sol .- f_est

    sugg = [d' * inv(C) * d for (d, C) in zip(diffs, f_covs)]

    @series begin
        label --> ""
        yscale --> :log10
        return ts[2:end], sugg
    end
    if !isnothing(interval)
        @series begin
            seriestype := :hline
            seriescolor := :black
            linestyle := :dash
            label := "$(interval*100)% Interval"
            p = 1-interval
            return [quantile(Chisq(d), 1-(p/2)), quantile(Chisq(d), p/2)]
        end
    end
end

@userplot Plot_Residuals
@recipe function f(pr::Plot_Residuals)
    sol = pr.args[1]
    d = sol.solver.constants.d

    accepted_proposals = [p for p in sol.proposals if p.accept]
    measurements = [p.measurement for p in accepted_proposals]
    times = [p.t for p in accepted_proposals]
    res = [z.μ' * z.Σ^(-1) * z.μ for z in measurements] ./ d

    yscale = all(res .> 0) ? :log10 : :identity
    yscale --> yscale
    yguide --> "residual"

    @series begin
        # marker --> :x
        label --> "residual"
        return times, res
    end
    @series begin
        seriestype := :hline
        linestyle := :dash
        seriescolor := 3
        label := nothing
        return [1]
    end

    # Plot std of measurements
    msmnt_vars = [diag(m.Σ) for m in measurements]
    msmnt_stds = [norm(sqrt.(S)) for S in msmnt_vars]
    @series begin
        label := "Measurement std"
        seriescolor := 2
        times, 3msmnt_stds
    end

    # Plot rejected proposals
    rejected_proposals = [p for p in sol.proposals if !p.accept]
    if length(rejected_proposals) > 0
        times = [p.t for p in rejected_proposals]
        measurements = [p.measurement for p in rejected_proposals]
        res = [z.μ' * z.Σ^(-1) * z.μ for z in measurements] ./ d
        steps = [p.dt for p in rejected_proposals]
        # rejected_t, rejected_h, rejected_errors = collect(zip(sol.rejected...))
        @series begin
            seriestype := :scatter
            seriescolor := :black
            marker := :x
            label := "rejected"
            return times .+ steps, res
        end
    end
end

@userplot Plot_Sigmas
@recipe function f(ps::Plot_Sigmas)
    sol = ps.args[1]

    # Plot Sigmas
    accepted_proposals = [p for p in sol.proposals if p.accept]
    ts = [p.t for p in accepted_proposals]
    σ²_dynamic = [p.σ²
                  # dynamic_sigma_estimation(sol.solver.sigma_estimator; p...)
                  for p in accepted_proposals]
    if !all(σ²_dynamic .== 1)
        σ² = σ²_dynamic
    else
        σ²_static = [static_sigma_estimation(sol.solver.sigma_estimator, sol.solver, accepted_proposals[1:i])
                     for i in 1:length(accepted_proposals)]
        σ² = σ²_static
    end

    label --> "σ²"
    yguide --> "σ²"
    # marker --> :x
    yscale --> :log10
    legend --> :bottomright

    return ts, σ²
end

function hairer_plot(sol; analytic=nothing, tol=nothing, f_ylims=nothing, title=nothing)
    analytic = isnothing(analytic) ? solve(sol.prob) : analytic

    p_f = plot_solution(sol)
    plot!(p_f, title=title, ylabel="f(t)")
    plot_analytic_solution!(p_f, sol, analytic)

    p_h = plot_stepsizes(sol)
    plot_sigmas!(twinx(p_h), sol)

    # 3: Errors
    p_err = plot_errors(sol, analytic)

    p_res = plot_residuals(sol)

    return plot(
        p_f,
        p_h,
        p_err,
        p_res,
        xlims=(sol.t[1], sol.t[end]),
        layout=(4,1),
        size=(600,800))
end

max_stars_repo_name: nathanaelbosch/capos

max_issues_repo_name: nathanaelbosch/capos

max_forks_repo_name: nathanaelbosch/capos


text:
#= hausdorff.jl
Philip Barrett, pobarrett@gmail.com
30may2016, Chicago

Computes the hausdorff distance between polygons (and points)
=#

"""
    hausdorff

Computes the hausdorff distance between two polygons, or two arrays of polygons
"""
function hausdorff( poly1::Polygon, poly2::Polygon )
  return hausdorffCore( poly1.pts, poly2 )
end

function hausdorff( poly1::Array{Polygon,1}, poly2::Array{Polygon,1} )
  N = length(poly1)
  length(poly2) == N || error("Polygon arrays must have same diemnsion")
      # Error handling
  return [ hausdorffCore( poly1[i].pts, poly2[i] ) for i in 1:N ]
end

function hausdorff( poly1::Array{Polygon,2}, poly2::Array{Polygon,2} )
  N, M = size(poly1)
  size(poly2) == (N, M) || error("Polygon arrays must have same diemnsion")
      # Error handling
  return [ hausdorffCore( poly1[i,j].pts, poly2[i,j] ) for i in 1:N, j in 1:M ]
end
"""
    function hausdorffCore( newpts::Matrix, poly::Polygon  )

Core part of any hausdorff distance routine.  Measures the distance between a matrix of points newpts and a polygon poly
"""
function hausdorffCore( newpts::Matrix, poly::Polygon )

  pts, dirs, dists = poly.pts, poly.dirs, poly.dists
      # Unpack the polygon object
  N = size(newpts)[1]
      # number of points
  dline, dvert, newdist, hausdorff = 0.0, 0.0, 0.0, 0.0
      # initialize the distance measures
  for i in 1:N
    dline = setsNearestLine( vec(newpts[i,:]), dirs, dists )
        # Distance to nearest line
    dvert = setsNearestVert( vec(newpts[i,:]), pts )
        # Distance to nearest vertex
    newdist = min( dline, dvert )
        # Minimum distance
    hausdorff  = max( hausdorff, newdist )
        # Update hausdorff dist via minmax
  end
  return hausdorff
end

function setsNearestLine( pt::Vector, dirs::Matrix, dists::Vector )

  N::Int = length(dists)
      # Number of sides
  dff = dists - dirs * pt
  absdiff = abs( dff )
      # Difference between the projection directions and the boundaries (and absolute value)
  nearpts = ones( N ) * pt' + dirs .* ( dff * ones( 1, 2 ) )
      # Recover the closest points on the normal surfaces
  idx = setsInsidePtsIdx( nearpts, dirs, dists )
      # Test whether these points are inside the Polygon
  if !any(idx)
    return Inf
        # Outsize return if none of the nearest points are in the polygon
  end
  distinside = [ idx[i] ? absdiff[i] : Inf for i in 1:N ]
      # Compute the minimum distances to the points on the lines that are
      # inside the polygon
  return minimum( distinside )
      # Return the minimum
end

function setsInsidePtsIdx( cands::Matrix, dirs::Matrix,
                            dists::Vector )
  N::Int = size(cands)[1]
      # Number of candidate points
  D = dirs * cands'  - dists * ones( 1, N )
      # Negative if inside the set
  return all( D .<= 0, 1 )
end

function setsNearestVert( cand::Vector, pts::Matrix )
  N = size(pts)[1]
      # Nubmer of points
  return minimum( [ norm( cand' - pts[i,:] ) for i in 1:N ] )
end

# double sets_nearest_vert_cpp_eigen( VectorXd pt, MatrixXd mB ){
# // C++ translation of sets.nearest.vert
#
#   int iB = mB.rows() ;
#         // Number of rows in B
#   VectorXd vDists( iB ) ;
#         // Initialize the distance vector
#   for( int iII=0; iII<iB; iII++ ){
#     vDists[ iII ] = ( pt.transpose() - mB.row( iII ) ).norm() ;
#           // Compute the distance to pt
#   }
#   return vDists.minCoeff() ;
# }


text:
"""
Uses `hunting-circular-orbits.jl` and the Page & Thorne (1974) eq. (12)
to simulate a flux model.
Works for a < 0.97, then the orbit finder starts to panic close to the ISCO.
"""

# using Revise

using GeodesicTracer
using ComputedGeodesicEquations
using GeodesicBase
using AccretionFormulae

import CarterBoyerLindquist # for rms function
import LinearAlgebra

using StaticArrays
using Interpolations

# include("/Users/lucyackland-snow/Desktop/DevEnv1/dev/AccretionFormulae/src/johannsen.jl")
include("hunting-circular-orbits.jl")

# for Kerr, this is just `u[2]`, but to illustrate
# how you could do this for other metrics
function sqrt_g_tilde(m, u)
    metric = GeodesicBase.metric(m, u)
    metric_no_θ = metric[setdiff(1:4, 3), setdiff(1:4, 3)]

    sqrt(-LinearAlgebra.det(metric_no_θ))
end
"""
Johannsen 2014 eq 28
"""
function sqrt_g_tilde(r; a=0.97, M=0, ϵ_3=0, α_13=0, α_22=0, α_52=0)
    sqrt_g_tilde =   (r/abs(1 + (α_13 * M^3)/(r^3) -
                        (α_22 * a^2 * M^2)/(r^4) +
                        (α_13 * a^2 * M^3)/(r^5))) *
                √(((1 + (ϵ_3 * M^3)/(r^3))^3)/
                (1 + (α_52 * M^2)/(r^2)))
end

# gradient function
#
# we just compute discrete derivatives, since this is
# sufficient when we have small `r`.
# using interpolations can make the resolution better
# and increase the accuracy, but this is a good quick-and-dirty way
# to see if our results are going in the right direction
function ∇r(array, r)
    deriv = diff(array) ./ diff(r)
    # ensure same dimension using a lazy extension
    push!(deriv, deriv[end])
    deriv
end

"""
    f_approx(Lz, ∇rwt, ∇rLz_array, wt_array)
A discretized version of `f` from Page & Thorne (1974) eq. (12):
```math
f(r_n) \\approx - \\left. \\left[
   \\frac{1}{L_z} \\left(\\frac{\\partial \\dot{u}^t}{\\partial r} \\right)
\\right] \\right\\rvert_{r = r_n}
\\sum_{i=0}^{n} \\left\\{ \\left. \\left[
   \\frac{1}{\\dot{u}^t} \\left( \\frac{\\partial L_z}{\\partial r}\\right)
\\right] \\right\\rvert_{r = r_i}
\\right\\}
```
"""
function f_approx(Lz, ∇rwt, ∇rLz_array, wt_array, r; a=0.97, M=1.0, ϵ_3=0, α_13=0, α_22=0, α_52=0)
    -(∇rwt / Lz) * sum(∇rLz_array ./ wt_array) * r/sqrt_g_tilde(r; a, M, ϵ_3, α_13, α_22, α_52)
end

# use this to see where things break down / if they break down
function plot_derivative_test(r_range, res)
    lzs = res[:, 2]
    # naive derivatives
    derivs = ∇r(lzs, r_range)

    plot(r_range, lzs, label = "Lz", c = :black, legend = false)
    plot!(r_range, derivs, label = "derivative", c = :black, style = :dot, lw = 1.2)
end

# test plot
function f_approx_test(r_range, res; a=0.97, M=1.0, ϵ_3=0, α_13=0, α_22=0, α_52=0)
    wts = res[:, 4]
    Lzs = res[:, 2]

    ∇rwt_array = ∇r(wts, r_range)
    ∇rLz_array = ∇r(Lzs, r_range)

    fs = map(2:length(r_range)) do i
        @views f_approx(Lzs[i], ∇rwt_array[i], ∇rLz_array[1:i], wts[1:i], r_range[i]; a=a, M=M, ϵ_3=ϵ_3, α_13=α_13, α_22=α_22, α_52=α_52)
    end

    # we rescale here for a better comparison
    fs = fs ./ maximum(fs)

    # we lose a point, so select all but last `r_range`
    # since our derivative function needs at least 2 points  work
    # plot(r_range[2:end], fs, c = :black, xlabel = "r", ylabel = "f", label = "simulated f")
end

# this works up until about 0.97 then, the orbit search
# starts to break down

function init(;m, a=0.97, M=1.0, disc_radius=50.0, ϵ_3=0, α_13=0, α_22=0, α_52=0)
    # m = BoyerLindquist(M = M, a = a)
    # m = JohannsenAD(M=1.0, a=a)
    u = @SVector [0.0, 13.0, deg2rad(90), 0.0]

    # making the steps too big over-estimates the gradient
    r_range = CarterBoyerLindquist.rms(m.M, m.a):0.01:disc_radius
    res = find_orbit_range(;m, r_range = r_range, a = m.a)

    # creating an interpolated function of f

    fs = f_approx_test(r_range, res; a=a, M=M, ϵ_3=ϵ_3, α_13=α_13, α_22=α_22, α_52=α_52)
    return fs, r_range
end

function f_approx_func(r, fs, r_range; a=0.97, M=1.0)

    r_range_short = collect(r_range)
    pop!(r_range_short)
    interp_linear_f = LinearInterpolation(r_range_short, fs,extrapolation_bc=Line())
    interp_linear_f(r)
end

G = 6.67e-11
c = 3e8
L_☼ = 3.8e26
M_☼ = 1.99e30
σ_SB = 5.67e-8
η = 0.1

function diss_approx(mdot, r, a_star, M, fs, r_range)
    diss_approx = ((c^6) / (G^2)) * mdot * f_approx_func(r, fs, r_range; a=a_star, M=M) / (4 * π * r)
end

function temperature_approx(r, a_star, M, fs, r_range)
    m_dot = AccretionFormulae.mdot(M)
    temperature_approx = (diss_approx(m_dot, r, a_star, M, fs, r_range) / σ_SB)^(1 / 4)
end

function observed_temperature_approx(r, a_star, M, g, fs, r_range)
    T = temperature_approx(r, a_star, M, fs, r_range)
    observed_temperature_approx = g * T
end


text:
@testset "authentication" begin
    @test true
end


text:
using Statistics
using Printf
using WeightedEnsemble

include("doublewell_setup.jl");

# number of coarse steps in WE
n_we_steps = 10;
# number of time steps during mutation step
nΔt_coarse = nΔt ÷ n_we_steps;
# number of samples in coarse matrix
n_samples_per_bin = 10^2;
# ensemble size
n_particles = 10^2;

# define bin structure
voronoi_pts = [[x] for x in LinRange(a - 0.1, b + 0.1, 21)];
B₀, bin_id, rebin! = setup_Voronoi_bins(voronoi_pts);

# define the mutation mapping
opts = MDOptions(n_iters = nΔt_coarse, n_save_iters = nΔt_coarse)
mutation! = x -> sample_trajectory!(x, sampler, options = opts);

selection! = (E, B, t) -> uniform_selection!(E, B);
we_sampler = WEsampler(mutation!, selection!, rebin!);

f = x -> Float64(-0.1 < x[1] < 0.1); # define observable
f1 = x -> Float64(-0.1 < x[1] < 0.1); # define observables
f2 = x -> Float64(x[1] > 1.1);
obs = (f1, f2);

# set up ensemble
E₀ = Dirac_to_Ensemble(x₀, n_particles);
rebin!(E₀, B₀, 0);

Random.seed!(200)
f_trajectory = WeightedEnsemble.trun_we_observables(E₀, B₀, we_sampler, n_we_steps, (f,));


text:
function construct_icd9(icd9::AbstractString)
    system = "ICD9"
    value = icd9_add_dots(icd9)
    return DiagnosisClass(; system, value)
end

function construct_ccs9(ccs9::AbstractString)
    system = "CCS9"
    value = ccs9
    return DiagnosisClass(; system, value)
end