# # Hypercube Examples (Bravais Lattice)

# ## Preamble
using LatticeTools
using Plots

mkpath("example_hypercube")

function draw_hypercube(hypercube::Hypercube, coordinates::AbstractMatrix{<:Integer})
    xlim = (minimum(coordinates[1,:]) - 3.5, maximum(coordinates[1,:]) + 3.5)
    ylim = (minimum(coordinates[2,:]) - 3.5, maximum(coordinates[2,:]) + 3.5)
    boundary = let
        r0 = [0,0]
        r1 = hypercube.shape_matrix[:,1]
        r2 = hypercube.shape_matrix[:,2]
        hcat(r0, r1, r1 .+ r2, r2, r0)
    end
    fig = plot(boundary[1,:], boundary[2,:], label="", size=(400, 400))
    for Ri in Iterators.product(-1:1, -1:1)
        alpha = (Ri == (0,0)) ? 1.0 : 0.1
        Rr = hypercube.shape_matrix * [Ri...]
        scatter!(coordinates[1,:] .+ Rr[1], coordinates[2,:] .+ Rr[2], aspect_ratio=1, markersize=12, xlim=xlim, ylim=ylim, markeralpha=alpha, label="")
        annotation = []
        for i in 1:size(coordinates, 2)
            x = coordinates[1,i] + Rr[1]
            y = coordinates[2,i] + Rr[2]
            if xlim[1] < x < xlim[2] && ylim[1] < y < ylim[2]
                push!(annotation, (x, y, text("$i", 8, :black, :center)))
            end
        end
        if !isempty(annotation)
            annotation = [annotation...]
            annotate!(annotation)
        end
    end
    fig
end

# ## (4,-4) x (4,4)
size_matrix = [ 4 4; -4 4]
hypercube = Hypercube(size_matrix)
generator_translations = find_generators(hypercube)
coordinates = generate_coordinates(hypercube, generator_translations)
coordmat = hcat(coordinates...)
println("All elements")
for (i, c) in enumerate(coordinates)
    println("$i : $c")
end
println("Generator translations")
for (it, t) in enumerate(eachcol(generator_translations))
    println("t($it) = $t")
end
draw_hypercube(hypercube, coordmat)
savefig("example_hypercube/bravais_(4,-4)x(4,4).svg")

# ![](example_hypercube//bravais_(4,-4)x(4,4).svg)

# ## (2,-2) x (2,4)
size_matrix = [ 2 2; -2 4]
hypercube = Hypercube(size_matrix)
generator_translations = find_generators(hypercube)
coordinates = generate_coordinates(hypercube, generator_translations)
coordmat = hcat(coordinates...)
draw_hypercube(hypercube, coordmat)
println("All elements")
for (i, c) in enumerate(coordinates)
    println("$i : $c")
end
println("Generator translations")
for (it, t) in enumerate(eachcol(generator_translations))
    println("t($it) = $t")
end
draw_hypercube(hypercube, coordmat)
savefig("example_hypercube/bravais_(2,-2)x(2,4).svg")

# ![](example_hypercube//bravais_(2,-2)x(2,4).svg)


using ProgressBars, Plots, StatsBase, Statistics, LaTeXStrings, JLD

function RandomWalker(P,TimeSteps)
    X = 0
    for i in 1:TimeSteps
        X += sample([1,-1], Weights([P,1-P]))
    end
    return X
end

function Linear_fit(X, Y)
    A = [hcat(X) reshape(ones(length(X)), length(X), 1)]
    b = reshape(hcat(Y), length(Y), 1)
    line = (A \ b)
    return line
end

PList = [0.1,0.3,0.5,0.7,0.9]
N = 100000
TimeSteps = 10000
RES = [[] for i = 1:5]

for n in 1:5
    P = PList[n]
    for i in ProgressBar(1:N)
        push!(RES[n],RandomWalker(P,TimeSteps))
    end
    push!(Avg,mean(RES[n]))
end

save("../../Data/Q4/Q4-RW.jld", "Data", RES)

Line = round.(Linear_fit(PList,Avg))

p1 = histogram(RES[1],bins = 50, label = L"P = %$(PList[1]), \ l = \tau = 1", xlabel=L"\langle X_{final} \rangle_{(P)}", ylabel = L"Number", c = :steelblue)
p2 = histogram(RES[2],bins = 50, label = L"P = %$(PList[2]), \ l = \tau = 1", xlabel=L"\langle X_{final} \rangle_{(P)}", ylabel = L"Number", c = :purple)
p3 = histogram(RES[3],bins = 50, label = L"P = %$(PList[3]), \ l = \tau = 1", xlabel=L"\langle X_{final} \rangle_{(P)}", ylabel = L"Number", c = :red)
p4 = histogram(RES[4],bins = 50, label = L"P = %$(PList[4]), \ l = \tau = 1", xlabel=L"\langle X_{final} \rangle_{(P)}", ylabel = L"Number", c = :gold)
p5 = histogram(RES[5],bins = 50, label = L"P = %$(PList[5]), \ l = \tau = 1", xlabel=L"\langle X_{final} \rangle_{(P)}", ylabel = L"Number", c = :green)
p6 = begin
    plot(PList,Avg,label = L"Y = %$(round(Line[1],digits=3)) X %$(round(Line[2],digits=3))", c = :black)
    scatter!(PList,Avg,label = L"Data\ Points",legend = 140,xlabel = L"P",ylabel = L"\langle \overline{X_{final}} \rangle_{(P)}",c = :steelblue)
end

Pmain1 = begin
    plot(plot(plot(p1,p2,layout=(1,2)),plot(p3,p4,layout=(1,2)),plot(p5,p6,layout=(1,2)),layout=(3,1)))
    plot!(plot_title = L"Distribution\ of\ X_{final}\ (N = %$N\ , Time\ Steps=%$TimeSteps)",size = (1000,1000))
end
savefig(Pmain1,"../../Figs/Q4/Q4-tot.pdf")


function readfile(filename, row::Int64)
    return readdlm(filename)[:,row]
end

function pn(x::Y, xvalues::T, yvalues::T , n::Int64) where {T, Y}
    @assert length(xvalues)==length(yvalues)
    l = length(xvalues)
    denom = 1.0
    numer = 1.0
    for i in 1:l
        if i != n
            denom = denom*(xvalues[n]-xvalues[i])
        end
    end
    for i in 1:l
        if i != n
            numer = numer*(x-xvalues[i])
        end
    end
    return yvalues[n]*numer/denom
end

function lagrangeinterpolation(x::Y, xvalues::T,yvalues::T) where {T,Y}
    @assert length(xvalues)==length(yvalues)
    l = length(xvalues)
    sumpn = 0.0
    for i in 1:l
        sumpn = sumpn + pn(x, xvalues, yvalues, i)
    end
    return sumpn
end


#' # Optimizing water networks with Knitro
#'
#' This notebook explains how to solve effectively NLP and MINLP
#' problem with the Knitro solver.
#'
#' The problem studied hereafter was originally designed by
#' Pierre Carpentier, for educational purposes.
#' The original code is available [here](http://perso.ensta-paristech.fr/~pcarpent/TP_Reseau/ENSMP/).
#' The author credits him for the following tutorial.
#'
#' # Description of the problem
#'
#' The tutorial studies the optimization of water flows in a water
#' distribution network (at steady state). The optimizer aims at
#' minimizing the network's energy while statisfying some linear constraints
#' corresponding to Kirchhoff's first law.
#'
#' Let $G = (N, E)$ be an oriented graph. We note $n = |E|$ the number of
#' arcs and $m = |N|$ the number of nodes in this graph.
#'
#' We suppose that the network has $r$ tanks storing some waters to fulfill
#' the demands in $d$ nodes (distinct from the tanks).
#' We split the set of nodes $N$ accordingly:
#' $$
#' N = N_r \cup N_d.
#' $$
#'
#' We suppose further that the graph is *connected*, implying
#' $$
#' n \geq m - 1 .
#' $$
#' We note $A \in R^{m \times n}$ the incidence matrix of the graph.
#'
#' Lets introduce physical variables to describe further the water network.
#' - the vector of resulting flows at nodes, denoted $f = (f_d, f_r)$. $f_d$ is given for demands' nodes.
#' - the vector of pressures at nodes, denoted $p = (p_d, p_r)$. Pressures $p_r$ are given for the reservoirs.
#' - the vector of resistances in all arcs, denoted $r$ (parameters of the problem)
#' - the vector of flows across arcs, denoted $q$.
#'
#' The decision variable is the vector of flows $q$.
#'
#' ### Constraints
#'
#' The first Kirchhoff law states that:
#' $$
#' A q -f = 0,
#' $$
#' as we suppose that no accumulation occurs in the nodes.
#' The second Kirchhoff law takes into account the losses in the pipes,
#' which is given on each arc by a function $\phi_{alpha}$ (corresponding
#' to the Colebrooks law):
#' $$
#' \phi_{\alpha}(q_\alpha) = r_\alpha q_\alpha | q_\alpha |
#' $$
#' The second Kirchhoff law writes, in a vectorial manner,
#' $$
#' A^\top p + r \circ q \circ | q | = 0
#' $$
#'
#' ### Objective
#' On each arc $\alpha$, we define the energy function $\Phi_\alpha$
#' as
#' $$
#' \Phi_\alpha(q_\alpha) = \dfrac{1}{3} r_\alpha q_\alpha^2 | q_\alpha |
#' $$
#' On the graph, the overall energy equates
#' $$
#' J(q, f_r) = \dfrac{1}{3} q^\top (r \circ q \circ | q |) + p_r^\top f_r
#' $$
#'
#' The global problem writes
#' $$
#' \min_{q} \dfrac{1}{3} q^\top (r \circ q \circ | q |) + p_r^\top f_r \qquad
#' s.t. \quad Aq  - f = 0
#' $$
#'
#' ### Reformulation
#' By applying some mathematical tricks, we are able to reformulate
#' the problem in the following manner.
#'
#' By considering some properties of the overall graph, we split the vector
#' of flows in two $q = (q_T, q_C)$, where $q_T$ depends linearly on $q_C$.
#' Then, we are able to prove the existence of a matrix $B$ and a fixed
#' vector $q_0$ such that the vector of flows on arcs writes
#' $$
#' q = q_0 + B \; q_C
#' $$
#' By using this formulation, we can reduce the dimension of the search
#' space by selecting as decision variable the subvector $q_C$ instead of $q$.
#'
#' Note that the problem becomes also constraint-free, as the vector $q$
#' statisfying the previous equation sastifies also the first Kirchhoff law.
#'
#' We reformulate the optimization problem as
#' $$
#' \min_{q_c} \dfrac{1}{3} q^\top (r \circ q \circ | q |) + p_r^\top f_r \qquad
#' s.t. \quad q = q_0 + B q_C
#' $$

#' ---
#' # Numerical resolution
#' We start by importing JuMP and Knitro.

using JuMP, KNITRO

#' Import data.
include("data.jl");

#' Plotting utilities
using Pkg

# Turned off by default.
PLOT_GRAPH = false
if haskey(Pkg.installed(), "PyPlot")
    PLOT_GRAPH = true
    include("utils.jl");
end

#' ---
#' ## First case study: the non-linear problem
#'
#' Define Knitro Optimizer with JuMP.
model = Model(with_optimizer(KNITRO.Optimizer, outlev=3))

#' Write non-linear optimization problem first.
# Parameters.
α1 = Ar' * pr
# Dimension of the problem.
nx = n - md
@variable(model, qc[1:nx])
# Add dummy variable in the model...
@variable(model, q[1:n])
# ... as we restrict the problem inside a 9-dimensional manifold:
@constraint(model, q .== q0 + B*qc)
@NLobjective(model, Min,
             sum(r[i] * abs(q[i]) * q[i]^2 / 3 + α1[i] * q[i] for i in 1:n))
optimize!(model)

#' Display results
if PLOT_GRAPH
    optimal_flow = JuMP.value.(q)
    fig = figure()
    plot_network(flow=optimal_flow)
    display(fig)
end

#' Observe that some arcs seem almost useless here...

#' ---
#' ## Extension to Mixed-integer non-linear programming

#' We were able to quickly formule the NLP problem with JuMP, and Knitro
#' finds a solution in few iterations.
#' We now modify slightly the problem. Consider the $n$ arcs inside the graph.
#' The questions are:
# - *how many arcs in $E$ can we remove before the problem becomes infeasible?*
# - *how much is the cost impacted by arc removals?*

#' We define hereafter the MINLP version of the problem.
# Add a maximum flows through the pipes
const QMAX = 10.

function load_mip_model!(model::JuMP.Model; nremovals=3)
    # Reconsider the NLP problem introduced previously.
    @variable(model, qc[1:nx])
    @variable(model, q[1:n])
    @constraint(model, q .== q0 + B*qc)

    # Introduce switch: $z$ is set to 0 if the corresponding arc is removed
    # from the graph.
    @variable(model, z[1:n], Bin)

    # Bounds the abs flows through the arcs by QMAX
    # The constraints write:
    # | q | <= q_{max} * z
    # Note that z_a = 0 implies that q_a = 0.
    @constraint(model,  q .<= QMAX * z)
    @constraint(model, -q .<= QMAX * z)
    # Ensure that we remove exactly nremovals arcs.
    @constraint(model, sum(z) == n - nremovals)

    # Same cost as previously.
    @NLobjective(model, Min,
                 sum(r[i] * abs(q[i]) * q[i]^2 / 3 + α1[i] * q[i] for i in 1:n))
    return
end

