using ReverseDiff
using MNIST

type Dataset
  X::AbstractArray
  y::AbstractArray
end

type MNISTData
  train::Dataset
  dev::Dataset
  test::Dataset
end

function onehot{T}(y::Vector{T}, n)
  yy = zeros(T, n, length(y))
  for i=1:length(y)
    label = y[i]
    yy[Int(label)+1, i] = 1
  end
  return yy
end

function load_mnist()
  X, y = MNIST.traindata()
  X_train, y_train = X[:, 1:50000], y[1:50000]
  X_dev, y_dev = X[:, 50001:end], y[50001:end]
  X_test, y_test = MNIST.testdata()
  return MNISTData(
    Dataset(X_train/255, onehot(y_train, 10)),
    Dataset(X_dev/255, onehot(y_dev, 10)),
    Dataset(X_test/255, onehot(y_test, 10))
  )
end

function sgd_update(v, Δv, α)
  for i in eachindex(v)
    v[i] -= α * Δv[i]
  end
end

data = load_mnist()

const BATCH_SIZE = 100
const STEPS_PER_EPOCH = floor(size(data.train.X, 2) / BATCH_SIZE)
const EPOCHS = 20
α = 1e-3

# parameters
W = randn(Float64, 10, 784) * 0.01
b = zeros(Float64, 10)

inputs = (zeros(Float64, 784, BATCH_SIZE), zeros(Float64, 10, BATCH_SIZE), W, b)
outputs = map(similar, inputs)

ReverseDiff.@forward minus_log(x) = -log(x)
softmax(x) = (exp_x = exp.(x); exp_x ./ sum(exp_x, 1))
cross_entropy_loss(x, y) = mean(sum(y .* minus_log.(x), 1))

model(x, y, W, b) = cross_entropy_loss(softmax(W*x .+ b), y)
predict(x, W, b) = W*x .+ b
accuracy(y, ŷ) = mean(mapslices(indmax, y, 1) .== mapslices(indmax, ŷ, 1))

@time model! = ReverseDiff.compile_gradient(model, inputs)

for i=1:EPOCHS
  @time for j=1:STEPS_PER_EPOCH
    s = Int((j-1) * BATCH_SIZE + 1)
    e = Int(j * BATCH_SIZE)
    X_batch = @view data.train.X[:, s:e]
    y_batch = @view data.train.y[:, s:e]
    inputs = (X_batch, y_batch, W, b)
    model!(outputs, inputs)
    sgd_update(W, outputs[3], α)
    sgd_update(b, outputs[4], α)
  end

  dev_acc = accuracy(predict(data.dev.X, W, b), data.dev.y)
  println("Epoch $(i), Dev accuracy = $(dev_acc)")
end

test_acc = accuracy(predict(data.test.X, W, b), data.test.y)
println("Test accuracy = $(test_acc)")


using Test
using ClimateMachine

let
    for N in 1:10
        for FloatType in (Float64, Float32)
            (ξ, ω) = ClimateMachine.Mesh.Elements.lglpoints(FloatType, N)
            I∫ =
                ClimateMachine.Mesh.Grids.indefinite_integral_interpolation_matrix(
                    ξ,
                    ω,
                )
            for n in 1:N
                if N == 1
                    @test sum(abs.(I∫ * ξ)) < 10 * eps(FloatType)
                else
                    @test I∫ * ξ .^ n ≈
                          (ξ .^ (n + 1) .- (-1) .^ (n + 1)) / (n + 1)
                end
            end
        end
    end
end


@doc raw"""
	peak_factor_cl56(m::S, r_ps::T, fas::FourierParameters, sdof::Oscillator; nodes::Int=35) where {S<:Real,T<:Real}

Peak factor computed using the Cartwright and Longuet-Higgins (1956) formulation.

Evaluates the integral to obtain the peak factor ``\psi``:
```math
\psi = \sqrt{2} \int_0^\infty 1 - \left( 1 - \xi \exp\left( -z^2 \right)\right)^{n_e} dz
```
where ``n_e`` is the number of extrema, ``\xi`` is the ratio ``n_z/n_e`` with ``n_z`` being the number of zero crossings.

The integral is evaluated using Gauss-Legendre integration -- and is suitable for use within an automatic differentiation environment.

See also: [`peak_factor`](@ref), [`peak_factor_dk80`](@ref)
"""
function peak_factor_cl56(m::S, r_ps::T, fas::FourierParameters, sdof::Oscillator; nodes::Int=35) where {S<:Real,T<:Real}
	# get the numbers of zero crossing and extrema
	rvt = RandomVibrationParameters(:CL56)
	n_z, n_e = zeros_extrema_numbers(m, r_ps, fas, sdof, rvt)
	# get the ratio of zero crossings to extrema
	ξ = n_z / n_e

	# compute the Gauss Legendre nodes and weights
	xi, wi = gausslegendre(nodes)

	z_min = 0.0
	z_max = 8.0
	dfac = (z_max-z_min)/2
	pfac = (z_max+z_min)/2

	# scale the nodes to the appropriate range
	zi = @. dfac * xi + pfac

	# define the integrand
	integrand(z) = 1.0 - (1.0 - ξ*exp(-z^2))^n_e

	int_sum = dfac * dot( wi, integrand.(zi) )

	pf = sqrt(2.0) * int_sum
	return pf
end

@doc raw"""
	peak_factor_cl56(m::S, r_ps::T, Dex::U, m0::V, fas::FourierParameters, sdof::Oscillator; nodes::Int=35) where {S<:Real,T<:Real,U<:Real,V<:Real}

Peak factor computed using the Cartwright and Longuet-Higgins (1956) formulation, using pre-computed `Dex` and `m0` values.

Evaluates the integral to obtain the peak factor ``\psi``:
```math
\psi = \sqrt{2} \int_0^\infty 1 - \left( 1 - \xi \exp\left( -z^2 \right)\right)^{n_e} dz
```
where ``n_e`` is the number of extrema, ``\xi`` is the ratio ``n_z/n_e`` with ``n_z`` being the number of zero crossings.

The integral is evaluated using Gauss-Legendre integration -- and is suitable for use within an automatic differentiation environment.

See also: [`peak_factor`](@ref), [`peak_factor_dk80`](@ref)
"""
function peak_factor_cl56(m::S, r_ps::T, Dex::U, m0::V, fas::FourierParameters, sdof::Oscillator; nodes::Int=35) where {S<:Real,T<:Real,U<:Real,V<:Real}
	# get all necessary spectral moments
	mi = spectral_moments([2, 4], m, r_ps, fas, sdof)
	m2 = mi[1]
	m4 = mi[2]
	# get the peak factor
	n_z = Dex * sqrt( m2 / m0 ) / π
	n_e = Dex * sqrt( m4 / m2 ) / π
	# get the ratio of zero crossings to extrema
	ξ = n_z / n_e

	# compute the Gauss Legendre nodes and weights
	xi, wi = gausslegendre(nodes)

	z_min = 0.0
	z_max = 8.0
	dfac = (z_max-z_min)/2
	pfac = (z_max+z_min)/2

	# scale the nodes to the appropriate range
	zi = @. dfac * xi + pfac

	# define the integrand
	integrand(z) = 1.0 - (1.0 - ξ*exp(-z^2))^n_e

	int_sum = dfac * dot( wi, integrand.(zi) )

	pf = sqrt(2.0) * int_sum
	return pf
end

@doc raw"""
	peak_factor_cl56(n_z::T, n_e::T; nodes::Int=35) where T<:Real

Peak factor computed using the Cartwright and Longuet-Higgins (1956) formulation, using pre-computed `n_z` and `n_e` values.

Evaluates the integral to obtain the peak factor ``\psi``:
```math
\psi = \sqrt{2} \int_0^\infty 1 - \left( 1 - \xi \exp\left( -z^2 \right)\right)^{n_e} dz
```
where ``n_e`` is the number of extrema, ``\xi`` is the ratio ``n_z/n_e`` with ``n_z`` being the number of zero crossings.

The integral is evaluated using Gauss-Legendre integration -- and is suitable for use within an automatic differentiation environment.

See also: [`peak_factor`](@ref), [`peak_factor_dk80`](@ref)
"""
function peak_factor_cl56(n_z::T, n_e::T; nodes::Int=35) where T<:Real
	# get the ratio of zero crossings to extrema
	ξ = n_z / n_e

	# compute the Gauss Legendre nodes and weights
	xi, wi = gausslegendre(nodes)

	z_min = 0.0
	z_max = 8.0
	dfac = (z_max-z_min)/2
	pfac = (z_max+z_min)/2

	# scale the nodes to the appropriate range
	zi = @. dfac * xi + pfac

	# define the integrand
	integrand(z) = 1.0 - (1.0 - ξ*exp(-z^2))^n_e

	int_sum = dfac * dot( wi, integrand.(zi) )

	pf = sqrt(2.0) * int_sum
	return pf
end

@doc raw"""
	peak_factor_cl56_gk(m::S, r_ps::T, fas::FourierParameters, sdof::Oscillator) where {S<:Real,T<:Real}

Peak factor computed using the Cartwright and Longuet-Higgins (1956) formulation, using adaptive Gauss-Kronrod integration.

Evaluates the integral to obtain the peak factor ``\psi``:
```math
\psi = \sqrt{2} \int_0^\infty 1 - \left( 1 - \xi \exp\left( -z^2 \right)\right)^{n_e} dz
```
where ``n_e`` is the number of extrema, ``\xi`` is the ratio ``n_z/n_e`` with ``n_z`` being the number of zero crossings.

The integral is evaluated using Gauss-Kronrod integration via the QuadGK.jl package -- and is not suitable for use within an automatic differentiation environment.

See also: [`peak_factor`](@ref), [`peak_factor_dk80`](@ref)
"""
function peak_factor_cl56_gk(m::S, r_ps::T, fas::FourierParameters, sdof::Oscillator) where {S<:Real,T<:Real}
	# get the numbers of zero crossing and extrema
	rvt = RandomVibrationParameters(:CL56)
	n_z, n_e = zeros_extrema_numbers(m, r_ps, fas, sdof, rvt)
	# get the ratio of zero crossings to extrema
	ξ = n_z / n_e

	# define the integrand
	integrand(z) = 1.0 - (1.0 - ξ*exp(-z^2))^n_e

	return sqrt(2.0) * quadgk(integrand, 0.0, Inf)[1]
end

"""
	peak_factor_integrand_cl56(z, n_z, n_e)

Integrand of the Cartwright Longuet-Higgins (1956) peak factor formulation.

See also: [`peak_factor_cl56`](@ref)
"""
function peak_factor_integrand_cl56(z, n_z, n_e)
  ξ = n_z / n_e
  return 1.0 - (1.0 - ξ*exp(-z^2))^n_e
end

"""
	vanmarcke_cdf(x, n_z::T, δeff::T) where T<:Real

CDF of the Vanmarcke (1975) peak factor distribution -- used within the Der Kiureghian (1980) peak factor expression.
"""
function vanmarcke_cdf(x, n_z::T, δeff::T) where T<:Real
	if x < eps()
		return zero(T)
	else
		xsq = x^2
		eδe = exp( -sqrt(π/2) * δeff * x )
		Fx = ( 1.0 - exp(-xsq/2) ) * exp( -n_z * (1.0 - eδe)/(exp(xsq/2) - 1.0) )
		return Fx
	end
end

"""
	vanmarcke_ccdf(x, n_z::T, δeff::T) where T<:Real

CCDF of the Vanmarcke (1975) peak factor distribution -- used within the Der Kiureghian (1980) peak factor expression.
"""
function vanmarcke_ccdf(x, n_z::T, δeff::T) where T<:Real
	return 1.0 - vanmarcke_cdf(x, n_z, δeff)
end

@doc raw"""
	peak_factor_dk80(m::S, r_ps::T, fas::FourierParameters, sdof::Oscillator; nodes::Int=30) where {S<:Real,T<:Real}

Peak factor computed using the Der Kiureghian (1980)/Vanmarcke (1975) formulation.

Evaluates the integral to obtain the peak factor ``\psi``:
```math
\psi = \int_0^\infty 1 - F_Z(z)dz
```
where ``F_Z(z)`` is the CDF defined by:
```math
F_Z(z) = \left[ 1-\exp\left(-\frac{z^2}{2}\right) \right] \exp\left[ -n_z \frac{1 - \exp\left( -\sqrt{\frac{\pi}{2}}\delta_{eff} z \right)}{\exp\left(\frac{z^2}{2}\right) - 1} \right]
```

The effective bandwidth ``\delta_{eff}=\delta^{1.2}`` where the bandwidth is:
```math
\delta = \sqrt{ 1 - \frac{m_1^2}{m_0 m_2} }
```

The integral is evaluated using Gauss-Legendre integration -- and is suitable for use within an automatic differentiation environment.

See also: [`peak_factor`](@ref), [`peak_factor_cl56`](@ref)
"""
function peak_factor_dk80(m::S, r_ps::T, fas::FourierParameters, sdof::Oscillator; nodes::Int=30) where {S<:Real,T<:Real}
	# compute first three spectral moments
	mi = spectral_moments([0, 1, 2], m, r_ps, fas, sdof)
	m0 = mi[1]
	m1 = mi[2]
	m2 = mi[3]

	# bandwidth, and effective bandwidth
	δ = sqrt( 1.0 - (m1^2 / (m0*m2)) )
	δeff = δ^1.2

	# excitation duration
	Dex = boore_thompson_2014(m, r_ps, fas)
	# number of zero crossings
	n_z = Dex * sqrt( m2/m0 ) / π

	# compute the Gauss Legendre nodes and weights
	xi, wi = gausslegendre(nodes)

	z_min = 0.0
	z_max = 6.0
	dfac = (z_max-z_min)/2
	pfac = (z_max+z_min)/2

	# scale the nodes to the appropriate range
	zi = @. dfac * xi + pfac
	# evaluate integrand
	yy = vanmarcke_ccdf.(zi, n_z, δeff)

	return int_sum = dfac * dot( wi, yy )
