module Maths
import FiniteDifferences
import LinearAlgebra: Tridiagonal, mul!, ldiv!
import SpecialFunctions: erf, erfc, gamma
import StaticArrays: SVector
import Random: AbstractRNG, randn, GLOBAL_RNG
import FFTW
import Luna
import Luna.Utils: saveFFTwisdom, loadFFTwisdom
import Roots: fzero
import Dierckx
import Peaks

#= Pre-created finite difference methods for speed.
   Use (order+6)th order central finite differences with 2 adaptive steps up to
   11th order. Above that creates overflow errors, so we cap it. =#
FDMs = [FiniteDifferences.central_fdm(min(order+6, 11), order, adapt=1) for order=1:7]

"""
    derivative(f, x, order::Integer)

Calculate derivative of function f(x) at value x using finite differences
"""
function derivative(f, x, order::Integer)
    if order == 0
        return f(x)
    else
        scale = abs(x) > 0 ? x : 1.0
        FDMs[order](y->f(y*scale), x/scale)/scale^order
    end
end

"""
    gauss(x, σ; x0=0, power=2)

Gaussian or hypergaussian function over `x` with std dev `σ`, power in the exponent of `power`,
centred at `x0`.
"""
gauss(x, σ; x0=0, power=2) = exp(-1/2 * ((x-x0)/σ)^power)

"""
    gauss(x; fwhm, x0=0, power=2)

Gaussian or hypergaussian function over `x` with FWHM `fwhm`, power in the exponent of `power`,
centred at `x0`.
"""
gauss(x; x0=0, power=2, fwhm) = gauss(x, fwhm_to_σ(fwhm; power=power), x0 = x0, power=power)

fwhm_to_σ(fwhm; power=2) = fwhm / (2 * (2 * log(2))^(1 / power))

"""
    gaussnorm(σ; power=2)
    gaussnorm(;fwhm, power=2)

Norm of a Gaussian with standard deviation `σ` or FWHM `fwhm` and power in the exponent
`power`.
"""
gaussnorm(σ; power=2) = 2σ*2^(1/power)*gamma((power+1)/power)
gaussnorm(;fwhm, power=2) = gaussnorm(fwhm_to_σ(fwhm; power=power); power=power)

"""
    randgauss([rng=GLOBAL_RNG], μ, σ, args...)

Generate random numbers with normal distribution centred on `μ` with standard deviation `σ`.

Any additional `args` are passed onto `randn` and can be used to specify the output dimensions.
"""
randgauss(μ, σ, args...) = randgauss(GLOBAL_RNG, μ, σ, args...)
randgauss(rng::AbstractRNG, μ, σ, args...) = σ*randn(rng, args...) .+ μ

"""
    moment(x::Vector, y::Vector, n = 1)

Calculate the nth moment of the vector y.
"""
function moment(x::Vector, y::Vector, n=1)
    if length(x) ≠ length(y)
        throw(DomainError(x, "x and y must have same length"))
    end
    return sum(x.^n .* y) / sum(y)
end

"""
    moment(x::Vector, y, n=1; dim=1)

Calculate the nth moment of multi-dimensional array y along dimension dim.
"""
function moment(x::Vector, y, n=1; dim=1)
    if size(y, dim) ≠ length(x)
        throw(DomainError(y, "y must be of same length as x along dim"))
    end
    xshape = ones(Integer, ndims(y))
    xshape[dim] = length(x)
    return sum(reshape(x, Tuple(xshape)).^n .* y, dims=dim) ./ sum(y, dims=dim)
end

"""
    rms_width(x::Vector, y; dim=1)

RMS width of distribution y on axis x
"""
function rms_width(x::Vector, y::Vector; dim=1)
    return sqrt(moment(x, y, 2) - moment(x, y, 1)^2)
end

function rms_width(x::Vector, y; dim=1)
    return sqrt.(moment(x, y, 2, dim=dim) - moment(x, y, 1, dim=dim).^2)
end

"""
    fwhm(x, y; method=:linear, baseline=false, minmax=:min, level=0.5)

Calculate the full width at half maximum (FWHM) of `y` on the axis `x`.

`method` can be `:spline` or `:nearest`. `:spline` uses a [`CSpline`](@ref), whereas
`:nearest` finds the closest values either side of the crossing point and interpolates linearly.

If `baseline` is true, the width is not taken at
`level * maximum(y)`, but at half of the span of `y`, `level * (maximum(y) - minimum(y))`.

`minmax` determines whether the FWHM is taken at the narrowest (`:min`) or the widest (`:max`)
point of y.

The default `level=0.5` requests the full width at half maximum. Setting `level` to something
different computes the corresponding width. E.g. `level=0.1` for the 10% width.
"""
function fwhm(x, y::AbstractVector; kwargs...)
    left, right = level_xings(x, y; kwargs...)
    return abs(right-left)
end

function fwhm(x, y::AbstractArray; dim=1, kwargs...)
    idxlo = CartesianIndices(size(y)[1:dim-1])
    idxhi = CartesianIndices(size(y)[dim+1:end])
    outshape = collect(size(y))
    outshape[dim] = 1
    out = zeros(eltype(y), Tuple(outshape))
    for hi in idxhi
        for lo in idxlo
            out[lo, 1, hi] = fwhm(x, y[lo, :, hi]; kwargs...)
        end
    end
    out
end

"""
    level_xings(x, y; method=:linear, baseline=false, minmax=:min)

Find crossings of the curve `y` on the axis `x` with the value `level * maximum(y)`.

`method` can be `:spline`, `:linear` or `:nearest`. `:spline` uses a [`CSpline`](@ref).
`:linear` finds the closest values either side of the crossing point and interpolates linearly.
`:nearest` just uses the closest values either side of the crossing point.

If `baseline` is true, the width is not taken at
`level * maximum(y)`, but at half of the span of `y`, `level * (maximum(y) - minimum(y))`.

`minmax` determines whether the crossings are taken at the narrowest (`:min`) or the widest (`:max`)
point of y.
"""
function level_xings(x::AbstractVector, y::AbstractVector;
                     method=:linear, baseline=false, minmax=:min, level=0.5)
    minmax in (:min, :max) || error("minmax has to be :min or :max")
    if baseline
        val = minimum(y) + level*(maximum(y) - minimum(y))
    else
        val = level*maximum(y)
    end
    if !(method in (:spline, :linear, :nearest))
        error("Unknown FWHM method $method")
    end
    maxidx = argmax(y)
    xmax = x[maxidx]
    if method in (:nearest, :linear)
        try
            if minmax == :min
                lefti = findlast((x .< xmax) .& (y .< val))
                righti = findfirst((x .> xmax) .& (y .< val))
                (method == :nearest) && return (x[lefti], x[righti])

                left = linterpx(x[lefti], x[lefti+1], y[lefti], y[lefti+1], val)
                right = linterpx(x[righti-1], x[righti], y[righti-1], y[righti], val)
            else
                lefti = findfirst((x .< xmax) .& (y .> val))
                righti = findlast((x .> xmax) .& (y .> val))
                (method == :nearest) && return (x[lefti], x[righti])

                left = linterpx(x[lefti-1], x[lefti], y[lefti-1], y[lefti], val)
                right = linterpx(x[righti], x[righti+1], y[righti], y[righti+1], val)
            end
            return left, right
        catch
            return NaN, NaN
        end
    elseif method == :spline
        #spline method
        try
            spl = BSpline(x, y .- val)
            r = roots(spl; maxn=10000)
            rleft = r[r .< xmax]
            rright = r[r .> xmax]
            if minmax == :min
                return maximum(rleft), minimum(rright)
            else
                return minimum(rleft), maximum(rright)
            end
        catch e
            return NaN, NaN
        end
    else
        error("Unknown level_xings method $method")
    end
end

"""
    linterpx(x1, x2, y1, y2, val)

Given two points on a straight line, `(x1, y1)` and `(x2, y2)`, calculate the value of `x`
at which this straight line intercepts with `val`.

# Examples
```jldoctest
julia> x1 = 0; x2 = 1; y1 = 0; y2 = 2; # y = 2x
julia> linterpx(x1, x2, y1, y2, 0.5)
0.25
```
"""
function linterpx(x1, x2, y1, y2, val)
    slope = (y2-y1)/(x2-x1)
    icpt = y1 - slope*x1
    (val-icpt)/slope
end

"""
    hwhm(f, x0=0; direction=:fwd)

Find the value `x` where the function `f(x)` drops to half of its maximum, which is located
at `x0`. For `direction==:fwd`, search in the region `x > x0`, for :bwd, search in `x < x0`.
"""
function hwhm(f, x0=0; direction=:fwd)
    m = f(x0)
    fhalf(x) = f(x) - 0.5*m
    if direction == :fwd
        xhw = fzero(fhalf, x0, Inf)
    elseif direction == :bwd
        xhw = fzero(fhalf, -Inf, x0)
    end
    return abs(xhw - x0)
end

"""
    cumtrapz!([out, ] y, x; dim=1)

Trapezoidal integration for multi-dimensional arrays or vectors, in-place or with output array.

If `out` is omitted, `y` is integrated in place. Otherwise the result is placed into `out`.

`x` can be an array (the x axis) or a number (the x axis spacing).
"""
function cumtrapz! end

function cumtrapz!(y, x; dim=1)
    idxlo = CartesianIndices(size(y)[1:dim-1])
    idxhi = CartesianIndices(size(y)[dim+1:end])
    _cumtrapz!(y, x, idxlo, idxhi)
end

"""
    _cumtrapz!([out, ] y, x, idxlo, idxhi)

Inner function for multi-dimensional `cumtrapz!` - uses 1-D routine internally
"""
function _cumtrapz!(y, x, idxlo, idxhi)
    for lo in idxlo
        for hi in idxhi
            cumtrapz!(view(y, lo, :, hi), x)
        end
    end
end

function cumtrapz!(y::T, x) where T <: Union{SubArray, Vector}
    tmp = y[1]
    y[1] = 0
    for i in 2:length(y)
        tmp2 = y[i]
        y[i] = y[i-1] + 1//2 * (tmp + tmp2) * _dx(x, i)
        tmp = tmp2
    end
end

function cumtrapz!(out, y, x; dim=1)
    idxlo = CartesianIndices(size(y)[1:dim-1])
    idxhi = CartesianIndices(size(y)[dim+1:end])
    _cumtrapz!(out, y, x, idxlo, idxhi)
end

function _cumtrapz!(out, y, x, idxlo, idxhi)
    for lo in idxlo
        for hi in idxhi
            cumtrapz!(view(out, lo, :, hi), view(y, lo, :, hi), x)
        end
    end
end

function cumtrapz!(out, y::Union{SubArray, Vector}, x)
    out[1] = 0
    for i in 2:length(y)
        out[i] = out[i-1]+ 1//2*(y[i-1] + y[i])*_dx(x, i)
    end
end

"""
    _dx(x, i)

Calculate the axis spacing at index `i` given an axis `x`. If `x` is a number, interpret this
as `δx` directly
"""
_dx(x, i) = x[i] - x[i-1]
_dx(δx::Number, i) = δx

"""
    cumtrapz(y, x; dim=1)

Calculate the cumulative trapezoidal integral of `y`.

`x` can be an array (the x axis) or a number (the x axis spacing).
"""
function cumtrapz(y, x; dim=1)
    out = similar(y)
    cumtrapz!(out, y, x; dim=dim)
    return out
end

"""
    normbymax(x, dims)

Return an array normalised by its maximum value
"""
function normbymax(x, dims)
    return x ./ maximum(x; dims=dims)
end

function normbymax(x)
    return x ./ maximum(x)
end

"""
    log10_norm(x)

Normalised log10 i.e. maximum of output is 0
"""
function log10_norm(x)
    return log10.(normbymax(x))
end

function log10_norm(x, dims)
    return log10.(normbymax(x, dims))
end

"""
    errfun_window(x, xmin, xmax, width)
    errfun_window(x, xmin, xmax, width_left, width_right)

Filtering window based on the error function. If `width_left` and `width_right` are given
separately, the smoothing width is different on each side.
"""
function errfun_window(x, xmin, xmax, width)
    return @. 0.5 * (erf((x - xmin) / width) + erfc((x - xmax) / width) - 1)
end

function errfun_window(x, xmin, xmax, width_left, width_right)
    return @. 0.5 * (erf((x - xmin) / width_left) + erfc((x - xmax) / width_right) - 1)
end

"""
    planck_taper(x, xmin, xmax, ε)

Planck taper window as defined in the paper (https://arxiv.org/pdf/1003.2939.pdf eq(7)).

#Arguments
-`xmin` : lower limit (window is 0 here)
-`xmax` : upper limit (window is 0 here)
-`ε` : fraction of window width over which to increase from 0 to 1
"""
function planck_taper(x, xmin, xmax, ε)
    x0 = (xmax + xmin) / 2
    xc = x .- x0
    X = (xmax - xmin)
    x1  = -X / 2
    x2 = -X / 2 * (1 - 2ε)
    x3 = X / 2 * (1 - 2ε)
    x4 = X / 2
    return _taper.(xc, x1, x2, x3, x4)
end

"""
    planck_taper(x, left0, left1, right1, right0)

Planck taper window, but finding the taper width by defining 4 points:
The window increases from 0 to 1 between `left0` and `left1`, and then drops again
to 0 between `right1` and `right0`.
"""
function planck_taper(x, left0, left1, right1, right0)
    x0 = (right0 + left0) / 2
    xc = x .- x0
    X = right0 - left0
    εleft = abs(left1 - left0) / X
    εright = abs(right0 - right1) / X
    x1  = -X / 2
    x2 = -X / 2 * (1 - 2εleft)
    x3 = X / 2 * (1 - 2εright)
    x4 = X / 2
    return _taper.(xc, x1, x2, x3, x4)
end

function _taper(xc, x1, x2, x3, x4)
    if xc < x1
        return zero(xc)
    elseif x1 < xc < x2
        z12 = (x2 - x1) / (xc - x1) + (x2 - x1) / (xc - x2)
        return 1 / (1 + exp(z12))
    elseif x2 <= xc <= x3
        return one(xc)
    elseif x3 < xc < x4
        z34 = (x3 - x4) / (xc - x3) + (x3 - x4) / (xc - x4)
        return 1 / (1 + exp(z34))
    else
        return zero(xc)
    end
end

"""
    hypergauss_window(x, xmin, xmax, power = 10)

Hypergaussian window function.
"""
function hypergauss_window(x, xmin, xmax, power = 10)
    fw = xmax - xmin
    x0 = (xmax + xmin) / 2
    return gauss(x, x0 = x0, fwhm = fw, power = power)
end

"""
    gabor(t, A, ts, fw)

Compute the Gabor transform (aka spectrogram or time-gated Fourier transform) of the vector
`A`, sampled on axis `t`, with windows centred at `ts` and a window FWHM of `fw`.
"""
function gabor(t::Vector, A::Vector, ts, fw)
    tmp = Array{eltype(A), 2}(undef, (length(t), length(ts)))
    for (ii, ti) in enumerate(ts)
        tmp[:, ii] = A .* Maths.gauss.(t; x0=ti, fwhm=fw)
    end
    _gaborFT(tmp)
end

