# Load the NPCF code
println("Loading NPCF module")
using NPCFs, Test, Distributed
Npart = 50 # number of particles
boxsize = 200 # input size

## Single thread runtest, comparing NPCF simple + pairwise outputs for different N and D values
println("Single-worker runtest")
for N in 2:5
    for D in 2:4
        println("Testing Cartesian N=$N, D=$D")
        npcf = NPCF(N=N,D=D,verb=true,periodic=true,volume=boxsize^D,coords="cartesian",r_min=20,r_max=95,nbins=6,lmax=1);

        pos = hcat(rand(Npart,D)*boxsize,ones(Npart)); # columns: [x, y, z, weight]

        npcf_output1 = compute_npcf_simple(pos, npcf);
        npcf_output2 = compute_npcf_pairwise(pos, npcf);

        # Compute the error
        err = sum(npcf_output1[:]-npcf_output2[:])/sum(npcf_output2[:]);
        println("Error is $err\n")

        @test abs(err)<1e-10
    end
end

## As above, but for spherical geometry in 2 dimensions
println("Spherical geometry runtest")
for N in 2:5
    println("Testing spherical N=$N, D=2")
    npcf = NPCF(N=N,D=2,verb=true,periodic=false,volume=4pi,coords="spherical",r_min=-0.8,r_max=0.8,nbins=6,lmax=1);

    phi_arr = rand(Npart)*2pi
    theta_arr = acos.(rand(Npart).*2 .-1)
    pos = hcat(phi_arr,theta_arr,ones(Npart));

    npcf_output1 = compute_npcf_simple(pos, npcf);
    npcf_output2 = compute_npcf_pairwise(pos, npcf);

    # Compute the error
    err = sum(npcf_output1[:]-npcf_output2[:])/sum(npcf_output2[:]);
    println("Error is $err\n")

    @test abs(err)<1e-10
end

## As above, but for an aperiodic geometry
println("Aperiodic box runtest")
for N in 2:5
    for D in 2:4
        println("Testing Cartesian N=$N, D=$D")
        npcf = NPCF(N=N,D=D,verb=true,periodic=false,volume=boxsize^D,coords="cartesian",r_min=20,r_max=95,nbins=6,lmax=1);

        pos = hcat(rand(Npart,D)*boxsize,ones(Npart)); # columns: [x, y, z, weight]

        npcf_output1 = compute_npcf_simple(pos, npcf);
        npcf_output2 = compute_npcf_pairwise(pos, npcf);

        # Compute the error
        err = sum(npcf_output1[:]-npcf_output2[:])/sum(npcf_output2[:]);
        println("Error is $err\n")

        @test abs(err)<1e-10
    end
end

## Multiple worker runtest
# Here we use N = D = 3, and run on two workers
println("Multi-worker runtest")
addprocs(2) # run on two workers
@everywhere using NPCFs
npcf = NPCF(N=3,D=3,verb=true,periodic=true,volume=boxsize^3,coords="cartesian",r_min=20,r_max=95,nbins=6,lmax=1);

pos = hcat(rand(Npart,3)*boxsize,ones(Npart)); # columns: [x, y, z, weight]

npcf_output1 = compute_npcf_simple(pos, npcf);
npcf_output2 = compute_npcf_pairwise(pos, npcf);

# Compute the error
err = sum(npcf_output1[:]-npcf_output2[:])/sum(npcf_output2[:]);
println("Error is $err\n")

@test abs(err)<1e-10

println("All tests passed!")


import Compat.LinearAlgebra: RealHermSymComplexHerm, normalize
import Base: reverse

convert(::Type{BitSet},r::IntCol) = BitSet(r)
BitSet(elems::Vararg{Int}) = BitSet(elems)

reverse(n::Number) = n

tname(T::Type) = T.name.name

function randomsmooth(n::Integer,m::Integer)
    R = Matrix{Float64}(undef,n,m)
    for i = 1:m
        R[:,i] = randomsmooth(n)
    end
    return R
end

function randomsmooth(n::Integer)
    x = transpose(0:1/(n-1):1)
    N = rand(4:7) # number of sines to use
    f = 2.0 .+ 10 .* rand(N) # frequencies
    A = normalize!((rand(N).-0.5)./f,2) # amplitudes
    ϕ = 2π.*rand(N) # phase
    return vec(sum(A.*sin.(f.*x.+ϕ),1))
end

gaussian(x::Real,σ::Real) = exp(-0.5*(x/σ)^2)
gaussianprime(x::Real,σ::Real) = -exp(-0.5*(x/σ)^2)*x/σ^2

function gaussian_fwhm(x::Real,w)
    # w is the full width half maximum (FWHM)
    return exp(-4ln2*(x/w)^2)
end

function gaussianprime_fwhm(x::Real,w)
    # w is the full width half maximum (FWHM)
    return -x*ln2*exp2(3-4(x/w)^2)/w^2
end

function sqrtfact(n::Integer)
    if 0 <= n < 128
        return sqrtfact_table[n+1]
    elseif n >= 128 # Stirling series for values not stored in the table
        return sqrt(sqrt(2π*n)*(n/ℯ)^n*(1.0+1/12n+1/288n^2))
    else
        throw(ArgumentError("n must be larger than or equal to 0"))
    end
end

# pairwise sum of squared differences (SSD)
function norm2_diff(A::Vector{T},B::Vector{T}) where {T}
    n = length(A)
    n == length(B) || throw(DimensionMismatch())
    n == 0 ? abs2(zero(T)-zero(T)) : norm2_diff(A, B, 1, n)
end

function norm2_diff(A::Vector, B::Vector, i1::Integer, n::Integer)
    if n < 128
        @inbounds s = abs2(A[i1] - B[i1])
        for i in i1+1 : i1+n-1
            @inbounds s += abs2(A[i] - B[i])
        end
        return s
    else
        n2 = n÷2
        return norm2_diff(A, B, i1, n2) + norm2_diff(A, B, i1+n2, n-n2)
    end
end

normalize(z::Complex) = z == 0 ? one(z) : z/abs(z)
normalize(z::Real) = one(z)

function unwrap!(p)
    n,m = size(p,1),size(p,2)
    n < 2 && return p
    for j = 1:m, i = 2:n
        d = p[i,j] - p[i-1,j]
        if abs(d) > π
            p[i,j] -= floor((d+π)/2π) * 2π
        end
    end
    return p
end

expim(H::RealHermSymComplexHerm) = expim!(Matrix{complex(eltype(H))}(H),copy(H))
expim!(H::RealHermSymComplexHerm) = expim!(Matrix{complex(eltype(H))}(H),H)

function expim!(R::Matrix,H::RealHermSymComplexHerm,Λ=Vector{real(eltype(R))}(undef,size(H,1)),U=Matrix(H),B=similar(R))
    # First decompose H into U*Λ*Uᴴ
    #F = eigfact!(H); copy!(Λ,F.values); copy!(U,F.vectors)
    hermfact!(Λ,U,H)
    # Calculate the imaginary exponential of each eigenvalue and multiply
    # each column of U to obtain B = U*exp(iΛ)
    n = length(Λ)
    @inbounds for j = 1:n
        a = cis(Λ[j])
        @simd for i = 1:n
            B[i,j] = a * U[i,j]
        end
    end
    # Finally multiply B by Uᴴ to obtain U*exp(iΛ)*Uᴴ = exp(i*H)
    @static if VERSION < v"0.7.0-"
        A_mul_Bc!(R,B,U)
    else
        mul!(R,B,adjoint(U))
    end
    return R
end

excited_sys(A::QuObject, subs::Vararg{Integer}) = excited_sys(dims(A), sys)
excited_sys(::Union{Dims{N},Val{N}}, subs::Vararg{Integer}) where N =
    ntuple(i->ifelse(i ∈ subs, 1, 0), Val(N))

lengthrange(a::Real,len::Integer) = UnitRange{typeof(a)}(a, oftype(a, a+len-1))

tensored_iterator(dims::Dims) = (revtuple(t) for t ∈ product(lengthrange.(0,revtuple(dims))...))


function test_cat()
  ϵ = eps(Float64)
  rtol = sqrt(ϵ)
  @testset ExtendedTestSet "Concatenation" begin
    A = simple_sparse_matrix(Float64, 100, 100)
    B = simple_matrix(Float64, 100, 10)
    C = simple_matrix(Float64, 100, 90)
    D = [A B C]
    Ao = LinearOperator(A)
    Bo = LinearOperator(B)
    Co = LinearOperator(C)
    Do = [Ao Bo Co]
    Do2 = LinearOperator(D)
    rhs = simple_vector(Float64, 200)

    @test(norm(Do * rhs - D * rhs) <= rtol * norm(D * rhs))  # throws InexactError()
    @test(norm(Do2 * rhs - D * rhs) <= rtol * norm(D * rhs))

    @test_throws LinearOperatorException [LinearOperator(ones(5, 5)) opEye(3)]

    # test vcat
    A = simple_sparse_matrix(Float64, 100, 100)
    B = simple_matrix(Float64, 10, 100)
    C = simple_matrix(Float64, 90, 100)
    D = [A; B; C]
    Ao = LinearOperator(A)
    Bo = LinearOperator(B)
    Co = LinearOperator(C)
    Do = [Ao; Bo; Co]
    Do2 = LinearOperator(D)
    rhs = simple_vector(Float64, 100)

    @test(norm(Do * rhs - D * rhs) <= rtol * norm(D * rhs))
    @test(norm(Do2 * rhs - D * rhs) <= rtol * norm(D * rhs))

    @test_throws LinearOperatorException [LinearOperator(ones(5, 5)); opEye(3)]
    K = [opEye(2) opZeros(2, 3); opZeros(3, 2) opEye(3)]
    v = simple_vector(Float64, 5)
    @test all(v .== K * v)

    K = [opEye(2); sparse(1.0I, 2, 2)]
    v = simple_vector(Float64, 2)
    @test all(K * v .== [v; v])
  end
end

test_cat()


using DynamicHMCModels

ProjDir = @__DIR__

df = DataFrame(CSV.read(joinpath(ProjDir, "..", "..", "data",  "Kline.csv"), delim=';'))

# New col logpop, set log() for population data
df[!, :logpop] = map((x) -> log(x), df[:, :population]);
df[!, :society] = 1:10;

Base.@kwdef mutable struct KlineModel{Ty <: AbstractVector,
  Tx <: AbstractMatrix, Ts <: AbstractVector}
    "Observations (total_tools)."
    y::Ty
    "Covariates (logpop)"
    x::Tx
    "Society"
    s::Ts
    "Number of observations (10)"
    N::Int
    "Number of societies (also 10)"
    N_societies::Int
end

function make_transformation(model::KlineModel)
    as( (β = as(Array, size(model.x, 2)), α = as(Array, model.N_societies), σ = asℝ₊) )
end

# Instantiate the model with data and inits.

N = size(df, 1)
N_societies = length(unique(df[:, :society]))
x = hcat(ones(Int64, N), df[:, :logpop]);
s = df[:, :society]
y = df[:, :total_tools]
model = KlineModel(; y=y, x=x, s=s, N=N, N_societies=N_societies)

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

function (model::KlineModel)(θ)
    @unpack y, x, s, N, N_societies = model   # data
    @unpack β, α, σ = θ  # parameters
    ll = 0.0
    ll += logpdf(Cauchy(0, 1), σ)
    ll += sum(logpdf.(Normal(0, σ), α)) # α[1:10]
    ll += logpdf.(Normal(0, 10), β[1]) # a
    ll += logpdf.(Normal(0, 1), β[2]) # a
    ll += sum(
      [loglikelihood(Poisson(exp(α[s[i]] + dot(x[i, :], β))), [y[i]]) for i in 1:N]
    )
    ll
end

println()
θ = (β = [1.0, 0.25], α = rand(Normal(0, 1), N_societies), σ = 0.2)
model(θ) |> display
println()

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

P = TransformedLogDensity(make_transformation(model), model)
∇P = ADgradient(:ForwardDiff, P);
results = mcmc_with_warmup(Random.GLOBAL_RNG, ∇P, 1000)
posterior = P.transformation.(results.chain)

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

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

# Set varable names

parameter_names = ["a", "bp", "sigma_society"]
pooled_parameter_names = ["a_society[$i]" for i in 1:10]

# Create a3d

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

chns = MCMCChains.Chains(a3d,
  vcat(parameter_names, pooled_parameter_names),
  Dict(
    :parameters => parameter_names,
    :pooled => pooled_parameter_names
  )
);

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

Empirical Posterior Estimates:
                            Mean                SD               Naive SE             MCSE            ESS
            a          1.076167468  0.7704872560 0.01218247319 0.0210530022 1000.000000
           bp         0.263056273  0.0823415805 0.00130193470 0.0022645077 1000.000000
  a_society.1   -0.191723568  0.2421382537 0.00382854195 0.0060563054 1000.000000
  a_society.2    0.054569029  0.2278506876 0.00360263570 0.0051693148 1000.000000
  a_society.3   -0.035935050  0.1926364647 0.00304584994 0.0039948433 1000.000000
  a_society.4    0.334355037  0.1929971201 0.00305155241 0.0063871707  913.029080
  a_society.5    0.049747513  0.1801287716 0.00284808595 0.0043631095 1000.000000
  a_society.6   -0.311903245  0.2096126337 0.00331426674 0.0053000536 1000.000000
  a_society.7    0.148637507  0.1744680594 0.00275858223 0.0047660246 1000.000000
  a_society.8   -0.164567976  0.1821341074 0.00287979309 0.0034297298 1000.000000
  a_society.9    0.277066965  0.1758237250 0.00278001719 0.0055844175  991.286501
 a_society.10   -0.094149204  0.2846206232 0.00450024719 0.0080735022 1000.000000
sigma_society    0.310352849  0.1374834682 0.00217380450 0.0057325226  575.187461
";

# Describe the chain

describe(chns) |> display
println()

# Describe the chain

describe(chns, sections=[:pooled])