end

@doc raw"""
	peak_factor_dk80(m::S, r_ps::T, Dex::U, m0::V, fas::FourierParameters, sdof::Oscillator; nodes::Int=30) where {S<:Real,T<:Real,U<:Real,V<:Real}

Peak factor computed using the Der Kiureghian (1980)/Vanmarcke (1975) formulation, using precomputed `Dex` and `m0`.

Evaluates the integral to obtain the peak factor ``\psi``:
```math
\psi = \int_0^\infty 1 - F_Z(z)dz
```
where ``F_Z(z)`` is the CDF defined by:
```math
F_Z(z) = \left[ 1-\exp\left(-\frac{z^2}{2}\right) \right] \exp\left[ -n_z \frac{1 - \exp\left( -\sqrt{\frac{\pi}{2}}\delta_{eff} z \right)}{\exp\left(\frac{z^2}{2}\right) - 1} \right]
```

The effective bandwidth ``\delta_{eff}=\delta^{1.2}`` where the bandwidth is:
```math
\delta = \sqrt{ 1 - \frac{m_1^2}{m_0 m_2} }
```

The integral is evaluated using Gauss-Legendre integration -- and is suitable for use within an automatic differentiation environment.

See also: [`peak_factor`](@ref), [`peak_factor_cl56`](@ref)
"""
function peak_factor_dk80(m::S, r_ps::T, Dex::U, m0::V, fas::FourierParameters, sdof::Oscillator; nodes::Int=30) where {S<:Real,T<:Real,U<:Real,V<:Real}
	# compute first three spectral moments
	mi = spectral_moments([1, 2], m, r_ps, fas, sdof)
	m1 = mi[1]
	m2 = mi[2]

	# bandwidth, and effective bandwidth
	δ = sqrt( 1.0 - (m1^2 / (m0*m2)) )
	δeff = δ^1.2

	# number of zero crossings
	n_z = Dex * sqrt( m2/m0 ) / π

	# compute the Gauss Legendre nodes and weights
	xi, wi = gausslegendre(nodes)

	z_min = 0.0
	z_max = 6.0
	dfac = (z_max-z_min)/2
	pfac = (z_max+z_min)/2

	# scale the nodes to the appropriate range
	zi = @. dfac * xi + pfac
	# evaluate integrand
	yy = vanmarcke_ccdf.(zi, n_z, δeff)

	return int_sum = dfac * dot( wi, yy )
end

@doc raw"""
	peak_factor_dk80_gk(m::S, r_ps::T, fas::FourierParameters, sdof::Oscillator) where {S<:Float64,T<:Real}

Peak factor computed using the Der Kiureghian (1980)/Vanmarcke (1975) formulation, using precomputed `Dex` and `m0`.

Evaluates the integral to obtain the peak factor ``\psi``:
```math
\psi = \int_0^\infty 1 - F_Z(z)dz
```
where ``F_Z(z)`` is the CDF defined by:
```math
F_Z(z) = \left[ 1-\exp\left(-\frac{z^2}{2}\right) \right] \exp\left[ -n_z \frac{1 - \exp\left( -\sqrt{\frac{\pi}{2}}\delta_{eff} z \right)}{\exp\left(\frac{z^2}{2}\right) - 1} \right]
```

The effective bandwidth ``\delta_{eff}=\delta^{1.2}`` where the bandwidth is:
```math
\delta = \sqrt{ 1 - \frac{m_1^2}{m_0 m_2} }
```

The integral is evaluated using adaptive Gauss-Kronrod integration from the QuadGK.jl package -- and is therefore not suitable for use within an automatic differentiation environment.

See also: [`peak_factor`](@ref), [`peak_factor_dk80`](@ref), [`peak_factor_cl56`](@ref)
"""
function peak_factor_dk80_gk(m::S, r_ps::T, fas::FourierParameters, sdof::Oscillator) where {S<:Float64,T<:Real}
	# compute first three spectral moments
	mi = spectral_moments([0, 1, 2], m, r_ps, fas, sdof)
	m0 = mi[1]
	m1 = mi[2]
	m2 = mi[3]

	# bandwidth, and effective bandwidth
	δ = sqrt( 1.0 - (m1^2 / (m0*m2)) )
	δeff = δ^1.2

	# excitation duration
	Dex = boore_thompson_2014(m, r_ps, fas)
	# number of zero crossings
	n_z = Dex * sqrt( m2/m0 ) / π

	integrand(x) = vanmarcke_ccdf(x, n_z, δeff)

	return quadgk(integrand, 0.0, Inf)[1]
end

@doc raw"""
	peak_factor(m::S, r_ps::T, fas::FourierParameters, sdof::Oscillator; pf_method::Symbol=:DK80) where {S<:Real,T<:Real}

Peak factor ``u_{max} / u_{rms}`` with a switch of `pf_method` to determine the approach adopted.
`pf_method` can currently be one of:
	- `:CL56` for Cartright Longuet-Higgins (1956)
	- `:DK80` for Der Kiureghian (1980), building on Vanmarcke (1975)

Defaults to `:DK80`.
"""
function peak_factor(m::S, r_ps::T, fas::FourierParameters, sdof::Oscillator, rvt::RandomVibrationParameters) where {S<:Real,T<:Real}
	if rvt.pf_method == :CL56
		return peak_factor_cl56(m, r_ps, fas, sdof)
	elseif rvt.pf_method == :DK80
		return peak_factor_dk80(m, r_ps, fas, sdof)
	else
		U = get_parametric_type(fas)
		return U(NaN)
	end
end

"""
	peak_factor(m::S, r_ps::T, Dex::U, m0::V, fas::FourierParameters, sdof::Oscillator, rvt::RandomVibrationParameters) where {S<:Real,T<:Real,U<:Real,V<:Real}

Peak factor u_max / u_rms with a switch of `pf_method` to determine the approach adopted.
`pf_method` can currently be one of:
	- `:CL56` for Cartright Longuet-Higgins (1956)
	- `:DK80` for Der Kiureghian (1980), building on Vanmarcke (1975)

Defaults to `:DK80`.
"""
function peak_factor(m::S, r_ps::T, Dex::U, m0::V, fas::FourierParameters, sdof::Oscillator, rvt::RandomVibrationParameters) where {S<:Real,T<:Real,U<:Real,V<:Real}
	if rvt.pf_method == :CL56
		return peak_factor_cl56(m, r_ps, Dex, m0, fas, sdof)
	elseif rvt.pf_method == :DK80
		return peak_factor_dk80(m, r_ps, Dex, m0, fas, sdof)
	else
		if S <: Dual
			return S(NaN)
		elseif V <: Dual
			return V(NaN)
		else
			return U(NaN)
		end
	end
end


using Combinatorics

function part1(input)
    D = parse_input(input)
    return foldl(max, (total_happiness(D, order) for order in permutations(1:8)))
end

function part2(input)
    D = parse_input(input)
    return foldl(max, (total_happiness2(D, order) for order in permutations(1:8)))
end

function parse_input(input)
    D = fill(0, 8, 8)
    lines = eachline(input)
    for i = 1:8
        for j = 1:7
            words = split(first(lines))
            happiness = parse(Int, words[4])
            if words[3] == "lose"
                happiness *= -1
            end
            D[i, j + (j >= i)] += happiness
            D[j + (j >= i), i] += happiness
        end
    end
    return D
end

total_happiness(D, order) = sum(D[order[i], order[mod1(i + 1, 8)]] for i = 1:8)
total_happiness2(D, order) = sum(D[order[i], order[i + 1]] for i = 1:7)


mutable struct FGMState{VT, IT}
    iter::Integer
    x::VT
    xₙ::VT
    y::VT
    yₙ::VT
    ∇y::VT
    ∇xₙ::VT
    t::VT
    stepstate::IT
end

struct FastGradientMethodIterator{MT, VT, IT, T}
    H::MT
    b::VT
    x₀::VT
    proj::Function
    step::IT
    L::T
end

function iterate( fgm::FastGradientMethodIterator )
    VT = typeof( fgm.x₀ )
    ET = eltype( fgm.x₀ )

    len = length( fgm.x₀ )
    zer = fill( zero(ET), (len, ) )

    # Setup the step length computation iterator
    (β, stepstate) = iterate( fgm.step )

    IT = typeof( stepstate )

    state = FGMState{VT, IT}( 0, fgm.x₀, fgm.x₀, fgm.x₀, fgm.x₀, zer, zer, zer, stepstate )

    # Run the first pass
    return iterate( fgm, state )
end

function iterate( fgm::FastGradientMethodIterator, state::FGMState )
    state.iter = state.iter + 1

    # Copy the previous values into the current values
    state.x = state.xₙ
    state.y = state.yₙ

    # Compute the gradient
    state.∇y = fgm.H*state.y + fgm.b

    # Compute the new step
    state.t = state.y - (1/fgm.L) * state.∇y

    # Project into the feasible space
    state.xₙ = fgm.proj( state.t )

    # Compute the step length and apply it
    (β, state.stepstate) = iterate( fgm.step, state.stepstate )

    state.yₙ = state.xₙ + β*( state.xₙ - state.x )

    # Compute the residual of the next point for the stopping condition computation
    state.∇xₙ = fgm.H*state.xₙ + fgm.b

    return (state, state)
end

function header()
    @printf "Iteration | Residual (Conjugate)\n"
end

function disp( state::FGMState )
    @printf "   %3d    |   %.4e\n" state.iter compute_conjugate( state )
end

"""
```julia
    fastgradientmethod( H::AbstractMatrix{T}, b::AbstractVector{T}, proj::Function; x₀::AbstractVector{T} = zeros( T, size( b ) ),
                                                                                L::Union{Nothing, T} = nothing,
                                                                                μ::Union{Nothing, T} = nothing,
                                                                                step::AbstractStep = ConstantStep(),
                                                                                stopconditions::Vector{SC} = [Best(1e-4)],
                                                                                maxiter::Integer = 100,
                                                                                dispinterval::Integer = 1
                      ) where {T <: Number, SC <: AbstractStopCondition}
```

Compute the solution to the quadratic program ``min xᵀHx + bᵀx s.t. x ∈ χ`` where the set constraint
``χ`` is supplied as a projection operation that maps a value of ``x`` to a value inside the set ``χ``.
The projection operator is specified as the function `proj`, which must take a single argument that is
a vector of type `T` (the current iterate) and return a vector of type `T` (the projection of the iterate
onto the set ``Χ``).

The behavior of the algorithm can be modified using the keyword arguments to set various algorithm parameters:
* x₀ - The initial iterate
* L - User-specified value for the maximum eigenvalue of `H`. If not specified, then the initial phase of the
      algorithm will perform an eigenvalue decomposition on `H` and set `L` to its maximum eigenvalue.
* μ - User-specified value for the minimum eigenvalue of `H`. If not specified, then the initial phase of the
      algorithm will perform an eigenvalue decomposition on `H` and set `L` to its minimum eigenvalue.
* step - How the step length should be computed at each iteration.
* stopconditions - A vector containing the stopping conditions that the algorithm uses to determine convergence.
                   A boolean OR of all the stopping conditions will be used, so when the first one is met, the
                   algorithm will terminate. Set to `nothing` to only use an iteration bound.
* maxiter - The maximum number of iterations to perform
* dispinterval - Print iteration information every `dispinterval` iterations. Set to 0 to disable printing.
"""
function fastgradientmethod( H::AbstractMatrix{T}, b::AbstractVector{T}, proj::Function; x₀::AbstractVector{T} = zeros( T, size( b ) ),
                                                                                L::Union{Nothing, T} = nothing,
                                                                                μ::Union{Nothing, T} = nothing,
                                                                                step::AbstractStep = ConstantStep(),
                                                                                stopconditions::Union{Vector{SC}, Nothing} = [Best(1e-4)],
                                                                                maxiter::Integer = 100,
                                                                                dispinterval::Integer = 1
                           ) where {T <: Number, SC <: AbstractStopCondition}

    # Compute the lipschitz and convexity parameter if they aren't provided
    if isnothing( L ) || isnothing( μ )
        eig = eigvals( H )

        isnothing( L ) && ( L = maximum( eig ) )
        isnothing( μ ) && ( μ = minimum( eig ) )
    end

    n = size( b, 1 )

    init!( step, L, μ )

    iter = 0

    # Create the actual iterator that we use
    MT = typeof( H )
    VT = typeof( b )
    IT = typeof( step )
    LT = typeof( L )

    iter = FastGradientMethodIterator{MT, VT, IT, LT}( H, b, x₀, proj, step, L )

    # Setup the stopping conditions
    if stopconditions != nothing
        for cond in stopconditions
            init!( cond, n, L, μ )
        end

        iter = halt( iter, stopconditions )
    end

    iter = take( iter, maxiter )
    iter = enumerate( iter )

    # Configure the printing of iteration information
    if dispinterval > 0
        iter = apply( iter, disp, dispinterval )
        header()
    end

    # Actually loop the algorithm
    (_, finalstate) = loop( iter )

    return (finalstate.xₙ, finalstate.iter)
