module TSNE

using LinearAlgebra, ProgressBars
using Statistics: mean, mean!
using Printf: @printf, @sprintf

export tsne

"""
    Hbeta!(P::AbstractVector, D::AbstractVector, beta::Number)

Compute the point perplexities `P` given its squared distances to the other
points `D` and the precsion of Gaussian distribution `beta`.
"""
function Hbeta!(P::AbstractVector, D::AbstractVector, beta::Number)
    @inbounds P .= exp.(D .* -beta)
    sumP = sum(P)
    H = log(sumP) + beta * dot(D, P) / sumP
    @inbounds P ./= sumP
    return H
end

"""
    x2p(D::AbstractMatrix{T}, tol::Number = 1e-5, perplexity::Number = 30.0)

Convert `n×n` squared distances matrix `D` into `n×n` perplexities matrix `P`.
Performs a binary search to get P-values in such a way that each conditional
Gaussian has the same perplexity.
"""
function x2p(X::AbstractMatrix{T}, tol::Number = 1e-5,
     perplexity::Number = 30.0; max_iter::Integer = 50,
	 progress = true) where T <: Number
    # Initializing some variables
    n, d = size(X)
    sum_X = sum(x->x^2, X, dims = 2)
	D = fill(zero(T), (n, n))
	BLAS.gemm!('N', 'T', -2.0, X, X, 0.0, D)
    D .+= sum_X .+ sum_X'
    P = fill(zero(T), n, n)
    beta = fill(one(T), n)
    logU = log(perplexity)
    Di = fill(zero(T), n)
    thisP = similar(Di)
    # Loop over all datapoints
	if progress
		pb = ProgressBar(1:n)
		set_description(pb, "Computing P-values...")
	else
		pb = 1:n
	end
    for i in pb
        # Compute the Gaussian kernel and entropy for the current precision
        betamin = 0.0
        betamax = Inf
        betai = 1.0
        copyto!(thisP, view(P, i, :))
        copyto!(Di, view(D, :, i))
        Di[i] = prevfloat(Inf) 		 # exclude D[i,i] from minimum(), yet make it finite and exp(-D[i,i])==0.0
        minD = minimum(Di)     		 # distance of i-th point to its closest neighbour
        @inbounds Di .-= minD  		 # entropy is invariant to offsetting Di, which helps to avoid overflow
        H = Hbeta!(thisP, Di, betai)

        Hdiff = H - logU
        tries = 0

        while abs(Hdiff) > tol && tries < max_iter
            # If not, increase or decrease precision
            if Hdiff > 0.0
                betamin = betai
                betai = isfinite(betamax) ? (betai + betamax) / 2 : betai * 2
            else
                betamax = betai
                betai = (betai + betamin) / 2
            end
            # Recompute the values
            H = Hbeta!(thisP, Di, betai)
            Hdiff = H - logU
            tries += 1
        end
        # Set the final row of P
        @inbounds P[i, :] .= thisP
        beta[i] = betai
    end
    # Return final P matrix
    @printf("Mean σ: %f\n", mean(sqrt.(1.0 ./ beta)))
    return P
end

