abstract type InnerProduct end

dot_arg(f,p::InnerProduct) = evaluate(f,p)
dot_arg(f::AbstractArray,p::InnerProduct) = [f(x) for x in evalpoints(p), f in f]
dot_arg(f::AbstractBasis,p::InnerProduct) = Matrix(f,evalpoints(p))

LinearAlgebra.norm(f, p::InnerProduct) = norm(dot_arg(f,p), p)
LinearAlgebra.dot(f,g,p::InnerProduct) = dot(dot_arg.((f,g),Ref(p))...,p)

"""
    WeightedL2(x,w) <: InnerProduct

# Example
```
p = WeightedL2([0,1],[2,3])
norm(identity,p) -> sqrt(0^2*2 + 1^2*3)
```
"""
struct WeightedL2{X,W} <: InnerProduct
    points::X
    weights::W
end

evalpoints(p::WeightedL2) = p.points
evaluate(f,p::WeightedL2) = f.(evalpoints(p))
dot_arg(f::AbstractArray{<:Number},p::WeightedL2) = f

LinearAlgebra.norm(f::AbstractVecOrMat{<:Number}, p::WeightedL2) = sqrt(dot(f,f,p))

function LinearAlgebra.dot(
    f::AbstractVecOrMat{<:Number},
    g::AbstractVecOrMat{<:Number},
    p::WeightedL2
)
    w = p.weights
    # Workaround for https://github.com/JuliaLang/julia/issues/35424
    length(w) == 0 && return zero(eltype(f))*zero(eltype(w))*zero(eltype(g))
    return f'*Diagonal(w)*g
end

using FastGaussQuadrature

LegendreInnerProduct(n) = WeightedL2(gausslegendre(n)...)
JacobiInnerProduct(n,a,b) = WeightedL2(gaussjacobi(n,a,b)...)


using GR
using LaTeXStrings

𝒩(μ, σ) = 1 / (σ * √(2π)) * exp.(-0.5 * ((x .- μ) / σ) .^ 2)

x = LinRange(-5, 5, 500);
y = hcat(𝒩(0, √0.2), 𝒩(0, √1), 𝒩(0, √5), 𝒩(-2, √0.5));

plot(x, y, xlabel=L"\mathcal{X}", ylabel=L"\mathcal{N}(\mu,\,\sigma^{2})", title=L"\frac{1}{\sigma\sqrt{2\pi}} e^{-\frac{1}{2} \left({\frac{x-\mu}{\sigma}}\right)^2}", labels=(L"\mu=0, \sigma^2=0.2", L"\mu=0, \sigma^2=1", L"\mu=0, \sigma^2=5", L"\mu={-2}, \sigma^2=0.5"), xlim=(-5.2, 5.2), ylim=(-0.05, 1.05), linewidth=3)


using ValidatedNumerics, ValidatedNumerics.RootFinding
using Test

@testset "Bisection tests" begin
    @testset "Single variable" begin
        f(x) = sin(x) + cos(x/2) + x/5

        roots = bisection(f, -∞..∞, tolerance=1e-3)
        roots2 = [root.interval for root in roots]

        @test roots2 ==  [  Interval(-0.8408203124999999, -0.8398437499999999),
                            Interval(3.6425781249999996, 3.6435546874999996),
                            Interval(6.062499999999999, 6.063476562499999)
                         ]

    end

    @testset "Two variables" begin
        @intervalbox g(x, y) = (x^2 + y^2 - 1, y - x)

        X = (-∞..∞) × (-∞..∞)
        roots = bisection(g, X, tolerance=1e-3)
        roots2 = [root.interval for root in roots]

        @test roots2 == [IntervalBox(Interval(-0.7080078124999999, -0.7070312499999999), Interval(-0.7080078124999999, -0.7070312499999999)),
 IntervalBox(Interval(-0.7080078124999999, -0.7070312499999999), Interval(-0.7070312499999999, -0.7060546874999999)),
 IntervalBox(Interval(-0.7070312499999999, -0.7060546874999999), Interval(-0.7080078124999999, -0.7070312499999999)),
 IntervalBox(Interval(0.7060546874999999, 0.7070312499999999), Interval(0.7070312499999999, 0.7080078124999999)),
 IntervalBox(Interval(0.7070312499999999, 0.7080078124999999), Interval(0.7060546874999999, 0.7070312499999999)),
 IntervalBox(Interval(0.7070312499999999, 0.7080078124999999), Interval(0.7070312499999999, 0.7080078124999999))]

    end

end


"""
    uniform_bsplines_eval_basis( spline_degree, normalized_offset, bspl )

# UNIFORM B-SPLINE FUNCTIONS

## Evaluate all non vanishing uniform B-Splines in unit cell.

Returns an array with the values of the b-splines of the
requested degree, evaluated at a given cell offset. The cell size is
normalized between 0 and 1, thus the offset given must be a number
between 0 and 1.

Output:

```math
bspl(1:d+1)= B_d(-(d+1)/2+d+x),...,B_d(-(d+1)/2+x)
```

with ``d``=`spline_degree` and ``x``=`normalized_offset`
where ``B_d=B_{d-1}*B_0`` and ``B_0=1_[-1/2,1/2]`` and `*` is convolution
the following FORTRAN code can be used for comparison with
[deboor](http://pages.cs.wisc.edu/~deboor/)

```fortran
do i=-d,d+1
    t(i+d+1)=real(i,8)
end do
call bsplvb(t,d+1,1,normalized_offset,d+1,out)
```

We also have the property (from the symmetry of the B-spline)
```math
out[1:d+1]= B_d(-(d+1)/2+xx),...,B_d(-(d+1)/2+d+xx),...,
```
where ``xx=1-`` `normalized_offset`

"""
function uniform_bsplines_eval_basis(spline_degree::Int, normalized_offset::Float64)
    @assert spline_degree >= 0
    @assert normalized_offset >= 0.0
    @assert normalized_offset <= 1.0

    bspl = zeros(Float64, spline_degree + 1)

    bspl[1] = 1.0
    @inbounds for j in 1:spline_degree
        xx = -normalized_offset
        j_real = Float64(j)
        inv_j = 1.0 / j_real
        saved = 0.0
        for r in 0:(j - 1)
            xx = xx + 1.0
            temp = bspl[r + 1] * inv_j
            bspl[r + 1] = saved + xx * temp
            saved = (j_real - xx) * temp
        end
        bspl[j + 1] = saved
    end

    return bspl
end

function uniform_bsplines_eval_basis!(
    bspl::Vector{Float64}, spline_degree::Int, normalized_offset::Float64
)
    bspl[1] = 1.0
    @inbounds for j in 1:spline_degree
        xx = -normalized_offset
        j_real = Float64(j)
        inv_j = 1.0 / j_real
        saved = 0.0
        for r in 0:(j - 1)
            xx = xx + 1.0
            temp = bspl[r + 1] * inv_j
            bspl[r + 1] = saved + xx * temp
            saved = (j_real - xx) * temp
        end
        bspl[j + 1] = saved
    end
end

function uniform_bsplines_eval_basis!(
    bspl::Array{Float64,2},
    spline_degree::Int,
    normalized_offset_x::Float64,
    normalized_offset_y::Float64,
)
    bspl[1, 1] = 1.0
    bspl[1, 2] = 1.0
    @inbounds for j in 1:spline_degree
        xx = -normalized_offset_x
        yy = -normalized_offset_y
        j_real = Float64(j)
        inv_j = 1.0 / j_real
        saved_x = 0.0
        saved_y = 0.0
        for r in 0:(j - 1)
            xx = xx + 1
            yy = yy + 1
            temp_x = bspl[r + 1, 1] * inv_j
            temp_y = bspl[r + 1, 2] * inv_j
            bspl[r + 1, 1] = saved_x + xx * temp_x
            bspl[r + 1, 2] = saved_y + yy * temp_y
            saved_x = (j_real - xx) * temp_x
            saved_y = (j_real - yy) * temp_y
        end
        bspl[j + 1, 1] = saved_x
        bspl[j + 1, 2] = saved_y
    end
end

export eval_uniform_periodic_spline_curve

"""
    eval_uniform_periodic_spline_curve( degree, scoef )

Evaluate uniform periodic spline curve defined by coefficients scoef at
knots (which are the grid points)

"""
function eval_uniform_periodic_spline_curve(degree::Int, scoef::Vector{Float64})

    # get bspline values at knots
    bspl = uniform_bsplines_eval_basis(degree, 0.0)
    n = length(scoef)
    sval = similar(scoef)
    for i in 1:n
        val = 0.0
        for j in 1:degree
            imj = mod1(i - j + n, n)
            val = val + bspl[j] * scoef[imj]
        end
        sval[i] = val
    end

    return sval
end


@doc """
      Heun's Method (with iteration)
      Supported by:
      -PhD Kelvyn B. Sánchez kelvyn.baruc@gmail.com
 """ ->
function heun(f, tspan, y0; h = 1, AbsTol = 0.00001)

t = tspan[1]:h:tspan[end]
n = length(t)

      if t[n] < tspan[end]
            t = [t, tspan[end]]
            n += 1
      end

y1 = f(t[1],y0)
eq = length(y1) # Test the number of equations in f
y = zeros(Int(n),eq)
y[1,:] = y0
y1 = zeros(Int(n),eq)
y2 = zeros(Int(n),eq)
y3 = zeros(Int(n),eq)

	for i=1:n-1 # Solve for each x
            for j=1:eq # Solve for each equation in f
                  s1 = f(t[i],y[i,:])
                  y1[i,j] = y[i,j] + s1[j] * (t[i+1] - t[i])

                  s2 = f(t[i+1],y1[i,:])
                  y2[i,j] = y[i,j] + ((s1[j] + s2[j])/2) * (t[i+1] - t[i])

                  err = 100
                  while err >= AbsTol
                        s1 = f(t[i],y[i,:])
                        s2 = f(t[i+1],y2[i,:])
                        y3[i,j] = y[i,j] + ((s1[j] + s2[j]) / 2) * (t[i+1] - t[i])
                        err = abs((y3[i,j] - y2[i,j])/y3[i,j]) * 100
                        y2[i,j] = y3[i,j]
                        y[i+1,j] = y3[i,j]
                  end
            end
	end
	return t, y
end


using Distributions
using LightBayes

function simulate(r::Float64, n::Int)
    # Simulate Δb for H0 & H1
    #
    # generate two clusters of samples, respectively centered
    # at (r, 0) and (-r, 0), where each cluster contains n
    # samples.
    #
    # This function computes the values of Δb, respectively
    # for the following two hypotheses:
    #
    # - H0: all samples are generated by a single component
    # - H1: each cluster of samples are generated by a different
    #       component
    #
    # Prior model: θ ~ N(0, σ0^2) with σ0 = 10.0
    # Likelihood:  x ~ N(θ, σ^2) with σ = 1.0
    #

    # data
    n = 1000
    X1 = [r, 0.0] .+ rand(2, 1000)
    X2 = [-r, 0.0] .+ rand(2, 1000)

    X = hcat(X1, X2)
    @assert size(X) == (2, 2n)

    # prior
    β0 = zeros(2)
    σ0 = 10.0
    κ0 = inv(σ0^2)
    pri = MvNormalCanon(β0, κ0)

    # likelihood model
    σ = 1.0
    g = IsoGaussModel(σ)

    # H0: single component
    pp0 = posterior(pri, suffstats(g, X))
    Δb0 = logpar(pp0) - logpar(pri)

    # H1: two components:
    pp1 = posterior(pri, suffstats(g, X1))
    pp2 = posterior(pri, suffstats(g, X2))
    Δb1 = (logpar(pp1) - logpar(pri)) + (logpar(pp2) - logpar(pri))

    return (Δb0, Δb1)
end

# main script

n = 1000

for r in [0.0, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0]
    Δb0, Δb1 = simulate(r, n)
    @printf("With r = %5.1f:  H0.Δb = %8.2f,  H1.Δb = %8.2f\n", r, Δb0, Δb1)
end


"""Compute the entropy of the specified dataset."""
function entropy(dataset::AbstractArray{T, 2}, target::Integer,
                c::Integer=2) where {T}
    class = dataset[:, target]
    n = length(class)
    uniques = unique(class)

    return -mapreduce(u -> begin
        subset = dataset[class .== u, :]
        p = size(subset, 1) / n
        return p * log2(p)
    end, +, uniques) / log2(c)
end

entropy(dataset::AbstractArray{T, 2}, target::String,
        header::AbstractArray{String}, c::Integer=2) where {T} =
    entropy(dataset, something(findfirst(isequal(target), header), 0), c)

"""Compute the information gain of the specified attribute."""
function information_gain(dataset::AbstractArray{T, 2}, attribute::Integer,
                        target::Integer, c::Integer=2) where {T}
    n = size(dataset, 1)
    attr = dataset[:, attribute]
    uniques = unique(attr)

    attr_sub_entropy = mapreduce(u -> begin
        subset = dataset[attr .== u, :]
        return size(subset, 1) / n * entropy(subset, target, c)
    end, +, uniques)

    return entropy(dataset, target, c) - attr_sub_entropy
end

information_gain(dataset::AbstractArray{T, 2}, attribute::String,
                target::String, header::AbstractArray{String},
                c::Integer=2) where {T} =
    information_gain(dataset,
            something(findfirst(isequal(attribute), header), 0),
            something(findfirst(isequal(target), header), 0), c)

information_gain(dataset::AbstractArray{T, 2}, attribute::String,
                target::Integer, header::AbstractArray{String},
                c::Integer=2) where {T} =
    information_gain(dataset,
            something(findfirst(isequal(attribute), header), 0), target, c)

information_gain(dataset::AbstractArray{T, 2}, attribute::Integer,
                target::String, header::AbstractArray{String},
                c::Integer=2) where {T} =
    information_gain(dataset, attribute,
            something(findfirst(isequal(target), header), 0), c)


immutable Binomial <: DiscreteUnivariateDistribution
    size::Int
    prob::Float64
    function Binomial(n::Real, p::Real)
    	if n <= 0
	    	error("size must be positive")
	    else
	    	if 0.0 <= p <= 1.0
	    		new(int(n), float64(p))
	    	else
	    		error("prob must be in [0, 1]")
			end
	    end
	end
end

Binomial(size::Integer) = Binomial(size, 0.5)
Binomial() = Binomial(1, 0.5)

@_jl_dist_2p Binomial binom

entropy(d::Binomial) = d.size * (-xlogx(1.0 - d.prob) - xlogx(d.prob))

insupport(d::Binomial, x::Number) = isinteger(x) && 0 <= x <= d.size

kurtosis(d::Binomial) = (1.0 - 2.0 * d.prob * (1.0 - d.prob)) / var(d)

mean(d::Binomial) = d.size * d.prob

median(d::Binomial) = iround(d.size * d.prob)

# TODO: May need to subtract 1 sometimes
modes(d::Binomial) = iround((d.size + 1.0) * d.prob)

function mgf(d::Binomial, t::Real)
	n, p = d.size, d.prob
	return (1.0 - p + p * exp(t))^n
end

function cf(d::Binomial, t::Real)
	n, p = d.size, d.prob
	return (1.0 - p + p * exp(im * t))^n
end

modes(d::Binomial) = iround([d.size * d.prob])

# TODO: rand() is totally screwed up

skewness(d::Binomial) = (1.0 - 2.0 * d.prob) / std(d)

var(d::Binomial) = d.size * d.prob * (1.0 - d.prob)


using Plots
using Distributions
using Random
gr(size=(300,300))

