text:
export FiniteDifferenceMethod, fdm, backward_fdm, forward_fdm, central_fdm, extrapolate_fdm

"""
    FiniteDifferences.DEFAULT_CONDITION

The default value for the [condition number](https://en.wikipedia.org/wiki/Condition_number)
of finite difference method. The condition number specifies the amplification of the ∞-norm
when passed to the function's derivative.
"""
const DEFAULT_CONDITION = 10

"""
    FiniteDifferences.DEFAULT_FACTOR

The default factor number. The factor number specifies the multiple to amplify the estimated
round-off errors by.
"""
const DEFAULT_FACTOR = 1

abstract type FiniteDifferenceMethod{P,Q} end

"""
    UnadaptedFiniteDifferenceMethod{P,Q} <: FiniteDifferenceMethod{P,Q}

A finite difference method that estimates a `Q`th order derivative from `P` function
evaluations. This method does not dynamically adapt its step size.

# Fields
- `grid::SVector{P,Int}`: Multiples of the step size that the function will be evaluated at.
- `coefs::SVector{P,Float64}`: Coefficients corresponding to the grid functions that the
    function evaluations will be weighted by.
- `coefs_neighbourhood::NTuple{3,SVector{P,Float64}}`: Sets of coefficients used for
    estimating the magnitude of the derivative in a neighbourhood.
- `bound_estimator::FiniteDifferenceMethod`: A finite difference method that is tuned to
    perform adaptation for this finite difference method.
- `condition::Float64`: Condition number. See See [`DEFAULT_CONDITION`](@ref).
- `factor::Float64`: Factor number. See See [`DEFAULT_FACTOR`](@ref).
- `max_range::Float64`: Maximum distance that a function is evaluated from the input at
    which the derivative is estimated.
- `∇f_magnitude_mult::Float64`: Internally computed quantity.
- `f_error_mult::Float64`: Internally computed quantity.
"""
struct UnadaptedFiniteDifferenceMethod{P,Q} <: FiniteDifferenceMethod{P,Q}
    grid::SVector{P,Int}
    coefs::SVector{P,Float64}
    coefs_neighbourhood::NTuple{3,SVector{P,Float64}}
    condition::Float64
    factor::Float64
    max_range::Float64
    ∇f_magnitude_mult::Float64
    f_error_mult::Float64
end

"""
    AdaptedFiniteDifferenceMethod{
        P, Q, E<:FiniteDifferenceMethod
    } <: FiniteDifferenceMethod{P,Q}

A finite difference method that estimates a `Q`th order derivative from `P` function
evaluations.

This method dynamically adapts its step size. The adaptation works by explicitly estimating
the truncation error and round-off error, and choosing the step size to optimally balance
those. The truncation error is given by the magnitude of the `P`th order derivative, which
will be estimated with another finite difference method (`bound_estimator`). This finite
difference method, `bound_estimator`, will be tasked with estimating the `P`th order
derivative in a _neighbourhood_, not just at some `x`. To do this, it will use a careful
reweighting of the function evaluations to estimate the `P`th order derivative at, in the
case of a central method, `x - h`, `x`, and `x + h`, where `h` is the step size. The
coeffients for this estimate, the _neighbourhood estimate_, are given by the three sets of
coeffients in `bound_estimator.coefs_neighbourhood`. The round-off error is estimated by the
round-off error of the function evaluations performed by `bound_estimator`. The trunction
error is amplified by `condition`, and the round-off error is amplified by `factor`. The
quantities `∇f_magnitude_mult` and `f_error_mult` are precomputed quantities that facilitate
the step size adaptation procedure.

# Fields
- `grid::SVector{P,Int}`: Multiples of the step size that the function will be evaluated at.
- `coefs::SVector{P,Float64}`: Coefficients corresponding to the grid functions that the
    function evaluations will be weighted by.
- `coefs_neighbourhood::NTuple{3,SVector{P,Float64}}`: Sets of coefficients used for
    estimating the magnitude of the derivative in a neighbourhood.
- `condition::Float64`: Condition number. See See [`DEFAULT_CONDITION`](@ref).
- `factor::Float64`: Factor number. See See [`DEFAULT_FACTOR`](@ref).
- `max_range::Float64`: Maximum distance that a function is evaluated from the input at
    which the derivative is estimated.
- `∇f_magnitude_mult::Float64`: Internally computed quantity.
- `f_error_mult::Float64`: Internally computed quantity.
- `bound_estimator::FiniteDifferenceMethod`: A finite difference method that is tuned to
    perform adaptation for this finite difference method.
"""
struct AdaptedFiniteDifferenceMethod{
    P, Q, E<:FiniteDifferenceMethod
} <: FiniteDifferenceMethod{P,Q}
    grid::SVector{P,Int}
    coefs::SVector{P,Float64}
    coefs_neighbourhood::NTuple{3,SVector{P,Float64}}
    condition::Float64
    factor::Float64
    max_range::Float64
    ∇f_magnitude_mult::Float64
    f_error_mult::Float64
    bound_estimator::E
end

"""
    FiniteDifferenceMethod(
        grid::AbstractVector{Int},
        q::Int;
        condition::Real=DEFAULT_CONDITION,
        factor::Real=DEFAULT_FACTOR,
        max_range::Real=Inf
    )

Construct a finite difference method.

# Arguments
- `grid::Vector{Int}`: The grid. See [`AdaptedFiniteDifferenceMethod`](@ref) or
    [`UnadaptedFiniteDifferenceMethod`](@ref).
- `q::Int`: Order of the derivative to estimate.

# Keywords
- `condition::Real`: Condition number. See [`DEFAULT_CONDITION`](@ref).
- `factor::Real`: Factor number. See [`DEFAULT_FACTOR`](@ref).
- `max_range::Real=Inf`: Maximum distance that a function is evaluated from the input at
    which the derivative is estimated.

# Returns
- `FiniteDifferenceMethod`: Specified finite difference method.
"""
function FiniteDifferenceMethod(
    grid::SVector{P,Int},
    q::Int;
    condition::Real=DEFAULT_CONDITION,
    factor::Real=DEFAULT_FACTOR,
    max_range::Real=Inf,
) where P
    _check_p_q(P, q)
    coefs, coefs_neighbourhood, ∇f_magnitude_mult, f_error_mult = _coefs_mults(grid, q)
    return UnadaptedFiniteDifferenceMethod{P,q}(
        grid,
        coefs,
        coefs_neighbourhood,
        condition,
        factor,
        max_range,
        ∇f_magnitude_mult,
        f_error_mult
    )
end
function FiniteDifferenceMethod(grid::AbstractVector{Int}, q::Int; kw_args...)
    return FiniteDifferenceMethod(SVector{length(grid)}(grid), q; kw_args...)
end

"""
    (m::FiniteDifferenceMethod)(f::Function, x::T) where T<:AbstractFloat

Estimate the derivative of `f` at `x` using the finite differencing method `m` and an
automatically determined step size.

# Arguments
- `f::Function`: Function to estimate derivative of.
- `x::T`: Input to estimate derivative at.

# Returns
- Estimate of the derivative.

# Examples

```julia-repl
julia> fdm = central_fdm(5, 1)
FiniteDifferenceMethod:
  order of method:       5
  order of derivative:   1
  grid:                  [-2, -1, 0, 1, 2]
  coefficients:          [0.08333333333333333, -0.6666666666666666, 0.0, 0.6666666666666666, -0.08333333333333333]

julia> fdm(sin, 1)
0.5403023058681607

julia> fdm(sin, 1) - cos(1)  # Check the error.
2.098321516541546e-14

julia> FiniteDifferences.estimate_step(fdm, sin, 1.0)  # Computes step size and estimates the error.
(0.001065235154086019, 1.9541865128909085e-13)
```
"""
# We loop over all concrete subtypes of `FiniteDifferenceMethod` for Julia v1.0 compatibility.
for T in (UnadaptedFiniteDifferenceMethod, AdaptedFiniteDifferenceMethod)
    @eval begin
        function (m::$T)(f::TF, x::Real) where TF<:Function
            x = float(x)  # Assume that converting to float is desired, if it isn't already.
            step = first(estimate_step(m, f, x))
            return m(f, x, step)
        end
        function (m::$T{P,0})(f::TF, x::Real) where {P,TF<:Function}
            # The automatic step size calculation fails if `Q == 0`, so handle that edge
            # case.
            return f(x)
        end
    end
end

"""
    (m::FiniteDifferenceMethod)(f::Function, x::T, step::Real) where T<:AbstractFloat

Estimate the derivative of `f` at `x` using the finite differencing method `m` and a given
step size.

# Arguments
- `f::Function`: Function to estimate derivative of.
- `x::T`: Input to estimate derivative at.
- `step::Real`: Step size.

# Returns
- Estimate of the derivative.

# Examples

```julia-repl
julia> fdm = central_fdm(5, 1)
FiniteDifferenceMethod:
  order of method:       5
  order of derivative:   1
  grid:                  [-2, -1, 0, 1, 2]
  coefficients:          [0.08333333333333333, -0.6666666666666666, 0.0, 0.6666666666666666, -0.08333333333333333]

julia> fdm(sin, 1, 1e-3)
 0.5403023058679624

julia> fdm(sin, 1, 1e-3) - cos(1)  # Check the error.
-1.7741363933510002e-13
```
"""
# We loop over all concrete subtypes of `FiniteDifferenceMethod` for 1.0 compatibility.
for T in (UnadaptedFiniteDifferenceMethod, AdaptedFiniteDifferenceMethod)
    @eval begin
        function (m::$T{P,Q})(f::TF, x::Real, step::Real) where {P,Q,TF<:Function}
            x = float(x)  # Assume that converting to float is desired, if it isn't already.
            fs = _eval_function(m, f, x, step)
            return _compute_estimate(m, fs, x, step, m.coefs)
        end
    end
end

function _eval_function(
    m::FiniteDifferenceMethod, f::TF, x::T, step::Real,
) where {TF<:Function,T<:AbstractFloat}
    return f.(x .+ T(step) .* m.grid)
end

function _compute_estimate(
    m::FiniteDifferenceMethod{P,Q},
    fs::SVector{P,TF},
    x::T,
    step::Real,
    coefs::SVector{P,Float64},
) where {P,Q,TF,T<:AbstractFloat}
    # If we substitute `T.(coefs)` in the expression below, then allocations occur. We
    # therefore perform the broadcasting first. See
    # https://github.com/JuliaLang/julia/issues/39151.
    _coefs = T.(coefs)
    return sum(fs .* _coefs) ./ T(step)^Q
end

# Check the method and derivative orders for consistency.
function _check_p_q(p::Integer, q::Integer)
    q >= 0 || throw(DomainError(q, "order of derivative (`q`) must be non-negative"))
    q < p || throw(DomainError(
        (q, p),
        "order of the method (`p`) must be strictly greater than that of the derivative " *
        "(`q`)",
    ))
end

function _coefs(grid, p, q)
    # For high precision on the `\`, we use `Rational`, and to prevent overflows we use
    # `BigInt`. At the end we go to `Float64` for fast floating point math, rather than
    # rational math.
    C = [Rational{BigInt}(g^i) for i in 0:(p - 1), g in grid]
    x = zeros(Rational{BigInt}, p)
    x[q + 1] = factorial(big(q))
    return SVector{p}(Float64.(C \ x))
end

const _COEFFS_MULTS_CACHE = Dict{
    Tuple{SVector,Integer},  # Keys: (grid, q)
    # Values: (coefs, coefs_neighbourhood, ∇f_magnitude_mult, f_error_mult)
    Tuple{SVector,Tuple{Vararg{SVector}},Float64,Float64}
}()

# Compute coefficients for the method and cache the result.
function _coefs_mults(grid::SVector{P, Int}, q::Integer) where P
    return get!(_COEFFS_MULTS_CACHE, (grid, q)) do
        # Compute coefficients for derivative estimate.
        coefs = _coefs(grid, P, q)
        # Compute coefficients for a neighbourhood estimate.
        if all(grid .>= 0)
            coefs_neighbourhood = (
                _coefs(grid .- 2, P, q),
                _coefs(grid .- 1, P, q),
                _coefs(grid, P, q)
            )
        elseif all(grid .<= 0)
            coefs_neighbourhood = (
                _coefs(grid, P, q),
                _coefs(grid .+ 1, P, q),
                _coefs(grid .+ 2, P, q)
            )
        else
            coefs_neighbourhood = (
                _coefs(grid .- 1, P, q),
                _coefs(grid, P, q),
                _coefs(grid .+ 1, P, q)
            )
        end
        # Compute multipliers.
        ∇f_magnitude_mult = sum(abs.(coefs .* grid .^ P)) / factorial(big(P))
        f_error_mult = sum(abs.(coefs))
        return coefs, coefs_neighbourhood, ∇f_magnitude_mult, f_error_mult
    end
end

function Base.show(
    io::IO,
    m::MIME"text/plain",
    x::FiniteDifferenceMethod{P, Q},
) where {P, Q}
    @printf io "FiniteDifferenceMethod:\n"
    @printf io "  order of method:       %d\n" P
    @printf io "  order of derivative:   %d\n" Q
    @printf io "  grid:                  %s\n" x.grid
    @printf io "  coefficients:          %s\n" x.coefs
end

"""
    function estimate_step(
        m::FiniteDifferenceMethod,
        f::Function,
        x::T
    ) where T<:AbstractFloat

Estimate the step size for a finite difference method `m`. Also estimates the error of the
estimate of the derivative.

# Arguments
- `m::FiniteDifferenceMethod`: Finite difference method to estimate the step size for.
- `f::Function`: Function to evaluate the derivative of.
- `x::T`: Point to estimate the derivative at.

# Returns
- `Tuple{<:AbstractFloat, <:AbstractFloat}`: Estimated step size and an estimate of the
    error of the finite difference estimate. The error will be `NaN` if the method failed
    to estimate the error.
"""
function estimate_step(
    m::UnadaptedFiniteDifferenceMethod, f::TF, x::T,
) where {TF<:Function,T<:AbstractFloat}
    step, acc = _compute_step_acc_default(m, x)
    return _limit_step(m, x, step, acc)
end
function estimate_step(
    m::AdaptedFiniteDifferenceMethod{P,Q}, f::TF, x::T,
) where {P,Q,TF<:Function,T<:AbstractFloat}
    ∇f_magnitude, f_magnitude = _estimate_magnitudes(m.bound_estimator, f, x)
    if ∇f_magnitude == 0.0 || f_magnitude == 0.0
        step, acc = _compute_step_acc_default(m, x)
    else
        step, acc = _compute_step_acc(m, ∇f_magnitude, eps(f_magnitude))
    end
    return _limit_step(m, x, step, acc)
