"""
    KmeansAlg(m::Int, metric::SemiMetric=SqEuclidean(); nMarkov = 10, tol = 1e-3)

k-Means [1] initialization on the data `X` taking `m` inducing points.
The seeding is computed via [2], `nMarkov` gives the number of MCMC steps for the seeding.

## Arguments
- `k::Int` : Number of inducing points
- `metric::SemiMetric` : Metric used to compute the distances for the k-means algorithm

## Keyword Arguments
- `nMarkov::Int` : Number of random steps for the seeding
- `tol::Real` : Tolerance for the `kmeans` algorithm

[1] Arthur, D. & Vassilvitskii, S. k-means++: The advantages of careful seeding. in Proceedings of the eighteenth annual ACM-SIAM symposium on Discrete algorithms 1027–1035 (Society for Industrial and Applied Mathematics, 2007).
[2] Bachem, O., Lucic, M., Hassani, S. H. & Krause, A. Fast and Provably Good Seedings for k-Means. Advances in Neural Information Processing Systems 29 55--63 (2016) doi:10.1109/tmtt.2005.863818.
"""
struct KmeansAlg{Tm,T} <: OffIPSA
    m::Int
    metric::Tm
    nMarkov::Int
    tol::T
    function KmeansAlg(
        m::Int, metric::SemiMetric=SqEuclidean(); nMarkov::Int=10, tol::T=1e-3
    ) where {T<:Real}
        m > 0 || throw(ArgumentError("The number of inducing points m should be positive"))
        tol > 0 || throw(ArgumentError("The tolerance tol should be positive"))
        return new{typeof(metric),T}(m, metric, nMarkov, tol)
    end
end

function inducingpoints(
    rng::AbstractRNG, alg::KmeansAlg, X::AbstractVector; weights=nothing, kwargs...
)
    if alg.m >= length(X)
        return edge_case(alg.m, length(X), X)
    end
    return kmeans_ip(
        rng, X, alg.m, alg.metric; nMarkov=alg.nMarkov, weights=weights, tol=alg.tol
    )
end

function Base.show(io::IO, alg::KmeansAlg)
    return print(io, "k-Means Selection of Inducing Points (m : $(alg.m))")
end

#Return K inducing points from X, m being the number of Markov iterations for the seeding
function kmeans_ip(
    rng::AbstractRNG,
    X::AbstractVector,
    nC::Int,
    metric::SemiMetric;
    nMarkov::Int=10,
    weights=nothing,
    tol=1e-3,
)
    C = kmeans_seeding(rng, X, nC, metric, nMarkov)
    C = reduce(hcat, C)
    kmeans!(reduce(hcat, X), C; weights=weights, tol=tol, distance=metric)
    return ColVecs(C)
end

"""Fast and efficient seeding for KMeans based on [`Fast and Provably Good Seeding for k-Means](https://las.inf.ethz.ch/files/bachem16fast.pdf)"""
function kmeans_seeding(
    rng::AbstractRNG,
    X::AbstractVector, # Input data
    nC::Integer, # Number of centroids
    metric::SemiMetric, # Metric used
    nMarkov::Integer, # Number of Markov iterations
)
    nSamples = length(X) # Number of input samples
    # Preprocessing,
    init = rand(rng, 1:nSamples) # Sample first random center
    C = collect.(X[init:init]) # Initialize the collection of centroids
    q = vec(pairwise(metric, X, C)) # Create the pairwise values between the data and the first centroid
    q = Weights(q / sum(q) .+ 1.0 / (2 * nSamples), 1) # Create weights to work with
    for i in 2:nC
        x = X[sample(rng, q)] # weighted sampling,
        mindist = mindistance(metric, x, C) # Find the closest centroid to the random sample x
        for j in 2:nMarkov # Iterate over nMarkov iterations
            y = X[sample(rng, q)] # Draw a new sample
            dist = mindistance(metric, y, C) # Find the closest centroid to the random sample y
            if (dist / mindist) > rand(rng) # Acceptance step to see if y is better than x
                x = y
                mindist = dist
            end
        end
        push!(C, x)
    end
    return C
end


# Load Julia packages (libraries)

using Pkg, DrWatson

using MonteCarloMeasurements
using StanSample
using StatisticalRethinking

df = CSV.read(sr_datadir("chimpanzees.csv"), DataFrame);

# Define the Stan language model

stan10_1s = "
data{
    int N;
    int pulled_left[N];
}
parameters{
    real a;
}
model{
    real p;
    a ~ normal( 0 , 10 );
    pulled_left ~ binomial( 1 , inv_logit(a) );
}
";

# Define the Stanmodel and set the output format to :mcmcchains.

m10_1s = SampleModel("m10.1s", stan10_1s);

# Input data for cmdstan

m10_1_data = Dict("N" => size(df, 1),
  "pulled_left" => df[!, :pulled_left]);

# Sample using cmdstan

rc10_1s = stan_sample(m10_1s, data=m10_1_data);

# Result rethinking

rethinking = "
  mean   sd  5.5% 94.5% n_eff Rhat
a 0.32 0.09 0.18  0.46   166    1
";

# Describe the draws

if success(rc10_1s)
    part10_1s = read_samples(m10_1s, :particles)
    part10_1s |> display
end


module ProperOrthogonalDecomposition
using Statistics
using LinearAlgebra

export  PODBasis,
        POD,
        PODsvd,
        PODsvd!,
        PODeigen,
        PODeigen!,
        modeConvergence,
        modeConvergence!

include("types.jl")
include("convergence.jl")
include("PODsvd.jl")
include("PODeigen.jl")

"""
    POD(X)

Computes the Proper Orthogonal Decomposition of `X`. Returns a PODBasis.

This method is just a thin wrapper around `PODsvd(X)`.
"""
function POD(X; subtractmean::Bool = false)
    PODsvd(X, subtractmean = subtractmean)
end

"""
    POD(X, W)

Computes the Proper Orthogonal Decomposition of `X` with weights `W`. Returns a PODBasis.

This method is just a thin wrapper around `PODsvd(X)`.
"""
function POD(X, W; subtractmean::Bool = false)
    PODsvd(X, W, subtractmean = subtractmean)
end

end # module


function compute_bc(params, α, n)
    (; r, D, W, dim) = params
    m = length(D)

    bc = zeros(m, 2)
    bc[:, 1] .= 1

    if α == 0
        return bc
    end

    ρ = r[1:m-1] ./ .√D[1:m-1]
    ρt = r[1:m-1] ./ .√D[2:m]

    for i = 1:m-1
        J, Y, dJ, dY = compute_JY(α * ρ[i], n, dim)
        J1, Y1, dJ1, dY1 = compute_JY(α * ρt[i], n, dim)

        A11 = -√(D[i] / D[i+1]) * dJ * Y1 + J * dY1 + √D[i] / W[i] * α * dJ * dY1
        A12 = -√(D[i] / D[i+1]) * dY * Y1 + Y * dY1 + √D[i] / W[i] * α * dY * dY1
        A21 = √(D[i] / D[i+1]) * dJ * J1 - J * dJ1 - √D[i] / W[i] * α * dJ * dJ1
        A22 = √(D[i] / D[i+1]) * dY * J1 - Y * dJ1 - √D[i] / W[i] * α * dY * dJ1

        z = α * ρt[i]
        if dim == 2
            Q = z * π / 2
        else
            Q = z^2
        end
        A = Q * [A11 A12; A21 A22]
        bc[i+1, :] = A * bc[i, :]
    end

    bc
end


include("euler/euler.jl")
using .BigIntegers: BigInteger
using .Numbers: digits_sum

compute(n::Int)::Int = digits_sum(prod(map(BigInteger, 1:n)))


compute(n::Int, k::Int)::Float64 = round(7(1 - prod(1 .- k ./ (6n ÷ 7 + 1:n))), digits=9)


function vector_to_symmetric_matrix_converter()
    E1 = [
        1.0 0.0
        0.0 0.0
        0.0 1.0
    ]
    E2 = [
        0.0 0.0
        0.0 1.0
        1.0 0.0
    ]
    return [E1, E2]
end

function make_row_matrix(matrix, vals)
    return hcat([v * matrix for v in vals]...)
end

function interpolation_matrix(vals, ndofs)
    return make_row_matrix(diagm(ones(ndofs)), vals)
end

function transform_gradient(gradf, jacobian)
    return gradf / Diagonal(jacobian)
end

function bilinear_form(basis, quad, stiffness, jacobian)
    dim = dimension(basis)
    nf = number_of_basis_functions(basis)
    ndofs = dim * nf
    matrix = zeros(ndofs, ndofs)
    vectosymmconverter = vector_to_symmetric_matrix_converter()
    detjac = prod(jacobian)
    for (p, w) in quad
        grad = transform_gradient(gradient(basis, p), jacobian)
        NK = zeros(3, 2nf)
        for k = 1:dim
            NK .+= make_row_matrix(vectosymmconverter[k], grad[:, k])
        end
        matrix .+= NK' * stiffness * NK * detjac * w
    end
    return matrix
end

function bilinear_form(basis, quad, stiffness, cellmap::CellMap)
    return bilinear_form(basis, quad, stiffness, jacobian(cellmap))
end

function mass_matrix(basis, quad, detjac::R, ndofs) where {R<:Real}
    nf = number_of_basis_functions(basis)
    totaldofs = ndofs * nf
    matrix = zeros(totaldofs, totaldofs)
    for (p, w) in quad
        vals = basis(p)
        N = interpolation_matrix(vals, ndofs)
        matrix .+= N' * N * detjac * w
    end
    return matrix
end

function mass_matrix(basis, quad, cellmap::CellMap, ndofs)
    detjac = determinant_jacobian(cellmap)
    return mass_matrix(basis, quad, detjac, ndofs)
end

function mass_matrix(basis, quad, scale::V, ndofs) where {V<:AbstractVector}
    nf = number_of_basis_functions(basis)
    totaldofs = ndofs * nf
    matrix = zeros(totaldofs, totaldofs)
    @assert length(scale) == length(quad)
    for (idx, (p, w)) in enumerate(quad)
        vals = basis(p)
        N = interpolation_matrix(vals, ndofs)
        matrix .+= N' * N * scale[idx] * w
    end
    return matrix
end

function component_mass_matrix(basis, quad, component, scale)
    nf = number_of_basis_functions(basis)
    dim = dimension(basis)
    @assert dim == length(component)
    totaldofs = dim * nf
    matrix = zeros(totaldofs, totaldofs)
    projector = component * component'
    for (p, w) in quad
        vals = basis(p)
        NI = interpolation_matrix(vals, dim)
        NP = make_row_matrix(projector, vals)
        matrix .+= NI' * NP * scale * w
    end
    return matrix
end

function face_traction_operator(basis, quad, normal, stiffness, facescale, cellmap)

    dim = dimension(basis)
    @assert length(normal) == dim
    nf = number_of_basis_functions(basis)
    ndofs = dim * nf
    matrix = zeros(ndofs, ndofs)

    vectosymmconverter = vector_to_symmetric_matrix_converter()
    N = sum([normal[k] * vectosymmconverter[k]' for k = 1:dim])
    jac = jacobian(cellmap)

    for (p, w) in quad
        grad = transform_gradient(gradient(basis, p), jac)
        vals = basis(p)

        NK = zeros(3, 2nf)
        for k = 1:dim
            NK .+= make_row_matrix(vectosymmconverter[k], grad[:, k])
        end
        NI = interpolation_matrix(vals, dim)
        matrix .+= NI' * N * stiffness * NK * facescale * w
    end
    return matrix
end

function face_traction_component_operator(
    basis,
    quad,
    component,
    normal,
    stiffness,
    facescale,
    cellmap,
)
    dim = dimension(basis)
    @assert length(normal) == dim
    @assert length(component) == dim

    nf = number_of_basis_functions(basis)
    ndofs = dim * nf
    matrix = zeros(ndofs, ndofs)

    vectosymmconverter = vector_to_symmetric_matrix_converter()
    N = sum([normal[k] * vectosymmconverter[k]' for k = 1:dim])
    projector = component * component'
    jac = jacobian(cellmap)

    for (p, w) in quad
        grad = transform_gradient(gradient(basis, p), jac)
        vals = basis(p)

        NK = sum([make_row_matrix(vectosymmconverter[k], grad[:, k]) for k = 1:dim])
        NI = interpolation_matrix(vals, dim)

        matrix .+= NI' * projector * N * stiffness * NK * facescale * w
    end
    return matrix
end

function component_linear_form(rhsfunc, basis, quad, component, cellmap, detjac)
    dim = dimension(basis)
    @assert length(component) == dim
    nf = number_of_basis_functions(basis)
    rhs = zeros(dim * nf)
    for (p, w) in quad
        rhsval = rhsfunc(cellmap(p))
        NI = interpolation_matrix(basis(p), dim)
        rhs .+= rhsval * NI' * component * detjac * w
    end
    return rhs
end

function component_linear_form(rhsval, basis, quad, component, detjac)
    dim = dimension(basis)
    @assert length(component) == dim
    nf = number_of_basis_functions(basis)
    rhs = zeros(dim * nf)
    for (p, w) in quad
        NI = interpolation_matrix(basis(p), dim)
        rhs .+= rhsval * NI' * component * detjac * w
    end
    return rhs
end

function linear_form(rhsfunc, basis, quad, cellmap)
    dim = dimension(basis)
    nf = number_of_basis_functions(basis)
    rhs = zeros(dim * nf)
    detjac = determinant_jacobian(cellmap)
    for (p, w) in quad
        vals = rhsfunc(cellmap(p))
        N = interpolation_matrix(basis(p), dim)
        rhs .+= N' * vals * detjac * w
    end
    return rhs
end

function linear_form(rhsfunc, basis, quad, cellmap, detjac)
    dim = dimension(basis)
    nf = number_of_basis_functions(basis)
    rhs = zeros(dim * nf)
    for (p, w) in quad
        vals = rhsfunc(cellmap(p))
        N = interpolation_matrix(basis(p), dim)
        rhs .+= N' * vals * detjac * w
    end
    return rhs
end

function constant_linear_form(rhsval, basis, quad, detjac)
    ndofs = length(rhsval)
    nf = number_of_basis_functions(basis)
    rhs = zeros(ndofs * nf)
    for (p, w) in quad
        vals = basis(p)
        NI = interpolation_matrix(basis(p), ndofs)
        rhs .+= NI' * rhsval * detjac * w
    end
    return rhs
end


function pow1_test()
    #=
        max z
    s.t.
        x^0.9 * y^(0.1) >= |z| (i.e (x, y, z) are in the 3d power cone with a=0.9)
        x == 2
        y == 1
    =#
    model = TestModel{Float64}()

    a = 0.9
    v = MOI.add_variables(model, 3)
    @test MOI.get(model, MOI.NumberOfVariables()) == 3

    vov = MOI.VectorOfVariables(v)

    vc = MOI.add_constraint(model, vov, MOI.PowerCone(a));

    cx = MOI.add_constraint(model, MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(1.0, v[1])], 0.), MOI.EqualTo(2.))
    cy = MOI.add_constraint(model, MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(1.0, v[2])], 0.), MOI.EqualTo(1.))

    MOI.set(model, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(), MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(1.0, v[3])], 0.0))
    MOI.set(model, MOI.ObjectiveSense(), MOI.MAX_SENSE)

    return model
