export asqp, solveKKT, getLinearIndependentConstraints

"""
function asqp(G,c,Ae,be,Ai,bi,x=0*c; maxIter, tolP, out)

Find optimizer for quadratic programming problem using active set method

min_x q(x) = 0.5\*x'\*G*x + c'\*x
subject to Ax >= b

HERE: Assume that active constraints are linearly independent.

	Input:
        1. G 	 - matrix of weights for x, nxn
        2. c 	 - linear vector of weights, nx1
        3. Ai/Ae - matrix of inequality/equality constraints, mxn or pxn
        4. bi/be - inequality/equality constraints to satisfy, mx1 or px1
        5. x 	 - initial starting point if provided or default to origin [0,0,...,0], nx1

	Keyword arguments:
        6. maxIter - number of iterations for finding feasible set and direction before ending asqp
        7. tolP    - tolerance for ending search for direction and asqp
        8. out	   - print intermediate outputs if out > 1

    Output:
        1. x 	- minimizer for quadratic programming problem, nx1
        2. iter - number of iterations to find minimizer
        3. Wci  - inequality working conditions at minimizer, m'xn

"""
function asqp(G,c,Ae,be,Ai,bi,xc=0*c; maxIter::Int64=5, tolP::Float64=1e-10,out::Int64 = 2)

	# if rank(G) < size(G,2)
	# 	Ginv = x->(G+1e-6*eye(size(G,2)))\x
	# else
	# 	Ginv = x->G\x
	#  end

    # println("########## ASQP start ########### \n Ae $Ae \n be $be \n Ai $Ai \n bi $bi \n G $G \n c $c \n x0 $xc")
    Wci = find(Ai*xc-bi .< 0)
	if (length(Wci)>0 && size(Ae,1)>0)  || length(Wci)>1
    	Wci = getLinearIndependentConstraints([Ae;Ai],[vec(collect(1:size(Ae,1)));  size(Ae,1)+Wci])
		Wci = Wci[size(Ae,1)+1:end]-size(Ae,1)
	end
    Aic = Ai[Wci,:]
    # println("Aic $Aic")

    iter = 0
    while true
        if out>1
            qc = .5*xc'*G*xc + c'*xc
            nlwc = length(Wci)
            println("i=$iter,\t q(x)= $qc, \t numel(Wci)=$nlwc")
        end

        # println(" Current xc = $xc")

        # update gradient of QP
        gc  = G*xc+c

        # update pc and lambda of inequality constraints
        pc,lami = solveKKT(xc,G,gc,Ae,be,Aic,bi[Wci])
        # pc,lami = solveSchur(xc,G,gc,Ae,be,Aic,bi[Wci],Ginv)

        # println("\n ##### iter $iter ##### \n pc $pc \t lami ", lami)

        if norm(pc)<tolP # if stepsize to x* is small
            # check Lagrange multipliers
            if all(lami.>=0) # we are done!
                break #return xc,iter,Wc
            else # remove constraint
                l_, idm = findmin(lami); # find inequality working set index corresponding to smallest value of multiplier
                Wci = deleteat!(Wci,idm) # remove one inequality constraint
                # println("removed one inequality constraint, new Wci $Wci")
                Aic  = Ai[Wci,:] # update active constraints
            end
        else
            # add one constraint to active set
            t2 = map(x -> min(0,x), (bi-Ai*xc)) # keep only negative entries that satisfy inequaity constraint, t2 has only entries <= 0, necessary to ensure alpha > 0
            t1 = (Ai*pc);
            t  = t2./t1;
            # println("bi-Ai*xc ", bi-Ai*xc)
            # println("\n t = t2./t1  $t \n t1 = Ai*pc $t1")
            # println("\n t2 = min(0,bi-Ai*xc) $t2")
            # println("t1 >= 0 \t", 1*(t1.>=0))
            t[find(t1 .>= 0)] = 1; # set entries in t equal 1 if t1 entries >= 0, necessary to ensure alpha > 0
            t[Wci]    = 1;
            # println("t[Wci]=1 $t")
            alpha = min(1.0,minimum(t)); # minimum for finding array min only, find step size
            if alpha<0
                error("alpha negative");
            end
            xc = xc + alpha*pc;
            # println("xc + alpha*pc $xc")
            if alpha<1 # add one blocking constraint
                idm = findin(t, alpha); # find indicies of entries in t == alpha
				# Wci = sort([Wci;idm[1]]); # sort to ensure index for inequality constraints are above equality constraints
				Wci = [Wci;idm[1]];
                Aic  = Ai[Wci,:] # found working set
            end
        end
        # check stopping
        if iter > maxIter
            break
        end
        iter = iter+1;
    end
    return xc, iter, Wci
end

"""
 function solveKKT(xc,G,gc,Ae,be,Ai,bi)

 Solves KKT system

       | G   Ae'  Ai' | |  p   |   | -gc |
 KKT = | Ae   0   0   | | -lame| = | -he |
       | Ai   0   0   | | -lami|   | -hi |

 gc = c+G*xc, he = Ae*xc-be, hi = Ai*xc-bi

 Input:
	 1. xc	  - current point, nx1
	 1. G 	  - matrix of weights for x, nxn
	 2. gc 	  - gradient of quadratic problem, nx1
	 3. Ai/Ae - matrix of inequality/equality constraints, mxn or pxn
	 4. bi/be - inequality/equality constraints to satisfy, mx1 or px1

Output:
	1. p 	- stepsize in direction of satisfying both inequality and equality constraints to update xc
	2. lami - Lagrange multipliers associated with inequaity constraints

"""
function solveKKT(xc,G,gc,Ae,be,Ai,bi)

    hi = Ai*xc-bi # h = 0, x0 = 0, b = 0
    he = Ae*xc-be

    KKT = [G Ae' Ai';
           Ae zeros(size(Ae,1), size(Ae,1)+size(Ai,1));
           Ai zeros(size(Ai,1), size(Ae,1)+size(Ai,1)) ]

	if rank(KKT) < size(KKT,1)
	   KKT = [G+1e-6*eye(size(G,1)) Ae' Ai';
	           Ae zeros(size(Ae,1), size(Ae,1)+size(Ai,1));
	           Ai zeros(size(Ai,1), size(Ae,1)+size(Ai,1)) ]
	end
    # println("solveKKT $KKT \n rank(KKT) ", rank(KKT))
    # println("\n ###### solveKKT ###### \n rank(KKT) $(rank(KKT)) \t size(KKT) ", size(KKT), " \t cond(KKT) ", cond(KKT), "\n xc $xc \n gc $gc \n G $G \n cond(G) ", cond(G)) #1e6 or larger

    rhs = [-gc; -he; -hi];

    try
        xsol = KKT\rhs;
        # println("inside try xsol $xsol")
        catch E
        println("ERROR!!!! \n rank(KKT) $(rank(KKT)) \t size(KKT) ", size(KKT), " \t cond(KKT) ", cond(KKT), "\n xc $xc \n gc $gc \n Ai $Ai \n bi $bi \n be $be \n Ae $Ae \n G $G \n rank(G) ", rank(G)) #1e6 or larger
        # JLD.save("asqp error case.jld","G", G, "xc", xc, "gc", gc, "Ai", Ai, "bi", bi, "Ae", Ae, "be", be)
        # MAT.matwrite("test.mat", Dict(
                #"G" => G,
                #"xc" => xc,
                #"gc" => gc,
                #"Ai" => Ai,
                #"bi" => bi,
                #"Ae" => Ae,
                #"be" => be,
                #"KKT" => KKT
                #))
        throw(E)
    end

    xsol = KKT\rhs
    p = xsol[1:length(gc)]
    lame = -xsol[length(gc)+(1:size(Ae,1))];
    lami = -xsol[length(gc)+size(Ae,1)+1:end];

    # println("\n p and lam $xsol \n Ae $Ae \n Ai ", Ai)

    return p,lami
end
#
#
# """
#  solveSchur(xc,G,gc,Ae,be,Ai,bi,Ginv=x->G\\x)
#
#  Solves KKT system
#
#        | G   Ae'  Ai' | |  p   |   | -gc |
#  KKT = | Ae   0   0   | | -lame| = | -he |
#        | Ai   0   0   | | -lami|   | -hi |
#
#  using Schur complement
#
#  he = Ae*xc-be, gc = c+G*xc, hi = Aic*xc-bic
#
# """
# function solveSchur(xc,G,gc,Ae,be,Ai,bi,Ginv=x->G\x)
#     # println("##### solve Schur #######")
#
#     A = [Ae;Ai]
#     b = [be;bi]
#     elen = length(be)
#     #println("A $A \t x $xc \t b $b \n G $G \t Wc $Wc")
#
#     h   = A*xc-b
#     if isempty(A)
#         lami = zeros(0)
#         p    = Ginv(-gc)
#     else
#         lam  = (A*(Ginv(A')))\(A*Ginv(gc)-h)  # note that lam has a minus sign in KKT system
# # println("lam = $lam")
# # println("A'*lam-gc=$(A'*lam-gc)")
#         p    = Ginv(A'*lam-gc)
#         lami = lam[elen+1:end]
#     end
#     return p, lami
# end

"""
function getLinearIndependentConstraints(A,Wc)

Find linearly independent rows of A (working set of constraints are independent of each other)

Input:
	1. A  - set of constraints in matrix form, matrix of integers
	2. Wc - current working set of constraints, array of integers

Output:
	1. Wc - updated working set of constraints, array of integers

"""

function getLinearIndependentConstraints(A,Wc)
    if length(Wc)==1 # only working constraint
        return Wc
    end
    # save('/tmp/getLinearIndependentConstraints.mat')

    # println("##### getLinearIndependentConstraints ###### \n A $A \n Wc $Wc")
    # make matrix smaller using its sparsity
    A  = A[Wc,:];
    id = find(0.< sum(abs.(A),1)); # sum along colums of A > 0
    # println("id $id \n A $A")
    A = A[:,id[:]]; # restrict to nonzero columns
    # println("A ", A)

    # check if rank is full
    if maximum(size(A)) <=3000
        if rank(full(A)) == length(Wc)
            #println("full rank A")
            return Wc
        end
    end

    debit = 1;
    credit = 2:size(A,1);

    # eliminate linear dependent constraints row by row
    Aac = A[debit,:];
    # println("Aac before credit ", Aac)

    while !(isempty(credit))
        # tic()
        At  = A[credit,:]';
        # println("At ", At, "\n Aac ", size(Aac)) #row vector always as columns
        xx = (Aac)\At; # solve system after adding new row (constraints)
        # println("xx ", xx)
        res = full(Aac*xx-At); # convert sparse to full matrix
        # println("res ", res)
        nres = sum(res.^2,1); # residual sum of squares along column
        mx,ii = findmax(nres); # find column with largest RSS
        # println("max is $mx")

        if mx>1e-3
            debit  = [debit;credit[ii]];
            # println("debit ", debit)
            credit = deleteat!(Vector(credit),ii) # remove row of constraints from inactive set
            # println("credit ", credit, "\n At[:,ii]' ", size(At[:,ii],2), "\ Aac ", Aac)
            Aac = [Aac'; At[:,ii]']; # update active set of constraints
            # println("credit ", credit, "\n Aac ", Aac)
        else
            break;
        end
        # toc()
    end

    Wc = Wc[debit]
    return Wc
end


module UnitfulAngles

import Unitful
import Dates
using Unitful: @unit, Quantity, NoDims, @u_str, uconvert, ustrip
import Base: sin, cos, tan, sec, csc, cot, asin, acos, atan, asec, acsc, acot, convert

######################### Angle units ##########################################
@unit turn          "τ"             Turn          2π*u"rad"     false
@unit doubleTurn    "§"             DoubleTurn    2turn         false
@unit halfTurn      "π"             HalfTurn      turn//2       false
@unit quadrant      "⦜"             Quadrant      turn//4       false
@unit sextant       "sextant"       Sextant       turn//6       false
@unit octant        "octant"        Octant        turn//8       false
@unit clockPosition "clockPosition" ClockPosition turn//12      false
@unit hourAngle     "hourAngle"     HourAngle     turn//24      false
@unit compassPoint  "compassPoint"  CompassPoint  turn//32      false
@unit hexacontade   "hexacontade"   Hexacontade   turn//60      false
@unit brad          "brad"          BinaryRadian  turn//256     false
@unit diameterPart  "diameterPart"  DiameterPart  1u"rad"/60    false # ≈ turn/377
@unit grad          "ᵍ"             Gradian       turn//400     false
@unit arcminute     "′"             Arcminute     u"°"//60      false # = turn/21,600
@unit arcsecond     "″"             Arcsecond     u"°"//3600    false # = turn/1,296,000

######################### Functions ############################################

# cos and sin have *pi versions, and *d versions
for _f in (:cos, :sin)
    @eval $_f(x::Quantity{T, NoDims, typeof(halfTurn)}) where {T} = $(Symbol("$(_f)pi"))(ustrip(x))
    @eval $_f(x::Quantity{T, NoDims, typeof(diameterPart)}) where {T} = $_f(ustrip(uconvert(u"rad", x)))
    for _u in (doubleTurn, turn, quadrant, sextant, octant, clockPosition, hourAngle, compassPoint, hexacontade, brad, grad, arcminute, arcsecond)
        @eval $_f(x::Quantity{T, NoDims, typeof($_u)}) where {T} = $(Symbol("$(_f)pi"))(ustrip(uconvert(halfTurn, x)))
    end
end

# These functions don't have *pi versions, but have *d versions
for _f in (:tan, :sec, :csc, :cot)
    @eval $_f(x::Quantity{T, NoDims, typeof(diameterPart)}) where {T} = $_f(ustrip(uconvert(u"rad", x)))
    for _u in (doubleTurn, turn, halfTurn, quadrant, sextant, octant, clockPosition, hourAngle, compassPoint, hexacontade, brad, grad, arcminute, arcsecond)
        @eval $_f(x::Quantity{T, NoDims, typeof($_u)}) where {T} = $(Symbol("$(_f)d"))(ustrip(uconvert(u"°", x)))
    end
end

# Inverse functions
for _f in (:acos, :asin, :atan, :asec, :acsc, :acot),
    _u in (diameterPart, u"°", u"rad", doubleTurn, turn, halfTurn, quadrant, sextant, octant, clockPosition, hourAngle, compassPoint, hexacontade, brad, grad, arcminute, arcsecond)
    @eval $_f(::typeof($_u), x::Number) = uconvert($_u, $_f(x)*u"rad")
end

for _u in (diameterPart, u"°", u"rad", doubleTurn, turn, halfTurn, quadrant, sextant, octant, clockPosition, hourAngle, compassPoint, hexacontade, brad, grad, arcminute, arcsecond)
    @eval atan(::typeof($_u), y::Number, x::Number) = uconvert($_u, atan(y, x)*u"rad")
end

_tohour(x::T) where {T <: Dates.TimePeriod} = x/convert(T, Dates.Hour(1)) # credit to kristoffer.carlsson see https://discourse.julialang.org/t/convert-time-interval-to-seconds/3806/2