end


#=

  Magic hexagon in ConstraintSolver.jl

  Prob023: Magic Hexagon
  http://www.comp.rgu.ac.uk/staff/ha/ZCSP/prob023/prob023.pdf
  http://www.cse.unsw.edu.au/~tw/csplib/prob/prob023/

  Model created by Hakan Kjellerstrand, hakank@gmail.com
  See also my Julia page: http://www.hakank.org/julia/

=#

using ConstraintSolver, JuMP
using Cbc, GLPK, Ipopt
const CS = ConstraintSolver
include("constraints_utils.jl")

function magic_hexagon(print_solutions=true,all_solutions=true,timeout=6)

    cbc_optimizer = optimizer_with_attributes(Cbc.Optimizer, "logLevel" => 0)
    glpk_optimizer = optimizer_with_attributes(GLPK.Optimizer)
    ipopt_optimizer = optimizer_with_attributes(Ipopt.Optimizer)

    model = Model(optimizer_with_attributes(CS.Optimizer,   "all_solutions"=> all_solutions,
                                                            # "all_optimal_solutions"=>all_solutions,
                                                            "logging"=>[],

                                                            "traverse_strategy"=>:BFS,
                                                            # "traverse_strategy"=>:DFS,
                                                            # "traverse_strategy"=>:DBFS,

                                                            # "branch_split"=>:Smallest,
                                                            # "branch_split"=>:Biggest,
                                                            "branch_split"=>:InHalf,

                                                            # https://wikunia.github.io/ConstraintSolver.jl/stable/options/#branch_strategy-(:Auto)
                                                            "branch_strategy" => :IMPS, # default
                                                            # "branch_strategy" => :ABS, # Activity Based Search
                                                            # "activity.decay" => 0.999, # default 0.999
                                                            # "activity.max_probes" => 10, # default, 10
                                                            # "activity.max_confidence_deviation" => 20, # default 20

                                                            # "simplify"=>false,
                                                            # "simplify"=>true, # default

                                                            "time_limit"=>timeout,

                                                            # "backtrack" => false, # default true
                                                            # "backtrack_sorting" => false, # default true

                                                            # "lp_optimizer" => cbc_optimizer,
                                                            # "lp_optimizer" => glpk_optimizer,
                                                            # "lp_optimizer" => ipopt_optimizer,
                                        ))

    n = 19
    @variable(model, 1 <= x[1:n] <= n, Int)
    A, B, C, D, E, F, G, H, I, J, K, L, M, N, O, P, Q, R, S = x

    @constraint(model, x in CS.AllDifferent())

    @constraint(model, A + B + C ==  38)
    @constraint(model, D + E + F + G ==  38)
    @constraint(model, H + I + J + K + L ==  38,)
    @constraint(model, M + N + O + P ==  38,)
    @constraint(model, Q + R + S ==  38,)
    @constraint(model, A + D + H ==  38,)
    @constraint(model, B + E + I + M ==  38,)
    @constraint(model, C + F + J + N + Q ==  38,)
    @constraint(model, G + K + O + R ==  38,)
    @constraint(model, L + P + S ==  38,)
    @constraint(model, C + G + L ==  38,)
    @constraint(model, B + F + K + P ==  38,)
    @constraint(model, A + E + J + O + S ==  38,)
    @constraint(model, D + I + N + R ==  38,)
    @constraint(model, H + M + Q ==  38,)

    @constraint(model, A < H)
    @constraint(model, A < C)
    @constraint(model, A < L)
    @constraint(model, A < Q)
    @constraint(model, A < S)
    @constraint(model, C < H)

    # @objective(model,Max,z)

    # Solve the problem
    optimize!(model)

    status = JuMP.termination_status(model)
    # println("status:$status")
    num_sols = 0
    if status == MOI.OPTIMAL
        num_sols = MOI.get(model, MOI.ResultCount())
        println("num_sols:$num_sols\n")
        if print_solutions
            for sol in 1:num_sols
                println("solution #$sol")
                x_val = convert.(Integer,JuMP.value.(x; result=sol))
                println("x:$x_val")

            end
        end
    else
        println("status:$status")
    end

    return status, num_sols
end

@time magic_hexagon()


using Combinatorics

"""
    all_backdoor_adjustment_sets(dag, x, y)

Obtain all sets of variables to be controlled to close backdoor paths from variable x to y.
"""
all_backdoor_adjustment_sets(dag::DAG, x::Symbol, y::Symbol)::Vector{Vector{Symbol}} =
    nodes_labels.((dag, ), all_backdoor_adjustment_sets(dag, node(dag, x), node(dag, y)))

function all_backdoor_adjustment_sets(dag::DAG, x::Int, y::Int)::Vector{Vector{Int}}
    is_valid_backdoor_adjustment_set(dag, x, y, Int[]) && return [Int[]]

    possible_vars = Int[]
    for (idx, parent) ∈ enumerate(bfs_parents(dag.graph, x))
        parent > 0 && continue
        idx == y && continue
        push!(possible_vars, idx)
    end

    valid_sets::Vector{Vector{Int}} = []
    for s ∈ powerset(possible_vars, 1)
        super_complete = Bool[]
        for vs ∈ valid_sets
            push!(super_complete, intersect(vs, s) == vs)
        end
        any(super_complete) && continue
        if is_valid_backdoor_adjustment_set(dag, x, y, s)
            push!(valid_sets, s)
        end
    end
    valid_sets
end

"""
    is_valid_backdoor_adjustment_set

Checks that given set z is a valid "backdoor adjustment set", i.e.
z being observed disconnects x from y.
"""
is_valid_backdoor_adjustment_set(dag::DAG, x::Int, y::Int, z::Vector{Int})::Bool =
    is_d_separated(dag, [x], [y], union(z, [x]), ignore_edges=[x => y])

export
    all_backdoor_adjustment_sets,
    is_valid_backdoor_adjustment_set


mutable struct CustomVariance
    meanratio::Float64
    sn::Float64 # sum nominator
    sd::Float64 # sum denominator
    n::Int # count
    m2::Float64 # var*(n-1)
    var::Float64 # weigthedVariance
    function CustomVariance(num, denom)
        n = length(num)
        @assert n == length(denom) # length(w)
        sn = sum(num)
        sd = sum(denom)
        meanratio = sn / sd
        m2 = zero(eltype(num))
        for i = 1:n
            @inbounds m2 += denom[i] * (num[i] / denom[i] - meanratio)^2
        end
        if n > 1
            var = m2 / (n - 1)
        else
            var = zero(eltype(num))
        end
        return new(meanratio, sn, sd, n, m2, var)
    end
    function CustomVariance()
        return new(0.0, 0.0, 0.0, 0, 0.0, 0.0)
    end
end

function removeSinglePoint!(a::CustomVariance, num, denom)
    sn = a.sn - num
    sd = a.sd - denom
    meanratio = sn / sd
    # merge: m2 = a.m2 + b.m2 + 2*(a.sn*(a.meanratio-meanratio)+§*(num/denom-meanratio)) + a.sd*(meanratio^2-a.meanratio^2) + denom*(meanratio^2-num/denom^2)
    # unmerge:
    a.m2 = a.m2 + 2 * (num * (a.meanratio - num / denom) + sn * (a.meanratio - meanratio)) + denom * ((num / denom)^2 - a.meanratio^2) + sd * (meanratio^2 - a.meanratio^2)
    a.sn = sn
    a.sd = sd
    a.n -= 1
    a.meanratio = meanratio
    a.var = a.m2 / (a.n - 1)
    return nothing
end

function addSinglePoint!(a::CustomVariance, num, denom)
    sn = a.sn + num
    sd = a.sd + denom
    a.n = a.n + 1
    meanratio = sn / sd
    a.m2 = a.m2 + 2 * (a.sn * (a.meanratio - meanratio) + num * (num / denom - meanratio)) + a.sd * (meanratio^2 - a.meanratio^2) + denom * (meanratio^2 - (num / denom)^2)
    a.sn = sn
    a.sd = sd
    a.meanratio = meanratio
    a.var = a.m2 / (a.n - 1)
    return nothing
end

function Base.merge!(a::CustomVariance, b::CustomVariance)
    # merges b into a
    # i.e. the points form b are 'added' to a
    sn = a.sn + b.sn
    sd = a.sd + b.sd
    n = a.n + b.n
    meanratio = sn / sd
    m2 = a.m2 + b.m2 + 2 * (a.sn * (a.meanratio - meanratio) + b.sn * (b.meanratio - meanratio)) + a.sd * (meanratio^2 - a.meanratio^2) + b.sd * (meanratio^2 - b.meanratio^2)
    var = m2 / (n - 1)
    a.sn = sn
    a.sd = sd
    a.n = n
    a.meanratio = meanratio
    a.m2 = m2
    a.var = var
    return nothing
end

function Base.merge(a::CustomVariance, b::CustomVariance)
    c = deepcopy(a)
    merge!(c, b)
    return c
end

"""
function unmerge!(a::CustomVariance,b::CustomVariance)
this function updates the mean and variance of a
The assumption is that the underlying data of b is a subset of a.
If this is not the case, the result of this function is most likely meaningless
That is to say "subtract" the series b from the series a and update the mean and variance of a accordingly
"""
function unmerge!(a::CustomVariance, b::CustomVariance)
    sn = a.sn - b.sn
    sd = a.sd - b.sd
    n = a.n - b.n
    meanratio = sn / sd
    # merge: m2 = a.m2 + b.m2 + 2*(a.sn*(a.meanratio-meanratio)+b.sn*(b.meanratio-meanratio)) + a.sd*(meanratio^2-a.meanratio^2) + b.sd*(meanratio^2-b.meanratio^2)
    # unmerge:
    m2 = a.m2 - b.m2 + 2 * (b.sn * (a.meanratio - b.meanratio) + sn * (a.meanratio - meanratio)) + b.sd * (b.meanratio^2 - a.meanratio^2) + sd * (meanratio^2 - a.meanratio^2)
    var = m2 / (n - 1)
    a.sn = sn
    a.sd = sd
    a.n = n
    a.meanratio = meanratio
    a.m2 = m2
    a.var = var
    return nothing
end


using SequentialLOB
using StylizedFacts
using Plots
using DelimitedFiles
using AdaptiveABC
using Distributions

# Read in Observed Market Data

observed_price_path = round.(readdlm("$(dirname(@__FILE__))/Original_Price_Bars_2300.csv", ',', Float64, '\n')[:,1], digits=2)
observed_log_returns = diff(log.(observed_price_path[:,1]))
observed_summary_stats = get_summary_stats(observed_log_returns)

observed_price_path

plt = plot(
    1:size(observed_log_returns, 1) + 1, observed_price_path, legend=false,
    xlab="Time", ylab="Mid-Price",
)

display(plt)
savefig(plt, "$(dirname(@__FILE__))/market_data_price_path.pdf")

market_data_stylized_facts = StylizedFactsPlot(observed_price_path)

# Plot Log Returns Hist

dist = fit(Normal, observed_log_returns)
plt = histogram(observed_log_returns, normalize=true, legend=false)
plot!(plt, x -> pdf(dist, x), xlim=xlims(), legend=false)
savefig(plt, "$(dirname(@__FILE__))/market_data_hist_log_returns.pdf")

# Plot Log Returns

plt = StylizedFacts.plot_log_returns(market_data_stylized_facts, "")
savefig(plt, "$(dirname(@__FILE__))/market_data_returns.pdf")

# Plot QQ Log Returns

plt = StylizedFacts.plot_qq_log_returns(market_data_stylized_facts, "")
savefig(plt, "$(dirname(@__FILE__))/market_data_qq_log_returns.pdf")

# Plot ACF Order Flow

plt = StylizedFacts.plot_acf_order_flow(market_data_stylized_facts, "")
savefig(plt, "$(dirname(@__FILE__))/market_data_acf_order_flow.pdf")

# Plot ACF Log Returns

plt = StylizedFacts.plot_acf_log_returns(market_data_stylized_facts, "")
savefig(plt, "$(dirname(@__FILE__))/market_data_acf_log_returns.pdf")

# Plot ACF Abs Log Returns

plt = StylizedFacts.plot_acf_abs_log_returns(market_data_stylized_facts, "")
savefig(plt, "$(dirname(@__FILE__))/market_data_acf_abs_log_returns.pdf")

