using SignalDecomposition
using Interpolations
using Plots

const _SDC = SignalDecomposition
const _INT = Interpolations

const freq = 50.0
const nw_id_default = 1

xfmr = Dict("voltage_harmonics" => [1,3],
            "current_harmonics" => [1,3], #this one has to match the data in the matpower file
            "N" => 50,
            "current_type" => :rectangular,
            "excitation_type" => :sigmoid,
            "inom" => 0.7,
            "ψmax" => 0.005,
            "voltage_type" => :rectangular,
            "dv" => [0.1,0.1],
            "vmin" => -[1.1,1.1],
            "vmax" => [1.1,1.1],
            "dθ" => [π/5,π/5],
            "θmin" => [0.0,0.0],
            "θmax" => [2π,2π])

excitation_flux_polar(V, θ, w, t) =
    sum(V[n] ./ w[n] .* sin.(w[n] .* t .+ θ[n]) for n in 1:length(V))
excitation_flux_rectangular(Vre, Vim, w, t) =
    excitation_flux_polar(hypot.(Vre,Vim), atan.(Vim,Vre), w, t)
excitation_current_sigmoid(inom, ψmax, ψ) =
    -inom .* log.(2.0 ./ (ψ ./ ψmax .+ 1.0) .- 1.0)

""
function sample_voltage_polar(voltage_harmonics, dv, vmin, vmax, dθ, θmin, θmax)
    S = reduce(vcat,[[vmin[ni]:dv[ni]:vmax[ni],θmin[ni]:dθ[ni]:θmax[ni]]
                      for (ni,nh) in enumerate(voltage_harmonics)])
    R = [1:length(s) for s in S]
    return S, R
end
function sample_voltage_rectangular(voltage_harmonics, dv, vmin, vmax)
    S = reduce(vcat,[[vmin[ni]:dv[ni]:vmax[ni],vmin[ni]:dv[ni]:vmax[ni]]
                      for (ni,nh) in enumerate(voltage_harmonics)])
    R = [1:length(s) for s in S]
    return S, R
end

xfmr_exc = xfmr

# HARMONICS
voltage_harmonics = xfmr_exc["voltage_harmonics"]
current_harmonics = xfmr_exc["current_harmonics"]

# TIME
N = xfmr_exc["N"]

# EXCITATION CURRENT
current_type = xfmr_exc["current_type"]

# EXCITATION FLUX
excitation_type = xfmr_exc["excitation_type"]
inom = xfmr_exc["inom"]
ψmax = xfmr_exc["ψmax"]

# EXCITATION VOLTAGE
voltage_type = xfmr_exc["voltage_type"]
dv   = xfmr_exc["dv"]
vmin = xfmr_exc["vmin"]
vmax = xfmr_exc["vmax"]
dθ   = xfmr_exc["dθ"]
θmin = xfmr_exc["θmin"]
θmax = xfmr_exc["θmax"]

# TIME
w = (2.0 * pi * freq) .* voltage_harmonics
t = 0.0:(1 / (freq * N * maximum(current_harmonics))):(5.0 / freq)

# DECOMPOSITION
fq = _SDC.Sinusoidal(freq .* current_harmonics)

# SAMPLE VOLTAGE
if voltage_type == :polar
    S, R = sample_voltage_polar(voltage_harmonics, dv, vmin, vmax, dθ, θmin, θmax)
elseif voltage_type == :rectangular
    S, R = sample_voltage_rectangular(voltage_harmonics, dv, vmin, vmax)
end
Ia = Dict(nh => zeros(R...) for nh in current_harmonics)
Ib = Dict(nh => zeros(R...) for nh in current_harmonics)

nr = [20,20,12,12]
sample = [S[ni][ns] for (ni,ns) in enumerate(nr)]

# if voltage_type == :polar
#     V, θ = sample[1:2:end], sample[2:2:end]
#     ψexc = excitation_flux_polar(V, θ, w, t)
# else voltage_type == :rectangular
Vre, Vim = sample[1:2:end], sample[2:2:end]
ψexc = excitation_flux_rectangular(Vre, Vim, w, t)
# end

# if excitation_type == :sigmoid
I_exc = excitation_current_sigmoid(inom, ψmax, ψexc)
# end

plot(t,ψexc)
plot(t,I_exc)

_SDC.decompose(t, I_exc, fq)

I, φ = fq.A[2:end], -fq.φ[2:end]


#!/usr/bin/julia
#
# Implementation of split step propagation method.
#
using Plots: Animation, plot, frame, gif;

export propagate, find_groundstate;

"""
    gen_opr(H::SeparableHamiltonian, par::SimParams)

Generate kinetic and potential operators in momentum and configuration space
(respectively) assuming that Strang splitting will be used (with the half
exponential factor on the potential term).

Assumes ħ = 1 if values of type `Unitful.Quantity` are not passed.
"""
function gen_opr(H::SeparableHamiltonian, par::SimParams)
  # Compensate for factor of -1 introduced by imaginary time propagation.
  if (imag(par.Δt) != 0) fac=-1
  else fac = 1;
  end

  # This version assumes ħ = 1 and that you'll resolve the unit conversions
  # yourself.
  return (exp.(ustrip.(-im*H.T*par.Δt*fac)), exp.(ustrip.(-0.5im*H.V*par.Δt*fac)));
end
function gen_opr(H::SeparableHamiltonian{T1}, par::SimParams{T2}) where
                 {N1<:Number,T1<:Unitful.Quantity{N1,Unitful.𝐌*Unitful.𝐋^2/Unitful.𝐓^2},
                  N2<:Number, T2<:Unitful.Quantity{N2,Unitful.𝐓}}
  # Compensate for factor of -1 introduced by imaginary time propagation.
  if (imag(par.Δt) != 0) fac=-1
  else fac = 1;
  end

  return (exp.(-im*H.T*par.Δt*fac/UnitfulAtomicHarmonic.ħ_u),
          exp.(-0.5im*H.V*par.Δt*fac/UnitfulAtomicHarmonic.ħ_u));
end

"""
    propagate(sim::SplitStep,
              xgrid::G=Grid(10, 0.01),
              kgrid::UniformGrid{M,1}=Grid(xgrid),
              par::SimParams=SimParams(0.05, 100, 1e-8),
              H::SeparableHamiltonian=SeparableHamiltonian(xgrid, kgrid, 0.5*xgrid.^2),
              ψ::WaveFunction{T,1,C,G}=WaveFunction(xgrid, exp.(-xgrid.^2/2));
              debug::Bool = true) where
        {T,C<:ConfigurationSpace,M<:MomentumSpace,G<:UniformGrid{C,1}}

Propagate the given wave function one time step using the split-step method.

# Arguments
- `sim::SplitStep`: the propagation method. Primarily used for dispatch, but can
  contain method specific simulation parameters.
- `xgrid::UniformGrid{ConfigurationSpace,1}`: grid of configuration space to evolve over.
- `kgrid::UniformGrid{M,1}`: 1D discretised momentum space.
- `par::SimParams`: Simulation parameters common to all simulation methods.
- `H::SeparableHamiltonian`: Hamiltonian to evolve system with respect to.
- `ψ::WaveFunction{T,1,C,G}`: 1D configuration space wave function to evolve.
- `debug::Bool`: If additional debug messages should be shown.
"""
function propagate(sim::SplitStep,
                   xgrid::G=Grid(10, 0.01),
                   kgrid::UniformGrid{M,1}=Grid(xgrid),
                   par::SimParams=SimParams(0.05, 100, 1e-8),
                   H::SeparableHamiltonian=SeparableHamiltonian(xgrid, kgrid, 0.5*xgrid.^2),
                   ψ::WaveFunction{T,1,C,G}=WaveFunction(xgrid, exp.(-xgrid.^2/2));
                   debug::Bool = false) where
        {T,C<:ConfigurationSpace,M<:MomentumSpace,G<:UniformGrid{C,1}}
  # Construct evolution operators from the passed Hamiltonian.
  (opT, opV) = gen_opr(H, par);

  # Work on a copy of the wave function passed with all units stripped.
  #ψ = WaveFunction{Complex{Float64},1,C,G}(ustrip.(ψi), ustrip(ψi.Δ));

  # Half-step in real space.
  ψ .*= opV;

  # FFT to momentum space.
  ϕ = transform(ψ, kgrid);

  # Full step in momentum space.
  ϕ .*= opT;

  # FFT back to position space.
  ψ = transform(ϕ, xgrid);

  # Half-step in real space.
  ψ .*= opV;

  # Renormalise wave function if running in imaginary time.
  if (ustrip(imag(par.Δt)) != 0)
    if debug println("Pre-normalisation norm was ", norm(ψ)); end
    normalise!(ψ);
  end

  # Return wave function of the same type as that passed.
  #return WaveFunction{T,1,C,G}(ψ, ψ.Δ);
  return ψ;
end

"""
    find_groundstate(sim::SplitStep,
                     xgrid::G,
                     kgrid::UniformGrid{M,1},
                     H::SeparableHamiltonian{C},
                     ψi::WaveFunction{T,1,C,G}=WaveFunction(xgrid, exp.(-xgrid.^2/2)),
                     par::SimParams=SimParams(0.1im, 250, 1e-8);
                     debug::Bool = false,
                     debugpath::String = "simrevamp/dev",
                     printstats::Bool = true) where
        {T,C<:ConfigurationSpace,M<:MomentumSpace,G<:UniformGrid{C,1}}

Determine the lowest energy eigenstate of the given potential.
"""
function find_groundstate(sim::SplitStep,
                         xgrid::G,
                         kgrid::UniformGrid{M,1},
                         H::SeparableHamiltonian{K,BT,B,C},
                         ψi::WaveFunction{T,1,C,G}=WaveFunction(xgrid, exp.(-xgrid.^2/2)),
                         par::SimParams=SimParams(0.1im, 250, 1e-8);
                         debug::Bool = false,
                         debugpath::String = "simrevamp/dev",
                         debugpltopts::SplitStepPlotOptions=SplitStepPlotOptions(true, 0.1),
                         printstats::Bool = true) where
        {T,K,BT,B,C<:ConfigurationSpace,M<:MomentumSpace,G<:UniformGrid{C,1}}

  if real(par.Δt) > 0*unit(par.Δt)
    throw(DomainError(par.Δt, "time step must be imaginary."));
  end

  # Create animation for convergence process if debugging.
  if debug
    anim = Animation();
    plt = SimPlot();
  end

  # Allow future expansion for determining multiple eigenstates.
  states = 1;

  # Define WaveBasis to store result in and make copy of passed WaveFunction for
  # operating on.
  Φ = WaveBasis{T}(xgrid, states);
  ψ = deepcopy(ψi);

  for i ∈ 1:states
    # Define vector to store wave function energy at each propagation step.
    E = zeros(typeof(1.0*unit(K)),par.steps+1);
    E[1] = energy(ψ, H);

    for j ∈ 1:par.steps
      if (debug && j%25 == 0) println("Iteration ", j); end

      ψ[:] = propagate(sim, xgrid, kgrid, par, H, ψ);
      E[j+1] = energy(ψ, H);

      if debug
        plot(plt, debugpltopts, xgrid, ψ, H,
            title=string("Step ", j, ", E = ", round(typeof(E[j+1]), E[j+1], sigdigits=6)));
        frame(anim);
      end

      # Stop propagation if energy has converged.
      if abs(E[j+1] - E[j]) < par.ϵ
        if printstats
          println("Converged to eigenstate ", i, " in ", j, " steps.");
          println("Energy found was ", round(typeof(E[j+1]), E[j+1], sigdigits=6));
        end
        break;
      end
    end
    normalise!(ψ);
    Φ[:,i] = ψ;
    #ψ[:] = orthog(ψi, Φ, i);
  end

  # Plot debug animation if requested.
  if debug
    check_path(debugpath);
    gif(anim, string(debugpath, gen_filename(par, xgrid, SplitStep(), debugpltopts), ".gif"), fps=16);
  end

  return Φ;
end


using CayleyDickson
using Base.Test: @test, @test_throws

@test begin
    l = Exo2Binion(Exo2Real(1,2,3,4))
    r = Exo2Binion(1,0,2,0,3,0,4)
    l == r
end

@test begin
    a = CayleyDickson.randomBigFloat()
    isreal(Exo2Binion(a))
end

@test begin
    a = 1
    !isreal(Exo2Binion(a, a, a, a))
end

@test begin
    z = random(Exo2Binion{Int})
    z == +(z)
end

@test begin
    a = CayleyDickson.randomBigFloat()
    real(Exo2Binion(a)) == a
end

@test begin
    io = IOBuffer()
    show(io, Exo2Binion(1.0, 2, 3, 4, 5, 6, 7, 8))
    l = String(take!(io))
    r = "[1: 1.0, i: 2.0, a: 3.0, ia: 4.0, b: 5.0, ib: 6.0, ab: 7.0, (ia)b: 8.0]"
    l == r
end

@test begin
    io = IOBuffer()
    show(io, Exo2Binion(1.0, 2, 3, 4, 5, 6, 7))
    l = String(take!(io))
    r = "[1: 1.0, i: 2.0, a: 3.0, ia: 4.0, b: 5.0, ib: 6.0, ab: 7.0, (ia)b: 0.0]"
    l == r
end

@test begin
    io = IOBuffer()
    show(io, Exo2Binion(1.0, 2, 3, 4, 5, 6))
    l = String(take!(io))
    r = "[1: 1.0, i: 2.0, a: 3.0, ia: 4.0, b: 5.0, ib: 6.0, ab: 0.0, (ia)b: 0.0]"
    l == r
end

@test begin
    io = IOBuffer()
    show(io, Exo2Binion(1.0, 2, 3, 4, 5))
    l = String(take!(io))
    r = "[1: 1.0, i: 2.0, a: 3.0, ia: 4.0, b: 5.0, ib: 0.0, ab: 0.0, (ia)b: 0.0]"
    l == r
end

@test begin
    io = IOBuffer()
    show(io, Exo2Binion(1.0, 2, 3, 4))
    l = String(take!(io))
    r = "[1: 1.0, i: 2.0, a: 3.0, ia: 4.0, b: 0.0, ib: 0.0, ab: 0.0, (ia)b: 0.0]"
    l == r
end

@test begin
    a = CayleyDickson.randomBigInt()
    b = CayleyDickson.randomBigFloat()
    c = CayleyDickson.randomBigInt()
    d = CayleyDickson.randomBigFloat()
    l = Exo2Binion(Exo1Binion(a, b, c, d))
    r = Exo2Binion(a, b, c, d)
    l == r
end

@test begin
    a = CayleyDickson.randomBigInt()
    b = CayleyDickson.randomBigFloat()
    l = Exo2Binion(Binion(a, b))
    r = Exo2Binion(a, b)
    l == r
end

@test begin
    a = CayleyDickson.randomBigInt()
    b = CayleyDickson.randomBigFloat()
    c = CayleyDickson.randomBigInt()
    l = Exo2Binion(a, b, c)
    r = Exo2Binion(a, b, c, 0)
    l == r
end

@test begin
    x = random(Exo2Binion{BigFloat})
    y = random(Exo2Binion{BigFloat})
    x != y
