#####
# StandardLine vs. StandardLine
#####

function is_colliding(a::StandardLine{T}, b::StandardLine{T}, pos_ba::Vector2D{T}) where {T}
    y = pos_ba[2]
    if iszero(y)
        x = pos_ba[1]
        half_length_a = get_half_length(a)
        half_length_b = get_half_length(b)
        return !((x + half_length_b <= -half_length_a) || (x - half_length_b >= half_length_a))
    else
        return false
    end
end

function is_colliding(l1::StandardLine{T}, l2::StandardLine{T}, pos::Vector2D{T}, dir::Vector2D{T}) where {T}
    half_length = get_half_length(l1)

    tail_l2, head_l2 = get_vertices(l2, pos, dir)

    x1 = tail_l2[1]
    y1 = tail_l2[2]

    x2 = head_l2[1]
    y2 = head_l2[2]

    t = y1 / (y1 - y2)

    if isfinite(t)
        x_intersection = x1 + t * (x2 - x1)
        return (zero(t) < t < one(t)) && (-half_length < x_intersection < half_length)
    else
        return (iszero(y1)) && (iszero(y2)) && !((half_length <= x1) || (x2 <= -half_length))
    end
end

#####
# StandardCircle vs. StandardPoint
#####

function is_inside(circle::StandardCircle{T}, pos::Vector2D{T}) where {T}
    radius = get_radius(circle)
    return LA.dot(pos, pos) < radius * radius
end

is_colliding(circle::StandardCircle{T}, point::StandardPoint{T}, pos::Vector2D{T}) where {T} = is_inside(circle, pos)
is_colliding(point::StandardPoint{T}, circle::StandardCircle{T}, pos::Vector2D{T}) where {T} = is_inside(circle, pos) # no need to reverse pos because of symmetry

is_colliding(circle::StandardCircle{T}, point::StandardPoint{T}, pos::Vector2D{T}, dir::Vector2D{T}) where {T} = is_inside(circle, pos)
is_colliding(point::StandardPoint{T}, circle::StandardCircle{T}, pos::Vector2D{T}, dir::Vector2D{T}) where {T} = is_inside(circle, pos) # no need to reverse pos because of symmetry

#####
# StandardCircle vs. StandardLine
#####

function get_projection(line::StandardLine{T}, pos::Vector2D{T}) where {T}
    half_length = get_half_length(line)
    x = pos[1]
    if x < -half_length
        return get_tail(line)
    elseif x > half_length
        return get_head(line)
    else
        return typeof(pos)(x, zero(x))
    end
end

function is_colliding(line::StandardLine{T}, circle::StandardCircle{T}, pos::Vector2D{T}) where {T}
    projection = get_projection(line, pos)
    vec = pos - projection
    radius = get_radius(circle)
    return LA.dot(vec, vec) < radius * radius
end

is_colliding(circle::StandardCircle{T}, line::StandardLine{T}, pos::Vector2D{T}) where {T} = is_colliding(line, circle, pos) # no need to reverse pos because of symmetry

is_colliding(line::StandardLine{T}, circle::StandardCircle{T}, pos::Vector2D{T}, dir::Vector2D{T}) where {T} = is_colliding(line, circle, pos)
is_colliding(circle::StandardCircle{T}, line::StandardLine{T}, pos::Vector2D{T}, dir::Vector2D{T}) where {T} = is_colliding(line, circle, invert(pos, dir)...)

#####
# StandardCircle vs. StandardCircle
#####

function is_colliding(c1::StandardCircle{T}, c2::StandardCircle{T}, pos::Vector2D{T}) where {T}
    r1 = get_radius(c1)
    r2 = get_radius(c2)
    r = r1 + r2
    return LA.dot(pos, pos) < r * r
end

is_colliding(c1::StandardCircle{T}, c2::StandardCircle{T}, pos::Vector2D{T}, dir::Vector2D{T}) where {T} = is_colliding(c1, c2, pos)

#####
# StandardRect vs. StandardPoint
#####

function is_inside(rect::StandardRect{T}, pos::Vector2D{T}) where {T}
    half_width = get_half_width(rect)
    half_height = get_half_height(rect)

    x = pos[1]
    y = pos[2]

    return (-half_width < x < half_width) && (-half_height < y < half_height)
end

is_colliding(rect::StandardRect{T}, point::StandardPoint{T}, pos::Vector2D{T}) where {T} = is_inside(rect, pos)
is_colliding(point::StandardPoint{T}, rect::StandardRect{T}, pos::Vector2D{T}) where {T} = is_inside(rect, pos) # no need to reverse pos because of symmetry

is_colliding(rect::StandardRect{T}, point::StandardPoint{T}, pos::Vector2D{T}, dir::Vector2D{T}) where {T} = is_inside(rect, pos)
is_colliding(point::StandardPoint{T}, rect::StandardRect{T}, pos::Vector2D{T}, dir::Vector2D{T}) where {T} = is_colliding(rect, point, invert(pos, dir)...)

#####
# StandardRect vs. StandardLine
#####

function separating_axis_exists(rect::StandardRect{T}, line::StandardLine{T}, pos::Vector2D{T}) where {T}
    half_height = get_half_height(rect)
    half_width = get_half_width(rect)

    x = pos[1]
    y = pos[2]
    half_length = get_half_length(line)

    return (y <= -half_height) || (y >= half_height) || (x + half_length <= -half_width) || (x - half_length >= half_width)
end

is_colliding(rect::StandardRect{T}, line::StandardLine{T}, pos::Vector2D{T}) where {T} = !separating_axis_exists(rect, line, pos)
is_colliding(line::StandardLine{T}, rect::StandardRect{T}, pos::Vector2D{T}) where {T} = is_colliding(rect, line, pos) # no need to reverse pos because of symmetry

function separating_axis_exists(rect::StandardRect{T}, line::StandardLine{T}, pos::Vector2D{T}, dir::Vector2D{T}) where {T}
    half_width = get_half_width(rect)
    half_height = get_half_height(rect)

    tail, head = get_vertices(line, pos, dir)

    min_x, max_x = minmax(tail[1], head[1])
    min_y, max_y = minmax(tail[2], head[2])

    return (max_x <= -half_width) || (min_x >= half_width) || (max_y <= -half_height) || (min_y >= half_height)
end

function separating_axis_exists(line::StandardLine{T}, rect::StandardRect{T}, pos::Vector2D{T}, dir::Vector2D{T}) where {T}
    half_length = get_half_length(line)

    bottom_left, bottom_right, top_right, top_left = get_vertices(rect, pos, dir)

    min_x, max_x = extrema((bottom_left[1], bottom_right[1], top_right[1], top_left[1]))
    min_y, max_y = extrema((bottom_left[2], bottom_right[2], top_right[2], top_left[2]))

    return (max_x <= -half_length) || (min_x >= half_length) || (max_y <= zero(T)) || (min_y >= zero(T))
end

is_colliding(rect::StandardRect{T}, line::StandardLine{T}, pos::Vector2D{T}, dir::Vector2D{T}) where {T} = !(separating_axis_exists(rect, line, pos, dir) || separating_axis_exists(line, rect, invert(pos, dir)...))
is_colliding(line::StandardLine{T}, rect::StandardRect{T}, pos::Vector2D{T}, dir::Vector2D{T}) where {T} = !(separating_axis_exists(line, rect, pos, dir) || separating_axis_exists(rect, line, invert(pos, dir)...))

#####
# StandardRect vs. StandardCircle
#####

get_projection(rect::StandardRect{T}, pos::Vector2D{T}) where {T} = clamp.(pos, get_bottom_left(rect), get_top_right(rect))

function is_colliding(rect::StandardRect{T}, circle::StandardCircle{T}, pos::Vector2D{T}) where {T}
    projection = get_projection(rect, pos)
    vec = pos - projection
    radius = get_radius(circle)
    return LA.dot(vec, vec) < radius * radius
end

is_colliding(circle::StandardCircle{T}, rect::StandardRect{T}, pos::Vector2D{T}) where {T} = is_colliding(rect, circle, pos) # no need to invert pos because of symmetry

is_colliding(rect::StandardRect{T}, circle::StandardCircle{T}, pos::Vector2D{T}, dir::Vector2D{T}) where {T} = is_colliding(rect, circle, pos)
is_colliding(circle::StandardCircle{T}, rect::StandardRect{T}, pos::Vector2D{T}, dir::Vector2D{T}) where {T} = is_colliding(rect, circle, invert(pos, dir)...)

#####
# StandardRect vs. StandardRect
#####

function separating_axis_exists(r1::StandardRect{T}, r2::StandardRect{T}, pos::Vector2D{T}) where {T}
    half_width_r1 = get_half_width(r1)
    half_height_r1 = get_half_height(r1)

    half_width_r2 = get_half_width(r2)
    half_height_r2 = get_half_height(r2)

    x = pos[1]
    y = pos[2]

    return (x + half_width_r2 <= -half_width_r1) || (x - half_width_r2 >= half_width_r1) || (y + half_height_r2 <= -half_height_r1) || (y - half_height_r2 >= half_height_r1)
end

is_colliding(a::StandardRect{T}, b::StandardRect{T}, pos_ba::Vector2D{T}) where {T} = !separating_axis_exists(a, b, pos_ba)

function separating_axis_exists(r1::StandardRect{T}, r2::StandardRect{T}, pos::Vector2D{T}, dir::Vector2D{T}) where {T}
    half_width_r1 = get_half_width(r1)
    half_height_r1 = get_half_height(r1)

    bottom_left_r2, bottom_right_r2, top_right_r2, top_left_r2 = get_vertices(r2, pos, dir)

    min_x_r2, max_x_r2 = extrema((bottom_left_r2[1], bottom_right_r2[1], top_right_r2[1], top_left_r2[1]))
    min_y_r2, max_y_r2 = extrema((bottom_left_r2[2], bottom_right_r2[2], top_right_r2[2], top_left_r2[2]))

    return ((half_width_r1 <= min_x_r2) || (max_x_r2 <= -half_width_r1) || (half_height_r1 <= min_y_r2) || (max_y_r2 <= -half_height_r1))
end

is_colliding(r1::StandardRect{T}, r2::StandardRect{T}, pos::Vector2D{T}, dir::Vector2D{T}) where {T} = !(separating_axis_exists(r1, r2, pos, dir) || separating_axis_exists(r2, r1, invert(pos, dir)...))


import Random.randperm
import Base.copy

include("MOKP.jl")
include("tools.jl")

mutable struct Individu
    x::Array{Int,1} # vecteur solution
    y::Array{Int,1} # valeurs sur les objectifs
    rang::Int        # rang
    listeDomines::Array{Individu,1}    # liste des individus dominés par celui-ci, vide au début
    nbDomine::Int                    # nombre d'individus qui dominent celui-ci, vide au début
    distanceCrowding::Float64  # valeur de la distance de Crowding
end

function copy(a::Individu)
    return Individu(copy(a.x), copy(a.y), copy(a.rang), copy(a.listeDomines), copy(a.nbDomine), copy(a.distanceCrowding))
end

function verification_NSGA(prob::_bi01IP, x::solution)
    poids = zeros(size(prob.A, 1))
    objectif = [0, 0]
    for i = 1:size(prob.A, 1)
        for j = 1:size(prob.A, 2)
            poids[i] += prob.A[i, j] * x.sol[j]
        end
    end
    for i = 1:size(prob.C, 2)
        objectif[1] += prob.C[1, i] * x.sol[i]
        objectif[2] += prob.C[2, i] * x.sol[i]
    end
    if sum(poids .<= prob.b) == length(prob.b)
        x.val_objectif = objectif
        x.cout = poids
        return true
    else
        return false
    end
end

function genererPopulation(nIndividus::Int, prob::_bi01IP)
    pop = Vector{Individu}(undef, nIndividus)
    len = 1

    x1 = Individu(pointsExtremes(prob, 1), [0, 0], 0, [], 0, 0.0)
    x2 = Individu(pointsExtremes(prob, 2), [0, 0], 0, [], 0, 0.0)
    x1 = updateY(x1, prob)
    x2 = updateY(x2, prob)

    pop[len] = x1
    len += 1
    pop[len] = x2
    len += 1

    delta2 = abs(x2.y[2] - x1.y[2])
    delta1 = abs(x2.y[1] - x1.y[1])

    step = delta2 / 10 / 2

    i = 1
    while i * step < delta2
        x = Individu(solutionEpsilon(prob, 1, x1.y[2] + i * step), [0, 0], 0, [], 0, 0.0)
        x = updateY(x, prob)
        pop[len] = x
        len += 1
        i += 1
    end
    step = delta1 / 10 / 2
    i = 1

    while i * step < delta1
        x = Individu(solutionEpsilon(prob, 2, x2.y[1] + i * step), [0, 0], 0, [], 0, 0.0)
        x = updateY(x, prob)
        pop[len] = x
        len += 1
        i += 1
    end

    while len <= nIndividus
        sol = solution([], zeros(size(prob.C, 2)), [])
        rand_indexes = randperm(length(sol.sol))
        idx = 1
        while verification(prob, sol)
            sol.sol[rand_indexes[idx]] = 1
            idx += 1
        end

        indiv = Individu(sol.sol, sol.val_objectif, 0, [], 0, 0.0)
        pop[len] = indiv
        len += 1
    end
    return pop
end

# Prends deux individus et renvoie true si le premier domine le second
function domine(p::Individu, q::Individu)
    if p.y[1] >= q.y[1] && p.y[2] >= q.y[2]
        return true
    end
    return false
end

# Prend un population et renvoie vrai s'il existe un Individu
#avec un rang non assigné
function rangZero(population)
    for p in population
        if p.rang == 0
            return true
        end
    end
    return false
end

# Prend une population et assigne un rang à chaque individus
function triNonDomine(population)
    for p in population
        empty!(p.listeDomines)
        p.nbDomine = 0
        p.rang = 0
    end

    for i = 1:length(population)
        for j = i+1:length(population)
            if domine(population[i], population[j])
                push!(population[i].listeDomines, population[j])
                population[j].nbDomine += 1
            elseif domine(population[j], population[i])
                push!(population[j].listeDomines, population[i])
                population[i].nbDomine += 1
            end
        end
        if population[i].nbDomine == 0
            population[i].rang = 1
        end
    end

    k = 2
    while rangZero(population)
        for p in population
            if p.rang == k - 1
                for q in p.listeDomines
                    q.nbDomine -= 1
                    if q.nbDomine == 0
                        q.rang = k
                    end
                end
            end
        end
        k += 1
    end
end

#Prend une population et renvoie le front maximum
function trouverMaxFront(population)
    max = 0
    for p in population
        if p.rang > max
            max = p.rang
        end
    end
    return max
