module LotkaVolterra2dProblem

    using GeometricIntegrators.Equations

    export lotka_volterra_2d_ode,  lotka_volterra_2d_pdae,
           lotka_volterra_2d_iode, lotka_volterra_2d_idae,
           lotka_volterra_2d_vode, lotka_volterra_2d_vdae,
           lotka_volterra_2d_dg, lotka_volterra_2d_vdae_slrk

    Δt = 0.01
    nt = 10

    const A1=1.0
    const A2=1.0
    const B1=1.0
    const B2=2.0

    const X0=1.0
    const Y0=1.0

    function ϑ₁(t, q)
        q[2] + log(q[2]) / q[1]
    end

    function ϑ₂(t, q)
        q[1]
    end

    function v₁(t, q)
        q[1] * (A2*q[2] - B2)
    end

    function v₂(t, q)
        q[2] * (B1 - A1*q[1])
    end

    const q₀=[X0, Y0]
    const v₀=[v₁(0, q₀), v₂(0, q₀)]

    function dϑ₁d₁(t, q)
        - log(q[2]) / q[1]^2
    end

    function dϑ₁d₂(t, q)
        1 + 1 / (q[1] * q[2])
    end

    function dϑ₂d₁(t, q)
        one(eltype(q))
    end

    function dϑ₂d₂(t, q)
        zero(eltype(q))
    end

    function d²ϑ₁d₁d₁(t, q)
        + 2 * log(q[2]) / q[1]^3
    end

    function d²ϑ₁d₁d₂(t, q)
        - 1 / (q[1]^2 * q[2])
    end

    function d²ϑ₁d₂d₁(t, q)
        - 1 / (q[1]^2 * q[2])
    end

    function d²ϑ₁d₂d₂(t, q)
        - 1 / (q[1] * q[2]^2)
    end

    function d²ϑ₂d₁d₁(t, q)
        zero(eltype(q))
    end

    function d²ϑ₂d₁d₂(t, q)
        zero(eltype(q))
    end

    function d²ϑ₂d₂d₁(t, q)
        zero(eltype(q))
    end

    function d²ϑ₂d₂d₂(t, q)
        zero(eltype(q))
    end

    function ϑ(t::Number, q::AbstractVector, Θ::AbstractVector)
        Θ[1] = ϑ₁(t,q)
        Θ[2] = ϑ₂(t,q)
        nothing
    end

    function ϑ(t::Number, q::AbstractVector)
        [ϑ₁(t,q), ϑ₂(t,q)]
    end

    function ω(t, q, Ω)
        Ω[1,1] = 0
        Ω[1,2] = dϑ₁d₂(t,q) - dϑ₂d₁(t,q)

        Ω[2,1] = dϑ₂d₁(t,q) - dϑ₁d₂(t,q)
        Ω[2,2] = 0

        nothing
    end

    function f₁(t, q, v)
        dϑ₁d₁(t,q) * v[1] + dϑ₂d₁(t,q) * v[2]
    end

    function f₂(t, q, v)
        dϑ₁d₂(t,q) * v[1] + dϑ₂d₂(t,q) * v[2]
    end

    function g₁(t, q, v)
        dϑ₁d₁(t,q) * v[1] + dϑ₁d₂(t,q) * v[2]
    end

    function g₂(t, q, v)
        dϑ₂d₁(t,q) * v[1] + dϑ₂d₂(t,q) * v[2]
    end

    function g̅₁(t, q, v)
        d²ϑ₁d₁d₁(t,q) * q[1] * v[1] + d²ϑ₁d₂d₁(t,q) * q[1] * v[2] + d²ϑ₂d₁d₁(t,q) * q[2] * v[1] + d²ϑ₂d₂d₁(t,q) * q[2] * v[2]
    end

    function g̅₂(t, q, v)
        d²ϑ₁d₁d₂(t,q) * q[1] * v[1] + d²ϑ₁d₂d₂(t,q) * q[1] * v[2] + d²ϑ₂d₁d₂(t,q) * q[2] * v[1] + d²ϑ₂d₂d₂(t,q) * q[2] * v[2]
    end

    function hamiltonian(t, q)
        A1*q[1] + A2*q[2] - B1*log(q[1]) - B2*log(q[2])
    end

    function dHd₁(t, q)
        A1 - B1 / q[1]
    end

    function dHd₂(t, q)
        A2 - B2 / q[2]
    end

    function lotka_volterra_2d_dH(t, q, dH)
        dH[1] = dHd₁(t, q)
        dH[2] = dHd₂(t, q)
        nothing
    end

    function lotka_volterra_2d_ϑ(t, q, Θ)
        ϑ(t, q, Θ)
    end

    function lotka_volterra_2d_ϑ(t, q, v, Θ)
        ϑ(t, q, Θ)
    end

    function lotka_volterra_2d_ϑ(κ, t, q, v, Θ)
        Θ[1] = (1-κ) * ϑ₁(t,q) - κ * f₁(t,q,q)
        Θ[2] = (1-κ) * ϑ₂(t,q) - κ * f₂(t,q,q)
        nothing
    end

    function lotka_volterra_2d_ω(t, q, Ω)
        ω(t, q, Ω)
    end

    function lotka_volterra_2d_v(t, q, v)
        v[1] = v₁(t, q)
        v[2] = v₂(t, q)
        nothing
    end

    function lotka_volterra_2d_v(t, q, p, v)
        lotka_volterra_2d_v(t, q, v)
    end

    function lotka_volterra_2d_v_ham(t, q, p, v)
        v[1] = 0
        v[2] = 0
        nothing
    end

    function lotka_volterra_2d_f(t::Real, q::Vector, v::Vector, f::Vector)
        f[1] = f₁(t,q,v) - dHd₁(t,q)
        f[2] = f₂(t,q,v) - dHd₂(t,q)
        nothing
    end

    function lotka_volterra_2d_f(κ::Real, t::Real, q::Vector, v::Vector, f::Vector)
        f[1] = (1-κ) * f₁(t,q,v) - κ * (g₁(t,q,v) + g̅₁(t,q,v)) - dHd₁(t,q)
        f[2] = (1-κ) * f₂(t,q,v) - κ * (g₂(t,q,v) + g̅₂(t,q,v)) - dHd₂(t,q)
        nothing
    end

    function lotka_volterra_2d_f_ham(t::Real, q::Vector, f::Vector)
        f[1] = - dHd₁(t,q)
        f[2] = - dHd₂(t,q)
        nothing
    end

    function lotka_volterra_2d_f_ham(t::Real, q::Vector, v::Vector, f::Vector)
        lotka_volterra_2d_f_ham(t, q, f)
    end

    function lotka_volterra_2d_g(t::Real, q::Vector, v::Vector, g::Vector)
        g[1] = f₁(t,q,v)
        g[2] = f₂(t,q,v)
        nothing
    end

    function lotka_volterra_2d_g(t::Real, q::Vector, p::Vector, v::Vector, g::Vector)
        lotka_volterra_2d_g(t, q, v, g)
    end

    function lotka_volterra_2d_g(κ::Real, t::Real, q::Vector, v::Vector, g::Vector)
        g[1] = (1-κ) * f₁(t,q,v) - κ * (g₁(t,q,v) + g̅₁(t,q,v))
        g[2] = (1-κ) * f₂(t,q,v) - κ * (g₂(t,q,v) + g̅₂(t,q,v))
        nothing
    end

    # function lotka_volterra_2d_g(κ::Real, t::Real, q::Vector, v::Vector, g::Vector)
    #     g[1] = (1-κ) * g₁(t,q,v) - κ * g̅₁(t,q,v) - κ * f₁(t,q,v)
    #     g[2] = (1-κ) * g₂(t,q,v) - κ * g̅₂(t,q,v) - κ * f₂(t,q,v)
    #     nothing
    # end

    function lotka_volterra_2d_g̅(t::Real, q::Vector, v::Vector, g::Vector)
        g[1] = g₁(t,q,v)
        g[2] = g₂(t,q,v)
        nothing
    end

    function lotka_volterra_2d_g̅(t::Real, q::Vector, p::Vector, v::Vector, g::Vector)
        lotka_volterra_2d_g̅(t, q, v, g)
    end

    function lotka_volterra_2d_u(t, q, v, u)
        u[1] = v[1]
        u[2] = v[2]
        nothing
    end

    function lotka_volterra_2d_u(t, q, p, v, u)
        lotka_volterra_2d_u(t, q, v, u)
    end

    function lotka_volterra_2d_ϕ(t, q, p, ϕ)
        ϕ[1] = p[1] - ϑ₁(t,q)
        ϕ[2] = p[2] - ϑ₂(t,q)
        nothing
    end

    function lotka_volterra_2d_ψ(t, q, p, v, f, ψ)
        ψ[1] = f[1] - g₁(t,q,v)
        ψ[2] = f[2] - g₂(t,q,v)
        nothing
    end

    function lotka_volterra_2d_ode(q₀=q₀)
        ODE(lotka_volterra_2d_v, q₀; h=hamiltonian)
    end

    function lotka_volterra_2d_iode(q₀=q₀, p₀=ϑ(0, q₀))
        IODE(lotka_volterra_2d_ϑ, lotka_volterra_2d_f,
             lotka_volterra_2d_g, q₀, p₀;
             h=hamiltonian, v=lotka_volterra_2d_v)
    end

    function lotka_volterra_2d_vode(q₀=q₀, p₀=ϑ(0, q₀))
        VODE(lotka_volterra_2d_ϑ, lotka_volterra_2d_f,
             lotka_volterra_2d_g, q₀, p₀;
             v=lotka_volterra_2d_v, Ω=lotka_volterra_2d_ω, ∇H=lotka_volterra_2d_dH)
    end

    function lotka_volterra_2d_idae(q₀=q₀, p₀=ϑ(0, q₀), λ₀=zero(q₀))
        IDAE(lotka_volterra_2d_ϑ, lotka_volterra_2d_f,
             lotka_volterra_2d_u, lotka_volterra_2d_g,
             lotka_volterra_2d_ϕ, q₀, p₀, λ₀;
             v=lotka_volterra_2d_v)
    end

    function lotka_volterra_2d_pdae(q₀=q₀, p₀=ϑ(0, q₀), λ₀=zero(q₀))
        PDAE(lotka_volterra_2d_v_ham, lotka_volterra_2d_f_ham,
             lotka_volterra_2d_u, lotka_volterra_2d_g,
             lotka_volterra_2d_ϕ, q₀, p₀, λ₀)
    end

    function lotka_volterra_2d_vdae(q₀=q₀, p₀=ϑ(0, q₀), λ₀=zero(q₀))
        VDAE(lotka_volterra_2d_ϑ, lotka_volterra_2d_f_ham,
             lotka_volterra_2d_g, lotka_volterra_2d_g̅,
             lotka_volterra_2d_ϕ, lotka_volterra_2d_ψ,
             q₀, p₀, λ₀; v=lotka_volterra_2d_v)
    end

    function lotka_volterra_2d_vdae_slrk(q₀=q₀, p₀=ϑ(0, q₀), λ₀=zero(q₀))
        VDAE(lotka_volterra_2d_ϑ, lotka_volterra_2d_f,
             lotka_volterra_2d_g, lotka_volterra_2d_g̅,
             lotka_volterra_2d_ϕ, lotka_volterra_2d_ψ,
             q₀, p₀, λ₀; v=lotka_volterra_2d_v)
    end

    function lotka_volterra_2d_dg(q₀=q₀; κ=0)
        lotka_volterra_2d_ϑ_κ(t, q, v, p) = lotka_volterra_2d_ϑ(κ, t, q, v, p)
        lotka_volterra_2d_f_κ(t, q, v, f) = lotka_volterra_2d_f(κ, t, q, v, f)
        lotka_volterra_2d_g_κ(t, q, λ, g) = lotka_volterra_2d_g(κ, t, q, λ, g)

        IODE(lotka_volterra_2d_ϑ_κ, lotka_volterra_2d_f_κ,
             lotka_volterra_2d_g_κ, q₀, p₀;
             v=lotka_volterra_2d_v)
    end

end


using Distributions, ColorSchemes
using Plots
gr()

export fire

"""
cannonball_trajectory(v::Real,θ::Real,t)

this function return the cannonball [x,y] coords at time t for
for initial data v,θ
"""
function cannonball_trajectory(θ::Real, v::Real, width::Real, t)
    g = 9.81
    x,y = 0.5*width .+ v*cos(θ).*t, v*sin(θ).*t - 0.5*g*t.^2
    return (x,y)
end # function cannonball_trajectory

"""
    random_tries(gen_size::Real; max_speed::Real = 20)

this function create random population of size = 'size' for the cannonball
for the velocity 'v' and angle 'θ'
v ∈ [0, typemax(Int64)]
θ ∈ [0, π]
"""
function random_tries(gen_size::Real; max_speed::Real = 20)::Matrix
    generation = Matrix{Real}(undef,gen_size,2)
    for ii in 1:gen_size
        generation[ii,1] = π * rand()
        generation[ii,2] = max_speed * rand()
    end
    return generation
end # function random_tries

"""
hit_coordinate(θ::Real,v::Real,width::Real)

this function find the coordinates of the cannonballs at the bag edge
"""
function hit_coordinate(θ::Real,v::Real,width::Real)
    x = 0.5 * width
    x_hit = width
    if θ > π/2
        x = -x
        x_hit = 0
    end
    t = x/(v*cos(θ))
    y = v*t*sin(θ) - 0.5 * 9.81 * t^2
    y = y < 0 ? 0 : y
    return x_hit, y
 end # function hit_coordinate

"""
    escaped(θ::Real, v::Real, width::Real, height::Real)

this function check if the cannonball escaped the bag
"""
function escaped(θ::Real, v::Real, width::Real, height::Real)
    x_hit, y_hit = hit_coordinate(θ, v, width)
    return (x_hit == 0 || x_hit == width) && y_hit > height
end # function escaped

"""
    cumulative_probabilities(results)

this function build a vector of comulative fittnes results
"""
function cumulative_probabilities(results)
    cprob = Vector{Real}(undef,length(results))
    total = 0.0
    for (ii,res) in enumerate(results)
        total += res[2]
        cprob[ii] = total
    end
    return cprob
end # function cumulative_probabilities