end

@test begin
    a = CayleyDickson.randomBigInt()
    z = random(Exo2Binion{BigInt})
    l = a + z
    r = z + a
    l == r
end

@test begin
    a = random(Exo1Binion{BigInt})
    z = random(Exo2Binion{BigInt})
    l = a + z
    r = z + a
    l == r
end

@test begin
    a = CayleyDickson.randomBigInt()
    z = random(Exo2Binion{BigInt})
    l = a - z
    r = -(z - a)
    l == r
end

@test begin
    a = random(Exo1Binion{BigInt})
    z = random(Exo2Binion{BigInt})
    l = a - z
    r = -(z - a)
    l == r
end

@test begin
    x = random(Exo2Binion{BigInt})
    y = random(Exo2Binion{BigInt})
    l = x + y
    r = y + x
    l == r
end

@test begin
    x = random(Exo2Binion{BigInt})
    y = random(Exo2Binion{BigInt})
    l = x * y
    r = y * x
    l != r
end

@test begin
    x = random(Exo2Binion{BigInt})
    y = random(Exo2Binion{BigInt})
    l = x - y
    r = -(y - x)
    l == r
end

@test begin
    x = random(Exo2Binion{BigInt})
    y = random(Exo2Binion{BigInt})
    z = random(Exo2Binion{BigInt})
    l = (x + y) + z
    r = x + (y + z)
    l == r
end

@test begin
    x = random(Exo2Binion{BigInt})
    y = random(Exo2Binion{BigInt})
    z = random(Exo2Binion{BigInt})
    l = (x * y) * z
    r = x * (y * z)
    l != r
end

@test begin
    x = random(Exo2Binion{BigInt})
    y = random(Exo2Binion{BigInt})
    l = alternatorL(x, y)
    r = zero(Exo2Binion{BigInt})
    l == r
end

@test begin
    x = random(Exo2Binion{BigInt})
    y = random(Exo2Binion{BigInt})
    l = alternatorR(x, y)
    r = zero(Exo2Binion{BigInt})
    l == r
end

@test begin
    x = random(Exo2Binion{BigInt})
    conj(conj(x)) == x
end

@test begin
    x = random(Exo2Binion{Rational{BigInt}})
    inv(inv(x)) == x
end

@test begin
    x = random(Exo2Binion{BigInt})
    y = random(Exo2Binion{BigInt})
    l = conj(x * y)
    r = conj(y) * conj(x)
    l == r
end

@test begin
    x = random(Exo2Binion{BigInt})
    y = random(Exo2Binion{BigInt})
    l = conj(x + y)
    r = conj(x) + conj(y)
    l == r
end

@test begin
    x = random(Exo2Binion{BigInt})
    y = random(Exo2Binion{BigInt})
    l = conj(x - y)
    r = conj(x) - conj(y)
    l == r
end

@test begin
    x = random(Exo2Binion{Rational{BigInt}})
    y = random(Exo2Binion{Rational{BigInt}})
    l = inv(x * y)
    r = inv(y) * inv(x)
    l == r
end

@test begin
    x = random(Exo2Binion{Rational{BigInt}})
    y = random(Exo2Binion{Rational{BigInt}})
    l = inv(x / y)
    r = y / x
    l == r
end

@test begin
    x = random(Exo2Binion{Rational{BigInt}})
    y = random(Exo2Binion{Rational{BigInt}})
    l = inv(x \ y)
    r = y \ x
    l == r
end

@test begin
    x = random(Exo2Binion{BigInt})
    y = random(Exo2Binion{BigInt})
    z = random(Exo2Binion{BigInt})
    l = x * (y + z)
    r = (x * y) + (x * z)
    l == r
end

@test begin
    a = CayleyDickson.randomBigInt()
    x = random(Exo2Binion{BigInt})
    y = random(Exo2Binion{BigInt})
    l = a * (x + y)
    r = (a * x) + (a * y)
    l == r
end

@test begin
    a = CayleyDickson.randomBigInt()
    x = random(Exo2Binion{BigInt})
    l = a * x
    r = x * a
    l == r
end

@test begin
    a = random(Exo1Binion{BigInt})
    x = random(Exo2Binion{BigInt})
    l = commutator(a, x)
    r = zero(Exo2Binion{BigInt})
    l != r
end

@test begin
    x = random(Exo2Binion{BigInt})
    y = random(Exo2Binion{BigInt})
    z = random(Exo2Binion{BigInt})
    l = (x + y) * z
    r = (x * z) + (y * z)
    l == r
end

@test begin
    x = random(Exo2Binion{BigInt})
    y = random(Exo2Binion{BigInt})
    z = random(Exo2Binion{BigInt})
    l = x * (y - z)
    r = (x * y) - (x * z)
    l == r
end

@test begin
    x = random(Exo2Binion{BigInt})
    y = random(Exo2Binion{BigInt})
    z = random(Exo2Binion{BigInt})
    l = (x - y) * z
    r = (x * z) - (y * z)
    l == r
end

@test begin
    x = random(Exo2Binion{BigInt})
    quadrance(x) > big(0) || quadrance(x) == big(0)
end

@test begin
    x = random(Exo2Binion{BigInt})
    y = random(Exo2Binion{BigInt})
    l = quadrance(x * y)
    r = quadrance(x) * quadrance(y)
    l == r
end

@test_throws ErrorException begin
    inv(Exo2Binion(0, 0.0))
end

@test_throws ErrorException begin
    random(Exo2Binion{Int}) / 0
end

@test_throws ErrorException begin
    0 \ random(Exo2Binion{Int})
end

@test begin
    a = CayleyDickson.randomBigRational()
    z = random(Exo2Binion{Rational{BigInt}})
    l = z / a
    r = a \ z
    l == r
end

@test begin
    a = CayleyDickson.randomBigRational()
    z = random(Exo2Binion{Rational{BigInt}})
    l = a / z
    r = z \ a
    l == r
end

@test begin
    z = random(Exo2Binion{BigInt})
    z == z + zero(Exo2Binion{BigInt})
end

@test begin
    z = random(Exo2Binion{BigInt})
    z == z * one(Exo2Binion{BigInt})
end

@test begin
    z = random(Exo2Binion{BigInt})
    z == z + zero(z)
end

@test begin
    z = random(Exo2Binion{BigInt})
    z == z * one(z)
end


@testset verbose=true "indexing" begin
    import SphericalFunctions: WignerHsize, WignerHrange, WignerHindex
    import SphericalFunctions: WignerDsize, WignerDrange, WignerDindex
    import SphericalFunctions: Ysize, Yrange, Yindex, deduce_limits
    import SphericalFunctions: theta_phi

    ell_max = 16
    ell_max_slow = ell_max ÷ 2

    @testset "WignerHrange" begin
        r1(mp_max, ell_max) = hcat([
            [ell, mp, m] for ell in 0:ell_max
            for mp in -min(ell, mp_max):min(ell, mp_max)
            for m in abs(mp):ell
        ]...)'
        for ell_max in 0:ell_max
            a = WignerHrange(ell_max)  # Implicitly, mp_max=ell_max
            b = r1(ell_max, ell_max)
            @test a == b
            for mp_max in 0:ell_max
                a = WignerHrange(ell_max, mp_max)
                b = r1(mp_max, ell_max)
                @test a == b
            end
        end
    end

    @testset "WignerHsize" begin
        for i in -5:-1
            @test WignerHsize(i) == 0
            for j in -5:5
                @test WignerHsize(i, j) == 0
            end
        end
        for ell_max in 0:ell_max
            @test WignerHsize(ell_max) == size(WignerHrange(ell_max, ell_max), 1)
            for mp_max in ell_max
                @test WignerHsize(ell_max, mp_max) == size(WignerHrange(ell_max, mp_max), 1)
            end
        end
    end

    @testset "WignerHindex" begin
        function fold_H_indices(ell, mp, m)
            if m < -mp
                if m < mp
                    return [ell, -mp, -m]
                else
                    return [ell, -m, -mp]
                end
            else
                if m < mp
                    return [ell, m, mp]
                else
                    return [ell, mp, m]
                end
            end
        end

        for ell_max_i in 0:ell_max_slow
            r = WignerHrange(ell_max_i)
            for ell in 0:ell_max_i
                for mp in -ell:ell
                    for m in -ell:ell
                        i = WignerHindex(ell, mp, m)
                        @test r[i, :] == fold_H_indices(ell, mp, m)
                    end
                end
            end
            for mp_max in 0:ell_max_i
                r = WignerHrange(ell_max_i, mp_max)
                for ell in 0:ell_max_i
                    for mp in -ell:ell
                        for m in -ell:ell
                            if abs(mp) > mp_max && abs(m) > mp_max
                                continue
                            end
                            i = WignerHindex(ell, mp, m, mp_max)
                            @test r[i, :] == fold_H_indices(ell, mp, m)
                        end
                    end
                end
            end
        end
    end

    @testset "WignerDrange" begin
        function r2(ell_min, mp_max, ell_max)
            a = [
                [ℓ, mp, m]
                for ℓ in ell_min:ell_max
                for mp in -min(ℓ, mp_max):min(ℓ, mp_max)
                for m in -ℓ:ℓ
            ]
            collect(transpose(reshape(collect(Iterators.flatten(a)), (3, length(a)))))
        end

        for ell_max in 0:ell_max_slow÷2
            let ell_min = 0
                a = WignerDrange(ell_max)  # Implicitly, mp_max=ell_max
                b = r2(ell_min, ell_max, ell_max)
                @test a == b
                for mp_max in 0:ell_max
                    a = WignerDrange(ell_max, mp_max)
                    b = r2(ell_min, mp_max, ell_max)
                    @test a == b
                end
            end
        end
    end

    @testset "WignerDsize" begin
        for ell_max in 0:ell_max
            let ell_min = 0
                for mp_max in 0:ell_max
                    a = WignerDsize(ell_max, mp_max)
                    b = size(WignerDrange(ell_max, mp_max), 1)
                    @test a == b
                end
            end
        end
        for ell_max in 0:ell_max
            let ell_min = 0
                for mp_max in 0:ell_max
                    a = WignerDsize(ell_max, mp_max)
                    b = size(WignerDrange(ell_max, mp_max), 1)
                    @test a == b
                end
            end
        end
        for ell_max in 0:ell_max
            let ell_min = 0
                for mp_max in [ell_max]
                    a = WignerDsize(ell_max, mp_max)
                    # a = WignerDsize_ellmin(ell_min, ell_max)
                    b = size(WignerDrange(ell_max, mp_max), 1)
                    @test a == b
                end
            end
        end
        for ell_max in 0:ell_max
            let ell_min = 0
                for mp_max in [ell_max]
                    a = WignerDsize(ell_max, mp_max)
                    # a = WignerDsize(ell_max)
                    b = size(WignerDrange(ell_max, mp_max), 1)
                    @test a == b
                end
            end
        end
    end

    @testset "WignerDindex" begin
        for ellmax in 0:ell_max_slow
            r = WignerDrange(ellmax)
            for ell in 0:ellmax
                for mp in -ell:ell
                    for m in -ell:ell
                        i = WignerDindex(ell, mp, m)
                        @test r[i, :] == [ell, mp, m]
                    end
                end
            end
            let ell_min = 0
                r = WignerDrange(ellmax)
                for ell in ell_min:ellmax
                    for mp in -ell:ell
                        for m in -ell:ell
                            i = WignerDindex(ell, mp, m)
                            @test r[i, :] == [ell, mp, m]
                        end
                    end
                end
                for mp_max in 0:ellmax
                    r = WignerDrange(ellmax, mp_max)
                    for ell in ell_min:ellmax
                        for mp in -min(ell, mp_max):min(ell, mp_max)
                            for m in -ell:ell
                                i = WignerDindex(ell, mp, m, mp_max)
                                @test r[i, :] == [ell, mp, m]
                            end
                        end
                    end
                end
            end
        end
    end

    @testset "Yrange" begin
        function r3(ell_min, ell_max)
            a = [
                [ℓ, m]
                for ℓ in ell_min:ell_max
                for m in -ℓ:ℓ
            ]
            collect(transpose(reshape(collect(Iterators.flatten(a)), (2, length(a)))))
        end
        for ell_max in 0:ell_max
            for ell_min in 0:ell_max
                a = Yrange(ell_min, ell_max)
                b = r3(ell_min, ell_max)
                @test a == b
            end
        end
    end

    @testset "Ysize" begin
        for ell_max in 0:ell_max
            for ell_min in 0:ell_max
                a = Ysize(ell_min, ell_max)
                b = size(Yrange(ell_min, ell_max))[1]
                @test a == b
            end
        end
    end

    @testset "deduce_limits" begin
        for ℓmax in 0:4096
            for ℓmin in 0:min(2, ℓmax)
                deduced = deduce_limits(Ysize(ℓmin, ℓmax), nothing)
                @test deduced == (ℓmin, ℓmax)
            end
        end
        for ℓmax in 0:ell_max
            for ℓmin in [0]
                deduced = deduce_limits(Ysize(ℓmin, ℓmax))
                @test deduced == (ℓmin, ℓmax)
            end
        end
        for ℓmax in 0:ell_max
            for ℓmin in 0:ℓmax
                deduced = deduce_limits(Ysize(ℓmin, ℓmax), ℓmin)
                @test deduced == (ℓmin, ℓmax)
            end
        end
        for ℓmax in 1:ell_max
            prev_size = Ysize(ℓmax-1)
            this_size = Ysize(ℓmax)
            if abs(prev_size - this_size) > 1
                mid_size = (prev_size + this_size) ÷ 2
                @test_throws ErrorException deduce_limits(mid_size)
            end
        end
    end

    @testset "Yindex" begin
        for ell_max in 0:ell_max
            let ell_min = 0
                r = Yrange(ell_min, ell_max)
                for ell in ell_min:ell_max
                    for m in -ell:ell
                        i = Yindex(ell, m)
                        @test r[i, :] == [ell, m]
                    end
                end
                s = Yrange(ell_max)
                @test array_equal(r, s)
            end
            for ell_min in 0:ell_max
                r = Yrange(ell_min, ell_max)
                for ell in ell_min:ell_max
                    for m in -ell:ell
                        i = Yindex(ell, m, ell_min)
                        @test r[i, :] == [ell, m]
                    end
                end
            end
        end
    end

    @testset "phi_theta <-> theta_phi" for T in [Float64, Float32, BigFloat]
        for nθ in 2:(2ell_max+1)
            for nϕ in 2:(2ell_max+1)
                θϕ = theta_phi(nθ, nϕ, T)
                ϕθ = permutedims(phi_theta(nϕ, nθ, T), [2, 1, 3])[:, :, 2:-1:1]
                @test array_equal(θϕ, ϕθ)
            end
        end
    end

end


using Random, IntervalLinearAlgebra, Plots

sizes = Iterators.flatten((2:9, 10:10:90, 100:100:1000)) |> collect

times = zeros(length(sizes), 4)

seed = 42
Random.seed!(seed)

random_interval_matrix(N) = Interval.(randn(N, N) .± abs.(randn(N, N)))