end

function _estimate_magnitudes(
    m::FiniteDifferenceMethod{P,Q}, f::TF, x::T,
) where {P,Q,TF<:Function,T<:AbstractFloat}
    step = first(estimate_step(m, f, x))
    fs = _eval_function(m, f, x, step)
    # Estimate magnitude of `∇f` in a neighbourhood of `x`.
    ∇fs = SVector{3}(
        _compute_estimate(m, fs, x, step, m.coefs_neighbourhood[1]),
        _compute_estimate(m, fs, x, step, m.coefs_neighbourhood[2]),
        _compute_estimate(m, fs, x, step, m.coefs_neighbourhood[3])
    )
    ∇f_magnitude = maximum(maximum.(abs, ∇fs))
    # Estimate magnitude of `f` in a neighbourhood of `x`.
    f_magnitude = maximum(maximum.(abs, fs))
    return ∇f_magnitude, f_magnitude
end

function _compute_step_acc_default(m::FiniteDifferenceMethod, x::T) where {T<:AbstractFloat}
    # Compute a default step size using a heuristic and [`DEFAULT_CONDITION`](@ref).
    return _compute_step_acc(m, m.condition, eps(T))
end

function _compute_step_acc(
    m::FiniteDifferenceMethod{P,Q}, ∇f_magnitude::Real, f_error::Real,
) where {P,Q}
    # Set the step size by minimising an upper bound on the error of the estimate.
    C₁ = f_error * m.f_error_mult * m.factor
    C₂ = ∇f_magnitude * m.∇f_magnitude_mult
    step = (Q / (P - Q) * (C₁ / C₂))^(1 / P)
    # Estimate the accuracy of the method.
    acc = C₁ * step^(-Q) + C₂ * step^(P - Q)
    return step, acc
end

function _limit_step(
    m::FiniteDifferenceMethod, x::T, step::Real, acc::Real,
) where {T<:AbstractFloat}
    # First, limit the step size based on the maximum range.
    step_max = m.max_range / maximum(abs.(m.grid))
    if step > step_max
        step = step_max
        acc = NaN
    end
    # Second, prevent very large step sizes, which can occur for high-order methods or
    # slowly-varying functions.
    step_default, _ = _compute_step_acc_default(m, x)
    step_max_default = 1000step_default
    if step > step_max_default
        step = step_max_default
        acc = NaN
    end
    return step, acc
end

for direction in [:forward, :central, :backward]
    fdm_fun = Symbol(direction, "_fdm")
    grid_fun = Symbol("_", direction, "_grid")
    @eval begin
        function $fdm_fun(
            p::Int,
            q::Int;
            adapt::Int=1,
            condition::Real=DEFAULT_CONDITION,
            factor::Real=DEFAULT_FACTOR,
            max_range::Real=Inf,
            geom::Bool=false
        )
            _check_p_q(p, q)
            grid = $grid_fun(p)
            geom && (grid = _exponentiate_grid(grid))
            coefs, coefs_nbhd, ∇f_magnitude_mult, f_error_mult = _coefs_mults(grid, q)
            if adapt >= 1
                bound_estimator = $fdm_fun(
                    # We need to increase the order by two to be able to estimate the
                    # magnitude of the derivative of `f` in a neighbourhood of `x`.
                    p + 2,
                    p;
                    adapt=adapt - 1,
                    condition=condition,
                    factor=factor,
                    max_range=max_range,
                    geom=geom
                )
                return AdaptedFiniteDifferenceMethod{p, q, typeof(bound_estimator)}(
                    grid,
                    coefs,
                    coefs_nbhd,
                    condition,
                    factor,
                    max_range,
                    ∇f_magnitude_mult,
                    f_error_mult,
                    bound_estimator
                )
            else
                return UnadaptedFiniteDifferenceMethod{p, q}(
                    grid,
                    coefs,
                    coefs_nbhd,
                    condition,
                    factor,
                    max_range,
                    ∇f_magnitude_mult,
                    f_error_mult
                )
            end
        end

        @doc """
    $($(Meta.quot(fdm_fun)))(
        p::Int,
        q::Int;
        adapt::Int=1,
        condition::Real=DEFAULT_CONDITION,
        factor::Real=DEFAULT_FACTOR,
        max_range::Real=Inf,
        geom::Bool=false
    )

Contruct a finite difference method at a $($(Meta.quot(direction))) grid of `p` points.

# Arguments
- `p::Int`: Number of grid points.
- `q::Int`: Order of the derivative to estimate.

# Keywords
- `adapt::Int=1`: Use another finite difference method to estimate the magnitude of the
    `p`th order derivative, which is important for the step size computation. Recurse
    this procedure `adapt` times.
- `condition::Real`: Condition number. See [`DEFAULT_CONDITION`](@ref).
- `factor::Real`: Factor number. See [`DEFAULT_FACTOR`](@ref).
- `max_range::Real=Inf`: Maximum distance that a function is evaluated from the input at
    which the derivative is estimated.
- `geom::Bool`: Use geometrically spaced points instead of linearly spaced points.

# Returns
- `FiniteDifferenceMethod`: The specified finite difference method.
        """ $fdm_fun
    end
end

_forward_grid(p::Int) = SVector{p}(0:(p - 1))

_backward_grid(p::Int) = SVector{p}((1 - p):0)

function _central_grid(p::Int)
    if isodd(p)
        return SVector{p}(div(1 - p, 2):div(p - 1, 2))
    else
        return SVector{p}((div(-p, 2):-1)..., (1:div(p, 2))...)
    end
end

function _exponentiate_grid(grid::SVector{P,Int}, base::Int=3) where P
    return sign.(grid) .* div.(base .^ abs.(grid), base)
end

function _is_symmetric(
    vec::SVector{P};
    centre_zero::Bool=false,
    negate_half::Bool=false
) where P
    half_sign = negate_half ? -1 : 1
    if isodd(length(vec))
        centre_zero && vec[end ÷ 2 + 1] != 0 && return false
        return vec[1:end ÷ 2] == half_sign .* reverse(vec[(end ÷ 2 + 2):end])
    else
        return vec[1:end ÷ 2] == half_sign .* reverse(vec[(end ÷ 2 + 1):end])
    end
end

function _is_symmetric(m::FiniteDifferenceMethod)
    grid_symmetric = _is_symmetric(m.grid, centre_zero=true, negate_half=true)
    coefs_symmetric =_is_symmetric(m.coefs, negate_half=true)
    return grid_symmetric && coefs_symmetric
end

"""
    extrapolate_fdm(
        m::FiniteDifferenceMethod,
        f::Function,
        x::Real,
        initial_step::Real=10,
        power::Int=1,
        breaktol::Real=Inf,
        kw_args...
    ) where T<:AbstractFloat

Use Richardson extrapolation to extrapolate a finite difference method.

Takes further in keyword arguments for `Richardson.extrapolate`. This method
automatically sets `power = 2` if `m` is symmetric and `power = 1`. Moreover, it defaults
`breaktol = Inf`.

# Arguments
- `m::FiniteDifferenceMethod`: Finite difference method to estimate the step size for.
- `f::Function`: Function to evaluate the derivative of.
- `x::Real`: Point to estimate the derivative at.
- `initial_step::Real=10`: Initial step size.

# Returns
- `Tuple{<:AbstractFloat, <:AbstractFloat}`: Estimate of the derivative and error.
"""
function extrapolate_fdm(
    m::FiniteDifferenceMethod,
    f::Function,
    x::Real,
    initial_step::Real=10;
    power::Int=1,
    breaktol::Real=Inf,
    kw_args...
) where T<:AbstractFloat
    (power == 1 && _is_symmetric(m)) && (power = 2)
    return extrapolate(
        step -> m(f, x, step),
        float(initial_step);
        power=power,
        breaktol=breaktol,
        kw_args...
    )
end


text:
# Barycenter solver

struct SinkhornBarycenterSolver{A<:Sinkhorn,M,CT,W,E<:Real,T<:Real,R<:Real,C1,C2}
    source::M
    C::CT
    eps::E
    w::W
    alg::A
    atol::T
    rtol::R
    maxiter::Int
    check_convergence::Int
    cache::C1
    convergence_cache::C2
end

function build_solver(
    μ::AbstractMatrix,
    C::AbstractMatrix,
    ε::Real,
    w::AbstractVector,
    alg::Sinkhorn;
    atol=nothing,
    rtol=nothing,
    check_convergence=10,
    maxiter::Int=1_000,
)
    # check that input marginals are balanced
    checkbalanced(μ)

    size2 = (size(μ, 2),)

    # compute type
    T = float(Base.promote_eltype(μ, one(eltype(C)) / ε))

    # build caches using SinkhornGibbsCache struct (since there is no dependence on ν)
    cache = build_cache(T, alg, size2, μ, C, ε)
    convergence_cache = build_convergence_cache(T, size2, μ)

    # set tolerances
    _atol = atol === nothing ? 0 : atol
    _rtol = rtol === nothing ? (_atol > zero(_atol) ? zero(T) : sqrt(eps(T))) : rtol

    # create solver
    solver = SinkhornBarycenterSolver(
        μ, C, ε, w, alg, _atol, _rtol, maxiter, check_convergence, cache, convergence_cache
    )
    return solver
end

"""
    sinkhorn_barycenter(μ, C, ε, w, alg = SinkhornGibbs(); kwargs...)

Compute the Sinkhorn barycenter for a collection of `N` histograms contained in the columns of `μ`, for a cost matrix `C` of size `(size(μ, 1), size(μ, 1))`, relative weights `w` of size `N`, and entropic regularisation parameter `ε`.
Returns the entropically regularised barycenter of the `μ`, i.e. the histogram `ρ` of length `size(μ, 1)` that solves

```math
\\min_{\\rho \\in \\Sigma} \\sum_{i = 1}^N w_i \\operatorname{OT}_{\\varepsilon}(\\mu_i, \\rho)
```

where ``\\operatorname{OT}_{ε}(\\mu, \\nu) = \\inf_{\\gamma \\Pi(\\mu, \\nu)} \\langle \\gamma, C \\rangle + \\varepsilon \\Omega(\\gamma)``
is the entropic optimal transport loss with cost ``C`` and regularisation ``\\epsilon``.
"""
function sinkhorn_barycenter(μ, C, ε, w, alg::Sinkhorn; kwargs...)
    solver = build_solver(μ, C, ε, w, alg; kwargs...)
    solve!(solver)
    return solution(solver)
end

ext: jl


text:
using Test
using Altro
using BenchmarkTools
using TrajectoryOptimization
using RobotDynamics
using RobotZoo
using StaticArrays, LinearAlgebra
using JLD2
using FileIO
using FiniteDiff
using Random
const RD = RobotDynamics
const TO = TrajectoryOptimization

TEST_TIME = false
Random.seed!(1)

@testset "Benchmark Problems" begin
    if !haskey(ENV, "CI")
        include("benchmark_problems.jl")
    end
    include("nl_cartpole.jl")
    include("cartpole.jl")
    include("quadrotor.jl")
    include(joinpath(@__DIR__,"..","examples","quickstart.jl"))
end

@testset "Solvers" begin
    include("constructors.jl")
    include("augmented_lagrangian_tests.jl")
    include("solve_tests.jl")
    include("finite_diff.jl")
    include("socp_test.jl")
end

@testset "Solver Options" begin
    include("solver_opts.jl")
    include("solver_stats.jl")
end


text:
@enum ColumnType::UInt8 begin
    ColumnTypeByte = 0
    ColumnTypeUByte = 1
    ColumnTypeBool = 2
    ColumnTypeShort = 3
    ColumnTypeUShort = 4
    ColumnTypeInt = 5
    ColumnTypeUInt = 6
    ColumnTypeLong = 7
    ColumnTypeULong = 8
    ColumnTypeFloat = 9
    ColumnTypeDouble = 10
    ColumnTypeString = 11
    ColumnTypeJson = 12
    ColumnTypeDateTime = 13
    ColumnTypeBinary = 14
end

FlatBuffers.@with_kw mutable struct Column
    name::String = ""
    type::ColumnType = 0
    title::String = ""
    description::String = ""
    width::Int32 = -1
    precision::Int32 = -1
    scale::Int32 = -1
    nullable::Bool = true
    unique::Bool = false
    primary_key::Bool = false
    metadata::String = ""
end
FlatBuffers.@ALIGN(Column, 1)
FlatBuffers.slot_offsets(::Type{T}) where {T <: Column} = [
    0x00000004, 0x00000006, 0x00000008, 0x0000000A,
    0x0000000C, 0x0000000E, 0x00000010, 0x00000012,
    0x00000014, 0x00000016, 0x00000018
]

Column(buf::AbstractVector{UInt8}) = FlatBuffers.read(Column, buf)
Column(io::IO) = FlatBuffers.deserialize(io, Column)

@enum GeometryType::UInt8 begin
    GeometryTypeUnknown = 0
    GeometryTypePoint = 1
    GeometryTypeLineString = 2
    GeometryTypePolygon = 3
    GeometryTypeMultiPoint = 4
    GeometryTypeMultiLineString = 5
    GeometryTypeMultiPolygon = 6
    GeometryTypeGeometryCollection = 7
    GeometryTypeCircularString = 8
    GeometryTypeCompoundCurve = 9
    GeometryTypeCurvePolygon = 10
    GeometryTypeMultiCurve = 11
    GeometryTypeMultiSurface = 12
    GeometryTypeCurve = 13
    GeometryTypeSurface = 14
    GeometryTypePolyhedralSurface = 15
    GeometryTypeTIN = 16
    GeometryTypeTriangle = 17
end

FlatBuffers.@with_kw mutable struct Crs
    org::String = ""
    code::Int32 = 0
    name::String = ""
    description::String = ""
    wkt::String = ""
    code_string::String = ""
end
FlatBuffers.@ALIGN(Crs, 1)
FlatBuffers.slot_offsets(::Type{T}) where {T <: Crs} = [
    0x00000004, 0x00000006, 0x00000008, 0x0000000A,
    0x0000000C, 0x0000000E
]

Crs(buf::AbstractVector{UInt8}) = FlatBuffers.read(Crs, buf)
Crs(io::IO) = FlatBuffers.deserialize(io, Crs)

