@doc raw"""
    gradient_descent(M, F, ∇F, x)

perform a gradient_descent ``x_{k+1} = \mathrm{retr}_{x_k} s_k\nabla f(x_k)`` with
different choices of ``s_k`` available (see `stepsize` option below).

# Input
* `M` – a manifold ``\mathcal M``
* `F` – a cost function ``F\colon\mathcal M\to\mathbb R`` to minimize
* `∇F` – the gradient ``\nabla F\colon\mathcal M\to T\mathcal M`` of F
* `x` – an initial value ``x ∈ \mathcal M``

# Optional
* `stepsize` – ([`ConstantStepsize`](@ref)`(1.)`) specify a [`Stepsize`](@ref)
  functor.
* `retraction_method` – (`ExponentialRetraction()`) a `retraction(M,x,ξ)` to use.
* `stopping_criterion` – ([`StopWhenAny`](@ref)`(`[`StopAfterIteration`](@ref)`(200), `[`StopWhenGradientNormLess`](@ref)`(10.0^-8))`)
  a functor inheriting from [`StoppingCriterion`](@ref) indicating when to stop.
* `return_options` – (`false`) – if activated, the extended result, i.e. the
    complete [`Options`](@ref) are returned. This can be used to access recorded values.
    If set to false (default) just the optimal value `x_opt` if returned
...
and the ones that are passed to [`decorate_options`](@ref) for decorators.

# Output
* `x_opt` – the resulting (approximately critical) point of gradientDescent
OR
* `options` - the options returned by the solver (see `return_options`)
"""
function gradient_descent(M::Manifold, F::TF, ∇F::TDF, x; kwargs...) where {TF,TDF}
    x_res = allocate(x)
    copyto!(x_res, x)
    return gradient_descent!(M, F, ∇F, x_res; kwargs...)
end
@doc raw"""
    gradient_descent!(M, F, ∇F, x)

perform a gradient_descent ``x_{k+1} = \mathrm{retr}_{x_k} s_k\nabla f(x_k)`` inplace of `x`
with different choices of ``s_k`` available.

# Input
* `M` – a manifold ``\mathcal M``
* `F` – a cost function ``F\colon\mathcal M\to\mathbb R`` to minimize
* `∇F` – the gradient ``\nabla F\colon\mathcal M\to T\mathcal M`` of F
* `x` – an initial value ``x ∈ \mathcal M``

For more options, especially [`Stepsize`](@ref)s for ``s_k``, see [`gradient_descent`](@ref)
"""
function gradient_descent!(
    M::Manifold,
    F::TF,
    ∇F::TDF,
    x;
    stepsize::Stepsize=ConstantStepsize(1.0),
    retraction_method::AbstractRetractionMethod=ExponentialRetraction(),
    stopping_criterion::StoppingCriterion=StopWhenAny(
        StopAfterIteration(200), StopWhenGradientNormLess(10.0^-8)
    ),
    direction=IdentityUpdateRule(),
    return_options=false,
    kwargs..., #collect rest
) where {TF,TDF}
    p = GradientProblem(M, F, ∇F)
    o = GradientDescentOptions(
        x;
        stopping_criterion=stopping_criterion,
        stepsize=stepsize,
        direction=direction,
        retraction_method=retraction_method,
    )
    o = decorate_options(o; kwargs...)
    resultO = solve(p, o)
    if return_options
        return resultO
    else
        return get_solver_result(resultO)
    end
end
#
# Solver functions
#
function initialize_solver!(p::P, o::O) where {P<:GradientProblem,O<:GradientDescentOptions}
    o.∇ = get_gradient(p, o.x)
    return o
end
function step_solver!(p::P, o::O, iter) where {P<:GradientProblem,O<:GradientDescentOptions}
    s, o.∇ = o.direction(p, o, iter)
    o.x = retract(p.M, o.x, -s .* o.∇, o.retraction_method)
    return o
end
get_solver_result(o::GradientDescentOptions) = o.x


function fib(n)
  x,y = (0,1)
  for i = 1:n x,y = (y, x+y) end
  x
end


function update_position!(x, p, fx, fp)
    i_update = (fx .< fp)
    p[i_update, :] = copy(x[i_update, :])
    fp[i_update, :] = fx[i_update, :]
end

