function isValid(input::String)
    checkNine(n) = n > 9 ? n - 9 : n

    input = filter(x -> x != ' ' ,collect(input))

    length(input) <= 1 && return false

    parsedArr = try
        map(x -> parse(Int,x), input)
    catch
        return false
    end

    flag = true
    result = 0
    for i in reverse(parsedArr)
        if flag
            result += i
        else
            result += checkNine(2 * i)
        end
        flag = !flag
    end

    0 == (result % 10)
end


using Pkg
Pkg.activate(pwd())

# # Neural networks
# ## Prepare data

using RDatasets

iris = dataset("datasets", "iris")
X = Matrix(iris[:, 1:4])
y = iris.Species

#### Exercise:
#Write the `split` function, which randomly splits the dataset and the labels into training
#and testing sets. Its input should be the dataset `X` and the labels `y`. It should have
#four outputs. Include 80% of data in the training set and 20% of data in the testing set
# by default.
#
#**Hints:**
#- Use the `randperm` function from the `Random` package.
#- While `y` can be assumed to a vector, `X` is a matrix or a more-dimensional array. Then
#it is beneficial to use the `selectdim` function to select subindices along the correct
#dimension.
#
#---
#### Solution:

using Random

function split(X, y::AbstractVector; dims=1, ratio_train=0.8, kwargs...)
    n = length(y)
    size(X, dims) == n || throw(DimensionMismatch("..."))

    n_train = round(Int, ratio_train*n)
    i_rand = randperm(n)
    i_train = i_rand[1:n_train]
    i_test = i_rand[n_train+1:end]

    return selectdim(X, dims, i_train), y[i_train], selectdim(X, dims, i_test), y[i_test]
end

#+

X_train, y_train, X_test, y_test = split(X, y)

# ---
#
# ### Exercise:
# Write the `normalize` functions as described above. It should have two inputs and two
# outputs. The keyword argument `dims` should also be included.
#
# **Hint**: check the help for the `mean` function.
#
# ---
# ### Solution:

using Statistics

function normalize(X_train, X_test; dims=1, kwargs...)
    col_mean = mean(X_train; dims)
    col_std = std(X_train; dims)

    return (X_train .- col_mean) ./ col_std, (X_test .- col_mean) ./ col_std
end

#+

X_train, X_test = normalize(X_train, X_test)

# ---
#
# ### Exercise:
# Write the `onehot` function that converts the labels `y` into their one-hot
# representation. The samples should be along the second dimension. Write the `onecold`
# function that converts the one-hot representation into the one-cold (original)
# representation. Both these functions need to have two arguments; the second one is
# `classes`, which equals `unique(y)`.
#
# Write a check that both functions work correctly.
#
# **Hints:**
# - The one-hot representation for a label has the size equalling to the number of classes.
# All entries besides one are zeros.
# - Since the one-hot representation represents probabilities, the prediction is the class
# with the highest probability.
#
# ---
# ### Solution:

function onehot(y, classes)
    y_onehot = falses(length(classes), length(y))
    for (i, class) in enumerate(classes)
        y_onehot[i, y .== class] .= 1
    end
    return y_onehot
end

onecold(y, classes) = [classes[argmax(y_col)] for y_col in eachcol(y)]

#+

classes = unique(y)
isequal(onecold(onehot(y, classes), classes), y)

# ---

using LinearAlgebra

function prepare_data(X, y; do_normal=true, do_onehot=true, kwargs...)
    X_train, y_train, X_test, y_test = split(X, y; kwargs...)

    if do_normal
        X_train, X_test = normalize(X_train, X_test; kwargs...)
    end

    classes = unique(y)

    if do_onehot
        y_train = onehot(y_train, classes)
        y_test = onehot(y_test, classes)
    end

    return X_train, y_train, X_test, y_test, classes
end

#+

Random.seed!(666)

iris = dataset("datasets", "iris")
X = Matrix(iris[:, 1:4])
y = iris.Species

