@testset "singular case" begin
# From Nocedal & Wright, p. 288-289

# Jacobian is singular at the starting point.
# Used to test the behavior of algorithms in that context.

function f_sinj!(F, x)
    F[1] = x[1]
    F[2] = 10*x[1]/(x[1]+convert(eltype(x), 0.1))+2*x[2]^2
end

function g_sinj!(J, x)
    J[1, 1] = 1
    J[1, 2] = 0
    J[2, 1] = 1/(x[1]+convert(eltype(x), 0.1))^2
    J[2, 2] = 4*x[2]
end

df = OnceDifferentiable(f_sinj!, g_sinj!, [3.0, 0.0], [3.0, 0.0])
df32 = OnceDifferentiable(f_sinj!, g_sinj!, [3.0f0, 0.0f0], [3.0f0, 0.0f0])

# Test disabled, not stable across runs
#r = nlsolve(df, [ 3.0; 0], method = :newton, ftol = 1e-5)
#@assert converged(r)
#@assert norm(r.zero) < 1e-5

r = nlsolve(df, [ 3.0; 0.0], method = :trust_region)
@assert converged(r)
@assert norm(r.zero) < 1e-6

r = nlsolve(df32, [3.0f0; 0.0f0], method = :trust_region)
@assert converged(r)
@assert norm(r.zero) < 1e-6

let a = rand(10)
    A = a*a'
    global f_let!, g_let!
    function f_let!(fvec, x)
        copyto!(fvec, A*x)
    end

    function g_let!(fjac, x)
        copyto!(fjac, A)
    end
end

df = OnceDifferentiable(f_let!, g_let!, rand(10), rand(10))
r = nlsolve(df, rand(10), method = :trust_region)
end


using StatsBase

# Distance function -------------------------------------------------

"""
    euclidean(a::AbstractVector, b::AbstractVector)

Calculate euclidean distance from two vectors. √∑(aᵢ - bᵢ)²
"""
function euclidean(a::AbstractVector{T},
                   b::AbstractVector{T}) where {T<:AbstractFloat}
    @assert length(a) == length(b)

    # euclidean(a, b) = √∑(aᵢ- bᵢ)²
    s = zero(T)
    @simd for i in eachindex(a)
        @inbounds s += (a[i] - b[i])^2
    end
    return √s
end

# KnnClassifier -----------------------------------------------------
struct KnnClassifier{T<:Real, S}
    X::AbstractMatrix{T}
    y::AbstractVector{S}
    K::Int

    # Internal constructor
    function KnnClassifier(X::AbstractMatrix{T}, y::AbstractVector{S}, K::Int) where {T<:Real, S}
        @assert size(X, 1) == size(y, 1)
        return new{T, S}(X, y, K)
    end
end

# External constructor
function KnnClassifier(Xv::AbstractVector{T}, y::AbstractVector{S}, K::Int) where {T<:Real, S}
    X = reshape(Xv, length(Xv), 1)
    return KnnClassifier(X, y, K)
end

function predict(model::KnnClassifier{T, S}, input::AbstractVector{T}) where {T<:Real, S}
    nl, nc = size(model.X, 1), size(model.X, 2)

    @assert size(input, 1) == nc

    # calculate all the distances
    dists = Vector{T}(undef, nl)
    for i in axes(model.X, 1)
        dists[i] = euclidean(input, model.X[i, :])
    end

    # sort index of distances and select the k nearest indexes
    idx = sortperm(dists)[1:model.K]

    # select the KNearestNeighbours
    kneighbours = model.y[idx]

    # get prediction
    return mode(kneighbours)
end

function predict(model::KnnClassifier{T, S}, input::AbstractMatrix{T}) where {T<:Real, S}
    nl, nc = size(model.X, 1), size(model.X, 2)
    nli = size(input, 1)

    @assert size(input, 2) == nc

    results = Vector{S}(undef, nli)
    for j in axes(input, 1)

        # calculate all the distances
        dists = Vector{T}(undef, nl)
        for i in axes(model.X, 1)
            dists[i] = euclidean(input[j, :], model.X[i, :])
        end

        # sort index of distances and select the k nearest indexes
        idx = sortperm(dists)[1:model.K]

        # select the KNearestNeighbours
        kneighbours = model.y[idx]

        # get prediction
        results[j] = mode(kneighbours)
    end

    return results
end

# KnnRegression -----------------------------------------------------
struct KnnRegression{T<:Real, S<:Real}
    X::AbstractMatrix{T}
    y::AbstractVector{S}
    K::Int

    # Internal constructor
    function KnnRegression(X::AbstractMatrix{T}, y::AbstractVector{S}, K::Int) where {T<:Real, S<:Real}
        @assert size(X, 1) == size(y, 1)
        return new{T, S}(X, y, K)
    end
end

# External constructor
function KnnRegression(Xv::AbstractVector{T}, y::AbstractVector{S}, K::Int) where {T<:Real, S<:Real}
    X = reshape(Xv, length(Xv), 1)
    return new{T, S}(X, y, K)
end

function predict(model::KnnRegression{T, S}, input::AbstractVector{T}, w::Bool=false) where {T<:Real, S<:Real}
    nl, nc = size(model.X, 1), size(model.X, 2)

    @assert size(input, 1) == nc

    # calculate all the distances
    dists = Vector{T}(undef, nl)
    for i in axes(model.X, 1)
        dists[i] = euclidean(input, model.X[i, :])
    end

    # sort index of distances and select the k nearest indexes
    idx = sortperm(dists)[1:model.K]

    # select the KNearestNeighbours
    y_kneighbours = model.y[idx]

    # get prediction by mean or weighted mean
    if w
        k_min_dists = sort(dists)[1:model.K]
        return mean(y_kneighbours, weights(k_min_dists))
    else
        return mean(y_kneighbours)
    end

end

function predict(model::KnnRegression{T, S}, input::AbstractMatrix{T}, w::Bool=false) where {T<:Real, S<:Real}
    nl, nc = size(model.X, 1), size(model.X, 2)
    nli = size(input, 1)

    @assert size(input, 2) == nc

    results = Vector{S}(undef, nli)
    for j in axes(input, 1)

        # calculate all the distances
        dists = Vector{T}(undef, nl)
        for i in axes(model.X, 1)
            dists[i] = euclidean(input[j, :], model.X[i, :])
        end

        # sort index of distances and select the k nearest indexes
        idx = sortperm(dists)[1:model.K]

        # select the KNearestNeighbours
        y_kneighbours = model.y[idx]

        # get prediction by mean or weighted mean
        if w
            k_min_dists = sort(dists)[1:model.K]
            results[j] = mean(y_kneighbours, weights(k_min_dists))
        else
            results[j] = mean(y_kneighbours)
        end
    end

    return results
end


using ADCME
using Plots
using PyCall
@pyimport scipy.interpolate as si
pyplot()

################# --- functions --- #################

# flux function non-convex
function f_w(sw, M=1)
    # u := saturation of water
    # M := ratio of viscosities
    return sw^2 / (sw^2 + ((1 - sw)^2 / M))
end

# gradient of non-convex flux
function f_w_grad(sw, M=1)
    return 2 * M * sw * (1-sw) / ((M+1) * sw^2 - 2 * sw + 1)^2
end

# analytical Buckley-Leverett solution to sw (sw := Saturation of water)
function sw(x, t, M=1)
    # x := spacial coord
    # t := time nodes
    # M := viscosity ratio μ_oil/μ_water
    s = range(0, 1, length=length(x))[end:-1:1]
    x_s = @. f_w_grad(s)

    # compute sasturation before shock
    dif = x_s[2] - x_s[1]
    for i in 2:length(x)
        if x_s[i+1] - x_s[i] < dif
            x_s = x_s[1:i]
            @goto stop
        else
            dif = x_s[i+1] - x_s[i]
        end
    end
    @label stop

    x_s_times = []
    [push!(x_s_times, x_s * k) for k in t ]

    sw_ = zeros(Nx*Nt)

    # Perform the interpolation
    for k in 1:length(t)
        bench = x_s_times[k][end]
        tmp = si.griddata(Array(x_s_times[k]), Array(s), Array(x))
        sw_[(k-1)*Nx+1:k*Nx] = @. ifelse(x <= bench, tmp, 0.0)
        convert(Array{Float64,1}, x_s_times[k])
    end

    return sw_
end

### --- numerical parameters --- ###

X = 1 # spacial domain from 0 to 1
T = 1 # time domain from 0 to 1

dx = 0.01
dt = 0.01
Nx = round(Int, X/dx + 1) # spacial nodes
Nt = round(Int, T/dt) # time nodes

# simple coordinates
x = range(0, X, step=dx) # x coords
t = range(dt, T, step=dt) # t steps