end

function pow2_test()
    #=
        max z
    s.t.
        x^0.9 * y^(0.1) >= |z| (i.e (x, y, z) are in the 3d power cone with a=0.9)
        x == 2
        y == 1
    =#
    a = 0.9
    model = TestModel{Float64}()
    v = MOI.add_variables(model, 3)
    @test MOI.get(model, MOI.NumberOfVariables()) == 3

    vov = MOI.VectorOfVariables(v)

    vc = MOI.add_constraint(model, MOI.VectorAffineFunction{Float64}(vov), MOI.PowerCone(a))

    cx = MOI.add_constraint(model, MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(1.0, v[1])], 0.), MOI.EqualTo(2.))
    cy = MOI.add_constraint(model, MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(1.0, v[2])], 0.), MOI.EqualTo(1.))

    MOI.set(model, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(), MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(1.0, v[3])], 0.0))
    MOI.set(model, MOI.ObjectiveSense(), MOI.MAX_SENSE)

    return model
end


using Test, PolynomialBases
using LinearAlgebra

ufunc(x) = sinpi(x)
uprim(x) = π*cospi(x)

function tolerance(p, T=Float64)
    if p <= 2
        5.
    elseif p <= 6
        0.5
    elseif p <= 8
        0.04
    elseif p <= 10
        0.002
    elseif p <= 12
        if T==Float32
            5.e-5
        else
            3.e-5
        end
    elseif p <= 14
        if T==Float32
            3.e-5
        else
            6.e-7
        end
    elseif p <= 16
        if T==Float32
            4.e-5
        else
            7.e-9
        end
    else
        if T==Float32
            6.e-5
        else
            7.e-11
        end
    end
end

# regression test
for basis_type in subtypes(PolynomialBases.NodalBasis{PolynomialBases.Line})
    println(devnull, "  ", basis_type(1))
    basis_type <: ClosedNewtonCotes && continue
    for p in 5:20, T in (Float32, Float64)
        basis = basis_type(p, T)
        u = ufunc.(basis.nodes)
        @test basis.D * u ≈ uprim.(basis.nodes) atol=tolerance(p, T)
    end
end

# compare direct derivative evaluation and matrix
for basis_type in subtypes(PolynomialBases.NodalBasis{PolynomialBases.Line})
    for p in 0:15
        basis = try
            basis_type(p)
        catch m
            isa(m, DomainError) && continue
            throw(m)
        end
        u = ufunc.(basis.nodes)
        xplot = range(-1, stop=1, length=100)
        u1 = derivative_at(xplot, u, basis)
        u2 = interpolate(xplot, basis.D*u, basis)
        @test u1 ≈ u2
    end

    for p in 0:15
        basis = try
            basis_type(p, Float32)
        catch m
            isa(m, DomainError) && continue
            throw(m)
        end
        u = ufunc.(basis.nodes)
        xplot = range(-1, stop=1, length=100)
        u1 = derivative_at(xplot, u, basis)
        u2 = interpolate(xplot, basis.D*u, basis)
        @test norm(u1 - u2, Inf) < 5.f-3
    end
end


using Base.Threads

function logsoftmax!(out::AbstractVecOrMat, xs::AbstractVecOrMat)
  @threads for j = 1:size(xs, 2)
    @inbounds begin
      xi_max = xs[1, j]
      for i = 1:size(out, 1)
        xi_max = max(xi_max, xs[i, j])
      end
      s = zero(eltype(out))
      for i = 1:size(out, 1)
        s += exp(xs[i, j] - xi_max)
      end
      for i = 1:size(out, 1)
        out[i, j] = xs[i, j] - log(s) - xi_max
      end
    end
  end
  out
end

logsoftmax!(xs) = logsoftmax!(xs, xs)
logsoftmax(xs) = logsoftmax!(similar(xs), xs)

∇logsoftmax(Δ, xs) = ∇softmax(Δ ./ max.(eps(eltype(xs)),softmax(xs)), xs)

"""
    logsoftmax(xs) = log.(exp.(xs) ./ sum(exp.(xs)))

logsoftmax computes the log of softmax(xs) and it is more numerically stable
than softmax function in computing the cross entropy loss.
"""
logsoftmax


using Distributions

# walker type, with a position and if it is a wanderer
#  mutable struct walker
#      x::Float64
#      y::Float64
#      loner::Bool
#  end

"""
    distance(walker1,walker2)

Computes de distance between two walkers.
"""
function distance(w1::walker,w2::walker)
    on_x = (w1.x - w2.x)^2
    on_y = (w1.y - w2.y)^2
    sqrt(on_x+on_y)
end

"""
    freeze(walker1,walker2,tol)

Changes the loner state of a walker if it's distance
to another one is < tol
"""
function freeze(w1::walker,w2::walker,tol::Float64)
    if w1.loner == false || w2.loner == false
        if distance(w1,w2) <= tol
            return w1.loner=w2.loner=0
        end
    end
end

"""
    look_around(element,W,tol)

Computes the distance of a walker to all the other walkers
in an array of them and freezes them if their distance to
the nearest one is < tol
"""
function look_around(element::Int64,W::Array,tol::Float64)
    not_me = W[1:end .!= element]
    map(x -> freeze(the_walkers[element],x,tol),not_me)
end

"""
    random_walk(w,min,max)

Makes a walker do a random walk inside a restriction box with
length Δ
"""
function random_walk(w::walker,min::Number,max::Number,Δ=0.1)
    Δx = rand([-Δ,0,Δ])
    Δy = rand([-Δ,0,Δ])
    if w.x + Δx > min && w.x + Δx < max
        w.x = w.x + Δx;
    end
    if w.y + Δy > min && w.y + Δy < max
        w.y = w.y + Δy;
    end
end

"""
    walker_set(count,min_box,max_box)

Gives a set of walkers  bounded in a box
"""
function walker_set(count::Int64,min_box,max_box)
    the_walkers = [walker(round(rand(Uniform(min_box,max_box));digits=1),
                          round(rand(Uniform(min_box,max_box));digits=1),
                          1) for i in 1:count]
    return the_walkers
end

"""
dla(walker_array,tol,min_box,max_box;max_no_mov)

Performs a DLA on a set of walkers in a bounding box,
makes them stop moving if their distance < tol to a
non-moving walker. Stops when you get the maximum of
non-moving ones of desire
"""
function dla(walker_array::Array,tol,min_box,max_box;max_no_mov=0)
    while sum([x.loner for x ∈ walker_array]) > max_no_mov
        global the_movers = findall(x -> x.loner == true, walker_array);
        for i ∈ the_movers
            look_around(i,walker_array,tol)
        end
        the_movers = findall(x -> x.loner == true, the_walkers);
        map(x -> random_walk(walker_array[x],min_box,max_box),the_movers);
    end
    walker_array
end

## Main program
#  the_walkers = [walker(round(rand(Uniform(-5,5));digits=1),
#                        round(rand(Uniform(-5,5));digits=1),
#                        1) for i in 1:5000]
#  unshift!(the_walkers,walker(0.,0.,0))
#
#  while sum([x.loner for x ∈ the_walkers]) > 0
#      global the_movers = findall(x -> x.loner == true, the_walkers);;
#      for i ∈ the_movers
#          look_around(i,the_walkers,0.1)
#      end
#      the_movers = findall(x -> x.loner == true, the_walkers);
#      map(x -> random_walk(the_walkers[x],-5,5),the_movers);
#  end
#
#  ## Plotting
#  the_x = [w.x for w in the_walkers]
#  the_y = [w.y for w in the_walkers]
#  p = scatter(the_x,the_y,markersize=3.,
#              leg=false,ticks=nothing,
#              markerstrokewidth=0.,border=false)
#  png(p,"/figs/initial")


"""
    regularized_fit_data

Returns a dictionary `D` with data for the regularized data fitting
example on page 329 of VMLS.  The items in the dictionary are:
* `D["xtrain"]`: vector of length 10
* `D["ytrain"]`: vector of length 10
* `D["xtest"]`: vector of length 20
* `D["ytest"]`: vector of length 20.

"""
function regularized_fit_data()
    X = [
    1.06830e-01  1.82422e+00;
    2.55780e-01  8.13549e-01;
    2.71120e-01  9.24060e-01;
    3.02990e-01  4.20670e-01;
    5.02040e-01  4.46178e-01;
    6.58450e-01 -3.73407e-02;
    6.84230e-01 -2.10935e-01;
    7.02590e-01 -1.03327e-01;
    7.34060e-01  3.32097e-01;
    9.40350e-01  2.29278e+00
    ];
    xtrain = X[:,1];
    ytrain = X[:,2];
    Xtest = [
    3.96674e-01 -1.05772e+00;
    7.77517e-01  8.79117e-01;
    1.18400e-01  1.98136e+00;
    2.23266e-01  8.67012e-01;
    9.01463e-01  2.13650e+00;
    3.58033e-01 -7.01948e-01;
    2.60402e-01  9.41469e-01;
    8.04281e-01  1.49755e+00;
    6.31664e-01  5.50205e-01;
    1.49704e-01  1.34245e+00;
    5.51333e-01  1.21445e+00;
    6.63999e-01  4.49111e-03;
    1.64948e-01  9.57535e-01;
    6.51698e-01  7.78640e-02;
    1.23026e-01  1.73558e+00;
    3.37066e-01 -3.25244e-01;
    8.32080e-02  2.56555e+00;
    2.04422e-01  1.10081e+00;
    9.78000e-01  2.70644e+00;
    4.03676e-01 -1.10049e+00
    ];
    xtest = Xtest[:,1];
    ytest = Xtest[:,2];
    return Dict([ ("xtrain", xtrain), ("ytrain", ytrain),
        ("xtest", xtest), ("ytest", ytest) ])
end


const L = 2*10^4
iclasslabel = ["Deficient", "Perfect", "Abundant"]
iclass = zeros(Int64, 3)
iclass[1] = one(Int64) #by convention 1 is deficient

for n in 2:L
    if isprime(n)
        iclass[1] += 1
    else
        iclass[sign(divisorsum(n)-n)+2] += 1
    end
end

println("Classification of integers from 1 to ", L)
for i in 1:3
    println("   ", iclasslabel[i], ", ", iclass[i])
end


using Plots
sphere(x) = sum(x.^2)

dpi = 600
x = y = range(-3, stop=3, length=200)

l = Vector{Float64}()
for i in x
   for j in y
       v = [i j]
       append!(l, sphere(v))
   end
end

surface(x,y,l, dpi=dpi)
png("./sphere_surface.png")
contour(x,y,l, dpi=dpi, fill = true)
png("./sphere_contour.png")


@testset "integral, constant" begin
    N = (10,20)
    v₀ = rand()
    val = fill(v₀, N)

    p₀ = [rand(), rand()]
    r = 5.0
    xₙ, xₚ = p₀[1]-r, p₀[1]+r
    yₙ, yₚ = p₀[2]-r, p₀[2]+r

    xbound = range(xₙ, xₚ, length=N[1]+1)
    ybound = range(yₙ, yₚ, length=N[2]+1)
    vxlbound = (xbound, ybound)

    vc = VoxelwiseConstant(val, vxlbound)

    @test integral(vc) ≈ v₀ * (xₚ-xₙ) * (yₚ-yₙ)
    @test (s = integrate(vc, p->norm(p .- p₀)≤r); abs(s.val - v₀ * π * r^2) ≤ s.err)
end  # @testset "integral, constant"

@testset "integral, non-constant" begin
    N = (10,20)
    val = rand(N...)

    p₀ = [rand(), rand()]
    r = 5.0
    xₙ, xₚ = p₀[1]-r, p₀[1]+r
    yₙ, yₚ = p₀[2]-r, p₀[2]+r

    xbound = range(xₙ, xₚ, length=N[1]+1)
    ybound = range(yₙ, yₚ, length=N[2]+1)
    vxlbound = (xbound, ybound)

    ∆x = (xₚ-xₙ) / N[1]
    ∆y = (yₚ-yₙ) / N[2]

    vc = VoxelwiseConstant(val, vxlbound)

    @test integral(vc) ≈ sum(val) * ∆x * ∆y
end  # @testset "integral, non-constant"


export LPLoss, LPPenalty,
        L1Loss, L1Penalty,
        L2Loss, L2Penalty,
        LogisticLoss, MultinomialLoss

# = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =
# No Loss / No Penalty
# = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =

(l::NoLoss)(a::AVR, b::AVR) = 0.0
(p::NoPenalty)(θ::AVR)      = 0.0

# = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =
# LP Losses and Penalties
# = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =

"""
$TYPEDEF

Scaled L-p loss of the residual.

``L(x, y) = ||x-y||_{p}^{p}/p``

The scaling simplifies expressions in the common L2 case.
"""
struct LPLoss{p} <: AtomicLoss where p <: Real end

"""
$TYPEDEF

Scaled L-p norm of the parameter vector.

``P(θ) = ||θ||_{p}^{p}/p``

The scaling simplifies expressions in the common L2 case.
"""
struct LPPenalty{p} <: AtomicPenalty where p <: Real end

# Useful Shortcuts
const L1Loss    = LPLoss{1}
const L1Penalty = LPPenalty{1}
const L2Loss    = LPLoss{2}
const L2Penalty = LPPenalty{2}
const LPCost{p} = Union{LPLoss{p},LPPenalty{p}}

const L1R = ScaledPenalty{L1Penalty}
const L2R = Union{NoPenalty,ScaledPenalty{L2Penalty}}
const ENR = Union{L1R,CompositePenalty}

"""
$SIGNATURES

Return the `p` in an `LPCost{p}`.
"""
getp(lpc::LPCost{p}) where p = p

"""
$SIGNATURES

Compute the lp norm to the `p`-th power of a vector given `p` scaled by `p`.
"""
function lp(v::AbstractVector{<:Real}, p)
    p == Inf && return maximum(v)
    p == 1   && return sum(abs.(v))
    p == 2   && return sum(abs2.(v)) / 2
    p  > 0   && return sum(abs.(v).^p) / p
    throw(DomainError("[lp] `p` has to be greater than 0"))
end

(l::LPLoss)(a::AVR, b::AVR) = lp(a .- b, getp(l))
(l::LPLoss)(r::AVR)         = lp(r, getp(l))
(p::LPPenalty)(θ::AVR)      = lp(θ, getp(p))

# = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =
# Logistic loss
# = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =

"""
$TYPEDEF

``L(x, y) = -∑logσ(xᵢyᵢ)``

where `logσ` is the log of the sigmoid function; `yᵢ ∈ {±1}`. In a logistic regression `x`
corresponds to `Xθ` where `X` is the design matrix and `θ` the vector of parameters.
See [`logsigmoid`](@ref).
"""
struct LogisticLoss <: AtomicLoss end

