using Plots
using LinearAlgebra

function plot_plane(center, spanning_vector_a, spanning_vector_b, color=:grey, ε = 1/2, frame=false, frame_color=:blue)
	x = [center[1] + ε * spanning_vector_a[1] + ε * spanning_vector_b[1],
	     center[1] - ε * spanning_vector_a[1] + ε * spanning_vector_b[1],
	     center[1] + ε * spanning_vector_a[1] - ε * spanning_vector_b[1],
	     center[1] - ε * spanning_vector_a[1] - ε * spanning_vector_b[1]]

	y = [center[2] + ε * spanning_vector_a[2] + ε * spanning_vector_b[2],
	     center[2] - ε * spanning_vector_a[2] + ε * spanning_vector_b[2],
	     center[2] + ε * spanning_vector_a[2] - ε * spanning_vector_b[2],
	     center[2] - ε * spanning_vector_a[2] - ε * spanning_vector_b[2]]

	z = [center[3] + ε * spanning_vector_a[3] + ε * spanning_vector_b[3],
	     center[3] - ε * spanning_vector_a[3] + ε * spanning_vector_b[3],
	     center[3] + ε * spanning_vector_a[3] - ε * spanning_vector_b[3],
	     center[3] - ε * spanning_vector_a[3] - ε * spanning_vector_b[3]]

	if frame
		path3d!(x,y,z;color=frame_color)
	end
	mesh3d!(x,y,z;color,connections=([0,3],[1,2],[3,0]))
end


using QuadGK, HCubature

"""
Create a new directory `dirname` in the current path.
Append a number to the name if the directory already exists.
"""
function make_new_dir(dirname::String)
    if ispath(dirname)
        k=1
        while ispath(dirname*"_$k")
            k+=1
        end
        return mkdir(dirname*"_$k")
    else
        return mkdir(dirname)
    end
end

"""
    nice_number(x,K::Int)

Return a string which is the rounded value of `x` with `K` (decimal) digits and
if necessary, pad `"0"` to the right so that there are exactly `K` decimal places.
"""
function nice_float2string(x,K::Int)
    integer_part = trunc(Int,x)
    y = x - integer_part
    decimal_part = rpad(round(Int,y*10^K),K,"0")
    return "$(integer_part).$decimal_part"
end

"""
Return the coefficients `c1,c2,λ` of the SOH model computed for the `model` `"Fokker-Planck"` (default)
or `"BGK"` with the concentration parameters `κ`.
These coefficients are given by

c₁ = ∫\\_[0,π] cos(θ)exp(κcos(θ))dθ / ∫\\_[0,π] exp(κcos(θ))dθ

c₂ = ∫\\_[0,π] cos(θ)sin(θ)g(θ)exp(κcos(θ))dθ / ∫\\_[0,π] sin(θ)g(θ)exp(κcos(θ))dθ

λ = 1/κ

where

g(θ) = θ/κ - π/κ * ∫\\_[0,θ] exp(-κcos(ψ))dψ / ∫\\_[0,π] exp(-κcos(ψ))dψ in the Fokker-Planck case

and

g(θ) = 1 in the BGK case.
"""
function coefficients_Vicsek(κ;model="Fokker-Planck")
    λ = 1/κ
    I1 = quadgk(θ->cos(θ)*exp(κ*cos(θ)),0,pi,rtol=1e-5)[1]
    Z1 = quadgk(θ->exp(κ*cos(θ)),0,pi,rtol=1e-5)[1]
    c1 = I1/Z1
    if model=="BGK"
        I2 = quadgk(θ->cos(θ)*sin(θ)^2*exp(κ*cos(θ)),0,pi,rtol=1e-5)[1]
        Z2 = quadgk(θ->sin(θ)^2*exp(κ*cos(θ)),0,pi,rtol=1e-5)[1]
        c2 = I2/Z2
        return c1,c2,λ
    elseif model=="Fokker-Planck"
        Zinv = quadgk(θ->exp(-κ*cos(θ)),0,pi,rtol=1e-5)[1]
        I21 = quadgk(θ->θ*cos(θ)*sin(θ)*exp(κ*cos(θ)),0,pi,rtol=1e-5)[1]/κ
        I22 = pi/(κ*Zinv) * hcubature(x->cos(x[1])*sin(x[1])*exp(κ*cos(x[1]))*exp(-κ*cos(x[2]))*(x[2]<x[1]),(0,0),(pi,pi),rtol=1e-3)[1]
        I2 = I21 - I22
        Z21 = quadgk(θ->θ*sin(θ)*exp(κ*cos(θ)),0,pi,rtol=1e-5)[1]/κ
        Z22 = pi/(κ*Zinv) * hcubature(x->sin(x[1])*exp(κ*cos(x[1]))*exp(-κ*cos(x[2]))*(x[2]<x[1]),(0,0),(pi,pi),rtol=1e-3)[1]
        Z2 = Z21 - Z22
        c2 = I2/Z2
        return c1,c2,λ
    else
        ArgumentError("Model not defined!")
    end
end


# calculate eigen values/vectors
# plot eigen vectors on the contour