for (j, mult_mode) in enumerate((:fast, :rank1, :slow))

    set_multiplication_mode(mult_mode)
    Random.seed!(seed)
    for (i, N) in enumerate(sizes)
        A = random_interval_matrix(N)
        B = random_interval_matrix(N)
        t = @benchmark $A * $B
        times[i, j] = minimum(t.times)
        @show N
    end

end

# normal multiplication with accurate rounding mode
set_multiplication_mode(:slow)
setrounding(Interval, :accurate)
Random.seed!(seed)
for (i, N) in enumerate(sizes)
    A = random_interval_matrix(N)
    B = random_interval_matrix(N)
    t = @benchmark $A * $B
    times[i, 4] = minimum(t.times)
    @show N
end

# back to default settings
setrounding(Interval, :tight)
set_multiplication_mode(:fast)

# floating point multiplication
times_float = zeros(size(sizes))
for (i, N) in enumerate(sizes)
    A = randn(N, N)
    B = randn(N, N)
    t = @benchmark $A * $B
    times_float[i] = minimum(t.times)
    @show N
end

# plotting
labels = ["fast" "rank1" "slow-tight" "slow-accurate"]
plot(sizes, times/1e6; label=labels, axis=:log, m=:auto, legend=:topleft)
plot!(sizes, times_float/1e6, label="Float64", m=:auto)
xlabel!("size")
ylabel!("time [ms]")
yticks!(10.0 .^(-3:5))

ratios = times ./ times_float
labels = ["fast" "rank1" "slow-tight" "slow-accurate"]
plot(sizes, ratios; label=labels, axis=:log, m=:auto, legend=:topleft)
xlabel!("size")
ylabel!("ratio")
yticks!(10.0 .^(-3:5))

## overestimate benchmarks

radii = 10.0 .^ (-15:15)
sizes = [10, 50, 100, 200]
Nsamples = 10

min_overestimate = zeros(length(radii), length(sizes))
max_overestimate = zeros(length(radii), length(sizes))
mean_overestimate = zeros(length(radii), length(sizes))

# function random_interval_matrix(N, r)
#     Ac = randn(N, N)
#     return Ac .± (r * abs.(Ac))
# end

random_interval_matrix(N, r) = randn(N, N) .± r*abs.(randn(N, N))
Random.seed!(seed)
for (i, N) in enumerate(sizes)
    for (j, r) in enumerate(radii)

        tmp_statistics = zeros(3)
        for _ in 1:Nsamples
            A = random_interval_matrix(N, r)
            B = random_interval_matrix(N, r)

            set_multiplication_mode(:fast)
            Cfast = A * B
            set_multiplication_mode(:slow)
            Cslow = A * B

            ratios = (diam.(Cfast)./diam.(Cslow) .- 1) * 100
            tmp_statistics .+= (minimum(ratios), mean(ratios), maximum(ratios))
        end

        min_overestimate[j, i] = tmp_statistics[1]/Nsamples
        mean_overestimate[j, i] = tmp_statistics[2]/Nsamples
        max_overestimate[j, i] = tmp_statistics[3]/Nsamples
        println("N=$(i)/$(length(sizes)) r=$(j)/$(length(radii))")
    end
end

for (i, N) in enumerate(sizes)
    plot(radii, max_overestimate[:, i], xaxis=:log, m=:o, label="max")
    plot!(radii, min_overestimate[:, i], m=:o, label="min")
    plot!(radii, mean_overestimate[:, i], m=:o, label="mean")
    xlabel!("radius")
    ylabel!("%-overestimate")
    title!("matrix size N=$N")
    savefig("accurate_overestimate_N$(N)_random_radius.pdf")
end


using Vega

scatterplot(x = rand(100), y = rand(100))

using Vega, Distributions

d1 = MultivariateNormal([0.0, 0.0], [1.0 0.9; 0.9 1.0])
d2 = MultivariateNormal([10.0, 10.0], [4.0 0.5; 0.5 4.0])
points = vcat(rand(d1, 500)', rand(d2, 500)')

x = points[:, 1]
y = points[:, 2]
group = vcat(ones(Int, 500), ones(Int, 500) + 1)

scatterplot(x = x, y = y, group = group)

x = [0,1,2,3,4,5,6,7,8,9,0,1,2,3,4,5,6,7,8,9]
y = [28, 43, 81, 19, 52, 24, 87, 17, 68, 49, 55, 91, 53, 87, 48, 49, 66, 27, 16, 15]
g = [0,0,0,0,0,0,0,0,0,0,1,1,1,1,1,1,1,1,1,1]

a = areaplot(x = x, y = y, group = g, stacked = true)

fruit = ["peaches", "plums", "blueberries", "strawberries", "bananas"];
bushels = [100, 32, 180, 46, 21];
piechart(x = fruit, y = bushels, holesize = 125)

x = Array(Int, 900)
y = Array(Int, 900)
color = Array(Float64, 900)
tmp = 0
for counter in 1:30
    for counter2 in 1:30
        tmp += 1
        x[tmp] = counter
        y[tmp] = counter2
        color[tmp] = rand()
    end
end
hm = heatmap(x = x, y = y, color = color)


using Test

using AlgebraicPetri
using AlgebraicPetri.Epidemiology
using AlgebraicPetri.BilayerNetworks
using LabelledArrays

using Catlab
using Catlab.CategoricalAlgebra
using Catlab.Graphics
import Catlab.CategoricalAlgebra: migrate!
using Catlab.WiringDiagrams
using Catlab.Programs.RelationalPrograms
display_uwd(ex) = to_graphviz(ex, box_labels=:name, junction_labels=:variable, edge_attrs=Dict(:len=>".75"));

# #### SIR Model:

# define model
sir = @relation (s,i,r) begin
    infection(s,i)
    recovery(i,r)
end
display_uwd(sir)
#-
#

# Extract the Petri network form of the model.

psir = apex(oapply_epi(sir))
psir
Graph(psir)

# Convert the Petri network into a bilayer network and draw it.
# This model uses a computation graph to express the computation of the vector field of the Petri net.

bnsir = BilayerNetwork()
lab_bnsir = LabelledBilayerNetwork()
migrate!(bnsir, psir)
migrate!(lab_bnsir, psir)
bnsir
@test typeof(to_graphviz(bnsir)) == Graph
@test typeof(to_graphviz(lab_bnsir)) == Graph

bnsir_test = @acset BilayerNetwork begin
    Qin = 3
    Qout = 3
    Win = 3
    Wa = 3
    Wn = 3
    Box = 2
    arg = [1,2,2]
    call = [1,1,2]
    efflux = [1,1,2]
    effusion = [1,2,2]
    influx = [1,1,2]
    infusion = [2,2,3]
end

@test bnsir == bnsir_test

######################
# Round Trip testing #
######################
function roundtrip(pn::AbstractPetriNet, bn::AbstractBilayerNetwork)
    roundtrippetri = PetriNet()
    migrate!(roundtrippetri, bn)
    pn_structure = PetriNet()
    copy_parts!(pn_structure, pn)
    return roundtrippetri, pn_structure
end

function roundtrip(pn::AbstractLabelledPetriNet, bn::AbstractLabelledBilayerNetwork)
    roundtrippetri = LabelledPetriNet()
    migrate!(roundtrippetri, bn)
    return roundtrippetri, pn
end

function test_roundtrip(pn::AbstractPetriNet, bn::AbstractBilayerNetwork)
    roundtrippetri, pn_structure = roundtrip(pn, bn)
    @test roundtrippetri == pn_structure
end

bnsir = BilayerNetwork()
migrate!(bnsir, psir)
test_roundtrip(psir, bnsir)

seir = @relation (s,e,i,r) begin
    exposure(s,i,e)
    illness(e,i)
    recovery(i,r)
end

pseir = apex(oapply_epi(seir))
bnseir = BilayerNetwork()
migrate!(bnseir, pseir)
test_roundtrip(pseir, bnseir)

lbnseir = LabelledBilayerNetwork()
migrate!(lbnseir, pseir)
test_roundtrip(pseir, lbnseir)

########################
# Compile and Evaluate #
########################

du = LVector(S=0.0,I=0.0,R=0.0)
u = LVector(S=10.0,I=1.0,R=0.0)
params = LVector(inf=0.1,rec=0.05)

bn_du =  AlgebraicPetri.BilayerNetworks.evaluate(lab_bnsir, u; inf=0.1,rec=0.05)
vectorfield(psir)(du, u, params, 0)

@test all(abs.(bn_du .- du[[:S,:I,:R]]) .< 1e-9)

eval(AlgebraicPetri.BilayerNetworks.compile(lab_bnsir, :du, :ϕ, :u; inf=0.1, rec=0.05))
comp_du = [0.0,0.0,0.0]
comp_ϕ = [0.0,0.0,0.0]

f!(comp_du, comp_ϕ, u, 0)

@test all(abs.(comp_du .- du[[:S,:I,:R]]) .< 1e-9)

p, params = params, nothing
eval(AlgebraicPetri.BilayerNetworks.compile(lab_bnsir, :du, :ϕ, :u, :p))
comp_du = [0.0,0.0,0.0]
comp_ϕ = [0.0,0.0,0.0]

f!(comp_du, comp_ϕ, u, p, 0)

@test all(abs.(comp_du .- du[[:S,:I,:R]]) .< 1e-9)

# Test for arbitrary variable names
eval(AlgebraicPetri.BilayerNetworks.compile(lab_bnsir, :a, :b, :c, :d))
comp_du = [0.0,0.0,0.0]
comp_ϕ = [0.0,0.0,0.0]

f!(comp_du, comp_ϕ, u, p, 0)

@test all(abs.(comp_du .- du[[:S,:I,:R]]) .< 1e-9)

##############
# Edge cases #
##############

function test_sir_equivalence(bn, pn)
  du = LVector(S=0.0,I=0.0,R=0.0)
  u = LVector(S=10.0,I=1.0,R=0.0)
  params = LVector(inf=0.1,rec=0.05)

  bn_du =  AlgebraicPetri.BilayerNetworks.evaluate(bn, u; inf=0.1,rec=0.05)
  vectorfield(pn)(du, u, params, 0)

  @test all(abs.(bn_du .- du[[:S,:I,:R]]) .< 1e-9)
end

# Ensure that the migration function properly balances the BLN
bnsir_edge = @acset LabelledBilayerNetwork begin
    Qin = 3
    Qout = 3
    Win = 3
    Wa = 2
    Wn = 2
    Box = 2
    arg = [1,2,2]
    call = [1,1,2]
    efflux = [1,2]
    effusion = [1,2]
    influx = [1,2]
    infusion = [2,3]
    parameter=[:inf, :rec]
    variable=[:S, :I, :R]
    tanvar=[:S, :I, :R]
end

test_sir_equivalence(bnsir_edge, psir)

edge_sir = LabelledPetriNet()
migrate!(edge_sir, bnsir_edge)
test_sir_equivalence(bnsir_edge, edge_sir)

bnsir_edge_add = @acset LabelledBilayerNetwork begin
    Qin = 3
    Qout = 3
    Win = 3
    Wa = 4
    Wn = 4
    Box = 2
    arg = [1,2,2]
    call = [1,1,2]
    efflux = [1,1,1,2]
    effusion = [1,2,2,2]
    influx = [1,1,1,2]
    infusion = [2,2,2,3]
    parameter=[:inf, :rec]
    variable=[:S, :I, :R]
    tanvar=[:S, :I, :R]
end

test_sir_equivalence(bnsir_edge_add, psir)

edge_sir = LabelledPetriNet()
migrate!(edge_sir, bnsir_edge_add)
test_sir_equivalence(bnsir_edge, edge_sir)

bnsir_migrate_break = @acset LabelledBilayerNetwork begin
    Qin = 3
    Qout = 3
    Win = 2
    Wa = 2
    Wn = 2
    Box = 2
    arg = [1,2]
    call = [1,1]
    efflux = [1,2]
    effusion = [1,2]
    influx = [1,2]
    infusion = [2,3]
    parameter=[:inf, :rec]
    variable=[:S, :I, :R]
    tanvar=[:S, :I, :R]
end
bn_error = "Mass action does not allow species $(bnsir_migrate_break[2, :variable]) to be "*
           "removed without contributing to the rate."
@test_throws ErrorException(bn_error) migrate!(edge_sir, bnsir_migrate_break)


using Plots; pyplot();
using Distributions
using LinearAlgebra
using KernelFunctions
# Translational invariants kernels

default(lw=3.0,titlefontsize=28,tickfontsize=18)

x₀ = 0.0; l = 0.1
n_grid = 101
fill(x₀,n_grid,1)
xrange = reshape(collect(range(-3,3,length=n_grid)),:,1)

k = transform(SqExponentialKernel(),1.0)
K1 = kernelmatrix(k,xrange,obsdim=1)
p = heatmap(K1,yflip=true,colorbar=false,framestyle=:none,background_color=RGBA(0.0,0.0,0.0,0.0))
savefig(joinpath(@__DIR__,"src","assets","heatmap_sqexp.png"))

k = @kernel Matern32Kernel FunctionTransform(x->(sin.(x)).^2)
K2 = kernelmatrix(k,xrange,obsdim=1)
p = heatmap(K2,yflip=true,colorbar=false,framestyle=:none,background_color=RGBA(0.0,0.0,0.0,0.0))
savefig(joinpath(@__DIR__,"src","assets","heatmap_matern.png"))

k = transform(PolynomialKernel(c=0.0,d=2.0), LinearTransform(randn(3,1)))
K3 = kernelmatrix(k,xrange,obsdim=1)
p = heatmap(K3,yflip=true,colorbar=false,framestyle=:none,background_color=RGBA(0.0,0.0,0.0,0.0))
savefig(joinpath(@__DIR__,"src","assets","heatmap_poly.png"))

k = 0.5*SqExponentialKernel()*transform(LinearKernel(),0.5) + 0.4*(@kernel Matern32Kernel() FunctionTransform(x->sin.(x)))
K4 = kernelmatrix(k,xrange,obsdim=1)
p = heatmap(K4,yflip=true,colorbar=false,framestyle=:none,background_color=RGBA(0.0,0.0,0.0,0.0))
savefig(joinpath(@__DIR__,"src","assets","heatmap_prodsum.png"))

plot(heatmap.([K1,K2,K3,K4],yflip=true,colorbar=false)...,layout=(2,2))
savefig(joinpath(@__DIR__,"src","assets","heatmap_combination.png"))

##

for k in [SqExponentialKernel,ExponentialKernel]
    K = kernelmatrix(k(),xrange,obsdim=1)
    v = rand(MvNormal(K+1e-7I))
    plot(xrange,v,lab="",title="f(x)",framestyle=:none) |> display
    savefig(joinpath(@__DIR__,"src","assets","GP_sample_$(k).png"))
    plot(xrange,kernel.(k(),x₀,xrange),lab="",ylims=(0,1.1),title="k(0,x)") |> display
    savefig(joinpath(@__DIR__,"src","assets","kappa_function_$(k).png"))
end