#' **Optimization related remark:** Note that the constraint
#' $$
#' | q | \leq q_{max} \; z
#' $$
#' could be reformulated as a complementarity constraint (also supported by Knitro
#' but not natively by JuMP).

#' **Physics related remark:** We are a bit loosy with the physics of the
#' problem if we remove the arcs in the previous manner. Imagine that for
#' a given arc $a \in E$, $z_a =0$ thus implying $q_a =0$. Then, if we
#' note $i$ and $j$ the two adjacent nodes, the physics tells us that
#' $$
#' p_i = 0 \, \quad p_j = 0 ,
#' $$
#' which is not the case if we solve the previous optimization problem.
#' However, the goal of this tutorial is purely pedagogical and as a
#' consequence we allow us to play a bit with the physics.

#' ### Default Knitro

#' Solve with default Knitro.
# Build non-linear solver.
model = Model(with_optimizer(KNITRO.Optimizer, outlev=3))
load_mip_model!(model)
@time JuMP.optimize!(model)

#' We plot the solution by displaying in red the arcs removed from the graph.
# Plot!
if PLOT_GRAPH
    optimal_flow = JuMP.value.(model[:q])
    fig = figure()
    plot_network(flow=optimal_flow)
    display(fig)
end

#' In the resolving, Knitro uses the Branch & Bound algorithm to find
#' the optimal solution corresponding to the (convex) MINLP problem.
#' We refer to the [documentation](https://www.artelys.com/docs/knitro/2_userGuide/minlp.html)
#' for further details.
#' The approach used here is different than in BONMIN (outer approximation).

#' Observe that Knitro with default options computes 165 nodes before
#' finding the solution.
#' Is it possible to find a better tuning for Knitro?

#' ### MIP-Tuner

#' Since Knitro 12.0, a MINLP tuner was added to compute the optimal
#' parameterization to solve a given MINLP problem with Knitro. The MINLP
#' tuner uses an exhaustive search to find the optimal setting.
#'
#' By default, the tuner tests 36 combinations.
model = Model(with_optimizer(KNITRO.Optimizer, outlev=3, tuner=1))
load_mip_model!(model)
@time JuMP.optimize!(model)

#' Setting `mip_branchrule=2` and `mip_selectrule=3` seems to give
#' better results, according to the tuner.
#' Note that setting `outlev` to 0 switch off Knitro's output.
mip_solver = with_optimizer(KNITRO.Optimizer, mip_knapsack=0, mip_branchrule=2,
                            mip_selectrule=3, outlev=0)
model = Model(mip_solver)
load_mip_model!(model)
@time JuMP.optimize!(model)

#' ### Quantify impact of arcs' removals

#' The fewer arcs, the harder the problem.
#' If `nremovals >= 10`, Knitro takes too long to solve the problem.
#' However, the problem remains feasible till n_removals = 13.
#' Compute for instance the solution of the following problem:
feas_model = Model(with_optimizer(KNITRO.Optimizer))
@variable(feas_model, qc[1:nx])
@variable(feas_model, q[1:n])
@constraint(feas_model, q .== q0 + B*qc)
@variable(feas_model, z[1:n], Bin)
@constraint(feas_model,  q .<= QMAX * z)
@constraint(feas_model, -q .<= QMAX * z)
@objective(feas_model, Min, sum(z))

#' We know study the evolution of the objective cost w.r.t. the
#' number of removals.
max_removals = 9

# Save results in some arrays.
cost_values = Float64[]

for nremove in 1:max_removals
    println("Remove ", nremove, " arcs from graph.")
    model = Model(mip_solver)
    load_mip_model!(model, nremovals=nremove)
    @time JuMP.optimize!(model)

    push!(cost_values, JuMP.objective_value(model))

    # Plot!
    if PLOT_GRAPH
        optimal_flow = JuMP.value.(model[:q])
        fig = figure()
        plot_network(flow=optimal_flow)
        title("Optimal solution with $nremove removals")
        display(fig)
    end
end

#' Plot evolution of costs w.r.t. number of removals.
if PLOT_GRAPH
    fig = figure()
    plot(1:max_removals, cost_values, lw=3, c="k")
    xlabel("#removals")
    ylabel("Objective value")
    display(fig)
end

#' We observe that removing up to 6 arcs does not impact significantly
#' the cost.


using StatsBase, Distributions, Plots; pyplot()

lambda, N = 1, 10^6
xGrid = 0:6

expDist = Exponential(1/lambda)
floorData = counts(convert.(Int,floor.(rand(expDist,N))), xGrid)/N
geomDist = Geometric(1-MathConstants.e^-lambda)

plot( xGrid, floorData,
	line=:stem, marker=:circle,
	c=:blue, ms=10, msw=0, lw=4,
	label="Floor of Exponential")
plot!( xGrid, pdf.(geomDist,xGrid),
	line=:stem, marker=:xcross,
	c=:red, ms=6, msw=0, lw=2,
	label="Geometric", ylims=(0,1),
	xlabel="x", ylabel="Probability")


"""Calculates the 2nd order centred gradient in x-direction on any grid (u,v,T or q).
The size of dudx must be m-1,n compared to m,n = size(u)"""
function ∂x!(dudx::AbstractMatrix,u::AbstractMatrix)
    m,n = size(dudx)
    @boundscheck (m+1,n) == size(u) || throw(BoundsError())

    @inbounds for j ∈ 1:n
        for i ∈ 1:m
            dudx[i,j] = u[i+1,j] - u[i,j]
        end
    end
end

"""Calculates the 2nd order centred gradient in y-direction on any grid (u,v,T or q).
The size of dudy must be m,n-1 compared to m,n = size(u)."""
function ∂y!(dudy::AbstractMatrix,u::AbstractMatrix)
    m,n = size(dudy)
    @boundscheck (m,n+1) == size(u) || throw(BoundsError())

    @inbounds for j ∈ 1:n
        for i ∈ 1:m
            dudy[i,j] = u[i,j+1] - u[i,j]
        end
    end
end

""" ∇² is the 2nd order centred Laplace-operator ∂/∂x^2 + ∂/∂y^2.
The 1/Δ²-factor is omitted and moved into the viscosity coefficient."""
function ∇²!(du::AbstractMatrix,u::AbstractMatrix)
    m, n = size(du)
    @boundscheck (m+2,n+2) == size(u) || throw(BoundsError())

    @inbounds for i ∈ 1:n
        for j ∈ 1:m
            du[j,i] = minus_4*u[j+1,i+1] + u[j,i+1] + u[j+2,i+1] + u[j+1,i] + u[j+1,i+2]
        end
    end
end


"""
    GaussianLikelihood(σ²)

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

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

GaussianLikelihood() = GaussianLikelihood(1e-6)

GaussianLikelihood(σ²::Real) = GaussianLikelihood([σ²])

@functor GaussianLikelihood

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

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

"""
    HeteroscedasticGaussianLikelihood(l::AbstractLink=ExpLink())

Heteroscedastic Gaussian likelihood.
This is a Gaussian likelihood whose mean and variance are functions of
latent processes.

```math
    p(y|[f, g]) = Normal(y | f, sqrt(l(g)))
```
On calling, this would return a normal distribution with mean `f` and variance `l(g)`.
Where `l` is link going from R to R^+
"""
struct HeteroscedasticGaussianLikelihood{Tl<:AbstractLink}
    invlink::Tl
end

HeteroscedasticGaussianLikelihood() = HeteroscedasticGaussianLikelihood(ExpLink())

function (l::HeteroscedasticGaussianLikelihood)(f::AbstractVector{<:Real})
    return Normal(f[1], sqrt(l.invlink(f[2])))
end

function (l::HeteroscedasticGaussianLikelihood)(fs::AbstractVector)
    return MvNormal(first.(fs), Diagonal(l.invlink.(last.(fs))))
end


using LinearAlgebra

function so3ToVec(so3mat)
    if( size(so3mat)!=(3,3) )
        println("!Warning! : Input matrix should be 3x3.")
        return
    end
    if( !((so3mat[1,1]==so3mat[2,2]==so3mat[3,3]==0)) )
        println("!Warning! : Diagonals should be zero.")
        return
    end
    if( so3mat[2,1]!=-so3mat[1,2] ||
        so3mat[3,1]!=-so3mat[1,3] ||
        so3mat[3,2]!=-so3mat[2,3] )
        println("!Warning! : Entered matrix isn't a skew symmetric matrix.")
        return
    end
    ret_vec = [so3mat[3,2];
               so3mat[1,3];
               so3mat[2,1] ]
        return ret_vec
end

function vecToSo3(omg)
    if( size(omg)!=(3,) )
        println("!Warning! : Input matrix should be 3x1.")
        return
    end

    ret_ss_matrix = [0      -omg[3]   omg[2];
                     omg[3]    0     -omg[1];
                    -omg[2]  omg[1]      0];
    return ret_ss_matrix
end

using LinearAlgebra

function so3ToVec(so3mat)
    if( size(so3mat)!=(3,3) )
        println("!Warning! : Input matrix should be 3x3.")
        return
    end
    if( !((so3mat[1,1]==so3mat[2,2]==so3mat[3,3]==0)) )
        println("!Warning! : Diagonals should be zero.")
        return
    end
    if( so3mat[2,1]!=-so3mat[1,2] ||
        so3mat[3,1]!=-so3mat[1,3] ||
        so3mat[3,2]!=-so3mat[2,3] )
        println("!Warning! : Entered matrix isn't a skew symmetric matrix.")
        return
    end
    ret_vec = [so3mat[3,2];
               so3mat[1,3];
               so3mat[2,1] ]
        return ret_vec
end

function VecToso3(omg)
    if( size(omg)!=(3,) )
        println("!Warning! : Input matrix should be 3x1.")
        return
    end

    ret_ss_matrix = [0      -omg[3]   omg[2];
                     omg[3]    0     -omg[1];
                    -omg[2]  omg[1]      0];
    return ret_ss_matrix
end
#****************************************************************#
# VeysiADN 17 Apr 2021
# Advanced Robotics Homework 4 Codes
#****************************************************************#
using LinearAlgebra
const lAlgebra = LinearAlgebra
#****************************************************************#
# Returns normalized version of given matrix/array
function Normalize(N)
    return N/lAlgebra.norm(N)
end
#****************************************************************#
# Returns 4x4 twist matrix ‘se3mat’ (an element of Lie Algebra se3)
# from given twist vector (= 6x1 vector) ‘xi
function VecTose3( xi )
    w_skew = [ 0   -xi[3]   xi[2];
              xi[3]   0    -xi[1];
             -xi[2]  xi[1]    0 ]
    return vcat(hcat(w_skew,xi[4:6]), zeros(1, 4))
end
#****************************************************************#
# Returns twist vector ‘xi’ from given twist matrix ‘se3mat’
function se3ToVec( se3mat )
    omg = [se3mat[3, 2], se3mat[1, 3],se3mat[2, 1]]
    v = se3mat[1:3, 4]
    return vcat(omg,v)
end
#****************************************************************#
# Calculates and return rotation matrix ‘R’ by
# evaluating matrix exponential of 3x3 skew symmetric matrix ‘so3mat’

function MatrixExp3(so3mat)
        omg = [so3mat[3, 2], so3mat[1, 3],so3mat[2, 1]]
        theta = lAlgebra.norm(omg)
        omg_n = so3mat / theta
        return lAlgebra.I + sin(theta) * omg_n + (1 - cos(theta)) * omg_n * omg_n
end

#****************************************************************#
# Calculates and return transformation matrix ‘g’ by
# evaluating matrix exponential of 4x4 twist matrix ‘se3mat’
function MatrixExp6( se3mat )
    omg = [se3mat[3, 2], se3mat[1, 3],se3mat[2, 1]]
    theta = lAlgebra.norm(omg)
    omg_n = se3mat[1:3, 1:3] / theta
    temp  = hcat(MatrixExp3(se3mat[1:3, 1:3]),
                         (lAlgebra.I * theta +
                          (1 - cos(theta)) * omg_n +
                          (theta - sin(theta)) * omg_n * omg_n) *
                         se3mat[1:3, 4] / theta)
    return    vcat(temp,[0 0 0 1])
end

#****************************************************************#
# VeysiADN 12 May 2021
# Advanced Robotics Homework 5 Codes
#****************************************************************#
using LinearAlgebra
const lAlgebra = LinearAlgebra
#****************************************************************#
# Computes exponential coordinate ‘so3mat’
# (skew symmetric matrix) of the rotation matrix R

function MatrixLog3(R)
    acosinput = (lAlgebra.tr(R) - 1) / 2
    if acosinput >= 1
        return zeros(3, 3)
    else
        theta = acos(acosinput)
        return theta / 2 / sin(theta) * (R - R')
    end
end
#****************************************************************#
# Computes exponential coordinate ‘se3mat’
# (twist matrix) of the rigid body transformation matrix T
function MatrixLog6(T)
    if( size(T)!=(4,4) )
        println("!Warning! : Transformation matrix should be 4x4.")
        return
    end
    R = T[1:3,1:3]
    t = T[1:3,4]
    omg_h = MatrixLog3(R)
    w = [omg_h[3,2];omg_h[1,3];omg_h[2,1]]
    if omg_h == zeros(3, 3)
        return vcat(hcat(zeros(3, 3), T[1:3, 4]), [0 0 0 0])
    else
        theta = acos((lAlgebra.tr(R) - 1) / 2)
        return vcat(hcat(omg_h,
                         (I - omg_h / 2 +
                          (1 / theta - 1 / tan(theta / 2) / 2) *
                          omg_h * omg_h / theta) * T[1:3, 4]),
                    [0 0 0 0])
    end
end
#****************************************************************#

#****************************************************************#
# VeysiADN 12 May 2021
# Advanced Robotics Homework 5 Codes
#****************************************************************#
using LinearAlgebra
#****************************************************************#
# Returns Adjoint matrix with given rigid body
# transformation matrix g
function Adjoint(g)

    R=g[1:3,1:3];
    p=g[1:3,4];
    p_hat = [ 0       -p[3]     p[2];
             p[3]       0      -p[1];
            -p[2]      p[1]       0 ]

    return [R zeros(3,3);p_hat*R R]
end
#****************************************************************#
# VeysiADN 4 Apr 2021
# Advanced Robotics Homework 3 Codes
#****************************************************************#
using LinearAlgebra
#****************************************************************#
# Returns 3 x 3 rotation matrix representing
# rotation about unit vector w by angle theta.
# theta should be in radian
function Rodrigues(w,theta)
    if( length(w) != 3)
        println("!Warning! : W matrix should have 3 elements.")
        return
    end
    if(size(w)==(1,3))
        w=w';
    end
    w_hat = [0      -w[3]    w[2];
             w[3]     0     -w[1];
            -w[2]    w[1]      0];
    rotation_matrix = cos(theta)*Matrix(1I,3,3) + sin(theta)*w_hat+(1-cos(theta))*w*w';
    return rotation_matrix;
end
#****************************************************************#
# Builds rigid transformation matrix T from given
#rotation matrix R and translation t
function RpToTrans(R,t)
    if( size(R)!=(3,3) )
        println("!Warning! : Rotation matrix should be 3x3.")
        return
    end
    if( length(t) != 3 )
        println("!Warning! : Translation matrix should have 3 elements.")
        return
    end
    if(size(t)==(1,3))
        t=t';
    end
    R_B_T = hcat(R,t) ;
    R_B_T = vcat(R_B_T,[0 0 0 1]) ;
    return R_B_T
end
#****************************************************************#
#Extracts rotation matrix R and translation t from given /
# rigid transformation matrix T
function TransToRp(T)
    if( size(T)!=(4,4) )
        println("!Warning! : Transformation matrix should be 4x4.")
        return
    end
    R = T[1:3,1:3]
    t = T[1:3,4]
    return R,t;
end
#****************************************************************#
# Inverts rigid transformation matrix T
function TransInv(T)
    if( size(T)!=(4,4) )
        println("!Warning! : W matrix should be 4x4.")
        return
    end
    if( det(T)==0 )
        println("!Warning! : Inverse doesn't exist (Determinant is equal to zero).")
        return
    end
    R,t = TransToRp(T);
    r_inv =hcat(R',-(R'*t))
    r_inv = vcat(r_inv,[0 0 0 1]);
    return r_inv;
end
#****************************************************************#


# nuclear Norm (times a constant)

"""
  NuclearNorm(λ::Real=1.0)

