using JuMP, EAGO

                     m = Model()

                     EAGO.register_eago_operators!(m)

                     @variable(m, -1 <= x[i=1:3] <= 1)

                     @variable(m, -2.1106331886844685 <= q <= -0.28632222807890817)

                     add_NL_constraint(m, :(softplus(-0.901869745284984 + -0.09697470954184428*softplus(-0.8276381457841007 + 0.36077566442210207*$(x[1]) + 0.4291228542230603*$(x[2]) + 0.43678775223595334*$(x[3])) + 0.09177414115023685*softplus(0.9224048735895125 + -0.5729901252145462*$(x[1]) + -0.6177157080444822*$(x[2]) + -0.37503852018425166*$(x[3]))) + softplus(-0.5704938352741911 + -0.31710381800885923*softplus(-0.8276381457841007 + 0.36077566442210207*$(x[1]) + 0.4291228542230603*$(x[2]) + 0.43678775223595334*$(x[3])) + -0.16638488372553573*softplus(0.9224048735895125 + -0.5729901252145462*$(x[1]) + -0.6177157080444822*$(x[2]) + -0.37503852018425166*$(x[3]))) - $q <= 0.0))

                     @objective(m, Min, q)

                     return m


function OLSestimator(y, x)
    estimate = inv(x' * x) * (x' * y)
    return estimate
end


function Jdist!(J::Vector{Float64},x::Vector{Float64},C::Float64,R0::Float64,
    R1::Float64,Ws::Float64,W1::Float64,rho::Float64,beta::Float64,c0::Float64,
    a1::Float64,a2::Float64,h::Float64)
    J[1] = -a1/(4*a2*Ws^4)*(0.125*(W1-Ws*x[1])+c0) - (-3*W1^4/(8192*Ws^4) +
        W1^3*x[1]^2/(2048*Ws^3) - W1^3*c0/(128*Ws^4) + 3*W1^2*x[1]^2/(4096*Ws^2) -
        3*W1^2*c0^2/(64*Ws^4) - 3*W1*x[1]^3/(2048*Ws) + 3*W1*x[1]^2*c0/(128*Ws^2) -
        3*W1*x[1]*c0^2/(32*Ws^3) + 5*x[1]^4/8192 - x[1]^3*c0/(64*Ws) +
        9*x[1]^2*c0^2/(64*Ws^2) - x[1]*c0^3/(2*Ws^3) + c0^4/(2*Ws^4));
end


##########################
# FunctionCell Interface #
##########################
"""
	whiten([;kwargs])

Constructs an untrained cell that when piped data inside, calculates a transform `W`
of the input data `X` so that `Wᵀcov(X)W=I`.

# Keyword arguments (same as in `MultivariateStats`)
  * `regcoef::Real` is the ratio of variaces preserved in the principal subspace
  * `mean` is the mean vector, can be `nothing`, `0` or precomputed mean vector (default `nothing`)

Read the `MultivariateStats.jl` documentation for more information.
"""
whiten(;kwargs...) = FunctionCell(whiten, (), ModelProperties(), kwtitle("Whitening", kwargs); kwargs...)

############################
# DataCell/Array Interface #
############################
"""
	whiten(x, [;kwargs])

Trains a function cell that when receiving data into will whiten the input
the observations based on the projection matrix W extracted from `x`.
"""
# Training
whiten(x::T where T<:CellData; kwargs...) = whiten(getx!(x); kwargs...)
whiten(x::T where T<:AbstractVector; kwargs...) = whiten(mat(x, LearnBase.ObsDim.Constant{2}()); kwargs...)
whiten(x::T where T<:AbstractMatrix; kwargs...) = begin

	whdata = fit(MultivariateStats.Whitening, getobs(x); kwargs...)

	# Build model properties
	modelprops = ModelProperties(nvars(x), nvars(x))

	# Returned trained cell
	FunctionCell(whiten, Model(whdata, modelprops), kwtitle("Whitening",kwargs))
end

# Execution
whiten(x::T where T<:CellData, model::Model{<:MultivariateStats.Whitening}) =
	datacell(whiten(getx!(x), model), gety(x))
whiten(x::T where T<:AbstractVector, model::Model{<:MultivariateStats.Whitening}) =
	whiten(mat(x, LearnBase.ObsDim.Constant{2}()), model)
whiten(x::T where T<:AbstractMatrix, model::Model{<:MultivariateStats.Whitening}) =
	MultivariateStats.transform(model.data, getobs(x))


@registerN  Krawchouk AbstractCDOP2 p 𝐍
export  Krawchouk

"""
     Krawchouk{p,𝐍}

Also spelled  Krawtchouk,  Kravhcuk,….

References: [Koekoek and Swarttouw §1.10](https://arxiv.org/pdf/math/9602214.pdf);  see  also  [Coleman](https://arxiv.org/pdf/1101.1798.pdf) for a different  parameterization.
"""
Krawchouk

abcde(::Type{<:Krawchouk{p,𝐍}})  where {p,𝐍} = NamedTuple{(:a,:b,:c,:d,:e)}((0, p-1, 0, 1, -p*𝐍))

basis_symbol(::Type{<:Krawchouk{p,𝐍}}) where {p,𝐍} = "kᵖ" * "₍" * sprint(io -> unicode_subscript(io, 𝐍)) * "₎"
Polynomials.domain(::Type{<:Krawchouk{p,𝐍}}) where {p,𝐍} = Polynomials.Interval(0, 𝐍)
weight_function(::Type{<:Krawchouk{p,𝐍}})  where {p,𝐍} = x  -> generalized_binomial(𝐍,x)  * p^x * (1-p)^(𝐍-x)

function k1k0(P::Type{<:Krawchouk}, n::Int)
    one(eltype(P))/(n+1)
end

function classical_hypergeometric(P::Type{<:Krawchouk{p,N}}, n::Int, x) where {p,N}
    monic = classical_hypergeometric(MonicKrawchouk{p,N},n,x)
    monic * kn(P,n)
end

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

@registerN MonicKrawchouk AbstractCDOP2 p 𝐍
export MonicKrawchouk
ϟ(::Type{<:MonicKrawchouk{p,𝐍}}) where {p,𝐍} = Krawchouk{p,𝐍}
ϟ(::Type{<:MonicKrawchouk{p,𝐍,T}}) where {p,𝐍,T} = Krawchouk{p,𝐍,T}
@register_monic(MonicKrawchouk)

function classical_hypergeometric(P::Type{<:MonicKrawchouk{p,N}}, n::Int, x) where {p,N}
    Pochhammer( -N,n) * p^n * pFq((-n, -x), -N, 1/p)
end


#############################
# Algorithms from Section 3 #
#############################
export MultivariateVandermondeMatrix, FindEquations

"""
FindEquations(data::Array{T,2},
              alg::Symbol,
              d::Int64,
              homogeneous_equations::Bool)
              where  {T<:Number}

Finds equations.
* data is a matrix whose colums are the data points in Ω.
* alg is the algorithm that should be used (one of :with_svd, :with_qr, :with_rref).
* d is the degree of the equations.
* homogeneous_equations = true restricts the search space to homogeneous polynomials.
* homogeneous_equations = false computes all polynomials of degree at most d.

Alternatively, you can specifiy the exponents of the monomials.

FindEquations(data::Array{T,2},
              alg::Symbol,
              exponents::Array{Array{Int64,1},1})
              where  {T<:Number}

Here, exponents is an array of exponents.
"""
function FindEquations(data::Array{T,2}, alg::Symbol, exponents::Array{Array{Int64,1},1}) where  {T<:Number}
    M = MultivariateVandermondeMatrix(data, exponents)
    FindEquations(M,alg)
end
function FindEquations(data::Array{T,2}, alg::Symbol, degree::Int64, homogeneous_equations::Bool) where  {T<:Number}
    @assert degree > 0 "The degree must be positive."
    M = MultivariateVandermondeMatrix(data, degree, homogeneous_equations)
    FindEquations(M, alg)
end

"""
MultivariateVandermondeMatrix(data::Array{T},
                              d::Int64,
                              homogeneous_equations::Bool)

Creates a multivariate Vandermonde.
* data is a matrix whose colums are the data points in Ω.
* d is the degree of the monomials.
* homogeneous_equations = true restricts the space of monomials to monomials of degree d.
* homogeneous_equations = false computes all monomials of degree at most d.

Alternatively, you can specifiy the exponents of the monomials.

FindEquations(data::Array{T,2},
              alg::Symbol,
              exponents::Array{Array{Int64,1},1})
              where  {T<:Number}

Here, exponents is an array of exponents.
"""
# MultivariateVandermondeMatrix struct
struct MultivariateVandermondeMatrix
    Vandermonde::Array
    exponents::Vector

    function MultivariateVandermondeMatrix(data::Array{T}, exponents::Vector) where {T<:Number}

        @assert length(unique(length.(exponents))) == 1 "Error: Exponents differ in size."

        n,m = size(data)
        N = length(exponents)
        v = veronese(exponents,T, n)
        U = zeros(T,m,N)
        for i=1:m
            U[i,:] = v(data[:,i])
        end
        new(U, exponents)
    end

    function MultivariateVandermondeMatrix(data::Array{T}, d::Int64,  homogeneous_equations::Bool) where {T<:Number}
        n=size(data,1)
        if homogeneous_equations
            exponents = collect(Combinatorics.multiexponents(n, d))
        else
            exponents = vcat(map(i -> collect(Combinatorics.multiexponents(n,-i)), -d:0)...)
        end
        MultivariateVandermondeMatrix(data, exponents)
    end
end

# This function creates a function v.
# v(x) is the array with all the monomials in the entries of x of degree d
function veronese(exponents::Vector, ::Type{T}, n::Int)  where {T<:Number}
    v = map(exponents) do e
        FP.Polynomial(reshape(e, n, 1), [one(T)])
    end
    function (x::Vector)
        FP.evaluate(v,x)
    end
end

"""
FindEquations(M::MultivariateVandermondeMatrix,
              alg::Symbol)

Finds equations on input a multivariate Vandermonde matrix M.
* alg is one of :with_svd, :with_qr, :with_rref.

You can specifiy the tolerance.

FindEquations(M::MultivariateVandermondeMatrix,
             alg::Symbol,
             tol::Float64)
"""
# function that gets the equations from a MultivariateVandermondeMatrix
function FindEquations(M::MultivariateVandermondeMatrix, alg::Symbol)

    if alg == :with_svd
        m, N = size(M.Vandermonde)
        s = LinearAlgebra.svd(M.Vandermonde, full = true)
        tol = max(m,N) * maximum(s.S) * eps(Float64)
        rk = sum(s.S .> tol)
        return Polynomials_from_coefficients(s.Vt[rk + 1:end,:], M.exponents)
    elseif alg == :with_qr
        return Polynomials_from_coefficients(with_qr(M), M.exponents)
    elseif alg == :with_rref
        return Polynomials_from_coefficients(with_rref(M), M.exponents)
    else
        println("Method $(alg) not known.")
    end

end

function FindEquations(M::MultivariateVandermondeMatrix, alg::Symbol, tol::Float64)

    if alg == :with_svd
        m, N = size(M.Vandermonde)
        s = LinearAlgebra.svd(M.Vandermonde, full = true)
        rk = sum(s.S .> tol)
        return Polynomials_from_coefficients(s.Vt[rk + 1:end,:], M.exponents)
    elseif alg == :with_qr
        return Polynomials_from_coefficients(with_qr(M, tol), M.exponents)
    elseif alg == :with_rref
        return Polynomials_from_coefficients(with_rref(M, tol), M.exponents)
    else
        println("Method $(alg) not known.")
    end

end

# Creates a polynomial from a coefficient vector
function Polynomials_from_coefficients(kernel::Matrix{T}, exponents::Vector) where {T<:Number}
    tol = 1e-10
    l = size(kernel,1)
    nvars = length(exponents[1])
    DynamicPolynomials.@polyvar x_[1:nvars]

    if l == 0
        return [0]
    else
        map(1:l) do i
            non_zero_coeffs = findall(x -> abs(x) > tol, kernel[i,:])
            if length(non_zero_coeffs) > 0
                monomial = map(c -> prod(map(i -> x_[i]^exponents[c][i], 1:nvars)), non_zero_coeffs)
                return LinearAlgebra.dot(kernel[i,non_zero_coeffs],  monomial)
            else
                return 0.0*x[1]
            end
        end
    end
end

function with_qr(M::MultivariateVandermondeMatrix)
    R = LinearAlgebra.qr(M.Vandermonde).R
    m, N = size(R)
    s = LinearAlgebra.svdvals(R)
    tol = max(m,N) * maximum(s) * eps(Float64)
    return kernel_qr(R, tol)
end

function with_qr(M::MultivariateVandermondeMatrix, tol::Float64)
    R = LinearAlgebra.qr(M.Vandermonde).R
    return kernel_qr(R, tol)
end