end

# Prend un frond et assigne la distance de corwding à chaque individu du front
function distanceCrowding(front)
    for i = 1:length(front[1].y)
        sort!(front, lt = (a, b) -> a.y[i] < b.y[i], rev = true)

        max = front[1].y[i]
        min = front[end].y[i]

        for j = 2:(length(front)-1)
            greaterNeighbor = front[j-1].y[i]
            lowerNeighbor = front[j+1].y[i]
            front[j].distanceCrowding += ((greaterNeighbor - lowerNeighbor) / (max - min))
        end
    end
end

#renvoie le gagnant du tournoi entre deux individus
function tournoi(a::Individu, b::Individu)
    if a.rang < b.rang
        return a
    elseif b.rang < a.rang
        return b
    elseif a.distanceCrowding < b.distanceCrowding
        return a
    elseif b.distanceCrowding < a.distanceCrowding
        return b
    end
    return rand((a, b))
end

#teste si deux individus sont égaux
function isequal(a::Individu, b::Individu)
    return a.x == b.x
end

# prends un individu et assigne le y à son vecteur x correspondant
# Si le vecteur ne respecte pas les conditions du probleme alors
# des xi aléatoires passent de 1 à 0 jusqu'à ce que l'individu respecte les conditions
function updateY(a::Individu, prob::_bi01IP)
    a.y[1], a.y[2] = 0, 0
    for i = 1:length(a.x)
        a.y[1] += prob.C[1, i] * a.x[i]
        a.y[2] += prob.C[2, i] * a.x[i]
    end
    poids = zeros(size(prob.A, 1))
    for i = 1:size(prob.A, 1)
        for j = 1:size(prob.A, 2)
            poids[i] += prob.A[i, j] * a.x[j]
        end
    end

    for i = 1:length(prob.b)
        while poids[i] > prob.b[i]
            pos = rand(1:length(a.x))
            if a.x[pos] == 1
                a.x[pos] = 0
                a.y[1] -= prob.C[1, pos]
                a.y[2] -= prob.C[2, pos]
                for j = 1:size(prob.A, 1)
                    poids[j] -= prob.A[j, pos]
                end
            end
        end
    end
    return a
end

# réalise un crossover entre 2 individus
function crossover(a::Individu, b::Individu, probCrossover::Float64, prob::_bi01IP)
    if rand() <= probCrossover
        i = rand(1:length(a.x))
        savea = copy(a.x[i:end])
        a.x[i:end] = b.x[i:end]
        b.x[i:end] = savea

        a = updateY(a, prob)
        b = updateY(b, prob)
    end
    return a, b
end

# réalise une simple mutation en inversant un bit d'indice aléatoire d'un individu a
function mutation(a::Individu, prob::_bi01IP, probMutation::Float64)
    if rand() <= probMutation
        i = rand(1:length(a.x))
        if a.x[i] == 1
            a.x[i] = 0
            a.y[1] -= prob.C[1, i]
            a.y[2] -= prob.C[2, i]
        else
            a.x[i] = 1
            a = updateY(a, prob)
        end
    end
    return a
end

function genererPopulationEnfants(nIndividus::Int, probCrossover::Float64, probMutation::Float64, population, prob::_bi01IP)
    enfants = Vector{Individu}(undef, nIndividus)
    len = 1
    for i = 1:2:nIndividus
        indiv1 = tournoi(population[rand(1:length(population))], population[rand(1:length(population))])
        indiv2 = tournoi(population[rand(1:length(population))], population[rand(1:length(population))])

        a = copy(indiv1)
        b = copy(indiv2)
        a, b = crossover(a, b, probCrossover, prob)

        a = mutation(a, prob, probMutation)
        b = mutation(b, prob, probMutation)

        enfants[len] = a
        len += 1
        enfants[len] = b
        len += 1
    end
    return enfants
end

#Selectionne une population de n individus ou n + epsilon parmis 2xn individus
function selectionIndividus(popTotale, nIndividus::Int)
    if popTotale[nIndividus].rang == 1
        debut = findfirst(a -> a.rang == 2, popTotale)
        fin = findlast(a -> a.rang == 2, popTotale)
        popsecond = view(popTotale, debut:fin)
        distanceCrowding(popsecond)
        sort!(popsecond, by = a -> a.distanceCrowding)
        epsilon = round(Int, length(popsecond) / 5)
        popFinale = union(view(popTotale, 1:debut-1), view(popsecond, 1:epsilon))
    else
        debut = findfirst(a -> a.rang == popTotale[nIndividus].rang, popTotale)
        fin = findlast(a -> a.rang == popTotale[nIndividus].rang, popTotale)
        popsecond = view(popTotale, debut:fin)
        distanceCrowding(popsecond)
        sort!(popsecond, by = a -> a.distanceCrowding)
        popFinale = union(view(popTotale, 1:debut-1), view(popsecond, 1:(nIndividus-debut+1)))
    end
    return popFinale
end

#
function kung(population::Array{Individu,1})
    SN::Vector{Individu} = []
    sort!(population, by = a -> a.y[1])
    push!(SN, population[1])
    minYFeas = population[1].y[2]
    for i = 2:length(population)
        if population[i].y[2] < minYFeas
            push!(SN, population[i])
            minYFeas = population[i].y[2]
        end
    end
    return SN
end

function nsga2(nIndividus::Int, nGenerations::Int, probCrossover::Float64, probMutation::Float64, prob::_bi01IP, timeAllowed)
    #print("temps de gen de pop :")
    population = genererPopulation(nIndividus, prob)
    triNonDomine(population)
    t = 0.0
    start = time()
    gen = 0
    while gen < nGenerations && t < timeAllowed
        enfants = genererPopulationEnfants(nIndividus, probCrossover, probMutation, population, prob)
        popTotale = union(enfants, population)
        triNonDomine(popTotale)
        sort!(popTotale, by = a -> a.rang)
        population = selectionIndividus(popTotale, nIndividus)
        t = time() - start
        gen += 1
    end
    SN = kung(population)
    # println("pts non dominés : ", length(SN))
    # println("generations : ", gen)
    # println("time : ", time() - start)
    return SN
end

function get_YN_nsga(SN::Vector{Individu})
    return map(x -> x.y, SN)
end


# for use with static, axis-symmetric metrics
# create new abstract type for easy re-definition
abstract type AbstractAutoDiffMetricParams{T} <: AbstractMetricParams{T} end
abstract type AbstractAutoDiffStaticAxisSymmetricParams{T} <:
              AbstractAutoDiffMetricParams{T} end

# interface
metric_components(m::AbstractAutoDiffStaticAxisSymmetricParams{T}, rθ) where {T} =
    error("Not implemented for $(typeof(m)).")

"""
    inverse_metric_components(g_comp)

Calculates ``g^{tt}``, ``g^{rr}``, ``g^{\\theta\\theta}``, ``g^{\\phi\\phi}``, ``g^{t\\phi}`` of a static,
axis-symmetric metric from ``g_{tt}``, ``g_{rr}``, ``g_{\\theta\\theta}``, ``g_{\\phi\\phi}``, ``g_{t\\phi}``
using a symbolically computed inverse matrix method.

## Developer notes

To recreate:

```julia
using Symbolics
@variables g[1:5] # non zero metric components
metric = [
    g[1] 0 0 g[5]
    0 g[2] 0 0
    0 0 g[3] 0
    g[5] 0 0 g[4]
]
inv(metric)
```
"""
@fastmath function inverse_metric_components(g_comp)
    @inbounds let g1 = g_comp[1],
        g2 = g_comp[2],
        g3 = g_comp[3],
        g4 = g_comp[4],
        g5 = g_comp[5]

        (
            (g2 * g3 * g4) / (g1 * g2 * g3 * g4 - (g5^2) * g2 * g3),
            (g1 * g3 * g4 - (g5^2) * g3) / (g1 * g2 * g3 * g4 - (g5^2) * g2 * g3),
            (g1 * g2 * g4 - (g5^2) * g2) / (g1 * g2 * g3 * g4 - (g5^2) * g2 * g3),
            (g1 * g2 * g3) / (g1 * g2 * g3 * g4 - (g5^2) * g2 * g3),
            (-g2 * g3 * g5) / (g1 * g2 * g3 * g4 - (g5^2) * g2 * g3)
        )
    end
end

"""
    calc_symb_geodesic_equation(g1inv, j1, j2, v)

Symbolically pre-computed geodesic equations for a static, axis-symmetric metric.

## Developer notes

To recreate:

```julia
using Symbolics, Tullio
@variables ginv[1:5], j1[1:5], j2[1:5], v[1:4] # non zero metric components
inverse_metric = [
    ginv[1] 0 0 ginv[5]
    0 ginv[2] 0 0
    0 0 ginv[3] 0
    ginv[5] 0 0 ginv[4]
]
j1_mat = [
    j1[1] 0 0 j1[5]
    0 j1[2] 0 0
    0 0 j1[3] 0
    j1[5] 0 0 j1[4]
]
j2_mat = [
    j2[1] 0 0 j2[5]
    0 j2[2] 0 0
    0 0 j2[3] 0
    j2[5] 0 0 j2[4]
]
j0 = zeros(Float64, (4, 4))
jacobian = (j0, j1_mat, j2_mat, j0)
# christoffel symbols
@tullio Γ[i, k, l] :=
    1 / 2 *
    inverse_metric[i, m] *
    (jacobian[l][m, k] + jacobian[k][m, l] - jacobian[m][k, l])
# compute geodesic equation
@tullio δxδλ[i] := -v[j] * Γ[i, j, k] * v[k]
```
"""
@fastmath function compute_geodesic_equation(ginv, j1, j2, v::SVector{4,T}) where {T}
    @inbounds @SVector [
        -2(0.5ginv[5] * j1[4] + 0.5ginv[1] * j1[5]) * v[2] * v[4] -
        2(0.5ginv[1] * j1[1] + 0.5ginv[5] * j1[5]) * v[1] * v[2] -
        2(0.5ginv[1] * j2[1] + 0.5ginv[5] * j2[5]) * v[1] * v[3] -
        2(0.5ginv[5] * j2[4] + 0.5ginv[1] * j2[5]) * v[3] * v[4],
        0.5(v[1]^2) * ginv[2] * j1[1] +
        0.5(v[3]^2) * ginv[2] * j1[3] +
        0.5(v[4]^2) * ginv[2] * j1[4] +
        ginv[2] * j1[5] * v[1] * v[4] - 0.5(v[2]^2) * ginv[2] * j1[2] -
        ginv[2] * j2[2] * v[2] * v[3],
        0.5(v[1]^2) * ginv[3] * j2[1] +
        0.5(v[2]^2) * ginv[3] * j2[2] +
        0.5(v[4]^2) * ginv[3] * j2[4] +
        ginv[3] * j2[5] * v[1] * v[4] - 0.5(v[3]^2) * ginv[3] * j2[3] -
        ginv[3] * j1[3] * v[2] * v[3],
        -2(0.5ginv[4] * j1[4] + 0.5ginv[5] * j1[5]) * v[2] * v[4] -
        2(0.5ginv[4] * j2[4] + 0.5ginv[5] * j2[5]) * v[3] * v[4] -
        2(0.5ginv[5] * j1[1] + 0.5ginv[4] * j1[5]) * v[1] * v[2] -
        2(0.5ginv[4] * j2[5] + 0.5ginv[5] * j2[1]) * v[1] * v[3]
    ]
end

@fastmath function constrain_time(g_comp, v, μ=0.0, positive::Bool = true)
    @inbounds begin
        discriminant = (
            -g_comp[1] * g_comp[2] * v[2]^2 - g_comp[1] * g_comp[3] * v[3]^2 -
            g_comp[1] * μ^2 - (g_comp[1] * g_comp[4] - g_comp[5]^2) * v[4]^2
        )
        if positive
            -(g_comp[5] * v[4] + √discriminant) / g_comp[1]
        else
            -(g_comp[5] * v[4] - √discriminant) / g_comp[1]
        end
    end
end

function constrain(m::AbstractAutoDiffStaticAxisSymmetricParams{T}, u, v; μ::T = 0.0) where {T}
    rθ = @SVector [u[2], u[3]]
    g_comps = metric_components(m, rθ)
    constrain_time(g_comps, v, μ)
end

@inbounds function geodesic_eq(
    m::AbstractAutoDiffStaticAxisSymmetricParams{T},
    u,
    v
) where {T}
    # get the only position components we need for this metric type
    rθ = @SVector [u[2], u[3]]
    # calculate all non-zero components
    g_comps = metric_components(m, rθ)
    # use AD to get derivatives
    jacs = ForwardDiff.vector_mode_jacobian(x -> metric_components(m, x), rθ)
    # calculate all non-zero inverse matric components
    ginv_comps = inverse_metric_components(g_comps)
    # calculate acceleration
    compute_geodesic_equation(ginv_comps, jacs[:, 1], jacs[:, 2], v)
end

export AbstractAutoDiffStaticAxisSymmetricParams


# Implementation of an elementary cellular automata
# according to https://mathworld.wolfram.com/ElementaryCellularAutomaton.html
# uses random initialization and uses wolfram codes to specify the rule
# Asynchronous update of the 1D lattice

using Agents, Random
using CairoMakie
using InteractiveDynamics
using CSV

"""
The automaton living in a 1D space
"""
mutable struct Cell <: AbstractAgent
    id::Int
    pos::Dims{1}
    status::Int # either 0 or 1, where 1 is 'alive'
end

"""
Returns an array with the rule for the next status configurations, 1-indexed.

Thus, the first position corresponds to the rule for the cell neighborhood status (0,0,0)
and the last to (1,1,1)
"""
function rule_from_code(wolfram_code)
    return digits(wolfram_code, base=2, pad=8)
end

"""
Takes the status of a neighborhood and returns the corresponding
index in the model rule (0,0,0) -> 1 (0,1,0) -> 3
"""
function configuration_index(cell_statuses)
    # takes the tuples and forms a string (0,1,1) -> "011"
    binary_code = string(cell_statuses...)
    # turns the string binary code into the integer
    index = parse(Int, binary_code, base=2) + 1
    return index
end

"""
Given a cell checks its neighbors and
decides the next status of the cell based on the model rule
"""
function next_status(cell, model)
    neighbors = collect(nearby_agents(cell, model))
    cell_statuses = (c.status for c in [neighbors[1], cell, neighbors[2]])
    index = configuration_index(cell_statuses)
    return model.rule[index]
end

