import StatsBase: countmap
function raw_majority_vote(labels)
    votes = countmap(labels)
    votes = sort(collect(votes), by=x->x[2], rev=true)
    winner = votes[1][1]
    return winner
end

@assert raw_majority_vote(['a', 'b', 'c', 'b']) == 'b'

function majority_vote(labels)
    """Assumes that labels are ordered from nearest to farthest."""
    vote_counts = countmap(labels)
    vote_counts = sort(collect(vote_counts), by=x->x[2], rev=true)
    winner = vote_counts[1][1]
    if(length(vote_counts) == 1)
        return winner
    end

    winner_count= vote_counts[1][2]
    second_count = vote_counts[2][2]
    # num_winners = length([count
                    #    for count in vote_counts.values()
                    #    if count == winner_count])

    if (winner_count != second_count)
        return winner                     # unique winner, so return it
    else
        return majority_vote(labels[1:length(labels)-1]) # try again without the farthest
    end
end

# Tie, so look at first 4, then 'b'
@assert majority_vote(['a', 'b', 'c', 'b', 'a'])== 'b'

using Distances

struct LabeledPoint
    point
    label
    LabeledPoint(point, label) = new(point, label)
end

function knn_classify(k, labeled_points, new_point)

    by_distance = sort(labeled_points, by = x -> SqEuclidean()(x.point, new_point))

    k_nearest_labels = [lp.label for lp in by_distance[1:k]]

    return majority_vote(k_nearest_labels)
end

function parse_iris_row(line)
    """
    sepal_length, sepal_width, petal_length, petal_width, class
    """
    row = split(line, ",")
    measurements = [parse(Float64, value) for value in row[1:end-1]]
    # class is e.g. "Iris-virginica"; we just want "virginica"
    label = split(row[end], "-")[end]

    return LabeledPoint(measurements, label)
end

data = open(f->read(f,String), "iris.data")
irislines = []
for i in split(data, "\n")
    push!(irislines,i)
end

iris_data = [parse_iris_row(line) for line in irislines]
using Random
Random.seed!(12)
function split_data(origdata, prob)
    """Split data into fractions [prob, 1 - prob]"""
    data = copy(origdata)                    # Make a shallow copy
    data = shuffle(data)              # because shuffle modifies the list.
    cut = Int(floor(length(data) * prob))       # Use prob to find a cutoff
    return data[1:cut], data[cut+1:end]     # and split the shuffled list there.
end

iris_train, iris_test = split_data(iris_data, 0.70)
@assert length(iris_train) == 0.7*150
@assert length(iris_test) == 0.3 * 150

confusion_matrix = Dict()
let num_correct = 0

    for iris in iris_test
        if(iris.label == "virginica")
            # continue
        end
        predicted = knn_classify(5, iris_train, iris.point)
        actual = iris.label

        if predicted == actual
            num_correct += 1
        end

        confusion_matrix[(predicted, actual)] = get(confusion_matrix,(predicted, actual),0) + 1
    end

    pct_correct = num_correct / length(iris_test)

    print(pct_correct)
    print( confusion_matrix)

end


using PoissonRandom
import Distributions
using Test, Statistics

n_tsamples = 10^5

function test_samples(rand_func,
                      distr::Distributions.DiscreteUnivariateDistribution,
                      n::Int;                                # number of samples to generate
                      q::Float64=1.0e-8,                     # confidence interval, 1 - q as confidence
                      verbose::Bool=false)                   # show intermediate info (for debugging)

    # The basic idea
    # ------------------
    #   Generate n samples, and count the occurences of each value within a reasonable range.
    #   For each distinct value, it computes an confidence interval of the counts
    #   and checks whether the count is within this interval.
    #
    #   If the distribution has a bounded range, it also checks whether
    #   the samples are all within this range.
    #
    #   By setting a small q, we ensure that failure of the tests rarely
    #   happen in practice.
    #
    λ = distr.λ
    n > 1 || error("The number of samples must be greater than 1.")
    0.0 < q < 0.1 || error("The value of q must be within the open interval (0.0, 0.1).")

    # determine the range of values to examine
    vmin = minimum(distr)
    vmax = maximum(distr)

    rmin = floor(Int,quantile(distr, 0.00001))::Int
    rmax = floor(Int,quantile(distr, 0.99999))::Int
    m = rmax - rmin + 1  # length of the range
    p0 = Distributions.pdf.((distr,), rmin:rmax)  # reference probability masses
    @assert length(p0) == m

    # determine confidence intervals for counts:
    # with probability q, the count will be out of this interval.
    #
    clb = Vector{Int}(undef,m)
    cub = Vector{Int}(undef,m)
    for i = 1:m
        bp = Distributions.Binomial(n, p0[i])
        clb[i] = floor(Int,quantile(bp, q/2))
        cub[i] = ceil(Int,Distributions.cquantile(bp, q/2))
        @assert cub[i] >= clb[i]
    end

    # generate samples
    samples = [rand_func(λ) for i in 1:n]
    @assert length(samples) == n

    # scan samples and get counts
    cnts = zeros(Int, m)
    for i = 1:n
        @inbounds si = samples[i]
        if rmin <= si <= rmax
            cnts[si - rmin + 1] += 1
        else
            vmin <= si <= vmax ||
                error("Sample value out of valid range.")
        end
    end

    # check the counts
    for i = 1:m
        verbose && println("v = $(rmin+i-1) ==> ($(clb[i]), $(cub[i])): $(cnts[i])")
        clb[i] <= cnts[i] <= cub[i] ||
            error("The counts are out of the confidence interval.")
    end
    return samples
end

println("testing count random sampler")
for λ in [0.2, 0.5, 1.0, 2.0, 5.0, 10.0, 15.0, 20.0, 30.0]
  test_samples(PoissonRandom.count_rand, Distributions.Poisson(λ), n_tsamples)
end

println("testing ad random sampler")
for λ in [5.0, 10.0, 15.0, 20.0, 30.0]
  test_samples(PoissonRandom.ad_rand, Distributions.Poisson(λ), n_tsamples)
end

println("testing mixed random sampler")
for λ in [5.0, 10.0, 15.0, 20.0, 30.0]
  test_samples(pois_rand, Distributions.Poisson(λ), n_tsamples)
end


"""
    drawpointorbitR2(f, x0, y0 [; preiterations, iterations])

Return the drawing of the orbit of a point \$(x_0,y_0)\$ under a function
    \$f:\\mathbb{R}^2\\rightarrow\\mathbb{R}^2\$,
    in a rectangular region in \$\\mathbb{R}^2\$.

The orbit of \$(x_0,y_0)\$ under \$f\$ is defined as
    \$o(x_0,y_0,f)=\\{(x_0,y_0),f(x_0,y_0),\\dots,f^n(x_0,y_0),\\dots\\}\$

#### Arguments
- `f::Function`: A function \$f:\\mathbb{R}^2\\rightarrow\\mathbb{R}^2\$.
- `x0::Real`: X coordinate of the point.
- `y0::Real`: Y coordinate of the point.
- `preiterations::Integer`: Number of first iterations to calculate but not to draw.
- `iterations::Integer`: Number of iterations to calculate (after `preiterations`).
"""
function drawpointorbitR2(f::Function, x0::Real, y0::Real;
    preiterations::Integer=0, iterations::Integer=100)

    # Verifying functions
    @assert typeof(f(1., 1.)) <: Tuple{Real,Real}

    SDDGraphics.newdrawing()
    SDDGraphics.updatecolorarray(iterations)

    xn,yn = x0,y0

    for n in 1:preiterations
        xn,yn = f(xn,yn)
    end # for n preiterations

    for n in 1:iterations
        if SDDGraphics.insiderectregion(xn,yn)
            SDDGraphics.color(n)
            SDDGraphics.drawpoint(xn,yn)
        end # if
        xn,yn = f(xn,yn)
    end # for n iterations

    SDDGraphics.drawing()

end # function drawpointorbitR2

"""
    drawpointssetorbitR2(f, ps [; preiterations, iterations])

Return the drawing of the orbit of a points set under a function
    \$f:\\mathbb{R}^2\\rightarrow\\mathbb{R}^2\$,
    in a rectangular region in \$\\mathbb{R}^2\$.

The orbit of the set \$P\$ under \$f\$ is defined as
    \$o(P,f)=\\bigcup_{n\\geq 0}f^n(P)\$.

#### Arguments
- `f::Function`: A function \$f:\\mathbb{R}^2\\rightarrow\\mathbb{R}^2\$.
- `ps::Array{Tuple{Real,Real},1}`: The set of points (array of pairs of real numbers).
- `preiterations::Integer`: Number of first iterations to calculate but not to draw.
- `iterations::Integer`: Number of iterations to calculate (after `preiterations`).
"""
function drawpointssetorbitR2(f::Function, ps::Array{Tuple{T,T},1};
    preiterations::Integer=0, iterations::Integer=100) where T <: Real

    # Verifying functions
    @assert typeof(f(1., 1.)) <: Tuple{Real,Real}

    SDDGraphics.newdrawing()
    SDDGraphics.updatecolorarray(iterations)

    pns = deepcopy(ps)

    for n in 1:preiterations
        for k in 1:length(pns)
            pns[k] = f(pns[k]...)
        end # for k pns
    end # for n preiterations

    for n in 1:iterations
        SDDGraphics.color(n)
        for k in 1:length(pns)
            if SDDGraphics.insiderectregion(pns[k]...)
                SDDGraphics.drawpoint(pns[k]...)
            end # if
            pns[k] = f(pns[k]...)
        end # for k pns
    end # for n iterations

    SDDGraphics.drawing()

end # function drawpointorbitR2

"""
    drawpointorbitC(f, z0 [; preiterations, iterations])

Return the drawing of the orbit of a point \$z_0\$ under a function
    \$f:\\mathbb{C}\\rightarrow\\mathbb{C}\$,
    in a rectangular region in \$\\mathbb{C}\$.

The orbit of \$z_0\$ under \$f\$ is defined as
    \$o(z_0,f)=\\{z_0,f(z_0),\\dots,f^n(z_0),\\dots\\}\$.

#### Arguments
- `f::Function`: A function \$f:\\mathbb{C}\\rightarrow\\mathbb{C}\$.
- `z0::Number`: The point.
- `preiterations::Integer`: Number of first iterations to calculate but not to draw.
- `iterations::Integer`: Number of iterations to calculate (after `preiterations`).
"""
function drawpointorbitC(f::Function, z0::Number;
    preiterations::Integer=0, iterations::Integer=100)

    # Verifying functions
    @assert typeof(f(1.0im)) <: Number

    SDDGraphics.newdrawing()
    SDDGraphics.updatecolorarray(iterations)

    zn = z0

    for n in 1:preiterations
        zn = f(zn)
    end # for n preiterations

    for n in 1:iterations
        if SDDGraphics.insiderectregion(zn)
            SDDGraphics.color(n)
            SDDGraphics.drawpoint(zn)
        end #if
        zn = f(zn)
    end # for n iterations

    SDDGraphics.drawing()

end # function drawpointorbitC

"""
    drawpointssetorbitC(f, ps [; preiterations, iterations])

Return the drawing of the orbit of a points set under a function
    \$f:\\mathbb{C}\\rightarrow\\mathbb{C}\$,
    in a rectangular region in \$\\mathbb{C}\$.

The orbit of the set \$P\$ under \$f\$ is defined as
    \$o(P,f)=\\bigcup_{n\\geq 0}f^n(P)\$.

#### Arguments
- `f::Function`: A function \$f:\\mathbb{C}\\rightarrow\\mathbb{C}\$.
- `ps::Array{Number,1}`: The set of points (array of numbers).
- `preiterations::Integer`: Number of first iterations to calculate but not to draw.
- `iterations::Integer`: Number of iterations to calculate (after `preiterations`).
"""
function drawpointssetorbitC(f::Function, ps::Array{T,1};
    preiterations::Integer=0, iterations::Integer=100) where T <: Number

    # Verifying functions
    @assert typeof(f(1.0im)) <: Number

    SDDGraphics.newdrawing()
    SDDGraphics.updatecolorarray(iterations)

    pns = deepcopy(ps)

    for n in 1:preiterations
        pns = f.(pns)
    end # for n preiterations

    for n in 1:iterations
        SDDGraphics.color(n)
        for k in 1:length(pns)
            if SDDGraphics.insiderectregion(pns[k])
                SDDGraphics.drawpoint(pns[k])
            end # if
            pns[k] = f(pns[k])
        end # for k pns
    end # for n iterations

    SDDGraphics.drawing()

end # function drawpointorbitR2

"""
    drawpointorbitpathR2(f, x0, y0 [; preiterations, iterations, kind])

Return the drawing of the orbit's path of a point \$(x_0,y_0)\$ under a function
\$f:\\mathbb{R}^2\\rightarrow\\mathbb{R}^2\$,
in a rectangular region in \$\\mathbb{R}^2\$.

The orbit of \$(x_0,y_0)\$ under \$f\$ is defined as
\$o(x_0,y_0,f)=\\{(x_0,y_0),f(x_0,y_0),\\dots,f^n(x_0,y_0),\\dots\\}\$.

#### Arguments
- `f::Function`: A function \$f:\\mathbb{R}^2\\rightarrow\\mathbb{R}^2\$.
- `x0::Real`: X coordinate of the point.
- `y0::Real`: Y coordinate of the point.
- `preiterations::Integer`: Number of first iterations to calculate but not to draw.
- `iterations::Integer`: Number of iterations to calculate (after `preiterations`).
- `kind::Symbol`: Kind of path: `:line` or `:arc`.
"""
function drawpointorbitpathR2(f::Function, x0::Real, y0::Real;
    preiterations::Integer=0, iterations::Integer=100, kind::Symbol=:line)

    drawsegment = SDDGraphics.drawlinesegment

    if kind == :arc
        SDDGraphics.supported(:drawarc)
        drawsegment = SDDGraphics.drawarc
    else
        SDDGraphics.supported(:drawlinesegment)
    end

    # Verifying functions
    @assert typeof(f(1., 1.)) <: Tuple{Real,Real}

    SDDGraphics.newdrawing()
    SDDGraphics.updatecolorarray(iterations)

    xn0,yn0 = x0,y0

    for n in 1:preiterations
        xn0,yn0 = f(xn0,yn0)
    end # for n preiterations

    for n in 1:iterations
        xn1,yn1 = f(xn0,yn0)
        SDDGraphics.color(n)
        drawsegment(xn0,yn0,xn1,yn1)
        xn0,yn0 = xn1,yn1
    end # for n iterations

    SDDGraphics.drawing()

end # function drawpointorbitpathR2