"""
    RSOCBridge{T, F, G}

The `RotatedSecondOrderCone` is `SecondOrderCone` representable; see [1, p. 104].
Indeed, we have ``2tu = (t/√2 + u/√2)^2 - (t/√2 - u/√2)^2`` hence
```math
2tu \\ge || x ||_2^2
```
is equivalent to
```math
(t/√2 + u/√2)^2 \\ge || x ||_2^2 + (t/√2 - u/√2)^2.
```
We can therefore use the transformation ``(t, u, x) \\mapsto (t/√2+u/√2, t/√2-u/√2, x)``.
Note that the linear transformation is a symmetric involution (i.e. it is its own transpose and its own inverse).
That means in particular that the norm is of constraint primal and duals are preserved by the tranformation.

[1] Ben-Tal, Aharon, and Arkadi Nemirovski. *Lectures on modern convex optimization: analysis, algorithms, and engineering applications*. Society for Industrial and Applied Mathematics, 2001.
"""
struct RSOCBridge{T, F, G} <: AbstractBridge
    soc::CI{F, MOI.SecondOrderCone}
end
function rotate_function(f::MOI.AbstractVectorFunction, T::Type)
    d = MOI.output_dimension(f)
    f_scalars = MOIU.eachscalar(f)
    t = f_scalars[1]
    u = f_scalars[2]
    x = f_scalars[3:d]
    s2 = √T(2)
    ts = MOIU.operate!(/, T, t, s2)
    us = MOIU.operate!(/, T, u, s2)
    # Cannot use `operate!` here since `ts` and `us` are needed for the next
    # line
    y  = ts - us
    z  = MOIU.operate!(+, T, ts, us)
    return MOIU.operate(vcat, T, z, y, x)
end
function bridge_constraint(::Type{RSOCBridge{T, F, G}}, model,
                           f::MOI.AbstractVectorFunction,
                           s::MOI.RotatedSecondOrderCone) where {T, F, G}
    soc = MOI.add_constraint(model, rotate_function(f, T),
                             MOI.SecondOrderCone(MOI.dimension(s)))
    return RSOCBridge{T, F, G}(soc)
end

function MOI.supports_constraint(::Type{RSOCBridge{T}},
                                ::Type{<:MOI.AbstractVectorFunction},
                                ::Type{MOI.RotatedSecondOrderCone}) where T
    return true
end
MOIB.added_constrained_variable_types(::Type{<:RSOCBridge}) = Tuple{DataType}[]
function MOIB.added_constraint_types(::Type{<:RSOCBridge{T, F}}) where {T, F}
    return [(F, MOI.SecondOrderCone)]
end
function concrete_bridge_type(::Type{<:RSOCBridge{T}},
                              G::Type{<:MOI.AbstractVectorFunction},
                              ::Type{MOI.RotatedSecondOrderCone}) where T
    S = MOIU.promote_operation(/, T, MOIU.scalar_type(G), T)
    Y = MOIU.promote_operation(-, T, S, S)
    Z = MOIU.promote_operation(+, T, S, S)
    F = MOIU.promote_operation(vcat, T, Z, Y, G)
    return RSOCBridge{T, F, G}
end

# Attributes, Bridge acting as a model
function MOI.get(b::RSOCBridge{T, F},
                 ::MOI.NumberOfConstraints{F, MOI.SecondOrderCone}) where {T, F}
    return 1
end
function MOI.get(b::RSOCBridge{T, F},
                 ::MOI.ListOfConstraintIndices{F, MOI.SecondOrderCone}) where {T, F}
    return [b.soc]
end

# References
function MOI.delete(model::MOI.ModelLike, bridge::RSOCBridge)
    MOI.delete(model, bridge.soc)
end

# Attributes, Bridge acting as a constraint
function MOI.get(model::MOI.ModelLike, attr::MOI.ConstraintFunction,
                 bridge::RSOCBridge{T, F, G}) where {T, F, G}
    # As it is an involution, we can just reapply the same transformation
    func = MOI.get(model, attr, bridge.soc)
    return MOIU.convert_approx(G, rotate_function(func, T))
end
function MOI.get(model::MOI.ModelLike, attr::MOI.ConstraintSet, bridge::RSOCBridge)
    set = MOI.get(model, attr, bridge.soc)
    return MOI.RotatedSecondOrderCone(MOI.dimension(set))
end
# As the linear transformation is a symmetric involution,
# the constraint primal and dual both need to be processed by reapplying the
# same transformation
function rotate_result(model,
                       attr::Union{MOI.ConstraintPrimal, MOI.ConstraintDual},
                       ci::MOI.ConstraintIndex)
    x = MOI.get(model, attr, ci)
    T = eltype(x)
    s2 = √T(2)
    return [x[1]/s2 + x[2]/s2; x[1]/s2 - x[2]/s2; x[3:end]]
end
# Need to define both `get` methods and redirect to `rotate_result` to avoid
# ambiguity in dispatch
function MOI.get(model::MOI.ModelLike, attr::MOI.ConstraintPrimal,
                 bridge::RSOCBridge)
    return rotate_result(model, attr, bridge.soc)
end
function MOI.get(model::MOI.ModelLike, attr::MOI.ConstraintDual,
                 bridge::RSOCBridge)
    return rotate_result(model, attr, bridge.soc)
end


# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.10.3
#   kernelspec:
#     display_name: Julia 1.6.4
#     language: julia
#     name: julia-1.6
# ---

# %% [markdown]
# # 通常の信頼区間とベイズ版信用区間の比較

# %%
using Distributions
using Roots
using Plots
using StatsFuns: logistic, logit

# \approx TAB → ≈
# \lessapprox TAB → ⪅
x ⪅ y = x < y || x ≈ y

# %%
# Bernoulli分布モデルでの通常の信頼区間

"""P値函数(exact版)"""
function pvalue_exact(n, p, k)
    bin = Binomial(n, p)
    p0 = pdf(bin, k)
    s = sum(pdf(bin, j) for j in support(bin) if pdf(bin, j) ⪅ p0)
    min(1, s)
end

"""P値函数(正規分布近似)"""
function pvalue_normal(n, p, k)
    bin = Binomial(n, p)
    normal = Normal(mean(bin), std(bin))
    min(1, 2cdf(normal, k), 2ccdf(normal, k))
end

"""信頼区間函数"""
function confidence_interval(pvalue_func, n, k; α = 0.05)
    f(t) = pvalue_func(n, logistic(t), k) - α
    CI = logistic.(find_zeros(f, -10.0, 10.0))
    if length(CI) < 2
        return 2k ≤ n ? [0.0, first(CI)] : [first(CI), 1.0]
    else
        return [first(CI), last(CI)]
    end
end

# %%
# Bernoulli分布モデルと共役事前分布に関するベイズ版信用区間

"""共役事後分布"""
posterior_dist(n, k; a = 1, b = a) = Beta(a + k, b + n - k)

"""ベイズ版P値函数"""
function pvalue_bayes(n, p, k; a = 1, b = a)
    posterior = posterior_dist(n, k; a, b)
    v0 = logpdf(posterior, p)
    f(t) = logpdf(posterior, logistic(t)) - v0
    m = params(posterior) |> ((α, β),) -> (α - 1)/(α + β - 2)
    if m ≤ 0
        s = ccdf(posterior, p)
    elseif m ≥ 1
        s = cdf(posterior, p)
    elseif p ≤ m
        q = logistic(find_zero(f, min(50, logit(m) + 1)))
        s = cdf(posterior, p) + ccdf(posterior, q)
    else
        q = logistic(find_zero(f, max(-50, logit(m) - 1)))
        s = cdf(posterior, q) + ccdf(posterior, p)
    end
    min(1, s)
end

"""ベイズ版信用区間函数"""
function credible_interval(n, k; α = 0.05, a = 1, b = a)
    confidence_interval((n, p, k) -> pvalue_bayes(n, p, k; a, b), n, k; α)
end

# %%
# 信頼区間とベイズ版信用区間の比較

function plot_cis(n, k; α = 0.05, a = 1, b = a)
    ci_exact = confidence_interval(pvalue_exact, n, k; α)
    ci_normal = confidence_interval(pvalue_normal, n, k; α)
    ci_bayes = credible_interval(n, k; α, a, b)

    bin = Binomial(n, k/n)
    μ, σ = mean(bin)/n, std(bin)/n
    p = range(max(1e-8, μ - 3.5σ), min(1-1e-8, μ + 4σ); length=400)
    plot(p, p -> pvalue_bayes(n, p, k; a, b); label="Bayesian p-value", c = 1, lw=1.5)
    plot!(p, p -> pvalue_normal(n, p, k); label="p-value (normal approx.)", c = 2, ls=:dash, lw=1.5)
    plot!(p, p -> pvalue_exact(n, p, k); label="p-value (exact)", c = 3, ls=:dashdot, lw=1.5)
    plot!(; ytick=0:0.05:1)
    plot!(ci_bayes, fill(1.2α, 2); label="Bayesian credible interval", c=1, lw=2)
    plot!(ci_normal, fill(1.0α, 2); label="conf. int. (normal approx.)", c=2, ls=:dash, lw=2)
    plot!(ci_exact, fill(0.8α, 2); label="confidence interval (exact)", c=3, ls=:dashdot, lw=2)
    title!("n = $n, k = $k, α = $α, prior = Beta($a, $b)"; titlefontsize=10)
end

# %%
plot_cis(10, 3)

# %%
plot_cis(100, 30)

# %%
plot_cis(1000, 300)

# %% [markdown]
# $n$ を大きくするにつれて通常の信頼区間とベイズ版信用区間は数値的によく一致するようになる。
#
# このようなことはシンプルなモデルでは一般的に広く成立している.

# %%
plot_cis(1000, 30)

# %%
plot_cis(1000, 3)

# %%
plot_cis(1000, 3; a = 0.5)

# %%
plot_cis(1000, 3; a = 0)

# %%

0024/confidence and credible intervals.jl


# ---
# title: 873. Length of Longest Fibonacci Subsequence
# id: problem873
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Array, Dynamic Programming
# link: <https://leetcode.com/problems/length-of-longest-fibonacci-subsequence/description/>
# hidden: true
# ---
#
# A sequence `X_1, X_2, ..., X_n` is _fibonacci-like_ if:
#
#   * `n >= 3`
#   * `X_i + X_{i+1} = X_{i+2}` for all `i + 2 <= n`
#
# Given a **strictly increasing**  array `A` of positive integers forming a
# sequence, find the **length** of the longest fibonacci-like subsequence of
# `A`.  If one does not exist, return 0.
#
# ( _Recall that a subsequence is derived from another sequence`A` by deleting
# any number of elements (including none) from `A`, without changing the order
# of the remaining elements.  For example, `[3, 5, 8]` is a subsequence of `[3,
# 4, 5, 6, 7, 8]`._)
#
#
#
# **Example 1:**
#
#
#
#     Input: [1,2,3,4,5,6,7,8]
#     Output: 5
#     Explanation: The longest subsequence that is fibonacci-like: [1,2,3,5,8].
#
#
# **Example 2:**
#
#
#
#     Input: [1,3,7,11,12,14,18]
#     Output: 3
#     **Explanation** :
#     The longest subsequence that is fibonacci-like:
#     [1,11,12], [3,11,14] or [7,11,18].
#
#
#
#
# **Note:**
#
#   * `3 <= A.length <= 1000`
#   * `1 <= A[0] < A[1] < ... < A[A.length - 1] <= 10^9`
#   * _(The time limit has been reduced by 50% for submissions in Java, C, and C++.)_
#
#
## @lc code=start
using LeetCode

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


# ---
# title: 910. Smallest Range II
# id: problem910
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Math, Greedy
# link: <https://leetcode.com/problems/smallest-range-ii/description/>
# hidden: true
# ---
#
# Given an array `A` of integers, for each integer `A[i]` we need to choose
# **either  `x = -K` or `x = K`**, and add `x` to `A[i] **(only once)**`.
#
# After this process, we have some array `B`.
#
# Return the smallest possible difference between the maximum value of `B` and
# the minimum value of `B`.
#
#
#
# **Example 1:**
#
#
#
#     Input: A = [1], K = 0
#     Output: 0
#     **Explanation** : B = [1]
#
#
# **Example 2:**
#
#
#
#     Input: A = [0,10], K = 2
#     Output: 6
#     **Explanation** : B = [2,8]
#
#
# **Example 3:**
#
#
#
#     Input: A = [1,3,6], K = 3
#     Output: 3
#     **Explanation** : B = [4,6,3]
#
#
#
#
# **Note:**
#
#   1. `1 <= A.length <= 10000`
#   2. `0 <= A[i] <= 10000`
#   3. `0 <= K <= 10000`
#
#
## @lc code=start
using LeetCode

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


function load_data(dataset::AbstractString)
    if dataset == "facebook_data"
        return load_facebook_data()
    end
end

function load_facebook_data()
    pathname = joinpath(dirname(dirname(@__FILE__)),"data")
    fname = joinpath(pathname, "facebook_data.txt")
    facebook_adj_matrix = fill(0, (4039, 4039))
    open(fname) do file
        for ln in eachline(file)
            pieces = split(ln,' ', keepempty=false)
            facebook_adj_matrix[(parse(Int64,pieces[1]) + 1), (parse(Int64,pieces[2]) + 1)] = 1
        end
    end

    facebook_adj_matrix = sparse(facebook_adj_matrix)
    facebook_adj_matrix = facebook_adj_matrix + facebook_adj_matrix'
    facebook_adj_matrix = min.(facebook_adj_matrix, 1)
    facebook_adj_matrix = MatrixNetwork(facebook_adj_matrix)

    return facebook_adj_matrix
end


using ChaosTools
using Test
using DynamicalSystemsBase, DelayEmbeddings
using StatsBase
standardize = DelayEmbeddings.standardize
test_value = (val, vmin, vmax) -> @test vmin <= val <= vmax

ti = time()
@testset "ChaosTools tests" begin

include("basins/basins_tests.jl")
include("basins/uncertainty_tests.jl")
include("basins/tipping_points_tests.jl")

include("orbitdiagrams/orbitdiagram_tests.jl")
include("orbitdiagrams/poincare_tests.jl")

include("chaosdetection/lyapunov_exponents.jl")
#include("chaosdetection/gali_tests.jl") # TODO: make those faster
include("chaosdetection/partially_predictable_tests.jl")
include("chaosdetection/01test.jl")
#include("chaosdetection/expansionentropy_tests.jl")  # TODO: make those faster