"""
    pca(X::AbstractMatrix, ndims::Integer = 50)

Run PCA on `X` to reduce the number of its dimensions to `ndims`.
"""
function pca(X::AbstractMatrix, ndims::Integer = 50)
    (n, d) = size(X)
    (d <= ndims) && return X
    Y = X .- mean(X, dims = 1)
    C = Symmetric((Y' * Y) ./ (n - 1))
    Ceig = eigen(C, (d - ndims + 1):d) # take eigvects for top ndims largest eigvals
    return Y * reverse(Ceig.vectors, dims = 2)
end

"""
    tsne(X::AbstractMatrix{T}, no_dims=2, initial_dims=50,
     max_iter::Integer=1000, perplexity=30.0) where T<:Number

Apply t-SNE (t-Distributed Stochastic Neighbor Embedding) to `X`,
i.e. embed its points into `ndims` dimensions preserving close neighbours.

Returns the `point×ndims` matrix of calculated embedded coordinates.

Different from original implementation: the default is not to use PCA for initialization.

### Arguments
	* `no_dims` determines the number of dimensions in the final t-SNE embedding
	* `initial_dims` the number of dimensions of the dataset after apllying
	PCA to initialize the solution
	* `max_iter` how many iterations of t-SNE to perform
	* `perplexity` the number of "effective neighbours" of a datapoint,
	usually increases with the amount of points in the dataset. Typical
	values are between 5 and 50, the default is 30
	* `min_gain`, `eta`, `cheat_scale`, `initial_momentum`, `final_momentum`,
	`stop_cheat_iter`, `momentum_switch_iter` low level parameters of t-SNE optimization
"""
function tsne(X::AbstractMatrix{T}, no_dims = 2, initial_dims::Integer = 50,
     max_iter::Integer = 1000, perplexity::Number = 30.0;
	 initial_momentum::Number = 0.5, final_momentum = 0.8, eta::Integer = 500,
	 min_gain::Number = 0.01, cheat_scale::Number = 4.0, progress = true,
	 stop_cheat_iter::Integer = 100, momentum_switch_iter::Integer = 20) where T <: Number

	X = pca(X, initial_dims)
    n, d = size(X)
    Y = randn(n, no_dims)			  # embedded final matrix
    dY = fill(zero(T), n, no_dims)    # gradient vector
    iY = fill(zero(T), n, no_dims)    # momentum vector
    gains = fill(one(T), n, no_dims)  # how much momentum is affected by gradient
    P = x2p(X, 1e-5, perplexity)
    P .+= P'						  # symmetriztion
    P .*= cheat_scale / sum(P)		  # early exaggeration + normalization
    P .= max.(P, 1e-12)

    # Pre-allocating some matrixes
	L = fill!(similar(P), zero(T))	  # temp matrix for Student-t and gradient steps
    Q = fill!(similar(P), zero(T))	  # temp matrix with low dimensional probabilities
	sum_Y = fill(zero(T), n)
	Y_mean = fill(zero(T), 1, no_dims)
	error = similar(P)
	last_error = NaN

    # Run iterations
	pb = progress ? ProgressBar(1:max_iter) : 1:max_iter
    for iter in pb
        # Compute pairwise affinities
        sum!(x->x^2, sum_Y, Y)
		# L = 2YY'
		BLAS.gemm!('N', 'T', -2.0, Y, Y, 0.0, L)
		# Student-t Distribution
		# L = 1/(1 + sum_Y + sum_Y' + L)
		@inbounds for j in 1:size(L, 2)
			@inbounds for i in 1:size(L, 1)
				L[i, j] = 1 / (1 + sum_Y[j] + sum_Y[i] + L[i, j])
			end
		end
		fill!(sum_Y, 0.0)
        inv_sum_Q = 1.0 / sum(L)
		@inbounds for j = 1:size(Q, 2)
			Pj = view(P, :, j)
			Qj = view(Q, :, j)
			Lj = view(L, :, j)
			# Diagonal should be zero
			Lj[j] = 0.0
			@inbounds for i = 1:size(Lj, 1)
				Qj[i] = ifelse(Lj[i] > 1e-12, Lj[i] * inv_sum_Q, 1e-12)
				# Reusing L for gradient step
				Lj[i] *= (Pj[i] - Qj[i])
				# Reusing sum_Y for column sums
				sum_Y[i] += Lj[i]
			end
		end
		# Compute gradient
		@inbounds for (i, ldiag) in enumerate(sum_Y)
			L[i, i] -= ldiag
		end
		# dY = -4LY
		BLAS.gemm!('N', 'N', -4.0, L, Y, 0.0, dY)
        # Perform the update
        momentum = iter <= momentum_switch_iter ? initial_momentum : final_momentum
        @inbounds for i in eachindex(gains)
            gains[i] = max(ifelse(((dY[i] > 0.) == (iY[i] > 0.)),
                                gains[i] * 0.8,
                                gains[i] + 0.2),
                                min_gain)
            iY[i] = momentum * iY[i] - eta * (gains[i] * dY[i])
            Y[i] += iY[i]
        end
        @inbounds Y .-= mean!(Y_mean, Y)
        # Compute current value of cost function
        if iter % 50 == 0
			map!((x, y)->x * log(x / y), error, P, Q)
			last_error = sum(error)
        end

		progress && set_description(pb, string(@sprintf("Error: %.4f", last_error)))

        # Stop lying about P-values
        if iter == stop_cheat_iter
            P ./= 4.0
        end
    end
    # Return solution
    return Y
end

end  # module TSNE


module TestMonoidalDirectedWiringDiagrams
using Test

using Catlab.Theories, Catlab.WiringDiagrams

# Categorical interface
#######################

# Category
#---------

# Generators
A, B = Ports([:A]), Ports([:B])
f = singleton_diagram(Box(:f,A,B))
g = singleton_diagram(Box(:g,B,A))
@test nboxes(f) == 1
@test boxes(f) == [ Box(:f,A,B) ]
@test nwires(f) == 2

# Composition
@test nboxes(compose(f,g)) == 2
@test boxes(compose(f,g)) == [ Box(:f,A,B), Box(:g,B,A) ]
@test nwires(compose(f,g)) == 3

# Domains and codomains
@test dom(f) == Ports([:A])
@test codom(f) == Ports([:B])
@test dom(compose(f,g)) == Ports([:A])
@test codom(compose(f,g)) == Ports([:A])

# Associativity
@test compose(compose(f,g),f) == compose(f,compose(g,f))

# Identity
@test compose(id(dom(f)), f) == f
@test compose(f, id(codom(f))) == f

# Symmetric monoidal category
#----------------------------

# Domains and codomains
@test dom(otimes(f,g)) == otimes(dom(f),dom(g))
@test codom(otimes(f,g)) == otimes(codom(f),codom(g))

# Associativity and unit
X, Y = Ports([:A,:B]), Ports([:C,:D])
I = munit(Ports)
@test otimes(X,I) == X
@test otimes(I,X) == X
@test otimes(otimes(X,Y),X) == otimes(X,otimes(Y,X))
@test otimes(otimes(f,g),f) == otimes(f,otimes(g,f))

# Braiding
@test compose(braid(X,Y),braid(Y,X)) == id(otimes(X,Y))

# Permutations
W = otimes(X,Y)
@test permute(W, [1,2,3,4]) == id(W)
@test permute(W, [1,2,3,4], inverse=true) == id(W)
@test permute(W, [3,4,1,2]) == braid(X,Y)
@test permute(W, [3,4,1,2], inverse=true) == braid(Y,X)
@test_throws AssertionError permute(W, [1,2])

# Diagonals
#----------

# Basic composition
d = WiringDiagram(dom(f), otimes(codom(f),codom(f)))
fv1 = add_box!(d, first(boxes(f)))
fv2 = add_box!(d, first(boxes(f)))
add_wires!(d, [
  (input_id(d),1) => (fv1,1),
  (input_id(d),1) => (fv2,1),
  (fv1,1) => (output_id(d),1),
  (fv2,1) => (output_id(d),2),
])
@test compose(mcopy(dom(f)), otimes(f,f)) == d

# Domains and codomains
@test dom(mcopy(Ports([:A]))) == Ports([:A])
@test codom(mcopy(Ports([:A]))) == Ports([:A,:A])
@test dom(mcopy(Ports([:A,:B]),3)) == Ports([:A,:B])
@test codom(mcopy(Ports([:A,:B]),3)) == Ports([:A,:B,:A,:B,:A,:B])

# Associativity
A = Ports([:A])
@test compose(mcopy(A), otimes(id(A),mcopy(A))) == mcopy(A,3)
@test compose(mcopy(A), otimes(mcopy(A),id(A))) == mcopy(A,3)

# Commutativity
@test compose(mcopy(A), braid(A,A)) == mcopy(A)

# Unitality
@test compose(mcopy(A), otimes(id(A),delete(A))) == id(A)

# Cartesian categories
A = Ports{CartesianCategory}([:A])
@test mcopy(A) == mcopy(Ports([:A]))
@test delete(A) == delete(Ports([:A]))
@test_throws MethodError mmerge(A)
@test_throws MethodError create(A)

# Codiagonals
#------------

# Domains and codomains
@test dom(mmerge(Ports([:A]))) == Ports([:A,:A])
@test codom(mmerge(Ports([:A]))) == Ports([:A])
@test dom(mmerge(Ports([:A,:B]),3)) == Ports([:A,:B,:A,:B,:A,:B])
@test codom(mmerge(Ports([:A,:B]),3)) == Ports([:A,:B])

# Associativity
A = Ports([:A])
@test compose(otimes(id(A),mmerge(A)), mmerge(A)) == mmerge(A,3)
@test compose(otimes(mmerge(A),id(A)), mmerge(A)) == mmerge(A,3)

# Commutativity
@test compose(braid(A,A), mmerge(A)) == mmerge(A)

# Unitality
@test compose(otimes(id(A),create(A)), mmerge(A)) == id(A)

# Cocartesian categories
A = Ports{CocartesianCategory}([:A])
@test mmerge(A) == mmerge(Ports([:A]))
@test create(A) == create(Ports([:A]))
@test_throws MethodError mcopy(A)
@test_throws MethodError delete(A)

# Bidiagonals
#------------

# Monoidal categories with bidiagonals, and non-naturality of explicit
# representation.
A = Ports{MonoidalCategoryWithBidiagonals}([:A])
@test boxes(mcopy(A)) == [ Junction(:A,1,2) ]
@test boxes(mcopy(otimes(A,A))) == repeat([ Junction(:A,1,2) ], 2)
@test compose(create(A), mcopy(A)) != create(otimes(A,A))
@test compose(mmerge(A), delete(A)) != delete(otimes(A,A))

# Biproduct categories, and naturality of implicit representation.
A = Ports{BiproductCategory}([:A])
@test compose(create(A), mcopy(A)) == create(otimes(A,A))
@test compose(mmerge(A), delete(A)) == delete(otimes(A,A))

# Dagger category
#----------------

f = singleton_diagram(DaggerSymmetricMonoidalCategory, Box(:f,[:A],[:B]))
g = singleton_diagram(DaggerSymmetricMonoidalCategory, Box(:g,[:B],[:A]))

@test boxes(dagger(f)) == [ BoxOp{:dagger}(Box(:f,[:A],[:B])) ]

# Domain and codomain
@test dom(dagger(f)) == codom(f)
@test codom(dagger(f)) == dom(f)

# Functoriality
@test is_permuted_equal(dagger(compose(f,g)), compose(dagger(g),dagger(f)), [2,1])
@test dagger(otimes(f,g)) == otimes(dagger(f),dagger(g))

# Involutivity
@test dagger(dagger(f)) == f
@test dagger(dagger(compose(f,g))) == compose(f,g)
@test dagger(dagger(otimes(f,g))) == otimes(f,g)

# Compact closed category
#------------------------

### Duals

A, B = [ Ports{CompactClosedCategory}([sym]) for sym in [:A, :B] ]
I = munit(typeof(A))

@test boxes(dunit(A)) == [ Junction(:A, [], [PortOp{:dual}(:A), :A]) ]
@test boxes(dcounit(A)) == [ Junction(:A, [:A, PortOp{:dual}(:A)], []) ]

# Domains and codomains
@test dom(dunit(A)) == I
@test codom(dunit(A)) == otimes(dual(A),A)
@test dom(dcounit(A)) == otimes(A,dual(A))
@test codom(dcounit(A)) == I

# Coherence relations
@test is_permuted_equal(dunit(otimes(A,B)),
  compose(dunit(B), otimes(id(dual(B)), dunit(A), id(B))), [2,1])
@test is_permuted_equal(dcounit(otimes(A,B)),
  compose(otimes(id(A), dcounit(B), id(dual(A))), dcounit(A)), [2,1])

### Adjoint mates

f = singleton_diagram(CompactClosedCategory, Box(:f,[:A],[:B]))
g = singleton_diagram(CompactClosedCategory, Box(:g,[:B],[:A]))

@test boxes(mate(f)) == [ BoxOp{:mate}(Box(:f,[:A],[:B])) ]

# Domain and codomain
@test dom(mate(f)) == dual(codom(f))
@test codom(mate(f)) == dual(dom(f))

# Functoriality
@test is_permuted_equal(mate(compose(f,g)), compose(mate(g),mate(f)), [2,1])
@test is_permuted_equal(mate(otimes(f,g)), otimes(mate(g),mate(f)), [2,1])

# Involutivity
@test mate(mate(f)) == f
@test mate(mate(compose(f,g))) == compose(f,g)
@test mate(mate(otimes(f,g))) == otimes(f,g)

# Traced monoidal category
#-------------------------

const TracedMon = TracedMonoidalCategory
A, B, X, Y = [ Ports{TracedMon}([sym]) for sym in [:A,:B,:X,:Y] ]
I = munit(typeof(A))
f = singleton_diagram(TracedMon, Box(:f, [:X,:A], [:X,:B]))

# Domain and codomain
@test dom(trace(X, A, B, f)) == A
@test codom(trace(X, A, B, f)) == B

# Naturality
g = singleton_diagram(TracedMon, Box(:g, [:A], [:A]))
h = singleton_diagram(TracedMon, Box(:h, [:B], [:B]))
@test trace(X, compose(id(X)⊗g, f, id(X)⊗h)) == compose(g, trace(X,f), h)

# Stength, aka superposing
g = singleton_diagram(TracedMon, Box(:g, [:C], [:C]))
@test trace(X, otimes(f,g)) == otimes(trace(X,f), g)

# Symmetry sliding
f = singleton_diagram(TracedMon, Box(:f, [:X,:Y,:A], [:Y,:X,:B]))
@test trace(X⊗Y, compose(f, braid(Y,X)⊗id(B))) ==
  trace(Y⊗X, compose(braid(Y,X)⊗id(A), f))

# Vanishing
f = singleton_diagram(TracedMon, Box(:f, [:X,:Y,:A], [:X,:Y,:B]))
@test trace(X⊗Y, f) == trace(Y, trace(X, f))
f = singleton_diagram(TracedMon, Box(:f, [:A], [:B]))
@test trace(I, f) == f

# Yanking
@test trace(X, braid(X,X)) == id(X)

# Bicategory of relations
#------------------------

A, B = [ Ports{BicategoryRelations}([sym]) for sym in [:A, :B] ]
R = singleton_diagram(BicategoryRelations, Box(:R,[:A],[:B]))
S = singleton_diagram(BicategoryRelations, Box(:S,[:A],[:B]))

# Domains and codomains
@test dom(meet(R,S)) == A
@test codom(meet(R,S)) == B
@test dom(top(A,B)) == A
@test codom(top(A,B)) == B

# Units and counits
@test dunit(A) == merge_junctions(compose(create(A), mcopy(A)))
@test dcounit(A) == merge_junctions(compose(mmerge(A), delete(A)))

# Abelian bicategory of relations
#--------------------------------

A, B = [ Ports{AbelianBicategoryRelations}([sym]) for sym in [:A, :B] ]
R = singleton_diagram(AbelianBicategoryRelations, Box(:R,[:A],[:B]))
S = singleton_diagram(AbelianBicategoryRelations, Box(:S,[:A],[:B]))

# Domains and codomains.
@test dom(plus(A)) == dom(mmerge(A)) && codom(plus(A)) == codom(mmerge(A))
@test dom(coplus(A)) == dom(mcopy(A)) && codom(coplus(A)) == codom(mcopy(A))
@test dom(zero(A)) == dom(create(A)) && codom(zero(A)) == codom(create(A))
@test dom(cozero(A)) == dom(delete(A)) && codom(cozero(A)) == codom(delete(A))

@test dom(join(R,S)) == A
@test codom(join(R,S)) == B
@test dom(bottom(A,B)) == A
@test codom(bottom(A,B)) == B

# Distinctness of "multiplicative" and "additive" (co)diagonals.
@test plus(A) != mmerge(A)
@test coplus(A) != mcopy(A)
@test zero(A) != create(A)
@test cozero(A) != delete(A)

# Junctions
###########

# Add and remove junctions
#-------------------------

A, B, C = [ Ports([sym]) for sym in [:A, :B, :C] ]
f = singleton_diagram(Box(:f, [:A], [:B]))
g = singleton_diagram(Box(:g, [:B], [:C]))

# Copies.
d = compose(f, mcopy(B))
junctioned = compose(f, junction_diagram(B,1,2))
@test add_junctions(d) == junctioned
@test rem_junctions(junctioned) == d

d = compose(mcopy(A), otimes(f,f))
junctioned = compose(junction_diagram(A,1,2), otimes(f,f))
@test is_permuted_equal(add_junctions(d), junctioned, [3,1,2])
@test rem_junctions(junctioned) == d

# Merges.
d = compose(mmerge(A), f)
junctioned = compose(junction_diagram(A,2,1), f)
@test is_permuted_equal(add_junctions(d), junctioned, [2,1])
@test rem_junctions(junctioned) == d

d = compose(otimes(f,f), mmerge(B))
junctioned = compose(otimes(f,f), junction_diagram(B,2,1))
@test add_junctions(d) == junctioned
@test rem_junctions(junctioned) == d

# Deletions.
d = compose(f, delete(B))
junctioned = compose(f, junction_diagram(B,1,0))
@test add_junctions(d) == junctioned
@test rem_junctions(junctioned) == d

# Creations.
d = compose(create(A), f)
junctioned = compose(junction_diagram(A,0,1), f)
@test is_permuted_equal(add_junctions(d), junctioned, [2,1])
@test rem_junctions(junctioned) == d

# Copies, merges, deletions, and creations, all at once.
d = compose(create(A), f, mcopy(B), mmerge(B), g, delete(C))
junctioned = compose(junction_diagram(A,0,1), f, junction_diagram(B,1,2),
                     junction_diagram(B,2,1), g, junction_diagram(C,1,0))
actual = add_junctions(d)
# XXX: An isomorphism test would be more convenient.
perm = [ findfirst([b] .== boxes(actual)) for b in boxes(junctioned) ]
@test is_permuted_equal(actual, junctioned, perm)
@test rem_junctions(junctioned) == d

# Simplify junctions
#-------------------

A = Ports([:A])
j = junction_diagram

# Comonoid laws.
@test merge_junctions(j(A,1,2)⋅(j(A,1,2)⊗id(A))) ==
  j(A,1,3)⋅permute(A⊗A⊗A,[3,1,2]) # FIXME: Shouldn't need permutation.
@test merge_junctions(j(A,1,2)⋅(id(A)⊗j(A,1,2))) == j(A,1,3)
@test merge_junctions(j(A,1,2)⋅(j(A,1,0)⊗id(A))) == j(A,1,1)
@test merge_junctions(j(A,1,2)⋅(id(A)⊗j(A,1,0))) == j(A,1,1)

# Caps and cups.
@test merge_junctions(j(A,0,1)⋅j(A,1,2)) == j(A,0,2)
@test merge_junctions(j(A,2,1)⋅j(A,1,0)) == j(A,2,0)

# Zigzag laws.
@test merge_junctions((id(A)⊗j(A,0,2))⋅(j(A,2,0)⊗id(A))) == j(A,1,1)
@test merge_junctions((j(A,0,2)⊗id(A))⋅(id(A)⊗j(A,2,0))) == j(A,1,1)

end


###################################
## Example code
## for Bayesian Gaussin Mixture Model

using PyPlot, PyCall
push!(LOAD_PATH,".")
import GaussianMixtureModel

"""
Visualize data & estimation in 2D space.
"""
function visualize_2D(X::Matrix{Float64}, S::Matrix{Float64}, S_est::Matrix{Float64}, text)
    cmp = get_cmap("jet")

    K1 = size(S, 1)
    K2 = size(S_est, 1)
    col1 = [pycall(cmp.o, PyAny, Int(round(val)))[1:3] for val in linspace(0,255,K1)]
    col2 = [pycall(cmp.o, PyAny, Int(round(val)))[1:3] for val in linspace(0,255,K2)]

    f, (ax1, ax2) = subplots(1,2,num=text)
    f[:clf]()
    f, (ax1, ax2) = subplots(1,2,num=text)

    for k in 1 : K1
        ax1[:scatter](X[1, S[k,:].==1], X[2, S[k,:].==1], color=col1[k])
    end
    ax1[:set_title]("truth")

    for k in 1 : K2
        ax2[:scatter](X[1, S_est[k,:].==1], X[2, S_est[k,:].==1], color=col2[k])
    end

    ax2[:set_title]("estimation")
end

"""
Run a test script for 2D data clustering.
"""
function test_2D()
    ## set model
    D = 2 # data dimension
    K = 4 #  number of mixture components
    alpha = 100.0 * ones(K)
    beta = 0.1
    m = zeros(D)
    nu = D + 1.0
    W = eye(D)
    cmp = [GaussianMixtureModel.GW(beta, m, nu, W) for _ in 1 : K]
    bgmm = GaussianMixtureModel.BGMM(D, K, alpha, cmp)

    ## generate data
    N = 300
    gmm = GaussianMixtureModel.sample_GMM(bgmm)
    X, S = GaussianMixtureModel.sample_data(gmm, N)

    ## inference
    max_iter = 100
    tic()
    S_est, post_bgmm, VB = GaussianMixtureModel.learn_VI(X, bgmm, max_iter)
    #S_est, post_bgmm, VB = GaussianMixtureModel.learn_GS(X, bgmm, max_iter)
    #S_est, post_bgmm, VB = GaussianMixtureModel.learn_CGS(X, bgmm, max_iter)
    toc()

    ## plot
    visualize_2D(X, S, GaussianMixtureModel.winner_takes_all(S_est), "2D plot")

    # VB check
    figure("ELBO")
    clf()
    plot(VB)
    ylabel("ELBO")
    xlabel("iterations")
    show()
end

test_2D()


# transfer_operators.jl : tranform quantities from fine to coarse grid and back

# TransferKind
abstract type TransferKind end

struct Injection <: TransferKind end

struct FullWeighting <: TransferKind end

struct HalfWeighting <: TransferKind end

struct Cubic <: TransferKind end

Base.broadcastable(kind::TransferKind) = Ref(kind)

# TransferOperator
abstract type TransferOperator end

struct Restriction <: TransferOperator end

struct Interpolation <: TransferOperator end

# GridTransferOperator
struct GridTransferOperator{K<:TransferKind,O<:TransferOperator}
    kind::K
    op::O
end

# stencil definitions
stencil(op::GridTransferOperator{Injection,Restriction}) = @SVector [0,1.,0]

stencil(op::GridTransferOperator{FullWeighting,Restriction}) = @SVector [1/4,1/2,1/4]

stencil(op::GridTransferOperator{FullWeighting,Interpolation}) = @SVector [1/2,1,1/2]

stencil(op::GridTransferOperator{Cubic,Interpolation}) = @SVector [-1/16,0,9/16,1,9/16,0,-1/16]

# TODO: half weighting needs a 2d stencil...

# restriction
R₁(op::TransferKind,n::Int) = expand(GridTransferOperator(op,Restriction()),n) # 1d restriction operator
R(op::TransferKind,n::Int) = R₁(op,n)
R(op::TransferKind,n::Int...) = kron(R₁.(op,n)...)
restrict(u::Vector{T},op::TransferKind,n::Int...) where {T<:AbstractFloat} = R(op,n...)*u

# interpolation
P₁(op::TransferKind,n::Int) = expand(GridTransferOperator(op,Interpolation()),2n) # 1d interpolation operator
P(op::TransferKind,n::Int) = P₁(op,n)
P(op::TransferKind,n::Int...) = kron(P₁.(op,n)...)
interpolate(u::Vector{T},op::TransferKind,n::Int...) where {T<:AbstractFloat} = P(op,n...)*u

# prolongation = interpolation
prolongate(u,op) = interpolate(u,op)

# expand
function expand(op::GridTransferOperator{K,O} where {K,O},n)
    st = stencil(op)
    R = CartesianIndices((n-1,))
    I1, Iend = extrema(R)
    Is = Int[]
    Js = Int[]
    Vs = Float64[]
    @inbounds for I in R
        if iseven(I.I[1])
            is,js,vs = _stencil2mat(st,I,I1,Iend)
            push!(Is,is.>>1...)
            push!(Js,js...)
            push!(Vs,vs...)
        end
    end
    operator2matrix(op,Is,Js,Vs)
end

operator2matrix(op::GridTransferOperator{K,Interpolation} where {K},Is,Js,Vs) = sparse(Js,Is,Vs)
operator2matrix(op::GridTransferOperator{K,Restriction} where {K},Is,Js,Vs) = sparse(Is,Js,Vs)


"""
    get_proportion(for_prop_arr::AbstractVector{Float64})::AbstractVector{Float64}
Compute the proportional values for each element in the 1-Dim array.
returns a AbstractVector{Float64}
# Example
```
julia> get_proportion([1.0,2.0,3.0,4.0,5.0,34.0000034,2423,5656.98988])
=[0.00012301651427098885, 0.0002460330285419777, 0.00036904954281296654, 0.0004920660570839554, 0.0006150825713549442, 0.004182561903469769, 0.298069014078606, 0.6959031763038594]
```
"""
function get_proportion(for_prop_arr::AbstractVector{Float64})::AbstractVector{Float64}
    return for_prop_arr ./ sum(for_prop_arr)
end

"""
    get_proportion_round(for_prop_arr::AbstractVector{Float64};round_digits::Int64)::AbstractVector{Float64}
Compute the proportional values for each element in the array and round to a specified decimal digits.
If the sum of all the elements is not 1 , then the adjustment is done in the 1st largest element
returns a AbstractVector{Float64}
# Example
```
julia> get_proportion_round([1.33,1.33,1.23,1.1111],round_digits = 2)
=[0.26, 0.27, 0.25, 0.22]
```
"""
function get_proportion_round(for_prop_arr::AbstractVector{Float64};round_digits::Int64 = 2)::AbstractVector{Float64}
    temparr=  round.(for_prop_arr ./ sum(for_prop_arr), digits=round_digits, base = 10)
    diff = 1 - sum(temparr) #delta naming
    if diff != 0.0
        max = argmax(temparr)
        temparr[max] =  round(temparr[max] + diff, digits=round_digits, base = 10)
    end
    return temparr
end

"""
    get_proportion_round_add1251(for_prop_arr::AbstractVector{Float64};round_digits::Int64)::AbstractVector{Float64}
Compute the proportional values for each element in the array and round to a specified decimal digits.
If the sum of all the elements is not 1 , then add 1251 to the array and try looking for proportional correctness
after 2 itterations , the adjustment is done in the 1st largest element
returns a AbstractVector{Float64}
# Example
```
julia> get_proportion_round_add1251([1.33,1.33,1.23,1.1111],round_digits = 2)
=[0.26, 0.27, 0.25, 0.22]
```
"""
function get_proportion_round_add1251(for_prop_arr::AbstractVector{Float64};round_digits::Int64 = 2)::AbstractVector{Float64}
    temparr=  round.(for_prop_arr ./ sum(for_prop_arr), digits=round_digits, base = 10)
    diff = 1 - sum(temparr)
    if diff != 0.0
        check = 1
        while check < 2
            sqarr= (for_prop_arr) .+1251
            temparr=  round.(sqarr ./ sum(sqarr), digits=round_digits, base = 10)
            diff = 1 - sum(temparr)
            if diff != 0.00
                check = check + 1
                if check == 2
                    max = argmax(temparr)
                    temparr[max] = round(temparr[max]+diff,digits=round_digits, base = 10)
                end
            else
                break
            end
        end
    end
    return temparr
end


# problem 1a
N = 10
A = zeros(N,N)
for i in 1:N, j in 1:N
    abs(i-j)<=1 ? A[i,j]+=1 : nothing
    i==j ? A[i,j]-=3 : nothing
end
A

# problem 1b
sum(sum(1/j for j in 1:i) for i in 1:25)

#### Prepare Data

X = rand(1000, 3)               # feature matrix
a0 = rand(3)                    # ground truths
y = X * a0 + 0.1 * randn(1000);  # generate response

X2 = hcat(ones(1000),X)
lreg = X2\y

# problem 3
using DataFrames
data = DataFrame(X)
data[:y] = y
lm1 = fit(LinearModel, @formula(y ~ x1 + x2 + x3), data)
@test coef(lm1) ≈ lreg atol=1e-16

# problem 4
r = 2.9:.00005:4; numAttract = 150
steady = ones(length(r),1)*.25
for i=1:400 ## Get to steady state
  steady .= r.*steady.*(1-steady)
end
x = zeros(length(steady),numAttract)
x[:,1] = steady
@inbounds for i=2:numAttract ## Grab values at the attractor
  x[:,i] = r.*x[:,i-1].*(1-x[:,i-1])
end
using Plots
plot(collect(r),x,seriestype=:scatter,markersize=.002,legend=false)

macro ~(y,ex)
  new_ex = Meta.quot(ex)
  quote
    inner_ex = $(esc(new_ex))
    data_name = Symbol(string(inner_ex.args[end])[1])
    eval_ex = Expr(:(=),:data,data_name)
    eval(Main,eval_ex)
    new_X = Matrix{Float64}(size(data,1),length(inner_ex.args)-1)
    cur_spot = 0
    for i in 2:length(inner_ex.args)
      if inner_ex.args[i] == 1
        new_X[:,i-1] = ones(size(data,1))
      else
        col = parse(Int,string(string(inner_ex.args[i])[2]))
        new_X[:,i-1] = data[:,col]
      end
    end
    $(esc(y)),new_X
  end
end

y = rand(10)
X = rand(10,4)
y~1+X1+X2+X4

function solve_least_squares(y,X)
  X\y
end
solve_least_squares(tup::Tuple) = solve_least_squares(tup...)
solve_least_squares(y~1+X1+X2+X4)

function myquantile(d::UnivariateDistribution, q::Number)
    θ = mean(d)
    tol = Inf
    while tol > 1e-5
        θold = θ
        θ = θ - (cdf(d, θ) - q) / pdf(d, θ)
        tol = abs(θold - θ)
    end
    θ
end

for dist in [Gamma(5, 1), Normal(0, 1), Beta(2, 4)]
    @show myquantile(dist, .75)
    @show quantile(dist, .75)
    println()
end

using LightGraphs, Distributions
function mkTree(maxdepth::Int = 10, p::Float64 = 0.8, g::SimpleGraph = Graph(1), currhead::Int = 1)
    if (maxdepth <= 1) g
    else
       b = Binomial(2, p)
       nEdges = max(1, rand(b))
        for leaves in 1:nEdges
            add_vertex!(g)
            newnode = nv(g)
            add_edge!(g, currhead, newnode)
            mkTree(maxdepth-1, p, g, newnode)
        end
    end
    g
end

using Roots
f(x) = 10 - x + e*sin(x)
fzero(f,BigFloat(2.0))

f! = function (x,dx)
  dx[1] = x[1]   + x[2]   + x[3]^2 -12
  dx[2] = x[1]^2 - x[2]   + x[3]   - 2
  dx[3] = 2x[1]  - x[2]^2 + x[3]   - 1
end
using NLsolve
res = nlsolve(f!,[1.0;1.0;1.0])
res.zero
res = nlsolve(f!,[1.0;1.0;1.0],autodiff=true)
res.zero


println("\nTesting cost terms\n")

# Simple Terms

x = Variable(randn(10))
X = Variable(randn(3,4))
A = randn(4,10)
b = randn(4)

cf = norm(x, 0)
@test cf.lambda == 1
@test cf.f(~x) == norm(~x,0)

cf = 3*norm(x, 0)
@test cf.lambda == 3
@test cf.f(~x) == norm(~x,0)

cf = norm(x, 0) <= 3
@test cf.lambda == 1
@test cf.f(~x) == (IndBallL0(3))(~x)

cf = norm(x, 1)
@test cf.lambda == 1
@test cf.f(~x) == (NormL1())(~x)

cf = norm(x, 1) <= 1.5
@test cf.lambda == 1
@test cf.f(~x) == (IndBallL1(1.5))(~x)

cf = 10*norm(x, 1) <= 1.5
@test cf.lambda == 1
@test cf.f(~x) == (IndBallL1(1.5/10))(~x)

cf = norm(x)
@test cf.lambda == 1
@test cf.f(~x) == norm(~x)

cf = pi*norm(x,2)
@test cf.lambda - pi == 0
@test cf.f(~x) == norm(~x)

cf = 3*norm(X,2,1)
@test cf.lambda - 3 == 0
@test cf.f(~X) == sum(  sqrt.(sum((~X).^2, dims=1 )) )

cf = 4*norm(X,2,1,2)
@test cf.lambda - 4 == 0
@test cf.f(~X) == sum(  sqrt.(sum((~X).^2, dims=2 )) )

@test_throws ErrorException 4*norm(X,1,2)

cf = norm(x, 2) <= 2.3
@test cf.lambda == 1
@test cf.f(~x) == (IndBallL2(2.3))(~x)

cf = norm(x, 2) == 2.3
@test cf.lambda == 1
@test cf.f(~x) == (IndSphereL2(2.3))(~x)

cf = norm(x, Inf)
@test cf.lambda == 1
@test cf.f(~x) == norm(~x,Inf)

cf = norm(x, Inf) <= 5.0
@test cf.lambda == 1
@test cf.f(~x) == (IndBallLinf(5.0))(~x)

cf = x <= 3.0
@test cf.lambda == 1
@test cf.f(~x) == (IndBox(-Inf, 3.0))(~x)

cf = 3.0 <= x
@test cf.lambda == 1
@test cf.f(~x) == (IndBox(3.0, Inf))(~x)

cf = x >= 1.0
@test cf.lambda == 1
@test cf.f(~x) == (IndBox(1.0, Inf))(~x)

cf = 1.0 >= x
@test cf.lambda == 1
@test cf.f(~x) == (IndBox(-Inf,1.0))(~x)

cf = x in [-5.0, 5.0]
@test cf.lambda == 1
@test cf.f(~x) == (IndBox(-5.0, 5.0))(~x)

cf = norm(x, 2)^2
@test cf.lambda == 1
@test cf.f(~x) == norm(~x)^2

cf = 0.5*norm(x, 2)^2
@test cf.lambda == 0.5
@test cf.f(~x) == norm(~x)^2

cf = 7*(0.5*norm(x, 2))^2
@test cf.lambda == 7*0.25
@test cf.f(~x) == norm(~x)^2

cf = 2*rank(X) <= 6
@test cf.lambda == 1
@test cf.f(~X) == (IndBallRank(3))(~X)

cf = rank(X)
@test_throws MethodError cf.f(~X)

cf = norm(X,*)
U, S, V = svd(~X)
@test cf.lambda == 1
@test cf.f(~X) == sum(S)

cf = rank(X)
@test_throws MethodError cf.f(~X)

y = randn(size(~x))
cf = hingeloss(x,y)
@test cf.lambda == 1
@test cf.f(~x) == (HingeLoss(y))(~x)

y = randn(size(~x))
cf = sqrhingeloss(x,y)
@test cf.lambda == 1
@test cf.f(~x) == (SqrHingeLoss(y))(~x)

y = randn(size(~x))
cf = logisticloss(x,y)
@test cf.lambda == 1
@test cf.f(~x) == (LogisticLoss(y))(~x)

xp = Variable(rand(10))
bp = rand(Float64, size(~xp))
cf = crossentropy(xp,bp)
@test cf.lambda == 1
@test cf.f(~xp) == (CrossEntropy(bp))(~xp)

cf = logbarrier(x)
@test cf.lambda == 1
@test cf.f(~x) == (LogBarrier(1.0))(~x)

cf = maximum(x)
@test cf.lambda == 1
@test cf.f(~x) == (Maximum(1.0))(~x)

cf = sumpositive(x)
@test cf.lambda == 1
@test cf.f(~x) == (SumPositive())(~x)

a = 1.
cf = huberloss(x,a)
@test cf.lambda == 1
@test cf.f(~x) == (HuberLoss(a))(~x)

a = randn(size(x))
cf = dot(a,x)
@test cf.lambda == 1
@test cf.f(~x) == (Linear(a))(~x)

#IndBinary
lu = (-1.0,randn(length(~x)))
cf = x == lu
@test cf.lambda == 1
@test cf.f(~x) == (IndBinary(lu...))(~x)

# IndAffine (not working in julia < 1.1)
if VERSION.major >= 1 && VERSION.minor >= 1
  cf = A*x-b == 0
  @test cf.lambda == 1
  @test cf.f(~x) == (IndAffine(A,b))(~x)

  cf = (A*x == b)
  @test cf.lambda == 1
  @test cf.f(~x) == (IndAffine(A,-b))(~x)
end

cf = 2*norm(x,1)
ccf = conj(cf)
@test ccf.A == cf.A
@test ccf.f == Conjugate(Postcompose(NormL1(),2))
@test_throws ErrorException conj(norm(randn(2,10)*x,1))

cf = 2*norm(x,1)
ccf = smooth(cf,2.0)
@test ccf.A == cf.A
@test ccf.f(~x) == MoreauEnvelope(Postcompose(NormL1(),2),2.0)(~x)

# Summing terms

x = Variable(10)
cf = ls(x) + 10*norm(x, 1)
@test cf[1].lambda == 1
@test cf[1].f(~x) == 0.5*norm(~x)^2
@test cf[2].lambda == 10
@test cf[2].f(~x) == norm(~x,1)

x = Variable(10)
cf = () #empty cost function
cf += 10*norm(x, 1)
@test length(cf) == 1
@test cf[1].lambda == 10
@test cf[1].f(~x) == 10*norm(~x,1)

x = Variable(10)
cf = () #empty cost function
cf += ls(x) + 10*norm(x, 1)
@test cf[1].lambda == 1
@test cf[1].f(~x) == 0.5*norm(~x)^2
@test cf[2].lambda == 10
@test cf[2].f(~x) == norm(~x,1)

# More complex situations

x = Variable(10)
A = randn(5, 10)
y = Variable(7)
B = randn(5, 7)
b = randn(5)

cf = ls(A*x - b) + norm(x, 1)
@test cf[1].lambda == 1
@test cf[1].f(~x) == 0.5*norm(~x)^2
@test norm(affine(cf[1])*(~x) - (A*(~x)-b)) < 1e-12
@test cf[2].lambda == 1
@test cf[2].f(~x) == norm(~x,1)

cf = ls(A*x - B*y + b) + norm(y, 1) + 5*norm(y, 2)
@test cf[1].lambda == 1
@test cf[1].f(~x) == 0.5*norm(~x)^2
@test cf[2].lambda == 1
@test cf[2].f(~x) == norm(~x,1)
@test cf[3].lambda == 5
@test cf[3].f(~x) == norm(~x,2)

cf = 10*(ls(A*x - B*y + b) + norm(y, 1) + 5*norm(y, 2))
@test cf[1].lambda == 10
@test cf[1].f(~x) == 0.5*norm(~x)^2
@test cf[2].lambda == 10
@test cf[2].f(~x) == norm(~x,1)
@test cf[3].lambda == 50
@test cf[3].f(~x) == norm(~x,2)

cf = 0.5*norm(A*x - B*y + b, 2)^2 + norm(x, 1) + norm(y, 2)
@test cf[1].lambda == 0.5
@test cf[1].f(~x) == norm(~x)^2
@test cf[2].lambda == 1
@test cf[2].f(~x) == norm(~x,1)
@test cf[3].lambda == 1
@test cf[3].f(~x) == norm(~x,2)

# Properties
A = randn(5, 10)
u = Variable(5)
w = Variable(5)
z = Variable(5)

cf = norm(A*x + z)
@test StructuredOptimization.is_smooth(cf) == false
@test StructuredOptimization.is_smooth(cf^2) == true

cf = norm(w + z)^2
@test StructuredOptimization.is_smooth(cf) == true
@test StructuredOptimization.is_AcA_diagonal(cf) == false

cf = norm(x, 1) + norm(y, 2)
@test StructuredOptimization.is_smooth.(cf) == (false,false)
@test StructuredOptimization.is_AcA_diagonal.(cf) == (true,true)


"Implements Zeckendorf's theorem and outputs a Zeckendorf sequence {cᵪ | c ∈ {0, 1}}"
function zeckendorf{T <: Integer}(n::T)
    @assert n >= 0 "n is required to be larger or equal than zero"

    x, z = 2, 2
    fib_cache = [1, 1]

    # find fib(z-1) <= n <= fib(z)
    while z <= n
        x += 1
        z = fib_cache[x-1] + fib_cache[x-2]
        push!(fib_cache, z)
    end

    c = falses(x-1)

    for i in (x-1):-1:1
        if fib_cache[i] <= n
            n -= fib_cache[i]
            c[i] = true
        end
    end

    return c
end

"Generates a binary string from a boolean array in reverse order"
binary_string{T <: AbstractArray{Bool}}(A::T) = join(map(Int, reverse(A)))


# Normal distribution

import StatsFuns
export Normal

import Base: eltype

@measure Normal(μ,σ) ≃ (1/sqrt2π) * Lebesgue(X)

function logdensity(d::Normal{P} , x::X) where {P <: NamedTuple{(:μ, :σ)}, X}
    return - log(d.par.σ)  - (x - d.par.μ)^2 / (2 * d.par.σ^2)
end

# Standard normal

function logdensity(d::Normal{EmptyNamedTuple} , x::X) where {X}
    return - x^2 / 2
end

# Normal() = Normal{EmptyNamedTuple,Real}(NamedTuple())

sampletype(::Normal{NamedTuple{(),Tuple{}}}) = Real

Base.rand(μ::Normal{EmptyNamedTuple}) = randn()

≪(::Normal, ::Lebesgue{X}) where X <: Real = true
representative(::Normal) = Lebesgue(Real)


# Poisson

"""
poipmf!(lambda, prob; left = 0, right = length(prob)-1+left)
poipmf(lambda, right; left = 0)

Compute the p.m.f. of Poisson distribution with mean lambda.
`left` and `right` are the domain of Poisson distribution.
They should be choosen so that the total probability in the domain becomes 1.
In `poipmf!``, the p.m.f. is saved to the vector `prob`.
The right can be obtained from `rightbound`.
The retuen value is the normalizing constant so that the total sum of `prob` is 1, called the weight.
`poipmf` returns a tuple of (weight, prob).
"""

@origin (prob => left) function poipmf!(lambda::Tv, prob::Vector{Tv};
    left::Ti = 0, right::Ti = length(prob)-1+left) where {Tv, Ti}
    @inbounds begin
        log2piOver2::Tv = log(2*pi) / 2
        mode::Ti = floor(Ti, lambda)
        if mode >= 1
            prob[mode] = exp(-lambda + mode * log(lambda)
                - log2piOver2 - (mode + 1/2) * log(mode) + mode)
        else
            prob[mode] = exp(-lambda)
        end
        # down
        for j = mode:-1:left+1
            prob[j-1] = j / lambda * prob[j]
        end
        # up
        for j = mode:right-1
            prob[j+1] = lambda / (j+1) * prob[j]
        end
        # compute W
        weight::Tv = 0
        s::Ti = left
        t::Ti = right
        while s < t
            if prob[s] <= prob[t]
                weight += prob[s]
                s += 1
            else
                weight += prob[t]
                t -= 1
            end
        end
        weight += prob[s]
    end
end

function poipmf(lambda::Tv, right::Ti; left::Ti = 0) where {Tv, Ti}
    prob = Vector{Tv}(undef, right-left+1)
    weight = poipmf!(lambda, prob, left=left, right=right)
    (weight, prob)
end

"""
cpoipmf!(lambda, prob, cprob; left = 0, right = length(prob)-1+left)
cpoipmf(lambda, right; left = 0)

Compute the p.m.f. and complementary c.d.f. of Poisson distribution with mean lambda.
`left` and `right` are the domain of Poisson distribution.
They should be choosen so that the total probability in the domain becomes 1.
In `cpoipmf!``, the p.m.f. and c.d.f. are stored to `prob` and `cprob`, respectively.
The right can be obtained from `rightbound`.
The retuen value is the normalizing constant so that the total sum of `prob` is 1, called the weight.
`cpoipmf` returns a tuple of (weight, prob, cprob).
"""

@origin (prob => left, cprob => left) function cpoipmf!(lambda::Tv, prob::Vector{Tv}, cprob::Vector{Tv}; left::Ti = 0, right::Ti = length(prob)-1+left) where {Tv, Ti}
    weight::Tv = poipmf!(lambda, prob, left=left, right=right)
    @inbounds begin
        cprob[right] = 0
        for k = right-1:-1:left
            cprob[k] = cprob[k+1] + prob[k+1]
        end
        weight
    end
end

function cpoipmf(lambda::Tv, right::Ti; left::Ti = 0) where {Tv, Ti}
    prob = Vector{Tv}(undef, right-left+1)
    cprob = Vector{Tv}(undef, right-left+1)
    weight = cpoipmf!(lambda, prob, cprob, left=left, right=right)
    (weight, prob, cprob)
end

"""
rightbound(::Type{Ti} = Int, lambda::Tv, q::Tv = Tv(1.0e-8))

Compute the rightbound of Poisson distribution with mean lambda.
The rightbound finds a quantile so that the complementary c.d.f. becomes `q`.
"""

function rightbound(lambda::Tv, q::Tv = Tv(1.0e-8))::Int where {Tv}
    rightbound(Int, lambda, q)
end

function rightbound(::Type{Ti}, lambda::Tv, q::Tv = Tv(1.0e-8))::Ti where {Tv, Ti}
    z = cquantile(Normal(), q)
    if lambda < 3.0
        ll = exp(-lambda)
        total = ll
        right::Ti = 0
        while true
            right += 1
            ll *= lambda / right
            total += ll
            if total + q >= 1.0
                break
            end
        end
        right
    else
        right = floor(Ti, (z + sqrt(4.0 * lambda - 1.0))^2 / 4.0 + 1.0)
    end
end


"""
    GaussianLikelihood(σ²)

Gaussian likelihood with `σ²` variance. This is to be used if we assume that the
uncertainity associated with the data follows a Gaussian distribution.

```math
    p(y|f) = Normal(y | f, σ²)
```
On calling, this would return a normal distribution with mean `f` and variance σ².
"""
struct GaussianLikelihood{T<:Real}
    σ²::T
end

GaussianLikelihood() = GaussianLikelihood(1e-6)

@functor GaussianLikelihood

(l::GaussianLikelihood)(f::Real) = Normal(f, sqrt(l.σ²))

(l::GaussianLikelihood)(fs::AbstractVector{<:Real}) = MvNormal(fs, sqrt(l.σ²))

"""
    HeteroscedasticGaussianLikelihood(σ²)

Heteroscedastic Gaussian likelihood.
This is a Gaussian likelihood whose mean and the log of whose variance are functions of the
latent process.

```math
    p(y|[f, g]) = Normal(y | f, exp(g))
```
On calling, this would return a normal distribution with mean `f` and variance `exp(g)`.
"""
struct HeteroscedasticGaussianLikelihood end

(::HeteroscedasticGaussianLikelihood)(f::AbstractVector{<:Real}) = Normal(f[1], exp(f[2]))

(::HeteroscedasticGaussianLikelihood)(fs::AbstractVector) = MvNormal(first.(fs), exp.(last.(fs)))


function getPost(testResult)

    # Create a Dictionary containing the possible combinations:
    P = Dict()

    # Idea: solve by recursive calls of function,
    # constantly updating the prior --> not implemented

    deseaseFrequency = 0.001
    truePositives    = 0.99
    falsePositives   = 0.05

    ### Priors:
    P[":("] = deseaseFrequency
    P[":)"] = 1.0 - P[":("]

    ### Likelihoods
    # Positive test restuls, given status:
    P["+|:("] = truePositives
    P["+|:)"] = falsePositives
    # Negative test restuls, given status:
    P["-|:("] = 1.0 - truePositives
    P["-|:)"] = 1.0 - falsePositives

    # Update for each new measurement:
    for i in 1:length(testResult)
        # Marginal Likelihoods:
        P["+"] = P["+|:("] * P[":("] + P["+|:)"] * P[":)"]
        P["-"] = P["-|:("] * P[":("] + P["-|:)"] * P[":)"]

        ### Posteriors:
        P[":(|+"] = P["+|:("] * P[":("] / P["+"]
        P[":(|-"] = P["-|:("] * P[":("] / P["-"]
        P[":)|+"] = 1.0 - P[":(|+"]
        P[":)|-"] = 1.0 - P[":(|-"]

        # Update priors to the posteriors:
        P[":("] = P[string(":(|",testResult[i])]
        P[":)"] = 1.0 - P[":("]
    end

    # Desired Posteriors are the new Priors:
    # If no test result, then this also works, because the posterior = prior
    posteriors = [P[":("], P[":)"]]

    return posteriors
end

# Enter the received results:
myResult = "+-"

# Show the posteriors for each step when receiving test results,
# including "no result" at the start:
for i in 1:length(myResult)+1
    print("\nAccumulated Results: ")
    println(myResult[1:i-1])
    println("[P(:(|T), P(:)|T)] = ",getPost(myResult[1:i-1]))
    print("Sanity check: Sum(P) = ")
    println(sum(getPost(myResult[1:i-1])))
end


## --- Hafnium isotopes

    # Calculate the initial Hf ratio and epsilon Hf at time t Ma
    function eHf(Hf176_Hf177, Lu176_Hf177, t; eHfOnly::Bool=true)

        # Lutetium decay constant (Soderlund et al., 2004
        lambda = 1.867E-11

        # Present-day CHUR composition (Bouvier et al., 2008)
        CHUR_Hf176_Hf177 = 0.282785
        CHUR_Lu176_Hf177 = 0.0336

        # Calculate initial Hf ratio at time t
        Hf176_Hf177_t = Hf176_Hf177 .- Lu176_Hf177.*(exp.(t .* 10^6*lambda) .- 1)

        # Calculate CHUR Hf ratio at time t
        CHUR_Hf176_Hf177_t = CHUR_Hf176_Hf177 .- CHUR_Lu176_Hf177.*(exp.(t .* 10^6*lambda) .- 1)

        # Calculate corresponding epsilon Hf
        eHf=(Hf176_Hf177_t ./ CHUR_Hf176_Hf177_t .- 1) .* 10^4

        if eHfOnly
            return eHf
        else
            return (eHf, Hf176_Hf177_t)
        end
    end
    export eHf

    function bin_bsr_eHf(x,Hf176_Hf177,Lu176_Hf177,age,min,max,nbins,x_sigma,Hf176_Hf177_sigma,Lu176_Hf177_sigma,age_sigma,nresamples)
        data = hcat(x,Hf176_Hf177,Lu176_Hf177,age)
        sigma = hcat(x_sigma,Hf176_Hf177_sigma,Lu176_Hf177_sigma,age_sigma)

        means = Array{Float64}(undef,nbins,nresamples)
        c = Array{Float64}(undef,nbins)
        for i=1:nresamples
            dbs = bsresample(data,sigma,length(age))
            eHf_resampled = eHf(dbs[:,2], dbs[:,3], dbs[:,4])
            (c,m,s) = binmeans(dbs[:,1], eHf_resampled, min, max, nbins)
            means[:,i] = m
        end

        m = nanmean(means,dim=2)
        el = m - nanpctile(means,2.5,dim=2)
        eu = nanpctile(means,97.5,dim=2) - m

        return (c, m, el, eu)
    end
    export bin_bsr_eHf

## ---


# This file is a part of AstroLib.jl. License is MIT "Expat".
# Copyright (C) 2016 Mosè Giordano.

function _rhotheta{T<:AbstractFloat}(period::T, periastron::T, eccentricity::T,
                                     semimajor_axis::T, inclination::T,
                                     omega::T, omega2::T, epoch::T)
    rho = theta = -one(period)
    # See chapter 55.
    n = 360.0/period
    M = deg2rad(n*(epoch - periastron))
    E = kepler_solver(M, eccentricity)
    r  = semimajor_axis*(1.0 - eccentricity*cos(E))
    nu = trueanom(E, eccentricity)
    # Convert variables in radians.
    omega2      = deg2rad(omega2)
    inclination = deg2rad(inclination)
    omega       = deg2rad(omega)
    theta = omega + atan2(sin(nu + omega2)*cos(inclination), cos(nu + omega2))
    rho   = r*cos(nu + omega2)/cos(theta - omega)
    # Convert theta to degrees and for it to be in [0, 360) range.
    theta = cirrange(rad2deg(theta))
    return rho, theta
end

"""
    rhotheta(period, periastron, eccentricity, semimajor_axis, inclination, omega, omega2, epoch) -> rho, theta

### Purpose ###

Calculate the separation and position angle of a binary star.

### Explanation ###

This function will return the separation \$\\rho\$ and position angle
\$\\theta\$ of a visual binary star derived from its orbital elements.  The
algorithms described in the following book will be used: Meeus J., 1992,
Astronomische Algorithmen, Barth.  Compared to the examples given at page 400
and no discrepancy found.

### Arguments ###

* `period`: period [year]
* `periastro`: time of periastron passage [year]
* `eccentricity`: eccentricity of the orbit
* `semimajor_axis`: semi-major axis [arc second]
* `inclination`: inclination angle [degree]
* `omega`: node [degree]
* `omega2`: longitude of periastron [degree]
* `epoch`: epoch of observation [year]

All input parameters have to be scalars.

### Output ###

The 2-tuple \$(\\rho, \\theta)\$, where

* \$\\rho\$: separation [arc second]
* \$\\theta\$: position angle [degree]

### Example ###

Find the position of Eta Coronae Borealis at the epoch 2016

``` julia
ρ, θ = rhotheta(41.623, 1934.008, 0.2763, 0.907, 59.025, 23.717, 219.907, 2016)
# => (0.6351167848228113,214.42513388052114)
```

### Notes ###

Code of this function is based on IDL Astronomy User's Library.
"""
rhotheta(period::Real, periastron::Real, eccentricity::Real,
         semimajor_axis::Real, inclination::Real,
         omega::Real, omega2::Real, epoch::Real) =
             _rhotheta(promote(float(period), float(periastron),
                               float(eccentricity), float(semimajor_axis),
                               float(inclination), float(omega),
                               float(omega2), float(epoch))...)


#############################################################################
#############################################################################
#
# This file implements factorization
#
#############################################################################
#############################################################################

"""
Factors a polynomial over the field Z_p.

Returns a vector of tuples of (irreducible polynomials (mod p), multiplicity) such that their product of the list (mod p) is f. Irreducibles are fixed points on the function factor.
"""
function factor(f::Polynomial, prime::Int)::Vector{Tuple{Polynomial,Int}}
    #Cantor Zassenhaus factorization

    f_modp = mod(f, prime)
    degree(f_modp) ≤ 1 && return [(f_modp,1)]

    # make f primitive
    ff = prim_part(f_modp)(prime)
    # @show "after prim:", ff

     # make f square-free
    squares_poly = gcd(f, derivative(ff), prime)
    ff = (ff ÷ squares_poly)(prime)
    # @show "after square free:", ff

    # make f monic
    old_coeff = leading(ff).coeff
    ff = (ff ÷ old_coeff)(prime)
    # @show "after monic:", ff

    dds = dd_factor(ff, prime)

    ret_val = Tuple{Polynomial,Int}[]

    for (k,dd) in enumerate(dds)
        sp = dd_split(dd, k, prime)
        sp = map((p)->(p ÷ leading(p).coeff)(prime),sp) #makes the polynomials inside the list sp, monic
        for mp in sp
            push!(ret_val, (mp, multiplicity(f_modp,mp,prime)) )
        end
    end

    #Append the leading coefficient as well
    push!(ret_val, (leading(f_modp).coeff* one(Polynomial), 1) )

    return ret_val
end

"""
Expand a factorization.
"""
function expand_factorization(factorization::Vector{Tuple{Polynomial,Int}})::Polynomial
    length(factorization) == 1 && return first(factorization[1])^last(factorization[1])
    return *([first(tt)^last(tt) for tt in factorization]...)
end

"""
Compute the number of times g divides f
"""
function multiplicity(f::Polynomial, g::Polynomial, prime::Int)::Int
    degree(gcd(f, g, prime)) == 0 && return 0
    return 1 + multiplicity((f ÷ g)(prime), g, prime)
end

"""
Distinct degree factorization.

Given a square free polynomial `f` returns a list, `g` such that `g[k]` is a product of irreducible polynomials of degree `k` for `k` in 1,...,degree(f) ÷ 2, such that the product of the list (mod `prime`) is equal to `f` (mod `prime`).
"""
function dd_factor(f::Polynomial, prime::Int)::Array{Polynomial}
    x = x_poly()
    w = deepcopy(x)
    g = Array{Polynomial}(undef,degree(f)) #Array of polynomials indexed by degree

    #Looping over degrees
    for k in 1:degree(f)
        w = rem(pow_mod(w,prime,prime), f)(prime)
        g[k] = gcd(w - x, f, prime)
        f = (f ÷ g[k])(prime)
    end

    #edge case for final factor
    f != one(Polynomial) && push!(g,f)

    return g
end

"""
Distinct degree split.

Returns a list of irreducible polynomials of degree `d` so that the product of that list (mod prime) is the polynomial `f`.
"""
function dd_split(f::Polynomial, d::Int, prime::Int)::Vector{Polynomial}
    f = mod(f,prime)
    degree(f) == d && return [f]
    degree(f) == 0 && return []
    w = rand(Polynomial, degree = d, monic = true)
    w = mod(w,prime)
    n_power = (prime^d-1) ÷ 2
    g = gcd(pow_mod(w,n_power,prime) - one(Polynomial), f, prime)
    ḡ = (f ÷ g)(prime) # g\bar + [TAB]
    return vcat(dd_split(g, d, prime), dd_split(ḡ, d, prime) )
end


function set_up_operator(p::Param)

    # Variables
    rmin        = p.rmin
    rmax        = p.rmax
    rnodes      = p.rnodes

    thnodes     = p.thnodes

    #Create position vectors
    dr = (rmax - rmin)/(rnodes-1)
    r = [rmin + (i-1)*dr for i in 1:rnodes]

    dth = 2*pi/(thnodes)
    theta = [ (i-1)*dth for i in 1:thnodes]

    #Create Differential Operators
    Dr  =   Diff_Operator_2D(1 ,1, [dr,dth] , [rnodes,thnodes])
    Drr =   Diff_Operator_2D(2 ,1, [dr,dth] , [rnodes,thnodes])

    Dp  =  Diff_Operator_2D_Periodic(1 ,2, [dr,dth] , [rnodes,thnodes])
    Dpp =  Diff_Operator_2D_Periodic(2 ,2, [dr,dth] , [rnodes,thnodes])

    one_over_r2 = spzeros(rnodes*thnodes,rnodes*thnodes)

    for i in 1:rnodes*thnodes
        one_over_r2[i,i] = (1/r[1 + (i-1)%rnodes])^2
    end

    RHS_Opp = Drr + one_over_r2*Dpp + (0.25)*one_over_r2
    #sparse(Drr + (0.25)*one_over_r2*I  + one_over_r2*Dpp);

    av1 = zeros(rnodes*thnodes)
    av2 = zeros(rnodes*thnodes)

    #Operator structure
    return Opps(r, theta , av1 , av2  , Dr, Dp, RHS_Opp)

end


using Plots, LaTeXStrings; pyplot()

xGrid = 0:0.01:10
uGrid = 0:0.01:1
busy = 0.8

F(t)= t<=0 ? 0 : 1 - busy*exp(-(1-busy)t)

infimum(B) = isempty(B) ? Inf : minimum(B)
invF(u) = infimum(filter((x) -> (F(x) >= u),xGrid))

p1 = plot(xGrid,F.(xGrid), c=:blue, xlims=(-0.1,10), ylims=(0,1),
	xlabel=L"x", ylabel=L"F(x)")

p2 = plot(uGrid,invF.(uGrid), c=:blue, xlims=(0,0.95), ylims=(0,maximum(xGrid)),
	xlabel=L"u", ylabel=L"F^{-1}(u)")

plot(p1, p2, legend=:none, size=(800, 400))


"""
    sphere_mesh(ne, radius::FT, wtype::AbstractSphereWarp) where {FT<:AbstractFloat}

Generates a sphere mesh of a given `radius`, warp type and number of elements along each side
of a cubed surface. The sphere mesh is generated by applying the warping function to a cube panel mesh.
"""
function sphere_mesh(ne, radius::FT, wtype::AbstractSphereWarp) where {FT<:AbstractFloat}
    mesh = cube_panel_mesh(ne, FT)
    mesh.coordinates .-= FT(0.5)
    mesh.coordinates .*= (FT(2) * radius)
    for i = 1:mesh.nverts
        mesh.coordinates[i, :] .= cubed_sphere_warp(wtype, mesh.coordinates[i, :]...)
    end
    return mesh
end

"""
    cubed_sphere_warp(::EquiangularSphereWarp, a, b, c)

Given points `(a, b, c)` on the surface of a cube, warp the points out to a
spherical shell of radius `R` based on the equiangular gnomonic grid proposed by
[Ronchi1996]

Returns a tuple of (x1, x2, x3), where x1, x2, x3 are the corresponding coordinates on a sphere
of radius R.

The "Cubed Sphere": A New Method for the Solution of Partial Differential Equations in Spherical Geometry
C. RONCHI,*,1 R. IACONO,* AND P. S. PAOLUCCI†
JOURNAL OF COMPUTATIONAL PHYSICS 124, 93–114 (1996)
https://github.com/CliMA/ClimateMachine.jl/blob/master/src/Numerics/Mesh/Topologies.jl
"""
function cubed_sphere_warp(wtype::EquiangularSphereWarp, a, b, c)
    R = max(abs(a), abs(b), abs(c))
    function f(sR, ξ, η)
        X, Y = tan(π * ξ / 4), tan(π * η / 4)
        ζ1 = sR / sqrt(X^2 + Y^2 + 1)
        ζ2, ζ3 = X * ζ1, Y * ζ1
        ζ1, ζ2, ζ3
    end

    fdim = argmax(abs.((a, b, c)))
    if fdim == 1
        if a < 0
            # (-R, *, *) : formulas for Face I from Ronchi, Iacono, Paolucci (1996)
            #              but for us face IV of the developed net of the cube
            x1, x2, x3 = f(-R, b / a, c / a)
        elseif a > 0
            # ( R, *, *) : formulas for Face III from Ronchi, Iacono, Paolucci (1996)
            #              but for us face II of the developed net of the cube
            x1, x2, x3 = f(R, b / a, c / a)
        end
    elseif fdim == 2
        if b < 0
            # ( *,-R, *) : formulas for Face II from Ronchi, Iacono, Paolucci (1996)
            #              but for us face V of the developed net of the cube
            x2, x1, x3 = f(-R, a / b, c / b)
        elseif b > 0
            # ( *, R, *) : formulas for Face IV from Ronchi, Iacono, Paolucci (1996)
            #              but for us face III of the developed net of the cube
            x2, x1, x3 = f(R, a / b, c / b)
        end
    elseif fdim == 3
        if c > 0
            # ( *, *, R) : formulas for Face V from Ronchi, Iacono, Paolucci (1996)
            #              but for us face VI of the developed net of the cube
            x3, x2, x1 = f(R, b / c, a / c)
        elseif c < 0
            # ( *, *,-R) : formulas for Face VI from Ronchi, Iacono, Paolucci (1996)
            #              but for us face I of the developed net of the cube
            x3, x2, x1 = f(-R, b / c, a / c)
        end
    end

    return x1, x2, x3
end

"""
    cubed_sphere_warp(::EquidistantSphereWarp, a, b, c)

Returns a tuple of (x1, x2, x3), where x1, x2, x3 are the corresponding coordinates on a sphere
of radius R.

https://github.com/CliMA/ClimateMachine.jl/blob/master/src/Numerics/Mesh/Topologies.jl
"""
function cubed_sphere_warp(wtype::EquidistantSphereWarp, a, b, c)
    R = max(abs(a), abs(b), abs(c))
    r = hypot(a, b, c)

    return R * a / r, R * b / r, R * c / r
end

"""
    cube_panel_mesh(ne)

This function builds a cube panel mesh with a resolution of `ne` elements along each edge.

               v8 (xs,xe,xe)          v7 (xe,xe,xe)
                 o--------e11---------o
                /|                   /|
               / |                  / |
              /  |                 /  |
            e12  e8               e10 e7
            /    |               /    |
           /     |            v6/     |
       v5 o--------e9----------o      |
          |    v4o------e3-----|------o v3 (xe,xe,xs)
          |     /   o------->  |     /
          |    /   /           |    /
          e5  e4  /           e6   e2
          |  /   /             |  /
          | /   /              | /
          |/   o               |/
          o--------e1----------o
         v1                    v2
       (xs,xs,xs)               (xe,xs,xs)

       panel 1 => 1 4 3 2
       panel 2 => 2 3 7 6
       panel 3 => 3 4 8 7
       panel 4 => 1 5 8 4
       panel 5 => 1 2 6 5
       panel 6 => 5 6 7 8

       edge  1 => 1 2
       edge  2 => 2 3
       edge  3 => 3 4
       edge  4 => 4 1
       edge  5 => 1 5
       edge  6 => 2 6
       edge  7 => 3 7
       edge  8 => 4 8
       edge  9 => 5 6
       edge 10 => 6 7
       edge 11 => 7 8
       edge 12 => 8 5

            v8 +---e11---+ v7
               | ^       |
               | |       |
              e8 ^   3   e7
               | |       |
            v4 | o--<-o  | v3      v7        v8
  v8 +---e8----+---e3----+----e7---+---e11---+
     | ^       | o-->->  | o-->--> | ^       |
     | |       | |       | |       | |       |
   e12 ^  4   e4 ^  1   e2 ^  2   e10^  6   e12
     | |       | |       | |       | |       |
     | o-<-<-o | o       | o       | o<-<--o |
     +---e5----+---e1----+----e6---+----e9---+
    v5       v1| o-->--o | v2      v6        v5
               |       | |
              e5    5  | e6
               |       ↓ |
               |         |
               +---e9----+
              v5         v6
"""
function cube_panel_mesh(ne::I, ::Type{FT}) where {FT<:AbstractFloat,I<:Integer}
    xs, xe = FT(0), FT(1)
    nverts = (ne + 1)^3 - (ne - 1)^3
    nfaces = 12 * ne + 6 * (2 * ne * (ne - 1))
    nelems = 6 * ne * ne
    nbndry = 0

    nx = ne + 1

    nfaces_edg = 12 * ne

    emat = reshape(1:nelems, ne, ne, 6)
    ndmat = zeros(I, ne + 1, ne + 1)

    panel_verts = [
        1 2 3 1 1 5
        4 3 4 5 2 6
        3 7 8 8 6 7
        2 6 7 4 5 8
    ]
    panel_edges = [
        4 2 3 5 1 9
        3 7 8 12 6 10
        2 10 11 8 9 11
        1 6 7 4 5 12
    ]
    panel_edges_rev = Bool.([
        1 0 0 0 0 0
        1 0 0 1 0 0
        1 1 1 1 1 0
        1 1 1 0 1 0
    ])
    # node coordinates
    xc = range(xs, xe; step = FT(1 / ne)) # [xs, xs+Δ, xs+2Δ, ..., xe]
    xci = view(xc, 2:ne)                   # [xs+Δ, xs+2Δ, ..., xe-Δ]
    xcir = view(xc, ne:-1:2)                # reverse(xci)
    sc = ones(FT, ne - 1) * xs              # [xs, ...., xs]
    ec = ones(FT, ne - 1) * xe              # [xe, ...., xe]
    # where Δ = (xe-xs)/ne
    sc2 = ones(FT, (ne - 1) * (ne - 1)) * xs
    ec2 = ones(FT, (ne - 1) * (ne - 1)) * xe

    xci12 = repeat(xci, outer = ne - 1)
    xci1r2 = repeat(xcir, outer = ne - 1)
    xci21 = repeat(xci, inner = ne - 1)

    edge_nodes = reshape(1:((ne-1)*12), ne - 1, 12) .+ 8
    # coordinates
    coordinates = vcat(
        hcat(
            [xs, xe, xe, xs, xs, xe, xe, xs], # x1,
            [xs, xs, xe, xe, xs, xs, xe, xe], # x2,
            [xs, xs, xs, xs, xe, xe, xe, xe], # x3 vertex coordinates
        ),
        vcat(
            hcat(xci, sc, sc), # edge 1
            hcat(ec, xci, sc), # edge 2
            hcat(xcir, ec, sc), # edge 3
            hcat(sc, xcir, sc), # edge 4
            hcat(sc, sc, xci), # edge 5
            hcat(ec, sc, xci), # edge 6
            hcat(ec, ec, xci), # edge 7
            hcat(sc, ec, xci), # edge 8
            hcat(xci, sc, ec), # edge 9
            hcat(ec, xci, ec), # edge 10
            hcat(xcir, ec, ec), # edge 11
            hcat(sc, xcir, ec), # edge 12
        ),
        hcat(xci21, xci12, sc2), # panel 1
        hcat(ec2, xci12, xci21), # panel 2
        hcat(xci1r2, ec2, xci21), # panel 3
        hcat(sc2, xci21, xci12), # panel 4
        hcat(xci12, sc2, xci21), # panel 5
        hcat(xci12, xci21, ec2), # panel 6
    )
    face_interior = reshape(1:((ne-1)*(ne-1)), ne - 1, ne - 1)

    nfc1i = (nx - 2) * (nx - 1) # panels with normals along local first direction
    nfc2i = (nx - 1) * (nx - 2) # panels with normals along local second direction
    nfci = nfc1i + nfc2i
    fci1 = reshape(1:nfc1i, nx - 2, nx - 1)
    fci2 = reshape(1:nfc2i, nx - 1, nx - 2)

    fcmat1 = zeros(I, nx, nx - 1) # face numbering
    fcmat2 = zeros(I, nx - 1, nx) # for each panel

    edge_faces = reshape(1:(12*ne), ne, 12)

    face_verts = zeros(I, nfaces, 2)
    face_neighbors = zeros(I, nfaces, 5)
    face_boundary = Vector{I}(zeros(nfaces)) # all interior nodes (no boundaries)
    elem_verts = zeros(I, nelems, 4)
    elem_faces = zeros(I, nelems, 4)

    for sfc = 1:6
        ndmat[1, 1],
        ndmat[ne+1, 1],  # panel vertices
        ndmat[ne+1, ne+1],
        ndmat[1, ne+1] = panel_verts[:, sfc]

        ndmat[2:ne, 1] .=
            panel_edges_rev[1, sfc] ? reverse(edge_nodes[:, panel_edges[1, sfc]]) :
            edge_nodes[:, panel_edges[1, sfc]]  # panel edges

        ndmat[end, 2:ne] .=
            panel_edges_rev[2, sfc] ? reverse(edge_nodes[:, panel_edges[2, sfc]]) :
            edge_nodes[:, panel_edges[2, sfc]]

        ndmat[ne:-1:2, end] .=
            panel_edges_rev[3, sfc] ? reverse(edge_nodes[:, panel_edges[3, sfc]]) :
            edge_nodes[:, panel_edges[3, sfc]]

        ndmat[1, ne:-1:2, end] .=
            panel_edges_rev[4, sfc] ? reverse(edge_nodes[:, panel_edges[4, sfc]]) :
            edge_nodes[:, panel_edges[4, sfc]]

        offset = 8 + 12 * (ne - 1) + (sfc - 1) * (ne - 1) * (ne - 1) # interior
        ndmat[2:ne, 2:ne] .= face_interior .+ offset

        fcmat1[1, end:-1:1] .=
            panel_edges_rev[4, sfc] ? reverse(edge_faces[:, panel_edges[4, sfc]]) :
            edge_faces[:, panel_edges[4, sfc]]

        fcmat1[end, :] .=
            panel_edges_rev[2, sfc] ? reverse(edge_faces[:, panel_edges[2, sfc]]) :
            edge_faces[:, panel_edges[2, sfc]]

        fcmat2[:, 1] .=
            panel_edges_rev[1, sfc] ? reverse(edge_faces[:, panel_edges[1, sfc]]) :
            edge_faces[:, panel_edges[1, sfc]]

        fcmat2[end:-1:1, end] .=
            panel_edges_rev[3, sfc] ? reverse(edge_faces[:, panel_edges[3, sfc]]) :
            edge_faces[:, panel_edges[3, sfc]]

        off = ne * 12 + (sfc - 1) * nfci

        fcmat1[2:(end-1), :] .= fci1 .+ off
        fcmat2[:, 2:(end-1)] .= fci2 .+ (off + nfc1i)

        face_verts[fcmat1[:], 1] .= ndmat[:, 1:ne][:] # face nodes
        face_verts[fcmat1[:], 2] .= ndmat[:, 2:(ne+1)][:]
        face_verts[fcmat2[:], 1] .= ndmat[1:ne, :][:]
        face_verts[fcmat2[:], 2] .= ndmat[2:(ne+1), :][:]

        if sfc == 1
            bdy1 = emat[:, 1:1, 5]'
            bdy2 = emat[ne:-1:1, 1:1, 3]'
            bdy3 = emat[1:1, 1:ne, 4]'
            bdy4 = emat[:, 1:1, 2]
        elseif sfc == 2
            bdy1 = emat[ne:ne, :, 5]
            bdy2 = emat[1:1, :, 3]
            bdy3 = emat[:, ne:ne, 1]
            bdy4 = emat[ne:ne, :, 6]'
        elseif sfc == 3
            bdy1 = emat[ne:ne, :, 2]
            bdy2 = emat[:, ne:ne, 4]'
            bdy3 = emat[ne:ne, ne:-1:1, 1]'
            bdy4 = emat[ne:-1:1, ne:ne, 6]
        elseif sfc == 4
            bdy1 = emat[:, 1:1, 1]'
            bdy2 = emat[1:1, :, 6]
            bdy3 = emat[1:1, :, 5]'
            bdy4 = emat[ne:ne, :, 3]'
        elseif sfc == 5
            bdy1 = emat[:, 1:1, 4]'
            bdy2 = emat[1:1, :, 2]
            bdy3 = emat[1:1, :, 1]'
            bdy4 = emat[:, 1:1, 6]
        else # sfc == 6
            bdy1 = emat[ne:ne, :, 4]
            bdy2 = emat[:, ne:ne, 2]'
            bdy3 = emat[:, ne:ne, 5]
            bdy4 = emat[ne:-1:1, ne:ne, 3]
        end
        face_neighbors[fcmat1[:], 1] .= vcat(bdy1, emat[:, :, sfc])[:]
        face_neighbors[fcmat1[:], 3] .= vcat(emat[:, :, sfc], bdy2)[:]
        face_neighbors[fcmat2[:], 1] .= hcat(bdy3, emat[:, :, sfc])[:]
        face_neighbors[fcmat2[:], 3] .= hcat(emat[:, :, sfc], bdy4)[:]

        elem_verts[emat[:, :, sfc][:], :] .= hcat(
            ndmat[1:ne, 1:ne][:], # node numbers (local node 1)
            ndmat[2:(ne+1), 1:ne][:], # for each element (local node 2)
            ndmat[2:(ne+1), 2:(ne+1)][:], #(local node 3)
            ndmat[1:ne, 2:(ne+1)][:], # (local node 4)
        )

        elem_faces[emat[:, :, sfc][:], :] .= hcat(
            fcmat2[:, 1:(nx-1)][:], # (local face 1)
            fcmat1[2:nx, :][:],   # each element (local face 2)
            fcmat2[:, 2:nx][:], # (local face 3)
            fcmat1[1:(nx-1), :][:], # face numbers for (local face 4)
        )
    end

    ref_fc_verts = [
        1 2 3 4
        2 3 4 1
    ]

    for fc = 1:nfaces
        elems = (face_neighbors[fc, 1], face_neighbors[fc, 3])
        for e = 1:2
            el = elems[e]
            localface = findfirst(elem_faces[el, :] .== fc)
            face_neighbors[fc, 2+(e-1)*2] = localface
        end
        # setting up relative orientation
        or1 =
            elem_verts[elems[1], ref_fc_verts[:, face_neighbors[fc, 2]]] ==
            face_verts[fc, :] ? 1 : -1
        or2 =
            elem_verts[elems[2], ref_fc_verts[:, face_neighbors[fc, 4]]] ==
            face_verts[fc, :] ? 1 : -1
        face_neighbors[fc, 5] = or1 * or2
    end

    boundary_tags = sort(unique(face_boundary))
    face_boundary_tags = sort(face_boundary)
    face_boundary = sortperm(face_boundary)

    face_boundary_offset = I[1, length(face_boundary)+1]

    boundary_tag_names = (:interior,)

    # add unique vertex iterator information
    vtconn = map(i -> zeros(I, i), zeros(I, nverts))

    for el = 1:nelems
        for lv = 1:4
            vt = elem_verts[el, lv]
            push!(vtconn[vt], el, lv)
        end
    end
    unique_verts = I[]
    uverts_conn = I[]
    uverts_offset = I.([1])
    for vt = 1:nverts
        lconn = length(vtconn[vt])
        if lconn > 0
            push!(unique_verts, vt)
            push!(uverts_conn, vtconn[vt]...)
            push!(uverts_offset, uverts_offset[end] + lconn)
        end
    end

    return Mesh2D(
        nverts,
        nfaces,
        nelems,
        nbndry,
        coordinates,
        unique_verts,
        uverts_conn,
        uverts_offset,
        face_verts,
        face_neighbors,
        face_boundary,
        boundary_tags,
        boundary_tag_names,
        face_boundary_offset,
        elem_verts,
        elem_faces,
    )
end

function unfold_cube_panel_to_plane(
    x::FTA1D,
    y::FTA1D,
    z::FTA1D,
    panelno::Int,
) where {FT<:AbstractFloat,FTA1D<:AbstractVector{FT}}
    @assert 1 ≤ panelno ≤ 6
    n = length(x)
    xx = similar(x)
    yy = similar(x)

    if panelno == 1
        xx = x
        yy = y
    elseif panelno == 2
        xx = z .+ FT(1)
        yy = y
    elseif panelno == 3
        xx = x
        yy = z .+ FT(1)
    elseif panelno == 4
        xx = -z
        yy = y
    elseif panelno == 5
        xx = x
        yy = -z
    elseif panelno == 6
        xx = 3 .- x
        yy = y
    end

    return xx, yy
end


"""
    K, Cl = h2synthesize(P::ExtendedStateSpace, γ = nothing)

Synthesize H₂-optimal controller K and calculate the closed-loop transfer function from `w` to `z`.
Ref: Cha. 14.5 in Robust and Optimal Control.

If `γ = nothing`, use the formulas for H₂ in Ch 14.5. If γ is a large value, the H∞ formulas are used. As γ → ∞, these two are equivalent. The h∞ formulas do a coordinate transfromation that handles slightly more general systems so if you run into an error, it might be worth trying setting γ to something large, e.g., 1000.
"""
function h2synthesize(P::ExtendedStateSpace, γ = nothing)

    if γ === nothing
        X2, Y2, F2, L2 = _solvematrixequations2(P)
        Â2 = P.A + P.B2*F2 + L2*P.C2
        K = ss(Â2, -L2, F2, 0)
        return K, lft(ss(P), K)
    end

    P̄, Ltrans12, Rtrans12, Ltrans21, Rtrans21 = _transformp2pbar(P)
    X2, Y2, F2, L2 = _solvematrixequations(P̄, γ)

    _checkfeasibility(
        X2,
        Y2,
        γ,
        1e-3,
        10;
        verbose = false,
    ) || throw(DomainError("Not feasible for γ = $γ"))

    K = _synthesizecontroller(
        P̄,
        X2,
        Y2,
        F2,
        L2,
        γ,
        Ltrans12,
        Rtrans12,
        Ltrans21,
        Rtrans21,
    )
    Cl = lft(ss(P), K)
    K, Cl
end

function _solvematrixequations2(P::ExtendedStateSpace)

    A = P.A
    B1 = P.B1
    B2 = P.B2
    C1 = P.C1
    C2 = P.C2
    D11 = P.D11
    D12 = P.D12
    D21 = P.D21
    D22 = P.D22

    P1 = size(C1, 1)
    P2 = size(C2, 1)
    M1 = size(B1, 2)
    M2 = size(B2, 2)

    HX = [A zeros(size(A)); -C1'*C1 -A'] - [B2; -C1'*D12] * [D12'*C1 B2']
    HY = [A' zeros(size(A)); -B1*B1' -A] - [C2'; -B1*D21'] * [D21*B1' C2]

    # Solve matrix equations
    X2 = _solvehamiltonianare(HX)
    Y2 = _solvehamiltonianare(HY)

    F2 = -(B2'X2 + D12'C1)

    L2 = -(B1 * D21' + Y2 * C2')

    return X2, Y2, F2, L2
end


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

using Markdown
using InteractiveUtils

# ╔═╡ f29303e2-1d66-11eb-3313-f335871e9674
using Pkg, DrWatson

# ╔═╡ f2935b8a-1d66-11eb-02e8-3dacec840928
begin
	@quickactivate "StatisticsWithJuliaPlutoNotebooks"
	using LinearAlgebra, StatsBase, DataFrames
end

# ╔═╡ 7877839e-1d66-11eb-0a3c-5953d222ccbc
md"## Listing 1.8"

# ╔═╡ f293d704-1d66-11eb-36d8-4bb2f8dff940
md"##### Transition probability matrix."

# ╔═╡ f2a01cf8-1d66-11eb-3573-d91ed64fb7bc
P = [0.5 0.4 0.1;
     0.3 0.2 0.5;
     0.5 0.3 0.2];

# ╔═╡ f2a0c444-1d66-11eb-0f37-99e69755fe8f
md"##### First way."

# ╔═╡ f2a67846-1d66-11eb-022c-6f0a51c33a42
piProb1 = (P^100)[1,:]

# ╔═╡ f2b43238-1d66-11eb-3713-c5c8f2d4af76
md"##### Second way."

# ╔═╡ f2b4d63e-1d66-11eb-3d9f-65ba9b1ac506
begin
	A = vcat((P' - I)[1:2,:],ones(3)')
	b = [0 0 1]'
	piProb2 = (A\b)[:, 1]
end

# ╔═╡ f2c17bc8-1d66-11eb-0569-9b84c6b33cca
md"##### Third way."

# ╔═╡ f2c2e8fa-1d66-11eb-3421-3d76cd8c9a73
begin
	eigVecs = eigvecs(copy(P'))
	highestVec = eigVecs[:,findmax(abs.(eigvals(P)))[2]]
	piProb3 = Array{Float64}(highestVec)/norm(highestVec,1)
end

# ╔═╡ f2d180ba-1d66-11eb-216b-e3e640811323
md"##### Fourth way."

# ╔═╡ f2da9662-1d66-11eb-1f95-3da69865d370
begin
numInState = zeros(Int,3)
	state = 1
	N = 10^6
	for t in 1:N
		numInState[state] += 1
		global state = sample(1:3,weights(P[state,:]))
	end
	piProb4 = numInState/N
end

# ╔═╡ f2dbd9d2-1d66-11eb-12b9-af2d7f1f73da
DataFrame(:piProb1 => piProb1, :piProb2 => piProb2, :piProb3 => piProb3, :piProb4 => piProb4)

# ╔═╡ f2e37750-1d66-11eb-3f4e-e7b249fb2052
md"## End of listing 1.8"

# ╔═╡ Cell order:
# ╟─7877839e-1d66-11eb-0a3c-5953d222ccbc
# ╠═f29303e2-1d66-11eb-3313-f335871e9674
# ╠═f2935b8a-1d66-11eb-02e8-3dacec840928
# ╟─f293d704-1d66-11eb-36d8-4bb2f8dff940
# ╠═f2a01cf8-1d66-11eb-3573-d91ed64fb7bc
# ╟─f2a0c444-1d66-11eb-0f37-99e69755fe8f
# ╠═f2a67846-1d66-11eb-022c-6f0a51c33a42
# ╟─f2b43238-1d66-11eb-3713-c5c8f2d4af76
# ╠═f2b4d63e-1d66-11eb-3d9f-65ba9b1ac506
# ╟─f2c17bc8-1d66-11eb-0569-9b84c6b33cca
# ╠═f2c2e8fa-1d66-11eb-3421-3d76cd8c9a73
# ╟─f2d180ba-1d66-11eb-216b-e3e640811323
# ╠═f2da9662-1d66-11eb-1f95-3da69865d370
# ╠═f2dbd9d2-1d66-11eb-12b9-af2d7f1f73da
# ╟─f2e37750-1d66-11eb-3f4e-e7b249fb2052


#!/usr/bin/env julia
#
# Jerzy Wroczyński (nr. indeksu 250075)
# Obliczenia Naukowe, Laboratorium
# Lista 5.
#

include("utilities.jl")
using .utilities

include("blocksys.jl")
using .blocksys

include("generate-a-test.jl")
using .testdata

global l = 4

global n = parse(Int, ARGS[1])

let A, b
end

# whether to run with classic built-in algorithm
classic = "classic" in ARGS

# whether to run using a custom sparse matrix solution
custom = "custom" in ARGS

# whether to run improved algorithm with pivoting
pivot = "pivot" in ARGS

# whether to run using the LU decomposition
lu = "lu" in ARGS

println(stderr, "loading data...")
if n == 16
    # read the already pregenerated testfiles
    A = utilities.Amatrix("test-data/n16/A.txt", custom)
    b = utilities.bmatrix("test-data/n16/b.txt")
else
    pathA, pathb = testdata.generateTestData(n, l)
    A = utilities.Amatrix(pathA, custom)
    b = utilities.bmatrix(pathb)
end

let x̃
end

println(stderr, "finding solution...")
@time if classic
    println(stderr, "classic algorithm")
    x̃ = A\b
else
    println(stderr, "improved algorithm")
    let pivot
    end

    # by default don’t perform pivoting
    pivot_ = [i for i in 1:n]

    if pivot
        # calculate the pivot row permutation vector
        pivot_ = blocksys.partialPivot(A, n, l, true)
    end

    if lu
        LU, _ = blocksys.decomposeIntoLU(A, b, n, l, pivot_)
        x̃ = blocksys.solveFromLU(LU, b, n, l, pivot_)
    else
        # calculate the `x` vector
        x̃ = blocksys.gaussianElimination(A, b, n, l, pivot_)
    end
end

# print relative error
println("relative error = ", utilities.relativeError(x̃))


using ApproxFun, Base.Test

## PDEs

d=Interval()^2

x=Fun(identity,d[1]);y=Fun(identity,d[2])

#dirichlet(d) is u[-1,:],u[1,:],u[:,-1],u[:,1]
A=[dirichlet(d);lap(d)]
f=[real(exp(-1+1.0im*y));
                        real(exp(1+1im*y));
                        real(exp(x-1im));
                        real(exp(x+1im))]
S=schurfact(A,100)
@time S=schurfact(A,100)
u=S\f
u=S\f
@time u=S\f;
println("Laplace: should be ~0.014, 0.01")

d=Interval()^2
S=schurfact([neumann(d);lap(d)+100I],100)
@time S=schurfact([neumann(d);lap(d)+100I],100)
u=S\ones(4)
u=S\ones(4)
@time u=S\ones(4)
println("Neumann Helmholtz: should be ~0.016, 0.016")

dx=Interval(0.,1.);dt=Interval(0.0,0.54)
d=dx*dt

V=Fun(x->x^2,dx)

Dt=Derivative(d,[0,1]);Dx=Derivative(d,[1,0])

ϵ=0.1

u0=Fun(x->exp(-25*(x-.5)^2)*exp(-1.0im/(5*ϵ)*log(2cosh(5*(x-.5)))),dx)
L=1im*ϵ*Dt+.5*ϵ^2*Dx^2-V⊗1

PO=discretize([timedirichlet(d);L],50)
@time PO=discretize([timedirichlet(d);L],50)
u=PO\u0
@time    u=PO\u0

println("Schrodinger: should be ~0.013,0.015")


"""
    stumpff!(beta, x, Gz)

Calculates the values of the c-functions and returns the first four g-functions, which will be used to solve Kepler's equation.

```math
c_{0} = \\cos(\\sqrt{z}) \\hspace{3cm} c_{1} = \\frac{\\sin(\\sqrt{z})}{\\sqrt{z}} \\hspace{3cm} c_{n} = \\frac{1}{n!} - z c_{n-2} \\hspace{3cm} G_{n} = x^n c_{n}
```

# Args

* `beta`: a coefficient calculated in function `KeplerFlow`, which equals to `µ`/`a` (where `a` is the semi-major axis and `mu` the standard gravitational parameters of the bodies involved).
* `x`: current estimation of the universal anomaly to solve the universal Kepler equation.
* `Gz`: array of four elements, in which the first four g-functions will be saved.
"""
function stumpff!(beta, x, Gz)
    z = beta * (x^2)

    # C-funtzioak
    if (z < 0)
        Gz[1] = cosh(sqrt(abs(z)))
        Gz[2] = sinh(sqrt(abs(z)))/sqrt(abs(z))
    else
        Gz[1] = cos(sqrt(z))
        Gz[2] = sin(sqrt(z))/sqrt(z)
    end

    for i = 3:4
        Gz[i] = (1.0 - Gz[i-2])/z
    end

    # G-funtzioak
    for i = 2:4
        Gz[i] *= x^(i-1)
    end
end

"""
    keplerSolve!(x0, beta, eta0, zeta0, r0, Gz, dt, trace)

Solves the universal equation of kepler. In this equation, G_{2} and G_{3} are the third and fourth g-functions, respectively.

```math
r_{0}X + \\eta_{0}G_{2} + \\zeta_{0}G_{3} - dt = 0
```

# Args

* `x0`: initial estimation of the universal anomaly.
* `beta`: a coefficient calculated in function `KeplerFlow`, which equals to `µ`/`a` (where `a` is the semi-major axis and `mu` the standard gravitational parameters of the bodies involved).
* `eta0`: a coefficient that is needed to solve, calculated in function `KeplerFlow`.
* `zeta0`: another quantity of the equation it is needed to solve
* `r0`: initial distance between the two bodies.
* `Gz`: array of the first four g-functions. Instead of defining it on each execution of the function, the same array is used in the whole simulation. That is why it's passed as an argument.
* `dt`: time-step of the Kepler-flow.
* `trace`: boolean value that shows information of the execution if true.

# Returns

* `r`: distance between the two bodies after time-step dt.
"""
function keplerSolve!(x0, beta, eta0, zeta0, r0, Gz, dt, trace)

    xi = x0 # initial estimation
    r = r0
    imax = 50
    error = Inf
    tol = sqrt(eps(typeof(xi))) / 100

    for i in 1:imax

        stumpff!(beta, xi, Gz)
        fx = r0*xi + eta0*Gz[3] + zeta0*Gz[4] - dt
        r = r0 + eta0*Gz[2] + zeta0*Gz[3]
        eta = eta0*Gz[1] + zeta0*Gz[2]
        eps = fx / (r - (fx*eta)/(2*r))
        xi -= eps
        error = abs(eps)

        if trace
            println("### $i. Iteration ###")
            println("GZ: $Gz")
            println("R: $r")
            println("X$i: $xi")
            println("Epsilon: $eps\n")
        end

        if error < tol
            if trace
                println("Stop-condition achieved.")
                println("Universal anomaly X = $xi\n")
            end
            return r
        end

    end

    if trace
        println("Newton's method didn't converge.")
        println("Last calculated universal anomaly X = $xi\n")
    end

    return r
end

"""
    keplerFlow(r0_bek, v0_bek, dt, mu, Gz, trace = false)

Defines the position and speeds of a particle -after a time-step- orbiting another body given its initial position and speeds.

# Args

* `r0_bek`: position vector of the secondary body orbiting the primary body.
* `v0_bek`: velocity vector of the secondary body orbiting the primary body.
* `dt`: time-step, the function returns the new position and speed vectors after dt time has passed.
* `mu`: standard gravitational parameter of the two bodies.
* `r0`: initial distance between the two bodies.
* `Gz`: array of the first four g-functions. Instead of defining it on each execution of the function, the same array is used in the whole simulation. That is why it's passed as an argument.
* `trace`: boolean value that shows information of the execution if true.

# Returns

* `r_bek`: the position vector of the body after dt time from the initial values.
* `v_bek`: the position vector of the body after dt time from the initial values.
"""
function keplerFlow(r0_bek, v0_bek, dt, mu, Gz, trace = false)

    r0 = norm(r0_bek)
    v02 = dot(v0_bek,v0_bek)
    eta = dot(r0_bek,v0_bek)
    alpha = mu/r0
    beta = 2.0*alpha - v02
    zeta = mu - beta*r0

    # with beta, we can know the shape of the orbit
    if trace
        if beta > 0.0
            println("Orbit is elliptic.\n")
        elseif beta < 0.0
            println("Orbit is hyperbolic.\n")
        else beta == 0.0
            println("Orbit is parabolic.\n")
        end
    end

    ### Initial estimation ###
    x0 = (dt/r0) * (1.0 - eta/2.0)
    if trace
        println("X0 = $x0\n")
    end

    ### Solve equation ###
    r = keplerSolve!(x0, beta, eta, zeta, r0, Gz, dt, trace)

    rinv = 1.0/r
    ### New position vector ###
    f = -alpha*Gz[3]
    g = r0*Gz[2] + eta*Gz[3]

    r_bek = f * r0_bek + g * v0_bek + r0_bek

    ### New speed vector ###
    df = -alpha*Gz[2]*rinv
    dg = -mu*Gz[3]*rinv

    v_bek = df * r0_bek + dg * v0_bek + v0_bek

    if trace
        println("New position vector: R = $r0_bek")
        println("New speed vector: V = $v0_bek")
    end

    return r_bek, v_bek

end

"""
    keplerFlowWithKahan(r0_bek, v0_bek, dt, mu, Gz, trace, r_k, v_k)

Defines the position and speeds of a particle -after a time-step- orbiting another body given its initial position and speeds. When the new vectors are calculated, Kahan's algorithm is applied.

# Args

* `r0_bek`: position vector of the secondary body orbiting the primary body.
* `v0_bek`: speed vector of the secondary body orbiting the primary body.
* `dt`: time-step, the function returns the new position and speed vectors after dt time has passed.
* `mu`: standard gravitational parameter of the two bodies.
* `r0`: initial distance between the two bodies.
* `Gz`: array of the first four g-functions. Instead of defining it on each execution of the function, the same array is used in the whole simulation. That is why it's passed as an argument.
* `trace`: boolean value that shows information of the execution if true.
* `r_k`: lost information in the sum of the ``f`` and ``g`` coefficients is saved here and used in the next iteration (for positions).
* `v_k`: lost information in the sum of the ``\\dot{f}`` and ``\\dot{g}`` coefficients is saved here and used in the next iteration (for velocities).

# Returns

* `r_bek`: the position vector of the body after dt time-step from the initial values.
* `v_bek`: the velocity vector of the body after dt time-step from the initial values.
"""
function keplerFlowWithKahan(r0_bek, v0_bek, dt, mu, Gz, trace, r_k, v_k)

    r0 = norm(r0_bek)
    v02 = dot(v0_bek,v0_bek)
    eta = dot(r0_bek,v0_bek)
    alpha = mu/r0
    beta = 2.0*alpha - v02
    zeta = mu - beta*r0

    # with beta, we can know the shape of the orbit
    if trace
        if beta > 0.0
            println("Orbit is elliptic.\n")
        elseif beta < 0.0
            println("Orbit is hyperbolic.\n")
        else beta == 0.0
            println("Orbit is parabolic.\n")
        end
    end

    ### Initial estimation ###
    x0 = (dt/r0) * (1.0 - eta/2.0)
    if trace
        println("X0 = $x0\n")
    end

    ### Solve equation ###
    r = keplerSolve!(x0, beta, eta, zeta, r0, Gz, dt, trace)

    rinv = 1.0/r
    ### New position vector ###
    f = -alpha*Gz[3]
    g = r0*Gz[2] + eta*Gz[3]

    r0_bek_aux = copy(r0_bek)
    kahanSum!(r0_bek, f * r0_bek + g * v0_bek, r_k)

    ### New speed vector ###
    df = -alpha*Gz[2]*rinv
    dg = -mu*Gz[3]*rinv

    kahanSum!(v0_bek, df * r0_bek_aux + dg * v0_bek, v_k)

    if trace
        println("New position vector: R = $r0_bek")
        println("New speed vector: V = $v0_bek")
    end

    return r0_bek, v0_bek
end

"""
    keplerStep!(r, v, dt, mu, Gz, trace)

Drifts all particles/bodies under ``\\mathcal{H}_{\\mathrm{Kepler}}`` hamiltonian.

```math
\\mathcal{H}_{\\mathrm{Kepler}}(\\textbf{q}, \\textbf{p}) = \\sum_{i=2}^{N} \\frac{\\textbf{p'}^{2}_{i}}{2 m_{i}'} - \\sum_{i=2}^{N} \\frac{Gm_{i}^{'}M_{i}}{\\mid \\textbf{q}_{i}^{'} \\mid}
```

# Args

* `r`: NxD matrix, where N is the number of bodies and D the dimensions. Initially must contain the initial positions of the bodies, but the output positions will be saved here as well (NxD, where N is the number of bodies  and D is the number of dimensions).
* `v`: NxD matrix. Initially must contain the initial speeds of the bodies, but the output speeds will be saved here as well (NxD, where N is the number of bodies and D is the number of dimensions).
* `dt`: time-step between the initial and final positions and speeds.
* `mu`: N array. Standard gravitational parameter of the N bodies.
* `r0`: initial distance between the two bodies.
* `Gz`: array of the first four g-functions. Instead of defining it on each execution of the function, the same array is used in the whole simulation. That is why it's passed as an argument.
* `trace`: boolean value that shows information of the execution if true.
"""
function keplerStep!(r, v, dt, mu, Gz, trace)
    # drift all particles under H_Kepler
    for i in 2:size(r)[1]
        r[i,:], v[i,:] = keplerFlow(r[i,:], v[i,:], dt, mu[i], Gz, trace)
    end
end

"""
    keplerStepWithKahan!(r, v, dt, mu, Gz, trace, r_k, v_k)

Drifts all particles/bodies under ``\\mathcal{H}_{\\mathrm{Kepler}}`` hamiltonian.

```math
\\mathcal{H}_{\\mathrm{Kepler}}(\\textbf{q}, \\textbf{p}) = \\sum_{i=2}^{N} \\frac{\\textbf{p'}^{2}_{i}}{2 m_{i}'} - \\sum_{i=2}^{N} \\frac{Gm_{i}^{'}M_{i}}{\\mid \\textbf{q}_{i}^{'} \\mid}
```

# Args

* `r`: NxD matrix, where N is the number of bodies and D the dimensions. Initially must contain the initial positions of the bodies, but the output positions will be saved here as well.
* `v`: NxD matrix. Initially must contain the initial speeds of the bodies, but the output speeds will be saved here as well.
* `dt`: time-step between the initial and final positions and speeds.
* `mu`: array of N elements, containing standard gravitational parameters of the N bodies.
* `r0`: initial distance between the two bodies.
* `Gz`: array of the first four g-functions. Instead of defining it on each execution of the function, the same array is used in the whole simulation. That is why it's passed as an argument.
* `trace`: boolean value that shows information of the execution if true.
* `r_k`: auxiliary variable to save extra information of Kahan's algorithm (positions). Same dimension as `r`.
* `v_k`: auxiliary variable to save extra information of Kahan's algorithm (velocities). Same dimension as `v`.
"""
function keplerStepWithKahan!(r, v, dt, mu, Gz, trace, r_k, v_k)
    ### drift all particles under H_Kepler
    for i in 2:size(r)[1]
        r[i,:], v[i,:] =keplerFlowWithKahan(r[i,:], v[i,:], dt, mu[i], Gz, trace, r_k[i,:], v_k[i,:])
    end
end


using Plots
pyplot()

#---------------------------------------------------------------------#
#This code plots the Stability of for the 1D Wave Equation.
#Written by F.X. Giraldo on 03/2019
#           Department of Applied Maths
#           Naval Postgraduate School
#           Monterey; CA 93943-5216
#Variables:
#---------------------------------------------------------------------#
#method=input(" 1=FE & Upwind \n 2=FE & Centered \n 3=FE & Downwind \n 4=BE & Upwind \n 5=BE & Centered \n 6=BE & Downwind \n 7=TR & Upwind \n 8=TR & Centered \n 9=TR & Downwind \n Enter Method: ")
method=1

n=100
dtheta=2*pi/n
theta=0:dtheta:2*pi
if (method <= 3)
    cmax=1
elseif [method >= 4]
    cmax=10
end
m=10*cmax
dc=cmax/m
c=0:dc:cmax
em1=exp.(-im*theta)
ep1=exp.(+im*theta)

a=ones(size(em1))
#Plot Stability of Specific Equations
xmatrix=Array{Float64}(undef,n,m+2)
ymatrix=Array{Float64}(undef,n,m+2)
if method == 1 #FE + Upwind
    for j=1:m+1
        z=a - c[j]*(a - em1) #1st order upwind [explicit 1st order in time]
        xmatrix[:,j]=real(z)
        ymatrix[:,j]=imag(z)
    end
elseif method == 2 #FE + Centered
    for j=2:m+1
        z=a - 0.5*c[j]*(ep1 - em1); #2nd order centered [explicit 1st order in time]
        xmatrix[:,j]=real(z)
        ymatrix[:,j]=imag(z)
    end
elseif method == 3 #FE + Downwind
    for j=2:m+1
        z=a - c[j]*(ep1 - a); #1st order downwind [explicit 1st order in time]
        xmatrix[:,j]=real(z)
        ymatrix[:,j]=imag(z)
    end
elseif method == 4 #BE + Upwind
    for j=2:m+1
        z=a/(a + c[j]*(a - em1)) #1st order uwpind, (implicit 1st order in time)
        xmatrix[:,j]=real(z)
        ymatrix[:,j]=imag(z)
    end
elseif method == 5 #BE + Centered
    for j=2:m+1
        z=a/(a + 0.5*c[j]*(ep1 - em1)) #2nd order centered, (implicit 1st order in time)
        xmatrix[:,j]=real(z)
        ymatrix[:,j]=imag(z)
    end
elseif method == 6 #BE + Downwind
    for j=2:m+1
        z=a/(a + c[j]*(ep1 - a)) #1st order downwind, (implicit 1st order in time)
        xmatrix[:,j]=real(z)
        ymatrix[:,j]=imag(z)
    end
elseif method == 7 #TR + Upwind
    for j=2:m+1
        z=(a - 0.5*c[j]*(a - em1))./(a + 0.5*c[j]*(a - em1)) #1st order upwind [Trapezoidal 2nd order in time]
        xmatrix[:,j]=real(z)
        ymatrix[:,j]=imag(z)
    end
elseif method == 8 #TR + Centered
    for j=2:m+1
        z=(a - 0.25*c[j]*(ep1 - em1))./(a + 0.25*c[j]*(ep1 - em1)) #2nd order centered [TR 2nd order in time]
        xmatrix[:,j]=real(z)
        ymatrix[:,j]=imag(z)
    end
elseif method == 9 #TR + Downwind
    for j=2:m+1
        z=(a - 0.5*c[j]*(ep1 - a))./(a + 0.5*c[j]*(ep1 - a)) #1st order downwind [TR 2nd order in time]
        xmatrix[:,j]=real(z)
        ymatrix[:,j]=imag(z)
    end
end

#Plot Unit Circle
z=ep1
xmatrix[:,m+2]=real(z)
ymatrix[:,m+2]=imag(z)
#plot_handle=plot(xmatrix,ymatrix,leg=false,w=2,aspect_ratio=1)
plot(xmatrix,ymatrix,leg=false,w=1,line=(:dash,1),aspect_ratio=1)
plot!(real(z),imag(z),leg=false,w=1,color=[:black],aspect_ratio=1)
xaxis!("Re(z)")
yaxis!("Im(z)")
title!("Von Neumann Stability Analysis")
#display(plot(plot_handle))


using Test, TSVD, LinearAlgebra, SparseArrays

@testset "test tsvd with m = $m, n = $n, and p = $p" for
    (m, n, p) in ((10, 6, 0.8),
                  (6, 10, 0.8),
                  (10, 10, 0.8),
                  (100, 60, 0.1),
                  (60, 100, 0.1),
                  (100, 100, 0.1))

    mnp = round(Integer, m*n*p)

    for A in (randn(m, n),
              complex.(randn(m, n), randn(m, n)),
              sprandn(m, n, p),
              sparse(rand(1:m, mnp), rand(1:n, mnp), complex.(randn(mnp), randn(mnp)), m, n))

        Uf, sf, Vf = svd(Array(A))

        for k = 1:5
            U, s, V = TSVD.tsvd(A, k)
            @test norm(s - sf[1:k]) < sqrt(eps(real(eltype(A))))*mnp
            @test norm(abs.(U'Uf[:,1:k]) - I) < sqrt(eps(real(eltype(A))))*mnp
            @test norm(abs.(V'Vf[:,1:k]) - I) < sqrt(eps(real(eltype(A))))*mnp

            s, V = TSVD.tsvd2(A, k)
            # tmp = TSVD._teig(TSVD.AtA(A, randn(n)), k, debug = true)
            # @test norm(sqrt(reverse(tmp[1]))[1:k] - sf[1:k]) < sqrt(eps(real(eltype(A))))*mnp
            @test norm(s - sf[1:k]) < sqrt(eps(real(eltype(A))))*mnp
            @test norm(abs.(V'Vf[:,1:k]) - I) < sqrt(eps(real(eltype(A))))*mnp
        end
    end
end

@testset "Issue 9" begin
    data = rand(1:100, 50, 50)
    @test @inferred(TSVD.tsvd(data, 2))[2] ≈ svdvals(data)[1:2]
end

@testset "Issue 25. Degenerate basis" begin
  @test tsvd([1 0 ; 0 1], 1, initvec=[1.0; 0.0])[2] ≈ [1.0]
end


function entropy(partitionSet, N)
    return mapreduce(x -> -(length(x) / N) * log(length(x) / N), +, partitionSet)
end

function getPartitionSet(vector, nrOfClusters)
    d = Dict(zip(1:nrOfClusters,[Int[] for _ in 1:nrOfClusters]));
    for i in 1:length(vector)
        push!(d[vector[i]], i)
    end
    return collect(values(d))
end

function mutualInformation(partitionSetA, partitionSetB, N)

    anonFunc = (x) -> (
                intersection = size(intersect(partitionSetA[x[1]], partitionSetB[x[2]]),1);
                return intersection > 0 ?
                        (intersection/N) * log((intersection * N) / (size(partitionSetA[x[1]],1) * size(partitionSetB[x[2]],1))) : 0;
    )

    mapreduce(anonFunc, + , product(1:length(partitionSetA),1:length(partitionSetB)))
end

function normalizedMutualInformation(partitionSetA, partitionSetB, N)
    mutualInformation(partitionSetA, partitionSetB, N) / ((entropy(partitionSetA, N)  + entropy(partitionSetB, N)) / 2)
end


using Triangulation_diagrams
using PyPlot

N = [ 0  1  2  4  4  5  3  3  2  1  5  6  6  2  6  4
      2  1  0  0  2  1  1  3  2  3  3  0  4  4  2  4 ]

T = [  1  3  3  5  5
       3  5 12 12 13
      14 14  5 13 14
       9  8  6 15 16
      10  9  7 11  8
       2  7  4  6 11 ]

E = [ 14 10  1  2  3  4 12 15 13 16
      10  1  2  3  4 12 15 13 16 14 ]

QN = 6

tri = Triangulation(N, T, E, QN)

figure(1)
draw_triangles(tri)
enumerate_vertices(tri)
enumerate_triangles(tri)
enumerate_bdry_edges(tri, 0.2)
mark_first_local_vertex(tri, 0.3)
axis("off")
savefig("quad_triangulation.pdf")


# # Chebyshev center
# Boyd & Vandenberghe, "Convex Optimization"
# Joëlle Skaf - 08/16/05
#
# Adapted for Convex.jl by Karanveer Mohan and David Zeng - 26/05/14
#
# The goal is to find the largest Euclidean ball (i.e. its center and
# radius) that lies in a polyhedron described by affine inequalites in this
# fashion: $P = \{x : a_i'*x \leq b_i, i=1,\ldots,m \}$ where $x \in \mathbb{R}^2$.

using Convex
using LinearAlgebra
import SCS

# Generate the input data
a1 = [2; 1];
a2 = [2; -1];
a3 = [-1; 2];
a4 = [-1; -2];
b = ones(4, 1);

# Create and solve the model
r = Variable(1)
x_c = Variable(2)
p = maximize(r)
p.constraints += a1' * x_c + r * norm(a1, 2) <= b[1];
p.constraints += a2' * x_c + r * norm(a2, 2) <= b[2];
p.constraints += a3' * x_c + r * norm(a3, 2) <= b[3];
p.constraints += a4' * x_c + r * norm(a4, 2) <= b[4];
solve!(p, SCS.Optimizer; silent_solver = true)
p.optval

# Generate the figure
x = range(-1.5, stop = 1.5, length = 100);
theta = 0:pi/100:2*pi;
using Plots
plot(x, x -> -x * a1[1] / a1[2] + b[1] / a1[2])
plot!(x, x -> -x * a2[1] / a2[2] + b[2] / a2[2])
plot!(x, x -> -x * a3[1] / a3[2] + b[3] / a3[2])
plot!(x, x -> -x * a4[1] / a4[2] + b[4] / a4[2])
plot!(
    evaluate(x_c)[1] .+ evaluate(r) * cos.(theta),
    evaluate(x_c)[2] .+ evaluate(r) * sin.(theta),
    linewidth = 2,
)
plot!(
    title = "Largest Euclidean ball lying in a 2D polyhedron",
    legend = nothing,
)


"""
    function spearman_corr(;
        players::Array{Player},
        ranking_scheme1::String,
        ranking_scheme2::String
    )::Float64

Calculates the Spearman Correlation between two ranks.

Arguments
=========
- `players`:
    the players whose rankings to consider.
- `ranking_scheme1`:
    ranking scheme 1 to use.
- `ranking_scheme2`:
    ranking scheme 2 to use.
"""

function spearman_corr(;
    players::Array{Player},
    ranking_scheme1::String,
    ranking_scheme2::String
)::Float64
    rankings1 = map(player -> player.rankings[ranking_scheme1], players)
    rankings2 = map(player -> player.rankings[ranking_scheme2], players)
    return corspearman(rankings1, rankings2)
end

"""
    function kendtau_corr(;
        players::Array{Player},
        ranking_scheme1::String,
        ranking_scheme2::String
    )::Float64

Calculates the Kendall Tau Correlation between two ranks.

Arguments
=========
- `players`:
    the players whose rankings to consider.
- `ranking_scheme1`:
    ranking scheme 1 to use.
- `ranking_scheme2`:
    ranking scheme 2 to use.
"""

function kendtau_corr(;
    players::Array{Player},
    ranking_scheme1::String,
    ranking_scheme2::String
)::Float64
    rankings1 = map(player -> player.rankings[ranking_scheme1], players)
    rankings2 = map(player -> player.rankings[ranking_scheme2], players)
    return corkendall(rankings1, rankings2)
end


# # 220: 2D Nonlinear Poisson with boundary reaction and boundary species
# ([source code](SOURCE_URL))

module Example220_NonlinearPoisson2D_BoundarySpecies

using Printf
using VoronoiFVM
using ExtendableGrids
using GridVisualize

function main(;n=10,Plotter=nothing,verbose=false,unknown_storage=:sparse)

    h=1.0/convert(Float64,n)
    X=collect(0.0:h:1.0)
    Y=collect(0.0:h:1.0)

    grid=VoronoiFVM.Grid(X,Y)

    k=1.0
    eps=1.0
    physics=VoronoiFVM.Physics(
    num_species=3,
    breaction=function(f,u,node)
        if  node.region==2
            f[1]=k*(u[1]-u[3])
            f[3]=k*(u[3]-u[1])+ k*(u[3]-u[2])
            f[2]=k*(u[2]-u[3])
        end
    end,

    bstorage=function(f,u,node)
        if  node.region==2
            f[3]=u[3]
        end
    end,

    flux=function(f,u0,edge)
        u=unknowns(edge,u0)
        f[1]=eps*(u[1,1]-u[1,2])
        f[2]=eps*(u[2,1]-u[2,2])
    end,

    source=function(f,node)
        x1=node[1]-0.5
        x2=node[2]-0.5
        f[1]=exp(-20.0*(x1^2+x2^2))
    end,

    storage=function(f,u,node)
        f[1]=u[1]
        f[2]=u[2]
    end
    )

    sys=VoronoiFVM.System(grid,physics,unknown_storage=unknown_storage)

    enable_species!(sys,1,[1])
    enable_species!(sys,2,[1])
    enable_boundary_species!(sys,3,[2])

    function tran32!(a,b)
        a[1]=b[2]
    end

    bgrid2=subgrid(grid,[2],boundary=true,transform=tran32!)

    inival=unknowns(sys)
    inival.=0.0
    U=unknowns(sys)

    eps=1.0e-2

    control=VoronoiFVM.NewtonControl()
    control.verbose=verbose
    control.tol_linear=1.0e-5
    control.tol_relative=1.0e-5
    control.max_lureuse=0
    tstep=0.01
    time=0.0
    istep=0
    u5=0
    p=GridVisualizer(Plotter=Plotter,layout=(3,1))
    while time<1
        time=time+tstep
        solve!(U,inival,sys,control=control,tstep=tstep)
        inival.=U
        if verbose
            @printf("time=%g\n",time)
        end
        tstep*=1.0
        istep=istep+1
        U_bound=view(U[3,:],bgrid2)
        u5=U_bound[5]
        scalarplot!(p[1,1],grid,U[1,:],clear=true)
        scalarplot!(p[2,1],grid,U[2,:])
        scalarplot!(p[3,1],bgrid2,U_bound,show=true,flimits=(0,0.0025))
    end
    return u5
end

function test()
    main(unknown_storage=:sparse) ≈ 0.0020781361856598
    main(unknown_storage=:dense) ≈ 0.0020781361856598
end
end


@doc raw"""
    Segment2

An object `s` of the data type [`Segment2`](@ref) is a directed straight line
segment in the two-dimensional Euclidean plane ``Ε^2``, i.e. a straight line
segment ``[p, q]`` connecting two points ``p, q ∈ ℝ^2``.

The segment is topologically closed, i.e. the end points belong to it. Point `p`
is called the *source* and `q` is called the *target* of `s`. The length of `s`
is the Euclidean distance between `p` and `q`. Note that there is only a
function to compute the square of the length, because otherwise we had to
perform a square root operation which is expensive, and may not be exact.
"""
Segment2

"""
    Segment2(p::Point2, q::Point2)

Introduces a segment `s` with source `p` and target `q`.

The segment is directed from the source towards the target.
"""
Segment2(p::Point2, q::Point2)

"""
    ==(s::Segment2, q::Segment2)

Test for equality: Two segments are equal, iff their sources and targets are
equal.
"""
==(s::Segment2, q::Segment2)

"""
    source(s::Segment2)

The source of `s`.
"""
source(s::Segment2)

"""
    target(s::Segment2)

The target of `s`.
"""
target(s::Segment2)

"""
    min(s::Segment2)

Returns the point of `s` with lexicographically smallest coordinate.
"""
min(s::Segment2)

"""
    max(s::Segment2)

Returns the point of `s` with lexicographically largest coordinate.
"""
max(s::Segment2)

@doc raw"""
    vertex(s::Segment2, i::Integer)

Returns source of target or `s`: `vertex(s, 0)` returns the source of `s`,
`vertex(s, 1)` returns the target of `s`

The parameter `i` is taken module 2, which gives easy access to the other
vertex.
"""
vertex(s::Segment2, i::Integer)

"""
    point(s::Segment2, i::Integer)

Returns `vertex(s, i)`.
"""
point(s::Segment2, i::Integer)

"""
    squared_length(s::Segment2)

Returns the squared lenght of `s`.
"""
squared_length(s::Segment2)

"""
    direction(s::Segment2)

Returns the direction from source to target of `s`.
"""
direction(s::Segment2)

"""
    to_vector(s::Segment2)

Returns the vector `target(s) - source(s)`.
"""
to_vector(s::Segment2)

"""
    opposite(s::Segment2)

Returns a segment with source and target point interchanged.
"""
opposite(s::Segment2)

"""
    supporting_line(s::Segment2)

Returns the line `l` passing through `s`.
"""
supporting_line(s::Segment2)

"""
    is_degenerate(s::Segment2)

Segment `s` is degenerate, if source and target are equal.
"""
is_degenerate(s::Segment2)

"""
    is_horizontal(s::Segment2)

Segment `s` is horizontal, if both the source and target's ``y``-coordinate is
the same.
"""
is_horizontal(s::Segment2)

"""
    is_vertical(s::Segment2)

Segment `s` is vertical, if both the source and target's ``x``-coordinate is
the same.
"""
is_vertical(s::Segment2)

"""
    has_on(s::Segment2, p::Point2)

A point is on `s`, iff it is equal to the source or target of `s` or if it is
in the interior of `s`.
"""
has_on(s::Segment2, p::Point2)

"""
    collinear_has_on(s::Segment2, p::Point2)

Checks if point `p` is on segment `s`.

!!! info "Precondition"

    `p` is on the supporting line of `s`.
"""
collinear_has_on(s::Segment2, p::Point2)

"""
    bbox(s::Segment2)

Returns a bounding box containing `s`.
"""
bbox(s::Segment2)

"""
    transform(s::Segment2, t::AffTransformation2)

Returns the segment obtained by applying `t` on the source and the target of
`s`.
"""
transform(s::Segment2, t::AffTransformation2)


using UtilitiesForMRI, Random

# Cartesian domain
n = (128,128,128)
h = (1f0, 1f0, 1f0)
X = spatial_sampling(Float32, n; h=h)

# Cartesian sampling (dense)
phase_encoding = (1,2)
K = kspace_Cartesian_sampling(X; phase_encoding=phase_encoding)
# using PyPlot
# figure()
# plot3D(K[1][:,1],  K[1][:,2],  K[1][:,3], ".")
# plot3D(K[10][:,1], K[10][:,2], K[10][:,3], ".")
# plot3D(K[20][:,1], K[20][:,2], K[20][:,3], ".")
# plot3D(K[128*(10-1)+1][:,1], K[128*(10-1)+1][:,2], K[128*(10-1)+1][:,3], ".")
# plot3D(K[128*(20-1)+1][:,1], K[128*(20-1)+1][:,2], K[128*(20-1)+1][:,3], ".")
# plot3D(K[128^2][:,1], K[128^2][:,2], K[128^2][:,3], ".")

# Cartesian sampling (randomized)
subsampling = (1:128^2)[randperm(128^2)][1:32]
K = kspace_Cartesian_sampling(X; phase_encoding=(1,3), subsampling=subsampling)
# figure()
# for i = 1:32
#     plot3D(K[i][:,1], K[i][:,2], K[i][:,3], ".")
# end

# for i = 1:50:128^2
#     plot3D(K[i][:,1], K[i][:,2], K[i][:,3], ".")
# end