Returns the function `λ∑σ_i(X)`, where `σ_i(X)` is i-th singular value of matrix X.
"""

immutable NuclearNorm{R <: Real} <: ProximableConvex
  lambda::R
  function NuclearNorm(lambda::R)
    if lambda < 0
      error("parameter λ must be nonnegative")
    else
      new(lambda)
    end
  end
end

NuclearNorm{R <: Real}(lambda::R=1.0) = NuclearNorm{R}(lambda)

function (f::NuclearNorm){T <: RealOrComplex}(X::AbstractArray{T,2})
  U, S, V = svd(X);
  return f.lambda * sum(S);
end

function prox!{T <: RealOrComplex}(f::NuclearNorm, X::AbstractArray{T,2}, Y::AbstractArray{T,2}, gamma::Real=1.0)
  U, S, V = svd(X)

  for i in eachindex(S)
    S[i] = max(0, S[i] - f.lambda*gamma);
  end

  Y[:] = U * diagm(S) * V'
  return f.lambda * sum(S);
end

fun_name(f::NuclearNorm) = "nuclear norm"
fun_dom(f::NuclearNorm) = "AbstractArray{Real,2}, AbstractArray{Complex,2}"
fun_expr(f::NuclearNorm) = "X ↦ λ∑σ_i(X)"
fun_params(f::NuclearNorm) = "λ = $(f.lambda)"

function prox_naive{T <: RealOrComplex}(f::NuclearNorm, X::AbstractArray{T,2}, gamma::Real=1.0)
  U,S,V = svd(X)
  ftemp = NormL1(1.0)
  S_γ, fS_γ =  prox(ftemp,S,f.lambda*gamma)
  Y = U * diagm(S_γ) * V'
  return Y, f.lambda * sum(S_γ)
end


using DifferentialEquations
using DiffEqCallbacks
using Distributions
using StatsBase
using Random
using DataFrames
using StatsPlots
using BenchmarkTools

function rate_to_proportion(r::Float64,t::Float64)
    1-exp(-r*t)
end;

@enum InfectionStatus Susceptible Infected Recovered

Base.zero(::Type{InfectionStatus}) = Infected

function sir_abm!(du,u,p,t)
    (β,c,γ,δt) = p
    N = length(u)
    # Initialize du to u
    for i in 1:N
        du[i] = u[i]
    end
    for i in 1:N # loop through agents
        # If recovered
        if u[i]==Recovered
            continue
        # If susceptible
        elseif u[i]==Susceptible
            ncontacts = rand(Poisson(c*δt))
            while ncontacts > 0
                j = sample(1:N)
                if j==i
                    continue
                end
                a = u[j]
                if a==Infected && rand() < β
                    du[i] = Infected
                    break
                end
                ncontacts -= 1
            end
        # If infected
        else u[i]==Infected
            if rand() < γ
                du[i] = Recovered
            end
        end
    end
    nothing
end;

δt = 0.1
tf = 40.0
tspan = (0.0,tf);

β = 0.05
c = 10.0
γ = rate_to_proportion(0.25,δt)
p = [β,c,γ,δt];

N = 1000
I0 = 10
u0 = Array{InfectionStatus}(undef,N)
for i in 1:N
    if i <= I0
        s = Infected
    else
        s = Susceptible
    end
    u0[i] = s
end

Random.seed!(1234);

susceptible(u) = count(i == Susceptible for i in u)
infected(u) = count(i == Infected for i in u)
recovered(u) = count(i == Recovered for i in u);

saved_values = SavedValues(Float64, Tuple{Int64,Int64,Int64})
cb = SavingCallback((u,t,integrator)->(susceptible(u),infected(u),recovered(u)),
    saved_values,
    saveat=0:δt:tf);

prob_abm = DiscreteProblem(sir_abm!,u0,tspan,p);

sol_abm = solve(prob_abm,
    solver = FunctionMap(),
    dt = δt,
    callback = cb,
    dense = false,
    save_on = false);

df_abm = DataFrame(saved_values.saveval)
rename!(df_abm,[:S,:I,:R])
df_abm[!,:t] = saved_values.t;

@df df_abm plot(:t,
    [:S :I :R],
    label=["S" "I" "R"],
    xlabel="Time",
    ylabel="Number")

@benchmark solve(prob_abm,
    solver=FunctionMap,
    dt=δt,
    callback=cb,
    dense=false,
    save_on=false)


# Example 2.5 from A. Law (2013): Simulation Modeling and Analysis, 5ed

using EventSimulation
using Distributions   # Exponential distribution
using StatsBase       # weighted mean
using DataFrames      # nicer ouptut of results

mutable struct Job
    tin::Float64  # time when job entered the system
    tout::Float64 # time when job left the system
    ttrue::Float64 # time needed without cycling
    remaining::Float64 # time remaining to finish the job
end

mutable struct SharedComputer <: AbstractState
    a::Exponential{Float64}    # arrival time distribution
    p::Exponential{Float64}    # processing time distribution
    τ::Float64                 # switching time
    q::Float64                 # processing quantum time
    n::Int                     # number of terminals
    reps::Int                  # number of finished jobs to collect
    job_queue::Vector{Job}     # queue of jobs waiting for processing
    finished_jobs::Vector{Job} # list of finished jobs
    job::Union{Job, Nothing}  # current job on computer
    # this structure holds information for reporting
    # for each perior between events it holds the length of queue,
    # if the computer were busy and the length of time period
    state_chunks::Vector{Tuple{Int, Bool, Float64}}
    SharedComputer(a, p, τ, q, n, reps) =
        new(Exponential(a), Exponential(p), τ, q, n, reps,
            Job[], Job[], nothing, Tuple{Int, Bool, Float64}[])
end

function arrival(s)
    ttrue = rand(s.state.p)
    j = Job(s.now, NaN, ttrue, ttrue)
    isnothing(s.state.job) && register!(s, start_compute, 0.0)
    push!(s.state.job_queue, j)
end

function start_compute(s)
    j = popfirst!(s.state.job_queue)
    s.state.job = j
    processed = j.remaining < s.state.q ? j.remaining : s.state.q
    j.remaining -= processed
    register!(s, stop_compute, processed + s.state.τ)
end

function stop_compute(s)
    j = s.state.job
    s.state.job = nothing
    if j.remaining == 0
        j.tout = s.now
        push!(s.state.finished_jobs, j)
        register!(s, arrival, rand(s.state.a))
    else
        push!(s.state.job_queue, j)
    end
    isempty(s.state.job_queue) || register!(s, start_compute, 0.0)
    length(s.state.finished_jobs) == s.state.reps && terminate!(s)
end

monitor(s, Δ) =
    push!(s.state.state_chunks,
          (length(s.state.job_queue), !isnothing(s.state.job), Δ))

function run(a, p, τ, q, n, reps)
    sc = SharedComputer(a, p, τ, q, n, reps)
    s = Scheduler(sc, Float64, monitor)
    for _ in 1:n
        register!(s, arrival, rand(sc.a))
    end
    go!(s)
    chunks = s.state.state_chunks
    (n=n,
     response=mean(j.tout - j.tin for j in s.state.finished_jobs),
     queue=mean(getindex.(chunks, 1), Weights(getindex.(chunks, 3))),
     CPU=mean(getindex.(chunks, 2), Weights(getindex.(chunks, 3))))
end

DataFrame([run(25.0, 0.8, 0.015, 0.1, n, 1000) for n in 10:10:80])


"""
    logtp_equal_rates_equal_time(i, j)