"""
    drawpointorbitpathC(f, z0 [; preiterations, iterations, kind])

Return the drawing of the orbit's path of a point \$z_0\$ under a function
\$f:\\mathbb{C}\\rightarrow\\mathbb{C}\$,
in a rectangular region in \$\\mathbb{C}\$.

The orbit of \$z_0\$ under \$f\$ is defined as
\$o(z_0,f)=\\{z_0,f(z_0),\\dots,f^n(z_0),\\dots\\}\$.

#### Arguments
- `f::Function`: A function \$f:\\mathbb{C}\\rightarrow\\mathbb{C}\$.
- `z0::Number`: The point.
- `preiterations::Integer`: Number of first iterations to calculate but not to draw.
- `iterations::Integer`: Number of iterations to calculate (after `preiterations`).
- `kind::Symbol`: Kind of path: `:line` or `:arc`.
"""
function drawpointorbitpathC(f::Function, z0::Number;
    preiterations::Integer=0, iterations::Integer=100, kind::Symbol=:line)

    drawsegment = SDDGraphics.drawlinesegment

    if kind == :arc
        SDDGraphics.supported(:drawarc)
        drawsegment = SDDGraphics.drawarc
    else
        SDDGraphics.supported(:drawlinesegment)
    end

    # Verifying functions
    @assert typeof(f(1.0im)) <: Number

    SDDGraphics.newdrawing()
    SDDGraphics.updatecolorarray(iterations)

    zn0 = z0

    for n in 1:preiterations
        zn0 = f(zn0)
    end # for n preiterations

    for n in 1:iterations
        zn1 = f(zn0)
        SDDGraphics.color(n)
        drawsegment(zn0,zn1)
        zn0 = zn1
    end # for n iterations

    SDDGraphics.drawing()

end # function drawpointorbitpathC

"""
    drawpointorbitpathR(f, x0, y0 [; preiterations, iterations])

Return the drawing of the orbit's arc-path of a point \$x_0\$ under a function
\$f:\\mathbb{R}\\rightarrow\\mathbb{R}\$.

The orbit of \$x_0\$ under \$f\$ is defined as
\$o(x_0,f)=\\{x_0,f(x_0,\\dots,f^n(x_0,\\dots\\}\$.

#### Arguments
- `f::Function`: A function \$f:\\mathbb{R}^2\\rightarrow\\mathbb{R}^2\$.
- `x0::Real`: X coordinate of the point.
- `y0::Real`: Y coordinate of the point.
- `preiterations::Integer`: Number of first iterations to calculate but not to draw.
- `iterations::Integer`: Number of iterations to calculate (after `preiterations`).
"""
function drawpointorbitpathR(f::Function, x0::Real;
    preiterations::Integer=0, iterations::Integer=100)

    SDDGraphics.supported(:drawarc)

    # Verifying functions
    @assert typeof(f(1.)) <: Real

    SDDGraphics.newdrawing()
    SDDGraphics.updatecolorarray(iterations)

    xn0 = x0

    for n in 1:preiterations
        xn0 = f(xn0)
    end # for n preiterations

    for n in 1:iterations
        xn1 = f(xn0)
        SDDGraphics.color(n)
        SDDGraphics.drawarc(xn0,0,xn1,0)
        xn0 = xn1
    end # for n iterations

    SDDGraphics.drawing()

end # function drawpointorbitpathR2


import Gadfly

function plot1D(gp::GP; clim::Tuple{Float64, Float64}=(minimum(gp.X), maximum(gp.X)), CI::Float64=1.96, res::Int=1000)

        sx = (clim[2]-clim[1])/(res-1)
        x=collect(clim[1]:sx:clim[2])
        mu, Sigma = predict(gp, x)
        conf = CI*sqrt(Sigma)
        u = mu + conf
        l = mu - conf

        Gadfly.plot(Gadfly.layer(x=x, y=mu, ymin=l, ymax=u, Gadfly.Geom.line, Gadfly.Geom.ribbon),
                Gadfly.layer(x=gp.X,y=gp.y,Gadfly.Geom.point))
end

function plot2D(gp::GP; clim::Tuple{Float64, Float64, Float64, Float64} = (minimum(gp.X[1,:]), maximum(gp.X[1,:]),
                                                                   minimum(gp.X[2,:]), maximum(gp.X[2,:])),
                res::Int=50)
        sx = (clim[2]-clim[1])/(res-1)
        sy = (clim[4]-clim[3])/(res-1)
        A = Array(Float64,2,res^2)
        for (i,a) in enumerate(clim[1]:sx:clim[2])
            for (j,b) in enumerate(clim[3]:sy:clim[4])
                A[:,(i-1)*res+j] = [a,b]
            end
        end

        mu = predict(gp,A)[1]
        z= reshape(mu,res,res)
        Gadfly.plot(z=z,x=collect(clim[1]:sx:clim[2]),y=collect(clim[3]:sy:clim[4]),Gadfly.Geom.contour)
end

function Gadfly.plot(gp::GP; kwargs...)
    d, n = size(gp.X)
    if d>2
        error("Can only plot one or two dimensional Gaussian processes")
    elseif d==1
        plot1D(gp; kwargs...)
    elseif d==2
        plot2D(gp; kwargs...)
    end
end


using DSP, Random
export AutoRegressive

"""
    AutoRegressive(n, method = LPCLevinson())

Autoregressive surrogates of order-`n`. The autoregressive coefficients `φ` are estimated
using `DSP.lpc(x, n, method)`, and thus see the documentation of DSP.jl for possible
`method`s.

While these surrogates are obviously suited to test the null hypothesis whether the data
are coming from a autoregressive process, the Fourier Transform-based surrogates are
probably a better option. The current method is more like an explicit way to
produce surrogates for the same hypothesis by fitting a model.
It can be used as convient way to estimate
autoregressive coefficients and automatically generate surrogates based on them.

The coefficients φ of the autoregressive fit can be found by doing
```julia
sg = surrogenerator(x, AutoRegressive(n))
φ = sg.init.φ
```
"""
struct AutoRegressive{M} <: Surrogate
    n::Int
    m::M
end
AutoRegressive(n::Int) = AutoRegressive(n, LPCLevinson())

function surrogenerator(x, method::AutoRegressive, rng = Random.default_rng())
    φ, e = lpc(x, method.n, method.m)
    N, f = length(x), length(φ)
    d = Normal(0, std(x))

    init = (
        d = d,
        φ = φ,
        N = N,
    )

    # The surrogate
    T = typeof(d).parameters[1]
    s = zeros(T, N + f)

    return SurrogateGenerator(method, x, s, init, rng)
end

function (sg::SurrogateGenerator{<:AutoRegressive})()
    N, φ, d = sg.init.N, sg.init.φ, sg.init.d
    return autoregressive!(sg.s, d, N, φ, sg.rng)
end

function autoregressive!(s, d::Normal{T}, N, φ, rng) where {T}
    f = length(φ)
    @inbounds for i in 1:f; s[i] = rand(rng, d); end
    @inbounds for i in f+1:length(s)
        s̃ = zero(T)
        for j in 1:f
            s̃ += φ[j] * s[i - j]
        end
        s̃ += rand(rng, d)
        s[i] = s̃
    end
    return view(s, f+1:N+f)
end


function createPointsOcta(n::Int64)
    points,triangulation  = computexyz(n)
	return points,triangulation
end

function linspace(start,stop,n)
	return collect(range(start,stop=stop,length=n))
end

function computexyz(n::Int64,slerpflag=true::Bool)
	pt0 = [0.0, 0.0, 1.0]
	pt1 = [0.0, 1.0, 0.0]
	pt2 = [1.0, 0.0, 0.0]

	if slerpflag
		# use slerp instead of linspace
		pts01 = slerp(pt0,pt1,n)
		pts02 = slerp(pt0,pt2,n)
	else
		pts01 = linspace(pt0,pt1,n) #type is: n-element LinSpace{Array{Float64,1}}:
		pts02 = linspace(pt0,pt2,n)
	end

	nptsoctant = Int64(n*(n+1)/2)
	pts = zeros(3,nptsoctant)

	# Creates points in planar geometry
	counter = 0;
	for i=1:n
		if slerpflag
			if i==1
				tmp = pts01[:,1]
			else
				tmp = slerp(pts01[:,i],pts02[:,i],i)
			end
		else
			tmp =linspace(pts01[i],pts02[i],i)
		end
		for j=1:i
			counter += 1
			if slerpflag
				pts[:,counter] = tmp[:,j]
			else
				pts[:,counter] = tmp[j]
			end
		end

	end

	# Project points onto sphere
	for i=1:nptsoctant
		pts[:,i] = pts[:,i]/norm(pts[:,i])
	end

	# We now want to get the connectivity. Therefore enumerate all
	#the points (overkill from a computational point of view) and then
	#write the correct connectivity.

    # Matrix that assigns an ID to the points
   ids=zeros(Int64,n,n); # Too large, but not important.
   #Matrix that will later contain all triangles
   nTrianglesOctant = Int64(n*n-2*n+1);
   triangles = zeros(Int64,3,nTrianglesOctant);

	counter = 0
	for i=1:n
		for j=1:i
			counter +=1
			ids[i,j] = counter;
		end
	end

	# Now create triangles
	counter = 0
	tmp = zeros(Int64,1)
	for i=1:n
		for j=1:i-1
			tmp  = [ids[i,j],ids[i,j+1],ids[i-1,j]]
			counter += 1
			triangles[:,counter] = tmp;
		end
		if i<n
			for j=1:i-1
				tmp = [ids[i,j],ids[i,j+1],ids[i+1,j+1]]
				counter +=1
				triangles[:,counter] = tmp;
			end
		end
	end

	# Now we have the quadrature points and triangles for a single octant
	ptsAll = deepcopy(pts);
	#mat ptsAll = pts; // the quadrature points for all octants, but with duplicated #entries

	tmp = deepcopy(pts); tmp[1,:] *= -1.0
	ptsAll = deepcopy(hcat(ptsAll,tmp));

	tmp = deepcopy(pts); tmp[2,:] *= -1.0;
	ptsAll = deepcopy(hcat(ptsAll,tmp));

	tmp = deepcopy(pts); tmp[1,:] *= -1.0; tmp[2,:] *= -1.0;
	ptsAll = deepcopy(hcat(ptsAll,tmp));

	tmp = deepcopy(ptsAll); tmp[3,:] *= -1.0;
	ptsAll = deepcopy(hcat(ptsAll,tmp));

	trianglesAll = deepcopy(triangles);
	for i=1:7
		trianglesAll = (hcat(trianglesAll,triangles .+i .*nptsoctant))
	end
	ptsAll,triangulation =unique(ptsAll,trianglesAll)

	ptsAll = permutedims(ptsAll)
	triangulation = permutedims(triangulation)
	return ptsAll, triangulation
end


@testset "Utils" begin

    A = [missing 1.0; 2.0 missing]
    ClimateTools.replace_missing!(A)
    @test isnan(A[1,1])
    @test isnan(A[2,2])

    x = 20.0:0.5:150

    ω = ClimateTools.weight(x, frac=0.25, power=1.0)
    @test maximum(ω) == 1.0
    @test minimum(ω) == 0.0
    @test (ω[34]+ω[33])/2 == 0.5
    @test findfirst(ω .== 1.0) == 66 # idx where ω = 1.0 and cutoff is set to 0.25

    ω = ClimateTools.weight(x, frac=0.25, power=2.0)
    @test maximum(ω) == 1.0
    @test minimum(ω) == 0.0
    @test round((ω[34]+ω[33])/2, digits=2) == 0.25
    @test findfirst(ω .== 1.0) == 66 # idx where ω = 1.0 and cutoff is set to 0.25

end


"""
    VectorCartesianSet{N, TS}

Represents the Cartesian product `S` of several vector sets.
`S = S_1 × S_2 × S_3 ... × S_N`, such that a vector `x ∈ S ⇔ x[i:i+dim(x)] ∈ S[i] ∀ i`.
The dimension of the set is the sum of the dimensions of the individual vector sets.
"""
struct VectorCartesianSet{N, TS <: NTuple{N, MOI.AbstractVectorSet}} <: MOI.AbstractVectorSet
    sets::TS
    dims::NTuple{N, Int}
    start_indices::NTuple{N, Int}
end

function VectorCartesianSet(sets...)
    dims = MOI.dimension.(sets)
    start_indices = tuple(1, dims[1:end-1]...)
    for i in 2:length(sets)
        start_indices = Base.setindex(start_indices,
            start_indices[i-1] + dims[i-1],
            i,
        )
    end
    return VectorCartesianSet(sets, dims, start_indices)
end

MOI.dimension(set::VectorCartesianSet) = sum(set.dims)

function MOI.dual_set(set::VectorCartesianSet)
    dual_sets = MOI.dual_set.(set.sets)
    return VectorCartesianSet(dual_sets...)
end

MOI.dual_set_type(set::VectorCartesianSet) = typeof(MOI.dual_set(set))


using CSV, DataFrames, CmdStan, StanMCMCChains, MCMCChains
gr(size=(500,500));

pdir = @__DIR__
cd(pdir)

howell1 = CSV.read("Howell1.csv", delim=';')
df = convert(DataFrame, howell1);

df2 = filter(row -> row[:age] >= 18, df)
first(df2, 5)

heightsmodel = "
// Inferring a Rate
data {
  int N;
  real<lower=0> h[N];
}
parameters {
  real<lower=0> sigma;
  real<lower=0,upper=250> mu;
}
model {
  // Priors for mu and sigma
  mu ~ normal(178, 20);
  sigma ~ uniform( 0 , 50 );

  // Observed heights
  h ~ normal(mu, sigma);
}
";

stanmodel = Stanmodel(name="heights", model=heightsmodel,
output_format=:mcmcchains);

heightsdata = Dict("N" => length(df2[:height]), "h" => df2[:height]);

rc, chn, cnames = stan(stanmodel, heightsdata, pdir, diagnostics=false,
  CmdStanDir=CMDSTAN_HOME);

describe(chn)
describe(chn, section=:internals)


#= ===================
   LOGISTIC CLASSIFIER
   =================== =#

@with_kw_noshow mutable struct LogisticClassifier <: MMI.Probabilistic
    lambda::Real             = 1.0
    gamma::Real              = 0.0
    penalty::SymStr          = :l2
    fit_intercept::Bool      = true
    penalize_intercept::Bool = false
    solver::Option{Solver}   = nothing
    multi_class::Bool        = false
end

glr(m::LogisticClassifier) =
    LogisticRegression(m.lambda, m.gamma;
                       penalty=Symbol(m.penalty),
                       multi_class=m.multi_class,
                       fit_intercept=m.fit_intercept,
                       penalize_intercept=m.penalize_intercept)

descr(::Type{LogisticClassifier}) = "Classifier corresponding to the loss function ``L(y, Xθ) + λ|θ|₂²/2 + γ|θ|₁`` where `L` is the logistic loss."