# Fun conversion between time and angles
# NOTE: not sure whether to use `convert` or `uconvert`
for _u in (diameterPart, u"°", u"rad", doubleTurn, turn, halfTurn, quadrant, sextant, octant, clockPosition, hourAngle, compassPoint, hexacontade, brad, grad, arcminute, arcsecond)
    @eval begin
        function convert(::typeof($_u), t::Dates.Time)
            x = t - Dates.Time(0, 0, 0)
            uconvert($_u, _tohour(x)*hourAngle)
        end
    end
    @eval convert(::Type{Dates.Time}, x::Quantity{T, NoDims, typeof($_u)}) where {T} = Dates.Time(0, 0, 0) + Dates.Nanosecond(round(Int, ustrip(uconvert(hourAngle, x))*3600000000000))
end

# Enable precompilation with Unitful extended units
# http://ajkeller34.github.io/Unitful.jl/stable/extending/#precompilation
const localunits = Unitful.basefactors
function __init__()
    merge!(Unitful.basefactors, localunits)
    Unitful.register(UnitfulAngles)
end

end # module


using
  PyPlot,
  Printf,
  FourierFlows

using FFTW: rfft

import GeophysicalFlows.TwoDTurb

 dev = CPU()   # Device (CPU/GPU)

  nx = 256     # Resolution
  Lx = 2π      # Domain size
   ν = 1e-6    # Viscosity
  nν = 1       # Order of (hyper-)viscosity. nν=1 means Laplacian
  dt = 0.1     # Timestep
nint = 200     # Number of steps between plots
ntot = 10nint  # Number of total timesteps

# Define problem
prob = TwoDTurb.Problem(nx=nx, Lx=Lx, ν=ν, nν=nν, dt=dt, stepper="FilteredRK4", dev=dev)
TwoDTurb.set_zeta!(prob, ArrayType(dev)(rand(nx, nx)))

cl, vs, gr = prob.clock, prob.vars, prob.grid
x, y = gridpoints(gr)

"Plot the vorticity of the twodturb problem `prob`."
function makeplot!(ax, prob)
  sca(ax)
  cla()
  pcolormesh(x, y, vs.zeta)
  title("Vorticity")
  xlabel(L"x")
  ylabel(L"y")
  pause(0.01)
  nothing
end

# Step forward
fig1, ax = subplots(figsize=(8, 8))

while cl.step < ntot
  @time begin
    stepforward!(prob, nint)
    @printf("step: %04d, t: %6.1f", cl.step, cl.t)
  end

  TwoDTurb.updatevars!(prob)
  makeplot!(ax, prob)
end

# Plot the radial energy spectrum
E  = @. 0.5*(vs.u^2 + vs.v^2) # energy density
Eh = rfft(E)
kr, Ehr = FourierFlows.radialspectrum(Eh, gr, refinement=1)

fig2, axs = subplots(ncols=2, figsize=(8, 4))

sca(axs[1])
pcolormesh(x, y, vs.zeta)
xlabel(L"x")
ylabel(L"y")
title("Vorticity")

sca(axs[2])
plot(kr, abs.(Ehr))
xlabel(L"k_r")
ylabel(L"\int | \hat{E} | k_r \mathrm{d} k_{\theta}")
title("Radial energy spectrum")

xlim(0, nx/4)
axs[2].set_yscale("log")

tight_layout(w_pad=0.1)
show()


using StatsBase
using Test

@testset "StatsBase.Moments" begin
weight_funcs = (weights, aweights, fweights, pweights)

##### weighted var & std

x = [0.57, 0.10, 0.91, 0.72, 0.46, 0.0]
w = [3.84, 2.70, 8.29, 8.91, 9.71, 0.0]

@testset "Uncorrected with $f" for f in weight_funcs
    wv = f(w)
    m = mean(x, wv)

    # expected uncorrected output
    expected_var = sum(abs2.(x .- m), wv) / sum(wv)
    expected_std = sqrt.(expected_var)

    @testset "Variance" begin
        @test var(x, wv; corrected=false)           ≈ expected_var
        @test var(x, wv; mean=m, corrected=false)   ≈ expected_var
    end

    @testset "Standard Deviation" begin
        @test std(x, wv; corrected=false)           ≈ expected_std
        @test std(x, wv; mean=m, corrected=false)   ≈ expected_std
    end

    @testset "Mean and Variance" begin
        (m, v) = mean_and_var(x; corrected=false)
        @test m == mean(x)
        @test v == var(x; corrected=corrected=false)

        (m, v) = mean_and_var(x, wv; corrected=false)
        @test m == mean(x, wv)
        @test v == var(x, wv; corrected=false)
    end

    @testset "Mean and Standard Deviation" begin
        (m, s) = mean_and_std(x; corrected=false)
        @test m == mean(x)
        @test s == std(x; corrected=false)

        (m, s) = mean_and_std(x, wv; corrected=false)
        @test m == mean(x, wv)
        @test s == std(x, wv; corrected=false)
    end
end

# expected corrected output for (weights, aweights, fweights, pweights)
expected_var = [NaN, 0.0694434191182236, 0.05466601256158146, 0.06628969012045285]
expected_std = sqrt.(expected_var)

@testset "Corrected with $(weight_funcs[i])" for i in eachindex(weight_funcs)
    wv = weight_funcs[i](w)
    m = mean(x, wv)

    @testset "Variance" begin
        if isa(wv, Weights)
            @test_throws ArgumentError var(x, wv; corrected=true)
        else
            @test var(x, wv; corrected=true)           ≈ expected_var[i]
            @test var(x, wv; mean=m, corrected=true)   ≈ expected_var[i]
        end
    end

    @testset "Standard Deviation" begin
        if isa(wv, Weights)
            @test_throws ArgumentError std(x, wv; corrected=true)
        else
            @test std(x, wv; corrected=true)           ≈ expected_std[i]
            @test std(x, wv; mean=m, corrected=true)   ≈ expected_std[i]
        end
    end

    @testset "Mean and Variance" begin
        (m, v) = mean_and_var(x; corrected=true)
        @test m == mean(x)
        @test v == var(x; corrected=true)

        if isa(wv, Weights)
            @test_throws ArgumentError mean_and_var(x, wv; corrected=true)
        else
            (m, v) = mean_and_var(x, wv; corrected=true)
            @test m == mean(x, wv)
            @test v == var(x, wv; corrected=true)
        end
    end

    @testset "Mean and Standard Deviation" begin
        (m, s) = mean_and_std(x; corrected=true)
        @test m == mean(x)
        @test s == std(x; corrected=true)

        if isa(wv, Weights)
            @test_throws ArgumentError mean_and_std(x, wv; corrected=true)
        else
            (m, s) = mean_and_std(x, wv; corrected=true)
            @test m == mean(x, wv)
            @test s == std(x, wv; corrected=true)
        end
    end
end

x = rand(5, 6)
w1 = rand(5)
w2 = rand(6)