FlatBuffers.@with_kw mutable struct Header
    name::String = ""
    envelope::Vector{Float64} = []
    geometry_type::GeometryType = 0
    has_z::Bool = false
    has_m::Bool = false
    has_t::Bool = false
    has_tm::Bool = false
    columns::Vector{Column} = []
    features_count::UInt64 = 0
    index_node_size::UInt16 = 16
    crs::Crs = Crs()
    title::String = ""
    description::String = ""
    metadata::String = ""
end
FlatBuffers.@ALIGN(Header, 1)
FlatBuffers.slot_offsets(::Type{T}) where {T <: Header} = [
    0x00000004, 0x00000006, 0x00000008, 0x0000000A,
    0x0000000C, 0x0000000E, 0x00000010, 0x00000012,
    0x00000014, 0x00000016, 0x00000018, 0x0000001A,
    0x0000001C, 0x0000001E
]
FlatBuffers.root_type(::Type{T}) where {T <: Header} = true

Header(buf::AbstractVector{UInt8}) = FlatBuffers.read(Header, buf)
Header(io::IO) = FlatBuffers.deserialize(io, Header)


text:
#####
##### Object
#####

# Not to be confused with the "objects" in the `ObjectCache` above.
mutable struct Object{T}
    ptr::Ptr{Nothing}
    manager::T

    # Inner constructor to ensure finalizers are attached>
    function Object(ptr::Ptr{Nothing}, manager::T) where {T}
        object = new{T}(ptr, manager)
        if !isnull(ptr)
            unsafe_register!(manager, object)
            finalizer(free, object)
        end
        return object
    end
end

Base.pointer(object::Object) = object.ptr
unsafe_pointer(object::Object) = object.ptr
blockpointer(object::Object) = pointer_from_objref(object)
metadata(object::Object) = unsafe_block(unsafe_pointer(object))

free(object::Object) = free(manager(object), unsafe_pointer(object))
manager(object::Object) = object.manager
getid(o::Object) = getid(metadata(o))

"""
$(TYPEDSIGNATURES)

Allocate `bytes` from `objects`'s manager.
If `id` is not given, it will be selected automatically.
"""
function alloc(
    object::Object,
    bytes::Integer,
    priority::AllocationPriority = PreferLocal,
    id = getid(object.manager),
)
    return alloc(manager(object), bytes, priority, id)
end


text:
module ControlAllocationSwitching

using FlightSims
const FS = FlightSims
using FaultTolerantControl
const FTC = FaultTolerantControl
using Plots
using LinearAlgebra
using Transducers
using DifferentialEquations
using UnPack
using ReferenceFrameRotations
using JLD2, FileIO
using Printf
using NumericalIntegration

export draw_figures
export simulation

include("figures.jl")
include("simulation.jl")

end


text:
# Automatically generated using Clang.jl

@cenum cutensorOperator_t::UInt32 begin
    CUTENSOR_OP_IDENTITY = 1
    CUTENSOR_OP_SQRT = 2
    CUTENSOR_OP_RELU = 8
    CUTENSOR_OP_CONJ = 9
    CUTENSOR_OP_RCP = 10
    CUTENSOR_OP_ADD = 3
    CUTENSOR_OP_MUL = 5
    CUTENSOR_OP_MAX = 6
    CUTENSOR_OP_MIN = 7
    CUTENSOR_OP_UNKNOWN = 126
end

@cenum cutensorStatus_t::UInt32 begin
    CUTENSOR_STATUS_SUCCESS = 0
    CUTENSOR_STATUS_NOT_INITIALIZED = 1
    CUTENSOR_STATUS_ALLOC_FAILED = 3
    CUTENSOR_STATUS_INVALID_VALUE = 7
    CUTENSOR_STATUS_ARCH_MISMATCH = 8
    CUTENSOR_STATUS_MAPPING_ERROR = 11
    CUTENSOR_STATUS_EXECUTION_FAILED = 13
    CUTENSOR_STATUS_INTERNAL_ERROR = 14
    CUTENSOR_STATUS_NOT_SUPPORTED = 15
    CUTENSOR_STATUS_LICENSE_ERROR = 16
    CUTENSOR_STATUS_CUBLAS_ERROR = 17
    CUTENSOR_STATUS_CUDA_ERROR = 18
    CUTENSOR_STATUS_INSUFFICIENT_WORKSPACE = 19
    CUTENSOR_STATUS_INSUFFICIENT_DRIVER = 20
end

@cenum cutensorAlgo_t::Int32 begin
    CUTENSOR_ALGO_GETT = -4
    CUTENSOR_ALGO_TGETT = -3
    CUTENSOR_ALGO_TTGT = -2
    CUTENSOR_ALGO_DEFAULT = -1
end

@cenum cutensorComputeType_t::Int32 begin
    CUTENSOR_R_MIN_16F = 1
    CUTENSOR_C_MIN_16F = 2
    CUTENSOR_R_MIN_32F = 4
    CUTENSOR_C_MIN_32F = 8
    CUTENSOR_R_MIN_64F = 16
    CUTENSOR_C_MIN_64F = 32
    CUTENSOR_R_MIN_8U  = 64
    CUTENSOR_R_MIN_32U = 128
    CUTENSOR_R_MIN_8I  = 256
    CUTENSOR_R_MIN_32I = 512
end

@cenum cutensorWorksizePreference_t::UInt32 begin
    CUTENSOR_WORKSPACE_MIN = 1
    CUTENSOR_WORKSPACE_RECOMMENDED = 2
    CUTENSOR_WORKSPACE_MAX = 3
end

const cutensorHandle_t = Ptr{Cvoid}
const cutensorTensorDescriptor_t = Ptr{Cvoid}
const CUTENSOR_MAJOR = 1
const CUTENSOR_MINOR = 0
const CUTENSOR_PATCH = 0


text:
# The SteinMaternUnboundedPowerKernel has of the form:
#
# k(x, y) = ((1 + ||x||^alpha)*(1 + ||y||^alpha)) * km(x,y)
#
# where ||.|| is the L2 norm and km is the matern kernel.

type SteinMaternUnboundedPowerKernel <: SteinMaternWeightedKernel
    # the alpha parameter
    alpha::Float64
    # the gamma parameter
    gamma::Float64
end

# default is use gamma = sqrt{3} and alpha = 2.5
SteinMaternUnboundedPowerKernel() = SteinMaternUnboundedPowerKernel(2.5, sqrt(3.0))

# default is use gamma sqrt{3}
SteinMaternUnboundedPowerKernel(alpha::Float64) = SteinMaternUnboundedPowerKernel(alpha, sqrt(3.0))

# The term (1 + gamma*||z||^alpha)
function _Qz(ker::SteinMaternUnboundedPowerKernel, z::Array{Float64, 1})
    1.0 + norm(z,2)^ker.alpha
end

# The term ((1 + gamma*||x||^alpha)*(1 + gamma*||y||^alpha))^-1
function Q(ker::SteinMaternUnboundedPowerKernel, x::Array{Float64, 1}, y::Array{Float64, 1})
    _Qz(ker, x) * _Qz(ker, y)
end

function gradxQ(ker::SteinMaternUnboundedPowerKernel, x::Array{Float64,1}, y::Array{Float64,1})
    fy = _Qz(ker, y)
    mx = norm(x,2)^(ker.alpha - 2.0)
    ker.alpha * mx * fy * x
end

function gradxyQ(ker::SteinMaternUnboundedPowerKernel, x::Array{Float64,1}, y::Array{Float64,1})
    mx = norm(x,2)^(ker.alpha - 2)
    my = norm(y,2)^(ker.alpha - 2)
    ker.alpha^2 * mx * my * dot(x, y)
end


text:
### large number of functions for loading, analysing and plotting genomic data e.g. ChIP-seq, ATAC-seq etc.

function loaddatasets(file="chip_atac_datasets.tsv", projdir=getprojectdir())
    meta = CSV.read(joinpath(projdir, "data", file), DataFrame)

    sort!(meta, [:Study, :Target, :Cell, :Rep])

    meta = @subset(meta, :Target .!= "Input");
    meta.Index = 1:size(meta, 1)
    meta
end

#### Functions to generate pileups of Chip/ATAC at loci and cache results
pile(chrom, loc, FM::FragMatrixSingle{T}, fl=120) where{T}  = 1e+6*fragregion(chrom, loc, "+", FM, inc_meta_frag!, fl, 4, one)/totalfrags(FM)
pile(chrom, loc, FM::FragMatrixPair{T}, minfrag=0, maxfrag=1000) where{T}  = 1e+6*fragregion(chrom, loc, FM, inc_meta_frag!, minfrag, maxfrag, 4, one)/totalfrags(FM)

function loadpiles(file, labels)
    df = CSV.read(file, DataFrame)
    @assert names(df) == labels
    collect(eachcol(df))
end

function piles_saveload(chrom, loc, inds, labels, FM::Nothing, file; saveload=true)
    !isfile(file) && error("Error: precalcuated $file not found, load alignments to regenerate")
    if saveload == false
        @warn "Loading precalculated signal tracks set to false, but alignments not loaded, loading precalculated signals"
    end
    loadpiles(file, labels)
end
function piles_saveload(chrom, loc, inds, labels, FM, file; saveload=true)
    saveload && isfile(file) && return loadpiles(file, labels)
    P = @showprogress map(ind -> mapreduce(i -> vec(pile(chrom, loc, FM[i])), +, ind), inds)

    if saveload
        mkpath(dirname(file))
        CSV.write(file, DataFrame(P, labels), delim='\t', compress=true)
    end
    P
end
function piles(chrom, loc, inds, grouplabel, labels, FM; saveload=true, projdir=getprojectdir())
    file = joinpath(projdir, "results", "piles", string("pile_", grouplabel, "_", chrom, "_", first(loc), "_", last(loc), ".tsv.gz"))
    piles_saveload(chrom, loc, inds, labels, FM, file, saveload=saveload)
end
piles(regionlabel, grouplabel, datagroup::DataFrame, regioncoords, FM ; saveload=true) = piles(regioncoords[regionlabel]..., datagroup.Inds, grouplabel, datagroup.Label, FM, saveload=saveload)

### generate a vector plots of pile ups from different data sources
function prp(chrom, loc, datagroup, P, regioncoords; dx=50, annot=true, regspace=0, plotrec=true, plotline=false, cc=[], ylm=-1)

    labels = string.(datagroup.Cell, " ", datagroup.Target)

    pp = first(loc):dx:last(loc)
    phs = Plots.Plot[]
    for (i, p) in enumerate(P)
        c = isempty(cc) ? :auto : cc[i]

        q = plot()
        if ylm != -1
            yl = (0, ylm)
            plot!(ylims=(yl))
        else
            yl = ylims()
        end

        ### Mark pileup with ROI containing variants
        if plotrec && haskey(regioncoords, "roi")
            roi = last(regioncoords["roi"])
            plot!(rectangle(last(roi) - first(roi), yl[2] - yl[1], first(roi), 0), fillcolor=:lightgrey, linecolor=nothing, lab="", fillalpha=0.5)
        end
        if plotline && haskey(regioncoords, "criticalregion")
            cr = last(regioncoords["criticalregion"])
            vline!([first(cr), last(cr)], c=:black, lab="", ls=:dot)
        end
        plot!(pp, avh(p, dx), lab="", fill=0, c=c)
        plot!(ylims=(yl))

        annot && annotate!((first(loc), (yl[1] + yl[2])/2, text(labels[i], font(:left, 10, :middle, "helvetica"))))
        if i < length(P)
            plot!(q, xticks=false, bottom_margin=-3Measures.mm)
        end
        push!(phs, q)
    end
    phs
end

##################### functions to generate different genomics panels plots

####### histones #######
function calc_histonepanel(regionlabel, datagroup, regioncoords, FM, genecoords, saveload=true)
    datahistone = @subset(datagroup, occursin.(r"H3K27ac|H3K27me3|H3K4me1|H3K4me3", :Target), :Study .!= "PPCebola")

    cellorder = Dict("ES" => 1, "DE" => 2, "GT" => 3, "PP1" => 4, "PP2" => 5, "EndoCβH1" => 6, "islets" => 7, "Alpha" => 8, "Beta" => 9, "Exocrine" => 10)
    datahistone.CellOrder = getindex.(Ref(cellorder), datahistone.Cell)
    sort!(datahistone, [:Target, :CellOrder])
    datahistone.GroupIndex = 1:size(datahistone, 1)

    PP = piles(regionlabel, "histone", datahistone, regioncoords, FM, saveload=saveload)
    p = histonepanel(regioncoords[regionlabel]..., datahistone, PP, genecoords, regioncoords, cc=cgrad(:viridis, length(unique(datahistone.Cell)) + 2, categorical=true)[2:end-1])
end

function histonepanel(chrom, loc, datahistone, PH, genecoords, regioncoords; dxs=[500, 500, 1000, 500], cc=cgrad(:viridis, length(unique(datahistone.Cell)), categorical=true))
    histoneorder = ["H3K27ac", "H3K4me1", "H3K27me3"]
    ylms = [1.0, 1.5, 1.0, 1]
    phs = Vector{Plots.Plot}[]
    for (h, ylm, dx) in zip(histoneorder, ylms, dxs)
        dh = @subset(datahistone, :Target .== h)
        ps = prp(chrom, loc, dh, PH[dh.GroupIndex], regioncoords, dx=dx, annot=false, regspace=2000, cc=cc)

        for (p, l) in zip(ps, dh.Cell)
            plot!(p, ylims=[0, ylm], yticks=0:1)
            if l == "ES"
                plot!(title=h)
            end
            if h == first(histoneorder)
                plot!(p, ylabel=l, left_margin=5Measures.mm)
            end
        end

        push!(phs, ps)
    end

    mp = maximum(length.(phs))

    ### fill in grid with empty plots
    for ps in phs
        for i = (length(ps)+1):mp
            plot!(ps[end], xticks=false, bottom_margin=-3Measures.mm)
            p = plot(xlims=xlims(ps[end]), yaxis=false, yticks=false, framestyle=:none)
            push!(ps, p)
        end
    end

    for ps in phs
        plot!(ps[end], xticks=false, bottom_margin=-3Measures.mm)
        p = plotgenemodels(loc, @subset(genecoords, :GeneName .== "HK1"))
        plot!(ylims=(-1, 2), yticks=false, yaxis=false, xticks=false, xaxis=false, xlabel=string(chrom, ":", first(loc), "-", last(loc)))
        push!(ps, p)
    end

    plot(permutedims(reduce(hcat, phs))..., layout=(length(phs[1]), length(histoneorder)), link=:x, fontfamily="helvetica", xlims=(first(loc), last(loc)), size=(1000, 400), grid=false)