_gaborFT(x::Array{T, 2}) where T <: Real = FFTW.rfft(x, 1)
_gaborFT(x::Array{T, 2}) where T <: Complex = FFTW.fft(x, 1)

"""
    hilbert(x; dim=1)

Compute the Hilbert transform, i.e. find the analytic signal from a real signal.
"""
function hilbert(x::Array{T,N}; dim = 1) where T <: Real where N
    xf = FFTW.fft(x, dim)
    n1 = size(xf, dim)÷2
    n2 = size(xf, dim)
    idxlo = CartesianIndices(size(xf)[1:dim - 1])
    idxhi = CartesianIndices(size(xf)[dim + 1:end])
    xf[idxlo, 2:n1, idxhi] .*= 2
    xf[idxlo, (n1+1):n2, idxhi] .= 0
    return FFTW.ifft(xf, dim)
end

"""
    plan_hilbert!(x; dim=1)

Pre-plan a Hilbert transform.

Returns a closure `hilbert!(out, x)` which places the Hilbert transform of `x` in `out`.
"""
function plan_hilbert!(x; dim=1)
    loadFFTwisdom()
    FT = FFTW.plan_fft(copy(x), dim, flags=Luna.settings["fftw_flag"])
    saveFFTwisdom()
    xf = Array{ComplexF64}(undef, size(FT))
    idxlo = CartesianIndices(size(xf)[1:dim - 1])
    idxhi = CartesianIndices(size(xf)[dim + 1:end])
    n1 = size(xf, dim)÷2
    n2 = size(xf, dim)
    xc = complex(x)
    function hilbert!(out, x)
        copyto!(xc, x)
        mul!(xf, FT, xc)
        xf[idxlo, 2:n1, idxhi] .*= 2
        xf[idxlo, (n1+1):n2, idxhi] .= 0
        ldiv!(out, FT, xf)
    end
    return hilbert!
end

"""
    plan_hilbert(x; dim=1)

Pre-plan a Hilbert transform.

Returns a closure `hilbert(x)` which returns the Hilbert transform of `x` without allocation.

!!! warning
    The closure returned always returns a reference to the same array buffer, which could lead
    to unexpected results if it is called from more than one location. To avoid this the array
    should either: (i) only be used in the same code segment; (ii) only be used transiently
    as part of a larger computation; (iii) be copied.
"""
function plan_hilbert(x; dim=1)
    out = complex(x)
    hilbert! = plan_hilbert!(x, dim=dim)
    function hilbert(x)
        hilbert!(out, x)
    end
    return hilbert
end

"""
    oversample(t, x; factor::Int=4, dim=1)

Oversample (smooth) an array by 0-padding in the frequency domain. Works for both real-valued
and complex-valued arrays.
"""
function oversample(t, x::Array{<:Real, N}; factor::Int=4, dim=1) where N
    if factor == 1
        return t, x
    end
    xf = FFTW.rfft(x, dim)

    len = size(xf, dim)
    newlen_t = factor * length(t)
    if iseven(newlen_t)
        newlen_ω = Int(newlen_t / 2 + 1)
    else
        newlen_ω = Int((newlen_t + 1) / 2)
    end
    δt = t[2] - t[1]
    δto = δt / factor
    Nto = collect(0:newlen_t - 1)
    to = t[1] .+ Nto .* δto

    shape = collect(size(xf))
    shape[dim] = newlen_ω
    xfo = Array{eltype(xf), N}(undef, Tuple(shape))
    fill!(xfo, 0.0)
    idxlo = CartesianIndices(size(xfo)[1:dim - 1])
    idxhi = CartesianIndices(size(xfo)[dim + 1:end])
    xfo[idxlo, 1:len, idxhi] .= factor .* xf
    return to, FFTW.irfft(xfo, newlen_t, dim)
end