#= ======================
   MULTINOMIAL CLASSIFIER
   ====================== =#

@with_kw_noshow mutable struct MultinomialClassifier <: MMI.Probabilistic
    lambda::Real             = 1.0
    gamma::Real              = 0.0
    penalty::SymStr          = :l2
    fit_intercept::Bool      = true
    penalize_intercept::Bool = false
    solver::Option{Solver}   = nothing
end

glr(m::MultinomialClassifier) =
    MultinomialRegression(m.lambda, m.gamma;
                          penalty=Symbol(m.penalty),
                          fit_intercept=m.fit_intercept,
                          penalize_intercept=m.penalize_intercept)

descr(::Type{MultinomialClassifier}) =
    "Classifier corresponding to the loss function " *
    "``L(y, Xθ) + λ|θ|₂²/2 + γ|θ|₁`` where `L` is the multinomial loss."


module sl_model

using DifferentialEquations, Flux, Optim, DiffEqFlux, Plots

model_params=[0.1, 6.28, 0.1, 1.0];
all_ic=[[0.0,0.01], [π,0.2], [π/2.0, 0.1]];
tend=15.0
tstart=0.0
sampling=0.1

function sl_oscillator(du,u,p,t)
    γ, ω0, β, K=p
    ψ,R=u
    du[1]=ω0+K/2.0*sin(β)*(1+R^2)
    du[2]=-γ*R+K/2.0*cos(β)*R*(1-R^2)
end

function InitPlot()

    mycolor=[:blue, :red, :green, :purple, :black, :cyan, :orange]
    odedata=[solve(ODEProblem(sl_oscillator,η,(tstart,tend), model_params), Tsit5()) for η in all_ic];
    pl = plot(odedata[1],proj=:polar,vars=(1,2), color=:blue, legend=false)

    for k in 2:length(all_ic)
        plot!(pl, odedata[k], color=mycolor[k], vars=(1,2), proj=:polar)
    end
    title!(pl, "Oscillator Model")
    xlabel!(pl, "X")
    ylabel!(pl, "Y")
    savefig("SL_init.png")
    display(pl)

end

function predict_adjoint(param) # Our 1-layer neural network
    prob=ODEProblem(sl_oscillator,[0.1,0.1],(tstart,tend), model_params)
    Array(concrete_solve(prob,Tsit5(),param[1:2],param[3:end],saveat=tstart:sampling:tend))
end

# Generate some data to fit, and add some noise to it
data=predict_adjoint([0.0,0.08,0.1, 6.28, 0.1, 1.0])
σN=0.05
data+=σN*randn(size(data))
data=abs.(data)

function loss_adjoint(param)
  prediction = predict_adjoint(param)
  loss = sum(abs2,prediction - data)
  loss
end

function train_model(;pguess=[0.1,0.1,0.15, 6.0, 0.2, 0.8])
    println("The initial loss function is $(loss_adjoint(pguess)[1])")
    resinit=DiffEqFlux.sciml_train(loss_adjoint,pguess,ADAM(),maxiters=3000)
    res = DiffEqFlux.sciml_train(loss_adjoint,resinit.minimizer,BFGS(initial_stepnorm = 0.0001))
    println("The parameters are $(res.minimizer) with final loss value $(res.minimum)")
    return(res)
end

function plotFit(param)

    validationPlot(param, param[1:2])

end

function validationPlot(param, ic)

    lay = @layout [ a{0.6w} [b; c]]

    sol_fit=solve(ODEProblem(sl_oscillator,ic,(0.0,15.0),param[3:end]), Tsit5(), saveat=tstart:0.01:tend)
    sol_data=solve(ODEProblem(sl_oscillator,ic,(0.0,15.0),model_params), Tsit5(), saveat=tstart:0.01:tend)
    pl1=plot(sol_fit,lw=2, legend=false, vars=(1,2), color=:green, proj = :polar)
    plot!(pl1, sol_data, color=:blue, proj=:polar, vars=(1,2))

    #radius plot
    pl2=plot(sol_fit,lw=2, legend=false, vars=(0,2), color=:green)
    plot!(pl2, sol_data, color=:blue, vars=(0,2), lw=2)

    #angle plot
    pl3=plot(sol_fit.t,sin.(sol_fit[1,:]), lw=2, legend=false, color=:green)
    plot!(pl3, sol_data.t, sin.(sol_data[1,:]), color=:blue)

    pl=plot(pl1, pl2, pl3, layout=lay)
    savefig("Sl_valid_2.png")
    display(pl)

end

end #module

sl_model.InitPlot()
res=sl_model.train_model()
sl_model.plotFit(res.minimizer)
sl_model.validationPlot(res.minimizer, [π, 0.6])


function load_motor_data()
    df = CSV.read("data/motor.csv");
    x = df.times;
    y = df.accel;
    # Scaling the x-variables to lie in [0,1]
    L = maximum(x) - minimum(x);
    t = (x .- minimum(x))/L;
    return x, y, t
end

function alpha(p::Int)
    a = zeros(p, 1);
    for i = 0:p-1
        for j = 0:i
            for k = i:p-1
                a[i+1] = a[i+1] + (-1.0)^(k-j)/((k+j+1)*
                        factorial(j)*
                        factorial(p-1)*
                        factorial(p-1-k)*
                        factorial(i-j)*
                        factorial(k-i));
            end
        end
    end
    return a;
end

