text:
######################################################################################################
# Copy from https://github.com/xuedong/LinBAI.jl  by Xuedong Shang                                   #
#  * Fixed Xuedong's incorrect implementation of XYAdaptive using the version in Tanner Fiez's repo: #
#    https://github.com/fiezt/Transductive-Linear-Bandit-Code/blob/master/XY_ADAPTIVE.py             #
######################################################################################################
import Distributions: Multinomial;
using LinearAlgebra;
include("../utilities/regret.jl");
include("../utilities/tracking.jl");
include("../utilities/envelope.jl");
###################################################################### Ours ######################################################################
"""
Ours: Frank-Wolfe based Sampling
"""
struct FWSampling
end
long(sr::FWSampling) = "FW-Sampling";
abbrev(sr::FWSampling) = "FWS";
mutable struct FWSamplingState
    x;
    Vxinv;
    FWSamplingState(K, Vxinv) = new([1.0/K for i=1:K], Vxinv);
end
function start(sr::FWSampling, N, Vxinv)
    FWSamplingState(length(N), Vxinv);
end
function nextsample(sr::FWSamplingState, pep::LinearBestArm, N, S, Vinv)
    K = length(N); t = sum(N); hw = N./t; hμ = Vinv*S; hr = [hμ'pep.arms[k] for k=1:K]; hi = argmax(hr);
    r = t^(-9/10)/K; z = [0.0 for i=1:K];
    if !hμ_in_lambda(hr, hi, K) || is_complete_square(floor(Int, t/K))
        z = [1.0/K for i=1:K];
    else
        f, ∇f, fidx = compute_f_∇f_linear_bai(sr.x, hμ, r, pep.arms, sr.Vxinv)
        if length(fidx) == 1 # best challenger
            challenger_idx = argmax(∇f[fidx[1]]);
            z = [(challenger_idx==j) ? 1 : 0 for j=1:K];
        else # solve LP of the zero-sum matrix game
            Σ = [[(i==j) ? 1 : 0 for j=1:K]-sr.x for i=1:K];
            A = [[Σ[i]'∇f[j] for i=1:K] for j in fidx]; # construct payoff matrix
            z = solveZeroSumGame(A, K, length(fidx));
        end
    end
    # update
    setfield!(sr, :x, sr.x*((t-1.0)/t) + z*1.0/t);
    nextVxinv = sherman_morrison(sr.Vxinv*(t-1.0)/t, z[1]*pep.arms[1]/t);
    for k=2:K
        nextVxinv = sherman_morrison(nextVxinv, z[k]*pep.arms[k]/t);
    end
    setfield!(sr, :Vxinv, nextVxinv);
    return hi, argmax(sr.x ./ hw);
end
function nextsample(sr::FWSamplingState, pep::LinearThreshold, N, S, Vinv)
    K = length(N); t = sum(N); hw = N./t; hμ = Vinv*S; hr = [hμ'pep.arms[k] for k=1:K]; hi = argmax(hr);
    r = t^(-9/10)/K; z = [0.0 for i=1:K];
    if !hμ_in_lambda_threshold(hr, K, pep.threshold) || is_complete_square(t)
        z = [1.0/K for i=1:K];
    else
        f, ∇f, fidx = compute_f_∇f_linear_threshold(pep, sr.x, hμ, r, sr.Vxinv);
        if length(fidx) == 1 # best challenger
            challenger_idx = argmax(∇f[fidx[1]]);
            z = [(challenger_idx==j) ? 1 : 0 for j=1:K];
        else # solve LP of the zero-sum matrix game
            Σ = [[(i==j) ? 1 : 0 for j=1:K]-sr.x for i=1:K];
            A = [[Σ[i]'∇f[j] for i=1:K] for j in fidx]; # construct payoff matrix
            z = solveZeroSumGame(A, K, length(fidx));
        end
    end
    # update
    setfield!(sr, :x, sr.x*((t-1.0)/t) + z*1.0/t);
    nextVxinv = sherman_morrison(sr.Vxinv*(t-1.0)/t, z[1]*pep.arms[1]/t);
    for k=2:K
        nextVxinv = sherman_morrison(nextVxinv, z[k]*pep.arms[k]/t);
    end
    setfield!(sr, :Vxinv, nextVxinv);
    return hi, argmax(sr.x ./ hw);
end

###################################################################### Baselines ######################################################################
"""
Our implementation of Lazy T&S (Jedra and Proutiere, 2020)
with heuristic stopping time (see their appendix A.1)
and without their special design for coping with "many-arm" setting
"""
struct LazyTrackAndStop
end
long(sr::LazyTrackAndStop) = "LazyTaS";
abbrev(sr::LazyTrackAndStop) = "LT";
mutable struct LazyTrackAndStopState
    sumw;
    w;
    LazyTrackAndStopState(N) = new(zeros(length(N)), [1.0/length(N) for i=1:length(N)]);
end
function start(sr::LazyTrackAndStop, N)
    LazyTrackAndStopState(N);
end
function check_power2(t)
    exponent = floor(Int, log2(t));
    return (t == 2^exponent);
end
function β(sr::LazyTrackAndStop, A, c1, c2, c3, dim, δ, t) # threshold
    #return c2 * log(sqrt(det((1/(0.1*c1))*A + Matrix{Float64}(I,dim,dim))) / δ); # (Eq 9) in main paper of (Jedra and Proutiere, 2020)
    return c2 * (log(1/δ) + 0.5*log(t) + c3); #The version called modified threshold described in appendix A.1 of (Jedra and Proutiere, 2020)
end
function nextsample(sr::LazyTrackAndStopState, pep, N, S, Vinv, A, A0, i0, c0)
    K = length(N); t = sum(N); hμ = Vinv*S; dim = length(hμ); hi = argmax([hμ'pep.arms[k] for k=1:K]);
    if check_power2(t) # lazy update
        w = copy(sr.w);
        for i=1:1000 # the setting in Lazy T&S (Jedra and Proutiere, 2020)
            Vwinv = pinv(sum([w[k]*pep.arms[k]*(pep.arms[k]') for k=1:K]));
            _, ∇f, fidx = compute_f_∇f_linear_bai(w, hμ, 0, pep.arms, Vwinv); # simplified by our Proposition 2
            w_next = zeros(K);
            w_next[argmax(∇f[fidx[1]])] = 1.0;
            w = w*(i/(i+1)) + w_next/(i+1);
            if norm(w_next) / ((i+1)*norm(w)) < 0.001
                break
            end
        end
        setfield!(sr, :w, w);
    end
    setfield!(sr, :sumw, sr.sumw+sr.w);
     # we simplify the arm tracking rule, without implementing their special design for coping issues arised in "many-arm" setting
    (minimum(eigvals(A)) < c0 * sqrt(t)) ? arm = A0[i0+1] : arm = argmin(N - sr.sumw);

    return hi, arm, (i0+1) % dim; # tracking
end

"""
LineGame-C (Degenne et al. 2020) implemented by Xuedong Shang
"""
struct ConvexGame
    TrackingRule
end
long(sr::ConvexGame) = "ConvexGame-" * abbrev(sr.TrackingRule);
abbrev(sr::ConvexGame) = "CG-" * abbrev(sr.TrackingRule);
struct ConvexGameState
    h  # one online learner
    t  # tracking rule
    ConvexGameState(TrackingRule, P) = new(LinBAIAdaHedge(length(P)), TrackingRule(vec(P)))
end
function start(sr::ConvexGame, N, P)
    ConvexGameState(sr.TrackingRule, P)
end
# optimistic gradients
function optimistic_gradient(pep, hμ, t, P::Matrix, λs, Vinv)
    nb_I = nanswers(pep, hµ)
    K = narms(pep, hµ)
    grads = zeros(size(P))
    for k = 1:K
        arm = pep.arms[k]
        for i = 1:nb_I
            ref_value = (hµ .- λs[i, :])'arm
            confidence_width = log(t)
            deviation = sqrt(2 * confidence_width * ((arm') * Vinv * arm))
            ref_value > 0 ? grads[i, k] = 0.5 * (ref_value + deviation)^2 :
            grads[i, k] = 0.5 * (ref_value - deviation)^2
            grads[i, k] = min(grads[i, k], confidence_width)
        end
    end
    return grads
end
function nextsample(sr::ConvexGameState, pep, star, ξ, N, P, S, Vinv)
    nb_I = size(P)[1]
    K = size(P)[2]
    t = sum(N)
    hμ = Vinv * S # emp. estimates
    # query the learner
    vec_W = act(sr.h)
    W = permutedims(reshape(vec_W, (K, nb_I)), (2, 1))  # W is the I*K matrix of answers and pulls.
    # best response λ-player
    _, λs, (_, ξs) = glrt_cgc(pep, W, hμ)
    ∇ = vec(transpose(optimistic_gradient(pep, hμ, t, P, λs, Vinv)))
    incur!(sr.h, -∇)
    # tracking
    big_index = track(sr.t, vec(transpose(P)), vec_W)
    i = div(big_index - 1, K) + 1
    k = ((big_index - 1) % K) + 1
    return i, k
end

"""
LineGame (Degenne et al. 2020) implemented by Xuedong Shang
"""
struct LearnerK
    TrackingRule
end
long(sr::LearnerK) = "LinGame-" * abbrev(sr.TrackingRule);
abbrev(sr::LearnerK) = "Lk-" * abbrev(sr.TrackingRule);
struct LearnerKState
    hs  # I online learners
    t  # tracking rule
    LearnerKState(TrackingRule, N) = new(
        Dict{Int64,LinBAIAdaHedge}(),  # We could allocate one LinBAIAdaHedge for each answer, but for some problems there are 2^d answers.
        TrackingRule(vec(N)),
    )
end
function start(sr::LearnerK, N, P)
    LearnerKState(sr.TrackingRule, N)
end
# optimistic gradients
function optimistic_gradient(pep, hμ, t, N::Vector, λ, Vinv)
    K = length(pep.arms)
    grads = zeros(length(N))
    for k = 1:K
        arm = pep.arms[k]
        ref_value = (hµ .- λ)'arm
        confidence_width = log(t)
        deviation = sqrt(2 * confidence_width * ((arm') * Vinv * arm))
        ref_value > 0 ? grads[k] = 0.5 * (ref_value + deviation)^2 :
        grads[k] = 0.5 * (ref_value - deviation)^2
        grads[k] = min(grads[k], confidence_width)
    end
    return grads
end
function nextsample(sr::LearnerKState, pep, star, ξ, N, P, S, Vinv)
    t = sum(N)
    hμ = Vinv * S # emp. estimates
    nb_I = nanswers(pep, hµ)
    K = length(pep.arms)
    star = istar(pep, hµ)
    if !haskey(sr.hs, star)  # if we never saw that star, initialize an LinBAIAdaHedge learner
        sr.hs[star] = LinBAIAdaHedge(length(N))
    end
    # query the learner
    w = act(sr.hs[star])
    # best response λ-player
    _, (_, λ), (_, ξs) = glrt_cgc(pep, w, hμ)
    ∇ = optimistic_gradient(pep, hμ, t, N, λ, Vinv)
    incur!(sr.hs[star], -∇)
    # tracking
    k = track(sr.t, vec(N), w)
    return star, k
end

"""
Our implementation of XY-Adaptive (Soare et al. 2014) by merging Xuedong Shang' with Tanner Fiez' codes
"""
struct XYAdaptive end
long(sr::XYAdaptive) = "XY-Adaptive";
abbrev(sr::XYAdaptive) = "XY-A";
mutable struct XYAdaptiveState
    phase;
    XYAdaptiveState() = new(0);
end
function start(sr::XYAdaptive, N)
    XYAdaptiveState();
end
function randmin(vector, rank = 1)
   # returns an integer, not a CartesianIndex
    vector = vec(vector)
    Sorted = sort(vector, rev = false)
    m = Sorted[rank]
    Ind = findall(x -> x == m, vector)
    index = Ind[floor(Int, length(Ind) * rand())+1]
    return index
end
function drop_arms(Xactive, Vinv, μ, phase, δ)
    X = copy(Xactive)
    K = length(Xactive)
    for i = 1:K
        arm = X[i]
        for j = 1:K
            if j == i
                continue
            end
            arm_prime = X[j]
            y = arm_prime - arm
            # After replacing Xuedong's implementation with Fiez' implementation: https://github.com/fiezt/Transductive-Linear-Bandit-Code/blob/master/XY_ADAPTIVE.py#L125,
            # the below ↓ condition is consistent to Eq (11) in (Soare 2014).
            if (y' * Vinv * y * 2 * log(2*π^2*K^2*phase^2/(6*δ)))^0.5 <= y'μ
                filter!(x -> x ≠ arm, Xactive)
                break
            end
        end
    end
    return Xactive
end
function nextsample(sr::XYAdaptiveState, pep, N, S, Vinv, Xactive, α, ρ, ρ_old, t_old, δ)
    t = sum(N)
    hμ = Vinv * S # emp. estimates
    nb_I = nanswers(pep, hµ)
    K = length(Xactive)
    star = istar(pep, hµ)
    Y = build_gaps(Xactive)
    nb_gaps = length(Y)
    k = randmin([maximum([transpose(Y[i]) * sherman_morrison(Vinv, pep.arms[j]) * Y[i] for i = 1:nb_gaps]) for j = 1:nb_I])
    if ρ / t < α * ρ_old / t_old
        t_old = t;
        ρ_old = ρ;
        Xcopy = copy(pep.arms);
        Xactive = drop_arms(Xcopy, Vinv, hμ, sr.phase, δ);
        setfield!(sr, :phase, sr.phase+1);
    end
    ρ = maximum([transpose(Y[i]) * sherman_morrison(Vinv, pep.arms[k]) * Y[i] for i = 1:nb_gaps])
    return star, k, Xactive, ρ, ρ_old, t_old
end

"""
LinGapE (Xu et al. 2018) implemented by Xuedong Shang
"""
struct LinGapE end
long(sr::LinGapE) = "LinGapE";
abbrev(sr::LinGapE) = "LG";
function start(sr::LinGapE, N)
    return sr
end
function gap(arm1, arm2, μ)
    (arm1 - arm2)'μ
end
function confidence(arm1, arm2, Vinv)
    sqrt(transpose(arm1 - arm2) * Vinv * (arm1 - arm2))
end
function nextsample(sr::LinGapE, pep, star, ξ, N, S, Vinv, β)
    hμ = Vinv * S # emp. estimates
    nb_I = nanswers(pep, hµ)
    K = length(pep.arms)
    star = istar(pep, hµ)
    c_t = sqrt(2 * β)
    ucb, ambiguous = findmax([gap(pep.arms[i], pep.arms[star], hμ) +
                              confidence(pep.arms[i], pep.arms[star], Vinv) * c_t for i = 1:K])
    k = argmin([confidence(
        pep.arms[star],
        pep.arms[ambiguous],
        sherman_morrison(Vinv, pep.arms[i]),
    ) for i = 1:K])
    return star, k, ucb
end

"""
Uniform sampling implemented by Xuedong Shang
"""
struct RoundRobin end
long(sr::RoundRobin) = "Uniform";
abbrev(sr::RoundRobin) = "RR";
function start(sr::RoundRobin, N)
    return sr
end
function nextsample(sr::RoundRobin, pep, N, S, Vinv)
    return (1 + (sum(N) % length(N))) * ones(Int64, 2)
end


text:
"""
    HS071(
        enable_hessian::Bool,
        enable_hessian_vector_product::Bool = false,
    )

An AbstractNLPEvaluator for the problem:
```
min x1 * x4 * (x1 + x2 + x3) + x3
st  x1 * x2 * x3 * x4 >= 25
    x1^2 + x2^2 + x3^2 + x4^2 = 40
    1 <= x1, x2, x3, x4 <= 5
```
Start at (1,5,5,1)
End at (1.000..., 4.743..., 3.821..., 1.379...)
"""
struct HS071 <: MOI.AbstractNLPEvaluator
    enable_hessian::Bool
    enable_hessian_vector_product::Bool
    function HS071(
        enable_hessian::Bool,
        enable_hessian_vector_product::Bool = false,
    )
        return new(enable_hessian, enable_hessian_vector_product)
    end
end

function MOI.initialize(d::HS071, requested_features::Vector{Symbol})
    for feat in requested_features
        if !(feat in MOI.features_available(d))
            error("Unsupported feature $feat")
            # TODO: implement Jac-vec products for solvers that need them
        end
    end
end

function MOI.features_available(d::HS071)
    features = [:Grad, :Jac, :ExprGraph]
    if d.enable_hessian
        push!(features, :Hess)
    end
    if d.enable_hessian_vector_product
        return push!(features, :HessVec)
    end
    return features
end

function MOI.objective_expr(::HS071)
    return :(
        x[$(MOI.VariableIndex(1))] *
        x[$(MOI.VariableIndex(4))] *
        (
            x[$(MOI.VariableIndex(1))] +
            x[$(MOI.VariableIndex(2))] +
            x[$(MOI.VariableIndex(3))]
        ) + x[$(MOI.VariableIndex(3))]
    )
end

function MOI.constraint_expr(::HS071, i::Int)
    if i == 1
        return :(
            x[$(MOI.VariableIndex(1))] *
            x[$(MOI.VariableIndex(2))] *
            x[$(MOI.VariableIndex(3))] *
            x[$(MOI.VariableIndex(4))] >= 25.0
        )
    elseif i == 2
        return :(
            x[$(MOI.VariableIndex(1))]^2 +
            x[$(MOI.VariableIndex(2))]^2 +
            x[$(MOI.VariableIndex(3))]^2 +
            x[$(MOI.VariableIndex(4))]^2 == 40.0
        )
    else
        error("Out of bounds constraint.")
    end
end

MOI.eval_objective(::HS071, x) = x[1] * x[4] * (x[1] + x[2] + x[3]) + x[3]

function MOI.eval_constraint(::HS071, g, x)
    g[1] = x[1] * x[2] * x[3] * x[4]
    g[2] = x[1]^2 + x[2]^2 + x[3]^2 + x[4]^2
    return
end

function MOI.eval_objective_gradient(::HS071, grad_f, x)
    grad_f[1] = x[1] * x[4] + x[4] * (x[1] + x[2] + x[3])
    grad_f[2] = x[1] * x[4]
    grad_f[3] = x[1] * x[4] + 1
    grad_f[4] = x[1] * (x[1] + x[2] + x[3])
    return
end

function MOI.jacobian_structure(::HS071)
    return Tuple{Int64,Int64}[
        (1, 1),
        (1, 2),
        (1, 3),
        (1, 4),
        (2, 1),
        (2, 2),
        (2, 3),
        (2, 4),
    ]
end

function MOI.hessian_lagrangian_structure(d::HS071)
    @assert d.enable_hessian
    return Tuple{Int64,Int64}[
        (1, 1),
        (2, 1),
        (2, 2),
        (3, 1),
        (3, 2),
        (3, 3),
        (4, 1),
        (4, 2),
        (4, 3),
        (4, 4),
    ]
end

function MOI.eval_constraint_jacobian(::HS071, J, x)
    # Constraint (row) 1
    J[1] = x[2] * x[3] * x[4]  # 1,1
    J[2] = x[1] * x[3] * x[4]  # 1,2
    J[3] = x[1] * x[2] * x[4]  # 1,3
    J[4] = x[1] * x[2] * x[3]  # 1,4
    # Constraint (row) 2
    J[5] = 2 * x[1]  # 2,1
    J[6] = 2 * x[2]  # 2,2
    J[7] = 2 * x[3]  # 2,3
    J[8] = 2 * x[4]  # 2,4
    return
end

function MOI.eval_constraint_jacobian_transpose_product(::HS071, y, x, w)
    y[1] = (x[2] * x[3] * x[4]) * w[1] + (2 * x[1]) * w[2]
    y[2] = (x[1] * x[3] * x[4]) * w[1] + (2 * x[2]) * w[2]
    y[3] = (x[1] * x[2] * x[4]) * w[1] + (2 * x[3]) * w[2]
    y[4] = (x[1] * x[2] * x[3]) * w[1] + (2 * x[4]) * w[2]
    return
end

function MOI.eval_hessian_lagrangian(d::HS071, H, x, σ, μ)
    @assert d.enable_hessian
    # Again, only lower left triangle
    # Objective
    H[1] = σ * (2 * x[4])               # 1,1
    H[2] = σ * (x[4])               # 2,1
    H[3] = 0                          # 2,2
    H[4] = σ * (x[4])               # 3,1
    H[5] = 0                          # 3,2
    H[6] = 0                          # 3,3
    H[7] = σ * (2 * x[1] + x[2] + x[3])  # 4,1
    H[8] = σ * (x[1])               # 4,2
    H[9] = σ * (x[1])               # 4,3
    H[10] = 0                         # 4,4
    # First constraint
    H[2] += μ[1] * (x[3] * x[4])  # 2,1
    H[4] += μ[1] * (x[2] * x[4])  # 3,1
    H[5] += μ[1] * (x[1] * x[4])  # 3,2
    H[7] += μ[1] * (x[2] * x[3])  # 4,1
    H[8] += μ[1] * (x[1] * x[3])  # 4,2
    H[9] += μ[1] * (x[1] * x[2])  # 4,3
    # Second constraint
    H[1] += μ[2] * 2  # 1,1
    H[3] += μ[2] * 2  # 2,2
    H[6] += μ[2] * 2  # 3,3
    H[10] += μ[2] * 2
    return
end

function MOI.eval_hessian_lagrangian_product(d::HS071, h, x, v, σ, μ)
    @assert d.enable_hessian_vector_product
    # Objective
    h[1] =
        2.0 * x[4] * v[1] +
        x[4] * v[2] +
        x[4] * v[3] +
        (2.0 * x[1] + x[2] + x[3]) * v[4]
    h[2] = x[4] * v[1] + x[1] * v[4]
    h[3] = x[4] * v[1] + x[1] * v[4]
    h[4] = (2.0 * x[1] + x[2] + x[3]) * v[1] + x[1] * v[2] + x[1] * v[3]
    h .*= σ
    # First constraint
    h[1] +=
        μ[1] * (x[3] * x[4] * v[2] + x[2] * x[4] * v[3] + x[2] * x[3] * v[4])
    h[2] +=
        μ[1] * (x[3] * x[4] * v[1] + x[1] * x[4] * v[3] + x[1] * x[3] * v[4])
    h[3] +=
        μ[1] * (x[2] * x[4] * v[1] + x[1] * x[4] * v[2] + x[1] * x[2] * v[4])
    h[4] +=
        μ[1] * (x[2] * x[3] * v[1] + x[1] * x[3] * v[2] + x[1] * x[2] * v[3])
    # Second constraint
    h[1] += μ[2] * 2.0 * v[1]
    h[2] += μ[2] * 2.0 * v[2]
    h[3] += μ[2] * 2.0 * v[3]
    h[4] += μ[2] * 2.0 * v[4]
    return
end

"""
    FeasibilitySenseEvaluator(enable_hessian::Bool)

An AbstractNLPEvaluator for the problem:
```
Test for FEASIBILITY_SENSE.
Find x satisfying x^2 == 1.
```
"""
struct FeasibilitySenseEvaluator <: MOI.AbstractNLPEvaluator
    enable_hessian::Bool
end

function MOI.initialize(
    d::FeasibilitySenseEvaluator,
    requested_features::Vector{Symbol},
)
    for feat in requested_features
        if !(feat in MOI.features_available(d))
            error("Unsupported feature $feat")
            # TODO: implement Jac-vec and Hess-vec products
            # for solvers that need them
        end
    end
    return
end

function MOI.features_available(d::FeasibilitySenseEvaluator)
    if d.enable_hessian
        return [:Grad, :Jac, :Hess, :ExprGraph]
    else
        return [:Grad, :Jac, :ExprGraph]
    end
end

MOI.objective_expr(::FeasibilitySenseEvaluator) = :()

function MOI.constraint_expr(::FeasibilitySenseEvaluator, i::Int)
    @assert i == 1
    return :(x[$(MOI.VariableIndex(1))]^2 == 1)
end

MOI.eval_objective(d::FeasibilitySenseEvaluator, x) = 0.0

function MOI.eval_constraint(::FeasibilitySenseEvaluator, g, x)
    g[1] = x[1]^2
    return
end

function MOI.eval_objective_gradient(::FeasibilitySenseEvaluator, grad_f, x)
    grad_f[1] = 0.0
    return
end

function MOI.jacobian_structure(::FeasibilitySenseEvaluator)
    return Tuple{Int64,Int64}[(1, 1)]
end

function MOI.hessian_lagrangian_structure(d::FeasibilitySenseEvaluator)
    @assert d.enable_hessian
    return Tuple{Int64,Int64}[(1, 1)]
end

function MOI.eval_constraint_jacobian(::FeasibilitySenseEvaluator, J, x)
    J[1] = 2 * x[1]
    return
end

function MOI.eval_hessian_lagrangian(d::FeasibilitySenseEvaluator, H, x, σ, μ)
    @assert d.enable_hessian
    H[1] = 2 * μ[1] # 1,1
    return
end

function _test_HS071(model::MOI.ModelLike, config::Config, evaluator::HS071)
    @requires MOI.supports(model, MOI.NLPBlock())
    @requires _supports(config, MOI.optimize!)
    @requires MOI.supports_constraint(
        model,
        MOI.VariableIndex,
        MOI.LessThan{Float64},
    )
    @requires MOI.supports_constraint(
        model,
        MOI.VariableIndex,
        MOI.GreaterThan{Float64},
    )
    @requires MOI.supports(model, MOI.VariablePrimalStart(), MOI.VariableIndex)
    lb = [25.0, 40.0]
    ub = [Inf, 40.0]
    block_data = MOI.NLPBlockData(MOI.NLPBoundsPair.(lb, ub), evaluator, true)
    v = MOI.add_variables(model, 4)
    @test MOI.get(model, MOI.NumberOfVariables()) == 4
    l = [1.0, 1.0, 1.0, 1.0]
    u = [5.0, 5.0, 5.0, 5.0]
    start = [1, 5, 5, 1]
    for i in 1:4
        cub = MOI.add_constraint(model, v[i], MOI.LessThan(u[i]))
        # We test this after the creation of every `VariableIndex` constraint
        # to ensure a good coverage of corner cases.
        @test cub.value == v[i].value
        clb = MOI.add_constraint(model, v[i], MOI.GreaterThan(l[i]))
        @test clb.value == v[i].value
        MOI.set(model, MOI.VariablePrimalStart(), v[i], start[i])
    end
    MOI.set(model, MOI.NLPBlock(), block_data)
    MOI.set(model, MOI.ObjectiveSense(), MOI.MIN_SENSE)
    MOI.optimize!(model)
    @test MOI.get(model, MOI.TerminationStatus()) == config.optimal_status
    @test MOI.get(model, MOI.ResultCount()) >= 1
    @test MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT
    @test isapprox(
        MOI.get(model, MOI.ObjectiveValue()),
        17.014017145179164,
        config,
    )
    @test isapprox(
        MOI.get(model, MOI.VariablePrimal(), v),
        [1.0, 4.7429996418092970, 3.8211499817883077, 1.3794082897556983],
        config,
    )
    # TODO: Duals? Maybe better to test on a convex instance.
    return
end

"""
    test_nonlinear_hs071(model::MOI.ModelLike, config::Config)

Test the nonlinear HS071 problem.
"""
function test_nonlinear_hs071(model::MOI.ModelLike, config::Config)
    _test_HS071(model, config, HS071(true))
    return
end

function setup_test(
    ::typeof(test_nonlinear_hs071),
    model::MOIU.MockOptimizer,
    config::Config,
)
    config.optimal_status = MOI.LOCALLY_SOLVED
    MOI.Utilities.set_mock_optimize!(
        model,
        (mock) -> begin
            MOI.Utilities.mock_optimize!(
                mock,
                config.optimal_status,
                [
                    1.0,
                    4.7429996418092970,
                    3.8211499817883077,
                    1.379408289755698,
                ],
            )
            MOI.set(mock, MOI.ObjectiveValue(), 17.014017145179164)
        end,
    )
    flag = model.eval_objective_value
    model.eval_objective_value = false
    return () -> model.eval_objective_value = flag
end

"""
    test_nonlinear_hs071_no_hessian(model::MOI.ModelLike, config::Config)

Test the nonlinear HS071 problem without hessians.
"""
function test_nonlinear_hs071_no_hessian(model, config)
    _test_HS071(model, config, HS071(false))
    return
end

function setup_test(
    ::typeof(test_nonlinear_hs071_no_hessian),
    model::MOIU.MockOptimizer,
    config::Config,
)
    config.optimal_status = MOI.LOCALLY_SOLVED
    MOI.Utilities.set_mock_optimize!(
        model,
        (mock) -> begin
            MOI.Utilities.mock_optimize!(
                mock,
                config.optimal_status,
                [
                    1.0,
                    4.7429996418092970,
                    3.8211499817883077,
                    1.379408289755698,
                ],
            )
            MOI.set(mock, MOI.ObjectiveValue(), 17.014017145179164)
        end,
    )
    flag = model.eval_objective_value
    model.eval_objective_value = false
    return () -> model.eval_objective_value = flag
end

"""
    test_nonlinear_hs071_hessian_vector_product(
        model::MOI.ModelLike,
        config::Config,
    )

Test the nonlinear HS071 problem with a hessian vector product.
"""
function test_nonlinear_hs071_hessian_vector_product(model, config)
    _test_HS071(model, config, HS071(false, true))
    return
end

function setup_test(
    ::typeof(test_nonlinear_hs071_hessian_vector_product),
    model::MOIU.MockOptimizer,
    config::Config,
)
    config.optimal_status = MOI.LOCALLY_SOLVED
    MOI.Utilities.set_mock_optimize!(
        model,
        (mock) -> begin
            MOI.Utilities.mock_optimize!(
                mock,
                config.optimal_status,
                [
                    1.0,
                    4.7429996418092970,
                    3.8211499817883077,
                    1.379408289755698,
                ],
            )
            MOI.set(mock, MOI.ObjectiveValue(), 17.014017145179164)
        end,
    )
    flag = model.eval_objective_value
    model.eval_objective_value = false
    return () -> model.eval_objective_value = flag
end

"""
    test_nonlinear_hs071_NLPBlockDual(model::MOI.ModelLike, config::Config)

Test NLPBlockDual.
"""
function test_nonlinear_hs071_NLPBlockDual(model::MOI.ModelLike, config::Config)
    @requires MOI.supports(model, MOI.NLPBlock())
    @requires _supports(config, MOI.optimize!)
    @requires MOI.supports(model, MOI.VariablePrimalStart(), MOI.VariableIndex)
    @requires _supports(config, MOI.NLPBlockDual)
    n_variables = 4
    v = MOI.add_variables(model, n_variables)
    l = [1.1, 1.2, 1.3, 1.4]
    u = [5.1, 5.2, 5.3, 5.4]
    start = [2.1, 2.2, 2.3, 2.4]
    cl = MOI.add_constraint.(model, v, MOI.GreaterThan.(l))
    cu = MOI.add_constraint.(model, v, MOI.LessThan.(u))
    MOI.set.(model, MOI.VariablePrimalStart(), v, start)
    lb, ub = [25.0, 40.0], [Inf, 40.0]
    evaluator = MOI.Test.HS071(true)
    block_data = MOI.NLPBlockData(MOI.NLPBoundsPair.(lb, ub), evaluator, true)
    MOI.set(model, MOI.NLPBlock(), block_data)

    for sense in [MOI.MIN_SENSE, MOI.MAX_SENSE]
        MOI.set(model, MOI.ObjectiveSense(), sense)
        MOI.optimize!(model)
        # Get primal solution
        x = MOI.get(model, MOI.VariablePrimal(), v)
        # Get dual solution
        con_dual = MOI.get(model, MOI.NLPBlockDual())
        var_dual_lb = MOI.get(model, MOI.ConstraintDual(), cl)
        var_dual_ub = MOI.get(model, MOI.ConstraintDual(), cu)

        # Evaluate ∇f(x)
        g = zeros(n_variables)
        MOI.eval_objective_gradient(evaluator, g, x)
        # Evaluate ∑ μᵢ * ∇cᵢ(x)
        jtv = zeros(n_variables)
        MOI.eval_constraint_jacobian_transpose_product(
            evaluator,
            jtv,
            x,
            con_dual,
        )

        # Test that (x, μ, νl, νᵤ) satisfies stationarity condition
        # σ ∇f(x) - ∑ μᵢ * ∇cᵢ(x) - νₗ - νᵤ = 0
        σ = (sense == MOI.MAX_SENSE) ? -1.0 : 1.0
        @test isapprox(σ .* g, jtv .+ var_dual_ub .+ var_dual_lb, config)
    end
    return
end

function setup_test(
    ::typeof(test_nonlinear_hs071_NLPBlockDual),
    model::MOIU.MockOptimizer,
    config::Config,
)
    config.optimal_status = MOI.LOCALLY_SOLVED
    MOI.Utilities.set_mock_optimize!(
        model,
        (mock) -> begin
            MOI.Utilities.mock_optimize!(
                mock,
                config.optimal_status,
                [1.1, 1.5735520521364397, 2.6746837915674373, 5.4],
                (MOI.VariableIndex, MOI.GreaterThan{Float64}) =>
                    [28.590703804861093, 0.0, 0.0, 0.0],
                (MOI.VariableIndex, MOI.LessThan{Float64}) =>
                    [0.0, 0.0, 0.0, -5.582550550234756],
            )
            MOI.set(mock, MOI.NLPBlockDual(), [0.178761800, 0.985000823])
        end,
        (mock) -> begin
            MOI.Utilities.mock_optimize!(
                mock,
                config.optimal_status,
                [
                    4.567633003532132,
                    1.6613736769916652,
                    1.7612041983937636,
                    3.6434497419346723,
                ],
                (MOI.VariableIndex, MOI.GreaterThan{Float64}) =>
                    [0.0, 0.0, 0.0, 0.0],
                (MOI.VariableIndex, MOI.LessThan{Float64}) =>
                    [0.0, 0.0, 0.0, 0.0],
            )
            MOI.set(mock, MOI.NLPBlockDual(), [0.0, -5.008488315])
        end,
    )
    model.eval_variable_constraint_dual = false
    return () -> model.eval_variable_constraint_dual = true
end

"""
    test_nonlinear_objective_and_moi_objective_test(
        model::MOI.ModelLike,
        config::Config,
    )

Test that nonlinear objectives take precedence over MOI.ObjectiveFunction.
"""
function test_nonlinear_objective_and_moi_objective_test(
    model::MOI.ModelLike,
    config::Config,
)
    @requires MOI.supports(model, MOI.NLPBlock())
    @requires _supports(config, MOI.optimize!)
    @requires MOI.supports(model, MOI.VariablePrimalStart(), MOI.VariableIndex)
    lb = [1.0]
    ub = [1.0]
    block_data = MOI.NLPBlockData(
        MOI.NLPBoundsPair.(lb, ub),
        FeasibilitySenseEvaluator(true),
        true,
    )
    x = MOI.add_variable(model)
    @test MOI.get(model, MOI.NumberOfVariables()) == 1
    # Avoid starting at zero because it's a critical point.
    MOI.set(model, MOI.VariablePrimalStart(), x, 1.5)
    MOI.set(model, MOI.NLPBlock(), block_data)
    MOI.set(model, MOI.ObjectiveSense(), MOI.MAX_SENSE)
    f_x = MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(1.0, x)], 0.0)
    # This objective function should be ignored.
    MOI.set(model, MOI.ObjectiveFunction{typeof(f_x)}(), f_x)
    MOI.optimize!(model)
    @test MOI.get(model, MOI.TerminationStatus()) == config.optimal_status
    @test MOI.get(model, MOI.ResultCount()) >= 1
    @test MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT
    @test isapprox(MOI.get(model, MOI.ObjectiveValue()), 0.0, config)
    xv = MOI.get(model, MOI.VariablePrimal(), x)
    @test isapprox(abs(xv), 1.0, config)
    return
end

function setup_test(
    ::typeof(test_nonlinear_objective_and_moi_objective_test),
    model::MOIU.MockOptimizer,
    config::Config,
)
    config.optimal_status = MOI.LOCALLY_SOLVED
    MOI.Utilities.set_mock_optimize!(
        model,
        (mock) -> begin
            MOI.Utilities.mock_optimize!(mock, config.optimal_status, [-1.0]),
            MOI.set(mock, MOI.ObjectiveValue(), 0.0)
        end,
    )
    flag = model.eval_objective_value
    model.eval_objective_value = false
    return () -> model.eval_objective_value = flag
end

"""
    test_nonlinear_objective(model::MOI.ModelLike, config::Config)

Test a nonlinear objective only.
"""
function test_nonlinear_objective(model::MOI.ModelLike, config::Config)
    @requires MOI.supports(model, MOI.NLPBlock())
    @requires _supports(config, MOI.optimize!)
    @requires MOI.supports(model, MOI.VariablePrimalStart(), MOI.VariableIndex)
    lb = [1.0]
    ub = [1.0]
    block_data = MOI.NLPBlockData(
        MOI.NLPBoundsPair.(lb, ub),
        FeasibilitySenseEvaluator(true),
        true,
    )
    x = MOI.add_variable(model)
    @test MOI.get(model, MOI.NumberOfVariables()) == 1
    # Avoid starting at zero because it's a critical point.
    MOI.set(model, MOI.VariablePrimalStart(), x, 1.5)
    MOI.set(model, MOI.NLPBlock(), block_data)
    MOI.set(model, MOI.ObjectiveSense(), MOI.MAX_SENSE)
    MOI.optimize!(model)
    @test MOI.get(model, MOI.TerminationStatus()) == config.optimal_status
    @test MOI.get(model, MOI.ResultCount()) >= 1
    @test MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT
    @test isapprox(MOI.get(model, MOI.ObjectiveValue()), 0.0, config)
    xv = MOI.get(model, MOI.VariablePrimal(), x)
    @test isapprox(abs(xv), 1.0, config)
    return
end

function setup_test(
    ::typeof(test_nonlinear_objective),
    model::MOIU.MockOptimizer,
    config::Config,
)
    config.optimal_status = MOI.LOCALLY_SOLVED
    MOI.Utilities.set_mock_optimize!(
        model,
        (mock) -> begin
            MOI.Utilities.mock_optimize!(mock, config.optimal_status, [-1.0]),
            MOI.set(mock, MOI.ObjectiveValue(), 0.0)
        end,
    )
    flag = model.eval_objective_value
    model.eval_objective_value = false
    return () -> model.eval_objective_value = flag
end

"""
    test_nonlinear_without_objective(model::MOI.ModelLike, config::Config)

Test a nonlinear problem without an objective.
"""
function test_nonlinear_without_objective(model::MOI.ModelLike, config::Config)
    @requires MOI.supports(model, MOI.NLPBlock())
    @requires _supports(config, MOI.optimize!)
    @requires MOI.supports(model, MOI.VariablePrimalStart(), MOI.VariableIndex)
    lb = [1.0]
    ub = [1.0]
    block_data = MOI.NLPBlockData(
        MOI.NLPBoundsPair.(lb, ub),
        FeasibilitySenseEvaluator(true),
        false,
    )
    x = MOI.add_variable(model)
    @test MOI.get(model, MOI.NumberOfVariables()) == 1
    # Avoid starting at zero because it's a critical point.
    MOI.set(model, MOI.VariablePrimalStart(), x, 1.5)
    MOI.set(model, MOI.NLPBlock(), block_data)
    MOI.optimize!(model)
    @test MOI.get(model, MOI.TerminationStatus()) == config.optimal_status
    @test MOI.get(model, MOI.ResultCount()) >= 1
    @test MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT
    @test isapprox(MOI.get(model, MOI.ObjectiveValue()), 0.0, config)
    xv = MOI.get(model, MOI.VariablePrimal(), x)
    @test isapprox(abs(xv), 1.0, config)
    return
end

function setup_test(
    ::typeof(test_nonlinear_without_objective),
    model::MOIU.MockOptimizer,
    config::Config,
)
    config.optimal_status = MOI.LOCALLY_SOLVED
    MOI.Utilities.set_mock_optimize!(
        model,
        (mock) -> begin
            MOI.Utilities.mock_optimize!(mock, config.optimal_status, [-1.0]),
            MOI.set(mock, MOI.ObjectiveValue(), 0.0)
        end,
    )
    flag = model.eval_objective_value
    model.eval_objective_value = false
    return () -> model.eval_objective_value = flag
end

"""
    test_nonlinear_mixed_complementarity(model::MOI.ModelLike, config::Config)

Test the solution of the linear mixed-complementarity problem:
`F(x) complements x`, where `F(x) = M * x .+ q` and `0 <= x <= 10`.
"""
function test_nonlinear_mixed_complementarity(
    model::MOI.ModelLike,
    config::Config,
)
    @requires MOI.supports_constraint(
        model,
        MOI.VectorAffineFunction{Float64},
        MOI.Complements,
    )
    @requires _supports(config, MOI.optimize!)
    x = MOI.add_variables(model, 4)
    MOI.add_constraint.(model, x, MOI.Interval(0.0, 10.0))
    MOI.set.(model, MOI.VariablePrimalStart(), x, 0.0)
    M = Float64[0 0 -1 -1; 0 0 1 -2; 1 -1 2 -2; 1 2 -2 4]
    q = [2; 2; -2; -6]
    terms = MOI.VectorAffineTerm{Float64}[]
    for i in 1:4
        push!(
            terms,
            MOI.VectorAffineTerm(4 + i, MOI.ScalarAffineTerm(1.0, x[i])),
        )
        for j in 1:4
            iszero(M[i, j]) && continue
            push!(
                terms,
                MOI.VectorAffineTerm(i, MOI.ScalarAffineTerm(M[i, j], x[j])),
            )
        end
    end
    MOI.add_constraint(
        model,
        MOI.VectorAffineFunction(terms, [q; 0.0; 0.0; 0.0; 0.0]),
        MOI.Complements(8),
    )
    @test MOI.get(model, MOI.TerminationStatus()) == MOI.OPTIMIZE_NOT_CALLED
    MOI.optimize!(model)
    @test MOI.get(model, MOI.TerminationStatus()) == config.optimal_status
    x_val = MOI.get.(model, MOI.VariablePrimal(), x)
    @test isapprox(x_val, [2.8, 0.0, 0.8, 1.2], config)
    return
end

function setup_test(
    ::typeof(test_nonlinear_mixed_complementarity),
    model::MOIU.MockOptimizer,
    config::Config,
)
    config.optimal_status = MOI.LOCALLY_SOLVED
    MOI.Utilities.set_mock_optimize!(
        model,
        (mock) -> MOI.Utilities.mock_optimize!(
            mock,
            config.optimal_status,
            [2.8, 0.0, 0.8, 1.2],
        ),
    )
    return
end

"""
    test_nonlinear_qp_complementarity_constraint(
        model::MOI.ModelLike,
        config::Config,
    )

Test the solution of the quadratic program with complementarity constraints:

```
    min  (x0 - 5)^2 +(2 x1 + 1)^2
    s.t.  -1.5 x0 + 2 x1 + x2 - 0.5 x3 + x4 = 2
        x2 complements(3 x0 - x1 - 3)
        x3 complements(-x0 + 0.5 x1 + 4)
        x4 complements(-x0 - x1 + 7)
        x0, x1, x2, x3, x4 >= 0

```
which rewrites, with auxiliary variables

```
    min  (x0 - 5)^2 +(2 x1 + 1)^2
    s.t.  -1.5 x0 + 2 x1 + x2 - 0.5 x3 + x4 = 2  (cf1)
        3 x0 - x1 - 3 - x5 = 0                 (cf2)
        -x0 + 0.5 x1 + 4 - x6 = 0              (cf3)
        -x0 - x1 + 7 - x7 = 0                  (cf4)
        x2 complements x5
        x3 complements x6
        x4 complements x7
        x0, x1, x2, x3, x4, x5, x6, x7 >= 0

```
"""
function test_nonlinear_qp_complementarity_constraint(
    model::MOI.ModelLike,
    config::Config,
)
    @requires MOI.supports_constraint(
        model,
        MOI.VectorOfVariables,
        MOI.Complements,
    )
    @requires _supports(config, MOI.optimize!)
    x = MOI.add_variables(model, 8)
    MOI.set.(model, MOI.VariablePrimalStart(), x, 0.0)
    MOI.add_constraint.(model, x, MOI.GreaterThan(0.0))
    MOI.set(
        model,
        MOI.ObjectiveFunction{MOI.ScalarQuadraticFunction{Float64}}(),
        MOI.ScalarQuadraticFunction(
            MOI.ScalarQuadraticTerm.([2.0, 8.0], x[1:2], x[1:2]),
            MOI.ScalarAffineTerm.([-10.0, 4.0], x[[1, 2]]),
            26.0,
        ),
    )
    MOI.add_constraint(
        model,
        MOI.ScalarAffineFunction(
            MOI.ScalarAffineTerm.([-1.5, 2.0, 1.0, 0.5, 1.0], x[1:5]),
            0.0,
        ),
        MOI.EqualTo(2.0),
    )
    MOI.add_constraint(
        model,
        MOI.ScalarAffineFunction(
            MOI.ScalarAffineTerm.([3.0, -1.0, -1.0], x[[1, 2, 6]]),
            0.0,
        ),
        MOI.EqualTo(3.0),
    )
    MOI.add_constraint(
        model,
        MOI.ScalarAffineFunction(
            MOI.ScalarAffineTerm.([-1.0, 0.5, -1.0], x[[1, 2, 7]]),
            0.0,
        ),
        MOI.EqualTo(-4.0),
    )
    MOI.add_constraint(
        model,
        MOI.ScalarAffineFunction(
            MOI.ScalarAffineTerm.(-1.0, x[[1, 2, 8]]),
            0.0,
        ),
        MOI.EqualTo(-7.0),
    )
    MOI.add_constraint(
        model,
        MOI.VectorOfVariables([x[3], x[4], x[5], x[6], x[7], x[8]]),
        MOI.Complements(6),
    )
    @test MOI.get(model, MOI.TerminationStatus()) == MOI.OPTIMIZE_NOT_CALLED
    MOI.optimize!(model)
    @test MOI.get(model, MOI.TerminationStatus()) == config.optimal_status
    x_val = MOI.get.(model, MOI.VariablePrimal(), x)
    @test isapprox(x_val, [1.0, 0.0, 3.5, 0.0, 0.0, 0.0, 3.0, 6.0], config)
    @test isapprox(MOI.get(model, MOI.ObjectiveValue()), 17.0, config)
    return
end

function setup_test(
    ::typeof(test_nonlinear_qp_complementarity_constraint),
    model::MOIU.MockOptimizer,
    config::Config,
)
    config.optimal_status = MOI.LOCALLY_SOLVED
    MOI.Utilities.set_mock_optimize!(
        model,
        (mock) -> begin
            MOI.Utilities.mock_optimize!(
                mock,
                config.optimal_status,
                [1.0, 0.0, 3.5, 0.0, 0.0, 0.0, 3.0, 6.0],
            )
            MOI.set(mock, MOI.ObjectiveValue(), 17.0)
        end,
    )
    return
end

"""
    test_nonlinear_HS071_internal(::MOI.ModelLike, ::Config)

A test for the correctness of the HS071 evaluator.

This is mainly for the internal purpose of checking their correctness as
written. External solvers can exclude this test without consequence.
"""
function test_nonlinear_HS071_internal(::MOI.ModelLike, ::Config)
    d = MOI.Test.HS071(true, true)
    @test MOI.objective_expr(d) == :(
        x[$(MOI.VariableIndex(1))] *
        x[$(MOI.VariableIndex(4))] *
        (
            x[$(MOI.VariableIndex(1))] +
            x[$(MOI.VariableIndex(2))] +
            x[$(MOI.VariableIndex(3))]
        ) + x[$(MOI.VariableIndex(3))]
    )
    @test MOI.constraint_expr(d, 1) == :(
        x[$(MOI.VariableIndex(1))] *
        x[$(MOI.VariableIndex(2))] *
        x[$(MOI.VariableIndex(3))] *
        x[$(MOI.VariableIndex(4))] >= 25.0
    )
    @test MOI.constraint_expr(d, 2) == :(
        x[$(MOI.VariableIndex(1))]^2 +
        x[$(MOI.VariableIndex(2))]^2 +
        x[$(MOI.VariableIndex(3))]^2 +
        x[$(MOI.VariableIndex(4))]^2 == 40.0
    )
    @test_throws ErrorException MOI.constraint_expr(d, 3)
    MOI.initialize(d, [:Grad, :Jac, :ExprGraph, :Hess, :HessVec])
    @test :HessVec in MOI.features_available(d)
    x = ones(4)
    # f(x)
    @test MOI.eval_objective(d, x) == 4.0
    # g(x)
    g = zeros(2)
    MOI.eval_constraint(d, g, x)
    @test g == [1.0, 4.0]
    # f'(x)
    ∇f = fill(NaN, length(x))
    MOI.eval_objective_gradient(d, ∇f, x)
    @test ∇f == [4.0, 1.0, 2.0, 3.0]
    # Jacobian
    Js = MOI.jacobian_structure(d)
    J = fill(NaN, length(Js))
    MOI.eval_constraint_jacobian(d, J, x)
    @test J == [1, 1, 1, 1, 2, 2, 2, 2]
    # Hessian-lagrangian
    Hs = MOI.hessian_lagrangian_structure(d)
    H = fill(NaN, length(Hs))
    MOI.eval_hessian_lagrangian(d, H, x, 1.0, [1.0, 1.0])
    @test H == [4, 2, 2, 2, 1, 2, 5, 2, 2, 2]
    # Hessian-lagrangian-product
    Hv = fill(NaN, length(x))
    v = [1.0, 1.1, 1.2, 1.3]
    MOI.eval_hessian_lagrangian_product(d, Hv, x, v, 1.0, [1.0, 1.0])
    @test Hv == [15.1, 8.0, 8.1, 12.2]
    return
end

"""
    test_nonlinear_Feasibility_internal(::MOI.ModelLike, ::Config)

A test for the correctness of the FeasibilitySenseEvaluator evaluator.

This is mainly for the internal purpose of checking their correctness as
written. External solvers can exclude this test without consequence.
"""
function test_nonlinear_Feasibility_internal(::MOI.ModelLike, ::Config)
    d = MOI.Test.FeasibilitySenseEvaluator(true)
    @test MOI.objective_expr(d) == :()
    @test MOI.constraint_expr(d, 1) == :(x[$(MOI.VariableIndex(1))]^2 == 1.0)
    @test_throws AssertionError MOI.constraint_expr(d, 2)
    MOI.initialize(d, [:Grad, :Jac, :ExprGraph, :Hess])
    @test :Hess in MOI.features_available(d)
    x = [1.5]
    # f(x)
    @test MOI.eval_objective(d, x) == 0.0
    # g(x)
    g = zeros(1)
    MOI.eval_constraint(d, g, x)
    @test g == [1.5^2]
    # f'(x)
    ∇f = fill(NaN, length(x))
    MOI.eval_objective_gradient(d, ∇f, x)
    @test ∇f == [0.0]
    # Jacobian
    Js = MOI.jacobian_structure(d)
    J = fill(NaN, length(Js))
    MOI.eval_constraint_jacobian(d, J, x)
    @test J == [3.0]
    # Hessian-lagrangian
    Hs = MOI.hessian_lagrangian_structure(d)
    H = fill(NaN, length(Hs))
    MOI.eval_hessian_lagrangian(d, H, x, 1.0, [1.1])
    @test H == [2.2]
    return
end

"""
    InvalidEvaluator()

An AbstractNLPEvaluator for the problem: `min NaN`, where `eval_objective`
returns a `NaN`.

Solvers should return a `TerminationStatus` of `MOI.INVALID_MODEL`.
"""
struct InvalidEvaluator <: MOI.AbstractNLPEvaluator end

function MOI.initialize(d::InvalidEvaluator, requested_features::Vector{Symbol})
    for feat in requested_features
        @assert feat in MOI.features_available(d)
    end
    return
end

function MOI.features_available(::InvalidEvaluator)
    return [:Grad, :ExprGraph]
end

MOI.objective_expr(::InvalidEvaluator) = :(NaN)

MOI.eval_objective(::InvalidEvaluator, x) = NaN

function MOI.eval_objective_gradient(::InvalidEvaluator, grad_f, x)
    grad_f[1] = NaN
    return
end

"""
    test_nonlinear_InvalidEvaluator_internal(::MOI.ModelLike, ::Config)

A test for the correctness of the InvalidEvaluator evaluator.

This is mainly for the internal purpose of checking their correctness as
written. External solvers can exclude this test without consequence.
"""
function test_nonlinear_InvalidEvaluator_internal(::MOI.ModelLike, ::Config)
    d = InvalidEvaluator()
    @test MOI.objective_expr(d) == :(NaN)
    MOI.initialize(d, [:Grad, :ExprGraph])
    x = [1.5]
    # f(x)
    @test isnan(MOI.eval_objective(d, x))
    # f'(x)
    ∇f = fill(NaN, length(x))
    MOI.eval_objective_gradient(d, ∇f, x)
    @test isnan(∇f[1])
    return
end

"""
    test_nonlinear_invalid(model::MOI.ModelLike, config::Config)

Test that a nonlinear model returns the TerminationStatus `INVALID_MODEL` if a
NaN is detected in a function evaluation.
"""
function test_nonlinear_invalid(model::MOI.ModelLike, config::Config)
    @requires MOI.supports(model, MOI.NLPBlock())
    @requires _supports(config, MOI.optimize!)
    MOI.add_variable(model)
    MOI.set(model, MOI.ObjectiveSense(), MOI.MAX_SENSE)
    MOI.set(
        model,
        MOI.NLPBlock(),
        MOI.NLPBlockData(MOI.NLPBoundsPair[], InvalidEvaluator(), true),
    )
    MOI.optimize!(model)
    @test MOI.get(model, MOI.TerminationStatus()) == MOI.INVALID_MODEL
    return
end

function setup_test(
    ::typeof(test_nonlinear_invalid),
    model::MOIU.MockOptimizer,
    ::Config,
)
    MOI.Utilities.set_mock_optimize!(
        model,
        mock -> MOI.Utilities.mock_optimize!(mock, MOI.INVALID_MODEL, [NaN]),
    )
    return
end

ext: jl


text:
abstract type AbstractBlockMatrix{T} <: AbstractMatrix{T} end
struct BlockMatrix{T} <: AbstractBlockMatrix{T}
    blocks::Vector{Matrix{T}}
end
nblocks(bm::BlockMatrix) = length(bm.blocks)
block(bm::BlockMatrix, i::Integer) = bm.blocks[i]

function Base.size(bm::AbstractBlockMatrix)
    n = Compat.mapreduce(blk -> Compat.LinearAlgebra.checksquare(block(bm, blk)),
                         +, 1:nblocks(bm), init=0)
    return (n, n)
end
function Base.getindex(bm::AbstractBlockMatrix, i::Integer, j::Integer)
    (i < 0 || j < 0) && throw(BoundsError(i, j))
    for k in 1:nblocks(bm)
        blk = block(bm, k)
        n = size(blk, 1)
        if i <= n && j <= n
            return blk[i, j]
        elseif i <= n || j <= n
            return 0
        else
            i -= n
            j -= n
        end
    end
    i, j = (i, j) .+ size(bm)
    throw(BoundsError(i, j))
end
Base.getindex(A::AbstractBlockMatrix, I::Tuple) = getindex(A, I...)

mutable struct MockSDOptimizer{T} <: AbstractSDOptimizer
    nconstrs::Int
    blkdims::Vector{Int}
    constraint_constants::Vector{T}
    objective_coefficients::Vector{Tuple{T, Int, Int, Int}}
    constraint_coefficients::Vector{Vector{Tuple{T, Int, Int, Int}}}
    optimize!::Function # Function used to set dummy primal/dual values and statuses
    hasprimal::Bool
    hasdual::Bool
    terminationstatus::MOI.TerminationStatusCode
    primalstatus::MOI.ResultStatusCode
    dualstatus::MOI.ResultStatusCode
    X::BlockMatrix{T}
    Z::BlockMatrix{T}
    y::Vector{T}
end
MockSDOptimizer{T}() where T = MockSDOptimizer{T}(0,
                                                  Int[],
                                                  T[],
                                                  Tuple{T, Int, Int, Int}[],
                                                  Vector{Tuple{T, Int, Int, Int}}[],
                                                  (::MockSDOptimizer) -> begin end,
                                                  false,
                                                  false,
                                                  MOI.OPTIMIZE_NOT_CALLED,
                                                  MOI.NO_SOLUTION,
                                                  MOI.NO_SOLUTION,
                                                  BlockMatrix{T}(Matrix{T}[]),
                                                  BlockMatrix{T}(Matrix{T}[]),
                                                  T[])
mockSDoptimizer(T::Type) = SDOIOptimizer(MockSDOptimizer{T}(), T)

MOI.get(::MockSDOptimizer, ::MOI.SolverName) = "MockSD"

function MOI.empty!(mock::MockSDOptimizer{T}) where T
    mock.nconstrs = 0
    mock.blkdims = Int[]
    mock.constraint_constants = T[]
    mock.objective_coefficients = Tuple{T, Int, Int, Int}[]
    mock.constraint_coefficients = Vector{Tuple{T, Int, Int, Int}}[]
    mock.hasprimal = false
    mock.hasdual = false
    mock.terminationstatus = MOI.OPTIMIZE_NOT_CALLED
    mock.primalstatus = MOI.NO_SOLUTION
    mock.dualstatus = MOI.NO_SOLUTION
    mock.X = BlockMatrix{T}(Matrix{T}[])
    mock.Z = BlockMatrix{T}(Matrix{T}[])
    mock.y = T[]
end
coefficienttype(::MockSDOptimizer{T}) where T = T

getnumberofconstraints(optimizer::MockSDOptimizer) = optimizer.nconstrs
getnumberofblocks(optimizer::MockSDOptimizer) = length(optimizer.blkdims)
getblockdimension(optimizer::MockSDOptimizer, blk) = optimizer.blkdims[blk]
function init!(optimizer::MockSDOptimizer{T}, blkdims::Vector{Int},
               nconstrs::Integer) where T
    optimizer.nconstrs = nconstrs
    optimizer.blkdims = blkdims
    optimizer.constraint_constants = zeros(T, nconstrs)
    optimizer.objective_coefficients = Tuple{T, Int, Int, Int}[]
    optimizer.constraint_coefficients = map(i -> Tuple{T, Int, Int, Int}[], 1:nconstrs)
end

getconstraintconstant(optimizer::MockSDOptimizer, c) = optimizer.constraint_constants[c]
function setconstraintconstant!(optimizer::MockSDOptimizer, val, c::Integer)
    optimizer.constraint_constants[c] = val
end

getobjectivecoefficients(optimizer::MockSDOptimizer) = optimizer.objective_coefficients
function setobjectivecoefficient!(optimizer::MockSDOptimizer, val, blk::Integer, i::Integer, j::Integer)
    push!(optimizer.objective_coefficients, (val, blk, i, j))
end

getconstraintcoefficients(optimizer::MockSDOptimizer, c) = optimizer.constraint_coefficients[c]
function setconstraintcoefficient!(optimizer::MockSDOptimizer, val, c::Integer,
                                   blk::Integer, i::Integer, j::Integer)
    push!(optimizer.constraint_coefficients[c], (val, blk, i, j))
end

MOI.get(mock::MockSDOptimizer, ::MOI.TerminationStatus) = mock.terminationstatus
function MOI.set(mock::MockSDOptimizer, ::MOI.TerminationStatus,
                 value::MOI.TerminationStatusCode)
    mock.terminationstatus = value
end
MOI.get(mock::MockSDOptimizer, ::MOI.PrimalStatus) = mock.primalstatus
MOI.set(mock::MockSDOptimizer, ::MOI.PrimalStatus, value::MOI.ResultStatusCode) = (mock.primalstatus = value)
MOI.get(mock::MockSDOptimizer, ::MOI.DualStatus) = mock.dualstatus
MOI.set(mock::MockSDOptimizer, ::MOI.DualStatus, value::MOI.ResultStatusCode) = (mock.dualstatus = value)

getX(mock::MockSDOptimizer) = mock.X
getZ(mock::MockSDOptimizer) = mock.Z
gety(mock::MockSDOptimizer) = mock.y
function getprimalobjectivevalue(mock::MockSDOptimizer{T}) where T
    v = zero(T)
    for (α, blk, i, j) in mock.objective_coefficients
        v += α * block(mock.X, blk)[i, j]
        if i != j
            v += α * block(mock.X, blk)[j, i]
        end
    end
    v
end
function getdualobjectivevalue(mock::MockSDOptimizer{T}) where T
    v = zero(T)
    for c in 1:mock.nconstrs
        v += mock.constraint_constants[c] * mock.y[c]
    end
    v
end

function MOI.optimize!(mock::MockSDOptimizer)
    mock.hasprimal = true
    mock.hasdual = true
    mock.optimize!(mock)
end

function MOIU.set_mock_optimize!(mock::MockSDOptimizer, opts::Function...)
    mock.optimize! = MOIU.rec_mock_optimize(mock, opts...)
end
# TODO remove the following methods once it is defined for AbstractMockOptimizer in MOIU
function MOIU.rec_mock_optimize(mock::MockSDOptimizer, opt::Function, opts::Function...)
    # TODO replace mock.optimize! = ... by MOI.set(..., MOIU.MockOptimizeFunction, ...)
    # where MOIU.MockOptimizeFunction is a MockModelAttribute
    (mock::MockSDOptimizer) -> begin
        opt(mock)
        mock.optimize! = MOIU.rec_mock_optimize(mock, opts...)
    end
end
MOIU.rec_mock_optimize(mock::MockSDOptimizer, opt::Function) = opt

# TOD remove the following methods once it is defined for AbstractMockSDOptimizer in MOIU
function MOIU.mock_optimize!(mock::MockSDOptimizer, termstatus::MOI.TerminationStatusCode, primal, dual...)
    MOI.set(mock, MOI.TerminationStatus(), termstatus)
    MOIU.mock_primal!(mock, primal)
    MOIU.mock_dual!(mock, dual...)
end
# Default termination status
function MOIU.mock_optimize!(mock::MockSDOptimizer, primdual...)
    MOIU.mock_optimize!(mock, MOI.OPTIMAL, primdual...)
end
function MOIU.mock_optimize!(mock::MockSDOptimizer, termstatus::MOI.TerminationStatusCode)
    MOI.set(mock, MOI.TerminationStatus(), termstatus)
end

# Primal
function MOIU.mock_primal!(mock::MockSDOptimizer, primstatus::MOI.ResultStatusCode, varprim...)
    MOI.set(mock, MOI.PrimalStatus(), primstatus)
    MOIU.mock_varprimal!(mock, varprim...)
end
# Default primal status
MOIU.mock_primal!(mock::MockSDOptimizer, varprim::Vector) = MOIU.mock_primal!(mock, MOI.FEASIBLE_POINT, varprim)
function MOIU.mock_primal!(mock::MockSDOptimizer)
    # No primal solution
    mock.hasprimal = false
end

# Sets variable primal to varprim
function MOIU.mock_varprimal!(mock::MockSDOptimizer) end
function MOIU.mock_varprimal!(mock::MockSDOptimizer{T}, X::Vector{Matrix{T}}) where T
    mock.X = BlockMatrix{T}(X)
end
to_matrix(X::Matrix) = X
function to_matrix(X::Vector)
    @assert length(X) == 1
    reshape(X, 1, 1)
end
function MOIU.mock_varprimal!(mock::MockSDOptimizer{T}, X::Vector) where T
    MOIU.mock_varprimal!(mock, to_matrix.(X))
end

# Dual
function MOIU.mock_dual!(mock::MockSDOptimizer, dualstatus::MOI.ResultStatusCode, conduals...)
    MOI.set(mock, MOI.DualStatus(), dualstatus)
    MOIU.mock_condual!(mock, conduals...)
end
# Default dual status
function MOIU.mock_dual!(mock::MockSDOptimizer, conduals...)
    status = !mock.hasprimal || MOI.get(mock, MOI.PrimalStatus()) == MOI.INFEASIBLE_POINT ? MOI.INFEASIBILITY_CERTIFICATE : MOI.FEASIBLE_POINT
    MOIU.mock_dual!(mock, status, conduals...)
end
function MOIU.mock_dual!(mock::MockSDOptimizer)
    # No dual solution
    mock.hasdual = false
end

# Sets constraint dual to conduals
function MOIU.mock_condual!(mock::MockSDOptimizer) end
function MOIU.mock_condual!(mock::MockSDOptimizer{T}, y::Vector{T}) where T
    mock.y = y
    # blockdims can be negative for diagonal blocks. We allocate full blocks
    # here. As mock optimizers are used only for testing, we favor simplicity.
    mock.Z = BlockMatrix{T}(map(n -> zeros(T, abs(n), abs(n)), mock.blkdims))
    # FIXME shouldn't Z be defined as the opposite i.e. Z = C - sum y_i A_i >= 0
    # instead of sum y_i A_i - C <= 0 ?
    if mock.dualstatus != MOI.INFEASIBILITY_CERTIFICATE
        # Infeasibility certificate is a ray so we only take the homogeneous part
        # FIXME:check that this is also done IN MOIU.MockOptimizer
        for (α, blk, i, j) in mock.objective_coefficients
            block(mock.Z, blk)[i, j] -= α
            if i != j
                block(mock.Z, blk)[j, i] -= α
            end
        end
    end
    for (c, constraint_coefficients) in enumerate(mock.constraint_coefficients)
        for (α, blk, i, j) in constraint_coefficients
            block(mock.Z, blk)[i, j] += α * y[c]
            if i != j
                block(mock.Z, blk)[j, i] += α * y[c]
            end
        end
    end
end


text:
"""
Authors: Yaroslav Salii, 2020+
         Kara Ignatenko, 2021

Submodule responsible for grouping and aggregating tracts to various levels.

aggregation.jl
2021-08-13 v.0.1 First modular version
2021-09-09 v.0.2 add dummy partByTra to maintain dispatch interface
2021-11-23 v.0.3 move dispatchAggPart here
"""
module Aggregation
using DataFrames
using Statistics: mean

export aggByCty,
       aggBySte,
       aggByAP,
       partByTra,
       partByCty,
       partByAP,
       partBySte,
       dispatchAggPart

 #NB! :Name is a String
 """Return a name-generating function for partitions based on the column names `nms`"""
 function select_mkName(nms::Array{String})
    if ["Ste","Cty","Tra"] ⊆ nms #node's a census tract
         return r-> join([r.Ste,r.Cty,r.Tra],"~")
    elseif ["Ste","Cty"] ⊆ nms #node's a county
         return r-> join([r.Ste,r.Cty],"~")
    elseif  ["Ste"] ⊆ nms #node's a state
         return r-> string(r.Ste) #ensure it's output as a string
     elseif ["IATA_Code"] ⊆ nms #node is Voronoi cell around an AP
         return r -> r.IATA_Code
     else
        error("Nodes header not recognized. Can't generate a :Name without FIPS or IATA_Code.")
     end
 end

"""
By-state PRE-aggregation with dumb Euclidean centroid for geographical coordinates.
Input: a FluTE \$name-tracts.dat, a la `[:Ste,:Cty,:Tra,:Pop,:LAT,:LNG]``
"""
#TODO: put the by--end output into a variable and add post-processing (adding the ID)
function aggBySte(idf::DataFrame;make_names=true)
    gd = groupby(idf,[:Ste]) #group by U.S. State FIPS
    out = combine(gd, :Pop => sum, :LAT => mean, :LNG => mean, renamecols = false)
    if make_names && "Name" ∉ names(out)
        insertcols!(out, :Name => map( select_mkName(names(out)), eachrow(out)))
    end
end

"""
By-county PRE-aggregation with dumb Euclidean centroid for geographical coordinates.
Input: a FluTE \$name-tracts.dat, a la `[:Ste,:Cty,:Tra,:Pop,:LAT,:LNG]`
"""
function aggByCty(idf::DataFrame;make_names=true)
    gd = groupby(idf,[:Ste,:Cty]) #group by U.S. County FIPS, within the same State FIPS
    out = combine(gd, :Pop => sum, :LAT => mean, :LNG => mean, renamecols = false)
    if make_names && "Name" ∉ names(out)
        insertcols!(out, :Name => map( select_mkName(names(out)), eachrow(out)))
    end
end

"""
By-airport (`:IATA_Code`) PRE-aggregation, with dumb Euclidean centroid for geographical coordinates.
Input: a FluTE \$name-tracts.dat WITH dsg AP, a la `[:Name,:IATA_Code,:Pop,:LAT,:LNG]`
"""
function aggByAP(idf::DataFrame;make_names=true)
    gd = groupby(idf,[:IATA_Code]) #group by the tract's dsg AP :IATA_Code
    out = combine(gd, :Pop => sum, :LAT => mean, :LNG => mean, renamecols = false)
    if make_names && "Name" ∉ names(out)
        insertcols!(out, :Name => map( select_mkName(names(out)), eachrow(out)))
    end
    sort!(out,:IATA_Code) #ensure it's sorted by the name/IATA_Code
end

#---AUX::---PARTITION---AND---REVERSE------------------------#
"""Map county names in `pns` to lists of node indices in `ns`."""
function partByCty(ns::DataFrame,pns::DataFrame)
    Dict(x => filter(ri -> ns.Ste[ri]==split(x,"~")[1] &&
                        ns.Cty[ri]==split(x,"~")[2],
                        1:nrow(ns)) for x ∈ pns.Name)
end

#consider a unified approach to partBySte and partByAP;
"""Map state names in `pns` to lists of node indices in `ns`."""
function partBySte(ns::DataFrame,pns::DataFrame)
    Dict(x => filter(ri -> ns.Ste[ri]==x, 1:nrow(ns)) for x ∈ pns.Ste |> unique)
end

"""Map IATA codes in `pns` to lists of node indices in `ns`."""
function partByAP(ns::DataFrame,pns::DataFrame)
    Dict(x => filter(ri -> ns.IATA_Code[ri]==x, 1:nrow(ns)) for x ∈ pns.IATA_Code |> unique)
end

"""
Dummy partition: map each tract name to a 1-vector with its row number.
Arguments `ns` and `pns` are identical, wit `pns` kept to maintain the partByXX interface
"""
function partByTra(ns::DataFrame,pns::DataFrame)
    Dict( ns.Name[ix] => [ix] for ix ∈ 1:nrow(ns))
end

"""
Reverse-explode a `dict` as returned by `partBy...`, i.e. map indices to their
corresponding parition name.
"""
function revexplPart(dict::Dict)
    ( v => k  for k ∈ keys(dict) for v ∈ dict[k]) |> Dict
end

"""
Reverse-explode a `dict` as returned by `partBy...` and map node names to their
corresponding parition name.
ns MUST have [:Name]
"""
function revexplPart(dict::Dict, ns::DataFrame)
    ( ns.Name[v] => k  for k ∈ keys(dict) for v ∈ dict[k]) |> Dict
end

"""
Reverse-explode a `dict` as returned by `partBy...` and map node names to their
corresponding parition *index*.
ns MUST have [:Name]; ix:: Name => partition_index
"""
function revexplPart(dict::Dict, ns::DataFrame,ix::Dict)
    ( ns.Name[v] => ix[k]  for k ∈ keys(dict) for v ∈ dict[k]) |> Dict
end

"Dispatch aggregation and partition functions by aggregation type [tra | cty | ap | ste]"
function dispatchAggPart(agg::String)
    if agg == "tra"
        return (identity,partByTra)
    elseif agg == "cty"
        return (aggByCty,partByCty)
    elseif agg == "ap"
        return (aggByAP,partByAP)
    elseif agg == "ste"
        return (aggBySte,partBySte)
    else
        error("Unknown aggregation requested. Terminating.\n")
    end
end

end


text:
include("custom_determ.jl")
include("distributions.jl")
include("choice_at.jl")
include("call_at.jl")
include("map.jl")
include("unfold.jl")
include("recurse.jl")
include("dist.jl")

max_stars_repo_name: ali-ramadhan/Gen

max_issues_repo_name: ali-ramadhan/Gen

max_forks_repo_name: ali-ramadhan/Gen


text:
# Usage examples:
# dict = YAML.load_file("json/oco2.yaml"; dicttype=OrderedDict{String,Any})
# met = Dataset(oco_met_file)
# ocoData = Dataset(oco_file)
# geometry, ils, meteo, measurement = load_L1(dict,ocoData, met)

function load_L1(dict, L1::NCDataset, met::NCDataset)
    # Dicts with mostly NCVars but also some basic stuff:
    geometry    = Dict{String, Any}()
    ils         = Dict{String, Any}()
    meteo       = Dict{String, Any}()
    measurement = Dict{String, Any}()

    # Reading in Geometries
    loadEntries!(L1, dict, geometry, "Geometry")

    # Reading in Meterorology
    loadEntries!(met, dict, meteo, "Meteorology")

    # Reading in Meausurements
    loadEntries!(L1, dict, measurement, "Measurement")

    # Reading in Meausurements
    loadEntries!(L1, dict, ils, "ILS")

    return L1_OCO(geometry, ils, meteo, measurement)
end

function load_L1(dict::String, L1::String, met::String)
    ocoData = Dataset(L1);
    metData = Dataset(met);
    # Load dictionary:
    dictOCO2 = YAML.load_file(dict; dicttype=OrderedDict{String,Any});

    # Load L1 file (could just use filenames here as well later on)
    return load_L1(dictOCO2,ocoData, metData)
end

function loadEntries!(NC, dict, out_dict, name)
    if haskey(dict,name)
        for (key, value) in dict[name]
            if typeof(value) == String
                out_dict[key] = getNC_var(NC, value)
            else
                out_dict[key] = value
            end
        end
    end
end

function getMeasurement(oco::L1_OCO, bands::Tuple, indices::Tuple, GeoInd)
    @assert length(indices) == length(bands) "Length of bands and indices has to be identical"
    n = length(indices)
    # First band:
    band = oco.measurement["bands"][bands[1]]
    # Extended Dimensions (irrespective of sounding)
    extended_dims = [GeoInd[1],bands[1]];
    ind  = indices[1]
    bandIndices = (ind .- ind[1] .+ 1,)
    rad = oco.measurement[band][ind,GeoInd...]
    FT = typeof(rad[1])
    dispPoly = Polynomial(view(oco.ils["dispersion"], :, extended_dims...))
    ν = FT.(dispPoly.(indices[1]))
    # First ILS
    # First hard-coded:
    @show FT, typeof(ν)

    # Hard coded for now, needs to be changed later:
    grid_x = FT(-0.35e-3):FT(0.001*1e-3):FT(0.35e-3)
    ils_pixel   = prepare_ils_table(grid_x, oco.ils["ils_response"][:], oco.ils["ils_grid"][:],extended_dims)
    oco2_kernels = (VariableKernelInstrument(ils_pixel, ν, collect(ind .-1)),)
    # Concatenate rest (if applicable)
    for i=2:n
        band = oco.measurement["bands"][bands[i]]
        ind  = indices[i]
        bandIndices = (bandIndices..., ind .-ind[1] .+ 1 .+ bandIndices[end][end])
        rad = [rad; oco.measurement[band][ind,GeoInd...]]
        extended_dims = [GeoInd[1],bands[i]];
        dispPoly = Polynomial(view(oco.ils["dispersion"], :, extended_dims...))
        ν = [ν; FT.(dispPoly.(indices[i]))]
        # ILS kernels:
        grid_x = FT(-0.35e-3):FT(0.001*1e-3):FT(0.35e-3)
        ils_pixel   = prepare_ils_table(grid_x, oco.ils["ils_response"][:], oco.ils["ils_grid"][:],extended_dims)
        @show ind
        oco2_kernels = (oco2_kernels..., VariableKernelInstrument(ils_pixel, FT.(dispPoly.(indices[i])), collect(ind .-1)))
        @show VariableKernelInstrument(ils_pixel, FT.(dispPoly.(indices[i])), collect(ind .-1)).ν_out
    end

    #### Meteo stuff  ###
    p_half = reverse(oco.meteo["ak"] + oco.meteo["bk"] * oco.meteo["p_surf"][GeoInd...]/100)
    T   =    oco.meteo["T"][:,GeoInd...]
    q   =    oco.meteo["q"][:,GeoInd...]
    lat = oco.geometry["lat"][GeoInd...]
    lon = oco.geometry["lon"][GeoInd...]
    sza = oco.geometry["sza"][GeoInd...]
    ϕ   = oco.geometry["ϕ"][GeoInd...]
    vza = oco.geometry["vza"][GeoInd...]
    azi = oco.geometry["azi"][GeoInd...]

    # From OCO-2 ATBD Page 53 https://docserver.gesdisc.eosdis.nasa.gov/public/project/OCO/OCO_L1B_ATBD.pdf
    StokesCoef = [FT(0.5),FT(cosd(2ϕ)/2), FT(sind(2ϕ)/2), FT(0)]

    return MeasurementOCO(
        ν,
        bandIndices,
        rad,
        oco.measurement["radiance_o2"].attrib["Units"],
        "μm",
        lat,
        lon,
        vza,
        sza,
        azi,
        FT((1 + 0.0053024*sind(lat)^2 - 0.0000058*sin(2lat)^2) * 9.780318), # WELMEC formula
        ϕ,
        StokesCoef,
        FT.(p_half),
        FT.((p_half[1:end-1] + p_half[2:end])/2),
        T,
        q,
        oco2_kernels
    )
end


text:
using Test

include("../scripts/JuliaDSP.jl")
using .JuliaDSP

@testset "Convolution" begin
    a1 = [1, 5, 0, 3]
    a2 = [5, -3, -2, 5]
    a3 = [100, 200]

    @testset "Arrays of length 1" begin
        @test JuliaDSP.convolve([1], [1]) == reshape([1], 1, 1)
        @test JuliaDSP.convolve([100], [13]) == reshape([1300], 1, 1)
        @test JuliaDSP.convolve([90], [100]) == reshape([9000], 1, 1)
    end

    @testset "Commutative property" begin
        @test JuliaDSP.convolve(a1, a2) == [5 22 -17 10 16 -6 15]
        @test JuliaDSP.convolve(a2, a1) == [5 22 -17 10 16 -6 15]
    end

    @testset "Arrays are different sizes" begin
        @test JuliaDSP.convolve([1, 8, 6], [78]) == [78 624 468]
        @test JuliaDSP.convolve([1, 2], [1, 2, 3]) == [1 4 7 6]
        @test JuliaDSP.convolve(a1, a3) == [100  700 1000 300 600]
    end
end

@testset "Matrix multiplication" begin
    m1 = [1 2 3; 4 5 6]
    m2 = [7 8; 9 10; 11 12]
    m3 = [11 12; 4 5]
    m4 = [12 13; 1000 1000]

    @testset "0 case" begin
        @test JuliaDSP.matrixMultiply([], []) == reshape([], 0, 0)
        @test JuliaDSP.matrixMultiply([88], []) == reshape([], 0, 0)
        @test JuliaDSP.matrixMultiply([], [4]) == reshape([], 0, 0)
    end

    @testset "1 case" begin
        @test JuliaDSP.matrixMultiply([1], [3]) == reshape([3], 1, 1)
        @test JuliaDSP.matrixMultiply([45], [3]) == reshape([135], 1, 1)
        @test JuliaDSP.matrixMultiply([100], [100]) == reshape([10000], 1, 1)

    end

    @testset "Matrices are both square and the same size" begin
        @test JuliaDSP.matrixMultiply(m3, m4) == [12132 12143; 5048 5052]
    end

    @testset "Matrices are different sizes" begin
        @test JuliaDSP.matrixMultiply(m1, m2) == [58 64; 139 154]
    end
end


text:
using Test
using MathOptInterface
const MOI = MathOptInterface
const MOIU = MOI.Utilities

@testset "EmptyVector{$T}" for T in [Int, Float64]
    v = MOIU.EmptyVector{T}()
    @test size(v) == (0,)
    @test length(v) == 0
    @test isempty(v)
    @test eltype(v) == T
    @test iterate(v) === nothing
    c = collect(v)
    @test c isa Vector{T}
    @test isempty(c)
end

@testset "LazyMap{$T}" for T in [Int, Float64]
    v = MOIU.LazyMap{T}(x -> x^2, [2, 3])
    @test size(v) == (2,)
    @test length(v) == 2
    @test !isempty(v)
    @test eltype(v) == T
    c = collect(v)
    @test c isa Vector{T}
    @test c == [4, 9]
end


text:
#
# BufferedIO
#

const Plain = Union{Int16,Int32,Int64,Int128,UInt16,UInt32,UInt64,UInt128,Float16,Float32,
                    Float64}
const PlainType = Union{Type{Int16},Type{Int32},Type{Int64},Type{Int128},Type{UInt16},
                        Type{UInt32},Type{UInt64},Type{UInt128},Type{Float16},
                        Type{Float32},Type{Float64}}
const DEFAULT_BUFFER_SIZE = 1024

struct BufferedWriter <: IO
    f::IOStream
    buffer::Vector{UInt8}
    file_position::Int64
    position::Base.RefValue{Int}
end

function BufferedWriter(io::IOStream, buffer_size::Int)
    pos = position(io)
    skip(io, buffer_size)
    BufferedWriter(io, Vector{UInt8}(undef, buffer_size), pos, Ref{Int}(0))
end
Base.show(io::IO, ::BufferedWriter) = print(io, "BufferedWriter")

function finish!(io::BufferedWriter)
    f = io.f
    buffer = io.buffer
    io.position[] == length(buffer) ||
        error("buffer not written to end; position is $(io.position[]) but length is $(length(buffer))")
    seek(f, io.file_position)
    write(f, buffer)
    io.position[] = 0
    nothing
end

@inline function _write(io::BufferedWriter, x)
    position = io.position[]
    buffer = io.buffer
    n = sizeof(x)
    n + position <= length(buffer) || throw(EOFError())
    io.position[] = position + n
    unsafe_store!(Ptr{typeof(x)}(pointer(buffer, position+1)), x)
    # Base.show_backtrace(STDOUT, backtrace())
    # gc()
    return n
end
@inline Base.write(io::BufferedWriter, x::UInt8) = _write(io, x)
@inline Base.write(io::BufferedWriter, x::Int8) = _write(io, x)
@inline Base.write(io::BufferedWriter, x::Plain)  = _write(io, x)

function Base.unsafe_write(io::BufferedWriter, x::Ptr{UInt8}, n::UInt64)
    buffer = io.buffer
    position = io.position[]
    n + position <= length(buffer) || throw(EOFError())
    unsafe_copyto!(pointer(buffer, position+1), x, n)
    io.position[] = position + n
    # Base.show_backtrace(STDOUT, backtrace())
    # gc()
    return n
end

Base.position(io::BufferedWriter) = io.file_position + io.position[]

struct BufferedReader <: IO
    f::IOStream
    buffer::Vector{UInt8}
    file_position::Int64
    position::Base.RefValue{Int}
end

BufferedReader(io::IOStream) =
    BufferedReader(io, Vector{UInt8}(), position(io), Ref{Int}(0))
Base.show(io::IO, ::BufferedReader) = print(io, "BufferedReader")

function readmore!(io::BufferedReader, n::Int)
    f = io.f
    amount = max(bytesavailable(f), n)
    buffer = io.buffer
    oldlen = length(buffer)
    resize!(buffer, oldlen + amount)
    unsafe_read(f, pointer(buffer, oldlen+1), amount)
end

@inline function _read(io::BufferedReader, T::DataType)
    position = io.position[]
    buffer = io.buffer
    if length(buffer) - position < sizeof(T)
        readmore!(io, sizeof(T))
    end
    io.position[] = position + sizeof(T)
    unsafe_load(Ptr{T}(pointer(buffer, position+1)))
end
@inline Base.read(io::BufferedReader, T::Type{UInt8}) = _read(io, T)
@inline Base.read(io::BufferedReader, T::Type{Int8}) = _read(io, T)
@inline Base.read(io::BufferedReader, T::PlainType) = _read(io, T)

function Base.read(io::BufferedReader, ::Type{T}, n::Int) where T
    position = io.position[]
    buffer = io.buffer
    m = sizeof(T) * n
    if length(buffer) - position < m
        readmore!(io, m)
    end
    io.position[] = position + m
    arr = Vector{T}(undef, n)
    unsafe_copyto!(pointer(arr), Ptr{T}(pointer(buffer, position+1)), n)
    arr
end
Base.read(io::BufferedReader, ::Type{T}, n::Integer) where {T} =
    read(io, T, Int(n))

Base.position(io::BufferedReader) = io.file_position + io.position[]

function adjust_position!(io::BufferedReader, position::Integer)
    if position < 0
        throw(ArgumentError("cannot seek before start of buffer"))
    elseif position > length(io.buffer)
        readmore!(io, position - length(io.buffer))
    end
    io.position[] = position
end

Base.seek(io::BufferedReader, offset::Integer) =
    adjust_position!(io, offset - io.file_position)

Base.skip(io::BufferedReader, offset::Integer) =
    adjust_position!(io, io.position[] + offset)

finish!(io::BufferedReader) =
    seek(io.f, io.file_position + io.position[])

function truncate_and_close(io::IOStream, endpos::Integer)
    truncate(io, endpos)
    close(io)
end

# We sometimes need to compute checksums. We do this by first calling begin_checksum when
# starting to handle whatever needs checksumming, and calling end_checksum afterwards. Note
# that we never compute nested checksums, but we may compute multiple checksums
# simultaneously.

function begin_checksum_read(io::IOStream)
    BufferedReader(io)
end
function begin_checksum_write(io::IOStream, sz::Integer)
    BufferedWriter(io, sz)
end
function end_checksum(io::Union{BufferedReader,BufferedWriter})
    ret = Lookup3.hash(io.buffer, 1, io.position[])
    finish!(io)
    ret
end


text:
module GenomicStore

export get_track_interval
export get_chr_sizes_dict
export save_point_start_track
export save_interval_track
export save_bedgraph_track
export save_bed_track
export store_info_str
export store_info_dict
include("genomic-store.jl")

export getdb
export delete_track!
include("crud.jl")

export store_samples
include("store.jl")
end # module


text:
using SpecialFunctions, Documenter

makedocs(
    modules     = [SpecialFunctions],
    sitename    = "SpecialFunctions.jl",
    authors     = "Jeff Bezanson, Stefan Karpinski, Viral B. Shah, et al.",
    pages       = [
        "Home"      => "index.md",
        "Overview"  => "functions_overview.md",
        "List"      => "functions_list.md",
    ],
)

deploydocs(
    repo    = "github.com/JuliaMath/SpecialFunctions.jl.git",
    target  = "build",
)


text:
### A Pluto.jl notebook ###
# v0.17.3

using Markdown
using InteractiveUtils

# ╔═╡ f2224d36-7972-4e21-badc-3a401bf6f0bb
pwd()

# ╔═╡ 7d6d3864-3bff-414d-b077-ab59673f6d25
@__DIR__

# ╔═╡ 866f4346-58fa-11ec-2d52-1fe8cc15ffbd
begin
    import Pkg
    # activate the shared project environment
    Pkg.activate(Base.current_project())
    # instantiate, i.e. make sure that all packages are downloaded
    Pkg.instantiate()

    using Plots, PlutoUI
	using StanSample, StatisticalRethinking
end

# ╔═╡ 0c9c514e-598c-418f-90e5-4a74ffb5bdc0
begin
    import Pkg
    # activate a temporary environment
    Pkg.activate(mktempdir())
    Pkg.add([
        Pkg.PackageSpec(name="Plots", version="1"),
        Pkg.PackageSpec(name="PlutoUI", version="0.7"),
    ])
    using Plots, PlutoUI, LinearAlgebra
end

# ╔═╡ Cell order:
# ╠═866f4346-58fa-11ec-2d52-1fe8cc15ffbd
# ╠═f2224d36-7972-4e21-badc-3a401bf6f0bb
# ╠═7d6d3864-3bff-414d-b077-ab59673f6d25
# ╠═0c9c514e-598c-418f-90e5-4a74ffb5bdc0


text:
function get_obs(x::Attributes, visited, obs=Set())
    if x in visited; return; else; push!(visited, x); end
    union!(obs, values(x))
    return obs
end

function get_obs(x::Union{AbstractVector, Tuple}, visited, obs=Set())
    if x in visited; return; else; push!(visited, x); end
    isempty(x) && return
    for e in x
        get_obs(e, visited, obs)
    end
    return obs
end

get_obs(::Nothing, visited, obs=Set()) = obs
get_obs(::DataType, visited, obs=Set()) = obs
get_obs(::Method, visited, obs=Set()) = obs

function get_obs(x, visited, obs=Set())
    if x in visited; return; else; push!(visited, x); end
    for f in fieldnames(typeof(x))
        if isdefined(x, f)
            get_obs(getfield(x, f), visited, obs)
        end
    end
end

function get_obs(x::Union{AbstractDict, NamedTuple}, visited, obs=Set())
    if x in visited; return; else; push!(visited, x); end
    for (k, v) in pairs(x)
        get_obs(v, visited, obs)
    end
    return obs
end

function get_obs(x::Union{Observables.AbstractObservable, Observables.ObserverFunction}, visited, obs=Set())
    push!(obs, x)
    return obs
end

function get_obs(x::Axis, visited, obs=Set())
    if x in visited; return; else; push!(visited, x); end

    get_obs(x.attributes, visited, obs)
    get_obs(x.layoutobservables, visited, obs)
    get_obs(x.scene, visited, obs)
    get_obs(x.finallimits, visited, obs)
    get_obs(x.scrollevents, visited, obs)
    get_obs(x.keysevents, visited, obs)
    get_obs(x.mouseeventhandle.obs, visited, obs)
    return obs
end

function precompile_obs(x)
    obsies = get_obs(x, Base.IdSet())
    for obs in obsies
        Base.precompile(obs)
    end
end

function _precompile_()
    ccall(:jl_generating_output, Cint, ()) == 1 || return nothing
    precompile(Makie.backend_display, (CairoBackend, Scene))
    activate!()
    f, ax1, pl = scatter(1:4)
    f, ax2, pl = lines(1:4)
    Makie.colorbuffer(ax1.scene)
    Makie.colorbuffer(ax2.scene)
    precompile_obs(ax1)
    precompile_obs(ax2)
    return
end


text:
export  Viscoplasticity1D

mutable struct Viscoplasticity1D
    ρ::Float64 # density
    E::Float64 # Young's modulus
    # hardening parameter, yield function = |σ - q| - (σY + Kα)
    K::Float64
    B::Float64
    η::Float64
    σY::Float64
    α::Float64
    α_::Float64 # α to be updated in `commitHistory`
    σ0::Float64 # stress at last time step
    σ0_::Float64 # σ0 to be updated in `commitHistory
    q::Float64 # stress at last time step
    q_::Float64 # σ0 to be updated in `commitHistory`
end

function Viscoplasticity1D(prop::Dict{String, Any})
    ρ = prop["rho"];
    E = prop["E"]; B = prop["B"];
    K = prop["K"]; σY = prop["sigmaY"]
    η = prop["eta"]
    α = 0.0; α_ = 0.0
    σ0 = 0.0; σ0_ = 0.0
    q = 0.0; q_ = 0.0
    Viscoplasticity1D(ρ, E, K, B, η, σY, α, α_, σ0, σ0_, q, q_)
end

@doc """
    For Viscoplasticity material
    :param hysteresis_variables: [eps_vp, alpha, q], plastic strain, internal hardening variable, and back stress
    The yield condition is
        f = |sigma - q| - (σY + alpha * K)
    here K is the plastic modulus, , σY is the flow stress
    D_eps_vp = <phi(f)>/eta * df/dsigma
    D_alpha = |D_eps_p|
    D_q     = B * D_eps_vp

    here <phi(f)> = (f+|f|)/2

""" ->
function getStress(self::Viscoplasticity1D,  strain::Float64,  Dstrain::Float64, Δt::Float64)

    local dΔσdΔε
    ε = strain; ε0 = Dstrain
    σ0 = self.σ0;  α0 = self.α;  q0 = self.q
    E = self.E;    K = self.K;   B = self.B;
    η = self.η; σY = self.σY

    #trial stress
    Δγ = 0.0
    σ = σ0 + E*(ε-ε0)
    α = α0 + abs(Δγ)
    q = q0 + B*Δγ
    ξ = σ - q

    r2 = abs(ξ) - (σY + K*α)
    if r2 <= 0
        σ = σ0 + E*(ε-ε0)
        dΔσdΔε = E

    else

        Δγ = r2 * Δt / (η + Δt * (E + H + B))
        q += B * Δγ * sign(ξ)
        α += Δγ

        σ -= Δγ * E * sign(ξ)
        dΔσdΔε = E * (η + Δt*(B + K)) / (η + Δt*(B + E + K))
    end

    # #@show Δγ
    self.α_  = self.α + Δγ
    self.σ0_ = σ
    self.q_  = q

    return σ, dΔσdΔε
end

function getTangent(self::Viscoplasticity1D)
    error("Not implemented")
end

function commitHistory(self::Viscoplasticity1D)
    self.α = self.α_
    self.σ0 = self.σ0_
    self.q = self.q_
end


text:
using GeometricFlux
using GeometricFlux.Datasets
using CUDA
using Flux
using Flux: @functor
using FillArrays
using GraphSignals
using Graphs
using LinearAlgebra
using NNlib, NNlibCUDA
using SparseArrays: SparseMatrixCSC
using Statistics: mean
using Zygote
using Test

cuda_tests = [
    "cuda/conv",
    # "cuda/msgpass",
]

tests = [
    # "layers/gn",
    # "layers/msgpass",
    "layers/conv",
    "layers/pool",
    "layers/misc",
    "models",
]

if CUDA.functional()
    append!(tests, cuda_tests)
else
    @warn "CUDA unavailable, not testing GPU support"
end

@testset "GeometricFlux" begin
    for t in tests
        include("$(t).jl")
    end
end


text:
"""
Calculate Dynamic Programming using Bayesian approach
    dp_bayes(q1, q1L, q2L, times, cut)
    :param q1: vector (N)
    :param q1L: vector (q1L)
    :param q2L: vector (q2L)
    :param times: MCMC parameter number of simulations
    :param cut: MCMC parameter number of cuts

    :return MatchIn2: warping function
    :return NDist: minimal distance
    :return q2LL: re-parameterized q2L
"""
function dp_bayes(q1, q1L, q2L, times, cut)

    colnum = length(q1L);
    rownum = div(colnum,times);
    q2LLlen = (colnum-1)*times+1;
    q2LL = zeros(q2LLlen);
    tempspan1 = collect(1:q2LLlen)-1;
    temp_span2 = float(tempspan1);
    timesf = float(times);
    q2LL_time = (temp_span2*(1/timesf))+1;
    q2L_time1 = collect(1:colnum);
    q2L_time2 = float(q2L_time1);
    q2LL_i = interpolate((collect(1:length(q2L)),), q2L, Gridded(Linear()))
    q2LL = q2LL_i[q2LL_time];
    ID = zeros(rownum+1, colnum+1);
    S = zeros(ID);
    jend = 0; start = 0; end1 = 0; k = 0; index =0;
    interx = zeros(Int64, times-1);
    intery2 = zeros(times-1);
    intery = zeros(Int64, times-1);
    span1 = (1:(times-1));
    span2 = float(span1);
    q1x = Array(Float64, times-1);
    q2y = Array(Float64, times-1);

    for i in 1:(rownum-1)
        if (i*cut+1) < colnum
            jend = i*cut+1;
        else
            jend = colnum;
        end

        for j in i:jend-1
            start = max(i-1,j-cut);
            end1 = min(j-1, cut*(i-1));
            n = start-1+collect(1:(end1-start+1));
            k = length(n);
            interx = times*(i-1)+span1;
            Energy = Array(Float64, k);

            for m in 1:k
                jf = j;
                divisor = (jf-n[m])/times;
                intery2 = divisor.*span2+n[m];
                for l in 1:(times-1)
                    intery[l] = round(times*intery2[l]);
                    q1x[l] = q1L[interx[l]+1];
                    q2y[l] = q2LL[intery[l]+1];
                end
                Energy[m] = S[i,n[m]+1]+(q1[i]-sqrt(divisor)*q2L[n[m]+1])^2 +
                    (norm(q1x-sqrt(divisor)*q2y))^2;
            end
            min_val, index = findmin(Energy);
            loc = n[index];
            S[i+1,j+1] = min_val;
            ID[i+1,j+1] = loc;
        end
    end

    i = rownum;
    j = colnum;
    start = max(i-1,j-cut);
    end1 = min(j-1,cut*(i-1));
    n = start-1+collect(1:(end1-start+1));
    k = length(n);
    interx = times*(i-1)+span1;
    span2 = float(span1);
    Energy = Array(Float64,k);

    for m in 1:k
        jf = j;
        divisor = (jf-n[m])/times;
        intery2 = divisor.*span2+n[m];
        for l in 1:times-1
            intery[l] = round(times*intery2[l]);
            q1x[l] = q1L[interx[l]+1];
            if intery[l]+1 < q2LLlen
                q2y[l] = q2LL[intery[l]+1];
            else
                q2y[l] = 0.0;
            end
        end
        Energy[m] = S[i,n[m]+1]+(q1[i]-sqrt(divisor)*q2L[n[m]+1])^2 +
                    (norm(q1x-sqrt(divisor)*q2y))^2;
    end

    min_val, index = findmin(Energy);
    loc = n[index];
    S[i+1,j+1] = min_val;
    ID[i+1,j+1] = loc;

    path = Array(Int64,rownum);
    count = round(Integer,ID[i+1,j+1]);
    oldcount = 0;
    path[i] = count;

    while count>1
        i -= 1;
        oldcount = count;
        count = round(Integer, ID[i+1,oldcount+1]);
        path[i] = count;
    end

    MatchIn2 = path+1;
    NDist = S[rownum+1,colnum+1]/colnum;

    return MatchIn2, NDist, q2LL
end

"""
Calculate mean function using Bayesian Dynamic Programming
    DP_mean(f, times=5)
    :param f: array (M,N) of N functions
    :param times: MCMC parameter number of times to split

    Returns Dict containing
    :return distmfamily: distance matrix
    :return match_matrix: match_matrix
    :return position: position
    :return mu_5: mean function
    :return rtmatrix: rtmatrix
    :return sumdist: sum distance
    :return qt_fitted: qt fitted matrix
    :return esimator: estimator
    :return estimator2: estimator2
    :return regcurve: registered curves
"""
function DP_mean(f, times=5)
    cut = 5*times;
    iter = 20;
    timet = LinRange(0,1,size(f,1));
    group = collect(1:size(f, 2));
    m = length(timet) - 1;
    n = size(f, 2);
    qt_matrix = zeros(m, n);

    for j = 1:n
        qt_matrix[:, j] = sign(diff(f[:,j])).*sqrt(abs(diff(f[:,j]))./diff(timet));
        rescale = sqrt(m/sum(qt_matrix[:, j].^2));
        qt_matrix[:, j] = rescale.*qt_matrix[:, j];
    end
    row = collect(1:times:m);
    qt_fitted_matrix = zeros(m, n);

    search = Array(Float64, n);
    meanq = mean(qt_matrix, 2);
    for j = 1:n
        search[j] = Enorm(vec(qt_matrix[:, j] - meanq));
    end
    position = argmin(search);

    mu_5 = qt_matrix[:, position];
    mu_curve = zeros(iter, m+1);
    dist_matrix = zeros(iter+1, n);
    for j = 1:n
        dist_matrix[1, j] = (Enorm(vec(mu_5 - qt_matrix[:, j])))^2/m;
    end
    rtmatrix = zeros(m+1, n);
    match_matrix = zeros(length(row)+1, n);

    i = 1;
    diffdist = 1000.;
    while ((i<iter) & (diffdist > 0.001))
        for j = 1:n
            MatchIn2, NDist, q2LL = dp_bayes(mu_5[row], mu_5, qt_matrix[:, j],
                                             times, cut);
            match = [MatchIn2; m+1];
            f_i = interpolate(([row; m+1],), float(match), Gridded(Linear()))
            idy = round(Integer, f_i[1:m]);
            idy[idy .> m] = m;
            scale1 = sqrt(diff(match)*(1/times));
            scalevec = kron(scale1, ones(times));
            tmpidy = copy(idy) - 1;
            extended_idy = tmpidy[1:m]*times+1;
            q_fitted = scalevec.*q2LL[extended_idy];
            qt_fitted_matrix[:, j] = q_fitted;
            dist_matrix[i+1, j] = NDist;
            rtmatrix[:, j] = [idy; m+1];
            match_matrix[:, j] = match;
        end
        diffdist = abs(sum(dist_matrix[i+1, :]) - sum(dist_matrix[i, :]));
        mu_5 = mean(qt_fitted_matrix, 2);
        rescale = sqrt(m/sum(mu_5.^2));
        mu_5 *= rescale;
        i += 1;
    end

    estimator2 = copy(mu_5);
    invidy, revscalevec = findkarcherinv(match_matrix, times);
    invidy = invidy[1:m];
    invidy[invidy .>= m] = m;
    f_i = interpolate((collect(1:m),), vec(mu_5), Gridded(Linear()))
    mu_5 = revscalevec.*f_i[invidy];
    rescale = sqrt(m/sum(mu_5.^2));
    estimator = rescale.*copy(mu_5);
    reg_curve = zeros(m+1, n);
    for j = 1:n
        f_i = interpolate((collect(0:m),), f[:, j], Gridded(Linear()))
        tmp = f_i[LinRange(0, m, times*(m+1)-1)];
        reg_curve[:, j] = tmp[round(Integer, (rtmatrix[:,j] - 1)*times+1)];
    end
    crossmean = mean(reg_curve, 2);
    sumdist = sum(reg_curve, 2);

    out = Dict("distfamily" => dist_matrix, "match_matrix" => match_matrix,
               "position" => position, "mu_5" => mu_5,
               "rtmatrix" => rtmatrix, "sumdist" => sumdist,
               "qt_fitted" => qt_fitted_matrix, "estimator" => estimator,
               "estimator2" => estimator2, "regcurve" => reg_curve);
    return out
end


text:
using Test
using ExprParsers

# Standard Parsers
# ================

@testset "Block" begin
  parser = EP.Block(EP.anysymbol, :b, EP.Expr(:tuple, EP.anything))
  @test_throws EP.ParseError parse_expr(parser, quote
    5
    b
    (1,4,a)
  end)

  test_closure(parser, quote
    a
    b
    (1,4,a)
  end)

  parsed = parse_expr(EP.Block(), quote
    f(a) = A
    3
  end)
  parsed_exprs = filter(x -> !isa(x, LineNumberNode), parsed.exprs)
  @test parsed_exprs[1].head == :(=)
  @test parsed_exprs[1].args[1] == :(f(a))
  @test parsed_exprs[2] == 3
end

@testset "Arg" begin
  parser = EP.Arg()
  parsed = parse_expr(parser, Expr(:kw, :a, 4))
  @test parsed.name == :a
  @test parsed.type == Any
  @test parsed.default == 4
  test_closure(parser, Expr(:kw, :a, 4))

  @test_throws EP.ParseError parse_expr(parser, :(a = 5))

  expr = :(f(a::Int, b=5; c, d=6) = a + b + c + d)
  expr.args[1].args[3:end][2].head
  a, b = [parse_expr(parser, x) for x in expr.args[1].args[3:end]]
  c, d = [parse_expr(parser, x) for x in expr.args[1].args[2].args]
  @test a.name == :a
  @test a.default == EP.nodefault
  @test a.type == :Int

  @test b.name == :b
  @test b.default == 5
  @test b.type == Any

  @test c.name == :c
  @test c.default == EP.nodefault
  @test c.type == Any

  @test d.name == :d
  @test d.default == 6
  @test d.type == Any

  @testset "Vararg" begin
    expr = :(f(args...; kwargs...) = (args, kwargs))
    args_expr = expr.args[1].args[3]
    kwargs_expr = expr.args[1].args[2].args[1]
    args = parse_expr(parser, args_expr)
    kwargs = parse_expr(parser, kwargs_expr)
    @test args.type === Vararg
    @test args.name === :args
    @test args.default === EP.nodefault
    @test kwargs.type === Vararg
    @test kwargs.name === :kwargs
    @test kwargs.default === EP.nodefault

    test_closure(parser, args_expr)
    test_closure(parser, kwargs_expr)
  end

  parsed = EP.Arg_Parsed(name = :a)
  @test parse_expr(EP.Arg(), to_expr(parsed)) == parsed
end

@testset "Module" begin
  parser = EP.Module()
  expr = :(module A
    something
    3
  end)
  parsed = parse_expr(parser, expr)
  @test parsed.no_bare == true
  @test parsed.name == :A
  @test filter(x -> !isa(x, LineNumberNode), parse_expr(EP.Block(), parsed.body).exprs) == [:something, 3]

  test_closure(parser, expr)

  @test_throws EP.ParseError parse_expr(parser, :(a::Int))
  @test_throws EP.ParseError parse_expr(parser, quote
      module A
    end
  end)
end

@testset "Assignment" begin
  parser = EP.Assignment()
  parsed = parse_expr(parser, :(a::Int = 4))
  @test parsed.left == :(a::Int)
  @test parsed.right == 4

  test_closure(parser, :(a::Int = 4))

  @test_throws EP.ParseError parse_expr(parser, :(a::Int))

  parsed = EP.Assignment_Parsed(left = :left, right = :right)
  @test parse_expr(EP.Assignment(), to_expr(parsed)) == parsed
end

@testset "Call" begin
  parser = EP.Call()
  parsed = parse_expr(parser, :(f{A, B}(a, b, c...; d...)))
  @test parsed.name == :f
  @test parsed.curlies == [:A, :B]
  @test parsed.args == [:a, :b, :(c...)]
  @test parsed.kwargs == [:(d...)]

  test_closure(parser, :(f{A, B}(a, b, c...; d...)))

  @test_throws EP.ParseError parse_expr(parser, :(a = 4))

  parsed = EP.Call_Parsed(name = :f)
  @test parse_expr(EP.Call(), to_expr(parsed)) == parsed
end

@testset "Expr" begin
  expr = Expr(:head, :a1, 2, [3,4])
  parsed_expr = EP.Expr_Parsed(:head, [:a1, 2, [3,4]])
  @test parse_expr(EP.Expr(), expr) == parsed_expr
  @test_throws EP.ParseError parse_expr(EP.Expr(head = EP.AnyOf(:hi, :ho)), expr)
  @test parse_expr(EP.Expr(head = EP.AnyOf(:hi, :ho, :head)), expr) == parsed_expr
  @test_throws EP.ParseError parse_expr(EP.Expr(args = [1,2]), expr)
  @test parse_expr(EP.Expr(args = [:a1, 2, [3,4]]), expr) == parsed_expr
  @test parse_expr(EP.Expr(args = [:a1, 2, EP.anything]), expr) == parsed_expr

  # All Parsers should be Closure operators
  test_closure(EP.Expr(), expr)

  # second constructor
  parser = EP.Expr(quote
    a = 4
    b = 5
  end)
  test_closure(parser, quote
    a = 4
    b = 5
  end)

  # Expr should also support nested Parsers
  parser = EP.Expr(quote
    a = $(EP.Isa(Int))
    b = $(EP.anysymbol)
  end)
  parsed = parse_expr(parser, quote
    a = 4
    b = a
  end)

  test_closure(parser, quote
    a = 4
    b = a
  end)

  @test parsed.args.exprs[4].args[2] isa Base.Symbol

  @test_throws EP.ParseError parse_expr(parser, quote
    a = 4.0
    b = a
  end)
end

@testset "Signature" begin
  parser = EP.Signature()

  expr = :(
    f(a::String, b::Int...; c = 5, d...)
  )
  parsed = parse_expr(parser, expr)
  @test parsed.name == :f
  @test parsed.args == [:(a::String), :(b::Int...)]
  @test parsed.kwargs == [Expr(:kw, :c, 5), :(d...)]
  @test parsed.curlies == []
  @test parsed.wheres == []
  test_closure(parser, expr)

  expr = :(
    (a::String, b::Int...; c = 5, d...) where B where A
  )
  parsed = parse_expr(parser, expr)
  @test parsed.name == nothing
  @test parsed.args == [:(a::String), :(b::Int...)]
  @test parsed.kwargs == [Expr(:kw, :c, 5), :(d...)]
  @test parsed.curlies == []
  @test parsed.wheres == [:A, :B]
  test_closure(parser, expr)

  expr = :(
    f{A}(a::String, b::Int...; c = 5, d...) where {A, B}
  )
  parsed = parse_expr(parser, expr)
  @test parsed.name == :f
  @test parsed.args == [:(a::String), :(b::Int...)]
  @test parsed.kwargs == [Expr(:kw, :c, 5), :(d...)]
  @test parsed.curlies == [:A]
  @test parsed.wheres == [:A, :B]
  test_closure(parser, expr)

  @test_throws EP.ParseError parse_expr(parser, :(
    f = 5
  ))

  parsed = EP.Signature_Parsed()
  @test parse_expr(EP.Signature(), to_expr(parsed)) == parsed
end

@testset "Function" begin
  parser = EP.Function()

  expr = :(
    function f(a::String, b::Int...; c = 5, d...)
    end
  )
  parsed = parse_expr(parser, expr)
  @test parsed.name == :f
  @test parsed.args == [:(a::String), :(b::Int...)]
  @test parsed.kwargs == [Expr(:kw, :c, 5), :(d...)]
  @test parsed.curlies == []
  @test parsed.wheres == []
  test_closure(parser, expr)

  expr = :(
    function (a::String, b::Int...; c = 5, d...) where B where A
    end
  )
  parsed = parse_expr(parser, expr)
  @test parsed.name == nothing
  @test parsed.args == [:(a::String), :(b::Int...)]
  @test parsed.kwargs == [Expr(:kw, :c, 5), :(d...)]
  @test parsed.curlies == []
  @test parsed.wheres == [:A, :B]
  test_closure(parser, expr)

  expr = :(
    f{A}(a::String, b::Int...; c = 5, d...) where {A, B} = nothing
  )
  parsed = parse_expr(parser, expr)
  @test parsed.name == :f
  @test parsed.args == [:(a::String), :(b::Int...)]
  @test parsed.kwargs == [Expr(:kw, :c, 5), :(d...)]
  @test parsed.curlies == [:A]
  @test parsed.wheres == [:A, :B]
  test_closure(parser, expr)

  @test_throws EP.ParseError parse_expr(parser, :(
    f = 5
  ))

  parsed = EP.Function_Parsed()
  @test parse_expr(EP.Function(), to_expr(parsed)) == parsed
end

@testset "Macro" begin
  parser = EP.Macro()
  expr = :(@mymacro whatever 3)
  parsed = parse_expr(parser, expr)
  @test parsed.name == :mymacro
  @test parsed.args == [:whatever, 3]
  test_closure(parser, expr)

  parsed = EP.Macro_Parsed(name = :mymcaro)
  @test parse_expr(EP.Macro(), to_expr(parsed)) == parsed
end

@testset "NestedDot" begin
  parser = EP.NestedDot()
  expr = :(A{T}("hi").b.c.d)
  parsed = parse_expr(parser, expr)
  @test parsed.base == :(A{T}("hi"))
  @test parsed.properties == [:b, :c, :d]
  test_closure(parser, expr)

  @test_throws EP.ParseError parse_expr(parser, :(A{hi}))

  parsed = EP.NestedDot_Parsed(base = :a, properties = [:b, :c])
  @test parse_expr(EP.NestedDot(), to_expr(parsed)) == parsed
end

@testset "Reference" begin
  parser = EP.Reference()
  @test parse_expr(parser, :a).name == :a
  expr = :(A{T, S})
  parsed = parse_expr(parser, expr)
  @test parsed.name == :A
  @test parsed.curlies == [:T, :S]
  test_closure(parser, expr)

  expr = :(A.B.C{T, S})
  parsed = parse_expr(parser, expr)
  @test parsed.name == :(A.B.C)
  @test parsed.curlies == [:T, :S]
  test_closure(parser, expr)

  expr = :(A.B.C)
  parsed = parse_expr(parser, expr)
  @test parsed.name == :(A.B.C)
  @test parsed.curlies == []
  test_closure(parser, expr)

  @test_throws EP.ParseError parse_expr(parser, :(f(a, b)))

  parsed = EP.Reference_Parsed(name = :a)
  @test parse_expr(EP.Reference(), to_expr(parsed)) == parsed
end

@testset "Type" begin
  parser = EP.Type()
  expr = Any
  @test parse_expr(parser, expr).name == Any
  @test parse_expr(parser, expr).curlies == []
  @test parse_expr(parser, expr).wheres == []
  test_closure(parser, expr)

  expr = :Any
  @test parse_expr(parser, expr).name == :Any
  @test parse_expr(parser, expr).curlies == []
  @test parse_expr(parser, expr).wheres == []
  test_closure(parser, expr)

  expr = :(A{T, S})
  @test parse_expr(parser, expr).name == :A
  @test parse_expr(parser, expr).curlies == [:T, :S]
  @test parse_expr(parser, expr).wheres == []
  test_closure(parser, expr)

  expr = :(A{T, S} where S where T)
  @test parse_expr(parser, expr).name == :A
  @test parse_expr(parser, expr).curlies == [:T, :S]
  @test parse_expr(parser, expr).wheres == [:T, :S]
  test_closure(parser, expr)

  @test_throws EP.ParseError parse_expr(parser, :(f(a, b)))

  parsed = EP.Type_Parsed(name = :a)
  @test parse_expr(EP.Type(), to_expr(parsed)) == parsed
end

@testset "TypeAnnotation" begin
  parser = EP.TypeAnnotation()
  expr = :(::Any)
  @test parse_expr(parser, expr).name == nothing
  @test parse_expr(parser, expr).type == :Any
  test_closure(parser, expr)

  expr = :(something(hi)::Any)
  @test parse_expr(parser, expr).name == :(something(hi))
  @test parse_expr(parser, expr).type == :Any
  test_closure(parser, expr)

  @test_throws EP.ParseError parse_expr(parser, :a)

  parsed = EP.TypeAnnotation_Parsed(name = :a, type = Any)
  @test parse_expr(EP.TypeAnnotation(), to_expr(parsed)) == parsed
end

@testset "TypeRange" begin
  parser = EP.TypeRange()
  expr = :(A <: Any)
  @test parse_expr(parser, expr).lb == Union{}
  @test parse_expr(parser, expr).name == :A
  @test parse_expr(parser, expr).ub == :Any
  test_closure(parser, expr)

  expr = :(A >: Integer)
  @test parse_expr(parser, expr).lb == :Integer
  @test parse_expr(parser, expr).name == :A
  @test parse_expr(parser, expr).ub == Any
  test_closure(parser, expr)

  expr = :(Number <: B <: Any)
  @test parse_expr(parser, expr).lb == :Number
  @test parse_expr(parser, expr).name == :B
  @test parse_expr(parser, expr).ub == :Any
  test_closure(parser, expr)

  @test_throws EP.ParseError parse_expr(parser, :(a = 4))

  parsed = EP.TypeRange_Parsed(name = :T)
  @test parse_expr(EP.TypeRange(), to_expr(parsed)) == parsed
end


text:
using StreamGraphs
using Test
using Logging

# Testing tuples
gl = global_logger()
global_logger(ConsoleLogger(gl.stream, Logging.Error))

@testset "Sets Tests" begin
    @info "Check operations on sets"
    include("sets.jl")
end

@testset "Tuple Tests" begin
    @info "Check tuple equality approximation"
    include("tuples.jl")
end

@testset "Intervals Tests" begin
    @info "Check Intervals operations"
    include("intervals.jl")
end

@testset "Node Tests" begin
    @info "Check Nodes"
    include("nodes.jl")
end

@testset "Link Tests" begin
    @info "Check Links"
    include("links.jl")
end

@testset "LinkStream Tests" begin
    @info "Check LinkStreams"
    include("link_streams.jl")
end

@testset "StreamGraph Tests" begin
    @info "Check StreamGraphs"
    include("stream_graphs.jl")
end

@testset "TimeCursor Tests" begin
    @info "Check TimeCursors"
    include("time_cursors.jl")
end

@testset "Parser Tests" begin
    @info "Check Parsers"
    include("parsers.jl")
end

@testset "Jump Tests" begin
    @info "Check Jumps"
    include("jumps.jl")
end

@testset "Path Tests" begin
    @info "Check Paths"
    include("paths.jl")
end


text:
using Test
import IJulia: Comm, comm_target

@testset "comm" begin
    target = :notarget
    comm_id = "6BA197D8A67A455196279A59EB2FE844"
    comm = Comm(target, comm_id, false)
    @test :notarget == comm_target(comm)
    @test !comm.primary

    # comm_info_request in comm_manager.jl
    comms = Dict{String, Comm}(
        "id" => Comm(Symbol("jupyter.widget"), "id", false)
    )
    msg_content = Dict("target_name" => "jupyter.widget")
    reply = if haskey(msg_content, "target_name")
        let t, cb
            t = Symbol(msg_content["target_name"])
            cb(k, v) = comm_target(v) == t # For 0.6
            cb(kv) = comm_target(kv[2]) == t # For 0.7
            filter(cb, comms)
        end
    else
        comms
    end
    @test Dict{String,Comm} == typeof(reply)
    _comms = Dict{String,Dict{Symbol,Symbol}}()
    for (comm_id,comm) in reply
        local comm_id, comm
        _comms[comm_id] = Dict(:target_name => comm_target(comm))
    end
    @test Dict("id"=>Dict(:target_name=>Symbol("jupyter.widget"))) == _comms
    content = Dict(:comms => _comms)
    @test Dict(:comms=>Dict("id"=>Dict(:target_name=>Symbol("jupyter.widget")))) == content
end


text:
using modular_CMD
using Test

@testset "modular_CMD.jl" begin
    # Write your tests here.
end