### A Pluto.jl notebook ###
# v0.14.7

using Markdown
using InteractiveUtils

# ╔═╡ 6a41e9d0-c58c-11eb-078e-c36bf22a38d2
md"""
# Support Vector Machines

[Reference 1](https://scikit-learn.org/stable/modules/svm.html)\
[Reference 2](https://en.wikipedia.org/wiki/Support_vector_machine)\
[Reference 3](http://www.robots.ox.ac.uk/~az/lectures/ml/lect3.pdf)
"""

# ╔═╡ 9cfc2328-0e4f-4a86-944c-31c0c133fe9a
md"""
The SVM classification problem (directly taken from sklearn website):\

Given training vectors $x_i\in \mathbb{R}^P$, $i=1,...,n$, and a vector $y\in\{1,-1\}^n$\
The goal is to find $w\in \mathbb{R}^P$ and $b\in \mathbb{R}$ such that prediction given by $\text{sign}(w^T\phi(x)+b)$ is accurate for most samples\

Solve the following __primal problem__:\
$$\begin{align}
\min_{w,b,\zeta}\frac{1}{2}w^Tw+C\sum^n_{i=1}\zeta_i\\
\text{subject to }y_i(w^T\phi(x_i)+b)\geq1-\zeta_i,\\
\zeta_i\geq0,i=1,...,n
\end{align}$$\

Intuition is we are maximizing the margin (by minimizing $\|w\|^2=w^Tw$), while penalize when a sample is misclassified\

The __dual problem__ to primal is:\
$$\begin{align}
\min_\alpha\frac{1}{2}\alpha^TQ\alpha-e^T\alpha\\
\text{subject to }y^T\alpha=0,\\
0\leq\alpha_i\leq C,i=1,...,n
\end{align}$$\

 $e$ is vector of all ones, $Q$ is nxn positive semidefinite matrix, $Q_{ij}\equiv y_iy_jK(x_i,x_j)$, and $K(x_i,x_j)=\phi(x_i)^T\phi(x_j)$ is the kernel
$\alpha_i$ are dual coefficients, upper-bounded by $C$\
Highlighting the fact that training vectors are implicitly mapped to a higher (or infinite) dimensional space by function $\phi$\

After optimization problem solved, the __output__ for a given sample $x$ is:
$$\sum_{i\in SV}y_i\alpha_iK(x_i,x)+b$$\

The problem solved by `liblinear` for `LinearSVC` is a equivant form of primal problem:\
$$\min_{w,b}\frac{1}{2}w^Tw+C\sum_{i=1}\max(0,y_i(w^T\phi(x_i)+b))$$\
which does not involve inner products between samples, and therefore cannot apply kernel tricks\

$$C=\frac{1}{\text{alpha}}$$
"""

# ╔═╡ 30975066-eb17-4682-bda1-fd77f9d44734
function ingredients(path::String)
	# this is from the Julia source code (evalfile in base/loading.jl)
	# but with the modification that it returns the module instead of the last object
	name = Symbol(basename(path))
	m = Module(name)
	Core.eval(m,
        Expr(:toplevel,
             :(eval(x) = $(Expr(:core, :eval))($name, x)),
             :(include(x) = $(Expr(:top, :include))($name, x)),
             :(include(mapexpr::Function, x) = $(Expr(:top, :include))(mapexpr, $name, x)),
             :(include($path))))
	m
end

# ╔═╡ 9088c5d0-f07c-404f-8e08-2c8585f86a2f
import Random

# ╔═╡ fc9d6f80-70c1-4518-b9db-29c6e3e2d280
import Statistics

# ╔═╡ 6a0019bf-ab70-4095-ac52-7606e401d9a5
import LinearAlgebra

# ╔═╡ b8447232-6b1d-4b3a-8d11-a9a2da1cf31b
JuTools = ingredients("../tools.jl").JuTools

# ╔═╡ e4380307-4abe-4a20-9127-e608d1a5fa6e
md"""
### Linear SVM Implementation (Hinge Loss)

[Reference](https://stackoverflow.com/questions/48804198/soft-margin-in-linear-support-vector-machine-using-python)\

Cost function is:\
$$J=\frac{1}{2}w^Tw+\frac{C}{N}\sum^N_{i=1}\max\Big(0,1-y_i(w^T\phi(x_i)+b)\Big)$$

Gradient function for $w$ is:\
$$\frac{\partial J}{\partial w}=w + \frac{C}{N}\sum^N_{i=1}\begin{cases}
0 & y_i(w^T\phi(x_i)+b) \geq 1\\
-y_i\phi(x_i) & \text{otherwise}
\end{cases}$$

Gradient function for $b$ is:\
$$\frac{\partial J}{\partial b}=\frac{C}{N}\sum^N_{i=1}\begin{cases}
0 & y_i(w^T\phi(x_i)+b) \geq 1\\
-y_i & \text{otherwise}
\end{cases}$$
"""

# ╔═╡ 6841f5fe-45ad-4f71-84d5-368d196a40ab
# define a struct to store weights
# this should be returned by a training function
# alpha should be treated as constant
mutable struct WeightsLinearSVM
    C::AbstractFloat
    w::Array{T} where T<:AbstractFloat
    b::AbstractFloat
end

