# Filter prototypes, transformations, and transforms

using ..Windows

abstract type FilterType end

#
# Butterworth prototype
#

function Butterworth(::Type{T}, n::Integer) where {T<:Real}
    n > 0 || error("n must be positive")

    poles = zeros(Complex{T}, n)
    for i = 1:div(n, 2)
        w = convert(T, 2i-1)/2n
        pole = complex(-sinpi(w), cospi(w))
        poles[2i-1] = pole
        poles[2i] = conj(pole)
    end
    if isodd(n)
        poles[end] = -1
    end
    ZeroPoleGain{:s}(Complex{T}[], poles, one(T))
end

"""
    Butterworth(n)

``n`` pole Butterworth filter.
"""
Butterworth(n::Integer) = Butterworth(Float64, n)

#
# Chebyshev type I and II prototypes
#

function chebyshev_poles(::Type{T}, n::Integer, ε::Real) where {T<:Real}
    p = zeros(Complex{T}, n)
    μ = asinh(convert(T, 1)/ε)/n
    b = -sinh(μ)
    c = cosh(μ)
    for i = 1:div(n, 2)
        w = convert(T, 2i-1)/2n
        pole = complex(b*sinpi(w), c*cospi(w))
        p[2i-1] = pole
        p[2i] = conj(pole)
    end
    if isodd(n)
        w = convert(T, 2*div(n, 2)+1)/2n
        pole = b*sinpi(w)
        p[end] = pole
    end
    p
end

function Chebyshev1(::Type{T}, n::Integer, ripple::Real) where {T<:Real}
    n > 0 || error("n must be positive")
    ripple >= 0 || error("ripple must be non-negative")

    ε = sqrt(10^(convert(T, ripple)/10)-1)
    p = chebyshev_poles(T, n, ε)
    k = one(T)
    for i = 1:div(n, 2)
        k *= abs2(p[2i])
    end
    if iseven(n)
        k /= sqrt(1+abs2(ε))
    else
        k *= real(-p[end])
    end
    ZeroPoleGain{:s}(Complex{T}[], p, k)
end

"""
    Chebyshev1(n, ripple)

`n` pole Chebyshev type I filter with `ripple` dB ripple in
the passband.
"""
Chebyshev1(n::Integer, ripple::Real) = Chebyshev1(Float64, n, ripple)

function Chebyshev2(::Type{T}, n::Integer, ripple::Real) where {T<:Real}
    n > 0 || error("n must be positive")
    ripple >= 0 || error("ripple must be non-negative")

    ε = 1/sqrt(10^(convert(T, ripple)/10)-1)
    p = chebyshev_poles(T, n, ε)
    for i = 1:length(p)
        p[i] = inv(p[i])
    end

    z = zeros(Complex{T}, n-isodd(n))
    k = one(T)
    for i = 1:div(n, 2)
        w = convert(T, 2i-1)/2n
        ze = complex(zero(T), -inv(cospi(w)))
        z[2i-1] = ze
        z[2i] = conj(ze)
        k *= abs2(p[2i])/abs2(ze)
    end
    isodd(n) && (k *= -real(p[end]))

    ZeroPoleGain{:s}(z, p, k)
end

"""
    Chebyshev2(n, ripple)

`n` pole Chebyshev type II filter with `ripple` dB ripple in
the stopband.
"""
Chebyshev2(n::Integer, ripple::Real) = Chebyshev2(Float64, n, ripple)

#
# Elliptic prototype
#
# See Orfanidis, S. J. (2007). Lecture notes on elliptic filter design.
# Retrieved from http://www.ece.rutgers.edu/~orfanidi/ece521/notes.pdf

# Compute Landen sequence for evaluation of elliptic functions
function landen(k::Real)
    niter = 7
    kn = Vector{typeof(k)}(undef, niter)
    # Eq. (50)
    for i = 1:niter
        kn[i] = k = abs2(k/(1+sqrt(1-abs2(k))))
    end
    kn
end

# cde computes cd(u*K(k), k)
# sne computes sn(u*K(k), k)
# Both accept the Landen sequence as generated by landen above
for (fn, init) in ((:cde, :(cospi(u/2))), (:sne, :(sinpi(u/2))))
    @eval begin
        function $fn(u::Number, landen::Vector{T}) where T<:Real
            winv = inv($init)
            # Eq. (55)
            for i = length(landen):-1:1
                oldwinv = winv
                winv = 1/(1+landen[i])*(winv+landen[i]/winv)
            end
            w = inv(winv)
        end
    end
end

# sne inverse
function asne(w::Number, k::Real)
    oldw = NaN
    while w != oldw
        oldw = w
        kold = k
        # Eq. (50)
        k = abs2(k/(1+sqrt(1-abs2(k))))
        # Eq. (56)
        w = 2*w/((1+k)*(1+sqrt(1-abs2(kold)*w^2)))
    end
    2*asin(w)/π
end

function Elliptic(::Type{T}, n::Integer, rp::Real, rs::Real) where {T<:Real}
    n > 0 || error("n must be positive")
    rp > 0 || error("rp must be positive")
    rp < rs || error("rp must be less than rs")

    # Eq. (2)
    εp = sqrt(10^(convert(T, rp)/10)-1)
    εs = sqrt(10^(convert(T, rs)/10)-1)

    # Eq. (3)
    k1 = εp/εs
    k1 >= 1 && error("filter order is too high for parameters")

    # Eq. (20)
    k1′² = 1 - abs2(k1)
    k1′ = sqrt(k1′²)
    k1′_landen = landen(k1′)

    # Eq. (47)
    k′ = one(T)
    for i = 1:div(n, 2)
        k′ *= sne(convert(T, 2i-1)/n, k1′_landen)
    end
    k′ = k1′²^(convert(T, n)/2)*k′^4

    k = sqrt(1 - abs2(k′))
    k_landen = landen(k)

    # Eq. (65)
    v0 = -im/convert(T, n)*asne(im/εp, k1)

    z = Vector{Complex{T}}(undef, 2*div(n, 2))
    p = Vector{Complex{T}}(undef, n)
    gain = one(T)
    for i = 1:div(n, 2)
        # Eq. (43)
        w = convert(T, 2i-1)/n

        # Eq. (62)
        ze = complex(zero(T), -inv(k*cde(w, k_landen)))
        z[2i-1] = ze
        z[2i] = conj(ze)

        # Eq. (64)
        pole = im*cde(w - im*v0, k_landen)
        p[2i] = pole
        p[2i-1] = conj(pole)

        gain *= abs2(pole)/abs2(ze)
    end

    if isodd(n)
        pole = im*sne(im*v0, k_landen)
        p[end] = pole
        gain *= abs(pole)
    else
        gain *= 10^(-convert(T, rp)/20)
    end

    ZeroPoleGain{:s}(z, p, gain)
end

"""
    Elliptic(n, rp, rs)

`n` pole elliptic (Cauer) filter with `rp` dB ripple in the
passband and `rs` dB attentuation in the stopband.
"""
Elliptic(n::Integer, rp::Real, rs::Real) = Elliptic(Float64, n, rp, rs)

#
# Prototype transformation types
#

# returns frequency in half-cycles per sample ∈ (0, 1)
function normalize_freq(w::Real, fs::Real)
    w <= 0 && error("frequencies must be positive")
    f = 2*w/fs
    f >= 1 && error("frequencies must be less than the Nyquist frequency $(fs/2)")
    f
end

struct Lowpass{T} <: FilterType
    w::T
end

"""
    Lowpass(Wn[; fs])

Low pass filter with cutoff frequency `Wn`. If `fs` is not
specified, `Wn` is interpreted as a normalized frequency in
half-cycles/sample.
"""
Lowpass(w::Real; fs::Real=2) = Lowpass{typeof(w/1)}(normalize_freq(w, fs))

struct Highpass{T} <: FilterType
    w::T
end

"""
    Highpass(Wn[; fs])

High pass filter with cutoff frequency `Wn`. If `fs` is not
specified, `Wn` is interpreted as a normalized frequency in
half-cycles/sample.
"""
Highpass(w::Real; fs::Real=2) = Highpass{typeof(w/1)}(normalize_freq(w, fs))

struct Bandpass{T} <: FilterType
    w1::T
    w2::T
end

"""
    Bandpass(Wn1, Wn2[; fs])

Band pass filter with normalized pass band (`Wn1`, `Wn2`). If
`fs` is not specified, `Wn1` and `Wn2` are interpreted as
normalized frequencies in half-cycles/sample.
"""
function Bandpass(w1::Real, w2::Real; fs::Real=2)
    w1 < w2 || error("w1 must be less than w2")
    Bandpass{Base.promote_typeof(w1/1, w2/1)}(normalize_freq(w1, fs), normalize_freq(w2, fs))
end

struct Bandstop{T} <: FilterType
    w1::T
    w2::T
end

"""
    Bandstop(Wn1, Wn2[; fs])

Band stop filter with normalized stop band (`Wn1`, `Wn2`). If
`fs` is not specified, `Wn1` and `Wn2` are interpreted as
normalized frequencies in half-cycles/sample.
"""
function Bandstop(w1::Real, w2::Real; fs::Real=2)
    w1 < w2 || error("w1 must be less than w2")
    Bandstop{Base.promote_typeof(w1/1, w2/1)}(normalize_freq(w1, fs), normalize_freq(w2, fs))
end

#
# Prototype transformation implementations
#
# The formulas implemented here come from the documentation for the
# corresponding functionality in Octave, available at
# https://staff.ti.bfh.ch/sha1/Octave/index/f/sftrans.html
# The Octave implementation was not consulted in creating this code.

# Create a lowpass filter from a lowpass filter prototype
transform_prototype(ftype::Lowpass, proto::ZeroPoleGain{:s, Z, P, K}) where {Z, P, K} =
    ZeroPoleGain{:s, Z, P, K}(ftype.w * proto.z, ftype.w * proto.p,
              proto.k * ftype.w^(length(proto.p)-length(proto.z)))

# Create a highpass filter from a lowpass filter prototype
function transform_prototype(ftype::Highpass, proto::ZeroPoleGain{:s})
    z = proto.z
    p = proto.p
    k = proto.k
    nz = length(z)
    np = length(p)
    TR = Base.promote_eltype(z, p)
    newz = zeros(TR, max(nz, np))
    newp = zeros(TR, max(nz, np))

    num = one(eltype(z))
    for i = 1:nz
        num *= -z[i]
        newz[i] = ftype.w / z[i]
    end

    den = one(eltype(p))
    for i = 1:np
        den *= -p[i]
        newp[i] = ftype.w / p[i]
    end

    abs(real(num) - 1) < np*eps(real(num)) && (num = 1)
    abs(real(den) - 1) < np*eps(real(den)) && (den = 1)
    ZeroPoleGain{:s}(newz, newp, oftype(k, k * real(num)/real(den)))
end

# Create a bandpass filter from a lowpass filter prototype
function transform_prototype(ftype::Bandpass, proto::ZeroPoleGain{:s})
    z = proto.z
    p = proto.p
    k = proto.k
    nz = length(z)
    np = length(p)
    ncommon = min(nz, np)
    TR = Base.promote_eltype(z, p)
    newz = zeros(TR, 2*nz+np-ncommon)
    newp = zeros(TR, 2*np+nz-ncommon)
    for (oldc, newc) in ((p, newp), (z, newz))
        for i = 1:length(oldc)
            b = oldc[i] * ((ftype.w2 - ftype.w1)/2)
            pm = sqrt(b^2 - ftype.w2 * ftype.w1)
            newc[2i-1] = b + pm
            newc[2i] = b - pm
        end
    end
    ZeroPoleGain{:s}(newz, newp, oftype(k, k * (ftype.w2 - ftype.w1) ^ (np - nz)))
end

# Create a bandstop filter from a lowpass filter prototype
function transform_prototype(ftype::Bandstop, proto::ZeroPoleGain{:s})
    z = proto.z
    p = proto.p
    k = proto.k
    nz = length(z)
    np = length(p)
    npairs = nz+np-min(nz, np)
    TR = Base.promote_eltype(z, p)
    newz = Vector{TR}(undef, 2*npairs)
    newp = Vector{TR}(undef, 2*npairs)

    num = one(eltype(z))
    for i = 1:nz
        num *= -z[i]
        b = (ftype.w2 - ftype.w1)/2/z[i]
        pm = sqrt(b^2 - ftype.w2 * ftype.w1)
        newz[2i-1] = b - pm
        newz[2i] = b + pm
    end

    den = one(eltype(p))
    for i = 1:np
        den *= -p[i]
        b = (ftype.w2 - ftype.w1)/2/p[i]
        pm = sqrt(b^2 - ftype.w2 * ftype.w1)
        newp[2i-1] = b - pm
        newp[2i] = b + pm
    end

    # Any emaining poles/zeros are real and not cancelled
    npm = sqrt(-complex(ftype.w2 * ftype.w1))
    for (n, newc) in ((np, newp), (nz, newz))
        for i = n+1:npairs
            newc[2i-1] = -npm
            newc[2i] = npm
        end
    end

    abs(real(num) - 1) < np*eps(real(num)) && (num = 1)
    abs(real(den) - 1) < np*eps(real(den)) && (den = 1)
    ZeroPoleGain{:s}(newz, newp, oftype(k, k * real(num)/real(den)))
end

transform_prototype(ftype, proto::FilterCoefficients{:s}) =
    transform_prototype(ftype, convert(ZeroPoleGain, proto))

"""
    analogfilter(responsetype, designmethod)

Construct an analog filter. See below for possible response and
filter types.
"""
analogfilter(ftype::FilterType, proto::FilterCoefficients) =
    transform_prototype(ftype, proto)

# Bilinear transform
bilinear(f::FilterCoefficients{:s}, fs::Real) = bilinear(convert(ZeroPoleGain, f), fs)
function bilinear(f::ZeroPoleGain{:s,Z,P,K}, fs::Real) where {Z,P,K}
    ztype = typeof(0 + zero(Z)/fs)
    z = fill(convert(ztype, -1), max(length(f.p), length(f.z)))

    ptype = typeof(0 + zero(P)/fs)
    p = Vector{typeof(zero(P)/fs)}(undef, length(f.p))

    num = one(one(fs) - one(Z))
    for i = 1:length(f.z)
        z[i] = (2 + f.z[i] / fs)/(2 - f.z[i] / fs)
        num *= (2 * fs - f.z[i])
    end

    den = one(one(fs) - one(P))
    for i = 1:length(f.p)
        p[i] = (2 + f.p[i] / fs)/(2 - f.p[i]/fs)
        den *= (2 * fs - f.p[i])
    end

    ZeroPoleGain{:z}(z, p, f.k * real(num)/real(den))
end