Evaluate the logarithm of a transition probability of a simple birth and death
process when the death rate `μ` is equal to the birth rate `λ` and the time `t`
is equal to `1 / λ`. It is a very unlikely event (let's say impossible) but it
simplifies into a special case.
"""
@inline function logtp_equal_rates_equal_time(
  i::I,
  j::I
) where {
  I <: Integer
}
  lgamma(i + j) - lgamma(i) - lgamma(j + 1) - (i + j) * log(2)
end

"""
    logtp_equal_rates(i, j, t, λ)

Evaluate the logarithm of a transition probability of a simple birth and death
process.
"""
@inline function logtp_equal_rates(
  i::I,
  j::I,
  t::F,
  λ::F
)::F where {
  I <: Integer,
  F <: AbstractFloat
}
  a::F, b::F = (j <= i) ? (i, j) : (j, i)

  θ = λ * t

  lgamma(i + j) - lgamma(i) - lgamma(j + 1) + (i + j) * log(θ / (1 + θ)) +
  log_hypergeometric(a, b, logexpm1(-2 * log(θ)))
end

"""
    logtp_equal_rates_alternating(i, j, t, λ)

Evaluate the logarithm of a transition probability of a simple birth and death
process when terms are alternating in sign.
"""
@inline function logtp_equal_rates_alternating(
  i::I,
  j::I,
  t::F,
  λ::F
)::F where {
  I <: Integer,
  F <: AbstractFloat
}
  a::F, b::F = (j <= i) ? (i, j) : (j, i)

  θ = λ * t

  lgamma(i + j) - lgamma(i) - lgamma(j + 1) + (i + j) * log(θ / (1 + θ)) +
  log_meixner_ortho_poly(a, b, 2 * log(θ))
end

"""
    logtp_equal_rates_equal_time_extinction(i)

Evaluate the logarithm of a transition probability of a simple birth and death
process when the death rate `μ` is equal to the birth rate `λ`, time `t` is
equal to `1 / λ`, and 'j' is equal to zero. It is a very unlikely event
(let's say impossible) but it simplifies into a special case.
"""
@inline function logtp_equal_rates_equal_time_extinction(
  i::I
) where {
  I <: Integer
}
  - i * log(2)
end

"""
    logtp_equal_rates_extinction(i, t, λ)

Evaluate the logarithm of the probability of extinction for a simple birth and
death process when the death rate `μ` is equal to the birth rate `λ`.
"""
@inline function logtp_equal_rates_extinction(
  i::I,
  t::F,
  λ::F
)::F where {
  I <: Integer,
  F <: AbstractFloat
}
  i * log((λ * t) / (1 + λ * t))
end

"""
    logtp_pure_birth(i, j, t, λ)

Evaluate the logarithm of the transition probability of a pure-birth process.
"""
@inline function logtp_pure_birth(
  i::I,
  j::I,
  t::F,
  λ::F
)::F where {
  I <: Integer,
  F <: AbstractFloat
}
  lgamma(j) + (j - i) * logexpm1(λ * t) -
  (lgamma(i) + lgamma(j - i + 1) + j * λ * t)
end

"""
    logtp_pure_death(i, j, t, μ)

Evaluate the logarithm of the transition probability of a pure-death process.
"""
@inline function logtp_pure_death(
  i::I,
  j::I,
  t::F,
  μ::F
)::F where {
  I <: Integer,
  F <: AbstractFloat
}
  lgamma(i + 1) + (i - j) * logexpm1(μ * t) -
  (lgamma(j + 1) + lgamma(i - j + 1) + i * μ * t)
end

"""
    logtp_equal_time(i, j, λ, μ)

Evaluate the logarithm of a transition probability of a simple birth and death
process when the time `t` is equal to `log(λ / μ) / (λ - μ)`. It is a very
unlikely event (let's say impossible) but it simplifies into a special case.
"""
@inline function logtp_equal_time(
  i::I,
  j::I,
  λ::F,
  μ::F
)::F where {
  I <: Integer,
  F <: AbstractFloat
}
  lgamma(i + j) + i * log(μ) + j * log(λ) -
  (lgamma(i) + lgamma(j + 1) + (i + j) * log(λ + μ))
end

"""
    logtp_sum(i, j, t, λ, μ)

Evaluate the logarithm of a transition probability of a simple birth and death
process.
"""
@inline function logtp_sum(
  i::I,
  j::I,
  t::F,
  λ::F,
  μ::F
)::F where {
  I <: Integer,
  F <: AbstractFloat
}
  a::F, b::F = (j <= i) ? (i, j) : (j, i)

  θ = (λ - μ) * t
  ω = log(λ / μ)

  x = if λ < μ
    logexpm1(θ - ω) + log1mexp(θ + ω) - 2 * log1mexp(θ)
  else
    logexpm1(θ + ω) + log1mexp(θ - ω) - 2 * logexpm1(θ)
  end

  lgamma(i + j) - lgamma(i) - lgamma(j + 1) + j * ω +
  (i + j) * log(expm1(θ) / expm1(θ + ω)) +
  log_hypergeometric(a, b, x)
end

"""
    logtp_sum_alternating(i, j, t, λ, μ)

Evaluate the logarithm of a transition probability of a simple birth and death
process when terms are alternating in sign.
"""
@inline function logtp_sum_alternating(
  i::I,
  j::I,
  t::F,
  λ::F,
  μ::F
)::F where {
  I <: Integer,
  F <: AbstractFloat
}
  a::F, b::F = (j <= i) ? (i, j) : (j, i)

  θ = (λ - μ) * t
  ω = log(λ / μ)

  x = log(2 * λ * μ) +
      logexpm1(log1pexp(2 * θ) - (θ + log(F(2)))) -
      log((λ - μ)^2)

  lgamma(i + j) - lgamma(i) - lgamma(j + 1) + j * ω +
  (i + j) * log(expm1(θ) / expm1(θ + ω)) +
  log_meixner_ortho_poly(a, b, x)
end

"""
    logtp_equal_time_extinction(i, t, λ, μ)

Evaluate the logarithm of the probability of extinction for a simple birth and
death process when `t = log(λ / μ) / (λ - μ)`.
"""
@inline function logtp_equal_time_extinction(
  i::I,
  λ::F,
  μ::F
)::F where {
  I <: Integer,
  F <: AbstractFloat
}
  i * log(μ / (λ + μ))
end

"""
    logtp_extinction(i, t, λ, μ)

Evaluate the logarithm of the probability of extinction for a simple birth and
death process.
"""
@inline function logtp_extinction(
  i::I,
  t::F,
  λ::F,
  μ::F
)::F where {
  I <: Integer,
  F <: AbstractFloat
}
  i * log(expm1((λ - μ) * t) / expm1(log(λ / μ) + (λ - μ) * t))
end

"""
    trans_prob(i, j, t, λ, μ; log_value)

Evaluate the transition probability of a simple birth and death process, i.e.
the probability of moving from `i` to `j` in `t` time when the birth rate is
equal to `λ` and the death rate is equal to `μ`.

The minimum amount of units `i` allowed at time 0 is 1 but `j` is allowed to be
0. Time `t`, birth rate `λ`, and death rate `μ` cannot be negative.

Let ``α = (μ e^{(λ - μ) t} - μ) / (λ e^{(λ - μ) t} - μ)`` and
``β = (λ e^{(λ - μ) t} - λ) / (λ e^{(λ - μ) t} - μ)``. Transition probability
is equal to (Bailey, 1964):
``\\sum_{h = 0}^{\\min(i, j)} \\binom{i}{h} \\binom{i + j - h - 1}{i - 1}
α^{i - h} β^{j - h} (1 - α - β)^{h}``.

# References:

Bailey, N. T. J. (1964). The elements of stochastic processes with applications
to the natural sciences. Wiley, New York, NY, USA. ISBN 0-471-04165-3.
"""
function trans_prob(
  i::I,
  j::I,
  t::Real,
  η::Vector{R};
  log_value::Bool=true
) where {
  I <: Integer,
  R <: Real
}
  if i < 1
    msg = "Initial population size 'i' must be greater than zero."
    throw(DomainError(i, msg))
  end

  if j < 0
    msg = "Final population size 'j' must be greater than or equal to zero."
    throw(DomainError(j, msg))
  end

  if t < 0
    msg = "Time 't' must be greater than or equal to zero."
    throw(DomainError(t, msg))
  end

  if η[1] < 0
    msg = "Birth rate 'η[1]' must be greater than or equal to zero."
    throw(DomainError(η[1], msg))
  end

  if η[2] < 0
    msg = "Death rate 'η[2]' must be greater than or equal to zero."
    throw(DomainError(η[2], msg))
  end

  # promote values to the appropriate float
  F = float(R)
  t, λ, μ = F(t), F(η[1]), F(η[2])
  ϵ = floatmin(F)

  log_trans_prob::F = if (t < ϵ) || ((λ < ϵ) && (μ < ϵ))
    if i == j
      0
    else
      -Inf
    end
  elseif λ ≈ μ
    ξ = 1 / λ

    if j > 0
      if t ≈ ξ
        logtp_equal_rates_equal_time(i, j)
      elseif t < ξ
        logtp_equal_rates(i, j, t, λ)
      else
        logtp_equal_rates_alternating(i, j, t, λ)
      end
    else
      if t ≈ ξ
        logtp_equal_rates_equal_time_extinction(i)
      else
        logtp_equal_rates_extinction(i, t, λ)
      end
    end
  elseif μ < ϵ
    if j >= i
      logtp_pure_birth(i, j, t, λ)
    else
      -Inf
    end
  elseif λ < ϵ
    if j <= i
      logtp_pure_death(i, j, t, μ)
    else
      -Inf
    end
  else
    ξ = log(λ / μ) / (λ - μ)

    if j > 0
      if t ≈ ξ
        logtp_equal_time(i, j, λ, μ)
      elseif t < ξ
        logtp_sum(i, j, t, λ, μ)
      else
        logtp_sum_alternating(i, j, t, λ, μ)
      end
    else
      if t ≈ ξ
        logtp_equal_time_extinction(i, λ, μ)
      else
        logtp_extinction(i, t, λ, μ)
      end
    end
  end

  if log_value
    log_trans_prob
  else
    exp(log_trans_prob)
  end
end


export NormalNormalTransformer

@doc raw"""
    NormalNormalTransformer(normal_sampler, new_params, old_params=(0.0, 1.0))

Generate normal random variables from another normal sampler.

Transform samples from a normal distribution sampler into samples from a normal
distribution with different parameters. This uses the result that if ``Z`` is a
standard normal random variable (``Z \sim \text{Normal}(0, 1)``) then
``X =  \sigma Z + \mu`` has distribution ``\text{Normal}(\mu, \sigma^2)``.

# Arguments
* `normal_sampler::sampler`: a sampler producing normally distributed samples.
* `new_params::Tuple{Real, Real}`: a tuple (mu, sigma2) giving the mean and
    variance of the target distribution
* `old_params::Tuple{Real, Real}=(0.0, 1.0)`: a tuple (mu, sigma2) giving the
    mean and variance of the normal sampler. Defaults to a standard normal
    distribution.

# Examples
```julia
# Transforming ten standard normal random variables into N(1, 2)
s = BoxMullerTransformSampler()
t = NormalNormalTransformer(s, (1, 2))
transformed_samples = sample(t, 10)
```
"""
struct NormalNormalTransformer <: Sampler
    normal_sampler::Sampler
    new_params::NamedTuple{(:mu, :sigma2), Tuple{Real, Real}}
    old_params::NamedTuple{(:mu, :sigma2), Tuple{Real, Real}}
    function NormalNormalTransformer(
        normal_sampler::Sampler,
        new_params::Tuple{Real, Real},
        old_params::Tuple{Real, Real}=(0.0, 1.0)
    )
        # Name parameter tuples
        new_params = (mu = new_params[1], sigma2 = new_params[2])
        old_params = (mu = old_params[1], sigma2 = old_params[2])
        # Validation
        if !(old_params[:sigma2] > 0 && new_params[:sigma2] > 0)
            error("variance must be positive")
        end
        new(normal_sampler, new_params, old_params)
    end
end

@doc raw"""
    sample(s::NormalNormalTransformer, N)

Transform normal samples into normal samples with new parameters.

# Arguments
* `s::NormalNormalTransformer`: a normal-normal transformer.
* `N::Integer`: the number of samples to generate.

# Returns
* `samples::Array{Float64}`: a 2-D array in which each column is a sample.
"""
function sample(s::NormalNormalTransformer, N::Integer)
    old_samples = sample(s.normal_sampler, N)
    std_samples = (old_samples .- s.old_params[:mu]) ./
        sqrt(s.old_params[:sigma2])
    new_samples = sqrt(s.new_params[:sigma2]) .* std_samples .+
        s.new_params[:mu]

    return new_samples
end


using Primes

# These primes are good example primes, as they have p = ±4 (mod 9), so taking cube roots is easy in F_p^2
plist = [5, 13, 23, 31, 41, 59, 67]
bigplist = [big(2)^107 - 1, big(2)^521 - 1]

# The prime used in the OG SIDH paper. Does not have above property^
OGp = big(2)^63*big(3)^41*11-1

# If p = 2 (mod 3) we have an easy supersingular curve
# E = Affine_EC_TwiHes(F, inv(4*one(F)), zero(F))
# Or twisted to
# E = Affine_EC_TwiHes(F, sqrt(F(b))*inv(4*one(F)), zero(F))
# This is isomorphic to the Weierstrass curve y^2 = x^3 + 1

# Similarly (slightly worse though) If p = 3 (mod 4) we have an easy supersingular curve
# y^2 = x^3 - x
# Translating this to twisted Hessian form will be future work

# If p = 1 (mod 12), just run away, dont look back.

#Scratch all this, much easier than that...
#If p = 2 (mod 3), then 0 is a supersingular j-invariant
#If p = 3 (mod 4), then 1728 is a supersingular j-invariant

#Example:
OGp = big(2)^63*big(3)^41*11-1
F = Fpsquare(OGp)
E = TwiHesEC(F, inv(4*one(F)), zero(F))
#w = F([1850222081870264162797200520152908562431, 2673953603825912900316077551548884232036])
#MAYBE: all points of order two are of the form (X : 1 : 1)
#Want to find roots of x^3-a = 0, where a = 3700444163740528325594401040305817124855
#One lies in F_p, the other two are the roots of x^2 + 3700444163740528325594401040305817124861*x + 4
#Roots are r1 = 3700444163740528325594401040305817124862 + 1647463043911297475037754062791951339209*i
#and r2 = 3700444163740528325594401040305817124862 + 2052981119829230850556646977513865785654*i
#r1 = F([3700444163740528325594401040305817124862, 1647463043911297475037754062791951339209])
#P = E(-r1, one(F), one(F))
#E0 = Isogeny(E,P).codomain
#c = F([3398654792751657841130893587213062545659])*w
#Q = E0(zero(F), -w, one(F))
#E1 = Isogeny(E0, Q).codomain

#Simple example with isogeny over F_p
F = Zmod(31)
E = TwiHesEC(F, F(1), F(7))
listP = allPointsNaive(E)
phi = Isogeny(E, E(F(13), F(1), F(1)))

#Simple example with isogeny over F_p^2
F = Fpsquare(31)
E = TwiHesEC(F, F([4, 5]), F([1,2]))
listP = allPointsNaive(E)
phi = Isogeny(E, E(F([9,19]), F([1]), F([1])))


"""
    Module Basis1D
General math tools for 1D bases used by DG
"""

module Basis1D

export gauss_lobatto_quad, gauss_quad
export jacobiP, grad_jacobiP
export vandermonde1D, grad_vandermonde1D
export chebyshev
export dPhi1D

using LinearAlgebra
using SpecialFunctions

"""
    gauss_lobatto_quad(degree, α=0, β=0)
Initialize Gauss-Legendre-Lobatto (GLL) quadrature nodes and weights (x,w)
to be able to exactly integrate (α,β) Jacobi polynomial of given degree
"""
function gauss_lobatto_quad(degree, α=0, β=0)
    # Need n points to integrate polynomials of degree 2n-3
    n = ceil(Int, (degree+3)/2)
    p = n-1
    if (α!=0) && (β!=0)
        error("currently only allows alpha == beta == 0")
    end
    x = zeros(p+1)
    w = zeros(p+1)
    if p == 0
        x[1] = 0
        w[1] = 2
    elseif p == 1
        x[1] = -1.0
        x[2] = 1.0
        w[1] = 1.0
        w[2] = 1.0
    else
        xint, w = gauss_quad(2*p-3, α+1, β+1)
        x = [-1 transpose(xint) 1]

        V = vandermonde1D(p,x)
        w = vec(sum(inv(V*V'),dims=2))
    end
    return x[:],w[:]
end

"""
    gauss_quad(degree, α, β)
Initialize Gaussian quadrature nodes and weights (x,w) to be able to exactly
integrate Jacobi Polynomial (α,β) of given degree
"""
function gauss_quad(degree, α=0, β=0)
    # Need n points to integrate polynomials of degree 2n-1
    n = ceil(Int, (degree+1)/2)
    p = n-1
    if p == 0
        x = [-(α-β)/(α+β+2)]
        w = [2]
        return x, w
    end

    J = zeros(p+1, p+1)
    h₁ = @. 2*(0:p)+α+β
    J = diagm(0 => (@. -1/2*(α^2-β^2)/(h₁+2)/h₁), 1 => (@. 2/(h₁[1:p]+2)*sqrt((1:p)*((1:p)+α+β)*((1:p)+α)*((1:p)+β)/(h₁[1:p]+1)/(h₁[1:p]+3))))
    if α+β<10*eps()
        J[1,1] = 0.0
    end
    J = J + transpose(J) # Finalize symmetric tridiagonal Jacobi matrix

    x, V = eigen(J)
    w = @. transpose(V[1,:])^2*2^(α+β+1)/(α+β+1)*gamma(α+1)*gamma(β+1)/gamma(α+β+1)
    return x[:], w[:]
end

"""
    grad_jacobiP(p, r, α, β)
Evaluate the derivative of Jacobi Polynomial (α, β) of order p at nodes r
"""
function grad_jacobiP(p, r, α, β)
    dP = zeros(length(r))
    if p != 0
        dP = sqrt(p*(p+α+β+1))*jacobiP(p-1,r,α+1,β+1)
    end
    return dP
end

"""
    jacobiP(p, x, α, β)
Evaluate Jacobi Polynomial (α, β) of order p at points x
"""
function jacobiP(p, x, α, β)
    xp = x
    if size(xp, 2) == 1
        xp = transpose(xp)
    end

    PL = zeros(p+1,length(xp))
    γ₀ = 2^(α+β+1)/(α+β+1)*gamma(α+1)*gamma(β+1)/gamma(α+β+1)
    PL[1,:] .= 1.0/sqrt(γ₀)
    if p == 0
        P = transpose(PL)
        return P
    end
    γ₁ = (α+1)*(β+1)/(α+β+3)*γ₀
    PL[2,:] = ((α+β+2).*xp/2 .+ (α-β)/2)/sqrt(γ₁)
    if p == 1
        P = PL[p+1,:]
        return P
    end

    aold = 2/(2+α+β)*sqrt((α+1)*(β+1)/(α+β+3))

    for i = 1:p-1
        h₁ = 2i+α+β
        anew = 2/(h₁+2)*sqrt((i+1)*(i+1+α+β)*(i+1+α)*(i+1+β)/(h₁+1)/(h₁+3))
        bnew = -(α^2-β^2)/h₁/(h₁+2)
        PL[i+2,:] = 1/anew*(-aold*transpose(PL[i,:]).+(xp.-bnew).*transpose(PL[i+1,:]))
        aold = anew
    end

    P = PL[p+1,:]
    return P;
end

"""
    vandermonde1D(p, r)
Initialize the 1D Vandermonde matrix of order p Legendre polynomials at nodes r
"""
function vandermonde1D(p, r)
    V1D = zeros(length(r), p+1)
    for j = 1:p+1
        V1D[:,j] = jacobiP(j-1, r[:], 0, 0)
    end
    return V1D
end

"""
    grad_vandermonde1D(p, r)
Initialize the 1D gradient Vandermonde matrix of order p Legendre polynomials at nodes r
"""
function grad_vandermonde1D(p, r)
    V1D = zeros(length(r), p+1)
    for j = 1:p+1
        V1D[:,j] = grad_jacobiP(j-1, r[:], 0, 0)
    end
    return V1D
end

"""
    chebyshev(p)
Return the 1D Chebyshev nodes of order p on [-1,1]
"""
function chebyshev(p)
    cheby = cos.((p:-1:0)*π/p)
end

"""
    interpolation_matrix1D(x_from, x_to)
Compute an interpolation matrix from a set of 1D points to another set of 1D points.
Assumes that points interpolating from provide enough accuracy (aka - they are
well spaced out and of high enough order), and define an interval. Points
interpolating onto can be of any size, but must be defined on this same interval.
Interpolation matrix ∈ ℜ^(size of x_to x size of x_from)
"""
function interpolation_matrix1D(x_from, x_to)
    # Create nodal representation of reference bases
    order = size(x_from,1) - 1 # Assumes order = (size of x_from) - 1
    n_from = size(x_from,1)
    l_from = vandermonde1D(order, x_from)

    eye = diagm(0=>ones(n_from))
    V = reshape(l_from, n_from,n_from)
    coeffs_phi = V \ eye

    # Compute reference bases on the output points
    l_to = vandermonde1D(order, x_to)
    n_to = size(l_to,1)
    V_to = reshape(l_to, n_to,n_from)

    # Construct interpolation matrix
    Interp1D = V_to*coeffs_phi
end

"""
    dPhi1D(x_from, x_to)
Compute the gradient of the basis functions defined by 1D points on another
set of 1D points.

Assumes that points interpolating from provide enough accuracy (aka - they are
well spaced out and of high enough order), and define an interval. Points
interpolating onto can be of any size, but must be defined on this same interval.
dPhi ∈ ℜ^(size of x_to × size of x_from × 1)
"""
function dPhi1D(x_from, x_to)
    dim = 1
    # Create nodal representation of reference bases
    order = size(x_from,1) - 1 # Assumes order = (size of x_from) - 1
    n_from = size(x_from,1)
    l_from = vandermonde1D(order, x_from)

    eye = diagm(0=>ones(n_from))
    V = reshape(l_from, n_from,n_from)
    coeffs_phi = V \ eye

    # Compute derivative of reference bases on the output points
    dl_to = grad_vandermonde1D(order, x_to)
    n_to = size(dl_to,1)
    dV = reshape(dl_to, n_to,n_from,dim)
    # Construct gradient of phi = dV*coeffs_phi
    dPhi_to = Array{Float64,3}(undef, n_to,n_from,dim)
    for l = 1:dim
        dPhi_to[:,:,l] = dV[:,:,l]*coeffs_phi
    end
    return dPhi_to
end

end


function integrate(f, a, b, n)
	h = (b - a) / n
	x = zeros(1, 3)
	x[1] = f(a) + f(b)
	for i ∈ 1:(n - 1)
		X = a + i*h
		if iseven(i)
			x[3] += f(X)
		else
			x[2] += f(X)
		end
	end
	return h*(x[1] + 2*x[3] + 4*x[2])/3
end

f(x) = x
println(integrate(f, 0, 3, 100))


module SymRCM

# (C) 2020, Petr Krysl

export symrcm

using SparseArrays

"""
    adjgraph(A)

Compute the adjacency graph from a sparse matrix.

The sparse matrix `A` is assumed to be symmetric.
The results will be wrong if it isn't.

- `sortbydeg`: Should the neighbor lists be sorted by column degree? The default is
  `true`, but often results of very similar quality are obtained when this is
  set to `false` and the lists are not sorted. The second option is much
  faster, as the sorting is expensive.
"""
function adjgraph(A::SparseMatrixCSC; sortbydeg = true)
    colptr = A.colptr
    rowval = A.rowval
    ncols = length(colptr)-1
    neighbors = Vector{Vector{eltype(colptr)}}(undef, ncols)
    cdeg = diff(colptr) # the degree is colptr[j+1]-colptr[j]
    for j in 1:ncols
        cstart = colptr[j]
        jdeg = cdeg[j]
        neighbors[j] = [rowval[cstart+m-1] for m in 1:jdeg]
    end
    # All of these sorts can be done in parallel,  they are totally independent.
    # The question is when to switch over to parallel execution so as to
    # amortize the cost of starting up threads.
    if sortbydeg
        for j in 1:ncols
            sort!(neighbors[j], by = j -> cdeg[j])
        end
    end
    return neighbors
end

"""
    nodedegrees(adjgr::Vector{Vector{Int}})

Compute the degrees of the nodes in the adjacency graph.

conn = [9 1 8 4;
       1 3 2 8;
       8 2 7 5;
       2 6 7 7];
nfens = 9;
adjgr = adjgraph(conn, nfens)
nodedegrees(adjgr)

julia> degrees = node_degrees(adjgr)
9-element Array{Int64,1}:
 5
 6
 3
 3
 3
 2
 4
 7
 3
"""
function nodedegrees(adjgr::Vector{Vector{T}}) where {T}
    degrees = fill(0, length(adjgr))
    for k = 1:length(adjgr)
        degrees[k] = length(adjgr[k])
    end
    return degrees
end

"""
    symrcm(adjgr::Vector{Vector{T}}, degrees::Vector{T}) where {T}

Reverse Cuthill-McKee node-renumbering algorithm.
"""
function symrcm(adjgr::Vector{Vector{T}}, degrees::Vector{T}) where {T}
    @assert length(adjgr) == length(degrees)
    # Initialization
    n = length(adjgr)
    ndegperm = sortperm(degrees) # sorted nodal degrees
    inR = fill(false, n) # Is a node in the result list?
    inQ = fill(false, n) # Is a node in the queue?
    R = T[]
    sizehint!(R, n)
    Q = T[] # Node queue
    sizehint!(Q, n)
    while true
        P = zero(T) # Find the next node to start from
        while !isempty(ndegperm)
            i = popfirst!(ndegperm)
            if !inR[i]
                P = i
                break
            end
        end
        if P == zero(T)
            break # That was the last node
        end
        # Now we have a node to start from: put it into the result list
        push!(R, P); inR[P] = true
        empty!(Q) # empty the queue
        append!(Q, adjgr[P]); inQ[adjgr[P]] .= true # put adjacent nodes in queue
        while length(Q) >= 1
            C = popfirst!(Q) # child to put into the result list
            inQ[C] = false # make note: it is not in the queue anymore
            if !inR[C]
                push!(R, C); inR[C] = true
            end
            for i in adjgr[C] # add all adjacent nodes into the queue
                if (!inR[i]) && (!inQ[i]) # contingent on not being in result/queue
                    push!(Q, i); inQ[i] = true
                end
            end
        end
    end
    return reverse(R) # reverse the result list
end

"""
    symrcm(A::SparseMatrixCSC; sortbydeg = true)

Reverse Cuthill-McKee node-renumbering algorithm.

Compute the adjacency graph from a sparse matrix. The sparse matrix `A` is
assumed to be symmetric. The results will be wrong if it isn't.

- `sortbydeg`: Should the neighbor lists be sorted by column degree? The default is
  `true`, but often results of very similar quality are obtained when this is
  set to `false` and the lists are not sorted. The second option can be much
  faster, as the sorting is expensive when the neighbor lists are long.
"""
function symrcm(A::SparseMatrixCSC; sortbydeg = true)
    ag = adjgraph(A; sortbydeg = sortbydeg)
    nd = nodedegrees(ag)
    return symrcm(ag, nd)
end

end # module


using JuMP
using Clp

m = Model(Clp.Optimizer)

@variable(m, x₁ ≥ 0)
@variable(m, x₂ ≥ 0)
@variable(m, x₃ ≥ 0)

@objective(m, Max, 2x₁  + 4x₂ + 3x₃)
@constraint(m, 3x₁  + 4x₂ + 2x₃ ≤ 60)
@constraint(m, 2x₁  + x₂ + 2x₃ ≤ 40)
@constraint(m, x₁  + 3x₂ + 2x₃ ≤ 80)

print(m)

optimize!(m)

status = termination_status(m)

println("Solution status: ", status)

println("Objective value: ", objective_value(m))
println("x₁ = ", value(x₁))
println("x₂ = ", value(x₂))
println("x₃ = ", value(x₃))


using Distributions, LaTeXStrings, Random, Plots
Random.seed!(1234)

# Simulate n coin flips with a bias of 0.6
n = 1000
flips = rand(n) .< 0.6
p = sum(flips)
q = n - p

# Create plot of Beta(α,β) density
α = 10
β = 5

x_values = 0:0.001:1

# Get prior and posterior values
prior = [pdf(Beta(α,β),x) for x in x_values]
posterior = [pdf(Beta(α+p,β+q),x) for x in x_values]

# Create plots
plot(x_values, posterior, seriestype = :line, label = "posterior",
 yaxis = L"f(\theta|y)", xaxis = L"\theta", grid = false)

plot!(x_values, prior, seriestype = :line, label = "prior", grid = false)

posterior_quantile = cdf(Beta(α+p,β+q),0.65) - cdf(Beta(α+p,β+q),0.55)
confidence_interval = [quantile(Beta(α+p,β+q),0.025), quantile(Beta(α+p,β+q),0.975)]
bayes_mean = mean(Beta(α+p,β+q))

plot(x_values, normal_values)
savefig("beta_prior_coin_ex_2.svg")


using StatsBase
using Plots

Parameters = [[[0 0 ; 0 0.16], [0 ; 0]],
        [[0.85 0.04 ; -0.04 0.85], [0 ; 1.6]],
        [[0.20 -0.26 ; 0.23 0.22], [0 ; 1.6]],
        [[-0.15 0.28 ; 0.26 0.24], [0 ; 0.44]]]

f1(vector, Pars) = Pars[1][1] * vector + Pars[1][2]
f2(vector, Pars) = Pars[2][1] * vector + Pars[2][2]
f3(vector, Pars) = Pars[3][1] * vector + Pars[3][2]
f4(vector, Pars) = Pars[4][1] * vector + Pars[4][2]

function RandomBarnsleyFern(P, PointNums)
    Functions = [f1, f2, f3, f4]
    X = [0.0 for i = 1:PointNums]
    Y = [0.0 for i = 1:PointNums]
    for i in 1:PointNums
        point = [rand() ; rand()]
        for j in 1:P
            s = sample([1,2,3,4], Weights([0.01,0.85,0.07,0.07]))
            point = Functions[s](point, Parameters)
        end
        X[i]= point[1]
        Y[i]= point[2]
    end
    return X, Y
end

for i in [10, 100]
    for j in [10000,1000000]
        x,y = RandomBarnsleyFern(i, j)

        scatter(x,y,markersize = 10^4/j, legend = false, border=:none, color =:green,dpi=200, label = "Points number = $j Operations number = $i")
        savefig("C:\\Users\\Yaghoub\\Documents\\GitHub\\ComputationalPhysics-Fall2021\\ProblemSet1\\Figs\\Q6\\RandomBarnsleyFern-p$(i)num$(j).png")
    end
end


# ## Quiver plots

#md # ![FILE_NAME.png](images/FILE_NAME.png)

n = 7
f(x,y) = 1 / hypot(x, y)
x = repeat(-3:(2*3)/n:3, 1, n) |> vec
y = repeat(-3:(2*3)/n:3, 1, n)' |> vec
vx = f.(x,y) .* cos.(atan.(y,x)) |> vec
vy = f.(x,y) .* sin.(atan.(y,x)) |> vec

quiver(x, y, quiver=(vx, vy), aspect_ratio=:equal)


"""
Return the rotation angle around the
reference axis for the first vector to be
aligned to the second vector.

Arguments
----------
:Type{ProjectedRotationAngle}
    must be type ``ProjectedRotationAngle``

input:AbstractArray
    the vector to be rotated

target:AbstractArray
    the target vector to be aligned to

ref_axis:AbstractArray
    the reference axis around which input vector will be rotated

tol_near_zero=1e-15:AbstractFloat
    (optional) tolerance for the length of the input vectors to be close to zero
"""
function gage(::Type{ProjectedRotationAngle}, input::AbstractArray, target::AbstractArray,
    ref_axis::AbstractArray, tol_near_zero::AbstractFloat=1e-15)

    @debug @assert length(input) == length(target)

    ref_axis_magnitude = norm(ref_axis, 2)
    input_magnitude = norm(input, 2)
    target_magnitude = norm(target, 2)

    if ref_axis_magnitude < tol_near_zero # this means the input vector is the origin
        return 0. # no need to rotate
    else
        unit_ref_axis = ref_axis / ref_axis_magnitude
    end

    if input_magnitude < tol_near_zero # this means the input vector is the origin
        return 0. # no need to rotate
    else
        unit_input = input / input_magnitude
    end

    if target_magnitude < tol_near_zero # this means the target vector is the origin
        return 0. # no need to rotate
    else
        unit_target = target / target_magnitude
    end

    # get the projection along a direction normal to the ref_axis
    function normal_projection(input::AbstractArray, unit_ref_axis::AbstractArray)
        return input - dot(input, unit_ref_axis) * unit_ref_axis
    end

    return gage(RotationAngle, normal_projection(unit_input, unit_ref_axis),
        normal_projection(unit_target, unit_ref_axis))
end


module DiffSolve

export eulersolve, eulercromer, eulermid, verlet, velverlet, rk4

"""
Solve a system of first order ODEs dx/dt = f(t, x(t)) using the Euler method.
x(t) is a vector.

# Arguments
- `init`: initial x vector
- `func`: the function dx/dt = f(t, x(t))
- `start`: initial t
- `stop`: final t

# Returns
vector of all t and matrix containing x at all t (the (i, j) element in the matrix
represents the jth variable in the x vector x[j] at the ith time step t[i]).
"""
function eulersolve(func::Function, init::Vector{Float64}, start::Float64, stop::Float64;
        step::Float64 = 0.01)
    ts = collect(range(start, stop, step=step)) # exclucive range
    xs = Matrix{Float64}(undef, length(ts), length(init))
    xs[1, :] = init

    x = init
    for (i, t) in enumerate(@view ts[1:end-1])
        x += step * func(t, x)
        xs[i + 1, :] = x
    end

    return ts, xs
end

"""
Solve a system of first order ODEs dx/dt = f(t, x(t)) using the 4th order Runge–Kutta
method (a.k.a. RK4, or Classic Runge–Kutta). x(t) is a vector.

# Arguments
- `init`: initial x vector
- `func`: the function dx/dt = f(t, x(t))
- `start`: initial t
- `stop`: final t

# Returns
vector of all t and matrix containing x at all t (the (i, j) element in the matrix
represents the jth variable in the x vector x[j] at the ith time step t[i]).
"""
function rk4(func::Function, init::Vector{Float64}, start::Float64, stop::Float64;
        step::Float64 = 0.01)
    ts = collect(range(start, stop, step=step)) # exclucive range
    xs = Matrix{Float64}(undef, length(ts), length(init))
    xs[1, :] = init

    x = init
    for (i, t) in enumerate(@view ts[1:end-1])
        k1 = func(t, x)
        k2 = func(t + step / 2, x .+ step / 2 * k1)
        k3 = func(t + step / 2, x .+ step / 2 * k2)
        k4 = func(t + step, x .+ step * k3)

        x += step / 6 * (k1 + 2k2 + 2k3 + k4)
        xs[i + 1, :] = x
    end

    return ts, xs
end

"""
Solve a second order ODE d²x/dt² = a(t, x(t)) using the Euler–Cromer method
(a.k.a. the semi-implicit euler method, or the Newton–Størmer–Verlet).

# Arguments
- `initx`, `initv`: initial x (displacement) and dx/dt (velocity v)
- `acceleration`: the function d²x/dt² = a(t, x(t))
- `start`: initial t
- `stop`: final t

# Returns
t, x, and v (dx/dt) vectors.
"""
function eulercromer(acceleration::Function, initx::Float64, initv::Float64,
        start::Float64, stop::Float64; step::Float64 = 0.01)
    ts = collect(range(start, stop, step=step)) # exclucive range
    xs = Vector{Float64}(undef, length(ts))
    vs = Vector{Float64}(undef, length(ts))
    xs[1] = initx
    vs[1] = initv

    x = initx
    v = initv
    for (i, t) in enumerate(@view ts[1:end-1])
        v += step * acceleration(t, x)
        x += step * v

        xs[i + 1] = x
        vs[i + 1] = v
    end

    return ts, xs, vs
end

"""
Solve a second order ODE d²x/dt² = a(t, x(t)) using the Midpoint method
(a modified euler method).

# Arguments
- `initx`, `initv`: initial x (displacement) and dx/dt (velocity v)
- `acceleration`: the function d²x/dt² = a(t, x(t))
- `start`: initial t
- `stop`: final t

# Returns
t, x, and v (dx/dt) vectors. Due to the way this method works, the v vector is the
velocity at midpoints between the time steps.
"""
function eulermid(acceleration::Function, initx::Float64, initv::Float64,
        start::Float64, stop::Float64; step::Float64 = 0.01)
    ts = collect(range(start, stop, step=step)) # exclucive range
    xs = Vector{Float64}(undef, length(ts))
    vs = Vector{Float64}(undef, length(ts))
    xs[1] = initx
    vs[1] = initv - acceleration(start, initx) * step / 2

    x = initx
    v = vs[1]
    for (i, t) in enumerate(@view ts[1:end-1])
        v += step * acceleration(t, x)
        x += step * v

        xs[i + 1] = x
        vs[i + 1] = v
    end

    return ts, xs, vs
end

"""
Solve a second order ODE d²x/dt² = a(t, x(t)) using Verlet integration
(a.k.a. the Störmer–Verlet method).

# Arguments
- `initx`, `initv`: initial x (displacement) and dx/dt (velocity v)
- `acceleration`: the function d²x/dt² = a(t, x(t))
- `start`: initial t
- `stop`: final t

# Returns
t, x, and v (dx/dt) vectors.
"""
function verlet(acceleration::Function, initx::Float64, initv::Float64,
        start::Float64, stop::Float64; step::Float64 = 0.01)
    ts = collect(range(start, stop, step=step)) # exclucive range
    xs = Vector{Float64}(undef, length(ts))
    vs = Vector{Float64}(undef, length(ts))

    xs[1] = initx
    vs[1] = initv
    xs[2] = initx + initv * step + acceleration(start, initx) * step * step / 2

    for i in 2:length(xs)-1
        xs[i + 1] = 2 * xs[i] - xs[i - 1] + acceleration(ts[i], xs[i]) * step * step
        vs[i + 1] = (xs[i + 1] - xs[i]) / step
    end

    return ts, xs, vs
end

"""
Solve a second order ODE d²x/dt² = a(t, x(t)) using the velocity Verlet method

# Arguments
- `initx`, `initv`: initial x (displacement) and dx/dt (velocity v)
- `acceleration`: the function d²x/dt² = a(t, x(t))
- `start`: initial t
- `stop`: final t

# Returns
t, x, and v (dx/dt) vectors.
"""
function velverlet(acceleration::Function, initx::Float64, initv::Float64,
        start::Float64, stop::Float64; step::Float64 = 0.01)
    ts = collect(range(start, stop, step=step)) # exclucive range
    xs = Vector{Float64}(undef, length(ts))
    vs = Vector{Float64}(undef, length(ts))

    xs[1] = initx
    vs[1] = initv

    for i in 1:length(ts)-1
        accbefore = acceleration(ts[i], xs[i])
        xs[i + 1] = xs[i] + vs[i] * step + accbefore * step * step / 2
        accafter = acceleration(ts[i + 1], xs[i + 1])
        vs[i + 1] = vs[i] + (accbefore + accafter) * step / 2
    end

    return ts, xs, vs
end

end


using SpecialFunctions
using Base.MathConstants
using DataFrames, CSV
using Printf

"""
    gamma_derivatives()

Derive a table of derivatives of the Gamma function at 1, i.e., Gamma^{(m)}(a)

This code is used to create the results, which are included in the main code via a lookup table.

"""
function gamma_derivatives(k::Integer, x::Float64)
    #    general form from: http://erikerlandson.github.io/blog/2016/06/15/computing-derivatives-of-the-gamma-function/
    return SpecialFunctions.gamma(x) * D(k,0,x)
end

function D(k::Integer, n::Integer, x::Float64)
    # general form from: http://erikerlandson.github.io/blog/2016/06/15/computing-derivatives-of-the-gamma-function/
    if n<0
        throw(DomainError(n))
    elseif k<0
        throw(DomainError(k))
    end
    if k==0
        return 1
    elseif k==1
        return SpecialFunctions.polygamma(n,x)
    else
        total = D(k-1,n+1,x)
        for j=0:n
            total += binomial(n,j) * D(1,j,x) * D(k-1,n-j,x)
        end
        return total
    end
end

function g1(t::Int) # Crandall,2012, p.17
    # t derivate of Gamma function at 1
    PG21 = -2.4041138063191902 # SpecialFunctions.polygamma(2,1)
    PG41 = -24.8862661234409   # SpecialFunctions.polygamma(4,1)
    if t==0
        return 1.0
    elseif t==1
        return -γ
    elseif t==2
        # return γ^2 - γ + pi^2/6, Crandall seems to be wrong on this
        return γ^2 + pi^2/6  # from Mathematica
    elseif t==3
        # return -2*γ^3 + 9*γ^2 - (π^2+6)*γ + 3*π^2/2 - 4*SpecialFunctions.zeta(3), Crandall seems to be wrong on this
        return -γ^3  - π^2*γ/2 + PG21 # from Mathematica
    elseif t==4
        return γ^4 + γ^2*π^2 + 3*π^4/20 - 4*γ*PG21 # from Mathematica
    elseif t==5
        return -γ^5 - (20/12)*γ^3*π^2 - (9/12)*γ*π^4 + (10*γ^2 + (20/12)*π^2)*PG21 + PG41 # from Mathematica
    else
        return NaN
    end
end

function g2(t::Int) # http://erikerlandson.github.io/blog/2016/06/15/computing-derivatives-of-the-gamma-function/
    # t derivate of Gamma function at 1
    return gamma_derivatives(t,1.0)
end

# #### output a table
# K = 14
# k = collect(0:K)
# G1 = zeros(Float64, K+1)
# G2 = zeros(Float64, K+1)
# for i=1:K+1
#     G1[i] = g1(i-1)
#     G2[i] = g2(i-1)
#     @printf("G1[i] = %.20f, G2[i] = %.20f\n", G1[i], G2[i])
# end

# df = DataFrame( k=k, G1=G1, G2=G2 )
# CSV.write("../data/gamma_derivatives.csv", df)


# collatz conjenture

#Jl Plots

function collatz_next(n)
    if n % 2 == 0
        return n ÷ 2
    else
        return 3n+1
    end
end

function collatz_path(n)
    steps = Vector{Int}()
    while n != 1
        push!(steps, n)
        n = collatz_next(n)
    end
    push!(steps, 1)
    return steps
end

using Plots

x = 1:10000
y = length.(collatz_path.(x))
Plots.scatter(x, y, label="Stopping time", legend=:topleft, size=(800,600))

x = 2 .^ (0:13)
y = length.(collatz_path.(x))
Plots.scatter!(x, y, color=:red, label="for powers of two")


module DTWDTfunctions
export test, computeErrorFunction, accumulateErrorFunction, backtrackDistanceFunction, computeDTWerror

"""
 USAGE: err = computeErrorFunction( u1, u0, nSample, lag )

 INPUT:

   u1      = trace that we want to warp; size = (nsamp,1)

   u0      = reference trace to compare with: size = (nsamp,1)

   nSample = numer of points to compare in the traces

   lag     = maximum lag in sample number to search

   norm    = 'L2' or 'L1' (default is 'L2')

 OUTPUT:

    err = the 2D error function; size = (nsamp,2*lag+1)

 The error function is equation 1 in Hale, 2013.

 You could umcomment the L1 norm and comment the L2 norm if you want on Line 29

 Original by Di Yang
 Last modified by Dylan Mikesell (25 Feb. 2015)
"""
function computeErrorFunction(u1::Union{Array{Float32,1},Array{Float64,1}}, u0::Union{Array{Float32,1},Array{Float64,1}}, nSample::Int, lag::Int; norm::String="L2")

    if lag >= nSample
        error("computeErrorFunction:lagProblem ","lag must be smaller than nSample");
    end

    #using JLD2
    #@load "../exampledata/sineShiftData.jld2"
    #nSample = length(u0)
    #lag = 80
    #norm = "L2"

    # Allocate error function variable
    err = zeros(Float64, nSample, 2 * lag + 1 );

    #--------------------------------------------------------------------------
    # initial error calculation
    for ll = -lag:lag # loop over lags

        thisLag = ll + lag + 1;

        for ii = 1:nSample # loop over samples

            if ( ii + ll >= 1 && ii + ll <= nSample ) # skip corners for now, we will come back to these

                diff = u1[ii] - u0[ii + ll]; # sample difference

                if norm == "L2"
                        err[ii, thisLag] = diff^2; # difference squared error
                elseif norm == "L1"
                        err[ii, thisLag] = abs(diff); # absolute value errors
                else
                    error("norm type is not defined.")
                end

            end

        end

    end

    #--------------------------------------------------------------------------
    # Now fix corners with constant extrapolation
    for ll = -lag:lag # loop over lags

        thisLag = ll + lag + 1;

        for ii = 1:nSample # loop over samples

            if ( ii + ll < 1 ) # lower left corner (negative lag, early time)

                err[ii, thisLag] = err[-ll + 1, thisLag];

            elseif ( ii + ll > nSample ) # upper right corner (positive lag, late time)

                err[ii, thisLag] = err[nSample - ll, thisLag];

            end

        end

    end

    return err
end

"""

 USAGE: d = accumulation_diw_mod( dir, err, nSample, lag, b )

 INPUT:

   dir = accumulation direction ( dir > 0 = forward in time, dir <= 0 = backward in time)

   err = the 2D error function; size = (nsamp,2*lag+1)

   nSample = numer of points to compare in the traces

   lag = maximum lag in sample number to search

   b = strain limit (integer value >= 1)

 OUTPUT:

    d = the 2D distance function; size = (nsamp,2*lag+1)

 The function is equation 6 in Hale, 2013.

 Original by Di Yang Last modified by Dylan Mikesell (25 Feb. 2015)
"""
function accumulateErrorFunction(dir::Int, err::Array{Float64,2}, nSample::Int, lag::Int, b::Int)

    #nSample = length(u0)
    #lag = 80
    #dir = -1

    nLag = (2 * lag ) + 1; # number of lags from [ -lag : +lag ]

    # allocate distance matrix
    d = zeros(Float64, nSample, nLag);

    #--------------------------------------------------------------------------
    # Setup indices based on forward or backward accumulation direction
    #--------------------------------------------------------------------------
    if dir > 0            # FORWARD
        iBegin = 1;       # start index
        iEnd   = nSample; # end index
        iInc   = 1;       # increment
    else                  # BACKWARD
        iBegin = nSample; # start index
        iEnd   = 1;       # stop index
        iInc   = -1;      # increment
    end
    #--------------------------------------------------------------------------
    # Loop through all times ii in forward or backward direction

    for ii = iBegin:iInc:iEnd

        # min/max to account for the edges/boundaries
        ji = max(1, min(nSample, ii - iInc ));     # i-1 index
        jb = max(1, min(nSample, ii - iInc * b )); # i-b index

        # loop through all lags l
        for ll = 1:nLag

            # -----------------------------------------------------------------
            # check limits on lag indices
            lMinus1 = ll - 1; # lag at l-1

            if lMinus1 < 1  # check lag index is greater than 1
                lMinus1 = 1; # make lag = first lag
            end

            lPlus1 = ll + 1; # lag at l+1

            if lPlus1 > nLag # check lag index less than max lag
    #             lPlus1 = nLag - 1; # D.Y. version
                lPlus1 = nLag; # D.M. version
            end
            # -----------------------------------------------------------------

            # get distance at lags (ll-1, ll, ll+1)
            distLminus1 = d[jb, lMinus1]; # minus:  d( i-b, j-1 )
            distL       = d[ji, ll];      # actual: d( i-1, j   )
            distLplus1  = d[jb, lPlus1];  # plus:   d( i-b, j+1 )

            if ji != jb # equation 10 in Hale (2013)
                for kb = ji:-iInc:jb+iInc # sum errors over i-1:i-b+1
                    distLminus1 = distLminus1 + err[kb, lMinus1];
                    distLplus1  = distLplus1  + err[kb, lPlus1];
                end
            end

            # equation 6 (if b=1) or 10 (if b>1) in Hale (2013) after treating boundaries
            d[ii, ll] = err[ii, ll] + min(distLminus1, distL, distLplus1);
        end
    end
    return d
end

"""

 USAGE: stbar = backtrackDistanceFunction( dir, d, err, lmin, b )

 INPUT:

   dir   = side to start minimization ( dir > 0 = front, dir <= 0 =  back)

   d     = the 2D distance function; size = (nsamp,2*lag+1)

   err   = the 2D error function; size = (nsamp,2*lag+1)

   lmin  = minimum lag to search over

   b     = strain limit (integer value >= 1)

 OUTPUT:

   stbar = vector of integer shifts subject to |u(i)-u(i-1)| <= 1/b

 The function is equation 2 in Hale, 2013.

 Original by Di Yang Last modified by Dylan Mikesell (19 Dec. 2014)

"""
function backtrackDistanceFunction(dir::Int, d::Array{Float64,2}, err::Array{Float64,2}, lmin::Int, b::Int)

    #d = dist
    #dir = -1
    #lmin = -maxLag
    #b = 1

    nSample = size(d,1); # number of samples
    nLag    = size(d,2); # number of lags
    stbar   = zeros(Int64, nSample); # allocate

    #--------------------------------------------------------------------------
    # Setup indices based on forward or backward accumulation direction
    #--------------------------------------------------------------------------
    if dir > 0            # FORWARD
        iBegin = 1;       # start index
        iEnd   = nSample; # end index
        iInc   = 1;       # increment
    else                  # BACKWARD
        iBegin = nSample; # start index
        iEnd   = 1;       # stop index
        iInc   = -1;      # increment
    end
    #--------------------------------------------------------------------------
    # start from the end (front or back)
    ll0 = argmin(d[iBegin,:]); # find minimum accumulated distance at front or back depending on 'dir'
    stbar[iBegin] = ll0 + lmin - 1; # absolute value of integer shift
    #--------------------------------------------------------------------------
    # move through all time samples in forward or backward direction
    ii = iBegin;

    while ii != iEnd
        if ii == iBegin
            ll = ll0;
        else
            ll = ll_next
        end

        # min/max for edges/boundaries
        ji = max( 1, min(nSample, ii + iInc) );
        jb = max( 1, min(nSample, ii + iInc * b) );

        # -----------------------------------------------------------------
        # check limits on lag indices

        lMinus1 = ll - 1; # lag at l-1

        if lMinus1 < 1 # check lag index is greater than 1
            lMinus1 = 1; # make lag = first lag
        end

        lPlus1 = ll + 1; # lag at l+1

        if lPlus1 > nLag # check lag index less than max lag
            lPlus1 = nLag; # D.M. and D.Y. version
        end
        # -----------------------------------------------------------------
        # get distance at lags (ll-1, ll, ll+1)
        distLminus1 = d[jb, lMinus1]; # minus:  d( i-b, j-1 )
        distL       = d[ji, ll];      # actual: d( i-1, j   )
        distLplus1  = d[jb, lPlus1];  # plus:   d( i-b, j+1 )

        if ji != jb # equation 10 in Hale (2013)
            for kb = ji:iInc:jb-iInc # sum errors over i-1:i-b+1
                distLminus1 = distLminus1 + err[kb, lMinus1];
                distLplus1  = distLplus1  + err[kb, lPlus1];
            end
        end

        dl = min(distLminus1, distL, distLplus1); # update minimum distance to previous sample

        if ( dl != distL ) # then ll != ll and we check forward and backward
            if ( dl == distLminus1 )
                global ll_next = lMinus1;
            else # ( dl == lPlus1 )
                global ll_next = lPlus1;
            end
        else
            ll_next = ll
        end

        # assume ii = ii - 1
        ii += iInc; # previous time sample

        stbar[ii] = ll_next + lmin - 1; # absolute integer of lag
        # now move to correct time index, if smoothing difference over many
        # time samples using 'b'

        if ( ll_next == lMinus1 || ll_next == lPlus1 ) # check edges to see about b values
            if ( ji != jb ) # if b>1 then need to move more steps
                for kb = ji:iInc:jb - iInc
                    ii = ii + iInc; # move from i-1:i-b-1
                    stbar[ii] = ll_next + lmin - 1; # constant lag over that time
                end
            end
        end
    end

    return stbar

    #------------------------#
    # while loop is tricky with Julia
    # we modified ll to ll_next
    # this is the original source from dylanmikesell
    #------------------------#
    # ii = iBegin;
    # while (ii ~= iEnd)
    #
    #     % min/max for edges/boundaries
    #     ji = max( 1, min( [ nSample, ii + iInc ] ) );
    #     jb = max( 1, min( [ nSample, ii + iInc * b ] ) );
    #
    #     % -----------------------------------------------------------------
    #     % check limits on lag indices
    #     lMinus1 = ll - 1; % lag at l-1
    #
    #     if lMinus1 < 1 % check lag index is greater than 1
    #         lMinus1 = 1; % make lag = first lag
    #     end
    #
    #     lPlus1 = ll + 1; % lag at l+1
    #
    #     if lPlus1 > nLag % check lag index less than max lag
    #         lPlus1 = nLag; % D.M. and D.Y. version
    #     end
    #     % -----------------------------------------------------------------
    #
    #     % get distance at lags (ll-1, ll, ll+1)
    #     distLminus1 = d( jb, lMinus1 ); % minus:  d( i-b, j-1 )
    #     distL       = d( ji, ll );      % actual: d( i-1, j   )
    #     distLplus1  = d( jb, lPlus1 );  % plus:   d( i-b, j+1 )
    #
    #     if (ji ~= jb) % equation 10 in Hale (2013)
    #         for kb = ji : iInc : jb - iInc % sum errors over i-1:i-b+1
    #             distLminus1 = distLminus1 + err( kb, lMinus1 );
    #             distLplus1  = distLplus1  + err( kb, lPlus1  );
    #         end
    #     end
    #
    #     dl = min( [ distLminus1, distL, distLplus1 ] ); % update minimum distance to previous sample
    #
    #     if ( dl ~= distL ) % then ll ~= ll and we check forward and backward
    #         if ( dl == distLminus1 )
    #             ll = lMinus1;
    #         else % ( dl == lPlus1 )
    #             ll = lPlus1;
    #         end
    #     end
    #
    #     % assume ii = ii - 1
    #     ii = ii + iInc; % previous time sample
    #
    #     stbar(ii) = ll + lmin - 1; % absolute integer of lag
    #
    #     % now move to correct time index, if smoothing difference over many
    #     % time samples using 'b'
    #     if ( ll == lMinus1 || ll == lPlus1 ) % check edges to see about b values
    #         if ( ji ~= jb ) % if b>1 then need to move more steps
    #             for kb = ji : iInc : jb - iInc
    #                 ii = ii + iInc; % move from i-1:i-b-1
    #                 stbar(ii) = ll + lmin - 1; % constant lag over that time
    #             end
    #         end
    #     end
    # end

end

"""

 Compute the accumulated error along the warping path for Dynamic
 Time Warping.

 USAGE: function error = computeDTWerror( Aerr, u, lag0 )

 INPUT:

   Aerr = error MATRIX (equation 13 in Hale, 2013)

   u    = warping function (samples) VECTOR

   lag0 = value of maximum lag (samples) SCALAR

 Written by Dylan Mikesell Last modified: 25 February 2015
"""
function computeDTWerror(Aerr::Array{Float64,2}, u::Array{Int64,1}, lag0::Int)

    npts = length(u);

    if size(Aerr,1) != npts
        #println("Funny things with dimensions of error matrix: check inputs.");
        #Aerr = transpose(Aerr);
        error("Funny things with dimensions of error matrix: check inputs.")
    end

    error = 0; # initialize

    # accumulate error
    for ii = 1:npts
        idx = lag0 + 1 + u[ii]; # index of lag
        error += Aerr[ii,idx]; # sum error
    end

    return error

end

end


# indicator of the zero cone

"""
  IndZero()

Returns the indicator function of the zero point, or "zero cone", i.e.,
  `g(x) = 0 if x = 0, +∞ otherwise`
"""

immutable IndZero <: IndicatorConvexCone end

function (f::IndZero){T <: RealOrComplex}(x::AbstractArray{T})
  for k in eachindex(x)
    if x[k] != zero(T)
      return Inf
    end
  end
  return 0.0
end

function prox!{T <: RealOrComplex}(f::IndZero, x::AbstractArray{T}, y::AbstractArray{T}, gamma::Real=1.0)
  for k in eachindex(x)
    y[k] = zero(T)
  end
  return 0.0
end

fun_name(f::IndZero) = "indicator of the zero cone"
fun_dom(f::IndZero) = "AbstractArray{Real}, AbstractArray{Complex}"
fun_expr(f::IndZero) = "x ↦ 0 if all(x = 0), +∞ otherwise"
fun_params(f::IndZero) = "none"

function prox_naive(f::IndZero, x::AbstractArray, gamma::Real=1.0)
  return zero(x), 0.0
end


using Distributed

## Function to get the number of points inside a circle
function in_circle_points(n)
    points_in = 0 ## Points inside the circle
    for i = 1:n
        x, y = rand(), rand() ## Randomly generate x,y btwn 0 & 1
        ## Count if generated value is within radius of the circle
        ## int the unit square
        points_in += (x*x + y*y) <= 1
    end
    ## Return the number of points in
    return points_in
end

## Run code in parallel across n cores
function parallel_pi(n)
    ## Create the number of workers specified by the
    p = nworkers()
    ## Combine results via addition (+)
    n_in = @distributed (+) for i=1:p
        in_circle_points(n / p)
    end
    return 4 * n_in / n ## Approximation of pi
end

## Run code serially on 1 worker
function serial_pi(n)
    n_in = in_circle_points(n)
    return 4 * n_in / n
end


#= Fornberg algorithm
This implements the Fornberg (1988) algorithm (https://doi.org/10.1090/S0025-5718-1988-0935077-0)
and hermite-based finite difference Fornberg(2020) algorithm (https://doi.org/10.1093/imanum/draa006)
to obtain Finite Difference weights over arbitrary points to arbitrary order.
Inputs:
        order: The derivative order for which we need the coefficients
        x0   : The point in the array 'x' for which we need the coefficients
        x    : A dummy array with relative coordinates, e.g., central differences
               need coordinates centred at 0 while those at boundaries need
               coordinates starting from 0 to the end point
        dfdx : optional argument to consider weights of the first-derivative of function or not
                if    dfdx == false (default kwarg), implies Fornberg(1988)
                      dfdx == true,                  implies Fornberg(2020)
    Outputs:
        if dfdx == false (default kwarg),   _C : weights to approximate derivative of required order using function values only.
                                 else,   _D,_E : weights to approximate derivative of required order using function and its first-
                                                 derivative values respectively.
=#

function calculate_weights(order::Int, x0::T, x::AbstractVector; dfdx::Bool = false) where T<:Real
    N = length(x)
    @assert order < N "Not enough points for the requested order."
    M = order
    c1 = one(T)
    c4 = x[1] - x0
    C = zeros(T, N, M+1)
    C[1,1] = 1
    @inbounds for i in 1 : N-1
        i1 = i + 1
        mn = min(i, M)
        c2 = one(T)
        c5 = c4
        c4 = x[i1] - x0
        for j in 0 : i-1
            j1 = j + 1
            c3 = x[i1] - x[j1]
            c2 *= c3
            if j == i-1
                for s in mn : -1 : 1
                    s1 = s + 1
                    C[i1,s1] = c1*(s*C[i,s] - c5*C[i,s1]) / c2
                end
                C[i1,1] = -c1*c5*C[i,1] / c2
           end
            for s in mn : -1 : 1
                s1 = s + 1
                C[j1,s1] = (c4*C[j1,s1] - s*C[j1,s]) / c3
            end
            C[j1,1] = c4 * C[j1,1] / c3
        end
        c1 = c2
    end
    #=
        This is to fix the problem of numerical instability which occurs when the sum of the stencil_coefficients is not
        exactly 0.
        https://scicomp.stackexchange.com/questions/11249/numerical-derivative-and-finite-difference-coefficients-any-update-of-the-fornb
        Stack Overflow answer on this issue.
        http://epubs.siam.org/doi/pdf/10.1137/S0036144596322507 - Modified Fornberg Algorithm
    =#
    _C = C[:,end]
    if order != 0
        _C[div(N,2)+1] -= sum(_C)
    end
    if dfdx == false
        return _C
    else
        A = x .- x';
        s = sum(1 ./ (A + I(N)), dims = 1) .- 1;
        cp = factorial.(0:M);
        cc = C./cp'
        c̃ = zeros(N, M+2);
        for k in 1:M+1
           c̃[:,k+1] = sum(cc[:,1:k].*cc[:,k:-1:1], dims = 2);
        end
        E = c̃[:,1:M+1] - (x .- x0).*c̃[:,2:M+2];
        D = c̃[:,2:M+2] + 2*E.*s';
        D = D.*cp';
        E = E.*cp';

        _D = D[:,end];   _E = E[:,end]
        return _D, _E
    end
end


"""
Stationary vector with iterative methods
"""

"""
stguess(Q::MatT, ::Type{Tv} = Float64)::Vector{Tv}

Get a vector which is guessed as the stationary vector of CTMC.
This is used as the initial vector for the iterative methods.
"""

function stguess(Q::MatT, ::Type{Tv} = Float64)::Vector{Tv} where {Tv,MatT}
    m, n = size(Q)
    @assert m == n
    result = Vector{Tv}(undef, n)
    for (i,x) in enumerate(spdiag(Q))
        result[i] = 1/x
    end
    result ./= sum(result)
end

"""
stgs(Q::SparseCSC{Tv,Ti}, x0::Vector{Tv}=stguess(Q,Tv); maxiter=5000, steps=20, rtol::Tv=Tv(1.0e-6))

Get a stationary vector of CTMC.

Parameters:
- Q: CTMC Kernal
- x0: Initial vector for iteration
- maxiter: The maximum number of iteration. The algorithm stops when the number of iteration becomes maxiter.
- steps: The number of steps to check the convergence
- rtol: the tolerance error. When the relative errors of two successive vectors with steps attains rtol, the algorithm stops.
Return value:
A tuple of
- x: stationary vector
- conv: A boolean whether the algorithm converges or not
- iter: The number of iterations
- rerror: The relative error when the algorithm stops
"""

function stgs(Q::SparseMatrixCSC{Tv,Ti}; x0::Vector{Tv}=stguess(Q,Tv),
        maxiter=5000, steps=20, rtol::Tv=Tv(1.0e-6)) where {Tv,Ti}
    stgs(SparseCSC(Q), x0=x0, maxiter=maxiter, steps=steps, rtol=rtol)
end

function stgs(Q::SparseCSC{Tv,Ti}; x0::Vector{Tv}=stguess(Q,Tv),
        maxiter=5000, steps=20, rtol::Tv=Tv(1.0e-6)) where {Tv,Ti}
    m, n = size(Q)
    @assert m == n
    b = zeros(Tv, n)
    x = copy(x0)
    iter = 0
    conv = false
    rerror::Tv = 0
    prevx = similar(x)
    while true
        prevx .= x
        for i in 1:steps
            gsstep!(x, Q, b)
            x ./= sum(x)
        end
        # rerror = maximum(abs.((x - prevx) ./ x))
        rerror = maximum(abs.(x - prevx)) / maximum(x)
        iter += steps
        if rerror < rtol
            conv = true
            break
        end
        if iter >= maxiter
            break
        end
    end
    return x, conv, iter, rerror
end

"""
stpower(P::AbstractMatrix{Tv}, x0::Vector{Tv}=stguess(Q,Tv); maxiter=5000, steps=20, rtol::Tv=Tv(1.0e-6))

Get a stationary vector of DTMC with power method.

Parameters:
- P: The transition probability matrix for DTMC
- x0: Initial vector for iteration
- maxiter: The maximum number of iteration. The algorithm stops when the number of iteration becomes maxiter.
- steps: The number of steps to check the convergence
- rtol: the tolerance error. When the relative errors of two successive vectors with steps attains rtol, the algorithm stops.
Return value:
A tuple of
- x: stationary vector
- conv: A boolean whether the algorithm converges or not
- iter: The number of iterations
- rerror: The relative error when the algorithm stops
"""

function stpower(P::AbstractMatrix{Tv}; x0::Vector{Tv}=stguess(P,Tv),
    maxiter=5000, steps=20, rtol::Tv=Tv(1.0e-6)) where {Tv}
    m, n = size(P)
    @assert m == n
    Pdash = P'
    x = copy(x0)
    iter = 0
    conv = false
    rerror::Tv = 0
    prevx = similar(x)
    while true
        prevx .= x
        for i in 1:steps
            x = Pdash * x
            x ./= sum(x)
        end
        # rerror = maximum(abs.((x - prevx) ./ x))
        rerror = maximum(abs.(x - prevx)) / maximum(x)
        iter += steps
        if rerror < rtol
            conv = true
            break
        end
        if iter >= maxiter
            break
        end
    end
    return x, conv, iter, rerror
end


# project x onto matrix A's column space, given A is full column rank
function Proj(x,A)
    y = try
        pinv(A' * A) * (A' * x)
    catch err
        if err != Nothing
            A' * x
        end
    end
    return A * y
end

function wavelet_perp_Matrix(w,A)
    N, k = size(A)
    B = zeros(N, k - 1)
    if k == 1
        return
    end
    tmp = A'*w
    if tmp[1] > 1e-6
        M = Matrix{Float64}(I, k-1, k-1)
        return A*vcat((-tmp[2:end] ./ tmp[1])', M)
    end
    if tmp[end] > 1e-6
        M = Matrix{Float64}(I, k-1, k-1)
        return A*vcat(M,(-tmp[1:end-1] ./ tmp[end])')
    end
    ind = findall(tmp .== maximum(tmp))[1]
    rest_ind = setdiff([i for i in 1:k], ind)
    M = Matrix{Float64}(I, k-1, k-1)
    B = A * vcat(vcat(M[1:ind-1,:], (-tmp[rest_ind] ./ tmp[ind])'), M[ind:end,:])
    return B
end

# Gram-Schmidt Process Orthogonalization
function gram_schmidt(A; tol = 1e-12)
    # Input: matirx A
    # Output: orthogonalization matrix of A's column vectors

    # Convert the matrix to a list of column vectors
    a = [A[:,i] for i in 1:size(A,2)]

    # Start Gram-Schmidt process
    q = []
    complement_dim = 0
    for i = 1:length(a)
        qtilde = a[i]
        for j = 1:(i - 1)
            qtilde -= (q[j]' * a[i]) * q[j]
        end
        if norm(qtilde) < tol
            complement_dim = size(A,2) - i + 1
            break
        end
        push!(q, qtilde / norm(qtilde))
    end
    Q = zeros(size(A,1), length(q))
    for i = 1:length(q)
        Q[:,i] .= q[i]
    end
    return Q, complement_dim
end

#
"""
    mgslp(A::Matrix{Float64}; tol::Float64 = 1e-12, p::Float64 = 1.0)

Modified Gram-Schmidt Process Orthogonalization with ℓᵖ pivoting algorithm (MGSLp)

# Input Arguments
- `A::Matrix{Float64}`: whose column vectors are to be orthogonalized.

# Output Argument
- `A::Matrix{Float64}`: orthogonalization matrix of A's column vectors based on ℓᵖ pivoting.
"""
function mgslp(A::Matrix{Float64}; tol::Float64 = 1e-12, p::Float64 = 1.0)
    n = size(A, 2)
    # Convert the matrix to a list of column vectors
    a = [A[:, i] for i in 1:n]

    # Start modified Gram-Schmidt process
    q = []
    v = copy(a)
    vec_lp_norm = [norm(v[i], p) for i in 1:n]
    complement_dim = 0
    for i = 1:n
        # Pivoting based on minimum ℓᵖ-norm
        idx = findmin(vec_lp_norm)[2]
        v[i], v[idx + i - 1] = v[idx + i - 1], v[i]
        # Check the linear dependency
        if norm(v[i]) < tol
            complement_dim = n - i + 1
            break
        end
        qtilde = v[i] / norm(v[i], 2)
        vec_lp_norm = []
        for j = (i + 1):n
            v[j] .-= (qtilde' * v[j]) * qtilde
            push!(vec_lp_norm, norm(v[j], p))
        end
        push!(q, qtilde)
    end
    Q = zeros(size(A,1), length(q))
    for i = 1:length(q)
        Q[:,i] .= q[i]
    end
    return Q, complement_dim
end


export hindmarsh_rose

"""
Equations of motion for the Hindmarsh-Rose system.
"""
function eom_hindmarsh_rose(u, p, t)
    a, b, c, d, r, s, xᵣ, I = (p...,)
    x, y, z = (u...,)

	ϕ = -a*x^3 + b*x^2
	ψ = c - d*x^2
    dx = y + ϕ - z + I
	dy = ψ - y
	dz = r*(s*(x - xᵣ) - z)
    return SVector{3}(dx, dy, dz)
end

"""
	hindmarsh_rose(u₀, p)

Initialise a Hindmarsh-Rose system, which is a model of neuronal
spiking.

```math
\\begin{aligned}
\\dfrac{dx}{dt} &= y + \\phi(x) - z + I
\\dfrac{dy}{dt} &= \\psi(x) - y
\\dfrac{dz}{dt} &= r[s(x - x_R) - z],
\\end{aligned}
```
where

```math
\\begin{aligned}
\\phi(x) &= -ax^3+bx^2
\\psi(x) &= c - dx^2
\\end{aligned}
```

If parameters other than the defaults are to be used, they must be
provided as a vector `[a, b, c, d, r, s, xᵣ, I]`.
"""
function hindmarsh_rose(u₀, p)
    ContinuousDynamicalSystem(eom_hindmarsh_rose, u₀, p)
end
hindmarsh_rose(;u₀ = rand(3), a = 1, b = 3, c = 1, d = 5, r = 1e-3, s = 4, xᵣ = - 8/5, I = -8) =
    hindmarsh_rose(u₀, [a, b, c, d, r, s, xᵣ, I])