function kernel_qr(R::Array{T,2}, tol::Float64) where {T <: Number}
    n,m = size(R)

    if n>=m
        index = findall(x -> abs(x) < tol, [R[i,i] for i in 1:m])
        index2 = setdiff([i for i in 1:m], index)
        R_small = R[:,index2]

        kernel = zeros(eltype(R), length(index), m)

        for i = 1:length(index)
            kernel[i,index[i]] = one(eltype(R))
            kernel[i,index2] =  (-1) .* R_small\R[:,index[i]]
        end
        return kernel
    else
        index = findall(x -> abs(x) < tol, [R[i,i] for i in 1:n])
        index2 = setdiff([i for i in 1:m], index)
        R_small = R[:,index2]

        kernel = zeros(eltype(R), length(index), m)

        for i = 1:length(index)
            kernel[i,index[i]] = one(eltype(R))
            kernel[i,index2] =  (-1) .* R_small\R[:,index[i]]
        end
        R_new = [R; LinearAlgebra.pinv(kernel) * kernel]
        s = LinearAlgebra.svd(R_new, full = true)
        rk = sum(s.S .> tol)
        return  [kernel; s.Vt[rk + 1:end,:]]
    end
end

function with_rref(M::MultivariateVandermondeMatrix)
    R = RowEchelon.rref(M.Vandermonde)
    m, N = size(R)
    s = LinearAlgebra.svdvals(R)
    tol = max(m,N) * maximum(s) * eps(Float64)
    return kernel_rref(R, tol)
end
function with_rref(M::MultivariateVandermondeMatrix, tol::Float64)
    R = RowEchelon.rref(M.Vandermonde)
    return kernel_rref(R, tol)
end
function kernel_rref(R::Array{T,2}, tol::Float64) where {T <: Number}
    n,m = size(R)
    R = R[findall([LinearAlgebra.norm(R[i,:]) for i in 1:n] .> sqrt(m) * tol),:]
    rk = size(R,1)
    index = zeros(Int64, rk, 2)
    for i = 1:rk
        where_are_the_ones = findall(abs.(R[i,:]).> tol)
        index[i,:] = [i where_are_the_ones[1]]
    end
    pivots = setdiff([i for i in 1:m], index[:,2])

    kernel = zeros(eltype(R), m-rk, m)
    for i = 1:(m-rk)
        t = findall(index[:,2] .< pivots[i])
        kernel[i,pivots[i]] = 1
        for j in t
            kernel[i, index[j,2]] = - R[index[j,1],pivots[i]]
        end
    end
    return kernel
end


export TR_Cub
function TR_Cub(h :: AbstractNLPModel;
               t₀ :: Float64 = h.meta.x0[1],
               tol :: Float64 = 1e-7,
               maxiter :: Int = 50,
               verbose :: Bool = true,
               eps1 ::Float64 = 0.2,
               eps2 :: Float64 = 0.8,
               red :: Float64 = 0.5,
               aug :: Float64 = 2.0,
               Δ :: Float64 = 1.0)

    (length(h.meta.x0) > 1) && warn("Not a 1-D problem ")
    t = t₀; iter = 0;
    fₖ = obj(h, [t])[1]
    gₖ = grad(h, [t])[1]

    secₖ = 1.0
    dN = -gₖ #pour la première itération
    A = 0.5; B = 0.0

    q(d) = fₖ + gₖ * d + A * d^2 + B * d^3

    verbose && @printf(" iter  t         gₖ          Δ        pred         ared\n")
    verbose && @printf(" %4d %7.2e  %7.2e  %7.2e \n", iter,t,gₖ,Δ)

    while ((abs(gₖ) > tol) & (iter < maxiter)) | (iter == 0)

        if q(Δ) < q(-Δ)
            d = Δ
        else
            d = -Δ
        end

        cub(t) = fₖ + gₖ * t + A * t^2 + B * t^3
        #dcub(t)=gₖ+2*A*t+3*B*t^2
        dR = roots(Poly([gₖ, 2 * A, 3 * B]))

        if isreal(dR)
          dR = real(dR)
          dN2 = dR[1]
          if length(dR) > 1
            if (cub(dR[2]) < cub(dR[1]))
              dN2 = dR[2]
            end
          end
        else
          dN2 = -gₖ * s / y
        end

        if (abs(dN2) < Δ) & (q(d) > q(dN2))
            d = dN2
        end

        if d == 0.0
            return (t, fₖ, norm(gₖ, Inf), iter + 1, false, false, :FailedDirectionComputation, h.counters.neval_obj, h.counters.neval_grad, h.counters.neval_hess)
            #L'OUTIL DE CALCUL DES RACINES DU PACKAGE Polynomials ARRONDIE À 0 CERTAINES RACINES D'OÙ LA NOUVELLE CONDITION
        end

        ftestTR = obj(h, [t + d])[1]
        gtestTR = grad(h, [t + d])[1]

        pred = gₖ * d + A * d^2 + B * d^3

        if pred > - 1e-10
            ared = (gₖ + gtestTR) * d / 2
        else
            ared = ftestTR - fₖ
        end

        ratio = ared / pred
        if (ratio < eps1) & (abs(d) == Δ)
            Δ = red * Δ
        else
          #two points memory
          tpred = t; gkm1 = gₖ; fkm1 = fₖ
          t = t + d; gₖ = gtestTR; fₖ = ftestTR

          #cubic two points interpolation
          s = t - tpred; y = gₖ - gkm1

          α = -s
          z = gₖ + gkm1 + 3 * (fₖ - fkm1) / α
          discr = z^2 - gₖ * gkm1
          denom = gₖ + gkm1 + 2 * z
          B = 1 / 3 * (gₖ + gkm1 +2 * z) / (α * α)
          A = -(gₖ + z) / α

          if ratio > eps2
            Δ = aug * Δ
          end
        end

        iter += 1
        verbose && @printf(" %4d %7.2e  %7.2e  %7.2e %7.2e %7.2e\n", iter,t,gₖ,Δ,pred,ared)
    end

    status = :NotSolved
    (abs(gₖ) < tol) && (status = :Optimal)
    (iter >= maxiter) && (status = :Tired)
    tired = iter > maxiter
    optimal = abs(gₖ) < tol
    return (t, fₖ, norm(gₖ, Inf), iter, optimal, tired, status, h.counters.neval_obj, h.counters.neval_grad, h.counters.neval_hess)

end


module LogNormalTest

using Test
using ForneyLab
using ForneyLab: prod!, unsafeMean, unsafeLogMean, unsafeVar, unsafeLogVar, unsafeCov, unsafeLogCov, outboundType, isApplicable, dims
using ForneyLab: SPLogNormalOutNPP, VBLogNormalOut

@testset "dims" begin
    @test dims(ProbabilityDistribution(Univariate, LogNormal, m=1.0, s=1.0)) == 1
end

@testset "vague" begin
    @test vague(LogNormal) == ProbabilityDistribution(Univariate, LogNormal, m=1.0, s=huge)
end

@testset "unsafe mean and variance" begin
    @test unsafeMean(ProbabilityDistribution(Univariate, LogNormal, m=1.0, s=2.0)) == exp(1.0 + 0.5*2.0)
    @test unsafeLogMean(ProbabilityDistribution(Univariate, LogNormal, m=1.0, s=2.0)) == 1.0
    @test unsafeVar(ProbabilityDistribution(Univariate, LogNormal, m=1.0, s=2.0)) == exp(2.0*1.0 + 2.0)*(exp(2.0) - 1.0)
    @test unsafeLogVar(ProbabilityDistribution(Univariate, LogNormal, m=1.0, s=2.0)) == 2.0
    @test unsafeCov(ProbabilityDistribution(Univariate, LogNormal, m=1.0, s=2.0)) == exp(2.0*1.0 + 2.0)*(exp(2.0) - 1.0)
    @test unsafeLogCov(ProbabilityDistribution(Univariate, LogNormal, m=1.0, s=2.0)) == 2.0
end

@testset "log pdf" begin
    @test isapprox(logPdf(ProbabilityDistribution(Univariate, LogNormal, m=1.2, s=0.5),2), -1.522411904058978)
end

@testset "Gamma approximatons to LogNormal" begin
    @test ForneyLab.laplace(Gamma, ProbabilityDistribution(Univariate, LogNormal, m=0.0, s=2.0)) == ProbabilityDistribution(Univariate, Gamma, a=0.5, b=0.5)
end

@testset "prod!" begin
    @test ProbabilityDistribution(Univariate, LogNormal, m=1.0, s=2.0) * ProbabilityDistribution(Univariate, PointMass, m=1.0) == ProbabilityDistribution(Univariate, PointMass, m=1.0)
    @test_throws Exception ProbabilityDistribution(Univariate, LogNormal, m=1.0, s=2.0) * ProbabilityDistribution(Univariate, PointMass, m=-1.0)
    @test ProbabilityDistribution(Univariate, LogNormal, m=0.0, s=2.0) * ProbabilityDistribution(Univariate, Gamma, a=3.0, b=4.0) == ProbabilityDistribution(Univariate, Gamma, a=2.5, b=4.5)
end

#-------------
# Update rules
#-------------

@testset "SPLogNormalOutNPP" begin
    @test SPLogNormalOutNPP <: SumProductRule{LogNormal}
    @test outboundType(SPLogNormalOutNPP) == Message{LogNormal}
    @test isApplicable(SPLogNormalOutNPP, [Nothing, Message{PointMass}, Message{PointMass}])

    @test ruleSPLogNormalOutNPP(nothing, Message(Univariate, PointMass, m=1.0), Message(Univariate, PointMass, m=2.0)) == Message(Univariate, LogNormal, m=1.0, s=2.0)
end

@testset "VBLogNormalOut" begin
    @test VBLogNormalOut <: NaiveVariationalRule{LogNormal}
    @test outboundType(VBLogNormalOut) == Message{LogNormal}
    @test isApplicable(VBLogNormalOut, [Nothing, ProbabilityDistribution, ProbabilityDistribution])
    @test !isApplicable(VBLogNormalOut, [ProbabilityDistribution, ProbabilityDistribution, Nothing])

    @test ruleVBLogNormalOut(nothing, ProbabilityDistribution(Univariate, PointMass, m=1.5), ProbabilityDistribution(Univariate, PointMass, m=3.0)) == Message(Univariate, LogNormal, m=1.5, s=3.0)
end

@testset "averageEnergy and differentialEntropy" begin
    @test differentialEntropy(ProbabilityDistribution(Univariate, LogNormal, m=1.0, s=2.0)) == averageEnergy(LogNormal, ProbabilityDistribution(Univariate, LogNormal, m=1.0, s=2.0), ProbabilityDistribution(Univariate, PointMass, m=1.0), ProbabilityDistribution(Univariate, PointMass, m=2.0))
end

end #module


include("sec_6_2.jl")

using PyPlot

r = 3
M = 2r
ppI = 40
store = Store(α, r+1, M, ppI)

N = 5
t, U = FODEdG!(λ, tmax, f, u0, N, r, M, store)
τ = FractionalTimeDG.reconstruction_pts(T, r)
pcwise_t, pcwise_U = evaluate_pcwise_poly!(U, t, τ[0:r], store)
pcwise_u = [ u(t, λ, u0, f, 10, store) for t in pcwise_t ]
pcwise_err = pcwise_U - pcwise_u
Uhat = reconstruction(U, u0, store)
pcwise_t_fine, pcwise_Uhat = evaluate_pcwise_poly!(Uhat, t, ppI, store)
pcwise_u_fine = [ u(t, λ, u0, f, 10, store) for t in pcwise_t_fine ]
pcwise_rec_err = pcwise_Uhat - pcwise_u_fine
pcwise_rec_err[1,1] = NaN
a = FractionalTimeDG.Fourier_Legendre_coefs(r, t, store) do s
    return u(s, λ, u0, f, 10, store)
end