# Pre-warp filter frequencies for digital filtering
prewarp(ftype::Union{Lowpass, Highpass}) = (typeof(ftype))(prewarp(ftype.w))
prewarp(ftype::Union{Bandpass, Bandstop}) = (typeof(ftype))(prewarp(ftype.w1), prewarp(ftype.w2))
# freq in half-samples per cycle
prewarp(f::Real) = 4*tan(pi*f/2)

# Digital filter design
"""
    digitalfilter(responsetype, designmethod)

Construct a digital filter. See below for possible response and
filter types.
"""
digitalfilter(ftype::FilterType, proto::FilterCoefficients) =
    bilinear(transform_prototype(prewarp(ftype), proto), 2)

#
# Special filter types
#

"""
    iirnotch(Wn, bandwidth[; fs])

Second-order digital IIR notch filter [^Orfandis] at frequency `Wn` with
bandwidth `bandwidth`. If `fs` is not specified, `Wn` is
interpreted as a normalized frequency in half-cycles/sample.

[^Orfandis]: Orfanidis, S. J. (1996). Introduction to signal processing. Englewood Cliffs, N.J: Prentice Hall, p. 370.
"""
function iirnotch(w::Real, bandwidth::Real; fs=2)
    w = normalize_freq(w, fs)
    bandwidth = normalize_freq(bandwidth, fs)

    # Eq. 8.2.23
    b = 1/(1+tan(pi*bandwidth/2))
    # Eq. 8.2.22
    cosw0 = cospi(w)
    Biquad(b, -2b*cosw0, b, -2b*cosw0, 2b-1)
end

#
# FIR filter design
#

# Get length and alpha for Kaiser window filter with specified
# transition band width and stopband attenuation in dB
function kaiserord(transitionwidth::Real, attenuation::Real=60)
    n = ceil(Int, (attenuation - 7.95)/(π*2.285*transitionwidth))+1

    if attenuation > 50
        β = 0.1102*(attenuation - 8.7)
    elseif attenuation >= 21
        β = 0.5842*(attenuation - 21)^0.4 + 0.07886*(attenuation - 21)
    else
        β = 0.0
    end

    return n, β/π
end

struct FIRWindow{T}
    window::Vector{T}
    scale::Bool
end

"""
    FIRWindow(window; scale=true)

FIR filter design using window `window`, a vector whose length
matches the number of taps in the resulting filter.

If `scale` is `true` (default), the designed FIR filter is
scaled so that the following holds:

- For [`Lowpass`](@ref) and [`Bandstop`](@ref) filters, the frequency
  response is unity at 0 (DC).
- For [`Highpass`](@ref) filters, the frequency response is unity
  at the Nyquist frequency.
- For [`Bandpass`](@ref) filters, the frequency response is unity
  in the center of the passband.
"""
FIRWindow(window::Vector; scale::Bool=true) = FIRWindow(window, scale)

# FIRWindow(n::Integer, window::Function, args...) = FIRWindow(window(n, args...))
"""
    FIRWindow(; transitionwidth, attenuation=60, scale=true)

Kaiser window FIR filter design. The required number of taps is
calculated based on `transitionwidth` (in half-cycles/sample)
and stopband `attenuation` (in dB). `attenuation` defaults to
60 dB.
"""
FIRWindow(; transitionwidth::Real=throw(ArgumentError("must specify transitionwidth")),
          attenuation::Real=60, scale::Bool=true) =
    FIRWindow(kaiser(kaiserord(transitionwidth, attenuation)...), scale)

# Compute coefficients for FIR prototype with specified order
function firprototype(n::Integer, ftype::Lowpass)
    w = ftype.w

    [w*sinc(w*(k-(n-1)/2)) for k = 0:(n-1)]
end

function firprototype(n::Integer, ftype::Bandpass)
    w1 = ftype.w1
    w2 = ftype.w2

    [w2*sinc(w2*(k-(n-1)/2)) - w1*sinc(w1*(k-(n-1)/2)) for k = 0:(n-1)]
end

function firprototype(n::Integer, ftype::Highpass)
    w = ftype.w
    isodd(n) || throw(ArgumentError("FIRWindow highpass filters must have an odd number of coefficients"))

    out = [-w*sinc(w*(k-(n-1)/2)) for k = 0:(n-1)]
    out[div(n, 2)+1] += 1
    out
end

function firprototype(n::Integer, ftype::Bandstop)
    w1 = ftype.w1
    w2 = ftype.w2
    isodd(n) || throw(ArgumentError("FIRWindow bandstop filters must have an odd number of coefficients"))

    out = [w1*sinc(w1*(k-(n-1)/2)) - w2*sinc(w2*(k-(n-1)/2)) for k = 0:(n-1)]
    out[div(n, 2)+1] += 1
    out
end

scalefactor(coefs::Vector, ::Union{Lowpass, Bandstop}) = sum(coefs)
function scalefactor(coefs::Vector, ::Highpass)
    c = zero(coefs[1])
    for k = 1:length(coefs)
        c += ifelse(isodd(k), coefs[k], -coefs[k])
    end
    c
end
function scalefactor(coefs::Vector, ftype::Bandpass)
    n = length(coefs)
    freq = middle(ftype.w1, ftype.w2)
    c = zero(coefs[1])
    for k = 0:n-1
        c += coefs[k+1]*cospi(freq*(k-(n-1)/2))
    end
    c
end

function digitalfilter(ftype::FilterType, proto::FIRWindow)
    coefs = firprototype(length(proto.window), ftype)
    @assert length(proto.window) == length(coefs)
    out = coefs .* proto.window
    proto.scale ? rmul!(out, 1/scalefactor(out, ftype)) : out
end

# Compute FIR coefficients necessary for arbitrary rate resampling
function resample_filter(rate::AbstractFloat, Nϕ = 32, rel_bw = 1.0, attenuation = 60)
    f_nyq       = rate >= 1.0 ? 1.0/Nϕ : rate/Nϕ
    cutoff      = f_nyq * rel_bw
    trans_width = cutoff * 0.2

    # Determine resampling filter order
    hLen, α = kaiserord(trans_width, attenuation)

    # Round the number of taps up to a multiple of Nϕ.
    # Otherwise the missing taps will be filled with 0.
    hLen = Nϕ * ceil(Int, hLen/Nϕ)

    # Ensure that the filter is an odd length
    if (iseven(hLen))
        hLen += 1
    end

    # Design filter
    h = digitalfilter(Lowpass(cutoff), FIRWindow(kaiser(hLen, α)))
    rmul!(h, Nϕ)
end

# Compute FIR coefficients necessary for rational rate resampling
function resample_filter(rate::Union{Integer,Rational}, rel_bw = 1.0, attenuation = 60)
    Nϕ          = numerator(rate)
    decimation  = denominator(rate)
    f_nyq       = min(1/Nϕ, 1/decimation)
    cutoff      = f_nyq * rel_bw
    trans_width = cutoff * 0.2

    # Determine resampling filter order
    hLen, α = kaiserord(trans_width, attenuation)

    # Round the number of taps up to a multiple of Nϕ (same as interpolation factor).
    # Otherwise the missing taps will be filled with 0.
    hLen = Nϕ * ceil(Int, hLen/Nϕ)

    # Ensure that the filter is an odd length
    if (iseven(hLen))
        hLen += 1
    end

    # Design filter
    h = digitalfilter(Lowpass(cutoff), FIRWindow(kaiser(hLen, α)))
    rmul!(h, Nϕ)
end


push!(LOAD_PATH, "C:\\Users\\Zhang\\Documents\\GitHub\\pcreo\\src")

using LinearAlgebra: I, ×
using Random: shuffle!
using UUIDs

using MultiFloats
using NearestNeighbors: KDTree, knn
using ProgressMeter: @showprogress

using DZLinearAlgebra
using RieszEnergy

use_clean_multifloat_arithmetic()
setprecision(64 * 3)

function split_into_runs(items::AbstractArray{T,N};
                         by=identity, threshold=zero(T)) where {T,N}
    runs = Vector{T}[]
    @inbounds if !isempty(items)
        current_run = T[items[1]]
        current_score = by(items[1])
        for i = 2 : length(items)
            item = items[i]
            next_score = by(item)
            if abs(next_score - current_score) <= threshold
                push!(current_run, item)
            else
                push!(runs, current_run)
                current_run = T[item]
            end
            current_score = next_score
        end
        push!(runs, current_run)
    end
    return runs
end

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

const NUM_BITS = 90
const EPSILON = inv(Float64x3(2)^NUM_BITS)

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

function candidate_isometries(p1, p2, q1, q2)
    return [
        hcat(q1, q2, q1 × q2) / hcat(p1, p2, p1 × p2),
        hcat(q1, q2, q1 × q2) / hcat(p1, p2, p2 × p1),
        hcat(q1, q2, q1 × q2) / hcat(p2, p1, p1 × p2),
        hcat(q1, q2, q1 × q2) / hcat(p2, p1, p2 × p1),
    ]
end

function riesz_energy_distances!(
        result::AbstractVector{Tuple{T,Int,Int}},
        points::AbstractMatrix{T}) where {T<:Real}
    dim, num_points = size(points)
    energy = zero(T)
    k = 0
    @inbounds for i = 1 : num_points
        for j = i+1 : num_points
            norm_sq = zero(T)
            @simd ivdep for d = 1 : dim
                temp = points[d,i] - points[d,j]
                norm_sq += temp * temp
            end
            k += 1
            dist = renormalize(unsafe_sqrt(norm_sq))
            energy += inv(dist)
            result[k] = (dist, i, j)
        end
    end
    return energy
end

function isocmp(x::Tuple{T,Int,Int}, y::Tuple{T,Int,Int}) where {T<:Real}
    @inbounds return MultiFloats._lt(x[1], y[1])
end