function spline_kernel(t::Array{Float64}, p::Int)

    #TO-DO
    # Check inputs
    # Maybe the Vector(range) can be done more eleganly.

    if all(diff(t) .> 0)
        monotonic = 1;
    elseif all(diff(t) .< 0)
        monotonic = 0;
    else
        throw(DomainError("t is not strictly monotonic"));
    end

    if size(t,2) != 1 && size(t,1) == 1;
        t = t';
    end

    fp = factorial.(p-1:-1:0)
    a = alpha(p).*fp;
    if monotonic == 1;
        U = (repeat(t,1,p).^Vector(p-1:-1:0)')./fp';
        V = (repeat(t,1,p).^Vector(p:2*p-1)').*a';
    elseif monotonic == 0
        U = (repeat(t,1,p).^Vector(p:2*p-1)')./fp';
        V = (repeat(t,1,p).^Vector(p-1:-1:0)').*a';
    end

    return U,V;
end

function spline_kernel_matrix(U::Array{Float64}, V::Array{Float64})
    return tril(U*V') + triu(V*U',1);
end

function logp(Σ, T, y, σf, σn)
    m = rank(T);
    n = length(y);
    Ki = cholesky(σf^2*Σ  + σn^2*I);
    if LinearAlgebra.issuccess(Ki) == false
        throw(DomainError("Cholesky Factorization failed"));
    end
    A = T'*(Ki\T);
    C = (Ki\T)*(A\(T'*inv(Ki)));
    lp = -0.5*y'*(Ki\y)  +
          0.5*y'*(C*y)   +
         -0.5*logdet(Ki) +
         -0.5*logdet(A)  +
         -(n-m)*0.5*log(2*π);
    return lp;
end

function dlogp(Σ, T, y, σf, σn)
    m = rank(T);
    n = length(y);
    Ki = cholesky(σf^2*Σ  + σn^2*I);
    if LinearAlgebra.issuccess(Ki) == false
        throw(DomainError("Cholesky Factorization failed"));
    end
    A = T'*(Ki\T);
    # TO_DO: Solving for the identity?
    dKf = 2*σf*Σ;
    dAf = -T'*(Ki\dKf)*(Ki\T);
    dKn = 2*σn*I;
    dAn = -T'*(inv(Ki)*dKn)*(Ki\T);
    Pf = Ki\dKf*inv(Ki);
    Pn = Ki\(dKn*inv(Ki));
    G = T*(A\(T'*(Ki\y)));
    dlpf =  0.5*y'*(Pf*y)                +
            0.5*(-2*(y'*Pf)*G + G'*Pf*G) +
           -0.5*tr(Ki\dKf)               +
           -0.5*tr(A\dAf);
   dlpn =   0.5*y'*(Pn*y)                   +
            0.5*(-2*(y'*Pn)*G + G'*Pn*G)    +
           -0.5*tr(inv(Ki)*dKn)             +
           -0.5*tr(A\dAn);
    return [dlpf; dlpn];
end

## Auxillary functions
function create_T(t, m)
    T = ones(length(t), m);
    for ν = 2:m
        T[:, ν] = t.^(ν-1)/factorial(ν - 1);
    end
    return T;
end

## Leave-one-out (LOO), what to do with explicit basis functions?


using LatinSquares

if length(ARGS) == 0
	println("Usage: julia run_latin n")
	exit()
end
n = parse(ARGS[1])
println("n = $n")
println("Starting up")

# Here's a place to adjust Gurobi settings
# env = Gurobi.Env()
# setparam!(env,"Threads", 8)
# setparam!(env,"ConcurrentMIP", 2)

tic();
try
	A,B = ortho_latin(n)
	println("A = $A")
	println("B = $B")
end
toc()


using LinearAlgebra
using SparseArrays

function max_out_rate(chain)
    max = 0
    for src in nodes(chain.state_graph)
        total_rate = 0.0
        for edge in out_edges(chain.state_graph, src)
            rate = chain.trans_rates[edge]
            total_rate += rate
        end
        if total_rate > max
            max = total_rate
        end
    end
    return max
end

function trans_prob_matrix(chain, unif_rate)
    rows = Vector{Int}()
    cols = Vector{Int}()
    vals = Vector{Float64}()
    for src in nodes(chain.state_graph)
        total_rate = 0.0
        for edge in out_edges(chain.state_graph, src)
            dst = dst_node(chain.state_graph, edge)
            rate = chain.trans_rates[edge]
            total_rate += rate
            push!(rows, dst)
            push!(cols, src)
            push!(vals, rate / unif_rate)
        end
        push!(rows, src)
        push!(cols, src)
        self_prob = 1.0 - total_rate / unif_rate
        @assert self_prob >= 0
        push!(vals, self_prob)
    end
    dim = node_count(chain.state_graph)
    sparse(rows, cols, vals, dim, dim)
end

function trans_rate_matrix(chain)
    rows = Vector{Int}()
    cols = Vector{Int}()
    vals = Vector{Float64}()
    for src in nodes(chain.state_graph)
        total_rate = 0.0
        for edge in out_edges(chain.state_graph, src)
            dst = dst_node(chain.state_graph, edge)
            rate = chain.trans_rates[edge]
            total_rate += rate
            push!(rows, dst)
            push!(cols, src)
            push!(vals, rate)
        end
        push!(rows, src)
        push!(cols, src)
        push!(vals, -total_rate)
    end
    dim = node_count(chain.state_graph)
    sparse(rows, cols, vals, dim, dim)
end

export fintime_solve_prob, fintime_solve_cum, trans_rate_matrix, trans_prob_matrix, state_prob, state_cumtime

include("poisson_trunc.jl")

struct FintimeProbSolution
    prob::Vector{Float64}
end

struct FintimeCumSolution
    cumtime::Vector{Float64}
end

struct FintimeSolution
    prob::Vector{Float64}
    cumtime::Vector{Float64}
end

function state_prob(sol::Union{FintimeProbSolution, FintimeSolution}, state::Integer)
    return sol.prob[state]
end

function state_cumtime(sol::Union{FintimeCumSolution, FintimeSolution}, state::Integer)
    return sol.cumtime[state]
end

function fintime_solve(chain, init_prob, time)
    prob = fintime_solve_prob(chain, init_prob, time)
    cum = fintime_solve_cum(chain, init_prob, time)
    FintimeSolution(prob.prob, cum.cumtime)
end

function fintime_solve_prob(chain::ContMarkovChain, init_prob, time::Real; unif_rate_factor=1.05, tol=1e-6, ss_check_interval=10)
    @assert unif_rate_factor >= 1.0
    unif_rate = max_out_rate(chain) * unif_rate_factor
    P = trans_prob_matrix(chain, unif_rate)
    prob = fill(0.0, state_count(chain))
    for i in eachindex(init_prob)
        prob[i] = init_prob[i]
    end

    ltp, rtp, poi_probs = poisson_trunc_point(unif_rate * time, tol)

    checkpoint = copy(prob)
    prob_old = copy(prob)

    for k in 0:ltp-1
        prob, prob_old = prob_old, prob
        mul!(prob, P, prob_old)
        if (k + 1) % ss_check_interval == 0
            checkpoint .-= prob
            diff = maximum(abs, checkpoint)
            checkpoint .= prob
            if  diff < tol
                return FintimeProbSolution(prob)
            end
        end
    end
    term_sum = 0.0
    prob_t = fill(0.0, length(prob))
    for k in ltp:rtp
        term = poi_probs[k - ltp + 1]
        term_sum += term
        @. prob_t += term * prob

        prob, prob_old = prob_old, prob
        mul!(prob, P, prob_old)
        if (k - ltp + 1) % ss_check_interval == 0
            @. checkpoint -= prob
            diff = maximum(abs, checkpoint)
            @. checkpoint = prob
            if diff < tol
                @. prob_t += prob * (1.0 - term_sum)
                return FintimeProbSolution(prob_t)
            end
        end
    end

    return FintimeProbSolution(prob_t)
end

function fintime_solve_cum(chain::ContMarkovChain, init_prob, time::Real; unif_rate_factor=1.05, tol=1e-6, ss_check_interval=10)
    unif_rate = max_out_rate(chain) * unif_rate_factor
    P = trans_prob_matrix(chain, unif_rate)
    trans_rate_matrix(chain)
    Matrix{Float64}(P)

    prob = fill(0.0, state_count(chain))
    for i in eachindex(init_prob)
        prob[i] = init_prob[i]
    end
    prob_old = copy(prob)
    checkpoint = copy(prob)

    qt = unif_rate * time
    rtp, poi_probs = poisson_cum_rtp(qt, time, tol)
    right_cum = 1.0 - poi_probs[1]
    sum_right_cum = right_cum

    sol = fill(0.0, length(prob))
    @. sol += right_cum * prob
    ss_reached = false
    for i in 1:rtp
        prob, prob_old = prob_old, prob
        mul!(prob, P, prob_old)

        right_cum -= poi_probs[i + 1]
        sum_right_cum += right_cum

        @. sol += right_cum * prob
        if i % ss_check_interval == 0
            checkpoint .-= prob
            diff = maximum(abs, checkpoint)
            checkpoint .= prob
            if diff < tol
                ss_reached = true
                break;
            end
        end
    end
    sol .*= 1.0 / unif_rate
    if ss_reached
        @. sol += (time - sum_right_cum / unif_rate) * prob
    end
    return FintimeCumSolution(sol)
end


# By default, Julia/LLVM does not use fused multiply-add operations (FMAs).
# Since these FMAs can increase the performance of many numerical algorithms,
# we need to opt-in explicitly.
# See https://ranocha.de/blog/Optimizing_EC_Trixi for further details.
@muladd begin

#     CurvedSurface{RealT<:Real}
#
# Contains the data needed to represent a curve with data points (x,y) as a Lagrange polynomial
# interpolant written in barycentric form at a given set of nodes.
struct CurvedSurface{RealT<:Real}
  nodes               ::Vector{RealT}
  barycentric_weights ::Vector{RealT}
  coordinates         ::Array{RealT, 2} #[nnodes, ndims]
end

# evalute the Gamma curve interpolant at a particular point s and return the (x,y) coordinate
function evaluate_at(s, boundary_curve::CurvedSurface)

   @unpack nodes, barycentric_weights, coordinates = boundary_curve

   x_coordinate_at_s_on_boundary_curve = lagrange_interpolation(s, nodes, view(coordinates, :, 1),
                                                                barycentric_weights)
   y_coordinate_at_s_on_boundary_curve = lagrange_interpolation(s, nodes, view(coordinates, :, 2),
                                                                barycentric_weights)

   return x_coordinate_at_s_on_boundary_curve, y_coordinate_at_s_on_boundary_curve
end

# evalute the derivative of a Gamma curve interpolant at a particular point s
# and return the (x,y) coordinate
function derivative_at(s, boundary_curve::CurvedSurface)

   @unpack nodes, barycentric_weights, coordinates = boundary_curve

   x_coordinate_at_s_on_boundary_curve_prime = lagrange_interpolation_derivative(s, nodes,
                                                                                view(coordinates, :, 1),
                                                                                barycentric_weights)
   y_coordinate_at_s_on_boundary_curve_prime = lagrange_interpolation_derivative(s, nodes,
                                                                                view(coordinates, :, 2),
                                                                                barycentric_weights)
   return x_coordinate_at_s_on_boundary_curve_prime, y_coordinate_at_s_on_boundary_curve_prime
end

# Chebysehv-Gauss-Lobatto nodes and weights for use with curved boundaries
function chebyshev_gauss_lobatto_nodes_weights(n_nodes::Integer)

  # Initialize output
  nodes   = zeros(n_nodes)
  weights = zeros(n_nodes)

  # Get polynomial degree for convenience
  N = n_nodes - 1

  for j in 1:n_nodes
    nodes[j]   = -cospi( (j-1) / N )
    weights[j] = pi / N
  end
  weights[1]   = 0.5 * weights[1]
  weights[end] = 0.5 * weights[end]

  return nodes, weights
end

# Calculate Lagrange interpolating polynomial of a function f(x) at a given point x for a given
# node distribution.
function lagrange_interpolation(x, nodes, fvals, wbary)
# Barycentric two formulation of Lagrange interpolant
  numerator   = zero(eltype(fvals))
  denominator = zero(eltype(fvals))

  for j in eachindex(nodes)
    if isapprox(x, nodes[j], rtol=eps(x))
      return fvals[j]
    end
    t            = wbary[j] / ( x - nodes[j] )
    numerator   += t * fvals[j]
    denominator += t
  end

  return numerator/denominator
end

# Calculate derivative of a Lagrange interpolating polynomial of a function f(x) at a given
# point x for a given node distribution.
function lagrange_interpolation_derivative(x, nodes, fvals, wbary)

  at_node   = false
  numerator = zero(eltype(fvals))
  i         = 0

  for j in eachindex(nodes)
    if isapprox(x, nodes[j])
      at_node     = true
      p           = fvals[j]
      denominator = -wbary[j]
      i           = j
    end
  end

  if at_node
    for j in eachindex(nodes)
      if j != i
        numerator += wbary[j] * ( p - fvals[j] ) / ( x - nodes[j] )
      end
    end
  else
    denominator = zero(eltype(fvals))
    p = lagrange_interpolation(x, nodes, fvals, wbary)
    for j in eachindex(nodes)
      t            = wbary[j] / (x - nodes[j])
      numerator   += t * ( p - fvals[j] ) / ( x - nodes[j] )
      denominator += t
    end
  end

  return numerator/denominator # p_prime
end

end # @muladd


@testset "Triangulation: " begin
    P1 = SMatrix{3, 4, Float64}(
        0.999701, -0.0171452, 0.0174497,
        0.0174497, 0.999695, -0.0174524,
        -0.017145, 0.0177517, 0.999695,
        -500, -100, -100,
    )
    P2 = SMatrix{3, 4, Float64}(
        0.99969, 0.0177543, -0.0174497,
        -0.0174497, 0.999695, 0.0174524,
        0.0177543, -0.0171425, 0.999695,
        500, -100, -100,
    )
    K = SMatrix{3, 3, Float64}(
        7291.67, 0, 0,
        0, 7291.67, 0,
        639.5, 511.5, 1,
    )

    P1 = K * P1
    P2 = K * P2

    p1 = SVector{2, Float64}(146, 642.288)
    p2 = SVector{2, Float64}(1137.31, 385.201)

    t = triangulate(p1, p2, P1, P2)
    t /= t[4]
    @test isapprox(t, SVector{4, Float64}(0, 100, 10000, 1); atol=1e-1)
end

@testset "Triangulation: Horizontal setup" begin
    P1 = SMatrix{3, 4, Float64}(I)
    P2 = hcat(SMatrix{3, 3, Float64}(I), SVector{3, Float64}(-1000, 0, 0))
    K = SMatrix{3, 3, Float64}(
        7291.67, 0, 0,
        0, 7291.67, 0,
        639.5, 511.5, 1,
    )

    P1 = K * P1
    P2 = K * P2

    p1 = SVector{2, Float64}(1004.08, 511.5)
    p2 = SVector{2, Float64}(274.917, 511.5)

    t = triangulate(p1, p2, P1, P2)
    t /= t[4]
    @test isapprox(t, SVector{4, Float64}(500, 0, 10000, 1); atol=1e-1)
end


@doc """
### SimpleODEMod.calc_x2_4t3

  Calculates and returns x^2 + 4*t^3
"""->
function calc_x2_4t3(coords::AbstractArray{Tmsh}, params::ParamType2, t) where Tmsh

  x = coords[1]
  y = coords[2]

  u = x^2 + 4*(t^3)

  return u
end

@doc """
### SimpleODEMod.calc_x2_3t2

  Calculates and returns x^2 + 3*t^2
"""->
function calc_x2_3t2(coords::AbstractArray{Tmsh}, params::ParamType2, t) where Tmsh

  x = coords[1]
  y = coords[2]

  u = x^2 + 3*(t^2)

  return u
end

@doc """
### SimpleODEMod.calc_x2_t2

  Calculates and returns x^2 + t^2
"""->
function calc_x2_t2(coords::AbstractArray{Tmsh}, params::ParamType2, t) where Tmsh

  x = coords[1]
  y = coords[2]

  u = x^2 + t^2

  return u
end

@doc """
### SimpleODEMod.calc_4t3

  Calculates and returns 4t^3
  (du/dt of calc_x2_t4)
"""->
function calc_4t3(coords::AbstractArray{Tmsh}, params::ParamType2, t) where Tmsh

  x = coords[1]
  y = coords[2]

  u = 4*(t^3)

  return u
end

@doc """
### SimpleODEMod.calc_3t2

  Calculates and returns 3t^2
  (du/dt of calc_x2_t3)
"""->
function calc_3t2(coords::AbstractArray{Tmsh}, params::ParamType2, t) where Tmsh

  x = coords[1]
  y = coords[2]

  u = 3*(t^2)

  return u
end

@doc """
### SimpleODEMod.calc_2t

  Calculates and returns 2t
  (du/dt of calc_x2_t2)
"""->
function calc_2t(coords::AbstractArray{Tmsh}, params::ParamType2, t) where Tmsh

  x = coords[1]
  y = coords[2]

  u = 2*t

  return u
end


using DelimitedFiles
using PyPlot
using RAFF

datafile = "/tmp/output.txt"

fp = open(datafile, "r")

N = parse(Int, readline(fp))

M = readdlm(fp)

close(fp)

x = M[:, 1]
y = M[:, 2]
c = M[:, 3]

sol = [779.616, 5315.36, 0.174958, 2.52718]

n, model, modelstr = RAFF.model_list["logistic"]

modl1 = (x) -> model(x, sol)
modl2 = (x) -> model(x, θSol)

t = minimum(x):0.01:maximum(x)
PyPlot.plot(t, modl1.(t), "r", label="RAFF")

θSol = [1000.0, 5000.0, 0.2, 3.0]
PyPlot.plot(t, modl2.(t), "b--", label="True")

θLS = [-959.07, 8151.03, 0.0927191, 0.940711]
modl3 = (x) -> model(x, θLS)

PyPlot.plot(t, modl3.(t), "g-.", label="Least Squares")

PyPlot.legend(loc=4)

PyPlot.scatter(x[c .== 0.0], y[c .== 0.0], color=PyPlot.cm."Paired"(0.1), marker="o", s=50.0, linewidths=0.2,
               alpha=1.0)

PyPlot.scatter(x[c .== 1.0], y[c .== 1.0], color=PyPlot.cm."Paired"(0.5), marker="^", s=50.0, linewidths=0.2,
               alpha=1.0)

PyPlot.show()

PyPlot.savefig("/tmp/figure.png", DPI=100)


"""
    _rategradient(∂X, ∂Y)

Returns the rate of change in units of the values stored in the grid, and the
angle of the change in wind direction, *i.e.* an angle of 180 means that the
value is increasing *from* the south. When both ∂X and ∂Y are equal to 0, the
angle is assumed to be 0.
"""
function _rategradient(∂X::T, ∂Y::T) where {T <: Number}
    if ∂X == ∂Y == 0.0
        return (0.0, 0.0)
    end
    m = sqrt(∂X^2 + ∂Y^2)
    Δ = ∂X >= 0.0 ? 0.0 : 180.0
    θ = rad2deg(atan(∂X, ∂Y)) + Δ
    θ = ∂X > 0.0 ? θ + 180.0 : θ
    return (m, θ)
end

"""
    _rateofchange(x::Vector{T}, y::Vector{T}, z::Vector{T})

Rate of change for a series of three points, defined as a series of `x` and `y`
coordinates and a value `z`. Returns a rate of change (in units of `z`) and a
gradient in degrees.
"""
function _rateofchange(x::Vector{T}, y::Vector{T}, z::Vector{T}) where {T <: Number}

    # Check that all three vectors have the same length
    length(x) == length(y) || throw(DimensionMismatch("x and y must have the same length"))
    length(x) == length(z) || throw(DimensionMismatch("x and z must have the same length"))

    # Get the matrix of coefficients
    C = cat(y, x, fill(one(T), length(x)); dims=(2, 2))
    coeff = Base.inv(C) * z

    X = sum(C[:, 1]) / 3.0
    Y = sum(C[:, 2]) / 3.0

    ∂X = coeff[2] * Y + coeff[3]
    ∂Y = coeff[1] * X + coeff[3]

    # Rate of change and direction
    return _rategradient(∂X, ∂Y)
end

"""
    _rateofchange(A::Matrix{T}) where {T <: Number}

Rate of change for a series of a lattice points. Returns the rate of change and
the gradient for a 2x2 grid of numbers.
"""
function _rateofchange(x::Vector{T}, y::Vector{T}, z::Matrix{T};) where {T <: Number}
    size(z) == (2, 2) || throw(DimensionMismatch("the matrix z must have size (2,2)"))

    # The values in x and y should be sorted, but we'll still enforce it
    Δx = abs(first(diff(x)))
    Δy = abs(first(diff(y)))

    # We can get the values directly from the matrix
    Z₄, Z₁, Z₃, Z₂ = z

    ∂X = Z₂ - Z₁ + Δy * (Z₁ - Z₂ + Z₃ - Z₄)
    ∂Y = Z₄ - Z₁ + Δx * (Z₁ - Z₂ + Z₃ - Z₄)

    # Rate of change and direction
    return _rategradient(∂X, ∂Y)
end

_rateofchange(z::Matrix{T}) where {T <: Number} = _rateofchange(one(T) / 2, one(T) / 2, z)


import ChainRulesCore
import DifferentiableBackwardEuler
import DPFEHM
import GaussianRandomFields
import Optim
import PyPlot
import Random
import Zygote

Random.seed!(0)

#set up the grid
mins = [0, 0, 0]; maxs = [50, 50, 5]#size of the domain, in meters
ns = [100, 100, 10]#number of nodes on the grid
coords, neighbors, areasoverlengths, volumes = DPFEHM.regulargrid3d(mins, maxs, ns)#build the grid

#set up the boundary conditions
Qs = zeros(size(coords, 2))
injectionnode = 1#inject in the lower left corner
Qs[injectionnode] = 1e-4#m^3/s
dirichletnodes = Int[size(coords, 2)]#fix the pressure in the upper right corner
dirichleths = zeros(size(coords, 2))
dirichleths[size(coords, 2)] = 0.0

#set up the initial condition
h0 = zeros(size(coords, 2))

#set up the conductivity field
lambda = 50.0#meters -- correlation length of log-conductivity
sigma = 1.0#standard deviation of log-conductivity
mu = -9.0#mean of log conductivity -- ~1e-4 m/s, like clean sand here https://en.wikipedia.org/wiki/Hydraulic_conductivity#/media/File:Groundwater_Freeze_and_Cherry_1979_Table_2-2.png
cov = GaussianRandomFields.CovarianceFunction(2, GaussianRandomFields.Matern(lambda, 1; σ=sigma))
x_pts = range(mins[1], maxs[1]; length=ns[1])
y_pts = range(mins[2], maxs[2]; length=ns[2])
num_eigenvectors = 200
grf = GaussianRandomFields.GaussianRandomField(cov, GaussianRandomFields.KarhunenLoeve(num_eigenvectors), x_pts, y_pts)
logKs = zeros(reverse(ns)...)
logKs2d = mu .+ GaussianRandomFields.sample(grf)'#generate a random realization of the log-conductivity field
for i = 1:ns[3]#copy the 2d field to each of the 3d layers
	v = view(logKs, i, :, :)
	v .= logKs2d
end

#set up the storage
specificstorage = fill(0.1, size(coords, 2))

#plot the log-conductivity
fig, ax = PyPlot.subplots()
img = ax.imshow(logKs[1, :, :], origin="lower")
ax.title.set_text("Conductivity Field")
fig.colorbar(img)
display(fig)
println()
PyPlot.close(fig)

logKs2Ks_neighbors(Ks) = exp.(0.5 * (Ks[map(p->p[1], neighbors)] .+ Ks[map(p->p[2], neighbors)]))
function solveforh(logKs, dirichleths)
	@assert length(logKs) == length(Qs)
	Ks_neighbors = logKs2Ks_neighbors(logKs)
	p = [Ks_neighbors; dirichleths]
	h_gw = DifferentiableBackwardEuler.steps(h0, f_gw, f_gw_u, f_gw_p, f_gw_t, p, 0.0, 1.0 * 60 * 60 * 24 * 365 * 10; abstol=1e-1, reltol=1e-1)
	return h_gw
end
hflat2h3d(h) = reshape(h, reverse(ns)...)

function unpack(p)
	@assert length(p) == length(neighbors) + size(coords, 2)
	Ks = p[1:length(neighbors)]
	dirichleths = p[length(neighbors) + 1:length(neighbors) + size(coords, 2)]
	return Ks, dirichleths
end

#set up some functions needed by diffeq
function f_gw(u, p, t)
	Ks, dirichleths = unpack(p)
	return DPFEHM.groundwater_residuals(u, Ks, neighbors, areasoverlengths, dirichletnodes, dirichleths, Qs, specificstorage, volumes)
end
function f_gw_u(u, p, t)
	Ks, dirichleths = unpack(p)
	return DPFEHM.groundwater_h(u, Ks, neighbors, areasoverlengths, dirichletnodes, dirichleths, Qs, specificstorage, volumes)
end
function f_gw_p(u, p, t)
	Ks, dirichleths = unpack(p)
	J1 = DPFEHM.groundwater_Ks(u, Ks, neighbors, areasoverlengths, dirichletnodes, dirichleths, Qs, specificstorage, volumes)
	J2 = DPFEHM.groundwater_dirichleths(u, Ks, neighbors, areasoverlengths, dirichletnodes, dirichleths, Qs, specificstorage, volumes)
	return hcat(J1, J2)
end
f_gw_t(u, p, t) = zeros(length(u))

#now do a forward solve for the head
print("forward solve time")
@time h = solveforh(logKs, dirichleths)#solve for the head
#plot the head at the bottom of the domain
fig, ax = PyPlot.subplots()
img = ax.imshow(hflat2h3d(h[:, end])[1, :, :], origin="lower")
ax.title.set_text("Head")
fig.colorbar(img)
display(fig)
println()
PyPlot.close(fig)

#now compute the gradient of a function involving solveforh
gradient_node = div(size(coords, 2), 2) + 500
gradient_node_x = coords[1, gradient_node]
gradient_node_y = coords[2, gradient_node]
print("forward and gradient time")
@time grad = Zygote.gradient((x, y)->hflat2h3d(solveforh(x, y)[:, end])[gradient_node], logKs, dirichleths)#calculate the gradient (which involves a redundant calculation of the forward pass)
function_evaluation, back = Zygote.pullback((x, y)->hflat2h3d(solveforh(x, y)[:, end])[gradient_node], logKs, dirichleths)#this pullback thing lets us not redo the forward pass
print("gradient time")
@time grad = back(1.0)#compute the gradient of a function involving solveforh
#plot the gradient of the function w.r.t. the logK at the bottom of the domain
fig, ax = PyPlot.subplots()
img = ax.imshow(grad[1][1, :, :], origin="lower", extent=[mins[1], maxs[1], mins[2], maxs[2]])
ax.plot([gradient_node_x], [gradient_node_y], "r.", ms=10, alpha=0.5)
ax.title.set_text("Gradient of head at dot w.r.t. logK at bottom of domain")
fig.colorbar(img)
display(fig)
println()
PyPlot.close(fig)


using Cubature
using Distributions
using LaTeXStrings
using LinearAlgebra
using Plots

"""
Implmentation of the Monte Carlo Integrator with Adaptive Importance Sampling
from Ryu, E. and Boyd, S., Adaptive Importance Sampling via Stochastic Convex Programming,
arXiv preprint, 2015.

Assumes a functional form f : Rᵈ -> R.
Parameters:
- d: input dimension
- n_stages: number of iterations performed by the integrator
- f: pdf whose expectation should be integrated
- C: update step
- μ_glob: mean of f
"""
function ais_rb(
    d::Int,
    n_stages::Int,
    f::Function,
    C::AbstractFloat,
    μ_glob::AbstractVector,
)
    g_λ = (μ, Σ) -> MultivariateNormal(μ, Σ)
    μ = rand(d)
    Σ = I + zeros(d, d)
    S = inv(Σ)
    m = S * μ
    X_k, est = (zeros(d), zeros(d))
    error_list = Float64[]
    @inbounds for k in 1:n_stages
        mul!(μ, Σ, m)
        X_k .= rand(g_λ(μ, Σ))
        est .= (est .* (k - 1) .+ f(X_k) .* X_k / pdf(g_λ(μ, Σ), X_k)) / k
        m .-=
            C .* (norm(X_k)^2) .* (f(X_k)^2) ./ (2 * pdf(g_λ(μ, Σ), X_k) * sqrt(k)) *
            (μ - X_k)
        push!(error_list, norm(est .- μ_glob) / norm(μ_glob))
    end
    return est, error_list
end

"""
Implmentation of the Monte Carlo Integrator with Adaptive Importance Sampling
from  Oh, M-S and Berger, J., Adaptive importance sampling in Monte Carlo integration, Journal
of Statistical Computation and Simulation, 1992.

Assumes a functional form f : Rᵈ -> R.
Parameters:
- d: input dimension
- n_stages: number of iterations performed by the integrator
- f: pdf whose expectation should be integrated
- μ_glob: mean of f
- n_samples: list of number of samples used for each iteration
"""
function ais_ob(
    d::Int,
    n_stages::Int,
    f::Function,
    μ_glob::AbstractVector,
    n_samples::AbstractVector,
)
    g_λ = (λ, Σ) -> MultivariateNormal(λ, Σ)
    λ = rand(d)
    Σ = I + zeros(d, d)
    WΦ = zeros(n_stages, d + 1)
    error_list = Float64[]
    @inbounds for k in 1:n_stages
        w_k = θ -> f(θ) / pdf(g_λ(λ, Σ), θ)
        samples_k_list = [rand(g_λ(λ, Σ)) for n = 1:n_samples[k]]
        W_k = h -> sum((h.(samples_k_list)) .* (w_k.(samples_k_list)))
        WΦ[k, 1] = W_k(x -> one(eltype(x)))
        WΦ[k, 2:end] .= [W_k(x -> x[k-1]) for k = 2:d+1]
        sum_weights = sum(@view WΦ[:, 1])
        λ = [sum(@view WΦ[:, k]) ./ sum_weights for k = 2:d+1]
        push!(error_list, norm(λ - μ_d) / norm(μ_d))
    end
    return λ, error_list
end

function init_test(d::Int)
    μ_d = rand(d)
    A = rand(d, d)
    ## Trick to make Σ_d positive semidefinite
    Σ_d = A' * A
    f = x -> pdf(MvTDist(2, μ_d, Σ_d), x)
    return μ_d, Σ_d, f
end

function run_ais_rb_normal_test(d::Int, n_iters_test::Int, n_stages_rb::Int)
    error_list_rb_final = zeros(n_stages_rb)
    for i = 1:n_iters_test
        @time est, error_list = ais_rb(d, n_stages_rb, f, 0.01, μ_d)
        error_list_rb_final = error_list_rb_final .* (i - 1) / i + error_list ./ i
    end
    return error_list_rb_final
end

function run_ais_ob_normal_test(d::Int, n_iters_test::Int, n_stages::Int, n_samples_ob::AbstractVector)
    error_list_ob_final = zeros(n_stages_ob)
    for i = 1:n_iters_test
        @time est, error_list = ais_ob(d, n_stages_ob, f, μ_d, n_samples_ob)
        error_list_ob_final = error_list_ob_final .* (i - 1) / i + error_list ./ i
    end
    return error_list_ob_final
end

# Initialize integrand and test parameters
d = 2
μ_d, _, f = init_test(d)
n_iters_test = 10

# RB tests
n_stages_rb = 1000000
error_list_rb_final = run_ais_rb_normal_test(d, n_iters_test, n_stages_rb)

plot_gap = 10
plot(
    log10.([k for k = Int(n_stages_rb / 2):plot_gap:n_stages_rb]),
    log10.(error_list_rb_final[Int(n_stages_rb / 2):plot_gap:end]),
    xlabel = L"$\log(N_s)$",
    ylabel = L"$\log(\varepsilon_{rel})$",
    label = "Stochastic Convex Programming",
)

# OB tests
n_stages_ob = 1000
n_samples_ob = [Int(n_stages_ob / 2) for k = 1:n_stages_ob]
error_list_ob_final = run_ais_ob_normal_test(d, n_iters_test, n_stages_ob, n_samples_ob)

plot_gap = 10
plot(
    log10.([k for k = Int(n_stages_ob / 2):plot_gap:n_stages_ob]),
    log10.(error_list_ob_final[Int(n_stages_ob / 2):plot_gap:end]),
    xlabel = L"$\log(N_s)$",
    ylabel = L"$\log(\varepsilon_{rel})$",
    label = "Cumulative Average Parameter Update",
)


"""
    Laplace(μ,β)

The *Laplace distribution* with location `μ` and scale `β` has probability density function

```math
f(x; \\mu, \\beta) = \\frac{1}{2 \\beta} \\exp \\left(- \\frac{|x - \\mu|}{\\beta} \\right)
```

```julia
Laplace()       # Laplace distribution with zero location and unit scale, i.e. Laplace(0, 1)
Laplace(μ)      # Laplace distribution with location μ and unit scale, i.e. Laplace(μ, 1)
Laplace(μ, β)   # Laplace distribution with location μ and scale β

params(d)       # Get the parameters, i.e., (μ, β)
location(d)     # Get the location parameter, i.e. μ
scale(d)        # Get the scale parameter, i.e. β
```

External links

* [Laplace distribution on Wikipedia](http://en.wikipedia.org/wiki/Laplace_distribution)

"""
struct Laplace{T<:Real} <: ContinuousUnivariateDistribution
    μ::T
    θ::T
    Laplace{T}(µ::T, θ::T) where {T} = new{T}(µ, θ)
end

function Laplace(μ::T, θ::T; check_args=true) where {T <: Real}
    check_args && @check_args(Laplace, θ > zero(θ))
    return Laplace{T}(μ, θ)
end

Laplace(μ::Real, θ::Real) = Laplace(promote(μ, θ)...)
Laplace(μ::Integer, θ::Integer) = Laplace(float(μ), float(θ))
Laplace(μ::T) where {T <: Real} = Laplace(μ, one(T))
Laplace() = Laplace(0.0, 1.0, check_args=false)

const Biexponential = Laplace

@distr_support Laplace -Inf Inf

#### Conversions
function convert(::Type{Laplace{T}}, μ::S, θ::S) where {T <: Real, S <: Real}
    Laplace(T(μ), T(θ))
end
function convert(::Type{Laplace{T}}, d::Laplace{S}) where {T <: Real, S <: Real}
    Laplace(T(d.μ), T(d.θ), check_args=false)
end

#### Parameters

location(d::Laplace) = d.μ
scale(d::Laplace) = d.θ
params(d::Laplace) = (d.μ, d.θ)
@inline partype(d::Laplace{T}) where {T<:Real} = T

#### Statistics

mean(d::Laplace) = d.μ
median(d::Laplace) = d.μ
mode(d::Laplace) = d.μ

var(d::Laplace) = 2d.θ^2
std(d::Laplace) = sqrt2 * d.θ
skewness(d::Laplace{T}) where {T<:Real} = zero(T)
kurtosis(d::Laplace{T}) where {T<:Real} = 3one(T)

entropy(d::Laplace) = log(2d.θ) + 1

#### Evaluations

zval(d::Laplace, x::Real) = (x - d.μ) / d.θ
xval(d::Laplace, z::Real) = d.μ + z * d.θ

pdf(d::Laplace, x::Real) = exp(-abs(zval(d, x))) / 2scale(d)
logpdf(d::Laplace, x::Real) = - (abs(zval(d, x)) + log(2scale(d)))

cdf(d::Laplace, x::Real) = (z = zval(d, x); z < 0 ? exp(z)/2 : 1 - exp(-z)/2)
ccdf(d::Laplace, x::Real) = (z = zval(d, x); z > 0 ? exp(-z)/2 : 1 - exp(z)/2)
logcdf(d::Laplace, x::Real) = (z = zval(d, x); z < 0 ? loghalf + z : loghalf + log2mexp(-z))
logccdf(d::Laplace, x::Real) = (z = zval(d, x); z > 0 ? loghalf - z : loghalf + log2mexp(z))

quantile(d::Laplace, p::Real) = p < 1/2 ? xval(d, log(2p)) : xval(d, -log(2(1 - p)))
cquantile(d::Laplace, p::Real) = p > 1/2 ? xval(d, log(2(1 - p))) : xval(d, -log(2p))
invlogcdf(d::Laplace, lp::Real) = lp < loghalf ? xval(d, logtwo + lp) : xval(d, -(logtwo + log1mexp(lp)))
invlogccdf(d::Laplace, lp::Real) = lp > loghalf ? xval(d, logtwo + log1mexp(lp)) : xval(d, -(logtwo + lp))

function gradlogpdf(d::Laplace, x::Real)
    μ, θ = params(d)
    x == μ && error("Gradient is undefined at the location point")
    g = 1 / θ
    x > μ ? -g : g
end
heslogpdf(d::Laplace{T}, x::Real) where {T<:Real} = zero(T)

function mgf(d::Laplace, t::Real)
    st = d.θ * t
    exp(t * d.μ) / ((1 - st) * (1 + st))
end
function cf(d::Laplace, t::Real)
    st = d.θ * t
    cis(t * d.μ) / (1+st*st)
end

#### Sampling

rand(rng::AbstractRNG, d::Laplace) =
    d.μ + d.θ*randexp(rng)*ifelse(rand(rng, Bool), 1, -1)

#### Fitting

function fit_mle(::Type{<:Laplace}, x::Array)
    xc = copy(x)
    a = median!(xc)
    Laplace(a, StatsBase.mad!(xc, center=a))
end


# structs to represent state transition/occupation priors.

abstract type StatePrior end

Base.broadcastable(s::StatePrior) = Ref(s)

"""
    add(::StatePrior, state, n)

Update the sufficient statistics of the state prior, returning a new version of
the state prior and the index corresponding to the state (possibly changed in
the update).
"""
function add(::StatePrior, state, n) end

"""
    state_to_index(::StatePrior, state)

Convert a state to the index of the corresponding mixture component.  Defaults
to the identity mapping.
"""
state_to_index(::StatePrior, state) = state

function log_prior(p::StatePrior)
    cands = candidates(p)
    log_weights = log_prior.(p, cands)
    log_weights .-= StatsFuns.logsumexp(log_weights)
    return log_weights
end

function Base.rand(p::StatePrior)
    weights = Weights(exp.(log_prior(p)))
    sample(candidates(p), weights)
end

StatsBase.entropy(p::StatePrior) = entropy(exp.(log_prior(p)))

function simulate(p::StatePrior, n::Int)
    states = Int[]
    for _ in 1:n
        p, x = add(p, rand(p))
        push!(states, x)
    end
    p, states
end

struct ChineseRestaurantProcess <: StatePrior
    α::Float64
    N::Vector{Float64}
end

ChineseRestaurantProcess(α::Float64) = ChineseRestaurantProcess(α, Vector{Float64}())

candidates(crp::ChineseRestaurantProcess) = 1:length(crp.N)+1
function add(crp::ChineseRestaurantProcess, x::Int, n::Float64=1.)
    N = copy(crp.N)
    if x == length(N)+1
        push!(N, n)
    else
        N[x] += n
    end
    return ChineseRestaurantProcess(crp.α, N), x
end

log_prior(crp::ChineseRestaurantProcess, x::Int) =
    x == length(crp.N)+1 ? log(crp.α) : log(crp.N[x])

marginal_log_prior(crp::ChineseRestaurantProcess) =
    log_prior = sum(lgamma(n) for n in crp.N) + length(crp.N) * log(crp.α)

"""
    StickyCRP <: StatePrior

Represents a "sticky" Chinese Restaurant Process, where there's a constant,
additional probability that the state at time t is the same as time t-1.  That
is, with probability ρ, ``x_t = x_{t-1}``, and with probability ``(1-\rho)``, a
new ``x`` is sampled according to a CRP.

There's some additional book-keeping that's necessary, as "sticking" transitions
are not informative for the CRP prior, but the proportion of transitions that
are same-state are informative when re-sampling the stickiness parameter.  So we
need to track
1. The number of times a state ``j`` was visited when ``s=0`` (non-sticky).
2. The number of times a state ``j`` was visited when the previous state was
   also ``j``, _regardless_ of whether sticky or not.

The first is stored on `N`, and the second on `Nsame` (and the value of the last
`x` on `last`).

When generating candidates the value `x=0` is used as a sentinal for sticking.
"""
struct StickyCRP <: StatePrior
    α::Float64
    ρ::Float64                  # probability of "sticking"
    last::Int64
    N::Vector{Float64}          # number of non-sticky occuptations of state
    Nsame::Vector{Float64}      # number of transitions to same state (sticky or not)
end

StickyCRP(α::Float64, κ::Float64) = StickyCRP(α, κ, 1, Vector{Float64}(), Vector{Float64}())
Base.isempty(crp::StickyCRP) = isempty(crp.N)

state_to_index(crp::StickyCRP, state::Int) = state == 0 ? crp.last : state

candidates(crp::StickyCRP) = isempty(crp) ? (1:1) : 0:length(crp.N)+1
add(crp::StickyCRP, x::Int, n::Float64=1.0) = add(crp::StickyCRP, state_to_index(crp, x), x==0, n)
function add(crp::StickyCRP, x::Int, sticky::Bool, n::Float64)
    N = copy(crp.N)
    Nsame = copy(crp.Nsame)

    # new component
    if x == length(N)+1
        N = push!(N, 0.)
        Nsame = push!(Nsame, 0.)
    end

    # same transition -> add to Nsame
    if x == crp.last
        Nsame[x] += n
    end

    # sticky -> don't update N
    if !sticky
        N[x] += n
    end

    return StickyCRP(crp.α, crp.ρ, x, N, Nsame), x
end

function log_prior(crp::StickyCRP, x::Int)
    if x == 0
        # sticky: ∝ ρ / (1-ρ) * (∑N+α)
        return logit(crp.ρ) + log(sum(crp.N)+crp.α)
    elseif 0 < x ≤ length(crp.N)
        return log(crp.N[x])
    elseif x == length(crp.N)+1
        return log(crp.α)
    else
        throw(ArgumentError("state $x is invalid for sticky CRP with " *
                            "$(length(crp.N)) components (valid values are " *
                            "0..$(length(crp.N)+1)"))
    end
end

################################################################################
# a changepoint state prior

"""
    struct ChangePoint <: StatePrior

A changepoint prior.

# Fields

* logp::Float64 - log-probability of a change
* k::Int - number of states.
* n::Int - number of data points seen so far.
"""
struct ChangePoint <: Particles.StatePrior
    logp::Float64
    k::Int
    n::Int
end

function ChangePoint(p::Float64)
    0 ≤ p ≤ 1 || throw(ArgumentError("ChangePoint probability p must be between 0 and 1 (got p=$p)"))
    ChangePoint(log(p), 0, 0)
end

function add(cp::ChangePoint, state, n=1.)
    state == cp.k || state == cp.k+1 ||
        throw(ArgumentError("ChangePoint state must be same as last ($(cp.k)) or last+1, got $state"))
    ChangePoint(cp.logp, cp.k + (state>cp.k), cp.n+1), state
end

candidates(cp::ChangePoint) =
    max(cp.k,1):cp.k+1

# prior is 1 for starting transition (k=0).  then p for a change, 1-p for stay
log_prior(cp::ChangePoint, state::Int) =
    cp.k == 0 ? 0. :
    state == cp.k ? StatsFuns.log1mexp(cp.logp) : cp.logp

# marginal prior is 1 if there's just one data point.  if > 1, there are k-1
# changepoints, and n-1 opportunities for a changepoint.
marginal_log_prior(cp::ChangePoint) =
    cp.n == 1 ? 0. : cp.logp*(cp.k-1) + StatsFuns.log1mexp(cp.logp)*(cp.n-cp.k)

"""
    struct NStatePrior <: StatePrior

A state prior with a fixed, finite number of states.

# Fields
* `counts::Vector{Float64}`
* `total_count::Float64`
* `n::Int` number of states
"""
struct NStatePrior <: StatePrior
    counts::Vector{Float64}
    total_count::Float64
    n::Int
end

NStatePrior(counts::Vector{<:Real}) = NStatePrior(counts, sum(counts), length(counts))

function add(ns::NStatePrior, state, n=1.)
    state ∈ 1:ns.n || throw(ArgumentError("Invalid state for NStatePrior: $state " *
                                          "(valid states are 1:$(ns.n))"))
    counts = copy(ns.counts)
    counts[state] += n
    NStatePrior(counts, ns.total_count + n, ns.n), state
end

candidates(ns::NStatePrior) = 1:ns.n
log_prior(ns::NStatePrior, state) = log(ns.counts[state]) - log(ns.total_count)
marginal_log_prior(ns::NStatePrior) = 0.


#== # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# Description
#
#   Functions related to the analysis of the Right Ascension of the Ascending
#   Node (RAAN).
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # ==#

export compute_RAAN_lt, sim_RAAN_J2

"""
    compute_RAAN_lt(JD::Number, asc_node_lt::Number)

Compute the RAAN (0,2π) \\[rad] so that the orbit plane local time is
`asc_node_lt` [hour] at the Julian day `JD`.

"""
function compute_RAAN_lt(JD::Number, asc_node_lt::Number)
    # Get the sun position at noon (UT) represented in the Inertial ref. frame.
    s_i = sun_position_i(JD)

    # Get the desired angle between the Sun and the ascending node [deg].
    alpha = (asc_node_lt-12)*π/12

    # Get the right ascension of the Sun in the Inertial ref. frame. This is the
    # Sun apparent local time.
    SALT = atan(s_i[2],s_i[1])

    # Get the equation of time to compute the Sun mean local time [rad].
    eot = equation_of_time(JD)

    # Compute the Sun mean local time.
    SMLT = SALT + eot

    # Compute the desired RAAN in the interval 0, 2*pi.
    RAAN = mod(SMLT+alpha, 2π)
end

"""
    sim_RAAN_J2(a::Number, e::Number, i::Number, RAAN_0::Number, numDays::Integer)

Simulate the RAAN of an orbit with semi-major axis `a` [m], eccentricity `e`,
inclination `i` [rad], and initial RAAN `RAAN_0` [rad] considering J2
perturbations. The analysis is performed for `numDays` days.

# Returns

A `numDays` × 2 matrix in which the i-th line is:

    | day | RAAN (0,2π) [rad] |

"""
function sim_RAAN_J2(a::Number,
                     e::Number,
                     i::Number,
                     RAAN_0::Number,
                     numDays::Integer)

    # Initialization of variables.
    days = collect(0:1:numDays-1) # Vector of the days in which the RAAN will be
                                  # simulated.

    # RAAN rotation rate [rad/day].
    dOmega = dRAAN(a, e, i, :J2)*24*3600

    # Simulate the RAAN for each day considering just the J2 perturbations.
    RAAN = mod.(RAAN_0 .+ dOmega.*days, 2π)

    [days RAAN]
end


using Plots
plotlyjs()

include("../julia/QQplot.jl");
include("../julia/utils.jl");

x_true = readcsv("../data/spikes.csv")[:];
rs_estimated = readcsv("../data/firing_rate_estimated.csv")[:];

QQplot(rs_estimated, x_true)
title!("Renewal Gamma")
xlabel!("data")
ylabel!("Maximum likelihood")


# This file is a part of BAT.jl, licensed under the MIT License (MIT).

"""
    struct BAT.LogUniform{T<:Real} <: Distributions.Distribution{Univariate,Continuous}

*Experimental feature, not part of stable public API.*

The log-uniform distribution (reciprocal distribution) over an interval ``[a, b]``.

Constructors:

* ```LogUniform(a::Real, b::Real)```

See also [Reciprocal distribution on Wikipedia](https://en.wikipedia.org/wiki/Reciprocal_distribution).
"""
struct LogUniform{T<:Real} <: Distributions.Distribution{Univariate,Continuous}
    a::T
    b::T
    _normf::T
end

function LogUniform{T}(a::Real, b::Real) where {T<:Real}
    @argcheck 0 < T(a) < T(b)
    LogUniform{T}(a, b, inv(log(b / a)))
end

function LogUniform(a::Real, b::Real)
    T = promote_type(typeof(a), typeof(b))
    LogUniform{T}(a, b)
end

LogUniform() = LogUniform{Float64}(1, 2)

Base.minimum(d::LogUniform{T}) where T = d.a
Base.maximum(d::LogUniform{T}) where T = d.b

StatsBase.params(d::LogUniform{T}) where T = (d.a, d.b)
@inline Distributions.partype(::LogUniform{T}) where T = T

Distributions.location(d::LogUniform) = d.a
Distributions.scale(d::LogUniform) = d.b - d.a

Base.eltype(::Type{LogUniform{T}}) where T = T

Statistics.mean(d::LogUniform{T}) where T = (d.b - d.a) * d._normf
StatsBase.median(d::LogUniform{T}) where T = quantile(d, 1//2)
StatsBase.mode(d::LogUniform{T}) where T = d.a
StatsBase.modes(d::LogUniform{T}) where T = Fill(mode(d), 0)

Statistics.var(d::LogUniform{T}) where T = (d.b^2 - d.a^2)/2 * d._normf - mean(d)^2
StatsBase.std(d::LogUniform{T}) where T = sqrt(var(d))
# StatsBase.skewness(d::LogUniform{T}) where T = ...
# StatsBase.kurtosis(d::LogUniform{T}) where T = ...

# StatsBase.entropy(d::LogUniform{T}) where T = ...

function Distributions.logpdf(d::LogUniform{T}, x::U) where {T,U<:Real}
    log(pdf(d, x))
end

function Distributions.pdf(d::LogUniform{T}, x::U) where {T,U<:Real}
    R = promote_type(T,U)
    y = inv(x) * d._normf
    ifelse(insupport(d, x), R(y), zero(R))
end

Distributions.logcdf(d::LogUniform, x::Real) = log(cdf(d, x))

function Distributions.cdf(d::LogUniform{T}, x::U) where {T,U<:Real}
    R = promote_type(T,U)
    p = log(x / d.a) * d._normf
    ifelse(insupport(d, x), R(p), zero(R))
end

Distributions.logccdf(d::LogUniform, x::Real) = log(ccdf(d, x))

function Distributions.ccdf(d::LogUniform{T}, x::U) where {T,U<:Real}
    y = cdf(d, x)
    one(y) - y
end

function Distributions.quantile(d::LogUniform{T}, p::U) where {T,U<:Real}
   R = promote_type(T,U)
   x = exp(p / d._normf) * d.a
   convert(float(R), x)
end

function Distributions.cquantile(d::LogUniform{T}, p::U) where {T,U<:Real}
    y = quantile(d, p)
    one(y) - y
end

# Distributions.mgf(d::LogUniform, t::Real) = ...
# Distributions.cf(d::LogUniform, t::Real) = ...

# # Distributions.gradlogpdf(d::LogUniform, x::Real) = ...

# Implemented implicitly via quantile:
# Base.rand(rng::AbstractRNG, d::LogUniform)

Distributions.truncated(d::LogUniform, l::Real, u::Real) = LogUniform(promote(max(l, d.a), min(u, d.b))...)
Distributions.truncated(d::LogUniform, l::Integer, u::Integer) = LogUniform(promote(max(l, d.a), min(u, d.b))...)


"""
    function csi()
---
## INPUT
- `K` : kernel matrix n × n
- `Y` : target vector n × d
- `m` : maximal rank
- `κ` : trade-off between approximation of `K` and prediction of `Y` (suggested: `.99`)
- `centering` : `1` if centering, `0` otherwise (suggested: 1)
- `δ` : number of columns of cholesky performed in advance (suggested: `40`)
- `tol` : minimum gain at iteration (suggested: `1e-4`)

## OUTPUT
- `G` : Cholesky decomposition -> K(P,P) is approximated by G*G'
- `P` : permutation matrix
- `Q`, `R` : QR decomposition of `G` (or `center(G)` if centering)
- `error1` : `tr(K-G*G')/tr(K)` at each step of the decomposition
- `error2` : `||Y-Q*Q'*Y||_F^2 / ||Y||_F^2` at each step of the decomposition
- `predicted_gain` : predicted gain before adding each column
- `true_gain` : actual gain after adding each column
"""
function csi(K, Y, m; centering = true, κ = 0.99, δ = 40, tol = 1e-4)
    n = size(K, 1)
    d = size(Y, 2)
    size(Y, 1) == n && @error "The targets have the wrong number of data points"
    tdm = min(m + δ, n)         # try dimension
    G = zeros(n, tdm)           # Cholesky factor
    diagK = diag(K)
    P = 1:n
    Q = zeros(n, tdm)           # Q of QR factors
    R = zeros(tdm, tdm)         # R of QR factors
    traceK = tr(K)
    λ = (1 - κ) / traceK
    centering && (Y .-= mean(Y, dims=1))
    sumY² = sum(Y.^2)
    μ = κ / sumY²
    # error1(1) = traceK  # to be confirmed if a vector
    # error2(1) = sumY2;
    k = 0               # current index of the Cholesky decomposition
    kadv = 0            # current index of the look ahead steps
    Dadv = diagK
    D = diagK
    # approximation cost cached quantities
    A1 = zeros(n, 1)
    A2 = zeros(n, 1)
    A3 = zeros(n, 1)
    GTG = zeros(m)
    QTY = zeros(m, d)
    QTYYTQ = zeros(m, m)
    # makes sure that delta is smaller than n
    δ = min(δ, n);

    for i=1:delta # first performs δ steps of Cholesky and QR decomposition
        kadv = kadv + 1
    end
        #=
        # select best index
        diagmax = Dadv(kadv)
        jast = 1
        for j in 1 : n - kadv + 1
            if Dadv(j + kadv - 1) > diagmax / .99
                diagmax = Dadv(j + kadv - 1)
                jast = j
            end
        end
        if diagmax<1e-12
            # all pivots are too close to zero, stops
            # this can only happen if the matrix has rank less than delta
            kadv = kadv - 1;
            break;
        else
            jast = jast + kadv - 1;

            # permute indices
            P( [kadv jast] )        = P( [jast kadv] );
            Dadv( [kadv jast] )     = Dadv( [jast kadv] );
            D( [kadv jast] )        = D( [jast kadv] );
            A1( [kadv jast] )        = A1( [jast kadv] );
            G([kadv jast],1:kadv-1) = G([ jast kadv],1:kadv-1);
            Q([kadv jast],1:kadv-1) = Q([ jast kadv],1:kadv-1);

            # compute new Cholesky column
            G(kadv,kadv)=Dadv(kadv);
            G(kadv,kadv)=sqrt(G(kadv,kadv));
            newKcol = K(P(kadv+1:n),P(kadv));
            G(kadv+1:n,kadv)=1/G(kadv,kadv)*( newKcol - G(kadv+1:n,1:kadv-1)*(G(kadv,1:kadv-1))');

            # update diagonal
            Dadv(kadv+1:n) =  Dadv(kadv+1:n) - G(kadv+1:n,kadv).^2;
            Dadv(kadv) = 0;

            # performs QR
            if centering
                Gcol = G(:,kadv) - 1/n * repmat( sum(G(:,kadv) ,1),n,1 );
            else
                Gcol = G(:,kadv);
            end
            R(1:kadv-1,kadv) = Q(:,1:kadv-1)' * Gcol;
            Q(:,kadv) = Gcol - Q(:,1:kadv-1) * R(1:kadv-1,kadv);
            R(kadv,kadv) = norm(Q(:,kadv));
            Q(:,kadv) = Q(:,kadv) / R(kadv,kadv);

            # update cached quantities
            if centering
                GTG(1:kadv,kadv) = G(:,1:kadv)' * G(:,kadv);
            else
                GTG(1:kadv,kadv) = R(1:kadv,1:kadv)' * R(1:kadv,kadv);
            end
            GTG(kadv,1:kadv) = GTG(1:kadv,kadv)';
            QTY(kadv,:) = Q(:,kadv)' * Y(P,:);
            QTYYTQ(kadv,1:kadv) = QTY(kadv,:) * QTY(1:kadv,:)';
            QTYYTQ(1:kadv,kadv) = QTYYTQ(kadv,1:kadv)';

            # update costs
            A1(kadv:n) = A1(kadv:n) + GTG(kadv,kadv) * ( G(kadv:n,kadv).^2 );
            A1(kadv:n) = A1(kadv:n) + 2 * G(kadv:n,kadv) .* ( G(kadv:n,1:kadv-1) * GTG(1:kadv-1,kadv) );

        end
    end

% compute remaining costs for all indices
A2 = sum( ( G(:,1:kadv) * ( R(1:kadv,1:kadv)' * QTY(1:kadv,:) ) ).^2 , 2);
A3 = sum( ( G(:,1:kadv) * R(1:kadv,1:kadv)' ).^2 , 2);

% start main loop
while k<m,

    k = k +1;

    % compute the gains in approximation for all remaining indices
    dJK = zeros(n-k+1,1);
    for i=1:n-k+1
        kast = k+i-1;
        if D(kast)<1e-12,
            % this column is already generated by already
            % selected columns -> cannot be selected
            dJK(i)=-1e100;
        else
            dJK(i) = A1(kast);
            if kast > kadv,
                % add eta
                dJK(i) = dJK(i) + D(kast)^2 - (D(kast) - Dadv(kast))^2;
            end
            dJK(i) = dJK(i) / D(kast);
        end
    end

    dJY = zeros(n-k+1,1);
    if kadv>k
        for i=1:n-k+1
            kast = k+i-1;
            if A3(kast) < 1e-12
                dJY(i) = 0;
            else
                dJY(i) = A2(kast) / A3(kast);
            end
        end
    end

    % select the best column
    dJ = lambda * dJK + mu * dJY;
    diagmax = -1;
    jast = 0;
    for j=1:n-k+1
        if D(j+k-1)>1e-12
            if dJ(j) > diagmax / 0.9
                diagmax = dJ(j);
                jast = j;
            end
        end
    end
    if jast==0,
        % no more good indices, exit
        k=k-1;
        break;
    end
    jast = jast + k - 1;
    predicted_gain(k) = diagmax;

    % performs one cholesky + QR step:
    % if new pivot not already selected, use pivot
    % otherwise, select new look ahead index that maximize Dadv

    if jast > kadv,
        newpivot = jast;
        jast = kadv + 1;
    else
        a = 1e-12;
        b = 0;
        for j=1:n-kadv
            if Dadv(j+kadv)>a/.99
                a=Dadv(j+kadv);
                b = j+kadv;
            end
        end
        if b==0, newpivot = 0;
        else newpivot = b;
        end
    end

    if newpivot > 0
        % performs steps

        kadv = kadv + 1;

        % permute
        P( [kadv newpivot] ) = P( [newpivot kadv] );
        Dadv( [kadv newpivot] ) = Dadv( [newpivot kadv] );
        D( [kadv newpivot] ) = D( [newpivot kadv] );
        A1( [kadv newpivot] ) = A1( [newpivot kadv] );
        A2( [kadv newpivot] ) = A2( [newpivot kadv] );
        A3( [kadv newpivot] ) = A3( [newpivot kadv] );
        G([kadv newpivot],1:kadv-1)=G([ newpivot kadv],1:kadv-1);
        Q([kadv newpivot],1:kadv-1)=Q([ newpivot kadv],1:kadv-1);

        % compute new Cholesky column
        G(kadv,kadv)=Dadv(kadv);
        G(kadv,kadv)=sqrt(G(kadv,kadv));
        newKcol = K(P(kadv+1:n),P(kadv));
        G(kadv+1:n,kadv)=1/G(kadv,kadv)*( newKcol - G(kadv+1:n,1:kadv-1)*(G(kadv,1:kadv-1))');

        % update diagonal
        Dadv(kadv+1:n) =  Dadv(kadv+1:n) - G(kadv+1:n,kadv).^2;
        Dadv(kadv) = 0;

        % performs QR
        if centering
            Gcol = G(:,kadv) - 1/n * repmat( sum(G(:,kadv) ,1),n,1 );
        else
            Gcol = G(:,kadv);
        end
        R(1:kadv-1,kadv) = Q(:,1:kadv-1)' * Gcol;
        Q(:,kadv) = Gcol - Q(:,1:kadv-1) * R(1:kadv-1,kadv);
        R(kadv,kadv) = norm(Q(:,kadv));
        Q(:,kadv) = Q(:,kadv) / R(kadv,kadv);

        % update the cached quantities
        if centering
            GTG(k:kadv,kadv) = G(:,k:kadv)' * G(:,kadv);
        else
            GTG(k:kadv,kadv) = R(1:kadv,k:kadv)' * R(1:kadv,kadv);
        end
        GTG(kadv,k:kadv) = GTG(k:kadv,kadv)';
        QTY(kadv,:) = Q(:,kadv)' * Y(P,:);
        QTYYTQ(kadv,k:kadv) = QTY(kadv,:) * QTY(k:kadv,:)';
        QTYYTQ(k:kadv,kadv) = QTYYTQ(kadv,k:kadv)';

        % update costs
        A1(kadv:n) = A1(kadv:n) + GTG(kadv,kadv) * ( G(kadv:n,kadv).^2 );
        A1(kadv:n) = A1(kadv:n) + 2 * G(kadv:n,kadv) .* ( G(kadv:n,k:kadv-1) * GTG(k:kadv-1,kadv) );
        A3(kadv:n) = A3(kadv:n) + G(kadv:n,kadv).^2 * sum( R(k:kadv,kadv).^2, 1 );
        temp = R(k:kadv,kadv)' * R(k:kadv,k:kadv-1);
        A3(kadv:n) = A3(kadv:n) + 2 *  G(kadv:n,kadv) .* ( G(kadv:n,k:kadv-1) * temp' );
        temp = R(k:kadv,kadv)' *  QTYYTQ(k:kadv,k:kadv);
        temp1 = temp * R(k:kadv,kadv) ;
        A2(kadv:n) = A2(kadv:n) + G(kadv:n,kadv).^2 * temp1;
        temp2 = temp * R(k:kadv,k:kadv-1);
        A2(kadv:n) = A2(kadv:n) + 2 * G(kadv:n,kadv) .* ( G(kadv:n,k:kadv-1) * temp2' );

    end

    % permute pivots in the Cholesky and QR decomposition between p,q
    p = k;
    q = jast;
    if p<q

        % store some quantities
        Gbef = G(:,p:q);
        Gbeftotal = G(:,k:kadv);
        GTGbef = GTG(p:q,p:q);
        QTYYTQbef = QTYYTQ(p:q,k:kadv);
        Rbef = R(p:q,p:q);
        Rbeftotal = R(k:kadv,k:kadv);
        tempG = eye(q-p+1,q-p+1);
        tempQ = eye(q-p+1,q-p+1);

        for s=q-1:-1:p
            % permute indices
            P( [s s+1] ) = P( [s+1 s] );
            Dadv( [s s+1] ) = Dadv( [s+1 s] );
            D( [s s+1] ) = D( [s+1 s] );
            A1( [s s+1] ) = A1( [s+1 s] );
            A2( [s s+1] ) = A2( [s+1 s] );
            A3( [s s+1] ) = A3( [s+1 s] );
            G([s s+1],1:kadv)=G([s+1 s], 1:kadv);
            Gbef([s s+1],:)=Gbef([s+1 s], :);
            Gbeftotal([s s+1],:)=Gbeftotal([s+1 s], :);
            Q([s s+1],1:kadv)=Q([s+1 s] ,1:kadv);

            % update decompositions
            [Q1,R1] = qr2( G(s:s+1,s:s+1)');
            G(:,s:s+1) = G(:,s:s+1) * Q1;
            G(s,s+1)=0;
            R(1:kadv,s:s+1) = R(1:kadv,s:s+1) * Q1;
            [Q2,R2] = qr2( R(s:s+1,s:s+1) );
            R(s:s+1,1:kadv) = Q2' * R(s:s+1,1:kadv);
            Q(:,s:s+1) = Q(:,s:s+1) * Q2;
            R(s+1,s)=0;

            % update relevant quantities
            nonchanged = [k:s-1 s+2:kadv ];
            GTG(nonchanged,s:s+1) = GTG(nonchanged,s:s+1) * Q1;
            GTG(s:s+1,nonchanged) = GTG(nonchanged,s:s+1)';
            GTG(s:s+1,s:s+1) = Q1' * GTG(s:s+1,s:s+1) * Q1;
            QTY(s:s+1,:) = Q2' * QTY(s:s+1,:);
            QTYYTQ(nonchanged,s:s+1) = QTYYTQ(nonchanged,s:s+1) * Q2;
            QTYYTQ(s:s+1,nonchanged) = QTYYTQ(nonchanged,s:s+1)';
            QTYYTQ(s:s+1,s:s+1) = Q2' * QTYYTQ(s:s+1,s:s+1) * Q2;
            tempG(:,s-p+1:s-p+2) = tempG(:,s-p+1:s-p+2) * Q1;
            tempQ(:,s-p+1:s-p+2) = tempQ(:,s-p+1:s-p+2) * Q2;
        end

        % update costs
        tempG = tempG(:,1);
        tempGG = GTGbef * tempG ;
        A1(k:n) = A1(k:n) - 2 * G(k:n,k) .* ( Gbef(k:n,:) * tempGG );                % between p and q -> different
        A1(k:n) = A1(k:n) - 2 * G(k:n,k) .* ( G(k:n,k:p-1) * GTG(k:p-1,k) );         % below p
        A1(k:n) = A1(k:n) - 2 * G(k:n,k) .* ( G(k:n,q+1:kadv) * GTG(q+1:kadv,k) );   % above q
        tempQ = tempQ(:,1);
        temp = G(k:n,q+1:kadv) * R(k,q+1:kadv)';
        temp = temp + G(k:n,k:p-1) * R(k,k:p-1)';
        temp = temp + Gbef(k:n,:) * Rbef' * tempQ;
        A3(k:n) = A3(k:n) - temp.^2 ;
        A2(k:n) = A2(k:n) + temp.^2 * QTYYTQ(k,k);
        temp2 = ( tempQ' * QTYYTQbef ) * Rbeftotal;
        A2(k:n) = A2(k:n) - 2 * temp .* ( Gbeftotal(k:n,:) * temp2' );

    else

                % update costs
        A1(k:n) = A1(k:n) - 2 * G(k:n,k) .* ( G(k:n,k:kadv) * GTG(k:kadv,k) );
        A3(k:n) = A3(k:n) - ( G(k:n,k:kadv) * R(k,k:kadv)' ).^2;
        temp = G(k:n,k:kadv) * R(k,k:kadv)';
        A2(k:n) = A2(k:n) + ( temp ).^2 * QTYYTQ(k,k);
        temp2 = QTYYTQ(k,k:kadv) * R(k:kadv,k:kadv);
        A2(k:n) = A2(k:n) - 2 * temp .* ( G(k:n,k:kadv) * temp2' );
    end

    % update diagonal and other quantities (A1,B1)
    D(k+1:n) =  D(k+1:n) - G(k+1:n,k).^2;
    D(k) = 0 ;
    A1(k:n) = A1(k:n) + GTG(k,k) * ( G(k:n,k).^2 );

    % compute errors and true gains
    temp2 =  Q(:,k)' * Y(P,:);
    temp2 = sum(temp2.^2);
    temp1 = sum( G(:,k).^2 );
    true_gain(k) = temp1 * lambda + temp2 * mu;
    error1(k+1) = error1(k) - temp1;
    error2(k+1) = error2(k) - temp2;
    if true_gain(k) < tol, break; end
end

% reduce dimensions of decomposition
G=G(:,1:k);
Q=Q(:,1:k);
R=R(1:k,1:k);

% compute and normalize errors
error = lambda * error1 + mu * error2;
error1 = error1 / traceK;
error2 = error2 / sumY2;
    =#
end

"""
    function csi_gaussian(x, α, Y, m; centering=1, κ=.99, δ=40, tol=1e-4)
---
## INPUT
- `x` : input data `nvar` × `n`
- `α` : kernel parameter of the RBF kernel
- `Y` : target vector `n` × `d`
- `m` : maximal rank
- `κ` : trade-off between approximation of K and prediction of Y (suggested: .99)
- `centering` : 1 if centering, 0 otherwise (suggested: 1)
- `δ` : number of columns of cholesky performed in advance (suggested: 40)
- `tol` : minimum gain at iteration (suggested: 1e-4)
## OUTPUT
- `G` : Cholesky decomposition → `K(P,P)` is approximated by `G × G'`
- `P` : permutation matrix
- `Q`, `R` : QR decomposition of `G` (or center(`G`) if centering)
- `error1` : `tr(K - G × G') / tr(K)` at each step of the decomposition
- `error2` : `||Y - Q × Q' × Y||_F^2 / || Y || _F^2` at each step of the decomposition
- `predicted_gain` : predicted gain before adding each column
- `true_gain` : actual gain after adding each column

"""
function csi_gaussian(x, α, Y, m;
                      centering=1, κ=.99, δ=40, tol=1e-4)
    K = exp(-α * sqdist(x, x))
    G, P, Q, R, error1, error2, error, predicted_gain, true_gain =
        csi(K, centering, κ, δ, tol)
    # return G, P, Q, R, error1, error2, error, predicted_gain, true_gain
end

function center_m(M)
    n = size(M, 1)
    for i in 1:size(M, 2)
        M[:, i] .-= mean(M[:, i])
    end
    M
end

"""
    function sqdist(a, b)
---
Computes squared Euclidean distance matrix,
a rectangular matrix of pairwise distances
between points in `a` (given in columns) and points in `b`.

- NB: very fast implementation taken from Roland Bunschoten
- XJ: my tests showed that `aa` and `bb` need to be centered.
"""
function sqdist(a, b)
    aa = sum(a .* a, dims = 1)
    bb = sum(b .* b, dims = 1)
    ab = a'b
    d  = abs.(repeat(aa', outer = (1, size(bb, 2))) +
              repeat(bb, outer = (size(aa, 2), 1)) -
              2ab)
end

"""
    function demo_csi()
---
This is a Julia version of [`csi`](https://www.di.ens.fr/~fbach/csi/csi-1.0.zip).
Here, I try to reproduce the method they described for [**Predictive low-rank decomposition for kernel methods**](https://www.di.ens.fr/~fbach/ICML_2005_CSI_3.pdf).
"""
function demo_csi()
    n, m = 500, 100
    n *= 2                      # why not just n = 1000?
    @debug "Dimensions" n m

    tm = [1 -.5; -.5 1]         # tiny matrix
    x1 = hcat(2tm * randn(2, n ÷ 4) + repeat([-3; 5], outer = (1, n ÷ 4)),
              tm * randn(2, n ÷ 4) + repeat([3; -.5], outer = (1, n ÷ 4)))
    x2 = hcat(tm * randn(2, n ÷ 4) + repeat([2; 3], outer = (1, n ÷ 4)),
              tm * randn(2, n ÷ 4) + repeat([6; 0], outer = (1, n ÷ 4)))
    x = [x1 x2]
    y = [ones(1, n ÷ 2) zeros(1, n ÷ 2)
         zeros(1, n ÷ 2) ones(1, n ÷ 2)]

    rp = randperm(n)

    xtest = x[:, rp[n÷2:end]]
    ytest = y[:, rp[n÷2:end]]
    x = x[:, rp[1:n÷2]]
    y = y[:, rp[1:n÷2]]

    ix = y[1, :] .== 1
    scatter(x[1, ix], x[2, ix])
    ix = y[2, :] .== 1
    scatter!(x[1, ix], x[2, ix])
    savefig("t.pdf")

    α = .5
    κ = 1e-4
    @debug "par" κ
end


x = 3
println(2x^2 - 3x + 1)
# 10
println(1.5x^2 - .5x + 1)
# 13.0

println(zero(Float32))
# 0.0f0
println(zero(Int32))
# 0
println(one(Int32))
# 1
println(one(BigFloat))
# 1.000000000000000000000000000000000000000000000000000000000000000000000000000000


# Compression
# #+HTML: Source
# #+HTML:

# [[file:~/.julia/dev/MatrixProductStates/README.org::*Compression][Compression:1]]
function compress(ψ::MPS{L, T}, to_the::Right; Dcut::Int=typemax(Int)) where {L, T}
    tensors = Array{T, 3}[]

    B = ψ[1]
    d = length(B[1, 1, :])

    @cast Bm[(σ¹, a⁰), a¹] |= B[a⁰, a¹, σ¹]
    U, S, V = psvd(Bm, rank=Dcut)
    #S = S/√sum(S .^ 2)

    @cast A[a⁰, a¹, σ¹] |= U[(σ¹, a⁰), a¹] (σ¹:d)
    push!(tensors, A)

    for i ∈ 2:L
        B = ψ[i]
        d = length(B[1, 1, :])

        @tensor M[aⁱ⁻¹, aⁱ, σⁱ] := (Diagonal(S)*V')[aⁱ⁻¹, aⁱ⁻¹′] * B[aⁱ⁻¹′, aⁱ, σⁱ]
        @cast   Mm[(σⁱ, aⁱ⁻¹), aⁱ] |= M[aⁱ⁻¹, aⁱ, σⁱ]

        U, S, V = psvd(Mm, rank=Dcut)
        #S = S/√sum(S .^ 2)

        @cast A[aⁱ⁻¹, aⁱ, σⁱ] |= U[(σⁱ, aⁱ⁻¹), aⁱ] (σⁱ:d)
        push!(tensors, A)
    end
    MPS{L, T}(tensors), Left()
end

leftcanonical(ψ) = compress(ψ, right)[1]

function compress(ψ::MPS{L, T}, to_the::Left; Dcut::Int=typemax(Int)) where {L, T}
    tensors = Array{T, 3}[]

    A = ψ[L]
    d = length(A[1, 1, :])
    @cast Am[aᴸ⁻¹, (σᴸ, aᴸ)] |= A[aᴸ⁻¹, aᴸ, σᴸ]

    U, S, V = psvd(Am, rank=Dcut)
    #S = S/√sum(S .^ 2)

    @cast B[aᴸ⁻¹, aᴸ, σᴸ] |= V'[aᴸ⁻¹, (σᴸ, aᴸ)] (σᴸ:d)
    push!(tensors, B)

    for i ∈ (L-1):-1:1
        A = ψ[i]
        d = length(A[1, 1, :])
        @tensor M[aⁱ⁻¹, aⁱ, σⁱ]    := A[aⁱ⁻¹, aⁱ′, σⁱ] * (U * Diagonal(S))[aⁱ′, aⁱ]
        @cast   Mm[aⁱ⁻¹, (σⁱ, aⁱ)] |= M[aⁱ⁻¹, aⁱ, σⁱ]

        U, S, V = psvd(Mm, rank=Dcut)
        #S = S/√sum(S .^ 2)

        @cast B[aⁱ⁻¹, aⁱ, σⁱ] |= V'[aⁱ⁻¹, (σⁱ, aⁱ)] (σⁱ:d)
        push!(tensors, B)
    end
    MPS{L, T}(reverse(tensors)), Right()
end

rightcanonical(ψ) = compress(ψ, left)[1]

compress(ψ; Dcut) = compress(ψ, left, Dcut=Dcut)[1]
# Compression:1 ends here


# # fast algorithm from https://stackoverflow.com/questions/217578/how-can-i-determine-whether-a-2d-point-is-within-a-polygon
#
# struct Pt
#     x::Float64
#     y::Float64
# end
#
# # takes two points and generates a function giving the line. All points giving 0 from that function will be on the line
# _points_to_line(pt1::Pt, pt2::Pt) = (pt::Pt) -> (pt2.y - pt1.y)*pt.x + (pt2.x - pt1.x)*pt.y + (pt2.x*pt1.y - pt1.x*pt2.y)
# function _lineintersects(pt1_1, pt1_2, pt2_1, pt2_2,
#                         linefun1 = _points_to_line(pt1_1, pt1_2), linefun2 = _points_to_line(pt2_1, pt2_2))
#     (pt1_1 == pt1_2 || pt2_1 == pt2_2) && return 0 # line has length 0
#     tmp1, tmp2 = linefun1(pt2_1), linefun1(pt2_2)
#     sign(tmp1) == sign(tmp2) && return 0
#     tmp3, tmp4 = linefun2(pt1_1), linefun2(pt1_2)
#     sign(tmp3) == sign(tmp4) && return 0
#     return 1
# end
#
#
# function point_in_polygon(pt::Pt, polygon::AbstractVector{<:Pt},
#                         bbox = (extrema(p.x for p in polygon)..., extrema(p.y for p in polygon)...))
#
#     (pt.x < bbox[1] || pt.x > bbox[2] || pt.y < bbox[3] || pt.y > bbox[4] || bbox[1] == bbox[2] || bbox[3] == bbox[4]) && return false
#
#     raypt = Pt(bbox[1] - (bbox[2] - bbox[1]), pt.y)
#     rayline = _points_to_line(raypt, pt)
#     intersections = 0
#
#     @inbounds for i in eachindex(polygon)[2:end]
#         intersections += _lineintersects(raypt, pt, polygon[i-1], polygon[i], rayline)
#     end
#     intersections += _lineintersects(raypt, pt, last(polygon), first(polygon), rayline)
#     Bool(intersections%2)
# end


using Base.Test
using OptimTools

if !isdefined(:probs) || isempty(probs)
	println("getting test problems")
	include("getTestFunctions.jl")
end

his = []; flag = []; x0 = []; xsol = []

firstOrderMethods = (trdogleg,trcauchy)
# try stopping based on atol
for method in firstOrderMethods
	println("testing $(string(method))")
	for k=1:length(probs)
		prob = probs[k]

		println("\t\t$(prob[1])")

		fk   = prob[2]
		dfk  = prob[3]
		d2fk  = prob[4]
		xsol = float(prob[5])
		x0   = float(xsol) + 10*(1+norm(xsol))* randn(size(xsol))
		atol = 1e-4

		xtr,flag,his = method(fk,dfk,d2fk, x0,maxIter=1000000,atol=atol,out=1*(k==3))

		@test norm(dfk(xtr)) .< atol
	end
end


function F₀(x::AbstractArray{T}, α::T, λ::T) where T<:Real
    out = zero(T)
    @inbounds for i in 1:length(x)
        δᵢ = max(abs(x[i])-λ,0)
        out += δᵢ
    end
    out -= α+λ
    out
end

function F₁(x::AbstractArray{T}, α::T, λ::T) where T<:Real
    if λ==0
        return sum(abs,diff(x))
    else
        x_prox = fit(FusedLasso,x,λ).β
        return sum(abs,diff(x_prox))-α-λ
    end
end

function absepi_projection!(θ_proj::AbstractArray{T}, skip::Int, θ::AbstractArray{T}, α::T) where T<:Real
    copyto!(θ_proj,θ)
    n = length(θ)
    @views if sum(abs, θ[(skip+1):n]) <= α
        return α
    else
        λstar = find_zero(λ -> F₀(θ[(skip+1):n],α,λ),(0,maximum(abs,θ[(skip+1):n])+max(0,-α)),Bisection())
        @inbounds for i in 1:(n-skip)
            θ_proj[skip+i] = sign(θ[skip+i])*max(abs(θ[skip+i])-λstar,0)
        end
        return α + λstar
    end
end

function diffepi_projection!(θ_proj::AbstractArray{T}, skip::Int, θ::AbstractArray{T}, α::T, inv_DDᵀ_D::AbstractArray{T}) where T<:Real
    copyto!(θ_proj,θ)
    n = length(θ)
    if sum(abs, diff(θ[(skip+1):n])) <= α
        return α
    else
        λₘ = maximum(abs, inv_DDᵀ_D*θ[(skip+1):n])
        λstar = find_zero(λ -> F₁(θ[(skip+1):n],α,λ),(0,λₘ+max(0,-α)),Bisection())
        θ_proj[(skip+1):n] = fit(FusedLasso,θ[(skip+1):n],λstar).β
        return α + λstar
    end
end


# # Logarithmic number system

# Computing basic functions like `power`, `exp` and `besselj` is not trivial for reversible programming.
# There is no efficient constant memory algorithm using pure fixed point numbers only.
# For example, to compute `x ^ n` reversiblly with fixed point numbers,
# we need to allocate a vector of size $O(n)$.
# With logarithmic numbers, the above computation is straight forward.

using LogarithmicNumbers
using NiLang, NiLang.AD
using FixedPointNumbers

@i function i_power(y::T, x::T, n::Int) where T
    if !iszero(x)
        @routine begin
            lx ← one(ULogarithmic{T})
            ly ← one(ULogarithmic{T})
            ## convert `x` to a logarithmic number
            ## Here, `*=` is reversible for log numbers
            if x > 0
                lx *= convert(x)
            else
                lx *= convert(-x)
            end
            for i=1:n
                ly *= lx
            end
        end

        ## convert back to fixed point numbers
        y += convert(ly)
        if x < 0 && n%2 == 1
            NEG(y)
        end

        ~@routine
    end
end

# To check the function
i_power(Fixed43(0.0), Fixed43(0.4), 3)

# ## `exp` function as an example
# The following example computes `exp(x)`.

@i function i_exp(y!::T, x::T) where T<:Union{Fixed, GVar{<:Fixed}}
    @invcheckoff begin
        @routine begin
            s ← one(ULogarithmic{T})
            lx ← one(ULogarithmic{T})
            k ← 0
        end
        lx *= convert(x)
        y! += convert(s)
        @from k==0 while s.log > -20
            k += 1
            s *= lx / k
            y! += convert(s)
        end
        ~(@from k==0 while s.log > -20
            k += 1
            s *= x / k
        end)
        lx /= convert(x)
        ~@routine
    end
end

x = Fixed43(3.5)

# We can check the reversibility
out, _ = i_exp(Fixed43(0.0), x)
@assert out ≈ exp(3.5)

# Computing the gradients
_, gx = NiLang.AD.gradient(Val(1), i_exp, (Fixed43(0.0), x))
@assert gx ≈ exp(3.5)