for k in [GammaExponentialKernel(1.0,1.5)]
    sparse =1
    while !isposdef(kernelmatrix(k,xrange*sparse,obsdim=1) + 1e-5I); sparse += 1; end
    v = rand(MvNormal(kernelmatrix(k,xrange*sparse,obsdim=1)+1e-7I))
    plot(xrange,v,lab="",title="f(x)",framestyle=:none) |> display
    savefig(joinpath(@__DIR__,"src","assets","GP_sample_GammaExponentialKernel.png"))
    plot(xrange,kernel.(k,x₀,xrange),lab="",ylims=(0,1.1),title="k(0,x)") |> display
    savefig(joinpath(@__DIR__,"src","assets","kappa_function_GammaExponentialKernel.png"))
end

for k in [MaternKernel,Matern32Kernel,Matern52Kernel]
    K = kernelmatrix(k(),xrange,obsdim=1)
    v = rand(MvNormal(K+1e-7I))
    plot(xrange,v,lab="",title="f(x)",framestyle=:none) |> display
    savefig(joinpath(@__DIR__,"src","assets","GP_sample_$(k).png"))
    plot(xrange,kernel.(k(),x₀,xrange),lab="",ylims=(0,1.1),title="k(0,x)") |> display
    savefig(joinpath(@__DIR__,"src","assets","kappa_function_$(k).png"))
end

for k in [RationalQuadraticKernel]
    K = kernelmatrix(k(),xrange,obsdim=1)
    v = rand(MvNormal(K+1e-7I))
    plot(xrange,v,lab="",title="f(x)",framestyle=:none) |> display
    savefig(joinpath(@__DIR__,"src","assets","GP_sample_$(k).png"))
    plot(xrange,kernel.(k(),x₀,xrange),lab="",ylims=(0,1.1),title="k(0,x)") |> display
    savefig(joinpath(@__DIR__,"src","assets","kappa_function_$(k).png"))
end


export  r_scale,
        rm_logistic,
        rm_laguerre,
        rm_logisticsum,
        rm_hermite,
        rm_hermite_prob,
        rm_laguerre,
        rm_jacobi,
        rm_jacobi01,
        rm_hahn,
        rm_meixner_pollaczek,
        rm_legendre,
        rm_legendre01,
        rm_chebyshev1,
        rm_compute
"""
    r_scale(c::Float64,β::Vector{Float64},α::Vector{Float64})
Given the recursion coefficients `(α,β)` for a system of orthogonal polynomials that are orthogonal with respect to some positive weight ``m(t)``,
this function returns the recursion coefficients `(α_,β_)` for the scaled measure ``c m(t)`` for some positive ``c``.
"""
function r_scale(c::Float64,a::Vector{Float64},b::Vector{Float64})
    @assert c>0 "Measure can only be scaled by positive number (provided c=$c)"
    return a, [c*b[1]; b[2:end]]
end

"""
    rm_compute(weight::Function,lb::Float64,ub::Float64,Npoly::Int64=4,Nquad::Int64=10;quadrature::Function=clenshaw_curtis)
Given a positive `weight` function with domain `(lb,ub)`, i.e. a function ``w: [lb, ub ] \\rightarrow \\mathbb{R}_{\\geq 0}``,
this function creates `Npoly` recursion coefficients `(α,β)`.

The keyword `quadrature` specifies what quadrature rule is being used.
"""
function rm_compute(weight::Function,lb::Float64,ub::Float64,Npoly::Int64=4,Nquad::Int64=10;quadrature::Function=clenshaw_curtis,discretization::Function=stieltjes)
    @assert Npoly <= Nquad
    Npoly == 0 && return Array{Float64,1}(undef,0), Array{Float64,1}(undef,0)
    n,w= quadgp(weight,lb,ub,Nquad;quadrature=quadrature)
    a,b = discretization(Npoly,n,w)
    # display((b,maximum(b)))
    # the following assertions are taken from Theorem 1.28 from
    # Gautschi, W. "Orthogonal Polynomials: Computation and Approximation"
    @assert lb<minimum(a) && maximum(a)< ub "Not all recurrence coefficients α ($(minimum(a)),$(maximum(b))) are elements of the support s=($lb,$ub)."
    if length(b)>=2
        @assert 0<minimum(b) && maximum(b[2:end])<=max(lb^2,ub^2) "Not all recurrence coefficients β are elements of the interval (0, max(lb^2,ub^2)) = (0,$(max(lb^2,ub^2))), namely ($(minimum(b)), $(maximum(b)))."
    end
    return a,b
end

rm_compute(m::Measure,Npoly::Int64=4,Nquad::Int64=10;quadrature::Function=clenshaw_curtis,discretization::Function=stieltjes) = rm_compute(m.w,m.dom[1],m.dom[2],Npoly,Nquad,quadrature=quadrature,discretization=discretization)

##
function rm_logisticsum(n::Int,p1::Vector{Float64},p2::Vector{Float64};Mmax::Int=100,eps0::Real=1e-9)
    M0 = n
    Mcap = 0
    Mi = M0
    α, β = zeros(n), zeros(n)
    for i=1:Mmax
        nai, wai = quadrature_logistic(Mi,p1[1],p1[2],0.5) # 0.5 is needed to make the weight function a density, i.e. integral is 1.
        nbi, wbi = quadrature_logistic(Mi,p2[1],p2[2],0.5) # 0.5 is needed to make the weight function a density, i.e. integral is 1.
        ni, wi   = [nai; nbi], [wai; wbi]
        α_, β_ = stieltjes(n,ni,wi)
        err = maximum( abs(β_[k]-β[k])/β_[k] for k=1:n )
        # display("err = $err")
        if err<=eps0
            Mcap = i-1
            return α_, β_
        else
            α, β, = α_, β_
            i==1 ? Mi = M0+1 : Mi = Int(2^(floor((i-1)/5))*n)
        end
    end
    warn("Algorithm did not terminate after $Mmax iterations.")
end

"""
    rm_logistic(N::Int)
Creates `N` recurrence coefficients for monic polynomials that are orthogonal
on ``(-\\infty,\\infty)`` relative to ``w(t) = \\frac{\\mathrm{e}^{-t}}{(1 - \\mathrm{e}^{-t})^2}``
"""
function rm_logistic(N::Int)
    @assert N >= 0 "parameter(s) out of range."
    N == 0 && return Array{Float64,1}(undef,0), Array{Float64,1}(undef,0)
    zeros(N), pushfirst!(map(k->k^4*pi^2/(4*k^2-1),Base.OneTo(N-1)),1.)
end

"""
    rm_hermite(N::Int,mu::Float64)
    rm_hermite(N::Int)
Creates `N` recurrence coefficients for monic generalized Hermite polynomials
that are orthogonal on ``(-\\infty,\\infty)`` relative to ``w(t) = |t|^{2 \\mu} \\mathrm{e}^{-t^2}``

The call `rm_hermite(N)` is the same as `rm_hermite(N,0)`.
"""
function rm_hermite(N::Int,mu::Float64)
    @assert N >= 0 && mu > -0.5 "parameter(s) out of range."
    N == 0 && return Array{Float64,1}(undef,0), Array{Float64,1}(undef,0)
    m0 = mu != 0. ? gamma(mu + 0.5) : sqrt(π)
    N == 1 && return [0.], [m0]
    return zeros(N), pushfirst!(map(x->isodd(x) ? 0.5*x+mu : 0.5*x,1:N-1), m0)
end
rm_hermite(N::Int) = rm_hermite(N,0.)

"""
    rm_hermite_prob(N::Int)
Creates `N` recurrence coefficients for monic probabilists' Hermite polynomials
that are orthogonal on ``(-\\infty,\\infty)`` relative to ``w(t) = \\mathrm{e}^{-0.5t^2}``
"""
function rm_hermite_prob(N::Int)
    @assert N >= 0 "parameter(s) out of range."
    N == 0 && return Array{Float64,1}(undef,0), Array{Float64,1}(undef,0)
    # return zeros(N), [sqrt(2*pi); collect(1.:N-1) ]
    return zeros(N), pushfirst!( collect(1.:N-1), sqrt(2*pi))
end

"""
    rm_laguerre(N::Int,a::Float64)
    rm_laguerre(N::Int)
Creates `N` recurrence coefficients for monic generalized Laguerre polynomials
that are orthogonal on ``(0,\\infty)`` relative to ``w(t) = t^a \\mathrm{e}^{-t}``.

The call `rm_laguerre(N)` is the same as `rm_laguerre(N,0)`.
"""
function rm_laguerre(N::Int,a::Float64)
    @assert N >= 0 && a > -1. "parameter(s) out of range"
    N == 0 && return Array{Float64,1}(undef,0), Array{Float64,1}(undef,0)
    N==1 && return [a+1.], [gamma(a+1)]
    n = 1.:N-1
    return pushfirst!(map(x->2x+a+1.,n)::Vector{Float64},a+1.), pushfirst!(map(x->x^2+a*x,n)::Vector{Float64},gamma(a+1))
end
function rm_laguerre(N::Int)
    rm_laguerre(N,0.)
end

"""
    rm_jacobi(N::Int,a::Float64,b::Float64)
    rm_jacobi(N::Int,a::Float64)
    rm_jacobi(N::Int)
Creates `N` recurrence coefficients for monic Jacobi polynomials
that are orthogonal on ``(-1,1)`` relative to ``w(t) = (1-t)^a (1+t)^b``.

The call `rm_jacobi(N,a)` is the same as `rm_jacobi(N,a,a)` and `rm_jacobi(N)` the same as
`rm_jacobi(N,0,0)`.
"""
function rm_jacobi(N::Int,a::Float64,b::Float64)
    @assert N >= 0 && a >- 1. && b >- 1. "parameter(s) out of range"
    N == 0 && return Array{Float64,1}(undef,0), Array{Float64,1}(undef,0)
    nu = (b - a) / ( a+ b + 2.);
    mu::Float64 = a + b + 2. <= 128. ? 2^(a+b+1)*((gamma(a+1)*gamma(b+1))/gamma(a+b+2)) : exp((a+b+1)*log(2)+((log(gamma(a+1))+log(gamma(b+1)))-log(gamma(a+b+2))))
    N == 1 && return [nu], [mu]
    n = 1:N-1
    nab = map(x->2x+a+b,n)
    A = pushfirst!(map(x->(b^2-a^2)/(x*(x+2)),nab),nu)
    B1 = map(x->4*(x+a)*(x+b)*x*(x+a+b),n[2:N-1])
    B2 = map(x->(x^2.)*(x+1.)*(x-1.),nab[2:N-1])
    B3 = 4. * (a + 1) * (b + 1) / ((a + b + 2)^2 * (a + b + 3));
    A, pushfirst!(pushfirst!(B1./B2,B3),mu)::Vector{Float64}
end

rm_jacobi(N::Int,a::Float64) = rm_jacobi(N,a,a)
rm_jacobi(N::Int) = rm_jacobi(N,0.,0.)
"""
    rm_jacobi01(N::Int,a::Float64,b::Float64)
    rm_jacobi01(N::Int,a::Float64)
    rm_jacobi01(N::Int)

Creates `N` recurrence coefficients for monic Jacobi polynomials
that are orthogonal on ``(0,1)`` relative to ``w(t) = (1-t)^a t^b``.

The call `rm_jacobi01(N,a)` is the same as `rm_jacobi01(N,a,a)` and `rm_jacobi01(N)` the same as
`rm_jacobi01(N,0,0)`.
"""
function rm_jacobi01(N::Int,a::Float64,b::Float64)
    @assert N>=0 && a>-1. && b>-1. "parameter(s) out of range"
    N == 0 && return Array{Float64,1}(undef,0), Array{Float64,1}(undef,0)
    c, d = rm_jacobi(N,a,b)
    map(x->(1+x)/2,c), pushfirst!(0.25*d[2:N], d[1]/2^(a+b+1.))
end
rm_jacobi01(N::Int,a::Float64) = rm_jacobi01(N,a,a)
rm_jacobi01(N::Int) = rm_jacobi01(N,0.,0.)

"""
    rm_legendre(N::Int)

Creates `N` recurrence coefficients for monic Legendre polynomials
that are orthogonal on ``(-1,1)`` relative to ``w(t) = 1``.
"""
rm_legendre(N::Int) = rm_jacobi(N)

"""
    rm_legendre01(N::Int)

Creates `N` recurrence coefficients for monic Legendre polynomials
that are orthogonal on ``(0,1)`` relative to ``w(t) = 1``.
"""
rm_legendre01(N::Int) = rm_jacobi01(N)

function rm_chebyshev1(N::Int64)
    @assert N>=0 "N has to be non-negative"
    α = zeros(Float64,N)
    if N == 1
        return α, [pi]
    elseif N == 2
        return α, [pi; 0.5]
    else
        return α, pushfirst!(pushfirst!(0.25*ones(N-2),0.5),pi)
    end
end
#     rm_hahn(N::Int,a::Float64,b::Float64)
#     rm_hahn(N::Int,a::Float64)
#     rm_hahn(N::Int)
#
#  Creates `N` recurrence coefficients for monic Hahn polynomials.
#
#     ab=rm_hahn(N,a,b) generates the N+1 recurrence coefficients
#     for the monic Hahn polynomials with parameters a and b. These
#     are orthogonal on the discrete set of N+1 points 0,1,2,...,N
#     with weights (a+k choose k)(b+N-k choose N-k), k=0,1,2,...,N.
#     The N+1 alpha-coefficients are stored in the first column,
#     the N+1 beta-coefficients in the second column, of the
#     (N+1)x2 array ab. The call ab=rm_hahn(N,a) is the same as ab=
#     rm_hahn(N,a,a), and ab=HAHN(N) the same as ab=rm_hahn(N,0,0),
#     which produces the recurrence coefficients of the discrete
#     Chebyshev polynomials for the points 0,1,2,...,N.
# """
# function rm_hahn(N::Int,a::Float64,b::Float64)
#     N-=1
#     @assert N>=0 && a>=-1. && b>=-1. "parameter(s) out of range"
#     N==0 ? (return Array{Float64,1}(undef,0), Array{Float64,1}(undef,0)) : ()
#     ab = zeros(N+1,2)
#     n=1:N; ab[1,2]=prod(1 .+ (a+b+1.)./n);
#     if a+b == 0.
#       n=(1:N+1)';
#       ab[n,1]=((2*n .+ (a+b-1.))*N+(b-a)*n .+ a)./(2*(2*n .- 1));
#       n=(1:N)';
#       ab[n.+1,2]=.25*((N+1)^2)*(1 .+ a./n).*(1 .+ b./n).*(1 .-(n./(N+1)).^2)./(4 .- (1 ./n).^2);
#     elseif a+b+1. == 0.
#       n=(1:N+1)';
#       ab[n,1]=((2*(n .- 1).^2 .+ b)*N .+ (2*b+1)*(n .- 1).^2)./(4*(n .- 1).^2 .- 1);
#       n=(1:N)';
#       ab[n.+1,2]=(.25*((N+1)^2)).*(1 .+ a*(1 ./n)).*(1 .+ b./n).*(1 .- n./(N+1)).*(1 .+ (n .- 1)./(N+1))./(4 .- (1 ./n).^2);
#     else
#       n=(1:N+1)';
#       ab[n,1]=((n .+ (a+b)).*(n .+ a).*((N+1) .- n)./(2*n .+ (a+b))+(n .- 1).*(n .+ (b-1)).*(N .+ n .+ (a+b))./(2*n .+ (a+b-2)))./(2*n .+ (a+b-1));
#       n=(1:N)';
#       ab[n.+1,2]=((N+1)^2)*(1 .+ a./n).*(1 .+ b./n).*(1 .+ (a+b)./n).*(1 .- n./(N+1)).*(1 .+ (n .+ (a+b))./(N .+ 1))./(((2 .+ (a+b)./n).^2).*((2 .+ (a+b)./n).^2 .- (1 ./n).^2));
#     end
#     return ab[:,1], ab[:,2]
# end
# rm_hahn(N::Int,a::Float64) = rm_hahn(N,a,a)
# rm_hahn(N::Int) = rm_hahn(N,0.,0.)

