#
# bradi.jl --
#
# Find a global minimum of an univariate function with the "Bracket then Dig"
# method described in:
#
#     Ferréol Soulez, Éric Thiébaut, Michel Tallon, Isabelle Tallon-Bosc
#     and Paulo Garcia, "Optimal a posteriori fringe tracking in optical
#     interferometry", Proc. SPIE 9146, Optical and Infrared Interferometry
#     IV, 91462Y (July 24, 2014); doi:10.1117/12.2056590
#
#-------------------------------------------------------------------------------
#
# This file is part of OptimPackNextGen.jl which is licensed under the MIT
# "Expat" License.
#
# Copyright (C) 2014-2020, Éric Thiébaut.
#

module BraDi

using ..Brent: fmin_atol, fmin_rtol, fminbrkt

"""
    BraDi.minimize([T=Float64,] f, x) -> (xbest, fbest)
    BraDi.maximize([T=Float64,] f, x) -> (xbest, fbest)

finds the global minimum (resp. maximum) of an univariate function `f`.  The
argument `x` is an abstract vector of coordinates in monotonic order; `x[1]`
and `x[end]` are the endpoints of the global search interval and it is assumed
that no more than a single local minimum lies in any subinterval
`[x[i],x[i+2]]`.  The result is a tuple of 2-values: `xbest` the position of
the global minimum and `fbest` the function value at this position.

Optional argument `T` is the floating-point type used for the computations.

If specified, keywords `atol` and `rtol` set the absolute and relative
tolerances for the precision (see `Brent.fmin`).

For instance:

    BraDi.minimize(f, range(a,stop=b,length=n))

performs a global search in the closed interval `[a,b]` which is sampled by `n`
equalspaced points.

These functions implement the BraDi ("Bracket" then "Dig") algorithm described
in:

> Ferréol Soulez, Éric Thiébaut, Michel Tallon, Isabelle Tallon-Bosc and Paulo
> Garcia, "Optimal a posteriori fringe tracking in optical interferometry",
> Proc. SPIE 9146, Optical and Infrared Interferometry IV, 91462Y (July 24,
> 2014); doi:10.1117/12.2056590

See also `Brent.fmin` and `Step.minimize`.

"""
minimize(f, x::AbstractVector{<:Real}; kdws...) =
    minimize(Float64, f, x; kdws...)

maximize(f, x::AbstractVector{<:Real}; kdws...) =
    maximize(Float64, f, x; kdws...)

function minimize(::Type{T}, f, x::AbstractVector{<:Real};
                  atol::Real = fmin_atol(T),
                  rtol::Real = fmin_rtol(T)) where {T<:AbstractFloat}
    _minimize(T, f, x, T(atol), T(rtol))
end

function maximize(::Type{T}, f, x::AbstractVector{<:Real};
                  kdws...) where {T<:AbstractFloat}
    (xbest, fbest) = minimize(T, t -> -f(t), x; kdws...)
    return (xbest, -fbest)
end

function _minimize(::Type{T}, f, x::AbstractVector{<:Real},
                   atol::T, rtol::T) where {T<:AbstractFloat}
    xbest = xa = xb = xc = T(x[1])
    fbest = fa = fb = fc = T(f(xc))
    n = length(x)
    for j in 2 : n + 1
        xa, xb = xb, xc
        fa, fb = fb, fc
        if j ≤ n
            xc = T(x[j])
            fc = T(f(xc))
            if fc < fbest
                xbest = xc
                fbest = fc
            end
        end
        if fa ≥ fb ≤ fc
            # A minimum has been bracketed in [XA,XC].
            xm, fm = fminbrkt(f, xb, fb, xa, fa, xc, fc, atol, rtol)
            if fm < fbest
                xbest = xm
                fbest = fm
            end
        end
    end
    return (xbest, fbest)
end

@doc @doc(minimize) maximize

end # module


struct SVRG_basic_iterable{R<:Real,C<:RealOrComplex{R},Tx<:AbstractArray{C},Tf,Tg}
    F::Array{Tf}            # smooth term
    g::Tg                   # nonsmooth term
    x0::Tx                  # initial point
    N::Int                  # of data points in the finite sum problem
    L::Maybe{Union{Array{R},R}}  # Lipschitz moduli of nabla f_i
    μ::Maybe{Union{Array{R},R}}  # convexity moduli of the gradients
    γ::Maybe{R}             # stepsize
    m::Maybe{Int}           # number of inner loop updates
    plus::Bool              # for SVRG++ variant
end

mutable struct SVRG_basic_state{R<:Real,Tx}
    γ::R                    # stepsize
    m::Int                  # number of inner loop updates
    av::Tx                  # the running average
    z::Tx
    z_full::Tx              # the outer loop argument
    w::Tx                   # the inner loop variable
    ind::Array{Int}         # running idx set
    # some extra placeholders
    ∇f_temp::Tx             # placeholder for gradients
    temp::Tx
end

function SVRG_basic_state(γ::R, m, av::Tx, z::Tx, z_full::Tx, w::Tx, ind) where {R,Tx}
    return SVRG_basic_state{R,Tx}(γ, m, av, z, z_full, w, ind, copy(av), copy(av))
end

function Base.iterate(iter::SVRG_basic_iterable{R}) where {R}
    N = iter.N
    ind = collect(1:N)
    m = iter.m === nothing ? m = 2 * N : m = iter.m
    # updating the stepsize
    if iter.γ === nothing
        if iter.plus
            @warn "provide a stepsize γ"
            return nothing
        else
            if iter.L === nothing || iter.μ === nothing
                @warn "smoothness or convexity parameter absent"
                return nothing
            else
                L_M = maximum(iter.L)
                μ_M = maximum(iter.μ)
                γ = 1 / (10 * L_M)
                # condition Theorem 3.1
                rho = (1 + 4 * L_M * γ^2 * μ_M * (N + 1)) / (μ_M * γ * N * (1 - 4L_M * γ))
                if rho >= 1
                    @warn "convergence condition violated...provide a stepsize!"
                end
            end
        end
    else
        γ = iter.γ # provided γ
    end
    # initializing the vectors
    av = zero(iter.x0)
    for i = 1:N
        ∇f, ~ = gradient(iter.F[i], iter.x0)
        ∇f ./= N
        av .+= ∇f
    end
    z_full = copy(iter.x0)
    z = zero(av)
    w = copy(iter.x0)
    state = SVRG_basic_state(γ, m, av, z, z_full, w, ind)
    print()
    return state, state
end

function Base.iterate(iter::SVRG_basic_iterable{R}, state::SVRG_basic_state{R}) where {R}
    # The inner cycle
    for i in rand(state.ind, state.m)
        gradient!(state.temp, iter.F[i], state.z_full)
        gradient!(state.∇f_temp, iter.F[i], state.w)
        state.temp .-= state.∇f_temp
        state.temp .-= state.av
        state.temp .*= state.γ
        state.temp .+= state.w
        CIAOAlgorithms.prox!(state.w, iter.g, state.temp, state.γ)
        state.z .+= state.w   # keeping track of the sum of w's
    end
    # full update
    state.z_full .= state.w #state.z ./ state.m
    # iter.plus || (state.w .= state.z_full) # only for basic SVRG
    state.z = zero(state.z)  # for next iterate
    state.av .= state.z
    for i = 1:iter.N
        gradient!(state.∇f_temp, iter.F[i], state.z_full)
        state.∇f_temp ./= iter.N
        state.av .+= state.∇f_temp
    end
    # iter.plus && (state.m *= 2) # only for SVRG++

    return state, state
end

solution(state::SVRG_basic_state) = state.z_full


module juliaAllotter

using CSV
using LinearAlgebra
using Statistics
using JLD
include("core.jl")

import Random
Random.seed!(2^13 - 1)

courseFile = "data/courseFile.csv"
studentsFile = "data/studentsFile.csv"
choiceIndices = [0, 1, 2, 3, 4, 5]
choiceWeights = [10, 1, 2, 3, 4, 5]
costWeights = [-1.0, 1.0, 1.0]
iterations = 4900

coursesDF = CSV.read(courseFile)
studentsDF = CSV.read(studentsFile)

courseCount = length(coursesDF[:, 1])
studentCount = length(studentsDF[:, 1])
cpiArray = Array(studentsDF.CPI)
times = Array(coursesDF.CourseNeeds)

choiceIdx = makeArray(studentCount, studentsDF, courseCount, coursesDF, choiceIndices)
choiceWeights = makeArray(studentCount, studentsDF, courseCount, coursesDF, choiceWeights)

initialAllotment = allotment(studentCount, courseCount, times, Matrix{Float64}(I, studentCount, studentCount))
initialAllotment = squeezeAllotment(initialAllotment)

choiceGoodnessOld, cpiGoodnessOld = calcGoodness(initialAllotment, choiceWeights, cpiArray)

finalAllottment, utility = runMCMC(initialAllotment, iterations, studentCount, courseCount, costWeights, choiceIdx, choiceWeights, studentsDF, cpiArray, times)
choiceGoodnessNew, cpiGoodnessNew = calcGoodness(finalAllottment, choiceWeights, cpiArray)

writePerformance(finalAllottment, choiceGoodnessOld, cpiGoodnessOld, choiceGoodnessNew, cpiGoodnessNew, utility)

end # module


# This file contains the core functions of WignerSymbols and CG-coefficient.

"""
    HalfInt = Union{Integer, Rational}
Angular momentum quantum numbers may be half integers and integers. With `HalfInt` type, you can use both `2` and `3//2` as parameters.
But the parameter like `5//3`, who's denominator is not `2` while gives out error.
"""
const HalfInt = Union{Integer,Rational}

"""
basic CG coefficient calculation function
"""
function _dCG(dj1::BigInt, dj2::BigInt, dj3::BigInt, dm1::BigInt, dm2::BigInt, dm3::BigInt)
    check_jm(dj1, dm1) || error(jm_mismatching_msg(dj1, dm1))
    check_jm(dj2, dm2) || error(jm_mismatching_msg(dj2, dm2))
    check_jm(dj3, dm3) || error(jm_mismatching_msg(dj3, dm3))
    check_couple(dj1, dj2, dj3) || error(miss_couple_msg(dj1, dj2, dj3))
    dm1 + dm2 == dm3 || return zero(SqrtRational)
    J::BigInt = div(dj1 + dj2 + dj3, 2)
    Jm1::BigInt = J - dj1
    Jm2::BigInt = J - dj2
    Jm3::BigInt = J - dj3
    j1mm1::BigInt = div(dj1 - dm1, 2)
    j2mm2::BigInt = div(dj2 - dm2, 2)
    j3mm3::BigInt = div(dj3 - dm3, 2)
    j2pm2::BigInt = div(dj2 + dm2, 2)
    A = SqrtRational(
        binomial(dj1, Jm2) * binomial(dj2, Jm3) // (
            binomial(J + 1, Jm3) * binomial(dj1, j1mm1) * binomial(dj2, j2mm2) * binomial(dj3, j3mm3)
        )
    )
    B::BigInt = zero(BigInt)
    low::BigInt = max(zero(BigInt), j1mm1 - Jm2, j2pm2 - Jm1)
    high::BigInt = min(Jm3, j1mm1, j2pm2)
    for z in low:high
        B = -B + binomial(Jm3, z) * binomial(Jm2, j1mm1 - z) * binomial(Jm1, j2pm2 - z)
    end
    return A * (iphase(high) * B)
end