X_train, y_train, X_test, y_test, classes = prepare_data(X', y; dims=2)

#+

Random.seed!(666)
aux1 = prepare_data(X, y; dims=1)

#+

Random.seed!(666)
aux2 = prepare_data(X', y; dims=2)

#+

norm(aux1[1] - aux2[1]')

# ## Create the network

struct SimpleNet{T<:Real}
    W1::Matrix{T}
    b1::Vector{T}
    W2::Matrix{T}
    b2::Vector{T}
end

# ### Exercise:
# Write an outer constructor for `SimpleNet`. Its inputs should be three integers
# representing the input size of the three layers. All matrices should be initialized based
# on the normal distribution.
#
# **Hint**: think about the representation of the dense layer.
#
# ---
# ### Solution:

SimpleNet(n1, n2, n3) = SimpleNet(randn(n2, n1), randn(n2), randn(n3, n2), randn(n3))

# ---

Random.seed!(666)
m = SimpleNet(size(X_train,1), 5, size(y_train,1))

# ### Exercise:
# Write a functor `function (m::SimpleNet)(x)` which computes the prediction (forward pass)
# of the neural network `SimpleNet`.
#
# **Bonus**: try to make the functor work for both vectors (one sample) and matrices
# (multiple samples) `x`.
#
# ---
# ### Solution:

function (m::SimpleNet)(x)
    z1 = m.W1*x .+ m.b1
    a1 = max.(z1, 0)
    z2 = m.W2*a1 .+ m.b2
    return exp.(z2) ./ sum(exp.(z2), dims=1)
end

# ---

m(X_train[:,1:2])

# ## Train the network

function grad(m::SimpleNet, x::AbstractVector, y; ϵ=1e-10)
    z1 = m.W1*x .+ m.b1
    a1 = max.(z1, 0)
    z2 = m.W2*a1 .+ m.b2
    a2 = exp.(z2) ./ sum(exp.(z2), dims=1)
    l = -sum(y .* log.(a2 .+ ϵ))

    e_z2 = exp.(z2)
    l_part = (- e_z2 * e_z2' + Diagonal(e_z2 .* sum(e_z2))) / sum(e_z2)^2

    l_a2 = - y ./ (a2 .+ ϵ)
    l_z2 = l_part * l_a2
    l_a1 = m.W2' * l_z2
    l_z1 = l_a1 .* (a1 .> 0)
    l_x = m.W1' * l_z1

    l_W2 = l_z2 * a1'
    l_b2 = l_z2
    l_W1 = l_z1 * x'
    l_b1 = l_z1

    return l, l_W1, l_b1, l_W2, l_b2
end

#+

g_all = [grad(m, X_train[:,k], y_train[:,k]) for k in 1:size(X_train,2)]
typeof(g_all)

#+

function mean_tuple(d::AbstractArray{<:Tuple})
    Tuple([mean([d[k][i] for k in 1:length(d)]) for i in 1:length(d[1])])
end

#+

g_mean = mean_tuple(g_all)
typeof(g_mean)

# ### Exercise:
# Train the network with a gradient descent with stepsize ``\alpha=0.1`` for ``200``
# iterations. Save the objective value at each iteration and plot the results.
#
# ---
# ### Solution:

α = 1e-1
max_iter = 200
L = zeros(max_iter)

for iter in 1:max_iter
    grad_all = [grad(m, X_train[:,k], y_train[:,k]) for k in 1:size(X_train,2)]
    grad_mean = mean_tuple(grad_all)

    L[iter] = grad_mean[1]

    m.W1 .-= α*grad_mean[2]
    m.b1 .-= α*grad_mean[3]
    m.W2 .-= α*grad_mean[4]
    m.b2 .-= α*grad_mean[5]
end

# ---

using Plots

plot(L;
    xlabel="Iteration",
    ylabel="Loss function",
    label="",
    title="Loss function on the training set"
)

# ## Prediction
# ### Exercise:
#
# Write a function which predict the labels for samples. Show the accuracy on both training # and testing sets.
#
# ---
# ### Solution:

predict(X) = m(X)
accuracy(X, y) = mean(onecold(predict(X), classes) .== onecold(y, classes))

println("Train accuracy = ", accuracy(X_train, y_train))
println("Test accuracy = ", accuracy(X_test, y_test))


# NOTE: we now have an AbstractODEProblem
abstract type AbstractODEProblem end

struct ODEProblem{F,T<:Tuple{Number,Number},U<:AbstractVector,P<:AbstractVector} <: AbstractODEProblem
    f::F
    tspan::T
    u0::U
    θ::P
end

abstract type ODESolver end

struct Euler{T} <: ODESolver
    dt::T
end

function (solver::Euler)(prob::AbstractODEProblem, u, t)
    f, θ, dt  = prob.f, prob.θ, solver.dt
    (u + dt*f(u,θ), t+dt)
end
struct RK2{T} <: ODESolver
    dt::T
end
function (solver::RK2)(prob::AbstractODEProblem, u, t)
    f, θ, dt  = prob.f, prob.θ, solver.dt
    du = f(u,θ)
    uh = u + du*dt
    u + dt/2*(du + f(uh,θ)), t+dt
end

function solve(prob::AbstractODEProblem, solver::ODESolver)
    t = prob.tspan[1]; u = prob.u0
    us = [u]; ts = [t]
    while t < prob.tspan[2]
        (u,t) = solver(prob, u, t)
        push!(us,u)
        push!(ts,t)
    end
    # NOTE: we are not `reduce(hcat,us)`ing any more
    ts, us
end


"""
polyinterp(t,y)

Return a callable polynomial interpolant through the points in
vectors `t`,`y`. Uses the barycentric interpolation formula.
"""
function polyinterp(t,y)
    n = length(t)-1
    C = (t[end]-t[1]) / 4       # scaling factor to ensure stability
    tc = t/C

    # Adding one node at a time, compute inverses of the weights.
    omega = ones(n+1)
    for m = 1:n
        d = tc[1:m] .- tc[m+1]      # vector of node differences
        omega[1:m] = omega[1:m].*d  # update previous
        omega[m+1] = prod( -d )     # compute the new one
    end
    w = 1 ./ omega                  # go from inverses to weights

    p = function (x)
        # Compute interpolant.
        terms = @. w / (x - t)
        if any(isinf.(terms))       # divided by zero here
            # Apply L'Hopital's Rule exactly.
            idx = findfirst(x.==t)
            f = y[idx]
        else
            f = sum(y.*terms) / sum(terms)
        end
    end
    return p
end

"""
triginterp(t,y)

Return trigonometric interpolant for points defined by vectors `t`
and `y`.
"""
function triginterp(t,y)
    N = length(t)

    function trigcardinal(x)
        if isodd(N)      # odd
            tau = sin(N*pi*x/2) / (N*sin(pi*x/2))
        else             # even
            tau = sin(N*pi*x/2) / (N*tan(pi*x/2))
        end
        if isnan(tau)
            tau = 1
        end
        return tau
    end

    p = function (x)
        sum( y[k]*trigcardinal(x-t[k]) for k=1:N )
    end
    return p
end

"""
ccint(f,n)

Perform Clenshaw-Curtis integration for the function `f` on `n`+1
nodes in [-1,1]. Return integral and a vector of the nodes used.
Note: `n` must be even.
"""
function ccint(f,n)
    # Find Chebyshev extreme nodes.
    theta = [ i*pi/n for i=0:n ]
    x = -cos.(theta)

    # Compute the C-C weights.
    c = zeros(n+1)
    c[[1,n+1]] .= 1/(n^2-1)
    v = 1 .- 2*sum( cos.(2*k*theta[2:n])/(4*k^2-1) for k=1:n/2-1 )
    v -= cos.(n*theta[2:n])/(n^2-1)
    c[2:n] = 2*v/n

    # Evaluate integrand and integral.
    I = dot(c,f.(x))   # use vector inner product
    return I,x
end

"""
glint(f,n)

Perform Gauss-Legendre integration for the function `f` on `n` nodes
in (-1,1). Return integral and a vector of the nodes used.
"""
function glint(f,n)
    # Nodes and weights are found via a tridiagonal eigenvalue problem.
    beta = @. 0.5/sqrt(1-(2*(1:n-1))^(-2))
    T = diagm(-1=>beta,1=>beta)
    lambda,V = eigen(T)
    p = sortperm(lambda)
    x = lambda[p]           # nodes
    c = @. 2*V[1,p]^2       # weights

    # Evaluate the integrand and compute the integral.
    I = dot(c,f.(x))      # vector inner product
    return I,x
end

"""
intde(f,h,M)

Perform doubly-exponential integration of function `f` over
(-Inf,Inf), using discretization size `h` and truncation point `M`.
Return integral and a vector of the nodes used.
"""
function intde(f,h,M)
    # Find where to truncate the trapezoid sum.
    K = ceil( log(4/pi*log(2*M))/h )

    # Integrate by trapezoids in a transformed variable t.
    t = h*(-K:K)
    x = @. sinh(pi/2*sinh(t))
    dxdt = @. pi/2*cosh(t)*cosh(pi/2*sinh(t))

    I = h*dot(f.(x),dxdt)
    return I,x
end

"""
intsing(f,h,delta)

Integrate function `f` (possibly singular at 1 and -1) over
[-1+`delta`,1-`delta`] using discretization size `h`. Return
integral and a vector of the nodes used.
"""
function intsing(f,h,delta)
    # Find where to truncate the trapezoid sum.
    K = ceil(log(-2/pi*log(delta/2))/h)

    # Integrate over a transformed variable.
    t = h*(-K:K)
    x = @. tanh(pi/2*sinh(t))
    dxdt = @. pi/2*cosh(t) / (cosh(pi/2*sinh(t))^2)

    I = h*dot(f.(x),dxdt)
    return I,x
end


"""
**Nadam constructor (Nesterov Adam)**

```julia
     Nadam(;η::Real=0.001, β₁::Real=0.9, β₂::Real=0.999, ϵ::Real=10e-8)
```

Algorithm :

```math
\\begin{align*}
m_t =& \\beta_1 m_{t-1} + (1-\\beta_1)g_t\\\\
v_t =& \\beta_2 v_{t-1} + (1-\\beta_2)g_t^2\\\\
\\hat{m}_t =& \\frac{m_t}{1-\\beta_1^t}\\\\
\\hat{v}_t =& \\frac{v_t}{1-\\beta_2^t}\\\\
\\Delta x_t =&\\frac{\\eta}{\\sqrt{\\hat{v}_t+\\epsilon}}(\\beta_1\\hat{m}_t+\\frac{(1-\\beta_1)g_t}{1-\\beta_1^t})
\\end{align*}
```

[Algorithm Reference](http://cs229.stanford.edu/proj2015/054_report.pdf) and [other reference](http://www.cs.toronto.edu/~fritz/absps/momentum.pdf)
"""
mutable struct Nadam <: Optimizer
    opt_type::String
    t::Int64
    ϵ::Float64
    η::Float64
    β₁::Float64
    β₂::Float64
    m_t::AbstractArray
    v_t::AbstractArray
end

"Construct Nadam optimizer"
function Nadam(;η::Real=0.001, β₁::Real=0.9, β₂::Real=0.999, ϵ::Real=10e-8)
    @assert η > 0.0 "η must be greater than 0"
    @assert β₁ > 0.0 "β₁ must be greater than 0"
    @assert β₂ > 0.0 "β₂ must be greater than 0"
    @assert ϵ > 0.0 "ϵ must be greater than 0"

    Nadam("Nadam", 0, ϵ, η, β₁, β₂, [], [])
end

params(opt::Nadam) = "ϵ=$(opt.ϵ), η=$(opt.η), β₁=$(opt.β₁), β₂=$(opt.β₂)"

function update(opt::Nadam, g_t::AbstractArray{T}) where {T<:Real}
    # resize biased moment estimates if first iteration
    if opt.t == 0
        opt.m_t = zero(g_t)
        opt.v_t = zero(g_t)
    end

    # update timestep
    opt.t += 1

    # update biased first moment estimate
    opt.m_t = opt.β₁ * opt.m_t + (one(T) - opt.β₁) * g_t

    # update biased second raw moment estimate
    opt.v_t = opt.β₂ * opt.v_t + (one(T) - opt.β₂) * ((g_t) .^2)

    # compute bias corrected first moment estimate
    m̂_t = opt.m_t / (one(T) - opt.β₁^opt.t)

    # compute bias corrected second raw moment estimate
    v̂_t = opt.v_t / (one(T) - opt.β₂^opt.t)

    # apply update
    ρ = opt.η ./ (sqrt.(v̂_t .+ opt.ϵ))
    ρ .*= (opt.β₁ * m̂_t + (one(T) - opt.β₁) * g_t / (one(T) - opt.β₁^opt.t))

    return ρ
end


"""
    Hm... Power method should work well for symmatric tridiagonal matrix that is
    always changing its size... In fact it does, let's try inverse iterations
    instead.
"""

"""
    Left and right boundary can actually be infinite.
    Because of left and right bound, we have to restrict it to type Abstract Floats so that it's
    comparable, which is going to be used a lot for search the eigen values of the matrix.
"""
function InversePowerIterSearch(
        A::AbstractMatrix{T},
        left_bound::T,
        right_bound::T
    ) where {T<: AbstractFloat}
    # Get the mid point.
    @assert size(A, 1) == size(A, 2) "It has to be square matrix."
    @assert left_bound != right_bound "The interval must be non-degenerate."
    if left_bound > right_bound
        left_bound, right_bound = right_bound, left_bound
    end
    if left_bound == -Inf && right_bound == Inf
        error("Left and right bound cannot both be Inf, then it's unbounded. ")
    end
    v = randn(T, size(A, 1))    # previous iterations
    u = similar(v)                 # current iterations
    Au = similar(v)
    t = similar(v)              # temp
    MaxItr = 100
    while MaxItr > 0
        # Choose the right point.
        if left_bound == -Inf || right_bound == Inf
            θ = left_bound == -Inf ? 2*left_bound : 2*right_bound
            if θ == -Inf || θ == Inf
                error("Inverse Iteration Igen Search failed, bound limit reached for half opened interval.  ")
            end
        else
            θ = (left_bound + right_bound)/2
        end
        println("θ: $θ ∈ [$left_bound, $right_bound]")
        # Inverse Power iterations.
        λ = θ
        while λ == θ || norm(Au - λ*u, 1) > 1e-10
            println(λ)
            t .= u
            u .= (A - λ*I)\v
            u ./= norm(u)
            if norm(u) == NaN
                u .= v
                println("coverted to $λ")
                break
            else
                v .= t
                Au = A*u
                λ = dot(u, Au)
            end
        end
        # Interval shrink.
        if λ > left_bound && λ < right_bound

            return λ, u
        else
            if λ <= left_bound
                left_bound = θ

            else
                right_bound = θ
            end
        end
        MaxItr -= 1
    end
    if MaxItr == 0
       error("Inverse Power iterations method reached maximal iterations without convergence. ")
    end
end

# Basic Tests and shit
using LinearAlgebra
A = Diagonal(sort(randn(10)))
display(A)
λ, v = InversePowerIterSearch(A, A[1, 1], A[3, 3])
print("λ = $λ")
# λ, v = InversePowerIterSearch(A, -Inf, A[2, 2])
# print("λ = $λ")


"""
    fakedata_gaussian(pts_per_clust::Int = 10, nclust::Int = 2, xmin = 0.0, xmax = nclust * 7.0, nx = round(Int, (xmax - xmin) * 10), σ = 1.0, amin = 1.0, amax = 2.0)

DOCSTRING

# Arguments:
- `pts_per_clust`: DESCRIPTION
- `nclust`: DESCRIPTION
- `xmin`: DESCRIPTION
- `xmax`: DESCRIPTION
- `nx`: DESCRIPTION
- `σ`: DESCRIPTION
- `amin`: DESCRIPTION
- `amax`: DESCRIPTION
"""
function fakedata_gaussian(
    pts_per_clust::Int = 10,
    nclust::Int        = 2,
    xmin               = 0.0,
    xmax               = nclust * 7.0,
    nx                 = round(Int, (xmax - xmin) * 10),
    σ                  = 1.0,
    amin               = 1.0,
    amax               = 2.0,
)
    x = LinRange(xmin, xmax, nx)
    μ = LinRange(xmin, xmax, nclust + 2)[2:end-1]
    a = LinRange(xmin, xmax, pts_per_clust + 2)[2:end-1]

    npts = pts_per_clust * nclust
    data = zeros(nx, npts)
    labels = zeros(Int, npts)

    i = 1
    for c = 1:nclust
        for n = 1:pts_per_clust
            data[:, i] = gauss_func.(a[n], μ[c], σ, x)
            labels[i] = c
            i += 1
        end
    end

    return data, labels
end

gauss_func(a, μ, σ, x) = a * exp((-(x - μ) .^ 2) / (2 * σ^2))


# A version of the conjugate gradient method for linesearch.

"""A version of the conjugate gradient method for linesearch.
CG(A, b, ϵa, ϵr, itmax; quad) solves the symmetric linear system 'A * x = b'
If quad = true, the values of the quadratic model are computed
A can be positive definite or not.
"""
function CGlin(A, b, ϵa::Float64=1e-8, ϵr::Float64=1e-6, itmax::Int=0, ε::Float64=1e-6; quad::Bool=false)
    n = size(b, 1) # size of the problem
    (size(A, 1) == n & size(A, 2) == n) || error("Inconsistent problem size")
    @info(loggerCGlin, @sprintf("CGlin: system of %d equations in %d variables", n, n))

    x = zeros(n) # initial estimation x = 0
    xNorm = 0.0
    xNorms = [xNorm] # Values of ‖x‖
    r = -b # initial residual r = Ax-b = -b
    d = b # first descent direction
    rNorm = norm(r, 2)
    ϵ = ϵa + ϵr * rNorm

    iter = 0
    itmax == 0 && (itmax = 2 * n)

    if quad
        q = 0.0
        qvalues = [q] # values of the quadratic model
        @info(loggerCGlin, @sprintf("%5s %10s %10s\n", "Iter", "‖r‖", "q"))
        @info(loggerCGlin, @sprintf("    %d    %8.1e    %8.1e", iter, rNorm, q))
    end

    solved = rNorm ≤ ϵ
    tired = iter ≥ itmax

    while ! (solved || tired)
        iter += 1
        Ad = A * d
        dAd = dot(d, Ad)

        # if the model is not convexe, the algorithm stops
        if dAd ≤ ε * dot(d, d)
            @debug(loggerCGlin, @sprintf("non positive curvature dAd = %8.1e", dAd))
            iter == 1 && return b
            return x
        end

        α = rNorm^2 / dAd # step for x estimation
        x = x + α * d # new estimation
        xNorm = norm(x, 2)
        push!(xNorms, xNorm)
        roldNorm = rNorm
        r = r + α * Ad # new residual
        rNorm = norm(r, 2)

        if quad
            q = -dot(b, x) + 0.5 * dot(x, A * x)
            push!(qvalues, q)
            @info(loggerCGlin, @sprintf("    %d    %8.1e    %8.1e", iter, rNorm, q))
        end

        solved = rNorm ≤ ϵ
        tired = iter ≥ itmax

        (solved || tired) && continue
        β = rNorm^2 / roldNorm^2 # step for the next descent direction
        d = -r + β * d # new descent direction

        end
        return x
    end


@doc doc"""
Approximation de la solution du sous-problème ``q_k(s) = s^{t}g + (1/2)s^{t}Hs``
        avec ``s=-t g_k,t > 0,||s||< \delta_k ``

# Syntaxe
```julia
s1, e1 = Pas_De_Cauchy(gradient,Hessienne,delta)
```

# Entrées
 * **gradfk** : (Array{Float,1}) le gradient de la fonction f appliqué au point ``x_k``
 * **hessfk** : (Array{Float,2}) la Hessienne de la fonction f appliqué au point ``x_k``
 * **delta**  : (Float) le rayon de la région de confiance

# Sorties
 * **s** : (Array{Float,1}) une approximation de la  solution du sous-problème
 * **e** : (Integer) indice indiquant l'état de sortie:
        si g != 0
            si on ne sature pas la boule
              e <- 1
            sinon
              e <- -1
        sinon
            e <- 0

# Exemple d'appel
```julia
g1 = [0; 0]
H1 = [7 0 ; 0 2]
delta1 = 1
s1, e1 = Pas_De_Cauchy(g1,H1,delta1)
```
"""
function Pas_De_Cauchy(g,H,delta)

  a = (1/2)*g'*H*g
  b = norm(g)

  if (b > 0)
    t = (b^2)/(2*a)
    s = -t*g
    e = 1
    if norm(s)>delta
        t = delta/b
        s = -t*g
        e = -1
    end
  else
    t = 0
    s = -t*g
    e = 0
  end

  return s, e
end

2A/S7/UE Optimisation et RO/Optimisation/Optinum/src/Pas_De_Cauchy.jl


@doc raw"""
    rigid_body_velocity(vel_P, ω, r_PQ)

Calculate rigid solid velocity field.

Return velocity of a point Q of a rigid solid given the velocity of a point P (vel_P), the
rotational velocity of the solid (ω) and the relative position of Q with respect to P.

If the reference frame 1 is attached to the solid and the velocity is calculated with
respect to reference frame 0:

``v_{10}^{Q} = v_{10}^{P} + \omega_{10} \times r^{PQ}``

being:
- ``v_{10}^{Q}`` the velocity of point Q, fixed to 1, wrt 0
- ``\omega_{10}`` the angular velocity of the solid 1 wrt 0
- ``r^{PQ}`` the position of Q wrt P (``r^{Q}-r^{P}``)

Every vector needs to be expressed in the same coordinate system.

# References

1. Stevens, B. L., Lewis, F. L., (1992). Aircraft control and simulation: dynamics, controls
   design, and autonomous systems. John Wiley & Sons. (Section 1.3, page 26)
"""
function rigid_body_velocity(vel_P, ω, r_PQ)
    vel_Q = vel_P + ω × r_PQ
    return vel_Q
end

@doc raw"""
    rigid_body_acceleration(acc_P, ω, ω_dot, r_PQ)

Calcualte rigid body acceleration field.

Return the acceleration of a point Q of a rigid solid given the acceleration of a point P
(acc_P), the rotational velocity of the solid (ω), the rotational acceleration of the solid
(ω_dot) and the relative position of Q with respect to P.

``a_{10}^{Q} = a_{10}^{P} + \omega_{10} \times (\omega_{10} \times r^{PQ}) + \dot{\omega}_{10} \times r^{PQ}``

being:
- ``a_{10}^{Q}`` the acceleration of point Q, fixed to 1, wrt 0
- ``\omega_{10}`` the angular velocity of the solid 1 wrt 0
- ``\dot{\omega}_{10}`` the angular acceleration of the solid 1 wrt 0
- ``r^{PQ}`` the position of Q wrt P (``r^{Q}-r^{P}``)

# References

1. Stevens, B. L., Lewis, F. L., (1992). Aircraft control and simulation: dynamics, controls
   design, and autonomous systems. John Wiley & Sons. (Section 1.3, Formaula 1.3-14c, page 26)
"""
function rigid_body_acceleration(acc_P, ω, ω_dot, r_PQ)
    acc_Q = acc_P + ω × (ω × r_PQ) + ω_dot × r_PQ
    return acc_Q
end

"""
    pqr_2_ψθϕ_dot(p, q, r, ψ, θ, ϕ)

Transform body angular velocity (p, q, r) [rad/s] to Euler angles rates (ψ_dot, θ_dot,
ϕ_dot) [rad/s] given the euler angles (θ, ϕ) [rad] using kinematic angular equations.

# See also

[`ψθϕ_dot_2_pqr`](@ref), [`pqr_2_quat_dot`](@ref)

# References

1. Stevens, B. L., Lewis, F. L., & Johnson, E. N. (2015). Aircraft control and simulation:
   dynamics, controls design, and autonomous systems. John Wiley & Sons. Equation (1.4-4) (page 20)
"""
function pqr_2_ψθϕ_dot(p, q, r, θ, ϕ)

    sθ, cθ = sin(θ), cos(θ)
    sϕ, cϕ = sin(ϕ), cos(ϕ)

    ψ_dot = (q * sϕ + r * cϕ) / cθ
    θ_dot = q * cϕ - r * sϕ
    # ϕ_dot = p + (q * sϕ + r * cϕ) * tan(θ)
    ϕ_dot = p + ψ_dot * sθ

    return [ψ_dot, θ_dot, ϕ_dot]
end

"""
    ψθϕ_dot_2_pqr(ψ_dot, θ_dot, ϕ_dot, ψ, θ, ϕ)

Transform Euler angles rates (ψ_dot, θ_dot, ϕ_dot) [rad/s] to body angular velocity (p, q,
r) [rad/s] given the euler angles (θ, ϕ) [rad] using kinematic angular equations.

# See also

[`pqr_2_ψθϕ_dot`](@ref)

# References

1. Stevens, B. L., Lewis, F. L., & Johnson, E. N. (2015). Aircraft control and simulation:
   dynamics, controls design, and autonomous systems. John Wiley & Sons. Equation (1.4-3) (page 20)
"""
function ψθϕ_dot_2_pqr(ψ_dot, θ_dot, ϕ_dot, θ, ϕ)

    sθ, cθ = sin(θ), cos(θ)
    sϕ, cϕ = sin(ϕ), cos(ϕ)

    p = ϕ_dot - sθ * ψ_dot
    q = cϕ * θ_dot + sϕ*cθ * ψ_dot
    r = -sϕ * θ_dot + cϕ*cθ * ψ_dot

    return [p, q, r]
end

"""
    pqr_2_quat_dot(p, q, r, q0, q1, q2, q3)

Transform body angular velocity (p, q, r) [rad/s] to quaternion rates [1/s].

# See also

[`pqr_2_ψθϕ_dot`](@ref)

# References

1. Stevens, B. L., Lewis, F. L., & Johnson, E. N. (2015). Aircraft control and simulation:
   dynamics, controls design, and autonomous systems. John Wiley & Sons. Equation (1.8-15)
   (page 51).
"""
function pqr_2_quat_dot(p, q, r, q0, q1, q2, q3)

    Ω = [
        0 -p -q -r;
        p  0  r -q;
        q -r  0  p;
        r  q -p  0;
    ]
    q = [q0; q1; q2; q3]
    q_dot = 0.5 * Ω * q

    return q_dot
end

"""
    uvw_to_tasαβ(u, v, w)

Calculate true air speed (TAS), angle of attack (α) and angle of side slip (β) from velocity
expressed in body axis.

# Notes

This function assumes that u, v, w are the body components of the aerodynamic speed. This is
not true in genreal (wind speed different from zero), as u, v, w represent velocity with
respect to an inertial reference frame.

# References

1. Stevens, B. L., Lewis, F. L., & Johnson, E. N. (2015). Aircraft control and simulation:
   dynamics, controls design, and autonomous systems. John Wiley & Sons. Equation (2.3-6b)
   (page 78).
"""
function uvw_to_tasαβ(u, v, w)

    tas = sqrt(u*u + v*v + w*w)
    α = atan(w, u)
    β = asin(v / tas)
    return [tas, α, β]
end

"""
    uvw_dot_to_tasαβ_dot(u, v, w, u_dot, v_dot, w_dot)

Calculate time derivatives of velocity expressed as TAS, AOA, AOS.

# Notes
Note that tas here is not necessarily true air speed. Could also be inertial speed in the
direction of airspeed. It will concide whith TAS for no wind.

# See also

[`tasαβ_dot_to_uvw_dot`](@ref)

# References
1. Morelli, Eugene A., and Vladislav Klein. Aircraft system identification: Theory and
   practice. Williamsburg, VA: Sunflyte Enterprises, 2016. Equation 3.33 (page 44).
2. Stevens, B. L., Lewis, F. L., & Johnson, E. N. (2015). Aircraft control and simulation:
   dynamics, controls design, and autonomous systems. John Wiley & Sons. Equation (2.3-10)
   (page 81).
"""
function uvw_dot_to_tasαβ_dot(u, v, w, u_dot, v_dot, w_dot)

    # [1] 3.31
    tas = sqrt(u*u + v*v + w*w)
    # [1] 3.33a
    tas_dot = (u * u_dot + v * v_dot + w * w_dot) / tas
    # [1] 3.33b
    β_dot = ((u*u + w*w) * v_dot - v * (u * u_dot + w * w_dot)) / (tas^2 * sqrt(u*u + w*w))
    # [2] 2.3.10b
    # β_dot = (v_dot * tas - v * tas_dot) / (tas * sqrt(u*u + w*w))
    # [1] 3.33c
    α_dot = (w_dot * u - w * u_dot) / (u*u + w*w)

    return [tas_dot, α_dot, β_dot]
end

"""
    tasαβ_dot_to_uvw_dot(tas, α, β, tas_dot, α_dot, β_dot)

Obatain body velocity derivatives given velocity in wind axis and its derivatives.

# See also

[`uvw_dot_to_tasαβ_dot`](@ref)

# References

1. Morelli, Eugene A., and Vladislav Klein. Aircraft system identification: Theory and
   practice. Williamsburg, VA: Sunflyte Enterprises, 2016. Derived from equation 3.32
   (page 44).
"""
function tasαβ_dot_to_uvw_dot(tas, α, β, tas_dot, α_dot, β_dot)
    u_dot = tas_dot * cos(α) * cos(β) - tas * (α_dot * sin(α) * cos(β) + β_dot * cos(α) * sin(β))
    v_dot = tas_dot * sin(β) + tas * β_dot * cos(β)
    w_dot = tas_dot * sin(α) * cos(β) + tas * (α_dot * cos(α) * cos(β) - β_dot * sin(α) * sin(β))
    return [u_dot, v_dot, w_dot]
end

@doc raw"""
    rate_of_climb_constrain_no_wind(γ, α, β, ϕ)

Calculate pitch angle (θ rad) to obtain a flight path angle (γ rad) at certain angle of
attack (α rad), angle of sideslip (β rad) and roll (ϕ rad).

Inertial velocity and aerodynamic velocity are equivalent in the absence of wind:

`` v_{cg-e}^e = R_{eb} R_{bw} v_{cg-air}^w ``

# References

1. Stevens, B. L., Lewis, F. L., (1992). Aircraft control and simulation: dynamics, controls
   design, and autonomous systems. John Wiley & Sons. (Section 3.6, equation 3.6-3,
   page 187)
"""
function rate_of_climb_constrain_no_wind(γ, α, β, ϕ)
    a = cos(α) * cos(β)
    b = sin(ϕ) * sin(β) + cos(ϕ) * sin(α) * cos(β)
    sq = sqrt(a^2 - sin(γ)^2 + b^2)
    θ = (a * b + sin(γ) * sq) / (a^2 - sin(γ)^2)
    θ = atan(θ)
    return θ
end


#=
  From https://math.stackexchange.com/questions/3633307/bayes-rule-broken
  """
  Bayes' Rule broken?!?!

  This question has been driving me CRAZY for 4 days now. The question comes from
  the textbook 'One Thousand Exercises in Probability', specifically Exercise 3 in
  section 1.4. The solution does not make sense! The question goes as follows:

     'A man possesses five coins, two double-headed, two normal and one double-tailed.
     The man shuts his eyes, picks a coin at random, and tosses the coin. He opens
     his eyes, sees a head: what is the probability the lower face is also a head?.

  The book gives an answer that is 2/3
  ....
  """

  Distributions of variable coin (num:0)
  1.00000 =>    6585  (0.658500)
  2.00000 =>    3415  (0.341500)

  Distributions of variable flip_is_head (num:0)
  1.00000 =>    6585  (0.658500)
  0.00000 =>    3415  (0.341500)

   Cf ~/blog/five_coins.blog
      ~/webppl/family_out_problem.wppl
=#

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

@model function five_coins()
    doubleHead = 1
    normalCoin = 2
    doubleTail = 3

    head = 1
    tail = 2

    coin ~ Categorical([2/5,2/5,1/5]) # [doubleHead,normalCoin,doubleTail]

    #  Toss the coin
    toss ~
        (coin == doubleHead) ? Dirac(head) :
        (coin == doubleTail) ? Dirac(tail) :
        Categorical([1/2,1/2]) # [head,tail]

    #  Flip the coin (i.e. turn it around)
    flip ~ Dirac((coin == doubleHead) ? head :
                 (coin == doubleTail) ? tail :
                 (toss == head) ? tail : head
                 )

    #  We observe that the toss give a head.
    true ~ Dirac(toss == head);

    #  What is the probability that the other side (flip side)
    #  is head?
    flip_is_head ~ Dirac(flip == head)
end

model = five_coins()

num_chains = 4

# chns = sample(model, Prior(), 10_000)
# chns = sample(model, MH(), 10_000)
chns = sample(model, PG(5), 10_000)
# chns = sample(model, SMC(), 10_000)
# chns = sample(model, IS(), 10_000)

display(chns)
# display(plot(chns))

show_var_dist_pct(chns,:coin)
show_var_dist_pct(chns,:flip_is_head)


"""
    fixedStateCV(
        states :: Vector{<:AbstractMatrix},
        kraus_ops :: Vector{<:AbstractMatrix}
    ) :: Tuple{Float64, Vector{Matrix{ComplexF64}}}

For a fixed ensemble of `states` and quantum channel described by `kraus_ops`,
the communication value (CV) is evaluated by maximizing over all POVM measurements.
This optimization is expressed in primal form as the semidefinite program:

```math
\\begin{matrix}
    \\text{fixedStateCV}(\\mathcal{N})&= \\max_{\\{\\Pi_x\\}_{x}} \\sum_x \\text{Tr}\\left[\\Pi_x\\mathcal{N}(\\rho_x)\\right] \\\\
    & \\\\
    & \\text{s.t.} \\quad \\sum_x\\Pi_x = \\mathbb{I} \\;\\; \\text{and} \\;\\; \\Pi_x \\geq 0
\\end{matrix}
```

where each state ``\\rho_x`` satisfies ``\\text{Tr}[\\rho_x] = 1`` and ``\\rho_x \\geq 0``.
The channel ``\\mathcal{N}`` is applied to each state as
``\\mathcal{N}(\\rho_x) = \\sum_j k \\rho_x k^\\dagger``.

# Returns

A `Tuple`, `(cv, opt_povm)` where `cv` is the evaluated communication value and
`opt_povm` is the optimal POVM measurement.
"""
function fixedStateCV(
    states :: Vector{<:AbstractMatrix},
    kraus_ops :: Vector{<:AbstractMatrix}
) :: Tuple{Float64, Vector{Matrix{ComplexF64}}}
    d = size(kraus_ops[1],1)
    n = length(states)

    # add povm variables and constraints
    povm_vars = map(i -> HermitianSemidefinite(d), 1:n)
    constraints = sum(map(Πy -> real(Πy), povm_vars)) == Matrix{Float64}(I, d, d)
    constraints += sum(map(Πy -> imag(Πy), povm_vars)) == zeros(Float64, d, d)

    # apply channel to states
    evolved_states = map(ρ ->  sum(k -> k * ρ * k' , kraus_ops), states)

    # maximize CV over POVM measurements
    objective = maximize(real(sum(tr.(povm_vars .* evolved_states))), constraints)
    qsolve!(objective)

    cv = objective.optval
    opt_povm = map(Πy -> Πy.value, povm_vars)

    return cv, opt_povm
end

"""
    fixedMeasurementCV(
        povm :: Vector{<:AbstractMatrix},
        kraus_ops :: Vector{<:AbstractMatrix}
    ) :: Tuple{Float64, Vector{Matrix{ComplexF64}}}

For a fixed `povm` measurement and quantum channel described by `kraus_ops`, the
communication value (CV) and optimal state encodings are computed.
The fixed measurement CV is evaluated as

```math
\\text{fixedMeasurementCV}(\\mathcal{N}) = \\sum_y ||\\mathcal{N}^{\\dagger}(\\Pi_y)||_{\\infty}
```

where ``||\\mathcal{N}^{\\dagger}(\\Pi_y)||_{\\infty}`` is the largest eigenvalue
of the POVM element ``\\Pi_y`` evolved by the adjoint channel,
``\\mathcal{N}^{\\dagger}(\\Pi_y) = \\sum_j k^{\\dagger}_j \\Pi_y k_j``.
The states which maximize the CV are simply the eigenvectors corresponding to the
largest eigenvalue of each respective POVM element.

# Returns

A `Tuple`, `(cv, opt_states)` where `cv` is the communication value and
`opt_states` is the set of optimal states.
"""
function fixedMeasurementCV(
    povm :: Vector{<:AbstractMatrix},
    kraus_ops :: Vector{<:AbstractMatrix}
) :: Tuple{Float64, Vector{Matrix{ComplexF64}}}
    d = size(kraus_ops[1],2)
    n = length(povm)

    # apply adjoint channel to POVM
    evolved_povm = map(Π ->  sum(k -> k' * Π * k , kraus_ops), povm)

    opt_states = Vector{Matrix{ComplexF64}}(undef, n)
    cv = 0
    for i in 1:n
        povm_el = evolved_povm[i]
        vals, vecs = eigen(povm_el)

        opt_states[i] = vecs[:,end] * vecs[:,end]'
        cv += real(vals[end])
    end

    return cv, opt_states
end

"""
    seesawCV(
        init_states :: Vector{<:AbstractMatrix},
        kraus_ops :: Vector{<:AbstractMatrix},
        num_steps :: Int64;
        verbose :: Bool = false
    )

Performs the see-saw optimization technique to maximize the communication
value (CV) of the channel described by `kraus_ops` over all states and measurements.
This iterative and biconvex optimization technique combines coordinate ascent
maximization with semidefinite programming.
The number of iterations is determined by `num_steps` where each iteration
consists of a two-step procedure:

1. The POVM measurement is optimized with respect to a fixed state ensemble
   using the [`fixedStateCV`](@ref) function.
2. The state ensemble is optimized with respect to a fixed povm state using
   the [`fixedMeasurementCV`](@ref) function.

This procedure is initialized with `init_states` and after many iterations, a
local maximum of the CV is found.
The `verbose` keyword argument can be used to print out the CV evaluated in each
step.

The see-saw method has shown success in similar encoding/decoding
optimization problems in quantum information, *e.g.*,
[https://arxiv.org/abs/quant-ph/0307138v2](https://arxiv.org/abs/quant-ph/0307138v2)
and [https://arxiv.org/abs/quant-ph/0606078v1](https://arxiv.org/abs/quant-ph/0606078v1).
We note that our implementation is quite distinct from previous works, however,
the core iterative approach remains the same.

# Returns

A `Tuple` containing the following data in order:

1. `max_cv_tuple :: Tuple`, `(max_cv, opt_states, opt_povm)` A 3-tuple containing the maximal
   communication value and the optimal states/POVM that achieve this value.
2. `cvs :: Vector{Float64}`, A list of each evaluated CV. Since states and measurements
   are optimized in each iteration, we have `length(cvs) == 2 * num_steps`.
3. `opt_ensembles :: Vector{Vector{Matrix{ComplexF64}}}`, A list of state ensembles
   optimized in each step, where `length(opt_ensembles) == num_steps`.
4. `opt_povms :: Vector{Vector{Matrix{ComplexF64}}}`, A list of POVM measurements
   optimized in each step, where `length(opt_povms) == num_steps`.

!!! warning "Optimum Not Guaranteed"
    This function is not guaranteed to find a global or local optima. However,
    `seesawCV` will always provide a lower bound on the communication value.
"""
function seesawCV(
    init_states :: Vector{<:AbstractMatrix},
    kraus_ops :: Vector{<:AbstractMatrix},
    num_steps :: Int64;
    verbose :: Bool = false
) :: Tuple
    opt_ensembles = Vector{Vector{Matrix{ComplexF64}}}(undef, num_steps)
    opt_povms = Vector{Vector{Matrix{ComplexF64}}}(undef, num_steps)
    cvs = zeros(Float64, 2*num_steps)
    cv_id = 1

    max_cv_tuple = (1, [], [])

    for i in 1:num_steps
        # maximizing CV over POVMs
        opt_states = (i == 1) ? init_states : opt_ensembles[i-1]

        fixed_state_cv, opt_povm = fixedStateCV(opt_states, kraus_ops)

        cvs[cv_id] = fixed_state_cv
        cv_id += 1
        opt_povms[i] = opt_povm

        if fixed_state_cv > max_cv_tuple[1]
            max_cv_tuple = (fixed_state_cv, opt_states, opt_povm)
        end

        # maximizing CV over states
        fixed_povm_cv, opt_states = fixedMeasurementCV(opt_povm, kraus_ops)

        cvs[cv_id] = fixed_povm_cv
        cv_id += 1
        opt_ensembles[i] = opt_states

        if fixed_povm_cv > max_cv_tuple[1]
            max_cv_tuple = (fixed_povm_cv, opt_states, opt_povm)
        end

        if verbose
            println("i = ", i)
            println("fixed_state_cv = ", fixed_state_cv)
            println("fixed_povm_cv = ", fixed_povm_cv)
            println("max_cv = ", max_cv_tuple[1])
        end
    end

    return max_cv_tuple, cvs, opt_ensembles, opt_povms
end


#=
Section Allocation

   Suppose you have $n$ students in a class who need to be assigned to
      $m$ discussion sections. Each student needs to be assigned to exactly
      one section. Each discussion section should have between 6 and 10
      students. Suppose an $n \times m$ preference matrix $P$ is given,
      where $P_{ij}$ gives student $i$'s ranking for section $j$
      (1 would mean it is the student's top choice, 10,000 or a large number
      would mean the student can not attend that section).

   The goal will be to get an allocation matrix $X$, where $X_{ij} = 1$
      if student $i$ is assigned to section $j$ and $0$ otherwise.

  source: https://github.com/JuliaOpt/Convex.jl/blob/master/examples/
                             section_allocation/section_allocation.ipynb
=#

using Convex, GLPKMathProgInterface

include("data.jl"); # data.jl has our preference matrix, P

X = Variable(size(P), :Bin)

#=
   We want every student to be assigned to exactly one section.
   So, every row must have exactly one non-zero entry.
   In other words, the sum of all the columns for every row is 1.
   We also want each section to have between 6 and 10 students,
       so the sum of all the rows for every column should be between
       these constraints = [sum(X, 2) == 1, sum(X, 1) <= 10, sum(X, 1) >= 6]

   Our objective is simple sum(X .* P), which can be more efficiently
       represented as vec(X)' * vec(P).
   Since each entry of X is either 0 or 1, this is basically summing up
       the rankings of students that were assigned to them.
   If all students got their first choice, this value will be the number
       of students since the ranking of the first choice is 1.
=#

p = minimize(vec(X)' * vec(P), constraints)

solve!(p, GLPKSolverMIP())


####
#### Parametric gates
####

"""
    Rphi(ϕ)

The phase shift gate ``R_ϕ``.
This is equivalent to qiskit's `u1` gate.
"""
function Rphi(ϕ)
    return [1 0; 0 cis(ϕ)]
end

# Could use complex(reverse(sincospi(z))...)
"""
    _exp_ipi(z)

Return `exp(im * pi * z)`. This uses the accurate `cospi` and `sinpi` functions.
"""
_exp_ipi(z) = cispi(z)

"""
    Rphipi(z)

Return the ``R_ϕ`` gate for ``ϕ = zπ``. This is more accurate
than `Rphi(pi * z)`.
"""
function Rphipi(z)
    return [1 0; 0 _exp_ipi(z)]
end

"""
    U2(ϕ, λ)

u2 gate. Need to put math definitions in here.
"""
function U2(ϕ, λ)
    return [1 -cis(λ); cis(ϕ) cis(ϕ + λ)] / sqrt(2)
end

# This does not appar any faster than computing each separately
# function RXalt(θ)
#     (s, c) = sincos(θ/2)
#     si = - im * s
#     return [c si; si c]
# end

function RX(θ)
    c = cos(θ/2)
    s = - im * sin(θ/2)
    return [c s; s c]
end

function RXpi(z)
    c = cospi(z/2)
    s = -im * sinpi(z/2)
    return [c s; s c]
end

function RY(θ)
    c = cos(θ/2)
    s = sin(θ/2)
    return [c -s; s c]
end

function RYpi(z)
    c = cospi(z/2)
    s = sinpi(z/2)
    return [c -s; s c]
end

function RZ(θ)
    return [cis(-θ/2) 0; 0 cis(θ/2)]
end

function RZpi(z)
    return [_exp_ipi(-z/2) 0; 0 _exp_ipi(z/2)]
end

"""
    Ualt(θ, ϕ, λ)

Matrix from SU(2).
Alternative parameterization.
"""
function Ualt(θ, ϕ, λ)
    c = cos(θ/2)
    s = sin(θ/2)
    fpl =  (ϕ + λ) / 2
    cfpl = cos(fpl)
    sfpl = sin(fpl)
    fml =  (ϕ - λ) / 2
    cfml = cos(fml)
    sfml = sin(fml)
    f00 = complex(cfpl, -sfpl) * c
    f01 = complex(-cfml, sfml) * s
    f10 = complex(cfml, sfml) * s
    f11 = complex(cfpl, sfpl) * c
    return [f00 f01; f10 f11]
end

"""
    U(θ, ϕ, λ)

Matrix from SU(2). This is
the same as qiskit's U or U3.
"""
function U(θ, ϕ, λ)
    c = cos(θ/2)
    s = sin(θ/2)
    fpl =  (ϕ + λ)
    cfpl = cos(fpl)
    sfpl = sin(fpl)
    f00 = c
    f01 = -s * complex(cos(λ), sin(λ))
    f10 = s * complex(cos(ϕ), sin(ϕ))
    f11 = c * complex(cfpl, sfpl)
    return [f00 f01; f10 f11]
end

"""
    Ualtpi(θ, ϕ, λ)

Matrix from SU(2), with `θ`, `ϕ`, and `λ` given as multiples of `π`.
This is more accurate than `Ualt`.
"""
function Ualtpi(θ, ϕ, λ)
    c = cospi(θ/2)
    s = sinpi(θ/2)
    fpl =  (ϕ + λ) / 2
    fml =  (ϕ - λ) / 2
    cfpl = cospi(fpl)
    sfpl = sinpi(fpl)
    cfml = cospi(fml)
    sfml = sinpi(fml)
    f00 = complex(cfpl, -sfpl) * c
    f01 = complex(-cfml, sfml) * s
    f10 = complex(cfml, sfml) * s
    f11 = complex(cfpl, sfpl) * c
    return [f00 f01; f10 f11]
end

@doc raw"""
    R(θ, ϕ)

The R gate.
```math
R(\theta, \phi) = e^{-i \frac{\theta}{2} (\cos{\phi} x + \sin{\phi} y)}
```
"""
function R(θ, ϕ)
    c = cos(θ/2)
    s = sin(θ/2)
    return [c -im*exp(-im*ϕ)*s;
            -im*exp(im*ϕ)*s c]
end

@doc raw"""
    Rpi(θ, ϕ)

The R gate with θ and ϕ reduced by π. This is may be more accurate,
for example with integral and half-integral multiples of π.
```math
Rpi(\theta, \phi) = e^{-i \frac{\pi\theta}{2} (\cos{\pi\phi} x + \sin{\pi\phi} y)}
```
"""
function Rpi(θ, ϕ)
    c = cospi(θ/2)
    s = sinpi(θ/2)
    return [c -im*cispi(-ϕ)*s;
            -im*cispi(ϕ)*s c]
end

"""
    RXXYY(θ)

The R_XXYY gate, `exp(-im θ / 2 (X⊗X + Y⊗Y))`.
"""
function RXXYY(θ)
    _RXXYY(θ, cos, sin)
end

"""
    RXXYYpi(θ)

The R_XXYY gate with `θ` scaled by `π`. `exp(-im * pi * θ / 2 * (X⊗X + Y⊗Y))`
"""
function RXXYYpi(θ)
    _RXXYY(θ, cospi, sinpi)
end

function _RXXYY(θ, cfunc, sfunc)
    II = I2 ⊗ I2 / 2
    XX = X ⊗ X / 2
    YY = Y ⊗ Y / 2
    ZZ = Z ⊗ Z / 2
    return II + ZZ + cfunc(θ) * (II - ZZ) -im * sfunc(θ) * (XX + YY)
end


# # [Basic Epidemiology Models](@id epidemiology_example)
#
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/examples/epidemiology.ipynb)

using Petri
using LabelledArrays
using Plots
using DiffEqJump
using StochasticDiffEq
using OrdinaryDiffEq

# ### SIR Model
#
# The SIR model represents the epidemiological dynamics of an infectious disease
# that causes immunity in its victims. There are three *states:* `Suceptible
# ,Infected, Recovered`. These states interact through two *transitions*.
# Infection has the form `S+I -> 2I` where a susceptible person meets an
# infected person and results in two infected people. The second transition is
# recovery `I -> R` where an infected person recovers spontaneously.

S  = [:S,:I,:R]
Δ  = LVector(
       inf=(LVector(S=1, I=1), LVector(I=2)),
       rec=(LVector(I=1),      LVector(R=1)),
     )
sir = Petri.Model(S, Δ)

Graph(sir)

# Once a model is defined, we can define out initial parameters `u0`, a time
# span `tspan`, and the transition rates of the interactions `β`

u0 = LVector(S=990.0, I=10.0, R=0.0)
tspan = (0.0,40.0)
# add a dynamic transition rate for infection
# where the rate of infection decreases over time
# as is dependent on the current state of the system
β = LVector(inf=((u,t)->((3/sum(u))/(t+1))), rec=0.25);

# each transition rates can one of three options:
#
# - constant: `β = [.25]`
#   - where the rate is specified by a value of type `Number`
# - time dependent: `β = [t->((3/1000)/(t+1))]`
#   - where `t` is the current time step
# - state and time dependent: `β = [(u,t)->((3/sum(u))/(t+1))]`
#   - where `u` is the current state of `u0` and `t` is the current time step

# Petri.jl provides interfaces to StochasticDiffEq.jl, DiffEqJump.jl, and
# OrdinaryDiffEq.jl Here, we call the `JumpProblem` function that returns an
# DiffEqJump problem object that can be passed to the DiffEqJump solver which
# can then be plotted and visualized

prob = JumpProblem(sir, u0, tspan, β)
sol = DiffEqJump.solve(prob,SSAStepper())

plot(sol)

# Similarly, we can generated `SDEProblem` statements that can be used with
# StochasticDiffEq solvers

prob, cb = SDEProblem(sir, u0, tspan, β)
sol = StochasticDiffEq.solve(prob,LambaEM(),callback=cb)

plot(sol)

# Lastly, we can generated `ODEProblem` statements that can be used with
# OrdinOrdinaryDiffEq solvers

prob = ODEProblem(sir, u0, tspan, β)
sol = OrdinaryDiffEq.solve(prob,Tsit5(),reltol=1e-8,abstol=1e-8)

plot(sol)

# ### SEIR Model

S = [:S,:E,:I,:R]
Δ = LVector(
      exp=(LVector(S=1, I=1), LVector(I=1, E=1)),
      inf=(LVector(E=1),      LVector(I=1)),
      rec=(LVector(I=1),      LVector(R=1)),
    )
seir = Petri.Model(S, Δ)

Graph(seir)
#-
u0 = LVector(S=990.0, E=10.0, I=0.0, R=0.0)
tspan = (0.0,40.0)
β = LVector(exp=0.7/sum(u0), inf=0.5, rec=0.25)

prob, cb = SDEProblem(seir, u0, tspan, β)
sol = StochasticDiffEq.solve(prob,LambaEM(),callback=cb)

plot(sol)

# ### SEIRD Model

S = [:S,:E,:I,:R, :D]
Δ = LVector(
      exp=(LVector(S=1, I=1), LVector(I=1, E=1)),
      inf=(LVector(E=1),      LVector(I=1)),
      rec=(LVector(I=1),      LVector(R=1)),
      die=(LVector(I=1),      LVector(D=1)),
    )
seird = Petri.Model(S, Δ)

Graph(seird)
#-
u0 = LVector(S=990.0, E=10.0, I=0.0, R=0.0, D=0.0)
tspan = (0.0,40.0)
β = LVector(exp=0.9/sum(u0), inf=0.9, rec=0.25, die=0.03)

prob, cb = SDEProblem(seird, u0, tspan, β)
sol = StochasticDiffEq.solve(prob,LambaEM(),callback=cb)

plot(sol)


grid_size = 25
gamma_vals = linspace(0.05, 0.95, grid_size)
w_bar_vals = similar(gamma_vals)

mcm = McCallModel()

for (i, gamma) in enumerate(gamma_vals)
    mcm.gamma = gamma
    w_bar = compute_reservation_wage(mcm)
    w_bar_vals[i] = w_bar
end

plot(gamma_vals,
    w_bar_vals,
    lw=2,
    alpha=0.7,
    xlabel="job offer rate",
    ylabel="reservation wage",
    label=L"$\bar w$ as a function of $\gamma$")


module p3

# Return boundary conditions values and indices, ϕb and β, for problem 2,
# with N subdivisions
function boundaryconditions3(N::Integer)
    ϕb = zeros(N+1,N+1)                 # boundary conditions
    β  = zeros(Int,N+1,N+1)             # boundary condition locations

    # Conditions on the perimeter
    ϕb[:,1] = ϕb[:,N+1] = ϕb[1,:] = ϕb[N+1,:] = 0.0
    β[:,1]  = β[:,N+1]  = β[1,:]  = β[N+1,:]  = 1

    # Conditions on the interior
    # xmin < x < xmax  →  N*xmin + 1 < i < N*xmax + 1
    xmin = 0.5
    xmax = 0.75
    ymin = 0.625
    ymax = 0.875

    xs = int([ceil(N*xmin+1):floor(N*xmax+1)])
    ys = int([ceil(N*ymin+1):floor(N*ymax+1)])

    ϕb[xs,ys] = 0.0
    β[xs,ys]  = 1

    return ϕb, β
end

# Calculate the potential
function V(x::Real, y::Real, N::Integer, q::Real)

    a = b = 1:(N+1)
    a .-= (N*x + 0.5)           # avoid singularity
    b .-= (N*y + 0.5)
    a .*= a
    b .*= b
    return -N*q ./ √(a .+ b')

end

V(N::Integer, q::Real) = V(0.5, 0.25, N, q)

function minor(A::Matrix, β::Vector)

    size(A)[1] == size(A)[2] || error("Matrix must be square")

    N = sum(1 .- β)             # dimensions of new matrix

    B = Array(typeof(A[1]),N,N)

    ii = 1
    jj = 1

    for j in 1:size(A)[1]
        β[j] == 1 && continue
        for i in 1:size(A)[2]
            β[i] == 1 && continue
            B[ii,jj] = A[i,j]
            ii += 1
        end
        jj +=1
        ii = 1
    end

    return B
end

# repopulate a vector v with zeros for fixed parameters
function addwater(v::Vector, β::Vector)

    vf = Array(typeof(v[1]),length(β))
    ii = 1

    length(v) == sum(1-β) || error("v and β have different num free params")

    for i in 1:length(β)
        if β[i] == 1
            vf[i] = 0               # put zero for a fixed value
        elseif β[i] == 0
            vf[i] = v[ii]           # fill in the free value
            ii += 1                 # prep the next free value
        else
            error("β can only have 0 or 1 values")
        end
    end
    return vf

end

end


using Flux
using Statistics
using Plots
using DifferentialEquations

# Solve the ode given by $$ u' = \cos 2\pi t $$ by approximation of $u$ with a Neural Network
NNODE = Chain(x -> [x], # Take in a scalar and transform it into an array
              Dense(1,32,tanh),
              Dense(32,1),
              first) # Take first value, i.e. return a scalar
NNODE(1.0)

# construct function $g(t)$ to enforce initial condition $$ g(t) = t NNODE(t) + u(0)
u₀ = 1.0
g(t) = t*NNODE(t) + u₀

#Construct loss function to numericaly compute derivative and check against R.H.S.
ϵ = sqrt(eps(Float32))
loss() = mean(abs2(((g(t+ϵ)-g(t))/ϵ) - cos(2π*t)) for t in 0:1f-2:1f0)

opt = Flux.Descent(0.01)
data = Iterators.repeated((), 5000)
iter = 0
cb = function () #callback function to observe training
    global iter += 1
    if iter % 500 == 0
        display(loss())
    end
end
display(loss())
Flux.train!(loss, Flux.params(NNODE), data, opt; cb=cb)

# compare against true solution
t = 0:1e-3:1.0

plot(t, g.(t), label="NN")
plot!(t, u₀ .+ sin.(2π.*t)/2π, label="True Solution")


const AbstractVec = SVector{3, T} where T <: AbstractFloat
const Vec = AbstractVec{Float64}

function normalize(v::AbstractVec)::AbstractVec
    v / norm(v)
end

@inline function x(v::AbstractVec)::AbstractFloat
    v[1]
end

@inline function y(v::AbstractVec)::AbstractFloat
    v[2]
end

@inline function z(v::AbstractVec)::AbstractFloat
    v[3]
end

function random_on_unit_disk()::AbstractVec
    v::AbstractVec = Vec(randn(), randn(), 0.0)
    normalize(v)
end

function random_in_unit_disk()::AbstractVec
    v::AbstractVec = random_on_unit_disk()
    r::AbstractFloat = rand()
    sqrt(r) * v
end

function random_on_unit_sphere()::AbstractVec
    v::AbstractVec = Vec(randn(), randn(), randn())
    normalize(v)
end

function random_in_unit_sphere()::AbstractVec
    v::AbstractVec = random_on_unit_disk()
    r::AbstractFloat = rand()
    cbrt(r) * v
end

function reflect(v::AbstractVec, n::AbstractVec)::AbstractVec
    v::AbstractVec = normalize(v)
    v - 2dot(v, n)n
end

function refract(v::AbstractVec, n::AbstractVec, index::AbstractFloat)::Tuple{Bool, Union{Nothing, AbstractVec}}
    v::AbstractVec = normalize(v)
    dt::AbstractFloat = dot(v, n)
    discriminant::AbstractFloat = 1.0 - index^2 * (1.0 - dt^2)
    if discriminant > 0
        refracted::AbstractVec = index * (v - n * dt) - n * sqrt(discriminant)
        true, refracted
    end
    false, nothing
end


function nominal_diameter(method,P,Hn,Vdotn)
        if method == "Bier" # default
            D = 0.176*(P/Hn)^0.466
        elseif method == "Sarkaria"
            D = 0.71*P^0.43/Hn^0.65
        elseif method == "Warnick"
            D = 0.72*P^0.43/Hn^0.63
        elseif method == "Moffat"
            D = 0.52*P^0.43/Hn^0.6
        # methods not depending on efficiency
        elseif method == "USBR"
            D = 1.517*Vdotn^0.5/Hn^0.25
        elseif method == "Fahlbusch"
            D = 1.12*Vdotn^0.45/Hn^0.12
        end
    return D*1.5
end
# for darcys friction factor Sami and Jain
function f_Darcy(NRe, pipe_roughness,D)
    f = 1/(1.8*log10(6.9/NRe+(pipe_roughness/D/3.7)^1.1))^2
    return f
end
# for Reynolds number
function renoldsnum(rho,v,D,mu)
    return rho*v*D/mu
end
#


using LinearAlgebra

"""
    angularWaveVectorComponent(n, L)

Returns the angular wave vector component for a rectangular room given the
mode index and room length along the required component cartesian direction, in
meters.
# Example
```jldoctest
julia> nx = 7
7

julia> Lx = 5
5

julia> kx = angularWaveVectorComponent(nx, Lx)
4.39822971502571
```
"""
function angularWaveVectorComponent(n::Integer, L::Real)
    return n * π / L
end

"""
    angularWaveVector(nx, ny, nz, Lx, Ly, Lz)

Returns the angular wave vector of a rectangular room mode. Mode numbers along
the cartesian directions are nx ny and nz, the room sizes along the cartesian
directions are Lx, Ly and Lz meters.
# Example
```jldoctest
julia> angularWaveVector(1, 4, 7, 5.0, 4.0, 3.0)
3-element Array{Float64,1}:
 0.6283185307179586
 3.141592653589793
 7.330382858376184
```
"""
function angularWaveVector(
    nx::Integer, ny::Integer, nz::Integer,
    Lx::Real, Ly::Real, Lz::Real
    )

    return [
        angularWaveVectorComponent(nx, Lx);
        angularWaveVectorComponent(ny, Ly);
        angularWaveVectorComponent(nz, Lz)
        ]

end

"""
    modeFrequency(nx, ny, nz, Lx, Ly, Lz, c)

Returns the modal frequency of a mode characterized by indeces nx, ny and nz
along the cartesian directions for a rectangular room of sizes Lx, Ly and Lz
meters along the cartesian directions. c is the speed of sound in air, by
default 343.0 meters per second (speed at standard conditions).
# Example
```jldoctest
julia> modeFrequency(1, 4, 7, 5.0, 4.0, 3.0)
436.7174156260672

julia> modeFrequency(1, 4, 7, 5.0, 4.0, 3.0, 340.0)
432.89772977511035
```
"""
function modeFrequency(
    nx::Integer, ny::Integer, nz::Integer,
    Lx::Real, Ly::Real, Lz::Real,
    c::Real = 343.0
    )

    return norm(angularWaveVector(nx, ny, nz, Lx, Ly, Lz)) * c / 2π

end

"""
    mode(x, y, z, nx, ny, nz, Lx, Ly, Lz)

Returns the modal shape at coordinates in x, y, z for mode characterized by
indeces nx, ny and nz along the cartesian directions for a rectangular room of
sizes Lx, Ly and Lz meters along the cartesian directions. Result returned
normalized between 1 and -1, arbitrary units.
# Example
```jldoctest
julia> mode(2.5, 2.0, 1.5, 1, 4, 7, 5.0, 4.0, 3.0)
0.0

julia> mode.([3.0; 2.5], [1.0; 2.0], [0.75; 1.5], 1, 4, 7, 5.0, 4.0, 3.0)
2-element Array{Float64,1}:
 0.2185080122244105
 0.0
```
"""
function mode(
    x::Real, y::Real, z::Real,
    nx::Integer, ny::Integer, nz::Integer,
    Lx::Real, Ly::Real, Lz::Real
    )

return cospi(nx * x / Lx) * cospi(ny * y / Ly) * cospi(nz * z / Lz)

end

"""
    meshGrid(x, y, z)

Returns the 3-D grid coordinates defined by the vectors x, y, and z.
The grid is represented by three matrices of size length(y) by length(x) by
length(z). This is similar to MATLAB's meshgrid.
# Example
```jldoctest
julia> x = [1; 2]
2-element Array{Int64,1}:
 1
 2

julia> y = [3; 4]
2-element Array{Int64,1}:
 3
 4

julia> z = [5; 6]
2-element Array{Int64,1}:
 5
 6

julia> X, Y, Z = meshGrid(x, y, z)
([1 2; 1 2]

[1 2; 1 2], [3 3; 4 4]

[3 3; 4 4], [5 5; 5 5]

[6 6; 6 6])
```
"""
function meshGrid(x::AbstractArray, y::AbstractArray, z::AbstractArray)

    X = zeros(eltype(x), length(y), length(x), length(z))
    fX(i, j, k) = X[i, j, k] = x[j]
    [fX(i, j, k) for i in 1:length(y), j in 1:length(x), k in 1:length(z)]

    Y = zeros(eltype(y), length(y), length(x), length(z))
    fY(i, j, k) = Y[i, j, k] = y[i]
    [fY(i, j, k) for i in 1:length(y), j in 1:length(x), k in 1:length(z)]

    Z = zeros(eltype(z), length(y), length(x), length(z))
    fZ(i, j, k) = Z[i, j, k] = z[k]
    [fZ(i, j, k) for i in 1:length(y), j in 1:length(x), k in 1:length(z)]

    return X, Y, Z

end

"""
    indexGrid(Nx, Ny, Nz)

Returns a 3-D grid up to indeces Nx, Ny and Nz along the 3 cartesian
directions. The grid is represented by three matrices of size Ny by Nx by Nz.
# Example
```jldoctest
julia> Gx, Gy, Gz = indexGrid(3, 4, 5)
([1 2 3; 1 2 3; 1 2 3; 1 2 3]

[1 2 3; 1 2 3; 1 2 3; 1 2 3]

[1 2 3; 1 2 3; 1 2 3; 1 2 3]

[1 2 3; 1 2 3; 1 2 3; 1 2 3]

[1 2 3; 1 2 3; 1 2 3; 1 2 3], [1 1 1; 2 2 2; 3 3 3; 4 4 4]

[1 1 1; 2 2 2; 3 3 3; 4 4 4]

[1 1 1; 2 2 2; 3 3 3; 4 4 4]

[1 1 1; 2 2 2; 3 3 3; 4 4 4]

[1 1 1; 2 2 2; 3 3 3; 4 4 4], [1 1 1; 1 1 1; 1 1 1; 1 1 1]

[2 2 2; 2 2 2; 2 2 2; 2 2 2]

[3 3 3; 3 3 3; 3 3 3; 3 3 3]

[4 4 4; 4 4 4; 4 4 4; 4 4 4]

[5 5 5; 5 5 5; 5 5 5; 5 5 5])

```
"""
function indexGrid(Nx::Integer, Ny::Integer, Nz::Integer)
    return meshGrid(0:Nx, 0:Ny, 0:Nz)
end

"""
    roomAxes(rx, ry, rz, Lx, Ly, Lz)

Returns the 1-D axes coordinates along the sides of a rectangular room with
reference to a corner. The room has sizes Lx, Ly and Lz meters along the
cartesian directions, while rx, ry and rz express the step of the axes along
each cartesian direction.
"""
function roomAxes(
    rx::Real, ry::Real, rz::Real,
    Lx::Real, Ly::Real, Lz::Real
    )

    x = range(0.0, stop = Lx, length = round(Integer, 1 + Lx / rx))
    y = range(0.0, stop = Ly, length = round(Integer, 1 + Ly / ry))
    z = range(0.0, stop = Lz, length = round(Integer, 1 + Lz / rz))

    return x, y, z

end

"""
    roomGrid(rx, ry, rz, Lx, Ly, Lz)

Returns the 3-D grid coordinates inside a rectangular room with reference to a
corner. The room has sizes Lx, Ly and Lz meters along the cartesian directions,
while rx, ry and rz express the step of the grid along each cartesian direction.
"""
function roomGrid(
    rx::Real, ry::Real, rz::Real,
    Lx::Real, Ly::Real, Lz::Real
    )

    x, y, z = roomAxes(rx, ry, rz, Lx, Ly, Lz)

    return meshGrid(x, y, z)

end


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

n, N, alpha = 3, 10^7, 0.1

myT(nObs) = rand(Normal())/sqrt(rand(Chisq(nObs-1))/(nObs-1))
mcQuantile = quantile([myT(n) for _ in 1:N],alpha)
analyticQuantile = quantile(TDist(n-1),alpha)

println("Quantile from Monte Carlo: ", mcQuantile)
println("Analytic qunatile: ", analyticQuantile)

xGrid = -5:0.1:5
plot(xGrid, pdf.(Normal(), xGrid), c=:black, label="Normal Distribution")
scatter!(xGrid, pdf.(TDist(1) ,xGrid),
	c=:blue, msw=0, label="DOF = 1")
scatter!(xGrid, pdf.(TDist(3), xGrid),
	c=:red, msw=0, label="DOF = 3")
scatter!(xGrid, pdf.(TDist(100),xGrid),
	c=:green, msw=0, label="DOF = 100",
	xlims=(-4,4), ylims=(0,0.5), xlabel="X", ylabel="Density")


"""
    rkf45(f,a,b,y0,M,delta)
berisi program untuk mencari solusi persamaan
differensial `y' = f(t,y)` dengan `y(a) = y0` pada interval `[a, b]`. Program ini secara
default berisi 6 masukan, yaitu fungsi `f(t,y)`,
titik ujung interval penyelesaian `[a,b]`,
nilai awal `y0`, jumlah sub-interval `M` dan nilai toleransi error `delta`

# Example
```jl
julia> f(t,y) = 1+y^2;

julia> sol = rkf45(f,0,1.4,0,7,2e-5)
14×2 Array{Float64,2}:
 0.0     0.0
 0.2     0.20271
 0.6     0.684165
 0.8     1.02968
 ⋮
 1.35    4.45586
 1.375   5.04273
 1.3875  5.39534
 1.4     5.79895
```
return solusi masalah nilai awal `sol`.
"""
function rkf45(f,a,b,y0,M,delta)
  M = round(M)
  a2=1/4;b2=1/4;a3=3/8;b3=3/32;c3=9/32;
  a4=12/13;b4=1932/2197;c4=-7200/2197;d4=7296/2197;
  a5=1;b5=439/216;c5=-8;d5=3680/513;e5=-845/4104;
  a6=1/2;b6=-8/27;c6=2;d6=-3544/2565;e6=1859/4104;f6=-11/40;
  r1=1/360;r3=-128/4275;r4=-2197/75240;r5=1/50;r6=2/55;
  n1=25/216;n3=1408/2565;n4=2197/4104;n5=-1/5;
  big=1e15;
  h=(b-a)/M;
  hmin=h/64;
  hmax=h*64;
  maxi=200;
  j=1;
  Y = y0;
  T = a;
  br= b-0.001*abs(b);
  err=NaN
  while T[j]<b
    if (T[j]+h)>br;h=b-T[j];end
    #% Hitung koefisien
    k1=h*f(T[j],Y[j]);
    y2=Y[j]+b2*k1;
    k2=h*f(T[j]+a2*h,y2);
    y3=Y[j]+b3*k1+c3*k2;
    k3=h*f(T[j]+a3*h,y3);
    y4=Y[j]+b4*k1+c4*k2+d4*k3;
    k4=h*f(T[j]+a4*h,y4);
    y5=Y[j]+b5*k1+c5*k2+d5*k3+e5*k4;
    k5=h*f(T[j]+a5*h,y5);
    y6=Y[j]+b6*k1+c6*k2+d6*k3+e6*k4+f6*k5;
    k6=h*f(T[j]+a6*h,y6);
    err=abs(r1*k1+r3*k3+r4*k4+r5*k5+r6*k6);
    ynew=Y[j]+n1*k1+n3*k3+n4*k4+n5*k5;
    #% Perbarui ukuran langkah
    if (err<delta) || (h<2*hmin)
      Y = [Y; ynew];
      if (T[j]+h)>br
        T = [T; b];
      else
        T = [T; T[j]+h];
      end
      j=j+1;
    end
    if (err==0)
      s=0;
    else
      s=0.84*(delta*h/err)^(0.25);
    end
    if (s<0.75)&&(h>2*hmin);h = h/2;end
    if (s>1.50)&&(2*h<hmax);h = 2*h;end
    if abs(Y[j])>big || maxi==j;break;end
  end
  sol = [T Y];
  return sol
end


export getEdgeLength, getFaceArea, getCellVolume

function getEdgeLength(mesh::TetraMesh)
# get edge length

if isempty(mesh.L)

  ne = mesh.ne
  L = zeros(ne)

  for i=1:ne
    i1 = mesh.edges[i,1]
    i2 = mesh.edges[i,2]

    L[i] = getEdgeLength(mesh.Points[i1,:],  mesh.Points[i2,:])
  end

  mesh.L = L

end

return mesh.L
end

function getFaceArea(mesh::TetraMesh)
# get face size

if isempty(mesh.F)
  nf = mesh.nf
  F = zeros(nf)
  for i=1:nf
    i1 = mesh.faces[i,1]
    i2 = mesh.faces[i,2]
    i3 = mesh.faces[i,3]

    F[i] = getFaceArea(mesh.Points[i1,:], mesh.Points[i2,:], mesh.Points[i3,:])
  end

mesh.F = F
end

return mesh.F
end

function getCellVolume(mesh::TetraMesh)
# get volume

if isempty(mesh.V)

  nc = mesh.nc
  T  = mesh.Tetras
  V = zeros(nc)

  for i=1:nc
    i1 = T[i,1]
    i2 = T[i,2]
    i3 = T[i,3]
    i4 = T[i,4]

    V[i] = getCellVolume(mesh.Points[i1,:], mesh.Points[i2,:], mesh.Points[i3,:], mesh.Points[i4,:])
  end

  mesh.V = V

end

return mesh.V
end

function getEdgeLength(P1::Array{Float64,1}, P2::Array{Float64,1})
	return sqrt((P1[1] - P2[1])^2 + (P1[2] - P2[2])^2 + (P1[3] - P2[3])^2)
end

function getFaceArea(P1::Array{Float64,1}, P2::Array{Float64,1}, P3::Array{Float64,1})
	return norm(cross(P3 - P1, P3 - P2))/2
end

function getCellVolume(P1::Array{Float64,1}, P2::Array{Float64,1}, P3::Array{Float64,1}, P4::Array{Float64,1})
	return 1/6.0*abs(dot(P4 - P3, cross(P4 - P1, P4 - P2)))
end


using StaticArrays: SArray, @SMatrix
using LinearAlgebra: I, norm

# This matrix-based implementation is very cool (IMO)
# and only ~25% slower than the reference implementation for addition, multiplication, subtraction
#
# We get a 10x speedup for division and abs (although, for abs in particular, base widens to guard against overflow)
struct ComplexNumber{T<:Real} <: Number
    mat::SArray{Tuple{2,2},T,2,4}
    ComplexNumber(mat) = new{eltype(mat)}(mat)
    ComplexNumber(real::T1,imag::T2) where {T1 <: Real, T2 <: Real} = new{promote_type(T1,T2)}(@SMatrix([real  -imag; imag real]))
end

Base.:*(l::ComplexNumber,r::ComplexNumber) = ComplexNumber(l.mat * r.mat)
Base.:+(l::ComplexNumber,r::ComplexNumber) = ComplexNumber(l.mat + r.mat)
Base.:-(l::ComplexNumber,r::ComplexNumber) = ComplexNumber(l.mat - r.mat)
Base.:/(l::ComplexNumber,r::ComplexNumber) = ComplexNumber(l.mat * inv(r.mat))
Base.:*(z::ComplexNumber,r::Real) = ComplexNumber(r*z.mat)
Base.:*(r::Real, z::ComplexNumber) = z*r
Base.:+(z::ComplexNumber,r::Real) = ComplexNumber(r*I + z.mat)
Base.:+(r::Real, z::ComplexNumber) = z + r

Base.show(io::IO, z::ComplexNumber) = print(io, "ComplexNumber(", Re(z), ", ", Im(z), ")")
Base.:≈(l::ComplexNumber,r::ComplexNumber) = l.mat ≈ r.mat # I had to make one of the division tests ≈ to pass it
Base.:(==)(l::ComplexNumber,r::ComplexNumber) = l.mat == r.mat

Base.abs(z::ComplexNumber) = norm(z.mat[:,1])

# Sorry about the capital letter. That's what maths does.
@inline Re(z::ComplexNumber) = @inbounds z.mat[1,1]
@inline Im(z::ComplexNumber) = @inbounds z.mat[2,1]

Base.real(z::ComplexNumber) = Re(z)
Base.imag(z::ComplexNumber) = Im(z)

Base.conj(z::ComplexNumber) = ComplexNumber(transpose(z.mat))

Base.exp(z::ComplexNumber) = ComplexNumber(exp(z.mat))

# There's not an iota of a chance that ι will take off as the sensible symbol to use for the imaginary unit
const ι = jm = ComplexNumber(0,1)


# This file includes the MonteCarlo simulation of ASK modulation scheme and compares
# the numerical results with the theoretical results.

using DigitalCommunications
using Plots

# Simulation parameters
k = 3
M = 2^k
nsymbols = Int(1e6)
nbits = k * nsymbols
Ep = 100.                       # Modulating pulse energy.
ebno = collect(0 : 10)          # Snr per bit
esno = ebno .+ 10 * log10(k)    # Snr ber symbol

# Communcation system components
gen = BitGenerator(nbits)
coding = GrayCoding(M)
modulator = Modulator(PAM(M))
channel = AWGNChannel()
detector = Detector(alphabet(modulator))

# Monte Carlo simulation
message = coding(gen.bits)  # Message signal
symerr = zeros(length(esno))
for i in 1 : length(symerr)
    channel.esno = esno[i]  # Update channel snr
    mbar = gen.bits |> coding |> modulator |> channel |> detector  # Extracted message signal
    symerr[i] = sum(mbar .!= message) / length(message)  # Symbol error rate
end

# Plots
plt = plot(title="$M-PAM", xlabel="ebno [dB]", ylabel="Pe")
plot!(ebno, berpam.(ebno, M), marker=:circle, yscale=:log10, label="theoretical")
plot!(ebno, symerr, marker=:circle, yscale=:log10, label="montecarlo")


#==============================================================================
    Code for solving the Hamiltonian Jacboi Bellman for
	   an basic model: a neoclassical growth model: ρV(k) = max_{c} U(c) + V'(k)[F(k)-δk - c]
	   Where s(k) = F(k) - δk - c(k) and c(k) = U'^{-1}(V'(k))

	Translated Julia code from Matlab code by Ben Moll:
        http://www.princeton.edu/~moll/HACTproject.htm

		Updated to julia 1.0.0
==============================================================================#

using Distributions, Plots

σ= 2.0 #
ρ = 0.05 #the discount rate
δ = 0.05 # the depreciation rate
A = 1.0
α= 0.3

k_ss = (α*A/(ρ+δ))^(1/(1-α))

H= 150
k_min = 0.001*k_ss
k_max = 2.0*k_ss

k = LinRange(k_min, k_max, H)
k = convert(Array, k) # create grid for a values
dk = (k_max-k_min)/(H-1)

maxit = 1000
ε = 10e-6

dVf, dVb = [zeros(H,1) for i=1:2]

#initial guess for V
v0 = (A.*k.^α).^(1-σ)/(1-σ)/ρ
v= v0

dist=[]

for n=1:maxit
	V=v

    # forward difference
	dVf[1:H-1] = (V[2:H]-V[1:H-1])/dk
	dVf[H] = 0

	# backward difference
	dVb[2:H] = (V[2:H]-V[1:H-1])/dk
	dVb[1] = 0 # the boundary condition

	I_concave = dVb .> dVf

    # consumption and savings with forward difference
    cf = dVf.^(-1/σ)
    μ_f = A.*k.^α - δ.*k -cf

    # consumption and savings with backward difference
    cb = dVb.^(-1/σ)
    μ_b = A.*k.^α - δ.*k -cb

	c0 = A.*k.^α - δ.*k
    dV0 = c0.^(-σ)

    # Now to make a choice between forward and backward difference
    If = μ_f .> 0
    Ib = μ_b .< 0
    I0 = 1.0.-If-Ib
    Ib[1] = false
    If[1] = true
    Ib[H] = true
    If[H] = false

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

    global c = dV_Upwind.^(-1/σ)
    V_change = c.^(1-σ)/(1-σ) + dV_Upwind.*(A.*k.^α - δ.*k-c) -ρ.*V

	# update
	Δ = .9*dk/(findmax(A.*k.^α- δ.*k-c)[1])
	global v = v + Δ*V_change
	push!(dist,findmax(abs.(V_change))[1])
	if dist[n] .< ε
		println("Value Function Converged Iteration=")
		println(n)
		break
	end

end

plot(dist, grid=false,
		xlabel="Iteration", ylabel="||V^{n+1} - V^n||",
		ylims=(-0.001,0.030),
		legend=false, title="")
png("Convergence")

v_err = c.^(1-σ)/(1-σ) + dV_Upwind.*(A.*k.^α - δ.*k -c) - ρ.*v

plot(k, v_err, grid=false,
		xlabel="k", ylabel="Error in the HJB equation",
		xlims=(k_min,k_max),
		legend=false, title="")
png("HJB_error")

plot(k, v, grid=false,
		xlabel="k", ylabel="V(k)",
		xlims=(k_min,k_max),
		legend=false, title="")
png("Value_function_vs_k")

plot(k, c, grid=false,
		xlabel="k", ylabel="c(k)",
		xlims=(k_min,k_max),
		legend=false, title="")
png("c(k)_vs_k")

# approximation at the borrowing constraint
k_dot = (A.*k.^α - δ.*k -c)

plot(k, k_dot, grid=false,
		xlabel="k", ylabel="s(k)",
		xlims=(k_min,k_max), title="", label="s(k)", legend=:topright)
plot!(k, zeros(H,1), label="", line=:dash)
png("stateconstraint")


using Distributions
using StatsBase

function Binom(k, N, p)
    return pdf(Binomial(N, p), k)
end

function Binom_sum(N, p)
    if iseven(N)
        N -= 1
    end
    return sum([Binom(k, N, p) for k = N:-1:N/2])
end

function compute_analytical_accuracy(; N, p_GC_sig, p_GC_bkg, weight = [1, 1])
    p_α_sig = 1 - p_GC_sig
    p_α_bkg = 1 - p_GC_bkg
    return mean([Binom_sum(N, p_α_sig), Binom_sum(N, p_α_bkg)], weights(weight))
end

function compute_analytical_accuracies(;
    N_half,
    p_GC_sig = 0.4,
    p_GC_bkg = 0.5,
    weight = [1, 1],
)
    bases_included = 1:N_half
    analytical_accuracies = [
        compute_analytical_accuracy(
            N = i * 2,
            p_GC_sig = p_GC_sig,
            p_GC_bkg = p_GC_bkg,
            weight = weight,
        ) for i in bases_included
    ]
    return analytical_accuracies
end

#%%

using CairoMakie
using ColorSchemes

N_half = 38
x = 1:N_half

analytical_accuracies = compute_analytical_accuracies(;N_half=N_half)

colormap = [x for x in ColorSchemes.Set1_9.colors]

#%%

f = Figure()
ax = Axis(
    f[1, 1],
    title = "Analytical accuracy",
    xlabel = "# bases included (centered)",
    ylabel = "Accuracy",
    # limits = (0.5, half_seq_length + 0.5, 0.634, 0.701),
    # limits = (0.5, half_seq_length + 0.5, ylimits...),
    xticks = 1:2:N_half,
)

scatterlines!(
    ax,
    x,
    analytical_accuracies,
    color = colormap[3],
    markercolor = colormap[3],
    label = "Analytical",
)

axislegend(position = :rb)
f

#%%

compute_analytical_accuracy(N=1, p_GC_sig = 0.4, p_GC_bkg = 0.5)
compute_analytical_accuracy(N=2, p_GC_sig = 0.4, p_GC_bkg = 0.5)
compute_analytical_accuracy(N=3, p_GC_sig = 0.4, p_GC_bkg = 0.5)
compute_analytical_accuracy(N=4, p_GC_sig = 0.4, p_GC_bkg = 0.5)
compute_analytical_accuracy(N=5, p_GC_sig = 0.4, p_GC_bkg = 0.5)
compute_analytical_accuracy(N=6, p_GC_sig = 0.4, p_GC_bkg = 0.5)


"""
    dualgraph(dist::Matrix{Float64}; method::Symbol = :inverse, σ::Float64 = 1.0)

build the dual graph's weight matrix based on the given non-trivial eigenvector
metric.

# Input Arguments
- `dist::Matrix{Float64}`: eigenvector distance matrix
- `method::Symbol`: default is by taking inverse of the distance between
    eigenvectors. Ways to build the dual graph edge weights. Option: `:inverse`,
    `:gaussian`.
- `σ::Float64`: default is `1.0`. Gaussian variance parameter.

# Output Argument
- `G_star::GraphSig`: A `GraphSig` object containing the weight matrix of the
    dual graph.

"""
function dualgraph(dist::Matrix{Float64}; method::Symbol = :inverse, σ::Float64 = 1.0)
    N = Base.size(dist, 1)
    W_star = zeros(N, N)
    if method == :inverse
        for i = 1:(N - 1), j = (i + 1):N
            W_star[i, j] = 1 / dist[i, j]
        end
    elseif method == :gaussian
        for i = 1:(N - 1), j = (i + 1):N
            W_star[i, j] = exp(-dist[i, j] / σ^2)
        end
    else
        error("method must be :inverse or :gaussian.")
    end
    W_star = W_star + W_star'
    G_star = GraphSig(sparse(W_star); name = "dual graph")
    return G_star
end


using Random, Statistics
Random.seed!(1)
allData = [rand(24), rand(15), rand(73)]

xBarArray = mean.(allData)
nArray = length.(allData)
xBarTotal = mean(vcat(allData...))
L = length(nArray)

ssBetween=sum([nArray[i]*(xBarArray[i] - xBarTotal)^2 for i in 1:L])
ssWithin=sum([sum([(ob - xBarArray[i])^2 for ob in allData[i]]) for i in 1:L])
ssTotal=sum([sum([(ob - xBarTotal)^2 for ob in allData[i]]) for i in 1:L])

println("Sum of squares between groups: ", ssBetween)
println("Sum of squares within groups: ", ssWithin)
println("Sum of squares total: ", ssTotal)


import Random
using Plots

println("Pool Testing Optimization\nMonte Carlo Simulation")

min_total = 1
max_total = 100
min_infected = 0
iterations = 100

function run(total, infected, n_groups)
	# Trivial cases
	# No one is infected or testing everybody
	if infected==0 || n_groups==total
		return n_groups
	# Only 1 group (there is only 1 grouping step)
	elseif n_groups==1
		return total+1
	end

	people = zeros(Bool, total)
	people[1:infected] .= true
	Random.shuffle!(people)

	n_per_group = total ÷ n_groups
	n_plus_groups = total % n_groups

	groups = zeros(UInt, n_groups)

	for i in 1:infected
		j = 1+(Random.rand(UInt) % n_groups)
		groups[j] += 1
	end

	infected_groups = count(g -> g>0, groups)

	tests = n_groups

	for (index, value) in enumerate(groups)
		if value>0
			if index<=n_plus_groups
				tests += n_per_group+1
			else
				tests += n_per_group
			end
		end
	end

	#=
	println("Total   \t", total)
	println("Infected\t", infected)
	println("People\t", people)
	println("People per group\t", n_per_group)
	println("# groups with 1 more\t", n_plus_groups)
	println("Infected per group\t", groups)
	println("Number of infected groups\t", infected_groups)
	println("Number of tests\t", tests)
	=#

	return tests
end

n = binomial(max_total+2-min_infected, 2) - binomial(min_total+1-min_infected, 2)
results = Array{UInt, 2}(undef, n, 3)

i = 1
for total in min_total:max_total
	for infected in min_infected:total
		println("total ", total, "  infected ", infected)

		optimum_n_groups = total
		previous_sum_tests = iterations*total

		for n_groups in 1:total
			sum_tests = 0

			for iteration in 1:iterations
				tests = run(total, infected, n_groups)
				sum_tests += tests
				#println("total ", total, "  infected ", infected, "  n_groups ", n_groups, "  iteration ", iteration, "  tests ", tests)
			end

			if sum_tests<previous_sum_tests
				optimum_n_groups = n_groups
			end
		end

		global i
		results[i, :] = [total, infected, optimum_n_groups]
		i += 1
	end
end

println("Calculated")

matrix = zeros(Float32, (max_total+1-min_total, max_total+1-min_infected))
for k = 1:size(results, 1)
	local i = results[k, 1] + 1 - min_total
	local j = results[k, 2] + 1 - min_infected
	matrix[i, j] = results[k, 3]/results[k, 1]
end

#println("results ", results)
#println("matrix ", matrix)

heatmap(min_total:max_total, min_infected:max_total, matrix',
		xlabel="Total", ylabel="Infected",
		title="Optimal number of groups / Total",
		aspect_ratio=:equal,
		show=true)
println("Plotted")
readline()

println("End")


### A Pluto.jl notebook ###
# v0.14.2

using Markdown
using InteractiveUtils

# ╔═╡ 0a5335e0-ab9e-11eb-0a85-139ff5cf0b1d
md"# Chapter 02 from NNFS book"

# ╔═╡ 0c5d667f-0ea5-4d21-80ad-58d144470a73
md"### A Single Neuron (4 inputs, 1 output)"

# ╔═╡ d85b097d-c7c7-439f-80ec-a1b75c696cb1
md"""
for **each input** we need **one weigth** and for **each output** we need **one bias**.
So, for this simple example with 1 neuron with 4 inputs we are gonna have:

- **3** inputs
- **3** weights
- **1** bias.
"""

# ╔═╡ ef746509-f29b-42ae-b9ec-a2e18ad4b327
begin
	inputs = [1; 2; 3; 2.5]
	weights = [0.2; 0.8; -0.5; 1.0]
	bias = 2.0

	inputs'*weights + bias # or dot(inputs, weights) + bias
end

# ╔═╡ a3acafe8-6a4b-4438-84d9-16a00cb21997
md"### A Layer of 03 Neurons (4 inputs, 3 outputs)"

# ╔═╡ 9fbc1e4c-52c0-4b99-92dc-4969b55f645d
begin
	#inputs = [1, 2, 3, 2.5]
	weights1 = [0.2, 0.8, -0.5, 1.0]
	weights2 = [0.5, -0.91, 0.26, -0.5]
	weights3 = [-0.26, -0.27, 0.17, 0.87]
	bias1 = 2.0
	bias2 = 3.0
	bias3 = 0.5

	[inputs'*weights1 + bias1,
	inputs'*weights2 + bias2,
	inputs'*weights3 + bias3]
end

# ╔═╡ 3613059e-501d-4cb0-bd73-8e42e6395e16
md" with only one unique matrix of weights we use matrix multiplication from linear algebra"

# ╔═╡ 0e309e2a-59e8-4fce-819d-afeec8c21ca8
begin
	weights_array = [
		0.2    0.8  -0.5   1.0;
		0.5   -0.91  0.26 -0.5;
		-0.26 -0.27  0.17  0.87;
	]
	bias_vec = [2.0 3.0 0.5]

	(weights_array * inputs) .+ bias_vec
end

# ╔═╡ d27568ac-934c-4408-880d-4bd34ccead72
md"### A Layer of Neurons with multiple batch of Data"

# ╔═╡ 91d6b3e6-5124-4963-b6c2-f2c65b67b045
md"Now we are gonna need use transpose operation to match the right vectors in matrix multiplication"

# ╔═╡ 7ae39d13-e448-4c91-81be-d44a15f1a414
begin
	inputs_batch = [
		1.0 2.0 3.0 2.5
		2.0 5.0 -1.0 2.0
		-1.5 2.7 3.3 -0.8
	]

	(inputs_batch * weights_array') .+ bias_vec
end

# ╔═╡ Cell order:
# ╟─0a5335e0-ab9e-11eb-0a85-139ff5cf0b1d
# ╟─0c5d667f-0ea5-4d21-80ad-58d144470a73
# ╟─d85b097d-c7c7-439f-80ec-a1b75c696cb1
# ╠═ef746509-f29b-42ae-b9ec-a2e18ad4b327
# ╟─a3acafe8-6a4b-4438-84d9-16a00cb21997
# ╠═9fbc1e4c-52c0-4b99-92dc-4969b55f645d
# ╟─3613059e-501d-4cb0-bd73-8e42e6395e16
# ╠═0e309e2a-59e8-4fce-819d-afeec8c21ca8
# ╟─d27568ac-934c-4408-880d-4bd34ccead72
# ╟─91d6b3e6-5124-4963-b6c2-f2c65b67b045
# ╠═7ae39d13-e448-4c91-81be-d44a15f1a414


# function nint(x::Float64)
#     return round(Int, x)
# end

# function geo_coordinate(x, y)
#     PI = 3.141592

#     deg = nint(x)
#     m = x - deg
#     latitude = PI * (deg + 5.0 * m / 3.0) / 180.0

#     deg = nint(y)
#     m = y - deg
#     longitude = PI * (deg + 5.0 * m / 3.0) / 180.0

#     return latitude, longitude
# end

# function distance2D(xi, yi, xj, yj; dist="EUC_2D")
#     if dist == "EUC_2D"
#         xd = xi - xj
#         yd = yi - yj
#         return nint(sqrt(xd*xd + yd*yd))
#     elseif dist == "MAN_2D"
#         xd = abs(xi - xj)
#         yd = abs(yi - yj)
#         return nint(xd + yd)
#     elseif dist == "MAX_2D"
#         xd = abs(xi - xj)
#         yd = abs(yi - yj)
#         return max(nint(xd), nint(yd))
#     elseif dist == "GEO"
#         lat_i, long_i = geo_coordinate(xi, yi)
#         lat_j, long_j = geo_coordinate(xj, yj)
#         RRR = 6378.388
#         q1 = cos(long_i - long_j)
#         q2 = cos(lat_i - lat_j)
#         q3 = cos(lat_i + lat_j)
#         dij =  RRR * acos( 0.5*((1.0+q1)*q2 - (1.0-q1)*q3) ) + 1.0
#         return floor(Int, dij)
#     else
#         error("Distance function $dist is not supported.")
#     end
# end

# function dist_matrix(x::Vector{Float64}, y::Vector{Float64}; dist="EUC_2D")
#     n_nodes = length(x)
#     @assert length(x) == length(y)

#     M = Matrix{Int}(undef, n_nodes, n_nodes)

#     for i in 1:n_nodes
#         for j in i:n_nodes
#             if i == j
#                 M[i, j] = 0
#             else
#                 M[i, j] = distance2D(x[i], y[i], x[j], y[j]; dist=dist)
#                 M[j, i] = M[i, j]
#             end
#         end
#     end
#     return M
# end


#' ---
#' title: Geographical Clustering With Additional Constraint
#' ---

#' **Originally Contributed by**: Matthew Helm ([with help from Mathieu Tanneau on Julia Discourse](https://discourse.julialang.org/t/which-jump-jl-solver-for-this-problem/43350/17?u=mthelm85))

#' The goal of this exercise is to cluster $n$ cities into $k$ groups, minimizing the total pairwise distance between cities
#' *and* ensuring that the variance in the total populations of each group is relatively small.

#' For this example, we'll use the 20 most populous cities in the United States.

using Cbc
using DataFrames
using Distances
using JuMP
using LinearAlgebra

cities = DataFrame(
    city=[ "New York, NY", "Los Angeles, CA", "Chicago, IL", "Houston, TX", "Philadelphia, PA", "Phoenix, AZ", "San Antonio, TX", "San Diego, CA", "Dallas, TX", "San Jose, CA", "Austin, TX", "Indianapolis, IN", "Jacksonville, FL", "San Francisco, CA", "Columbus, OH", "Charlotte, NC", "Fort Worth, TX", "Detroit, MI", "El Paso, TX", "Memphis, TN"],
    population=[8405837,3884307,2718782,2195914,1553165,1513367,1409019,1355896,1257676,998537,885400,843393,842583,837442,822553,792862,792727,688701,674433,653450],
    lat=[40.7127,34.0522,41.8781,29.7604,39.9525,33.4483,29.4241,32.7157,32.7766,37.3382,30.2671,39.7684,30.3321,37.7749,39.9611,35.2270,32.7554,42.3314,31.7775,35.1495],
    lon=[-74.0059,-118.2436,-87.6297,-95.3698,-75.1652,-112.0740,-98.4936,-117.1610,-96.7969,-121.8863,-97.7430,-86.1580,-81.6556,-122.4194,-82.9987,-80.8431,-97.3307,-83.0457,-106.4424,-90.0489])

#' ### Model Specifics

#' We will cluster these 20 cities into 3 different groups and we will assume that the ideal or target population $P$ for a
#' group is simply the total population of the 20 cities divided by 3:

n = size(cities,1)
k = 3
P = sum(cities.population) / k

#' ### Obtaining the distances between each city

#' Let's leverage the *Distances.jl* package to compute the pairwise Haversine distance between each of the cities in our data
#' set and store the result in a variable we'll call `dm`:

dm = Distances.pairwise(Haversine(6372.8), Matrix(cities[:, [3,4]])', dims=2)

#' Our distance matrix is symmetric so we'll convert it to a `LowerTriangular` matrix so that we can better interpret the
#' objective value of our model (if we don't do this the total distance will be doubled):

dm = LowerTriangular(dm)

#' ### Build the model
#' Now that we have the basics taken  care of, we can set up our model, create decision variables, add constraints, and then
#' solve.

#' First, we'll set up a model that leverages the [Cbc](https://github.com/coin-or/Cbc) solver. Next, we'll set up a binary
#' variable $x_{i,k}$ that takes the value $1$ if city $i$ is in group $k$ and $0$ otherwise. Each city must be in a group, so
#' we'll add the constraint $\sum_kx_{i,k} = 1$ for every $i$.

model = Model(Cbc.Optimizer)

@variable(model, x[1:n, 1:k], Bin)

for i in 1:n
    @constraint(model, sum(x[i,:]) == 1)
end

#'The total population of a group $k$ is $Q_k = \sum_ix_{i,k}q_i$ where $q_i$ is simply the $i$th value from the `population`
#' column in our `cities` DataFrame. Let's add constraints so that $\alpha \leq (Q_k - P) \leq \beta$. We'll set $\alpha$
#' equal to -2,500,000 and $\beta$ equal to 2,500,000. By adjusting these thresholds you'll find that there is a tradeoff
#' between having relatively even populations between groups and having geographically close cities within each group. In
#' other words, the larger the absolute values of $\alpha$ and $\beta$, the closer together the cities in a group will be but
#' the variance between the group populations will be higher.

α = -2_500_000
β = 2_500_000

for i in 1:k
    @constraint(model, (x' * cities.population)[i] - P <= β)
    @constraint(model, (x' * cities.population)[i] - P >= α)
end

#' Now we need to add one last binary variable $z_{i,j}$ to our model that we'll use to compute the total distance between the
#' cities in our groups, defined as  $\sum_{i,j}d_{i,j}z_{i,j}$. Variable $z_{i,j}$ will equal $1$ if cities $i$ and $j$ are
#' in the same group, and $0$ if they are not in the same group.

#' To ensure that $z_{i,j} = 1$ if and only if cities $i$ and $j$ are in the same group, we add the constraints $z_{i,j} \geq
#' x_{i,k} + x_{j,k} - 1$ for every pair $i,j$ and every $k$:

@variable(model, z[1:n,1:n], Bin)

for k in 1:k, i in 1:n, j in 1:n
    @constraint(model, z[i,j] >= x[i,k] + x[j,k] - 1)
end

#' We can now add an objective to our model which will simply be to minimize the dot product of $z$ and our distance matrix,
#' `dm`. We can then call `optimize!` and review the results.

@objective(model, Min, dot(z,dm));

optimize!(model)

#' ### Reviewing the Results

#' Now that we have results, we can add a column to our `cities` DataFrame for the group and then loop through our $x$
#' variable to assign each city to its group. Once we have that, we can look at the total population for each group and also
#' plot the cities and their groups to verify visually that they are grouped by geographic proximity.

cities.group = zeros(n)

for i in 1:n, j in 1:k
    if round(value.(x)[i,j]) == 1.0
        cities.group[i] = j
    end
end

for group in groupby(cities, :group)
    @show group
    println("")
    @show sum(group.population)
    println("")
end

#' The populations of each group are fairly even and we can see from the plot below that the groupings look good in terms of
#' geographic proximity:

#' <img src="img/geo_clusters.png" style="width: auto; height: auto" alt="Geographic Clusters">


"""
    id(k::Int64) -> SparseMatrixCSC{ℤ₂, Int64}

Generate and return a `k × k` identity-matrix.
"""
function id(k::Int64)::SparseMatrixCSC{ℤ₂, Int64}
    I = sparse(zeros(ℤ₂, k, k))
    for i = 1:k
        I[i,i] = one(ℤ₂)
    end
    return I
end

"""
    anti_id(k::Int64) -> SparseMatrixCSC{ℤ₂, Int64}

Generate and return a `k × k` anti-identity-matrix.
"""
function anti_id(k::Int64)::SparseMatrixCSC{ℤ₂, Int64}
    antiid = ones(ℤ₂, k, k)
    for i = 1:k
        antiid[i,i] = zero(ℤ₂)
    end
    return sparse(antiid)
end


using Images, Colors

function filter_sobel(image::AbstractMatrix, direction::String="-")

  tipo = typeof(image)

  try
    image = convert(Array{Float32}, image)
  catch er
    error(er,"\nProblem converting the image into an array.\n")
  end

  if direction=="-"
    mask_x = [-1 0 1; -2 0 2; -1 0 1];
    mask_y = [1 2 1; 0 0 0; -1 -2 -1];

    Gx = imfilter(image,centered(mask_x));
    Gy = imfilter(image,centered(mask_y));

    d1 = size(Gx);
    d2 = size(Gy);

    if d1!=d2
      error("Problem during the function's execution.")
    end

    img_final = zeros(d1[1],d1[2])

    for i = 1:d1[1]
      for j = 1:d1[2]

        img_final[i,j] = sqrt((Gx[i,j]^2)+(Gy[i,j]^2))

        if img_final[i,j]>1.0
          img_final[i,j]=1.0;
        elseif img_final[i,j]<0.0
          img_final[i,j]=0.0;
        end

      end
    end
elseif direction=="W->E"
    mask = [1 2 1; 0 0 0; -1 -2 -1];
    img_final = imfilter(image,centered(mask));
    d1 = size(img_final);
    for i = 1:d1[1]
      for j = 1:d1[2]
        if img_final[i,j]>1.0
          img_final[i,j]=1.0;
        elseif img_final[i,j]<0.0
          img_final[i,j]=0.0;
        end
      end
    end
elseif direction=="N->S"
    mask = [-1 0 1; -2 0 2; -1 0 1];
    img_final = imfilter(image,centered(mask));
    d1 = size(img_final);
    for i = 1:d1[1]
      for j = 1:d1[2]
        if img_final[i,j]>1.0
          img_final[i,j]=1.0;
        elseif img_final[i,j]<0.0
          img_final[i,j]=0.0;
        end
      end
    end
  elseif direction=="NW->SE"
    mask = [0 1 2; -1 0 1; -2 -1 0];
    img_final = imfilter(image,centered(mask));
    d1 = size(img_final);
    for i = 1:d1[1]
      for j = 1:d1[2]
        if img_final[i,j]>1.0
          img_final[i,j]=1.0;
        elseif img_final[i,j]<0.0
          img_final[i,j]=0.0;
        end
      end
    end
  elseif direction=="SW->NE"
    mask = [-2 -1 0; -1 0 1; 0 1 2];
    img_final = imfilter(image,centered(mask));
    d1 = size(img_final);
    for i = 1:d1[1]
      for j = 1:d1[2]
        if img_final[i,j]>1.0
          img_final[i,j]=1.0;
        elseif img_final[i,j]<0.0
          img_final[i,j]=0.0;
        end
      end
    end
  else
    error("The position value must be an Text variable. Available values are: W->E, N->S, NW->SE, and SW->NE. This is not a mandatory field.\n");
  end

  img_final = convert(tipo, img_final)
  return(img_final)
end


module Quadratic

using STMO: Tracker, notrack, PathTrack, trace

export fquad, plot_quadratic, solve_quadratic, quadratic_ls, gradient_descent

using LinearAlgebra, Plots

"""Compute a quadratic function."""
fquad(x::Vector, P::Matrix, q::Vector, r::Real=0.0) = 0.5x' * P * x + q' * x + r

"""Plot a 1-D quadratic function."""
function plot_quadratic(p::Real, q::Real, r::Real, (xl, xu), stepsize=0.1; kwargs...)
    return plot(x -> p*x^2 + q*x+r, xl:stepsize:xu, xlabel="\$x\$"; kwargs...)
end

"""Plot a 2-D quadratic function."""
function plot_quadratic(P::AbstractMatrix, q::AbstractVector, r::Real,
                            (x1l, x1u), (x2l, x2u), stepsize=0.1; kwargs...)
    fun = (x1, x2) -> [x1,x2] |> x -> 0.5x' * P * x + q' * x + r
    return contour(x1l:stepsize:x1u, x2l:stepsize:x2u, fun, xlabel="\$x_1\$",
                                            ylabel="\$x_2\$"; kwargs...)
end

"""
    solve_quadratic(p::Real, q::Real[, r::Real=0.0])

Finds the minimizer of an 1-D quadratic system. If the quadratic term `p` is
negative, the function raises an error.

Inputs:
    - `p`, `q`, `r`: the terms of the 1-D quadratic system

Output:
    - xstar: the minimizer, a number
"""
function solve_quadratic(p::Real, q::Real, r::Real=0.0)
    @assert p > 0.0
    return - q / p
end

"""
    solve_quadratic(P::AbstractMatrix, q::AbstractVector, r::Real=0)

Finds the minimizer of an N-D quadratic system.
P is assumed to be a symmetric positive-definite matrix.

Inputs:
    - P, q, r: the terms of the nD quadratic system

Output:
    - xstar: the minimizer, an (n x 1) vector
"""
function solve_quadratic(P::AbstractMatrix, q::AbstractVector, r::Real=0.0)
    return - P \ q
end

"""
    quadratic_ls(P::AbstractMatrix, q::AbstractVector, Δx::AbstractVector,
                                    x::AbstractVector)

Find the exact step size that minimized a quadratic system in
a given point x for a given search direction Dx

Inputs:
    - P, q: the terms of the nD quadratic system
    - x: starting point
    - Δx: search direction

Output:
    - t: optimal step size
"""
function quadratic_ls(P::AbstractMatrix, q::AbstractVector, Δx::AbstractVector,
                                    x::AbstractVector)
    t = solve_quadratic(Δx' * P * Δx, Δx' * P * x + Δx ⋅ q)
    return t
end

"""
    gradient_descent(P::AbstractArray, q::AbstractVector,
            x₀::AbstractVector; β::Real=0.0, ϵ::Real=1e-6,
            tracker::Tracker=notrack)

Computes the minimizes of a quadratic system using gradient descent. Optionally
provide momentum.

Inputs:
    - P, q: the terms of the nD quadratic system
    - x₀: starting point
    - ϵ: convergence parameter
    - β: momentum parameter
    - tracker: object of the type `Tracker` to save the steps

Outputs:
    - xstar: the found minimum
"""
function gradient_descent(P::AbstractArray, q::AbstractVector,
            x₀::AbstractVector; β::Real=0.0, ϵ::Real=1e-6,
            tracker::Tracker=notrack)
    @assert 0 ≤ β < 1
    x = x₀  # initial value
    Δx = zero(x)  # pre-allocate a vector for the gradient
    while true
        Δx .= (1.0 - β) * (- P * x .- q) .+ β *  Δx
        if norm(Δx) < ϵ
            break
        end
        # determine stepsize using exact line search
        t = quadratic_ls(P, q, Δx, x)
        # perform step
        x .+= t * Δx
        trace(tracker, x)  # saves the steps
    end
    return x
end

# SIGNAL RECOVERY
# ---------------

module SignalRecovery
using LinearAlgebra: I

signalfun(x, n) = 3sin(x * 2 * pi / n) +
            2cos(x * 4 * pi / n) +
            sin(x * 4 * pi / n) + 0.8 * cos(x * 12 * pi / n)

"""
    generate_noisy_measurements(m, n; σ=1.0)

Generate noisy measurements according to some function f

Inputs:
    - m : number of observations
    - n : dim of x
    - σ : normally distributed noise (default = 1)

Output:
    - y : vector of noisy measurements
    - ind : vector of indices
"""
function generate_noisy_measurements(m, n; σ=1.0)
    ind = rand(1:n, m)
    y = signalfun.(ind, n) .+ randn(m) * σ
    return y, ind
end

"""
    make_connection_matrix(n, γ=100)

Generates the kernel matrix and the inverse kernel matrix

Uses γ as a characteristic length scale of the radial basis kernel,
assumes periodic boundaries.

Inputs:
    - n : number of points
    - γ : length scale

Output:
    - K
    - Kinv
"""
function make_connection_matrix(n, γ=100)
    # cyclic distance
    d(i, j) = min((i-j)^2, ((i-n)-j)^2, ((j-n)-i)^2)
    # distance matrix
    D = [d(i, j) for i in 1:n, j in 1:n]
    # kernel matrix
    K = exp.(-D / γ^2) + 1e-2I
    return K, inv(K)
end

"""
    make_bookkeeping(I, n)

Constructs the bookkeeping matrix R

Inputs:
    - ind : the indices of the measurements
    - n : dimensionality of signal vector

Output:
    - R : m x n bookkeeping matrix
"""
function make_bookkeeping(ind, n)
    m = length(ind)
    R = zeros(Int, m, n)
    for (i, j) in enumerate(ind)
        R[i,j] = 1
    end
    return R
end

export generate_noisy_measurements, make_connection_matrix, make_bookkeeping, signalfun

end # module SignalRecovery

export SignalRecovery

end  # module Quadratic


using Distributions
using Optim
using HTTP
using GLM
using LinearAlgebra
using Random
using Statistics
using DataFrames
using CSV
using FreqTables

 #:::::::::::::::::::::::::::::::::::::::::::::::::::
 # question 1
 #:::::::::::::::::::::::::::::::::::::::::::::::::::

url = "https://raw.githubusercontent.com/OU-PhD-Econometrics/fall-2021/master/ProblemSets/PS4-mixture/nlsw88t.csv"
df = CSV.read(HTTP.get(url).body,DataFrame)
X = [df.age df.white df.collgrad]
Z = hcat(df.elnwage1, df.elnwage2, df.elnwage3, df.elnwage4,
df.elnwage5, df.elnwage6, df.elnwage7, df.elnwage8)
y = df.occ_code

function mlogit_with_Z(theta, X, Z, y)

    alpha = theta[1:end-1]
    gamma = theta[end]
    K = size(X,2)
    J = length(unique(y))
    N = length(y)
    bigY = zeros(N,J)
    for j=1:J
        bigY[:,j] = y.==j
    end
    bigAlpha = [reshape(alpha,K,J-1) zeros(K)]

    T = promote_type(eltype(X),eltype(theta))
    num   = zeros(T,N,J)
    dem   = zeros(T,N)
    for j=1:J
        num[:,j] = exp.(X*bigAlpha[:,j] .+ (Z[:,j] .- Z[:,J])*gamma)
        dem .+= num[:,j]
    end

    P = num./repeat(dem,1,J)

    loglike = -sum( bigY.*log.(P) )

    return loglike
end
startvals = [2*rand(7*size(X,2)).-1; .1]
td = TwiceDifferentiable(theta -> mlogit_with_Z(theta, X, Z, y), startvals; autodiff = :forward)
# run the optimizer
theta_hat_optim_ad = optimize(td, startvals, LBFGS(), Optim.Options(g_tol = 1e-5, iterations=100_000, show_trace=true, show_every=50))
theta_hat_mle_ad = theta_hat_optim_ad.minimizer
# evaluate the Hessian at the estimates
H  = Optim.hessian!(td, theta_hat_mle_ad)
theta_hat_mle_ad_se = sqrt.(diag(inv(H)))
println([theta_hat_mle_ad theta_hat_mle_ad_se]) # these standard errors match Stata

#:::::::::::::::::::::::::::::::::::::::::::::::::::
# question 2
#:::::::::::::::::::::::::::::::::::::::::::::::::::

#Yes. Estimated gamma in PS3 is -0.09419383447879262. Yet gamma/100 is the change in utility with a 1% increase in expected wage.
#It should not be negative. Here I got a postitive number(1.3074804109794527) so it is more reasonable.

#:::::::::::::::::::::::::::::::::::::::::::::::::::
# question 3
#:::::::::::::::::::::::::::::::::::::::::::::::::::

#(a)
include("lgwt.jl") # make sure the function gets read in
# define distribution
d = Normal(0,1) # mean=0, standard deviation=1
# get quadrature nodes and weights for 7 grid points
nodes, weights = lgwt(7,-4,4)
# now compute the integral over the density and verify it
println(sum(weights.*pdf.(d,nodes)))
# now compute the expectation and verify it
println(sum(weights.*nodes.*pdf.(d,nodes)))

#(b)
#(1)
d = Normal(0,2)
nodes, weights = lgwt(7,-5,5)
println(sum(weights.*(nodes.^2).*pdf.(d,nodes)))
#(2)
d = Normal(0,2)
nodes, weights = lgwt(10,-5,5)
println(sum(weights.*(nodes.^2).*pdf.(d,nodes)))
#(3)
#As the number of quadrature points increase, the quadrature approximates closer to the true value.

#(c)
#(1)
d = Normal(0,2)
a = -5*2
b = 5*2
z = Uniform(a,b)
x = rand(z,1000000)
println((b-a)*mean((x.^2).*pdf.(d,x)))
#(2)
println((b-a)*mean(x.*pdf.(d,x)))
#(3)
println((b-a)*mean(pdf.(d,x)))
#(4)
x = rand(z,1000)
println((b-a)*mean((x.^2).*pdf.(d,x)))
println((b-a)*mean(x.*pdf.(d,x)))
println((b-a)*mean(pdf.(d,x)))
#when D = 1000000, the simulated integral approximates the true value better than that when D = 1000

function mixed_logit_normal(theta, X, Z, y)

    alpha = theta[1:end-2]
    mu_gamma = theta[end-1]
    sigma_gamma = theta[end]
    K = size(X,2)
    J = length(unique(y))
    N = length(y)
    bigY = zeros(N,J)
    for j=1:J
        bigY[:,j] = y.==j
    end
    bigAlpha = [reshape(alpha,K,J-1) zeros(K)]

    T = promote_type(eltype(X),eltype(theta))
    num   = zeros(T,N,J)
    dem   = zeros(T,N)
    d = Normal(mu_gamma,sigma_gamma)
    gammas, weights = lgwt(10,-4,4)
    M = size(gammas,1)
    F = ones(N,M)
    for i =1:M
        for j=1:J
            num[:,j] = exp.(X*bigAlpha[:,j] .+ (Z[:,j] .- Z[:,J]).*gammas[i])
            dem .+= num[:,j]
        end
        P = (num./repeat(dem,1,J))

        for j=1:J
            P[:,j]=P[:,j].^bigY[:,j]
            F[:,i].*=P[:,j]
        end
    end

    loglike = -sum([log(sum(weights.*F[i,:].*pdf.(d,gammas))) for i = 1:N])

    return loglike
end

#:::::::::::::::::::::::::::::::::::::::::::::::::::
# question 4
#:::::::::::::::::::::::::::::::::::::::::::::::::::

startvals = [2*rand(7*size(X,2)).-1; 0; 1]
obj = optimize(alpha->mixed_logit_normal(alpha, X, Z, y), startvals, LBFGS(), Optim.Options(g_tol=1e-5, iterations=100_000, show_trace=true, show_every=50))
beta = obj.minimizer[1:end-2]
mu_gamma = obj.minimizer[end-1]
sigma_gamma = obj.minimizer[end]
println('beta is ',beta)
println('mu_gamma is ',mu_gamma)
println('sigma_gamma is ', sigma_gamma)

#I also tried this code yet MethodError: no method matching Float64(::ForwardDiff.Dual{ForwardDiff.Tag{var"#7#8", Float64}, Float64, 12})
# startvals = [2*rand(7*size(X,2)).-1; 0; 1]
# td = TwiceDifferentiable(theta -> mixed_logit_normal(theta, X, Z, y), startvals; autodiff = :forward)
# theta_hat_optim_ad = optimize(td, startvals, LBFGS(), Optim.Options(g_tol = 1e-5, iterations=100000, show_trace=true, show_every=50))
# theta_hat_mle_ad = theta_hat_optim_ad.minimizer
# H  = Optim.hessian!(td, theta_hat_mle_ad)
# theta_hat_mle_ad_se = sqrt.(diag(inv(H)))
# println([theta_hat_mle_ad theta_hat_mle_ad_se]) # these standard errors match Stata

#:::::::::::::::::::::::::::::::::::::::::::::::::::
# question 5
#:::::::::::::::::::::::::::::::::::::::::::::::::::

function mixed_logit_normal2(theta, X, Z, y)

    alpha = theta[1:end-2]
    mu_gamma = theta[end-1]
    sigma_gamma = theta[end]
    K = size(X,2)
    J = length(unique(y))
    N = length(y)
    bigY = zeros(N,J)
    for j=1:J
        bigY[:,j] = y.==j
    end
    bigAlpha = [reshape(alpha,K,J-1) zeros(K)]

    T = promote_type(eltype(X),eltype(theta))
    num   = zeros(T,N,J)
    dem   = zeros(T,N)
    d = Normal(mu_gamma,sigma_gamma)
    gammas, weights = lgwt(10,-4,4)
    M = size(gammas,1)
    F = ones(N,M)
    for i =1:M
        for j=1:J
            num[:,j] = exp.(X*bigAlpha[:,j] .+ (Z[:,j] .- Z[:,J]).*gammas[i])
            dem .+= num[:,j]
        end
        P = (num./repeat(dem,1,J))

        for j=1:J
            P[:,j]=P[:,j].^bigY[:,j]
            F[:,i].*=P[:,j]
        end
    end

    a = -5*sigma_gamma
    b = 5*sigma_gamma
    z = Uniform(a,b)
    x = rand(z,1000000)
    loglike = -sum([log((b-a)*mean(F[i,:].*pdf.(d,gammas))) for i = 1:N])

    return loglike
end

startvals = [2*rand(7*size(X,2)).-1; 0; 1]
obj = optimize(alpha->mixed_logit_normal2(alpha, X, Z, y), startvals, LBFGS(), Optim.Options(g_tol=1e-5, iterations=100_000, show_trace=true, show_every=50))
beta = obj.minimizer[1:end-2]
mu_gamma = obj.minimizer[end-1]
sigma_gamma = obj.minimizer[end]
println('beta is ',beta)
println('mu_gamma is ',mu_gamma)
println('sigma_gamma is ', sigma_gamma)