"""
    rm_meixner_pollaczek(N::Int,lambda::Float64,phi::Float64)
    rm_meixner_pollaczek(N::Int,lambda::Float64)

 Creates `N` recurrence coefficients for monic
 Meixner-Pollaczek polynomials with parameters λ and ϕ. These are orthogonal on
 ``[-\\infty,\\infty]`` relative to the weight function ``w(t)=(2 \\pi)^{-1} \\exp{(2 \\phi-\\pi)t} |\\Gamma(\\lambda+ i t)|^2``.

 The call `rm_meixner_pollaczek(n,lambda)` is the same as `rm_meixner_pollaczek(n,lambda,pi/2)`.
"""
function rm_meixner_pollaczek(N::Int,lambda::Float64,phi::Float64)
    @assert N>=0 && lambda>0. && phi>0. "parameter(s) out of range"
    N == 0 && return Array{Float64,1}(undef,0), Array{Float64,1}(undef,0)
    n=1:N;
    sinphi=sin(phi); lam2=2*lambda;
    ab=zeros(Float64,N,2)
    if sinphi==1
      ab[:,1]=zeros(N);
    else
        ab[:,1]= -(n .+ (lambda-1))/tan(phi);
    end
    ab[:,2]= (n .- 1).*(n .+ (lam2-2))/(4*sinphi^2);
    ab[1,2]=gamma(lam2)/(2*sinphi)^lam2;
    return ab[:,1], ab[:,2]
end
rm_meixner_pollaczek(N::Int,lambda::Float64) = rm_meixner_pollaczek(N,lambda,pi/2)


include("../examples/tent.jl")
using JLD
σ_o = 0.1
function p_y_g_x(a)
		return log(1/sqrt(2π)/σ_o) -0.5*a*a/σ_o/σ_o
end
function obs(x)
    return x
end
function sir(y,x,s,N_thr=10)
	K = size(y)[1]
	N_p = size(x)[1]
    w = ones(N_p)./N_p
	new_pts = similar(x)
	for k = 1:K
    	for i = 1: N_p
				log(w[i]) = log(w[i]) + p_y_g_x(y[k] - obs(x[i]))
    	end
		@show sum(w)
		w ./= sum(w)
	    N_eff = 1.0/sum(w.*w)
        if (N_eff < N_thr)
	        #resample
            # multinomial resampling
			cdfw = cumsum(w)
			new_pts .= similar(x)
            for j = 1:N_p
			    r = rand()
			    for i = 1:N_p
		    	    if cdfw[i] >= r
					    new_pts[j] = x[i]
			            break
			        end
		        end
		    end
			x .= new_pts
	    end
		x .= next.(x,s)
    end
	return x, w
end
function assimilate()
    # Often much larger.
	Np = 100
	K = 10
	x = rand(Np)
	s = 0.001
	x_true = ones(K)
	x0_true = 2*pi*rand()
	x_true[1] = x0_true
	y = zeros(K)
	for i = 2:K
	    x_true[i] = next(x_true[i-1], s)
		y[i] = x_true[i] + σ_o*randn()
	end
	N_thr = 20
	# store trajectory of w and x.
	x, w = sir(y, x, s, N_thr)
	return x, w, x_true[K]
end


using Test, MPI
using ClimateMachine.Mesh.Topologies
using ClimateMachine.Mesh.Grids
using ClimateMachine.Mesh.Geometry
using StaticArrays

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

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

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

    Ne = (20, 2, 20)

    polynomialorder = 4

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

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

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

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


################################################################################
#                          TEST: Euler Angle and Axis
################################################################################

# Test constructions.
for k = 1:samples
    # Sample a random vector and angle.
    a = rand()
    v = rand(3)

    angleaxis = EulerAngleAxis(a,v)

    # Check if the representation was created correctly.
    @test angleaxis.a == a
    @test angleaxis.v == v

    # Test multiplication.
    a1  = -π + 2*π*rand()
    v1  = randn(3)
    v1 /= norm(v1)

    a2  = -π + 2*π*rand()
    v2  = randn(3)
    v2 /= norm(v2)

    ea1 = EulerAngleAxis(a1, v1)
    ea2 = EulerAngleAxis(a2, v2)
    ear = ea2*ea1

    q1  = angleaxis_to_quat(ea1)
    q2  = angleaxis_to_quat(ea2)
    qr  = q1*q2
    eaq = quat_to_angleaxis(qr)

    @test    ear.a ≈ eaq.a
    @test ear.v[1] ≈ eaq.v[1]
    @test ear.v[2] ≈ eaq.v[2]
    @test ear.v[3] ≈ eaq.v[3]

    # Test inversion.
    a   = -π + 2*π*rand()
    v   = randn(3)
    v  /= norm(v)
    ea  = EulerAngleAxis(a,v)
    ear = ea*inv_rotation(ea)

    @test norm(ear.v) ≈ 0 atol = 1e-7
    @test       ear.a ≈ 0 atol = 1e-7
end

# Test types.
@test typeof(EulerAngleAxis(1, [1;1;1])) <: EulerAngleAxis{Int}
@test typeof(EulerAngleAxis(1.f0, Float32[1;1;1])) <: EulerAngleAxis{Float32}
@test typeof(EulerAngleAxis(1.f0, [1;1;1])) <: EulerAngleAxis{Float32}

# Test multiplication that leads to unitary rotation.
let
    ea1 = EulerAngleAxis(pi, [1.0,0.0,0.0])
    ear = ea1*ea1

    @test norm(ear.v) ≈ 0
    @test       ear.a ≈ 0
end

# Test `Float16` and `Float32` multiplications.
let
    v   = SVector{3,Float32}(1,0,0)
    a   = Float32(1)
    ea1 = EulerAngleAxis(a,v)
    ear = ea1*ea1

    @test eltype(ear.a) == eltype(ear.v) == Float32

    v   = SVector{3,Float16}(1,0,0)
    a   = Float16(1)
    ea1 = EulerAngleAxis(a,v)
    ear = ea1*ea1

    @test eltype(ear.a) == eltype(ear.v) == Float16
end

# Test exceptions.
@test_throws Exception EulerAngleAxis(randn(), [1;1;1;1])
@test_throws Exception EulerAngleAxis(randn(), [1;1])
@test_throws Exception EulerAngleAxis(randn(), SVector{4,Int}[1;1;1;1])
@test_throws Exception EulerAngleAxis(randn(), SVector{4,Float32}[1;1;1;1])
@test_throws Exception EulerAngleAxis(randn(), SVector{4,Float64}[1;1;1;1])
@test_throws Exception EulerAngleAxis(randn(), SVector{2,Int}[1;1])
@test_throws Exception EulerAngleAxis(randn(), SVector{2,Float32}[1;1])
@test_throws Exception EulerAngleAxis(randn(), SVector{2,Float64}[1;1])


using TransformUtils

function uteOdom(x::Array{Float64,1}, vc::Float64, alpha::Float64, dt::Float64)
  L = 2.83; H=076; b=0.5;a=3.78;
  d = [(vc*cos(x[3])-vc/L*tan(x[3])*(a*sin(x[3])+b*cos(x[3])));
       (vc*sin(x[3])+vc/L*tan(x[3])*(a*cos(x[3])-b*sin(x[3])));
        vc/L*tan(alpha)]
  x += dt*d
  x[3] = wrapRad(x[3])
  return x
end

function allOdo(sensors::Array{Float64,2})
  odo = zeros(size(sensors,1),3)
  T0 = 0.0
  for i in 2:size(sensors,1)
    odo[i,:]=uteOdom(vec(odo[i-1,:]),sensors[i,2],sensors[i,3],sensors[i,1]-T0)
    T0 = sensors[i,1]
  end

  return odo
end

vc(v::Float64, alpha::Float64; L=2.80381, H=0.828329 ) = v./(1.0-tan(alpha)*H/L)
ve(v::Float64, alpha::Float64; L=2.80381, H=0.828329 ) = v.*(1.0-tan(alpha).*H/L)
dPhi(v::Float64, alpha::Float64; L=2.80381) = 1.0*v*tan(alpha)/L

function compensateRawsensors(drs::Array{Float64,1};
                    L=2.80381, H=0.828329, whlsf=0.94, strsf=1.0199, strbi=0.00159)
    return whlsf*drs[2], strsf*drs[3]+strbi
end

function uteOdomEasy(x::Array{Float64,1}, whlspd::Float64, strangl::Float64, dt::Float64; L=2.80381, H=0.828329)
  wTb = SE2(x)
  v = vc(whlspd,strangl,L=L,H=H)
  dph = dPhi(v,strangl,L=L)
  #dph = v*tan(strangl)
  bTnb = SE2(dt*[v;0.0;dph])
  wTnb = wTb*bTnb
  return vec(se2vee(wTnb))
end

# perform integration over all data, sensors is rows of data, by columns of sensors
function allOdoEasy(sensors::Array{Float64,2};
                    L=2.80381, H=0.828329, whlsf=0.94, strsf=1.0199, strbi=0.00159)
  #msensors = [sensors[:,1]';x[1]*sensors[:,2]';x[5]*sensors[:,3]'+x[2]]';
  odo = zeros(size(sensors,1),3)
  for i in 2:size(sensors,1)
    whlspd, strang = compensateRawsensors(vec(sensors[i,:]), L=L,H=H, whlsf=whlsf, strbi=strbi,strsf=strsf)
    odo[i,:] = uteOdomEasy(vec(odo[i-1,:]), whlspd, strang, (sensors[i,1]-sensors[i-1,1]), L=L, H=H )
  end

  return odo
end


@doc raw"""
    ∇distance(M,y,x[, p=2])

compute the (sub)gradient of the distance (squared)

```math
f(x) = \frac{1}{2} d^p_{\mathcal M}(x,y)
```

to a fixed point `y` on the manifold `M` and `p` is an
integer. The gradient reads

```math
  \nabla f(x) = -d_{\mathcal M}^{p-2}(x,y)\log_xy
```

for $p\neq 1$ or $x\neq  y$. Note that for the remaining case $p=1$,
$x=y$ the function is not differentiable. In this case, the function returns the
corresponding zero tangent vector, since this is an element of the subdifferential.

# Optional

* `p` – (`2`) the exponent of the distance,  i.e. the default is the squared
  distance
"""
∇distance(M,y,x,p::Int=2) = (p==2) ? -log(M,x,y) : -distance(M,x,y)^(p-2)*log(M,x,y)

@doc raw"""
    ∇u,∇v = ∇intrinsic_infimal_convolution_TV12(M,f,u,v,α,β)

compute (sub)gradient of the intrinsic infimal convolution model using the mid point
model of second order differences, see [`costTV2`](@ref), i.e. for some $f ∈ \mathcal M$
on a `PowerManifold` manifold $\mathcal M$ this function computes the (sub)gradient of

```math
E(u,v) =
\frac{1}{2}\sum_{i ∈ \mathcal G} d_{\mathcal M}(g(\frac{1}{2},v_i,w_i),f_i)
+ \alpha
\bigl(
\beta\mathrm{TV}(v) + (1-\beta)\mathrm{TV}_2(w)
\bigr),
```
where both total variations refer to the intrinsic ones, [`∇TV`](@ref) and
[`∇TV2`](@ref), respectively.
"""
function ∇intrinsic_infimal_convolution_TV12(M::mT,f,u,v,α,β) where {mT <: Manifold}
  c = mid_point(M,u,v,f)
  iL = log(M,c,f)
  return adjoint_differential_geodesic_startpoint(M,u,v,1/2,iL) + α*β*∇TV(M,u), adjoint_differential_geodesic_endpoint(M,u,v,1/2,iL) + α * (1-β) * ∇TV2(M,v)
end
@doc raw"""
    ∇TV(M,(x,y),[p=1])

compute the (sub) gradient of $\frac{1}{p}d^p_{\mathcal M}(x,y)$ with respect
to both $x$ and $y$.
"""
function ∇TV(M::MT, xT::Tuple{T,T}, p=1)where {MT <: Manifold, T}
  x = xT[1];
  y = xT[2];
  if p==2
      return (-log(M,x,y), -log(M,y,x))
  else
    d = distance(M,x,y);
    if d==0 # subdifferential containing zero
      return (zero_tangent_vector(M,x),zero_tangent_vector(M,y))
    else
      return (-log(M,x,y)/(d^(2-p)), -log(M,y,x)/(d^(2-p)))
    end
  end
end
@doc raw"""
    ξ = ∇TV(M,λ,x,[p])
Compute the (sub)gradient $\partial F$ of all forward differences orrucirng,
in the power manifold array, i.e. of the function

$F(x) = \sum_{i}\sum_{j ∈ \mathcal I_i} d^p(x_i,x_j)$

where $i$ runs over all indices of the `PowerManifold` manifold `M`
and $\mathcal I_i$ denotes the forward neighbors of $i$.

# Input
* `M` – a `PowerManifold` manifold
* `x` – a point.

# Ouput
* ξ – resulting tangent vector in $T_x\mathcal M$.
"""
function ∇TV(M::PowerManifold,x,p::Int=1)
    power_size = power_dimensions(M)
    rep_size = representation_size(M.manifold)
    R = CartesianIndices(Tuple(power_size))
    d = length(power_size)
    maxInd = last(R)
    X = zero_tangent_vector(M,x)
    c = costTV(M,x,p,0)
    for i in R # iterate over all pixel
        di = 0.
        for k in 1:d # for all direction combinations
            ek = CartesianIndex(ntuple(i  ->  (i==k) ? 1 : 0, d) ) #k th unit vector
            j = i+ek # compute neighbor
            if all( map(<=, j.I, maxInd.I)) # is this neighbor in range?
                if p != 1
                    g = (c[i]==0 ? 1 : 1/c[i]) .* ∇TV(M.manifold,(x[i],x[j]),p) # Compute TV on these
                else
                    g = ∇TV(M.manifold,(x[i],x[j]),p) # Compute TV on these
                end
                X[i] += g[1]
                X[j] += g[2]
            end
        end # directions
    end # i in R
    return X
end

