function tr_minimize!(objective, x0, approach, B0, options)
    T = eltype(x0)

    x, fx, ∇fx, z, fz, ∇fz, B = prepare_variables(objective, approach, x0, copy(x0), B0)
    p = copy(x)

    Δk = T(10.0)

    x, fx, ∇fx, z, fz, ∇fz, B, Δkp1, accept, is_converged = iterate!(p, x, fx, ∇fx, z, fz, ∇fz, B, Δk, approach, objective, options)

    iter = 0
    is_converged = false
    while iter <= options.maxiter && !is_converged
        iter += 1
        x, fx, ∇fx, z, fz, ∇fz, B, Δkp1, accept, is_converged = iterate!(p, x, fx, ∇fx, z, fz, ∇fz, B, Δkp1, approach, objective, options)
    end
    return z, fz, ∇fz, iter
end

function iterate!(p, x, fx, ∇fx, z, fz, ∇fz, Bx, Δk, approach, objective, options; scale=false)
    T = eltype(x)
    scheme, subproblemsolver = approach

    fx = fz
    copyto!(x, z)
    copyto!(∇fx, ∇fz)

    spr = subproblemsolver(∇fx, Bx, Δk, p; abstol=1e-10, maxiter=50)
    Δm = -spr.mz
    # Grab the model value, m. If m is zero, the solution, z, does not improve
    # the model value over x. If the model is not converged, but the optimal
    # step is inside the trust region and gives a zero improvement in the objec-
    # tive value, we may conclude that "something" is wrong. We might be at a
    # ridge (positive-indefinite case) for example, or the scaling of the model
    # is such that we cannot satisfy ||∇f|| < tol.
    if abs(spr.mz) < eps(T)
        # set flag to check for problems
    end

    # add Retract
    z .= x + p

    # Update before acceptance, to keep adding information about the hessian
    # even when the step is not "good" enough.
    fz, ∇fz, B = update_obj(objective, spr.p, spr.p, ∇fx, z, ∇fz, Bx, scheme, scale)

    # Δf is often called ared or Ared for actual reduction. I prefer "change in"
    # f, or Delta f.
    Δf = fx - fz

    # Calculate the ratio of actual improvement over predicted improvement.
    R = Δf/Δm

    Δkp1, reject_step = update_trust_region(spr, R, p)

    if reject_step
        z .= x
        fz = fx
        ∇fz .= ∇fx
    end

    is_converged = converged(z, ∇fz, options.g_tol)

    return x, fx, ∇fx, z, fz, ∇fz, B, Δkp1, reject_step, is_converged
end

function update_trust_region(spr, R, p)
    T = eltype(p)
    # Chosing a parameter > 0 might be preferable here. See p. 4 of Yuans survey
    # We want to avoid cycles, but we also need something that takes very small
    # steps when convergence is hard to achieve.
    α = T(1)/7 # acceptance ratio
    t2 = T(1)/4
    t3 = t2 # could differ!
    t4 = T(1)/2
    λ34 = T(0)/2
    γ = T(2.5) # gamma for grow
    λγ = T(0)/2 # distance along growing interval ∈ (0, 1]
    Δmax = T(10)^5 # restrict the largest step
    σ = T(1)/4

    Δk = spr.Δ
    # We accept all steps larger than α ∈ [0, 1/4). See p. 415 of [SOREN] and
    # p.79 as well as  Theorem 4.5 and 4.6 of [N&W]. A α = 0 might cycle,
    # see p. 4 of [YUAN].
    if !(α <= R)
        if spr.interior
            # If you reject an interior solution, make sure that the next
            # delta is smaller than the current step. Otherwise you waste
            # steps reducing Δk by constant factors while each solution
            # will be the same.
            Δkp1 = σ * norm(p)
        else
            Δkp1 = λ34*norm(p, 2)*t3 + (1-λ34)*Δk*t4
        end
        reject_step = true
    else
        # While we accept also the steps in the case that α <= Δf < t2, we do not
        # trust it too much. As a result, we restrict the trust region radius. The
        # new trust region radius should be set to a radius Δkp1 ∈ [t3*||d||, t4*Δk].
        # We use the number λ34 ∈ [0, 1] to move along the interval. [N&W] sets
        # λ34 _= 1 and t4 = 1/4, see Algorithm 4.1 on p. 69.
        if R < t2
            Δkp1 = λ34*norm(p, 2)*t3 + (1-λ34)*Δk*t4
        else
            Δkp1 = min(λγ*Δk+(1-λγ)*Δk*γ, Δmax)
        end
        reject_step = false
    end
    return Δkp1, reject_step
end


function test_pairwise(F::Function, a, b, s::Integer)
    s1 = minhash(F, a, s)
    s2 = minhash(F, b, s)
    @test intersectionlength(s1, s2) == length(intersect(s1.hashes, s2.hashes))
end

function test_multi(F::Function, x, s)
    sketches = [minhash(F, i, s) for i in x]
    m = intersectionlength(sketches)
    allgood = true
    for i in 1:length(sketches) - 1
        for j in i+1:length(sketches)
            tp = intersectionlength(sketches[i], sketches[j])
            allgood &= tp == m[j, i]
        end
    end
    @test allgood
end

@testset "Pairwise" begin
    test_pairwise(hash, 1:100, 60:200, 40)
    test_pairwise(hash, 1:100, 'A':'z', 40)
    test_pairwise(hash, 1:100, 'A':'z', 1)
    test_pairwise(hash, 'A':'g', 'W':'s', 10)
end

@testset "Multi" begin
    xs = [collect(Set(rand(1:250, 100))) for i in 1:5]
    test_multi(hash, xs, 50)
end

@testset "Misc" begin
    s = minhash('A':'z', 38)
    @test length(s) == 38
    s = minhash(5:55, 250)
    @test length(s) == 51
    @test !isempty(s)
    @test isempty(minhash(hash, 1:0, 100))
end


using PDMats: PDMat
using GaussianProcesses: make_posdef!

function chistat(gpT::GPE, gpC::GPE, Xb::AbstractMatrix)
    pred_T = predict_full(gpT, Xb)
    pred_C = predict_full(gpC, Xb)
    μ = pred_T[1].-pred_C[1]
    Σ = pred_T[2]+pred_C[2]
    return dot(μ, Σ \ μ)
end
function chistat(gpT::GPE, gpC::GPE, Xb::AbstractMatrix,
                 Σcliff::PDMat, cK_T::AbstractMatrix, cK_C::AbstractMatrix)
    μT = predict_mu(gpT, Xb, cK_T)
    μC = predict_mu(gpC, Xb, cK_C)
    μ = μT - μC
    return dot(μ, Σcliff \ μ)
end

function sim_chi_null!(
        gpT::GPE, gpC::GPE, gpNull::GPE,
        treat::BitVector, Xb::AbstractMatrix,
        Σcliff::PDMat, cK_T::AbstractMatrix, cK_C::AbstractMatrix)
    Ysim = prior_rand(gpNull)

    gpT.y = Ysim[treat]
    gpC.y = Ysim[.!treat]

    update_alpha!(gpT)
    update_alpha!(gpC)

    return chistat(gpT, gpC, Xb, Σcliff, cK_T, cK_C)
end
function nsim_chi(gpT::GPE, gpC::GPE, gpNull::GPE, Xb::AbstractMatrix, nsim::Int)
    gpT_mod = modifiable(gpT)
    gpC_mod = modifiable(gpC)

    _, Σcliff = cliff_face(gpT, gpC, Xb)
    cK_T = cov(gpT.kernel, Xb, gpT.x)
    cK_C = cov(gpC.kernel, Xb, gpC.x)

    treat = BitVector(undef, gpNull.nobs)
    treat[:] .= false
    treat[1:gpT.nobs] .= true
    kernel = gpT_mod.kernel
    mT = gpT_mod.mean
    mC = gpC_mod.mean
    Σraw, chol = make_posdef!(copy(Σcliff))
    PDcliff = PDMat(Σraw, chol)
    t_sims = [sim_chi_null!(gpT_mod, gpC_mod, gpNull, treat, Xb, PDcliff, cK_T, cK_C)
        for _ in 1:nsim];
    return t_sims
end

function boot_chi2test(gpT::GPE, gpC::GPE, Xb::AbstractMatrix, nsim::Int)
    gpNull = make_null(gpT, gpC)
    chi_obs = chistat(gpT, gpC, Xb)
    chi_sims = GeoRDD.nsim_chi(gpT, gpC, gpNull, Xb, nsim)
    return mean(chi_sims .> chi_obs)
end

function placebo_chi(angle::Float64, X::AbstractMatrix, Y::Vector,
                 kern::Kernel, mean::Mean, logNoise::Float64,
                 nsim::Int)
    shift = shift_for_even_split(angle, X)
    left = left_points(angle, shift, X)
    gp_left  = GPE(X[:,left],   Y[left], mean, kern, logNoise)
    gp_right = GPE(X[:,.!left], Y[.!left], mean, kern, logNoise)
    Xb = placebo_sentinels(angle, shift, X, 100)
    pval = boot_chi2test(gp_left, gp_right, Xb, nsim)
    return pval
end
function placebo_chi(angle::Float64, gp::GPE, nsim::Int; kwargs...)
    logNoise = convert(Float64, gp.logNoise)
    return placebo_chi(angle, gp.x, gp.y, gp.kernel, gp.mean, logNoise, nsim; kwargs...)
end


using DataStructures

function solve()
    n, m = [parse(Int, x) for x in split(readline())]
    ds = IntDisjointSets(n)
    for _=1:m
        a, b = [parse(Int, x) for x in split(readline())]
        union!(ds, a, b)
    end

    size = zeros(Int, n)
    for x=1:n
        root = find_root!(ds, x)
        size[root] += 1
    end
    maximum(size)
end

println(solve())


export pairwise_marginal, pairwise_MI

using LinearAlgebra: diagind, diag
using CUDA: CUDA, CuMatrix, CuVector, CuArray

xlogx(x) =
    iszero(x) ? zero(x) : x * log(x)

xlogy(x, y) =
    iszero(x) && !isnan(y) ? zero(x) : x * log(y)

#############################
# Mutual Information on Binary Data
#############################

function pairwise_marginal(data::Union{Matrix{Bool}, CuMatrix{Bool}, BitMatrix};
        weights::Union{Vector, CuVector, Nothing}=nothing,
        pseudocount=0.0,
        Float=Float64)
    N = isnothing(weights) ? size(data, 1) : sum(weights)
    D = size(data, 2)
    base = N + pseudocount
    if data isa CuMatrix{Bool} && !isnothing(weights)
        weights = CuVector(weights)
    end

    not_data = Float.(.!data)
    data = Float.(data)
    pxy = similar(data, D, D, 4)

    if isnothing(weights)
        pxy[:, :, 4] = transpose(data) * data
        pxy[:, :, 3] = transpose(data) * not_data
        pxy[:, :, 2] = transpose(not_data) * data
    else
        pxy[:, :, 4] = transpose(data) * (data .* weights)
        pxy[:, :, 3] = transpose(data) * (not_data .* weights)
        pxy[:, :, 2] = transpose(not_data) * (data .* weights)
    end
    pxy[:, :, 1] = N .- pxy[:, :, 4] .- pxy[:, :, 3] .- pxy[:, :, 2]

    joint_count = pseudocount / 4
    pxy .+= joint_count

    # diagonal is the marginal of p(x)
    diag = diagind(pxy[:, :, 1])
    @view(pxy[:,:, 1])[diag] .+= joint_count
    @view(pxy[:,:, 2])[diag] .-= joint_count
    @view(pxy[:,:, 3])[diag] .-= joint_count
    @view(pxy[:,:, 4])[diag] .+= joint_count

    if data isa CuMatrix
        CUDA.unsafe_free!(data)
        CUDA.unsafe_free!(not_data)
    end
    pxy ./= base
end

function pairwise_MI(data::Union{Matrix{Bool}, CuMatrix{Bool}, BitMatrix};
        weights=nothing,
        pseudocount=0.0,
        Float=Float64)
    pxy = pairwise_marginal(data; weights, pseudocount, Float)
    p0 = diag(pxy[:, :, 1])
    p1 = diag(pxy[:, :, 4])
    D = size(data, 2)
    pxpy = Float.(similar(data, D, D, 4))
    pxpy[:,:,1] = p0 * p0'
    pxpy[:,:,2] = p0 * p1'
    pxpy[:,:,3] = p1 * p0'
    pxpy[:,:,4] = p1 * p1'
    pxy_log_pxy = @. xlogx(pxy)
    pxy_log_pxpy = @. xlogy(pxy, pxpy)
    MI = dropdims(sum(pxy_log_pxy - pxy_log_pxpy, dims=3), dims=3)
    if data isa CuMatrix{Bool}
        CUDA.unsafe_free!(pxy)
        CUDA.unsafe_free!(pxpy)
        CUDA.unsafe_free!(pxy_log_pxy)
        CUDA.unsafe_free!(pxy_log_pxpy)
    end
    MI
end

#############################
# Mutual Information on Categorical Data
#############################

function pairwise_marginal(data::Matrix;
        weights::Union{Vector, Nothing}=nothing,
        num_cats=maximum(data),
        pseudocount=0.0,
        Float=Float64)
    @assert minimum(data) > 0 "Categorical data labels are assumed to be indexed starting 1"

    num_samples = size(data, 1)
    num_vars = size(data, 2)

    if weights === nothing
        weights = ones(Int32, num_samples)
    else
        pseudocount = pseudocount * sum(weights) / num_samples
    end

    num_vars = size(data, 2)
    pair_margs = Array{Float}(undef, num_vars, num_vars, num_cats, num_cats)
    Z = sum(weights) + pseudocount
    single_smooth = pseudocount / num_cats
    pair_smooth = single_smooth / num_cats

    for i = 1:num_vars
        # note: multithreading needs to be on inner loop for thread-safe copying across diagonal
        Threads.@threads for j = 1:num_vars
            if i<=j
                if i!=j
                    @inbounds pair_margs[i,j,:,:] .= pair_smooth
                else
                    @inbounds pair_margs[i,j,:,:] .= zero(eltype(pair_margs))
                    for l = 1:num_cats
                        @inbounds pair_margs[i,j,l,l] = single_smooth
                    end
                end
                @simd for k = 1:size(data,1) # @avx here gives incorrect results
                    @inbounds pair_margs[i,j,data[k,i],data[k,j]] += weights[k]
                end
                @inbounds pair_margs[i,j,:,:] ./= Z
            else
                @inbounds pair_margs[i,j,:,:] .= (@view pair_margs[j,i,:,:])'
            end
            nothing
        end
    end

    return pair_margs