end

####### pancreatic progenitors #######

function calc_progentior_panel(regionlabel, datagroup, regioncoords, FM, genecoords)
    dataprog = @subset(datagroup, .!occursin.(r"^H[23]", :Target), :DataType .== "ChIP", :Study .!= "EndoCChIP", :Cell .∉ Ref(["DE", "ES", "GT", "islets"]), :Target .∉ Ref(["CTCF"]),
    (:Study .!= "PPCebola") .| (:Target .== "TEAD1") .| (:Cell .== "LiverBud"))

    dataprog = [sort(@subset(dataprog, .!occursin.(r"FOX", :Target)), :Target) ; sort(@subset(dataprog, occursin.(r"FOX", :Target)), [:Target, order(:Cell, rev=true)])]
    dataprog.GroupIndex = 1:size(dataprog, 1)

    PP = piles(regionlabel, "pancprog", dataprog, regioncoords, FM)
    pancpropanel(regioncoords[regionlabel]..., dataprog, PP, genecoords, regioncoords)
    p = plot!(size=(492, 400))# xlims=(71065952-5000,71065952+5000))

end

function pancpropanel(chrom, loc, dataprog, PP, genecoords, hk1coords)

    yl = (0, 3)
    theme(:wong)
    cc= getindex.(Ref(Dict(u => i for (i, u) in enumerate(unique(dataprog.Target)))), dataprog.Target)
    cc = replace(cc, 4 => maximum(cc)+1)
    phs = prp(chrom, loc, dataprog, PP, hk1coords, dx=100, annot=false, ylm=yl[2], cc=cc)
    labels = string.(dataprog.Target, " ", ifelse.(occursin.("PP", dataprog.Cell), "PP", "Liver Bud"))

    for (i, p) in enumerate(phs)
        annotate!(p, [(first(loc), (yl[1] + yl[2])/2, text(labels[i], font(:left, 10, :middle, "helvetica")))])
        plot!(p, yticks=[0, 3])
    end

    plot!(phs[end], xticks=false, bottom_margin=-3Measures.mm)

    p = plotgenemodels(loc, @subset(genecoords, :GeneName .== "HK1"))
    plot!(ylims=(-1, 2), yticks=false, yaxis=false, xticks=false,  xaxis=false, xlabel=string(chrom, ":", first(loc), "-", last(loc)))
    push!(phs, p)

    plot(phs..., layout=(length(phs), 1), size=(700, 400), grid=false, link=:x, xlims=(first(loc), last(loc)))
end

####### islets #######
function calcisletpanel(regionlabel, datagroup, regioncoords, FM, genecoords)

    dataislets = @subset(datagroup, .!occursin.(r"^H[23]", :Target),
                                    :Study .∉ Ref(["DiffChip", "DiffChipATAC", "EndoCChIP", "PPCebola", "EndoCATAC"]),
                                    (:Study .!= "PancSNATAC") .| (:Cell .== "beta_1"))

    PP = piles(regionlabel, "islet_chip", dataislets, regioncoords, FM)

    ## sort islet data by region in ROI descending

    roi_iv = last(regioncoords["roi"]) .- first(last(regioncoords[regionlabel])) .+ 1
    atac_ind = dataislets.DataType .== "ATAC"

    si = sortperm(map(p -> maximum(p[roi_iv]), PP[.!atac_ind]), rev=true)
    PP = [PP[.!atac_ind][si] ; PP[atac_ind]]
    dataislets = [dataislets[.!atac_ind, :][si, :] ; dataislets[atac_ind, :]]

    isletpanel(regioncoords[regionlabel]..., dataislets, PP, genecoords, regioncoords)
    plot!(size=(480, 350))
end

function isletpanel(chrom, loc, datagroup, PP, genecoords, hk1coords)

    yl = (0, 10)
    theme(:wong)
    cc= getindex.(Ref(Dict(u => i for (i, u) in enumerate(unique(datagroup.Target)))), datagroup.Target)
    cc = replace(cc, 4 => maximum(cc)+1)
    phs = prp(chrom, loc, datagroup, PP, hk1coords, dx=100, annot=false, ylm=yl[2], cc=cc)
    labels = ifelse.(datagroup.DataType .== "ATAC", "snATAC beta-cells", datagroup.Target)

    for (i, p) in enumerate(phs)

        if datagroup.DataType[i] == "ATAC"
            ylt = (0, 3)
        else
            ylt = yl
        end
        plot!(p, yticks=[ylt...], ylims=ylt)
        annotate!(p, [(first(loc), (ylt[1] + ylt[2])/2, text(labels[i], font(:left, 10, :middle, "helvetica")))])
    end

    plot!(phs[end], xticks=false, bottom_margin=-3Measures.mm)

    p = plotgenemodels(loc, @subset(genecoords, :GeneName .== "HK1"))
    plot!(ylims=(-1, 2), yticks=false, yaxis=false, xticks=false, xaxis=false, xlabel=string(chrom, ":", first(loc), "-", last(loc)))
    push!(phs, p)

    plot(phs..., layout=(length(phs), 1), size=(700, 400), grid=false, link=:x, fontfamily="helvetica")
end

####### snATAC #######
function calc_snatachpanel(regionlabel, datagroup, regioncoords, FM, genecoords)
    datasnatac_hh = @subset(datagroup, :Target .== "ATAC", :Study .== "PancSNATAC", occursin.(r"alpha_1|beta_1|delta_1", :Cell))
    datasnatac_hl = @subset(datagroup, :Target .== "ATAC", :Study .== "PancSNATAC", occursin.(r"alpha_2|beta_2|delta_2", :Cell))

    PPH = piles(regionlabel, "snatac_hormone_high", datasnatac_hh, regioncoords, FM)
    PPL = piles(regionlabel, "snatac_hormone_low",  datasnatac_hl, regioncoords, FM)

    snatacpanel(regioncoords[regionlabel]..., datasnatac_hh, datasnatac_hl, PPH, PPL, genecoords, regioncoords)
end

function snatacpanel(chrom, loc, datasnatac_hh, datasnatac_hl, PPH, PPL, genecoords, hk1coords)

    cc = cgrad(:inferno, 4, categorical=false)
    phs_H = prp(hk1coords["genelocus"]..., datasnatac_hh, PPH, hk1coords, dx=100, ylm=5, cc=cc)
    phs_L = prp(hk1coords["genelocus"]..., datasnatac_hl, PPL, hk1coords, dx=100, ylm=5, cc=cc)

    plot!(phs_H[end], xticks=false, bottom_margin=-3Measures.mm)
    plot!(phs_L[end], xticks=false, bottom_margin=-3Measures.mm)
    p = plotgenemodels(last(hk1coords["genelocus"]), @subset(genecoords, :GeneName .== "HK1"))
    plot!(ylims=(-1, 2), yticks=false, yaxis=false, xticks=false,  xaxis=false)

    for (h, l) in zip(phs_L, phs_H)
        plot!(h, yticks=[0, 5])
        plot!(l, yticks=[0, 5])
    end
    push!(phs_L, p)
    push!(phs_H, deepcopy(p))

    plot(mapreduce(collect, vcat, zip(phs_H, phs_L))..., layout=(length(phs_L), 2), size=(1000, 200), grid=false, xlims=(first(loc), last(loc)), link=:x, fontfamily="helvetica")
end

####### atac diff #######
function calcatacdiff_subpanel(regionlabel, datagroup, regioncoords, FM, genecoords)

    data_atac_diff = @subset(datagroup, :DataType .== "ATAC", :Cell .∈ Ref(["ES", "DE", "GT", "PP1", "PP2", "beta_1"]))
    data_atac_diff = sort(data_atac_diff, order(:Cell, by=x -> ifelse(x == "ES", "A", ifelse(x == "EndoC", "Z", x))))

    PP = piles(regionlabel, "atac_diff", data_atac_diff, regioncoords, FM)
    cc = cgrad(:viridis, size(data_atac_diff, 1)+2, categorical=false)

    chrom, loc = regioncoords[regionlabel]
    phs = prp(chrom, loc, data_atac_diff, PP, regioncoords, dx=100, ylm=3.5, cc=cc[2:end-1], annot=false, plotrec=false, plotline=true)
    for (p, l) in zip(phs, data_atac_diff.Cell)
        plot!(p, title=l, yticks=false, left_margin=-1mm)
    end

    plot!(phs[3], xlabel=string(chrom, ":", loc[1], "-", loc[end]))
    plot!(phs[1], yticks=[0, 3], left_margin=5mm)

    p = plot(phs..., layout=(1, length(phs)), size=(325, 100), grid=false, right_margin=-1mm, yaxis=false, bottom_margin=5mm, fontfamily="helvetica",
            tick_direction=:out, xticks=([mean(loc)], ""), titlefont=font(12, "helvetica"))

end

function calc_atacdiff_supp(regionlabel, datagroup, regioncoords, FM, genecoords)

    data_atac_diff = @subset(datagroup, :DataType .== "ATAC", :Cell .∈ Ref(["ES", "DE", "GT", "PP1", "PP2"]))
    data_atac_diff = sort(data_atac_diff, order(:Cell, by=x -> ifelse(x == "ES", "A", ifelse(x == "EndoC", "Z", x))))

    PP = piles(regionlabel, "atac_diff", data_atac_diff, regioncoords, FM)
    cc = cgrad(:viridis, size(data_atac_diff, 1)+1, categorical=false)

    phs = prp(regioncoords[regionlabel]..., data_atac_diff, PP, regioncoords, dx=100, ylm=5, cc=cc)
    plot!(phs[end], xticks=false, bottom_margin=-3Measures.mm)
    for p in phs
        plot!(p, yticks=[0, 5])
    end

    p = plotgenemodels(last(regioncoords[regionlabel]), @subset(genecoords, :GeneName .== "HK1"))
    plot!(ylims=(-1, 2), yticks=false, yaxis=false, xticks=false,  xaxis=false)

    push!(phs, p)
    p = plot(phs..., layout=(length(phs), 1), size=(492, 280), grid=false, link=:x)
end


text:
@printf("\nTesting non linear operators\n")

# Sigmoid
n = 4
x = randn(n)
r = randn(n)
op = Sigmoid(Float64,(n,),2)

y, grad = test_NLop(op,x,r,verb)

n,m,l = 4,5,6
x = randn(n,m)
r = randn(n,m)
op = Sigmoid((n,m),2)

y, grad = test_NLop(op,x,r,verb)

n = 10
x = randn(n)
r = randn(n)
op = SoftMax(Float64,(n,))

y, grad = test_NLop(op,x,r,verb)

# SoftMax
n,m,l = 4,5,6
x = randn(n,m,l)
r = randn(n,m,l)
op = SoftMax(Float64,(n,m,l))

y, grad = test_NLop(op,x,r,verb)

# SoftPlus
n = 10
x = randn(n)
r = randn(n)
op = SoftPlus(Float64,(n,))

n,m,l = 4,5,6
x = randn(n,m,l)
r = randn(n,m,l)
op = SoftPlus(Float64,(n,m,l))

y, grad = test_NLop(op,x,r,verb)

# Exp
n,m,l = 4,5,6
x = randn(n,m,l)
r = randn(n,m,l)
op = Exp(n,m,l)
op = Exp(Float64,(n,m,l))

y, grad = test_NLop(op,x,r,verb)

## Sin
n,m,l = 4,5,6
x = randn(n,m,l)
r = randn(n,m,l)
op = Sin(n,m,l)
op = Sin(Float64,(n,m,l))

y, grad = test_NLop(op,x,r,verb)

## Cos
n,m,l = 4,5,6
x = randn(n,m,l)
r = randn(n,m,l)
op = Cos(n,m,l)
op = Cos(Float64,(n,m,l))

y, grad = test_NLop(op,x,r,verb)

# Atan
n = 10
x = randn(n)
r = randn(n)
op = Atan(n)
op = Atan(Float64,(n,))

y, grad = test_NLop(op,x,r,verb)

# Tanh
n = 10
x = randn(n)
r = randn(n)
op = Tanh(n)
op = Tanh(Float64,(n,))

y, grad = test_NLop(op,x,r,verb)

# Sech
n = 10
x = randn(n)
r = randn(n)
op = Sech(n)
op = Sech(Float64,(n,))

y, grad = test_NLop(op,x,r,verb)

# Pow
n = 10
x = randn(n)
r = randn(n)
op = Pow(Float64,(n,),2)

y, grad = test_NLop(op,x,r,verb)

n = 10
x = abs.(randn(n))
r = abs.(randn(n))
op = Pow(Float64,(n,),0.5)

y, grad = test_NLop(op,x,r,verb)


text:
using Test, MPI
using CLIMA.Mesh.Topologies
using CLIMA.Mesh.Grids
using CLIMA.Mesh.Geometry
using StaticArrays

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

@testset "LocalGeometry" begin
  FT = Float64
  ArrayType = Array

  xmin      = 0
  ymin      = 0
  zmin      = 0
  xmax      = 2000
  ymax      = 400
  zmax      = 2000

  Ne        = (20,2,20)

  polynomialorder = 4

  brickrange = (range(FT(xmin); length=Ne[1]+1, stop=xmax),
                range(FT(ymin); length=Ne[2]+1, stop=ymax),
                range(FT(zmin); length=Ne[3]+1, stop=zmax))
  topl = StackedBrickTopology(MPI.COMM_SELF, brickrange, periodicity = (true, true, false))

  grid = DiscontinuousSpectralElementGrid(topl,
                                          FloatType = FT,
                                          DeviceArray = ArrayType,
                                          polynomialorder = polynomialorder
                                          )

  S = ((xmax-xmin)/(polynomialorder*Ne[1]), (ymax-ymin)/(polynomialorder*Ne[2]), (zmax-zmin)/(polynomialorder*Ne[3]))
  Savg = cbrt(prod(S))
  M = SDiagonal(S.^-2)

  for e in 1:size(grid.vgeo,3)
    for n in 1:size(grid.vgeo,1)
      g = LocalGeometry(Val(polynomialorder), grid.vgeo, n, e)
      @test lengthscale(g) ≈ Savg
      @test Geometry.resolutionmetric(g) ≈ M
    end
  end
end

max_stars_repo_name: akshaysridhar/CLIMA

max_issues_repo_name: akshaysridhar/CLIMA