include("period_return/periodicity_tests.jl")
include("period_return/period_tests.jl")
include("period_return/yin_tests.jl")
# include("period_return/transit_time_tests.jl")

include("dimensions/dims.jl")
include("dimensions/correlationdim.jl")
include("nlts_tests.jl")
# include("dyca_tests.jl") # TODO: fix DyCa tests once and for
# all by using https://github.com/JuliaDynamics/JuliaDynamics/tree/master/timeseries

end

ti = time() - ti
println("\nTest took total time of:")
println(round(ti, digits = 33), " seconds or ",
round(ti/60, digits = 3), " minutes")


"""
	stroboscopicmap(ds::ContinuousDynamicalSystem, T; kwargs...)  → smap

Return a map (integrator) that produces iterations over a period `T` of the `ds`,
known as a stroboscopic map. See [Integrator API](@ref) for handling integrators.

See also [`poincaremap`](@ref).

## Keyword Arguments
* `u0`: initial state
* `diffeq` is a `NamedTuple` (or `Dict`) of keyword arguments propagated into
  `init` of DifferentialEquations.jl.

## Example
```julia
f = 0.27; ω = 0.1
ds = Systems.duffing(zeros(2); ω, f, d = 0.15, β = -1)
smap = stroboscopicmap(ds, 2π/ω; diffeq = (;reltol = 1e-8))
reinit!(smap, [1.0, 1.0])
u = step!(smap)
u = step!(smap, 4) # step 4 iterations forward
```
"""
function stroboscopicmap(ds::CDS, T; u0 = get_state(ds), diffeq = NamedTuple())
	integ = integrator(ds, u0; diffeq)
	return StroboscopicMap{typeof(integ), dimension(ds), typeof(T)}(integ, T)
end
struct StroboscopicMap{I, D, F} <: GeneralizedDynamicalSystem
	integ::I
	T::F
end
isdiscretetime(::StroboscopicMap) = true
DelayEmbeddings.dimension(::StroboscopicMap{I, D}) where {I, D} = D

integrator(p::StroboscopicMap) = p
function step!(smap::StroboscopicMap)
	step!(smap.integ, smap.T, true)
	return smap.integ.u
end
function step!(smap::StroboscopicMap, n::Int)
	for k in 1:n; step!(smap.integ, smap.T, true); end
	return smap.integ.u
end
function reinit!(smap::StroboscopicMap, u0)
	reinit!(smap.integ, u0)
	return
end
function get_state(smap::StroboscopicMap)
	return smap.integ.u
end

function Base.show(io::IO, smap::StroboscopicMap)
    println(io, "Iterator of the stroboscopic map")
    println(io,  rpad(" rule f: ", 14),     DynamicalSystemsBase.eomstring(smap.integ.f.f))
    println(io,  rpad(" Period: ", 14),     smap.T)
end

current_time(smap::StroboscopicMap) = current_time(smap.integ)
function (smap::StroboscopicMap)(t)
	if t == current_time(smap)
		return get_state(smap)
	else
		error("Can't extrapolate discrete systems!")
	end
end
integrator(pinteg::StroboscopicMap, args...; kwargs...) = pinteg


using LogDensityProblems, Test, Distributions, TransformVariables
import LogDensityProblems: capabilities, dimension, logdensity
using LogDensityProblems: logdensity_and_gradient, LogDensityOrder

import ForwardDiff, Flux, Zygote, TransformVariables, Random
using Parameters: @unpack

using Plots
gr()

Random.seed!(3)

####
#### transformed Bayesian problem
####

ProjDir = @__DIR__
cd(ProjDir)
include("../src/HMC.jl")
include("../src/generate_n_samples.jl")

# ### snippet 9.3

# Construct the logdensity problem

struct clip_9_3_model{TY <: AbstractVector, TX <: AbstractVector}
    "Observations."
    y::TY
    "Covariate"
    x::TX
end

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

function (problem:: clip_9_3_model)(θ)
    @unpack y, x, = problem    # extract the data
    @unpack muy, mux = θ     # works on the named tuple too
    ll = 0.0
    ll += loglikelihood(Normal(mux, 1), x)
    ll += loglikelihood(Normal(muy, 1), y)
    ll += logpdf(Normal(0, 1), mux)
    ll += logpdf(Normal(0, 1), muy)
    ll
end

# Instantiate the model with data and inits.

Random.seed!(3)

N = 100
x = rand(Normal(0, 1), N)
#=
x = [0.83975036,  0.70534183,  1.30596472, -1.38799622,  1.27291686,  0.18419277,
  0.75227990,  0.59174505, -0.98305260, -0.27606396, -0.87085102,  0.71871055,
  0.11065288, -0.07846677, -0.42049046, -0.56212588,  0.99751344, -1.10513006,
  -0.14228783,  0.31499490]
=#
y = rand(Normal(0, 1), N)
#=
y = [2.287247161, -1.196771682, -0.694292510, -0.412292951, -0.970673341, -0.947279945,
  0.748139340, -0.116955226,  0.152657626,  2.189978107,  0.356986230,  2.716751783,
  2.281451926,  0.324020540,  1.896067067,  0.467680511, -0.893800723, -0.307328300,
  -0.004822422,  0.988164149]
=#

p = clip_9_3_model(y, x)
θ = (muy = 0.0, mux=0.0)

# Write a function to return properly dimensioned transformation.

t = as((muy = asℝ, mux = asℝ))
#t1 = as((muy = asℝ₊, mux = asℝ₊))

problem_transformation(p::clip_9_3_model) = t

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

P = TransformedLogDensity(problem_transformation(p), p)
∇P = ADgradient(:ForwardDiff, P);

println()
p(θ) |> display
v, g = logdensity_and_gradient(∇P, [0.0, 0.0])
println()
v  |> display
g  |> display
value_U = -v
grad_U = -g
v, g = logdensity_and_gradient(∇P, [-0.1, 0.2])
println()
v  |> display
g  |> display
value_U = -v
grad_U = -g

@testset "transformed Bayesian problem 2D" begin

    for _ in 1:100
        x = random_arg(t)
        θ, lj = transform_and_logjac(t, x)
        px = logdensity(P, x)
        px2, ∇px = logdensity_and_gradient(∇P, x)
        @test px == px2
    end

end

println()
new_q, ptraj, qtraj, accept, dH = HMC(p, ∇P, 0.03, 11, [-0.1, 0.2])

new_q |> display
println()
ptraj |> display
println()
qtraj |> display
println()
dH |> display
println()

fig = generate_n_samples(p, ∇P, q =  [-0.1, 0.2]);
savefig(fig, "fig-9-3.pdf")


include("direct.jl")


using Test
using ModelingToolkit, OrdinaryDiffEq

include("../examples/rc_model.jl")

sys = structural_simplify(rc_model)
@test !isempty(ModelingToolkit.defaults(sys))
u0 = [
      capacitor.v => 0.0
      capacitor.p.i => 0.0
      resistor.v => 0.0
     ]
prob = ODEProblem(sys, u0, (0, 10.0))
sol = solve(prob, Rodas4())

@test sol[resistor.p.i] == sol[capacitor.p.i]
@test sol[resistor.n.i] == -sol[capacitor.p.i]
@test sol[capacitor.n.i] == -sol[capacitor.p.i]
@test iszero(sol[ground.g.i])
@test iszero(sol[ground.g.v])
@test sol[resistor.v] == sol[source.p.v] - sol[capacitor.p.v]

prob = ODAEProblem(sys, u0, (0, 10.0))
sol = solve(prob, Tsit5())

@test sol[resistor.p.i] == sol[capacitor.p.i]
@test sol[resistor.n.i] == -sol[capacitor.p.i]
@test sol[capacitor.n.i] == -sol[capacitor.p.i]
@test iszero(sol[ground.g.i])
@test iszero(sol[ground.g.v])
@test sol[resistor.v] == sol[source.p.v] - sol[capacitor.p.v]
#using Plots
#plot(sol)

include("../examples/serial_inductor.jl")
sys = structural_simplify(ll_model)
u0 = [
      inductor1.i => 0.0
      inductor2.i => 0.0
      inductor2.v => 0.0
     ]
prob = ODEProblem(sys, u0, (0, 10.0))
sol = solve(prob, Rodas4())

prob = ODAEProblem(sys, u0, (0, 10.0))
sol = solve(prob, Tsit5())


using MLDatasets
using BSON
using BSON: @load
using Plots
using Statistics