function pso(func, lb, ub; swarmsize=100, ω=0.5, ϕp=0.5, ϕg=0.5, maxiter=100,
             verbose=false, swarmIni=[], particle_output=false)

    obj = x -> func(x)

    minstep=1e-8
    minfunc=1e-8

    # Initialize the particle swarm
    vhigh = abs.(ub .- lb) .*0.2
    vlow = -vhigh .*0.2
    S = swarmsize
    D = length(lb)  # the number of dimensions each particle has

    if length(swarmIni)>=1
        x=swarmIni
    else
        x = lb' .+ rand(S, D) .* (ub .- lb)'  # particle positions
    end
    v = vlow' .+ rand(S, D) .* (vhigh .- vlow)'  # particle velocities
    p = zeros(S, D)  # best particle positions

    Vmin = ones(S,1) * vlow'
    Vmax = ones(S,1) * vhigh'

    fx = [obj(x[i, :]) for i = 1:S]  # current particle function values
    fp = ones(S) * Inf  # best particle function values

    g = copy(x[1, :])  # best swarm position
    fg = Inf  # best swarm position starting value

    # Store particle's best position (if constraints are satisfied)
    update_position!(x, p, fx, fp)

    # Update swarm's best position
    i_min = argmin(fp)
    if fp[i_min] < fg
        g = copy(p[i_min, :])
        fg = fp[i_min]
    end

    # Iterate until termination criterion met
    it = 1
    while it <= maxiter
        rp = rand(S, D)
        rg = rand(S, D)

        # Update the particles' velocities and positions
        ω = ω
        v = ω*v .+ ϕp*rp.*(p .- x) .+ ϕg*rg.*(g' .- x)
        v = min.(v,Vmax)
        v = max.(v,Vmin)

        x += v
        # Correct for bound violations
        maskl = x .< lb'
        masku = x .> ub'
        x = x.*(.~(maskl .| masku)) .+ lb'.*maskl .+ ub'.*masku

        # Update objectives and constraints
        for i = 1:S
            fx[i] = obj(x[i, :])
        end

        # Store particle's best position (if constraints are satisfied)
        update_position!(x, p, fx, fp)

        # Compare swarm's best position with global best position
        i_min = argmin(fp)
        if fp[i_min] < fg
            verbose && println("New best for swarm at iteration $(it): $(p[i_min, :]) $(fp[i_min])")

            p_min = copy(p[i_min, :])
            stepsize = √(sum((g .- p_min).^2))

            if abs.(fg .- fp[i_min]) <= minfunc
                verbose && println("Stopping search: Swarm best objective change less than $(minfunc)")
                return (g, fg, p, fp)
            end
            if stepsize <= minstep
                verbose && println("Stopping search: Swarm best position change less than $(minstep)")
                return (g, fg, p, fp)
            end

            g = copy(p_min)
            fg = fp[i_min]
        end

        verbose && println("Best after iteration $(it): $(g) $(fg)")
        it += 1
    end

    verbose && println("Stopping search: maximum iterations reached --> $(maxiter)")
    return particle_output ? (g, fg, p, fp) : (g, fg)
end


using Distributions
using Random
using LinearAlgebra
using PDMats
using Statistics
using Test
import JSON
import Distributions: _univariate, _multivariate, _rand_params

@testset "matrixvariates" begin
#=
    1. baseline tests
    2. compare 1 x 1 matrix-variate with univariate
    3. compare row/column matrix-variate with multivariate
    4. compare density evaluation against archived Stan output
    5. special, distribution-specific tests
    6. main testing method
    7. run matrix-variate unit tests
=#

#  =============================================================================
#  1. baseline test
#  =============================================================================

#  --------------------------------------------------
#  Check that a random draw from d has the right properties
#  --------------------------------------------------

function test_draw(d::MatrixDistribution, X::AbstractMatrix)
    @test size(d) == size(X)
    @test size(d, 1) == size(X, 1)
    @test size(d, 2) == size(X, 2)
    @test length(d) == length(X)
    @test rank(d) == rank(X)
    @test insupport(d, X)
    @test logpdf(d, X) ≈ log(pdf(d, X))
    @test logpdf(d, [X, X]) ≈ log.(pdf(d, [X, X]))
    @test loglikelihood(d, X) ≈ logpdf(d, X)
    @test loglikelihood(d, [X, X]) ≈ 2 * logpdf(d, X)
    if d isa MatrixFDist
        # Broken since `pdadd` is not defined for SubArray
        @test_broken loglikelihood(d, cat(X, X; dims=3)) ≈ 2 * logpdf(d, X)
    else
        @test loglikelihood(d, cat(X, X; dims=3)) ≈ 2 * logpdf(d, X)
    end
    nothing
end

test_draw(d::MatrixDistribution) = test_draw(d, rand(d))

#  --------------------------------------------------
#  Check that sample quantities are close to population quantities
#  --------------------------------------------------

function test_draws(d::MatrixDistribution, draws::AbstractArray)
    @test isapprox(mean(draws), mean(d), atol = 0.1)
    @test isapprox(cov(hcat(vec.(draws)...)'), cov(d) , atol = 0.1)
    nothing
end

function test_draws(d::LKJ, draws::AbstractArray)
    @test isapprox(mean(draws), mean(d), atol = 0.1)
    @test isapprox(var(draws), var(d) , atol = 0.1)
    nothing
end

function test_draws(d::MatrixDistribution, M::Integer)
    rng = MersenneTwister(123)
    @testset "Testing matrix-variates with $key" for (key, func) in
        Dict("rand(...)" => [rand, rand],
             "rand(rng, ...)" => [dist -> rand(rng, dist), (dist, n) -> rand(rng, dist, n)])
        test_draws(d, func[2](d, M))
    end
    @testset "Testing matrix-variates with $key" for (key, func) in
        Dict("rand!(..., false)" => [(dist, X) -> rand!(dist, X, false), rand!],
             "rand!(rng, ..., false)" => [(dist, X) -> rand!(rng, dist, X, false), (dist, X) -> rand!(rng, dist, X)])
        m = [Matrix{partype(d)}(undef, size(d)) for _ in Base.OneTo(M)]
        x = func[1](d, m)
        @test x ≡ m
        @test isapprox(mean(x) , mean(d) , atol = 0.1)
        m3 = Array{partype(d), 3}(undef, size(d)..., M)
        x = func[2](d, m3)
        @test x ≡ m3
        @test isapprox(mean(x) , mean(mean(d)) , atol = 0.1)
    end
    @testset "Testing matrix-variates with $key" for (key, func) in
        Dict("rand!(..., true)" => (dist, X) -> rand!(dist, X, true),
             "rand!(rng, true)" => (dist, X) -> rand!(rng, dist, X, true))
        m = Vector{Matrix{partype(d)}}(undef, M)
        x = func(d, m)
        @test x ≡ m
        @test isapprox(mean(x), mean(d) , atol = 0.1)
    end
    repeats = 10
    m = Vector{Matrix{partype(d)}}(undef, repeats)
    rand!(d, m)
    @test isassigned(m, 1)
    m1=m[1]
    rand!(d, m)
    @test m1 ≡ m[1]
    rand!(d, m, true)
    @test m1 ≢ m[1]
    m1 = m[1]
    rand!(d, m, false)
    @test m1 ≡ m[1]
    nothing
end

#  --------------------------------------------------
#  Check that the convert and partype methods work
#  --------------------------------------------------

function test_convert(d::MatrixDistribution)
    distname = getproperty(parentmodule(typeof(d)), nameof(typeof(d)))
    @test distname(params(d)...) == d
    @test d == deepcopy(d)
    for elty in (Float32, Float64, BigFloat)
        del1 = convert(distname{elty}, d)
        del2 = convert(distname{elty}, getfield.(Ref(d), fieldnames(typeof(d)))...)
        @test del1 isa distname{elty}
        @test del2 isa distname{elty}
        @test partype(del1) == elty
        @test partype(del2) == elty
    end
    nothing
end

#  --------------------------------------------------
#  Check that cov and var agree
#  --------------------------------------------------

function test_cov(d::MatrixDistribution)
    @test vec(var(d)) ≈ diag(cov(d))
    @test reshape(cov(d), size(d)..., size(d)...) ≈ cov(d, Val(false))
end

test_cov(d::LKJ) = nothing

#  --------------------------------------------------
#  Check dim
#  --------------------------------------------------

function test_dim(d::MatrixDistribution)
    @test dim(d) == size(d, 1)
    @test dim(d) == size(d, 2)
    @test dim(d) == size(mean(d), 1)
    @test dim(d) == size(mean(d), 2)
end

test_dim(d::Union{MatrixNormal, MatrixTDist}) = nothing

#  --------------------------------------------------
#  RUN EVERYTHING
#  --------------------------------------------------

function test_distr(d::MatrixDistribution, M::Integer)
    test_draw(d)
    test_draws(d, M)
    test_cov(d)
    test_convert(d)
    test_dim(d)
    nothing
end

function test_distr(dist::Type{<:MatrixDistribution},
                    n::Integer,
                    p::Integer,
                    M::Integer)
    d = dist(_rand_params(dist, Float64, n, p)...)
    test_distr(d, M)
    nothing
end

#  =============================================================================
#  2. test matrix-variate against the univariate it collapses to in the 1 x 1 case
#  =============================================================================

function test_against_univariate(D::MatrixDistribution, d::UnivariateDistribution)
    X = rand(D)
    x = X[1]
    @test logpdf(D, X) ≈ logpdf(d, x)
    @test pdf(D, X) ≈ pdf(d, x)
    @test mean(D)[1] ≈ mean(d)
    @test var(D)[1] ≈ var(d)
    nothing
end

function test_draws_against_univariate_cdf(D::MatrixDistribution, d::UnivariateDistribution)
    α = 0.025
    M = 100000
    matvardraws = [rand(D)[1] for m in 1:M]
    @test pvalue_kolmogorovsmirnoff(matvardraws, d) >= α
    nothing
end

test_draws_against_univariate_cdf(D::LKJ, d::UnivariateDistribution) = nothing

function test_against_univariate(dist::Type{<:MatrixDistribution})
    D = dist(_rand_params(dist, Float64, 1, 1)...)
    d = _univariate(D)
    test_against_univariate(D, d)
    test_draws_against_univariate_cdf(D, d)
    nothing
end

#  =============================================================================
#  3. test matrix-variate against the multivariate it collapses to in the row/column case
#  =============================================================================

function test_against_multivariate(D::MatrixDistribution, d::MultivariateDistribution)
    X = rand(D)
    x = vec(X)
    @test logpdf(D, X) ≈ logpdf(d, x)
    @test pdf(D, X) ≈ pdf(d, x)
    @test vec(mean(D)) ≈ mean(d)
    @test cov(D) ≈ cov(d)
    nothing
end

function test_against_multivariate(dist::Union{Type{MatrixNormal}, Type{MatrixTDist}}, n::Integer, p::Integer)
    D = dist(_rand_params(dist, Float64, n, 1)...)
    d = _multivariate(D)
    test_against_multivariate(D, d)
    D = dist(_rand_params(dist, Float64, 1, p)...)
    d = _multivariate(D)
    test_against_multivariate(D, d)
    nothing
end

test_against_multivariate(dist::Type{<:MatrixDistribution}, n::Integer, p::Integer) = nothing

#  =============================================================================
#  4. test density evaluation against archived output from Stan
#  =============================================================================

function jsonparams2dist(dist::Type{Wishart}, dict)
    ν = dict["params"][1][1]
    S = zeros(Float64, dict["dims"]...)
    S[:] = Vector{Float64}( dict["params"][2] )
    return Wishart(ν, S)
end

function jsonparams2dist(dist::Type{InverseWishart}, dict)
    ν = dict["params"][1][1]
    S = zeros(Float64, dict["dims"]...)
    S[:] = Vector{Float64}( dict["params"][2] )
    return InverseWishart(ν, S)
end

function jsonparams2dist(dist::Type{LKJ}, dict)
    d = dict["params"][1][1]
    η = dict["params"][2][1]
    return LKJ(d, η)
end

function unpack_matvar_json_dict(dist::Type{<:MatrixDistribution}, dict)
    d = jsonparams2dist(dist, dict)
    X = zeros(Float64, dict["dims"]...)
    X[:] = Vector{Float64}(dict["X"])
    lpdf = dict["lpdf"][1]
    return d, X, lpdf
end

function test_against_stan(dist::Type{<:MatrixDistribution})
    filename = joinpath(@__DIR__, "ref", "matrixvariates", "jsonfiles", "$(dist)_stan_output.json")
    stan_output = JSON.parsefile(filename)
    K = length(stan_output)
    for k in 1:K
        d, X, lpdf = unpack_matvar_json_dict(dist, stan_output[k])
        @test isapprox(logpdf(d, X), lpdf, atol = 1e-10)
        @test isapprox(logpdf(d, [X, X]), [lpdf, lpdf], atol=1e-8)
        @test isapprox(pdf(d, X), exp(lpdf), atol = 1e-6)
        @test isapprox(pdf(d, [X, X]), exp.([lpdf, lpdf]), atol=1e-6)
    end
    nothing
end

function test_against_stan(dist::Union{Type{MatrixNormal}, Type{MatrixTDist}, Type{MatrixBeta}, Type{MatrixFDist}})
    nothing
end

#  =============================================================================
#  5. special, distribution-specific tests
#  =============================================================================

test_special(dist::Type{<:MatrixDistribution}) = nothing

function test_special(dist::Type{MatrixNormal})
    D = MatrixNormal(_rand_params(MatrixNormal, Float64, 2, 2)...)
    @testset "X ~ MN(M, U, V) ⟺ vec(X) ~ N(vec(M), V ⊗ U)" begin
        d = vec(D)
        test_against_multivariate(D, d)
    end
    @testset "MatrixNormal mode" begin
        @test mode(D) == D.M
    end
    @testset "PDMat mixing and matching" begin
        n = 3
        p = 4
        M = randn(n, p)
        u = rand()
        U_scale = ScalMat(n, u)
        U_dense = Matrix(U_scale)
        U_pd    = PDMat(U_dense)
        U_pdiag = PDiagMat(u*ones(n))
        v = rand(p)
        V_pdiag = PDiagMat(v)
        V_dense = Matrix(V_pdiag)
        V_pd    = PDMat(V_dense)
        UV = kron(V_dense, U_dense)
        baseeval = logpdf(MatrixNormal(M, U_dense, V_dense), M)
        for U in [U_scale, U_dense, U_pd, U_pdiag]
            for V in [V_pdiag, V_dense, V_pd]
                d = MatrixNormal(M, U, V)
                @test cov(d) ≈ UV
                @test logpdf(d, M) ≈ baseeval
            end
        end
    end
    nothing
end

function test_special(dist::Type{Wishart})
    n = 3
    M = 5000
    α = 0.05
    ν, Σ = _rand_params(Wishart, Float64, n, n)
    d = Wishart(ν, Σ)
    H = rand(d, M)
    @testset "deprecations" begin
        for warn in (true, false)
            @test @test_deprecated(Wishart(n - 1, Σ, warn)) == Wishart(n - 1, Σ)
        end
    end
    @testset "meanlogdet" begin
        @test isapprox(Distributions.meanlogdet(d), mean(logdet.(H)), atol = 0.1)
    end
    @testset "H ~ W(ν, Σ), a ~ p(a) independent ⟹ a'Ha / a'Σa ~ χ²(ν)" begin
        q = MvTDist(10, randn(n), rand(d))
        ρ = Chisq(ν)
        A = rand(q, M)
        z = [A[:, m]'*H[m]*A[:, m] / (A[:, m]'*Σ*A[:, m]) for m in 1:M]
        @test pvalue_kolmogorovsmirnoff(z, ρ) >= α
    end
    @testset "H ~ W(ν, I) ⟹ H[i, i] ~ χ²(ν)" begin
        κ = n + 1
        ρ = Chisq(κ)
        g = Wishart(κ, ScalMat(n, 1))
        mymats = zeros(n, n, M)
        for m in 1:M
            mymats[:, :, m] = rand(g)
        end
        for i in 1:n
            @test pvalue_kolmogorovsmirnoff(mymats[i, i, :], ρ) >= α / n
        end
    end
    @testset "Check Singular Branch" begin
        # Check that no warnings are shown: #1410
        rank1 = @test_logs Wishart(n - 2, Σ)
        rank2 = @test_logs Wishart(n - 1, Σ)
        test_draw(rank1)
        test_draw(rank2)
        test_draws(rank1, rand(rank1, 10^6))
        test_draws(rank2, rand(rank2, 10^6))
        test_cov(rank1)
        test_cov(rank2)

        X = H[1]
        @test Distributions.singular_wishart_logkernel(d, X) ≈ Distributions.nonsingular_wishart_logkernel(d, X)
        @test Distributions.singular_wishart_logc0(n, ν, d.S, rank(d)) ≈ Distributions.nonsingular_wishart_logc0(n, ν, d.S)
        @test logpdf(d, X) ≈ Distributions.singular_wishart_logkernel(d, X) + Distributions.singular_wishart_logc0(n, ν, d.S, rank(d))
    end
    nothing
end

function test_special(dist::Type{MatrixTDist})
    @testset "MT(v, M, vΣ, Ω) → MN(M, Σ, Ω) as v → ∞" begin
        n, p = (6, 3)
        M, Σ, Ω = _rand_params(MatrixNormal, Float64, n, p)
        MT = MatrixTDist(1000, M, 1000Σ, Ω)
        MN = MatrixNormal(M, Σ, Ω)
        A = rand(MN)
        @test isapprox(logpdf(MT, A), logpdf(MN, A), atol = 0.1)
    end
    @testset "PDMat mixing and matching" begin
        n = 3
        p = 4
        ν = max(n, p) + 1
        M = randn(n, p)
        u = rand()
        U_scale = ScalMat(n, u)
        U_dense = Matrix(U_scale)
        U_pd    = PDMat(U_dense)
        U_pdiag = PDiagMat(u*ones(n))
        v = rand(p)
        V_pdiag = PDiagMat(v)
        V_dense = Matrix(V_pdiag)
        V_pd    = PDMat(V_dense)
        UV = kron(V_dense, U_dense) ./ (ν - 2)
        baseeval = logpdf(MatrixTDist(ν, M, U_dense, V_dense), M)
        for U in [U_scale, U_dense, U_pd, U_pdiag]
            for V in [V_pdiag, V_dense, V_pd]
                d = MatrixTDist(ν, M, U, V)
                @test cov(d) ≈ UV
                @test logpdf(d, M) ≈ baseeval
            end
        end
    end
    nothing
end

function test_special(dist::Type{LKJ})
    @testset "LKJ mode" begin
        @test mode(LKJ(5, 1.5)) == mean(LKJ(5, 1.5))
        @test_throws ArgumentError mode( LKJ(5, 0.5) )
    end
    @testset "LKJ marginals" begin
        d = 4
        η = abs(3randn())
        G = LKJ(d, η)
        M = 10000
        α = 0.05
        L = sum(1:(d - 1))
        ρ = Distributions._marginal(G)
        mymats = zeros(d, d, M)
        for m in 1:M
            mymats[:, :, m] = rand(G)
        end
        for i in 1:d
            for j in 1:i-1
                @test pvalue_kolmogorovsmirnoff(mymats[i, j, :], ρ) >= α / L
            end
        end
    end
    @testset "LKJ integrating constant" begin
        #  =============
        #  odd non-uniform
        #  =============
        d = 5
        η = 2.3
        lkj = LKJ(d, η)
        @test Distributions.lkj_vine_loginvconst(d, η) ≈ Distributions.lkj_onion_loginvconst(d, η)
        @test Distributions.lkj_onion_loginvconst(d, η) ≈ Distributions.lkj_loginvconst_alt(d, η)
        @test lkj.logc0 == -Distributions.lkj_onion_loginvconst(d, η)
        #  =============
        #  odd uniform
        #  =============
        d = 5
        η = 1.0
        lkj = LKJ(d, η)
        @test Distributions.lkj_vine_loginvconst(d, η) ≈ Distributions.lkj_onion_loginvconst(d, η)
        @test Distributions.lkj_onion_loginvconst(d, η) ≈ Distributions.lkj_onion_loginvconst_uniform_odd(d)
        @test Distributions.lkj_vine_loginvconst(d, η) ≈ Distributions.lkj_vine_loginvconst_uniform(d)
        @test Distributions.lkj_onion_loginvconst(d, η) ≈ Distributions.lkj_loginvconst_alt(d, η)
        @test Distributions.lkj_onion_loginvconst(d, η) ≈ Distributions.corr_logvolume(d)
        @test lkj.logc0 == -Distributions.lkj_onion_loginvconst_uniform_odd(d)
        #  =============
        #  even non-uniform
        #  =============
        d = 6
        η = 2.3
        lkj = LKJ(d, η)
        @test Distributions.lkj_vine_loginvconst(d, η) ≈ Distributions.lkj_onion_loginvconst(d, η)
        @test Distributions.lkj_onion_loginvconst(d, η) ≈ Distributions.lkj_loginvconst_alt(d, η)
        @test lkj.logc0 == -Distributions.lkj_onion_loginvconst(d, η)
        #  =============
        #  even uniform
        #  =============
        d = 6
        η = 1.0
        lkj = LKJ(d, η)
        @test Distributions.lkj_vine_loginvconst(d, η) ≈ Distributions.lkj_onion_loginvconst(d, η)
        @test Distributions.lkj_onion_loginvconst(d, η) ≈ Distributions.lkj_onion_loginvconst_uniform_even(d)
        @test Distributions.lkj_vine_loginvconst(d, η) ≈ Distributions.lkj_vine_loginvconst_uniform(d)
        @test Distributions.lkj_onion_loginvconst(d, η) ≈ Distributions.lkj_loginvconst_alt(d, η)
        @test Distributions.lkj_onion_loginvconst(d, η) ≈ Distributions.corr_logvolume(d)
        @test lkj.logc0 == -Distributions.lkj_onion_loginvconst_uniform_even(d)
    end
    @testset "check integrating constant as a volume" begin
        #  d = 2: Lebesgue measure of the set of correlation matrices is 2.
        volume2D = 2
        @test volume2D ≈ exp( Distributions.lkj_onion_loginvconst(2, 1) )
        @test 1 / volume2D ≈ exp( LKJ(2, 1).logc0 )
        #  d = 3: Lebesgue measure of the set of correlation matrices is π²/2.
        #  See here: https://www.jstor.org/stable/2684832
        volume3D = 0.5π^2
        @test volume3D ≈ exp( Distributions.lkj_onion_loginvconst(3, 1) )
        @test 1 / volume3D ≈ exp( LKJ(3, 1).logc0 )
        #  d = 4: Lebesgue measure of the set of correlation matrices is (32/27)π².
        #  See here: https://doi.org/10.4169/amer.math.monthly.123.9.909
        volume4D = (32 / 27)*π^2
        @test volume4D ≈ exp( Distributions.lkj_onion_loginvconst(4, 1) )
        @test 1 / volume4D ≈ exp( LKJ(4, 1).logc0 )
    end
    nothing
end

#  =============================================================================
#  6. main method
#  =============================================================================

function test_matrixvariate(dist::Type{<:MatrixDistribution},
                            n::Integer,
                            p::Integer,
                            M::Integer)
    test_distr(dist, n, p, M)
    test_against_univariate(dist)
    test_against_multivariate(dist, n, p)
    test_against_stan(dist)
    test_special(dist)
    nothing
end

#  =============================================================================
#  7. run unit tests for matrix-variate distributions
#  =============================================================================

matrixvariates = [(MatrixNormal, 2, 4, 10^5),
                  (Wishart, 2, 2, 10^6),
                  (InverseWishart, 2, 2, 10^6),
                  (MatrixTDist, 2, 4, 10^5),
                  (MatrixBeta, 3, 3, 10^5),
                  (MatrixFDist, 3, 3, 10^5),
                  (LKJ, 3, 3, 10^5)]

for distribution in matrixvariates
    dist, n, p, M = distribution
    println("    testing $(dist)")
    @testset "$(dist)" begin
        test_matrixvariate(dist, n, p, M)
    end
end
end


# This contains code that was formerly a part of Julia. License is MIT: http://julialang.org/license

import Base.Order: Forward, Ordering, lt

# Heap operations on flat arrays
# ------------------------------

# Binary heap indexing
heapleft(i::Integer) = 2i
heapright(i::Integer) = 2i + 1
heapparent(i::Integer) = div(i, 2)

# Binary min-heap percolate down.
function percolate_down!(xs::AbstractArray, i::Integer, x=xs[i], o::Ordering=Forward, len::Integer=length(xs))
    @inbounds while (l = heapleft(i)) <= len
        r = heapright(i)
        j = r > len || lt(o, xs[l], xs[r]) ? l : r
        if lt(o, xs[j], x)
            xs[i] = xs[j]
            i = j
        else
            break
        end
    end
    xs[i] = x
end

percolate_down!(xs::AbstractArray, i::Integer, o::Ordering, len::Integer=length(xs)) = percolate_down!(xs, i, xs[i], o, len)

# Binary min-heap percolate up.
function percolate_up!(xs::AbstractArray, i::Integer, x=xs[i], o::Ordering=Forward)
    @inbounds while (j = heapparent(i)) >= 1
        if lt(o, x, xs[j])
            xs[i] = xs[j]
            i = j
        else
            break
        end
    end
    xs[i] = x
end

percolate_up!(xs::AbstractArray{T}, i::Integer, o::Ordering) where {T} = percolate_up!(xs, i, xs[i], o)

"""
    heappop!(v, [ord])

Given a binary heap-ordered array, remove and return the lowest ordered element.
For efficiency, this function does not check that the array is indeed heap-ordered.
"""
function heappop!(xs::AbstractArray, o::Ordering=Forward)
    x = xs[1]
    y = pop!(xs)
    if !isempty(xs)
        percolate_down!(xs, 1, y, o)
    end
    x
end

"""
    heappush!(v, x, [ord])

Given a binary heap-ordered array, push a new element `x`, preserving the heap property.
For efficiency, this function does not check that the array is indeed heap-ordered.
"""
function heappush!(xs::AbstractArray, x, o::Ordering=Forward)
    push!(xs, x)
    percolate_up!(xs, length(xs), x, o)
    xs
end

# Turn an arbitrary array into a binary min-heap in linear time.
"""
    heapify!(v, ord::Ordering=Forward)

In-place [`heapify`](@ref).
"""
function heapify!(xs::AbstractArray, o::Ordering=Forward)
    for i in heapparent(length(xs)):-1:1
        percolate_down!(xs, i, o)
    end
    xs
end

"""
    heapify(v, ord::Ordering=Forward)

Returns a new vector in binary heap order, optionally using the given ordering.
```jldoctest
julia> a = [1,3,4,5,2];

julia> heapify(a)
5-element Array{Int64,1}:
 1
 2
 4
 5
 3

julia> heapify(a, Base.Order.Reverse)
5-element Array{Int64,1}:
 5
 3
 4
 1
 2
```
"""
heapify(xs::AbstractArray, o::Ordering=Forward) = heapify!(copyto!(similar(xs), xs), o)

"""
    isheap(v, ord::Ordering=Forward)

Return `true` if an array is heap-ordered according to the given order.

```jldoctest
julia> a = [1,2,3]
3-element Array{Int64,1}:
 1
 2
 3

julia> isheap(a,Base.Order.Forward)
true

julia> isheap(a,Base.Order.Reverse)
false
```
"""
function isheap(xs::AbstractArray, o::Ordering=Forward)
    for i in 1:div(length(xs), 2)
        if lt(o, xs[heapleft(i)], xs[i]) ||
           (heapright(i) <= length(xs) && lt(o, xs[heapright(i)], xs[i]))
            return false
        end
    end
    true
end


@testset "Det" begin
    Mod.@polyvar x y

    @test det([x 1; y 1]) == x - y
    @test det([x 1; x 1]) == 0
    @test det([x 1 1 1; x y 1 2; 0 0 0 1; x 0 y 0]) == -x*y^2 + 2*x*y - x
end


@inline _F(t, x, P) = B(t, P)*x + β(t, P)
@inline _dHinv(t, K, P) = B(t, P)*K + K*B(t, P)' - a(t, P)
@inline _dK(t, K, P) = B(t, P)*K + K*B(t, P)' + a(t, P)
@inline _dPhi(t, Phi, P) = B(t, P)*Phi

"""
    gpHinv!(K::SamplePath, P, KT=zero(T))

Precompute ``K = H^{-1}`` from ``(d/dt)K = BK + KB' + a`` for a guided proposal.
"""
gpHinv!(K::SamplePath{T}, P, KT=zero(T)) where {T} = _solvebackward!(R3(), _dHinv, K, KT, P)
gpH♢! = gpHinv!

"""
gpV!(K::SamplePath, P, KT=zero(T))

Precompute `V` from ``(d/dt)V = BV + β``, ``V_T = v`` for a guided proposal.
"""
gpV!(V::SamplePath{T}, P, v::T) where {T} = _solvebackward!(R3(), _F, V, v, P)

gpmu(tt, u, P) = solve(R3(), _F, tt, u, P)
gpK(tt, u, P) = solve(R3(), _dK, tt, u, P)

"""
    fundamental_matrix(tt, P)

Compute fundamental solution.
"""
fundamental_matrix(tt, P, Phi0 = one(outertype(P))) = solve(R3(), _dPhi, tt, Phi0, P)


using Symbolics, SparseArrays, LinearAlgebra, BenchmarkTools

nfact = 10
nitem = 5

nobs = nitem*nfact
ntot = nobs+nfact

npar = Int64(nobs + (nfact+1)*nfact/2 + nobs - nfact)

@Symbolics.variables x[1:npar]

x = [x[i] for i in 1:size(x, 1)]
#cov = [cov[i,j] for i in 1:size(cov, 1) for j in 1:size(cov, 2)]
S = zeros(Num, ntot, ntot)
S[diagind(S)[1:nobs]] .= x[1:nobs]
xind = nobs
for i in 1:nfact
    for j in i:nfact
        xind = xind+1
        ind = nobs+i
        jnd = nobs+j
        S[ind, jnd] = x[xind]
        S[jnd, ind] = x[xind]
    end
end

F = zeros(Num, nobs, ntot)

F[diagind(F)] .= 1

A = zeros(Num, ntot, ntot)

for i in 1:nfact
    for j in 2:nitem
        xind = xind+1
        jnd = i+nobs
        ind = j+(i-1)*nitem
        A[ind, jnd] = x[xind]
    end
end

for i in 1:nfact
    A[1+(i-1)*nitem, nobs+i] = 1
end

S = sparse(S)

#F
F = sparse(F)

#A
A = sparse(A)

par = [x[i] for i = 1:npar]
invia = I + A + A^2 + A^3 + A^4
imp_cov_sym = F*invia*S*permutedims(invia)*permutedims(F)
imp_cov_sym = Array(imp_cov_sym)
imp_cov_sym = Symbolics.simplify.(imp_cov_sym)

imp_fun = eval(Symbolics.build_function(imp_cov_sym, par)[2])

Σ = zeros(nobs, nobs)
randpar = rand(npar)

imp_fun(Σ, randpar)

Σ

using BenchmarkTools

@benchmark imp_fun(Σ, randpar)

########## Benchmark Gradients
using sem, Optim

obs_cov = rand(nobs, nobs); obs_cov = obs_cov*obs_cov'
semobserved = SemObsCommon(data = obs_cov)
start_val_ml = randpar
loss_ml = Loss([SemML(semobserved, [0.0], similar(start_val_ml))])

# imply
imply_ml = ImplySymbolic(A, S, F, x, startval)

grad_ml = sem.∇SemML_2()

diff_ana_ml =
    SemAnalyticDiff(
        LBFGS(),# P = one(rand(31, 31))),
        Optim.Options(
            ;f_tol = 1e-10,
            x_tol = 1.5e-8),
            (grad_ml,))

model_ml_ana = Sem(semobserved, imply_ml, loss_ml, diff_ana_ml)

solution_ana_ml = sem_fit(model_ml_ana)

startval = ones(size(randpar))

model_ml_ana(startval)

gradient = similar(randpar)
model_ml_ana(startval, gradient)

using FiniteDiff
gradient ≈ FiniteDiff.finite_difference_gradient(model_ml_ana, startval)

isposdef(Symmetric(model_ml_ana.imply.imp_cov))

@benchmark model_ml_ana(startval, gradient)

@benchmark model_ml_ana(startval)


# Cross-check against data published in Frost 2017 PRB.

@testset "FrostPolaronMobility2017" begin

T=300

#Ts,Kμs, Hμs, FHIPμs, ks, Ms, As, Bs, Cs, Fs, Taus
#effectivemass=0.12 # the bare-electron band effective-mass.
# --> 0.12 for electrons and 0.15 for holes, in MAPI. See 2014 PRB.
# MAPI  4.5, 24.1, 2.25THz - 75 cm^-1 ; α=
MAPIe=polaronmobility(T, 4.5, 24.1, 2.25, 0.12)
MAPIh=polaronmobility(T, 4.5, 24.1, 2.25, 0.15)

# Hellwarth mobility
@test MAPIe.Hμ[1] ≈ 136.42 rtol=0.02
# Test variational parameters
@test MAPIe.v[1] ≈ 19.86 rtol=0.02
@test MAPIe.w[1] ≈ 16.96 rtol=0.02

# Same for the MAPI holes @ 300 K
@test MAPIh.Hμ[1] ≈ 94.15 rtol=0.02
@test MAPIh.v[1] ≈ 20.09  rtol=0.02
@test MAPIh.w[1] ≈ 16.81  rtol=0.02

end


using MOCNeutronTransport
using StaticArrays
@testset "Point_2D" begin
    @testset "$F" for F in [Float32, Float64, BigFloat]
        @testset "Operators" begin
            p₁ = Point_2D{F}(1, 2)
            p₂ = Point_2D{F}(2, 4)

            # Point_2D isapprox
            p = Point_2D{F}(1, 2 - 10*eps(F))
            @test F(2) ≈ 2 - 10*eps(F)
            @test p ≈ p₁

            # Point_2D addition
            p = p₁ + p₂
            @test p == [3, 6]

            # Point_2D subtraction
            p = p₁ - p₂
            @test p == [-1, -2]

            # Cross product
            p₁ = Point_2D{F}(2, 3)
            p₂ = Point_2D{F}(5, 6)
            v = p₁ × p₂
            @test v ≈ -3

            # Dot product
            @test p₁ ⋅ p₂ ≈ 10 + 18

            # Number addition
            p₁ = Point_2D{F}(1, 2)
            p₂ = Point_2D{F}(2, 4)
            p = p₁ .+ 1
            @test p  ≈ [2, 3]

            # Subtraction
            p = p₁ .- 1
            @test p == [0, 1]

            # Multiplication
            p = 4*p₁
            @test p == [4, 8]

            # Division
            p = p₁/4
            @test p == [1//4, 1//2]

            # Unary -
            p = -p₁
            @test p == [-1, -2]

            # Matrix multiplcation
            A = SMatrix{2, 2, F, 4}(1, 2, 3, 4)
            p = Point_2D{F}(1, 2)
            q = A*p
            @test q[1] ≈ 7
            @test q[2] ≈ 10
            @test typeof(q) == Point_2D{F}
        end

        @testset "Methods" begin
            p₁ = Point_2D{F}(1, 2)
            p₂ = Point_2D{F}(2, 4)

            # distance
            @test distance(p₁, p₂) ≈ sqrt(5)

            # distance²
            @test distance²(p₁, p₂) ≈ 5

            # norm
            @test norm(p₁) ≈ sqrt(5)
            @test norm(p₂) ≈ sqrt(4 + 16)

            # norm²
            @test norm²(p₁) ≈ 5
            @test norm²(p₂) ≈ 4 + 16

            # midpoint
            mp = midpoint(p₁, p₂)
            @test mp[1] ≈ 3//2
            @test mp[2] ≈ 3

            # sort_points
            p₁ = Point_2D{F}(1, 0)
            p₂ = Point_2D{F}(2, 0)
            p₃ = Point_2D{F}(3, 0)
            points = [p₃, p₁, p₂]
            points_sorted = sortpoints(Point_2D{F}(0, 0), points)
            @test points_sorted[1] == p₁
            @test points_sorted[2] == p₂
            @test points_sorted[3] == p₃
            # mutating
            sortpoints!(Point_2D{F}(0, 0), points)
            @test points[1] == p₁
            @test points[2] == p₂
            @test points[3] == p₃
        end
    end
end


using Base.Test
using QuantumOptics

@testset "schroedinger" begin

N = 3
Ncutoff = 2
T = [0.:0.1:1.;]

Ω = [0 2 3;
     2 0 1;
     3 1 0]

ω = [1., 1.2, 1.5]

basis_fock = FockBasis(Ncutoff)
basis = tensor([basis_fock for i=1:N]...)

a = destroy(basis_fock)
at = create(basis_fock)
I = identityoperator(basis_fock)

psi0 = tensor([coherentstate(basis_fock, i%Ncutoff) for i=1:N]...)

# Interaction picture
Hrot = SparseOperator[]
for i=1:N, j=1:N
    if i==j
        continue
    end
    h = embed(basis, [i,j], [a, Ω[i,j]*at])
    push!(Hrot, h)
end
Hrot = sum(Hrot)

# Schroedinger picture
function f(t, psi)
    H = SparseOperator[embed(basis, i, ω[i]*at*a) for i=1:N]
    for i=1:N, j=1:N
        if i==j
            continue
        end
        h = embed(basis, [i,j], [a, exp(1im*(ω[i]-ω[j])*t)*Ω[i,j]*at])
        push!(H, h)
    end
    sum(H)
end

tout, psi_rot_t = timeevolution.schroedinger(T, psi0, Hrot)
tout, psi_t = timeevolution.schroedinger_dynamic(T, psi0, f)

for (i, t) in enumerate(tout)
    R = prod([embed(basis, i, expm(1im*ω[i]*t*full(at*a))) for i=1:N])
    psi_rot = psi_rot_t[i]
    psi = psi_t[i]
    # @test abs(dagger(psi_rot)*R*psi) < 1e-5
    rho = dm(psi)
    rho_rot = dm(psi_rot)
    @test tracedistance(rho_rot, full(R)*rho*dagger(full(R))) < 1e-5
end

function fout(t, psi)
 deepcopy(psi)
end
t_fout, psi_fout = timeevolution.schroedinger(T, psi0, Hrot; fout=fout)
@test t_fout == tout && psi_fout == psi_rot_t

end # testset


using EffectiveWaves
using LinearAlgebra
using Plots
pyplot()

medium = Medium(ρ=1.0, c=1.0)
species = [Specie(ρ=0.1,r=0.4, c=0.2, volfrac=0.3)]

ωs = [1.0]
ω = ωs[1]
k = ω/medium.c
hankel_order = ho = 2

kps = wavenumbers(ω, medium, species; symmetry = :plane, dim = 3,
    hankel_order = hankel_order, num_wavenumbers=17,
    tol = 1e-8,
    mesh_points = 14, mesh_size = 2.0)

# Doesnt really work
# R_kps = wavenumbers(ω, medium, species; dim = 3, hankel_order = hankel_order,
#     num_wavenumbers=27, tol = 1e-8, mesh_points = 10, mesh_size = 2.0)
R_kps = Complex{Float64}[2.30521+1.85168im, -1.46227+2.03392im, -1.45281+2.03493im, 10.9431+5.60625im, -10.0589+5.61005im, -10.0562+5.61218im, 19.1007+6.88548im, -18.2139+6.88838im, 27.0945+7.7147im, 35.0209+8.32989im, 35.0222+8.33306im, 42.9157+8.82717im, 50.7889+9.23896im, 58.6476+9.58905im, 58.649+9.59209im, 66.5001+9.90123im, 74.3447+10.1762im, 82.1845+10.4237im, 90.0205+10.6489im]

wavesystem2D = wavematrix2D(ωs[1], medium, species; hankel_order=hankel_order);
det2D(kp) = det(wavesystem2D(kp))
kps2 = wavenumbers(ω, medium, species; dim = 2, hankel_order = hankel_order,
    num_wavenumbers=27, tol = 1e-8, mesh_points = 10, mesh_size = 2.0);

wavesystemPlane = wavematrix3DPlane(ωs[1], medium, species; hankel_order=hankel_order)
detP(kp) = det(wavesystemPlane(kp))
detP.(kps)

inds = findall( abs.(detP.(kps)) .< 1e-5);
kps = kps[inds]

wavesystem3D = wavematrix3D(ωs[1], medium, species; hankel_order=hankel_order);
detR(kp) = (det(wavesystem3D(kp)))^(1/2)
detR.(kps)

MRsvd = svd(wavesystem3D(kps[2]))
MRsvd = svd(wavesystem3D(kps[1]))
MRsvd.S
wave_vec =  MRsvd.V[:,end]
wave_mat = reshape(wave_vec,((ho+1)^2,(ho+1)^2))


function run(s)
    lines = split(s, "\n")
    constraints = [(idx, parse(Int, el)) for (idx, el) in enumerate(split(lines[2], ",")) if el != "x"]
    ts::Int = 0
    jump::Int = 1
    current_idx::Int = 1
    max_idx = length(constraints)
    while true
        offset, mul = constraints[current_idx]
        if (ts + (offset - 1)) % mul == 0
            current_idx += 1
            jump *= mul
        end
        if current_idx > max_idx
            return ts
        end
        ts += jump
    end
    return 0
end

#########################################

function main()
    res, time, memory = @timed run(ARGS[1])
    println("_duration:$(time * 1000)")
    println(res)
end

main()


### from RosettaCode
### http://rosettacode.org/wiki/Factors_of_an_integer#Julia
function factors(n)
    f = [one(n)]
    for (p,e) in factor(n)
        f = reduce(vcat, f, [f*p^j for j in 1:e])
    end
    return length(f) == 1 ? [one(n), n] : sort!(f)
end

for x in 2:24
	println(x, " ", 10*sum(factors(x)))
end

count = 100000
x = 0
while x < 34000000
	count += 1
	x = 10*sum(factors(count))
	if count % 10000 == 0
		println(count, " ", x)
	end
end
println(count, " ", x)

### second way
### allocate M houses and add a loop over elves adding presents to each house

# allocate the houses
houses = 11*ones(Int64, 3400000)
for elf in 2:length(houses)
	houses[elf:elf:min(50*elf, 3400000)] += 11*elf
end

for (idx, h) in enumerate(houses)
	if h > 34000000
		println(idx)
		break
	end
end


using CUDA
using MPI
using Test
using LinearAlgebra
using Random
using StaticArrays
using KernelAbstractions: CPU, CUDADevice
using ClimateMachine
using ClimateMachine.SystemSolvers
using ClimateMachine.MPIStateArrays
using Random

Random.seed!(1235)

let
    ClimateMachine.init()
    mpicomm = MPI.COMM_WORLD
    ArrayType = ClimateMachine.array_type()
    n = 100
    T = Float64
    A = rand(n, n)
    scale = 2.0
    ϵ = 0.1
    # Matrix 1
    A = A' * A .* ϵ + scale * I

    err_thresh = sqrt(eps(T))

    # Matrix 2
    # A = Diagonal(collect(1:n) * 1.0)
    positive_definite = minimum(eigvals(A)) > eps(1.0)
    @test positive_definite

    b = ones(n) * 1.0
    mulbyA!(y, x) = (y .= A * x)

    tol = sqrt(eps(T))
    method(b, tol) = ConjugateGradient(b, max_iter = n)
    linearsolver = method(b, tol)

    x = ones(n) * 1.0
    x0 = copy(x)
    iters = linearsolve!(mulbyA!, nothing, linearsolver, x, b; max_iters = Inf)
    exact = A \ b
    x0 = copy(x)

    @testset "Array test" begin
        @test norm(x - exact) / norm(exact) < err_thresh
        @test norm(A * x - b) / norm(b) < err_thresh
    end

    # Testing for CUDA CuArrays
    if CUDA.has_cuda_gpu()
        at_A = ArrayType(A)
        at_b = ArrayType(b)
        at_x = ArrayType(ones(n) * 1.0)
        mulbyat_A!(y, x) = (y .= at_A * x)
        at_method(b, tol) = ConjugateGradient(b, max_iter = n)
        linearsolver = at_method(at_b, tol)
        iters = linearsolve!(
            mulbyat_A!,
            nothing,
            linearsolver,
            at_x,
            at_b;
            max_iters = n,
        )
        exact = at_A \ at_b

        @testset "CuArray test" begin
            @test norm(at_x - exact) / norm(exact) < err_thresh
            @test norm(at_A * at_x - at_b) / norm(at_b) < err_thresh
        end
    end

    mpi_b = MPIStateArray{T}(mpicomm, ArrayType, 4, 4, 4)
    mpi_x = MPIStateArray{T}(mpicomm, ArrayType, 4, 4, 4)
    mpi_A = ArrayType(randn(4^3, 4^3))
    mpi_A .= mpi_A' * mpi_A

    function mpi_mulby!(x, y)
        fy = y.data[:]
        fx = mpi_A * fy
        x.data[:] .= fx[:]
        return nothing
    end

    mpi_b.data[:] .= ArrayType(randn(4^3))
    mpi_x.data[:] .= ArrayType(randn(4^3))

    mpi_method(mpi_b, tol) = ConjugateGradient(mpi_b, max_iter = n)
    linearsolver = mpi_method(mpi_b, tol)
    iters = linearsolve!(
        mpi_mulby!,
        nothing,
        linearsolver,
        mpi_x,
        mpi_b;
        max_iters = n,
    )

    exact = mpi_A \ mpi_b[:]

    @testset "MPIStateArray test" begin
        @test norm(mpi_x.data[:] - exact) / norm(exact) < err_thresh
        mpi_Ax = MPIStateArray{T}(mpicomm, ArrayType, 4, 4, 4)
        mpi_mulby!(mpi_Ax, mpi_x)
        @test norm(mpi_Ax - mpi_b) / norm(mpi_b) < err_thresh
    end

    # ## More Complex Example
    function closure_linear_operator!(A, tup)
        function linear_operator!(y, x)
            alias_x = reshape(x, tup)
            alias_y = reshape(y, tup)
            for i6 in 1:tup[6]
                for i4 in 1:tup[4]
                    for i2 in 1:tup[2]
                        for i1 in 1:tup[1]
                            tmp = alias_x[i1, i2, :, i4, :, i6][:]
                            tmp2 = A[i1, i2, i4, i6] * tmp
                            alias_y[i1, i2, :, i4, :, i6] .=
                                reshape(tmp2, (tup[3], tup[5]))
                        end
                    end
                end
            end
        end
    end

    tup = (3, 4, 7, 2, 20, 2)

    B = [
        randn(tup[3] * tup[5], tup[3] * tup[5])
        for i1 in 1:tup[1], i2 in 1:tup[2], i4 in 1:tup[4], i6 in 1:tup[6]
    ]
    columnwise_A = [
        B[i1, i2, i4, i6] * B[i1, i2, i4, i6]' + 10I
        for i1 in 1:tup[1], i2 in 1:tup[2], i4 in 1:tup[4], i6 in 1:tup[6]
    ]
    columnwise_inv_A = [
        inv(columnwise_A[i1, i2, i4, i6])
        for i1 in 1:tup[1], i2 in 1:tup[2], i4 in 1:tup[4], i6 in 1:tup[6]
    ]
    columnwise_linear_operator! = closure_linear_operator!(columnwise_A, tup)
    columnwise_inverse_linear_operator! =
        closure_linear_operator!(columnwise_inv_A, tup)

    mpi_tup = (tup[1] * tup[2] * tup[3], tup[4], tup[5] * tup[6])
    b = randn(mpi_tup)
    x = randn(mpi_tup)

    linearsolver = ConjugateGradient(
        x,
        max_iter = tup[3] * tup[5],
        dims = (3, 5),
        reshape_tuple = tup,
    )

    iters =
        linearsolve!(columnwise_linear_operator!, nothing, linearsolver, x, b)
    x_exact = copy(x)
    columnwise_inverse_linear_operator!(x_exact, b)

    @testset "Columnwise test" begin
        @test norm(x - x_exact) / norm(x_exact) < err_thresh
    end
end

nothing


function closer(a::Aircraft, b::Aircraft, p::Aircraft)
	if (a.dyn.x - p.dyn.x)^2 + (a.dyn.y - p.dyn.y)^2 <
		(b.dyn.x - p.dyn.x)^2 + (b.dyn.y - p.dyn.y)^2
		return true
	else
		return false
	end
end

function farther(a::Aircraft, b::Aircraft, p::Aircraft)
	if (a.dyn.x - p.dyn.x)^2 + (a.dyn.y - p.dyn.y)^2 >
		(b.dyn.x - p.dyn.x)^2 + (b.dyn.y - p.dyn.y)^2
		return true
	else
		return false
	end
end

"""
	sort the Aircraft in an array
	using quick sort, returns ascending order
"""
function sort_agents!(p::Aircraft, agents_arr::Vector{Aircraft}, lo::Int64, hi::Int64)
    i, j = lo, hi
    while i < hi
        pivot = agents_arr[(lo + hi) >>> 1]
        while i <= j
            while closer(agents_arr[i], pivot, p); i += 1; end
            while farther(agents_arr[j], pivot, p); j -= 1; end
            if i <= j
                agents_arr[i], agents_arr[j] = agents_arr[j], agents_arr[i]
                i, j = i + 1, j - 1
            end
        end
        if lo < j; sort_agents!(p, agents_arr, lo, j); end
        lo, j = i, hi
    end
end

function get_section_id(ac::Aircraft, intruder::Aircraft, num_sections::Int64)
	ac_pos_angle = norm_angle(norm_angle(atan2(intruder.dyn.y - ac.dyn.y, intruder.dyn.x - ac.dyn.x)) - norm_angle(ac.dyn.heading))
	# ac_pos_angle = norm_angle(atan2(intruder.dyn.y - ac.dyn.y, intruder.dyn.x - ac.dyn.x))
	section = Int(floor(ac_pos_angle / (2 * pi / num_sections)) + 1)
	return section
end

function get_state(ac, intruder)
	intruder_pos_angle = atan2(intruder.dyn.y - ac.dyn.y, intruder.dyn.x - ac.dyn.x)

	return [ft2m(sqrt((ac.dyn.x - intruder.dyn.x)^2 + (ac.dyn.y - intruder.dyn.y)^2)),
			norm_angle(intruder_pos_angle - ac.dyn.heading),
			norm_angle(intruder.dyn.heading - ac.dyn.heading),
			ft2m(ac.dyn.v),
			ft2m(intruder.dyn.v)]
end

function get_full_and_low_fi_state(correction_type::Symbol, ac::Aircraft, All_AC::Vector{Aircraft}, num_sections::Int64)

	intruders = get_observation!(ac, All_AC, ac.sensor) # all intruders are within the sensing range

	if correction_type == :closest || policy
		# Method using num_sections closest intruders
		full_state = zeros(num_sections * STATE_DIM + AUG_STATE_DIM)
		hi = length(intruders)
		lo = (hi == 0) ? 0 : 1
		if hi > 0
			# # 1. Using multiple close intruders
			# # Ascending order in terms of distance from ego_agent
			sort_agents!(ac, intruders, lo, hi)
			# # 2. Using multiple dangerous intruders
			# # Ascending order in terms of maximum q_vals wrt ego_agent
			# # The minimum maximum q_vals => the most dangerous
			# sort_agents_q_vals!(env, intruders, lo, hi)

			for i in 1 : num_sections
				if i <= length(intruders)
					full_state[STATE_DIM * (i - 1) + 1 : STATE_DIM * i] .= get_state(ac, intruders[i])
				else
					full_state[STATE_DIM * (i - 1) + 1 : STATE_DIM * i] .= terminal_state
				end
			end
		else # no intruder
			full_state[1 : STATE_DIM * num_sections] .= TERM_VAR
		end

	elseif correction_type == :sector
		# # Method using num_sections sections
		section_ac_arr = Dict{Int64, Vector{Aircraft}}()
		for isection in 1:num_sections
			section_ac_arr[isection] = Aircraft[]
		end

		for int in intruders
			if int.id != ac.id
				section_id = get_section_id(ac, int, num_sections)
				push!(section_ac_arr[section_id], int)
			end
		end

		full_state = zeros(num_sections * STATE_DIM + AUG_STATE_DIM)

		# sort ac in each section
		for isection in 1:num_sections
			hi = length(section_ac_arr[isection])
			lo = (hi == 0) ? 0 : 1
			if hi > 0
				sort_agents!(ac, section_ac_arr[isection], lo, hi)
				# resulting section_ac_arr is in ascending order,
				# grab ac that is closest to the ego ac for each section
				# to update the full_state
				full_state[STATE_DIM * (isection - 1) + 1 : STATE_DIM * isection] .=
					get_state(ac, section_ac_arr[isection][1]) # [m]
			else
				# no intruder in the section
				# set the corresponding state var to terminal values
				full_state[STATE_DIM * (isection - 1) + 1 : STATE_DIM * isection] .=
					[TERM_VAR, TERM_VAR, TERM_VAR, TERM_VAR, TERM_VAR]
			end
		end
		#
	end # if correction_type

	# append the destination information to the full_state
	# 1. deviation from the destination, measured by rad
	deviation = norm_angle(ac.dyn.heading -
		norm_angle(atan2(ac.y_dest - ac.dyn.y, ac.x_dest - ac.dyn.x)))
	if deviation > π
		deviation = deviation - 2 * π # ∈ [-π, π]
		# deviation = 2 * π - deviation # ∈ [0, π], deprecated
	end

	# 2. current distance to destination
	dist_to_dest = sqrt((ac.dyn.y - ac.y_dest)^2 + (ac.dyn.x - ac.x_dest)^2) # [ft]
	# 3. previous distance to destination
	prev_dist_to_dest = ac.dist_to_dest # [ft]
	## update ego_agent_dist_to_dest
	ac.dist_to_dest = dist_to_dest # [ft]

	# update state entries
	full_state[end - AUG_STATE_DIM + 1 : end] .= [deviation, ft2m(dist_to_dest), ft2m(prev_dist_to_dest)] # [m]

	# find the lowfi_state from the full_state
	closest_intruder_section_id = indmin(full_state[[(i - 1) * STATE_DIM + IND_rho for i in 1:num_sections]])
	lowfi_state = full_state[STATE_DIM * (closest_intruder_section_id - 1) + 1 : STATE_DIM * closest_intruder_section_id]

	return full_state, lowfi_state
end

function normalize_obs!(ac::Aircraft, num_sections::Int64, obs::Vector{Float64})
	for isection in 0 : num_sections - 1
		if obs[isection * STATE_DIM + IND_rho] == TERM_VAR
			obs[isection * STATE_DIM + 1 : isection * STATE_DIM + STATE_DIM] *= 0 # if terminal, set empty
			# obs[isection * STATE_DIM + 1 : isection * STATE_DIM + STATE_DIM] = 1 # if terminal, set to 1
		else
			obs[isection * STATE_DIM + IND_rho] /= SENSING_RANGE
			obs[isection * STATE_DIM + IND_theta] /= (2 * pi)
			obs[isection * STATE_DIM + IND_phi] /= (2 * pi)
			# obs[isection * STATE_DIM + IND_v_own] = (obs[isection * STATE_DIM + IND_v_own] - V_MIN) / (V_MAX - V_MIN + 1e-12)
			obs[isection * STATE_DIM + IND_v_own] /= V_MAX
			obs[isection * STATE_DIM + IND_v_int] /= V_MAX
		end
	end
	obs[IND_deviation] /= pi
	obs[IND_dist_to_dest] /= ft2m(ac.nominal_route_len)
	obs[IND_prev_dist_to_dest] /= ft2m(ac.nominal_route_len)
end

"""
	update the observation of an Aircraft object with Nearest_Intruder_Tracker
"""
function get_observation!(ac::Aircraft, All_AC::Array{Aircraft, 1}, ::Nearest_Intruder_Tracker)
	# search for the closest intruder from ownship
	intruders = Aircraft[]
	closest_intruder = nothing
	closest_dist = Inf
	num_intruders_at_this_step = 0
	for other_ac in All_AC
		if other_ac.id != ac.id
			dist = sqrt((ac.dyn.x - other_ac.dyn.x)^2 + (ac.dyn.y - other_ac.dyn.y)^2)
			if dist <= ac.sensor.sensing_range
				push!(intruders, other_ac)
				num_intruders_at_this_step += 1
				ac.cumulative_num_ac_encountered += 1
				if dist < closest_dist
					closest_dist = dist
					closest_intruder = other_ac
				end
			end
		end
	end

	push!(ac.num_ac_encountered_arr, num_intruders_at_this_step)

	if closest_intruder == nothing
		ac.sensor.obs = []
	else
		ac.sensor.obs = [construct_obs(closest_dist, ac, closest_intruder)]
	end

	return intruders
end

"""
	Update the observation of an Aircraft object with Multi_Threat_Tracker
"""
function get_observation!(ac::Aircraft, All_AC::Array{Aircraft, 1}, ::Multi_Threat_Tracker)
	ac.sensor.obs = [] # dump previous obs
	intruders = Aircraft[]
	for other_ac in All_AC
		if other_ac.id != ac.id
			dist = sqrt((ac.dyn.x - other_ac.dyn.x)^2 + (ac.dyn.y - other_ac.dyn.y)^2)
		    if dist <= ac.sensor.sensing_range
				obs = construct_obs(dist, ac, other_ac)
		    	push!(ac.sensor.obs, obs)
		    	push!(intruders, other_ac)
		    end
		end
	end
	ac.cumulative_num_ac_encountered += length(intruders)
	push!(ac.num_ac_encountered_arr, length(intruders))
	return intruders
end

"""
	Helper: construct the Observation object
"""
function construct_obs(dist::Float64, ac::Aircraft, other_ac::Aircraft)
	rho = dist
	intruder_pos_angle = atan2(other_ac.dyn.y - ac.dyn.y, other_ac.dyn.x - ac.dyn.x)
	theta = intruder_pos_angle - ac.dyn.heading
	psi = other_ac.dyn.heading - ac.dyn.heading
	# constrain theta and psi
	if theta > pi
		theta = theta - 2 * pi
	elseif theta < - pi
		theta = theta + 2 * pi
	end
	if psi > pi
		psi = psi - 2 * pi
	elseif psi < - pi
		psi = psi + 2 * pi
	end

	sos = ac.dyn.v
	soi = other_ac.dyn.v
	tau = 0 # coaltitude
	pRA = ac.cas.Advisory_to_Ind_dict[ac.advisory]

	obs = Observation(rho, theta, psi, sos, soi, tau, pRA, intruder_pos_angle)
	return obs
end


include("../examples/mpec/bard1.jl")
include("../examples/mpec/dempe.jl")
include("../examples/mpec/gnash1m.jl")

include("../examples/mpec/bard1_test.jl")


export ARC_Cub
function ARC_Cub(h :: AbstractNLPModel;
                 t₀ :: Float64 = h.meta.x0[1],
                 tₘ :: Float64 = 100.0,
                 tol :: Float64 = 1e-7,
                 maxiter :: Int = 50,
                 verbose :: Bool = false)

    (length(h.meta.x0) > 1) && warn("Not a 1-D problem ")
    # ARC parameters
    eps1 = 0.2; eps2 = 0.8
    red = 0.5; aug = 20
    Δ = 1.0
    t = t₀

    iter = 0;
    fₖ = obj(h,[t])[1]
    gₖ = grad(h, [t])[1]

    secₖ = 1.0
    dN=-gₖ #pour la première iter
    A=0.5
    B=0.0
    q(d) = fₖ + gₖ * d + A * d^2 + B * d^3

    verbose && @printf(" iter  t         gₖ          Δ        pred         ared\n")
    verbose && @printf(" %4d %7.2e  %7.2e  %7.2e \n", iter, t, gₖ, Δ)

    while ((abs(gₖ[1])>tol) & (iter < maxiter)) || (iter == 0)

      #Step computation
      Quad(t) = fₖ + gₖ * t + A * t^2 + B * t^3 + (1 / (4 * Δ)) * t^4
      dQuad(t) = gₖ + 2 * A * t + 3 * B * t^2 + (1 / Δ) * t^3

      #dR=PolynomialRoots.roots([gₖ,2*A,3*B,(1/Δ)])
      #TECHNIQUEMENT LE CALCUL DES RACINES PAR PolynomialRoots À LA MÊME PRÉCISION QUE SCILAB, MAIS IL DONNE UN SEG FAULT (comme pour TR_Cub)
      p = Poly([gₖ, 2 * A, 3 * B, (1 / Δ)])
      dR = roots(p)
      vmin=Inf
      dN=0
      for i=1:length(dR)
        rr = dR[i]
        if isreal(rr)
          rr = real(rr)
          vact = Quad(rr)
          if rr * gₖ[1] < 0.0
            if vact < vmin
              dN = rr
              vmin = vact
            end
          end
        end
      end

      d = dN

      if d == 0.0
        return (t, fₖ, norm(gₖ, Inf), iter + 1, false, false, :FailedDirectionComputation, h.counters.neval_obj, h.counters.neval_grad, h.counters.neval_hess)
        #L'OUTIL DE CALCUL DES RACINES DU PACKAGE Polynomials ARRONDIE À 0 CERTAINES RACINES D'OÙ LA NOUVELLE CONDITION
      end

      ftestTR=obj(h, [t + d])[1]
      gtestTR=grad(h, [t + d])[1]

      pred = gₖ[1] * d + A *d ^2 + B * d^3
      if pred > -1e-10
        ared=(gₖ[1] + gtestTR[1]) * d / 2
      else
        ared = ftestTR - fₖ
      end

      ratio = ared / pred
      if ratio < eps1
        Δ = red * Δ
      else
        #two point memory
        tpred = t
        gkm1 = gₖ
        fkm1 = fₖ

        t = t + d
        gₖ = gtestTR
        fₖ = ftestTR

        #Cubic two points interpolation
        s = t - tpred
        y = gₖ - gkm1
        α = -s
        z = gₖ + gkm1 + 3 *( fₖ - fkm1) / α
        discr= z^2 - gₖ * gkm1
        denom = gₖ + gkm1 + 2 * z
        B = (1 / 3) * (gₖ + gkm1 + 2 * z) / ( α * α)
        A = -(gₖ + z) / α
        if ratio > eps2
          Δ = aug * Δ
        end
      end
      iter += 1
      verbose && @printf(" %4d %7.2e  %7.2e  %7.2e %7.2e %7.2e\n", iter, t, gₖ, Δ, pred, ared)
    end
    status = :NotSolved
    (abs(gₖ) < tol) && (status = :Optimal)
    (iter >= maxiter) && (status = :Tired)
    tired = iter > maxiter
    optimal = abs(gₖ) < tol
    return (t, fₖ, norm(gₖ, Inf), iter, optimal, tired, status, h.counters.neval_obj, h.counters.neval_grad, h.counters.neval_hess)

end


using ITensors, Test

@testset "README Examples" begin
  @testset "ITensor Basics" begin
    i = Index(3)
    j = Index(5)
    k = Index(2)
    l = Index(7)

    A = ITensor(i, j, k)
    B = ITensor(j, l)

    A[i(1), j(1), k(1)] = 11.1
    A[i(2), j(1), k(2)] = -21.2
    A[k(1), i(3), j(1)] = 31.1  # can provide Index values in any order
    # ...

    # Contract over shared index j
    C = A * B

    @test hasinds(C, i, k, l) == true

    D = randomITensor(k, j, i) # ITensor with random elements

    # Add two ITensors
    # must have same set of indices
    # but can be in any order
    R = A + D
  end

  @testset "SVD of a Matrix" begin
    i = Index(10)
    j = Index(20)
    M = randomITensor(i, j)
    U, S, V = svd(M, i)
    @test norm(M - U * S * V) < 1E-12
  end

  @testset "SVD of a Tensor" begin
    i = Index(4, "i")
    j = Index(4, "j")
    k = Index(4, "k")
    l = Index(4, "l")
    T = randomITensor(i, j, k, l)
    U, S, V = svd(T, i, k)
    @test hasinds(U, i, k)
    @test hasinds(V, j, l)
    @test norm(T - U * S * V) < 1E-12
  end

  @testset "Making Tensor Indices" begin
    i = Index(3)     # Index of dimension 3
    @test dim(i) == 3     # dim(i) = 3

    ci = copy(i)
    @test ci == i    # true

    j = Index(5, "j") # Index with a tag "j"

    @test j != i     # false

    s = Index(2, "n=1,Site") # Index with two tags,
    # "Site" and "n=1"
    @test hastags(s, "Site") # hastags(s,"Site") = true
    @test hastags(s, "n=1")  # hastags(s,"n=1") = true

    i1 = prime(i) # i1 has a "prime level" of 1
    # but otherwise same properties as i
    @test i1 != i # false, prime levels do not match
  end

  @testset "DMRG" begin
    N = 100
    sites = siteinds("S=1", N)

    # Input operator terms which define
    # a Hamiltonian matrix, and convert
    # these terms to an MPO tensor network
    ampo = AutoMPO()
    for j in 1:(N - 1)
      add!(ampo, "Sz", j, "Sz", j + 1)
      add!(ampo, 0.5, "S+", j, "S-", j + 1)
      add!(ampo, 0.5, "S-", j, "S+", j + 1)
    end
    H = MPO(ampo, sites)

    # Create an initial random matrix product state
    psi0 = randomMPS(sites)

    sweeps = Sweeps(2)
    maxdim!(sweeps, 10, 20, 100, 100, 200)
    cutoff!(sweeps, 1E-10)

    # Run the DMRG algorithm, returning energy
    # (dominant eigenvalue) and optimized MPS
    energy, psi = dmrg(H, psi0, sweeps; outputlevel=0)
    #println("Final energy = $energy")
  end
end

nothing


using Distances
using Clustering
using Statistics
using LinearAlgebra
using Distances
using RGBDSegmentation: CDNImage
"""
```
function initial_centroids(plain_data,  k::Integer)
```
Return a ``9 \times k`` matrix containing the centroids of each  cluster.

Each centroid is represented as a 9-dimensional vector in which the
first three components represent the color components,
the 4:6 indices represent the spatial components and
the last three component represent the normal component.

Each slice of the array are normalized. i.e.
``norm(c[1:3, j]) = 1``

"""
function initial_centroids(plain_data,   k::Integer)
    centroids = zeros(9, k)
    indices = [1:3, 4:6, 7:9]
    for i=1:3
        model = kmeans(
            plain_data[indices[i], :],
            k,
            maxiter=500)
        labels = assignments(model)
        for j=1:k
            centroids[indices[i], j] = mean(
                plain_data[indices[i], findall(labels.==j)],
                dims=2)
	    centroids[indices[i], j] /= norm(centroids[indices[i], j])
        end
    end
    return centroids
end
function indmin(x)
    (vals, cidxs) = findmin(x, dims=2)
    idx = [cidx[2] for cidx in cidxs]
    return (vals,idx)
end
"""
```
function initial_clusters(p_data::Array{T, 2},
                          k::Integer) where T <: Number
```
Return the initial assigments using a combined distance among the patterns of the three images.
"""
function initial_clusters(p_data::Array{T, 2},
                          k::Integer) where T <: Number

    centroids = initial_centroids(p_data, k)
    indices = [1:3, 4:6, 7:9]
    diff = 10000;
    epsilon = 0.1;
    value = 100;
    iteration = 1;
    maxnumIter = 50;
    clust = nothing;
    while (diff > epsilon) && (maxnumIter>iteration)
        iteration = iteration+1;
        oldvalue = value;
        distance_to_center = zeros(size(p_data,2), k)
        distances = [Euclidean(), Euclidean(), CosineDist()]
        for i=1:3, j=1:k
            distance_to_center[:,j] += colwise(
                distances[i],
                p_data[indices[i],:],
                vec(centroids[indices[i],j])
            )
        end
        # assign points to nearest cluster (based on minimum distance)
        (distMin,clust) = indmin(distance_to_center)
        value = sum(distMin);
        # compute new cluster centroids
	for h=1:k
    	  try
	    cluster_indices = vec(clust.==h)
            centroids[1:3, h] = mean(
                p_data[1:3, cluster_indices],
                dims=2)
            centroids[4:6, h] = mean(p_data[4:6,cluster_indices],
                                    dims=2)
            centroids[7:9, h] = sum(p_data[7:9,cluster_indices],
                                    dims=2)
            centroids[7:9, h] = centroids[7:9, h]./ norm(centroids[7:9, h]);
	   catch e

                succ = 0;
         	return;
	   end
	end
	indices_invalid_centroid = unique([id[2] for id in vec(findall(isnan.(centroids)))])
    	if (!isempty(indices_invalid_centroid))
            centroids = centroids[:, setdiff(1:k, indices_invalid_centroid)]
	    k = k - length(indices_invalid_centroid);
	end
	diff = abs(value - oldvalue);
    end
    return clust
end


using NCDatasets
using Formatting
using PyPlot
using Statistics

if length(ARGS) >= 1
    fn = ARGS[1]
end

const Re = 6371000.0
const cp = 1004.0

ds = Dataset(fn, "r")

lat = ds["lat"][:]
lon = ds["lon"][:]
lev = ds["lev"][:]
ilev = ds["ilev"][:]

deg2rad = π/180.0

cross_section = (lon[2] - lon[1]) * deg2rad * Re * cos.(lat * deg2rad)

tot_prec = ds["PRECC"][:] + ds["PRECL"][:]
tot_prec = view(tot_prec, :, :, 1)
tot_prec = view(mean(tot_prec, dims=(1,)), 1, :)       # take sum over lon
close(ds)

fig, ax = plt[:subplots](1, 1, figsize=(6,4))

ax[:plot]([-90, 90], [0, 0], linewidth=2, color="#cccccc", dashes=(5,3))
ax[:plot](lat, tot_prec * 86400.0 * 365 * 1000.0, linewidth=2, color="black")
ax[:set_xlabel]("Latitude [deg]")
ax[:set_ylabel]("Annual Precipiation [mm / yr]")

ax[:set_title](fn)
ax[:set_xlim](-90, 90)
ax[:set_ylim](-100, 2500)

plt[:show]()


import StatsBase: sample, WeightVec

# test logistic loss

## generate data
srand(1);
m,n,k = 1000,1000,3;
kfit = k+1
# variance of measurement
sigmasq = .1

# coordinates of covariates
X_real = randn(m,k)
# directions of observations
Y_real = randn(k,n)

XY = X_real*Y_real;
A = zeros(Int, (m, n))
logistic(x) = 1/(1+exp(-x))
for i=1:m
	for j=1:n
		A[i,j] = (logistic(XY[i,j]) >= rand()) ? true : false
	end
end

# and the model
losses = LogisticLoss()
rx, ry = QuadReg(.1), QuadReg(.1);
glrm = GLRM(A,losses,rx,ry,kfit)
gg = GGLRM(A, losses, rx, ry, kfit)
fit!(glrm);
fit!(gg);
@test_approx_eq_eps(whole_objective(gg, gg.X'gg.Y), whole_objective(gg, glrm.X'glrm.Y), 0.1*length(A))


using LinearAlgebra, SparseArrays, Statistics, StatsBase

# This expresses the probability the frequency of cells that carry variant a certain variant c_j is at least f.
φ(f,u,ρ,c) = (u / (u-log(1+(ρ-1)*f)))^c

# This calculates the expected number of cells in the population at a given time (measured here based on the discussion
# of the paper (see Equations 9 and 10 for details)
N(s,ρ,u) = exp((ρ/u -1/u)*s) + 2*sinh((1/u - ρ/u)*s)/ρ

function gmap(seed,u,ρ, debug = false)
    # Calculation of the approximate mapping g(c_j)
    # For details see Equation (10) of the manuscript
    if debug
        println("exp part = ",exp((ρ/u - 1/u)*seed))
        println("sinh part = ",2*sinh((1/u - ρ/u)*seed)/ρ)
    end
    return seed == 0 ? 0 : ceil(u*( exp((ρ/u - 1/u)*seed) + 2*sinh((1/u - ρ/u)*seed)/ρ))
end

function popEstimate(seed,u,ρ, debug = false)
    # This function returns an estimate of the expected number of cells in the population,
    # at the time  each of the seeding-cell-specific variants appeared. For details see the manuscript (Equations 9 and 10)
    if debug
        println("exp part = ",exp((ρ/u - 1/u)*seed))
        println("sinh part = ",2*sinh((1/u - ρ/u)*seed)/ρ)
    end
    return exp((ρ/u - 1/u)*seed) + 2*sinh((1/u - ρ/u)*seed)/ρ
end

function calculate_sampling_probabilities_populations(u, ts1, ts2, t_det, ρ, ρ1, ρ2 )
    # This refers to the 3rd population based seeding strategy
    # For context see the related discussion in Section 4.1.3 of the paper.
    p0 = zeros(t_det)
    p1 = zeros(t_det)
    p2 = zeros(t_det)
    for k=1:t_det
        if k < ts1
            N0 = N(k,ρ,u)
            p0[k] = 1
            continue
        end
        if k < ts2
            N0 = N(k,ρ,u)
            N1 = N(k-ts1,ρ1,u)
            p0[k] = N0 / (N0 + N1)
            p1[k] = N1 / (N0 + N1)
            continue
        end
        N0 = N(k,ρ,u)
        N1 = N(k-ts1,ρ1,u)
        N2 = N(k-ts2,ρ2,u)
        p0[k] = (N0) / (N0+N1+N2)
        p1[k] = (N1) / (N0+N1+N2)
        p2[k] = (N2) / (N0+N1+N2)
    end
    p0, p1, p2
end

function Bm(f,u,ρ,seed, debug = false)
    Bm = 0
    for j=0:seed
        Ptemp = 1.0
        if j != seed
            p = (1-φ(f,u,ρ,gmap(j+1,u,ρ))) - (1-φ(f,u,ρ,gmap(j,u,ρ)))
        else
            p = φ(f,u,ρ,gmap(j,u,ρ))
        end
        Bm += (seed-j)*p
        if debug
            println("P(U_j+1) = ", (1-φ(f,u,ρ,gmap(j+1,u,ρ))))
            println("P(U_j) = ", (1-φ(f,u,ρ,gmap(j,u,ρ))))
            println("Bm = ", Bm)
        end
    end
    return Bm
end

function BmGeneric(d, k_max, debug = false)
    Bm = 0
    BmPMF = zeros(k_max+1)
    for j=0:k_max
        if j != k_max
            p = (1-d[j+2]) - (1-d[j+1])
        else
            p = d[k_max]
        end
        Bm += (k_max-j)*p
        BmPMF[k_max-j+1] = p
    end
    return Bm, BmPMF
end

function BmDist(f,u,ρ,seed, debug = false)
    # Calculates the B_md distribution (See equation 6, and the related discussion in the paper)
    Bm = []
    for j=0:seed
        Ptemp = 1.0
        if j != seed
            p = (1-φ(f,u,ρ,gmap(j+1,u,ρ))) - (1-φ(f,u,ρ,gmap(j,u,ρ)))
        else
            p = φ(f,u,ρ,gmap(j,u,ρ))
        end
        append!(Bm,p)
        if debug
            println("P(U_j+1) = ", (1-φ(f,u,ρ,gmap(j+1,u,ρ))))
            println("P(U_j) = ", (1-φ(f,u,ρ,gmap(j,u,ρ))))
            println("Bm = ", Bm)
        end
    end
    return Bm
end

function BmVar(f,u,ρ,seed)
    # To measure the innate variance of the random variable Bmd that follows the above distribution BmDist
    Bm = 0
    Bm2 = 0
    for j=0:seed
        Ptemp = 1.0
        if j != seed
            p = (1-φ(f,u,ρ,gmap(j+1,u,ρ))) - (1-φ(f,u,ρ,gmap(j,u,ρ)))
        else
            p = φ(f,u,ρ,gmap(j,u,ρ))
        end
        Bm += (seed-j)*p
        Bm2 += (seed-j)^2*p
    end
    return Bm2 - Bm^2
end

# The following are simple Helper functions necessary to facilitate the production of the plots

function kmaxEstimate(population,u,ρ)
    k = 1
    while exp((ρ/u - 1/u)*k) + 2*sinh((1/u - ρ/u)*k)/ρ < population
        k += 1
    end
    return k
end

function findtd(u, ρ, population)
    s = 1
    while exp((ρ / u - 1 / u) * s) + 2 * sinh((1 / u - ρ / u) * s) / ρ < population
        s += 1
    end
    return s
end

function finds_i(u, ρ, time, population)
    x = 0.0
    ρ1 = ρ / (x + 1)
    while exp((ρ1 / u - 1 / u) * time) + 2 * sinh((1 / u - ρ1 / u) * time) / ρ1 < population
        x += 0.0005
        ρ1 = ρ / (x + 1)
    end
    return x
end

function findts_i(u, ρ, td0, population)
    t = 0
    while exp((ρ / u - 1 / u) * t) + 2 * sinh((1 / u - ρ / u) * t) / ρ < population
        t += 1
    end
    return td0 - t # returns the type origination time measured in [0..td]  (integer value, i.e., the number of mutations...)
end


@testset "distances.jl" begin
    p1 = Point(30.0, 100.0)
    p2 = Point(50.0, 210.0)
    p3 = Point(40.0, 180.0)

    ls = Arc(p1,p2)
    lss = Arcs([p1, p2, p3])

    @test (angular_distance(p1, p2)) ≈ 8773e3/6371008.7714 atol = 1e3
    @test (angular_length(ls)) ≈ 8773e3/6371008.7714 atol = 1e3
    @test (angular_distance(p1, p1)) == 0.0
    @test (angular_length(lss)) ≈ (8773e3+2585e3)/6371008.7714 atol = 1e3

    # Example http://edwilliams.org/avform.htm Cross track error
    pLAX = Point(33.95,-118.4)
    pJFK = Point(40.63333, -73.78333)
    pD = Point(34.5, -116.5)

    dist_LAX2D = angular_distance(pLAX, pD)
    @test dist_LAX2D ≈ rad2deg(0.02905) atol = 0.1
    bearing_LAX2D = azimuth(pLAX, pD)
    @test bearing_LAX2D ≈ 70.17 atol = 0.01
    bearing_LAX2JFK = azimuth(pLAX, pJFK)

    @test angular_distance(dist_LAX2D, bearing_LAX2D, bearing_LAX2JFK) ≈
    rad2deg(0.00216747) atol = 0.0001
    @test angular_distance(pD, pLAX, bearing_LAX2JFK) ≈
    rad2deg(0.00216747) atol = 0.0001
    @test angular_distance(pD, Line(pLAX, bearing_LAX2JFK)) ≈
    rad2deg(0.00216747) atol = 0.0001

    @test along_line_angular_distance(rad2deg(0.02905), rad2deg(0.00216747)) ≈
    rad2deg(0.0289691) atol = 0.00001
    @test along_line_angular_distance(pD, Line(pLAX, bearing_LAX2JFK)) ≈
    rad2deg(0.0289691) atol = 0.00001

    px1 = Point(-10.0, -20.0)
    px2 = Point(10.0, -20.0)
    px3a = Point(0.0, -10.0)
    @test angular_distance(px3a, px1, px2) ≈ 10.0 atol = 0.1
    px3b = Point(15.0, -20.0)
    @test angular_distance(px3b, px1, px2) ≈ 5.0 atol = 0.1
    px3c = Point(-17.0, -20.0)
    @test angular_distance(px3c, px1, px2) ≈ 7.0 atol = 0.1
    @test angular_distance(px3c, Arc(px1, px2)) ≈ 7.0 atol = 0.1

    px4 = Point(0.0, 0.0)
    px5 = Point(10.0, 0.0)
    arc = Arc(px4, px5)
    @test (angular_distance(Point(-5.0, 0.0), arc)) ≈ 556.0*1000/6371008.7714 atol = 100.0
    @test (angular_distance(Point(-5.0, 5.0), arc)) ≈ 785.8*1000/6371008.7714 atol = 100.0
    @test (angular_distance(Point(0.0, 5.0), arc)) ≈ 556.0*1000/6371008.7714 atol = 100.0
    @test (angular_distance(Point(5.0, 5.0), arc)) ≈ 553.9*1000/6371008.7714 atol = 100.0
    @test (angular_distance(Point(10.0, 5.0), arc)) ≈ 547.5*1000/6371008.7714 atol = 100.0
    @test (angular_distance(Point(15.0, 0.0), arc)) ≈ 556.0*1000/6371008.7714 atol = 100.0
    @test (angular_distance(Point(15.0, -5.0), arc)) ≈ 776.9*1000/6371008.7714 atol = 100.0
    @test (angular_distance(Point(10.0, -5.0), arc)) ≈ 547.5*1000/6371008.7714 atol = 100.0
    @test (angular_distance(Point(5.0, -5.0), arc)) ≈ 553.9*1000/6371008.7714 atol = 100.0
    @test (angular_distance(Point(0.0, -5.0), arc)) ≈ 556.0*1000/6371008.7714 atol = 100.0
    @test (angular_distance(Point(5.0, 0.0), arc)) ≈ 0.0*1000/6371008.7714 atol = 100.0

    @test (angular_distance(Point(-5.0, 0.0), px4)) ≈ 556.0*1000/6371008.7714 atol = 100.0
    @test (angular_distance(Point(-5.0, 5.0), px4)) ≈ 785.8*1000/6371008.7714 atol = 100.0
    @test (angular_distance(Point(0.0, 5.0), px4)) ≈ 556.0*1000/6371008.7714 atol = 100.0
    @test (angular_distance(Point(5.0, 5.0), arc)) ≈ 553.9*1000/6371008.7714 atol = 100.0
    @test (angular_distance(Point(10.0, 5.0), px5)) ≈ 547.5*1000/6371008.7714 atol = 100.0
    @test (angular_distance(Point(15.0, 0.0), px5)) ≈ 556.0*1000/6371008.7714 atol = 100.0
    @test (angular_distance(Point(15.0, -5.0), px5)) ≈ 776.9*1000/6371008.7714 atol = 100.0
    @test (angular_distance(Point(10.0, -5.0), px5)) ≈ 547.5*1000/6371008.7714 atol = 100.0
    @test (angular_distance(Point(5.0, -5.0), arc)) ≈ 553.9*1000/6371008.7714 atol = 100.0
    @test (angular_distance(Point(0.0, -5.0), px4)) ≈ 556.0*1000/6371008.7714 atol = 100.0
    @test (angular_distance(Point(5.0, 0.0), arc)) ≈ 0.0*1000/6371008.7714 atol = 100.0

    @test (angular_length(Polygon(p1, [p1, p2, p3, p1]))) ≈ (8773e3+
    2585e3)/6371008.7714+(angular_distance(p1, p3)) atol = 1e3

end


module DoublePendulum

using NeuralNetworkAnalysis, LaTeXStrings
using NeuralNetworkAnalysis: Specification
import Plots, DifferentialEquations
using Plots: plot, plot!, savefig, font, Measures.mm

# problem

m = 0.5
L = 0.5
c = 0.0
g = 1.0
gL = g/L
mL = m*L^2

@taylorize function double_pendulum!(dx, x, p, t)
    x₁, x₂, x₃, x₄, T₁, T₂ = x

    Δ12 = x₁ - x₂
    ★ = cos(Δ12)
    x3sin12 = x₃^2 * sin(Δ12)
    x4sin12 = x₄^2 * sin(Δ12) / 2
    gLsin1 = gL * sin(x₁)
    gLsin2 = gL * sin(x₂)
    T1_frac = (T₁ - c * x₃) / (2 * mL)
    T2_frac = (T₂ - c * x₄) / mL
    bignum = x3sin12 - ★ * (gLsin1 - x4sin12 + T1_frac) + gLsin2 + T2_frac
    denom = ★^2 / 2 - 1

    dx[1] = x₃
    dx[2] = x₄
    dx[3] = ★ * bignum / (2 * denom) - x4sin12 + gLsin1 + T1_frac
    dx[4] = - bignum / denom
end;

controller = read_nnet_sherlock(@modelpath "" "controller_Double-Pendulum_mr_sherlock")

X₀ = BallInf(fill(1.15, 4), 0.15)
X₀_small = BallInf(fill(1.000005, 4), 0.000005)
U₀ = ZeroSet(2)
vars_idx = Dict(:state_vars=>1:4, :control_vars=>5:6)
ivp = @ivp(x' = double_pendulum!(x), dim: 6, x(0) ∈ X₀ × U₀)
ivp_small = @ivp(x' = double_pendulum!(x), dim: 6, x(0) ∈ X₀_small × U₀)

period = 0.02
k = 20
T = k * period
T_warmup = 2 * period

prob = ControlledPlant(ivp, controller, vars_idx, period)
prob_small = ControlledPlant(ivp_small, controller, vars_idx, period)

# simulation

println("simulation")
res = @timed simulate(prob, T=T, trajectories=10, include_vertices=true)
sim = res.value
print_timed(res)

# reachability analysis

alg = TMJets20(abstol=1e-9, orderT=8, orderQ=1)
alg_nn = Ai2()

function benchmark(; T=T, silent::Bool=false)
    silent || println("flowpipe construction")
    res = @timed solve(prob, T=T, alg_nn=alg_nn, alg=alg)
    sol = res.value
    silent || print_timed(res)

    silent || println("flowpipe construction (small initial states)")
    res = @timed solve(prob_small, T=T, alg_nn=alg_nn, alg=alg)
    sol_small = res.value
    silent || print_timed(res)

    return sol, sol_small
end

benchmark(T=T_warmup, silent=true)
res = @timed benchmark()
sol, sol_small = res.value

# plotting

vars = (1, 2)
fig = plot(xlab=L"\theta", ylab=L"\theta'", leg=:topleft,
		   legendfontsize=5,
		   tickfont=font(25, "Times"),
		   guidefontsize=25,
		   xguidefont=font(25, "Times"),
		   yguidefont=font(25, "Times"),
		   xticks=([1, 1.25, 1.5, 1.75, 2], ["1.0", "1.25", "1.5", "1.75", "2.0"]),
		   bottom_margin=0mm, left_margin=0mm, right_margin=4mm, top_margin=0mm, size=(500, 450))
plot!(fig, sol, vars=vars, color=:orange, lab="")
plot!(project(X₀, vars), c=:blue, lab=L"X_0", alpha=0.6)
plot!(fig, sol_small, vars=vars, color=:red, lab="")
plot!(project(X₀_small, vars), c=:cyan, lab=L"X_0'", alpha=1)
plot_simulation!(fig, sim; vars=vars, color=:black, lab="")
savefig("doublependulum_more_robust_x1x2_juliareach.png")

end
nothing


function discrete_phs(N,a,b)
	Nx = N   	#    energy variables: x
	Ne = Nx+1   # co-energy variables: e

	ei,we = lgwt(Ne,a,b)    # discretization of co-energy variables
	#ei = collect(linspace(a,b,Ne,))
	if Nx > 1
		xi,wx = lgwt(Nx,a,b)
		#xi = collect(linspace(a,b,Nx))
	else
		xi,wx = (a+b)/2, (b-a);
	end
	#M = massmatrix(ei,xi,we);
	xquad, wquad = lgwt(Ne,a,b)
	M = massmatrix(ei,xi, xquad, wquad);
	D = dermatrix(ei,xi,1);
	p0 = map(i->leg_pol(a,ei,i), 1:length(ei))
	pL = map(i->leg_pol(b,ei,i), 1:length(ei))

	F = [M 0*M;
	     pL' pL'*0;
		 0*M M;
		 0*p0' p0'];
	E = -[D*0 D;
	      pL'*0 -pL';
		  D D*0;
		  p0' p0'*0];
		  J = E/F
    phJ = -[zeros(N,N) J[1:N,(N+2):end-1];
		          J[N+2:end-1,1:N] zeros(N,N)];
    phB = -[J[[1:N;N+2:2*N+1],N+1] J[[1:N;(N+2):(2*N+1)],end]];
    phD = [J[N+1,N+1] J[N+1,end];
        J[end,N+1] J[end,end]];

	if N > 1
		Q = massmatrix(xi,xi,xquad,wquad)
		#Q = sparse(diagm(wx[:],0));
		Q = blkdiag(Q,Q);
		else
		#Q = diagm([wx; wx][:]);
	end

	pflow = Collocation_points(xi, wx)
	peffort = Collocation_points(ei, we)

	p = Phs(phJ, phB, phD, Q)
	#p.Q = Q
	p.disc_data = Discrete_domain(pflow, peffort)
	#func(x :: Array) = (0.5*x'*p.Q*x)[1]
	#p.Hamiltonian = func
	#p.GradHam = ForwardDiff.gradient(p.Hamiltonian)
	p
end

function discrete_phs2(N,a,b)
	Nx = N
	Ne = Nx+2
	ei,wi = lgwt(Ne,a,b)
	if Nx > 2
		xi,wx = lgwt(Nx,a,b)
	else
		xi,wx = (a+b)/2, (b-a);
	end
	M = massmatrix(ei,xi,wi);
	D = dermatrix(ei,xi,2);
	leg_pol_der = i -> (algo_diff(x->leg_pol(x,ei,i), 1))
	p0 = map(i->leg_pol(a,ei,i), 1:length(ei))
	p0d = map(i->leg_pol_der(i)(a), 1:length(ei))
	pL = map(i->leg_pol(b,ei,i), 1:length(ei))
	pLd = map(i->leg_pol_der(i)(b), 1:length(ei))

	F = [M 0*M;
	     pL' pL'*0;
	     pLd' pL'*0
		 0*M M;
	     0*pLd' p0d'
		 0*p0' p0'];
	E = [D*0 D;
	      pL'*0 -pLd';
	      pL'*0  pL';
	       -D D*0;
	       p0' p0'*0;
	      -p0d' p0'*0];
	J = E/F
	neworder = [1:N; (1:N) .+ (N+2); N+1:N+2; 2*N+3:2*N+4];
    Jnew = J[neworder,neworder]
    phJ = - Jnew[1:2*N,1:2*N];
    phB = - Jnew[1:2*N, 2*N .+ (1:4)];
    phD = Jnew[2*N .+ (1:4), 2*N .+ (1:4)];

	if N > 1
		#Q = (diagm(wx[:],0));
		Q = (diagm(0 => wx[:]));
		Q = blkdiag(Q,Q);
		else
		Q = diagm([wx; wx][:]);
	end

	pflow = Collocation_points(xi, wx)
	peffort = Collocation_points(ei, wi)

	p = Phs(phJ, phB, phD, Q)
	#p.Q = Q
	p.disc_data = Discrete_domain(pflow, peffort)
	#func(x :: Array) = (0.5*x'*p.Q*x)[1]
	#p.Hamiltonian = func
	#p.GradHam = ForwardDiff.gradient(p.Hamiltonian)
	p
end

function discrete_phs2_distports(N,K,a,b,ad,bd)
	Nx = N
	Ne = Nx+2
	zei,we = lgwt(Ne,a,b)
	zvi,wv = lgwt(K,ad,bd)
	if Nx > 2
		zxi,wx = lgwt(Nx,a,b)
	else
		zxi,wx = (a+b)/2, (b-a);
	end
	M = massmatrix(zei,zxi,we);

	leg_pol_der2 = (z,zi,i) -> algo_diff(x->leg_pol(x,zi,i),2)(z);
	# M_theta_phizz can be exactly obtained by quadrature if K >= N-3:
	if K < N-3
		warn("K < N-3: The distributed ports matrix is not exactly computed by quadrature")
	end
	M_theta_phizz = massmatrix(zvi,zxi,wv, leg_pol_der2);
	M_phi = diagm(0 => wx[:]);
	B_bar = M_phi \ M_theta_phizz;
	D = dermatrix(zei,zxi,2);
	leg_pol_der = i -> (algo_diff(x->leg_pol(x,zei,i), 1))
	p0 = map(i->leg_pol(a,zei,i), 1:length(zei))
	p0d = map(i->leg_pol_der(i)(a), 1:length(zei))
	pL = map(i->leg_pol(b,zei,i), 1:length(zei))
	pLd = map(i->leg_pol_der(i)(b), 1:length(zei))

	F = [M 0*M;
	     p0' pL'*0;
	     p0d' pL'*0
		 0*M M;
	     0*pLd' pLd'
		 0*p0' pL'];
	F = blkdiag(F,Matrix(1I, K, K));
	E = [D*0 D;
	      pL'*0 p0d';
	      pL'*0  -p0';
	       -D D*0;
	       -pL' p0'*0;
	      pLd' p0'*0];
	E = blkdiag(E, zeros(K,K));
	E[1:size(D,1),end-K+1:end] = - B_bar;
    E[end-K+1:end,1:size(D,2)] =  B_bar' * M;
	J = E/F
	neworder = [1:N; (1:N) .+ (N+2); N+1:N+2; 2*N+3:2*N+4; (2*N+4) .+ (1:K)];
    Jnew = J[neworder,neworder]
    phJ = - Jnew[1:2*N,1:2*N];
    phB = - Jnew[1:2*N, 2*N .+ (1:4)];
	phBdistributed = - Jnew[1:2*N, (2*N + 4) .+ (1:K)];
    phD = Jnew[2*N .+ (1:4), 2*N .+ (1:4)];

	if N > 1
		Q = (diagm(0 => wx[:]));
		Q = blkdiag(Q,Q);
		else
		Q = diagm([wx; wx][:]);
	end

	pflow = Collocation_points(zxi, wx)
	peffort = Collocation_points(zei, we)
	pdistributed = Collocation_points(zvi, wv)

	p = Phs(phJ, phB, phD, Q)
	p.Bd = phBdistributed
	#p.Q = Q
	p.disc_data = Discrete_domain(pflow, peffort)
	#func(x :: Array) = (0.5*x'*p.Q*x)[1]
	#p.Hamiltonian = func
	#p.GradHam = ForwardDiff.gradient(p.Hamiltonian)
	p
end


@doc raw"""
    βdifferential_geodesic_startpoint(κ,t,d)

weights for the [`jacobi_field`](@ref) corresponding to the differential of the geodesic
with respect to its start point $D_x g(t;p,q)[X]$.
They are

```math
\beta(\kappa) = \begin{cases}
\frac{\sinh(d(1-t)\sqrt{-\kappa})}{\sinh(d\sqrt{-\kappa})}
&\text{ if }\kappa < 0,\\
1-t & \text{ if } \kappa = 0,\\
\frac{\sin((1-t)d\sqrt{\kappa})}{\sinh(d\sqrt{\kappa})}
&\text{ if }\kappa > 0.
\end{cases}
```

Due to a symmetry agrument, these are also used to compute $D_q g(t; p,q)[\eta]$

# See also

[`differential_geodesic_endpoint`](@ref), [`differential_geodesic_startpoint`](@ref), [`jacobi_field`](@ref)
"""
function βdifferential_geodesic_startpoint(κ, t, d)
    (κ == 0) && return 1.0 - t
    (d == 0) && return 1.0 - t
    (κ < 0) && return sinh(sqrt(-κ) * (1.0 - t) * d) / sinh(sqrt(-κ) * d)
    return (κ > 0) && return sin(sqrt(κ) * (1.0 - t) * d) / sin(sqrt(κ) * d)
end
@doc raw"""
    βdifferential_exp_basepoint(κ,t,d)

weights for the [`jacobi_field`](@ref) corresponding to the differential of the geodesic
with respect to its start point $D_p \exp_p X [Y]$. They are

```math
\beta(\kappa) = \begin{cases}
\cosh(\sqrt{-\kappa})&\text{ if }\kappa < 0,\\
1 & \text{ if } \kappa = 0,\\
\cos(\sqrt{\kappa}) &\text{ if }\kappa > 0.
\end{cases}
```

# See also

[`differential_exp_basepoint`](@ref), [`jacobi_field`](@ref)
"""
function βdifferential_exp_basepoint(κ, ::Number, d)
    (κ < 0) && return cosh(sqrt(-κ) * d)
    (κ > 0) && return cos(sqrt(κ) * d)
    return 1.0
end
@doc raw"""
    βdifferential_exp_argument(κ,t,d)

weights for the [`jacobi_field`](@ref) corresponding to the differential of the geodesic
with respect to its start point $D_X \exp_p X[Y]$. They are

```math
$\beta(\kappa) = \begin{cases}
\frac{\sinh(d\sqrt{-\kappa})}{d\sqrt{-\kappa}}&\text{ if }\kappa < 0,\\
1 & \text{ if } \kappa = 0,\\
\frac{\sin(d\sqrt{\kappa})}{\sqrt{d\kappa}}&\text{ if }\kappa > 0.
\end{cases}
```

# See also

[`differential_exp_argument`](@ref), [`jacobi_field`](@ref)
"""
function βdifferential_exp_argument(κ, ::Number, d)
    ((κ == 0) || (d == 0)) && return 1.0
    (κ < 0) && return sinh(sqrt(-κ) * d) / (d * sqrt((-κ)))
    return (κ > 0) && return sin(sqrt(κ) * d) / (d * sqrt(κ))
end
@doc raw"""
    βdifferential_log_basepoint(κ,t,d)

weights for the [`jacobi_field`](@ref) corresponding to the differential of the geodesic
with respect to its start point $D_p \log_p q[X]$. They are

```math
\beta(\kappa) = \begin{cases}
-\sqrt{-\kappa}d\frac{\cosh(d\sqrt{-\kappa})}{\sinh(d\sqrt{-\kappa})}&\text{ if }\kappa < 0,\\
-1 & \text{ if } \kappa = 0,\\
-\sqrt{\kappa}d\frac{\cos(d\sqrt{\kappa})}{\sin(d\sqrt{\kappa})}&\text{ if }\kappa > 0.
\end{cases}
```

# See also

[`differential_log_argument`](@ref), [`differential_log_argument`](@ref), [`jacobi_field`](@ref)
"""
function βdifferential_log_basepoint(κ, ::Number, d)
    ((d == 0) || (κ == 0)) && return -1.0
    (κ < 0) && return -sqrt(-κ) * d * cosh(sqrt(-κ) * d) / sinh(sqrt(-κ) * d)
    return (κ > 0) && return -sqrt(κ) * d * cos(sqrt(κ) * d) / sin(sqrt(κ) * d)
end
@doc raw"""
    βdifferential_log_argument(κ,t,d)

weights for the JacobiField corresponding to the differential of the logarithmic
map with respect to its argument $D_q \log_p q[X]$. They are

```math
\beta(\kappa) = \begin{cases}
\frac{ d\sqrt{-\kappa} }{\sinh(d\sqrt{-\kappa})}&\text{ if }\kappa < 0,\\
1 & \text{ if } \kappa = 0,\\
\frac{ d\sqrt{\kappa} }{\sin(d\sqrt{\kappa})}&\text{ if }\kappa > 0.
\end{cases}
```

# See also

[`differential_log_basepoint`](@ref), [`jacobi_field`](@ref)
"""
function βdifferential_log_argument(κ, ::Number, d)
    ((d == 0) || (κ == 0)) && return 1.0
    (κ < 0) && return sqrt(-κ) * d / sinh(sqrt(-κ) * d)
    return sqrt(κ) * d / sin(sqrt(κ) * d)
end

@doc raw"""
    Y = adjoint_Jacobi_field(M, p, q, t, X, β)

Compute the AdjointJacobiField $J$ along the geodesic $γ_{p,q}$ on the manifold
$\mathcal M$ with initial conditions (depending on the application)
$X ∈ T_{γ_{p,q}(t)}\mathcal M$ and weights $β$. The result is a vector
$Y ∈ T_p\mathcal M$. The main difference to [`jacobi_field`](@ref) is the,
that the input `X` and the output `Y` switched tangent spaces.
For detais see [`jacobi_field`](@ref)
"""
function adjoint_Jacobi_field(M::Manifold, p, q, t, X, β=βdifferential_geodesic_startpoint)
    x = shortest_geodesic(M, p, q, t)
    B = get_basis(M, p, DiagonalizingOrthonormalBasis(log(M, p, q)))
    V = get_vectors(M, p, B)
    Θ = vector_transport_to.(Ref(M), Ref(p), V, Ref(x), Ref(ParallelTransport()))
    # Decompose wrt. frame, multiply with the weights from w and recompose with Θ.
    Y = sum(
        (inner.(Ref(M), Ref(x), Ref(X), Θ)) .*
        (β.(B.data.eigenvalues, Ref(t), Ref(distance(M, p, q)))) .* V,
    )
    return Y
end
function adjoint_Jacobi_field(
    M::Circle{ℝ}, p::Real, q::Real, t::Real, X::Real, β=βdifferential_geodesic_startpoint
)
    x = shortest_geodesic(M, p, q, t)
    B = get_basis(M, p, DiagonalizingOrthonormalBasis(log(M, p, q)))
    V = get_vectors(M, p, B)
    Θ = vector_transport_to.(Ref(M), Ref(p), V, Ref(x), Ref(ParallelTransport()))
    # Decompose wrt. frame, multiply with the weights from w and recompose with Θ.
    Y = sum(
        (inner.(Ref(M), Ref(x), Ref(X), Θ)) .*
        (β.(B.data.eigenvalues, Ref(t), Ref(distance(M, p, q)))) .* V,
    )[1]
    return Y
end
function adjoint_Jacobi_field(
    M::AbstractPowerManifold, p, q, t, X, β=βdifferential_geodesic_startpoint
)
    rep_size = representation_size(M.manifold)
    Y = allocate_result(M, adjoint_Jacobi_field, p, X)
    for i in get_iterator(M)
        #lY = adjoint_Jacobi_field(M.manifold, p[i], q[i], t, X[i], β)
        Y[M, i] = adjoint_Jacobi_field(M.manifold, p[M, i], q[M, i], t, X[M, i], β)
    end
    return Y
end

@doc raw"""
    Y = jacobi_field(M, p, q, t, X, β)
compute the Jacobi jield $J$ along the geodesic $γ_{p,q}$ on the manifold $\mathcal M$ with
initial conditions (depending on the application) $X ∈ T_p\mathcal M$ and weights $β$. The
result is a tangent vector `Y` from $T_{γ_{p,q}(t)}\mathcal M$.

# See also

[`adjoint_Jacobi_field`](@ref)
"""
function jacobi_field(M::Manifold, p, q, t, X, β=βdifferential_geodesic_startpoint)
    x = shortest_geodesic(M, p, q, t)
    B = get_basis(M, p, DiagonalizingOrthonormalBasis(log(M, p, q)))
    V = get_vectors(M, p, B)
    Θ = vector_transport_to.(Ref(M), Ref(p), V, Ref(x), Ref(ParallelTransport()))
    Y = zero_tangent_vector(M, p)
    # Decompose wrt. frame, multiply with the weights from w and recompose with Θ.
    Y .= sum(
        (inner.(Ref(M), Ref(p), Ref(X), V)) .*
        (β.(B.data.eigenvalues, Ref(t), Ref(distance(M, p, q)))) .* Θ,
    )
    return Y
end
function jacobi_field(
    M::Circle{ℝ}, p::Real, q::Real, t::Real, X::Real, β=βdifferential_geodesic_startpoint
)
    x = shortest_geodesic(M, p, q, t)
    B = get_basis(M, p, DiagonalizingOrthonormalBasis(log(M, p, q)))
    V = get_vectors(M, p, B)
    Θ = vector_transport_to.(Ref(M), Ref(p), V, Ref(x), Ref(ParallelTransport()))
    Y = zero_tangent_vector(M, p)
    # Decompose wrt. frame, multiply with the weights from w and recompose with Θ.
    Y = sum(
        (inner.(Ref(M), Ref(p), Ref(X), V)) .*
        (β.(B.data.eigenvalues, Ref(t), Ref(distance(M, p, q)))) .* Θ,
    )[1]
    return Y
end
function jacobi_field(
    M::AbstractPowerManifold, p, q, t, X, β=βdifferential_geodesic_startpoint
)
    rep_size = representation_size(M.manifold)
    Y = allocate_result(M, adjoint_Jacobi_field, p, X)
    for i in get_iterator(M)
        Y[M, i] = jacobi_field(M.manifold, p[M, i], q[M, i], t, X[M, i], β)
    end
    return Y
end


"""
A node type that applies the swing equation and voltage dynmaics
to the frequency/angle dynamics.
In the following, we followed the implementation of the Extended Model/ third
order model including voltage dynamics according to K. Schmietendorf et. al.
"Self-Organized Synchronization and Voltage Stability in Networks of Synchronous
Machines", 2013.
The third order model is a special case of the fourth order model with Ed = 0
and Xd = Xd_dash.

Additionally to u, it has the internal dynamic variable ω representing the
frequencyof the rotator relative to the grid frequency Ω, i.e. the real
frequency ωr of the rotator is given as ωr=Ω+ω.

Keyword Arguments

H: inertia constant (given in [s]
P: active (real) power output, also called the mechanical torque applied to the shaft, given in [pu]
D: damping coefficient, (given in [s], see P. Sauer, eq. (5.156) where the damping torque is equal Dω)
Ω: rated frequency in [1/s] of the power grid, often 2π⋅50Hz
T_d_dash: time constant of d-axis, given in [s], see P. Sauer, chapter 3.7, p. 54 for a general explanation on time constants
X_d_dash: transient reactance of d-axis, given in [pu]
X_d: reactance of d-, given in [pu]
X_d: reactance of q-axis, given in [pu]
E_f: scaled field voltage, which, if set equal to 1.0 pu, gives 1.0 pu open-circuit terminal voltage. The physical device that provides the value of `E_f` is called the exciter (according to P. Sauer, p. 65)
"""

using PowerDynamics
using NetworkDynamics
import PowerDynamics: dimension, symbolsof, construct_vertex
import PowerDynamics: AbstractNode
import Base: @__doc__

@DynamicNode ThirdOrderEq(H, P, D, Ω, E_f,T_d_dash, X_q_dash, X_d_dash, X_d) begin
    @assert H > 0 "inertia (H) should be >0"
    @assert D >= 0 "damping (D) should be >=0"
    @assert T_d_dash > 0 "time constant of d-axis (T_d_dash) should be >0"
    @assert X_q_dash >= 0
    @assert X_d_dash >= 0 "transient reactance of d-axis (X_d_dash) should be >=0"
    @assert X_d >= 0 "reactance of d-axis (X_d_dash) should be >=0"
    Ω_H = Ω / (2 * H)
end [[θ,dθ],[ω, dω]] begin
    println(e_s, e_d)
    i_c = 1im * i * exp(-1im * θ)
    e_c = 1im * u * exp(-1im * θ)
    p = real(u * conj(i))
    e_d = 0                         # simplification from K.Schmietendorf paper
    e_q = imag(e_c)
    i_d = real(i_c)
    i_q = imag(i_c)

    dθ = ω
    dω = (P - D * ω - p - (X_q_dash - X_d_dash) * i_d * i_q) * Ω_H
    de_q = (1 / T_d_dash) * (E_f - e_q + i_d * (X_d - X_d_dash))
    # -> u = e_q * exp(1im * θ)
    du = de_q * exp(1im * θ) + u * 1im * ω

    #de_c = 1im*de_q
    #du = -1im * de_c * exp(1im * θ)+ u * 1im * ω
    #dω = (P - D * ω - p - (X_q_dash - X_d_dash) * i_d * i_q) * Ω_H
end

export ThirdOrderEq


using LinearAlgebra

function ESPRIT(Y,K,m=0)
    # Estimation of signal parameters via rotational invariance techniques
    # (ESPRIT) algorithm for line spectral estimation from measurements.
    # Can handle multiple snapshots by constructing a low rank Hankel matrix
    # from the data.
    #
    # Inputs:
    # Y: Data matrix
    # K: number of sinusoids
    # m: Hankel parameter of the Hankel matrix of the data
    #
    # Parameters:
    # N = size(Y,1): number of measurements per snapshot
    # L = size(Y,2): number of snapshots
    #
    # Output:
    # f_hat: estimate of the frequencies of the K sinusoids in [rad]

    (N,L) = size(Y)

    #C = X*X'/L

    if m == 0
    m = ceil(Int,N/2)
    end

    C = zeros(m,m)
    for j = 1:L, i = m:N
        C = C + Y[i-m+1:i,j]*Y[i-m+1:i,j]'/(L*N)
    end
    N = m

    evdC = eigen(C)
    p = sortperm(evdC.values,rev=true)[1:K]

    S = evdC.vectors[:,p]

    phi = S[1:N-1,:]\S[2:N,:]

    f_hat = ((angle.(eigvals(phi)))/(2*pi) .+ 1).%1
    #f_hat = real(log.(eigvals(phi))/(1im*2*pi))
end


@testset "polyagamma" begin
    @test mean(PolyaGamma(1, 0)) == 1 / 4
    @test mean(PolyaGamma(1, 2.0)) == tanh(1.0) / 4

    for (b, c) in ((1, 0), (1, 2.0), (3, 0), (3, 2.5), (3, 3.2), (1.2, 3.2))
        p = PolyaGamma(b, c)
        @test logpdf(p, rand(p)) isa Real
        @test mean(rand(p, 10000)) ≈ mean(p) atol = 1e-2
    end

    p = PolyaGamma(1, 0)
    @test Distributions.insupport(p, 0)
    @test !Distributions.insupport(p, -1)
    @test minimum(p) === 0
    @test maximum(p) == Inf
    @test Distributions.params(p) == (1, 0)
end


import StatsBase

single_estimate_funcs = [
    StatsBase.var,
    StatsBase.std,
    StatsBase.middle,
    StatsBase.median,
    StatsBase.mean,
    StatsBase.genmean,
    StatsBase.genvar,
    StatsBase.harmmean,
    StatsBase.mode,
    StatsBase.percentile,
    StatsBase.quantile,
    StatsBase.rle,
    StatsBase.sem,
    StatsBase.span,
    StatsBase.summarystats,
    StatsBase.totalvar,
    StatsBase.kurtosis,
    StatsBase.moment,
    StatsBase.skewness,
    StatsBase.renyientropy
]

n = 10

udata = UncertainValueDataset(example_uvals)

@testset "Single-estimate statistic for dataset" begin
    @testset "$(single_estimate_funcs[i])" for i = 1:length(single_estimate_funcs)
        f = single_estimate_funcs[i]

        if f == StatsBase.summarystats
            @test f(udata, n) isa Vector{StatsBase.SummaryStats{T}} where T
        elseif f == StatsBase.percentile
            @test f(udata, 10, n) isa Vector{T} where T <: Real
            @test f(udata, [10, 20], n) isa Vector{Vector{T}} where T <: Real
        elseif f == StatsBase.quantile
            @test f(udata, 0.1, n) isa Vector{T} where T <: Real
            @test f(udata, [0.2, 0.5], n) isa Vector{Vector{T}} where T <: Real
        elseif f == StatsBase.moment
            @test f(udata, 1, n) isa Vector{T} where T <: Real
            @test f(udata, 2, n) isa Vector{T} where T <: Real
        elseif f == StatsBase.genmean
            @test f(udata, 4, n) isa Vector{T} where T <: Real
        elseif f == StatsBase.rle
            @test rle(udata, n) isa Vector{Tuple{Vector{T1}, Vector{T2}}} where {T1, T2}
        elseif f == StatsBase.span
            @test f(udata, n) isa Vector{T} where {T <: AbstractRange{S}} where S
        elseif f == StatsBase.renyientropy
            @test f(udata, 1, n) isa Vector{T} where T <: Real
        else
            @test f(udata, n) isa Vector{T} where T <: Real
        end
    end;
end


## Preloads
using Statistics
using FFTW
using Plots
using BenchmarkTools
using Profile
using LinearAlgebra

# %% Code Cell
function fink_filter_bank_fast(J, L)

    # -------- set parameters
    dθ   = π/8        # 8 angular bins hardwired
    nx   = 256
    dx   = nx/2-1

    # -------- allocate output array of zeros
    filt = zeros(256, 256, J, L)

    # -------- allocate theta and logr arrays
    logr = zeros(nx, nx)
    θ    = zeros(nx, nx)

    for l = 0:L-1
        θ_l = dθ*l

    # -------- allocate anggood BitArray
        anggood = falses(nx, nx)

    # -------- loop over pixels
        for x = 1:nx
            sx = mod(x+dx,nx)-dx -1    # define sx,sy so that no fftshift() needed
            for y = 1:nx
                sy = mod(y+dx,nx)-dx -1
                θ_pix  = mod(atan(sy, sx)+π -θ_l, 2*π)
                θ_good = abs(θ_pix-π) <= dθ

            # If this is a pixel we might use, calculate log2(r)
                if θ_good
                    anggood[y, x] = θ_good
                    θ[y, x]       = θ_pix
                    r = sqrt(sx^2 + sy^2)
                    logr[y, x] = log2(max(1,r))
                end
            end
        end
        angmask = findall(anggood)
    # -------- compute the wavelet in the Fourier domain
    # -------- the angular factor is the same for all j
        F_angular = cos.((θ[angmask].-π).*4)

    # -------- loop over j for the radial part
        for j = 0:J-1
            jrad  = 7-j
            Δj    = abs.(logr[angmask].-jrad)
            rmask = (Δj .<= 1)

    # -------- radial part
            F_radial = cos.(Δj[rmask] .* (π/2))
            ind      = angmask[rmask]
            filt[ind,j+1,l+1] = F_radial .* F_angular[rmask]
        end
    end
    return filt
end

fink_filter_set = fink_filter_bank_fast(8,8)
test_img = ones(256,256)
copyto!(test_img,fink_filter_set[:,:,1,3])

# %% Code Cell

function speedy_DHC(image::Array{Float64,2}, filter_set::Array{Float64,4})
    FFTW.set_num_threads(2)

    # array sizes
    (Nx, Ny)  = size(image)
    (_,_,J,L) = size(fink_filter_set)

    out_coeff = []
    # allocate coeff arrays
    S0  = zeros(2)
    S1  = zeros(J, L)
    S20 = zeros(Float64, J*L, J*L)
    S12 = zeros(Float64, J*L, J*L)
    im_rd_0_1  = zeros(Float64, Nx, Ny, J, L)
    im_fd_0_1 = zeros(Float64, Nx, Ny, J, L)

    ## 0th Order
    S0[1]   = mean(image)
    norm_im = image.-S0[1]
    S0[2]   = BLAS.dot(norm_im,norm_im)/(Nx*Ny)
    norm_im ./= sqrt(Nx*Ny*S0[2])

    append!(out_coeff,S0[:])

    ## 1st Order
    im_fd_0 = fft(norm_im)

    Atmp = zeros(ComplexF64,Nx,Ny)
    Btmp = zeros(Float64,Nx,Ny)
    ## Main 1st Order and Precompute 2nd Order
    @views for l = 1:L
        for j = 1:J
            @inbounds copyto!(Atmp, im_fd_0 .* filter_set[:,:,j,l])
            Btmp .= abs.(Atmp)
            @inbounds im_fd_0_1[:,:,j,l] .= Btmp
            @inbounds S1[j,l] +=BLAS.dot(Btmp,Btmp)
            ifft!(Atmp)
            @inbounds im_rd_0_1[:,:,j,l] .= abs.(Atmp)
        end
    end
    append!(out_coeff,S1[:])

    #DPF version of 2nd Order
    Amat    = reshape(im_fd_0_1, Nx*Ny, J*L)
    mul!(S12,Amat',Amat)

    Amat    = reshape(im_rd_0_1, Nx*Ny, J*L)
    mul!(S20,Amat',Amat)

    append!(out_coeff,S20[:])
    append!(out_coeff,S12[:])

    return out_coeff
end

temp = speedy_DHC(test_img,fink_filter_set)

@time speedy_DHC(test_img,fink_filter_set)

Profile.clear()
@profile speedy_DHC(test_img,fink_filter_set)

Juno.profiler()

BenchmarkTools.DEFAULT_PARAMETERS.seconds = 30
@benchmark speedy_DHC(test_img,fink_filter_set)

function speedy_DHC_old(image::Array{Float64,2}, filter_set::Array{Float64,4})
    FFTW.set_num_threads(2)

    # array sizes
    (Nx, Ny)  = size(image)
    (_,_,J,L) = size(fink_filter_set)

    out_coeff = []
    # allocate coeff arrays
    S0  = zeros(2)
    S1  = zeros(J, L)
    S20 = zeros(Float64, J, L, J, L)
    S12 = zeros(Float64, J, L, J, L)
    im_rd_0_1  = zeros(Float64,    Nx, Ny, J, L)
    im_fdf_0_1 = zeros(Float64,    Nx, Ny, J, L)
    #im_fd_0_1  = zeros(ComplexF64, Nx, Ny, J, L)

    ## 0th Order
    S0[1]   = mean(image)
    norm_im = image.-S0[1]
    S0[2]   = sum(norm_im .* norm_im)/(Nx*Ny)
    norm_im ./= sqrt(Nx*Ny*S0[2])

    # Was this intentional to set this back to image?
    norm_im = image

    append!(out_coeff,S0[:])

    ## 1st Order
    im_fd_0 = fft(norm_im)

    foo = zeros(ComplexF64,Nx,Ny)
    Btmp = zeros(Float64,Nx,Ny)
    ## Main 1st Order and Precompute 2nd Order
    for l = 1:L
        for j = 1:J
            #@inbounds had!(im_fd_0_1[:,:,j,l],filter_set[:,:,j,l]) #wavelet already in fft domain not shifted
            # @inbounds Btmp .= abs.(im_fd_0_1[:,:,j,l])
            # We don't need an fftshift()
            # @inbounds im_fdf_0_1[:,:,j,l] .= fftshift(Btmp)
            copyto!(foo, im_fd_0 .* filter_set[:,:,j,l])
            Btmp .= abs.(foo)
            @inbounds im_fdf_0_1[:,:,j,l] .= Btmp
            # I think you are doing a 256x256 matrix multiplication here
            #@inbounds S1[j,l]+=BLAS.dot(Btmp,Btmp) #normalization choice arb to make order unity
            # You mean this:  (and BLAS won't speed this up much)
            # S1[j,l] = sum(Btmp.*Btmp)
            S1[j,l] = sum([x^2 for x in Btmp])  #slightly faster? (no malloc)
            #@inbounds Atmp .= ifft(foo)
            @inbounds im_rd_0_1[:,:,j,l] .= abs.(ifft(foo))

        end
    end
    append!(out_coeff,S1[:])

    #DPF version of 2nd Order
    Amat    = reshape(im_fdf_0_1, Nx*Ny, J*L)
    S12 = reshape(Amat' * Amat, J, L, J, L)

    Amat    = reshape(im_rd_0_1, Nx*Ny, J*L)
    S20 = reshape(Amat' * Amat, J, L, J, L)

    append!(out_coeff,S20)
    append!(out_coeff,S12)

    return out_coeff
end

temp = speedy_DHC_old(test_img,fink_filter_set)

@time speedy_DHC_old(test_img,fink_filter_set)

Profile.clear()
@profile speedy_DHC_old(test_img,fink_filter_set)

Juno.profiler()

BenchmarkTools.DEFAULT_PARAMETERS.seconds = 30
@benchmark speedy_DHC_old(test_img,fink_filter_set)