# 1D array of all possible coordinte pairs
# will also be all colocation points for PDE regularization
x_star = repeat(x, Nt)
t_star = collect(Iterators.flatten(repeat(t', Nx)))

################# --- anal solution --- #################

sw_anal = sw(x, t)
sw_anal_tf = constant(sw_anal)

################# --- neural net --- #################

# training input for neural network
x_star_tf = constant(x_star)
t_star_tf = constant(t_star)
#train_input = constant(hcat(train_XX, train_TT))

# layers
config = [20, 20, 20, 20, 20, 20, 20, 20, 1]
# xavier initialization
θ = Variable(fc_init([2; config]))

# neural network function
sw_nn = squeeze(fc(hcat(x_star_tf, t_star_tf), config, θ)) + 1

# derivatives for PDE
sw_t_nn = tf.gradients(sw_nn, t_star_tf)[1]
sw_x_nn = tf.gradients(sw_nn, x_star_tf)[1]
sw_xx_nn = tf.gradients(sw_x_nn, x_star_tf)[1]
f_w_nn = f_w(sw_nn)
f_w_x_nn = tf.gradients(f_w_nn, x_star_tf)[1]

### --- loss --- ###

# loss funciton: (IC/BC - analitical_solution) + (sw_t + f_w(sw) - ϵ * sw_xx)
# sw_t       := time derivative of water saturation
# f_w(sw)    := flux funciton dependend on water saturation
# ϵ         := 1/Pe
# sw_xx      := second spatial derivative of water saturation

ϵ = constant(2.5*10e-3)

# Loss with diffusion
loss = (sum((sw_nn[1:Nx] - sw_anal_tf[1:Nx])^2) / Nx +
        sum((sw_nn[Nx+1:Nx:end] - sw_anal_tf[Nx+1:Nx:end])^2)/ (Nt - 1) +
        sum((sw_t_nn + f_w_x_nn - ϵ * sw_xx_nn)^2) / (Nx * Nt))
# Loss without diffusion
#= loss = (sum((u_nn[1:Nx] - u_anal_train[1:Nx])^2) / Nx +
        sum((u_nn[Nx+1:Nx:end] - u_anal_train[Nx+1:Nx:end])^2)/ (Nt - 1) +
        sum((u_nn[Nx*2-1:Nx:end] - u_anal_train[Nx*2-1:Nx:end])^2)/ (Nt - 1) +
        sum((u_t_nn + f_w_x_nn)^2) / (Nx * Nt)) =#

################# --- training --- #################

### optimising

losses = []
# callback function for optimisation
i = 0
function cb()
    global i
    if i % 10 == 0
        println("Epoch: $i")
    end
    #push!(losses, loss)
    i += 1
end

sess = Session(); init(sess)

# Scipy optimizer
opt = ScipyOptimizerInterface(loss; method="L-BFGS-B", options=Dict("maxiter"=> 3000, "ftol"=>1e-12, "gtol"=>1e-12))
ScipyOptimizerMinimize(sess, opt; loss_callback=cb)

#BFGS!(sess, loss)

### --- predict --- ###

sw_test_nn = run(sess, squeeze(fc(hcat(x_star_tf, t_star_tf), config, θ))+1)

error = sum(@. sqrt((sw_test_nn - sw_anal).^2)) / sum(@. sqrt(sw_anal.^2))
println("\nFinal Loss:\t", run(sess, loss))
println("sw error:\t", error)

### --- Plots --- ###

println("\nPrinting images...")
close("all")

# times at which to plot
t_plot = [0.25, 0.5, 0.75]
sw_anal = reshape(sw_anal, (Nx, Nt))
sw_test_nn = reshape(sw_test_nn, (Nx, Nt))

for i in 1:length(t_plot)
    idx = round(Int, t_plot[i] / dt * Nt*dt)
    println("0.$idx")
    plot(x, sw_anal[:, idx], label="anal")
    plot!(x, sw_test_nn[:, idx], label="NN", line=:dash)
    name = "./results/0.$idx.png"
    savefig(name)
end

### --- Save NN --- ###

println("saving NN...")
# note if you run script multiple times -> number epochs adds up
name = "./model.mat"
save(sess, name)

sess.close()


using Plots

backtrg = [[0.0; 0.0],[10.0; 0.0],[5.0; 8.66],[0.0; 0.0]]

lasttringles = [[[0.0; 0.0],[10.0; 0.0],[5.0; 8.66],[0.0; 0.0]]]
for i in 1:8
    appendTringles = []
    for points in lasttringles
        delta = points[1]
        Trgpoints = (points - fill(delta,4)) / 2
        movecor1 = Trgpoints[2]
        movecor2 = Trgpoints[3]
        Trgpoints += fill(delta,4)
        append!(appendTringles, [Trgpoints])
        append!(appendTringles, [Trgpoints + fill(movecor1, 4)])
        append!(appendTringles, [Trgpoints + fill(movecor2, 4)])
    end
    lasttringles = appendTringles
end

plot(hcat(backtrg...)[1,:], hcat(backtrg...)[2,:], legend = false, border=:none, dpi = 500)

for points in lasttringles
    plot!(hcat(points...)[1,:], hcat(points...)[2,:], legend = false, fill = (0, 0.5, :green), dpi = 500)
end
plot!(dpi = 500)
savefig("C:\\Users\\Yaghoub\\Documents\\GitHub\\A-few-fractals-in-Julia\\Fractals-Fig\\ST.png")


using DrWatson
@quickactivate "StatReth"

# %%
using Distributions
using Plots
using StatsPlots

# %% 4.1
pos = [sum(rand(Uniform(-1, 1), 16)) for _ in 1:1000]

# %% 4.2
prod(1 .+ rand(Uniform(0.0, 0.1), 12))

# %% 4.3
growth = [prod(1 .+ rand(Uniform(0.0, 0.1), 12)) for _ in 1:10_000]
density(growth)

# %% 4.4
big = [prod(1 .+ rand(Uniform(0.0, 0.5), 12)) for _ in 1:10_000]
small = [prod(1 .+ rand(Uniform(0.0, 0.01), 12)) for _ in 1:10_000]

# %% 4.5
log_big = [prod(1 .+ log.(rand(Uniform(0.0, 0.5), 12))) for _ in 1:10_000]
density(log_big)
density(log_big, xscale = :log10)


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

#####
##### Some multivariate test problems
#####

####
#### generic code
####

"""
A type for test functions with uniform bounds.
All subtypes (test functions) support:
1. being called with a vector of arbitrary length
2. `minimum_location`,
3. `lower_bounds`, `upper_bounds` (defined via `lu_bounds`),
4. having a global minimum *value* of `0`.
See [`TEST_FUNCTIONS`](@ref).
"""
abstract type UniformBounds end

lower_bounds(f::UniformBounds, n::Integer) = fill(Float64(lu_bounds(f)[1]), n)

upper_bounds(f::UniformBounds, n::Integer) = fill(Float64(lu_bounds(f)[2]), n)

####
#### shifted quadratic
####

"A test function with global minimum at [a, …]. For sanity checks."
struct ShiftedQuadratic{T <: Real} <: UniformBounds
    a::T
end

(f::ShiftedQuadratic)(x) = (a = f.a; sum(x -> abs2(x - a), x))

minimum_location(f::ShiftedQuadratic, n::Integer) = fill(f.a, n)

lu_bounds(f::ShiftedQuadratic) = (f.a - 50, f.a + 100)

const SHIFTED_QUADRATIC = ShiftedQuadratic(0.5)

####
#### Griewank
####

"""
The Griewank problem.
From Ali, Khompatraporn, and Zabinsy (2005, p 559).
"""
struct Griewank{T <: Real} <: UniformBounds
    a::T
end

function (f::Griewank)(x)
    sum(abs2, x) / f.a - prod(((i, x),) -> cos(x / √i), enumerate(x)) + 1
end

minimum_location(::Griewank, n::Integer) = zeros(n)

lu_bounds(::Griewank) = (-50, 100)

const GRIEWANK = Griewank(200.0)

####
#### Levy Montalvo 2.
####

"""
Levi and Montalvo 2 problem.
From Ali, Khompatraporn, and Zabinsy (2005, p 662).
"""
struct LevyMontalvo2 <: UniformBounds end

function (::LevyMontalvo2)(x)
    xn = last(x)
    0.1 * abs2(sinpi(3 * first(x))) + abs2(xn -  1) * (1 + abs2(sinpi(2 * xn))) +
        sum(@. abs2($(x[1:(end - 1)]) - 1) * (1 + abs2(sinpi(3 * $(x[2:end])))))
end

minimum_location(::LevyMontalvo2, n::Integer) = ones(n)

lu_bounds(::LevyMontalvo2) = (-10, 15)

const LEVY_MONTALVO2 = LevyMontalvo2()

####
#### Rastrigin
####

"""
Rastrigin problem.
From Ali, Khompatraporn, and Zabinsy (2005, p 665).
"""
struct Rastrigin <: UniformBounds end

function (::Rastrigin)(x)
    10 * length(x) + sum(@. abs2(x) - 10 * cospi(2 * x))
end

minimum_location(::Rastrigin, n::Integer) = zeros(n)

lu_bounds(::Rastrigin) = (-5.12, 5.12)

const RASTRIGIN = Rastrigin()

####
#### Rosenbrock
####

"""
Rosenbrock problem.
From Ali, Khompatraporn, and Zabinsy (2005, p 666).
"""
struct Rosenbrock <: UniformBounds end

function (::Rosenbrock)(x)
    x1 = x[1:(end - 1)]
    x2 = x[2:end]
    sum(@. 100 * abs2(x2 - abs2(x1)) + abs2(x1 - 1))
end

minimum_location(::Rosenbrock, n::Integer) = ones(n)

lu_bounds(::Rosenbrock) = (-30, 30)

const ROSENBROCK = Rosenbrock()

####
#### helper code
####

"A tuple of all test functions."
const TEST_FUNCTIONS = (SHIFTED_QUADRATIC, GRIEWANK, LEVY_MONTALVO2, RASTRIGIN, ROSENBROCK)


@testset "Test composition order" begin
  Qx, x = FlintQQ["x"]
  K, a = NumberField(x^6 + 108, "a")
  A = automorphisms(K)
  for f in A
    for g in A
      @test g(f(a)) == compose(f, g)(a)
      @test g(f(a)) == (f * g)(a)
    end
  end
end


using  LinearAlgebra
const  IdentityMatrix = UniformScaling{Bool}

function clip(s::Real, min_thresh::Real, max_thresh::Real, min_new::Real = min_thresh, max_new::Real = max_thresh)
	s = ifelse(s < min_thresh, min_new, ifelse(s > max_thresh, max_new, s))
end

function scaled_norm(E::IdentityMatrix, v::Array, p::Real = 2)
	E.λ ? norm(v, p) : zero(eltype(v))
end

function scaled_norm(E::Diagonal, v::Array{T}, p::Real = 2) where{T}
	if p == 2
		return scaled_norm2(E, v)
	elseif p == Inf
		return scaled_norm_Inf(E, v)
	elseif p == 1
		return scaled_norm1(E, v)
	else
		throw(ArgumentError("bad norm specified"))
	end
end

function scaled_norm2(E::Diagonal, v::Array)
	sum_sq  = zero(eltype(v))
	for i = 1:length(v)
		sum_sq += (E.diag[i] * v[i])^2
	end
	return sqrt(sum_sq)
end

function scaled_norm_Inf(E::Diagonal, v::Array)
	norm  = zero(eltype(v))
	for i = 1:length(v)
		norm = max(norm, abs(E.diag[i] * v[i]))
	end
	return norm
end

function scaled_norm1(E::Diagonal, v::Array)
	norm  = zero(eltype(v))
	for i = 1:length(v)
		norm += abs(E.diag[i] * v[i])
	end
	return norm
end

function col_norms!(v::Array{Tf, 1},
	A::Matrix{Tf};
	reset::Bool = true) where{Tf <: AbstractFloat}

	if(reset) v.= 0 end

	for i = 1:size(A, 2)
		v[i] = max(v[i], norm(view(A, :, i), Inf))
	end
	return v
end

function col_norms!(v::Array{Tf, 1},
	A::SparseMatrixCSC{Tf,Ti}; reset::Bool = true) where{Tf <: AbstractFloat, Ti <: Integer}

	if(reset) v.= 0 end

	for i = 1:A.n
		@inbounds for j = A.colptr[i]:(A.colptr[i + 1] - 1)
		@inbounds v[i] = max(v[i], abs(A.nzval[j]))
	end
	end
	return v
end

function row_norms!(v::Array{Tf, 1},
	A::Matrix{Tf};
	reset::Bool = true) where{Tf <: AbstractFloat}

	if(reset) v.= 0 end

	for i = 1:size(A, 1)
		v[i] = max(v[i], norm(view(A, i, :), Inf))
	end
	return v
end

function row_norms!(v::Array{Tf, 1},
	A::SparseMatrixCSC{Tf, Ti};
	reset::Bool = true) where{Tf <: AbstractFloat, Ti <: Integer}

	if(reset) v.= 0 end

	@inbounds for i = 1:(A.colptr[end] - 1)
		@inbounds v[A.rowval[i]] = max(v[A.rowval[i]], abs(A.nzval[i]))
	end
	return v
end

function lmul!(L::Diagonal, M::SparseMatrixCSC)

	#NB : Same as:  @views M.nzval .*= D.diag[M.rowval]
	#but this way allocates no memory at all and
	#is marginally faster
	m, n = size(M)
	(m == length(L.diag)) || throw(DimensionMismatch())

	@inbounds for i = 1:(M.colptr[end] - 1)
	@inbounds M.nzval[i] *= L.diag[M.rowval[i]]
	end
	return M
end

lmul!(L::IdentityMatrix, M::AbstractMatrix) = L.λ ? M : M .= zero(eltype(M))

function rmul!(M::SparseMatrixCSC, R::Diagonal)

	m, n = size(M)
	(n == length(R.diag)) || throw(DimensionMismatch())

	@inbounds for i = 1:n, j = M.colptr[i]:(M.colptr[i + 1] - 1)
	@inbounds M.nzval[j] *= R.diag[i]
	end
	return M
end

rmul!(M::AbstractMatrix, R::IdentityMatrix) = R.λ ? R : R .= zero(eltype(R))

function lrmul!(L::Diagonal, M::SparseMatrixCSC, R::Diagonal)

	Mnzval  = M.nzval
	Mrowval = M.rowval
	Mcolptr = M.colptr
	Rd      = R.diag
	Ld      = L.diag

	m, n = size(M)
	(m == length(Ld) && n == length(Rd)) || throw(DimensionMismatch())

	@inbounds for i = 1:n, j = Mcolptr[i]:(Mcolptr[i + 1] - 1)
	@inbounds Mnzval[j] *= Ld[Mrowval[j]] * Rd[i]
	end
	return M
end

lrmul!(L::IdentityMatrix,
	M::AbstractMatrix,
	R::IdentityMatrix) = (L.λ && R.λ) ? M : M .= zero(eltype(M))

lrmul!(L::Diagonal,
	M::SparseMatrixCSC,
	R::IdentityMatrix) = R.λ ? lmul!(L, M) : M .= zero(eltype(M))

lrmul!(L::Diagonal,
	M::AbstractMatrix,
	R::Diagonal) = LinearAlgebra.lmul!(L, LinearAlgebra.rmul!(M, R))

lrmul!(L::Diagonal,
	M::AbstractMatrix,
	R::IdentityMatrix) = R.λ ? LinearAlgebra.lmul!(L, M) : M .= zero(eltype(M))

lrmul!(L::IdentityMatrix,
	M::AbstractMatrix,
	R::Diagonal) = L.λ ? LinearAlgebra.rmul!(M, R) : M .= zero(eltype(M))

"""
    symmetrize_upper!(A)

Symmetrizes the matrix A by calculating A = 0.5 * (A + A') but only performs the operation on the upper triangular part.
"""
function symmetrize_upper!(A::AbstractMatrix)
	n = size(A, 1)
	@assert(size(A, 1) == size(A, 2), "Matrix is not square.")
	@inbounds for j in 1:n, i in 1:j
 		A[i, j] = (A[i, j] + A[j, i]) / 2
	end
	nothing
end

"""
    symmetrize_full!(A)

Symmetrizes the matrix A by calculating A = 0.5 * (A + A') and storing the result in-place.
"""
function symmetrize_full!(A::AbstractMatrix)
	n = size(A, 1)
	@assert(size(A, 1) == size(A, 2), "Matrix is not square.")
	@inbounds for j in 1:n, i in 1:j
 		A[i, j] = (A[i, j] + A[j, i]) / 2
 		A[j, i] = A[i, j]
	end
	nothing
end

# this function assumes real symmetric X and only considers the upper triangular part
function is_pos_sem_def(X, tol)
    # set option 'N' to only compute eigenvalues, s is ordered from min to max
   s, U = LAPACK.syevr!('N', 'A', 'U', X, 0.0, 0.0, 0, 0, -1.0);
   return s[1] >= -tol
end

function is_neg_sem_def(X, tol)
    # set option 'N' to only compute eigenvalues, s is ordered from min to max
   s, U = LAPACK.syevr!('N', 'A', 'U', X, 0.0, 0.0, 0, 0, -1.0);
   return s[end] <= tol
end


# Kolmogorov distribution
# defined as the sup_{t \in [0,1]} |B(t)|, where B(t) is a Brownian bridge
# used in the Kolmogorov--Smirnov test for large n.

immutable Kolmogorov <: ContinuousUnivariateDistribution
end

insupport(::Kolmogorov, x::Real) = zero(x) <= x < Inf
insupport(::Type{Kolmogorov}, x::Real) = zero(x) <= x < Inf

mean(d::Kolmogorov) = 0.5*√2π*log(2.0)
var(d::Kolmogorov) = pi*pi/12.0 - 0.5*pi*log(2.0)^2
# TODO: higher-order moments also exist, can be obtained by differentiating series

# cdf and ccdf are based on series truncation.
# two different series are available, e.g. see:
#   N. Smirnov, "Table for Estimating the Goodness of Fit of Empirical Distributions",
#   The Annals of Mathematical Statistics , Vol. 19, No. 2 (Jun., 1948), pp. 279-281
#   http://projecteuclid.org/euclid.aoms/1177730256
# use one series for small x, one for large x
# 5 terms seems to be sufficient for Float64 accuracy
# some divergence from Smirnov's table in 6th decimal near 1.0 (e.g. 1.04): occurs in
# both series so assume error in table.

function cdf(d::Kolmogorov,x::Real)
    if x <= 0.0
        return 0.0
    elseif x > 1.0
        return 1.0-ccdf(d,x)
    end
    s = 0.0
    for i = 1:5
        s += exp(-((2*i-1)*π/x)^2/8.0)
    end
    √2π*s/x
end

function ccdf(d::Kolmogorov,x::Real)
    if x <= 1.0
        return 1.0-cdf(d,x)
    end
    s = 0.0
    for i = 1:5
        s += (iseven(i) ? -1 : 1)*exp(-2.0*(i*x)^2)
    end
    2.0*s
end

# TODO: figure out how best to truncate series
function pdf(d::Kolmogorov,x::Real)
    if x <= 0.0
        return 0.0
    elseif x <= 1.0
        c = π/(2.0*x)
        s = 0.0
        for i = 1:20
            k = ((2*i-1)*c)^2
            s += (k-1.0)*exp(-k/2.0)
        end
        return √2π*s/x^2
    else
        s = 0.0
        for i = 1:20
            s += (iseven(i) ? -1 : 1)*i^2*exp(-2.0*(i*x)^2)
        end
        return 8.0*x*s
    end
end

# Devroye IV.5

function rand(d::Kolmogorov)
    # p = cdf(d,t)
    if rand() < p
        # left interval
        while true
            # generate truncated gamma variate
            # g = rand(Truncated(Gamma(1.5,1.0),0.0,t))
            tp = pi*pi/(8*t*t)
            acc = false
            while !acc
                e0, e1 = Base.Random.randmtzig_exprnd(), Base.Random.randmtzig_exprnd()
                e0 = e= / (1- 1/(2*tp))
                e1 = 2*e1
                g = tp + e0
                acc = (e0*e0 <= tp*e1*(g+tp)) || (g/tp - 1 - log(g/tp) <= e1)
            end

            x = pi/sqrt(8g)
            w = 0
            z = 1/(2*g)
            p = exp(-g)
            n = 1
            q = 1
            u = rand()
            while u >= w
                w += z*q
                if u >= w
                    return w
                end
                n += 2
                q = p^(n*n-1)
                w -= n*n*q
            end
        end
    else
        e = Base.Random.randmtzig_exprnd()
        u = rand()
        x = sqrt(t*t+e/2)
        w = 0
        n = 1
        z = exp(-2*x*x)
        while u > w
            n += 1
            w += n*n*z^(n*n-1)
            if u >= w
                return x
            end
            n += 1
            w -= n*n*z^(n*n-1)
        end
    end
end


## SMA performance: buffer vector version
using TimeSeries, MarketTechnicals
using Reactive

nma = 50
ohlc_BA = TimeSeries.readtimearray("../../data/OHLC_BA.csv")
sma_ta = sma(ohlc_BA["Close"], nma)

s_ohlc = Reactive.Input(ohlc_BA.values[nma,:])

function smabuffer!(buffer, val_ohlc)
    push!(buffer, val_ohlc[4]) # add new close price
    shift!(buffer) # remove the earliest element
    return buffer
end
s_sma = lift(mean, foldl(smabuffer!, ohlc_BA["Close"].values[1:nma], s_ohlc))

vsig = [s_sma.value]
function vsigupd!(vsig::Vector{Float64}, nma::Int64, nmax::Int64)
    for i in (nma + 1):nmax
        push!(s_ohlc, ohlc_BA.values[i,:])
        push!(vsig, s_sma.value)
    end
end

# verify sma values
vsigupd!(vsig, nma, length(ohlc_BA))
println("check if SMA values are correct:")
println(mean(abs(vsig - sma_ta.values)) < 1e-7)
println("")

println("SMA foldl-buffer version:")
println("reactive signal updates timing")
expr_run = parse("@time vsigupd!(vsig, nma, length(ohlc_BA))")
for i = 1:5
    eval(expr_run)
end

### results
# elapsed time: 0.214774855 seconds (21274016 bytes allocated, 19.48% gc time)

Profile.clear()  # in case we have any previous profiling data
@profile vsigupd!(vsig, nma, length(ohlc_BA))
Profile.print()


"""
    ngwp_analysis(G::GraphSig, wavelet_packet::Array{Float64,3})

For a GraphSig object `G`, generate the matrix of NGWP expansion coefficients.

# Input Arguments
- `G::GraphSig`: an input GraphSig object
- `wavelet_packet::Array{Float64,3}`: the varimax wavelets packet.

# Output Argument
- `dmatrix::Array{Float64,3}`: the expansion coefficients matrix.

"""
function ngwp_analysis(G::GraphSig, wavelet_packet::Array{Float64,3})
    f = G.f
    fcols = size(f, 2)
    (N, jmax, ) = Base.size(wavelet_packet)

    dmatrix = zeros(N, jmax, fcols)
    dmatrix[:, 1, :] = f

    for j = 2:jmax
        for i = 1:N
            dmatrix[i, j, :] = f' * wavelet_packet[i, j, :]
        end
    end

    return dmatrix
end

"""
    const_proj_wavelets(𝚽,vlist,elist; method = "Modified Gram-Schmidt with Lp Pivoting")

construct projection wavelets, i.e., project δ ∈ vlist onto span({φⱼ| j ∈ elist}).

# Input Arguments
- `𝚽::Matrix{Float64}`: graph Laplacian eigenvectors 𝚽
- `vlist::Array{Int}`: the list of considered node indices.
- `elist::Array{Int}`: the list of considered eigenvector indices.
- `method::Symbol`: default is `:MGSLp`. other options: `:IP` (Iterative-Projection),
    `GS` (Gram Schmidt).

# Output Argument
- `Wav::Matrix{Float64}`: a matrix whose columns are projected wavelet vectors.

"""
function const_proj_wavelets(𝚽, vlist, elist; method = :MGSLp)
    if length(vlist) == 1
        return 𝚽[:, elist]
    end
    N = size(𝚽, 1)
    m = length(vlist)
    Wav = zeros(N, m)

    B = 𝚽[:, elist]

    if method == :IP
        for k in 1:length(vlist)
            wavelet = Proj(spike(vlist[k], N), B)
            Wav[:, k] .= wavelet ./ norm(wavelet)
            B = wavelet_perp_Matrix(wavelet, B)
        end
    elseif method == :GS
        P = B * B'
        for k in 1:length(vlist)
            wavelet = P * spike(vlist[k], N)
            Wav[:,k] .= wavelet ./ norm(wavelet)
        end
        Wav, complement_dim = gram_schmidt(Wav)
        if complement_dim != 0
            complement_space = B * nullspace(Wav' * B)
            Wav = hcat(Wav, complement_space)
        end
    elseif method == :MGSLp
        P = B * B'
        for k in 1:length(vlist)
            wavelet = P * spike(vlist[k], N)
            Wav[:,k] .= wavelet ./ norm(wavelet)
        end
        Wav, complement_dim = mgslp(Wav)
        if complement_dim != 0
            complement_space = B * nullspace(Wav' * B)
            Wav = hcat(Wav, complement_space)
        end
    else
        error("Do not support method ", method)
    end

    return Wav
end

"""
    function NGWP_jkl(GP_star::GraphPart, drow::Int, dcol::Int)

Generate the (j,k,l) indices for the NGWP basis vector corresponding to the coefficient dmatrix(drow,dcol)

### Input Arguments
* `GP_star::GraphPart`: a GraphPart object of the dual grpah
* `drow::Int`: the row of the expansion coefficient
* `dcol::Int`: the column of the expansion coefficient

### Output Argument
* `j`: the level index of the expansion coefficient
* `k`: the subregion in dual graph's index of the expansion coefficient
* `l`: the tag of the expansion coefficient
"""
function NGWP_jkl(GP_star::GraphPart, drow::Int, dcol::Int)
    (j, k, l) = GHWT_jkl(GP_star, drow, dcol)
    return j,k,l
end


# # Entries
#
# The key ingredient for data representations are `AxisEntries`.
#
# ## The `AxisEntries` type
#
# An `AxisEntries` object is
# made of four components:
# - axis,
# - entries,
# - scales,
# - labels.

using AlgebraOfGraphics, CairoMakie
using AlgebraOfGraphics: CategoricalScale, ContinuousScale
resolution = (600, 600)
fig = Figure(; resolution)
N = 11
rg = range(1, 2, length=N)
ae = AxisEntries(
    Axis(fig[1, 1]),
    [
        Entry(
            plottype=Scatter,
            positional=(rg, cosh.(rg)),
            named=(color=1:N, marker=fill("b", N));
            attributes=Dict(:markersize => 15)
        ),
        Entry(
            plottype=Scatter,
            positional=(rg, sinh.(rg)),
            named=(color=1:N, marker=fill("c", N));
            attributes=Dict(:markersize => 15)
        ),
    ],
    Dict(
        1 => ContinuousScale(identity, (0, 4)),
        2 => ContinuousScale(identity, (0, 4)),
        :color => ContinuousScale(identity, (1, N)),
        :marker => CategoricalScale(["a", "b", "c"], [:circle, :utriangle, :dtriangle]),
    ), # scales
    Dict(
        1 => "x",
        2 => "y",
        :color => "identity",
        :marker => "function"
    ), # labels
)
plot!(ae)
fig


include("euler/euler.jl")
using .Primes: prime_numbers

compute(n::Int)::Int = prod([prime ^ trunc(Int, log(prime, n)) for prime ∈ prime_numbers(n)])


using JuMP, VariationalInequality
using Test

# https://supernet.isenberg.umass.edu/articles/SPE_Model_Information_Asymmetry_in_Quality.pdf
# Problem (15), Data in Table 1, Example 1
m = 2; n = 1

model = VIPModel()

@variable(model, s[i=1:m] >=0)
@variable(model, d[j=1:n] >=0)
@variable(model, Q[i=1:m, j=1:n] >= 0)
@variable(model, q[i=1:m] >= 0)

@constraint(model, supply[i=1:m], s[i] == sum(Q[i,j] for j in 1:n))
@constraint(model, demand[j=1:n], d[j] == sum(Q[i,j] for i in 1:m))

as = [5; 2]
bs = [5; 10]
@mapping(model, pi[i=1:m], as[i] * s[i] + q[i] + bs[i])

ac = [1; 2]
bc = [15; 20]
@mapping(model, c[i=1:m, j=1:n], ac[i,j] * Q[i,j] + bc[i,j] )

ad = [2]
bd = [100]
@NLexpression(model, qhat[j=1:n], sum(q[i]*Q[i,j] for i in 1:m) / ( sum(Q[i,j] for i in 1:m) + 1e-6 ) )
@mapping(model, nrho[j=1:n], ad[j] * d[j] - qhat[j] - bd[j] )

aq = [5; 10]
@mapping(model, Fq[i=1:m], aq[i] * q[i] - pi[i] )

@innerproduct(model, pi, s)
@innerproduct(model, c, Q)
@innerproduct(model, nrho, d)
@innerproduct(model, Fq, q)

for i=1:m, j=1:n
    setvalue(Q[i,j], 1.0)
end

sol1, Fval1, gap1 = solveVIP(model, algorithm=:fixed_point, max_iter=10000, step_size=0.1, tolerance=1e-10)
@assert 0<= gap1 < 1e-6

@show sol1
@show gap1

@show sol1[Q[1,1]]
@show sol1[Q[2,1]]
@show sol1[q[1]]
@show sol1[q[2]]


#   Unit examples of the FourierAnalysis Package for julia language
#   v 0.0.1 - last update 24th of September 2019
#
#   MIT License
#   Copyright (c) 2019, Marco Congedo, CNRS, Grenobe, France:
#   https://sites.google.com/site/marcocongedo/home

# ? CONTENTS :
#   This example shows how to compute univariate measures in
#   the time-frequency domain.

#   ~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~  #
#                                                                             #
#   ~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~  #

using FourierAnalysis, FFTW, LinearAlgebra, Statistics, Plots

# add module for reading the two EEG text files to be used ater
push!(LOAD_PATH, @__DIR__)
using IOtxt

# Get EEG file names with complete path (they have extension .txt)
S=getFilesInDir(@__DIR__; ext=(".txt",))

# read the two EEG data files and put them in a Matrix object
X1=readEEG(S[1])
X2=readEEG(S[2])

#############################################################################
# Univariate measures
# see Congedo, 2018: https://hal.archives-ouvertes.fr/hal-01868538/document

sr, wl, bandwidth=128, 512, 2
Pz=15
𝐱=[X1[:, Pz], X2[:, Pz]] # get the two times-series at electrode Pz
𝐘=TFanalyticsignal(𝐱, sr, wl, bandwidth; fmax=32, nonlinear=false)
𝐀=TFamplitude(𝐘)

# Mean Amplitude Eq. 0.1 (MAmp)
# compute the MAmp averaging in a TF region from a TFAnalyticSignalVector object
MAmp=meanAmplitude(𝐘, (8, 12), (1, 512); mode=mean)
# compute the MAmp averaging in a TF region from a TFAmplitudeVector object
MAmp=meanAmplitude(𝐀, (8, 12), (1, 512); mode=mean)
# compute the MAmp averaging in a TF region directly from data
MAmp=meanAmplitude(𝐱, sr, wl, (8, 12), (1, 512), bandwidth; mode=mean)

# compute the MAmp in a TF region from a TFAnalyticSignalVector object
MAmp=meanAmplitude(𝐘, (8, 12), (1, 512); mode=extract)
# compute the MAmp in a TF region from a TFAmplitudeVector object
MAmp=meanAmplitude(𝐀, (8, 12), (1, 512); mode=extract)
# compute the MAmp in a TF region directly from data
MAmp=meanAmplitude(𝐱, sr, wl, (8, 12), (1, 512), bandwidth; mode=extract)
# NB the Analytic Signal or the amplitude objects must be linear (note the nonlinear=false above)

# All these operation can be obtained on smoothed Amplitude, e.g.,
MAmp=meanAmplitude(𝐱, sr, wl, (8, 12), (1, 512), bandwidth;
                   mode=extract,
                   fsmoothing=hannSmoother,
                   tsmoothing=hannSmoother)

# The same operations can be done with the other univariate measures. For example:
# Concentration Eq. 0.2 (Con)
# compute the Con averaging in a TF region from a TFAnalyticSignalVector object
Con=concentration(𝐘, (8, 12), (1, 512); mode=mean)
# compute the Con in a TF region directly from data
Con=concentration(𝐱, sr, wl, (8, 12), (1, 512), bandwidth; mode=extract)
# Mean Direction Eq. 0.3 (MDir)
# compute the MDir averaging in a TF region directly from data
MDir=meanDirection(𝐱, sr, wl, (8, 12), (1, 512), bandwidth; mode=mean)
# compute the MDir in a TF region from a TFAnalyticSignalVector object
MDir=meanDirection(𝐘, (8, 12), (1, 512); mode=extract)

# and for the non-linear counterpart:
# Phase Concentration Eq. 0.4 (PCon)
# compute the Con in a TF region directly from data
Con=concentration(𝐱, sr, wl, (8, 12), (1, 512), bandwidth; mode=extract, nonlinear=true)
# Phase Mean Direction Eq. 0.5 (PMDir)
# compute the MDir averaging in a TF region directly from data
MDir=meanDirection(𝐱, sr, wl, (8, 12), (1, 512), bandwidth; mode=mean, nonlinear=true)

# If you try to compute a non-linear measure from a linear AnalyticSignal object
# you will get en error (see the REPL)
Con=concentration(𝐘, (8, 12), (1, 512); mode=mean, nonlinear=true)

# In order to compute those quantities from Analytic Signal objects
# first we need to compute a non-linear Analytic Signal objects:
𝐘=TFanalyticsignal(𝐱, sr, wl, bandwidth; fmax=32, nonlinear=true)
# then
# Phase Concentration Eq. 0.4 (PCon)
Con=concentration(𝐘, (8, 12), (1, 512); mode=mean, nonlinear=true)
# compute the MDir in a TF region from a TFAnalyticSignalVector object
MDir=meanDirection(𝐘, (8, 12), (1, 512); mode=extract, nonlinear=true)
# NB the Analytic Signal or the amplitude objects must be non-linear (note the nonlinear=true above)

############################################################################


using Statistics
using Test
using DelimitedFiles
import MLJBase
const Mlj = MLJBase
using StableRNGs
using BetaML

#TESTRNG = FIXEDRNG # This could change...
TESTRNG = StableRNG(123)

println("*** Testing Perceptron algorithms...")

# ==================================
# TEST 1: Normal perceptron
println("Going through Test1 (normal Perceptron)...")

perceptronData     = readdlm(joinpath(@__DIR__,"data/binary2DData.csv"),'\t')
x = copy(perceptronData[:,[2,3]])
y = convert(Array{Int64,1},copy(perceptronData[:,1]))
ntrain = Int64(round(length(y)*0.8))
xtrain = x[1:ntrain,:]
ytrain = y[1:ntrain]
xtest = x[ntrain+1:end,:]
ytest = y[ntrain+1:end]
classes = unique(y)
out      = perceptron(xtrain, ytrain, shuffle=false,nMsgs=0)
ŷtrain    = predict(xtrain,out.θ,out.θ₀,out.classes)
ϵtrain    = error(ytrain, mode(ŷtrain))
ŷtest    = predict(xtest,out.θ,out.θ₀,classes)
outTest  = perceptron(xtrain, ytrain, shuffle=false,nMsgs=0,returnMeanHyperplane=true)
ŷavgtest = predict(xtest,outTest.θ,outTest.θ₀,outTest.classes)
ϵ        = error(ytest, mode(ŷtest))
ϵavg     = error(ytest, mode(ŷavgtest))
@test ϵ    < 0.03
@test ϵavg < 0.2

println("Testing multiple classes...")

#xtrain = [3 4 5; 5 3 5; 3 7 2; 8 5 3; 4 2 3; 3 2 1; 8 3 4; 3 5 1; 1 9 3; 4 2 1]
xtrain = rand(TESTRNG,100,3)
ytt    = [(0.5*x[1]+0.2*x[2]^2+0.3*x[3]+1) for (i,x) in enumerate(eachrow(xtrain))]
ytrain = [i > median(ytt)*1.1 ? "big" :  i > median(ytt)*0.9 ? "avg" : "small" for i in ytt]
#xtest  = [2 2 3; 3 2 2; 4 1 2; 4 3 2; 3 7 2]
xtest = rand(TESTRNG,20,3)
ytt2   = [(0.5*x[1]+0.2*x[2]^2+0.3*x[3]+1) for (i,x) in enumerate(eachrow(xtest))]
ytest  = [i > median(ytt2)*1.1 ? "big" :  i > median(ytt2)*0.9 ? "avg" : "small" for i in ytt2]

out    = perceptron(xtrain,  ytrain, shuffle=false,nMsgs=0)
ŷtrain = predict(xtrain,out.θ,out.θ₀,out.classes)
ŷtest  = predict(xtest,out.θ,out.θ₀,out.classes)
ϵtrain = error(ytrain, mode(ŷtrain))
ϵtest  = error(ytest, mode(ŷtest))

@test ϵtrain  < 0.4
@test ϵavg    < 0.4

# ==================================
# Test 2: Kernel Perceptron
# ==================================
println("Going through Test2 (Kernel Perceptron)...")

xtrain = [3 4 5; 5 3 5; 3 7 2; 8 5 3; 4 2 3; 3 2 1; 8 3 4; 3 5 1; 1 9 3; 4 2 1]
ytt    = [(0.5*x[1]+0.2*x[2]^2+0.3*x[3]+1) for (i,x) in enumerate(eachrow(xtrain))]
ytrain = [i > median(ytt) ? 1 : -1 for i in ytt]
xtest  = [ 3 7 2; 2 2 3; 3 2 2; 4 1 2; 4 3 2;]
#xtest = xtrain
ytt2    = [(0.5*x[1]+0.2*x[2]^2+0.3*x[3]+1) for (i,x) in enumerate(eachrow(xtest))]
ytest = [i > median(ytt2) ? 1 : -1 for i in ytt2]
#out   = kernelPerceptron(xtrain, ytrain, K=polynomialKernel,rShuffle=true,nMsgs=100)
#ŷtest = predict(xtest,out[1][1],out[1][2],out[1][3], K=polynomialKernel)
out   = kernelPerceptronBinary(xtrain, ytrain, K=radialKernel,shuffle=false,nMsgs=0,α=ones(Int64,length(ytrain)))
# the same: out   = kernelPerceptronBinary(xtrain, ytrain, K=radialKernel,shuffle=false,nMsgs=0)
ŷtest = predict(xtest,out.x,out.y,out.α, K=out.K)
ϵ = error(ytest, ŷtest)
ŷtestExpected = [1,-1,-1,-1,-1]
@test ϵ ≈ 0.2
#@test any(isapprox(ŷtestExpected,ŷtest))
@test any(ŷtestExpected == ŷtest )

# Multiclass..
outMultiClass   = kernelPerceptron(xtrain, ytrain, K=radialKernel,shuffle=false,nMsgs=0)
ŷtest = predict(xtest,outMultiClass.x,outMultiClass.y,outMultiClass.α, outMultiClass.classes,K=outMultiClass.K)
ϵ = error(ytest, mode(ŷtest))
ŷtestExpected = [1,-1,-1,-1,-1]
@test ϵ ≈ 0.2
#@test any(isapprox(ŷtestExpected,ŷtest))
@test any(ŷtestExpected == mode(ŷtest) )

xtrain = rand(TESTRNG,100,3)
ytt    = [(0.5*x[1]+0.2*x[2]^2+0.3*x[3]+1) for (i,x) in enumerate(eachrow(xtrain))]
ytrain = [i > median(ytt)*1.1 ? "big" :  i > median(ytt)*0.9 ? "avg" : "small" for i in ytt]
#xtest  = [2 2 3; 3 2 2; 4 1 2; 4 3 2; 3 7 2]
xtest = rand(TESTRNG,20,3)
ytt2   = [(0.5*x[1]+0.2*x[2]^2+0.3*x[3]+1) for (i,x) in enumerate(eachrow(xtest))]
ytest  = [i > median(ytt2)*1.1 ? "big" :  i > median(ytt2)*0.9 ? "avg" : "small" for i in ytt2]

out    = kernelPerceptron(xtrain,  ytrain, shuffle=false,nMsgs=0,T=1000)
ŷtrain = predict(xtrain,out.x,out.y,out.α, out.classes,K=out.K)
ŷtest  = predict(xtest,out.x,out.y,out.α, out.classes,K=out.K)
ϵtrain = error(ytrain, mode(ŷtrain))
ϵtest  = error(ytest, mode(ŷtest))

@test ϵtrain  < 0.1
@test ϵtest   < 0.8

# ==================================
# Test 3: Pegasos
# ==================================
println("Going through Test3 (Pegasos)...")

perceptronData     = readdlm(joinpath(@__DIR__,"data/binary2DData.csv"),'\t')
x = copy(perceptronData[:,[2,3]])
y = convert(Array{Int64,1},copy(perceptronData[:,1]))
xtrain = x[1:160,:]
ytrain = y[1:160]
xtest  = x[161:end,:]
ytest  = y[161:end]

out   = pegasos(xtrain, ytrain, shuffle=false, nMsgs=0)
ŷtest = predict(xtest,out.θ,out.θ₀,out.classes)
outAvg   = pegasos(xtrain, ytrain, shuffle=false, nMsgs=0, returnMeanHyperplane=true)
ŷavgtest = predict(xtest,outAvg.θ,outAvg.θ₀,outAvg.classes)
ϵ = error(ytest, mode(ŷtest))
ϵavg = error(ytest, mode(ŷavgtest))
@test ϵ ≈ 0.025
@test ϵavg ≈ 0.1

println("Testing pegasos with multiple classes...")

#xtrain = [3 4 5; 5 3 5; 3 7 2; 8 5 3; 4 2 3; 3 2 1; 8 3 4; 3 5 1; 1 9 3; 4 2 1]
xtrain = rand(TESTRNG,100,3)
ytt    = [(0.5*x[1]+0.2*x[2]^2+0.3*x[3]+1) for (i,x) in enumerate(eachrow(xtrain))]
ytrain = [i > median(ytt)*1.1 ? "big" :  i > median(ytt)*0.9 ? "avg" : "small" for i in ytt]
#xtest  = [2 2 3; 3 2 2; 4 1 2; 4 3 2; 3 7 2]
xtest = rand(TESTRNG,20,3)
ytt2   = [(0.5*x[1]+0.2*x[2]^2+0.3*x[3]+1) for (i,x) in enumerate(eachrow(xtest))]
ytest  =  [i > median(ytt2)*1.1 ? "big" :  i > median(ytt)*0.9 ? "avg" : "small" for i in ytt2]

out    = pegasos(xtrain,  ytrain, shuffle=false,nMsgs=0)
ŷtrain = predict(xtrain,out.θ,out.θ₀,out.classes)
ŷtest  = predict(xtest,out.θ,out.θ₀,out.classes)
ϵtrain = error(ytrain, mode(ŷtrain))
ϵtest  = error(ytest, mode(ŷtest))

@test ϵtrain  <= 0.8 # this relation is not linear, normal error is big
@test ϵtest   <= 0.8

# ==================================
# NEW TEST
println("Testing classification of the sepal database using perceptron algorithms...")
iris = readdlm(joinpath(@__DIR__,"data","iris_shuffled.csv"),',',skipstart=1)
x = convert(Array{Float64,2}, iris[:,1:4])
y = convert(Array{String,1}, iris[:,5])

ntrain = Int64(round(size(x,1)*0.8))
xtrain = x[1:ntrain,:]
ytrain = y[1:ntrain]
xtest = x[ntrain+1:end,:]
ytest = y[ntrain+1:end]

model = perceptron(xtrain,ytrain)
ŷtrain = predict(xtrain,model.θ,model.θ₀,model.classes)
@test accuracy(mode(ŷtrain),ytrain) >= 0.79
ŷtest = predict(xtest,model.θ,model.θ₀,model.classes)
@test accuracy(mode(ŷtest),ytest)  >= 0.9

model = kernelPerceptron(xtrain,ytrain)
ŷtrain = predict(xtrain,model.x,model.y,model.α,model.classes)
@test accuracy(mode(ŷtrain),ytrain) >= 0.9
ŷtest = predict(xtest,model.x,model.y,model.α,model.classes)
@test accuracy(mode(ŷtest),ytest)  >= 0.9

model = pegasos(xtrain,ytrain)
ŷtrain = predict(xtrain,model.θ,model.θ₀,model.classes)
@test accuracy(mode(ŷtrain),ytrain) >= 0.64
ŷtest = predict(xtest,model.θ,model.θ₀,model.classes)
@test accuracy(mode(ŷtest),ytest)  >= 0.76

# ==================================
# NEW TEST
println("Testing MLJ interface for Perceptron models....")

X, y                           = Mlj.@load_iris

model                          = PerceptronClassifier(rng=copy(TESTRNG))
regressor                      = Mlj.machine(model, X, y)
(fitresult, cache, report)     = Mlj.fit(model, 0, X, y)
yhat                           = Mlj.predict(model, fitresult, X)
@test Mlj.mean(Mlj.LogLoss(tol=1e-4)(yhat, y)) < 3.1

model                          = KernelPerceptronClassifier(rng=copy(TESTRNG))
regressor                      = Mlj.machine(model, X, y)
(fitresult, cache, report)     = Mlj.fit(model, 0, X, y)
yhat                           = Mlj.predict(model, fitresult, X)
@test Mlj.mean(Mlj.LogLoss(tol=1e-4)(yhat, y)) < 0.5

model                          = PegasosClassifier(rng=copy(TESTRNG))
regressor                      = Mlj.machine(model, X, y)
(fitresult, cache, report)     = Mlj.fit(model, 0, X, y)
yhat                           = Mlj.predict(model, fitresult, X)
@test Mlj.mean(Mlj.LogLoss(tol=1e-4)(yhat, y)) < 1.3


module DemonstratingLinearInterpolatorsInterpolations

using Printf
using LazyAlgebra
using LazyAlgebra: Direct, Adjoint
using LinearInterpolators
using InterpolationKernels
using LinearInterpolators.Interpolations

using PyCall
# pygui(:gtk); # can be: :wx, :gtk, :qt
using LaTeXStrings
import PyPlot
const plt = PyPlot

function rundemos(::Type{T} = Float64) where {T<:AbstractFloat}

    z = T[0.5, 0.3, 0.1, 0.0, -0.2, -0.7, -0.7, 0.0, 1.7, 1.9, 2.1]

    # 1-D example
    t = range(-3, 14, length=2000);
    I0 = SparseInterpolator(RectangularSpline(T, Flat), t, length(z))
    I1 = SparseInterpolator(LinearSpline(T, Flat), t, length(z))
    I2 = SparseInterpolator(CatmullRomSpline(T, SafeFlat), t, length(z))
    I3 = SparseInterpolator(LanczosKernel(T, 8), t, length(z))
    plt.figure(1)
    plt.clf()
    plt.plot(t, I0(z), color="darkgreen",
             linewidth=1.5, linestyle="-", label="Rectangular spline");
    plt.plot(t, I1(z), color="darkred",
                 linewidth=1.5, linestyle="-", label="Linear spline");
    plt.plot(t, I2(z), color="orange",
             linewidth=1.5, linestyle="-", label="Catmull-Rom");
    plt.plot(t, I3(z), color="violet",
             linewidth=1.5, linestyle="-", label="Lanczos 8");
    plt.plot(1:length(z), z, "o", color="darkblue", label="Points")
    plt.legend()
    plt.title("Interpolations with Cardinal Kernels")

    # Interpolation + smoothing
    t = range(-3, 14, length=2000);
    I5 = SparseInterpolator(QuadraticSpline(T), t, length(z))
    I6 = SparseInterpolator(CubicSpline(T, Flat), t, length(z))
    I7 = SparseInterpolator(MitchellNetravaliSpline(T, Flat), t, length(z))
    plt.figure(2)
    plt.clf()
    plt.plot(t, I6(z), color="darkred",
                 linewidth=1.5, linestyle="-", label="Cubic B-spline");
    plt.plot(t, I5(z), color="darkgreen",
             linewidth=1.5, linestyle="-", label="Quadratic B-spline");
    #plt.plot(t, I7(z), color="orange",
    #         linewidth=1.5, linestyle="-", label="Mitchell & Netravali");
    plt.plot(1:length(z), z, "o", color="darkblue", label="Points")
    plt.legend()
    plt.title("Interpolations with Smoothing Kernels")

    # Test conversion to a sparse matrix.
    S1 = sparse(I1)
    S2 = sparse(I2)
    println("max. error 1: ", maximum(abs.(I1(z) - S1*z)))
    println("max. error 2: ", maximum(abs.(I2(z) - S2*z)))

    # 2-D example
    t = reshape(range(-3, 14, length=20*21), (20,21));
    I = SparseInterpolator(CatmullRomSpline(T, Flat), t, length(z))
    plt.figure(3)
    plt.clf()
    plt.imshow(I(z));
end
#rundemos()
#------------------------------------------------------------------------------

const CPU = 2.94e9

function printtest(prefix::String, count::Integer, tup)
    elt, mem, gct = tup[2:4]
    @printf("%s   %6.0f cycles (Mem. = %8d bytes   GC = %3.0f%%)\n",
            prefix, CPU*elt/count, mem, 1e2*gct/elt)
end

function testdirect!(dst::Vector{T}, ker::Kernel{T,S,B},
                     x::AbstractVector{T}, src::Vector{T},
                     cnt::Integer) where {T,S,B}
    for k in 1:cnt
        apply!(dst, Direct, ker, x, src)
    end
    return dst
end

function testadjoint!(dst::Vector{T}, ker::Kernel{T,S,B},
                      x::AbstractVector{T}, src::Vector{T},
                      cnt::Integer) where {T,S,B}
    for k in 1:cnt
        apply!(dst, Adjoint, ker, x, src)
    end
    return dst
end

function runtests(::Type{T} = Float64,
                  len::Integer = 1000, cnt::Integer = 100) where T<:AbstractFloat
    dim = 500
    w = 10
    K01 = RectangularSpline(T, Flat)
    K02 = RectangularSpline(T, SafeFlat)
    K03 = LinearSpline(T, Flat)
    K04 = LinearSpline(T, SafeFlat)
    K05 = QuadraticSpline(T, Flat)
    K06 = QuadraticSpline(T, SafeFlat)
    K07 = CardinalCubicSpline(T, -1, Flat)
    K08 = CardinalCubicSpline(T, -1, SafeFlat)
    K09 = CatmullRomSpline(T, Flat)
    K10 = CatmullRomSpline(T, SafeFlat)
    x = rand(T(1 - w):T(dim + w), len)
    y = randn(T, dim)
    z = Array{T}(undef, len)
    # compile methods
    testdirect!(z,  K01, x, y, 3)
    testadjoint!(y, K01, x, z, 3)
    testdirect!(z,  K02, x, y, 3)
    testadjoint!(y, K02, x, z, 3)
    testdirect!(z,  K03, x, y, 3)
    testadjoint!(y, K03, x, z, 3)
    testdirect!(z,  K04, x, y, 3)
    testadjoint!(y, K04, x, z, 3)
    testdirect!(z,  K05, x, y, 3)
    testadjoint!(y, K05, x, z, 3)
    testdirect!(z,  K06, x, y, 3)
    testadjoint!(y, K08, x, z, 3)
    testdirect!(z,  K07, x, y, 3)
    testadjoint!(y, K07, x, z, 3)
    testdirect!(z,  K08, x, y, 3)
    testadjoint!(y, K08, x, z, 3)
    testdirect!(z,  K09, x, y, 3)
    testadjoint!(y, K09, x, z, 3)
    testdirect!(z,  K10, x, y, 3)
    testadjoint!(y, K10, x, z, 3)
    # warm-up
    tmp = randn(T, 1_000_000)
    # run tests
    n = cnt*len
    printtest("RectangularSpline   Flat     direct ", n, @timed(testdirect!(z,  K01, x, y, cnt)))
    printtest("RectangularSpline   Flat     adjoint", n, @timed(testadjoint!(y, K01, x, z, cnt)))
    printtest("RectangularSpline   SafeFlat direct ", n, @timed(testdirect!(z,  K02, x, y, cnt)))
    printtest("RectangularSpline   SafeFlat adjoint", n, @timed(testadjoint!(y, K02, x, z, cnt)))
    println()
    printtest("LinearSpline        Flat     direct ", n, @timed(testdirect!(z,  K03, x, y, cnt)))
    printtest("LinearSpline        Flat     adjoint", n, @timed(testadjoint!(y, K03, x, z, cnt)))
    printtest("LinearSpline        SafeFlat direct ", n, @timed(testdirect!(z,  K04, x, y, cnt)))
    printtest("LinearSpline        SafeFlat adjoint", n, @timed(testadjoint!(y, K04, x, z, cnt)))
    println()
    printtest("QuadraticSpline     Flat     direct ", n, @timed(testdirect!(z,  K05, x, y, cnt)))
    printtest("QuadraticSpline     Flat     adjoint", n, @timed(testadjoint!(y, K05, x, z, cnt)))
    printtest("QuadraticSpline     SafeFlat direct ", n, @timed(testdirect!(z,  K06, x, y, cnt)))
    printtest("QuadraticSpline     SafeFlat adjoint", n, @timed(testadjoint!(y, K06, x, z, cnt)))
    println()
    printtest("CardinalCubicSpline Flat     direct ", n, @timed(testdirect!(z,  K07, x, y, cnt)))
    printtest("CardinalCubicSpline Flat     adjoint", n, @timed(testadjoint!(y, K07, x, z, cnt)))
    printtest("CardinalCubicSpline SafeFlat direct ", n, @timed(testdirect!(z,  K08, x, y, cnt)))
    printtest("CardinalCubicSpline SafeFlat adjoint", n, @timed(testadjoint!(y, K08, x, z, cnt)))
    println()
    printtest("CatmullRomSpline    Flat     direct ", n, @timed(testdirect!(z,  K09, x, y, cnt)))
    printtest("CatmullRomSpline    Flat     adjoint", n, @timed(testadjoint!(y, K09, x, z, cnt)))
    printtest("CatmullRomSpline    SafeFlat direct ", n, @timed(testdirect!(z,  K10, x, y, cnt)))
    printtest("CatmullRomSpline    SafeFlat adjoint", n, @timed(testadjoint!(y, K10, x, z, cnt)))

end

end # module


module MultiObjectiveProblems

import JuMP
using OSQP
import ForwardDiff
using Parameters: @with_kw

import NLopt;    # for ZDT4 local critical set

export MOP, SamplingFunction, FixedPointSet, Constraints, Box
export num_objectives, num_vars;
export get_objectives, get_vector_objective, get_gradients, get_omega_function, constraints,
    get_pareto_set, get_pareto_front, get_points, get_scatter_arrays, get_scatter_points,
    get_random_point, get_ideal_point, get_critical_set, get_critical_front,
    get_starting_point;
export sampling_methods;

abstract type Constraints end;
struct Box <:Constraints
    lb :: Vector{Float64}
    ub :: Vector{Float64}
end

@doc """
        abstract type MOP end;

    A supertype for all problem types.
    Each problem must implement the following methods:
    * `num_vars`
    * `num_objectives`
    * `get_objectives`
"""
abstract type MOP end;
Broadcast.broadcastable(m::M where M<:MOP) = Ref(m)

# required
@doc "Return on array of multivariate objective functions."
function get_objectives( :: MOP ) end;
@doc "Return the number of decision variables."
function num_vars( :: MOP) end;
@doc "Return the number of objective functions."
function num_objectives( :: MOP ) end;

# non-mandatory but generically implemented

@doc "Return a constraint object for `mop` or nothing."
function constraints( mop :: M where M <: MOP ) end;
function get_vector_objective( mop :: M where M<:MOP )
    func_array = get_objectives( mop );
    return (x::Vector{R} where R<:Real) -> [ f(x) for f in func_array ];
end

function get_gradients( mop :: M where M <: MOP )
    return [ (x :: Vector{R} where R<:Real ) -> ForwardDiff.gradient( f, x )
        for f in get_objectives(mop) ]
end

function get_omega_function( mop :: M where M <: MOP)
    omega_function = function( x :: Vector{R} where {R<:Real} )
        try
            ∇F = transpose( hcat( ( ∇f(x) for ∇f in get_gradients( mop ) )... ) );

            prob = JuMP.Model( OSQP.Optimizer );

            JuMP.set_silent(prob);
            JuMP.set_optimizer_attribute(prob,"eps_rel",1e-5);
            JuMP.set_optimizer_attribute(prob,"eps_abs", 1e-5);
            JuMP.set_optimizer_attribute(prob,"polish",true);

            JuMP.@variable(prob, α );     # negative of marginal problem value
            JuMP.@variable(prob, d[1:num_vars(mop)] );   # direction vector

            JuMP.@objective(prob, Min, α);

            JuMP.@constraint(prob, descent_contraints, ∇F*d .<= α);
            JuMP.@constraint(prob, norm_constraints, -1.0 .<= d .<= 1.0);

            global_constraints = constraints(mop);
            if isa( global_constraints, Box )
                JuMP.@constraint(prob, box_constraints,  global_constraints.lb .<= x .+ d .<= global_constraints.ub );
            end

            JuMP.optimize!(prob)
            ω = -JuMP.value(α);
            return ω
        catch e
            return -Inf
        end
    end
    return omega_function
end

function get_random_point( mop:: M where M <: MOP )
    if isa( constraints(mop), Box )
        lb = constraints(mop).lb;
        w = constraints(mop).ub .- lb;
        return lb .+ w .* rand( length(lb) );
    end
end

@doc """
Get a starting point for the problem.
If the problem `mop` does not specify its own function, then
a random point is returned for box constrained problems and `nothing` else.
"""
get_starting_point( mop :: M where M<: MOP, args...; kwargs...) = get_random_point(mop);

# non-mandatory and not implemented by default
@doc "Return the ideal (or an utopia) point of the problem."
function get_ideal_point( ::MOP ) end;

@doc """
    get_pareto_set( mop )

Get an object representing the Pareto Set of `mop`.
Possible return types are `FixedPointSet` and `SamplingFunction`.

See also: [`get_points`](@ref), [`get_scatter_points`](@ref), [`SamplingFunction`](@ref),
[`FixedPointSet`](@ref)
"""
function get_pareto_set( :: MOP ) end;

@doc """
    get_pareto_front( mop )

Get an object representing the Pareto Front of `mop`.
Possible return types are `FixedPointSet` and `SamplingFunction`.

See also: [`get_points`](@ref), [`get_scatter_points`](@ref), [`SamplingFunction`](@ref),
[`FixedPointSet`](@ref)
"""
function get_pareto_front( :: MOP ) end;

function get_critical_set( :: MOP ) end;
function get_critical_front( :: MOP ) end;

@doc "Abstract Super Type for Pareto Set and Pareto Frontier."
abstract type CompareSet end;
function get_points( ::CompareSet, ::Int; kwargs...) end;

@doc "A wrapper around a `list_of_points`."
struct FixedPointSet <: CompareSet
    list_of_points :: Vector{Vector{R}} where R<:Real;
end;     # for calculated point clouds

@doc """
    A `SamplingFunction` enables you to sample points from the Pareto Set
    or the Pareto Front.
    The basic usage for `samp_func :: SamplingFunction` is
    ```
    num_points = 10;
    list_of_points = get_points( samp_func, num_points );
    ```

    Check `sampling_methods(samp_func)` to see whether multiple sampling techniques
    are available.
    You can pass a sampling method via `method` keyword argument:
    ```
    list_of_points = get_points( samp_func, num_points; method = :regular );
    ```

    See also: [`sampling_methods`](@ref), [`get_points`](@ref), [`get_scatter_points`](@ref),
"""
@with_kw struct SamplingFunction <: CompareSet
    mop :: M where M <: MOP;
    type :: Symbol;
    methods :: Vector{Symbol};

    @assert type ∈ [:ParetoSet, :ParetoFront, :ParetoCriticalSet, :ParetoCriticalFront]
end

@doc "Return a list of Symbols that can be pased to `get_points` for sampling from
a `SamplingFunction`."
sampling_methods( samp_func :: SamplingFunction ) = samp_func.methods;

@doc """
    get_points( samp_func, N; kwargs... )

Sample `N` points using `samp_func::SamplingFunction`.
Additional `kwargs` are passed down to the problem-specfic implementations.
"""
function get_points( samp_func :: SamplingFunction, N :: Int; kwargs... )
    return get_points( samp_func, samp_func.mop, Val(samp_func.type), N; kwargs...);
end

@doc "
    get_points( point_set, N; kwargs... )

Get all points using `point_set::FixedPointSet`."
function get_points( point_set :: FixedPointSet, N :: Int; kwargs... )
    println("Retrieving *all* points from a `FixedPointSet`, all other arguments are ignored.")
    return point_set.list_of_points;
end

@doc "Like `get_points` but returns data suited for plotting.\n
Specify variable axes by kwarg `dims` which default to `[1,2]`.

See also: [`get_points`](@ref)"
function get_scatter_points( samp_func :: SamplingFunction, N :: Int; dims = [1,2], kwargs... )
    get_scatter_arrays( get_points( samp_func, N; kwargs... ); dims = dims )
end

# not called directly
function get_points( samp_func :: SamplingFunction, mop :: M where M<:MOP,
        ::Val{:ParetoFront}, N; kwargs... )
    pset = get_points( samp_func, samp_func.mop, Val(:ParetoSet), N; kwargs...);
    F = get_vector_objective( mop );
    return F.(pset)
end

# not called directly
function get_points( samp_func :: SamplingFunction, mop :: M where M<:MOP,
    ::Val{:ParetoCriticalFront}, N; kwargs... )
pset = get_points( samp_func, samp_func.mop, Val(:ParetoCriticalSet), N; kwargs...);
F = get_vector_objective( mop );
return F.(pset)
end

# also not called directly
@doc "Return x and y data suited for scatter plots when `arr` is a list of vectors."
function get_scatter_arrays( arr :: Vector{Vector{R}} where R <: Real;
        dims :: Vector{Int} = [1, 2] )
    X = [ point[dims[1]] for point in arr ];
    Y = [ point[dims[2]] for point in arr ];
    if length( dims ) == 2
        return X, Y
    else
        Z = [ point[dims[3]] for point in arr ];
        return X,Y,Z
    end
end

include("two_parabolas.jl");
include("ZDT_problems.jl");
include("DTLZ_problems.jl");
include("lovison_problems.jl")
include("MHT_problems.jl")

end


# construct constant uniform distribution (case 1)
function constant_distr()

  # define constant function
  function uniform(x, p)
    return 0.5f0
  end

  # binary samples
  samples = [0, 1]
  return uniform, samples
end

# construct distribution with jump at p_crit (case 2)
function jump_distr(;p_crit=1.0f0)

  # define step function
  function stepfunc(x, p, p_crit)
    if p <= p_crit
      if x == zero(eltype(x))
        return one(eltype(p))
      else
        return zero(eltype(p))
      end
    else
      if x == zero(eltype(x))
        return zero(eltype(p))
      else
        return one(eltype(p))
      end
    end
  end

  # binary samples
  samples = [0, 1]
  return (x, p) -> stepfunc(x, p, p_crit), samples
end

# construct ``tilted'' distribution with jump at p_crit (case 3)
function tilt_distr(p_range; p_crit = 2.0f0)

  # tilt function
  function tilt(p, p_crit)
    a_left = -1.0f0
    b_left = 5.0f0
    a_right = -1.0f0
    b_right = 2.0f0
    if p <= p_crit
      x = a_left*p+b_left
    else
      x = a_right*p+b_right
    end

    return x
  end

  # sample from tilt function
  samples = map(p -> tilt(p, p_crit), p_range)
  function tilt_distr(x ,p, p_range)
    target = samples[findall(x -> x==p, p_range)[1]]
    if target == x
      return one(eltype(p))
    else
      return zero(eltype(p))
    end
  end
  return (x, p) -> tilt_distr(x, p, p_range), samples
end

# boltzmann factor
function boltzmann_factor_energy(energy, T)
  if T == Inf
    return one(eltype(T))
  else
    return exp(-energy/T)
  end
end

# probability associated with energy at given temperature based on Boltzmann distribution
function get_probabilities(p_range, energies)
  probs = zeros(eltype(p_range[1]), length(p_range), size(energies)[1])
  for indx_p in 1:length(p_range)
    p = p_range[indx_p]
    partition_func = zero(eltype(p))
    for i in 1:size(energies)[1]
      energy = energies[i, 1]
      num = energies[i, 2]
      probs[indx_p, i] = num*boltzmann_factor_energy(energy, p)
      partition_func += probs[indx_p, i]
    end
    if partition_func == Inf
      probs[indx_p, :] = zeros(eltype(p_range[1]), length(probs[indx_p, :]))
    else
      probs[indx_p, :] = probs[indx_p, :]./partition_func
    end
  end
  return probs
end

# construct probability distribution based on individual probabilities
function thermal_distr(x, p, probs, p_range)
  if p == Inf
    p_indx = length(p_range)
  else
    p_indx = Int(round((p-p_range[1])/(p_range[2]-p_range[1])))+1
  end

  prob_distr = probs[p_indx, Int(x)]
  return prob_distr
end

# construct probability distribution of energy in ising model based on exact enumeration of all configurations
function ising_exact_distr(energies, p_range)
  probs = get_probabilities(p_range, energies)
  return (x, p) -> thermal_distr(x ,p, probs, p_range)
end

# probability for specific sample at fixed value of tuning parameter
function distr_approx_x_p(full_data, p_range, x, p)
  p_indx = Int(round((p-p_range[1])/(p_range[2]-p_range[1])))+1
  return full_data[p_indx, x]
end

# construct probability distribution based on Monte Carlo samples
function distr_approx(energies, unique_energies, numbers, p_range)
  full_data =zeros(eltype(p_range[1]),(length(p_range),length(unique_energies)))
  for i in 1:length(p_range)
    numtot = zero(eltype(p_range[1]))
    for j in 1:length(unique_energies)
      number_list = numbers[i]
      energy_list = energies[i]
      sample = unique_energies[j]
      indices = findall(x -> x==sample, energy_list)
      if length(indices) == 0
        full_data[i, j] = zero(eltype(sample))
      else
        full_data[i, j] = number_list[indices[1]]
        numtot += number_list[indices[1]]
      end
    end
    full_data[i, :] = full_data[i, :]/numtot
  end

  return (x, p) -> distr_approx_x_p(full_data, p_range, x, p), collect(1:length(unique_energies))
end


"""conditional_expectation = -cvar = -expected_shortfall"""
function conditional_expectation!(model, w, formulation::AbstractSampleBased;
    sample_probability = fill(1/formulation.number_of_samples,formulation.number_of_samples),
    quantile = 0.95
)
    # parameters
    numA = formulation.number_of_assets
    numS = formulation.number_of_samples
    α = quantile
    # dual variables
    @variable(model, z)
    @variable(model, y[i=1:numS] >= 0)

    @constraints(
        model,
        begin
            ys[s=1:numS], y[s] >= z - sum(formulation.sampled_returns[s, :]'w)
        end
    )
    return z - sum(sample_probability'y) / (1.0 - α)
end

function po_max_conditional_expectation_limit_predicted_return(formulation::AbstractSampleBased, minimal_return;
    rf::Real = 0.0, current_wealth::Real = 1.0,
    model::JuMP.Model = base_model(formulation.number_of_assets; current_wealth=current_wealth),
    w = model[:w],
    quantile = 0.95, kwargs...
)
    # auxilary variables
    sum_invested = create_sum_invested_variable(model, w)

    # model
    @constraint(model, predicted_portfolio_return!(model, w, formulation) + rf * (current_wealth - sum_invested) >= minimal_return * current_wealth)

    # objective function
    @objective(model, Max, conditional_expectation!(model, w, formulation; quantile = quantile, kwargs...))

    return model
end

function po_max_predicted_return_limit_conditional_expectation(formulation::AbstractSampleBased, minimal_return;
    rf::Real = 0.0, current_wealth::Real = 1.0,
    model::JuMP.Model = base_model(formulation.number_of_assets; current_wealth=current_wealth),
    w = model[:w],
    quantile = 0.95, kwargs...
)
    # auxilary variables
    sum_invested = create_sum_invested_variable(model, w)

    # model
    @constraint(model, conditional_expectation!(model, w, formulation; quantile = quantile, kwargs...) >= minimal_return)

    # objective function
    @objective(model, Max, predicted_portfolio_return!(model, w, formulation) + rf * (current_wealth - sum_invested))

    return model
end


# Implement number base conversion

function all_your_base(digits, base_in, base_out)

end


#!/usr/bin/env julia

# Calculate max count for some huffman tree depth.

function fibonacci(n)
	f = zeros(Int, n)
	f[1] = 1
	f[2] = 1
	for i in 3:n
		f[i] = f[i-1] + f[i-2]
	end
	f
end

if length(ARGS) ≠ 1
	println("""
Usage:
	huffman_tree_depth.jl <block_len_log2>

""")
	exit(1)
end

block_len_log2 = parse(Int, ARGS[1])
@show block_len_log2

block_len = 1 << block_len_log2

max_depth = 0
for d in 1:32
	f = fibonacci(d+1)
	max_freq = cumsum(f)[end]
	if max_freq >= block_len
		max_depth = d
		break
	end
end

@show max_depth


using Distributions, RecursiveArrayTools, NLopt, Random
Random.seed!(123)

original_solution = VectorOfArray([(sol(t[i])) for i in 1:length(t)])
original_solution_matrix_form = convert(Array,original_solution)
weight = zeros(size(original_solution_matrix_form)[1],size(original_solution_matrix_form)[2])
error = zeros(size(original_solution_matrix_form)[1],size(original_solution_matrix_form)[2])

for i in 1:size(original_solution_matrix_form)[1]
 for j in 1:size(original_solution_matrix_form)[2]
  tmp = rand()
  weight[i,j] = tmp*tmp
  d = Normal(0,tmp)
  error[i,j] = rand(d)
 end
end

weighted_data = original_solution_matrix_form + error

weighted_cost_function = build_loss_objective(prob1,Tsit5(),
          L2Loss(t,weighted_data,data_weight=weight),maxiters=10000,verbose=false)
opt = Opt(:LN_COBYLA, 1)
min_objective!(opt, weighted_cost_function)
(minf,minx,ret) = NLopt.optimize(opt,[1.3])
@test minx[1] ≈ 1.5 atol=1e-2


function spgr(; M0, angle, TR, R1, R2star = 0.0, TE = 0.0)
    return @. M0 * sin(angle) * exp(-R2star * TE) * (1 - exp(-R1 * TR)) /
              (1 - cos(angle) * exp(-R1 * TR))
end

function conc_to_R1(Ct; r1, R10)
    return @. R10 + r1 * Ct
end

function R1_to_concentration(R1; R10, r1)
    return @. (R1 - R10) / r1
end

function R1_to_signal(R1; M0, angle, TR)
    return spgr(; M0, angle, TR, R1)
end

function signal_to_R1(signal; R10, angle, TR, BAF::Int = 1, mask = [true])
    if typeof(signal) <: AbstractVector
        signal = reshape(signal, 1, length(signal))
        input_was_vector = true
    else
        input_was_vector = false
    end
    volume_size = size(signal)[1:end-1]
    R1 = fill(NaN, size(signal)...)
    resolved_R10 = unify_size(R10, volume_size)
    resolved_mask = resolve_mask_size(mask, volume_size)

    for voxel in findall(resolved_mask)
        R1[voxel, :] = _signal_to_R1(signal[voxel, :], resolved_R10[voxel], angle, TR, BAF)
    end
    if input_was_vector
        R1 = R1[:]
    end
    return R1
end

function _signal_to_R1(signal, R10, angle, TR, BAF)
    S_pre = mean(signal[1:BAF])
    cosα = cos(angle)
    E1 = exp(-TR * R10)
    s = signal ./ S_pre
    log_term = @. (1 - s + s * E1 - cosα * E1) / (1 - s * cosα + s * E1 * cosα - E1 * cosα)
    R1 = [val > 0 ? (-1 / TR) * log(val) : NaN for val in log_term]
    return R1
end

function concentration_to_signal(Ct; r1, M0, angle, TR, R10)
    R1 = concentration_to_R1(Ct; r1, R10)
    signal = R1_to_signal(R1; M0, angle, TR)
    return signal
end

function signal_to_concentration(S; R10, angle, TR, r1, BAF::Int = 1, mask = true)
    R1 = signal_to_R1(S; R10, angle, TR, BAF, mask)
    Ct = R1_to_concentration(R1; R10, r1)
    return Ct
end

function fit_relaxation_despot(;
    signal::AbstractArray,
    angles::AbstractVector,
    TR::Real,
    mask = true,
)
    (signal, angles, mask, num_angles, volume_size) =
        resolve_relaxation_inputs(; signal, angles, mask)
    M0, T1 = (fill(NaN, volume_size...) for _ = 1:2)

    for idx in eachindex(IndexCartesian(), mask)
        if mask[idx] == false
            continue
        end
        M0[idx], T1[idx] = _despot(signal[idx, :], angles, TR)
    end
    return (; est = (; M0, T1))
end

# Refs: Deoni, S. C. L., Peters, T. M., & Rutt, B. K. (2005). MRM 53(1), 237–241. https://doi.org/10.1002/mrm.20314
function _despot(signal::AbstractVector, alpha::AbstractVector, TR)
    x = signal ./ tan.(alpha)
    y = signal ./ sin.(alpha)
    numerator = sum(x .* y) - sum(x) .* sum(y) / length(alpha)
    denominator = sum(x .^ 2) - sum(x) .^ 2 / length(alpha)
    slope = numerator ./ denominator
    intercept = mean(y) - slope * mean(x)
    M0 = intercept / (1 - slope)
    if slope > 0
        T1 = -TR / log(slope)
    else
        T1 = NaN
    end
    return M0, T1
end

function fit_relaxation_novifast(;
    signal::AbstractArray,
    angles::AbstractVector,
    TR::Real,
    mask = true,
)
    (signal, angles, mask, num_angles, volume_size) =
        resolve_relaxation_inputs(; signal, angles, mask)
    M0, T1 = (fill(NaN, volume_size...) for _ = 1:2)

    for idx in eachindex(IndexCartesian(), mask)
        if mask[idx] == false
            continue
        end
        M0[idx], T1[idx] = _novifast(signal[idx, :], angles, TR)
    end
    return (; est = (; M0, T1))
end

# Refs: Ramos-Llorden, G., Vegas-Sanchez-Ferrero, G., Bjork, M., Vanhevel, F., Parizel, P. M., San Jose Estepar, R., … Sijbers, J. (2018). IEEE Transactions on Medical Imaging, 37(11), 2414–2427. https://doi.org/10.1109/TMI.2018.2833288
function _novifast(
    signal::AbstractVector,
    alpha::AbstractVector,
    TR;
    initialvalues = (M0 = 5000, T1 = 1500),
    maxiter = 10,
    tol = 1e-6,
)
    est =
        [initialvalues.M0 * (1 * exp(-TR / initialvalues.T1)), exp(-TR / initialvalues.T1)]
    sinα = sin.(alpha)
    cosα = cos.(alpha)
    k = 0
    relerr = 1e10
    while k <= maxiter && relerr > tol
        previous_est = copy(est)

        commondenom = 1 ./ (1 .- est[2] .* cosα)
        a = signal .* cosα .* commondenom
        b = sinα .* commondenom
        ahat = est[1] .* b .* cosα .* commondenom
        svec = signal .* commondenom

        svec_b = svec' * b
        b_a = b' * a
        svec_ahat = svec' * ahat
        a_ahat = a' * ahat
        b_b = b' * b
        b_ahat = b' * ahat

        num1 = svec_b * a_ahat - b_a * svec_ahat
        num2 = b_b * svec_ahat - svec_b * b_ahat
        denom = b_b * a_ahat - b_ahat * b_a
        est = [num1 / denom, num2 / denom]

        k = k + 1
        rel_err = norm(est .- previous_est) / norm(previous_est)
    end
    M0 = est[1] / (1 - est[2])
    if est[2] > 0
        T1 = -TR / log(est[2])
    else
        T1 = NaN
    end
    return M0, T1
end

function fit_relaxation_nls(;
    signal::AbstractArray,
    angles::AbstractVector,
    TR::Real,
    initialvalues = (M0 = 5000.0, T1 = 1500.0),
    mask = true,
)
    (signal, angles, mask, num_angles, volume_size) =
        resolve_relaxation_inputs(; signal, angles, mask)
    M0, T1 = (fill(NaN, volume_size...) for _ = 1:2)

    @. model(x, p) = _spgr(x, p[1], 1 / p[2], TR)
    for idx in eachindex(IndexCartesian(), mask)
        if mask[idx] == false
            continue
        end
        fit = curve_fit(model, angles, signal[idx, :], [initialvalues.M0, initialvalues.T1])
        M0[idx], T1[idx] = fit.param
    end
    return (; est = (; M0, T1))
end

# [Todo] Replace with spgr() --- but check for speed difference first
function _spgr(angle, M0, R1, TR)
    return @. M0 * sin(angle) * (1 - exp(-R1 * TR)) / (1 - cos(angle) * exp(-R1 * TR))
end

const relaxation_dict = Dict{Symbol,Function}(
    :nls => fit_relaxation_nls,
    :despot => fit_relaxation_despot,
    :novifast => fit_relaxation_novifast,
)

function fit_relaxation(fittype::Symbol = :novifast; kwargs...)
    return relaxation_dict[fittype](; kwargs...)
end


function findk(f, A, B, C, delta, maxk = 170)
    # Compute the total degree of the truncated bivariate Taylor expansion
    # for approximating f{A, B}(C) with accuracy delta, for triangular A and B
    #
    #--------------------------INPUT-------------------------------------------
    #
    # f			handle function for evaluating f and its derivatives: f(x, y, i, j)
    # A,B		mxm and nxn triangular matrix coefficients (likely the Schur forms)
    # C 		mxn right hand side
    # delta     upper bound for the magnitude of the residual
    # maxk		(optional) maximum value for k, default = 170
    #
    #--------------------------------------------------------------------------
    #

    (m, n) = size(C);
    lA = diag(A); lB = diag(B);
    l = sum(lA) / m; mu = sum(lB) / n;
    nrmC = norm(C);
    nrmA = norm(diag(A) .- l, Inf); # opnorm(A - l * one(A), 1);
    nrmB = norm(diag(B) .- mu, Inf); # opnorm(B - mu * one(B), 1);

    k = 0;

    ff = norm([ f(x,y,0,0) for x = lA, y = lB ], Inf)

    while k < maxk
        k = k + 1;
        y = findk_res(f, l, mu, k, nrmA, nrmB);
        if y < delta * ff
            mx = -Inf;
            for x = lA  # TODO it might be done more efficiently if f(x, y, i, j) allows vectorization
                for y = lB
                    tmp = findk_res(f, x, y, k, nrmA, nrmB);
                    if tmp > mx
                        mx = tmp;
                    end
                end
            end
            if mx < delta * ff
                break
            end
        end
    end

    return k
end

#------------------Auxiliary function for findk --------------------------------
function findk_res(f, x, y, k, nrmA, nrmB)
    yy = 0.0;

    fj = 1.0;
    fk = convert(Float64, factorial(big(k+1)));

    for j = 0 : k + 1
        # yy = yy + abs(f(x, y, k + 1 - j, j)) / factorial(k + 1 - j) * nrmA^(k + 1 - j) / factorial(j) * nrmB^j;
        yy = yy + abs(f(x, y, k + 1 - j, j)) / fk * nrmA^(k + 1 - j) / fj * nrmB^j;

        # Update factorials for the next round
        fj = (j+1) * fj;
        fk = fk / (k+1+j);
    end

    return yy;
end

function taylor_eval(f, A, B, C, delta)
    # Returns [X, digits]
    # Evaluate the truncated bivariate taylor expansion of accuracy delta of f{A, B}(C)
    # for triangular coefficients A and B
    #
    #--------------------------INPUT-------------------------------------------
    #
    # f			handle function for evaluating f and its derivatives: f(x, y, i, j)
    # A,B		mxm and nxn triangular matrix coefficients (likely the Schur forms)
    # C 		mxn right hand side
    # delta     upper bound for the magnitude of the residual
    #
    #--------------------------------------------------------------------------
    #

    (m, n) = size(C);
    l = sum(diag(A)) / m; mu = sum(diag(B)) / n;
    k = findk(f, A, B, C, delta);
    digits = k;
    A = A - l * one(A);
    B = B - mu * one(B);

    coef = zeros(ComplexF64, k, k)
    for i = 1 : k
        for j = 1 : k
            if i + j <= k + 1
                coef[i, j] = f(l, mu, i - 1, j - 1)
            end
        end
    end

    # XX at the generic step is A^i / i! * C
    XX = C
    X = zero(C)

    for i = 1 : k

        YY = XX
        for j = 1 : k + 1 - i
            X = X + coef[i,j] * YY
            YY = YY * B / j
        end

        XX = A * XX / i
    end

    return X, digits;
end


#    Problem : GROUP A
#    *********
#    The chained Rosenbrock function as given by Fletcher.
#
#    Original SIF Source:
#    The second problem given by
#    R. Fletcher,
#    "An optimal positive definite update for sparse Hessian matrices"
#    Numerical Analysis report NA/145, University of Dundee, 1992.
#
#    FLETCHR.SIF classification OUR2-AN-V-0
#
#    The Number of variables is variable.
#
# Daniel Henderson, 08/2021

f = (x) -> begin
    fx = 0.0
    for i in firstindex(x):lastindex(x)-1
        γ = x[i+1] - x[i] + 1 - x[i]^2
        fx += 100*γ^2
    end
    return fx
end

g! = (g, x) -> begin
    for i in firstindex(x):lastindex(x)-1
        g[i] += 200*(1 + 2x[i])(-1+x[i]+x[i]^2-x[i+1])
        g[i+1] += 200*(1 - x[i] - x[i]^2 + x[i+1])
    end
    return g
end

fg! = (g, x) -> begin
    fx = 0.0
    for i in firstindex(x):lastindex(x)-1
        γ = x[i+1] - x[i] + 1 - x[i]^2
        fx += 100*γ^2
        g[i] += 200*(-1 - 2x[i])*(1 - x[i] - x[i]^2 + x[1+i])
        g[i+1] += 200*(1 - x[i] - x[i]^2 + x[1+i])
    end
    return fx, g
end

init = (n::Int=1000) -> begin
	n < 2 && @warn("FLETCHCR: number of variables must be ≥ 2")
	n = max(n, 2)

    return n, zeros(n)
end

@warn "TODO: FLETCHCR, determine starting iterate - see whiteboard"
# TestSet["FLETCHCR"] = UncProgram("FLETCHCR", f, g!, fg!, init)


function forward_nolog!(α, c, obs_density::AbstractMatrix, hmm::HMM)
    T, S = size(obs_density, 1), nb_states(hmm)
    π0, P = initial_distribution(hmm), transition_matrix(hmm)
    for i = 1:S
        α[1, i] = π0[i] * obs_density[1, i]
    end
    c[1] = sum(@view α[1, :])  # scaling
    for i = 1:S
        α[1, i] /= c[1]
    end
    for t = 1:T-1
        for j = 1:S
            α[t+1, j] = sum(α[t, i] * P[i, j] for i = 1:S) * obs_density[t+1, j]
        end
        c[t+1] = sum(@view α[t+1, :])  # scaling
        for j = 1:S
            α[t+1, j] /= c[t+1]
        end
    end
    for t = 1:T
        if all_zeros(@view α[t, :])
            throw(OverflowError("Probabilities are too small in forward step."))
        end
    end
    return nothing
end

function forward_log!(logα, obs_logdensity::AbstractMatrix, hmm::HMM)
    T, S = size(obs_logdensity, 1), nb_states(hmm)
    logπ0, logP = log.(initial_distribution(hmm)), log.(transition_matrix(hmm))
    for i = 1:S
        logα[1, i] = logπ0[i] + obs_logdensity[1, i]
    end
    for t = 1:T-1
        for j = 1:S
            logα[t+1, j] = logsumexp(logα[t, :] + logP[:, j]) + obs_logdensity[t+1, j]  # TODO
        end
    end
    for t = 1:T
        if all_nan(@view logα[t, :])
            throw(OverflowError("Log probabilities are too small in forward step."))
        end
    end
    return nothing
end

function backward_nolog!(β, c, obs_density::AbstractMatrix, hmm::HMM)
    T, S = size(obs_density, 1), nb_states(hmm)
    P = transition_matrix(hmm)

    for i = 1:S
        β[T, i] = 0.0
    end
    for t = T-1:-1:1
        for i = 1:S
            β[t, i] = sum(P[i, j] * obs_density[t+1, j] * β[t+1, j] for j = 1:S) / c[t]
        end
    end
    for t = 1:T
        if all_zeros(@view β[t, :])
            throw(OverflowError("Log probabilities are too small in backward step."))
        end
    end
    return nothing
end

function backward_log!(logβ, obs_logdensity::AbstractMatrix, hmm::HMM)
    T, S = size(obs_logdensity, 1), nb_states(hmm)
    logP = log.(transition_matrix(hmm))

    for i = 1:S
        logβ[T, i] = 0.0
    end
    for t = T-1:-1:1
        for i = 1:S
            logβ[t, i] = logsumexp(logP[i, :] .+ obs_logdensity[t+1, :] .+ logβ[t+1, :])  # TODO
        end
    end
    for t = 1:T
        if all_nan(@view logβ[t, :])
            throw(OverflowError("Log probabilities are too small in backward step."))
        end
    end
    return nothing
end

function forward_backward_nolog!(α, β, c, γ, ξ, obs_density::AbstractMatrix, hmm::HMM)
    T, S = size(obs_density, 1), nb_states(hmm)

    forward_nolog!(α, c, obs_density, hmm)
    backward_nolog!(β, c, obs_density, hmm)

    for t = 1:T
        for i = 1:S
            γ[t, i] = α[t, i] * β[t, i]
        end
        sumγ = sum(@view γ[t, :])
        for i = 1:S
            γ[t, i] /= sumγ
        end
    end

    for t = 1:T-1
        for i = 1:S, j = 1:S
            ξ[t, i, j] = α[t, i] * P[i, j] * obs_density[t+1, j] * β[t+1, j]
        end
        sumξ = sum(@view ξ[t, :, :])
        for i = 1:S, j = 1:S
            ξ[t, i, j] /= sumξ
        end
    end

    logL = sum(log.(c))  # TODO: @view

    return log(L)
end

function forward_backward_log!(
    logα,
    logβ,
    logγ,
    logξ,
    obs_logdensity::AbstractMatrix,
    hmm::HMM,
)
    T, S = size(obs_logdensity, 1), nb_states(hmm.transitions)
    logP = log.(transition_matrix(hmm))

    forward_log!(logα, obs_logdensity, hmm)
    backward_log!(logβ, obs_logdensity, hmm)

    for t = 1:T
        for i = 1:S
            logγ[t, i] = logα[t, i] + logβ[t, i]
        end
        logsumγ = logsumexp(@view logγ[t, :])
        for i = 1:S
            logγ[t, i] -= logsumγ
        end
    end

    for t = 1:T-1
        for i = 1:S, j = 1:S
            logξ[t, i, j] = logα[t, i] + logP[i, j] + obs_logdensity[t+1, j] + logβ[t+1, j]
        end
        logsumξ = logsumexp(@view logξ[t, :, :])
        for i = 1:S, j = 1:S
            logξ[t, i, j] -= logsumξ
        end
    end

    logL = logsumexp(logα[T, :])

    return logL
end

## Likelihood of observations

function update_obs_density!(
    obs_density::AbstractMatrix,
    hmm::HMM,
    observations::AbstractVector,
)
    T, S = length(observations), nb_states(hmm)
    for t = 1:T
        for s = 1:S
            obs_density[t, s] = density(emission(hmm, s), observations[t])
        end
    end
    for t = 1:T
        if all_zeros(@view obs_logdensity[t, :])
            throw(OverflowError("Densities are too small for observations."))
        end
    end
end

function update_obs_logdensity!(
    obs_logdensity::AbstractMatrix,
    hmm::HMM,
    observations::AbstractVector,
)
    T, S = length(observations), nb_states(hmm)
    for t = 1:T
        for s = 1:S
            obs_logdensity[t, s] = logdensity(emission(hmm, s), observations[t])
        end
    end
    for t = 1:T
        if all_minus_inf(@view obs_logdensity[t, :])
            throw(OverflowError("Log-densities are too small for observations."))
        end
    end
end

function baum_welch_nolog!(
    α,
    β,
    c,
    γ,
    ξ,
    obs_density::AbstractMatrix,
    hmm::HMM{Tr,Em},
    observations::AbstractVector;
    iterations,
) where {Tr,Em}
    logL_evolution = Float64[]
    for _ = 1:iterations
        update_obs_density!(obs_density, hmm, observations)
        logL = forward_backward_nolog!(α, β, c, γ, ξ, obs_density, hmm)
        push!(logL_evolution, logL)
        new_transitions = fit(Tr, γ = γ, ξ = ξ)
        new_emissions = [fit(Em, observations, γ[:, s]) for s = 1:nb_states(hmm)]  # TODO: @view
        hmm = HMM(new_transitions, new_emissions)
    end
    return hmm, logL_evolution
end

function baum_welch_log!(
    logα,
    logβ,
    logγ,
    logξ,
    obs_logdensity::AbstractMatrix,
    hmm::HMM{Tr,Em},
    observations::AbstractVector;
    iterations,
) where {Tr,Em}
    logL_evolution = Float64[]
    for _ = 1:iterations
        update_obs_logdensity!(obs_logdensity, hmm, observations)
        logL = forward_backward_log!(logα, logβ, logγ, logξ, obs_logdensity, hmm)
        push!(logL_evolution, logL)
        new_transitions = fit(Tr, exp.(logγ), exp.(logξ))
        new_emissions = [fit(Em, observations, exp.(logγ[:, s])) for s = 1:nb_states(hmm)]  # TODO: @view
        hmm = HMM(new_transitions, new_emissions)
    end
    return hmm, logL_evolution
end

## Non mutating version

function baum_welch(hmm::HMM, observations::AbstractVector; iterations, log = true)
    T, S = length(observations), nb_states(hmm)
    if log
        logα = Matrix{Float64}(undef, T, S)
        logβ = Matrix{Float64}(undef, T, S)
        logγ = Matrix{Float64}(undef, T, S)
        logξ = Array{Float64,3}(undef, T - 1, S, S)
        obs_logdensity = Matrix{Float64}(undef, T, S)
        hmm_est, logL_evolution = baum_welch_log!(
            logα,
            logβ,
            logγ,
            logξ,
            obs_logdensity,
            hmm,
            observations;
            iterations = iterations,
        )
    else
        α = Matrix{Float64}(undef, T, S)
        β = Matrix{Float64}(undef, T, S)
        c = Vector{Float64}(undef, T)
        γ = Matrix{Float64}(undef, T, S)
        ξ = Array{Float64,3}(undef, T - 1, S, S)
        obs_density = Matrix{Float64}(undef, T, S)
        hmm_est, logL_evolution = baum_welch_nolog!(
            α,
            β,
            c,
            γ,
            ξ,
            obs_density,
            hmm,
            observations;
            iterations = iterations,
        )
    end
    return hmm_est, logL_evolution
end


using NURBS.Bases

export Curve, BCurve, NCurve, evalpoint, evaltangent, evalderiv1

struct Curve{B<:Basis1D, T<:Real} <: Object1D
    basis::B
    d1b::B
    points::Array{T, 2}

    function Curve(basis::B, points::Array{T, 2}) where {B<:Basis1D, T<:Real}
        d1b = deriv(basis)
        new{B, T}(basis, d1b, points)
    end
end

const BCurve = Curve{BSplineBasis}
const NCurve = Curve{NURBSBasis}

function evalpoint(crv::Curve, u)
    return crv.basis(u, crv.points)
end

function evaltangent(crv::Curve, u)
    t = crv.d1b(u, crv.points)'
    t ./ norm(t)
end

function evalderiv1(crv::Curve, u)
    return crv.d1b(u, crv.points)'
end


using Statistics, LinearAlgebra

# bonds of a petersen graph
const L1 = [(1, 6), (2, 7), (3, 8), (4, 9), (5, 10),
    (1, 2), (2, 3), (3, 4), (4, 5), (1, 5), (6, 8),
    (8, 10), (7, 10), (7, 9), (6, 9)]

# disconnected bonds of a petersen graph
const L2 = [(1, 3), (1, 4), (1, 7), (1, 8), (1, 9),
    (1, 10), (2, 4), (2, 5), (2, 6), (2, 8), (2, 9),
    (2, 10), (3, 5), (3, 6), (3, 7), (3, 9), (3, 10),
    (4, 6), (4, 7), (4, 8), (4, 10), (5, 6), (5, 7),
    (5, 8), (5, 9), (6, 7), (6, 10), (7, 8), (8, 9),
    (9, 10)]

function myvar_and_mean(v)
    mv  = Statistics.mean(v)
    sum((v .- mv).^2)./(length(v)-1), mv
end

function relu(x::T) where T
    x > 0 ? x : zero(T)
end

function embedding_loss(x)
    a = [LinearAlgebra.norm(x[:,i]-x[:,j]) for (i, j) in L1]
    b = [LinearAlgebra.norm(x[:,i]-x[:,j]) for (i, j) in L2]
    va, ma = myvar_and_mean(a)
    vb, mb = myvar_and_mean(b)
    va + vb + exp(relu(-mb + ma + 0.1)) - 1
end


"""
Simplest example of model with 1 variable, 1 shock and 1 transition equation.
It shows the boilerplate code for creating models.
"""
module E1

using ModelBaseEcon

model = Model()

@parameters model begin
    α = 0.5
    β = 0.5
end

@variables model y

@shocks model y_shk

@autoexogenize model y = y_shk

@equations model begin
    y[t] = α * y[t - 1] + β * y[t + 1] + y_shk[t]
end

@initialize model

end


# ------------------------------------------------------------------------------
# diffusion operator example
# ------------------------------------------------------------------------------

using LinearAlgebra

# setup
mesh = Mesh2D(1.0, 1.0)
basis = TensorH1LagrangeBasis(4, 4, 1, 2)

# weak form
function diffusionweakform(du::Array{Float64}, w::Array{Float64})
    dv = du * w[1]
    return [dv]
end

# diffusion operator
inputs = [
    OperatorField(basis, [EvaluationMode.gradient]),
    OperatorField(basis, [EvaluationMode.quadratureweights]),
]
outputs = [OperatorField(basis, [EvaluationMode.gradient])]
diffusion = Operator(diffusionweakform, mesh, inputs, outputs)

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

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


###############################################################################
#
# Surface roughness
#
###############################################################################
"""
    aerodynamic_roughness(h_can::FT) where {FT<:AbstractFloat}

Calculate the aerodynamic roughness for wind, given
- `h_can` Average canopy height
"""
function aerodynamic_roughness(h_can::FT) where {FT<:AbstractFloat}
    return h_can * FT(0.07)
end

###############################################################################
#
# Zero plane displacement
#
###############################################################################
"""
    zero_plane_displacement(h_can::FT) where {FT<:AbstractFloat}

Calculate the zero plane displacement height, given
- `h_can` Average canopy height
"""
function zero_plane_displacement(h_can::FT) where {FT<:AbstractFloat}
    return h_can * 2/3
end

###############################################################################
#
# Calculate wind speed at different height
#
###############################################################################
"""
    wind_speed(us::FT, z0::FT, d::FT, z::FT) where {FT<:AbstractFloat}
    wind_speed(wr::FT, zr::FT, z0::FT, d::FT, z::FT) where {FT<:AbstractFloat}

Calculate the wind speed at height `z`, given
- `us` Friction velocity `[m s⁻¹]`
- `d` Zero plane displacement `[m]`
- `z0` Aerodynamic roughnes `[m]`
- `z` Height `[m]`
- `wr` Wind speed at reference height `[m s⁻¹]`
- `zr` Reference height `[m s⁻¹]`

The equation used is
```math
u(z) = \\dfrac{us}{κ} ⋅ \\log \\left( \\dfrac{z - d}{z0} \\right)
```
where `κ` is the Von Kármán constant.
"""
function wind_speed(
            us::FT,
            z0::FT,
            d::FT,
            z::FT
) where {FT<:AbstractFloat}
    return us / K_VON_KARMAN(FT) * log((z-d)/z0)
end

function wind_speed(
            wr::FT,
            zr::FT,
            z0::FT,
            d::FT,
            z::FT
) where {FT<:AbstractFloat}
    return wr * log((z-d)/z0) / log((zr-d)/z0)
end

###############################################################################
#
# Update wind speed at different canopy layer
#
###############################################################################
"""
    wind_speed!(spac::SPACMono{FT}, us::FT) where {FT<:AbstractFloat}
    wind_speed!(spac::SPACMono{FT}, wr::FT, zr::FT) where {FT<:AbstractFloat}

Update wind speed profile in the canopy, given
- `spac` [`SPACMono`](@ref) type struct
- `us` Friction velocity `[m s⁻¹]`
- `wr` Wind speed at reference height `[m s⁻¹]`
- `zr` Reference height `[m s⁻¹]`
"""
function wind_speed!(
            spac::SPACMono{FT},
            us::FT
) where {FT<:AbstractFloat}
    @unpack air_bounds, n_canopy, wind_d, wind_z0, wind_zs = spac;
    spac.winds .= wind_speed.(us, wind_z0, wind_d, wind_zs);

    return nothing
end

function wind_speed!(
            spac::SPACMono{FT},
            wr::FT,
            zr::FT
) where {FT<:AbstractFloat}
    @unpack air_bounds, n_canopy, wind_d, wind_z0, wind_zs = spac;
    spac.winds .= wind_speed.(wr, zr, wind_z0, wind_d, wind_zs);

    return nothing
end


# TODO: remove initial values for all inputs once IO handling in MTK is in place
"""
    Constant(val; name)

Outputs a constant value `val`.
"""
function Constant(val; name)
    @variables y(t)=val [output=true]
    @parameters val=val
    eqs = [
        y ~ val
    ]
    ODESystem(eqs, t, name=name)
end

"""
    Integrator(; k=1, name)

Outputs `y = ∫k*u dt`, corresponding to the transfer function `1/s`.
"""
function Integrator(; k=1, name)
    @variables x(t)=0 u(t)=0 [input=true] y(t)=0 [output=true]
    @parameters k=k
    eqs = [
        Dₜ(x) ~ k*u
        y ~ x
    ]
    ODESystem(eqs, t, name=name)
end

"""
    Derivative(; k=1, T, name)

Outputs an approximate derivative of the input. The transfer function of this block is
```
k       k
─ - ──────────
T    2 ⎛    1⎞
    T ⋅⎜s + ─⎟
       ⎝    T⎠
```
and a state-space realization is given by `ss(-1/T, 1/T, -k/T, k/T)`
where `T` is the time constant of the filter.
A smaller `T` leads to a more ideal approximation of the derivative.
"""
function Derivative(; k=1, T, name)
    @variables x(t)=0 u(t)=0 [input=true] y(t)=0 [output=true]
    @parameters T=T k=k
    eqs = [
        Dₜ(x) ~ (u - x) / T
        y ~ (k/T)*(u - x)
    ]
    ODESystem(eqs, t, name=name)
end

"""
    FirstOrder(; k=1, T, name)

A first-order filter with a single real pole in `s = -T` and gain `k`. The transfer function
is given by `Y(s)/U(s) = `
```
   k
───────
sT + 1
```
"""
function FirstOrder(; k=1, T, name)
    @variables x(t)=0 u(t)=0 [input=true] y(t) [output=true]
    @parameters T=T k=k
    eqs = [
        Dₜ(x) ~ (-x + k*u) / T
        y ~ x
    ]
    ODESystem(eqs, t, name=name)
end

"""
    SecondOrder(; k=1, w, d, name)

A second-order filter with gain `k`, a bandwidth of `w` rad/s and relative damping `d`. The transfer function
is given by `Y(s)/U(s) = `
```
      k*w^2
─────────────────
s² + 2d*w*s + w^2
```
Critical damping corresponds to `d=1`, which yields the fastest step response without overshoot, d < 1` results in an under-damped filter while `d > 1` results in an over-damped filter.
`d = 1/√2` corresponds to a Butterworth filter of order 2 (maximally flat frequency response).
"""
function SecondOrder(; k=1, w, d, name)
    @variables x(t)=0 xd(t)=0 u(t)=0 [input=true] y(t) [output=true]
    @parameters k=k w=w d=d
    eqs = [
        Dₜ(x) ~ xd
        Dₜ(xd) ~ w*(w*(k*u - x) - 2*d*xd)
        y ~ x
    ]
    ODESystem(eqs, t, name=name)
end

"""
    PID(; k, Ti=false, Td=false, wp=1, wd=1, Ni, Nd=12, y_max=Inf, y_min=-y_max, gains = false, name)

Proportional-Integral-Derivative (PID) controller with output saturation, set-point weighting and integrator anti-windup.

This block has inputs
- `u_r` reference/set point
- `u_y` measurement signal
and output `y` corresponding to the control signal.

The equation for the control signal is roughly
```
k(ep + 1/Ti * ∫e + 1/Td * d/dt(ed))
e = u_r - u_y
ep = wp*u_r - u_y
ed = wd*u_r - u_y
```
where the transfer function for the derivative includes additional filtering, see `? Derivative` for more details.

- `k`: Proportional gain
- `Ti`: Integrator time constant. Set to `false` to turn off integral action.
- `Td`: Derivative time constant. Set to `false` to turn off derivative action.
- `wp`: Set-point weighting in the proportional part.
- `wd`: Set-point weighting in the derivative part.
- `Nd`: Derivative limit, limits the derivative gain to Nd/Td. Reasonable values are ∈ [8, 20]. A higher value gives a better approximation of an ideal derivative at the expense of higher noise amplification.
- `Ni`: `Ni*Ti` controls the time constant `Tₜ` of anti-windup tracking. A common (default) choice is `Tₜ = √(Ti*Td)` which is realized by `Ni = √(Td / Ti)`. Anti-windup can be effectively turned off by setting `Ni = Inf`.
`gains`: If `gains = true`, `Ti` and `Td` will be interpreted as gains with a fundamental PID transfer function on parallel form `ki=Ti, kd=Td, k + ki/s + kd*s`
"""
function PID(; k, Ti=false, Td=false, wp=1, wd=1,
    Ni    = Ti == 0 ? Inf : √(max(Td / Ti, 1e-6)),
    Nd    = 12,
    y_max = Inf,
    y_min = y_max > 0 ? -y_max : -Inf,
    gains = false,
    name
)
    if gains
        Ti = k / Ti
        Td = Td / k
    end
    0 ≤ wp ≤ 1 || throw(ArgumentError("wp out of bounds, got $(wp) but expected wp ∈ [0, 1]"))
    0 ≤ wd ≤ 1 || throw(ArgumentError("wd out of bounds, got $(wd) but expected wd ∈ [0, 1]"))
    Ti ≥ 0     || throw(ArgumentError("Ti out of bounds, got $(Ti) but expected Ti ≥ 0"))
    Td ≥ 0     || throw(ArgumentError("Td out of bounds, got $(Td) but expected Td ≥ 0"))
    y_max ≥ y_min || throw(ArgumentError("y_min must be smaller than y_max"))

    @variables x(t)=0 u_r(t)=0 [input=true] u_y(t)=0 [input=true] y(t) [output=true] e(t)=0 ep(t)=0 ed(t)=0 ea(t)=0

    @named D = Derivative(k = Td, T = Td/Nd) # NOTE: consider T = max(Td/Nd, 100eps()), but currently errors since a symbolic variable appears in a boolean expression in `max`.
    if isequal(Ti, false)
        @named I = Gain(false)
    else
        @named I = Integrator(k = 1/Ti)
    end
    @named sat = Saturation(; y_min, y_max)
    derivative_action = Td > 0
    @parameters k=k Td=Td wp=wp wd=wd Ni=Ni Nd=Nd # TODO: move this line above the subsystem definitions when symbolic default values for parameters works. https://github.com/SciML/ModelingToolkit.jl/issues/1013
    # NOTE: Ti is not included as a parameter since we cannot support setting it to false after this constructor is called. Maybe Integrator can be tested with Ti = false setting k to 0 with IfElse?

    eqs = [
        e ~ u_r - u_y # Control error
        ep ~ wp*u_r - u_y  # Control error for proportional part with setpoint weight
        ea ~ sat.y - sat.u # Actuator error due to saturation
        I.u ~ e + 1/(k*Ni)*ea  # Connect integrator block. The integrator integrates the control error and the anti-wind up tracking. Note the apparent tracking time constant 1/(k*Ni), since k appears after the integration and 1/Ti appears in the integrator block, the final tracking gain will be 1/(Ti*Ni)
        sat.u ~ derivative_action ? k*(ep + I.y + D.y) : k*(ep + I.y) # unsaturated output = P + I + D
        y ~ sat.y
    ]
    systems = [I, sat]
    if derivative_action
        push!(eqs, ed ~ wd*u_r - u_y)
        push!(eqs, D.u ~ ed) # Connect derivative block
        push!(systems, D)
    end
    ODESystem(eqs, t, name=name, systems=systems)
end

"""
    StateSpace(A, B, C, D=0; x0=zeros(size(A,1)), name)

A linear, time-invariant state-space system on the form.
```
ẋ = Ax + Bu
y = Cx + Du
```
Transfer functions can also be simulated by converting them to a StateSpace form.
"""
function StateSpace(A, B, C, D=0; x0=zeros(size(A,1)), name)
    nx = size(A,1)
    nu = size(B,2)
    ny = size(C,1)
    if nx == 0
        length(C) == length(B) == 0 || throw(ArgumentError("Dimension mismatch between A,B,C matrices"))
        return Gain(D; name=name)
    end
    if B isa AbstractVector
        B = reshape(B, length(B), 1)
    end
    if D == 0
        D = zeros(ny, nu)
    end
    @variables x[1:nx](t)=x0 u[1:nu](t)=0 [input=true] y[1:ny](t)=C*x0 [output=true]
    x = collect(x) # https://github.com/JuliaSymbolics/Symbolics.jl/issues/379
    u = collect(u)
    y = collect(y)
    # @parameters A=A B=B C=C D=D # This is buggy
    eqs = [
        Dₜ.(x) .~ A*x .+ B*u
        y      .~ C*x .+ D*u
    ]
    ODESystem(eqs, t, name=name)
end


using CompScienceMeshes
using BEAST
using Base.Test

m = meshrectangle(1.0, 1.0, 0.5)
b = meshsegment(1.0, 0.3, 3)

edges = skeleton(m, 1)

vps = cellpairs(m, edges)
rwg = raviartthomas(m, vps)

nt = ntrace(rwg, b)

#nt = BEAST.ntrace2(rwg, b)

n = 0
Σ = geometry(nt)
on_bnd = overlap_gpredicate(b)
for i in eachindex(rwg.fns)
    length(nt.fns[i]) == 1 || continue
    n += 1
    c = nt.fns[i][1].cellid
    edge = chart(Σ, Σ.faces[c])
    @test on_bnd(edge)
    @test nt.fns[i][1].refid == 1
    @test nt.fns[i][1].coeff == 2.0
end

@test n == 2

## test the scalar trace of a lgrange basis
using CompScienceMeshes
using BEAST
using Base.Test
m = meshrectangle(1.0, 1.0, 0.5)
b = meshsegment(1.0, 0.3, 3)

X = lagrangec0d1(m,b)

@test numfunctions(X) == 2
@test length(X.fns[1]) == 3
@test length(X.fns[2]) == 6

Y = strace(X,b)
A = find(length(f) for f in Y.fns)

@test length(A) == 1
@test length(Y.fns[1]) == 2
@test Y.fns[1][1].coeff ≈ 1
@test Y.fns[1][2].coeff ≈ 1


@testset "$(rpad("Vandermonde Matrix Tests",80))" begin

    n = 11
    x = collect(range(0., stop=1., length=n))
    y = rand(n)

    V = vandermonde_matrix(x)
    Vinv = vandermonde_matrix_inverse(x)

    x = V\y
    a = *(Vinv, y)

    @test maximum(abs.((a .- x) ./ a)) < eps(Float32)

end


@testset "reshape" begin
    p = randn(2,3)
    x = Var(p)
    y = reshape(x, 6)
    @test reshape(p, 6) == y.data
    gradient!(y)
    @test size(x) == size(x.grad)
end

@testset "matmul" begin
    w = Var([1 2 3; 4 5 6])
    x = Var(transpose(w.data))
    y = w * x
    @test y.data == [14 32; 32 77]

    x = randn(3, 2)
    w = randn(2, 3)
    f = x -> x * Var(w)
    @test gradcheck(f, x)

    x_data = randn(10, 1)
    w_data = randn(1, 5)
    f = w -> matmul(Var(x_data), w)
    @test gradcheck(f, w_data)
end