"""
    RelativeEntropyBridge{T}

The `RelativeEntropyCone` is representable with exponential cone and LP
constraints, since ``u \\ge \\sum_{i=1}^n w_i \\log (\\frac{w_i}{v_i})`` if and
only if there exists a vector ``y`` such that ``u \\ge \\sum_i y_i`` and
``y_i \\ge w_i \\log (\\frac{w_i}{v_i})`` or equivalently
``v_i \\ge w_i \\exp (\\frac{-y_i}{w_i})`` or equivalently
``(-y_i, w_i, v_i) \\in ExponentialCone``, for all ``i``.
"""
struct RelativeEntropyBridge{T,F,G,H} <: AbstractBridge
    y::Vector{MOI.VariableIndex}
    ge_index::MOI.ConstraintIndex{F,MOI.GreaterThan{T}}
    exp_indices::Vector{MOI.ConstraintIndex{G,MOI.ExponentialCone}}
end

function bridge_constraint(
    ::Type{RelativeEntropyBridge{T,F,G,H}},
    model::MOI.ModelLike,
    f::H,
    s::MOI.RelativeEntropyCone,
) where {T,F,G,H}
    f_scalars = MOI.Utilities.eachscalar(f)
    d = MOI.dimension(s)
    v_dim = div(d - 1, 2)
    y = MOI.add_variables(model, v_dim)
    rhs = MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(one(T), y), zero(T))
    ge_index = MOI.Utilities.normalize_and_add_constraint(
        model,
        MOI.Utilities.operate(-, T, f_scalars[1], rhs),
        MOI.GreaterThan(zero(T)),
        allow_modify_function = true,
    )
    w_start = 1 + v_dim
    exp_funcs = [
        MOI.Utilities.operate(
            vcat,
            T,
            MOI.Utilities.operate(-, T, y[i]),
            f_scalars[w_start+i],
            f_scalars[1+i],
        ) for i in 1:v_dim
    ]
    exp_indices = [
        MOI.add_constraint(model, exp_func_i, MOI.ExponentialCone()) for
        exp_func_i in exp_funcs
    ]
    return RelativeEntropyBridge{T,F,G,H}(y, ge_index, exp_indices)
end

function MOI.supports_constraint(
    ::Type{RelativeEntropyBridge{T}},
    ::Type{<:MOI.AbstractVectorFunction},
    ::Type{MOI.RelativeEntropyCone},
) where {T}
    return true
end

function MOI.Bridges.added_constrained_variable_types(
    ::Type{<:RelativeEntropyBridge},
)
    return Tuple{Type}[]
end

function MOI.Bridges.added_constraint_types(
    ::Type{RelativeEntropyBridge{T,F,G,H}},
) where {T,F,G,H}
    return Tuple{Type,Type}[(F, MOI.GreaterThan{T}), (G, MOI.ExponentialCone)]
end

function concrete_bridge_type(
    ::Type{<:RelativeEntropyBridge{T}},
    H::Type{<:MOI.AbstractVectorFunction},
    ::Type{MOI.RelativeEntropyCone},
) where {T}
    S = MOI.Utilities.scalar_type(H)
    F = MOI.Utilities.promote_operation(-, T, S, S)
    G = MOI.Utilities.promote_operation(
        vcat,
        T,
        MOI.Utilities.promote_operation(-, T, MOI.VariableIndex),
        S,
    )
    return RelativeEntropyBridge{T,F,G,H}
end

function MOI.get(bridge::RelativeEntropyBridge, ::MOI.NumberOfVariables)::Int64
    return length(bridge.y)
end

function MOI.get(bridge::RelativeEntropyBridge, ::MOI.ListOfVariableIndices)
    return copy(bridge.y)
end

function MOI.get(
    ::RelativeEntropyBridge{T,F},
    ::MOI.NumberOfConstraints{F,MOI.GreaterThan{T}},
)::Int64 where {T,F}
    return 1
end

function MOI.get(
    bridge::RelativeEntropyBridge{T,F,G},
    ::MOI.NumberOfConstraints{G,MOI.ExponentialCone},
)::Int64 where {T,F,G}
    return length(bridge.y)
end

function MOI.get(
    bridge::RelativeEntropyBridge{T,F},
    ::MOI.ListOfConstraintIndices{F,MOI.GreaterThan{T}},
) where {T,F}
    return [bridge.ge_index]
end

function MOI.get(
    bridge::RelativeEntropyBridge{T,F,G},
    ::MOI.ListOfConstraintIndices{G,MOI.ExponentialCone},
) where {T,F,G}
    return copy(bridge.exp_indices)
end

function MOI.delete(model::MOI.ModelLike, bridge::RelativeEntropyBridge)
    for exp_index_i in bridge.exp_indices
        MOI.delete(model, exp_index_i)
    end
    MOI.delete(model, bridge.ge_index)
    MOI.delete(model, bridge.y)
    return
end

function MOI.get(
    model::MOI.ModelLike,
    ::MOI.ConstraintFunction,
    bridge::RelativeEntropyBridge{T,F,G,H},
) where {T,F,G,H}
    func = MOI.Utilities.zero_with_output_dimension(G, 1 + 2 * length(bridge.y))
    MOI.Utilities.operate_output_index!(
        +,
        T,
        1,
        func,
        MOI.get(model, MOI.ConstraintFunction(), bridge.ge_index),
    )
    w_start = 1 + length(bridge.y)
    for i in eachindex(bridge.y)
        exp_func_i = MOI.Utilities.eachscalar(
            MOI.get(model, MOI.ConstraintFunction(), bridge.exp_indices[i]),
        )
        MOI.Utilities.operate_output_index!(-, T, 1, func, exp_func_i[1])
        MOI.Utilities.operate_output_index!(+, T, 1 + i, func, exp_func_i[3])
        MOI.Utilities.operate_output_index!(
            +,
            T,
            w_start + i,
            func,
            exp_func_i[2],
        )
    end
    return MOI.Utilities.convert_approx(
        H,
        MOI.Utilities.remove_variable(func, bridge.y),
    )
end

function MOI.get(
    ::MOI.ModelLike,
    ::MOI.ConstraintSet,
    bridge::RelativeEntropyBridge,
)
    return MOI.RelativeEntropyCone(1 + 2 * length(bridge.y))
end

function MOI.supports(
    model::MOI.ModelLike,
    attr::Union{MOI.ConstraintPrimalStart,MOI.ConstraintDualStart},
    ::Type{RelativeEntropyBridge{T,F,G,H}},
) where {T,F,G,H}
    ci_1 = MOI.ConstraintIndex{F,MOI.GreaterThan{T}}
    ci_2 = MOI.ConstraintIndex{G,MOI.ExponentialCone}
    return MOI.supports(model, attr, ci_1) && MOI.supports(model, attr, ci_2)
end

function MOI.get(
    model::MOI.ModelLike,
    attr::Union{MOI.ConstraintPrimal,MOI.ConstraintPrimalStart},
    bridge::RelativeEntropyBridge{T},
) where {T}
    primal = zeros(T, 1 + 2 * length(bridge.y))
    primal[1] = MOI.get(model, attr, bridge.ge_index)
    w_start = 1 + length(bridge.y)
    for i in eachindex(bridge.y)
        exp_primal_i = MOI.get(model, attr, bridge.exp_indices[i])
        primal[1] -= exp_primal_i[1]
        primal[1+i] = exp_primal_i[3]
        primal[w_start+i] = exp_primal_i[2]
    end
    return primal
end

function MOI.set(
    model::MOI.ModelLike,
    attr::MOI.ConstraintPrimalStart,
    bridge::RelativeEntropyBridge{T},
    value,
) where {T}
    v_dim = length(bridge.y)
    v_value = value[2:(v_dim+1)]
    w_value = value[(v_dim+2):end]
    y_value = [w_i * log(w_i / v_i) for (v_i, w_i) in zip(v_value, w_value)]
    MOI.set(
        model,
        attr,
        bridge.ge_index,
        value[1] - reduce(+, y_value, init = zero(T)),
    )
    for i in 1:v_dim
        MOI.set(
            model,
            attr,
            bridge.exp_indices[i],
            [-y_value[i], w_value[i], v_value[i]],
        )
        MOI.set(model, MOI.VariablePrimalStart(), bridge.y[i], y_value[i])
    end
    return
end

# Given a is dual on u - sum(y) >= 0 and (b_i, c_i, d_i) is dual on
# (-y_i, w_i, v_i) in ExponentialCone, the dual on
# (u, v, w) in RelativeEntropyCone is (a, d_i, c_i).
function MOI.get(
    model::MOI.ModelLike,
    attr::Union{MOI.ConstraintDual,MOI.ConstraintDualStart},
    bridge::RelativeEntropyBridge{T},
) where {T}
    dual = zeros(T, 1 + 2 * length(bridge.y))
    dual[1] = MOI.get(model, attr, bridge.ge_index)[1]
    w_start = 1 + length(bridge.y)
    for i in eachindex(bridge.y)
        exp_dual_i = MOI.get(model, attr, bridge.exp_indices[i])
        dual[1+i] = exp_dual_i[3]
        dual[w_start+i] = exp_dual_i[2]
    end
    return dual
end

# Given constraint dual start of (u, v, w), constraint dual on GreaterThan is u
# and on exponential cone constraint i is (r_i, w_i, v_i), but since y_i is
# free, its dual is 0, so we have -r_i + u == 0 hence r_i = u.
# Note: alternatively, we could use the Lambert W function to calculate
# r_i = exp(W(-w_i / (-v_i * e))) * (-v_i * e), but this is more complicated.
function MOI.set(
    model::MOI.ModelLike,
    ::MOI.ConstraintDualStart,
    bridge::RelativeEntropyBridge,
    value,
)
    u = value[1]
    MOI.set(model, MOI.ConstraintDualStart(), bridge.ge_index, u)
    w_start = 1 + length(bridge.y)
    for i in eachindex(bridge.y)
        MOI.set(
            model,
            MOI.ConstraintDualStart(),
            bridge.exp_indices[i],
            [u, value[w_start+i], value[1+i]],
        )
    end
    return
end


# PSO

# Import Packages

using Interpolations
using PyPlot

pygui(true)

# Form Camber

x = y = 0: 10: 100
z = rand(length(x), length(y)) .* 10
# fig1 = surf(x, y, z)
xq = yq = 0: 100
interp_cubic = CubicSplineInterpolation((x, y), z)
zq = interp_cubic(xq, yq)
fig1 = surf(xq, yq, zq)

# Create Pops

N = 100
iter = 1000
ω = 0.8
c1 = 1
c2 = 1
poplimit = [x[1], x[end]]
vlimit = [-3, 3]

pop = zeros(N, 3)
v = zeros(N, 2)
localBest = zeros(N, 3)
globalBest = zeros(N, 3)

for ii = 1: N

    global pop, v, localBest

    pop[ii, 1: 2] = rand(1, 2) * 100
    pop[ii, 3] = interp_cubic(pop[ii, 1], pop[ii, 2])
    localBest[ii, :] = pop[ii, :]
    v[ii, :] = rand(1, 2) * 2 .* vlimit[1] .+ vlimit[2]

end

scatter3D(pop[:, 1], pop[:, 2], pop[:, 3])

for jj = 1: 2

    for kk = 1: N

        global pop, v, localBest, globalBest

        # Update Pops

        popnew = zeros(1, 3)
        popnew[1: 2] = pop[kk, 1: 2] .+ v[kk, :]

        if (poplimit[1] ≤ popnew[1, 1] ≤ poplimit[2]) && (poplimit[1] ≤ popnew[1, 2] ≤ poplimit[2])

            popnew[3] = interp_cubic(popnew[1], popnew[2])
            pop[kk, :] = popnew[1, :]

            if pop[kk, 3] ≥ localBest[kk, 3]

                localBest[kk, :] = pop[kk, :]

                if localBest[kk, 3] ≥ globalBest[3]

                    globalBest = localBest[kk, :]

                end

            end

        else

            popnew[3] = 0
            pop[kk, :] = popnew[1, :]

        end

        # Update v

        vnew = zeros(1, 2)
        vrand = rand(2, 2) .* 2 .- 1
        vchange = [c1 * (globalBest[1] - pop[kk, 1]) c2 * (localBest[kk, 1] - pop[kk, 1]); c1 * (globalBest[2] - pop[kk, 2]) c2 * (localBest[kk, 2] - pop[kk, 2])]
        vnew = ω * v[kk, :] .+ vrand .* vchange

        if (vlimit[1] ≤ vnew[1, 1] ≤ vlimit[2]) && (vlimit[1] ≤ vnew[1, 2] ≤ vlimit[2])

            v[kk, :] = vnew[1, :]

        end

    end

end

scatter3D(globalBest[1], globalBest[2], globalBest[3])

println("max is at $(globalBest)")


__precompile__(true)

module IntervalContractors

export plus_rev, minus_rev, inv_rev,
       mul_rev, div_rev, power_rev,
       sign_rev, max_rev, min_rev,
        sqr_rev, sqrt_rev, abs_rev,
        exp_rev, exp2_rev, exp10_rev, expm1_rev,
        log_rev, log2_rev, log10_rev, log1p_rev,
        sin_rev, cos_rev, tan_rev,
        asin_rev, acos_rev, atan_rev,
        sinh_rev, cosh_rev, tanh_rev,
        asinh_rev, acosh_rev, atanh_rev,
        mul_rev_IEEE1788

using IntervalArithmetic

const half_pi = IntervalArithmetic.half_pi(Float64)  # interval
const two_pi = IntervalArithmetic.two_pi(Float64)  # interval

#
# Base.:∪(f::Function, g::Function) = X -> ( f(X) ∪ g(X) )
# Base.:∩(f::Function, g::Function) = X -> ( f(X) ∩ g(X) )  # or f(g(X)) for contractors

include("arithmetic.jl")
include("transformations.jl")
include("powers.jl")
include("exponential.jl")
include("trig.jl")
include("inverse_trig.jl")
include("hyperbolic.jl")
include("inverse_hyperbolic.jl")
include("extrema.jl")

"""
Dictionary mapping functions to their reverse functions.
"""
const reverse_operations = Dict(
                    :+     => :plus_rev,
                    :-     => :minus_rev,
                    :*     => :mul_rev,
                    :/     => :div_rev,
                    :^     => :power_rev,
                    :()    => :()   # empty operation
                    )

for f in (  :sqrt, :sqr, :abs,
            :exp, :exp2, :exp10, :expm1,
            :log, :log2, :log10, :log1p,
            :sin, :cos, :tan,
            :asin, :acos, :atan,
            :sinh, :cosh, :tanh,
            :asinh, :acosh, :atanh,
            :inv, :sign, :max, :min)
    reverse_operations[f] = Symbol(f, "_rev")
end

end


module TensorValues

using StaticArrays
using Base: @propagate_inbounds, @pure

export MultiValue
export TensorValue
export VectorValue