function prepare_data(file_name)
    X = Iris.features()
    y_orig = Iris.labels()

    y_name = unique(y_orig)
    y = zeros(length(y_orig))
    for i in 1:length(y_name)
        y[y_orig .== y_name[i]] .= i
    end

    X = Matrix(X')
    y = Int64.(y[:])

    BSON.bson(file_name, X=X, y=y, y_name=y_name)
end

function normalize_data(X)
    col_mean = mean(X, dims=1)
    col_std = std(X, dims=1)
    return (X .- col_mean) ./ col_std
end

file_name = joinpath("src", "lecture_08", "data", "iris.bson")
prepare_data(file_name)


# ---
# title: 719. Find K-th Smallest Pair Distance
# id: problem719
# author: Tian Jun
# date: 2020-10-31
# difficulty: Hard
# categories: Array, Binary Search, Heap
# link: <https://leetcode.com/problems/find-k-th-smallest-pair-distance/description/>
# hidden: true
# ---
#
# Given an integer array, return the k-th smallest **distance** among all the
# pairs. The distance of a pair (A, B) is defined as the absolute difference
# between A and B.
#
# **Example 1:**
#
#
#
#     Input:
#     nums = [1,3,1]
#     k = 1
#     **Output: 0**
#     Explanation:
#     Here are all the pairs:
#     (1,3) -> 2
#     (1,1) -> 0
#     (3,1) -> 2
#     Then the 1st smallest distance pair is (1,1), and its distance is 0.
#
#
# **Note:**
#
#   1. `2 <= len(nums) <= 10000`.
#   2. `0 <= nums[i] < 1000000`.
#   3. `1 <= k <= len(nums) * (len(nums) - 1) / 2`.
#
#
## @lc code=start
using LeetCode

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


# Kirsten Van Huffel & Tristan Vanneste

module BeesAlgorithm

# you have to import everything you need for your module to work
# if you use a new package, don't forget to add it in the package manager

# export all functions that are relevant for the user
export ArtificialBeeColonization, sphere, ackley, rosenbrock, branin, rastrigine

"""
    initialize_population(D::Number, bounds_lower::Vector, bounds_upper::Vector, Np::Number)

This function generates 'Np' random solutions (food sources) within the domain
of the variables to form an initial population for the ABC algorithm.

Input
- 'D': number of decision variables in the function to be minimized
- 'bounds_lower': lower bounds of variables
- 'bounds_upper': upper bounds of variables
- 'Np': number of food sources/employed bees/onlooker bees

Output
- 'population': a random population of the size 'Np' with 'D' variables

## Examples

```julia-repl
julia> bounds_lower = [-5,-5,-5,-5]
julia> bounds_upper = [5,5,5,5]
julia> D=4
julia> n=9
julia> initialize_population(D, bounds_lower, bounds_upper, n)
9-element Array{Any,1}:
 [-2, 2, 5, -5]
 [2, 1, -5, 1]
 [4, 1, 2, 2]
 [-3, -4, -2, -2]
 [1, 2, 4, 3]
 [4, -2, 3, -4]
 [1, 1, 0, -1]
 [2, -2, -2, -5]
 [-3, 2, 3, -3]
```
"""
function initialize_population(D::Number, bounds_lower::Vector, bounds_upper::Vector, Np::Number)
    population = []
    for i in 1:Np
        food_source = collect(rand(bounds_lower[i]:bounds_upper[i]) for i in 1:D)
        push!(population, food_source)
    end
    return population
end

"""
    compute_objective(solutions, f::Function)

Compute the objective values for an objective function 'f'. This objective value is minimized in the ABC algorithm.

Input
- 'solutions': population of solutions or a single solution
- 'f': the function that you want to use for computing objective values

Output
- 'output': objective values

## Examples

```julia-repl
julia> D=4
julia> n=9
julia> bounds_lower = [-5,-5,-5,-5]
julia> bounds_upper = [5,5,5,5]
julia> population = initialize_population(D, bounds_lower, bounds_upper, n)
julia> compute_objective(population, sphere)
9-element Array{Any,1}:
 34
 75
 18
 51
 42
 51
 31
 42
 59

 julia> compute_objective(population, ackley)
 9-element Array{Any,1}:
  8.836638915350669
 11.587599478917705
  6.914978162949289
 10.20776787857683
  9.538956614465143
 10.20776787857683
  8.538932726524155
  9.538956614465143
 10.722262633846869
```
"""
function compute_objective(solutions, f::Function)
    return f.(solutions)
end

"""
    compute_fitness(objective_values)

This functions computes the fitness of each solution in the population.
The fitness is computed as ``1/(1+objective values)`` if f(x) >= 0 and
as ``1 + absolute(objective values)`` if f(x) < 0.
The bigger the objective values the smaller the fitness values.

Input
- 'objective_values': objective values

Output
- 'fitness_values': fitness values corresponding to the input values (objective values)

## Examples

```julia-repl
julia> objective_values = [4 5 -1]
julia> compute_fitness(objective_values)
3-element Array{Any,1}:
 0.2
 0.16666666666666666
 2
```
"""
function compute_fitness(objective_values)

    fitness_values = zeros(length(objective_values),1)
    fitness_values[objective_values .>= 0] = 1 ./ (1 .+ objective_values[objective_values .>= 0])
    fitness_values[objective_values .< 0] = 1 .+ abs.(objective_values[objective_values .< 0])
    return fitness_values

end

"""
    foodsource_info_prob(fitness_values)

This function measures the food source information as probabilities.
The food source information is computed as following: ``0.9 * (fitness / maximum(fitness )) + 0.1``.

Input
- 'fitness_values': fitness values

Output
- 'probabilities': probabilities of the food source information

## Examples

```julia-repl
julia> fitness_values = [0.2 0.16 1]
julia> foodsource_info_prob(fitness_values)
3-element Array{Any,1}:
 0.28
 0.24400000000000002
 1.0
```
"""
function foodsource_info_prob(fitness_values)

    probabilities = 0.9 .* (fitness_values ./ maximum(fitness_values)) .+ 0.1
    return probabilities
end

"""
    create_newsolution(solution::Vector, population, bounds_lower::Vector, bounds_upper::Vector)

Creates new solution by changing one variable using a random partner solution.
Firstly, a random variable in the solution to be modified will be selected.
Secondly,  a random partner in the current population is chosen, and a random
variable of this partner solution is picked. Based on the value of these two variables,
a new value is interpolated and assigned to the first random variable of the original solution.

Input
- 'solutions': solution
- 'population': population of solutions
- 'bounds_lower': lower bounds of variables
- 'bounds_upper': upper bounds of variables

Output
- 'solution_new': new solution with one variable changed using a random variable in a random partner solution

## Examples

```julia-repl
julia> solution = [4 0 1 4]
julia> bounds_lower = [-5,-5,-5,-5]
julia> bounds_upper = [5,5,5,5]
julia> D=4
julia> n=9
julia> population = initialize_population(D, bounds_lower, bounds_upper, n)
julia> create_newsolution(solution, population, bounds_lower, bounds_upper)
1×4 Array{Float64,2}:
 -1.49748  0.0  1.0  4.0
```
"""
function create_newsolution(solution::Vector, population, bounds_lower::Vector, bounds_upper::Vector)
    # select random variable to change
    randomvar1_index = rand(1:size(solution)[1], 1)

    # select partner solution to generate new solution
    randompartner_index = rand(1:size(population)[1], 1)

    # select random variable in partner solution to interpolate with

    randompartner = population[randompartner_index, :][1]
    randomvar2_index = rand(1:length(randompartner), 1)

    # create new food location
    phi = rand()*2-1 #random number between -1 and 1
    global solution_new = float(deepcopy(solution))
    a = solution[randomvar1_index]
    b = randompartner[randomvar2_index]
    solution_new[randomvar1_index] = a + phi*(a - b)

    # check if lower bound is violated
    if solution_new[randomvar1_index] < bounds_lower[randomvar1_index]
        solution_new[randomvar1_index] = bounds_lower[randomvar1_index]
    end

    # check if upper bound is violated
    if solution_new[randomvar1_index] > bounds_upper[randomvar1_index]
        solution_new[randomvar1_index] = bounds_upper[randomvar1_index]
    end
return solution_new
end

"""
    employed_bee_phase(population, bounds_lower::Vector, bounds_upper::Vector,
                            trial::Array, Np::Number, f::Function)

In this function, the employed bee phase is implemented.
Employed bees try to identify better food source than the one they were associated previously.
A new solution is generated using a partner solution.
Thereafter, greedy selection is performed, meaning that a new solution only  will be accepted if it is better than the current solution.
Every bee in the swarm will explores one food source. All solutions get an opportunity to generate a new solution in the employed bee phase.

Input
- 'population': population of solutions
- 'bounds_lower': lower bounds of variables
- 'bounds_upper': upper bounds of variables
- 'trial': Vector of current trials of solutions in the population
- 'Np': number of food sources/employed bees/onlooker bees
- 'f': the function that you want to use for computing objective values

Output
- 'population new evolved': newly generated population of solutions
- 'fitness new evolved': fitness values of the new population
- 'objective new evolved': objective values of the new population
- 'trial': updated trials of new solutions in the new population
    When original solution has failed to generate better solution, trial counter is increased by 1 unit
    When better solution has been found, the trial counter for this new solution is set to zero

## Examples

```julia-repl
julia> trial = zeros(size(population)[1])
julia> bounds_lower = [-5,-5,-5,-5]
julia> bounds_upper = [5,5,5,5]
julia> D=4
julia> N=9
julia> population = initialize_population(D, bounds_lower, bounds_upper, n)
julia> population_new_evolved, fitness_new_evolved, objective_new_evolved, trial = employed_bee_phase(population, bounds_lower, bounds_upper, trial,n ,sphere)
julia> population_new_evolved
9-element Array{Any,1}:
 [5.0, -5.0, 5.0, 2.748420414594148]
 [2, 4, 0, -4]
 [-5.0, 2.0, 0.0, -1.1761590179019343]
 [1, 5, 5, -2]
 [2.0, -1.0, -3.0, 0.75942632274392]
 [1, 5, -1, 1]
 [0.0, 3.1592990065770192, -2.0, -3.0]
 [-4, 3, 3, 2]
 [3, 2, -3, -3]
```
"""
function employed_bee_phase(population, bounds_lower::Vector, bounds_upper::Vector,
                                trial::Array, Np::Number, f::Function)
    population_new = []

    # create new food sources
    for i in 1:Np
        solution = population[i, :][1]
        solution_new = solution
        while solution_new == solution
            solution_new = create_newsolution(solution, population, bounds_lower, bounds_upper)
        end
        append!(population_new, [solution_new])
    end

    # evaluate fitness old and new population
    objective_values_old = compute_objective(population, f)
    fitness_old = compute_fitness(objective_values_old)
    objective_values_new = compute_objective(population_new, f)
    fitness_new = compute_fitness(objective_values_new)

    # perform greedy selection
    population_new_evolved = []
    fitness_new_evolved = []
    objective_new_evolved = []

    for j in 1:Np
        if fitness_new[j] > fitness_old[j]
            append!(population_new_evolved, [population_new[j]])
            append!(fitness_new_evolved, fitness_new[j])
            append!(objective_new_evolved, objective_values_new[j])
            trial[j] = 0
        else
            append!(population_new_evolved, [population[j]])
            append!(fitness_new_evolved, fitness_old[j])
            append!(objective_new_evolved, objective_values_old[j])
            trial[j] += 1
        end
    end

    return population_new_evolved, fitness_new_evolved, objective_new_evolved, trial
end

"""
    onlooker_bee_phase(population, bounds_lower::Vector,
                                    bounds_upper::Vector,
                                    trial::Array, Np::Number,
                                    f::Function)

This function implements the onlooker bee phase.
In the onlooker bee phase, a food source is selected for further exploitation with a probability related to the nectar amount,
i.e. a solution with higher fitness will have a higher probability to be chosen.
Fitter solutions may undergo multiple onlooker bee explorations. As in the employed bee phase,
new solutions are generated using a partner solution and greedy selection is performed.
In contrast to the employed bee phase, not every food source will be explored,
since every onlooker bee will explore a certain food source with a certain probability (depending on nectar amount).

Input
- 'population': population of solutions
- 'bounds_lower': lower bounds of variables
- 'bounds_upper': upper bounds of variables
- 'trial': Vector of current trials of solutions in the population
- 'Np': number of food sources/employed bees/onlooker bees
- 'f': the function that you want to use for computing objective values

Output
- 'population new evolved': newly generated population of solutions
- 'fitness new evolved': fitness values of the new population
- 'objective new evolved': objective values of the new population
- 'trial': updated trials of new solutions in the new population
    When original solution has failed to generate better solution, trial counter is increased by 1 unit
    When better solution has been found, the trial counter for this new solution is set to zero

## Examples

```julia-repl
julia> bounds_lower = [-5,-5,-5,-5]
julia> bounds_upper = [5,5,5,5]
julia> D=4
julia> N=9
julia> population = initialize_population(D, bounds_lower, bounds_upper, N)
julia> trial = zeros(size(population)[1])
julia> population_new_evolved, fitness_new_evolved, objective_new_evolved, trial = onlooker_bee_phase(population, bounds_lower, bounds_upper, trial,N ,sphere)
julia> population_new_evolved
9-element Array{Any,1}:
 [-1.0, -5.0, 0.5313751547457599, 5.0]
 [-1, -3, -4, -2]
 [0, 1, -3, 5]
 [1.0, -4.0, -5.0, 1.8770812747894876]
 [-0.944880850158663, -1.2865870734899967, 0.04370050469571929, 1.0]
 [4, -3, 5, 5]
 [0.0, 1.0208794819836884, 0.0, -3.0]
 [-2, 2, 5, -1]
 [-5, -3, -4, 5]
```
"""
function onlooker_bee_phase(population, bounds_lower::Vector, bounds_upper::Vector,
                    trial::Array, Np::Number, f::Function)
    m = 0 # onlooker bee
    n = 1 # food source

    objective_values = compute_objective(population,f)
    fitness = compute_fitness(objective_values)
    # first calculate the probability values
    proba = foodsource_info_prob(fitness)

    while m <= Np # we want for every onlooker bee a new solution
        r = rand()
        if r <= proba[n]
            solution = population[n, :][1] # solution n
            objective_values_old = compute_objective([solution], f)
            fitness_old = compute_fitness(objective_values_old)[1]

            solution_new = solution
            while solution_new == solution
                solution_new = create_newsolution(solution, population, bounds_lower, bounds_upper)
            end

            objective_values_new = compute_objective([solution_new], f)
            fitness_new = compute_fitness(objective_values_new)[1]
            if fitness_new > fitness_old # if this get accepted
                population[n, :] = [solution_new]
                trial[n]=0
            else
                trial[n] += 1
            end
            m = m + 1
        end
        # if the rand < proba is not sattisfied
        n = n +1
        if n > Np
            n = 1
        end
    end
    objective_new_evolved = compute_objective(population,f)
    fitness_new_evolved = compute_fitness(objective_new_evolved)

    return population, fitness_new_evolved, objective_new_evolved, trial
end

"""
    scouting_phase(population, bounds_lower::Vector, bounds_upper::Vector,
                    D::Number ,trials::Array, fitness, objective, limit::Number, f::Function)

This function implements the scouting phase.
If the value of the trial counter for a certain solution is greater than fixed limit, then a solution can enter the scout phase.
The latter food source is then considered as exhausted and will therefore be abandoned by the bees.
The current optimal solution will be stored separately to avoid losing it in the scouting phase.
After discarding the exhausted solution, a new random solution is generated and the trial counter of this solution is reset to zero.

Input
- 'population': population of solutions
- 'bounds_lower': lower bounds of variables
- 'bounds_upper': upper bounds of variables
- 'D': number of decision variables in the function to be minimized
- 'trials': current trial of solutions
- 'fitness': fitness values
- 'objective': objective values
- 'limit': limit value
- 'f': the function that you want to use for computing objective values

Output
- 'population new evolved': newly generated population of solutions
- 'fitness new evolved': fitness values of the new population
- 'objective new evolved': objective values of the new population
- 'trial': updated trials of new solutions in the new population
    When original solution has failed to generate better solution, trial counter is increased by 1 unit
    When better solution has been found, the trial counter for this new solution is set to zero

## Examples

```julia-repl
julia> trial = ones(size(population)[1])
julia> bounds_lower = [-5,-5,-5,-5]
julia> bounds_upper = [5,5,5,5]
julia> D=4
julia> N=9
julia> population = initialize_population(D, bounds_lower, bounds_upper, n)
julia> population
9-element Array{Any,1}:
 [3, 1, 2, -2]
 [-5, 4, 4, 2]
 [-4, 5, -2, -4]
 [5, 5, -5, 0]
 [-2, -4, 5, 5]
 [1, 2, 5, -2]
 [1, -1, 5, -5]
 [-2, -5, -4, -5]
 [5, -1, -1, -1]
julia> objective= compute_objective(population,sphere)
julia> fitness = compute_fitness(objective)
julia> population_new_evolved, fitness_new_evolved, objective_new_evolved, trial  = scouting_phase(population, bounds_lower, bounds_upper, D, trial, fitness, objective, 0, sphere)
julia> population_new_evolved
9-element Array{Any,1}:
 [3, 1, 2, -2]
 [-5, 4, 4, 2]
 [-2.959166963468496, 2.2098469985145712, -0.2708630367391418, 4.89872222904274]
 [5, 5, -5, 0]
 [-2, -4, 5, 5]
 [1, 2, 5, -2]
 [1, -1, 5, -5]
 [-2, -5, -4, -5]
 [5, -1, -1, -1]
```
"""
function scouting_phase(population, bounds_lower::Vector, bounds_upper::Vector,
                        D::Number ,trials::Array, fitness, objective, limit::Number, f::Function)

    # check whether the trial vector exceed the limit value and importantly where
    index_exceed = trials .> limit
    if sum(index_exceed) >= 1 # there is minimal one case where we exceed the limit
        if sum(maximum(trials) .== trials) > 1 # multiple cases have the same maximum so chose randomly

            possible_scoutings = findall(trials .== maximum(trials))
            idx = rand(1:size(possible_scoutings)[1])
            scouting_array = possible_scoutings[idx]
        else # only one array has a maximum => chose this one

            scouting_array = argmax(trials)

        end
        pop = population[scouting_array]
        fit = fitness[scouting_array]
        obj = objective[scouting_array]
        trail = trials[scouting_array]

        #creating random population
        sol_new = bounds_lower + (bounds_upper-bounds_lower) .* rand(D) # -5 *(10*rand)
        new_obj = compute_objective([sol_new],f)
        new_fit = compute_fitness(new_obj)

        # replacing the new population
        population[scouting_array] = sol_new

        fitness[scouting_array] = new_fit[1]
        objective = float(objective)
        objective[scouting_array] = new_obj[1]
        trials[scouting_array] = 0

    end
    return population, fitness, objective, trials
end

"""
    ArtificialBeeColonization(D::Number, bounds_lower::Vector, bounds_upper::Vector,
                                S::Number, T::Number, limit::Number, f::Function)

This functions runs the Artificial Bee Colony Algorithm.
First a initialize population is made. For each iteration (T) the ABC algorithm undergoes the employed bee phase,
onlooker bee phase and possibly the scouting phase.
During the process, the solution with the highest fitness value is stored separately and updated after every iteration and scouting phase.

Input
- 'D': number of decision variables in the function to be minimized
- 'bounds_lower': lower bounds of variables
- 'bounds_upper': upper bounds of variables
- 'S': swarm size (S/2 = Np, the number of employed bees, onlooker bees and scouting bees)
- 'T': number of cycles
- 'limit': decides when scouts phase needs to be executed (often taken Np*D)
- 'f': the function that you want to use for computing objective values (this function will be minimized)

Output
- 'optimal_solution': gives a vector of the size D with the optimal solution
- 'populations': all populations that were computed during the ABC algorithm
- 'best fitness tracker': a vector with for each iteration the best fitness value so far encountered.

## Examples

```julia-repl
julia> S = 24
julia> bounds_lower = [-100,-100,-100,-100]
julia> bounds_upper = [100,100,100,100]
julia> D=4
julia> limit = D * (S/2)
julia> T = 500
julia> optimal_solution,populations,best_fitness_tracker = ArtificialBeeColonization(D, bounds_lower, bounds_upper, S, T, limit, sphere)
julia> optimal_solution
4-element Array{Float64,1}:
  1.4537737816170574e-9
 -9.379333377223156e-9
 -2.478809612571768e-9
  3.9716934385662925e-11
```
"""
function ArtificialBeeColonization(D::Number, bounds_lower::Vector, bounds_upper::Vector,
                                    S::Number, T::Number, limit::Number, f::Function)
    @assert D > 0 "D must be positive"
    @assert bounds_lower <= bounds_upper "Lower bounds must be smaller than upper bounds"
    @assert length(bounds_lower) == length(bounds_upper) == D  "The length of the lower bounds must be equal to the length of the upper bounds and the number of decision variables"
    @assert iseven(S) "The particle swarm size must be an even number"
    @assert S > 0 "Particle swarm size must be a positive number"
    @assert T > 0 "Number of cycles must be positive"

    Np = Int8(S/2) # number of food sources/employed bees/onlooker bees

    # initialize population
    population = initialize_population(D, bounds_lower, bounds_upper, Np)

    # calculate objective values and fitness values for population
    objective_values = compute_objective(population, f)
    fitness_values = compute_fitness(objective_values)

    # initialize trial vector for population
    trial = zeros(Np, 1)
    best_fitness = 0
    best_fitness_tracker = []
    optimal_solution = []
    populations = []
    for iterations in 1:T

        ## EMPLOYED BEE PHASE
        population, fitness_values, objective_values, trial = employed_bee_phase(population,
                                                                    bounds_lower, bounds_upper, trial, Np, f)

        ## ONLOOKER BEE PHASE
        population, fitness_values, objective_values, trial = onlooker_bee_phase(population,
                                                                    bounds_lower, bounds_upper, trial, Np, f)

        ## SCOUTING PHASE
        if maximum(fitness_values) > best_fitness
            best_fitness = maximum(fitness_values)
            ind = argmax(fitness_values)
            optimal_solution = population[ind]

        end

        population, fitness_values, objective_values, trial = scouting_phase(population, bounds_lower, bounds_upper,
                                                                    D,trial, fitness_values, objective_values, limit, f)

        if maximum(fitness_values) > best_fitness
            best_fitness = maximum(fitness_values)
            ind = argmax(fitness_values)
            optimal_solution = population[ind]

        end
        populations = append!(populations, [population])
        best_fitness_tracker = append!(best_fitness_tracker, best_fitness)
    end

    return optimal_solution, populations, best_fitness_tracker
end

"""
    sphere(x)

This is computing the sphere function value for the vector x.

Input
- 'x': input vector for the sphere function

Output
- 'output': output value for the sphere function

## Examples

```julia-repl
julia> sphere(4)
16
julia> sphere([4,5])
41
```
"""
function sphere(x)
    return sum(x.^2)
end

"""
    ackley(x; a=20, b=0.2, c=2π)

This is computing the ackley function value for the vector x.

Input
- 'x': input vector for the ackley function

Output
- 'output': output value for the ackley function

## Examples

```julia-repl
julia> ackley(4)
11.013420717655569
julia> ackley([4,5])
11.913518152857637
```
"""
function ackley(x; a=20, b=0.2, c=2π)
    d = length(x)
    return -a * exp(-b*sqrt(sum(x.^2)/d)) -
        exp(sum(cos.(c .* x))/d) + a + exp(1)
end

"""
    rosenbrock(x; a=1, b=5)

This is computing the rosenbrock function value for the vector x.
Watch out! This function always needs a 2-element Array as input

Input
- 'x': input vector for the rosenbrock function

Output
- 'output': output value for the rosenbrock function

## Examples

```julia-repl
julia> rosenbrock([4,5])
12109
```
"""
function rosenbrock(x; a=1, b=100)
    # 2 dimensions!
    @assert length(x) == 2 "The rosenbrock function is restricted to two decision variables"
    return (a-x[1])^2 + b*(x[2]-x[1]^2)^2
end

"""
    branin((x1, x2); a=1, b=5.1/(4pi^2), c=5/pi, r=6, s=10, t=1/8pi)

This is computing the branin function value for the vector x.
Watch out! This function always needs a 2-element Array as input

Input
- 'x': input vector for the branin function

Output
- 'output': output value for the branin function

## Examples

```julia-repl
julia> branin([4,5])
14.608661704375713
```
"""
function branin(x; a=1, b=5.1/(4pi^2), c=5/pi, r=6, s=10, t=1/8pi)
    # 2 dimensions!
    @assert length(x) == 2 "The branin function is restricted to two decision variables"
    return a * (x[2] - b * x[1]^2 + c * x[1] - r)^2 + s * (1 - t) * cos(x[1]) + s
end
"""
    rastrigine(x; A=10)

This is computing the rastrigine function value for the input values of x.

Input

- 'x': input vector for the rastrigine function

Output

- 'output': output value for the rastrigine function

## Examples

```julia-repl
julia> rastrigine(4)
16.0
julia> rastrigine([4,5])
41.0
```
"""
function rastrigine(x; A=10)
    return length(x) * A + sum(x.^2 .- A .* cos.(2pi .* x))
end

end


module SearchSortedNearest

export searchsortednearest

# from: https://discourse.julialang.org/t/findnearest-function/4143/5
"""
    searchsortednearest(a, x; by=<transform>, lt=<comparison>, rev=false)

Find the index of (sorted) collection `a` that is nearest (smallest absolute difference) to `x`.
Ties go to the smallest index.

# Examples

    using SearchSortedNearest

    searchsortednearest(1:10, 1.1) == 1
    searchsortednearest(1:10, 1.9) == 2
"""
function searchsortednearest(a, x; by=identity, lt=isless, rev=false)
    idx = searchsortedfirst(a, x; by, lt, rev)
    idx == 1 && return idx
    idx > length(a) && return length(a)
    a[idx] == x && return idx
    idx = lt(abs(by(a[idx]) - by(x)), abs(by(a[idx - 1]) - by(x))) ? idx : idx - 1
    return rev ? length(a) - idx + 1 : idx
end

end


# # Three Dimensions
#
#md # [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__/notebooks/solid_body_rotation.ipynb)
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/notebooks/solid_body_rotation.ipynb)
#
# Simulate the trajectory of an individual point, first in a perfectly circular flow (a.k.a. solid body rotation). Then add a convergent term to obtain a spiraling trajectory, and a constant vertical velocity for the third dimension. These simple flow configurations can be thought of as idealized models e.g. ocean meso-scale eddies.
#
# For additional documentation e.g. see :
# [1](https://JuliaClimate.github.io/IndividualDisplacements.jl/dev/),
# [2](https://JuliaClimate.github.io/MeshArrays.jl/dev/),
# [3](https://docs.juliadiffeq.org/latest/solvers/ode_solve.html),
# [4](https://en.wikipedia.org/wiki/Displacement_(vector))
#
# Exercise examples:
# - make the sinking velocity decrease with time (hint: it increases in the original notebook)
# - change the number of times the particle goes around the origin
#
# ![solid body rotation](https://github.com/JuliaClimate/IndividualDisplacements.jl/raw/master/examples/figs/SolidBodyRotation.gif)

