# Logistic Regression Module with CUDA support
# Logistic Regression Module
module LogRegGPU
export train, predict, predict_proba, scale, cost, accuracy

using CUDA
import Random
import Statistics

CUDA.allowscalar(false)

"""
Helper function to scale input `X` to unit variance\\
`X` is required to have dimension 2
"""
function scale(X::Array)::Array
    @assert ndims(X) == 2
    u = Statistics.mean(X, dims=1)
    s = Statistics.std(X, dims=1)
    res = (X .- u) ./ s
    return res
end

"""
Sigmoid function on array `Z`
"""
function sigmoid(Z::CuArray)::CuArray
    denom = 1 .+ (Float32(MathConstants.e) .^ (-Z))
    return 1 ./ denom
end

"""
Cost function for logistic Regression\\
If `X` is shape (M, N)\\
`y` should be (M,)\\
`beta` can be (N,) or (N+1,)
------
May not be used
"""
function cost(X::CuArray, y::CuArray, beta::CuArray)::AbstractFloat
    @assert ndims(X) == 2
    @assert ndims(y) == 1
    @assert ndims(beta) == 1
    @assert size(X)[1] == size(y)[1]
    @assert size(X)[2] == size(beta)[1] || size(X)[2] == size(beta)[1] - 1
    if size(X)[2] < size(beta)[1]
        X = hcat(X, ones(size(X)[1]))
    end
    m = size(X)[1]
    X_combined = X * beta
    prob = sigmoid(X_combined)
    vec = y .* (log.(prob)) .+ (1 .- y) .* (log.(1 .- prob))
    cost = -1 / Float32(m) * sum(vec)
    return cost
end

function cost(X::Array, y::Array, beta::Array)::AbstractFloat
    @assert ndims(X) == 2
    @assert ndims(y) == 1
    @assert ndims(beta) == 1
    @assert size(X)[1] == size(y)[1]
    @assert size(X)[2] == size(beta)[1] || size(X)[2] == size(beta)[1] - 1
    if size(X)[2] < size(beta)[1]
        X = hcat(X, ones(size(X)[1]))
    end
    X = CUDA.CuArray(Float32.(X))
    y = CUDA.CuArray(Float32.(y))
    beta = CUDA.CuArray(Float32.(beta))
    m = size(X)[1]
    X_combined = X * beta
    prob = sigmoid(X_combined)
    vec = y .* (log.(prob)) .+ (1 .- y) .* (log.(1 .- prob))
    cost = -1 / float(m) * sum(vec)
    CUDA.reclaim()
    return cost
end

"""
Predict function for Logistic Regression\\
If `X` is shape (M, N)\\
`beta` can be (N,) or (N+1,)\\
Returns 1d array of real probabilities
"""
function predict_proba(X::CuArray, beta::CuArray)::CuArray
    @assert ndims(X) == 2
    @assert ndims(beta) == 1
    @assert size(X)[2] == size(beta)[1] || size(X)[2] == size(beta)[1] - 1
    if size(X)[2] < size(beta)[1]
        X = hcat(X, ones(size(X)[1]))
    end
    X_combined = X * beta
    prob = sigmoid(X_combined)
    return prob
end

function predict_proba(X::Array, beta::Array)::Array
    @assert ndims(X) == 2
    @assert ndims(beta) == 1
    @assert size(X)[2] == size(beta)[1] || size(X)[2] == size(beta)[1] - 1
    if size(X)[2] < size(beta)[1]
        X = hcat(X, ones(size(X)[1]))
    end
    X = CUDA.CuArray(Float32.(X))
    beta = CUDA.CuArray(Float32.(beta))
    X_combined = X * beta
    prob = sigmoid(X_combined)
    prob_cpu = Array(prob)
    CUDA.reclaim()
    return prob_cpu
end

"""
Predict function for Logistic Regression\\
If `X` is shape (M, N)\\
`beta` can be (N,) or (N+1,)\\
Returns 1d array of 0,1
"""
function predict(X::CuArray, beta::CuArray)::CuArray
    @assert ndims(X) == 2
    @assert ndims(beta) == 1
    @assert size(X)[2] == size(beta)[1] || size(X)[2] == size(beta)[1] - 1
    if size(X)[2] < size(beta)[1]
        X = hcat(X, ones(size(X)[1]))
    end
    X_combined = X * beta
    prob = sigmoid(X_combined)
    real_prob::CuArray = map(m -> m >= 0.5 ? 1.0 : 0.0, prob)
    return real_prob
end

function predict(X::Array, beta::Array)::Array
    @assert ndims(X) == 2
    @assert ndims(beta) == 1
    @assert size(X)[2] == size(beta)[1] || size(X)[2] == size(beta)[1] - 1
    if size(X)[2] < size(beta)[1]
        X = hcat(X, ones(size(X)[1]))
    end
    X = CUDA.CuArray(Float32.(X))
    beta = CUDA.CuArray(Float32.(beta))
    X_combined = X * beta
    prob = sigmoid(X_combined)
    prob_cpu = Array(prob)
    real_prob::Array = map(m -> m >= 0.5 ? 1.0 : 0.0, prob_cpu)
    CUDA.reclaim()
    return real_prob
end

"""
Learning function for Gradient Descent\\
If `X` is shape (M, N)\\
`y` should be (M,)\\
`beta` should be (N,)
------
Returns `nothing`
------
Note: `beta` will be updated inplace
"""
function learn!(X::CuArray, y::CuArray, beta::CuArray, momentum::CuArray, alpha::AbstractFloat)
    @assert ndims(X) == 2
    @assert ndims(y) == 1
    @assert ndims(beta) == 1
    @assert size(X) == (size(y)[1], size(beta)[1])
    predictions = predict_proba(X, beta)
    offset = predictions .- y
    gradients = X' * offset
    gradients .= gradients ./ size(X)[1]
    gradients .= gradients .* alpha
    momentum .= gradients .+ (0.9 .* momentum)
    beta .= beta .- momentum
    return nothing
end

"""
Training function for Logisitic Regression\\
Implemented using Gradient Descent\\
If `X` is shape (M, N)\\
`y` should be (M,)\\
`max_iter` should be >= 0
------
Returns `beta` in shape (N+1,)

------
Set `early_stop` to `false`, to force run maximum iteractions\\
Set `random_weights` to `false` to initialize weights of 0.0
"""
function train(X::Array, y::Array; learning_rate::AbstractFloat=0.1, max_iter::Integer=1000,
        n_iter_no_change::Integer=5, tol::AbstractFloat=0.001, verbose::Bool=false,
        shuffle::Bool=true, early_stop::Bool=true, random_weights::Bool=true)::Array
    @assert ndims(X) == 2
    @assert ndims(y) == 1
    @assert size(X)[1] == size(y)[1]
    @assert max_iter >= 0
    @assert n_iter_no_change >= 0
    @assert tol >= 0.0
    tol = Float32(tol)
    X = hcat(X, ones(size(X)[1])) # for constant multiplication
    X .= Float32.(X)
    y = Float32.(y)
    if shuffle
        JuTools.shuffle_data!(X, y)
    end
    X = CUDA.CuArray(X)
    y = CUDA.CuArray(y)
    beta = nothing
    if random_weights
        beta = CuArray(Random.randn(size(X)[2]))
    else
        beta = CuArray(zeros(size(X)[2]))
    end
    momentum = CuArray(zeros(size(X)[2]))
    best_cost = nothing
    n_cost_no_change = n_iter_no_change
    for i = 1:max_iter
        if n_cost_no_change <= 0 && early_stop
            break
        end
        learn!(X, y, beta, momentum, learning_rate)
        new_cost = cost(X, y, beta)
        if verbose
            acc = accuracy(predict(X, beta), y)
            println("Iter: ", i)
            println("Cost = ", new_cost)
            println("Accuracy = ", acc)
            println()
        end
        if early_stop
            if best_cost === nothing || isnan(best_cost)
                best_cost = new_cost
            else
                if new_cost > best_cost - tol
                    n_cost_no_change -= 1
                else
                    best_cost = min(new_cost, best_cost)
                    n_cost_no_change = n_iter_no_change
                end
            end
        end
    end
    beta_cpu = Array(beta)
    CUDA.reclaim()
    return beta_cpu
end

"""
Accuracy function for Logisitic Regression\\
`y_pred` and `y_real` should have same shape, 1d array
------
Returns accuracy as float
"""
function accuracy(y_pred::Array, y_real::Array)::AbstractFloat
    @assert ndims(y_pred) == ndims(y_real) == 1
    @assert size(y_pred) == size(y_real)
    acc = Statistics.mean(float.(y_pred) .== float.(y_real))
    return acc
end

function accuracy(y_pred::CuArray, y_real::CuArray)::AbstractFloat
    @assert ndims(y_pred) == ndims(y_real) == 1
    @assert size(y_pred) == size(y_real)
    acc = Statistics.mean(y_pred .== y_real)
    return acc
end

end


function randOrth(d)
    q = Matrix(qr(randn(d,d)).Q)
    q = q * Diagonal(rand([1;-1],d))
    return q
end

"""
    Q = randOrth(d, pert)

Produces a random orthogonal matrix close to I.
The smaller pert is, the closer to I.
The diagonals are always positive.
"""
function randOrth(d, pert)
    q = Matrix(qr(eye(d)+pert*randn(d,d)).Q)
    q = q * Diagonal(sign.(diag(q)))
    return q
end

function randFullBDD(n,d)
    A = triu([randOrth(d) for i in 1:n, j in 1:n],1)
    for i in 1:n
        A[i,i] = eye(d)*(0.001 + (n-1)/2)
    end
    B = BlockCSC(A)
    B = B + transpose(B)
end

function randOrthWt(a,d)
    #deg = sum(a,dims=1)
    n = size(a,1)
    ai,aj,av = findnz(triu(a))
    vq = [randOrth(d)*av[i] for i in 1:length(av)]

    ai = vcat(ai,[n])
    aj = vcat(aj,[n])
    vq = vcat(vq,[zeros(d,d)])

    B = BlockCSC(sparse([ai;aj],[aj;ai],[vq;collect.(transpose.(vq))],n,n))
    #B = B + transpose(B)

    E = one(typeof(B[1,1]))

    sp = BlockCSC(sparse(1:n,1:n,[E*sum(opnorm.(a[:,i])) for i in 1:n]))

    return B + sp
end