function ellipse!(μ,Σ;kw...)
    for i=1:3
        plot!([tuple((i*sqrt(Σ)*[cos(t),sin(t)]+μ)...) for t in range(0,stop=2π,length=100)];
              seriestype=:shape,linealpha=0.5,linecolor=:gray,fillalpha=[0.4,0.2,0.1][i],kw...)
    end
    current()
end

function mixtureplot(X₁,X₂,μ₀,Σ₀,μ₁,Σ₁,Π)
    scatter(X₁,X₂;fillalpha=0.5,markerstrokewidth=0.5,
        marker_z=Π,mc=ColorGradient([:lightblue,:orange]),
        colorbar=:false)
    ellipse!(μ₀,Σ₀,fillcolor=:lightblue)
    ellipse!(μ₁,Σ₁,fillcolor=:orange)
    plot!(;bg=:transparent,xlims=(-4.5,8.5),ylims=(-3.5,10.5),
        leg=false,ticks=:none,ratio=:equal)
end

function scalein(x)
    (x-1/2)^101/(0.5^101)
end

Random.seed!(123);
n = 100
α = 0.4
𝒩₀ = MvNormal([1,1],[2.0 1.0; 1.0 2.0])
𝒩₁ = MvNormal([3.0,7.0],[1.5 0; 0 0.5])
X₁ = zeros(n)
X₂ = zeros(n)
Z = zeros(Bool,n)
for i=1:n
    Z[i] = rand(Bernoulli(α))
    X₁[i],X₂[i] = Z[i] ? rand(𝒩₁) : rand(𝒩₀)
end

α = 0.6
μ₀ = [3.0,3.0]
μ₁ = [1.0,6.0]
Σ₀ = 1.0*Matrix(I, 2, 2)
Σ₁ = 1.0*Matrix(I, 2, 2)
mixtureplot(X₁,X₂,μ₀,Σ₀,μ₁,Σ₁)

Π = [α*pdf(MvNormal(μ₁,Σ₁),[x₁,x₂]) /
       ((1-α)*pdf(MvNormal(μ₀,Σ₀),[x₁,x₂]) +
        α*pdf(MvNormal(μ₁,Σ₁),[x₁,x₂])) for (x₁,x₂) in zip(X₁,X₂)];