@testset "Uncorrected with $f" for f in weight_funcs
    wv1 = f(w1)
    wv2 = f(w2)
    m1 = mean(x, wv1, dims=1)
    m2 = mean(x, wv2, dims=2)

    expected_var1 = sum(abs2.(x .- m1) .* w1, dims = 1) ./ sum(wv1)
    expected_var2 = sum(abs2.(x .- m2) .* w2', dims = 2) ./ sum(wv2)
    expected_std1 = sqrt.(expected_var1)
    expected_std2 = sqrt.(expected_var2)

    @testset "Variance" begin
        @test var(x, wv1, 1; corrected=false) ≈ expected_var1
        @test var(x, wv2, 2; corrected=false) ≈ expected_var2
        @test var(x, wv1, 1; mean=m1, corrected=false) ≈ expected_var1
        @test var(x, wv2, 2; mean=m2, corrected=false) ≈ expected_var2
    end

    @testset "Standard Deviation" begin
        @test std(x, wv1, 1; corrected=false)          ≈ expected_std1
        @test std(x, wv2, 2; corrected=false)          ≈ expected_std2
        @test std(x, wv1, 1; mean=m1, corrected=false) ≈ expected_std1
        @test std(x, wv2, 2; mean=m2, corrected=false) ≈ expected_std2
    end

    @testset "Mean and Variance" begin
        for d in 1:2
            (m, v) = mean_and_var(x, d; corrected=false)
            @test m == mean(x, dims=d)
            @test v == var(x, dims=d, corrected=false)
        end

        (m, v) = mean_and_var(x, wv1, 1; corrected=false)
        @test m == mean(x, wv1, dims=1)
        @test v == var(x, wv1, 1; corrected=false)

        (m, v) = mean_and_var(x, wv2, 2; corrected=false)
        @test m == mean(x, wv2, dims=2)
        @test v == var(x, wv2, 2; corrected=false)
    end

    @testset "Mean and Standard Deviation" begin
        for d in 1:2
            (m, s) = mean_and_std(x, d; corrected=false)
            @test m == mean(x, dims=d)
            @test s == std(x, dims=d; corrected=false)
        end

        (m, s) = mean_and_std(x, wv1, 1; corrected=false)
        @test m == mean(x, wv1, dims=1)
        @test s == std(x, wv1, 1; corrected=false)

        (m, s) = mean_and_std(x, wv2, 2; corrected=false)
        @test m == mean(x, wv2, dims=2)
        @test s == std(x, wv2, 2; corrected=false)
    end
end

@testset "Corrected with $f" for f in weight_funcs
    wv1 = f(w1)
    wv2 = f(w2)
    m1 = mean(x, wv1, dims=1)
    m2 = mean(x, wv2, dims=2)

    if !isa(wv1, Weights)
        expected_var1 = sum(abs2.(x .- m1) .* w1, dims = 1) .* StatsBase.varcorrection(wv1, true)
        expected_var2 = sum(abs2.(x .- m2) .* w2', dims = 2) .* StatsBase.varcorrection(wv2, true)
        expected_std1 = sqrt.(expected_var1)
        expected_std2 = sqrt.(expected_var2)
    end

    @testset "Variance" begin
        if isa(wv1, Weights)
            @test_throws ArgumentError var(x, wv1, 1; corrected=true)
        else
            @test var(x, wv1, 1; corrected=true) ≈ expected_var1
            @test var(x, wv2, 2; corrected=true) ≈ expected_var2
            @test var(x, wv1, 1; mean=m1, corrected=true) ≈ expected_var1
            @test var(x, wv2, 2; mean=m2, corrected=true) ≈ expected_var2
        end
    end

    @testset "Standard Deviation" begin
        if isa(wv1, Weights)
            @test_throws ArgumentError std(x, wv1, 1; corrected=true)
        else
            @test std(x, wv1, 1; corrected=true)          ≈ expected_std1
            @test std(x, wv2, 2; corrected=true)          ≈ expected_std2
            @test std(x, wv1, 1; mean=m1, corrected=true) ≈ expected_std1
            @test std(x, wv2, 2; mean=m2, corrected=true) ≈ expected_std2
        end
    end

    @testset "Mean and Variance" begin
        for d in 1:2
            (m, v) = mean_and_var(x, d; corrected=true)
            @test m == mean(x, dims=d)
            @test v == var(x, dims=d, corrected=true)
        end

        if isa(wv1, Weights)
            @test_throws ArgumentError mean_and_var(x, wv1, 1; corrected=true)
        else
            (m, v) = mean_and_var(x, wv1, 1; corrected=true)
            @test m == mean(x, wv1, dims=1)
            @test v == var(x, wv1, 1; corrected=true)

            (m, v) = mean_and_var(x, wv2, 2; corrected=true)
            @test m == mean(x, wv2, dims=2)
            @test v == var(x, wv2, 2; corrected=true)
        end
    end

    @testset "Mean and Standard Deviation" begin
        for d in 1:2
            (m, s) = mean_and_std(x, d; corrected=true)
            @test m == mean(x, dims=d)
            @test s == std(x, dims=d, corrected=true)
        end

        if isa(wv1, Weights)
            @test_throws ArgumentError mean_and_std(x, wv1, 1; corrected=true)
        else
            (m, s) = mean_and_std(x, wv1, 1; corrected=true)
            @test m == mean(x, wv1, dims=1)
            @test s == std(x, wv1, 1; corrected=true)

            (m, s) = mean_and_std(x, wv2, 2; corrected=true)
            @test m == mean(x, wv2, dims=2)
            @test s == std(x, wv2, 2; corrected=true)
        end
    end
end

@testset "Skewness and Kurtosis with $f" for f in weight_funcs
    wv = f(ones(5) * 2.0)

    @test skewness(1:5)             ≈  0.0
    @test skewness([1, 2, 3, 4, 5]) ≈  0.0
    @test skewness([1, 2, 2, 2, 5]) ≈  1.1731251294063556
    @test skewness([1, 4, 4, 4, 5]) ≈ -1.1731251294063556

    @test skewness([1, 2, 2, 2, 5], wv) ≈ 1.1731251294063556

    @test kurtosis(1:5)             ≈ -1.3
    @test kurtosis([1, 2, 3, 4, 5]) ≈ -1.3
    @test kurtosis([1, 2, 3, 3, 2]) ≈ -1.1530612244897953

    @test kurtosis([1, 2, 3, 4, 5], wv) ≈ -1.3
end

@testset "General Moments with $f" for f in weight_funcs
    x = collect(2.0:8.0)
    @test moment(x, 2) ≈ sum((x .- 5).^2) / length(x)
    @test moment(x, 3) ≈ sum((x .- 5).^3) / length(x)
    @test moment(x, 4) ≈ sum((x .- 5).^4) / length(x)
    @test moment(x, 5) ≈ sum((x .- 5).^5) / length(x)

    @test moment(x, 2, 4.0) ≈ sum((x .- 4).^2) / length(x)
    @test moment(x, 3, 4.0) ≈ sum((x .- 4).^3) / length(x)
    @test moment(x, 4, 4.0) ≈ sum((x .- 4).^4) / length(x)
    @test moment(x, 5, 4.0) ≈ sum((x .- 4).^5) / length(x)

    w = f([1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0])
    x2 = collect(2.0:6.0)
    @test moment(x, 2, w) ≈ sum((x2 .- 4).^2) / 5
    @test moment(x, 3, w) ≈ sum((x2 .- 4).^3) / 5
    @test moment(x, 4, w) ≈ sum((x2 .- 4).^4) / 5
    @test moment(x, 5, w) ≈ sum((x2 .- 4).^5) / 5
end

end # @testset "StatsBase.Moments"


"""
Luby Transform Code

# References

- David J. C. MacKay. 2002. Information Theory, Inference & Learning Algorithms (pp. 589-596).
  Cambridge University Press, New York, NY, USA.
- https://franpapers.com/en/algorithmic/2018-introduction-to-fountain-codes-lt-codes-with-python/
"""
module LubyTransformCode

include("lehmer.jl")
include("soliton.jl")
include("core.jl")
include("encode.jl")
include("decode.jl")

end


immutable NaiveSolver{N, T} <: MathProgBase.AbstractMathProgSolver
    rate::T
    max_step::T
    precondition_divisors::SVector{N, T}
    iteration_limit::Int
    gradient_convergence_tolerance::T
end

NaiveSolver(num_vars;
    rate=1.0,
    max_step=1.0,
    precondition_divisors::AbstractVector=ones(SVector{num_vars, Float64}),
    iteration_limit=100,
    gradient_convergence_tolerance=1e-3) =
NaiveSolver{num_vars, Float64}(rate,
    max_step,
    convert(SVector{num_vars, Float64}, precondition_divisors),
    iteration_limit,
    gradient_convergence_tolerance)

type NaiveModel{N, T} <: MathProgBase.AbstractNonlinearModel
    solver::NaiveSolver{N, T}
    x::MVector{N, T}
    x_lb::SVector{N, T}
    x_ub::SVector{N, T}
    g::MVector{N, T}
    objval::T
    problem::MathProgBase.AbstractNLPEvaluator
    status::Symbol

    NaiveModel(solver::NaiveSolver{N, T}) = new(solver)
end

NaiveModel{N, T}(solver::NaiveSolver{N, T}) = NaiveModel{N, T}(solver)

MathProgBase.NonlinearModel{N, T}(solver::NaiveSolver{N, T}) = NaiveModel{N, T}(solver)

function MathProgBase.loadproblem!{N, T}(model::NaiveModel{N, T},
    num_vars::Integer, num_constraints::Integer,
    x_lb::AbstractVector, x_ub::AbstractVector,
    g_lb::AbstractVector, g_ub::AbstractVector,
    sense::Symbol,
    problem::MathProgBase.AbstractNLPEvaluator)

    @assert sense == :Min
    @assert num_vars == N
    @assert num_constraints == 0

    model.x_lb = convert(SVector{N, T}, x_lb)
    model.x_ub = convert(SVector{N, T}, x_ub)
    model.x = zeros(MVector{N, T})
    model.g = zeros(MVector{N, T})
    model.objval = zero(T)
    model.problem = problem
end

MathProgBase.setwarmstart!(model::NaiveModel, x) = copy!(model.x, x)

function MathProgBase.optimize!{N, T}(model::NaiveModel{N, T})
    model.status = :UserLimit
    conditioned_max_step = model.solver.max_step .* model.solver.precondition_divisors
    conditioned_rate = model.solver.rate .* model.solver.precondition_divisors
    for i in 1:model.solver.iteration_limit
        MathProgBase.eval_grad_f(model.problem, model.g, model.x)

        if maximum(abs.(model.g)) < model.solver.gradient_convergence_tolerance
            model.status = :Optimal
            break
        end

        step = max.(min.(conditioned_max_step, conditioned_rate .* model.g), -conditioned_max_step)
        model.x .= max.(min.(model.x .- step, model.x_ub), model.x_lb)

    end
end

MathProgBase.status(model::NaiveModel) = model.status
MathProgBase.getsolution(model::NaiveModel) = model.x


using Test
using QuantLib
using Dates

euribor3m = euribor_index(QuantLib.Time.TenorPeriod(Dates.Month(3)))

todaysDate = Date(2006, 5, 23)
set_eval_date!(settings, todaysDate)
calendar = euribor3m.fixingCalendar
fixingDays = euribor3m.fixingDays
settlementDate = QuantLib.Time.advance(Dates.Day(fixingDays), calendar, todaysDate)

threeMonthFraQuote = zeros(10)
threeMonthFraQuote[2] = 0.030
threeMonthFraQuote[3] = 0.031
threeMonthFraQuote[4] = 0.032
threeMonthFraQuote[7] = 0.033
threeMonthFraQuote[10] = 0.034

## QUOTES ##
h1x4 = Quote(threeMonthFraQuote[2])
h2x5 = Quote(threeMonthFraQuote[3])
h3x6 = Quote(threeMonthFraQuote[4])
h6x9 = Quote(threeMonthFraQuote[7])
h9x12 = Quote(threeMonthFraQuote[10])

## RATE HELPERS ##
fraDayCount = euribor3m.dc
convention = euribor3m.convention
endOfMonth = euribor3m.endOfMonth

fra1x4 = FraRateHelper(h1x4, 1, 4, fixingDays, calendar, convention, endOfMonth, fraDayCount)
fra2x5 = FraRateHelper(h2x5, 2, 5, fixingDays, calendar, convention, endOfMonth, fraDayCount)
fra3x6 = FraRateHelper(h3x6, 3, 6, fixingDays, calendar, convention, endOfMonth, fraDayCount)
fra6x9 = FraRateHelper(h6x9, 6, 9, fixingDays, calendar, convention, endOfMonth, fraDayCount)
fra9x12 = FraRateHelper(h9x12, 9, 12, fixingDays, calendar, convention, endOfMonth, fraDayCount)

fraInstruments = FraRateHelper[fra1x4, fra2x5, fra3x6, fra6x9, fra9x12]
termStructureDC = QuantLib.Time.ISDAActualActual()
tol = 1.0e-15

fraTermStructure = PiecewiseYieldCurve(settlementDate, fraInstruments, termStructureDC, QuantLib.Math.LogLinear(), Discount(), tol, IterativeBootstrap())

euribor3m = update_termstructure(euribor3m, fraTermStructure)

fraCalendar = euribor3m.fixingCalendar
fraBusinessDayConvention = euribor3m.convention
fraFwdType = LongPosition()
fraNotional = 100.0
monthsToStart = [1, 2, 3, 6, 9]
fraTermMonths = 3
fraValueDate = QuantLib.Time.advance(Dates.Month(monthsToStart[1]), fraCalendar, settlementDate, fraBusinessDayConvention)
fraMaturityDate = QuantLib.Time.advance(Dates.Month(fraTermMonths), fraCalendar, fraValueDate, fraBusinessDayConvention)
fraStrikeRate = threeMonthFraQuote[monthsToStart[1] + 1]
myFRA = ForwardRateAgreement(fraValueDate, fraMaturityDate, fraFwdType, fraStrikeRate, fraNotional, euribor3m, fraTermStructure)

# test calculations
@test forward_rate(myFRA).rate == 0.030000000000000172
@test spot_value(myFRA) == 99.73470289713767
@test forward_value(myFRA) == 100.76666666666667
@test implied_yield(myFRA, spot_value(myFRA), forward_value(myFRA), settlementDate, SimpleCompounding(), fraDayCount).rate == 0.030039933543589872

# other tests
@test maturity_date(fra1x4) == Date(2006, 9, 26)
@test QuantLib.implied_quote(fraTermStructure.instruments[1]) == 0.030000000000000172


using Distributions, Random
using Test, LinearAlgebra, PDMats

#  dimensions

n = 3
p = 6

#  parameters

M    = randn(n, p)
U    = rand(InverseWishart(n + 2, Matrix(1.0I, n, n)))
V    = rand(InverseWishart(p + 2, Matrix(1.0I, p, p)))
UF32 = Matrix{Float32}(U)
VBF  = Matrix{BigFloat}(V)
PDU  = PDMat(U)
PDV  = PDMat(V)
m    = randn(n)
w    = randn(p)
μ    = randn()
σ    = 2.0rand()
σ²   = σ ^ 2

#  distribution instances

D = MatrixNormal(M, U, V)  #  n x p
G = MatrixNormal(p, n)     #  p x n
L = MatrixNormal(reshape(m, n, 1), U, reshape([σ²], 1, 1))  #  n x 1
H = MatrixNormal(reshape(w, 1, p), reshape([σ²], 1, 1), V)  #  1 x p
K = MatrixNormal(reshape([μ], 1, 1), reshape([σ], 1, 1), reshape([σ], 1, 1))  #  1 x 1

d = vec(D) # MvNormal(vec(M), V ⊗ U)
g = MvNormal( Matrix(1.0I, p*n, p*n) )
l = MvNormal(m, σ² * U)
h = MvNormal(w, σ² * V)
k = Normal(μ, σ)

#  draws

A = rand(D)
B = rand(G)
X = rand(L)
Y = rand(H)
Z = rand(K)

a = vec(A)
b = vec(B)
x = vec(X)
y = vec(Y)
z = Z[1, 1]

@testset "Check all MatrixNormal constructors" begin
    @test MatrixNormal(M, U,        PDV.chol) isa MatrixNormal
    @test MatrixNormal(M, U,        PDV)      isa MatrixNormal
    @test MatrixNormal(M, PDU.chol, V)        isa MatrixNormal
    @test MatrixNormal(M, PDU.chol, PDV.chol) isa MatrixNormal
    @test MatrixNormal(M, PDU.chol, PDV)      isa MatrixNormal
    @test MatrixNormal(M, PDU,      V)        isa MatrixNormal
    @test MatrixNormal(M, PDU,      PDV.chol) isa MatrixNormal
    @test MatrixNormal(M, PDU,      PDV)      isa MatrixNormal
end

@testset "MatrixNormal promotion during construction" begin
    R = MatrixNormal(M, UF32, VBF)
    @test partype(R) == BigFloat
end

@testset "MatrixNormal construction errors" begin
    @test_throws ArgumentError MatrixNormal(M, V, U)
    @test_throws ArgumentError MatrixNormal(M, U, U)
end

@testset "MatrixNormal params" begin
    MM, PDUU, PDVV = params(D)
    @test MM == M
    @test U == PDUU.mat
    @test V == PDVV.mat
end

@testset "MatrixNormal size" begin
    @test size(D) == (n, p)
    @test size(G) == (p, n)
    @test size(L) == (n, 1)
    @test size(H) == (1, p)
    @test size(K) == (1, 1)
end

@testset "MatrixNormal rank" begin
    @test rank(D) == min(n, p)
    @test rank(G) == min(n, p)
    @test rank(L) == 1
    @test rank(H) == 1
    @test rank(K) == 1

    @test rank(D) == rank(rand(D))
    @test rank(G) == rank(rand(G))
    @test rank(L) == rank(rand(L))
    @test rank(H) == rank(rand(H))
    @test rank(K) == rank(rand(K))
end

@testset "MatrixNormal insupport" begin
    @test insupport(D, rand(D))
    @test insupport(G, rand(G))
    @test insupport(L, rand(L))
    @test insupport(H, rand(H))
    @test insupport(K, rand(K))

    @test !insupport(D, rand(G))
    @test !insupport(D, M + M * im)
    @test !insupport(L, rand(H))
    @test !insupport(H, rand(L))
end

@testset "MatrixNormal mean" begin
    @test mean(D) == M
    @test mean(G) == zeros(p, n)
    @test mean(L) == reshape(m, n, 1)
    @test mean(H) == reshape(w, 1, p)
    @test mean(K) == reshape([μ], 1, 1)
end

@testset "MatrixNormal mode" begin
    @test mode(D) == M
    @test mode(G) == zeros(p, n)
    @test mode(L) == reshape(m, n, 1)
    @test mode(H) == reshape(w, 1, p)
    @test mode(K) == reshape([μ], 1, 1)
end

@testset "MatrixNormal logpdf" begin
    @test logpdf(D, A) ≈ logpdf(d, a)
    @test logpdf(G, B) ≈ logpdf(g, b)
    @test logpdf(L, X) ≈ logpdf(l, x)
    @test logpdf(H, Y) ≈ logpdf(h, y)
    @test logpdf(K, Z) ≈ logpdf(k, z)
end

@testset "MatrixNormal sample moments" begin
    @test isapprox(mean(rand(D, 100000)), mean(D) , atol = 0.1)
    @test isapprox(cov(hcat(vec.(rand(D, 100000))...)'), kron(V, U) , atol = 0.1)
end

@testset "MatrixNormal conversion" for elty in (Float32, Float64, BigFloat)
    Del1 = convert(MatrixNormal{elty}, D)
    Del2 = convert(MatrixNormal{elty}, M, PDU, PDV, D.logc0)

    @test partype(Del1) == elty
    @test partype(Del2) == elty
end


clear()
@testset "check pooling op's gradient at single dim" begin
    for d in [1 2 3]
        for pool in [maximum minimum sum mean linearpool exppool]
            @testset "check $pool op's gradient at dim = $d" begin
                DIMS = d
                TYPE = Array{Float64};
                # [1] prepare input data and its label
                inputdims = 64;
                timeSteps = 16;
                batchsize = 32;
                x = Variable(rand(inputdims, timeSteps, batchsize); type=TYPE, keepsgrad=true,);
                if d==1;l = Variable(rand(1, timeSteps, batchsize); type=TYPE);end
                if d==2;l = Variable(rand(inputdims, 1, batchsize); type=TYPE);end
                if d==3;l = Variable(rand(inputdims, timeSteps, 1); type=TYPE);end

                # [2] forward and backward propagation
                COST1 = cost(mseLoss(pool(x; dims=DIMS), l));
                backward();

                # [3] with a samll change of a weight
                GRAD = x.delta[1];
                DELTA = 1e-6;
                x.value[1] += DELTA;

                # [4] forward and backward propagation
                COST2 = cost(mseLoss(pool(x; dims=DIMS), l));
                backward();

                # [5] check if the auto-grad is true or not
                dLdW = (COST2 - COST1)/DELTA;   # numerical gradient
                err  = abs((dLdW-GRAD)/(GRAD+eps(Float32)))*100;  # relative error in %
                err  = err < 1e-1 ? 0.0 : err;
                @test err<1.0
            end
        end
    end
end

@testset "check pooling op's gradient at mutiple dims" begin
    for pool in [maximum minimum sum mean linearpool exppool]
        @testset "check $pool op's gradient" begin
            DIMS = (1,2)
            TYPE = Array{Float64};

            # [1] prepare input data and its label
            inputdims = 64;
            timeSteps = 16;
            batchsize = 32;
            x = Variable(rand(inputdims, timeSteps, batchsize); type=TYPE,keepsgrad=true);
            l = Variable(rand(1,         1,         batchsize); type=TYPE);

            # [2] forward and backward propagation
            COST1 = cost(mseLoss(pool(x; dims=DIMS), l));
            backward();

            # [3] with a samll change of a weight
            GRAD = x.delta[1];
            DELTA = 1e-6;
            x.value[1] += DELTA;

            # [4] forward and backward propagation with a samll change of a weight
            COST2 = cost(mseLoss(pool(x; dims=DIMS), l));
            backward();

            # [5] check if the auto-grad is true or not
            dLdW = (COST2 - COST1)/DELTA;   # numerical gradient
            err  = abs((dLdW-GRAD)/(GRAD+eps(Float64)))*100;  # relative error in %
            err  = err < 1e-1 ? 0.0 : err;
            @test err<1.0
        end
    end
end

@testset "check maxmin and minmax op's gradient at mutiple dims" begin
    for pool in [maxmin minmax]
        @testset "check $pool op's gradient" begin
            DIM1 = 1
            DIM2 = 2
            TYPE = Array{Float64};

            # [1] prepare input data and its label
            inputdims = 64;
            timeSteps = 16;
            batchsize = 32;
            x = Variable(rand(inputdims, timeSteps, batchsize); type=TYPE,keepsgrad=true);
            l = Variable(rand(1,         1,         batchsize); type=TYPE);

            # [2] forward and backward propagation
            COST1 = mseLoss(pool(x; dims1=DIM1, dims2=DIM2), l) |> cost
            backward();

            # [3] with a samll change of a weight
            GRAD = x.delta[1];
            DELTA = 1e-6;
            x.value[1] += DELTA;

            # [4] forward and backward propagation with a samll change of a weight
            COST2 = mseLoss(pool(x; dims1=DIM1, dims2=DIM2), l) |> cost
            backward();

            # [5] check if the auto-grad is true or not
            dLdW = (COST2 - COST1)/DELTA;   # numerical gradient
            err  = abs((dLdW-GRAD)/(GRAD+eps(Float64)))*100;  # relative error in %
            err  = err < 1e-1 ? 0.0 : err;
            @test err<1.0
        end
    end
end


#############################################################################
# extreme_point.jl
# Compute the of the base polyhedron for an given ordering, i.e., the elements
# of the result is calculated in the order of the given ordering.
#############################################################################

export extreme_point

function extreme_point(F::Function, w::AbstractArray,
                       V = collect(1:length(w))::AbstractArray)
  @assert F([])[1] == 0 # f should be 0 at the empty set
  w = Int.(w)
  n = length(w)
  @assert length(V) == n
  x = zeros(n)
  for ii = 1:n
    x[V[w[ii]]] = F(V[w[1: ii]])[1] - F(V[w[1: ii-1]])[1]
  end
  return x
end

extreme_point(F::SubmodFunc, w::AbstractArray) = extreme_point(x -> evaluate(F, x), w)


using LaTeXStrings
using Plots
using Plots.PlotMeasures
pyplot()

wp = JvWorker(grid_size=25)
v_init = wp.x_grid .* 0.5

f(x) = bellman_operator(wp, x)
V = compute_fixed_point(f, v_init, max_iter=300)

s_policy, ϕ_policy = bellman_operator(wp, V, ret_policies=true)

# === plot solution === #
p = plot(wp.x_grid, [ϕ_policy s_policy V],
         title=["ϕ policy" "s policy" "value function"],
         color=[:orange :blue :green],
         xaxis=("x", (0.0, maximum(wp.x_grid))),
         yaxis=((-0.1, 1.1)), size=(800, 800),
         legend=false, layout=(3, 1),
         bottom_margin=Measures.Length(:mm, 20))


function compute_frequencies(Z::Matrix, W::Vector, q)
    N, M = size(Z)
    length(W) == M || error("incompatible length of W")
    Pij = zeros(q-1, q-1, N, N)
    for i in 1:N
        for j in i:N
            for s in 1:M
                coli = Z[i,s]
                colj = Z[j,s]
                (coli == q || colj == q) && continue
                Pij[colj,coli,j,i] += W[s]
                Pij[coli,colj,i,j] = Pij[colj,coli,j,i]
            end
        end
    end
    Pi = zeros(q-1,N)
    for i in 1:N
        for s in 1:M
            coli = Z[i,s]
            (coli == q) && continue
            Pi[coli,i] += W[s]
        end
    end
    return Pij,Pi
end

compute_frequencies(Z,W) = compute_frequencies(Z,W,maximum(Z))
compute_frequencies(Z) = compute_frequencies(Z,ones(size(Z,2))/size(Z,2),maximum(Z))

function add_pseudocount(Pi_true::Array{Float64,2}, Pij_true::Array{Float64,4},pc)
    0 <= pc <= 1 || error("pseudocount $pc should be ∈ [0,1]")
    s,N = size(Pi_true)
    q = s+1
    Pipc  = (1.0-pc)*Pi_true  .+ pc/q
    Pijpc = (1.0-pc)*Pij_true .+ pc/q^2

    for i in 1:N
        Pijpc[:,:,i,i] .= (1-pc) * Pij_true[:,:,i,i]
        for a in 1:s
            Pijpc[a,a,i,i] = Pipc[a,i]
        end
    end
    Pijpc,Pipc
end

function Cij(Pi,Pij)
    q,N = size(Pi)
    cij = Matrix{Float64}(undef, N*q,N*q)
    mi = Vector{Float64}(undef,N*q)
    for i in 1:N
        rowblk = ((i-1)*q+1):i*q
        vpi = @view Pi[:,i]
        mi[rowblk].= vpi
        for j in 1:N
            vpj = @view Pi[:,j]
            vpij = @view Pij[:,:,i,j]
            colblk = ((j-1)*q+1):j*q
            cij[rowblk,colblk] .= vpij - vpi*vpj'
        end
    end
    return cij,mi
end

"""
    covariance_matrix(Z, W; pc::Real=0)
Compute the covariance matrix from numerical alignment `Z` and weights `W`. The output is a `N(q-1) × N(q-1)`
(skipping color `q`) symmetric matrix.
## Keywords arguments:
*  `pc` in [0,1]: pseudocount [default = `0`]
end
"""
function covariance_matrix(Z, W; pc::Real=0)
    0<=pc<=1 || error("pc = $pc should be ∈ [0,1]")
    Pij,Pi = compute_frequencies(Z,W)
    if pc > 0
        Pijpc,Pipc=add_pseudocount(Pi,Pij,pc)
        return Cij(Pipc,Pijpc)

    end
    return Cij(Pi,Pij)
end

covariance_matrix(Z; kwds...) = covariance_matrix(Z,ones(size(Z,2))/size(Z,2); kwds...)


## Julia 1.1.0

# Initialize
using Plots
plotlyjs() # Backend PlotlyJS

# Archimedean Spiral
# r = a + bΘ

a = 1.0
b = 2.0
angles = 8π
Θ = range(0, stop = angles, length = 1000)
r = a .+ b .* Θ

plot(Θ, r, proj=:polar)
savefig("./png/archimedean-spiral.png")

steps = 100
anim = @animate for i = 1:steps
    Θ = range(0, stop = angles * i / steps, length = 10 * i)
    r = a .+ b .* Θ
    plot(Θ, r, proj=:polar)
end

gif(anim, "./gif/archimedean-spiral.gif", fps = 30)


using Manopt, Manifolds, ManifoldsBase, Test
using LinearAlgebra: Diagonal, dot, eigvals, eigvecs

@testset "Conjugate Gradient coefficient rules" begin
    M = Euclidean(2)
    F(M, x) = norm(x)^2
    gradF(::Euclidean, x) = 2 * x
    P = GradientProblem(M, F, gradF)
    x0 = [0.0, 1.0]
    sC = StopAfterIteration(1)
    s = ConstantStepsize(1.0)
    retr = ExponentialRetraction()
    vtm = ParallelTransport()

    grad_1 = [1.0, 1.0]
    δ1 = [0.0, 2.0]
    grad_2 = [1.0, 1.5]
    δ2 = [0.5, 2.0]
    diff = grad_2 - grad_1

    dU = SteepestDirectionUpdateRule()
    O = ConjugateGradientDescentOptions(M, x0, sC, s, dU, retr, vtm, zero_vector(M, x0))
    @test O.coefficient(P, O, 1) == 0

    dU = ConjugateDescentCoefficient()
    O = ConjugateGradientDescentOptions(M, x0, sC, s, dU, retr, vtm, zero_vector(M, x0))
    O.gradient = grad_1
    O.δ = δ1
    # for the first case we get zero
    @test O.coefficient(P, O, 1) == 0.0
    O.gradient = grad_2
    O.δ = δ2
    @test O.coefficient(P, O, 2) == dot(grad_2, grad_2) / dot(-δ2, grad_1)

    dU = DaiYuanCoefficient()
    O = ConjugateGradientDescentOptions(M, x0, sC, s, dU, retr, vtm)
    O.gradient = grad_1
    O.δ = δ1
    # for the first case we get zero
    @test O.coefficient(P, O, 1) == 0.0
    O.gradient = grad_2
    O.δ = δ2
    @test O.coefficient(P, O, 2) == dot(grad_2, grad_2) / dot(δ2, grad_2 - grad_1)

    dU = FletcherReevesCoefficient()
    O = ConjugateGradientDescentOptions(M, x0, sC, s, dU, retr, vtm)
    O.gradient = grad_1
    O.δ = δ1
    # for the first case we get zero
    @test O.coefficient(P, O, 1) == 1.0
    O.gradient = grad_2
    O.δ = δ2
    @test O.coefficient(P, O, 2) == dot(grad_2, grad_2) / dot(grad_1, grad_1)

    dU = HagerZhangCoefficient()
    O = ConjugateGradientDescentOptions(M, x0, sC, s, dU, retr, vtm)
    O.gradient = grad_1
    O.δ = δ1
    # for the first case we get zero
    @test O.coefficient(P, O, 1) == 0.0
    O.gradient = grad_2
    O.δ = δ2
    denom = dot(δ1, diff)
    ndiffsq = dot(diff, diff)
    @test O.coefficient(P, O, 2) ==
        dot(diff, grad_2) / denom - 2 * ndiffsq * dot(δ1, grad_2) / denom^2

    dU = HeestenesStiefelCoefficient()
    O = ConjugateGradientDescentOptions(M, x0, sC, s, dU, retr, vtm)
    O.gradient = grad_1
    O.δ = δ1
    @test O.coefficient(P, O, 1) == 0.0
    O.gradient = grad_2
    O.δ = δ2
    @test O.coefficient(P, O, 2) == dot(diff, grad_2) / dot(δ1, diff)

    dU = LiuStoreyCoefficient()
    O = ConjugateGradientDescentOptions(M, x0, sC, s, dU, retr, vtm)
    O.gradient = grad_1
    O.δ = δ1
    @test O.coefficient(P, O, 1) == 0.0
    O.gradient = grad_2
    O.δ = δ2
    @test O.coefficient(P, O, 2) == -dot(grad_2, diff) / dot(δ1, grad_1)

    dU = PolakRibiereCoefficient()
    O = ConjugateGradientDescentOptions(M, x0, sC, s, dU, retr, vtm)
    O.gradient = grad_1
    O.δ = δ1
    @test O.coefficient(P, O, 1) == 0.0
    O.gradient = grad_2
    O.δ = δ2
    @test O.coefficient(P, O, 2) == dot(grad_2, diff) / dot(grad_1, grad_1)
end
@testset "Conjugate Gradient runs – Low Rank matrix approx" begin
    A = Diagonal([2.0, 1.1, 1.0])
    M = Sphere(size(A, 1) - 1)
    F(::Sphere, x) = x' * A * x
    gradF(M, x) = project(M, x, 2 * A * x) # project the Euclidean gradient

    x0 = [2.0, 0.0, 2.0] / sqrt(8.0)
    x_opt = conjugate_gradient_descent(
        M,
        F,
        gradF,
        x0;
        stepsize=ArmijoLinesearch(),
        coefficient=FletcherReevesCoefficient(),
        stopping_criterion=StopAfterIteration(15),
    )
    @test isapprox(F(M, x_opt), minimum(eigvals(A)); atol=2.0 * 1e-2)
    @test isapprox(x_opt, eigvecs(A)[:, size(A, 1)]; atol=3.0 * 1e-1)
    x_opt2 = conjugate_gradient_descent(
        M,
        F,
        gradF,
        x0;
        stepsize=ArmijoLinesearch(),
        coefficient=FletcherReevesCoefficient(),
        stopping_criterion=StopAfterIteration(15),
        return_options=true,
    )
    @test get_solver_result(x_opt2) == x_opt
end


using Base.LinAlg
import Base.LinAlg.givensAlgorithm
# Below code stolen from newer Base

"""
    lowrankupdate!(C::Cholesky, v::StridedVector) -> CC::Cholesky

Update a Cholesky factorization `C` with the vector `v`. If `A = C[:U]'C[:U]` then `CC = cholfact(C[:U]'C[:U] + v*v')` but the computation of `CC` only uses `O(n^2)` operations. The input factorization `C` is updated in place such that on exit `C == CC`. The vector `v` is destroyed during the computation.
"""
function lowrankupdate!(C::Cholesky, v::StridedVector)
    A = C.factors
    n = length(v)
    if size(C, 1) != n
        throw(DimensionMismatch("updating vector must fit size of factorization"))
    end
    if C.uplo == 'U'
        conj!(v)
    end

    for i = 1:n

        # Compute Givens rotation
        c, s, r = givensAlgorithm(A[i,i], v[i])

        # Store new diagonal element
        A[i,i] = r

        # Update remaining elements in row/column
        if C.uplo == 'U'
            for j = i + 1:n
                Aij = A[i,j]
                vj  = v[j]
                A[i,j]  =   c*Aij + s*vj
                v[j]    = -s'*Aij + c*vj
            end
        else
            for j = i + 1:n
                Aji = A[j,i]
                vj  = v[j]
                A[j,i]  =   c*Aji + s*vj
                v[j]    = -s'*Aji + c*vj
            end
        end
    end
    return C
end

"""
    lowrankdowndate!(C::Cholesky, v::StridedVector) -> CC::Cholesky

Downdate a Cholesky factorization `C` with the vector `v`. If `A = C[:U]'C[:U]` then `CC = cholfact(C[:U]'C[:U] - v*v')` but the computation of `CC` only uses `O(n^2)` operations. The input factorization `C` is updated in place such that on exit `C == CC`. The vector `v` is destroyed during the computation.
"""
function lowrankdowndate!(C::Cholesky, v::StridedVector)
    A = C.factors
    n = length(v)
    if size(C, 1) != n
        throw(DimensionMismatch("updating vector must fit size of factorization"))
    end
    if C.uplo == 'U'
        conj!(v)
    end

    for i = 1:n

        Aii = A[i,i]

        # Compute Givens rotation
        s = conj(v[i]/Aii)
        s2 = abs2(s)
        if s2 > 1
            throw(LinAlg.PosDefException(i))
        end
        c = sqrt(1 - abs2(s))

        # Store new diagonal element
        A[i,i] = c*Aii

        # Update remaining elements in row/column
        if C.uplo == 'U'
            for j = i + 1:n
                vj = v[j]
                Aij = (A[i,j] - s*vj)/c
                A[i,j] = Aij
                v[j] = -s'*Aij + c*vj
            end
        else
            for j = i + 1:n
                vj = v[j]
                Aji = (A[j,i] - s*vj)/c
                A[j,i] = Aji
                v[j] = -s'*Aji + c*vj
            end
        end
    end
    return C
end


"""
    Standard()
Standard Scaler. This scaler scale data as:
``\\tilde{\\boldsymbol{x}} = \\frac{x_{i}-\\mu}{\\sigma}``

# Example
```jldoctest preprocessing
julia> x = [
    16.862463771320925 68.10823385851712
    15.382965696961577 65.4313485700859
    8.916228406218375 53.92034559524475
    10.560285659132695 59.17305391117168
    12.142253214135884 62.28708207525656
    5.362107221163482 43.604947901567414
    13.893239446341777 62.44348617377496
    11.871357065173395 60.28433066289655
    29.83792267802442 69.22281924803998
    21.327107214235483 70.15810991597944
    23.852372696012498 69.81780163668844
    26.269031430914108 67.61037566099782
    22.78907104644012 67.78105545358633
    26.73342178134947 68.59263965946904
    9.107259141706415 56.565383817343495
    29.38551885863976 68.1005579469209
    7.935966787763017 53.76264777936664
    29.01677894379809 68.69484161138638
    6.839609488194577 49.69794758177567
    13.95215840314148 62.058116579899085]; #These data are also used to explanations of other functions.

julia> t = [169.80980778351542, 167.9081124078835, 152.30845618985222, 160.3110300206261, 161.96826472170756, 136.02842285615077, 163.98131131382686, 160.117817321485, 172.22758529098235, 172.21342437006865, 171.8939175591617, 169.83018083884602, 171.3878062674257, 170.52487535026015, 156.40282783981309, 170.6488327896672, 151.69267899906185, 172.32478221316322, 145.14365314788827, 163.79383292080666];

julia> scaler = Standard()
Standard(Float64[])

julia> fit!(scaler, x)
2×2 Matrix{Float64}:
 19.0591   64.467
  6.95818   6.68467

julia> transform(scaler, x)
20×2 Matrix{Float64}:
  0.0310374   0.44579
  0.0104337   0.218714
  1.01139     0.710027
  1.37285     0.954091
 -0.895893   -0.270589
  0.983361    0.47397
  1.39855     0.645624
 -1.31861    -0.901482
  0.0147702   0.241708
  0.29675     0.483076
 -0.338824   -0.104812
  0.432442    0.387922
  0.860418    0.567095
 -0.495306    0.140871
  0.963084    0.552767
 -1.38901    -1.05926
  0.469323    0.719196
  0.0669475   0.512023
 -1.47583    -1.44017
 -1.99789    -3.27656
```
"""
mutable struct Standard
    p::AbstractVecOrMat
    Standard() = new(Array{Float64}(undef, 0))
end

"""
    fit!(scaler, x; dims=1)
fit the scaler with `x`. `dims` is the dimension of the number of data.
"""
function fit!(scaler::Standard, x; dims=1)
    scaler.p = vcat(mean(x, dims=dims), std(x, dims=dims))
end

ss(x, m, s) = @. (x-m)/s

"""
    transform(scaler, x; dims=1)
transform data with scaler. `dims` is the dimension of the number of data.
"""
function transform(scaler::Standard, x; dims=1)
    p = scaler.p
    check_size(x, p, dims)
    y = similar(x, Float32)
    if dims == 1
        for i in 1 : size(x, 2)
            y[:, i] = ss(x[:, i], p[:, i]...)
        end
    elseif dims == 2
        for i in 1 : size(x, 1)
            y[i, :] = ss(x[i, :], p[:, i]...)
        end
    end
    return y
end

"""
    fit_transform!(scaler, x; dims=1)
fit scaler with `x`, and transform `x`. `dims` is the dimension of the number of data.
"""
function fit_transform!(scaler::Standard, x; dims=1)
    fit!(scaler, x, dims=dims)
    transform(scaler, x, dims=dims)
end

iss(x, m, s) = @. x*s+m

"""
    inv_transform(scaler, x; dims=1)
Convert `x` in reverse. `dims` is the dimension of the number of data.
"""
function inv_transform(scaler::Standard, x; dims=1)
    p = scaler.p
    check_size(x, p, dims)
    y = similar(x, Float32)
    if dims == 1
        for i in 1 : size(x, 2)
            y[:, i] = iss(x[:, i], p[:, i]...)
        end
    elseif dims == 2
        for i in 1 : size(x, 1)
            y[i, :] = iss(x[i, :], p[:, i]...)
        end
    end
    return y
end


"Global Lax Friedrichs Flux"
function glf_num_flux(ul, ur, f, α)
    return 0.5*(f(ul)+f(ur))-α/2*(ur-ul)
end

"Flux for advection equation (simple upwinding)"
advection_num_flux(ul, ur, f, α) = ul

"""
Returns the Rusanov interface flux for the Euler equations
V. V. Rusanov, Calculation of Interaction of Non-Steady Shock Waves with
Obstacles, J. Comput. Math. Phys. USSR, 1, pp. 267-279, 1961.
"""
function rusanov_euler_num_flux(ul, ur, f, α)
  @assert(size(ul,1) == size(ur,1),
  "ul and ur vector have different number of components")
  F = zeros(ul)

  # Fisical variables and sound speeds
  ρL = ul[1];vL=ul[2]/ρL;EL = ul[3]
  pL = (γ - 1)*(EL - 0.5*ρL*vL*vL)
  aL = sqrt(γ*max(pL/ρL,1.0))

  ρR = ur[1];vR=ur[2]/ρR;ER = ur[3]
  pR = (γ - 1)*(ER - 0.5*ρR*vR*vR)
  aR = sqrt(γ*max(pR/ρR,1.0))

  # Find the maximum eigenvalue for each interface
  maxρ = max(abs(vL)+aL, abs(vR)+aR)

  # f₁ = ρ*u
  F[1] = 0.5*(ρL*vL + ρR*vR - maxρ*(ρR - ρL))
  # f₂ = rho*u*u+p
  F[2] = 0.5*(ρL*vL*vL + pL + ρR*vR*vR +
                   pR - maxρ*(ρR*vR - ρL*vL))
  # f₃ = (E+p)*u
  F[3] = 0.5*((EL + pL)*vL + (ER + pR)*vR - maxρ*(ER - EL))
  return F
end


cd("INSERT_PATH")

using Distributions
using PolynomialRoots

include("ToolBox.jl")
include("predictionGSBR.jl")

# define map parameters
theta = [0.05, 2.55, 0., -0.99]
x0 = 1.
dataSeed = 2
n = 114

# noise mixture: f_{2,l}
w1 = 0.6
lam1 = 1e05
lam2 = 1e-02 * lam1

x = genData(n,theta,x0,w1,lam1,lam2,dataSeed)
# Simulation parameters, T: prediction horizon [T=0 --> Plain GSBReconstruction]
degree,maxiter,burnin,seed,T = 5,500000,250000,3,0

@time thetas, zp, sx0 = predictionGSBR(x,degree,maxiter,burnin,seed,T);


using Distributions, ForwardDiff
using KernelGoodnessOfFit

# setup the problem
rbf_ = GaussianRBF(1.0)

μ = 0.0
σ² = 1.0
p_ = MultivariateNormal([μ, 0], [σ² 0; 0 σ²])  # true
q_ = MultivariateNormal([μ, 5.0], [σ² 0; 0 σ²])  # false

vs_ = [0.0 ; 0.0]
vs_ = reshape(vs_, 2, 1)

n_ = 100
xs_ = rand(p_, n_)

# simple example
T = KernelGoodnessOfFit.KSDTest(rbf_, 10, 0.05)

@test perform(T, q_, xs_).result == :reject
@test perform(T, p_, xs_).result == :accept


#https://en.wikipedia.org/wiki/Rotation_matrix#In_three_dimensions

function euler_rotation_matrix(θ)
	a = θ[1]
	b = θ[2]
	c = θ[3]

	[cos(a) * cos(b) (cos(a) * sin(b) * sin(c) - sin(a) * cos(c)) (cos(a) * sin(b) * cos(c) + sin(a) * sin(c));
	 sin(a) * cos(b) (sin(a) * sin(b) * sin(c) + cos(a) * cos(c)) (sin(a) * sin(b) * cos(c) - cos(a) * sin(c));
	 -sin(b) cos(b) * sin(c) cos(b) * cos(c)]
end


# Keep track of parameter estimates
m_x = zeros(T+1,)
w_x = .1 *ones(T+1,)
a_τ = 1. *ones(T+1,)
b_τ = .5 *ones(T+1,)

for k = 1:T

    # Initialize recognition distribution marginals
    marginals = Dict(:x_kmin1 => ProbabilityDistribution(Univariate, GaussianMeanPrecision, m=m_x[k], w=w_x[k]),
                     :x_k => ProbabilityDistribution(Univariate, GaussianMeanPrecision, m=m_x[k], w=w_x[k]),
                     :τ => ProbabilityDistribution(Univariate, Gamma, a=a_τ[k], b=b_τ[k]))

    # Initialize data
    data = Dict(:y_k => Y[k],
                :m_x_kmin1 => m_x[k],
                :w_x_kmin1 => w_x[k],
                :τ1 => a_τ[k],
                :τ2 => b_τ[k])

    # Iterate updates
    for n = 1:10
        stepx!(data, marginals)
        stepτ!(data, marginals)
    end

    # Store estimates
    m_x[k+1] = mean(marginals[:x_k])
    w_x[k+1] = inv(cov(marginals[:x_k]))
    a_τ[k+1] = marginals[:τ].params[:a]
    b_τ[k+1] = marginals[:τ].params[:b]

end

# Generate plot of state estimates
p1 = scatter(1:T, Y, color="black", markersize=6, label="observations", xlabel="time (t)", ylabel="signal")
plot!(1:T, m_x[2:end], ribbon=[sqrt.(inv.(w_x[2:end])) sqrt.(inv.(w_x[2:end]))], color="blue", label="inferred states")

# Generate plot of process noise precision estimates
m_τ = a_τ ./ b_τ
s_τ = sqrt.(a_τ ./ b_τ.^2)
p2 = plot(1:T, m_τ[2:end], ribbon=[s_τ[2:end] s_τ[2:end]], color="blue", xlabel="time (t)", legend=:bottomright, ylabel="τ", label="inferred precision")

# Actually plot
p3 = plot(p1,p2, layout=(2,1), size=(800,600))


using Base.Test
using Elemental

m, n = 10, 10

for T in (Float32, Float64, Complex{Float32}, Complex{Float64})
    for mat in (Elemental.Matrix, Elemental.DistMatrix)
        A = mat(T)
        Elemental.gaussian!(A, m, n)

        @test norm(A, 1) ≈ norm(T[A[i,j] for i = 1:m, j = 1:n], 1)
        @test norm(A) ≈ norm(T[A[i,j] for i = 1:m, j = 1:n])
        @test norm(A, Inf) ≈ norm(T[A[i,j] for i = 1:m, j = 1:n], Inf)
        @test countnz(A) ≈ countnz(T[A[i,j] for i = 1:m, j = 1:n])
    end
end


using IntervalArithmetic, IntervalRootFinding
using ForwardDiff
using Base.Test

const D = IntervalRootFinding.derivative

setprecision(Interval, Float64)
float_pi = @interval(pi)

setprecision(Interval, 10000)
big_pi = @interval(pi)
# Using precision "only" 256 leads to overestimation of the true roots for `cos`
# i.e the Newton method gives more accurate results!

half_pi = big_pi / 2
three_halves_pi = 3*big_pi/2

@testset "Testing newton1d" begin

    f(x) = exp(x^2) - cos(x)
    f′(x) = 2*x*exp(x^2) + sin(x)
    f1(x) = x^4 - 10x^3 + 35x^2 - 50x + 24
    f1′(x) = 4x^3 - 30x^2 + 70x - 50

    for autodiff in (false, true)
        if autodiff
            rts1 = newton1d(sin, -5..5)
            rts2 = newton1d(f, -∞..∞)
            rts3 = newton1d(f1, -10..10)

        else
            rts1 = newton1d(sin, cos, -5..5)
            rts2 = newton1d(f, f′, -∞..∞)
            rts3 = newton1d(f1, f1′, -10..10)
        end

        @test length(rts1) == 3
        L = [(-π.. -π), (0..0), (π..π)]
        for i = 1:length(rts1)
            @test L[i] == rts1[i].interval && :unique == rts1[i].status
        end

        @test length(rts2) == 1
        @test (0..0) == rts2[1].interval && :unique == rts2[1].status

        @test length(rts3) == 4
        L = [1, 2, 3, 4]
        for i = 1:length(rts3)
            @test L[i] in rts3[i].interval && :unique == rts3[i].status
        end
    end
end


#Wave equation Differential Equation
function Zilndroa_Diffeq!(df,f,param,t)

    #Get Operators and variables
    Opps        = param[1]
    vars        = param[2]

    r = Opps.rcoord
    θ = Opps.tcoord

    RHS = Opps.DiffO_1
    Dr  = Opps.Dr
    Dθ  = Opps.Dθ

    dr_ψ  = Opps.aux_vex_1
    dθ_ψ  = Opps.aux_vex_2

    #Fields of our problem (Julia changes the rhs if you change the lhs)
    ψ   = f.x[1]
    ϕ   = f.x[2]

    dψ  = df.x[1]
    dϕ  = df.x[2]

    mul!(dr_ψ , Dr , ψ)
    mul!(dθ_ψ , Dθ , ψ)

    #Differential equation
    mul!(dψ , I ,  ϕ)
    mul!(dϕ , RHS ,  ψ)

    Nr      = vars.rnodes
    Ntheta  = vars.thnodes
    Ω       = vars.Ω
    #Boundary conditions

    for i in 1:Ntheta

        #Inner BC (Impedance)
        idx = (i-1)*Nr+1
        dψ[idx] = Z(t , θ[i] ,  vars ) * ( dr_ψ[idx] - 0.5 * ψ[idx] / r[1] ) - Ω * dθ_ψ[idx]
        #dψ[idx] = 0

        #Outer BC (reflective)
        idx = i*Nr
        dψ[idx] = 0
    end

end

function solve_zilindroa_system(p_in::Param)

    #
    # Setup Routine
    #

    #Welcome message
    welcome_message();

    #Start time
    init_time = time()
    max_runtime = time_string_to_seconds(p_in.max_runtime)

    #Setup routine setup check
    if( !setup_folder(p_in) )
        print_exit_message(  time() - init_time )
        return false
    end

    #Read parameters
    p = p_in
    print_parameters_to_screen(p)
    println("OK")

    #Storage file
    println("Creating HDF5 file...")

    fdata   = "./data"
    fname   = fdata * p.folder *  "$(p.fname).h5"
    fid = h5open(fname, "w")

    save_parameters_to_file( p , fid )
    println("OK")

    #Create operators structures
    println("Setting up Operators...")
    Operators = set_up_operator(p)
    println("OK")

    #Time variables
    tmin = p.t_sim_init
    tmax = p.t_sim_final
    tspan       = ( tmin , tmax )

    iter        = 0
    iter_save   = 0
    every       = floor(Int , p.out_every_t/p.deltat + 1.0)
    save_time   = 0

    #Initial Conditions
    println("Setting up inicial configuration...")
    ψ , dψ = get_gaussian_pulse( p , Operators)

    U = ArrayPartition( ψ , dψ )
    println("OK")

    my_params = ( Operators , p)

    #Problem ODE and integrator
    prob = ODEProblem(  Zilndroa_Diffeq! , U , tspan , my_params )

    integrator = init( prob , RK4() , save_everystep=false , dt=p.deltat , adaptive=false )

    #Create a group for the functions
    g1 =  create_group(fid , "ψ");
    g2 =  create_group(fid , "dψ");

    #Save first configuration of field

    dset_name = string("iteration_",iter_save)

    data1 = integrator.u.x[1]
    dset1, dtype1  = create_dataset(g1, dset_name , data1 )
    write_dataset(dset1, dtype1 , data1)

    data2 = integrator.u.x[2]
    dset2, dtype2  = create_dataset(g2, dset_name , data2 )
    write_dataset(dset2, dtype2 , data2)

    #Add HDF5.attributes time to both datasets
    HDF5.attributes(dset1)["time"] =  integrator.t
    HDF5.attributes(dset2)["time"] =  integrator.t

    #Progress bar and time evolution
    println("Starting up simulation...")
    prog1 = Progress( Int(floor(p.t_sim_final-p.t_sim_init))*100 )

    for (u,t) in tuples(integrator)

        #Increase iteration counter
        iter += 1

        #Step integrator
        DifferentialEquations.step!(integrator)

        #Update bar
        ProgressMeter.update!(prog1, Int(floor(integrator.t*100)))

        #If save iteration
        if iter % every == 0

            iter_save += 1
            save_time = integrator.t

            #data set names
            dset_name = string("iteration_",iter_save)

            #data set for both displacement and derivative
            data1 = integrator.u.x[1]
            dset1, dtype1  = create_dataset(g1, dset_name , data1 )
            write_dataset(dset1, dtype1 , data1)

            data2 = integrator.u.x[2]
            dset2, dtype2  = create_dataset(g2, dset_name , data2 )
            write_dataset(dset2, dtype2 , data2)

            #Add HDF5.attributes time to both datasets
            HDF5.attributes(dset1)["time"] =  integrator.t
            HDF5.attributes(dset2)["time"] =  integrator.t

        end

        #Check if time has exceeded
        if( time_since(init_time) > max_runtime )
            println("\nMaximum allowed runtime reached , Ending Simulation!")
            break
        end

    end

    #Add time information to file
    HDF5.attributes(fid)["max_iter"]    = iter_save
    HDF5.attributes(fid)["t_min"]       = p.t_sim_init
    HDF5.attributes(fid)["t_max"]       = save_time
    HDF5.attributes(fid)["out_every_t"] = p.out_every_t
    HDF5.attributes(fid)["dt"]          = p.deltat

    close(fid)
    println("Saving Files...")

    print_exit_message(  time() - init_time )

    return true

end


# Cost functions

mse(ŷ, y) = sum((ŷ .- y).^2)/length(y)

crossentropy(ŷ::AbstractVecOrMat, y::AbstractVecOrMat) =
  -sum(y .* log.(ŷ)) / size(y, 2)

@deprecate logloss(x, y) crossentropy(x, y)

function logitcrossentropy(logŷ::AbstractVecOrMat, y::AbstractVecOrMat)
  logŷ = logŷ .- maximum(logŷ, 1)
  ypred = logŷ .- log.(sum(exp.(logŷ), 1))
  -sum(y .* ypred) / size(y, 2)
end


# Load modules
using Plots, LaTeXStrings
pyplot()
using ThinFilmsTools

function main()
    # Define beam
    λ = LinRange(400,700,500) # wavelength range [nm]
    θ = [0.] # Angle of incidence
    beam = PlaneWave(λ, θ)
    # Define layers for two systems
    layers1 = [
        LayerTMMO(RIdb.air(beam.λ)),
        LayerTMMO(RIdb.dummy(beam.λ,1.38,0.0); type=:OT, d=1/4.),
        LayerTMMO(RIdb.dummy(beam.λ,2.15,0.0); type=:OT, d=1/2.),
        LayerTMMO(RIdb.bk7(beam.λ)),
    ]
    layers2 = [
        LayerTMMO(RIdb.air(beam.λ)),
        LayerTMMO(RIdb.dummy(beam.λ,1.38,0.0); type=:OT, d=1/4.),
        LayerTMMO(RIdb.dummy(beam.λ,1.90,0.0); type=:OT, d=1/2.),
        LayerTMMO(RIdb.bk7(beam.λ)),
    ]
    # Reference wavelenth [nm]
    λ0 = 509.8
    # call main script
    sol1 = tmm_optics(beam, layers1; λ0=λ0)
    sol2 = tmm_optics(beam, layers2; λ0=λ0)
    return sol1, sol2
end

sol1, sol2 = main()

# plot the R, T and A spectra
plot(Spectrum1D(),
     sol1.beam.λ, sol1.Spectra.Rs,
     label=L"Ta$_2$O$_5$", line=(:solid),
     ylims=(0.0, 0.025),
     xlims=(sol1.beam.λ[1], sol1.beam.λ[end]),
     yaxis=("Reflectance"),
);
plot!(Spectrum1D(),
      sol2.beam.λ, sol2.Spectra.Rs,
      label=L"Y$_2$O$_3$",
      line=(:dashdot),
);
gui()

# plot the refractive index profile
plot(RIprofile(), sol1, title=L"Ta$_2$O$_5$")
gui()
### Note that if you choose θ outside the range it will show the EMF for one extrema. Same for λ.

plot(RIprofile(), sol2, title=L"Y$_2$O$_3$")
gui()
### Note that if you choose θ outside the range it will show the EMF for one extrema. Same for λ.


using SpecialFunctions

# https://github.com/JuliaMath/SpecialFunctions.jl/blob/d18ff04178dd37a60cc716a60ab3caf1a6903f43/src/gamma.jl#L20
@cufunc function SpecialFunctions.digamma(x::T) where {T<:Real}
    ψ = zero(x)
    if x ≤ 0 # reflection formula
        ψ -= π / tan(π * x)
        x = 1 - x
    end

    if x < 7
        # shift using recurrence formula
        ν = one(x)
        n = 7 - floor(x)
        while ν <= n - 1
            ψ -= inv(x + ν)
            ν += one(x)
        end
        ψ -= inv(x)
        x += n
    end
    t = inv(x)
    ψ += log(x) - t / 2
    t *= t # 1/z^2
    # the coefficients here are Float64(bernoulli[2:9] .// (2*(1:8)))
    ψ -= (
        t *
        @evalpoly(
            t,
            T(0.08333333333333333), T(-0.008333333333333333),
            T(0.003968253968253968), T(-0.004166666666666667),
            T(0.007575757575757576), T(-0.021092796092796094),
            T(0.08333333333333333), T(-0.4432598039215686)
        )
    )
    return ψ
end

function _cutrigamma(x::T) where {T<:Real}
    ψ = zero(x)
    if x < 8
        # shift using recurrence formula
        n = 8 - CUDA.floor(x)
        ψ += inv(x)^2
        ν = one(x)
        while ν <= n - 1
            ψ += inv(x + ν)^2
            ν += one(x)
        end
        x += n
    end
    t = inv(x)
    w = t * t # 1/z^2
    ψ += t + w / 2
    # the coefficients here are Float64(bernoulli[2:9])
    p = @evalpoly(
        w,
        T(0.16666666666666666), T(-0.03333333333333333),
        T(0.023809523809523808), T(-0.03333333333333333),
        T(0.07575757575757576), T(-0.2531135531135531),
        T(1.1666666666666667), T(-7.092156862745098)
    )
    ψ += t * w * p
    return ψ
end

@cufunc function SpecialFunctions.trigamma(x)
    if x ≤ 0 # reflection formula
        return (π / CUDA.sin(π * x))^2 - _cutrigamma(1 - x)
    else
        return _cutrigamma(x)
    end
end

@cufunc SpecialFunctions.loggamma(x) = CUDA.lgamma(x)

@cufunc_register SpecialFunctions.logbeta
@cufunc_function SpecialFunctions.logbeta(a, b) = (
      SpecialFunctions.loggamma(a)
    + SpecialFunctions.loggamma(b)
    - SpecialFunctions.loggamma(a + b)
)

@cufunc SpecialFunctions.erf(x) = CUDA.erf(x)
@cufunc SpecialFunctions.erfc(x) = CUDA.erfc(x)
@cufunc SpecialFunctions.erfcx(x) = CUDA.erfcx(x)
@cufunc SpecialFunctions.erfinv(x) = CUDA.erfinv(x)
@cufunc SpecialFunctions.erfcinv(x) = CUDA.erfcinv(x)

#####################
### Distributions ###
#####################
using StatsFuns
import StatsFuns: gammalogpdf, zval, xval, normlogpdf, normcdf, log2π, invsqrt2

@cufunc StatsFuns.gammalogpdf(k::Real, θ::Real, x::Number) = -loggamma(k) - k * log(θ) + (k - 1) * log(x) - x / θ

@cufunc function StatsFuns.xval(μ::Real, σ::Real, z::Number)
    if isinf(z) && iszero(σ)
        μ + one(σ) * z
    else
        μ + σ * z
    end
end
@cufunc StatsFuns.zval(μ::Real, σ::Real, x::Number) = (x - μ) / σ

# logpdf
@cufunc_register StatsFuns.normlogpdf
@cufunc_function StatsFuns.normlogpdf(z::Number) = -(abs2(z) + log2π)/2
@cufunc_function function StatsFuns.normlogpdf(μ::Real, σ::Real, x::Number)
    z = zval(μ, σ, x)
    normlogpdf(z) - log(σ)
end

# cdf
@cufunc_register StatsFuns.normcdf
@cufunc_function StatsFuns.normcdf(z::Number) = erfc(-z * invsqrt2)/2
@cufunc_function function StatsFuns.normcdf(μ::Real, σ::Real, x::Number)
    z = zval(μ, σ, x)
    normcdf(z)
end

# logcdf
@cufunc_register StatsFuns.normlogcdf
@cufunc_function function StatsFuns.normlogcdf(z::Number)
    invsqrt2 = inv(sqrt(oftype(z, 2)))
    if z < -1.0
        log(SpecialFunctions.erfcx(-z * invsqrt2)/2) - abs2(z)/2
    else
        SpecialFunctions.log1p(-SpecialFunctions.erfc(z * invsqrt2)/2)
    end
end
@cufunc_function function StatsFuns.normlogcdf(μ::Real, σ::Real, x::Number)
    z = zval(μ, σ, x)
    normlogcdf(z)
end

# truncated normlogpdf
@cufuncf function truncatednormlogpdf(μ, σ, x, lb, ub)
    logtp = StatsFuns.normlogcdf(μ, σ, ub) - StatsFuns.normlogcdf(μ, σ, lb)
    # TODO: deal with outside of boundary
    StatsFuns.normlogpdf(μ, σ, x) - logtp
    # TODO: seems like there's something messed up with the way we return `Inf`
    # if lb <= x <= ub
    #     StatsFuns.normlogpdf(μ, σ, x) - logtp
    # else
    #     TF = float(eltype(x))
    #     -TF(Inf)
    # end
end


@testset "CUFFT" begin

if !isdefined(CuArrays, :CUFFT)
@warn "Not testing CUFFT"
else
using CuArrays.CUFFT
@info "Testing CUFFT $(CUFFT.version())"

# notes:
#   plan_bfft does not need separate testing since it is used by plan_ifft

using FFTW

N1 = 8
N2 = 32
N3 = 64
N4 = 8

MYRTOL = 1e-5
MYATOL = 1e-8

# out-of-place
function dotest1(X::AbstractArray{T,N}) where {T <: Complex,N}
    fftw_X = fft(X)
    d_X = CuArray(X)
    p = plan_fft(d_X)
    d_Y = p * d_X
    Y = collect(d_Y)
    @test isapprox(Y, fftw_X, rtol = MYRTOL, atol = MYATOL)

    pinv = plan_ifft(d_Y)
    d_Z = pinv * d_Y
    Z = collect(d_Z)
    @test isapprox(Z, X, rtol = MYRTOL, atol = MYATOL)

    pinv2 = inv(p)
    d_Z = pinv2 * d_Y
    Z = collect(d_Z)
    @test isapprox(Z, X, rtol = MYRTOL, atol = MYATOL)
end

function dotest1(X::AbstractArray{T,N}) where {T <: Real,N}
    fftw_X = rfft(X)
    d_X = CuArray(X)
    p = plan_rfft(d_X)
    d_Y = p * d_X
    Y = collect(d_Y)
    @test isapprox(Y, fftw_X, rtol = MYRTOL, atol = MYATOL)

    pinv = plan_irfft(d_Y,size(X,1))
    d_Z = pinv * d_Y
    Z = collect(d_Z)
    @test isapprox(Z, X, rtol = MYRTOL, atol = MYATOL)

    pinv2 = inv(p)
    d_Z = pinv2 * d_Y
    Z = collect(d_Z)
    @test isapprox(Z, X, rtol = MYRTOL, atol = MYATOL)

    pinv3 = inv(pinv)
    d_W = pinv3 * d_X
    W = collect(d_W)
    @test isapprox(W, Y, rtol = MYRTOL, atol = MYATOL)
end

# in-place
function dotest2(X::AbstractArray{T,N}) where {T <: Complex,N}
    fftw_X = fft(X)
    d_X = CuArray(X)
    p = plan_fft!(d_X)
    p * d_X
    Y = collect(d_X)
    @test isapprox(Y, fftw_X, rtol = MYRTOL, atol = MYATOL)

    pinv = plan_ifft!(d_X)
    pinv * d_X
    Z = collect(d_X)
    @test isapprox(Z, X, rtol = MYRTOL, atol = MYATOL)
end

# no inplace rfft for now

# batch transforms
function dotest3(X::AbstractArray{T,N},region) where {T <: Complex,N}
    fftw_X = fft(X,region)
    d_X = CuArray(X)
    p = plan_fft(d_X,region)
    d_Y = p * d_X
    Y = collect(d_Y)
    @test isapprox(Y, fftw_X, rtol = MYRTOL, atol = MYATOL)

    pinv = plan_ifft(d_Y,region)
    d_Z = pinv * d_Y
    Z = collect(d_Z)
    @test isapprox(Z, X, rtol = MYRTOL, atol = MYATOL)
end

function dotest3(X::AbstractArray{T,N},region) where {T <: Real,N}
    fftw_X = rfft(X,region)
    d_X = CuArray(X)
    p = plan_rfft(d_X,region)
    d_Y = p * d_X
    Y = collect(d_Y)
    @test isapprox(Y, fftw_X, rtol = MYRTOL, atol = MYATOL)

    pinv = plan_irfft(d_Y,size(X,region[1]),region)
    d_Z = pinv * d_Y
    Z = collect(d_Z)
    @test isapprox(Z, X, rtol = MYRTOL, atol = MYATOL)
end

@testset "FFT" for (rtype,ctype) in [(Float32,ComplexF32), (Float64,ComplexF64)]

@testset "1D FFT" begin
    dims = (N1,)
    X = rand(ctype, dims)
    dotest1(X)
end
@testset "1D inplace FFT" begin
    dims = (N1,)
    X = rand(ctype, dims)
    dotest2(X)
end

@testset "2D FFT" begin
    dims = (N1,N2)
    X = rand(ctype, dims)
    dotest1(X)
end
@testset "2D inplace FFT" begin
    dims = (N1,N2)
    X = rand(ctype, dims)
    dotest2(X)
end

@testset "Batch 1D FFT" begin
    dims = (N1,N2)
    X = rand(ctype, dims)
    dotest3(X,1)

    dims = (N1,N2)
    X = rand(ctype, dims)
    dotest3(X,2)

    dims = (N1,N2)
    X = rand(ctype, dims)
    dotest3(X,(1,2))
end

@testset "3D FFT" begin
    dims = (N1,N2,N3)
    X = rand(ctype, dims)
    dotest1(X)
end
@testset "3D inplace FFT" begin
    dims = (N1,N2,N3)
    X = rand(ctype, dims)
    dotest2(X)
end

@testset "Batch 2D FFT (in 3D)" begin
    dims = (N1,N2,N3)
    for region in [(1,2),(2,3),(1,3)]
        X = rand(ctype, dims)
        dotest3(X,region)
    end

    X = rand(ctype, dims)
    @test_throws ArgumentError dotest3(X,(3,1))
end

@testset "Batch 2D FFT (in 4D)" begin
    dims = (N1,N2,N3,N4)
    for region in [(1,2),(1,4),(3,4)]
        X = rand(ctype, dims)
        dotest3(X,region)
    end
    for region in [(1,3),(2,3),(2,4)]
        X = rand(ctype, dims)
        @test_throws ArgumentError dotest3(X,region)
    end

end

@testset "1D real FFT" begin
    X = rand(rtype, N1)
    dotest1(X)
end

@testset "Batch 1D real FFT" begin
    dims = (N1,N2)
    X = rand(rtype, dims)
    dotest3(X,1)

    dims = (N1,N2)
    X = rand(rtype, dims)
    dotest3(X,2)

    dims = (N1,N2)
    X = rand(rtype, dims)
    dotest3(X,(1,2))
end

@testset "2D real FFT" begin
    X = rand(rtype, N1,N2)
    dotest1(X)
end

@testset "Batch 2D real FFT (in 3D)" begin
    dims = (N1,N2,N3)
    for region in [(1,2),(2,3),(1,3)]
        X = rand(rtype, dims)
        dotest3(X,region)
    end

    X = rand(rtype, dims)
    @test_throws ArgumentError dotest3(X,(3,1))
end

@testset "Batch 2D real FFT (in 4D)" begin
    dims = (N1,N2,N3,N4)
    for region in [(1,2),(1,4),(3,4)]
        X = rand(rtype, dims)
        dotest3(X,region)
    end
    for region in [(1,3),(2,3),(2,4)]
        X = rand(rtype, dims)
        @test_throws ArgumentError dotest3(X,region)
    end
end

@testset "3D real FFT" begin
    X = rand(rtype, N1, N2, N3)
    dotest1(X)
end

end # testset FFT

# integer array arguments
function dotest5(X::AbstractArray{T,N}) where {T <: Complex,N}
    fftw_X = fft(X)
    d_X = CuArray(X)
    p = plan_fft(d_X)
    d_Y = p * d_X
    Y = collect(d_Y)
    @test isapprox(Y, fftw_X, rtol = MYRTOL, atol = MYATOL)
    d_Y = fft(d_X)
    Y = collect(d_Y)
    @test isapprox(Y, fftw_X, rtol = MYRTOL, atol = MYATOL)
end

function dotest5(X::AbstractArray{T,N}) where {T <: Real,N}
    fftw_X = rfft(X)
    d_X = CuArray(X)
    p = plan_rfft(d_X)
    d_Y = p * d_X
    Y = collect(d_Y)
    @test isapprox(Y, fftw_X, rtol = MYRTOL, atol = MYATOL)
    d_Y = rfft(d_X)
    Y = collect(d_Y)
    @test isapprox(Y, fftw_X, rtol = MYRTOL, atol = MYATOL)
end

@testset "Int FFT" for (rtype,ctype) in [(Int32,Complex{Int32}), (Int64,Complex{Int64})]

@testset "1D FFT" begin
    dims = (N1,)
    X = rand(ctype, dims)
    dotest5(X)
end

@testset "1D real FFT" begin
    X = rand(rtype, N1)
    dotest5(X)
end

end # testset int FFT

end

end


using Test
using ApproxManifoldProducts
using Manifolds

##

@testset "test getManifoldPartial on Euclidean(N)" begin

##

M = TranslationGroup(3)

@test getManifoldPartial(M, [1;2;3])[1] == TranslationGroup(3)
@test getManifoldPartial(M, [2;3])[1] == TranslationGroup(2)

@test getManifoldPartial(M, [1;2;3], zeros(3))[1] == TranslationGroup(3)
@test isapprox( getManifoldPartial(M, [1;2;3], zeros(3))[2], zeros(3) )

@test getManifoldPartial(M, [2;3], zeros(3))[1] == TranslationGroup(2)
@test isapprox( getManifoldPartial(M, [2;3], zeros(3))[2], zeros(2))

##
end

@testset "test getManifoldPartial on Circle()" begin

##

M = Circle()

@test getManifoldPartial(M, [1])[1] == Circle()
@test_throws ErrorException getManifoldPartial(M, [2;])

@test getManifoldPartial(M, [1], [0;])[1] == Circle()
@test getManifoldPartial(M, [1], [0;])[2] == [0]

##
end

@testset "test getManifoldPartial on Rotations(2)" begin

##

M = Rotations(2)

@test getManifoldPartial(M, [1])[1] == Rotations(2)
@test_throws ErrorException getManifoldPartial(M, [2;])

@test getManifoldPartial(M, [1], [1 0; 0 1])[1] == Rotations(2)
@test getManifoldPartial(M, [1], [1 0; 0 1])[2] == [1 0; 0 1]

##
end

@testset "test getManifoldPartial on SpecialEuclidean(2)" begin

##

M = SpecialEuclidean(2)

@test getManifoldPartial(M, [1;2;3])[1] == SpecialEuclidean(2)

@test getManifoldPartial(M, [1;])[1] == TranslationGroup(1)
@test getManifoldPartial(M, [2;])[1] == TranslationGroup(1)
@test getManifoldPartial(M, [1;2])[1] == TranslationGroup(2)

@test getManifoldPartial(M, [3;])[1] == SpecialOrthogonal(2)

@test getManifoldPartial(M, [1;3;])[1] == ProductManifold(TranslationGroup(1), SpecialOrthogonal(2))

repr = ProductRepr([0.0; 0], [1 0; 0 1.0])

@test getManifoldPartial(M, [1;2;3], repr)[1] == SpecialEuclidean(2)
@test getManifoldPartial(M, [1;2;3], repr)[2] == repr

@test getManifoldPartial(M, [1;], repr)[1] == TranslationGroup(1)
@test getManifoldPartial(M, [1;], repr)[2] == [0.0;]

@test getManifoldPartial(M, [2;], repr)[1] == TranslationGroup(1)
@test getManifoldPartial(M, [2;], repr)[2] == [0.0;]

@test getManifoldPartial(M, [1;2], repr)[1] == TranslationGroup(2)
@test getManifoldPartial(M, [1;2], repr)[2] == [0.0;0]

@test getManifoldPartial(M, [3;], repr)[1] == SpecialOrthogonal(2)
@test getManifoldPartial(M, [3;], repr)[2] == repr.parts[2]

@test getManifoldPartial(M, [1;3;], repr)[1] == ProductManifold(TranslationGroup(1), SpecialOrthogonal(2))
r_repr = getManifoldPartial(M, [1;3;], repr)[2]
@test r_repr isa ProductRepr
@test r_repr.parts[1] == [0.0;]
@test r_repr.parts[2] == [1 0; 0 1.0]

##
end

@testset "Reminder, getManifoldPartial on Sphere(2) [TBD]" begin

##

@error "Must fix Sphere(2) partial test"
@test_broken false

##
end

@testset "test getPoints under partial with representation" begin

##

N = 100
M = SpecialEuclidean(2)
u0 = ProductRepr([0.0; 0], [1 0; 0 1.0])

pts = [exp(M, u0, hat(M, u0, [10 .+ randn(2);randn()])) for i in 1:N]

P = manikde!(M, pts)

P12 = marginal(P, [1;2])

p12 = getPoints(P12)

@test length(p12) == N
@test length(p12[1]) == 2
@test_broken P12.manifold isa TranslationGroup(2)

##
end

@testset "test marginal of marginal (partial) helper" begin
##

M = TranslationGroup(3)
pts = [randn(3) for _ in 1:75]

X = manikde!(M, pts, partial=[1;3])

X_ = marginal(X, [3])

ps3 = getPoints(X_)

for (i,pt) in enumerate(pts)
  @test isapprox(ps3[i][1], pt[3])
end

try
  M = TranslationGroup(4)
  # check the constructor when only a few points are available
  X = manikde(M, pts, partial=[1;3;4])
catch
  @test_broken false
end

##
end

@testset "test replace overloads full and partial/marginal" begin
##

N = 10
M = TranslationGroup(3)
pts0 = [zeros(3) for _ in 1:N]
X0 = manikde!(M, pts0, bw=zeros(3))

pts = [randn(3) for _ in 1:N]
X = manikde!(M, pts)

##

X_ = replace(X0, X)
@test isapprox(X_, X)

##

X = manikde!(M, pts, partial=[1;3])
X_ = replace(X0, X)

# check metadata
@test  isapprox( getBW(X_, false)[[1;3],1], getBW(X,  false)[[1;3],1] )
@test !isapprox( getBW(X_, false)[[1;3],1], getBW(X0, false)[[1;3],1] )

@test  isapprox( X_.infoPerCoord[[1;3]], X.infoPerCoord[[1;3]] )

@test !isPartial(X_)

# check points
x0 = getPoints(X0)
x  = getPoints(X, false)
x_ = getPoints(X_)
for (i,pt) in enumerate(x0)
  # partial of X does not replace
  @test isapprox(pt[2], x_[i][2])
  @test isapprox(x[i][[1;3]], x_[i][[1;3]])
end

## must also test replace for partial into different partial

pts3 = [randn(3) for _ in 1:N]
X3 = manikde!(M, pts3, partial=[3;])
# and replace partial/marginal values
X__ = replace(X, X3)

@test isPartial(X__)
@test X__._partial == [1;3]

x__ = getPoints(X__, false)
for (i,pt) in enumerate(x)
  @test isapprox(pt[1:2], x__[i][1:2])
  @test isapprox(pts3[i][3], x__[i][3])
end

## union of two partials over all dimensions should drop the partial status

pts12 = [randn(3) for _ in 1:N]
X12 = manikde!(M, pts12, partial=[1;2], infoPerCoord=2*ones(3))

X_np = replace(X12, X3)

@test !isPartial(X_np)
@test isapprox( X_np.infoPerCoord, [2;2;1] )

##
end

#


export repressilator

using LabelledArrays, DynamicalSystemsBase, DelayEmbeddings

@inline @inbounds function eom_repressilator(u, p, t)
    α, α₀, n, β = p.α, p.α₀, p.n, p.β

    # pᵢ := concentration of protein repressor i
    # mᵢ := concentration of mRNA associated with pᵢ
    m₁, m₂, m₃, p₁, p₂, p₃ = u

    ṁ₁ = -m₁ + α/(1 + p₃^n) + α₀
    ṁ₂ = -m₂ + α/(1 + p₁^n) + α₀
    ṁ₃ = -m₃ + α/(1 + p₂^n) + α₀
    ṗ₁ = -β*(p₁ - m₁)
    ṗ₂ = -β*(p₂ - m₂)
    ṗ₃ = -β*(p₃ - m₃)

    return SVector{6}(ṁ₁, ṁ₂, ṁ₃, ṗ₁, ṗ₂, ṗ₃)
end

"""
    repressilator(;u₀ = rand(6), α = 10.0, α₀ = 0.0, β = 100.0,
        n = 2) → ContinuousDynamicalSystem

A six-dimensional repressilator (or repression-driven oscillator) from Elowitz & Leibler
(2000)[^Elowitz2000]. The equations are scaled to be non-dimensional.

Used in Sun & Bollt (2014) to study the performance of the causation entropy algorithm.

[^Elowitz2000]: Elowitz, M. B., & Leibler, S. (2000). A synthetic oscillatory network of
    transcriptional regulators. Nature, 403(6767), 335-338.
[^Sun2014]: Sun, J., Cafaro, C., & Bollt, E. M. (2014). Identifying the coupling structure
    in complex systems through the optimal causation entropy principle. Entropy, 16(6),
    3416-3433.
"""
function repressilator(;u₀ = rand(6), α = 10.0, α₀ = 0.0, β = 100.0, n = 2)

    p = @LArray [α, α₀, β, n] (:α, :α₀, :β, :n)
    ContinuousDynamicalSystem(eom_repressilator, u₀, p)
end


function second_moment(q::Normal)
    return abs2(mean(q)) + var(q)
end

function second_moment(q::Normal, y::Real)
    return abs2(mean(q) - y) + var(q)
end

# This is not exactly an approximation but corresponds anyway
# to the expectation of the function σ(f)
function approx_expected_logistic(μ, c)
    lower, upper = LogExpFunctions._logistic_bounds(μ)
    return μ < lower ? zero(μ) : (μ > upper ? one(μ) : exp(μ / 2) * sech(c / 2) / 2)
end


"""
```
SeisDiff(d ; <keyword arguments>)
```

Apply differentiation to seismic traces in freq. It can also be
used to apply a phase rotation

# Arguments
* `d`: Input 2D data array in tx domain. First dimension is time.

# Keyword arguments
* `delay=0.1`: desired time delay in seconds.
* `pow=-2`: order of derivative
* `rot=0`: constant phase shift or rotation

# Example
```
julia> d,extent = SeisLinearEvents(d); SeisPlot(d);
julia> d2 = SeisDiff(d;dt=0.004,pow=1); SeisPlot(d);
```
"""

function SeisDiff(d;dt=0.001,pow=-2,rot=0)

	nt = size(d,1)
	D = fft(d,1)
	dw = 2.*pi/nt/dt
	nw = Int(nt/2) + 1
	eps = pow < 0 ? 10.0 : 0.0
	for iw=1:nw
		D[iw,:] *= exp(rot*1im)*((iw*dw) + eps)^pow
	end
	# symmetries
	for iw=nw+1:nt
		D[iw,:] = conj(D[nt-iw+2,:])
	end
	d = real(ifft(D,1))
	return d

end

function SeisDiff(in,h::Array{Header,1};pow=-2,rot=0)

	out = differentiate_traces(in;dt=h[1].d1,pow=pow,rot=rot)
	return out,h

end

function SeisDiff(in::AbstractString,out::AbstractString;pow=-2,rot=0)

	@compat parameters = Dict(:pow=>pow,:rot=>rot)
	SeisProcess(in,out,[SeisDiff],[parameters])

end


# test on partial products with SpecialEuclidean(2)

using Manifolds
using ApproxManifoldProducts
using Test
# using Random
using FileIO

##

@testset "partial product with a SpecialEuclidean(2)" begin
##

datafile = joinpath(@__DIR__, "testdata", "partialtest.jld2")
dict = load(datafile)
pts1 = dict["pts1"]
pts2 = dict["pts2"]

randU = Float64[]
randN = Float64[]

len = length(pts1)

# define test manifold
M = SpecialEuclidean(2)
e0 = identity_element(M)

# p1 full SpecialEuclidean(2)
p1 = manikde!(M, pts1)
p1_ = marginal(p1,[1;2])

# p2 only Translation(2) part
_p2 = manikde!(M, pts2)
p2 = marginal(_p2, [1;2])

# product of full and marginal
# p12 = p1*p2

## =======================FIRST PRODUCT================================================

selectedLabels=Vector{Vector{Int}}()
p12 = manifoldProduct([p1; p2]; addEntropy=false,
                                recordLabels=true,
                                selectedLabels=selectedLabels,
                                _randU=randU,
                                _randN=randN  )
#
selectedLabels
#

@test !isPartial(p12)

p12_ = marginal(p12, [1;2])
@test isPartial(p12_)

_p12_ = manikde!(TranslationGroup(2), getPoints(p12_))

## intermediate test, check product of selected kernels match what is in the marginal

for sidx = 1:len
  bw1 = getBW(p1)[:,1] .^2
  bw2 = getBW(p2, false)[:,1] .^2

  u1 = pts1[selectedLabels[sidx][1]]
  u2 = pts2[selectedLabels[sidx][2]]

  u12, = calcProductGaussians(M, [u1,u2], [bw1,bw2]);

  @test isapprox( u12.parts[1], getPoints(p12)[sidx].parts[1] )
end

## now check the marginal dimensions only

# now do submanifold dimensions separately as reference test -- should get a similar result
pts1_ = getPoints(p1_)
pts2_ = getPoints(p2)

## Do the translation part separate

for sidx = 1:len
  bw1 = getBW(p1_)[:,1] .^2
  bw2 = getBW(p2)[:,1] .^2

  u1 = pts1_[selectedLabels[sidx][1]]
  u2 = pts2_[selectedLabels[sidx][2]]

  u12, = calcProductGaussians(TranslationGroup(2), [u1,u2], [bw1,bw2])

  @test isapprox( u12, getPoints(p12)[sidx].parts[1] )
end

## =======================SECOND PRODUCT================================================
## NEW PRODUCT ON ONLY THE TRANSLATION PART, COMPARE WITH ABOVE RESULT

p1__ = manikde!(TranslationGroup(2), pts1_)
p2__ = manikde!(TranslationGroup(2), pts2_)

# p12__ = p1__*p2__
selectedLabels__=Vector{Vector{Int}}()
p12__ = manifoldProduct([p1__; p2__]; addEntropy=false,
                                      recordLabels=true,
                                      selectedLabels=selectedLabels__,
                                      _randU=randU,
                                      _randN=randN  )
#
selectedLabels__

## compare calcProduct of full*partial with selection from partial*partial

sidx = 1
for sidx = 1:len

bw1 = getBW(p1)[:,1] .^2
bw2 = getBW(p2, false)[:,1] .^2

# full-partial points, but selected from partial-partial product
u1 = pts1[selectedLabels__[sidx][1]]
u2 = pts2[selectedLabels__[sidx][2]]

u12, = calcProductGaussians(M, [u1,u2], [bw1,bw2]);
u12_, = calcProductGaussians(TranslationGroup(2), [u1.parts[1],u2.parts[1]], [bw1[1:2],bw2[1:2]]);

@test isapprox( u12.parts[1], u12_ )
@test isapprox( getPoints(p12__)[sidx], u12_ )

end

##

@test mmd(_p12_, p12__) < 0.001

##
end

## DEBUG PLOTTING ==========================================================================

## Plots showing the problem, p12 was wrong!!

# using Cairo, RoMEPlotting
# Gadfly.set_default_plot_size(35cm,20cm)

# n=10; plotKDE([p1_;p2; p12], levels=3, selectedPoints=selectedLabels[n:n])
# n=3; plotKDE([p1__;p2__; p12__], levels=3, selectedPoints=selectedLabels__[n:n])
# plotKDE([p1__; p2__; p12__])

# plotKDE([_p12_; p12__])

##


"""
    NoiseSource(eltype, samplerate, nchannels, std=1)

NoiseSource is a multi-channel noise signal generator. The noise on each channel is independent.

Inputs
------
* `samplerate` specifies the sample rate of the signal specified in Hz.
* `nchannels` specifies the number of channels of the signal.
* `std` specifies the desired standard deviation of the signal.

Output
------
* SampleSource object

Example
-------
```julia
source_object = NoiseSource(Float64, 48u"kHz", 2, 0.3)
wn = read(source_object, 480)         # Specify number of samples of signal to generate
wn = read(source_object, 50u"ms")     # Specify length of time of signal to generate
```
"""
mutable struct NoiseSource{T} <: SampleSource
    samplerate::Float64
    nchannels::Int64
    std::Float64

    function NoiseSource(eltype, samplerate::Number, nchannels::Int, std::Number=1)
        new{eltype}(samplerate, nchannels, std)
    end
    function NoiseSource(eltype, samplerate::Unitful.Frequency, nchannels::Int, std::Number=1)
        samplerate = ustrip(uconvert(u"Hz", samplerate))
        NoiseSource(eltype, samplerate, nchannels, std)
    end
end

Base.eltype(::NoiseSource{T}) where T = T
nchannels(source::NoiseSource) = source.nchannels
samplerate(source::NoiseSource) = source.samplerate

function unsafe_read!(source::NoiseSource, buf::Array, frameoffset, framecount)
    buf[1+frameoffset:framecount+frameoffset, 1:source.nchannels] = source.std .* randn(framecount, source.nchannels)
    framecount
end


# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    Quantile(distribution=Normal())

The quantile transform to a given `distribution`.

# Examples

```julia
Quantile()
Quantile(Normal())
```
"""
struct Quantile{D} <: Colwise
  dist::D
end

Quantile() = Quantile(Normal())

assertions(::Type{Quantile}) = [assert_continuous]

isrevertible(::Type{<:Quantile}) = true

colcache(::Quantile, x) = EmpiricalDistribution(x)

function colapply(transform::Quantile, x, c)
  origin, target = c, transform.dist
  qqtransform(x, origin, target)
end

function colrevert(transform::Quantile, y, c)
  origin, target = transform.dist, c
  qqtransform(y, origin, target)
end

# transform samples from original to target distribution
function qqtransform(samples, origin, target)
  # avoid evaluating the quantile at 0 or 1
  T = eltype(samples)
  low = zero(T) + T(0.001)
  high = one(T) - T(0.001)
  map(samples) do sample
    prob = cdf(origin, sample)
    quantile(target, clamp(prob, low, high))
  end
end


abstract type AbstractChain end

# Chain interface
Base.getindex(ch::AbstractChain, i::Integer) = throw(MethodError(getindex, (typeof(ch),Integer)))
Base.length(ch::AbstractChain) = throw(MethodError(length, (typeof(ch),)))
Base.eltype(ch::AbstractChain) = throw(MethodError(eltype, (typeof(ch),)))
function Base.iterate(ch::AbstractChain, state=1)
    state > length(ch) && return nothing
    return (ch[state], state+1)
end

"""Empty Chain"""
struct EmptyChain <: AbstractChain end
Base.length(ch::EmptyChain) = 0

""" k-Chain is a formal linear combination of k-cells.

    Given a set E, we can construct a free R-module M that has E ⊆ M as a basis M, the module M is the module of the formal linear combinations of elements of E, or free module over E, R^{(E)}, s.t. given a finite subset {X_1, ..., X_n} of E, a formal linear combination of X_1, ..., X_n is

    a_1 X_1 + ··· + a_n X_n,

    where the a_i ∈ R.
"""
mutable struct Chain{PID, IX<:Integer} <: AbstractChain
    dim::Int
    coefs::Vector{PID} # coefficients from PID
    elems::Vector{IX}  # indexes to elements of E
end
Chain(coefs::Vector{PID}, elems::Vector{IX}) where {PID, IX<:Integer} = Chain{PID, IX}(0, coefs, elems)
Chain(dim::Int, ::Type{PID}, ::Type{IX}) where {PID, IX<:Integer} = Chain{PID, IX}(dim, PID[], IX[])
Chain(dim::Int, ::Type{PID}) where {PID} = Chain(dim, PID, Int)
Chain(::Type{PID}) where {PID} = Chain(0,PID)
dim(ch::Chain) = ch.dim

# implement interface
Base.getindex(ch::Chain, i::Integer) = (ch.coefs[i], ch.elems[i])
Base.length(ch::Chain) = length(ch.coefs)
Base.eltype(ch::Chain{PID, IX}) where {PID, IX<:Integer} = (PID, IX)
Base.iszero(ch::Chain) = length(ch.elems) == 0

function Base.show(io::IO, ch::Chain)
    if iszero(ch)
        print(io, "0")
    else
        for (i,(c,el)) in enumerate(zip(ch.coefs, ch.elems))
            i != 1 && print(io, " + ")
            print(io, "$c[$el]")
        end
    end
end

function Base.push!(ch::Chain{PID, IX}, coef::PID, idx::IX) where {PID, IX<:Integer}
    push!(ch.coefs, coef)
    push!(ch.elems, idx)
    return ch
end
Base.push!(ch::Chain{PID, IX}, e::Tuple{PID,IX}) where {PID, IX<:Integer} = push!(ch, e[1], e[2])
Base.push!(ch::Chain{PID, IX}, e::Pair{PID,IX}) where {PID, IX<:Integer}  = push!(ch, e[1], e[2])
function Base.append!(a::Chain, b::Chain)
    append!(a.coefs, b.coefs)
    append!(a.elems, b.elems)
    return a
end
+(ch::Chain{PID}, e::Tuple{PID,Integer}) where {PID} = push!(ch, e[1], e[2])

+(a::Chain{PID, IX}, b::Chain{PID, IX}) where {PID, IX<:Integer} = Chain{PID, IX}(dim(a), vcat(a.coefs, b.coefs), vcat(a.elems, b.elems))
-(a::Chain{PID, IX}, b::Chain{PID, IX}) where {PID, IX<:Integer} = Chain{PID, IX}(dim(a), vcat(a.coefs, -b.coefs), vcat(a.elems, b.elems))

*(ch::Chain{PID, IX}, r::PID) where {PID, IX<:Integer} = Chain{PID, IX}(dim(ch), ch.coefs*r, ch.elems)
*(r::PID, c::Chain{PID}) where {PID} = c*r

function simplify(ch::Chain)
    PID, IX = eltype(ch)
    tmpChain = Dict{Integer,PID}()
    for (c,el) in ch
        if haskey(tmpChain, el)
            tmpChain[el] += c
        else
            tmpChain[el] = c
        end
    end
    cc = Chain(dim(ch), PID, IX)
    for (el,c) in tmpChain
        c != zero(PID) && push!(cc, (c,el))
    end
    return cc
end


using TropicalYao.Reversible.LogLikeNumbers
using Test

@testset "i_einsum" begin
	N = 3
	a = Tropical.(randn(N, N))
	b = Tropical.(randn(N, N))
	c = Tropical.(zeros(N, N))
	i_einsum!(((1,2), (2,3)), (a, b), (1,3), c)
	@test c ≈ a*b
end


@testset "Angles" begin
  @test ∠(P2(0,1),P2(0,0),P2(1,0)) ≈ T(-π/2)
  @test ∠(P2(1,0),P2(0,0),P2(0,1)) ≈ T(π/2)
  @test ∠(P2(-1,0),P2(0,0),P2(0,1)) ≈ T(-π/2)
  @test ∠(P2(0,1),P2(0,0),P2(-1,0)) ≈ T(π/2)
  @test ∠(P2(0,-1),P2(0,0),P2(1,0)) ≈ T(π/2)
  @test ∠(P2(1,0),P2(0,0),P2(0,-1)) ≈ T(-π/2)
  @test ∠(P2(0,-1),P2(0,0),P2(-1,0)) ≈ T(-π/2)
  @test ∠(P2(-1,0),P2(0,0),P2(0,-1)) ≈ T(π/2)
end


# This file is part of IntegerSequences.
# Copyright Peter Luschny. License is MIT.

(@__DIR__) ∉ LOAD_PATH && push!(LOAD_PATH, (@__DIR__))

module GeneralizedFibonacci
using Nemo, SeqUtils

export ModuleGeneralizedFibonacci
export Multinacci, V309896, V006053, V188021, V231181

"""

* The classical Fibonacci numbers are exported from the module Fibonacci. See I000045, F000045, L000045, V000045, R000045 and is000045.

* Fibonacci(n) is defined as the number of compositions of n with no part equal to 1. They are the special case Fibonacci(n) = Multinacci(2, n).

```
[m\n] 0  1  2   3   4   5    6    7     8     9    10    11
----------------------------------------------------------------
[0]   1, 0, 0,  0,  0,  0,   0,   0,    0,    0,    0,    0, ...
[1]   1, 1, 1,  1,  1,  1,   1,   1,    1,    1,    1,    1, ...
[2]   1, 1, 2,  3,  5,  8,  13,  21,   34,   55,   89,  144, ...
[3]   1, 1, 3,  4,  9, 14,  28,  47,   89,  155,  286,  507, ...
[4]   1, 1, 4,  5, 14, 20,  48,  75,  165,  274,  571,  988, ...
[5]   1, 1, 5,  6, 20, 27,  75, 110,  275,  429, 1001, 1637, ...
[6]   1, 1, 6,  7, 27, 35, 110, 154,  429,  637, 1638, 2548, ...
```

*  Multinacci, V309896, V006053, V188021, V231181
"""
const ModuleGeneralizedFibonacci = ""

"""

Return the ``n``-th term of the ``m``-th Multinacci sequence.
"""
function Multinacci(m::Int, n::Int)::fmpz
    haskey(CacheMultinacci, (m, n)) && return CacheMultinacci[(m, n)]
    n <  0 && return fmpz(0); n == 0 && return fmpz(1); m == 0 && return fmpz(m^n)
    a = sum((-1)^j*binomial(ZZ(m-1-j), ZZ(j ) )*Multinacci(m, n-1-2*j) for j ∈ 0:div(m-1,2))
    b = sum((-1)^j*binomial(ZZ(m-1-j), ZZ(j+1))*Multinacci(m, n-2-2*j) for j ∈ 0:div(m-2,2))
    r = a + b
    CacheMultinacci[(m, n)] = r
end

"""

Return the ``k``-th term of the ``n``-th Multinacci sequence.
"""
V309896(n, k) = Multinacci(n, k)

const CacheMultinacci = Dict{Tuple{Int, Int}, fmpz}()

"""

Return the ``n``-th term of the third Multinacci sequence.
"""
V006053(n) = Multinacci(3, n)

"""

Return the ``n``-th term of the fourth Multinacci sequence.
"""
V188021(n) = Multinacci(4, n)

"""

Return the ``n``-th term of the fiveth Multinacci sequence.
"""
V231181(n) = Multinacci(5, n)

#START-TEST-########################################################

using Test

function test()
    @testset "Multinacci" begin
        @test Multinacci(0, 0) == 1
        @test [Multinacci(2, n) for n ∈ 0:9] == [1, 1, 2, 3, 5,  8, 13, 21, 34,  55]
        @test [V006053(n) for n ∈ 0:9] == [1, 1, 3, 4, 9, 14, 28, 47, 89, 155]
        @test [V188021(n) for n ∈ 0:9] == [1, 1, 4, 5, 14, 20, 48, 75, 165, 274]
        @test [V231181(n) for n ∈ 0:9] == [1, 1, 5, 6, 20, 27, 75, 110, 275,429]
        @test [Multinacci(6, n) for n ∈ 0:9] == [1, 1, 6, 7, 27, 35, 110, 154,429,637]
        @test [Multinacci(7, n) for n ∈ 0:9] == [1, 1, 7, 8, 35, 44, 154, 208,637,910]
        @test [Multinacci(8, n) for n ∈ 0:9] == [1, 1, 8, 9, 44, 54,208,273,910, 1260]
    end
end

function demo()
    for n ∈ 0:9
        print("$n ")
        [Multinacci(n, k) for k ∈ 0:9] |> Println
    end
end

function perf()
    @time (for n ∈ 0:1000 Multinacci(3, n) end)
end

function main()
    test()
    demo()
    perf()
end

main()

end # module