"""
    B = randOrthWtPotential(a,d, pert)

Produces a bdd matrix with orthogonal matrices, slightly perturbed from a potential.
If the potential is zero, then the resulting matrix is singular.
Note that if there are very few edges, then the matrix can be singular anyway.
This motivates randBDDWtPotential, which is never singular.
"""
function randOrthWtPotential(a,d, pert)
    #deg = sum(a,dims=1)
    n = size(a,1)
    ai,aj,av = findnz(triu(a))

    vr = [randOrth(d) for i in 1:n]

    vq = [-vr[ai[i]]*vr[aj[i]]'*randOrth(d,pert)*av[i] for i in 1:length(av)]

    ai = vcat(ai,[n])
    aj = vcat(aj,[n])
    vq = vcat(vq,[zeros(d,d)])

    B = BlockCSC(sparse(ai,aj,vq,n,n))
    B = B + transpose(B)

    dg = [eye(d)*sum(opnorm.(B[:,i])) for i in 1:n]
    B = B + BlockCSC(sparse(Diagonal(dg)))
end

function randBDDWtPotential(a,d, pert)
    #deg = sum(a,dims=1)
    n = size(a,1)
    ai,aj,av = findnz(triu(a))

    vr = [randOrth(d) for i in 1:n]

    vq = [(-vr[ai[i]]*vr[aj[i]]' + randn(d,d)*pert)*av[i] for i in 1:length(av)]

    ai = vcat(ai,[n])
    aj = vcat(aj,[n])
    vq = vcat(vq,[zeros(d,d)])

    B = BlockCSC(sparse(ai,aj,vq,n,n))
    B = B + transpose(B)

    dg = [eye(d)*sum(opnorm.(B[:,i])) for i in 1:n]
    B = B + BlockCSC(sparse(Diagonal(dg)))
end

function randBDDWt(a,d)
    n = size(a,1)
    ai,aj,av = findnz(triu(a))
    vq = [randn(d,d) for i in 1:length(av)]
    ai = vcat(ai,[n])
    aj = vcat(aj,[n])
    vq = vcat(vq,[zeros(d,d)])
    B = BlockCSC(sparse(ai,aj,vq,n,n))
    B = B + transpose(B)

    dg = [eye(d)*sum(opnorm.(B[:,i])) for i in 1:n]
    B = B + BlockCSC(sparse(Diagonal(dg)))
end


## Show the effect of seeded PageRank on MNIST

## Load the MNIST Graph and draw it
##
using DelimitedFiles
using SparseArrays
using LinearAlgebra
ei,ej = copy.(eachcol(Int.(readdlm("mnist-train-4.edges"))))
xy = readdlm("mnist-train-4.xy")
n = maximum(ei)
A = sparse(ei,ej,1,n,n)
@assert issymmetric(A)

##
## Simple PageRank
function simplepagerank(A,α,v)
  @assert(0 ≤ α < 1, "needs probably α")
  @assert(all(vi -> vi ≥ 0, v), "needs non-negative v")
  v = v ./ sum(v) # we can normalize for them.
  d = vec(sum(A,dims=2)) # compute the degrees
  x = copy(v) # start of with v
  nsteps = 2*ceil(Int,log(eps(1.0))/log(α)) # upper bound on steps
  for i=1:nsteps
    x = α*(A'*(x./d)) .+ (1-α).*v
  end
  return x/sum(x) # renormalize to probability
end

## Let's show global pagerank first
pr = simplepagerank(A,0.85,ones(n)./n)
## Now show the entire graph
using GraphRecipes
using LinearAlgebra
using Plots

ei, ej = findnz(A)
graphplot(ei, ej, x =xy[:,1], y=xy[:,2],
  markercolor=:black, markerstrokecolor=:white,
  size=(1200,1200),dpi=300,
  markersize=0, linecolor=1, linealpha=0.2, linewidth=0.5,
  markeralpha=0.2,colorbar=false,
  axis_buffer=0.02, background=nothing)
p=sortperm(pr)
scatter!(xy[p,1],xy[p,2],marker_z=-abs.(log10.(pr[p])).^0.5,markerstrokewidth=0,alpha=0.5)
savefig("pagerank-mnist.png")
## now show Seeded PageRank
v = zeros(n)
v[1] = 1
pr = simplepagerank(A,0.85,v)

using GraphRecipes
using LinearAlgebra
using Plots

ei, ej = findnz(A)
graphplot(ei, ej, x =xy[:,1], y=xy[:,2],
  markercolor=:black, markerstrokecolor=:white,
  size=(1200,1200),dpi=300,
  markersize=0, linecolor=1, linealpha=0.2, linewidth=0.5,
  markeralpha=0.2,colorbar=false,
  axis_buffer=0.02, background=nothing)
p=sortperm(pr)
scatter!(xy[p,1],xy[p,2],marker_z=-abs.(log10.(pr[p])).^0.5,markerstrokewidth=0,alpha=0.5)
savefig("seeded-pagerank-mnist.png")

##
## Try the same set of nodes labeled 1 as we used for the least squares exampmle.
using MLDatasets
train_x, train_y = MNIST.traindata()
nlabels = 10
S = findall(train_y .== 1)[1:nlabels]
##
v = zeros(n)
v[S] .= 1/length(S)
pr = simplepagerank(A,0.85,v)
ei, ej = findnz(A)
graphplot(ei, ej, x =xy[:,1], y=xy[:,2],
  markercolor=:black, markerstrokecolor=:white,
  size=(1200,1200),dpi=300,
  markersize=0, linecolor=1, linealpha=0.2, linewidth=0.5,
  markeralpha=0.2,colorbar=false,
  axis_buffer=0.02, background=nothing)
p=sortperm(pr)
scatter!(xy[p,1],xy[p,2],marker_z=-abs.(log10.(pr[p])).^0.5,markerstrokewidth=0,alpha=0.5)
savefig("seeded-pagerank-mnist-set-1.png")

##
v = zeros(n)
v[S] .= 1/length(S)
pr = simplepagerank(A,0.5,v)
ei, ej = findnz(A)
graphplot(ei, ej, x =xy[:,1], y=xy[:,2],
  markercolor=:black, markerstrokecolor=:white,
  size=(1200,1200),dpi=300,
  markersize=0, linecolor=1, linealpha=0.2, linewidth=0.5,
  markeralpha=0.2,colorbar=false,
  axis_buffer=0.02, background=nothing)
p=sortperm(pr)
scatter!(xy[p,1],xy[p,2],marker_z=-abs.(log10.(pr[p])).^0.5,markerstrokewidth=0,alpha=0.5)
savefig("seeded-pagerank-mnist-set-1-alpha-half.png")

##
v = zeros(n)
v[S] .= 1/length(S)
pr = simplepagerank(A,0.99,v)
ei, ej = findnz(A)
graphplot(ei, ej, x =xy[:,1], y=xy[:,2],
  markercolor=:black, markerstrokecolor=:white,
  size=(1200,1200),dpi=300,
  markersize=0, linecolor=1, linealpha=0.2, linewidth=0.5,
  markeralpha=0.2,colorbar=false,
  axis_buffer=0.02, background=nothing)
p=sortperm(pr)
scatter!(xy[p,1],xy[p,2],marker_z=-abs.(log10.(pr[p])).^0.5,markerstrokewidth=0,alpha=0.5)
savefig("seeded-pagerank-mnist-set-1-alpha-99.png")


"""
High-order singular value decomposition (HO-SVD).
"""
function hosvd(tensor::AbstractArray{T,N}, core_dims::NTuple{N, Int};
               pad_zeros::Bool=false, compute_error::Bool=false) where {T,N}
    pad_zeros || _check_tensor(tensor, core_dims)

    factors = map(1:N) do i
        X = _col_unfold(tensor, i)
        f = eigen(Symmetric(X'X), max(1, size(X,2)-core_dims[i]+1):size(X,2)).vectors
        if pad_zeros && size(f, 2) < core_dims[i] # fill missing factors with zeros
            f = hcat(f, zeros(T, size(tensor, i), core_dims[i]-size(f, 2)))
        end
        mapslices(_check_sign, f, dims=1)
    end

    res = Tucker(factors, tensorcontractmatrices(tensor, factors))
    compute_error && _set_rel_residue(res, tensor)
    return res
end

hosvd(tensor::AbstractArray, r::Int;
      pad_zeros::Bool=false, compute_error::Bool=false) =
    hosvd(tensor, ntuple(_ -> r, ndims(tensor)),
          pad_zeros=pad_zeros, compute_error=compute_error)


using QuantEcon

"""
The data for a consumption problem, including some default values.
"""
struct ConsumptionProblem{TF<:AbstractFloat}
    beta::TF
    y::Vector{TF}
    b0::TF
    P::Matrix{TF}
end

"""
Parameters
----------

beta : discount factor
P    : 2x2 transition matrix
y    : Array containing the two income levels
b0   : debt in period 0 (= state_1 debt level)
"""
function ConsumptionProblem(;
                 beta = 0.96,
                 y = [2.0, 1.5],
                 b0 = 3.0,
                 P = [0.8 0.2;
                      0.4 0.6])

    ConsumptionProblem(beta, y, b0, P)
end

"""
Computes endogenous values for the complete market case.

Parameters
----------

cp : instance of ConsumptionProblem

Returns
-------

    c_bar : constant consumption
    b1    : rolled over b0
    b2    : debt in state_2

associated with the price system

    Q = beta * P

"""

function consumption_complete(cp::ConsumptionProblem)

    beta, P, y, b0 = cp.beta, cp.P, cp.y, cp.b0  # Unpack

    y1, y2 = y          # extract income levels
    b1 = b0             # b1 is known to be equal to b0
    Q = beta * P        # assumed price system

    # Using equation (7) calculate b2
    b2 = (y2 - y1 - (Q[1, 1] - Q[2, 1] - 1) * b1)/(Q[1, 2] + 1 - Q[2, 2])

    # Using equation (5) calculae c_bar
    c_bar = y1 - b0 + ([b1 b2] * Q[1, :] )[1]

    return c_bar, b1, b2
end

"""
Computes endogenous values for the incomplete market case.

Parameters
----------

cp : instance of ConsumptionProblem
N_simul : Integer

"""
function consumption_incomplete(cp::ConsumptionProblem;
                                N_simul::Integer=150)

    beta, P, y, b0 = cp.beta, cp.P, cp.y, cp.b0  # Unpack
    # For the simulation define a quantecon MC class
    mc = MarkovChain(P)

    # Useful variables
    y = y''
    v = inv(eye(2) - beta * P) * y

    # Simulat state path
    s_path = simulate(mc, N_simul, init = 1)

    # Store consumption and debt path
    b_path, c_path = ones(N_simul + 1), ones(N_simul)
    b_path[1] = b0

    # Optimal decisions from (12) and (13)
    db = ((1 - beta) * v - y) / beta

    for (i, s) in enumerate(s_path)
        c_path[i] = (1 - beta) * (v - b_path[i] * ones(2, 1))[s, 1]
        b_path[i + 1] = b_path[i] + db[s, 1]
    end

    return c_path, b_path[1:end-1], y[s_path], s_path
end


### A Pluto.jl notebook ###
# v0.14.5

using Markdown
using InteractiveUtils

# ╔═╡ 6e58a278-ede3-440f-b5ad-c23d03833918
using LinearAlgebra

# ╔═╡ 81465bdf-4eb3-4dc5-8f60-2ccdca956e73
md"#### Sigmoid"

# ╔═╡ 78b943f0-ae7b-11eb-3de1-d5cfb31433c6
function σ(x)
	t = exp(-abs(x))
	return ifelse(x ≥ 0, inv(1 + t), t / (1 + t))
end

# ╔═╡ 6fe8d0a2-fd10-4097-95c8-fea57d2a32e8
const sigmoid = σ

# ╔═╡ 112f4771-632e-4c35-81bd-ec7f09b40dc0
md"#### ReLU"

# ╔═╡ 257e3608-759e-49b2-aa9f-88211bba198b
relu(x) = max(zero(x), x)

# ╔═╡ 5b4291e3-1f67-46ac-89f5-d00aaefdac8a
sigmoid(4) == σ(4)

# ╔═╡ 45ba9913-5316-457e-9f97-2622661397e1
begin
	inputs = [1, 2, 3, 2.5]
	weights = [0.2, 0.8, -0.5, 1.0]
	bias = 2.0

	(
		σ(weights ⋅ inputs + bias),
		relu(weights ⋅ inputs + bias)
	)
end

# ╔═╡ e2f8442b-5ccd-431a-8881-3e746b0eb047
begin
	#inputs = [1, 2, 3, 2.5]
	weights_matrix = [
		 0.2   0.8  -0.5   1.0
		 0.5  -0.91  0.26 -0.5
		-0.26 -0.27  0.17  0.87
	]

	bias_vec = [2.0, 3.0, 0.5]

	(
		σ.(weights_matrix * inputs + bias_vec),
		relu.(weights_matrix * inputs + bias_vec)
	)
end

# ╔═╡ 0e4a6376-d3fa-47fc-b140-ddb65ee0c592
md"#### Softmax"

# ╔═╡ afbe3c9c-6db4-415c-921d-cdc4033e9b06
function softmax(x::Vector{T}) where {T<:Real}
	m = maximum(x)
	exp_val = exp.(x .- m)
	s = sum(exp_val)
	exp_val ./ s
end

# ╔═╡ 662251f6-03ae-4ea6-8bd1-e90ec8a10805
function softmax(x::Matrix{T}) where {T<:Real}
	m = maximum.(eachcol(x))
	exp_val = exp.(x .- m')
	s = sum(eachrow(exp_val))
	return exp_val ./ s'
end

# ╔═╡ 383304f8-eb30-45be-99c4-542132ecc78a
begin
	#inputs_batch = [
	#	[1, 2, 3, 2.5] [2, 5, -1, 2] [-1.5, 2.7, 3.3, -0.8]
	#]
	inputs_batch = [
		1 2 3 2.5
		2 5 -1 2
		-1.5 2.7 3.3 -0.8
	]

	(
		softmax(weights_matrix * inputs + bias_vec),
		softmax(weights_matrix * inputs_batch' .+ bias_vec) # orientacao coluna
	)
end

# ╔═╡ Cell order:
# ╠═6e58a278-ede3-440f-b5ad-c23d03833918
# ╟─81465bdf-4eb3-4dc5-8f60-2ccdca956e73
# ╠═78b943f0-ae7b-11eb-3de1-d5cfb31433c6
# ╠═6fe8d0a2-fd10-4097-95c8-fea57d2a32e8
# ╟─112f4771-632e-4c35-81bd-ec7f09b40dc0
# ╠═257e3608-759e-49b2-aa9f-88211bba198b
# ╠═5b4291e3-1f67-46ac-89f5-d00aaefdac8a
# ╠═45ba9913-5316-457e-9f97-2622661397e1
# ╠═e2f8442b-5ccd-431a-8881-3e746b0eb047
# ╟─0e4a6376-d3fa-47fc-b140-ddb65ee0c592
# ╠═afbe3c9c-6db4-415c-921d-cdc4033e9b06
# ╠═662251f6-03ae-4ea6-8bd1-e90ec8a10805
# ╠═383304f8-eb30-45be-99c4-542132ecc78a


using LinearAlgebra

# Linear Combination
a=[1,2]; b=[3,4];
alpha = -0.5; beta = 1.5;
c = alpha*a + beta*b

# Define function
function lincomb(coeff, vectors)
    n = length(vectors[1])
    a = zeros(n);
    for i = 1:length(vectors)
        a = a + coeff[i] * vectors[i];
    end
return a
end

#Alternate definition
function lincomb(coeff, vectors)
    return sum( coeff[i] * vectors[i]
        for i = 1:length(vectors) )
    end

# Use funtion
lincomb( ( -0.5, 1.5), ( [1, 2], [ 3, 4]) )


using LazIO
using LasIO
using FileIO
using StaticArrays
using LinearAlgebra
using BenchmarkTools
using DataFrames
using PointCloudRasterizers

"Store coordinates"
function coordinates(points,h)
    coords = SArray{Tuple{9},Float64,1,9}[]
    for (ind,p) in enumerate(points)
        x = xcoord(p,h)
        y = ycoord(p,h)
        z = zcoord(p,h)
        inten = intensity(p)
        ret = return_number(p)
        nret = number_of_returns(p)
        classificc = classification(p)
        gps = gps_time(p)
        push!(coords,(ind, x, y, z, inten, ret, nret, classificc, gps))
    end
    return coords
end

"Convert Cartesian to Spherical coordinates"
function cart2pol(x::Float64, y::Float64, z::Float64)
    cart = Tuple{Float64,Float64,Float64}
    ρ = sqrt(x^2 + y^2 + z^2) #distance
    θ = atand(y,x) #horizontal angle
    ϕ = acosd(z/ρ) #incidence angle
    cart = (ρ,θ,ϕ)
    return cart
end

"Convert Spherical to Cartesian coordinates"
function pol2cart(ρ::Float64, θ::Float64, ϕ::Float64)
    pol = Tuple{Float64,Float64,Float64}
    x = (ρ * sind(ϕ) * cosd(θ))
    y = (ρ * sind(ϕ) * sind(θ))
    z = (ρ * cosd(ϕ))
    return pol = (x,y,z)
end

const c_air, c_water = 2.99*10^8, 2.25*10^8
const n_air_water = c_water/c_air

"Correct an underwater point according to a watersurface point from the refraction factor by using Snell's Law."
function refraction_correction(watersurface::SArray{Tuple{4},Float64,1,4},underwater::SArray{Tuple{4},Float64,1,4})
    x0, y0, z0 = watersurface[2:4]
    x1, y1, z1 = underwater[2:4]
    Δx, Δy, Δz = (x1-x0), (y1-y0), (z1-z0)
    ind = underwater[1]
    #Calculate the distance of the two points (ρ), the vertical-incidence angle (ϕ) and horizontal angle (θ) between the two points
    ρ,θ,ϕ = cart2pol(Δx, Δy, Δz)

    #Calculate the underwater angle due to the refraction effect.
    #Using Snell's Law: n(air)*sin(ϕ_air) = n(water)*sin(ϕ_water),
    #where ϕ_air: the angle of incidence (I know that the angle of the scan is 20)
    #      ϕ_water: the angle of refraction
    #      n(air): the refraction idex of medium containing the incident ray, value=1.000293
    #      n(water): the refraction idex of medium containing the transmitted ray, value=1.333

    #check the height values of the water-surface and underwater point
    if z0 == z1
        error("the water-surface and underwater point are at the same height level")
    elseif z0 > z1
        ϕ_water = asind(n_air_water*sind(ϕ)) #ϕ_new is the new vertical
    else
        error("the underwater point can not be higher than the watersurface point")
    end

    #Distance under the water has been influenced by the different speed of light in the water.
    #In a specific moment, ρ = c_air * t and ρ_water = c_water * t.
    #So, by diving these two equations, the ratio (ρ1/ρ2) = n_air_water
    ρ_water = (ρ*n_air_water)

    #Calculate the correction in the 3d space using the pol2cart function
    xcor, ycor, zcor = pol2cart(ρ_water, θ, ϕ_water)

    #Define the new corrected coordinates of the points
    xnew = x0 + xcor
    ynew = y0 + ycor
    znew = z0 - zcor

    #Return the new corrected point
    corrections = @SVector[xcor, ycor, zcor]
    correct_point = @SVector[ind, xnew, ynew ,znew]

    return corrections, correct_point
end

"Group the laser pulses based on their GPS time"
function laserpulses(points)
    println("Execution Time of laser pulses " * string(length(points)) * " points")
    @time begin
        laser_pulses = SArray{Tuple{}}[] #specify the type of array
        sorted_gps = sort!(points, by = x -> x[9], rev = false)
        pp = [tuple(p[1],p[2],p[3],p[4],p[5],p[6],p[7],p[8],p[9]) for p in sorted_gps]
        df = DataFrame(pp[:]) #create a dataframe
        groups = groupby(df,:9) #group by gps time
        for gr in groups
            if size(gr,1) > 1
                push!(laser_pulses,[gr])
            end
        end
    end
    return laser_pulses
end

"Calculate the corrected/refracted points for every pulse"
function corrected_points(laser_pulses,coords)
    corrected_laser_pulses = [] #specify the type of array

    INDEXES_corr = [] #indexes of corrected_points
    INDEXES_un_corr = [] #indexes of non-corrected points

    for every_pulse in laser_pulses
        indexes = getindex(every_pulse)[1]
        ind_surf = Int64(indexes[1]) #water surface index
        p_surf = coords[ind_surf] #water surface point
        p_surf = @SVector[ind_surf,p_surf[2],p_surf[3],p_surf[4]] #create Arrays
        corrected_points = SArray{Tuple{4},Float64,1,4}[] #store the corrected points for each pulse!
        not_corrected = SArray{Tuple{4},Float64,1,4}[] #store the un-corrected points due to the Exception ERRORsfor each pulse!

        push!(corrected_points,(p_surf)) # add the water surface point
        push!(not_corrected,(p_surf)) # add the water surface point

        for ind_under = 2:length(indexes) #iterate throught the inder water points
            n = Int64(indexes[ind_under])
            p_under = coords[n]
            p_under = @SVector[n,p_under[2],p_under[3],p_under[4]]
            try
                correction, corre_pt = refraction_correction(p_surf,p_under)
                push!(INDEXES_corr,p_under[1])
                push!(corrected_points,(corre_pt)) # I will classify as un-refracted with classif.code = 19
            catch
                push!(INDEXES_un_corr,p_under[1])
                push!(not_corrected,(p_under)) # I will classify as un-refracted with classif.code = 20
            end
        end

        if length(corrected_points) > 1 #check if there is NOT only water surface point
            push!(corrected_laser_pulses,[corrected_points])
        end
    end
    return corrected_laser_pulses, INDEXES_corr, INDEXES_un_corr
end

#Write a new classified Pointcloud
function write_pointcloud(ds,header,pp,indexes_corr,index_un_corr)
    n = length(indexes_corr) + length(index_un_corr)
    println("Execution Time of writing points " * string(n) * " points")
    @time  begin
        laz_out = joinpath(path, "myoutput_" * string(length(pp)) * "_points.laz")
        LazIO.write(laz_out, ds.header) do io
            for (ind,p) in enumerate(ds)
                if ind in indexes_corr
                    p.classification = UInt8(19)
                elseif ind in indexes_un_corr #this elseif doesn't work!!why???
                    p.classification = UInt8(20)
                else
                    p.classification = UInt8(9)
                end
                LazIO.writepoint(io, p)
            end
        end
    end
end

#MAIN#
const path = "C:/Users/alexandr/OneDrive - Stichting Deltares/Desktop/NL3_subset/"
const filename = "NL3_clipped_water_cropheight_pointformat3_classification.laz" #classification = 9 (water)
const lazinput = path * filename

#Open LAZ pointcloud
open_file = File{format"LAZ_"}(lazinput)
h, points = load(open_file)
n = length(points)

#dataset
ds = LazIO.open(lazinput)

#subset
pp = points[1500000:1520000]
#Array with coordinates
coords = coordinates(pp,h)
#Separate laser pulses
laser_pulses = laserpulses(coords)
#Return corrected points
corrected_laser_pulses, indexes_corr, indexes_un_corr  = corrected_points(laser_pulses,coords)
#write pointcloud
write_pointcloud(ds,h,pp,indexes_corr,indexes_un_corr)


using Pkg
Pkg.activate(@__DIR__)
Pkg.instantiate()

using LinearAlgebra, Symbolics, DifferentialEquations, JLD2

struct DoublePendulum{T}
	m1::T
	m2::T
	l1::T
	l2::T
end

n = 2 # number of states
model = DoublePendulum(1.0, 1.0, 1.0, 1.0) # model

# kinematics
function kinematics_1(model::DoublePendulum, q)
	θ1, θ2 = q

	[0.5 * model.l1 * sin(θ1);
	 -0.5 * model.l1 * cos(θ1)]
end

function kinematics_2(model::DoublePendulum, q)
	θ1, θ2 = q

	[model.l1 * sin(θ1) + 0.5 * model.l2 * sin(θ1 + θ2);
	 -model.l1 * cos(θ1) - 0.5 * model.l2 * cos(θ1 + θ2)]
end

# fast kinematics functions
@variables q[1:n]
@variables q̇[1:n]

k1 = kinematics_1(model, q)
k2 = kinematics_2(model, q)

k1_exp = Symbolics.build_function(k1, q)
k2_exp = Symbolics.build_function(k2, q)

k1_func = eval(k1_exp[1])
k2_func = eval(k2_exp[1])

# kinematics Jacobians
j1 = Symbolics.jacobian(k1, q, simplify = true)
j2 = Symbolics.jacobian(k2, q, simplify = true)

j1_exp = Symbolics.build_function(j1, q)
j2_exp = Symbolics.build_function(j2, q)

j1_func = eval(j1_exp[1])
j2_func = eval(j2_exp[1])

# Lagrangian
function lagrangian(model, q, q̇)
	L = 0.0

	# mass 1
	v1 = j1_func(q) * q̇
	L += 0.5 * model.m1 * transpose(v1) * v1 		# kinetic energy
	L -= model.m1 * 9.81 * k1_func(q)[2]            # potential energy

	# mass 2
	v2 = j2_func(q) * q̇
	L += 0.5 * model.m2 * transpose(v2) * v2
	L -= model.m2 * 9.81 * k2_func(q)[2]

	return L
end

# fast Lagrangian
L = lagrangian(model, q, q̇)

dLq = Symbolics.gradient(L, q, simplify = true)
dLq̇ = Symbolics.gradient(L, q̇, simplify = true)
ddL = Symbolics.hessian(L, [q; q̇], simplify = true)

# mass matrix
M = ddL[n .+ (1:n), n .+ (1:n)]
M = simplify.(M)

# dynamics bias
C = ddL[n .+ (1:n), 1:n] * q̇ - dLq
C = simplify.(C)

# dynamics
# ẋ = [q̇; M \ (-1.0 * C)]
ẋ = [q̇; M \ (-0.5 * q̇ -1.0 * C)]
ẋ = simplify.(ẋ)

ẋ_exp = Symbolics.build_function(ẋ, q, q̇)
dynamics = eval(ẋ_exp[1])

# save dynamics function
path = joinpath(pwd(), "dynamics.jld2")
# @save path ẋ_exp
# @load path ẋ_exp

# DifferentialEquations.jl
function dynamics!(ẋ, x, p, t)
	ẋ .= dynamics(view(x, 1:n), view(x, n .+ (1:n)))
end

# simulate
x0 = [0.5 * π; 0.0; 0.0; 0.0]
tspan = (0.0, 10.0)
dt = 0.01
prob = ODEProblem(dynamics!, x0, tspan)
sol = solve(prob, Tsit5(), adaptive = false, dt = dt)

# MeshCat.jl
include(joinpath(pwd(), "visuals.jl"))
vis = Visualizer()
render(vis)
visualize_double_pendulum!(vis, model, sol.u, Δt = dt)


mutable struct UncertaintyTrapEcon{TF<:AbstractFloat, TI<:Integer}
    a::TF          # Risk aversion
    γ_x::TF        # Production shock precision
    ρ::TF          # Correlation coefficient for θ
    σ_θ::TF        # Standard dev of θ shock
    num_firms::TI  # Number of firms
    σ_F::TF        # Std dev of fixed costs
    c::TF          # External opportunity cost
    μ::TF          # Initial value for μ
    γ::TF          # Initial value for γ
    θ::TF          # Initial value for θ
    σ_x::TF        # Standard deviation of shock
end

function UncertaintyTrapEcon(;a::AbstractFloat=1.5, γ_x::AbstractFloat=0.5,
                             ρ::AbstractFloat=0.99, σ_θ::AbstractFloat=0.5,
                             num_firms::Integer=100, σ_F::AbstractFloat=1.5,
                             c::AbstractFloat=-420.0, μ_init::AbstractFloat=0.0,
                             γ_init::AbstractFloat=4.0,
                             θ_init::AbstractFloat=0.0)
    σ_x = sqrt(a / γ_x)
    UncertaintyTrapEcon(a, γ_x, ρ, σ_θ, num_firms, σ_F, c, μ_init,
                        γ_init, θ_init, σ_x)

end

function ψ(uc::UncertaintyTrapEcon, F::Real)
    temp1 = -uc.a * (uc.μ - F)
    temp2 = 0.5 * uc.a^2 * (1 / uc.γ + 1 / uc.γ_x)
    return (1 / uc.a) * (1 - exp(temp1 + temp2)) - uc.c
end

"""
Update beliefs (μ, γ) based on aggregates X and M.
"""
function update_beliefs!(uc::UncertaintyTrapEcon, X::Real, M::Real)
    # Simplify names
    γ_x, ρ, σ_θ = uc.γ_x, uc.ρ, uc.σ_θ

    # Update μ
    temp1 = ρ * (uc.γ * uc.μ + M * γ_x * X)
    temp2 = uc.γ + M * γ_x
    uc.μ =  temp1 / temp2

    # Update γ
    uc.γ = 1 / (ρ^2 / (uc.γ + M * γ_x) + σ_θ^2)
end

update_θ!(uc::UncertaintyTrapEcon, w::Real) =
    (uc.θ = uc.ρ * uc.θ + uc.σ_θ * w)

"""
Generate aggregates based on current beliefs (μ, γ).  This
is a simulation step that depends on the draws for F.
"""
function gen_aggregates(uc::UncertaintyTrapEcon)
    F_vals = uc.σ_F * randn(uc.num_firms)

    M = sum(ψ.(uc, F_vals) .> 0)  # Counts number of active firms
    if M > 0
        x_vals = uc.θ + uc.σ_x * randn(M)
        X = mean(x_vals)
    else
        X = 0.0
    end
    return X, M
end


using LinearAlgebra
include("gates.jl")

# common states
ket_0 = Complex[1; 0]
ket_1 = Complex[0; 1]
ket_plus = Complex[1; 1] ./ sqrt(2)
ket_minus = Complex[1; -1] ./ sqrt(2)

# TODO
# struct Op
#     U::Matrix{Complex}
#     isUnitary::Bool
# end

"""
    conj_transpose(s)

conj_transpose is a clearer name for adjoint for students.
it is identical to `s'`
"""
function conj_transpose(s::Union{AbstractVecOrMat})
    conj.(transpose(s))
end

"""
    commutator(A, B)

commutator(A, B) := A * B - B * A
Takes in two operators, if the result is zero, the operators commute.
"""
function commutator(A, B)
    A * B - B * A
end

"""
    ops_commute(A, B)

ops_commute(A, B) := commutator(A,B) ≈ fill(0., (dim, dim))
for ops to commute they must be square and equal dim, so dim
    is arbitrarily size(A, 1)
"""
function ops_commute(A, B)
    if size(A) != size(B)
        return false
    end
    dim = size(A, 1)
    commutator(A,B) ≈ fill(0., (dim, dim))
end

    """
    is_unitary(s)

A unitary matrix U is defined by U†U = UU† = I.
returns true if `s` is unitary.
Note `s` must be square.
"""
# TODO: fix 1/sqrt(2) and normalization
function is_unitary(s)
    LinearAlgebra.checksquare(s)
    s * conj_transpose(s) ≈ Id
end

"""
    is_hermitian(s)

A hermitian matrix U is defined by U = †U.
returns true if s is hermitian.
"""
# throws error if not hermitian
function is_hermitian(s)
    LinearAlgebra.checksquare(s);
    s ≈ conj_transpose(s)
end

"""
    is_normed(s)

if sum(s .^ 2)) ≈ 1., returns true
"""
# throws error if s is not normalized
function is_normed(s)
    mag = convert.(Float64, sum(s .^ 2))
    @assert mag ≈ 1.
end

"""
    is_linear(s)

"""
function is_linear()

end

"""
    inner(a, b)

given two equal dimension column vectors, a and b,
returns the inner product <a|b> where <a| == conj_transpose(a)
this is a duplicate function of dot(a, b) or (a ⋅ b).
"""
# use ⋅ ?
function inner(a, b)
    conj_transpose(a) * b
end

"""
    outer(a, b)

given two equal dimension column vectors, |a> and |b>,
returns the inner product |a><b| where <b| == conj_transpose(b)
note, this function does NOT return the same as kron(a, b).
outer returns a matrix, whereas kron returns a vector
"""
# todo splat n ops
# add assert equal dims
function outer(a, b)
    a * conj_transpose(b)
end

"""
    Cgate(U)

Cgate(U) creates a controlled U gate, given a unitary U.
"""
Cgate(U) = kron(outer(ket_0, ket_0), Id) + kron(outer(ket_1, ket_1), U)

"""
    bell_states()

returns the four bell states as vectors in order of binary increasing
qubit input states (00, 01, 10, 11)
"""
function bell_states()
    # # bell states
    # 00: 00 + 11
    # 10: 00 - 11
    # 01: 01 + 10
    # 11: 01 - 10
    β_00 = Cgate(X) * kron(H * ket_0, ket_0)
    β_01 = Cgate(X) * kron(H * ket_0, ket_1)
    β_10 = Cgate(X) * kron(H * ket_1, ket_0)
    β_11 = Cgate(X) * kron(H * ket_1, ket_1)
    β_00, β_01, β_10, β_11
end

"""
    cauchy_schwarz(u, v)

given arbitrary column vectors `|u⟩, |v⟩`, we want to show that,

⟨u|u⟩⟨v|v⟩ ≥ |⟨u|v⟩|²
"""

function cauchy_schwarz(u, v)

end
# function eulerform(U)
    # thm
    # U single qubit unitary ⟹ ∃ α, β, γ, δ ∈ ℝ, st
    # U = exp(iα) * [exp(-i\)]

# end


module Math3D

using LinearAlgebra

export orthogonalize, orthogonalize!
export rotx, roty, rotz, rotxyz, rot, axisangle, trans, perspective, ortho

function orthogonalize(v::Vector, n::Vector)
	nUnit = normalize(n)
	v - nUnit * dot(v, nUnit)
end

function orthogonalize!(v::Vector, n::Vector)
	nUnit = normalize(n)
	v -= nUnit * dot(v, nUnit)
end

function rotx(m::Matrix, angle::Real)
	s = sin(angle)
	c = cos(angle)
	m[1,1] = 1
	m[2,1] = m[3,1] = m[1,2] = m[1,3] = 0
	m[2,2] = m[3,3] = c
	m[3,2] = s
	m[2,3] = -s
	return m
end

function roty(m::Matrix, angle::Real)
	s = sin(angle)
	c = cos(angle)
	m[1,1] = m[3,3] = c
	m[2,1] = m[1,2] = m[3,2] = m[2,3] = 0
	m[2,2] = 1
	m[3,1] = -s
	m[1,3] = s
	return m
end

function rotz(m::Matrix, angle::Real)
	s = sin(angle)
	c = cos(angle)
	m[1,1] = m[2,2] = c
	m[2,1] = s
	m[1,2] = -s
	m[3,1] = m[3,2] = m[1,3] = m[2,3] = 0
	m[3,3] = 1
	return m
end

# rotation with angles around each of the axes
function rotxyz(m::Matrix, xAngle::Real, yAngle::Real, zAngle::Real)
	# Tait-Bryan angles transform in the order X1Y2Z3
	sx = sin(xAngle)
	cx = cos(xAngle)
	sy = sin(yAngle)
	cy = cos(yAngle)
	sz = sin(zAngle)
	cz = cos(zAngle)

	m[1,1] = cy*cz
	m[2,1] = cx*sz+cz*sx*sy
	m[3,1] = sx*sz-cx*cz*sy

	m[1,2] = -cy*sz
	m[2,2] = cx*cz-sx*sy*sz
	m[3,2] = cz*sx+cx*sy*sz

	m[1,3] = sy
	m[2,3] = -cy*sx
	m[3,3] = cx*cy

	return m
end

rotx(angle::Real) = rotx(Array{typeof(angle)}(undef, 3, 3), angle)
roty(angle::Real) = roty(Array{typeof(angle)}(undef, 3, 3), angle)
rotz(angle::Real) = rotz(Array{typeof(angle)}(undef, 3, 3), angle)
rotxyz(xAngle::Real, yAngle::Real, zAngle::Real) = rotxyz(Array{typeof(xAngle)}(undef, 3, 3), xAngle, yAngle, zAngle)

function rot(m::Matrix, axis::Vector, angle::Real)
	u = normalize(axis[1:3])
	ux = u[1]
	uy = u[2]
	uz = u[3]
	s = sin(angle)
	c = cos(angle)

	m[1,1] = c+ux*ux*(1-c)
	m[2,1] = uy*ux*(1-c)+uz*s
	m[3,1] = uz*ux*(1-c)-uy*s

	m[1,2] = ux*uy*(1-c)-uz*s
	m[2,2] = c+uy*uy*(1-c)
	m[3,2] = uz*uy*(1-c)+ux*s

	m[1,3] = ux*uz*(1-c)+uy*s
	m[2,3] = uy*uz*(1-c)-ux*s
	m[3,3] = c+uz*uz*(1-c)

	return m
end

rot(axis::Vector, angle::Real) = rot(Array{typeof(angle)}(undef, 3, 3), axis, angle)

function axisangle(m::Matrix{T}) where T
	t = m[1,1] + m[2,2] + m[3,3]
	anglecos = (t-1)/2
	anglecos >= 1 - eps(T) && return T[0,0,1], zero(T)
	axis = [m[3,2] - m[2,3], m[1,3] - m[3,1], m[2,1] - m[1,2]]
	if anglecos <= -1 + eps(T)
		squares = T[(m[1,1]+1)/2, (m[2,2]+1)/2, (m[3,3]+1)/2]
		axis /= 4
		i = indmax(squares)
		axis[i] = squares[i]
		maxel = sqrt(squares[i])
		axis /= maxel
		return axis,T(pi)
	end
	axislen = norm(axis)
	anglesin = axislen/2
	angle = atan2(anglesin, anglecos)
	axis /= axislen
	axis, angle
end

function trans(m::Matrix, t::Vector)
	m[1,4] = t[1]
	m[2,4] = t[2]
	m[3,4] = t[3]
	return m
end

trans(t::Vector) = trans(Matrix{eltype(t)}(I, 4, 4), t)

function scale(m::Matrix, s::Vector)
  m[1,1] = s[1]
  m[2,2] = s[2]
  m[3,3] = s[3]
  m[1,2] = m[1,3] = m[2,1] = m[2,3] = m[3,1] = m[3,2] = 0
end

scale(s::Vector) = diagm(0 => vcat(s, 1))

function perspective(m::Matrix, left::Real, right::Real, top::Real, bottom::Real, near::Real, far::Real; leftHanded = false)
	zsign = leftHanded ? -1 : 1
	m[1:4, 1:4] = [ 2near/(right-left)				   0  (right+left)/(right-left)*zsign					  0 ;
									 0  2near/(top-bottom)  (top+bottom)/(top-bottom)*zsign					  0 ;
									 0				   0	 (-far-near)/(far-near)*zsign  -2far*near/(far-near) ;
									 0				   0						 -1*zsign					  0 ]
	return m
end

perspective(left::Real, right::Real, top::Real, near::Real, far::Real; leftHanded = false) = perspective(Array{typeof(near)}(undef, 4, 4), left, right, top, bottom, near, far, leftHanded = leftHanded)
perspective(m::Matrix, width::Real, height::Real, near::Real, far::Real; leftHanded = false) = perspective(m, -0.5width, 0.5width, -0.5height, 0.5height, near, far, leftHanded = leftHanded)
perspective(width::Real, height::Real, near::Real, far::Real; leftHanded = false) = perspective(Array{typeof(near)}(undef, 4, 4), width, height, near, far, leftHanded = leftHanded)
function persp_horizontal_fov(m::Matrix, hfov::Real, w_h_ratio::Real, near::Real, far::Real; leftHanded = false)
	width = 2near*tan(hfov/2)
	perspective(m, width, width / w_h_ratio, near, far, leftHanded = leftHanded)
end
persp_horizontal_fov(hfov::Real, w_h_ratio::Real, near::Real, far::Real; leftHanded = false) = persp_horizontal_fov(Array{typeof(near)}(undef, 4, 4), hfov, w_h_ratio, near, far, leftHanded = leftHanded)
function persp_vertical_fov(m::Matrix, vfov::Real, w_h_ratio::Real, near::Real, far::Real; leftHanded = false)
	height = 2near*tan(vfov/2)
	perspective(m, w_h_ratio * height, height, near, far, leftHanded = leftHanded)
end
persp_vertical_fov(vfov::Real, w_h_ratio::Real, near::Real, far::Real; leftHanded = false) = persp_vertical_fov(Array{typeof(near)}(undef, 4, 4), vfov, w_h_ratio, near, far, leftHanded = leftHanded)

function ortho(m::Matrix, left::Real, right::Real, top::Real, bottom::Real, near::Real, far::Real; leftHanded = false)
	zsign = leftHanded ? -1 : 1
	m[1:4, 1:4] = [ 2/(right-left)			   0				   0 (-right-left)/(right-left) ;
								 0  2/(top-bottom)				   0 (-top-bottom)/(top-bottom) ;
								 0			   0 -2/(far-near)*zsign	 (-far-near)/(far-near) ;
								 0			   0				   0						  1 ]
	return m
end

ortho(left::Real, right::Real, top::Real, bottom::Real, near::Real, far::Real; leftHanded = false) = ortho(Array{typeof(near)}(undef, 4, 4), left, right, top, bottom, near, far, leftHanded = leftHanded)
ortho(m::Matrix, width::Real, height::Real, near::Real, far::Real; leftHanded = false) = ortho(m, -0.5width, 0.5width, -0.5height, 0.5height, near, far, leftHanded = leftHanded)
ortho(width::Real, height::Real, near::Real, far::Real; leftHanded = false) = ortho(Array{typeof(near)}(undef, 4, 4), width, height, near, far, leftHanded = leftHanded)

end


"""
    magic_number(n::Int64))

A magic number has to build as the power of 5 or the sum of unique powers of 5. First few
magic numbers are:
    1. 5
    2. 5^2 = 25
    3. 5 + 25 = 30
    4. 5^3 = 125
    5. 125 + 5 = 130

# Arguments
- `n::Int64`: Integer value of the number, which has to become a magic number

# Examples
```julia-repl
julia> import ClassicAlgorithmsCollections
julia> ClassicAlgorithmsCollections.magic_number(5)
130
```
"""
function magic_number(n::Int64)

    pow = 1
    result = 0

    while n != 0
        pow *= 5
        if n & 1 % Bool
            result += pow
        end
        n >>= 1
    end
    return result
end


function compute_log_joint(observations, partition, tau0, tau1, sigma, theta)
  n = length(observations)
  k = length(partition)
  prob = k*log(sigma) + lgamma(theta) + lgamma(theta/sigma + k) - lgamma(theta/sigma) - lgamma(theta + n)
  for cluster in partition
    prob += lgamma(length(cluster) - sigma) - lgamma(1 - sigma)
    prob += compute_log_conditonal_observations(observations, cluster, tau0, tau1)
  end
  prob
end

function compute_log_conditonal_observations(observations, cluster, tau0, tau1)
  nl = length(cluster)
  prob = (nl/2)*log(tau1) - (nl/2)*log(2*pi) + 0.5*log(tau0) + 0.5*log(tau0+nl)
  prob += -tau1/2*(sum(observations)) + 0.5*(tau0*mu_0+tau1*sum(observations[cluster]))^2/(tau0+nl*tau1)
  prob
end

# Test of similarity between distributions
function correct_posterior(empirical_probs, data, partitions, τ0, τ1, σ, θ)
    true_log_probs = map(p -> compute_log_joint(data, p, τ0, τ1, σ, θ), partitions)
    true_probs = exp.(true_log_probs)
    true_probs /= sum(true_probs)

    empirical_probs /= sum(empirical_probs)

    # compare distribitions
    # L2
    L2 = sum((empirical_probs - true_probs).^2)

    # Discrepancy
    discr = maximum(abs.(empirical_probs - true_probs))
    return L2, discr
end


"""
    NeuralNetworkKernel()

Kernel of a Gaussian process obtained as the limit of a Bayesian neural network with a
single hidden layer as the number of units goes to infinity.

# Definition

Consider the single-layer Bayesian neural network
``f \\colon \\mathbb{R}^d \\to \\mathbb{R}`` with ``h`` hidden units defined by
```math
f(x; b, v, u) = b + \\sqrt{\\frac{\\pi}{2}} \\sum_{i=1}^{h} v_i \\mathrm{erf}\\big(u_i^\\top x\\big),
```
where ``\\mathrm{erf}`` is the error function, and with prior distributions
```math
\\begin{aligned}
b &\\sim \\mathcal{N}(0, \\sigma_b^2),\\\\
v &\\sim \\mathcal{N}(0, \\sigma_v^2 \\mathrm{I}_{h}/h),\\\\
u_i &\\sim \\mathcal{N}(0, \\mathrm{I}_{d}/2) \\qquad (i = 1,\\ldots,h).
\\end{aligned}
```
As ``h \\to \\infty``, the neural network converges to the Gaussian process
```math
g(\\cdot) \\sim \\mathcal{GP}\\big(0, \\sigma_b^2 + \\sigma_v^2 k(\\cdot, \\cdot)\\big),
```
where the neural network kernel ``k`` is given by
```math
k(x, x') = \\arcsin\\left(\\frac{x^\\top x'}{\\sqrt{\\big(1 + \\|x\\|^2_2\\big) \\big(1 + \\|x'\\|_2^2\\big)}}\\right)
```
for inputs ``x, x' \\in \\mathbb{R}^d``.[^CW]

[^CW]: C. K. I. Williams (1998). Computation with infinite neural networks.
"""
struct NeuralNetworkKernel <: Kernel end

function (κ::NeuralNetworkKernel)(x, y)
    return asin(dot(x, y) / sqrt((1 + sum(abs2, x)) * (1 + sum(abs2, y))))
end

function kernelmatrix(::NeuralNetworkKernel, x::ColVecs, y::ColVecs)
    validate_inputs(x, y)
    X_2 = sum(x.X .* x.X; dims=1)
    Y_2 = sum(y.X .* y.X; dims=1)
    XY = x.X' * y.X
    return asin.(XY ./ sqrt.((X_2 .+ 1)' * (Y_2 .+ 1)))
end

function kernelmatrix(::NeuralNetworkKernel, x::ColVecs)
    X_2_1 = sum(x.X .* x.X; dims=1) .+ 1
    XX = x.X' * x.X
    return asin.(XX ./ sqrt.(X_2_1' * X_2_1))
end

function kernelmatrix_diag(::NeuralNetworkKernel, x::ColVecs)
    x_2 = vec(sum(x.X .* x.X; dims=1))
    return asin.(x_2 ./ (x_2 .+ 1))
end

function kernelmatrix_diag(::NeuralNetworkKernel, x::ColVecs, y::ColVecs)
    validate_inputs(x, y)
    x_2 = vec(sum(x.X .* x.X; dims=1) .+ 1)
    y_2 = vec(sum(y.X .* y.X; dims=1) .+ 1)
    xy = vec(sum(x.X' .* y.X'; dims=2))
    return asin.(xy ./ sqrt.(x_2 .* y_2))
end

function kernelmatrix(::NeuralNetworkKernel, x::RowVecs, y::RowVecs)
    validate_inputs(x, y)
    X_2 = sum(x.X .* x.X; dims=2)
    Y_2 = sum(y.X .* y.X; dims=2)
    XY = x.X * y.X'
    return asin.(XY ./ sqrt.((X_2 .+ 1) * (Y_2 .+ 1)'))
end

function kernelmatrix(::NeuralNetworkKernel, x::RowVecs)
    X_2_1 = sum(x.X .* x.X; dims=2) .+ 1
    XX = x.X * x.X'
    return asin.(XX ./ sqrt.(X_2_1 * X_2_1'))
end

function kernelmatrix_diag(::NeuralNetworkKernel, x::RowVecs)
    x_2 = vec(sum(x.X .* x.X; dims=2))
    return asin.(x_2 ./ (x_2 .+ 1))
end

function kernelmatrix_diag(::NeuralNetworkKernel, x::RowVecs, y::RowVecs)
    validate_inputs(x, y)
    x_2 = vec(sum(x.X .* x.X; dims=2) .+ 1)
    y_2 = vec(sum(y.X .* y.X; dims=2) .+ 1)
    xy = vec(sum(x.X .* y.X; dims=2))
    return asin.(xy ./ sqrt.(x_2 .* y_2))
end

Base.show(io::IO, ::NeuralNetworkKernel) = print(io, "Neural Network Kernel")


function ldl_decomp(A)
    # assumption: the matrix A is positive definite (Hermitian)
    # not explicitly checked but assumed: n_row == n_col
    # only the lower left part is used from A

    (n_row, n_col) = size(A)
    L = UnitLowerTriangular( zeros(n_row, n_col) )
    d = zeros(n_col)

    # initial d element
    @inbounds d[1] = A[1,1]
    @inbounds for jj = 1:(n_col-1)

        # determine the next column of L
        @inbounds for ii = (jj+1):n_row
            internal_sum = zero(eltype(A))
            @inbounds @simd for kk = 1:(jj-1)
                @fastmath internal_sum += d[kk] * L[jj,kk] * L[ii,kk]
            end
            @fastmath L[ii, jj] = (A[ii, jj] - internal_sum)/ d[jj]
        end

        # determine the next d element
        diag_sum = zero(eltype(A))
        @inbounds @simd for kk = 1:jj
            @fastmath diag_sum += d[kk] * L[jj+1, kk]^2
        end
        @fastmath d[jj+1] = A[jj+1, jj+1] - diag_sum
    end

    return (L, d)
end

function linsolve_ldl(A, b)

    (L, d) = ldl_decomp(A)
    y = triangular_solve_lower_colwise(L, b)
    @fastmath y ./= d
    x = triangular_solve_upper_colwise_transposed(L, y)

    return x
end


### A Pluto.jl notebook ###
# v0.14.8

using Markdown
using InteractiveUtils

# ╔═╡ 86a927d5-0b05-4a7a-83c4-0f98fb15509e
begin
	using PlutoUI
	using Plots
end

# ╔═╡ 5d15bae3-d408-4f74-98a3-8b5e53562770
md"""
# Chapter 4
#### Nicetak
 $(import Dates; Dates.format(Dates.today(), Dates.DateFormat("U d, Y")))
"""

# ╔═╡ 435ece74-d26a-11eb-1bd1-5b65d058f3a1
md"""
## 4.1
トリボナッチ数列の第$N$項の値を求める再起関数を設計してください.
"""

# ╔═╡ 4a01190f-1e0e-4371-9847-e66eeab849b9
function f₁(n)
	if n == 0
		return 0
	elseif n == 1
		return 0
	elseif n == 2
		return 1
	end

	return f₁(n - 1) + f₁(n - 2) + f₁(n - 3)
end

# ╔═╡ 0224c51b-3028-45f0-bf6c-7423aac643a6
with_terminal() do
	for n ∈ 0:9
		println("n = $n: $(f₁(n))")
	end
end

# ╔═╡ aa51aa81-b8be-4e65-b259-e816655d01c9
plot(n -> (@timed f₁(n)).time, 1:33,
	title = "Computational Time of Naive Tribonacci",
	xlabel = "n",
	ylabel = "Seconds",
	legend = false
)

# ╔═╡ 5d8bedcf-6f20-4be3-961b-3b81ef486209
md"""
## 4.2
問題4.1で設計した再帰関数をメモ化によって効率化してください.
また, メモ化を実施後の計算量を評価してください.
"""

# ╔═╡ 1f806382-2995-4da2-b33d-c0403c530163
function f₂(n, memo = -ones(Int, n))
	if n == 0
		return 0
	elseif n == 1
		return 0
	elseif n == 2
		return 1
	end

	if memo[n] != -1
		return memo[n]
	end

	res = f₂(n - 1, memo) + f₂(n - 2, memo) + f₂(n - 3, memo)
	memo[n] = res
	return res
end

# ╔═╡ 645cdd0a-97c5-4a90-a70e-09a3571142db
with_terminal() do
	for n ∈ 0:9
		println("n = $n: $(f₂(n))")
	end
end

# ╔═╡ d65730a4-ec5f-436f-876b-fc98eb203107
md"""
計算量は$T_0, \dots, T_{N-1}$までメモ化されている時, $T_N$の計算量が$O(1)$であることから, $O(N)$.
"""

# ╔═╡ 3ac6e8f6-a3e9-4d81-a271-5457b20477bc
plot(n -> median((@timed f₂(n)).time for _ in 1:100), 1:3000,
	title = "Computational Time of Memorized Tribonacci",
	xlabel = "n",
	ylabel = "Seconds",
	legend = false
)

# ╔═╡ f6e73830-ad09-4718-868d-26f6a2c11b16
md"""
## 4.3
フィボナッチ数列の一般項が
$F_N = \frac{1}{\sqrt{5}}\left[\left(\frac{1 + \sqrt{5}}{2}\right)^N -
\left(\frac{1 - \sqrt{5}}{2}\right)^N \right]$
で表されることを示してください.
"""

# ╔═╡ 65503b27-14d4-42c8-abec-70060988d465
md"""
**答**
$F_0 = 1, F_1 = 1$は自明. 一般項を$F_{N-1}$まで仮定して,

$\begin{aligned}
\sqrt{5}(F_{N-1} + F_{N-2}) &= \left(1 + \frac{2}{1 + \sqrt{5}}\right)
\left(\frac{1 + \sqrt{5}}{2}\right)^{N-1} -
\left(1 + \frac{2}{1 - \sqrt{5}}\right)
\left(\frac{1 - \sqrt{5}}{2}\right)^{N-1} \\
&= \left(\frac{1 + \sqrt{5}}{2}\right)^N -
\left(\frac{1 - \sqrt{5}}{2}\right)^N.
\end{aligned}$
よって,
$F_N = F_{N - 1} + F_{N - 2}$.
帰納法から, 命題は示された.
"""

# ╔═╡ ae54b3e7-e04c-457f-839d-8aeecb6eba0d
md"""
## 4.4
コード 4.5で示したアルゴリズムの計算量が $O\left(\left(\frac{1 + \sqrt{5}}{2}\right)^N\right)$ で与えられることを示してください.
"""

# ╔═╡ e0cd89ce-9a49-4d51-9c36-1b467b3e4add
md"""
**答** コード4.5のアルゴリズムの計算量が$O(f(N))$であると仮定する.
定義から, $f(0) = f(1) = 1$は自明. また, Big Oの計算規則とコードの定義から,

$f(N) = f(N - 1) + f(N - 2).$

これはフィボナッチ数列と一致し, 一般項は問4.3で示された通りである.
支配項は$\left(\frac{1 + \sqrt{5}}{2}\right)^N$であるから, 命題は示された.
"""

# ╔═╡ 0303b656-d949-4365-b884-bf7bbc026a6d
md"""
## 4.6
部分和問題に対する再帰関数を用いる計算量$O(2^N)$のコード4.9に対してメモ化して, $O(NW)$ の計算量で動作するようにしてください.
"""

# ╔═╡ 14bcda24-9431-468a-ac5b-4bf6feb1e8e8
function f₆(as, w, i = length(as), memo = -ones(Int, length(as), w))
	if w < 0
		return false
	elseif w == 0
		return true
	elseif i == 0
		return false
	end

	if memo[i, w] == 1
		return true
	elseif memo[i, w] == 0
		return false
	end

	if f₆(as, w, i - 1, memo)
		memo[i, w] = 1
		return true
	elseif f₆(as, w - as[i], i - 1, memo)
		memo[i, w] = 1
		return true
	else
		memo[i, w] = 0
		return false
	end

end

# ╔═╡ 501824d0-017f-4d05-a477-cd7f4cef1df8
with_terminal() do
	println("Given a = [2, 4, 8]")
	for w ∈ 1:16
		println("w = $w: $(f₆([2, 4, 8], w))")
	end
end

# ╔═╡ Cell order:
# ╟─5d15bae3-d408-4f74-98a3-8b5e53562770
# ╠═86a927d5-0b05-4a7a-83c4-0f98fb15509e
# ╟─435ece74-d26a-11eb-1bd1-5b65d058f3a1
# ╠═4a01190f-1e0e-4371-9847-e66eeab849b9
# ╟─0224c51b-3028-45f0-bf6c-7423aac643a6
# ╟─aa51aa81-b8be-4e65-b259-e816655d01c9
# ╟─5d8bedcf-6f20-4be3-961b-3b81ef486209
# ╠═1f806382-2995-4da2-b33d-c0403c530163
# ╟─645cdd0a-97c5-4a90-a70e-09a3571142db
# ╟─d65730a4-ec5f-436f-876b-fc98eb203107
# ╟─3ac6e8f6-a3e9-4d81-a271-5457b20477bc
# ╟─f6e73830-ad09-4718-868d-26f6a2c11b16
# ╟─65503b27-14d4-42c8-abec-70060988d465
# ╟─ae54b3e7-e04c-457f-839d-8aeecb6eba0d
# ╟─e0cd89ce-9a49-4d51-9c36-1b467b3e4add
# ╟─0303b656-d949-4365-b884-bf7bbc026a6d
# ╠═14bcda24-9431-468a-ac5b-4bf6feb1e8e8
# ╟─501824d0-017f-4d05-a477-cd7f4cef1df8


### A Pluto.jl notebook ###
# v0.12.7

using Markdown
using InteractiveUtils

# ╔═╡ 6fd616f6-1ea1-11eb-3814-8bfb4a096c49
using Pkg, DrWatson

# ╔═╡ 6ffe2628-1ea1-11eb-24ea-57f985146a72
begin
	@quickactivate "StatisticsWithJuliaPlutoNotebooks"
	using Distributions, Random, StatsBase, DataFrames, Plots
	Random.seed!(1)
end;

# ╔═╡ ed174bc4-1ea0-11eb-1e2f-a32874cec549
md"## Listing 3.17"

# ╔═╡ 70179b62-1ea1-11eb-13ce-7fb4d54bad4b
function rouletteSpins(r,p)
    x = 0
    wins = 0
    while true
        x += 1
        if rand() < p
            wins += 1
            if wins == r
                return x
            end
        end
    end
end

# ╔═╡ 6b47d86c-2163-11eb-157b-136057260d22
begin
	r, p, N = 5, 18/37,10^6
	xGrid = r:r+15

	mcEstimate = counts([rouletteSpins(r,p) for _ in 1:N],xGrid)/N

	nbDist = NegativeBinomial(r,p)
	nbPmf = [pdf(nbDist,x-r) for x in xGrid]
end;

# ╔═╡ 6b480968-2163-11eb-0f1c-d18166a51f5e
begin
	plot( xGrid, mcEstimate,
		line=:stem, marker=:circle, c=:blue,
		ms=10, msw=0, lw=4, label="MC estimate")
	plot!( xGrid, nbPmf, line=:stem,
		 marker=:xcross, c=:red, ms=6, msw=0, lw=2, label="PMF",
		 xlims=(0,maximum(xGrid)), ylims=(0,0.2),
		 xlabel="x", ylabel="Probability")
end

# ╔═╡ 70182cd8-1ea1-11eb-094d-8d0d49cf15f3
md"## End of listing 3.17"

# ╔═╡ Cell order:
# ╟─ed174bc4-1ea0-11eb-1e2f-a32874cec549
# ╠═6fd616f6-1ea1-11eb-3814-8bfb4a096c49
# ╠═6ffe2628-1ea1-11eb-24ea-57f985146a72
# ╠═70179b62-1ea1-11eb-13ce-7fb4d54bad4b
# ╠═6b47d86c-2163-11eb-157b-136057260d22
# ╠═6b480968-2163-11eb-0f1c-d18166a51f5e
# ╟─70182cd8-1ea1-11eb-094d-8d0d49cf15f3


using LinearAlgebra
using StatsBase

"""
    shift_to_non_negative(matrix::Array)

Returns a matrix in which values are non-negative. This is done by finding the
minimal value in the input matrix and adding its absolute value to the matrix
elements.
"""
function shift_to_non_negative(matrix::Array)

    min_val = findmin(matrix)[1]
    if min_val < 0
        return matrix .-= min_val
    else
        return matrix
    end
end

"""
    normalize_to_01(matrix::Array; use_factor=false, norm_factor=256)

Returns a matrix which values are in range [0, 1]. If 'use_factor' is set to
'true' then values are normalized to 'norm_factor' (by default set to 256).

If the values in the input matrix are below 0, then they are shifted so that only positive numbers are in
the matrix (the values are normalized to new maximal value or norm_factor).
"""
function normalize_to_01(matrix::Array; use_factor = false, norm_factor = 256)
    normalized_matrix = copy(matrix)

    min_val = findmin(normalized_matrix)[1]
    if min_val < 0
        normalized_matrix .+= abs(min_val)
    else
        normalized_matrix .-= abs(min_val)
    end

    max_val = findmax(normalized_matrix)[1]

    if use_factor
        if max_val > norm_factor
            @warn "Maximal values exceed \'norm_factor\'."
        end
        normalized_matrix = normalized_matrix ./ norm_factor
    else
        normalized_matrix = normalized_matrix ./ max_val
    end

    return normalized_matrix
end

# function symmetrize_image(image)
"""
    function diagonal_symmetrize(image::Matrix; below_over_upper::Bool=false)

Takes an 'image' in the form of a matrix and return a copy which is symmetric
with respect to diagonal- values above diagonal are copied over values below the
diagonal. This can be inverted by setting 'below_over_upper=true'.

If the input matrix is not square, then square matrix is created by taking
matrix of k times 'k' elements, 'k=min(r,c)' where 'r' is number of rows and 'c'
is number of columns.
"""
function diagonal_symmetrize(image::Matrix; below_over_upper::Bool = false)
    w, h = size(image)
    mat_size = findmin([w, h])[1]

    img = copy(image[1:mat_size, 1:mat_size])

    # Get all cartesian indices from input matrix
    matrix_indices = CartesianIndices((1:mat_size, 1:mat_size))
    # Filter out indices below diagonal
    if below_over_upper
        matrix_indices = findall(x -> x[1] > x[2], matrix_indices)
    else
        matrix_indices = findall(x -> x[2] > x[1], matrix_indices)
    end

    # how many elements are above diagonal
    repetition_number = Int(ceil((mat_size * (mat_size - 1)) / 2))

    # Substitute elements
    for k = 1:repetition_number
        # n_pos = matrix_indices[k]
        mat_ind = matrix_indices[k]
        # ordered_matrix[mat_ind] = k
        img[mat_ind[2], mat_ind[1]] = img[mat_ind]
    end

    try
        checksquare(img)
    catch err
        if isa(err, DimensionMismatch)
            @error "Resulting matrix is not a square matrix"
            throw(err)
        end
    end
    # issymmetric(Float64.(img))
    return img
end

# =====
# matrix ordering

"""
    function get_ordered_matrix(in_matrix::Matrix;
                                assign_same_values::Bool = false,
                                force_symmetry::Bool = false,
                                small_dist_grouping::Bool = false,
                                min_dist::Number = 1e-16,
                                total_dist_groups::Int = 0,
                                ordering_start::Int=1)

Takes a @input_matrix and returns ordered form of this matrix.
The ordered form is a matrix which elements represent ordering from smallest to
highest values in @input_matrix.

If @input_matrix is symmetric, then ordering happens only with upper diagonal.
Lower diagonal is symetrically copied from values above diagonal.

By default, if there is a geoup of entriess with the same value, they all are
assigned with the same ordering number. This can be changed with
@assign_same_values parameter.

Symetry ordering can be froced with @force_symmetry parameter.

By setting 'small_dist_grouping' to true, all the values that difference is
lower than 'min_dist', will be assigned with the same order number.

# Examples
```julia-repl
julia> a = [0 11 12;
            11 0 13;
            12 13 0];
julia> get_ordered_matrix(a)
3×3 Array{Int64,2}:
 0  1  2
 1  0  3
 2  3  0
```

```julia-repl
julia> b = [38 37 36 30;
            37 34 30 32;
            36 30 31 30;
            30 32 30 29]
julia> get_ordered_matrix(b; assign_same_values=false)
4×4 Array{Int64,2}:
0  6  5  2
6  0  1  4
5  1  0  3
2  4  3  0

julia> get_ordered_matrix(b; assign_same_values=true)
4×4 Array{Int64,2}:
0  4  3  1
4  0  1  2
3  1  0  1
1  2  1  0
```
"""
function get_ordered_matrix(in_matrix::Matrix;
                            assign_same_values::Bool = false,
                            force_symmetry::Bool = false,
                            small_dist_grouping::Bool = false,
                            min_dist::Number = 1e-16,
                            total_dist_groups::Int = 0,
                            ordering_start::Int=1)

    # TODO Symmetry must be forced for matrix in which there are NaN elements- needs
    #   to be further investigated
    # TODO not working for negative only values
    # TODO check for square matrix

    # ==
    mat_size = size(in_matrix)
    ord_mat = zeros(Int, mat_size)

    # how many elements are above diagonal
    if issymmetric(in_matrix) || force_symmetry
        matrix_indices =
            generate_indices(mat_size, symmetry_order = true, include_diagonal = false)
        do_symmetry = true
    else
        matrix_indices = generate_indices(mat_size, symmetry_order = false)
        do_symmetry = false
    end
    total_elements = length(matrix_indices)

    # Collect vector of indices
    all_ind_collected = arr_to_vec(matrix_indices)

    # Sort indices vector according to inpu array
    # TODO Cant this be done with sortperm? in_matrix > UpperTriangular |> sortperm
    index_sorting = sort_indices_by_values(in_matrix, all_ind_collected)

    ordering_number = ordering_start
    for k = 1:total_elements
        # global ordering_number
        next_sorted_pos = index_sorting[k]
        mat_ind = matrix_indices[next_sorted_pos]

        if assign_same_values && k != 1
            prev_sorted_pos = index_sorting[k-1]
            prev_mat_ind = matrix_indices[prev_sorted_pos]

            cond1 = in_matrix[prev_mat_ind] == in_matrix[mat_ind]
            cond2 = small_dist_grouping
            cond3 = abs(in_matrix[prev_mat_ind] - in_matrix[mat_ind]) < min_dist

            if cond1 || (cond2 && cond3)
                ordering_number -= 1
            end
        end
        set_values!(ord_mat, mat_ind, ordering_number; do_symmetry = do_symmetry)
        ordering_number += 1
        # else
        #     set_values!(ord_mat, mat_ind, ordering_number; do_symmetry=do_symmetry)
        #     ordering_number+=1
        # end
    end

    return ord_mat
end

# TODO this one has to be specified for 3 dim matrix
function get_ordered_matrix(input_array::Array{Any,3}; do_slices = true, dims = 0)
    arr_size = size(input_array)
    out_arr = zeros(Int, arr_size)

    if do_slices
        # param check
        if dims > length(arr_size)
            throw(DomainError("Given dimension is greater than total size of array."))
        elseif dims > 0
            throw(DomainError("Given dimension must be positive value."))
        elseif dims <= 3
            throw(DomainError("Given dimension must be lower than 3."))
        end

        for dim = 1:arr_size[dims]
            if dims == 1
                out_arr[dim, :, :] = get_ordered_matrix(input_array[dim, :, :])
            elseif dims == 2
                out_arr[:, dim, :] = get_ordered_matrix(input_array[:, dim, :])
            elseif dims == 3
                out_arr[:, :, dim] = get_ordered_matrix(input_array[:, :, dim])
            end
        end
    else
        out_arr = get_ordered_matrix(input_array)
    end
end

function get_ordered_matrix(input_array::Array)
    out_array = copy(input_array)
    arr_size = size(input_array)
    total_elements = length(input_array)

    # Collect vector of indices
    all_ind_collected = collect(reshape(generate_indices(arr_size), (length(input_array))))

    # Sort indices vector according to inpu array
    index_sorting = sort!(
        [1:total_elements;],
        by = i -> (input_array[all_ind_collected][i], all_ind_collected[i]),
    )

    for k = 1:total_elements
        target = index_sorting[k]
        out_array[target] = k
    end

    return out_array
end

# Care must be taken so that values from 'input_matrix' are within distance
# groups, otherwise error is thrown.
"""
    function group_distances(input_matrix::Array, total_dist_groups::Int)

Takes a matrix and rearranges values into 'total_dist_groups' number of groups.
Every group is assigned with number value from range '<0,1>'.
"""
function group_distances(input_matrix::Array, total_dist_groups::Int)

    normed_matrix = normalize_to_01(input_matrix)
    target_matrix = copy(normed_matrix)

    h, w = size(input_matrix)

    if h * w < total_dist_groups
        throw(DomainError("Total number of groups exceed total number of entries in input matrix"))
    end

    total_borders = total_dist_groups + 1

    range_val = collect(range(0, 1, length = total_borders))

    for k = 2:total_borders
        indices = findall(x -> x >= range_val[k-1] && x <= range_val[k], normed_matrix)
        target_matrix[indices] .= range_val[k]
    end
    unique(target_matrix)

    # Sets last range to values smaller than unity, just in case this might cause trobules
    # normed_matrix[normed_matrix .> range_val[end-1]] .= 0.99
    return target_matrix
end

"""
    generate_indices(matrix_size::Tuple; symmetry_order::Bool=false, include_diagonal::Bool=true)

Return all the possible indices of the matrix of size 'matrix_size'.
'matrix_size' may be a tuple or a series of integer arguments corresponding to
the lengths in each dimension.

If 'symetry_order' is set to'true', then only indices of values below diagonal
are returned.
"""
function generate_indices(matrix_size::Tuple;
                            symmetry_order::Bool = false,
                            include_diagonal::Bool = true)

    # Get all cartesian indices from input matrix
    matrix_indices = CartesianIndices(matrix_size)
    # Filter out indices below diagonal
    if symmetry_order
        matrix_indices = findall(x -> x[1] <= x[2], matrix_indices)
    else
        matrix_indices = findall(x -> true, matrix_indices)
    end

    if !include_diagonal
        filter!(x -> x[1] != x[2], matrix_indices)
    end

    return matrix_indices
end

"""
    generate_indices(matrix_size::Int; symmetry_order::Bool=false, include_diagonal::Bool=true)

Generate indices for a matrix of given dimensions. 'generate_indices' is a
series of integer arguments corresponding to the lengths in each dimension.
"""
function generate_indices(matrix_size::Int;
                            symmetry_order::Bool = false,
                            include_diagonal::Bool = true)
    return generate_indices(
        (matrix_size, matrix_size);
        symmetry_order = symmetry_order,
        include_diagonal = include_diagonal,
    )
end

"""
    arr_to_vec(some_array::Array)

Takes an array and reshapes it into a vector.
"""
function arr_to_vec(some_array::Array)
    return collect(reshape(some_array, length(some_array)))
end

function cartesianInd_to_vec(some_array::CartesianIndices)
    return collect(reshape(some_array, length(some_array)))
end

"""
    sort_indices_by_values(values_matrix::T, index_vector) where {T<:VecOrMat}

Sorts the 'index_vector' according to corresponding values in the 'values_matrix'
and returns a Vector of intigers which is an list of ordering of
'sorted index_vector'.
"""
function sort_indices_by_values(values_matrix::T, index_vector) where {T<:VecOrMat}
    if !isa(index_vector, Vector)
        throw(TypeError(
            :sort_indices_by_values,
            "\'index_vector\' must be a vector, otherwise an ordering list can no be created!",
            Vector,
            typeof(index_vector),
        ))
    end
    total_elements = length(index_vector)
    return sort!(
        [1:total_elements;],
        by = i -> (values_matrix[index_vector][i], index_vector[i]),
    )
end

"""
    set_values!(input_matrix::Matrix, position::CartesianIndex, target_value::Number; do_symmetry=false)

Assigns 'target_value' to indices at 'input_matrix[position[1], position[2]]'.
If 'do_symmetry' is set to 'true', then the 'target_value' is also assigned at
position 'input_matrix[position[2], position[1]]'.
"""
function set_values!(input_matrix::Matrix,
                        position::CartesianIndex,
                        target_value::Number;
                        do_symmetry::Bool = false)
    input_matrix[position[1], position[2]] = target_value
    if do_symmetry
        input_matrix[position[2], position[1]] = target_value
    end
    return input_matrix
end

# matrix ordering
# =====

function get_high_dim_ordered_matrix(input_matrix)
    matrix_size = size(input_matrix)
    ordered_matrix_3D = zeros(Int, matrix_size)

    for slice = 1:matrix_size[1]
        ordered_matrix_3D[slice, :, :] = get_ordered_matrix(input_matrix[slice, :, :])
    end
    return ordered_matrix_3D
end

"""
    reduce_arrs_to_min_len(arrs)

Takes vector of vectors of different length and returns array of arrays which
are of the same length. Length in the output is the shortest vector length from
the input- values above this size are discarded.
"""
function reduce_arrs_to_min_len(arrs::Array)
    @debug "Argument specific"
    new_arr = copy(arrs)

    simulation = size(new_arr, 1)
    min_size = Inf
    for m = 1:simulation
        @debug "Simulation number" m
        current_size = size(new_arr[m], 1)
        @debug "Current size: " current_size
        if convert(Float64, current_size) < min_size
            min_size = current_size
            @debug "min size changed to: " min_size
        end
    end
    # min_size = Int.(min_size)
    @debug "Concatenating"
    for m = 1:simulation
        new_arr[m] = new_arr[m][1:min_size, :]
    end
    min_size = Inf
    return new_arr
end

"""
    increase_arrs_to_max_len(arrs)

Takes vector of vectors of different length and returns array of arrays which
are of the same length. Length in the output is the longest vector length from
the input- values above this size are discarded.
"""
function increase_arrs_to_max_len(arrs)
    new_arr = copy(arrs)

    simulation = size(new_arr, 1)
    max_size = 0
    for m = 1:simulation
        @debug "Simulation number" m
        current_size = size(new_arr[m], 1)
        @debug "Current size: " current_size
        if convert(Float64, current_size) > max_size
            max_size = current_size
            @debug "min size changed to: " max_size
        end
    end
    # max_size = Int.(max_size)
    @debug "Concatenating"
    for m = 1:simulation
        correct_len_arr = zeros(Int, max_size, 3)
        correct_len_arr[1:size(arrs[m], 1), :] = new_arr[m][:, :]
        new_arr[m] = correct_len_arr
    end
    # min_size = Inf
    return new_arr
end


export GammaShapeRate

import Distributions: Gamma, shape, rate
import SpecialFunctions: loggamma, digamma, gamma
import StatsFuns: log2π

struct GammaShapeRate{T <: Real} <: ContinuousUnivariateDistribution
    a :: T
    b :: T
end

GammaShapeRate(a::Real, b::Real)       = GammaShapeRate(promote(a, b)...)
GammaShapeRate(a::Integer, b::Integer) = GammaShapeRate(float(a), float(b))
GammaShapeRate(a::Real)                = GammaShapeRate(a, one(a))
GammaShapeRate()                       = GammaShapeRate(1.0, 1.0)

Distributions.@distr_support GammaShapeRate 0 Inf

Distributions.support(dist::GammaShapeRate) = Distributions.RealInterval(minimum(dist), maximum(dist))

Distributions.shape(dist::GammaShapeRate)  = dist.a
Distributions.rate(dist::GammaShapeRate)   = dist.b
Distributions.scale(dist::GammaShapeRate)  = inv(dist.b)
Distributions.mean(dist::GammaShapeRate)   = shape(dist) / rate(dist)
Distributions.var(dist::GammaShapeRate)    = shape(dist) / abs2(rate(dist))
Distributions.params(dist::GammaShapeRate) = (shape(dist), rate(dist))

function Distributions.entropy(dist::GammaShapeRate)
    a, b = params(dist)
    return a - log(b) + loggamma(a) + (1-a)*digamma(a)
end

function logmean(dist::GammaShapeRate)
    a, b = params(dist)
    return digamma(a) - log(b)
end

function loggammamean(dist::GammaShapeRate)
    a, b = params(dist)
    return 0.5 * (log2π - (digamma(a) - log(b))) + mean(dist) * (-1 + digamma(a + 1) - log(b))
end

function meanlogmean(dist::GammaShapeRate)
    a, b = params(dist)
    return mean(dist) * (digamma(a + 1) - log(b))
end

Base.eltype(::GammaShapeRate{T}) where T = T

Base.convert(::Type{ GammaShapeRate{T} }, a::Real, b::Real) where { T <: Real } = GammaShapeRate(convert(T, a), convert(T, b))

vague(::Type{ <: GammaShapeRate }) = GammaShapeRate(1.0, tiny)

prod_analytical_rule(::Type{ <: GammaShapeRate }, ::Type{ <: GammaShapeRate }) = ProdAnalyticalRuleAvailable()

function prod(::ProdAnalytical, left::GammaShapeRate, right::GammaShapeRate)
    T = promote_type(eltype(left), eltype(right))
    return GammaShapeRate(shape(left) + shape(right) - one(T), rate(left) + rate(right))
end

Distributions.pdf(dist::GammaShapeRate, x::Real)    = (rate(dist)^shape(dist)) / gamma(shape(dist)) * x^(shape(dist)-1)*exp(-rate(dist) * x)
Distributions.logpdf(dist::GammaShapeRate, x::Real) = shape(dist) * log(rate(dist)) - loggamma(shape(dist)) + (shape(dist)-1)*log(x) - rate(dist)*x


using LabelledArrays

export rossler_lorenz

@inline @inbounds function eom_rossler_lorenz(u, p, t)
    c_xy, a₁, a₂, a₃, b₁, b₂, b₃ = (p...,)
    x1, x2, x3, y1, y2, y3 = u[1], u[2], u[3], u[4], u[5], u[6]

    dx1 = -a₁*(x2 + x3)
    dx2 = a₁*(x1 + a₂*x2)
    dx3 = a₁*(a₂ + x3*(x1 - a₃))
    dy1 = b₁*(-y1 + y2)
    dy2 = b₂*y1 - y2 - y1*y3 + c_xy*(x2^2)
    dy3 = y1*y2 - b₃*y3

    return SVector{6}(dx1, dx2, dx3, dy1, dy2, dy3)
end

"""
    rossler_lorenz(;u₀ = rand(6), a₁ = -6, a₂ = 6, a₃ = 2.0,
        b₁ = 10, b₂ = 28, b₃ = 8/3, c_xy = 1) → ContinuousDynamicalSystem

Initialise a Rössler-Lorenz system consisting of two independent 3D subsystems:
one Rössler system and one Lorenz system. They are coupled such that the
second component (`x₂`) of the Rössler system unidirectionally forces the
second component (`y₂`) of the Lorenz system.

The parameter `c_xy` controls the coupling strength. The implementation here also
allows for tuning the parameters of each subsystem by introducing the constants
`a₁`, `a₂`, `a₃`, `b₁`, `b₂`, `b₃`. Default values for these parameters are
as in [1].

## Equations of motion

The dynamics is generated by the following vector field

```math
\\begin{aligned}
\\dot x_1 &= a_1(x_2 + x_3) \\\\
\\dot x_2 &= a_2(x_1 + 0.2x_2) \\\\
\\dot x_3 &= a_2(0.2 + x_3(x_1 - a_3)) \\\\
\\dot y_1 &= b_1(y_2 - y_1) \\\\
\\dot y_2 &= y_1(b_2 - y_3) - y_2 +c_{xy}(x_2)^2 \\\\
\\dot y_3 &= y_1 y_2 - b_3y_3
\\end{aligned}
```

with the coupling constant ``c_{xy} \\geq 0``.

## References

1. Krakovská, Anna, et al. "Comparison of six methods for the detection of causality in a
    bivariate time series." Physical Review E 97.4 (2018):042207.
    [https://journals.aps.org/pre/abstract/10.1103/PhysRevE.97.042207](https://journals.aps.org/pre/abstract/10.1103/PhysRevE.97.042207)
"""
function rossler_lorenz(;u₀ = rand(6), a₁ = -6, a₂ = 6, a₃ = 2.0,
    b₁ = 10, b₂ = 28, b₃ = 8/3, c_xy = 1)

    p = @LArray [c_xy, a₁, a₂, a₃, b₁, b₂, b₃] (:c_xy, :a₁, :a₂, :a₃, :b₁, :b₂, :b₃)
    ContinuousDynamicalSystem(eom_rossler_lorenz, u₀, p)
end


"""
- second-order (so)
    ||z1|| <= z2

A Semismooth Newton Method for Fast, Generic Convex Programming
    https://arxiv.org/abs/1705.00772

# todo: try backtracking linesearch for cone
    https://digital.lib.washington.edu/researchworks/bitstream/handle/1773/43279/Dueri_washington_0250E_19426.pdf?isAllowed=y&sequence=1
"""

# second-order cone
function κ_so(z)
    z1 = z[1:end-1]
    z2 = z[end]

    z_proj = zero(z)

    if norm(z1) <= z2
        z_proj = copy(z)
    elseif norm(z1) <= -z2
        z_proj = zero(z)
    else
        a = 0.5 * (1.0 + z2 / norm(z1))
        z_proj[1:end-1] = a * z1
        z_proj[end] = a * norm(z1)
    end

    return z_proj
end

# second-order cone Jacobian
function Jκ_so(z)
    z1 = z[1:end-1]
    z2 = z[end]
    m = length(z)

    if norm(z1) <= z2
        return Diagonal(ones(m))
    elseif norm(z1) <= -z2
        return Diagonal(zeros(m))
    else
        D = zeros(m, m)
        for i = 1:m
            if i < m
                D[i, i] = 0.5 + 0.5 * z2 / norm(z1) - 0.5 * z2 * ((z1[i])^2.0) / norm(z1)^3.0
            else
                D[i, i] = 0.5
            end
            for j = 1:m
                if j > i
                    if j < m
                        D[i, j] = -0.5 * z2 * z1[i] * z1[j] / norm(z1)^3.0
                        D[j, i] = -0.5 * z2 * z1[i] * z1[j] / norm(z1)^3.0
                    elseif j == m
                        D[i, j] = 0.5 * z1[i] / norm(z1)
                        D[j, i] = 0.5 * z1[i] / norm(z1)
                    end
                end
            end
        end
        return D
    end
end

function cone_product(z, s)
    if length(z) > 1
        z0 = z[1]
        z1 = z[2:end]

        s0 = s[1]
        s1 = s[2:end]
        return [z' * s; z0 * s1 + s0 * z1]
    else
        return [z' * s]
    end
end

function e_soc(n)
    e = zeros(n)
    e[1] = 1.0
    return e
end


#

using Gadfly, Colors
using Distributions
using FFTW
using DSP

p1 = Normal(-pi/2, 0.5)
p2 = Normal(pi/2, 0.5)
p12_t = Normal(0.0, 1.0/sqrt(8.0))

x = range(-pi, stop=pi, length=1024)

g1 = (x)->pdf(p1,x)
g2 = (x)->pdf(p2,x)
g12_t = (x)->pdf(p12_t,x)

y1 = g1.(x)
y2 = g2.(x)
y12_t = g12_t.(x)

y12_t .*= 0.5
temp = fftshift(y12_t)
y12_t .+= temp

plot(y=y12_t, Geom.line)

Y12_t = fftshift(fft(y12_t))

pl = Gadfly.plot(
layer(y=abs.(Y1), Geom.line, Theme(default_color=colorant"red")),
layer(y=abs.(Y2), Geom.line, Theme(default_color=colorant"blue")),
layer(y=abs.(Y12_t), Geom.line, Theme(default_color=colorant"magenta"))
)

pl |> SVG("/tmp/test.svg",100cm,100cm)

y12_tr = ifft(ifftshift(Y12_t))

plot(y=abs.(y12_tr), Geom.line)

plot(
layer(x=x, y=y1, Geom.line, Theme(default_color=colorant"red")),
layer(x=x, y=y2, Geom.line, Theme(default_color=colorant"blue"))
)

using Makie

scene = lines(x, abs.(y12_tr), color = :blue)
# scatter!(scene, x, y1, color = :red, markersize = 0.1)

Y1 = fftshift(fft(y1))
Y2 = fftshift(fft(y2))
Y12_t = fft(y12_t)

plot(y=abs.(Y1), Geom.line)
plot(y=abs.(Y2), Geom.line)

Y12 = conv((Y1), (Y2))

st = 1
off = 2046
y12 = ifft(ifftshift(Y12[st:(st+off)]))

plot(y=abs.(y12), Geom.line)

# plot(x=real.(Y1), y=imag.(Y1), Geom.point)
# plot(x=real.(Y2), y=imag.(Y2), Geom.point)

plot(y=abs.(Y12), Geom.line)

##

#YY1 = [Y1; Y1; Y1]
YY2 = [ifftshift(Y2); ifftshift(Y2); ifftshift(Y2)]

plot(y=abs.(YY1), Geom.line)

YY12 = conv(ifftshift(Y1), (YY2))

plot(y=abs.(YY12), Geom.line)

st = 1024
off = 2047
yy12 = ifft(ifftshift(YY12[st:(st+off)]))

pl = plot(y=abs.(yy12), Geom.line)

##

using Cairo, Fontconfig
using Gadfly

pl |> PDF("/tmp/test.pdf", 30cm, 20cm)

run(`evince /tmp/test.pdf`)

##

plot(x=real.(Y12_t), y=imag.(Y12_t), Geom.point)

# brute force product
y12_d = y1 .* y2
y12_d ./= sum(y12_d)*(x[2]-x[1])

plot(
layer(x=x, y=y1, Geom.line, Theme(default_color=colorant"red")),
layer(x=x, y=y2, Geom.line, Theme(default_color=colorant"blue")),
layer(x=x, y=y12_d, Geom.line, Theme(default_color=colorant"magenta"))
)

# calculating the square using convolution

y1sqrt = sqrt.(y1)

Y1sqrt = fft(y1sqrt)

Y11 = conv(Y1sqrt, conj(Y1sqrt))

##

st = 1
off = 512
y11 = ifft(Y11[st:(st+off)])

plot(
layer(y=y1, Geom.line, Theme(default_color=colorant"red")),
layer(y=abs.(y11)./(sum(abs.(y11))*(x[2]-x[1])), Geom.line, Theme(default_color=colorant"blue"))
# layer(x=x, y=y12_d, Geom.line, Theme(default_color=colorant"magenta"))
)

## Linear convolution in Euclid space

xx = 0:127;

Y = (x::Real; mu::Real=0.0) -> pdf(Normal(mu, 1.0), x)

y1 = Y.(xx, mu=10)
y2 = Y.(xx, mu=20)

fY1 = fft(y1)
fY2 = fft(y2)

fYY = fY1.*fY2
yy = abs.(ifft(fYY))

plot(
  layer(
    x=xx,
    y=y1,
    Geom.line
  ),
  layer(
    x=xx,
    y=y2,
    Geom.line,
    Theme(default_color=colorant"deepskyblue")
  ),
  layer(
    x=xx,
    y=yy,
    Geom.line,
    Theme(default_color=colorant"magenta")
  )
)

## Test with Sine function

x = range(-7.5, stop=7.5, length=2048)
y = sin.(10.0*x) + sin.(100.0*x)

Y = fft(y)

Ys = fftshift(Y)

plot(y=abs.(Ys), Geom.line)

#


# N = 13195 # 29
 N = 600851475143 # 6857
ubound = Int(ceil(sqrt(N)))
if ubound%2 == 0
	ubound -= 1
end
lpf = 1

function isPrime(n::Int)
	for i in 2:ceil(sqrt(n))
		if n%i==0
			return false
		end
	end
	return true
end

while ubound > 1
	if N%ubound==0 && isPrime(ubound)
		lpf = ubound
		break
	end
	ubound -= 2
end
println(N, '\t', lpf)


"""
diffper(n,xspan)

Construct 2nd-order differentiation matrices for functions with
periodic end conditions, using `n` unique nodes in the interval
`xspan`. Return a vector of nodes and the  matrices for the first
and second derivatives.
"""
function diffper(n,xspan)
    a,b = xspan
    h = (b-a)/n
    x = @. a + h*(0:n-1)   # nodes, omitting the repeated data

    # Construct Dx by diagonals, then correct the corners.
    dp = fill(0.5/h,n-1)        # superdiagonal
    dm = fill(-0.5/h,n-1)       # subdiagonal
    Dx = diagm(-1=>dm,1=>dp)
    Dx[1,n] = -1/(2*h)
    Dx[n,1] = 1/(2*h)

    # Construct Dxx by diagonals, then correct the corners.
    d0 =  fill(-2/h^2,n)        # main diagonal
    dp =  ones(n-1)/h^2         # superdiagonal and subdiagonal
    Dxx = diagm(-1=>dp,0=>d0,1=>dp)
    Dxx[1,n] = 1/(h^2)
    Dxx[n,1] = 1/(h^2)

    return x,Dx,Dxx
end


function get_energy( p::Param, every::Int , tlimit::Float64)

    #Get variables
    xmin        = p.xmin
    xmax        = p.xmax
    Nx          = p.xnodes

    ymin        = p.ymin
    ymax        = p.ymax
    Ny          = p.ynodes

    #Discretize space
    dx      = (xmax - xmin)/(Nx-1)
    x       = [xmin + (i-1)*dx for i in 1:Nx]

    dy      = (ymax - ymin)/(Ny-1)
    y       = [ymin + (i-1)*dy for i in 1:Ny]

    if(p.Boundaries == :periodic )
        dx      = (xmax - xmin)/(Nx)
        x       = [xmin + (i-1)*dx for i in 1:Nx]

        dy      = (ymax - ymin)/(Ny)
        y       = [ymin + (i-1)*dy for i in 1:Ny]
    end

    #Create Differential Operators
    Dx  =   Diff_Operator_2D(1 ,1, [dx,dy] , [Nx,Ny])
    Dxx =   Diff_Operator_2D(2 ,1, [dx,dy] , [Nx,Ny])

    Dy  =   Diff_Operator_2D(1 ,2, [dx,dy] , [Nx,Ny])
    Dyy =   Diff_Operator_2D(2 ,2, [dx,dy] , [Nx,Ny])

    if(p.Boundaries == :periodic )
        Dx  =   Diff_Operator_2D_Periodic(1 ,1, [dx,dy] , [Nx,Ny])
        Dxx =   Diff_Operator_2D_Periodic(2 ,1, [dx,dy] , [Nx,Ny])

        Dy  =  Diff_Operator_2D_Periodic(1 ,2, [dx,dy] , [Nx,Ny])
        Dyy =  Diff_Operator_2D_Periodic(2 ,2, [dx,dy] , [Nx,Ny])
    end

    #Get max time iteration
    max_sim_iter , tmax = get_t_max(p)
    max_plot_iter = max_sim_iter

    if(tlimit <=  tmax)
        max_plot_iter = floor(Int , max_sim_iter * tlimit/tmax )
    else
        println("ATTENTION: max time input is larger than simulation time, using tmax")
    end
    #Mas user iteration

    #Potential Matrix
    Vpot = potential_matrix(p)

    # Total volume of space
    Volume = ((ymax-ymin)*(xmax-xmin))

    #Storage vectors
    E_vector = Float64[]
    t_vector = Float64[]

    t = 0
    E = 0

    #Loop to get energy
    for i in 1:every:max_plot_iter

        #Time and fields
        t, _ , _ , ψ , dψ = get_fields(p,i)

        #Volume element
        dV = dx*dy
        E = 0

        #Gradient term (possibly wrong)
         vx = Dx * reshape(ψ, Nx*Ny)
         vy = Dy * reshape(ψ, Nx*Ny)

        #Potential term
        v2 = Vpot * reshape(ψ, Nx*Ny)

        #Add all terms
        for i in 1:(Nx*Ny) #length(ψ)
            E+= vx[i]^2 + vy[i]^2 + dψ[i]^2 + (p.μ)^2*ψ[i]^2 + v2[i]^2
        end

        #Add to storage vectors  and normalize
        push!(t_vector, t)
        push!(E_vector, dV * E / Volume )
    end

    return t_vector , E_vector

end

function get_energy( fname::String, every::Int , tlimit::Float64)

    #Get Variables
    fid         = h5open( fname , "r")

    xmin = read_attribute(fid, "xmin")
    xmax = read_attribute(fid, "xmax")
    Nx   = read_attribute(fid, "xnodes")

    ymin = read_attribute(fid, "ymin")
    ymax = read_attribute(fid, "ymax")
    Ny   = read_attribute(fid, "ynodes")

    #Discretize space
    dx      = (xmax - xmin)/(Nx-1)
    x       = [xmin + (i-1)*dx for i in 1:Nx]

    dy      = (ymax - ymin)/(Ny-1)
    y       = [ymin + (i-1)*dy for i in 1:Ny]

    #Create Differential Operators
    Dx  =   Diff_Operator_2D(1 ,1, [dx,dy] , [Nx,Ny])
    Dxx =   Diff_Operator_2D(2 ,1, [dx,dy] , [Nx,Ny])

    Dy  =   Diff_Operator_2D(1 ,2, [dx,dy] , [Nx,Ny])
    Dyy =   Diff_Operator_2D(2 ,2, [dx,dy] , [Nx,Ny])

    #Get max time iteration
    max_sim_iter , tmax = get_t_max(fname)
    max_plot_iter = max_sim_iter

    if(tlimit <=  tmax)
        max_plot_iter = floor(Int , max_sim_iter * tlimit/tmax )
    else
        println("ATTENTION: max time input is larger than simulation time, using tmax")
    end
    #Mas user iteration

    #Potential Matrix
   # Vpot = potential_matrix(p)

    # Total volume of space
    Volume = ((ymax-ymin)*(xmax-xmin))

    #Storage vectors
    E_vector = Float64[]
    t_vector = Float64[]

    t = 0
    E = 0

    #Loop to get energy
    for i in 1:every:max_plot_iter

        #Time and fields
        t, _ , _ , ψ , dψ = get_fields(fname,i)

        #Volume element
        dV = dx*dy
        E = 0

        #Gradient term (possibly wrong)
         vx = Dx * reshape(ψ, Nx*Ny)
         vy = Dy * reshape(ψ, Nx*Ny)

        #Potential term
       # v2 = Vpot * reshape(ψ, Nx*Ny)

        #Add all terms
        for i in 1:(Nx*Ny) #length(ψ)
            E+= vx[i]^2 + vy[i]^2 + dψ[i]^2 # + (p.μ)^2*ψ[i]^2  # + v2[i]^2
        end

        #Add to storage vectors  and normalize
        push!(t_vector, t)
        push!(E_vector, dV * E / Volume )
    end

    return t_vector , E_vector

end


# Simple example of an inhomogeneous heat equation:
#
#           y
#             |
#             |     u_y = 0
#          Ly ----------------------------------
#             |                                |
#             |                                |
#             |                                |
#     u_x = 0 |    u_t - a(u_xx + u_yy) = f    | u_x = 0
#             |                                |
#             |                                |
#             |                                |
#           -----------------------------------|-----  x
#                         u_y = 0              Lx
#
# Here, a>0, the initial data is u0 = 0 and the source term is
#
#      f(x, y, t) = t cos(pi x / Lx) cos(pi y / Ly)
#
# The exact solution is
#
#                    lambda t - 1 - exp(-lambda t)
#      u(x, y, t) =  ----------------------------- phi(x,y)
#                          lambda^2
#
# where
#
#      phi(x,y) = cos(pi x / Lx ) cos(pi y / Ly)
#      lambda = a pi^2 ( 1 / Lx^2 + 1 / Ly^2 )
#
using FinElt
using FinElt.PlanarPoisson
include("ode23s.jl")

const a = 0.1
const Lx = 1.0
const Ly = 1.0
const lambda = a * pi^2 * ( 1/Lx^2 + 1/Ly^2 )
const T = 2.0

function phi(x)
    return cos(pi*x[1]/Lx)* cos(pi*x[2]/Ly)
end

function exact_u(x, t)
    v = ( lambda*t - 1 + exp(-lambda*t) ) / lambda^2
    return v * phi(x)
end

function f(x, t)
    return t * phi(x)
end

function RHS(t, u, S, mesh, dof)
    F = assembled_vector("Omega", source_times_func!,
                         x -> f(x,t), mesh, dof)
    return F - S * u
end

maxnorm(x) = norm(x, Inf)

start = time()
mesh = read_msh_file("../rectangle/rect2.msh")
dof = degrees_of_freedom(mesh, String[])
M = assembled_matrix("Omega", func_times_func!, 1.0, mesh, dof)
S = assembled_matrix("Omega", grad_dot_grad!, a, mesh, dof)
u0 = zeros(length(dof.freenode))
t, u = ode23s((t,u) -> RHS(t, u, S, mesh, dof), u0, [0,T];
             reltol=1.0e-4, abstol=1.0e-4, mass=M, norm=maxnorm)

N = length(t)
err = zeros(N)
for n = 1:N
    un = get_nodal_vals(x->exact_u(x,t[n]), mesh)
    err[n] = maxnorm(un-u[n])
end

finish = time()
elapsed = finish - start

figure(1)
plot(t, err)
xlabel(L"$t$")
title("Error solving the inhomogeneous heat equation")
grid(true)


#=
This first example comes almost straight from the documentation for DiffEqFlux

https://github.com/SciML/DiffEqFlux.jl#optimizing-parameters-of-an-ode-for-an-optimal-control-problem
https://julialang.org/blog/2019/01/fluxdiffeq/

=#

module lotka_volterra

using DifferentialEquations, Flux, Optim, DiffEqFlux, Plots

u0 = [1.0,1.0]
tstart=0.0
tend=10.0
sampling=0.1

model_params= [1.5,1.0,3.0,1.0]

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

function InitPlot(;vars=(0,1))
    #Simple function to plot the dynamics
    tspan=(tstart,tend)
    prob = ODEProblem(model,u0,tspan,model_params)
    sol = solve(prob,Tsit5())
    plot(sol, lw=2, legend=false, vars=(0,1))
end

function predict_adjoint(param) # Our 1-layer neural network
    prob=ODEProblem(model,[1.0,0.0],(tstart,tend), model_params)
    Array(concrete_solve(prob,Tsit5(),param[1:2],param[3:end],saveat=tstart:sampling:tend, abstol=1e-8,reltol=1e-6))
end

# Generate some data to fit, and add some noise to it
data=predict_adjoint([1.0,1.0,1.5,1.0,3.0,1.0])
σN=0.1
data+=σN*randn(size(data))
data=abs.(data) #Keep measurements positive

# Returning more than just the loss function breaks the Flux optim
function loss_adjoint(param)
  prediction = predict_adjoint(param)
  loss = sum(abs2,prediction - data)
  loss
end;

#Test

function train_model(;pguess=[0.8,1.2,1.2,1.0,2.9,1.1])
    println("The initial loss function is $(loss_adjoint(pguess)[1])")
    #Train the ODE
    resinit=DiffEqFlux.sciml_train(loss_adjoint,pguess,ADAM(), maxiters=3000)
    res = DiffEqFlux.sciml_train(loss_adjoint,resinit.minimizer,BFGS(initial_stepnorm = 1e-5))
    println("The parameters are $(res.minimizer) with final loss value $(res.minimum)")
    return(res)
end

function plotFit(param)

    tspan=(tstart,tend)
    sol_fit=solve(ODEProblem(model,param[1:2],tspan,param[3:end]), Tsit5())

    tgrid=tstart:sampling:tend
    pl=plot(sol_fit, lw=2, legend=false)
    scatter!(pl,tgrid, data[1,:], color=:blue)
    scatter!(pl,tgrid, data[2,:], color=:red)
    xlabel!(pl,"Time")
    ylabel!(pl,"Population")
    title!(pl,"Model Parameter Fits")
    savefig("Lotka_Volterra_ParamFit.png")
    display(pl)
end

function validationPlot(param, ic)
    tspan=(tstart,tend)
    sol_fit=solve(ODEProblem(model,ic,tspan,param[3:end]), Tsit5())
    sol_actual=solve(ODEProblem(model,ic,tspan,model_params), Tsit5(), saveat=0.0:0.1:10.0)

    pl=scatter(sol_actual, lw=2.0, color=:blue, vars=(0,1))
    scatter!(sol_actual, color=:red, vars=(0,2))
    plot!(pl, sol_fit, lw=2, legend=false, color=:blue, vars=(0,1))
    plot!(pl, sol_fit, lw=2, color=:red, vars=(0,2))
    xlabel!(pl,"Rabbits")
    ylabel!(pl,"Lynx")
    title!(pl,"Validation Plot")
    savefig("Lotka_Volterra_Validation_Plot.png")
    display(pl)
end

end #module

lotka_volterra.InitPlot()

resL=lotka_volterra.train_model()

lotka_volterra.plotFit(resL.minimizer)

lotka_volterra.validationPlot(resL.minimizer, [6.,6.])


#=
Algorithm RSVD.
given a matrix A \in \reals^{m \times n}, a target rank k and an oversampling
      parameter p (e.g., p = 10 is a good choice).
stage A. Find an approximate range.
    1. form an n \times (k+p) Gaussian random matrix G.
    2. form the sample matrix Y = AG.
    3. orthonormalize the columns of Y via a QR factorization.
stage B. Form a specific factorization.
    4. form the (k+p) \times n matrix B = Q'A.
    5. form the SVD of the (small) matrix B as B = \hat{U} D V'.
    6. form U = Q \hat{U}.
return matrices U, D, and V as an approximate rank (k+p) SVD of A.
=#
using Random, LinearAlgebra, Plots, Statistics
# ---------------------------------------------------------------------------- #
function vignette_rsvd()
rng = Random.seed!(2);              # set seed for reproducibility
n_sims = 10;                        # number of simulations
n_subs = 25;                        # number of subsamples
m = 2000;                           # rows of matrix
n = 20*ceil(log(m));                # columns of matrix
# n_sims = 50; n_subs = 25; m = 1000; n = 15*ceil(log(m)); # alt. run parameters
# n_sims = 1;  n_subs = 25; m = 5000; n = 15*ceil(log(m)); # alt. run parameters
k = ceil(log(m));                   # target rank
p = max(ceil(log(m)), 10);          # oversampling parameter
n = Int16(n); k = Int16(k); p = Int16(p);   # convert floats to integers
# ---------------------------------------------------------------------------- #
times = zeros(n_sims, 2);           # bookkeeping for run times
fro_mean = zeros(n_sims, n_subs);   # bookkeeping for mean Frobenius norm
op_mean = zeros(n_sims, n_subs);    # bookkeeping for mean operator norm
norm_bound = zeros(n_sims, 2);      # bookkeeping for theoretrical norm
fro_lo = zeros(n_sims, 1);          # bookkeeping for Frobenius norm bound
op_lo = zeros(n_sims, 1);           # bookkeeping for operator norm bound
# ---------------------------------------------------------------------------- #
for j in 1:n_sims
    # "data" matrix
    A = [2*ones(m,2)+rand(m,2) randn(m, k-2) 0.01*randn(m, n-k)]/sqrt(m);
    # NOTE: the matrix A has k columns that lead to "important" singular values.
    # The remaing n-k columns correspond to fast-decaying singular values.
    for i in 1:n_subs
        # stage A of randomized SVD
        G = randn(n, k+p);                              # Gaussian random matrix
        Y = A*G;                                        # sample matrix
        F = qr(Y); Q = Matrix(F.Q);                     # orthonormalize Y
        # stage B of randomized SVD
        B = Q'*A;                                       # form small matrix
        U_B, D_B, V_B = svd(B);                         # SVD of B
        U = Q*U_B;                                      # rank k matrix
        # bookkeeping and comparisons
        U_A, D_A, V_A = svd(A);                         # SVD of A
        fro_mean[j,i] = norm(A - Q*Q'*A, 2);
        op_mean[j,i] = opnorm(A - Q*Q'*A);
    end # end of inner simulation loop (i.e., simulation for a fixed matrix A)
    fro_lo[j,1] = sum( D_A[k+1:min(m,n)].^2 )^0.5;
    op_lo[j,1] =  D_A[k+1];
    norm_bound[j, 1] = (1 + k / (p-1))^(0.5) * fro_lo[j,1];     # Frobenius
    norm_bound[j, 2] = (1 + sqrt(k / (p-1))) * op_lo[j,1] +     # operator
                            exp(1) * (sqrt(k+p) / p) * fro_lo[j,1];
end # end of outer simulation loop
# ---------------------------------------------------------------------------- #
# plot showing singular value decay
p3 = plot(D_A, yscale = :log10, linecolor = :blue,
            marker = :circle, markercolor = :blue, label = "full",
            title = "sing. value decay (final sim.)");
p4 = plot(D_A[1:k+p], yscale = :log10, linecolor = :blue,
            marker = :circle, markercolor = :blue, label = "full",
            title = "sing. value comparison (final sim.)")
plot!(D_B, yscale = :log10, linecolor = :red, linestyle = :dash,
            marker = :x, markerstrokecolor = :red, label = "randomized");

# ---------------------------------------------------------------------------- #
# plots of average Frobenius and operator norms vs. theoretical bounds
p1 = plot(norm_bound[:,1],linecolor = :blue, marker = :circle,
            markercolor = :blue, label = "upper",
            title = "E-Y bounds: Frobenius");
plot!(mean(fro_mean, dims = 2),linecolor = :red, linestyle = :dash,
     marker = :x, markerstrokecolor = :red, label = "mean");
     plot!(fro_lo, linecolor = :blue, linestyle = :dot,
           marker = :star8, markercolor = :blue, label = "lower",
           legend = :bottomright);
p2 = plot(norm_bound[:,2],linecolor = :blue, marker = :circle,
            markercolor = :blue, label = "upper",
            title = "E-Y bounds: operator");
plot!(mean(op_mean, dims = 2),linecolor = :red, linestyle = :dash,
     marker = :x, markerstrokecolor = :red, label = "mean");
plot!(op_lo,linecolor = :blue, linestyle = :dot,
      marker = :star8, markercolor = :blue, label = "lower",
      legend = :bottomright);
# return summary plot as "output" of the function
plot(p3, p4, p1, p2, layout=(2,2))
end # end of function
# ---------------------------------------------------------------------------- #
vignette_rsvd()


# Some misc functions related to quadratic optimization problems
#
import LinearAlgebra: Symmetric, UniformScaling

@doc raw"""
```julia
check_first_order(∇f::AbstractVector{<:Real},
                  xstar::AbstractVector{<:Real},
                  bc::BoundConstraints{<:Real,1})

check_first_order(Q::Symmetric{<:Real},
                  q::AbstractVector{<:Real},
                  xstar::AbstractVector{<:Real},
                  bc::BoundConstraints{<:Real,1})
```

Check First-Order Conditions
(see [Bound Constrained Optimization slides](https://wiki.mcs.anl.gov/leyffer/images/0/01/07-bndCons.pdf))

If ``x^\star=\arg\min f(x), x\in[l,u]`` then:

```math
\partial_i f(x^\star) = \left\{\begin{array}{ll}
\ge 0, & \text{if } x^\star[i] = l[i] \\
= 0, & \text{if } l[i] \le x^\star[i] \le u[i] \\
\le 0, & \text{if } x^\star[i] = u[i] \\
\end{array}
\right.
```

This is equivalent to:
```math
x^\star = P_{[l,u]}(x^\star-\nabla f(x^\star))
```

According to the previous result, this function returns:
```math
\max \mid x^\star - P_{[l,u]}(x^\star-(Q.x^\star+q)) \mid
```

For a local stationary point this quantity must be null

The second function is a wrapper that computes ``∇f=Q.x^\star+q``
"""
function check_first_order(∇f::AbstractVector{<:Real},
                           xstar::AbstractVector{<:Real},
                           bc::BoundConstraints{<:Real,1})

    # the condition assumes feasible constraints
    @assert xstar ∈ bc

    @assert size(∇f)==size(xstar)

    # TODO: refactoring to avoid mem alloc
    # (-> inline computation of max(...sequence...)
    v = xstar-∇f
    v = project!(v,bc)

    maximum(abs.(xstar .- v))
end

function check_first_order(Q::Symmetric{<:Real},
                           q::AbstractVector{<:Real},
                           xstar::AbstractVector{<:Real},
                           bc::BoundConstraints{<:Real,1})
    # TODO: for BlasFloat use BLAS... to avoid unnecessary memory
    # allocs...
    ∇f=Q*xstar+q
    check_first_order(∇f,xstar,bc)
end


module TheAlgorithms

# Usings/Imports (keep sorted)
using DataFrames
using DifferentialEquations
using GLM
using LinearAlgebra
using Plots
using Random

## Exports
# Please keep the folders/functions sorted

# Exports: knapsack
export ZeroOnePack!
export CompletePack!

# Exports: math
export abs_max
export abs_min
export abs_val
export area_circle
export area_ellipse
export area_heron_triangle
export area_parallelogram
export area_rectangle
export area_rhombus
export area_square
export area_trapezium
export area_triangle
export ceil_val, floor_val
export collatz_sequence
export euler_method
export is_armstrong
export line_length
export mean
export median
export mode
export SIR # TODO: make the name lowercase if possible
export surfarea_cube
export surfarea_sphere
export trapazoidal_area

# Exports: matrix
export determinant
export lu_decompose

# Exports: project-rosalind
export count_nucleotides
export dna2rna
export reverse_complement

# Exports: searches
export binarysearch
export search

# Exports: statistics
export OLSbeta
export PearsonCorrelation
# Exports: strings
export is_palindrome

# Exports: scheduling
export fcfs

## Includes
# Please keep the folders/files sorted (by dependencies then alphabetical order)

# Includes: knapsack
include("knapsack/knapsack.jl")

# Includes: math
include("math/abs.jl")
include("math/area_under_curve.jl")
include("math/armstrong_number.jl")
include("math/area.jl")
include("math/average_mean.jl")
include("math/average_median.jl")
include("math/average_mode.jl")
include("math/ceil_floor.jl")
include("math/collatz_sequence.jl")
include("math/euler_method.jl")
include("math/line_length.jl")
include("math/sir_model.jl")

# Includes: matrix
include("matrix/lu_decompose.jl") # used by determinant.jl
include("matrix/determinant.jl")

# Includes: project-rosalind
include("project-rosalind/count_nucleotide.jl")
include("project-rosalind/dna2rna.jl")
include("project-rosalind/reverse_complement.jl")

# Includes: search
include("searches/binary_search.jl")
include("searches/linear_search.jl")

# Includes: statistics
include("statistics/ordinary_least_squares.jl")

# Includes: strings
include("strings/is_palindrome.jl")

# Includes: scheduling
include("scheduling/fcfs.jl")

end


### A Pluto.jl notebook ###
# v0.12.4

using Markdown
using InteractiveUtils

# ╔═╡ f252364e-f2e3-11ea-09b5-ff1a685abd56
using Pkg, DrWatson

# ╔═╡ c126b36a-f2e3-11ea-1b0f-a7d38ecffcd8
md"## Clip-03-01t.jl"

# ╔═╡ f2526f60-f2e3-11ea-2d32-276d0b91d917
@quickactivate "StatisticalRethinkingTuring"

# ╔═╡ f252ea76-f2e3-11ea-093a-01c585ee2ea1
md"### snippet 3.1"

# ╔═╡ f25a9eec-f2e3-11ea-1611-9129817fe7a6
begin
	Pr_Positive_Vampire = 0.95
	Pr_Positive_Mortal = 0.01
	Pr_Vampire = 0.001
	Pr_Positive = Pr_Positive_Vampire * Pr_Vampire + Pr_Positive_Mortal * (1 - Pr_Vampire)
	Pr_Vampire_Positive = Pr_Positive_Vampire * Pr_Vampire / Pr_Positive
	Pr_Vampire_Positive
end

# ╔═╡ f25b24a2-f2e3-11ea-302d-f3e7f49f645b
md"## End of clip-03-01t.jl"

# ╔═╡ Cell order:
# ╟─c126b36a-f2e3-11ea-1b0f-a7d38ecffcd8
# ╠═f252364e-f2e3-11ea-09b5-ff1a685abd56
# ╠═f2526f60-f2e3-11ea-2d32-276d0b91d917
# ╟─f252ea76-f2e3-11ea-093a-01c585ee2ea1
# ╠═f25a9eec-f2e3-11ea-1611-9129817fe7a6
# ╟─f25b24a2-f2e3-11ea-302d-f3e7f49f645b