"""
use CG coefficient to calculate 3j-symbol
"""
function _d3j(dj1::BigInt, dj2::BigInt, dj3::BigInt, dm1::BigInt, dm2::BigInt, dm3::BigInt)
    iphase(dj1 + div(dj3 + dm3, 2)) * SqrtRational(1 // (dj3 + 1)) * _dCG(dj1, dj2, dj3, -dm1, -dm2, dm3)
end

"""
basic 6j-symbol calculation funciton
"""
function _d6j(dj1::BigInt, dj2::BigInt, dj3::BigInt, dj4::BigInt, dj5::BigInt, dj6::BigInt)
    check_couple(dj1, dj2, dj3) || error(miss_couple_msg(dj1, dj2, dj3))
    check_couple(dj1, dj5, dj6) || error(miss_couple_msg(dj1, dj5, dj6))
    check_couple(dj4, dj2, dj6) || error(miss_couple_msg(dj4, dj2, dj6))
    check_couple(dj4, dj5, dj3) || error(miss_couple_msg(dj4, dj5, dj3))
    j123::BigInt = div(dj1 + dj2 + dj3, 2)
    j156::BigInt = div(dj1 + dj5 + dj6, 2)
    j426::BigInt = div(dj4 + dj2 + dj6, 2)
    j453::BigInt = div(dj4 + dj5 + dj3, 2)
    jpm123::BigInt = div(dj1 + dj2 - dj3, 2)
    jpm132::BigInt = div(dj1 + dj3 - dj2, 2)
    jpm231::BigInt = div(dj2 + dj3 - dj1, 2)
    jpm156::BigInt = div(dj1 + dj5 - dj6, 2)
    jpm426::BigInt = div(dj4 + dj2 - dj6, 2)
    jpm453::BigInt = div(dj4 + dj5 - dj3, 2)
    A = SqrtRational(
        binomial(j123 + 1, dj1 + 1) * binomial(dj1, jpm123) // (
            binomial(j156 + 1, dj1 + 1) * binomial(dj1, jpm156) *
            binomial(j453 + 1, dj4 + 1) * binomial(dj4, jpm453) *
            binomial(j426 + 1, dj4 + 1) * binomial(dj4, jpm426)
        )
    )
    A = A * Rational{BigInt}(1 // (dj4 + 1))
    B::BigInt = zero(BigInt)
    low::BigInt = max(j123, j453, j426, j156)
    high::BigInt = min(jpm123 + j453, jpm132 + j426, jpm231 + j156)
    for n = low:high
        B = -B + binomial(n + 1, j123 + 1) * binomial(jpm123, n - j453) *
                binomial(jpm132, n - j426) * binomial(jpm231, n - j156)
    end
    return A * (iphase(high) * B)
end

"""
basic 9j-symbol calculation funciton
"""
function _d9j(dj1::BigInt, dj2::BigInt, dj3::BigInt,
              dj4::BigInt, dj5::BigInt, dj6::BigInt,
              dj7::BigInt, dj8::BigInt, dj9::BigInt)
    check_couple(dj1, dj2, dj3) || error(miss_couple_msg(dj1, dj2, dj3))
    check_couple(dj4, dj5, dj6) || error(miss_couple_msg(dj4, dj5, dj6))
    check_couple(dj7, dj8, dj9) || error(miss_couple_msg(dj7, dj8, dj9))
    check_couple(dj1, dj4, dj7) || error(miss_couple_msg(dj1, dj4, dj7))
    check_couple(dj2, dj5, dj8) || error(miss_couple_msg(dj2, dj5, dj8))
    check_couple(dj3, dj6, dj9) || error(miss_couple_msg(dj3, dj6, dj9))
    j123::BigInt = div(dj1 + dj2 + dj3, 2)
    j456::BigInt = div(dj4 + dj5 + dj6, 2)
    j789::BigInt = div(dj7 + dj8 + dj9, 2)
    j147::BigInt = div(dj1 + dj4 + dj7, 2)
    j258::BigInt = div(dj2 + dj5 + dj8, 2)
    j369::BigInt = div(dj3 + dj6 + dj9, 2)
    pm123::BigInt = div(dj1 + dj2 - dj3, 2)
    pm132::BigInt = div(dj1 + dj3 - dj2, 2)
    pm231::BigInt = div(dj2 + dj3 - dj1, 2)
    pm456::BigInt = div(dj4 + dj5 - dj6, 2)
    pm465::BigInt = div(dj4 + dj6 - dj5, 2)
    pm564::BigInt = div(dj5 + dj6 - dj4, 2)
    pm789::BigInt = div(dj7 + dj8 - dj9, 2)
    pm798::BigInt = div(dj7 + dj9 - dj8, 2)
    pm897::BigInt = div(dj8 + dj9 - dj7, 2)
    P0_nu::BigInt = binomial(j123 + 1, dj1 + 1) * binomial(dj1, pm123) *
                    binomial(j456 + 1, dj5 + 1) * binomial(dj5, pm456) *
                    binomial(j789 + 1, dj9 + 1) * binomial(dj9, pm798)
    P0_de::BigInt = binomial(j147 + 1, dj1 + 1) * binomial(dj1, div(dj1 + dj4 - dj7, 2)) *
                    binomial(j258 + 1, dj5 + 1) * binomial(dj5, div(dj2 + dj5 - dj8, 2)) *
                    binomial(j369 + 1, dj9 + 1) * binomial(dj9, div(dj3 + dj9 - dj6, 2))
    P0 = SqrtRational(P0_nu // P0_de)
    dtl::BigInt = max(abs(dj2 - dj6), abs(dj4 - dj8), abs(dj1 - dj9))
    dth::BigInt = min(dj2 + dj6, dj4 + dj8, dj1 + dj9)
    PABC::Rational{BigInt} = zero(Rational{BigInt})
    for dt::BigInt = dtl:2:dth
        j19t::BigInt = div(dj1 + dj9 + dt, 2)
        j26t::BigInt = div(dj2 + dj6 + dt, 2)
        j48t::BigInt = div(dj4 + dj8 + dt, 2)
        Pt_de::BigInt = binomial(j19t + 1, dt + 1) * binomial(dt, div(dj1 + dt - dj9, 2)) *
                        binomial(j26t + 1, dt + 1) * binomial(dt, div(dj2 + dt - dj6, 2)) *
                        binomial(j48t + 1, dt + 1) * binomial(dt, div(dj4 + dt - dj8, 2))
        Pt_de *= (dt + 1)^2
        xl::BigInt = max(j123, j369, j26t, j19t)
        xh::BigInt = min(pm123 + j369, pm132 + j26t, pm231 + j19t)
        At::BigInt = zero(BigInt)
        for x = xl:xh
            At = -At + binomial(x + 1, j123 + 1) * binomial(pm123, x - j369) * binomial(pm132, x - j26t) * binomial(pm231, x - j19t)
        end
        yl::BigInt = max(j456, j26t, j258, j48t)
        yh::BigInt = min(pm456 + j26t, pm465 + j258, pm564 + j48t)
        Bt::BigInt = zero(BigInt)
        for y = yl:yh
            Bt = -Bt + binomial(y + 1, j456 + 1) * binomial(pm456, y - j26t) * binomial(pm465, y - j258) * binomial(pm564, y - j48t)
        end
        zl::BigInt = max(j789, j19t, j48t, j147)
        zh::BigInt = min(pm789 + j19t, pm798 + j48t, pm897 + j147)
        Ct::BigInt = zero(BigInt)
        for z = zl:zh
            Ct = -Ct + binomial(z + 1, j789 + 1) * binomial(pm789, z - j19t) * binomial(pm798, z - j48t) * binomial(pm897, z - j147)
        end
        PABC += (iphase(xh + yh + zh) * At * Bt * Ct) // Pt_de
    end
    return P0 * (iphase(dth) * PABC)
end

"""
    dCG(dj1::Integer, dj2::Integer, dj3::Integer, dm1::Integer, dm2::Integer, dm3::Integer)
CG coefficient function with double angular monentum number parameters, so that the parameters can be integer.
You can calculate `dCG(1, 1, 2, 1, 1, 2)` to calculate the real `CG(1//2, 1//2, 1, 1/2, 1//2, 1)`
"""
dCG(dj1::Integer, dj2::Integer, dj3::Integer, dm1::Integer, dm2::Integer, dm3::Integer) = _dCG(BigInt.((dj1, dj2, dj3, dm1, dm2, dm3))...)

"""
    d3j(dj1::Integer, dj2::Integer, dj3::Integer, dm1::Integer, dm2::Integer, dm3::Integer)
3j-symbol function with double angular monentum number parameters, so that the parameters can be integer.
"""
d3j(dj1::Integer, dj2::Integer, dj3::Integer, dm1::Integer, dm2::Integer, dm3::Integer) = _d3j(BigInt.((dj1, dj2, dj3, dm1, dm2, dm3))...)

"""
    d6j(dj1::Integer, dj2::Integer, dj3::Integer, dj4::Integer, dj5::Integer, dj6::Integer)
6j-symbol function with double angular monentum number parameters, so that the parameters can be integer.
"""
d6j(dj1::Integer, dj2::Integer, dj3::Integer, dj4::Integer, dj5::Integer, dj6::Integer) = _d6j(BigInt.((dj1, dj2, dj3, dj4, dj5, dj6))...)

"""
    d9j(dj1::Integer, dj2::Integer, dj3::Integer,
        dj4::Integer, dj5::Integer, dj6::Integer,
        dj7::Integer, dj8::Integer, dj9::Integer)
9j-symbol function with double angular monentum number parameters, so that the parameters can be integer.
"""
d9j(dj1::Integer, dj2::Integer, dj3::Integer,
    dj4::Integer, dj5::Integer, dj6::Integer,
    dj7::Integer, dj8::Integer, dj9::Integer) = _d9j(BigInt.((dj1, dj2, dj3, dj4, dj5, dj6, dj7, dj8, dj9))...)

@doc raw"""
    CG(j1::HalfInt, j2::HalfInt, j3::HalfInt, m1::HalfInt, m2::HalfInt, m3::HalfInt)
CG coefficient ``\langle j_1m_1 j_2m_2 | j_3m_3 \rangle``
"""
CG(j1::HalfInt, j2::HalfInt, j3::HalfInt, m1::HalfInt, m2::HalfInt, m3::HalfInt) = _dCG(BigInt.((2j1, 2j2, 2j3, 2m1, 2m2, 2m3))...)

@doc raw"""
    threeJ(j1::HalfInt, j2::HalfInt, j3::HalfInt, m1::HalfInt, m2::HalfInt, m3::HalfInt)
Wigner 3j-symbol
```math
\begin{pmatrix}
j_1 & j_2 & j_3 \\
m_1 & m_2 & m_3
\end{pmatrix}
```
"""
threeJ(j1::HalfInt, j2::HalfInt, j3::HalfInt, m1::HalfInt, m2::HalfInt, m3::HalfInt) = _d3j(BigInt.((2j1, 2j2, 2j3, 2m1, 2m2, 2m3))...)

@doc raw"""
    sixJ(j1::HalfInt, j2::HalfInt, j3::HalfInt, j4::HalfInt, j5::HalfInt, j6::HalfInt)
Wigner 6j-symbol
```math
\begin{Bmatrix}
j_1 & j_2 & j_3 \\
j_4 & j_5 & j_6
\end{Bmatrix}
```
"""
sixJ(j1::HalfInt, j2::HalfInt, j3::HalfInt, j4::HalfInt, j5::HalfInt, j6::HalfInt) = _d6j(BigInt.((2j1, 2j2, 2j3, 2j4, 2j5, 2j6))...)

@doc raw"""
    nineJ(j1::HalfInt, j2::HalfInt, j3::HalfInt,
          j4::HalfInt, j5::HalfInt, j6::HalfInt,
          j7::HalfInt, j8::HalfInt, j9::HalfInt)
Wigner 9j-symbol
```math
\begin{Bmatrix}
j_1 & j_2 & j_3 \\
j_4 & j_5 & j_6 \\
j_7 & j_8 & j_9
\end{Bmatrix}
```
"""
nineJ(j1::HalfInt, j2::HalfInt, j3::HalfInt,
      j4::HalfInt, j5::HalfInt, j6::HalfInt,
      j7::HalfInt, j8::HalfInt, j9::HalfInt) = _d9j(BigInt.((2j1, 2j2, 2j3, 2j4, 2j5, 2j6, 2j7, 2j8, 2j9))...)


struct LinearObservation{T}
    H::Array{T}
    R::Array{T}
end

function observe(s, lo::LinearObservation)
    return lo.H * s + Gaussian(lo.R)
end


@testset "domain.jl" begin
    # Closed interval domain
    a, b = 0.0, 1.0
    domain = ClosedIntervalDomain(a, b)
    @test domain.left == a
    @test domain.right == b

    N = 10
    x = discretize(domain, N)
    @test length(x) == N + 1
    @test minimum(x) == a
    @test maximum(x) == b

    # Periodic interval domain
    a, b = -1.0, 1π
    domain = PeriodicIntervalDomain(a, b)
    @test domain.left == a
    @test domain.right == b

    N = 15
    x = discretize(domain, N)
    @test length(x) == N
    @test minimum(x) == a + (b - a) / N
    @test maximum(x) == b

    # Concrete unknown domain type
    domain = UnknownDomain()
    @test_throws Exception discretize(domain, 10)
end


using SumProductNetworks

# Load an automatic differentiation (AD) package, this case we use ForwardDiff.
# Note that using a reverse mode AD is more efficient.
using ForwardDiff

# Loading some utility functions.
using Random, LinearAlgebra

# Generate some training data.
N = 1000
D = 4
x = rand(1:2, N, D)

# Helper function to construct an SPN structure.
function buildSPN(T::Type{<:Real}, x; K = 2)
    N, D = size(x)
    root = FiniteSumNode{T}(;D = D);

    for k in 1:K
        add!(root, FiniteProductNode(D = D), map(T, log(1.0/K)))
        for d in 1:D
            add!(root[k], FiniteSumNode{T}(D = D))
            for l in 1:2
                add!(root[k][d], IndicatorNode(l, d), map(T, log(0.5)))
            end
        end
    end
    updatescope!(root)
    return SumProductNetwork(root)
end

# Number of children under a sum node.
K = 2

# Construct the SPN.
spn = buildSPN(Real, x, K = K)

# Collect all sum nodes.
snodes = filter(n -> isa(n, SumNode), values(spn))

# Helper function used by ForwardDiff.
function f(θ)

    N, D = size(x)
    c = 1
    for i in 1:length(snodes)
        K = length(snodes[i])

        # Extraction of parameters.
        ϕ = θ[c:(c+K-1)]

        # Copy parameters to SPN node.
        snodes[i].logweights[:] = log.(projectToPositiveSimplex!(ϕ))

        c += length(snodes[i])
    end

    # Return llh.
    return -mean(logpdf(spn.root, x))
end

# Create an initial guess for all parameters and project them to the positive simplex.
q = mapreduce(n -> projectToPositiveSimplex!(rand(length(n))), vcat, snodes)

@info "Number of parameters to optimize: $(length(q))"

# Configure the AD.
chunk = ForwardDiff.Chunk(10)
config = ForwardDiff.GradientConfig(f, q, chunk)

# Construct AD function.
g(θ) = ForwardDiff.gradient!(similar(θ), f, θ, config)

# Number of iterations.
τ = 100

# log likelihood values.
ℓ = zeros(τ)

@info "Starting optimization for $(τ) iterations..."
@info "Initial log likelihood: $(-f(q))"

for i in 1:τ

    # Save LLH.
    ℓ[i] = -f(q)
    # Gradient step.
    δq = g(q)
    q .-= 0.1 * δq
end

@info "Log likelihood of optimized model: $(last(ℓ))"


using LinearAlgebra
using Distributions
using ForwardDiff
using Pathfinder
using Test

@testset "multi path pathfinder" begin
    @testset "MvNormal" begin
        n = 10
        nruns = 20
        ndraws = 1000_000
        ndraws_per_run = ndraws ÷ nruns
        Σ = rand_pd_mat(Float64, n)
        μ = randn(n)
        d = MvNormal(μ, Σ)
        logp(x) = logpdf(d, x)
        ∇logp(x) = ForwardDiff.gradient(logp, x)
        x₀s = [rand(Uniform(-2, 2), n) for _ in 1:nruns]
        q, ϕ, component_ids = multipathfinder(
            logp, ∇logp, x₀s, ndraws; ndraws_elbo=100, ndraws_per_run=ndraws_per_run
        )
        @test q isa MixtureModel
        @test ncomponents(q) == nruns
        @test Distributions.component_type(q) <: MvNormal
        @test ϕ isa AbstractMatrix
        @test size(ϕ) == (n, ndraws)
        @test component_ids isa Vector{Int}
        @test length(component_ids) == ndraws
        @test extrema(component_ids) == (1, nruns)
        μ_hat = mean(ϕ; dims=2)
        Σ_hat = cov(ϕ .- μ_hat; dims=2, corrected=false)
        # adapted from the MvNormal tests
        # allow for 15x disagreement in atol, since this method is approximate
        multiplier = 15
        for i in 1:n
            atol = sqrt(Σ[i, i] / ndraws) * 8 * multiplier
            @test μ_hat[i] ≈ μ[i] atol = atol
        end
        for i in 1:n, j in 1:n
            atol = sqrt(Σ[i, i] * Σ[j, j] / ndraws) * 10 * multiplier
            @test isapprox(Σ_hat[i, j], Σ[i, j], atol=atol)
        end
    end
end


mutable struct ADM{U, O} <: AbstractSubspaceOptimizer
    λ::U
    R::O
end

get_threshold(opt::ADM) = opt.λ
set_threshold!(opt::ADM, λ) = (opt.λ = λ; return)

"""
    ADM()
    ADM(λ = 0.1)

Optimizer for finding a sparse basis vector in a subspace based on [this paper](https://arxiv.org/pdf/1412.4659.pdf).
`λ` is the weight for the soft-thresholding operation.
"""
function ADM(λ::T = 0.1) where T <: Real
    f = NormL1(λ)
    return ADM(λ, f)
end

# ADM algorithm
function fit!(q::AbstractArray{T, 1}, Y::AbstractArray, opt::ADM; maxiter::Int64= 10, tol::T = eps(eltype(q))) where T <: Real

    x = Y*q
    q_ = deepcopy(q)
    iters_ = 0

    while iters_ <= maxiter
        iters_ += 1
        prox!(x, opt.R, Y*q)
        mul!(q, Y', x)
        normalize!(q, 2)

        if norm(q - q_) < tol
            break
        else
            q_ .= q
        end
    end

    return iters_
end

# ADM initvary
function fit!(q::AbstractArray{T, 2}, Y::AbstractArray, opt::ADM; maxiter::Int64= 10, tol::T = eps(eltype(q))) where T <: Real
    iters_ = Inf
    i_ = Inf

    @inbounds for i in 1:size(q, 2)
        i_ = fit!(q[:, i], Y, opt, maxiter = maxiter, tol = tol)
        if iters_ > i_
            iters_ = i_
        end
    end

    return iters_
end

# ADM pareto
function fit!(X::AbstractArray, A::AbstractArray, Y::AbstractArray, opt::ADM; rtol = 0.99, maxiter::Int64 = 1, convergence_error::T = eps(), f::Function = (xi, theta)->[norm(xi, 0); norm(theta'*xi, 2)], g::Function = x->norm(x),) where T <: Real
    # Return just the best candidate for the subspace optimization

    θ = zeros(eltype(A), size(X, 1), size(A, 2))
    θ[size(A, 1)+1:end, :] .= A

    fg(xi, theta) = (g∘f)(xi, theta)

    iters = Inf

    @inbounds for i in 1:size(Y, 1)
        for j in 1:size(A, 2)
            @views θ[1:size(A, 1), j] .= Y[i, j].*A[:, j]
        end

        N = nullspace(θ', rtol = rtol)
        Q = deepcopy(N) # Deepcopy for inplace
        # Find sparse vectors in nullspace
        # Calls effectively the ADM algorithm with varying initial conditions
        fit!(Q, N', opt, maxiter = maxiter)

        # Find sparse vectors in nullspace
        # Calls effectively the ADM algorithm with varying initial conditions
        iters_ = fit!(Q, N', opt, maxiter = maxiter, tol = convergence_error)
        iters_ < iters ? iters = iters_ : nothing

        # Compute pareto front
        for (j, ξ) in enumerate(eachcol(Q))
            if j == 1
                X[:, i] .= ξ
            else
                evaluate_pareto!(view(X, :, i), view(ξ, :), fg, view(θ, :, :))
            end
        end

        X[abs.(X[:, i]) .< get_threshold(opt), i] .= zero(eltype(X))
        @views X[:, i] .= X[:, i] ./ maximum(abs, X[:, i])
    end

    return iters
end


include("gemm.jl")

using LinearAlgebra

let
    nsamples = 10
    ns = 48:48:1500
    for iter = 1:length(ns)
        n = m = k = ns[iter]
        C = rand(n,m)
        A = rand(n,k)
        B = rand(k,m)

        # julia mul!
        t_mul! = floatmax(Float64)
        for s = 1:nsamples
            t = @elapsed mul!(C,A,B)
            t_mul! = min(t,t_mul!)
        end
        GFLOPS_mul![iter] = 2*n*m*k / t_mul! / 1e9

        t_ijp! = floatmax(Float64)
        for s = 1:nsamples
            t = @elapsed mygemm_ijp!(C,A,B)
            t_ijp! = min(t,t_ijp!)
        end

    end
    df = DataFrame(m = ns,

    )
end


" Rock-salt lattice "
rock_salt(T::Type=Float64; unit=u"nm") = Crystal(
    eye(T, 3) * unit,
    position=T[0.5, 0.5, 0.5] * unit,
    position=T[0, 0, 0] * unit,
    species=['A', 'B']
)

" Zinc-blende lattice "
zinc_blende(T::Type=Float64; unit=u"nm") = Crystal(
    T[0 0.5 0.5; 0.5 0 0.5; 0.5 0.5 0] * unit,
    tposition=T[0 0 0; 0.25 0.25 0.25] * unit,
    species=['A', 'B']
)

" Diamond lattice "
diamond(T::Type=Float64; unit=u"nm") = Crystal(
    T[0 0.5 0.5; 0.5 0 0.5; 0.5 0.5 0] * unit,
    tposition=T[0 0 0; 0.25 0.25 0.25] * unit
)

" Wurtzite lattice "
wurtzite(T::Type=Float64; unit=u"nm") = Crystal(
    T[0.5 0.5 0; -sqrt(3)/2 sqrt(3)/2 0; 0 0 1] * unit,
    tposition=T[0.5  sqrt(3)/6 0;
                0.5 -sqrt(3)/6 0.5;
                0.5  sqrt(3)/6 0.25;
                0.5 -sqrt(3)/6 0.75] * unit,
    species=['A', 'A', 'B', 'B']
)


@info "loading packages"
using Distributed
addprocs(30; exeflags="--project")

@everywhere begin

using TicketToRide
using Flux
using Flux.Optimise
using TicketToRide: train!
using LuxurySparse
using LinearAlgebra

function task(n, nlayers; verbose=false, nprune=10, nepochs=10, niteration=100, relative_error=true, least_prune=10)
    opt = Optimise.Descent(1e-2)
    circuit = variational_circuit(n, nlayers);
    ham = heisenberg1D(n);
    dispatch!(circuit, :random);
    circuit, history = prune_train(opt, circuit, ham;
        nprune=nprune, nepochs=nepochs,
        niteration=niteration,
        relative_error=relative_error,
        least_prune=least_prune);

    opt = Optimise.Descent(1e-3)
    circuit = variational_circuit(n, nlayers);
    ham = heisenberg1D(n);
    dispatch!(circuit, :random);
    circuit, history = prune_train(opt, circuit, ham;
        nprune=nprune, nepochs=nepochs,
        niteration=niteration,
        relative_error=relative_error,
        least_prune=least_prune);

    return circuit, history
end

function run_task(total, n, nlayers; verbose=false, nprune=10, nepochs=10, niteration=100, relative_error=true, least_prune=10)
    results = pmap(1:total) do k
        task(n, nlayers;
        nprune=nprune, nepochs=nepochs,
        niteration=niteration,
        relative_error=relative_error,
        least_prune=least_prune)
    end

    min_energy = 0
    circuit, history = nothing, nothing
    for each in results
        if each[2][end] < min_energy
            circuit = each[1]
            history = each[2]
        end
    end
    return circuit, history
end

end

using JLD2
# @time c, h = run_task(2, 2, 2; nepochs=2, nprune=5, niteration=2, least_prune=5)

@time c, h = run_task(nprocs(), 10, 2; nepochs=2, nprune=5, niteration=2, least_prune=5)
c, h = run_task(nprocs(), 10, 100; nepochs=20, nprune=50, niteration=2000, least_prune=5)

jldopen("reinitialize-data-10-10.jld", "w+") do f
    f["circuit"] = c
    f["history"] = h
end

# opt = Optimise.Descent(1e-3)
# circuit = variational_circuit(10, 100);
# nparameters(circuit)
# dispatch!(circuit, :random);
# history = train!(opt, 2000, heisenberg1D(10), circuit; verbose=true)
# prune_train(opt, circuit, heisenberg1D(10);
#     nprune=50, nepochs=20,
#     niteration=500,
#     relative_error=true,
#     least_prune=5);

# nparameters(circuit)


abstract type AbstractHyperplane <: AbstractSpatial end

abstract type AbstractLine <: AbstractHyperplane end
abstract type AbstractPlane <: AbstractHyperplane end

for (type, supertype) in zip([:Line, :Plane], [:AbstractLine, :AbstractPlane])

    @eval struct $type{N, T} <: $supertype

        point::SVector{N, T}
        vector::SVector{N, T}

        function $type(point::SVector{N, T}, vector::SVector{N, T}; kwargs...) where {N, T}

            if is_zero(vector; kwargs...)
                throw(ArgumentError("The vector must have a non-zero magnitude."))
            end

            new{N, T}(point, vector)
        end
    end

    @eval function $type(point::AbstractVector, vector::AbstractVector; kwargs...)

        point_static, vector_static = _convert_to_svector([point, vector])

        return $type(point_static, vector_static; kwargs...)
    end
end

for (type, alias) in zip(
    (:AbstractLine, :AbstractPlane),
    (:(:direction), :(:normal)),
)
    @eval function Base.getproperty(obj::$type, symbol::Symbol)

        if symbol === $alias
            return obj.vector
        end

        return getfield(obj, symbol)
    end
end


using DiagonalSBP
using SparseArrays: sparsevec, spzeros, sparse
using LinearAlgebra: I, eigvals
using Logging: @info
using Printf: @sprintf
using PGFPlotsX: @pgf, Plot, Axis, Table, pgfsave, SemiLogYAxis, LogLogAxis
using ForwardDiff: derivative

function noncharacteristic_operator(p, N, R)
  # Total number of points
  Nq = N + 1

  # Get the various derivative operators
  (D, S0, SN, HI, H, x) = DiagonalSBP.D2(p, N; xc = (0, 1))

  # Form the SBP inner product matrix
  A = -H * D + SN - S0

  L¹ = sparsevec([ 1], [1], Nq)'
  L² = sparsevec([Nq], [1], Nq)'

  α = -(1 - R) / (1 + R)

  # Solving the system in first order form with [u; v]
  Ouu = spzeros(Nq, Nq)
  Ouv = I
  Ovu = -A
  Ovv = α * L²' * L² + α * L¹' * L¹

  # Form the operator
  O = [Ouu Ouv;
       HI*Ovu HI*Ovv]

  return (O, x, H)
end

function characteristic_operator(p, N, R)
  # Total number of points
  Nq = N + 1

  # Get the various derivative operators
  (D, S0, SN, HI, H, x) = DiagonalSBP.D2(p, N; xc = (0, 1))
  B¹, B² = S0[1:1, :], SN[Nq:Nq, :]

  # Form the SBP inner product matrix
  A = -H * D + SN - S0

  L¹ = sparse([1], [ 1], [1], 1, Nq)
  L² = sparse([1], [Nq], [1], 1, Nq)
  n¹ = -1
  n² = +1

  θ_R = DiagonalSBP.D2_remainder_parameters(p).θ_R
  θ_H = DiagonalSBP.D2_remainder_parameters(p).θ_H
  Γ = N * (1 / θ_R + 1 / θ_H)

  # Solving the system in first order form with [u; v; ϕ¹; ϕ²]
  Ouu = spzeros(Nq, Nq)
  Ouv = I
  Ouϕ¹ = spzeros(Nq, 1)
  Ouϕ² = spzeros(Nq, 1)

  Ovu  = -A
  Ovv  = spzeros(Nq, Nq)
  Ovϕ¹ = spzeros(Nq, 1)
  Ovϕ² = spzeros(Nq, 1)

  Oϕ¹u  = spzeros(1, Nq)
  Oϕ¹v  = spzeros(1, Nq)
  Oϕ¹ϕ¹ = spzeros(1, 1)
  Oϕ¹ϕ² = spzeros(1, 1)

  Oϕ²u  = spzeros(1, Nq)
  Oϕ²v  = spzeros(1, Nq)
  Oϕ²ϕ¹ = spzeros(1, 1)
  Oϕ²ϕ² = spzeros(1, 1)

  Ovu  += ((1 - R) / 2) * L¹' * (n¹ * B¹ - Γ * L¹) + n¹ * B¹' * L¹
  Ovv  += -((1 - R) / 2) * L¹' * L¹
  Ovϕ¹ += ((1 - R) / 2) * Γ * L¹' - n¹ * B¹'

  Oϕ¹u  += ((1 + R) / 2) * (Γ * L¹ - n¹ * B¹)
  Oϕ¹v  += ((1 + R) / 2) * L¹
  Oϕ¹ϕ¹ += sparse([1], [1], [-Γ * (1 + R) / 2], 1, 1)

  Ovu  +=  ((1 - R) / 2) * L²' * (n² * B² - Γ * L²) + n² * B²' * L²
  Ovv  += -((1 - R) / 2) * L²' * L²
  Ovϕ² +=  ((1 - R) / 2) * Γ * L²' - n² * B²'

  Oϕ²u  += ((1 + R) / 2) * (Γ * L² - n² * B²)
  Oϕ²v  += ((1 + R) / 2) * L²
  Oϕ²ϕ² += sparse([1], [1], [-Γ * (1 + R) / 2], 1, 1)

  # Form the operator
  O = [    Ouu  Ouv  Ouϕ¹  Ouϕ²
       HI*[Ovu  Ovv  Ovϕ¹  Ovϕ²]
           Oϕ¹u Oϕ¹v Oϕ¹ϕ¹ Oϕ¹ϕ²
           Oϕ²u Oϕ²v Oϕ²ϕ¹ Oϕ²ϕ²]

  return (O, x, H)
end

function convergence_test(p, Ns, R, t_final = 1.0)
  @info "sbp order = $p"

  # Pulse center
  μ = 1 / 2

  # Pulse width
  w = 1 / 15

  # Analytic solution
  # f(x) = exp(-((x - μ) / w)^2) / 2
  f(x) = 0 < x < 1 ? sin(2 * pi * x)^6 : 0

  ue(x, t) = (f(x - t) + f(x + t)) + R * (f(2 - x - t) + f(-x + t))
  ∂t_ue(x, t) = derivative(t->ue(x, t), t)

  # storage for error
  err_NC  = zeros(length(Ns))
  err_C   = zeros(length(Ns))

  # loop over all N
  for (iter, N) = enumerate(Ns)
    @info @sprintf " level %d with N = %4d" iter N

    (NC, x, H) = noncharacteristic_operator(p, N, R)
    q0 = [ue.(x, 0); ∂t_ue.(x, 0)]

    qf = exp(t_final * Matrix(NC)) * q0

    Δu = ue.(x, t_final) - qf[1:N+1]
    Δv = ∂t_ue.(x, t_final) - qf[(N+2):(2N+2)]

    err_NC[iter] = sqrt(Δu' * H * Δu)
    if iter > 1
      rate = ((log(err_NC[iter-1]) - log(err_NC[iter])) /
              (log(Ns[iter]) - log(Ns[iter-1])))
      @info @sprintf """ Non-Characteristic
      error = %.2e
      rate  = %.2e""" err_NC[iter] rate
    else
      @info @sprintf """ Non-Characteristic
      error = %.2e""" err_NC[iter]
    end

    (C, x, H) = characteristic_operator(p, N, R)
    q0 = [ue.(x, 0); ∂t_ue.(x, 0); 0; 0]

    qf = exp(t_final * Matrix(C)) * q0

    Δu = ue.(x, t_final) - qf[1:N+1]
    Δv = ∂t_ue.(x, t_final) - qf[(N+2):(2N+2)]

    err_C[iter] = sqrt(Δu' * H * Δu)
    if iter > 1
      rate = ((log(err_C[iter-1]) - log(err_C[iter])) /
              (log(Ns[iter]) - log(Ns[iter-1])))
      @info @sprintf """ Characteristic
      error = %.2e
      rate  = %.2e""" err_C[iter] rate
    else
      @info @sprintf """ Characteristic
      error = %.2e""" err_C[iter]
    end
  end
  return err_C, err_NC
end

function make_spectrum_plots(p, N, R)
  h = 1 / N
  (NC, _...) = noncharacteristic_operator(p, N, R)
  (C, _...)  = characteristic_operator(p, N, R)

  ev_NC = h * eigvals(Matrix(NC))
  ev_C  = h * eigvals(Matrix(C))

  ymin, ymax, xmin, xmax = -2.5, 2.5, -3.35, 1e-1

  pts = nothing
  if !all(xmin .< real.(ev_NC) .< xmax)
    pts = ev_NC[findall(.!(xmin .< real.(ev_NC) .< xmax))]
    @assert length(pts) == 2
    @assert pts[1] ≈ pts[2]
    @assert abs(imag(pts[1])) < 1e-14
    ev_NC = [ev_NC[xmin .< real.(ev_NC) .< xmax]; -3.1]
  end
  @assert all(xmin .< real.(ev_C ) .< xmax)
  @assert all(ymin .< imag.(ev_NC) .< ymax)
  @assert all(ymin .< imag.(ev_C ) .< ymax)

  xmin = xmin - 0.2

  @pgf a = Axis(
            {
             ymin = ymin,
             ymax = ymax,
             xmin = xmin,
             xmax = xmax,
             legend_entries = {"non-characteristic", "characteristic"},
             title = "R = $R, N = $N",
             xlabel = "\$\\textrm{Re}(h \\lambda)\$",
             ylabel = "\$\\textrm{Im}(h \\lambda)\$",
            }
           )
  @pgf push!(a, Plot(
                     {
                      only_marks,
                      color = "red",
                      mark = "x",
                     },
                     Table(real.(ev_NC), imag.(ev_NC)))
            )
  @pgf push!(a, Plot(
                     {
                      only_marks,
                      color = "blue",
                      mark = "+",
                     },
                     Table(real.(ev_C), imag.(ev_C))),
            )
  if !isnothing(pts)
    @pgf push!(a,
               [@sprintf """
                \\node[anchor=south] () at (axis cs:-3.1,0){\$%.2e\$};
                \\node[anchor=west] (source) at (axis cs:-3.1,0){};
                \\node (destination) at (axis cs:-3.35,0){};
                \\draw[->](source)--(destination);
                """ real(pts[1])])
  end

  pgfsave("bc_spectra_$R.tex", a)
  pgfsave("bc_spectra_$R.pdf", a)
end

function compare_eigs(p, N)
  h = 1 / N
  Rs = -0.95:0.05:0.95
  real_C = zeros(length(Rs))
  real_NC = zeros(length(Rs))
  imag_C = zeros(length(Rs))
  imag_NC = zeros(length(Rs))
  for (iter, R) = enumerate(Rs)
    (NC, _...) = noncharacteristic_operator(p, N, R)
    (C, _...)  = characteristic_operator(p, N, R)

    ev_NC = h * eigvals(Matrix(NC))
    ev_C  = h * eigvals(Matrix(C))

    real_C[iter] = minimum(real.(ev_C))
    real_NC[iter] = minimum(real.(ev_NC))
    imag_C[iter] = maximum(imag.(ev_C))
    imag_NC[iter] = maximum(imag.(ev_NC))
  end

  @pgf a = SemiLogYAxis(
                {
                 xmin = -1,
                 xmax =  1,
                 xlabel = "\$R\$",
                 ylabel = "max\$(-\\textrm{Re}(h \\lambda)\$",
                 legend_entries = {"non-characteristic", "characteristic"},
                }
               )
  @pgf push!(a, Plot(
                     {
                      only_marks,
                      color = "red",
                      mark = "x",
                     },
                     Table(Rs, -real_NC)
                    )
            )
  @pgf push!(a, Plot(
                     {
                      only_marks,
                      color = "blue",
                      mark = "+",
                     },
                     Table(Rs, -real_C)
                    )
            )
  pgfsave("real_eig.tex", a)
  pgfsave("real_eig.pdf", a)
end

#=
for R in (-0.99, -0.9, 0, 0.9, 1.0)
  make_spectrum_plots(4, 50, R)
end
compare_eigs(4, 50)
=#

let
  Ns = 17 * 2 .^ (0:5)
  hs = 1 ./ Ns

  tri_x = [hs[end-1], hs[end], hs[end-1], hs[end-1]]
  tri_y = [hs[end-1], hs[end], hs[end  ], hs[end-1]] ./ hs[end-1]
  scale = Dict()

  scale[(-0.99, 2)] = 0.001
  scale[( 0.0 , 2)] = 0.0002
  scale[( 0.99, 2)] = 0.001

  scale[(-0.99, 4)] = 1e-6
  scale[( 0.0 , 4)] = 5e-7
  scale[( 0.99, 4)] = 2e-6

  scale[(-0.99, 6)] = 2e-7
  scale[( 0.0 , 6)] = 1e-7
  scale[( 0.99, 6)] = 1e-8
  for R in (-0.99, 0.0, 0.99)
    C2, NC2 = convergence_test(2, Ns, R, 0.9)
    C4, NC4 = convergence_test(4, Ns, R, 0.9)
    C6, NC6 = convergence_test(6, Ns, R, 0.9)

    @pgf a = LogLogAxis(
                        {
                         ymax = 1,
                         ymin = 1e-11,
                         xmin = hs[end]/2,
                         xmax = 2hs[1],
                         xlabel = "\$h\$",
                         ylabel = "\$\\|\\Delta u\\|_{H}\$",
                        }
                       )
    tc2 = 0.6*NC2[end-1] * tri_y.^2
    @pgf push!(a, Plot(
                       {
                        color = "black",
                       },
                       Table(tri_x, tc2)
                      )
              )
    tc4 = 0.6*NC4[end-1] * tri_y.^4
    @pgf push!(a, Plot(
                       {
                        color = "black",
                       },
                       Table(tri_x, 0.6*NC4[end-1] * tri_y.^4)
                      )
              )
    tc6 = 0.6*NC6[end-1] * tri_y.^6
    @pgf push!(a, Plot(
                       {
                        color = "black",
                       },
                       Table(tri_x, 0.6*NC6[end-1] * tri_y.^5)
                      )
              )
    #=
    @pgf push!(a,
              [@sprintf("""
               \\node[anchor=north] () at (axis cs:%e, %e){\$s = 2\$};
               \\node[anchor=north] () at (axis cs:%e, %e){\$s = 4\$};
               \\node[anchor=north] () at (axis cs:%e, %e){\$s = 5\$};
               """,
              (hs[end-1] + hs[end]) / 2, tc2[2],
              (hs[end-1] + hs[end]) / 2, tc4[2],
              (hs[end-1] + hs[end]) / 2, tc6[2],
               )])
    =#
    @pgf push!(a,
               [@sprintf("""
                \\node[anchor=west] () at (axis cs:%e, %e){\$2\$};
                \\node[anchor=west] () at (axis cs:%e, %e){\$4\$};
                \\node[anchor=west] () at (axis cs:%e, %e){\$5\$};
                """,
                hs[end-1], exp((log(tc2[1]) + log(tc2[2]))/2),
                hs[end-1], exp((log(tc4[1]) + log(tc4[2]))/2),
                hs[end-1], exp((log(tc6[1]) + log(tc6[2]))/2),
               )])

    @pgf push!(a, Plot(
                       {
                        color = "red",
                        mark = "o",
                       },
                       Table(hs, C2)
                      )
              )
    @pgf push!(a, Plot(
                       {
                        color = "red",
                        mark = "x",
                       },
                       Table(hs, NC2)
                      )
              )
    @pgf push!(a, Plot(
                       {
                        color = "blue",
                        mark = "o",
                       },
                       Table(hs, C4)
                      )
              )
    @pgf push!(a, Plot(
                       {
                        color = "blue",
                        mark = "x",
                       },
                       Table(hs, NC4)
                      )
              )
    @pgf push!(a, Plot(
                       {
                        color = "green",
                        mark = "o",
                       },
                       Table(hs, C6)
                      )
              )
    @pgf push!(a, Plot(
                       {
                        color = "green",
                        mark = "x",
                       },
                       Table(hs, NC6)
                      )
              )
    pgfsave("convergence_$R.tex", a)
    pgfsave("convergence_$R.pdf", a)
  end
end


################################################################################
# Copyright (C) Marcus Greiff 2020
#
# @file example_FOF_attitude.jl
# @author Marcus Greiff
# @date July 2020
#
# This example demonstrates the filtered output feedback for the attitude
# dynamics with errors related to distances on SO(3), and a closed loop
# simulation, visualizing the state trajectories along with the observer states
# and the state trajectories. SImilar to the FSF attitude examples, it also
# plots the associated Lyapunov function and distances as computed in the
# controller. Contrary to the FSF controllers, this FOF controller has memory,
# by virtue of the integration of the observer dynamcis - consequently, we can
# only call the function once pwe time-step, which means the ay multistep-method
# has to be avoided in the simulation of the closed loop system in Julia.
# Consequently, the system is simulated with a constant step Euler method
# instead of Tsit5().
#
# Furthermore, the method does not have access to the true system states, so
# The lyapunov function and errors have to be evaluated externally in Julia,
# instead of in C. So contrary to the FSF feedback controllers, the estimated
# state (updated by the FOF controller) is stored in the State structure,
# instead of the actual system state.
#
################################################################################
using DifferentialEquations
using LinearAlgebra
using LaTeXStrings
using AerialVehicleControl

showPlot = [1,2];     # Set to 0 for no plot, or any number in (1,2)
savePlot = false; # Set to true if the plot is to be saved
namePlot = "continuous_SO3"

# Recompile the C-code
recompile()

# Wrapper for the calling of the FOF controller from C
function update_control(R::ref_state_qw_t, S::dyn_state_qw_t, C::con_state_qw_fof_t, dt::Cdouble)
    status = ccall((:update_attitude_FOF_SO3_continuous, AerialVehicleControl.CONT_LIB_PATH),
        Cint,
        (Ref{ref_state_qw_t}, Ref{dyn_state_qw_t}, Ref{con_state_qw_fof_t}, Cdouble,),
        R,
        S,
        C,
        dt)
    return status;
end

# Wrapper for the updating of the measurements in the controller structure
function update_measurements(
    C::con_state_qw_fof_t,
    y0m::matrix_double_t,
    y1m::matrix_double_t,
    y2m::matrix_double_t,
    y3m::matrix_double_t
)
    status = ccall((:update_attitude_FOF_SO3_continuous_measurements, AerialVehicleControl.CONT_LIB_PATH),
        Cint,
        (Ref{con_state_qw_fof_t}, Ref{matrix_double_t}, Ref{matrix_double_t},Ref{matrix_double_t}, Ref{matrix_double_t},),
        C,
        y0m,
        y1m,
        y2m,
        y3m)
    return status;
end

# Right-hand side in the ODE used for numerical integration. Here, we have
#
# x[1:4]   = q     (quaternion, integrated in Julia)
# x[5:7]   = w     (att. rates, integrated in Julia)
# x[8:11]  = qr    (quaternion reference, integrated in Julia)
# x[12:14] = wr    (att. rates reference, integrated in Julia)
# x[15:18] = qhat  (quaternion estimate, integrated in C using CG)
# x[19:21] = what  (att. rates estimate, integrated in C using CG)
# x[22:24] = taur  (torques reference, computed in Julia)
# x[25:27] = tau   (torques actuation, computed in C)
# x[28]    = prevt (time at the previous call of the OF-controller)
# x[29]    = lyap  (Lyapunov function, computed in Julia)
# x[30]    = Psi1  (Estimate attitude error, computed in Julia)
# x[31]    = Psi2  (Tracking attitude error, computed in Julia)
#
function odefun!(dx,x,params,t)
    # Extract the states
    q     = x[1:4];
    w     = x[5:7];
    qr    = x[8:11];
    wr    = x[12:14];
    tprev = x[28];

    S = params[1];
    C = params[2];

    # Extract the parameters
    J = ntuple_2_mat(C.inertia, 3, 3)

    # Reference torques driving the reference dynamics
    taur = [sin(t), sin(2*t), sin(3*t)];

    # Reference dynamics
    dqr, dwr = attitude_dynamics(qr, wr, J, taur)

    # Create state structs for ccall
    R = ref_state_qw_t(hcat(qr), hcat(wr), hcat(dwr), 0.1)

    # Compute control signal through a ccall to the FSF configured on SO3
    if tprev != t
        timestep = Cdouble(t - tprev);
        #println("Computed new control action at a time t=$t, dt = $timestep")

        # Update the measurements in the controller object
        v1 = [C.globalDirections[1][i] for i = 1:3];
        v2 = [C.globalDirections[2][i] for i = 1:3];
        v3 = [C.globalDirections[3][i] for i = 1:3];
        y0 = hcat(w);
        y1 = hcat(quat_2_SO3(hcat(q))'*v1);
        y2 = hcat(quat_2_SO3(hcat(q))'*v2);
        y3 = hcat(quat_2_SO3(hcat(q))'*v3);
        y0m = matrix_double_t(y0);
        y1m = matrix_double_t(y1);
        y2m = matrix_double_t(y2);
        y3m = matrix_double_t(y3);
        @assert(1 == update_measurements(C, y0m, y1m, y2m, y3m));

        # Update the controller
        @assert(1 == update_control(R, S, C, timestep));
    end

    # Actuate the dynamics with the computed control signal
    tau = [C.torques[i] for i in 1:3]
    # Physical dynamics
    dq , dw  = attitude_dynamics( q,  w, J, tau )

    # Externalize state derivatives (the observer dynamics are simulated in the FOF controller)
    dx[1:4]   = dq;
    dx[5:7]   = dw;
    dx[8:11]  = dqr;
    dx[12:14] = dwr;

    # Normalize the quaternions which are simulated by DifferentialEquations.jl
    x[1:4]    = q  / norm(q );
    x[8:11]   = qr / norm(qr);

    # Output the estimates
    qhat = [S.quaternion[i] for i = 1:4];
    what = [S.attrates[i] for i = 1:3];
    x[15:18] = qhat;
    x[19:21] = what;

    # Externalize other signals for visualization
    x[22:24]  = taur;
    x[25:27]  = tau;
    x[28]     = t;   # Set the current time so that we can compute the dt

    # Evaluate distances and the Lyapunov fucntion
    Rmat   = quat_2_SO3(hcat(q))
    Rhat   = quat_2_SO3(hcat(qhat))
    Rr     = quat_2_SO3(hcat(qr))
    Re     = Rr*Rmat';
    Rtilde = Rhat*Rmat';
    we     = wr   - w;
    wtilde = what - w;

    v1 = [C.globalDirections[1][i] for i = 1:3];
    v2 = [C.globalDirections[2][i] for i = 1:3];
    v3 = [C.globalDirections[3][i] for i = 1:3];

    erre1 = Re * Rtilde' * v1 - v1;
    erre2 = Re * Rtilde' * v2 - v2;
    erre3 = Re * Rtilde' * v3 - v3;

    errt1 = Rtilde' * v1 - v1;
    errt2 = Rtilde' * v2 - v2;
    errt3 = Rtilde' * v3 - v3;

    x[29] = (+(
               +erre1'*erre1*C.gain_ki[1]
               +erre2'*erre2*C.gain_ki[2]
               +erre3'*erre3*C.gain_ki[3]
             ) / 2.0
             +we'*J*we
             +(
               +errt1'*errt1*C.gain_ki[1]
               +errt2'*errt2*C.gain_ki[2]
               +errt3'*errt3*C.gain_ki[3]
             ) / 2.0
             +(1.0/(2.0 * C.gain_cw)) * wtilde'*J*wtilde)

    x[30] = tr(Matrix{Float64}(I, 3, 3) - Rhat'*Rmat) / 2.0
    x[31] = tr(Matrix{Float64}(I, 3, 3) - Rr'*Rmat) / 2.0
end

function initialize_example()
    # Randomize the initial system states
    q0 = randn(4,1); q0 = q0./norm(q0); w0 = 2*randn(3,1);

    # Randomize an initial reference of states
    qr0 = randn(4,1); qr0 = qr0./norm(qr0); wr0 = 0*rand(3,1); ar0 = 0.0*rand(3,1); fr0 = 0.5;
    R = ref_state_qw_t(hcat(qr0), hcat(wr0), hcat(ar0), fr0)

    # Randomize a set of observer states
    qhat0 = randn(4,1); qhat0 = qhat0./norm(qhat0); what0 = 2*randn(3,1);
    S = dyn_state_qw_t(hcat(qhat0), hcat(what0))

    # Randomize the controller struct
    J  = rand_PSD(3); J  = J ./norm(J );
    Kw = rand_PSD(3); Kw = Kw./norm(Kw); Kw = 2*Kw;
    ki = (rand(1,3) .+ 0.5); ki = ki./norm(ki); ki = 2*ki;
    Cw = rand_PSD(3); Cw = Cw./norm(Cw); Cw = 2*Cw;
    cw, cR = 2.0, 2.0;
    a, b, c, d = 1.0, 2.0, 3.0, 4.0;
    Y0 = rand(3,1) .- 0.5;
    Yi = rand(3,3) .- 0.5;
    Vi = rand(3,3) .- 0.5;
    Vi = Matrix{Float64}(I, 3, 3);
    # Normalize the directions in Vi
    for i = 1:3
        Vi[:,i] = Vi[:,i] / norm(Vi[:,i])
    end

    # Form controller object
    C  = con_state_qw_fof_t(J, Kw, ki, Cw, cw, cR, a, b, c, d, Y0, Yi, Vi);

    # Initial states for the debugging and logging of signals
    output = zeros(10,1);

    x0    = vcat(q0, w0, qr0, wr0, qhat0, what0, output)
    tspan = (0.0,20.0)
    return x0, tspan, R, S, C
end

# Compute the solution
x0, tspan, _, S, C = initialize_example()
prob               = ODEProblem(odefun!, x0, tspan, [S, C]);
sol                = solve(prob, Euler(), adaptive=false, dt = 1e-4)

# Visualization of the results
if 1 in showPlot
    using Plots
    pqqr = plot(sol, vars = (0,1:4),   color=:black, style=:solid, linewidth=2, xaxis="Time (t)",yaxis="Attitude [.]", label=[L"q(t)" nothing nothing nothing])
    plot!(      sol, vars = (0,8:11),  color=:blue,  style=:dash,  linewidth=2, xaxis="Time (t)",yaxis="Attitude [.]", label=[L"q_r(t)" nothing nothing nothing])
    plot!(      sol, vars = (0,15:18), color=:red,   style=:dot,   linewidth=2, xaxis="Time (t)",yaxis="Attitude [.]", label=[L"\hat{q}(t)" nothing nothing nothing])
    pwwr = plot(sol, vars = (0,5:7),   color=:black, style=:solid, linewidth=2, xaxis="Time (t)",yaxis="Attitude rate [rad/s]", label=[L"\omega(t)" nothing nothing])
    plot!(      sol, vars = (0,12:14), color=:blue,  style=:dash,  linewidth=2, xaxis="Time (t)",yaxis="Attitude rate [rad/s]", label=[L"\omega_{r}(t)" nothing nothing])
    plot!(      sol, vars = (0,19:21), color=:red,   style=:dot,   linewidth=2, xaxis="Time (t)",yaxis="Attitude rate [rad/s]", label=[L"\hat{\omega}(t)" nothing nothing])
    pttr = plot(sol, vars = (0,22:24), color=:black, style=:solid, linewidth=2, xaxis="Time (t)",yaxis="Torque [Nm]",           label=[L"\tau_r(t)" nothing nothing])
    plot!(      sol, vars = (0,25:27), color=:blue,  style=:solid, linewidth=2, xaxis="Time (t)",yaxis="Torque [Nm]",           label=[L"\tau(t)" nothing nothing])
    plot(pqqr, pwwr, pttr, layout=(3,1), size=(1000,750))
    gui()
    if savePlot
        savefig(joinpath(AerialVehicleControl.CONT_LIB_DOCS, "$(namePlot)_FOF_states.png"))
    end
end

if 2 in showPlot
    using Plots
    plotattd = plot(sol, vars = (0,30), color=:black, linewidth=2, xaxis="Time (t)",yaxis="Attitude error",    label=L"\Psi(\hat{R}(t), R(t))")
    plot!(          sol, vars = (0,31), color=:blue,  linewidth=2, xaxis="Time (t)",yaxis="Attitude error",    label=L"\Gamma(R_r(t), R(t))")
    plotlyap    = plot(sol, vars = (0,29), color=:black, linewidth=2, xaxis="Time (t)",yaxis="Lyapunov function", label=L"\mathcal{V}(t)")
    plotlyaplog = plot(sol.t, log10.(abs.(sol[29,:])), color=:black, linewidth=2, xaxis="Time (t)",yaxis="Lyapunov function", label=L"log_{10}(\mathcal{V}(t))")
    plot(plotattd, plotlyap, plotlyaplog, layout=(3,1), size=(1000,500))
    gui()
    if savePlot
        savefig(joinpath(AerialVehicleControl.CONT_LIB_DOCS, "$(namePlot)_FOF_errors.png"))
    end
end


## test_plot.jl : test plotting commands

@testset "Random field plotting    " begin

    ## 1D ##
    m = Matern(0.75,1,p=1.)
    cov = CovarianceFunction(1,m)
    pts = 0:0.01:1
    grf1 = GaussianRandomField(cov,KarhunenLoeve(500),pts)
    plot(grf1)
    plot_eigenvalues(grf1)
    plot_eigenfunction(grf1, 3)
    plot_eigenfunction(grf1, 25)

    ## 2D ##
    e = Exponential(0.01)
    cov = CovarianceFunction(2,e)
    pts = 0:0.05:1
    grf2 = GaussianRandomField(cov,KarhunenLoeve(10),pts,pts)
    surface(grf2)
    heatmap(grf2)
    contourf(grf2)
    contour(grf2)
    plot_eigenvalues(grf2)
    plot_eigenfunction(grf1, 2)
    plot_eigenfunction(grf1, 10)

    ## Separable ##
    e1 = Exponential(0.1, p=1)
    sp1 = SeparableCovarianceFunction(e1)
    pts = 0:0.01:1
    sg1 = GaussianRandomField(sp1,KarhunenLoeve(10),pts)
    plot_eigenvalues(sg1)
    plot_eigenfunction(sg1, 1)
    plot_eigenfunction(sg1, 4)
    plot_eigenfunction(sg1, 7)

    e2 = Exponential(0.01, p=1)
    sp2 = SeparableCovarianceFunction(e1,e2)
    sg2 = GaussianRandomField(sp2,KarhunenLoeve(10),pts,pts)
    plot_eigenvalues(sg2)
    plot_eigenfunction(sg2, 2)
    plot_eigenfunction(sg2, 5)
    plot_eigenfunction(sg2, 8)

end


#using Revise
using IsopycnalSurfaces, Test

@testset "IsopycnalSurfaces.jl" begin
    #############################
    # Test dedup! function
    @testset "deduplication" begin
        na = 10
        for ndupes = 1:na-1
            println(ndupes)
            a = sort(randn(na))
            b = randn(na)
            # make duplicates at end
            counter = 0
            while counter < ndupes
                a[end-counter-1] = a[end-counter]
                counter += 1
            end
            dedup!(a,b)
            println(a)
            println(b)
            println(@test issorted(a))
        end
    end

    @testset "column" begin
        ################################
        # Idealized mapping onto sigma1
        pz = collect(0.:500.:4000.) # pressure levels
        nz = length(pz)

        θz = collect(range(20,stop=10,length=nz))
        Sz = collect(range(36,stop=35,length=nz))
        ztest = sort(rand(2:8,2))
        σ₁true = sigma1column(θz[ztest],Sz[ztest],pz[ztest])

        σ₁grid = range(minimum(σ₁true),stop=maximum(σ₁true),length=20)
        σ₁=sigma1column(θz,Sz,pz)

        @testset "column_spline" begin
            sgood = findall(minimum(σ₁) .<= σ₁grid .<= maximum(σ₁))
            pσ = var2sigmacolumn(σ₁,pz,σ₁grid[sgood],splorder=3)
            @test isapprox(pσ[begin],pz[ztest[begin]])
            @test isapprox(pσ[end],pz[ztest[end]])

        end

        @testset "column_linear" begin
            sgood = findall(minimum(σ₁) .<= σ₁grid .<= maximum(σ₁))
            pσ = var2sigmacolumn(σ₁,pz,σ₁grid[sgood],linearinterp=true)

            @test isapprox(pσ[begin],pz[ztest[begin]])
            @test isapprox(pσ[end],pz[ztest[end]])
        end

        @testset "EOS" begin
            # just check that this runs
            # would be better to show some similarity in results

            println("Using the EOS from MITgcmTools.jl:")
            display(sigma2column(θz,Sz,pz,"JMD95")) # using the EOS from MITgcmTools.jl
            println("")
            println("Default, using the EOS from PhysOcean.jl/EOS80.jl:")
            display(sigma2column(θz,Sz,pz))  # default, using the EOS from PhysOcean.jl/EOS80.jl

        end # EOS

        # test input from a 3D array
        @testset "3d_array" begin
            eoses = ("EOS80","JMD95")
            for eos in eoses
                ztest = sort(rand(2:8,2))
                σ₁true = sigma1column(θz[ztest],Sz[ztest],pz[ztest],eos)

                # put θ,S,p into 3D array
                nx = 10; ny = 10;
                θ = Array{Float64,3}(undef,nx,ny,nz)
                S = Array{Float64,3}(undef,nx,ny,nz)
                p = Array{Float64,3}(undef,nx,ny,nz)

                [θ[i,j,k] = θz[k] for i = 1:nx for j = 1:ny for k = 1:nz ]
                [S[i,j,k] = Sz[k] for i = 1:nx for j = 1:ny for k = 1:nz ]
                [p[i,j,k] = pz[k] for i = 1:nx for j = 1:ny for k = 1:nz ]

                σ₁grid = collect(range(minimum(σ₁true),stop=maximum(σ₁true),length=20))

                vars = Dict("θ" => θ, "Sp" => S)

                @testset "3d_array_spline" begin
                    p₀ = 1000
                    varsσ = vars2sigma1(vars,pz,σ₁grid,splorder=3,eos=eos)
                    xx = rand(1:nx); yy = rand(1:ny)
                    @test isapprox(varsσ["p"][xx,yy,begin],pz[ztest[begin]])
                    @test isapprox(varsσ["p"][xx,yy,end],pz[ztest[end]])
                end

                @testset "3d_array_linear" begin
                    #splorder = 3
                    varsσ = vars2sigma1(vars,pz,σ₁grid,linearinterp=true,eos=eos)
                    xx = rand(1:nx); yy = rand(1:ny)
                    @test isapprox(varsσ["p"][xx,yy,begin],pz[ztest[begin]])
                    @test isapprox(varsσ["p"][xx,yy,end],pz[ztest[end]])
                end
            end
        end
    end # column
end


# sumS and sumS2 are padded by 1 pixel at the start of each dimension, to avoid writting code to specifically deal with the first line and first column of the images.

# 2D SEARCH AREA SUM-TABLES.

function znccSumArrays!( search::AbstractArray{T,2}, sumS::A{U,2}, sumS2::A{U,2}
                       ) where {T<:Real, U<:Real}
    h, w = size(search);
    @inbounds for c in 2:w+1
        for r in 2:h+1
			 sumS[r,c] = search[h-r+2,w-c+2]   +  sumS[r-1,c] +  sumS[r,c-1] -  sumS[r-1,c-1];
            sumS2[r,c] = search[h-r+2,w-c+2]^2 + sumS2[r-1,c] + sumS2[r,c-1] - sumS2[r-1,c-1];
        end
    end
end

function znccSumValues( sumS::Array{T,2}, sumS2::Array{T,2},
                        row, col, ss, si ) where {T<:Real}

    pad = 1;
    row = row + pad;
    col = col + pad;
    r0, r1 = max( 1, row - si[1] ), min( ss[1] + pad, row );
    c0, c1 = max( 1, col - si[2] ), min( ss[2] + pad, col );

    opS  = 0.0
    opS += sumS[ r1, c1 ];
    opS -= sumS[ r1, c0 ];
    opS -= sumS[ r0, c1 ];
    opS += sumS[ r0, c0 ];

    opS2  = 0.0
    opS2 += sumS2[ r1, c1 ];
    opS2 -= sumS2[ r1, c0 ];
    opS2 -= sumS2[ r0, c1 ];
    opS2 += sumS2[ r0, c0 ];

    return opS, opS2
end

# 3D SEARCH VOLUME SUM-TABLES.

# optimized sumarea, reusing previous operations
function znccSumArrays!( SA::AbstractArray{U,3}, S::Array{T,3}, S2::Array{T,3}
                       ) where {T<:Real,U<:Real}

     h,  w,  d = size(SA);
    @inbounds for z in 2:d+1
        for c in 2:w+1
            topSum  = 0.0;
            topSum2 = 0.0;
            for r in 2:h+1
                SAVal  = SA[h-r+2,w-c+2,d-z+2]
                SAVal2 = SAVal*SAVal;

                 S[r,c,z] = SAVal  +  S[r,c-1,z] +  S[r,c,z-1] -  S[r,c-1,z-1] +  topSum;
                S2[r,c,z] = SAVal2 + S2[r,c-1,z] + S2[r,c,z-1] - S2[r,c-1,z-1] + topSum2;

                 topSum  += SAVal;
                topSum2  += SAVal2;
            end
        end
    end
end

function znccSumValues( sumS::Array{T,3}, sumS2::Array{T,3},
                        row, col, zet, ss, si ) where {T<:Real}
    pad = 1;
    row = row + pad;
    col = col + pad;
    zet = zet + pad;
    r0, r1 = max( 1, row - si[1] ), min( ss[1] + pad, row );
    c0, c1 = max( 1, col - si[2] ), min( ss[2] + pad, col );
    z0, z1 = max( 1, zet - si[3] ), min( ss[3] + pad, zet );

    opS  = 0.0;
    opS += sumS[ r1, c1, z1 ]
    opS -= sumS[ r0, c1, z1 ]
    opS -= sumS[ r1, c0, z1 ]
    opS -= sumS[ r1, c1, z0 ]
    opS += sumS[ r1, c0, z0 ]
    opS += sumS[ r0, c1, z0 ]
    opS += sumS[ r0, c0, z1 ]
    opS -= sumS[ r0, c0, z0 ]

    opS2  = 0.0;
    opS2 += sumS2[ r1, c1, z1 ]
    opS2 -= sumS2[ r0, c1, z1 ]
    opS2 -= sumS2[ r1, c0, z1 ]
    opS2 -= sumS2[ r1, c1, z0 ]
    opS2 += sumS2[ r1, c0, z0 ]
    opS2 += sumS2[ r0, c1, z0 ]
    opS2 += sumS2[ r0, c0, z1 ]
    opS2 -= sumS2[ r0, c0, z0 ]

    return opS, opS2
end


# Accelerated inner product

C = zeros(m,p)

for j=1:m
    for l=1:p
        C[j,l] = dot(conj(view(A,j,:)),view(B,:,l))
    end
end

Code/Numerical Analysis & Machine Learning/Numerical Linear Algebra/InnerProduct.jl


abstract type AbstractTransform end

struct Output{T}
    z::Vector{T}
    rz::Vector{T}
end

function display(o::Output)
    println("    z: Transformed Observations")
    display(o.z)
    println("\n    rz: Reduced Transformed Observations")
    display(o.rz)
    return nothing
end

function output(W, g, λ, y, p)
    z = (x -> transform(g, λ, x)).(y)
    @views rz = z[p+1:end] .- W' * z[1:p]
    return Output(z, rz)
end
function output(T, n, m)
    z = Matrix{T}(undef, n)
    rz = Matrix{T}(undef, m)
    return Output(z, rz)
end
function set!(o::Output, W, g, λ, y)
    p = length(o.z) - length(o.rz)
    o.z .= (x -> radial(g, λ, x)).(y)
    @views o.rz .= o.z[:, p+1:end] .- W * o.z[:, 1:p]
    return o
end

logjacobian(g::AbstractTransform, λ::Tuple, y) = sum(x -> log∂(g, λ, x), y)
transform(g::AbstractTransform, λ::Tuple, y) = transform(g, λ..., y)
log∂(g::AbstractTransform, λ::Tuple, y) = log∂(g, λ..., y)

struct Identity <: AbstractTransform end

transform(::Identity, y) = y
log∂(::Identity, _) = 0

struct BoxCox <: AbstractTransform end

transform(::BoxCox, λ::Number, y) = λ == 0 ? log(y) : expm1(log(y) * λ) / λ
log∂(::BoxCox, λ::Number, y) = (λ - 1) * log(y)

struct YeoJohnson <: AbstractTransform end

function transform(::YeoJohnson, λ::Number, y)
    λ′ = y >= 0 ? λ : 2 - λ
    return sign(y) * (λ′ == 0 ? log1p(abs(y)) : expm1(log1p(abs(y)) * λ′) / λ′)
end
log∂(::YeoJohnson, λ::Number, y) = sign(y) * (λ - 1) *  log(abs(y))

struct SumTanh <: AbstractTransform end

function transform(::SumTanh, α::Vector, β::Vector, c::Vector, y)
    return sum(α .* tanh.(β .* (y .+ c)))
end
function log∂(::SumTanh, α::Vector, β::Vector, c::Vector, y)
    return sum(α .* β .* sech.(β .* (y .+ c)) .^ 2)
end


mutable struct MipDiscretSolver
    inst::Instance
    loglevel::Int         # niveau de verbosité
    bestsol::Solution     # meilleure Solution rencontrée
    # Les attributs spécifiques au modèle
    model::Model  # Le modèle MIP
    x         # vecteur des variables d'atterrissage
    b         # vecteur des variables de précédence (before)
    cost      # variable du cout de la solution
    costs     # variables du cout de chaque avion
    # Le constructeur
    function MipDiscretSolver(inst::Instance)
        this=new()
        this.inst = inst
        this.loglevel = Args.get("level")
        # this.bestsol = Solution(inst) # plante dans le proto car timing non dispo !
        this.bestsol = Solution(inst, update=false) # pas besoin de résoudre le timing !
        this.model = new_model(mode=:mip, log_level=this.loglevel)
        return this
    end
end

"""Création du modèle discrétisé (plus lent, mais polyvalent)"""
function buildSeqataModel(sv::MipDiscretSolver)

    # nb d'avion, pour abréger l'écriture
    planes = sv.inst.planes
    n = sv.inst.nb_planes
    model = sv.model

    # Quelques raccourcis car utilisés un peu partout
    lbmin = lb_min(sv.inst)
    ubmax = ub_max(sv.inst)

    ##########################################################
    # Création des variables
    ##########################################################

    # y: indicateur des dates d'atterrissage
    # y=1 si l'avion i atterrit à la date t dans [0,T-1]
    # @variable(model, y[p in planes, 1:ubmax], Bin) ORI pour indice Int et Plane
    @variable(model, y[p in planes, lbmin:ubmax], Bin)

    # variable dérivée : date d'atterrissage effective de l'avion i
    @expression(model, x[p in planes], sum(t*y[p,t] for t in lbmin:ubmax))
    sv.x = x

    # Coût de chaque avion compte tenu de sa date d'atterrissage
    @expression(model, costs[p in planes],
            sum(get_cost(p,t)*y[p,t] for t in p.lb:p.ub))

    # Objectif Minimiser le coût de pénalité total
    @expression(model, total_cost,
                sum(get_cost(p,t)*y[p,t] for p in planes for t in lbmin:ubmax))

    sv.costs = costs
    sv.cost = total_cost

    ##########################################################
    # Définition de l'objectif
    ##########################################################

    @objective(model, Min, total_cost)

    ##########################################################
    # Création des contraintes
    ##########################################################

    # Contrainte sur les bornes de la date d'atterrissage
    # AMPL: C1 {i in 1..n} : E[i] <= x[i] <= L[i];
    @constraint(model, c1[p in planes], p.lb <= x[p] <= p.ub)

    # Chaque avion ne peut avoir qu'une seule date d'atterrissage.
    # AMPL: C2 {i in 1..n} : sum{t in T_SET} y[i,t] = 1;
    @constraint(model, c2[p in planes], sum(y[p,t] for t in lbmin:ubmax) == 1)

    # Deux dates d'atterrissage ne doivent pas être trop rapprochées
    # Principe :
    # - on prend toutes les paires d'avions possibles
    # - pour toutes les dates d'atterissage possibles de p1
    #   - on interdit a p2 d'atterrir dans la fenêtre trop près avant
    #
    for p1 in planes, p2 in planes
        if !(p1.id<p2.id) continue end # évite de traiter le cas symétrique
        # ici p1 peut être avant ou après p2
        for t_i in p1.lb:p1.ub
            t_lb_j = max(p2.lb, t_i - get_sep(sv.inst, p2, p1) + 1) # p2 avant p1
            t_ub_j = min(p2.ub, t_i + get_sep(sv.inst, p1, p2) - 1) # p2 après p1
            for t_j in t_lb_j:t_ub_j
                # Si besoin de nommé les contraintes : c3[p1.id,p2.id,t_i,t_j]
                @constraint(model, y[p1,t_i] + y[p2,t_j] <= 1)
            end
        end
    end
end

# Résoud le problème complet : calcul l'ordre, le timing (dates d'atterrissage)
# et le cout total de la solution optimale au problème SEQATA (coûts linéaires)
# Cette fonction crée le modèle PLNE, le résoud puis met l'objet solution à
# jour.
#
function solve(sv::MipDiscretSolver)
    ln2("BEGIN solve(MipDiscretSolver)")
    ln2( "="^60 )

    lg2("Construction du modèle ($(ms())) ... ")
    buildSeqataModel(sv)
    ln2("fait ($(ms())).")

    lg2("Lancement de la résolution ($(ms())) ... ")
    optimize!(sv.model)
    ln2("fait ($(ms())).")

    lg2("Test de validité de la solution ($(ms())) ... ")
    if  JuMP.termination_status(sv.model) != MOI.OPTIMAL
        print("ERREUR : pas de solution pour :\n    ")
        @show JuMP.termination_status(sv.model)
        println( to_s(sv.bestsol) )
        exit(1)
    end

    lg2("Exploitation des résultats (mise à jour de la solution)($(ms())) ... ")
    # Extraction des valeurs entières des variables d'atterrissage

    # Il reste maintenant à mettre à jour notre objet solution à partir du
    # résultat de la résolution MIP
    for (i, p) in enumerate(sv.inst.planes)
        sv.bestsol.planes[i] = p
        sv.bestsol.x[i] = round(Int, value(sv.x[p]))  # value() retourne Float64 !
        sv.bestsol.costs[i] = value(sv.costs[p])
    end
    sv.bestsol.cost = round(value(sv.cost), digits=Args.args[:cost_precision])

    # On trie juste la solution par date d'atterrissage croissante des avions
    # pour améliorer la présentation de la solution
    sort!(sv.bestsol)
    ln2("fait. ($(ms()))")

    ln2("END solve(MipDiscretSolver)")
end


#!/usr/bin/julia

# Copyright (C) 2021, 2022 Matheus Fernandes Bigolin <mfrdrbigolin@disroot.org>
# SPDX-License-Identifier: MIT

"Day Fifteen, Chiton."

include("Utils.jl")

# N.b.: external package, used for the PriorityQueue data structure.
using DataStructures

using DelimitedFiles
using .Utils

"Find the `cavern` path with the lowest risk; then, return this total risk."
function solve(cavern :: Array{Int8, 2}) :: Int
    sizepaths = size(cavern)

    unvisited = [(i,j) for j in 1:sizepaths[2] for i in 1:sizepaths[1]]

    distance = PriorityQueue(zip(unvisited, fill(Inf, length(cavern))))
    distance[(1,1)] = 0

    current = (1,1)

    # Dijkstra’s algorithm.
    while current != sizepaths
        i, j = current
        neighbors = Array{Tuple{Int64, Int64}}([])

        for idx in [(i - 1, j), (i + 1, j), (i, j - 1), (i, j + 1)]
            idxs = CartesianIndices(cavern)
            if CartesianIndex(idx) in idxs && idx in keys(distance)
                push!(neighbors, idx)
            end
        end

        for neighbor in neighbors
            rel_dist = cavern[neighbor...] + peek(distance)[2]
            if rel_dist < distance[neighbor]
                distance[neighbor] = rel_dist
            end
        end

        dequeue!(distance)

        current = peek(distance)[1]
    end

    return distance[sizepaths]
end

"Enlarge `cavern` to its actual dimension."
function enlarge(cavern :: Array{Int8, 2}) :: Array{Int8, 2}
    pattern = reshape([Int8(i + j) for i in 0:4 for j in 0:4], (5,5))

    enlarged_pattern = repeat(pattern, inner=size(cavern))
    repeated_cavern = repeat(cavern, 5, 5)

    return [Int8(mod1(i, 9)) for i in repeated_cavern + enlarged_pattern]
end

function main()
    usage_and_exit(length(ARGS) != 1)

    data = collect.(readdlm(first(ARGS), String))

    cavern = hcat([parse.(Int8, path) for path in data]...)
    large_cavern = enlarge(cavern)

    println(solve(cavern))
    println(solve(large_cavern))
end

if abspath(PROGRAM_FILE) == @__FILE__
    main()
end


using JuMP, GLPK, Ipopt
using DualDecomposition

const DD = DualDecomposition
const parallel = DD.parallel

"""
a: interest rate
π: unit stock price
ρ: unit dividend price

K: number of stages
L: number of stock types
2^L scenarios in each stage
2^L^(K-1)=16 scenarios in total
ρ = 0.05 * π
bank: interest rate 0.01
stock1: 1.03 or 0.97
stock2: 1.06 or 0.94
...

b_k: initial asset (if k=1) and income (else)
B_k: money in bank
x_{k,l}: number of stocks to buy/sell (integer)
y_{k,l}: total stocks

deterministic model:

    max     B_K+∑_{l=1}^{L}π_{K,l}y_{K,l}

    s.t.    B_1+∑_{l=1}^{L}π_{1,l}x_{1,l} = b_1

            b_k+(1+a)B_{k-1}+∑_{l=1}^{L}ρ_{k,l}y_{k-1,l} = B_k+∑_{l=1}^{L}π_{k,l}x_{k,l}, ∀ k=2,…,K

            y_{1,l} = x_{1,l}, ∀ l=1,…,L

            y_{k-1,l}+x_{k,l} = y_{k,l}, ∀ k=2,…,K, l=1,…,L

            x_{k,l} ∈ ℤ , ∀ k=1,…,K, l=1,…,L

            y_{k,l} ≥ 0, ∀ k=1,…,K, l=1,…,L

            B_k ≥ 0, ∀ k=1,…,K.
"""
const K = 3
const L = 2
const a = 0.01
const b_init = 100  # initial capital
const b_in = 30   # income

# iteratively add nodes
# root node
function create_nodes()::DD.Tree
    ξ = Dict{Symbol, Union{Float64,<:AbstractArray{Float64}}}(:π => ones(L))
    tree = DD.Tree(ξ)

    #subproblem formulation
    function subproblem_builder(tree::DD.Tree, subtree::DD.SubTree, node::DD.SubTreeNode)
        mdl = subtree.model
        x = @variable(mdl, x[l=1:L], Int, base_name="n1_x")
        #x = @variable(mdl, x[l=1:L], base_name="n1_x")

        y = @variable(mdl, y[l=1:L] >= 0, base_name="n1_y")
        DD.set_output_variable!(node, :y, y)

        B = @variable(mdl, B >= 0, base_name="n1_B")
        DD.set_output_variable!(node, :B, B)

        π = DD.get_scenario(node)[:π]
        @constraints(mdl,
            begin
                B + sum( π[l] * x[l] for l in 1:L) == b_init
                [l=1:L], y[l] - x[l] == 0
            end
        )
        DD.set_stage_objective(node, 0.0)

        JuMP.unregister(mdl, :x)
        JuMP.unregister(mdl, :y)
        JuMP.unregister(mdl, :B)
    end

    DD.set_stage_builder!(tree, 1, subproblem_builder)

    create_nodes!(tree, 1)
    return tree
end

# child nodes
function create_nodes!(tree::DD.Tree, pt::Int)
    for scenario = 1:2^L
        prob = 1/2^L
        π = get_realization(DD.get_scenario(tree, pt)[:π], scenario)
        ξ = Dict{Symbol, Union{Float64,<:AbstractArray{Float64}}}(:π => π)
        id = DD.add_child!(tree, pt, ξ, prob)

        #subproblem formulation
        function subproblem_builder(tree::DD.Tree, subtree::DD.SubTree, node::DD.SubTreeNode)
            mdl = subtree.model
            id = DD.get_id(node)
            #x = @variable(mdl, x[l=1:L], Int, base_name = "n$(id)_x")
            x = @variable(mdl, x[l=1:L], base_name = "n$(id)_x")

            y = @variable(mdl, y[l=1:L] >= 0, base_name = "n$(id)_y")
            DD.set_output_variable!(node, :y, y)

            B = @variable(mdl, B >= 0, base_name = "n$(id)_B")
            DD.set_output_variable!(node, :B, B)

            y_ = @variable(mdl, y_[l=1:L] >= 0, base_name = "n$(id)_y_")
            DD.set_input_variable!(node, :y, y_)

            B_ = @variable(mdl, B_ >= 0, base_name = "n$(id)_B_")
            DD.set_input_variable!(node, :B, B_)

            π = DD.get_scenario(node)[:π]
            pt = DD.get_parent(node)
            ρ = DD.get_scenario(tree, pt)[:π] * 0.05
            @constraint(mdl, B + sum( π[l] * x[l] - ρ[l] * y_[l] for l in 1:L) - (1+a) * B_ == b_in)
            @constraint(mdl, [l=1:L], y[l] - x[l] - y_[l] == 0)

            #dummy bound for input variables to avoid subproblem unboundedness
            @constraint(mdl, [l=1:L], y_[l] <= 500)
            @constraint(mdl, B_ <= 500)
            if DD.get_stage(node) < K
                DD.set_stage_objective(node, 0.0)
            else
                DD.set_stage_objective(node, -(B + sum( π[l] * y[l] for l in 1:L )))
            end
            JuMP.unregister(mdl, :x)
            JuMP.unregister(mdl, :y)
            JuMP.unregister(mdl, :B)
            JuMP.unregister(mdl, :y_)
            JuMP.unregister(mdl, :B_)
        end

        DD.set_stage_builder!(tree, id, subproblem_builder)
        if DD.get_stage(tree, id) < K
            create_nodes!(tree, id)
        end
    end
end

# construct realization event
function get_realization(ξ::Array{Float64,1}, scenario::Int)::Array{Float64,1}
    ret = ones(L)
    multipliers = digits(scenario - 1, base=2, pad=L)*2 - ones(L)
    for l = 1:L
        ret[l] = ξ[l] * (1 + multipliers[l] * l * 0.03)
    end
    return ret
end

tree = create_nodes()
#node_cluster = DD.decomposition_not(tree)
node_cluster = DD.decomposition_scenario(tree)
#node_cluster = DD.decomposition_temporal(tree) #There is a DUAL_INFEASIBLE issue

# Number of block components
NS = length(node_cluster)

# Initialize MPI
parallel.init()

# Create DualDecomposition instance.
algo = DD.LagrangeDual()

# partition scenarios into processes
parallel.partition(NS)

coupling_variables = Vector{DD.CouplingVariableRef}()
models = Dict{Int,JuMP.Model}()

for block_id in parallel.getpartition()
    nodes = node_cluster[block_id]
    subtree = DD.create_subtree!(tree, block_id, coupling_variables, nodes)
    set_optimizer(subtree.model, GLPK.Optimizer)
    DD.add_block_model!(algo, block_id, subtree.model)
    models[block_id] = subtree.model
end

# Set nonanticipativity variables as an array of symbols.
DD.set_coupling_variables!(algo, coupling_variables)

# Lagrange master method
LM = DD.BundleMaster(BM.ProximalMethod, optimizer_with_attributes(Ipopt.Optimizer, "print_level" => 0))
#LM = DD.BundleMaster(BM.TrustRegionMethod, GLPK.Optimizer)

# Solve the problem with the solver; this solver is for the underlying bundle method.
DD.run!(algo, LM)

# Write timing outputs to files
DD.write_all(algo)

# Finalize MPI
parallel.finalize()


module LocalPolynomialFitting

export lpparams, lpfit,
       optbandn,
       epanwin

"""Epanechnikov's window"""
epanwin(n::Integer) = 3/4 * (1 - linspace(-1,1,n+2).^2)[2:end-1]

include("lanshammar.jl")

end


include("../src/RandomField.jl")
using .RandomField

# chose input parameters:
n 		= 1000		# number of grid points
n_real	= 10000		# number of realizations
l 		= 30.0		# real space box is [-l,l]
f 		= gaus_rf	# spectral density

var_y(ν) = 2*sqrt(2pi)*(l/pi)*f(ν)

# create structures (allocate memory)
Y = Vector{ComplexF64}(undef,n)	# fourier realization
Z = Vector{ComplexF64}(undef,n)	# real realization

M_est = zeros(ComplexF64,n) # est. mean
C_est = zeros(ComplexF64,n)	# est. stationary covariance
F_est = zeros(ComplexF64,n)	# est. spectral density
C_tru = Vector{ComplexF64}(undef,n) # true stationary covariance
F_tru = Vector{ComplexF64}(undef,n)	# true spectral density

T, V = create_grids(n,l)

# draw realization
draw_realization!(Y,var_y,V)
ift!(Z,Y,n,l)

# zeroth order moment
for i=1:n_real
	draw_realization!(Y,var_y,V)
	ift!(Z,Y,n,l)
	M_est .+= Z
end
M_est ./= n_real

# first order moments
for i=1:n_real
	draw_realization!(Y,var_y,V)
	ift!(Z,Y,n,l)
	@. F_est += sqrt(1/8pi)*(pi/l) * Y * conj(Y)
	@. C_est += 0.5 * Z * conj(Z[div(n,2)])
end
F_est ./= n_real
C_est ./= n_real

@. F_tru = f(V)
ift!(C_tru,F_tru,n,l)

# print to file:
out = open("../data/output.csv", "w")
print(out, "T, V, ReY, ReZ, ImZ, M_est, C_est, C_tru, F_est, F_tru\n")
for i in 1:n
	print(out,
		"$(T[i]),"*
		"$(V[i]),"*
		"$(real(Y[i])),"*
		"$(real(Z[i])),"*
		"$(imag(Z[i])),"*
		"$(real(M_est[i])),"*
		"$(real(C_est[i])),"*
		"$(real(C_tru[i])),"*
		"$(real(F_est[i])),"*
		"$(real(F_tru[i]))\n")
end
close(out)


using ITensors
using ITensorsGrad
using OptimKit
using Random
using Zygote

using ITensors: data

Random.seed!(1243)

N = 20
χ = 20

s = siteinds("S=1/2", N; conserve_qns=true)
os = OpSum()
for n in 1:(N - 1)
  os .+= 0.5, "S+", n, "S-", n + 1
  os .+= 0.5, "S-", n, "S+", n + 1
  os .+= "Sz", n, "Sz", n + 1
end
Hmpo = MPO(os, s)
ψ₀mps = randomMPS(s, n -> isodd(n) ? "↑" : "↓"; linkdims=χ)

H = data(Hmpo)
ψ₀ = data(ψ₀mps)

# The Rayleigh quotient to minimize
function E(H::Vector{ITensor}, ψ::Vector{ITensor})
  N = length(ψ)
  ψdag = dag.(addtags.(ψ, "bra"; tags="Link"))
  ψ′dag = prime.(ψdag)
  e = ITensor(1.0)
  for n in 1:N
    e = e * ψ′dag[n] * H[n] * ψ[n]
  end
  norm = ITensor(1.0)
  for n in 1:N
    norm = norm * ψdag[n] * ψ[n]
  end
  return e[] / norm[]
end

@show sum(flux, ψ₀)
@show E(H, ψ₀)

E(ψ) = E(H, ψ)
∇E(ψ) = E'(ψ)
fg(ψ) = (E(ψ), ∇E(ψ))

linesearch = HagerZhangLineSearch(; c₁=0.1, c₂=0.9, ϵ=1e-6, θ=1 / 2, γ=2 / 3, ρ=5.0)
algorithm = LBFGS(5; maxiter=100, gradtol=1e-8, linesearch=linesearch, verbosity=2)
ψ, fψ, gψ, numfg, normgradhistory = optimize(fg, ψ₀, algorithm)

sweeps = Sweeps(5)
setmaxdim!(sweeps, χ)
fψ_dmrg, ψ_dmrg = dmrg(Hmpo, ψ₀mps, sweeps)

time_Eψ = @elapsed E(ψ)
time_∇Eψ = @elapsed E'(ψ)
@show time_∇Eψ / time_Eψ;
@show fψ_dmrg
@show E(H, ψ)


path = "docs/src/assets/"

using ISA, Plots
ψ₀ = AMFMcomp(t->exp(-t^2),t->200.0,0.0)
ψ₁ = AMFMcomp(t->1.0,t->100*t+50,0.1)
ψ₂ = AMFMcomp(t->1.0+0.8*cos(2t),t->100 + 7.5*sin(t),π)
fs = 16_000
t = 0:1/fs:3
𝚿₀ = numComp( ψ₀(t), fs )
𝚿₁ = numComp( ψ₁(t), fs )
𝚿₂ = numComp( ψ₂(t), fs )
𝐳 = numModel([𝚿₀,𝚿₁,𝚿₂])
plot(𝐳)
png(path*"IS_num_exAMFMmodel1.png")


using  Statistics,StatsBase,Distributions,DataFrames,CSV,Plots,Match,Zygote
pyplot()

# Helper functions:
include("routines.jl")

## TEST SET CONSTRUCTION AND ANALYSIS
df=CSV.read("JigsawPred.csv",DataFrame;delim=",",header=false)

dfm=Matrix(df)
Ntest=size(df,1)
c_pred=zeros(Int,Ntest,6) # multi-label problem with 6 labels
for n=1:Ntest
    for k = 1:6
        if dfm[n,k]>0.5
            c_pred[n,k]=1
        end
    end
end
p1_pred = 1.0*dfm # for this classifier, the predictions are essentially deterministic

df2=CSV.read("JigsawTrue.csv",DataFrame;delim=",",header=false)
dfm=Matrix(df2)
c_true=zeros(Int,Ntest,6) # multi-label problem with 6 labels
for n=1:Ntest
    for k = 1:6
        if dfm[n,k]>0.5
            c_true[n,k]=1
        end
    end
end

F1Test_micro_true,F1Test_macro_true,accuracy_individual,accuracy_joint = MultiMetrics(c_pred,c_true)

# Bernoulli Sampling approximation:
lambda=0.9
p1_true_guess = lambda*p1_pred .+ (1-lambda)*0.5 # guess of the true label on test set

# Define the parameters for the Optimal Sampling Distribution:
htilde,dummy = happroxFn_multi_micro(c_pred,p1_true_guess)

Mvals=LinRange(10,Ntest,10) # expected number of samples

percent=zeros(length(Mvals))

acc_mean=zeros(length(Mvals))
acc_lower=zeros(length(Mvals))
acc_upper=zeros(length(Mvals))

f1_micro_mean=zeros(length(Mvals))
f1_micro_lower=zeros(length(Mvals))
f1_micro_upper=zeros(length(Mvals))

f1_macro_mean=zeros(length(Mvals))
f1_macro_lower=zeros(length(Mvals))
f1_macro_upper=zeros(length(Mvals))

for experiment=1:length(Mvals)

    global M = Mvals[experiment] # average number of elements in each test set replicate
    global b = optimalb(htilde,M)
    global s = (rand(Ntest) .< b)
    global total_samples_used = sum(s)

    ToLabel=findall(s.==1)

    F1hat_micro_mean,F1hat_micro_var=FhatBS_micro(c_pred[ToLabel,:],c_true[ToLabel,:],b[ToLabel])
    F1hat_macro_mean,F1hat_macro_var=FhatBS_macro(c_pred[ToLabel,:],c_true[ToLabel,:],b[ToLabel])

    # convert this to confidence bounds by fitting this to say a beta distribution:

    global confidence_level=0.9

    println("-----------------------------------------")
    println("total samples used  = $total_samples_used")
    println("% of test data      = $(total_samples_used/Ntest)")
    percent[experiment]=total_samples_used/Ntest

    println("")

    println("true F1             = $(F1Test_micro_true)")
    local mn,lower,upper = confidenceBeta(F1hat_micro_mean,F1hat_micro_var,confidence_level)
    println("estimated F1        = $mn [$lower,$upper]")

    f1_micro_mean[experiment]=F1hat_micro_mean
    f1_micro_lower[experiment]=lower
    f1_micro_upper[experiment]=upper

    println("")

    println("true F1             = $(F1Test_macro_true)")
    local mn,lower,upper = confidenceBeta(F1hat_macro_mean,F1hat_macro_var,confidence_level)
    println("estimated F1        = $mn [$lower,$upper]")

    f1_macro_mean[experiment]=F1hat_macro_mean
    f1_macro_lower[experiment]=lower
    f1_macro_upper[experiment]=upper

end

mp=maximum(percent)

f1_true = F1Test_micro_true
ribbon_upper=f1_micro_upper - f1_micro_mean
ribbon_lower=f1_micro_mean - f1_micro_lower
plot(reuse=false,show=true)
plot_f1_micro=plot!(percent,f1_micro_mean,ribbon=(ribbon_lower,ribbon_upper),label="$(confidence_level) confidence interval",framestyle = :box,minorgrid=true)
plot!(plot_f1_micro)
plot!([0,mp],[f1_true,f1_true],label="true value",legend=:bottomright)#,ylims=[0.5*f1_true,1.5*f1_true])
title!("F1 Micro")
xlabel!("fraction of test data labelled")

f1_true = F1Test_macro_true
ribbon_upper=f1_macro_upper - f1_macro_mean
ribbon_lower=f1_macro_mean - f1_macro_lower
plot(reuse=false,show=true)
plot_f1_macro=plot!(percent,f1_macro_mean,ribbon=(ribbon_lower,ribbon_upper),label="$(confidence_level) confidence interval",framestyle = :box,minorgrid=true)
plot!(plot_f1_macro)
plot!([0,mp],[f1_true,f1_true],label="true value",legend=:bottomright)#,ylims=[0.5*f1_true,1.5*f1_true])
title!("F1 Macro")
xlabel!("fraction of test data labelled")


using GeoPhyInv
using LinearAlgebra
using Random

mgrid = [range(0.0, stop=10.,step=0.1), range(0.0, stop=30.,step=0.2)];
ageom=AGeom(mgrid, SSrcs(10), Srcs(10), Recs(10));

tgrid=range(0, stop=1.0, step=0.1);

srcwav=SrcWav(tgrid, ageom, [:P]);

Random.randn!(srcwav[3][:P]);

x=randn(length(tgrid));
update!(srcwav, [:P,], x);

x1=randn(length(tgrid));
x2=randn(length(tgrid));
update!(srcwav[1], [:P,], x1);
update!(srcwav[2], [:P,], x2);

rmul!(srcwav, 2.0);


module Net

export sctenifoldnet, tenrnet, manialn, drgenes, pcnet, tensordecomp, tensordecomp2

const NCOMP1,NCOMP2=3,5
const NLAYERS,NCELLS=10,500

using Statistics, LinearAlgebra, Distributions, Random
using MultipleTesting, Random, SparseArrays, TSVD, TensorToolbox
# import KrylovKit

include("support.jl")
include("pcnet.jl")                   # include the contents of other files in the module
include("tenrnet.jl")
include("tensordecomp.jl")
include("manialn.jl")
include("drgenes.jl")

function sctenifoldnet(X::AbstractMatrix{T}, Y::AbstractMatrix{T}; donorm::Bool=true) where T<:Real
    Z0,_=tenrnet(X,donorm=donorm)
    Z1,_=tenrnet(Y,donorm=donorm)
    Z0=0.5*(Z0+Z0')
    Z1=0.5*(Z1+Z1')
    d=manialn(Z0,Z1)
    fc,p,adjp=drgenes(d)
    return d,fc,p,adjp
end

end


using RayTracer
using Parameters

@with_kw mutable struct Particle{Dim}
    x⃗::Point{Dim}
    v⃗::Vec{Dim}
    m::Float = 1.0
end
@with_kw struct ExternalForces{Dim}
    F⃗::Vec{Dim}
    α::Float = 0.0
end

function step!(particle::Particle,forces::ExternalForces;dt = 1.0)
    @unpack x⃗,v⃗,m = particle
    @unpack F⃗,α = forces
    particle.x⃗ = x⃗ + dt .* v⃗
    particle.v⃗ = v⃗ + dt .* (F⃗/m - α .* v⃗)
    return particle.x⃗
end

using Plots

plt = plot(
    1,
    xlim = (-3, 3),
    ylim = (0, 3),
    title = "Universe",
    marker = 2,
    xlabel = "x",
    ylabel = "y",
)
projectile = Particle(x⃗ =  Point(-3, 0),v⃗ =  Vec(2, 3))
forces = ExternalForces(F⃗ = Vec(0, -1.0), α = .3)

# build an animated gif by pushing new points to the plot, saving every 2th frame
@gif for i=1:50
    push!(plt, (projectile.x⃗.coords)...)
    step!(projectile,forces,dt = 0.1)
end every 2


# Objects in the game and their supporting functions (update!, collide!, ...)

# Abstract positions and dimensions in arbitrary coordinate systems for the
# game. Concrete Pos and Dims must always be used with the same CoordType.
# Can convert between them (possibly in reference to the Camera).
abstract type AbstractPos{CoordType} end
abstract type AbstractDims{CoordType} end

""" WorldCoords are the absolute space of the game, independent of camera. """
struct WorldCoords end

"""
    WorldPos(5.0,-200.0)
x,y float coordinates in the game world (not necessarily the same as pixel
coordinates on the screen).
"""
struct WorldPos <: AbstractPos{WorldCoords}  # 0,0 == middle
    x::Float64
    y::Float64
end
toWorldPos(p::WorldPos, c) = p
"""
    WorldDims(5.0,-200.0)
w,h float dimensions in the game world (not necessarily the same as pixel
coordinates on the screen).
"""
struct WorldDims <: AbstractDims{WorldCoords}  # 0,0 == middle
    w::Float64
    h::Float64
end
toWorldDims(d::WorldDims, c) = d

"""
    Vector2D(-2.5,1.0)
x,y vector representing direction in the game world. Could represent a velocity,
a distance, etc. Subtracting two `WorldPos`itions results in a `Vector2D`.
"""
struct Vector2D   # TODO: consider making this abstract and parameterized on CoordType.
    x::Float64
    y::Float64
end
import Base.*, Base./, Base.-, Base.+
+(a::Vector2D, b::Vector2D) = Vector2D(a.x+b.x, a.y+b.y)
-(a::Vector2D, b::Vector2D) = Vector2D(a.x-b.x, a.y-b.y)
*(a::Vector2D, x::Number) = Vector2D(a.x*x, a.y*x)
*(x::Number, a::Vector2D) = a*x
/(a::Vector2D, x::Number) = Vector2D(a.x/x, a.y/x)
+(a::P, b::Vector2D) where {P<:AbstractPos} = P(a.x+b.x, a.y+b.y)
-(a::P, b::Vector2D) where {P<:AbstractPos} = P(a.x-b.x, a.y-b.y)
+(a::Vector2D, b::P) where {P<:AbstractPos} = P(a.x+b.x, a.y+b.y)
-(a::Vector2D, b::P) where {P<:AbstractPos} = P(a.x-b.x, a.y-b.y)
-(a::AbstractPos, b::AbstractPos) = Vector2D(a.x-b.x, a.y-b.y)
-(x::AbstractPos) = AbstractPos(-x.x, -x.y)
-(x::Vector2D) = Vector2D(-x.x, -x.y)

magSqrd(v::Vector2D) = v.x^2 + v.y^2
magnitude(v::Vector2D) = sqrt(magSqrd(v))
unitVec(v::Vector2D) = v / magnitude(v)


function render!(fbuffer, hbuffer, cmap, hmap::Matrix{T}, pos::Tuple{T,T,T}, θ::T) where T<:Real
    # draw background (light blue sky)
    fill!(fbuffer, RGB24(0.529f0, 0.808f0, 0.980f0))
    fill!(hbuffer, WIN_HEIGHT)
    # precompute some constants
    w = Base.multiplicativeinverse(MAP_WIDTH%UInt)
    h = Base.multiplicativeinverse(MAP_HEIGHT%UInt)
    # compute orientation angle θ sine and cosine
    sθ, cθ = sincos(θ)
    # draw from front to back (low z coordinate to high z coordinate)
    @inbounds for z in 1 : VIEW_DIST
        # find view-line on map, this calculation corresponds to a field of view of 90°
        px = pos[1] - (cθ+sθ)*z; py = pos[2] - (cθ-sθ)*z
        # segment the line in equal increments
        dx = 2z*cθ/WIN_WIDTH; dy = -2z*sθ/WIN_WIDTH
        # compute local z-scale factor
        scale = ZSCALE/z
        # raster horizontal view-line and draw the vertical for each segment
        for j in 1 : WIN_WIDTH
            # compute map indices from px & py
            pxi = UInt(1) + rem(unsafe_trunc(Int, px)%UInt, w)
            pyi = UInt(1) + rem(unsafe_trunc(Int, py)%UInt, h)
            # compute the on-screen height of the feature
            height = unsafe_trunc(Int, muladd(pos[3]-hmap[pxi,pyi], scale, HORIZON))
            # assign color to column in framebuffer
            color = cmap[pxi,pyi]
            for i in max(height,1) : hbuffer[j]
                fbuffer[j,i] = color
            end
            # keep track of feature height
            hbuffer[j] = min(height, hbuffer[j])
            # move to next view-line feature position
            px += dx; py += dy
        end
    end
end

# "fisheye camera"
#ϕ = range(-π/4,π/4,length=WIN_WIDTH)
# compute position of feature
#if z > VIEW_DIST÷2
#    sθ, cθ = sincos(θ-ϕ[j])
#    px = pos[1] - sθ*z; py = pos[2] - cθ*z
#end


module HierarchicalGeometry

# Write your package code here.
using StaticArrays
using LinearAlgebra
using Polyhedra
using LazySets
# using QHull

"""
    cross_product_operator(x)

Computes `Ωₓ` such that `Ωₓy = x ⨯ y`
"""
cross_product_operator(x) = SMatrix{3,3}(
    [0.0    -x[3]   x[2];
     x[3]   0.0     -x[1];
     -x[2]  x[1]    0.0]
)

abstract type OverapproxModel end

export
    RegularPolyhedronOverapprox

"""
    RegularPolyhedronOverapprox{D,N}

Used to overrapproximate convex sets with a pre-determined set of support
vectors.
"""
struct RegularPolyhedronOverapprox{D,N} <: OverapproxModel
    support_vectors::NTuple{N,SVector{D,Float64}}
end
get_support_vecs(model::RegularPolyhedronOverapprox) = [v for v in model.support_vectors]

"""
    RegularPolyhedronOverapprox(dim::Int,N::Int,epsilon=0.1)

Construct a regular polyhedron overapproximation model by specifying the number
of dimensions and the number of support vectors to be arranged radially about
the axis formed by the unit vector along each dimension.
epsilon if the distance between support vector v and an existing support vector
is less than epsilon, the new vector will not be added.
"""
function RegularPolyhedronOverapprox(dim::Int,N::Int,epsilon=0.1)
    vecs = Vector{SVector{dim,Float64}}()
    # v - the initial support vector for a given dimension
    v = zeros(dim)
    v[end] = 1.0
    for i in 1:dim
        d = zeros(dim)
        d[i] = 1.0
        A = cross_product_operator(d)
        for j in 1:N
            v = normalize(exp(A*j*2*pi/N)*v)
            add = true
            for vp in vecs
                if norm(v-vp) < epsilon
                    add = false
                    break
                end
            end
            if add
                @show v
                push!(vecs,v)
            end
        end
        v = d
    end
    RegularPolyhedronOverapprox(tuple(vecs...))
end

export
    equatorial_overrapprox_model

"""
    equatorial_overapprox_model(lat_angles=[-π/4,0.0,π/4],lon_angles=collect(0:π/4:2π),epsilon=0.1)

Returns a RegularPolyhedronOverapprox model generated with one face for each
combination of pitch and yaw angles specified by lat_angles and lon_angles,
respectively. There are also two faces at the two poles
"""
function equatorial_overapprox_model(lat_angles=[-π/4,0.0,π/4],lon_angles=collect(0:π/4:2π),epsilon=0.1)
    vecs = Vector{SVector{3,Float64}}()
    push!(vecs,[0.0,0.0,1.0])
    push!(vecs,[0.0,0.0,-1.0])
    for phi in lat_angles
        for theta in lon_angles
            v = normalize([
                cos(phi)*cos(theta),
                cos(phi)*sin(theta),
                sin(phi)
            ])
            add = true
            for vp in vecs
                if norm(v-vp) < epsilon
                    add = false
                    break
                end
            end
            if add
                push!(vecs,v)
            end
        end
    end
    RegularPolyhedronOverapprox(tuple(vecs...))
end

function LazySets.overapproximate(lazy_set,model::RegularPolyhedronOverapprox{N},epsilon::Float64=0.1) where {D,N}
    halfspaces = map(v->LazySets.HalfSpace(v,ρ(v,lazy_set)),get_support_vecs(model))
    sort!(halfspaces; by=h->ρ(h.a, lazy_set))
    poly = HPolyhedron()
    while !isempty(halfspaces) #&& !isbounded(poly)
        poly = intersection(poly,halfspaces[1])
        deleteat!(halfspaces,1)
    end
    @assert isbounded(poly)
    hpoly = convert(HPolytope,poly)
    # vpoly = tovrep(hpoly)
    # continue adding new halfspace constraints until the distance "clipped off"
    # by the new halfspace constraint is less than epsilon
    # while !isempty(halfspaces)
    #     for v in vertices_list(hpoly)
    #
    #     end
    # end
    hpoly
end

# """
#     sort_facet_vtxs
#
# Returns a counterclockwise ordering of 'vtx_ids' around the face defined by
# the halfspace associated with h_idx. It is assumed that the vertices form a
# convex polygon
# """
# function sort_facet_vtxs(poly,h_idx,vtx_idxs)
#     if !isempty(vtx_idxs)
#         vtxs = map(v->get(poly,v), vtx_idxs)
#         h = get(poly,h_idx) # halfspace
#         ctr = sum(vtxs)
#         # axes
#         a1 = normalize(vtxs[1]-ctr)
#         a2 = normalize(cross(h.a,a1))
#         # angles = map(v->atan(dot(a2,v),dot(a1,v)))
#         return sort(vtx_idxs;by=v->atan(dot(a2,get(poly,v)),dot(a1,get(poly,v))))
#     end
#     return vtx_idxs
# end

# """
#     polytope_vtxs_and_edges
#
# Identify the edges of each facet of a polytope
# """
# function polytope_vtxs_and_edges(hpolytope::AbstractPolyhedron)
#     poly = polyhedron(hpolytope)
#     vtxs = collect(points(poly))
#     edges = Vector{Tuple{Int,Int}}()
#     for idx in eachindex(halfspaces(poly))
#         vtx_idxs = incidentpointindices(poly,idx)
#         sorted_vtx_ids = sort_facet_vtxs(poly,idx,vtx_idxs)
#         for i in 1:length(sorted_vtx_ids)-1
#             v1 = sorted_vtx_ids[i]
#             v2 = sorted_vtx_ids[i+1]
#             push!(edges, (v1.value, v2.value))
#         end
#         push!(edges,(sorted_vtx_ids[end].value,sorted_vtx_ids[1].value))
#     end
#     return vtxs, edges
# end
#
# function write_vertices(filename,vtxs::Vector{Vector{Float64}})
#     open(filename, "w") do io
#         for v in vtxs
#             for c in v
#                 print(io,c," ")
#             end
#             print(io,"\n")
#         end
#     end
# end
#
# function write_edges(filename,edge_list::Vector{Tuple{Int,Int}})
#     open(filename, "w") do io
#         for e in edge_list
#             for v in e
#                 print(io,v," ")
#             end
#             print(io,"\n")
#         end
#     end
# end
#
# function write_vtxs_and_edges(filename,poly)
#     vtx_file = join([filename,".nodes"])
#     edge_file = join([filename,".edges"])
#     vtxs, edge_list = polytope_vtxs_and_edges(poly)
#     write_vertices(vtx_file,vtxs)
#     write_edges(edge_file,edge_list)
# end

end


#验证所推导的求傅里叶系数公式正确性
#比较variance大小
using LinearAlgebra
using StatsBase
using Plots
using Notifier

function pt(arg...)
    flush(stdout)
    println(arg...)
end
pt()

M=4
m=9
p=normalize!(rand(M),1)#态对应的概率幅list
function A(x)
    sum(p .* cos.(Array(1:M) .* 2x))/2+0.5
end
xk=Array(1:m-1) .* pi ./ m
function est2(i)
    4/m*(1+sum(A.(xk) .* cos.(2i .* xk)))
end
function on(a,b)
    a%b==0 ? b-1 : -1
end
function est3(i)
    #跟est2本质一样，但舍弃est2中的浮点运算
    oo=0
    for j=1:M
        oo+=p[j]*(on(i+j,m)+on(i-j,m))
    end
    4/m*(1+1/2*on(i,m)+1/4*oo)
end
for i=1:M
    pt()
    est2(i)|>pt
    p[i]|>pt
    est3(i)|>pt
end
function D1(N)
    o=0
    for i=1:M
        o+=16/m^2/N*sum(A.(xk) .* (1 .- A.(xk)) .* cos.(2i .* xk).^2)
    end
    o
end
function D2(N)
    o=0
    for i=1:M
        o+=p[i]*(1-p[i])/N/(m-1)
    end
    o
end
pt("variance")
D1(3)|>pt
D2(3)|>pt
# plot(D1.(1:100))
# plot!(D2.(1:100))
# for i=1:m,j=1:m
#     pt(i," ",j)
#     sum(cos.(2i .* Array(1:m-1) .* pi ./ m) .* cos.(2j .* Array(1:m-1) .* pi ./ m))|>pt
# end


using Test
using Statistics
using VectorizedStatistics

using VectorizationBase
#VectorizationBase.vsum(x::Float64) = x

@testset "Vreducibles" begin include("testVreducibles.jl") end
@testset "ArrayStats" begin include("testArrayStats.jl") end
@testset "Correlation and Covariance" begin include("testCovCor.jl") end

@testset "Sorting" begin include("testSorting.jl") end


using DelimitedFiles

function solve(input)
    raw = readdlm(input,'\n')
    coords = vcat(map(parseCoords,raw)...)
    best_time = 0
    smaller_height = height(coords,0)
    for time in 1:10460
        current_height = height(coords,time)
        if current_height < smaller_height
            smaller_height=current_height
            best_time = time
        end
    end
    @show best_time
    current = copy(coords[:,1:2])+best_time*coords[:,3:4]
    x_min,x_max = extrema(current[:,1])
    y_min, y_max = extrema(current[:,2])
    w = x_max-x_min+1
    h = y_max-y_min+1
    grid = fill(' ',(h,w))

    for point in 1:size(current)[1]
        grid[current[point,:][2]-y_min+1,current[point,:][1]-x_min+1]='#'
    end

    for line in 1:size(grid)[1]
        println(join(grid[line,:]))
    end

end

function parseCoords(s)
    m = match(r"position\=\<(?<x>.+),(?<y>.+)\> velocity\=\<(?<vx>.+),(?<vy>.+)\>", s)
    x = tryparse(Int64,m[:x])
    y = tryparse(Int64,m[:y])
    vx = tryparse(Int64,m[:vx])
    vy = tryparse(Int64,m[:vy])
    return [x y vx vy]
end

function height(coords,time)
    ymin,ymax = extrema(coords[:,2] + coords[:,4]*time)
    return ymax-ymin
end

solve("input.txt")


"""
    replica_fidelity(df::DataFrame; p_field = :hproj, skip = 0)
Compute the fidelity of the average coefficient vector and the projector defined in
`p_field` from the result of replica [`lomc!()`](@ref) passed as  argument `df`,
using replicas `_1` and `_2`.
Calls [`ratio_of_means()`](@ref) to perform a blocking analysis
on a ratio of the means of separate time series and returns a
[`RatioBlockingResult`](@ref).
The first `skip` steps in the time series are skipped.

The fidelity of states `|ψ⟩` and `|ϕ⟩` is defined as
```math
F(ψ,ϕ) = \\frac{|⟨ψ|ϕ⟩|^2}{⟨ψ|ψ⟩⟨ϕ|ϕ⟩} .
```
Specifically, `replica_fidelity` computes
```math
F(\\mathbf{v},⟨\\mathbf{c}⟩) =
    \\frac{⟨(\\mathbf{c}_1⋅\\mathbf{v})(\\mathbf{v}⋅\\mathbf{c}_1)⟩}
    {⟨\\mathbf{c}_1⋅\\mathbf{c}_1⟩} ,
```
where `v` is the projector specified by `p_field`, which is assumed to be normalised to
unity with the two-norm (i.e. `v⋅v == 1`), and ``\\mathbf{c}_1`` and ``\\mathbf{c}_2``
are two replica coefficient vectors.
"""
function replica_fidelity(df::DataFrame; p_field = :hproj, skip = 0,  args...)
    p_field_1 = Symbol(p_field, :_1)
    p_field_2 = Symbol(p_field, :_2)
    fid_num = conj(getproperty(df, p_field_1)) .* getproperty(df, p_field_2)
    fid_num = fid_num[skip+1:end]
    # denominator
    fid_den = df.c1_dot_c2[skip+1:end]

    return ratio_of_means(fid_num, fid_den; args...)
end