(::LogisticLoss)(x::AVR, y::AVR) = -sum(logsigmoid.(x .* y))

"""
$TYPEDEF

``L(P, y) = ∑log Zᵢ - ∑∑ 1(yᵢ=j)Pᵢⱼ``

where `P` is a matrix where each row contains class probabilities, `yᵢ ∈ {1, 2, ..., K}`
corresponding to column indices and,

``Zᵢ = ∑ exp(Pᵢ)``

In a multinomial regression, `P` corresponds to `XW` where `X` is the design matrix and `W` the
matrix of size `p * K` where each column corresponds to the parameters corresponding to that class.
"""
struct MultinomialLoss <: AtomicLoss end

(::MultinomialLoss)(P::Matrix{<:Real}, y::Vector{Int}) = begin
    L = 0.0
    @inbounds for i in eachindex(y)
        Pᵢ = P[i, :]
        m  = maximum(Pᵢ)
        sᵢ = sum(exp.(Pᵢ .- m)) # avoid overflow
        logZᵢ = log(sᵢ) + m
        L += logZᵢ - P[i, y[i]]
    end
    return L
end
(l::MultinomialLoss)(y::Vector{Int}, P::Matrix{<:Real}) = l(P, y)


using CreditMetrics
using Statistics

# ========================================================
# Using CreditMetrics for simulating a single counterparty
# ========================================================

# Probabilities of being in each rating class after a period
# Note that ProbabilityVector is very unforgiving when checking if probabilities sum to one
pv = ProbabilityVector([0.05, 0.15, 0.25, 0.25, 1.0 - (0.05 + 0.15 + 0.25 + 0.25)])

# Value of the exposure to the counterparty in each rating class after a period
vals = [-1.0, -0.2, -0.1, 0.0, 0.05]

# ---------------------------------------------
# Random outcomes without macroeconomic effects
# ---------------------------------------------

# Compute a vector of values in a number of scenarios
my_base_outcomes = outcome_value(pv, vals, rand(1_000_000))

# Analytically predicted mean value from probabilities and values
meanvalue(pv, vals)

# Observed mean value is the same
mean(my_base_outcomes)

# Observed standard deviation
std(my_base_outcomes)

# --------------------------------------------------------------
# Single counterparty, but in a -2 sigma macroeconomic scenarios
# --------------------------------------------------------------

# Correlation to macroeconomic environment
rho = 0.3

# Convert to z (Gauss) based thresholds
link = NormalThresholdLink{Float64}()
thr = RealLineThresholds(UnitIntervalThresholds(pv), link)

# Deform thresholds according to macroeconomic environment
dthr = deform(thr, 2.0 * rho, sqrt(1 - rho^2))

# At this time we can either convert back to unit interval thresholds or stay in transformed coordinates
# Normally, it will be more efficient to just procede in transformed coordinates
my_stressed_outcomes = outcome_value(dthr, vals, randn(1_000_000))

# Observed mean value differs from original mean due to macroeconomic stress
mean(my_stressed_outcomes)

# Observed standard deviation
std(my_stressed_outcomes)

# Simulating over macroeconomic scenarios, but taking mean wrt the idiosyncratic effect

# For each scenario, compute value after deformation
my_macro_outcomes = deformation_value(thr, vals, link, -rho * randn(1_000_000), sqrt(1 - rho^2))

# Observed mean corresponds to unstressed mean
mean(my_macro_outcomes)

# Observed unexpected loss at 0.05%
quantile(my_macro_outcomes, 0.0005) - mean(my_macro_outcomes)


"""
    Multivector4{T <: Real} <: AbstractMultivector{T}

An immutable 16-tuple of real numbers that represents a member
of a 4-dimensional multivector algebra.

Each `Multivector4` has the form
```math
    p+qD
```
where ``p`` and ``q`` are 3-dimensional multivectors.
"""
struct Multivector4{T <: Real} <: AbstractMultivector{T}
    l::Multivector3{T}
    r::Multivector3{T}

    Multivector4{U}(l::Multivector3{U}, r::Multivector3{U}) where {U <: Real} = new(l, r)
end

function Multivector4(x::Multivector3{T}, y::Multivector3{T}) where {T <: Real}
    Multivector4{T}(x, y)
end

function Multivector4(x::Multivector3)
    Multivector4(x, zero(x))
end

function Multivector4(x::Multivector2)
    Multivector4(Multivector3(x), Multivector3(zero(x)))
end

function Multivector4(x::Multivector1)
    Multivector4(Multivector3(x), Multivector3(zero(x)))
end

function show(io::IO, z::Multivector4)
    print(io, "[1: ")
    print(io, z.l.l.l.l)
    print(io, ", A: ")
    print(io, z.l.l.l.r)
    print(io, ", B: ")
    print(io, z.l.l.r.l)
    print(io, ", AB: ")
    print(io, z.l.l.r.r)
    print(io, ", C: ")
    print(io, z.l.r.l.l)
    print(io, ", AC: ")
    print(io, z.l.r.l.r)
    print(io, ", BC: ")
    print(io, z.l.r.r.l)
    print(io, ", (AB)C: ")
    print(io, z.l.r.r.r)
    print(io, ", D: ")
    print(io, z.r.l.l.l)
    print(io, ", AD: ")
    print(io, z.r.l.l.r)
    print(io, ", BD: ")
    print(io, z.r.l.r.l)
    print(io, ", (AB)D: ")
    print(io, z.r.l.r.r)
    print(io, ", CD: ")
    print(io, z.r.r.l.l)
    print(io, ", (AC)D: ")
    print(io, z.r.r.l.r)
    print(io, ", (BC)D: ")
    print(io, z.r.r.r.l)
    print(io, ", ((AB)C)D: ")
    print(io, z.r.r.r.r)
    print(io, "]")
end

function real(z::Multivector4)
    real(z.l)
end

function unreal(z::Multivector4)
    vcat(unreal(z.l), asarray(z.r))
end

function zero(z::Multivector4{T}) where {T <: Real}
    Multivector4{T}(zero(z.l), zero(z.r))
end

function zero(::Type{Multivector4{T}}) where {T <: Real}
    Multivector4{T}(zero(T))
end

function one(z::Multivector4{T}) where {T <: Real}
    Multivector4{T}(one(z.l))
end

function one(::Type{Multivector4{T}}) where {T <: Real}
    Multivector4{T}(one(T))
end

"""
    conj{T <: Real}(z::Multivector4{T})

The `Multivector4` conjugate.
This operation is an involution.
"""
function conj(z::Multivector4{T}) where {T <: Real}
    Multivector4{T}(conj(z.l), -z.r)
end

"""
    cloak{T <: Real}(z::Multivector4{T})

The cloak conjugate changes the sign of even blades.
This operation is equivalent to `-dagger(z)` and thus is also an involution.
"""
function cloak(z::Multivector4{T}) where {T <: Real}
    Multivector4{T}(cloak(z.l), dagger(z.r))
end

"""
    dagger{T <: Real}(z::Multivector4{T})

The dagger conjugate changes the sign of odd blades.
This operation is an involution.
"""
function dagger(z::Multivector4{T}) where {T <: Real}
    Multivector4{T}(dagger(z.l), cloak(z.r))
end

"""
    star{T <: Real}(z::Multivector4{T})

Returns the Hodge star conjugate.
This operation is not an involution, but `star(star(z)) = dagger(z)`.
"""
function star(z::Multivector4{T}) where {T <: Real}
    Multivector4{T}(dagger(star(z.r)), star(z.l))
end

"""
    selfstar(z::Multivector4)

The self-star-conjugate part.
This operation is idempotent.
"""
function selfstar(z::Multivector4)
    (z + star(z) + dagger(z) + star(dagger(z))) / 4
end

"""
    antiselfstar(z::Multivector4)

The anti-self-star-conjugate part.
This operation is idempotent.
"""
function antiselfstar(z::Multivector4)
    (z - star(z) + dagger(z) - star(dagger(z))) / 4
end

function (+)(x::Multivector4, y::Multivector4)
    Multivector4(x.l + y.l, x.r + y.r)
end

function (+)(x::Multivector4, y::Multivector3)
    Multivector4(x.l + y, x.r)
end

function (+)(x::Multivector3, y::Multivector4)
    Multivector4(x + y.l, y.r)
end

function (+)(x::Multivector4, y::Multivector2)
    Multivector4(x.l + y, x.r)
end

function (+)(x::Multivector2, y::Multivector4)
    Multivector4(x + y.l, y.r)
end

function (+)(x::Multivector4, y::Multivector1)
    Multivector4(x.l + y, x.r)
end

function (+)(x::Multivector1, y::Multivector4)
    Multivector4(x + y.l, y.r)
end

function (+)(z::Multivector4, a::Real)
    Multivector4(z.l + a, z.r)
end

function (+)(a::Real, z::Multivector4)
    Multivector4(z.l + a, z.r)
end

function (-)(z::Multivector4{T}) where {T <: Real}
    Multivector4{T}(-z.l, -z.r)
end

function (-)(x::Multivector4, y::Multivector4)
    Multivector4(x.l - y.l, x.r - y.r)
end

function (-)(x::Multivector4, y::Multivector3)
    Multivector4(x.l - y, x.r)
end

function (-)(x::Multivector3, y::Multivector4)
    Multivector4(x - y.l, -y.r)
end

function (-)(x::Multivector4, y::Multivector2)
    Multivector4(x.l - y, x.r)
end

function (-)(x::Multivector2, y::Multivector4)
    Multivector4(x - y.l, -y.r)
end

function (-)(x::Multivector4, y::Multivector1)
    Multivector4(x.l - y, x.r)
end

function (-)(x::Multivector1, y::Multivector4)
    Multivector4(x - y.l, -y.r)
end

function (-)(z::Multivector4, a::Real)
    Multivector4(z.l - a, z.r)
end

function (-)(a::Real, z::Multivector4)
    Multivector4(a - z.l, -z.r)
end

"""
    (*)(x::Multivector4, y::Multivector4)

Wedge product of 4-dimensional multivectors.
This operation is non-commutative and non-associative.
"""
function (*)(x::Multivector4, y::Multivector4)
    Multivector4(x.l * y.l,  (y.r * x.l) + (x.r * conj(y.l)))
end

function (*)(x::Multivector4, y::Multivector3)
    Multivector4(x.l * y, x.r * conj(y))
end

function (*)(x::Multivector3, y::Multivector4)
    Multivector4(x * y.l,  y.r * x)
end

function (*)(x::Multivector4, y::Multivector2)
    Multivector4(x.l * y, x.r * conj(y))
end

function (*)(x::Multivector2, y::Multivector4)
    Multivector4(x * y.l,  y.r * x)
end

function (*)(x::Multivector4, y::Multivector1)
    Multivector4(x.l * y, x.r * conj(y))
end

function (*)(x::Multivector1, y::Multivector4)
    Multivector4(x * y.l,  y.r * x)
end

"""
    (*)(z::Multivector4, a::Real)
    (*)(a::Real, z::Multivector4)

Scaling and/or reflection of a `Multivector4` by a real number.
"""
function (*)(z::Multivector4, a::Real)
    Multivector4(z.l * a, z.r * a)
end

function (*)(a::Real, z::Multivector4)
    Multivector4(a * z.l, a * z.r)
end

function (/)(x::Multivector4, y::Multivector4)
    if iszerodivisor(y)
        error(ZeroDivisorDenominator)
    end

    x * inv(y)
end

function (\)(y::Multivector4, x::Multivector4)
    if iszerodivisor(y)
        error(ZeroDivisorDenominator)
    end

    inv(y) * x
end

function (/)(z::Multivector4, a::Real)
    Multivector4(z.l / a, z.r / a)
end

function (\)(a::Real, z::Multivector4)
    Multivector4(a \ z.l, a \ z.r)
end

function random(::Type{Multivector4{T}}) where T <: Real
    Multivector4{T}(random(Multivector3{T}), random(Multivector3{T}))
end


abstract type RandomGraphType end
struct ErdosRenyi <: RandomGraphType end
struct RandomGeometric <: RandomGraphType end

# --------------------------------------------------------------------------- #
#						       Random Graph Models
# --------------------------------------------------------------------------- #

