# Identity matrix
eye(m::Int64) = Diagonal(ones(m))

unpack(s, f::Symbol) = reduce(hcat, getproperty.(s, f))

function unpack(s, f::Vector{Symbol})
    A = []
    for i in f
        push!(A, reduce(hcat, getproperty.(s, i)))
    end
    return A
end

unpack_matrix(s, f) = getproperty.(s, f)

function unpack_matrix(s, f::Vector{Symbol})
    A = []
    for i in f
        push!(A, getproperty.(s, i))
    end
    return A
end

_return_type(P::Matrix, C) = Matrix(C)
_return_type(P, C) = Diagonal(C)

_square_root_filter(filter::CholeskyFactorization, P) = _return_type(P, cholesky(P).L)

function _square_root_filter(filter::ModifiedCholeskyFactorization, P)
	F = _modified_cholesky(P)
	_return_type(P, F.U * sqrt.(F.D))
end

function _square_root_filter(filter::SVDFactorization, P)
	_C = svd(Matrix(P))
	return _return_type(P, _C.U * sqrt(Diagonal(_C.S)))
end

"""

	F = _modified_cholesky(P)

Computes modified Cholesky factors F.U and F.D of a symmetric positive definite matrix P,
such that F.U is unit upper triangular and F.D is diagonal. Thus, and P = F.U*F.D*F.U'.

Reference: Grewal - Kalman filtering theory and practice

"""
function _modified_cholesky(P::AbstractMatrix)
	isposdef(P) || throw("input matrix must be positive-definite. Modified Cholesky factorization failed.")
	n = size(P, 1)
	_P = (P .+ P') ./ 2.0 # take symmetric part
	U, D = similar(_P), similar(_P)
	for j in n:-1:1
   		for i in j:-1:1
			s = P[i, j]
			for k in j+1:n
				s -= U[i, k]*D[k, k]*U[j, k]
			end
			if i==j
				D[j, j] = s
				U[j, j] = 1.0
			else
				U[i, j] = s / D[j, j]
			end
		end
	end
	return (U=U, D=D)
end


function load_distributions_samples()
return quote

using .Distributions: Distribution, Uniform, Normal
import .Distributions

# ======================================================
# Sampling from a uniform distribution on balls/spheres
# ======================================================

"""
    _sample_unit_nsphere_muller!(D::Vector{Vector{N}}, n::Int, p::Int;
                                 rng::AbstractRNG=GLOBAL_RNG,
                                 seed::Union{Int, Nothing}=nothing) where {N}

Draw samples from a uniform distribution on an ``n``-dimensional unit
sphere using Muller's method.

### Input

- `D`    -- output, vector of points
- `n`    -- dimension of the sphere
- `p`    -- number of random samples
- `rng`  -- (optional, default: `GLOBAL_RNG`) random number generator
- `seed` -- (optional, default: `nothing`) seed for reseeding

### Output

A vector of `nsamples` vectors.

### Algorithm

This function implements Muller's method of normalised Gaussians [1] to uniformly
sample over the ``n``-dimensional sphere ``S^n`` (which is the bounding surface
of the ``n``-dimensional unit ball).

Given ``n`` canonical Gaussian random variables ``Z₁, Z₂, …, Z_n``, the
distribution of the vectors

```math
\\dfrac{1}{α}\\left(z₁, z₂, …, z_n\\right)^T,
```
where ``α := \\sqrt{z₁² + z₂² + … + z_n²}``, is uniform over ``S^n``.

[1] Muller, Mervin E. *A note on a method for generating points uniformly on
    n-dimensional spheres.* Communications of the ACM 2.4 (1959): 19-20.
"""
function _sample_unit_nsphere_muller!(D::Vector{Vector{N}}, n::Int, p::Int;
                                      rng::AbstractRNG=GLOBAL_RNG,
                                      seed::Union{Int, Nothing}=nothing) where {N}
    rng = reseed(rng, seed)
    Zdims = [Normal() for _ in 1:n] # normal distributions, one for each dimension
    v = Vector{N}(undef, n) # sample direction
    @inbounds for j in 1:p
        α = zero(N)
        for i in 1:n
            v[i] = rand(rng, Zdims[i])
            α += v[i]^2
        end
        D[j] = v ./ sqrt(α)
    end
    return D
end

"""
    _sample_unit_nball_muller!(D::Vector{Vector{N}}, n::Int, p::Int;
                               rng::AbstractRNG=GLOBAL_RNG,
                               seed::Union{Int, Nothing}=nothing) where {N}

Draw samples from a uniform distribution on an ``n``-dimensional unit ball
using Muller's method.

### Input

- `D`    -- output, vector of points
- `n`    -- dimension of the ball
- `p`    -- number of random samples
- `rng`  -- (optional, default: `GLOBAL_RNG`) random number generator
- `seed` -- (optional, default: `nothing`) seed for reseeding

### Output

A vector of `nsamples` vectors.

### Algorithm

This function implements Muller's method of normalised Gaussians [1] to
sample from the interior of the ball.

Given ``n`` Gaussian random variables ``Z₁, Z₂, …, Z_n`` and a uniformly
distributed random variable ``r`` with support in ``[0, 1]``, the distribution
of the vectors

```math
\\dfrac{r^{1/n}}{α} \\left(z₁, z₂, …, z_n\\right)^T,
```
where ``α := \\sqrt{z₁² + z₂² + … + z_n²}``, is uniform over the
``n``-dimensional unit ball.

[1] Muller, Mervin E. *A note on a method for generating points uniformly on
    n-dimensional spheres.* Communications of the ACM 2.4 (1959): 19-20.
"""
function _sample_unit_nball_muller!(D::Vector{Vector{N}}, n::Int, p::Int;
                                    rng::AbstractRNG=GLOBAL_RNG,
                                    seed::Union{Int, Nothing}=nothing) where {N}

    rng = reseed(rng, seed)
    Zdims = [Normal() for _ in 1:n] # normal distributions, one for each dimension
    Zrad = Uniform() # distribution to pick random radius
    one_over_n = one(N)/n
    v = Vector{N}(undef, n) # sample direction
    @inbounds for j in 1:p
        α = zero(N)
        for i in 1:n
            v[i] = rand(rng, Zdims[i])
            α += v[i]^2
        end
        r = rand(rng, Zrad)
        β = r^one_over_n / sqrt(α)
        D[j] = v .* β
    end
    return D
end

# ========================
# Sampling from a LazySet
# ========================

"""
    Sampler

Abstract type for defining new sample methods.

### Notes

All subtypes should implement a `_sample!` method.
"""
abstract type Sampler end

"""
    sample(X::LazySet{N}, num_samples::Int;
           rng::AbstractRNG=GLOBAL_RNG,
           seed::Union{Int, Nothing}=nothing) where {N}

Sampling of an arbitrary `LazySet` `X`.

### Input

- `X`           -- lazyset
- `num_samples` -- number of random samples
- `sampler`     -- Sampler used (default: `RejectionSampler`)
- `rng`         -- (optional, default: `GLOBAL_RNG`) random number generator
- `seed`        -- (optional, default: `nothing`) seed for reseeding

### Output

A vector of `num_samples` vectors.
If `num_samples` is not passed one sample as a single vector is returned.

### Algorithm

See the documentation of `RejectionSampler`.
"""
function sample(X::LazySet{N}, num_samples::Int;
                sampler=RejectionSampler,
                rng::AbstractRNG=GLOBAL_RNG,
                seed::Union{Int, Nothing}=nothing) where {N<:Real}
    require(:Distributions; fun_name="sample")
    @assert isbounded(X) "this function requires that the set `X` is bounded"*
                            ", but it is not"

    D = Vector{Vector{N}}(undef, num_samples) # preallocate output
    _sample!(D, sampler(X); rng=rng, seed=seed)
    return D
end

# without argument, returns a single element (instead of a singleton)
function sample(X::LazySet{N}; kwargs...) where {N<:Real}
    return sample(X, 1; kwargs...)[1]
end

# =====================
# Rejection Sampling
# =====================

"""
    RejectionSampler{S<:LazySet, D<:Distribution} <: Sampler

Type used for rejection sampling of an arbitrary `LazySet` `X`.

### Fields

- `X`           -- lazyset
- `box_approx`  -- Distribution from which the sample is drawn

### Algorithm

Draw a sample ``x`` from a uniform distribution of a box-overapproximation of the
original set ``X`` in all ``n`` dimensions. The function rejects a drawn sample ``x``
and redraws as long as the sample is not contained in the original set ``X``,
i.e., ``x ∉ X``.
"""
struct RejectionSampler{S<:LazySet, D<:Distribution} <: Sampler
    X::S
    box_approx::Vector{D}
end

function RejectionSampler(X, distribution=Uniform)
    B = box_approximation(X)
    canonical_support = hcat(low(B), high(B))
    dims = size(canonical_support, 1)
    box_approx = [distribution(canonical_support[i,:]...) for i = 1:dims]
    return RejectionSampler(X, box_approx)
end

"""
    _sample!(D::Vector{Vector{N}},
             sampler::RejectionSampler;
             rng::AbstractRNG=GLOBAL_RNG,
             seed::Union{Int, Nothing}=nothing) where {N<:Real}

Sample points using rejection sampling.

### Input

- `D`           -- output, vector of points
- `sampler`     -- Sampler from which the points are sampled
- `rng`         -- (optional, default: `GLOBAL_RNG`) random number generator
- `seed`        -- (optional, default: `nothing`) seed for reseeding

### Output

A vector of `num_samples` vectors.
"""
function _sample!(D::Vector{Vector{N}},
                  sampler::RejectionSampler;
                  rng::AbstractRNG=GLOBAL_RNG,
                  seed::Union{Int, Nothing}=nothing) where {N<:Real}
    rng = reseed(rng, seed)
    @inbounds for i in 1:length(D)
        w = rand.(Ref(rng), sampler.box_approx)
        while w ∉ sampler.X
            w = rand.(Ref(rng), sampler.box_approx)
        end
        D[i] = w
    end
    return D
end

end # quote
end # function load_distributions_samples()


using Random, PyPlot, StatsBase, Combinatorics
Random.seed!(1)
function bruteSetsProbabilityAllMiss(n)
    omega = collect(permutations(1:n))
    matchEvents = []

    for i in 1:n
        event = []
        for p in omega
            if p[i] == i
                push!(event,p)
            end
        end
        push!(matchEvents,event)
    end

    noMatch = setdiff(omega,union(matchEvents...))
    return length(noMatch)/length(omega)
end

function formulaCalcAllMiss(n)
    return sum([(-1)^k/factorial(k) for k in 0:n])
end

function mcAllMiss(n,N)
    function envelopeStuffer()
        envelopes = Random.shuffle!(collect(1:n))
        return sum([envelopes[i] == i for i in 1:n]) == 0
    end

    data = [envelopeStuffer() for _ in 1:N]
    return sum(data)/N
end

N = 10^6

println("n\tBrute Force\tFormula\t\tMonte Carlo\t\tAnalytic",)
for n in 1:8
    bruteForce = bruteSetsProbabilityAllMiss(n)
    fromFormula = formulaCalcAllMiss(n)
    fromMC = mcAllMiss(n,N)
    println(n,"\t",round(bruteForce,digits=4),"\t\t",round(fromFormula,digits=4),
	"\t\t",round(fromMC,digits=4),"\t\t",round(1/MathConstants.e,digits=4))
end


function four1(data::Array{T,1}, isign::Int64 = 1) where {T<:Number}
	nn = length(data)
	nnlog = log2(nn)
	if ceil(nnlog) != nnlog
		error("Length of data vector must be a power of 2.")
	end
	j::Int64 = 1
	res::Array{Complex{Float64},1} = copy(data)

	for i = 1:nn
		if j > i
			res[j], res[i] = data[i], data[j]
		end
		m = nn/2
		while m >= 2 && j > m
			j -= m
			m = m/2
		end
		j += m
	end

	mmax = 1;
	while nn > mmax
		istep = 2*mmax
		theta = isign*(pi/mmax)
		wp = -2*sin(theta/2)^2 + sin(theta)*im
		w = 1 + 0im

		for m = 1:mmax
			for i = m:istep:nn
				j = i + mmax
				temp = w*res[j]
				res[j] = res[i] - temp
				res[i] += temp
			end
			w += w*wp
		end
		mmax = istep
	end

	res
end

input_size = 2048
Fs = 1000
T = 1/Fs           # Sampling period

t = T*(0:(input_size-1))
X = 0.7*sin.(2*pi*50*t) + sin.(2*pi*120*t)

@time Y = four1(X, 1)

;


my_number = 1.2            # Assignment

pi                         # Pre-defined variable

My_number                  # Error - was defined without capital M

δ = 0.001                  # Unicode characters allowed (type \delta and press tab)

my_number_2 = (my_number + 1) * δ     # Use pre-defined variables in expressions

counter = 1               # counter = 1

counter += 1              # counter = counter + 1

counter ^= 4              # counter = counter ^ 4

angle = 60                # degrees

angle *= π / 180          # convert to radians

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

println(360 / 2*π)  # Left first - equals 180*π
println(360 / 2π)   # Right first - equals 180/π

println(2^2x)       # 2^(2x), not (4)x
println(2x^2)       # 2*(x^2), not (2x)^2

#=
   Experiments with spheres
=#

radius1 = 100     # Outer sphere radius
radius2 = 50      # Inner sphere radius

# Compute the volume of each sphere
vol1 = 4π * radius1^3 / 3
vol2 = 4π * radius2^3 / 3

# Compute the volume of a hollow sphere
vol_hollow = vol1 - vol2


### A Pluto.jl notebook ###
# v0.14.4

using Markdown
using InteractiveUtils

# ╔═╡ 892fb737-6fc0-4438-b586-039a440b01f9
using LinearAlgebra, SparseArrays

# ╔═╡ 069cf190-77cc-4757-9db3-1945ddd07438
using DifferentialEquations, Plots

# ╔═╡ 7bc3e874-b7be-11eb-0792-cf4e21e8c8b2
md"# Heat equation with Proportional control

Consider a one-dimensional rod $\Omega = (0,L)$ and time $t \in (0,T).$
"

# ╔═╡ 1158de2e-b7a6-43f3-a4e6-2ef9598a4ca5
L = 0.1 # Length of rod

# ╔═╡ 30cebec3-9c70-4c6a-94a8-fdcbe6cda783
Tf = 2000.0 # Final simulation time

# ╔═╡ 93fb0f90-c454-410f-aa5e-cd8e6c055dc0
md"The rod has the physical properties
- thermal conductivity $\lambda$,
- specific heat capacity $c$,
- mass density $\rho$.

The diffusivity constant is given by $\alpha = \frac{\lambda}{c \rho}$. The heat equation is noted as

$\dot{\vartheta}(t,x) = \alpha \frac{\partial^2}{\partial x^2} \vartheta(t,x)$

for $(t,x) \in (0,T) \times \Omega$ with initial condition

$\vartheta(0,x) = \vartheta_{0}(x)$

for $x \in \overline{\Omega}$."

# ╔═╡ b294bbfc-45e7-4a15-a839-5ceb56746be7
begin
	λ = 45.0;    # Thermal conductivity
	ρ = 7800.0;  # Mass density
	cap = 480.0; # Specific heat capacitivity
end

# ╔═╡ 543b3f13-4cef-46f1-9141-172b81e7d48b
α = λ/(cap * ρ)   # Diffusivity

# ╔═╡ 1deb117d-21f7-4329-a7db-2be26e43129a
md"## Boundary conditions

Only on the right side of the rod, heat transfer

$\phi_{t}(t,x=L) = - h (\vartheta(t,x=L) - \vartheta_{amb})$

and heat radiation

$\phi_{r}(t,x=L) = - \epsilon \varrho (\vartheta(t,x=L)^4 - \vartheta_{amb}^4)$

is assumed to operate. Parameter $\vartheta_{amb}$ describes the ambient temperature, $h>0$ is called the heat transfer coefficient, $\epsilon \in (0,1)$ is called emissivity and $\varrho\approx 5.67 \cdot 10^{-8}$ is known as Stefan-Boltzmann constant. For simplicity, $k=\epsilon \cdot \varrho$.

The sum of $\phi_{t}$ and $\phi_{r}$

$\phi_{out}(t,L) = \phi_{t}(t,L) + \phi_{r}(t,L) = - h (\vartheta(t,L) - \vartheta_{amb}) - k (\vartheta(t,L)^4 - \vartheta_{amb}^4)$

is the flux from the rod to the environment."

# ╔═╡ 78292b42-a203-496c-a143-def3ff128f69
begin
	h = 10.0; # Heat transfer coefficient
	ϵ = 0.6;  # Emissivity
	sb = 5.67*10^(-8) # Stefan-Boltzmann constant
	k = ϵ * sb; 	  # Radiation coefficient
end

# ╔═╡ 0ecea33b-bf24-48eb-b4b4-bbf288e62e6b
θamb = 298.0 # Ambient temperature in Kelvin

# ╔═╡ 3a059f15-d850-468a-bb1d-f42ef5b24609
ϕout(θ) = -h * (θ - θamb)  - k*(θ^4 - θamb^4)

# ╔═╡ 7202aa44-b393-4e0b-9c34-2a9b9376194c
md"### Induced heat

On the left side of the rod, a heat source is assumed as

$\phi_{in}(t,x=0) = b ~ u(t)$

with constant $b>0$, input signal

$u(t) = K_{p} ~ e(t)$

and error

$e(t) = y_{ref}(t) - y(t).$

The output is considered as the temperature on the right side, e.g.

$y(t) = \vartheta(t, x=L)$

and the reference is as an arbitrary fixed value

$y_{ref}(t) = y_{ref} > 0.$

Thus, the natural (Robin) boundary conditions are noted as

$\lambda ~ \left.\frac{\partial}{\partial x} \vartheta(\cdot, x) \right\rvert_{x = 0} \cdot \vec{n} ~=~ \phi_{in}(t,0) = b ~ u(t)$

and

$\lambda ~  \left. \frac{\partial}{\partial x} \vartheta(\cdot, x) \right\rvert_{x = L} \cdot \vec{n} ~=~ \phi_{out}(t,x) = - h (\vartheta(t,L) - \vartheta_{amb}) - k (\vartheta(t,L)^4 - \vartheta_{amb}^4)$

with outer [normal vector](https://en.wikipedia.org/wiki/Normal_(geometry)) $\vec{n}$ on the left or right boundary. Here, the normal vector $\vec{n} = -1$ on the left boundary at $x=0$, and  $\vec{n} = 1$ on the right boundary at $x=L$."

# ╔═╡ 2775433b-d8b1-48be-9b60-08641db15803
b = 1;

# ╔═╡ 4c59909b-ffb6-4dfc-94ab-142d41deeeab
Kp = 10^3  # Proportional gain

# ╔═╡ 4f0c3671-d3bf-4e36-a1f3-8e8d0ac0e454
yref = 400.0 # Reference temperature

# ╔═╡ 23343fa0-6033-498d-b089-d43e4382e332
u_in(err) = Kp * err # input signals

# ╔═╡ daa17037-fd28-460c-97f5-826742b9a419
md"### Spatial approximation

The one-dimensional rod is discretized as a one-dimensional grid with N points and the finite discretization is noted by

$\Delta x = \frac{L}{N-1}.$

The resulting grid points have the position $x^{0} = 0$, $x^{n} = n~\Delta x$ and $x^{N-1} = L$. The second order derivative is approximated using the [Taylor series](https://en.wikipedia.org/wiki/Taylor_series) as

$\left. \frac{\partial^2 f(x)}{\partial x^2} \right\rvert_{x = \tilde{x}} \approx \frac{1}{\Delta x^2} \left[ f(\tilde{x} - \Delta x) - 2 f(\tilde{x}) + f(\tilde{x} + \Delta x)  \right].$

This [Finite Difference](https://en.wikipedia.org/wiki/Finite_difference_method) scheme is written in matrix-vector form as

$\frac{\partial^2 \vartheta(t,x)}{\partial x^2} \approx \frac{1}{\Delta x^2}
\begin{pmatrix}
-2 & 1 & 0 & \cdots &  & 0 \\
1 & -2 & 1 & 0  & \cdots & \vdots \\
0 & 1 & -2 & 1 & \ddots &  \\
 & & \ddots & \ddots & \ddots &  & \\
& & & 1 & -2 & 1 \\
& & &  & 1 & -2 \\
\end{pmatrix}
~
\begin{pmatrix}
\vartheta(t, x^{0}) \\
\vartheta(t, x^{1}) \\
\vdots \\
\vartheta(t, x^{N-1}) \\
\end{pmatrix}.$
"

# ╔═╡ 0309bf5e-4033-4832-86ed-146aea0863c0
N = 101 # Number of grid elements

# ╔═╡ 5c9ec071-6140-4af6-a418-df79c10daf5e
Δx = L/(N-1) # Finite discretization

# ╔═╡ 905337c5-9eee-4b67-b215-bfd7ca905e4c
# Diffusion matrix
M = spdiagm(-1 => ones(N-1), 0 => -2*ones(N), 1 => ones(N-1));

# ╔═╡ f0e9f197-303b-48a6-a540-9b11513050d4
Matrix(M)[1:5,1:5]

# ╔═╡ 0e155080-2cfe-45c0-8166-ba91f95d6f8b
md"### Approximated boundary conditions

The spatial approximation of the heat equation

$\frac{\partial^2 \vartheta(t,x)}{\partial x^2} \approx \frac{1}{\Delta x^2} \left( \vartheta(\cdot, x^{n-1}) - 2 ~ \vartheta(\cdot, x^{n}) + \vartheta(\cdot, x^{n+1}) \right)$

can not be evaluated directly at the grid points $x^{0}$ and $x^{N-1}$ because they depend on values from the not-existing grid points $x^{-1}$ and $x^{N}$. Therefore, the Neumann boundary condition is used to find $\vartheta(\cdot, x^{-1})$ and $\vartheta(\cdot, x^{N})$.

The Robin boundary condition states that the heat flux at both sides is described by

$\lambda \left. \frac{\partial \vartheta(\cdot,x)}{\partial x}  \right\rvert_{x = 0} \cdot \vec{n} = -1 \cdot \lambda \frac{\vartheta(\cdot, x^{1}) - \vartheta(\cdot, x^{-1})}{2 \Delta x}   = \phi_{in}(t,x^{0})$

on the left boundary and

$\lambda \left. \frac{\partial \vartheta(\cdot,x)}{\partial x}  \right\rvert_{x = L} \cdot \vec{n} = \lambda \frac{\vartheta(\cdot, x^{N}) - \vartheta(\cdot, x^{N-2})}{2 \Delta x}  = \phi_{out}(t,x^{N-1})$

on the right boundary. Reformulating both equations lead to

$\vartheta(\cdot, x^{-1}) = \vartheta(\cdot, x^{1}) + \frac{2 \Delta x}{\lambda} ~ \phi_{in}(t,x^{0})$

and

$\vartheta(\cdot, x^{N}) = \vartheta(\cdot, x^{N-2})  + \frac{2 \Delta x}{\lambda} ~ \phi_{out}(t,x^{N-1}).$

On the left boundary one yields

$\left. \frac{\partial^2 \vartheta(t,x)}{\partial x^2} \right\rvert_{x = 0} \approx \frac{1}{\Delta x^2} \left( - 2 ~ \vartheta(\cdot, x^{0}) + 2 ~ \vartheta(\cdot, x^{1}) \right) + \frac{2}{\lambda ~ \Delta x} \phi_{in}(t,x^{0})$

and analog on the right boundary

$\left. \frac{\partial^2 \vartheta(t,x)}{\partial x^2} \right\rvert_{x = 0} \approx \frac{1}{\Delta x^2} \left(2 ~ \vartheta(\cdot, x^{N-2}) - 2 ~ \vartheta(\cdot, x^{N-1}) \right) + \frac{2}{\lambda ~ \Delta x} \phi_{out}(t,x^{N-1}).$

The induced and emitted flux are noted as

$\Phi_{in}(t) =
\begin{pmatrix}
b  \\
0 \\
\vdots \\
0 \\
\end{pmatrix}
~ u(t)

\quad \text{and} \quad

\Phi_{out}(t) =
\begin{pmatrix}
0 \\
\vdots \\
0 \\
1 \\
\end{pmatrix}
~\phi_{out}(t, x^{N-1})$

and the diffusion matrix is described by

$M =
\begin{pmatrix}
-2 & 2 & 0 & \cdots &  & 0 \\
1 & -2 & 1 & 0  & \cdots & \vdots \\
0 & 1 & -2 & 1 & \ddots &  \\
 & & \ddots & \ddots & \ddots &  & \\
& & & 1 & -2 & 1 \\
& & &  & 2 & -2 \\
\end{pmatrix}.$
"

# ╔═╡ c61962e8-86bd-45ed-a62d-ffc812d8e35e
# First row
M[1,2] = 2;

# ╔═╡ 33baab03-fbc1-4944-bc16-e4022b051c69
# Last row
M[end,end-1] = 2;

# ╔═╡ facdbabf-6e2b-418a-a07e-36188244e585
md"### Heat equation in State-space representation

After the spatial approximation the heat equation has a form of

$\dot{\theta}(t) =  \frac{\alpha}{\Delta x^2} ~ M ~ \theta(t) + \frac{2 \alpha}{\Delta x} ~ \Phi_{in}(t) + \frac{2 \alpha}{\Delta x} ~ \Phi_{out}(t) \tag{1}$

with M as diffusion matrix, temperature vector $\quad \theta(t) = \left( \vartheta(t, x^{0}), \cdots, \vartheta(t, x^{N-1}) \right)^{\top}$.

The original ODE $(1)$ is recasted in [state-space representation](https://en.wikipedia.org/wiki/State-space_representation) as

$\dot{\theta}(t) = A ~ \theta(t) + B ~ u(t) + E ~ w(t)$
$y(t) = C \theta(t)$

with
- system matrix $A =  \frac{\alpha}{\Delta x^2} ~ M$,
- input vector
$B = \frac{2 ~ \alpha}{\lambda ~ \Delta x}
\begin{pmatrix}
b \\
0 \\
\vdots \\
0
\end{pmatrix},$

- disturbance vector
$E = \frac{2 ~ \alpha}{\lambda ~ \Delta x}
\begin{pmatrix}
0 \\
\vdots \\
0 \\
1
\end{pmatrix},$

- emitted flux $w(t) =  \phi_{out}(t,x^{N-1})$ and
- output vector $C = (0, \cdots, 0, 1)$

to describe the input-state-output behaviour including the integral controller. The emitted heat is considered as a disturbance. The PI controller is introduced as

$u(t) = K_{p} ~ e(t)$

with error
$e(t) = y_{ref}(t) - y(t) = y_{ref}(t) - C ~ \theta(t)$

The thermal dynamics with PI control is described by

$\dot{\theta}(t) = A ~ \theta(t) + B ~ K_{p} ~ e(t) + E ~ w(t).$

This is an ordinary differential equation (ODE) that can be solved with common solvers like forward Euler method or Runge-Kutta scheme.

"

# ╔═╡ 5ef85c8f-cb69-42df-a098-ab98903a5f32
A = α/(Δx^2) * M

# ╔═╡ 25bb45ab-55af-4b8b-a11b-05be68631b78
begin
	B = spzeros(N); # Input vector
	B[1] = 2 * (α/(λ*Δx)) * b
	E = spzeros(N); # Disturbance vector
	E[end] = 2 * (α/(λ*Δx)) * b
	C = spzeros(1,N); # Output vector
	C[end] = 1;
end

# ╔═╡ 81a2d204-7386-4ed7-af89-7c076940817b
w(θ) = ϕout(θ)

# ╔═╡ 113d8c87-f4e6-4f1b-b188-879e7ad1e146
function heat_eq!(dx,x,p,t)
    global u_hist

    err = yref - x[end]
    u = u_in(err)

    dx .= A * x + B * u + E * ϕout( x[end] )    # Integration of temperatures

  end

# ╔═╡ d9138271-79fc-4b7c-8501-e8d8f3304cd5
md"### Initial Conditions

For simplicity, the initial data of the heat equation is assumed as $\vartheta_{0}(x)= 1000$ Kelvin for all $x \in \left[0,L\right]$. This function is approximated to gain the initial conditions of the ODE

$\theta(0) = \left( 300, \cdots , 300 \right)^{\top}.$
"

# ╔═╡ db6248c7-24b8-4f44-82d3-4892203297b5
θ₀ = 300.0 * ones(N)

# ╔═╡ 837632eb-b4db-47ec-8e4a-32a0fe81060b
md"## Simulation"

# ╔═╡ 3f382fe5-2aa6-4d82-b939-977985ae8bbb
Δt = 10^(-2) # Sampling time

# ╔═╡ 02ba374a-15b4-4481-baf6-c1a3ce45ed5e
tspan = (0.0, Tf)

# ╔═╡ 0db82493-4dda-4992-aa5b-63734b556a90
prob = ODEProblem( heat_eq!, θ₀, tspan ) # ODE Problem

# ╔═╡ f7dcb8a6-0eb4-4762-942e-53362d4fe1ba
sol = solve(prob,Euler(),dt=Δt,progress=true, saveat=1.0, save_start=true); # Solving the ODE

# ╔═╡ 9c38b2b3-1622-41f3-a2c2-ff79b552e454
# 1-dimensional grid
xspan = 0 : Δx : L

# ╔═╡ ba66c778-6c1c-4e32-934d-e8aca0eb1c66
heatmap(sol.t, xspan, sol[1:end,1:end], xaxis="Time [s]", yaxis="Position [m]", title="Evolution of temperature")

# ╔═╡ 3dcfed19-8235-4309-b1c7-8d4e30ba6701
plot(sol.t,[sol[1,:], sol[end-1,:]], label=["Left" "Right"], title="Temperature at the left/right end", legend=:bottomright)

# ╔═╡ Cell order:
# ╟─7bc3e874-b7be-11eb-0792-cf4e21e8c8b2
# ╠═1158de2e-b7a6-43f3-a4e6-2ef9598a4ca5
# ╠═30cebec3-9c70-4c6a-94a8-fdcbe6cda783
# ╟─93fb0f90-c454-410f-aa5e-cd8e6c055dc0
# ╠═b294bbfc-45e7-4a15-a839-5ceb56746be7
# ╠═543b3f13-4cef-46f1-9141-172b81e7d48b
# ╟─1deb117d-21f7-4329-a7db-2be26e43129a
# ╠═78292b42-a203-496c-a143-def3ff128f69
# ╠═0ecea33b-bf24-48eb-b4b4-bbf288e62e6b
# ╠═3a059f15-d850-468a-bb1d-f42ef5b24609
# ╟─7202aa44-b393-4e0b-9c34-2a9b9376194c
# ╠═2775433b-d8b1-48be-9b60-08641db15803
# ╠═4c59909b-ffb6-4dfc-94ab-142d41deeeab
# ╠═4f0c3671-d3bf-4e36-a1f3-8e8d0ac0e454
# ╠═23343fa0-6033-498d-b089-d43e4382e332
# ╟─daa17037-fd28-460c-97f5-826742b9a419
# ╠═892fb737-6fc0-4438-b586-039a440b01f9
# ╠═0309bf5e-4033-4832-86ed-146aea0863c0
# ╠═5c9ec071-6140-4af6-a418-df79c10daf5e
# ╠═905337c5-9eee-4b67-b215-bfd7ca905e4c
# ╠═f0e9f197-303b-48a6-a540-9b11513050d4
# ╟─0e155080-2cfe-45c0-8166-ba91f95d6f8b
# ╠═c61962e8-86bd-45ed-a62d-ffc812d8e35e
# ╠═33baab03-fbc1-4944-bc16-e4022b051c69
# ╟─facdbabf-6e2b-418a-a07e-36188244e585
# ╠═5ef85c8f-cb69-42df-a098-ab98903a5f32
# ╠═25bb45ab-55af-4b8b-a11b-05be68631b78
# ╠═81a2d204-7386-4ed7-af89-7c076940817b
# ╠═113d8c87-f4e6-4f1b-b188-879e7ad1e146
# ╟─d9138271-79fc-4b7c-8501-e8d8f3304cd5
# ╠═db6248c7-24b8-4f44-82d3-4892203297b5
# ╠═837632eb-b4db-47ec-8e4a-32a0fe81060b
# ╠═069cf190-77cc-4757-9db3-1945ddd07438
# ╠═3f382fe5-2aa6-4d82-b939-977985ae8bbb
# ╠═02ba374a-15b4-4481-baf6-c1a3ce45ed5e
# ╠═0db82493-4dda-4992-aa5b-63734b556a90
# ╠═f7dcb8a6-0eb4-4762-942e-53362d4fe1ba
# ╠═9c38b2b3-1622-41f3-a2c2-ff79b552e454
# ╠═ba66c778-6c1c-4e32-934d-e8aca0eb1c66
# ╠═3dcfed19-8235-4309-b1c7-8d4e30ba6701


using Random, Distributions, KernelDensity, Plots; pyplot()
Random.seed!(0)

mu1, sigma1 = 10, 5
mu2, sigma2 = 40, 12
dist1, dist2 = Normal(mu1,sigma1), Normal(mu2,sigma2)
p = 0.3
mixRv() = (rand() <= p) ? rand(dist1) : rand(dist2)
mixPDF(x) = p*pdf(dist1,x) + (1-p)*pdf(dist2,x)

n = 2000
data = [mixRv() for _ in 1:n]

kdeDist = kde(data)

xGrid = -20:0.1:80
pdfKDE = pdf(kdeDist,xGrid)

plot(xGrid, pdfKDE, c=:blue, label="KDE PDF")
stephist!(data, bins=50, c=:black, normed=:true, label="Histogram")
p1 = plot!(xGrid, mixPDF.(xGrid), c=:red, label="Underlying PDF",
    xlims=(-20,80), ylims=(0,0.035), legend=:topleft,
    xlabel="X", ylabel = "Density")

hVals = [0.5,2,10]
kdeS = [kde(data,bandwidth=h) for h in hVals]
plot(xGrid, pdf(kdeS[1],xGrid), c = :green, label= "h=$(hVals[1])")
plot!(xGrid, pdf(kdeS[2],xGrid), c = :blue, label= "h=$(hVals[2])")
p2 = plot!(xGrid, pdf(kdeS[3],xGrid), c = :purple, label= "h=$(hVals[3])",
    xlims=(-20,80), ylims=(0,0.035), legend=:topleft,
    xlabel="X", ylabel = "Density")
plot(p1,p2,size = (800,400))


function sigmoid(z)
  1 ./ (1 .+ e .^ (-z))
end

function sigmoid_grad(z)
  sigmoid(z) .* (1 - sigmoid(z))
end

function init_theta(num_input, num_output, nums_hidden)
  num_units = [num_input, nums_hidden..., num_output]
  layers = length(num_units)
  Theta = [rand(num_units[i], num_units[i-1]+1) - 0.5 for i = 2:layers]
end

function cost(Y, Y_pred)
  len = size(Y, 1)
 J = (-1.0 / len ) * sum((Y .* log(Y_pred)) + ((1-Y) .* log(1-Y_pred)))
end

function forward_prop(X, Theta)
  A = {}
  H = {}
  push!(H, X)
  push!(A, zeros(1, 1))

  for t=1:length(Theta)
    # add bias
    H[t] = [ones(size(H[t], 1), 1) H[t]]
    # pre-activation
    push!(A, H[t] * Theta[t]')
    # activation
    push!(H, sigmoid(A[t+1]))
  end
  A, H, H[end]
end

function back_prop(X, Y, Theta)
  len = size(X, 1)
  T = length(Theta)

  Theta_grad = {}
  for i=1:T
    push!(Theta_grad, zeros(size(Theta[i])))
  end

  A, H, Y_pred = forward_prop(X, Theta)

  delta_N = {}

  for t=1:T
    push!(delta_N, {})
  end

  delta = Y_pred - Y

  delta_N[T] = delta

  for t=0:T-2
    delta = (delta*Theta[T-t][:, 2:end]) .* sigmoid_grad(A[T-t])
    delta_N[T-t-1] = delta
  end

  for t=1:T
    Theta_grad[t] = Theta_grad[t] + delta_N[t]' * H[t]
  end

  Y_pred, Theta_grad
end

function fit(X, Y, Theta, learning_rate, epochs)
  J_list = zeros(epochs)
  for i=1:epochs
    Y_pred, Theta_grad = back_prop(X, Y, Theta)
    J_list[i] = cost(Y_train, Y_pred)
    for t=1:length(Theta)
      Theta[t] = Theta[t] - learning_rate*Theta_grad[t]
    end
  end
  Theta, J_list
end

function predict(X, Theta)
  _, _, y = forward_prop(X, Theta)
  y
end


include("../../util/src/util_module.jl")

module dev

module matrix_vector_prod
using LinearAlgebra

function matrix_vector_prod_default(A::Matrix{T}, x::Vector{T}) where {T<:Float64}
    sz_a = size(A)
    sz_b = size(x)

    if length(sz_a) != 2
        error("A is not a matrix")
    end
    if length(sz_b) != 1
        error("B is not a vector")
    end
    if sz_a[2] != sz_b[1]
        error("A and B are not compatible (size mismatch) ")
    end
    c = A * x
    return c
end

function matrix_vector_prod_inner(A, x)
    # A in R^m-by-p
    # x in R^p-by-1
    # c = A * x in R^m-by-1
    sz_a = size(A)
    sz_b = size(x)

    if length(sz_a) != 2
        error("A is not a matrix")
    end
    if length(sz_b) != 1
        error("B is not a vector")
    end
    if sz_a[2] != sz_b[1]
        error("A and B are not compatible (size mismatch) ")
    end

    c = zeros(eltype(A), sz_a[1])
    @inbounds @simd for ii = 1:sz_a[1]
        # c[ii] = A[ii, :]' * x
        c[ii] =  dot( @view(A[ii, :]), x)
    end
    return c
end

function matrix_vector_prod_inner_loop(A, x)
    # A in R^m-by-p
    # x in R^p-by-1
    # c = A * x in R^m-by-1
    sz_a = size(A)
    sz_b = size(x)

    if length(sz_a) != 2
        error("A is not a matrix")
    end
    if length(sz_b) != 1
        error("B is not a vector")
    end
    if sz_a[2] != sz_b[1]
        error("A and B are not compatible (size mismatch) ")
    end

    c = zeros(eltype(A), sz_a[1])
    @inbounds @simd for ii = 1:sz_a[1]
        inner_sum = zero(eltype(x))
        @inbounds @simd for jj = 1:sz_a[2]
            inner_sum += A[ii, jj] * x[jj]
        end
        c[ii] = inner_sum
    end
    return c
end

function matrix_vector_prod_outer(A, x)
    # A in R^m-by-p
    # x in R^p-by-1
    # c = A * x in R^m-by-1
    sz_a = size(A)
    sz_b = size(x)

    if length(sz_a) != 2
        error("A is not a matrix")
    end
    if length(sz_b) != 1
        error("B is not a vector")
    end
    if sz_a[2] != sz_b[1]
        error("A and B are not compatible (size mismatch) ")
    end

    c = zeros(eltype(A), sz_a[1])
    @inbounds @simd for jj = 1:sz_a[2]
        c .+= @view( A[:, jj] ) .* x[jj]
    end
    return c
end

function matrix_vector_prod_outer_loop(A, x)
    # A in R^m-by-p
    # x in R^p-by-1
    # c = A * x in R^m-by-1
    sz_a = size(A)
    sz_b = size(x)

    if length(sz_a) != 2
        error("A is not a matrix")
    end
    if length(sz_b) != 1
        error("B is not a vector")
    end
    if sz_a[2] != sz_b[1]
        error("A and B are not compatible (size mismatch) ")
    end

    c = zeros(eltype(A), sz_a[1])
    @inbounds @simd for jj = 1:sz_a[2]
        @inbounds @simd for ii = 1:sz_a[1]
            c[ii] += A[ii, jj] * x[jj]
        end
    end
    return c
end

function matrix_mult_inner(A, B)
    # A in R^m-by-p
    # B in R^p-by-n
    # C = A * B in R^m-by-n
    # [ma, p] = size(A)
    # [, p] = size(A)
    sz_a = size(A)
    sz_b = size(B)

    if length(sz_a) != 2
        error("A is not a matrix")
    end
    if length(sz_b) != 2
        error("B is not a matrix")
    end
    if sz_a[2] != sz_b[1]
        error("A and B are not compatible (size mismatch) ")
    end

    C = zeros(eltype(A), sz_a[1], sz_b[2])
    for ii = 1:sz_a[1]
        for jj = 1:sz_b[2]
            C[ii, jj] = dot( A[ii, :],  B[:, jj] )
            # C[ii, jj] = A[ii, :]' *  B[:, jj]
        end
    end
    return C
end

function matrix_mult_outer(A, B)
    # A in R^m-by-p
    # B in R^p-by-n
    # C = A * B in R^m-by-n
    # [ma, p] = size(A)
    # [, p] = size(A)
    sz_a = size(A)
    sz_b = size(B)

    if length(sz_a) != 2
        error("A is not a matrix")
    end
    if length(sz_b) != 2
        error("B is not a matrix")
    end
    if sz_a[2] != sz_b[1]
        error("A and B are not compatible (size mismatch) ")
    end

    C = zeros(eltype(A), sz_a[1], sz_b[2])
    for kk = 1:sz_a[2]
        C += A[:, kk] * B[kk, :]'
    end
    return C
end

function matrix_mult_default(A, B)
    # A in R^m-by-p
    # B in R^p-by-n
    # C = A * B in R^m-by-n
    # [ma, p] = size(A)
    # [, p] = size(A)
    sz_a = size(A)
    sz_b = size(B)

    if length(sz_a) != 2
        error("A is not a matrix")
    end
    if length(sz_b) != 2
        error("B is not a matrix")
    end
    if sz_a[2] != sz_b[1]
        error("A and B are not compatible (size mismatch) ")
    end
    C = A * B
    return C
end

function matrix_mult_inner_loop(A, B)
    # A in R^m-by-p
    # B in R^p-by-n
    # C = A * B in R^m-by-n
    # [ma, p] = size(A)
    # [, p] = size(A)
    sz_a = size(A)
    sz_b = size(B)

    if length(sz_a) != 2
        error("A is not a matrix")
    end
    if length(sz_b) != 2
        error("B is not a matrix")
    end
    if sz_a[2] != sz_b[1]
        error("A and B are not compatible (size mismatch) ")
    end
    C = zeros(eltype(A), sz_a[1], sz_b[2])

    for ii = 1:sz_a[1]
        for jj = 1:sz_b[2]
            inner_sum = zero(eltype(A))
            for kk = 1:sz_a[2]
                inner_sum += A[ii, kk] *  B[kk, jj]
            end
            C[ii, jj] = inner_sum
        end
    end
    return C
end

function matrix_mult_outer_loop(A, B)
    #
    # Algo probably wrong
    #
    # A in R^m-by-p
    # B in R^p-by-n
    # C = A * B in R^m-by-n
    # [ma, p] = size(A)
    # [, p] = size(A)
    sz_a = size(A)
    sz_b = size(B)

    if length(sz_a) != 2
        error("A is not a matrix")
    end
    if length(sz_b) != 2
        error("B is not a matrix")
    end
    if sz_a[2] != sz_b[1]
        error("A and B are not compatible (size mismatch) ")
    end
    C = zeros(eltype(A), sz_a[1], sz_b[2])

    #=
    for kk = 1:sz_a[2]
        for ii = 1:sz_a[1]
            C[ii, kk] += A[ii, kk] * B[kk, ii]
        end
    end
    =#
    for ii = 1:sz_b[2]
        for kk = 1:sz_a[2]
            C[:, ii] += A[:, kk] * B[kk, ii]
        end
    end

    return C
end

end # dev

module devtest
using ..dev
using BenchmarkTools

function benchmark_matrix_vector_prod_default(A, x)
    b = @benchmark dev.matrix_vector_prod_default($A, $x)

    println("=== matrix_vector_prod_default ===")
    show(
            stdout,
            MIME("text/plain"),
            b )
    println("\n\n")
end

function benchmark_matrix_vector_prod_inner(A, x)
    b = @benchmark dev.matrix_vector_prod_inner($A, $x)

    println("=== matrix_vector_prod_inner ===")
    show(
            stdout,
            MIME("text/plain"),
            b )
    println("\n\n")
end

function benchmark_matrix_vector_prod_outer(A, x)
    b = @benchmark dev.matrix_vector_prod_outer($A, $x)

    println("=== matrix_vector_prod_outer ===")
    show(
            stdout,
            MIME("text/plain"),
            b )
    println("\n\n")
end

function benchmark_matrix_vector_prod_inner_loop(A, x)
    b = @benchmark dev.matrix_vector_prod_inner_loop($A, $x)

    println("=== matrix_vector_prod_inner_loop ===")
    show(
            stdout,
            MIME("text/plain"),
            b )
    println("\n\n")
end

function benchmark_matrix_vector_prod_outer_loop(A, x)
    b = @benchmark dev.matrix_vector_prod_outer_loop($A, $x)

    println("=== matrix_vector_prod_outer_loop ===")
    show(
            stdout,
            MIME("text/plain"),
            b )
    println("\n\n")
end

function benchmark_all()

    # size definitions
    # m = 1000
    # p = 5000

    # m = 10
    # p = 500000

    m = 100000
    p = 50

    # allocate the matrices
    A = rand(m, p)
    x = rand(p)

    benchmark_matrix_vector_prod_default(A, x)
    benchmark_matrix_vector_prod_inner(A, x)
    benchmark_matrix_vector_prod_outer(A, x)
    benchmark_matrix_vector_prod_inner_loop(A, x)
    benchmark_matrix_vector_prod_outer_loop(A, x)

end

end

module devtest_matrix
    using ..dev
    using BenchmarkTools

function benchmark_matrix_mult_inner(A, B)
    b = @benchmark dev.matrix_mult_inner($A, $B)

    println("=== matrix_mult_inner ===")
    show(
            stdout,
            MIME("text/plain"),
            b )
    println("\n\n")
   end

function benchmark_matrix_mult_outer(A, B)
    b = @benchmark dev.matrix_mult_outer($A, $B)

    println("=== matrix_mult_outer ===")
    show(
            stdout,
            MIME("text/plain"),
            b )
    println("\n\n")
end

function benchmark_matrix_mult_default(A, B)
    b = @benchmark dev.matrix_mult_default($A, $B)

    println("=== matrix_mult_default ===")
    show(
            stdout,
            MIME("text/plain"),
            b )
    println("\n\n")
end

function benchmark_matrix_mult_inner_loop(A, B)
    b = @benchmark dev.matrix_mult_inner_loop($A, $B)

    println("=== matrix_mult_inner_loop ===")
    show(
            stdout,
            MIME("text/plain"),
            b )
    println("\n\n")
end

function benchmark_matrix_mult_outer_loop(A, B)
    b = @benchmark dev.matrix_mult_outer_loop($A, $B)

    println("=== matrix_mult_outer_loop ===")
    show(
            stdout,
            MIME("text/plain"),
            b )
    println("\n\n")
end

function benchmark_all()
    # size definitions
    m = 100
    p = 500
    n = 100

    # allocate the matrices
    A = rand(m, p)
    B = rand(p, n)

    benchmark_matrix_mult_inner(A, B)
    benchmark_matrix_mult_outer(A, B)
    benchmark_matrix_mult_default(A, B)
    benchmark_matrix_mult_inner_loop(A, B)
    # benchmark_matrix_mult_outer_loop(A, B)
end

    # multiply
    #=
    Ci = dev.matrix_mult_inner(A, B)
    Co = dev.matrix_mult_outer(A, B)
    Cd = A * B;

    println( all( abs.( Ci .- Co ) .<= 1e-10 ) ) =#

end

#=
module plottest

using PyPlot
PyPlot.pygui(true)

x = LinRange(-2, 2, 30) # [1, 2, 3]
y = LinRange(-4, 1, 30) # [10, 20, 30]
z = [i.^2 + j.^2 for i in x, j in y]
z2 = [50, 51, 52]

PyPlot.figure()
# PyPlot.plot3D(x, y, z2)
PyPlot.plot_surface(x, y, z, alpha = 0.5)

end =#


using Distributions, Plots; pyplot()

triangDist = TriangularDist(0,2,1)
xGrid = 0:0.1:2
N = 10^6
inverseSampledData = quantile.(triangDist,rand(N))

histogram( inverseSampledData, bins=30, normed=true,
	ylims=(0,1.1), label="Inverse transform data")
plot!( xGrid, pdf.(triangDist,xGrid), c=:red, lw=4,
	xlabel="x", label="PDF", ylabel = "Density", legend=:topright)


## compose_matrix.jl : methods to compose example matrices

"""
anisotropic2d(ϵ, n, m)

2d anisotropic Laplacian on [0,1]^2 using n x m points and anisotropy number ϵ
"""
anisotropic2d(ϵ, n, m) = stencil2mat(SArray{Tuple{3,3}}([0 -m^2 0 -ϵ*n^2 2(ϵ*n^2+m^2) -ϵ*n^2 0 -m^2 0]), n, m)

"""
anisotropic2d(ϵ, β, n, m)

2d rotated anisotropic Laplacian on [0,1]^2 using n x m points with anisotropy number ϵ and angle β
"""
anisotropic2d(ϵ, β, n, m) = stencil2mat(SArray{Tuple{3,3}}([0 0 0 -n^2*(ϵ*cos(β)^2+sin(β)^2) 2n^2*(ϵ*cos(β)^2+sin(β)^2) -n^2*(ϵ*cos(β)^2+sin(β)^2) 0 0 0]), n, m) + stencil2mat(SArray{Tuple{3,3}}([n*m/2*((ϵ-1)*cos(β)*sin(β)) 0 -n*m/2*((ϵ-1)*cos(β)*sin(β)) 0 0 0 -n*m/2*((ϵ-1)*cos(β)*sin(β)) 0 n*m/2*((ϵ-1)*cos(β)*sin(β))]), n, m) + stencil2mat(SArray{Tuple{3,3}}([0 -m^2*(ϵ*sin(β)^2+cos(β)^2) 0 0 2m^2*(ϵ*sin(β)^2+cos(β)^2) 0 0 -m^2*(ϵ*sin(β)^2+cos(β)^2) 0]), n, m)

"""
anisotropic3d(ϵ₁, ϵ₂, n, m, l)

3d anisotropic Laplacian on [0,1]^3 using n x m x l points and anisotropy numbers ϵ₁ and ϵ₂
"""
anisotropic3d(ϵ₁, ϵ₂, n, m, l) = stencil2mat(SArray{Tuple{3,3,3}}([0 0 0 0 -l^2 0 0 0 0 0 -ϵ₂*m^2 0 -ϵ₁*n^2 2(ϵ₁*n^2+ϵ₂*m^2+l^2) -ϵ₁*n^2 0 -ϵ₂*m^2 0 0 0 0 0 -l^2 0 0 0 0]), n, m, l)


"""
    _orientation(x::Tuple{<:Number, <:Number}, y::Tuple{<:Number, <:Number}, z::Tuple{<:Number, <:Number})

Utility function for `convex_hull`. Given an ordered triplet, this function returns if three
points are collinear, oriented in clock-wise or anticlock-wise direction.
"""
function _orientation(x::Tuple{<:Number, <:Number}, y::Tuple{<:Number, <:Number}, z::Tuple{<:Number, <:Number})

    a = ((y[2] - x[2]) * (z[1] - y[1]) - (y[1] - x[1]) * (z[2] - y[2]))

    if isapprox(a, 0, atol=1e-6)
        return 0 # collinear
    elseif a > 0
        return 1 # clock-wise
    else
        return 2 # counterclock-wise
    end

end

"""
    _lt_filter(a::Tuple{<:Number, <:Number}, b::Tuple{<:Number, <:Number})

Utility function for sorting step in `convex_hull`. Given two points, `a` and `b`, this function
returns true if `a` has larger polar angle (counterclock-wise direction) than `b` w.r.t. first point `chull_p1`.
"""
function _lt_filter(a::Tuple{<:Number, <:Number}, b::Tuple{<:Number, <:Number})
    o = QCO._orientation(chull_p1, a, b)

    if o == 0
        if LA.norm(collect(chull_p1 .- b))^2 >= LA.norm(collect(chull_p1 .- a))^2
            return true
        else
            return false
        end
    elseif o == 2
        return true
    else
        return false
    end

end

"""
    convex_hull(points::Vector{T}) where T<:Tuple{Number, Number}

Graham's scan algorithm to compute the convex hull of a finite set of `n` points in a plane
with time complexity `O(n*log(n))`. Given a vector of tuples of co-ordinates, this function returns a
vector of tuples of co-ordinates which form the convex hull of the given set of points.

Sources: https://doi.org/10.1016/0020-0190(72)90045-2
         https://en.wikipedia.org/wiki/Graham_scan
"""
function convex_hull(points::Vector{T}) where T<:Tuple{Number, Number}

    num_points = size(points)[1]

    if num_points == 0
        Memento.error(_LOGGER, "Atleast one point is necessary for evaluating the convex hull")
    elseif num_points <= 2
        return points
    end

    min_y = points[1][2]
    min = 1

    # Bottom-most or else the left-most point when tie
    for i in 2:num_points
        y = points[i][2]
        if (y < min_y) || (y == min_y && points[i][1] < points[min][1])
            min_y = y
            min = i
        end
    end

    # Placing the bottom/left-most point at first position
    points[1], points[min] = points[min], points[1]

    global chull_p1 = points[1]

    points = Base.sort(points, lt = QCO._lt_filter)

    # If two or more points are collinear with chull_p1, remove all except the farthest from chull_p1
    arr_size = 2
    for i in 2:num_points
        while (i < num_points) && (QCO._orientation(chull_p1, points[i], points[i+1]) == 0)
            i += 1
        end
        points[arr_size] = points[i]
        arr_size += 1
    end

    if arr_size <= 3
        return []
    end

    chull = []
    append!(chull, [points[1]])
    append!(chull, [points[2]])
    append!(chull, [points[3]])

    for i in 4:(arr_size-1)
        while (size(chull)[1] > 1) && (QCO._orientation(chull[end - 1], chull[end], points[i]) != 2)
            pop!(chull)
        end
        append!(chull, [points[i]])
    end

    return chull
end


## Exercise 4-7
## Make a more general version of circle called arc that takes an additional parameter angle, which determines what fraction of a circle to draw. angle is in units of degrees, so when angle = 360, arc should draw a complete circle.
using ThinkJulia

println("Ans: ")

function arc(turtle::Turtle, radius::Int, angle)
    len = 2 * π * radius * angle / 360  # finding the length of the arc 2πr * a/360
    n = Int(trunc(len) / 3) + 1 # number of times the turtle will rotate and draw

    len /= n
    angle /= n

    for i in 1:n
        forward(turtle, len)
        turn(turtle, -angle)
    end
end

@svg begin
    arc(Turtle(), 200, 270)
end

println("End.")


using Symbolics
using SatisfiabilityInterface

vars = @variables x, y, z, w

constraints = [
    x ∈ 1:2
    y ∈ 2:5
    z ∈ 3:7
    z == x * y
]

prob = DiscreteCSP(constraints)

solve(prob)

all_solutions(prob)

###

vars = @variables x, y, z, w

constraints = [
    x ∈ 1:2
    y ∈ 2:5
    z ∈ 3:7
    w ∈ 4:6
    z == x + y
    w == z + x
]

prob = DiscreteCSP(constraints)
solve(prob)
all_solutions(prob)

###

constraints = [
    x ∈ 1:2
    y ∈ 2:5
    z ∈ 3:7
    z == 2x
]

all_solutions(DiscreteCSP(constraints))

###

vars = @variables x, y, z

# constraints = [
#     x ∈ 1:4
#     y ∈ 2:5
#     z == x + y
#     z ≤ 4
# ]

### Subset sum Problem

# https://en.wikipedia.org/wiki/Subset_sum_problem

# Given a set of integers like {−7, −3, −2, 5, 8},
# do a subset sum to 0?

s = [-7, -3, -2, 5, 8]
n = length(s)
x = [Num(Variable(:x, i)) for i in 1:n]

constraints = [
    [x[i] ∈ (0:1) for i in 1:n]
    sum(x[i] * s[i] for i in 1:n) == 0
    sum(x) ≠ 0   # at least one is not 0
    ]

prob = DiscreteCSP(constraints)

all_solutions(prob)

#####

@variables x, y, z

constraints = [
    x ∈ 1:2
    y ∈ 2:5
    z ∈ 3:7
    z == 2x
]

prob = DiscreteCSP(constraints)
all_solutions(prob)

## Example from https://github.com/newptcai/BeeEncoder.jl

function make_vector(name::Symbol, range)
    return Num.(Variable.(name, range))
end

@variables x, y, z, w
b = make_vector(:b, 1:4)
# @variables b[1:4]

¬(x::Num) = 1 - x   # assumes a Boolean variable

constraints = [
    x ∈ 0:5
    y ∈ -4:9
    z ∈ -5:10
    w ∈ 0:10

    [bb ∈ 0:1 for bb ∈ b]

    x + y == z

    b[1] == ¬b[2]
    b[2] == 1

    ¬b[1] + b[2] + ¬b[3] + b[4] == w
]

prob = DiscreteCSP(constraints)
solve(prob)

# @time length(all_solutions(prob))   # 1184 solutions in 38 seconds

###

constraints = [
    x ∈ -5:5
    y ∈ -4:9
    x - y^2 == 3
    # sum(b) == w

]

prob = DiscreteCSP(constraints)
all_solutions(prob)

###

@variables x, y

constraints = [
    x ∈ -5:5
    y ∈ -4:9
    10x + y == 9
    # sum(b) == w
]

prob = DiscreteCSP(constraints)
all_solutions(prob)


#=
It is possible to show that the square root of two can be expressed as an infinite continued fraction.

√ 2 = 1 + 1/(2 + 1/(2 + 1/(2 + ... ))) = 1.414213...

By expanding this for the first four iterations, we get:

1 + 1/2 = 3/2 = 1.5
1 + 1/(2 + 1/2) = 7/5 = 1.4
1 + 1/(2 + 1/(2 + 1/2)) = 17/12 = 1.41666...
1 + 1/(2 + 1/(2 + 1/(2 + 1/2))) = 41/29 = 1.41379...

The next three expansions are 99/70, 239/169, and 577/408, but the eighth expansion, 1393/985, is the first example where the number of digits in the numerator exceeds the number of digits in the denominator.

In the first one-thousand expansions, how many fractions contain a numerator with more digits than denominator?
=#
function calc()
  p = BigInt(1)
  q = BigInt(1)
  count = 1
  total = 0
  while count <= 1000
    np = p + 2q
    nq = p + q
    p,q = np,nq
    if length(string(p)) > length(string(q))
      total = total + 1
    end
    count = count + 1
  end
  total
end
@time println(calc())


# nuclear Norm (times a constant)

export NuclearNorm

"""
**Nuclear norm**

    NuclearNorm(λ=1.0)

Returns the function
```math
f(X) = \\|X\\|_* = λ ∑_i σ_i(X),
```
where `λ` is a positive parameter and ``σ_i(X)`` is ``i``-th singular value of matrix ``X``.
"""
struct NuclearNorm{R <: Real} <: ProximableFunction
  lambda::R
  function NuclearNorm{R}(lambda::R) where {R <: Real}
    if lambda < 0
      error("parameter λ must be nonnegative")
    else
      new(lambda)
    end
  end
end

is_convex(f::NuclearNorm) = true

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

function (f::NuclearNorm{R})(X::AbstractMatrix{T}) where {R <: Real, T <: Union{R, Complex{R}}}
  F = svd(X);
  return f.lambda * sum(F.S);
end

function prox!(Y::AbstractMatrix{T}, f::NuclearNorm{R}, X::AbstractMatrix{T}, gamma::R=one(R)) where {R <: Real, T <: Union{R, Complex{R}}}
  F = svd(X)
  S_thresh = max.(zero(R), F.S .- f.lambda*gamma)
  rankY = findfirst(S_thresh .== zero(R))
  if rankY === nothing
    rankY = minimum(size(X))
  end
  Vt_thresh = view(F.Vt, 1:rankY, :)
  U_thresh = view(F.U, :, 1:rankY)
  # TODO: the order of the following matrix products should depend on the shape of x
  M = S_thresh[1:rankY] .* Vt_thresh
  mul!(Y, U_thresh, M)
  return f.lambda * sum(S_thresh);
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(f::NuclearNorm, X::AbstractMatrix{T}, gamma=1.0) where T
  F = svd(X)
  S = max.(0, F.S .- f.lambda*gamma)
  Y = F.U * (Diagonal(S) * F.Vt)
  return Y, f.lambda * sum(S)
end


using Plots
using Printf

"""
    untildivergence(f, x, maxiter=50)

Iterate `f` on `x` until its module is greater than 2 or `maxiter` iterations have been done.

Returns the number of iterations.
"""
function untildivergence(f, x, maxiter = 50)
    i = 0
    while abs(x) < 2 && (i += 1) <= maxiter
        x = f(x)
    end
    return i
end

"""
    julia_set_wh(f; kwargs)

Approximate the Julia set of `f` in a grid using the escape time algorithm.

The grid is defined by its bounds `xmin`, `xmax`, `ymin` and `ymax` as well as its width `w` and height `h` in pixels. The optional parameter `maxiter` is passed on to `untildivergence`.
"""
function julia_set_wh(f; xmin::Real=-2, xmax::Real=2, ymin::Real=-2, ymax::Real=2, w::Int=1000, h::Int=1000, maxiter::Int=50)
    img = Array{Int,2}(undef, h, w)
    x = LinRange(xmin, xmax, w)
    y = LinRange(ymin, ymax, h)
    # creating labels for use with `heatmap`
    xticks = LinRange(1,w,11)
    yticks = LinRange(1,h,11)
    xlabels = map(z->@sprintf("%.2f",z), LinRange(xmin,xmax,11))
    ylabels = map(z->@sprintf("%.2f",z), LinRange(ymin,ymax,11))

    for i in 1:w, j in 1:h
        img[j,i] = untildivergence(f, x[i]+im*y[j], maxiter)
    end
    return img, xticks, xlabels, yticks, ylabels
end

"""
    julia_set(f, cellsize; kwargs)

Approximate the Julia of `f` set in a grid of cells of length `cellsize`.

The grid is defined by the same parameters as in `julia_set_wh`, except its width and height in pixels are computed with `cellsize`.
"""
function julia_set(f, cellsize::Real; xmin=-2, xmax=2, ymin=-2, ymax=2, maxiter=50)
    w = Integer(round(abs(xmax-xmin)/cellsize))
    h = Integer(round(abs(ymax-ymin)/cellsize))
    return julia_set_wh(f, xmin=xmin, xmax=xmax, ymin=ymin, ymax=ymax, w=w, h=h, maxiter=maxiter)
end

"""
    mandelbrot_set_wh(f; kwargs)

Approximate the Mandelbrot set of `f` in a grid using the escape time algorithm.

The grid is defined by its bounds `xmin`, `xmax`, `ymin` and `ymax` as well as its width `w` and height `h` in pixels. The optional parameter `maxiter` is passed on to `untildivergence`.
"""
function mandelbrot_set_wh(f; xmin::Real=-2, xmax::Real=2, ymin::Real=-2, ymax::Real=2, w::Integer=1000, h::Integer=1000, maxiter::Integer=50)
    img = Array{Int,2}(undef, h, w)
    x = LinRange(xmin, xmax, w)
    y = LinRange(ymin, ymax, h)
    # creating labels for use with `heatmap`
    xticks = LinRange(1,w,11)
    yticks = LinRange(1,h,11)
    xlabels = map(z->@sprintf("%.2f",z), LinRange(xmin,xmax,11))
    ylabels = map(z->@sprintf("%.2f",z), LinRange(ymin,ymax,11))

    for i in 1:w, j in 1:h
        img[j,i] = untildivergence(f(x[i]+im*y[j]), 0, maxiter)
    end
    return img, xticks, xlabels, yticks, ylabels
end

"""
    mandelbrot_set(f, cellsize; kwargs)

Approximate the Mandelbrot of `f` set in a grid of cells of length `cellsize`.

The grid is defined by the same parameters as in `julia_set_wh`, except its width and height in pixels are computed with `cellsize`.
"""
function mandelbrot_set(f, cellsize; xmin=-2, xmax=2, ymin=-2, ymax=2, maxiter=50)
    w = Integer(round(abs(xmax-xmin)/cellsize))
    h = Integer(round(abs(ymax-ymin)/cellsize))
    return mandelbrot_set_wh(f, xmin=xmin, xmax=xmax, ymin=ymin, ymax=ymax, w=w, h=h, maxiter=maxiter)
end

# the traditional julia/mandelbrot function
f(c) = z -> z^2 + c

img, xticks, xlabels, yticks, ylabels = mandelbrot_set(f, 0.00005, xmin=-2.01, xmax = -1.8, ymin = -0.1, ymax = 0.1)

img, xticks, xlabels, yticks, ylabels = julia_set(f(im), 0.0005, xmin=-2, xmax=2, ymin=-2,ymax=2)

heatmap(img,
        xticks = (xticks, xlabels),
        yticks = (yticks, ylabels),
        aspect_ratio=:equal
        )

# alternative, for better looking images
using Images, Colors
g = Gray.(img/51)

g = Gray.(.!(img .>= 50))

mset = img .>= 50

#save("julia_minus2.png", g)

include("FractalDimensions.jl")

B = boundary(mset)
Gray.(B)
points = array_to_points(B)

d, N, ε = box_counting_dimension(points, 10)

C, r = correlation_dimension(points)

loglog_regression(C,r)


### Gali's Book Chapter 3 model
### We use this as a consistency check of gensys implementation
### By Daniel Coutinho and Gilberto Boaretto

#### Parameters from Gali's book example, chapter 3, page 52, first edition

bet = 0.99
sig = 1
phi = 1
alfa = 1/3
epsilon = 6
theta = 2/3
phi_pi = 1.5
phi_y = 0.5/4
rho_v = 0.5

true_pars = (alfa = alfa, bet = bet, epsilon = epsilon, theta = theta, sig = sig, s2 = 1, phi = phi, phi_pi = phi_pi, phi_y = phi_y, rho_v = rho_v)

THETA = (1-alfa)/(1-alfa+alfa*epsilon)
lamb = (1-theta)*(1-bet*theta)/theta*THETA
kappa = lamb*(sig+(phi+alfa)/(1-alfa))

## Gali's give analytical expressions for the effects of a shock in the monetary policy over inflation and output gap,chapter 3, page 51, first edition

LAMBDA_v = ((1-bet*rho_v)*(sig*(1-rho_v)+phi_y)+kappa*(phi_pi-rho_v))^(-1)

y_irf(v) = -(1-bet*rho_v)*LAMBDA_v*v
pi_irf(v) = -kappa*LAMBDA_v*v
r_irf(v) = sig*(1-rho_v)*(1-bet*rho_v)*LAMBDA_v*v
i_irf(v) = (sig*(1-rho_v)*(1-bet*rho_v) - rho_v*kappa)*LAMBDA_v*v
v_irf(v,uu) = rho_v*v + uu

## Test drive: lets see if this gives the same IRFs that we get in Gali's book
## This is just for debugging, so comment it out when using it to check gensys consistency

#using Plots

irfs_true = zeros(15,5)

irfs_true[1,1] = 0

uu = 0.25

for t in 2:15
    irfs_true[t,4] = v_irf(irfs_true[t-1,4],uu)
    irfs_true[t,1] = pi_irf(irfs_true[t,4])
    irfs_true[t,2] = y_irf(irfs_true[t,4])
    irfs_true[t,3] = i_irf(irfs_true[t,4])
    irfs_true[t,5] = r_irf(irfs_true[t,4])
    global uu = 0
end

#plot(irfs[2:13,1], label = "v")
#plot(irfs[2:13,2]*4, label = "pi")
#plot(irfs[2:13,3], label = "y")
#plot(4*irfs[2:13,4], label = "i")
#plot(4*irfs[2:13,5], label = "r")


# AAA algorithm from the paper "The AAA Alorithm for Rational Approximation"
# by Y. Nakatsukasa, O. Sete, and L.N. Trefethen, SIAM Journal on Scientific Computing
# 2018

struct AAAapprox{T <: AbstractArray} <: BRInterp
    x::T
    f::T
    w::T
    errvec::T
end

(a::AAAapprox)(zz) = reval(zz, a.x, a.f, a.w)

# Handle function inputs as well
function aaa(Z::AbstractArray{T,1}, F::S;  tol=1e-13, mmax=100, verbose=false, clean=true) where {T, S<:Function}
    aaa(Z, F.(Z), tol=tol, mmax=mmax, verbose=verbose, clean=clean)
end

"""aaa  rational approximation of data F on set Z
        r = aaa(Z, F; tol, mmax, verbose, clean)

 Input: Z = vector of sample points
        F = vector of data values at the points in Z
        tol = relative tolerance tol, set to 1e-13 if omitted
        mmax: max type is (mmax-1, mmax-1), set to 100 if omitted
        verbose: print info while calculating default = false
        clean: detect and remove Froissart doublets default = true

 Output: r = an AAA approximant as a callable struct with fields
         z, f, w = vectors of support pts, function values, weights
         errvec = vector of errors at each step

 NOTE: changes from matlab version:
         switched order of Z and F in function signature
         added verbose and clean boolean flags
         pol, res, zer = vectors of poles, residues, zeros are now only calculated on demand
         by calling prz(z::AAAapprox)
"""
function aaa(Z::AbstractArray{T,1}, F::AbstractArray{S,1}; tol=1e-13, mmax=100,
             verbose=false, clean=true) where {S, T}
    # filter out any NaN's or Inf's in the input
    keep = isfinite.(F)
    F = F[keep]
    Z = Z[keep]

    M = length(Z)                    # number of sample points
    mmax = min(M, mmax)              # max number of support points

    reltol = tol*norm(F, Inf)
    SF = spdiagm(M, M, 0 => F)       # left scaling matrix

    F, Z = promote(F, Z)
    P = promote_type(S, T)

    J = [1:M;]
    z = P[]                          # support points
    f = P[]                          # function values at support points
    C = P[]
    w = P[]

    errvec = P[]
    R = F .- mean(F)
    @inbounds for m = 1:mmax
        j = argmax(abs.(F .- R))               # select next support point
        push!(z, Z[j])
        push!(f, F[j])
        deleteat!(J, findfirst(isequal(j), J))   # update index vector

        # next column of Cauchy matrix
        C = isempty(C) ? reshape((1 ./ (Z .- Z[j])), (M,1)) : [C (1 ./ (Z .- Z[j]))]

        Sf = diagm(f)                         # right scaling matrix
        A = SF * C - C * Sf                   # Loewner matrix
        G = svd(A[J, :])

        # A[J, :] might not have full rank, so svd can return fewer than m columns
        #w = G.V[:, m]
        w = G.V[:, end]                       # weight vector = min sing vector

        N = C * (w .* f)                      # numerator
        D = C * w
        R .= F
        R[J] .= N[J] ./ D[J]                  # rational approximation

        err = norm(F - R, Inf)
        verbose && println("Iteration: ", m, "  err: ", err)
        errvec = [errvec; err]                # max error at sample points
        err <= reltol && break                # stop if converged
    end
    r = AAAapprox(z, f, w, errvec)

    # remove Frois. doublets if desired.  We do this in place
    if clean
        pol, res, zer = prz(r)            # poles, residues, and zeros
        ii = findall(abs.(res) .< 1e-13)  # find negligible residues
        length(ii) != 0 && cleanup!(r, pol, res, zer, Z, F)
    end
    return r
end

function prz(r::AAAapprox)
    z, f, w = r.x, r.f, r.w
    m = length(w)
    B = diagm(ones(m+1))
    B[1, 1] = 0.0
    E = [0.0  transpose(w); ones(m) diagm(z)]
    pol, _ = eigen(E, B)
    pol = pol[isfinite.(pol)]
    dz = 1e-5 * exp.(2im*pi*[1:4;]/4)

    # residues
    res = r(pol .+ transpose(dz)) * dz ./ 4

    E = [0 transpose(w .* f); ones(m) diagm(z)]
    zer, _ = eigen(E, B)
    zer = zer[isfinite.(zer)]
    pol, res, zer
end

function reval(zz, z, f, w)
    # evaluate r at zz
    zv = size(zz) == () ? [zz] : vec(zz)
    CC = 1.0 ./ (zv .- transpose(z))           # Cauchy matrix
    r = (CC * (w .* f)) ./ (CC * w)            # AAA approx as vector
    r[isinf.(zv)] .= sum(f .* w) ./ sum(w)

    ii = findall(isnan.(r))                    # find values NaN = Inf/Inf if any
    @inbounds for j in ii
        if !isnan(zv[j]) && ((v = findfirst(isequal(zv[j]), z)) !== nothing)
            r[j] = f[v]  # force interpolation there
        end
    end
    r = size(zz) == () ? r[1] : reshape(r, size(zz))             # the AAA approximation
end

# Only calculate the updated z, f, and w
function cleanup!(r, pol, res, zer, Z, F)
    z, f, w = r.x, r.f, r.w
    m = length(z)
    M = length(Z)
    ii = findall(abs.(res) .< 1e-13)  # find negligible residues
    ni = length(ii)
    ni == 0 && return
    println("$ni Froissart doublets. Number of residues = ", length(res))

    @inbounds for j = 1:ni
        azp = abs.(z .- pol[ii[j]] )
        jj = findall(isequal(minimum(azp)), azp)
        deleteat!(z, jj)    # remove nearest support points
        deleteat!(f, jj)
    end

    @inbounds for j = 1:length(z)
        jj = findall(isequal(z[j]), Z)
        deleteat!(F, jj)
        deleteat!(Z, jj)
    end
    m = m - length(ii)
    SF = spdiagm(M-m, M-m, 0 => F)
    Sf = diagm(f)
    C = 1 ./ (Z .- transpose(z))
    A = SF*C - C*Sf
    G = svd(A)
    w[:] .= G.V[:, m]
    println("cleanup: ", size(z), "  ", size(f), "  ", size(w))
    return nothing
end


#############################
##  Numerical experiments associated with section 5.4.1
#############################
## Nonlinear equations are treated by first defining an objective function, then using standard iterative methods
## such as Newton's method via NLsolve.jl. Using more sophisticated methods, linesearches etc. can sometimes result in significant speed boosts.
##

using NLsolve, ApproxFun, MultivariateOrthogonalPolynomials, BandedMatrices, BlockBandedMatrices, Plots
using SparseVolterraExamples

#####################################################
## Problem in Equation (25), u_1(x) = exp(x)
#####################################################
####
## Define g(x), the nonlinearity and the kernel
g(x) = exp(x)+1/3*x*(1-exp(3*x))
    gFun = Fun(g,Jacobi(0,1,0..1))
    f(u) = u^3
    K(x,y) = x
    n = 15 # Desired polynomial order of approximation for accuracy
####
## Define the objective function and compute the solution using a standard Newton method
objective(x) = triNonLinearVolterraObjective(x,f,K,gFun,n)
sol = nlsolve(objective,zeros(n),method=:newton,ftol=1e-8)
####
## Plot function aganinst analytic solution
plot(Fun(Jacobi(0,1,0..1),sol.zero),grid=false,legend=:topleft,xlabel="x",ylabel="u(x)",label="sparse method")
plot!(exp,0,1,grid=false,legend=:topleft,xlabel="x",ylabel="u(x)",label="analytic exp(x)")
####
## Or plot error. A simple loop over polynomial degree can then produce Figure 7(a).
plot(Fun(Jacobi(0,1,0..1),sol.zero)-Fun(x->exp(x),Jacobi(0,1,0..1)),grid=false,legend=:bottomleft,xlabel="x",ylabel="error",label=false)

#####################################################
## Problem in Equation (26), u_2(x) = sin(x)
#####################################################
####
## Define g(x), the nonlinearity and the kernel
g(x) = sin(x)+1/4*sin(x)^2-1/4*x^2
    gFun = Fun(g,Jacobi(0,1,0..1))
    f(u) = u^2
    K(x,y) = x-y
    n = 14 # Desired polynomial order of approximation for accuracy
####
## Define the objective function and compute the solution using a standard Newton method
objective(x) = triNonLinearVolterraObjective(x,f,K,gFun,n)
sol = nlsolve(objective,zeros(n),method=:newton,ftol=1e-15)
####
## Plot function aganinst analytic solution
plot(Fun(Jacobi(0,1,0..1),sol.zero),grid=false,legend=:topleft,xlabel="x",ylabel="u(x)",label="sparse method")
plot!(sin,0,1,grid=false,legend=:topleft,xlabel="x",ylabel="u(x)",label="analytic sin(x)")
####
## Or plot error. A simple loop over polynomial degree can then produce Figure 7(b).
plot(Fun(Jacobi(0,1,0..1),sol.zero)-Fun(x->sin(x),Jacobi(0,1,0..1)),grid=false,legend=:bottomleft,xlabel="x",ylabel="error",label=false)

examples/Section 5-4-1 - Nonlinear Volterra integral equations.jl


using LinearAlgebra

"""
    is_schur(A::AbstractMatrix)

Check if a square matrix is Schur, i.e. the eigenvalues lie within the unit circle.

### Input
- `A` -- square matrix

### Ouput
Boolean value indicating if the matrix `A` is Schur.
"""
is_schur(A::AbstractMatrix) = maximum(abs.(eigvals(A))) < 1.0


export solarposition, zenith, zenith_azimuth, sines_and_cosines

# Calculate the position of the sun as (zenith,azimuth) in degrees for a location
# on earth given as (latitude,longitude) in degrees. The DateTime should be given
# as Universal Time, e.g. UTC.
function solarposition_wiki(datetime::DateTime, latitude, longitude)
	# Position of the Sun as seen from Earth.
	# Ported from Matlab, original helpstring below.

	#   This is the most basic algorithm. It is documented in Seinfeld &
	#   Pandis, Duffie & Beckman and Wikipedia.
	#
	# [ANGLES,PROJECTION] = SOLARPOSITION(DATE,TIME,LATITUDE,LONGITUDE,TIME_ZONE)
	# returns ZENITH & AZIMUTH for all DATE & TIME pairs at LATITUDE, LONGITUDE.
	# ANGLES = [ZENITH,AZIMUTH] and PROJECTION = [PHI_X, PHI_Y]
	# PHI_X is projection on x-z plane & PHI_Y is projection on y-z plane.
	# DATETIME can be string, vector [YEAR, MONTH, DAY, HOURS, MINUTES, SECONDS],
	#   cellstring or matrix N x [YEAR, MONTH, DAY, HOURS, MINUTES, SECONDS] for N
	#   times.
	# LATITUDE [degrees] and LONGITUDE [degrees] are the coordinates of the site.
	# TIME_ZONE [hours] of the site.
	# ROTATION [degrees] clockwise rotation of system relative to north.
	# DST [logical] flag for daylight savings time, typ. from March to November
	#   in the northern hemisphere.
	#
	# References:
	# http://en.wikipedia.org/wiki/Solar_azimuth_angle
	# http://en.wikipedia.org/wiki/Solar_elevation_angle
	#
	# Mark A. Mikofski
	# Copyright (c) 2013

	# Equation of Time
	# "Alternative calculation" on wikipedia:
	# https://en.wikipedia.org/wiki/Equation_of_time#Calculating_the_equation_of_time
	# Supposedly accurate to 6 seconds (0.1 minutes)

	year = Dates.year(datetime)
	# date in days starting at zero on 1 January (the subtraction produces a difference in milliseconds)
	d_ms = datetime - DateTime(year,1,1)
	d = d_ms.value/1000/3600/24

	# Earth's mean angular orbital velocity [degrees/day]
	w = 360/365.24
	# angle the earth moves on its orbit at its average speed from the December solstice
	# 10 is the approximate number of days from the December solstice to January 1
	a = w * (d+10)
	# The angle the Earth moves from the solstice to date D, including a first-order correction
	# for the Earth's orbital eccentricity, 0.0167. The number 2 is the number of days from 1 January
	# to the date of the Earth's perihelion.
	b = a + (360/pi*0.0167)*sind(w*(d-2))
	# C is the difference between the angles moved at mean speed, and at the corrected speed projected
	# onto the equatorial plane, and divided by 180 to get the difference in "half turns". The value 23.44°
	# is the obliquity (tilt) of the Earth's axis. The subtraction gives the conventional sign to the equation of time.
	c = (a - atand(tand(b)/cosd(23.44))) / 180
	# For any given value of x, arctan x (sometimes written as tan−1 x) has multiple values, differing from each other
	# by integer numbers of half turns. This may cause C to be wrong by an integer number of half turns. The excess
	# half turns are removed in the next step of the calculation to give the equation of time:
	ET = 720*(c - round(c))

	# approximate solar time [hours]
	solarTime = 24*(d - floor(d)) + longitude*24/360 + ET/60
	t_h = 15*(solarTime - 12) # [degrees] hour angle

	# declination [degrees]
	# accurate to 0.2 degrees, see https://en.wikipedia.org/wiki/Position_of_the_Sun#Calculations
	delta = -asind(sind(23.44)*cosd(b))
	zenith = acosd(sind(latitude)*sind(delta) + cosd(latitude)*cosd(delta)*cosd(t_h)) # [degrees] zenith

	# azimuth [0, 180], absolute value measured from due south, so east = west = 90,
	# south = 0, north = 180
	cos_phi = (cosd(zenith)*sind(latitude) - sind(delta)) ./ (sind(zenith)*cosd(latitude)) # cosine(azimuth)
	phi_south = acosd(clamp(cos_phi,-1,1))
	#phi_south = acosd(cos_phi);

	# azimuth [0, 360], measured clockwise from due north,
	# so east = 90, south = 180, and west = 270 degrees
	azimuth = 180 + sign(t_h)*phi_south # Shift domain to 0-360 deg

	#angles = [theta, phi]; # [degrees] zenith, azimuth
	return zenith, azimuth, solarTime, delta
end

function solarposition_wiki_faster(datetime::DateTime, latitude, longitude)
	year = Dates.year(datetime)
	d_ms = datetime - DateTime(year,1,1)
	d = d_ms.value/1000/3600/24
	w = 360/365.24
	a = w * (d+10)
	b = a + (360/pi*0.0167)*sind(w*(d-2))
	c = (a - atand(tand(b)/cosd(23.44))) / 180
	ET = 720*(c - round(c))
	solarTime = 24*(d - floor(d)) + longitude*24/360 + ET/60
	t_h = 15*(solarTime - 12) # [degrees] hour angle
	delta = -asind(sind(23.44)*cosd(b))
	slat, clat = sind(latitude), cosd(latitude)
	sdel, cdel = sind(delta), cosd(delta)
	czen = slat*sdel + clat*cdel*cosd(t_h)
	zenith = acosd(czen) # [degrees] zenith

	# azimuth [0, 180], absolute value measured from due south, so east = west = 90,
	# south = 0, north = 180
	cos_phi = (czen*slat - sdel) / (sind(zenith)*clat) # cosine(azimuth)
	phi_south = acosd(clamp(cos_phi,-1,1))
	#phi_south = acosd(cos_phi);

	# azimuth [0, 360], measured clockwise from due north,
	# so east = 90, south = 180, and west = 270 degrees
	azimuth = 180 + sign(t_h)*phi_south # Shift domain to 0-360 deg

	#angles = [theta, phi]; # [degrees] zenith, azimuth
	return zenith, azimuth, solarTime, delta
end

# Splits the DateTime into (year, month, day) (as integers) and the hour (as a floating point)
function splitdatetime(datetime::DateTime)
	# date = floor(datetime, Dates.Day)
	y, m, d = Dates.yearmonthday(datetime)
	return y, m, d, (datetime - DateTime(y,m,d)).value/1000/3600
end

# time correction term, difference between Terrestrial Time and Universal Time, see Grena (2012) section 2.1.
# Conversion to TT is necessary since UT follows the rotation of the earth, which is not uniform.
# t [days] days from January 1, 2060 (in UT)
# returns Δτ [seconds]
differenceTTUT(t) = 96.4 + 0.00158*t

# Calculate time variable needed for Grena (2012) algorithms, see section 3.1.
# The time t is the input DateTime expressed as days from January 1, 2060 (in UT).
# returns t [days]
calctime(datetime::DateTime) = (datetime-DateTime(2060,1,1)).value/1000/3600/24

# Same result as above, but this is the function in Grena's paper. The Julia version is faster.
function calctime_grena(datetime::DateTime)
	y, m, d, h = splitdatetime(datetime)
	if m <= 2
		m += 12
		y -= 1
	end
	return trunc(365.25*(y-2000)) + trunc(30.6001*(m+1)) - trunc(0.01*y) + d + 0.0416667*h - 21958
end

# Full version of algorithm 1 of Grena (2012), see section 3.2.
# Calculate the position of the sun as (zenith,azimuth) in radians for a location
# on earth given as (latitude,longitude) in degrees. The DateTime should be given
# as Universal Time, e.g. UTC.
function solarposition_grena1(datetime::DateTime, latitude, longitude)
	ω = 0.017202786 # [1/day]
	t = calctime(datetime)
	t_e = t + 1.1574e-5 * differenceTTUT(t)
	s1, c1 = sin(ω*t_e), cos(ω*t_e)
	s2, c2 = 2*s1*c1, (c1+s1)*(c1-s1)
	α = mod(-1.38880 + 1.72027920e-2*t_e + 3.199e-2*s1 - 2.65e-3*c1 + 4.050e-2*s2 + 1.525e-2*c2, 2π)
	δ = 6.57e-3 + 7.347e-2*s1 - 3.9919e-1*c1 + 7.3e-4*s2 - 6.60e-3*c2
	H = mod(1.75283 + 6.3003881*t + deg2rad(longitude) - α + π, 2π) - π
	sϕ, cϕ = sind(latitude), cosd(latitude)
	sδ, cδ = sin(δ), cos(δ)
	sH, cH = sin(H), cos(H)
	se0 = sϕ*sδ + cϕ*cδ*cH
	ep = asin(se0) - 4.26e-5*sqrt(1-se0^2)  #cos(e0)
	azimuth = atan(sH, cH*sϕ - sδ*cϕ/cδ)	#azimuth = 0 towards south and positive direction towards west
	zenith = π/2 - ep
	return zenith, azimuth				# radians
	# return zenith, azimuth, δ, H		# radians
	# return rad2deg(zenith), rad2deg(mod(azimuth+π,2π)), rad2deg(H)/15+12, rad2deg(δ) # degrees, same azimuth origin as solarposition0
end

# First half of algorithm 1 of Grena (2012), see section 3.2.
# Calculate the position of the sun as (declination, hour angle) in radians for a location
# on earth given by its longitude in degrees. The DateTime should be given
# as Universal Time, e.g. UTC.
function solarposition(datetime::DateTime, longitude)
	ω = 0.017202786 # [1/day]
	t = calctime(datetime)
	t_e = t + 1.1574e-5 * differenceTTUT(t)
	s1, c1 = sin(ω*t_e), cos(ω*t_e)
	s2, c2 = 2*s1*c1, (c1+s1)*(c1-s1)
	α = mod(-1.38880 + 1.72027920e-2*t_e + 3.199e-2*s1 - 2.65e-3*c1 + 4.050e-2*s2 + 1.525e-2*c2, 2π)
	δ = 6.57e-3 + 7.347e-2*s1 - 3.9919e-1*c1 + 7.3e-4*s2 - 6.60e-3*c2
	H = mod(1.75283 + 6.3003881*t + deg2rad(longitude) - α + π, 2π) - π
	return δ, H
end

# Full version of algorithm 1 of Grena (2012), see section 3.2.
# Calculate the position of the sun as (zenith, azimuth) in degrees for a location
# on earth given as (latitude,longitude) in degrees. The DateTime should be given
# as Universal Time, e.g. UTC.
function solarposition(datetime::DateTime, latitude, longitude)
	δ, H = solarposition(datetime, longitude)
	zenith, azimuth = zenith_azimuth(latitude, sines_and_cosines(δ, H)...)
	return rad2deg(zenith), rad2deg(azimuth)
end

sines_and_cosines(δ, H) = sin(δ), cos(δ), sin(H), cos(H)

function zenith_azimuth(latitude, sδ, cδ, sH, cH)
	sϕ, cϕ = sind(latitude), cosd(latitude)
	se0 = sϕ*sδ + cϕ*cδ*cH
	azimuth = atan(sH, cH*sϕ - sδ*cϕ/cδ)	# azimuth = 0 towards south and positive direction towards west
	zenith = π/2 - asin(se0) + 4.26e-5*sqrt(1-se0^2)  # sqrt(1-se0^2) = cos(e0), but slightly faster
	return zenith, azimuth		# radians
end

function zenith(latitude, sδ, cδ, sH, cH)
	sϕ, cϕ = sind(latitude), cosd(latitude)
	se0 = sϕ*sδ + cϕ*cδ*cH
	return π/2 - asin(se0) + 4.26e-5*sqrt(1-se0^2)  # sqrt(1-se0^2) = cos(e0), but slightly faster
end

const solarconstant = 1361  # W/m2

# Annual variation of solar insolation outside the atmosphere (W/m2):
# https://www.itacanet.org/the-sun-as-a-source-of-energy/part-2-solar-energy-reaching-the-earths-surface/#2.1.-The-Solar-Constant
solarinsolation(dt::DateTime) = solarconstant * (1 + 0.034*cos(2*π*(dt - DateTime(year(dt))).value/1000/3600/24/365.25))

# dt = DateTime(2019,4,15,13,46)
# δ, H = solarposition(dt, 12)	# return rad2deg(zenith), rad2deg(mod(azimuth+π,2π)), rad2deg(H)/15+12, rad2deg(δ)

# pos = sines_and_cosines(δ, H)
# zen = zenith(58, pos...)


using StatsBase, Distributions, Plots; pyplot()

function prn(lambda)
    k, p = 0, 1
    while p > MathConstants.e^(-lambda)
        k += 1
        p *= rand()
    end
    return k-1
end

xGrid, lambda, N = 0:16, 5.5, 10^6

pDist = Poisson(lambda)
bPmf = pdf.(pDist,xGrid)
data = counts([prn(lambda) for _ in 1:N],xGrid)/N

plot( xGrid, data,
	line=:stem, marker=:circle,
	c=:blue, ms=10, msw=0, lw=4, label="MC estimate")
plot!( xGrid, bPmf, line=:stem,
	marker=:xcross, c=:red, ms=6, msw=0, lw=2, label="PMF",
	ylims=(0,0.2), xlabel="x", ylabel="Probability of x events")


function identical_matrix(n::Integer)
    ret = zeros(n, n)
    for itr1 in 1 : n
        ret[itr1, itr1] = 1
    end
    return ret
end

function exchange_row!(mat::Matrix, a::Integer, b::Integer)
    # if $a$ or $b$ are out of bound, do nothing
    if a > size(mat, 1) || b > size(mat, 1)
        return mat
    end
    # exchange row $a$ and row $b$
    for itr1 in 1 : size(mat, 2)
        mat[a, itr1], mat[b, itr1] = mat[b, itr1], mat[a, itr1]
    end
    return mat
end

function similarity_transformation!(
        mat::Matrix, a::Integer, b::Integer)
    # if $a$ or $b$ are out of bound, do nothing
    if a > size(mat, 1) || b > size(mat, 1)
        return mat
    end
    # taking advantage of the matrix, the transformation
    # is just exchanging part of the row $a$ and row $b$
    for itr1 in 1 : (min(a, b) - 1)
        mat[a, itr1], mat[b, itr1] = mat[b, itr1], mat[a, itr1]
    end
    return mat
end

function plu(mat::Matrix)
    # make a copy of the original matrix, which will finally become $U$
    U = float(copy(mat))
    # # U = copy(mat)
    # it's assumed that the matrix has $m$ rows and $n$ columns
    m, n = size(U)
    # initialize two identical matrices as $P$ and $L$
    P = identical_matrix(m)
    L = identical_matrix(m)
    # $i$ and $j$ are used to represent iterators of row and column,
    # respectively
    i = j = 1
    while i <= m && j <= n
        # find the entry with largest absolute value
        imax, vmax = 0.0, 0.0
        for itr1 in i : m
            if vmax < abs(U[itr1, j])
                imax, vmax = itr1, abs(U[itr1, j])
            end
        end
        (vmax > 0) ? (k = imax) : (j += 1; continue)
        # left multiply $U$ and $P$ with a permutation matrix
        if k != i
            exchange_row!(U, i, k)
            exchange_row!(P, i, k)
        end
        # left and right multiply $L$ with the same matrix
        similarity_transformation!(L, i, k)
        # `forward elimination'
        for itr1 in (i + 1) : m
            scaler = U[itr1, j] / U[i, j]
            for itr2 in j : n
                U[itr1, itr2] -= U[i, itr2] * scaler
            end
            # right multiply $L$ with an elimination matrix
            L[itr1, i] = scaler
        end
        i += 1
        j += 1
    end
    return P, L, U
end


using ThreeBodyProblem
using DifferentialEquations
using Plots
using LinearAlgebra
# using ForwardDiff

# The first step is to define the system. ThreeBodyProblem.jl has built-in functions for common systems like Sun/Jupiter.
sys = sun_jupiter()

## There's even a recipe to plot the system for us!
plot(sys)

## We can see the Lagrange Points, but where are the Sun and Jupiter? Turns out that the distance between them is so much larger than the radius of either one that it makes it hard to see them on a truly scaled picture. Let's scale the Sun and Jupiter to make them visible.
plot(sys, scaled = true)

## That's better! Now let's compute the L1 and L2 Lagrange points of our system

L1, L2 = computeLpts(sys)
# Declare state vectors for L1 and L2 (with zero velocity)
rv1 = [L1; zeros(3)] # state at Lagrange point L1 (zero velocity added on)
rv2 = [L2; zeros(3)] # state at Lagrange point L2

## The Lagrange points are equilibrium points in our dynamics. This means that an object placed there perfectly will stay there forever. But, Lagrange points L1, L2, and L3 are unstable, meaning if you perturb the object slightly, it may fall away. This is how we will generate the invariant manifolds.
## We first need to linearize the dynamics about these Lagrange points to determine which directions are stable and unstable

Φ₀ = I(6) # The initial State Transition Matrix (STM)
w₀ = vcat(rv1,reshape(Φ₀,36,1)) # Augmented state with state and STM
# Set up the ODE problem
tspan = (0.,1.)
prob = ODEProblem(CR3BPstm!,w₀,tspan,sys)
sol = solve(prob, reltol=1e-6)
Φₜ = Matrix(reshape(sol.u[end][7:42],6,6))

# last time step should match very closely with rv1
rvₜ = sol.u[end][1:6]
err = norm(rv1 - rvₜ)

# Compute the eigenvectors of the Monodromy Matrix
Λ,V = eigen(Φₜ,sortby=isreal) # Λ is vector of eigenvalues, V is matrix of eigenvectors
Yw = real(V[:,findall(isreal, Λ)]) # Eigenvectors corresponding to real eigenvalues
Λ = Λ[findall(isreal, Λ)] # Purely real eigenvalues (have 0.0 imaginary component)
Yws = Yw[:,findmin(real(Λ))[2]] # Eigenvector associated with stable eigenvalue λ < 1
Ywu = Yw[:,findmax(real(Λ))[2]] # Eigenvector associated with unstable eigenvalue λ > 1

## We perturb our states in the stable and unstable directions
α = 1e-6

# L1 points
rv1up = rv1 + α*Ywu # unstable manifold + side
rv1un = rv1 - α*Ywu # unstable manifold - side
rv1sp = rv1 + α*Yws # stable manifold + side
rv1sn = rv1 - α*Yws # stable manifold - side

# Integrate trajectories for 2 Jupiter periods
tf = 2*2π
tspan_forward = (0.,tf) # make sure to add decimal points so tspan contains floating point values, not integers
tspan_back = (0.,-tf)
myreltol = 1e-12

# Set up the ODE Problems
prob1up = ODEProblem(CR3BPdynamics!,rv1up,tspan_forward,sys) # Unstable positive
prob1un = ODEProblem(CR3BPdynamics!,rv1un,tspan_forward,sys) # Unstable negative
prob1sp = ODEProblem(CR3BPdynamics!,rv1sp,tspan_back,sys) # Stable positive
prob1sn = ODEProblem(CR3BPdynamics!,rv1sn,tspan_back,sys) # Stable negative

# Solutions to the ODEs
sol1up = solve(prob1up, reltol=myreltol)
sol1un = solve(prob1un, reltol=myreltol)
sol1sp = solve(prob1sp, reltol=myreltol)
sol1sn = solve(prob1sn, reltol=myreltol)

plot(sys,scaled=true)
plot!(sol1up,vars=(1,2),label="Wu+",linecolor=:red)
plot!(sol1un,vars=(1,2),label="Wu-",linecolor=:magenta)
plot!(sol1sp,vars=(1,2),label="Ws+",linecolor=:blue)
plot!(sol1sn,vars=(1,2),label="Ws-",linecolor=:cyan)
plot!(aspect_ratio=1,ylims=[-1,1],xlims=[-0.6,1.3],legend=:topleft,flip=false)

## We can see two strands moving in towards the sun, but the rest is a jumbled mess. Let's zoom in to see what's going on. Also, I need to make sure things are scaled correctly, or Jupiter will look way too big.
plot(sys)
plot!(sol1up,vars=(1,2),label="Wu+",linecolor=:red)
plot!(sol1un,vars=(1,2),label="Wu-",linecolor=:magenta)
plot!(sol1sp,vars=(1,2),label="Ws+",linecolor=:blue)
plot!(sol1sn,vars=(1,2),label="Ws-",linecolor=:cyan)
plot!(aspect_ratio=1,ylims=[-0.04,0.04],xlims=[0.93,1.06],legend=:outerright,flip=false)

# L2 points
rv2up = rv2 + α*Ywu # unstable manifold + side
rv2un = rv2 - α*Ywu # unstable manifold - side
rv2sp = rv2 + α*Yws # stable manifold + side
rv2sn = rv2 - α*Yws # stable manifold - side

## Now we look at the L2 point manifolds
prob2up = ODEProblem(CR3BPdynamics!,rv2up,tspan_forward,sys) # Unstable positive
prob2un = ODEProblem(CR3BPdynamics!,rv2un,tspan_forward,sys) # Unstable negative
prob2sp = ODEProblem(CR3BPdynamics!,rv2sp,tspan_back,sys) # Stable positive
prob2sn = ODEProblem(CR3BPdynamics!,rv2sn,tspan_back,sys) # Stable negative

sol2up = solve(prob2up, reltol=1e-6)
sol2un = solve(prob2un, reltol=1e-6)
sol2sp = solve(prob2sp, reltol=1e-6)
sol2sn = solve(prob2sn, reltol=1e-6)

plot(sys, scaled=true)
plot!(sol2up,vars=(1,2),label="Wu+",linecolor=:red)
plot!(sol2un,vars=(1,2),label="Wu-",linecolor=:magenta)
plot!(sol2sp,vars=(1,2),label="Ws+",linecolor=:blue)
plot!(sol2sn,vars=(1,2),label="Ws-",linecolor=:cyan)
plot!(aspect_ratio=1,ylims=[-1.0,1.0],xlims=[-0.3,1.5],legend=:topright,flip=false)
## This time there are two strands moving out away from the Sun, along with another jumble around Jupiter. Let's zoom in again to take a look.
plot(sys)
plot!(sol2up,vars=(1,2),label="Wu+",linecolor=:red)
plot!(sol2un,vars=(1,2),label="Wu-",linecolor=:magenta)
plot!(sol2sp,vars=(1,2),label="Ws+",linecolor=:blue)
plot!(sol2sn,vars=(1,2),label="Ws-",linecolor=:cyan)
plot!(aspect_ratio=1,ylims=[-0.04,0.04],xlims=[0.93,1.07],legend=:outerright,flip=false)


################################################################################
#
# The LLL-Reduction Algorithm
#
# Implementration of algorithm described in :
#  Lenstra, A.K., Lenstra, H.W. & Lovász, L. Math. Ann. (1982)
#            261: 515. https://doi.org/10.1007/BF01457454
#
################################################################################

"""
  lll(b₀::AbstractMatrix{Float64})::AbstractMatrix{Float64}

Performs Lenstra-Lenstra-Lovasv reduction on the n vectors in b₀.

Parameters
----------
- b₀ : vector of n vectors of n Float64

Returns
-------
- The LLL reduction of the vectors in `b₀`
"""
function lll(b₀::Vector{Vector{Float64}})::Vector{Vector{Float64}}
    b = deepcopy(b₀) # preserve argument by working on a copy
    n = length(b)

    # initialize mu and B with zeros
    μ = zeros(n,n)
    B = zeros(n)

    # function (*) of page 521
    function ☼(l)
        abs(μ[k,l]) < 0.5 && return
        r = round(μ[k,l])
        b[k] .-= r * b[l]
        μ[k,1:l-1] .-= r * μ[l,1:l-1]
        μ[k,l] -= r
    end

    b₂ = similar(b)
    for t in 1:n
        b₂[t] = copy(b[t])
        for j in 1:t-1
            μ[t,j] = dot(b[t], b₂[j]) / B[j]
            b₂[t] .-= μ[t,j] * b₂[j]
        end
        B[t] = dot(b₂[t], b₂[t])
    end

    k = 2
    while k <= n
        ☼(k-1)
        if B[k] >= (0.75 - μ[k,k-1]^2) * B[k-1]
            foreach(☼, k-2:-1:1)
            k += 1
        else
            μₒ = μ[k,k-1]
            Bₒ = B[k] + μₒ^2 * B[k-1]
            μ[k,k-1] = μₒ * B[k-1] / Bₒ
            B[k] = B[k] * B[k-1] / Bₒ
            B[k-1] = Bₒ

            bₒ = b[k] ; b[k] = b[k-1] ; b[k-1] = bₒ
            for j in 1:k-2
                tmp = μ[k,j] ; μ[k,j] = μ[k-1,j] ; μ[k-1,j] = tmp
            end
            for t in k+1:n
                tmp = μ[k,k-1] - μₒ * μ[t,k]
                μ[t,k-1] = μ[t,k] + μ[k,k-1] * tmp
                μ[t,k]   = tmp
            end

            k = max(2, k-1)
        end
    end
    b
end