figure(4)
trans_pcwise_t = copy(pcwise_t')
trans_pcwise_err = copy(pcwise_err')
line0 = semilogy(pcwise_t_fine, abs.(pcwise_rec_err), "k:")
line1 = semilogy(trans_pcwise_t[:,1], abs.(trans_pcwise_err[:,1]), "o",
                 markersize=3)
line2 = semilogy(trans_pcwise_t[:,2], abs.(trans_pcwise_err[:,2]), "v",
                 markersize=3)
line3 = semilogy(trans_pcwise_t[:,3], abs.(trans_pcwise_err[:,3]), "+",
                 markersize=4)
line4 = semilogy(trans_pcwise_t[:,4], abs.(trans_pcwise_err[:,4]), "x",
                 markersize=4)
legend((line0[1], line1[1], line2[1], line3[1], line4[1]),
       (L"$\widehat E(t)$", L"$j=0$", L"j=1", L"j=2", L"j=3"))
xlabel(L"$t$", size=12)
ylabel("Error")
grid(true)
savefig("fig4.eps")


# Repeated-measures Designs
#
# author: yizhan miao
# email: yzmiao@protonmail.com
# last update: Oct 26 2018

"""One-way Repeated Measures ANOVA

Syntax:
- `RESULT, _F, _p = rm_anova(Y::Array{T, 1}, S::Array{D, 1}, F::Array{D, 1};name::Symbol=:F, preview::Bool=true) where {T <: Number, D <: Integer}`

Arguments:
- `Y`: data set in 1d vector
- `S`: corresponding subject marker, as Array{Integer, 1}
- `F`: corresponding factor marker, as Array{Integer, 1}

Keywords:
- `name`: the name of the factor [default: `:F`]
- `preview`: print anova table [default: `true`]

Returns:
- ANOVA Result
- F statistic
- p value

Example:
```julia
julia> rm_anova_Y = [
    -4.4297,  4.7513,  3.2971, -1.4606,
     4.8458,  5.1163,  6.0739, -1.9225,
     6.1542, 10.4794, 12.4438,  9.2150];
julia> rm_anova_S = [1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4];
julia> rm_anova_F = [1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3];
julia> ANOVA.rm_anova(rm_anova_Y, rm_anova_S, rm_anova_F, name=:factor);
```
```
One-way repeated measures ANOVA
-----------------------------------------------
source  SS      df    MS      F       Prob
-----------------------------------------------
factor  169.437 2     84.718  12.364  0.00744
error   41.112  6     6.852   NaN     NaN
S       74.253  3     24.751  NaN     NaN
total   284.802 11    NaN     NaN     NaN
-----------------------------------------------
```
"""
function rm_anova(Y::Array{T, 1}, S::Array{D, 1}, F::Array{D, 1};name::Symbol=:F, preview::Bool=true) where {T <: Number, D <: Integer}
    # check balance
    _factor = unique(F)
    _n_factor = zeros(Int64, length(_factor))
    for (idx, each) in enumerate(_factor)
        _n_factor[idx] = F[F .== each] |> length
    end

    _subject = unique(S)
    _n_subject = zeros(Int64, length(_subject))
    for (idx, each) in enumerate(_subject)
        _n_subject[idx] = S[S .== each] |> length
    end

    @assert(length(unique(_n_factor)) == 1, "input data is not balanced across factors.")
    @assert(length(unique(_n_subject)) == 1, "input data is not balanced across subject: $(unique(_n_subject)).")
    # TODO: support unbalanced ANOVA

    _mean_total = mean(Y)
    _mean_treat = [mean(Y[F .== _F]) for _F in _factor]
    _mean_subject = [mean(Y[S .== _S]) for _S in _subject]

    _SS_total = (Y .- _mean_total).^2 |> sum
    _SS_treat = sum((_mean_treat .- _mean_total).^2) * length(_subject)
    _SS_subject = sum((_mean_subject .- _mean_total) .^ 2) * length(_factor)
    _SS_SF = _SS_total - _SS_treat - _SS_subject

    _df_total = length(Y) - 1
    _df_treat = length(_factor) - 1
    _df_subject = length(_subject) - 1
    _df_SF = _df_total - _df_treat - _df_subject

    ##### rm anova #####
    _MS_treat = _SS_treat / _df_treat
    _MS_subject = _SS_subject / _df_subject
    _MS_SF = _SS_SF / _df_SF

    _F = _MS_treat / _MS_SF
    _F_dist = FDist(_df_treat, _df_SF)

    _p = ccdf(_F_dist, _F)

    TOTAL = anova_entry(:total, _SS_total, _df_total, NaN, NaN, NaN)
    TREAT = anova_entry(name, _SS_treat, _df_treat, _MS_treat, _F, _p)
    SF    = anova_entry(:error, _SS_SF, _df_SF, _MS_SF, NaN, NaN)
    SUBJECT = anova_entry(:S, _SS_subject, _df_subject, _MS_subject, NaN, NaN)

    RESULT = anova_result("One-way repeated measures ANOVA", [TREAT, SF, SUBJECT, TOTAL])

    preview&&preview_anova_result(RESULT)

    RESULT, _F, _p

end

"""Two-way Repeated Measures ANOVA

Syntax:
- `Result = rm_anova2(Y::Array{T, 1}, S::Array{D, 1}, A::Array{D, 1}, B::Array{D, 1};names::Tuple{Symbol, Symbol}=(:A, :B), preview::Bool=true) where {T <: Number, D <: Integer}`

Arguments:
- `Y`: data set in 1d vector
- `S`: corresponding subject marker, as Array{Integer, 1}
- `A`: corresponding factor A marker, as Array{Integer, 1}
- `B`: corresponding factor B marker, as Array{Integer, 1}

Keywords:
- `names`: the name of the factor [default: `(:A, :B)`]
- `preview`: print anova table [default: `true`]

Return:
- ANOVA Result

"""
function rm_anova2(Y::Array{T, 1}, S::Array{D, 1}, A::Array{D, 1}, B::Array{D, 1};names::Tuple{Symbol, Symbol}=(:A, :B), preview::Bool=true) where {T <: Number, D <: Integer}
    # check balance
    _factor_A = unique(A)
    _n_factor_A = zeros(Int64, length(_factor_A))
    for (idx, each) in enumerate(_factor_A)
        _n_factor_A[idx] = A[A .== each] |> length
    end

    _factor_B = unique(B)
    _n_factor_B = zeros(Int64, length(_factor_B))
    for (idx, each) in enumerate(_factor_B)
        _n_factor_B[idx] = B[B .== each] |> length
    end

    _subject = unique(S)
    _n_subject = zeros(Int64, length(_subject))
    for (idx, each) in enumerate(_subject)
        _n_subject[idx] = S[S .== each] |> length
    end

    @assert(length(unique(_n_factor_A)) == 1, "input data is not balanced across factor A.")
    @assert(length(unique(_n_factor_B)) == 1, "input data is not balanced across factor B.")
    @assert(length(unique(_n_subject)) == 1,  "input data is not balanced across subject.")
    #TODO: support unbalanced ANOVA

    _mean_total = mean(Y)
    _mean_A = [mean(Y[A .== _A]) for _A in _factor_A]
    _mean_B = [mean(Y[B .== _B]) for _B in _factor_B]
    _mean_BA = [mean(Y[(A .== _A) .& (B .== _B)]) for _A in _factor_A, _B in _factor_B]
    _mean_S = [mean(Y[S .== _S]) for _S in _subject]

    _mean_AS = [mean(Y[(A .== _A) .& (S .== _S)]) for _A in _factor_A, _S in _subject]
    _mean_BS = [mean(Y[(B .== _B) .& (S .== _S)]) for _B in _factor_B, _S in _subject]

    _SS_total = (Y .- _mean_total).^2 |> sum
    _SS_A = sum((_mean_A .- _mean_total) .^ 2) * length(_subject) * length(_factor_B)
    _SS_B = sum((_mean_B .- _mean_total) .^ 2) * length(_subject) * length(_factor_A)
    _SS_BA = sum((_mean_BA .- _mean_total) .^ 2) * length(_subject) - _SS_A - _SS_B
    _SS_S = sum((_mean_S .- _mean_total).^2) * length(_factor_A) * length(_factor_B)
    _SS_AS = sum((_mean_AS .- _mean_total) .^ 2) * length(_factor_B) - _SS_A - _SS_S
    _SS_BS = sum((_mean_BS .- _mean_total) .^ 2) * length(_factor_A) - _SS_B - _SS_S
    _SS_BAS = _SS_total - _SS_A - _SS_B - _SS_BA - _SS_S - _SS_AS - _SS_BS

    _df_total = length(Y) - 1
    _df_A = length(_factor_A) - 1
    _df_B = length(_factor_B) - 1
    _df_S = length(_subject) - 1
    _df_BA = _df_A * _df_B
    _df_AS = _df_A * _df_S
    _df_BS = _df_B * _df_S
    _df_BAS = _df_BA * _df_S

    _MS_A = _SS_A / _df_A
    _MS_B = _SS_B / _df_B
    _MS_BA = _SS_BA / _df_BA
    _MS_S = _SS_S / _df_S
    _MS_AS = _SS_AS / _df_AS
    _MS_BS = _SS_BS / _df_BS
    _MS_BAS = _SS_BAS / _df_BAS

    ##### ANOVA #####
    _F_A = _MS_A / _MS_AS
    _F_B = _MS_B / _MS_BS
    _F_BA = _MS_BA / _MS_BAS

    _p_A = ccdf(FDist(_df_A, _df_AS), _F_A)
    _p_B = ccdf(FDist(_df_B, _df_BS), _F_B)
    _p_BA = ccdf(FDist(_df_BA, _df_BAS), _F_BA)

    TOTAL = anova_entry(:total, _SS_total, _df_total, NaN, NaN, NaN)
    _S    = anova_entry(:S, _SS_S, _df_S, _MS_S, NaN, NaN)
    _A    = anova_entry(names[1], _SS_A, _df_A, _MS_A, _F_A, _p_A)
    _B    = anova_entry(names[2], _SS_B, _df_B, _MS_B, _F_B, _p_B)
    _BA   = anova_entry(Symbol("$(names[2])*$(names[1])"), _SS_BA, _df_BA, _MS_BA, _F_BA, _p_BA)
    _AS   = anova_entry(Symbol("S*$(names[1])"), _SS_AS, _df_AS, _MS_AS, NaN, NaN)
    _BS   = anova_entry(Symbol("S*$(names[2])"), _SS_BS, _df_BS, _MS_BS, NaN, NaN)
    _BAS   = anova_entry(Symbol("S*$(names[2])*$(names[1])"), _SS_BAS, _df_BAS, _MS_BAS, NaN, NaN)

    RESULT = anova_result("Two-way repeated measures ANOVA", [_S, _A, _B, _BA, _AS, _BS, _BAS, TOTAL])

    preview&&preview_anova_result(RESULT)

    RESULT
end


using PkgSRA
using Plots, Random
using Flux, DiffEqFlux, OrdinaryDiffEq, Statistics
using DSP
pyplot()
Random.seed!(11)

#####
##### Generate true dynamics matrix and data
#####
# Note: continuous version

# Generate the control signal starts
# include("example_forcing_functions.jl")
include("example_falling_ball.jl")
u0 = u1;
# dat_no_wall = solve_ball_system()
# plot(dat_no_wall[1,:], label="Ball position")

# dat = solve_deformed_ball_system()

U_func(X) = U_func_spring(X);
dat = solve_ball_system(U_func_time_trivial,
                        U_func)
# dat = Array(sol)

# Add a row of constants to the data (not actually needed here)
n, m = size(dat)
dat_and_const = vcat(dat, ones(1,m))

#####
##### Mini-SRA: initial guess for dynamics vs. control
#####
dat_grad = numerical_derivative(dat, ts)
# Plot this numerical derivative vs. the truth
plot(dat_grad[2,:], label="Numerical derivative")
# plot!(core_dyn_true(dat)[2,:], label="True derivative")

U_guess = sra_dmd(dat_and_const, dat_grad,
                    quantile_threshold=0.8)[1]
A, B = dmdc(dat_and_const, dat_grad, U_guess)
# Loop a couple of times
for i in 1:20
    pt = plot(U_guess[2,:], label="Residual")
    # plot!(U_func.(ts), label="True")
    title!("Iteration $i")
    display(pt)
    # Update guess for control signal
    global U_guess = calculate_residual(dat_and_const,
                                        dat_grad,
                                        A)
    sparsify_signal!(U_guess, quantile_threshold=0.8)
    # Update guess for dynamics
    global A, B = dmdc(dat_and_const, dat_grad, U_guess)
end

# Filter the signal to be easier for the NN to learn
# myfilter = digitalfilter(Lowpass(0.5),Butterworth(2));
# U_filter = zeros(size(U_guess))
# for i in 1:size(U_guess, 1)
#     U_filter[i,:] = filtfilt(myfilter, U_guess[i,:]);
# end
# U_filter *= maximum(abs.(U_guess))/maximum(abs.(U_filter))
U_true = zeros(size(dat))
for i in 1:size(dat,2)
    U_true[:,i] = U_func(dat[:,i])
end
# U_true = [U_func(d) for d in dat]
plot(U_true[2,:], label="Control Signal")
plot!(U_guess[2,:], label="Raw Guessed Control Signal")
# plot!(U_filter[2,:], label="Filtered Guessed Control Signal")
U_filter = U_guess

#####
##### Learn a NN that can reconstruct the SRA control signal
#####
# Now, a function of phase space
nn_dim = 24
U_dim = 1 # Expect that the true dimensionality is 1
# Note: need to initialize the first layer so that the whole space is used
ctr_dyn = Chain(Dense(n, nn_dim),
                Dense(nn_dim, nn_dim, σ),
                Dense(nn_dim, U_dim),
                Dense(U_dim,2)) # Equivalent to the B matrix
ps = Flux.params(ctr_dyn)
loss_U_filter() = sum(abs2,ctr_dyn(dat) .- U_filter)

loss_to_use = loss_U_filter

cb_ctr = function (i=2)
    pt = plot(Flux.data(ctr_dyn(dat)[i,:]), label="Prediction")
    plot!(U_filter[i,:], label="Guess for control signals")
    display(pt)
    return pt
end
cb_ctr(2)

# Initial fast learning
tol = 200
dat_dummy = Iterators.repeated((), 100)
opt = ADAM(0.1)
while loss_to_use() > 5*tol
    Flux.train!(loss_to_use, ps, dat_dummy, opt, cb = ()->@show loss_to_use())
    cb_ctr(2)
end
opt = ADAM(0.01)
while loss_to_use() > tol
    Flux.train!(loss_to_use, ps, dat_dummy, opt, cb = ()->@show loss_to_use())
    cb_ctr(2)
end
println("Finished learning control signal")

#####
##### As a sanity check, do the NN syntax with the true dynamics
#####
# B_true = Float32.([0.; 1.])
# function dudt_true(u::AbstractArray,p,t)
#     core_dyn_true(u) + B_true*Float32(U_func(t))
# end
# p = param(Float32[0.0]) # Unused
# prob_true = ODEProblem(dudt_true,u0,tspan,p)
# sol = diffeq_rd(p,prob_true,Tsit5(), saveat=ts)
# plot(Flux.data.(Array(sol))[2,:], label="Ideal Simulation")
# plot!(dat[2,:], label="Data")
# title!("True dynamics")
#
# best_loss = sum(abs2,Flux.data.(Array(sol)) .- dat)
# println("The best loss achievable is $best_loss")

#####
##### DEFINE the unified mixed NODE
#####
# Part 1: Linear dynamics; don't predict time
# Note: use actual constructor, not the convenience one
core_dyn = Dense(param(Float32.(A[:, 1:n])),
                 param(Float32.(A[:, end])),
                 identity)
# Part 2: Neural network for control signal; initialized above
p = param(Float32[1.0; 1.0])

function dudt_(u::AbstractArray,p,t)
    core_dyn(u) + p.*ctr_dyn(u)
end

prob = ODEProblem(dudt_,u0,tspan,p)
# Check the initial solution
# sol = diffeq_rd(p,prob,Tsit5(), saveat=ts)
# plot(Flux.data.(Array(sol))[2,:], label="Simulation")
# plot!(dat[2,:], label="Data")

_x = param(u0) # The initial condition can be learned too

function predict_rd()
    Flux.Tracker.collect(diffeq_rd(p,prob,Tsit5(),u0=_x, saveat=ts))
end
function loss_rd(λ1=0.0001, λ2=0.001)
    return sum(abs2,predict_rd() .- dat) +      # Good fit
            λ1*loss_U_filter() +                # Control signal close to guess
            λ2*sum(abs,ctr_dyn(dat))   # L1 norm on ctr
end

cb_plot = function (i=1)
  # display(loss_rd())
  pt = plot(ts, Flux.data(predict_rd()[i,:]), label="Prediction")
  plot!(ts, dat[i,:], label="Data")
  display(pt)
  return pt
end

#####
##### SOLVE the mixed NODE
#####
ps = params(core_dyn,ctr_dyn,p);

# Display the ODE with the current parameter values.
cb_plot(1)
# Initial quick learning
opt = ADAM(1e-6)
dat_dummy = Iterators.repeated((), 10)
Flux.train!(loss_rd, ps,
            dat_dummy, opt, cb = ()->@show loss_rd())
cb_plot(1)
# Run for a long time
opt = ADAM(1e-9)
tol = 50
dat_dummy = Iterators.repeated((), 5)
while loss_rd() > tol
    Flux.train!(loss_rd, ps,
                dat_dummy, opt, cb = ()->@show loss_rd())
    cb_plot(1)
    break
end

#####
##### Instead: SOLVE the mixed NODE via alternating
#####

# Pure loss
loss_dat() = sum(abs2,predict_rd() .- dat)
loss_ctr(λ=0.1) = loss_dat() + λ*sum(abs,ctr_dyn(dat))

function plot_dat_and_ctr()
    pt1 = cb_plot();
    pt2 = cb_ctr();
    pt = plot(pt1, pt2, layout=@layout [a;b]);
    display(pt)
end

ps_core = params(core_dyn);
ps_ctr = params(ctr_dyn,p);

plot_dat_and_ctr()
dat_dummy = Iterators.repeated((), 10)
tol = 150;
opt = ADAM(1e-10)
while loss_dat() > tol
    println("Learning core dynamics; direct L2 loss")
    Flux.train!(loss_dat, ps_core,
                dat_dummy, opt, cb = ()->@show loss_dat())
    plot_dat_and_ctr()

    println("Learning Control signal; direct L2 loss with L1 penalty")
    Flux.train!(loss_ctr, ps_ctr,
                dat_dummy, opt, cb = ()->@show loss_dat())
    plot_dat_and_ctr()

    # break
end


using Base.Test

include("dpmm.jl")
include("smc_samplers.jl")

@testset "CRP tests" begin

    @testset "assignment book-keeping" begin
        crp = CRPState()

        incorporate!(crp, 1, next_new_cluster(crp))
        c1 = assignment(crp, 1)
        @test c1 == 1

        incorporate!(crp, 2, next_new_cluster(crp))
        c2 = assignment(crp, 2)
        @test c2 == 2

        incorporate!(crp, 3, c2)
        @test counts(crp, c1) == 1
        @test counts(crp, c2) == 2
        @test next_new_cluster(crp) == c2 + 1
        @test crp.next_cluster == c2 + 2

        unincorporate!(crp, 1)
        @test !has_cluster(crp, c1)
        @test counts(crp, c2) == 2
        @test next_new_cluster(crp) == c1

        unincorporate!(crp, 2)
        @test counts(crp, c2) == 1

        unincorporate!(crp, 3)
        @test !has_cluster(crp, c2)
        @test next_new_cluster(crp) == c2

        # this is never decremented
        @test crp.next_cluster == c2 + 2
    end

    @testset "log_probability" begin

        alpha = 0.25
        crp = CRPState()
        @test log_probability(crp, alpha) == 0.0

        # [1]
        incorporate!(crp, 1, next_new_cluster(crp))
        @test log_probability(crp, alpha) == 0.0

        # join the same table as 1: [1,2]
        incorporate!(crp, 2, assignment(crp, 1))
        actual = log_probability(crp, alpha)
        expected = log(1. / (1. + alpha))
        @test isapprox(actual, expected)

        # join a new table: [1] [2]
        unincorporate!(crp, 2)
        incorporate!(crp, 2, next_new_cluster(crp))
        actual = log_probability(crp, alpha)
        expected = log(alpha / (1. + alpha))
        @test isapprox(actual, expected)

        # join same table as 1: [1, 3] [2]
        incorporate!(crp, 3, assignment(crp, 1))
        actual = log_probability(crp, alpha)
        expected = log((alpha / (1.+alpha)) *
                       (1./(2.+alpha)))
        @test isapprox(actual, expected)

        # join same table as 1 [1, 3, 4] [2]
        incorporate!(crp, 4, assignment(crp, 1))
        actual = log_probability(crp, alpha)
        expected = log((alpha / (1.+alpha)) *
                       (1./(2.+alpha)) *
                       (2./(3.+alpha)))
        @test isapprox(actual, expected)
    end

end

@testset "NIGN" begin

    # arbitrary values
    m = 1.23
    r = 2.41
    nu = 5.44
    s = 4.2
    alpha = nu/2.
    beta = s/2.
    params = NIGNParams(m, r, nu, s)

    nign = NIGN()
    actual = log_probability_density(nign, params)
    @test abs(actual) < 1e-15

    # https://www.cs.ubc.ca/~murphyk/Papers/bayesGauss.pdf
    x1 = -4.1
    x2 = 2.22
    incorporate!(nign, x1)
    incorporate!(nign, x2)
    (mu_actual, r_actual, s_actual, nu_actual) = posterior_params(nign, params)
    # m, r, s, nu
    avg = mean([x1, x2])
    sum_sq = sum(([x1, x2] - avg) .* ([x1,x2] - avg))
    n = 2
    mu_expect = (r * m + n * avg) / (r + n)
    r_expect = r + n
    alpha_expect = alpha + n/2.
    beta_expect = beta + 0.5 * sum_sq  + r*n*((avg-m)*(avg-m))/(2*(r+n))
    nu_expect = alpha_expect*2.
    s_expect = beta_expect*2.
    @test isapprox(mu_actual, mu_expect)
    @test isapprox(r_actual, r_expect)
    @test isapprox(s_actual, s_expect)
    @test isapprox(nu_actual, nu_expect)

    # should go back to empty again
    unincorporate!(nign, x1)
    unincorporate!(nign, x2)
    actual = log_probability_density(nign, params)
    @test abs(actual) < 1e-15

end

@testset "DPMM SMC" begin

    @testset "extreme easy two-cluster dataset" begin
        srand(1)

        # generate dataset
        num_data = 50
        data = Array{Float64,1}(num_data)
        cluster_1_members = []
        cluster_2_members = []
        for i=1:num_data
            if rand() < 0.5
                mu = -1.0
                push!(cluster_1_members, i)
            else
                mu = 1.0
                push!(cluster_2_members, i)
            end
            std = 0.01
            data[i] = rand(Normal(mu, std))
        end

        # set custom values of hypers to match extreme dataset
        m_prior = Normal(0, 1) # mean cluster center
        r_prior = Gamma(1, 1) # r * data-precision is the precision of cluster centers
        nu_prior = Gamma(1, 1)# expected data-precision of cluster is nu/s
        s_prior = Gamma(1, 1)
        nign_params_prior = NIGNParamsPrior(m_prior, r_prior, nu_prior, s_prior)
        alpha_prior = Gamma(1, 1)
        hypers = Hyperparameters(nign_params_prior, alpha_prior)

        gibbs_sweeps_per_datum = 1
        num_particles = 1
        scheme = make_optimal_proposal_smc_scheme(data, hypers,
                               gibbs_sweeps_per_datum, num_particles)
        state, lml_estimate = smc(scheme)

        # check that there are two clusters, and that clustering
        # is correct
        @test length(state.components) == 2
        cluster_1 = assignment(state.crp, cluster_1_members[1])
        cluster_2 = assignment(state.crp, cluster_2_members[2])
        @test cluster_1 != cluster_2
        for i in cluster_1_members[2:end]
            @test assignment(state.crp, i) == cluster_1
        end
        for i in cluster_2_members[2:end]
            @test assignment(state.crp, i) == cluster_2
        end

    end

    @testset "prior proposal weight" begin
        alpha = 0.01
        m = 0.9
        r = 1.1
        nu = 1.2
        s = 1.3
        nign_params = NIGNParams(m, r, nu, s)
        alpha_prior = Gamma(1, 1)
        nign_params_prior = NIGNParamsPrior(Normal(0, 1), Gamma(1, 1), Gamma(1, 1), Gamma(1, 1))
        hypers = Hyperparameters(nign_params_prior, alpha_prior)
	    state = DPMMState(alpha, CRPState(), OrderedDict{Int,NIGN}(),
                      nign_params, hypers)

        srand(1)

        x1 = 4.2
        ll1 = add_datum_prior!(state, 1, x1)
        ll1_rev = remove_datum_prior!(state, 1, x1)
        @test isapprox(ll1, ll1_rev)

        x2 = -3.1
        ll1_new = add_datum_prior!(state, 1, x1)
        @test isapprox(ll1, ll1_new)
        ll2 = add_datum_prior!(state, 2, x2)
        # NOTE: we seeded so this should remain true, if it doesn't change the
        # seed or make alpha even smaller
        @test assignment(state.crp, 1) == assignment(state.crp, 2)
        component = NIGN()
        incorporate!(component, x1)
        incorporate!(component, x2)
        total_logp = log_probability_density(component, nign_params)
        @test isapprox(total_logp, ll1_new + ll2)

        ll1_rev = remove_datum_prior!(state, 1, x1)
        ll2_rev = remove_datum_prior!(state, 2, x2)
        @test isapprox(total_logp, ll1_rev + ll2_rev)
    end

    @testset "optimal proposal weight" begin
        alpha = 2.45
        m = 0.9
        r = 1.1
        nu = 1.2
        s = 1.3
        nign_params = NIGNParams(m, r, nu, s)
        alpha_prior = Gamma(1, 1)
        nign_params_prior = NIGNParamsPrior(Normal(0, 1), Gamma(1, 1), Gamma(1, 1), Gamma(1, 1))
        hypers = Hyperparameters(nign_params_prior, alpha_prior)
	    state = DPMMState(alpha, CRPState(), OrderedDict{Int,NIGN}(),
                      nign_params, hypers)

        srand(1)

        x1 = 4.2
        ll1 = add_datum_conditional!(state, 1, x1)
        ll1_rev = remove_datum_conditional!(state, 1, x1)
        @test isapprox(ll1, ll1_rev)
        component = NIGN()
        incorporate!(component, x1)
        @test isapprox(ll1, log_probability_density(component, nign_params))

        ll1_new = add_datum_conditional!(state, 1, x1)
        x2 = 1.5
        @test isapprox(ll1, ll1_new)
        ll2 = add_datum_conditional!(state, 2, x2)
        ll2_rev = remove_datum_conditional!(state, 2, x2)
        @test isapprox(ll2, ll2_rev)

        # explicitly sum over the two possible assignments
        # p(x_2 = alone | x_1) p(y_2 | x_1, x_2 = alone, y_1)
        component_alone = NIGN()
        incorporate!(component_alone, x2)
        ll2_alone = log_probability_density(component_alone, nign_params)
        ll2_alone += log(alpha / (1. + alpha))

        # p(x_2 = together | x_1) p(y_2 | x_1, x_2 = together, y_1)
        component_together = NIGN()
        incorporate!(component_together, x1)
        component_old_logpdf = log_probability_density(component_together, nign_params)
        incorporate!(component_together, x2)
        ll2_together = log_probability_density(component_together, nign_params) - component_old_logpdf
        ll2_together += log(1. / (1. + alpha))

        # p(y_2 | x_1, y_1) = sum of two terms above
        ll2_expected = logsumexp([ll2_alone, ll2_together])
        @test isapprox(ll2, ll2_expected)
    end

end


module Test_cplyap

using LinearAlgebra
using MatrixEquations
using Test

@testset "Testing positive continuous Lyapunov equation solvers" begin

n = 30
m = 5
p = 3
Ty = Float64

@testset "Positive continuous Lyapunov equations" begin

reltol = eps(float(100))
a = -1; b = 2im; @time u = plyapc(a,b)
@test abs(a*u*u'+u*u'*a'+b*b') < reltol

reltol = eps(float(100f0))
a = -1f0; b = 2f0im; @time u = plyapc(a,b)
@test abs(a*u*u'+u*u'*a'+b*b') < reltol

for Ty in (Float64, Float32)

ar = rand(Ty,n,n)
ar = ar-2*norm(ar)*Matrix(I,n,n)
ac = rand(Ty,n,n)+im*rand(Ty,n,n)
ac = ac-2*norm(ac)*Matrix(I,n,n)
br = rand(Ty,n,m)
brw = rand(Ty,n,n+m)
bc = br+im*rand(Ty,n,m)
bcw = brw+im*rand(Ty,n,n+m)
cr = rand(Ty,p,n)
crt = rand(Ty,n+p,n)
cc = cr+im*rand(Ty,p,n)
cct = crt+im*rand(Ty,n+p,n)
Ty == Float64 ? reltol = eps(float(100)) : reltol = eps(100*n*one(Ty))

@time u = plyapc(ar,br);
x = u*u'; @test norm(ar*x+x*ar'+br*br')/norm(x)/norm(ar) < reltol

@time u = plyapc(ar,brw);
x = u*u'; @test norm(ar*x+x*ar'+brw*brw')/norm(x)/norm(ar) < reltol

@time u = plyapc(ar,I,br);
x = u*u'; @test norm(ar*x+x*ar'+br*br')/norm(x)/norm(ar) < reltol

@time u = plyapc(ar',cr')
x = u'*u; @test norm(ar'*x+x*ar+cr'*cr)/norm(x)/norm(ar) < reltol

@time u = plyapc(ar',crt')
x = u'*u; @test norm(ar'*x+x*ar+crt'*crt)/norm(x)/norm(ar) < reltol

@time u = plyapc(ac,bc);
x = u*u'; @test norm(ac*x+x*ac'+bc*bc')/norm(x)/norm(ac) < reltol

@time u = plyapc(ac,bcw);
x = u*u'; @test norm(ac*x+x*ac'+bcw*bcw')/norm(x)/norm(ac) < reltol

@time u = plyapc(ac',cc');
x = u'*u; @test norm(ac'*x+x*ac+cc'*cc)/norm(x)/norm(ac) < reltol

@time u = plyapc(ac',cct');
x = u'*u; @test norm(ac'*x+x*ac+cct'*cct)/norm(x)/norm(ac) < reltol

@time u = plyapc(ac',I,cc');
x = u'*u; @test norm(ac'*x+x*ac+cc'*cc)/norm(x)/norm(ac) < reltol

@time u = plyapc(ac,br);
x = u*u'; @test norm(ac*x+x*ac'+br*br')/norm(x)/norm(ac) < reltol

@time u = plyapc(ac',cr');
x = u'*u; @test norm(ac'*x+x*ac+cr'*cr)/norm(x)/norm(ac) < reltol

@time u = plyapc(ar,bc);
x = u*u'; @test norm(ar*x+x*ar'+bc*bc')/norm(x)/norm(ar) < reltol

@time u = plyapc(ar',cc');
x = u'*u; @test norm(ar'*x+x*ar+cc'*cc)/norm(x)/norm(ar)  < reltol

end
end

@testset "Positive continuous generalized Lyapunov equations" begin

reltol = eps(float(100))
a = -1; ee = 4; b = 2im; @time u = plyapc(a,ee,b)
@test abs(a*u*u'*ee'+ee*u*u'*a'+b*b') < reltol

reltol = eps(float(100f0))
a = -1f0; ee = 4f0; b = 2f0im; @time u = plyapc(a,ee,b)
@test abs(a*u*u'*ee'+ee*u*u'*a'+b*b')  < reltol

for Ty in (Float64, Float32)

ar = rand(Ty,n,n)
ar = ar-2*norm(ar)*Matrix(I,n,n)
er = rand(n,n)
ar = er*ar
ac = rand(Ty,n,n)+im*rand(Ty,n,n)
ac = ac-2*norm(ac)*Matrix(I,n,n)
ec = er+im*rand(Ty,n,n)
ac = ec*ac
br = rand(Ty,n,m)
brw = rand(Ty,n,n+m)
bc = br+im*rand(Ty,n,m)
bcw = brw+im*rand(Ty,n,n+m)
cr = rand(Ty,p,n)
crt = rand(Ty,n+p,n)
cc = cr+im*rand(Ty,p,n)
cct = crt+im*rand(Ty,n+p,n)
Ty == Float64 ? reltol = eps(float(100)) : reltol = eps(100*n*one(Ty))

@time u = plyapc(ar,er,br);
x = u*u'; @test norm(ar*x*er'+er*x*ar'+br*br')/norm(x)/norm(ar) < reltol

@time u = plyapc(ar,er,brw);
x = u*u'; @test norm(ar*x*er'+er*x*ar'+brw*brw')/norm(x)/norm(ar) < reltol

@time u = plyapc(ar',er',cr');
x = u'*u; @test norm(ar'*x*er+er'*x*ar+cr'*cr)/norm(x)/norm(ar) < reltol

@time u = plyapc(ar',er',crt');
x = u'*u; @test norm(ar'*x*er+er'*x*ar+crt'*crt)/norm(x)/norm(ar) < reltol

@time u = plyapc(ac,ec,bc);
x = u*u'; @test norm(ac*x*ec'+ec*x*ac'+bc*bc')/norm(x)/norm(ac)/norm(ec) < reltol

@time u = plyapc(ac,ec,bcw);
x = u*u'; @test norm(ac*x*ec'+ec*x*ac'+bcw*bcw')/norm(x)/norm(ac)/norm(ec) < reltol