function erdos_renyi(n,p)
    A = sprand(n,n,p)
    U = triu(A,1)
    is,js,_ = findnz(max.(U,U'))
    return sparse(is,js,ones(length(is)),n,n)
end

function erdos_renyi(n;degreedist=LogNormal(log(5),1))

    println("made it")
    # form the edges for sparse
    ei = Int[]
    ej = Int[]

    for i=1:n
      deg = ceil(Int,minimum((rand(degreedist),n-1)))
      neighbors = sample(1:n,deg+1, replace = false)

      for j in neighbors
        if i != j
          push!(ei,i)
          push!(ej,j)
        end
      end
    end

    A = sparse(ei,ej,1.0,n,n)

    return max.(A,A')
end

#--                             David's graph code                          --#

function random_geometric_graph(n,k)
  xy = rand(2,n)
  T = BallTree(xy)
  idxs = knn(T, xy, minimum((k,n)))[1]
  # form the edges for sparse
  ei = Int[]
  ej = Int[]
  for i=1:n
    for j=idxs[i]
      if i != j
        push!(ei,i)
        push!(ej,j)
      end
    end
  end
  A = sparse(ei,ej,1.0,n,n)

  return max.(A,A')
end

function spatial_network(n::Integer, d::Integer; degreedist=LogNormal(log(5),1))
  xy, ei, ej = spatial_graph_edges(n, d;degreedist=degreedist)
  A = sparse(ei,ej,1.0,n,n)
  return max.(A,A')
end

function spatial_graph_edges(n::Integer,d::Integer;degreedist=LogNormal(log(5),1))
  xy = rand(d,n)
  T = BallTree(xy)
  # form the edges for sparse
  ei = Int[]
  ej = Int[]
  for i=1:n
    deg = ceil(Int,minimum((rand(degreedist),n-2)))
    idxs, dists = knn(T, xy[:,i], deg+1)
    for j in idxs
      if i != j
        push!(ei,i)
        push!(ej,j)
      end
    end
  end
  return xy, ei, ej
end


"""
    QDHT{p, n}(R, N; dim=1)
    QDHT([p, ] R, N; dim=1)
    QDHT(p, [n, ] R, N; dim=1)

`p`-th order quasi-discrete Hankel transform over aperture radius `R` with `N` samples
which transforms along dimension `dim`. If not given, `p` defaults to 0, and `n` the
spherical dimension defaults to 1 (cylindrical).

After:

[1] L. Yu, M. Huang, M. Chen, W. Chen, W. Huang, and Z. Zhu, Optics Letters 23 (1998)

[2] M. Guizar-Sicairos and J. C. Gutiérrez-Vega, JOSA A 21, 53 (2004)

[3] H. F. Johnson, Comput. Phys. Commun., 43, 2 (1987)

but with some alterations:

The transform matrix T is not the same as C/T defined in [1, 2] but is more like the form
used in equation 14 of [3].
Instead of dividing by ``J_1(α_{pn}) J_1(α_{pm})`` we divide by ``J_1(α_{pn})^2``.
This cancels out the factor between ``f`` and ``F`` so we do not have to mutltiply (divide)
by ``J_1(α_{pn})`` (``J_1(α_{pm})``) before and after applying the transform matrix.

Follows [`AbstractFFT`](https://github.com/JuliaMath/AbstractFFTs.jl) approach of applying
fwd and inv transform with `mul` and `ldiv`.

To calculate radial integrals of functions sampled using `QDHT`, use [`integrateR`](@ref)
and [`integrateK`](@ref).

The type of the coefficients is inferred from the type of `R` (but is promoted to be at
least `Float`), so for arbitrary precision use `QDHT([p, ] BigFloat(R), ...)`.
"""
struct QDHT{p, n, T<:Real}
    N::Int # Number of samples
    T::Array{T, 2} # Transform matrix
    j1sq::Array{T, 1} # (j₁)² factors
    K::T # Highest spatial frequency
    k::Vector{T} # Spatial frequency grid
    R::T # Aperture size (largest real-space coordinate)
    r::Vector{T} # Real-space grid
    scaleR::Vector{T} # Scale factor for real-space integration
    scaleK::Vector{T} # Scale factor for frequency-space integration
    scaleRK::T # Scalar factor multiplied/divided during fwd/rev transform
    dim::Int # Dimension along which to transform
end

function QDHT{p, n}(R, N; dim=1) where {p, n}
    R = float(R)
    pf = oftype(R, p)
    cn = oftype(R, sphbesselj_scale(n))
    roots = sphbesselj_zero.(pf, n, 1:N) # type of sphbesselj_zero is inferred from first argument
    S = sphbesselj_zero(pf, n, N+1)
    r = roots .* R/S # real-space vector
    K = S/R # Highest spatial frequency
    k = roots .* K/S # Spatial frequency vector
    j₁sq = abs2.(sphbesselj.(p+1, n, roots))
    T = 2 * cn * S^(-(n + 1) / 2) * sphbesselj.(p, n, (roots .* roots') ./ S) ./ j₁sq' # Transform matrix

    scaleR = 2 * cn^2 / K^(n + 1) ./ j₁sq # scale factor for real-space integration
    scaleK = 2 * cn^2 / R^(n + 1) ./ j₁sq # scale factor for reciprocal-space integration
    scaleRK = (R / K) ^ ((n + 1) / 2)
    QDHT{p, n, eltype(T)}(N, T, j₁sq, K, k, R, r, scaleR, scaleK, scaleRK, dim)
end

QDHT{p}(R, N; dim=1) where {p} = QDHT{p, 1}(R, N; dim=dim)
QDHT(R, N; dim=1) = QDHT{0}(R, N; dim=dim)
QDHT(p, R, N; dim=1) = QDHT{p}(R, N; dim=dim)
QDHT(p, n, R, N; dim=1) = QDHT{p, n}(R, N; dim=dim)

"""
    order(q::QDHT)

Order of transform `q`.
"""
order(::QDHT{p}) where {p} = p

"""
    sphericaldim(q::QDHT)

Spherical dimension ``n`` of transform `q`.

Dimension of the ``n``-sphere over whose radial coordinate the bases of the transform `q`
are orthogonal polynomials. For the circle, ``n=1``, and for the sphere, ``n=2``.

See [`sphbesselj`](@ref).
"""
sphericaldim(::QDHT{p,n}) where {p,n} = n

"""
    inv(Q::QDHT) -> QDHT

Compute the `QDHT` for the inverse of the Hankel transform `Q`, that is, the transform such
that `Q * A == inv(Q) \\ A`
"""
function Base.inv(Q::T) where {T<:QDHT}
    return T(
        Q.N,
        Q.T,
        Q.j1sq,
        Q.R,
        Q.r,
        Q.K,
        Q.k,
        Q.scaleK,
        Q.scaleR,
        inv(Q.scaleRK),
        Q.dim,
    )
end

"""
    reshape(Q::QDHT, N::Int) -> QDHT

Compute a new `QDHT` with the same order, spherical dimension, and radius but with `N`
points.
"""
function Base.reshape(Q::QDHT{p, n}, N::Int) where {p, n}
    N == Q.N && return Q
    return QDHT{p, n}(Q.R, N, dim=Q.dim)
end

"
    mul!(Y, Q::QDHT, A)

Calculate the forward quasi-discrete Hankel transform of array `A` using the QDHT `Q`
and store the result in `Y`.

# Examples
```jldoctest
julia> q = QDHT{0, 1}(1e-2, 8); A = exp.(-q.r.^2/(1e-3*q.R)); Y = similar(A);
julia> mul!(Y, q, A)
8-element Array{Float64,1}:
  4.326937831591551e-6
  2.3341589529175126e-6
  7.689558743828849e-7
  1.546419420523699e-7
  1.8999259906096856e-8
  1.4159642663129888e-9
  7.013670190083954e-11
 -6.07681871673291e-13
```
"
function mul!(Y, Q::QDHT, A)
    dot!(Y, Q.T, A, dim=Q.dim)
    Y .*= Q.scaleRK
end

"
    ldiv!(Y, Q::QDHT, A)

Calculate the inverse quasi-discrete Hankel transform of array `A` using the QDHT `Q`
and store the result in `Y`.

# Examples
```jldoctest
julia> q = QDHT{0, 1}(1e-2, 8); A = exp.(-q.r.^2/(1e-3*q.R)); Y = similar(A);
julia> mul!(Y, q, A);
julia> YY = similar(Y); ldiv!(YY, q, Y);
julia> YY ≈ A
true
```
"
function ldiv!(Y, Q::QDHT, A)
    dot!(Y, Q.T, A, dim=Q.dim)
    Y ./= Q.scaleRK
end

"""
    *(Q::QDHT, A)

Calculate the forward quasi-discrete Hankel transform of array `A` using the QDHT `Q`.

# Examples
```jldoctest
julia> q = QDHT{0, 1}(1e-2, 8); A = exp.(-q.r.^2/(1e-3*q.R));
julia> q*A
8-element Array{Float64,1}:
  4.326937831591551e-6
  2.3341589529175126e-6
  7.689558743828849e-7
  1.546419420523699e-7
  1.8999259906096856e-8
  1.4159642663129888e-9
  7.013670190083954e-11
 -6.07681871673291e-13
```
"""
function *(Q::QDHT, A)
    out = similar(A)
    mul!(out, Q, A)
end

"""
    \\(Q::QDHT, A)

Calculate the inverse quasi-discrete Hankel transform of array `A` using the QDHT `Q`.

# Examples
```jldoctest
julia> q = QDHT{0, 1}(1e-2, 8); A = exp.(-q.r.^2/(1e-3*q.R));
julia> Ak = q*A;
julia> q \\ Ak ≈ A
true
```
"""
function \(Q::QDHT, A)
    out = similar(A)
    ldiv!(out, Q, A)
end

"""
    integrateR(A, Q::QDHT; dim=1)

Radial integral of `A`, over the aperture of `Q` in real space.

Assuming `A` contains samples of a function `f(r)` at sample points `Q.r`, then
`integrateR(A, Q)` approximates ∫f(r)r dr from r=0 to r=∞.

!!! note
    `integrateR` and `integrateK` fulfill Parseval's theorem, i.e. for some array `A`,
    `integrateR(abs2.(A), q)` and `integrateK(abs2.(q*A), q)` are equal, **but**
    `integrateR(A, q)` and `integrateK(q*A, q)` are **not** equal.

!!! warning
    using `integrateR` to integrate a function (i.e. `A` rather than `abs2(A)`) is only
    supported for the 0th-order QDHT. For more details see [Derivations](@ref).

# Examples
```jldoctest
julia> q = QDHT{0, 1}(10, 128); A = exp.(-q.r.^2/2);
julia> integrateR(abs2.(A), q) ≈ 0.5 # analytical solution of ∫exp(-r²)r dr from 0 to ∞
true
```
"""
integrateR(A, Q::QDHT; dim=1) = dimdot(Q.scaleR, A; dim=dim)

"""
    integrateK(Ak, Q::QDHT; dim=1)

Radial integral of `A`, over the aperture of `Q` in reciprocal space.

Assuming `A` contains samples of a function `f(k)` at sample points `Q.k`, then
`integrateR(A, Q)` approximates ∫f(k)k dk from k=0 to k=∞.

!!! note
    `integrateR` and `integrateK` fulfill Parseval's theorem, i.e. for some array `A`,
    `integrateR(abs2.(A), q)` and `integrateK(abs2.(q*A), q)` are equal, **but**
    `integrateR(A, q)` and `integrateK(q*A, q)` are **not** equal.

# Examples
```jldoctest
julia> q = QDHT{0, 1}(10, 128); A = exp.(-q.r.^2/2);
julia> integrateR(abs2.(A), q) ≈ 0.5 # analytical solution of ∫exp(-r²)r dr from 0 to ∞
true
julia> Ak = q*A;
julia> integrateK(abs2.(Ak), q) ≈ 0.5 # Same result
true

```
"""
integrateK(Ak, Q::QDHT; dim=1) = dimdot(Q.scaleK, Ak; dim=dim)

"""
    onaxis(Ak, Q::QDHT; dim=Q.dim)

Calculate on-axis sample in space (i.e. at r=0) from transformed array `Ak`.

!!! note
    `onaxis` is currently only supported for 0-order transforms

# Examples
```jldoctest
julia> q = QDHT{0, 1}(10, 128); A = exp.(-q.r.^2/2);
julia> onaxis(q*A, q) ≈ 1 # should be exp(0) = 1
true
```
"""
function onaxis(Ak, Q::QDHT{0,n}; dim = Q.dim) where {n}
    j00ocn = inv(gamma((n + 1) / 2) * 2^((n - 1) / 2))
    return j00ocn .* integrateK(Ak, Q; dim = dim)
end
onaxis(Ak, Q::QDHT{0,1}; dim=Q.dim) = J₀₀ .* integrateK(Ak, Q; dim=dim)

"""
    symmetric(A, Q::QDHT)

Create symmetric array from samples in `A`, including on-axis sample.

Given `A`, sampled at `[r₁, r₂, r₃, ...]`, generates array sampled at
`[...-r₃, -r₂, -r₁, 0, r₁, r₂, r₃...]`

# Examples
```jldoctest
julia> q = QDHT{0, 1}(10, 128); A = exp.(-q.r.^2);
julia> As = symmetric(A, q);
julia> size(As)
(257,)
julia> As[1:128] == A[128:-1:1]
true
julia> As[129] ≈ 1 # should be exp(0) = 1
true
julia> As[130:end] == A
true
```
"""
function symmetric(A, Q::QDHT; dim=Q.dim)
    order(Q) == 0 || throw(DomainError("`symmetric` is only supported for 0-order transforms"))
    s = size(A)
    N = s[dim]
    idxlo = CartesianIndices(s[1:(dim - 1)])
    idxhi = CartesianIndices(s[(dim + 1):end])
    A0 = asarray(onaxis(Q * A, Q))
    out = cat(view(A, idxlo, N:-1:1, idxhi), A0, view(A, idxlo, :, idxhi); dims=dim)
    return out
end

"""
    Rsymmetric(Q::QDHT)

Create radial coordinate array to go along with `symmetric(A, Q::QDHT)`.

# Examples
```jldoctest
julia> q = QDHT{0, 1}(10, 4);
julia> q.r
4-element Array{Float64,1}:
 1.6106347946239767
 3.697078919099734
 5.795844623798052
 7.8973942990196395
julia> Rsymmetric(q)
9-element Array{Float64,1}:
 -7.8973942990196395
 -5.795844623798052
 -3.697078919099734
 -1.6106347946239767
  0.0
  1.6106347946239767
  3.697078919099734
  5.795844623798052
  7.8973942990196395
```
"""
Rsymmetric(Q::QDHT) = vcat(-view(Q.r, lastindex(Q.r):-1:firstindex(Q.r)), 0, Q.r)

"""
    oversample(A, Q::QDHT; factor::Int=4) -> Tuple{AbstractArray,QDHT}
    oversample(A, Q::QDHT, QNew::QDHT) -> AbstractArray

Oversample (smooth) the array `A`, which is sampled with the `QDHT` `Q`, by a `factor`.

If calling the function many times, it is more efficient to precompute `QNew` with
[`reshape`](@ref) and then provide it to the function.

This works like Fourier-domain zero-padding: a new `QDHT` is created with the same radius,
but `factor` times more points. The existing array is transformed and placed onto this
new spatial frequency grid, and the rest filled with zeros. Transforming back yields the
same shape in space but with more samples.

!!! note
    Unlike in zero-padding using FFTs, the old and oversampled **spatial** grids do not
    have any sampling points in common.
"""
function oversample(A, Q::QDHT; factor::Int=4)
    QNew = reshape(Q, Q.N * factor)
    return oversample(A, Q, QNew), QNew
end
function oversample(A, Q::QDHT, QNew::QDHT)
    Ak = Q * A
    Ako = padzeros(Ak, QNew.N; dim = QNew.dim)
    Ao = QNew \ Ako
    return Ao
end


abstract Model # this is either an LSTM or RNN

type NNMatrix # Neural net layer's weights & gradients
    n::Int
    d::Int
    w::Matrix{Float64} # matrix of weights
    dw::Matrix{Float64} # matrix of gradients
    NNMatrix(n::Int) = new(n, 1, zeros(n,1), zeros(n,1))
    NNMatrix(n::Int, d::Int) = new(n, d, zeros(n,d), zeros(n,d))
    NNMatrix(n::Int, d::Int, w::Matrix{Float64}, dw::Matrix{Float64}) = new(n, d, w, dw)
    NNMatrix(w::Array) = new(size(w,1), size(w,2), w, zeros(size(w,1),size(w,2)))
    function NNMatrix(w::Matrix{Float64}, dw::Matrix{Float64})
      if size(w) != size(dw)
        error("Matrix and Gradient Matrix sizes do not match");
      end
      new(size(w,1), size(w,2), w, dw)
    end
end

randNNMat(n::Int, d::Int, std::FloatingPoint=1.) = NNMatrix(n, d, randn(n,d)*std, zeros(n,d))
onesNNMat(n::Int, d::Int) = NNMatrix(n, d, ones(n, d), zeros(n, d))

function softmax(m::NNMatrix)
    out = NNMatrix(m.n,m.d)
    maxval = maximum(m.w)
    out.w[:] = exp(m.w - maxval)
    out.w /= sum(out.w)
    return out
end


using LaTeXStrings, PyPlot, MosekTools, MarkovBounds
# use Catalyst to define the reaction network under consideration
# in this case, we investigate simple Michaelis-Menten kinetics
rn = @reaction_network begin
    10.0, S + E --> SE
    0.2, SE --> S + E
    4.0, SE --> P + E
    1.3, P --> S
end
S, E, SE, P = species(rn)
# specify initial molecular counts in system
x0 = Dict(S => 10, E => 10, SE => 0, P => 0)
x_scales = Dict(S => 5, E => 5, SE => 5, P => 5)
# compute bounds on mean molecular counts of the product counts at steady state
# for different truncation orders
orders = 2:6
sol = Dict()
for m in orders
    sol[m] = stationary_mean(rn, [S,P], x0, m, Mosek.Optimizer, auto_scaling = false, scales = x_scales)
end

for s in [S, P]
    fig, ax = subplots()
    ax.plot(orders, [sol[m].bounds[s][1] for m in orders], marker = "*", color="blue", label="lower bound")
    ax.plot(orders, [sol[m].bounds[s][2] for m in orders], marker = "*", color="red", label="upper bound")
    ax.set_xlabel("truncation order")
    ax.set_ylabel(string(L"⟨",s.f.name,L"_{\infty}⟩"))
    ax.legend()
    display(fig)
end

# same for variance
sol = Dict()
for m in orders
    sol[m] = stationary_variance(rn, [S,P], x0, m, Mosek.Optimizer, auto_scaling = false, scales = x_scales)
end

for s in [S, P]
    fig, ax = subplots()
    ax.plot(orders, [sol[m].bounds[s] for m in orders], marker = "*", color="red", label="upper bound")
    ax.set_xlabel("truncation order")
    ax.set_ylabel(string(L"Var(",s.f.name,L"_{\infty})"))
    ax.legend()
    display(fig)
end

# transient moments
nT = 5
traj = []
Tf_range = [0.025, 0.05, 0.075, 0.1, 0.15, 0.2, 0.25, 0.3,
            0.4, 0.5, 0.625, 0.75, 1.0, 1.25, 1.5, 2.0, 3.0, 4.0, 5.0]
order = 4
for Tf in Tf_range
    trange = range(Tf/nT, stop=Tf, length=nT)
    sol = transient_mean(rn, [S,P], x0, order, trange, Mosek.Optimizer, scales = x_scales)
    push!(traj, sol)
end

for s in [S,P]
    fig, ax = subplots()
    ax.plot(Tf_range, [sol.bounds[s][1] for sol in traj], color="blue", label = "lower bound")
    ax.plot(Tf_range, [sol.bounds[s][2] for sol in traj], color="red", label = "upper bound")
    ax.set_xlabel("time")
    ax.set_ylabel(string(L"⟨", s.f.name,L"⟩"))
    ax.legend()
    display(fig)
end

# variance
nT = 5
traj = []
Tf_range = [0.025, 0.05, 0.075, 0.1, 0.15, 0.2, 0.25, 0.3,
            0.4, 0.5, 0.625, 0.75, 1.0, 1.25, 1.5, 2.0, 3.0, 4.0, 5.0]
order = 4
for Tf in Tf_range
    trange = range(Tf/nT, stop=Tf, length=nT)
    sol = transient_variance(rn, [P], x0, order, trange, Mosek.Optimizer, scales =  x_scales)
    push!(traj, sol)
end

ps = []
for s in [P]
    fig, ax = subplots()
    ax.plot(Tf_range, [sol.bounds[s][1] for sol in traj], color="red", label = "upper bound")
    ax.set_xlabel("time")
    ax.set_ylabel(string(L"Var(", s.f.name,L")"))
    ax.legend()
    display(fig)
end

## volume of confidence ellipsoid
orders = 2:8
sol = Dict()
for m in orders
    sol[m] = stationary_confidence_ellipsoid(rn, x0, [S,P], m, Mosek.Optimizer; scales = x_scales)
end

fig, ax = subplots()
ax.plot(orders, [sol[m].bounds[[S,P]] for m in orders], marker="*", color="red", label = "upper bound")
ax.set_xlabel("order")
ax.set_ylabel("Volume of S-P confidence ellipsoids")
ax.legend()
display(fig)


# Canonical form of multivariate normal

import NumericExtensions
import NumericExtensions.ScalMat
import NumericExtensions.PDiagMat
import NumericExtensions.PDMat

using Distributions
using Base.Test

##### construction, basic properties, and evaluation

h = [1., 2., 3.]
dv = [1.2, 3.4, 2.6]
J = [4. -2. -1.; -2. 5. -1.; -1. -1. 6.]

x1 = [3.2, 1.8, 2.4]
x = rand(3, 100)

# SGauss

gs = IsoNormalCanon(h, 2.0)
@test isa(gs, IsoNormalCanon)
@test dim(gs) == 3
@test mean(gs) == mode(gs) == h / 2.0
@test !gs.zeromean
@test invcov(gs) == diagm(fill(2.0, 3))
@test cov(gs) == diagm(fill(0.5, 3))
@test var(gs) == diag(cov(gs))
@test_approx_eq entropy(gs) 0.5 * logdet(2π * e * cov(gs))

gsz = IsoNormalCanon(3, 2.0)
@test isa(gsz, IsoNormalCanon)
@test dim(gsz) == 3
@test mean(gsz) == zeros(3)
@test gsz.zeromean

# DGauss

gd = DiagNormalCanon(h, dv)
@test isa(gd, DiagNormalCanon)
@test dim(gd) == 3
@test_approx_eq mean(gd) h ./ dv
@test !gd.zeromean
@test invcov(gd) == diagm(dv)
@test_approx_eq cov(gd) diagm(1.0 / dv)
@test_approx_eq var(gd) diag(cov(gd))
@test_approx_eq entropy(gd) 0.5 * logdet(2π * e * cov(gd))

gdz = DiagNormalCanon(dv)
@test isa(gdz, DiagNormalCanon)
@test dim(gdz) == 3
@test mean(gdz) == zeros(3)
@test gdz.zeromean

# Gauss

gf = MvNormalCanon(h, J)
@test isa(gf, MvNormalCanon)
@test dim(gf) == 3
@test_approx_eq mean(gf) J \ h
@test !gf.zeromean
@test invcov(gf) == J
@test_approx_eq cov(gf) inv(J)
@test_approx_eq var(gf) diag(cov(gf))
@test_approx_eq entropy(gf) 0.5 * logdet(2π * e * cov(gf))

gfz = MvNormalCanon(J)
@test isa(gfz, MvNormalCanon)
@test dim(gfz) == 3
@test mean(gfz) == zeros(3)
@test gfz.zeromean

# conversion

us = convert(IsoNormal, gs)
gs2 = convert(IsoNormalCanon, us)
@test isa(us, IsoNormal)
@test isa(gs2, IsoNormalCanon)
@test dim(gs2) == dim(gs)
@test_approx_eq gs2.h gs.h
@test_approx_eq gs2.J.value gs.J.value

ud = convert(DiagNormal, gd)
gd2 = convert(DiagNormalCanon, ud)
@test isa(ud, DiagNormal)
@test isa(gd2, DiagNormalCanon)
@test dim(gd2) == dim(gd)
@test_approx_eq gd2.h gd.h
@test_approx_eq gd2.J.diag gd.J.diag

uf = convert(MvNormal, gf)
gf2 = convert(MvNormalCanon, uf)
@test isa(uf, MvNormal)
@test isa(gf2, MvNormalCanon)
@test dim(gf2) == dim(gf)
@test_approx_eq gf2.h gf.h
@test_approx_eq gf2.J.mat gf.J.mat

# logpdf evaluation

@test_approx_eq logpdf(gs, x1) logpdf(us, x1)
@test_approx_eq logpdf(gs, x)  logpdf(us, x)
@test_approx_eq logpdf(gsz, x1) logpdf(convert(IsoNormal, gsz), x1)
@test_approx_eq logpdf(gsz, x)  logpdf(convert(IsoNormal, gsz), x)

@test_approx_eq logpdf(gd, x1) logpdf(ud, x1)
@test_approx_eq logpdf(gd, x)  logpdf(ud, x)
@test_approx_eq logpdf(gdz, x1) logpdf(convert(DiagNormal, gdz), x1)
@test_approx_eq logpdf(gdz, x)  logpdf(convert(DiagNormal, gdz), x)

@test_approx_eq logpdf(gf, x1) logpdf(uf, x1)
@test_approx_eq logpdf(gf, x)  logpdf(uf, x)
@test_approx_eq logpdf(gfz, x1) logpdf(convert(MvNormal, gfz), x1)
@test_approx_eq logpdf(gfz, x)  logpdf(convert(MvNormal, gfz), x)

# sampling

x = rand(gs)
@test isa(x, Vector{Float64})
@test length(x) == dim(gs)

x = rand(gd)
@test isa(x, Vector{Float64})
@test length(x) == dim(gd)

x = rand(gf)
@test isa(x, Vector{Float64})
@test length(x) == dim(gf)

x = rand(gsz)
@test isa(x, Vector{Float64})
@test length(x) == dim(gsz)

x = rand(gdz)
@test isa(x, Vector{Float64})
@test length(x) == dim(gdz)

x = rand(gfz)
@test isa(x, Vector{Float64})
@test length(x) == dim(gfz)

n = 10
x = rand(gs, n)
@test isa(x, Matrix{Float64})
@test size(x) == (dim(gs), n)

x = rand(gd, n)
@test isa(x, Matrix{Float64})
@test size(x) == (dim(gd), n)

x = rand(gf, n)
@test isa(x, Matrix{Float64})
@test size(x) == (dim(gf), n)

x = rand(gsz, n)
@test isa(x, Matrix{Float64})
@test size(x) == (dim(gsz), n)

x = rand(gdz, n)
@test isa(x, Matrix{Float64})
@test size(x) == (dim(gdz), n)

x = rand(gfz, n)
@test isa(x, Matrix{Float64})
@test size(x) == (dim(gfz), n)


module PhysicsFormulae

using LinearAlgebra
include("./Initial.jl")
using .Initial

const G = 6.67408e-11

# Calcualte the difference in velocities of two nodes due to gravitational force
function acceleration(node1, node2)
    # First we calculate the distance
    displacement = node1.position - node2.position
    distance = sum(square.(displacement))

    dv = normalize(displacement)
    dv = ((G * node2.mass) / square(distance)) .* dv

    return dv
end

# Calcualte the velocities of two nodes after collusion
function bounce(node1, node2)
    m₁, m₂ = node1.mass, node2.mass
    v₁, v₂ = node1.velocity, node2.velocity

    vf₁ = v₁ .* ((m₁ - m₂) / (m₂ + m₁))
    vf₁ += v₂ .* ((2 * m₂) / (m₂ + m₁))

    vf₂ = v₁ .* ((2 * m₁) / (m₂ + m₁))
    vf₂ += v₂ .* ((m₂ - m₁) / (m₂ + m₁))

    return vf₁, vf₂
end

# Calculate the combined velocity of two nodes after collusion
function combinedVelocity(node1, node2)
    m₁, m₂ = node1.mass, node2.mass
    v₁, v₂ = node1.velocity, node2.velocity

    vf = (m₁ .* v₁ + m₂ .* v₂) / (m₁ + m₂)

    return vf
end

function collideIfColliding(node1, node2)
    p₁, p₂ = node1.position, node2.position
    r₁, r₂ = node1.radius, node2.radius

    minDistanceForCollusion = r₁ + r₂
    distanceBetweenNodes = sum(square.(p₁ - p₂))

    if distanceBetweenNodes <= minDistanceForCollusion
        if onCollusion == "bounce"
            bounce(node1, node2)
        else
            combinedVelocity(node1, node2)
        end
    end
end

square(n) = n^2

export acceleration, bounce, combinedVelocity, collideIfColliding
end


using SurrealNumbers

n = 3
x = convert.(SurrealFinite, 0:n)
A = zeros(SurrealFinite, n+1, n+1)
for i=1:n+1
    for j=1:n+1
        A[i,j] = x[i] + x[j]
    end
end
float.(A)
A
generation.(A)
A .== canonicalise.(A)

n = 2
x = convert.(SurrealFinite, (0:2^n)/2^n)
A = zeros(SurrealFinite, 2^n+1, 2^n+1)
for i=1:2^n +1
    for j=1:2^n +1
        A[i,j] = x[i] + x[j]
    end
end
float.(A)
A
generation.(A)
A .== canonicalise.(A)

n = 3
x = convert.(SurrealFinite, 0:n)
A = zeros(SurrealFinite, n+1, n+1)
for i=1:n+1
    for j=1:n+1
        A[i,j] = x[i] * x[j]
    end
end
float.(A)
generation.(A)
A .== canonicalise.(A)

# analyse.(A)

# function analyse(s::SurrealFinite)
#     println("depth = ", generation(s) )

# end


module math

import builtin

self.epsilon = 0.000001
self.pi = 3.14159265359
self.e = 2.71828182845

def min(args) {
	m = args[0]
  for i in args {
    if (i < m) {
      m = i
    }
  }
  m
}

def max(args) {
	m = args[0]
  for i in args {
    if (i > m) {
      m = i
    }
  }
  m
}

def pow(num, power) {
  builtin.pow__(num, power)
}

; if 1 arg is specified, defaults to natural log.
def log(args) {
  builtin.log__(args)
}

def ln(x) {
  builtin.log__(x)
}

def log10(x) {
  log(10, x)
}

def sqrt(x) pow(x, 0.5)


##

using ACEatoms, Test
ES = ACEatoms.Electrostatics

##

@info("Total dipole moment test")
X = [0.0 0.0 0.0; 1.0 0.0 0.0]
Q = [-1.0; 1.0]
MU = [0.0 0.0 0.0; 0.0 1.0 -2.0]
tot_dipole = ES.get_dipole(X, Q, MU)
println(@test(isapprox(tot_dipole, [0.208194 1.0 -2.0], atol=1e-6) ))

##

@info("Soft charge-charge interaction energy test")
X = [0.0 0.0; 1.0 0.0]
q1 = +1.0
q2 = -1.0
E = ES.soft_coulomb(X[1, :], q1, X[2, :], q2, 1.0)
println(@test( isapprox(E, -14.39964547, atol=1e-6) ))

##

@info("Soft charge-dipole interaction energy test")
X = [0.0 0.0; 1.0 0.0]
q1 = 1.0
dip = [1.0 0.0]
E = ES.soft_q_μ(X[1, :], q1, X[2, :], dip, 1.0)
println(@test( isapprox(E, -2.99792458, atol=1e-6) ))

##


# Laplace distribution

export Laplace

@parameterized Laplace() ≪ (1/2) * Lebesgue(ℝ)

Laplace(nt::NamedTuple{(:μ,:σ)}) = Affine(nt, Laplace())
Laplace(nt::NamedTuple{(:μ,:ω)}) = Affine(nt, Laplace())
Laplace(nt::NamedTuple{(:σ,)}) = Affine(nt, Laplace())
Laplace(nt::NamedTuple{(:ω,)}) = Affine(nt, Laplace())
Laplace(nt::NamedTuple{(:μ,)}) = Affine(nt, Laplace())

@affinepars Laplace

function logdensity(d::Laplace{()} , x)
    return -abs(x)
end

Base.rand(rng::AbstractRNG, μ::Laplace{()}) = rand(rng, Dists.Laplace())

≪(::Laplace, ::Lebesgue{X}) where X <: Real = true

TV.as(::Laplace) = asℝ

distproxy(::Laplace{()}) = Dists.Laplace()


#=
An optimization oracle is a linear optimization oracle for the set S, i.e., it solves
    z^\ast(c) :=    min c^Tw
                    s.t. w \in S
The solution of the above problem returned by the oracle is referred to as w^\ast(\cdot) in
the paper

A specific implementation of an optimization oracle should take a single input argument c,
and should return the pair (z^\ast(c), w^\ast(c))

Below are several basic examples
=#

"""
    simplex_oracle(c)

Optimization oracle when S is the unit simplex
"""
function simplex_oracle(c::Vector{Float64})
    (z, ind) = findmin(c)
    w = zeros(length(c))
    w[ind] = 1
    return (z, w)
end

"""
    svm_oracle(c)

Optimization oracle for the SVM example where S = [-1/2, +1/2].

# Arguments
- `c` is assumed to be a vector of length 1, NOT a scalar
"""
function svm_oracle(c::Vector{Float64})
    c_val = c[1]
    if c_val > 0
        w = [-0.5]
        z = (-0.5)*c_val
    else
        w = [0.5]
        z = 0.5*c_val
    end

    return (z, w)
end


# Test statistics related functions

using NumericExtensions
using Base.Test

# data

x = randn(3, 4)
y = randn(3, 4)

x1 = randn(6)
y1 = randn(6)

x2 = randn(5, 6)
y2 = randn(5, 6)

x3 = randn(3, 4, 5)
y3 = randn(3, 4, 5)

x4 = randn(3, 4, 5, 3)
y4 = randn(3, 4, 5, 3)

# varm

safe_varm(x, u) = sum(abs2(x - u)) / (length(x) - 1)
safe_varm(x, u, d::Int) = sum(abs2(x .- u), d) / (size(x, d) - 1)

@test_approx_eq varm(x, mean(y)) safe_varm(x, mean(y))
@test_approx_eq varm(x1, mean(y1,1), 1) safe_varm(x1, mean(y1,1), 1)

@test_approx_eq varm(x2, mean(y2,1), 1) safe_varm(x2, mean(y2,1), 1)
@test_approx_eq varm(x2, mean(y2,2), 2) safe_varm(x2, mean(y2,2), 2)

@test_approx_eq varm(x3, mean(y3,1), 1) safe_varm(x3, mean(y3,1), 1)
@test_approx_eq varm(x3, mean(y3,2), 2) safe_varm(x3, mean(y3,2), 2)
@test_approx_eq varm(x3, mean(y3,3), 3) safe_varm(x3, mean(y3,3), 3)

@test_approx_eq varm(x4, mean(y4,1), 1) safe_varm(x4, mean(y4,1), 1)
@test_approx_eq varm(x4, mean(y4,2), 2) safe_varm(x4, mean(y4,2), 2)
@test_approx_eq varm(x4, mean(y4,3), 3) safe_varm(x4, mean(y4,3), 3)
@test_approx_eq varm(x4, mean(y4,4), 4) safe_varm(x4, mean(y4,4), 4)

# stdm

safe_stdm(x, u) = sqrt(safe_varm(x, u))
safe_stdm(x, u, d::Int) = sqrt(safe_varm(x, u, d))

@test_approx_eq stdm(x, mean(y)) safe_stdm(x, mean(y))
@test_approx_eq stdm(x1, mean(y1,1), 1) safe_stdm(x1, mean(y1,1), 1)

@test_approx_eq stdm(x2, mean(y2,1), 1) safe_stdm(x2, mean(y2,1), 1)
@test_approx_eq stdm(x2, mean(y2,2), 2) safe_stdm(x2, mean(y2,2), 2)

@test_approx_eq stdm(x3, mean(y3,1), 1) safe_stdm(x3, mean(y3,1), 1)
@test_approx_eq stdm(x3, mean(y3,2), 2) safe_stdm(x3, mean(y3,2), 2)
@test_approx_eq stdm(x3, mean(y3,3), 3) safe_stdm(x3, mean(y3,3), 3)

@test_approx_eq stdm(x4, mean(y4,1), 1) safe_stdm(x4, mean(y4,1), 1)
@test_approx_eq stdm(x4, mean(y4,2), 2) safe_stdm(x4, mean(y4,2), 2)
@test_approx_eq stdm(x4, mean(y4,3), 3) safe_stdm(x4, mean(y4,3), 3)
@test_approx_eq stdm(x4, mean(y4,4), 4) safe_stdm(x4, mean(y4,4), 4)

# var

safe_var(x) = safe_varm(x, mean(x))
safe_var(x, d::Int) = safe_varm(x, mean(x, d), d)

@test_approx_eq var(x) safe_var(x)
@test_approx_eq var(x1, 1) safe_var(x1, 1)

@test_approx_eq var(x2, 1) safe_var(x2, 1)
@test_approx_eq var(x2, 2) safe_var(x2, 2)

@test_approx_eq var(x3, 1) safe_var(x3, 1)
@test_approx_eq var(x3, 2) safe_var(x3, 2)
@test_approx_eq var(x3, 3) safe_var(x3, 3)

@test_approx_eq var(x4, 1) safe_var(x4, 1)
@test_approx_eq var(x4, 2) safe_var(x4, 2)
@test_approx_eq var(x4, 3) safe_var(x4, 3)
@test_approx_eq var(x4, 4) safe_var(x4, 4)

r = zeros(size(x2, 2)); var!(r, x2, 1)
@test_approx_eq r vec(var(x2, 1))

# std

safe_std(x) = sqrt(var(x))
safe_std(x, d) = sqrt(var(x, d))

@test_approx_eq std(x) safe_std(x)
@test_approx_eq std(x1, 1) safe_std(x1, 1)

@test_approx_eq std(x2, 1) safe_std(x2, 1)
@test_approx_eq std(x2, 2) safe_std(x2, 2)

@test_approx_eq std(x3, 1) safe_std(x3, 1)
@test_approx_eq std(x3, 2) safe_std(x3, 2)
@test_approx_eq std(x3, 3) safe_std(x3, 3)

@test_approx_eq std(x4, 1) safe_std(x4, 1)
@test_approx_eq std(x4, 2) safe_std(x4, 2)
@test_approx_eq std(x4, 3) safe_std(x4, 3)
@test_approx_eq std(x4, 4) safe_std(x4, 4)

r = zeros(size(x2, 2)); std!(r, x2, 1)
@test_approx_eq r vec(std(x2, 1))

# logsumexp

safe_logsumexp(x) = log(sum(exp(x)))
safe_logsumexp(x, d::Int) = log(sum(exp(x), d))

@test_approx_eq logsumexp(x) safe_logsumexp(x)
@test_approx_eq logsumexp(x1, 1) safe_logsumexp(x1, 1)
@test_approx_eq logsumexp(x2, 1) safe_logsumexp(x2, 1)
@test_approx_eq logsumexp(x2, 2) safe_logsumexp(x2, 2)
@test_approx_eq logsumexp(x3, 1) safe_logsumexp(x3, 1)
@test_approx_eq logsumexp(x3, 2) safe_logsumexp(x3, 2)
@test_approx_eq logsumexp(x3, 3) safe_logsumexp(x3, 3)

r = zeros(size(x2, 2)); logsumexp!(r, x2, 1)
@test_approx_eq r vec(logsumexp(x2, 1))

@test_approx_eq logsumexp(x + 1000.) logsumexp(x) + 1000.
@test_approx_eq logsumexp(x2 + 1000., 1) logsumexp(x2, 1) + 1000.
@test_approx_eq logsumexp(x2 + 1000., 2) logsumexp(x2, 2) + 1000.
@test_approx_eq logsumexp(x3 + 1000., 2) logsumexp(x3, 2) + 1000.

# softmax

safe_softmax(x) = exp(x) / sum(exp(x))
safe_softmax(x, d::Int) = exp(x) ./ sum(exp(x), d)

@test_approx_eq softmax(x) safe_softmax(x)
@test_approx_eq softmax(x1, 1) safe_softmax(x1, 1)
@test_approx_eq softmax(x2, 1) safe_softmax(x2, 1)
@test_approx_eq softmax(x2, 2) safe_softmax(x2, 2)
@test_approx_eq softmax(x3, 1) safe_softmax(x3, 1)
@test_approx_eq softmax(x3, 2) safe_softmax(x3, 2)
@test_approx_eq softmax(x3, 3) safe_softmax(x3, 3)

r = zeros(size(x2)); softmax!(r, x2, 1)
@test_approx_eq r softmax(x2, 1)

@test_approx_eq softmax(x + 1000.) softmax(x)
@test_approx_eq softmax(x2 + 1000., 1) softmax(x2, 1)
@test_approx_eq softmax(x2 + 1000., 2) softmax(x2, 2)
@test_approx_eq softmax(x3 + 1000., 2) softmax(x3, 2)


using BSplineKit
using BSplineKit.BSplines: multiplicity

using BandedMatrices
using LinearAlgebra
using Test

eval_poly(x::Number, P) = @evalpoly(x, P...)

# Test a polynomial of degree k - 1.
# The splines should approximate the polynomial (and its derivatives) perfectly.
# This is also used to test the SplineInterpolations module.
function test_polynomial(x, ::BSplineOrder{k}) where {k}
    # Coefficients of polynomial of degree k - 1 (see ?evalpoly).
    # P(x) = -1 + 2x - 3x^2 + 4x^3 - ...
    P = ntuple(d -> (-d)^d, Val(k))
    P′ = ntuple(d -> d * P[d + 1], Val(k - 1))    # derivative
    Pint = (0, ntuple(d -> P[d] / d, Val(k))...)  # antiderivative

    # Interpolate polynomial at `x` locations.
    y = eval_poly.(x, Ref(P))
    itp = @inferred interpolate(x, y, BSplineOrder(k))

    @test startswith(repr(itp), "SplineInterpolation containing")

    S = spline(itp)
    @test length(S) == length(x)

    @testset "Interpolations" begin
        @test itp.(x) ≈ y

        let x = (x[2] + x[3]) / 3
            @test itp(x) == S(x)  # these are equivalent
        end

        @test order(itp) === order(S)
        @test basis(itp) === basis(S)
        @test knots(itp) === knots(S)
        @test coefficients(itp) === coefficients(S)
        @test integral(itp) == integral(S)
        @test diff(itp, Derivative(1)) == diff(S, Derivative(1)) == Derivative() * itp

        # "incompatible lengths of B-spline basis and collocation points"
        @test_throws(
            DimensionMismatch,
            SplineInterpolation(undef, basis(S), x[1:4], eltype(S)),
        )

        let B = basis(S)
            xs = collocation_points!(Vector{Float32}(undef, length(B)), B)
            itp = @inferred SplineInterpolation(undef, B, xs)
            @test eltype(itp) === Float32
        end

        # "input data has incorrect length"
        @test_throws DimensionMismatch interpolate!(itp, rand(length(x) - 1))
    end

    S′ = diff(S, Derivative(1))
    @test Derivative() * S == S′  # alternative notation
    Sint = integral(S)

    a, b = boundaries(basis(S))

    @test Sint(a) == 0  # this is an arbitrary choice
    Pint_a = @evalpoly(a, Pint...)

    # Compare values on a finer grid.
    let Nx = 9 * length(S) + 42
        x = LinRange(a, b, Nx)
        @test all(@evalpoly(x, P...) ≈ S(x) for x in x)
        @test all(@evalpoly(x, P′...) ≈ S′(x) for x in x)
        @test all(@evalpoly(x, Pint...) - Pint_a ≈ Sint(x) for x in x)
    end

    nothing
end

function test_splines(B::BSplineBasis, knots_in)
    k = order(B)
    t = knots(B)

    @testset "Knots (k = $k)" begin
        let (ka, kb) = multiplicity.(Ref(t), (1, length(t)))
            @test ka == kb == k
        end

        @test @views all(t[1:k] .== knots_in[1])
        @test @views all(t[(end - k + 1):end] .== knots_in[end])
        @test @views t[(k + 1):(end - k)] == knots_in[2:(end - 1)]
    end

    @testset "B-splines (k = $k)" begin
        N = length(B)
        @test_throws DomainError evaluate(B, 0, 0.2)
        @test_throws DomainError evaluate(B, N + 1, 0.2)

        @test_throws BoundsError B[0]
        @test_throws BoundsError B[N + 1]

        # In segment [t[k + 1], t[k + 2]], non-zero B-splines are 2:(k + 1).
        @test @inferred(nonzero_in_segment(B, k + 1)) == 2:(k + 1)

        # Fewer B-splines near the boundaries. This is usually not a problem
        # because knots are repeated there...
        @test nonzero_in_segment(B, 1) == 1:1
        @test nonzero_in_segment(B, N) == (N - k + 1):N
        @test nonzero_in_segment(B, N + 1) == (N - k + 2):N
        @test nonzero_in_segment(B, N + k - 1) == N:N
        @test isempty(nonzero_in_segment(B, N + k))

        # These intervals are not defined
        @test isempty(nonzero_in_segment(B, 0))
        @test isempty(nonzero_in_segment(B, length(knots(B))))

        # Verify values at the boundaries.
        a, b = boundaries(B)
        @test evaluate(B, 1, a) == 1.0
        @test evaluate(B, N, b) == 1.0

        @test startswith(
            repr(B),
            "$N-element BSplineBasis of order $k",
        )

        xs = range(a, b; length = 4N + 1)
        us = similar(xs)
        vs = similar(us)
        i = N >> 2
        evaluate!(us, B, i, xs)
        @test us == evaluate(B, i, xs)
        @test us == BasisFunction(B, i).(xs)
        @test us == B[i].(xs)
        @test us == map!(B[i], vs, xs)
        @test us ≈ B[i, Float32].(xs)

        @testset "iterate" begin
            for (i, b) in enumerate(B)
                @test b === B[i]
            end
        end
    end

    xcol = collocation_points(B, method=Collocation.AvgKnots())
    @test xcol[1] == knots_in[1]
    @test xcol[end] == knots_in[end]

    C = @inferred collocation_matrix(B, xcol)

    let N = size(C, 1)
        T = eltype(C)
        l, u = bandwidths(C)
        @test startswith(
            sprint(show, MIME("text/plain"), C),
            "$N×$N CollocationMatrix{$T} with bandwidths ($l, $u):\n",
        )
    end

    @testset "LU factorisation" begin
        T = eltype(C)
        @test T === Float64
        @test C isa CollocationMatrix{T}
        F = @inferred factorize(C)
        @test F isa LU{T, <:CollocationMatrix{T}}
        @test F.L * F.U ≈ C
        @test F.P == I  # permutation matrix = identity matrix (no pivoting)
        @test F.p == 1:size(C, 1)  # no row permutations

        y = sin.(xcol)
        u = F \ y
        @test C * u ≈ y

        let v = similar(u, length(u) - 2)
            @test_throws DimensionMismatch ldiv!(v, F, y)
        end
    end

    @testset "Spline (k = $k)" begin
        @testset "Polynomial" begin
            test_polynomial(xcol, BSplineOrder(k))
        end

        # Generate data at collocation points and get B-spline coefficients.
        ucol = cos.(xcol)
        coefs = C \ ucol
        @test C * coefs ≈ ucol

        @inferred Spline(B, coefs)
        S = Spline(B, coefs)
        @test length(S) == length(B)
        @test all(S.(xcol) .≈ ucol)
        @test coefficients(S) === coefs
        @test diff(S, Derivative(0)) === S
        @test_nowarn show(devnull, S)
        @test Splines.parent_spline(S) === S

        # Broadcasting
        f(S, x) = S(x)
        @test f.(S, xcol) == S.(xcol)

        # Create new spline, then compare it to S.
        let P = Spline(undef, B)
            cp = coefficients(P)
            fill!(cp, 0)
            @test P != S
            @test !(P ≈ S)

            copy!(cp, coefs)  # copy coefficients of S
            @test P == S
            @test P ≈ S

            @test copy(S) == P
        end
    end

    nothing
end

function test_splines(::BSplineOrder{k}) where {k}
    gauss_lobatto_points(N) = [-cos(π * n / N) for n = 0:N]
    breaks = (
        -1:0.05:1,
        gauss_lobatto_points(10 + k),
    )
    for x in breaks
        @inferred BSplineBasis(BSplineOrder(k), copy(x))
        @inferred BSplineBasis(BSplineOrder(k), copy(x), augment=Val(false))
        @inferred BSplineBasis(BSplineOrder(k), copy(x), augment=Val(true))
        @inferred (() -> BSplineBasis(k, copy(x)))()
        g = BSplineBasis(k, copy(x))

        @testset "BSplineBasis equality" begin
            h = BSplineBasis(k, copy(x))
            @test g == h

            h = BSplineBasis(BSplineOrder(k + 1), copy(x))
            @test g != h

            h = BSplineBasis(k, x .+ 1)
            @test g != h
        end

        @test order(g) == k
        test_splines(g, x)
    end
    nothing
end

@testset "Splines" begin
    test_splines(BSplineOrder(4))
    test_splines(BSplineOrder(5))
end


# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/NodeNumbering.jl/blob/master/LICENSE

"""
    adjacency_visualization(RCM_adjacency::Dict{Int, Vector{Int}})

Visualizes adjacency graph Dicts as matrices. Result can be plotted using `PyPlot` function `matshow()`.
A part of the Reverse Cuthill-McKee Algorithm.

# Examples

```jldoctest
julia> RCM_adjacency = Dict{Int, Vector{Int}}(
                       7 => [5, 6, 9],
                       9 => [7, 8],
                       4 => [2, 6],
                       2 => [4, 3, 1],
                       3 => [5, 6, 2, 1],
                       8 => [6, 9],
                       5 => [7, 3],
                       6 => [7, 3, 8, 4],
                       1 => [3, 2]);

julia> adjacency_visualization(RCM_adjacency)
9×9 Array{Float64,2}:
 1.0  1.0  1.0  0.0  0.0  0.0  0.0  0.0  0.0
 1.0  1.0  1.0  1.0  0.0  0.0  0.0  0.0  0.0
 1.0  1.0  1.0  0.0  1.0  1.0  0.0  0.0  0.0
 0.0  1.0  0.0  1.0  0.0  1.0  0.0  0.0  0.0
 0.0  0.0  1.0  0.0  1.0  0.0  1.0  0.0  0.0
 0.0  0.0  1.0  1.0  0.0  1.0  1.0  1.0  0.0
 0.0  0.0  0.0  0.0  1.0  1.0  1.0  0.0  1.0
 0.0  0.0  0.0  0.0  0.0  1.0  0.0  1.0  1.0
 0.0  0.0  0.0  0.0  0.0  0.0  1.0  1.0  1.0

```
"""
function adjacency_visualization(RCM_adjacency::Dict{Int, Vector{Int}})
    m=eye(length(RCM_adjacency))
    for (k, v) in RCM_adjacency
        for i in v
          m[k, i]=1
        end
    end
    return m
end


using OneDmin
using Test
using Printf
using Logging

fh(x) = - ( x + sin(x)) * exp(-x^2)
x0 = big(0.2); x1 = big(1.0)

using ForwardDiff

gh(x) = ForwardDiff.derivative(fh, x)  # first derivative
Hh(x) = ForwardDiff.derivative(gh, x)  # second derivative
Ch(x) = ForwardDiff.derivative(Hh, x)  # third derivative

setprecision(4096)

function test(algo::Function, x, xm)
    h = fh(x)
    dh = gh(x)
    d2h = Hh(x)
    d3h = Ch(x)

    hm = fh(xm)
    dhm = gh(xm)
    d2hm = Hh(xm)

    it = 0
    maxit = 25
    with_logger(Logging.NullLogger()) do
        @info "it  x  |h'|"
        while abs(dh) > BigFloat("1.0e-1200")
            xp = algo(x,xm,h,hm,dh,dhm,d2h,d2hm,d3h,0)
            x, xm = xp, x

            hm = h
            dhm = dh
            d2hm = d2h

            h = fh(x)
            dh = gh(x)
            d2h = Hh(x)
            d3h = Ch(x)

            it += 1
            if (it > maxit)  break  end
            @info @sprintf("%2d",it), @sprintf("%f",x), @sprintf("%.2e",abs(dh))
        end
    end
    return it

end

iter1 = Dict(base .=> [16,12,8])
iter2 = Dict(quintic .=> [8,7])
iter3 = Dict(cubic1 .=> [8,8,11])
iter4 = Dict(cubic0 .=> [11,11,15])

iter = merge(iter1, iter2, iter3, iter4)

@testset "OneDmin.jl" begin
    xm = x1; x = x0

    @info "Testing algos in base."
    for algo in base
        @info "testing" algo
        it = test(algo, x, xm)
        @test it == iter[algo]
    end

    @info "Testing algos in quintic."
    for algo in quintic
        @info "testing" algo
        it = test(algo, x, xm)
        @test it == iter[algo]
    end

    @info "Testing algos in cubic1."
    for algo in cubic1
        @info "testing" algo
        it = test(algo, x, xm)
        @test it == iter[algo]
    end

    @info "Testing algos in cubic0."
    for algo in cubic0
        @info "testing" algo
        it = test(algo, x, xm)
        @test it == iter[algo]
    end

    include("test1D.jl")

end


using DecisionMakingUtils
using OnlineStats
using Test

@testset "ZeroOneNormalization Test" begin
    k = 3
    T = Float32
    low = -collect(T, 1:k)
    high = collect(T, 1:k)
    ranges = hcat(low, high)
    n = ZeroOneNormalization(ranges)
    x = zeros(T, k)
    buff = zero(x)
    y = n(x)
    @test typeof(y) == typeof(x)
    @test size(y) == size(x)

    @test all(y .≥ 0.0)
    @test all(y .≤ 1.0)
    y = n(low)
    @test all(y .≥ 0.0)
    @test all(y .≤ 1.0)
    y = n(high)
    @test all(y .≥ 0.0)
    @test all(y .≤ 1.0)
    n(buff, high)
    @test all(isapprox.(buff, y))

    @. low -= 1
    @. high += 1
    n(buff, low)
    @test all(buff .< 0.0)
    n(buff, low, fit=true)
    @test all(buff .≥ 0.0)
    n(buff, high)
    @test all(buff .> 1.0)
    n(buff, high, fit=true)
    @test all(buff .≤ 1.0)
end

@testset "PosNegNormalization Test" begin
    k = 3
    T = Float32
    low = -collect(T, 1:k)
    high = collect(T, 1:k)
    ranges = hcat(low, high)
    n = PosNegNormalization(ranges)

    x = zeros(T, k)
    buff = zero(x)
    y = n(x)
    @test typeof(y) == typeof(x)
    @test size(y) == size(x)
    @test all(y .≥ -1.0)
    @test all(y .≤ 1.0)
    y = n(low)

    @test all(y .≥ -1.0)
    @test all(y .≤ 1.0)
    y = n(high)
    @test all(y .≥ -1.0)
    @test all(y .≤ 1.0)
    n(buff, high)
    @test all(isapprox.(buff, y))

    @. low -= 1
    @. high += 1
    n(buff, low)
    @test all(buff .< -1.0)
    n(buff, low, fit=true)
    @test all(buff .≥ -1.0)
    n(buff, high)
    @test all(buff .> 1.0)
    n(buff, high, fit=true)
    @test all(buff .≤ 1.0)
end

@testset "GaussianNormalization Test" begin
    k = 3
    T = Float32
    x1 = collect(T, 1:k)
    x2 = collect(T, 1:k) .* 2
    x3 = collect(T, 1:k) .* 3

    n = GaussianNormalization(k)

    x = zeros(T, k)
    buff = zero(x)
    y = n(x)
    @test typeof(y) == typeof(x)
    @test size(y) == size(x)

    @test all(isapprox.(y, 0.0))
    @. x = 1
    y = n(x)
    @test all(isapprox.(y, 1.0))

    n(buff, x)
    @test all(isapprox.(buff, y))

    v = k*KahanVariance()
    x1 = collect(T, 1:k)
    x2 = collect(T, 1:k) .* 2
    x3 = collect(T, 1:k) .* 3

    fit!(v, x1)
    @. y = (x1 - mean.(v.stats)) / std.(v.stats)
    n(buff, x1, fit=true)
    @test all(isapprox(buff, y))

    fit!(v, x2)
    @. y = (x2 - mean.(v.stats)) / std.(v.stats)
    n(buff, x2, fit=true)
    @test all(isapprox(buff, y))

    fit!(v, x3)
    @. y = (x3 - mean(v.stats)) / std(v.stats)
    n(buff, x3, fit=true)
    @test all(isapprox(buff, y))
end


using Pkg
Pkg.activate(@__DIR__)
Pkg.instantiate()

using Logging
using Statistics
using Gadfly
using Cairo
using DataFrames
using DifferentialEquations
import Distributions: Normal, Uniform

struct IVP
    f::Function
    ini::Vector
    t::AbstractArray
    phi::Vector
    data::Vector
    IVP(f, ini, t, phi, data) = new(f, ini, t, phi, data)
    IVP(f, ini, t, phi) = new(f, ini, t, phi,
                              (solve_ivp(new(f, ini, t, phi, []))
                               |> transpose_vector))
end

function get_ivps()::Dict
    var_ini = 0.01
    ivps = Dict()

    function fic(dz_dt, z, phi, t)
        r_1 = phi[1]*z[1]
        r_2 = phi[2]*z[2]

        dz_dt[1] = -r_1
        dz_dt[2] = r_1 - r_2
    end

    ivps[:fic] = Dict()
    ini = [1., 0.]
    t = range(0.0, 1.0, length=10)
    phi = [5.0035, 1.0]

    ivps[:fic][:nom] = IVP(fic, ini, t, phi)
    ivps_test = IVP(fic, add_noise(ini, var_ini), t, phi)
    ivps[:fic][:test] = IVP(fic, ivps_test.ini, t, phi,
                            add_noise(ivps_test.data, 0.1))

    ivps[:fic][:ss] = IVP(fic, ini, t,
                          [5.011535232847903,
                           0.9490955193480339], [])
    ivps[:fic][:ssr] = IVP(fic, ini, t,
                           [5.662374165061859,
                            4.600853424625171], [])
    ivps[:fic][:ds] = IVP(fic, ini, t,
                          [4.767811504376019,
                           0.9620285702765027], [])
    ivps[:fic][:dss] = IVP(fic, ini, t,
                           [5.057700471184158,
                            1.1660869766794006], [])

    function ccgo(dz_dt, z, phi, t)
        r_1 = phi[1]*z[1]^2
        r_2 = phi[2]*z[2]
        r_3 = phi[3]*z[1]^2

        dz_dt[1] = - r_1 - r_3
        dz_dt[2] = r_1 - r_2
    end

    ivps[:ccgo] = Dict()
    ini = [1., 0.]
    t = range(0.0, 0.95, length=10)
    phi = [12.214, 7.9798, 2.2216]

    ivps[:ccgo][:nom] = IVP(ccgo, ini, t, phi)
    ivps_test = IVP(ccgo, add_noise(ini, var_ini), t, phi)
    ivps[:ccgo][:test] = IVP(ccgo, ivps_test.ini, t, phi,
                            add_noise(ivps_test.data, 0.1))

    ivps[:ccgo][:ss] = IVP(ccgo, ini, t,
                           [10.601173576959617,
                            7.024603700551421,
                            5.296705518161371], [])
    ivps[:ccgo][:ssr] = IVP(ccgo, ini, t,
                            [12.694091102822359,
                             7.7461038473135275,
                             2.3847262326540015], [])
    ivps[:ccgo][:ds] = IVP(ccgo, ini, t,
                           [3.261725472047141,
                            9.713554040149868,
                            15.845961062140784], [])
    ivps[:ccgo][:dss] = IVP(ccgo, ini, t,
                            [11.701892961901253,
                             7.718104344633692,
                             2.748291378149207], [])

    function fhn(dz_dt, z, phi, t)
        a, b, c = phi
        V, R = z
        u = 0
        dz_dt[1] = c*(V - (V^3)/3 + R) + u
        dz_dt[2] = -(1/c)*(V - a + b*R)
    end

    ivps[:fhn] = Dict()
    ini = [-1.0, 1.0]
    t = range(0.0, 20.0, length=10)
    phi = [0.2, 0.2, 3.0]

    ivps[:fhn][:nom] = IVP(fhn, ini, t, phi)
    ivps_test = IVP(fhn, add_noise(ini, var_ini, false), t, phi)
    ivps[:fhn][:test] = IVP(fhn, ivps_test.ini, t, phi,
                            add_noise(ivps_test.data, 0.01, false))

    ivps[:fhn][:ss] = IVP(fhn, ini, t,
                          [1.0006533637635783e-5,
                           36762.161906015695,
                           29720.269687392283], [])
    ivps[:fhn][:ssr] = IVP(fhn, ini, t,
                           [0.1973538954650349,
                            0.11390293281470151,
                            3.0111993244825737], [])
    ivps[:fhn][:ds] = IVP(fhn, ini, t,
                          [0.06942005151231272,
                           0.4361123400524667,
                           1.696854260779007], [])
    ivps[:fhn][:dss] = IVP(fhn, ini, t,
                           [0.17695871191283072,
                            0.18851523532377604,
                            3.0158958746888795], [])
    return ivps
end

function transpose_vector(v::Vector{Vector{T}})::Vector{Vector{T}} where T
    [[v[i][j] for i in 1:length(v)] for j in 1:length(v[1])]
end

function solve_ivp(ivp::IVP)::Vector
    prob = ODEProblem(ivp.f, ivp.ini,
                      (ivp.t[1], ivp.t[end]), ivp.phi)
    osol  = solve(prob, AutoVern7(Rodas5()), saveat=ivp.t)
    return osol.u
end

function add_noise(data::Vector{T}, percent::Float64,
                   non_negative::Bool=true)::Vector{T} where T
    if percent == 0.0
        return data
    end

    epsilon_map = function(x)
        (isa(x, Vector{Vector{T}} where T) ?
         map(epsilon_map, x) :
         repeat([0.01*mean(x)], length(x)))
    end

    epsilon_iter = (map(epsilon_map, data)
                    |> Iterators.Flatten
                    |> Iterators.Stateful)

    noise_map = function(x)
        if isa(x, Vector)
            map(noise_map, x)
        else
            sigma = abs(percent*x + iterate(epsilon_iter)[1])
            d = Normal(0,sigma)
            noisy_x = x + rand(d)
            non_negative ? abs(noisy_x) : noisy_x
        end
    end

    noisy_data = map(noise_map, data)

    return noisy_data
end

"""
Normalized Mean Square Error
"""
function nmse(data::Vector{T}, data_est::Vector{T})::T where T
    normalizer = abs2(maximum(data) - minimum(data))
    res = mean(abs2.(data_est-data))/normalizer
    return res
end

"""
Normalized Root Mean Square Error
"""
function nrmse(data_nominal::Vector, data_estimated::Vector)::Float64
    state_pairs = zip(data_nominal, data_estimated)
    total_nmse = map(pair -> nmse(pair...),
                     state_pairs)
    # (mean) Normalized Root Mean Squared Error
    nrmse = mean(total_nmse) |> mean |> sqrt
    return nrmse
end

function gen_plt_err(ivp_comp::IVP,
                     ivp_est::IVP)::NamedTuple

    ini_cond = getindex.(ivp_comp.data, 1)
    ivp_est = IVP(ivp_est.f, ini_cond, ivp_est.t, ivp_est.phi)

    df_est = map(data -> DataFrame(x=ivp_est.t,
                                   y=data,
                                   Legend=repeat(["Estimated"],
                                                 length(ivp_est.t))),
                 ivp_est.data)

    df_comp = map(data -> DataFrame(x=ivp_comp.t,
                                   y=data,
                                   Legend=repeat(["Data"],
                                                 length(ivp_comp.t))),
                 ivp_comp.data)

    df_pairs = zip(df_comp, df_est)

    gen_layer(df) = layer(vcat(df...), x=:x, y=:y,
                          color=:Legend,
                          linestyle=:Legend,
                          Geom.line, Geom.point)

    layers = map(gen_layer, df_pairs)

    gen_plot(layers) = plot((layers...),
                            Guide.XLabel("t"),
                            Guide.YLabel("x"))
    plt = gen_plot(layers)

    nrmse_result = nrmse(ivp_comp.data, ivp_est.data)

    return (plt=plt, err=nrmse_result)
end

function get_res_from_dict(ivps::Dict)::Vector{NamedTuple}
    # Unpack first Dictionary layer into Vector of NamedTuple
    isest(x) = x != :nom && x != :test

    est_ivps = map(k -> (ivp=ivps[k], method=string(k)),
                   (keys(ivps) |> collect |> x -> filter(isest, x)))

    nom_test_ivps = map(k -> (ivp=ivps[k], method=string(k)),
                        (keys(ivps) |> collect |> x -> filter(!isest, x)))

    compare_pairs = [(i, j) for i in nom_test_ivps for j in est_ivps]

    results_list = map(x -> (gen_plt_err(x[1].ivp, x[2].ivp)
                             |> res -> (plt=res.plt,
                                        err=res.err,
                                        method=x[2].method,
                                        plt_type=x[1].method)),
                             compare_pairs)
    return results_list
end

function save_results(results::NamedTuple, path::AbstractString)::String
    base_path = joinpath(path, results.model)

    # Save plots
    map(x ->
        draw(PDF(joinpath(base_path, "$(x.plt_type)_$(x.method).pdf")), x.plt),
        results.res)

    # Save info
    io = joinpath(path, "info.log") |> f -> open(f, "a+")
    with_logger(SimpleLogger(io)) do
        map(x ->
            @info("NRMSE", x.plt_type, results.model, x.method, x.err),
            results.res)
    end
    close(io)
    return base_path
end

function main(args::Array{<:String})::Nothing
    path = string(args[1])

    str_mkdir(x) = (x
                    |> string
                    |> y -> joinpath(path, y)
                    |> z -> (isdir(z) ?
                             z : mkdir(z)))

    ivps = get_ivps()

    folders = map(str_mkdir,
                  (keys(ivps) |> collect))

    res_list = map(k -> (res=get_res_from_dict(ivps[k]),
                         model=string(k)),
                   (keys(ivps) |> collect))

    saved_files = map(r -> save_results(r, path), res_list)

    with_logger(ConsoleLogger()) do
        @debug res_list
        @debug saved_files
    end
    nothing
end

main(ARGS)


using CPLEX, Base.Test

@testset "LP 03" begin
    # Test get/set objective coefficients in LP

    env = CPLEX.Env()

    # original model
    #
    #   maximize  2x + 2y
    #
    # s.t. 0.2 <= x, y <= 1
    #

    model = CPLEX.cplex_model(env;
      name="lp_03",
      sense=:Max,
      f=[2.0, 2.0],
      lb=[0.2, 0.2],
      ub=[1.0, 1.0])

    lb_ = CPLEX.get_varLB(model)
    ub_ = CPLEX.get_varUB(model)
    c_ = CPLEX.get_obj(model)

    @test lb_ == [0.2, 0.2]
    @test ub_ == [1.0, 1.0]
    @test c_ == [2.0, 2.0]

    CPLEX.optimize!(model)

    sol = CPLEX.get_solution(model)
    @test sol[1] ≈ 1
    @test sol[2] ≈ 1
    @test CPLEX.get_objval(model) ≈ 4

    # change objective (warm start)
    #
    # maximize x - y
    #
    # s.t. 0.2 <= x, y <= 1
    #

    CPLEX.set_obj!(model, [1, -1])

    c_ = CPLEX.get_obj(model)
    @test c_ == [1.0, -1.0]

    CPLEX.optimize!(model)

    sol = CPLEX.get_solution(model)
    @test sol[1] ≈ 1
    @test sol[2] ≈ 0.2
    @test CPLEX.get_objval(model) ≈ 0.8
end


leftfactorial(n::Integer) = n ≤ 0 ? zero(n) : sum(factorial, 0:n-1)

@show leftfactorial.(0:10)
@show ndigits.(leftfactorial.(big.(1000:1000:10_000)))


# The formulas for `initialize_dual_potential!` are in Section 6.1.2 of SFVTP19.
function initialize_dual_potential!(::KL{ρ}, a::DiscreteMeasure) where {ρ}
    c = -ρ * log(length(a.log_density))
    a.dual_potential .= c
end

function initialize_dual_potential!(::TV{ρ}, a::DiscreteMeasure) where {ρ}
    c = -ρ * sign(log(length(a.log_density)))
    a.dual_potential .= c
end

# Specialized implementation for the KL-divergence
# Prop. 12 of SFVTP19.
function sinkhorn_divergence!(
    D::KL{ρ},
    C,
    a::DiscreteMeasure,
    b::DiscreteMeasure,
    ϵ = 1e-1;
    kwargs...,
) where {ρ}
    f = a.dual_potential
    g = b.dual_potential
    scaled_exp = z -> exp(-z / ρ)
    unbalanced_sinkhorn!(D, C, a, a, ϵ; kwargs...)
    term_aa = fdot(scaled_exp, a.density, f)
    unbalanced_sinkhorn!(D, C, b, b, ϵ; kwargs...)
    term_bb = fdot(scaled_exp, b.density, g)
    unbalanced_sinkhorn!(D, C, a, b, ϵ; kwargs...)
    term_ab = fdot(scaled_exp, a.density, f) + fdot(scaled_exp, b.density, g)
    return -(ρ + ϵ / 2) * (term_ab - term_aa - term_bb)
end

# Needed to avoid ambiguity errors
sinkhorn_divergence!(
    D::KL{ρ},
    C::AbstractMatrix,
    a::DiscreteMeasure,
    b::DiscreteMeasure,
    ϵ = 1e-1;
    kwargs...,
) where {ρ} = throw(ArgumentError("Must pass a cost function `C`, not a cost matrix."))


using RegERMs, Base.Test

# Seed random generator to avoid random effects in testing
srand(1)

X = [1 1; 2 2;  1 -1]
y = [-1; -1; 1]

@test_approx_eq_eps optimize(SVM(X, y), optimizer=:l_bfgs).theta [-0.1499984861841396,-0.3499984861803246] 1e-5
@test_approx_eq_eps optimize(SVM(X, y; λ=1.0), optimizer=:l_bfgs).theta [1.34394e-16,-1.0] 1e-5
@test_approx_eq_eps optimize(SVM(X, y; λ=10.1), optimizer=:l_bfgs).theta [1.34394e-16,-1.0] 1e-5
@test_approx_eq_eps optimize(SVM(X, y), optimizer=:sgd).theta [-0.1499984861841396,-0.3499984861803246] 5e-2
@test_approx_eq_eps optimize(SVM(X, y; λ=1.0), optimizer=:sgd).theta [1.34394e-16,-1.0] 5e-2
@test_approx_eq_eps optimize(SVM(X, y; λ=10.0), optimizer=:sgd).theta [1.34394e-16,-1.0] 5e-2
@test_approx_eq_eps optimize(SVM(X, y, L1RDAParameters(0.1,1.0,1.0)), optimizer=:l1_rda).theta [-14.811388300841895,-37.39005079444413] 5e-2
@test_approx_eq_eps optimize(SVM(X, y, L1RDAParameters(1.0,1.0,1.0)), optimizer=:l1_rda).theta [0.0,-8.866306299725341] 5e-2
@test_approx_eq_eps optimize(SVM(X, y, L1RDAParameters(10.0,1.0,1.0)), optimizer=:l1_rda).theta [0.0,0.0] 5e-2
model1 = optimize(SVM(X, y, L1RDAParameters(0.1,1.0,1.0)), optimizer=:l1_rda)
model2 = optimize(SVM(X, y, λ=10.0), optimizer=:sgd)

@test predict(model1, X) == y
@test predict(model2, X) == y
show(IOBuffer(), SVM(X, y))

@test_throws DimensionMismatch SVM(X', y)
@test_throws ArgumentError SVM(X, [3; 3; 2])
@test_throws ArgumentError optimize(SVM(X, y, L1RDAParameters(-0.1,1.0,1.0)), optimizer=:l1_rda)
@test_throws ArgumentError optimize(SVM(X, y, L1RDAParameters(0.1,-1.0,1.0)), optimizer=:l1_rda)
@test_throws ArgumentError optimize(SVM(X, y, L1RDAParameters(0.1,1.0,-1.0)), optimizer=:l1_rda)


"""
    Δθ = angle_in_plane(vec, mat)
    Δθ = angle_in_plane(vec1, vec2, mat)
    Δθ = angle_in_plane(vec, orbit)
    Δθ = angle_in_plane(vec, vec, orbit)
    Δθ = angle_in_plane(vec, orbit, t)
    Δθ = angle_in_plane(orbit1, orbit2, t)

Computes the angle between two vectors after projecting them to the orbital plane defined by an orbit's
basis matrix `mat`.

If only one vector is provided, the angle between the vector `vec` and the orbit's periapsis
is found.

If a vector, orbit, and time `t` are provided, the in-plane angle between the vector and the orbit's
position at `t` is found.

If two orbits and a time `t` are provide, the in-plane angle between the two orbits' positions at time
`t` is found.
"""
function angle_in_plane(vec::AbstractVector{<:Real}, mat::Union{Rotation{3,Float64},AbstractMatrix}=MRP(1I))
    planevec = mat\vec
    return atan(planevec[2],planevec[1])
end

angle_in_plane(vec::AbstractVector{<:Real}, vecref::AbstractVector{<:Real}, mat::Union{Rotation{3,Float64},AbstractMatrix}=MRP(1I)
    ) = wrap_angle(angle_in_plane(vec,mat) - angle_in_plane(vecref,mat))
angle_in_plane(vec::AbstractVector{<:Real}, vecref::AbstractVector{<:Real}, xdir::AbstractVector{<:Real}, ydir::AbstractVector{<:Real}
    ) = angle_in_plane(vec, vecref, basis_rotation(xdir,ydir))

# angle between the vector `vec` projected to the orbital plane and the periapsis of `orb`
angle_in_plane(vec::AbstractVector{<:Real}, orb::Orbit
    ) = angle_in_plane(vec, orb.basis)

# angle between the vectors `vec` and `vecref` after projecting both to the orbital plane of `orb`
angle_in_plane(vec::AbstractVector{<:Real}, vecref::AbstractVector{<:Real}, orb::Orbit
    ) = angle_in_plane(vec, vecref, orb.basis)

# angle between the vector `vec` projected to the orbital plane and the orbit's position at time `t`
angle_in_plane(vec::AbstractVector{<:Real}, orb::Orbit, t::Real
    ) = wrap_angle(angle_in_plane(vec, time_state_vector(t, orb)[1], orb))

# angle between the orbital positions of `orb` from `orbref` at time `t`, after projecting the position `orb` to the orbital plane of `orbref`
angle_in_plane(orb::Orbit, orbref::Orbit, t::Real
    ) = wrap_angle(angle_in_plane(time_state_vector(t, orb)[1], orbref, t))


using MambaModels

# Data

globe_toss = Dict{Symbol, Any}(
  :w => [6, 7, 5, 6, 6],
  :n => [9, 9, 9, 9, 9]
)
globe_toss[:N] = length(globe_toss[:w]);

# Model Specification

model = Model(
  w = Stochastic(1,
    (n, p, N) ->
      UnivariateDistribution[Binomial(n[i], p) for i in 1:N],
    false
  ),
  p = Stochastic(() -> Beta(1, 1))
);

# Initial Values

inits = [
  Dict(:w => globe_toss[:w], :n => globe_toss[:n], :p => 0.5),
  Dict(:w => globe_toss[:w], :n => globe_toss[:n], :p => rand(Beta(1, 1)))
];

# Sampling Scheme

scheme = [NUTS(:p)]
setsamplers!(model, scheme);

# MCMC Simulations

chn = mcmc(model, globe_toss, inits, 10000, burnin=2500, thin=1, chains=2);

# Describe draws

describe(chn)

# Convert to MCMCChains.Chains object

chn2 = MCMCChains.Chains(chn.value, String.(chn.names))

# Describe the MCMCChains

MCMCChains.describe(chn2)

# End of `02/m2.1m.jl`