"""
Initializes the ABM

# Arguments
- `n_cells::Int`: total of cell automaton
- `wolfram_code::Int`: wolfram code for the rule
- `seed::Int`: random seed
"""
function build_model(; n_cells = 100, wolfram_code=30, seed = 30)
    space = GridSpace((n_cells,); metric=:chebyshev)

    properties = Dict(:rule => rule_from_code(wolfram_code),)

    model = ABM(
        Cell,
        space; properties,
        scheduler=Schedulers.randomly,
        rng = MersenneTwister(seed))

    for x in 1:n_cells
        cell = Cell(nextid(model), (x,), rand([0,1]))
        add_agent_pos!(cell, model)
    end

    return model
end

"""
Asynchronous update of the cells
"""
function cell_step!(cell, model)
    cell.status = next_status(cell, model)
end

# Initialize model
model = build_model(n_cells=100, wolfram_code=30)

# Runs the model and collects data
data, _ = run!(model, cell_step!, 100; adata=[:status]);

# The data contains the step, id/position and status of the cell (1/0)
data

CSV.write("ca_data_async.csv", data);

# Lets plot the time evolution in the y axis
heatmap(data.id, data.step, data.status, colormap=:Blues_3)


using NonlinearEigenproblemsTest
using NonlinearEigenproblems
using Test
using LinearAlgebra

@bench @testset "cd player native" begin
    nep = nep_gallery("nlevp_native_cd_player")
    n = size(nep, 1)

    (Λ,V)=polyeig(nep)

    (λ,v)=quasinewton(nep,λ=Λ[1],v=V[:,1],logger=displaylevel, armijo_factor=0.5,armijo_max=10, tol=1e-10)
    verify_lambdas(1, nep, λ, v, 1e-10)

    (λ,v)=resinv(nep,λ=Λ[2],v=V[:,2],logger=displaylevel, armijo_factor=0.5,armijo_max=10, tol=1e-10)
    verify_lambdas(1, nep, λ, v, 1e-10)

    @testset "Errors thrown" begin
        @test_throws MethodError nep_gallery("nlevp_native_cd_player", 15)
        @test_throws MethodError nep_gallery("nlevp_native_cd_player", t=15)
    end

end


abstract Model # this is either an LSTM or RNN

type NNMatrix # Neural net layer's weights & gradients
    n::Int
    d::Int
    w::Matrix{Float64} # matrix of weights
    dw::Matrix{Float64} # matrix of gradients
    NNMatrix(n::Int) = new(n, 1, zeros(n,1), zeros(n,1))
    NNMatrix(n::Int, d::Int) = new(n, d, zeros(n,d), zeros(n,d))
    NNMatrix(n::Int, d::Int, w::Array, dw::Array) = new(n, d, w, dw)
end

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

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


using ChaosTools, DynamicalSystemsBase
using Test
using ChaosTools.Distances: Cityblock, Euclidean

test_value = (val, vmin, vmax) -> @test vmin <= val <= vmax

println("\nTesting nonlinear timeseries analysis...")

@testset "Broomhead-King" begin
    using Random
    Random.seed!(1234);
    ds = Systems.gissinger(ones(3)) # standard initial condition

    data = trajectory(ds, 1000.0; Δt = 0.05)
    x = data[1:end-1, 1]
    s = x .+ 0.01rand(length(x))

    Ux, Σx = broomhead_king(x, 40)
    Us, Σs = broomhead_king(s, 40)
    R = embed(x, 2, 30)

    for j in 10:40
        @test Σx[j] < Σs[j]
    end

    newcoords = Dataset(Us[:, 1], Us[:, 2])
    newcoordsclean = Dataset(Ux[:, 1], Ux[:, 2])

    Dnew = generalized_dim(newcoords)
    DR = generalized_dim(R)
    DC = generalized_dim(newcoordsclean)
    @test abs(Dnew - 0.1 - DR) < 0.2 # subtract 0.1 for "added dimensionality"
    @test abs(DC - DR) < 0.2
end


using JuMP
include("../src/JOH.jl")

function one_hot_encode(index, group, n_groups)
    enc = zeros(n_groups)
    enc[index] = 1
    map(elem->vcat(elem.resources, enc), group)
end

function jump_model(problem::MMKP_Data, num_threads)::Model
    model = Model()
	JOH.Matheur.set_threads!(model, num_threads)

    #get matrix of resource requirements per item
    n_groups = length(problem.groups)
    groups = map(t -> one_hot_encode(t[1], t[2], n_groups), enumerate(problem.groups))
    flat_groups = vcat(groups...)
    resource_mat = hcat(flat_groups...)

    #add the one-hot dimensions to the resource bounds
    res_lims = vcat(problem.resource_bounds, ones(n_groups))

    #extract the profit
    profit = vcat(map(g->[i.profit for i in g], problem.groups)...)

    #make the model
    @variable(model, x[1:length(profit)], Bin)
    @objective(model, Max, sum(profit .* x))
    for (i, row) in enumerate(eachrow(resource_mat))
        @constraint(model, sum(row .* x) <= res_lims[i])
    end

    model
end


include(joinpath(@__DIR__, "../../src/FLS.jl"))

fxb = Symbol("f×ᵦg")

eqschema = @acset LabeledGraph begin
    V = 3
    E = 4
    vlabel = [:A,:B,fxb]
    elabel = [:f,:g, :π1, :π2]
    src = [1,1,3,3]
    tgt = [2,2,1,1]
end

eqconed = @acset LabeledGraph begin
    V = 3
    E = 2
    vlabel = [:A,:A,:B]
    elabel = [:f,:g]
    src = [1,2]
    tgt = [3,3]
end

eqfls = FLS(:equalizer, eqschema, [
    Cone(eqconed, fxb, [1=>:π1,2=>:π2])], []);


##  SCRIPT NAME:   equalityConstraints1.jl
##  DATE:          March 2015
#
#   This script returns the equality constraints subject to which the
#   family maximises its objective function in each period. If a constraint
#   is of the following form:
#   x + y = a
#   then this must be returned by this script as:
#   result[i] = x + y - a
#   which is set to 0 by the optimiser. The family problem is subject to the
#   following constraints:
#   (a) sequential budget constraint
#
#
##  ALEXANDROS THELOUDIS, UCL

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

function equalityConstraints1(choice_var::Vector, grad::Vector, A0, wH0, wW0, p0)

  # Declare choice variables:
  A!    = choice_var[1]
  consH = choice_var[2]
  consW = choice_var[3]
  kons  = choice_var[4]

  # Calculate gradient:
  if length(grad) > 0
    grad[1] = -1/(1+r)
    grad[2] = -p0
    grad[3] = -p0
    grad[4] = -1
  end

  # Write equality constraints and return results:
  # (a) Sequential budget constraint:
  A0 + wH0*maxAnnWorkingHrsH + wW0*maxAnnWorkingHrsW - p0*(consH + consW) - kons - A!/(1+r)

end


using SchwarzChristoffel
using Test

using StaticArrays

@testset "integral" begin
    f = SchwarzChristoffelDerivate(SVector(0),SVector(0)); @test integral(f,1,1) ≈ 1
    f = SchwarzChristoffelDerivate(SVector(0),SVector(0)); @test integral(f,1,2) ≈ 2
    f = SchwarzChristoffelDerivate(SVector(1),SVector(0)); @test integral(f,1,2) ≈ 1

    f = SchwarzChristoffelDerivate(SVector(0),SVector(-1)); @test integral(f,1,1) ≈ 0.5
    f = SchwarzChristoffelDerivate(SVector(0),SVector(-1)); @test integral(f,1,2) ≈ 2
    f = SchwarzChristoffelDerivate(SVector(1),SVector(-1)); @test integral(f,1,2) ≈ 0.5

    f = SchwarzChristoffelDerivate(SVector(0),SVector(0.5)); @test integral(f,1,1) ≈ 2
    f = SchwarzChristoffelDerivate(SVector(0),SVector(0.5)); @test integral(f,1,2) ≈ 2*sqrt(2)
    f = SchwarzChristoffelDerivate(SVector(1),SVector(0.5)); @test integral(f,1,2) ≈ 2

    f = SchwarzChristoffelDerivate(SVector(-1,0),SVector(-1,-1)); @test integral(f,2,1) ≈ 5/6
    f = SchwarzChristoffelDerivate(SVector(-1,0),SVector(-1,-1)); @test integral(f,2,2) ≈ 14/3

    f = SchwarzChristoffelDerivate(SVector(-1e-8,0),SVector(0.5,0)); @test integral(f,2,1) ≈ (sqrt(1e8+1)-1)/5000

    f = SchwarzChristoffelDerivate(SVector(0,1),SVector(0,0)); @test segment(f,1) ≈ 1
end


using SymbolicUtils, SpecialFunctions
using SymbolicUtils.Rewriters
using Symbolics

SEMIINTRULES = [
        # CONSTANTS AND POWERS
        @acrule sqrt(~x) => sqrt(pi)/2*~x
        #@acrule((~x) => 4*~x^(3/2)/(3*sqrt(pi)))
        @acrule (~x)^(~α) => gamma(~α+1)/gamma(~α+3/2)*(~x)^(~α+1/2)

        # BINOMIALS
        @acrule sqrt(1+~x) => sqrt(~x/pi) + ((1+~x)*atan(sqrt(~x)))/sqrt(pi)
        @acrule sqrt(1-~x) => sqrt(~x/pi) + ((1-~x)*atanh(sqrt(~x)))/sqrt(pi)
        @acrule 1/sqrt(1+~x) => 2/sqrt(pi)*atan(sqrt(~x))
        @acrule 1/sqrt(1-~x) => 2/sqrt(pi)*atanh(sqrt(~x))
        @acrule 1/(1+~x) => 2*asinh(sqrt(~x))/sqrt(pi*(1+~x))
        @acrule 1/(1-~x) => 2*asin(sqrt(~x))/sqrt(pi*(1-~x))
        @acrule 1/(1+~x)^(3/2) => 2/(1+~x)*sqrt(~x/pi)
        @acrule 1/(1-~x)^(3/2) => 2/(1-~x)*sqrt(~x/pi)
        @acrule 1/(1-~x)^(~p) => -incomplete_beta(~x, 1/2, ~p-1/2)/(sqrt(pi)*(1-~x)^(~p-1/2))
        @acrule 1/sqrt(~x*(1-~x)) => 2*ellipk(~x)/sqrt(pi)
        @acrule 1/sqrt(~x*(1+~x)) => 2*ellipk(~x/(1+~x))/sqrt(pi*(1+~x))
        @acrule sqrt(~x/(1-~x)) => 2/sqrt(pi)*(ellipk(~x)-ellipe(~x))
        @acrule 1/(sqrt(~x)*(1+~x)) => sqrt(pi/(1+~x))
        @acrule 1/(sqrt(~x)*(1-~x)) => sqrt(pi/(1-~x))
        @acrule sqrt(~x)/(~x) => -sqrt(pi/(1+~x))+sqrt(pi)
        @acrule sqrt(~x)/(-~x) => sqrt(pi/(1-~x))-sqrt(pi)
        @acrule sqrt(~x)/(1-~x)^2 => sqrt(pi)*~x/(2*(1-~x)^(3/2))
        @acrule sqrt((1-~x)/~x) => 2/sqrt(pi)*ellipe(~x)
        @acrule sqrt((1+~x)/~x) => 2*sqrt((1+~x)/pi)*ellipe(~x/(1+~x))
        @acrule  (~x)^(~p)/(1+~x)^~r => gamma(~p+1)/gamma(~p+3/2)*(~x)^(~p+1/2)* Hypergeometric1F1(~r, ~p+1, p+3/2, -~x)
        @acrule  (~x)^(~p)/(1-~x)^~r => gamma(~p+1)/gamma(~p+3/2)*(~x)^(~p+1/2)* Hypergeometric1F1(~r, ~p+1, p+3/2, ~x)

        # EXPONENTIAL AND RELATED FUNCTIONS
        @acrule exp(~x) => exp(~x)*erf(sqrt(~x))
        @acrule exp(-~x) => 2/sqrt(pi)*dawson(sqrt(~x))
        @acrule exp(~x)*erf(sqrt(~x)) => exp(~x)-1
        @acrule dawson(~x) => 1/2*sqrt(pi)*(1-exp(-~x))
        @acrule exp(~x)*erfc(sqrt(~x)) => 1-exp(~x)*erfc(sqrt(~x))
        @acrule exp(~x)*erfc(-sqrt(~x)) => exp(~x)*erfc(sqrt(~x))-1
        @acrule erf(sqrt(~x)) => ~x*exp(-1/2*~x)*(besseli(1, 1/2*~x)+besseli(0, 1/2*~x))
        @acrule exp(~x)/sqrt(~x) => sqrt(pi)*exp(1/2*~x)*besseli(0, 1/2*~x)
        @acrule exp(-~x)/sqrt(~x) => sqrt(pi)*exp(-1/2*~x)*besseli(0, 1/2*~x)
        @acrule cosh(sqrt(~x))/sqrt(~x) => sqrt(pi)*besseli(0, sqrt(~x))
        @acrule (1-cos(sqrt(~x)))/~x => sqrt(pi*(~x))*Struve(1, sqrt(~x))
        @acrule (1-cosh(sqrt(~x)))/~x => sqrt(pi*(~x))*MStruve(1, sqrt(~x))
        @acrule sin(~x) => sin(~x-1/4*pi)+sqrt(2)*AuxiliaryFresnelSin(sqrt(2*~x/pi))
        @acrule cos(~x) => cos(~x-1/4*pi)-sqrt(2)*AuxiliaryFresnelCos(sqrt(2*~x/pi))

        @acrule sinh(~x) => (exp(~x)*erf(sqrt(~x)))/2-dawson(sqrt(~x))/sqrt(pi)
        @acrule cosh(~x) => (exp(~x)*erf(sqrt(~x)))/2+dawson(sqrt(~x))/sqrt(pi)
        @acrule asin(sqrt(~x))/sqrt(1-~x) => -sqrt(pi)/2*log(1-~x)
        @acrule atan(sqrt(~x)) => sqrt(pi)*(sqrt(1+~x)-1)
        @acrule asinh(sqrt(~x))/sqrt(1+~x) => sqrt(pi)/2*log(1+~x)
        @acrule atanh(sqrt(~x)) => sqrt(pi)*(1-sqrt(1-~x))

        # BESSEL AND STRUVE FUNCTIONS
        @acrule besselj(0, sqrt(~x)) => sin(sqrt(~x))/sqrt(pi*(~x))
        @acrule besselj(1, sqrt(~x)) => 2*(1-cos(sqrt(~x)))/sqrt(pi*~x)
        @acrule besselj(~v, sqrt(~x))/(~x)^(~v/2) => (sqrt(2)*Struve(~v-1/2, sqrt(~x)))/(~x)^(~v/2-1/4)

        @acrule sqrt(~x)*besselj(1, sqrt(~x)) => (2*sin(sqrt(~x))-2*sqrt(~x)*cos(sqrt(~x)))/sqrt(pi)
        @acrule (~x)^(~v/2)*besselj(~v, sqrt(~x)) => (~x)^((~v)/2+1/4)*besselj(~v+1/2, sqrt(~x))
        @acrule besseli(0, sqrt(~x)) => sinh(sqrt(~x))/sqrt(pi*~x)
        @acrule besseli(1, sqrt(~x))/sqrt(~x) => 2*(1-cosh(sqrt(~x)))/sqrt(pi*~x)
        @acrule besseli(~v, sqrt(~x))/(~x)^(~v/2) => (sqrt(2)*Legendre(~v-1/2, sqrt(~x)))/(~x)^(~v/2-1/4)
        @acrule sqrt(~x)*besseli(1, sqrt(~x)) => (2*sqrt(~x)*cosh(sqrt(~x))-2*sinh(sqrt(~x)))/sqrt(pi)
        @acrule (~x)^(~v/2)*besseli(~v, sqrt(~x)) => sqrt(2)*(~x)^((~v)/2+1/4)*besseli(~v+1/2, sqrt(~x))
        @acrule exp(-~x)*besseli(0, ~x) => erf(sqrt(2*~x))/sqrt(2)
        @acrule Struve(0, sqrt(~x)) => 2*(1-cos(sqrt(~x)))/sqrt(pi)

        @acrule  Struve(0, sqrt(~x))/(~x) => 2*SinIntegral(sqrt(~x))/(sqrt(pi*(~x)))
        @acrule sqrt(~x)*Struve(1, sqrt(~x)) =>(~x+2-2*sqrt(~x)*sin(sqrt(~x))-2*cos(sqrt(~x)))/sqrt(pi)
        @acrule (~x)^(~v/2)*Struve(~v, sqrt(~x)) => sqrt(2)*(~x)^((~v)/2+1/4)*Struve(~v+1/2, sqrt(~x))

        @acrule MStruve(0, sqrt(~x)) => (2*(cosh(sqrt(~x))-1))/sqrt(pi)
        @acrule MStruve(0, sqrt(~x))/(~x) => (HyperSinIntegral(sqrt(~x)))/(sqrt(pi)*(~x))
        @acrule sqrt(~x)*MStruve(1, sqrt(~x)) => (2-~x+2*sqrt(~x)*sinh(sqrt(~x))-2*cosh(sqrt(~x)))/sqrt(pi)
        @acrule  (~x)^(~v/2)*MStruve(~v, sqrt(~x)) => sqrt(2)*(~x)^((~v)/2+1/4)*MStruve(~v+1/2, sqrt(~x))

        # MISCELLANEOUS FUNCTIONS
        @acrule log(~x) => 2*sqrt(~x/pi)*(log(4*~x-2))
        @acrule sqrt(~x)*log(~x) => (~x*sqrt(pi))/2*(log(~x/4)+1)
        @acrule log(~x)/sqrt(~x) => sqrt(pi)*log(~x/4)

        @acrule ellipk(~x) => sqrt(pi)*asin(sqrt(~x))
        @acrule ellipe(~x) => sqrt(pi*(~x)*(1-~x))/2 + sqrt(pi)*asin(sqrt(~x))/2
]