@time u = plyapc(ac',ec',cc');
x = u'*u; @test norm(ac'*x*ec+ec'*x*ac+cc'*cc)/norm(x)/norm(ac) < reltol

@time u = plyapc(ac',ec',cct');
x = u'*u; @test norm(ac'*x*ec+ec'*x*ac+cct'*cct)/norm(x)/norm(ac) < reltol

@time u = plyapc(ar,er,bc);
x = u*u'; @test norm(ar*x*er'+er*x*ar'+bc*bc')/norm(x)/norm(ar) < reltol

@time u = plyapc(ar',er',cc');
x = u'*u; @test norm(ar'*x*er+er'*x*ar+cc'*cc)/norm(x)/norm(ar) < reltol

@time u = plyapc(ac,ec,br);
x = u*u'; @test norm(ac*x*ec'+ec*x*ac'+br*br')/norm(x)/norm(ac)/norm(ec) < reltol

@time u = plyapc(ac',ec',cr');
x = u'*u; @test norm(ac'*x*ec+ec'*x*ac+cr'*cr)/norm(x)/norm(ac) < reltol

end
end

@testset "Positive 2x2 continuous Lyapunov equations" begin

A = [-1.1 1.; -1. -1.]
E = [1. 1.; 0. 1.]
R = [1. 1.; 0. 1.]
reltol = eps(float(100))