α = sum(Π)/n
μ₀ = [(1 .- Π) ⋅ X₁, (1 .- Π) ⋅ X₂] / sum(1 .- Π)
μ₁ = [Π ⋅ X₁, Π ⋅ X₂] / sum(Π)
Σ₀ = Matrix(Hermitian(sum((1-π)*([x₁,x₂] - μ₀) * ([x₁,x₂] - μ₀)' for (x₁,x₂,π) in zip(X₁,X₂,Π))/sum(1 .- Π)))
Σ₁ = Matrix(Hermitian(sum(π*([x₁,x₂] - μ₁) * ([x₁,x₂] - μ₁)' for (x₁,x₂,π) in zip(X₁,X₂,Π))/sum(Π)))
Π = [α*pdf(MvNormal(μ₁,Σ₁),[x₁,x₂]) /
       ((1-α)*pdf(MvNormal(μ₀,Σ₀),[x₁,x₂]) +
        α*pdf(MvNormal(μ₁,Σ₁),[x₁,x₂])) for (x₁,x₂) in zip(X₁,X₂)];

mixtureplot(X₁,X₂,μ₀,Σ₀,μ₁,Σ₁,Π)


using OrdinaryDiffEq
using StaticArrays
using Plots

"""
    Duffing

A Duffing equation with periodic forcing of the form

```math
mx″ + cx′ + kx + μx³ = F₀cos(ωt)
```
"""
Base.@kwdef mutable struct Duffing
    m::Float64 = 1.0
    c::Float64 = 0.01
    k::Float64 = 2.0
    μ::Float64 = 1
    F₀::Float64 = 0.05
    ω::Float64 = 1.0
end

"""
    rhs_dt(u, p::Duffing, t)

Return the right-hand side of the Duffing.
"""
function rhs_dt(u, p::Duffing, t)
    x = u[1]
    x′ = u[2]
    f = p.F₀ * cos(p.ω * t)
    # Use a static vector (SVector) for speed; it's not necessary for any other reason
    return SVector(x′, (f - p.c * x′ - p.k * x - p.μ * x^3) / p.m)
end

"""
    integrate(u, p::Duffing, t)

Return a solution of the Duffing equation over the time interval `t` starting at the initial
condition `u` with the parameters of the Duffing equation taken from `p`.
"""
function integrate(u, p::Duffing, t)
    # Explicitly convert everything to Float64 in case integer values are used
    # Use SVectors for speed (not necessary)
    prob = ODEProblem(rhs_dt, SVector{2, Float64}(u), Float64.(t), p)
    return solve(prob, Tsit5())
end

"""
    example()

Run an example Duffing equation simulation and plot the results.
"""
function example()
    p = Duffing()
    sol = integrate([0, 0], p, (0, 100))
    plot(sol.t, sol[1,:])
end

# Code to benchmark the timings

# using BenchmarkTools
# p = Duffing()
# @benchmark integrate($([0, 0]), $p, (0, 100))

# Output on my laptop

# BenchmarkTools.Trial:
#   memory estimate:  60.34 KiB
#   allocs estimate:  462
#   --------------
#   minimum time:     73.100 μs (0.00% GC)
#   median time:      75.700 μs (0.00% GC)
#   mean time:        90.955 μs (3.73% GC)
#   maximum time:     3.807 ms (96.36% GC)
#   --------------
#   samples:          10000
#   evals/sample:     1


# 1D diffusion problem

# Packages and inclusions
using ModelingToolkit,DiffEqOperators,LinearAlgebra,Test,OrdinaryDiffEq, DomainSets
using ModelingToolkit: Differential

# Tests
@testset "Test 00: Dt(u(t,x)) ~ Dxx(u(t,x))" begin
    # Method of Manufactured Solutions
    u_exact = (x,t) -> exp.(-t) * cos.(x)

    # Parameters, variables, and derivatives
    @parameters t x
    @variables u(..)
    Dt = Differential(t)
    Dxx = Differential(x)^2

    # 1D PDE and boundary conditions
    eq  = Dt(u(t,x)) ~ Dxx(u(t,x))
    bcs = [u(0,x) ~ cos(x),
           u(t,0) ~ exp(-t),
           u(t,Float64(π)) ~ -exp(-t)]

    # Space and time domains
    domains = [t ∈ Interval(0.0,1.0),
               x ∈ Interval(0.0,Float64(π))]

    # PDE system
    pdesys = PDESystem(eq,bcs,domains,[t,x],[u(t,x)])

    # Method of lines discretization
    dx = range(0.0,Float64(π),length=30)
    order = 2
    discretization = MOLFiniteDifference([x=>dx],t)
    discretization_edge = MOLFiniteDifference([x=>dx],t;grid_align=edge_align)
    # Explicitly specify order of centered difference
    discretization_centered = MOLFiniteDifference([x=>dx],t;centered_order=order)
    # Higher order centered difference
    discretization_centered_order4 = MOLFiniteDifference([x=>dx],t;centered_order=4)

    for disc in [discretization, discretization_edge, discretization_centered, discretization_centered_order4]
        # Convert the PDE problem into an ODE problem
        prob = discretize(pdesys,disc)

        # Solve ODE problem
        sol = solve(prob,Tsit5(),saveat=0.1)

        if disc.grid_align == center_align
            x = dx[2:end-1]
        else
            x = (dx[1:end-1]+dx[2:end])/2
        end
        t = sol.t

        # Test against exact solution
        for i in 1:length(sol)
            exact = u_exact(x, t[i])
            u_approx = sol.u[i]
            @test all(isapprox.(u_approx, exact, atol=0.01))
        end
    end
end

@testset "Test 01: Dt(u(t,x)) ~ D*Dxx(u(t,x))" begin
    # Parameters, variables, and derivatives
    @parameters t x D
    @variables u(..)
    Dt = Differential(t)
    Dxx = Differential(x)^2

    # 1D PDE and boundary conditions
    eq  = Dt(u(t,x)) ~ D*Dxx(u(t,x))
    bcs = [u(0,x) ~ -x*(x-1)*sin(x),
           u(t,0) ~ 0.0,
           u(t,1) ~ 0.0]

    # Space and time domains
    domains = [t ∈ Interval(0.0,1.0),
               x ∈ Interval(0.0,1.0)]

    # PDE system
    pdesys = PDESystem(eq,bcs,domains,[t,x],[u(t,x)],[D=>10.0])

    # Method of lines discretization
    dx = 0.1
    order = 2
    discretization = MOLFiniteDifference([x=>dx],t)

    # Convert the PDE problem into an ODE problem
    prob = discretize(pdesys,discretization)

    # Solve ODE problem
    sol = solve(prob,Tsit5(),saveat=0.1)

    # Test
    n = size(sol,1)
    t_f = size(sol,3)
    @test sol[end] ≈ zeros(n) atol = 0.001;
end

# @test_set "Test 02: Dt(u(t,x)) ~ Dx(D(t,x))*Dx(u(t,x))+D(t,x)*Dxx(u(t,x))" begin
@test_broken begin
    # Parameters, variables, and derivatives
    @parameters t x
    @variables u(..) D(..)
    Dt = Differential(t)
    Dx = Differential(x)
    Dxx = Differential(x)^2

    D = (t,x) -> 0.999 + 0.001 * t * x
    DxD = expand_derivatives(Dx(D(t,x)))

    # 1D PDE and boundary conditions

    eq  = [ Dt(u(t,x)) ~ DxD*Dx(u(t,x))+D(t,x)*Dxx(u(t,x)),]

    bcs = [u(0,x) ~ -x*(x-1)*sin(x),
           u(t,0) ~ 0.0,
           u(t,1) ~ 0.0]

    # Space and time domains
    domains = [t ∈ Interval(0.0,1.0),
               x ∈ Interval(0.0,1.0)]

    # PDE system
    pdesys = PDESystem(eq,bcs,domains,[t,x],[u(t,x)])

    # Method of lines discretization
    dx = 0.1
    order = 2
    discretization = MOLFiniteDifference([x=>dx],t)

    # Convert the PDE problem into an ODE problem
    prob = discretize(pdesys,discretization)

    # Solve ODE problem
    sol = solve(prob,Tsit5(),saveat=0.1)

    # Test
    n = size(sol,1)
    t_f = size(sol,3)
    @test_broken sol[:,1,t_f] ≈ zeros(n) atol=0.01;
end

@testset "Test 03: Dt(u(t,x)) ~ Dxx(u(t,x)), homogeneous Neumann BCs" begin
    # Method of Manufactured Solutions
    u_exact = (x,t) -> exp.(-t) * cos.(x)

    # Parameters, variables, and derivatives
    @parameters t x
    @variables u(..)
    Dt = Differential(t)
    Dx = Differential(x)
    Dxx = Differential(x)^2

    # 1D PDE and boundary conditions
    eq  = Dt(u(t,x)) ~ Dxx(u(t,x))
    bcs = [u(0,x) ~ cos(x),
           Dx(u(t,0)) ~ 0,
           Dx(u(t,Float64(pi))) ~ 0]

    # Space and time domains
    domains = [t ∈ Interval(0.0,1.0),
               x ∈ Interval(0.0,Float64(pi))]

    # PDE system
    pdesys = PDESystem(eq,bcs,domains,[t,x],[u(t,x)])

    # Method of lines discretization
    dx = range(0.0,Float64(π),length=300)
    order = 2
    discretization = MOLFiniteDifference([x=>dx],t)
    discretization_edge = MOLFiniteDifference([x=>dx],t;grid_align=center_align)

    # Convert the PDE problem into an ODE problem
    for disc in [discretization, discretization_edge]
        prob = discretize(pdesys,disc)

        # Solve ODE problem
        sol = solve(prob,Tsit5(),saveat=0.1)

        if disc.grid_align == center_align
            x_sol = dx[2:end-1]
        else
            x_sol = (dx[1:end-1]+dx[2:end])/2
        end
        t_sol = sol.t

        # Test against exact solution
        for i in 1:length(sol)
            exact = u_exact(x_sol, t_sol[i])
            u_approx = sol.u[i]
            @test all(isapprox.(u_approx, exact, atol=0.01))
            @test sum(u_approx) ≈ 0 atol=1e-10
        end
    end
end

@testset "Test 03a: Dt(u(t,x)) ~ Dxx(u(t,x)), Neumann BCs" begin
    # Method of Manufactured Solutions
    u_exact = (x,t) -> exp.(-t) * sin.(x)

    # Parameters, variables, and derivatives
    @parameters t x
    @variables u(..)
    Dt = Differential(t)
    Dx = Differential(x)
    Dxx = Differential(x)^2

    # 1D PDE and boundary conditions
    eq  = Dt(u(t,x)) ~ Dxx(u(t,x))
    bcs = [u(0,x) ~ sin(x),
           Dx(u(t,0)) ~ exp(-t),
           Dx(u(t,Float64(pi))) ~ -exp(-t)]

    # Space and time domains
    domains = [t ∈ Interval(0.0,1.0),
               x ∈ Interval(0.0,Float64(pi))]

    # PDE system
    pdesys = PDESystem(eq,bcs,domains,[t,x],[u(t,x)])

    # Method of lines discretization
    dx = range(0.0,Float64(π),length=30)
    order = 2
    discretization = MOLFiniteDifference([x=>dx],t)
    discretization_edge = MOLFiniteDifference([x=>dx],t;grid_align=edge_align)

    # Convert the PDE problem into an ODE problem
    for disc ∈ [discretization, discretization_edge]
        prob = discretize(pdesys,disc)

        # Solve ODE problem
        sol = solve(prob,Tsit5(),saveat=0.1)

        if disc.grid_align == center_align
            x = dx[2:end-1]
        else
            x = (dx[1:end-1]+dx[2:end])/2
        end
        t = sol.t

        # Test against exact solution
        # exact integral based on Neumann BCs
        integral_u_exact = t -> sum(sol.u[1] * dx[2]) + 2 * (exp(-t) - 1)
        for i in 1:length(sol)
            exact = u_exact(x, t[i])
            u_approx = sol.u[i]
            @test all(isapprox.(u_approx, exact, atol=0.01))
            # test mass conservation
            integral_u_approx = sum(u_approx * dx[2])
            @test integral_u_exact(t[i]) ≈ integral_u_approx atol=1e-13
        end
    end
end

@testset "Test 04: Dt(u(t,x)) ~ Dxx(u(t,x)), Neumann + Dirichlet BCs" begin
    # Method of Manufactured Solutions
    u_exact = (x,t) -> exp.(-t) * sin.(x)

    # Parameters, variables, and derivatives
    @parameters t x
    @variables u(..)
    Dt = Differential(t)
    Dx = Differential(x)
    Dxx = Differential(x)^2

    # 1D PDE and boundary conditions
    eq  = Dt(u(t,x)) ~ Dxx(u(t,x))
    bcs = [u(0,x) ~ sin(x),
           u(t,0) ~ 0.0,
           Dx(u(t,Float64(pi))) ~ -exp(-t)]

    # Space and time domains
    domains = [t ∈ Interval(0.0,1.0),
               x ∈ Interval(0.0,Float64(pi))]

    # PDE system
    pdesys = PDESystem(eq,bcs,domains,[t,x],[u(t,x)])

    # Method of lines discretization
    dx = range(0.0,Float64(π),length=30)
    order = 2
    discretization = MOLFiniteDifference([x=>dx],t)

    # Convert the PDE problem into an ODE problem
    prob = discretize(pdesys,discretization)

    # Solve ODE problem
    sol = solve(prob,Tsit5(),saveat=0.1)
    x = dx[2:end-1]
    t = sol.t

    # Test against exact solution
    for i in 1:length(sol)
        exact = u_exact(x, t[i])
        u_approx = sol.u[i]
        @test all(isapprox.(u_approx, exact, atol=0.01))
    end
end

@testset "Test 05: Dt(u(t,x)) ~ Dxx(u(t,x)), Robin BCs" begin
    # Method of Manufactured Solutions
    u_exact = (x,t) -> exp.(-t) * sin.(x)

    # Parameters, variables, and derivatives
    @parameters t x
    @variables u(..)
    Dt = Differential(t)
    Dx = Differential(x)
    Dxx = Differential(x)^2

    # 1D PDE and boundary conditions
    eq  = Dt(u(t,x)) ~ Dxx(u(t,x))
    bcs = [u(0,x) ~ sin(x),
           u(t,-1.0) + 3Dx(u(t,-1.0)) ~ exp(-t) * (sin(-1.0) + 3cos(-1.0)),
           4u(t,1.0) + Dx(u(t,1.0)) ~ exp(-t) * (4sin(1.0) + cos(1.0))]

    # Space and time domains
    domains = [t ∈ Interval(0.0,1.0),
               x ∈ Interval(-1.0,1.0)]

    # PDE system
    pdesys = PDESystem(eq,bcs,domains,[t,x],[u(t,x)])

    # Method of lines discretization
    dx = 0.01
    order = 2
    discretization = MOLFiniteDifference([x=>dx],t)
    discretization_edge = MOLFiniteDifference([x=>dx],t;grid_align=edge_align)

    for disc ∈ [discretization, discretization_edge]
        # Convert the PDE problem into an ODE problem
        prob = discretize(pdesys,disc)

        # Solve ODE problem
        sol = solve(prob,Tsit5(),saveat=0.1)
        x = (-1:dx:1)
        if disc.grid_align == center_align
            x = x[2:end-1]
        else
            x = (x[1:end-1].+x[2:end])/2
        end
        t = sol.t

        # Test against exact solution
        for i in 1:length(sol)
            exact = u_exact(x, t[i])
            u_approx = sol.u[i]
            @test all(isapprox.(u_approx, exact, atol=0.1))
        end
    end
end

@testset "Test 06: Dt(u(t,x)) ~ Dxx(u(t,x)), time-dependent Robin BCs" begin
    # Method of Manufactured Solutions
    u_exact = (x,t) -> exp.(-t) * sin.(x)

    # Parameters, variables, and derivatives
    @parameters t x
    @variables u(..)
    Dt = Differential(t)
    Dx = Differential(x)
    Dxx = Differential(x)^2

    # 1D PDE and boundary conditions
    eq  = Dt(u(t,x)) ~ Dxx(u(t,x))
    bcs = [u(0,x) ~ sin(x),
          t^2 * u(t,-1.0) + 3Dx(u(t,-1.0)) ~ exp(-t) * (t^2 * sin(-1.0) + 3cos(-1.0)),
          4u(t,1.0) + t * Dx(u(t,1.0)) ~ exp(-t) * (4sin(1.0) + t * cos(1.0))]

    # Space and time domains
    domains = [t ∈ Interval(0.0,1.0),
               x ∈ Interval(-1.0,1.0)]

    # PDE system
    pdesys = PDESystem(eq,bcs,domains,[t,x],[u(t,x)])

    # Method of lines discretization
    dx = 0.01
    order = 2
    discretization = MOLFiniteDifference([x=>dx],t)

    # Convert the PDE problem into an ODE problem
    prob = discretize(pdesys,discretization)

    # Solve ODE problem
    sol = solve(prob,Rodas4(),reltol=1e-6,saveat=0.1)

    x = (-1:dx:1)
    t = sol.t

    # Test against exact solution
    for i in 1:length(sol)
       exact = u_exact(x, t[i])
       # Due to structural simplification
       # [u2 -> u(n-1), u(1), u(n)]
       # Will be fixed by sol[u]
       u_approx = [sol.u[i][end-1];sol.u[i][1:end-2];sol.u[i][end]]
       @test all(isapprox.(u_approx, exact, atol=0.01))
    end
end

@testset "Test 07: Dt(u(t,r)) ~ 1/r^2 * Dr(r^2 * Dr(u(t,r))) (Spherical Laplacian)" begin
    # Method of Manufactured Solutions
    # general solution of the spherical Laplacian equation
    # satisfies Dr(u(t,0)) = 0
    u_exact = (r,t) -> exp.(-t) * sin.(r) ./ r

    # Parameters, variables, and derivatives
    @parameters t r
    @variables u(..)
    Dt = Differential(t)
    Dr = Differential(r)

    # 1D PDE and boundary conditions

    eq  = Dt(u(t,r)) ~ 1/r^2 * Dr(r^2 * Dr(u(t,r)))
    bcs = [u(0,r) ~ sin(r)/r,
           Dr(u(t,0)) ~ 0,
           u(t,1) ~ exp(-t) * sin(1)]
        #    Dr(u(t,1)) ~ -exp(-t) * sin(1)]

    # Space and time domains
    domains = [t ∈ Interval(0.0,1.0),
               r ∈ Interval(0.0,1.0)]

    # PDE system
    pdesys = PDESystem(eq,bcs,domains,[t,r],[u(t,r)])

    # Method of lines discretization
    dr = 0.1
    order = 2
    discretization = MOLFiniteDifference([r=>dr],t)
    prob = discretize(pdesys,discretization)

    # Solve ODE problem
    sol = solve(prob,Tsit5(),saveat=0.1)

    r = (0:dr:1)[2:end-1]
    t = sol.t

    # Test against exact solution
    for i in 1:length(sol)
        exact = u_exact(r, t[i])
        u_approx = sol.u[i]
        @test all(isapprox.(u_approx, exact, atol=0.01))
    end
end

@testset "Test 08: Dt(u(t,r)) ~ 4/r^2 * Dr(r^2 * Dr(u(t,r))) (Spherical Laplacian)" begin
    # Method of Manufactured Solutions
    # general solution of the spherical Laplacian equation
    # satisfies Dr(u(t,0)) = 0
    u_exact = (r,t) -> exp.(-4t) * sin.(r) ./ r

    # Parameters, variables, and derivatives
    @parameters t r
    @variables u(..)
    Dt = Differential(t)
    Dr = Differential(r)

    # 1D PDE and boundary conditions

    eq  = Dt(u(t,r)) ~ 4/r^2 * Dr(r^2 * Dr(u(t,r)))
    bcs = [u(0,r) ~ sin(r)/r,
           Dr(u(t,0)) ~ 0,
           u(t,1) ~ exp(-4t) * sin(1)]

    # Space and time domains
    domains = [t ∈ Interval(0.0,1.0),
               r ∈ Interval(0.0,1.0)]

    # PDE system
    pdesys = PDESystem(eq,bcs,domains,[t,r],[u(t,r)])

    # Method of lines discretization
    dr = 0.1
    order = 2
    discretization = MOLFiniteDifference([r=>dr],t)
    prob = discretize(pdesys,discretization)

    # Solve ODE problem
    sol = solve(prob,Tsit5(),saveat=0.1)

    r = (0:dr:1)[2:end-1]
    t = sol.t

    # Test against exact solution
    for i in 1:length(sol)
        exact = u_exact(r, t[i])
        u_approx = sol.u[i]
        @test all(isapprox.(u_approx, exact, atol=0.01))
    end
end

@testset "Test 10: linear diffusion, two variables, mixed BCs" begin
    # Method of Manufactured Solutions
    u_exact = (x,t) -> exp.(-t) * cos.(x)
    v_exact = (x,t) -> exp.(-t) * sin.(x)

    # Parameters, variables, and derivatives
    @parameters t x
    @variables u(..) v(..)
    Dt = Differential(t)
    Dx = Differential(x)
    Dxx = Dx^2

    # 1D PDE and boundary conditions
    eqs = [Dt(u(t,x)) ~ Dxx(u(t,x)),
           Dt(v(t,x)) ~ Dxx(v(t,x))]
    bcs = [u(0,x) ~ cos(x),
           v(0,x) ~ sin(x),
           u(t,0) ~ exp(-t),
           Dx(u(t,1)) ~ -exp(-t) * sin(1),
           Dx(v(t,0)) ~ exp(-t),
           v(t,1) ~ exp(-t) * sin(1)]

    # Space and time domains
    domains = [t ∈ Interval(0.0,1.0),
               x ∈ Interval(0.0,1.0)]

    # PDE system
    pdesys = PDESystem(eqs,bcs,domains,[t,x],[u(t,x),v(t,x)])

    # Method of lines discretization
    l = 100
    dx = range(0.0,1.0,length=l)
    order = 2
    discretization = MOLFiniteDifference([x=>dx],t)

    # Convert the PDE problem into an ODE problem
    prob = discretize(pdesys,discretization)

    # Solve ODE problem
    sol = solve(prob,Tsit5(),saveat=0.1)

    x_sol = dx[2:end-1]
    t_sol = sol.t

    # Test against exact solution
    for i in 1:length(sol)
        @test all(isapprox.(u_exact(x_sol, t_sol[i]), sol.u[i][1:l-2], atol=0.01))
        @test all(isapprox.(v_exact(x_sol, t_sol[i]), sol.u[i][l-1:end], atol=0.01))
    end
end

@testset "Test 11: linear diffusion, two variables, mixed BCs, with parameters" begin
    @parameters t x
    @parameters Dn, Dp
    @variables u(..) v(..)
    Dt = Differential(t)
    Dx = Differential(x)
    Dxx = Differential(x)^2

    eqs  = [Dt(u(t,x)) ~ Dn * Dxx(u(t,x)) + u(t,x)*v(t,x),
            Dt(v(t,x)) ~ Dp * Dxx(v(t,x)) - u(t,x)*v(t,x)]
    bcs = [u(0,x) ~ sin(pi*x/2),
        v(0,x) ~ sin(pi*x/2),
        u(t,0) ~ 0.0, Dx(u(t,1)) ~ 0.0,
        v(t,0) ~ 0.0, Dx(v(t,1)) ~ 0.0]

    domains = [t ∈ Interval(0.0,1.0),
            x ∈ Interval(0.0,1.0)]

    pdesys = PDESystem(eqs,bcs,domains,[t,x],[u(t,x),v(t,x)],[Dn=>0.5, Dp=>2])
    discretization = MOLFiniteDifference([x=>0.1],t)
    prob = discretize(pdesys,discretization)
    @test prob.p == [0.5,2]
    # Make sure it can be solved
    sol = solve(prob,Tsit5())
end

@testset "Test 12: linear diffusion, two variables, mixed BCs, different independent variables" begin
    # Method of Manufactured Solutions
    u_exact = (x,t) -> exp.(-t) * cos.(x)
    v_exact = (y,t) -> exp.(-t) * sin.(y)

    # Parameters, variables, and derivatives
    @parameters t x y
    @variables u(..) v(..)
    Dt = Differential(t)
    Dx = Differential(x)
    Dxx = Dx^2
    Dy = Differential(y)
    Dyy = Dy^2

    # 1D PDE and boundary conditions
    eqs = [Dt(u(t,x)) ~ Dxx(u(t,x)),
           Dt(v(t,y)) ~ Dyy(v(t,y))]
    bcs = [u(0,x) ~ cos(x),
           v(0,y) ~ sin(y),
           u(t,0) ~ exp(-t),
           Dx(u(t,1)) ~ -exp(-t) * sin(1),
           Dy(v(t,0)) ~ exp(-t),
           v(t,2) ~ exp(-t) * sin(2)]

    # Space and time domains
    domains = [t ∈ Interval(0.0,1.0),
               x ∈ Interval(0.0,1.0),
               y ∈ Interval(0.0,2.0)]

    # PDE system
    pdesys = PDESystem(eqs,bcs,domains,[t,x,y],[u(t,x),v(t,y)])

    # Method of lines discretization
    l = 100
    dx = range(0.0,1.0,length=l)
    dy = range(0.0,2.0,length=l)
    order = 2
    discretization = MOLFiniteDifference([x=>dx,y=>dy],t)

    # Convert the PDE problem into an ODE problem
    prob = discretize(pdesys,discretization)

    # Solve ODE problem
    sol = solve(prob,Tsit5(),saveat=0.1)

    x_sol = dx[2:end-1]
    y_sol = dy[2:end-1]
    t_sol = sol.t

    # Test against exact solution
    for i in 1:length(sol)
        @test all(isapprox.(u_exact(x_sol, t_sol[i]), sol.u[i][1:l-2], atol=0.01))
        @test all(isapprox.(v_exact(y_sol, t_sol[i]), sol.u[i][l-1:end], atol=0.01))
    end
end

@testset "Test 13: one linear diffusion with mixed BCs, one ODE" begin
    # Method of Manufactured Solutions
    u_exact = (x,t) -> exp.(-t) * sin.(x)
    v_exact = (t) -> exp.(-t)

    @parameters t x
    @variables u(..) v(..)
    Dt = Differential(t)
    Dx = Differential(x)
    Dxx = Dx^2

    # 1D PDE and boundary conditions
    eqs = [Dt(u(t,x)) ~ Dxx(u(t,x)),
           Dt(v(t)) ~ -v(t)]
    bcs = [u(0,x) ~ sin(x),
           v(0) ~ 1,
           u(t,0) ~ 0,
           Dx(u(t,1)) ~ exp(-t) * cos(1)]

    # Space and time domains
    domains = [t ∈ Interval(0.0,1.0),
               x ∈ Interval(0.0,1.0)]

    # PDE system
    pdesys = PDESystem(eqs,bcs,domains,[t,x],[u(t,x),v(t)])

    # Method of lines discretization
    l = 100
    dx = range(0.0,1.0,length=l)
    order = 2
    discretization = MOLFiniteDifference([x=>dx],t)

    # Convert the PDE problem into an ODE problem
    prob = discretize(pdesys,discretization)

    # Solve ODE problem
    sol = solve(prob,Tsit5(),saveat=0.1)

    x_sol = dx[2:end-1]
    t_sol = sol.t

    # Test against exact solution
    for i in 1:length(sol)
        @test all(isapprox.(u_exact(x_sol, t_sol[i]), sol.u[i][1:length(x_sol)], atol=0.01))
        @test v_exact(t_sol[i]) ≈ sol.u[i][end] atol=0.01
    end
end

@testset "Test error 01: Test Invalid Centered Order" begin
    # Method of Manufactured Solutions
    u_exact = (x,t) -> exp.(-t) * cos.(x)

    # Parameters, variables, and derivatives
    @parameters t x
    @variables u(..)
    Dt = Differential(t)
    Dxx = Differential(x)^2

    # 1D PDE and boundary conditions
    eq  = Dt(u(t,x)) ~ Dxx(u(t,x))
    bcs = [u(0,x) ~ cos(x),
           u(t,0) ~ exp(-t),
           u(t,Float64(π)) ~ -exp(-t)]

    # Space and time domains
    domains = [t ∈ Interval(0.0,1.0),
               x ∈ Interval(0.0,Float64(π))]

    # PDE system
    pdesys = PDESystem(eq,bcs,domains,[t,x],[u(t,x)])

    # Method of lines discretization
    dx = range(0.0,Float64(π),length=30)

    # Explicitly specify and invalid order of centered difference
    for order in 1:6
        discretization = MOLFiniteDifference([x=>dx],t;centered_order=order)
        if order % 2 != 0
            @test_throws ArgumentError discretize(pdesys,discretization)
        else
            discretize(pdesys,discretization)
        end
    end
end


function BondNSxPs(m,b0,b1,b2,tau,b3=0.0,tau2=1.0)
#BondNSxPs    Extended Nelson and Siegel (1987) spot rate, forward rate, and discount function
#
#
#
#
#  Usage:    (s,f,d) = BondNSxPs(m,b0,b1,b2,tau[,b3[,tau2]])
#
#  Input:    m         NxK matrix  times to maturity in years (don't use less than 1e-8 or so)
#            b0        scalar, beta0 parameter in NS
#            b1        scalar, beta1 parameter in NS
#            b2        scalar, beta2 parameter in NS
#            tau       scalar, tau   parameter in NS
#            b3        (optional) scalar, beta2 parameter in extended NS [0]
#            tau2      (optional) scalar, tau2  parameter in extended NS [1]
#
#  Output:   s         NxK matrix,  spot rate (continously compounded)
#            f         NxK matrix,  forward rate (continously compounded)
#            d         NxK matrix,  discount function
#
#
#
#  Note:     For more details, see
#            (a) Svensson (1995), "Estimating Forward Interest Rates with
#                the Extended Nelson & Siegel Method," Quarterly Review,
#                Sveriges Riksbank, 1995:3, 13-26.
#            (b) Soderlind and Svensson (1997), "New Techniques to Extract
#                Market Expectations from Financial Instruments,"
#                Journal of Monetary Economics 40, 383-429.
#
#
#
#  Paul.Soderlind@unisg.ch, April 2002, to Julia 2015
#----------------------------------------------------------------------------*/

                                                 #forward rate
  f =  b0 .+ b1*exp.(-m/tau) + b2*(m/tau).*exp.(-m/tau) + b3*(m/tau2).*exp.(-m/tau2)

  s =  b0 .+ b1* (1 .- exp.(-m/tau)) ./(m/tau) +
             b2*((1 .- exp.(-m/tau)) ./(m/tau)  - exp.(-m/tau)) +
             b3*((1 .- exp.(-m/tau2))./(m/tau2) - exp.(-m/tau2))        #spot rate

  d = exp.(-s.*m)             #discount function

  return s,f,d

end
#------------------------------------------------------------------------------

#------------------------------------------------------------------------------
function BondNSxEstPs(par0,Q,tm,c,s0,ytmLoss=0,weight=1.0)
#BondNSxEstPs  Estimates parameters in (extended) Nelson-Siegel yield curve model by
#              non-linear least squares (minimizing squared differences between
#              actual and fitted bond prices or yield to maturities).
#
#
#
#  Usage:    NSb = BondNSxEstPs(par0,Q,tm,c,s0[,ytmLoss[,weight]])
#
#
#  Input:    par0       4x1 or 6x1 vector, initial guess of paremeters
#                         if 4x1: standard NS with par0 = [b0,b1,b2,tau]
#                         if 6x1: extended NS with par0 = [b0,b1,b2,tau,b3,tau2]
#            Q          nx1 vector, data on bond prices, eg. 1.01
#            tm         nx1 vector, data on time to maturity (in years), eg. 2.54
#            c          nx1 vector, data on bond coupons, eg. 0.06
#            s0         scalar, restricted value of short rate:
#                       if [] no restriction, else b0 + b1 = s0 is imposed
#            ytmLoss    scalar, scalar, 0: mimimize squared price errors;
#                                       1: minimize squared ytm errors (default 0)
#            weight     scalar or nx1 vector, weights in the loss function
#
#  Output:   NSb        4x1 or 6x1 vector, estimated parameters
#
#
#
#  Note:     For more details, see
#            (a) Svensson (1995), "Estimating Forward Interest Rates with
#                the Extended Neolson & Siegel Method," Quarterly Review,
#                Sveriges Riksbank, 1995:3, 13-26.
#            (b) Soderlind and Svensson (1997), "New Techniques to Extract
#                Market Expectations from Financial Instruments,"
#                Journal of Monetary Economics 40, 383-429.
#
#
#  Uses:  Optim
#
#
#
#  Paul.Soderlind@unisg.ch, April 2002, to Julia Nov 2015
#------------------------------------------------------------------------------

  Qtc = [Q tm c]

  if in(length(par0),[4,6]) && !isempty(s0)
      par0 = deleteat!(par0,2)   #b1 = s0 - b0 is imposed by BondNSxLossPs
  end

  Sol = optimize(b->BondNSxLossPs(b,Qtc,s0,ytmLoss,weight),par0)
  NSb = Optim.minimizer(Sol)
  if !Optim.converged(Sol)
    println("no convergence")
    return
  end

  if length(NSb) == 3
    NSb[[1,3]] = abs.(NSb[[1,3]])
    NSb        = [NSb[1]; s0-NSb[1]; NSb[2:3]]
  elseif length(NSb) == 4
    NSb[[1,4]] = abs.(NSb[[1,4]])
  elseif length(NSb) == 5
    NSb[[1,3,5]] = abs.(NSb[[1,3,5]])
    NSb = [NSb[1]; s0-NSb[1]; NSb[2:5]]
  elseif length(NSb) == 6
    NSb[[1,4,6]] = abs.(NSb[[1,4,6]])
  end

  return NSb

end
#----------------------------------------------------------------------------

#------------------------------------------------------------------------------
function BondNSxLossPs(b,Qtc,s0,ytmLoss=0,weight=1.0)
#BondNSxLossPs    Defines loss function for bond prices in extended Nelson-Siegel model.
#                 Used for estimation of the parameters in the model by minimizing
#                 Loss (squared price deviations, possibly with weights).
#
#
#  Usage:    Loss = BondNSxLossPs(b,Qtc,s0,ytmLoss,weight)    or
#                 = BondNSxLossPs(b,Qtc,s0,ytmLoss)           or
#                 = BondNSxLossPs(b,Qtc,s0)
#
#  Input:    b         3x1, 4x1, 5x1, or 6x1 vector with parameters in Nelson-Siegel model
#                        if 3x1: [b0,b2,tau], NS with restriction that b1 = g_s0 - b0
#                        if 4x1: [b0,b1,b2,tau], NS without restrictions
#                        if 5x1: [b0,b2,tau,b3,tau2], extended NS with restriction that b1 = g_s0 - b0
#                        if 6x1: [b0,b1,b2,tau,b3,tau2], extended NS without restrictions
#            Qtc       nx3 matrix, data on [bond prices,time to maturity,coupons]. Use ytm instead
#                      of bond prices if ytmLoss==1
#            s0        scalar, restricted value of short rate
#            ytmLoss   scalar, 0: mimimize squared price errors;
#                      1: minimize squared ytm errors (default 0)
#            weight    scalar or nx1 vector, weights in the loss function
#
#  Output:   Loss    scalar, sum of squared differences between implied and actual bond
#                    prices (or yields)
#
#  Calls on: BondNSxPs, BondYieldToMatPs (if ytmLoss==1)
#
#------------------------------------------------------------------------------

  (Q,tm,c) = [Qtc[:,i] for i=1:3]       #data on bond prices, time to maturity, coupons
  n        = length(c)           #number of bonds

  if length(b) == 3             #standard Nelson-Siegel with restriction b1 = s0-b0
    (b0,b1,b2,tau,b3,tau2) = (abs(b[1]),s0-abs(b[1]),b[2],abs(b[3]),0.0,1.0)
  elseif length(b) == 4         #standard Nelson-Siegel
   (b0,b1,b2,tau,b3,tau2)  = (abs(b[1]),b[2],b[3],abs(b[4]),0.0,1.0)
  elseif length(b) == 5         #extended Nelson-Siegel with restriction b1 = s0 - b0
    (b0,b1,b2,tau,b3,tau2) = (abs(b[1]),s0-abs(b[1]),b[2],abs(b[3]),b[4],abs(b[5]))
  elseif length(b) == 6         #extended Nelson-Siegel
   (b0,b1,b2,tau,b3,tau2)  = (abs(b[1]),b[2],b[3],abs(b[4]),b[5],abs(b[6]))
  end

  QNS = fill(NaN,n)
  for i = 1:n                            #loop over bonds
    ti     = filter(z->z>0,mod(tm[i],1):tm[i])      #time to coupon payments, >0 only
    d      = BondNSxPs(ti,b0,b1,b2,tau,b3,tau2)[3]  #NSx: spot, forward, discount fn
    QNS[i] = sum(d.*c[i]) + d[end]                  #fitted bond price
    if ytmLoss == 1
      QNS[i] = BondYieldToMatPs(QNS[i],c[i],ti,1,1,0.05,1e-7)  #fitted ytm
    end
  end
  Loss = 1.0 + 100*sum( weight.*(QNS - Q).^2 )  #weighted sum of squared deviations of fitted from actual

  return Loss

end
#------------------------------------------------------------------------------


using Random, Statistics, LinearAlgebra, Test, FillArrays, Printf, TotalLeastSquares, StatsBase
using RandomizedLinAlg
Random.seed!(0)

@testset "TotalLeastSquares" begin

    @testset "TLS" begin
        @info "Testing TLS"
        x   = randn(3)
        A   = randn(50,3)
        σa  = 1
        σy  = 0.01
        An  = A + σa*randn(size(A))
        y   = A*x
        yn  = y + σy*randn(size(y))
        Qaa = σa^2*Eye(prod(size(A)))
        Qay = 0Eye(prod(size(A)),length(y))
        Qyy = σy^2*Eye(prod(size(y)))

        x̂ = An\yn
        @printf "Least squares error: %25.3e %10.3e %10.3e, Norm: %10.3e\n" (x-x̂)... norm(x-x̂)
        @test norm(x-x̂) < 1

        x̂ = wls(An,yn,Qyy)
        @printf "Weigthed Least squares error: %16.3e %10.3e %10.3e, Norm: %10.3e\n" (x-x̂)... norm(x-x̂)
        @test norm(x-x̂) < 1

        x̂ = tls(An,yn)
        @printf "Total Least squares error: %19.3e %10.3e %10.3e, Norm: %10.3e\n" (x-x̂)... norm(x-x̂)
        @test norm(x-x̂) < 1

        x̂ = wtls(An,yn,Qaa,Qay,Qyy,iters=10)
        @printf "Weighted Total Least squares error: %10.3e %10.3e %10.3e, Norm: %10.3e\n" (x-x̂)... norm(x-x̂)
        println("----------------------------")
        @test norm(x-x̂) < 1

        @test tls(An,yn) ≈ tls!([An yn], size(An,2))

        rowC = rowcovariance([[σa^2*Eye(3) zeros(3); zeros(1,3) σy^2] for _ in 1:50])
        @test rowC[1] ≈ Qaa
        @test rowC[2] ≈ Qay
        @test rowC[3] ≈ Qyy

    end

    @testset "Robust PCA" begin
        @info "Testing Robust PCA"
        D = [0.462911    0.365901  0.00204357    0.692873    0.935861;
        0.0446199    0.108606   0.0664309   0.0736707    0.264429;
        0.320581    0.287788    0.073133    0.188872    0.526404;
        0.356266    0.197536 0.000718338    0.513795    0.370094;
        0.677814    0.011651    0.818047   0.0457694    0.471477]

        A = [0.462911   0.365901  0.00204356   0.345428   0.623104;
        0.0446199  0.108606  0.0429271    0.0736707  0.183814;
        0.320581   0.203777  0.073133     0.188872   0.472217;
        0.30725    0.197536  0.000717701  0.201626   0.370094;
        0.234245   0.011651  0.103622     0.0457694  0.279032]

        E = [0.0        0.0        0.0        0.347445  0.312757 ;
        0.0        0.0        0.0235038  0.0       0.0806151;
        0.0        0.0840109  0.0        0.0       0.0541868;
        0.0490157  0.0        6.5061e-7  0.312169  0.0      ;
        0.443569   0.0        0.714425   0.0       0.192445]

        Â, Ê,_,_ = rpca(D, nonnegE=true, nonnegA=true, verbose=false)

        @test Â ≈ A atol=1.0e-6
        @test Ê ≈ E atol=1.0e-6
        @test norm(D - (Â + Ê))/norm(D) < sqrt(eps())

        Â, Ê,_,_ = rpca(D, nonnegE=false, nonnegA=false, verbose=false)
        @test norm(D - (Â + Ê))/norm(D) < sqrt(eps())

        @testset "Missing values" begin
            @info "Testing Missing values"

            res = map(1:20) do _
                N = 500
                y = sin.(0.1 .* (1:N)) .+ 0.1*randn(N)
                miss = rand(N) .< 0.1
                yn = y .+ miss .* 1e2
                yf = lowrankfilter(yn,40)
                mean(abs2,y-yf)/mean(abs2,y)
            end
            @test mean(res) < 0.025

        end

        @testset "Complex" begin
            @info "Testing Complex"
            u = randn(ComplexF64, 100)
            v = randn(ComplexF64, 20)
            E = randn(ComplexF64,100,20) .* 10 .* (rand.() .< 0.01)
            A = u*v'
            D = A .+ E
            Ah, Eh, s, sv = rpca(D)
            @test sum(abs2, Eh-E)/sum(abs2, E) < 1e-5
            @test sum(abs2, Ah-A)/sum(abs2, A) < 1e-5

        end

    end

@testset "rtls" begin
    @info "Testing rtls"
    passes = map(1:1000) do _
        x   = randn(3)
        A   = randn(50,3)
        σ   = 5
        An  = A + σ*randn(size(A)) .* (rand(size(A)...) .< 0.1)
        y   = A*x
        yn  = y + σ*randn(size(y)) .* (rand(size(y)...) .< 0.1)

        AA  = [An yn]
        Ah,Eh,_,_ = rpca(AA, verbose=false)

        sum(abs2, Ah - [A y])/sum(abs2,[A y]) < sum(abs2, AA - [A y])/sum(abs2,[A y])
    end
    @show mean(passes)
    @test mean(passes) > 0.9

    passes = map(1:1000) do _
        x   = randn(3)
        A   = randn(50,3)
        σ   = 50
        An  = A + σ*randn(size(A)) .* (rand(size(A)...) .< 0.1)
        y   = A*x
        yn  = y + σ*randn(size(y)) .* (rand(size(y)...) .< 0.1)

        x̂t = tls(An,yn)
        x̂r = rtls(An,yn)

        norm(x-x̂r) < norm(x-x̂t)
    end
    @show mean(passes)
    @test mean(passes) > 0.9

    # @testset "Elastic Net rtls" begin
    #     @info "Testing Elastic Net rtls"
    #     passes = map(1:500) do _
    #         x     = randn(10)
    #         A     = randn(100,3) * randn(3,10)
    #         σs    = 50
    #         σd    = 0.5
    #         S     = σs*randn(size(A)) .* (rand(size(A)...) .< 0.1)
    #         D     = σd*randn(size(A))
    #         N     = S + D
    #         An    = A + N
    #
    #         Ah,Eh,_,_ = rpca(An, verbose=false)
    #
    #         λ = 1/sqrt(100)
    #         μ = 0.001λ
    #         Ah2,Eh2,_,_ = rpca(An, verbose=false, proxE = ElasticNet(λ, μ))
    #
    #         sum(abs2, Ah - A)/sum(abs2,A) < sum(abs2, An - A)/sum(abs2,A),
    #         sum(abs2, Eh2 - N)/sum(abs2,N) < sum(abs2, Eh - N)/sum(abs2,N)
    #     end
    #     @show mean(getindex.(passes, 1))
    #     @test mean(getindex.(passes, 1)) > 0.9
    #     @show mean(getindex.(passes, 2))
    #     @test mean(getindex.(passes, 2)) > 0.9
    #
    #
    # end

end

# @testset "rtls ga" begin
#     @info "Testing rtls ga"
#
#     passes = map(1:1000) do _
#         x   = randn(5)
#         A   = randn(500,5)
#         σ   = 50
#         An  = A + σ*randn(size(A)) .* (rand(size(A)...) .< 0.01)
#         y   = A*x
#         yn  = y + σ*randn(size(y)) .* (rand(size(y)...) .< 0.01)
#
#         x̂t = tls(An,yn)
#         x̂r = rtls_ga(An,yn, μ=entrywise_median)
#
#         norm(x-x̂r) < norm(x-x̂t)
#     end
#     @show mean(passes)
#     @test mean(passes) > 0.7
#
# end

@testset "soft hankel" begin
    @info "Testing soft hankel"

    @test hankel(1:20,2) == [1:19 2:20]
    @test hankel(1:20,3,2) == [1:2:17 2:2:18 3:2:19]

    A = hankel(1:8,4)
    @test ishankel(A)
    An = A + 0.1randn(size(A))
    @test !ishankel(An)
    Anc = copy(An)
    TotalLeastSquares.soft_hankel!(An, 0.1)
    @test sum(abs2,An-A) < sum(abs2,Anc-A)

    An = -A + 0.1randn(size(A))
    Anc = copy(An)
    TotalLeastSquares.soft_hankel!(An, 0.1)
    @test sum(abs2,An+A) < sum(abs2,Anc+A)

    passes = map(1:100) do _
        y = randn(20)
        yn = y .+ 0.1 .* randn.()
        A = hankel(y,3)
        An = hankel(yn,3)
        Anc = copy(An)
        A1,E1,_,_ = rpca(An, verbose=false, nukeA=false)
        A2,E2,_,_ = rpca(An, verbose=false, nukeA=false, hankel=true)
        sum(abs2,A2-A) < sum(abs2,A1-A)
    end
    @show mean(passes)
    @test mean(passes) > 0.7

    @info "Small random Hankel"
    passes = map(1:500) do _
        y = randn(100)
        A = hankel(y,5)
        @test ishankel(A)
        A1,E1,_,_ = rpca(A, verbose=false, nukeA=false)
        A2,E2,_,_ = rpca(A, verbose=false, nukeA=false, hankel=true)
        @test ishankel(A2)
        @test ishankel(E2)
        mean(abs2,A2-A) < mean(abs2,A1-A)
    end
    @show mean(passes)
    @test mean(passes) > 0.8

    @info "Big random Hankel"
    passes = map(1:10) do _
        y = randn(1000)
        A = hankel(y,50)
        @test ishankel(A)
        A1,E1,_,_ = rpca(A, verbose=false, nukeA=false)
        A2,E2,_,_ = rpca(A, verbose=false, nukeA=false, hankel=true)
        @test ishankel(A2)
        @test ishankel(E2)
        mean(abs2,A2-A) < mean(abs2,A1-A)
    end
    @show mean(passes)
    @test mean(passes) >= 0.8

end

@testset "lowrankfilter" begin
    @info "Testing lowrankfilter"
    qn(x) = x ./ quantile(abs.(x), 0.9)
    T = 1000
    t = 1:T
    y = sin.(0.1 .* t) |> qn

    H = hankel(y, 2)
    @test ishankel(H)
    @test unhankel(H) == y

    H = hankel(y,2,2)
    yh = unhankel(H,2,T)
    @test yh == y

    H = hankel(y,5,2)
    yh = unhankel(H,2,T)
    @test yh[1:end-1] ≈ y[1:end-1]

    y2 = randn(T)
    H = hankel([y y2],5,2)
    yh = unhankel(H,2,T,2)
    @test yh[1:end-1,:] ≈ [y y2][1:end-1,:]

    n = 20randn(T) .* (rand(T) .< 0.01) + 0.1randn(T)
    yf = lowrankfilter(y+n) |> qn
    @test mean(abs2, y-yf)/mean(abs2, n) < 0.001
    @show mean(abs2, y-yf)/mean(abs2, n)

    # Randomized lin alg
    yf = lowrankfilter(y+n, opnorm=x->rnorm(x,10)) |> qn
    @test mean(abs2, y-yf)/mean(abs2, n) < 0.001
    @show mean(abs2, y-yf)/mean(abs2, n)

    yf = lowrankfilter(y+n, svd=rsvd_fnkz) |> qn
    @test mean(abs2, y-yf)/mean(abs2, n) < 0.05
    @show mean(abs2, y-yf)/mean(abs2, n)

    yf = lowrankfilter(y+n, opnorm=x->rnorm(x,10), svd=rsvd_fnkz) |> qn
    @test mean(abs2, y-yf)/mean(abs2, n) < 0.05
    @show mean(abs2, y-yf)/mean(abs2, n)

    yf = lowrankfilter(y+n, opnorm=x->rnorm(x,10), svd=rsvd_fnkz, maxrank=5) |> qn
    @test mean(abs2, y-yf)/mean(abs2, n) < 0.05
    @show mean(abs2, y-yf)/mean(abs2, n)

    # Only SSA
    n = randn(T)
    yf = lowrankfilter(y+n, sv=2) |> qn
    @test mean(abs2, y-yf)/mean(abs2, n) < 0.05
    @show mean(abs2, y-yf)/mean(abs2, n)

    @testset "Multi channel" begin
        @info "Testing Multi channel"

        T = 1000
        t = 1:T
        y1 = sin.(0.1 .* t)
        y2 = sin.(0.3 .* t)
        y = [y1 y2+0.5y1]
        yn = y .+ 0.1 .* randn.() .+ (rand.() .< 0.01) .* 5 .* randn.()
        yf = lowrankfilter(yn, 20)
        @show mean(abs2,y-yf)/mean(abs2,y-yn)
        @test mean(abs2,y-yf)/mean(abs2,y-yn) < 0.05

        # There shoud be benefit to joint filtering since the two signals are correlated
        yf1 = lowrankfilter(yn[:,1], 20)
        @test mean(abs2,y1-yf1)/mean(abs2,y1-yf[:,1]) > 1.02

        # Same freq, different phase
        y1 = sin.(0.1 .* t)
        y2 = sin.(0.1 .* t .+ 0.01)
        y = [y1 y2]
        yn = y .+ 0.1 .* randn.() .+ (rand.() .< 0.01) .* 5 .* randn.()
        yf = lowrankfilter(yn, 20)
        @show mean(abs2,y-yf)/mean(abs2,y-yn)
        @test mean(abs2,y-yf)/mean(abs2,y-yn) < 0.02

        yf1 = lowrankfilter(yn[:,1], 20)
        @show mean(abs2,y1-yf1)/mean(abs2,y1-yf[:,1])
        @test mean(abs2,y1-yf1)/mean(abs2,y1-yf[:,1]) > 1.2

    end

end

@testset "rpca_ga" begin
    @info "Testing rpca_ga"

    Random.seed!(1)
    for r = 1:10, ϵ = exp10.(LinRange(-8, 0, 20))
        s = svd(randn(10,40))
        u,v = s.U[:,1:r], s.V[:,1:r]
        A = u*Diagonal(10*(1:r))*v' .+ ϵ .* randn.()
        Q = rpca_ga(A, r, verbose=false)
        @test rank([Q u], atol=ϵ) == r # Q and u should span the same space, but they may not share any other properties
        @test norm(Q'Q - I) < sqrt(eps())
    end

    # Below is the reverse dimensions of above
    for r = 1:10, ϵ = exp10.(LinRange(-8, 0, 20))
        s = svd(randn(40,10))
        u,v = s.U[:,1:r], s.V[:,1:r]
        A = u*Diagonal(10*(1:r))*v' .+ ϵ .* randn.()
        Q = rpca_ga(A, r, verbose=false)
        @test rank([Q u], atol=ϵ) == r # Q and u should span the same space, but they may not share any other properties
        @test norm(Q'Q - I) < sqrt(eps())
    end

    @testset "entrywise trimmed mean" begin
        @info "Testing entrywise trimmed mean"

        U = randn(10,10)
        s = zeros(10)
        w = ones(10)
        m1 = TotalLeastSquares.μ!(s,w,U)
        @test m1 ≈ mean(U, dims=2)

        m2 = TotalLeastSquares.entrywise_trimmed_mean(s,w,U,0)
        @test m2 ≈ mean(U, dims=2)

        w = randn(10)
        m1 = TotalLeastSquares.μ!(s,w,U)
        @test m1 ≈ sum(U .* w', dims=2)./ sum(w)

        m2 = TotalLeastSquares.entrywise_trimmed_mean(s,w,U,0)
        @test m2 ≈ sum(U .* w', dims=2)./ sum(w)

        w = ones(10)
        m2 = TotalLeastSquares.entrywise_trimmed_mean(s,w,U,0.1)
        for i in eachindex(m2)
            @test m2[i] ≈ mean(StatsBase.trim(U[i,:], prop=0.1))
        end

        r = 3; ϵ = 1e-8
        Random.seed!(1)
        passes = map(Iterators.product(1:4, exp10.(LinRange(-8, -1, 3)))) do (r,ϵ)
            # @show r, ϵ
            s = svd(randn(10,1000))
            u,v = s.U[:,1:r], s.V[:,1:r]
            # A = u*Diagonal(10*(1:r))*v' .+ ϵ .* randn.()
            A = u*Diagonal(s.S[1:r])*v' .+ ϵ .* randn.()
            A .+= 1000*randn.() .* (rand.() .< 0.01)

            Qtm = rpca_ga(A, r, verbose=false, μ = TotalLeastSquares.entrywise_trimmed_mean, iters=120)
            Qm = rpca_ga(A, r, verbose=false, iters=120)
            sum(svdvals([Qtm u])[r+1:2r]) < sum(svdvals([Qm u])[r+1:2r])
            # @test norm(Q'Q - I) < r*sqrt(eps()) # This test should pass but doesn't
        end
        @show mean(passes)
        @test mean(passes) > 0.8

        passes = map(Iterators.product(1:4, exp10.(LinRange(-8, -1, 3)))) do (r,ϵ)
            # @show r, ϵ
            s = svd(randn(10,1000))
            u,v = s.U[:,1:r], s.V[:,1:r]
            # A = u*Diagonal(10*(1:r))*v' .+ ϵ .* randn.()
            A = u*Diagonal(s.S[1:r])*v' .+ ϵ .* randn.()
            A .+= 1000*randn.() .* (rand.() .< 0.01)

            Qmed = rpca_ga(A, r, verbose=false, μ = TotalLeastSquares.entrywise_median, iters=120)
            Qm = rpca_ga(A, r, verbose=false, iters=120)
            sum(svdvals([Qmed u])[r+1:2r]) < sum(svdvals([Qm u])[r+1:2r])
            # @test norm(Q'Q - I) < r*sqrt(eps())  # This test should pass but doesn't
        end
        @show mean(passes)
        @test mean(passes) > 0.9
    end
end

@testset "irls" begin
    @info "Testing irls"

    N = 10000
    n = 20
    A = randn(N, n)
    x = randn(n)
    y = A*x
    e = 100*(rand(N) .< 0.01) + 0.1randn(N)
    @time xh = irls(A,y+e, verbose=true,iters=100)
    @test norm(x-xh)/norm(x) < 0.01
    xhls = A\(y+e)
    @test norm(x-xh) < norm(x-xhls)
end

@testset "sls" begin
    @info "Testing sls"

    for i = 1:100
        x = TotalLeastSquares.proj_simplex!(randn(i))
        @test sum(x) ≈ 1
        @test all(x .>= 0)
    end

    s1(x) = x ./= sum(x)
    ##
    n = 10
    N = 2
    results = map(1:100) do _
        a = [(rand(n)) for _ in 1:N]
        A = reduce(hcat,a)
        x = s1(rand(N))
        y = A*x
        yn = y .+ 0.01.*randn.()
        xh1 = TotalLeastSquares.proj_simplex!(A\yn)

        xh2 = sls(A,yn,verbose=false)
        @test sum(xh2) ≈ 1

        norm(x-xh2) < norm(x-xh1) || norm(yn-A*xh2) < norm(yn-A*xh1)
    end
    @test mean(results) >= 0.98
end

@testset "flts" begin
    @info "Testing flts"
    # Example from paper
    x = 10randn(1000)
    a, b = 1, 2
    y = a*x .+ b
    y[801:end] = 5randn(200)
    x[801:end] = 5randn(200) .+ 50
    xb = [x ones(1000)]
    # default
    res = flts(xb,y, verbose = true)
    @test a ≈ res[1]
    @test b ≈ res[2]
    # define % of outliers
    res1 = flts(xb,y, N = 10, outliers = 0.20, verbose = true)
    @test a ≈ res1[1]
    @test b ≈ res1[2]
    # define length of set H and check the subset
    (H, res2, Q) = flts(xb,y, N = 10, h = 800, return_set = true, verbose = true)
    @test a ≈ res2[1]
    @test b ≈ res2[2]
    @test all(H .< 801)
    @test Q + 1 ≈ 1
end
end


"""
    natural_eigdist(𝚽, 𝛌, Q; α = 1.0, T = :Inf,
                    input_format = :zero_measures, distance = :DAG,
                    edge_weight = 1, edge_length = 1)

compute natural distances between graph Laplacian eigenvectors.

# Input Arguments
- `𝚽::Matrix{Float64}`: matrix of (weighted) graph Laplacian eigenvectors.
- `𝛌::Vector{Float64}`: vector of eigenvalues.
- `Q::Matrix{Float64}`: unweighted incidence matrix of the graph.
- `α::Float64`: ROT parameter. (default: `1.0`)
- `T::Any`: TSD parameter, i.e., the stopping time T in K_functional (default: `:Inf`)
- `input_format::Symbol`: options: `:zero_measures`, `:pmf1` and `:pmf2` (default: `:zero_measures`)
- `distance::Symbol`: options: `:ROT`, `:HAD`, `:DAG` and `:TSD` (default: `:DAG`)
- `edg_length::Any`: vector of edge lengths (default: 1 represents unweighted graphs)
- `edge_weight::Any`: the weights vector, which stores the affinity weight of
    each edge (default: `1` represents unweighted graphs).

# Output Argument
- `dis::Matrix{Float64}`: the distance matrix, dis[i,j] = d(𝜙ᵢ₋₁, 𝜙ⱼ₋₁).

"""
function natural_eigdist(𝚽, 𝛌, Q; α = 1.0, T = :Inf,
                         input_format = :zero_measures, distance = :DAG,
                         edge_weight = 1, edge_length = 1)
    N = size(Q, 1)
    P = deepcopy(𝚽)
    if input_format == :zero_measures
        P[:, 1] .= 0
    elseif input_format == :pmf1
        P = P.^2
    elseif input_format == :pmf2
        P = exp.(P) ./ sum(exp.(P), dims = 1)
    else
        @error("input_format does not support $(input_format)!")
        return
    end

    if distance == :ROT
        D = eigROT_Distance(P, Q; edge_length = edge_length, α = α)
    elseif distance == :HAD
        D = eigHAD_Distance(𝚽, 𝛌)
    elseif distance == :DAG
        D = eigDAG_Distance(𝚽, Q, N; edge_weight = edge_weight)
    elseif distance == :TSD
        tL = Q * Q'
        t𝛌, t𝚽 = eigen(Matrix(tL))
        D = eigTSD_Distance(P, t𝚽, t𝛌, Q; length = edge_length, T = T)
    else
        error("distance does not support $(distance)!")
        return
    end

    return D
end


@testset "lanczos" begin
	@testset "givens" begin
		symOrth = KrylovMethods.symOrtho

		v     = ([0;randn()],[randn();0],randn(2))
		for k=1:3
			println("Givens rotation for v=[$(v[k][1]),$(v[k][2])]")
			c,s,r = symOrth(v[k][1],v[k][2])
			@test abs(c.^2 + s.^2 - 1) < 1e-14
			Q = [c s; -s c]
			w = [norm(v[k]);0]
			@test norm(Q*v[k]-w) < 1e-13
		end
	end
	@testset "bidiag" begin
		println("=== Testing Lanczos Bidiagonalization ===")
		A = sprandn(10,10,.4) + I
		p = randn(10)
		k = 9

		U,B,V = lanczosBidiag(A,p,k)

		@test norm(Matrix(U'*A*V-B))/norm(Matrix(A)) < 1e-15
	end

	@testset "tridiag" begin
		# test stopping for low-rank matrix
		A = randn(100,10)
		H = A*A'
		T,V = lanczosTridiag(H,ones(100),100,doReorth=true)
		@test size(T,1)==11
		@test size(V,2)==11

		# A   = matrixdepot("rosser",4)
		# Aff  = x -> A*x
		# T1,V1 = lanczosTridiag(A,ones(size(A,1)),size(A,1),tol=1e-10,doReorth=true)
		# T2,V2 = lanczosTridiag(Aff,ones(size(A,1)),size(A,1),tol=1e-10,doReorth=false)
		# sT1  = svd(Matrix(T1)).S
		# sT2  = svd(Matrix(T2)).S
		# sA   = svd(Matrix(A)).S
		#
		# @test abs(sT1[1]-sA[1])/sA[1] < 1e-4
		# @test abs(sT2[1]-sA[1])/sA[1] < 1e-4
		# B = V1'*V1
		# @test norm(B - Diagonal(diag(B))) < 1e-10
		# @test norm(diag(B)- ones(size(B,1))) < 1e-10
	end
end


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

"""
    euclidsq(x, y)

Return the Euclidean distance between two indexable objects.
"""
euclidsq(x, y) = sum((x[i] - y[i])^2 for i in eachindex(x))

# Support matrix data in a GPU and AD compatible way

"""
    euclidsq(X::T, Y::T) where {T<:AbstractMatrix}

Return the Euclidean distance between columns in two matrices.
"""
function euclidsq(X::T, Y::T) where {T<:AbstractMatrix}
    XiXj = transpose(X) * Y
    x² = sum(X .^ 2; dims=1)
    y² = sum(Y .^ 2; dims=1)
    transpose(x²) .+ y² - 2XiXj
end

"""
    euclidsq(X::T) where {T<:AbstractMatrix}

Effective version of `euclidsq(X, X)`.
"""
function euclidsq(X::T) where {T<:AbstractMatrix}
    XiXj = transpose(X) * X
    x² = sum(X .^ 2; dims=1)
    transpose(x²) .+ x² - 2XiXj
end

"""
    gaussian_gramian(xs, ys, σ=1)

Gramian matrix for samples `xs` and `ys` using a Gaussian kernel
kernel with bandwidth `σ`.
"""
gaussian_gramian(xs; kwargs...) = gaussian_gramian(xs, xs; kwargs...)
gaussian_gramian(xs, ys; σ=1) =
  [exp(-euclidsq(x, y) / 2σ^2) for x in xs, y in ys]

function gaussian_gramian(X::T, Y::T; σ=1) where {T<:AbstractMatrix}
    gaussian_gramian(euclidsq(X, Y), σ)
end

gaussian_gramian(esq, σ::AbstractFloat) = exp.(-esq ./ 2σ^2)

"""
    safe_diagm(mat, a)

Generate a squared matrix whose diagonal is `a` that is
compatible to perform addition on `mat`. It hebaves
differently based on `mat` is on CPU or GPU.

It is compatible with
- CuArrays.jl (see lib/cuarrays.jl)
- Zygote.jl (see lib/zygote.jl)
"""
safe_diagm(mat, a) = a * I


# These methods are used solely to help the auto-arima algorithm
# Kwiatkowski, D.;
# Phillips, P. C. B.;
# Schmidt, P.;
# Shin, Y. (1992)
# "Testing the null hypothesis
#  of stationarity against the
#  alternative of a unit root"

function calc_lag_kpss(lags::Bool, n::Int)
    return lags ? trunc(Int, 4*(n/100)^0.25) : trunc(Int, 12*(n/100)^0.25)
end

function σ²_estimator(l::Int, n::Int, ϵ::Vector{Fl}) where Fl
    return ((1/n)*sum(ϵ[t]^2 for t in 1:n)) + ((2/n)*sum((1-(s/(l + 1))) * (sum(ϵ[t]*ϵ[t - s] for t in s + 1:n)) for s in 1:l))
end

function p_value_weighted_average(η::Fl, p1::Fl, cv1::Fl,
                                  p2::Fl, cv2::Fl) where Fl
    return ((p1*(cv1 - cv2)) - ((cv1 - η)*(p1 - p2)))/(cv1 - cv2)
end

function p_value_from_η(η::Fl, crit_vals::Vector{Fl}) where Fl
    if η >= crit_vals[1]
        p_value = 0.01
    elseif η < crit_vals[1] && η >= crit_vals[2]
        p_value = p_value_weighted_average(η, 0.01, crit_vals[1], 0.025, crit_vals[2])
    elseif η < crit_vals[2] && η >= crit_vals[3]
        p_value = p_value_weighted_average(η, 0.025, crit_vals[2], 0.05, crit_vals[3])
    elseif η < crit_vals[3] && η <= crit_vals[4]
        p_value = p_value_weighted_average(η, 0.05, crit_vals[3], 0.1, crit_vals[4])
    else
        p_value = 0.1
    end
    return p_value
end

function calc_η_level(y::Vector{Fl}, lags::Bool) where Fl
    n = length(y)
    ϵ = y .- (mean(y))
    l = calc_lag_kpss(lags, n)
    η = (1/n^2)*sum(sum((((sum(ϵ[i] for i in 1:t))^2)/σ²_estimator(l, n, ϵ)) for t in 1:n))
    return η
end

function kpss_test(y::Vector{Fl}; lags::Bool = true) where Fl
    η = calc_η_level(y, lags)
    crit_vals = [0.739, 0.574, 0.463, 0.347]
    p_value = p_value_from_η(η, crit_vals)
    return p_value
end


using Test

function naive!(C, A, B)
    row, col = size(C)
    colA = size(A)[2]
    for i in 1:row
        for j in 1:col
            for k in 1:colA
                C[i, j] = C[i, j] + A[i, k] * B[k, j]
            end
        end
    end
end

function unroll_k!(C, A, B)
    row, col = size(C)
    colA = size(A)[2]
    for i in 1:row
        for j in 1:col
            for k in 1:2:colA
                C[i, j] = C[i, j] + A[i, k] * B[k, j] + A[i, k+1] * B[k+1, j]
            end
        end
    end
end

function unroll_j!(C, A, B)
    row, col = size(C)
    colA = size(A)[2]
    for i in 1:row
        for j in 1:2:col
            for k in 1:colA
                C[i,   j] = C[i, j] + A[i, k] * B[k, j]
                C[i, j+1] = C[i, j+1] + A[i, k] * B[k, j+1]
            end
        end
    end
end

function unroll_i!(C, A, B)
    row, col = size(C)
    colA = size(A)[2]
    for i in 1:2:row
        for j in 1:col
            for k in 1:colA
                C[i,   j] = C[i,   j] + A[i,   k] * B[k, j]
                C[i+1, j] = C[i+1, j] + A[i+1, k] * B[k, j]
            end
        end
    end
end

function unroll_ij!(C, A, B)
    row, col = size(C)
    colA = size(A)[2]
    for i in 1:2:row
        for j in 1:2:col
            for k in 1:colA
                C[i ,  j] = C[i,   j] + A[i, k] * B[k, j]
                C[i+1, j] = C[i+1, j] + A[i+1, k] * B[k ,j]
                C[i, j+1] = C[i, j+1] + A[i, k] * B[k, j+1]
                C[i+1, j+1] = C[i+1, j+1] + A[i+1, k] * B[k, j+1]
            end
        end
    end
end

function main(N)
    A = reshape(convert.(Float64, collect(1:N^2)), N, N)
    B = reshape(convert.(Float64, collect(1:N^2)), N, N)
    C = zeros(N, N)
    @time refC = A * B
    C = zeros(N, N)
    @time naive!(C, A, B)
    C = zeros(N ,N)
    @time unroll_k!(C, A, B)
    @test C ≈ refC
    C = zeros(N, N)
    @time unroll_i!(C, A, B)
    @test C ≈ refC
    C = zeros(N, N)
    @time unroll_j!(C, A, B)
    @test C ≈ refC
    C = zeros(N, N)
    @time unroll_ij!(C, A, B)
    @test C ≈ refC
end


module Mandelbrot
export mandelbrot_image, mandelbrot

using Base.Threads

function mandelbrot(c::Complex, max_iter::Int)
    z = 0.0 + 0.0im
    for i in 1:max_iter
        z = z^2 + c
        if abs(z) > 2
            return max_iter - i
        end
    end
    return 0
end

function mandelbrot_image(width, height, max_iter = 1000; rmin = -2, rmax = 1, imin = -1, imax = 1)
    A = [
        Complex((j/width * (rmax - rmin) + rmin), (i/height * (imax - imin) + imin))
        for i in 1:height, j in 1:width
    ]

    set = Array{Int}(undef, height, width)

    @time @inbounds @threads for i in 1:height
        for j in 1:width
            set[i, j] = mandelbrot(A[i, j], max_iter)
        end
    end

    return set
end

end


firstPayment = 3000
i = 0.12
m = 12

# (A/P, i, N)
function capitalRecoveryFactor(i, N)
	(i * (1 + i)^N) / ((i+1)^N - 1)
end

# Find: How much the guy paid the student
P = firstPayment / (1 + i/m)^(m * 1/12)
A = P * capitalRecoveryFactor(i/m, 12)
C = A * 1 / capitalRecoveryFactor(i/m, 5) + A + 1000

println("Guy paid: ", C)


function hermite_polynomial_coefficient_matrix(left,right)
    [1.0   left   left^2   left^3
     0.0   1.0   2.0*left  3.0*(left^2)
     1.0   right  right^2  right^3
     0.0   1.0   2.0*right 3.0*right^2]
end

function hermite_polynomial_coefficients(left,right)
    M = hermite_polynomial_coefficient_matrix(left,right)
    coeffs = M\I
end

function hermite_polynomials(x,left,right)
    coeffs = hermite_polynomial_coefficients(left,right)
    order,numbasis = size(coeffs)
    return [polynomial_from_coefficients(x,coeffs[:,i]) for i = 1:numbasis]
end


"""
FLOYDWARSHALL
-------------
    compute all shortest paths using the Floyd-Warshall algorithm.

    (D,P) = floydwarshall(A) returns the shortest distance matrix between all pairs
    of nodes in the graph A in matrix D.  If A has a negative weight cycle, then this
    algorithm will throw an error. P is the matrix of predecessors.

Functions
---------
- (D,P) = floydwarshall(A::MatrixNetwork)
- (D,P) = floydwarshall{T}(A::SparseMatrixCSC{T,Int64})

Example
-------
~~~
A = load_matrix_network("all_shortest_paths_example")
(D,P) = floydwarshall(A)
~~~
"""
:floydwarshall

## setup functions:

function floydwarshall_phase1(A::MatrixNetwork)
    (nzi,nzj,nzv) = csr_to_sparse(A.rp,A.ci,A.vals,A.n)
    return (nzi,nzj,nzv,A.n)
end

function floydwarshall_phase1{T}(A::SparseMatrixCSC{T,Int64})
    (ri,ci,ai) = findnz(A)
    return (ri,ci,ai,A.n)
end

function floydwarshall_phase2{T}(ri::Vector{Int64},ci::Vector{Int64},ai::Vector{T},n::Int64)

    nz = length(ai)
    D = Inf*ones(Int64,n,n)

    #TODO: check: always compute P or give the option of just computing D?

    P = zeros(Int64,n,n)
    # initialize the distance and predecessor matrix
    for ei = 1:nz
        i = ri[ei]
        j = ci[ei]
        v = ai[ei]
        if v < D[i,j]
            D[i,j] = v
            P[i,j] = i
        end
    end

    ids = sub2ind((n,n),1:n,1:n)
    D[ids] = 0 # set diagonal to 0

    for k=1:n
        for i=1:n
            for j=1:n
                if D[i,k]+D[k,j] < D[i,j]
                    D[i,j] = D[i,k]+D[k,j]
                    P[i,j] = P[k,j]
                end
            end
        end
    end

    if any(diag(D).<0)
        warn("floydwarshall:negativeCycle","negative weight cycle detected")
    end

    return (D,P)
end

## floyd warshall
function floydwarshall(A::MatrixNetwork)
    (nzi,nzj,nzv,n) = floydwarshall_phase1(A)
    (D,P) = floydwarshall_phase2(nzi,nzj,nzv,n)
    return (D,P)
end

function floydwarshall{T}(A::SparseMatrixCSC{T,Int64})
    (nzi,nzj,nzv,n) = floydwarshall_phase1(A)
    (D,P) = floydwarshall_phase2(nzi,nzj,nzv,n)
    return (D,P)
end


using Turing, Distributions, DifferentialEquations

# Import MCMCChain, Plots, and StatsPlots for visualizations and diagnostics.
using MCMCChains, Plots, StatsPlots

# Set a seed for reproducibility.
using Random
Random.seed!(14);

function lotka_volterra(du,u,p,t)
  x, y = u
  α, β, γ, δ  = p
  du[1] = (α - β*y)x # dx =
  du[2] = (δ*x - γ)y # dy =
end
p = [1.5, 1.0, 3.0, 1.0]
u0 = [1.0,1.0]
prob1 = ODEProblem(lotka_volterra,u0,(0.0,10.0),p)
sol = solve(prob1,Tsit5())
plot(sol)

sol1 = solve(prob1,Tsit5(),saveat=0.1)
odedata = Array(sol1) + 0.8 * randn(size(Array(sol1)))
plot(sol1, alpha = 0.3, legend = false); scatter!(sol1.t, odedata')

Turing.setadbackend(:forwarddiff)

@model function fitlv(data, prob1)
    σ ~ InverseGamma(2, 3) # ~ is the tilde character
    α ~ truncated(Normal(1.5,0.5),0.5,2.5)
    β ~ truncated(Normal(1.2,0.5),0,2)
    γ ~ truncated(Normal(3.0,0.5),1,4)
    δ ~ truncated(Normal(1.0,0.5),0,2)

    p = [α,β,γ,δ]
    prob = remake(prob1, p=p)
    predicted = solve(prob,Tsit5(),saveat=0.1)

    for i = 1:length(predicted)
        data[:,i] ~ MvNormal(predicted[i], σ)
    end
end

model = fitlv(odedata, prob1)

# This next command runs 3 independent chains without using multithreading.
chain = mapreduce(c -> sample(model, NUTS(.65),1000), chainscat, 1:3)

plot(chain)

pl = scatter(sol1.t, odedata');

chain_array = Array(chain)
for k in 1:300
    resol = solve(remake(prob1,p=chain_array[rand(1:1500), 1:4]),Tsit5(),saveat=0.1)
    plot!(resol, alpha=0.1, color = "#BBBBBB", legend = false)
end
# display(pl)
plot!(sol1, w=1, legend = false)

@model function fitlv2(data, prob1) # data should be a Vector
    σ ~ InverseGamma(2, 3) # ~ is the tilde character
    α ~ truncated(Normal(1.5,0.5),0.5,2.5)
    β ~ truncated(Normal(1.2,0.5),0,2)
    γ ~ truncated(Normal(3.0,0.5),1,4)
    δ ~ truncated(Normal(1.0,0.5),0,2)

    p = [α,β,γ,δ]
    prob = remake(prob1, p=p)
    predicted = solve(prob,Tsit5(),saveat=0.1)

    for i = 1:length(predicted)
        data[i] ~ Normal(predicted[i][2], σ) # predicted[i][2] is the data for y - a scalar, so we use Normal instead of MvNormal
    end
end

model2 = fitlv2(odedata[2,:], prob1)

Threads.nthreads()

# This next command runs 3 independent chains with multithreading.
chain2 = sample(model2, NUTS(.45), MCMCThreads(), 5000, 3, progress=false)

pl = scatter(sol1.t, odedata');
chain_array2 = Array(chain2)
for k in 1:300
    resol = solve(remake(prob1,p=chain_array2[rand(1:12000), 1:4]),Tsit5(),saveat=0.1)
    # Note that due to a bug in AxisArray, the variables from the chain will be returned always in
    # the order it is stored in the array, not by the specified order in the call - :α, :β, :γ, :δ
    plot!(resol, alpha=0.1, color = "#BBBBBB", legend = false)
end
#display(pl)
plot!(sol1, w=1, legend = false)

function delay_lotka_volterra(du, u, h, p, t)
   x, y = u
   α, β, γ, δ = p
   du[1] = α * h(p, t-1; idxs=1) - β * x * y
   du[2] = -γ * y + δ * x * y
   return
end

p = (1.5,1.0,3.0,1.0)
u0 = [1.0; 1.0]
tspan = (0.0,10.0)
h(p, t; idxs::Int) = 1.0
prob1 = DDEProblem(delay_lotka_volterra,u0,h,tspan,p)

sol = solve(prob1,saveat=0.1)
ddedata = Array(sol)
ddedata = ddedata + 0.5 * randn(size(ddedata))

scatter(sol.t, ddedata'); plot!(sol)

Turing.setadbackend(:forwarddiff)
@model function fitlv(data, prob1)

    σ ~ InverseGamma(2, 3)
    α ~ Truncated(Normal(1.5,0.5),0.5,2.5)
    β ~ Truncated(Normal(1.2,0.5),0,2)
    γ ~ Truncated(Normal(3.0,0.5),1,4)
    δ ~ Truncated(Normal(1.0,0.5),0,2)

    p = [α,β,γ,δ]

    #prob = DDEProblem(delay_lotka_volterra,u0,_h,tspan,p)
    prob = remake(prob1, p=p)
    predicted = solve(prob,saveat=0.1)
    for i = 1:length(predicted)
        data[:,i] ~ MvNormal(predicted[i], σ)
    end
end;
model = fitlv(ddedata, prob1)

chain = sample(model, NUTS(.65), MCMCThreads(), 300, 3, progress=true)
plot(chain)

chain

pl = scatter(sol.t, ddedata')
chain_array = Array(chain)
for k in 1:100
    resol = solve(remake(prob1,p=chain_array[rand(1:450),1:4]),Tsit5(),saveat=0.1)
    # Note that due to a bug in AxisArray, the variables from the chain will be returned always in
    # the order it is stored in the array, not by the specified order in the call - :α, :β, :γ, :δ

    plot!(resol, alpha=0.1, color = "#BBBBBB", legend = false)
end
#display(pl)
plot!(sol)

using Zygote, DiffEqSensitivity
Turing.setadbackend(:zygote)
prob1 = ODEProblem(lotka_volterra,u0,(0.0,10.0),p)

@model function fitlv(data, prob)
    σ ~ InverseGamma(2, 3)
    α ~ truncated(Normal(1.5,0.5),0.5,2.5)
    β ~ truncated(Normal(1.2,0.5),0,2)
    γ ~ truncated(Normal(3.0,0.5),1,4)
    δ ~ truncated(Normal(1.0,0.5),0,2)
    p = [α,β,γ,δ]
    prob = remake(prob, p=p)

    predicted = solve(prob,saveat=0.1)
    for i = 1:length(predicted)
        data[:,i] ~ MvNormal(predicted[i], σ)
    end
end;
model = fitlv(odedata, prob1)
chain = sample(model, NUTS(.65),1000)

@model function fitlv(data, prob)
    σ ~ InverseGamma(2, 3)
    α ~ truncated(Normal(1.5,0.5),0.5,2.5)
    β ~ truncated(Normal(1.2,0.5),0,2)
    γ ~ truncated(Normal(3.0,0.5),1,4)
    δ ~ truncated(Normal(1.0,0.5),0,2)
    p = [α,β,γ,δ]
    prob = remake(prob, p=p)
    predicted = solve(prob,saveat=0.1,sensealg=InterpolatingAdjoint(autojacvec=ReverseDiffVJP(true)))
    for i = 1:length(predicted)
        data[:,i] ~ MvNormal(predicted[i], σ)
    end
end;
model = fitlv(odedata, prob1)
@time chain = sample(model, NUTS(.65),1000)

u0 = [1.0,1.0]
tspan = (0.0,10.0)
function multiplicative_noise!(du,u,p,t)
  x,y = u
  du[1] = p[5]*x
  du[2] = p[6]*y
end
p = [1.5,1.0,3.0,1.0,0.1,0.1]

function lotka_volterra!(du,u,p,t)
  x,y = u
  α,β,γ,δ = p
  du[1] = dx = α*x - β*x*y
  du[2] = dy = δ*x*y - γ*y
end

prob_sde = SDEProblem(lotka_volterra!,multiplicative_noise!,u0,tspan,p)

ensembleprob = EnsembleProblem(prob_sde)
@time data = solve(ensembleprob,SOSRI(),saveat=0.1,trajectories=1000)
plot(EnsembleSummary(data))

Turing.setadbackend(:forwarddiff)
@model function fitlv(data, prob)
    σ ~ InverseGamma(2,3)
    α ~ truncated(Normal(1.3,0.5),0.5,2.5)
    β ~ truncated(Normal(1.2,0.25),0.5,2)
    γ ~ truncated(Normal(3.2,0.25),2.2,4.0)
    δ ~ truncated(Normal(1.2,0.25),0.5,2.0)
    ϕ1 ~ truncated(Normal(0.12,0.3),0.05,0.25)
    ϕ2 ~ truncated(Normal(0.12,0.3),0.05,0.25)
    p = [α,β,γ,δ,ϕ1,ϕ2]
    prob = remake(prob, p=p)
    predicted = solve(prob,SOSRI(),saveat=0.1)

    if predicted.retcode != :Success
        Turing.acclogp!(_varinfo, -Inf)
    end
    for j in 1:length(data)
        for i = 1:length(predicted)
            data[j][i] ~ MvNormal(predicted[i],σ)
        end
    end
end;

model = fitlv(data, prob_sde)
chain = sample(model, NUTS(0.25), 5000, init_theta = [1.5,1.3,1.2,2.7,1.2,0.12,0.12])
plot(chain)

if isdefined(Main, :TuringTutorials)
    Main.TuringTutorials.tutorial_footer(WEAVE_ARGS[:folder], WEAVE_ARGS[:file])
end


module problem119

description = """
Digit power sum

The number 512 is interesting because it is equal to the sum of its digits raised to some power: 5 + 1 + 2 = 8, and 83 = 512. Another example of a number with this property is 614656 = 284.
We shall define an to be the nth term of this sequence and insist that a number must contain at least two digits to have a sum.
You are given that a2 = 512 and a10 = 614656.

Find a30.
"""

# sum the digits of an integer by repeated div/mod by 10
function digitsum(n :: Integer)
  total = 0
  while n > 0
    total += n % 10
    n = div(n, 10)
  end
  return total
end

# Find all numbers with are powers of their digit sums, up to a maximum number of digits
# Insight: the digit sum for any number < 10^N can be at most N*9
#          so just consider powers of i from 2 to 9N, checking digitsums
function digitPowerSums(maxDigits = 15)
  results = BigInt[]
  xMax = BigInt(10)^maxDigits
  for i = 2 : 9 * maxDigits
    x = BigInt(i)
    while x < xMax
      if x > 10 && digitsum(x) == i
        push!(results, x)
      end
      x *= i
    end
  end
  sort!(results)
  return results
end

using Base.Test

@test digitsum(1234) == (1+2+3+4)
a = digitPowerSums()
@test a[2] == 512
@test a[10] == 614656

end


# rmsd100
# root mean-squared difference between two arrays in %

export rmsd100

using LinearAlgebra: norm

"""
    rmsd = rmsd100(x, y ; mask)

Compute 100 * RMSD (root mean squared difference) between `x` and `y` within domain mask.

in
- `x` : array
- `y` : another array of same size

option:
- `mask::Array{Bool}` : domain over which to compute the RMSE; default `trues(size(x))`

out
- rmsd : rmsd of `x` vs `y` within `mask` in %
"""
function rmsd100(x::AbstractArray{<:Number}, y::AbstractArray{<:Number}
        ; mask::AbstractArray{Bool} = trues(size(x)))
    return 100 * norm(x[mask] - y[mask]) / sqrt(sum(mask))
end


# This file is part of the IntervalArithmetic.jl package; MIT licensed

"""
    setdiff(x::Interval, y::Interval)

Calculate the set difference `x ∖ y`, i.e. the set of values
that are inside the interval `x` but not inside `y`.

Returns an array of intervals.
The array may:

- be empty if `x ⊆ y`;
- contain a single interval, if `y` overlaps `x`
- contain two intervals, if `y` is strictly contained within `x`.
"""
function setdiff(x::Interval{T}, y::Interval{T}) where T
    intersection = x ∩ y

    isempty(intersection) && return [x]
    intersection == x && return [emptyinterval(T)]

    if x.lo == intersection.lo
        if isfinite(intersection.hi)
            return [unsafe_interval(intersection.hi, x.hi)]
        else
            return [emptyinterval(T)]
        end
    end
    if x.hi == intersection.hi
        if isfinite(intersection.lo)
            return [unsafe_interval(x.lo, intersection.lo)]
        else
            return [emptyinterval(T)]
        end
    end

    return [unsafe_interval(x.lo, y.lo),
            unsafe_interval(y.hi, x.hi)]
end
setdiff(x::ValidInterval, ::ValidInterval) = [x]

"""
    distance(a::Interval, b::Interval)

Hausdorff distance between two intervals.
"""
distance(a::Interval, b::Interval) = max(abs(a.lo-b.lo), abs(a.hi-b.hi))

const where_bisect = 0.49609375

"""
    bisect(X::Interval, α=0.49609375)

Split the interval `X` at position α; α=0.5 corresponds to the midpoint.
Returns a tuple of the new intervals.
"""
function bisect(X::Interval, α=where_bisect)
    @assert 0 ≤ α ≤ 1

    m = mid(X, α)

    return (unsafe_interval(X.lo, m), unsafe_interval(m, X.hi))
end

"""
    mid(a::Interval, α=0.5)

Find an intermediate point at a relative position `α`` in the interval `a`.
The default is the true midpoint at `α = 0.5`.

Assumes 0 ≤ α ≤ 1.

Warning: if the parameter `α = 0.5` is explicitly set, the behavior differs
from the default case if the provided `Interval` is not finite, since when
`α` is provided `mid` simply replaces `+Inf` (respectively `-Inf`) by `prevfloat(+Inf)`
(respecively `nextfloat(-Inf)`) for the computation of the intermediate point.
"""
function mid(a::Interval{T}, α) where T

    lo = (a.lo == -Inf ? nextfloat(T(-Inf)) : a.lo)
    hi = (a.hi == +Inf ? prevfloat(T(+Inf)) : a.hi)

    β = T(α)

    midpoint = β * (hi - lo) + lo
    isfinite(midpoint) && return midpoint
    #= Fallback in case of overflow: hi - lo == +Inf.
       This case can not be the default one as it does not pass several
       IEEE1788-2015 tests for small floats.
    =#
    return (1 - β) * lo + β * hi
end

mid(::EmptyInterval{T}, α) where T = T(NaN)


using SMC
using PyPlot
using Base.Test

@testset "logsumexp tests" begin

    # on vector
    vec = [1., 2.2, 3.1]
    actual = logsumexp(vec)
    expected = log(sum(exp.(vec)))
    @test isapprox(actual, expected)

    # on 2D array
    arr = [[1. 2. 3.]; [3.3 4.1 5.2]]
    # sum over rows
    actual = logsumexp(arr, 1)
    expected = [3.3955454645979626 4.215519523179755 5.305083319768697]
    @test isapprox(actual, expected)

    # sum over columns
    actual = logsumexp(arr, 2)
    expected = [3.40760596444438 5.593689178496464]'
    @test isapprox(actual, expected)
end

@testset "hmm tests" begin

    initial_state_prior = [0.4, 0.6]
    transition_model = [0.7 0.3;
                        0.1 0.9]
    observation_model = [0.6 0.3 0.1;
                         0.1 0.4 0.5]
    hmm = HiddenMarkovModel(initial_state_prior,
                            transition_model,
                            observation_model)
    observations = [1, 3, 2]

    @testset "forward probabilities" begin
        # p(x_t | y_{1:t-1}) for each t = 1,..,T
        actual = SMC.hmm_forward_pass(hmm, observations)
        # number of time steps x number of states
        @test size(actual) == (3, 2)

        # check against manually computed values
        expected = begin

            # p(x_2 = 1, y_1 = 1, y_2 = 3) = p(x_1 = 1, x_2 = 1, y_1 = 1, y_2 = 3) + p(x_1 = 2, x_2 = 1, y_1 = 1, y_2 = 3)
            shared_x2_eq_1 = ((0.4 * 0.7 * 0.6 * 0.1) + (0.6 * 0.1 * 0.1 * 0.1))

            # p(x_2 = 2, y_1 = 1, y_2 = 3) = p(x_1 = 1, x_2 = 2, y_1 = 1, y_2 = 3) + p(x_1 = 2, x_2 = 2, y_1 = 1, y_2 = 3)
            shared_x2_eq_2 = ((0.4 * 0.3 * 0.6 * 0.5) + (0.6 * 0.9 * 0.1 * 0.5)) #

            f11 = 0.4 # p(x_1 = 1)
            f12 = 0.6 # p(x_1 = 2)

            # p(x_2 = 1, y_1 = 1) = p(x_1 = 1, x_2 = 1, y_1 = 1) + p(x_1 = 2, x_2 = 1, y_1 = 1)
            f21 = (0.4 * 0.7 * 0.6) + (0.6 * 0.1 * 0.1)

            # p(x_2 = 2, y_1 = 1) = p(x_1 = 1, x_2 = 2, y_1 = 1) + p(x_1 = 2, x_2 = 2, y_1 = 1)
            f22 = (0.4 * 0.3 * 0.6) + (0.6 * 0.9 * 0.1)

            # p(x_3 = 1, y_1 = 1, y_2 = 3) = p(x_2 = 1, x_3 = 1, y_1 = 1, y_2 = 3) + p(x_2 = 2, x_3 = 1, y_1 = 1, y_2 = 3)
            f31 = shared_x2_eq_1 * 0.7 + shared_x2_eq_2 * 0.1

            # p(x_3 = 2, y_1 = 1, y_2 = 3) = p(x_2 = 1, x_3 = 2, y_1 = 1, y_2 = 3) + p(x_2 = 2, x_3 = 2, y_1 = 1, y_2 = 3)
            f32 = shared_x2_eq_1 * 0.3 + shared_x2_eq_2 * 0.9

            [f11 f12; f21 f22; f31 f32]
        end
        @test isapprox(actual, expected)

        # check log space agrees
        @test isapprox(exp.(SMC.hmm_log_forward_pass(hmm, observations)), expected)
    end

    @testset "marginal likelihoods" begin
        actual = SMC.hmm_marginal_likelihood(hmm, observations)

        expected = (
              (0.4 * 0.7 * 0.7 * 0.6 * 0.1 * 0.3) + # x = (1, 1, 1)
              (0.4 * 0.7 * 0.3 * 0.6 * 0.1 * 0.4) + # x = (1, 1, 2)
              (0.4 * 0.3 * 0.1 * 0.6 * 0.5 * 0.3) + # x = (1, 2, 1)
              (0.4 * 0.3 * 0.9 * 0.6 * 0.5 * 0.4) + # x = (1, 2, 2)
              (0.6 * 0.1 * 0.7 * 0.1 * 0.1 * 0.3) + # x = (2, 1, 1)
              (0.6 * 0.1 * 0.3 * 0.1 * 0.1 * 0.4) + # x = (2, 1, 2)
              (0.6 * 0.9 * 0.1 * 0.1 * 0.5 * 0.3) + # x = (2, 2, 1)
              (0.6 * 0.9 * 0.9 * 0.1 * 0.5 * 0.4)) # x = (2, 2, 2)

        @test isapprox(actual, expected)

        # check that log space agrees
        @test isapprox(actual, exp(hmm_log_marginal_likelihood(hmm, observations)))
    end

    @testset "posterior sampling" begin
        particle = hmm_posterior_sample(hmm, observations)

        # smoke test: check the type and dimensions and other basic properties
        @test typeof(particle) == Array{Int,1}
        @test length(particle) == 3
        @test all(particle .>= 1)
        @test all(particle .<= 2)

        # TODO: test the distribution? Maybe next year :)

    end

    @testset "sanity checks" begin

        # sanity check on extreme observation
        srand(1)
        hmm = HiddenMarkovModel([0.5, 0.5],                     # prior
                                [0.5 0.5; 0.5 0.5],             # transition
                                [0.9999 0.0001; 0.0001 0.9999]) # observation
        obs = [1]
        num_samples = 100
        num_1 = 0
        for i = 1:num_samples
            particle = hmm_posterior_sample(hmm, obs)
            num_1 += (particle[1] == 1)
        end
        @test num_1 >= 95

        # sanity check on extreme prior and transition
        srand(1)
        hmm = HiddenMarkovModel([0.0001, 0.9999],                     # prior
                                [0.9999 0.0001; 0.0001 0.9999],             # transition
                                [0.5 0.5; 0.5 0.5]) # observation
        obs = [1, 2]
        num_samples = 100
        num_1 = 0
        for i = 1:num_samples
            particle = hmm_posterior_sample(hmm, obs)
            num_1 += (particle[2] == 1)
        end
        @test num_1 <= 5

        # a symmetric extreme distribution
        hmm = HiddenMarkovModel([0.5, 0.5],         # prior
                                [0.9999 0.0001; 0.0001 0.9999],  # transition
                                [0.9999 0.0001; 0.0001 0.9999]) # observation

        # switch happens at step 9 -> 10
        obs_a = [1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2]
        obs_b = [2, 2, 2, 2, 2, 2, 2, 2, 2, 1, 1, 1, 1, 1, 1, 1, 1]

        # test for symmetry
        @test isapprox(hmm_log_marginal_likelihood(hmm, obs_a),
                       hmm_log_marginal_likelihood(hmm, obs_b))

        # check samples
        srand(1)
        num_samples = 100

        num_start_1 = 0
        num_end_2 = 0
        for i = 1:num_samples
            particle = hmm_posterior_sample(hmm, obs_a)
            num_start_1 += all(particle[1:9] .== 1)
            num_end_2 += all(particle[10:end] .== 2)
        end
        @test num_start_1 >= 95
        @test num_end_2 >= 95

        srand(1)
        num_start_2 = 0
        num_end_1 = 0
        for i = 1:num_samples
            particle = hmm_posterior_sample(hmm, obs_b)
            num_start_2 += all(particle[1:9] .== 2)
            num_end_1 += all(particle[10:end] .== 1)
        end
        @test num_start_2 >= 95
        @test num_end_1 >= 95
    end

end

@testset "smc tests" begin

    hmm = HiddenMarkovModel(
            [0.4, 0.6], # prior
            [0.7 0.3;
             0.1 0.9], # transition
            [0.6 0.3 0.1;
             0.1 0.4 0.5]) # observation
    observations = [1, 3, 2]

    # SMC with prior proposal
    srand(1)
    num_particles = 10000
    scheme = HMMPriorSMCScheme(hmm, observations, num_particles)
    output, log_ml_estimate = smc(scheme)
    @test length(output) == 3
    expected = hmm_log_marginal_likelihood(hmm, observations)
    @test isapprox(log_ml_estimate, expected, atol=0.1, rtol=0)

    # conditional SMC (take output from above as input)
    log_ml_estimate = conditional_smc(scheme, output)
    @test isapprox(log_ml_estimate, expected, atol=0.1, rtol=0)

    # SMC with conditional proposal
    srand(1)
    num_particles = 10000
    scheme = HMMConditionalSMCScheme(hmm, observations, num_particles)
    output, log_ml_estimate = smc(scheme)
    @test length(output) == 3
    expected = hmm_log_marginal_likelihood(hmm, observations)
    @test isapprox(log_ml_estimate, expected, atol=0.02, rtol=0)

    # conditional SMC (take output from above as input)
    log_ml_estimate = conditional_smc(scheme, output)
    @test isapprox(log_ml_estimate, expected, atol=0.1, rtol=0)
end

@testset "smc sandwiching the evidence with exact sampling" begin

    hmm = HiddenMarkovModel(
            [0.4, 0.6], # prior
            [0.7 0.3;
             0.1 0.9], # transition
            [0.6 0.3 0.1;
             0.1 0.4 0.5]) # observation
    observations = [1, 3, 2]

    # true marginal likelihood
    expected = hmm_log_marginal_likelihood(hmm, observations)

    function test_bounds(smc_scheme_constructor, all_num_particles::Array{Int,1},
                         num_reps::Int)
        lower_bound_raw_data = Array{Float64, 2}(length(all_num_particles), num_reps)
        upper_bound_raw_data = Array{Float64, 2}(length(all_num_particles), num_reps)
        for (i, num_particles) in enumerate(all_num_particles)
            scheme = HMMConditionalSMCScheme(hmm, observations, num_particles)
            for j = 1:num_reps
                output, lml_estimate = smc(scheme)
                lower_bound_raw_data[i, j] = lml_estimate
            end
            for j = 1:num_reps
                output = hmm_posterior_sample(hmm, observations)
                lml_estimate = conditional_smc(scheme, output)
                upper_bound_raw_data[i, j] = lml_estimate
            end
        end

        lower_bound_means = mean(lower_bound_raw_data, 2)
        lower_bound_stds = std(lower_bound_raw_data, 2)
        upper_bound_means = mean(upper_bound_raw_data, 2)
        upper_bound_stds = std(upper_bound_raw_data, 2)

        # they should be upper and lower bounds in all cases
        for i=1:length(all_num_particles)
            @test lower_bound_means[i] <= expected
            @test upper_bound_means[i] >= expected
        end

        # the estimates should get tighter and less noisy with more particles
        for i=1:length(all_num_particles)-1
            @test lower_bound_means[i] <= lower_bound_means[i+1]
            @test upper_bound_means[i] >= upper_bound_means[i+1]
            @test lower_bound_stds[i] >= lower_bound_stds[i+1]
            @test upper_bound_stds[i] >= upper_bound_stds[i+1]
        end

        # the final estimates should both be tight
        @test isapprox(lower_bound_means[end], expected, atol=0.01, rtol=0)
        @test isapprox(upper_bound_means[end], expected, atol=0.01, rtol=0)

    end

    # SMC with prior proposal
    srand(1)
    test_bounds(HMMConditionalSMCScheme, [1, 3, 10, 30, 100], 1000)

    # SMC with conditional proposal
    srand(1)
    test_bounds(HMMPriorSMCScheme, [1, 3, 10, 30, 100], 1000)

end

@testset "rendering hmm" begin
    hmm = HiddenMarkovModel(
        [0.4, 0.6], # prior
        [0.7 0.3;
         0.1 0.9], # transition
        [0.6 0.3 0.1;
         0.1 0.4 0.5]) # observation

    hmm_fname = "$(tempname()).png"
    latent_states_fname = "$(tempname()).png"
    observations_fname = "$(tempname()).png"

    # just check that rendering works in the current environment without crashign
    plt[:figure]()
    render_hmm!(hmm)
    plt[:savefig](hmm_fname)

    plt[:figure]()
    render_hmm_states!(hmm, [1, 1, 1, 2, 2, 1, 2, 1, 2, 2, 2], true, [1, 10], [1, 2])
    plt[:savefig](latent_states_fname)

    plt[:figure]()
    render_hmm_observations!(hmm, [1, 1, 1, 3, 3, 1, 2, 1, 2, 2, 2])
    plt[:savefig](observations_fname)

    Base.Filesystem.rm(hmm_fname)
    Base.Filesystem.rm(latent_states_fname)
    Base.Filesystem.rm(observations_fname)
end


function piecewise(x)
	IfElse.ifelse(x[1] < 0.5, 0.0,
		IfElse.ifelse(x[1] < 2.0, m_ss * x[1] - 0.5 * m_ss,
			-0.250 * m_ss * (x[1] - 2.0) + 1.5 * m_ss))
end

function d_piecewise(x)
	IfElse.ifelse(x[1] < 0.5, 0.0,
		IfElse.ifelse(x[1] < 2.0, m_ss,
			-0.250 * m_ss))
end

function mirror_and_repeat(x; l::T = 5.0) where {T}
	# l = length of the pattern  ---^_/^-
	# we repeat it this way ---^_/^--^\_^------^_/^--^\_^---
	# section ≡ even -> normal pattern
	# section ≡ odd -> mirrored pattern
	section = floor(x / l)
	x_rel = x % l
	mir = (section % 2) == 1
	if !mir
		x_mir = x_rel
	elseif mir # mirroring
		x_mir = l - x_rel
	end
	mult = mir == 1 ? -1 : 1
	return x_mir, mult
end

# smooth kinks w/ cubic polynomial
poly(a, z) = a[4] + a[3] * z + a[2] * z^2.0 + a[1] * z^3.0
d_poly(a, z) = a[3] + 2.0 * a[2] * z + 3.0 * a[1] * z^2.0

function generate_piecewise_terrain(mss; repeat::Bool = false)
	# piece 1
	m1 = 0.0
	x1 = 0.4
	y1 = m1 * x1

	m2 = m_ss
	x2 = 0.6
	y2 = m2 * 0.1

	A1 = [x1^3.0 x1^2.0 x1 1.0;
	      x2^3.0 x2^2.0 x2 1.0;
		  3.0 * x1^2.0 2.0 * x1 1.0 0.0;
		  3.0 * x2^2.0 2.0 * x2 1.0 0.0]
	b1 = [y1; y2; m1; m2]

	a1 = A1 \ b1

	@assert isapprox(poly(a1, x1) - y1, 0.0, atol = 1.0e-8)
	@assert isapprox(poly(a1, x2) - y2, 0.0, atol = 1.0e-8)

	# piece 2
	m1 = m_ss
	x1 = 1.4
	y1 = m_ss * x1

	m2 = -0.250 * m_ss
	x2 = 1.6
	y2 = m_ss * 1.5 + m2 * 0.1

	A2 = [x1^3.0 x1^2.0 x1 1.0;
	      x2^3.0 x2^2.0 x2 1.0;
		  3.0 * x1^2.0 2.0 * x1 1.0 0.0;
		  3.0 * x2^2.0 2.0 * x2 1.0 0.0]
	b2 = [y1; y2; m1; m2]

	a2 = A2 \ b2

	@assert isapprox(poly(a2, x1) - y1, 0.0, atol = 1.0e-8)
	@assert isapprox(poly(a2, x2) - y2, 0.0, atol = 1.0e-8)

	if !repeat
		function piecewise_smoothed(x)
			IfElse.ifelse(x[1] < 0.4, 0.0,
				IfElse.ifelse(x[1] < 0.6, poly(a1, x[1]),
					IfElse.ifelse(x[1] < 1.9, m_ss * x[1] - 0.5 * m_ss,
						IfElse.ifelse(x[1] < 2.1, poly(a2, x[1] - 0.5),
						-0.250 * m_ss * (x[1] - 2.0) + 1.5 * m_ss))))
		end

		function d_piecewise_smoothed(x)
			IfElse.ifelse(x[1] < 0.4, 0.0,
				IfElse.ifelse(x[1] < 0.6, d_poly(a1, x[1]),
					IfElse.ifelse(x[1] < 1.9, m_ss,
						IfElse.ifelse(x[1] < 2.1, d_poly(a2, x[1] - 0.5),
						-0.250 * m_ss))))
		end
		return piecewise_smoothed, d_piecewise_smoothed
	elseif repeat
		@show "repp"
		function piecewise_repeat_smoothed(x)
			x_mir, mult = mirror_and_repeat(x[1]; l = 5.0)
			IfElse.ifelse(x_mir < 0.4, mult * 0.0,
				IfElse.ifelse(x_mir < 0.6, mult * poly(a1, x_mir),
					IfElse.ifelse(x_mir < 1.9, mult * (m_ss * x_mir - 0.5 * m_ss),
						IfElse.ifelse(x_mir < 2.1, mult * poly(a2, x_mir - 0.5),
						mult * (-0.250 * m_ss * (x_mir - 2.0) + 1.5 * m_ss)))))
		end

		function d_piecewise_repeat_smoothed(x)
			x_mir, mult = mirror_and_repeat(x[1]; l = 5.0)
			IfElse.ifelse(x_mir < 0.4, 0.0,
				IfElse.ifelse(x_mir < 0.6, d_poly(a1, x_mir),
					IfElse.ifelse(x_mir < 1.9, m_ss,
						IfElse.ifelse(x_mir < 2.1, d_poly(a2, x_mir - 0.5),
						-0.250 * m_ss))))
		end
		return piecewise_repeat_smoothed, d_piecewise_repeat_smoothed
	end
end

# x = range(-1.0, stop = 4.0, length = 1000)
# plot(x, piecewise.(x))#, aspect_ratio = :equal)
# plot!(x, piecewise_smoothed.(x))#, aspect_ratio = :equal)
#
# plot(x, d_piecewise.(x), aspect_ratio = :equal)
# plot!(x, d_piecewise_smoothed.(x), aspect_ratio = :equal)

m_ss = tan(deg2rad(10.0)) # 10 degree slope
p1, dp1 = generate_piecewise_terrain(m_ss, repeat = false)
piecewise1_2D_lc = Environment{R2, LinearizedCone}(p1, dp1)
plot(-1:0.01:15, piecewise1_2D_lc.surf.(-1:0.01:15))
plot!(-1:0.01:15, piecewise1_2D_lc.surf_grad.(-1:0.01:15))

m_ss = tan(deg2rad(-15.0)) # 10 degree slope
p2, dp2 = generate_piecewise_terrain(m_ss)
piecewise2_2D_lc = Environment{R2, LinearizedCone}(p2, dp2)
# plot(-1:0.01:15, piecewise2_2D_lc.surf.(-1:0.01:15))
# plot!(-1:0.01:15, piecewise2_2D_lc.surf_grad.(-1:0.01:15))

# vis = Visualizer()
# open(vis)
# plot_surface!(vis, piecewise1_2D_lc, n = 100)
# plot_surface!(vis, piecewise2_2D_lc, n = 100)