module CalcPlotEigen
    using Plots, Base.Iterators, LinearAlgebra
    pyplot()

    function pdf(x, mu, cov_mat)
        # determinant
        det_mat = det(cov_mat)

        diff_x = x - mu
        return exp(-((diff_x'/cov_mat)*diff_x)/2)/(2*pi*sqrt(det_mat))
    end

    function main(is_test=false)
        # x-y grid
        vx = 0:200
        vy = 0:100

        # mu, covariance matrix
        mu = [150; 50]
        cov = [100 -25*sqrt(3); -25*sqrt(3) 50]

        # calculate eigen values/vectors
        eig_vals, eig_vecs = eigen(cov)
        println("Eigen Values: $(eig_vals)")
        println("Eigen Vectors: $(eig_vecs)")
        println("Eigen Vector1: $(eig_vecs[:, 1])")
        println("Eigen Vector2: $(eig_vecs[:, 2])")

        # long or short
        if eig_vals[1] > eig_vals[2]
            long_val = eig_vals[1]
            long_vec = eig_vecs[:, 1]
            short_val = eig_vals[2]
            short_vec = eig_vecs[:, 2]
        else
            long_val = eig_vals[2]
            long_vec = eig_vecs[:, 2]
            short_val = eig_vals[1]
            short_vec = eig_vecs[:, 1]
        end
        println("Long Value: $(long_val)")
        println("Long Vector: $(long_vec)")
        println("Short Value: $(short_val)")
        println("Short Vector: $(short_vec)")

        # plot
        z = [pdf([x; y], mu, cov) for x in vx, y in vy]
        contour(vx, vy, z', c=:haline, aspect_ratio=:equal)
        long = 2 * sqrt(long_val) * long_vec
        short = 2 * sqrt(short_val) * short_vec
        quiver!([mu[1]], [mu[2]], quiver=([long[1]], [long[2]]),
                c=:red, aspect_ratio=:equal)
        quiver!([mu[1]], [mu[2]], quiver=([short[1]], [short[2]]),
                c=:blue, aspect_ratio=:equal)

        if is_test == false
            save_path = joinpath(split(@__FILE__, "src")[1], "src/prob_stats/multi_dim_gauss_dist/eigen/contour_eigen.png")
            savefig(save_path)
        end
    end
end


@registerN Hahn AbstractCDOP3 α β 𝐍
export Hahn

"""
    Hahn{α,β,𝐍}

References: [Koekoek and Swarttouw §1.5](https://arxiv.org/pdf/math/9602214.pdf)

!!! note
    In  [Koekoek and Swarttouw](https://arxiv.org/pdf/math/9602214.pdf) sections 1.3, 1.4, and 1.6  are other  Hahn-type polynomials,  not  implemented here.

"""
Hahn

abcde(::Type{<:Hahn{α,β,𝐍}})  where {α,β,𝐍} = NamedTuple{(:a,:b,:c,:d,:e)}((1, -(β+𝐍+1), 0,  α+β+2, -𝐍*(α+1)))

basis_symbol(::Type{<:Hahn{α,β,𝐍}}) where {α,β,𝐍} = "Q⁽ᵅᵝ⁾"
Polynomials.domain(::Type{<:Hahn{α, β, 𝐍}}) where {α, β,  𝐍} = Polynomials.Interval(0, 𝐍)
weight_function(::Type{<:Hahn{α,β,𝐍}}) where {α,β,𝐍} = x -> generalized_binomial(α+x,x) * generalized_binomial(𝐍+β-x,𝐍-x)

# use   recurrence relation  1.5.3 of Koekoek and Swarttouw
# -xQᵢ = AQᵢ₊₁ -  (A+C)Qᵢ + CQᵢ₋₁
#  or Qᵢ₊₁ = (1/A ⋅ x - (1 + C/A)) Qᵢ + C/A ⋅ Qᵢ₋₁
# Using kᵢ = A⁻¹ᵢ₋₁ ⋅ A⁻¹ᵢ₋₂ ⋯ A⁻¹₀; so kᵢ₊₁/kᵢ = A⁻¹ᵢ, kᵢ₊₁/kᵢ₋ᵢ = A⁻¹ᵢA⁻¹ᵢ₋₁
#
# function Aᵢ⁻¹(n,α,β,𝐍)
#     num = (2n+α+β + 1) * (2n+α+β+2)
#     den = (n+α + β + 1)*(n + α + 1) *(𝐍-n)
#     num/den
# end

# define kn, k1k_1 through defaults
function k1k0(P::Type{<:Hahn{α,β, 𝐍}}, n::Int) where {α,β,𝐍}

    num  = (2*n + α + β + 1)*(2*n + α + β + 2)
    den = (-n + 𝐍)*(n + α + 1)*(n + α + β + 1)

    -(one(eltype(P))*num)/den

end

# end

function classical_hypergeometric(::Type{<:Hahn{α,β,𝐍}}, n::Int, x)  where {α,β,𝐍}
    pFq((-n, -x, n+1+α+β),  (α+1, -𝐍),  1)
end


using TravelingSalesmanExact, GLPK
using TravelingSalesmanHeuristics
using TravelingSalesmanExact: ATT, euclidean_distance
using TravelingSalesmanBenchmarks
using Plots
using Printf, Random
gr(fmt=:svg)

function get_plt_coords(cities)
    n = length(cities)
    inc(a) = a == n ? one(a) : a + 1
    return [cities[inc(j)][1] for j = 0:n], [cities[inc(j)][2] for j = 0:n]
end
TravelingSalesmanExact.plot_cities(cities; kwargs...) = plot(get_plt_coords(cities)...; kwargs...)
plot_cities!(cities; kwargs...) = plot!(get_plt_coords(cities)...; kwargs...)

function plot_tours(cities, pairs; kwargs...)
    plts = []
    colors = sequential_palette(0, length(pairs)+1)[2:end]
    for (index, (tour, label)) in enumerate(pairs)
        plt = plot_cities(cities[tour], label = label, linewidth = 2,
                            color = colors[index])
        plot_cities!(cities; title="Comparison of tours", st=:scatter, label="City locations", markersize = 5, kwargs...)
        push!(plts, plt)
    end
    return plts
end

Random.seed!(155);
N = 50
cities = [100*rand(2) for _ = 1:N]
cost = [ euclidean_distance(cities[i], cities[j]) for i=1:N, j=1:N ]

t_exact, c_exact = get_optimal_tour(cost, with_optimizer(GLPK.Optimizer))

c(q) = solve_tsp(cost; quality_factor = q)[2]

qs = range(10, stop = 100, step = 10)

plot(qs, c, xlabel="quality", ylabel="Cost", label="solve_tsp 1", title="random cities")
for j = 2:5
    plot!(qs, c, label="solve_tsp $j")
end
hline!([c_exact], label="Exact cost")

t_heuristic_100, c_heuristic_100 = solve_tsp(cost; quality_factor = 100)
t_heuristic_40, c_heuristic_40 = solve_tsp(cost; quality_factor = 40)

plts = plot_tours(cities, [
    (t_heuristic_40, @sprintf("Heuristic; quality 40, cost=%.2f",   c_heuristic_40)),
    (t_heuristic_100, @sprintf("Heuristic; quality 100, cost=%.2f", c_heuristic_100)),
    (t_exact, @sprintf("Optimal tour; cost=%.2f", c_exact))
    ])
for plt in plts
    display(plt)
end

function compare_cities(N)
    cities = [100*rand(2) for _ = 1:N]
    cost = [ euclidean_distance(cities[i], cities[j]) for i=1:N, j=1:N ]
    t_exact, c_exact = get_optimal_tour(cost, with_optimizer(GLPK.Optimizer))

    c(q) = solve_tsp(cost; quality_factor = q)[2]

    qs = range(10, stop = 100, step = 10)

    line_plt = plot(qs, c, xlabel="quality", ylabel="Cost", label="solve_tsp 1", title="random cities")
    for j = 2:5
        plot!(qs, c, label="solve_tsp $j")
    end
    hline!([c_exact], label="Exact cost")

    t_heuristic_100, c_heuristic_100 = solve_tsp(cost; quality_factor = 100)
    t_heuristic_40, c_heuristic_40 = solve_tsp(cost; quality_factor = 40)
    city_plts = plot_tours(cities, [
    (t_heuristic_40, @sprintf("Heuristic; quality 40, cost=%.2f",   c_heuristic_40)),
    (t_heuristic_100, @sprintf("Heuristic; quality 100, cost=%.2f", c_heuristic_100)),
    (t_exact, @sprintf("Optimal tour; cost=%.2f", c_exact))
    ])
    return line_plt, city_plts...
end

for j = 1:5
    plts = compare_cities(50)
    for plt in plts
        display(plt)
    end
end

file = @isdefined(WEAVE_ARGS) ? WEAVE_ARGS[:file] : nothing
TravelingSalesmanBenchmarks.bench_footer(file)


using Tables
using DataFrames
using StatsPlots
using BenchmarkTools

function sir_ode!(du,u,p,t)
    (S,I,R) = u
    (β,c,γ) = p
    N = S+I+R
    @inbounds begin
        du[1] = -β*c*I/N*S
        du[2] = β*c*I/N*S - γ*I
        du[3] = γ*I
    end
    nothing
end;

t0 = 0.0
δt = 0.1
tmax = 40.0
u0 = [990.0,10.0,0.0] # S,I,R
p = [0.05,10.0,0.25]; # β,c,γ

function euler(f, u0, p, δt, t0, tmax)
    t = t0 # Initialize time
    u = copy(u0) # Initialize struct parametric inherited
    du = zeros(length(u0)) # Initialize derivatives
    f(du,u,p,t)
    sol = [] # Store output
    times = [] # Store times
    push!(sol,copy(u))
    push!(times,t)
    # Main loop
    while t < tmax
        t = t + δt # Update time
        u .= u .+ du.*δt # Update state
        sir_ode!(du,u,p,t) # Update derivative
        push!(sol,copy(u)) # Store output
        push!(times,t) # Store time
    end
    sol = hcat(sol...) # Convert to matrix
    return times, sol
end;

times, sol = euler(sir_ode!, u0, p, δt, t0, tmax);

df = DataFrame(Tables.table(sol'))
rename!(df,["S","I","R"])
df[!,:t] = times;

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

@benchmark euler(sir_ode!, u0, p, δt, t0, tmax)


# Nonlinear curve fitting.

"""
# Nonlinear multi-variable least squares

Uses a Newton like algorithm to compute the least squares fit
of a model

## Parameters:

 * `x` an array where each colum is a variable
 * `fun` The function that should be fitted (more later on)
 * `a0` An initial guess for each fitting parameter
 * `eps` Convergence criterion
 * `maxiter` Maximum number of iterations

## Specifying the model

The model is specified in argument `fun` that should be callable according
to the following signature:

`r = fun(x, a)`

where both `x` and `a` are vectors

The way the algorithm is implemented allows for the function to be implicit. This means
that we are trying to fit the relationship `fun(x, a) = 0`.

## Initial guess

The initial guess is important and should be as close as possible to the expeted values. It should, at least, give an order of magnitude of each parameter. Zero values are not recommended since in this
case no order of magnitude is available and it is assumed. The initial step is assumed to be `a0/10`.

## Convergence criterion

The initial guess `a0` will probably not fit the data very well. The largest value
of `maxerr = maxabs(fun(x, a0))` (maximum error) is used as a reference and
every time a new approximation is computed, if the change in paramaters is compared to this
reference error. If it is small, `maxabs(da) < eps*maxerr`, the algorithm converged.

## Return values

A tuple containing:

 * The estimated coefficients
 * A `Bool` stating whether the algorithm converged
 * Number of iterations it took to converge.

"""
function nonlinear_fit(x, fun, a0, eps=1e-8, maxiter=200)

    na = length(a0)
    np = size(x, 1)
    nv = size(x, 2)
    for i in 1:na
        if a0[i] == 0
            a0[i] = 0.01
        end
    end

    da = a0 / 10
    a1 = zeros(na)
    a = zeros(na)
    for k = 1:na
        a1[k] = a0[k] - da[k]
    end
    xp = zeros(nv)
    A = zeros(np, na)
    r0 = zeros(np)
    r1 = zeros(np)
    for p = 1:np
        for k = 1:nv
            xp[k] = x[p,k]
        end
        r0[p] = fun(xp, a0)
        r1[p] = fun(xp, a1)
    end
    maxerr = maximum(abs, [maximum(abs, r0), maximum(abs, r1)])
    iter = 1
    convergence = false
    for i = 1:maxiter
        iter = i
        for p = 1:np
            for k = 1:nv
                xp[k] = x[p,k]
            end
            for k = 1:na
                aa = a1[k]
                a1[k] = a1[k] + da[k]
                r = fun(xp, a1)
                a1[k] = aa

                A[p,k] = -(r1[p] - r) / da[k]
            end
        end
        da = A \ r1
        for k = 1:na
            a1[k] = a1[k] - da[k]
        end
        for p = 1:np
            r0[p] = r1[p]
            for k = 1:nv
                xp[k] = x[p,k]
            end
            r1[p] = fun(xp, a1)
        end
        if maximum(abs, r1) < eps * maxerr
            convergence = true
            break
        end
    end

    return a1, convergence, iter
end

"""
   a = gauss_newton_fit(x, y, fun, ∇fun!, a0[[, eps,] maxiter])

Gauss-Newton nonlinear least squares. Given vectors `x` and `y`, the tries to fit parameters `a` to
a function `f` using least squares approximation:

 ``y = f(x, a₁, ..., aₙ)``

For more general approximations, see [`gauss_newton_fit`](@ref).

### Arguments:

 * `x` Vector with x values
 * `y` Vector with y values
 * `fun` a function that is called as `fun(x, a)` where `a` is a vector of parameters.
 * `∇fun!` A function that calculares the derivatives with respect to parameters `a`
 * `a0` Vector with the initial guesses of the parameters
 * `eps` Maximum residual for convergence
 * `maxiter` Maximum number of iterations for convergence

## Return value

A vector with the convrged array. If no convergence is achieved, the function throws an error.

## Specification of the fitting function

The function that should be fitted shoud be specified by Julia funcion with the following signature:

```julia
fun(x::T, a::AbstractVector{T}) where {T<:Number}
```

The derivatives with respect to each fitting parameter `a[i]` should have the following signature:

```julia
∇fun!(x::T, a::AbstractVector{T}, df::AbstractVector{T}) where {T<:Number}
```

No return value is expected and the derivatives are returned in argument `df`.

## Initial approximation (guess)

If the initial approximation is not good enough, divergence is possible.

**Careful** with parameters close to 0. The initial guess should never be 0.0 because the initial
value of the parameter is used as reference value for computing resiuduals.

## Convergence criteria

The argumento `maxiter` specifies the maximum number of iterations that should be carried out.
At each iteration,

``aₖⁿ⁺¹ = aₖⁿ + δₖ``

Convergence is achieved when

``|δᵏ / aₖ⁰| < ε``

## Example
```julia
x = 1.0:10.0
a = [3.0, 2.0, 1.0]
y = a[1] + a[2]*x + a[3]*x^2
fun(x, a) = a[1] + a[2]*x + a[3]*x^2

function ∇fun!(x, a, df)
    df[1] = 1.0
    df[2] = x
    df[3] = x^2
end

a = gauss_newton_fit(x, y, fun, ∇fun!, [0.5, 0.5, 0.5], 1e-8, 30)
```
"""
function gauss_newton_fit(x::AbstractVector{T}, y::AbstractVector{T}, fun, ∇fun!, a0::AbstractVector{T},
                          eps=1e-8, maxiter=200) where {T<:Number}
    P = length(x) # Number of points
    N = length(a0) # Number of parameters

    xi = zero(T)
    df = zeros(T, N)
    a = zeros(T, N)
    for i in 1:N
        a[i] = a0[i]
        if a[i] == 0
            a[i] = 0.01
        end
    end

    A = zeros(T, N, N)
    b = zeros(T, N)

    δref = abs.(a)
    maxerr = zero(T)
    for iter = 1:maxiter
        A .= zero(T)
        b .= zero(T)
        for i in 1:P
            xi = x[i]
            yi = y[i]
            f = fun(xi, a) - yi
            ∇fun!(xi, a, df)

            # Assemble LHS
            for k in 1:N
                for j in 1:N
                    A[j,k] += df[j] * df[k]
                end
            end
            # Assemble RHS
            for j in 1:N
                b[j] -= f * df[j]
            end
        end

        δ = A\b
        a .+= δ

        # Verify convergence:
        maxerr = maximum(abs, δ./δref)
        if maxerr < eps
            return(a)
        end

    end

    error("gauss_newton_fit failed to converge in $maxiter iterations with relative residual of $maxerr !")

    return(a)
end

"""
   a = gauss_newton_generic_fit(x, y, fun, ∇fun!, a0[[, eps,] maxiter])

Gauss-Newton nonlinear least squares. Given matrix `x` the function tries to fit parameters `a` to
an implicit function `f` using least squares approximation:

``f(x₁,..., xₘ, a₁, ..., aₙ) = 0``

This function is very similar to [`gauss_newton_fit`](@ref) but more generic. It doesn't limit the number of
independent variables and doesn't require a `y` LHS.

### Arguments:

 * `x` Matrix where each column is a variable
 * `
 * `fun` a function that is called as `fun(x, a)` where `a` is a vector of parameters.
 * `∇fun!` A function that calculares the derivatives with respect to parameters `a`
 * `a0` Vector with the initial guesses of the parameters
 * `eps` Maximum residual for convergence
 * `maxiter` Maximum number of iterations for convergence

## Return value

A vector with the convrged array. If no convergence is achieved, the function throws an error.

## Specification of the fitting function

The function that should be fitted shoud be specified by Julia funcion with the following signature:

```julia
fun(x::T, a::AbstractVector{T}) where {T<:Number}
```

The derivatives with respect to each fitting parameter `a[i]` should have the following signature:

```julia
∇fun!(x::T, a::AbstractVector{T}, df::AbstractVector{T}) where {T<:Number}
```

No return value is expected and the derivatives are returned in argument `df`.

## Initial approximation (guess)

If the initial approximation is not good enough, divergence is possible.

**Careful** with parameters close to 0. The initial guess should never be 0.0 because the initial
value of the parameter is used as reference value for computing resiuduals.

## Convergence criteria

The argumento `maxiter` specifies the maximum number of iterations that should be carried out.
At each iteration,

``aₖⁿ⁺¹ = aₖⁿ + δₖ``

Convergence is achieved when

``|δᵏ / aₖ⁰| < ε``

## Example
```julia
x = 1.0:10.0
a = [3.0, 2.0, 1.0]
y = a[1] + a[2]*x + a[3]*x^2
fun(x, a) = a[1] + a[2]*x + a[3]*x^2

function ∇fun!(x, a, df)
    df[1] = 1.0
    df[2] = x
    df[3] = x^2
end

a = gauss_newton_fit(x, y, fun, ∇fun!, [0.5, 0.5, 0.5], 1e-8, 30)
```

"""
function gauss_newton_generic_fit(x::AbstractMatrix{T}, fun, ∇fun!, a0::AbstractVector{T},
                          eps=1e-8, maxiter=200) where {T<:Number}
    P = size(x, 1) # Number of points
    M = size(x, 2) # Number of columns (variables)
    N = length(a0) # Number of parameters

    xi = zeros(T, M)
    df = zeros(T, N)
    a = zeros(T, N)
    for i in 1:N
        a[i] = a0[i]
        if a[i] == 0
            a[i] = 0.01
        end
    end

    A = zeros(T, N, N)
    b = zeros(T, N)

    δref = abs.(a)
    maxerr = zero(T)
    for iter = 1:maxiter
        A .= zero(T)
        b .= zero(T)
        for i in 1:P
            for k in 1:M
                xi[k] = x[i,k]
            end
            f = fun(xi, a)

            ∇fun!(xi, a, df)
            # Assemble LHS
            for k in 1:N
                for j in 1:N
                    A[j,k] += df[j] * df[k]
                end
            end
            # Assemble RHS
            for j in 1:N
                b[j] -= f * df[j]
            end
        end

        δ = A\b
        a .+= δ

        # Verify convergence:
        maxerr = maximum(abs, δ./δref)
        if maxerr < eps
            return(a)
        end

    end

    error("gauss_newton_fit failed to converge in $maxiter iterations with relative residual of $maxerr !")

    return(a)
end

"""
   a = secant_nls_fit(x, y, fun, ∇fun!, a0[[, eps,] maxiter])

Secant/Gauss-Newton nonlinear least squares. DOESN'T NEED A DERIVATIVE FUNCTION. Given vectors `x` and `y`, the tries to fit parameters `a` to
a function `f` using least squares approximation:

 ``y = f(x, a₁, ..., aₙ)``

For more general approximations, see [`gauss_newton_fit`](@ref).

### Arguments:

 * `x` Vector with x values
 * `y` Vector with y values
 * `fun` a function that is called as `fun(x, a)` where `a` is a vector of parameters.
 * `∇fun!` A function that calculares the derivatives with respect to parameters `a`
 * `a0` Vector with the initial guesses of the parameters
 * `eps` Maximum residual for convergence
 * `maxiter` Maximum number of iterations for convergence

## Return value

A vector with the convrged array. If no convergence is achieved, the function throws an error.

## Specification of the fitting function

The function that should be fitted shoud be specified by Julia funcion with the following signature:

```julia
fun(x::T, a::AbstractVector{T}) where {T<:Number}
```

The derivatives with respect to each fitting parameter `a[i]` should have the following signature:

```julia
∇fun!(x::T, a::AbstractVector{T}, df::AbstractVector{T}) where {T<:Number}
```

No return value is expected and the derivatives are returned in argument `df`.

## Initial approximation (guess)

If the initial approximation is not good enough, divergence is possible.

**Careful** with parameters close to 0. The initial guess should never be 0.0 because the initial
value of the parameter is used as reference value for computing resiuduals.

## Convergence criteria

The argumento `maxiter` specifies the maximum number of iterations that should be carried out.
At each iteration,

``aₖⁿ⁺¹ = aₖⁿ + δₖ``

Convergence is achieved when

``|δᵏ / aₖ⁰| < ε``

## Example
```julia
x = 1.0:10.0
a = [3.0, 2.0, 1.0]
y = a[1] + a[2]*x + a[3]*x^2
fun(x, a) = a[1] + a[2]*x + a[3]*x^2

a = secant_nls_fit(x, y, fun, ∇fun!, [0.5, 0.5, 0.5], 1e-8, 30)
```
"""

function secant_nls_fit(x::AbstractVector{T}, y::AbstractVector{T}, fun, aguess::AbstractVector{T},
                          eps=1e-8, maxiter=200) where {T<:Number}
    P = length(x) # Number of points
    N = length(aguess) # Number of parameters

    xi = zero(T)
    df = zeros(T, N)
    a = zeros(T, N)
    for i in 1:N
        a[i] = aguess[i]
        if a[i] == 0
            a[i] = 0.01
        end
    end

    δ = a .* (one(T)/20)
    f1 = zeros(T, P)
    a .+= δ
    A = zeros(T, N, N)
    b = zeros(T, N)

    δref = abs.(a)
    maxerr = zero(T)
    for iter = 1:maxiter

        A .= zero(T)
        b .= zero(T)
        f1 .= fun.(x, Ref(a))
        for i in 1:P
            xi = x[i]
            yi = y[i]
            f = f1[i] - yi
            for k in 1:N
                a[k] -= δ[k]
                df[k] = (f1[i] - fun(xi, a)) / δ[k]
                a[k] += δ[k]
            end
            # Assemble LHS
            for k in 1:N
                for j in 1:N
                    A[j,k] += df[j] * df[k]
                end
            end
            # Assemble RHS
            for j in 1:N
                b[j] -= f * df[j]
            end
        end
        δ = A\b
        a .+= δ
        # Verify convergence:
        maxerr = maximum(abs, δ./δref)
        if maxerr < eps
            return(a)
        end

    end

    error("gauss_newton_fit failed to converge in $maxiter iterations with relative residual of $maxerr !")

    return(a)
end


# Generate a random graph with n vertices where each edge is included with probability p.
function erdos_renyi_graph(g::GT, n::Integer, p::Real; has_self_loops=false) where {GT<:AbstractGraph}
    for i=1:n
        start_ind = is_directed(g) ? 1 : i
        for j=start_ind:n
            if(rand() <= p && (i != j || has_self_loops))
                add_edge!(g, i, j)
            end
        end
    end
    return g
end

# Convenience function with a default graph type.
function erdos_renyi_graph(n::Integer, p::Real; is_directed=true, has_self_loops=false)
    g = simple_inclist(n, is_directed=is_directed)
    erdos_renyi_graph(g, n, p, has_self_loops=has_self_loops)
end

# Generate a 'small world' random graph based on the Watts-Strogatz model.
# Written with much reference to the implementation from GraphStream <http://graphstream-project.org>.
# The resulting graph has n vertices,
#   Each vertex has a base degree of k  (n > k, k >= 2, k must be even.)
#   There is a beta chance of each edge being 'rewired'
function watts_strogatz_graph(g::GT, n::Integer, k::Integer, beta::Real) where {GT<:AbstractGraph}
    # When attributes are restored, we can enable placing the vertices around a circle.
    # This is nice for explicit visualization of the model.
    #space = linspace(0,2*pi,n+1)
    #for v in vertices(g)
    #    attrs = attributes(v)
    #    x = round((cos(space[i]) + 1)*100,2)
    #    y = round((sin(space[i]) + 1)*100,2)
    #    attrs["pos"] = "$(x),$(y)"
    #end

    # Link each node to the k/2 nodes next to it in each direction.
    # For each vertex in the graph:
    for i in 1:n
        # For the next k/2 higher indexed nodes:
        for j in 1:div(k,2)
            # Either make an edge to the neighbor or a long edge across the graph.
            if rand() > beta
                add_edge!(g, i, ((i+j - 1) % n) + 1 )
            else
                # Add a random link across the graph and delete the link to the neighbor.
                while true
                    target = rand(1:(n-1))
                    if (target >= i)
                        target += 1
                    end
                    if !(target in out_neighbors(i,g))
                        add_edge!(g,i, target)
                        break
                    end
                end
            end
        end
    end
    g
end

# Convenience function with a default graph type.
function watts_strogatz_graph(n::Integer, k::Integer, beta::Real)
    g = simple_inclist(n, is_directed=false)
    watts_strogatz_graph(g,n,k,beta)
end


using Test, Polyhedra

@testset "./src/LocalPolytope/facets.jl" begin

using BellScenario

@testset "facets()" begin
    @testset "3-simplex normalized" begin
        vertices = [[1,0],[0,1],[0,0]]
        facet_dict = LocalPolytope.facets(vertices)

        @test facet_dict["equalities"] == Vector{Vector{Int64}}[]
        @test facet_dict["facets"] == [[-1,0,0],[0,-1,0],[1,1,1]]
        @test length(keys(facet_dict)) == 2
    end

    @testset "3-simplex generalized" begin
        vertices = [[1,0,0],[0,1,0],[0,0,1]]
        facet_dict = LocalPolytope.facets(vertices)

        @test facet_dict["equalities"] == [[1,1,1,1]]
        @test facet_dict["facets"] == [[0,-1,0,0],[0,0,-1,0],[0,1,1,1]]
        @test length(keys(facet_dict)) == 2
    end

    @testset "CH inequalities" begin
        scenario = BipartiteNonSignaling(2,2,2,2)
        vertices = LocalPolytope.vertices(scenario)
        facet_dict = LocalPolytope.facets(vertices)

        @test facet_dict["equalities"] == Vector{Vector{Int64}}[]
        @test facet_dict["facets"] == [
            [0, 0, 0, 0, -1, 0, 0, 0, 0], [0, 0, 0, 0, 0, -1, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, -1, 0, 0], [0, 0, 0, 0, 0, 0, 0, -1, 0],
            [-1, 0, 0, 0, 0, 1, 0, 0, 0], [-1, 0, 0, 0, 1, 0, 0, 0, 0],
            [0, -1, 0, 0, 0, 0, 0, 1, 0], [0, -1, 0, 0, 0, 0, 1, 0, 0],
            [0, 0, -1, 0, 0, 0, 1, 0, 0], [0, 0, -1, 0, 1, 0, 0, 0, 0],
            [0, 0, 0, -1, 0, 0, 0, 1, 0], [0, 0, 0, -1, 0, 1, 0, 0, 0],
            [-1, 0, -1, 0, 1, 1, 1, -1, 0], [-1, 0, 0, -1, 1, 1, -1, 1, 0],
            [0, -1, -1, 0, 1, -1, 1, 1, 0], [0, -1, 0, -1, -1, 1, 1, 1, 0],
            [0, 1, 0, 1, 0, 0, 0, -1, 1], [0, 1, 1, 0, 0, 0, -1, 0, 1],
            [1, 0, 0, 1, 0, -1, 0, 0, 1], [1, 0, 1, 0, -1, 0, 0, 0, 1],
            [0, 1, 0, 1, 1, -1, -1, -1, 1], [0, 1, 1, 0, -1, 1, -1, -1, 1],
            [1, 0, 0, 1, -1, -1, 1, -1, 1], [1, 0, 1, 0, -1, -1, -1, 1, 1]
        ]
        @test length(keys(facet_dict)) == 2
    end
end

@testset "facets(::Polyhedron)" begin
    scenario = LocalSignaling(3,3,2)
    local_poly = LocalPolytope.vrep(scenario)

    @test !hrepiscomputed(local_poly)

    facets = LocalPolytope.facets(local_poly)
    @test length(facets) == 15

    @test hrepiscomputed(local_poly)

    match_vertices = LocalPolytope.vertices(scenario)
    match_facets = LocalPolytope.facets(match_vertices)["facets"]

    @test facets == match_facets
end

end


@testset "Information Measures" begin

    s = res"ARNDCQEGHILKMFPSTWYV-"
    r = reverse(s)
    g = res"GGGGGGGGGGGGGGGGGGGGG"
    Pg = probabilities(g)
    Pgg = probabilities(g,g)
    Pggg = probabilities(g,g,g)
    Ps = probabilities(s)
    Pss = probabilities(s,s)
    Psr = probabilities(s,r)
    Psss = probabilities(s,s,s)
    count_random = count(Residue[], Residue[], pseudocounts=AdditiveSmoothing(1.0))
    prob_random = probabilities(Residue[], Residue[], pseudocounts=AdditiveSmoothing(1.0))

    @testset "Entropy" begin

        @testset "H(X)" begin

            @test entropy(Pg) ≈ 0.0
            @test entropy(Ps) ≈ log(20.0)
            @test entropy(Ps, 2) ≈ log(2, 20.0)
        end

        @testset "Joint Entropy: H(X,Y)" begin

            @test entropy(Pgg) ≈ 0.0
            @test entropy(Pss) ≈ log(20.0)
            @test entropy(Psr) ≈ log(19.0)
            @test entropy(Pss, 2) ≈ log(2, 20.0)
            @test entropy(count_random) ≈ log(400.0)
            @test entropy(prob_random) ≈ log(400.0)
        end

        @testset "Joint Entropy: H(X,Y,Z)" begin

            @test entropy(Pggg) ≈ 0.0
            @test entropy(Psss) ≈ log(20)
        end

        @testset "Using counts" begin

            @test entropy(count(g,g)) ≈ entropy(Pgg)
            @test entropy(count(s,s)) ≈ entropy(Pss)
            @test entropy(count(s,r)) ≈ entropy(Psr)
        end
    end

    @testset "Kullback-Leibler" begin

        @test kullback_leibler(Ps, [1.0/20.0 for i in 1:20]) ≈ 0.0
        @test kullback_leibler(Ps, Ps) ≈ 0.0
        @test kullback_leibler(Ps, BLOSUM62_Pi) ≈ kullback_leibler(Ps)
        @test kullback_leibler(Ps, BLOSUM62_Pi) ≈ mapreduce(i->0.05*log(0.05/BLOSUM62_Pi[i]),+,1:20)
        @test kullback_leibler(Psr, Psr) ≈ 0.0
    end

    @testset "Mutual Information" begin

        @test mutual_information(Pgg) ≈ 0.0

        P = gettablearray(prob_random)
        Q = getmarginalsarray(prob_random)[:,1] * getmarginalsarray(prob_random)[:,2]' # 0.002500000000000001
        @test mutual_information(prob_random) ≈ sum(P .* log.(P./Q)) # So, it's != 0.0
        @test mutual_information(count_random) ≈ 0.0 # It's more accurate

        Pgs = probabilities(g,s)
        P = gettablearray(Pgs)
        Q = getmarginalsarray(Pgs)[:,1] * getmarginalsarray(Pgs)[:,2]' # 0.05000000000000002
        @test mutual_information(Pgs) ≈ mapreduce(x -> isnan(x) ? 0.0 : x, +, P.*log.(P./Q))
        @test mutual_information(count(g,s)) ≈ 0.0 # It's more accurate

        # MI(X,Y)=H(X)+H(Y)-H(X,Y)

        @test mutual_information(Psr) ≈ (marginal_entropy(Psr,1) +
                                         marginal_entropy(Psr,2) - entropy(Psr))

        @test mutual_information(Pss) ≈ (marginal_entropy(Pss,1) +
                                         marginal_entropy(Pss,2) - entropy(Pss))

        @test mutual_information(Psr,2) ≈ (marginal_entropy(Psr,1,2) +
                                           marginal_entropy(Psr,2,2) - entropy(Psr,2))

        @test mutual_information(Pss,20) ≈ (marginal_entropy(Pss,1,20) +
                                            marginal_entropy(Pss,2,20) - entropy(Pss,20))

        @test isapprox(mutual_information(prob_random), 0.0, atol=1e-15)
        @test mutual_information(count_random) ≈ 0.0
        @test isapprox(mutual_information(count_random), marginal_entropy(count_random,1) +
                       marginal_entropy(count_random,2) - entropy(count_random), atol=1e-13)

        @testset "Using counts" begin

            @test mutual_information(count_random) ≈ 0.0
            @test mutual_information(count(g,g)) ≈ mutual_information(Pgg)
            @test mutual_information(count(s,s)) ≈ mutual_information(Pss)
            @test mutual_information(count(s,r)) ≈ mutual_information(Psr)
        end

        @testset "MI(X,Y,Z)" begin

            @test mutual_information(Pggg) ≈ 0.0
            @test mutual_information(probabilities(g,s,r)) ≈ 0.0
            @test mutual_information(probabilities(s,s,r)) ≈ mutual_information(count(s,s,r))

            # MI(X,Y,Z) = H(X) + H(Y) + H(Z) - H(X,Y) - H(X,Z) - H(Y,Z) + H(X,Y,Z)
            @test mutual_information(Psss) ≈ mutual_information(count(s,s,s))
            @test mutual_information(Psss) ≈ (marginal_entropy(Psss, 1) +
                                              marginal_entropy(Psss, 2) +
                                              marginal_entropy(Psss, 3) -
                                              entropy(Pss) - entropy(Pss) -
                                              entropy(Pss) + entropy(Psss))

            # MI(X,Y,Z) <= min{ H(X,Y), H(X,Z), H(Y,Z) }
            @test mutual_information(Psss) ≈ mutual_information(Pss)
        end

        @testset "Pairwise Gap Percentage" begin

            @test gap_union_percentage(count(res"AA--", res"--AA", alphabet=GappedAlphabet())) ≈ 100.0
            @test gap_intersection_percentage(count(res"AA--", res"--AA", alphabet=GappedAlphabet())) ≈ 0.0

            @test gap_union_percentage(count(res"AA--", res"--AA", alphabet=GappedAlphabet(), weights=Weights([.25, .25, .25, .25]))) ≈ 100.0
            @test gap_intersection_percentage(count(res"AA--", res"--AA", alphabet=GappedAlphabet(), weights=Weights([.25, .25, .25, .25]))) ≈ 0.0

            @test gap_union_percentage(count(res"AAA-", res"AA--", alphabet=GappedAlphabet())) ≈ 50.0
            @test gap_intersection_percentage(count(res"AAA-", res"AA--", alphabet=GappedAlphabet())) ≈ 25.0

            @test gap_union_percentage(count(res"AAA-", res"AA--", alphabet=GappedAlphabet(), weights=Weights([.2, .2, .2, .4]))) ≈ 60.0
            @test gap_intersection_percentage(count(res"AAA-", res"AA--", alphabet=GappedAlphabet(), weights=Weights([.2, .2, .2, .4]))) ≈ 40.0
        end
    end
end


module problem117

description = """
Red, green, and blue tiles
Using a combination of black square tiles and oblong tiles chosen from:
 red tiles measuring two units,
 green tiles measuring three units,
 and blue tiles measuring four units,
 it is possible to tile a row measuring five units in length in exactly fifteen different ways.

How many ways can a row measuring fifty units in length be tiled?"""

# Fill function, based on problems 114, 115, 116
function F(n :: Integer, memo = zeros(Int, n))
  # Base case: no space remaining in the row
  if n < 2 return 0 end

  # Recursive case, count 1 for each placement of block at positions 1:n-m+1, and all ways of
  # filling the remaining space at the end of the row after each placement, memoizing results.
  if memo[n] == 0
    total = 0
    for m=2:4, blockpos = 1:(n-m+1)
      total += 1 + F(n - blockpos - m + 1, memo)
    end
    memo[n] = total
  end
  return memo[n]
end

function mixed_fill_count(n :: Integer)
  return 1 + F(n) # +1 for empty row being a valid solution
end

using Base.Test

@test mixed_fill_count(5) == 15

end


#All these examples of path are in 4 stages

using Distributions, Random
rng = MersenneTwister(01012019);

"""
	gaussian_path1D()

Returns a '4x1' dimensional array of Gaussian random walk
"""
function gaussian_path1D()
    return vcat(0.0, cumsum(randn(rng, 3, 1), dims = 1)) #4 stages
end

"""
	gaussian_path2D()

Returns a '4x2' dimensional array of Gaussian random walk
"""
function gaussian_path2D()
    gsmatrix = randn(rng, 4, 2) * [1.0 0.0 ; 0.9 0.3] #will create an (dimension x nstages) matrix
    gsmatrix[1,:] .= 0.0
    return cumsum(gsmatrix .+ [1.0 0.0], dims = 1)
end

"""
	running_maximum1D()

Returns a '4x1' dimensional array of Running Maximum process.
"""
function running_maximum1D()
    rmatrix = vcat(0.0, cumsum(randn(rng, 3, 1), dims = 1))
    for i = 2 : 4
        rmatrix[i] = max.(rmatrix[i-1], rmatrix[i])
    end
    return rmatrix
end

"""
	running_maximum2D()

Returns a '4x2' dimensional array of Running Maximum process.
"""
function running_maximum2D()
    rmatrix = vcat(0.0, cumsum(randn(rng, 3, 1), dims = 1))
    rmatrix2D = zeros(4, 2)
    rmatrix2D[:,1] .= vec(rmatrix)
    for j = 2 : 2
        for i = 2 : 4
            rmatrix2D[i,j] = max.(rmatrix[i-1], rmatrix[i])
        end
    end
    return rmatrix2D * [1.0 0.0; 0.9 0.3]
end

"""
	path()

Returns a sample of stock prices following the a simple random random process.
"""
function path()
    return  100 .+ 50 * vcat(0.0, cumsum(randn(rng, 3, 1), dims = 1))
end


# In this exercise we will try to predict the quality class of wines given some chemical characteristics

# In detail, the attributes of this dataset are:
#   1) Alcohol
#   2) Malic acid
#   3) Ash
#   4) Alcalinity of ash
#   5) Magnesium
#   6) Total phenols
#   7) Flavanoids
#   8) Nonflavanoid phenols
#   9) Proanthocyanins
#   10) Color intensity
#   11) Hue
#   12) OD280/OD315 of diluted wines
#   13) Proline

# Further information concerning this dataset can be found online on the [UCI Machine Learning Repository dedicated page](https://archive.ics.uci.edu/ml/datasets/wine) or in particular on [this file](https://archive.ics.uci.edu/ml/machine-learning-databases/wine/wine.names)

# Our prediction concerns the quality class of the wine (1, 2 or 3) that is given as first column of the data.

# 1) Start by setting the working directory to the directory of this file and activate it. If you have the provided `Manifest.toml` file in the directory, just run `Pkg.instantiate()`, otherwise manually add the packages Pipe, HTTP, Plots and BetaML.
# Also, seed the random seed with the integer `123`.

# 2) Load the packages/modules DelimitedFiles, Pipe, HTTP, Plots, BetaML

# 3) Load from internet or from local file the input data as a Matrix.
# You can use `readdlm`` using the comma as field separator.
dataURL="https://archive.ics.uci.edu/ml/machine-learning-databases/wine/wine.data"

# 4) Now create the X matrix of features using the second to final columns of the data you loaded above and the Y vector by taking the 1st column. Transform the Y vector to a vector of integers using the `Int()` function (broadcasted). Make shure you have a 178×13 matrix and a 178 elements vector

# 5) Partition the data in (xtrain,xval) and (ytrain,yval) keeping 80% of the data for training and reserving 35% for testing. Keep the default option to shuffle the data, as the input data isn't.

# 6) As the output is multinomial we need to encode ytrain. We use the `oneHotEncoder()` function to make `ytrain_oh`

# 7) Define a Neural Network model with the following characteristics:
#   - 3 dense layers with respectively 13, 20 and 3 nodes and activation function relu
#   - a `VectorFunctionLayer` with 3 nodes and `softmax` as activation function
#   - `crossEntropy` as the neural network cost function

# 8) Train your model using `ytrain` and a scaled version of `xtrain`` (where all columns have zero mean and 1 standard deviaiton) for 100 epochs and use a batch size of 6 records.
# Save the output of your training function to `trainingLogs`

# 9) Predict the training labels ŷtrain and the validation labels ŷval. Recall you did the training on the scaled features!

# 10) Compute the train and test accuracies using the function `accuracy`

# 11) Compute and print a ConfutionMatrix of the validation data true vs. predicted

# 12) Run the following commands to plots the average loss per epoch
plot(trainingLogs.ϵ_epochs)

# 13) (Optional) Run the same workflow without scaling the data or using `squaredCost` as cost function. How this affect the quality of your predictions ?


export color, chromatic_number

"""
`color(G,k)`: Return a `k`-coloring of `G` (or error if none exists).
"""
function color(G::SimpleGraph, k::Int)
    VV = vlist(G)
    EE = elist(G)
    n = NV(G)
    m = NE(G)

    if k<1
        error("Number of colors must be positive")
    end

    err_msg = "This graph is not " * string(k) * " colorable"

    # Special cases that don't require integer programming

    # see if we know the chromatic number already
    if cache_check(G,:chromatic_number)
        chi = cache_recall(G,:chromatic_number)
        if k<chi
            error(err_msg)
        end
    end

    result = Dict{vertex_type(G),Int}()

    if k==1
        if NE(G) > 0
            error(err_msg)
        end
        for v in VV
            result[v] = 1
        end
        return result
    end

    if k==2
        return two_color(G)
    end

    MOD = Model(get_solver())

    @variable(MOD, x[VV,1:k], Bin)

    for v in VV
        @constraint(MOD, sum(x[v,i] for i=1:k) == 1)
    end

    for e in EE
        u = e[1]
        v = e[2]
        for i=1:k
            @constraint(MOD, x[u,i] + x[v,i] <= 1)
        end
    end

    optimize!(MOD)
    status = Int(termination_status(MOD))

    if status != 1
        error(err_msg)
    end

    X = value.(x)

    for v in VV
        for c = 1:k
            if X[v,c] > 0
                result[v] = c
            end
        end
    end

    return result
end

function chromatic_number(G::SimpleGraph, verb::Bool=false)::Int
    if cache_check(G,:chromatic_number)
        return cache_recall(G,:chromatic_number)
    end

    if NV(G) == 0
        return 0
    end
    if NE(G) == 0
        return 1
    end

    # lower bound: larger of clique size or n/alpha
    n = NV(G)
    alf = length(max_indep_set(G))
    lb1 = Int(floor(n/alf))
    lb2 = length(max_clique(G))
    lb = max(lb1,lb2)

    # upper bound: try a random greedy coloring
    f = greedy_color(G)
    ub = maximum(values(f))

    k =  chromatic_number_work(G,lb,ub,verb)
    cache_save(G,:chromatic_number,k)
    return k
end

function chromatic_number_work(G::SimpleGraph, lb::Int, ub::Int, verb::Bool)::Int
    if verb
        print("$lb <= chi(G) <= $ub")
    end
    if lb == ub
        if verb
            println("\tand we're done")
        end
        return lb
    end

    mid = Int(floor((ub+lb)/2))

    if verb
        print("\tlooking for a $mid coloring")
    end

    try
        f = color(G,mid)  # success
        if verb
            println("\tfound")
        end
        return chromatic_number_work(G,lb,mid,verb)
    catch
    end
    if verb
        println("\tno such coloring")
    end
    return chromatic_number_work(G,mid+1,ub,verb)
end


@model function gmm_tarray(x, K)
    N = length(x)

    # Cluster centers.
    μ ~ filldist(Normal(), K)

    # Cluster association prior.
    w ~ Dirichlet(K, 1.0)

    # Cluster assignments.
    z = tzeros(Int, N)
    for n = 1:N
        z[n] ~ DiscreteNonParametric(1:K, w)
    end

    # Observations.
    for n = 1:N
        x[n] ~ Normal(μ[z[n]], 1.0)
    end
end

gmm_tarray_example(x = [0.1, -0.05, 1.0], K = 2) = gmm_tarray(x, K)

# K clusters, each one around i for i = 1:K with variance 0.5
@model function hmm_tarray(x, K, ::Type{T}=Float64) where {T<:Real}
    N = length(x)

    # State sequence.
    s = tzeros(Int, N)

    # Emission matrix.
    m = Vector{T}(undef, K)

    # Transition matrix.
    T = Vector{Vector{T}}(undef, K)

    # Assign distributions to each element
    # of the transition matrix and the
    # emission matrix.
    for i = 1:K
        T[i] ~ Dirichlet(K, 1.0)
        m[i] ~ Normal(i, 0.5)
    end

    # Observe each point of the input.
    s[1] ~ Categorical(K)
    x[1] ~ Normal(m[s[1]], 0.1)

    for i = 2:N
        s[i] ~ Categorical(T[s[i-1]])
        x[i] ~ Normal(m[s[i]], 0.1)
    end
end

hmm_tarray_example(x = [0.1, -0.05, 1.0], K = 2) = hmm_tarray(x, K)

function stickbreak(v)
    K = length(v) + 1
    cumprod_one_minus_v = cumprod(1 .- v)

    return map(1:K) do k
        if k == 1
            v[1]
        elseif k == K
            cumprod_one_minus_v[K - 1]
        else
            v[k] * cumprod_one_minus_v[k - 1]
        end
    end
end

@model function imm_stick_tarray(y, α, K)
    N = length(y)
    crm = DirichletProcess(α)

    v ~ filldist(StickBreakingProcess(crm), K - 1)
    w = stickbreak(v)

    # Cluster assignments
    z = tzeros(Int, N)
    for n = 1:N
        z[n] ~ Categorical(w)
    end

    # Cluster centers
    L = identity(K)
    μ ~ filldist(Normal(), L)

    # Observations
    for n = 1:N
        y[n] ~ Normal(μ[z[n]], 1.0)
    end
end

imm_stick_tarray_example(y = data_neal, α = α_neal, K = 10) = imm_stick_tarray()


using LinearAlgebra
using SpecialFunctions
using OffsetArrays

const Γ = gamma

# Coefficients in the Jacobi polynomial recurrence relations.
a(α, β, n) = 2/(2n+α+β) * √(n * (n+α+β) * (n+α) * (n+β) / (2n+α+β-1) / (2n+α+β+1))
b(α, β, n) = -(α^2 - β^2) / (2n+α+β) / (2n+α+β+2)

" Evaluates the Jacobi polynomial Pₙᵅᵝ(x). "
function jacobi(x, α, β, n::Int)
    Pᵅᵝ = n <= 1 ? OffsetArray(zeros(2), 0:1) : OffsetArray(zeros(n+1), 0:n)
    Pᵅᵝ[0] = √(2.0^-(α+β+1) * Γ(α+β+2) / Γ(α+1) / Γ(β+1))
    Pᵅᵝ[1] = Pᵅᵝ[0]/2 * √((α+β+3) / (α+1) / (β+1)) * ((α+β+2)*x + α - β)
    for n′ in 1:n-1
        Pᵅᵝ[n′+1] = ((x - b(α,β,n′)) * Pᵅᵝ[n′] - a(α,β,n′) * Pᵅᵝ[n′-1]) / a(α, β, n′+1)
    end
    return Pᵅᵝ[n]
end

" Evaluates the derivative of the Jacobi polynomial Pₙᵅᵝ(x). "
∂jacobi(x, α, β, n::Int) = n == 0 ? 0 : √(n * (n+α+β+1)) * jacobi(x, α+1, β+1, n-1)

" Guassian quadrature points and weights for the Jacobi polynomial Pₙᵅᵝ. "
function jacobi_gauss_quadrature(α, β, N)
    N == 0 && return [(α-β) / (α+β+2)], [2]

    # Form symmetric matrix from recurrence.
    dv = OffsetArray(zeros(N+1), 0:N)  # diagonal vector
    ev = OffsetArray(zeros(N+1), 0:N)  # sub/super-diagonal vector

    for n in 0:N
        dv[n] = b(α, β, n)
        ev[n] = a(α, β, n)
    end

    # Create full matrix combining the two.
    # Need to pass arrays that are not offset.
    J = SymTridiagonal(dv[0:N], ev[1:N])
    (α + β) ≈ 0 && (J[1, 1] = 0)

    # Compute quadrature points and weights by eigenvalue solve.
    x, V = eigen(J)
    w = @. V[1, :]^2 * 2^(α+β+1) / (α+β+1)
    @. w *= factorial(α) * factorial(β) / factorial(α+β)

    return x, w
end

" Guass-Labatto quadrature points for the Jacobi Polynomial Pₙᵅᵝ. "
function jacobi_gauss_lobatto(α, β, N)
    N == 0 && error("What are you doing? Gauss-Lobatto points only make sense if N >= 1.")
    N == 1 && return [-1, 1]

    x = zeros(N+1)
    x[1], x[N+1] = -1, 1

    x_GQ, _ = jacobi_gauss_quadrature(α+1, β+1, N-2)
    x[2:N] .= x_GQ

    return x
end


function bs(O::Option)
    d1 = (log(O.s / O.k) + (O.r + O.σ * O.σ / 2) * O.t) / (O.σ * √O.t)
    d2 = d1 - O.σ * √O.t

    if O.call_or_put == "Put"
        return CND(-d2) * O.k * exp(-O.r * O.t) - CND(-d1) * O.s
    else
        return O.s * CND(d1) - exp(-O.r * O.t) * O.k * CND(d2)
    end
end

function CND(x::Float64)
    a1 = 0.31938153
    a2 = -0.356563782
    a3 = 1.781477937
    a4 = -1.821255978
    a5 = 1.330274429
    l = abs(x)
    k = 1 / (1 + 0.2316419 * l)

    CND = 1 - 1 / sqrt(2 * π) * exp(-l^2 / 2) * (a1 * k + a2 * k^2 + a3 * k^3 + a4 * k^4 + a5 * k^5)

    if x < 0
        return 1 - CND
    end

    return CND
end


### 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

"""


function predict_svm(θ,x)
	return dot(θ,x) > 0 ? 1 : -1
end

function predict_logreg(theta,x)
	prediction =  1 / ( 1 + exp(dot(-theta,x)))
	if prediction  < 0.5
		return -1
	else
		return 1
	end
end

function evaluate(X,y,theta,mode,verbose=false)
	positiveHits = 0
	for i in 1:size(X,1)
		if mode == "logreg"
			label = predict_logreg(theta,X[i:i,:])
		elseif mode == "svm"
			label = predict_svm(theta,X[i:i,:])
		else
			error("Prediction mode not aviable")
		end
		if label == y[i]
			positiveHits += 1
		end
		if label == 1
			label = "+1"
		end
		labelString = "-1.0"
		if y[i] == 1.0
			labelString = "+1.0"
		end
		if verbose
			@printf("Prediction : %s | Label : %s | Accuracy %f \n",label,labelString,positiveHits / i)
		end
	end
	return positiveHits / size(X,1)
end


# 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)