function isisomorphic(p::Matrix{T}, q::Matrix{T};
                      verbose::Bool=false) where {T <: Real}
    if size(p) != size(q)
        if verbose
            println("DIFFERENT SIZES")
        end
        return false
    end
    dim, num_points = size(p)
    num_pairs = div(num_points * (num_points - 1), 2)
    threshold = EPSILON * dim * num_points
    p_dist = Vector{Tuple{T,Int,Int}}(undef, num_pairs)
    q_dist = Vector{Tuple{T,Int,Int}}(undef, num_pairs)
    Ep = riesz_energy_distances!(p_dist, p)
    Eq = riesz_energy_distances!(q_dist, q)
    if !(abs((Ep - Eq) / (Ep + Eq)) <= threshold)
        if verbose
            println("DIFFERENT ENERGIES: ", abs((Ep - Eq) / (Ep + Eq)))
        end
        return false
    end
    sort!(p_dist; lt=isocmp)
    sort!(q_dist; lt=isocmp)
    @inbounds for i = 1 : num_pairs
        if !(abs(p_dist[i][1] - q_dist[i][1]) <= threshold)
            if verbose
                println("DIFFERENT DISTANCES: ",
                        abs(p_dist[i][1] - q_dist[i][1]))
            end
            return false
        end
    end
    p_runs = split_into_runs(p_dist; by=x->x[1], threshold=threshold)
    q_runs = split_into_runs(q_dist; by=x->x[1], threshold=threshold)
    filter!(run -> !(abs(run[end][1] - 2) <= threshold), p_runs)
    filter!(run -> !(abs(run[end][1] - 2) <= threshold), q_runs)
    if length.(p_runs) != length.(q_runs)
        if verbose
            println("DIFFERENT BARCODES")
        end
        return false
    end
    _, run_idx = findmin(length.(p_runs))
    _, pi, pj = p_runs[run_idx][1]
    for (_, qi, qj) in q_runs[run_idx]
        for R in candidate_isometries(p[:,pi], p[:,pj], q[:,qi], q[:,qj])
            @assert maximum(abs.(R' * R - I)) < threshold
            p_tree = KDTree(R * p)
            idxs, dists = knn(p_tree, q, 1)
            idxs = vcat(idxs...)
            dists = vcat(dists...)
            @assert length(idxs) == num_points
            @assert length(dists) == num_points
            if issetequal(idxs, 1:num_points) && (maximum(dists) < threshold)
                return true
            end
        end
    end
    if verbose
        println("DID NOT FIND ISOMETRY")
    end
    return false
end

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

const SOURCE_DIR = "C:/Users/Zhang/PCREO/PCREO"
const DATABASE_DIR = "C:/PCREO-DB"
const DATABASE_CACHE = Dict{Int,Vector{Tuple{Float64x3,String}}}()

function build_database_cache(num_points::Int)
    @assert !haskey(DATABASE_CACHE, num_points)
    println("Building database cache for size: ", num_points)
    identifier = "PCREO-" * lpad(num_points, 8, '0')
    cache_entry = Tuple{Float64x3,String}[]
    for foldername in readdir(DATABASE_DIR)
        if occursin(identifier, foldername)
            folderpath = joinpath(DATABASE_DIR, foldername)
            @assert isdir(folderpath)
            exemplarpath = joinpath(DATABASE_DIR,
                foldername, first(readdir(folderpath)))
            @assert isfile(exemplarpath)
            exemplar = Float64x3.(BigFloat.(readlines(exemplarpath)))
            @assert length(exemplar) % 3 == 0
            num_points = div(length(exemplar), 3)
            exemplar = reshape(exemplar, (3, num_points))
            @assert norm(riesz_gradient(exemplar)) <= EPSILON * length(exemplar)
            energy = riesz_energy(exemplar)
            push!(cache_entry, (energy, folderpath))
        end
    end
    DATABASE_CACHE[num_points] = cache_entry
end

function find_in_database(points::Matrix{Float64x3})
    dim, num_points = size(points)
    threshold = EPSILON * dim * num_points
    Ep = riesz_energy(points)
    if !haskey(DATABASE_CACHE, num_points)
        build_database_cache(num_points)
    end
    @assert haskey(DATABASE_CACHE, num_points)
    for (Eq, folderpath) in DATABASE_CACHE[num_points]
        if abs((Ep - Eq) / (Ep + Eq)) <= threshold
            @assert isdir(folderpath)
            exemplarpath = joinpath(folderpath, first(readdir(folderpath)))
            @assert isfile(exemplarpath)
            exemplar = Float64x3.(BigFloat.(readlines(exemplarpath)))
            @assert length(exemplar) % 3 == 0
            exemplar = reshape(exemplar, (3, div(length(exemplar), 3)))
            if isisomorphic(points, exemplar)
                return folderpath
            end
        end
    end
    return nothing
end

function add_to_database_entry(sourcepath::String, folderpath::String)
    mv(sourcepath, joinpath(folderpath, basename(sourcepath)))
end

function create_database_entry(points::Matrix{Float64x3}, sourcepath::String)
    dim, num_points = size(points)
    @assert haskey(DATABASE_CACHE, num_points)
    energy = riesz_energy(points)
    foldername = "PCREO-$(lpad(num_points, 8, '0'))-$(uppercase(string(uuid4())))"
    folderpath = joinpath(DATABASE_DIR, foldername)
    push!(DATABASE_CACHE[num_points], (energy, folderpath))
    mkdir(folderpath)
    add_to_database_entry(sourcepath, folderpath)
end

function main()
    for filename in readdir(SOURCE_DIR)
        path = joinpath(SOURCE_DIR, filename)
        if isfile(path)
            println("Reading: ", path)
            lines = readlines(path)
            @assert length(lines) % 3 == 0
            num_points = div(length(lines), 3)
            points = reshape(Float64x3.(BigFloat.(lines)), (3, num_points))
            if norm(riesz_gradient(points)) <= EPSILON * length(points)
                folderpath = find_in_database(points)
                if isnothing(folderpath)
                    println("Creating new database entry.")
                    create_database_entry(points, path)
                else
                    println("Matched existing database entry: ", folderpath)
                    add_to_database_entry(path, folderpath)
                end
            else
                println("WARNING: Gradient norm exceeds threshold.")
            end
        end
    end
end

# function readpoints(path)
#     lines = readlines(path)
#     @assert length(lines) % 3 == 0
#     reshape(Float64x3.(BigFloat.(lines)),
#             (3, div(length(lines), 3)))
# end

# function main()
#     for foldername in readdir(DATABASE_DIR)
#         println("Verifying directory: ", foldername)
#         folderpath = joinpath(DATABASE_DIR, foldername)
#         @assert isdir(folderpath)
#         exemplarpaths = [joinpath(folderpath, exemplarname)
#                          for exemplarname in readdir(folderpath)]
#         points = readpoints.(exemplarpaths)
#         sizes = unique!(size.(points))
#         @assert length(sizes) == 1
#         dim, num_points = sizes[1]
#         threshold = EPSILON * dim * num_points
#         Emin, Emax = extrema(riesz_energy.(points))
#         @assert (Emax - Emin) / (Emax + Emin) < threshold
#         for i = 1 : length(points)
#             for j = i+1 : length(points)
#                 if !isisomorphic(points[i], points[j]; verbose=true)
#                     println("ERROR:")
#                     println("    ", exemplarpaths[i])
#                     println("    ", exemplarpaths[j])
#                     println(isisomorphic(readpoints(exemplarpaths[i]), readpoints(exemplarpaths[j])))
#                 end
#             end
#         end
#     end
# end

main()


# Example minimum variance analysis (MVA) from Galileo satellite Ganymede flyby
# magnetometer data.
# G8 inbound magnetopausee LMN for interval 15:51:40-15:52:20
# I cannot repeat the results of Figure 6 in [Kivelson 1998]
# -0.72 -0.55 0.43 from [Kivelson 1998]
# -0.942 -0.187 0.278 from elliptic empirical model
#
# Hongyang Zhou, hyzhou@umich.edu 01/10/2020

using CSV, DataFrames, CodecZlib, Dates, PyPlot, Mmap

"""
    mva_analysis(filename::AbstractString, index_::UnitRange, DoPlot=false)

This function is currently only written for the specific input file.
"""
function mva_analysis(filename::AbstractString, index_::UnitRange, DoPlot=false)
   # Load magnetometer data Bx, By, Bz
   if filename[end-1:end] == "gz"
      df = CSV.File(transcode(GzipDecompressor, Mmap.mmap(filename));
         header=2, delim=" ", ignorerepeated=true) |> DataFrame
   else
      df = CSV.File(filename; header=2, delim=" ", ignorerepeated=true)
   end

   # inbound crossing interval
   F = mva(df.Bx[index_], df.By[index_], df.Bz[index_])

   t = DateTime.(df.yr, df.month, df.day, df.hr, df.min, floor.(Int, df.sec),
      floor.(Int, 1e3 .* (df.sec - floor.(df.sec))) )

   if DoPlot
      # Original Cartesian coordinates
      figure(figsize=(12,5))
      subplot(411)
      plot(t, df.Bx, label=L"B_x")
      legend()
      subplot(412)
      plot(t, df.By, label=L"B_y")
      legend()
      subplot(413)
      plot(t, df.Bz, label=L"B_z")
      legend()
      subplot(414)
      plot(hypot.(df.Bx, df.By, df.Bz), label=L"B")
      legend()
      tight_layout()

      # Coordinate transformation
      id_ = 2000:2900
      BL = sum(F.vectors[:,1]' .* [df.Bx[id_] df.By[id_] df.Bz[id_]];dims=2)
      BM = sum(F.vectors[:,2]' .* [df.Bx[id_] df.By[id_] df.Bz[id_]];dims=2)
      BN = sum(F.vectors[:,3]' .* [df.Bx[id_] df.By[id_] df.Bz[id_]];dims=2)
      figure(figsize=(12,5))
      subplot(411)
      plot(t[id_], BL, label=L"B_L")
      legend()
      subplot(412)
      plot(t[id_], BM, label=L"B_M")
      legend()
      subplot(413)
      plot(t[id_], BN, label=L"B_N")
      legend()
      subplot(414)
      plot(t[id_], hypot.(BL, BM, BN), label=L"B")
      legend()
      tight_layout()
   end

   F
end


module StandardYoungTableaux

import Distributions

export IPartition
export StandardYoungTableau
export dualSYT
export SYT2ballot
export ballot2SYT
export allSYTx
export allBallotSequences
export firstSYT
export nextSYT
export hooks
export hooklengths
export countSYTx
export RS
export SYT2YoungPath
export YoungPath2SYT
export randomSYT
export randomYoungPath

"""
    IPartition(partition)

Defines an integer partition.

# Argument

- `partition`: vector of positive integers in decreasing order
"""
struct IPartition
  n::Int64
  partition::Vector{Int64}

  function IPartition(partition::Vector{Int64}, check::Bool = true)
    partition = filter(!=(0), partition)
    if check
      if any(diff(partition) .> 0)
        throw(ArgumentError("Partition entries must be decreasing."))
      end
      if !isempty(partition)
        partition[end] >= 1 ||
          throw(ArgumentError("Found non-positive entry in partition."))
      end
    end
    return new(sum(partition), partition)
  end
end

function _jaggedTranspose(aa)
  out = Vector{Vector{Int64}}(undef, 0)
  while !isempty(aa)
    push!(out, map(x -> x[1], aa))
    aa = filter(x -> !isempty(x), map(x -> x[2:end], aa))
  end
  return out
end

function _checkSYTrows(tableau)
  all(map(row -> all(diff(row) .> 0), tableau))
end

function _isSYT(tableau) # _isSYT([[1], [2,3]]) is true ! - does not check shape
  entries = vcat(tableau...)
  return issetequal(entries, 1:length(entries)) &&
         _checkSYTrows(tableau) &&
         _checkSYTrows(_jaggedTranspose(tableau))
end

"""
    StandardYoungTableau(tableau)

Defines a standard Young tableau.

# Argument
- `tableau`: a vector of vectors of positive integers
"""
struct StandardYoungTableau
  shape::IPartition
  tableau::Vector{Vector{Int64}}

  function StandardYoungTableau(tableau::Vector{Vector{Int64}}, check::Bool = true)
    shape = map(length, tableau)
    if check
      if any(shape .== 0) || any(diff(shape) .> 0)
        throw(ArgumentError("Invalid shape."))
      end
      if !_isSYT(tableau)
        throw(ArgumentError("Not a valid tableau."))
      end
    end
    return new(IPartition(shape, false), tableau)
  end
end

"""
    dualSYT(syt)

Dual (conjugate) standard Young tableau.

# Argument
- `syt`: a standard Young tableau
"""
function dualSYT(syt::StandardYoungTableau)
  return StandardYoungTableau(_jaggedTranspose(syt.tableau), false)
end

function _isBallotSequence(s::Vector{Int64}) # https://oeis.org/A000085
  if s[1] != 1 || any(s .<= 0)
    return false
  end
  for k = 1:length(s)
    init = s[1:k]
    for r = 1:maximum(init)
      if length(findall(==(r), init)) < length(findall(==(r + 1), init))
        return false
      end
    end
  end
  return true
end

function _syt2ballot(tableau)
  N = sum(map(length, tableau))
  a = Int64[]
  for i = 1:N
    k = 0
    j = Int64[]
    while isempty(j)
      k += 1
      j = findall(==(i), tableau[k])
      if !isempty(j)
        push!(a, k)
      end
    end
  end
  return a
end

"""
    SYT2ballot(syt)

The ballot sequence corresponding to a standard Young tableau.

# Argument
- `syt`: a standard Young tableau
"""
function SYT2ballot(syt::StandardYoungTableau)
  return _syt2ballot(syt.tableau)
end

function _ballot2syt(a::Vector{Int64})
  return map(i -> findall(==(i), a), 1:maximum(a))
end

"""
    ballot2SYT(b)

The standard Young tableau corresponding to a ballot sequence.

# Argument
- `b`: a ballot sequence (vector of integers)
"""
function ballot2SYT(b::Vector{Int64})
  if !_isBallotSequence(b)
    error("Not a ballot sequence.")
  end
  return StandardYoungTableau(_ballot2syt(b), false)
end

function _ballot(N, lambda, a, more)
  #N = sum(lambda)
  it = N
  @inbounds if more
    lambda = zeros(Int64, N)
    lambda[1] = 1
    isave = 0
    for i = 2:N
      lambda[a[i]] += 1
      if a[i] < a[i-1]
        isave = i
        break
      end
    end
    if isave == 0
      return (a = a, more = false)
    end
    it = lambda[a[isave]+1]
    for j = 0:(N-1)
      i = N - j
      if lambda[i] == it
        a[isave] = i
        lambda[i] -= 1
        it = isave - 1
        break
      end
    end
  end # end of 'if more'
  lambda0 = zeros(Int64, N)
  @inbounds lambda0[1:length(lambda)] = lambda
  lambda = lambda0
  k = 1
  ir = 1
  @inbounds while true
    if N < ir
      break
    end
    if lambda[ir] != 0
      a[k] = ir
      lambda[ir] -= 1
      k += 1
      ir += 1
      continue
    end
    if it < k
      break
    end
    ir = 1
  end
  if N == 1
    return (a = a, more = false)
  end
  @inbounds for j = 1:(N-1)
    if a[j+1] < a[j]
      return (a = a, more = true)
    end
  end
  return (a = a, more = false)
end

"""
    allBallotSequences(lambda)

All ballot sequences associated to a given integer partition.

# Argument
- `lambda`: an integer partition
"""
function allBallotSequences(lambda::IPartition)
  N = lambda.n
  partition = lambda.partition
  (a, more) = _ballot(N, partition, zeros(Int64, N), false)
  out = [copy(a)]
  while more
    (a, more) = _ballot(N, partition, a, true)
    push!(out, copy(a))
  end
  return out
end

"""
    allSYTx(lambda)

All standard Young tableaux of a given shape.

# Argument
- `lambda`: an integer partition

# Example

    lambda = IPartition([3, 1])
    allSYTx(lambda)
"""
function allSYTx(lambda::IPartition)
  N = lambda.n
  partition = lambda.partition
  (a, more) = _ballot(N, partition, zeros(Int64, N), false)
  out = [StandardYoungTableau(_ballot2syt(a), false)]
  while more
    (a, more) = _ballot(N, partition, a, true)
    push!(out, StandardYoungTableau(_ballot2syt(a), false))
  end
  return out
end

function _islastsyt(tableau)
  a = _syt2ballot(tableau)
  N = length(a)
  for j = 1:(N-1)
    if a[j+1] < a[j]
      return false
    end
  end
  return true
end

function _firstsyt(lambda, a)
  lambda = filter(!=(0), lambda)
  N = sum(lambda)
  if N == 1
    return [[1]]
  end
  it = N
  lambda0 = zeros(Int64, N)
  lambda0[1:length(lambda)] = lambda
  lambda = lambda0
  k = 1
  ir = 1
  while true
    if N < ir
      break
    end
    if lambda[ir] != 0
      a[k] = ir
      lambda[ir] -= 1
      k += 1
      ir += 1
      continue
    end
    if it < k
      break
    end
    ir = 1
  end
  return _ballot2syt(a)
end

function _nextsyt(tableau)
  if _islastsyt(tableau)
    return nothing
  end
  a = _syt2ballot(tableau)
  N = sum(map(length, tableau))
  lambda = zeros(Int64, N)
  lambda[1] = 1
  isave = 0
  for i = 2:N
    lambda[a[i]] += 1
    if a[i] < a[i-1]
      isave = i
      break
    end
  end
  if isave == 0
    return _ballot2syt(a)
  end
  it = lambda[a[isave]+1]
  for j = 0:(N-1)
    i = N - j
    if lambda[i] == it
      a[isave] = i
      lambda[i] -= 1
      it = isave - 1
      break
    end
  end
  return _firstsyt(lambda, a)
end

"""
    firstSYT(lambda)

The 'first' standard Young tableau of a given shape.

# Argument
- `lambda`: an integer partition
"""
function firstSYT(lambda::IPartition)
  return StandardYoungTableau(
    _firstsyt(lambda.partition, zeros(Int64, lambda.n)),
    false
  )
end

"""
    nextSYT(syt)

The 'next' standard Young tableau.

# Argument
- `syt`: a standard Young tableau

# Example

    y = firstSYT(IPartition([4, 2]))
    nextSYT(y)
"""
function nextSYT(syt::StandardYoungTableau)
  succ = _nextsyt(syt.tableau)
  if isnothing(succ)
    return nothing
  end
  return StandardYoungTableau(succ, false)
end

function _dualPartition(lambda)
  return map(i -> length(findall(>=(i), lambda)), 1:lambda[1])
end

"""
    hooks(lambda)

The hooks of a partition.

# Argument
- `lambda`: an integer partition
"""
function hooks(lambda::IPartition)
  partition = lambda.partition
  dlambda = _dualPartition(partition)
  l = length(partition)
  out = Vector{Vector{Vector{Int64}}}(undef, l)
  for i = 1:l
    out[i] = Vector{Vector{Int64}}(undef, partition[i])
    for j = 1:partition[i]
      out[i][j] = [0; partition[i] - j + 1]
    end
  end
  for j = 1:length(dlambda)
    for i = 1:dlambda[j]
      out[i][j][1] = dlambda[j] - i + 1
    end
  end
  return out
end

"""
    hooklengths(lambda)

The hook lengths of a partition.

# Argument
- `lambda`: an integer partition
"""
function hooklengths(lambda::IPartition)
  h = hooks(lambda)
  map(x -> vcat((map(y -> sum(y) - 1, x))...), h)
end

"""
    countSYTx(lambda)

Number of standard Young tableaux of a given shape.

# Argument
- `lambda`: an integer partition, the shape

# Example

    lambda = IPartition([4, 2])
    countSYTx(lambda)
    length(allSYTx(lambda))
"""
function countSYTx(lambda::IPartition)
  numterms = Int64[]
  denterms = filter(>=(2), vcat(hooklengths(lambda)...))
  for i = 2:lambda.n
    idx = findfirst(==(i), denterms)
    if isnothing(idx)
      push!(numterms, i)
    else
      deleteat!(denterms, idx)
    end
  end
  return numerator(prod(numterms) // prod(denterms))
end

function _bump(P, Q, e, i)
  if isempty(P)
    return (P = [[e]], Q = [[i]])
  end
  p = P[1]
  if e > last(p)
    P[1] = vcat(p, e)
    Q[1] = vcat(Q[1], i)
    return (P = P, Q = Q)
  else
    j = findfirst(!, p .< e)
    w = p[j]
    P[1][j] = e
    b = _bump(P[2:end], Q[2:end], w, i)
    return (P = vcat([P[1]], b.P), Q = vcat([Q[1]], b.Q))
  end
end

"""
    RS(sigma)

Pair of standard Young tableaux given by the Robinson-Schensted(-Knuth) correspondence from a permutation.

# Argument
- `sigma`: a permutation given as a vector of integers

# Example

    RS([3, 4, 1, 2])
"""
function RS(sigma)
  if !issetequal(sigma, 1:length(sigma))
    error("`sigma` is not a permutation")
  end
  out = _bump([], [], sigma[1], 1)
  for i = 2:length(sigma)
    out = _bump(out.P, out.Q, sigma[i], i)
  end
  return (
    P = StandardYoungTableau(out.P, false),
    Q = StandardYoungTableau(out.Q, false)
  )
end

function _whichRow(tableau, x)
  idx = findfirst(==(x), tableau[1])
  i = 1
  while isnothing(idx)
    i += 1
    idx = findfirst(==(x), tableau[i])
  end
  return i
end

"""
    SYT2YoungPath(syt)

Path of the Young graph corresponding to a standard Young tableau.

# Argument
- `syt`: a standard Young tableau

# Example

    y = StandardYoungTableau([[1,3,4], [2]])
    SYT2YoungPath(y)
"""
function SYT2YoungPath(syt::StandardYoungTableau)
  N = syt.shape.n
  path = Vector{Vector{Int64}}(undef, N)
  path[1] = vcat(1, zeros(Int64, length(syt.tableau) - 1))
  for k = 2:N
    part = copy(path[k-1])
    i = _whichRow(syt.tableau, k)
    part[i] += 1
    path[k] = part
  end
  return map(x -> IPartition(x, false), path)
end

"""
    YoungPath2SYT(path)

Standard Young tableau corresponding to a path of the Young graph.

# Argument
- `path`: a vector of integers partitions forming a path of the Young graph

# Example

    y = StandardYoungTableau([[1,3,4], [2]])
    path = SYT2YoungPath(y)
    YoungPath2SYT(path)
"""
function YoungPath2SYT(path::Vector{IPartition})
  Ns = map(x -> x.n, path)
  N = length(path)
  if any(Ns .!= 1:N)
    error(
      "Invalid path. The n-th element of the path must be a partition of n."
    )
  end
  if N == 1
    return StandardYoungTableau([[1]], false)
  end
  incrmts = diff(map(x -> length(x.partition), path))
  if any((incrmts .!= 0) .* (incrmts .!= 1))
    error("Invalid path.")
  end
  l = length(path[end].partition)
  path2 = zeros(Int64, l, N)
  for i = 1:N
    part_i = path[i].partition
    path2[1:length(part_i), i] = part_i
  end
  diffs = diff(path2, dims = 2)
  check = mapslices(
    x -> sum(x) == 1 && all((x .== 0) .| (x .== 1)), diffs, dims = 1
  )
  if !all(check)
    error("invalid path")
  end
  tableau = vcat([[1]], repeat([Int64[]], l - 1))
  for i = 1:(N-1)
    idx = findfirst(==(1), diffs[:, i])
    tableau[idx] = vcat(tableau[idx], i + 1)
  end
  return StandardYoungTableau(tableau, false)
end

"""
    randomSYT(lambda)

Uniformly random standard Young tableau of a given shape.

# Argument
- `lambda`: an integer partition, the shape
"""
function randomSYT(lambda::IPartition)
  partition = copy(lambda.partition)
  N = lambda.n
  a = zeros(Int64, N)
  i = 0
  k = 0
  while true
    i += 1
    for j in 1:partition[i]
      a[j] += 1
      k += 1
    end
    if N <= k
      break
    end
  end
  for m in 1:N
    local i, j
    while true
      i = rand(1:a[1])
      j = rand(1:partition[1])
      if i <= a[j] && j <= partition[i]
        break
      end
    end
    while true
      ih = a[j] + partition[i] - i - j
      if ih == 0
        break
      end
      k = rand(1:ih)
      if k <= (partition[i] - j)
        j += k
      else
        i += k - partition[i] + j
      end
    end
    partition[i] -= 1
    a[j] -= 1
    a[N - m + 1] = i
  end
  return StandardYoungTableau(_ballot2syt(a), false)
end

function _connectedPartitions(partition0)
  partition = copy(partition0)
  partition[1] += 1
  out = [partition]
  for i in 2:length(partition0)
    if partition0[i] < partition0[i - 1]
      partition = copy(partition0)
      partition[i] += 1
      push!(out, partition)
    end
  end
  return vcat(out, [vcat(partition0, 1)])
end

"""
    randomYoungPath(n)

Samples a path of the Young graph according to the Plancherel growth process.

# Argument
`n`: a positive integer, the size of the path to be sampled

# Example

    path = randomYoungPath(5)
    YoungPath2SYT(path)
"""
function randomYoungPath(n::Int64)
  path = [[1]]
  for i in 1:(n-1)
    set = _connectedPartitions(path[i])
    f_nu = countSYTx(IPartition(path[i], false))
    probs = map(p -> countSYTx(IPartition(p, false)), set) ./ ((i + 1) * f_nu)
    k = rand(Distributions.Categorical(probs))
    push!(path, set[k])
  end
  return map(x -> IPartition(x, false), path)
end

end


using GeoPhyInv
using Base.Test
using BenchmarkTools
using Calculus

const JD=GeoPhyInv.Data

fields=[:P]
# testing misfit and gradient calculation for TD objects
tgrid1=range(0., stop=1., length=5);
tgrid2=range(0., stop=1., length=10);

ageom=GeoPhyInv.Acquisition.AGeom_fixed(10,10,10,10,10,10,10,10)

x=GeoPhyInv.Data.TD_ones(fields,tgrid1,ageom)

yvec=randn(length(x))
xvec=similar(yvec)

@testset "testing copy!" begin
	@btime copy!(x, yvec);

	@btime copy!(xvec, x);

	@test xvec == yvec
end

@testset "simple LS error: x and y same time grid" begin
	y=GeoPhyInv.Data.TD_ones(fields,tgrid1,ageom)

	randn!(x)
	randn!(y)

	pa=GeoPhyInv.Data.P_misfit(x,y);
	@time GeoPhyInv.Data.func_grad!(pa,:dJx);
	gg1=vec(pa.dJx)

	xvec=vec(pa.x)
	gg2=Calculus.gradient(x -> (copy!(pa.x, x); return GeoPhyInv.Data.func_grad!(pa)), xvec)

	# check gradient with Finite Differencing
	@test gg1 ≈ gg2
end

#=

rrrrrr

# loop over same time grid and different time grid (interp_flag on/off)
for y in [GeoPhyInv.Data.TD_ones(fields,tgrid2,ageom), GeoPhyInv.Data.TD_ones(fields,tgrid1,ageom)]
	println("#########################################")

	# generate some random data
	randn!(y.d[1,1])
	randn!(x.d[1,1])

	for func_attrib in [:cls]
		coup=GeoPhyInv.Coupling.TD_delta(y.tgrid, [0.1,0.1], 0.0,  x.fields, x.ageom)
		randn!(coup.ssf[1,1])
		pa=GeoPhyInv.Data.P_misfit(x,y, func_attrib=func_attrib, coup=coup);

		xvec=vec(pa.x.d[1,1])
		gg2=Calculus.gradient(x -> err(x), xvec)

		@time GeoPhyInv.Data.func_grad!(pa,:dJx);
		gg1=vec(pa.dJx.d[1,1])
		# check gradient with Finite Differencing
		@test gg1 ≈ gg2

		function errw(xvec)

			for i in eachindex(pa.coup.ssf[1,1])
				pa.coup.ssf[1,1][i]=xvec[i]
			end

			return GeoPhyInv.Data.func_grad!(pa)
		end

		xvec=vec(pa.coup.ssf[1,1])
		gg2=Calculus.gradient(x -> errw(x), xvec)

		@time GeoPhyInv.Data.func_grad!(pa,:dJssf);
		gg1=vec(pa.dJssf[1,1])
		# check gradient with Finite Differencing
		@test gg1 ≈ gg2

	end
end

=#


"""
	RandomVariable(dist::Sampleable{Univariate}, name::Symbol)

Defines a random variable, with a univariate distribution from Distributions.jl and a name.

# Examples

```jldoctest
julia> RandomVariable(Normal(), :x)
RandomVariable(Normal{Float64}(μ=0.0, σ=1.0), :x)

julia> RandomVariable(Exponential(1), :x)
RandomVariable(Exponential{Float64}(θ=1.0), :x)
```
"""
struct RandomVariable <: RandomUQInput
    dist::Sampleable{Univariate}
    name::Symbol
end

"""
	sample(rv::RandomVariable, n::Integer)

Generates n samples from a random variable. Returns a DataFrame.

# Examples

```jldoctest
x = RandomVariable(Normal(), :x)
sample(x, 3)

# output

3×1 DataFrame
│ Row │ x          │
│     │ Float64    │
├─────┼────────────┤
│ 1   │ 1.00763    │
│ 2   │ -0.643618  │
│ 3   │ -0.0826132 │
```

See also: [`RandomVariable`](@ref)
"""
function sample(rv::RandomVariable, n::Integer=1)
    return DataFrame(rv.name => rand(rv.dist, n))
end

function to_physical_space!(rv::RandomVariable, x::DataFrame)
    x[!, rv.name] = quantile.(rv.dist, cdf.(Normal(), x[:, rv.name]))
    return nothing
end

function to_standard_normal_space!(rv::RandomVariable, x::DataFrame)
    x[!, rv.name] = quantile.(Normal(), cdf.(rv.dist, x[:, rv.name]))
    return nothing
end

mean(rv::RandomVariable) = DataFrame(rv.name => Distributions.mean(rv.dist))
mean(rvs::Array{RandomVariable}) = mapreduce(mean, hcat, rvs)

dimensions(rv::RandomVariable) = 1


using Mimi

page_years = [2020, 2030, 2040, 2050, 2075, 2100, 2150, 2200, 2250, 2300]
page_year_0 = 2015

function getpageindexfromyear(year)
    i = findfirst(isequal(year), page_years)
    if i == 0
        error("Invalid PAGE year: $year.")
    end
    return i
end

function getperiodlength(year)      # same calculations made for yagg_periodspan in the model
    i = getpageindexfromyear(year)

    if year==page_years[1]
        start_year = page_year_0
    else
        start_year = page_years[i - 1]
    end

    if year == page_years[end]
        last_year = page_years[end]
    else
        last_year = page_years[i + 1]
    end

    return (last_year - start_year) / 2
end

"""Applies undiscounting factor to get the SCC."""
function undiscount_scc(m::Model, year::Int)
    df = m[:EquityWeighting, :df_utilitydiscountfactor]
    consfocus0 = m[:GDP, :cons_percap_consumption_0][1]
    consfocus = m[:GDP, :cons_percap_consumption][:, 1]
    emuc = m[:EquityWeighting, :emuc_utilityconvexity]
    sccii = getpageindexfromyear(year)

    return df[sccii] * ((consfocus[sccii] / consfocus0)^-emuc)
end

@defcomp ExtraEmissions begin
    e_globalCO2emissions = Parameter(index=[time],unit="Mtonne/year")
    pulse_size = Parameter()
    pulse_year = Parameter()
    e_globalCO2emissions_adjusted = Variable(index=[time],unit="Mtonne/year")

    function run_timestep(p, v, d, t)
        if gettime(t) == p.pulse_year
            v.e_globalCO2emissions_adjusted[t] = p.e_globalCO2emissions[t] + p.pulse_size / getperiodlength(p.pulse_year)
        else
            v.e_globalCO2emissions_adjusted[t] = p.e_globalCO2emissions[t]
        end
    end
end

"""
compute_scc(m::Model = get_model(); year::Union{Int, Nothing} = nothing, eta::Union{Float64, Nothing} = nothing, prtp::Union{Float64, Nothing} = nothing)

Computes the social cost of CO2 for an emissions pulse in `year` for the provided Mimi-PAGE model.
If no model is provided, the default model from main_model.get_model() is used.
Discounting scheme can be specified by the `eta` and `prtp` parameters, which will update the values of emuc_utilitiyconvexity and ptp_timepreference in the model.
If no values are provided, the discount factors will be computed using the default PAGE values of emuc_utilitiyconvexity=1.1666666667 and ptp_timepreference=1.0333333333.
"""
function compute_scc(m::Model = get_model(); year::Union{Int, Nothing} = nothing, eta::Union{Float64, Nothing} = nothing, prtp::Union{Float64, Nothing} = nothing, pulse_size = 100000.)
    year === nothing ? error("Must specify an emission year. Try `compute_scc(m, year=2020)`.") : nothing
    !(year in page_years) ? error("Cannot compute the scc for year $year, year must be within the model's time index $page_years.") : nothing

    eta == nothing ? nothing : update_param!(m, :emuc_utilityconvexity, eta)
    prtp == nothing ? nothing : update_param!(m, :ptp_timepreference, prtp * 100.)

    mm = get_marginal_model(m, year=year, pulse_size=pulse_size)   # Returns a marginal model that has already been run
    scc = mm[:EquityWeighting, :td_totaldiscountedimpacts] / undiscount_scc(mm.base, year)

    return scc
end

"""
compute_scc_mm(m::Model = get_model(); year::Union{Int, Nothing} = nothing, eta::Union{Float64, Nothing} = nothing, prtp::Union{Float64, Nothing} = nothing)

Returns a NamedTuple (scc=scc, mm=mm) of the social cost of carbon and the MarginalModel used to compute it.
Computes the social cost of CO2 for an emissions pulse in `year` for the provided Mimi-PAGE model.
If no model is provided, the default model from main_model.get_model() is used.
Discounting scheme can be specified by the `eta` and `prtp` parameters, which will update the values of emuc_utilitiyconvexity and ptp_timepreference in the model.
If no values are provided, the discount factors will be computed using the default PAGE values of emuc_utilitiyconvexity=1.1666666667 and ptp_timepreference=1.0333333333.
"""
function compute_scc_mm(m::Model = get_model(); year::Union{Int, Nothing} = nothing, eta::Union{Float64, Nothing} = nothing, prtp::Union{Float64, Nothing} = nothing, pulse_size = 100000.)
    year === nothing ? error("Must specify an emission year. Try `compute_scc(m, year=2020)`.") : nothing
    !(year in page_years) ? error("Cannot compute the scc for year $year, year must be within the model's time index $page_years.") : nothing

    eta == nothing ? nothing : update_param!(m, :emuc_utilityconvexity, eta)
    prtp == nothing ? nothing : update_param!(m, :ptp_timepreference, prtp * 100.)

    mm = get_marginal_model(m, year=year, pulse_size=pulse_size)   # Returns a marginal model that has already been run
    scc = mm[:EquityWeighting, :td_totaldiscountedimpacts] / undiscount_scc(mm.base, year)
    scc_disaggregated = mm[:EquityWeighting, :addt_equityweightedimpact_discountedaggregated] / undiscount_scc(mm.base, year)

    return (scc = scc, scc_disaggregated = scc_disaggregated, mm = mm)
end

"""
get_marginal_model(m::Model = get_model(); year::Union{Int, Nothing} = nothing)
Returns a Mimi MarginalModel where the provided m is the base model, and the marginal model has additional emissions of CO2 in year `year`.
If no Model m is provided, the default model from main_model.get_model() is used as the base model.
Note that the returned MarginalModel has already been run.
"""
function get_marginal_model(m::Model = get_model(); year::Union{Int, Nothing} = nothing, pulse_size = 100000.)
    year === nothing ? error("Must specify an emission year. Try `get_marginal_model(m, year=2020)`.") : nothing
    !(year in page_years) ? error("Cannot add marginal emissions in $year, year must be within the model's time index $page_years.") : nothing

    mm = create_marginal_model(m, pulse_size)

    add_comp!(mm.marginal, ExtraEmissions, :extra_emissions; after=:co2emissions)
    connect_param!(mm.marginal, :extra_emissions => :e_globalCO2emissions, :co2emissions => :e_globalCO2emissions)
    set_param!(mm.marginal, :extra_emissions, :pulse_size, pulse_size)
    set_param!(mm.marginal, :extra_emissions, :pulse_year, year)

    connect_param!(mm.marginal, :CO2Cycle => :e_globalCO2emissions, :extra_emissions => :e_globalCO2emissions_adjusted)

    run(mm)

    return mm
end

function compute_scc_mcs(m::Model, samplesize::Int; year::Union{Int, Nothing} = nothing, eta::Union{Float64, Nothing} = nothing, prtp::Union{Float64, Nothing} = nothing, pulse_size = 100000.)
    # Setup of location of final results
    scc_results = zeros(samplesize)

    function mc_scc_calculation(sim_inst::SimulationInstance, trialnum::Int, ntimesteps::Int, ignore::Nothing)
        marginal = sim_inst.models[1]
        marg_damages = marginal[:EquityWeighting, :td_totaldiscountedimpacts] / undiscount_scc(mm.base, year)
        scc_results[trialnum] = marg_damages
    end

    # get simulation
    mcs = getsim()

    # Setup models
    eta == nothing ? nothing : update_param!(m, :emuc_utilityconvexity, eta)
    prtp == nothing ? nothing : update_param!(m, :ptp_timepreference, prtp * 100.)

    mm = get_marginal_model(m, year=year, pulse_size=pulse_size)   # Returns a marginal model that has already been run

    # Run
    res = run(mcs, mm, samplesize; post_trial_func=mc_scc_calculation)

    scc_results
end


"""si basa sul file voxellizationProva ma usa il mio package invece di LinearAlgebraicAlgebra
e verrà usato per l'analisi delle prestazioni una volta implementata la parallellizazzione"""

using CuboidGrids, SparseArrays, DataStructures
using Distributed
Lar = CuboidGrids

using ViewerGL
GL = ViewerGL

# Dati da larGrid

function CV2FV( v::Array{Int64} )
	faces = [
		[v[1], v[2], v[3], v[4]], [v[5], v[6], v[7], v[8]],
		[v[1], v[2], v[5], v[6]],	[v[3], v[4], v[7], v[8]],
		[v[1], v[3], v[5], v[7]], [v[2], v[4], v[6], v[8]]]
end

function CV2EV( v::Array{Int64} )
	edges = [
		[v[1], v[2]], [v[3], v[4]], [v[5], v[6]], [v[7], v[8]], [v[1], v[3]], [v[2], v[4]],
		[v[5], v[7]], [v[6], v[8]], [v[1], v[5]], [v[2], v[6]], [v[3], v[7]], [v[4], v[8]]]
end

function K( CV )	#costruisce e ritorna una matrice sparsa a partire dalla collezione passata(matrice dei vertici, degli edges, delle facce, etc)
	I = vcat( [ [k for h in CV[k]] for k=1:length(CV) ]...) #righe
	J = vcat(CV...)											#colonne
	Vals = [1 for k=1:length(I)]							#valori della matrice
	return sparse(I,J,Vals)									#costruzione della matrice sparsa
end

V,CV = Lar.cuboidGrid([10,20,10])
#pointcloud,V,CV = random3cells([40,20,10],4_000)

VV = [[v] for v=1:size(V,2)]
FV = convert(Array{Array{Int64,1},1}, collect(Set(cat(pmap(CV2FV,CV))))) #verts faces
EV = convert(Array{Array{Int64,1},1}, collect(Set(cat(pmap(CV2EV,CV))))) #verts edges

#Cotruzine delle varie matrici tramite la funzione K
M_0 = K(VV)
M_1 = K(EV)
M_2 = K(FV)
M_3 = K(CV)

∂_1 = M_0 * M_1'
∂_2 = (M_1 * M_2') .÷ 2 #	.÷ sum(M_1,dims=2)
s = sum(M_2,dims=2)
∂_3 = (M_2 * M_3')
∂_3 = ∂_3 ./	s
∂_3 = ∂_3 .÷ 1	#	.÷ sum(M_2,dims=2)

S2 = sum(∂_3,dims=2)
inner = [k for k=1:length(S2) if S2[k]==2]
outer = setdiff(collect(1:length(FV)), inner)

# point cloud
#GL.VIEW([GL.GLPoints(convert(Lar.Points,pointcloud'), GL.COLORS[1])])

# voxels
GL.VIEW([GL.GLLar2gl(V,CV)])

#=
# edges
GL.VIEW([ GL.GLGrid(V,EV,GL.Point4d(1,1,1,1))
          GL.GLAxis(GL.Point3d(-1,-1,-1),GL.Point3d(1,1,1)) ]);

# inner faces
GL.VIEW([ GL.GLGrid(V,FV[inner],GL.Point4d(1,1,1,1))
          GL.GLAxis(GL.Point3d(-1,-1,-1),GL.Point3d(1,1,1)) ]);

# outer faces
GL.VIEW([ GL.GLGrid(V,FV[outer],GL.Point4d(1,1,1,1))
          GL.GLAxis(GL.Point3d(-1,-1,-1),GL.Point3d(1,1,1)) ]);

=#

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

V,CV = Lar.cuboidGrid([5,20,10])
#pointcloud,V,CV = random3cells([40,20,10],4_000)

VV = [[v] for v=1:size(V,2)]
FV = convert(Array{Array{Int64,1},1}, collect(Set(cat(pmap(CV2FV,CV)))))
EV = convert(Array{Array{Int64,1},1}, collect(Set(cat(pmap(CV2EV,CV)))))

M_0 = K(VV)
M_1 = K(EV)
M_2 = K(FV)
M_3 = K(CV)

∂_1 = M_0 * M_1'
∂_2 = (M_1 * M_2') .÷ 2 #	.÷ sum(M_1,dims=2)
s = sum(M_2,dims=2)
∂_3 = (M_2 * M_3')
∂_3 = ∂_3 ./	s
∂_3 = ∂_3 .÷ 1	#	.÷ sum(M_2,dims=2)

S2 = sum(∂_3,dims=2)
inner = [k for k=1:length(S2) if S2[k]==2]
outer = setdiff(collect(1:length(FV)), inner)

# point cloud
#GL.VIEW([GL.GLPoints(convert(Lar.Points,pointcloud'), GL.COLORS[1])])

# voxels
GL.VIEW([GL.GLLar2gl(V,CV)])


using optim_mnfctr
using Test

l = [0 3 3 2 1 1 2 3 2 1 2 3;
     1 0 2 1 1 2 3 4 2 3 4 1;
     1 2 0 2 1 2 3 4 4 3 1 6;
     1 2 3 0 2 1 2 3 2 5 8 4;
     1 2 3 4 0 2 3 4 3 6 8 4;
     1 2 3 4 5 0 1 2 4 5 6 9;
     1 2 3 4 5 6 0 3 4 5 6 5;
     1 2 3 4 5 6 7 0 4 5 6 7;
     1 2 3 4 5 0 1 2 0 5 6 3;
     1 2 3 4 5 6 0 3 4 0 6 2;
     1 2 3 4 5 6 7 0 4 5 0 4;
     2 3 4 5 6 7 8 9 1 2 3 4]

m = [0 3 3 2 1 1 2 3 2 1 2;
     1 0 2 1 1 2 3 4 2 3 4;
     1 2 0 2 1 2 3 4 4 3 1;
     1 2 3 0 2 1 2 3 2 5 8;
     1 2 3 4 0 2 3 4 3 6 8;
     1 2 3 4 5 0 1 2 4 5 6;
     1 2 3 4 5 6 0 3 4 5 6;
     1 2 3 4 5 6 7 0 4 5 6;
     1 2 3 4 5 0 1 2 0 5 6;
     1 2 3 4 5 6 0 3 4 0 6;
     1 2 3 4 5 6 7 0 4 5 0]

n = [0 3 3 2 1 2 4 3 2;
     1 0 2 1 1 0 10 4 3;
     1 2 0 2 1 2 2 2 5;
     1 2 3 0 2 1 4 5 4;
     1 2 3 4 0 3 5 6 5;
     1 2 3 4 5 6 7 3 4;
     2 3 4 5 6 7 8 3 4;
     3 4 5 6 7 8 9 8 7;
     1 2 3 4 5 6 7 8 9]

matrix = n

@testset "optim_mnfctr test 1" begin

    # min_threaded, min_sequence_threaded = @btime run_threaded($n)
    min_threaded, min_sequence_threaded = run_threaded(n)

    # min_sequential, min_sequence_sequential = @btime run_sequential($n)
    min_sequential, min_sequence_sequential = run_sequential(n)

    # @assert min_threaded == min_sequential && min_sequence_threaded == min_sequence_sequential
    @test min_threaded == min_sequential && min_sequence_threaded == min_sequence_sequential

    # min_threaded_no_simd, min_sequence_threaded_no_simd = @btime run_threaded_no_simd(m)
    # min_threaded, min_sequence_threaded = @btime run_threaded(m)
    # min_threaded, min_sequence_threaded = @btime run_threaded(l)

    #print(min_)
    #print(min_sequence)

end

#=
To DO
- compare solutions with package TravellingsalesemanExact
- add heuristucs for non-exact solution (n > 13)
=#


# This implements the variant when A is dense and m < n

using LinearAlgebra

struct IndGraphFat{T <: RealOrComplex} <: IndGraph
  m::Int
  n::Int
  A::Array{T,2}
  AA::Array{T,2}
  tmp::Array{T,1}
  # tmpx::Array{T,1}
  F::Cholesky{T, Array{T, 2}} # LL factorization
end

function IndGraphFat(A::Array{T,2}) where {T <: RealOrComplex}
  m, n = size(A)
  AA = A * A'

  F = cholesky(AA + I)
  #normrows = vec(sqrt.(sum(abs2.(A), 2)))

  IndGraphFat(m, n, A, AA, Array{T, 1}(undef, m), F)
end

function prox!(
    x::AbstractArray{T},
    y::AbstractArray{T},
    f::IndGraphFat,
    c::AbstractArray{T},
    d::AbstractArray{T},
    gamma=1.0
    ) where {T <: RealOrComplex}

  # y .= f.F \ (f.A * c + f.AA * d)
  mul!(f.tmp, f.A, c)
  mul!(y, f.AA, d)
  y .+= f.tmp
  ldiv!(f.F, y)

  # f.A' * (d - y) + c # note: for complex the complex conjugate is used
  copyto!(f.tmp, d)
  f.tmp .-= y
  mul!(x, adjoint(f.A), f.tmp)
  x .+= c
  return 0.0
end

function (f::IndGraphFat)(x::AbstractArray{T}, y::AbstractArray{T}) where
    {T <: RealOrComplex}
  # the tolerance in the following line should be customizable
  mul!(f.tmp, f.A, x)
  f.tmp .-= y
  if norm(f.tmp, Inf) <= 1e-10
    return 0.0
  end
  return +Inf
end

fun_name(f::IndGraphFat) = "Indicator of an operator graph defined by dense full row rank matrix"
# fun_dom(f::IndGraph) = "AbstractArray{Real,1}, AbstractArray{Complex,1}"
# fun_expr(f::IndGraph) = "x,y ↦ 0 if Ax = y, +∞ otherwise"
# fun_params(f::IndGraph) =
#   string( "A = ", typeof(f.A), " of size ", size(f.A))

function prox_naive(
    f::IndGraphFat,
    c::AbstractArray{T},
    d::AbstractArray{T},
    gamma=1.0
  ) where {T <: RealOrComplex}

  y = f.F \ (f.A * c + f.AA * d)
  return c + f.A' * (d - y), y, 0.0
end

function (f::IndGraphFat)(xy::AbstractVector{T}) where
  {T <: RealOrComplex}
  x, y = splitinput(f, xy)
  return f(x, y)
end

(f::IndGraphFat)(xy::Tuple{AbstractVector{T}, AbstractVector{T}}) where
  {T <: RealOrComplex} = f(xy[1], xy[2])


"""
Computes the maximum flow between the source and target vertexes in a flow
graph using [Dinic\'s Algorithm](https://en.wikipedia.org/wiki/Dinic%27s_algorithm)
Returns the value of the maximum flow as well as the final flow matrix.

Use a default capacity of 1 when the capacity matrix isn\'t specified.

Requires arguments:
residual_graph::DiGraph                # the input graph
source::Int                            # the source vertex
target::Int                            # the target vertex
capacity_matrix::AbstractArray{T,2}    # edge flow capacities
"""

function dinic_impl{T<:Number}(
    residual_graph::DiGraph,               # the input graph
    source::Int,                           # the source vertex
    target::Int,                           # the target vertex
    capacity_matrix::AbstractArray{T,2}    # edge flow capacities
    )
    n = nv(residual_graph)                     # number of vertexes

    flow_matrix = zeros(T, n, n)           # initialize flow matrix
    P = zeros(Int, n)                      # Sharable parent vector

    flow = 0

    while true
        augment = blocking_flow!(residual_graph, source, target, capacity_matrix, flow_matrix, P)
        augment == 0 && break
        flow += augment
    end
    return flow, flow_matrix
end

"""
Uses BFS to identify a blocking flow in
the input graph and then backtracks from the targetto the source, aumenting flow
along all possible paths.

Requires arguments:
residual_graph::DiGraph                # the input graph
source::Int                            # the source vertex
target::Int                            # the target vertex
capacity_matrix::AbstractArray{T,2}    # edge flow capacities
flow_matrix::AbstractArray{T,2}        # the current flow matrix
"""
function blocking_flow!{T<:Number}(
    residual_graph::DiGraph,               # the input graph
    source::Int,                           # the source vertex
    target::Int,                           # the target vertex
    capacity_matrix::AbstractArray{T,2},   # edge flow capacities
    flow_matrix::AbstractArray{T,2},       # the current flow matrix
    )
    P = zeros(T, nv(residual_graph))
    return blocking_flow!(residual_graph,
                          source,
                          target,
                          capacity_matrix,
                          flow_matrix,
                          P)
end

"""blocking_flow!
Preallocated version of blocking_flow.Uses BFS to identify a blocking flow in
the input graph and then backtracks from the target to the source, aumenting flow
along all possible paths.

Requires arguments:
residual_graph::DiGraph                # the input graph
source::Int                            # the source vertex
target::Int                            # the target vertex
capacity_matrix::AbstractArray{T,2}    # edge flow capacities
flow_matrix::AbstractArray{T,2}        # the current flow matrix
P::AbstractArray{Int, 1}               # Parent vector to store Level Graph
"""

function blocking_flow!{T<:Number}(
    residual_graph::DiGraph,               # the input graph
    source::Int,                           # the source vertex
    target::Int,                           # the target vertex
    capacity_matrix::AbstractArray{T,2},   # edge flow capacities
    flow_matrix::AbstractArray{T,2},       # the current flow matrix
    P::AbstractArray{Int, 1}               # Parent vector to store Level Graph
    )
    n = nv(residual_graph)                     # number of vertexes

    fill!(P, -1)
    P[source] = -2

    Q = [source]
    sizehint!(Q, n)

    while length(Q) > 0                   # Construct the Level Graph using BFS
        u = pop!(Q)
        for v in fadj(residual_graph, u)
            if P[v] == -1 && capacity_matrix[u,v] > flow_matrix[u,v]
                P[v] = u
                unshift!(Q, v)
            end
        end
    end

    P[target] == -1 && return 0                    # BFS couldn't reach the target

    total_flow = 0

    for bv in badj(residual_graph, target)    # Trace all possible routes to source
        flow = typemax(T)
        v = target
        u = bv
        while v != source
            if u == -1                    # Vertex unreachable from source
                flow = 0
                break
            else
                flow = min(flow, capacity_matrix[u,v] - flow_matrix[u,v])
                v = u
                u = P[u]
            end
        end

        flow == 0 && continue                      # Flow cannot be augmented along path

        v = target
        u = bv
        while v != source             # Augment flow along path
            flow_matrix[u,v] += flow
            flow_matrix[v,u] -= flow
            v = u
            u = P[u]
        end

        total_flow += flow
    end
    return total_flow
end


using Test

using iLQGames:
    LinearSystem,
    discretize,
    linearize,
    linearize_discrete,
    LinearizationStyle,
    TrivialLinearization

using StaticArrays

A = @SMatrix [1 2; 3 4]
B = @SMatrix [1; 3]
ls = LinearSystem{0}(A, B)
ΔT = 0.1241241

# some test point
x0 = @SVector [1., 2.]
u0 = @SVector [2.]
t0 = .11

d1 = discretize(ls, Val{ΔT}())

@testset "linearization style" begin
    @test LinearizationStyle(ls) isa TrivialLinearization
    @test LinearizationStyle(d1) isa TrivialLinearization
end

@testset "linearize" begin
    # linearizing a linear system should not make a difference
    ls_lin = linearize(ls, x0, u0, t0)
    @test ls.A == ls_lin.A
    @test ls.B == ls_lin.B
end;


using Test

using MarketData
using TimeSeries

using MarketTechnicals

@testset "Momentum" begin

@testset "rsi" begin
    # TTR: RSI(x["Close"])
    # TTR: 45.95363 50.50249 48.97830 48.83191 ...
    @test values(rsi(cl, wilder=true))[1:4] ≈ [45.95363, 50.50249, 48.97830, 48.83191] atol=.01
    # TTR value is 55.84922
    @test values(rsi(cl))[end] ≈ 55.849 atol=.01
    # TTR value is 55.95932
    @test values(rsi(cl, wilder=true))[end] ≈ 55.959 atol=.01
    # TTR value is 56.21947
    @test values(rsi(cl, 10))[end] ≈ 56.219 atol=.01
    @test values(rsi(ohlc))[end, :] ≈ [68.030, 61.872, 70.062, 55.849] atol=.01

    # seed is included (TODO: debate amongst yourselves)
    @test timestamp(rsi(cl))[1]              == Date(2000,1,24)
    @test timestamp(rsi(cl, wilder=true))[1] == Date(2000,1,24)
    @test timestamp(rsi(cl, 10))[1]          == Date(2000,1,18)
    @test timestamp(rsi(ohlc))[1]            == Date(2000,1,24)

    @test timestamp(rsi(cl))[end]              == Date(2001,12,31)
    @test timestamp(rsi(cl, wilder=true))[end] == Date(2001,12,31)
    @test timestamp(rsi(cl, 10))[end]          == Date(2001,12,31)
    @test timestamp(rsi(ohlc))[end]            == Date(2001,12,31)
end

@testset "macd" begin
    @test isapprox(values(macd(cl))[end, 1], -0.020, atol=.01)
    @test isapprox(values(macd(cl))[end, 2], 0.421, atol=.01) # TTR value with percent=FALSE is 0.421175152
    @test isapprox(values(macd(cl))[end, 3], 0.441, atol=.01) # TTR value with percent=FALSE is 0.4414275
    @test timestamp(macd(cl))[end] == Date(2001,12,31)
end

@testset "macd multi-column TimeArray" begin
    # multi-column TimeArray
    # TTR: MACD(..., maType="EMA", percent=0)
    ta = macd(ohlc[:Open, :Close])
    @test colnames(ta)[1:2]  == [:Open_macd, :Close_macd]
    @test timestamp(ta)[end] == Date(2001, 12, 31)
    @test isapprox(values(ta)[end, 3], 0.44254569, atol=.01)    # Open_dif
    @test isapprox(values(ta)[end, 5], 4.536854e-01, atol=.01)  # Open_signal
    @test isapprox(values(ta)[end, 4], 0.421175152, atol=.01)   # Close_dif
    @test isapprox(values(ta)[end, 6], 4.414275e-01, atol=.01)  # Close_signal
end

@testset "chaikinoscillator" begin
    """
    Quote from TTR

    > EMA(adl, 3) - EMA(adl, 10)
        [1]            NA            NA            NA            NA            NA
        [6]            NA            NA            NA            NA  -6851466.867
    [11]  -5508824.158  -4145747.583  -8413321.022 -10025864.828 -10655603.670
    [16]  -8761003.151  -8025814.774  -6995436.368  -6906221.064  -4303730.935
    [21]  -3660985.797  -3451576.257  -2336638.695  -1217362.729    546925.363
    [26]   1788691.016   1349030.229   1738675.716   1239751.751   1946952.206
    [31]   2893048.890   2658896.120   2715328.628   1937437.546   1910085.631
    [36]   2145262.005   1978248.295   1075784.464   1048400.357    814497.124
    [41]   2914695.065   2944188.520   3555908.642   3085073.413   2268377.529
    [46]   1942893.354   1778694.035   1810347.568   1241765.961   -261885.061
    [51]  -1237656.884   -576205.604    613370.430    649457.675   1738353.505
    [56]   3591576.201   2828408.458   1949186.773   1159350.740    626861.789
    [61]    183217.498  -1132617.797   -712700.508   -691386.657    -57610.324
    [66]    774998.562    561961.806   1084391.707    634202.420   -523959.226
    [71]  -1968845.762  -2496129.556  -3167709.694  -2000329.427   -230334.824
    [76]   -577397.260  -1747313.599   -806782.684    617588.662    443152.062
    [81]   1214832.531   1341044.310   1592063.283    936678.383    186056.580
    [86]  -1209996.220  -1516911.058  -1995687.520  -2776094.760  -4090715.405
    """

    ta = chaikinoscillator(ohlcv)
    @test colnames(ta)     == [:chaikinoscillator]
    @test meta(ta)         == meta(ta)
    @test timestamp(ta)[1] == timestamp(ohlcv)[10]
    @test isapprox(values(ta)[1], -6851466.867, atol=.01)
    @test isapprox(values(ta)[2], -5508824.158, atol=.01)
    @test isapprox(values(ta)[3], -4145747.583, atol=.01)
end

@testset "cci" begin
    # TTR::CCI value is -38.931614
    @test isapprox(values(cci(ohlc))[1]  , -38.931614, atol=.01)
    # TTR::CCI value is 46.3511339
    @test isapprox(values(cci(ohlc))[end], 46.3511339, atol=.01)
    @test timestamp(cci(ohlc))[end] == Date(2001, 12, 31)

    ta = cci(TimeArray(collect(Date(2011, 1, 1):Day(1):Date(2011, 1, 30)),
                       fill(42, (30, 4)), [:Open, :High, :Low, :Close]))
    @test all(values(ta) .== 0)
end

@testset "aroon" begin
    """
    Quote from TTR

    > aroon(x[c('High', 'Low')], 25)
            aroonUp aroonDn oscillator
    [24,]      NA      NA         NA
    [25,]      NA      NA         NA
    [26,]      48      28         20
    [27,]      44      24         20
    [28,]      40      20         20
    [29,]      36      16         20
    [30,]      32      12         20
    """

    ta = aroon(ohlc)
    @test colnames(ta)  == [:up, :down, :osc]
    @test timestamp(ta) == timestamp(ohlc[25:end])
    @test isapprox(values(ta)[2, 1], 48)
    @test isapprox(values(ta)[2, 2], 28)
    @test isapprox(values(ta)[2, 3], 20)
end

@testset "roc" begin
    ta = roc(cl, 3)
    @test colnames(ta) == [:Close_roc_3]
    @test isapprox(values(ta)[1], -0.15133107021618722, atol=.01)
    @test isapprox(values(ta)[2], -0.02926829268292683, atol=.01)
    @test isapprox(values(ta)[3], -0.06009615384615385, atol=.01)
end

@testset "adx" begin
    ta = adx(ohlc)
    @test colnames(ta) == [:adx, :dx, :di_plus, :di_minus]
    @test isapprox(values(ta)[1, 1], 10.5998, atol=.01)
    @test isapprox(values(ta)[1, 2], 0.3916,  atol=.01)
    @test isapprox(values(ta)[1, 3], 23.6226, atol=.01)
    @test isapprox(values(ta)[1, 4], 23.4383, atol=.01)
    @test timestamp(ta)[1] == Date(2000, 2, 10)
end

@testset "stochasticoscillator" begin
    """
    Quote from TTR
    > stoch(x[, c("High", "Low", "Close")], maType=SMA)
                fastK      fastD      slowD
    [18,] 0.671428571 0.69466667 0.67441270
    [19,] 0.432000000 0.59342857 0.64901587
    [20,] 0.492857143 0.53209524 0.60673016
    [21,] 0.392857143 0.43923810 0.52158730
    [22,] 0.217586207 0.36776683 0.44636672
    [23,] 0.326296296 0.31224655 0.37308382
    [24,] 0.500000000 0.34796083 0.34265807
    [25,] 0.724444444 0.51691358 0.39237365
    [26,] 0.754814815 0.65975309 0.50820917
    [27,] 0.801061008 0.76010676 0.64559114
    [28,] 0.839964633 0.79861349 0.73949111
    """

    ta = stochasticoscillator(ohlc)
    @test colnames(ta)  == [:fast_k, :fast_d, :slow_d]
    @test timestamp(ta) == timestamp(ohlc[18:end])
    @test isapprox(values(ta)[1, 1], 67.142857, atol=.01)
    @test isapprox(values(ta)[1, 2], 69.466667, atol=.01)
    @test isapprox(values(ta)[1, 3], 67.441270, atol=.01)
end

end  # @testset "Momentum"


using Plexifications
using Base.Test: @test, @test_throws

@test_throws ErrorException begin
    inv(TriComplex(0, 0.0))
end

@test begin
    iszerodivisor(TriComplex(0,1,1,0))
end

@test begin
    !iszerodivisor(TriComplex(1,2,3,4))
end

@test begin
    a = Plexifications.randomBigFloat()
    isreal(TriComplex(a))
end

@test begin
    a = 1
    !isreal(TriComplex(a, a))
end

@test begin
    z = random(TriComplex{Int})
    z == +(z)
end

@test begin
    a = Plexifications.randomBigFloat()
    real(TriComplex(a)) == a
end

@test begin
    io = IOBuffer()
    show(io, TriComplex(1, 2, 3, 4, 5, 6, 7, 8))
    l = String(take!(io))
    r = "[1: 1, H: 2, G: 3, GH: 4, F: 5, FH: 6, FG: 7, FGH: 8]"
    l == r
end

@test begin
    x = random(TriComplex{BigFloat})
    y = random(TriComplex{BigFloat})
    x != y
end

@test begin
    a = Plexifications.randomBigInt()
    z = random(TriComplex{BigInt})
    l = a + z
    r = z + a
    l == r
end

@test begin
    a = Plexifications.randomBigInt()
    z = random(TriComplex{BigInt})
    l = a - z
    r = -(z - a)
    l == r
end

@test begin
    x = random(TriComplex{BigInt})
    y = random(TriComplex{BigInt})
    l = x + y
    r = y + x
    l == r
end

@test begin
    x = random(BiComplex{BigInt})
    y = random(TriComplex{BigInt})
    l = x + y
    r = y + x
    l == r
end

@test begin
    x = random(TriComplex{BigInt})
    y = random(TriComplex{BigInt})
    l = x * y
    r = y * x
    l == r
end

@test begin
    x = random(BiComplex{BigInt})
    y = random(TriComplex{BigInt})
    l = x * y
    r = y * x
    l == r
end

@test begin
    x = random(TriComplex{BigInt})
    y = random(TriComplex{BigInt})
    l = x - y
    r = -(y - x)
    l == r
end

@test begin
    x = random(BiComplex{BigInt})
    y = random(TriComplex{BigInt})
    l = x - y
    r = -(y - x)
    l == r
end

@test begin
    x = random(TriComplex{BigInt})
    y = random(TriComplex{BigInt})
    z = random(TriComplex{BigInt})
    l = (x + y) + z
    r = x + (y + z)
    l == r
end

@test begin
    x = random(TriComplex{BigInt})
    y = random(TriComplex{BigInt})
    z = random(TriComplex{BigInt})
    l = (x * y) * z
    r = x * (y * z)
    l == r
end

@test begin
    x = random(TriComplex{BigInt})
    barH(barH(x)) == x
end

@test begin
    x = random(TriComplex{BigInt})
    barG(barG(x)) == x
end

@test begin
    x = random(TriComplex{BigInt})
    barF(barF(x)) == x
end

@test begin
    x = random(TriComplex{BigInt})
    barG(barH(x)) == barH(barG(x))
end

@test begin
    x = random(TriComplex{BigInt})
    barF(barH(x)) == barH(barF(x))
end

@test begin
    x = random(TriComplex{BigInt})
    barG(barF(x)) == barF(barG(x))
end

@test begin
    x = random(TriComplex{Rational{BigInt}})
    inv(inv(x)) == x
end

@test begin
    x = random(TriComplex{BigInt})
    y = random(TriComplex{BigInt})
    l = barH(x * y)
    r = barH(x) * barH(y)
    l == r
end

@test begin
    x = random(TriComplex{BigInt})
    y = random(TriComplex{BigInt})
    l = barH(x + y)
    r = barH(x) + barH(y)
    l == r
end

@test begin
    x = random(TriComplex{BigInt})
    y = random(TriComplex{BigInt})
    l = barH(x - y)
    r = barH(x) - barH(y)
    l == r
end

@test begin
    x = random(TriComplex{BigInt})
    y = random(TriComplex{BigInt})
    l = barG(x * y)
    r = barG(x) * barG(y)
    l == r
end

@test begin
    x = random(TriComplex{BigInt})
    y = random(TriComplex{BigInt})
    l = barG(x + y)
    r = barG(x) + barG(y)
    l == r
end

@test begin
    x = random(TriComplex{BigInt})
    y = random(TriComplex{BigInt})
    l = barG(x - y)
    r = barG(x) - barG(y)
    l == r
end

@test begin
    x = random(TriComplex{BigInt})
    y = random(TriComplex{BigInt})
    l = barF(x * y)
    r = barF(x) * barF(y)
    l == r
end

@test begin
    x = random(TriComplex{BigInt})
    y = random(TriComplex{BigInt})
    l = barF(x + y)
    r = barF(x) + barF(y)
    l == r
end

@test begin
    x = random(TriComplex{BigInt})
    y = random(TriComplex{BigInt})
    l = barF(x - y)
    r = barF(x) - barF(y)
    l == r
end

@test begin
    x = random(TriComplex{Rational{BigInt}})
    y = random(TriComplex{Rational{BigInt}})
    l = inv(x * y)
    r = inv(y) * inv(x)
    l == r
end

@test begin
    x = random(TriComplex{Rational{BigInt}})
    y = random(TriComplex{Rational{BigInt}})
    l = inv(x / y)
    r = y / x
    l == r
end

@test begin
    x = random(TriComplex{Rational{BigInt}})
    y = random(TriComplex{Rational{BigInt}})
    l = inv(x \ y)
    r = y \ x
    l == r
end

@test begin
    x = random(TriComplex{BigInt})
    y = random(TriComplex{BigInt})
    z = random(TriComplex{BigInt})
    l = x * (y + z)
    r = (x * y) + (x * z)
    l == r
end

@test begin
    a = Plexifications.randomBigInt()
    x = random(TriComplex{BigInt})
    y = random(TriComplex{BigInt})
    l = a * (x + y)
    r = (a * x) + (a * y)
    l == r
end

@test begin
    a = Plexifications.randomBigInt()
    x = random(TriComplex{BigInt})
    l = a * x
    r = x * a
    l == r
end

@test begin
    x = random(TriComplex{BigInt})
    y = random(TriComplex{BigInt})
    z = random(TriComplex{BigInt})
    l = (x + y) * z
    r = (x * z) + (y * z)
    l == r
end

@test begin
    x = random(TriComplex{BigInt})
    y = random(TriComplex{BigInt})
    z = random(TriComplex{BigInt})
    l = x * (y - z)
    r = (x * y) - (x * z)
    l == r
end

@test begin
    x = random(TriComplex{BigInt})
    y = random(TriComplex{BigInt})
    z = random(TriComplex{BigInt})
    l = (x - y) * z
    r = (x * z) - (y * z)
    l == r
end

@test begin
    x = random(TriComplex{BigInt})
    norm(x) > big(0)
end

@test begin
    x = random(TriComplex{BigInt})
    y = random(TriComplex{BigInt})
    l = norm(x * y)
    r = norm(x) * norm(y)
    l == r
end

@test_throws ErrorException begin
    inv(TriComplex(0, 0.0))
end

@test_throws ErrorException begin
    random(TriComplex{Int}) / 0
end

@test_throws ErrorException begin
    0 \ random(TriComplex{Int})
end

@test begin
    a = Plexifications.randomBigRational()
    z = random(TriComplex{Rational{BigInt}})
    l = z / a
    r = a \ z
    l == r
end

@test begin
    a = Plexifications.randomBigRational()
    z = random(TriComplex{Rational{BigInt}})
    l = a / z
    r = z \ a
    l == r
end

@test begin
    z = random(TriComplex{BigInt})
    z == z + zero(TriComplex{BigInt})
end

@test begin
    z = random(TriComplex{BigInt})
    z == z * one(TriComplex{BigInt})
end

@test begin
    z = random(TriComplex{BigInt})
    z == z + zero(z)
end

@test begin
    z = random(TriComplex{BigInt})
    z == z * one(z)
end

@test begin
    x = random(TriComplex{BigInt})
    y = random(TriComplex{BigInt})
    l = commutator(x, y)
    r = zero(TriComplex{BigInt})
    l == r
end

@test begin
    x = random(TriComplex{BigInt})
    y = random(TriComplex{BigInt})
    l = alternatorL(x, y)
    r = zero(TriComplex{BigInt})
    l == r
end

@test begin
    x = random(TriComplex{BigInt})
    y = random(TriComplex{BigInt})
    l = alternatorR(x, y)
    r = zero(TriComplex{BigInt})
    l == r
end

@test begin
    x = random(TriComplex{BigInt})
    y = random(TriComplex{BigInt})
    l = flexator(x, y)
    r = zero(TriComplex{BigInt})
    l == r
end

@test begin
    x = random(TriComplex{BigInt})
    y = random(TriComplex{BigInt})
    z = random(TriComplex{BigInt})
    l = associator(x, y, z)
    r = zero(TriComplex{BigInt})
    l == r
end

@test begin
    x = random(TriComplex{BigInt})
    y = random(TriComplex{BigInt})
    z = random(TriComplex{BigInt})
    l = jacobiator(x, y, z)
    r = 3 * x * y * z
    l == r
end

@test begin
    l = TriComplex(Cmplex(1,2))
    r = TriComplex(1,2)
    l == r
end

@test begin
    l = TriComplex(BiComplex(1,2))
    r = TriComplex(1,2,0)
    l == r
end

@test begin
    l = TriComplex(1, 2, 3.0)
    r = TriComplex(1.0, 2.0, 3.0)
    l == r
end

@test begin
    l = TriComplex(1, 2, 3, 4.0)
    r = TriComplex(1.0, 2.0, 3.0, 4.0)
    l == r
end

@test begin
    l = TriComplex(1, 2, 3, 4, 5.0)
    r = TriComplex(1.0, 2.0, 3.0, 4.0, 5.0)
    l == r
end

@test begin
    l = TriComplex(1, 2, 3, 4, 5, 6.0)
    r = TriComplex(1.0, 2.0, 3.0, 4.0, 5.0, 6.0)
    l == r
end

@test begin
    l = TriComplex(1, 2, 3, 4, 5, 6, 7.0)
    r = TriComplex(1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0)
    l == r
end

@test begin
    l = TriComplex(1, 2, 3, 4, 5, 6, 7, 8.0)
    r = TriComplex(1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0)
    l == r
end


@testset "kernelproduct" begin
    rng = MersenneTwister(123456)
    x = rand(rng)*2
    v1 = rand(rng, 3)
    v2 = rand(rng, 3)

    k1 = LinearKernel()
    k2 = SqExponentialKernel()
    k3 = RationalQuadraticKernel()
    X = rand(rng, 2,2)

    k = KernelProduct(k1,k2)
    ks1 = 2.0*k1
    ks2 = 0.5*k2
    @test length(k) == 2
    @test string(k) == (
        "Product of 2 kernels:\n\tLinear Kernel (c = 0.0)\n\tSquared " *
        "Exponential Kernel"
    )
    @test k(v1, v2) == (k1 * k2)(v1, v2)
    @test (k * k3)(v1,v2) ≈ (k3 * k)(v1, v2)
    @test (k1 * k2)(v1, v2) == KernelProduct(k1, k2)(v1, v2)
    @test (k * ks1)(v1, v2) ≈ (ks1 * k)(v1, v2)
    @test (k * k)(v1, v2) == KernelProduct([k1, k2, k1, k2])(v1, v2)
    @test KernelProduct([k1, k2]) == KernelProduct((k1, k2)) == k1 * k2

    @test (KernelProduct([k1, k2]) * KernelProduct([k2, k1])).kernels == [k1, k2, k2, k1]
    @test (KernelProduct([k1, k2]) * k3).kernels == [k1, k2, k3]
    @test (k3 * KernelProduct([k1, k2])).kernels == [k3, k1, k2]

    @test (KernelProduct((k1, k2)) * KernelProduct((k2, k1))).kernels == (k1, k2, k2, k1)
    @test (KernelProduct((k1, k2)) * k3).kernels == (k1, k2, k3)
    @test (k3 * KernelProduct((k1, k2))).kernels == (k3, k1, k2)

    # Standardised tests.
    TestUtils.test_interface(k, Float64)
    test_ADs(
        x->SqExponentialKernel() * LinearKernel(c= x[1]), rand(1);
        ADs = [:ForwardDiff, :ReverseDiff, :Zygote],
    )

    test_params(k1 * k2, (k1, k2))
end


### A Pluto.jl notebook ###
# v0.11.14

using Markdown
using InteractiveUtils

# ╔═╡ 06b847e0-fcb2-11ea-1880-0543a2db1b1c
begin
	using Dierckx
	using Plots
end

# ╔═╡ 24e001e0-fcb2-11ea-0559-1b3da141b5a3
rAti = [0, 10, 20, 50, 100, 200, 500, 1e3, 2e3, 5e3, 1e4]

# ╔═╡ 9670e220-fcb2-11ea-37be-4f9bfa6437b1
zAti = [sin.(r/1e3) for r ∈ rAti]

# ╔═╡ a228e9f0-fcb2-11ea-1cd0-a7ef10379149
plot(rAti, zAti)

# ╔═╡ 2d598b10-fcb3-11ea-1579-4709585b2f96
# itpAti = LinearInterpolation(rAti, zAti)
splAti = Spline1D(rAti, zAti)

# ╔═╡ de70a04e-fcb3-11ea-1e68-3be191449dfb
# zAtiFcn(r) = itpAti(r)
zAtiFcn(r) = splAti(r)

# ╔═╡ 42fec5be-fcb3-11ea-2d2b-7df1958c48cf
begin
	r = range(0, rAti[end], length = 101)
	scatter(r, zAtiFcn)
	plot!(rAti, zAti)
end

# ╔═╡ fae441b0-fcb3-11ea-2baf-09505779f9a3
begin
	function InterpolatingFunction(rng, val)
		Itp = LinearInterpolation(rng, val)
		return ItpFcn(r) = Itp(r)
	end

	function InterpolatingFunction(rng, dpt, val)
		Itp = LinearInterpolation((dpt, rng), val)
		return ItpFcn(r, z) = Itp(z, r)
	end
end

# ╔═╡ 4ac0e5d0-fcb4-11ea-1477-0be39085827b
zAtiFcn2 = InterpolatingFunction(rAti, zAti)

# ╔═╡ 5cf89810-fcb4-11ea-194b-235bb415165c
scatter(r, zAtiFcn2)

# ╔═╡ 66914250-fcb4-11ea-20e3-07319a79a510
begin
	rc = rAti
	zc = [0, 10, 20, 100, 500, 1e3]
# 	c = 1550 .+ sin.(rc'/1e3) .+ sin.(zc/1e3)
	c = [1550 + sin(r/1e3) + sin(z/1e3) for z ∈ zc, r ∈ rc]
end

# ╔═╡ b45b8130-fcb4-11ea-0801-bd0afb438aca
heatmap(rc, zc, c)

# ╔═╡ cd252d60-fcb4-11ea-3112-d7f36a0f4c04
begin
	cItp = LinearInterpolation((zc, rc), c)
	cFcn(r, z) = cItp(z, r)
end

# ╔═╡ f3661d90-fcb4-11ea-0352-a5284b9dc242
begin
	z = range(0, zc[end], length = 101)
	heatmap(r, z, cFcn)
end

# ╔═╡ 28f73070-fcb5-11ea-0148-91f62d5ef3d7
begin
	zInd = 6
	zVal = zc[zInd]
	scatter(rc, c[zInd, :])
	plot!(r, r -> cFcn(r, zVal))
end

# ╔═╡ 5472ef00-fcb5-11ea-0640-672ede9c99b3
cFcn2 = InterpolatingFunction(rc, zc, c)

# ╔═╡ bd6c4be0-fcb6-11ea-1343-c1584a8f72d9
heatmap(r, z, cFcn2)

# ╔═╡ 984269a2-fcb9-11ea-04f1-0927fe41dfa3
cMat = vcat(hcat(0, transpose(rc)), hcat(zc, c))

# ╔═╡ 118e5a80-fcba-11ea-3a27-671c712ec008
begin
	rVec = [r for r ∈ cMat[1, 2:end]]
	zVec = [z for z ∈ cMat[2:end, 1]]
	cVec = cMat[2:end, 2:end]
end

# ╔═╡ b9c1e4ae-fcba-11ea-2d72-4778c0a74263
size(rVec), size(zVec), size(cVec)

# ╔═╡ 6adf09e0-fcba-11ea-2d8b-e5bbbece2d99
cFcn3 = InterpolatingFunction(rVec, zVec, cVec)

# ╔═╡ 4648ee10-fcbb-11ea-2247-735e9cb6e01d
heatmap(r, z, cFcn3)

# ╔═╡ Cell order:
# ╠═06b847e0-fcb2-11ea-1880-0543a2db1b1c
# ╠═24e001e0-fcb2-11ea-0559-1b3da141b5a3
# ╠═9670e220-fcb2-11ea-37be-4f9bfa6437b1
# ╠═a228e9f0-fcb2-11ea-1cd0-a7ef10379149
# ╠═2d598b10-fcb3-11ea-1579-4709585b2f96
# ╠═de70a04e-fcb3-11ea-1e68-3be191449dfb
# ╠═42fec5be-fcb3-11ea-2d2b-7df1958c48cf
# ╠═fae441b0-fcb3-11ea-2baf-09505779f9a3
# ╠═4ac0e5d0-fcb4-11ea-1477-0be39085827b
# ╠═5cf89810-fcb4-11ea-194b-235bb415165c
# ╠═66914250-fcb4-11ea-20e3-07319a79a510
# ╠═b45b8130-fcb4-11ea-0801-bd0afb438aca
# ╠═cd252d60-fcb4-11ea-3112-d7f36a0f4c04
# ╠═f3661d90-fcb4-11ea-0352-a5284b9dc242
# ╠═28f73070-fcb5-11ea-0148-91f62d5ef3d7
# ╠═5472ef00-fcb5-11ea-0640-672ede9c99b3
# ╠═bd6c4be0-fcb6-11ea-1343-c1584a8f72d9
# ╠═984269a2-fcb9-11ea-04f1-0927fe41dfa3
# ╠═118e5a80-fcba-11ea-3a27-671c712ec008
# ╠═b9c1e4ae-fcba-11ea-2d72-4778c0a74263
# ╠═6adf09e0-fcba-11ea-2d8b-e5bbbece2d99
# ╠═4648ee10-fcbb-11ea-2247-735e9cb6e01d


using GLMakie

xs = range(0, 10, length = 25)
ys = range(0, 15, length = 25)

fig = Figure()

sl = Slider(fig[2, 1], range = 0:0.1:3, startvalue = 1)

zs = lift(sl.value) do s
    convert(Float64,s).*[cos(x) * sin(y) for x in xs, y in ys]
end

titleString = lift(sl.value) do s
    "Step: "*string(s)
end

ax = Axis(fig[1, 1], title = titleString)

hm=heatmap!(xs, ys, zs,colormap = Reverse(:Spectral))
Colorbar(fig[1, 2],hm,label = "Values")
fig


function maxdiff(x::AbstractArray, y::AbstractArray)
    return mapreduce((a, b) -> abs(a - b), max, x, y)
end


########################
# Lorenz System  (MIMO)#
# by VSSM              #
########################

###########################################
# Controll parameters for VSSM controller #
###########################################

K=500
w=1

#############################################
# Parameters to design a Nominal Trajectory #
#############################################

A₁=2
ω₁=0.5
A₂=3
ω₂=0.7
A₃=1
ω₃=1

#################
# Time variable #
#################

# cycle time in seconds
δt=1e-3
# The lenght of the simulation in seconds
LONG=Int(2e5)

Λ=12

##########################
# Exact model parameters #
##########################
σₑ=5
βₑ=8/3
ρₑ=40

############################
# Approx. model parameters #
############################
σₐ=5
βₐ=8/3
ρₐ=40

#Arrays

xN=zeros(LONG)
xN_p=zeros(LONG)
yN=zeros(LONG)
yN_p=zeros(LONG)
zN=zeros(LONG)
zN_p=zeros(LONG)

x_p_des=zeros(LONG)
y_p_des=zeros(LONG)
z_p_des=zeros(LONG)
u=zeros(LONG)
t=zeros(LONG)

S_x=zeros(LONG)
S_p_x=zeros(LONG)
S_y=zeros(LONG)
S_p_y=zeros(LONG)
S_z=zeros(LONG)
S_p_z=zeros(LONG)

x=zeros(LONG)
y=zeros(LONG)
z=zeros(LONG)
x_p=zeros(LONG)
y_p=zeros(LONG)
z_p=zeros(LONG)

#initial conditions
hint_x=0
hint_y=0
hint_z=0

x[1]=A₁*sin(ω₁*δt)
y[1]=A₂*sin(ω₂*δt)
z[1]=A₃*sin(ω₃*δt)

##############
# SIMULATION #
##############

for i=1:LONG-1 #i counts the cycles of the simulation
  t[i]=δt*i
  #create Nominal Trajectory
  xN[i]=A₁*sin(ω₁*t[i])
  xN_p[i]=A₁*ω₁*cos(ω₁*t[i])
  yN[i]=A₂*sin(ω₂*t[i])
  yN_p[i]=A₂*ω₂*cos(ω₂*t[i])
  zN[i]=A₃*sin(ω₃*t[i])
  zN_p[i]=A₃*ω₃*cos(ω₃*t[i])
  # Compute  error (h) and error integral (hint) and the first derivative of the error (h_p)
  h_x=xN[i]-x[i]
  h_y=yN[i]-y[i]
  h_z=zN[i]-z[i]

  hint_x=hint_x+δt*h_x
  hint_y=hint_y+δt*h_y
  hint_z=hint_z+δt*h_z

  h_p_x=xN_p[i]-x_p[i]
  h_p_y=yN_p[i]-y_p[i]
  h_p_z=zN_p[i]-z_p[i]

  #Error metric (S), and the first derivative of the Error metric (S_p)
  S_x[i]=Λ*hint_x+h_x
  S_p_x[i]=Λ*h_x+h_p_x
  S_y[i]=Λ*hint_y+h_y
  S_p_y[i]=Λ*h_y+h_p_y
  S_z[i]=Λ*hint_z+h_z
  S_p_z[i]=Λ*h_z+h_p_z
  #Kinematic Block (PID type feedback)
  x_p_des[i]=xN_p[i]+Λ*h_x+K*tanh(S_x[i]/w)
  y_p_des[i]=yN_p[i]+Λ*h_y+K*tanh(S_y[i]/w)
  z_p_des[i]=zN_p[i]+Λ*h_z+K*tanh(S_z[i]/w)
  # Create the control signal
  u[i]=x_p_des[i]-σₐ*(y[i]-x[i])
  # Get the system "reaction" (the exact model of the system)
  x_p[i]=σₑ*(y[i]-x[i])+u[i]
  y_p[i]=x[i]*(ρₑ-z[i])-y[i]
  z_p[i]=x[i]*y[i]-βₑ*z[i]
  # Integrate back to x, y, z
  x[i+1]=x[i]+δt*x_p[i]
  y[i+1]=y[i]+δt*y_p[i]
  z[i+1]=z[i]+δt*z_p[i]
end #for

using PyPlot

figure(1)
grid("on")
title("Nominal and Realized trajectorys")
xlabel(L"Time \enspace [s]")
ylabel(L"q \enspace [m]")
plot(t[1:LONG-1],xN[1:LONG-1],color="#7684FF")
#plot(t[3:LONG-1],yN[3:LONG-1],color="#2A40FF")
#plot(t[3:LONG-1],zN[3:LONG-1],color="#3B427F")

plot(t[1:LONG-1],x[1:LONG-1],color="#FFAA41","r--")
#plot(t[3:LONG-1],y[3:LONG-1],color="#FF9F28","r--")
#plot(t[3:LONG-1],z[3:LONG-1],color="#B2690E","r--")

figure(2)
grid("on")
title("Tracking Error")
xlabel(L"Time \enspace [s]")
ylabel(L"error \enspace [m]")
plot(t[1:LONG-1],xN[1:LONG-1]-x[1:LONG-1],color="red")

figure(3)
grid("on")
title("Phase Space of Error Metric (S_p vs S)")
plot(S_x[1:LONG-1],S_p_x[1:LONG-1],color="blue")

figure(4)
grid("on")
title("Phase Space of motion (x_p vs x)")
plot(xN[1:LONG-1],xN_p[1:LONG-1],color="red")
plot(x[1:LONG-1],x_p[1:LONG-1],color="blue","r--")


"""
```
orientations = corner_orientations(img)
orientations = corner_orientations(img, corners)
orientations = corner_orientations(img, corners, kernel)
```

Returns the orientations of corner patches in an image. The orientation of a corner patch
is denoted by the orientation of the vector between intensity centroid and the corner.
The intensity centroid can be calculated as `C = (m01/m00, m10/m00)` where mpq is defined as -

    `mpq = (x^p)(y^q)I(y, x) for each p, q in the corner patch`

The kernel used for the patch can be given through the `kernel` argument. The default kernel used is a gaussian kernel of size `5x5`.
"""
function corner_orientations(img::AbstractArray{T, 2}, corners::Keypoints, kernel::Array{K, 2}) where {T<:Gray, K<:Real}
    h, w = size(kernel)
    pre_y = ceil(Int, (h - 1) / 2)
    pre_x = ceil(Int, (w - 1) / 2)
    post_y = floor(Int, (h - 1) / 2)
       post_x = floor(Int, (w - 1) / 2)
       img_padded = parent(padarray(img, Fill(0, (pre_y, pre_x), (post_y, post_x))))
    orientations = Float64[]
    for c in corners
        m10 = zero(T)
        m01 = zero(T)
        for i in 1:w
            col_sum = zero(T)
            for j in 1:h
                pixel = img_padded[c + CartesianIndex(j - 1, i - 1)] * kernel[j, i]
                m01 += pixel * (j - pre_y - 1)
                col_sum += pixel
            end
            m10 += col_sum * (i - pre_x - 1)
        end
        push!(orientations, atan(Float64(m01), Float64(m10)))
    end
    orientations
end

corner_orientations(img::AbstractArray{T, 2}, corners::Keypoints, kernel::Array{K, 2}) where {T, K<:Real} = corner_orientations(Gray.(img), corners, kernel)

function corner_orientations(img::AbstractArray, kernel::Array{K, 2}) where K<:Real
    corners = imcorner(img)
    corner_indexes = Keypoints(corners)
    corner_orientations(img, corner_indexes, kernel)
end

function corner_orientations(img::AbstractArray)
    corners = imcorner(img)
    corner_indexes = Keypoints(corners)
    kernel = Kernel.gaussian((2,2), (5, 5))
    kernel /= maxfinite(kernel)
    corner_orientations(img, corner_indexes, parent(kernel))
end

function corner_orientations(img::AbstractArray, corners::Keypoints)
    kernel = Kernel.gaussian((2,2), (5, 5))
    kernel /= maxfinite(kernel)
    corner_orientations(img, corners, parent(kernel))
end


using PyCall
BayesOpt = pyimport("bayes_opt")
BayesOpt.BayesianOptimization
so = pyimport("scipy.optimize")

function black_box_function(x)
    # Function with unknown internals we wish to maximize.
    return x[1] ^ 2 + x[2] ^2
end

so.minimize(black_box_function,[1.,1.],method="Nelder-Mead")

bds = Dict('x' => (2, 4), 'y' => (-3, 3))
bds = Dict("x"=> (2,2))
bds = Dict(-2,2)
rates = []
Dict(bds)
Dict(i => bd for (i,bd) in enumerate(bds))

optimizer = BayesOpt.BayesianOptimization(
    f=black_box_function,
    pbounds=bds,
    verbose=2, # verbose = 1 prints only when a maximum is observed, verbose = 0 is silent
    random_state=1)

py"""
    from bayes_opt import BayesianOptimization

    def black_box_function(x, y):
        # Function with unknown internals we wish to maximize.
        return -x ** 2 - (y - 1) ** 2 + 1

    # Bounded region of parameter space
    pbounds = {'x': (2, 4), 'y': (-3, 3)}

    optimizer = BayesianOptimization(
        f=black_box_function,
        pbounds=pbounds,
        verbose=2, # verbose = 1 prints only when a maximum is observed, verbose = 0 is silent
        random_state=1,
    )

    optimizer.maximize(
        init_points=2,
        n_iter=3,
    )
"""
optimizer.maximize(
    init_points=2,
    n_iter=3)

py"""
import numpy as np

def sinpi(x):
    return np.sin(np.pi * x)
"""

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

skopt = pyimport("skopt")
noise_level = 0.1

function f(x,ct)
    x[1]^2 + x[2]^2 + ct
end

g = (x) -> f(x,ct)
ct = 1.0
# Plot f(x) + contours
x = range(-2, stop=2, length=400) #.reshape(-1, 1)

fx = [f(x_i) for x_i in x]
plot(x, fx, label="True (unknown)")
plt.fill(np.concatenate([x, x[::-1]]),
        np.concatenate(([fx_i - 1.9600 * noise_level for fx_i in fx],
                        [fx_i + 1.9600 * noise_level for fx_i in fx[::-1]])),
        alpha=.2, fc="r", ec="None")
plt.legend()
plt.grid()
plt.show()

res = skopt.gp_minimize(g,                  # the function to minimize
                  [(-2.0, 2.0),(-2.0, 2.0)],      # the bounds on each dimension of x
                  acq_func="EI",      # the acquisition function
                  n_calls=50,         # the number of evaluations of f
                  n_random_starts=5,  # the number of random initialization points
                  noise=0.,       # the noise level (optional)
                  random_state=1234)   # the random seed

res["x"]


"""
    generate_sample_gibbs_perfect(
        [rng::Random.AbstractRNG,]
        ising::Ising{T}
    )::Vector{T} where {T<:Int}

Generate an exact realization of the [`PRS.Ising`](@ref) model using a tailored implementation of the perfect Gibbs sampler of [FeGuYi19](@cite).
"""
function generate_sample_gibbs_perfect(
    rng::Random.AbstractRNG,
    ising::Ising{T}
)::Vector{T} where {T<:Int}

    n = LG.nv(ising.graph)
    R = Set{T}(1:n)
    state = Vector{T}(undef, n)
    generate_sample!(rng, state, R, ising)

    while !isempty(R)
        i = rand(rng, R)
        if bayes_filter(rng, ising, state, i, R)
            generate_sample_conditional!(rng, state, i, ising)
            delete!(R, i)
        else
            union!(R, LG.neighbors(ising.graph, i))
        end
    end

    return state
end

function generate_sample_gibbs_perfect(ising::Ising)
    return generate_sample_gibbs_perfect(Random.default_rng(), ising)
end

"""
    bayes_filter(
        rng::Random.AbstractRNG,
        ising::Ising{T},
        state::AbstractVector{T},
        i::T,
        R::Set{T}
    )::Bool

This function is used as a subroutine of [`PRS.generate_sample_gibbs_perfect`](@ref).
"""
function bayes_filter(
    rng::Random.AbstractRNG,
    ising::Ising{T},
    state::AbstractVector{T},
    i::T,
    R::Set{T}
)::Bool where {T<:Int}
    ∂i_R̄ = setdiff(LG.neighbors(ising.graph, i), R)
    isempty(∂i_R̄) && return true

    Xi_J = state[i] * ising.J
    log_acc_ratio = Xi_J * (-sign(Xi_J) * length(∂i_R̄) - sum(state[j] for j in ∂i_R̄))
    return log(rand(rng)) < log_acc_ratio
end