"""
    selection(generation::Matrix,width::Real)

The function appends the y-coordinate where each ball hits the bag
edge to a list, The cumulative_probabilities function then stacks up the running
total of these, giving you your roulette wheel.
"""
function selection(generation::Matrix,width::Real)
    results = [hit_coordinate(θ,v,width) for (θ,v) in zip(generation[:,1],generation[:,2])]
    return cumulative_probabilities(results)
end # function selection

"""
    choose(choices)

The function spins the roulette wheel by picking a random number and
seeing which solution it corresponds to.
"""
function choose(choices)
    p = rand(Uniform(0,choices[end]))
    for (ind,val) in enumerate(choices)
        val >= p && return ind
    end
end # function choose

"""
crossover(generation::Matrix, Width::Real)

Armed with two parents, breed new solutions Not all genetic algorithms
use two parents; some just use one, and there are others that use more.
However, two is a conventional approach.
Notice this makes a new generation of the same size as the last generation.
"""
function crossover(generation::Matrix, width::Real)
    choices = selection(generation,width)
    next_generation = Matrix{Real}(undef,size(generation))
    for (ind,val) in enumerate(eachrow(generation))
        mum = generation[choose(choices),:]
        dad = generation[choose(choices),:] # TODO - why not taking out the pair choosen for mum?
        next_generation[ind,:] = collect(breed(mum,dad))
    end
    return next_generation
end # function crossover

"""
    breed(mum::Vector, dad::Vector)

In order to breed, the information is split—half from one parent and half from
another. Since there are two bits of information, each child will have a velocity
from one parent and an angle from the other. There is only one way to
split this.
In this example, the parents produce a single child.
"""
function breed(mum::Vector, dad::Vector)
    return mum[1],dad[2]
end # function breed

"""
    mutate(generation::Matrix; v_max::Real = 30, v_min::Real = 10)

the function change everything in a population, to keep it general.
the function uses probabilistic mutation, that is you want something to happen
on average, one time out of ten, One approac, is to get a random number between
0 and 1 and do that “something” if you get less than 0.1.
How do you mutate the value? It’s traditional to add or subtract a small
amount or scale (multiply) a little for real numbers. in our case
add a random number to the angle, but only use this if it stays between
0 and 180 degrees. and for the velocity, scale by something random between 0.9
and 1.1, cunningly avoiding the problem of potentially getting zero
"""
function mutate(generation::Matrix; v_max::Real = 30, v_min::Real = 10)
    for (ind,val) in enumerate(eachrow(generation))
        θ = val[1]
        v = val[2]
        if rand() < 0.1
            θ_new = θ + (rand(Uniform(-10,10)) |> deg2rad)
            if 0 < θ_new < π
                θ = θ_new
            end
        end
        if rand() < 0.1
            v *= rand(Uniform(0.9,1.1))
            if v > v_max
                v = v_max
            elseif v < v_min
                v = v_min
            end
        end
        generation[ind,:] = [θ,v]
    end
end # function mutate

"""
    display_epoch!(p::Plots.Plot, generation::Matrix, width::Real, height::Real)

this function plot the current epoch
"""
function display_epoch!(p::Plots.Plot, generation::Matrix, width::Real, height::Real)
    pcolor = rand(ColorSchemes.Accent_5.colors)
    markertype = rand(Plots.supported_markers())
    for (θ,v) in eachrow(generation)
        tend = (2*v*sin(θ))/9.81
        x,y = cannonball_trajectory(θ,v,width,range(0,tend[1],length=200))
        if escaped(θ, v, width, height)
            plot!(p, x, y, color = pcolor)
        else
            plot!(p, x, y, color = :red)
        end
        scatter!(p, x[1:Int(0.1*length(y)):end], y[1:Int(0.1*length(y)):end], marker = markertype, color = pcolor)
    end
end # function display_epoch

"""
    display_epoch(p::Plots.Plot, generation::Matrix, width::Real, height::Real)

this function plot the current epoch
"""
function display_epoch(pin::Plots.Plot, generation::Matrix, width::Real, height::Real)
    p = deepcopy(pin)
    for (θ,v) in eachrow(generation)
        tend = (2*v*sin(θ))/9.81
        x,y = cannonball_trajectory(θ,v,width,range(0,tend[1],length=200))
        if escaped(θ, v, width, height)
            plot!(p, x, y, color = :blue)
            scatter!(p, x[1:Int(0.1*length(y)):end], y[1:Int(0.1*length(y)):end], marker = :circle, color = :blue)
        else
            plot!(p, x, y, color = :red)
            scatter!(p, x[1:Int(0.1*length(y)):end], y[1:Int(0.1*length(y)):end], marker = :circle, color = :red)
        end
    end
    return p
end # function display_epoch

"""
    fire()

the main cannonball function
"""
function fire()
    epochs = 100
    items = 100
    height = 10
    width = 10
    max_cennonball_speed = 20
    box =( x = vec([0 0 width width]), y = vec([height 0 0 height]))
    generation = random_tries(items, max_speed = max_cennonball_speed)
    generation0 = deepcopy(generation) # save to contrast with last epoch
    # p = plot(box.x, box.y, legend = false, lw = 3, aspect_ratio=:equal)
    pbox = plot(box.x, box.y, legend = false, lw = 3)
    titles = ["First generation","Last generation"]
    for ii in 1:epochs
        results = []
        generation = crossover(generation, width)
        mutate(generation)
    end
    genPlot = (generation0,generation)
    pout = Vector{Any}(undef,length(genPlot))
    for (ind,gen) in enumerate(genPlot)
        pout[ind] = display_epoch(pbox, gen, width, height)
        title!(pout[ind],titles[ind])
        xlabel!(pout[ind],"X")
        ylabel!(pout[ind],"Y")
    end
    p = plot(pout[1],pout[2],layout = (2,1))
    display(p)
end # function fire


using CompScienceMeshes, BEAST
using LinearAlgebra
using Profile
using StaticArrays

function tau(x::SVector{U,T}) where {U,T}
    4.0-1.0
end

ttrc = X->ttrace(X,y)

T= tetmeshsphere(1.0,0.25)
X = nedelecc3d(T)
y = boundary(T)
@show numfunctions(X)

ϵ, μ, ω = 1.0, 1.0, 1.0; κ, η = ω * √(ϵ*μ), √(μ/ϵ)
ϵ_r = 4.0
χ = tau
K, I, B = VIE.singlelayer2(wavenumber=κ, tau=χ), Identity(), VIE.boundary2(wavenumber=κ, tau=χ)
E = VIE.planewave(direction=ẑ, polarization=x̂, wavenumber=κ)
H = -1/(im*κ*η)*curl(E)

@hilbertspace j
@hilbertspace k
α = ϵ_r
eq = @varform α*I[j,k]-K[j,k]+B[ttrc(j),k] == E[j]

evie = @discretise eq j∈X k∈X
u = solve(evie)

#postprocessing
Φ, Θ = [0.0], range(0,stop=2π,length=100)
pts = [point(cos(ϕ)*sin(θ), sin(ϕ)*sin(θ), cos(θ)) for θ in Θ for ϕ in Φ ]
ffe = potential(VIE.farfield(wavenumber=κ, tau=χ), pts, u, X)
ff = ffe

using Plots

#Farfield
plot(xlabel="theta")
plot!(Θ, norm.(ff), label="far field", title="E-VIE")

#Nearfield
Z = range(-1,1,length=100)
Y = range(-1,1,length=100)
nfpoints = [point(0,y,z) for  y in Y, z in Z]

Enear = BEAST.grideval(nfpoints,u,X)
Enear = reshape(Enear,100,100)

contour(real.(getindex.(Enear,1)))
heatmap(Z, Y, real.(getindex.(Enear,1)), clim=(0.0,1.0))


# ------------------------------------------------------------------------------
# lumped BDDC on macro elements example
# ------------------------------------------------------------------------------

using LinearAlgebra

# setup
mesh = Mesh2D(1.0, 1.0)
p = 1
numberelements1d = 4

# operator
diffusion = GalleryMacroElementOperator("diffusion", p + 1, p + 2, numberelements1d, mesh)

# BDDC smoother
bddc = DirichletBDDC(diffusion)

# compute operator symbols
A = computesymbols(bddc, [0.2], [π, π])
eigenvalues = real(eigvals(A))

# ------------------------------------------------------------------------------


using Test
using TestSetExtensions
using LinearAlgebra
using Random
using Qaintessent

##==----------------------------------------------------------------------------------------------------------------------

@testset ExtendedTestSet "statevector" begin

    @testset "test statevector int constructor" begin
        N = rand(1:10)
        s = Statevector(N)
        @test s.N == N
        @test s.state == zeros(ComplexF64, 2^N)
    end

    @testset "test statevector vector constructor" begin
        N = rand(1:10)
        ψ = randn(ComplexF64, 2^N)
        s = Statevector(ψ)
        @test s.N == N
        @test s.state == ψ
    end

    @testset "test statevector indexing" begin
        N = rand(1:10)
        ψ = randn(ComplexF64, 2^N)
        s = Statevector(ψ)
        r = rand(1:2^N)
        @test s[r] == ψ[r]
    end

    @testset "test statevector attributes" begin
        N = rand(1:10)
        ψ = randn(ComplexF64, 2^N)
        s = Statevector(ψ)
        @test size(s) == size(ψ)
        @test length(s) == length(ψ)
        @test IndexStyle(typeof(s)) == IndexLinear()
    end

    @testset ExtendedTestSet "util reduced density matrix for statevector" begin
        N = 4
        ψ = Statevector(randn(ComplexF64, 2^N))
        χ = Statevector(randn(ComplexF64, 2^N))
        # full density matrix |ψ⟩⟨χ|
        ρ = reshape(kron(conj(χ.state), ψ.state), 2^N, 2^N)
        id = Matrix{ComplexF64}(I, 2, 2)
        A = randn(ComplexF64, 2, 2)
        B = randn(ComplexF64, 2, 2)
        @testset ExtendedTestSet "reduced density matrix correctness" begin
            @test sum(kron(A, B) .* rdm(N, (4, 2), ψ, χ)) ≈ sum(kron(B, id, A, id) .* ρ)
        end

        @testset ExtendedTestSet "reduced density matrix exceptions" begin
            @test_throws ErrorException("Need at least one wire to act on.") rdm(N, (), ψ, χ)
            @test_throws ErrorException("Number of gate wires cannot be larger than total number of wires.") rdm(N, (1, 2, 3, 4, 5), ψ, χ)
            @test_throws ErrorException("Wire indices must be unique.") rdm(N, (2, 2), ψ, χ)
            @test_throws ErrorException("Wire index cannot be smaller than 1.") rdm(N, (-1, 2), ψ, χ)
            @test_throws ErrorException("Wire index cannot be larger than total number of wires.") rdm(N, (5, 1), ψ, χ)
        end
    end

    @testset ExtendedTestSet "util reduced density matrix for statevector" begin
        N = 4
        ψ = Statevector(randn(ComplexF64, 2^N))
        χ = Statevector(randn(ComplexF64, 2^N))
        # full density matrix |ψ⟩⟨χ|
        ρ = reshape(kron(conj(χ.state), ψ.state), 2^N, 2^N)
        id = Matrix{ComplexF64}(I, 2, 2)
        A = randn(ComplexF64, 2, 2)
        B = randn(ComplexF64, 2, 2)
        @testset ExtendedTestSet "reduced density matrix correctness" begin
            @test sum(kron(A, B) .* rdm(N, (4, 2), ψ, χ)) ≈ sum(kron(B, id, A, id) .* ρ)
        end

        @testset ExtendedTestSet "reduced density matrix exceptions" begin
            @test_throws ErrorException("Need at least one wire to act on.") rdm(N, (), ψ, χ)
            @test_throws ErrorException("Number of gate wires cannot be larger than total number of wires.") rdm(N, (1, 2, 3, 4, 5), ψ, χ)
            @test_throws ErrorException("Wire indices must be unique.") rdm(N, (2, 2), ψ, χ)
            @test_throws ErrorException("Wire index cannot be smaller than 1.") rdm(N, (-1, 2), ψ, χ)
            @test_throws ErrorException("Wire index cannot be larger than total number of wires.") rdm(N, (5, 1), ψ, χ)
        end
    end

end


"""
Allometry. Scaling rules that relate size to mass
"""
abstract type AbstractAllometry end

update_height!(o, u) = update_height!(allometry_pars(o), o, u)
update_height!(f::Nothing, o, u) = nothing
update_height!(f::AbstractAllometry, o, u) =
    set_height!(o, allometric_height(f, mass(o, u)))

"""
    SqrtAllometry(β0, β)

Height is given by the square root of mass above the initial mass `β0`,
multiplied by the scalar `β1`.

$(FIELDDOCTABLE)
"""
@flattenable @columns struct SqrtAllometry{B0,B1} <: AbstractAllometry
    # Field       | Flatn | Default | Unit | Bounds         | Log  | Description
    β0::B0        | false | 1e-4*24 | g    | (1e-6, 10.00)  | true | "Intercept. Mass at Om"
    β1::B1        | true  | 0.1     | m    | (1e-2, 1.0)    | true | "Scalar for conversion to meters"
end

allometric_height(f::SqrtAllometry, mass) =
    sqrt((mass - f.β0) / unit(f.β0)) * f.β1

"""
    Allometry(β0, β, α)

Simple allometric relationship between mass and height

$(FIELDDOCTABLE)
"""
@flattenable @columns struct Allometry{B0,B1,A} <: AbstractAllometry
    β0::B0        | false | 1e-4*24 | g    | (1e-6, 10.00)  | true | "Intercept. Mass at Om"
    β1::B1        | true  | 0.1     | m    | (1e-5, 10.00)  | true | "Scalar for conversion to meters"
    α::A          | true  | 0.1     | _    | (1e-3, 100.00) | true | "Exponent relating mass to vertical dimension"
end

allometric_height(f::Allometry, mass) =
    f.β1 * ((max(zero(mass), mass - f.β0)) / unit(f.β0))^f.α

"""
    FixedAllometry(height)

Height is fixed at `height`, independent of mass.

$(FIELDDOCTABLE)
"""
@flattenable @columns struct FixedAllometry{M} <: AbstractAllometry
    height::M     | true  | 1.0     | m    | (1e-3, 100.00) | true | "Fixed height or depth"
end

allometric_height(f::FixedAllometry, mass) = f.height