plt = StylizedFacts.plot_all_stylized_facts(market_data_stylized_facts)
display(plt)
savefig(plt, "$(dirname(@__FILE__))/market_data_stylised_facts.pdf")

# Plot Stylised Facts from Calibrated Models

# Constants
num_paths = 1; M = 400 ; T = 2299 ; p₀ = 238.75 ; L = 200 ; λ = 1.0

# Calibrated Parameters LOB Model

# ABC-PMC BBWM: D, σ, nu, μ
seed = 731986389128937
D, σ, nu, μ = [1.79, 0.401, 0.533, 1.518]
α = 0.0
model = SLOB(num_paths, T, p₀, M, L, D, σ, nu, α, SourceTerm(λ, μ))
calibrated_price_path = model(seed)[:,1]

calibrated_stylized_facts = StylizedFactsPlot(calibrated_price_path)

# Plot Price num_paths
plt = plot(1:size(calibrated_stylized_facts.price_path, 1), calibrated_stylized_facts.price_path, legend=false,
    xlab="Time", ylab="Mid-Price")
savefig(plt, "$(dirname(@__FILE__))/lob_abc_pmc_wb_price_path.pdf")

# Plot Log Returns Hist

dist = fit(Normal, calibrated_stylized_facts.log_returns)
plt = histogram(calibrated_stylized_facts.log_returns, normalize=true, legend=false)
plot!(plt, x -> pdf(dist, x), xlim=xlims(), legend=false)
savefig(plt, "$(dirname(@__FILE__))/lob_abc_pmc_wb_hist_log_returns.pdf")

# Plot Log Returns
plt = StylizedFacts.plot_log_returns(calibrated_stylized_facts, "")
display(plt)
savefig(plt, "$(dirname(@__FILE__))/lob_abc_pmc_wb_log_returns.pdf")

# Plot QQ Log Returns

plt = StylizedFacts.plot_qq_log_returns(calibrated_stylized_facts, "")
savefig(plt, "$(dirname(@__FILE__))/lob_abc_pmc_wb_qq_log_returns.pdf")

# Plot ACF Order Flow

plt = StylizedFacts.plot_acf_order_flow(calibrated_stylized_facts, "")
savefig(plt, "$(dirname(@__FILE__))/lob_abc_pmc_wb_acf_order_flow.pdf")

# Plot ACF Log Returns

plt = StylizedFacts.plot_acf_log_returns(calibrated_stylized_facts, "")
savefig(plt, "$(dirname(@__FILE__))/lob_abc_pmc_wb_acf_log_returns.pdf")

# Plot ACF Abs Log Returns

plt = StylizedFacts.plot_acf_abs_log_returns(calibrated_stylized_facts, "")
savefig(plt, "$(dirname(@__FILE__))/lob_abc_pmc_wb_acf_abs_log_returns.pdf")

plt = StylizedFacts.plot_all_stylized_facts(calibrated_stylized_facts)
display(plt)
savefig(plt, "$(dirname(@__FILE__))/lob_abc_pmc_wb_stylised_facts.pdf")

# ABC-PMC MADWE: D, σ, nu, μ
# seed = 13718273
# D, σ, nu, μ = [1.784, 1.455, 0.496, 1.547]
# α = 0.0
# model = SLOB(num_paths, T, p₀, M, L, D, σ, nu, α, SourceTerm(λ, μ))
# calibrated_price_path = model(seed)[:,1]

# calibrated_stylized_facts = StylizedFactsPlot(calibrated_price_path)
# plt = StylizedFacts.plot_all_stylized_facts(calibrated_stylized_facts)
# display(plt)
# savefig(plt, "$(dirname(@__FILE__))/lob_abc_pmc_we_stylised_facts.pdf")

# Calibrated Parameters SLOB Model

# ABC-PMC BBWM: D, σ, nu, μ, α_slob
seed = 189313131
D, σ, nu, μ, α = [2.087, 0.421, 0.603, 1.202, 106.504]
model = SLOB(num_paths, T, p₀, M, L, D, σ, nu, α, SourceTerm(λ, μ))
calibrated_price_path = model(seed)[:,1]

calibrated_stylized_facts = StylizedFactsPlot(calibrated_price_path)

# Plot Price num_paths
plt = plot(1:size(calibrated_stylized_facts.price_path, 1), calibrated_stylized_facts.price_path, legend=false,
    xlab="Time", ylab="Mid-Price")
savefig(plt, "$(dirname(@__FILE__))/slob_abc_pmc_wb_price_path.pdf")

# Plot Log Returns Hist

dist = fit(Normal, calibrated_stylized_facts.log_returns)
plt = histogram(calibrated_stylized_facts.log_returns, normalize=true, legend=false)
plot!(plt, x -> pdf(dist, x), xlim=xlims(), legend=false)
savefig(plt, "$(dirname(@__FILE__))/slob_abc_pmc_wb_hist_log_returns.pdf")

# Plot Log Returns
plt = StylizedFacts.plot_log_returns(calibrated_stylized_facts, "")
display(plt)
savefig(plt, "$(dirname(@__FILE__))/slob_abc_pmc_wb_log_returns.pdf")

# Plot QQ Log Returns

plt = StylizedFacts.plot_qq_log_returns(calibrated_stylized_facts, "")
savefig(plt, "$(dirname(@__FILE__))/slob_abc_pmc_wb_qq_log_returns.pdf")

# Plot ACF Order Flow

plt = StylizedFacts.plot_acf_order_flow(calibrated_stylized_facts, "")
savefig(plt, "$(dirname(@__FILE__))/slob_abc_pmc_wb_acf_order_flow.pdf")

# Plot ACF Log Returns

plt = StylizedFacts.plot_acf_log_returns(calibrated_stylized_facts, "")
savefig(plt, "$(dirname(@__FILE__))/slob_abc_pmc_wb_acf_log_returns.pdf")

# Plot ACF Abs Log Returns

plt = StylizedFacts.plot_acf_abs_log_returns(calibrated_stylized_facts, "")
savefig(plt, "$(dirname(@__FILE__))/slob_abc_pmc_wb_acf_abs_log_returns.pdf")

plt = StylizedFacts.plot_all_stylized_facts(calibrated_stylized_facts)
display(plt)
savefig(plt, "$(dirname(@__FILE__))/slob_abc_pmc_wb_stylised_facts.pdf")

# # ABC-PMC MADWE: D, σ, nu, μ, α
# seed = 63612536215
# D, σ, nu, μ, α = [2.087, 0.421, 0.603, 1.202, 106.504]
# model = SLOB(num_paths, T, p₀, M, L, D, σ, nu, α, SourceTerm(λ, μ))
# calibrated_price_path = model(seed)[:,1]

# calibrated_stylized_facts = StylizedFactsPlot(calibrated_price_path)
# plt = StylizedFacts.plot_all_stylized_facts(calibrated_stylized_facts)
# display(plt)
# savefig(plt, "$(dirname(@__FILE__))/slob_abc_pmc_we_stylised_facts.pdf")


using OrdinaryDiffEq, DifferentialEquations

tspan = (0.0f0, 200.0f0)

"""
FitzHugh Nagumo oscillator model from:
    https://en.wikipedia.org/wiki/FitzHugh%E2%80%93Nagumo_model
"""
function fhn_system(du, u, p, t;
                     U_func_time=U_func_time_trivial,
                     U_func_space=U_func_space_trivial)
    Ft = U_func_time(t)
    Fs = U_func_space(u)
    a, b, I_external, τ = p
    v, ω = u

    du[1] = v - (v^3)/3 - ω + I_external
    du[2] = (v + a - b*ω) / τ

    # Add controllers in
    du .+= Ft .+ Fs
end

# Generate data, with a time component
p = [0.7, 0.8, 0.5, 12.5]
u0 = [1.0, 0.0]
ts = range(tspan[1], tspan[2], length=5001)

function solve_fhn_system(;U_func_time=U_func_time_trivial,
                              U_func_space=U_func_space_trivial)
    prob = ODEProblem((du, u, p, t)->fhn_system(du, u, p, t,
                                U_func_time=U_func_time,
                                U_func_space=U_func_space),
                    u0, tspan, p)
    sol = solve(prob, Tsit5(), saveat=ts);
    return sol
end

#####
##### Turing function for Bayesian parameter estimation
#####
# Unforced version for use with Turing.jl
# function vdp_system(u, p, t)
#     u = convert.(eltype(p),u)
#     beta, gamma = p
#     S, I, R = u
#
#     N = sum([S,I,R])
#     tmp = beta*I*S/N
#     du = [-tmp; tmp - gamma*I; gamma*I]
#     return du
# end

#####
##### True model in SINDy syntax
#####
a, b, I, τ = p
#      x    y   c   xx xy yy xxx  xxy xyy yyy
 A = [[1   -1   I   0  0  0  -1/3 0   0   0];
      [1/τ -b/τ a/τ 0  0  0  0    0   0   0]]
n = size(A, 1)
sindy_library = Dict(
    "cross_terms"=>[2,3],
    "constant"=>nothing
);
core_dyn_true = sindyModel(ts, A,
                            convert_string2function(sindy_library),
                            ["v", "ω"])

# Actually export
export ts, solve_fhn_system, fhn_system, core_dyn_true


module SpectralEnvelope

using FFTW
using Polynomials
using LinearAlgebra
using Statistics

function detrend(x,y; deg = 1)
    fit = polyfit(x,y,deg)
    detrended = y - fit(x)
    return detrended
end

"""
smoothens the given series by mixen each points with it's neighboors.

    input:
    - data: the series you want to smooth
    - m : the numbers of points to be involved in the mixing process.

    output:
    - smoothened series, shorter of m points than the original.

"""

function smooth(data::Array{Float64,1}; m = 5)
    if m != 0
        smoothed_data = Float64[]
        weigths = Int[]
        for w in collect(-m:m)
            p = abs(2*w)+1
            if p >= 10
                append!(weigths,p)
            else
                append!(weigths,10)
            end
        end
        p_tot = sum([1/w for w in weigths])
        for i in 1:length(data)
            if i > m && i < length(data) - m
                smoothed_Iw = 0
                for (l,k) in enumerate(collect(-m:m))
                    smoothed_Iw += (1/weigths[l])*data[i+k]
                end
                append!(smoothed_data, smoothed_Iw/p_tot)
            end
        end
        return smoothed_data
    elseif m == 0
        return data
    end
end

"""
One-hot encodes the time-series. For k categories,
each category gets a k-1 unit vector, except the last one which
is mapped to zeros(k - 1).

    Example:

    vectorize([1,2,3,1,2]) returns :
    [[1,0],[0,1],[0,0],[1,0]]

"""

function vectorize(data)
    categories = unique(data)
    deleteat!(categories, length(categories))
    sorted_data = zeros(length(categories),length(data))
    for (t,d) in enumerate(data)
        for (i,j) in enumerate(categories)
            if d == j
                sorted_data[i,t] = 1
            end
        end
    end
    return sorted_data
end

"""
Splits the time-series into overlapping equal-lengthed chunks of given size.

    input :
    - x : the time-series
    - window : the size of the equal-lengthed chunks
    - step : size of the step by which the window is sliding along the time-series
    the smaller the value of size, the bigger the overlap between the different chunks.

    output :
    -  overlapping equal-lengthed windowed time-series
"""

function partitioning(x,window::Int,step::Int)
    return [x[i:i+window] for i in 1:step:length(x) if  i + window <= length(x)]
end

function partitioning(x::Array{Float64,2},window::Int,step::Int)
    return [x[:,i:i+window-1] for i in 1:step:length(x[1,:]) if  i + window <= length(x[1,:])]
end

"""
returns the power spectrum of a given time serie,

    input :
    - time serie
    - window for averaging (default = length(time serie)/10)
    - step for the window's sliding

    output :
    - value of the power spectrum

You will have to provide the frequencies yourself
"""

function power_spectrum(x::Array{Float64,1}, window::Int, step::Int)
    ts = partitioning(x,window,step)
    ps = [real(fft(i .- mean(i)).*conj(fft(i .- mean(i)))) for i in ts]
    pxx = mean(ps)
    return [pxx[i] for i in 1:div(window,2)]
end