# ╔═╡ a093820c-20a0-4af1-98b3-ea92ad37bc77
# define cost function for linear SVM
# assum Y_data is {-1, 1}
function cost(X_data::Array{T} where T<:Number, Y_data::Array{T} where T<:Number, weights::WeightsLinearSVM)::AbstractFloat
    @assert ndims(Y_data) == ndims(weights.w) == 1
    @assert size(X_data) == (size(Y_data)[1], size(weights.w)[1])
    loss_w = 0.5 * (weights.w' * weights.w)
    loss_inner = 1.0 .- Y_data .* vec(X_data * weights.w .+ weights.b)
    loss_inner .= map(m->max(0.0,m), loss_inner)
    loss = loss_w + weights.C * sum(loss_inner) / size(X_data)[1]
    return loss
end

# ╔═╡ 8c3fa961-d88d-4e28-98e9-18b8804c568e
X_data, Y_data = JuTools.data_generate_linear_2d()

# ╔═╡ 6341ae48-9cda-4019-bfe6-443caf8788aa
Y_data .= Y_data .* 2.0 .- 1.0 # convert from {0,1} to {-1,1}

# ╔═╡ 500703d2-e315-44c7-8d6c-c36f40b815df
X_train, X_test, Y_train, Y_test = JuTools.split_data(X_data, Y_data)

# ╔═╡ 7531e45c-95fd-4664-8c37-3364aae7020b
@show size(X_train)

# ╔═╡ 747fabf5-fd7b-40cc-b5da-b15f6925cff0
@show size(X_test)

# ╔═╡ 9504574b-2281-4f8a-83d5-3eea38acec53
@show size(Y_train)

# ╔═╡ c496efb2-5a50-4131-bd65-e9025a3273a4
@show size(Y_test)

# ╔═╡ 46e9b187-d445-4efe-b4f2-ad824e5ec832
weight_test = WeightsLinearSVM(1.0, Random.randn(size(X_data)[2]), Random.randn())

# ╔═╡ 8adbc614-9440-42d3-9672-f3b816e114a4
# define the learning function (gradient descent)
function learn!(X_data::Array{T} where T<:Number, Y_data::Array{T} where T<:Number, weights::WeightsLinearSVM, alpha::AbstractFloat)
    @assert ndims(Y_data) == ndims(weights.w) == 1
    @assert size(X_data) == (size(Y_data)[1], size(weights.w)[1])
    # compute deciding feature
    decide = (Y_data .* (X_data * weights.w .+ weights.b)) .< 1 # (? < 1) will be 1, otherwise 0
    # update w
    gradient_w = weights.w .+ (weights.C / size(X_data)[1]) .* vec(-(Y_data .* decide)' * X_data)
    gradient_w .= gradient_w .* alpha
    weights.w .= weights.w .- gradient_w
    # update b
    gradient_b = (weights.C / size(X_data)[1]) * sum(-(Y_data .* decide))
    gradient_b *= alpha
    weights.b = weights.b - gradient_b
    return nothing
end

# ╔═╡ 633040d0-651b-427f-a5a8-83124bbea43b
# define prediction function
function predict_proba(X_predict::Array{T} where T<:Number, weights::WeightsLinearSVM)::Array
    @assert ndims(X_predict) == 2
    @assert size(X_predict)[2] == size(weights.w)[1]
    prediction = vec(X_predict * weights.w .+ weights.b)
    return prediction
end

# ╔═╡ ea07e163-d2a0-4447-9239-64d1aee8b3a0
# output prediction is in {-1, 1}
function predict(X_predict::Array{T} where T<:Number, weights::WeightsLinearSVM)::Array
    @assert ndims(X_predict) == 2
    @assert size(X_predict)[2] == size(weights.w)[1]
    prediction = vec(X_predict * weights.w .+ weights.b)
    prediction .= map(m -> m >= 0 ? 1.0 : -1.0, prediction)
    return prediction
end

# ╔═╡ 77bd6c88-335f-46a9-a46b-3d3596d54bf9
md"""
### SVM (with various kernels)

[Reference](https://en.wikipedia.org/wiki/Quadratic_programming)\

For solving SVM minimization problem:\
$$\begin{align}
\min_\alpha\frac{1}{2}\alpha^TQ\alpha-e^T\alpha\\
\text{subject to }y^T\alpha=0,\\
0\leq\alpha_i\leq C,i=1,...,n
\end{align}$$\

where $e$ is vector of all ones, $Q$ is nxn positive semidefinite matrix, $Q_{ij}\equiv y_iy_jK(x_i,x_j)$, and $K(x_i,x_j)=\phi(x_i)^T\phi(x_j)$ is the kernel
$\alpha_i$ are dual coefficients, upper-bounded by $C$\

We'll be using a Quatratic Programming technique: [Sequential Minimal Optimization](https://en.wikipedia.org/wiki/Sequential_minimal_optimization)(SMO)\

My implementation is inspired from [this blog](https://jonchar.net/notebooks/SVM/), whose code is originally from [this paper](https://www.researchgate.net/publication/234786663_Fast_Training_of_Support_Vector_Machines_Using_Sequential_Minimal_Optimization)\

The problem can also be written into (objective function):\
$$\begin{align}
\max_\alpha e^T\alpha-\frac{1}{2}\alpha^TQ\alpha\\
\text{subject to }y^T\alpha=0,\\
0\leq\alpha_i\leq C,i=1,...,n
\end{align}$$\

The kernel functions that I'm going to implement are (using sklearn names):
* Linear Kernel: $\langle x_i,x_j \rangle$
* Polynomial Kernel: $(\gamma\langle x_i,x_j \rangle + r)^d$ ($d$ is degree, and $r$ is coeficient, $\gamma$ is a parameter)
* Rbf Kernel: $\exp(-\gamma\|x_i-x_j\|^2)$, where $\gamma$ is a parameter
* Sigmoid Kernel: $\tanh(\gamma\langle x_i,x_j \rangle + r)$, where $r$ and $\gamma$ are parameters
"""

# ╔═╡ 111304b2-bfbb-4875-838a-93c153d452be
# define a struct to store information
mutable struct WeightsSVM
    C::AbstractFloat                # constraint
    b::AbstractFloat                # threshold
    gamma::AbstractFloat            # parameter used for polynomial, rbf, and sigmoid kernels
    r::AbstractFloat                # parameter used for polynomial, and sigmoid kernels
    d::AbstractFloat                # parameter used for polynomial kernel
    tol_alpha::AbstractFloat        # tolerance for alpha
    tol_error::AbstractFloat        # tolerance for error
    alpha::Array{T} where T<:Number # alpha array
    error::Array{T} where T<:Number # array for error cache
    kernel::String                  # kernel function name
end

# ╔═╡ 3f6728b3-fb75-4783-ab6e-96ae6301f7cb
# linear kernel
function kernel_linear(X1::Array{T} where T<:Number, X2::Array{T} where T<:Number)::Array
    @assert ndims(X1) == ndims(X2) == 2
    @assert size(X1)[2] == size(X2)[2]
    result = X1 * X2'
    return result
end

# ╔═╡ e3b0cd51-75bc-41aa-96eb-251b376a183d
# polynomial kernel
function kernel_polynomial(X1::Array{T} where T<:Number, X2::Array{T} where T<:Number;
        d::AbstractFloat=1.0, r::AbstractFloat=0.0, gamma::AbstractFloat=1.0)::Array
    @assert ndims(X1) == ndims(X2) == 2
    @assert size(X1)[2] == size(X2)[2]
    result = (gamma .* (X1 * X2') .+ r) .^ d
    return result
end

# ╔═╡ 4a0ee69a-6151-423b-9227-65fa4e9f2b5e
# rbf kernel
function kernel_rbf(X1::Array{T} where T<:Number, X2::Array{T} where T<:Number; gamma::AbstractFloat=1.0)::Array
    @assert ndims(X1) == ndims(X2) == 2
    @assert size(X1)[2] == size(X2)[2]
    result = (sum(X1 .^ 2, dims=2) * ones(size(X2)[1])') .+ (ones(size(X1)[1]) * sum(X2 .^ 2, dims=2)') .- 2.0 .* (X1 * X2')
    result .= broadcast(m->max(0.0, m), result) # ignore very small negative outputs, due to precision
    result .= sqrt.(result)
    result .= (-gamma) .* result
    result .= exp.(result)
    return result
end

# ╔═╡ 85939fbf-de18-46d8-9708-dec8a85f49dd
# sigmoid kernel
function kernel_sigmoid(X1::Array{T} where T<:Number, X2::Array{T} where T<:Number; gamma::AbstractFloat=1.0, r::AbstractFloat=0.0)::Array
    @assert ndims(X1) == ndims(X2) == 2
    @assert size(X1)[2] == size(X2)[2]
    result = gamma .* (X1 * X2') .+ r
    result .= tanh.(gamma)
    return result
end

# ╔═╡ 99d0a23e-e06b-4c13-89a7-62aeebf7c559
# because we call it cost function, we will use the original formula
function cost(X_data::Array{T} where T<:Number, weights::WeightsSVM)::AbstractFloat
    @assert ndims(X_data) == ndims(weights.alpha) + 1 == 2
    @assert size(X_data)[1] == size(weights.alpha)[1]
    result = nothing
    if weights.kernel == "linear"
        result = kernel_linear(X_data, X_data)
    elseif weights.kernel == "polynomial"
        result = kernel_polynomial(X_data, X_data, d=weights.d, r=weights.r, gamma=weights.gamma)
    elseif weights.kernel == "rbf"
        result = kernel_rbf(X_data, X_data, gamma=weights.gamma)
    elseif weights.kernel == "sigmoid"
        result = kernel_sigmoid(X_data, X_data, gamma=weights.gamma, r=weights.r)
    else
        throw(ArgumentError("Error: kernel function $weights.kernel is not recognized"))
    end
    result = 0.5 * (weights.alpha' * result * weights.alpha) - sum(weights.alpha)
    return result
end

# ╔═╡ 50e8d679-d4fc-4ed8-9868-b59cc4151d2c
cost(X_data, Y_data, weight_test)

# ╔═╡ ece26477-6e53-4536-a6a8-61947cb2815e
# define learning each step function
# update weights in place, and return num steps
function learn_step!(X_data::Array{T} where T<:Number, Y_data::Array{T} where T<:Number,
        weights::WeightsSVM, id1::Integer, id2::Integer)::Integer
    @assert ndims(X_data) == ndims(Y_data) + 1 == 2
    @assert size(X_data)[1] == size(Y_data)[1]
    @assert size(X_data)[1] == size(weights.alpha)[1]
    @assert size(weights.error) == size(weights.alpha)
    @assert id1 >= 1
    @assert id2 >= 1
    # if choosing same alpha, skip
    if id1 == id2
        return 0
    end
    # prepare data
    alpha1 = weights.alpha[id1]
    alpha2 = weights.alpha[id2]
    Y1 = Y_data[id1]
    Y2 = Y_data[id2]
    error1 = weights.error[id1]
    error2 = weights.error[id2]
    # compute L & H
    L = nothing
    H = nothing
    if Y1 != Y2
        L = max(0.0, alpha2 - alpha1)
        H = min(weights.C, weights.C + alpha2 - alpha1)
    else
        L = max(0.0, alpha1 + alpha2 - weights.C)
        H = min(weights.C, alpha1 + alpha2)
    end
    if L == H
        return 0
    end
    # compute kernel results and 2nd derivative eta
    k11 = nothing
    k12 = nothing
    k22 = nothing
    n_features = size(X_data)[2]
    X_id1 = reshape(X_data[id1, :], (1, n_features))
    X_id2 = reshape(X_data[id2, :], (1, n_features))
    if weights.kernel == "linear"
        k11 = kernel_linear(X_id1, X_id1)[1]
        k12 = kernel_linear(X_id1, X_id2)[1]
        k22 = kernel_linear(X_id2, X_id2)[1]
    elseif weights.kernel == "polynomial"
        k11 = kernel_polynomial(X_id1, X_id1, d=weights.d, r=weights.r, gamma=weights.gamma)[1]
        k12 = kernel_polynomial(X_id1, X_id2, d=weights.d, r=weights.r, gamma=weights.gamma)[1]
        k22 = kernel_polynomial(X_id2, X_id2, d=weights.d, r=weights.r, gamma=weights.gamma)[1]
    elseif weights.kernel == "rbf"
        k11 = kernel_rbf(X_id1, X_id1, gamma=weights.gamma)[1]
        k12 = kernel_rbf(X_id1, X_id2, gamma=weights.gamma)[1]
        k22 = kernel_rbf(X_id2, X_id2, gamma=weights.gamma)[1]
    elseif weights.kernel == "sigmoid"
        k11 = kernel_sigmoid(X_id1, X_id1, gamma=weights.gamma, r=weights.r)[1]
        k12 = kernel_sigmoid(X_id1, X_id2, gamma=weights.gamma, r=weights.r)[1]
        k22 = kernel_sigmoid(X_id2, X_id2, gamma=weights.gamma, r=weights.r)[1]
    else
        throw(ArgumentError("Error: kernel function $weights.kernel is not recognized"))
    end
    eta = 2 * k12 - k11 - k22
    # compute new alpha2 (a2)
    a2 = nothing
    if eta < 0.0
        a2 = alpha2 - Y2 * (error1 - error2) / eta
        a2 = min(a2, H)
        a2 = max(a2, L)
    else
        weights.alpha[id2] = L
        Lobj = -cost(X_data, weights)
        weights.alpha[id2] = H
        Hobj = -cost(X_data, weights)
        weights.alpha[id2] = alpha2
        if Lobj > (Hobj + weights.tol_alpha)
            a2 = L
        elseif Lobj < (Hobj - weights.tol_alpha)
            a2 = H
        else
            a2 = alpha2
        end
    end
    # push to 0 or C
    if a2 < 1e-8
        a2 = 0.0
    elseif a2 > (weights.C - 1e-8)
        a2 = weights.C
    end
    # skip if cannot be optimized
    if abs(a2 - alpha2) < weights.tol_alpha * (a2 + alpha2 + weights.tol_alpha)
        return 0
    end
    # compute new alpha1 (a1)
    a1 = alpha1 + (Y1 * Y2) * (alpha2 - a2)
    if a1 < 0.0
        a2 += (Y1 * Y2) * a1
        a1 = 0.0
    elseif a1 > weights.C
        a2 += (Y1 * Y2) * (a1 - weights.C)
        a1 = weights.C
    end
    # update threshold
    b1 = error1 + Y1 * (a1 - alpha1) * k11 + Y2 * (a2 - alpha2) * k12 + weights.b
    b2 = error2 + Y1 * (a1 - alpha1) * k12 + Y2 * (a2 - alpha2) * k22 + weights.b
    b_new = nothing
    if 0 < a1 < weights.C
        b_new = b1
    elseif 0 < a2 < weights.C
        b_new = b2
    else
        b_new = (b1 + b2) * 0.5
    end
    # update error cache
    non_optimized_ids = [i for i in 1:size(X_data)[1] if (i != id1 && i != id2 && (0 < weights.alpha[i] < weights.C))]
    kerr1 = nothing
    kerr2 = nothing
    if weights.kernel == "linear"
        kerr1 = vec(kernel_linear(X_id1, X_data[non_optimized_ids, :]))
        kerr2 = vec(kernel_linear(X_id2, X_data[non_optimized_ids, :]))
    elseif weights.kernel == "polynomial"
        kerr1 = vec(kernel_polynomial(X_id1, X_data[non_optimized_ids, :], d=weights.d, r=weights.r, gamma=weights.gamma))
        kerr2 = vec(kernel_polynomial(X_id2, X_data[non_optimized_ids, :], d=weights.d, r=weights.r, gamma=weights.gamma))
    elseif weights.kernel == "rbf"
        kerr1 = vec(kernel_rbf(X_id1, X_data[non_optimized_ids, :], gamma=weights.gamma))
        kerr2 = vec(kernel_rbf(X_id2, X_data[non_optimized_ids, :], gamma=weights.gamma))
    elseif weights.kernel == "sigmoid"
        kerr1 = vec(kernel_sigmoid(X_id1, X_data[non_optimized_ids, :], gamma=weights.gamma, r=weights.r))
        kerr2 = vec(kernel_sigmoid(X_id2, X_data[non_optimized_ids, :], gamma=weights.gamma, r=weights.r))
    end
    weights.error[non_optimized_ids] .= weights.error[non_optimized_ids] .+
        ((Y1*(a1-alpha1)) .* kerr1) .+ ((Y2*(a2-alpha2)) .* kerr2) .+ (b_new - weights.b)
    weights.error[id1] = 0.0
    weights.error[id2] = 0.0
    # update alpha and b
    weights.b = b_new
    weights.alpha[id1] = a1
    weights.alpha[id2] = a2
    return 1
end

# ╔═╡ 59c861e3-8c4d-450d-b727-c59cdf17a84f
# now define the learning function
function learn!(X_data::Array{T} where T<:Number, Y_data::Array{T} where T<:Number,
        weights::WeightsSVM, id::Integer; verbose::Bool=false)::Integer
    @assert ndims(X_data) == ndims(Y_data) + 1 == 2
    @assert size(X_data)[1] == size(Y_data)[1]
    @assert size(X_data)[1] == size(weights.alpha)[1]
    @assert size(weights.error) == size(weights.alpha)
    @assert id >= 1
    Y = Y_data[id]
    alpha = weights.alpha[id]
    error = weights.error[id]
    r = error * Y
    if ((r < -weights.tol_error) && (alpha < weights.C)) || ((r > weights.tol_error) && (alpha > 0))
        alpha_target = [i for (i, m) in enumerate(weights.alpha) if (0.0 < m < weights.C)]
        # try argmax E1 - E2
        new_id = -1
        tmax = 0
        if verbose
            println("Trying argmax(abs(E1 - E2))")
        end
        for i in alpha_target
            tmp = abs(error - weights.error[i])
            if(tmp > tmax)
                tmax = tmp
                new_id = i
            end
        end
        if new_id >= 1
            step = learn_step!(X_data, Y_data, weights, id, new_id)
            if step > 0
                return step
            end
        end
        # loop non-bound alphas, randomly
        if verbose
            println("Trying random non-bound alphas")
        end
        for new_id in alpha_target[Random.randperm(length(alpha_target))]
            step = learn_step!(X_data, Y_data, weights, id, new_id)
            if step > 0
                return step
            end
        end
        # else loop all alphas, randomly
        if verbose
            println("Trying random remaining alphas")
        end
        for new_id in Random.randperm(length(weights.alpha))
            if new_id in alpha_target
                continue # skip the alpha ids that already looked at
            end
            step = learn_step!(X_data, Y_data, weights, id, new_id)
            if step > 0
                return step
            end
        end
    end
    return 0
end

# ╔═╡ 882b95f5-6261-42d4-91b2-16a7ce9c1aa4
# implement predict functions
function predict_proba(X_predict::Array{T} where T<:Number, X_data::Array{T} where T<:Number,
        Y_data::Array{T} where T<:Number, weights::WeightsSVM)::Array
    @assert ndims(X_predict) == ndims(X_data) == ndims(Y_data) + 1 == 2
    @assert size(X_predict)[2] == size(X_data)[2]
    result = nothing
    if weights.kernel == "linear"
        result = kernel_linear(X_data, X_predict)
    elseif weights.kernel == "polynomial"
        result = kernel_polynomial(X_data, X_predict, d=weights.d, r=weights.r, gamma=weights.gamma)
    elseif weights.kernel == "rbf"
        result = kernel_rbf(X_data, X_predict, gamma=weights.gamma)
    elseif weights.kernel == "sigmoid"
        result = kernel_sigmoid(X_data, X_predict, gamma=weights.gamma, r=weights.r)
    else
        throw(ArgumentError("Error: kernel function $weights.kernel is not recognized"))
    end
    prediction = vec((weights.alpha .* Y_data)' * result)
    return prediction
end

# ╔═╡ dde248cd-6fac-4c59-877d-58e5b8f06685
function predict(X_predict::Array{T} where T<:Number, X_data::Array{T} where T<:Number,
        Y_data::Array{T} where T<:Number, weights::WeightsSVM)::Array
    @assert ndims(X_predict) == ndims(X_data) == ndims(Y_data) + 1 == 2
    @assert size(X_predict)[2] == size(X_data)[2]
    result = nothing
    if weights.kernel == "linear"
        result = kernel_linear(X_data, X_predict)
    elseif weights.kernel == "polynomial"
        result = kernel_polynomial(X_data, X_predict, d=weights.d, r=weights.r, gamma=weights.gamma)
    elseif weights.kernel == "rbf"
        result = kernel_rbf(X_data, X_predict, gamma=weights.gamma)
    elseif weights.kernel == "sigmoid"
        result = kernel_sigmoid(X_data, X_predict, gamma=weights.gamma, r=weights.r)
    else
        throw(ArgumentError("Error: kernel function $weights.kernel is not recognized"))
    end
    prediction = vec((weights.alpha .* Y_data)' * result)
    prediction .= map(m -> m >= 0 ? 1.0 : -1.0, prediction)
    return prediction
end

# ╔═╡ aafac217-757c-442a-88b1-17e8895f703a
# training function for linear SVM
# assume Y_data is in {-1, 1}
# this function is similar to the training function for Logistic Regression (Both are gradient descent)
function train_linear(X_data::Array{T} where T<:Number, Y_data::Array{T} where T<:Number, C::AbstractFloat;
        learning_rate::AbstractFloat=0.1, max_iter::Integer=1000, n_iter_no_change::Integer=5, tol::AbstractFloat=0.001,
        verbose::Bool=false, shuffle::Bool=true, early_stop::Bool=true)::WeightsLinearSVM
    @assert ndims(X_data) == ndims(Y_data) + 1 == 2
    @assert size(X_data)[1] == size(Y_data)[1]
    @assert max_iter >= 0
    @assert n_iter_no_change >= 0
    @assert tol >= 0
    X_data = Float64.(X_data)
    Y_data = Float64.(Y_data)
    if shuffle
        JuTools.shuffle_data!(X_data, Y_data)
    end
    # is it better to use zero weights than normal weights ?
    weights = WeightsLinearSVM(C, Random.randn(size(X_data)[2]), Random.randn())
    best_cost = nothing
    n_cost_no_change = n_iter_no_change
    for i in 1:max_iter
        if n_cost_no_change <= 0 && early_stop
            break
        end
        learn!(X_data, Y_data, weights, learning_rate)
        new_cost = cost(X_data, Y_data, weights)
        if verbose
            acc = JuTools.compute_accuracy(predict(X_data, weights), Y_data)
            println("Iter: $i")
            println("Cost = $new_cost")
            println("Accuracy = $acc")
            println()
        end
        if early_stop
            if best_cost === nothing || isnan(best_cost)
                best_cost = new_cost
            else
                if new_cost > best_cost - tol
                    n_cost_no_change -= 1
                else
                    best_cost = min(new_cost, best_cost)
                    n_cost_no_change = n_iter_no_change
                end
            end
        end
    end
    return weights
end

# ╔═╡ ec256a57-ea44-4f73-af5b-f5e09ba2fec1
weights = train_linear(X_train, Y_train, 1.0, learning_rate=0.005, max_iter=20, tol=0.001, verbose=true)

# ╔═╡ fb710b22-79b0-477c-b2a6-f85fe8b4271a
JuTools.compute_accuracy(predict(X_test, weights), Y_test)

# ╔═╡ 47846484-1824-41b3-b26d-e4e3953fd079
# finally implement the training function
function train(X_data::Array{T} where T<:Number, Y_data::Array{T} where T<:Number, C::AbstractFloat;
        tol_alpha::AbstractFloat=0.001, tol_error::AbstractFloat=0.001, kernel::String="rbf", gamma::String="scale",
        degree::AbstractFloat=1.0, coef::AbstractFloat=0.0, verbose::Bool=false)::WeightsSVM
    @assert ndims(X_data) == ndims(Y_data) + 1 == 2
    @assert size(X_data)[1] == size(Y_data)[1]
    X_data = Float64.(X_data)
    Y_data = Float64.(Y_data)
    # gamma is computed the same way sklearn does
    gamma_num = nothing
    if gamma == "scale"
        gamma_num = 1.0 / (size(X_data)[2] * Statistics.var(X_data))
    elseif gamma == "auto"
        gamma_num = 1.0 / size(X_data)[2]
    else
        throw(ArgumentError("Error: gamma $gamma is not recognized, possible values are 'scale' and 'auto'"))
    end
    weights = WeightsSVM(C, 0.0, gamma_num, coef, degree, tol_alpha, tol_error, Float64.(zeros(size(X_data)[1])), -copy(Y_data), kernel)
    num_changed = 0
    examine_all = true
    total_steps = 0
    while (num_changed > 0) || examine_all
        num_changed = 0
        if examine_all
            if verbose
                println("Scanning all training data")
            end
            for i in 1:size(X_data)[1]
                step = learn!(X_data, Y_data, weights, i, verbose=verbose)
                num_changed += step
                if step > 0 && verbose
                    obj = -cost(X_data, weights)
                    println("1 step further, objective = $obj")
                end
            end
        else
            if verbose
                println("Scanning data whose alpha is not at limit")
            end
            alpha_target = [i for (i, m) in enumerate(weights.alpha) if (m != 0.0 && m != weights.C)]
            for i in alpha_target
                step = learn!(X_data, Y_data, weights, i, verbose=verbose)
                num_changed += step
                if step > 0 && verbose
                    obj = -cost(X_data, weights)
                    println("1 step further, objective = $obj")
                end
            end
        end
        if examine_all
            examine_all = false
        elseif num_changed <= 0
            examine_all = true
        end
        total_steps += num_changed
    end
    if verbose
        println("Training Complete\nTotal steps: $total_steps")
    end
    return weights
end

# ╔═╡ fdd56685-18d1-47b5-9b9c-27090422f250
weights_svm = train(X_train, Y_train, 10.0, kernel="rbf", verbose=false, gamma="auto", coef=0.0, degree=2.0)

# ╔═╡ 1b11d952-a35f-4778-8a7f-779eb72f0436
JuTools.compute_accuracy(predict(X_test, X_train, Y_train, weights_svm), Y_test)

# ╔═╡ Cell order:
# ╟─6a41e9d0-c58c-11eb-078e-c36bf22a38d2
# ╟─9cfc2328-0e4f-4a86-944c-31c0c133fe9a
# ╟─30975066-eb17-4682-bda1-fd77f9d44734
# ╠═9088c5d0-f07c-404f-8e08-2c8585f86a2f
# ╠═fc9d6f80-70c1-4518-b9db-29c6e3e2d280
# ╠═6a0019bf-ab70-4095-ac52-7606e401d9a5
# ╠═b8447232-6b1d-4b3a-8d11-a9a2da1cf31b
# ╟─e4380307-4abe-4a20-9127-e608d1a5fa6e
# ╠═6841f5fe-45ad-4f71-84d5-368d196a40ab
# ╠═a093820c-20a0-4af1-98b3-ea92ad37bc77
# ╠═8c3fa961-d88d-4e28-98e9-18b8804c568e
# ╠═6341ae48-9cda-4019-bfe6-443caf8788aa
# ╠═500703d2-e315-44c7-8d6c-c36f40b815df
# ╠═7531e45c-95fd-4664-8c37-3364aae7020b
# ╠═747fabf5-fd7b-40cc-b5da-b15f6925cff0
# ╠═9504574b-2281-4f8a-83d5-3eea38acec53
# ╠═c496efb2-5a50-4131-bd65-e9025a3273a4
# ╠═46e9b187-d445-4efe-b4f2-ad824e5ec832
# ╠═50e8d679-d4fc-4ed8-9868-b59cc4151d2c
# ╠═8adbc614-9440-42d3-9672-f3b816e114a4
# ╠═633040d0-651b-427f-a5a8-83124bbea43b
# ╠═ea07e163-d2a0-4447-9239-64d1aee8b3a0
# ╠═aafac217-757c-442a-88b1-17e8895f703a
# ╠═ec256a57-ea44-4f73-af5b-f5e09ba2fec1
# ╠═fb710b22-79b0-477c-b2a6-f85fe8b4271a
# ╟─77bd6c88-335f-46a9-a46b-3d3596d54bf9
# ╠═111304b2-bfbb-4875-838a-93c153d452be
# ╠═3f6728b3-fb75-4783-ab6e-96ae6301f7cb
# ╠═e3b0cd51-75bc-41aa-96eb-251b376a183d
# ╠═4a0ee69a-6151-423b-9227-65fa4e9f2b5e
# ╠═85939fbf-de18-46d8-9708-dec8a85f49dd
# ╠═99d0a23e-e06b-4c13-89a7-62aeebf7c559
# ╠═ece26477-6e53-4536-a6a8-61947cb2815e
# ╠═59c861e3-8c4d-450d-b727-c59cdf17a84f
# ╠═882b95f5-6261-42d4-91b2-16a7ce9c1aa4
# ╠═dde248cd-6fac-4c59-877d-58e5b8f06685
# ╠═47846484-1824-41b3-b26d-e4e3953fd079
# ╠═fdd56685-18d1-47b5-9b9c-27090422f250
# ╠═1b11d952-a35f-4778-8a7f-779eb72f0436


########################
###### Data Types ######
########################

# (1 + 2)::Int
# 3

function f(y::Int)
    x::Int8 = 100
    x + y
end

x = f(20) # 120

typeof(x) # Int8

############################
###### Abstract Types ######
############################

# abstract type «name» end
# abstract type «name» <: «supertype» end

abstract type Number end
abstract type Real     <: Number end
abstract type AbstractFloat <: Real end
abstract type Integer  <: Real end
abstract type Signed   <: Integer end
abstract type Unsigned <: Integer end

# The Number type is a direct child type of Any, and Real is its child.

#############################
###### Primitive Types ######
#############################

# primitive type «name» «bits» end
# primitive type «name» <: «supertype» «bits» end

# primitive type Float16 <: AbstractFloat 16 end
# primitive type Float32 <: AbstractFloat 32 end
# primitive type Float64 <: AbstractFloat 64 end

# primitive type Bool <: Integer 8 end
# primitive type Char <: AbstractChar 32 end

# primitive type Int8    <: Signed   8 end
# primitive type UInt8   <: Unsigned 8 end
# primitive type Int16   <: Signed   16 end
# primitive type UInt16  <: Unsigned 16 end
# primitive type Int32   <: Signed   32 end
# primitive type UInt32  <: Unsigned 32 end
# primitive type Int64   <: Signed   64 end
# primitive type UInt64  <: Unsigned 64 end
# primitive type Int128  <: Signed   128 end
# primitive type UInt128 <: Unsigned 128 end

####################
###### Struct ######
####################

struct Foo
    bar  # Any
    baz::Int
    qux::Float64
end

foo = Foo("Hello", 23, 1.5)

typeof(foo) # Foo

fieldnames(Foo) # (:bar, :baz, :qux)

foo.bar # "Hello"

# Composite objects declared with struct are immutable; they cannot be modified after construction.

mutable struct Bar
    baz
    qux::Float64
end

bar = Bar("Hello", 1.5)

bar.qux = 2.0

bar.baz = 1//2

#########################
###### Type Unions ######
#########################

IntOrString = Union{Int,AbstractString}

1 :: IntOrString

"Hello" :: IntOrString

# 1.0 :: IntOrString - ERROR

##############################
###### Parametric Types ######
##############################

# An important and powerful feature of Julia's type system is that it is parametric: types can
# take parameters, so that type declarations actually introduce a whole family of new types – one
# for each possible combination of parameter values.

struct Point{T}
    x::T
    y::T
end

# Point{Float64}
# Point{AbstractString}

# Point{Float64} <: Point
# true

# Point{AbstractString} <: Point
# true

# Point{Float64} <: Point{Real}
# false

# !!!! Even though Float64 <: Real we DO NOT have Point{Float64} <: Point{Real}. !!!!

# Define a method that accepts all arguments of type Point{T} where T is a subtype of Real
function norm(p::Point{<:Real})
    sqrt(p.x^2 + p.y^2)
end

p = Point{Float64}(1.0, 2.0)

typeof(p) # Point{Float64}(1.0, 2.0)

p1 = Point(1,2)

typeof(p1) # Point{Int64}(1, 2)

# p3 = Point(1, 2.5) !!!! Error !!!!

#######################################
###### Parametric Abstract Types ######
#######################################

abstract type Pointy{T} end

struct Point{T} <: Pointy{T}
    x::T
    y::T
end

# Point{Float64} <: Pointy{Float64}
# true

# Point{Float64} <: Pointy{Real}
# false

# Point{Float64} <: Pointy{<:Real}
# true

struct DiagPoint{T} <: Pointy{T}
    x::T
end

# Now both Point{Float64} and DiagPoint{Float64} are implementations of the Pointy{Float64} abstraction,
# and similarly for every other possible choice of type T. This allows programming to a common interface
# shared by all Pointy objects, implemented for both Point and DiagPoint

struct Point{T<:Real} <: Pointy{T}
    x::T
    y::T
end

# Real-world example of how all this parametric type machinery can be useful

struct Rational{T<:Integer} <: Real
    num::T
    den::T
end

# It only makes sense to take ratios of integer values, so the parameter type T is restricted to
# being a subtype of Integer, and a ratio of integers represents a value on the real number line,
# so any Rational is an instance of the Real abstraction.

#########################
###### Tuple Types ######
#########################

struct Tuple2{A,B}
    a::A
    b::B
end

typeof((1,"foo",2.5)) # Tuple{Int64, String, Float64}

# Vararg Tuple types

mytupletype = Tuple{AbstractString,Vararg{Int}}
# Tuple{AbstractString, Vararg{Int64, N} where N}

isa(("1",), mytupletype)
# true

# Named Tuple Types

typeof((a=1,b="hello"))
# NamedTuple{(:a, :b), Tuple{Int64, String}}

# The @NamedTuple macro provides a more convenient struct-like syntax for declaring NamedTuple types
# via key::Type declarations, where an omitted ::Type corresponds to ::Any.

@NamedTuple{a::Int, b::String}
# NamedTuple{(:a, :b), Tuple{Int64, String}}

@NamedTuple begin
    a::Int
    b::String
end

# NamedTuple{(:a, :b), Tuple{Int64, String}}

############################
###### UnionAll Types ######
############################

const T1 = Array{Array{T,1} where T, 1}
# Vector{Vector{T} where T} (alias for Array{Array{T, 1} where T, 1})

const T2 = Array{Array{T, 1}, 1} where T
# Array{Vector{T}, 1} where T

# Type T1 defines a 1-dimensional array of 1-dimensional arrays; each of the inner arrays consists
# of objects of the same type, but this type may vary from one inner array to the next. On the other hand,
# type T2 defines a 1-dimensional array of 1-dimensional arrays all of whose inner arrays must have the
# same type. Note that T2 is an abstract type, e.g., Array{Array{Int,1},1} <: T2, whereas T1 is a
# concrete type. As a consequence, T1 can be constructed with a zero-argument constructor a=T1() but
# T2 cannot.

Vector{T} = Array{T, 1}

# This is equivalent to const Vector = Array{T,1} where T.
# Writing Vector{Float64} is equivalent to writing Array{Float64,1}

####################################
###### Type{T} type selectors ######
####################################

# For each type T, Type{T} is an abstract parametric type whose only instance is the object T.

isa(Float64, Type{Float64})
# true

isa(Real, Type{Float64})
# false

isa(Real, Type{Real})
# true

# In other words, isa(A, Type{B}) is true if and only if A and B are the same object and that object is a type.

struct WrapType{T}
    value::T
end

WrapType(Float64) # default constructor, note DataType
# WrapType{DataType}(Float64)

WrapType(::Type{T}) where T = WrapType{Type{T}}(T)
# WrapType

WrapType(Float64) # sharpened constructor, note more precise Type{Float64}
# WrapType{Type{Float64}}(Float64)

####################################
###### Custom pretty-printing ######
####################################

struct Polar{T<:Real} <: Number
    r::T
    Θ::T
end

Polar(r::Real,Θ::Real) = Polar(promote(r,Θ)...)
# Polar


# calculate the volume of a sphere with radius r
# Arguments:  r = radius of the sphere
# Returns:    Volume

function sphereVolume(r)
  return 4/3 * pi * r^3
end


# Computes the logarithmic mean: (aR-aL)/(LOG(aR)-LOG(aL)) = (aR-aL)/LOG(aR/aL)
# Problem: if aL~= aR, then 0/0, but should tend to --> 0.5*(aR+aL)
#
# introduce xi=aR/aL and f=(aR-aL)/(aR+aL) = (xi-1)/(xi+1)
# => xi=(1+f)/(1-f)
# => Log(xi) = log(1+f)-log(1-f), and for small f (f^2<1.0E-02) :
#
#    Log(xi) ~=     (f - 1/2 f^2 + 1/3 f^3 - 1/4 f^4 + 1/5 f^5 - 1/6 f^6 + 1/7 f^7)
#                  +(f + 1/2 f^2 + 1/3 f^3 + 1/4 f^4 + 1/5 f^5 + 1/6 f^6 + 1/7 f^7)
#             = 2*f*(1           + 1/3 f^2           + 1/5 f^4           + 1/7 f^6)
#  (aR-aL)/Log(xi) = (aR+aL)*f/(2*f*(1 + 1/3 f^2 + 1/5 f^4 + 1/7 f^6)) = (aR+aL)/(2 + 2/3 f^2 + 2/5 f^4 + 2/7 f^6)
#  (aR-aL)/Log(xi) = 0.5*(aR+aL)*(105/ (105+35 f^2+ 21 f^4 + 15 f^6)
function ln_mean(value1, value2)
  epsilon_f2 = 1.0e-4
  ratio = value2 / value1
  # f2 = f^2
  f2 = (ratio * (ratio - 2) + 1) / (ratio * (ratio + 2) + 1)
  if f2 < epsilon_f2
    return (value1 + value2) * 52.5 / (105 + f2 * (35 + f2 * (21 + f2 * 15)))
  else
    return (value2 - value1) / log(ratio)
  end
end


using OrdinaryDiffEq
using ForwardDiff
using DataFrames
using LinearAlgebra
using StatsPlots

function sir_rates(u,p)
  (S,I,R) = u
  (β,c,γ) = p
  N = S+I+R
  infection = β*c*I/N*S
  recovery = γ*I
  [infection,recovery]
end;

sir_transitions = float.([[-1  1  0];
                          [ 0 -1  1]])

δt = 0.1
tmax = 40.0
tspan = (0.0,tmax)
t = 0.0:δt:tmax;

u0 = [990.0,10.0,0.0]; # S,I.R
c0 = zeros(3,3) # covariance matrix
uc0 = vcat(diagm(0=>u0), c0)

p = [0.05,10.0,0.25]; # β,c,γ

rates = sir_rates
transitions = sir_transitions
nrates, nstates = size(transitions);

function ode(du, u, p, t)
    du[1:3] = transitions' * rates(u, p)
end;

prob_ode = ODEProblem(ode, u0 ,tspan, p);

sol_ode = solve(prob_ode, Tsit5(), saveat=t);

df_ode = DataFrame(sol_ode(t)')
df_ode[!,:t] = t;

@df df_ode plot(:t,
    [:x1 :x2 :x3],
    label=["S" "I" "R"],
    xlabel="Time",
    ylabel="Number",
    c=[:blue :red :green])

function lna(du, u, p, t)
        mean_vec = diag(u)
        covar_mx = u[nstates+1:nstates*2,:]
        reaction_rates = rates(mean_vec, p)
        reaction_rates_jac = ForwardDiff.jacobian(
            y -> rates(y, p),
            mean_vec)
        A = transitions' * reaction_rates_jac
        du[1:nstates, :] .= diagm(0 => transitions'*reaction_rates)
        du[nstates + 1:end, :] .= A*covar_mx + covar_mx*A' + transitions' * diagm(0 => reaction_rates) * transitions
end;

prob_lna = ODEProblem(lna, uc0, tspan, p);

sol_lna = solve(prob_lna, Tsit5(),saveat=t);

mean_traj = Array{Float64,2}(undef, nstates, length(t))
covar_traj = Array{Array{Float64,2},1}(undef, length(t))
for j in 1:length(t)
    mean_traj[:,j] = diag(sol_lna[j][1:nstates,1:nstates])
    covar_traj[j] = sol_lna[j][nstates+1:end, 1:nstates]
end;

var_traj = zeros(nstates, length(t))
for (idx, elt) in enumerate(covar_traj)
    var_traj[:, idx] = diag(elt)
end
sd_traj = 1.96 .* sqrt.(var_traj);

plot(t, mean_traj',
     ribbon=sd_traj',
     label = ["S" "I" "R"],
     xlabel = "Time",
     ylabel = "Number")


#=

  Gaussian mixture model.

  cf ~/blog/gaussian_mixture_model.blog
     ~/psi/gaussian_mixture_model.psi
     ~/webppl/gaussian_mixture_model.wppl

=#

using Turing, StatsPlots, DataFrames
include("jl_utils.jl")

@model function gaussian_mixture_model()
    x = [0.2,1.0,0.5,0.6]
    n = length(x)

    p ~ Beta(0.5, 1)
    a ~ Uniform(-1, 1)
    b ~ Uniform(-1, 1)

    z = tzeros(n)
    x = tzeros(n)
    for i in 1:n
        z[i] ~ Bernoulli(p)
        if z[i] == 1
            x[i] ~ Normal(a, 1.0)
        else
            x[i] ~ Normal(b, 1.0)
        end
    end

    return a > b

end

model = gaussian_mixture_model()

num_chains = 4

# chains = sample(model, Prior(), 10_000)
# chains = sample(model, MH(), MCMCThreads(), 100_000, num_chains)
# chains = sample(model, MH(), MCMCThreads(), 10_000, num_chains)
# chains = sample(model, MH(), 10_000)

# chains = sample(model, PG(15), MCMCThreads(), 1_000, num_chains)

chains = sample(model, SMC(1000), MCMCThreads(), 10_000, num_chains)
# chains = sample(model, SMC(1000), 10_000)
# chains = sample(model, IS(), 10_000)
#
#chains = sample(model, Gibbs(HMC(0.1,5,:a,:b),PG(15,:p)), 10_000)
# chains = sample(model, Gibbs(NUTS(1000,0.65,:a,:b),PG(15,:p)), 10_000)
# chains = sample(model, Gibbs(HMC(0.1,5,:a,:b),SMC(1000,:p)), 10_000) # Nope

display(chains)
# display(plot(chains))

# show_var_dist_pct(chains,:a,20)
# show_var_dist_pct(chains,:b,20)
# show_var_dist_pct(chains,:p,20)

genq = generated_quantities(model, chains)
show_var_dist_pct(genq,20)


"""

trueG2_test() は，分割表のマス目のうちに 0 があるような場合でも，正確に lambda=0.0 にした場合の，正しい答えを返す。

"""

using Distributions
function trueG2_test(x; correct=false)
    # セルに 0 があっても正しい答えを出す
    # correct=true で，連続性の補正も行うことができる
    ln(n) = sum(n .== 0 ? 0 : n .* log(n) for n in vcat(n...))
    nrows, ncols = size(x)
    n = sum(x) # 全サンプルサイズ
    n1 = sum(x, dims=2) # 行和
    n2 = sum(x, dims=1) # 列和
    G2 = 2*(ln(x) - ln(n1) - ln(n2) + ln(n)) # G 統計量
    correct && (G2 /= 1 + (n * sum(1 ./ n1) - 1) * (n * sum(1 ./ n2) - 1) / (6n * nrows * ncols)) # 連続性の補正
    df = (nrows - 1) * (ncols - 1) # G の自由度
    p = ccdf(Chisq(df), G2)
    name = correct ? "corrected G-sq." : "G-sq."
    println("$name = $G2,  df = $df,  p value = $p")
end

"""
使用法

julia> z = [4 5 2 0
            0 7 6 1
            1 0 3 1]
3×4 Matrix{Int64}:
 4  5  2  0
 0  7  6  1
 1  0  3  1

julia> trueG2_test(z)
G-sq. = 15.364591286599591,  df = 6,  p value = 0.017602886503051453

また，ピアソンの χ2
検定では 2×2 分割表の場合にしか連続性の補正は行えないが，G2 検定は分割表の大きさに関わらず，連続性の補正を行うことができる。

上に示した関数 trueG2_test() で correct=true を指定すればよい。

julia> trueG2_test(z, correct=true) # 連続性の補正
corrected G-sq. = 13.776490649307341,  df = 6,  p value = 0.032235015595581236

"""


# Define a "model" type, that just needs a predict function
type GenericModel
	predict # Function that makes predictions
end

type LinearModel
	predict # Funcntion that makes predictions
	w # Weight vector
end

type CompressModel
	compress # Function that compresses
	expand # Function that de-compresses
	W # weight matrix
end

# Function to compute the mode of a vector
function mode(x)
	# Returns mode of x
	# if there are multiple modes, returns the smallest
	x = sort(x[:]);

	commonVal = [];
	commonFreq = 0;
	x_prev = NaN;
	freq = 0;
	for i in 1:length(x)
		if(x[i] == x_prev)
			freq += 1;
		else
			freq = 1;
		end
		if(freq > commonFreq)
			commonFreq = freq;
			commonVal = x[i];
		end
		x_prev = x[i];
	end
	return commonVal
end

# Return element-wise log, but set log(0)=0
function log0(x)
	y = copy(x)
	y[y.==0] = 1
	return log.(y)
end

# Return squared Euclidean distance all pairs of rows in X1 and X2
function distancesSquared(X1,X2)
	(n,d) = size(X1)
	(t,d2) = size(X2)
	assert(d==d2)
	return X1.^2*ones(d,t) + ones(n,d)*(X2').^2 - 2X1*X2'
end

# Subtract mean of each column and divide by standard deviation
# (or call it with mu and sigma to use these specific mean/std)
function standardizeCols(X;mu=[],sigma=[])
	(n,d) = size(X)

	if isempty(mu)
		mu_j = mean(X,1)
	else
		mu_j = mu
	end

	Xstd = zeros(n,d)
	for j in 1:d
		Xstd[:,j] = X[:,j] - mu_j[j]
	end

	if isempty(sigma)
		sigma_j = std(Xstd,1)
	else
		sigma_j = sigma
	end

	for j in 1:d
		Xstd[:,j] /= sigma_j[j]
	end

	if isempty(mu) & isempty(sigma)
		return (Xstd,mu_j,sigma_j)
	else
		return Xstd
	end
end

### A function to compute the gradient numerically
function numGrad(func,x)
	n = length(x);
	delta = 2*sqrt(1e-12)*(1+norm(x));
	g = zeros(n);
	e_i = zeros(n)
	for i = 1:n
		e_i[i] = 1;
		(fxp,) = func(x + delta*e_i)
		(fxm,) = func(x - delta*e_i)
		g[i] = (fxp - fxm)/2delta;
		e_i[i] = 0
	end
	return g
end

### Check if number is a real-finite number
function isfinitereal(x)
	return (imag(x) == 0) & (!isnan(x)) & (!isinf(x))
end

function dijkstra(D,s,t)
	# Finds shortest path from 's' to 't' with non-negative edge weights D
  # - set the edge weight D[i,j] to Inf if there is no edge between nodes 'i' and 'j' in the graph
	# (not really a fast implementation)

	n = size(D,1)

	distances = fill(Inf,n)
	distances[s] = 0

	visited = fill(false,n)

	while true

		# Choose next node to visit
		minDist = Inf
		i = []
		for j in 1:n
			if !visited[j] & (distances[j] < minDist)
				i = j
				minDist = distances[j]
			end
		end

		if isinf(minDist)
			# We have nowhere left to visit, there is no path
			break
		end

		# "Visit" the node: update neighbour distances
		for j in 1:n
			if distances[i] + D[i,j] < distances[j]
				distances[j] = distances[i] + D[i,j]
			end
		end
		visited[i] = true

		# If we just visited the solution
		if i == t
			break
		end
	end
	return distances[t]
end


function area(r)
    return (r^2 * asin(sqrt(r^2 - 1)/r)/2) - sqrt(r^2 - 1) - (r^2 * asin(1/r)/2) + 1
end

function prob(k)
    if k == 1
        return area(k + 1/2)/k^2
    else
        return (area(k + 1/2) - area(k - 1/2))/k^2
    end
end

function main(limit)
    expectation = sum(k * prob(k) for k in 1:limit)
    println(expectation)
end

main(100000)


using Contour, Gadfly, DataFrames

using OneDimensional: goldenSectionSearch, fibSearch

f(x, y) = x^4 + y^4 - 5 * (x * y - x^2 * y^2)
fdx(x, y) = 4x^3 - 5y + 10 * x * y^2
fdy(x, y) = 4y^3 - 5x + 10 * y * x^2
▽(x, y) = [fdx(x, y), fdy(x, y)]
L(x, y, λ) = [x, y] .- (λ .* ▽(x, y))

function simpleGradientDescent(f, ▽, x0, y0, ɛ) # down gradient with exponentital slow step
  callCnt, gradCallCnt, xs, ys = 1, 0, [x0], [y0]

  λ = 0.05
  fprev = f(x0, y0)
  while true
    xk, yk = xs[end], ys[end]
    xnew, ynew = L(xk, yk, λ)
    gradCallCnt += 1
    if xnew < 0 || 1 < xnew || ynew < 0 || 1 < ynew
      λ *= 0.75
      continue
    end
    fnew = f(xnew, ynew)
    callCnt += 1
    if fnew > fprev
      λ *= 0.75
      continue
    end
    fprev = fnew
    push!(xs, xnew)
    push!(ys, ynew)
    if norm([xs[end] - xs[end - 1], ys[end] - ys[end - 1]]) < ɛ
      break
    end
  end
  xs, ys, callCnt, gradCallCnt
end

function fastGradienDescent(f, ▽, x0, y0, ɛ) # fastest descent gradient
  callCnt, gradCallCnt, xl, yl = 1, 0, [x0], [y0]

  left, right = 0, 1
  while true
    x_left, y_left = xl[end], yl[end]
    g = ▽(x_left, y_left)
    gradCallCnt += 1
    ax = [(left - x_left) / -g[1], (right - x_left) / -g[1]]
    ay = [(left - y_left) / -g[2], (right - y_left) / -g[2]]
    sort!(ax)
    sort!(ay)
    a = [max(ax[1], ay[1]), min(ax[2], ay[2])]
    λ, cnt, _, _, _ = fibSearch(arg -> f(x_left - arg * g[1], y_left - arg * g[2]), a[1], a[2], ɛ)
    callCnt += cnt
    xnew, ynew = x_left - λ * g[1], y_left - λ * g[2]
    push!(xl, xnew)
    push!(yl, ynew)
    if norm([xl[end] - xl[end - 1], yl[end] - yl[end - 1]]) < ɛ
      break
    end
  end
  xl, yl, callCnt, gradCallCnt
end

ɛ = 0.0001
x0, y0 = 0.2, 0.8
sgx, sgy, cnt1, gcnt1 = simpleGradientDescent(f, ▽, x0, y0, ɛ)
fgx, fgy, cnt2, gcnt2 = fastGradienDescent(f, ▽, x0, y0, ɛ)

@printf("Обычный градиентный спуск: %d вычислений функции, %d вычислений градиента\n", cnt1, gcnt1)
@printf("Наискорейший спуск: %d вычислений функции, %d вычислений градиента\n", cnt2, gcnt2)

Gadfly.set_default_plot_size(18cm, 18cm)
plot(layer(x=[0:0.01:1], y=[0:0.01:1], z=f, Geom.contour(levels=15)),
     layer(x=sgx, y=sgy, Theme(default_color=color("orange")), Geom.point, Geom.line),
     layer(x=fgx, y=fgy, Theme(default_color=color("red")), Geom.point, Geom.line(preserve_order=true)),
     Guide.xticks(ticks=[0:0.2:1.0]), Guide.yticks(ticks=[0:0.2:1.0]),
     Guide.title("Линии уровня функции x^4 + y^4 - 5(xy - x^2 * y^2)"))


#==============================================================================

          Solving the HJB with a state constraint using
                the explicit method

               Translated from matlab code on Ben Moll's website:
                              http://www.princeton.edu/~moll/HACTproject.htm

==============================================================================#

using LinearAlgebra, SparseArrays, Plots

γ = 2 # parameter from CRRA utility
r = 0.03 # the interest rate
ρ = 0.05 # the discount rate

z1=.1
z2 =.2
z= [z1 z2]

λ1 = 0.02
λ2 = 0.03

λ = [λ1 λ2]

H = 500 # number of points in the grid space
amin=-0.02
amax=2

a = LinRange(amin,amax,H)
da = (amax-amin)/(H-1)

aa = [a a]
zz = ones(H,1)*z

maxit = 20000 # the maximum number of iterations we allow the finite differencing algorithm
crit = 10^(-6) # our critical value

dVf=zeros(H,2)
dVb = zeros(H,2)
c= zeros(H,2)

# Inital Guess
v0=zeros(H,2)
v0[:,1] = (z[1] .+ r.*a).^(1-γ)/(1-γ)/ρ
v0[:,2] = (z[2] .+ r.*a).^(1-γ)/(1-γ)/ρ
global v=v0

dist=[]

# The finite differeing loop

for n in 1:maxit
    global V=v
    # Forward differencing
    dVf[1:H-1,:]=(V[2:H,:]-V[1:H-1,:])/da
    dVf[H,:] = zeros(1,2) # will never be used
    # Backward differencing
    dVb[2:H,:]=(V[2:H,:]-V[1:H-1,:])/da
    dVb[1,:]= (z .+r.*amin).^(-γ)

    # Find consumption and savings with the forward difference
    cf = dVf.^(-1/γ)
    sf = zz +r.*aa -cf

    # Find consumption and savings with the backward difference
    cb= dVb.^(-1/γ)
    sb = zz +r.*aa -cb

    # Find consumption and savings at steady state
    c0 = zz + r.*aa
    dV0 = c0.^(-γ)

    # Now implement the upwind scheme in order to select the best differencing method
    If = sf .> 0
    Ib = sb .< 0
    I0 = (1 .- If - Ib)

    # Make sure the backward difference is using at the grid max
    Ib[H,:] = ones(1,2)
    If[H,:] = zeros(1,2)

    # State constraint at amin is automatically implemented

    global dV_Upwind = dVf.*If + dVb.*Ib + dV0.*I0

    global c = dV_Upwind.^(-1/γ)
    global V_switch = [V[:,2] V[:,1]]

    V_change = c.^(1-γ)/(1-γ) + dV_Upwind.*(zz + r.*aa -c) + ones(H,1)*λ.*(V_switch-V)-ρ.*V

    Δ = .9*da/(findmax(z2 .+r.*a)[1])
    global v = v  + Δ*V_change

    push!(dist, findmax(abs.(V_change))[1])

    if dist[n] < crit
        println("Value Function converged, Iteration=")
        println(n)
        break
    end
end

# Graphs
V_err = c.^(1-γ)/(1-γ) + dV_Upwind.*(zz +r.*aa -c) + ones(H,1)*λ.*(V_switch-V)-ρ.*V
adot = zz + r.*aa -c

plot(dist[:], xlabel="Iteration", ylabel="\$\\lvert V^{n+1}-V^{n}\\rvert\$")

plot(a,V_err, xlabel="k",ylabel="Error in HJB Equation", xlims=(amin,amax), legend=false)

plot(a,V, xlabel="a", ylabel="\$V_{i}(a)\$", xlims=(amin,amax),legend=false)

plot(a,c,xlabel="a", ylabel="\$c_{i}(a)\$", xlims=(amin,amax), legend=false)

plot(a, adot, xlabel="a", ylabel="\$s_{i}(a)\$", xlims=(amin,amax))

# Approximation at the borrowing constraint
u1=(z1+r*amin)^(-γ)
u2 = c[1,2]^(-γ)
u11 = -γ*(z1+r*amin)^(-γ-1)
ν = sqrt(-2*((ρ-r)*u1 + λ1*(u1-u2))/u11)
s_approx = -ν*(a.-amin).^(1/2)

plot(a, adot[:,1], xlabel="a", ylabel="\$s_{i}(a)\$", xlims=(amin,amax), label="\$s_{1}(a)\$")
plot!(a, adot[:,2], label="\$s_{2}(a)\$")
plot!(a, zeros(H,1), line=:dash, label="", color=:black)
plot!(a, s_approx, label="Approximatation")


#=

@authors : Spencer Lyon <spencer.lyon@nyu.edu>, John Stachurski

=#

using QuantEcon
using Distributions

"""
The Lucas asset pricing model --- parameters and grid data
"""
struct LucasTree{TF<:AbstractFloat}
    γ::TF                 # coefficient of risk aversion
    β::TF                 # Discount factor in (0, 1)
    α::TF                 # Correlation coefficient in the shock process
    σ::TF                 # Volatility of shock process
    ϕ::Distribution       # Distribution for shock process
    grid::Vector{TF}      # Grid of points on which to evaluate prices
    shocks::Vector{TF}    # Draws of the shock
    h::Vector{TF}         # The h function represented as a vector
end

"""
Constructor for the Lucas asset pricing model
"""
function LucasTree(;γ::AbstractFloat=2.0,
                β::AbstractFloat=0.95,
                α::AbstractFloat=0.9,
                σ::AbstractFloat=0.1,
                grid_size::Integer=100)

    ϕ = LogNormal(0.0, σ)
    shocks = rand(ϕ, 500)

    # == build a grid with mass around stationary distribution == #
    ssd = σ / sqrt(1 - α^2)
    grid_min, grid_max = exp(-4 * ssd), exp(4 * ssd)
    grid = collect(linspace(grid_min, grid_max, grid_size))

    # == set h(y) = β * int u'(G(y,z)) G(y,z) ϕ(dz) == #
    h = similar(grid)
    for (i, y) in enumerate(grid)
        h[i] = β * mean((y^α .* shocks).^(1 - γ))
    end

    return LucasTree(γ,
                     β,
                     α,
                     σ,
                     ϕ,
                     grid,
                     shocks,
                     h)
end

"""
The approximate Lucas operator, which computes and returns updated function
Tf on the grid points.
"""
function lucas_operator(lt::LucasTree, f::Vector)

    # == unpack names == #
    grid, α, β, h = lt.grid, lt.α, lt.β, lt.h
    z = lt.shocks

    Af = LinInterp(grid, f)

    Tf = [h[i] + β * mean(Af.(grid[i]^α.*z)) for i in 1:length(grid)]
    return Tf
end

"""
Compute the equilibrium price function associated with Lucas tree `lt`
"""
function solve_lucas_model(lt::LucasTree;
                           tol::AbstractFloat=1e-6,
                           max_iter::Integer=500)

    # == simplify notation == #
    grid, γ = lt.grid, lt.γ

    i = 0
    f = zeros(grid)  # Initial guess of f
    error = tol + 1

    while (error > tol) && (i < max_iter)
        f_new = lucas_operator(lt, f)
        error = maximum(abs, f_new - f)
        f = f_new
        i += 1
    end

    # p(y) = f(y) * y ^ γ
    price = f .* grid.^γ

    return price
end


import Base.intersect

complex_empty = Set{Complex{Float64}}()
identical_warn = "Intersection of identical clines returns the empty set"

"""
    intersect(C::Cline,D::Cline)
    C ∩ D

Find the points common to the clines `C` and `D` as a set of complex numbers.

Note that if `C==D` a warning is issued and the emptyset is returned.
"""
function intersect(C::Circle, D::Circle)::Set{Complex{Float64}}
    if C == D
        @warn identical_warn
        return complex_empty
    end

    z = center(C)
    r = radius(C)

    w = center(D)
    s = radius(D)

    d = abs(z - w)

    if d > r + s || d + s < r || d + r < s
        return Set{Complex{Float64}}()
    end

    θ = acos((r^2 + d^2 - s^2) / (2 * r * d))
    u = (w - z) / abs(w - z)
    uu = u * exp(θ * im)
    a = z + uu * r
    uu = u * exp(-θ * im)
    b = z + uu * r
    return Set([a, b])
end

function intersect(L::Line, M::Line)::Set{Complex{Float64}}
    if L == M
        @warn identical_warn
        return complex_empty
    end

    # check if parallel
    m1 = slope(L)
    m2 = slope(M)
    if (isinf(m1) && isinf(m2)) || (abs(m1 - m2) < get_tolerance())
        return complex_empty
    end

    a = L.a
    b = L.b

    ax, ay = reim(a)
    bx, by = reim(b)

    aa = M.a
    bb = M.b

    aax, aay = reim(aa)
    bbx, bby = reim(bb)

    MAT = zeros(2, 2)
    MAT[1, 1] = ax - bx
    MAT[1, 2] = bbx - aax
    MAT[2, 1] = ay - by
    MAT[2, 2] = bby - aay

    RHS = [bbx - bx, bby - by]
    st = MAT \ RHS

    s = st[1]
    z = s * L.a + (1 - s) * L.b
    return Set(z)
end

function intersect(C::Cline, D::Cline)::Set{Complex{Float64}}
    if C == D
        @warn identical_warn
        return complex_empty
    end

    F = LFT(C)  # map C to x-axis
    G = inv(F)
    DD = F(D)  # apply that to D and find where DD crosses the x-axis

    # if the images of D is a circle
    if isa(DD, Circle)
        x, y = reim(center(DD))
        r = radius(DD) + get_tolerance()

        if r < abs(y)
            return complex_empty = Set{Complex{Float64}}()
        end

        δ = sqrt(r^2 - y^2)
        a = G(x + δ)
        b = G(x - δ)
        if abs(a - b) < get_tolerance()  # if the points are too close, we say they're the same
            ab = (a + b) / 2
            return Set{Complex{Float64}}(ab)
        else
            return Set{Complex{Float64}}([a, b])
        end
    end

    # otherwise, DD, the image of D, is a line

    X = Line(0, 1)  # x-axis
    S = X ∩ DD
    return G.(S)

    # x, y = reim(DD.a)
    # xx, yy = reim(DD.b)

    # if abs(y - yy) < get_tolerance()  # DD is a horizontal line
    #     return complex_empty()
    # end

    # if abs(x - xx) < get_tolerance()  # DD is a vertical line
    #     z = G((x + xx) / 2)
    #     return Set{Complex{Float64}}(z)
    # end

    # slope = (y - yy) / (x - xx)
    # δ = y / slope
    # z = G(x - δ)
    # return Set{Complex{Float64}}(z)
end


test = (5764801, 17807724)
puzzle = (8987316, 14681524)
modulus = 20201227
seed = 7

mods = let mods = Vector(undef, modulus - 1)
    mods[1] = seed
    for i in 2:modulus - 1
        mods[i] = (7 * mods[i - 1]) % modulus
    end
    mods
end
@assert mods[end] == 1

function q1(card, door)::Int
    loops = (findfirst(==(card), mods), findfirst(==(door), mods))
    return mods[(*(loops...) - 1) % (modulus - 1) + 1]
end

@assert q1(test...) == 14897079

println("Q1: ", q1(puzzle...))


function penstock(Hn, Vdotn, H, theta; method="Bier",eta_nominal=0.9,pipe_roughness=0.05)
    g,rho,mu=parameters()
    P = eta_nominal*rho*g*Vdotn*Hn/1000;
    D = nominal_diameter(method,P,Hn,Vdotn)
    A = pi*D^2/4
    v = Vdotn/A
    NRe = renoldsnum(rho,v,D,mu)
    f = f_Darcy(NRe, pipe_roughness,D)
    L = H/sind(theta)
    hL = f*L*v^2/(2*g*D)
    hGrad = hL/sqrt(L^2-H^2)
    gradAngle = atan(hGrad)
    return sort((Dict("a. Nominal head"=>Hn,
                "b. Nominal dischare"=>Vdotn,
                "c. Nominal hydraulic efficiency"=>eta_nominal,
                "d. Nominal power output"=>P,
                "e. Design based on"=>method,
                "f. Economical diameter"=>D,
                "g. Cross-section"=>A,
                "h. Height"=>H,
                "i. Length"=>L,
                "j. Hydraulic crosssection"=>A,
                "k. Angle of inclination"=>theta,
                "l. Average velocity"=>v,
                "m. Nominal Reynolds number"=>NRe,
                "n. Darcy's friction factor"=>f,
                "o. Nominal allowed head loss"=>hL,
                "p. Hydraulic gradient"=>hGrad,
                "q. Pipe roughness"=>pipe_roughness,
                "r. Gradient angle"=>gradAngle)))
end

function headraceNsurgetank(Hn, Vdotn, H, theta, Hns; method="Bier",eta_nominal=0.9,pipe_roughness=0.05)
    g,rho,mu=parameters()
    headrace_dict = penstock(Hn, Vdotn, H, theta,method=method,eta_nominal=eta_nominal,pipe_roughness=pipe_roughness)
    # surge tank depends on headrace tunnel so
    n = 1.6 # safety factor for Thoma criteria
    # Acessing data of headrace tunnel from dictionary
    Dt,Lt,At,ft,vt = headrace_dict["f. Economical diameter"],headrace_dict["i. Length"],headrace_dict["g. Cross-section"],headrace_dict["n. Darcy's friction factor"],headrace_dict["l. Average velocity"]
    # calculating headrace tunnel friction factor
    beta = ft*Lt/(2*g*Dt)
    As = n*(Lt*At)/(2*g*beta*(Hn-beta*vt^2))
    Ds = sqrt(4*As/pi)
    Ymax = vt*sqrt((Lt*At)/(g*As))
    Hs = 1.2*(Hns+Ymax)
    vs = Vdotn/As
    surgetank_dict = sort((Dict("s. Diameter of surge tank"=>Ds,
                    "t. Cross-section of surge tank"=>As,
                    "u. Maximum height of water oscillation"=>Hns+Ymax,
                    "v. Height of surge tank"=>Hs,
                    "w. Allowed velocity of water inside surge tank"=>vs)))
    return sort(merge(headrace_dict,surgetank_dict))
end
function f2x1y(x,y)
    return 2*x+1*y
end


# # Trade-off curves
using Random
Random.seed!(1)
m = 25;
n = 10;
A = randn(m, n);
b = randn(m, 1);

#-

using Convex, SCS, LinearAlgebra

gammas = exp10.(range(-4, stop = 2, length = 100));

x_values = zeros(n, length(gammas));
x = Variable(n);
for i in 1:length(gammas)
    cost = sumsquares(A * x - b) + gammas[i] * norm(x, 1)
    problem = minimize(cost, [norm(x, Inf) <= 1])
    solve!(problem, SCS.Optimizer; silent_solver = true)
    x_values[:, i] = evaluate(x)
end

#-

# Plot the regularization path.

using Plots
plot(
    title = "Entries of x vs lambda",
    xaxis = :log,
    xlabel = "lambda",
    ylabel = "x",
)
for i in 1:n
    plot!(gammas, x_values[i, :], label = "x$i")
end
plot!()


"""
    PeriodicKernel(; r::AbstractVector=ones(Float64, 1))

Periodic kernel with parameter `r`.

# Definition

For inputs ``x, x' \\in \\mathbb{R}^d``, the periodic kernel with parameter ``r_i > 0`` is
defined[^DM] as
```math
k(x, x'; r) = \\exp\\bigg(- \\frac{1}{2} \\sum_{i=1}^d \\bigg(\\frac{\\sin\\big(\\pi(x_i - x'_i)\\big)}{r_i}\\bigg)^2\\bigg).
```

[^DM]: D. J. C. MacKay (1998). Introduction to Gaussian Processes.
"""
struct PeriodicKernel{T} <: SimpleKernel
    r::Vector{T}
    function PeriodicKernel(; r::AbstractVector{<:Real}=ones(Float64, 1))
        @check_args(PeriodicKernel, r, all(ri > zero(ri) for ri in r), "r > 0")
        return new{eltype(r)}(r)
    end
end

PeriodicKernel(dims::Int) = PeriodicKernel(Float64, dims)

"""
    PeriodicKernel([T=Float64, dims::Int=1])

Create a [`PeriodicKernel`](@ref) with parameter `r=ones(T, dims)`.
"""
PeriodicKernel(T::DataType, dims::Int=1) = PeriodicKernel(; r=ones(T, dims))

@functor PeriodicKernel

metric(κ::PeriodicKernel) = Sinus(κ.r)

kappa(::PeriodicKernel, d::Real) = exp(-0.5d)

function Base.show(io::IO, κ::PeriodicKernel)
    return print(io, "Periodic Kernel, length(r) = $(length(κ.r))")
end


using ComponentArrays
using DifferentialEquations
using Parameters: @unpack

tspan = (0.0, 20.0)

## Lorenz system
function lorenz!(D, u, p, t; f=0.0)
    @unpack σ, ρ, β = p
    @unpack x, y, z = u

    D.x = σ*(y - x)
    D.y = x*(ρ - z) - y - f
    D.z = x*y - β*z
    return nothing
end

lorenz_p = (σ=10.0, ρ=28.0, β=8/3)
lorenz_ic = ComponentArray(x=0.0, y=0.0, z=0.0)
lorenz_prob = ODEProblem(lorenz!, lorenz_ic, tspan, lorenz_p)

## Lotka-Volterra system
function lotka!(D, u, p, t; f=0.0)
    @unpack α, β, γ, δ = p
    @unpack x, y = u

    D.x =  α*x - β*x*y + f
    D.y = -γ*y + δ*x*y
    return nothing
end

lotka_p = (α=2/3, β=4/3, γ=1.0, δ=1.0)
lotka_ic = ComponentArray(x=1.0, y=1.0)
lotka_prob = ODEProblem(lotka!, lotka_ic, tspan, lotka_p)

## Composed Lorenz and Lotka-Volterra system
function composed!(D, u, p, t)
    c = p.c #coupling parameter
    @unpack lorenz, lotka = u

    lorenz!(D.lorenz, lorenz, p.lorenz, t, f=c*lotka.x)
    lotka!(D.lotka, lotka, p.lotka, t, f=c*lorenz.x)
    return nothing
end

comp_p = (lorenz=lorenz_p, lotka=lotka_p, c=0.01)
comp_ic = ComponentArray(lorenz=lorenz_ic, lotka=lotka_ic)
comp_prob = ODEProblem(composed!, comp_ic, tspan, comp_p)

## Solve problem
# We can solve the composed system...
comp_sol = solve(comp_prob)

# ...or we can unit test one of the component systems
lotka_sol = solve(lotka_prob)


function dbscan(arr::Array{Float64,2} , ϵ::Number, minPoints::Int)

	nPoints = size(arr,1)
	dists = fill(0.0, nPoints, nPoints)
	connectedPaths = fill(0, nPoints, nPoints)

	seeds = Int[]
	counts = Int[]
	clusters = zeros(Int, nPoints)
	visited = zeros(Bool, nPoints)

	# Make superpixel distance matrix and connectivity matrix
	for i in  1:nPoints
		for j in 1:nPoints
			dists[i,j] = euclidean(arr[i,:],arr[j,:])
		end
	end

	reachable = findall(x->x <= ϵ, dists)
	connectedPaths[reachable] .= 1

	function regionQuery(idx)
		neighbors = findall(x -> x == 1, connectedPaths[:,idx])
		return neighbors
	end

	k = 0

	for pt in 1:nPoints
		if clusters[pt] == 0 && !visited[pt]
			visited[pt] == true
			neighbors = regionQuery(pt)
			if length(neighbors) >= minPoints
				k += 1
				clusters[pt] = k
				cnt = 1
				while !isempty(neighbors)
					q = popfirst!(neighbors)
					if !visited[q]
						visited[q] = true
						qNeighbors = regionQuery(q)
						for x in qNeighbors
							if clusters[x] == 0
								push!(neighbors, x)
							end
						end
					end
					if clusters[q] == 0
						clusters[q] = k
						cnt += 1
					end
				end
				push!(seeds, pt)
				push!(counts, cnt)

			end
		end
	end

	return clusters
end


# Canonical Correlation Analysis

#### CCA Type

struct CCA{T<:Real}
    xmean::Vector{T}  # sample mean of X: of length dx (can be empty)
    ymean::Vector{T}  # sample mean of Y: of length dy (can be empty)
    xproj::Matrix{T}  # projection matrix for X, of size (dx, p)
    yproj::Matrix{T}  # projection matrix for Y, of size (dy, p)
    corrs::Vector{T}  # correlations, of length p

    function CCA(xm::Vector{T},
                 ym::Vector{T},
                 xp::Matrix{T},
                 yp::Matrix{T},
                 crs::Vector{T}) where T<:Real

        dx, px = size(xp)
        dy, py = size(yp)

        isempty(xm) || length(xm) == dx ||
            throw(DimensionMismatch("Incorrect length of xmean."))

        isempty(ym) || length(ym) == dy ||
            throw(DimensionMismatch("Incorrect length of ymean."))

        px == py ||
            throw(DimensionMismatch("xproj and yproj should have the same number of columns."))

        length(crs) == px ||
            throw(DimensionMismatch("Incorrect length of corrs."))

        new{T}(xm, ym, xp, yp, crs)
    end
end

## properties

xindim(M::CCA) = size(M.xproj, 1)
yindim(M::CCA) = size(M.yproj, 1)
outdim(M::CCA) = size(M.xproj, 2)

xmean(M::CCA) = fullmean(xindim(M), M.xmean)
ymean(M::CCA) = fullmean(yindim(M), M.ymean)

xprojection(M::CCA) = M.xproj
yprojection(M::CCA) = M.yproj

correlations(M::CCA) = M.corrs

## use

xtransform(M::CCA, X::AbstractVecOrMat{<:Real}) = transpose(M.xproj) * centralize(X, M.xmean)
ytransform(M::CCA, Y::AbstractVecOrMat{<:Real}) = transpose(M.yproj) * centralize(Y, M.ymean)

## show & dump

function show(io::IO, M::CCA)
    print(io, "CCA (xindim = $(xindim(M)), yindim = $(yindim(M)), outdim = $(outdim(M)))")
end

function dump(io::IO, M::CCA)
    show(io, M)
    println(io)
    println(io, "correlations: ")
    printvecln(io, correlations(M))
    println(io, "xmean:")
    printvecln(io, xmean(M))
    println(io, "ymean:")
    printvecln(io, ymean(M))
    println(io, "xprojection:")
    printarrln(io, xprojection(M))
    println(io, "yprojection:")
    printarrln(io, yprojection(M))
end

#### Perform CCA on data

## ccacov

function ccacov(Cxx::DenseMatrix{T},
                Cyy::DenseMatrix{T},
                Cxy::DenseMatrix{T},
                xmean::Vector{T},
                ymean::Vector{T},
                p::Int) where T<:Real

    # argument checking
    dx, dx2 = size(Cxx)
    dy, dy2 = size(Cyy)
    dx == dx2 || error("Cxx must be a square matrix.")
    dy == dy2 || error("Cyy must be a square matrix.")
    size(Cxy) == (dx, dy) ||
        throw(DimensionMismatch("size(Cxy) should be equal to (dx, dy)"))

    isempty(xmean) || length(xmean) == dx ||
        throw(DimensionMismatch("Incorrect length of xmean."))

    isempty(ymean) || length(ymean) == dy ||
        throw(DimensionMismatch("Incorrect length of ymean."))

    1 <= p <= min(dx, dy) ||
        throw(DimensionMismatch(""))

    _ccacov(Cxx, Cyy, Cxy, xmean, ymean, p)
end

function _ccacov(Cxx, Cyy, Cxy, xmean, ymean, p::Int)
    dx = size(Cxx, 1)
    dy = size(Cyy, 1)

    # solve Px and Py

    if dx <= dy
        # solve Px: (Cxy * inv(Cyy) * Cyx) Px = λ Cxx * Px
        # compute Py: inv(Cyy) * Cyx * Px

        G = cholesky(Cyy) \ Cxy'
        Ex = eigen(Symmetric(Cxy * G), Symmetric(Cxx))
        ord = sortperm(Ex.values; rev=true)
        vx, Px = extract_kv(Ex, ord, p)
        Py = qnormalize!(G * Px, Cyy)
    else
        # solve Py: (Cyx * inv(Cxx) * Cxy) Py = λ Cyy Py
        # compute Px: inv(Cx) * Cxy * Py

        H = cholesky(Cxx) \ Cxy
        Ey = eigen(Symmetric(Cxy'H), Symmetric(Cyy))
        ord = sortperm(Ey.values; rev=true)
        vy, Py = extract_kv(Ey, ord, p)
        Px = qnormalize!(H * Py, Cxx)
    end

    # compute correlations
    # Note: Px' * Cxx * Px == I
    #       Py' * Cyy * Py == I
    crs = coldot(Px, Cxy * Py)

    # construct CCA model
    CCA(xmean, ymean, Px, Py, crs)
end

## ccasvd

function ccasvd(Zx::DenseMatrix{T},
                Zy::DenseMatrix{T},
                xmean::Vector{T},
                ymean::Vector{T},
                p::Int) where T<:Real

    dx, n = size(Zx)
    dy, n2 = size(Zy)
    n == n2 ||
        throw(DimensionMismatch("Zx and Zy must have the same number of columns."))

    isempty(xmean) || length(xmean) == dx ||
        throw(DimensionMismatch("Incorrect length of xmean."))

    isempty(ymean) || length(ymean) == dy ||
        throw(DimensionMismatch("Incorrect length of ymean."))

    1 <= p <= min(dx, dy) ||
        throw(DimensionMismatch(""))

    _ccasvd(Zx, Zy, xmean, ymean, p)
end

# The implementation is partly based on:
#
#   David Weenink.
#   Canonical Correlation Analysis.
#   Institute of Phonetic Sciences, Univ. of Amsterdam,
#   Proceedings 25 (2003), 81-99.
#
#   Note: in this paper, each row is considered as an observation.
#   The algorithm is adpated to the column-major format here.
#
function _ccasvd(Zx::DenseMatrix{T}, Zy::DenseMatrix{T}, xmean::Vector{T}, ymean::Vector{T}, p::Int) where T<:Real
    # svd factorization of Z

    n = size(Zx, 2)

    # svd decomposition
    Sx = svd(Zx)
    Sy = svd(Zy)
    S = svd!(Sx.Vt * transpose(Sy.Vt)) # svd of Vx * Vy'

    # compute Px and Py
    ord = sortperm(S.S; rev=true)
    si = ord[1:p]
    Px = rmul!(Sx.U, Diagonal(1.0 ./ Sx.S)) * S.U[:, si]
    Py = rmul!(Sy.U, Diagonal(1.0 ./ Sy.S)) * S.V[:, si]

    # scale so that Px' * Cxx * Py == I
    #           and Py' * Cyy * Py == I,
    #
    # with Cxx = Zx * Zx' / (n - 1)
    #      Cyy = Zy * Zy' / (n - 1)
    #
    rmul!(Px, sqrt(n-1))
    rmul!(Py, sqrt(n-1))

    # compute correlations
    crs = rmul!(coldot(Zx'Px, Zy'Py), one(T)/(n-1))

    # construct CCA model
    CCA(xmean, ymean, Px, Py, crs)
end

## interface functions

function fit(::Type{CCA}, X::AbstractMatrix{T}, Y::AbstractMatrix{T};
             outdim::Int=min(min(size(X)...), min(size(Y)...)),
             method::Symbol=:svd,
             xmean=nothing,
             ymean=nothing) where T<:Real

    dx, n = size(X)
    dy, n2 = size(Y)

    n2 == n ||
        throw(DimensionMismatch("X and Y should have the same number of columns."))

    (n >= dx && n >= dy) ||
        @warn("CCA would be numerically instable when n < dx or n < dy.")

    xmv = preprocess_mean(X, xmean)
    ymv = preprocess_mean(Y, ymean)

    Zx = centralize(X, xmv)
    Zy = centralize(Y, ymv)

    if method == :cov
        Cxx = rmul!(Zx*transpose(Zx), inv(n - 1))
        Cyy = rmul!(Zy*transpose(Zy), inv(n - 1))
        Cxy = rmul!(Zx*transpose(Zy), inv(n - 1))
        M = ccacov(Cxx, Cyy, Cxy, xmv, ymv, outdim)
    elseif method == :svd
        M = ccasvd(Zx, Zy, xmv, ymv, outdim)
    else
        error("Invalid method name $(method)")
    end

    return M::CCA
end


"""
    Base.rand(randomwalk::RandomWalk{<:AbstractArray{<:DataType}, <:AbstractArray{<:Real},
                                     <:Sampleable}, ntimesteps::Integer) -> Matrix{Int}

Simulate `randomwalk` for `ntimesteps` time-steps over an integer lattice, returning one
sample path (i.e. trajectory) of length `ntimesteps`.

The time steps represented in the sample path are `0, 1, ..., ntimesteps-1`, for we start
with an initial state.

##### Complexity
- Time complexity: `O(n)`.
"""
function Base.rand(randomwalk::RandomWalk{<:AbstractArray{<:DataType},
                                          <:AbstractArray{<:Real},
                                          <:Sampleable}, ntimesteps::Integer)
    # allocate memory
    realizationmatrix = Matrix{Int}(undef, ntimesteps,
                                                   length(randomwalk.initialstate))
    realizationmatrix[1,:] = randomwalk.initialstate
    # draw a random step
    change = rand(randomwalk.stepsampler)
    for step in 2:ntimesteps
        realizationmatrix[step,:] = realizationmatrix[step-1,:] + change
        change = rand(randomwalk.stepsampler)
    end
    return realizationmatrix
end

"""
    Base.rand(randomwalk::RandomWalk{<:DataType, <:Real, <:Sampleable,
              ntimesteps::Integer) -> Vector{Int}

Simulate a one-dimensional `randomwalk` for `ntimesteps` time-steps, returning one sample
path (i.e. trajectory) of length `ntimesteps`.

The time steps represented in the sample path are `0, 1, ..., ntimesteps-1`, for we start
with an initial state.

##### Complexity
- Time complexity: `O(n)`.
"""
function Base.rand(randomwalk::RandomWalk{<:DataType, <:Real, <:Sampleable},
                   ntimesteps::Integer)
    # allocate memory
    realizationvec = Vector{eltype(randomwalk)}(undef, ntimesteps)
    realizationvec[1] = randomwalk.initialstate
    # draw a random step
    change = rand(randomwalk.stepsampler)[1]
    for step in 2:ntimesteps
        realizationvec[step] = realizationvec[step-1] + change
        change = rand(randomwalk.stepsampler)[1]
    end
    return realizationvec
end

"""
    Base.rand(randomwalk::RandomWalk{<:AbstractGraph, <:Integer, <:Sampleable},
              ntimesteps::Integer) -> Vector{Int}

Simulate `randomwalk` for `ntimesteps` time-steps over a graph, returning one sample path
(i.e. trajectory) of length `ntimesteps`.

The time steps represented in the sample path are `0, 1, ..., ntimesteps-1`, for we start
with an initial state.

##### Complexity
- Time complexity: `O(n)`.
"""
function Base.rand(randomwalk::RandomWalk{<:AbstractGraph, <:Integer, <:Sampleable},
                   ntimesteps::Integer)
    # allocate memory
    realizationvec = Vector{eltype(randomwalk)}(undef, ntimesteps)
    realizationvec[1] = randomwalk.initialstate
    for step in 2:ntimesteps
        # draw a neighbor at random
        realizationvec[step] = NeighborSampler(randomwalk.statespace,
                                               realizationvec[step-1]) |> rand
    end
    return realizationvec
end


#!jl # # EFTfitter.jl - BLUE Example
#jl # EFTfitter.jl - BLUE Example

#~When using multiple measurements of a single observable and a uniform prior for
#~the parameter representing the combined value, the combination of measurements performed with EFTfitter.jl
#~yields the same results as the Best Linear Unbiased Estimator (BLUE) method.

#~Here, we demonstrate this by using the examples of the paper
#~"How to combine correlated estimates of a single physical quantity"
#~by L. Lyons, D. Gibaut and P. Clifford (https://www.sciencedirect.com/science/article/pii/0168900288900186).
#~All numbers are taken from the example on charm particle lifetime experiments in section 5.
#~A factor of 10^13 is applied for convenience.

using EFTfitter
using BAT
using IntervalSets
using Statistics
using StatsBase
using LinearAlgebra
using Plots

#~We need one parameter for the best estimator and choose
#~a uniform distribution in the range 8 to 14 as prior:
parameters = BAT.NamedTupleDist(
    τ = 8..14,
)

#~When combining multiple measurements of the same observable,
#~only a function returning the combination parameter is needed:
estimator(params) = params.τ

#~In Eq. (17') of the reference paper the following covariance matrix is given:
covariance = [2.74 1.15 0.86 1.31;
              1.15 1.67 0.82 1.32;
              0.86 0.82 2.12 1.05;
              1.31 1.32 1.05 2.93]

#~For using this in EFTfitter.jl, we first need to convert the covariance matrix
#~into a correlation matrix and the corresponding uncertainty values:
corr, unc = EFTfitter.cov_to_cor(covariance)

measurements = (
    τ1 = Measurement(estimator,  9.5, uncertainties = (stat=unc[1],) ),
    τ2 = Measurement(estimator, 11.9, uncertainties = (stat=unc[2],) ),
    τ3 = Measurement(estimator, 11.1, uncertainties = (stat=unc[3],) ),
    τ4 = Measurement(estimator,  8.9, uncertainties = (stat=unc[4],) ),
)

correlations = (
    stat = Correlation(corr),
)

#~construct an `EFTfitterModel`:
model = EFTfitterModel(parameters, measurements, correlations)
posterior = PosteriorDensity(model);

#~sample the posterior with BAT.jl:
algorithm = MCMCSampling(mcalg =MetropolisHastings(), nsteps = 10^6, nchains = 4)
samples = bat_sample(posterior, algorithm).result

#~plot the posterior distribution for the combination parameter τ:
plot(samples, :τ, mean=true)
#md # ![blue plots](plots/plot_blue.png)

#~print numerical results of combination:
println("Mode: $(mode(samples).τ)")
println("Mean: $(mean(samples).τ) ± $(std(samples).τ)")
#md ```
#md Mode: 11.15985
#md Mean: 11.15471 ± 0.80180
#md ```

#!jl # ### Comparison with BLUE method
#jl # comparison with BLUE method:
blue = BLUE(model)
println("BLUE: $(blue.value) ± $(blue.unc)")
println("BLUE weights: $(blue.weights)")
#md ```
#md BLUE: 11.15983 ± 1.28604
#md BLUE weights: [0.145, 0.470, 0.347, 0.038]
#md ```


export birdcageSensitivity

"""
    birdcageSensitivity(N::Int64, ncoils::Int64, relative_radius::Float64)

Computes the sensitivity maps for each coils that are arranged
in a birdcage manner.
"""
function birdcageSensitivity(N::Int64,
                             ncoils::Int64,
                             relative_radius::Float64)

  out = zeros(ComplexF64, N, N, 1, ncoils)
  for c=0:ncoils-1
    coilx = relative_radius*cos(c*(2*pi/ncoils))
    coily = relative_radius*sin(c*(2*pi/ncoils))
    coil_phase = -c*(2*pi/ncoils)

    for y=0:N-1
      y_co = (y - (N/2))/(N/2) - coily
      for x=0:N-1
        x_co = (x - (N/2))/(N/2) - coilx
        rr = sqrt(x_co^2+y_co^2)
        phi = atan(x_co, -y_co) + coil_phase
        out[x+1,y+1, 1, c+1] = 1/(rr) * exp(1im*phi)
      end
    end
  end
 return out
end


function generate_synthetic_data(n, param_frac, sigma; signal_to_noise=10,
    fixed_rank_sparse=false, fixed_rank=1, fixed_sparse=0)
    """
    This function samples synthetic data by adding noise to the sum of a
    random low rank matrix and a random sparse matrix.

    :param n: The row and column dimension of the data to be sampled (Int64).
    :param param_frac: Parameter that controls the rank and sparsity of the
                       sampled low rank and sparse matrix when
                       fixed_rank_sparse=false (Float64).
    :param sigma: Parameter that controls the absolute magnitude of the noise
                  (Float64).
    :param signal_to_noise: Parameter that controls the signal to noise ratio of
                            the sampled data (Float64).
    :param fixed_rank_sparse: If true, use fixed_rank and fixed_sparse as the
                              rank and sparsity respectively of the low rank
                              matrix and the sparse matrix to be sampled (Bool).
    :param fixed_rank: Rank of the low rank matrix to be sampled when
                       fixed_rank_sparse=true.
    :param fixed_sparse: Sparsity of the sparse matrix to be sampled when
                         fixed_rank_sparse=true

    :return: This function returns 5 values:
             1) The sampled n x n data matrix.
             2) The sampled n x n low rank matrix.
             3) The sampled n x n sparse matrix.
             4) The sparsity level of the sparse matrix.
             5) The rank of the low rank matrix.
    """

    if fixed_rank_sparse
        k_sparse = fixed_sparse
        k_rank = fixed_rank
    else
        k_sparse = Int64(floor(param_frac * n^2))
        k_rank = Int64(floor(param_frac * n))
    end

    # Sample a rank k_rank matrix
    d = Normal(0, signal_to_noise * sigma / n^0.5)
    X = rand(d, (n, k_rank))
    Y = rand(d, (n, k_rank))
    L_0 = X * Y'

    # Sample a sparsity k_sparse matrix
    d = Uniform(-5, 5)
    E = rand(d, (n, n))
    S_0 = zeros(n, n)

    indexes = []
    for i = 1:n
        for j = 1:n
            append!(indexes, [(i, j)])
        end
    end

    select_indexes = sample(indexes, k_sparse, replace=false)
    for index in select_indexes
        S_0[index[1], index[2]] = E[index[1], index[2]]
    end

    data = L_0 + S_0

    # Add noise to the sum of the sampled low rank and sparse matrices
    d = Normal(0, sigma^2)
    noise = rand(d, (n, n))
    data = data + noise

    return data, L_0, S_0, k_sparse, k_rank

end;

function generate_synthetic_data_symmetric(n, param_frac, sigma;
    signal_to_noise=10, fixed_rank_sparse=false, fixed_rank=1, fixed_sparse=0)
    """
    This function samples symmetric synthetic data by adding symmetric noise to
    the sum of a random symmetric low rank matrix and a random symmetric sparse
    matrix.

    :param n: The row and column dimension of the data to be sampled (Int64).
    :param param_frac: Parameter that controls the rank and sparsity of the
                       sampled low rank and sparse matrix when
                       fixed_rank_sparse=false (Float64).
    :param sigma: Parameter that controls the absolute magnitude of the noise
                  (Float64).
    :param signal_to_noise: Parameter that controls the signal to noise ratio of
                            the sampled data (Float64).
    :param fixed_rank_sparse: If true, use fixed_rank and fixed_sparse as the
                              rank and sparsity respectively of the low rank
                              matrix and the sparse matrix to be sampled (Bool).
    :param fixed_rank: Rank of the low rank matrix to be sampled when
                       fixed_rank_sparse=true.
    :param fixed_sparse: Sparsity of the sparse matrix to be sampled when
                         fixed_rank_sparse=true

    :return: This function returns 5 values:
             1) The sampled n-by-n data matrix.
             2) The sampled n-by-n low rank matrix.
             3) The sampled n-by-n sparse matrix.
             4) The sparsity level of the sparse matrix.
             5) The rank of the low rank matrix.
    """

    if fixed_rank_sparse
        k_sparse = fixed_sparse
        k_rank = fixed_rank
    else
        k_sparse = Int64(floor(param_frac * n^2))
        k_rank = Int64(floor(param_frac * n))
    end

    # Sample a rank k_rank symmetric matrix
    d = Normal(0, signal_to_noise * sigma / n^0.5)
    X = rand(d, (n, k_rank))
    L_0 = X * X'

    # Sample a sparsity k_sparse symmetric matrix
    d = Uniform(-5, 5)
    E = rand(d, (n, n))
    S_0 = zeros(n, n)

    indexes = []
    for i = 1:(n-1)
        for j = (i+1):n
            append!(indexes, [(i, j)])
        end
    end

    select_indexes = sample(indexes, Int64(floor(k_sparse/2)), replace=false)
    for index in select_indexes
        S_0[index[1], index[2]] = E[index[1], index[2]]
        S_0[index[2], index[1]] = S_0[index[1], index[2]]
    end

    if k_sparse % 2 == 1
        index = sample(1:n, 1)
        S_0[index, index] = E[index, index]
    end

    data = L_0 + S_0

    # Add symmetric noise to the sum of the sampled low rank and sparse matrices
    d = Normal(0, sigma^2)
    noise = rand(d, (n, n))
    for i = 1:(n-1)
        for j = (i+1):n
            noise[j, i] = noise[i, j]
        end
    end
    data = data + noise

    return data, L_0, S_0, k_sparse, k_rank

end;


# In this tutorial, we will learn
#  - How to solve nonlinear multi-field PDEs in Gridap
#  - How to build FE spaces whose functions have zero mean value
#
# ## Problem statement
#
# The goal of this last tutorial is to solve a nonlinear multi-field PDE. As a model problem, we consider a well known benchmark in computational fluid dynamics, the lid-driven cavity for the incompressible Navier-Stokes equations. Formally, the PDE we want to solve is: find the velocity vector $u$ and the pressure $p$ such that
#
# ```math
# \left\lbrace
# \begin{aligned}
# -\Delta u + \mathit{Re}\ (u\cdot \nabla)\ u + \nabla p = 0 &\text{ in }\Omega,\\
# \nabla\cdot u = 0 &\text{ in } \Omega,\\
# u = g &\text{ on } \partial\Omega,
# \end{aligned}
# \right.
# ```
#
# where the computational domain is the unit square $\Omega \doteq (0,1)^d$, $d=2$, $\mathit{Re}$ is the Reynolds number (here, we take $\mathit{Re}=10$), and $(w \cdot \nabla)\ u = (\nabla u)^t w$  is the well known convection operator. In this example, the driving force is the Dirichlet boundary velocity $g$, which is a non-zero horizontal velocity with a value of $g = (1,0)^t$ on the top side of the cavity, namely the boundary $(0,1)\times\{1\}$, and $g=0$ elsewhere on $\partial\Omega$. Since we impose Dirichlet boundary conditions on the entire boundary $\partial\Omega$, the mean value of the pressure is constrained to zero in order have a well posed problem,
#
# ```math
# \int_\Omega q \ {\rm d}\Omega = 0.
# ```
#
# ## Numerical Scheme
#
# In order to approximate this problem we chose a formulation based on inf-sub stable $Q_k/P_{k-1}$ elements with continuous velocities and discontinuous pressures (see, e.g., [1] for specific details). The interpolation spaces are defined as follows.  The velocity interpolation space is
#
# ```math
# V \doteq \{ v \in [C^0(\Omega)]^d:\ v|_T\in [Q_k(T)]^d \text{ for all } T\in\mathcal{T} \},
# ```
# where $T$ denotes an arbitrary cell of the FE mesh $\mathcal{T}$, and $Q_k(T)$ is the local polynomial space in cell $T$ defined as the multi-variate polynomials in $T$ of order less or equal to $k$ in each spatial coordinate. Note that, this is the usual continuous vector-valued Lagrangian FE space of order $k$ defined on a mesh of quadrilaterals or hexahedra.  On the other hand, the space for the pressure is
#
# ```math
# \begin{aligned}
# Q_0 &\doteq \{ q \in Q: \  \int_\Omega q \ {\rm d}\Omega = 0\}, \text{ with}\\
# Q &\doteq \{ q \in L^2(\Omega):\ q|_T\in P_{k-1}(T) \text{ for all } T\in\mathcal{T}\},
# \end{aligned}
# ```
# where $P_{k-1}(T)$ is the polynomial space of multi-variate polynomials in $T$ of degree less or equal to $k-1$. Note that functions in $Q_0$ are strongly constrained to have zero mean value. This is achieved in the code by removing one degree of freedom from the (unconstrained) interpolation space $Q$ and  adding a constant to the computed pressure so that the resulting function has zero mean value.
#
# The weak form associated to these interpolation spaces reads: find $(u,p)\in U_g \times Q_0$ such that $[r(u,p)](v,q)=0$ for all $(v,q)\in V_0 \times Q_0$
# where $U_g$ and $V_0$ are the set of functions in $V$ fulfilling the Dirichlet boundary condition $g$ and $0$  on $\partial\Omega$ respectively. The weak residual $r$ evaluated at a given pair $(u,p)$ is the linear form defined as
#
# ```math
# [r(u,p)](v,q) \doteq a((u,p),(v,q))+ [c(u)](v),
# ```
# with
# ```math
# \begin{aligned}
# a((u,p),(v,q)) &\doteq \int_{\Omega} \nabla v \cdot \nabla u \ {\rm d}\Omega - \int_{\Omega} (\nabla\cdot v) \ p \ {\rm d}\Omega + \int_{\Omega} q \ (\nabla \cdot u) \ {\rm d}\Omega,\\
# [c(u)](v) &\doteq \int_{\Omega} v 	\cdot \left( (u\cdot\nabla)\ u \right)\ {\rm d}\Omega.\\
# \end{aligned}
# ```
# Note that the bilinear form $a$ is associated with the linear part of the PDE, whereas $c$ is the contribution to the residual resulting from the convective term.
#
# In order to solve this nonlinear weak equation with a Newton-Raphson method, one needs to compute the Jacobian associated with the residual $r$. In this case, the Jacobian $j$ evaluated at a pair $(u,p)$ is the bilinear form defined as
#
# ```math
# [j(u,p)]((\delta u, \delta p),(v,q)) \doteq a((\delta u,\delta p),(v,q))  + [{\rm d}c(u)](\delta u,v),
# ```
# where ${\rm d}c$ results from the linearization of the convective term, namely
# ```math
# [{\rm d}c(u)](\delta u,v) \doteq \int_{\Omega} v \cdot \left( (u\cdot\nabla)\ \delta u \right) \ {\rm d}\Omega + \int_{\Omega} v \cdot \left( (\delta u\cdot\nabla)\ u \right)  \ {\rm d}\Omega.
# ```
# The implementation of this numerical scheme is done in Gridap by combining the concepts previously seen for single-field nonlinear PDEs  and linear multi-field problems.
#
# ## Discrete model
#
# We start with the discretization of the computational domain. We consider a $100\times100$ Cartesian mesh of the unit square.

using Gridap
n = 100
domain = (0,1,0,1)
partition = (n,n)
model = CartesianDiscreteModel(domain,partition)

# For convenience, we create two new boundary tags,  namely `"diri1"` and `"diri0"`, one for the top side of the square (where the velocity is non-zero), and another for the rest of the boundary (where the velocity is zero).

labels = get_face_labeling(model)
add_tag_from_tags!(labels,"diri1",[6,])
add_tag_from_tags!(labels,"diri0",[1,2,3,4,5,7,8])

# ## FE spaces
#
# For the velocities, we need to create a conventional vector-valued continuous Lagrangian FE space. In this example, we select a second order interpolation.

D = 2
order = 2
reffeᵤ = ReferenceFE(lagrangian,VectorValue{2,Float64},order)
V = TestFESpace(model,reffeᵤ,conformity=:H1,labels=labels,dirichlet_tags=["diri0","diri1"])

# The interpolation space for the pressure is built as follows

reffeₚ = ReferenceFE(lagrangian,Float64,order-1;space=:P)
Q = TestFESpace(model,reffeₚ,conformity=:L2,constraint=:zeromean)

# With the options `:Lagrangian`, `space=:P`, `valuetype=Float64`, and `order=order-1`, we select the local polynomial space $P_{k-1}(T)$ on the cells $T\in\mathcal{T}$. With the symbol `space=:P` we specifically chose a local Lagrangian interpolation of type "P". Without using `space=:P`, would lead to a local Lagrangian of type "Q" since this is the default for quadrilateral or hexahedral elements. On the other hand, `constraint=:zeromean` leads to a FE space, whose functions are constrained to have mean value equal to zero, which is just what we need for the pressure space. With these objects, we build the trial multi-field FE spaces

uD0 = VectorValue(0,0)
uD1 = VectorValue(1,0)
U = TrialFESpace(V,[uD0,uD1])
P = TrialFESpace(Q)

Y = MultiFieldFESpace([V, Q])
X = MultiFieldFESpace([U, P])

# ## Triangulation and integration quadrature
#
# From the discrete model we can define the triangulation and integration measure

degree = order
Ωₕ = Triangulation(model)
dΩ = Measure(Ωₕ,degree)

# ## Nonlinear weak form
#
# The different terms of the nonlinear weak form for this example are defined following an approach similar to the one discussed for the $p$-Laplacian equation, but this time using the notation for multi-field problems.

const Re = 10.0
conv(u,∇u) = Re*(∇u')⋅u
dconv(du,∇du,u,∇u) = conv(u,∇du)+conv(du,∇u)

# The bilinear form reads
a((u,p),(v,q)) = ∫( ∇(v)⊙∇(u) - (∇⋅v)*p + q*(∇⋅u) )dΩ

# The nonlinear term and its Jacobian are given by
c(u,v) = ∫( v⊙(conv∘(u,∇(u))) )dΩ
dc(u,du,v) = ∫( v⊙(dconv∘(du,∇(du),u,∇(u))) )dΩ

# Finally, the Navier-Stokes weak form residual and Jacobian can be defined as
res((u,p),(v,q)) = a((u,p),(v,q)) + c(u,v)
jac((u,p),(du,dp),(v,q)) = a((du,dp),(v,q)) + dc(u,du,v)

# With the functions `res`, and `jac` representing the weak residual and the Jacobian, we build the nonlinear FE problem:
op = FEOperator(res,jac,X,Y)

# ## Nonlinear solver phase
#
# To finally solve the problem, we consider the same nonlinear solver as previously considered for the  $p$-Laplacian equation.

using LineSearches: BackTracking
nls = NLSolver(
  show_trace=true, method=:newton, linesearch=BackTracking())
solver = FESolver(nls)

# In this example, we solve the problem without providing an initial guess (a default one equal to zero will be generated internally)

uh, ph = solve(solver,op)

# Finally, we write the results for visualization (see next figure).

writevtk(Ωₕ,"ins-results",cellfields=["uh"=>uh,"ph"=>ph])

# ![](../assets/inc_navier_stokes/ins_solution.png)
#
#  ## References
#
#  [1] H. C. Elman, D. J. Silvester, and A. J. Wathen. *Finite elements and fast iterative solvers: with applications in incompressible fluid dynamics*. Oxford University Press, 2005.


# Simulate a Poisson point process on a triangle
# Author: H. Paul Keeler, 2019.
# Website: hpaulkeeler.com
# Repository: github.com/hpaulkeeler/posts
# For more details, see the post:
# hpaulkeeler.com/testing-the-julia-language-with-point-process-simulations/

#Note: Need the .+ for adding a scalar to an array
#Also need . for sqrt, exp, cos, sin etc and assinging scalars to arrays
#Don't confuse 1.-V with 1 .-V for array V
#Best to use vectors instead of 1-D matrices eg x=rand(n),  NOT x=rand(n,1).

using Distributions #for random simulations
using Plots #for plotting

#Simulation window parameters -- points A,B,C of a triangle
xA=0;xB=1;xC=1; #x values of three points
yA=0;yB=0;yC=1; #y values of three points;

#Point process parameters
lambda=100; #ntensity (ie mean density) of the Poisson process

#calculate sides of trinagle
a=sqrt((xA-xB)^2+(yA-yB)^2);
b=sqrt((xB-xC)^2+(yB-yC)^2);
c=sqrt((xC-xA)^2+(yC-yA)^2);
s=(a+b+c)/2; #calculate semi-perimeter

#Use Herron's forumula to calculate area
areaTotal=sqrt(s*(s-a)*(s-b)*(s-c)); #area of triangle

#Simulate a Poisson point process
numbPoints=rand(Poisson(areaTotal*lambda)); #Poisson number of points
U=rand(numbPoints);#uniform random variables
V=rand(numbPoints);#uniform random variables

xx=sqrt.(U).*xA .+sqrt.(U).*(1 .-V).*xB .+sqrt.(U).*V.*xC;#x coordinates of points
yy=sqrt.(U).*yA .+sqrt.(U).*(1 .-V).*yB .+sqrt.(U).*V.*yC;#y coordinates of points

#Plotting
plot1=scatter(xx,yy,xlabel ="x",ylabel ="y", leg=false);
display(plot1);


function compute(coins::Vector{Int}, n::Int)::Int
    combinations = vcat([1], repeat([0], n))
    for i ∈ 1:length(coins)
        for j ∈ coins[i] + 1:n + 1
            combinations[j] += combinations[j - coins[i]]
        end
    end
    return combinations[end]
end


"""

- `⁺pᵢ` : Particle i at next step
- `pᵢ`  : Particle i
- `pⱼ`  : Particle j
- `ϵ`   : Softening parameter
"""
function evaluate_Euler!(⁺aᵢ, pᵢ, pⱼ, ϵ)
    r = pⱼ.r .- pᵢ.r
    r⁻³ = (norm(r)^2 + ϵ^2)^(-3/2)

    @. ⁺aᵢ += G * pⱼ.m * r⁻³ * r
end

function evaluate_Euler!(ps, ϵ)
    for pᵢ in ps
        pᵢ.a .= 0.
        for pⱼ in ps
            pᵢ != pⱼ && evaluate_Euler!(pᵢ.a, pᵢ, pⱼ, ϵ)
        end
    end
end

function update!(p::AbstractParticle, Δt)
    @. p.v += p.a * Δt
    @. p.r += p.v * Δt
end

function update!(ps, Δt)
    for p in ps
        update!(p, Δt)
    end
end

function run_Euler(ps, params_sim)
    @unpack ϵ, Δt, t_end = params_sim

    times = (0:Δt:t_end)
    ts, rs, vs, as = prep_snapshot(ps, params_sim)

    for (i, t) in enumerate(times)
        evaluate_Euler!(ps, ϵ)
        save_snapshot!(ts, rs, vs, as, i, t, ps, params_sim)
        update!(ps, Δt)
    end
    ts, rs, vs, as
end


## Exercise 2-3
## Practice using the Julia REPL as a calculator:

## 1. The volume of a sphere with radius r is (4/3)πr^3. What is the volume of a sphere with radius 5?
println("Ans 1:")
println("the volume of sphere with radius 5 is: " * string((4/3) * π * 5 ^ 3))

## 2. Suppose the cover price of a book is $ 24.95, but bookstores get a 40 % discount. Shipping costs $ 3 for the first copy and 75 cents for each additional copy. What is the total wholesale cost for 60 copies?
println("Ans 2:")
cover_price = 24.95
n_copies = 60
shipping_cost = 3 + (n_copies - 1) * 0.75
books_price = n_copies * cover_price * 0.6
total = shipping_cost + books_price
println(total)

## 3. If I leave my house at 6:52 am and run 1 mile at an easy pace (8:15 per mile), then 3 miles at tempo (7:12 per mile) and 1 mile at easy pace again, what time do I get home for breakfast?
println("Ans 3:")
easy_pace_time = 8*60 + 15
tempo_pace_time = 7*60 + 12

total_time = 2 * easy_pace_time + 3 * tempo_pace_time

seconds = total_time % 60
minutes = total_time ÷ 60

println("you will reach home at 6:52 + ", minutes, " minutes and ", seconds, " seconds")
println("End.")