U, scale, β, α = plyap2(A, R, adj = true, disc = false)
X = U'*U; @test norm(A'*X+X*A+scale^2*R'*R)/max(1,norm(X))/norm(A) < reltol &&
                norm(β*U-U*A)/max(1,norm(U))/norm(A) < reltol &&
                norm(α*U/scale - scale*R)/max(1,norm(scale*R)) < reltol

U, scale, β, α = plyap2(A, R, adj = false, disc = false)
X = U*U'; @test norm(A*X+X*A'+scale^2*R*R')/max(1,norm(X))/norm(A) < reltol &&
                norm(U*β-A*U)/max(1,norm(U))/norm(A) < reltol &&
                norm(U*α/scale - scale*R)/max(1,norm(scale*R)) < reltol

U, scale, β, α = pglyap2(A, E, R, adj = true, disc = false)
X = U'*U; @test norm(A'*X*E+E'*X*A+scale^2*R'*R)/max(1,norm(X))/norm(A) < reltol &&
                norm(β*U*E-U*A)/max(1,norm(U))/norm(A) < reltol &&
                norm(α*U*E/scale - scale*R)/max(1,norm(scale*R)) < reltol

U, scale, β, α = pglyap2(A, E, R, adj = false, disc = false)
X = U*U'; @test norm(A*X*E'+E*X*A'+scale^2*R*R')/max(1,norm(X))/norm(A) < reltol &&
                norm(A*U-E*U*β)/max(1,norm(U))/norm(A) < reltol &&
                norm(E*U*α/scale - scale*R)/max(1,norm(scale*R)) < reltol

A = [-1.1 1.; -1. -1.]/10; A = convert(Matrix{Float32},A)
E = [1. 1.; 0. 1.]; E = convert(Matrix{Float32},E)
R = [1. 1.; 0. 1.]; R = convert(Matrix{Float32},R)
reltol = eps(100f0)

U, scale, β, α = plyap2(A, R, adj = true, disc = false)
X = U'*U; @test norm(A'*X+X*A+scale^2*R'*R)/max(1,norm(X))/norm(A) < reltol &&
                norm(β*U-U*A)/max(1,norm(U))/norm(A) < reltol &&
                norm(α*U/scale - scale*R)/max(1,norm(scale*R)) < reltol

U, scale, β, α = plyap2(A, R, adj = false, disc = false)
X = U*U'; @test norm(A*X+X*A'+scale^2*R*R')/max(1,norm(X))/norm(A) < reltol &&
                norm(U*β-A*U)/max(1,norm(U))/norm(A) < reltol &&
                norm(U*α/scale - scale*R)/max(1,norm(scale*R)) < reltol

U, scale, β, α = pglyap2(A, E, R, adj = true, disc = false)
X = U'*U; @test norm(A'*X*E+E'*X*A+scale^2*R'*R)/max(1,norm(X))/norm(A) < reltol &&
                norm(β*U*E-U*A)/max(1,norm(U))/norm(A) < reltol &&
                norm(α*U*E/scale - scale*R)/max(1,norm(scale*R)) < reltol

U, scale, β, α = pglyap2(A, E, R, adj = false, disc = false)
X = U*U'; @test norm(A*X*E'+E*X*A'+scale^2*R*R')/max(1,norm(X))/norm(A) < reltol &&
                norm(A*U-E*U*β)/max(1,norm(U))/norm(A) < reltol &&
                norm(E*U*α/scale - scale*R)/max(1,norm(scale*R)) < reltol

end

@testset "Continuous positive Lyapunov equations - Schur form" begin

for Ty in (Float64, Float32)

ar = rand(Ty,n,n)
ar = ar-2*norm(ar)*Matrix(I,n,n)
as,  = schur(ar)
ac = rand(Ty,n,n)+im*rand(Ty,n,n)
ac = ac-2*norm(ac)*Matrix(I,n,n)
acs,  = schur(ac)
br = rand(Ty,n,m)
bc = br+im*rand(Ty,n,m)
cr = rand(Ty,p,n)
cc = cr+im*rand(Ty,p,n)
Ty == Float64 ? reltol = eps(float(100)) : reltol = eps(100*n*one(Ty))

@time u = plyaps(as,br);
x = u*u'; @test norm(as*x+x*as'+br*br')/norm(x)/norm(as) < reltol

@time u = plyaps(as,ar);
x = u*u'; @test norm(as*x+x*as'+ar*ar')/norm(x)/norm(as) < reltol

@time u = plyaps(as,I,br);
x = u*u'; @test norm(as*x+x*as'+br*br')/norm(x)/norm(as) < reltol

@time u = plyaps(as',cr');
x = u'*u; @test norm(as'*x+x*as+cr'*cr)/norm(x)/norm(as) < reltol

@time u = plyaps(as',ar');
x = u'*u; @test norm(as'*x+x*as+ar'*ar)/norm(x)/norm(as) < reltol

@time u = plyaps(acs,bc);
x = u*u'; @test norm(acs*x+x*acs'+bc*bc')/norm(x)/norm(as) < reltol

@time u = plyaps(acs',cc');
x = u'*u; @test norm(acs'*x+x*acs+cc'*cc)/norm(x)/norm(as) < reltol