#nb # %% {"slideshow": {"slide_type": "slide"}, "cell_type": "markdown"}
# # 1 Problem Configuration
#
# Here we set up software, grid, flow fields, initial conditions.
#
# ### 1.1 Import Software

using IndividualDisplacements, DataFrames
p=dirname(pathof(IndividualDisplacements))
include(joinpath(p,"../examples/flow_fields.jl"))

#nb # %% {"slideshow": {"slide_type": "subslide"}, "cell_type": "markdown"}
# ### 1.2  Flow Fields
#
# The `simple_flow_field` function (defined in `helper_functions.jl`) defines a simple
# three-dimensional flow field. Exercise: locate `simple_flow_field` and modify the
# flow field parameters (e.g. intensity and sign of the convergent term).

np,nz=16,4 #gridded domain size (horizontal and vertical)

u,v,w=solid_body_rotation(np,nz) #staggered velocity arrays

𝐹=FlowFields(u,u,v,v,0*w,1*w,[0,19.95*2*pi]); #FlowFields data structure

#nb # %% {"slideshow": {"slide_type": "slide"}, "cell_type": "markdown"}
# ### 1.3 Initialize Individuals
#
# Let's just set up one individual at [np*1/3,np*1/3,nz*1/3] in the three-dimensional
# space where the flow fields have been configured

(x,y,z)=(np*1/3,np*1/3,nz*1/3)

𝐼=Individuals(𝐹,x,y,z)

#nb # %% {"slideshow": {"slide_type": "slide"}, "cell_type": "markdown"}
# ### 1.4 A Closer Look (optional)
#
# The above `Individuals` constructor wraps up 𝐹, the initial position, and other needed components
# within 𝐼. **At this point, you can either jump to section 2 or read through this section**
# to learn more about how the details as needed e.g. if you wanted to overide default options
# that were selected for you by the section 1.3 constructor.
#
# Initial position is

📌=[x,y,z]

# and the data structure ([DataFrame](http://juliadata.github.io/DataFrames.jl/stable/))
# to record properties along the individual's path accordingly.

🔴 = DataFrame(ID=Int[], x=Float64[], y=Float64[], z=Float64[], t=Float64[])

# It is the postprocessing function's responsibility to provide the record. It is thus
# important that this intermediary (`postproc`) be consistent with the solver setup (`sol`)
# and the expected record format (`🔴`).

function postproc(sol,𝐹::FlowFields;id=missing,𝑇=missing)
    df=postprocess_xy(sol,𝐹,id=id,𝑇=𝑇)
    #add third coordinate
    z=sol[3,:]
    df.z=z[:]
    return df
end

# The velocity function `🚄` relies only on flow fields obtained from
# `𝐹` (which is defined above) to interpolate velocity at the specified
# space-time position (e.g. those of individuals).

🚄 = dxdt!

# Now that every thing needed to carry out the computation is in place,
# we wrap up the problem configuration in a struct (`Individuals`) which
# links to the initial positions, flow fields, etc. all that will be
# necessary to compute trajectories over time (`∫!(𝐼,𝑇)`).

#assemble as a NamedTuple:
I=(position=📌,record=🔴,velocity=🚄,
postprocessing=postproc,parameters=𝐹)

#construct Individuals from NamedTuple:
𝐼=Individuals(I)

#nb # %% {"slideshow": {"slide_type": "slide"}, "cell_type": "markdown"}
# ## 2 Trajectory Simulations
#
# The `∫!` function call below returns the final positions & updates `𝐼.📌` accordingly. It also records properties observed along the trajectory in `𝐼.🔴`.
# Simple methods to visualize the individual trajectory (plot or movie) are provided at the end.

#nb # %% {"slideshow": {"slide_type": "slide"}, "cell_type": "markdown"}
# ### 2.1 Compute Trajectories

𝑇=(0.0,𝐼.𝑃.𝑇[2])
∫!(𝐼,𝑇)

#nb # %% {"slideshow": {"slide_type": "slide"}, "cell_type": "markdown"}
# ### 2.2 Visualize Trajectories
#
# - define `myplot` convenience function
# - single plot example using `myplot`
# - (generate animation using `myplot`)

#md p=dirname(pathof(IndividualDisplacements))
#md include(joinpath(p,"../examples/recipes_plots.jl"));
#md nt=length(𝐼.🔴.x)

#md myplot(i)=plot(𝐼.🔴.x[1:i],𝐼.🔴.y[1:i],𝐼.🔴.z[1:i],linewidth=2,arrow = 2,
#md     title="Solid body rotation / Spiral example",leg=false,
#md     xaxis="x",yaxis="y",zaxis="z",xlims=(0,np),ylims=(0,np));

#nb # %% {"slideshow": {"slide_type": "subslide"}}
# Single plot example:

#md plt=myplot(nt)
#md scatter!(plt,[📌[1]],[📌[2]],[📌[3]])
#md #scatter!(plt,[𝐼.🔴.x[end]],[𝐼.🔴.y[end]],[𝐼.🔴.z[end]])
#md scatter!(plt,[𝐼.📌[1]],[𝐼.📌[2]],[𝐼.📌[3]])

#nb # %% {"slideshow": {"slide_type": "subslide"}}
# Animation example:
#
# ```
# p=Int(ceil(nt/100))
# anim = @animate for i ∈ 1:p:nt
#     myplot(i)
# end
#
# pth=tempdir()*"/"
# gif(anim, pth*"SolidBodyRotation.gif", fps = 15)
# ```


# ---
# title: 1491. Average Salary Excluding the Minimum and Maximum Salary
# id: problem1491
# author: zhwang
# date: 2022-03-18
# difficulty: Easy
# categories: Array, Sort
# link: <https://leetcode.com/problems/average-salary-excluding-the-minimum-and-maximum-salary/description/>
# hidden: true
# ---
#
# Given an array of **unique** integers `salary` where `salary[i]` is the salary
# of the employee `i`.
#
# Return the average salary of employees excluding the minimum and maximum
# salary.
#
#
#
# **Example 1:**
#
#
#
#     Input: salary = [4000,3000,1000,2000]
#     Output: 2500.00000
#     Explanation: Minimum salary and maximum salary are 1000 and 4000 respectively.
#     Average salary excluding minimum and maximum salary is (2000+3000)/2= 2500
#
#
# **Example 2:**
#
#
#
#     Input: salary = [1000,2000,3000]
#     Output: 2000.00000
#     Explanation: Minimum salary and maximum salary are 1000 and 3000 respectively.
#     Average salary excluding minimum and maximum salary is (2000)/1= 2000
#
#
# **Example 3:**
#
#
#
#     Input: salary = [6000,5000,4000,3000,2000,1000]
#     Output: 3500.00000
#
#
# **Example 4:**
#
#
#
#     Input: salary = [8000,9000,2000,3000,6000,1000]
#     Output: 4750.00000
#
#
#
#
# **Constraints:**
#
#   * `3 <= salary.length <= 100`
#   * `10^3 <= salary[i] <= 10^6`
#   * `salary[i]` is unique.
#   * Answers within `10^-5` of the actual value will be accepted as correct.
#
#
## @lc code=start
using LeetCode