mass(o, u) = u[:V] * w_V(o)


module HALeqO

export haleqo

using SolverCore, NLPModels
using LinearAlgebra, SparseArrays
using HSL, PositiveFactorizations

"""
    normsq(v)
"""
normsq(v) = dot(v, v)

"""
    merit(y, yhat, μ, fx, cx)
"""
merit(y, yhat, μ, fx, cx) =
    fx + 0.25 * μ * normsq(y) + normsq(cx + μ .* (yhat - 0.5 * y)) / μ

"""
    haleqo( nlp )
"""
function haleqo(
    nlp::AbstractNLPModel;
    x::AbstractVector = copy(nlp.meta.x0),
    y::AbstractVector = copy(nlp.meta.y0),
    tol::Real = 1e-8,
    μ::Real = 0.1,
    max_iter::Int = 3000,
    max_time::Real = 300.0,
    max_eval::Int = 100000,
    use_filter::Bool = false,
)

    start_time = time()
    nx, ny = nlp.meta.nvar, nlp.meta.ncon
    T = eltype(x)

    if has_bounds(nlp) || inequality_constrained(nlp)
        error("Problem has inequalities, can't solve it")
    end

    # parameters
    θ = 0.5
    ls_beta = 0.5
    κμminus = 0.1
    ls_eta = 1e-4
    μmin = 1e-16

    # initialization
    iter = 0
    eltime = 0.0
    status = :unknown
    fx = obj(nlp, x)
    cx = cons(nlp, x)
    yhat = copy(y)
    subres = zeros(T, nx + ny)
    subres[1:nx] .= grad(nlp, x) + jtprod(nlp, x, y)
    subres[nx+1:nx+ny] .= cx
    optimality = norm(subres[1:nx], Inf)
    cviolation = norm(cx, Inf)
    cviol__old = cviolation
    residy = cviolation
    is_solved = optimality ≤ tol && cviolation ≤ tol
    is_infeasible = false
    is_tired = iter ≥ max_iter || eltime ≥ max_time || neval_obj(nlp) > max_eval
    dir = zeros(T, nx + ny)
    xold = zeros(T, nx)
    yold = zeros(T, ny)

    rhoBM = (1.0 / μ) * max(1.0, fx) / max(1.0, 0.5*normsq(cx)) # without abs()
    rhoBM = max(1e-4, min(rhoBM, 1e4))
    μ = 1.0 / rhoBM

    if use_filter
        phi_beta = 0.1 # 0 < beta \le 1
        phi_V(cviol, optim) = cviol + phi_beta * optim
        phi_O(cviol, optim) = cviol * phi_beta + optim
        phi_theta = 0.1 # 0 < theta < 1
        phi_V_max = phi_V(cviolation, optimality)
        phi_O_max = phi_O(cviolation, optimality)
        phi_V_max *= phi_theta
        phi_O_max *= phi_theta
    end

    @info log_header(
        [:iter, :fx, :cviol, :optim, :μ, :resy],
        [Int, T, T, T, T, T],
        hdr_override = Dict(:fx => "f(x)", :optim => "‖∇L‖", :cviol => "‖c(x)‖"),
    )
    @info log_row(Any[iter, fx, cviolation, optimality])

    while !(is_solved || is_infeasible || is_tired)

        # search direction
        # compute Newton direction for the regularized sub-problem, possibly with
        # some perturbation of the Hessian matrix to obtain a descent direction
        # for the merit function
        H = hess(nlp, x, y)
        H = Symmetric(SparseMatrixCSC(Matrix(cholesky(Positive, H, Val{false}))))
        J = jac(nlp, x)

        KKT = [H J'; J UniformScaling(-μ)]
        dir .= -subres
        LDLT = Ma57(KKT)
        ma57_factorize(LDLT)
        ma57_solve!(LDLT, dir)

        # gradient of merit function
        # ∇x merit = resx + (2/μ) resy
        # ∇y merit = - resy
        xold .= subres[1:nx] + (2.0 / μ) .* jtprod(nlp, x, subres[nx+1:nx+ny])
        # slope of merit along search direction
        # slope = ∇merit ⋅ dir
        slope = dot(dir[1:nx], xold) - dot(dir[nx+1:nx+ny], subres[nx+1:nx+ny])
        if slope ≥ 0.0
            status = :not_desc
            break
        end

        # line-search
        # backtracking line-search along the primal-dual augmented Lagrangian
        # merit function with Armijo sufficient decrease condition
        xold .= x
        yold .= y
        mold = merit(y, yhat, μ, fx, cx)
        slope *= ls_eta
        τ = one(T)
        x .+= dir[1:nx]
        y .+= dir[nx+1:nx+ny]
        while true
            fx = obj(nlp, x)
            cons!(nlp, x, cx)
            m = merit(y, yhat, μ, fx, cx)
            # check Armijo's condition
            if m ≤ mold + τ * slope
                break
            end
            τ *= ls_beta
            x .= xold + τ .* dir[1:nx]
            y .= yold + τ .* dir[nx+1:nx+ny]
        end

        # evaluate residuals and check termination
        subres[1:nx] .= grad(nlp, x) + jtprod(nlp, x, y)
        optimality = norm(subres[1:nx], Inf)
        cviolation = norm(cx, Inf)
        is_solved = optimality ≤ tol && cviolation ≤ tol

        iter += 1
        if !is_solved
            eltime = time() - start_time
            is_tired = iter ≥ max_iter || eltime ≥ max_time || neval_obj(nlp) > max_eval
            if !is_tired
                is_infeasible =
                    μ < μmin && norm(jtprod(nlp, x, cx), Inf) < cviolation * √tol
                if !is_infeasible
                    subres[nx+1:nx+ny] .= cx + μ .* (yhat - y)
                    residy = norm(subres[nx+1:nx+ny], Inf)
                end
            end
        end

        # sub-problem update
        if optimality ≤ tol && residy ≤ tol
            iter_type = :M
            # check improvement in constraint violation
            if cviolation > θ * cviol__old
                # update dual regularization parameter
                μ *= κμminus
            end
            # update dual estimate
            yhat .= y
            # re-evaluate some quantities
            subres[nx+1:nx+ny] .= cx
            cviol__old = cviolation
            residy = cviolation
            @info log_row(Any[iter, fx, cviolation, optimality, μ, "$(iter_type)"])
        elseif use_filter && phi_V(cviolation, optimality) ≤ phi_V_max
            iter_type = :V
            phi_V_max *= phi_theta
            yhat .= y
            subres[nx+1:nx+ny] .= cx
            residy = cviolation
            @info log_row(Any[iter, fx, cviolation, optimality, μ, "$(iter_type)"])
        elseif use_filter && phi_O(cviolation, optimality) ≤ phi_O_max
            iter_type = :O
            phi_O_max *= phi_theta
            yhat .= y
            subres[nx+1:nx+ny] .= cx
            residy = cviolation
            @info log_row(Any[iter, fx, cviolation, optimality, μ, "$(iter_type)"])
        else
            iter_type = :F
            @info log_row(Any[iter, fx, cviolation, optimality, μ, residy])
        end

    end

    status = if is_solved
        :first_order
    elseif is_infeasible
        :infeasible
    elseif is_tired
        if iter ≥ max_iter
            :max_iter
        elseif eltime ≥ max_time
            :max_time
        elseif neval_obj(nlp) > max_eval
            :max_eval
        end
    end

    # output
    eltime = time() - start_time
    return GenericExecutionStats(
        status,
        nlp,
        solution = x,
        objective = fx,
        multipliers = y,
        dual_feas = optimality,
        primal_feas = cviolation,
        iter = iter,
        elapsed_time = eltime,
    )
end

end # module


using PyCall
@pyimport matplotlib.pyplot as plt
include("Gaussian.jl")

# Make a gaussian
gaussian = Gaussian([1.0, 0.5, 0.0, 0.0, 1.0, pi/6])

# Define a grid
xmin, xmax, ymin, ymax = -5.0, 5.0, -5.0, 5.0
ni, nj = 201, 201
x = Array(linspace(xmin, xmax, nj))
y = Array(linspace(ymax, ymin, ni))
dx, dy = (xmax - xmin)/(nj - 1), (ymax - ymin)/(ni - 1)
dA = dx*dy

# Evaluate the gaussian on the grid
f = Array(Float64, (ni, nj))
for(j in 1:nj)
	for(i in 1:ni)
		f[i, j] = evaluate(gaussian, x[j], y[i])
	end
end

println(sum(f)*dA)

plt.imshow(f, interpolation="nearest")
plt.show()


type ForwardRateAgreement{DC <: DayCount, BC <: BusinessCalendar, C <: BusinessDayConvention, Y <: YieldTermStructure, P <: PositionType} <: AbstractForward
  lazyMixin::LazyMixin
  underlyingIncome::Float64
  underlyingSpotValue::Float64
  dc::DC
  calendar::BC
  convention::C
  settlementDays::Int
  payoff::ForwardTypePayoff
  valueDate::Date
  maturityDate::Date
  discountCurve::Y
  fraType::P
  forwardRate::InterestRate
  strikeForwardRate::InterestRate
  notionalAmount::Float64
  iborIndex::IborIndex
end

function ForwardRateAgreement(valueDate::Date, maturityDate::Date, position::PositionType, strikeForward::Float64, notionalAmount::Float64,
                              iborIndex::IborIndex, discountCurve::YieldTermStructure)
  calendar = iborIndex.fixingCalendar
  convention = iborIndex.convention
  settlementDays = iborIndex.fixingDays
  maturityDate = adjust(calendar, convention, maturityDate)
  fixingDate = advance(Dates.Day(-settlementDays), calendar, valueDate)
  forwardRate = InterestRate(fixing(iborIndex, iborIndex.ts, fixingDate), iborIndex.dc, SimpleCompounding(), Once())
  strikeForwardRate = InterestRate(strikeForward, iborIndex.dc, SimpleCompounding(), Once())

  strike = notionalAmount * compound_factor(strikeForwardRate, valueDate, maturityDate)
  payoff = ForwardTypePayoff(position, strike)

  return ForwardRateAgreement(LazyMixin(), 0.0, 0.0, iborIndex.dc, calendar, convention, settlementDays, payoff, valueDate, maturityDate, discountCurve,
                              position, forwardRate, strikeForwardRate, notionalAmount, iborIndex)
end

function spot_value(fra::ForwardRateAgreement)
  calculate!(fra)
  return fra.notionalAmount * compound_factor(forward_rate(fra), fra.valueDate, fra.maturityDate) * discount(fra.discountCurve, fra.maturityDate)
end

function forward_rate(fra::ForwardRateAgreement)
  calculate!(fra)
  return fra.forwardRate
end

function forward_value(fra::ForwardRateAgreement)
  calculate!(fra)
  return (fra.underlyingSpotValue - fra.underlyingIncome) / discount(fra.discountCurve, fra.maturityDate)
end

spot_income(::ForwardRateAgreement, ::YieldTermStructure) = 0.0

function implied_yield(fra::ForwardRateAgreement, underlyingSpotValue::Float64, forwardValue::Float64, settlementDate::Date, comp::CompoundingType, dc::DayCount)
  t = year_fraction(dc, settlementDate, fra.maturityDate)
  compoundingFactor = forwardValue / (underlyingSpotValue - spot_income(fra, fra.discountCurve))

  return implied_rate(compoundingFactor, dc, comp, t, Annual())
end

function perform_calculations!(fra::ForwardRateAgreement)
  fixingDate = advance(Dates.Day(-fra.settlementDays), fra.calendar, fra.valueDate)
  fra.forwardRate = InterestRate(fixing(fra.iborIndex, fra.iborIndex.ts, fixingDate), fra.iborIndex.dc, SimpleCompounding(), Once())
  fra.underlyingSpotValue = spot_value(fra)
  fra.underlyingIncome = 0.0
  # NPV calc

  return fra
end


mutable struct FastNode{G<:Group} <: OnlineStat{XY}
    stats::Vector{G}
    ids::Vector{Int}  # self, left, right
    j::Int
    at::Float64
    ig::Float64
end
function FastNode(p, nkeys=2; stat = FitNormal(), id=1)
    FastNode([p * stat for k in 1:nkeys], [id], 0, 0.0, -Inf)
end
Base.show(io::IO, o::FastNode) =
    print(io, "FastNode | $(nkeys(o)) keys × $(nvars(o)) vars | at=$((o.j, o.at))")

function _fit!(o::FastNode, xy)
    x, y = xy
    y in 1:nkeys(o) || error("y must be an integer in 1:nkeys")
    _fit!(o.stats[y], x)
end
nobs(o::FastNode) = sum(nobs, o.stats)
probs(o::FastNode) = nobs.(o.stats) ./ nobs(o)
nkeys(o::FastNode) = length(o.stats)
nvars(o::FastNode) = length(o.stats[1])
Base.getindex(o::FastNode, i) = [stat[i] for stat in o.stats]

_merge!(o::FastNode, o2::FastNode) = _merge!.(o.stats, o2.stats)

function fakedata(::Type{FastNode}, n, p)
    x = randn(n, p)
    y = [(rand() > 1 /(1 + exp(xb))) + 1 for xb in x * (1:p)]
    x, y
end
function classify(o::FastNode)
    out = 1
    n = nobs(o.stats[1])
    for j in 2:nkeys(o)
        n2 = nobs(o.stats[j])
        if n2 > n
            out = j
            n = n2
        end
    end
    out
end

function whichchild(o::FastNode, x::VectorOb)
    o.j == 0 && return rand(o.ids[2:3])
    x[o.j] < o.at ? o.ids[2] : o.ids[3]
end

impurity(p) = entropy(p, 2)

# tree needs `length` method
function split!(o::FastNode, tree)
    n = nobs(o)
    nl = zeros(nkeys(o))  # "prob" left
    nr = zeros(nkeys(o))  # "prob" right
    imp_before = impurity(probs(o))
    ig = -Inf
    ind = 0
    at = -Inf
    sc = zeros(9 * nkeys(o))
    for j in 1:nvars(o)
        stats_j = o[j]
        k = 0
        for stat in stats_j # split candidates
            μ = mean(stat)
            σ = std(stat)
            sc[k+1] = μ - 2σ
            sc[k+2] = μ - 1.5σ
            sc[k+3] = μ - σ
            sc[k+4] = μ - .5σ
            sc[k+5] = μ
            sc[k+6] = μ + .5σ
            sc[k+7] = μ + σ
            sc[k+8] = μ + 1.5σ
            sc[k+9] = μ + 2σ
            k += 9
        end
        for loc in sc
            for k in 1:nkeys(o)
                stat_kj = stats_j[k]  # summarizes variable j for label k
                # @show cdf(stat_kj, loc), std(stat_kj), nobs(stat_kj)
                nl[k] = cdf(stat_kj, loc) * nobs(stat_kj)
                nr[k] = nobs(stat_kj) - nl[k]
            end
            imp_l = impurity(nl ./ sum(nl))
            imp_r = impurity(nr ./ sum(nr))
            imp_after = smooth(imp_l, imp_r, sum(nr) / n)
            new_ig = imp_before - imp_after
            if new_ig > ig
                ig = new_ig
                ind = j
                at = loc
            end
        end
    end
    o.j = ind
    o.at = at
    o.ig = ig
    d = length(tree)
    push!(o.ids, d + 1)
    push!(o.ids, d + 2)
    FastNode(nvars(o), nkeys(o); id = d + 1), FastNode(nvars(o), nkeys(o); id = d + 2)
end

#-----------------------------------------------------------------------# FastTree
"""
    FastTree(p::Int, nclasses=2; stat=FitNormal(), maxsize=5000, splitsize=1000)

Calculate a decision tree of `p` predictors variables and classes `1, 2, …, nclasses`.
Nodes split when they reach `splitsize` observations until `maxsize` nodes are in the tree.
Each variable is summarized by `stat`, which can be `FitNormal()` or `Hist(nbins)`.

# Example

    x = randn(10^5, 10)
    y = rand([1,2], 10^5)

    o = fit!(FastTree(10), zip(eachrow(x),y))

    xi = randn(10)
    classify(o, xi)
"""
struct FastTree{T<:FastNode} <: OnlineStat{XY}
    tree::Vector{T}
    maxsize::Int
    splitsize::Int
end
function FastTree(p, nkeys=2; stat = FitNormal(), maxsize=5000, splitsize=1000)
    tree = [FastNode(p, nkeys; stat=stat)]
    FastTree(tree, maxsize, splitsize)
end
function Base.show(io::IO, o::FastTree)
    print(io, "FastTree(n=", nobs(o))
    print(io, ", size=", length(o.tree))
    print(io, ", maxsize=", o.maxsize)
    print(io, ", splitsize=", o.splitsize)
    print(io, ")")
end
nobs(o::FastTree) = nobs(o.tree[1])
nkeys(o::FastTree) = nkeys(o.tree[1])
nvars(o::FastTree) = nvars(o.tree[1])
Base.length(o::FastTree) = length(o.tree)
function _fit!(o::FastTree, xy)
    x, y = xy
    node = whichleaf(o, x)
    _fit!(node, xy)
    if length(o.tree) < o.maxsize && nobs(node) > o.splitsize
        left, right = split!(node, o.tree)
        push!(o.tree, left)
        push!(o.tree, right)
    end
end

function whichleaf(o, x::VectorOb)
    node = o.tree[1]
    while length(node.ids) > 1
        node = o.tree[whichchild(node, x)]
    end
    node
end

classify(o::FastTree, x::VectorOb) = classify(whichleaf(o, x))
function classify(o::FastTree, x::AbstractMatrix)
    [classify(o, xi) for xi in OnlineStatsBase.eachrow(x)]
end

#-----------------------------------------------------------------------# FastForest
"""
    FastForest(p, nkeys=2; stat=FitNormal(), kw...)

Calculate a random forest where each variable is summarized by `stat`.

# Keyword Arguments

- `nt=100)`: Number of trees in the forest
- `b=floor(Int, sqrt(p))`: Number of random features for each tree to receive
- `maxsize=1000`: Maximum size for any tree in the forest
- `splitsize=5000`: Number of observations in any given node before splitting
- `λ = .05`: Probability that each tree is updated on a new observation

# Example

    x, y = randn(10^5, 10), rand(1:2, 10^5)

    o = fit!(FastForest(10), zip(eachrow(x),y))

    classify(o, x[1,:])
"""
mutable struct FastForest{T<:FastTree} <: OnlineStat{XY}
    forest::Vector{T}
    subsets::Matrix{Int}
    p::Int
    λ::Float64
    n::Int
end
function FastForest(p, nkeys=2; stat = FitNormal(), maxsize=1000, splitsize = 5000,
                    nt=100, b=floor(Int, sqrt(p)), λ = 5 * (1 / nt))
    forest = [FastTree(b, nkeys; stat=stat, maxsize=maxsize, splitsize=splitsize) for i in 1:nt]
    subsets = fill(0, b, nt)
    for j in 1:size(subsets, 2)
        subsets[:, j] = sample(1:p, b; replace=false)
    end
    FastForest(forest, subsets, p, λ, 0)
end
nkeys(o::FastForest) = maximum(nkeys(tree) for tree in o.forest)
function Base.show(io::IO, o::FastForest)
    print(io, "FastForest(")
    print(io, "n=", nobs(o))
    print(io, ", nt=", length(o.forest))
    print(io, ", b=", size(o.subsets, 1))
    print(io, ", avg_size=", mean(length, o.forest))
    print(io, ")")
end
function _fit!(o::FastForest, xy)
    x, y = xy
    o.n += 1
    for (tree, subset) in zip(o.forest, OnlineStatsBase.eachcol(o.subsets))
        rand() < o.λ && fit!(tree, (x[subset], y))
    end
end

function _classify(o::FastForest, x::VectorOb, buffer::Vector{Int})
    buffer[:] .= 0
    for (i, tree) in enumerate(o.forest)
        buffer[classify(tree, x[o.subsets[:, i]])] += 1
    end
    findmax(buffer)[2]
end

classify(o::FastForest, x::VectorOb) = _classify(o, x, zeros(Int, nkeys(o)))
function classify(o::FastForest, x::AbstractMatrix)
    buffer = zeros(Int, nkeys(o))
    [_classify(o, xi, buffer) for xi in OnlineStatsBase.eachrow(x)]
end


#greedy algorithm

function greedy{T<:Real}(distmat::Matrix{T})
  dmsz = size(distmat)
  numCities = dmsz[1]
  ordDist = sortperm(vec(distmat))
  G = Graph(numCities)
  k = 1
  #main loop
  while ne(G) < numCities-1
    dnode,anode = ind2sub(dmsz,ordDist[k])
    if dnode != anode
      d°,a° = degree(G,dnode),degree(G,anode)
      if d° < 2 && a° < 2
        if d° == 0 && a° == 0
          add_edge!(G,dnode,anode)
        elseif independent(G,dnode,anode)
          add_edge!(G,dnode,anode)
        end
      end
    end
    k+=1
  end
  #final edge addition
  add_edge!(G,findin(degree(G),1)...)
  return saw(G,1,numCities)
end


@testset "ExtremelySimpleNL" begin
# input
s1 = lorenzx + 0.1noise # input
s2 = roesslerz + 0.1noise # input
k = 5 # backward dim
ℓ = 5 # forward dim
ε = 0.3 # range
w = 1 # theiler window

for (name, s) in zip(("lorenz", "roessler"), (s1, s2))

    τ = estimate_delay(s, "mi_min") #  5 # delaytime
    method = ExtremelySimpleNL(k, ℓ, τ, w, ε)

    x, r = decompose(s, method)

    err = nrmse(x, s)
    @test err < 0.2

    # figure()
    # title(name)
    # plot(s; alpha = 0.75, label = "original")
    # plot(x; alpha = 0.75, ls = "dashed", label = "noiseless, err=$err")
    # legend()
end
end

@testset "ManifoldProjection" begin
    @testset "Henon" begin
        he = Systems.henon()
        tr = trajectory(he, 10000; Ttr = 100)
        Random.seed!(151521)
        z = tr[:, 1]
        s = z .+ randn(10001)*0.1*std(z)
        # s = lorenzx + 0.1noise
        m = 5
        w = 1 # theiler window
        metric = Euclidean()
        searchtype = 30
        Q = 2
        Q = [2, 2, 2, 3, 3, 3, 3]

        x, r = decompose(s, ManifoldProjection(m, Q, searchtype))

        err = nrmse(x, tr[:, 1])
        @test err < 0.1

        # figure()
        # plot(s, label = "input")
        # plot(z, color = "k", ls = "dashed", label = "real")
        # plot(x, alpha = 0.5, label = "output")

        # figure()
        # scatter(s[1:end-1], s[2:end], s = 0.1, label = "input")
        # scatter(z[1:end-1], z[2:end], s = 0.1, label = "real", color = "k")
        # scatter(x[1:end-1], x[2:end], s = 0.1, label = "output", alpha = 0.5)
        # legend()
    end

    @testset "lorenz" begin
        lo = Systems.lorenz()
        tr = trajectory(lo, 1000; Ttr = 100, dt = 0.1)
        Random.seed!(151521)
        z = tr[:, 1]
        s = z .+ randn(10001)*0.2*std(z)
        m = 13
        w = 1 # theiler window
        metric = Euclidean()
        k = 30
        Q = 2
        Q = [2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4]

        x, r = decompose(s, ManifoldProjection(m, Q, 30))
        err = nrmse(x, z)
        @test err < 0.2

        # figure()
        # plot(s, label = "input")
        # plot(z, color = "k", ls = "dashed", label = "real")
        # plot(x, alpha = 0.5, label = "output")
        # estimate_delay(z, "mi_min")
        # xlim(0, 200)
        #
        # τ = 1
        # d = 2
        # figure()
        # plot(columns(embed(s, d, τ))..., marker = "o", ms = 1, lw = 0.1, label = "input")
        # plot(columns(embed(z, d, τ))..., marker = "o", ms = 1, lw = 0.1, label = "real", color = "k")
        # plot(columns(embed(x, d, τ))..., marker = "o", ms = 1, lw = 0.1, label = "real", alpha = 0.5)
    end
end


using LinearAlgebra
using ArnoldiMethod: lock!

function example()
    n = 10
    Q = qr(rand(n, n)).Q * Matrix(I, n, n)
    R = diagm(0 => 1:10) .+ triu(randn(n, n) .* 0.01, 1)
    R[2,3] = 1.0
    R[3,2] = -1.0

    # Some Schur decomp
    A = Q * R * Q'

    lock!(R, Q, 1, [8,9,10])
    truncate!()

    @show norm(A * Q - Q * R)

    return R

end


type SDF_MAXWELL_VD <: SDFDynamicSystem
    omg::Float64
    zta::Float64
    xi::Float64
    kpa::Float64
    alpha::Float64
    u0::Float64
    v0::Float64
    ud0::Float64

    SDF_MAXWELL_VD(omg::Float64, zta::Float64, xi::Float64, kpa::Float64, alpha::Float64) = new(omg, zta, xi, kpa, alpha, 0.0, 0.0, 0.0)
    SDF_MAXWELL_VD(omg::Float64, zta::Float64, xi::Float64, kpa::Float64) = new(omg, zta, xi, kpa, 1.0, 0.0, 0.0, 0.0)
end

function systemEquation(s::SDF_MAXWELL_VD, p::Function)
    function dydt(y::AbstractArray{Float64,1}, t::Float64)
        u, v, ud = y
        omg2 = s.omg*s.omg
        dvdt = p(t)-2.0*s.zta*s.omg*v-omg2*u-s.kpa*omg2*(u-ud)
        if s.alpha == 1.0
            duddt = s.kpa*s.omg/2.0/s.xi*(u-ud)
        else
            duddt = sign(u-ud)*abs(s.kpa*s.omg/2.0/s.xi*(u-ud))^(1.0/s.alpha)
        end
        [v; dvdt; duddt]
    end
    dydt
end

function response(s::SDF_MAXWELL_VD, p::Function, t::AbstractArray{Float64,1})
    y0 = [s.u0; s.v0; s.ud0]

    f = systemEquation(s, p)

    y = ode4(f, y0, t)

    u = y[:,1]
    v = y[:,2]
    ud = y[:,3]
    fdz = 2.0*s.zta*s.omg*v
    fdx = s.kpa*s.omg*s.omg*(u-ud)
    a = map(p,t)-fdz-fdx-s.omg*s.omg*u

    [u v a fdz fdx]
end


using Plots, LaTeXStrings
gr(colorbar=false,size=(300,300),legend=false,grid=false)

using FourierGPE

## setup
L = (60.0,60.)
N = (512,512)
sim = Sim(L,N)
@unpack_Sim sim

# individual data files for each time
# nfiles = true
# filename = "test"
# savedir = "data"
# path = joinpath(@__DIR__,savedir)

μ = 1.0
g = 0.01
γ = 0.0
tf = 2pi
Nt = 150
ti = 0.0
t = LinRange(ti,tf,Nt)

## Bogoliubov state
x,y = X
kx,ky = K

f(k) = 1 + 4/k^2
u(k) = 0.5*(sqrt(f(k))+1)/f(k)^(1/4)
v(k) = 0.5*(sqrt(f(k))-1)/f(k)^(1/4)
lam = 0.01
bog(x,k) = u(k)*exp(im*k*x) - conj(v(k))*exp(-im*k*x)
ψb(x,y,k) = sqrt(μ/g)*(complex(one(x)) + lam*bog(x,k))*one(y)

kb = kx[20]
ψi = ψb.(x,y',kb)
ϕi = kspace(ψi,sim)

# Set time evolution and pack
# reltol = 1e-7
# alg = Vern7()

@pack_Sim! sim

## Evolve in k space
@time sol = runsim(sim)

z = abs2.(xspace(sol[end-1],sim)); heatmap(x,y,z)

## movie
anim = @animate for i in eachindex(t)
    ψ = xspace(sol[i],sim)
    n = abs2.(ψ)
    heatmap(x,y,n,aspectratio=1,transpose=true)
    xlims!(-L[1]/2,L[1]/2)
    ylims!(-L[2]/2,L[2]/2)
    xlabel!(L"x/\xi"); ylabel!(L"y/\xi")
end

gif(anim, "./examples/2dBogoliubov.gif", fps = 25)

## slice movie
anim = @animate for i in eachindex(t)
    z = abs2.(xspace(sol[i],sim))[:,256]
    plot(x,z)
    xlims!(-L[1]/2,L[1]/2)
    xlabel!(L"x/\xi")
    ylabel!(L"n(x)")
end

gif(anim, "./examples/2dBogoliubovSlice.gif", fps = 25)

## energy densities

function showenergies(ψ)
    et,ei,ec = energydecomp(ψ)
    p1 = heatmap(x,y,log10.(ei),aspectratio=1,transpose=true)
    xlabel!(L"x/\xi");ylabel!(L"y/\xi")
    title!("Incompressible")
    p2 = heatmap(x,y,log10.(ec),aspectratio=1,transpose=true)
    xlabel!(L"x/\xi");ylabel!(L"y/\xi")
    title!("Compressible")
    p = plot(p1,p2,size=(600,300))
    return p
end

K2 = k2(K)

anim = @animate for i in eachindex(t)
    ψ = xspace(sol[i],sim)
    psi = XField(ψ,X,K,K2)
    showenergies(psi)
end

gif(anim,"./examples/2dbogenergies.gif",fps = 25)

## energy totals
function xenergy(ϕ,sim,t)
    @unpack g,X = sim; x,y = X
    ψ = xspace(ϕ,sim)
    @. ψ *= 0.5*g*abs2(ψ) + V(x,y',t)
    return kspace(ψ,sim)
end

function gpenergy(ϕ,sim,t)
    @unpack μ,γ,espec,K = sim; kx,ky = K
    dkx,dky = kx[2]-kx[1],ky[2]-ky[1]
    χ = xenergy(ϕ,sim,t)
    H = @. espec*abs2(ϕ) + conj(ϕ)*χ
    return sum(H)*dkx*dky |> real
end

function qpressure(ϕ,sim)
    @unpack g,X,K,espec = sim; x,y = X; kx,ky = K
    dkx,dky = kx[2]-kx[1],ky[2]-ky[1]
    ψ = xspace(ϕ,sim)
    sqrtn = abs.(ψ)
    sqrtnk = kspace(sqrtn,sim)
    Eqp = @. espec*abs2(sqrtnk)
    return sum(Eqp)*dkx*dky |> real
end

H = zero(t)
Ex = zero(t)
Ei = zero(t)
Ec = zero(t)
Eqp = zero(t)
Et = zero(t)
Natoms = zero(t)
dx = diff(x)[1]; dy = diff(y)[1]
dkx = diff(kx)[1]; dky = diff(ky)[1]
for (i,t) in enumerate(t)
    ϕ = sol[i]
    ψ = xspace(ϕ,sim)
    psi = XField(ψ,X,K,K2)
    et,ei,ec = energydecomp(psi)
    Ex[i] = sum(conj(ϕ).*xenergy(ϕ,sim,t))*dkx*dky |> real
    Ei[i] = sum(ei)*dx*dy |> real
    Ec[i] = sum(ec)*dx*dy |> real
    Eqp[i] = qpressure(ϕ,sim)
    Et[i] = Ei[i] + Ec[i] + Eqp[i] + Ex[i]
    Natoms[i] = sum(abs2.(ψ))*dx*dy
    H[i] = gpenergy(ϕ,sim,t)
end

## compressible energy conservation
plot(t,Et./Natoms,label=L"E_t",size=(400,200),legend=true)
plot!(t,Ei./Natoms,label=L"E_i",legend=:bottomright)
plot!(t,Ec./Natoms,label=L"E_c")
plot!(t,H./Natoms,label=L"H")
plot!(t,Eqp./Natoms,label=L"E_{qp}")
plot!(t,Ex./Natoms,label=L"E_{x}")
xlabel!(L"t")


module SymmetrySnap
    using LinearAlgebra
    using Random
    using DelimitedFiles
    using MinkowskiReduction

    """
        getLatticeDiff(oldA, a, oldB, b, oldC, c, output)

    Save the change in lattice to output vector

    # Arguments
    - `oldA, oldB, oldC::Vector{Float64}`: Old lattice vectors
    - `a, b, c::Vector{Float64}`: New lattice vectors
    - `output::Vector{String}`: Output list
    """
    function getLatticeDiff(oldA, a, oldB, b, oldC, c, numSymmetries, output)
        push!(output, "Old Lattice -> New Lattice : L2 Norm, SymOps = $(numSymmetries)")
        push!(output, join([oldA, a, norm(oldA - a)]," -> "," : "))
        push!(output, join([oldB, b, norm(oldB - b)]," -> "," : "))
        push!(output, join([oldC, c, norm(oldC - c)]," -> "," : "))
        push!(output, "")
        return output
    end

    """
        getAtomDiff(oldAtomPos, atomPos, symmetries, output)

    Save the change in atom positions and number of symmetries found to output vector

    # Arguments
    - `oldAtomPos::Vector{Vector{Float64}}`: Old atom positions
    - `atomPos::Vector{Vector{Float64}}`: New atom positions
    - `symmetries::Int8`: Number of symmetries found
    - `output::Vector{String}`: Output list
    """
    function getAtomDiff(oldAtomPos, atomPos, symmetries, output)
        push!(output, "Old Atom Position -> New Atom Position : L2 Norm, SymOps")
        for i in 1:length(atomPos)
            push!(output, join([join([oldAtomPos[i], atomPos[i], norm(oldAtomPos[i] - atomPos[i])]," -> "," : "),symmetries[i]],", "))
        end
        return output
    end

    """
        mapAtomsToCell(a, b, c, latVecInv, atomPositions)

    Map given atom positions into positions inside the unit cell.

    # Arguments
    - `a,b,c::Vector{Float64}`: lattice vectors
    - `latVecInv::Array{Float64, 2}`: Inverse of lattice vectors
    - `atomPositions::Vector{Vector{Float64}}`

    # Examples
    ```jldoctest
    julia> mapAtomsToCell([1,0,0], [0,1,0], [0,0,1], [1, 0, 0; 0, 1, 0; 0, 0, 1], [1.5,-0.5,3.5])
    [0.5, 0.5, 0.5]
    ```
    """
    function mapAtomsToCell(a, b, c, latVecInv, atomPositions)
        newPositions = []
        for atom in atomPositions
            addTo = true
            shift = latVecInv * atom
            for i in 1:3
                shift[i] = floor(shift[i])
            end
            newPos = atom - a*shift[1] - b*shift[2] - c*shift[3]
            for newAtom in newPositions
                if isapprox(newPos, newAtom)
                    addTo = false
                    break
                end
            end
            if addTo == true
                push!(newPositions, newPos)
            end
        end
        return newPositions
    end

    """
        toDirect(a, b, c, atomPositions)

    Convert atom positions in Cartesian coordinates to direct coordinates.

    # Arguments
    - `a,b,c::Vector{Float64}`: lattice vectors
    - `atomPositions::Vector{Vector{Float64}}`: atom positions in Cartesian coordinates

    # Examples

    """
    function toDirect(a, b, c, atomPositions)
        latVecs = Array{Float64, 2}(undef, 3, 3)
        latVecs[1:end,1] = a
        latVecs[1:end,2] = b
        latVecs[1:end,3] = c

        newPositions = Vector{Vector{Float64}}(undef, length(atomPositions))
        for i in 1:length(atomPositions)
            newPositions[i] = \(latVecs, atomPositions[i])
        end
        return newPositions
    end

    """
        toCartesian(a, b, c, atomPositions)

        Convert atom positions in direct coordinates to Cartesian coordinates.

        # Arguments
        - `a,b,c::Vector{Float64}`: lattice vectors
        - `atomPositions::Vector{Vector{Float64}}`: atom positions in direct coordinates

        # Examples

    """
    function toCartesian(a, b, c, atomPositions)
        latVecs = Array{Float64, 2}(undef, 3, 3)
        latVecs[1:end,1] = a
        latVecs[1:end,2] = b
        latVecs[1:end,3] = c

        newPositions = Vector{Vector{Float64}}(undef, length(atomPositions))
        for i in 1:length(atomPositions)
            newPositions[i] = latVecs * atomPositions[i]
        end
        return newPositions
    end

    """
        test_atom(a, b, c, lat_vec_inv, symmetry_op, atom_1_pos, atom_1_index, atom_2_pos, atom_2_index, natoms[, epsilon=1e-2])

    Test if atom 1 maps to the position of atom 2 when the symmetry operator is applied. If valid, returns matrix A that represents the mapping from atom 1 to atom 2.

    # Arguments
    - `a,b,c::Vector{Float64}`: lattice vectors
    - `latVecInv::Array{Float64, 2}`: inverse lattice vectors matrix
    - `symmetryOp::Array{Float64, 2}`: symmetry operator
    - `atom_1_index, atom_2_index::{Int}`: indexes of atoms being compared
    - `atom_1_pos, atom_2_pos::Vector{Float64}`: positions of atoms being compared
    - `nAtoms::Int8`: Total number of atoms in the system
    - `epsilon::Float64`: Largest normed difference between positions for a valid symmetry operator

    # Examples

    """
    function testAtom(a, b, c, latVecInv, symmetryOp, atom1Pos, atom1Index, atom2Pos, atom2Index, nAtoms, shift, epsilon=1e-2)
        # In the case of a different crystal structure, A would be the lattice vectors as rows
        m = symmetryOp * atom1Pos #what is m??

        if shift != nothing
            m = m .+ shift
        end

        compareCell = latVecInv * m
        for i in 1:3
            compareCell[i] = floor(compareCell[i])
        end

        transformedAtom = symmetryOp * atom1Pos .- a*compareCell[1] .- b*compareCell[2] .- c*compareCell[3]
        if shift == nothing
            shift = atom2Pos .- transformedAtom
        end

        transformedAtom = transformedAtom .+ shift

        if norm(atom2Pos .- symmetryOp * atom1Pos .- shift) < epsilon || norm(atom2Pos .- transformedAtom) < epsilon
            A = zeros(Float64, 3, 3*nAtoms+1)
            if norm(atom2Pos .- symmetryOp * atom1Pos .- shift) < epsilon
                compareCell = [0,0,0]
            end
            for i in 1:3
                for j in 1:3
                    if symmetryOp[i,j] != 0
                        A[i,3*(atom1Index-1)+j] += symmetryOp[i,j]
                    end
                    if -Matrix(I, 3, 3)[i,j] != 0
                        A[i,3*(atom2Index-1)+j] += -Matrix(I, 3, 3)[i,j]
                    end
                end
            end
            A[:,3*nAtoms+1] = a*compareCell[1] .+ b*compareCell[2] .+ c*compareCell[3] .- shift
            A[abs.(A) .< 1e-6] .= 0
            if isapprox(A, zero(A), atol = 1e-10)
                return nothing, shift, true
            end
            return A, shift, true
        else
            return nothing, shift, false
        end
    end

    """
        runSymmetry(a,b,c,latVecInv, symmetryOp, atomPositions[, epsilon=1e-2])

    Test the given symmetry operator on each atom in atomPositions to verify if it is a true symmetry, and return the list of symmetry equivalent linear algebra problems.

    # Arguments
    - `a,b,c::Vector{Float64}`: lattice vectors
    - `latVecInv::Array{Float64, 2}`: inverse lattice vectors matrix
    - `symmetryOp::Array{Float64, 2}`: symmetry operator
    - `atomPositions::Vector{Vector{Float64}}`: atom positions
    - `epsilon::Float64`: Largest normed difference between positions for a valid symmetry operator

    # Examples

    """
    function runSymmetry(a, b, c, latVecInv, symmetryOp, atomPositions, epsilon=1e-2)
        # Which atom to shift by
        for i in 1:length(atomPositions)
            As = Vector{Array{Float64, 2}}(undef, 0)
            global shift = nothing
            indexFound = []
            # Iterate over all atoms for atom 1
            for j in 1:length(atomPositions)
                global symmFound = false
                # Iterate over all atoms for atom 2
                for k in 1:length(atomPositions)
                    if i+k-1 in indexFound
                        continue
                    end
                    A, shift, symmFound = testAtom(a, b, c, latVecInv, symmetryOp, atomPositions[j], j, atomPositions[(i+k-2)%length(atomPositions)+1], (i+k-2)%length(atomPositions)+1, length(atomPositions), shift, epsilon)
                    if A != nothing
                        push!(As, A)
                        push!(indexFound, i+k-1)
                    end
                    if symmFound
                        push!(indexFound, i+k-1)
                        break
                    end
                end
                if !symmFound
                    break
                end
            end
            if symmFound
                return As
            end
        end
        return nothing
    end

    """
        symmetrize(a, b, c, lat_vec_inv, atom_positions[, epsilon=1e-2])

    Using the valid symmetry operators found, solve the \\(Ax=b\\) problem to get the smallest correction in atom positions needed to align with the improved lattice vectors.

    # Arguments
    - `a,b,c::Vector{Float64}`: lattice vectors
    - `latVecInv::Array{Float64, 2}`: inverse lattice vectors matrix
    - `atomPositions::Vector{Vector{Float64}}`: atom positions
    - `epsilon::Float64`: Largest normed difference between positions for a valid symmetry operator

    # Examples

    """
    function symmetrize(a, b, c, latVecInv, atomPositions, symmetries, epsilon=1e-2)
        As = Vector{Array{Float64, 2}}(undef, 0)
        symm_num = 0
        for i in 1:length(symmetries)
            newA = runSymmetry(a, b, c, latVecInv, symmetries[i], atomPositions, epsilon)
            if newA != nothing
                symm_num += 1
                check = false
                for new in newA
                    for A in As
                        if isapprox(A,new, atol = 1e-10)
                            check = true
                            break
                        end
                    end
                    if !check
                        push!(As,new)
                    end
                end
            end
        end
        if length(As) == 0
            return atomPositions, zeros(length(atomPositions))
        end

        A = zeros(Float64, 3*length(As), 3*length(atomPositions))
        B = zeros(Float64, 3*length(As))
        for i in 0:length(As)-1
            A[3*i+1:3*i+3,:] = As[i+1][1:end,1:end .!= length(As[i+1][1,1:end])]
            B[3*i+1:3*i+3] = As[i+1][1:end,end]
        end

        aSVD = svd(A)
        newA = zeros(3*length(As), 3*length(atomPositions))
        for i in 1:length(aSVD.S)
            uArray = aSVD.U[1:end, i]
            vArray = aSVD.Vt[i, 1:end]
            part = aSVD.S[i]*(aSVD.U[1:end, i]*transpose(aSVD.Vt[i, 1:end]))
            part[abs.(part) .< 1e-6] .= 0
            newA += part
        end
        A = newA

        xUser = vcat(atomPositions...)
        xDiff = xUser - \(A,B)
        null = transpose(LinearAlgebra.nullspace(A))

        for i in 1:length(null[1:end, 1])
            xDiff -= null[i,1:end]*dot(xDiff,null[i,1:end])/norm(null[i,1:end])
        end
        xCorrected = xUser - xDiff
        newXs = Array{Array{Float64, 1}, 1}(undef, length(atomPositions))
        for i in 1:length(atomPositions)
            newX = Array{Float64, 1}(undef, 3)
            for j in 1:3
                newX[j] = xCorrected[3*(i-1) + j]
            end
            newXs[i] = newX
        end

        return a, b, c, newXs, repeat([symm_num], length(atomPositions))
    end

    """
        getSymmetryOps(a,b,c,latVecInv)

    Finds the applicable symmetry operators for the system by solving the following equation,

    \\(R = A B^{-1}\\),

    where \\(A\\) is the original lattice vectors, \\(B\\) is the new lattice vectors, and \\(R\\) is the symmetry operator being tested.

    If \\(R R^T = I\\), \\(R\\) is a valid symmetry operator.

    # Arguments
    - `a,b,c::Vector{Float64}`: lattice vectors
    - `latVecInv::Array{Float64, 2}`: inverse lattice vectors matrix

    # Examples

    """
    function getSymmetryOps(a, b, c, latVecInv, atomPositions, epsilon=1e-2)
        translations = [[0,0,1],[0,1,0],[1,0,0],[0,0,-1],[0,-1,0],[-1,0,0],[1,1,0],[1,-1,0],[-1,1,0],
                        [-1,-1,0],[1,0,1],[1,0,-1],[-1,0,1],[-1,0,-1],[0,1,1],[0,1,-1],[0,-1,1],[0,-1,-1],
                        [1,1,1],[1,1,-1],[1,-1,1],[-1,1,1],[1,-1,-1],[-1,1,-1],[-1,-1,1],[-1,-1,-1]]
        latticePoints = []
        for v in translations
            push!(latticePoints,v[1]*a+v[2]*b+v[3]*c)
        end

        oldLattice = Array{Float64, 2}(undef, 3, 3)
        oldLattice[1:end, 1] = a
        oldLattice[1:end, 2] = b
        oldLattice[1:end, 3] = c

        newLattices = Vector{Array{Float64, 2}}(undef, 15600)
        transformations = Vector{Array{Float64, 2}}(undef, 15600)
        index = 1
        for i in 1:length(latticePoints)
            for j in 1:length(latticePoints)
                for k in 1:length(latticePoints)
                    if i == k || j == k || i == j
                        continue
                    end
                    newLattice = Array{Float64, 2}(undef, 3, 3)
                    newLattice[1:end, 1] = latticePoints[i]
                    newLattice[1:end, 2] = latticePoints[j]
                    newLattice[1:end, 3] = latticePoints[k]
                    newLattices[index] = newLattice

                    transformLattice = Array{Float64, 2}(undef, 3, 3)
                    transformLattice[1, 1:end] = translations[i]
                    transformLattice[2, 1:end] = translations[j]
                    transformLattice[3, 1:end] = translations[k]
                    transformations[index] = transformLattice
                    index += 1
                end
            end
        end

        allOperators = []
        allErrors = []
        for i in 1:length(newLattices)
            operator = newLattices[i]*latVecInv
            operatorT = transpose(operator)
            push!(allOperators, operator)
            push!(allErrors, norm(operator*operatorT - Matrix(I,3,3)))
        end

        p = sortperm(allErrors)
        sortOperators = allOperators[p]
        sortErrors = allErrors[p]
        sortTransforms = transformations[p]

        # Symmetrix Matrix, so only 6 elements: xx, yy, zz, xy, xz, yz
        elementMap = Dict(
            (1, 1) => 1,
            (2, 2) => 2,
            (3, 3) => 3,
            (1, 2) => 4,
            (2, 1) => 4,
            (1, 3) => 5,
            (3, 1) => 5,
            (2, 3) => 6,
            (3, 2) => 6
            )
        numOps = [2, 3, 4, 6, 8, 12, 16, 24, 48]
        lattices = []
        errors = []
        opCounts = []
        for numOp in numOps
            A = zeros(6*numOp, 6)
            for opNum in 0:numOp-1
                for i in 1:3
                    for j in i:3
                        A[6*opNum + elementMap[(i, j)], elementMap[(i, j)]] -= 1
                        for k in 1:3
                            for l in 1:3
                                A[6*opNum + elementMap[(i, j)], elementMap[(k, l)]] += sortTransforms[opNum+1][i, k]*sortTransforms[opNum+1][j, l]
                            end
                        end
                    end
                end
            end

            xUser = transpose(oldLattice)*oldLattice

            xUser = [xUser[1, 1], xUser[2, 2], xUser[3, 3], xUser[1, 2], xUser[1, 3], xUser[2, 3]]
            diff = xUser

            null = transpose(nullspace(A))

            for i in 1:length(null[1:end, 1])
                diff -= null[i,1:end]*dot(diff,null[i,1:end])/norm(null[i,1:end])
            end

            xCorrected = xUser - diff

            xCorrFull = [xCorrected[1] xCorrected[4] xCorrected[5];xCorrected[4] xCorrected[2] xCorrected[6];xCorrected[5] xCorrected[6] xCorrected[3]]

            newSVD = svd(xCorrFull)

            newSMat = Diagonal(sqrt.(newSVD.S))
            newVtMat = newSVD.Vt

            findUSVD = nothing
            try
                findUSVD = svd(oldLattice*inv(newSMat*newVtMat))
            catch
                continue
            end
            newUMat = findUSVD.U * findUSVD.Vt

            push!(lattices, newUMat*newSMat*newVtMat)
            push!(errors, norm(oldLattice - lattices[end]))
            push!(opCounts, numOp)
        end

        lattice = Array{Float64, 2}(undef, 3, 3)
        numValidOps = 0
        for i in 1:length(errors)+1
            if i == length(errors)+1 || errors[i] > epsilon
                if i == 1
                    i == 2
                end
                lattice = lattices[i-1]
                numValidOps = opCounts[i-1]
                break
            end
        end

        transformation = lattice*latVecInv
        for i in 1:length(atomPositions)
            atomPositions[i] = transformation*atomPositions[i]
        end

        latVecInv = inv(lattice)
        a = lattice[1:end, 1]
        b = lattice[1:end, 2]
        c = lattice[1:end, 3]

        validOperators = []
        for i in 1:numValidOps
            newLattice = transpose(sortTransforms[i]*transpose(lattice))
            push!(validOperators, newLattice*latVecInv)
        end

        return a, b, c, latVecInv, atomPositions, validOperators
    end

    """
        symmSnap(a, b, c, atomPositions, atomTypes[, isDirect = false, fout = "symm_snap.out", epsilon=1e-2])

    Get corrected/reduced lattice vectors, and corrected atom positions to correspond with a Bravais lattice.

    # Arguments
    - `a,b,c::Vector{Float64}`: lattice vectors
    - `atomPositions::Vector{Vector{Float64}}`: atom positions
    - `atomTypes::_______`:
    - `isDirect::Bool`: indicate if atomPositions are already in direct coordinates
    - `fout::String`: file name to write out additional information to
    - `epsilon::Float64`: Largest normed difference between positions for a valid symmetry operator

    # Examples

    """
    function symmSnap(a, b, c, atomPositions, atomTypes, isDirect = false, fout = "symm_snap.out", epsilon=1e-2)
        if isDirect
            atomPositions = toCartesian(a, b, c, atomPositions)
        end

        lattice = Array{Float64, 2}(undef, 3, 3)
        lattice[1:end, 1] = a
        lattice[1:end, 2] = b
        lattice[1:end, 3] = c
        latVecInv = inv(lattice)
        atomPositions = mapAtomsToCell(a, b, c, latVecInv, atomPositions) #new atom positions, but mapped to the unit cell

        lattice = Array{Float64, 2}(undef, 3, 3)
        lattice[1:end, 1] = a
        lattice[1:end, 2] = b
        lattice[1:end, 3] = c
        latVecInv = inv(lattice)

        oldA = deepcopy(a)
        oldB = deepcopy(b)
        oldC = deepcopy(c)

        a, b, c, latVecInv, atomPositions, symmetries = getSymmetryOps(a, b, c, latVecInv, atomPositions)

        output = getLatticeDiff(oldA, a, oldB, b, oldC, c, length(symmetries), [])

        oldAtomPos = deepcopy(atomPositions)

        index = 1
        num_symm = Vector{Int8}(undef, length(atomPositions))

        for num in atomTypes
            a, b, c, atomPositions[index:index+num-1], num_symm[index:index+num-1] = symmetrize(a, b, c, latVecInv, atomPositions[index:index+num-1], symmetries, epsilon)
            index += num
        end

        if isDirect
            atomPositions = toDirect(a, b, c, atomPositions)
        end

        output = getAtomDiff(oldAtomPos, atomPositions, num_symm, output)

        writedlm(fout, output)

        if dot(a, cross(b, c)) < 0
            tempa = a
            a = b
            b = tempa
        end

        return a, b, c, atomPositions
    end

    function snapPoscar(in_file = "POSCAR", out_file = "POSCAR.fixed", log_file = "symm_snap.out", epsilon=1e-2)
        fin=readlines(in_file)
        a=[parse(Float64,i) for i in split(fin[3])]
        b=[parse(Float64,i) for i in split(fin[4])]
        c=[parse(Float64,i) for i in split(fin[5])]
        atom_types=[parse(Int8,i) for i in split(fin[6])]
        is_direct = lowercase(fin[7][1]) == "d" ? true : false
        atom_positions=[[parse(Float64,i) for i in split(fin[j])[1:3]] for j in 8:length(fin)]

        a, b, c, atom_positions = SymmetrySnap.symmSnap(a, b, c, atom_positions, atom_types, is_direct, log_file, epsilon)

        fin[3] = join(a, " ")
        fin[4] = join(b, " ")
        fin[5] = join(c, " ")

        writedlm(out_file, fin)
    end
end


using Gurobi
using HSL
⁺ = 1; ⁻ = 2; ± = 1:2
negate(i::Integer) = i==⁺ ? ⁻ : ⁺
split_matrix(C) = (max.(C, 0.0), max(-C, 0.0))

# Describes the system Δ₁y = AΔ₁x = (B⁺ - B⁻)Δ₁u + (C⁺ - C⁻)Δ₂{f}
function brouwer_certificate_LP(sys, B, C, Qlim, thermalrate, T, K, dV, Iter, limitQQ, limitIQ;
  region = :default, objective = :default)
  ny = size(B,1)
  nu = size(B,2)
  nf = size(C,2)

  B⁺,B⁻ = split_matrix(B)
  C⁺,C⁻ = split_matrix(C)

  m = Model(solver = GurobiSolver(Method=2,Crossover=0, NodefileStart=0.1))#,outlev=1))

  @variables m begin
    Δ₁y[1:ny,±] >= 0.0, (start = 0.0)
    Δ₂f[1:nf,±] >= 0.0, (start = 0.0)
    Δ₁u[1:nu,±] >= 0.0, (start = 0.0)
  end
  for k in 1:ny, s in ±
    s̄ = negate(s)
    @constraints m begin
      Δ₁y[k,s] >=
        sum(B⁺[k,j]*Δ₁u[j,s] + B⁻[k,j]*Δ₁u[j,s̄] for j=1:nu) +
        sum(C⁺[k,j]*Δ₂f[j,s] + C⁻[k,j]*Δ₂f[j,s̄] for j=1:nf)
    end
  end
  add_Δ₁y_constraints_LP(sys, m, thermalrate, T, K, dV, limitIQ)
  add_Δ₂f_constraints_LP(sys, m, Qlim, T, K, limitQQ)

  if region == :cone
    @constraint(m, [k in 1:nu], Δ₁u[k,⁻] == 0.0)
  elseif region == :uniform
    @constraint(m, [k in 1:nu, s in ±], Δ₁u[k,s] == Δ₁u[1,⁺])
  elseif region == :custom
    add_Δ₁u_constraints(m, sys)
  end # otherwise, don't impose any additional constraints

  if objective == :custom
    add_objective(m, sys)
  elseif objective == :gauss
    s2 = √2
    @NLexpression(m, lpr[k in 1:nu], log(erf(Δ₁u[k,⁺]/s2) + erf(Δ₁u[k,⁻]/s2)))
    @variable(m, logprob)
    @NLconstraint(m, logprob == sum(lpr[k] for k in 1:nu) - 0*nu*log(2))
    @objective(m, Max, logprob)
  elseif objective == :default
    @objective(m, Max, Δ₁u[1,⁺])
  end

  return m
end

export brouwer_certificate_LP


"""
    RMSProp(::Vector{XVariable}; lr=1e-2, inertia=0.99, eps=1e-8, L1decay=0.001, L2decay=0.01)

Implements RMSProp algorithm.
"""
mutable struct RMSProp <: Optimizer
    xparams::Vector{XVariable}
    w::Vector
    lr::AbstractFloat
    ϵ::AbstractFloat
    inertia::AbstractFloat
    L1decay::AbstractFloat
    L2decay::AbstractFloat
    name::String
    function RMSProp(xparams::Vector{XVariable}; lr=1e-2, inertia=0.99, eps=1e-8, L1decay=0.001, L2decay=0.01)
        n = length(xparams)
        w = Vector(undef,n)
        for i = 1:n
            c , θ = xparams[i]
            w[i] = Zeros(typeof(ᵛ(θ)), θ.shape)
        end
        new(xparams, w, lr, eps, inertia, L1decay, L2decay, "RMSProp")
    end
end

function Base.show(io::IO, O::RMSProp)
    print("RMSProp(lr=$(O.lr), ϵ=$(O.ϵ), inertia=$(O.inertia), L1decay=$(O.L1decay), L2decay=$(O.L2decay))")
end

function update!(O::RMSProp; clipfn::Function=LPInfNormClip, clipvalue=10.0)
    w = O.w
    ϵ = O.ϵ
    μ = - O.lr
    ρ = O.inertia
    λ₁ = O.L1decay
    λ₂ = O.L2decay

    for i = 1:length(O.xparams)
        c , θ = O.xparams[i]
        ∇ = clipfn(setNanInfZero(δ(θ)), clipvalue)
        𝒗 = ᵛ(θ)
        @. w[i] = ρ * w[i] + (1-ρ) * ∇ * ∇
        if c == 'w'
            if λ₁==0 && λ₂==0
                @. 𝒗 += μ / (sqrt(w[i])+ϵ) * ∇
            elseif λ₁==0 && λ₂!=0
                @. 𝒗 += μ / (sqrt(w[i])+ϵ) * (∇ + λ₂ * 𝒗)
            elseif λ₁!=0 && λ₂==0
                @. 𝒗 += μ / (sqrt(w[i])+ϵ) * (∇ + λ₁ * sign(𝒗))
            else  # λ₁!=0 && λ₂!=0
                @. 𝒗 += μ / (sqrt(w[i])+ϵ) * (∇ + λ₁ * sign(𝒗) + λ₂ * 𝒗)
            end
        else
            @. 𝒗 += μ / (sqrt(w[i])+ϵ) * ∇
        end
    end
end


"""
Vanilla RNN, i.e. ⤦\n
    hᵗ = f(w*xᵗ + u*hᵗ⁻¹ .+ b)
"""
mutable struct RNN <: Block
    w::VarOrNil # input to hidden weights
    b::VarOrNil # bias of hidden units
    u::VarOrNil # recurrent weights
    f::Function # activation function
    h::Any      # hidden variable
    function RNN(inputSize::Int, hiddenSize::Int, fn::Function=relu; type::Type=Array{Float32})
        T = eltype(type)
        w = randn(T, hiddenSize, inputSize) .* sqrt( T(2/inputSize) )
        b = zeros(T, hiddenSize, 1)
        u = randdiagonal(T, hiddenSize; from=-0.1990, to=0.1997)
        new(Variable{type}(w,true,true,true),
            Variable{type}(b,true,true,true),
            Variable{type}(u,true,true,true), fn, nothing)
    end
    function RNN(fn::Function)
        new(nothing, nothing, nothing, fn, nothing)
    end
end

function clone(this::RNN; type::Type=Array{Float32})
    cloned = RNN(this.f)
    cloned.w = clone(this.w, type=type)
    cloned.b = clone(this.b, type=type)
    cloned.u = clone(this.u, type=type)
    return cloned
end

mutable struct RNNs <: Block
    layers::Vector{RNN}
    function RNNs(topology::Vector{Int}, fn::Array{F}; type::Type=Array{Float32}) where F
        n = length(topology) - 1
        layers = Vector{RNN}(undef, n)
        for i = 1:n
            layers[i] = RNN(topology[i], topology[i+1], fn[i]; type=type)
        end
        new(layers)
    end
end

Base.getindex(m::RNNs,     k...) =  m.layers[k...]
Base.setindex!(m::RNNs, v, k...) = (m.layers[k...] = v)
Base.length(m::RNNs)       = length(m.layers)
Base.lastindex(m::RNNs)    = length(m.layers)
Base.firstindex(m::RNNs)   = 1
Base.iterate(m::RNNs, i=firstindex(m)) = i>length(m) ? nothing : (m[i], i+1)

function Base.show(io::IO, m::RNN)
    SIZE = size(m.w)
    TYPE = typeof(m.w.value)
    print(io, "RNN($(SIZE[2]), $(SIZE[1]), $(m.f); type=$TYPE)")
end

function Base.show(io::IO, m::RNNs)
    print(io, "RNNs\n      (\n          ")
    join(io, m.layers, ",\n          ")
    print(io, "\n      )")
end

function resethidden(m::RNN)
    m.h = nothing
end

function resethidden(model::RNNs)
    for m in model
        resethidden(m)
    end
end

function forward(m::RNN, x::Variable{T}) where T
    f = m.f  # activition function
    w = m.w  # input's weights
    b = m.b  # input's bias
    u = m.u  # memory's weights
    h = m.h ≠ nothing ? m.h : Variable{T}(Zeros(T, size(w,1), size(x,2)))
    x = f(matAddVec(w*x + u*h, b))
    m.h = x
    return x
end

function forward(model::RNNs, x::Variable)
    for m in model
        x = forward(m, x)
    end
    return x
end

function predict(m::RNN, x::T) where T
    f = m.f        # activition function
    w = m.w.value  # input's weights
    b = m.b.value  # input's bias
    u = m.u.value  # memory's weights
    h = m.h ≠ nothing ? m.h : Zeros(T, size(w,1), size(x,2))
    x = f(w*x + u*h .+ b)
    m.h = x
    return x
end

function predict(model::RNNs, x)
    for m in model
        x = predict(m, x)
    end
    return x
end

"""
    unbiasedof(m::RNN)

unbiased weights of RNN block
"""
function unbiasedof(m::RNN)
    weights = Vector(undef, 1)
    weights[1] = m.w.value
    return weights
end

function weightsof(m::RNN)
    weights = Vector(undef,3)
    weights[1] = m.w.value
    weights[2] = m.b.value
    weights[3] = m.u.value
    return weights
end

"""
    unbiasedof(model::RNNs)

unbiased weights of RNNs block
"""
function unbiasedof(model::RNNs)
    weights = Vector(undef, 0)
    for m in model
        append!(weights, unbiasedof(m))
    end
    return weights
end

function weightsof(m::RNNs)
    weights = Vector(undef,0)
    for i = 1:length(m)
        append!(weights, weightsof(m[i]))
    end
    return weights
end

function gradsof(m::RNN)
    grads = Vector(undef,3)
    grads[1] = m.w.delta
    grads[2] = m.b.delta
    grads[3] = m.u.delta
    return grads
end

function gradsof(m::RNNs)
    grads = Vector(undef,0)
    for i = 1:length(m)
        append!(grads, gradsof(m[i]))
    end
    return grads
end

function zerograds!(m::RNN)
    for v in gradsof(m)
        v .= zero(v)
    end
end

function zerograds!(m::RNNs)
    for v in gradsof(m)
        v .= zero(v)
    end
end

function paramsof(m::RNN)
    params = Vector{Variable}(undef,3)
    params[1] = m.w
    params[2] = m.b
    params[3] = m.u
    return params
end

function xparamsof(m::RNN)
    xparams = Vector{XVariable}(undef,3)
    xparams[1] = ('w', m.w)
    xparams[2] = ('b', m.b)
    xparams[3] = ('u', m.u)
    return xparams
end

function paramsof(m::RNNs)
    params = Vector{Variable}(undef,0)
    for i = 1:length(m)
        append!(params, paramsof(m[i]))
    end
    return params
end

function xparamsof(m::RNNs)
    xparams = Vector{XVariable}(undef,0)
    for i = 1:length(m)
        append!(xparams, xparamsof(m[i]))
    end
    return xparams
end

function nparamsof(m::RNN)
    lw = length(m.w)
    lb = length(m.b)
    lu = length(m.u)
    return (lw + lb + lu)
end

function bytesof(model::RNN, unit::String="MB")
    n = nparamsof(model) * elsizeof(model.w)
    return blocksize(n, uppercase(unit))
end

function nparamsof(m::RNNs)
    num = 0
    for i = 1:length(m)
        num += nparamsof(m[i])
    end
    return num
end

function bytesof(model::RNNs, unit::String="MB")
    n = nparamsof(model) * elsizeof(model[1].w)
    return blocksize(n, uppercase(unit))
end

function to(type::Type, m::RNN)
    m.w = to(type, m.w)
    m.b = to(type, m.b)
    m.u = to(type, m.u)
    return m
end

function to!(type::Type, m::RNN)
    m = to(type, m)
    return nothing
end

function to(type::Type, m::RNNs)
    for layer in m
        layer = to(type, layer)
    end
    return m
end

function to!(type::Type, m::RNNs)
    for layer in m
        to!(type, layer)
    end
end


import .Classification: fit!
using .Classification: SGDContext, StochasticGradientDescent
using  .NeuralNetwork: NeuralNet, backpropagation, evaluate!

function fit!(net::NeuralNet, context::SGDContext, X, y, labels)
    optim = StochasticGradientDescent(context)
    res = fit!(net, optim, X, y, labels)
    optim
end

function fit!(net::NeuralNet, optim::StochasticGradientDescent,
              X::Matrix{<:Real}, y::Matrix{<:Real}, labels)
    weave!(net, X, y)
    numsamples = size(X, 2)
    optim.loss = zeros(Float32, optim.context.maxiterations)
    batchsize = optim.context.batchsize
    if numsamples % batchsize != 0
        @warn "The minibatch size does not divide number of samples. Leftover $(numsamples%batchsize) dropped"
    end

    partitions = partition(numsamples, batchsize)
    for epoch in 1:optim.context.maxiterations
        shuffle!(X, y, labels)
        for batch in partitions
            minibatch!(net, optim, X, y, batch)
        end
        optim.iterations += 1
        evaluate!(net, optim, X, labels)
    end
end

function minibatch!(net::NeuralNet, optim::StochasticGradientDescent, X, y, batch)
    ∇b = [zeros(Float32, size(layer.b)) for layer in net.layers]
    ∇W = [zeros(Float32, size(layer.W)) for layer in net.layers]
    for i in batch
        Δ∇b, Δ∇W = backpropagation(net, X[:, i], y[:, i])
        ∇b .+= Δ∇b
        ∇W .+= Δ∇W
    end
    η = optim.context.learningrate
    B = length(batch)
    for (layer, ∂b, ∂W) in zip(net.layers, ∇b, ∇W)
        layer.b .-= η./B .* ∂b
        layer.W .-= η./B .* ∂W
    end
end

function shuffle!(X::Matrix{<:Real}, y::Matrix{<:Real})
    n = size(X, 2)
    j = 0
    for i in n:-1:2
        @fastmath j = 1 + abs(rand(Int))%n
        swapcols!(X, i, j)
        swapcols!(y, i, j)
    end
end

function shuffle!(X::Matrix{<:Real}, y::Matrix{<:Real}, z)
    n = size(X, 2)
    j = 0
    for i in n:-1:2
        @fastmath j = 1 + abs(rand(Int))%n
        swapcols!(X, i, j)
        swapcols!(y, i, j)
        swapcols!(z, i, j)
    end
end

function swapcols!(X, i, j)
    for k in 1:size(X, 1)
        X[k, i], X[k, j] = X[k, j], X[k, i]
    end
end


Pkg.add("DecisionTree")

using RDatasets: dataset
using DecisionTree

iris = dataset("datasets", "iris")
features = convert(Array, iris[:, 1:4]);
labels = convert(Array, iris[:, 5]);

# train full-tree classifier
model = build_tree(labels, features)
# prune tree: merge leaves having >= 90% combined purity (default: 100%)
model = prune_tree(model, 0.9)
# pretty print of the tree, to a depth of 5 nodes (optional)
print_tree(model, 5)
# apply learned model
apply_tree(model, [5.9,3.0,5.1,1.9])
# get the probability of each label
apply_tree_proba(model, [5.9,3.0,5.1,1.9], ["setosa", "versicolor", "virginica"])
# run n-fold cross validation for pruned tree,
# using 90% purity threshold pruning, and 3 CV folds
accuracy = nfoldCV_tree(labels, features, 0.9, 3)

# train random forest classifier
# using 2 random features, 10 trees, 0.5 portion of samples per tree (optional), and a maximum tree depth of 6 (optional)
model = build_forest(labels, features, 2, 10, 0.5, 6)
# apply learned model
apply_forest(model, [5.9,3.0,5.1,1.9])
# get the probability of each label
apply_forest_proba(model, [5.9,3.0,5.1,1.9], ["setosa", "versicolor", "virginica"])
# run n-fold cross validation for forests
# using 2 random features, 10 trees, 3 folds and 0.5 of samples per tree (optional)
accuracy = nfoldCV_forest(labels, features, 2, 10, 3, 0.5)

# train adaptive-boosted stumps, using 7 iterations
model, coeffs = build_adaboost_stumps(labels, features, 7);
# apply learned model
apply_adaboost_stumps(model, coeffs, [5.9,3.0,5.1,1.9])
# get the probability of each label
apply_adaboost_stumps_proba(model, coeffs, [5.9,3.0,5.1,1.9], ["setosa", "versicolor", "virginica"])
# run n-fold cross validation for boosted stumps, using 7 iterations and 3 folds
accuracy = nfoldCV_stumps(labels, features, 7, 3)

n, m = 10^3, 5 ;
features = randn(n, m);
weights = rand(-2:2, m);
labels = features * weights;

# train regression tree, using an averaging of 5 samples per leaf (optional)
model = build_tree(labels, features, 5)
# apply learned model
apply_tree(model, [-0.9,3.0,5.1,1.9,0.0])
# run n-fold cross validation, using 3 folds, averaging of 5 samples per leaf (optional)
# returns array of coefficients of determination (R^2)
r2 = nfoldCV_tree(labels, features, 3, 5)

# train regression forest, using 2 random features, 10 trees,
# averaging of 5 samples per leaf (optional), 0.7 of samples per tree (optional)
model = build_forest(labels,features, 2, 10, 5, 0.7)
# apply learned model
apply_forest(model, [-0.9,3.0,5.1,1.9,0.0])
# run n-fold cross validation on regression forest
# using 2 random features, 10 trees, 3 folds, averaging of 5 samples/leaf (optional),
# and 0.7 porition of samples per tree (optional)
# returns array of coefficients of determination (R^2)
r2 = nfoldCV_forest(labels, features, 2, 10, 3, 5, 0.7)


function optimal_simple(d::Divergence, m::DRO)
    p = zero(m.q)
    p[m.Imax] .= m.q[m.Imax]
    return p./sum(m.q[m.Imax])
end

# ----------------------------------------------------------------------------------------------------------
# Kullback-Leibler divergence
# ----------------------------------------------------------------------------------------------------------
"""
    KullbackLeibler

An empty structure representing Kullback-Leibler divergence.
"""
struct KullbackLeibler <: Divergence end

"""
    generate(d::KullbackLeibler)

Returns the generating function `ϕ(t) = t⋅log(t)` of the Kullback-Leibler divergence.
The generating function for `t == 0` is defined as `0`.
"""
generate(d::KullbackLeibler) = ϕ(t) = iszero(t) ? zero(t) : t*log(t)

"""
    name(d::KullbackLeibler)

Returns the full name of the Kullback-Leibler divergence.
"""
name(d::KullbackLeibler) = "Kullback-Leibler divergence"

"""
    bounds(d::KullbackLeibler, m::DRO)

Returns lower and upper bound for finding the root of the function `h` using the bisection method.
"""
bounds(d::KullbackLeibler, m::DRO) = (0, (m.cmax - m.cmin)/m.ε)

function h(d::KullbackLeibler, m::DRO, λ::Real)
    p_hat = m.q.*exp.(m.c./λ)
    val   = p_hat' * (m.c./λ .- log(sum(p_hat)) .- m.ε)
    return isnan(val) || val == - typemax(val) ? typemax(val) : val
end

"""
    function optimal(s::Our, d::KullbackLeibler, m::DRO; kwargs...)

Returns the optimal solution of the DRO model 'm' with Kullback-Leibler divergence.
"""
function optimal(s::Our, d::KullbackLeibler, m::DRO; kwargs...)
    p = optimal_simple(d, m)
    ϕ = generate(d)

    if sum(m.q .* ϕ.(p ./ m.q)) <= m.ε
        return p
    else
        λ = bisection(d, m; kwargs...)
        p = m.q.*exp.(m.c./λ)
        return p./sum(p)
    end
end

# ----------------------------------------------------------------------------------------------------------
# Burg entropy
# ----------------------------------------------------------------------------------------------------------
"""
    Burg

An empty structure representing Burg entropy.
"""
struct Burg <: Divergence end

"""
    generate(d::Burg, t)

Returns the generating function `ϕ(t) = - log(t)` of the Burg entropy.
"""
generate(d::Burg) = ϕ(t) = - log(t)

"""
    name(d::Burg)

Returns the full name of the Burg entropy.
"""
name(d::Burg) = "Burg entropy"

"""
    bounds(d::Burg, m::DRO)

Returns lower and upper bound for finding the root of the function `h` using the bisection method.
"""
bounds(d::Burg, m::DRO) = (m.cmax, m.cmax + (m.cmax - m.cmin)/m.ε)

function h(d::Burg, m::DRO, λ::Real)
    val = sum(m.q.*log.(λ .- m.c)) + log(sum(m.q./(λ .- m.c))) - m.ε
    return isnan(val) ? typemax(val) : val
end

"""
    function optimal(s::Our, d::Burg, m::DRO; kwargs...)

Returns the optimal solution of the DRO model 'm' with Burg entropy.
"""
function optimal(s::Our, d::Burg, m::DRO; kwargs...)
    p = optimal_simple(d, m)
    ϕ = generate(d)

    if sum(m.q .* ϕ.(p ./ m.q)) <= m.ε
        return p
    else
        λ = bisection(d, m; kwargs...)
        p = m.q./(λ .- m.c)
        return p./sum(p)
    end
end

# ----------------------------------------------------------------------------------------------------------
# Hellinger distance
# ----------------------------------------------------------------------------------------------------------
"""
    Hellinger

An empty structure representing Hellinger distance.
"""
struct Hellinger <: Divergence end

"""
    generate(d::Hellinger, t)

Returns the generating function `ϕ(t) = (√t - 1)²`  of the Hellinger distance.
"""
generate(d::Hellinger) = ϕ(t) = (sqrt(t) - 1)^2

"""
    name(d::Hellinger)

Returns the full name of the Hellinger distance.
"""
name(d::Hellinger) = "Hellinger distance"

"""
    bounds(d::Hellinger, m::DRO)

Returns lower and upper bound for finding the root of the function `h` using the bisection method.
"""
bounds(d::Hellinger, m::DRO) = (m.cmax, m.cmax + (2 - m.ε)*(m.cmax - m.cmin)/m.ε)

function h(d::Hellinger, m::DRO, λ::Real)
    val =  2*sum(m.q./(λ .- m.c)) - (2 - m.ε)*sqrt(sum(m.q./((λ .- m.c).^2)))
    return isnan(val) ? - typemax(val) : val
end

"""
    function optimal(s::Our, d::Hellinger, m::DRO; kwargs...)

Returns the optimal solution of the DRO model 'm' with Hellinger distance.
"""
function optimal(s::Our, d::Hellinger, m::DRO; kwargs...)
    p = optimal_simple(d, m)
    ϕ = generate(d)

    if sum(m.q .* ϕ.(p ./ m.q)) <= m.ε
        return p
    else
        λ = bisection(d, m; kwargs...)
        p = m.q./(λ .- m.c).^2
        return p./sum(p)
    end
end

# ----------------------------------------------------------------------------------------------------------
# χ²-distance
# ----------------------------------------------------------------------------------------------------------
"""
    ChiSquare

An empty structure representing χ²-distance.
"""
struct ChiSquare <: Divergence end

"""
    generate(d::ChiSquare, t)

Returns the generating function `ϕ(t) = (t - 1)²/t`  of the χ²-distance.
"""
generate(d::ChiSquare) = ϕ(t) = ((t - 1)^2)/t

"""
    name(d::ChiSquare)

Returns the full name of the χ²-distance.
"""
name(d::ChiSquare) = "χ²-distance"

"""
    initial(d::ChiSquare, m::DRO)

Returns the initial point for finding the root of the function `h` using the newton method.
"""
function initial(d::ChiSquare, m::DRO)
    f(μ)  = h(d, m, μ)
    λ_min = m.cmax
    q     = 0.01
    f_val = f(λ_min + q)
    add_eval()

    while f_val <= 0
        f == 0 && return λ_min + q

        q    /= 10
        f_val = f(λ_min + q)
        add_eval()
    end
    return λ_min + q
end

function h(d::ChiSquare, m::DRO, λ::Real)
    val = sum(m.q .* sqrt.(λ .- m.c))*sum(m.q./sqrt.(λ .- m.c)) - 1 - m.ε
    return isnan(val) ? - typemax(val) : val
end

function ∇h(d::ChiSquare, m::DRO, λ::Real)
    λc = λ .- m.c
    return (sum(m.q ./ sqrt.(λc))^2 - sum(m.q .* sqrt.(λc))*sum(m.q ./ (λc.^(3/2))))/2
end

"""
    function optimal(s::Our, d::ChiSquare, m::DRO; kwargs...)

Returns the optimal solution of the DRO model 'm' with χ²-distance.
"""
function optimal(s::Our, d::ChiSquare, m::DRO; kwargs...)
    p = optimal_simple(d, m)
    ϕ = generate(d)

    if sum(m.q .* ϕ.(p ./ m.q)) <= m.ε
        return p
    else
        λ = newton(d, m; kwargs...)
        p = m.q./sqrt.(λ .- m.c)
        return p./sum(p)
    end
end

# ----------------------------------------------------------------------------------------------------------
# Modified χ²-distance
# ----------------------------------------------------------------------------------------------------------
"""
    ModifiedChiSquare

An empty structure representing modified χ²-distance.
"""
struct ModifiedChiSquare <: Divergence end

"""
    generate(d::ModifiedChiSquare, t)

Returns the generating function `ϕ(t) = (t - 1)²`  of the modified χ²-distance.
"""
generate(d::ModifiedChiSquare) = ϕ(t) = (t - 1)^2

"""
    name(d::ModifiedChiSquare)

Returns the full name of the modified χ²-distance.
"""
name(d::ModifiedChiSquare) = "Modified χ²-distance"

"""
    initial(d::ModifiedChiSquare, m::DRO)

Returns the initial point for finding the root of the function `h` using the newton method.
"""
function initial(d::ModifiedChiSquare, m::DRO)
    f(λ)  = h(d, m, λ)
    λ     = - m.cmax + 10
    f_val = f(λ)
    add_eval()

    while f_val >= 0
        f == 0 && return λ

        λ    *= 10
        f_val = f(λ)
        add_eval()
    end
    return λ
end

function h(d::ModifiedChiSquare, m::DRO, λ::Real)
    val = sum(m.q .* max.(λ .+ m.c, 0).^2) - (1 + m.ε) * (sum(m.q .* max.(λ .+ m.c, 0)))^2
    return isnan(val) ? - typemax(val) : val
end

function ∇h(d::ModifiedChiSquare, m::DRO, λ::Real)
    I = findall(λ .+ m.c .>= 0)
    return 2 * sum(m.q[I] .* (λ .+ m.c[I])) * (1 - (1 + m.ε) * sum(m.q[I]))
end

"""
    function optimal(s::Our, d::ModifiedChiSquare, m::DRO; kwargs...)

Returns the optimal solution of the DRO model 'm' with modified χ²-distance.
"""
function optimal(s::Our, d::ModifiedChiSquare, m::DRO; kwargs...)
    p = optimal_simple(d, m)
    ϕ = generate(d)

    if sum(m.q .* ϕ.(p ./ m.q)) <= m.ε
        return p
    else
        λ = newton(d, m; kwargs...)
        p = m.q.*max.(λ .+ m.c, 0)
        return p./sum(p)
    end
end


using GeneralQP
using JuMP, Gurobi
using Random

rng = MersenneTwister(123)
n = 300
m = 500
P = randn(rng, n, n); P = (P + P')/2
q = randn(rng, n)
A = randn(rng, m, n); b = randn(rng, m)

# Choose Gurobi's primal simplex method
model = Model(solver=GurobiSolver(Presolve=0, Method=0))
@variable(model, x[1:size(A, 2)])
@constraint(model, A*x - b .<=0)
status = JuMP.solve(model)

x_init = getvalue(x)  # Initial point to be passed to our solver

using Main.GeneralQP

Main.GeneralQP.solve(P, q, A, b, x_init; verbosity=1)


"""
    em_coefficients(prim, E, B, mr, ld, rL, dt)

Calculate electromagnetic coeffcients in hyperbolic Maxwell's equations
"""
function em_coefficients(
    prim::AbstractMatrix{X},
    E::AbstractVector{Y},
    B::AbstractVector{Z},
    mr,
    lD,
    rL,
    dt,
) where {X<:Real,Y<:Real,Z<:Real}

    if eltype(prim) <: Int
        A = zeros(9, 9)
        b = zeros(9)
    else
        A = similar(prim, 9, 9)
        A .= 0.0
        b = similar(prim, 9)
        b .= 0.0
    end

    A[1, 1] = -1.0 / (2.0 * rL)
    A[2, 2] = -1.0 / (2.0 * rL)
    A[3, 3] = -1.0 / (2.0 * rL)
    A[4, 1] = mr / (2.0 * rL)
    A[5, 2] = mr / (2.0 * rL)
    A[6, 3] = mr / (2.0 * rL)
    A[7, 1] = 1.0 / (dt)
    A[8, 2] = 1.0 / (dt)
    A[9, 3] = 1.0 / (dt)

    A[1, 4] = 1.0 / (dt)
    A[1, 5] = -B[3] / (2.0 * rL)
    A[1, 6] = B[2] / (2.0 * rL)
    A[2, 4] = B[3] / (2.0 * rL)
    A[2, 5] = 1.0 / (dt)
    A[2, 6] = -B[1] / (2.0 * rL)
    A[3, 4] = -B[2] / (2.0 * rL)
    A[3, 5] = B[1] / (2.0 * rL)
    A[3, 6] = 1.0 / (dt)

    A[4, 7] = 1.0 / (dt)
    A[4, 8] = mr * B[3] / (2.0 * rL)
    A[4, 9] = -mr * B[2] / (2.0 * rL)
    A[5, 7] = -mr * B[3] / (2.0 * rL)
    A[5, 8] = 1.0 / (dt)
    A[5, 9] = mr * B[1] / (2.0 * rL)
    A[6, 7] = mr * B[2] / (2.0 * rL)
    A[6, 8] = -mr * B[1] / (2.0 * rL)
    A[6, 9] = 1.0 / (dt)

    A[7, 4] = prim[1, 1] / (2.0 * rL * lD^2)
    A[8, 5] = prim[1, 1] / (2.0 * rL * lD^2)
    A[9, 6] = prim[1, 1] / (2.0 * rL * lD^2)
    A[7, 7] = -(prim[1, 2] * mr) / (2.0 * rL * lD^2)
    A[8, 8] = -(prim[1, 2] * mr) / (2.0 * rL * lD^2)
    A[9, 9] = -(prim[1, 2] * mr) / (2.0 * rL * lD^2)

    b[1] =
        prim[2, 1] / (dt) + E[1] / (2.0 * rL) - B[2] * prim[4, 1] / (2.0 * rL) +
        B[3] * prim[3, 1] / (2.0 * rL)
    b[2] =
        prim[3, 1] / (dt) + E[2] / (2.0 * rL) - B[3] * prim[2, 1] / (2.0 * rL) +
        B[1] * prim[4, 1] / (2.0 * rL)
    b[3] =
        prim[4, 1] / (dt) + E[3] / (2.0 * rL) - B[1] * prim[3, 1] / (2.0 * rL) +
        B[2] * prim[2, 1] / (2.0 * rL)
    b[4] =
        prim[2, 2] / (dt) - mr * E[1] / (2.0 * rL) + mr * B[2] * prim[4, 2] / (2.0 * rL) -
        mr * B[3] * prim[3, 2] / (2.0 * rL)
    b[5] =
        prim[3, 2] / (dt) - mr * E[2] / (2.0 * rL) + mr * B[3] * prim[2, 2] / (2.0 * rL) -
        mr * B[1] * prim[4, 2] / (2.0 * rL)
    b[6] =
        prim[4, 2] / (dt) - mr * E[3] / (2.0 * rL) + mr * B[1] * prim[3, 2] / (2.0 * rL) -
        mr * B[2] * prim[2, 2] / (2.0 * rL)
    b[7] =
        E[1] / (dt) - prim[1, 1] * prim[2, 1] / (2.0 * rL * lD^2) +
        prim[1, 2] * prim[2, 2] * mr / (2.0 * rL * lD^2)
    b[8] =
        E[2] / (dt) - prim[1, 1] * prim[3, 1] / (2.0 * rL * lD^2) +
        prim[1, 2] * prim[3, 2] * mr / (2.0 * rL * lD^2)
    b[9] =
        E[3] / (dt) - prim[1, 1] * prim[4, 1] / (2.0 * rL * lD^2) +
        prim[1, 2] * prim[4, 2] * mr / (2.0 * rL * lD^2)

    return A, b

end