@time u = plyaps(acs',I,cc');
x = u'*u; @test norm(acs'*x+x*acs+cc'*cc)/norm(x)/norm(as) < reltol

F = UpperTriangular(rand(Ty,n,n))
R = copy(F)
@time plyapcs!(as,R,adj = true);
x = R'*R; @test norm(as'*x+x*as+F'*F)/norm(x)/norm(as) < reltol

F = UpperTriangular(rand(Ty,n,n))
R = copy(F)
@time plyapcs!(as,I,R,adj = true);
x = R'*R; @test norm(as'*x+x*as+F'*F)/norm(x)/norm(as) < reltol

F = UpperTriangular(rand(Ty,n,n))
R = copy(F)
@time plyapcs!(as,R,adj = false);
x = R*R'; @test norm(as*x+x*as'+F*F')/norm(x)/norm(as) < reltol

F = UpperTriangular(rand(Ty,n,n)+im*rand(Ty,n,n))
R = copy(F)
@time plyapcs!(acs,R,adj = true);
x = R'*R; @test norm(acs'*x+x*acs+F'*F)/norm(x)/norm(acs) < reltol

F = UpperTriangular(rand(Ty,n,n)+im*rand(Ty,n,n))
R = copy(F)
@time plyapcs!(acs,R,adj = false);
x = R*R'; @test norm(acs*x+x*acs'+F*F')/norm(x)/norm(acs) < reltol

er = rand(Ty,n,n)
ar = er*ar
as, es = schur(ar,er)
ec = er+im*rand(Ty,n,n)
ac = ec*ac
acs, ecs = schur(ac,ec)

@time u = plyaps(as,es,br);
x = u*u'; @test norm(as*x*es'+es*x*as'+br*br')/norm(x)/norm(as) < reltol

@time u = plyaps(as,es,ar);
x = u*u'; @test norm(as*x*es'+es*x*as'+ar*ar')/norm(x)/norm(as) < reltol

@time u = plyaps(as',es',cr');
x = u'*u; @test norm(as'*x*es+es'*x*as+cr'*cr)/norm(x)/norm(as) < reltol

@time u = plyaps(as',es',ar');
x = u'*u; @test norm(as'*x*es+es'*x*as+ar'*ar)/norm(x)/norm(as) < reltol

@time u = plyaps(acs,ecs,bc);
x = u*u'; @test norm(acs*x*ecs'+ecs*x*acs'+bc*bc')/norm(x)/norm(as) < reltol

@time u = plyaps(acs',ecs',cc');
x = u'*u; @test norm(acs'*x*ecs+ecs'*x*acs+cc'*cc)/norm(x)/norm(as) < reltol

F = UpperTriangular(rand(Ty,n,n))
R = copy(F)
@time plyapcs!(as,es,R,adj = true);
x = R'*R; @test norm(as'*x*es+es'*x*as+F'*F)/norm(x)/norm(as) < reltol

F = UpperTriangular(rand(Ty,n,n))
R = copy(F)
@time plyapcs!(as,es,R,adj = false);
x = R*R'; @test norm(as*x*es'+es*x*as'+F*F')/norm(x)/norm(as) < reltol

F = UpperTriangular(rand(Ty,n,n)+im*rand(Ty,n,n))
R = copy(F)
@time plyapcs!(acs,ecs,R,adj = false);
x = R*R'; @test norm(acs*x*ecs'+ecs*x*acs'+F*F')/norm(x)/norm(acs) < reltol

F = UpperTriangular(rand(Ty,n,n)+im*rand(Ty,n,n))
R = copy(F)
@time plyapcs!(acs,ecs,R,adj = true);
x = R'*R; @test norm(acs'*x*ecs+ecs'*x*acs+F'*F)/norm(x)/norm(acs) < reltol

end

end
end

end


# testing ERB filterbank design
fb = make_erb_filterbank(16000, 23, 150)
fb_matlab = readcsv(open(joinpath(dirname(@__FILE__), "data", "ERB_filter_coeffs.csv")))
@test length(fb.filters) == size(fb_matlab, 1)
for k=1:length(fb.filters)
	sos = fb.filters[k]
	biquads = sos.biquads
	g = sos.g
	for b=1:4
		@test_approx_eq biquads[b].b0 fb_matlab[k,1]
		@test_approx_eq biquads[b].b2 fb_matlab[k,6]
		@test_approx_eq biquads[b].a1 fb_matlab[k,8]
		@test_approx_eq biquads[b].a2 fb_matlab[k,9]
	end
	@test_approx_eq biquads[1].b1 fb_matlab[k,2]
	@test_approx_eq biquads[2].b1 fb_matlab[k,3]
	@test_approx_eq biquads[3].b1 fb_matlab[k,4]
	@test_approx_eq biquads[4].b1 fb_matlab[k,5]
	@test_approx_eq 1/g fb_matlab[k,10]
end

# testing ERB filterbank response
x = zeros(10000)
x[1] = 1.0
y = filt(fb, x)
y_matlab = readcsv(open(joinpath(dirname(@__FILE__), "data", "ERB_filter_response.csv")))
@test_approx_eq y y_matlab'


# This diagram shows what is meant by "lower", "upper", and "large":
#      +1   +1
#       |    |
# upper |    |
#       |    |
#      -1    |
#            | large
#      +1    |
#       |    |
# lower |    |
#       |    |
#      -1   -1
#
# That is, we are only concerned with 2:1 subdivision of a surface/element.

# Calculate forward projection matrix for discrete L2 projection from large to upper
#
# Note: This is actually an interpolation.
function calc_forward_upper(n_nodes)
  # Calculate nodes, weights, and barycentric weights
  nodes, weights = gauss_lobatto_nodes_weights(n_nodes)
  wbary = barycentric_weights(nodes)

  # Calculate projection matrix (actually: interpolation)
  operator = zeros(n_nodes, n_nodes)
  for j in 1:n_nodes
    poly = lagrange_interpolating_polynomials(1/2 * (nodes[j] + 1), nodes, wbary)
    for i in 1:n_nodes
      operator[j, i] = poly[i]
    end
  end

  return operator
end

# Calculate forward projection matrix for discrete L2 projection from large to lower
#
# Note: This is actually an interpolation.
function calc_forward_lower(n_nodes)
  # Calculate nodes, weights, and barycentric weights
  nodes, weights = gauss_lobatto_nodes_weights(n_nodes)
  wbary = barycentric_weights(nodes)

  # Calculate projection matrix (actually: interpolation)
  operator = zeros(n_nodes, n_nodes)
  for j in 1:n_nodes
    poly = lagrange_interpolating_polynomials(1/2 * (nodes[j] - 1), nodes, wbary)
    for i in 1:n_nodes
      operator[j, i] = poly[i]
    end
  end

  return operator
end

# Calculate reverse projection matrix for discrete L2 projection from upper to large (Gauss version)
#
# Note: To not make the L2 projection exact, first convert to Gauss nodes,
# perform projection, and convert back to Gauss-Lobatto.
function calc_reverse_upper(n_nodes, ::Val{:gauss})
  # Calculate nodes, weights, and barycentric weights for Legendre-Gauss
  gauss_nodes, gauss_weights = gauss_nodes_weights(n_nodes)
  gauss_wbary = barycentric_weights(gauss_nodes)

  # Calculate projection matrix (actually: discrete L2 projection with errors)
  operator = zeros(n_nodes, n_nodes)
  for j in 1:n_nodes
    poly = lagrange_interpolating_polynomials(1/2 * (gauss_nodes[j] + 1), gauss_nodes, gauss_wbary)
    for i in 1:n_nodes
      operator[i, j] = 1/2 * poly[i] * gauss_weights[j]/gauss_weights[i]
    end
  end

  # Calculate Vandermondes
  lobatto_nodes, lobatto_weights = gauss_lobatto_nodes_weights(n_nodes)
  gauss2lobatto = polynomial_interpolation_matrix(gauss_nodes, lobatto_nodes)
  lobatto2gauss = polynomial_interpolation_matrix(lobatto_nodes, gauss_nodes)

  return gauss2lobatto * operator * lobatto2gauss
end

# Calculate reverse projection matrix for discrete L2 projection from lower to large (Gauss version)
#
# Note: To not make the L2 projection exact, first convert to Gauss nodes,
# perform projection, and convert back to Gauss-Lobatto.
function calc_reverse_lower(n_nodes, ::Val{:gauss})
  # Calculate nodes, weights, and barycentric weights for Legendre-Gauss
  gauss_nodes, gauss_weights = gauss_nodes_weights(n_nodes)
  gauss_wbary = barycentric_weights(gauss_nodes)

  # Calculate projection matrix (actually: discrete L2 projection with errors)
  operator = zeros(n_nodes, n_nodes)
  for j in 1:n_nodes
    poly = lagrange_interpolating_polynomials(1/2 * (gauss_nodes[j] - 1), gauss_nodes, gauss_wbary)
    for i in 1:n_nodes
      operator[i, j] = 1/2 * poly[i] * gauss_weights[j]/gauss_weights[i]
    end
  end

  # Calculate Vandermondes
  lobatto_nodes, lobatto_weights = gauss_lobatto_nodes_weights(n_nodes)
  gauss2lobatto = polynomial_interpolation_matrix(gauss_nodes, lobatto_nodes)
  lobatto2gauss = polynomial_interpolation_matrix(lobatto_nodes, gauss_nodes)

  return gauss2lobatto * operator * lobatto2gauss
end

# Calculate reverse projection matrix for discrete L2 projection from upper to large (Gauss-Lobatto
# version)
function calc_reverse_upper(n_nodes, ::Val{:gauss_lobatto})
  # Calculate nodes, weights, and barycentric weights
  nodes, weights = gauss_lobatto_nodes_weights(n_nodes)
  wbary = barycentric_weights(nodes)

  # Calculate projection matrix (actually: discrete L2 projection with errors)
  operator = zeros(n_nodes, n_nodes)
  for j in 1:n_nodes
    poly = lagrange_interpolating_polynomials(1/2 * (nodes[j] + 1), nodes, wbary)
    for i in 1:n_nodes
      operator[i, j] = 1/2 * poly[i] * weights[j]/weights[i]
    end
  end

  return operator
end

# Calculate reverse projection matrix for discrete L2 projection from lower to large (Gauss-Lobatto
# version)
function calc_reverse_lower(n_nodes, ::Val{:gauss_lobatto})
  # Calculate nodes, weights, and barycentric weights
  nodes, weights = gauss_lobatto_nodes_weights(n_nodes)
  wbary = barycentric_weights(nodes)

  # Calculate projection matrix (actually: discrete L2 projection with errors)
  operator = zeros(n_nodes, n_nodes)
  for j in 1:n_nodes
    poly = lagrange_interpolating_polynomials(1/2 * (nodes[j] - 1), nodes, wbary)
    for i in 1:n_nodes
      operator[i, j] = 1/2 * poly[i] * weights[j]/weights[i]
    end
  end

  return operator
end


module sequenceCompare
    export DTWDistance

    # http://en.wikipedia.org/wiki/Dynamic_time_warping
    function DTWDistance{T}(s::Array{T}, t::Array{T})
        n, m = size(s,2), size(t, 2)
        DTW = zeros(n+1,m+1)
        #println("sizeof DTW: ", size(DTW))
        DTW[:,1], DTW[1,:] = Inf, Inf
        DTW[1,1] = 0

        for i=2:n+1
            for j=2:m+1
                #println("i=", i, "/", n, " j=", j, "/", m, " size(s)=", size(s), " size(t)=", size(t))
                cost = sum((s[:,i-1] - t[:,j-1]).^2)

                DTW[i,j] = cost + min(DTW[i-1, j  ], #insertion
                                      DTW[i  , j-1], #deletion
                                      DTW[i-1, j-1]) #match
                #println("cost is ", cost, " i=", i, "/", n, " j=", j, "/", m, " DTW=", DTW[i,j])
            end
        end

        return DTW[n+1,m+1]
    end

    # note. there is also a windowed version for better performance!  See the wikipedia article
    function DTWDistance{T}(s::Array{T}, t::Array{T}, w::Int64)
        n, m = size(s,2), size(t, 2)
        w = max(w, abs(n-m))
        DTW = zeros(n+1,m+1)
        #println("sizeof DTW: ", size(DTW))
        DTW[:,:] = Inf
        DTW[1,1] = 0

        for i=2:n+1
            for j=max(2, i-w):min(m+1, i+w)
                #println("i=", i, "/", n, " j=", j, "/", m, " size(s)=", size(s), " size(t)=", size(t))
                cost = sum((s[:,i-1] - t[:,j-1]).^2)

                DTW[i,j] = cost + min(DTW[i-1, j  ], #insertion
                                      DTW[i  , j-1], #deletion
                                      DTW[i-1, j-1]) #match
                #println("cost is ", cost, " i=", i, "/", n, " j=", j, "/", m, " DTW=", DTW[i,j])
            end
        end

        return DTW[n+1,m+1]
    end

    function HaversineDistance{T}(lat1::T, lon1::T, lat2::T, lon2::T)
        REarth = 6371.
        lat, lon = (lat1-lat2)*π/180., (lon1-lon2)*π/180.
        lat1, lat2 = lat1*π/180., lat2*π/180.
        a = sin(lat/2)*sin(lat/2)+cos(lat1)*cos(lat2)*sin(lon/2)*sin(lon/2)
        d = 2*atan2(sqrt(a), sqrt(1-a))
        d = REarth*d

        return d
    end

    function RMSE{T}(pred::T, real::T)
        return sqrt(mean((pred-real).*(pred-real)))
    end

    function RMSLE{T}(pred::T, real::T)
    return sqrt(mean((log(pred+1)-log(real+1)).^2))
    end

    function meanHaversideDistance{T}(lat1::T, lon1::T, lat2::T, lon2::T)
        return mean(HaversideDistance{T}(lat1, lon1, lat2, lon2))
    end

    function TravelTimePredict(submission, answers)
        return RMSLE(submission, answers)
    end

    function DestinationPredict(submission, answers)
        lat_sub, lon_sub = submission[:,1], submission[:,2]
        lat_real, lon_real = answers[:,1], answers[:,2]
        return meanHaversineDistance(lat_sub, lon_sub, lat_real, lon_real)
    end