export inner, outer, meas
export det, inv, tr, dot, norm
export mutable
export trace
export symmetic_part

import Base: show
import Base: zero, one
import Base: +, -, *, /, \, ==, ≈
import Base: conj
import Base: sum, maximum, minimum
import Base: getindex, iterate, eachindex
import Base: size, length, eltype
import Base: reinterpret
import Base: convert
import Base: CartesianIndices
import Base: LinearIndices
import Base: adjoint

import LinearAlgebra: det, inv, tr, dot, norm

include("Types.jl")

include("Indexing.jl")

include("Operations.jl")

include("Reinterpret.jl")

end # module


module HierarchialPerformanceTest

using HypothesisTests
using Statistics

const PerfMatrix = Array{Float64,2}

function hptRankSumTest(x::AbstractVector{S}, y::AbstractVector{T}, num_runs::Int) where {S<:Real,T<:Real}
    # Custom threshold as specified by the paper
    if num_runs < 12
        ExactMannWhitneyUTest(x, y)
    else
        ApproximateMannWhitneyUTest(x, y)
    end
end

function hptSignedRankTest(x::AbstractVector{T}, num_tests::Int) where T<:Real
    # Custom threshold as specified by the paper
    if num_tests < 25
        ExactSignedRankTest(x)
    else
        ApproximateSignedRankTest(x)
    end
end

function hpt(a::PerfMatrix, b::PerfMatrix; tail=:right)
    # Enforce identical and sane matrix sizes
    if size(a) != size(b)
        error("mismatching number of benchmarks/runs")
    elseif length(a) == 0
        error("no benchmarks given")
    end

    num_tests = size(a, 1) # n
    num_runs = size(a, 2) # m

    # Max significance for the *null* hypothesis, not the alternative
    maxNullSignificance = if (num_runs >= 5) 0.05 else 0.10 end

    # Calculate performance delta per benchmark
    test_deltas = Array{Float64}(undef, num_tests) # dT
    for test = 1:num_tests
        aRuns = a[test,:]
        bRuns = b[test,:]

        # Use SHT to determine which side is better
        shtP = pvalue(hptRankSumTest(aRuns, bRuns, num_runs), tail=tail)

        # Check is inverted because shtP = probability of null hypothesis (equal)
        test_deltas[test] = if (shtP < maxNullSignificance)
            median(aRuns) - median(bRuns)
        else
            0
        end
    end

    # Evaluate general performance using the selected SHT
    1 - pvalue(hptSignedRankTest(test_deltas, num_tests), tail=tail)
end

function hptSpeedup(a::PerfMatrix, b::PerfMatrix, minConfidence::Float64 = 0.95, digits::Int = 3, increment::Int = 1)
    # Calculate an extra digit internally to facilitate rounding
    baseSpeedup = 10 ^ digits
    # Use int for storing speedup to mitigate floating-point precision loss
    intSpeedup = baseSpeedup
    speedup = 1.0
    tail = :right

    # Invert if slower
    if hpt(a, b) < minConfidence
        increment *= -1
        tail = :left
    end

    while true
        confidence = hpt(a / speedup, b, tail=tail)
        if confidence >= minConfidence
            # Increment int and recalculate float to mitigate floating-point precision loss
            intSpeedup += increment
            speedup = intSpeedup / baseSpeedup
        else
            break
        end
    end

    speedup
end

function calcScore(results::PerfMatrix, reference::PerfMatrix, refScore::Int = 1000)
    # 95%-confidence speedup
    # We swap a and b because HPT considers higher values to be better, but
    # since our values are times rather than performance, lower is better
    round(Int, hptSpeedup(reference, results) * refScore)
end

export PerfMatrix, hpt, hptSpeedup, calcScore

end


using Revise
using AdFem
using PyPlot
using SparseArrays

cmesh = CrackMesh(50,50,1/50,24)
mmesh = cmesh.mesh
β = 1/4; γ = 1/2
a = b = 0.1

NT = 500
Δt = 2/NT

Ca = spzeros(cmesh.nslip, 2mmesh.ndof)
ls = zeros(cmesh.nslip)
for i = 1:cmesh.nslip
    u = cmesh.upper[i+1]
    l = cmesh.lower[i+1]
    Ca[i, u] = 1.0
    Ca[i, l] = -1.0
    ls[i] = cmesh.h/5
end
Ca = constant(Ca)

bdedge = cmesh.bdedge

ρ = ones(get_ngauss(mmesh))

M = compute_fem_mass_matrix1(ρ, mmesh)
Z = spzero(mmesh.ndof, mmesh.ndof)
M = [M Z;Z M]

ν = 0.1ones(get_ngauss(mmesh))
E = ones(get_ngauss(mmesh))

D = compute_plane_strain_matrix(E, ν)
K = compute_fem_stiffness_matrix(D, mmesh)

Nbd = size(bdedge, 1)
ρe = ones(4Nbd)
mu = 1/(2*(1+0.3))
lambda = 0.3/(1+0.3)/(1-2*0.3)
cs = ones(4Nbd) * sqrt(mu)
cp = ones(4Nbd) * sqrt(lambda + 2mu)
C = compute_absorbing_boundary_condition_matrix(ρe, cs, cp, bdedge, mmesh)

function ricker(;dt = 0.05, f0 = 5.0)
    nw = 2/(f0*dt)
    nc = floor(Int, nw/2)
    t = dt*collect(-nc:1:nc)
    b = (π*f0*t).^2
    w = @. (1 - 2b)*exp(-b)
end

F = zeros(NT, 2mmesh.ndof+cmesh.nslip)
r = ricker()
xy = gauss_nodes(mmesh)
S = @. exp(-20*((xy[:,1]-0.5)^2+(xy[:,2]-0.5)^2))

for i = 1:min(length(r), NT)
    f1 = S * r[i]
    F[i,:] = [compute_fem_source_term(zero(f1), zero(f1), mmesh);ls]
end

F = constant(F)

en = ExplicitNewmark(M, C, K, Δt, factorize_A = false)
en.A = factorize([en.A Ca'
        Ca spzero(cmesh.nslip)])
en.C = [en.C spzero(2mmesh.ndof, cmesh.nslip)
        spzero(cmesh.nslip, 2mmesh.ndof) spzero(cmesh.nslip)]
en.B = [en.B spzero(2mmesh.ndof, cmesh.nslip)
    spzero(cmesh.nslip, 2mmesh.ndof) spzero(cmesh.nslip)]

u = TensorArray(NT+1)
u = write(u, 1, zeros(2mmesh.ndof+cmesh.nslip))
u = write(u, 2, zeros(2mmesh.ndof+cmesh.nslip))

function condition(i, arr...)
    i <= NT+1
end
function body(i, arr...)
    u = arr[1]
    u0 = read(u, i-2)
    u1 = read(u, i-1)
    u = write(u, i, step(en, u0, u1, F[i-1]))
    i+1, u
end
i = constant(3, dtype = Int32)
_, out = while_loop(condition, body, [i, u])
U = stack(out)

sess = Session(); init(sess)
uarr = run(sess, U)

m = mean(uarr[:, 1:mmesh.ndof])
s = std(uarr[:, 1:mmesh.ndof])
vmin = -5s
vmax = 5s
function update(i)
    clf()
    visualize_scalar_on_fem_points(uarr[i,1:mmesh.ndof], mmesh, vmin = vmin, vmax = vmax)
    title("i = $i")
end
p = animate(update, 1:(NT÷20):NT)
saveanim(p, "udir.gif")

m = mean(uarr[:, mmesh.ndof+1:2mmesh.ndof])
s = std(uarr[:, mmesh.ndof+1:2mmesh.ndof])
vmin = m - 2s
vmax = m + 2s
function update(i)
    clf()
    visualize_scalar_on_fem_points(uarr[i,mmesh.ndof+1:2mmesh.ndof], mmesh, vmin = vmin, vmax = vmax)
    title("i = $i")
end
p = animate(update, 1:(NT÷20):NT)
saveanim(p, "vdir.gif")


function interpolate(c::EllipticalArc, p0, t)
    ϕ, rx, ry = c.angle, c.r1, c.r2
    θ = c.a1 + (c.a2 - c.a1) * t
    m = Mat2(cos(ϕ), sin(ϕ), -sin(ϕ), cos(ϕ))
    m * Point(rx * cos(θ), ry * sin(θ)) + c.c
end

function interpolate(c::CurveTo, p0, t)
    p1, p2, p3 = c.c1, c.c2, c.p
    (1 - t)^3 * p0 + 3(t - 2t^2 + t^3) * p1 + 3(t^2 -t^3) * p2 + t^3 * p3
end

discretize!(v::Vector{<:AbstractPoint}, c::Union{MoveTo, LineTo}) = push!(v, c.p)
function discretize!(v::Vector{<:AbstractPoint}, c::ClosePath)
    if v[end] !== v[begin]
        push!(v, v[begin])
    end
end

function discretize!(v::Vector{<:AbstractPoint}, c::Union{EllipticalArc, CurveTo})
    N0 = length(v)
    p0 = v[end]
    N = 100
    resize!(v, N0 + N)
    dt = 1.0/N
    for (i, t) in enumerate(dt:dt:1.0)
        v[N0 + i] = interpolate(c, p0, t)
    end
end

export discretize
function discretize(path::BezierPath)
    T = typeof(path.commands[1].p) #XXX BezierPath{T} ?
    v = Vector{T}()
    for c in path.commands
        discretize!(v, c)
    end
    return v
end

#=
using LinearAlgebra: norm, cross, ⋅, normalize

"""continous parametric from of bezier path"""
function (path::BezierPath)(t)
    N = length(path.commands) - 1
    t = t * N
    seg = round(Int, t % N) + 1
    tseg = rem(t, N)
    p0 = endpos(path.commands[seg]) # get last point from prev
    c = path.commands[seg+1]
    if c isa ClosePath
        return interpolate(LineTo(endpos(path.commands[1])), p0, tseg)
    else
        return interpolate(path.commands[seg+1], p0, tseg)
    end
end

endpos(c::MoveTo) = c.p
endpos(c::LineTo) = c.p
endpos(c::CurveTo) = c.p
endpos(c::EllipticalArc) = c(zero(Point2), 1)

interpolate(c::MoveTo, p0, t) = endpos(c)
interpolate(c::LineTo, p0, t) = p0 + (c.p - p0) * t

# experiments on adaptive discretization
function tangent(c::CurveTo, p0, t)
    p1, p2, p3 = c.c1, c.c2, c.p
    normalize(-3(1 - t)^2 * p0 + 3(1 - 4t + 3t^2) * p1 + 3(2t -3t^2) * p2 + 3t^2 * p3)
end

function discretize_adaptive!(v::Vector{<:AbstractPoint}, c::Union{EllipticalArc, CurveTo})
    α = 2 / 360 * 2π
    dtstep = 0.01

    p0 = v[end]

    t = 0.0
    while true
        oldT = tangent(c, p0, t)
        T = oldT
        steps = 0
        while abs(oldT ⋅ T) > cos(α)
            t += dtstep
            steps += 1
            T = tangent(c, p0, t)
        end
        if steps == 1 # if a single step was to much than adjust stepsize
            dtstep = dtstep/2
            continue
        end
        t >= 1.0 && break

        t -= dtstep
        push!(v, interpolate(c, p0, t))
    end
    push!(v, interpolate(c, p0, 1.0))
end
=#