"""
    periodogram_matrix(ts::Array{Float64,2}, demean = false; m=2)

Computes the k x k periodogram matrix (k number of categories).
If 'average' true, segments data and computes the average periodogram matrix over all segments.
m is the smoothing parameter.
"""
function periodogram_matrix(ts::Array{Float64,2}, average = false; m=2)
    if average == false
        periodo = zeros(length(ts[:,1]), length(ts[:,1]), length(ts[1,:]))
        DFT = zeros(Complex{Float64}, length(ts[:,1]), length(ts[1,:]))
        for i in 1:length(ts[:,1])
            DFT[i,:] = fft(ts[i,:] .- mean(ts[i,:]))
        end
        for i in 1:length(ts[:,1])
            for j in 1:length(ts[:,1])
                tmp = smooth(real((1/length(ts[1,:]))*DFT[i,:].*conj(DFT[j,:])); m = m)
                for w in 1:length(ts[1,:])-(2*m+1)
                    periodo[i,j,w] = tmp[w]
                end
            end
        end
        return periodo
    elseif average == true
        window = Int(floor(length(ts[1,:])/2))
        step = Int(floor(length(ts[1,:])/4))
        dims = length(ts[:,1])
        periodo = zeros(length(ts[:,1]), length(ts[:,1]), window)
        seg_ts = partitioning(ts, window, step)
        DFT = []
        for el in seg_ts
            tmp = zeros(Complex, dims, window)
            for i in 1:length(el[:,1])
                tmp[i,:] = fft(el[i,:] .- mean(el[i,:]) )
            end
            push!(DFT, tmp)
        end
        for el in DFT
            for i in 1:dims
                for j in 1:dims
                    tmp = smooth(real((1/window)*el[i,:].*conj(el[j,:])); m = m)
                    for w in 1:window-(2*m+1)
                        periodo[i,j,w] += tmp[w]
                    end
                end
            end
        end
        return periodo./length(seg_ts)
    end
end

"""
Computes the covariance-variance matrix of a given time-series.
"""
function varcov(ts::Array{Float64,2})
    cov_matrix = zeros(length(ts[:,1]), length(ts[:,1]))
    for i in 1:length(ts[:,1])
        for j in 1:length(ts[:,1])
            cov_matrix[i,j] = sum((1/(length(ts[1,:])-1))*(ts[i,:].-mean(ts[i,:])).*(ts[j,:].-mean(ts[j,:])))
        end
    end
    return cov_matrix
end

"""
Computes the spectral envelope of the given time-series.

    input:

    - ts: the time series to analyse
    - m : the degree of smoothing wished.
          m corresponds to the number of neighbooring points that are mixed
          with given point to realize the smoothing.

    output:

    - Frequencies : list of points corresponding to the involved frequencies.
                    contained in [0,0.5]
    - amplitude : values of the spectral envelope for each given frequency point.
    - eigenvectors : the optimal scaling for the different categories, for each frequency point.
    - categories : the list of categories present in the data.

"""
function spectral_envelope(ts; m = 3)
    result = Float64[]
    vec = vectorize(ts)
    eigvec = zeros(length(vec[:,1]), trunc(Int,length(ts)/2))
    S = inv(sqrt(varcov(vec)))
    # alternative definition of S
    # S = inv(varcov(vec))
    p = periodogram_matrix(vec; m = m)
    f_len = trunc(Int,length(p[1,1,:])/2)
    for i in 1:f_len
        ev = findmax(real.(eigvals(S*p[:,:,i]*S/f_len)))
        # ev = findmax(real.(eigvals(S*p[:,:,i]/f_len)))
        append!(result, ev[1])
        b = S*eigvecs(S*p[:,:,i])[ev[2],:]
        b = b./sqrt(sum(b.^2))
        eigvec[:,i] = real.(b)
    end
    return collect(1:length(result))/length(p[1,1,:]), result[1:end], eigvec
end

"""
     get_mappings(data, freq; m = 3)

Returns the optimal mappings corresponding to the frequency 'freq'.
Prints the position and strengh of peak at 'f' for control purposes.
input :
    - data : input categorical time series
    - freq : the frequency where one wants the optimal mappings
    - m : smoothing parameters (see 'spectral_envelope')
returns :
    - mappings : the mapping at the peak associated with the desired frequency.

example:
    # we have some data, plotting the spectral envelope, we see a peak at 0.33
    m = get_mappings(data, 0.33)
"""
function get_mappings(data, freq; m = 3)
    categories = unique(data)
    f, se, eigvecs = spectral_envelope(data; m = m)
    delta = f[2] - f[1]
    window = Int(div(0.04*length(f),1))
    peak_pos = findmin([abs(freq - delta*i) for i in 1:length(f)])[2]
    true_pos = findmax(se[peak_pos-window:peak_pos+window])[2] + peak_pos - window - 1
    mappings = Dict(categories[i] => round(eigvecs[i,true_pos]; digits = 2) for i in 1:length(eigvecs[:,1]))
    if length(categories) != length(eigvecs[:,1])
        push!(mappings, unique(data)[end] => 0.0)
    end
    println("position of peak: ",round(f[true_pos],digits = 2)," strengh of peak: ",round(se[true_pos], digits = 2))
    return mappings
end

function findmax_in(xserie,yserie,xlim)
    range = []
    real_pos = []
    for (index,value) in enumerate(xserie)
        if value>= xlim[1] && value <= xlim[2]
            append!(range,yserie[index])
            append!(real_pos,index)
        end
    end
    max,pos = findmax(range)
    return max, xserie[real_pos[pos]], real_pos[pos]
end

"""
    apply_mapping(ts, mapping)

Transforms the input time-series 'ts' according to the mapping 'mapping'.
Typically, 'mapping' would be the output of the 'get_mappings' function, although you can provide your own.
'mapping' should be a Dict specifying each value and what it should be mappied to. Example:
{"A" : 0.54, "G" : 0.62, "T" : -0.57, "C" : 0.0}

Returns an array which correspond to the mapped series according to 'mapping'
"""
function apply_mapping(ts, mapping)
    mapped_series = deepcopy(ts)
    for (index, value) in enumerate(mapped_series)
        mapped_series[index] = mapping[value]
    end
    return mapped_series
end

export spectral_envelope, get_mappings, detrend, smooth, power_spectrum, apply_mapping

end


using SchumakerSpline
using HCubature
using Distributions
using Random
using Optim
using FixedPointAcceleration
delta = 0.2
beta = 0.95
periodic_income = 1.0
shock_var = 1.0
shock_process = LogNormal(0.0, shock_var)
BudgetStateSpace = vcat( collect(0:0.015:periodic_income), collect(1.05:0.05:(3*periodic_income)))
InitialGuess = sqrt.(BudgetStateSpace)

function ValueGivenShock(Budget::Float64, epsilon::Float64, NextValueFunction::Schumaker)
    opt = optimize(x ->  -1.0*(epsilon*(x^delta) + beta*evaluate(NextValueFunction, Budget - x + periodic_income)), 0.0, Budget)
    return -1.0 * opt.minimum
end

function ExpectedUtility(Budget::Float64, NextValueFunction::Schumaker)
    if Budget > 0.00001
        integ = hcubature(epsilon-> ValueGivenShock(Budget, epsilon[1], NextValueFunction)* pdf(shock_process, epsilon[1]), [quantile(shock_process,0.0001)], [quantile(shock_process, 0.9999)])
        return integ[1]
    else
        return beta * evaluate(NextValueFunction, periodic_income)
    end
end

function OneIterateBudgetValues(BudgetValues::Array{Float64,1})
    NextValueFunction = Schumaker(BudgetStateSpace, BudgetValues)
    new_budget_values = zeros(length(BudgetStateSpace))
    for i in 1:length(BudgetStateSpace)
        new_budget_values[i] = ExpectedUtility(BudgetStateSpace[i], NextValueFunction)
    end
    return new_budget_values
end

function new_point(level::Float64, stepp::Float64, max_grad::Float64, grad_frac::Float64)
    expp = exp(grad_frac)
    frac = expp / (expp+1)
    if isinf(max_grad)
        return level + expp
    else
        return level + stepp * max_grad * frac
    end
end

function invert_new_point(level::Float64, previous_level::Float64, previous_stepp::Float64, previous_max_grad::Float64)
    rise = level - previous_level
    if isinf(previous_max_grad)
        return log(rise)
    else
        R = rise/(previous_stepp * previous_max_grad)
        return log(R/(1-R))
    end
end

function ShapeToBud(ShapeVec)
    lenlen = length(ShapeVec)
    BudgetValues = Array{Float64,1}(undef, lenlen)
    if lenlen > 0 BudgetValues[1] = ShapeVec[1] end
    if lenlen > 1 BudgetValues[2] = new_point(BudgetValues[1], 1.0,Inf,ShapeVec[2]) end
    if lenlen > 2
        for i in 3:lenlen
            stepp = BudgetStateSpace[i] - BudgetStateSpace[i-1]
            previous_grad = (BudgetValues[i-1]-BudgetValues[i-2])/(BudgetStateSpace[i-1] - BudgetStateSpace[i-2])
            BudgetValues[i] = new_point(BudgetValues[i-1], stepp, previous_grad, ShapeVec[i])
        end
    end
    return BudgetValues
end

function BudToShape(BudgetValues)
    lenlen = length(BudgetValues)
    new_shape_vec = Array{Float64,1}(undef, lenlen)
    if lenlen > 0 new_shape_vec[1] = BudgetValues[1] end
    if lenlen > 1 new_shape_vec[2] = invert_new_point(BudgetValues[2],BudgetValues[1], 1.0,Inf) end
    if lenlen > 2
        for i in 3:lenlen
            stepp = BudgetStateSpace[i] - BudgetStateSpace[i-1]
            previous_grad = (BudgetValues[i-1]-BudgetValues[i-2])/(BudgetStateSpace[i-1] - BudgetStateSpace[i-2])
            new_shape_vec[i] = invert_new_point(BudgetValues[i],BudgetValues[i-1], stepp, previous_grad)
        end
    end
    return new_shape_vec
end

function Reparameterised_FP_Vector(ShapeVec)
    BudgetValues = ShapeToBud(ShapeVec)
    new_Budget_Values = OneIterateBudgetValues(BudgetValues)
    new_shape_vec = BudToShape(new_Budget_Values)
    return new_shape_vec
end
function shapeconvergence(Inputs, Outputs)
    return maximum(abs.(ShapeToBud(Inputs) .- ShapeToBud(Outputs)))
end

# Testing that the conversions work properly.
fp = fixed_point(OneIterateBudgetValues, InitialGuess; PrintReports = true, MaxIter = Integer(1))
shape_guess = BudToShape(InitialGuess)
fp_reparam = fixed_point(Reparameterised_FP_Vector, shape_guess; PrintReports = true, MaxIter = Integer(1), ConvergenceMetric = shapeconvergence)
iterated = ShapeToBud(fp_reparam.Outputs_[:,1])
sum(abs.(fp.Outputs_[:,1] .- iterated) .> 1e-10)  == 0

# fixed point acceleration with the reparameterised version.
fp = fixed_point(OneIterateBudgetValues, InitialGuess; PrintReports = true, ConvergenceMetricThreshold = 1e-06)
fpfp = fp.FixedPoint_
fp_ander = fixed_point(Reparameterised_FP_Vector, shape_guess; PrintReports = true, ReportingSigFig = Integer(10), ConvergenceMetricThreshold = 1e-06, ConvergenceMetric = shapeconvergence)
iterated = ShapeToBud(fp_ander.FixedPoint_)
sum(abs.(fpfp .- iterated) .> 1e-4)  == 0
fp_aitken = fixed_point(Reparameterised_FP_Vector, shape_guess; PrintReports = true, ReportingSigFig = Integer(10), ConvergenceMetricThreshold = 1e-06, Algorithm = Aitken, ConvergenceMetric = shapeconvergence)
iterated = ShapeToBud(fp_aitken.FixedPoint_)
sum(abs.(fpfp .- iterated) .> 1e-4)  == 0
fp_newton = fixed_point(Reparameterised_FP_Vector, shape_guess; PrintReports = true, ReportingSigFig = Integer(10), ConvergenceMetricThreshold = 1e-06, Algorithm = FixedPointAcceleration.Newton, ConvergenceMetric = shapeconvergence)
iterated = ShapeToBud(fp_newton.FixedPoint_)
sum(abs.(fpfp .- iterated) .> 1e-4)  == 0
fp_sea = fixed_point(Reparameterised_FP_Vector, shape_guess; PrintReports = true, ReportingSigFig = Integer(10), ConvergenceMetricThreshold = 1e-06, Algorithm = SEA,  ConvergenceMetric = shapeconvergence)
iterated = ShapeToBud(fp_sea.FixedPoint_)
sum(abs.(fpfp .- iterated) .> 1e-4)  == 0


# symmetric, skew-symmetric and symmetric checks
"""
    symmetric(::SecondOrderTensor)
    symmetric(::FourthOrderTensor)

Computes the (minor) symmetric part of a second or fourth order tensor.
Return a `SymmetricTensor`.

# Examples
```jldoctest
julia> A = rand(Tensor{2,2})
2×2 Tensor{2, 2, Float64, 4}:
 0.590845  0.566237
 0.766797  0.460085

julia> symmetric(A)
2×2 SymmetricTensor{2, 2, Float64, 3}:
 0.590845  0.666517
 0.666517  0.460085
```
"""
@inline symmetric(S1::SymmetricTensors) = S1

@inline function symmetric(S::Tensor{2, dim}) where {dim}
    SymmetricTensor{2, dim}(@inline function(i, j) @inbounds i == j ? S[i,j] : (S[i,j] + S[j,i]) / 2 end)
end