average_salary_excluding_mami(salary::Vector{Int}) = (sum(salary) - minimum(salary) - maximum(salary)) / (length(salary) - 2)
## @lc code=end

## @lc test=start
@testset "1491.average-salary-excluding-the-minimum-and-maximum-salary.jl" begin
    @test average_salary_excluding_mami([4000,3000,1000,2000]) ≈ 2500.00000
    @test average_salary_excluding_mami([1000,2000,3000]) ≈ 2000.00000
    @test average_salary_excluding_mami([6000,5000,4000,3000,2000,1000]) ≈ 3500.00000
    @test average_salary_excluding_mami([8000,9000,2000,3000,6000,1000]) ≈ 4750.00000
end
## @lc test=end

## @lc info=start
# link: [solution to 1491]()
## @lc info=end


using Test, Batched

@testset "checking BatchedMatrix matmul" begin

    A = BatchedMatrix(rand(3, 4, 2, 3))
    B = BatchedMatrix(rand(4, 3, 2, 3))

    C = A * B

    vA = A.parent
    vB = B.parent
    vC = zeros(3, 3, 2, 3)

    for i in 1:2
        for j in 1:3
            vC[:, :, i, j] = vA[:, :, i, j] * vB[:, :, i, j]
        end
    end

    @test vC ≈ C.parent

end

@testset "checking BatchedMatrix matmul (transposed)" begin

    A = BatchedMatrix(rand(3, 4, 2, 3))
    B = BatchedMatrix(rand(3, 4, 2, 3))

    C = A * transpose(B)

    vA = A.parent
    vB = B.parent
    vC = zeros(3, 3, 2, 3)

    for i in 1:2
        for j in 1:3
            vC[:, :, i, j] = vA[:, :, i, j] * transpose(vB[:, :, i, j])
        end
    end

    @test vC ≈ C.parent

end

@testset "checking BatchedMatrix matmul (scale)" begin

    A = BatchedMatrix(rand(3, 4, 2, 3))
    B = BatchedUniformScaling(rand(2, 3))

    C = B * A

    vA = A.parent
    vB = B.scalars
    vC = zeros(3, 4, 2, 3)

    for i in 1:2
        for j in 1:3
            vC[:, :, i, j] = vB[i, j] * vA[:, :, i, j]
        end
    end

    @test vC ≈ C.parent

end

@testset "checking BatchedMatrix matmul (scale & transpose)" begin

    A = BatchedMatrix(rand(3, 4, 2, 3))
    B = BatchedUniformScaling(rand(2, 3))

    C = B * transpose(A)

    vA = A.parent
    vB = B.scalars
    vC = zeros(4, 3, 2, 3)

    for i in 1:2
        for j in 1:3
            vC[:, :, i, j] = vB[i, j] * transpose(vA[:, :, i, j])
        end
    end

    @test vC ≈ C.parent
end


export rearrange
export drop_powers
export get_averaged_equations
export d
export substitute_all
export get_all_terms
export var_name

"The derivative of f w.r.t. x of degree deg"
function d(f::Num, x::Num, deg=1)
    isequal(deg,0) ? f : (Differential(x)^deg)(f)
end
d(funcs::Vector{Num}, x::Num, deg=1) = [d(f, x, deg) for f in funcs]

"Declare a variable in the the currently active namespace"
 function declare_variable(name::String)
    var_sym = Symbol(name)
    @eval($(var_sym) = first(@variables $var_sym))
    return eval(var_sym)
 end

 "Declare a variable that is a function of another variable in the the current namespace"
 function declare_variable(name::String, independent_variable::Num)
    # independent_variable = declare_variable(independent_variable) convert string into Num
    var_sym = Symbol(name)
    new_var = @variables $var_sym(independent_variable)
    @eval($(var_sym) = first($new_var)) # store the variable under "name" in this namespace
    return eval(var_sym)
 end

 "Return the name of a variable (excluding independent variables)"
 var_name(x::Num) = var_name(x.val)
 var_name(x::Term) = String(Symbolics._toexpr(x).args[1])
 var_name(x::Sym) = String(x.name)

 """
 $(TYPEDSIGNATURES)

 Perform substitutions in `rules` on `x`.
 `include_derivatives=true` also includes all derivatives of the variables of the keys of `rules`.
 """
 function substitute_all(x::Union{Num, Equation}, rules::Dict; include_derivatives=true)
    if include_derivatives
        rules = merge(rules, Dict([Differential(var) => Differential(rules[var]) for var in keys(rules)]))
    end
    return substitute(x, rules)
 end

 "Variable substitution - dictionary"
function substitute_all(dict::Dict, rules::Dict)
    new_keys = substitute_all.(keys(dict), rules)
    new_values = substitute_all.(values(dict), rules)
    return Dict(zip(new_keys, new_values))
end

substitute_all(v::Union{Array{Num}, Array{Equation}}, rules::Union{Dict, Pair, Vector}) = [substitute_all(x, rules) for x in v]
substitute_all(x::Union{Num, Equation}, rules::Union{Pair, Vector, Dict}) = substitute_all(x, Dict(rules))
substitute_all(x, rules::OrderedDict) = substitute_all(x, Dict(rules))
substitute_all(x::Complex{Num}, rules) = substitute_all(Num(x.re.val.arguments[1]), rules)
substitute_all(x, rules) = substitute_all(Num(x), rules)

"""
$(SIGNATURES)
Remove parts of `expr` where the combined power of `vars` is => `deg`.

# Example
```julia-repl
julia> @variables x,y;
julia>drop_powers((x+y)^2, x, 2)
y^2 + 2*x*y
julia>drop_powers((x+y)^2, [x,y], 2)
0
julia>drop_powers((x+y)^2 + (x+y)^3, [x,y], 3)
x^2 + y^2 + 2*x*y
```
"""
function drop_powers(expr::Num, vars::Vector{Num}, deg::Int)
    @variables ϵ
    subs_expr = deepcopy(expr)
    rules = Dict([var => ϵ*var for var in unique(vars)])
    subs_expr = Symbolics.expand(substitute_all(subs_expr, rules))
    max_deg = max_power(subs_expr, ϵ)
    removal = Dict([ϵ^d =>Num(0) for d in deg:max_deg])
    res = substitute_all(substitute_all(subs_expr, removal), Dict(ϵ => Num(1)))
    Symbolics.expand(res)
    #res isa Complex ? Num(res.re.val.arguments[1]) : res
end

drop_powers(expr::Vector{Num}, var::Num, deg::Int) = [drop_powers(x, var, deg) for x in expr]

# calls the above for various types of the first argument
drop_powers(eq::Equation, var, deg) = drop_powers(eq.lhs, var, deg) .~ drop_powers(eq.lhs, var, deg)
drop_powers(eqs::Vector{Equation}, var, deg) = [drop_powers(eq.lhs, var, deg) .~ drop_powers(eq.rhs, var, deg) for eq in eqs]
drop_powers(expr, var::Num, deg::Int) = drop_powers(expr, [var], deg)
drop_powers(x, vars, deg) = drop_powers(Num(x), vars, deg)

flatten(a) = collect(Iterators.flatten(a))

###
# STUFF BELOW IS MAINLY FOR FOURIER-TRANSFORMING
###

get_independent(x::Num, t::Num) = get_independent(x.val, t)
get_independent(x::Complex{Num}, t::Num) = get_independent(x.re, t) + im*get_independent(x.im, t)
get_independent(x::Union{Term, Sym}, t::Num) = !is_function(x, t) ? x : 0
get_independent(x::Add, t::Num) = sum([get_independent(arg, t) for arg in arguments(x)])
get_independent(x::Pow, t::Num) = !is_function(x.base, t) && !is_function(x.exp, t) ? x : 0
get_independent(x::Div, t::Num) = !is_function(x.den, t) ? get_independent(x.num, t) / x.den : 0
get_independent(x::Mul, t::Num) = prod([get_independent(arg, t) for arg in arguments(x)])
get_independent(v::Vector{Num}, t::Num) = [get_independent(el, t) for el in v]
get_independent(x, t::Num) = x

#=
function get_independent(expr::SymbolicUtils.Add, t::Num)
    primitive_terms = keys(expr.dict)
    total = expr.coeff
    for term in primitive_terms
        total += get_independent(term, t) * expr.dict[term]
    end
    return total
end
=#

"Return all the terms contained in `x`"
get_all_terms(x::Num) = unique(_get_all_terms(Symbolics.expand(x).val))
get_all_terms(x::Equation) = unique(cat(get_all_terms(Num(x.lhs)), get_all_terms(Num(x.rhs)), dims=1))

_get_all_terms(x::Mul) = Num.(SymbolicUtils.arguments(x))
_get_all_terms(x::Div) = Num.([_get_all_terms(x.num)..., _get_all_terms(x.den)...])
_get_all_terms(x) = Num(x)

function _get_all_terms(x::Add)::Vector{Num}
    list = []
    for term in keys(x.dict)
        list = cat(list, _get_all_terms(term), dims=1)
    end
    return list
end

function is_harmonic(x::Num, t::Num)
    all_terms = get_all_terms(x)
    t_terms = setdiff(all_terms, get_independent(all_terms, t))
    isempty(t_terms) && return true
    trigs = is_trig.(t_terms)

    if !prod(trigs)
        return false
    else
        powers = [max_power(first(term.val.arguments), t) for term in t_terms[trigs]]
        return unique(powers) == [1]
    end
end

is_harmonic(x::Equation, t::Num) = is_harmonic(x.lhs, t) && is_harmonic(x.rhs, t)
is_harmonic(x, t) = is_harmonic(Num(x), Num(t))