function targeted_predictors(y::Array{Float64,1}, w::Matrix{Float64}, x::Matrix{Float64}; thresholding::String="hard")
    # Bai and Ng (2008) -> regress y on w and x and keep only the predictors x which are significant
    # this functions returns their column numbers
    # note that w should maybe consist of a constant (vector of 1s) and lags of y
    #predictors = DataFrame(hcat(y, w, x))  # unfortunately using formulae and the GLM package is a bit clumsy here
    #lm(Formula(convert(Symbol, "x1"), convert(Symbol, join(names(predictors[2:end]), "+"))), predictors)
    #lm(*("x1~", join(names(predictors[2:end]), "+")), predictors)

    design_matrix = hcat(w, x)
    if thresholding == "hard"  # TODO: hard thresholding does not work properly, only admits one variable plus the white covariance has negative diagonal elements
        coefficients = inv(design_matrix'*design_matrix)*(design_matrix'*y)  # OLS estimate
        residuals = y - design_matrix * coefficients
        # calcualte variance-covariance matrix according to White(1980) TODO: argument to choose if robust or not
        coefficient_covariance = inv(design_matrix'design_matrix)*(design_matrix'*diagm(residuals.^2)*design_matrix)*inv(design_matrix'design_matrix)
        num_negative = sum(diag(coefficient_covariance) .< 0)
        if num_negative > 0
            for i in 1:size(coefficient_covariance, 1)
                if coefficient_covariance[i, i] < 0
                   coefficient_covariance[i, i] = -coefficient_covariance[i, i]
               end
            end
            t_stat = coefficients./sqrt(diag(coefficient_covariance))
            println("there were $num_negative negative coefficient variances which have been set positive. THIS IS UNACCEPTABLE!")
        else
            t_stat = coefficients./sqrt(diag(coefficient_covariance))
        end
        critical_value = quantile(TDist(length(y)-length(coefficients)), 0.975)
        println(sum(abs(t_stat[size(w, 2)+1:end]) .> critical_value))
        return abs(t_stat[size(w, 2)+1:end]) .> critical_value  # return significant t-stats associated with x variables only (Bai, Ng, 2008)
    end
    if thresholding == "soft"  # lasso coefficients which are non 0 are used
        #res=lars(design_matrix, y, intercept=true, standardize=true, use_gram=true, maxiter=typemax(Int), lambda_min=0.0, verbose=true)
        res = GLMNet.glmnetcv(design_matrix, y)  # TODO: GLMNet should not be necessary here...
        betas = res.path.betas[:, indmin(res.meanloss)]
        return abs(betas[size(w, 2)+1:end]) .> 0
    end
end


using ExtremeLearning

# Testing a simple binary adder

const nhidden = 100

x = [1.0 1.0; 0.0 1.0; 0.0 0.0; 1.0 0.0]
y = [0.0 1.0; 1.0 0.0; 0.0 0.0; 1.0 0.0]

elm = ELM(nhidden,x,y)
fit!(elm, x, y)

new_x = [1.0 1.0; 0.0 1.0; 0.0 0.0; 1.0 0.0]
y_pred = predict(elm, new_x)

@test y_pred[1, 1] < 0.1
@test y_pred[1, 2] > 0.9
@test y_pred[2, 1] > 0.9
@test y_pred[2, 2] < 0.1
@test y_pred[3, 1] < 0.1
@test y_pred[3, 2] < 0.1
@test y_pred[4, 1] > 0.9
@test y_pred[4, 2] < 0.1


struct FreeFlow <: AbstractFlowDynamics end

function collision_time(b::MovableBall, Π::FixedPlane, ::FreeFlow)
    @unpack r, x, v = b
    @unpack n, p = Π

    if v ⋅ n < 0  # travelling in wrong direction
        return Inf
    end

    return (-r - (x - p)⋅n) / (v ⋅ n)
end

function collision_time(b1::MovableBall, b2::MovableBall, ::FreeFlow)
    Δx = b1.x - b2.x
    Δv = b1.v - b2.v

    b = Δx⋅Δv

    b > 0 && return Inf  # moving away from one another

    a = normsq(Δv)
    c = normsq(Δx) - (b1.r + b2.r)^2

    discriminant = b^2 - a*c

    if discriminant ≥ 0
        d = √discriminant

        # t1 = (-b + d) / a
        t2 = (-b - d) / a

        if t2 > 0
            return t2
        end
    end

    return Inf   # no collision
end

#
# function find_collision(b::MovableBall{N,T}, box::RectangularBox{N,T}, ::FreeFlow) where {N,T}
#
#     walls = box.walls
#
#     tol = 1e-10
#
#     min_collision_time = Inf
#     which = -1
#
#     for i in 1:length(walls)
#         t = collision_time(b, walls[i])
#         if tol < t < min_collision_time
#             min_collision_time = t
#             which = i
#         end
#     end
#
#     return (which, min_collision_time)
# end

"Flow ball for a time t"
function flow!(b::MovableBall, t, ::FreeFlow)
    b.x += b.v * t
end


function compute_offset(p::Integer, segment_lo::Integer)
    p² = p * p
    if p² >= segment_lo
        byte = cld(p², 30)

        # Wheel index stores both the index of the prime number mod 30
        # and the index of the active multiple. We start crossing off
        # p * p, so that would be to_idx(p % 30) twice. We combine those values
        # as a number between 1 ... 64.
        wheel = to_idx(p % 30)
        wheel_index = 8 * (wheel - 1) + wheel

        return byte, wheel_index
    else
        # p * q will be the first number to cross off
        q = cld(segment_lo, p)
        q_quot, q_rem = divrem(q, 30)

        remainders = (1, 7, 11, 13, 17, 19, 23, 29, 31)
        i = 1
        while remainders[i] < q_rem
            i += 1
        end

        # maybe wrap around
        q_rem = i == 9 ? 1 : remainders[i]
        q_quot = i == 9 ? q_quot + 1 : q_quot
        i = i == 9 ? 1 : i

        # Our actual first acceptable multiple
        q = 30q_quot + q_rem

        byte = cld(p * q, 30)
        wheel_index = 8 * (to_idx(p % 30) - 1) + i

        return byte, wheel_index
    end
end

struct Siever
    prime_div_30::Int

    # byte_index is the integer range 30(byte_index - 1) up to 30byte_index - 1
    byte_index::Int

    # Stores the next prime number to be crossed off.
    # If `p` is the prime number and `q` the next multiple to be stored
    # Wheel index 8 * to_idx(p % 30) * to_idx(q % 30)
    wheel_index::Int

    function Siever(p::Int, segment_lo::Int)
        byte, wheel = compute_offset(p, segment_lo)
        return new(p ÷ 30, byte, wheel)
    end

    Siever(prime_div_30, byte_index, wheel_index) = new(prime_div_30, byte_index, wheel_index)
end

function Base.show(io::IO, p::Siever)
    print(io, 30p.prime_div_30 + ps[(p.wheel_index - 1) ÷ 8 + 1], " (", p.byte_index, ", ", p.wheel_index, ")")
end


"""
    assemble_matrices(model)

Assemble finite element matrices.
"""
function assemble_matrices(model)
    (; mesh, D, T₂, ρ) = model

    # Deduce sizes
    ncompartment, nboundary = size(mesh.facets)
    npoint_cmpts = size.(mesh.points, 2)

    # Assemble finite element matrices
    volumes = zeros(ncompartment)
    for icmpt = 1:ncompartment
        points = mesh.points[icmpt]
        elements = mesh.elements[icmpt]
        fevolumes, = get_mesh_volumes(points, elements)
        volumes[icmpt] = sum(fevolumes)
    end

    # Assemble finite element matrices compartment-wise
    M_cmpts = []
    S_cmpts = []
    Mx_cmpts = [[] for _ = 1:3]
    G = []
    for icmpt = 1:ncompartment
        # Finite elements
        points = mesh.points[icmpt]
        elements = mesh.elements[icmpt]
        fevolumes, = get_mesh_volumes(points, elements)

        # Assemble mass, stiffness and flux matrices
        push!(M_cmpts, assemble_mass_matrix(elements', fevolumes))
        push!(S_cmpts, assemble_stiffness_matrix(elements', points', D[icmpt]))

        # Assemble first order product moment matrices
        for dim = 1:3
            push!(Mx_cmpts[dim], assemble_mass_matrix(elements', fevolumes, points[dim, :]))
        end

        # Compute surface integrals
        push!(G, zeros(npoint_cmpts[icmpt], 3))
        for iboundary = 1:nboundary
            facets = mesh.facets[icmpt, iboundary]
            if !isempty(facets)
                # Get facet normals
                _, _, normals = get_mesh_surfacenormals(points, elements, facets)

                # Surface normal weighted flux matrix (in each canonical direction)
                for dim = 1:3
                    Q = assemble_flux_matrix(facets', points', normals[dim, :])
                    G[icmpt][:, dim] += sum(Q, dims = 2)
                end
            end
        end

    end

    # Assemble global finite element matrices
    M = blockdiag(M_cmpts...)
    S = blockdiag(S_cmpts...)
    R = blockdiag((M_cmpts ./ T₂)...)
    Mx = [blockdiag(Mx_cmpts[dim]...) for dim = 1:3]
    Q_blocks = assemble_flux_matrices(mesh.points, mesh.facets)
    Q = couple_flux_matrix(model, Q_blocks, false)

    (; M, S, R, Mx, Q, M_cmpts, S_cmpts, Mx_cmpts, G, volumes)
end


"""
Created on Wed Apr 15 2020

@author: Yoann Pradat

    CentraleSupelec
    MICS laboratory
    9 rue Juliot Curie, Gif-Sur-Yvette, 91190 France

Implementation of the Maximum-Equalization algorithms and its variants under various constraints and for the β
divergence as presented in C. Fevotte & J. Idier, "Algorithms for nonnegative matrix factorization
with the beta-divergence ", Neural Compuation, 2011.
"""

abstract type AbstractFIParams <: AbstractNMFParams end

"""
    FIParams([,β, scale_W_iter, scale_W_last, α_H, α_W, l₁ratio_H, l₁ratio_W, θ, alg, seed])

Local parameters passed to `nmf_FI`.

# Arguments
- `β::Real`               : parameter of the β-divergence
- `scale_W_iter::Bool`    : scale the components of W to have unit l₁ norm during iterations
- `scale_W_last::Bool`     : scale the components of W to have unit l₁ norm after last iteration
- `α_H::Real`             : constant that multiplies the regularization term on H
- `α_W::Real`             : constant that multiplies the regularization term on W
- `l₁ratio_H::Real`       : l₁/l₂ regularization balance for H (in [0; 1]). 1 is l₁ only, 0 is l₂ only
- `l₁ratio_W::Real`       : l₁/l₂ regularization balance for W (in [0; 1]). 1 is l₁ only, 0 is l₂ only
- `θ::Real`               : θ for mixed update between ME and H. Used only in `nmf_ME`.
- `alg::Symbol`           : symbol to choose the algorithm. Choose
    + :mm for Maximization-Minimization
    + :h  for Heuristic (Maximization-Minimization without exponent)
    + :me for Maximization-Equalization
"""
mutable struct FIParams <: AbstractFIParams
    β::Real
    scale_W_iter::Bool
    scale_W_last::Bool
    α_H::Real
    α_W::Real
    l₁ratio_H::Real
    l₁ratio_W::Real
    θ::Real
    alg::Symbol

    function FIParams(;
        β::Real               = 2,
        scale_W_iter::Bool    = true,
        scale_W_last::Bool    = true,
        α_H::Real             = 0,
        α_W::Real             = 0,
        l₁ratio_H::Real       = 0,
        l₁ratio_W::Real       = 0,
        θ::Real               = 0.95,
        alg::Symbol           = :mm)

        valid_α_H = α_H >= 0
        valid_α_W = α_W >= 0
        valid_l₁ratio_H = l₁ratio_H >= 0 && l₁ratio_H <= 1
        valid_l₁ratio_W = l₁ratio_W >= 0 && l₁ratio_W <= 1
        valid_θ = θ >= 0 && θ <= 1
        valid_alg = alg in [:mm :me :h]

        valid_α_H || throw(ArgumentError("α_H must be > 0"))
        valid_α_W || throw(ArgumentError("α_W must be > 0"))
        valid_l₁ratio_H || throw(ArgumentError("l₁ratio_H must be in [0,1]"))
        valid_l₁ratio_W || throw(ArgumentError("l₁ratio_W must be in [0,1]"))
        valid_θ  || throw(ArgumentError("θ must be between 0 and 1"))
        valid_alg || throw(ArgumentError("alg must be one of :mm, :h or :me"))

        new(β, scale_W_iter, scale_W_last, α_H, α_W, l₁ratio_H, l₁ratio_W, θ, alg)
    end
end

# Overload copy from Base to allow copies FIParams structs
Base.copy(params::FIParams) = FIParams(
        β            = params.β,
        scale_W_iter = params.scale_W_iter,
        scale_W_last = params.scale_W_iter,
        α_W          = params.α_W,
        α_H          = params.α_H,
        l₁ratio_H    = params.l₁ratio_H,
        l₁ratio_W    = params.l₁ratio_W,
        θ            = params.θ,
        alg          = params.alg
    )

"""
    _γ_β(β)

Compute the exponent used in update rules of Maximization-Minimization algorithm.
"""
function _γ_β(β::Real)
    if β < 1
	γ = 1.0 / (2.0-β)
    elseif  β > 2
	γ = 1.0 / (β-1.0)
    else
	γ = 1
    end
    γ
end

"""
    _β_cost(V, W, H, params)

Compute the cost function that is minimized by the NMF alogrithm running with parameters params. It adds up
the β_divergence and regularization terms.

See also: [`β_divergence`](@ref), [`FIParams`](@ref)
"""
function _β_cost(V::Matrix{T}, W::Matrix{T}, H::Matrix{T}, params::FIParams) where T <: Real
    regu_H = params.α_H * (params.l₁ratio_H * norm(H, 1) + (1. - params.l₁ratio_H) * norm(H, 2)) / length(H)
    regu_W = params.α_W * (params.l₁ratio_W * norm(W, 1) + (1. - params.l₁ratio_W) * norm(W, 2)) / length(W)
    β_div  = β_divergence(V, W*H, params.β)

    β_div + regu_H + regu_W
end

"""
    _update_nmf_MM!(V, WH, W, H, params)

Update the H and W matrices in that order using the update rules of Maximization-Minimization algorithm.
"""
function _update_nmf_MM!(V::Matrix{T}, WH::Matrix{T}, W::Matrix{T}, H::Matrix{T}, params::FIParams) where T <: Real
    #### compute regularization
    #### for J_H = ||H||₁, ∇J_H = J_{K,N} with J the matrix filled with ones
    #### for J_H = 0.5*||H||₂, ∇J_H = H
    regu_H = H -> params.α_H * (params.l₁ratio_H * ones(size(H)) + (1. - params.l₁ratio_H) * H) / length(H)
    regu_W = W -> params.α_W * (params.l₁ratio_W * ones(size(W)) + (1. - params.l₁ratio_W) * W) / length(W)

    β = params.β

    #### exponent for multiplicative updates
    γ = _γ_β(β)

    if β > 2
        #### update rules for β > 2
        W .= W .* (max.((V .* WH .^ (β-2)) * transpose(H) - regu_W(W), 0) ./ (WH .^ (β-1) * transpose(H))) .^ γ
        WH .= W * H

        H .= H .* (max.(transpose(W) * (V .* WH .^ (β-2)) - regu_H(H), 0) ./ (transpose(W) * WH .^ (β-1))) .^ γ
        WH .= W * H

    elseif β == 2
        #### update rules for β = 2 can be written more efficiently
        W .= W .* (max.(V * transpose(H) - regu_W(W), 0) ./ (W * (H * transpose(H))))
        H .= H .* (max.((transpose(W) * V - regu_H(H)), 0) ./ ((transpose(W) * W) * H))
        WH .= W * H

    elseif 1 < β < 2
        #### update rules for 1 < β < 2
        #### the update of h_kn is obtained by solving
        #### \sum_f w_{fk} [ (\tilde{v}_f \frac{h_k}{\tilde{h}_k})^{\beta-1} - v_f (\tilde{v}_f \frac{h_k}{\tilde{h}_k})^{\beta-2} ] + \alpha_h \nabla_{h_k} J_k (h_k) = 0 with J_k the component of J_h(h) = \sum_k J_{k}(h_k)

        if params.α_W == 0
            W .= W .* (((V .* WH .^ (β-2)) * transpose(H)) ./ (WH .^ (β-1) * transpose(H)))
        else
            #### solving a X^2 + b X - c = 0 with X = \frac{h_k}{\tilde{h}_k}^{\beta-1}
            a = WH .^ (β-1) * transpose(H)
            b = regu_W(W)
            c = (V .* WH .^ (β-2)) * transpose(H)
            Δ =  b .^ 2 + 4a .* c

            W .= W .* ((- b + sqrt.(Δ)) ./ (2a)) .^ (1. / (β-1))
        end

        WH .= W * H

        if params.α_H == 0
            H .= H .* ((transpose(W) * (V .* WH .^ (β-2))) ./ (transpose(W) * WH .^ (β-1)))
        else
            #### see above
            a = transpose(W) * WH .^ (β-1)
            b = regu_H(H)
            c = transpose(W) * (V .* WH .^ (β-2))
            Δ =  b .^ 2 + 4a .* c

            H .= H .* ((- b + sqrt.(Δ)) ./ (2 .* a)) .^ (1. / (β-1))
        end

        WH .= W * H

    elseif β <= 1
        #### update rules for β <= 1
        W .= W .* (((V .* WH .^ (β-2)) * transpose(H)) ./ (WH .^ (β-1) * transpose(H) + regu_W(W))) .^ γ
        WH .= W * H

        H .= H .* ((transpose(W) * (V .* WH .^ (β-2))) ./ (transpose(W) * WH .^ (β-1) + regu_H(H))) .^ γ
        WH .= W * H
    end

    #### avoid underflow
    W .= max.(W, eps(T))
    H .= max.(H, eps(T))
    WH .= W * H

    if params.scale_W_iter
        scale_col_W!(W, H, 1)
    end
end

"""
    _update_nmf_H!(V, WH, W, H, params)

Update the H and W matrices in that order using the update rules of Heuristic algorithm.
"""
function _update_nmf_H!(V::Matrix{T}, WH::Matrix{T}, W::Matrix{T}, H::Matrix{T}, params::AbstractFIParams) where T <: Real
    #### compute regularization
    #### for J_H = ||H||₁, ∇J_H = J_{K,N} with J the matrix filled with ones
    #### for J_H = 0.5*||H||₂, ∇J_H = H
    regu_H = H -> params.α_H * (params.l₁ratio_H * ones(size(H)) + (1. - params.l₁ratio_H) * H) / length(H)
    regu_W = W -> params.α_W * (params.l₁ratio_W * ones(size(W)) + (1. - params.l₁ratio_W) * W) / length(W)

    if params.β > 2
        #### update rules for β > 2
        W .= W .* (((V .* WH .^ (params.β-2)) * transpose(H) - regu_W(W)) ./ (WH .^ (params.β-1) * transpose(H)))
        WH .= W * H

        H .= H .* ((transpose(W) * (V .* WH .^ (params.β-2)) - regu_H(H)) ./ (transpose(W) * WH .^ (params.β-1)))
        WH .= W * H

    elseif params.β == 2
        #### update rules for β = 2 can be written more efficiently
        W .= W .* ((V * transpose(H) - regu_W(W)) ./ (W * (H * transpose(H))))
        H .= H .* ((transpose(W) * V - regu_H(H)) ./ ((transpose(W) * W) * H))
        WH .= W * H

    elseif 1 < params.β < 2
        #### update rules for 1 < β < 2
        #### the update of h_kn is obtained by solving
        #### \sum_f w_{fk} [ (\tilde{v}_f \frac{h_k}{\tilde{h}_k})^{\beta-1} - v_f (\tilde{v}_f \frac{h_k}{\tilde{h}_k})^{\beta-2} ] + \alpha_h \nabla_{h_k} J_k (h_k) = 0 with J_k the component of J_h(h) = \sum_k J_{k}(h_k)

        if params.α_W == 0
            W .= W .* (((V .* WH .^ (params.β-2)) * transpose(H)) ./ (WH .^ (params.β-1) * transpose(H)))
        else
            #### solving a X^2 + b X - c = 0 with X = \frac{h_k}{\tilde{h}_k}^{\beta-1}
            a = WH .^ (params.β-1) * transpose(H)
            b = regu_W(W)
            c = (V .* WH .^ (params.β-2)) * transpose(H)
            Δ =  b .^ 2 + 4a .* c

            W .= W .* ((- b + sqrt.(Δ)) ./ (2 .* a)) .^ (1. / (params.β-1))
        end

        WH .= W * H

        if params.α_H == 0
            H .= H .* ((transpose(W) * (V .* WH .^ (params.β-2))) ./ (transpose(W) * WH .^ (params.β-1)))
        else
            #### see above
            a = transpose(W) * WH .^ (params.β-1)
            b = regu_H(H)
            c = transpose(W) * (V .* WH .^ (params.β-2))
            Δ =  b .^ 2 + 4a .* c

            H .= H .* ((- b + sqrt.(Δ)) ./ (2a)) .^ (1. / (params.β-1))
        end

        WH .= W * H

    elseif params.β <= 1
        #### update rules for β <= 1
        W .= W .* (((V .* WH .^ (params.β-2)) * transpose(H)) ./ (WH .^ (params.β-1) * transpose(H) + regu_W(W)))
        WH .= W * H

        H .= H .* ((transpose(W) * (V .* WH .^ (params.β-2))) ./ (transpose(W) * WH .^ (params.β-1) + regu_H(H)))
        WH .= W * H
    end

    #### avoid underflow
    W .= max.(W, eps(T))
    H .= max.(H, eps(T))
    WH .= W * H

    if params.scale_W_iter
        scale_col_W!(W, H, 1)
    end
end

"""
    _update_nmf_ME!(V, WH, W, H, params)

Update the H and W matrices in that order using the update rules of Maximization-Equalization algorithm.
"""
function _update_nmf_ME!(V::Matrix{T}, WH::Matrix{T}, W::Matrix{T}, H::Matrix{T}, params::FIParams) where T <: Real
    #### compute regularization
    #### for J_H = ||H||₁, ∇J_H = J_{K,N} with J the matrix filled with ones
    #### for J_H = 0.5*||H||₂, ∇J_H = H
    regu_H = H -> params.α_H * (params.l₁ratio_H * ones(size(H)) + (1. - params.l₁ratio_H) * H) / length(H)
    regu_W = W -> params.α_W * (params.l₁ratio_W * ones(size(W)) + (1. - params.l₁ratio_W) * W) / length(W)

    if params.β == -1
        #### ME always defined
        W_h = W .* (((V .* WH .^ (params.β-2)) * transpose(H)) ./ (WH .^ (params.β-1) * transpose(H) + regu_W(W)))
        W .= ((W_h .^ 2 + 8 * W_h .* W) .^ .5 + W_h) / 4
        WH .= W * H

        H_h = H .* ((transpose(W) * (V .* WH .^ (params.β-2))) ./ (transpose(W) * WH .^ (params.β-1) + regu_H(H)))
        H .= ((H_h.^2 + 8 * H_h .* H) .^ .5 + H_h) / 4
        WH .= W * H

    elseif params.β == 0
        #### ME always defined, equal to H
        W .= W .* (((V .* WH .^ (params.β-2)) * transpose(H)) ./ (WH .^ (params.β-1) * transpose(H) + regu_W(W)))
        WH .= W * H

        H .= H .* ((transpose(W) * (V .* WH .^ (params.β-2))) ./ (transpose(W) * WH .^ (params.β-1) + regu_H(H)))
        WH .= W * H

    elseif params.β == 0.5
        #### ME always defined
        W_h = W .* (((V .* WH .^ (params.β-2)) * transpose(H)) ./ (WH .^ (params.β-1) * transpose(H) + regu_W(W)))
        W .= (sqrt.(W + 8 * W_h) - sqrt.(W)) .^ 2 / 4
        WH .= W * H

        H_h = H .* ((transpose(W) * (V .* WH .^ (params.β-2))) ./ (transpose(W) * WH .^ (params.β-1) + regu_H(H)))
        H .= (sqrt.(H + 8 * H_h) - sqrt.(H)) .^2 / 4
        WH .= W * H

    elseif params.β == 1.5
        #### mixed update prolonged ME and MM

        #### compute W update from Heuristic algorithm
        if params.α_W == 0
            W_h = W .* (((V .* WH .^ (params.β-2)) * transpose(H)) ./ (WH .^ (params.β-1) * transpose(H)))
        else
            a = WH .^ (params.β-1) * transpose(H)
            b = regu_W(W)
            c = (V .* WH .^ (params.β-2)) * transpose(H)
            Δ =  b .^ 2 + 4a .* c

            W_h = W .* ((- b + sqrt.(Δ)) ./ (2a)) .^ (1. / (params.β-1))
        end

        #### compute mixed update
        DW = 12W_h - 3W
        IW = 3W_h .< W
        W .= real.((sqrt.(Complex.(DW)) - sqrt.(W)) .^2 / 4)
        W[IW] .= 0
        W .= W_h + params.θ * (W - W_h)
        WH .= W * H

        #### compute H update from Heuristic algorithm
        if params.α_H == 0
            H_h = H .* ((transpose(W) * (V .* WH .^ (params.β-2))) ./ (transpose(W) * WH .^ (params.β-1)))
        else
            a = transpose(W) * WH .^ (params.β-1)
            b = regu_H(H)
            c = transpose(W) * (V .* WH .^ (params.β-2))
            Δ =  b .^ 2 + 4a .* c

            H_h = H .* ((- b .+ sqrt.(Δ)) ./ (2a)) .^ (1. / (params.β-1))
        end

        #### compute mixed update
        DH = 12H_h - 3H
        IH = 3H_h .< H
        H .= real.((sqrt.(Complex.(DH)) - sqrt.(H)) .^2 / 4)
        H[IH] .= 0
        H .= H_h + params.θ * (H - H_h)
        WH .= W * H

    elseif params.β == 2
        #### mixed update prolonged ME and MM
        W_h = W .* ((V * transpose(H) - regu_W(W)) ./ (W * (H * transpose(H))))
        W .= 2W_h - W
        W[W .<= 0] .= 0
        W .= W_h + params.θ * (W - W_h)
        WH .= W * H

        H_h = H .* ((transpose(W) * V - regu_H(H)) ./ ((transpose(W) * W) * H))
        H .= 2H_h - H
        H[H .<= 0] .= 0
        H .= H_h + params.θ * (H - H_h)
        WH .= W * H

    elseif params.β == 3
        #### mixed update prolonged ME and MM
        W_h = W .* (((V .* WH .^ (params.β-2)) * transpose(H) - regu_W(W)) ./ (WH .^ (params.β-1) * transpose(H)))
        DW = W .* (12W_h - 3W)
        IW = 3W_h .< W
        W .= real.(sqrt.(Complex.(DW)) - W) / 2
        W[IW] .= 0
        W .= W_h + params.θ * (W - W_h)
        WH .= W * H

        H_h = H .* ((transpose(W) * (V .* WH .^ (params.β-2)) - regu_H(H)) ./ (transpose(W) * WH .^ (params.β-1)))

        DH = H .* (12H_h - 3H)
        IH = 3H_h .< H
        H .= real.(sqrt.(Complex.(DH)) - H) / 2
        H[IH] .= 0
        H .= H_h + params.θ * (H - H_h)

        WH .= W * H
    else
        throw(ArgumentError("ME algorithm not implemented for β=$(params.β)"))
    end

    #### avoid underflow
    W .= max.(W, eps(T))
    H .= max.(H, eps(T))
    WH .= W * H

    if params.scale_W_iter
        scale_col_W!(W, H, 1)
    end
end

"""
    _nmf_FI(V, W_init, H_init, global_params, local_params)

Helper function for `nmf_FI`. Runs the FI algorithms with given intial matrices W and H.
"""
function _nmf_FI(V::Matrix{T}, W_init::Matrix{T}, H_init::Matrix{T}, global_params::AbstractNMFParams, local_params::AbstractFIParams) where T <: Real
    W  = W_init[:,:]
    H  = H_init[:,:]
    WH = W * H

    #### metrics and stopping_crit
    func_cost          = (V, W, H) -> _β_cost(V, W, H, local_params)
    func_divg          = (V, W, H) -> β_divergence(V, W * H, local_params.β)
    metrics            = NMFMetrics(V, W_init, H_init, func_cost, func_divg)
    stopping_crit_met  = false
    iteration_is_sane  = true

    #### scale W before first iter if specified
    if local_params.scale_W_iter
        scale_col_W!(W, H, 1)
    end

    while !stopping_crit_met && iteration_is_sane
        #### update WH, W, H
        if local_params.alg == :mm
            _update_nmf_MM!(V, WH, W, H, local_params)
        elseif local_params.alg == :h
            _update_nmf_H!(V, WH, W, H, local_params)
        elseif local_params.alg == :me
            _update_nmf_ME!(V, WH, W, H, local_params)
        end

        #### update metrics
        _update_metrics!(metrics, V, W, H, global_params.stopping_crit)

        #### check sanity of iteration looking for potential divisions by zeros
        iteration_is_sane = _check_sanity_iter(W, H)

        #### check the stopping criterion
        stopping_crit_met = _check_stopping_crit(metrics, global_params)

        #### print intermediate if verbose ≥ 2
        if global_params.verbose ≥ 2 && global_params.max_iter < 10
            @printf("iter %d/%d | cost: %.5g\n", metrics.cur_iter, global_params.max_iter, metrics.cur_cost)
        elseif global_params.verbose ≥ 2 && (metrics.cur_iter % (global_params.max_iter ÷ 10) == 0)
            @printf("iter %d/%d | cost: %.5g\n", metrics.cur_iter, global_params.max_iter, metrics.cur_cost)
        end
    end

    if local_params.scale_W_last
        scale_col_W!(W, H, 1)
    end

    #### convergence information
    if !iteration_is_sane
        @warn "Stopped nmf because sanity check not passed at iteration" metrics.cur_iter
    else
        if global_params.verbose ≥ 1 && metrics.cur_iter == global_params.max_iter && global_params.stopping_crit != :none
            @warn "Maximum number of iterations reached"
        end
        if global_params.verbose ≥ 1 && metrics.cur_iter < global_params.max_iter
            @info "Stopping criterion $(global_params.stopping_crit) met at iter $(metrics.cur_iter)"
        end
    end

    NMFResults{T}(W, H, metrics, global_params, local_params, stopping_crit_met)
end

"""
    nmf_FI(V, global_params, local_params)

Run NMF algorithm using the set of input parameters. For the parameters, see [`Params`](@ref) and [`FIParams`](@ref).

# References
- C. Fevotte & J. Idier, "Algorithms for nonnegative matrix factorization
with the beta-divergence ", Neural Compuation, 2011.
"""
function nmf_FI(V::Matrix{T}, global_params::AbstractNMFParams, local_params::AbstractFIParams) where T <: Real
    # initalization
    W_init, H_init = initialize_nmf(
        V,
        rank = global_params.rank,
        init = global_params.init,
        dist = global_params.dist,
        seed = global_params.seed
    )

    # algorithm
    _nmf_FI(V, W_init, H_init, global_params, local_params)
end


# Used by fptas below
struct ScaleParams
    t::Vector{Int}
    ft::Vector{Float64}
    Ū::Int
    infty::Int

    function ScaleParams(mkt::VariedCostsMarket, ε::Float64)
        @assert 0 < ε < 1
        P = max(0,
            ceil(Int, log2(Int(mkt.m)^2 / (ε * sum(mkt.ft))))
        )
        t = mkt.t .* 2^P
        ft = mkt.f .* t
        Ū = ceil(Int, sum(ft))
        infty = sum(mkt.g) + 1

        return new(t, ft, Ū, infty)
    end
end

# Used by fptas below
@inbounds function G_recursor!(
    G_dict::Dict{Tuple{Int,Int},Int},
    j::Int,
    v::Int,
    mkt::VariedCostsMarket,
    sp::ScaleParams
)::Int
    return get(G_dict, (j, v)) do
        if v ≤ 0
            return 0
        elseif j == 0 || sp.t[j] < v || v ≥ sp.Ū
            return sp.infty
        else
            jmo = j - 1
            if mkt.f[j] < 1
                # Clamping prevents over/underflow: for any v<0 or v≥Ū the function
                # is trivially defined, so recording any more extreme number is meaningless
                v_minus_Δ = floor(Int, clamp((v - sp.ft[j]) / mkt.omf[j], -1, sp.Ū))

                push!(G_dict, (j, v) => min(
                    G_recursor!(G_dict, jmo, v, mkt, sp),
                    mkt.g[j] + G_recursor!(G_dict, jmo, v_minus_Δ, mkt, sp)
                ))
                return G_dict[(j, v)]
            else
                push!(G_dict, (j, v) => min(G_recursor!(G_dict, jmo, v, mkt, sp), mkt.g[j]))
                return G_dict[(j, v)]
            end
        end
    end
end

# Function I stole from Base.Sort that does midpoint using bitshift.
# Performance-optimized but safe only if `lo <= hi`.
midpoint(lo::T, hi::T) where {T<:Integer} = lo + ((hi - lo) >>> 0x01)

"""
    optimalportfolio_fptas(mkt, ε; verbose=false) -> X, v

Use the fully polynomial-time approximation scheme to produce a `1-ε`-optimal portfolio `X`,
with valuation `v`, for the [`VariedCostsMarket`](@ref) defined by `mkt`.
"""
function optimalportfolio_fptas(
    mkt::VariedCostsMarket,
    ε::Float64;
    verbose::Bool=false
)::Tuple{Vector{Int},Float64}
    sp = ScaleParams(mkt, ε)

    G_dict = Dict{Tuple{Int,Int},Int}()
    sizehint!(G_dict, mkt.m * mkt.m ÷ 2)

    # Binary search for max{ w :  G_recursor!(G_dict, mkt.m, w, mkt, sp ≤ H }
    # Should be able to use something like this with a proper by kw in searchsortedlast
    # vs = 0:sp.Ū
    # searchsortedlast(vs, by = ...)

    v = 0
    v_UB = sp.Ū
    @inbounds while v + 1 < v_UB
        # mid = (v + v_UB) ÷ 2
        mid = midpoint(v, v_UB)

        if G_recursor!(G_dict, mkt.m, mid, mkt, sp) > mkt.H
            v_UB = mid
        else
            v = mid
        end
    end

    X = Int[]
    @inbounds for j in reverse(1:mkt.m)
        # G_recursor!(G_dict, j, v, mkt, sp) < sp.infty &&
        if G_recursor!(G_dict, j, v, mkt, sp) < G_recursor!(G_dict, j - 1, v, mkt, sp)
            push!(X, j)
            v = floor(Int, clamp((v - sp.ft[j]) / mkt.omf[j], -1, sp.Ū))
        end
    end

    if verbose
        G_table = Array{Union{Missing,Int}}(missing, mkt.m, v_UB)
        for (j, v) in keys(G_dict)
            if 0 < j ≤ mkt.m && 0 < v ≤ v_UB
                G_table[j, v] = G_dict[(j, v)]
            end
        end

        issorted(mkt.perm) || @warn "t not sorted; table rows will differ from input"
        display(G_table)
    end

    return mkt.perm[X], valuation_nopermute_sorted(sort(X), mkt)
end


using Surrogates
using LinearAlgebra
using Flux
using Flux: @epochs
using PolyChaos

#######SRBF############
##### 1D #####

lb = 0.0
ub = 15.0
objective_function = x -> 2*x+1
x = [2.5,4.0,6.0]
y = [6.0,9.0,13.0]

# In 1D values of p closer to 2 make the det(R) closer and closer to 0,
#this does not happen in higher dimensions because p would be a vector and not
#all components are generally C^inf
p = 1.99
a = 2
b = 6

#Using Kriging
my_k_SRBF1 = Kriging(x,y,lb,ub)
surrogate_optimize(objective_function,SRBF(),a,b,my_k_SRBF1,UniformSample())

#Using RadialBasis
my_rad_SRBF1 = RadialBasis(x,y,a,b,rad = linearRadial())
surrogate_optimize(objective_function,SRBF(),a,b,my_rad_SRBF1,UniformSample())

my_wend_1d = Wendland(x,y,lb,ub)
surrogate_optimize(objective_function,SRBF(),a,b,my_wend_1d,UniformSample())

my_earth1d = EarthSurrogate(x,y,lb,ub)
surrogate_optimize(objective_function,SRBF(),a,b,my_earth1d,LowDiscrepancySample(2))

##### ND #####
objective_function_ND = z -> 3*norm(z)+1
lb = [1.0,1.0]
ub = [6.0,6.0]
x = sample(5,lb,ub,SobolSample())
y = objective_function_ND.(x)
p = [1.5,1.5]
theta = [1.0,1.0]

#Kriging

my_k_SRBFN = Kriging(x,y,lb,ub)
#Every optimization method now returns the y_min and its position
x_min, y_min = surrogate_optimize(objective_function_ND,SRBF(),lb,ub,my_k_SRBFN,UniformSample())

#Radials
lb = [1.0,1.0]
ub = [6.0,6.0]
x = sample(5,lb,ub,SobolSample())
objective_function_ND = z -> 3*norm(z)+1
y = objective_function_ND.(x)
my_rad_SRBFN = RadialBasis(x,y,lb,ub,rad = linearRadial())
surrogate_optimize(objective_function_ND,SRBF(),lb,ub,my_rad_SRBFN,UniformSample())

# Lobachevsky
x = sample(5,lb,ub,UniformSample())
y = objective_function_ND.(x)
alpha = [2.0,2.0]
n = 4
my_loba_ND = LobachevskySurrogate(x,y,lb,ub)
surrogate_optimize(objective_function_ND,SRBF(),lb,ub,my_loba_ND,UniformSample())
#Linear
lb = [1.0,1.0]
ub = [6.0,6.0]
x = sample(500,lb,ub,SobolSample())
objective_function_ND = z -> 3*norm(z)+1
y = objective_function_ND.(x)
my_linear_ND = LinearSurrogate(x,y,lb,ub)
surrogate_optimize(objective_function_ND,SRBF(),lb,ub,my_linear_ND,SobolSample(),maxiters=15)

#=
#SVM
lb = [1.0,1.0]
ub = [6.0,6.0]
x = sample(5,lb,ub,SobolSample())
objective_function_ND = z -> 3*norm(z)+1
y = objective_function_ND.(x)
my_SVM_ND = SVMSurrogate(x,y,lb,ub)
surrogate_optimize(objective_function_ND,SRBF(),lb,ub,my_SVM_ND,SobolSample(),maxiters=15)
=#

#Inverse distance surrogate
lb = [1.0,1.0]
ub = [6.0,6.0]
x = sample(5,lb,ub,SobolSample())
objective_function_ND = z -> 3*norm(z)+1
my_p = 2.5
y = objective_function_ND.(x)
my_inverse_ND = InverseDistanceSurrogate(x,y,lb,ub,p=my_p)
surrogate_optimize(objective_function_ND,SRBF(),lb,ub,my_inverse_ND,SobolSample(),maxiters=15)

#SecondOrderPolynomialSurrogate
lb = [0.0,0.0]
ub = [10.0,10.0]
obj_ND = x -> log(x[1])*exp(x[2])
x = sample(15,lb,ub,UniformSample())
y = obj_ND.(x)
my_second_order_poly_ND = SecondOrderPolynomialSurrogate(x,y,lb,ub)
surrogate_optimize(obj_ND,SRBF(),lb,ub,my_second_order_poly_ND,SobolSample(),maxiters=15)

####### LCBS #########
######1D######
objective_function = x -> 2*x+1
lb = 0.0
ub = 15.0
x = [2.0,4.0,6.0]
y = [5.0,9.0,13.0]
p = 1.8
a = 2.0
b = 6.0
my_k_LCBS1 = Kriging(x,y,lb,ub)
surrogate_optimize(objective_function,LCBS(),a,b,my_k_LCBS1,UniformSample())

##### ND #####
objective_function_ND = z -> 3*norm(z)+1
x = [(1.2,3.0),(3.0,3.5),(5.2,5.7)]
y = objective_function_ND.(x)
p = [1.2,1.2]
theta = [2.0,2.0]
lb = [1.0,1.0]
ub = [6.0,6.0]

#Kriging
my_k_LCBSN = Kriging(x,y,lb,ub)
surrogate_optimize(objective_function_ND,LCBS(),lb,ub,my_k_LCBSN,UniformSample())

##### EI ######

###1D###
objective_function = x -> (x+1)^2 - x + 2# Minimum of this function is at x = -0.5, y = -2.75
true_min_x = -0.5
true_min_y = objective_function(true_min_x)
lb = -5
ub = 5
x = sample(5, lb, ub, SobolSample())
y = objective_function.(x)
my_k_EI1 = Kriging(x,y,lb,ub; p = 2)
surrogate_optimize(objective_function,EI(),lb,ub,my_k_EI1, SobolSample(),maxiters=200,num_new_samples=155)

# Check that EI is correctly minimizing the objective
y_min, index_min = findmin(my_k_EI1.y)
x_min = my_k_EI1.x[index_min]
@test norm(x_min - true_min_x) < 0.05 * norm(ub .- lb)
@test abs(y_min - true_min_y) < 0.05 * (objective_function(ub) - objective_function(lb))

###ND###
objective_function_ND = z -> 3*norm(z)+1 # this is minimized at x = (0, 0), y = 1
true_min_x = (0.0, 0.0)
true_min_y = objective_function_ND(true_min_x)
x = [(1.2,3.0),(3.0,3.5),(5.2,5.7)]
y = objective_function_ND.(x)
min_y = minimum(y)
p = [1.2,1.2]
theta = [2.0,2.0]
lb = [-1.0,-1.0]
ub = [6.0,6.0]

#Kriging
my_k_EIN = Kriging(x,y,lb,ub)
surrogate_optimize(objective_function_ND,EI(),lb,ub,my_k_EIN,SobolSample())

# Check that EI is correctly minimizing instead of maximizing
y_min, index_min = findmin(my_k_EIN.y)
x_min = my_k_EIN.x[index_min]
@test norm(x_min .- true_min_x) < 0.05 * norm(ub .- lb)
@test abs(y_min - true_min_y) < 0.05 * (objective_function_ND(ub) - objective_function_ND(lb))

###ND with SectionSampler###
# We will make sure the EI function finds the minimum when constrained to a specific slice of 3D space
objective_function_section = x -> x[1]^2 + x[2]^2 + x[3]^2 # this is minimized at x = (0, 0, 0), y = 0

# We will constrain x[2] to some value
x2_constraint = 2.0
true_min_x = (0.0, x2_constraint, 0.0)
true_min_y = objective_function_section(true_min_x)

sampler = SectionSample([NaN64, x2_constraint, NaN64], SobolSample())
lb = [-1.0,x2_constraint,-1.0]
ub = [6.0,x2_constraint,6.0]
x = sample(5, lb, ub, sampler)
y = objective_function_section.(x)
p = [1.2,1.2,1.2]
theta = [2.0,2.0,2.0]

#Kriging
my_k_EIN_section = Kriging(x,y,lb,ub; p = p)
# Constrain our sampling to the plane where x[2] = 1
surrogate_optimize(objective_function_section,EI(),lb,ub,my_k_EIN_section, sampler)

# Check that EI is correctly minimizing instead of maximizing
y_min, index_min = findmin(my_k_EIN_section.y)
x_min = my_k_EIN_section.x[index_min]
@test norm(x_min .- true_min_x) < 0.05 * norm(ub .- lb)
@test abs(y_min - true_min_y) < 0.05 * (objective_function_section(ub) - objective_function_section(lb))

## DYCORS ##

#1D#
objective_function = x -> 3*x+1
x = [2.1,2.5,4.0,6.0]
y = objective_function.(x)
p = 1.9
lb = 2.0
ub = 6.0

my_k_DYCORS1 = Kriging(x,y,lb,ub,p=1.9)
surrogate_optimize(objective_function,DYCORS(),lb,ub,my_k_DYCORS1,UniformSample())

my_rad_DYCORS1 = RadialBasis(x,y,lb,ub,rad = linearRadial())
surrogate_optimize(objective_function,DYCORS(),lb,ub,my_rad_DYCORS1,UniformSample())

#ND#
objective_function_ND = z -> 2*norm(z)+1
x = [(2.3,2.2),(1.4,1.5)]
y = objective_function_ND.(x)
p = [1.5,1.5]
theta = [2.0,2.0]
lb = [1.0,1.0]
ub = [6.0,6.0]

my_k_DYCORSN = Kriging(x,y,lb,ub)
surrogate_optimize(objective_function_ND,DYCORS(),lb,ub,my_k_DYCORSN,UniformSample(),maxiters=30)

my_rad_DYCORSN = RadialBasis(x,y,lb,ub,rad = linearRadial())
surrogate_optimize(objective_function_ND,DYCORS(),lb,ub,my_rad_DYCORSN,UniformSample(),maxiters=30)

my_wend_ND = Wendland(x,y,lb,ub)
surrogate_optimize(objective_function_ND,DYCORS(),lb,ub,my_wend_ND,UniformSample(),maxiters=30)

### SOP ###
# 1D
objective_function = x -> 3*x+1
x = sample(20,1.0,6.0,SobolSample())
y = objective_function.(x)
p = 1.9
lb = 1.0
ub = 6.0
num_centers = 2
my_k_SOP1 = Kriging(x,y,lb,ub,p=1.9)
surrogate_optimize(objective_function,SOP(num_centers),lb,ub,my_k_SOP1,SobolSample(),maxiters=60)
#ND
objective_function_ND = z -> 2*norm(z)+1
x = [(2.3,2.2),(1.4,1.5)]
y = objective_function_ND.(x)
p = [1.5,1.5]
theta = [2.0,2.0]
lb = [1.0,1.0]
ub = [6.0,6.0]
my_k_SOPND = Kriging(x,y,lb,ub)
num_centers = 2
surrogate_optimize(objective_function_ND,SOP(num_centers),lb,ub,my_k_SOPND,SobolSample(),maxiters=20)

#multi optimization
#=
f  = x -> [x^2, x]
lb = 1.0
ub = 10.0
x  = sample(100, lb, ub, SobolSample())
y  = f.(x)
my_radial_basis_smb = RadialBasis(x, y, lb, ub, rad = linearRadial())
surrogate_optimize(f,SMB(),lb,ub,my_radial_basis_ego,SobolSample())

f  = x -> [x^2, x]
lb = 1.0
ub = 10.0
x  = sample(100, lb, ub, SobolSample())
y  = f.(x)
my_radial_basis_rtea = RadialBasis(x, y, lb, ub, rad = linearRadial())
Z = 0.8 #percentage
K = 2 #number of revaluations
p_cross = 0.5 #crossing vs copy
n_c = 1.0 # hyperparameter for children creation
sigma = 1.5 # mutation
surrogate_optimize(f,RTEA(Z,K,p_cross,n_c,sigma),lb,ub,my_radial_basis_rtea,SobolSample())
=#


#!/usr/bin/env julia
# -*- coding: UTF-8 -*-
# __julia-version__ = 1.7.2
# __author__        = "Max Kannenberg"
# __copyright__     = "2022"
# __license__       = "ISC"

#########################
## literature the driving dynamics equations are based on:
##
## @incollection{Bruenger:2014, % Chapter 4
##   author    = {Brünger, Olaf and Dahlhaus, Elias},
##   year      = {2014},
##   title     = {Running Time Estimation},
##   pages     = {65--90},
##   booktitle = {Railway Timetabling \& Operations.},
##   editora   = {Hansen, Ingo A.},
##   editorb   = {Pachl, Jörn},
##   isbn      = {978-3-777-10462-1},
##   publisher = {Eurailpress DVV Media Group},
## }
## @Book{Wende:2003,
##   author    = {Wende, Dietrich},
##   date      = {2003},
##   title     = {Fahrdynamik des Schienenverkehrs},
##   isbn      = {978-3-322-82961-0},
##   publisher = {Springer-Verlag},
## }
#########################

approxLevel = 6
v00 = 100/3.6     # velocity factor (in m/s)

## calculate forces

#TODO: replace the ? ? ?
"""
    calcTractionUnitResistance(v, train)

Calculate the vehicle resistance for the traction unit of the `train` dependend on the velocity `v`.

...
# Arguments
- `v::AbstractFloat`: the current velocity in m/s.
- `train::Train`: ? ? ?
...

# Examples
```julia-repl
julia> calcTractionUnitResistance(30.0, ? ? ?)
? ? ?
```
"""
function calcTractionUnitResistance(v::AbstractFloat, train::Train)
    # equation is based on [Wende:2003, page 151]
    f_Rtd0 = train.f_Rtd0 # coefficient for basic resistance due to the traction units driving axles (in ‰)
    f_Rtc0 = train.f_Rtc0 # coefficient for basic resistance due to the traction units carring axles (in ‰)
    F_Rt2  = train.F_Rt2  # coefficient for air resistance of the traction units (in N)
    m_td   = train.m_td   # mass on the traction unit's driving axles (in kg)
    m_tc   = train.m_tc   # mass on the traction unit's carrying axles (in kg)

    F_R_tractionUnit = f_Rtd0/1000 * m_td * g + f_Rtc0/1000 * m_tc * g + F_Rt2 * ((v + Δv_air) /v00)^2   # vehicle resistance of the traction unit (in N)   # /1000 because of the unit ‰
    # TODO: use calcForceFromCoefficient? F_R_tractionUnit = calcForceFromCoefficient(f_Rtd0, m_td) + calcForceFromCoefficient(f_Rtc0, m_tc) + F_Rt2 * ((v + Δv_air) /v00)^2       # vehicle resistance of the traction unit (in N)
    return F_R_tractionUnit
    #TODO: same variable name like in the rest of the tool? return R_traction
    #TODO: just one line? return train.f_Rtd0/1000*train.m_td*g+train.f_Rtc0/1000*train.m_tc*g+train.F_Rt2*((v+train.Δv_air)/v00)^2    # /1000 because of the unit ‰
end #function calcTractionUnitResistance

"""
TODO
calculate and return the wagons vehicle resistance dependend on the velocity
"""
function calcWagonsResistance(v::AbstractFloat, train::Train)
    # equation is based on a combination of the equations of Strahl and Sauthoff [Wende:2003, page 153] with more detailled factors (Lehmann, page 135)
    f_Rw0  = train.f_Rw0  # coefficient for basic resistance of the set of wagons (consist)  (in ‰)
    f_Rw1  = train.f_Rw1  # coefficient for the consists resistance to rolling (in ‰)
    f_Rw2  = train.f_Rw2  # coefficient fo the consistsr air resistance (in ‰)
    m_w    = train.m_w    # mass of the set of wagons (consist)  (in kg)

    F_R_wagons = m_w *g *(f_Rw0/1000 + f_Rw1/1000 *v /v00 + f_Rw2/1000 * ((v + Δv_air) /v00)^2)     # vehicle resistance of the wagons (in N)      # /1000 because of the unit ‰
# TODO: use calcForceFromCoefficient?    F_R_wagons = calcForceFromCoefficient(f_Rw0, m_w) + calcForceFromCoefficient(f_Rw1, m_w) *v /v00 + calcForceFromCoefficient(f_Rw2, m_w) * ((v + Δv_air) /v00)^2     # vehicle resistance of the wagons (in N)
    return F_R_wagons
end #function calcWagonsResistance

function calcForceFromCoefficient(f_R::Real, m::Real)
    # equation is based on [Wende:2003, page 8]

    # f_R: specific resistance (in ‰)
    # m: vehicle's mass (in kg)

    F_R = f_R /1000 *m *g     # Resisting Force (in N)  # /1000 because of the unit ‰
    return F_R
end #function calcForceFromCoefficient

function calcAcceleration(F_T::Real, F_R::Real, m_train::Real, ξ_train::Real)
    # equation is based on [Bruenger:2014, page 72] with a=dv/dt

    # F_T: tractive effort (in N)
    # F_R: resisting forces (in N)
    # m_train: train's mass (in kg)
    # ξ_train: train's rotation mass factor (without unit)

    a = (F_T - F_R) /m_train /ξ_train      # acceleration (in m/s)
    return a
end #function calcAcceleration

function calc_Δs_with_Δt(Δt::Real, a_prev::Real, v_prev::Real)
    # equation is based on [Wende:2003, page 37]

    # Δt: time step (in s)
    # a_prev: acceleration from previous data point
    # v_prev: velocitiy from previous data point
    Δs = Δt * (2*v_prev + Δt*a_prev) /2        # step size (in m)
    return Δs
end #function calc_Δs_with_Δt

function calc_Δs_with_Δv(Δv::Real, a_prev::Real, v_prev::Real)
    # equation is based on [Wende:2003, page 37]

    # Δv: velocity step (in m/s)
    # a_prev: acceleration from previous data point
    # v_prev: velocitiy from previous data point
    Δs = ((v_prev + Δv)^2 - v_prev^2)/2/a_prev      # step size (in m)
    return Δs
end #function calc_Δs_with_Δv

function calc_Δt_with_Δs(Δs::Real, a_prev::Real, v_prev::Real)
    # equation is based on [Wende:2003, page 37]

    # Δs: distance step (in m)
    # a_prev: acceleration from previous data point
    # v_prev: velocitiy from previous data point

    Δt = sign(a_prev) *sqrt((v_prev /a_prev)^2 + 2 *Δs /a_prev) - v_prev /a_prev       # step size (in m/s)
    return Δt
end #function calc_Δt_with_Δs

function calc_Δt_with_Δv(Δv::Real, a_prev::Real)
    # equation is based on [Wende:2003, page 37]

    # Δv: velocity step (in m/s)
    # a_prev: acceleration from previous data point
    Δt = Δv /a_prev        # step size (in s)
    return Δt
end #function calc_Δt_with_Δv

function calc_Δt_with_constant_v(Δs::Real, v::Real)
    # equation is based on [Wende:2003, page 37]

    # Δs: distance step (in m)
    # v: constant velocity (in m/s)
    Δt = Δs /v        # step size (in s)
    return Δt
end #function calc_Δt_with_constant_v

function calc_Δv_with_Δs(Δs::Real, a_prev::Real, v_prev::Real)
    # equation is based on [Wende:2003, page 37]

    # Δs: distance step (in m)
    # a_prev: acceleration from previous data point
    # v_prev: velocitiy from previous data point
    Δv = sqrt(v_prev^2 + 2*Δs*a_prev) - v_prev      # step size (in m/s)
    return Δv
end #function calc_Δv_with_Δs

function calc_Δv_with_Δt(Δt::Real, a_prev::Real)
    # equation is based on [Wende:2003, page 37]

    # Δt: time step (in s)
    # a_prev: acceleration from previous data point
    Δv = Δt * a_prev        # step size (in m/s)
    return Δv
end #function calc_Δv_with_Δt

function calc_ΔW(F_T_prev::Real, Δs::Real)
    # equation is based on [Wende:2003, page 17]

    # F_T_prev: tractive force from previous data point
    # Δs: distance step
    ΔW = F_T_prev * Δs      # mechanical work in this step (in Ws)
    return ΔW
end #function calc_ΔW

function calc_ΔE(ΔW::Real)
    # simplified equation
    # TODO!
    # ΔW: mechanical work in this step (in Ws)
    ΔE = ΔW                 # energy consumption in this step (in Ws)
    return ΔE
end #function calc_ΔW

function calcBrakingDistance(v_start::Real, v_end::Real, a_braking::Real)
    # equation is based on [Wende:2003, page 37]

    # v_start: velocity at the start of braking (in m/s)
    # v_end: target velocity at the end of braking (in m/s)
    # a_braking: constant braking acceleration (in m/s^2)
    s_braking = (v_end^2 - v_start^2) /2 /a_braking             # braking distance (in m)
    # TODO: also possible: calc_Δs_with_Δv(v_end-v_start, a_braking, v_start)
#    return max(0.0, ceil(s_braking, digits=approxLevel))         # ceil is used to be sure that the train stops at s_exit in spite of rounding errors
    return max(0.0, ceil(s_braking, digits=approxLevel +1))         # ceil is used to be sure that the train stops at s_exit in spite of rounding errors
end #function calcBrakingDistance

function calcBrakingStartVelocity(v_end::Real, a_braking::Real, s_braking::Real)
    # equation is based on [Wende:2003, page 37]

    # v_end: target velocity at the end of braking (in m/s)
    # a_braking: constant braking acceleration (in m/s^2)
    # s_braking: braking distance (in Ws)
    v_start = sqrt(v_end^2 - 2*a_braking *s_braking)          # braking start velocity (in m/s)
#    return floor(v_start, digits=approxLevel)
    return floor(v_start, digits=approxLevel +1)
end #function calcBrakingStartVelocity

function calcBrakingAcceleration(v_start::Real, v_end::Real, s_braking::Real)
    # equation is based on [Wende:2003, page 37]

    # v_start: braking start velocity (in m/s)
    # v_end: target velocity at the end of braking (in m/s)
    # s_braking: braking distance (in Ws)
    a_braking = (v_end^2 - v_start^2) /2 /s_braking       # constant braking acceleration (in m/s^2)
    return a_braking
end #function calcBrakingAcceleration


module RatFun
using Base, ApproxFun, RecipesBase

import ApproxFun: evaluate, dimension, domain, setdomain, PiecewiseSpace, DiracSpace,
                    PointSpace, ∞, FunTypes, components
import Base: +, -, *, /, getindex, broadcast, //, numerator, denominator

export RationalFun, inv

struct RationalFun{F1<:Fun,F2<:Fun}
    p::F1
    q::F2
end

//(a::FunTypes, b::FunTypes) = RationalFun(a,b)

domain(r::RationalFun) = domain(r.p)
setdomain(r::RationalFun,d) = RationalFun(setdomain(r.p,d),setdomain(r.q,d))

function evaluate(r::RationalFun,x)
    (r.p)(x)/(r.q)(x)
end

(r::RationalFun)(x) = evaluate(r,x)

numerator(r::RationalFun) = r.p
denominator(r::RationalFun) = r.q

*(r1::RationalFun,r2::RationalFun) = RationalFun(r1.p*r2.p,r1.q*r2.q)
*(r::RationalFun,a::Union{Number,Fun}) = RationalFun(r.p*a,r.q)
*(a::Union{Number,Fun},r::RationalFun) = RationalFun(a*r.p,r.q)

Base.inv(r::RationalFun) = RationalFun(r.q,r.p)

components(r::RationalFun) = RationalFun.(components(r.p),components(r.q))

broadcast(::typeof(/), r1::RationalFun, r2::RationalFun) = r1.*inv(r2)
broadcast(::typeof(/), a, r::RationalFun) = a.*inv(r)

(/)(r1::RationalFun,r2::RationalFun) = r1*inv(r2)
(/)(a,r::RationalFun) = a*inv(r)

(/)(r::RationalFun,a) = RationalFun(r.p,r.q*a)
broadcast(::typeof(/), r::RationalFun, a) = (1/a).*r

for op = (:+,:-)
  @eval begin
      $op(r1::RationalFun,r2::RationalFun) = RationalFun($op((r1.p*r2.q),(r2.p*r1.q)),r1.q*r2.q)
      broadcast(::typeof($op),r1::RationalFun,r2::RationalFun) = $op(r1,r2)
  end
end

Base.convert(::Type{Fun},r::RationalFun) = r.p/r.q

# The padding in this function can be improved
# No support for functions with poles within the domain
function plotptsvals(r::RationalFun)
    p = r.p
    q = r.q
    plen = ncoefficients(p)
    qlen = ncoefficients(q)
    if dimension(space(p)) == ∞ && dimension(space(q)) == ∞
        p=pad(p,10plen+10*qlen+1000)
        q=pad(q,10plen+10*qlen+1000)
        r = RationalFun(p,q)
    elseif dimension(space(p)) == ∞ && dimension(space(q)) < ∞
        p=pad(p,10plen+10*dimension(space(q))+1000)
        q=pad(q,10plen+10*dimension(space(q))+1000)
        r = RationalFun(p,q)
    elseif dimension(space(q)) == ∞ && dimension(space(p)) < ∞
        q=pad(q,dimension(space(p))+10qlen+1000)
        p=pad(p,dimension(space(p))+10qlen+1000)
        r = RationalFun(p,q)
    else # both finite dimensional.
        @assert dimension(space(p)) == dimension(space(q))
        p=pad(p,dimension(space(p)))
        q=pad(q,dimension(space(q)))
        r = RationalFun(p,q)
    end
    return points(r.p),values(r.p)./values(r.q)
end

@recipe function f(g::RationalFun{Fun{S,T,V1}}) where {S,T<:Real,V1}
    plotptsvals(g)
end

@recipe function f(x::AbstractVector{T},g::RationalFun{Fun{S,T,V1}}) where {S,T<:Real,V1}
    x,g(x)
end

@recipe function f(r::RationalFun{Fun{S1,T1,V1},
                                  Fun{S2,T2,V2}}) where {S1<:PiecewiseSpace,T1<:Real,S2<:PiecewiseSpace,T2<:Real,V1,V2}
    vp = components(r.p)
    vq = components(r.q)

    if !isempty(vp)
        @series begin
            primary --> true
            RationalFun(vp[1],vq[1])
        end
    end

    for k=2:length(vp)
        @series begin
            primary := false
            RationalFun(vp[k],vq[k])
        end
    end
end

# For dirac space, we draw a dotted line extending to infinity
@recipe function f(r::RationalFun{Fun{S1,T1,V1},
                                  Fun{S2,T2,V2}}) where {S1<:DiracSpace,T1<:Real,S2<:PointSpace,T2<:Real,V1,V2}
    p = r.p
    q = r.q
    pts=space(p).points
    n=length(pts)
    ws=pad(p.coefficients./q.coefficients,length(pts))
    @series begin
        primary --> true
        ones(2)*pts[1],[0,1]*ws[1]
    end

    if length(ws) > 1
        @series begin
            primary := false
            ones(2)*pts[2:end]',[0,1]*ws[2:end]'
        end
    end

    @series begin
        primary := false
        linestyle := :dot
        ones(2)*pts',[1,2]*ws'
    end
end

# for PointSpace, we draw just a line
@recipe function f(r::RationalFun{Fun{S1,T1,V1},
                                  Fun{S2,T2,V2}}) where {S1<:PointSpace,T1<:Real,S2<:PointSpace,T2<:Real,V1,V2}
    p = r.p
    q = r.q
    pts=space(p).points
    n=length(pts)
    ws=pad(p.coefficients./q.coefficients,length(pts))

    @series begin
        primary --> true
        ones(2)*pts[1],[0,1]*ws[1]
    end

    @series begin
        primary := false
        ones(2)*pts[2:end]',[0,1]*ws[2:end]'
    end
end

end #module


using Turing, MCMCChains
using Distributions
using ParetoSmooth

samples = randn(100, 1, 1)
data = randn(50)
chain = Chains(samples)

compute_loglike(μ, data) = logpdf(Normal(μ, 1), data)
compute_loglike(μ, σ, data) = logpdf(Normal(μ, σ), data)

pll1 = pointwise_log_likelihoods(compute_loglike, chain, data)
size(pll1) |> display


module ManifoldMeasures

using KeywordCalls,
    LinearAlgebra, Manifolds, MeasureTheory, Random, SpecialFunctions, StatsFuns
using Distributions: Distributions
using Manifolds: ℝ

include("specialfunctions.jl")
include("primitive.jl")
include("stiefel/primitive.jl")
include("stiefel/matrixlangevin.jl")
include("stiefel/matrixbingham.jl")
include("stiefel/macg.jl")

# primitive
export Hausdorff, Haar

# stiefel
export MatrixLangevin, MatrixVonMisesFisher, MatrixBingham, MatrixAngularCentralGaussian

end


## L-Shaped Benders Decomposition for Stochastic Programming
## Edward J. Xu <edxu96@outlook.com>
## August 11, 2019

"MILP model for L-Shaped Benders decomposition (LSBD) in matrix form."
mutable struct ModelLSBD
  n_x
  vec_min_y
  vec_max_y
  vec_c
  vec_f
  vec_b
  mat_aCap
  mat_bCap
  solution::Union{Solution, Missing}

  function ModelLSBD(
      vec_min_y::Array{Int64,1}, vec_max_y::Array{Int64,1},
      vec_c::Array{Int64,2}, vec_b::Array{Int64,2},
      vec_f::Array{Int64,2}, mat_aCap::Array{Int64,2},
      mat_bCap::Array{Int64,2}
      )
    checkListLength(vec_min_y, vec_max_y, 'vec_min_y', 'vec_max_y')
    check_col_vec(vec_c, "vec_c")
    check_col_vec(vec_f, "vec_f")
    check_col_vec(vec_b, "vec_b")
    check_mat_match(vec_c, mat_aCap, row, col, "vec_c", "mat_aCap")
    check_mat_match(vec_f, mat_bCap, row, col, "vec_f", "mat_bCap")
    check_mat_match(mat_aCap, mat_bCap, row, row, "mat_aCap", "mat_bCap")
    check_mat_match(mat_aCap, vec_b, row, row, "mat_aCap", "vec_b")
    new(vec_c, mat_aCap, vec_f, mat_bCap, vec_b, missing)
  end
end

"First two variables are updated."
function solve_master!(obj_mas, model_mas, e1_mat, e2, opt_cut::Bool)
  if opt_cut
    @constraint(model_mas, (e1_mat * vec_y)[1] + q >= e2)
  else
    ## Add feasible cut Constraints
    @constraint(model_mas, (e1_mat * vec_y)[1] >= e2)
    @constraint(model_mas, (e1_mat * vec_y)[1] + q >= e2)
  end

  optimize!(model_mas)
  obj_mas = objective_value(model_mas)
end

"First two variables are updated."
function solve_sub!(vec_ubar, obj_sub, vec_ybar, n_constraint, vec_h , mat_t,
    mat_w, vec_c)

  model_sub = Model(with_optimizer(GLPK.Optimizer))
  @variable(model_sub, vec_u[1: n_constraint] >= 0)
  @objective(model_sub, Max, (transpose(vec_h - mat_t * vec_ybar) * vec_u)[1])
  cons_dual = @constraint(model_sub, transpose(mat_w) * vec_u .<= vec_c)
  sol = optimize!(model_sub)

  print("  Sub Problem")
  vec_ubar = value.(vec_u)
  if sol == :OPTIMAL
    bool_sub = true
    obj_sub = objective_value(model_sub)
    vec_result_x = dual(cons_dual)
  elseif sol == :DUAL_INFEASIBLE # ???
    print("Not solved optimally because the feasible set is unbounded.\n")
    bool_sub = false
    obj_sub = objective_value(model_sub)
    vec_result_x = repeat([NaN], length(vec_c))
  elseif sol == :INFEASIBLE
    print("Not solved optimally because of infeasibility. Something is ",
      "wrong.\n")
    bool_sub = false
    obj_sub = NaN
    vec_ubar = hcat(vec_ubar)
    vec_result_x = hcat(repeat([NaN], length(vec_c)))
  end

  return bool_sub, vec_result_x
end

"""
Solve the sub-problem with ray.
  First two variables are updated.
"""
function solve_ray!(vec_ubar, obj_ray, vec_ybar, n_constraint, vec_h, mat_t,
    mat_w)

  model_ray = Model(with_optimizer(GLPK.Optimizer))
  @variable(model_ray, vec_u[1: n_constraint] >= 0)
  @objective(model_ray, Max, 1)
  @constraint(model_ray, (transpose(vec_h - mat_t * vec_ybar) * vec_u)[1] == 1)
  @constraint(model_ray, transpose(mat_w) * vec_u .<= 0)
  optimize!(model_ray)

  print("  Ray Problem\n")
  vec_ubar = value.(vec_u)
  obj_ray = objective_value(model_ray)
end

"""
L-Shaped Benders Decomposition for Stochastic Programming without Integer
  Variables in Second Stage.
"""
function lshaped(; n_x, vec_min_y, vec_max_y, vec_f, vec_pi, mat_c, mat_h,
    mat3_t, mat3_w, epsilon=1e-6, timesIterationMax=100)

  println("Begin L-Shaped Benders Decomposition")
  n_y = length(vec_min_y)
  num_s = length(mat3_t[:, 1, 1])
  n_constraint = length(mat3_w[1, :, 1])
  mod_mas = set_mod_mas(n_y, vec_min_y, vec_max_y)

  let
    ub = Inf  # upper bound
    lb = - Inf  # lower bound

    # initial value of master variables
    mat_uBar = zeros(num_s, n_constraint, 1)
    vec_ybar = zeros(n_y, 1)
    vec_result_x = zeros(num_s, Int8(n_x / num_s), 1)
    obj_sub = 0
    obj_sub_s = zeros(num_s)
    timesIteration = 1

    dict_obj_mas = Dict()
    dict_q = Dict()
    dict_obj_sub = Dict()
    dict_obj_ray = Dict()
    dict_ub = Dict()
    dict_lb = Dict()

    # Must make sure "result_q == obj_sub" in the final iteration
    # while ((ub - lb > epsilon) && (timesIteration <= timesIterationMax))  !!!
    while check_whe_continue(ub, lb, epsilon, result_q, obj_sub,
        timesIteration, timesIterationMax)
      ## 1. Solve sub/ray problem for each scenario
      bool_sub_s = trues(num_s)
      for s = 1: num_s
        bool_sub_s[s], vec_result_x[s, :, :] = solve_sub!(mat_uBar[s, :, :],
          obj_sub_s[s], vec_ybar, n_constraint, mat_h[s, :, :], mat3_t[s, :, :],
          mat3_w[s, :, :], mat_c[s, :, :])
        if !(bool_sub_s[s])
          solve_ray!(mat_uBar[s, :, :], obj_sub_s[s], vec_ybar, n_constraint,
            mat_h[s, :, :], mat3_t[s, :, :], mat3_w[s, :, :])
        end
      end
      obj_sub = (transpose(vec_pi) * obj_sub_s)[1]
      ub = min(ub, obj_sub + (transpose(vec_f) * vec_ybar)[1])

      ## 2. Add optimal cut to master problem
      e1_mat = sum(vec_pi[s] * (transpose(mat_uBar[s, :]) * mat3_t[s, :, :])[1]
        for s = 1: num_s)
      e2 = sum(vec_pi[s] * (transpose(mat_uBar[s, :]) * mat_h[s, :])[1] for
        s = 1: num_s)

      solve_master!(obj_mas, e1_mat, e2, bool_sub_s[1])
      vec_ybar = value.(vec_y)

      ## 3. Compare the bounds and decide whether to stop
      lb = max(lb, obj_mas)
      result_q = value.(q)

      dict_ub[timesIteration] = ub
      dict_lb[timesIteration] = lb
      dict_obj_mas[timesIteration] = obj_mas
      dict_obj_sub[timesIteration] = obj_sub
      dict_q[timesIteration] = result_q

      if bool_sub_s[1]
        println("------------------ Result in $(timesIteration)-th Iteration with Sub ",
            "-------------------\n", "ub: $(round(ub, digits = 5)), ",
            "lb: $(round(lb, digits = 5)), obj_mas: $(round(obj_mas, digits = 5)), ",
            "q: $result_q, obj_sub: $(round(obj_sub, digits = 5)).")
      else
        println("------------------ Result in $(timesIteration)-th Iteration with Ray ",
            "-------------------\n", "ub: $(round(ub, digits = 5)), ",
            "lb: $(round(lb, digits = 5)), obj_mas: $(round(obj_mas, digits = 5)), ",
            "q: $result_q, obj_ray: $(round(obj_ray, digits = 5)).")
      end
      timesIteration += 1
    end

    println("obj_mas: $(objective_value(model_mas))")
    println("-------------------------------------------------------------------------\n",
        "------------------------------ 2/4. Result ------------------------------\n",
        "-------------------------------------------------------------------------")
    println("ub: $(round(ub, digits = 5)), lb: $(round(lb, digits = 5)), ",
        "difference: $(round(ub - lb, digits = 5))")
    println("vec_x: $vec_result_x")
    vec_result_y = value.(vec_y)
    result_q = value.(q)
    println("vec_y: $vec_result_y")
    println("result_q: $result_q")

    println("-------------------------------------------------------------------------\n",
        "------------------------- 3/4. Iteration Result -------------------------\n",
        "-------------------------------------------------------------------------")
    # Initialize
    seq_timesIteration = collect(1: (timesIteration - 1))
    vec_ub = zeros(timesIteration - 1)
    vec_lb = zeros(timesIteration - 1)
    obj_sub_sRay = zeros(timesIteration - 1)
    vec_obj_mas = zeros(timesIteration - 1)
    vec_q = zeros(timesIteration - 1)
    vec_type = repeat(["ray"], (timesIteration - 1))

    #
    for i = 1: (timesIteration - 1)
      vec_obj_mas[i] = round(dict_obj_mas[i], digits = 5)
      vec_ub[i] = round(dict_ub[i], digits = 5)
      vec_lb[i] = round(dict_lb[i], digits = 5)
      vec_q[i] = round(dict_q[i], digits = 5)
      if haskey(dict_obj_sub, i)
        vec_type[i] = "sub"
        obj_sub_sRay[i] = round(dict_obj_sub[i], digits = 5)
      else
        obj_sub_sRay[i] = round(dict_obj_ray[i], digits = 5)
      end
    end
    table_iterationResult = hcat(seq_timesIteration, vec_ub, vec_lb,
      vec_obj_mas, vec_q, vec_type, obj_sub_sRay)
    pretty_table(table_iterationResult,
      ["Seq", "ub", "lb", "obj_mas", "q", "sub/ray", "obj_sub/ray"],
      compact; alignment=:l)
  end

  println("-------------------------------------------------------------------------\n",
      "-------------------------- 4/4. Nominal Ending --------------------------\n",
      "-------------------------------------------------------------------------\n")
end


### gls.jl ---  Perform Generalised Lomb–Scargle periodogram
#
# Copyright (C) 2017 Mosè Giordano.
#
# Maintainer: Mosè Giordano <mose AT gnu DOT org>
# Keywords: periodogram, lomb scargle
#
# This file is a part of LombScargle.jl.
#
# License is BSD 3-clause "New" or "Revised".
#
### Code:

struct GLSPlan{T,A,B<:AbstractVector{T},C,D,E,F,G} <: PeriodogramPlan
    times::A
    signal::B
    freq::C
    sumw::D
    w::E
    y::B
    YY::T
    noise::F
    norm::Symbol
    P::G
end

struct GLSPlan_fit_mean{T,A,B<:AbstractVector{T},C,D,E,F,G,H} <: PeriodogramPlan
    times::A
    signal::B
    freq::C
    sumw::D
    w::E
    y::B
    Y::F
    YY::T
    noise::G
    norm::Symbol
    P::H
end

# Generalised Lomb–Scargle algorithm: this takes into account uncertainties and
# fit the mean of the signal.  This is implemented following the recipe by
# * Zechmeister, M., Kürster, M. 2009, A&A, 496, 577  (URL:
#   http://dx.doi.org/10.1051/0004-6361:200811296,
#   Bibcode: http://adsabs.harvard.edu/abs/2009A%26A...496..577Z)
# In addition, some tricks suggested by
# * Press, W. H., Rybicki, G. B. 1989, ApJ, 338, 277 (URL:
#   http://dx.doi.org/10.1086/167197,
#   Bibcode: http://adsabs.harvard.edu/abs/1989ApJ...338..277P)
# to make computation faster are adopted.

function _generalised_lombscargle!(P, freqs, times, y, w, Y, YY, nil)
    @inbounds Threads.@threads for n in eachindex(freqs)
        ω = freqs[n] * 2 * pi
        # Find τ for current angular frequency
        C = S = CS = CC = nil
        @inbounds for i in eachindex(times)
            W    = w[i]
            s, c = sincos(ω * times[i])
            CS  += W*c*s
            CC  += W*c*c
            C   += W*c
            S   += W*s
        end
        CS -= C*S
        SS  = 1 - CC - S*S
        CC -= C*C
        ωτ   = atan(2CS, CC - SS) / 2
        # Now we can compute the power
        YC_τ = YS_τ = CC_τ = nil
        @inbounds for i in eachindex(times)
            W     = w[i]
            s, c  = sincos(ω*times[i] - ωτ)
            YC_τ += W*y[i]*c
            YS_τ += W*y[i]*s
            CC_τ += W*c*c
        end
        # "C_τ" and "S_τ" are computed following equation (7) of Press &
        # Rybicki, this formula simply comes from angle difference trigonometric
        # identities.
        sin_ωτ, cos_ωτ = sincos(ωτ)
        C_τ    = C*cos_ωτ + S*sin_ωτ
        S_τ    = S*cos_ωτ - C*sin_ωτ
        YC_τ  -= Y*C_τ
        YS_τ  -= Y*S_τ
        SS_τ   = 1 - CC_τ - S_τ*S_τ
        CC_τ  -= C_τ*C_τ
        P[n] = (abs2(YC_τ)/CC_τ + abs2(YS_τ)/SS_τ)/YY
    end
    return P
end

function _generalised_lombscargle!(P, freqs, times, y, w, YY, nil)
    @inbounds Threads.@threads for n in eachindex(freqs)
        ω = freqs[n] * 2 * pi
        # Find τ for current angular frequency
        C = S = CS = CC = nil
        @inbounds for i in eachindex(times)
            W    = w[i]
            s, c = sincos(ω * times[i])
            CS  += W*c*s
            CC  += W*c*c
        end
        ωτ   = atan(2CS, 2CC - 1) / 2
        # Now we can compute the power
        YC_τ = YS_τ = CC_τ = nil
        @inbounds for i in eachindex(times)
            W     = w[i]
            s, c  = sincos(ω*times[i] - ωτ)
            YC_τ += W*y[i]*c
            YS_τ += W*y[i]*s
            CC_τ += W*c*c
        end
        # P[n] should be (abs2(YC_τ)/CC_τ + abs2(YS_τ)/SS_τ)/YY, but we guard the values of
        # CC_τ and SS_τ in case they're zeros (note that SS_τ = 1 - CC_τ).
        frac_C = ifelse(iszero(CC_τ), nil, abs2(YC_τ) / CC_τ)
        frac_S = ifelse(CC_τ == 1,    nil, abs2(YS_τ) / (1 - CC_τ))
        P[n] = (frac_C + frac_S)/YY
    end
    return P
end

_periodogram!(P::AbstractVector, times, p::GLSPlan_fit_mean) =
    _generalised_lombscargle!(P, p.freq, times, p.y, p.w, p.Y, p.YY, zero(eltype(p.P)))
_periodogram!(p::GLSPlan_fit_mean) = _periodogram!(p.P, p.times, p)
_periodogram!(P::AbstractVector, times, p::GLSPlan) =
    _generalised_lombscargle!(P, p.freq, times, p.y, p.w, p.YY, zero(eltype(p.P)))
_periodogram!(p::GLSPlan) = _periodogram!(p.P, p.times, p)


using Plots,DelimitedFiles,TriplotRecipes

f(x,y) = exp(0.1*sin(5.1*x + -6.2*y) + 0.3*cos(4.3*x + 3.4*y))

x,y = eachcol(readdlm(joinpath(@__DIR__,"dolphin.xy")))
t = readdlm(joinpath(@__DIR__,"dolphin.t"),Int)' .+ 1
z = f.(x,y)

plot(aspect_ratio=:equal,size=(800,720))
tripcolor!(x,y,z,t,color=:magma)
trimesh!(x,y,t,fillalpha=0.0,linecolor=:white)
savefig("dolphin.png")


@testset "dijkstra" begin
    (A,xy,labels) = load_matrix_network_metadata("airports")
    A = -A; # fix funny encoding of airport data
    lax = 247; rst = 355

    (d,pred) = dijkstra(A,lax)
    @test maximum(d) == 540
    @test pred[end] == lax

    (d,pred) = dijkstra(MatrixNetwork(A),lax)
    @test maximum(d) == 540
    @test pred[end] == lax

    @test_throws ErrorException dijkstra(-sparse(1.0I,2,2), 1)
end


function decompose_RX_on_q3()

    println(">>>>> RX Gate on third qubit using U3Gate <<<<<")

    return Dict{String, Any}(

    "num_qubits" => 3,
    "maximum_depth" => 3,
    "elementary_gates" => ["U3_3", "Identity"],
    "target_gate" => QCOpt.kron_single_qubit_gate(3, QCOpt.RXGate(π/4), "q3"),
    "objective" => "minimize_depth",
    "decomposition_type" => "exact",

    "U3_θ_discretization" => [0, π/4],
    "U3_ϕ_discretization" => [0, -π/2],
    "U3_λ_discretization" => [0, π/2],
    )

end

function decompose_toffoli()

    function toffoli_circuit()
        # [(depth, gate)]
        return [(1, "T_1"),
                (2, "T_2"),
                (3, "H_3"),
                (4, "CNot_2_3"),
                (5, "Tdagger_3"),
                (6, "CNot_1_3"),
                (7, "T_3"),
                (8, "CNot_2_3"),
                (9, "Tdagger_3"),
                (10, "CNot_1_3"),
                (11, "T_3"),
                (12, "H_3"),
                (13, "CNot_1_2"),
                (14, "Tdagger_2"),
                (15, "CNot_1_2")
                ]
    end

    println(">>>>> Toffoli gate <<<<<")

    return Dict{String, Any}(

    "num_qubits" => 3,
    "maximum_depth" => 15,
    "elementary_gates" => ["T_1", "T_2", "T_3", "H_3", "CNot_1_2", "CNot_1_3", "CNot_2_3", "Tdagger_1", "Tdagger_2", "Tdagger_3", "Identity"],
    "target_gate" => QCOpt.ToffoliGate(),
    "objective" => "minimize_depth",
    "decomposition_type" => "exact",

    "input_circuit" => toffoli_circuit(),
    "set_cnot_lower_bound" => 6,
    "set_cnot_upper_bound" => 6,
    )

end

function decompose_toffoli_using_kronecker()

    println(">>>>> Toffoli gate using Kronecker <<<<<")

    return Dict{String, Any}(

    "num_qubits" => 3,
    "maximum_depth" => 12,
    "elementary_gates" => ["T_3", "H_3", "CNot_1_2", "CNot_1_3", "CNot_2_3", "Tdagger_3", "I_1xT_2xT_3", "CNot_1_2xH_3", "T_1xTdagger_2xI_3", "Identity"],
    "target_gate" => QCOpt.ToffoliGate(),
    "objective" => "minimize_depth",
    "decomposition_type" => "exact",

    "set_cnot_lower_bound" => 6,
    )

end

function decompose_toffoli_with_controlled_gates()

    # Reference: https://doi.org/10.1109/TCAD.2005.858352
    println(">>>>> Toffoli gate with controlled gates <<<<<")

    return Dict{String, Any}(

    "num_qubits" => 3,
    "maximum_depth" => 5,
    "elementary_gates" => ["CV_1_3", "CV_2_3", "CV_1_2", "CVdagger_1_3", "CVdagger_2_3", "CVdagger_1_2", "CNot_2_1", "CNot_1_2", "Identity"],
    # "elementary_gates" => ["CV_1_3", "CV_2_3", "CVdagger_1_3", "CNot_1_2", "CNot_2_1", "Identity"],
    "target_gate" => QCOpt.ToffoliGate(),
    "objective" => "minimize_depth",
    "decomposition_type" => "exact",
    )

end

function decompose_CNot_1_3()

    return Dict{String, Any}(

    "num_qubits" => 3,
    "maximum_depth" => 8,
    # "elementary_gates" => ["CNot_1_2", "CNot_2_3", "Identity"],
    "elementary_gates" => ["H_1", "H_3", "H_2", "CNot_2_1", "CNot_3_2", "Identity"],
    "target_gate" => QCOpt.get_full_sized_gate("CNot_1_3", 3),
    "objective" => "minimize_depth",
    )

end

function decompose_FredkinGate()

    return Dict{String, Any}(
    "num_qubits" => 3,
    "maximum_depth" => 7,
    # Reference: https://doi.org/10.1103/PhysRevA.53.2855
    "elementary_gates" => ["CV_1_2", "CV_2_3", "CV_1_3", "CVdagger_1_2", "CVdagger_2_3", "CVdagger_1_3", "CNot_1_2", "CNot_3_2", "CNot_2_3", "CNot_1_3", "Identity"],
    "target_gate" => QCOpt.CSwapGate(), #also Fredkin
    "objective" => "minimize_depth"
    )

end

function decompose_toffoli_left()

    # Reference: https://arxiv.org/pdf/0803.2316.pdf

    function target_gate()
        H_3 = QCOpt.get_full_sized_gate("H_3", 3)
        Tdagger_3 = QCOpt.get_full_sized_gate("Tdagger_3", 3)
        T_3 = QCOpt.get_full_sized_gate("T_3", 3)
        cnot_23 = QCOpt.get_full_sized_gate("CNot_2_3", 3)
        CNot_1_3 = QCOpt.get_full_sized_gate("CNot_1_3", 3)

        return H_3 * cnot_23 * Tdagger_3 * CNot_1_3 * T_3 * cnot_23 * Tdagger_3
    end

    return Dict{String, Any}(

    "num_qubits" => 3,
    "maximum_depth" => 7,
    "elementary_gates" => ["H_3", "T_1", "T_2", "T_3", "Tdagger_1", "Tdagger_2", "Tdagger_3", "CNot_1_2", "CNot_2_3", "CNot_1_3", "Identity"],
    "target_gate" => target_gate(),
    "objective" => "minimize_depth"
    )

end

function decompose_toffoli_right()

    # Reference: https://arxiv.org/pdf/0803.2316.pdf

    function target_gate()
        H_3 = QCOpt.get_full_sized_gate("H_3", 3)
        Tdagger_2 = QCOpt.get_full_sized_gate("Tdagger_2", 3)
        T_1 = QCOpt.get_full_sized_gate("T_1", 3)
        T_2 = QCOpt.get_full_sized_gate("T_2", 3)
        T_3 = QCOpt.get_full_sized_gate("T_3", 3)
        CNot_1_2 = QCOpt.get_full_sized_gate("CNot_1_2", 3)
        CNot_1_3 = QCOpt.get_full_sized_gate("CNot_1_3", 3)

        return CNot_1_3 * T_2 * T_3 * CNot_1_2 * H_3 * T_1 * Tdagger_2 * CNot_1_2
    end

    return Dict{String, Any}(

    "num_qubits" => 3,
    "maximum_depth" => 8,
    "elementary_gates" => ["H_3", "T_1", "T_2", "T_3", "Tdagger_1", "Tdagger_2", "Tdagger_3", "CNot_1_2", "CNot_2_3", "CNot_1_3", "Identity"],
    "target_gate" => target_gate(),
    "objective" => "minimize_depth"
    )

end

function decompose_miller()
    # Reference: https://doi.org/10.1109/TCAD.2005.858352

    function target_gate()
        CV_1_3 = QCOpt.get_full_sized_gate("CV_1_3", 3)
        CV_2_3 = QCOpt.get_full_sized_gate("CV_2_3", 3)
        CVdagger_2_3 = QCOpt.get_full_sized_gate("CVdagger_2_3", 3)
        CNot_1_2 = QCOpt.get_full_sized_gate("CNot_1_2", 3)
        CNot_3_1 = QCOpt.get_full_sized_gate("CNot_3_1", 3)
        CNot_3_2 = QCOpt.get_full_sized_gate("CNot_3_2", 3)

        return CNot_3_1 * CNot_3_2 * CV_2_3 * CNot_1_2 * CVdagger_2_3 * CV_1_3 * CNot_3_1 * CNot_1_2
    end

    return Dict{String, Any}(

    "num_qubits" => 3,
    "maximum_depth" => 8,
    "elementary_gates" => ["CV_1_3", "CV_2_3", "CVdagger_2_3", "CNot_1_2", "CNot_3_1", "CNot_3_2", "Identity"],
    "target_gate" => target_gate(),
    "objective" => "minimize_depth"
    )

end

function decompose_relative_toffoli()
    #Reference: https://arxiv.org/pdf/1508.03273.pdf

    return Dict{String, Any}(

        "num_qubits" => 3,
        "maximum_depth" => 9,
        "elementary_gates" => ["H_3", "T_3", "Tdagger_3", "CNot_1_2", "CNot_2_3", "CNot_1_3", "Identity"],
        "target_gate" => QCOpt.RCCXGate(),
        "objective" => "minimize_depth",

        "set_cnot_upper_bound" => 3
        )

end

function decompose_margolus()
    #Reference: https://arxiv.org/pdf/quant-ph/0312225.pdf

    return Dict{String, Any}(

        "num_qubits" => 3,
        "maximum_depth" => 7,
        "elementary_gates" => ["RY_3", "CNot_1_2", "CNot_2_3", "CNot_1_3", "Identity"],
        "target_gate" => QCOpt.MargolusGate(),
        "objective" => "minimize_depth",

        "RY_discretization" => -π:π/4:π,
        # "set_cnot_lower_bound" => 3,
        # "set_cnot_upper_bound" => 3
        )

end

function decompose_CiSwap()
    #Reference: https://doi.org/10.1103/PhysRevResearch.2.033097

    return Dict{String, Any}(

        "num_qubits" => 3,
        "maximum_depth" => 12,
        "elementary_gates" => ["H_3", "T_3", "Tdagger_3", "CNot_3_2", "CNot_2_3", "CNot_1_3", "Identity"],
        "target_gate" => QCOpt.CiSwapGate(),
        "objective" => "minimize_depth",
        )

end


using OrdinaryDiffEq, Flux, DiffEqSensitivity, DiffEqCallbacks, Test
using Random

Random.seed!(1234)

u0 = Float32[2.; 0.]
datasize = 100
tspan = (0.0f0,10.5f0)
dosetimes = [1.0,2.0,4.0,8.0]

function affect!(integrator)
    integrator.u = integrator.u.+1
end
cb_ = PresetTimeCallback(dosetimes,affect!,save_positions=(false,false))
function trueODEfunc(du,u,p,t)
    du .= -u
end
t = range(tspan[1],tspan[2],length=datasize)

prob = ODEProblem(trueODEfunc,u0,tspan)
ode_data = Array(solve(prob,Tsit5(),callback=cb_,saveat=t))
dudt2 = Chain(Dense(2,50,tanh),
             Dense(50,2))
p,re = Flux.destructure(dudt2) # use this p as the initial condition!

function dudt(du,u,p,t)
    du[1:2] .= -u[1:2]
    du[3:end] .= re(p)(u[1:2]) #re(p)(u[3:end])
end
z0 = Float32[u0;u0]
prob = ODEProblem(dudt,z0,tspan)

affect!(integrator) = integrator.u[1:2] .= integrator.u[3:end]
cb = PresetTimeCallback(dosetimes,affect!,save_positions=(false,false))

function predict_n_ode()
    _prob = remake(prob,p=p)
    Array(solve(_prob,Tsit5(),u0=z0,p=p,callback=cb,saveat=t,sensealg=ReverseDiffAdjoint()))[1:2,:]
    #Array(solve(prob,Tsit5(),u0=z0,p=p,saveat=t))[1:2,:]
end

function loss_n_ode()
    pred = predict_n_ode()
    loss = sum(abs2,ode_data .- pred)
    loss
end
loss_n_ode() # n_ode.p stores the initial parameters of the neural ODE

cba = function (;doplot=false) #callback function to observe training
  pred = predict_n_ode()
  display(sum(abs2,ode_data .- pred))
  # plot current prediction against data
  #pl = scatter(t,ode_data[1,:],label="data")
  #scatter!(pl,t,pred[1,:],label="prediction")
  #display(plot(pl))
  return false
end
cba()

ps = Flux.params(p)
data = Iterators.repeated((), 200)
Flux.train!(loss_n_ode, ps, data, ADAM(0.05), cb = cba)
@test loss_n_ode() < 0.4