@doc raw"""
    ξ = forward_logs(M,x)

compute the forward logs $F$ (generalizing forward differences) orrucirng,
in the power manifold array, the function

```math
$F_i(x) = \sum_{j ∈ \mathcal I_i} \log_{x_i} x_j,\quad i  ∈  \mathcal G,
```

where $\mathcal G$ is the set of indices of the `PowerManifold` manifold `M`
and $\mathcal I_i$ denotes the forward neighbors of $i$.

# Input
* `M` – a `PowerManifold` manifold
* `x` – a point.

# Ouput
* `ξ` – resulting tangent vector in $T_x\mathcal M$ representing the logs, where
  $\mathcal N$ is thw power manifold with the number of dimensions added to `size(x)`.
"""
function forward_logs(M::PowerManifold{𝔽,TM,TSize,TPR}, p) where {𝔽,TM,TSize,TPR}
    power_size = power_dimensions(M)
    R = CartesianIndices(Tuple(power_size))
    d = length(power_size)
    sX = size(p)
    maxInd = [last(R).I...] # maxInd as Array
    if d > 1
        d2 = fill(1,d+1)
        d2[d+1] = d
    else
        d2 = 1
    end
    N = PowerManifold(M.manifold, TPR(), power_size..., d)
    xT = repeat(p,inner=d2)
    X = zero_tangent_vector(N,xT)
    for i in R # iterate over all pixel
        for k in 1:d # for all direction combinations
            I = [i.I...] # array of index
            J = I .+ 1 .* (1:d .== k) #i + e_k is j
            if all( J .<= maxInd ) # is this neighbor in range?
                j = CartesianIndex{d}(J...) # neigbbor index as Cartesian Index
                set_component!(
                    N,
                    X,
                    log(M.manifold, get_component(M,p,i), get_component(M,p,j)),
                    i,
                    k,
                ) # Compute log and store in kth entry
            end
        end # directions
    end # i in R
    return X
end

@doc raw"""
    ∇TV2(M,(x,y,z),p)

computes the (sub) gradient of $\frac{1}{p}d_2^p(x,y,z)$ with respect
to $x$, $y$, and $z$, where $d_2$ denotes the second order absolute difference
using the mid point model, i.e. let
```math
  \mathcal C = \bigl\{ c ∈  \mathcal M \ |\ g(\tfrac{1}{2};x_1,x_3) \text{ for some geodesic }g\bigr\}
```
denote the mid points between $x$ and $z$ on the manifold $\mathcal M$.
Then the absolute second order difference is defined as

```math
d_2(x,y,z) = \min_{c ∈ \mathcal C_{x,z}} d(c,y).
```

While the (sub)gradient with respect to $y$ is easy, the other two require
the evaluation of an [`adjoint_Jacobi_field`](@ref).
See [Illustration of the Gradient of a Second Order Difference](@ref secondOrderDifferenceGrad)
for its derivation.
"""
function ∇TV2(M::MT, xT, p::Number=1) where {MT <: Manifold}
  x = xT[1];
  y = xT[2];
  z = xT[3];
  c = mid_point(M,x,z,y) # nearest mid point of x and z to y
  d = distance(M,y,c)
  innerLog = -log(M,c,y)
  if p==2
      return ( adjoint_differential_geodesic_startpoint(M,x,z,1/2,innerLog), -log(M,y,c), adjoint_differential_geodesic_endpoint(M,x,z,1/2,innerLog))
  else
    if d==0 # subdifferential containing zero
      return (zero_tangent_vector(M,x),zero_tangent_vector(M,y),zero_tangent_vector(M,z))
    else
      return ( adjoint_differential_geodesic_startpoint(M,x,z,1/2,innerLog/(d^(2-p))), -log(M,y,c)/(d^(2-p)), adjoint_differential_geodesic_endpoint(M,x,z,1/2,innerLog/(d^(2-p))) )
    end
  end
end
@doc raw"""
    ∇TV2(M,q [,p=1])

computes the (sub) gradient of $\frac{1}{p}d_2^p(x_1,x_2,x_3)$
with respect to all $x_1,x_2,x_3$ occuring along any array dimension in the
point `x`, where `M` is the corresponding `PowerManifold`.
"""
function ∇TV2(M::PowerManifold, q, p::Int=1)
    power_size = power_dimensions(M)
    rep_size = representation_size(M.manifold)
    R = CartesianIndices(Tuple(power_size))
    d = length(power_size)
    minInd, maxInd = first(R), last(R)
    X = zero_tangent_vector(M,q)
    c = costTV2(M,q,p,false)
    for i in R # iterate over all pixel
        di = 0.
        for k in 1:d # for all direction combinations (TODO)
            ek = CartesianIndex(ntuple(i  ->  (i==k) ? 1 : 0, d) ) #k th unit vector
            jF = i+ek # compute forward neighbor
            jB = i-ek # compute backward neighbor
            if all( map(<=, jF.I, maxInd.I) ) && all( map(>=, jB.I, minInd.I)) # are neighbors in range?
                if p != 1
                    g = (c[i] == 0 ? 1 : 1/c[i]) .* ∇TV2(M.manifold,(q[jB],q[i],q[jF]),p) # Compute TV2 on these
                else
                    g = ∇TV2(M.manifold,(q[jB],q[i],q[jF]),p) # Compute TV2 on these
                end
                set_component!(M,X,g[1],jB)
                set_component!(M,X,g[2],i)
                set_component!(M,X,g[3],jF)
            end
        end # directions
    end # i in R
    return X
end


import Random
import PyPlot

using LinearAlgebra
import Interpolations

#include("../src/PatchMixtureKriging.jl")
import PatchMixtureKriging

PyPlot.close("all")

Random.seed!(25)

fig_num = 1

##### regression
#θ = Spline34KernelType(0.2)
θ = PatchMixtureKriging.BrownianBridge10(1.0)
#θ = BrownianBridge20(1.0)
#θ = BrownianBridge1ϵ(4.5)
#θ = BrownianBridge2ϵ(2.5)

σ² = 1e-5
N = 15

# if 0 and 1 are included, we have posdef error, or rank = N - 2.
x_range = LinRange(1e-5, 1.0-1e-5, N)

X = collect( [x_range[n]] for n = 1:N )
coordwrapper = xx->PatchMixtureKriging.convert2itpindex(xx, X[1], X[end], [length(X)])

f = xx->sinc(4*xx)*xx^3
y = f.(x_range)

# check posdef.
K = PatchMixtureKriging.constructkernelmatrix(X, θ)
println("rank(K) = ", rank(K))

println("isposdef = ", isposdef(K))

# fit RKHS.
η = PatchMixtureKriging.RKHSProblemType( zeros(Float64,length(X)),
                     X,
                     θ,
                     σ²)
                     PatchMixtureKriging.fitRKHS!(η, y)

# query.
Nq = 100
xq_range = LinRange(0.0, 1.0, Nq)
xq = collect( [xq_range[n]] for n = 1:Nq )

f_xq = f.(xq_range) # generating function.

yq = Vector{Float64}(undef, Nq)
PatchMixtureKriging.query!(yq,xq,η)

# Visualize regression result.
PyPlot.figure(fig_num)
fig_num += 1

PyPlot.plot(X, y, ".", label = "observed")
PyPlot.plot(xq, yq, label = "query - RKHS")

PyPlot.plot(xq, f_xq, label = "true")

title_string = "1-D RKHS demo"
PyPlot.title(title_string)
PyPlot.legend()
##### end of regression


import IntervalArithmetic: diam, isinterior
import Base: iterate, eltype, IteratorSize, copy

export branch_and_prune, Bisection, Newton, RootSearch
export step!

diam(x::Root) = diam(x.interval)

struct RootSearchState{T <: Union{Interval,IntervalBox}}
    working::Vector{T}
    outputs::Vector{Root{T}}
end
RootSearchState(region::T) where {T<:Union{Interval,IntervalBox}} =
    RootSearchState([region], Root{T}[])

copy(state::RootSearchState) =
    RootSearchState(copy(state.working), copy(state.outputs))

"""
    RootSearch{R <: Union{Interval,IntervalBox}, S <: Contractor, T <: Real}

Type implementing the `Base.Iterator` interface to the branch and prune routine.
Returns the `RootSearchState` at each iteration. Note: Each iteration mutates
the `RootSearchState`. Use `copy(state::RootSearchState)` to create an
independent instance if necessary.
"""
struct RootSearch{R <: Union{Interval,IntervalBox}, S <: Contractor, T <: Real}
    region::R
    contractor::S
    tolerance::T
end

eltype(::Type{T}) where {R, T <: RootSearch{R}} = RootSearchState{R}
IteratorSize(::Type{T}) where {T <: RootSearch} = Base.SizeUnknown()

# function start(iter::RootSearch)
#     state = RootSearchState(iter.region)
#     sizehint!(state.outputs, 100)
#     sizehint!(state.working, 1000)
#     return state
# end

function iterate(iter::RootSearch)
    state = RootSearchState(iter.region)
    sizehint!(state.outputs, 100)
    sizehint!(state.working, 1000)
    return state, state
end

"""
    step!(state::RootSearchState, contractor, tolerance)

Progress `state` by treating one of its `working` regions. Note: `state.working`
is always modified. If a root is found, it is added to `state.outputs`.
"""
function step!(state::RootSearchState, contractor, tolerance)
    X = pop!(state.working)
    status, output = contractor(X, tolerance)
    if status == :empty
        return nothing
    elseif status == :unique
        push!(state.outputs, Root(output, :unique))
    elseif diam(output) < tolerance
        push!(state.outputs, Root(output, :unknown))
    else # branch
        X1, X2 = bisect(X)
        push!(state.working, X1, X2)
    end
    return nothing
end

# function next(iter::RootSearch, state::RootSearchState)
#     step!(state, iter.contractor, iter.tolerance)
#     return state, state
# end

function iterate(iter::RootSearch, state::RootSearchState)

    isempty(state.working) && return nothing

    step!(state, iter.contractor, iter.tolerance)
    return state, state
end

# done(iter::RootSearch, state::RootSearchState) = isempty(state.working)

"""
    branch_and_prune(X, contract, tol=1e-3)

Generic branch and prune routine for finding isolated roots using the `contract`
function as the contractor.

Inputs:
- `X`: `Interval` or `IntervalBox`
- `contract`: function that determines the status of a given box `X`. It
    returns the new box and a symbol indicating the status. Current possible
    values are of type `Bisection` or `Newton`

"""
function branch_and_prune(X, contractor, tol=1e-3)
    iter = RootSearch(X, contractor, tol)
    local output
    # complete iteration
    for state in iter
        output = state.outputs
    end
    return output
end

export recursively_branch_and_prune

function recursively_branch_and_prune(h, X, contractor=BisectionContractor, final_tol=1e-14)
    tol = 2
    roots = branch_and_prune(h, X, IntervalRootFinding.BisectionContractor, tol)

    while tol > 1e-14
       tol /= 2
       roots = branch_and_prune(h, roots, IntervalRootFinding.BisectionContractor, tol)
    end

    return roots
end

const IntervalLike{T} = Union{Interval{T}, IntervalBox{T}}
const NewtonLike = Union{Type{Newton}, Type{Krawczyk}}

"""
    roots(f, X, contractor, tol=1e-3)
    roots(f, deriv, X, contractor, tol=1e-3)

Uses a generic branch and prune routine to find in principle all isolated roots of a function
`f:R^n → R^n` in a box `X`, or a vector of boxes.

Inputs:
- `f`: function whose roots will be found
- `X`: `Interval` or `IntervalBox`
- `contractor`: function that, when applied to the function `f`, determines
    the status of a given box `X`. It returns the new box and a symbol indicating
    the status. Current possible values are `Bisection`, `Newton` and `Krawczyk`
- `deriv` ; explicit derivative of `f` for `Newton` and `Krawczyk`
"""
function roots(f, X::IntervalLike{T}, ::Type{Bisection}, tol::Float64=1e-3) where {T}
    branch_and_prune(X, Bisection(f), tol)
end

function roots(f, X::Interval{T}, C::NewtonLike, tol::Float64=1e-3) where {T}

    deriv = x -> ForwardDiff.derivative(f, x)

    roots(f, deriv, X, C, tol)
end

function roots(f, deriv, X::Interval{T}, C::NewtonLike, tol::Float64=1e-3) where {T}
    branch_and_prune(X, C(f, deriv), tol)
end

function roots(f, X::IntervalBox{T}, C::NewtonLike, tol::Float64=1e-3) where {T}

    deriv = x -> ForwardDiff.jacobian(f, x)

    roots(f, deriv, X, C, tol)
end

function roots(f, deriv, X::IntervalBox{T}, C::NewtonLike, tol::Float64=1e-3) where {T}
    branch_and_prune(X, C(f, deriv), tol)
end

roots(f, r::Root, contractor::Type{C}, tol::Float64=1e-3) where {C<:Contractor} =
    roots(f, r.interval, contractor, tol)
roots(f, deriv, r::Root, contractor::Type{C}, tol::Float64=1e-3) where {C<:Contractor} =
    roots(f, deriv, r.interval, contractor, tol)

# Acting on a Vector:

# TODO: Use previous status information about roots:
roots(f, V::Vector{Root{T}}, contractor::Type{C}, tol::Float64=1e-3) where {T, C<:Contractor} =
    vcat(roots.(f, V, contractor, tol)...)
roots(f, deriv, V::Vector{Root{T}}, contractor::Type{C}, tol::Float64=1e-3) where {T, C<:Contractor} =
    vcat(roots.(f, deriv, V, contractor, tol)...)

# Complex:

function roots(f, Xc::Complex{Interval{T}}, contractor::Type{C},
        tol::Float64=1e-3) where {T, C<:Contractor}

    g = realify(f)
    Y = IntervalBox(reim(Xc)...)
    rts = roots(g, Y, contractor, tol)

    return [Root(Complex(root.interval...), root.status) for root in rts]
end

function roots(f, Xc::Complex{Interval{T}}, C::NewtonLike, tol::Float64=1e-3) where {T}

    g = realify(f)

    g_prime = x -> ForwardDiff.jacobian(g, x)

    Y = IntervalBox(reim(Xc)...)
    rts = roots(g, g_prime, Y, C, tol)

    return [Root(Complex(root.interval...), root.status) for root in rts]
end

function roots(f, deriv, Xc::Complex{Interval{T}}, C::NewtonLike, tol::Float64=1e-3) where {T}

    g = realify(f)

    g_prime = realify_derivative(deriv)

    Y = IntervalBox(reim(Xc)...)
    rts = roots(g, g_prime, Y, C, tol)

    return [Root(Complex(root.interval...), root.status) for root in rts]
end

# Default
roots(f::F, deriv::F, X, tol::Float64=1e-15) where {F<:Function} = roots(f, deriv, X, Newton, tol)
roots(f::F, X, tol::Float64=1e-15) where {F<:Function} = roots(f, X, Newton, tol)


#---------------------------------------------------------
# # [A coupled PO₄–POP model](@id P-model)
#---------------------------------------------------------

#md # [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__/notebooks/3_Pmodel.ipynb)
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/notebooks/3_Pmodel.ipynb)

#md # !!! note
#md #     All the AIBECS tutorials and how-to guides are available as Jupyter notebooks.
#md #     You can view them with [nbviewer](https://nbviewer.jupyter.org/)or execute them online with [binder](https://mybinder.org/) by clicking on the badges above!
#md #     (Note that binder can be slow to launch and its memory caps can be a problem when running.)