end


using StanSample, DiffEqBayesStan, MCMCChains
using OrdinaryDiffEq
using ParameterizedFunctions
using RecursiveArrayTools
using Distributions
using StatsPlots

ProjDir = @__DIR__

println("Four parameter case")

f1 = @ode_def begin
  dx = a*x - b*x*y
  dy = -c*y + d*x*y
end a b c d
u0 = [1.0,1.0]
tspan = (0.0,10.0)
p = [1.5,1.0,3.0,1.0]
prob1 = ODEProblem(f1,u0,tspan,p)
sol = solve(prob1,Tsit5())
t = collect(range(1,stop=10,length=50))
randomized = VectorOfArray([(sol(t[i]) + .3randn(2)) for i in 1:length(t)])
data_array = convert(Array, randomized)

priors = [truncated(Normal(1.5,0.1),0.5,2),truncated(Normal(1.0,0.1),0.5,1.5),
          truncated(Normal(3.0,0.1),0.5,4),truncated(Normal(1.0,0.1),0.5,2)]

bayesian_result = stan_inference(prob1, t, data_array, priors;
  nchains=[4], num_samples=1000, num_warmup=1000,
  vars =(DiffEqBayesStan.StanODEData(), InverseGamma(4,1)),
  output_format=:dataframe)

if success(bayesian_result.return_code)
  df = bayesian_result.chains

  println("\ntheta means:\n",
    [mean(df.theta_1), mean(df.theta_2), mean(df.theta_3), mean(df.theta_4)])
  println("\ntheta vars:\n",
    [var(df.theta_1), var(df.theta_2), var(df.theta_3), var(df.theta_4)])
  println()

  dfs = read_samples(bayesian_result.model; output_format=:dataframes)
  for j in 1:4
    for i in 1:4
      if i == 1
        plot(dfs[i][:, "theta_$j"], lab="chain[$j]")
      else
        plot!(dfs[i][:, "theta_$j"], lab="chain[$j]")
      end
    end
    savefig(joinpath(ProjDir, "tmp", "theta_$j.png"))
  end

  sdf = read_summary(bayesian_result.model)
  sdf |> display
  println("\n")

end


using Flux, DiffEqFlux, OrdinaryDiffEq, Optim

#A desired MWE for now, not a test yet.

function f(du,u,p,t)
    y₁,y₂,y₃ = u
    k₁,k₂,k₃ = p
    du[1] = -k₁*y₁ + k₃*y₂*y₃
    du[2] =  k₁*y₁ - k₃*y₂*y₃ - k₂*y₂^2
    du[3] =  y₁ + y₂ + y₃ - 1
    nothing
end

u₀ = [1.0, 0, 0]
M = [1. 0  0
    0  1. 0
    0  0  0]
tspan = (0.0,10.0)
p = [0.04,3e7,1e4]
func = ODEFunction(f,mass_matrix=M)
prob = ODEProblem(f,u₀,tspan,(0.04,3e7,1e4))
sol = solve(prob,Rodas5())

dudt2 = Chain(x -> x.^3,Dense(3,50,tanh),Dense(50,2))

ndae = NeuralDAE(dudt2, (u,p,t) -> [u[1] + u[2] + u[3] - 1], tspan, M, Rodas5())

ndae(u₀)

function predict_n_dae(p)
    ndae(u₀,p)
end

function loss(p)
    pred = predict_n_dae(p)
    loss = sum(abs2,sol .- pred)
    loss,pred
end

p = p .+ rand(3) .* p

res = DiffEqFlux.sciml_train(loss, p, BFGS(initial_stepnorm = 0.0001))


using MOSLab
using Plots
using LaTeXStrings
NpolyDoping = 5.5e17 # Gate Npoly doping concentration in cm^-3
N_a = 5.5e17 # Bulk doping concentration in cm^-3
E_a = 0.0# The ground state energy of the acceptor in eV
t_ox = 3.8e-7 # oxide thickness in cm
t_si = 1e-6
CList2 = reshape( range(colorant"blue", stop=colorant"red",length=7), 1, 7 );

MOSf(T) = MOSStructure(NPoly(NpolyDoping),SiO2(),SemiconductorData(T,BoltzmanDist(),PSilicon(N_a,E_a)),t_ox,t_si) ## Calculate Parameters of a MOS Structure the given parameters using Boltzman Distribution at temperature T
transistor_model(T) = BSIM6Model(MOSf(T),1.0e-4,1.0e-4)
Idf(Vg,Vd,T) = Id(Vg,Vd,0.0,transistor_model(T)) ## Calculate the Current using the ACM Model of a transistor having the parameters from MOSf(T) and W =1.0 um, L = 1.0 um, similar contructors are available for the other models
plot()
for Vgv in 0.2:0.2:1.8
    plot!(0:0.05:1.8,x->1e6*Idf(Vgv,x,300.0),label="\$V_{GB} = $(Vgv) V\$") # plot Id, Vd characteristics for different VGB
end
PSPG = plot!(xlabel=L"V_{DS} \ [V]",ylabel=L"I_{DS} \ [\mu A]",legend=:topleft)

plot()
for (i,Tv) in enumerate(LinRange(200,500,5))
    plot!(0:0.05:1.0,x->1e6*Idf(x,0.1,Tv),label="\$T = $(Tv) K\$",linecolor=CList2[i]) # plot Id, Vg characteristics for different Temperatures
end
PSPD = plot!(xlabel=L"V_{GS} \ [V]",ylabel=L"I_{DS} \ [\mu A]",legend=:topleft,yaxis=:log10)
plot(PSPG,PSPD) # plot both graphs side by side
savefig("TransistorCurves.png")

gmIdf(Vg,Vd,T) = gm(Vg,Vd,0.0,transistor_model(T))/Idf(Vg,Vd,T)
plot()
for (i,Tv) in enumerate(LinRange(200,500,5))
    plot!(0:0.01:1.0,x->gmIdf(x,0.1,Tv),label="\$T = $(Tv) K\$",linecolor=CList2[i]) # plot Id, Vg characteristics for different Temperatures
end
PSPD = plot!(xlabel=L"V_{GS} \ [V]",ylabel=L"I_{DS} \ [\mu A]",legend=:topright)
savefig("gmId.png")


using Test, Random, LinearAlgebra, SparseArrays

using Yao
using QuAlgorithmZoo

@testset "QFT" begin
    include("QFT.jl")
end

@testset "CircuitBuild" begin
    include("CircuitBuild.jl")
end

@testset "RotBasis" begin
    include("RotBasis.jl")
end

@testset "Grover" begin
    include("Grover.jl")
end

@testset "PhaseEstimation" begin
    include("PhaseEstimation.jl")
end

@testset "HHL" begin
    include("HHL.jl")
end
@testset "diff Eq" begin
    include("lin_diffEq_test.jl")
end

@testset "QCOptProblem" begin
    include("QCOptProblem.jl")
end

@testset "hamiltonian solvers" begin
    include("hamiltonian_solvers.jl")
end

@testset "hadamard test" begin
    include("HadamardTest.jl")
end

@testset "Sequence" begin
    include("Sequence.jl")
end

@testset "Diff" begin
    include("Diff.jl")
end


abstract Covariance
import Distances

abstract Kernel

function flatten_hyp(m)
	lik_hyp = Float64[]
	for n in names(m.lik)
		push!(lik_hyp, m.lik.(n)...)
	end
	cov_hyp = Float64[]
	for n in names(m.k)
		push!(cov_hyp, m.k.(n)...)
	end
	vcat(lik_hyp, cov_hyp)

end

function inflate_hyp(m, hyp_vec)
	liks = hyp_vec[1:length(m.lik)]
	covs = hyp_vec[length(m.lik) + 1:end]
	typeof(m.k)(covs), typeof(m.lik)(liks)
end

type SqExpIso <: Kernel
	log_l
	log_s
end
SqExpIso(v::Array) = SqExpIso(v[1], v[2])

function kernel{T <: FloatingPoint}(h::SqExpIso, X1::Array{T, 2},
						X2::Array{T, 2})
	l2 = exp(2h.log_l)
	s2 = exp(2h.log_s)
	s2 * exp(-Distances.pairwise(Distances.SqEuclidean(), X1', X2') / (2l2))
end
kernel{T <: FloatingPoint}(h::SqExpIso, X::Array{T, 2}) = kernel(h, X, X)

type SqExpIsoPeriodic <: Kernel
	log_l
	log_s
	log_period
end
SqExpIsoPeriodic(v::Array) = SqExpIsoPeriodic(v[1], v[2], v[3])

function kernel{T <: FloatingPoint}(h::SqExpIsoPeriodic, X1::Array{T, 2},
						X2::Array{T, 2})
	l2 = exp(2h.log_l)
	s2 = exp(2h.log_s)
	period1 = exp(h.log_period)
	s2 *
	exp(-Distances.pairwise(Distances.SqEuclidean(),
		cos(2pi * X1 / period1)', cos(2pi * X2 / period1)') / (2l2))
end
kernel{T <: FloatingPoint}(h::SqExpIsoPeriodic, X::Array{T, 2}) = kernel(h, X, X)

type SqExpIso2Periodic <: Kernel
	log_l
	log_s
	log_period1
	log_period2
end
SqExpIso2Periodic(v::Array) = SqExpIso2Periodic(v[1], v[2], v[3], v[4])

function kernel{T <: FloatingPoint}(h::SqExpIso2Periodic, X1::Array{T, 2},
						X2::Array{T, 2})
	l2 = exp(2h.log_l)
	s2 = exp(2h.log_s)
	period1 = exp(h.log_period1)
	period2 = exp(h.log_period2)
	s2 *
	(exp(-Distances.pairwise(Distances.SqEuclidean(),
		cos(2pi * X1 / period1)', cos(2pi * X2 / period1)') / (2l2)) .+
	exp(-Distances.pairwise(Distances.SqEuclidean(),
		cos(2pi * X1 / period2)', cos(2pi * X2 / period2)') / (2l2)))
end
kernel{T <: FloatingPoint}(h::SqExpIso2Periodic, X::Array{T, 2}) = kernel(h, X, X)

type SqExpIso3Periodic <: Kernel
	log_l
	log_s
	log_period1
	log_period2
	log_period3
end
SqExpIso3Periodic(v::Array) = SqExpIso(v[1], v[2], v[3], v[4], v[5])

function kernel{T <: FloatingPoint}(h::SqExpIso3Periodic, X1::Array{T, 2},
						X2::Array{T, 2})
	l2 = exp(2h.log_l)
	s2 = exp(2h.log_s)
	period1 = exp(h.log_period1)
	period2 = exp(h.log_period2)
	period3 = exp(h.log_period3)
	s2 *
	(exp(-Distances.pairwise(Distances.SqEuclidean(),
		cos(2pi * X1 / period1)', cos(2pi * X2 / period1)') / (2l2)) .+
	exp(-Distances.pairwise(Distances.SqEuclidean(),
		cos(2pi * X1 / period2)', cos(2pi * X2 / period2)') / (2l2)) .+
	exp(-Distances.pairwise(Distances.SqEuclidean(),
		cos(2pi * X1 / period3)', cos(2pi * X2 / period3)') / (2l2)))
end
kernel{T <: FloatingPoint}(h::SqExpIso3Periodic, X::Array{T, 2}) = kernel(h, X, X)

type SqExpIso4Periodic <: Kernel
	log_l
	log_s
	log_period1
	log_period2
	log_period3
	log_period4
end
SqExpIso4Periodic(v::Array) = SqExpIso(v[1], v[2], v[3], v[4],
	v[5], v[6], v[7])