end

function pairwise_marginal(data::CuMatrix;
        weights::Union{CuVector, Nothing}=nothing,
        num_cats=maximum(data),
        pseudocount=0.0,
        Float=Float64)

    @assert minimum(data) > 0 "Categorical data labels are assumed to be indexed starting 1"

    num_samples = size(data, 1)
    num_vars = size(data, 2)

    if weights === nothing
        weights = ones(Int32, num_samples)
    else
        pseudocount = pseudocount * sum(weights) / num_samples
    end

    num_vars = size(data,2)
    pair_margs = CuArray{Float}(undef, num_vars, num_vars, num_cats, num_cats)
    Z = sum(weights) + pseudocount
    single_smooth = pseudocount / num_cats
    pair_smooth = single_smooth / num_cats

    data_device = CUDA.cudaconvert(data)
    weights_device = CUDA.cudaconvert(weights)
    pair_margs_device = CUDA.cudaconvert(pair_margs)

    var_indices = CuArray(1:num_vars)
    CUDA.@sync begin
        broadcast(var_indices, var_indices') do i,j
            if i <= j
                if i!=j
                    @inbounds pair_margs_device[i,j,:,:] .= pair_smooth
                else
                    @inbounds pair_margs_device[i,j,:,:] .= zero(Float)
                    for l = 1:num_cats
                        @inbounds pair_margs_device[i,j,l,l] = single_smooth
                    end
                end
                for k = 1:size(data_device,1)
                    pair_margs_device[i,j,data_device[k,i],data_device[k,j]] += weights_device[k]
                end
            end
            nothing
        end
        pair_margs ./= Z
        broadcast(var_indices, var_indices') do i,j
            if i > j
                for l = 1:num_cats, m = 1:num_cats
                    @inbounds pair_margs_device[i,j,l,m] = pair_margs_device[j,i,m,l]
                end
            end
            nothing
        end
    end
    return pair_margs
end

function pairwise_MI(data::Matrix;
            weights::Union{Vector, Nothing} = nothing,
            num_cats=maximum(data),
            pseudocount=0.0,
            Float=Float64)
    num_samples = size(data, 1)
    num_vars = size(data, 2)

    if weights === nothing
        weights = ones(Int32, num_samples)
    else
        pseudocount = pseudocount * sum(weights) / num_samples
    end

    # `joint_cont[i, j, k, w]' is the total weight of samples whose i- and j-th variable are k and w, respectively
    joint_cont = pairwise_marginal(data; weights, num_cats, pseudocount)

    # `marginal_cont[i, j]' is the total weight of sample whose i-th variable is j
    marginal_cont = zeros(Float, num_vars, num_cats)
    for i = 1:num_vars
        for j = 1:num_cats
            @inbounds marginal_cont[i,j] = joint_cont[i,i,j,j]
        end
    end

    # Compute mutual information
    MI = zeros(Float, num_vars, num_vars)
    for var1_idx = 1 : num_vars
        for var2_idx = var1_idx : num_vars
            @inbounds MI[var1_idx, var2_idx] = sum(
                xlogx.(joint_cont[var1_idx, var2_idx, :, :]) .-
                xlogy.(joint_cont[var1_idx, var2_idx, :, :], marginal_cont[var1_idx, :] .* marginal_cont[var2_idx, :]'))
        end
    end

    for var1_idx = 2 : num_vars
        for var2_idx = 1 : var1_idx - 1
            @inbounds MI[var1_idx, var2_idx] = MI[var2_idx, var1_idx]
        end
    end

    MI
end

# TODO categorical pairwise MI on GPU


"""
    test_bqg_rossbywave(stepper, dt, nsteps, dev; kwargs...)

Evolves a Rossby wave and compares with the analytic solution.
"""
function test_bqg_rossbywave(stepper, dt, nsteps, dev::Device=CPU())
    nx = 64
    Lx = 2π
     β = 2.0
     μ = 0.0
     ν = 0.0
     T = Float64

  # the following if statement is called so that all the cases of
  # Problem() fuction are tested
  if stepper=="ForwardEuler"
    eta = zeros(dev, T, (nx, nx))
  else
    eta(x, y) = 0 * x
  end

  prob = BarotropicQG.Problem(dev; nx=nx, Lx=Lx, eta=eta, β=β, μ=μ, ν=ν, stepper=stepper, dt=dt)
  sol, clock, vars, params, grid = prob.sol, prob.clock, prob.vars, prob.params, prob.grid

  x, y = gridpoints(grid)

  # the Rossby wave initial condition
   ampl = 1e-2
  kwave = 3 * 2π/grid.Lx
  lwave = 2 * 2π/grid.Ly
      ω = -params.β * kwave / (kwave^2 + lwave^2)
     ζ0 = @. ampl * cos(kwave * x) * cos(lwave * y)
    ζ0h = rfft(ζ0)

  BarotropicQG.set_zeta!(prob, ζ0)

  stepforward!(prob, nsteps)
  dealias!(sol, grid)
  BarotropicQG.updatevars!(prob)

  ζ_theory = @. ampl * cos(kwave * (x - ω/kwave * clock.t)) * cos(lwave * y)

  return isapprox(ζ_theory, vars.zeta, rtol=grid.nx * grid.ny * nsteps * 1e-12)
end

"""
    test_bqg_stochasticforcing_energy_budgets(dev; kwargs...)

Tests if the energy budget is closed for a BarotropicQG problem with stochastic forcing.
"""
function test_bqg_stochasticforcing_energybudget(dev::Device=CPU(); n=256, dt=0.01, L=2π, ν=1e-7, nν=2, μ=1e-1, T=Float64)
  n, L  = 256, 2π
  ν, nν = 1e-7, 2
  μ = 1e-1
  dt, tf = 0.005, 0.1/μ
  nt = round(Int, tf/dt)

  grid = TwoDGrid(dev, n, L)
  x, y = gridpoints(grid)
  K  = @. sqrt(grid.Krsq)                      # a 2D array with the total wavenumber

  # Forcing with spectrum ~ exp(-(k-kf)²/(2 dkf²))
  kf, dkf = 12.0, 2.0
  ε = 0.1

  forcing_spectrum = zeros(dev, T, (grid.nkr, grid.nl))
  @. forcing_spectrum = exp(-(K - kf)^2 / (2 * dkf^2))
  @. forcing_spectrum = ifelse(grid.Krsq < 2^2,  0, forcing_spectrum)   # no power at low wavenumbers
  @. forcing_spectrum = ifelse(grid.Krsq > 20^2, 0, forcing_spectrum)   # no power at high wavenumbers
  ε0 = parsevalsum(forcing_spectrum .* grid.invKrsq / 2, grid) / (grid.Lx * grid.Ly)
  forcing_spectrum .= ε / ε0 * forcing_spectrum

  Random.seed!(1234)

  function calcF!(Fh, sol, t, clock, vars, params, grid)
    eta = ArrayType(dev)(exp.(2π * im * rand(T, size(sol))) / sqrt(clock.dt))
    CUDA.@allowscalar eta[1, 1] = 0
    @. Fh = eta * sqrt(forcing_spectrum)
    return nothing
  end

  prob = BarotropicQG.Problem(dev; nx=n, Lx=L, ν=ν, nν=nν, μ=μ, dt=dt,
   stepper="RK4", calcF=calcF!, stochastic=true)

  BarotropicQG.set_zeta!(prob, 0*x)

  E = Diagnostic(BarotropicQG.energy,             prob, nsteps=nt)
  D = Diagnostic(BarotropicQG.energy_dissipation, prob, nsteps=nt)
  R = Diagnostic(BarotropicQG.energy_drag,        prob, nsteps=nt)
  W = Diagnostic(BarotropicQG.energy_work,        prob, nsteps=nt)
  diags = [E, D, W, R]

  stepforward!(prob, diags, nt)

  BarotropicQG.updatevars!(prob)

  dEdt_numerical = (E[2:E.i] - E[1:E.i-1]) / prob.clock.dt

  # If the Ito interpretation was used for the work
  # then we need to add the drift term
  # dEdt_computed = W[2:E.i] + ε - D[1:E.i-1] - R[1:E.i-1]      # Ito
  dEdt_computed = W[2:E.i] - D[1:E.i-1] - R[1:E.i-1]        # Stratonovich

  return isapprox(dEdt_numerical, dEdt_computed, rtol=1e-3)
end

"""
    test_bqg_deterministicforcing_energy_budgets(dev; kwargs...)

Tests if the energy budget is closed for a BarotropicQG problem with deterministic forcing.
"""
function test_bqg_deterministicforcing_energybudget(dev::Device=CPU(); n=256, dt=0.01, L=2π, ν=1e-7, nν=2, μ=1e-1)
  n, L  = 256, 2π
  ν, nν = 1e-7, 2
  μ = 1e-1
  dt, tf = 0.005, 0.1/μ
  nt = round(Int, tf/dt)

  grid = TwoDGrid(dev, n, L)
  x, y = gridpoints(grid)
  k₀, l₀ = 2π/grid.Lx, 2π/grid.Ly

  # Forcing = 0.01cos(4x)cos(5y)cos(2t)
  f = @. 0.01 * cos(4k₀*x) * cos(5l₀*y)
  fh = rfft(f)

  function calcF!(Fh, sol, t, clock, vars, params, grid)
    @. Fh = fh * cos(2t)
    return nothing
  end

  prob = BarotropicQG.Problem(dev; nx=n, Lx=L, ν=ν, nν=nν, μ=μ, dt=dt,
   stepper="RK4", calcF=calcF!, stochastic=false)

  BarotropicQG.set_zeta!(prob, 0*x)

  E = Diagnostic(BarotropicQG.energy,             prob, nsteps=nt)
  D = Diagnostic(BarotropicQG.energy_dissipation, prob, nsteps=nt)
  R = Diagnostic(BarotropicQG.energy_drag,        prob, nsteps=nt)
  W = Diagnostic(BarotropicQG.energy_work,        prob, nsteps=nt)
  diags = [E, D, W, R]

  stepforward!(prob, diags, round(Int, nt))

  BarotropicQG.updatevars!(prob)

  dEdt_numerical = (E[3:E.i] - E[1:E.i-2]) / (2 * prob.clock.dt)
  dEdt_computed  = W[2:E.i-1] - D[2:E.i-1] - R[2:E.i-1]

  return isapprox(dEdt_numerical, dEdt_computed, rtol=1e-4)
end

"""
    test_bqg_stochasticforcing_enstrophy_budgets(dev; kwargs...)

Tests if the enstrophy budget is closed for a BarotropicQG problem with stochastic forcing.
"""
function test_bqg_stochasticforcing_enstrophybudget(dev::Device=CPU(); n=256, dt=0.01, L=2π, ν=1e-7, nν=2, μ=1e-1, T=Float64)
  n, L  = 256, 2π
  ν, nν = 1e-7, 2
  μ = 1e-1
  dt, tf = 0.005, 0.1/μ
  nt = round(Int, tf/dt)

  grid = TwoDGrid(dev, n, L)
  x, y = gridpoints(grid)
  K  = @. sqrt(grid.Krsq)                      # a 2D array with the total wavenumber

  # Forcing with spectrum ~ exp(-(k-kf)²/(2 dkf²))
  kf, dkf = 12.0, 2.0
  εᶻ = 0.1

  forcing_spectrum = zeros(dev, T, (grid.nkr, grid.nl))
  @. forcing_spectrum = exp(-(K - kf)^2 / (2 * dkf^2))
  @. forcing_spectrum = ifelse(grid.Krsq < 2^2,  0, forcing_spectrum)   # no power at low wavenumbers
  @. forcing_spectrum = ifelse(grid.Krsq > 20^2, 0, forcing_spectrum)   # no power at high wavenumbers
  εᶻ0 = parsevalsum(forcing_spectrum / 2, grid) / (grid.Lx * grid.Ly)
  forcing_spectrum .= εᶻ / εᶻ0 * forcing_spectrum

  Random.seed!(1234)

  function calcF!(Fh, sol, t, clock, vars, params, grid)
    eta = ArrayType(dev)(exp.(2π * im * rand(T, size(sol))) / sqrt(clock.dt))
    CUDA.@allowscalar eta[1, 1] = 0
    @. Fh = eta * sqrt(forcing_spectrum)
    return nothing
  end

  prob = BarotropicQG.Problem(dev; nx=n, Lx=L, ν=ν, nν=nν, μ=μ, dt=dt,
   stepper="RK4", calcF=calcF!, stochastic=true)

  BarotropicQG.set_zeta!(prob, 0*x)

  Z = Diagnostic(BarotropicQG.enstrophy,             prob, nsteps=nt)
  D = Diagnostic(BarotropicQG.enstrophy_dissipation, prob, nsteps=nt)
  R = Diagnostic(BarotropicQG.enstrophy_drag,        prob, nsteps=nt)
  W = Diagnostic(BarotropicQG.enstrophy_work,        prob, nsteps=nt)
  diags = [Z, D, W, R]

  stepforward!(prob, diags, nt)

  BarotropicQG.updatevars!(prob)

  dZdt_numerical = (Z[2:Z.i] - Z[1:Z.i-1]) / prob.clock.dt

  # If the Ito interpretation was used for the work
  # then we need to add the drift term
  # dZdt_computed = W[2:Z.i] + εᶻ - D[1:Z.i-1] - R[1:Z.i-1]      # Ito
  dZdt_computed = W[2:Z.i] - D[1:Z.i-1] - R[1:Z.i-1]        # Stratonovich

  return isapprox(dZdt_numerical, dZdt_computed, rtol=1e-3)
end

"""
    test_bqg_deterministicforcing_enstrophy_budgets(dev; kwargs...)

Tests if the enstrophy budget is closed for a BarotropicQG problem with deterministic forcing.
"""
function test_bqg_deterministicforcing_enstrophybudget(dev::Device=CPU(); n=256, dt=0.01, L=2π, ν=1e-7, nν=2, μ=1e-1)
  n, L  = 256, 2π
  ν, nν = 1e-7, 2
  μ = 1e-1
  dt, tf = 0.005, 0.1/μ
  nt = round(Int, tf/dt)

  grid = TwoDGrid(dev, n, L)
  x, y = gridpoints(grid)
  k₀, l₀ = 2π/grid.Lx, 2π/grid.Ly

  # Forcing = 0.01cos(4x)cos(5y)cos(2t)
  f = @. 0.01 * cos(4k₀*x) * cos(5l₀*y)
  fh = rfft(f)

  function calcF!(Fh, sol, t, clock, vars, params, grid)
    @. Fh = fh * cos(2t)
    return nothing
  end

  prob = BarotropicQG.Problem(dev; nx=n, Lx=L, ν=ν, nν=nν, μ=μ, dt=dt,
   stepper="RK4", calcF=calcF!, stochastic=false)

  BarotropicQG.set_zeta!(prob, 0*x)

  Z = Diagnostic(BarotropicQG.enstrophy,             prob, nsteps=nt)
  D = Diagnostic(BarotropicQG.enstrophy_dissipation, prob, nsteps=nt)
  R = Diagnostic(BarotropicQG.enstrophy_drag,        prob, nsteps=nt)
  W = Diagnostic(BarotropicQG.enstrophy_work,        prob, nsteps=nt)
  diags = [Z, D, W, R]

  stepforward!(prob, diags, round(Int, nt))

  BarotropicQG.updatevars!(prob)

  dZdt_numerical = (Z[3:Z.i] - Z[1:Z.i-2]) / (2 * prob.clock.dt)
  dZdt_computed  = W[2:Z.i-1] - D[2:Z.i-1] - R[2:Z.i-1]

  return isapprox(dZdt_numerical, dZdt_computed, rtol=1e-4)
end

"""
    test_bqg_nonlinearadvection(dt, stepper, dev; kwargs...)

Tests the advection term in the BarotropicQG module by timestepping a
test problem with timestep dt and timestepper identified by the string stepper.
The test problem is derived by picking a solution ζf (with associated
streamfunction ψf) for which the advection term J(ψf, ζf) is non-zero. Next, a
forcing Ff is derived according to Ff = ∂ζf/∂t + J(ψf, ζf) - ν∇²ζf. One solution
to the vorticity equation forced by this Ff is then ζf. (This solution may not
be realized, at least at long times, if it is unstable.)
"""
function test_bqg_advection(dt, stepper, dev::Device=CPU(); n=128, L=2π, ν=1e-2, nν=1, μ=0.0)
  n, L  = 128, 2π
  ν, nν = 1e-2, 1
   μ = 0.0
  tf = 1.0
  nt = round(Int, tf/dt)

  grid = TwoDGrid(dev, n, L)
  x, y = gridpoints(grid)

  ψf = @. sin(2x) * cos(2y) + 2sin(x) * cos(3y)
  qf = @. -8sin(2x) * cos(2y) - 20sin(x) * cos(3y)

  Ff = @. -(
    ν*( 64sin(2x) * cos(2y) + 200sin(x) * cos(3y) )
    + 8 * ( cos(x) * cos(3y) * sin(2x) * sin(2y) - 3cos(2x) * cos(2y) * sin(x) * sin(3y) )
    )

  Ffh = rfft(Ff)

  function calcF!(Fh, sol, t, cl, v, p, g)
    Fh .= Ffh
    return nothing
  end

  prob = BarotropicQG.Problem(dev; nx=n, Lx=L, ν=ν, nν=nν, μ=μ, dt=dt, stepper=stepper, calcF=calcF!)

  BarotropicQG.set_zeta!(prob, qf)

  stepforward!(prob, round(Int, nt))

  BarotropicQG.updatevars!(prob)

  return isapprox(prob.vars.q, qf, rtol=rtol_barotropicQG)
end

"""
    test_bqg_energyenstrophy(dev)

Tests the energy and enstrophy function for a BarotropicQG problem.
"""
function test_bqg_energyenstrophy(dev::Device=CPU())
  nx, Lx  = 64, 2π
  ny, Ly  = 64, 3π
  grid = TwoDGrid(dev, nx, Lx, ny, Ly)
  k₀, l₀ = 2π/Lx, 2π/Ly # fundamental wavenumbers
  x, y = gridpoints(grid)

  energy_calc = 29/9
  enstrophy_calc = 2701/162

    eta = @. cos(10k₀*x) * cos(10l₀*y)
   psi0 = @. sin(2k₀*x) * cos(2l₀*y) + 2sin(k₀*x) * cos(3l₀*y)
  zeta0 = @. - ((2k₀)^2+(2l₀)^2) * sin(2k₀*x) * cos(2l₀*y) - (k₀^2+(3l₀)^2) * 2sin(k₀*x) * cos(3l₀*y)

  prob = BarotropicQG.Problem(dev; nx=nx, Lx=Lx, ny=ny, Ly=Ly, eta = eta, stepper="ForwardEuler")
  sol, cl, v, p, g = prob.sol, prob.clock, prob.vars, prob.params, prob.grid
  BarotropicQG.set_zeta!(prob, zeta0)
  BarotropicQG.updatevars!(prob)

  energyzeta0 = BarotropicQG.energy(prob)
  enstrophyzeta0 = BarotropicQG.enstrophy(prob)

  return isapprox(energyzeta0, energy_calc, rtol=rtol_barotropicQG) && isapprox(enstrophyzeta0, enstrophy_calc, rtol=rtol_barotropicQG) &&
  BarotropicQG.addforcing!(prob.timestepper.N, sol, cl.t, cl, v, p, g)==nothing
end

"""
    test_bqg_problemtype(dev, T)

Tests the BarotropicQG problem constructor for different DataType `T`.
"""
function test_bqg_problemtype(dev, T)
  prob = BarotropicQG.Problem(dev; T=T)

  A = ArrayType(dev)

  return (typeof(prob.sol)<:A{Complex{T}, 2} && typeof(prob.grid.Lx)==T && eltype(prob.grid.x)==T && typeof(prob.vars.u)<:A{T, 2})
end


@doc raw"""
## low voltage gated potassiumd channel

### simple version

$$i_{\text{ltk}} = g_{\text{ltk}} w^4 z (V - E_{\text{K}})$$
- subtype:
    - `kv1`: $\zeta=0.5$
    - `ikcnq`: $\zeta=1.0$
    - `ia`: $\zeta=0.0$

<!-- TODO: ### complex version -->

- modified from `inf_tau_w_ltk_rm.m`, `inf_tau_z_ltk_rm.m` and `I_ltk_rm.m`
"""
function low_voltage_gated_potassium(g::Real; iltkcomplex=false, subtype=:kv1)
    _param = (
        kv1 =   (0.5, 1000),
        ikcnq = (1.0, 10000),
        ia =    (0.0, 10)
    )
    _zeta, _camp = _param[subtype]

    if iltkcomplex
        #TODO
        return SimpleIonChannel("ltk_cmplx", :potassium,
            g, 0, 0,
            Kinetics(), Kinetics())
    else
        # Change back to 100 for kv1 like
        _w_tau = (V) -> 100*(6*exp((V+60)/6)+16*exp(-(V+60)/45))^(-1)+1.5 # IKL
        _w_infty = (V) -> (1+exp(-(V+44.5)/8.4))^(-1/4)
        _w = Kinetics(4, _w_infty, _w_tau)

        _z_tau = (V) -> _camp*(exp((V+60)/20)+exp(-(V+60)/8))^(-1)+50;
        _z = Kinetics(1, -71.0, 10.0, zeta=_zeta, _tau=_z_tau, state=:inactivation)

        return SimpleIonChannel("ltk", :potassium,
            g, _w, _z)
    end
end

@doc raw"""
## high voltage gated potassium channel

"""
function high_voltage_gated_potassium(g::Real; phi=0.85)
    _n_tau = (V) -> 100*(11*exp((V+60)/24)+21*exp(-(V+60)/23))^(-1)+0.7;
    _n_infty = (V) -> (1+exp(-(V+15)/5))^(-1/2)
    _n = HHModel.Kinetics(2, _n_infty, _n_tau)

    _p_tau = (V) -> 100*(4*exp((V+60)/32)+5*exp(-(V+60)/22))^(-1)+5;
    _p = HHModel.Kinetics(1, -23, 6, _tau=_p_tau)

    HHModel.ComplexIonChannel("htk", :potassium,
           g, [0.85, 0.15],
           [(_n, HHModel.Kinetics()),
            (_p, HHModel.Kinetics())]
    )
end

@doc raw"""
## hudgkin huxley sodium current

$$i_{\text{Na}} = g_{\text{Na}} m^3 h (V - E_{\text{Na}})$$

- modified from `inf_tau_m_rm.m`, `inf_tau_h_rm.m` and `I_na_rm.m`
"""
function hh_sodium(g::Real)
    _m_tau = (V) -> 10 / (5*exp((V+60)/18)+36*exp(-(V+60)/25))+0.04
    _m = Kinetics(3, -38.0, 7.0, _tau = _m_tau)

    _h_tau = (V) -> 100 / (7*exp((V+60)/11)+10*exp(-(V+60)/25))+0.6
    _h = Kinetics(1, -65.0, 6.0, _tau=_h_tau, state=:inactivation)

    SimpleIonChannel("ina", :sodium,
        g, _m, _h)
end

@doc raw"""
## hh potassium channel

$$i_{\text{K}} = g_{\text{K}} a ^ 4 b c (V - E_{\text{K}})$$
"""
function hh_potassium(g::Real)
    ik_rule = (knt, var) -> begin
        (a,b,c) = var
        a ^ 4 * b * c
    end

    _a_tau = (V) -> 100*(7*exp((V+60)/14)+29*exp(-(V+60)/24))^(-1)+0.1
    _a_infty = (V) -> (1+exp(-(V+31)/6))^(-1/4)
    _a = HHModel.Kinetics(4, _a_infty, _a_tau)

    _b_tau = (V) -> 1000*(14*exp((V+60)/27)+29*exp(-(V+60)/24))^(-1)+1;
    _b_infty = (V) -> (1+exp((V+66)/7))^(-1/2);
    _b = HHModel.Kinetics(1, _b_infty, _b_tau)

    _c_tau = (V) -> 90*(1+exp(-(V+66)/17))^(-1)+10;
    _c_infty = (V) -> (1+exp((V+66)/7))^(-1/2);
    _c = HHModel.Kinetics(1, _c_infty, _c_tau)

    HHModel.GenericIonChannel("ik", :potassium,
        g, ik_rule,
        [_a, _b, _c]
    );
end

@doc raw"""
## Ih current
$$i_{\text{h}} = g_{\text{h}} r (V - E_{\text{h}})$$

- modified from `inf_tau_r_rm.m` and `I_h_rm.m`
"""
function ihcurrent(g::Real)
    a = 86.86183;
     b = 5.76769;
     c = 3393.34171;
     d = 6.67393;
     e = 295.95798;
    # _r_tau = (V) -> (1/c) * 1/(exp(-(V+a)/b) + exp((V+a)/d))+e;
    _r_tau = (V) -> ((1/c)*1/(exp(-(V+a)/b)) + (exp((V+a)/d)))+e;
    _r = Kinetics(1, -96.11, 8.1, _tau=_r_tau, state=:inactivation)

    SimpleIonChannel("ih", :ih,
        g, Kinetics(), _r)
end

@doc raw"""
## leakage current
$$i_{\text{leak}} = g_{\text{leak}} (V - E_{\text{leak}})$$
"""
function leakage(g::Real)
    SimpleIonChannel("leakage", :leak,
        g, Kinetics(), Kinetics())
end


function test_addExp()
	# sequence of arguments should not matter:
	@test addExp(3.0, 5.0) ≈ addExp(5.0, 3.0)
	@test addExp(3.0, 5.0) ≈ log(exp(3)+exp(5))
end

test_addExp()


export bissect
function bissect(h :: AbstractNLPModel;
                t1 :: Float64 = h.meta.x0[1],
                t2 :: Float64 = 100.0,
                tol :: Float64 = 1e-7,
                maxiter :: Int = 50,
                verbose :: Bool = true)

    (length(h.meta.x0) > 1) && warn("Not a 1-D problem ")
    (tₐ, tᵦ) = trouve_intervalle(h, t1, 1.0)
    gₐ = grad(h, [tₐ])[1]
    gᵦ = grad(h, [tᵦ])[1]
    iter = 0

    verbose &&
        @printf(" iter        tₐ        tᵦ         gₐ        gᵦ        \n")
    verbose &&
        @printf(" %7.2e %7.2e  %7.2e  %7.2e  %7.2e\n", iter, tₐ, tᵦ, gₐ, gᵦ)
    while ((max(abs(gᵦ), abs(gₐ))) > tol) && (iter < maxiter) || iter == 0
      tₚ = (tₐ + tᵦ) / 2
      gₚ = grad(h, [tₚ])[1]
      if gₚ <= 0.0
        tₐ = tₚ
        gₐ = gₚ
      else
        tᵦ = tₚ
        gᵦ = gₚ
      end

      iter += 1
      verbose &&
        @printf(" %7.2e %7.2e  %7.2e  %7.2e  %7.2e\n", iter, tₐ, tᵦ, gₐ, gᵦ)
    end

    topt = (tₐ + tᵦ) / 2.0
    gopt = grad(h, [topt])[1]
    status = :NotSolved
    (abs(gopt) < tol) && (status = :Optimal)
    (iter >= maxiter) && (status = :Tired)
    tired = iter > maxiter
    optimal = abs(gopt) < tol
    return (topt, obj(h, [topt])[1], norm(gopt, Inf), iter, optimal, tired, status, h.counters.neval_obj, h.counters.neval_grad, h.counters.neval_hess)
end


"""
    clusterkcenters(X::AbstractMatrix, kcluster::Int; nReplicates::Int=10) -> F

Perform clustering with K-center algorithm. Input data `X` should belong to AbstractMatrix type
and its columns corresponds to variables, rows are frames. Also, the number of clusters `kcluster` should be specified.

Returns a NamedTuple object `F` which contains cluster index for each sample in `F.indexOfCluster`,
the coordinates of cluster centers in `F.center`, the indices of cluster centers in `F.indexOfCenter`,
and distances of samples from the nearest centers in `F.distanceFromCenter`.

#  Example
```julia-repl
julia> using MDToolbox, Plots
julia> X = rand(1000, 2)
julia> F = clusterkcenters(X, 3)
julia> scatter(X[:, 1], X[:, 2], c=F.indexOfCluster)
```
# References
```
This function uses the method described in
[1] S. Dasgupta and P. M. Long, J. Comput. Syst. Sci. 70, 555 (2005).
[2] J. Sun, Y. Yao, X. Huang, V. Pande, G. Carlsson, and L. J. Guibas, Learning 24, 2 (2009).
```
"""
function clusterkcenters(t::AbstractMatrix, kcluster::Int; nReplicates::Int=10)
    nframe = size(t, 1)
    #ta2, com = decenter(ta)
    dim = size(t, 2)

    indexOfCenter = zeros(Int64, kcluster)
    indexOfCluster = zeros(Int64, nframe)
    distanceFromCenter = zeros(Float64, nframe, 1)

    indexOfCenter_out = similar(indexOfCenter)
    indexOfCluster_out = similar(indexOfCluster)
    distanceFromCenter_out = similar(distanceFromCenter)
    center_out = similar(t, kcluster, dim)
    distanceMax_out = Inf64

    for ireplica = 1:nReplicates
        # create the center of the 1st cluster
        indexOfCenter[1] = rand(1:nframe)
        # first, all points belong to the 1st cluster
        indexOfCluster .= ones(Int64, nframe)
        # distances of the points from the 1st center
        ref = t[indexOfCenter[1]:indexOfCenter[1], :]
        distanceFromCenter .= sum((ref .- t).^2, dims=2)

        for i = 2:kcluster
            indexOfCenter[i] = argmax(distanceFromCenter[:])
            ref = t[indexOfCenter[i]:indexOfCenter[i], :]
            dist = sum((ref .- t).^2, dims=2)
            index = dist[:] .< distanceFromCenter[:]
            if any(index)
                # updated if the dist to a new cluster is smaller than the previous one
                distanceFromCenter[index] .= dist[index]
                indexOfCluster[index] .= i
            end
        end

        distanceMax = maximum(distanceFromCenter)
        if (ireplica == 1) | (distanceMax < distanceMax_out)
          distanceMax_out = distanceMax
          indexOfCenter_out = indexOfCenter
          indexOfCluster_out = indexOfCluster
          center_out = t[indexOfCenter, :]
          distanceFromCenter_out .= distanceFromCenter
        end
        Printf.@printf("%d iteration  distance_max = %f  kcluster = %d\n", ireplica, sqrt(distanceMax_out), kcluster)
    end
    distanceFromCenter_out .= sqrt.(distanceFromCenter_out)

    return (indexOfCluster=indexOfCluster_out, center=center_out, distanceFromCenter=distanceFromCenter, indexOfCenter=indexOfCenter_out)
end

"""
    clusterkcenters(ta::AbstractMatrix, kcluster::Int; nReplicates::Int=10) -> F

Perform clustering with K-center algorithm for a TrjArray variable `ta`.
"""
function clusterkcenters(ta::TrjArray, kcluster::Int; nReplicates::Int=10)
    X = ta.xyz
    F = clusterkcenters(X, kcluster, nReplicates=nReplicates)
    center_ta = ta[F.indexOfCenter, :]
    return (indexOfCluster=F.indexOfCluster, center=center_ta, distanceFromCenter=F.distanceFromCenter, indexOfCenter=F.indexOfCenter)
end

"""
    clustercutoff(X::AbstractMatrix, rcut=10.0) -> F

Perform clustering using the given cutoff radius.

#  Example
```julia-repl
julia> using MDToolbox, Plots
julia> X = rand(1000, 2)
julia> F = clustercutoff(X, 0.5)
julia> scatter(X[:, 1], X[:, 2], c=F.indexOfCluster)
```
# References
```
This function uses the method described in
[1] X. Daura, K. Gademann, B. Jaun, D. Seebach, W. F. van Gunsteren, and A. E. Mark, Angewandte 38, 236–240 (1999).
```
"""
function clustercutoff(t::AbstractMatrix, rcut=10.0)
    nframe = size(t, 1)
    rcut2 = rcut^2

    indexOfCluster = zeros(Int64, nframe)
    id = indexOfCluster .== 0

    center = similar(t)
    k = 1

    while any(id)
        x = view(t, id, :)
        index = view(indexOfCluster, id)

        n_max = 0
        iframe_max = 0
        for iframe = 1:size(x, 1)
            dist = sum((x[iframe:iframe, :] .- x).^2, dims=2)
            n = sum(dist .< rcut2)
            if n_max < n
                n_max = n
                iframe_max = iframe
            end
        end

        center[k:k, :] .= x[iframe_max:iframe_max, :]
        dist = sum((center[k:k, :] .- x).^2, dims=2)
        index[dist[:] .< rcut2] .= k
        k += 1

        id = indexOfCluster .== 0
    end

    return (indexOfCluster=indexOfCluster, center=center[1:(k-1), :])
end

"""
    compute_cov(ta::AbstractMatrix, lagtime::Int=0) -> cov

Compute a variance-covariance or time-lagged covariance matrix from input data `X`
Input data `X` should belong to AbstractMatrix type and its columns corresponds to variables, rows are frames.
Optional input is the lagtime=`lagtime` for the calculation of the covariance matrix (default is `lagtime=0`).

#  Example
```julia-repl
julia> X = rand(1000, 100)
julia> cov = compute_cov(X)
```
"""
function compute_cov(X::AbstractMatrix; lagtime::Int=0)
    nframe = size(X, 1)
    X_centerized = X .- mean(X, dims=1)
    cov = X_centerized[1:(end-lagtime), :]' * X_centerized[(1+lagtime):end, :] ./ (nframe - 1)
end

"""
    rsvd(X::AbstractMatrix; k::Number=10) -> F

Perform the randomized SVD for input data `X`.
Input data `X` should belong to AbstractMatrix type and its columns corresponds to variables, rows are frames.
Users can specify the dimension `k` of subspace onto which the data is randomly projected (by default `k=10`).

Returns a NamedTuple object `F` whose members are same as usual SVD, `F.V`, `F.S`, `F.U`.

#  Example
```julia-repl
julia> X = randn(rand(1000, 10))
julia> F = rsvd(X)
```

# References
```
Halko, Nathan, Per-Gunnar Martinsson, and Joel A. Tropp.
"Finding structure with randomness: Probabilistic algorithms for constructing approximate matrix decompositions."
SIAM review 53.2 (2011): 217-288.
```
"""
function rsvd(A::AbstractMatrix, k::Number=10)
    T = eltype(A)
    m, n = size(A)
    l = k + 2
    @assert l < n
    G = randn(T, n, l)
    #G .= G ./ sqrt.(sum(G.^2, dims=1))
    H = A * G
    @assert m > l
    Q = Matrix(qr!(H).Q)
    @assert m == size(Q, 1)
    @assert l == size(Q, 2)
    T = A' * Q
    V, σ, W = svd(T)
    U = Q * W
    λ = σ .* σ ./ m
    return (V=V[:,1:k], S=λ[1:k], U=U[:,1:k])
end

"""
    pca(X::AbstractMatrix; k=dimension) -> F

Perform principal component analysis (PCA). PCA captures degrees of freedom which have the largest variances.
Input data `X` should belong to AbstractMatrix type and its columns corresponds to variables, rows are frames.

Returns a NamedTuple object `F` which contains the principal components in `F.projection`,
the prncipal modes in the columns of the matrix `F.mode`, and the variances of principal components in `F.variance`.

If `k=dimension` is specified, the randomized SVD is used for reducing memory.
This algorithm first project the data into a randomly selected k+2-dimensional space,
then PCA is performed in the projected data. See the references for details.
Note that if the dimension of `X` is larger than 5000, the randomized SVD is forcibly used with `k=1000`.

#  Example
```julia-repl
julia> using MDToolbox, Plots, Statistics
julia> X = cumsum(rand(1000, 10))
julia> F = pca(X)
julia> plot(F.projection[:, 1], F.projection[:, 2])
```

# References
```
Halko, N., Martinsson, P.-G., Shkolnisky, Y. & Tygert, M.
An Algorithm for the Principal Component Analysis of Large Data Sets.
SIAM J. Sci. Comput. 33, 2580–2594 (2011).
```
"""
function pca(X::AbstractMatrix; k=nothing)
    is_randomized_pca = false
    if !isnothing(k)
        is_randomized_pca = true
        @printf "Warning: a randomized SVD approximation is used for PCA with reduced dimension k = %d\n" k
    end
    dim = size(X, 2)
    if !is_randomized_pca & (dim > 5000)
        k = 1000
        is_randomized_pca = true
        @printf "Warning: the dimension of input data is too large (dim > 5000)\n"
        @printf "Warning: randomized SVD approximation is used for PCA with reduced dimension k = %d\n" k
    end
    if is_randomized_pca
        nframe = size(X, 1)
        X_centerized = X .- mean(X, dims=1)
        F = rsvd(X_centerized, k)
        variance = F.S.^2 ./ (nframe - 1)
        mode = F.V
        projection = X_centerized * mode
    else
        # eigendecomposition of covariance matrix
        covar = compute_cov(X)
        F = eigen(covar, sortby = x -> -x)
        #lambda = F.values[end:-1:1]
        #W = F.vectors[:, end:-1:1]
        variance = F.values
        mode = F.vectors
        # projection
        X_centerized = X .- mean(X, dims=1)
        projection = X_centerized * mode
    end

    return (projection=projection, mode=mode, variance=variance)
end

"""
    tica(X::AbstractMatrix, lagtime::Int=1) -> F

This routine performs time-structure based Independent Component Analysis (tICA).
tICA captures degrees of freedom which are most important in the sense that their motions are very slow.
`X` belongs to AbstractMatrix type and its columns corresponds to variables, rows are frames.
User should specify the lagtime for the calculation of the time-lagged covariance matrix.
Returns a NamedTuple object `F` which contains independent components in `F.projection`,
the independent modes in the columns of the matrix `F.mode`, and the eigenvalues in `F.variance`.

If the dimension of `X` is larger than 5000, the randomized SVD approximation is forcibly used with `k=1000` to reduce memory.

#  Example
```julia-repl
julia> using MDToolbox, Plots, Statistics
julia> X = cumsum(rand(1000, 10))
julia> F = tica(X, 10)
julia> plot(F.projection[:, 1], F.projection[:, 2])
```

# References
```
Naritomi, Y. & Fuchigami, S.
Slow dynamics in protein fluctuations revealed by time-structure based independent component analysis:
The case of domain motions.
The Journal of Chemical Physics 134, 065101 (2011).
```
"""
function tica(X::AbstractMatrix, lagtime::Int=1)
    # standard and time-lagged covariance matrices
    covar0 = compute_cov(X, lagtime=0)
    covar  = compute_cov(X, lagtime=lagtime)

    # symmetrize the time-lagged covariance matrix
    covar .= 0.5 .* (covar .+ covar')

    # calc pseudo-inverse of covar0
    #covar0_inv = pinv(covar0)

    # remove degeneracy for solving the generalized eigenvalue problem
    #F = eigen(covar0, sortby = x -> -x)
    F = pca(X)
    pvariance = F.variance
    pmode = F.mode
    index = pvariance .> 10^(-6)
    pmode = pmode[:, index]
    covar0 = pmode'*covar0*pmode
    covar = pmode'*covar*pmode

    # solve the generalized eigenvalue problem
    F = eigen(covar, covar0, sortby = x -> -x)
    variance = F.values
    mode = F.vectors

    # projection
    mode = pmode * mode
    X_centerized = X .- mean(X, dims=1)
    projection = X_centerized * mode

    # normalize mode vectors
    fac = sqrt.(sum(mode.^2, dims=1))
    mode .= mode ./ fac

    return (projection=projection, mode=mode, variance=variance)
end


# Load sample, Calculate and Plot STA

# load sample (stim and ρ)
include("sample.jl")

# Calculate STA
include("sta.jl")
num_timesteps = 150

sta_val = sta(stim, ρ, num_timesteps)

# Plot
using PyPlot

time = 1-length(sta_val):0

plot(time*2, sta_val, color="red", linewidth=1.0, linestyle="-")

xlabel("Time (ms)");
ylabel("Stimulus");
title("Spike Triggered Average")


using Test
using MDCDL
using Random

@testset "Permutation" begin
    Random.seed!(39328582)
    @testset "inversion" begin
        for d = 1:5
            for n = 1:20
                df = (rand(1:4, d)...,)
                x = rand((4 .* df)...)

                pvx = MDCDL.mdarray2polyphase(x, df)

                nBlocks = pvx.nBlocks
                expctd = copy(pvx.data)
                actual = MDCDL.irotatedimspv(MDCDL.rotatedimspv(pvx.data, nBlocks[1]), nBlocks[1])
                @test actual == expctd
            end
        end
    end

    @testset "cyclic group" begin
        for d = 1:5
            for n = 1:20
                df = (rand(1:4, d)...,)
                x = rand((4 .* df)...)

                expctd = copy(x)

                pvx = MDCDL.mdarray2polyphase(x, df)
                pvdata = pvx.data
                nBlocks = pvx.nBlocks

                for blk in nBlocks
                    pvdata = MDCDL.rotatedimspv(pvdata, blk)
                end

                rpvx = MDCDL.PolyphaseVector(pvdata, nBlocks)
                actual = MDCDL.polyphase2mdarray(rpvx, df)
                @test actual == expctd
            end
        end
    end

    @testset "pv vs md" begin
        for d = 1:5
            for n = 1:20
                df = (fill(1,d)...,)
                szx = rand(1:16, d)
                x = rand(szx...)

                expctd = permutedims(x, [collect(2:d)..., 1])

                pvx = MDCDL.mdarray2polyphase(x, df)
                nBlocks = pvx.nBlocks

                permdata = MDCDL.rotatedimspv(pvx.data, nBlocks[1])
                permpvx = PolyphaseVector(permdata, ([nBlocks[2:d]..., nBlocks[1]]...,))

                actual = MDCDL.polyphase2mdarray(permpvx, df)
                @test actual == expctd
            end
        end
    end

    @testset "PolyphaseVector Type" begin
        for d = 1:5
            for n = 1:20
                df = (fill(1,d)...,)
                szx = rand(1:16, d)
                x = rand(szx...)
                nrot = rand(1:d)

                expctd = permutedims(x, circshift(collect(1:d), -nrot))

                pvx = MDCDL.mdarray2polyphase(x, df)
                permpvx = foldl(1:nrot, init = pvx) do prevpvx, k
                    MDCDL.rotatedimspv(prevpvx)
                end
                actual = MDCDL.polyphase2mdarray(permpvx, df)
                @test actual == expctd
            end
        end
    end
end


using CellListMap
using StaticArrays
import Random

#
# In this test we compute the minimum distance between two independent sets of particles
#
function nearest_neighbor(;N1=1_500,N2=1_500_000,parallel=true,x=nothing,y=nothing)

    # Number of particles, sides and cutoff
    sides = @SVector [250,250,250]
    cutoff = 10.
    box = Box(sides, cutoff)

    # Particle positions
    Random.seed!(321)
    if x === nothing
        x = [ sides .* rand(SVector{3,Float64}) for i in 1:N1 ]
    end
    if y === nothing
        y = [ sides .* rand(SVector{3,Float64}) for i in 1:N2 ]
    end

    # Initialize auxiliary linked lists
    cl = CellList(x, y, box, parallel=parallel)

    # Function that keeps the minimum distance
    f(i, j, d2, mind) = d2 < mind[3] ? (i, j, d2) : mind

    # We have to define our own reduce function here (for the parallel version)
    function reduce_mind(output, output_threaded)
        mind = output_threaded[1]
        for i in 2:length(output_threaded)
            if output_threaded[i][3] < mind[3]
                mind = output_threaded[i]
            end
        end
        return mind
    end

    # Initialize
    mind = (0, 0, +Inf)

    # Run pairwise computation
    mind = map_pairwise(
        (x, y, i, j, d2, mind) -> f(i, j, d2, mind),
        mind,box,cl;reduce=reduce_mind, parallel=parallel
    )

    # Take the square root of the minimum distance to return
    return (mind[1], mind[2], sqrt(mind[3]))

end


using Base.Test
import WaveletScattering: default_motherfrequency, isdyadic, issteerable,
    Gammatone, MexicanHat, Morlet

@test_approx_eq default_motherfrequency(Morlet(), 1) 0.39

@test Gammatone() == Gammatone(4)
@test !isdyadic(Gammatone())
@test_throws MethodError issteerable(Gammatone())

@test isdyadic(MexicanHat())
@test !issteerable(MexicanHat())

@test !isdyadic(Morlet())
@test issteerable(Morlet())


v=:v≖Cross(:A,:B)
e1=[1,0,0];e2=[0,1,0];e3=[0,0,1]
A=:A≖3*:c*Vec(e2)
B=:B≖0.5*Vec(e3)
c=:c≖299792458
vs=v&[A,B,c]
@test vs.rhs==4.49688687e8*Vec(e1)

using LinearAlgebra
v1=Vec([1,2,3]);v2=Vec([3,2,1]);
@test dot(v1,v2)==10


@testset verbose=true "Conversion" begin
    @testset verbose=true "$T" for T in FloatTypes
        using EllipsisNotation
        using Random
        Random.seed!(4321)

        @testset "Array types" begin
            dumb_from_float_array(f) = mapslices(x->Quaternion{T}(x...), f, dims=(1,))[1, ..]
            dumb_from_float_array(f::Vector) = Quaternion{T}(f...)
            function dumb_to_float_array(q)
                f = Array{T}(undef, (4, size(q)...))
                f[1, ..] = getproperty.(q, :w)
                f[2, ..] = getproperty.(q, :x)
                f[3, ..] = getproperty.(q, :y)
                f[4, ..] = getproperty.(q, :z)
                f
            end
            for dims in [(), (5,), (4, 3), (2, 3, 4)]
                if dims == ()
                    q = [randn(Quaternion{T})]
                else
                    q = randn(Quaternion{T}, dims)
                end
                f = to_float_array(q)
                @test size(f) == (4, size(q)...)
                @test f == dumb_to_float_array(q)
                @test q == dumb_from_float_array(f)
                @test q == from_float_array(f)
            end
            q = randn(Quaternion{T})
            f = to_float_array(q)
            @test f isa Vector{T}
            @test eltype(f) === T
            @test size(f) == (4,)
            @test f == q.components
        end

        @testset "Euler angles" begin
            N = 5_000

            random_angles = [2π .* rand(T, 3) .- π for _ in 1:5000]
            for i in 1:N
                α, β, γ = random_angles[i]
                q1 = from_euler_angles(α, β, γ)
                q2 = exp(imz*α/2) * exp(imy*β/2) * exp(imz*γ/2)
                @test q1 ≈ q2 atol=10eps(T)
            end
            q1 = from_euler_angles.(random_angles)
            q2 = map((αβγ)->(exp(αβγ[1]*imz/2)*exp(αβγ[2]*imy/2)*exp(αβγ[3]*imz/2)), random_angles)
            @test maximum(abs, q1 .- q2) < 10eps(T)

            random_rotors = randn(Rotor{T}, N)
            for i in 1:N
                q1 = random_rotors[i]
                q2 = from_euler_angles(to_euler_angles(q1))
                @test q1 ≈ q2 atol=250eps(T)
            end
            q1 = random_rotors
            q2 = from_euler_angles.(to_euler_angles.(random_rotors))
            @test maximum(abs, q1 .- q2) < 250eps(T)
        end

        @testset "Euler phases" begin
            N = 5_000
            ϵ = (T === Float16 ? 20eps(T) : 10eps(T))

            dumb_to_euler_phases(α, β, γ) = [exp(α*im), exp(β*im), exp(γ*im)]

            random_angles = [[3(π*rand(T))-π, π*rand(T), 3(π*rand(T))-π] for _ in 1:5000]
            for i in 1:N
                α, β, γ = random_angles[i]
                z1 = dumb_to_euler_phases(α, β, γ)

                q1 = from_euler_angles(α, β, γ)
                q2 = from_euler_phases(z1...)
                @test min(abs(q1-q2), abs(q1+q2)) < ϵ
                q3 = from_euler_phases(z1)
                @test q2 == q3

                z2 = to_euler_phases(q1)
                @test z1[2] ≈ z2[2] atol=ϵ
                if abs(β) < ϵ  # Typically only happens with Float16
                    @test z1[2] ≈ one(T) atol=ϵ
                    @test z2[2] ≈ one(T) atol=ϵ
                    @test z1[1]*z1[3] ≈ z2[1]*z2[3] atol=ϵ
                elseif abs(β-π) < ϵ  # Typically only happens with Float16
                    @test z1[2] ≈ -one(T) atol=ϵ
                    @test z2[2] ≈ -one(T) atol=ϵ
                    @test z1[1]*z1[3] ≈ z2[1]*z2[3] atol=ϵ
                else
                    @test z1[1] ≈ z2[1] atol=ϵ
                    @test z1[3] ≈ z2[3] atol=ϵ
                end

                q4 = from_euler_angles(α, zero(T), γ)
                @test iszero(q4.x)
                @test iszero(q4.y)
                z3 = dumb_to_euler_phases(α, zero(T), γ)
                z4 = to_euler_phases(q4)
                @test z3[2] ≈ one(T) atol=ϵ
                @test z4[2] ≈ one(T) atol=ϵ
                @test z3[1]*z3[3] ≈ z4[1]*z4[3] atol=ϵ

                q6 = from_euler_angles(α, T(π), γ)
                @test q6.w ≈ zero(T) atol=ϵ
                @test q6.z ≈ zero(T) atol=ϵ
                z5 = dumb_to_euler_phases(α, T(π), γ)
                z6 = to_euler_phases(q6)
                @test z5[2] ≈ -one(T) atol=ϵ
                @test z6[2] ≈ -one(T) atol=ϵ
                @test z5[1]*conj(z5[3]) ≈ z6[1]*conj(z6[3]) atol=(T === Float16 ? 20ϵ : ϵ)

                q7 = Quaternion(0, q6.x, q6.y, 0)
                z7 = to_euler_phases(q7)
                @test z7[2] == -one(T)
                @test z7[1]*conj(z7[3]) ≈ exp((α-γ)*im) atol=10eps(T)
            end
        end

        @testset "Spherical coordinates" begin
            N = 5_000

            random_angles = [2π .* rand(T, 3) .- π for _ in 1:5000]
            for i in 1:N
                α, β, γ = random_angles[i]
                q1 = from_spherical_coordinates(β, α)
                q2 = exp(imz*α/2) * exp(imy*β/2) * exp(imz*γ/2)
                @test q1*imz*inv(q1) ≈ q2*imz*inv(q2) atol=10eps(T)
            end

            random_rotors = randn(Rotor{T}, N)
            for i in 1:N
                q1 = random_rotors[i]
                q2 = from_spherical_coordinates(to_spherical_coordinates(q1))
                @test q1*imz*inv(q1) ≈ q2*imz*inv(q2) atol=420eps(T)
            end
        end

        @testset "Rotation matrices" begin
            if isbitstype(T)
                for _ in 1:5_000
                    q1 = randn(Rotor{T})
                    q2 = from_rotation_matrix(to_rotation_matrix(q1))
                    @test min(abs(q1-q2), abs(q1+q2)) < 50eps(T)
                end
            end
        end
    end
end


### utility functions ###

function logsumexp(w)
    offset = maximum(w)
    exp_w = exp.(w .- offset)
    s = sum(exp_w)
    log(s) + offset
end


# find midpoint of a range without causing an overflow to Inf
function robustmidpoint(a::Float64,b::Float64)
	l = min(a,b)
	u = max(a,b)
	if sign(l) == sign(u)
		m = l + (u-l)/2
	else
		m = (u+l)/2
	end
	m
end


# dimensions
nx = 2#12
nu = 1#4
nw = 0

# variables
@variables y[1:nx], x[1:nx], u[1:nu], w[1:nw]

y1 = ones(nx)
x1 = ones(nx)
u1 = ones(nu)
w1 = ones(nw)

xu1 = [x1; u1]
xuy1 = [x1; u1; y1]

# objective

function objective(x, u)
    Q = Diagonal(ones(nx))
    q = ones(nx)

    R = Diagonal(0.1 * ones(nu))
    r = zeros(nu)

    return log(transpose(x) * Q * x) + exp(dot(q, x)) + sin(transpose(u) * R * u)
end

obj = objective(x, u)
obj_x = Symbolics.gradient(obj, x)
obj_u = Symbolics.gradient(obj, u)
obj_xu = Symbolics.gradient(obj, [x; u])

obj_xx = Symbolics.hessian(obj, x)
obj_uu = Symbolics.hessian(obj, u)
obj_xuxu = Symbolics.hessian(obj, [x; u])

obj_xx_sp = Symbolics.sparsehessian(obj, x)
obj_uu_sp = Symbolics.sparsehessian(obj, u)
obj_xuxu_sp = Symbolics.sparsehessian(obj, [x; u])
typeof(obj_func[1])
obj_func = Symbolics.build_function([obj], x, u)
obj_x_func = Symbolics.build_function(obj_x, x, u)
obj_u_func = Symbolics.build_function(obj_u, x, u)
obj_xu_func = Symbolics.build_function(obj_xu, x, u)

obj_xx_func = Symbolics.build_function(obj_xx_sp.nzval, x, u)
obj_uu_func = Symbolics.build_function(obj_uu_sp.nzval, x, u)
obj_xuxu_func = Symbolics.build_function(obj_xuxu_sp.nzval, x, u)

f = eval(obj_func[1])
@benchmark f($x1, $u1)
@benchmark objective($x1, $u1)

fx = eval(obj_x_func[1])
@benchmark fx($x1, $u1)
@benchmark ForwardDiff.gradient(a -> objective(a, $u1), $x1)

fu = eval(obj_u_func[1])
@benchmark fu($x1, $u1)
@benchmark ForwardDiff.gradient(a -> objective($x1, a), $u1)

fxu = eval(obj_xu_func[1])
@benchmark fxu($x1, $u1)

fxx = eval(obj_xx_func[1])
@benchmark fxx($x1, $u1)

fuu = eval(obj_uu_func[1])
@benchmark fuu($x1, $u1)

fxuxu = eval(obj_xuxu_func[1])
@benchmark fxuxu($x1, $u1)

obj_1 = zeros(1)
obj_x1 = zeros(nx)
obj_u1 = zeros(nu)
obj_xu1 = zeros(nx + nu)

obj_xx1 = zeros(length(obj_xx_sp.nzval))
obj_uu1 = zeros(length(obj_uu_sp.nzval))
obj_xuxu1 = zeros(length(obj_xuxu_sp.nzval))

f! = eval(obj_func[2])
@benchmark f!($obj_1, $x1, $u1)

fx! = eval(obj_x_func[2])
@benchmark fx!($obj_x1, $x1, $u1)
@benchmark ForwardDiff.gradient!($obj_x1, a -> objective(a, $u1), $x1)

fu! = eval(obj_u_func[2])
@benchmark fu!($obj_u1, $x1, $u1)
@benchmark ForwardDiff.gradient!($obj_u1, a -> objective($x1, a), $u1)

fxu! = eval(obj_xu_func[2])
@benchmark fxu!($obj_xu1, $x1, $u1)

fxx! = eval(obj_xx_func[2])
@benchmark fxx!($obj_xx1, $x1, $u1)
# @benchmark ForwardDiff.hessian!($obj_xx1, a -> objective(a, $u1), $x1)

fuu! = eval(obj_uu_func[2])
@benchmark fuu!($obj_uu1, $x1, $u1)
# @benchmark ForwardDiff.hessian!($obj_uu1, a -> objective($x1, a), $u1)

fxuxu! = eval(obj_xuxu_func[2])
@benchmark fxuxu!($obj_xuxu1, $x1, $u1)

# discrete dynamics
function quaternion_rotation_matrix(q)
	r, i, j, k  = q

	r11 = 1.0 - 2.0 * (j^2.0 + k^2.0)
	r12 = 2.0 * (i * j - k * r)
	r13 = 2.0 * (i * k + j * r)

	r21 = 2.0 * (i * j + k * r)
	r22 = 1.0 - 2.0 * (i^2.0 + k^2.0)
	r23 = 2.0 * (j * k - i * r)

	r31 = 2.0 * (i * k - j * r)
	r32 = 2.0 * (j * k + i * r)
	r33 = 1.0 - 2.0 * (i^2.0 + j^2.0)

	[r11 r12 r13;
     r21 r22 r23;
     r31 r32 r33]
end

function mrp_quaternion_map(mrp)
	n2 = mrp[1] * mrp[1] + mrp[2] * mrp[2] + mrp[3] * mrp[3]
    M = 2 / (1 + n2)
    return [(1 - n2) / (1 + n2), M * mrp[1], M * mrp[2], M * mrp[3]]
end

mrp_rotation_matrix(mrp) = quaternion_rotation_matrix(mrp_quaternion_map(mrp))

function dynamics(z, u, w)
#     mass = 1.0
#     lc = 1.0
#     gravity = 9.81
#     damping = 0.1

#     [z[2], (u[1] / ((mass * lc * lc)) - gravity * sin(z[1]) / lc - damping * z[2] / (mass * lc * lc))]
# end

    # states
    x = z[1:3]
    r = z[4:6]
    v = z[7:9]
    ω = z[10:12]

    # controls
    w1 = u[1]
    w2 = u[2]
    w3 = u[3]
    w4 = u[4]

    # forces
    kf = 0.1
    F1 = kf * w1
    F2 = kf * w2
    F3 = kf * w3
    F4 = kf * w4

    F = [0.0, 0.0, F1 + F2 + F3 + F4] # total rotor force in body frame

    # moments
    km = 0.01
    M1 = km * w1
    M2 = km * w2
    M3 = km * w3
    M4 = km * w4

    L = 0.25
    τ = [L * (F2 - F4), L * (F3 - F1), (M1 - M2 + M3 - M4)] # total rotor torque in body frame

    gravity = 9.81
    mass = 1.0
    inertia = Diagonal(ones(3))
    inertia_inv = Diagonal(ones(3))

    [v;
     0.25 * ((1.0 - r' * r) * ω - 2.0 * cross(ω, r) + 2.0*(ω' * r) * r);
     [0.0, 0.0, -9.81] + (1.0 / mass) * mrp_rotation_matrix(r) * F;
     inertia_inv * (τ - cross(ω, inertia * ω))]
end

h = 0.1
function discrete_dynamics(y, x, u, w)
    y - (x + h * dynamics(y, u, w))
end

dyn = discrete_dynamics(y, x, u, w)
dyn = simplify.(dyn)

# dense
dyn_y = Symbolics.jacobian(dyn, y)
dyn_x = Symbolics.jacobian(dyn, x)
dyn_u = Symbolics.jacobian(dyn, u)

dyn_xuy = Symbolics.jacobian(dyn, [x; u; y])

dyn_func = Symbolics.build_function(dyn, y, x, u, w)
dyn_y_func = Symbolics.build_function(dyn_y, y, x, u, w)
dyn_x_func = Symbolics.build_function(dyn_x, y, x, u, w)
dyn_u_func = Symbolics.build_function(dyn_u, y, x, u, w)
dyn_xuy_func = Symbolics.build_function(dyn_xuy, y, x, u, w)

# sparse
dyn_y_sp = Symbolics.sparsejacobian(dyn, y)
dyn_x_sp = Symbolics.sparsejacobian(dyn, x)
dyn_u_sp = Symbolics.sparsejacobian(dyn, u)
dyn_xuy_sp = Symbolics.sparsejacobian(dyn, [x; u; y])

dyn_y_func_sp = Symbolics.build_function(dyn_y_sp.nzval, y, x, u, w)
dyn_x_func_sp = Symbolics.build_function(dyn_x_sp.nzval, y, x, u, w)
dyn_u_func_sp = Symbolics.build_function(dyn_u_sp.nzval, y, x, u, w)
dyn_xuy_func_sp = Symbolics.build_function(dyn_xuy_sp.nzval, y, x, u, w)

d = eval(dyn_func[1])
@benchmark d($y1, $x1, $u1, $w1)
@benchmark discrete_dynamics($y1, $x1, $u1, $w1)

dy = eval(dyn_y_func[1])
@benchmark dy($y1, $x1, $u1, $w1)
@benchmark ForwardDiff.jacobian(a -> discrete_dynamics(a, $x1, $u1, $w1), $y1)
dy_sp = eval(dyn_y_func_sp[1])
@benchmark dy_sp($y1, $x1, $u1, $w1)

dx = eval(dyn_x_func[1])
@benchmark dx($y1, $x1, $u1, $w1)
@benchmark ForwardDiff.jacobian(a -> discrete_dynamics($y1, a, $u1, $w1), $x1)
dx_sp = eval(dyn_x_func_sp[1])
@benchmark dx_sp($y1, $x1, $u1, $w1)

du = eval(dyn_u_func[1])
@benchmark du($y1, $x1, $u1, $w1)
@benchmark ForwardDiff.jacobian(a -> discrete_dynamics($y1, $x1, a, $w1), $u1)
du_sp = eval(dyn_u_func_sp[1])
@benchmark du_sp($y1, $x1, $u1, $w1)

dxuy = eval(dyn_xuy_func[1])
@benchmark dxuy($y1, $x1, $u1, $w1)
dxuy_sp = eval(dyn_xuy_func_sp[1])
@benchmark dxuy_sp($y1, $x1, $u1, $w1)

dyn_1 = zeros(nx)
dyn_y1 = zeros(nx, nx)
dyn_x1 = zeros(nx, nx)
dyn_u1 = zeros(nx, nu)
dyn_xuy1 = zeros(nx, nx + nx + nu)

dyn_y1_sp = zeros(length(dyn_y_sp.nzval))
dyn_x1_sp = zeros(length(dyn_x_sp.nzval))
dyn_u1_sp = zeros(length(dyn_u_sp.nzval))
dyn_xuy1_sp = zeros(length(dyn_xuy_sp.nzval))

d! = eval(dyn_func[2])
@benchmark d!($dyn_1, $y1, $x1, $u1, $w1)

dy! = eval(dyn_y_func[2])
@benchmark dy!($dyn_y1, $y1, $x1, $u1, $w1)
@benchmark ForwardDiff.jacobian!($dyn_y1, a -> discrete_dynamics(a, $x1, $u1, $w1), $y1)
dy_sp! = eval(dyn_y_func_sp[2])
@benchmark dy_sp!($dyn_y1_sp, $y1, $x1, $u1, $w1)

dx! = eval(dyn_x_func[2])
@benchmark dx!($dyn_x1, $y1, $x1, $u1, $w1)
@benchmark ForwardDiff.jacobian!($dyn_x1, a -> discrete_dynamics($y1, a, $u1, $w1), $x1)
dx_sp! = eval(dyn_x_func_sp[2])
@benchmark dx_sp!($dyn_x1_sp, $y1, $x1, $u1, $w1)

du! = eval(dyn_u_func[2])
@benchmark du!($dyn_u1, $y1, $x1, $u1, $w1)
@benchmark ForwardDiff.jacobian!($dyn_u1, a -> discrete_dynamics($y1, $x1, a, $w1), $u1)
du_sp! = eval(dyn_u_func_sp[2])
@benchmark du_sp!($dyn_u1_sp, $y1, $x1, $u1, $w1)

dxuy! = eval(dyn_xuy_func[2])
@benchmark dxuy!($dyn_xuy1, $y1, $x1, $u1, $w1)
dxuy_sp! = eval(dyn_xuy_func_sp[2])
@benchmark dxuy_sp!($dyn_xuy1_sp, $y1, $x1, $u1, $w1)
dxuy_sp!(dyn_xuy1_sp, y1, x1, u1, w1)
dyn_xuy1_sp

@variables λ[1:nx]
dyn_λ = dot(λ, dyn)
dyn_λ_xuy = Symbolics.hessian(dyn_λ, [x; u; y])
dyn_λ_xuy_sp = Symbolics.sparsehessian(dyn_λ, [x; u; y])

dyn_λ_xuy_func = Symbolics.build_function(dyn_λ_xuy, y, x, u, w, λ)
dyn_λ_xuy_func_sp = Symbolics.build_function(dyn_λ_xuy_sp.nzval, y, x, u, w, λ)

λ1 = ones(nx)
dyn_λ_xuy1 = zeros(nx + nx + nu, nx + nx + nu)
dyn_λ_xuy1_sp = zeros(length(dyn_λ_xuy_sp.nzval))

dλ_xuy! = eval(dyn_λ_xuy_func[2])
@benchmark dλ_xuy!($dyn_λ_xuy1, $y1, $x1, $u1, $w1, $λ1)

dλ_xuy_sp! = eval(dyn_λ_xuy_func_sp[2])
@benchmark dλ_xuy_sp!($dyn_λ_xuy1_sp, $y1, $x1, $u1, $w1, $λ1)

# state constraint
function state_constraint(x)
    x - ones(nx)
end

# control constraint
function control_constraint(u)
    [u - ones(nu); ones(nu) - u]
end

# state control constraint
function state_control_constraint(x, u)

end

# indices
T = 10
nx_dim = [nx for t = 1:T]
nu_dim = [nu for t = 1:T-1]
nw_dim = [nw for t = 1:T-1]
nxu_dim = [[nx + nu for t = 1:T-1]..., nx]

idx_x = [collect((t > 1 ? sum(nx_dim[1:t-1]) + sum(nu_dim[1:t-1]) : 0) .+ (1:nx_dim[t])) for t = 1:T]
idx_u = [collect((t > 1 ? sum(nx_dim[1:t-1]) + sum(nu_dim[1:t-1]) : 0) + nx_dim[t] .+ (1:nu_dim[t])) for t = 1:T-1]
idx_xu = [collect((t > 1 ? sum(nx_dim[1:t-1]) + sum(nu_dim[1:t-1]) : 0) .+ (1:nx_dim[t] + nu_dim[t])) for t = 1:T-1]
idx_xuy = [collect((t > 1 ? sum(nx_dim[1:t-1]) + sum(nu_dim[1:t-1]) : 0) .+ (1:nx_dim[t] + nu_dim[t] + nx_dim[t+1])) for t = 1:T-1]

typeof(idx_x)
num_var = sum(nx_dim) + sum(nu_dim)
z1 = randn(num_var)

# objective
obj_grad = zeros(num_var)

xv = [z1[idx_x[t]] for t = 1:T]
uv = [[z1[idx_u[t]] for t = 1:T-1]..., zeros(0)]
wv = [zeros(nw) for t = 1:T-1]

objv = [0.0]
gradv = [[obj_grad[idx_xu[t]] for t = 1:T-1]..., obj_grad[idx_x[T]]]

function eval_obj!(J, x, u, T)
    for t = 1:T
        f!(J, x[t], u[t])
    end
    return J
end

eval_obj!(objv, xv, uv, T)
@benchmark eval_obj!($objv, $xv, $uv, $T)

function eval_obj_grad!(grad, x, u, T)
    for t = 1:T
        fxu!(grad[t], x[t], u[t])
    end
end

eval_obj_grad!(gradv, xv, uv, T)
@benchmark eval_obj_grad!($gradv, $xv, $uv, $T)

# discrete dynamics
nd_dim = nx_dim[2:T]
num_con = sum(nx_dim[2:T])
idx_dyn = [(t > 1 ? sum(nx_dim[1 .+ (1:t-1)]) : 0) .+ (1:nx_dim[t+1]) for t = 1:T-1]
con1 = zeros(num_con)
conv = [con1[idx_dyn[t]] for t = 1:T-1]

function eval_con!(con, x, u, w, T)
    for t = 1:T-1
        d!(con[t], x[t+1], x[t], u[t], w[t])
    end
end

d!(conv[t], xv[t+1], xv[t], uv[t], wv[t])
eval_con!(conv, xv, uv, wv, T)
@benchmark eval_con!($conv, $xv, $uv, $wv, $T)

function eval_con!(con, x, u, w, T)
    for t = 1:T-1
        d!(con[t], x[t+1], x[t], u[t], w[t])
    end
end

num_jac = length(dyn_xuy_sp.nzval) * (T - 1)
idx_jac = [(t - 1) * length(dyn_xuy_sp.nzval) .+ (1:length(dyn_xuy_sp.nzval)) for t = 1:T-1]
jac1 = zeros(num_jac)
jacv = [jac1[idx_jac[t]] for t = 1:T-1]

function eval_jac!(jac, x, u, w, T)
    for t = 1:T-1
        dxuy_sp!(jac[t], x[t+1], x[t], u[t], w[t])
    end
end

dxuy_sp!(jacv[t], xv[t+1], xv[t], uv[t], wv[t])
eval_jac!(jacv, xv, uv, wv, T)
@benchmark eval_jac!($jacv, $xv, $uv, $wv, $T)

I_jac, J_jac, V_jac = findnz(dyn_xuy_sp)

function sparsity_jac(jac, nd_dim, nxu_dim, T;
    row_shift = 0, col_shift = 0)

    row = Int[]
    col = Int[]

    I, J, V = findnz(jac)

    for t = 1:T-1
        push!(row, (I .+ row_shift)...)
        push!(col, (J .+ col_shift)...)

        row_shift += nd_dim[t]
        col_shift += nxu_dim[t]
    end

    return collect(zip(row, col))
end

sparsity_jac(dyn_xuy_sp, nd_dim, nxu_dim, T)


using Markdown
using InteractiveUtils

using Pkg, DrWatson

begin
    #@quickactivate "StatisticalRethinkingStan"
    using StanSample
    using StatisticalRethinking
end

begin
    N = 100
    df = DataFrame(
        :h0 => rand(Normal(10, 2 ), N),
        :treatment => vcat(zeros(Int, Int(N/2)), ones(Int, Int(N/2)))
    )
    df[!, :fungus] =
        [rand(Binomial.(1, 0.5 .- 0.4 .* df.treatment[i]), 1)[1] for i in 1:N]
    df[!, :h1] =
        [df[i, :h0] + rand(Normal(5 - 3 * df.fungus[i]), 1)[1] for i in 1:N]
    data = (N = nrow(df), h0 = df.h0, h1 = df.h1,
        fungus = df.fungus, treatment = df.treatment)
end;

stan6_6 = "
data {
  int <lower=1> N;
  vector[N] h0;
  vector[N] h1;
}
parameters{
  real<lower=0> p;
  real<lower=0> sigma;
}
transformed parameters {
  vector[N] mu;
  for ( i in 1:N ) {
    mu[i] = h0[i] * p;
  }
}

model {
  p ~ lognormal(0, 0.25);
  sigma ~ exponential(1);
  h1 ~ normal(mu, sigma);
}
generated quantities {
    vector[N] log_lik;
    for (i in 1:N)
        log_lik[i] = normal_lpdf(h1[i] | mu[i], sigma);
}
";

begin
    m6_6_data = Dict(
      :N => nrow(df),
      :h0 => df[:, :h0],
      :h1 => df[:, :h1]
    )
    m6_6s = SampleModel("m6.6s", stan6_6)
    rc6_6s = stan_sample(m6_6s; data)
end

stan6_7 = "
data {
  int <lower=1> N;
  vector[N] h0;
  vector[N] h1;
  vector[N] treatment;
  vector[N] fungus;
}
parameters{
  real a;
  real bt;
  real bf;
  real<lower=0> sigma;
}
transformed parameters {
  vector[N] mu;
  vector[N] p;
  for ( i in 1:N ) {
    p[i] = a + bt*treatment[i] + bf*fungus[i];
    mu[i] = h0[i] * p[i];
  }
}
model {
  a ~ lognormal(0, 0.2);
  bt ~ normal(0, 0.5);
  bf ~ normal(0, 0.5);
  sigma ~ exponential(1);
  h1 ~ normal(mu, sigma);
}
generated quantities {
    vector[N] log_lik;
    for ( i in 1:N )
        log_lik[i] = normal_lpdf( h1[i] | mu[i], sigma);
}
";

begin
    m6_7s = SampleModel("m6.7s", stan6_7)
    rc6_7s = stan_sample(m6_7s; data)
end

if success(rc6_7s)
    waic(m6_7s)
end

if success(rc6_7s)
    post6_7s_df = read_samples(m6_7s, :dataframe)
    b6_7s = post6_7s_df[:, [:a, :bt, :bf, :sigma]]
    p6_7s = b6_7s.a .+ b6_7s.bt * df.treatment' + b6_7s.bf * df.fungus'
    mu6_7s = p6_7s .* df.h0'
    log_lik = logpdf.(Normal.(mu6_7s, post6_7s_df.sigma),  df.h1')
    waic(log_lik)
end

stan6_8 = "
data {
  int <lower=1> N;
  vector[N] h0;
  vector[N] h1;
  vector[N] treatment;
}
parameters{
  real a;
  real bt;
  real<lower=0> sigma;
}
transformed parameters {
  vector[N] mu;
  vector[N] p;
  for ( i in 1:N ) {
    p[i] = a + bt*treatment[i];
    mu[i] = h0[i] * p[i];
  }
}
model {
  a ~ lognormal(0, 0.2);
  bt ~ normal(0, 0.5);
  sigma ~ exponential(1);
  h1 ~ normal(mu, sigma);
}
generated quantities {
    vector[N] log_lik;
    for (i in 1:N)
        log_lik[i] = normal_lpdf(h1[i] | mu[i], sigma);
}
";

begin
    m6_8s = SampleModel("m6.8s", stan6_8)
    rc6_8s = stan_sample(m6_8s; data)
end

# ╔═╡ 425a70dc-6bb3-11eb-1a7e-39f3133b7173
if success(rc6_6s) && success(rc6_7s) && success(rc6_8s)
    df_waic = compare([m6_6s, m6_7s, m6_8s], :waic)
end

# ╔═╡ 904b7004-6d69-11eb-07fe-65b5bb8c4dd2
if success(rc6_6s) && success(rc6_7s) && success(rc6_8s)
    df_psis = compare([m6_6s, m6_7s, m6_8s], :psis)
end

# ╔═╡ ee0899d8-68e1-11eb-2679-e304ef15e9e4
begin
    loo6_7s, loos6_7s, pk6_7s = psisloo(m6_7s)
    pk_plot(pk6_7s)
end


abstract type Simplexer end

struct AffineSimplexer <: Simplexer
    a::Float64
    b::Float64
end
AffineSimplexer(;a = 0.025, b = 0.5) = AffineSimplexer(a, b)

function simplexer(S::AffineSimplexer, initial_x::Tx) where Tx
    n = length(initial_x)
    initial_simplex = Tx[copy(initial_x) for i = 1:n+1]
    for j = 1:n
        initial_simplex[j+1][j] = (1+S.b) * initial_simplex[j+1][j] + S.a
    end
    initial_simplex
end

abstract type NMParameters end

struct AdaptiveParameters <: NMParameters
    α::Float64
    β::Float64
    γ::Float64
    δ::Float64
end

AdaptiveParameters(;  α = 1.0, β = 1.0, γ = 0.75 , δ = 1.0) = AdaptiveParameters(α, β, γ, δ)
parameters(P::AdaptiveParameters, n::Integer) = (P.α, P.β + 2/n, P.γ - 1/2n, P.δ - 1/n)

struct FixedParameters <: NMParameters
    α::Float64
    β::Float64
    γ::Float64
    δ::Float64
end

FixedParameters(; α = 1.0, β = 2.0, γ = 0.5, δ = 0.5) = FixedParameters(α, β, γ, δ)
parameters(P::FixedParameters, n::Integer) = (P.α, P.β, P.γ, P.δ)

struct NelderMead{Ts <: Simplexer, Tp <: NMParameters} <: ZerothOrderOptimizer
    initial_simplex::Ts
    parameters::Tp
end

"""
# NelderMead
## Constructor
```julia
NelderMead(; parameters = AdaptiveParameters(),
             initial_simplex = AffineSimplexer())
```

The constructor takes 2 keywords:

* `parameters`, an instance of either `AdaptiveParameters` or `FixedParameters`,
and is used to generate parameters for the Nelder-Mead Algorithm
* `initial_simplex`, an instance of `AffineSimplexer`

## Description
Our current implementation of the Nelder-Mead algorithm is based on [1] and [3].
Gradient-free methods can be a bit sensitive to starting values and tuning parameters,
so it is a good idea to be careful with the defaults provided in Optim.jl.

Instead of using gradient information, Nelder-Mead is a direct search method. It keeps
track of the function value at a number of points in the search space. Together, the
points form a simplex. Given a simplex, we can perform one of four actions: reflect,
expand, contract, or shrink. Basically, the goal is to iteratively replace the worst
point with a better point. More information can be found in [1], [2] or [3].

## References
- [1] Nelder, John A. and R. Mead (1965). "A simplex method for function minimization". Computer Journal 7: 308–313. doi:10.1093/comjnl/7.4.308
- [2] Lagarias, Jeffrey C., et al. "Convergence properties of the Nelder–Mead simplex method in low dimensions." SIAM Journal on Optimization 9.1 (1998): 112-147
- [3] Gao, Fuchang and Lixing Han (2010). "Implementing the Nelder-Mead simplex algorithm with adaptive parameters". Computational Optimization and Applications. doi:10.1007/s10589-010-9329-3
"""
function NelderMead(; kwargs...)
    KW = Dict(kwargs)
    if haskey(KW, :initial_simplex) || haskey(KW, :parameters)
        initial_simplex, parameters = AffineSimplexer(), AdaptiveParameters()
        haskey(KW, :initial_simplex) && (initial_simplex = KW[:initial_simplex])
        haskey(KW, :parameters) && (parameters = KW[:parameters])
        return NelderMead(initial_simplex, parameters)
    else
        return NelderMead(AffineSimplexer(), AdaptiveParameters())
    end
end

Base.summary(::NelderMead) = "Nelder-Mead"

# centroid except h-th vertex
function centroid!(c::AbstractArray{T}, simplex, h=0) where T
    n = length(c)
    fill!(c, zero(T))
    @inbounds for i in 1:n+1
        if i != h
            xi = simplex[i]
            c .+= xi
        end
    end
    rmul!(c, T(1)/n)
end

centroid(simplex, h) = centroid!(similar(simplex[1]), simplex, h)

nmobjective(y::Vector, m::Integer, n::Integer) = sqrt(var(y) * (m / n))

function print_header(method::NelderMead)
    @printf "Iter     Function value    √(Σ(yᵢ-ȳ)²)/n \n"
    @printf "------   --------------    --------------\n"
end

function Base.show(io::IO, trace::OptimizationTrace{T, NelderMead}) where T
    @printf io "Iter     Function value    √(Σ(yᵢ-ȳ)²)/n \n"
    @printf io "------   --------------    --------------\n"
    for state in trace.states
        show(io, state)
    end
    return
end

function Base.show(io::IO, t::OptimizationState{NelderMead})
    @printf io "%6d   %14e    %14e\n" t.iteration t.value t.g_norm
    if !isempty(t.metadata)
        for (key, value) in t.metadata
            @printf io " * %s: %s\n" key value
        end
    end
    return
end

mutable struct NelderMeadState{Tx, T, Tfs} <: ZerothOrderState
    x::Tx
    m::Int
    simplex::Vector{Tx}
    x_centroid::Tx
    x_lowest::Tx
    x_second_highest::Tx
    x_highest::Tx
    x_reflect::Tx
    x_cache::Tx
    f_simplex::Tfs
    nm_x::T
    f_lowest::T
    i_order::Vector{Int}
    α::T
    β::T
    γ::T
    δ::T
    step_type::String
end

function initial_state(method::NelderMead, options, d, initial_x)
    T = eltype(initial_x)
    n = length(initial_x)
    m = n + 1
    simplex = simplexer(method.initial_simplex, initial_x)
    f_simplex = zeros(T, m)

    value!!(d, first(simplex))

    @inbounds for i in 1:length(simplex)
        f_simplex[i] = value(d, simplex[i])
    end
    # Get the indices that correspond to the ordering of the f values
    # at the vertices. i_order[1] is the index in the simplex of the vertex
    # with the lowest function value, and i_order[end] is the index in the
    # simplex of the vertex with the highest function value
    i_order = sortperm(f_simplex)

    α, β, γ, δ = parameters(method.parameters, n)

NelderMeadState(similar(initial_x), # Variable to hold final minimizer value for MultivariateOptimizationResults
          m, # Number of vertices in the simplex
          simplex, # Maintain simplex in state.simplex
          centroid(simplex,  i_order[m]), # Maintain centroid in state.centroid
          similar(initial_x), # Store cache in state.x_lowest
          similar(initial_x), # Store cache in state.x_second_highest
          similar(initial_x), # Store cache in state.x_highest
          similar(initial_x), # Store cache in state.x_reflect
          similar(initial_x), # Store cache in state.x_cache
          f_simplex, # Store objective values at the vertices in state.f_simplex
          T(nmobjective(f_simplex, n, m)), # Store nmobjective in state.nm_x
          f_simplex[i_order[1]], # Store lowest f in state.f_lowest
          i_order, # Store a vector of rankings of objective values
          T(α),
          T(β),
          T(γ),
          T(δ),
          "initial")
end

function update_state!(f::F, state::NelderMeadState{T}, method::NelderMead) where {F, T}
    # Augment the iteration counter
    shrink = false
    n, m = length(state.x), state.m

    centroid!(state.x_centroid, state.simplex, state.i_order[m])
    copyto!(state.x_lowest, state.simplex[state.i_order[1]])
    copyto!(state.x_second_highest, state.simplex[state.i_order[n]])
    copyto!(state.x_highest, state.simplex[state.i_order[m]])
    state.f_lowest = state.f_simplex[state.i_order[1]]
    f_second_highest = state.f_simplex[state.i_order[n]]
    f_highest = state.f_simplex[state.i_order[m]]

    # Compute a reflection
    @inbounds for j in 1:n
        state.x_reflect[j] = state.x_centroid[j] + state.α * (state.x_centroid[j]-state.x_highest[j])
    end

    f_reflect = value(f, state.x_reflect)
    if f_reflect < state.f_lowest
        # Compute an expansion
        @inbounds for j in 1:n
            state.x_cache[j] = state.x_centroid[j] + state.β *(state.x_reflect[j] - state.x_centroid[j])
        end
        f_expand = value(f, state.x_cache)

        if f_expand < f_reflect
            copyto!(state.simplex[state.i_order[m]], state.x_cache)
            @inbounds state.f_simplex[state.i_order[m]] = f_expand
            state.step_type = "expansion"
        else
            copyto!(state.simplex[state.i_order[m]], state.x_reflect)
            @inbounds state.f_simplex[state.i_order[m]] = f_reflect
            state.step_type = "reflection"
        end
        # shift all order indeces, and wrap the last one around to the first
        i_highest = state.i_order[m]
        @inbounds for i = m:-1:2
            state.i_order[i] = state.i_order[i-1]
        end
        state.i_order[1] = i_highest
    elseif f_reflect < f_second_highest
        copyto!(state.simplex[state.i_order[m]], state.x_reflect)
        @inbounds state.f_simplex[state.i_order[m]] = f_reflect
        state.step_type = "reflection"
        sortperm!(state.i_order, state.f_simplex)
    else
        if f_reflect < f_highest
            # Outside contraction
            @simd for j in 1:n
                @inbounds state.x_cache[j] = state.x_centroid[j] + state.γ * (state.x_reflect[j]-state.x_centroid[j])
            end
            f_outside_contraction = value(f, state.x_cache)
            if f_outside_contraction < f_reflect
                copyto!(state.simplex[state.i_order[m]], state.x_cache)
                @inbounds state.f_simplex[state.i_order[m]] = f_outside_contraction
                state.step_type = "outside contraction"
                sortperm!(state.i_order, state.f_simplex)

            else
                shrink = true
            end
        else # f_reflect > f_highest
            # Inside constraction
            @simd for j in 1:n
                @inbounds state.x_cache[j] = state.x_centroid[j] - state.γ *(state.x_reflect[j] - state.x_centroid[j])
            end
            f_inside_contraction = value(f, state.x_cache)
            if f_inside_contraction < f_highest
                copyto!(state.simplex[state.i_order[m]], state.x_cache)
                @inbounds state.f_simplex[state.i_order[m]] = f_inside_contraction
                state.step_type = "inside contraction"
                sortperm!(state.i_order, state.f_simplex)
            else
                shrink = true
            end
        end
    end

    if shrink
        for i = 2:m
            ord = state.i_order[i]
            copyto!(state.simplex[ord], state.x_lowest + state.δ*(state.simplex[ord]-state.x_lowest))
            state.f_simplex[ord] = value(f, state.simplex[ord])
        end
        step_type = "shrink"
        sortperm!(state.i_order, state.f_simplex)
    end

    state.nm_x = nmobjective(state.f_simplex, n, m)
    false
end

function after_while!(f, state, method::NelderMead, options)
    sortperm!(state.i_order, state.f_simplex)
    x_centroid_min = centroid(state.simplex, state.i_order[state.m])
    f_centroid_min = value(f, x_centroid_min)
    f_min, i_f_min = findmin(state.f_simplex)
    x_min = state.simplex[i_f_min]
    if f_centroid_min < f_min
        x_min = x_centroid_min
        f_min = f_centroid_min
    end
    f.F = f_min
    state.x .= x_min
end
# We don't have an f_x_previous in NelderMeadState, so we need to special case these
pick_best_x(f_increased, state::NelderMeadState) = state.x
pick_best_f(f_increased, state::NelderMeadState, d) = value(d)

function assess_convergence(state::NelderMeadState, d, options)
    g_converged = state.nm_x <= options.g_tol # Hijact g_converged for NM stopping criterior
    return false, false, g_converged, g_converged, false
end

function initial_convergence(d, state::NelderMeadState, method::NelderMead, initial_x, options)
    nmobjective(state.f_simplex, state.m, length(initial_x)) < options.g_tol
end

function trace!(tr, d, state, iteration, method::NelderMead, options)
    dt = Dict()
    if options.extended_trace
        dt["centroid"] = copy(state.x_centroid)
        dt["step_type"] = state.step_type
    end
    update!(tr,
    iteration,
    state.f_lowest,
    state.nm_x,
    dt,
    options.store_trace,
    options.show_trace,
    options.show_every,
    options.callback)
end


using EquivalenceClassesConstructor
using Test

using Printf
using DataStructures
using JLD
using Lazy
include("../src/Types.jl")
include("../src/GraphTools.jl")
include("../src/helpers.jl")
include("../src/IO.jl")

pred1 = Predicate((x,y) -> x == y)
pred2 = Predicate((x,y) -> x != y)
pred3 = Predicate((x,y) -> all(x .== -y))
pred4 = Predicate((x,y) -> x == y+3)

m1 = Mapping(x -> [x])
m2 = Mapping(x -> [-x])
vl1 = collect(1:5)
vl2 = collect(-3:5)
cl1 = find_classes(m1, vl1)
cl2 = find_classes(m2, vl2)
cl2_res_index  = Dict(3 => UInt32(1), -3 => UInt32(1), 2 => UInt32(2), -2=> UInt32(2), 1 => UInt32(3), -1 => UInt32(3), 0 => UInt32(4), 4 => UInt32(8), 5 => UInt32(9))
cl2_res_direct = Dict(3 => -3, -3 => -3, 2 => -2, -2=>-2, 1 => -1, -1 => -1, 0 => 0, 4 => 4, 5 => 5)

@testset "EquivalenceClassesConstructor.jl" begin
    @testset "helpers" begin
        include("./helpers.jl")
    end
    @testset "Types" begin
        include("./Types.jl")
    end
    @testset "graph" begin
        include("./GraphTools.jl")
    end
    @testset "IO" begin
        include("./IO.jl")
    end
end


using MethodOfParticularSolutions
using Nemo
using Test
using Printf

include("sphericaltriangle.jl")
include("lshape.jl")


@testset "Manopt Cyclic Proximal Point" begin
    using Dates
    n = 100
    N = PowerManifold(Circle(), n)
    f = artificial_S1_signal(n)
    F = x -> costL2TV(N, f, 0.5, x)
    proxes = ((λ, x) -> prox_distance(N, λ, f, x), (λ, x) -> prox_TV(N, 0.5 * λ, x))
    o = cyclic_proximal_point(
        N,
        F,
        proxes,
        f;
        λ=i -> π / (2 * i),
        stopping_criterion=StopAfterIteration(100),
        debug=[
            DebugIterate(), " ", DebugCost(), " ", DebugProximalParameter(), "\n", 10000
        ],
        record=[RecordProximalParameter(), RecordIterate(f), RecordCost()],
        return_options=true,
    )
    fR = get_solver_result(o)
    fR2 = cyclic_proximal_point(
        N, F, proxes, f; λ=i -> π / (2 * i), stopping_criterion=StopAfterIteration(100)
    )
    @test fR == fR2
    rec = get_record(o)
    @test F(f) > F(fR)
    #
    o = CyclicProximalPointOptions(f, StopAfterIteration(1), i -> π / (2 * i))
    p = ProximalProblem(N, F, proxes, [1, 2])
    @test_throws ErrorException get_proximal_map(p, 1.0, f, 3)
    @test_throws ErrorException ProximalProblem(N, F, proxes, [1, 2, 2])
    order1 = collect(1:3)
    Manopt.update_cpp_order!(order1, 3, 0, RandomEvalOrder())
    @test all(isinteger.(order1))
    @test minimum(order1) == 1
    @test maximum(order1) == 3

    # i=0 introduce new fixed ranom order
    order2 = collect(1:3)
    Manopt.update_cpp_order!(order2, 3, 0, FixedRandomEvalOrder())
    @test all(isinteger.(order2))
    @test minimum(order2) == 1
    @test maximum(order2) == 3
    order3 = copy(order2)
    Manopt.update_cpp_order!(order3, 3, 1, FixedRandomEvalOrder())
    @test order3 == order2 # only update on i=0

    @test Manopt.update_cpp_order!(collect(1:3), 3, 0, LinearEvalOrder()) == 1:3
    @test Manopt.update_cpp_order!(collect(1:3), 3, 1, LinearEvalOrder()) == 1:3
end


@testset "Test the LS version of the OMP 0 relative error" begin

    ψ = randn(100, 20)
    u = randn(100)
    m, n = size(ψ)

    ctrue = ψ\u

    # Force a greedy approach over the entire set of indices
    # With LS
    idxls, cls, ϵls = lsomp(ψ, u; invert = true, verbose = true, ϵrel = eps())

    @test norm(ctrue[idxls]-cls)<1e-14

    # Residual error must decrease
    @test all(ϵls[2:end]-ϵls[1:end-1] .< 0) == true

    for k=1:n
        cverif = view(ψ,:, idxls[1:k])\u
        @test abs(ϵls[k+1] - norm(u - view(ψ,:, idxls[1:k])*cverif))<1e-14
    end
end

@testset "Test the LS version of the OMP given relative error" begin
    ψ = randn(100, 20)
    u = randn(100)
    m, n = size(ψ)

    ctrue = ψ\u

    ϵrel = 0.1*1.00 + 0.9*norm(u - ψ*ctrue)/norm(u)

    # Force a greedy approach over the entire set of indices
    # With LS
    idxls, cls, ϵls = lsomp(ψ, u; invert = true, verbose = true, ϵrel = ϵrel)

    # Residual error must decrease
    @test all(ϵls[2:end]-ϵls[1:end-1] .< 0) == true

    @test ϵls[end]<ϵrel*norm(u)
    @test ϵls[end-1]>ϵrel*norm(u)

    for k=1:size(idxls,1)
        cverif = view(ψ,:, idxls[1:k])\u
        @test abs(ϵls[k+1] - norm(u - view(ψ,:, idxls[1:k])*cverif))<1e-14
    end
end

@testset "Test the LS version of the OMP given a relative number of features" begin
    ψ = randn(100, 20)
    u = randn(100)
    m, n = size(ψ)
    maxterms = 10
    ctrue = ψ\u

    # Force a greedy approach over the entire set of indices
    # With LS
    idxls, cls, ϵls = lsomp(ψ, u; invert = true, verbose = true, maxterms = maxterms)

    # Residual error must decrease
    @test all(ϵls[2:end]-ϵls[1:end-1] .< 0) == true

    @test size(idxls, 1) == maxterms
    @test size(cls, 1) == maxterms
    @test size(ϵls, 1) == maxterms + 1

    for k=1:size(idxls,1)
        cverif = view(ψ,:, idxls[1:k])\u
        @test abs(ϵls[k+1] - norm(u - view(ψ,:, idxls[1:k])*cverif))<1e-14
    end
end

@testset "Test the QR version of the OMP" begin

    ψ = randn(100, 20)
    u = randn(100)
    m, n = size(ψ)
    ctrue = ψ\u

    # Force a greedy approach over the entire set of indices
    # With QR
    idxqr, cqr, ϵqr = qromp(ψ, u; invert = true, verbose = true, ϵrel = eps())

    @test norm(ctrue[idxqr]-cqr)<1e-14

    # Residual error must decrease
    @test all(ϵqr[2:end]-ϵqr[1:end-1] .< 0) == true

    for k=1:n
        cverif = view(ψ,:, idxqr[1:k])\u
        @test abs(ϵqr[k+1] - norm(u - view(ψ,:, idxqr[1:k])*cverif))<1e-14
    end
end

@testset "Test the QR version of the OMP given relative error" begin
    ψ = randn(100, 20)
    u = randn(100)
    m, n = size(ψ)

    ctrue = ψ\u

    ϵrel = 0.1*1.00 + 0.9*norm(u - ψ*ctrue)/norm(u)

    # Force a greedy approach over the entire set of indices
    # With QR
    idxqr, cqr, ϵqr = qromp(ψ, u; invert = true, verbose = true, ϵrel = ϵrel)

    # Residual error must decrease
    @test all(ϵqr[2:end]-ϵqr[1:end-1] .< 0) == true

    @test ϵqr[end]<ϵrel*norm(u)
    @test ϵqr[end-1]>ϵrel*norm(u)

    for k=1:size(idxqr,1)
        cverif = view(ψ,:, idxqr[1:k])\u
        @test abs(ϵqr[k+1] - norm(u - view(ψ,:, idxqr[1:k])*cverif))<1e-14
    end
end

@testset "Test the QR version of the OMP given a number of features" begin
    ψ = randn(100, 20)
    u = randn(100)
    m, n = size(ψ)

    maxterms = 10

    ctrue = ψ\u

    ϵrel = 0.1*1.00 + 0.9*norm(u - ψ*ctrue)/norm(u)

    # Force a greedy approach over the entire set of indices
    # With QR
    idxqr, cqr, ϵqr = qromp(ψ, u; invert = true, verbose = true, maxterms = maxterms)

    # Residual error must decrease
    @test all(ϵqr[2:end]-ϵqr[1:end-1] .< 0) == true

    @test size(idxqr, 1) == maxterms
    @test size(cqr, 1) == maxterms
    @test size(ϵqr, 1) == maxterms + 1

    for k=1:size(idxqr,1)
        cverif = view(ψ,:, idxqr[1:k])\u
        @test abs(ϵqr[k+1] - norm(u - view(ψ,:, idxqr[1:k])*cverif))<1e-14
    end
end

@testset "Check that LS and QR versions are consistent" begin

    ψ = randn(100, 20)
    u = randn(100)
    ctrue = ψ\u

    # Force a greedy approach over the entire set of indices
    # With LS
    idxls, cls, ϵls = lsomp(ψ, u; invert = true, verbose = true, ϵrel = eps())

    # With QR
    idxqr, cqr, ϵqr = qromp(ψ, u; invert = true, verbose = true, ϵrel = eps())

    @test norm(ctrue[idxls]-cls)<1e-14
    @test norm(ctrue[idxqr]-cqr)<1e-14
    # @test all(ϵls - ϵqr .>= 0) == true
end


function simulated_annealing(fcn::Function,
                             x0::Array,
                             args...;
                             neighbor!::Function   = Optim.default_neighbor!,
                             iterations::Int       = 1000,
                             store_trace::Bool     = false,
                             show_trace::Bool      = false,
                             extended_trace::Bool  = false,
                             temperature::Function = Optim.log_temperature,
                             kwargs...)

    Optim.optimize(fcn, x0,
                   method = SimulatedAnnealing(neighbor = neighbor!,temperature = temperature),
                   iterations = iterations, store_trace = store_trace, show_trace = show_trace,
                   extended_trace = extended_trace)

end

function log_temperature(t::Real; initial_temperature::Real = 1.0)
    return initial_temperature/log(1+t)
end

function exponential_temperature(t::Real; initial_temperature::Real = 1.0, α::Real = .99)
    if !( 0.0 < α < 1.0)
        error("α must be in (0,1)")
    end
    return initial_temperature*α^t
end

function linear_temperature(t::Real;  initial_temperature::Real = 1, β = .005)
    return max(initial_temperature - t*β, 0.0)
end