function oversample(t, x::Array{<:Complex,N}; factor::Int=4, dim=1) where N
    if factor == 1
        return t, x
    end
    xf = FFTW.fftshift(FFTW.fft(x, dim), dim)

    len = size(xf, dim)
    newlen = factor * length(t)
    δt = t[2] - t[1]
    δto = δt / factor
    Nto = collect(0:newlen - 1)
    to = t[1] .+ Nto .* δto

    sidx  = (newlen - len)//2 + 1
    iseven(newlen) || (sidx -= 1//2)
    iseven(len) || (sidx += 1//2)
    startidx = Int(sidx)
    endidx = startidx+len-1

    shape = collect(size(xf))
    shape[dim] = newlen
    xfo = Array{eltype(xf), N}(undef, Tuple(shape))
    fill!(xfo, 0.0)
    idxlo = CartesianIndices(size(xfo)[1:dim - 1])
    idxhi = CartesianIndices(size(xfo)[dim + 1:end])
    xfo[idxlo, startidx:endidx, idxhi] .= factor .* xf
    return to, FFTW.ifft(FFTW.ifftshift(xfo, dim), dim)
end

"""
    aitken_accelerate(f, x0; n0 = 0, rtol = 1e-6, maxiter = 10000)

Find limit of a series by Aitken acceleration.
"""
function aitken_accelerate(f, x0; n0 = 0, rtol = 1e-6, maxiter = 10000)
    n = n0
    x0 = f(x0, n)
    x1 = f(x0, n + 1)
    x2 = f(x1, n + 2)
    Ax = aitken(x0, x1, x2)
    success = false
    while ~success && n < maxiter
        n += 1
        Axprev = Ax
        x0 = x1
        x1 = x2
        x2 = f(x2, n + 2)
        Ax = aitken(x0, x1, x2)

        if 2 * abs(Ax - Axprev) / abs(Ax + Axprev) < rtol
            success = true
        end
    end
    return Ax, success, n
end

function aitken(x0, x1, x2)
    den = (x0 - x1) - (x1 - x2)
    return x0 - (x1 - x0)^2 / den
end

"""
    converge_series(f, x0; n0 = 0, rtol = 1e-6, maxiter = 10000)

Find limit of series by brute force. The iteration is stopped when the relative change in
accumulated value is smaller than `rtol`.
"""
function converge_series(f, x0; n0 = 0, rtol = 1e-6, maxiter = 10000)
    n = n0
    x1 = x0
    success = false
    while ~success && n < maxiter
        x1 = f(x0, n)

        if 2 * abs(x1 - x0) / abs(x1 + x0) < rtol
            success = true
        end

        n += 1
        x0 = x1
    end
    return x1, success, n
end

"""
    check_spline_args(x, y)

Ensure that the x array contains unique and sorted values (while preserving the
relaionship between the x and y values).
"""
function check_spline_args(x, y)
    if any(diff(x) .== 0)
        error("entries in x must be unique")
    end
    if !issorted(x)
        idcs = sortperm(x)
        x = x[idcs]
        y = y[idcs]
    end
    x, y
end

"""
     make_spline_ifun(x, ifun)

If `ifun != nothing` then `ifun(x0)` should return the index of the first element in x
which is bigger than x0. Otherwise, it defaults two one of two options:
1. If `x` is uniformly spaced, the index is calculated based on the spacing of `x`
2. If `x` is not uniformly spaced, a `FastFinder` is used instead.
"""
function make_spline_ifun(x, ifun)
    if ifun === nothing
        δx = x[2] - x[1]
        if all(diff(x) .≈ δx)
            # x is uniformly spaced - use fast lookup
            xmax = maximum(x)
            xmin = minimum(x)
            N = length(x)
            ffast(x0) = x0 <= xmin ? 2 :
                        x0 >= xmax ? N :
                        ceil(Int, (x0-xmin)/(xmax-xmin)*(N-1))+1
            ifun = ffast
        else
            # x is not uniformly spaced - use brute-force lookup
            ifun = FastFinder(x)
        end
    end
    ifun
end

"""
    CSpline

Simple cubic spline, see e.g.:
http://mathworld.wolfram.com/CubicSpline.html Boundary
conditions extrapolate with initially constant gradient
"""
struct CSpline{Tx,Ty,Vx<:AbstractVector{Tx},Vy<:AbstractVector{Ty}, fT}
    x::Vx
    y::Vy
    D::Vy
    ifun::fT
    bounds_error::Bool
end

# make broadcast like a scalar
Broadcast.broadcastable(c::CSpline) = Ref(c)

"""
    CSpline(x, y [, ifun]; bounds_error=false)

Construct a `CSpline` to interpolate the values `y` on axis `x`.

If given, `ifun(x0)` should return the index of the first element in x which is bigger
than x0. Otherwise, it defaults two one of two options:
1. If `x` is uniformly spaced, the index is calculated based on the spacing of `x`
2. If `x` is not uniformly spaced, a `FastFinder` is used instead.
"""
function CSpline(x, y, ifun=nothing; bounds_error=false)
    x, y = check_spline_args(x, y)
    R = similar(y)
    R[1] = y[2] - y[1]
    for i in 2:(length(y)-1)
        R[i] = y[i+1] - y[i-1]
    end
    R[end] = y[end] - y[end - 1]
    @. R *= 3
    d = fill(4.0, size(y))
    d[1] = 2.0
    d[end] = 2.0
    dl = fill(1.0, length(y) - 1)
    M = Tridiagonal(dl, d, dl)
    D = M \ R
    CSpline(x, y, D, make_spline_ifun(x, ifun), bounds_error)
end

"""
    (c::CSpline)(x0)

Evaluate the `CSpline` at coordinate `x0`
"""
function (c::CSpline)(x0)
    if c.bounds_error
        x0 < c.x[1] && throw(DomainError("CSpline evaulated out of bounds, $x0 < $(c.x[1])"))
        x0 > c.x[end] && throw(DomainError("CSpline evaulated out of bounds, $x0 > $(c.x[end])"))
    end
    i = c.ifun(x0)
    x0 == c.x[i] && return c.y[i]
    x0 == c.x[i-1] && return c.y[i-1]
    t = (x0 - c.x[i - 1])/(c.x[i] - c.x[i - 1])
    (c.y[i - 1]
        + c.D[i - 1]*t
        + (3*(c.y[i] - c.y[i - 1]) - 2*c.D[i - 1] - c.D[i])*t^2
        + (2*(c.y[i - 1] - c.y[i]) + c.D[i - 1] + c.D[i])*t^3)
end

"""
    linterp(x, x1, y1, x2, y2)

Linear interpolation of `y` at position `xp`, between points `(x1, y1)` and `(x2, y2)`.
For `xp` outside interval `[x1, x2]` this corresponds to linear extrapolation.
"""
function linterp(xp, x1, y1, x2, y2)
    t = (xp - x1) / (x2 - x1)
    (1 - t) * y1 + t * y2
end

"""
    LinTerp

Linear interpolation.
"""
struct LinTerp{Tx,Ty,Vx<:AbstractVector{Tx},Vy<:AbstractVector{Ty}, fT}
    x::Vx
    y::Vy
    ifun::fT
end

# make  broadcast like a scalar
Broadcast.broadcastable(l::LinTerp) = Ref(l)

"""
    LinTerp(xp, xs, ys)

Construct a linear interpolator over `xs`, `ys`.
"""
function LinTerp(xs, ys)
    LinTerp(xs, ys, make_spline_ifun(xs, nothing))
end

"""
    (l::LinTerp)(x)

Evaluate a linear interpolator `LinTerp` at point `x`.
"""
function (l::LinTerp)(x)
    i = l.ifun(x)
    linterp(x, l.x[i - 1], l.y[i - 1], l.x[i], l.y[i])
end

"""
    fftfreq(x)

Calculate frequency vector k from samples x for complex-valued FFT
"""
function fftfreq(x)
    Dx = abs(x[2] - x[1])
    all(diff(x) .≈ Dx) || error("x must be spaced uniformly")
    N = length(x)
    n = collect(range(0, length=N))
    (n .- N/2) .* 2π/(N*Dx)
end

"""
    rfftfreq(x)

Calculate frequency vector k from samples x for rFFT
"""
function rfftfreq(x)
    Dx = abs(x[2] - x[1])
    all(diff(x) .≈ Dx) || error("x must be spaced uniformly")
    Nx = length(x)
    Nf = length(x)÷2 + 1
    n = collect(range(0, length=Nf))
    n .* 2π/(Nx*Dx)
end

rfftnorm(δt) = 2*fftnorm(δt)
fftnorm(δt) = δt/sqrt(2π)

"""
    FastFinder

Callable type which accelerates index finding for the case where inputs are usually in order.
"""
mutable struct FastFinder{xT, xeT}
    x::xT
    mi::xeT
    ma::xeT
    N::Int
    ilast::Int
    xlast::xeT
end

"""
    FastFinder(x)

Construct a `FastFinder` to find indices in the array `x`.

!!! warning
    `x` must be sorted in ascending order for `FastFinder` to work.
"""
function FastFinder(x::AbstractArray)
    issorted(x) || error("Input array for FastFinder must be sorted in ascending order.")
    if any(diff(x) .== 0)
        error("Entries in input array for FastFinder must be unique.")
    end
    FastFinder(x, x[1], x[end], length(x), 0, typemin(x[1]))
end

"""
    (f::FastFinder)(x0)

Find the first index in `f.x` which is larger than `x0`.

This is similar to [`findfirst`](@ref), but it starts at the index which was last used.
If the new value `x0` is close to the previous `x0`, this is much faster than `findfirst`.
"""
function (f::FastFinder)(x0::Number)
    # Default cases if we're out of bounds
    if x0 <= f.x[1]
        f.xlast = x0
        f.ilast = 1
        return 2
    elseif x0 >= f.x[end]
        f.xlast = x0
        f.ilast = f.N
        return f.N
    end
    if f.ilast == 0 # first call -- f.xlast is not set properly so comparisons won't work
        # return using brute-force method instead
        f.ilast = findfirst(x -> x>x0, f.x)
        return f.ilast
    end
    if x0 == f.xlast # same value as before - no work to be done
        return f.ilast
    elseif x0 < f.xlast # smaller than previous value - go through array backwards
        f.xlast = x0
        for i = f.ilast:-1:1
            if f.x[i] < x0
                f.ilast = i+1 # found last idx where x < x0 -> at i+1, x > x0
                return i+1
            end
        end
        # we only get to this point if we haven't found x0 - return the lower bound (2)
        f.ilast = 1
        return 2
    else # larger than previous value - just pick up where we left off
        f.xlast = x0
        for i = f.ilast:f.N
            if f.x[i] > x0
                f.ilast = i
                return i
            end
        end
        # we only get to this point if we haven't found x0 - return the upper bound (N)
        f.ilast = f.N
        return f.N
    end
end

struct RealBSpline{sT,hT,fT}
    rspl::sT
    h::hT
    hh::hT
    ifun::fT
end

struct CmplxBSpline{sT}
    rspl::sT
    ispl::sT
end

Broadcast.broadcastable(rs::RealBSpline) = Ref(rs)
Broadcast.broadcastable(cs::CmplxBSpline) = Ref(cs)

"""
    BSpline(x, y)

Construct a `RealBSpline` or `CmplxSpline` of given `order` (1 to 5, default=3)
to interpolate the values `y` on axis `x`.

If given, `ifun(x0)` should return the index of the first element in x which is bigger
than x0. Otherwise, it defaults two one of two options:
1. If `x` is uniformly spaced, the index is calculated based on the spacing of `x`
2. If `x` is not uniformly spaced, a `FastFinder` is used instead.

# Note
For accurate derivatives use 5th order splines
"""
function BSpline(x::AbstractVector{Tx}, y::AbstractVector{T}; ifun = nothing, order=3) where {Tx <: Real, T <: Real}
    check_spline_args(x, y)
    rspl = Dierckx.Spline1D(x, y, bc="extrapolate", k=order, s=0.0)
    h = zeros(eltype(y), order + 1)
    hh = similar(h)
    RealBSpline(rspl, h, hh, make_spline_ifun(Dierckx.get_knots(rspl), ifun))
end

function BSpline(x::AbstractVector{Tx}, y::AbstractVector{T}; kwargs...) where {Tx <: Real, T <: Complex}
    CmplxBSpline(BSpline(x, real(y); kwargs...),
                 BSpline(x, imag(y); kwargs...))
end

"""
    (cs::RealBSpline)(x)

Evaluate the `RealBSpline` at coordinate(s) `x`
"""
function (rs::RealBSpline)(x)
    splev!(rs.h, rs.hh, rs.rspl.t, rs.rspl.c, rs.rspl.k, x, rs.ifun)
end

"""
    (cs::CmplxSpline)(x)

Evaluate the `CmplxBSpline` at coordinate(s) `x`
"""
function (cs::CmplxBSpline)(x)
    complex(cs.rspl(x), cs.ispl(x))
end

"""
    derivative(rs::RealBSpline, x, order::Integer)

Calculate derivative of the spline `rs`. For `1 <= order < k` (where `k` is the spline order)
this uses an optimised routine.
For `order >= k` this falls back to the generic finite difference based method.
"""
function derivative(rs::RealBSpline, x, order::Integer)
    if order == 0
        return rs(x)
    elseif order < rs.rspl.k
        return Dierckx.derivative(rs.rspl, x, order)
    else
        invoke(derivative, Tuple{Any,Any,Integer}, rs, x, order)
    end
end

"""
    derivative(cs::CmplxBSpline, x, order::Integer)

Calculate derivative of the spline `cs`. For `1 <= order < k` (where `k` is the spline order)
this uses an optimised routine.
For `order >= k` this falls back to the generic finite difference based method.
"""
function derivative(cs::CmplxBSpline, x, order::Integer)
    complex(derivative(cs.rspl, x, order), derivative(cs.ispl, x, order))
end

"""
    differentiate_spline(rs::RealBSpline; order::Integer=1)

Return a new spline which is the derivative of `rs` and has the same support.
Higher orders are obtained by repeated spline differentiation.
"""
function differentiate_spline(rs::RealBSpline, order::Integer)
    x = Dierckx.get_knots(rs.rspl)
    dspl = rs
    for i = 1:order
        dy = derivative.(dspl, x, 1)
        dspl = BSpline(x, dy, order=rs.rspl.k)
    end
    dspl
end

"""
    roots(rs::RealBSpline)

Find the roots of the spline `rs`.
"""
function roots(rs::RealBSpline; maxn::Int=128)
    Dierckx.roots(rs.rspl; maxn=maxn)
end

"""
    splev!(h, hh, t, c, k, x, ifun)

Evaluate a spline s(x) of degree k, given in its b-spline representation.

# Arguments
- `h::ArrayPT<:Real,1}`: work space
- `hh::ArrayPT<:Real,1}`: work space
- `t::Array{T<:Real,1}`: the positions of the knots
- `c::Array{T<:Real,1}`: the b-spline coefficients
- `k::Integer`: the degree of s(x)
- `x::Real`: the point to evaluate at
- `ifun::Function`: a function to find the index `i` s.t. `t[i] > x`

Adapted from splev.f from Dierckx
http://www.netlib.org/dierckx/index.html
"""
function splev!(h, hh, t, c, k, x, ifun)
    k1 = k + 1
    k2 = k1 + 1
    nk1 = length(t) - k1
    tb = t[k1]
    te = t[nk1 + 1]
    l = k1
    l1 = l + 1
    # search for knot interval t(l) <= arg < t(l+1)
    l = ifun(x) - 1 + k
    # evaluate the non-zero b-splines at x.
    fpbspl!(h, hh, t, k, x, l)
    # find the value of s(x) at x.
    sp = 0.0
    ll = l - k1
    for j = 1:k1
        ll = ll + 1
        sp += c[ll]*h[j]
    end
    sp
end

"""
    fpbspl!(h, hh, t, k, x, l)

Evaluate the (k+1) non-zero b-splines of
degree k at t[l] <= x < t[l+1] using the stable recurrence
relation of de boor and cox.

# Arguments
- `h::ArrayPT<:Real,1}`: work space
- `hh::ArrayPT<:Real,1}`: work space
- `t::Array{T<:Real,1}`: the positions of the knots
- `k::Integer`: the degree of s(x)
- `x::Real`: the point to evaluate at
- `l::Integer`: the active knot location: `t[l] <= x < t[l+1]`

Adapted from fpbspl.f from Dierckx
http://www.netlib.org/dierckx/index.html
"""
function fpbspl!(h, hh, t, k, x, l)
    h[1] = one(x)
    for j = 1:k
        for i = 1:j
            hh[i] = h[i]
        end
        h[1] = 0.0
        for i = 1:j
            li = l+i
            lj = li-j
            if t[li] != t[lj]
                f = hh[i]/(t[li] - t[lj])
                h[i] = h[i] + f*(t[li] - x)
                h[i+1] = f*(x - t[lj])
            else
                h[i+1] = 0.0
            end
        end
    end
    return h
end

"""
    pkfw(x, y, pki; level=0.5, skipnonmono=true, closest=5)

Find the full width of a peak in `y` over `x` centred at index `pki`.

The default `level=0.5` requests the full width at half maximum. Setting `level` to something
different computes the corresponding width. E.g. `level=0.1` for the 10% width.

`skipnonmono=true` skips peaks which are not monotonically increaing/decreasing before/after the peak.

`closest=5` sets the minimum number of indices for the full width.
"""
function pkfw(x, y, pki; level=0.5, skipnonmono=true, closest=5)
    val = level*y[pki]
    iup = findnext(x -> x < val, y, pki)
    if iup == nothing
        iup = length(x)
    end
    idn = findprev(x -> x < val, y, pki)
    if idn == nothing
        idn = 1
    end
    if skipnonmono
        if any(diff(y[pki:iup]) .> 0)
            return missing
        end
        if any(diff(y[idn:pki]) .< 0)
            return missing
        end
    end
    if (iup - idn) < closest
        return missing
    end
    up = Maths.linterpx(x[iup - 1], x[iup], y[iup - 1], y[iup], val)
    dn = Maths.linterpx(x[idn], x[idn + 1], y[idn], y[idn + 1], val)
    return up - dn
end

"""
    findpeaks(x, y; threshold=0.0, filterfw=true)

Find isolated peaks in a signal `y` over `x` and return their value, FWHM and index.
`threshold=0.0` allows filtering peaks above a threshold value.
If `filterfw=true` then only peaks with a clean FWHM are returned.
"""
function findpeaks(x, y; threshold=0.0, filterfw=true)
    pkis = Peaks.argmaxima(y, 10)
    proms = Peaks.peakproms(pkis, y)
    pks = [(peak=y[pki], fw=pkfw(x, y, pki), position=x[pki], index=pki) for pki in pkis]
    # filter out peaks with missing fws
    if filterfw
        pks = filter(x -> !(x.fw === missing), pks)
    end
    # filter out peaks below threshold
    pks = filter(x -> x.peak > threshold, pks)
end

end


### A Pluto.jl notebook ###
# v0.12.6

using Markdown
using InteractiveUtils

# ╔═╡ 971fc5f0-1c52-11eb-27da-43003bbda193
begin
	using DifferentialEquations
	using Plots
end

# ╔═╡ 79e6b4e0-1c56-11eb-30a7-b78aa49dc552
function msd_sys(du,u,p,t)
  # Manual implementation of state transition matrix
  du[1] = u[2]
  du[2] = -p[1]*u[1] - p[2]*u[2]

  # Add acceleration due to gravity
  du[2] += p[3]

  # Add disturbance in the force
  du[2] +=  p[4]*sin(p[5]*t)

  # Add control input
  du[2] += p[6]
end

# ╔═╡ 943468fe-1c56-11eb-07ed-9d4f3578e7c6
begin
	g = 9.8 # gravity in metres per second per second
	mass = 0.1
	k_s = 2.5 # spring constant
	k_d = 0.1 # damping coefficient

	u0 = [-1., 0.]

	f_smc = 10.0

	p_simple = [k_s/mass, k_d/mass, g, 0, 0, 0]
	p_dist = [k_s/mass, k_d/mass, g, 5., 10., 0]
	p_smc = [k_s/mass, k_d/mass, g, 5., 10., 0.]

	tspan = (0.0,10.0)
end

# ╔═╡ a705ea40-1c56-11eb-1a72-136337ca099f
function smc_affect!(msd_ss_model)
    sigma = 20.0*msd_ss_model.u[1] + 1.0*msd_ss_model.u[2]
    if sigma < 0
        msd_ss_model.p[6] = +f_smc/mass
    else
        msd_ss_model.p[6] = -f_smc/mass
    end
end

# ╔═╡ b9fb9550-1c56-11eb-1080-b7ec4f1598fe
pcb = PeriodicCallback(smc_affect!, 10/10000)

# ╔═╡ b0c2973e-1c56-11eb-063b-cf48350c9806
begin
	prob_simple = ODEProblem(msd_sys,u0,tspan,p_simple)
	prob_dist = ODEProblem(msd_sys,u0,tspan,p_dist)
	prob_smc = ODEProblem(msd_sys,u0,tspan,p_smc)
end

# ╔═╡ dbb46550-1c56-11eb-32c5-3dd9bceda11a
begin
    raw_sol = solve(prob_simple,Tsit5());
    dist_sol = solve(prob_dist,Tsit5());
    smc_sol = solve(prob_smc,Tsit5(),callback=pcb);
end

# ╔═╡ fc2b3f70-1c56-11eb-1322-b9d6b5728032
begin
	plot(raw_sol,vars=(0,1),w=1,label="Raw", layout = (2,1), subplot=1)
	plot!(dist_sol,vars=(0,1),w=1,label="Disturbed")
	plot!(smc_sol,vars=(0,1),w=1,label="SMC")
	plot!(raw_sol,vars=(1,2),w=1,label="Raw",subplot=2)
	plot!(dist_sol,vars=(1,2),w=1,label="Disturbed",subplot=2)
	plot!(smc_sol,vars=(1,2),w=1,label="SMC",subplot=2,xlim=(-1.1,+1.5),ylim=(-5,11))
end

# ╔═╡ Cell order:
# ╠═971fc5f0-1c52-11eb-27da-43003bbda193
# ╠═79e6b4e0-1c56-11eb-30a7-b78aa49dc552
# ╠═943468fe-1c56-11eb-07ed-9d4f3578e7c6
# ╠═a705ea40-1c56-11eb-1a72-136337ca099f
# ╠═b9fb9550-1c56-11eb-1080-b7ec4f1598fe
# ╠═b0c2973e-1c56-11eb-063b-cf48350c9806
# ╠═dbb46550-1c56-11eb-32c5-3dd9bceda11a
# ╠═fc2b3f70-1c56-11eb-1322-b9d6b5728032


# ---
# title: Unsupervised DDVFA Example
# id: ddvfa_unsupervised
# cover: ../assets/ddvfa.png
# date: 2021-11-30
# author: "[Sasha Petrenko](https://github.com/AP6YC)"
# julia: 1.6
# description: This demo shows how to use DDVFA for unsupervised learning by clustering Iris samples.
# ---

# DDVFA is an unsupervised clustering algorithm by definition, so it can be used to cluster a set of samples all at once in batch mode.

# We begin with importing AdaptiveResonance for the ART modules and MLDatasets for loading some data.
using AdaptiveResonance
using MLDatasets

# We will download the Iris dataset for its small size and benchmark use for clustering algorithms.
Iris.download(i_accept_the_terms_of_use=true)
features, labels = Iris.features(), Iris.labels()

# Next, we will instantiate a DDVFA module.
# We could create an options struct for reuse with `opts=opts_DDVFA(...)`, but for now we will use the direct keyword arguments approach.
art = DDVFA(rho_lb=0.6, rho_ub=0.75)

# To train the module on the training data, we use `train!`.
# The train method returns the prescribed cluster labels, which are just what the algorithm believes are unique/separate cluster.
# This is because we are doing *unsupervised* learning rather than supervised learning with known labels.
y_hat_train = train!(art, features)

# Though we could inspect the unique entries in the list above, we can see the number of categories directly from the art module.
art.n_categories

# Because DDVFA actually has FuzzyART modules for F2 nodes, each category has its own category prototypes.
# We can see the total number of weights in the DDVFA module by summing `n_categories` across all F2 nodes.
total_vec = [art.F2[i].n_categories for i = 1:art.n_categories]
total_cat = sum(total_vec)


struct FixedNesterovSequence{R} end

"""
    FixedNesterovSequence(R::Type)

Construct the standard sequence of extrapolation coefficients, of type `R`, used in Nesterov-accelerated gradient methods.

Argument `R` should be a floating point type, and specifies the type of the elements in the sequence.

See also: [`SimpleNesterovSequence`](@ref), [`ConstantNesterovSequence`](@ref).
"""
FixedNesterovSequence(R) = FixedNesterovSequence{R}()

function Base.iterate(::FixedNesterovSequence{R}, t=R(1)) where R
    t_next = (1 + sqrt(1 + 4 * t^2)) / 2
    return (t - 1) / t_next, t_next
end

Base.IteratorSize(::Type{<:FixedNesterovSequence}) = Base.IsInfinite()
Base.eltype(::Type{FixedNesterovSequence{R}}) where R = R

struct SimpleNesterovSequence{R} end

"""
    SimpleNesterovSequence(R::Type)

Construct an iterator yielding the sequence `(k - 1) / (k + 2)`, of type `R`, for `k >= 1`.

This is the "simplified" sequence of extrapolation coefficients used in Nesterov-accelerated gradient methods.
Argument `R` should be a floating point type, and specifies the type of the elements in the sequence.

See also: [`FixedNesterovSequence`](@ref), [`ConstantNesterovSequence`](@ref).
"""
SimpleNesterovSequence(R) = SimpleNesterovSequence{R}()

Base.iterate(::SimpleNesterovSequence{R}, k=1) where {R} = R(k - 1) / (k + 2), k + 1

Base.IteratorSize(::Type{<:SimpleNesterovSequence}) = Base.IsInfinite()
Base.eltype(::Type{SimpleNesterovSequence{R}}) where R = R

"""
    ConstantNesterovSequence(m::R, stepsize::R)

Construct a constant iterator yielding `(1 - sqrt(m * stepsize)) / (1 + sqrt(m * stepsize))` with type `R`.

This sequence of extrapolation coefficients is commonly used in Nesterov-accelerated gradient methods
when applied to strongly convex functions.

See also: [`FixedNesterovSequence`](@ref), [`SimpleNesterovSequence`](@ref).
"""
function ConstantNesterovSequence(m::R, stepsize::R) where R
    k_inverse = m * stepsize
    return repeated((1 - sqrt(k_inverse)) / (1 + sqrt(k_inverse)))
end

mutable struct AdaptiveNesterovSequence{R}
    m::R
    stepsize::R
    theta::R
end

"""
    AdaptiveNesterovSequence(m::R)

Construct an object that yields a sequence of extrapolation coefficients, commonly used in Nesterov-accelerated gradient methods.

Argument `m` is usually the convexity modulus of the interested function, so it should be zero for barely convex functions,
and greater than zero for strongly convex ones. The extrapolation coefficients can be extracted by calling [`next!`](@ref) on the
object.

Note that this type, as opposed to e.g. [`FixedNesterovSequence`](@ref), is not a standard Julia iterator, in that it requires
the sequence of stepsizes (input to [`next!`](@ref)) in order to produce the output sequence.

Special cases: if `next!` is invoked with a constant stepsize `gamma`, then
- If `m == 0`, the output sequence will be the same as for `FixedNesterovSequence(R)`.
- If `m > 0`, the output sequence will be the same as for `ConstantNesterovSequence(m, gamma)`.

See also: [`next!`](@ref).
"""
AdaptiveNesterovSequence(m::R) where R = AdaptiveNesterovSequence{R}(m, -R(1), -R(1))

"""
    next!(seq::AdaptiveNesterovSequence{R}, stepsize::R)

Push a new stepsize value into the sequence, and return the next extrapolation coefficient.

See also: [`AdaptiveNesterovSequence`](@ref).
"""
function next!(seq::AdaptiveNesterovSequence{R}, stepsize::R) where R
    if seq.stepsize < 0
        seq.stepsize = stepsize
        seq.theta = seq.m > 0 ? sqrt(seq.m * stepsize) : R(1)
    end
    b = seq.theta^2 / seq.stepsize - seq.m
    delta = b^2 + 4 * (seq.theta^2) / (seq.stepsize * stepsize)
    theta = stepsize * (-b + sqrt(delta)) / 2
    beta = stepsize * (seq.theta) * (1 - seq.theta) / (seq.stepsize * theta + stepsize * seq.theta ^ 2)
    seq.stepsize = stepsize
    seq.theta = theta
    return beta
end

struct NesterovExtrapolation{S} end

NesterovExtrapolation(S) = NesterovExtrapolation{S}()
NesterovExtrapolation() = NesterovExtrapolation{SimpleNesterovSequence}()

acceleration_style(::Type{<:NesterovExtrapolation}) = NesterovStyle()

initialize(::NesterovExtrapolation{S}, x) where S = Iterators.Stateful(S{real(eltype(x))}())


#
# Factors Map Tests
#
# log, fill, abs, etc

@testset "Factors Map" begin

@testset "normalize" begin
    @testset "1" begin
        ϕ = Factor([:A, :B], [1 2; 3 4])
        ϕ2 = normalize(ϕ, p=1)

        @test_approx_eq ϕ2.potential [0.1 0.2; 0.3 0.4]
        @test ϕ.potential == [1 2; 3 4]

        @test_throws ArgumentError normalize(ϕ, :waldo)

        normalize!(ϕ, p=2)
        @test_approx_eq ϕ.potential [1 2; 3 4]./sqrt(30)
    end

    @testset "2" begin
        ϕ = Factor(-30)
        normalize!(ϕ)
        @test ϕ.potential == reshape([-1.0], ())

        ϕ = Factor(:A, [2, 3])
        normalize!(ϕ, p=1)
        @test_approx_eq ϕ.potential [0.4, 0.6]
    end

    @testset "3" begin
        ϕ = Factor(1:24, :X=>3, :Y=>2, :Z=>4)
        ϕ2 = normalize(ϕ, :Y)
        @test_approx_eq sum(ϕ2.potential, 2) fill(1.0, (3, 1, 4))

        ϕ3 = normalize(ϕ, [:X, :Z])
        @test_approx_eq sum(ϕ3.potential, [1, 3]) fill(1.0 ,(1, 2, 1))
    end
end

@testset "map" begin
    @testset "1" begin
        ϕ = Factor(1:16, :X=>4, :Y=>4)
        ϕ2 = map(x -> x+10, ϕ)
        @test ϕ.potential + 10 == ϕ2.potential

        map!(x-> 2*x - 10, ϕ)
        @test ϕ.potential == 2*reshape(1:16, (4,4) ) - 10
    end

    @testset "2" begin
        ϕ = Factor(:X, 1:4)
        map!(x -> x + 2, ϕ)
        @test ϕ.potential == [3, 4, 5, 6]
    end

    @testset "3" begin
        ϕ = Factor(16)
        ϕ2 = map(log10, ϕ)
        @test ϕ.potential == reshape([16], ())
        @test ϕ2.potential == reshape([log10(16)], ())
    end
end

@testset "rand!" begin
    @testset "1" begin
        ϕ = Factor(16, :X=>3, :Y=>4)
        rand!(ϕ)
        @test ϕ.potential != 16
    end

    @testset "2" begin
        ϕ = Factor(3)
        randn!(ϕ)
        @test ϕ.potential != 3
    end
end

@testset "fill!" begin
    @testset "1" begin
        ϕ = Factor(16, :X=>3, :Y=>4)
        fill!(ϕ, 2016)
        @test all(ϕ.potential .== 2016)
    end

    @testset "2" begin
        ϕ = Factor(3)
        fill!(ϕ, 19)
        @test ϕ.potential == reshape([19], ())
    end
end

@testset "log, exp, etc." begin
    @testset "1" begin
        ϕ = Factor(:X=>3, :Y=>4)
        rand!(ϕ)

        @test_approx_eq ϕ.potential exp(log(ϕ)).potential
        @test_approx_eq ϕ.potential log10(exp10(ϕ)).potential
    end

    @testset "2" begin
        ϕ = Factor(58)
        @test_approx_eq log2(ϕ).potential log2(58)
        @test_approx_eq secd(ϕ).potential secd(58)
    end
end

end


#=
Created on Thurday 16 Jan 2020
Last update:

@author: Michiel Stock
michielfmstock@gmail.com

Solution of the Monge problem
=#

using STMO
using STMO.OptimalTransport
using Plots, LaTeXStrings

p = scatter(X1[:,1], X1[:,2], color=myorange, label="location cells at t1")
xlabel!("\$x\$")
ylabel!("\$y\$")
scatter!(X2[:,1], X2[:,2], color=mygreen, label="location cells at t2")

C = dist(X1, X2)

perm, cost = monge_brute_force(C)

for (i, j) in enumerate(perm)
    plot!(p, [X1[i,1], X2[j,1]], [X1[i,2], X2[j,2]], label="", color=myred,
                alpha=0.8, lw=2)
end

title!(p, "Monge solution with cost=$cost")

savefig(p, "figures/mongesol.png")


using MatrixHelper
using Base.Test

# write your own tests here
a = rand(21,21)
asp = sortpermrows(a)
as = sortrows(a)

@test a[asp,:] == as
for i in 1:21
	@test selectpermrows(a, i) == asp[i]
end
@test selectpermrows(a, 6:12) == asp[6:12]

asp = sortpermcols(a)
as = sortcols(a)
@test a[:,asp] == as
for i in 1:21
	@test selectpermcols(a, i) == asp[i]
end
@test selectpermcols(a, 6:12) == asp[6:12]


#=

Tests for Discrete Decision  Processes (DDP)

Original Python Author: Daisuke Oyama
Authors: Spencer Lyon and Matthew McKay

Tests for markov/ddp.jl

=#

@testset "Testing markov/dpp.jl" begin

    #-Setup-#

    # Example from Puterman 2005, Section 3.1
    beta = 0.95

    # Formulation with Dense Matrices R: n x m, Q: n x m x n
    n, m = 2, 2  # number of states, number of actions
    R = [5.0 10.0; -1.0 -Inf]
    Q = Array{Float64}(n, m, n)
    Q[:, :, 1] = [0.5 0.0; 0.0 0.0]
    Q[:, :, 2] = [0.5 1.0; 1.0 1.0]

    ddp0 = DiscreteDP(R, Q, beta)

    # Formulation with state-action pairs
    L = 3  # Number of state-action pairs
    s_indices = [1, 1, 2]
    a_indices = [1, 2, 1]
    R_sa = [R[1, 1], R[1, 2], R[2, 1]]
    Q_sa = spzeros(L, n)
    Q_sa[1, :] = Q[1, 1, :]
    Q_sa[2, :] = Q[1, 2, :]
    Q_sa[3, :] = Q[2, 1, :]
    ddp0_sa = DiscreteDP(R_sa, Q_sa, beta, s_indices, a_indices)

    # List of ddp formulations
    ddp0_collection = (ddp0, ddp0_sa)

    # Maximum Iteration and Epsilon for Tests
    max_iter = 200
    epsilon = 1e-2

    # Analytical solution for beta > 10/11, Example 6.2.1
    v_star = [(5-5.5*beta)/((1-0.5*beta)*(1-beta)), -1/(1-beta)]
    sigma_star = [1, 1]

    @testset "bellman_operator methods" begin
        # Check both Dense and State-Action Pair Formulation
        for ddp in ddp0_collection
        	@test isapprox(bellman_operator(ddp, v_star), v_star)
    	end
    end

    @testset "RQ_sigma" begin
        nr, nc = size(R)
        # test for DDP
        sigmas = ([1, 1], [1, 2], [2, 1], [2, 2])
        for sig in sigmas
            r, q = RQ_sigma(ddp0, sig)

            for i_r in 1:nr
                @test r[i_r] == ddp0.R[i_r, sig[i_r]]
                for i_c in 1:length(sig)
                    @test vec(q[i_c, :]) == vec(ddp0.Q[i_c, sig[i_c], :])
                end
            end
        end

        # TODO: add test for DDPsa
    end

    @testset "compute_greedy methods" begin
        # Check both Dense and State-Action Pair Formulation
        for ddp in ddp0_collection
        	@test compute_greedy(ddp, v_star) == sigma_star
    	end
    end

    @testset "evaluate_policy methods" begin
        # Check both Dense and State-Action Pair Formulation
        for ddp in ddp0_collection
        	@test isapprox(evaluate_policy(ddp, sigma_star), v_star)
        end
    end

    @testset "methods for subtypes != (Float64, Int)" begin
        float_types = [Float16, Float32, Float64, BigFloat]
        int_types = [Int8, Int16, Int32, Int64, Int128,
                     UInt8, UInt16, UInt32, UInt64, UInt128]

        for ddp in ddp0_collection
            for f in (bellman_operator, compute_greedy)
                for T in float_types
                    f_f64 = f(ddp, [1.0, 1.0])
                    f_T = f(ddp, ones(T, 2))
                    @test isapprox(f_f64, convert(Vector{eltype(f_f64)}, f_T))
                end

                # only Integer subtypes can be Rational type params
                # NOTE: Only the integer types below don't overflow for this example
                for T in [Int64, Int128]
                    @test f(ddp, [1//1, 1//1]) == f(ddp, ones(Rational{T}, 2))
                end
            end

            for T in float_types, S in int_types
                v = ones(T, 2)
                s = ones(S, 2)
                # just test that we can call the method and the result is
                # deterministic
                @test bellman_operator!(ddp, v, s) == bellman_operator!(ddp, v, s)
            end

            for T in int_types
                s = T[1, 1]
                @test isapprox(evaluate_policy(ddp, s), v_star)
            end
        end
    end

    @testset "compute_greedy! changes ddpr.v" begin
        res = solve(ddp0, VFI)
        res.Tv[:] = 500.0
        compute_greedy!(ddp0, res)
        @test maximum(abs, res.Tv - 500.0) > 0
    end

    @testset "value_iteration" begin
        # Check both Dense and State-Action Pair Formulation
        for ddp_item in ddp0_collection
            # Compute Result
            res = solve(ddp_item, VFI)
            v_init = [0.0, 0.0]
            res_init = solve(ddp_item, v_init, VFI; epsilon=epsilon)

            # Check v is an epsilon/2-approxmation of v_star
            @test maximum(abs, res.v - v_star) < epsilon/2
            @test maximum(abs, res_init.v - v_star)    < epsilon/2

            # Check sigma == sigma_star.
            # NOTE we need to convert from linear to row-by-row index
            @test res.sigma == sigma_star
            @test res_init.sigma == sigma_star
        end
    end

    @testset "policy_iteration" begin
        # Check both Dense and State-Action Pair Formulation
        for ddp_item in ddp0_collection
            res = solve(ddp_item, PFI)
            v_init = [0.0, 1.0]
            res_init = solve(ddp_item, v_init, PFI)

            # Check v == v_star
            @test isapprox(res.v, v_star)
            @test isapprox(res_init.v, v_star)

            # Check sigma == sigma_star
            @test res.sigma == sigma_star
            @test res_init.sigma == sigma_star
        end
    end

    @testset "DiscreteDP{Rational,_,_,Rational} maintains Rational" begin
        ddp_rational = DiscreteDP(map(Rational{BigInt}, R),
                                  map(Rational{BigInt}, Q),
                                  map(Rational{BigInt}, beta))
        # do minimal number of iterations to avoid overflow
        vi = Rational{BigInt}[1//2, 1//2]
        @test eltype(solve(ddp_rational, VFI; max_iter=1, epsilon=Inf).v) == Rational{BigInt}
        @test eltype(solve(ddp_rational, vi, PFI; max_iter=1).v) == Rational{BigInt}
        @test eltype(solve(ddp_rational, vi, MPFI; max_iter=1, k=1, epsilon=Inf).v) == Rational{BigInt}
    end

    @testset "DiscreteDP{Rational{BigInt},_,_,Rational{BigInt}}  works" begin
        ddp_rational = DiscreteDP(map(Rational{BigInt}, R),
                                  map(Rational{BigInt}, Q),
                                  map(Rational{BigInt}, beta))
        # do minimal number of iterations to avoid overflow
        r1 = solve(ddp_rational, PFI)
        r2 = solve(ddp_rational, MPFI)
        r3 = solve(ddp_rational, VFI)
        @test maximum(abs, r1.v-v_star) < 1e-13
        @test r1.sigma == r2.sigma
        @test r1.sigma == r3.sigma
        @test r1.mc.p == r2.mc.p
        @test r1.mc.p == r3.mc.p
    end

    @testset "modified_policy_iteration" begin
        for ddp_item in ddp0_collection
            res = solve(ddp_item, MPFI)
            v_init = [0.0, 1.0]
            res_init = solve(ddp_item, v_init, MPFI)

                    # Check v is an epsilon/2-approxmation of v_star
            @test maximum(abs, res.v - v_star) < epsilon/2
            @test maximum(abs, res_init.v - v_star) < epsilon/2

            # Check sigma == sigma_star
            @test res.sigma == sigma_star
            @test res_init.sigma == sigma_star

            #Test Modified Policy Iteration k0
            k = 0
            res = solve(ddp_item, MPFI; max_iter=max_iter, epsilon=epsilon, k=k)

            # Check v is an epsilon/2-approxmation of v_star
            @test maximum(abs, res.v - v_star) < epsilon/2

            # Check sigma == sigma_star
            @test res.sigma == sigma_star
        end
    end

    @testset "DDPsa constructor" begin
        @testset "feasbile action pair" begin
            _R = [1.0, 0.0, 0.0, 1.0]
            _Q = fill(1/3, 4, 3)
            _s_ind = [1, 1, 3, 3]
            _a_ind = [1, 2, 1, 2]
            @test_throws ArgumentError DiscreteDP(_R, _Q, beta, _s_ind, _a_ind)
        end

        _R, _Q = R_sa, Q_sa
        _s_ind = [1, 1, 2]
        _a_ind = [1, 2, 1]

        @testset "beta in [0, 1)" begin
            @test_throws ArgumentError DiscreteDP(_R, _Q, -eps(), _s_ind, _a_ind)
            @test_throws ArgumentError DiscreteDP(_R, _Q, 1.0, _s_ind, _a_ind)
            @test_throws ArgumentError DiscreteDP(_R, _Q, 1+eps(), _s_ind, _a_ind)
        end

        @testset "argument sizes" begin
            # NQ != 2
            @test_throws ArgumentError DiscreteDP(_R, rand(4, 3, 1), beta, _s_ind, _a_ind)

            # NR != 1
            @test_throws ArgumentError DiscreteDP(rand(4, 1), _Q, beta, _s_ind, _a_ind)

            # incorrect lengths
            @test_throws ArgumentError DiscreteDP(rand(2), _Q, beta, _s_ind, _a_ind)
            @test_throws ArgumentError DiscreteDP(_R, rand(5, 2), beta, _s_ind, _a_ind)
            @test_throws ArgumentError DiscreteDP(_R, _Q, beta, rand(1:3, 2), _a_ind)
            @test_throws ArgumentError DiscreteDP(_R, _Q, beta, _s_ind, rand(1:3, 2))
        end

        @testset "duplicate sa pair" begin
            @test_throws ArgumentError DiscreteDP(_R, _Q, beta, _s_ind, [1, 1, 2])
        end
    end

    @testset "DDP constructor" begin
        @testset "beta in [0, 1)" begin
            @test_throws ArgumentError DiscreteDP(R, Q, -eps())
            @test_throws ArgumentError DiscreteDP(R, Q, 1.0)
            @test_throws ArgumentError DiscreteDP(R, Q, 1+eps())
        end

        @testset "feasbile action pair" begin
            #Dense Matrix
            n, m = 2, 2
            _R = [-Inf -Inf; 1.0 2.0]

            _Q = Array{Float64}(n, m, n)
            _Q[:, :, 1] = [0.5 0.0; 0.0 0.0]
            _Q[:, :, 2] = [0.5 1.0; 1.0 1.0]
            _beta = 0.95

            @test_throws ArgumentError DiscreteDP(_R, _Q, _beta)
        end

        @testset "R, Q sizes" begin
            # NQ != 3
            @test_throws ArgumentError DiscreteDP(R, zeros(2, 2), beta)

            # NR != 2
            @test_throws ArgumentError DiscreteDP(zeros(1), Q, beta)

            # incompatible dimensions
            @test_throws ArgumentError DiscreteDP(zeros(2, 3), Q, beta)
            @test_throws ArgumentError DiscreteDP(R, zeros(2, 3, 2), beta)
        end
    end

    @testset "ddp_negative_inf_error()" begin
        # Dense Matrix
        n, m = 3, 2
        R = [0 1;
             0 -Inf;
            -Inf -Inf]
        Q = fill(1.0/n, n, m, n)
        beta = 0.95

        @test_throws ArgumentError DiscreteDP(R, Q, beta)

        # State-Action Pair Formulation
        #
        # s_indices = [0, 0, 1, 1, 2, 2]
        # a_indices = [0, 1, 0, 1, 0, 1]
        # R_sa = reshape(R, n*m)
        # Q_sa_dense = reshape(Q, n*m, n)          #TODO: @sglyon Not sure how to reshape in Julia
        #
        # @test_throws ArgumentError DiscreteDP(R_sa, Q_sa, beta, s_indices, a_indices)
    end

end # end @testset


#--------------------------------------------------------------------
# DNSS.jl
# Soham 03-2022
# Simulate the collapse of a self-gravitating scalar field
# in Vaidya spacetime
#--------------------------------------------------------------------

using ForwardDiff, QuadGK

function psi(x::T)::T where {T<:Real}
    (a, b) = (-0.5, 0.5)
    # (a, b) = (-1, 1)
    if (a <= x <= b)
        return x * (1-(x/a))^8  * (1 - (x/b))^8
    else
        return 0.0
    end
end

function psi_spherical(u::T, v::T)::T where {T<:Number}
    t = v + u
    r = v - u
    α = π/2 - k*t
    if  r == 0.0
        return cos(k*t)
    else
        return cos(k*t) * (sin(k*r) / (k*r))
    end
end

function checkinitialdata(a::Field{S}) where {S}
    PS = derivative(a.space)
    dψ = Field(PS, dpsi)
    r  = Field(PS, grr)
    D  = derivative(PS)
    @show norm(-(2 / a) * (D * a) + 4 * π * r * (dψ)^2)
end

function dpsi(x::T)::T where {T<:Real}
    return ForwardDiff.derivative(psi, x)
end

function guv(v::T)::T where {T<:Real}
    integral, error = quadgk(x->(x * dpsi(x)^2), -1, v, rtol=1e-12)
    return exp(2π * integral)
end

function resolve(U::NTuple{3, Field{S}}, var::Symbol)::NTuple{3, Field{S}} where {S}
    PS = first(U).space
    @show PS
    D = derivative(PS)
    I = identity(PS)
    B = incomingboundary(PS)

    if var == :a
        (a, η, ψ) = U
        L = (D*η)*D + (4π * η * (D*ψ)^2) * I
        lna = linsolve(L + B, B*log(a))
        a   = exp(lna)
        return (a, η, ψ)
    end
end

# TODO: Test the simplest initial data for spherically symmetric wave
A  = ChebyshevGL{U, 20, Float64}(0.0, 1.0)
B  = ChebyshevGL{V, 20, Float64}(0.0, 1.0)
PS = ProductSpace(A, B)

k  = 0.01
ψ  = extractUboundary(Field(PS, psi_spherical), :incoming)
r  = extractUboundary(Field(PS, (u,v)->v-u), :incoming)
a  = extractUboundary(Field(PS, (u,v)->2), :incoming)
@show norm(lineconstraints((a, r, ψ)))


using Polyhedra
using Base.Test

A = [1 1; -1 0; 0 -1]
b = [1, 0, 0]
linset = IntSet([1])

@test_throws ErrorException SimpleHRepresentation(A, [0, 0], linset)
@test_throws ErrorException SimpleHRepresentation(A, [0, 0], IntSet([4]))
ine = SimpleHRepresentation(A, b, linset)
@test fulldim(ine) == 2
@test eltype(ine) == Int

V = [0 1; 1 0]
@test_throws ErrorException SimpleVRepresentation(V, [1 0 0], IntSet([]), IntSet([]))
@test_throws ErrorException SimpleVRepresentation(V, [1 1], IntSet([]), IntSet([2]))
@test_throws ErrorException SimpleVRepresentation(V, [1 1], IntSet([4]), IntSet([]))
@test_throws ErrorException SimpleVRepresentation(V, IntSet([4]))
ext = SimpleVRepresentation(V)
@test fulldim(ext) == 2
@test eltype(ext) == Int


#= hausdorff.jl
Philip Barrett, pobarrett@gmail.com
30may2016, Chicago

Computes the hausdorff distance between polygons (and points)
=#

"""
    hausdorff

Computes the hausdorff distance between two polygons, or two arrays of polygons
"""
function hausdorff( poly1::Polygon, poly2::Polygon )
  return hausdorffCore( poly1.pts, poly2 )
end

function hausdorff( poly1::Array{Polygon,1}, poly2::Array{Polygon,1} )
  N = length(poly1)
  length(poly2) == N || error("Polygon arrays must have same diemnsion")
      # Error handling
  return [ hausdorffCore( poly1[i].pts, poly2[i] ) for i in 1:N ]
end

function hausdorff( poly1::Array{Polygon,2}, poly2::Array{Polygon,2} )
  N, M = size(poly1)
  size(poly2) == (N, M) || error("Polygon arrays must have same diemnsion")
      # Error handling
  return [ hausdorffCore( poly1[i,j].pts, poly2[i,j] ) for i in 1:N, j in 1:M ]
end
"""
    function hausdorffCore( newpts::Matrix, poly::Polygon  )

Core part of any hausdorff distance routine.  Measures the distance between a matrix of points newpts and a polygon poly
"""
function hausdorffCore( newpts::Matrix, poly::Polygon )

  pts, dirs, dists = poly.pts, poly.dirs, poly.dists
      # Unpack the polygon object
  N = size(newpts)[1]
      # number of points
  dline, dvert, newdist, hausdorff = 0.0, 0.0, 0.0, 0.0
      # initialize the distance measures
  for i in 1:N
    dline = setsNearestLine( vec(newpts[i,:]), dirs, dists )
        # Distance to nearest line
    dvert = setsNearestVert( vec(newpts[i,:]), pts )
        # Distance to nearest vertex
    newdist = min( dline, dvert )
        # Minimum distance
    hausdorff  = max( hausdorff, newdist )
        # Update hausdorff dist via minmax
  end
  return hausdorff
end

function setsNearestLine( pt::Vector, dirs::Matrix, dists::Vector )

  N::Int = length(dists)
      # Number of sides
  dff = dists - dirs * pt
  absdiff = abs( dff )
      # Difference between the projection directions and the boundaries (and absolute value)
  nearpts = ones( N ) * pt' + dirs .* ( dff * ones( 1, 2 ) )
      # Recover the closest points on the normal surfaces
  idx = setsInsidePtsIdx( nearpts, dirs, dists )
      # Test whether these points are inside the Polygon
  if !any(idx)
    return Inf
        # Outsize return if none of the nearest points are in the polygon
  end
  distinside = [ idx[i] ? absdiff[i] : Inf for i in 1:N ]
      # Compute the minimum distances to the points on the lines that are
      # inside the polygon
  return minimum( distinside )
      # Return the minimum
end

function setsInsidePtsIdx( cands::Matrix, dirs::Matrix,
                            dists::Vector )
  N::Int = size(cands)[1]
      # Number of candidate points
  D = dirs * cands'  - dists * ones( 1, N )
      # Negative if inside the set
  return all( D .<= 0, 1 )
end

function setsNearestVert( cand::Vector, pts::Matrix )
  N = size(pts)[1]
      # Nubmer of points
  return minimum( [ norm( cand' - pts[i,:] ) for i in 1:N ] )
end

# double sets_nearest_vert_cpp_eigen( VectorXd pt, MatrixXd mB ){
# // C++ translation of sets.nearest.vert
#
#   int iB = mB.rows() ;
#         // Number of rows in B
#   VectorXd vDists( iB ) ;
#         // Initialize the distance vector
#   for( int iII=0; iII<iB; iII++ ){
#     vDists[ iII ] = ( pt.transpose() - mB.row( iII ) ).norm() ;
#           // Compute the distance to pt
#   }
#   return vDists.minCoeff() ;
# }


using ChebyshevQuantum
using Test

tol_val = 1e-11

@testset "test interp" begin

    function f(x)
        return exp(-0.2*sin(3*x))
    end

    c = ChebyshevQuantum.Interp.make_cheb(f , a=2, b = 3)

    @test isapprox(c(2.3), f(2.3), atol=tol_val)
    @test isapprox(c(2.0), f(2.0), atol=tol_val)
    @test isapprox(c(2.712), f(2.712), atol=tol_val)

    function f2(x)
        return log(x) + x^2
    end

    c2 = ChebyshevQuantum.Interp.make_cheb(f2 , a=2, b = 3)

    @test isapprox(c2(2.3), f2(2.3), atol=tol_val)
    @test isapprox(c2(2.0), f2(2.0), atol=tol_val)
    @test isapprox(c2(2.712), f2(2.712), atol=tol_val)

end

@testset "test functions var N" begin

    function f1(x)
        return exp(-0.2*sin(3*x))
    end
    function f2(x)
        return exp(-0.3*cos(3*x)) + 20
    end

    c1 = ChebyshevQuantum.Interp.make_cheb(f1 , a=2, b = 3)
    c2 = ChebyshevQuantum.Interp.make_cheb(f2 , a=2, b = 3)

    x=2.712
    @test isapprox( (c1+c2)(x), f1(x)+f2(x) , atol=tol_val)
    @test isapprox( (c1-c2)(x), f1(x)-f2(x) , atol=tol_val)
    @test isapprox( (c1^2)(x), f1(x)^2 , atol=tol_val)
    @test isapprox( (c1/c2)(x), f1(x)/f2(x) , atol=tol_val)

    @test isapprox( (-c1)(x), -f1(x) , atol=tol_val)
    @test isapprox( (2*c1)(x), 2*f1(x) , atol=tol_val)
    @test isapprox( (c1*2)(x), 2*f1(x) , atol=tol_val)
    @test isapprox( (2+c1)(x), 2+f1(x) , atol=tol_val)
    @test isapprox( (c1+2)(x), f1(x)+2 , atol=tol_val)

    @test isapprox( (2 / c2)(x), 2 / f2(x) , atol=tol_val)
    @test isapprox( (2 - c2)(x), 2 - f2(x) , atol=tol_val)

end

@testset "test functions fixed N" begin

    function f1(x)
        return exp(-0.2*sin(3*x))
    end
    function f2(x)
        return exp(-0.3*cos(3.5*x)) + 20
    end

    c1 = ChebyshevQuantum.Interp.make_cheb(f1 , a=2, b = 3, N=20, thr=-1)
    c2 = ChebyshevQuantum.Interp.make_cheb(f2 , a=2, b = 3, N=20, thr=-1)

    x=2.712
    @test isapprox( (c1+c2)(x), f1(x)+f2(x) , atol=tol_val)
    @test isapprox( (c1-c2)(x), f1(x)-f2(x) , atol=tol_val)
    @test isapprox( (c1^2)(x), f1(x)^2 , atol=tol_val)
    @test isapprox( (c1/c2)(x), f1(x)/f2(x) , atol=tol_val)

    @test isapprox( (-c1)(x), -f1(x) , atol=tol_val)
    @test isapprox( (2*c1)(x), 2*f1(x) , atol=tol_val)
    @test isapprox( (c1*2)(x), 2*f1(x) , atol=tol_val)
    @test isapprox( (2+c1)(x), 2+f1(x) , atol=tol_val)
    @test isapprox( (c1+2)(x), f1(x)+2 , atol=tol_val)

    @test isapprox( (2 / c2)(x), 2 / f2(x) , atol=tol_val)
    @test isapprox( (2 - c2)(x), 2 - f2(x) , atol=tol_val)

end

@testset "function returns complex number " begin

    function f(x)
        return exp(-1.0im*x)
    end

    c = ChebyshevQuantum.Interp.make_cheb(f , a=-pi, b = pi)

    x= pi/8
    @test isapprox( real(c(x)), real(f(x)) , atol=tol_val)
    @test isapprox( imag(c(x)), imag(f(x)) , atol=tol_val)

    x= 0.1234
    @test isapprox( real(c(x)), real(f(x)) , atol=tol_val)
    @test isapprox( imag(c(x)), imag(f(x)) , atol=tol_val)

    x= 0.5678
    @test isapprox( real(c)(x), real(f(x)) , atol=tol_val)
    @test isapprox( imag(c)(x), imag(f(x)) , atol=tol_val)

end


using QuantEcon
using Statistics
using LinearAlgebra

include(string(pwd(),"/src/gensys.jl"))

function log_like_dsge(par,data;kalman_tol = 1e-10)
    #order to par
    #alfa
    #beta
    #epsilon
    #theta
    #sig
    #sigma: this is the std dev of the innovation
    #phi
    #phi_pi
    #phi_y
    #rho_v

    #data will have t x p dimension: lines are periods p are variables

    alfa = par[1]
    bet = par[2]
    epsilon = par[3]
    theta = par[4]
    sig = par[5]
    #par 6 is coded bellow see line 56
    phi = par[7]
    phi_pi = par[8]
    phi_y = par[9]
    rho_v = par[10]

    THETA = (1-alfa)/(1-alfa+alfa*epsilon)
    lamb = (1-theta)*(1-bet*theta)/theta*THETA
    kappa = lamb*(sig+(phi+alfa)/(1-alfa))

    nobs = size(data,1)

    GAMMA_0 = [bet  0    0  0;
               1    sig  0  0;
               0    0    0  0;
               0    0    0  1]

    GAMMA_1 = [ 1       -kappa  0   0;
                0        sig    1   0;
               -phi_pi  -phi_y  1  -1;
                0        0      0   rho_v]

    PSI = [0; 0; 0; 1]

    PI = [bet  0;
          1    sig;
          0    0;
          0    0]

    p = size(GAMMA_1,1) #number of endogenous vars

    sol = gensys(GAMMA_0,GAMMA_1,PSI,PI; verbose = false)
    if sum(sol.eu) != 2
        return -99999999999
    end

    #Sig = zeros(p,p)
    #Sig[4,4] = par[6]

    G = zeros(1,p)
    G[1,2] = 1

    A = sol.Theta1
    R = [0] .+ 1e-8
    Q = par[6]^2*sol.Theta2*sol.Theta2'

    kalman_res = Kalman(A,G,Q,R) #create a Kalman filter instance

    y_mean = mean(data,dims=1)
    y_var = var(data,dims=1)
    y_var = reshape(y_var,size(y_var,2))
    #y_var = repeat([1],p)

    x_hat = repeat(y_mean,p) #initial mean of the state
    x_var = diagm(repeat(y_var,p))#variance initial of state
    #x_var = x_var*x_var'
    set_state!(kalman_res,x_hat,x_var)

    fit = zeros(nobs,4)

    llh = zeros(nobs)

    for j in 1:(nobs-1)
        med = kalman_res.cur_x_hat
        fit[j,:] = med
        varian = kalman_res.cur_sigma
        #println(det(varian))
        eta = data[j+1,:] - kalman_res.G*med #mean loglike
        P = kalman_res.G*varian*kalman_res.G' + kalman_res.R
        teste_cond = 1/cond(P)
        if teste_cond < kalman_tol
            llh[j] = -500
        #elseif det(P) <= 0
        #    llh[j] = -500
        else
            llh[j] = -(p*log(2*pi) + logdet(P) .+ eta'*inv(P)*eta)/2
            QuantEcon.update!(kalman_res,data[j+1,:]) #updating the kalman estimates
        end
        #println(kalman_res.cur_sigma)
        #println(det(varian))
        #println(P)
    end
    llh = llh[10:length(llh)]
    return sum(llh)
end


using ITensors
using LinearAlgebra
using QuadGK

function ising_mpo(
  pair_sₕ::Pair{<:Index,<:Index},
  pair_sᵥ::Pair{<:Index,<:Index},
  β::Real,
  J::Real=1.0;
  sz::Bool=false,
)
  sₕ, sₕ′ = pair_sₕ
  sᵥ, sᵥ′ = pair_sᵥ
  @assert dim(sₕ) == dim(sᵥ)
  d = dim(sₕ)
  T = ITensor(sₕ, sₕ′, sᵥ, sᵥ′)
  for i in 1:d
    T[i, i, i, i] = 1.0
  end
  if sz
    T[1, 1, 1, 1] = -1.0
  end
  s̃ₕ, s̃ₕ′, s̃ᵥ, s̃ᵥ′ = sim.((sₕ, sₕ′, sᵥ, sᵥ′))
  T̃ = T * δ(sₕ, s̃ₕ) * δ(sₕ′, s̃ₕ′) * δ(sᵥ, s̃ᵥ) * δ(sᵥ′, s̃ᵥ′)

  # Alternative method
  #Q = [exp(β * J) exp(-β * J); exp(-β * J) exp(β * J)]
  #X = √Q

  f(λ₊, λ₋) = [
    (λ₊ + λ₋)/2 (λ₊ - λ₋)/2
    (λ₊ - λ₋)/2 (λ₊ + λ₋)/2
  ]
  λ₊ = √(exp(β * J) + exp(-β * J))
  λ₋ = √(exp(β * J) - exp(-β * J))
  X = f(λ₊, λ₋)
  Xₕ = itensor(vec(X), s̃ₕ, sₕ)
  Xₕ′ = itensor(vec(X), s̃ₕ′, sₕ′)
  Xᵥ = itensor(vec(X), s̃ᵥ, sᵥ)
  Xᵥ′ = itensor(vec(X), s̃ᵥ′, sᵥ′)
  return T̃ * Xₕ′ * Xᵥ′ * Xₕ * Xᵥ
end

function ising_mpo(sₕ::Index, sᵥ::Index, args...; kwargs...)
  return ising_mpo(sₕ => sₕ', sᵥ => sᵥ', args...; kwargs...)
end

function ising_mpo_dual(
  sh::Tuple{Index,Index}, sv::Tuple{Index,Index}, β::Real, J::Real=1.0
)
  d = dim(sh[1])
  T = ITensor(sh[1], sh[2], sv[1], sv[2])
  sig(s) = 1.0 - 2.0 * (s - 1)
  E0 = -4.0
  for s1 in 1:d, s2 in 1:d, s3 in 1:d, s4 in 1:d
    E = sig(s1) * sig(s2) + sig(s2) * sig(s3) + sig(s3) * sig(s4) + sig(s4) * sig(s1)
    val = exp(-β * (E - E0))
    T[sh[1] => s1, sv[2] => s2, sh[2] => s3, sv[1] => s4] = val
  end
  return T
end

function ising_partition(sh, sv, β)
  ny, nx = size(sh)
  T = Matrix{ITensor}(undef, ny, nx)
  for iy in 1:ny, ix in 1:nx
    ixp = per(ix + 1, nx)
    iyp = per(iy + 1, ny)
    T[iy, ix] = ising_mpo(sh[iy, ix] => sh[iy, ixp], sv[iy, ix] => sv[iyp, ix], β)
  end
  return T
end

#
# Exact results
#

const βc = 0.5 * log(√2 + 1)

function ising_free_energy(β::Real, J::Real=1.0)
  k = β * J
  c = cosh(2 * k)
  s = sinh(2 * k)
  xmin = 0.0
  xmax = π
  integrand(x) = log(c^2 + √(s^4 + 1 - 2 * s^2 * cos(x)))
  integral, err = quadgk(integrand, xmin, xmax)::Tuple{Float64,Float64}
  return -(log(2) + integral / π) / (2 * β)
end

function ising_magnetization(β::Real)
  β > βc && return (1 - sinh(2 * β)^(-4))^(1 / 8)
  return 0.0
end


using Images
using Colors
using FileIO
using LinearAlgebra
using FileIO
using Printf

const Vec2 = Array{Float64, 1}
const Vec3 = Array{Float64, 1}

# function RadicalInverse_VdC(bits::Int16)
# {
#     bits = (bits << 16u) | (bits >> 16u);
#     bits = ((bits & 0x55555555u) << 1u) | ((bits & 0xAAAAAAAAu) >> 1u);
#     bits = ((bits & 0x33333333u) << 2u) | ((bits & 0xCCCCCCCCu) >> 2u);
#     bits = ((bits & 0x0F0F0F0Fu) << 4u) | ((bits & 0xF0F0F0F0u) >> 4u);
#     bits = ((bits & 0x00FF00FFu) << 8u) | ((bits & 0xFF00FF00u) >> 8u);
#     return float(bits) * 2.3283064365386963e-10;
# }

function VanDerCorpus(n::Int64, base::Int64)::Float64
    invBase = 1.0 / float(base);
    denom   = 1.0;
    result  = 0.0;

    for i = 0:31
        if(n > 0)
            denom   = mod(float(n), 2.0);
            result += denom * invBase;
            invBase = invBase / 2.0;
            n       = convert(Int64, floor(float(n) / 2.0));
        end
    end

    return result;
end

function Hammersley(i::Int64, N::Int64)::Vec2
    return [convert(Float64,i) / convert(Float64, N), VanDerCorpus(i, 2)]
end

function ImportanceSampleGGX(Xi::Vec2, roughness::Float64, N::Vec3)::Vec3
    a = roughness*roughness;

    phi = 2.0 * pi * Xi[1];
    cosTheta = sqrt((1.0 - Xi[2]) / (1.0 + (a*a - 1.0) * Xi[2]));
    sinTheta = sqrt(1.0 - cosTheta*cosTheta);

    #  from spherical coordinates to cartesian coordinates
    H = [cos(phi) * sinTheta, sin(phi) * sinTheta, cosTheta]

    #  from tangent-space vector to world-space sample vector
    up = abs(N[3]) < 0.999 ? [0.0, 0.0, 1.0] : [1.0, 0.0, 0.0];
    tangent = normalize(cross(up, N));
    bitangent = cross(N, tangent);

    sampleVec = tangent * H[1] + bitangent * H[2] + N * H[3];
    return normalize(sampleVec);
end

function GeometrySchlickGGX(NdotV::Float64, roughness::Float64)::Float64
    a = roughness;
    k = (a * a) / 2.0;

    nom = NdotV;
    denom = NdotV * (1.0 - k) + k;

    return nom / denom;
end

function GeometrySmith(roughness::Float64, NoV::Float64, NoL::Float64)::Float64
    ggx2 = GeometrySchlickGGX(NoV, roughness);
    ggx1 = GeometrySchlickGGX(NoL, roughness);
    return ggx1 * ggx2;
end

function integrateBRDF(NdotV::Float64, roughness::Float64, samples::Int64)::Vec2
    V = [sqrt(1.0 - NdotV * NdotV), 0.0, NdotV]
    A = 0.0;
    B = 0.0;
    N = [0.0, 0.0, 1.0];

    for i = 0:samples-1
	Xi = Hammersley(i, samples);
	H = ImportanceSampleGGX(Xi, roughness, N);
	L = normalize(2.0 * dot(V, H) * H - V);

	NoL = max(L[3], 0.0);
	NoH = max(H[3], 0.0);
	VoH = max(dot(V, H), 0.0);
	NoV = max(dot(N, V), 0.0);

	if NoL > 0.0
	    G = GeometrySmith(roughness, NoV, NoL);

	    G_Vis = (G * VoH) / (NoH * NoV);
	    Fc = (1.0 - VoH)^5.0;

	    A += (1.0 - Fc) * G_Vis;
	    B += Fc * G_Vis;
	end
    end

    return [A / convert(Float64, samples), B / convert(Float64, samples)]
end

function writeLine(x, data, size, samples)
    for y = 0:size-1
	NoV = (y + 0.5) * (1.0 / size);
	roughness = (x + 0.5) * (1.0 / size);
        rg = integrateBRDF(NoV, roughness, samples)
        data[x+1,y+1] = RGB{N0f16}(rg[1], rg[2], 0.0)
    end
end

function main()
    samples = 1000
    size = 512
    progress::Int64 = 0

    data = rand(RGB{N0f16}, size, size)
    for x = 0:size-1
        writeLine(x, data, size, samples)
        progressNew = convert(Int64, ceil(100 * x / size))

        if (progressNew > progress)
            progress = progressNew
            if (progress % 10 == 0)
                @printf("%d%% ", progress)
            end
        end
    end

    save(File(format"PNG", "brdfLUT.png"), data)

    @printf("\nSaved to brdfLUT.png (%u samples)\n", samples)
end

@time main()


include("../utils.jl")

@testset "Symplectic" begin
    @testset "Real" begin
        Sp_2 = Symplectic(2 * 1)
        Metr_Sp_2 = MetricManifold(Sp_2, RealSymplecticMetric())

        p_2 = [0.0 1.0/2.0; -2.0 -2.0]
        X1 = [
            -0.121212 0.121212
            0.969697 -1.0
        ]
        X2 = [
            0.0 0.0
            0.0 -1.0
        ]

        Sp_6 = Symplectic(6)
        points = [
            [
                1 1 3 0 0 0
                -1 -2 -4 0 0 0
                3 0 7 0 0 0
                3 4 10 14 5 -6
                -5 -9 -19 7 2 -3
                0 0 0 -2 -1 1
            ],
            [
                0 0 0 -2 1 5
                0 0 0 0 -1 -2
                0 0 0 -1 -2 -3
                -1 2 -1 9 -4 -21
                -7 11 -5 -2 10 24
                3 -4 2 -2 -7 -13
            ],
            [
                6 -1 11 -10 8 6
                2 0 3 23 -14 -14
                -3 0 -5 -1 5 1
                0 0 0 0 -1 0
                0 0 0 5 -3 -3
                0 0 0 3 -4 -2
            ],
            [
                2 -1 5 -4 4 2
                2 0 3 11 -6 -6
                -1 0 -2 -4 7 3
                0 0 0 0 -1 0
                0 0 0 2 -1 -1
                0 0 0 3 -4 -2
            ],
        ]

        large_tr_norm_points = [
            [
                0 -3 -5 0 0 0
                -2 -3 4 0 0 0
                -3 -5 5 0 0 0
                11 27 -5 5 -2 1
                -37 -29 117 40 -15 9
                22 30 -47 -27 10 -6
            ],
            [
                0 0 0 7 -2 1
                0 0 0 19 -5 3
                0 0 0 -6 2 -1
                -1 1 8 -14 4 -2
                0 -1 -2 -94 26 -15
                -1 -2 3 45 -11 7
            ],
        ]

        @testset "Basics" begin
            @test repr(Sp_2) == "Symplectic{$(2), ℝ}()"
            @test representation_size(Sp_2) == (2, 2)
            @test base_manifold(Sp_2) === Sp_2

            @test is_point(Sp_2, p_2)
            @test_throws DomainError is_point(Sp_2, p_2 + I, true)

            @test is_vector(Sp_2, p_2, X1; atol=1.0e-6)
            @test is_vector(Sp_2, p_2, X2; atol=1.0e-12)
            @test is_vector(Sp_2, p_2, X1 + X2; atol=1.0e-6)
            @test_throws DomainError is_vector(Sp_2, p_2, X1 + [0.1 0.1; -0.1 0.1], true)
        end
        @testset "Symplectic Inverse" begin
            I_2n = Array(I, 2, 2)
            @test Manifolds.symplectic_inverse_times(Sp_2, p_2, p_2) == I_2n
            @test Manifolds.symplectic_inverse_times!(Sp_2, copy(p_2), p_2, p_2) == I_2n
            @test inv(Sp_2, p_2) * p_2 == I_2n
            @test inv!(Sp_2, copy(p_2)) * p_2 == I_2n
        end
        @testset "Embedding" begin
            x = [0.0 1.0/2.0; -2.0 -2.0]
            y = similar(x)
            z = embed(Sp_2, x)
            @test z == x

            Y = similar(X1)
            embed!(Sp_2, Y, p_2, X1)
            @test Y == X1
        end
        @testset "Retractions and Exponential Mapping" begin
            q_exp = [
                -0.0203171 0.558648
                -1.6739 -3.19344
            ]
            @test isapprox(exp(Sp_2, p_2, X2), q_exp; atol=1.0e-5)
            @test isapprox(
                retract(Sp_2, p_2, X2, ExponentialRetraction()),
                q_exp;
                atol=1.0e-5,
            )

            q_cay = [
                0.0 0.5
                -2.0 -3.0
            ]
            @test retract(Sp_2, p_2, X2) == q_cay
            @test retract(Sp_2, p_2, X2, CayleyRetraction()) == q_cay

            X_inv_cayley_retraction = inverse_retract(Sp_2, p_2, q_cay)
            X_inv_cayley_retraction_2 =
                inverse_retract(Sp_2, p_2, q_cay, CayleyInverseRetraction())
            @test X_inv_cayley_retraction == X_inv_cayley_retraction_2
            @test X_inv_cayley_retraction ≈ X2
        end
        @testset "Riemannian metric" begin
            X1_p_norm = 0.49259905148939337
            @test norm(Sp_2, p_2, X1) == X1_p_norm
            @test norm(Sp_2, p_2, X1) == √(inner(Sp_2, p_2, X1, X1))

            X2_p_norm = 1 / 2
            @test norm(Sp_2, p_2, X2) == X2_p_norm
            @test norm(Sp_2, p_2, X2) == √(inner(Sp_2, p_2, X2, X2))

            q_2 = retract(Sp_2, p_2, X2, ExponentialRetraction())
            approximate_p_q_geodesic_distance = 0.510564444555605
            @test isapprox(
                distance(Sp_2, p_2, q_2),
                approximate_p_q_geodesic_distance;
                atol=1e-14,
            )

            # Project tangent vector into (T_pSp)^{\perp}:
            Extended_Sp_2 = MetricManifold(get_embedding(Sp_2), ExtendedSymplecticMetric())
            proj_normal_X2 = Manifolds.project_normal!(Extended_Sp_2, copy(X2), p_2, X2)
            @test isapprox(proj_normal_X2, zero(X2); atol=1.0e-16)

            # Project Project matrix A ∈ ℝ^{2 × 2} onto (T_pSp):
            A_2 = [5.0 -21.5; 3.14 14.9]
            A_2_proj = similar(A_2)
            Manifolds.project!(Extended_Sp_2, A_2_proj, p_2, A_2)
            @test is_vector(Sp_2, p_2, A_2_proj; atol=1.0e-16)

            # Change representer of A onto T_pSp:
            @testset "Change Representer" begin
                A_2_representer = change_representer(
                    MetricManifold(get_embedding(Sp_2), ExtendedSymplecticMetric()),
                    EuclideanMetric(),
                    p_2,
                    A_2,
                )
                @test isapprox(
                    inner(Sp_2, p_2, A_2_representer, X1),
                    tr(A_2' * X1);
                    atol=1.0e-12,
                )
                @test isapprox(
                    inner(Sp_2, p_2, A_2_representer, X2),
                    tr(A_2' * X2);
                    atol=1.0e-12,
                )
                @test isapprox(
                    inner(Sp_2, p_2, A_2_representer, A_2),
                    norm(A_2)^2;
                    atol=1.0e-12,
                )
            end
        end
        @testset "Generate random points/tangent vectors" begin
            M_big = Symplectic(20)
            p_big = rand(M_big)
            @test is_point(M_big, p_big, true; atol=1.0e-12)
            X_big = rand(M_big; vector_at=p_big)
            @test is_vector(M_big, p_big, X_big, true; atol=1.0e-12)
        end
        @testset "test_manifold(Symplectic(6), ...)" begin
            @testset "Type $(Matrix{Float64})" begin
                type = Matrix{Float64}
                pts = convert.(type, points)
                test_manifold(
                    Sp_6,
                    cat(pts, large_tr_norm_points; dims=1);
                    retraction_methods=[CayleyRetraction(), ExponentialRetraction()],
                    default_retraction_method=CayleyRetraction(),
                    default_inverse_retraction_method=CayleyInverseRetraction(),
                    test_inplace=true,
                    is_point_atol_multiplier=1.0e8,
                    is_tangent_atol_multiplier=1.0e6,
                    retraction_atol_multiplier=1.0e4,
                    test_project_tangent=true,
                    test_injectivity_radius=false,
                    test_exp_log=false,
                    test_representation_size=true,
                )
            end

            TEST_FLOAT32 && @testset "Type $(Matrix{Float32})" begin
                type = Matrix{Float64}
                pts = convert.(type, points)
                test_manifold(
                    Sp_6,
                    pts;
                    retraction_methods=[CayleyRetraction(), ExponentialRetraction()],
                    default_retraction_method=CayleyRetraction(),
                    default_inverse_retraction_method=CayleyInverseRetraction(),
                    test_inplace=true,
                    is_point_atol_multiplier=1.0e8,
                    is_tangent_atol_multiplier=1.0e6,
                    retraction_atol_multiplier=1.0e4,
                    test_project_tangent=true,
                    test_injectivity_radius=false,
                    test_exp_log=false,
                    test_representation_size=true,
                )
            end

            TEST_STATIC_SIZED && @testset "Type $(MMatrix{6, 6, Float64, 36})" begin
                type = MMatrix{6,6,Float64,36}
                pts = convert.(type, points)
                test_manifold(
                    Sp_6,
                    pts;
                    retraction_methods=[CayleyRetraction(), ExponentialRetraction()],
                    default_retraction_method=CayleyRetraction(),
                    default_inverse_retraction_method=CayleyInverseRetraction(),
                    test_inplace=true,
                    is_point_atol_multiplier=1.0e7,
                    is_tangent_atol_multiplier=1.0e6,
                    retraction_atol_multiplier=1.0e4,
                    test_project_tangent=false, # Cannot solve 'sylvester' for MMatrix-type.
                    test_injectivity_radius=false,
                    test_exp_log=false,
                    test_representation_size=true,
                )
            end
        end

        @testset "Gradient Computations" begin
            test_f(p) = tr(p)
            Q_grad = SymplecticMatrix(points[1])
            analytical_grad_f(p) = (1 / 2) * (p * Q_grad * p * Q_grad + p * p')

            p_grad = convert(Array{Float64}, points[1])
            fd_diff = RiemannianProjectionBackend(Manifolds.FiniteDifferencesBackend())

            @test isapprox(
                Manifolds.gradient(Sp_6, test_f, p_grad, fd_diff),
                analytical_grad_f(p_grad);
                atol=1.0e-9,
            )
            @test isapprox(
                Manifolds.gradient(Sp_6, test_f, p_grad, fd_diff; extended_metric=false),
                analytical_grad_f(p_grad);
                atol=1.0e-9,
            )

            grad_f_p = similar(p_grad)
            Manifolds.gradient!(Sp_6, test_f, grad_f_p, p_grad, fd_diff)
            @test isapprox(grad_f_p, analytical_grad_f(p_grad); atol=1.0e-9)

            Manifolds.gradient!(
                Sp_6,
                test_f,
                grad_f_p,
                p_grad,
                fd_diff;
                extended_metric=false,
            )
            @test isapprox(grad_f_p, analytical_grad_f(p_grad); atol=1.0e-9)
        end
    end

    @testset "SymplecticMatrix" begin
        # TODO: Test for different type matrices.
        @test SymplecticMatrix() == SymplecticMatrix(1)
        Sp_4 = Symplectic(4)
        pQ_1 = [
            0 0 -2 3
            0 0 1 -1
            -1 -1 0 0
            -3 -2 4 -4
        ]
        pQ_2 = [
            0 0 -2 5
            0 0 1 -2
            -2 -1 0 0
            -5 -2 4 -8
        ]
        p_odd_row = [
            0 0 -1 3
            0 0 1 0
            -2 -1 0 0
        ]
        p_odd_col = [
            0 -2 5
            0 1 -2
            -1 0 0
            -2 4 -8
        ]
        Q = SymplecticMatrix(pQ_1, pQ_2)
        Q2 = SymplecticMatrix(1)

        @testset "Type Basics" begin
            @test Q == Q2
            @test ndims(Q) == 2
            @test copy(Q) == Q
            @test eltype(SymplecticMatrix(1 // 1)) == Rational{Int64}
            @test convert(SymplecticMatrix{Float64}, Q) == SymplecticMatrix(1.0)
            @test "$Q" == "SymplecticMatrix{Int64}(): 1*[0 I; -I 0]"
            @test (
                "$(SymplecticMatrix(1 + 2im))" ==
                "SymplecticMatrix{Complex{Int64}}(): (1 + 2im)*[0 I; -I 0]"
            )
        end

        @testset "Matrix Operations" begin
            @test -Q == SymplecticMatrix(-1)
            @test (2 * Q) * (5 // 6) == SymplecticMatrix(5 // 3)

            @testset "Powers" begin
                @test inv(Q) * Q == I
                @test (
                    inv(SymplecticMatrix(-4.0 + 8im)) * SymplecticMatrix(-4.0 + 8im) ==
                    UniformScaling(1.0 + 0.0im)
                )
                @test Q * Q == -I
                @test Q^2 == -I
                @test Q^3 == -Q
                @test Q^4 == I
            end
            @testset "Addition (subtraction)" begin
                @test Q + Q == 2 * Q
                @test Q - SymplecticMatrix(1.0) == SymplecticMatrix(0.0)
                @test Q + pQ_1 == [
                    0 0 -1 3
                    0 0 1 0
                    -2 -1 0 0
                    -3 -3 4 -4
                ]
                @test Q - pQ_1 == [
                    0 0 3 -3
                    0 0 -1 2
                    0 1 0 0
                    3 1 -4 4
                ]
                @test pQ_1 - Q == [
                    0 0 -3 3
                    0 0 1 -2
                    0 -1 0 0
                    -3 -1 4 -4
                ]
                @test (pQ_1 + Q) == (Q + pQ_1)

                @test_throws ArgumentError Q + p_odd_row
            end
            @testset "Transpose-Adjoint" begin
                @test Q' == SymplecticMatrix(-1)
                @test transpose(SymplecticMatrix(10)) == SymplecticMatrix(-10)
                @test transpose(SymplecticMatrix(1 - 2.0im)) == SymplecticMatrix(-1 + 2.0im)
                @test adjoint(Q) == -Q
                @test adjoint(SymplecticMatrix(1 - 2.0im)) == SymplecticMatrix(-1 - 2.0im)
                @test adjoint(SymplecticMatrix(-1im)) == SymplecticMatrix(-1im)
                @test adjoint(SymplecticMatrix(2.0)) == SymplecticMatrix(-2.0)
            end
            @testset "Inplace mul!" begin
                z1 = [1 + 2im; 1 - 2im]
                @test lmul!(Q, copy(z1)) == Q * z1
                @test lmul!(Q, copy(p_odd_col)) == [
                    -1 0 0
                    -2 4 -8
                    0 2 -5
                    0 -1 2
                ]
                @test_throws ArgumentError lmul!(Q, copy(p_odd_row))

                @test rmul!(copy(z1'), Q) == z1' * Q
                @test rmul!(copy(p_odd_row), Q) == [
                    1 -3 0 0
                    -1 0 0 0
                    0 0 -2 -1
                ]
                @test_throws ArgumentError rmul!(copy(p_odd_col), Q)

                z1_copy = copy(z1)
                mul!(z1_copy, Q, z1)
                @test z1_copy == Q * z1

                @test_throws ArgumentError mul!(copy(p_odd_row), Q, p_odd_row)
                @test_throws ArgumentError mul!(copy(p_odd_row), p_odd_col, Q)

                @test_throws ArgumentError Q * p_odd_row
                @test_throws ArgumentError p_odd_col * Q

                mul!(z1_copy', z1', Q)
                @test z1_copy' == (z1' * Q)
            end
        end

        @testset "Symplectic Inverse Ops." begin
            @test ((Q' * pQ_1' * Q) * pQ_1 - I) == zeros(eltype(pQ_1), size(pQ_1)...)
        end
    end
end


using Plots, GLPlot; GLPlot.init()
using Reactive, GeometryTypes, Colors, GLAbstraction

function mandelbulb{T}(x0::T,y0::T,z0::T, n, iter)
    x,y,z = x0,y0,z0
    for i=1:iter
        r = sqrt(x*x + y*y + z*z)
        theta = atan2(sqrt(x*x + y*y) , z)
        phi = atan2(y,x)
        rn = r^n
        x1 = rn * sin(theta*n) * cos(phi*n) + x0
        y1 = rn * sin(theta*n) * sin(phi*n) + y0
        z1 = rn * cos(theta*n) + z0
        (x1*x1 + y1*y1 + z1*z1) > n && return T(i)
        x,y,z = x1,y1,z1
    end
    T(iter)
end
dims = (50, 50, 50)

x, y, z = ntuple(3) do i
    # linearly spaced array (not dense) from -1 to 1
    reshape(linspace(-1f0, 1f0, dims[i]), ntuple(j-> j == i ? dims[i] : 1, 3))
end
volume = zeros(Float32, dims)
# create two sliders to interact with the 2 parameters of the mandelbulb function
itslider = GLPlot.play_widget(1:15)
nslider = GLPlot.play_widget(linspace(1f0, 30f0, 100))

# register a callback to the sliders with map ("map over updates")
mandelvol_s = map(nslider, itslider) do n, it
    volume .= mandelbulb.(x, y, z, n, it)
end

glplot(mandelvol_s, color_norm = Vec2f0(0, 50))

function iso_particle(v, isoval)
    particles = Point3f0[]
    sz = 1f0 ./ (Point3f0(size(v)) - 1f0)
    @inbounds for z = 1:size(v, 3), y = 1:size(v, 2), x = 1:size(v, 1)
        if v[x, y, z] ≈ isoval
            push!(particles, (Point3f0(x, y, z) - 1f0) .* sz)
        end
    end
    particles
end
using GLAbstraction, Meshing
isoval = GLPlot.play_widget(1:20)

task, particles = GLPlot.async_map2(iso_particle, Point3f0[], mandelvol_s, isoval)

# share buffer for colors
pp = GLBuffer(particles)
glplot((Sphere(Point2f0(0), 0.005f0), pp), color = pp)

mesh = map(mandelvol_s, isoval) do vol, iso
    mesh = GLNormalMesh(vol, Float32(iso))
    if isempty(vertices(mesh))
        return GLNormalMesh(HyperRectangle(Vec3f0(0), Vec3f0(1)))
    end
    bb = AABB(mesh.vertices)
    mini = minimum(bb)
    w = widths(bb)
    map!(mesh.vertices) do v
        (v - mini) ./ w
    end
    mesh
end
glplot(mesh)


export SoftmaxLoss, getSoftMaxLoss, getLabels, getMisfit

"""
sofmax loss function
"""
mutable struct SoftmaxLoss{T}
    shift :: T
    theta :: T
    addBias::Bool
end

function getSoftMaxLoss(TYPE::Type,shift=zero(TYPE), theta=convert(TYPE,1e-3), addBias=true)
	return SoftmaxLoss{TYPE}(shift,theta,addBias);
end

import Base.display
display(this::SoftmaxLoss) = println("SoftmaxLoss(shift=$(this.shift),theta=$(this.theta),addBias=$(this.addBias))")

function getMisfit(this::SoftmaxLoss{T},W::Array{T},Y::Array{T},C::Array{T},doDY=true,doDW=true) where {T}

    dWF = zeros(T,0); d2WF = zeros(T,0); dYF =zeros(T,0); d2YF = zeros(T,0)

    szY  = size(Y)
    nex  = szY[2]
    if this.addBias
        Y   = [Y; ones(T,1,nex)]
    else
        Y   = copy(Y)  # make sure we don't overwrite features
    end

    szW  = [size(C,1),size(Y,1)]

    W   = reshape(W,tuple(szW...))
    Y   -= this.shift;

    S   = exp.(W*Y)

    Cp   = getLabels(this,S)
    err  = countnz(C-Cp)/2
    F    = -sum(sum(C.*(W*Y))) + sum(log.(sum(S,1)))
    para = [F,nex,err]
    F   /= nex

    if doDW || doDY
        dF   = -C + S./sum(S,1)
        d2F  = U -> this.theta *U + (U.*S)./sum(S,1) -
                    S.*(repmat(sum(S.*U,1)./sum(S,1).^2,size(S,1),1))
    end
    if doDW
        dWF  = vec(dF*(Y'/nex))
        matW      = W -> reshape(W,tuple(szW...))
        d2WFmv    = U -> vec((d2F(matW(U/nex)*Y)*Y'))
        d2WF      = LinearOperator(prod(szW),prod(szW),true,true,d2WFmv)
    end
    if doDY
        if this.addBias
            W = W[:,1:end-1]
        end
        dYF  =   vec(W'*dF)/nex

        WI     = T -> W*T
        WIT    = T -> W'*T
        matY   = Y -> reshape(Y,szY)
        d2YFmv = T -> vec(WIT(((d2F(WI(matY(T/nex)))))))

        d2YF = LinearOperator(prod(szY),prod(szY),true,true,d2YFmv)
    end
    return F, para, dWF, d2WF, dYF, d2YF
end

function getLabels(this::SoftmaxLoss,W::Array{T},Y=nothing) where {T}
    if Y==nothing
        S = W
        nex = size(S,2)
    else
        ny = size(Y); nf = ny[1]; nex = ny[2]
        if this.addBias
            W      = reshape(W,:,nf+1)
            Y     = [Y; ones(T,1,nex)]
        end
        Y     = Y - this.shift
        S      = exp.(W*Y)
    end
    P      = S./sum(S,1)
    Cp     = zeros(T,size(P))
    for i=1:size(P,2)
        Cp[indmax(P[:,i]),i] = 1
    end
    return Cp
end


using TullioStatistics
using Test
using Statistics
using StableRNGs

rng = StableRNG(111)

A = rand(rng,100);

@test TullioStatistics.mean(A) ≈ mean(A)


using AbstractGPs
using ConjugateComputationVI
using Distributions
using Optim
using ParameterHandling
using Plots
using Random
using RDatasets
using StatsBase
using StatsFuns
using Zygote

using ConjugateComputationVI: GaussHermiteQuadrature, UnivariateFactorisedLikelihood

# Adjoint for the Poisson logpdf.
# log(λ^x exp(-λ) / x!) =
# x log(λ) - λ - log(x!)
# dλ = x / λ - 1
Zygote.@adjoint function StatsFuns.poislogpdf(λ::Float64, x::Union{Float64, Int})
    function poislogpdf_pullback(Δ::Real)
        return Δ * (x / λ - 1), nothing
    end
    return StatsFuns.poislogpdf(λ, x), poislogpdf_pullback
end

# Download the data.
data = dataset("boot", "coal")

# Bin the data
h = fit(Histogram, data.Date, range(minimum(data.Date), maximum(data.Date); length=200))

# Construct data set to learn on.
x = collect(only(h.edges)[1:end-1])
y = h.weights

# Specify a model.
θ_init = (scale=positive(1.0), stretch=positive(1e-3))
θ_init_flat, unflatten = ParameterHandling.flatten(θ_init)

function build_latent_gp(θ::AbstractVector{<:Real})
    return build_latent_gp(ParameterHandling.value(unflatten(θ)))
end
function build_latent_gp(θ::NamedTuple)
    gp = GP(θ.scale * SEKernel() ∘ ScaleTransform(θ.stretch))
    lik = UnivariateFactorisedLikelihood(f -> Poisson(exp(f)))
    return LatentGP(gp, lik, 1e-9)
end

f_approx_post, results_summary = ConjugateComputationVI.optimize_elbo(
    build_latent_gp,
    GaussHermiteQuadrature(10),
    x,
    y,
    θ_init_flat,
    BFGS(
        alphaguess = Optim.LineSearches.InitialStatic(scaled=true),
        linesearch = Optim.LineSearches.BackTracking(),
    ),
    Optim.Options(
        show_trace = true,
        iterations=25,
        f_calls_limit=50,
    ),
);

function approx_post_marginal_samples(x::AbstractVector, N::Int)
    return map(marginals(f_approx_post(x))) do latent_marginal
        f = rand(latent_marginal, N)
        return exp.(f)
    end
end

function approx_post_95_CI(x::AbstractVector, N::Int)
    samples = approx_post_marginal_samples(x, N)
    return quantile.(samples, Ref((0.025, 0.5, 0.975)))
end

# Plot the predictions.
qs = approx_post_95_CI(x, 10_000);
p1 = plot(
    x, getindex.(qs, 1);
    linealpha=0,
    fillrange=getindex.(qs, 3),
    label="95% CI",
    fillalpha=0.3,
    color=:blue,
);
scatter!(p1, x, y; markersize=2, label="Observations");

p2 = plot(f_approx_post(x, 1e-6); ribbon_scale=3, color=:blue, label="approx post latent");
sampleplot!(f_approx_post(x, 1e-6), 10; color=:blue);

plot(p1, p2; layout=(2, 1))


using Distributions

module BayesChp
export bchp

function hazard_func(r, lambda)
    (1.0 / lambda) * ones(size(r))
end

function gauss_pdf(ms, val, sigma = 1)
  res = Float64[]
  for m in ms
    d = Normal(m, sigma)
    pdfgauss = pdf(d, val)
    push!(res, pdfgauss)
  end
  res
end
#println( gauss_pdf([1,2,3], 2.9) )

function estimate_mus(x)
  n = length(x)
  mus = Float64[]
  for i = 1:(n-1)
    subvec = x[(n-i) : end]
    push!(mus, mean(subvec))
  end
  #if n == 1
  #  push!(mus, x[1])
  #end
  mus
end
r = estimate_mus([1, 1, 5, 6])
print(r)

function bchp(X, mu0, mdist)
  T = length(X)
  R = zeros(T+1, T+1)
  R[1,1] = 1
  maxes = zeros(T+1)
  for t = 1 : T-1

    mus = estimate_mus(X[1:(t+1)])
    predprobs = gauss_pdf(mus, X[t], 1.0)

    H = hazard_func(collect(1:t), mdist)

    # Growth probabilities
    R[2:t+1, t + 1] = R[1:t, t] .* predprobs .* (1 - H)

    # Probability that there was a changepoint
    R[1, t+1] = sum( R[1:t, t] .* predprobs .* H )

    # Normalize
    R[:, t+1] = R[:, t+1] ./ sum(R[:, t+1])
    #
    maxes[t] = indmax(R[:, t])
  end
  R
  maxes
end

end
sig = readdlm("D:/MyTemp/gitlocal/recurrency-code/src.jl/sigtest.txt")

R = bchp(sig, 1.1, 100)
writedlm("D:/MyTemp/gitlocal/recurrency-code/src.jl/probs100.txt", R)

# R = bchp(sig, 1.1, 50)
# writedlm("D:/MyTemp/gitlocal/recurrency-code/src.jl/probs50.txt", R)
#
# R = bchp(sig, 1.1, 200)
# writedlm("D:/MyTemp/gitlocal/recurrency-code/src.jl/probs200.txt", R)


# Solves the Galewsky test case for the shallow water equations on a sphere
# of physical radius 6371220m. Involves a shear flow instability of a zonal
# jet triggered by an initial gravity wave.
# reference:
#   Galewsky, Scott and Polvani (2004) Tellus, 56A 429-440

const H₀ = 10000.0

function uθ(θϕr)
  θ,ϕ,r = θϕr
  ϵ     = 1.0e-8
  umax  = 80.0
  ϕ₁    = π/7
  ϕ₂    = π/2 - ϕ₁
  en    = exp(-4.0/((ϕ₂ - ϕ₁)*(ϕ₂ - ϕ₁)))
  u     = 0.0
  if ϕ > ϕ₁ + ϵ && ϕ < ϕ₂ - ϵ
    u = (umax/en)*exp(1.0/((ϕ - ϕ₁)*(ϕ - ϕ₂)))
  end
  u
end

# Initial velocity
function u₀(xyz)
  θϕr = xyz2θϕr(xyz)
  u   = uθ(θϕr)
  spherical_to_cartesian_matrix(θϕr)⋅VectorValue(u,0,0)
end

# Initial fluid depth
function h₀(xyz)
  θϕr   = xyz2θϕr(xyz)
  x,y,z = xyz
  θ,ϕ,r = θϕr
  h     = H₀
  hh    = 120.0
  α     = 1.0/3.0
  β     = 1.0/15.0
  ϕ₂    = π/4
  ni    = 1000
  ϕₚ    = 0.0
  dϕ    = abs(ϕ/ni)
  sgn   = 1.0
  if ϕ < 0.0
    sgn = -1.0
  end
  for i in 1:ni
    ϕₚ   = ϕₚ + sgn*dϕ
    _θϕr = VectorValue(θ,ϕₚ,r)
    u    = uθ(_θϕr)
    _f   = 2.0*Ωₑ*sin(ϕₚ)
    h    = h - rₑ*u*(_f + tan(ϕₚ)*u/rₑ)*dϕ/g
  end
  h = h + hh*cos(ϕ)*exp(-1.0*(θ/α)*(θ/α))*exp(-1.0*((ϕ₂ - ϕ)/β)*((ϕ₂ - ϕ)/β))
  h
end

# Topography
function topography(xyz)
  0.0
end


    A = [1 2 ; 3 4]
    B = [10, 10]
    sum!(B, A)
    sum!(A, B)

    using ReverseDiffSource

        Nn = [20,10,1]
        function ann(l1, l2, x1)
            x2 = l1 * x1
            x2 = max(0., x2) #relu
            x3 = l2 * x2
            x3 = (max(0., x3[1]))^2 #relu
        end

        # ls = Array[ nn.ws...]
        ls = init(Nn)
        dann = rdiff(ann, (ls..., ones(20) ))
        out = dann(ls..., ones(20) )

    nn = S.NN((ls[1], ls[2]), (S.Relu(), S.Relu()))
    ts2 = S.cycle(nn, ones(20), zeros(1))

        out[2]
        ts2.dws[1]

        out[3]
        ts2.dws[2]

        ts2.ds[end]

        out[3] ./ ts2.dws[2]

    r1 = S.calc(nn, ones(20))[1] # 1.16217
    r2 = ann(ls..., ones(20))
    r1-r2

    ls[1][9,1] += 1e-10
    (ann(ls..., ones(20)) - r1) * 1e10  # -0.8600

    nn.ws[1][9,1] += 1e-10
    (S.calc(nn, ones(20)) - r2) * 1e10  # -0.8600

    reload("SimpleNeuralNets"); S=SimpleNeuralNets

    nn = S.NN([10,10,10,10,1], [ S.Tanh(), S.Tanh()] )
    ts = S.TrainState(nn, 10)

    size(nn)
    size(ts)

    ts2 = S.cycle(nn, ones(20,10), ones(5,10))
    S.cycle!(ts, nn, ones(10,10), ones(1,10))

    S.depth(nn)

    s0 = S.calc(nn, rand(20))
    copy!(s0.ss[1], ones(20))
    S.compatible(nn, s0)
    S.calc!(nn, s0)

    s0 = S.calc(nn, rand(20,10))
    copy!(s0.ss[1], ones(20))
    S.compatible(nn, s0)
    S.calc!(nn, s0)

    a = (1,2,3)
    typeof(a)
    a[1]
    a[1] = 2

    nn = S.NN((randn(1,1),), (S.Relu(),))
    S.calc(nn, [1.])

    nn = S.NN((randn(1,1),), (S.SoftRelu(),))
    S.calc(nn, [1.])

########### ANN3 #################################"

    import Bokeh
    # using ReverseDiffSource
    reload("SimpleNeuralNets"); s = SimpleNeuralNets

    f(x) = exp(-(x-1)^2/5)*cos(x)*sign(x)

    xs = collect(-5:0.001:5)
    X = Float64[ x for x in xs ]'
    Y = Float64[ f(x) for x in xs ]'

    nn = s.NN([1,10,10,10,1], [ s.Relu(), s.Tanh()] )

    ts = s.TrainingSet(X, Y)
    any(ts.weights .!= 1.0)

    λ0, μ = 1e-3, 0.9
    s.sgd(nn, ts, λ0, μ, freq=200, steps=1000)

    ts2 = s.TrainingSet(X, Y, SqEuclidean(), rand(size(X,2)))
    s.sgd(nn, ts2, λ0, μ, freq=200, steps=1000)

    s.colwise(a,b,c)
    length(a.weights)

    px = s.calc(nn, X)
    sum(s.colwise(ts2.metric, px, ts2.Y))

    λ0, μ = 1e-3, 0.9
    s.sgd(nn, λ0, μ, X, Y, freq=200, steps=10000)
    λ0, μ = 1e-5, 0.999
    s.sgd(nn, λ0, μ, X, Y, freq=200, steps=10000)

    Bokeh.hold(true)
    Bokeh.plot(f, -5:5, "k-")
    px = collect(-5:0.01:5)
    py = s.calc(nn, px')'
    Bokeh.plot(px, py, "b-")
    Bokeh.showplot("/tmp/plot.html")
    Bokeh.hold(false)

    nn = s.NN([1,10,10,10,1], [ s.SoftRelu(), s.Tanh()] )
    λ0, μ = 1e-3, 0.9
    s.sgd(nn, λ0, μ, X, Y, freq=200, steps=10000)
    λ0, μ = 1e-3, 0.
    s.sgd(nn, λ0, μ, X, Y, freq=1000, steps=10000)

    Bokeh.hold(true)
    Bokeh.plot(f, -5:5, "k-")
    px = collect(-4:0.1:4)
    py = s.calc(nn, px')'
    Bokeh.plot(px, py, "b-")
    Bokeh.showplot("/tmp/plot.html")
    Bokeh.hold(false)