"""
    minorsymmetric(::FourthOrderTensor)

Compute the minor symmetric part of a fourth order tensor, return a `SymmetricTensor{4}`.
"""
@inline function minorsymmetric(S::Tensor{4, dim}) where {dim}
    SymmetricTensor{4, dim}(
        @inline function(i, j, k, l)
            @inbounds if i == j && k == l
                return S[i,j,k,l]
            else
                return (S[i,j,k,l] + S[j,i,k,l] + S[i,j,k,l] + S[i,j,l,k]) / 4
            end
        end
    )
end

@inline minorsymmetric(S::SymmetricTensors) = S

@inline symmetric(S::Tensor{4}) = minorsymmetric(S)

"""
    majorsymmetric(::FourthOrderTensor)

Compute the major symmetric part of a fourth order tensor.
"""
@inline function majorsymmetric(S::T) where {T <: FourthOrderTensor}
    get_base(T)(
        @inline function(i, j, k, l)
            @inbounds if i == j == k == l || i == k && j == l
                return S[i,j,k,l]
            else
                return (S[i,j,k,l] + S[k,l,i,j]) / 2
            end
        end
    )
end

# Unsafe versions that just truncates
@inline unsafe_symmetric(S::SymmetricTensors) = S
@inline function unsafe_symmetric(S::Tensor{2, dim}) where {dim}
    SymmetricTensor{2, dim}(@inline function(i, j) @inbounds S[i,j] end)
end
@inline function unsafe_symmetric(S::Tensor{4, dim}) where {dim}
    SymmetricTensor{4, dim}(@inline function(i, j, k, l) @inbounds return S[i,j,k,l] end)
end