max_forks_repo_name: akshaysridhar/CLIMA


text:
"""
	SModelsOptions
SModelsOptions is a struct that contains options necessary to train the surrogate
model
"""
struct SModelsOptions
  sModelType::Symbol          # The type of model to use for the surrogate model
	classifierType::Symbol 			# The type of classifier to use
  desiredMaxPerErrorRegressor::Float64 # The desired percentage error of the regressor after which training is stopped
	desiredAccuracyClassifier::Float64 # The desired percentage error of the classifier after which training is stopped
	desiredMinObs::Int64			# Minimum number of points in the train sample before stopping
  nWorkers::Int64          # To store the number of workers
  nbBatches::Int64         # Total number of batches before stopping
  batchSizeWorker::Int64   # Number of evals done by each worker at each round (multiple of nworkers())
  batchSize::Int64         # batchSize = nworkers()*batchSizeWorker
  maxEvals::Int64          # maxEvals = nbBatches*batchSize
  testTrainRatio::Float64  # size of the test sample relative to the train one
  gridType::Symbol         # Type of grid used when sampling from the parameter space
	penaltyValue::Float64		 # Output value when the model fails
	nonConvergenceFlag::Int64	# Flag to indicate non convergence of the model
	convergenceFlag::Int64		# Flag to indicate convergence of the model
end

function SModelsOptions( ; sModelType::Symbol = :MLPRegressor,
													classifierType::Symbol = :MLPClassifier,
                          desiredMaxPerErrorRegressor::Float64 = 0.5,
													desiredAccuracyClassifier::Float64 = 0.9,
													desiredMinObs::Int64 = 100, # The desired percentage error after which training is stopped
                          nWorkers::Int64 = nworkers(), # To store the number of workers
                          nbBatches::Int64 = 10,        # Total number of batches before stopping
                          batchSizeWorker::Int64 = 10,  # Number of evals done by each worker at each round (multiple of nworkers())
                          batchSize::Int64 = nworkers()*batchSizeWorker,         # batchSize = nworkers()*batchSizeWorker
                          maxEvals::Int64 = nbBatches*batchSize,        # maxEvals = nbBatches*batchSize
                          testTrainRatio::Float64 = 0.25,  # size of the test sample relative to the train one
                          gridType::Symbol = :uniform,
													penaltyValue::Float64 = 999.0,
													nonConvergenceFlag::Int64 = 0,	# Flag to indicate non convergence of the model
													convergenceFlag::Int64 = 1)		# Flag to indicate convergence of the model)

    SModelsOptions(sModelType,
									classifierType,
                  desiredMaxPerErrorRegressor,
									desiredAccuracyClassifier,
									desiredMinObs,
                  nWorkers,
                  nbBatches,
                  batchSizeWorker,
                  batchSize,
                  maxEvals,
                  testTrainRatio,
                  gridType,
									penaltyValue,
									nonConvergenceFlag,
									convergenceFlag)
end
"""
	SModelsProblem
SModelsProblem is a mutable struct that caries all the information needed to
build a surrogate model.
"""
mutable struct SModelsProblem
  modelFunction::Function      #function f:x -> y that we are trying to approximate
  lowerBound::Array{Float64,1} #lower bound for the parameter space
  upperBound::Array{Float64,1} #upper bound for the parameter space
  dimX::Int64                  #dimension of the input parameter
  dimY::Int64                  #dimension of the output vector
  options::SModelsOptions      #options
  trainingSuccessful::Bool     #desiredPerError reached?
  scaler::PyObject             #to rescale the data for the regressor
	scalerRobust::PyObject       #to rescale the data for the regressor
end

function SModelsProblem( ;modelFunction::Function = default_function,    #function f:x -> y that we are trying to approximate
                          lowerBound::Array{Float64,1} = zeros(1), #lower bound for the parameter space
                          upperBound::Array{Float64,1} = ones(1),                   #upper bound for the parameter space
                          dimX::Int64 = 1,                #dimension of the input parameter
                          dimY::Int64 = 1,                 #dimension of the output vector
                          options::SModelsOptions = SModelsOptions(),     #options
                          trainingSuccessful::Bool = false,     #desiredPerError reached?
                          scaler::PyObject = StandardScaler(),
													scalerRobust::PyObject = StandardScaler())

  SModelsProblem(modelFunction,      #function f:x -> y that we are trying to approximate
                lowerBound, #lower bound for the parameter space
                upperBound, #upper bound for the parameter space
                dimX,                #dimension of the input parameter
                dimY,                 #dimension of the output vector
                options,      #options
                trainingSuccessful,     #desiredPerError reached?
                scaler,
								scalerRobust)
end

"""
	default_function(x)
Function x->x. Used to initialize functions.
"""
function default_function(x)
	println("default_function, returns input")
	return x
end


text:
module TestView

using Test
using LazyGroupBy

@testset begin
    @test view.(grouped(isodd, [0, 7, 3, 1, 5, 9, 4, 3, 0, 5])) ==
          Dict(false => [0, 4, 0], true => [7, 3, 1, 5, 9, 3, 5])
end

@testset "mutation" begin
    xs = [0, 7, 3, 1, 5, 9, 4, 3, 0, 5]
    gs = view.(grouped(isodd, xs))
    @test gs == Dict(false => [0, 4, 0], true => [7, 3, 1, 5, 9, 3, 5])
    gs[false][1] = 10
    @test xs == [10, 7, 3, 1, 5, 9, 4, 3, 0, 5]
    gs[false][2:end] .= 20
    @test xs == [10, 7, 3, 1, 5, 9, 20, 3, 20, 5]
    gs[true] .= 30
    @test xs == [10, 30, 30, 30, 30, 30, 20, 30, 20, 30]
end

end  # module


text:
include("/Users/pstey/projects_code/ARules/src/frequent_itemset_tree6.jl")

n = 50
m = 12                # NOTE: most impactful for runtime complexity (num. items in transactions)
mx_depth = 10        # max depth of itemset tree (max size of transactions explored)
t_arr = [sample(1:100, m, replace = false) for _ in 1:n];
t_arr = map(t -> sort!(t), t_arr)

itree2 = frequent_item_tree(t_arr, 1, mx_depth)

n = 100
m = 12                # NOTE: most impactful for runtime complexity (num. items in transactions)
mx_depth = 10        # max depth of itemset tree (max size of transactions explored)
t_arr = [sample(1:100, m, replace = false) for _ in 1:n];
t_arr = map(t -> sort!(t), t_arr)

Profile.clear_malloc_data()
itree2 = frequent_item_tree(t_arr, 1, mx_depth);


text:
using Polynomials
using LinearOperators
using PolynomialOperators

type PadeExponentiator{T<:Number} <: Exponentiator
    p::Poly
    q::Poly
    tmp::Matrix{T}