"Convert all sin/cos terms in `x` into exponentials."
function trig_to_exp(x::Num)
    all_terms = get_all_terms(x)
    trigs = filter(z -> is_trig(z), all_terms)

    rules = []
    for trig in trigs
        is_pow = trig.val isa Pow # trig is either a trig or a power of a trig
        power = is_pow ? trig.val.exp : 1
        arg = is_pow ? arguments(trig.val.base)[1] : arguments(trig.val)[1]
        type = is_pow ? operation(trig.val.base) : operation(trig.val)

        if type == cos
            term = Complex{Num}((exp(im*arg) + exp(-im*arg))^power * (1//2)^power,0)
        elseif type == sin
            term = (1*im^power)* Complex{Num}( ((exp(-im*arg) - exp(im*arg)))^power *(1//2)^power, 0)
        end
        # avoid Complex{Num} where possible as this causes bugs
        # instead, the Nums store SymbolicUtils Complex types
        term = Num(term.re.val + im*term.im.val)
        append!(rules, [trig => term])
    end

    result = Symbolics.substitute(x, Dict(rules))
    #result = result isa Complex ? Num(first(result.re.val.arguments)) : result
    result = Num(result)
    return result
end

"Return true if `f` is a function of `var`."
is_function(f, var) = any(isequal.(get_variables(f), var))

"Return true if `f` is a sin or cos."
function is_trig(f::Num)
    f = f.val isa Pow ? f.val.base : f.val
    f isa Term && SymbolicUtils.operation(f) ∈ [cos,sin] && return true
    return false
end

"A vector of Sym(0) of length n"
Num_zeros(n::Int64) = [Num(0) for k in 1:n]
Num_zeros(vec::Vector{Any}) = Num_zeros(length(vec))

"""
$(TYPEDSIGNATURES)
Returns the coefficient of cos(ωt) in `x`.
"""
function fourier_cos_term(x, ω, t)
    _fourier_term(x,ω, t, cos)
end

"""
$(TYPEDSIGNATURES)
Returns the coefficient of sin(ωt) in `x`.
"""
function fourier_sin_term(x, ω, t)
    _fourier_term(x,ω, t, sin)
end

function _fourier_term(x::Equation, ω, t, f)
    _fourier_term(x.lhs, ω, t, f) .~ _fourier_term(x.rhs, ω, t, f)
end

"Return the coefficient of f(ωt) in `x` where `f` is a cos or sin."
function _fourier_term(x, ω, t, f)
    term = x * f(ω * t)
    term = trig_reduce(term)
    indep = get_independent(term, t)
    ft = Num(simplify_complex(Symbolics.expand(indep)))
    !isequal(ω, 0) ? 2*ft : ft # extra factor in case ω = 0 !
end

"Expand all sin/cos powers in `x`."
function trig_reduce(x)
    x = Symbolics.expand(x) # open all brackets
    x = trig_to_exp(x)
    x = expand_all(x) # expand products of exponentials
    x= simplify_exp_products(x) # simplify products of exps
    x = exp_to_trig(x)
    x
end

"Return the highest power of `y` occuring in the term `x`."
function max_power(x::Num, y::Num)
    terms = get_all_terms(x)
    powers = power_of.(terms, y)
    maximum(powers)
end

max_power(x::Vector{Num}, y::Num) = maximum(max_power.(x, y))
max_power(x::Complex, y::Num) = maximum(max_power.([x.re, x.im], y))
max_power(x, t) = max_power(Num(x), Num(t))

"Return the power of `y` in the term `x`"
function power_of(x::Num, y::Num)
    y.val isa Sym ? nothing : error("power of " * string(y) * " is ambiguous")
    power_of(x.val, y.val)
end

power_of(x::Pow, y::Sym) = isequal(x.base, y) ? x.exp : 0
power_of(x::Sym, y::Sym) = isequal(x, y) ? 1 : 0
power_of(x::Term, y::Sym) = 0
power_of(x, y) = 0


module FisherPlot

using LaTeXStrings
using CairoMakie
using Makie
using Base: @kwdef

function ellipseparameterization(a::Float64, b::Float64, θ::Float64)
    t = LinRange(0,2π, 200)
    x = Array(a .* sin.(θ) .* cos.(t) + b .* cos.(θ) .* sin.(t))
    y = Array(a .* cos.(θ) .* cos.(t) - b .* sin.(θ) .* sin.(t))
    return x, y
end

function gaussian(μ::Float64, σ::Float64, x)
    return 1/(sqrt(2π*σ^2))*exp(-0.5*(x-μ)^2/σ^2)
end

function ellipseparameters(covmatrix::Matrix{Float64}, i::Int64, j::Int64)
    σi = sqrt(covmatrix[i,i])
    σj = sqrt(covmatrix[j,j])
    σij = covmatrix[i,j]
    θ = (atan(2σij,(σi^2-σj^2)))/2
    a = sqrt((σi^2+σj^2)/2+sqrt(((σi^2-σj^2)^2)/4+σij^2))
    if i == j
        b = 0.
    else
        b = sqrt((σi^2+σj^2)/2-sqrt(((σi^2-σj^2)^2)/4+σij^2))
    end
    return σi, σj, a, b, θ
end

function preparecanvas(LaTeX_array, limits, ticks, probes, colors, PlotPars::Dict)
    matrix_dim = length(LaTeX_array)
    #TODO: add the textsize to PlotPars
    figure = Makie.Figure(textsize = 40, font = PlotPars["font"])
    for i in 1:matrix_dim
        for j in 1:i
            if i == j
                ax = Axis(figure[i,i],
                    width = PlotPars["sidesquare"], height = PlotPars["sidesquare"],
                    xticklabelsize = PlotPars["dimticklabel"],
                    yticklabelsize = PlotPars["dimticklabel"], yaxisposition = (:right),
                    xlabel = L"%$((LaTeX_array)[i])", xlabelsize = PlotPars["parslabelsize"],
                    ylabel = L"P/P_\mathrm{max}",ylabelsize = PlotPars["PPmaxlabelsize"], yticks = [0,1],
                    xticklabelrotation = PlotPars["xticklabelrotation"],
                    xticks = ([ticks[i,1], 0.5*(ticks[i,1]+ticks[i,2]), ticks[i,2]],
                        [string(myi) for myi in round.([ticks[i,1], 0.5*(ticks[i,1]+ticks[i,2]), ticks[i,2]], sigdigits = 3)]))
                Makie.ylims!(ax, (-0.0,1.05))
                Makie.xlims!(ax, (limits[i,1],limits[i,2]))
                Makie.hideydecorations!(ax, ticks = false, ticklabels = false, label = false)
                if i != matrix_dim
                    ax.alignmode = Mixed(right = MakieLayout.Protrusion(0), bottom = MakieLayout.Protrusion(0), top= MakieLayout.Protrusion(0))
                    hidexdecorations!(ax, ticks = true, ticklabels = true,  label = true)
                else
                    hidexdecorations!(ax, ticks = false, ticklabels = false,  label = false)
                end
            else
                ax = Axis(figure[i,j], width = PlotPars["sidesquare"], height = PlotPars["sidesquare"],
                    xticklabelsize = PlotPars["dimticklabel"], yticklabelsize = PlotPars["dimticklabel"],
                    ylabel = L"%$(LaTeX_array[i])", xlabel = L"%$((LaTeX_array)[j])",
                    ylabelsize = PlotPars["parslabelsize"], xlabelsize = PlotPars["parslabelsize"], xticklabelrotation = PlotPars["xticklabelrotation"],
                    yticks = ([ticks[i,1], 0.5*(ticks[i,1]+ticks[i,2]), ticks[i,2]],
                        [string(myi) for myi in round.([ticks[i,1], 0.5*(ticks[i,1]+ticks[i,2]), ticks[i,2]], sigdigits = 3)]),
                    xticks = ([ticks[j,1], 0.5*(ticks[j,1]+ticks[j,2]), ticks[j,2]],
                        [string(myi) for myi in round.([ticks[j,1], 0.5*(ticks[j,1]+ticks[j,2]), ticks[j,2]], sigdigits = 3)]),
                yticklabelpad=8)
                Makie.ylims!(ax, (limits[i,1],limits[i,2]))
                Makie.xlims!(ax, (limits[j,1],limits[j,2]))
                if i == matrix_dim
                    hidexdecorations!(ax, ticks = false, ticklabels = false,  label = false)
                else
                    hidexdecorations!(ax, ticks = true, ticklabels = true,  label = true)
                end
                if j == 1
                    hideydecorations!(ax, ticks = false, ticklabels = false,  label = false)
                    Legend(figure[1,matrix_dim],
                    [PolyElement(color = color, strokecolor = color, strokewidth = 1) for color in colors],
                    probes,
                    tellheight = false, tellwidth = false, rowgap = 10,
                    halign = :right, valign = :top, framecolor = :black, labelsize =55, patchsize = (70, 40), framevisible = true)
                else
                    hideydecorations!(ax, ticks = true, ticklabels = true,  label = true)
                    ax.alignmode = Mixed(right = MakieLayout.Protrusion(0), bottom = MakieLayout.Protrusion(0), top = MakieLayout.Protrusion(0))
                end
            end
        end
    end
    Makie.resize!(figure.scene, figure.layout.layoutobservables.reportedsize[]...)
    return figure
end

function drawgaussian!(canvas, σ, i, central, color)
    ax = canvas[i,i]
    x = Array(LinRange(-4σ+central,4σ+central, 200))
    Makie.lines!(ax, x, gaussian.(central, σ, x)./gaussian.(central, σ, central), color = color, linewidth = 4)
    Makie.band!(ax, x, 0, gaussian.(central, σ, x)./gaussian.(central, σ, central) , color=(color, 0.2))
    x = Array(LinRange(-σ+central,σ+central, 200))
    Makie.band!(ax, x, 0, gaussian.(central, σ, x)./gaussian.(central, σ, central) , color=(color, 0.4))
end

function drawellipse!(canvas, i, j, x, y, central_values, color)
    ax = canvas[i,j]

    Makie.lines!(ax, x .+ central_values[j], y .+ central_values[i], color = color, linewidth = 4)
    Makie.lines!(ax, 3x .+ central_values[j], 3y .+ central_values[i], color = color, linewidth = 4)

    Makie.band!(ax, x .+ central_values[j], 0, y .+ central_values[i], color=(color, 0.4))
    Makie.band!(ax, 3x .+ central_values[j], 0, 3y .+ central_values[i], color=(color, 0.2))
end

function paintcorrmatrix!(canvas, central_values, corr_matrix, color)
    for i in 1:length(central_values)
        for j in 1:i
            if i == j
                drawgaussian!(canvas, sqrt(corr_matrix[i,i]), i, central_values[i], color)
            else
                σi, σj, a, b, θ = ellipseparameters(corr_matrix, i,j)
                x,y = ellipseparameterization(a, b, θ)
                drawellipse!(canvas, i, j, x, y, central_values, color)
            end
        end
    end
end

function save(filename, obj)
    Makie.save(filename, obj)
end

end # module


using Plots

function midpoint_method(f::Function, x, u, t, h)
    k1, y = f(x, u, t)
    k2, _ = f(x + h/2*k1, u, t + h/2)
    k3, _ = f(x - h*k1 + 2*h*k2, u, t + h)
    x_VPG = x + h*k2
    d_VPG = h/6*(k1 - 2*k2 + k3)
    return x_VPG, y, d_VPG
end

function pt1(x, u)
    Tm = 10
    xdot = -1/Tm * x + 1/Tm * u
    y = x
    return xdot, y
end

function subtractor(u1, u2)
    return u1 - u2
end

function hysteresis(u)
    h_e = 0.085
    h_a = 0.065
    if u >= h_e
        y = 1;
    elseif u <= -h_e
        y = -1;
    else
        y = 0;
    end

    if i > 1
        if hys_state[i - 1] == 1 && u <= h_a
            y = 0;
        elseif hys_state[i - 1] == -1 && u >= -h_a
            y = 0;
        end
    end
    append!(hys_state, y)
    return y
end

function sys_PRM(x, u, t)
    _, y1::Float64 = pt1(x, 0)
    y2::Float64 = subtractor(u, y1)
    y3::Float64 = hysteresis(y2)
    xdot1, _ = pt1(x, y3)
    return xdot1, [y1 y2 y3]
end

t_step = 1;
Tm = 10
t0 = 0
tf = 20
epsilon = 1e-3
h_max = 2*Tm
h_min = 12*epsilon
u_values = []
x_values = []
y_values = Vector{Float64}()
d_values = []
t_values = []
h_values = []
hys_state = []
x = 0
h = 0.2
t = 0
i = 1
while t <= tf
    if t < t_step
        global u = 0
    else
        global u = 0.17
    end

    global x, y, d = midpoint_method(sys_PRM, x, u, t, h);
    global h = 0.9*h*min(max((epsilon/max(abs(d)))^(1/3), 0.3), 2)
    global h = min(max(h, h_min), h_max)

    append!(x_values, x)
    append!(y_values, y)
    append!(d_values, d)
    append!(h_values, h)
    append!(t_values, t)

    global t = t + h
    global i = i + 1
end

y_values = reshape(y_values, (3, :))

p1 = plot(t_values, y_values[1, :])
p2 = plot(t_values, y_values[2, :])
p3 = plot(t_values, y_values[3, :])
p4 = plot(t_values, h_values)
p5 = plot(t_values, d_values)

plot(p1, p2, p3, p4, p5, layout=(5, 1))


#++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
# This example demonstrates how to use Knitro to solve the following
# simple mathematical program with equilibrium/complementarity
# constraints(MPEC/MPCC).
#
#  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
#
# The complementarity constraints must be converted so that one
# nonnegative variable complements another nonnegative variable.
#
#  min  (x0 - 5)^2 +(2 x1 + 1)^2
#  s.t.  -1.5 x0 + 2 x1 + x2 - 0.5 x3 + x4 = 2  (c0)
#        3 x0 - x1 - 3 - x5 = 0                 (c1)
#        -x0 + 0.5 x1 + 4 - x6 = 0              (c2)
#        -x0 - x1 + 7 - x7 = 0                  (c3)
#        x2 complements x5
#        x3 complements x6
#        x4 complements x7
#        x0, x1, x2, x3, x4, x5, x6, x7 >= 0
#
# The solution is(1, 0, 3.5, 0, 0, 0, 3, 6), with objective value 17.
#++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

using JuMP, MathOptInterface
using KNITRO
using Test

const MOI = MathOptInterface
b = [-3.0, 4.0, 7.0]
M = [3.0   -1.0   0.0   0.0   0.0;
     -1.0   0.5   0.0   0.0   0.0;
     -1.0  -1.0   0.0   0.0   0.0]

model = JuMP.Model(with_optimizer(KNITRO.Optimizer))

@variable(model, x[1:5] >= 0)
@constraint(model, -1.5 * x[1] + 2.0 * x[2] + x[3] - 0.5 * x[4] + x[5] == 2.0)

@constraint(model, [M*x + b; x[3:5]] in MOI.Complements(3))
@objective(model, Min, (x[1] - 5.0)^2 + (2.0 * x[2] + 1.0)^2)

JuMP.optimize!(model)

@test JuMP.has_values(model)
@test JuMP.termination_status(model) == MOI.LOCALLY_SOLVED
@test JuMP.primal_status(model) == MOI.FEASIBLE_POINT
@test JuMP.value.(x) ≈ [1.0, 0.0, 3.5, 0.0, 0.0] atol=1e-3


using CPLEX

include(joinpath(Pkg.dir("MathProgBase"),"test","linprog.jl"))
linprogtest(CplexSolver(CPX_PARAM_PREIND=0, CPX_PARAM_LPMETHOD=2))

include(joinpath(Pkg.dir("MathProgBase"),"test","mixintprog.jl"))
mixintprogtest(CplexSolver())

include(joinpath(Pkg.dir("MathProgBase"),"test","quadprog.jl"))
quadprogtest(CplexSolver())
socptest(CplexSolver())

include(joinpath(Pkg.dir("MathProgBase"),"test","linproginterface.jl"))
linprogsolvertest(CplexSolver())

include(joinpath(Pkg.dir("MathProgBase"),"test","conicinterface.jl"))
coniclineartest(CplexSolver())
conicSOCtest(CplexSolver())
conicSOCRotatedtest(CplexSolver())
conicSOCINTtest(CplexSolver())


using LightGraphs
using LinearAlgebra
using StatsBase
using GraphPlot
using Optim
using ReverseDiff
using ForwardDiff

include("src/estimating_functions.jl")
include("src/generators.jl")

n = 50
beta = zeros(n)
for i in sample(collect(1:n), Int(floor(n / 10)))
    beta[i] = sqrt(log(n))
end
mu = -log(n)

g = sparse_beta_graph(beta, mu)

layout = (args...) -> spring_layout(args...; C=9)
gplot(g, layout=layout)

# Parameter Estimation

function sparse_beta_loglik(g::Graph, mu, beta)
    beta = exp.(log.(beta))
    tot = 0
    for i in 1:n
        tot -= degree(g)[i] * beta[i]
        for j in i:n
            tot += log(1 + exp(mu + beta[i] + beta[j]))
        end
    end
    return -(g.ne * mu) + tot
end

function sparse_beta_loglik_rp(g::Graph, mu, gamma)
    beta = exp.(gamma)
    tot = 0
    for i in 1:n
        tot -= degree(g)[i] * beta[i]
        for j in i:n
            tot += log(1 + exp(mu + beta[i] + beta[j]))
        end
    end
    return -(g.ne * mu) + tot
end

f = mubeta -> sparse_beta_loglik_rp(g, mubeta[1], mubeta[2:end])

tape = ReverseDiff.GradientTape(f, vcat(-1.5, repeat([1], n)))
(storage, params) -> ReverseDiff.gradient!(storage, tape, params)

optimize(
    mubeta -> sparse_beta_loglik_rp(g, mubeta[1], mubeta[2:end]),
    (storage, params) -> ReverseDiff.gradient!(storage, tape, params),
    # repeat([0], n+1),
    # repeat([Inf], n+1),
    vcat(-1.5, repeat([1], n)),
    LBFGS(),
    Optim.Options(show_every=1))

# # optimize(mubeta -> sparse_beta_loglik(g, mubeta[1], mubeta[2:end]), vcat(-1.5, repeat([0], n)))

# tape = ReverseDiff.GradientTape(mubeta -> sparse_beta_loglik(g, mubeta[1], mubeta[2:end]), vcat(-1.5, repeat([0], n)))
# f(mubeta) = sparse_beta_loglik(g, mubeta[1], mubeta[2:end])
# f_tape = ReverseDiff.GradientTape(f, vcat(-1.5, repeat([0], n)))
# g!(storage, params) = ReverseDiff.gradient!(storage, f_tape, params)

# optimize(
#     f,
#     g!(storage, params),
#     vcat(-1.5, repeat([0], n)),
#     LBFGS(),
#     Optim.Options(
#         show_every=1))


using Test
using DECAES
using DECAES.NNLS
using LinearAlgebra

function verify_solution(m, n)
    A = randn(m, n)
    b = randn(m)
    work = NNLSWorkspace(A, b)
    nnls!(work)

    GC.@preserve work begin
        x  = NNLS.solution(work)
        w  = NNLS.dual(work)
        i₊ = NNLS.components(work)
        m₊ = NNLS.ncomponents(work)
        m₀ = m - m₊
        i₀ = setdiff(1:length(x), i₊)
        x₊, x₀ = x[i₊], x[i₀]
        w₀, w₋ = w[i₊], w[i₀]
        A₊, A₀ = A[:,i₊], A[:,i₀]
        L  = NNLS.choleskyfactor(work, Val(:L))
        U  = NNLS.choleskyfactor(work, Val(:U))

        # Internals
        @test isperm(work.idx)

        # Solution
        @test all(>(0), x₊)
        @test all(==(0), x₀)
        @test x₊ ≈ (A₊'A₊) \ (A₊'b)
        @test x₊ ≈ A₊ \ b

        # Dual (i.e. gradient)
        if m₊ < m
            # Less components than rows; in general solution is not exact and gradient has negative components
            @test all(<(0), w₋)
            @test all(==(0), w₀)
            @test w₋ ≈ A₀'b - A₀'A₊ * x₊
        else
            # At least as many components as rows; solution is exact, gradient is zero
            @test all(==(0), w)
        end

        # Cholesky factors
        @test A₊'A₊ ≈ L'L
        @test A₊'A₊ ≈ U*U'
    end
    work
end

@testset "NNLS" begin
    for m in [1,2,5,8,16,25,32,50], n in [1,2,5,8,16,25,32,50]
        verify_solution(m, n)
    end
end


using Contour, GLVisualize, GeometryTypes, GLAbstraction, Colors, FileIO

if !isdefined(:runtests)
    window = glscreen()
    timesignal = bounce(linspace(0.0, 1.0, 360))
end
# create a rotation from the time signal
rotation = map(timesignal) do t
    rotationmatrix_z(Float32(t*2pi)) # -> 4x4 Float32 rotation matrix
end

xrange = -5f0:0.02f0:5f0
yrange = -5f0:0.02f0:5f0

z = Float32[sin(1.3*x)*cos(0.9*y)+cos(.8*x)*sin(1.9*y)+cos(y*.2*x) for x in xrange, y in yrange]
mini = minimum(z)
maxi = maximum(z)
color_ramp = map(x->RGBA{Float32}(x, 1.0), colormap("Blues"))

for h in mini:0.2f0:maxi
    c = contour(xrange, yrange, z, h)
    for elem in c.lines
        points = map(elem.vertices) do p
            Point3f0(p, h)
        end
        line_renderable = visualize(
            points, :lines,
            color=color_lookup(color_ramp, h, mini, maxi),
            model=rotation
        )
        _view(line_renderable, window, camera=:perspective)
    end
end

if !isdefined(:runtests)
    renderloop(window)
end


"""
    LVmodel(N::Int64 = 20, tSpan::Vector{<:Real} = [0, 50]; β = 1, h = 0.1, random = false)

Generates a Lotka-Volterra model for N species over the given time span.
The output is a time series of abundance values that can be stored as two separate objects:
`t, Xapp = LVmodel()`

# Arguments
- `N::Int64`: the size of the model community, that is, the number of species in the system.
- `tSpan::Vector{<:Real}`: the time span along which computation is performed.
- `β::Real`: the order of derivation for all the differential equations of the LV model.
- `h::Real`: the step size for correction.
- `random::Bool`: whether the function should return a randomised (`true`) or constant (`false`) model.
"""
function LVmodel(N::Int64 = 20, tSpan::Vector{<:Real} = [0, 50]; β = 1, h = 0.1, random = false)

    # order of derivatives
    β = ones(N) * β

    if random == false

        Random.seed!(1234)

    end

    # initial abundances
    X0 = 2 * rand(N)

    if random == false

        Random.seed!(1234)

    end

    # define system of differential equations
    par = [2, 2 * rand(N), rand(N), 4 * rand(N, N), N]

    # evaluate numerical solution
    t, Xapp = FDEsolver(lv, tSpan, X0, β, par, h = h)

    return t, Xapp'

end

function lv(t, x, par)

    l = par[1] # Hill coefficient
    b = par[2] # growth rates
    k = par[3] # death rates
    K = par[4] # inhibition matrix
    N = par[5] # number of species

    F = zeros(N)

    for i in 1:N

        # inhibition functions
        f = prod(K[i, 1:end .!= i].^l ./ (K[i, 1:end .!= i] .^ l .+ x[1:end .!= i].^l))

        # ode
        F[i] = x[i] .* (b[i] .* f .- k[i] .* x[i])

    end

    return F

end

"""
    abundance_plot(se::SummarizedExperiment, assay_name::String)

Plots a time series of the abundance for each feature (row) of an assay in
a SummarizedExperiment object. The corresponding should include a column
named 'time' containing an array of sampling times

# Arguments
- `se::SummarizedExperiment`: the experiment object of interest.
- `assay_name::String`: the name of the assay to base the visualisation on.
"""
function abundance_plot(se::SummarizedExperiment, assay_name::String)

    if sum("time" .== names(coldata(se))) == 0

        error("coldata(se) should include a column named 'time' containing an array of sampling times; for example, add it with coldata(se).time = 1:10 if there are 10 samples.")

    end

    if sum(ismissing.(assay(se, assay_name))) != 0

        error("the input assay contains rows with missing values; they should be either removed or replaced with reasonable numerical values")

    end

    labels = reshape(se.rowdata.name, (1, length(se.rowdata.name)))

    p = plot(se.coldata.time, Number.(assay(se, assay_name))',
             label = labels, legend_position = :outerleft,
             xaxis = "Time", yaxis = "Abundance")

    return(p)

end

function size(se::SummarizedExperiment)

    nrow(se.rowdata), nrow(se.coldata)

end

nrow(se::SummarizedExperiment) = nrow(se.rowdata)
ncol(se::SummarizedExperiment) = nrow(se.coldata)

size(se::SummarizedExperiment, dim::Int64) = ifelse(dim == 1, nrow(se), ncol(se))


using StochasticProcesses

bm = BrownianMotion()
t = linspace(0, 1, 10000)
sim(bm, t, 100)

println("sim(BrownianMotion) 10000x10000:")
for i in 1:5
  @time sim(bm, t, 10000)
end


using CollectiveSpins, QuantumOptics
using Test

@testset "reducedspin" begin

# Test creation
N = 4
b_full = SpinBasis(1//2)^N
b_red1 = CollectiveSpins.ReducedSpinBasis(N, 1)
b_red2 = CollectiveSpins.ReducedSpinBasis(N, 2)

# Operators
σp = [embed(b_full, i, sigmap(SpinBasis(1//2))) for i=1:N]

# Get all single excitation and two excitation states
GS_full = tensor([spindown(SpinBasis(1//2)) for i=1:N]...)
single_exc = Ket[GS_full]
two_exc = Ket[]
for i=1:N
    psi_ = σp[i]*GS_full
    push!(single_exc, normalize(psi_))
    for j=i+1:N
        push!(two_exc, normalize(σp[j]*psi_))
    end
end

# Subspace and projectors
b_sub1 = SubspaceBasis(b_full, single_exc)
b_sub2 = SubspaceBasis(b_full, [single_exc; two_exc])
@test length(b_red1) == length(b_sub1)
@test length(b_red2) == length(b_sub2)

P1 = projector(b_sub1, b_full)
P2 = projector(b_sub2, b_full)

psi1 = P1*single_exc[2]
psi2 = CollectiveSpins.reducedspinstate(b_red1, [1])
@test psi1 != psi2
@test psi1.data == psi2.data

psi1 = P2*two_exc[end] # Corresponds to state where N-1 and N are excited
psi2 = CollectiveSpins.reducedspinstate(b_red2, [N-1,N])
@test psi1 != psi2
@test psi1.data == psi2.data

# Test matrix elements on reduced bases
sz(i) = embed(b_full, i, sigmaz(SpinBasis(1//2)))
sp(i) = embed(b_full, i, sigmap(SpinBasis(1//2)))
for i=1:N
    @test CollectiveSpins.reducedsigmap(b_red1, i).data == sparse(P1*sp(i)*P1').data
    @test CollectiveSpins.reducedsigmam(b_red1, i).data == sparse(P1*sp(i)'*P1').data
    @test CollectiveSpins.reducedsigmap(b_red2, i).data == sparse(P2*sp(i)*P2').data
    @test CollectiveSpins.reducedsigmam(b_red2, i).data == sparse(P2*sp(i)'*P2').data
    @test CollectiveSpins.reducedsigmaz(b_red1, i).data == sparse(P1*sz(i)*P1').data
    @test CollectiveSpins.reducedsigmaz(b_red2, i).data == sparse(P2*sz(i)'*P2').data
end

i=2
j=4
tmp = CollectiveSpins.reducedsigmap(b_red2,i)*CollectiveSpins.reducedsigmap(b_red2,j)*CollectiveSpins.reducedspinstate(b_red2,[])
ind = findfirst(!iszero, tmp.data)
@test b_red2.indexMapper[findfirst(y->y[2]==ind, b_red2.indexMapper)][1] == [i,j]

# Test transitions
@test CollectiveSpins.reducedspintransition(b_red2, [i,j], [i-1,j-1]) == CollectiveSpins.reducedspintransition(b_red2, [i,j], [j-1,i-1])
@test_throws BoundsError CollectiveSpins.reducedspintransition(b_red2, [i+1,j+1], [1,1])

# Test a full example
N = 8     # Number of spins
M = 3     # Maximum number of excitations
a = 0.3   # spin-spin distance
eta = 0.5 # Drive amplitude
geometry = CollectiveSpins.geometry.chain(a, N)

e_list = [[0,0,1] for i=1:N]   # Quantization axis
S = CollectiveSpins.SpinCollection(geometry, e_list)

b = CollectiveSpins.ReducedSpinBasis(N, M)

# Drive
sx(j) = CollectiveSpins.reducedsigmax(b, j)
H_drive = eta*sum(sx(j) for j=1:N)

# Dipole-Dipole
H_dipole = CollectiveSpins.reducedspin.Hamiltonian(S, M)
Gammas, Jumps = CollectiveSpins.reducedspin.JumpOperators(S, M)

# Time Evolution
T = [0:0.05:3.;]
psi0 = CollectiveSpins.reducedspinstate(b, []) # start in the ground state
tout, rho = QuantumOptics.timeevolution.master(T, psi0, H_drive + H_dipole, Jumps; rates=Gammas)

# Test sigmaz
N = 4
M = 4
b_red = CollectiveSpins.ReducedSpinBasis(N,M)
SZ = sum(CollectiveSpins.reducedsigmaz(b_red,j) for j=1:N)

@test sum(SZ.data)==0
sz_diag = [SZ.data[i,i] for i=1:2^N]
for m=0:N
    inds = findall(isequal(-N+2m),sz_diag)
    @test length(inds) == binomial(N,m)
    state = CollectiveSpins.reducedspinstate(b_red, [1:m;])
    @test expect(SZ,state) == -N+2m
end

# Test with lower bound
b_red1 = CollectiveSpins.ReducedSpinBasis(N, 1, 1)
b_red2 = CollectiveSpins.ReducedSpinBasis(N, 2, 1)

# Subspace and projectors
tmp = copy(single_exc)
single_exc = tmp[2:end] # Remove ground state
b_sub1 = SubspaceBasis(b_full, single_exc)
b_sub2 = SubspaceBasis(b_full, [single_exc; two_exc])
@test length(b_red1) == length(b_sub1)
@test length(b_red2) == length(b_sub2)

P1 = projector(b_sub1, b_full)
P2 = projector(b_sub2, b_full)

psi1 = P1*single_exc[1]
psi2 = CollectiveSpins.reducedspinstate(b_red1, [1])
@test psi1 != psi2
@test psi1.data == psi2.data

psi1 = P2*two_exc[end] # Corresponds to state where N-1 and N are excited
psi2 = CollectiveSpins.reducedspinstate(b_red2, [N-1,N])
@test psi1 != psi2
@test psi1.data == psi2.data

# Test matrix elements on reduced bases
sz(i) = embed(b_full, i, sigmaz(SpinBasis(1//2)))
sp(i) = embed(b_full, i, sigmap(SpinBasis(1//2)))
for i=1:N
    @test CollectiveSpins.reducedsigmap(b_red2, i).data == sparse(P2*sp(i)*P2').data
    @test CollectiveSpins.reducedsigmam(b_red2, i).data == sparse(P2*sp(i)'*P2').data
    @test CollectiveSpins.reducedsigmaz(b_red1, i).data == sparse(P1*sz(i)*P1').data
    @test CollectiveSpins.reducedsigmaz(b_red2, i).data == sparse(P2*sz(i)'*P2').data
end

SZ1 = sum(CollectiveSpins.reducedsigmaz(b_red1, i) for i=1:N)
@test all([SZ1.data[i,i] for i=1:N] .== -N+2)
SZ2 = sum(CollectiveSpins.reducedsigmaz(b_red2, i) for i=1:N)
@test all([SZ2.data[i,i] for i=1:length(b_red2)] .== [[-2.0 for i=1:N]; zeros(binomial(N,2))])

end #testset