"""
    skew(::SecondOrderTensor)

Computes the skew-symmetric (anti-symmetric) part of a second order tensor, returns a `Tensor{2}`.
"""
@inline skew(S1::Tensor{2}) = (S1 - S1') / 2
@inline skew(S1::SymmetricTensor{2,dim,T}) where {dim, T} = zero(Tensor{2,dim,T})

# Symmetry checks
@inline LinearAlgebra.issymmetric(t::Tensor{2, 1}) = true
@inline LinearAlgebra.issymmetric(t::Tensor{2, 2}) = @inbounds t[1,2] == t[2,1]

@inline function LinearAlgebra.issymmetric(t::Tensor{2, 3})
    return @inbounds t[1,2] == t[2,1] && t[1,3] == t[3,1] && t[2,3] == t[3,2]
end

function isminorsymmetric(t::Tensor{4, dim}) where {dim}
    @inbounds for l in 1:dim, k in l:dim, j in 1:dim, i in j:dim
        if t[i,j,k,l] != t[j,i,k,l] || t[i,j,k,l] != t[i,j,l,k]
            return false
        end
    end
    return true
end

isminorsymmetric(::SymmetricTensor{4}) = true

function ismajorsymmetric(t::FourthOrderTensor{dim}) where {dim}
    @inbounds for l in 1:dim, k in l:dim, j in 1:dim, i in j:dim
        if t[i,j,k,l] != t[k,l,i,j]
            return false
        end
    end
    return true
end

LinearAlgebra.issymmetric(t::Tensor{4}) = isminorsymmetric(t)

LinearAlgebra.issymmetric(::SymmetricTensors) = true


using LAR
using PyCall
@pyimport larlib as p

function cmerge(models)
	V = hcat([models[k][1] for k=1:length(models)]...)
	shifts = [0]
	append!(shifts, [size(models[h][1],2) for h in 1:length(models)])
	FV = hcat([models[k][2]+shifts[k] for k=1:length(models)]...)
	EV = hcat([models[k][3]+shifts[k] for k=1:length(models)]...)
	return V,FV,EV
end

function boxes2lar(boxes)
	V = Array{Float64,1}[]
	EV = Array{Int64,1}[]
	FV = Array{Int64,1}[]
	for k=1:size(boxes,2)
		xm,ym,xM,yM = boxes[:,k]
		push!(V,[xm,ym])
		push!(V,[xM,yM])
		push!(V,[xm,yM])
		push!(V,[xM,ym])
		push!(EV,[ 4(k-1)+1, 4(k-1)+3 ])
		push!(EV,[ 4(k-1)+1, 4(k-1)+4 ])
		push!(EV,[ 4(k-1)+2, 4(k-1)+3 ])
		push!(EV,[ 4(k-1)+2, 4(k-1)+4 ])
		push!(FV,[ 4(k-1)+1, 4(k-1)+2, 4(k-1)+3, 4(k-1)+4])
	end
	return hcat(V...), hcat(EV...), hcat(FV...)
end

function boxes3lar(boxes)
	V = Array{Float64,1}[]
	EV = Array{Int64,1}[]
	FV = Array{Int64,1}[]
	for k=1:size(boxes,2)
		xm,ym,zm,xM,yM,zM = boxes[:,k]
		push!(V,[xm,ym,zm])
		push!(V,[xM,yM,zm])
		push!(V,[xm,yM,zm])
		push!(V,[xM,ym,zm])
		push!(V,[xm,ym,zM])
		push!(V,[xM,yM,zM])
		push!(V,[xm,yM,zM])
		push!(V,[xM,ym,zM])

		push!(EV,[ 8(k-1)+1, 8(k-1)+3 ])
		push!(EV,[ 8(k-1)+1, 8(k-1)+4 ])
		push!(EV,[ 8(k-1)+2, 8(k-1)+3 ])
		push!(EV,[ 8(k-1)+2, 8(k-1)+4 ])
		push!(EV,[ 8(k-1)+4+1, 8(k-1)+4+3 ])
		push!(EV,[ 8(k-1)+4+1, 8(k-1)+4+4 ])
		push!(EV,[ 8(k-1)+4+2, 8(k-1)+4+3 ])
		push!(EV,[ 8(k-1)+4+2, 8(k-1)+4+4 ])
		push!(EV,[ 8(k-1)+1, 8(k-1)+4+1 ])
		push!(EV,[ 8(k-1)+2, 8(k-1)+4+2 ])
		push!(EV,[ 8(k-1)+3, 8(k-1)+4+3 ])
		push!(EV,[ 8(k-1)+4, 8(k-1)+4+4 ])

		push!(FV,[ 8(k-1)+1, 8(k-1)+2, 8(k-1)+3, 8(k-1)+4])
		push!(FV,[ 8(k-1)+4+1, 8(k-1)+4+2, 8(k-1)+4+3, 8(k-1)+4+4])
		push!(FV,[ 8(k-1)+1, 8(k-1)+3, 8(k-1)+4+1, 8(k-1)+4+3])
		push!(FV,[ 8(k-1)+1, 8(k-1)+4, 8(k-1)+4+1, 8(k-1)+4+4])
		push!(FV,[ 8(k-1)+2, 8(k-1)+3, 8(k-1)+4+2, 8(k-1)+4+3])
		push!(FV,[ 8(k-1)+2, 8(k-1)+4, 8(k-1)+4+2, 8(k-1)+4+4])
	end
	return hcat(V...), hcat(EV...), hcat(FV...)
end

params = PyObject(pyeval("list([1.,0.,0.,0.1,  0.,1.,0.,0.1,  0.,0.,1.,0.1, 0.,0.,0.,0.1, 100.])"))
glass = p.MATERIAL(params)

function submanifoldMapping(V::Array{Float64,2},FV::Array{Int64,2},pivotFace::Int64)
	FW = [FV[:,k] for k=1:size(FV,2)]
	return submanifoldMapping(V,FW,pivotFace)
end

#function submanifoldMapping(V::Array{Float64,2},FV::Array{Array{Int64,1},1},pivotFace::Int64)
#    tx,ty,tz = V[:,FV[pivotFace][1]]
#    T = eye(4)
#    T[1,4], T[2,4], T[3,4] = -tx,-ty,-tz
#    facet = [ V[:,v] - [tx,ty,tz] for v in FV[pivotFace] ]
#    normal = normalize(p.COVECTOR(facet)[2:end])
#    a = normal
#    b = Float64[0,0,1]
#    if norm(cross(a,b)) < 10^-5.
#    	axis = normalize(cross(a,b))
#    else
#    	axis = b
#    end
#    angle = atan2(norm(cross(a,b)), dot(a,b))
#    # general 3D rotation (Rodrigues' rotation formula)
#    M = eye(4)
#    Cos, Sin = cos(angle), sin(angle)
#    I, u = eye(3), axis
#    Ux = [0        -u[3]      u[2];
#          u[3]        0      -u[1];
#         -u[2]      u[1]        0]
#    UU = [u[1]*u[1]    u[1]*u[2]    u[1]*u[3] ;
#          u[2]*u[1]    u[2]*u[2]    u[2]*u[3] ;
#          u[3]*u[1]    u[3]*u[2]    u[3]*u[3] ]
#    M[1:3,1:3] = Cos*I + Sin*Ux + (1.0-Cos)*UU
#    transform = M *  T
#    return transform
#end

function submanifoldMapping(V::Array{Float64,2},FV::Array{Array{Int64,1},1},pivotFace::Int64)
    tx,ty,tz = V[:,FV[pivotFace][1]]
    T = eye(4)
    T[1,4], T[2,4], T[3,4] = -tx,-ty,-tz
    facet = [ V[:,v] - [tx,ty,tz] for v in FV[pivotFace] ]
	centroid = p.CCOMB(PyObject(facet))
	u1 = facet[1]-centroid
	u2 = facet[2]-centroid
	u3 = cross(u1,u2)
	M = eye(4)
	M[1:3,1:3] = inv([u1 u2 u3])
    transform = M *  T
    return transform
end


cd(@__DIR__); include("setups/grid7x3.jl");
pyplot(dpi = 200)

# DAG pseudo-metric
distDAG = eigDAG_Distance(𝚽, Q, N)

# MDS embedding into R²
D = distDAG
E = transform(fit(MDS, D, maxoutdim=2, distances=true))

# set up all heatmap plots' positions
dx = 0.01; dy = dx;
xej = zeros(Nx, N); yej=zeros(Ny, N);
a = 5.0; b = 9.0;
for k = 1:N
    xej[:,k] = LinRange(E[1,k] - Ny * a * dx, E[1, k] + Ny * a * dx, Nx)
    yej[:,k] = LinRange(E[2,k] - a * dy, E[2, k] + a * dy, Ny)
end

##
plot()
for k=1:N
    heatmap!(xej[:, k], yej[:, k], reshape(𝚽[:, k], (Nx, Ny))', c = :viridis,
                colorbar = false, ratio = 1, annotations = (xej[4, k],
                yej[3, k] + b*dy, text(latexstring("\\varphi_{",
                string(eig2dct[k, 1]), ",", string(eig2dct[k, 2]), "}"), 10)))
end
plt = plot!(xlim = [-1.4, 1.3], ylim = [-1.4, 1.3], grid = false, clims = (-0.4, 0.4))
# first level partition
p1x = [-0.2, 1.0, NaN]; p1y = [1.3, -1.0, NaN];
plot!(p1x, p1y, c = :red, legend = false, width = 3)
# second level partition
p2x = [-1.0, 0.2, NaN, 0.4, 1.2, NaN]; p2y = [-0.8, 0.45, NaN, 0.25, 0.2, NaN];
plot!(p2x, p2y, c=:orange, legend = false, width = 2)
plt = current()
savefig(plt, "../figs/Grid7x3_DAG_2levels_partition.png")


"""
    trackpath(H::AbstractHomotopy, startvalue, algorithm::PredictorCorrector, start, finish; kwargs...)

Tracks the path ``z(t)`` which is implicitly defined via ``H(x,t)`` from
`z(start)=H(startvalue, start)` to `z(finish)= H(result, finish)` with the given
PredictorCorrector `algorithm`.

We reformulate the path ``z(t)`` as a path z(γ(s)) where γ: [1, 0] -> line from start to finish

## Optional arguments
* `maxiterations=10000`
* `tolerance=1e-4`: The tolerance used for the correction step
* `refinement_tolerance=1e-8`: The result will be refined to the given accuracy
* `initial_steplength=0.01`: The initial_steplength (w.r.t s)
* `correction_step_maxiterations=3`
* `successfull_steps_until_steplength_increase=3`
* `steplength_increase_factor=2.0`
* `steplength_decrease_factor=0.5`
* `record_trace=false`
* `record_steps=false`
* `debug=false`: Print *(quite a lot)* different values to debug behaviour.
"""
function trackpath(
    H::AbstractHomotopy{T},
    startvalue::Vector{T},
    algorithm::AbstractPredictorCorrectorAlgorithm,
    start::S,
    finish::S;kwargs...
) where {S<:Number,T<:Number}
    J_H = differentiate(H)
    ∂H∂t = ∂t(H)
    trackpath(H, J_H, ∂H∂t, startvalue, algorithm, start, finish; kwargs...)
end

function trackpath(
    H::AbstractHomotopy{T},
    J_H,
    ∂H∂t,
    startvalue::Vector{T},
    algorithm::AbstractPredictorCorrectorAlgorithm,
    start::S,
    finish::S;
    maxiterations=10000,
    tolerance=1e-4,
    refinement_tolerance=1e-12,
    initial_steplength=0.05,
    successfull_steps_until_steplength_increase=3,
    steplength_increase_factor=2.0,
    steplength_decrease_factor=0.5,
    record_trace=false,
    record_steps=false,
    correction_step_maxiterations=3,
    debug=false
) where {S<:Number,T<:Number}
    # We want to track a path z(t) from z(start) to z(finish).
    # We reformulate it as a problem z(γ(s)) where γ: [1, 0] -> line from start to finish
    γlength = finish - start
    γ(s) = start + (1 - s) * γlength

    # some setup before we start the main loop
    x = copy(startvalue)
    u = similar(x)

    steps::Vector{S} = [];
    trace::Vector{Vector{T}} = [];

    s = 1.0
    steplength = initial_steplength
    sucessive_successes = 0
    k = 0

    while s > 0 && k < maxiterations
        if record_steps
            push!(steps, s)
        end
        if record_trace
            push!(trace, x)
        end

        Δs= min(steplength, s)

        u .= predict(algorithm, H, J_H, ∂H∂t, x, γ(s), Δs * γlength)

        if debug
            println("")
            println("t: $t Δt: $Δt, iteration: $k")
            println("x: $x, norm: $(norm(evaluate(H,x,t)))")
            println("predicted: $x, norm: $(norm(evaluate(H,x,t)))")
        end

        converged = correct!(u, algorithm, H, J_H, x, γ(s - Δs), tolerance, correction_step_maxiterations)

        if converged
            x .= u
            s -= Δs

            sucessive_successes += 1
            if sucessive_successes == successfull_steps_until_steplength_increase
                steplength *= steplength_increase_factor
                sucessive_successes = 0
            end
        else
            if debug
                println("correction failed!")
            end
            sucessive_successes = 0
            steplength *= steplength_decrease_factor
        end

        k += 1
    end

    if s ≈ 0 && norm(evaluate(H, x, γ(0.0))) < tolerance
        refinement_converged = correct!(u, algorithm, H, J_H, x, γ(0.0), refinement_tolerance, 10)
        if refinement_converged
            x .= u
        else

        end
    end

    if k >= maxiterations
        retcode = :MaxIterations
    elseif norm(evaluate(H, x, γ(0.0))) > refinement_tolerance
        retcode = :Diverged
    else
        retcode = :Success
    end

    PathResult(x, retcode, startvalue, k, γ(s), steps, trace)
end


# functions to create parameters

dx_i/dt = r_i*X_i + sum(F_i(X_j)) - sum(F_j(X_i)) - m_i*X_i

dx_i/dt = fun1(X_i) + fun2(X_i) - fun3(X_i) - fun4(X_i)

function lorenz_added!(dx, x, p, t)
    dx[1] = sigma*(x[2]-x[3])
    dx[2] = x[2]*(rho - x[3]) - x[2]
    dx[3] = x[1]*x[2] - beta*x[3]
end

# dx is just a vector
# populate the vector via a loop through the functions where
# you look at each species
for i in species
    # get abundance from somewhere
    dx[i] = fun1(X_i) + fun2(X_i) - fun3(X_i) - fun4(X_i)
end

function sigma_fun(x::Array, sigma::Float)

    return (sigma*(x[2] - x[1]))

end

x0 = [1.0;0.0;0.0]

# dx1 = r1*x1 + sum(g12*x2, g13*x3) - sum(g21*x1, g31*x1) - m1*x1

x_metrics = [[2, 2], [1, 3], [1, 2.5]]
x_vals = [1, 2, 3]
x0 = [100, 200, 75]

"""
Function defines the density-independent growth of xi.
"""
function growth(r, i)
    return r*x[i]
end

"""
Determines for any focal x (i) and any other x (j) if the focal x will gain
biomass from j based on some arbitrary metric (metric defined in the x_metrics
vector), and if xi does gain from xj, then gives a g value (rate at which the
gain happens)
"""
function get_g(i, j)
    if x_metrics[i][2] > x_metrics[j][2]
        g = 0.5
    else
        g = 0
    end
end

"""
Figures out how much biomass is actually taken from any xj and given to a given
xi.
"""
function add(i, j)
    return get_g(i, j)*x[j]
end

"""
For a specific xi, gets all other x's that it can consume biomass from, gets the
biomass from each, and adds it to xi's biomass.
"""
function add_all(i)
    other_x = filter(x-> x ≠ i, x_vals)
    vals = []
    for j in other_x
        append!(vals, add(i, j))
    end

    return sum(vals)
end

"""
For any given xi, finds the other xj's that can remove biomass from xi, gets how
much biomass each xj will take from xi, and removes all of those values from xi.
"""
function remove_all(i)
    other_x = filter(x -> x ≠ i, x_vals)
    vals = []
    for j in other_x
        append!(vals, add(j, i))
    end

    return sum(vals)
end

function death(i)
    m = 1
    m*x[i]
end
death(i)
function eqns!(dx, x, p, t)
    dx = [[],[],[]]
    for i in x_vals
        dx[i] = growth(all_x[i][1], i) + add_all(i) - remove_all(i) - death(i)
    end
end

using DifferentialEquations
x0 = x
tspan = (0.0, 100.0)
prob = ODEProblem(eqns!, x0, tspan)
sol = solve(prob)

# dx1/dt = r1*x1 + x1(a12*x2) - a21*x1 - m1*x1
# dx2/dt = r2*x2 + x2(a21*x1) - a12*x2 - m2*x2

#dx/dt = indep_growth + dep_growth - dep_death - indep_death

function indep_growth(i, x)
    #define r
    if i == 1
        r = 1.5
        k = 5000
    else
        r = 1.9
        k = 2000
    end
    # return product of r and abundance of focal x
    return (r*x[i]*(1-(x[i]/k)))
end

function dep_growth(i, x)
    # determine if a12 or a21 should be used
    if i == 1 # use a12
        a = 0.010975
        k = 5000
        return x[i]*(a*x[2])*(1-(x[i]/k))
    else
        a = 0.01
        k = 2000
        return x[i]*(a*x[1])*(1-(x[i]/k))
    end
end

function dep_death(i, x)
    if i == 1
        a = 0.01
        return x[i]*a*x[2]
    else
        a = 0.010975
        return x[i]*a*x[1]
    end
    return death
end

function indep_death(i, x)
    if i == 1
        m = 1.4
    else
        m = 1.3
    end

    return m*x[i]
end

dx = zeros(2)
function eqns!(dx, x, p, t)
    for i in 1:2
        dx[i] = indep_growth(i, x) + dep_growth(i, x) - dep_death(i, x) - indep_death(i, x)
    end
end

r1 = 1.5
r2 = 1.9
a12 = 0.010975
a21 = 0.01
m1 = 1.4
m2 = 1.3
k1 = 5000
k2 = 2000
function eqns!(dx, x, p, t)
    dx[1] = r1*x[1]*(1-(x[1]/k1)) + x[1]*(a12*x[2])*(1-(x[1]/k1)) -
            x[1]*(a21*x[2]) - m1*x[1]
    dx[2] = r2*x[2]*(1-(x[2]/k2)) + x[2]*(a21*x[1])*(1-(x[2]/k2)) -
            x[2]*a12*x[1] - m2*x[2]
end

x0 = [100.0; 1000.0]
tspan = (0.0, 10.0)
prob = ODEProblem(eqns!, x0, tspan)
sol2 = solve(prob)
plot(sol2)

sol1[:] == sol2[:]


module Sudoku

using DancingLinks

export create_sudoku, solve_sudoku!, check_solution

const L = 3
const N = L * L

# For each label, there are 9 positions to place. There are 9 labels in each sub-grid. And there
# are 9 sub-grids in total.
const MAX_N_ROWS = N * N * N

# From https:#en.wikipedia.org/wiki/Exact_cover#Sudoku:
#
# - Row-Column: Each intersection of a row and column must contain exactly one number.
# - Row-Number: Each row must contain each number exactly once.
# - Column-Number: Each column must contain each number exactly once.
# - Box-Number: Each box must contain each number exactly once.
const MAX_N_COLS = N * N * 4

# `+ 1` for the root. `* 4` because there are 4 data objects linked together in each row.
const MAX_N = 1 + MAX_N_COLS + MAX_N_ROWS * 4

const Coord = Tuple{Int, Int}

struct SudokuSolver
    d::DLinks
    names::Vector{Tuple{Coord, Int}}  # the metadata, for outputing result
end

function create_sudoku(s::Array{Int})::SudokuSolver
    @assert size(s) == (N, N) "Invalid input"

    # allocates dancing links
    d = let
        L = Vector{Int}(undef, MAX_N)
        R = Vector{Int}(undef, MAX_N)
        U = Vector{Int}(undef, MAX_N)
        D = Vector{Int}(undef, MAX_N)

        h = 0
        C = Vector{Int}(undef, MAX_N)
        S = zeros(Int, MAX_N_COLS)

        DLinks(L, R, U, D, h, C, S)
    end
    names = Vector{Tuple{Coord, Int}}(undef, MAX_N)

    init_headers!(d, MAX_N_COLS)
    i = d.h

    for r in 1:N
        for c in 1:N
            cons = compute_constrait_indices(r, c)

            n_list = if s[r, c] == 0  # empty
                1:N
            else
                n = s[r, c]
                n:n
            end

            for n in n_list
                n_ofs = [n, n, n, 1]
                for ii in 1:4
                    DancingLinks.insert!(d,
                        i + ii, i + mod1(ii - 1, 4), i + mod1(ii + 1, 4),
                        cons[ii] + n_ofs[ii])
                    names[i + ii] = ((r, c), n)
                end
                i += 4
            end
        end
    end

    SudokuSolver(d, names)
end

const N2 = N * N

@inline function compute_constrait_indices(r::Int, c::Int)::Tuple{Int, Int, Int, Int}
    offset = 0
    row_cons = offset + N * (r - 1)

    offset += N2
    col_cons = offset + N * (c - 1)

    offset += N2
    box_cons = offset + N * ((r - 1) ÷ L * L + (c - 1) ÷ L)

    offset += N2
    r_c_cons = offset + ((r - 1) * N + (c - 1))

    (row_cons, col_cons, box_cons, r_c_cons)
end

function solve_sudoku!(sudoku::SudokuSolver)::Union{Array{Int}, Nothing}
    sudoku_solution = Int[]
    solved = dlx!(sudoku.d, sudoku_solution)
    if solved
        sol = zeros(Int, N, N)
        for r in sudoku_solution
            (coord, n) = sudoku.names[r]
            sol[coord...] = n
        end
        sol
    else
        println(stderr, "No solution.")
        nothing
    end
end

const ALL_NUMBERS = Set(1:N)

function check_solution(s::Array{Int}, puzzle::Array{Int})::Bool
    @assert size(s) == (N, N) "Invalid input"

    # Rows
    if !all(map(s -> s == ALL_NUMBERS, mapslices(Set, s, dims=2)))
        return false
    end

    # Columns
    if !all(map(s -> s == ALL_NUMBERS, mapslices(Set, s, dims=1)))
        return false
    end

    # Boxes
    for r in 1:L
        for c in 1:L
            if Set(view(s, (1:L) .+ L * (r - 1), (1:L) .+ L * (c - 1))) != ALL_NUMBERS
                return false
            end
        end
    end

    # Conforms to the completed grids
    all(@. (s == puzzle) | (puzzle == 0))
end

end  # module


@testset "Slow Manifold System" begin
  function slow_manifold(du, u, p, t)
    du[1] = p[1]*u[1]
    du[2] = p[2]*(u[2]-u[1]^2) + p[3]*sin(t^2)
  end

  u0 = [3.0; -2.0]
  tspan = (0.0, 5.0)
  p = [-2.0; -0.5; 1.0]

  problem = ODEProblem(slow_manifold, u0, tspan, p)
  solution = solve(problem, Tsit5(), saveat = 0.01)

  ufun(u,p,t) = sin(t^2)

  prob = ContinuousDataDrivenProblem(solution, U = ufun)

  @variables u[1:2] y[1:1] t
  Ψ = Basis([u; u[1]^2; y], u, controls = y, iv = t)

  for alg in [DMDPINV(); DMDSVD(); TOTALDMD()]
    res = solve(prob, Ψ, alg, digits = 1)
    b = result(res)
    m = metrics(res)
    @test isapprox(eigvals(b), [2*p[1]; p[1]; p[2]], atol = 1e-1)
    @test all(m[:L₂] .< 3e-1)

    # TODO This does not work right now, but it should
    #sdict = Dict([y[1] => sin(t^2)])

    #for (i,eq) in enumerate(equations(b))
    #  lhs, rhs = eq.lhs, eq.rhs
    #  b[i] = Num(lhs) ~ substitute(Num(rhs), sdict)
    #end

    #_prob = ODEProblem((args...)->b(args...), u0, tspan, parameters(res))
    #_sol = solve(_prob, Tsit5(), saveat = solution.t)
    #@test norm(solution - _sol)/size(X, 2) < 5e-1
  end
end

@testset "Slow Manifold Discrete System" begin
  function slow_manifold(du, u, p, t)
    du[1] = p[1]*u[1]
    du[2] = p[2]*u[2]+(p[1]^2 - p[2])*u[1]^2 + p[3]*sin(t^2)
  end

  u0 = [3.0; -2.0]
  tspan = (0.0, 10.0)
  p = [-0.8; -0.7; 1.0]

  problem = DiscreteProblem(slow_manifold, u0, tspan, p)
  solution = solve(problem, FunctionMap())

  ufun(u,p,t) = sin(t^2)

  ddprob = DiscreteDataDrivenProblem(solution, U = ufun)

  @variables u[1:2] y[1:1] t

  Ψ = Basis([u[1]; u[1]^2; u[2]-u[1]^2; y], u, controls = y, iv = t)

  for alg in [DMDPINV(); DMDSVD(); TOTALDMD()]
    res = solve(ddprob, Ψ, alg, digits = 1)
    b = result(res)
    m = metrics(res)
    @test isapprox(eigvals(b), [p[1]; p[2]; p[1]^2], atol = 1e-1)
    @test all(m[:L₂] .< 3e-1)
  end
end


type Linear{A<:AbstractArray} <: Layer{A, A}
    output :: A
    gradInput :: A
    weight :: A
    bias :: A
    gradWeight :: A
    gradBias :: A
end

# input and output are 1d
Linear(T::Type, batch_size, input_size, output_size) =
    Linear(Array(T, (output_size, batch_size)),
           Array(T, (input_size, batch_size)),
           Array(T, (output_size, input_size)),
           Array(T, (output_size, 1)),
           Array(T, (output_size, input_size)),
           Array(T, (output_size, 1)))

function updateOutput!{A}(self :: Linear{A}, input :: A)
    broadcast!(+, self.output, self.weight * input, self.bias)
    return self.output
end

function updateGradInput!{A}(self :: Linear{A}, input :: A,
                                gradOutput :: A)
    self.gradInput = transpose(self.weight) * gradOutput
    return self.gradInput
end

function accGradParameters!{A}(self :: Linear{A}, input :: A,
                                  gradOutput :: A, scale :: eltype(A))
    self.gradWeight = self.gradWeight + scale * gradOutput * transpose(input)
    self.gradBias = self.gradBias + scale * sum(gradOutput, 2)
end

export Linear


using NeuralPDE, Flux, ModelingToolkit, GalacticOptim, Optim, DiffEqFlux
using Quadrature,Cubature,Cuba
using Plots

# Physical and numerical parameters (fixed)
nu = 0.07
nx = 10001 #101
x_max = 2.0 * pi
dx = x_max / (nx - 1.0)
nt = 2 #10
dt = dx * nu
t_max = dt * nt

# Analytic function
analytic_sol_func(t, x) = -2*nu*(-(-8*t + 2*x)*exp(-(-4*t + x)^2/(4*nu*(t + 1)))/
                          (4*nu*(t + 1)) - (-8*t + 2*x - 12.5663706143592)*
                          exp(-(-4*t + x - 6.28318530717959)^2/(4*nu*(t + 1)))/
                          (4*nu*(t + 1)))/(exp(-(-4*t + x - 6.28318530717959)^2/
                          (4*nu*(t + 1))) + exp(-(-4*t + x)^2/(4*nu*(t + 1)))) + 4

function burgers(strategy, minimizer)

    @parameters x t
    @variables u(..)
    Dt = Differential(t)
    Dx = Differential(x)
    Dxx = Differential(x)^2

    eq = Dt(u(x, t)) + u(x, t) * Dx(u(x, t))  ~ nu * Dxx(u(x, t))

    bcs = [u(x,0.0) ~ analytic_sol_func(x, 0.0),
           u(0.0,t) ~ u(x_max,t)]

    domains = [x ∈ IntervalDomain(0.0, x_max),
               t ∈ IntervalDomain(0.0, t_max)]

    chain = FastChain(FastDense(2,16,Flux.σ),FastDense(16,16,Flux.σ),FastDense(16,1))
    discretization = PhysicsInformedNN(chain, strategy)

    indvars = [x, t]   #physically independent variables
    depvars = [u]      #dependent (target) variable

    dim = length(domains)

    losses = []
    error = []
    times = []

    dx_err = 0.00005

    error_strategy = GridTraining(dx_err)

    initθ = Float64.(DiffEqFlux.initial_params(chain))
    eltypeθ = eltype(initθ)
    parameterless_type_θ = DiffEqBase.parameterless_type(initθ)
    phi = NeuralPDE.get_phi(chain,parameterless_type_θ)
    derivative = NeuralPDE.get_numeric_derivative()

    _pde_loss_function = NeuralPDE.build_loss_function(eq,indvars,depvars,
                                             phi,derivative,nothing,chain,initθ,error_strategy)

    bc_indvars = NeuralPDE.get_variables(bcs,indvars,depvars)
    _bc_loss_functions = [NeuralPDE.build_loss_function(bc,indvars,depvars,
                                              phi,derivative,nothing,chain,initθ,error_strategy,
                                              bc_indvars = bc_indvar) for (bc,bc_indvar) in zip(bcs,bc_indvars)]

    train_sets = NeuralPDE.generate_training_sets(domains,dx_err,[eq],bcs,eltypeθ,indvars,depvars)
    train_domain_set, train_bound_set = train_sets

    pde_loss_function = NeuralPDE.get_loss_function(_pde_loss_function,
                                          train_domain_set[1],eltypeθ,
                                          parameterless_type_θ,error_strategy)

    bc_loss_functions = [NeuralPDE.get_loss_function(loss,set,
                                                 eltypeθ, parameterless_type_θ,
                                                 error_strategy) for (loss, set) in zip(_bc_loss_functions,train_bound_set)]

    loss_functions = [pde_loss_function; bc_loss_functions]
    loss_function__ = θ -> sum(map(l->l(θ) ,loss_functions))

    function loss_function_(θ,p)
            return loss_function__(θ)
    end

    cb = function (p,l)

        timeCounter = 0.0
        deltaT_s = time_ns() #Start a clock when the callback begins, this will evaluate questo misurerà anche il calcolo degli uniform error

        ctime = time_ns() - startTime - timeCounter #This variable is the time to use for the time benchmark plot
        append!(times, ctime/10^9) #Conversion nanosec to seconds
        append!(losses, l)
        append!(error, loss_function__(p))
        #println(length(losses), " Current loss is: ", l, " uniform error is, ", loss_function__(p))

        timeCounter = timeCounter + time_ns() - deltaT_s #timeCounter sums all delays due to the callback functions of the previous iterations

        return false
    end

    @named pde_system = PDESystem(eq, bcs, domains, indvars, depvars)

    discretization = NeuralPDE.PhysicsInformedNN(chain, strategy, init_params =initθ)
    prob = NeuralPDE.discretize(pde_system,discretization)

    startTime = time_ns() #Fix initial time (t=0) before starting the training

    if minimizer == "both"
        res = GalacticOptim.solve(prob, ADAM(); cb = cb, maxiters=5)
        prob = remake(prob,u0=res.minimizer)
        res = GalacticOptim.solve(prob, BFGS(); cb = cb, maxiters=15)
    else
        res = GalacticOptim.solve(prob, minimizer; cb = cb, maxiters=500)
    end

    phi = discretization.phi

    params = res.minimizer

    return [error, params, times, losses]
end

# Settings:
#maxIters = [(0,0,0,0,0,0,20000),(300,300,300,300,300,300,300)] #iters

strategies = [NeuralPDE.QuadratureTraining()]

strategies_short_name = ["QuadratureTraining"]

minimizers = [ADAM(),
              ADAM(0.000005),
              ADAM(0.0005),
              RMSProp(),
              RMSProp(0.00005),
              RMSProp(0.05),
              BFGS(),
              LBFGS()]

minimizers_short_name = ["ADAM",
                         "ADAM(0.000005)",
                         "ADAM(0.0005)",
                         "RMS",
                         "RMS(0.00005)",
                         "RMS(0.05)",
                         "BFGS",
                         "LBFGS"]

# Run models
error_res =  Dict()
params_res = Dict()
times = Dict()
losses_res = Dict()

print("Starting run \n")

for min in 1:length(minimizers) # minimizer
      for strat in 1:length(strategies) # strategy
            #println(string(strategies_short_name[1], "  ", minimizers_short_name[min]))
            res = burgers(strategies[strat], minimizers[min])
            push!(error_res, string(strat,min)     => res[1])
            push!(params_res, string(strat,min) => res[2])
            push!(times, string(strat,min)        => res[3])
            push!(losses_res, string(strat,min)        => res[4])
      end
end

#PLOT ERROR VS ITER: to compare to compare between minimizers, keeping the same strategy (easily adjustable to compare between strategies)
error_iter = Plots.plot(1:length(error_res["11"]), error_res["11"], yaxis=:log10, title = string("Burger error vs iter"), ylabel = "Error", label = string(minimizers_short_name[1]), ylims = (0.0001,1))
plot!(error_iter, 1:length(error_res["12"]), error_res["12"], yaxis=:log10, label = string(minimizers_short_name[2]))
plot!(error_iter, 1:length(error_res["13"]), error_res["13"], yaxis=:log10, label = string(minimizers_short_name[3]))
plot!(error_iter, 1:length(error_res["14"]), error_res["14"], yaxis=:log10, label = string(minimizers_short_name[4]))
plot!(error_iter, 1:length(error_res["15"]), error_res["15"], yaxis=:log10, label = string(minimizers_short_name[5]))
plot!(error_iter, 1:length(error_res["16"]), error_res["16"], yaxis=:log10, label = string(minimizers_short_name[6]))
plot!(error_iter, 1:length(error_res["17"]), error_res["17"], yaxis=:log10, label = string(minimizers_short_name[7]))
plot!(error_iter, 1:length(error_res["18"]), error_res["18"], yaxis=:log10, label = string(minimizers_short_name[8]))

Plots.plot(error_iter)

#Use after having modified the analysis setting correctly --> Error vs iter: to compare different strategies, keeping the same minimizer
#error_iter = Plots.plot(1:length(error_res["11"]), error_res["11"], yaxis=:log10, title = string("Burger error vs iter"), ylabel = "Error", label = string(strategies_short_name[1]), ylims = (0.0001,1))
#plot!(error_iter, 1:length(error_res["21"]), error_res["21"], yaxis=:log10, label = string(strategies_short_name[2]))
#plot!(error_iter, 1:length(error_res["31"]), error_res["31"], yaxis=:log10, label = string(strategies_short_name[3]))
#plot!(error_iter, 1:length(error_res["41"]), error_res["41"], yaxis=:log10, label = string(strategies_short_name[4]))
#plot!(error_iter, 1:length(error_res["51"]), error_res["51"], yaxis=:log10, label = string(strategies_short_name[5]))
#plot!(error_iter, 1:length(error_res["61"]), error_res["61"], yaxis=:log10, label = string(strategies_short_name[6]))
#plot!(error_iter, 1:length(error_res["71"]), error_res["71"], yaxis=:log10, label = string(strategies_short_name[7]))

#PLOT ERROR VS TIME: to compare to compare between minimizers, keeping the same strategy
error_time = plot(times["11"], error_res["11"], yaxis=:log10, label = string(minimizers_short_name[1]),title = string("Burger error vs time"), ylabel = "Error", size = (1500,500))
plot!(error_time, times["12"], error_res["12"], yaxis=:log10, label = string(minimizers_short_name[2]))
plot!(error_time, times["13"], error_res["13"], yaxis=:log10, label = string(minimizers_short_name[3]))
plot!(error_time, times["14"], error_res["14"], yaxis=:log10, label = string(minimizers_short_name[4]))
plot!(error_time, times["15"], error_res["15"], yaxis=:log10, label = string(minimizers_short_name[5]))
plot!(error_time, times["16"], error_res["16"], yaxis=:log10, label = string(minimizers_short_name[6]))
plot!(error_time, times["17"], error_res["17"], yaxis=:log10, label = string(minimizers_short_name[7]))
plot!(error_time, times["18"], error_res["18"], yaxis=:log10, label = string(minimizers_short_name[7]))

Plots.plot(error_time)

#Use after having modified the analysis setting correctly --> Error vs time: to compare different strategies, keeping the same minimizer
#error_time = plot(times["11"], error_res["11"], yaxis=:log10, label = string(strategies_short_name[1]),title = string("Burger error vs time"), ylabel = "Error", size = (1500,500))
#plot!(error_time, times["21"], error_res["21"], yaxis=:log10, label = string(strategies_short_name[2]))
#plot!(error_time, times["31"], error_res["31"], yaxis=:log10, label = string(strategies_short_name[3]))
#plot!(error_time, times["41"], error_res["41"], yaxis=:log10, label = string(strategies_short_name[4]))
#plot!(error_time, times["51"], error_res["51"], yaxis=:log10, label = string(strategies_short_name[5]))
#plot!(error_time, times["61"], error_res["61"], yaxis=:log10, label = string(strategies_short_name[6]))
#plot!(error_time, times["71"], error_res["71"], yaxis=:log10, label = string(strategies_short_name[7]))

using SciMLBenchmarks
SciMLBenchmarks.bench_footer(WEAVE_ARGS[:folder],WEAVE_ARGS[:file])