end
function PadeExponentiator{T<:Number}(m::Integer, n::Integer,
                                      v::StridedArray{T})
    pade = Pade(Poly(1.//convert(Vector{Int},gamma(1:17)),"x"), m, n)
    PadeExponentiator(pade.p, pade.q, Matrix{eltype(v)}(length(v), 2))
end

function (PE::PadeExponentiator)(Ω::AbstractLinearOperator, τ::Number,
                                 v::StridedArray, w::StridedArray)
    PE.tmp[:,2] = polyval(PE.p, τ*Ω, view(PE.tmp, :, 1))*v
    solve!(w, polyval(PE.q, τ*Ω, view(PE.tmp, :, 1)), PE.tmp[:,2])[1]
end

export PadeExponentiator


text:
type ConvergenceSimulation{SolType<:DESolution}
  solutions::Array{SolType}
  errors
  N
  auxdata
  𝒪est
  convergence_axis
end

function ConvergenceSimulation(solutions,convergence_axis;auxdata=nothing,additional_errors=nothing)
  N = size(solutions,1)
  uEltype = eltype(solutions[1].u[1])
  errors = Dict() #Should add type information
  for k in keys(solutions[1].errors)
    errors[k] = reshape(uEltype[sol.errors[k] for sol in solutions],size(solutions)...)
  end
  if additional_errors != nothing
    for k in keys(additional_errors)
      errors[k] = additional_errors[k]
    end
  end
  𝒪est = Dict(map(calc𝒪estimates,errors))
  𝒪esttmp = Dict() #Makes Dict of Any to be more compatible
  for (k,v) in 𝒪est
    if length(v)==1 push!(𝒪esttmp,Pair(k,v[1]))
    else push!(𝒪esttmp,Pair(k,v))
    end
  end
  𝒪est = 𝒪esttmp
  return(ConvergenceSimulation(solutions,errors,N,auxdata,𝒪est,convergence_axis))
end

function test_convergence(dts::AbstractArray,prob::AbstractSDEProblem,alg;numMonte=10000,save_timeseries=true,timeseries_steps=1,timeseries_errors=save_timeseries,adaptive=false,kwargs...)
  N = length(dts)
  is = repmat(1:N,1,numMonte)'
  solutions = pmap((i)->solve(prob,alg;dt=dts[i],save_timeseries=save_timeseries,timeseries_steps=timeseries_steps,adaptive=adaptive,timeseries_errors=timeseries_errors,kwargs...),is)
  if typeof(prob) <: SDEProblem
    solutions = convert(Array{SDESolution},solutions)
  elseif typeof(prob) <: SDETestProblem
    solutions = convert(Array{SDETestSolution},solutions)
  end
  solutions = reshape(solutions,numMonte,N)
  auxdata = Dict("dts" =>  dts)
  # Now Calculate Weak Errors
  additional_errors = Dict()
  # Final
  m_final = mean([s[end] for s in solutions],1)
  m_final_analytic = mean([s.u_analytic[end] for s in solutions],1)
  additional_errors[:weak_final] = mean.(abs.(m_final - m_final_analytic))
  if timeseries_errors
    l2_tmp = Vector{eltype(solutions[1][1])}(size(solutions,2))
    max_tmp = Vector{eltype(solutions[1][1])}(size(solutions,2))
    for i in 1:size(solutions,2)
      solcol = @view solutions[:,i]
      m_errors = [mean([solcol[j][i] for j in 1:length(solcol)]) for i in 1:length(solcol[1])]
      m_errors_analytic = [mean([solcol[j].u_analytic[i] for j in 1:length(solcol)]) for i in 1:length(solcol[1])]
      ts_weak_errors = [abs.(m_errors[i] - m_errors_analytic[i]) for i in 1:length(m_errors)]
      ts_l2_errors = [sqrt.(sumabs2(err)/length(err)) for err in ts_weak_errors]
      l2_tmp[i] = sqrt(sumabs2(ts_l2_errors)/length(ts_l2_errors))
      max_tmp[i] = maximum([maximum(err) for err in ts_weak_errors])
    end
    additional_errors[:weak_l2] = l2_tmp
    additional_errors[:weak_l∞] = max_tmp
  end
  ConvergenceSimulation(solutions,dts,auxdata=auxdata,additional_errors=additional_errors)
end

function test_convergence(dts::AbstractArray,prob::AbstractODEProblem,alg;save_timeseries=true,adaptive=false,kwargs...)
  N = length(dts)
  solutions = [solve(prob,alg;dt=dts[i],save_timeseries=save_timeseries,adaptive=adaptive,kwargs...) for i=1:N]
  auxdata = Dict(:dts =>  dts)
  ConvergenceSimulation(solutions,dts,auxdata=auxdata)
end

function test_convergence(dts::AbstractArray,dxs::AbstractArray,prob::AbstractHeatProblem,convergence_axis;T=1,alg=:Euler)
  if length(dts)!=length(dxs) error("Lengths of dts!=dxs. Invalid convergence simulation") end
  solutions = [solve(parabolic_squaremesh([0 1 0 1],dxs[i],dts[i],T,:dirichlet),prob,alg=alg) for i in eachindex(dts)]
  auxdata = Dict(
            :dts => [sol.fem_mesh.dt for sol in solutions],
            :dxs => [sol.fem_mesh.dx for sol in solutions],
            :Δμs => [sol.fem_mesh.μ  for sol in solutions],
            :Δνs => [sol.fem_mesh.ν  for sol in solutions])
  return(ConvergenceSimulation(solutions,convergence_axis,auxdata=auxdata))
end

function test_convergence(dxs::AbstractArray,prob::AbstractPoissonProblem)
  solutions = [solve(notime_squaremesh([0 1 0 1],dxs[i],:dirichlet),prob,solver=:Direct) for i in eachindex(dxs)]
  auxdata = Dict("dxs" => [sol.fem_mesh.dx for sol in solutions])
  return(ConvergenceSimulation(solutions,dxs,auxdata=auxdata))
end

function calc𝒪estimates(error::Pair)
  key = error.first
  error =error.second
  if ndims(error)>1 error=mean(error,1) end
  S = Vector{eltype(error)}(length(error)-1)
  for i=1:length(error)-1
    S[i] = log2(error[i+1]/error[i])
  end
  return(Pair(key,abs.(mean(S,1))))
end

"""
length(simres::ConvergenceSimulation)

Returns the number of simultations in the Convergence Simulation
"""
Base.length(sim::ConvergenceSimulation) = sim.N
Base.endof( sim::ConvergenceSimulation) = length(sim)
Base.getindex(sim::ConvergenceSimulation,i::Int) = sim.solutions[i]
Base.getindex(sim::ConvergenceSimulation,i::Int,I::Int...) = sim.solutions[i][I]

#=
function print(io::IO, sim::ConvergenceSimulation)
  println(io,"$(typeof(sim)) of length $(length(sim)).")
  print(io,"Convergence Estimates:")
  for (k,v) in sim.𝒪est
    print(" ($k,$v)")
  end
  println(io,"\n-----------Errors-----------")
  for (k,v) in sim.errors
    println(io,"$k: $v")
  end
end

function show(io::IO,sim::ConvergenceSimulation)
  println(io,"$(typeof(sim)) of length $(length(sim)).")
  print(io,"Convergence Estimates:")
  for (k,v) in sim.𝒪est
    print(io," ($k,$v)")
  end
end
=#

max_stars_repo_name: davidmuaddib/DiffEqDevTools

max_issues_repo_name: davidmuaddib/DiffEqDevTools

max_forks_repo_name: davidmuaddib/DiffEqDevTools


text:
"""

This module borrows heavily from pyamg. See [https://github.com/pyamg/pyamg/blob/master/pyamg/amg_core/smoothed_aggregation.h](https://github.com/pyamg/pyamg/blob/master/pyamg/amg_core/smoothed_aggregation.h).
"""
module Aggregation

"""
    standard_aggregation(A :: SparseMatrixCSC)

Returns a vector `v` where each entry in the vector gives the aggregate that
that vertex belongs to. Ex: v[1] = 2 means that vertex 1 belongs to aggregate
2.
"""
function standard_aggregation(A :: SparseMatrixCSC)
    m, n = size(A)
    rows = rowvals(A)
    vals = nonzeros(A)

    # hold aggregates. -1 is unaggregated
    aggregates = fill(-1, n)

    # find initial aggregates
    current_aggregate = 1
    for i in 1:n
        if aggregates[i] == -1 # not yet aggregated
            neighbors = rows[nzrange(A, i)]
            if all(x -> aggregates[x] == -1, neighbors) # no neighbors are aggregated
                aggregates[i] = current_aggregate
                aggregates[neighbors] = current_aggregate
                current_aggregate += 1
            end
        end
    end

    # Pull neighboring nodes into aggregates
    for i in 1:n
        if aggregates[i] == -1
            # find max weight edge
            max_weight = 0
            ind = -1
            for j in nzrange(A, i)
                if rows[j] != i && aggregates[rows[j]] != -1
                    w = vals[j]
                    if w > max_weight
                        max_weight = w
                        ind = rows[j]
                    end
                end
            end

            if ind == -1
                error("isolated node")
            else
                aggregates[i] = aggregates[ind]
            end
        end
    end

    aggregates
end

end


text:
"""
n = DMatrix!(n)
n = DMatrix!(n, (:mode=>:symbolic))
--------------------------------------------------------------------------------------\n
Author: Huckleberry Febbo, Graduate Student, University of Michigan
Date Create: 1/15/2017, Last Modified: 12/06/2019 \n
--------------------------------------------------------------------------------------\n
"""
function DMatrix!(n::NLOpt, kwargs...)         #TODO make IMatrix and option
  kw = Dict(kwargs);

  if !haskey(kw,:mode); mode=:defaut;
  else; mode = get(kw,:mode,0);
  end

  if mode==:defaut
    check_ts = zeros(Float64,n.ocp.Ni);     # check input
    for int in 1:n.ocp.Ni; check_ts[int]=maximum(n.ocp.ts[int]); end
    if maximum(check_ts) < 10*eps()
      error("\n ts is full of zeros: make sure that you call create_intervals() first to calculate ts! \n
                NOTE: This may have occured because  (:finalTimeDV => true) and the final time dv is not working properly!! \n")
    end
    D = [zeros((n.ocp.Nck[int]+1),(n.ocp.Nck[int]+1)) for int in 1:n.ocp.Ni];
    for int in 1:n.ocp.Ni
        D[int] = polyDiff(n.ocp.ts[int]) # +1 is already appended onto ts
    end
    n.ocp.DMatrix = [zeros((n.ocp.Nck[int]),(n.ocp.Nck[int]+1)) for int in 1:n.ocp.Ni];
    DM = [zeros((n.ocp.Nck[int]),(n.ocp.Nck[int])) for int in 1:n.ocp.Ni];
    n.ocp.IMatrix = [zeros((n.ocp.Nck[int]),(n.ocp.Nck[int]+1)) for int in 1:n.ocp.Ni];
    for int in 1:n.ocp.Ni
        n.ocp.DMatrix[int] = D[int][1:end-1,:];    # [Nck]X[Nck+1]
        if n.s.ocp.integrationScheme==:lgrImplicit
          DM[int] = n.ocp.DMatrix[int][:,2:end];   # [Nck]X[Nck]
          n.ocp.IMatrix[int] = inv(DM[int]);       # I = inv(D[:,2:N_k+1])
        end
    end
  elseif mode==:symbolic # for validation only, too slow otherwise
    error(" \n cannot precompile with SymPy \n
             so this fucntion was turned off for typical use!! \n
               -> do a (using SymPy) in NLOptControl.jl then remove this error message and rerun \n")
    Dsym = [Array{Any}(undef, n.ocp.Nck[int],n.ocp.Nck[int]+1) for int in 1:n.ocp.Ni];
    n.ocp.DMatrix = [Array{Any}(undef, n.ocp.Nck[int],n.ocp.Nck[int]+1) for int in 1:n.ocp.Ni]
    test = [Array{Any}(undef, n.ocp.Nck[int]+1) for int in 1:n.ocp.Ni];
    val = 1; # since this is always = 1 this funtion is useful for testing, without scaling the problem from [-1,1] this was useful becuase tf was a design variable
    tf = Sym("tf")
    createIntervals!(n,tf); # gives symbolic expression
    for int in 1:n.ocp.Ni
      for i in 1:n.ocp.Nck[int]+1
        test[int][i] =  n.ocp.ts[int][i](tf=>val)
      end
    end
    for int in 1:n.ocp.Ni
        for idx in 1:n.ocp.Nck[int]+1
            for j in 1:n.ocp.Nck[int]
                f = lagrange_basis_poly(tf, test[int], idx)
                Dsym[int][j,idx] = diff(f,tf) # symbolic differentiation --> slow but useful TODO include this in test functions
                n.ocp.DMatrix[int][j,idx] = Dsym[int][j,idx](tf=>test[int][j])
            end
        end
    end
  end
  nothing
end

"""
scale_w(w,ta,tb)
--------------------------------------------------------------------------------------\n
Author: Huckleberry Febbo, Graduate Student, University of Michigan
Date Create: 12/23/2017, Last Modified: 9/18/2017 \n
--------------------------------------------------------------------------------------\n
"""
function scale_w(w,ta,tb)
  (tb - ta)/2*w
end

"""
n = create_intervals(n)
--------------------------------------------------------------------------------------\n
Author: Huckleberry Febbo, Graduate Student, University of Michigan
Date Create: 12/23/2017, Last Modified: 12/06/2019 \n
--------------------------------------------------------------------------------------\n
"""
function createIntervals!(n::NLOpt)
    tm = range(-1,1;length=n.ocp.Ni+1)     # create mesh points
    di = 2/n.ocp.Ni                           # interval size
    # go through each mesh interval creating time intervals; map [tm[i-1],tm[i]] --> [-1,1]
    n.ocp.ts = [[scale_tau(n.ocp.tau[int],tm[int],tm[int+1]);di*int-1] for int in 1:n.ocp.Ni]
    n.ocp.ws = [scale_w(n.ocp.w[int],tm[int],tm[int+1]) for int in 1:n.ocp.Ni]
    nothing
end

# NOTE this function was used for testing, but is currently depreciated. When it is used again figure out why and explain why
# di = (tf + 1).n.ocp.Ni
function createIntervals!(n::NLOpt, tf)
    tm = range(-1,1; length=n.ocp.Ni+1)       # create mesh points
    di = (tf + 1)/n.ocp.Ni                      # interval size
    # go through each mesh interval creating time intervals; map [tm[i-1],tm[i]] --> [-1,1]
    n.ocp.ts = [[scale_tau(n.ocp.tau[int],tm[int],tm[int+1]);di*int-1] for int in 1:n.ocp.Ni]
    n.ocp.ws = [scale_w(n.ocp.w[int],tm[int],tm[int+1]) for int in 1:n.ocp.Ni]
    nothing
end

"""
L = lagrange_basis_poly!(x,x_data,L)
--------------------------------------------------------------------------------------\n
Author: Huckleberry Febbo, Graduate Student, University of Michigan
Date Create: 12/26/2016, Last Modified: 12/02/2019
Last Modified By: Alexander Buck
Citations: This function was influenced by the lagrange() function [located here](https://github.com/pjabardo/Jacobi.jl/blob/master/src/gauss_quad.jl)
--------------------------------------------------------------------------------------\n
# Input Arguments
* `x`: point to approximate function value at
* `x_data`: x data to used calculate basis polynomials
* `L`: array of interest

# Output Arguments
* `L`: Lagrange basis polynomials

A basic description of Lagrange interpolating polynomials is provided [here](http://127.0.0.1:8000/lagrange_poly.html#lagrange-poly)

# TODO get ride of Nc and replace it with Nck[int]
"""
function lagrange_basis_poly!(x::AbstractArray{T},x_data,L::AbstractArray{T}) where {T<:Number}
    Nc = length(x_data) - 1
    ns = length(x);
    L = zeros(Float64,Nc+1,ns);
    for idx in 1:Nc+1
      for j in 1:ns
        L[idx,j] = lagrange_basis_poly(x[j],x_data,Nc,idx);
      end
    end
    return L
end
lagrange_basis_poly(x::AbstractArray{T},x_data,Nc) where {T<:Number} = lagrange_basis_poly!(x::AbstractArray{T},x_data,Nc,zeros(x))

"""
D = polyDiff(x);
--------------------------------------------------------------------------\n
Last modifed for julia on 12/06/2019 by Mattias Fält\n
Original Author: JJ.A.C. Weideman, S.C. Reddy 1998\n
Original Function Name: poldif.m  |  Source: [matlabcentral](https://www.mathworks.com/matlabcentral/fileexchange/29-dmsuite)\n
https://pdfs.semanticscholar.org/bae2/1eb9458f194887bc8d7808383f56d7f4dca0.pdf
https://math.stackexchange.com/questions/1105160/evaluate-derivative-of-lagrange-polynomials-at-construction-points
--------------------------------------------------------------------------\n
# Input Arguments
* `x::Vector`: Vector of N distinct nodes.

# Output Arguments
* `D::Array{Float64,2}`: Differention Matrix

"""
function polyDiff(x)  #TODO get ride of B stuff
  N = length(x);
  x = x[:];                     # Make sure x is a column vector
  M = 1;                        # assume
  alpha = ones(N,1);
  B = zeros(M,N);
  L= Matrix{Bool}(LinearAlgebra.I, N, N);
  XX = repeat(x,1,N);
  DX = XX-XX';                 # DX contains entries x(k)-x(j).
  DX[L] = ones(N,1);           # Put 1's one the main diagonal.
  c = alpha.*prod(DX,dims=2);       # Quantities c(j).
  C = repeat(c,1,N);
  C = C./C';                   # Matrix with entries c(k)/c(j).
  Z = 1 ./ DX;                   # Z contains entries 1/(x(k)-x(j))
  Z[L] = zeros(N,1);           # with zeros on the diagonal.
  X = Z';                      # X is same as Z', but with
  # diagonal entries removed.
  flag = trues(size(X));
  flag[L] .= false;
  X = X[flag];
  X = reshape(X,N-1,N);
  Y = ones(N-1,N);             # Initialize Y and D matrices.
  D = Matrix{Float64}(LinearAlgebra.I,N,N);                  # Y is matrix of cumulative sums,
  temp = reshape(B[1,:],1,N)
  Y   = cumsum([temp; Y[1:N-1,:].*X], dims=1);     # Diagonals
  D   = Z.*(C.*repeat(LinearAlgebra.diag(D),1,N) - D);   # Off-diagonals
  D[L]   = Y[N,:];                         # Correct the diagonal
  return D
end


text:
using Test
using QuantumOptics

@testset "sortedindices" begin

s = QuantumOptics.sortedindices

@test s.complement(6, [1, 4]) == [2, 3, 5, 6]

@test s.remove([1, 4, 5], [2, 4, 7]) == [1, 5]
@test s.remove([1, 4, 5, 7], [2, 4, 7]) == [1, 5]
@test s.remove([1, 4, 5, 8], [2, 4, 7]) == [1, 5, 8]

@test s.shiftremove([1, 4, 5], [2, 4, 7]) == [1, 3]
@test s.shiftremove([1, 4, 5, 7], [2, 4, 7]) == [1, 3]
@test s.shiftremove([1, 4, 5, 8], [2, 4, 7]) == [1, 3, 5]

@test s.reducedindices([3, 5], [2, 3, 5, 6]) == [2, 3]
x = [3, 5]
s.reducedindices!(x, [2, 3, 5, 6])
@test x == [2, 3]

@test_throws AssertionError s.check_indices(5, [1, 6])
@test_throws AssertionError s.check_indices(5, [0, 2])
@test s.check_indices(5, Int[]) == nothing
@test s.check_indices(5, [1, 3]) == nothing
@test s.check_indices(5, [3, 1]) == nothing

@test_throws AssertionError s.check_sortedindices(5, [1, 6])
@test_throws AssertionError s.check_sortedindices(5, [3, 1])
@test_throws AssertionError s.check_sortedindices(5, [0, 2])
@test s.check_sortedindices(5, Int[]) == nothing
@test s.check_sortedindices(5, [1, 3]) == nothing

@test s.check_embed_indices([1,[3,5],10,2,[70,11]]) == true
@test s.check_embed_indices([1,3,1]) == false
@test s.check_embed_indices([1,[10,11],7,[3,1]]) == false
@test s.check_embed_indices([[10,3],5,6,[3,7]]) == false
@test s.check_embed_indices([]) == true
end # testset


text:
module RungeKutta

using DataStructures, FastAnonymous, Docile

export make_rkfe, show_rkfe,
       make_rk3, show_rk3,
       make_rk4a, show_rk4a,
       make_rk4b, show_rk4b

@doc* """
Runge-Kutta integration of ODEs
Based on Haskell code by Uwe Hollerbach.

dy
-- = f(t, y)
dt

y_{n+1} = y_n + h sum_{i=1}^s b_i k_i
k_i = f(t_n + c_i h, y_n + h sum_{j=1}^s a_{ij} k_j)

"Butcher Tableau" is

 c_1  a_11 a_12 ... a_1s
 c_2  a_21 a_22 ... a_2s
 ...  ...
 c_s  a_s1 a_s2 ... a_ss
      b_1  b_2  ... b_s

This module implements a method that can do a generic tableau, then
specializes with different tableaux to let the user pick a specific
method. Adaptive step-size methods, see below, add a row of d_j
coefficients and use that to report the error:

 e_{n+1} = h sum_{i=1}^s d_i k_i

 non-adaptive solvers:
	rkfe, rk3, rk4a, rk4b

 adaptive solvers:
	rkhe, rkbs, rkf45, rkck, rkdp, rkf78, rkv65

 auxiliary non-adaptive solvers (error estimators from the adaptive ones):
	rkhe_aux, rkbs_aux, rkf45_aux, rkck_aux, rkdp_aux, rkf78_aux, rkv65_aux

use rk4[ab] if you don't need an adaptive solver, rkdp or rkv65 if you do;
or use what you need if you're an expert.

DO NOT USE rkfe EXCEPT FOR DEMONSTRATIONS OF INSTABILITY!
(Or if you're an expert.)

Reference: E. Hairer, S. P. Norsett, G. Wanner,
Solving Ordinary Differential Equations I: Nonstiff Problems
(second revised edition, 1993).
"""

function rat (n :: Int)
    Rational (n, 1)
end

function isempty (lst:: LinkedList{Rational{Int}})
    return isa (lst, Nil)
end

function isempty (lst:: LinkedList{Int})
    return isa (lst, Nil)
end

function isempty (lst:: LinkedList{Float64})
    return isa (lst, Nil)
end

function isempty (lst:: LinkedList{Any})
    return isa (lst, Nil)
end

function mapPartial(f::Base.Callable, l::Cons)
    first = f(l.head)
    l2 = (first == nothing) ? nil() : cons(first, nil(typeof(first)))
    for h in l.tail
        fh = f(h)
        if (!(fh == nothing))
            l2 = cons(fh, l2)
        end
    end
    reverse(l2)
end

function foldl (f,init,v)
    ax = init
    for x in v
        ax = f(x,ax)
    end
    return ax
end

function zip (l1 :: LinkedList, l2 :: LinkedList)
    cons ((head(l1),head(l2)),zip(tail(l1),tail(l2)))
end

function zip (l1 :: Nil, l2 :: Nil)
    nil()
end

function foldl1 (f, v)
    if isa (v, Cons)
        let h = head (v), t = tail (v)
            if isempty (t)
                h
            else
                foldl (f, f(h, head(t)), tail(t))
            end
        end
    else
        throw (DomainError ())
    end
end

function normalize (r)
    n = num(r)
    d = den(r)
    if (n == 0)
        return 0
    elseif (d == 1)
        return r
    else
        c = gcd (n,d)
        if (c == d)
            return div(n,c)
        else
            return (div (n,c)) // (div(d,c))
        end
    end
end

# Stores a vector of rational numbers as a common denominator, then a vector
# of numerators, all stored as doubles: this lets us write the values in
# the source as exact rational numbers, yet we can take advantage of not
# having to do too many conversions and also of not having to multiply
# by 0 or 1 in some cases.

type RCL
    den  :: Float64
    nums ::  Union (LinkedList{Float64}, Nil{Any})
end

function ratToRCL (rs :: Nil{Any})
    return RCL(1.0, nil())
end

function ratToRCL (rs:: Cons{Rational{Int}})
    ds = map (den, rs)
    dp = foldl1 (*, ds)
    ns = map (x -> num (dp * x), rs)
    g  = foldl (gcd, dp, ns)
    return RCL(float64(dp / g), map (x -> float64 (x / g),  ns))
end

function ratToRCLs (x)
    map (ratToRCL, x)
end

function ratToReals (x:: LinkedList{Rational{Int}})
    map (float64, x)
end

##   A helper function to take the difference of two arrays of
##   rationals: we don't want to use (-) or (.-) because that requires
##   the array to be of the same size.  We want implicit zeros at the
##   end, as far as is necessary.

## fun diffs ([], []) = []
##  | diffs (xs, []) = xs
##  | diffs ([], xs) = map negate xs
##  | diffs (x::xs, y::ys) = (x +/ (negate y)) :: (diffs (xs,ys))

# Helper function to sum a list of K_i, skipping unnecessary
# multiplications

function m_scale (sc_fn)
    (sv) ->
    begin
        (sv[1] == 0.0) ? Nothing() : sc_fn (sv[1],sv[2])
    end
end

function m_sum (sum_fn)
    (x,ax) -> (x == nothing) ? ax : sum_fn (x,ax)
end

# Helper function to generate a list of K_i

function k_sum (sc_fn, sum_fn, h)
    f = ((rcl, ks) ->
         begin
             d     = rcl.den
             ns    = rcl.nums
             ns_ks = zip (ns, ks)
             return sc_fn (h / d, foldl1 (m_sum (sum_fn), mapPartial (m_scale (sc_fn), ns_ks)))
         end)
    return f
end

function gen_ks (ksum_fn,sum_fn,der_fn,h,old,cs,ar)
    (tn,yn) = old
    ks1 = nil()
    yn1 = yn
    for (a,c) in zip (ar,cs)
        yn1 = isempty (ks1) ? yn : sum_fn (yn, ksum_fn (a,ks1))
        v   = der_fn (tn + c*h, yn1)
        ks1 = cat(ks1, list(v))
    end
    return ks1
end

@doc* """
This is the first core routine.

Its arguments are:

   c table (specified internally)
   a table (specified internally)
   b table (specified internally)

user-specified arguments:

   scale function to scale a Y state vector ::
      (real * a -> a)

   sum function to add two Y state vectors ::
      (a * a -> a)

   derivative function F ::
      (real * a -> a)

   step size H ::
      real

   current state (T,Y) ::
      (real, a)

   and the return value is the new state Y_new
"""

function core1 (cl, al, bl)
    f = @anon ((sc_fn, sum_fn, der_fn) ->
               ((h) ->
                ((tn,yn) ->
                 begin
                     ksum = k_sum (sc_fn,sum_fn,h)
                     ks = gen_ks (ksum, sum_fn, der_fn, h, (tn,yn), cl, al)
                     return sum_fn (yn, ksum (bl, ks))
                 end)))
    return f
end

function core2 (cl, al, bl, dk)
    f = @anon ((sc_fn, sum_fn, der_fn) ->
               ((h) ->
                begin
                    ksum = k_sum (sc_fn,sum_fn,h)
                    ((old) ->
                     begin
                         ks   = gen_ks (ksum, sum_fn, der_fn, h, old, cl, al)
                         return (sum_fn (yn, ksum (bl, ks)), ksum (dl, ks))
                     end)
                end))

    return f
end

function list_show{T}(l::LinkedList{T},show,sep,startstr,endstr)
    if isa(l,Nil)
        if is(T,Any)
            "nil()"
        else
            string ("nil(", T, ")")
        end
    else
        string (startstr,
                head(l),
                sep,
                list_show(tail (l),show,sep,"",""),
                endstr)
    end
end

function def_list_show (show,xs)
    list_show (xs,show,",","[","]")
end

function rcl_show (rcl)
    d = rcl.den
    ns = rcl.nums
    string ("<", d, ", ", def_list_show (string,ns), ">")
end

function rk_show1 (title,cs,ar,bs)
    string (title, ":\ncs:\t", def_list_show (string,cs),
            "\nas:\t", list_show (ar,rcl_show,"\n\t","",""),
            "\nbs:\t", rcl_show (bs))
end

function rk_show2 (title,cs,ar,bs,ds)
    string (title, ":\nds:\t", rcl_show (ds),
            "\ncs:\t", def_list_show (string,cs),
            "\nbs:\t", rcl_show (bs),
            "\nas:\t", list_show (ar,rcl_show,"\n\t","",""))
end

##   Some specific explicit methods, taken from
##   "List of Runge-Kutta methods" at Wikipedia

## forward Euler: unconditionally unstable: don't use this!

cs_fe = ratToReals (list (rat (0)))
as_fe = ratToRCLs (list (list ()))
bs_fe = ratToRCL  (list (rat (1)))

make_rkfe = () -> core1 (cs_fe, as_fe, bs_fe)
show_rkfe = rk_show1 ("Forward Euler", cs_fe, as_fe, bs_fe)

## Kutta's third-order method:

cs_rk3 = ratToReals (list (rat (0), 1//2, rat (1)))
as_rk3 = ratToRCLs (list (list (), list (1//2), list (rat (-1), rat (2))))
bs_rk3 = ratToRCL  (list (1//6, 2//3, 1//6))
make_rk3 = () -> core1 (cs_rk3, as_rk3, bs_rk3)
show_rk3 = rk_show1 ("Kutta's third-order method", cs_rk3, as_rk3, bs_rk3)

## Classic fourth-order method

cs_rk4a = ratToReals (list (rat (0), 1//2, 1//2, rat (1)))
as_rk4a = ratToRCLs (list (list (), list (1//2), list (rat (0), 1//2), list (rat (0), rat (0), rat (1))))
bs_rk4a = ratToRCL  (list (1//6, 1//3, 1//3, 1//6))
make_rk4a = () -> core1 (cs_rk4a, as_rk4a, bs_rk4a)
show_rk4a = rk_show1 ("Classic fourth-order method", cs_rk4a, as_rk4a, bs_rk4a)

## Kutta's other fourth-order method... "The first [above] is more
## popular, the second is more precise." (Hairer, Norsett, Wanner)

cs_rk4b = ratToReals (list (rat (0), 1//3, 2//3, rat (1)))
as_rk4b = ratToRCLs (list (list (), list (1//3), list (-1//3, rat (1)), list (rat (1), rat (-1), rat (1))))
bs_rk4b = ratToRCL  (list (1//8, 3//8, 3//8, 1//8))
make_rk4b = () -> core1 (cs_rk4b, as_rk4b, bs_rk4b)
show_rk4b = rk_show1 ("Kutta's other classic fourth-order method", cs_rk4b, as_rk4b, bs_rk4b)

end # module


text:
function _extract_lp_solution_sub(graph, x)
    solution = Edge{Int}[]
    for e in edges(graph)
        if JuMP.value(x[e]) >= 0.5
            push!(solution, e)
        elseif JuMP.value(x[reverse(e)]) >= 0.5
            push!(solution, reverse(e))
        end
    end
    return solution
end

function _extract_lp_solution(i::SpanningTreeInstance{Int, Maximise}, x)
    return _extract_lp_solution_sub(i.graph, x)
end

function _extract_lp_solution(
    i::MinimumBudget{SpanningTreeInstance{Int, Maximise}, Int},
    x,
)
    return _extract_lp_solution_sub(i.instance.graph, x)
end


text:
using Random, Distributions, Plots

function cum_miss(; nobs, nvar, p)
    b = Bernoulli(p)
    sample = rand(b, nobs*nvar)
    sample = reshape(sample, (nobs, nvar))
    patterns = [sample[i, :] for i in 1:size(sample, 1)]

    #binomial = sum.(patterns)
    #histogram(binomial)

    remember = Vector{BitArray{1}}()
    rows = [Vector{Int64}(undef, 0) for i = 1:size(patterns, 1)]

    for i = 1:size(patterns, 1)
        unknown = true
        for j = 1:size(remember, 1)
            if patterns[i] == remember[j]
                push!(rows[j], i)
                unknown = false
            end
        end
        if unknown
            push!(remember, patterns[i])
            push!(rows[size(remember, 1)], i)
        end
    end

    rows = rows[1:length(remember)]
    n_pattern = size(rows, 1)

    # sort by number of missings
    sort_n_miss = sortperm(sum.(remember))
    remember = remember[sort_n_miss]
    remember_cart = findall.(!, remember)
    rows = rows[sort_n_miss]

    pattern_n_obs = size.(rows, 1)
    pattern_nvar_obs = length.(remember_cart)

    obs_pat = pattern_n_obs.*pattern_nvar_obs
    sort_n_miss = sortperm(obs_pat, rev = true)
    obs_pat = obs_pat[sort_n_miss]
    pattern_n_obs = pattern_n_obs[sort_n_miss]

    #pattern_n_obs_cum = pattern_n_obs_cum[sort_n_miss]
    pattern_nvar_obs = pattern_nvar_obs[sort_n_miss]
    #complexity = complexity[sort_n_miss]

    pattern_n_obs_cum = cumsum(pattern_n_obs)

    complexity = pattern_nvar_obs.^3
    obs_pat_cum = cumsum(obs_pat)

    p1 = plot(1:n_pattern, obs_pat_cum, title = "observed datapoints cum.", ylims = (0, Inf))
    p2 = plot(1:n_pattern, pattern_n_obs, title = "N observed")
    p3 = plot(1:n_pattern, pattern_n_obs_cum, title = "N observed cum.", ylims = (0, Inf))
    p4 = plot(1:n_pattern, pattern_nvar_obs, title = "N variables")
    p5 = plot(1:n_pattern, complexity, title = "N variables ^3")

    plot(p1, p2, p3, p4, layout = (2,2))
end

cum_miss(;nobs= 1000, nvar= 10, p= 0.1)

cum_miss(;nobs= 1000, nvar= 100, p= 0.1)

cum_miss(;nobs= 5000, nvar= 100, p= 0.1)

cum_miss(;nobs= 1000, nvar= 100, p= 0.01)

cum_miss(;nobs= 10000, nvar= 100, p= 0.01)

cum_miss(;nobs= 1000, nvar= 10, p= 0.9)

a = [Base.binomial(100, k) for k = 0:5]

function sim_dropout(nwaves, dist)
    pattern = falses(nwaves)
    sample = rand(dist, nwaves)
    t_dropout = findfirst(sample)
    !isnothing(t_dropout) ?
        pattern[t_dropout:nwaves] .= true :
        nothing
    pattern[1] = false
    return pattern
end

function get_miss(nvar, dist, miss_wave)
    if miss_wave
        pattern = trues(nvar)
    else
        pattern = rand(dist, nvar)
    end
    return pattern
end

function pattern_long(nvar, dist, miss_waves)
    pattern =
        [get_miss(nvar, dist, miss_wave) for miss_wave in miss_waves]
    pattern = vcat(pattern...)
    return(pattern)
end

function cum_miss_dropout(; nobs, nvar, nwaves, p, p_dropout)
    dist1 = Bernoulli(p_dropout)
    dist2 = Bernoulli(p)
    miss_waves = [sim_dropout(nwaves, dist1) for i in 1:nobs]

    patterns = [pattern_long(nvar, dist2, miss_waves[i]) for i in 1:nobs]

    #return patterns
    #binomial = sum.(patterns)
    #histogram(binomial)

    remember = Vector{BitArray{1}}()
    rows = [Vector{Int64}(undef, 0) for i = 1:size(patterns, 1)]

    for i = 1:size(patterns, 1)
        unknown = true
        for j = 1:size(remember, 1)
            if patterns[i] == remember[j]
                push!(rows[j], i)
                unknown = false
            end
        end
        if unknown
            push!(remember, patterns[i])
            push!(rows[size(remember, 1)], i)
        end
    end

    rows = rows[1:length(remember)]
    n_pattern = size(rows, 1)

    # sort by number of missings
    sort_n_miss = sortperm(sum.(remember))
    remember = remember[sort_n_miss]
    remember_cart = findall.(!, remember)
    rows = rows[sort_n_miss]

    pattern_n_obs = size.(rows, 1)
    pattern_nvar_obs = length.(remember_cart)

    obs_pat = pattern_n_obs.*pattern_nvar_obs
    sort_n_miss = sortperm(obs_pat, rev = true)
    obs_pat = obs_pat[sort_n_miss]
    pattern_n_obs = pattern_n_obs[sort_n_miss]

    #pattern_n_obs_cum = pattern_n_obs_cum[sort_n_miss]
    pattern_nvar_obs = pattern_nvar_obs[sort_n_miss]
    #complexity = complexity[sort_n_miss]

    pattern_n_obs_cum = cumsum(pattern_n_obs)

    complexity = pattern_nvar_obs.^3
    obs_pat_cum = cumsum(obs_pat)

    p1 = plot(1:n_pattern, obs_pat_cum, title = "observed datapoints cum.", ylims = (0, Inf))
    p2 = plot(1:n_pattern, pattern_n_obs, title = "N observed")
    p3 = plot(1:n_pattern, pattern_n_obs_cum, title = "N observed cum.", ylims = (0, Inf))
    p4 = plot(1:n_pattern, pattern_nvar_obs, title = "N variables")
    p5 = plot(1:n_pattern, complexity, title = "N variables ^3")

    plot(p1, p2, p3, p4, layout = (2,2))
end

cum_miss_dropout(nobs = 1000, nvar = 10, nwaves = 10, p = 0.05, p_dropout = 0.1)

dist = Binomial(10, 0.1)


text:
using Match

abstract type Grid end
abstract type Grid2 <: Grid end
abstract type Grid3 <: Grid end

const golden_angle = 2.39996322972865332

"""
    Grid(dr::Float64, symm::Symbol)::Grid

Define a grid with a given characterstic length ``\\text{d}r`` and symmetry type
`symm`. Characterstic length means that one particle occupies a volume
``\\text{d}r^n``, where ``n`` is the dimension. Supported symmetry types are

2D:
1) `:square` (square grid)
2) `:hexagonal` (hexagrid, dual of isometric grid)
3) `:vogel` (grid based on Vogel's spiral)

3D:
1) `:cubic` (cubic grid)
"""
function Grid(dr::Float64, symm::Symbol)::Grid
    return @match symm begin
		:square 	=> Squaregrid(dr)
		:hexagonal  => Hexagrid(dr)
        :vogel      => VogelGrid(dr)
        :cubic      => CubicGrid(dr)
		_ 			=> @error("Unsupported grid type: "*string(symm))
    end
end

function dimension(::Grid2)::Int64
    return 2
end

function dimension(::Grid3)::Int64
    return 3
end

mutable struct Squaregrid <: Grid2
    dr::Float64
end

function covering(grid::Squaregrid, s::Shape)::Vector{RealVector}
    xs = RealVector[]
    rect = boundarybox(s)
    i_min = Int64(floor(rect.x1_min/grid.dr))
    j_min = Int64(floor(rect.x2_min/grid.dr))
    i_max = Int64(ceil(rect.x1_max/grid.dr))
    j_max = Int64(ceil(rect.x2_max/grid.dr))
	for i in i_min:i_max, j in j_min:j_max
        x = RealVector(i*grid.dr, j*grid.dr, 0.)
        if is_inside(x, s)
            push!(xs, x)
        end
	end
    return xs
end

mutable struct Hexagrid <: Grid2
    dr::Float64
    a::Float64
    b::Float64
    Hexagrid(dr::Float64) = new(dr, (4/3)^(1/4)*dr, (3/4)^(1/4)*dr)
end

function covering(grid::Hexagrid, s::Shape)::Vector{RealVector}
    xs = RealVector[]
    rect = boundarybox(s)
    i_min = Int64(floor(rect.x1_min/grid.a)) - 1
    j_min = Int64(floor(rect.x2_min/grid.b))
    i_max = Int64(ceil(rect.x1_max/grid.a))
    j_max = Int64(ceil(rect.x2_max/grid.b))
	for i in i_min:i_max, j in j_min:j_max
        x1 = (i + (j%2)/2)*grid.a
        x2 = j*grid.b
        x = RealVector(x1, x2, 0.)
        if is_inside(x, s)
            push!(xs, x)
        end
	end
	return xs
end

mutable struct VogelGrid <: Grid2
    center::RealVector
    k::Float64
    dr::Float64
    VogelGrid(dr::Float64) = begin
        k = dr/sqrt(pi)
        return new(RealVector(0.,0.,0.), k, dr)
    end
end

function covering(vg::VogelGrid, s::Shape)::Vector{RealVector}
    bb = boundarybox(s)
    dl = RealVector(bb.x1_min, bb.x2_min, 0.)
    dr = RealVector(bb.x1_max, bb.x2_min, 0.)
    ur = RealVector(bb.x1_max, bb.x2_max, 0.)
    ul = RealVector(bb.x1_min, bb.x2_max, 0.)
    R = 0.0
    for corner in [dl, dr, ur, ul]
        R = max(R, norm(corner - vg.center))
    end
    N = (R/vg.k)^2
    xs = RealVector[]
    for n in 1:N
        x = vg.center + vg.k*sqrt(n)*RealVector(cos(n*golden_angle), sin(n*golden_angle), 0.)
        if is_inside(x, s)
            push!(xs, x)
        end
    end
    return xs
end

mutable struct CubicGrid <: Grid3
    dr::Float64
end

function covering(grid::CubicGrid, s::Shape)::Vector{RealVector}
    xs = RealVector[]
    box = boundarybox(s)
    i_min = Int64(floor(box.x1_min/grid.dr))
    j_min = Int64(floor(box.x2_min/grid.dr))
    k_min = Int64(floor(box.x3_min/grid.dr))
    i_max = Int64(ceil(box.x1_max/grid.dr))
    j_max = Int64(ceil(box.x2_max/grid.dr))
    k_max = Int64(ceil(box.x3_max/grid.dr))
	for i in i_min:i_max, j in j_min:j_max, k in k_min:k_max
        x = RealVector(i*grid.dr, j*grid.dr, k*grid.dr)
        if is_inside(x, s)
            push!(xs, x)
        end
	end
    return xs
end

"""
    generate_particles!(sys::ParticleSystem,
                        grid::Grid,
                        geometry::Shape,
                        constructor::Function)

Create particles using `constructor(x::RealVector)::AbstractParticle` at every `grid` point
inside a given shape.

"""
function generate_particles!(sys::ParticleSystem, grid::Grid, geometry::Shape, constructor::Function)
    xs = covering(grid, geometry)
    for x in xs
        push!(sys.particles, constructor(x))
    end
end


text:
create(c::Client, ::Type{Ban}, g::AbstractGuild, u::User; kwargs...) = create_guild_ban(c, g.id, u.id; kwargs...)

retrieve(c::Client, ::Type{Ban}, g::AbstractGuild, u::User) = get_guild_ban(c, g.id, u.id)
retrieve(c::Client, ::Type{Ban}, g::AbstractGuild) = get_guild_bans(c, g.id)

delete(c::Client, b::Ban, g::AbstractGuild) = remove_guild_ban(c, g.id, b.user.id)


text:
# ---
# title: 303. Range Sum Query - Immutable
# id: problem303
# author: Qling
# date: 2021-03-01
# difficulty: Easy
# categories: Dynamic Programming
# link: <https://leetcode.com/problems/range-sum-query-immutable/description/>
# hidden: true
# ---
#
# Given an integer array `nums`, find the sum of the elements between indices
# `i` and `j` `(i ≤ j)`, inclusive.
#
# Implement the `NumArray` class:
#
#   * `NumArray(int[] nums)` Initializes the object with the integer array `nums`.
#   * `int sumRange(int i, int j)` Return the sum of the elements of the `nums` array in the range `[i, j]` inclusive (i.e., `sum(nums[i], nums[i + 1], ... , nums[j])`)
#
#
#
# **Example 1:**
#
#
#
#     **Input**
#     ["NumArray", "sumRange", "sumRange", "sumRange"]
#     [[[-2, 0, 3, -5, 2, -1]], [0, 2], [2, 5], [0, 5]]
#     **Output**
#     [null, 1, -1, -3]
#
#     **Explanation**
#     NumArray numArray = new NumArray([-2, 0, 3, -5, 2, -1]);
#     numArray.sumRange(0, 2); // return 1 ((-2) + 0 + 3)
#     numArray.sumRange(2, 5); // return -1 (3 + (-5) + 2 + (-1))
#     numArray.sumRange(0, 5); // return -3 ((-2) + 0 + 3 + (-5) + 2 + (-1))
#
#
#
#
# **Constraints:**
#
#   * `0 <= nums.length <= 104`
#   * `-105 <= nums[i] <= 105`
#   * `0 <= i <= j < nums.length`
#   * At most `104` calls will be made to `sumRange`.
#
#
## @lc code=start
using LeetCode

struct NumArray
    sums::Vector
    function NumArray(nums::Vector)
        n = length(nums)
        sums = nums[:]

        for i = 2:n
            sums[i] = sums[i-1] + nums[i]
        end
        new(sums)
    end
end

function sum_range(nums::NumArray, i::Int, j::Int)
    i == j && return 0
    i == 1 && return nums.sums[j]
    return nums.sums[j] - nums.sums[i-1]
end
## @lc code=end


text:
module MovingFrame
using LinearAlgebra
using Statistics: mean, var, std
using GaussianProcesses
using GeometryBasics
using GSL
using PyCall
using SparseArrays
using DataFrames
using CSV
using FileIO: load, save

function __init__()
    copy!(igl, pyimport_conda("igl", "igl","conda-forge"))
    copy!(scipy, pyimport_conda("scipy", "scipy"))

end

include("./utils_movement.jl")
include("./smoothing.jl")
export Smoother, GaussianProcessSmoother, smoothing, fit, predict
include("./frenet_serret.jl")
export frenet_serret
include("./parallel_transport.jl")
export parallel_transport

include("spharm.jl")
export Plm,
    sph_Plm,
    Plm_array,
    Plm_array_index,
    real_spharm,
    real_spharm_array,
    j2lm,
    lm2j,
    spharm_coefs
include("shape_descriptor.jl")
export spharm_descriptor,
    reconstruct_from_spharm_coefs, reconstruct_mesh, sphere_parameterization
include("utils_shape.jl")
export cart2sph,
    sph2cart,
    mesh_centering,
    volume_normalizing,
    mesh_volume,
    create_mesh,
    centroid,
    internal_angles,
    double_surface_area,
    surface_area,
    mesh_reorientation

include("sphere.jl")
include("python_igl.jl")
export igl_gaussian_curvature, igl_mean_curvature
include("curvature.jl")
export mass_matrix_barycentric,
    mass_matrix_barycentric2,
    cotangent_laplacian,
    gaussian_curvature,
    mean_curvature,
    spherization

include("fileIO.jl")
export load_mesh, save_mesh, create_group_df

end # module


text:
# Day 16: Packet Decoder
function hex2bits(hex) # "BITS" representation of our data
    # Parse as Int8 as there is not Int4 datatype, but only keep last 4 digits.
    return Vector{Char}(
        join(bitstring(i)[5:8] for i in parse.(Int8, split(hex, ""), base=16))
    )
end
load(::Day{16}, path) = hex2bits(readline(path))

function bin2num(binary_string) # independent of string length
    return sum(c == '1' ? 2^(i - 1) : 0 for (i, c) in enumerate(reverse(binary_string)))
end

# Part 1 & 2 use the same parser
function parse_transmission!(bitstream)
    version_sum = 0

    function parse_packet!(bitstream)
        bits = splice!(bitstream, 1:6)
        version = bin2num(bits[1:3])
        type_id = bin2num(bits[4:6])
        version_sum += version

        # Every parsed packet returns its value and length
        if type_id == 4 # packet is a literal value
            # Iterate over value sub-packets
            is_last_group = false
            val_bits = Vector{Char}()
            subpacket_count = 0

            while !is_last_group
                is_last_group = (popfirst!(bitstream) == '0')
                append!(val_bits, splice!(bitstream, 1:4))
                subpacket_count += 1
            end

            packet_length = 6 + subpacket_count * 5
            value = bin2num(val_bits)
        else
            # Packet is an operator acting on subpackets
            length_type_id = popfirst!(bitstream)
            subpackets = Vector{Tuple{Int,Int}}()
            if length_type_id == '0'
                # Parse subpacket by length of bits
                subpacket_length = bin2num(splice!(bitstream, 1:15))
                subpacket_bits = splice!(bitstream, 1:subpacket_length)
                while !isempty(subpacket_bits)
                    push!(subpackets, parse_packet!(subpacket_bits))
                end
                packet_length = 6 + subpacket_length
            else
                # Parse subpackets by number
                subpacket_number = bin2num(splice!(bitstream, 1:11))
                while length(subpackets) < subpacket_number
                    push!(subpackets, parse_packet!(bitstream))
                end
                packet_length = 6 + sum(last.(subpackets))
            end
            # Apply operators to subpackets
            if type_id == 0
                value = sum(first.(subpackets))
            elseif type_id == 1
                value = prod(first.(subpackets))
            elseif type_id == 2
                value = minimum(first.(subpackets))
            elseif type_id == 3
                value = maximum(first.(subpackets))
            elseif type_id == 5
                value = Int(first.(subpackets[1]) > first.(subpackets[2]))
            elseif type_id == 6
                value = Int(first.(subpackets[1]) < first.(subpackets[2]))
            elseif type_id == 7
                value = Int(first.(subpackets[1]) == first.(subpackets[2]))
            end
        end

        return value, packet_length
    end

    value, _ = parse_packet!(bitstream)
    return version_sum, value
end

# Part 1 – return sum of versions
solve1(::Day{16}, data) = first(parse_transmission!(deepcopy(data)))
# Part 2 – return value
solve2(::Day{16}, data) = last(parse_transmission!(deepcopy(data)))

# No test for for both parts, just pass:
testresult1(::Day{16}) = 0
test1(::Day{16}) = 0
testresult2(::Day{16}) = 0
test2(::Day{16}) = 0


text:
@testset "Lines on a Quintic surface in 3-space" begin
    @testset "Lines on a quintic surface in 3-space" begin
        sys, q₀, gamma = fano_quintic()
        @time res =
            solve(sys; gamma = gamma, target_parameters = q₀, start_system = :total_degree)
        @test nsolutions(res) == 2875

        @time poly_res = solve(sys; target_parameters = q₀, start_system = :polyhedral)
        @test nsolutions(poly_res) == 2875
    end

end