# In this tutorial we will explicitly simulate 2 tracers whose distributions control and feed back on each other.

# We consider a simple model for the cycling of phosphorus with 2 state variables consisting of phosphate (PO₄) AKA dissolved inorganic phosphorus (DIP) and particulate organic phosphorus (POP).
# The dissolved phases are transported by advection and diffusion whereas the particulate phase sinks rapidly down the water column without any appreciable transport by the circulation.
#
# The governing equations that couple the 3D concentration fields of DIP and POP, denoted $x_\mathsf{DIP}$ and $x_\mathsf{POP}$, respectively, are:
#
# $$\left[\frac{\partial}{\partial t} + \nabla \cdot (\boldsymbol{u} + \mathbf{K}\nabla )\right] x_\mathsf{DIP} = -U(x_\mathsf{DIP}) + R(x_\mathsf{POP}),$$
#
# and
#
# $$\left[\frac{\partial}{\partial t} + \nabla \cdot \boldsymbol{w}\right] x_\mathsf{POP} = U(x_\mathsf{DIP}) - R(x_\mathsf{POP}).$$

# The $\nabla \cdot \left[ \boldsymbol{u} - \mathbf{K} \nabla \right]$ and $\nabla \cdot \boldsymbol{w}$ operators represent the ocean circulation and the sinking of particles, respectively.
# ([Tracer transport operators are described in the documentation](@ref tracer-transport-operators).)

# The function $U$ represents the biological uptake of DIP by phytoplankton, which we model here as
#
# $$U(x_\mathsf{DIP}) = \frac{x_\mathsf{DIP}}{\tau_\mathsf{DIP}} \, \frac{x_\mathsf{DIP}}{x_\mathsf{DIP} + k} \, (z < z_0),$$
#
# with the timescale, $\tau$, the half-saturation rate $k$, and the depth $z_0$ as parameters.

# The function $R$ defines the remineralization rate of POP, which converts POP back into DIP.
# For the remineralization, we simply use a linear rate constant, i.e.,
#
# $$R(x_\mathsf{POP}) = \frac{x_\mathsf{POP}}{\tau_\mathsf{POP}}.$$

# We start by telling Julia we want to use the AIBECS and the OCIM0.1 circulation for DIP.

using AIBECS
grd, T_OCIM = OCIM0.load()
T_DIP(p) = T_OCIM

# For the sinking of particles, we use the `transportoperator` function

T_POP(p) = transportoperator(grd, z -> w(z,p))

# for which we need to define the sinking speed `w(z,p)` as a function of depth `z` and of the parameters `p`.
# Following the assumption that $w(z) = w_0 + w' z$ increases linearly with depth, we write it as

function w(z,p)
    @unpack w₀, w′ = p
    return @. w₀ + w′ * z
end

# ##### Uptake (DIP → POP)

# For the uptake, $U$, we write

z = depthvec(grd)
function U(x,p)
    @unpack τ_DIP, k, z₀ = p
    return @. x/τ_DIP * x/(x+k) * (z≤z₀) * (x≥0)
end

# where we have "unpacked" the parameters to make the code clearer and as close to the mathematical equation as possible.
# (Note we have also added a constraint that `x` must be positive for uptake to happen.)

# ##### Remineralization (POP → DIP)

# For the remineralization, $R$, we write

function R(x,p)
    @unpack τ_POP = p
    return x / τ_POP
end

# ##### Net sources and sinks

# We lump the sources and sinks into `G` functions for DIP and POP.

function G_DIP(DIP, POP, p)
    @unpack DIP_geo, τ_geo = p
    return @. -$U(DIP,p) + $R(POP,p) + (DIP_geo - DIP) / τ_geo
end
function G_POP(DIP, POP, p)
    @unpack τ_geo = p
    return @. $U(DIP,p) - $R(POP,p) - POP / τ_geo
end

# where we have imposed a slow restoring of DIP to the global mean `DIP_geo` to prescribe the global mean concentration.
# (The `$` signs in front of `U` and `R` protect them from the broadcast macro `@.`)

# We now define and build the parameters.

# In this tutorial we will specify some initial values for the parameters
# and also include units.

import AIBECS: @units, units
import AIBECS: @initial_value, initial_value
using Unitful: m, d, s, yr, Myr, mol, mmol, μmol, μM
@initial_value @units struct PmodelParameters{U} <: AbstractParameters{U}
    w₀::U       |  0.64 | m/d
    w′::U       |  0.13 | m/d/m
    τ_DIP::U    | 230.0 | d
    k::U        |  6.62 | μmol/m^3
    z₀::U       |  80.0 | m
    τ_POP::U    |   5.0 | d
    τ_geo::U    |   1.0 | Myr
    DIP_geo::U  |  2.12 | mmol/m^3
end

# Finally, thanks to the initial values we provided, we can instantiate the parameter vector succintly as

p = PmodelParameters()

# We generate the state function `F`,

nb = sum(iswet(grd))
F = AIBECSFunction((T_DIP, T_POP), (G_DIP, G_POP), nb)

# generate the steady-state problem,

@unpack DIP_geo = p
x = DIP_geo * ones(2nb) # initial guess
prob = SteadyStateProblem(F, x, p)

# and solve it

sol = solve(prob, CTKAlg()).u

# We can look at different the DIP and POP fields using the Plots.jl recipes.

DIP, POP = state_to_tracers(sol, grd) # unpack tracers

# First, let's look at the mean profile

using Plots
plothorizontalmean(DIP * (mol/m^3) .|> μM, grd)

# We can plot the concentration of DIP at a given depth via, e.g.,

plothorizontalslice(DIP * (mol/m^3) .|> μM, grd, depth=1000m, color=:viridis)

# Or have a look at a map of the uptake at the surface

plotverticalintegral(U(DIP,p) * (mol/m^3/s) .|> mmol/yr/m^3, grd, color=:algae)

# Or look at what is exported below 500 m

plothorizontalslice(POP .* w(z,p) * (mol/m^3*m/s) .|> mmol/yr/m^2, grd, depth=500m, color=:inferno, rev=true)

# Now let's make our model a little fancier and use a fine topographic map to refine the remineralization profile.
# For this, we will use the ETOPO dataset, which can be downloaded by AIBECS via

f_topo = ETOPO.fractiontopo(grd)

# `f_topo` is the fraction of sediments located in each wet box of the `grd` grid.
# We can use it to redefine the transport operator for sinking particles to take into consideration the subgrid topography,
# such that the fine-resolution sediments intercept settling POP.

T_POP2(p) = transportoperator(grd, z -> w(z,p); frac_seafloor=f_topo)

# With this new vertical transport for POP, we can recreate our problem, solve it again

F2 = AIBECSFunction((T_DIP, T_POP2), (G_DIP, G_POP), nb)
prob2 = SteadyStateProblem(F2, x, p)
sol2 = solve(prob2, CTKAlg()).u
DIP2, POP2 = state_to_tracers(sol2, grd) # unpack tracers

# and check the difference

plotzonalaverage((DIP2 - DIP) ./ DIP .|> u"percent", grd, color=:balance, clim=(-5,5))

# This zonal average shows how much DIP is redistributed as it is prevented from sinking out of the surface layers with the new subgrid parameterization.

# Let's look at the vertical average.

plotverticalaverage((DIP2 - DIP) ./ DIP .|> u"percent", grd, color=:balance, clim=(-10,10))

# This shows minor changes on the order of 1%, on the global scale,
# except along the coasts, which retain much more DIP with the subgrid topography.


export Ellipsoid

mutable struct Ellipsoid{N,N²} <: Shape{N,N²}
    c::SVector{N,Float64}  # center of ellipsoid
    ri2::SVector{N,Float64}  # inverse squares of "radii" (semi-axes) in axis directions
    p::SMatrix{N,N,Float64,N²}  # projection matrix to Ellipsoid coordinates; must be orthonormal (see surfpt_nearby)
    Ellipsoid{N,N²}(c,ri2,p) where {N,N²} = new(c,ri2,p)  # suppress default outer constructor
end

Ellipsoid(c::SVector{N,<:Real},
          r::SVector{N,<:Real},
          axes::SMatrix{N,N,<:Real}=SMatrix{N,N,Float64}(I)
          ) where {N} =
    Ellipsoid{N,N*N}(c, float.(r).^-2, inv(axes ./ sqrt.(sum(abs2,axes,dims=Val(1)))))

Ellipsoid(c::AbstractVector{<:Real},  # center of ellipsoid
          r::AbstractVector{<:Real},  # ""
          axes::AbstractMatrix{<:Real}=Matrix{Float64}(I,length(c),length(c))) =  # columns are axes vector; assumed orthogonal
    (N = length(c); Ellipsoid(SVector{N}(c), SVector{N}(r), SMatrix{N,N}(axes)))

Ellipsoid(b::Cuboid{N,N²}) where {N,N²} = Ellipsoid{N,N²}(b.c, (b.r).^-2, b.p)

Base.:(==)(b1::Ellipsoid, b2::Ellipsoid) = b1.c==b2.c && b1.ri2==b2.ri2 && b1.p==b2.p
Base.isapprox(b1::Ellipsoid, b2::Ellipsoid) = b1.c≈b2.c && b1.ri2≈b2.ri2 && b1.p≈b2.p
Base.hash(b::Ellipsoid, h::UInt) = hash(b.c, hash(b.ri2, hash(b.p, hash(:Ellipsoid, h))))

level(x::SVector{N,<:Real}, b::Ellipsoid{N}) where {N} = √dot((b.p * (x - b.c)).^2, b.ri2) - 1.0

function surfpt_nearby(x::SVector{N,<:Real}, b::Ellipsoid{N}) where {N}
    if x == b.c
        _m, i = findmax(b.ri2)
        nout = b.p[i,:]  # assume b.p is orthogonal
        return b.c + nout/√b.ri2[i], nout
    end

    # For a given point x and equation of ellipsoid f(x) = 1, find t such that x₀ = x + t*∇f(x)
    # is on the ellipsoid.  Eventually this reduces to a quadratic equation for t.  The
    # following is evaluation of the quadratic formula in a numerically stable way.
    px = b.p * (x - b.c)  # in ellipsoid's coordinates
    px2 = px.^2

    px²r⁻² = px2 ⋅ b.ri2
    px²r⁻⁴ = px2 ⋅ (b.ri2.^2)
    px²r⁻⁶ = px2 ⋅ (b.ri2.^3)

    q24 = (px²r⁻² - 1) / px²r⁻⁴
    q64 = px²r⁻⁶ / px²r⁻⁴

    t = -q24 / (1 + √(1 - q24 * q64))

    # From t, recover x₀ = x + t*∇f(x).
    px₀ = (t*b.ri2 .+ 1) .* px  # surface point in ellipsoid coordinates

    # Transform back to the original coordinates.
    x₀ = b.p' * px₀ + b.c
    nout = normalize(b.p' * (px .* b.ri2))

    return x₀, nout
end

function boundpts(b::Ellipsoid{N}) where {N}
    # Return the points tangential to the bounding box.
    # For N = 3, it returns three points at which the direction normals are +x, +y, +z
    # directions, respectively.

    r2 = 1 ./ b.ri2
    ndir = b.p  # Cartesian directions in ellipsoid coordinates: b.p * I

    # In the ellipsoid coordinates, the point on the ellipsoid where the direction normal is
    # n is (n .* r2) / sqrt(r2' * n.^2).  Below is the broadcasted version of this over a
    # matrix n, whose each column is a direction normal.  Once calculated, we need to
    # change the coordinates back to the original coordinates.
    M = b.p' * ((ndir .* r2) ./ sqrt.(r2' * ndir.^2))

    return M
end

function bounds(b::Ellipsoid{N}) where {N}
    M = boundpts(b)

    # Note that when M does not include NaN, we can simply set m = diag(M), because the
    # first (second, third) column of M is the point on the ellipsoid at which the normal
    # vector is +x (+y, +z).  Therefore, the x (y, z) point of the first (second, third)
    # column has the largest x (y, z) coordinate.
    #
    # However, if one of a, b, c is zero, the shape is a disk.  Then one column of M can be
    # completely filled with NaN.  This column must not be counted as a bounding point, so
    # we apply NaN-ignoring maximum by StaticArrays.reducedim along the row direction.
    #
    # For the efficient implementation of NaN-ignoring maximum, see
    # https://discourse.julialang.org/t/inconsistency-between-findmax-and-maximum-with-respect-to-nan/4214/8
    m = reduce((x,y) -> x<y ? y : x, M, init=-Inf, dims=Val(2))[:,1]

    return (b.c-m,b.c+m)
end


function numlattices(rcutmax, pbc, a, b, c)

    d = (length(c) < 3) ? 2 : 3

    m = 0
    n = 0
    p = 0

    if pbc[1] == 1
        m = Int64(round(rcutmax/a[1]))
    end
    if pbc[2] == 1
        n = Int64(round(rcutmax/b[2]))
    end
    if d==3
        if pbc[3] == 1
            p = Int64(round(rcutmax/c[3]))
        end
    end

    return m, n, p
end

function referenceorigins(m, n, p, d)

xref = zeros(d, (2*m+1)*(2*n+1)*(2*p+1))
for i = 1:(2*p+1) # z
    for j = 1:(2*n+1) # y
        for k = 1:(2*m+1) # x
            t = k + (2*m+1)*(j-1) + (2*m+1)*(2*n+1)*(i-1)
            xref[1, t] = k-m-1;
            xref[2, t] = j-n-1;
            if d==3
                xref[3, t] = i-p-1;
            end
        end
    end
end

return xref

end

function latticeorigins(a, b, c, m, n, p)

    d = (length(c) < 3) ? 2 : 3

    xo = referenceorigins(m, n, p, d)

    if d==2
        B2R, R2B = domainmaps(a, b);
    else
        B2R, R2B = domainmaps(a, b, c);
    end

    latorigins = R2B*xo;

    return latorigins
end

function latticecoords(x, rcutmax, pbc, a, b, c)

    m, n, p = numlattices(rcutmax, pbc, a, b, c)

    latorigins = latticeorigins(a, b, c, m, n, p)
    ind = (m+1) + (2*m+1)*(n) + (2*m+1)*(2*n+1)*(p)

    d = size(x,1)
    nx = size(x,2)
    nl = size(latorigins,2)  # number of lattices

    y = zeros(d, nx*nl)
    for j = 1:nx
        for k = 1:d
            y[k,j] = x[k,j]
        end
    end
    q = nx;
    for i = 1:nl
        if i != ind
            for j = 1:nx
                q = q + 1
                for k = 1:d
                    y[k,q] = latorigins[k,i] + x[k,j]
                end
            end
        end
    end

    alist = zeros(Int64,nx*nl)
    for i = 1:nl
        for j = 1:nx
            alist[j + nx*(i-1)] = j
        end
    end

    # get
    neighi, neighj, neighnum = neighborlist(y, rcutmax);

    neighi = neighi[neighi .<= nx];
    ne = length(neighi);

    # list of neighbors for each atom i in x
    neighlist = neighj[1:ne];

    # number of neighbors for each atom i in x
    neighnum = neighnum[1:nx];
    neighnum = [0; cumsum(neighnum)];

    return y, alist, neighlist, neighnum