function kernel{T <: FloatingPoint}(h::SqExpIso4Periodic, X1::Array{T, 2},
						X2::Array{T, 2})
	l2 = exp(2h.log_l)
	s2 = exp(2h.log_s)
	period1 = exp(h.log_period1)
	period2 = exp(h.log_period2)
	period3 = exp(h.log_period3)
	period4 = exp(h.log_period4)
	s2 *
	(exp(-Distances.pairwise(Distances.SqEuclidean(),
		cos(2pi * X1 / period1)', cos(2pi * X2 / period1)') / (2l2)) .+
	exp(-Distances.pairwise(Distances.SqEuclidean(),
		cos(2pi * X1 / period2)', cos(2pi * X2 / period2)') / (2l2)) .+
	exp(-Distances.pairwise(Distances.SqEuclidean(),
		cos(2pi * X1 / period3)', cos(2pi * X2 / period3)') / (2l2)) .+
	exp(-Distances.pairwise(Distances.SqEuclidean(),
		cos(2pi * X1 / period4)', cos(2pi * X2 / period4)') / (2l2)))
end
kernel{T <: FloatingPoint}(h::SqExpIso4Periodic, X::Array{T, 2}) = kernel(h, X, X)

type MaternIso32 <: Kernel
	log_l
	log_s
end
MaternIso32(v::Array) = MaternIso32(v[1], v[2])
function kernel{T <: FloatingPoint}(h::MaternIso32, X1::Array{T, 2},
						X2::Array{T, 2})
	l = exp(h.log_l)
	s = exp(h.log_s)
	s3dol = sqrt(3) * Distances.pairwise(Distances.Euclidean(), X1', X2') / l
	s * (1 .+ s3dol) .* exp(-s3dol)
end
kernel{T <: FloatingPoint}(h::MaternIso32, X::Array{T, 2}) = kernel(h, X, X)

type MaternARD32 <: Kernel
	log_ls::Array{Float64,}
	log_s
end
MaternARD32(v::Array) = MaternARD32(v[1:end-1], v[end])
function kernel{T <: FloatingPoint}(h::MaternARD32, X1::Array{T, 2},
						X2::Array{T, 2})
	ls2 = 2exp(2h.log_ls[:])
	s = exp(h.log_s)
	s3dol = sqrt(3) * Distances.pairwise(Distances.WeightedEuclidean(1./ls2), X1', X2')
	s * (1 .+ s3dol) .* exp(-s3dol)
end
kernel{T <: FloatingPoint}(h::MaternARD32, X::Array{T, 2}) = kernel(h, X, X)


using ForwardDiff
using RootSolvers

function T2_h_i(β, α::FT, i::IT, N::IT) where {FT, IT}
  a = α; b = β
  g = (b+FT(1))/(b-FT(1))
  t = (FT(i)/FT(N) - a)/(FT(1)-a)
  f = ((b+FT(2)*a)*g^t - b + FT(2)*a)/((FT(2)*a+FT(1))*(FT(1)+g^t))
  return f
end

T2_root_Δh_near_hmax(β, hmin::FT,hmax::FT,α::FT,N::IT,Δh::FT) where {FT,IT} =
  T2_h_i(β,α,N,N) - T2_h_i(β,α,N-1,N) - Δh/(hmax-hmin)

T2_root_Δh_near_hmin(β, hmin::FT,hmax::FT,α::FT,N::IT,Δh::FT) where {FT,IT} =
  T2_h_i(β,α,1,N) - T2_h_i(β,α,0,N) - Δh/(hmax-hmin)

"""
    β_Δh_small

Compute `β` such that `N` grid points
between `hmin` and `hmax` stretch so
that the smallest `Δh` matches the one
given.

!!! warn
    If `N` is small, `β->∞`, and the
    given `Δh` may not be satisfied.
    `N` can be increased to fix this.
"""
function β_Δh_small(hmin::FT, hmax::FT, N::IT, Δh::FT) where {FT, IT}
  tol = eps(FT)
  β_initial = FT(1) + tol # Initial guess
  root_closure(β) = T2_root_Δh_near_hmax(β, hmin, hmax, FT(0), N, Δh)
  args = (NewtonsMethodAD(β_initial), CompactSolution(), SolutionTolerance(eps(FT)))
  sol = find_zero(β -> root_closure(β), args...)
  return sol.root
end

"""
    β_Δh_big

Compute `β` such that `N` grid points
between `hmin` and `hmax` stretch so
that the largest `Δh` matches the one
given.

!!! warn
    If `N` is small, `β->∞`, and the
    given `Δh` may not be satisfied.
    `N` can be increased to fix this.
"""
function β_Δh_big(hmin::FT, hmax::FT, N::IT, Δh::FT) where {FT, IT}
  tol = eps(FT)
  β_initial = FT(1) + tol # Initial guess
  root_closure(β) = T2_root_Δh_near_hmin(β, hmin, hmax, FT(0), N, Δh)
  args = (NewtonsMethodAD(β_initial), CompactSolution(), SolutionTolerance(eps(FT)))
  sol = find_zero(β -> root_closure(β), args...)
  return sol.root
end

"""
    β_Δh_both

Compute `β` such that `N` grid points
between `hmin` and `hmax` stretch so
that the smallest `Δh` matches the one
given.

!!! warn
    If `N` is small, `β->∞`, and the
    given `Δh` may not be satisfied.
    `N` can be increased to fix this.
"""
function β_Δh_both(hmin::FT, hmax::FT, N::IT, Δh::FT) where {FT, IT}
  tol = eps(FT)
  β_initial = FT(1) + tol # Initial guess
  root_closure(β) = T2_root_Δh_near_hmax(β, hmin, hmax, FT(1/2), N, Δh)
  args = (NewtonsMethodAD(β_initial), CompactSolution(), SolutionTolerance(eps(FT)))
  sol = find_zero(β -> root_closure(β), args...)
  return sol.root
end


function euclidean(v1::Array{T1,1}, v2::Array{T2,1})::Float64 where {T1 <: Number,T2 <: Number}
    (v1 .- v2).^2.0 |> sum |> sqrt
end

function euclidean(v1::Array{T,1})::Float64 where T <: Number
    v2 = zeros(Float64, length(v1))
    return euclidean(v1, v2)
end

function normalize(v1::Array{T,1})::Array{Float64,1} where T <: Number
    return v1 ./ euclidean(v1)
end

function normalize(data::DataFrame)::DataFrame
    df = similar(data)
    _, p = size(df)
    for i in 1:p
        df[:, i] = normalize(data[:, i])
    end
    return df
end

function apply_columns(fs::Array{Function,1}, data::Union{DataFrame,Array{T,2}} where T <: Number)
    _, m = size(data)
    return [fs[i](data[:,i]) for i in 1:m]
end

function apply_columns(f::Function, data::Union{DataFrame,Array{T,2}} where T <: Number)
    return [f(c) for c in eachcol(data)]
end

function colmins(data::DataFrame)
    return apply_columns(minimum, data)
end

function colmaxs(data::DataFrame)
    return apply_columns(maximum, data)
end

function colsums(data)::Array{Number,1}
    return apply_columns(sum, data)
end

function colmeans(data)::Array{Number,1}
    return apply_columns(mean, data)
end

function apply_rows(f::Function, data::Union{DataFrame,Array{T,2}} where T <: Number)
    return [f(c) for c in eachrow(data)]
end

function rowmins(data::DataFrame)::Array{Float64,1}
    return apply_rows(minimum, data)
end

function rowmaxs(data::DataFrame)::Array{Float64,1}
    return apply_rows(maximum, data)
end

function rowsums(data)::Array{Float64,1}
    return apply_rows(sum, data)
end

function rowmeans(data)::Array{Float64,1}
    return apply_rows(mean, data)
end

function unitize(v::Array{T,1})::Array{Float64,1} where {T <: Number}
    return v ./ sum(v)
end

function Base.:*(w::Array{Float64,1}, data::DataFrame)::DataFrame
    newdf = copy(data)
    _, p = size(newdf)
    for i in 1:p
        newdf[!, i] = convert.(Float64, newdf[!,i])
        newdf[:, i] = w[i] .* data[:, i]
    end
    return newdf
end

function Base.:-(r1::DataFrameRow, r2::DataFrameRow)::Array{Number,1}
    # v1 = convert(Array{Float64,1}, r1)
    v1 = Vector{Float64}(r1)
    # v2 = convert(Array{Float64,1}, r2)
    v2 = Vector{Float64}(r2)
    return v1 .- v2
end

function Base.:-(r1::Array{T,1}, r2::DataFrameRow)::Array{T,1} where T <: Number
    # v2 = convert(Array{Float64,1}, r2)
    v2 = Vector{Float64}(r2)
    return r1 .- v2
end

function Base.:-(r1::DataFrameRow, r2::Array{T,1})::Array{T,1} where T <: Number
    v1 = convert(Array{Float64,1}, r1)
    return v1 .- r2
end

function makeDecisionMatrix(mat::Array{T,2}; names::Union{Nothing,Array{String,1}}=nothing)::DataFrame where {T <: Number}
    _, m = size(mat)
    df = DataFrame()
    for i in 1:m
        if names isa Nothing
            name = string("Crt", i)
        else
            name = names[i]
        end
        df[:,Symbol(name)] = convert(Array{Float64,1}, mat[:, i])
    end
    return df
end

function Base.minimum(df::DataFrame)
    df |> Matrix |> minimum
end

function reverseminmax(fns::Array{Function,1})::Array{Function,1}
    newfs = map(x -> if x == minimum maximum else minimum end, fns)
    return newfs
end

function makeminmax(fns::Array{K,1} where K)::Array{Function,1}
    return convert(Array{Function,1}, fns)
end


using CrystallographyBase: Lattice
using LinearAlgebra: I

export distortby, distort, strainstate

# See https://link.springer.com/content/pdf/10.1007%2F978-3-7091-0382-1_7.pdf and https://doi.org/10.2138/am-1997-1-207
distortby(lattice::Lattice, strain::TensorStrain) =
    Lattice((I + strain.data) * lattice.data)
distortby(lattice::Lattice, strain::EngineeringStrain) =
    distortby(lattice, TensorStrain(strain))
const distort = distortby  # For the sake of compatibility

strainstate(old::Lattice, new::Lattice) = TensorStrain(new.data / old.data - I)


using JLD2

function oscillator()
    A = [0.0 1; -4π 0]
    X0 = BallInf([0.0, 10.0], 0.1)
    X = Universe(2)
    P = @ivp(x' = A*x, x ∈ X, x(0) ∈ X0)
    return P
end

function freedom()
    m1 = 1.0
    m2 = 1.0
    k1 = 1e4
    k2 = 1.0

    M = [m1 0; 0 m2]
    K = [(k1+k2) -k2; -k2 k2]

    sys = SecondOrderLinearContinuousSystem(M, zeros(size(M)), K)

    n = 2
    A = [zeros(n, n) Matrix(1.0I, n, n)
         -inv(M)*K    zeros(n, n)       ]

    X0 = Hyperrectangle([1.0, 10.0, 0.0, 0.0], [0.1, 0.5, 0.5, 0.5])
    P = @ivp(x' = Ax, x(0) ∈ X0)
    P = normalize(P)
    return P
end

function iss()
    ISS_path = joinpath(@__DIR__, "ISS.jld2")
    @load ISS_path A B

    A = Matrix(A)  # some methods cannot deal with sparse matrices
    U = Hyperrectangle(low=[0.0, 0.8, 0.9], high=[0.1, 1., 1.])
    U = ConstantInput(linear_map(B, U))
    X = Universe(270)
    X0 = BallInf(zeros(size(A, 1)), 0.0001)  # -0.0001 <= xi <= 0.0001 for all i
    P = @ivp(x' = A*x + u, x(0) ∈ X0, u ∈ U, x ∈ X)
    return P
end


####
# LCX Oracle tests
####

facts("suppFcnTest LCX") do
	srand(8675309); data = randn(100, 2)

	#VG There is a disturbing dependency on the solver here!!!!
	w = LCXOracle(data, .1, 0; Gamma =.35)
	zstar, ustar = suppFcn([1, 1], w, :Min)
	@fact zstar --> roughly(-3.1455496412236243)
	@fact ustar[1] --> roughly(-1.313810356077738)
	@fact ustar[2] --> roughly(-1.8317392851458862)

	zstar, ustar = suppFcn([1, 1], w, :Max)
	@fact zstar --> roughly(3.447833040263504)
	@fact ustar[1] --> roughly(1.5387299947906286)
	@fact ustar[2] --> roughly(1.9091030454728752)

	zstar, ustar = suppFcn([-1, 1], w, :Min)
	@fact zstar --> roughly(-2.8182398717617616)
	@fact ustar[1] --> roughly(1.3068007701076567)
	@fact ustar[2] --> roughly(-1.5114391016541049)

	zstar, ustar = suppFcn([-1, 1], w, :Max)
	@fact zstar --> roughly(3.157077158917673)
	@fact ustar[1] --> roughly(-1.7212256391067058)
	@fact ustar[2] --> roughly(1.4358515198109671)
end

facts("portTest LCX") do
	srand(8675309); data = randn(500, 2)
	w = LCXOracle(data, .1, .2, Gamma=.2, ab_cut_tol=5e-6)
	portTest(w, -1.4293015271256384, [0.5424436347803758, 0.4575563652196242])
end