INTRULES = Chain(SEMIINTRULES)

"""
# Symbolic fractional integral

    semiint(fun)

```semiint``` uses SymbolicUtils.jl and Symbolics.jl to compute symbolic fractional integral.

### Example

```julia-repl
julia> using SymbolicUtils
julia> @syms x
julia> semiint(x^4)
0.45851597901024005(x^4.5)
```
"""
semiint(x) = INTRULES(x)


##############################################################################
##
## Combination behaves like [A B C ...] without forming it
##
##############################################################################

struct Combination{T} <: AbstractMatrix{T}
    A::Tuple
    cumlength::Vector{Int}
end

function Combination(A::Union{AbstractVector{T}, AbstractMatrix{T}}...) where {T}
    Combination{T}(A, cumsum([size(x, 2) for x in A]))
end

Base.size(c::Combination) = (size(c.A[1], 1), c.cumlength[end])
Base.size(c::Combination, i::Integer) = size(c)[i]

function Base.view(c::Combination, ::Colon, j)
    index = searchsortedfirst(c.cumlength, j)
    newj = index == 1 ? j : j - c.cumlength[index-1]
    view(c.A[index], :, newj)
end

##############################################################################
##
## Crossprod computes [A B C ...]' [A B C ...] without forming it
##
##############################################################################
crossprod(A::AbstractMatrix) = A'A
function crossprod(A::AbstractMatrix, B::AbstractMatrix)
    u11, u12, u22 = A'A, A'B, B'B
    hvcat(2, u11, u12,
             u12', u22)
end

function crossprod(A::AbstractMatrix, B::AbstractMatrix, C::AbstractMatrix)
    u11, u12, u13 = A'A, A'B, A'C
    u22, u23 = B'B, B'C
    u33 = C'C
    hvcat(3, u11,  u12,  u13,
             u12', u22,  u23,
             u13', u23', u33)
end

##############################################################################
##
## Returns base of [A B C ...]
##
## TODO: You could protect against roundoff error by using a controlled sum algorithm (similar to sum_kbn) to compute elements of X'X, then converting to BigFloat before factoring.
##
##
##############################################################################

# rank(A) == rank(A'A)
function basecol(X::AbstractMatrix...; factorization = :Cholesky)
    cholm = cholesky!(Symmetric(crossprod(X...)), Val(true); tol = -1, check = false)
    r = 0
    if size(cholm, 1) > 0
        r = sum(diag(cholm.factors) .> size(X[1],1)^2 * eps())
        # used to be r = rank(cholm) but does not work wiht very high regressors at the same time as intercept
    end
    invpermute!(1:size(cholm, 1) .<= r, cholm.piv)
end

function getcols(X::AbstractMatrix,  basecolX::AbstractVector)
    sum(basecolX) == size(X, 2) ? X : X[:, basecolX]
end


#-----------------------------------------------------------------------
"""
        rms( x::Vector )

Computes the `rms` value for a vectors.
"""
rms=x->sqrt(dot(x,x)/length(x))            # Root Mean Square calculator

#-----------------------------------------------------------------------
"""
        PrintErr(  h::RealVec , v::RealVec  , e::RealVec , thresh::Tuple , pf::String , units::String )

This function prints the errors, their rms and 95% CL values, and the
maximum error.

"""
function PrintErr(  h::RealVec , v::RealVec  , e::RealVec ,
                    thresh::Tuple , pf::String , units::String )
    hrms = rms(h)
    vrms = rms(v)
    erms = rms(e)

    n=length(h)

    h95 = h[Int(0.05*n)]
    v95 = v[Int(0.05*n)]
    e95 = e[Int(0.05*n)]

    hcthr = count(x->x>thresh[1], h)
    vcthr = count(x->x>thresh[2], v)
    ecthr = count(x->x>thresh[3], e)

    open(pf*"-res.txt","a") do io
        println(io,"hrms =","\t",hrms,"\t",
                "# Horizontal error rms (",units,")")
        println(io,"vrms =","\t",vrms,"\t",
                "# Vertical error rms (",units,")")
        println(io,"erms =","\t",erms,"\t",
                "# Full error rms (",units,")")
        println(io,"h95 =","\t",h95,"\t",
                "# Horizontal error 95% CL (",units,")")
        println(io,"v95 =","\t",v95,"\t",
                "# Vertical error 95% CL (",units,")")
        println(io,"e95 =","\t",e95,"\t",
                "# Full error 95% CL (",units,")")
        println(io,"hcthr =","\t",hcthr,"\t",
                "# Horizontal errors above threshold of ",thresh[1]," ",
                units)
        println(io,"vcthr =","\t",vcthr,"\t",
                "# Vertical errors above threshold of ",thresh[2]," ",
                units)
        println(io,"ecthr =","\t",ecthr,"\t",
                "# Full errors above threshold of ",thresh[3]," ",units)
        println(io,"hmax =","\t",h[1],"\t",
                "# Largest horizontal error (",units,")")
        println(io,"vmax =","\t",v[1],"\t",
                "# Largest vertical error (",units,")")
        println(io,"emax =","\t",e[1],"\t",
                "# Largest total error (",units,")")
    end

    return (hrms,vrms,erms,h95,v95,e95,hcthr,vcthr,ecthr)
end     #---------------------------------------------------------------

#-----------------------------------------------------------------------
"""
PlotErr(   err::RealVec  , lines::Tuple , histpar::Tuple , xr::Tuple , yr::Tuple , labels::Tuple , plotname::String )

This function plots the errors.

"""
function PlotErr(   err::RealVec  , lines::Tuple , histpar::Tuple ,
                    xr::Tuple , yr::Tuple , labels::Tuple ,
                    plotname::String )
    plot!(vline([lines[1],lines[2]]))
    histogram!(
        err ,
        alpha=0.50 ,
        bins=histpar[2]:histpar[1]:histpar[3],
        xlims=xr ,
        ylims=yr ,
        xlabel=labels[1] ,
        ylabel=labels[2] ,
        leg=false ,
        fontfamily="Times"
    )
    savefig(plotname)
end     #---------------------------------------------------------------


# ==============================================================================
#                        size of model
# ==============================================================================
nz = 101; nx = 641; nstep=4000;
# grid cell size
 deltaZ = 10.0        ;  deltaX = 10.0;
hdeltaZ = deltaZ / 2.0; hdeltaX = deltaX / 2.0;
# time step size
f0  = 7.0; dt = 0.001;# source parameter
# ==============================================================================
#                        elastic model parameters
# ==============================================================================
vp = 3300.0; vs = vp / 1.732; density = 2800.0;
courantNumber = vp * dt * sqrt(1.0/deltaZ^2 + 1.0/deltaX^2)
if courantNumber > 1.0
   println("Courant number is: $courantNumber")
   error("time step is too large, simulation will be unstable")
end

# arrays for elastic model parameter
lambda = zeros(nz+2, nx+2);
mu     = zeros(nz+2, nx+2);
rho    = zeros(nz+2, nx+2);

for j = 2 : nx+1
    for i = 2 : nz+1
           rho[i,j] = density
            mu[i,j] = density * vs * vs
        lambda[i,j] = density * (vp*vp - 2.0*vs*vs)
    end
end

# half grid elastic model parameters
hLambdaZ    = 0.0
hLambda2MuZ = 0.0
hMuZ        = 0.0
hMuX        = 0.0
hRhoZX      = 0.0

# arrays for wave field variables
vz = zeros(nz+2, nx+2);
vx = zeros(nz+2, nx+2);
sigmazz = zeros(nz+2, nx+2);
sigmaxx = zeros(nz+2, nx+2);
sigmazx = zeros(nz+2, nx+2);

# temperal variables to save spatial derivative
vdvzdz = 0.0
vdvxdx = 0.0
vdvxdz = 0.0
vdvzdx = 0.0
vdsigmazzdz = 0.0
vdsigmazxdx = 0.0
vdsigmaxxdx = 0.0
vdsigmazxdz = 0.0

# memory variables for CPML, it can be improved to save memory !!!
mdvzdz = zeros(nz+2, nx+2)
mdvxdx = zeros(nz+2, nx+2)
mdvxdz = zeros(nz+2, nx+2)
mdvzdx = zeros(nz+2, nx+2)
mdsigmazzdz = zeros(nz+2, nx+2)
mdsigmazxdx = zeros(nz+2, nx+2)
mdsigmaxxdx = zeros(nz+2, nx+2)
mdsigmazxdz = zeros(nz+2, nx+2)

# ==============================================================================
#                      prepare for CPML damping profile
# ==============================================================================
# thickness of CPML in number of grid points
nPointPml = 10;
thickPmlZ = nPointPml * deltaZ;
thickPmlX = nPointPml * deltaX;

# parameter for CPML profile
nPower   = 2.0
kmax     = 1.0
alphaMax = pi * f0

# damping profile, need to be improved  !!!
# (k, d, alpha, a, b) at grid point, # (hk, hd, halpha, ha, hb) at half grid point along Z direction
 kz=ones(nz);  dz=zeros(nz);  alphaz=zeros(nz);  az=zeros(nz);  bz=zeros(nz);
hkz=ones(nz); hdz=zeros(nz); halphaz=zeros(nz); haz=zeros(nz); hbz=zeros(nz);

# along X direction
 kx=ones(nx);  dx=zeros(nx);  alphax=zeros(nx);  ax=zeros(nx);  bx=zeros(nx);
hkx=ones(nx); hdx=zeros(nx); halphax=zeros(nx); hax=zeros(nx); hbx=zeros(nx);

# specify the obsorbing boundary, origin specify the computation domain, z
reflectCoef=1e-3;
d0Z = -(nPower + 1) * vp * log(reflectCoef) / (2.0 * thickPmlZ);
d0X = -(nPower + 1) * vp * log(reflectCoef) / (2.0 * thickPmlX);
locationZ=0.0; locationX=0.0; abscissa=0.0; normalizedAbscissa=0.0;
originTop = thickPmlZ; originBottom=(nz-1)*deltaZ - thickPmlZ;

# vertical damping profile
for i = 1 : nz
    locationZ = deltaZ * (i-1)

    # upper profile at grid point
    abscissa  = originTop - locationZ
    if abscissa > 0.0
       normalizedAbscissa = abscissa / thickPmlZ
       dz[i] = d0Z * normalizedAbscissa^nPower
       kz[i] = 1.0
       alphaz[i] = alphaMax * (1.0 - normalizedAbscissa) + 0.1 * alphaMax
    end

    # upper profile at half grid point
    abscissa = originTop - (locationZ + hdeltaZ)
    if abscissa > 0.0
       normalizedAbscissa = abscissa / thickPmlZ
       hdz[i] = d0Z * normalizedAbscissa^nPower
       hkz[i] = 1.0
       halphaz[i] = alphaMax * (1.0 - normalizedAbscissa) + 0.1 * alphaMax
    end

    # bottom profile at grid point
    abscissa = locationZ - originBottom
    if abscissa > 0.0
       normalizedAbscissa = abscissa / thickPmlZ
       dz[i] = d0Z * normalizedAbscissa^nPower
       kz[i] = 1.0
       alphaz[i] = alphaMax * (1.0 - normalizedAbscissa) + 0.1 * alphaMax
    end

    # bottom profile at half grid point
    abscissa = (locationZ + hdeltaZ) - originBottom
    if abscissa > 0.0
       normalizedAbscissa = abscissa / thickPmlZ
       hdz[i] = d0Z * normalizedAbscissa^nPower
       hkz[i] = 1.0
       halphaz[i] = alphaMax * (1.0 - normalizedAbscissa) + 0.1 * alphaMax
    end

    bz[i] = exp(- ( dz[i]/ kz[i] +  alphaz[i]) * dt)
   hbz[i] = exp(- (hdz[i]/hkz[i] + halphaz[i]) * dt)

    if (abs( dz[i]) > eps(Float32))
        az[i] =  dz[i] / ( kz[i]*( dz[i]+ kz[i]* alphaz[i])) * ( bz[i]-1.0)
    end
    if (abs(hdz[i]) > eps(Float32))
       haz[i] = hdz[i] / (hkz[i]*(hdz[i]+hkz[i]*halphaz[i])) * (hbz[i]-1.0)
    end
end

# damping in x direction
thickPmlX = nPointPml * deltaX; originLeft= thickPmlX; originRight =(nx-1)*deltaX - thickPmlX;
for j = 1 : nx
    locationX = deltaX * (j-1)
    # left damping profile at grid point
    abscissa = originLeft - locationX
    if (abscissa > 0.0)
       normalizedAbscissa = abscissa / thickPmlX
       dx[j] = d0X * normalizedAbscissa^nPower
       kx[j] = 1.0
       alphax[j] = alphaMax * (1.0 - normalizedAbscissa) + 0.1 * alphaMax
    end

    # left profile at half grid point
    abscissa = originLeft - (locationX + hdeltaX)
    if abscissa > 0.0
       normalizedAbscissa = abscissa / thickPmlX
       hdx[j] = d0X * normalizedAbscissa^nPower
       hkx[j] = 1.0
       halphax[j] = alphaMax * (1.0 - normalizedAbscissa) + 0.1 * alphaMax
    end

    # right damping profile at grid point
    abscissa = locationX - originRight
    if abscissa > 0.0
       normalizedAbscissa = abscissa / thickPmlX
       dx[j] = d0X * normalizedAbscissa^nPower
       kx[j] = 1.0
       alphax[j] = alphaMax * (1.0 - normalizedAbscissa) + 0.1 * alphaMax
    end
    # bottom profile at half grid point
    abscissa = (locationX + hdeltaX) - originRight
    if abscissa > 0.0
       normalizedAbscissa = abscissa / thickPmlX
       hdx[j] = d0Z * normalizedAbscissa^nPower
       hkx[j] = 1.0
       halphax[j] = alphaMax * (1.0 - normalizedAbscissa) + 0.1 * alphaMax
    end

    bx[j] = exp(- ( dx[j]/ kx[j] +  alphax[j]) * dt)
   hbx[j] = exp(- (hdx[j]/hkx[j] + halphax[j]) * dt)

    if (abs( dx[j]) > eps(Float32))
        ax[j] =  dx[j] / ( kx[j]*( dx[j]+ kx[j]* alphax[j])) * ( bx[j]-1.0)
    end
    if (abs(hdx[j]) > eps(Float32))
       hax[j] = hdx[j] / (hkx[j]*(hdx[j]+hkx[j]*halphax[j])) * (hbx[j]-1.0)
    end

end

isz = 20; isx = floor(Int64, nx/2)+1;
# ==============================================================================
#                             start computation
# ==============================================================================
@time ccall((:onetimestepforward_, "/Users/wenlei/Documents/ElasticWave/test/test.so"),
      Void, (Ptr{Int64}, Ptr{Int64}, Ptr{Int64},
             Ptr{Float64}, Ptr{Float64}, Ptr{Float64},
             Ptr{Array{Float64,2}}, Ptr{Array{Float64,2}}, Ptr{Array{Float64,2}},
             Ptr{Float64}, Ptr{Float64}, Ptr{Float64}, Ptr{Float64}, Ptr{Float64},
             Ptr{Array{Float64,2}}, Ptr{Array{Float64,2}}, Ptr{Array{Float64,2}}, Ptr{Array{Float64,2}}, Ptr{Array{Float64,2}},
             Ptr{Float64}, Ptr{Float64}, Ptr{Float64}, Ptr{Float64}, Ptr{Float64}, Ptr{Float64}, Ptr{Float64}, Ptr{Float64},
             Ptr{Array{Float64,2}}, Ptr{Array{Float64,2}}, Ptr{Array{Float64,2}}, Ptr{Array{Float64,2}}, Ptr{Array{Float64,2}}, Ptr{Array{Float64,2}}, Ptr{Array{Float64,2}}, Ptr{Array{Float64,2}},
             Ptr{Array{Float64,1}}, Ptr{Array{Float64,1}}, Ptr{Array{Float64,1}}, Ptr{Array{Float64,1}}, Ptr{Array{Float64,1}}, Ptr{Array{Float64,1}},
             Ptr{Array{Float64,1}}, Ptr{Array{Float64,1}}, Ptr{Array{Float64,1}}, Ptr{Array{Float64,1}}, Ptr{Array{Float64,1}}, Ptr{Array{Float64,1}}),
             &nz, &nx, &nstep,
             &deltaZ, &deltaX, &dt,
             lambda, mu, rho,
             &hLambdaZ, &hLambda2MuZ, &hMuZ, &hMuX, &hRhoZX,
             sigmazz, sigmaxx, sigmazx, vz, vx,
             &vdvzdz, &vdvxdx, &vdvxdz, &vdvzdx, &vdsigmazzdz, &vdsigmazxdx, &vdsigmaxxdx, &vdsigmazxdz,
              mdvzdz,  mdvxdx,  mdvxdz,  mdvzdx,  mdsigmazzdz,  mdsigmazxdx,  mdsigmaxxdx,  mdsigmazxdz,
             az, bz, kz, haz, hbz, hkz,
             ax, bx, kx, hax, hbx, hkx)

             m = 7; n=6;
             A = zeros(Int64, m+2, n+2);
             ccall((:assign_, "/Users/wenlei/Documents/ElasticWave/test/test.so"),
                   Void, (Ptr{Array{Int64,2}}, Ptr{Int64}, Ptr{Int64}),
                          A                  , &m        , &n        )

# crazy


using PMDPs, PMDPs.LP
using JuMP, GLPK

@testset "LP" begin

    pp = simple_pp()
    traces = [simple_short_trace(pp), empty_trace(pp),]
    exhaust = simple_exhaust_capacity_trace(pp)
    diff = simple_trace_different_utility_revenue_optim_allocs(pp)
    mg = PMDPs.PMDPg(pp)

    # revenue
    r, u, alloc, action_seq, requests = LP.MILP_hindsight_pricing(mg, traces[1])
    @test @ntuple(r,u,alloc,action_seq) == (r=30.,u=3, alloc=[1.,1.,1.],action_seq=[10.,10.,10.])

    r, u, alloc, action_seq, requests = LP.MILP_hindsight_pricing(mg, traces[2])
    @test @ntuple(r,u,alloc,action_seq) == (r=0.,u=0, alloc=[],action_seq=[])

    r, u, alloc, action_seq, requests = LP.MILP_hindsight_pricing(mg, exhaust)
    @test @ntuple(r,u,alloc,action_seq) == (r=55.,u=6, alloc=[1., 1., 1., 0., 0.,1., 1.],
                action_seq = [10.0, 10.0, 10.0, PMDPs.REJECT_ACTION, PMDPs.REJECT_ACTION, 15.0, 10.0])

    r, u, alloc, action_seq, requests = LP.MILP_hindsight_pricing(mg, diff)
    @test @ntuple(r,u,alloc,action_seq) == (r=30.,u=4, alloc=[1., 1., 1., 0., 0.],
                action_seq = [10.0, 10., 10., PMDPs.REJECT_ACTION,  PMDPs.REJECT_ACTION])

    # utilization
    obj_val, u, alloc, action_seq, requests = LP.MILP_hindsight_pricing(mg, traces[1]; objective=:utilization)
    @test @ntuple(obj_val,u,alloc,action_seq) == (obj_val=3.,u=3, alloc=[1.,1.,1.],action_seq=[10.,10.,10.])

    obj_val, u, alloc, action_seq, requests = LP.MILP_hindsight_pricing(mg, traces[2]; objective=:utilization)
    @test @ntuple(obj_val,u,alloc,action_seq) == (obj_val=0.,u=0, alloc=[],action_seq=[])

    obj_val, u, alloc, action_seq, requests = LP.MILP_hindsight_pricing(mg, exhaust; objective=:utilization)
    @test @ntuple(obj_val,u,alloc,action_seq) == (obj_val=6.,u=6, alloc=[1., 1., 1., 0., 0.,1., 1.],
                action_seq = [10.0, 10.0, 10.0, PMDPs.REJECT_ACTION, PMDPs.REJECT_ACTION, 15.0, 10.0])

    obj_val, u, alloc, action_seq, requests = LP.MILP_hindsight_pricing(mg, diff; objective=:utilization)
    @test @ntuple(obj_val,u,alloc,action_seq) == (obj_val=5., u=5, alloc=[1., 1., 0., 1., 0.,],
                action_seq = [10., 10., PMDPs.REJECT_ACTION, 10., PMDPs.REJECT_ACTION])
end


using LinearAlgebra
using StaticArrays
using PyPlot
const Vec2f = SVector{2, Float64}
const Vec4f = SVector{4, Float64}

@inline function bisection_newton(f, df, left::Float64, right::Float64, eps=0.05, itr_max=20)::Float64
    x_est = (left + right)*0.5
    for itr = 1:itr_max
        f_df = f(x_est)/df(x_est)
        if (right-(x_est-f_df))*((x_est-f_df)-left)<0.0 || abs(f_df)<(right-left)/4.0
            f(x_est)>0 ? right = x_est : left = x_est
            f_df = x_est - (right + left)*0.5
        end
        x_est -= f_df
        abs(f_df)<eps && break
    end
    return x_est
end

@inline function cost_optimal(s0::Vec4f, s1::Vec4f)
    @views a = Vec2f(s0[1:2])
    @views b = Vec2f(s0[3:4])
    @views c = Vec2f(s1[1:2])
    @views d =Vec2f(s1[3:4])
    return cost_optimal(a, b, c, d)
end

@inline function cost_optimal(x0::Vec2f, v0::Vec2f, x1::Vec2f, v1::Vec2f)
    x01 = x1 - x0
    v01 = v1 - v0
    p = -4*(dot(v0, v0)+dot(v1, v1)+dot(v0, v1))
    q = 24*dot(v0+v1, x01)
    r = -36*dot(x01, x01)
    cost(t) = t + dot(v01, 4.0*v01/t-6(-v0*t+x01)/t^2)+dot(-6*v01/t^2+12(x01-v0*t)/t^3, -v0*t+x01)
    @fastmath d_cost(t) = t^4+p*t*t+q*t+r # df(t)/dt
    @fastmath dd_cost(t) = 4.0*t*t*t+2*p*t+q # ddf(t)/dt^2

    t_min = 0.0
    t_max = 10.0
    t_star = bisection_newton(d_cost, dd_cost, t_min, t_max)
    return cost(t_star), t_star
end

# see ICRA paper: D.J.Webb et al Kinodynamic RRT* (2013)
function forward_reachable_box(x0::Vec2f, v0::Vec2f, r::Float64)
    @fastmath tau_x_plus = 2/3*(-(v0.^2).+r + v0.*sqrt.((v0.^2).+r))
    @fastmath tau_x_minus = 2/3*(-(v0.^2).+r - v0.*sqrt.((v0.^2).+r))
    @fastmath xmax = v0.*tau_x_plus + x0 + sqrt.(1/3*(tau_x_plus.^2).*(-tau_x_plus.+r))
    @fastmath xmin = v0.*tau_x_minus + x0 - sqrt.(1/3*(tau_x_minus.^2).*(-tau_x_minus.+r))

    @fastmath tau_v_plus = 0.5*r
    @fastmath vmax = v0 .+ sqrt.(tau_v_plus.*(-tau_v_plus.+r))
    @fastmath vmin = v0 .- sqrt.(tau_v_plus.*(-tau_v_plus.+r))
    return xmin, xmax, vmin, vmax
end

function backward_reachable_box(x0::Vec2f, v0::Vec2f, r::Float64)
    @fastmath tau_x_plus = 2/3*((v0.^2).-r + v0.*sqrt.((v0.^2).+r))
    @fastmath tau_x_minus = 2/3*((v0.^2).-r - v0.*sqrt.((v0.^2).+r))
    @fastmath xmax = v0.*tau_x_plus + x0 + sqrt.(1/3*(tau_x_plus.^2).*(tau_x_plus.+r))
    @fastmath xmin = v0.*tau_x_minus + x0 - sqrt.(1/3*(tau_x_minus.^2).*(tau_x_minus.+r))

    @fastmath tau_v_plus = 0.5*r
    @fastmath vmax = v0 .+ sqrt.(tau_v_plus.*(tau_v_plus.+r))
    @fastmath vmin = v0 .- sqrt.(tau_v_plus.*(tau_v_plus.+r))
    return xmin, xmax, vmin, vmax
end

"""
function filter_freachable_exact(s_set, s_c, r)
    s_set_filtered = Vec4f[]
    for s in s_set
        @views xq = Vec2f(s[1:2])
        @views vq = Vec2f(s[3:4])
        @views xc = Vec2f(s_c[1:2])
        @views vc = Vec2f(s_c[3:4])
        cost_optimal(xc, vc, xq, vq)<r && push!(s_set_filtered, s)
    end
    return s_set_filtered
end
"""

function filter_reachable(Sset::Vector{Vec4f}, idxset::Vector{Int64},
                          s_c::Vec4f, r::Float64, ForR::Symbol)
    s_set_filtered = Vec4f[]
    @views x_c = Vec2f(s_c[1:2])
    @views v_c = Vec2f(s_c[3:4])
    xmin, xmax, vmin, vmax = ((ForR==:F) ?
                              forward_reachable_box(x_c, v_c, r) :
                              backward_reachable_box(x_c, v_c, r))

    function isinside(s::Vec4f)
        @inbounds !(xmin[1]<s[1]<xmax[1]) && return false
        @inbounds !(xmin[2]<s[2]<xmax[2]) && return false
        @inbounds !(vmin[1]<s[3]<vmax[1]) && return false
        @inbounds !(vmin[2]<s[4]<vmax[2]) && return false
        return true
    end
    idx_filter = Int64[]
    dist_filter = Float64[]
    time_filter = Float64[]
    for idx in idxset
        @inbounds s = Sset[idx]
        if isinside(s)
            cost, tau = (ForR==:B) ? cost_optimal(s, s_c) : cost_optimal(s_c, s)
            if cost < r
                push!(idx_filter, idx)
                push!(dist_filter, cost)
                push!(time_filter, tau)
            end
        end
    end
    return idx_filter, dist_filter, time_filter
end

# see ICRA paper: D.J.Webb et al Kinodynamic RRT* (2013), Eq.(20)
function gen_trajectory(s0::Vec4f, s1::Vec4f, tau, N_split = 10)
    @views x0 = Vec2f(s0[1:2])
    @views v0 = Vec2f(s0[3:4])
    @views x1 = Vec2f(s1[1:2])
    @views v1 = Vec2f(s1[3:4])
    x01 = x1 - x0
    v01 = v1 - v0
    d = Vec4f([-6*v01/tau^2 + 12*(-tau*v0+x01)/tau^3;
               4*v01/tau - 6*(-tau*v0+x01)/tau^2])
    @inline function f(t)
        s = t - tau
        eye = Matrix{Int}(I, 2, 2)
        @fastmath M_left = vcat(hcat(eye, eye*s), hcat(eye*0, eye))
        @fastmath M_right = vcat(hcat(eye*(-s^3)*(1.0/6.0), eye*s^2*0.5),
                       hcat(eye*(-s^2*0.5), eye*s))
        return  Vec4f(M_left*s1 + M_right*d)
    end
    waypoints = Vec4f[]
    for n in 1:N_split+1
        t = (n-1)*tau/N_split
        push!(waypoints, f(t))
    end
    return waypoints
end

function show_trajectory(s0, s1, tau, N_split = 20, c_ = :gray, linewidth_ = 0.5)
    waypoints =  gen_trajectory(s0, s1, tau, N_split)
    M = zeros(4, N_split+1)
    for i in 1:N_split+1
        @inbounds M[:, i] = waypoints[i]
    end
    @views plot(M[1, :], M[2, :], c=c_, linewidth = linewidth_)
end

function test()
    N = 10^4
    s_set = Vec4f[]
    for i = 1:N
        push!(s_set, Vec4f(rand()-0.5, rand()-0.5, 3*rand()-0.5, 3*rand()-0.5))
    end
    s_c = Vec4f(+0, +0, 0.2, 0.2)
    idxset = [x for x in 1:N]
    idx_for, trash, _ = filter_reachable(s_set, idxset, s_c, 1.2, :F)
    idx_back, trash = filter_reachable(s_set, idxset, s_c, 1.2, :B)
    s_for = s_set[idx_for]
    s_back = s_set[idx_back]

    setA = zeros(2, length(idx_for))
    setB = zeros(2, length(idx_back))
    for i = 1:length(idx_for)
        setA[:, i] = s_for[i][1:2]
    end
    for i = 1:length(idx_back)
        setB[:, i] = s_back[i][1:2]
    end
    scatter(setA[1, :], setA[2, :], s=3)
    scatter(setB[1, :], setB[2, :], s=3)

    xlim(-0.6, 0.6)
    ylim(-0.6, 0.6)
    return setA
end

function test_connection()
    """
    s0 = Vec4f(0.17, 0.08, 0.19, -0.06)
    s1 = Vec4f(0.22, 0.03, 0.139, -0.2719)
    """
    s0 = Vec4f(0.0, 0.0, 0.3, 0.0)
    s1 = Vec4f(0.1, 0.1, -0.1, -0.01)
    show_trajectory(s0, s1, 0.5)
end
#test_connection()


mutable struct Translational{T,Nc} <: Joint{T,Nc}
    vertices::NTuple{2,SVector{3,T}}
    V12::SMatrix{2,3,T,6}
    V3::Adjoint{T,SVector{3,T}}
    cid::Int64

    function Translational{T,Nc}(body1::AbstractBody{T}, body2::AbstractBody{T};
        p1::AbstractVector{T} = zeros(3), p2::AbstractVector{T} = zeros(3), axis::AbstractVector{T} = zeros(3)) where {T,Nc}

        vertices = (p1, p2)
        if norm(axis) != 0
            axis = axis / norm(axis)
        end
        A = svd(skew(axis)).Vt # in frame of body1
        V12 = A[1:2,:]
        V3 = axis' # instead of A[3,:] for correct sign: abs(axis) = abs(A[3,:])
        cid = body2.id

        new{T,Nc}(vertices, V12, V3, cid), body1.id, body2.id
    end
end

Translational0 = Translational{T,0} where T
Translational1 = Translational{T,1} where T
Translational2 = Translational{T,2} where T
Translational3 = Translational{T,3} where T


function calc_P_W(reactor, R_0, B_0)
  nu_term = ( 1 + reactor.nu_n ) ^ 2 * ( 1 + reactor.nu_T ) ^ 2
  nu_term /= ( 1 + 2 * ( reactor.nu_n + reactor.nu_T ) )

  K_W = (
    2.824e-3 * nu_term * (reactor.f_D^2) * reactor.epsilon * reactor.kappa ^0.4
  )

  tmp_n_20 = calc_n_20(reactor, R_0, B_0)
  tmp_T_k = calc_T_k(reactor, R_0, B_0)

  K_W * tmp_n_20^2 * tmp_T_k^2 * R_0
end


function test_hailstone_nexthailstone()
    @test_throws DomainError nexthailstone(0)
    @test nexthailstone(1) == 4
    @test nexthailstone(2) == 1
end

function test_hailstone_hailstonelength()
    @test_throws DomainError hailstonelength(0)
    @test hailstonelength(1) == 1
    @test hailstonelength(13) == 10
end

function test_hailstone_allhailstone()
    @test_throws DomainError allhailstone(0)
    @test collect(allhailstone(1)) == [1]
    @test collect(allhailstone(13)) == [13, 40, 20, 10, 5, 16, 8, 4, 2, 1]
    @test eltype(allhailstone(1)) == typeof(1)
    @test eltype(allhailstone(13)) == typeof(13)
end

function test_hailstone_all()
    print(rpad("Math.NumberTheory.Hailstone...", 50, ' '))

    test_hailstone_nexthailstone()
    test_hailstone_hailstonelength()
    test_hailstone_allhailstone()

    println("PASS")
end


# ---
# title: 1220. Count Vowels Permutation
# id: problem1220
# author: zhwang
# date: 2022-01-22
# difficulty: Hard
# categories: Dynamic Programming
# link: <https://leetcode.com/problems/count-vowels-permutation/description/>
# hidden: true
# ---
#
# Given an integer `n`, your task is to count how many strings of length `n` can
# be formed under the following rules:
#
#   * Each character is a lower case vowel (`'a'`, `'e'`, `'i'`, `'o'`, `'u'`)
#   * Each vowel `'a'` may only be followed by an `'e'`.
#   * Each vowel `'e'` may only be followed by an `'a'` or an `'i'`.
#   * Each vowel `'i'` **may not** be followed by another `'i'`.
#   * Each vowel `'o'` may only be followed by an `'i'` or a `'u'`.
#   * Each vowel `'u'` may only be followed by an `'a'.`
#
# Since the answer may be too large, return it modulo `10^9 + 7.`
#
#
#
# **Example 1:**
#
#
#
#     Input: n = 1
#     Output: 5
#     Explanation: All possible strings are: "a", "e", "i" , "o" and "u".
#
#
# **Example 2:**
#
#
#
#     Input: n = 2
#     Output: 10
#     Explanation: All possible strings are: "ae", "ea", "ei", "ia", "ie", "io", "iu", "oi", "ou" and "ua".
#
#
# **Example 3:  **
#
#
#
#     Input: n = 5
#     Output: 68
#
#
#
# **Constraints:**
#
#   * `1 <= n <= 2 * 10^4`
#
#
## @lc code=start
using LeetCode

function count_vowel_permutation(n::Int)::Int
    table = Dict{Char,Int}(i => 1 for i in "aeiou")
    quotient = 10^9 + 7
    for _ in 1:(n - 1)
        a = table['e'] + table['i'] + table['u']
        e = table['a'] + table['i']
        i = table['e'] + table['o']
        o = table['i']
        u = table['i'] + table['o']
        table['a'], table['e'], table['i'], table['o'], table['u'] = map(
            x -> mod(x, quotient), [a, e, i, o, u]
        )
    end
    return mod(sum(values(table)), quotient)
end

## @lc code=end


using JeszenszkiBasis
using Test

@testset "JeszenszkiBasis" begin

@testset "Szbasis" begin
    @testset for (K, N, D) in [(50, 0, 1), (1, 1, 1), (2, 3, 4), (3, 2, 6), (4, 4, 35)]
        sb = Szbasis(K, N)

        @test sb.K == K
        @test sb.N == N
        @test sb.D == D
        @test num_vectors(N, K) == D
        @test size(sb.vectors) == (sb.K, sb.D)
        @test length(sb) == D
        @test site_max(sb) == N

        # Occupations are correct.
        @test all(sum(sb.vectors, dims=1) .== N)

        # Each vector is unique.
        @test sb.vectors == unique(sb.vectors, dims=2)

        # Iteration works, serial numbers match up.
        d = 0
        for (i, v) in enumerate(sb)
            @test v in sb
            @test v == sb.vectors[:, i]
            @test v == sb[i]
            @test i == serial_num(sb, v)
            d += 1
        end
        @test d == D

        # Invalid vectors.
        v = zeros(Int, K)
        if N > 0
            @test !(v in sb)
        else
            @test v in sb
        end
        for n in [-1, N+1]
            v[1] = n
            @test !(v in sb)
        end
    end

    @test_throws DomainError Szbasis(0, 5)
    @test_throws DomainError Szbasis(5, -1)
end

@testset "RestrictedSzbasis" begin
    @testset for (K, N, M, D) in [(50, 0, 0, 1), (1, 1, 1, 1), (2, 3, 2, 2), (3, 2, 1, 3), (4, 4, 2, 19)]
        sb = RestrictedSzbasis(K, N, M)

        @test sb.K == K
        @test sb.N == N
        @test sb.M == M
        @test sb.D == D
        @test num_vectors(N, K, M) == D
        @test size(sb.vectors) == (sb.K, sb.D)
        @test length(sb) == D
        @test site_max(sb) == M

        # Occupations are correct.
        @test all(sum(sb.vectors, dims=1) .== N)

        # Each vector is unique.
        @test sb.vectors == unique(sb.vectors, dims=2)

        # Iteration works, serial numbers match up.
        d = 0
        for (i, v) in enumerate(sb)
            @test v in sb
            @test v == sb.vectors[:, i]
            @test v == sb[i]
            @test i == serial_num(sb, v)
            d += 1
        end
        @test d == D

        # Invalid vectors.
        v = zeros(Int, K)
        if N > 0
            @test !(v in sb)
        else
            @test v in sb
        end
        for n in [-1, N+1, M+1]
            v[1] = n
            @test !(v in sb)
        end
    end

    @test_throws DomainError RestrictedSzbasis(2, 5, 2)
end

@testset "sub_serial_num" begin
    let sb = Szbasis(3, 2)
        # Split each vector into 2 sites on the left and 1 on the right.
        left = []
        right = []

        for v in sb
            # Try range indexing.
            push!(left, sub_serial_num(sb, v[1:2]))
            # Try a view.
            push!(right, sub_serial_num(sb, view(v, 3:3)))
        end

        @test sort(left) == [1, 2, 3, 4, 5, 6]
        @test sort(right) == [1, 1, 1, 2, 2, 3]
    end

    let sb = RestrictedSzbasis(5, 2, 2)
        # Split each vector into 2 sites on the left and 3 on the right.
        left = []
        right = []

        for v in sb
            # Try range indexing.
            push!(left, sub_serial_num(sb, v[1:2]))
            # Try a view.
            push!(right, sub_serial_num(sb, view(v, 3:5)))
        end

        @test sort(left) == [1, 1, 1, 1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 5, 6]
        @test sort(right) == [1, 1, 1, 2, 2, 3, 3, 4, 4, 5, 6, 7, 8, 9, 10]
    end
end

@testset "string" begin
    @testset for (s1, v, s2) in [(sz"", Int[], ""), (sz"0", [0], "0"), (sz"123", [123], "123"), (sz"0 0", [0, 0], "0 0"), (sz"1 23 456", [1, 23, 456], "1 23 456"), (sz"0 0 1 0", [0, 0, 1, 0], "0 0 1 0")]
        @test s1 == v
        @test s2 == to_str(v)
        @test s2 == to_str(v[1:end])
        @test s2 == to_str(@view v[1:end])
    end
end

end


# Line drawing algorithm

"""
    drawline(func::Function, ci1::CartesianIndex{2}, ci2::CartesianIndex{2}, eachstep::Bool=false)

At each step along the line from `ci1` to `ci2` call `func` with a single argument, the row, column coordinates
of the pixel as a `Tuple{Int, Int}`.

If `eachstep=true` then `func` is called once when the row changes and once when the column changes.  When
`eachstep=false` then `func` is called less frequently and both row and col may change between calls.
"""
function drawline(
    func::Function,
    ci1::CartesianIndex{2},
    ci2::CartesianIndex{2},
    eachstep::Bool = false,
)
    r, c = ci1.I
    dr, dc = -abs(ci2.I[1] - ci1.I[1]), abs(ci2.I[2] - ci1.I[2])
    sr, sc = ci1.I[1] < ci2.I[1] ? 1 : -1, ci1.I[2] < ci2.I[2] ? 1 : -1
    err = dc + dr
    func(ci1.I)
    while !((c == ci2.I[2]) && (r == ci2.I[1]))
        e2 = 2 * err
        if e2 >= dr
            err += dr
            c += sc
            if eachstep
                func((r, c))
            end
        end
        if e2 <= dc
            err += dc
            r += sr
            if eachstep
                func((r, c))
            end
        end
        if !eachstep
            func((r, c))
        end
    end
end

"""
    drawray(
        func::Function,  # returns Bool
        ci1::CartesianIndex{2}, # starting point
        ci2::CartesianIndex{2}, # towards this point
        bounds::CartesianIndices{2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}},  # stay within
        eachstep::Bool=false
    )
    function drawray(
        func::Function,
        ci1::CartesianIndex{2},
        ci2::CartesianIndex{2},
        eachstep::Bool = false,
    )

    Draw a ray starting at `ci1` towards `ci2` and call `func` with a single argument, the row, column coordinates
of the pixel as a `Tuple{Int, Int}`.  `func` must return a `Bool` - true to continue drawing the ray and false to
stop drawing the ray.  The first version checks the bounds each step and the second doesn't.

If `eachstep=true` then `func` is called once when the row changes and once when the column changes.  When
`eachstep=false` then `func` is called less frequently and both row and col may change between calls.
"""

function drawray(
    func::Function,
    ci1::CartesianIndex{2},
    ci2::CartesianIndex{2},
    bounds::CartesianIndices{2, Tuple{Vararg{OrdinalRange{Int64, Int64}, 2}}},
    eachstep::Bool = false,
)
    r, c = ci1.I
    dr, dc = -abs(ci2.I[1] - ci1.I[1]), abs(ci2.I[2] - ci1.I[2])
    sr, sc = ci1.I[1] < ci2.I[1] ? 1 : -1, ci1.I[2] < ci2.I[2] ? 1 : -1
    err = dc + dr
    inb(r,c) = ((r in bounds.indices[1]) && (c in bounds.indices[2]))
    if !(inb(r,c) && func((r,c)))
        return nothing
    end
    while true
        e2 = 2 * err
        if e2 >= dr
            err += dr
            c += sc
            if eachstep && !(inb(r,c) && func((r, c)))
                return nothing
            end
        end
        if e2 <= dc
            err += dc
            r += sr
            if eachstep && !(inb(r,c) && func((r, c)))
                return nothing
            end
        end
        if (!eachstep) && !(inb(r,c) && func((r, c)))
            return nothing
        end

    end
end

function drawray(
    func::Function,
    ci1::CartesianIndex{2},
    ci2::CartesianIndex{2},
    eachstep::Bool = false,
)
    r, c = ci1.I
    dr, dc = -abs(ci2.I[1] - ci1.I[1]), abs(ci2.I[2] - ci1.I[2])
    sr, sc = ci1.I[1] < ci2.I[1] ? 1 : -1, ci1.I[2] < ci2.I[2] ? 1 : -1
    err = dc + dr
    if !func((r,c))
        return nothing
    end
    while true
        e2 = 2 * err
        if e2 >= dr
            err += dr
            c += sc
            if eachstep && !func((r, c))
                return nothing
            end
        end
        if e2 <= dc
            err += dc
            r += sr
            if eachstep && !func((r, c))
                return nothing
            end
        end
        if (!eachstep) && !func((r, c))
            return nothing
        end

    end
end


export getInterpolation

function getInterpolation(AT::SparseMatrixCSC, S::SparseMatrixCSC, coloring::Array, n::Int64)
	S.nzval .= 1.0
    S .= AT .* S
	## Pp - P.colptr
	## P.j - P.rowval
	## Here we build PT.
	Pp = getInterpolation1(AT, S, coloring, n) ## Essentially counts non-zeros per row.
	Pp .= Pp .+ 1
	#(Pp, Pj, Px) = getInterpolation2(AT, S, Pp, coloring, n)
	(Pp, Pj, Px) = getDirectInterpolation2(AT, S, Pp, coloring, n)
	Pj .= Pj .+ 1
	R = SparseMatrixCSC(maximum(Pj), n, Pp, Pj, Px)
	P = sparse(R')
	return P, R
end

function getInterpolation1(AT::SparseMatrixCSC, S::SparseMatrixCSC, coloring::Array, n::Int64)
	Pp = zeros(Int, size(AT.colptr))
	nz = 0
	for i=1:n
		if (coloring[i] == 1)
			nz += 1
		else
			for j in nzrange(S, i)
				if S.rowval[j] != i && coloring[S.rowval[j]] == 1
					nz += 1
				end
			end
		end
		Pp[i+1] = nz
	end
	return Pp
end

# Direct interpolation.
# A little faster
# Discussed here : https://computation.llnl.gov/projects/hypre-scalable-linear-solvers-multigrid-methods/yang1.pdf
# and on "Multigrid" book.
# First one I implemented for testing. Implementation from PyAMG (originally in C++).

function getDirectInterpolation2(AT::SparseMatrixCSC, S::SparseMatrixCSC, Pp::Array, coloring::Array, n::Int64)
	Px = zeros(Float64, Pp[end] - 1)
	Pj = zeros(Int, Pp[end] - 1)

	for i = 1:n
		# Make coarse node the i row of identity
		if coloring[i] == 1
			Pj[Pp[i]] = i
			Px[Pp[i]] = 1
			continue
		end

		# get positive and negative sums of nodes that are strongly connected to i
		(sum_strong_pos, sum_strong_neg) = getStrongSum(S, coloring, i)
		# get sum a_ij for j in 1:n , split to positive and negative
		(sum_all_pos, sum_all_neg, a_ii) = getAllSum(AT, coloring, i)

		i_alpha = sum_all_neg / sum_strong_neg
		i_beta = sum_all_pos / sum_strong_pos

		if sum_strong_pos == 0
			a_ii += sum_all_pos
			i_beta = 0.0
		end

		neg = -1 * i_alpha / a_ii
		pos = -1 * i_beta / a_ii

		nz = Pp[i]

		for j in nzrange(S, i)
			if coloring[S.rowval[j]] == 1 && S.rowval[j] != i
				Pj[nz] = S.rowval[j]
				if S.nzval[j] > 0
					Px[nz] = pos * S.nzval[j]
				else
					Px[nz] = neg * S.nzval[j]
				end
				nz += 1
			end
		end
	end

	# for node i in F : Pi = sum(w_ij*e_j). Here we have the transpose
	sum_map = zeros(Int, n)
	p_sum = 0
	for i = 1:n
		sum_map[i] = p_sum
		p_sum += coloring[i]
	end
	Pj .= sum_map[Pj]

   return Pp, Pj, Px
end

# get sum(a_ii + sum(a_in)) for n in weakly connected set
function getDenominator(AT::SparseMatrixCSC, S::SparseMatrixCSC, i::Int64)
	denominator = 0
	for j in nzrange(AT, i)
		denominator += AT.nzval[j]
	end

	for j in nzrange(S,i)
		if S.rowval[j] != i
			denominator -= S.nzval[j]
		end
	end
	return denominator
end

# Sum for strong coarse neighbours, returns tow sums : one for positive elements, and the other for negative elements
function getStrongSum(S::SparseMatrixCSC, coloring::Array, i::Int64)
	sum_positive = 0.0
	sum_negative = 0.0
	for j in nzrange(S, i)
		if coloring[S.rowval[j]] == 1 && S.rowval[j] != i
			if S.nzval[j] > 0
				sum_positive += S.nzval[j]
			else
				sum_negative += S.nzval[j]
			end
		end
	end
	sum_positive, sum_negative
end

function getAllSum(AT::SparseMatrixCSC, coloring::Array, i::Int64)
	sum_all_pos = 0.0
	sum_all_neg = 0.0
	diag = 0.0
	for j in nzrange(AT, i)
		if AT.rowval[j] == i
			diag += AT.nzval[j]
		else
			if AT.nzval[j] < 0
				sum_all_neg += AT.nzval[j]
			else
				sum_all_pos += AT.nzval[j]
			end
		end
	end
	sum_all_pos, sum_all_neg, diag
end

function getInnerDenominator(S::SparseMatrixCSC, coloring::Array, i::Int64, m::Int64)
	inner_denominator = 0
	# inner_denominator is sum of a_mk such that:
	# k is in C_i (coarse strong connection) and m is in F_i (fine strong connection)
	for k in nzrange(S,S.rowval[m])
		for kk in nzrange(S,i)
			if S.rowval[k] == S.rowval[kk] && coloring[S.rowval[k]] == 1
				inner_denominator+=S.nzval[k]
			end
		end
	end
	return inner_denominator
end

## This is the textbook version, discussed in "A Multigrid Tutorial".
# We calculate the i component of the operator :
# For every coarse node, the row is unity (e_i)
# For every free node, we calculate the weight and put sum(w_ij * e_j)

function getInterpolation2(AT::SparseMatrixCSC, S::SparseMatrixCSC, Pp::Array, coloring::Array, n::Int64)

	Px = zeros(Float64, Pp[end] - 1)
	Pj = zeros(Int, Pp[end] - 1)

	for i=1:n
		# Make coarse node the i row of identity
		## coloring = 1: C node, 0: F node
		if coloring[i] == 1
			Pj[Pp[i]] = i
			Px[Pp[i]] = 1
			continue
		end
		nz = Pp[i]
		# denominator holds all a_ij for weakly connected i-j.
		# with sparse arrays it's easier to take all, and subtract strong connections.
		i_denominator = getDenominator(AT, S, i)

		assert(i_denominator != 0)
		#j_numerator : AT[j][i] + ( AT[m][i] * AT[j][m] /  ( AT[k][m] for k in nz(S,i) if coloring(k) == 1) for m in nz(S,i) if coloring (k == 0))

		for j in nzrange(S,i)
			if S.rowval[j] == i || coloring[S.rowval[j]] == 0
				continue
			end

			Pj[nz] = S.rowval[j]

			j_numerator = S.nzval[j]
			for m in nzrange(S,i)
				# Ignore strong coarse connections
				if S.rowval[m] == i || coloring[S.rowval[m]] == 1
					continue
				end
				# check that AT[m][i], AT[j][m] != 0, otherwise there's no contribution to the numerator.
				inner_denominator = getInnerDenominator(S, coloring, i, m)
				assert(inner_denominator != 0)
				for l in nzrange(AT, S.rowval[m])
					if AT.rowval[l] == S.rowval[j]
						j_numerator += (S.nzval[m] * AT.nzval[l]) / inner_denominator
					end
				end
			end
			Px[nz] = - (j_numerator / i_denominator)
			nz += 1
		end
	end

	# for node i in F : Pi = sum(w_ij*e_j). Here we have the transpose.
	sum_map = zeros(Int, n)
	p_sum = 0
	for i = 1:n
		sum_map[i] = p_sum
		p_sum += coloring[i]
	end
	Pj .= sum_map[Pj]
	return Pp, Pj, Px
end


#
# Algorithm of:
#
# NONMONOTONE SPECTRAL PROJECTED GRADIENT METHODS ON CONVEX SETS
# ERNESTO G. BIRGIN, JOSÉ MARIO MARTÍNEZ, AND MARCOS RAYDAN
# SIAM J. O. PTIM. Vol. 10, No. 4, pp. 1196-1211
#
# Originally implemented by J. M. Martínez (IMECC - UNICAMP)
# Translated to Julia by L. Martínez (IQ-UNICAMP)
#
"""
```
spgbox!(f, g!, x::AbstractVector{Real}; lower=..., upper=..., options...)`
```

```
spgbox!(f, g!, lower::AbstractVector{Real}, upper::AbstractVector{Real}, x::AbstractVector{Real}; options...)`
```

Minimizes function `f` starting from initial point `x`, given the function to compute the gradient, `g!`. `f` must be of the form `f(x)`, and `g!` of the form `g!(g,x)`, where `g` is the gradient vector to be modified. It modifies the `x` vector, which will contain the best solution found (see `spgbox` for a non-mutating alternative).

Optional lower and upper box bounds can be provided using optional arguments `lower` and `upper`, which can be provided as the fourth and fifth arguments or with keyword parameters.

Returns a structure of type `SPGBoxResult`, containing the best solution found in `x` and the final objective function in `f`.

# Examples
```jldocstest
julia> f(x) = x[1]^2 + x[2]^2

julia> function g!(g,x)
         g[1] = 2*x[1]
         g[2] = 2*x[2]
       end
```
## Without bounds

```jldocstest
julia> x = rand(2)

julia> spgbox!(f,g!,x)

 SPGBOX RESULT:

 Convergence achieved.

 Final objective function value = 0.0
 Best solution found = [ 0.0, 0.0]
 Projected gradient norm = 0.0

 Number of iterations = 2
 Number of function evaluations = 3
```

## With bounds

```jldocstest
julia> x = 2 .+ rand(2)

julia> spgbox!(f,g!,x,lower=[2.,-Inf])

 SPGBOX RESULT:

 Convergence achieved.

 Final objective function value = 4.0
 Best solution found = [ 2.0, 0.0]
 Projected gradient norm = 0.0

 Number of iterations = 2
 Number of function evaluations = 3
```
"""
function spgbox!(
  f::Function,
  g!::Function,
  lower::Union{Nothing,AbstractVector{Float64}},
  upper::Union{Nothing,AbstractVector{Float64}},
  x::AbstractVector{Float64};
  eps::Float64 = 1.e-5,
  nitmax::Int = 100,
  nfevalmax::Int = 1000,
  m::Int = 10,
  vaux::VAux = VAux(length(x),m),
  iprint::Int = 0,
  project_x0::Bool = true
)
  return spgbox!(
    f,g!,x,lower=lower,upper=upper,
    eps=eps, nitmax=nitmax, nfevalmax=nfevalmax, m=m,
    vaux=vaux, iprint=iprint, project_x0=project_x0
  )
end
#
# call with lower and upper as keyword parameters
#
function spgbox!(
  f::Function,
  g!::Function,
  x::AbstractVector{Float64};
  lower::Union{Nothing,AbstractVector{Float64}} = nothing,
  upper::Union{Nothing,AbstractVector{Float64}} = nothing,
  eps::Float64 = 1.e-5,
  nitmax::Int = 100,
  nfevalmax::Int = 1000,
  m::Int = 10,
  vaux::VAux = VAux(length(x),m),
  iprint::Int = 0,
  project_x0::Bool = true
)

  # Number of variables
  n = length(x)

  # Auxiliary arrays (associate names and check dimensions)
  g = vaux.g
  xn = vaux.xn
  gn = vaux.gn
  fprev = vaux.fprev
  @assert length(g) == n "Auxiliar gradient vector `g` must be of the same length as `x`"
  @assert length(xn) == n "Auxiliar vector `xn` must be of the same length as `x`"
  @assert length(gn) == n "Auxiliar vector `gn` must be of the same length as `x`"
  @assert length(fprev) == m "Auxiliar vector `fprev` must be of length `m`"

  # Check if bounds are defined, project or not the initial point on them
  if !isnothing(lower)
    @assert length(lower) == n "Lower bound vector `lower` must be of the same length than x, got: $(length(lower))"
    if project_x0
      @. x = max(x,lower)
    else
      for i in eachindex(x)
        if x[i] < lower[i]
          error(" Initial value of variable $i smaller than lower bound, and `project_x0` is set to `false`. ")
        end
      end
    end
  end
  if !isnothing(upper)
    @assert length(upper) == n "Upper bound vector `upper` must be of the same length than `x`, got: $(length(upper))"
    if project_x0
      @. x = min(x,upper)
    else
      for i in eachindex(x)
        if x[i] > lower[i]
          error(" Initial value of variable $i greater than upper bound, and `project_x0` is set to `false`. ")
        end
      end
    end
  end

  # Objective function at initial point
  nfeval = 1
  fcurrent = f(x)
  g!(g,x)
  gnorm = pr_gradnorm(g,x,lower,upper)

  tspg = 1.
  for i in eachindex(fprev)
    fprev[i] = fcurrent
  end

  # Iteration counter
  nit = 0
  while nit < nitmax

    if iprint > 0
      println("----------------------------------------------------------- ")
      println(" Iteration: ", nit )
      println(" x = ", x[begin], " ... ", x[end] )
      println(" Objective function value = ", fcurrent)
    end

    # Compute gradient norm
    gnorm = pr_gradnorm(g,x,lower,upper)

    if iprint > 0
      println(" ")
      println(" Norm of the projected gradient = ", gnorm)
      println(" Number of function evaluations = ", nfeval)
    end

    # Stopping criteria
    if gnorm <= eps
      ierr = 0
      return SPGBoxResult(x,fcurrent,gnorm,nit,nfeval,ierr)
    end
    if nfeval >= nfevalmax
      ierr = 2
      return SPGBoxResult(x,fcurrent,gnorm,nit,nfeval,ierr)
    end

    t = tspg
    fref = maximum(fprev)

    if iprint > 2
      println(" fref = ", fref)
      println(" fprev = ", fprev)
      println(" t = ", t)
    end

    fn = +Inf
    while( fn > fref )
      for i in eachindex(x)
        xn[i] = x[i] - t*g[i]
        if !isnothing(upper)
          xn[i] = min(xn[i],upper[i])
        end
        if !isnothing(lower)
          xn[i] = max(xn[i],lower[i])
        end
      end
      if iprint > 2
        println(" xn = ", xn[begin], " ... ", xn[end])
      end
      nfeval = nfeval + 1
      fn = f(xn)
      if iprint > 2
        println(" f[end] = ", fn, " fref = ", fref )
      end
      # Maximum number of function evaluations achieved
      if nfeval > nfevalmax
        ierr = 2
        return SPGBoxResult(x,fcurrent,gnorm,nit,nfeval,ierr)
      end
      # Reduce region
      t = t/2
    end

    g!(gn,xn)
    num = 0.
    den = 0.
    for i in eachindex(x)
      num = num + (xn[i]-x[i])^2
      den = den + (xn[i]-x[i])*(gn[i]-g[i])
    end
    if den <= 0.
      tspg = 100.
    else
      tspg =  min(1.e3,num/den)
    end
    fcurrent = fn
    for i in eachindex(x)
      x[i] = xn[i]
      g[i] = gn[i]
    end
    for i in firstindex(fprev):lastindex(fprev)-1
      fprev[i] = fprev[i+1]
    end
    fprev[end] = fcurrent
    nit = nit + 1
  end

  # Maximum number of iterations achieved
  ierr = 1
  return SPGBoxResult(x,fcurrent,gnorm,nit,nfeval,ierr)

end

"""

```
spgbox(f, g!, x::AbstractVector{Real}; lower=..., upper=..., options...)`
```

```
spgbox(f, g!, lower::AbstractVector{Real}, upper::AbstractVector{Real}, x::AbstractVector{Real}; options...)`
```

Minimizes function `f` starting from initial point `x`, given the function to compute the gradient, `g!`. `f` must be of the form `f(x)`, and `g!` of the form `g!(g,x)`, where `g` is the gradient vector to be modified. This call *does not* mutate the `x` vector, instead it will create a (deep)copy of it (see `spgbox!` to the inplace alternative).

Optional lower and upper box bounds can be provided using optional arguments `lower` and `upper`, which can be provided as the fourth and fifth arguments or with keyword parameters.

Returns a structure of type `SPGBoxResult`, containing the best solution found in `x` and the final objective function in `f`.

"""
function spgbox(
  f::Function,
  g!::Function,
  lower::Union{Nothing,AbstractVector{Float64}},
  upper::Union{Nothing,AbstractVector{Float64}},
  x::AbstractVector{Float64};
  eps::Float64 = 1.e-5,
  nitmax::Int = 100,
  nfevalmax::Int = 1000,
  m::Int = 10,
  vaux::VAux = VAux(length(x),m),
  iprint::Int = 0,
  project_x0::Bool = true
)
  x0 = deepcopy(x)
  return spgbox!(
    f,g!,x0,lower=lower,upper=upper,
    eps=eps, nitmax=nitmax, nfevalmax=nfevalmax, m=m,
    vaux=vaux, iprint=iprint, project_x0=project_x0
  )
end
#
# call with lower and upper as keyword parameters
#
function spgbox(
  f::Function,
  g!::Function,
  x::AbstractVector{Float64};
  lower::Union{Nothing,AbstractVector{Float64}} = nothing,
  upper::Union{Nothing,AbstractVector{Float64}} = nothing,
  eps::Float64 = 1.e-5,
  nitmax::Int = 100,
  nfevalmax::Int = 1000,
  m::Int = 10,
  vaux::VAux = VAux(length(x),m),
  iprint::Int = 0,
  project_x0::Bool = true
)
  x0 = deepcopy(x)
  return spgbox!(
    f,g!,x0,lower=lower,upper=upper,
    eps=eps, nitmax=nitmax, nfevalmax=nfevalmax, m=m,
    vaux=vaux, iprint=iprint, project_x0=project_x0
  )
end


"""
    gradient_optimizer(obj_fun, x0; kwargs...)

Performs a very rudimentary version of gradient descent/ascent. Follows basic
pattern of NLopt package: `obj_fun` is a function which accepts arguments `x`, `g`; it
returns the objective function at `x` and sets the gradient `g` in place. The
value `x0` is an initial guess. Default behavior is to descend along the
gradient to minimize the function, setting keyword argument `objective` to
`:max` changes this behavior (default is `:min`). Returns a tuple
`(F_opt, x_opt, stop)` where `stop` is an array of the stopping criteria met.

List of keywords and default values:
 * `objective = :min` minimize `obj_fun`, setting to `:max` will maximize
 * `verbose = 0` Options are 0 (no output), 1 (some output), 2 (all the output)
 * `print_eval = 100` In maximum verbosity, waits this many evaluations before printing
 * `lr = 1.0` A "learning rate", the gradient is scaled by this much.
 * `adjust_lr = lr > 1` If true, every time a smaller value of `obj_fun` is reached, `lr` is halved.

Stopping criteria (except vanishing gradient) can be set with keyword arguments.
 * `maxeval = 1000` stop after evaluating `obj_fun` this many times
 * `xtol_abs = 0.0` stop when `|x_cur - x_prev| < xtol_abs`
 * `xtol_rel = 0.0` stop when `|x_cur - x_prev|/|x_cur| < xtol_rel`
 * `ftol_abs = 0.0` stop when `|F_cur - F_prev| < ftol_abs`
 * `ftol_rel = 0.0` stop when `|F_cur - F_prev|/|F_cur| < ftol_rel`

"""
function gradient_optimizer(obj_fun, x0; objective=:min,
    verbose::Int=0, print_eval::Int=100,
    lr::Float64=1.0, adjust_lr::Bool=(lr > 1.0),
    maxeval::Int=1000,
    xtol_abs::Float64=0.0, xtol_rel::Float64=0.0,
    ftol_abs::Float64=0.0, ftol_rel::Float64=0.0,
    kwargs...)

    # println("The type of obj_fun is $(typeof(obj_fun))")
    # obj_fun = (x,g) -> obj_fun(x,g)
    # println("Now the type of obj_fun is $(typeof(obj_fun))")
    # bookkeeping set up
    n_vars = length(x0)
    dkwargs = Dict(kwargs)
    obj_sign = (objective == :min ? -1 : 1)
    stopping_criteria = [:maxeval_reached,
                        :opt_reached,
                        :xtol_abs_reached,
                        :xtol_rel_reached,
                        :ftol_abs_reached,
                        :ftol_rel_reached]
    stop = falses(stopping_criteria)

    # gradient algorithm set up
    g_cur = zeros(x0); g_prev = zeros(x0); g_opt = zeros(x0)
    x_cur = zeros(x0); x_prev = zeros(x0); x_opt = zeros(x0)
    F_cur = 0.0; F_prev = 0.0; F_opt = -obj_sign * Inf;
    dF_cur = 0.0; dx_cur = 0.0; dx_next = 0.0
    eval = 0; eval_opt = 0
    x_prev[:] = x0[:]

    whoopscount = 0

    if verbose > 0
        println("gradient_optimizer: $(obj_sign > 0 ? "maximizing" : "minimizing") function of $n_vars variable$(n_vars > 1 ? "s" : "")")
    end
    if verbose > 1
        println("\tlr           : $lr ($(adjust_lr ? "will adjust" : "no adjust"))")
        println("\tmaxeval      : $maxeval")
        println("\txtol abs/rel : $xtol_abs / $xtol_rel")
        println("\tftol abs/rel : $ftol_abs / $ftol_rel")
        println()
        println("\teval\tF(x_cur)\t\tdF\t|dx_next|\tlr")
    end
    while stop.chunks[1] == 0
        # get current value of x_cur from x_prev and gradient
        x_cur[:] = x_prev[:] + obj_sign * lr * g_prev[:]
        # evaluate at new x_cur
        F_cur = obj_fun(x_cur, g_cur)
        dF_cur = F_cur - F_prev
        dx_next = lr * norm(g_cur) # size of next step

        # update local extremum
        if F_cur * obj_sign > F_opt * obj_sign
            # if we're more extreme than current best, update
            g_opt[:] = g_cur[:]
            x_opt[:] = x_cur[:]
            F_opt = F_cur
            eval_opt = eval
        end

        # print if necessary
        if verbose > 1 && mod(eval, print_eval) == 0
            println("\t$eval\t$F_cur\t$dF_cur\t$dx_next\t$lr$(eval_opt == eval ? " *" : "")")
        end

        # update the learning rate if we overshot
        if adjust_lr && (dF_cur * obj_sign < 0)
            whoopscount += 1
            lr /= 2
        end

        # increment counter; eval is now the index of the next step.
        eval += 1

        #push cur into prev slot
        x_prev[:] = x_cur[:]
        g_prev[:] = g_cur[:]
        F_prev = F_cur

        # Check stopping conditions
        # 1 maxeval
        # 2 opt reached (dF = 0)
        # 3 xtol_abs reached (lr * norm(g_cur) < xtol_abs)
        # 4 xtol_rel reached (lr * norm(g_cur) / norm(x_cur) < xtol_rel)
        # 5 ftol_abs reached (dF_cur < ftol_abs)
        # 6 ftol_rel reached (|dF_cur / F_cur| < ftol_rel )
        Fstep_rel = F_cur == 0 ? Inf : abs(dF_cur / F_cur)
        xstep_rel = norm(x_cur) == 0 ? Inf : abs(dx_next / norm(x_cur))
        stop = stop | [eval > maxeval,
                        dF_cur == 0.0,
                        dx_next < xtol_abs,
                        xstep_rel < xtol_rel,
                        abs(dF_cur) < ftol_abs,
                        Fstep_rel < ftol_rel]
    end
    if verbose > 0
        println("gradient_optimizer: Stopping criteria reached: $(stopping_criteria[stop])")
        println("gradient_optimizer: optimal value found: $F_opt")
    end
    if verbose > 1
        println("gradient_optimizer: opt. value reached at evaluation $eval_opt / $maxeval")
        println("gradient_optimizer: Learning rate adjusted $whoopscount time$(whoopscount != 1 ? "s" : "")")
        println("gradient_optimizer: final learning rate: $lr")
    end

    return F_opt, x_opt, stopping_criteria[stop]
end

"""
    brute_force_optimizer(obj_fun, d; objective=:min, kwargs...)

Perform a brute force search over a `d`-dimensional cube for the optimal value
of `obj_fun`. Returns `F_opt, x_opt, [:brute_force_search]` to mirror output
from `gradient_optimizer`

Some kwargs that can be passed:
 * `ranges::Vector{Iterable}` A vector of `d` iterables explicitly defining the search space
 * `save_fun::Array{Float64,d}` a matrix that will be modified in-place with the function values over the whole cube. Size must match the `length` of each iterable in `ranges`

WARNING: Currently only accepts functions which take exactly 3 parameters,
because that's all I need it for. And let's be real, you shouldn't be doing this
for large `d` anyway. (The further reason I am not generalizing is that I'll be
using this on a machine running Julia v0.5, so no Itertools package, and don't
feel like writing my own product iterator.)

Current functionality is limited: Must pass `ranges = [iter1, iter2, iter3]` as
a kwarg, this will define the cube in which search takes place. Each `iter` can
be any iterator that is indexed by a range of the form `1:N` and returns `Real`
numbers (and reports `length(iter) = N`)

"""
function brute_force_optimizer(obj_fun, d::Integer; verbose=0,
    objective=:min, kwargs...)
    if d == 3
        return brute_force_optimizer3(obj_fun, d; verbose=verbose, objective=objective, kwargs...)
    else
        error("brute_force_optimizer: Can't actually handle dimensionality other than 3 right now, sorry.")
        #MAYBEDO expand to any dimensionality
    end
end
function brute_force_optimizer3(obj_fun, d; verbose=0,
    objective=:min, save_fun::Array{Float64}=[], kwargs...)

    d = 3

    dkwargs = Dict(kwargs)
    obj_sign = (objective == :min) ? -1 : 1
    # lb = length(lb) < d ? lb * ones(d) : [lb...]
    # ub = lengtH(ub) < d ? ub * ones(d) : [ub...]

    #MAYBEDO implement lots of options to allow automated creation of ranges
    #from, e.g. lb and ub.

    # for now, make life easy and assume explicit ranges have been passed as
    # `ranges` kwarg.
    x_opt = zeros(d)
    F_opt = -obj_sign * Inf
    F_cur = 0.0
    ranges = pop!(dkwargs, :ranges)
    if verbose > 0
        println("brute_force_optimizer: $(obj_sign > 0 ? "maximizing" : "minimizing") function of $d variables")
        for k = 1:d
            println("\trange $k: $(length(ranges[k])) options")
        end
        if length(save_fun) > 0
            println("\tWriting function values to matrix")
        end
    end
    for (i1,x1) in enumerate(ranges[1])
        for (i2,x2) in enumerate(ranges[2])
            for (i3,x3) in enumerate(ranges[3])
                # score[i,j,k] = obj_fun([x1,x2,x3])
                F_cur = obj_fun([x1,x2,x3])
                if F_cur * obj_sign > F_opt * obj_sign
                    F_opt = F_cur
                    x_opt = [x1,x2,x3]
                end
                if length(save_fun) > 0
                    save_fun[i1,i2,i3] = F_cur
                end
            end
        end
    end
    return F_opt, x_opt, [:brute_force_search]
end