end


# Using Base modules.
using Random
# Load the distributions library.
using Distributions

mutable struct UGMM_CAVI <: InferenceAlgorithm
    X # Data
    K::Int       # number of iterations
    N::Int
    sigma2
    m
    phi
    s2
end

function UGMM_CAVI(X, num_components::Int, n_samples::Int, sigma2=1)
    K=num_components
    N=n_samples
    rng = RandomDevice()
    phi = rand(Dirichlet(K*N, (rand(rng, Float32,(1))*rand(rng, 1:9))[1]),(K));
    m = rand(rng,floor(minimum(X)):floor(maximum(X)), (1,K))
    m += maximum(X)*rand(rng, Float32, (1,K))
    s2 = rand(rng, Float32, (1, K))
    print("\nInit mean\n")
    println(m)
    print("\nInit s2\n")
    println(s2)
    UGMM_CAVI(X, num_components, n_samples, sigma2, m, phi, s2)

end

function get_elbo(alg::UGMM_CAVI)
    t1 = log.(alg.s2) - alg.m./(alg.sigma2)
    t1 = sum(t1)
    t2 = -0.5.*(reshape(alg.X, (alg.N*alg.K, 1)).^2 * ones(1, alg.K) + ones(alg.N*alg.K, 1)*((alg.s2)+(alg.m).^2 ))
    t2 .+= reshape(alg.X, (alg.N*alg.K, 1))*(alg.m)
    t2 .-= log.(alg.phi)
    t2 .+= alg.phi
    t2 = sum(t2)
    return t1 + t2
end

function _update_phi(alg::UGMM_CAVI)
    t1 = reshape(alg.X, (alg.N*alg.K, 1))*alg.m
    t2 = -(0.5*alg.m.^2 + 0.5*alg.s2)
    exponent = t1 .+ t2
    phi = exp.(exponent)
    alg.phi = phi ./ sum(phi,dims=2)
end
function _update_mu(alg::UGMM_CAVI)
    alg.m = sum((alg.phi.*reshape(alg.X, (alg.N*alg.K, 1))),dims=1) .* (1 ./alg.sigma2 .+ sum(alg.phi,dims=1)).^(-1)
    @assert size(alg.m)[2] == alg.K
    #print(self.m)
    alg.s2 = (1 ./alg.sigma2 .+ sum(alg.phi,dims=1)).^(-1)
    @assert size(alg.s2)[2] == alg.K
end

function _cavi_step(alg::UGMM_CAVI)
    _update_phi(alg)
    _update_mu(alg)
end

function fit(alg::UGMM_CAVI; max_iter=100, tol=1e-20, print_interval=50)
    elbo_values = [get_elbo(alg)]
    m_history = [alg.m]
    s2_history = [alg.s2]
    for iter_ in 1:(max_iter+1)
        _cavi_step(alg)
        append!(m_history, [alg.m])
        append!(s2_history, [alg.s2])
        append!(elbo_values, [get_elbo(alg)])

        if iter_ % print_interval == 0
            println(iter_, " Means: ", m_history[iter_], " ELBO: ", elbo_values[end])
        end
        if abs(elbo_values[end-1] - elbo_values[end]) < tol
            i=iter_
            elbo_val = elbo_values[end]
            println("ELBO converged with $elbo_val at iteration $i")
            break
        end
        if iter_ == max_iter
            elbo_val = elbo_values[end]
            print("ELBO ended with ll $elbo_val")
        end
    end
    elbo_values, m_history ,s2_history
end


#=
This code slightly modified from original code by Luke Stagner (license below):

The MIT License (MIT)
Copyright (c) 2015 Luke Stagner
Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

Original code found at: https://gist.github.com/lstagner/83f91cc138d4767523a3
=#
# module CubicSplines
#
# using Parameters
# using LinearAlgebra

export CubicSpline, MonotoneCubicSpline, interpolate, deriv, hess, extrapolate

struct CubicSpline{T} # Cubic Hermite Spline
    n::Int        # Number of knots
    x::Vector{T} # x position of knots
    y::Vector{T} # y position of knots
    m::Vector{T} # Tangent at knots
end

function CubicSpline(x::Vector{T}, y::Vector{T}) where {T<:Real}
    issorted(x) || throw(ArgumentError("x points must be in ascending order"))
    nx = length(x)
    m = zeros(nx)
    m[1] = (y[2] - y[1]) / (x[2] - x[1])
    for i = 2:nx-1
        m[i] =
            0.5 * (
                (y[i+1] - y[i]) / (x[i+1] - x[i]) +
                (y[i] - y[i-1]) / (x[i] - x[i-1])
            )
    end
    m[nx] = (y[nx] - y[nx-1]) / (x[nx] - x[nx-1])
    return CubicSpline(nx, x, y, m)
end
"""
  specify tangents explicitly
"""
function CubicSpline(x::Vector{T}, y::Vector{T}, m::Vector{T}) where {T<:Real}
    issorted(x) || throw(ArgumentError("x points must be in ascending order"))
    nx = length(x)
    return CubicSpline(nx, x, y, m)
end

function MonotoneCubicSpline(x::Vector{T}, y::Vector{T}) where {T<:Real}
    issorted(x) || throw(ArgumentError("x points must be in ascending order"))
    nx = length(x)
    m = zeros(nx)
    m[1] = (y[2] - y[1]) / (x[2] - x[1])

    for i = 2:nx-1
        hi = x[i+1] - x[i]
        hi_1 = x[i] - x[i-1]
        di = (y[i+1] - y[i]) / hi
        di_1 = (y[i] - y[i-1]) / hi_1
        m[i] = sign(di) != sign(di_1) ? 0.0 :
            3 *
        (hi_1 + hi) *
        (((2 * hi + hi_1) / di_1) + ((hi + 2 * hi_1) / di))^(-1.0)
    end

    m[nx] = (y[nx] - y[nx-1]) / (x[nx] - x[nx-1])

    return CubicSpline(nx, x, y, m)
end

function interpolate(S::CubicSpline, x::Real)
    xr = searchsorted(S.x, x)
    i1 = max(1, xr.stop)
    i2 = min(length(S.x), xr.start)
    if i1 != i2
        dx = (S.x[i2] - S.x[i1])
        t = (x - S.x[i1]) / dx
        h00 = 2 * t^3 - 3 * t^2 + 1
        h10 = t^3 - 2 * t^2 + t
        h01 = -2 * t^3 + 3 * t^2
        h11 = t^3 - t^2
        return h00 * S.y[i1] +
               h10 * dx * S.m[i1] +
               h01 * S.y[i2] +
               h11 * dx * S.m[i2]
    else
        return S.y[i1]
    end
end
interpolate(S::CubicSpline, x::Vector{T}) where {T<:Real} =
    map(t -> interpolate(S, t), x)

function interpolate(x1,x2,t)
    (1-t)*x1 .+ t*x2
end
function interpolate(x::Vector{T}, t::Real) where {T}
    idx = searchsorted(0:length(x)-1,t)
    # i = min(max(idx.stop,1),length(x)-1)
    # j = max(min(idx.start,length(x)),min(length(x),2))
    # For clipped interpolation:
    i = min(max(idx.stop,1),length(x))
    j = min(max(idx.start,1),length(x))
    interpolate(x[i],x[j],t+1-i)
end

function deriv(S::CubicSpline, x::Real)
    xr = searchsorted(S.x, x)
    i1 = max(1, xr.stop)
    i2 = min(length(S.x), xr.start)
    if i1 != i2
        dx = (S.x[i2] - S.x[i1])
        t = (x - S.x[i1]) / dx
        h00 = (6 * t^2 - 6 * t) / dx
        h10 = (3 * t^2 - 4 * t + 1) / dx
        h01 = (-6 * t^2 + 6 * t) / dx
        h11 = (3 * t^2 - 2 * t) / dx
        return h00 * S.y[i1] +
               h10 * dx * S.m[i1] +
               h01 * S.y[i2] +
               h11 * dx * S.m[i2]
    else
        return S.m[i1]
    end
end
deriv(S::CubicSpline, x::Vector{T}) where {T<:Real} = map(t -> deriv(S, t), x)

function hess(S::CubicSpline, x::Real)
    xr = searchsorted(S.x, x)
    i1 = max(1, xr.stop)
    i2 = min(length(S.x), xr.start)
    if i1 != i2
        dx = (S.x[i2] - S.x[i1])
        t = (x - S.x[i1]) / dx
        h00 = (12 * t - 6) / (dx * dx)
        h10 = (6 * t - 4) / (dx * dx)
        h01 = (-12 * t + 6) / (dx * dx)
        h11 = (6 * t - 2) / (dx * dx)
        return h00 * S.y[i1] +
               h10 * dx * S.m[i1] +
               h01 * S.y[i2] +
               h11 * dx * S.m[i2]
    else
        return 0.0
    end
end
hess(S::CubicSpline, x::Vector{T}) where {T<:Real} = map(t -> hess(S, t), x)

function extrapolate(
    S::CubicSpline{T},
    x::Union{T,Vector{T}},
    order::Int = 0,
) where {T<:Real}
    nx = length(x)
    yout = zeros(nx)
    if order == 0
        coeffs = [0.0, 0.0]
    elseif order == 1
        coeffs = [1.0, 0.0]
    elseif order >= 2
        coeffs = [1.0, 1 / 2]
    end
    xrange = extrema(S.x)
    for (i, xi) in enumerate(x)
        if (xi < xrange[1] || xi > xrange[2])
            if xi < xrange[1]
                idx = 1
            elseif xi > xrange[2]
                idx = length(S.x)
            end
            dy_dx = deriv(S, S.x[idx])[1]
            d2y_dx2 = hess(S, S.x[idx])[1]
            dx = (x[i] - S.x[idx])
            dy = coeffs[1] * dy_dx * dx + coeffs[2] * d2y_dx2 * dx^2
            yout[i] = S.y[idx] + dy # first order extension
        else
            # normal case
            yout[i] = interpolate(S, xi)[1]
        end
    end
    return yout
end

export ParametricSpline

@with_kw struct ParametricSpline{N,T}
    splines::NTuple{N,CubicSpline{T}} = NTuple{0,CubicSpline{Float64}}()
end
function ParametricSpline(t::Vector{T}, x::Array{T,2}) where {T<:Real}
    ParametricSpline(tuple(map(i -> CubicSpline(t, x[:, i]), 1:size(x, 2))...))
end
function ParametricSpline(
    t::Vector{T},
    x::Array{T,2},
    v::Array{T,2},
) where {T<:Real} # with prescribed knots
    ParametricSpline(tuple(map(
        i -> CubicSpline(t, x[:, i], v[:, i]),
        1:size(x, 2),
    )...))
end
interpolate(S::ParametricSpline, t::T) where {T<:Real} =
    [map(s -> interpolate(s, t), S.splines)...]
deriv(S::ParametricSpline, t::T) where {T<:Real} =
    [map(s -> deriv(s, t), S.splines)...]
hess(S::ParametricSpline, t::T) where {T<:Real} =
    [map(s -> hess(s, t), S.splines)...]

interpolate(S::ParametricSpline, t::T) where {T<:Vector} =
    map(k -> interpolate(S, k), t)
deriv(S::ParametricSpline, t::T) where {T<:Vector} = map(k -> deriv(S, k), t)
hess(S::ParametricSpline, t::T) where {T<:Vector} = map(k -> hess(S, k), t)

# interpolate(S::P,t::V) where {P<:ParametricSpline,V<:Vector} = map(k->interpolate(S,k), t)
# deriv(S::P,t) where {P<:ParametricSpline}                 = map(s->deriv(s,t), S.splines)
# hess(S::P,t) where {P<:ParametricSpline}                  = map(s->hess(s,t), S.splines)
#
# interpolate(S::P,t) where {P<:ParametricSpline}           = map(k->interpolate(S,k), S.splines)
# deriv(S::P,t) where {P<:ParametricSpline}                 = map(s->deriv(s,t), S.splines)
# hess(S::P,t) where {P<:ParametricSpline}                  = map(s->hess(s,t), S.splines)
export tangent_rate, curve_turn_rate

"""
    `tangent_rate(S::ParametricSpline,t::Real)`
    Computes the rate (w.r.t. t) at which the tangent vector to a curve is
    changing.
"""
function tangent_rate(S::ParametricSpline, t::Real)
    dr = deriv(S, t) # position first derivative w.r.t time
    ddr = hess(S, t) # position first derivative w.r.t time
    ds = sqrt(dot(dr, dr)) # arc length first derivative
    dds = sum(dr) / ds # arc length second derivative
    dT = (1 / (ds^2)) * (ds * ddr .- dds * dr)
end

"""
    `curve_turn_rate(S::ParametricSpline,t::Real)`
    computes the turn rate (w.r.t. t) of the tangent vector to a curve.
"""
function curve_turn_rate(S::ParametricSpline, t::Real)
    dr = map(s -> deriv(s, t), S.splines) # position first derivative w.r.t t
    dT = tangent_rate(S, t) # tangent vector first derivative w.r.t. t
    ω = dr × dT
end

extrapolate(S::ParametricSpline, t::T, order = 0) where {T<:Real} =
    map(s -> extrapolate(s, t, order)[1], S.splines)
extrapolate(S::ParametricSpline, t_vec::T, order = 0) where {T<:Vector} =
    map(k -> extrapolate(S, k, order), t_vec)

# end # module


using Statistics
import .Plots; const pl = Plots

export plot_gp_intervals

"""
    plot_gp_intervals(gp::GaussianProcess, x::AbstractMatrix,
                      y::Union{AbstractMatrix,Nothing} = nothing; input_dim = 1,
                      output_dim = 1)

Creates a simple 95% uncertainty interval plot for the GP at a set of ordered
data points ``x``. Optionally also plots a set of data points ``y``. For
multivariateand vector GPs, plots the dimensions specified by
`input_dim` and `output_dim`.
"""
function plot_gp_intervals(gp::GaussianProcess, x::AbstractMatrix, y::Union{AbstractMatrix,Nothing} = nothing; input_dim = 1, output_dim = 1)
  f = gp(x)[output_dim,:,:]
  n_samples = min(size(f,ndims(f)),32)

  x_d = x[input_dim,:]
  m = mean(f;dims=2) |> vec
  u = mapslices(v->quantile(v,0.975), f; dims=2) |> vec
  u_2sd = mean(f;dims=2) + mapslices(v->1.96*sqrt(var(v)), f; dims=2) |> vec
  l = mapslices(v->quantile(v,0.025), f; dims=2) |> vec
  l_2sd = mean(f;dims=2) - mapslices(v->1.96*sqrt(var(v)), f; dims=2) |> vec

  p = pl.plot(size=(800,600))

  y isa Nothing ? () : pl.scatter!(p, x_d, y[output_dim,:], label=nothing)
  pl.plot!(p, x_d, m, ribbon = (abs.(l.-m),abs.(u.-m)), label=nothing)
  for i in 1:n_samples
    pl.plot!(p, x_d, f[:,i], color=:gray, label=nothing)
  end

  p
end