#-------------------------------------------------------------------------------
#conversions between spectral units

export ν2f, f2ν, ν2λ, λ2ν, λ2f, f2λ

"""
Convert wavenumber [cm``^{-1}``] to frequency [1/s]
"""
ν2f(ν) = 100.0*𝐜*ν

"""
Convert frequency [1/s] to wavenumber [cm``^{-1}``]
"""
f2ν(f) = f/(100.0*𝐜)

"""
Convert wavenumber [cm``^{-1}``] to wavelength [m]
"""
ν2λ(ν) = 0.01/ν

"""
Convert wavelength [m] to wavenumber [cm``^{-1}``]
"""
λ2ν(λ) = 0.01/λ

"""
Convert wavelength [m] to frequency [1/s]
"""
λ2f(λ) = 𝐜/λ

"""
Convert frequency [1/s] to wavelength [m]
"""
f2λ(f) = f/𝐜

#-------------------------------------------------------------------------------
export planck, normplanck, dplanck, stefanboltzmann, equilibriumtemperature

"""
    planck(ν, T)

Compute black body intensity [W/m``^2``/cm``^{-1}``/sr] using [Planck's law](https://en.wikipedia.org/wiki/Planck%27s_law)

# Arguments
* `ν`: wavenumger [cm``^{-1}``]
* `T`: temperature [Kelvin]
"""
function planck(ν, T)
    νₘ = 100.0*ν #convert from cm^-1 to m^-1
    x = 𝐡*𝐜*νₘ/(𝐤*T) #exponent
    p = 2*𝐡*𝐜^2*νₘ^3 #prefactor
    #result, converting back to cm^-1 with the factor of 100
    100.0*p/(exp(x) - 1.0)
end

"""
    normplanck(ν, T)

Compute black body intensity [W/m``^2``/cm``^{-1}``/sr] using [Planck's law](https://en.wikipedia.org/wiki/Planck%27s_law), normalized by the power emitted per unit area at the given temperature ([`stefanboltzmann`](@ref)),

```
B(ν,T)/σT^4
```

yielding units of 1/cm``^{-1}``/sr.

# Arguments
* `ν`: wavenumger [cm``^{-1}``]
* `T`: temperature [Kelvin]
"""
normplanck(ν, T) = planck(ν, T)/stefanboltzmann(T)

"""
    dplanck(ν, T)

Evaluates ``\\frac{\\partial B}{\\partial T}``
"""
function dplanck(ν, T)
    νₘ = 100.0*ν #convert from cm^-1 to m^-1
    x = 𝐡*𝐜*νₘ/(𝐤*T) #exponent
    y = exp(x)
    if isinf(y)
        return zero(y) #fail accurately, Inf/(Inf - 1)^2 → 0
    end
    p = 2*𝐡^2*𝐜^3*νₘ^4/(𝐤*T^2) #prefactor
    #result, converting back to cm^-1 with the factor of 100
    100.0*p*y/(y - 1)^2
end

"""
    stefanboltzmann(T)

Compute black body radiation power using the [Stefan-Boltzmann](https://en.wikipedia.org/wiki/Stefan%E2%80%93Boltzmann_law) law, ``σT^4`` [W/m``^2``].
"""
stefanboltzmann(T) = 𝛔*(T^4)

"""
    equilibriumtemperature(F, A)

Compute the [planetary equilibrium temperature](https://en.wikipedia.org/wiki/Planetary_equilibrium_temperature), or equivalent blackbody temperature of a planet.

``(\\frac{(1 - A)F}{4\\sigma})^{1/4}``

# Arguments
* `F`: stellar flux [W/m``^2``]
* `A`: albedo
"""
equilibriumtemperature(F, A) = ((1 - A)*F/(4*𝛔))^(1/4)

"""
    equilibriumtemperature(L, A, R)

Compute the [planetary equilibrium temperature](https://en.wikipedia.org/wiki/Planetary_equilibrium_temperature), or equivalent blackbody temperature of a planet.

``(\\frac{(1 - A)L}{16 \\sigma \\pi R^2})^{1/4}``

# Arguments
* `L`: stellar luminosity [W]
* `A`: albedo
* `R`: orbital distance [m]
"""
equilibriumtemperature(L, A, R) = (L*(1 - A)/(16*𝛔*π*R^2))^(1/4)

#-------------------------------------------------------------------------------
export dτdP, transmittance, schwarzschild, absorption, emission

"""
    dτdP(σ, g, μ)

Evaluate the differential increase in optical depth in pressure coordinates, equivalent to the [`schwarzschild`](@ref) equation without Planck emission.

``\\frac{dτ}{dP} = σ\\frac{\\textrm{N}_A}{g μ}``

where ``N_A`` is Avogadro's number.

# Arguments
* `σ`: absorption cross-section [cm``^2``/molecule]
* `g`: gravitational acceleration [m/s``^2``]
* `μ`: mean molar mass [kg/mole]
"""
dτdP(σ, g, μ) = 1e-4*σ*𝐍𝐚/(μ*g)

"""
    transmittance(τ)

Evaluate transmittance from optical depth, ``t = e^{-τ}``
"""
transmittance(τ) = exp(-τ)

"""
    schwarzschild(I, ν, σ, T, P)

Evaluate the [Schwarzschild differential equation](https://en.wikipedia.org/wiki/Schwarzschild%27s_equation_for_radiative_transfer) for radiative transfer with units of length/height [m] and assuming the ideal gas law.

``\\frac{dI}{dz} = σ\\frac{P}{k_B T}[B_ν(T) - I]``

where ``B_ν`` is [`planck`](@ref)'s law.

# Arguments
* `I`: radiative intensity [W/m``^2``/cm``^{-1}``/sr]
* `ν`: radiation wavenumber [cm``^{-1}``]
* `σ`: absorption cross-section [cm``^2``/molecule]
* `T`: temperature [K]
* `P`: pressure [Pa]
"""
schwarzschild(I, ν, σ, T, P) = 1e-4*σ*(P/(𝐤*T))*(planck(ν,T) - I)

"""
    schwarzschild(I, ν, σ, g, μ, T)

Evaluate the [Schwarzschild differential equation](https://en.wikipedia.org/wiki/Schwarzschild%27s_equation_for_radiative_transfer) for radiative transfer with pressure units [Pa] and assuming the ideal gas law.

``\\frac{dI}{dP} = σ\\frac{\\textrm{N}_A}{g μ}[B_ν(T) - I]``

where ``B_ν`` is [`planck`](@ref)'s law and ``N_A`` is Avogadro's number.

# Arguments
* `I`: radiative intensity [W/m``^2``/cm``^{-1}``/sr]
* `ν`: radiation wavenumber [cm``^{-1}``]
* `σ`: absorption cross-section [cm``^2``/molecule]
* `g`: gravitational acceleration [m/s``^2``]
* `μ`: mean molar mass [kg/mole]
* `T`: temperature [K]
"""
schwarzschild(I, ν, σ, g, μ, T) = 1e-4*σ*(𝐍𝐚/(μ*g))*(planck(ν,T) - I)

"""
    absorption(I, σ, g, μ)

Evaluate the [Schwarzschild differential equation](https://en.wikipedia.org/wiki/Schwarzschild%27s_equation_for_radiative_transfer) for radiative transfer with pressure units [Pa] and assuming the ideal gas law, but **without Planck emission**. This can be useful if the wavenumber is far away from the emission region of the temperature and the Planck emission is known to be negligible.

``\\frac{dI}{dP} = -σ\\frac{\\textrm{N}_A}{g μ}I``

where ``N_A`` is Avogadro's number.

# Arguments
* `I`: radiative intensity [W/m``^2``/cm``^{-1}``/sr]
* `σ`: absorption cross-section [cm``^2``/molecule]
* `g`: gravitational acceleration [m/s``^2``]
* `μ`: mean molar mass [kg/mole]
"""
absorption(I, σ, g, μ) = -1e-4*σ*(𝐍𝐚/(μ*g))*I

"""
    emission(I, ν, σ, g, μ, T)

Evaluate the [Schwarzschild differential equation](https://en.wikipedia.org/wiki/Schwarzschild%27s_equation_for_radiative_transfer) for radiative transfer with pressure units [Pa] and assuming the ideal gas law, but **without Planck emission**. This can be useful if the wavenumber is far away from the emission region of the temperature and the Planck emission is known to be negligible.

``\\frac{dI}{dP} = -σ\\frac{\\textrm{N}_A}{g μ}B_ν(T)``

where ``B_ν`` is [`planck`](@ref)'s law

# Arguments
* `I`: radiative intensity [W/m``^2``/cm``^{-1}``/sr]
* `σ`: absorption cross-section [cm``^2``/molecule]
* `g`: gravitational acceleration [m/s``^2``]
* `μ`: mean molar mass [kg/mole]
"""
emission(ν, σ, g, μ, T) = 1e-4*σ*(𝐍𝐚/(μ*g))*planck(ν,T)


using DataFrames
using Distributions
using StatsPlots

function sir(beta,gamma,N,S0,I0,R0,tf)
    t = 0
    S = S0
    I = I0
    R = R0
    ta=DataArray(Float64,0)
    Sa=DataArray(Float64,0)
    Ia=DataArray(Float64,0)
    Ra=DataArray(Float64,0)
    while t < tf
        push!(ta,t)
        push!(Sa,S)
        push!(Ia,I)
        push!(Ra,R)
        pf1 = beta*S*I
        pf2 = gamma*I
        pf = pf1+pf2
        dt = rand(Exponential(1/pf))
        t = t+dt
        if t>tf
            break
        end
        ru = rand()
        if ru<(pf1/pf)
            S=S-1
            I=I+1
        else
            I=I-1
            R=R+1
        end
    end
    results = DataFrame()
    results[:time] = ta
    results[:S] = Sa
    results[:I] = Ia
    results[:R] = Ra
    return(results)
end

srand(42)

sir_out = sir(0.1/1000,0.05,1000,999,1,0,200);

head(sir_out)

# Plot
@df sir_out plot(:time, [:S :I :R], xlabel="Time",ylabel="Number")


"""
    LogTransform <: Transform

Logarithmically transform the data through: sign(x) * log(|x| + 1).

This allows transformations of all real numbers, not just positive ones.
"""
struct LogTransform <: Transform end
cardinality(::LogTransform) = OneToOne()

_logtransform(x) = sign(x) * log(abs(x) + one(x))
_invlogtransform(x) = sign(x) * (exp(sign(x) * x) - one(x))

function _apply(A::AbstractArray, transform::LogTransform; inverse=false, kwargs...)
    inverse && return _invlogtransform.(A)
    return _logtransform.(A)
end

"""
    InverseHyperbolicSine <: Transform

Logarithmically transform the data through: log(x + √(x² + 1)).

This is the inverse hyperbolic sine.
"""
struct InverseHyperbolicSine <: Transform end
cardinality(::InverseHyperbolicSine) = OneToOne()

function _apply(A::AbstractArray, transform::InverseHyperbolicSine; inverse=false, kwargs...)
    inverse && return sinh.(A)
    return asinh.(A)
end


nprocs()==1 && addprocs()

@everywhere logis(x) = 1 ./ (1 .+ exp(-x))
@everywhere dlogis(x) = logis(x) .* (1 .- logis(x))

@everywhere function initParams(x,y,n_hidden,max_iter)
    w1 = randn(n_hidden,size(x,1))
    w2 = randn(n_hidden,n_hidden)
    w3 = randn(size(y,1),n_hidden)
    B1 = rand(n_hidden,1)
    B2 = rand(n_hidden,1)
    errors = zeros(max_iter)
    outputs = zeros(size(y,2),max_iter)
    return (w1,w2,w3,errors,outputs,B1,B2)
end

@everywhere function ffward(x,y,w1,w2,w3)
    a1 = w1 * x
    h1 = logis(a1)
    a2 = w2 * h1
    h2 = logis(a2)
    ay = w3 * h2
    yhat = logis(ay)
    err = yhat - y
    return(a1,h1,a2,h2,ay,yhat,err)
end

@everywhere function updateW!(w1,w2,w3,x,h1,h2,δa1,δa2,err;lr=1)
    w1 .-= δa1 * x' .* lr
    w2 .-= δa2 * h1' .* lr
    w3 .-= err * h2' .* lr
end

@everywhere function bp(x,y,n_hidden,max_iter)
    (w1,w2,w3,errors,outputs,_,_)=initParams(x,y,n_hidden,max_iter)
    for i = 1:max_iter
        (a1,h1,a2,h2,ay,yhat,err) = ffward(x,y,w1,w2,w3)
        δa2 = (w3' * err) .* dlogis(a2)
        δa1 = (w2' * δa2) .* dlogis(a1)
        updateW!(w1,w2,w3,x,h1,h2,δa1,δa2,err)
        errors[i] = sum(abs(err))
        outputs[:,i] = yhat
    end
    return (outputs,errors)
end

@everywhere function fa(x,y,n_hidden,max_iter)
    (w1,w2,w3,errors,outputs,B1,B2)=initParams(x,y,n_hidden,max_iter)
    for i = 1:max_iter
        (a1,h1,a2,h2,ay,yhat,err) = ffward(x,y,w1,w2,w3)
        δa2 = (B2.*err) .* dlogis(a2)
        δa1 = (B1.*δa2) .* dlogis(a1)
        updateW!(w1,w2,w3,x,h1,h2,δa1,δa2,err)
        errors[i] = sum(abs(err))
        outputs[:,i] = yhat
    end
    return (outputs,errors)
end

@everywhere function dfa(x,y,n_hidden,max_iter)
    (w1,w2,w3,errors,outputs,B1,B2)=initParams(x,y,n_hidden,max_iter)
    for i = 1:max_iter
        (a1,h1,a2,h2,ay,yhat,err) = ffward(x,y,w1,w2,w3)
        δa2 = (B2.*err) .* dlogis(a2)
        δa1 = (B1.*err) .* dlogis(a1)
        updateW!(w1,w2,w3,x,h1,h2,δa1,δa2,err)
        errors[i] = sum(abs(err))
        outputs[:,i] = yhat
    end
    return (outputs,errors)
end

@everywhere function ifa(x,y,n_hidden,max_iter)
    (w1,w2,w3,errors,outputs,B1,_)=initParams(x,y,n_hidden,max_iter)
    for i = 1:max_iter
        (a1,h1,a2,h2,ay,yhat,err) = ffward(x,y,w1,w2,w3)
        δa1 = (B1.*err) .* dlogis(a1)
        δa2 = (w2 * δa1) .* dlogis(a2)
        updateW!(w1,w2,w3,x,h1,h2,δa1,δa2,err)
        errors[i] = sum(abs(err))
        outputs[:,i] = yhat
    end
    return (outputs,errors)
end

# Example: XOR problem

x = [1 1 ; 0 1 ; 1 0 ; 0 0]'
y = [0 1 1 0]
n_hidden = 15; # Number of hidden units
max_iter = 500; #Number of learning steps
trials = 30;

dfa_err=@parallel (hcat) for _=1:trials
    (_,dfa_errors) = dfa(x,y,n_hidden,max_iter)
    dfa_errors
end
(dfa_yhat,_) = dfa(x,y,n_hidden,max_iter)

bp_err = @parallel (hcat) for _=1:trials
    (_,bp_errors) = bp(x,y,n_hidden,max_iter)
    bp_errors
end
(bp_yhat,_) = bp(x,y,n_hidden,max_iter)

ifa_err=@parallel (hcat) for _=1:trials
    (_,ifa_errors) = ifa(x,y,n_hidden,max_iter)
    ifa_errors
end
(ifa_yhat,_) = ifa(x,y,n_hidden,max_iter);

fa_err=@parallel (hcat) for _=1:trials
    (_,fa_errors) = fa(x,y,n_hidden,max_iter)
    fa_errors
end
(fa_yhat,_) = fa(x,y,n_hidden,max_iter);

using Plots
l = @layout [
    a b
    c d
    e f
]
gr()
p1=plot(dfa_yhat',label=["1,1" "0,1" "1,0" "0,0"],title="DFA")
p2=plot(ifa_yhat',label=["1,1" "0,1" "1,0" "0,0"],title="IFA")
p3=plot(fa_yhat',label=["1,1" "0,1" "1,0" "0,0"],title="FA")
p4=plot(bp_yhat',label=["1,1" "0,1" "1,0" "0,0"],title="BP")
p5=plot([dfa_err[:,1] ifa_err[:,1] fa_err[:,1] bp_err[:,1]],label=["DFA" "IFA" "FA" "BP"],title="First Trial Error")
p6=plot([mean(dfa_err,2) mean(ifa_err,2) mean(fa_err,2) mean(bp_err,2)],label=["DFA" "IFA" "FA" "BP"],title="Mean Trials Error")
plot(p1,p2,p3,p4,p5,p6,size=(900,900),layout=l)


export proxL1!, proxL1, normL1

"""
    proxL1!(x::Array{T}, λ::Float64; sparseTrafo::Trafo=nothing, kargs...) where T

performs soft-thresholding - i.e. proximal map for the Lasso problem.

# Arguments:
* `x::Array{T}`                 - Vector to apply proximal map to
* `λ::Float64`                  - regularization paramter
* `sparseTrafo::Trafo=nothing`  - sparsifying transform to apply
"""
function proxL1!(x::AbstractArray{Tc}, λ::T; sparseTrafo::Trafo=nothing, kargs...) where {T, Tc <: Union{T, Complex{T}}}
  ε = eps(T)

  if sparseTrafo != nothing
    z = sparseTrafo*x
    z .= max.((abs.(z).-λ),0) .* (z.+ε)./(abs.(z).+ε)
    x .= adjoint(sparseTrafo)*z
  else
    x .= max.((abs.(x).-λ),0) .* (x.+ε)./(abs.(x).+ε)
  end

  return x
end

"""
    normL1(x::Array{T}, λ::Float64; sparseTrafo::Trafo=nothing, kargs...) where T

returns the value of the L1-regularization term.
Arguments are the same as in `proxL1!`
"""
function normL1(x::T, λ::Float64; sparseTrafo::Trafo=nothing, kargs...) where T<:AbstractArray
  if sparseTrafo != nothing
    l1Norm = λ*norm(sparseTrafo*x,1)
  else
    l1Norm = λ*norm(x,1)
  end
  return l1Norm
end


export findβ, microcanonical_ensemble, canonical_ensemble

# stupid implementation of softmax - vals should be regularized beforehand
function _softmax(vals)
    w = exp.(vals)
    w ./= sum(w)
end

"""
    _weights_canonical(β, evals)
Compute the distribution of eigenstates with given energies at temperature 1/β.
P(E) ∝ ∑ᵢexp(-β Eᵢ)

Energy values are assumed to be sorted!
"""
function _weights_canonical(evals, β)
    # assume sorted evals!
    regulator = β < 0 ? evals[end] : evals[1]
    _softmax(@. -β * (evals - regulator))
end

_ΔE(evals, E_0, β) = abs((evals ⋅ _weights_canonical(evals, β)) - E_0)

# TODO: How to optimize best? The function is pretty well behaved...
"""
    findβ(evals, E₀; β_0=0)
Compute the inverse temperature β s. t. the thermal expectation value of the energy matches E₀.
"""
findβ(evals, E_0; β_0 = 0) = Optim.optimize(β -> _ΔE(evals, E_0, β[1]), [float(β_0)], Optim.Newton(linesearch = LineSearches.HagerZhang())).minimizer[1]

"""
    canonical_ensemble(evals, E₀; β₀ = 0)
Compute the weights of the state assuming a canonical ensemble of temperature 1/β. β₀ is taken as starting value for the computation.
A state is assigned the weight ∝ exp(-β Eᵢ)

Energy values are assumed to be sorted!
"""
canonical_ensemble(evals, E_0; β_0 = 0) = _weights_canonical(evals, findβ(evals, E_0; β_0))

"""
    microcanonical_ensemble(evals, E₀[, ΔE])
Compute the weights of the state assuming a microcanonical ensemble. Default size of the window is 0.5% of the spectral width.
A state is in the ensemble if |Eᵢ-E_0| < ΔE.
If the ensemble would be empty, choose the closest state (in energy).

Energy values are assumed to be sorted!
"""
function microcanonical_ensemble(evals, E_0, ΔE)::Vector{Float64}
    micro = @. (evals < E_0+ΔE) & (evals > E_0-ΔE)
    s = sum(micro)
    if s > 0
        # enforce float conversion for type stability
        micro ./ float(s)
    else
        micro = zeros(Float64, length(micro))
        smaller = searchsortedlast(evals, E_0)
        ind = abs(evals[smaller]-E_0) < abs(evals[smaller+1]-E_0) ? smaller : smaller+1
        micro[ind] = 1
        micro
    end
end

microcanonical_ensemble(evals, E_0) = microcanonical_ensemble(evals, E_0, 0.005*(evals[end]-evals[1]))


@generated apply(a,b) = :( apply(a,b,Val( $( hasmethod(*,Tuple{a,b}) ) )) )
apply(a,b,::Val{true}) = a*b
apply(a,b,::Val{false}) = a(b)

"""
    tucker(C,B)

Apply `B[k]` to the `k`th dimension of tensor `C`.

# Examples
```
julia> C = rand(2); B = (rand(3,2),);
       tucker(C,B) == B[1]*C
true

julia> C = rand(2,3); B = (rand(4,2),rand(5,3));
       tucker(C,B) == B[1]*C*B[2]'
true
```
"""
function tucker end

@generated function tucker(
    C::AbstractArray{<:Any,N},
    B::NTuple{N,Any}
) where {N}
    quote
        C_0 = C
        Base.Cartesian.@nexprs $N k->begin
            tmp_k = reshape(C_{k-1},(size(C_{k-1},1),prod(Base.tail(size(C_{k-1})))))
            tmp_k = apply(B[k],tmp_k)
            C_k = Array(reshape(transpose(tmp_k),(Base.tail(size(C_{k-1}))...,size(tmp_k,1))))
            # https://github.com/JuliaLang/julia/issues/30988
        end
        return $(Symbol("C_",N))
    end
end
tucker(C::AbstractArray{<:Any,1}, B::NTuple{1,Any}) = apply(B[1],C)
tucker(C::AbstractArray{<:Any,2}, B::NTuple{2,Any}) = transpose(apply(B[2],transpose(apply(B[1],C))))


# ---
# title: 1227. Airplane Seat Assignment Probability
# id: problem1227
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Math, Dynamic Programming, Brainteaser
# link: <https://leetcode.com/problems/airplane-seat-assignment-probability/description/>
# hidden: true
# ---
#
# `n` passengers board an airplane with exactly `n` seats. The first passenger
# has lost the ticket and picks a seat randomly. But after that, the rest of
# passengers will:
#
#   * Take their own seat if it is still available,
#   * Pick other seats randomly when they find their seat occupied
#
# What is the probability that the n-th person can get his own seat?
#
#
#
# **Example 1:**
#
#
#
#     Input: n = 1
#     Output: 1.00000
#     Explanation: The first person can only get the first seat.
#
# **Example 2:**
#
#
#
#     Input: n = 2
#     Output: 0.50000
#     Explanation: The second person has a probability of 0.5 to get the second seat (when first person gets the first seat).
#
#
#
#
# **Constraints:**
#
#   * `1 <= n <= 10^5`
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


# # Sparse matrices
#
# Source to source automatic differentiation is useful in differentiating sparse matrices. It is a well-known problem that sparse matrix operations can not benefit directly from generic backward rules for dense matrices because general rules do not keep the sparse structure.
# In the following, we will show that reversible AD can differentiate the Frobenius dot product between two sparse matrices with the state-of-the-art performance. Here, the Frobenius dot product is defined as \texttt{trace(A'B)}.
# Its native Julia (irreversible) implementation is `SparseArrays.dot`.
#
# The following is a reversible counterpart

using NiLang, NiLang.AD
using SparseArrays

@i function idot(r::T, A::SparseMatrixCSC{T},B::SparseMatrixCSC{T}) where {T}
    m ← size(A, 1)
    n ← size(A, 2)
    @invcheckoff branch_keeper ← zeros(Bool, 2*m)
    @safe size(B) == (m,n) || throw(DimensionMismatch("matrices must have the same dimensions"))
    @invcheckoff @inbounds for j = 1:n
        ia1 ← A.colptr[j]
        ib1 ← B.colptr[j]
        ia2 ← A.colptr[j+1]
        ib2 ← B.colptr[j+1]
        ia ← ia1
        ib ← ib1
        @inbounds for i=1:ia2-ia1+ib2-ib1-1
            ra ← A.rowval[ia]
            rb ← B.rowval[ib]
            if (ra == rb, ~)
                r += A.nzval[ia]' * B.nzval[ib]
            end
            ## b move -> true, a move -> false
            branch_keeper[i] ⊻= @const ia == ia2-1 || ra > rb
            ra → A.rowval[ia]
            rb → B.rowval[ib]
            if (branch_keeper[i], ~)
                INC(ib)
            else
                INC(ia)
            end
        end
        ~@inbounds for i=1:ia2-ia1+ib2-ib1-1
            ## b move -> true, a move -> false
            branch_keeper[i] ⊻= @const ia == ia2-1 || A.rowval[ia] > B.rowval[ib]
            if (branch_keeper[i], ~)
                INC(ib)
            else
                INC(ia)
            end
        end
    end
    @invcheckoff branch_keeper → zeros(Bool, 2*m)
end

# Here, the key point is using a \texttt{branch\_keeper} vector to cache branch decisions.

# The time used for a native implementation is

using BenchmarkTools
a = sprand(1000, 1000, 0.01);
b = sprand(1000, 1000, 0.01);
@benchmark SparseArrays.dot($a, $b)

# To compute the gradients, we wrap each matrix element with `GVar`, and send them to the reversible backward pass

out! = SparseArrays.dot(a, b)
@benchmark (~idot)($(GVar(out!, 1.0)),
        $(GVar.(a)), $(GVar.(b)))

# The time used for computing backward pass is approximately 1.6 times Julia's native forward pass.
# Here, we have turned off the reversibility check off to achieve better performance.
# By writing sparse matrix multiplication and other sparse matrix operations reversibly,
# we will have a differentiable sparse matrix library with proper performance.

# See my another blog post for [reversible sparse matrix multiplication](https://nextjournal.com/giggle/how-to-write-a-program-differentiably).


plotvector(x, str="k-",label="", alpha=1.0) = plot3D([0;x[1]],[0;x[2]],[0;x[3]], str,alpha=alpha,label=label)

proj(v, x) = vec( dot(v, x) / dot(v, v) * v )

unitvector(x) = x / norm(x)


"""
Created on Mon Apr 13 2020

@author: Yoann Pradat

Plot divergences
"""

include("./src/VNMF.jl")
using .VNMF
using LaTeXStrings
using Plots

#### # Plot β-divergence for different values of β
#### #################################################################################################################

βys = [-1, 0, 1, 2]
βzs = [-0.5, 0.5, 1.5, 2.5]
titles = [L"(a) \beta < 0", L"(b) 0 \leq \beta < 1", L"(c) 1 \leq \beta < 2", L"(d) 2 \leq \beta"]

h = 0.05
p = plot(layout=(2,2), dpi=300, size=(1000, 600))

for sp=1:4
    βy = βys[sp]
    βz = βzs[sp]

    if βz < 1
	X = h:h:8
    else
	X = h:h:4
    end

    Y = [β_divergence(hcat([1.]), hcat([x]), βy) for x in X]
    Z = [β_divergence(hcat([1.]), hcat([x]), βz) for x in X]

    plot!(
	X,
	hcat(Y, Z),
	title = titles[sp],
	labels = ["\$\\beta = $(βy)\$" "\$\\beta = $(βz)\$" ],
	ls    = [:solid :dash],
	lw    = 1,
	color = "black",
	subplot = sp
    )

    plot!(
	xlabel = "",
	xlims  = (0,X[end]),
	xticks = 0:X[end]/4:X[end],
	ylims  = (0,1.5),
	yticks = 0:0.5:1.5,
	subplot = sp
    )
end

folder   = "demo/plot/divergence"
mkpath(folder)
filename = "beta_divergences.pdf"
filepath = joinpath(folder, filename)

savefig(p, filepath)
println("plot saved at $(filepath)")


const where_bisect = 0.49609375

"""
    bisect(X::Interval, α=0.49609375)

Split the interval `X` at position α; α=0.5 corresponds to the midpoint.
Returns a tuple of the new intervals.
"""
function bisect(X::Interval, α=where_bisect)
    @assert 0 ≤ α ≤ 1

    m = mid(X, α)

    return (Interval(X.lo, m), Interval(m, X.hi))
end

"""
    bisect(X::IntervalBox, α=0.49609375)

Bisect the `IntervalBox` `X` at position α ∈ [0,1] along its longest side.
"""
function bisect(X::IntervalBox, α=where_bisect)
    i = argmax(diam.(X))  # find longest side

    return bisect(X, i, α)
end

"""
    bisect(X::IntervalBox, i::Integer, α=0.49609375)

Bisect the `IntervalBox` in side number `i`.
"""
function bisect(X::IntervalBox, i::Integer, α=where_bisect)

    x1, x2 = bisect(X[i], α)

    X1 = setindex(X, x1, i)
    X2 = setindex(X, x2, i)

    return (X1, X2)
end


using PyPlot, Dierckx

cubehermitev1(x) = 2*x^3-3*x^2+1
cubehermitem1(x) = x^3-2*x^2+x
cubehermitev2(x) = -2*x^3+3*x^2
cubehermitem2(x) = (x-1)*x^2

@vectorize_1arg Number cubehermitev1
@vectorize_1arg Number cubehermitem1
@vectorize_1arg Number cubehermitev2
@vectorize_1arg Number cubehermitem2

function _plothermite()
    clf()
    x = linspace(0,1,300)
    plot(x, cubehermitev1(x), label="v1", color="r")
    plot(x, cubehermitem1(x), label="m1", color="g")
    plot(x, cubehermitev2(x), label="v2", color="orange")
    plot(x, cubehermitem2(x), label="m2", color="b")
    legend(loc="center right")
end

"""A callable object that acts as a monotone cubic spline, with a
choice of behaviors at the boundaries."""
type MonotoneSpline
    x      ::Vector{Real}
    y      ::Vector{Real}
    m      ::Vector{Real}
    dx     ::Vector{Real}
    bc     ::String
    function MonotoneSpline(x::Vector, y::Vector, m::Vector,
        dx::Vector; bc="error")
        new(x, y, m, dx, bc)
    end
end

function MonotoneSpline(x::Vector, y::Vector; bc="error")

    n=length(x)
    @assert n==length(y)

    secantslopes = zeros(Float64, n-1)
    for i=1:n-1
        secantslopes[i] = (y[i+1]-y[i])/(x[i+1]-x[i])
    end
    dx = x[2:end] .- x[1:end-1]
    m = zeros(Float64, n)
    m[2:n-1] = 0.5*(secantslopes[2:end] + secantslopes[1:end-1])
    m[1] = secantslopes[1]
    m[n] = secantslopes[n-1]

    for k=1:n-1
        if secantslopes[k] == 0
            m[k] = m[k+1] = 0
        else
            alpha = m[k] / secantslopes[k]
            beta = m[k+1] / secantslopes[k]
            if alpha < 0 || beta < 0
                throw(ArgumentError)
            end
            rsq = alpha^2 + beta^2
            if rsq>9
                tau = sqrt(9.0/rsq)
            else
                tau = 1.0
            end
            m[k] = tau*alpha*secantslopes[k]
            m[k+1] = tau*beta*secantslopes[k]
        end
    end
    return MonotoneSpline(x, y, m, dx; bc=bc)
end

function Base.call(spl::MonotoneSpline, z::Number)
    if z < spl.x[1]
        if spl.bc=="error"
            println("Oh no! 1 spl.bc=$(spl.bc)")
            throw(ArgumentError)
        elseif spl.bc=="constant"
            return spl.y[1]
        else
            return spl.y[1]+spl.m[1]*(z-spl.x[1])
        end
    elseif z > spl.x[end]
        if spl.bc=="error"
            println("Oh no! 2 spl.bc=$(spl.bc)")
            throw(ArgumentError)
        elseif spl.bc=="constant"
            return spl.y[end]
        else
            return spl.y[end]+spl.m[end]*(z-spl.x[end])
        end
    end

    bin = searchsortedfirst(spl.x, z)-1
    if bin<1; bin=1; end
    t = (z-spl.x[bin])/spl.dx[bin]
    return (cubehermitev1(t)*spl.y[bin] + cubehermitem1(t)*spl.m[bin]*spl.dx[bin] +
    cubehermitev2(t)*spl.y[bin+1] + cubehermitem2(t)*spl.m[bin+1]*spl.dx[bin])
end
Base.call(spl::MonotoneSpline, z::AbstractArray) = map(spl, z)

"""A callable object that acts as a monotone cubic spline in log-log space,
with a choice of behaviors at the boundaries."""
type MonotoneSplineLogLog
    x      ::Vector{Real}
    y      ::Vector{Real}
    ms     ::MonotoneSpline
    function MonotoneSplineLogLog(x::Vector, y::Vector; bc="error")
        ms = MonotoneSpline(log(x), log(y); bc=bc)
        new(x, y, ms)
    end
end

Base.call(spl::MonotoneSplineLogLog, z::Number) = exp(spl.ms(log(z)))
Base.call(spl::MonotoneSplineLogLog, z::AbstractArray) = map(spl, z)

"""Return the MonotoneSplineLogLog that has the same knots as c1 but passes
through the control points that result from applying c2 on top of c1."""
function compose_splinelog(c1::MonotoneSplineLogLog, c2::MonotoneSplineLogLog)
    x = c1.x
    y = c2(c1(x))
    MonotoneSplineLogLog(x, y; bc="extrapolate")
end

"""Return the MonotoneSpline that has the same knots as c1 but passes
through the control points that result from applying c2 on top of c1."""
function compose_spline(c1::MonotoneSpline, c2::MonotoneSpline)
    x = c1.x
    y = c2(c1(x))
    MonotoneSpline(x, y; bc="extrapolate")
end

function test1()
    x = collect(0:9)
    y = [0,1,3,10,10,11,12,14,14,16]
    f = MonotoneSpline(x, y; bc="extrapolate")
    f2 = Dierckx.Spline1D(float(x), float(y))
    a = linspace(-2,10,500)
    clf()
    plot(x,y,"ok",label="Input data")
    plot(a, f(a), "r", label="Monotone spline")
    plot(a, evaluate(f2,a), "b", label="Dierckx.Spline1D")
    legend(loc="lower right")
    return f
end

function test2()
    model(x) = 1000*((x./1000).^0.6 ) + x
    x = collect(linspace(0,8000,12))
    y = model(x)
    f = MonotoneSpline(x, y; bc="extrapolate")
    g = MonotoneSplineLogLog(x, y; bc="extrapolate")
    s(z) = evaluate(Dierckx.Spline1D(x, y), z)
    clf()
    a = linspace(0,8000,600)
    subplot(211)
    plot(a, model(a), "k", lw=2)
    plot(a, f(a), "r")
    plot(a, g(a), "b")
    plot(a, s(a), "g")
    plot(x,y,"ok")

    subplot(212)
    plot(a, model(a)-f(a), "r", label="Monotone spline")
    plot(a, model(a)-g(a), "b", label="Log-space mono spline")
    plot(a, model(a)-s(a), "g", label="Dierckx.Spline1D")
    legend()
end


using Revise

includet("src/Spring.jl")
IArray = InternalArray

# 3. Validating the solution --------------------------------------------------------------

const step = 1e-3
const nit = Int(40/step)  # round to avoid small rounding errors

# We define a problem with no friction
prob = SpringProblem(p0 = 0.8, v0 = 0, m = 1, k = 1.5)

x = eulermethod(prob, step, nit)

timeline = [prob.t0 + i*step for i in 1:length(x)]
plot(timeline, IArray(x, 1), label = "standard")

# 5. Finding the frequency of the unforced movement ---------------------------------------

# Empirically, We can see in the plot that the distance between two
# equal values (!= 0) is aproximately 5.
# Let's get the actual distance from our solution.

p0 = 1       # Reference point
P0 = -1      # Next equal point
for i = 50:length(x) # We avoid the first numbers because they are very close
    if norm(x[p0][1]-x[i][1]) < 1e-3
        P0 = i
        break
    end
end

if P0 != -1
    period = timeline[P0]-timeline[p0] # Is similar to what we observed
    freq = 1/period
end

# 6. Using this frequency to force it -----------------------------------------------------

cond_prob = SpringProblem(p0 = 0.8, v0 = 0, m = 1, k = 1.5, A = 1, ω = 2*π*freq)

y = eulermethod(cond_prob, step, nit)
plot!(timeline, IArray(y, 1), label = "conditioned")

savefig("media/spring_frequency_unforced.png")


#-----------------------------------------------------------
# RBF basis calculation
#-----------------------------------------------------------

function calculaterbfbasis!(Φ, x, gridpoints, r=1.0)

    M = length(gridpoints)

    N = length(x)

    for m in 1:M

        for n in 1:N

            @inbounds Φ[n,m] = exp(-norm(gridpoints[m] - x[n])^2 / r^2)

        end

    end

    nothing

end

function calculaterbfbasis(x, gridpoints, r=1.0)

    M = length(gridpoints)

    N = length(x)

    Φ = zeros(N, M)

    calculaterbfbasis!(Φ, x, gridpoints, r)

    return Φ

end


import Base: \
import LinearAlgebra: lu

using Base: OneTo
using LinearAlgebra: givensAlgorithm

"""
    rotate_right!(R, from, to, Q) → nothing

Rotates the diagonal elements between R[from:to,from:to]
such that the eigenvalue at R[to,to] → R[from,from]
and all the others move R[i,i] → R[i+1,i+1].

Assumes that `from` and `to` refer to the first index of the 2×2 block in the
case of complex conjugate pairs of eigenvalues.
"""
function rotate_right!(R, from::Integer, to::Integer, Q = NotWanted())
    i = to

    while i > from
        # Let's just do this check always; sometimes a block might split.
        curr_11 = is_start_of_11_block(R, i)
        prev_11 = is_end_of_11_block(R, i-1)
        j = prev_11 ? i - 1 : i - 2
        swap!(R, j, prev_11, curr_11, Q)
        i = j
    end

    nothing
end

"""
    rotate_left!(R, from, to, Q) → nothing

Rotates the diagonal elements between R[from:to,from:to]
such that the eigenvalue at R[from,from] → R[to,to]
and all the others move R[i,i] → R[i-1,i-1].

Assumes that `from` and `to` refer to the first index of the 2×2 block in the
case of complex conjugate pairs of eigenvalues.
"""
function rotate_left!(R, from::Integer, to::Integer, Q = NotWanted())
    i = from

    while true
        # Let's just do this check always; sometimes a block might split.
        curr_11 = is_start_of_11_block(R, i)
        j = curr_11 ? i + 1 : i + 2
        j > to && break
        next_11 = is_start_of_11_block(R, j)

        swap!(R, i, curr_11, next_11, Q)
        i = next_11 ? i + 1 : i + 2
    end

    nothing
end

struct CompletelyPivotedLU{T,N,TA<:SMatrix{N,N,T},TP}
    A::TA
    p::TP
    q::TP
    singular::Bool
end

struct CompletePivoting end

"""
    lu(A, CompletePivoting) → CompletelyPivotedLU(B, p, q)

Computes the LU factorization of A using complete pivoting.
"""
function lu(A::SMatrix{N,N,T}, ::Type{CompletePivoting}) where {N,T}
    # A, p and q should allocate, but escape analysis will eliminate this!
    if isbitstype(T)
        A = MMatrix(A)
    else
        A = SizedMatrix(A)
    end
    q = @MVector fill(N, N)
    p = @MVector fill(N, N)
    singular = false

    # Maybe I should consider doing this recursively.
    for k = OneTo(N - 1)
        # Find max value in sub-part.
        m, n, maxval = 1, 1, zero(real(T))
        for j = k:N, i = k:N
            if abs(A[i,j]) > maxval
                m, n, maxval = i, j, abs(A[i,j])
            end
        end

        # Store the p and q
        p[k] = m
        q[k] = n

        # Swap row and col
        for j = k:N
            A[k, j], A[m, j] = A[m, j], A[k, j]
        end

        for j = k:N
            A[j, k], A[j, n] = A[j, n], A[j, k]
        end

        Akk = A[k,k]

        # It has actually happened :(
        if iszero(Akk)
            singular = true
            break
        end

        for i = k+1:N
            A[i, k] /= Akk
        end

        for j = k+1:N
            Akj = A[k,j]

            for i = k+1:N
                A[i,j] -= A[i,k] * Akj
            end
        end
    end

    if iszero(A[N,N])
        singular = true
    end

    # Back to immutable land!
    return CompletelyPivotedLU(SMatrix(A), SVector(p), SVector(q), singular)
end

function (\)(LU::CompletelyPivotedLU{T,N}, b::SVector{N}) where {T,N}
    if isbitstype(T)
        x = MVector(b)
    else
        x = SizedVector(b)
    end

    # x ← L \ (P * b)
    for i = OneTo(N)
        x[i], x[LU.p[i]] = x[LU.p[i]], x[i]
        for j = i+1:N
            x[j] -= LU.A[j, i] * x[i]
        end
    end

    # x ← Q * (U \ x)
    for i = N:-1:1
        for j = N:-1:i+1
            x[i] -= LU.A[i, j] * x[j]
        end
        x[i] /= LU.A[i,i]
        x[i], x[LU.q[i]] = x[LU.q[i]], x[i]
    end

    # Go back to immutable land!
    SVector(x)
end

@inline sylvsystem(A::SMatrix{1,1,T}, B::SMatrix{2,2,T}) where {T} =
    @SMatrix [A[1,1]-B[1,1] -B[2,1]       ;
              -B[1,2]        A[1,1]-B[2,2]]

@inline sylvsystem(A::SMatrix{2,2,T}, B::SMatrix{1,1,T}) where {T} =
    @SMatrix [A[1,1]-B[1,1] A[1,2]       ;
              A[2,1]        A[2,2]-B[1,1]]

@inline sylvsystem(A::SMatrix{2,2,T}, B::SMatrix{2,2,T}) where {T} =
    @SMatrix [A[1,1]-B[1,1] A[1,2]        -B[2,1]       T(0)         ;
              A[2,1]        A[2,2]-B[1,1] T(0)          -B[2,1]      ;
              -B[1,2]       T(0)          A[1,1]-B[2,2] A[1,2]       ;
              T(0)          -B[1,2]       A[2,1]        A[2,2]-B[2,2]]

"""
    sylv(A, B, C) → X, singular

Solve A * X - X * B = C for X, where A and B are 1×1 or 2×2 matrices.

It works by recasting the Sylvester equation to a linear system
(I ⊗ A + Bᵀ ⊗ I) vec(X) = vec(C) of size 2 or 4, which is then solved by
Gaussian elimination with complete pivoting.

If the eigenvalues of A and B are equal, then `singular = true`.
"""
@inline function sylv(A::SMatrix{N,N,T}, B::SMatrix{M,M,T}, C::SMatrix{N,M,T}) where {T,N,M}
    fact = lu(sylvsystem(A, B), CompletePivoting)
    rhs = SVector{N*M,T}(C)
    SMatrix{N,M,T}(fact \ rhs), fact.singular
end

"""
    swap22_rotations(X) → c₁, s₁, c₂, s₂, c₃, s₃, c₄, s₄

Construct two sets of two Given's rotations that transform
```
  -x₁₁ -x₁₂
  -x₂₁ -x₂₂
   1    .
   .    1
```
to upper triangular form:
```
  x x  ⇒  x x  ⇒  * *  ⇒  x x  ⇒  x x
  x x  ⇒  * *  ⇒  . *  ⇒  . x  ⇒  . *
  1 .  ⇒  . *  ⇒  . x  ⇒  . *  ⇒  . .
  . 1  ⇒  . 1  ⇒  . 1  ⇒  . .  ⇒  . .
```
"""
function swap22_rotations(X::SMatrix{2,2,T}) where {T}
    # Upper triangulize first column of X
    c₁, s₁, nrm₁ = givensAlgorithm(-X[2,1], T(1))
    c₂, s₂, nrm₂ = givensAlgorithm(-X[1,1], nrm₁)

    # Apply the Givens rotations to the second column of X
    X22 = c₁ * -X[2,2] # + s₁ * 0
    X32 = -s₁' * -X[2,2] # + c₁ * 0
    # X12 = c₂ * -X[1,2] + s₂ * X22 # aint gonna need it!
    X22 = -s₂' * -X[1,2] + c₂ * X22

    # Upper triangularize the second column of X
    c₃, s₃, nrm₃ = givensAlgorithm(X32, T(1))
    c₄, s₄, nrm₄ = givensAlgorithm(X22, nrm₃)

    return c₁, s₁, c₂, s₂, c₃, s₃, c₄, s₄
end

"""
    swap12_rotations(X) → c₁, s₁, c₂, s₂

Construct two Given's rotations that transform
```
  -x₁₁ -x₁₂
   1    .
   .    1
```
to upper triangular form:
```
  x x  ⇒  * *  ⇒  x x
  1 .  ⇒  . *  ⇒  . *
  . 1  ⇒  . 1  ⇒  . .
```
"""
function swap12_rotations(X::SMatrix{1,2,T}) where {T}
    # Upper triangulize first column of X
    c₁, s₁, nrm₁₁ = givensAlgorithm(-X[1,1], T(1))

    # Apply the Givens rotations to the second column of X
    # X12 = c₁₁ * -X[1,2] # + s₁₁ * 0 # ain't gonna need it!
    X22 = -s₁' * -X[1,2] # + c₁₁ * 0

    # Upper triangularize the second column of X
    c₂, s₂, nrm₁₂ = givensAlgorithm(X22, T(1))

    return c₁, s₁, c₂, s₂
end

"""
    swap21_rotations(X) → c₁, s₁, c₂, s₂

Construct two Given's rotations that transform
```
  -x₁₁
  -x₂₁
   1
```
to upper triangular form:
```
  x  ⇒  x  ⇒  x
  x  ⇒  *  ⇒  .
  1  ⇒  .  ⇒  .
```
"""
function swap21_rotations(X::SMatrix{2,1,T}) where {T}
    c₁, s₁, nrm₁ = givensAlgorithm(-X[2,1], T(1))
    c₂, s₂, nrm₂ = givensAlgorithm(-X[1,1], nrm₁)
    return c₁, s₁, c₂, s₂
end

"""
    swap22!(R, i) → R

Swap a 2×2 block with a 2×2 block in R[i:i+4,i:i+4] via unitary transformations.
```
    i
    ↓
i → x x * *  ⇒  y y * *
    x x * *  ⇒  y y * *
    . . y y  ⇒  . . x x
    . . y y  ⇒  . . x x
```
Swapping means similarity transformation.
"""
function swap22!(R::AbstractMatrix{T}, i::Integer, Q = NotWanted()) where {T}
    m, n = size(R)

    @inbounds begin

        # Copy the upper triangular blocks + connection between them.
        A = @SMatrix [R[i+0,i+0] R[i+0,i+1]; R[i+1,i+0] R[i+1,i+1]]
        B = @SMatrix [R[i+2,i+2] R[i+2,i+3]; R[i+3,i+2] R[i+3,i+3]]
        C = @SMatrix [R[i+0,i+2] R[i+0,i+3]; R[i+1,i+2] R[i+1,i+3]]

        # A * X - X * B = C
        X, singular = sylv(A, B, C)

        # No need to swap if eigenvalues are indistinguishable
        singular && return R

        # Rotations that upper triangularize X
        c₁,s₁, c₂,s₂, c₃,s₃, c₄,s₄ = swap22_rotations(X)
        G₁ = Rotation3(c₁, s₁, c₂, s₂, i+0)
        G₂ = Rotation3(c₃, s₃, c₄, s₄, i+1)

        # Apply to R
        lmul!(G₁, R, i, n)
        rmul!(R, G₁, 1, i+3)
        lmul!(G₂, R, i, n)
        rmul!(R, G₂, 1, i+3)

        # Zero out things.
        R[i+2,i+0] = zero(T)
        R[i+3,i+0] = zero(T)
        R[i+2,i+1] = zero(T)
        R[i+3,i+1] = zero(T)

        # Accumulate
        rmul!(Q, G₁)
        rmul!(Q, G₂)
    end

    R
end

"""
    swap21!(R, i) → R

Swap a 1×1 block with a 2×2 block in R[i:i+3,i:i+3] via unitary transformations
```
    i
    ↓
i → x x *  ⇒  y * *
    x x *  ⇒  . x x
    . . y  ⇒  . x x
```
Swapping means similarity transformation.
"""
function swap21!(R::AbstractMatrix{T}, i::Integer, Q = NotWanted()) where {T}
    m, n = size(R)

    @inbounds begin

        # Copy the upper triangular blocks + connection between them.
        A = @SMatrix [R[i+0,i+0] R[i+0,i+1]; R[i+1,i+0] R[i+1,i+1]]
        B = @SMatrix [R[i+2,i+2]]
        C = @SMatrix [R[i+0,i+2]; R[i+1,i+2]]

        # A * X - X * B = C
        X, singular = sylv(A, B, C)

        # No need to swap if eigenvalues are indistinguishable
        singular && return R

        # Rotations that upper triangularize X
        c₁,s₁, c₂,s₂ = swap21_rotations(X)
        G₁ = Rotation3(c₁, s₁, c₂, s₂, i)

        # Apply rotations
        lmul!(G₁, R, i, n)
        rmul!(R, G₁, 1, i+2)

        # Zero out things.
        R[i+1,i+0] = zero(T)
        R[i+2,i+0] = zero(T)

        # Accumulate
        rmul!(Q, G₁)
    end

    R
end

"""
    swap12!(R, i) → R

Swap a 1×1 block with a 2×2 block in R[i:i+3,i:i+3] via unitary transformations
```
    i
    ↓
i → y * *  ⇒  x x *
    . x x  ⇒  x x *
    . x x  ⇒  . . y
```
Swapping means similarity transformation.

TODO: we can optimize this swap to use just one Rotation3 rather than two
Rotation2's.
"""
function swap12!(R::AbstractMatrix{T}, i::Integer, Q = NotWanted()) where {T}
    m, n = size(R)

    @inbounds begin

        # Copy the upper triangular blocks + connection between them.
        A = @SMatrix [R[i+0,i+0]]
        B = @SMatrix [R[i+1,i+1] R[i+1,i+2]; R[i+2,i+1] R[i+2,i+2]]
        C = @SMatrix [R[i+0,i+1] R[i+0,i+2]]

        # A * X - X * B = C
        X, singular = sylv(A, B, C)

        # No need to swap if eigenvalues are indistinguishable
        singular && return R

        # Rotations that upper triangularize X
        c₁,s₁, c₂,s₂ = swap12_rotations(X)
        G₁ = Rotation2(c₁, s₁, i+0)
        G₂ = Rotation2(c₂, s₂, i+1)

        # Apply rotations
        lmul!(G₁, R, i, n)
        rmul!(R, G₁, 1, i+2)
        lmul!(G₂, R, i, n)
        rmul!(R, G₂, 1, i+2)

        # Zero out things.
        R[i+2,i+0] = zero(T)
        R[i+2,i+1] = zero(T)

        # Accumulate
        rmul!(Q, G₁)
        rmul!(Q, G₂)
    end

    R
end

"""
    swap11!(R, i) → R

Swap a 1×1 block with a 1×1 block in R[i:i+1,i:i+1] via unitary transformations
    i
    ↓
i → x *  ⇒  y *
    . y  ⇒  . x
"""
function swap11!(R::AbstractMatrix, i::Integer, Q = NotWanted())
    m, n = size(R)

    @inbounds begin
        R₁₁ = R[i+0,i+0]
        R₁₂ = R[i+0,i+1]
        R₂₂ = R[i+1,i+1]

        # Turns out the Sylvester equation is not so hard to solve in this case.
        G, = get_rotation(R₁₂, R₂₂ - R₁₁, i)

        # Miniscule optimization by not touching R[i:i+1,i:i+1]
        lmul!(G, R, i+2, n)
        rmul!(R, G, 1, i-1)
        R[i+0,i+0] = R₂₂
        R[i+1,i+1] = R₁₁

        # Accumulate
        rmul!(Q, G)
    end

    R
end

"""
    swap!(R, i, curr_11, next_11, Q) → nothing

Swap the two consecutive blocks of the Schur form starting at index i.
"""
function swap!(R::AbstractMatrix, i::Integer, curr_11::Bool, next_11::Bool, Q = NotWanted())
    if curr_11
        if next_11
            swap11!(R, i, Q)
        else
            swap12!(R, i, Q)
        end
    else
        if next_11
            swap21!(R, i, Q)
        else
            swap22!(R, i, Q)
        end
    end
end

@inline is_start_of_11_block(R, i) = i == size(R, 2) || @inbounds(iszero(R[i+1,i]))
@inline is_end_of_11_block(R, i) = i == 1 || @inbounds(iszero(R[i,i-1]))


include("poly_factorization_project.jl")
println("Welcome to the Great Polynomial Factoriser")
println("")
println("")
println("----- Example Polynomials -----")
x = x_poly()
p1 = 2x^3 + 4x^2 - 3x
p2 = -6x^4 - 4x^2 - 3x + 3
println(p1)
println(p2)
println("")
println("mod 7")
x = x_poly(7)
p1 = 2x^5 + 5x - 3
p2 = -6x^4 - 4x^2 - 3x + 3
println(p1)
println(p2)
println("")
println("mod 17")
x = x_poly(17)
p1 = 16x^3 - 3
p2 = -6x^8 - 4x^2 - 3x + 1
println(p1)
println(p2)
println("")
println("mod 101")
x = x_poly(101)
a = 99x^8+ 50x^2 - 3
b = 3x^7+ 5x^7 + 10x^2 + 60x
@show a
@show b

println("")

println("----- Basic Operations -----")
@show a+b
@show a-b
@show a*b
@show a÷b
@show a^3
println("")
println("d/dx(", "$a) = ", derivative(a))
println("")
println("-----Factorisation Example------")
p = (7x^3 + 2x^2 + 8x + 1)*(x^2+x+1)

println("The Factorisation of: ", p," is:")
println("")
factorization = factor(p,true)
println("Here it is:")
pretty_factor(factorization)
pr = expand_factorization(factorization)
println("")
println("Reconstructing: ", pr)

println(" ")
println("---- Extended Euclid Algorithm on x^2 + 50 and 10x^3+10 (mod 101) ----")
egcd = extended_euclid_alg(x^2 + 50,10x^3+10)
pretty_print_egcd((((x^2 + 50)*(10x^3+10)).polynomial,b.polynomial), egcd)

println("")
println("----- Bechmarking -----")
println("Time taken to raise $a to the 1000th power:")
@time l = a^1000
println("")
multbenchZ()
println("")
multbenchZp()
println("")
powbench()
println("")

println(":)")


module geometry

export is_in_pipe, is_in_half_space, is_in_sphere

# Returns distance from cylinder center, squared
# query: DenseVector for which to compute cylindrical norm
# d₀: Direction of the cylinder
# center: Center of cylinder goes through this point
function cylindrical_norm_squared{T}(query::DenseVector{T}, d₀::DenseVector{T}, center::DenseVector{T})
    @assert length(d₀) == length(center) == length(query)
    @assert norm(d₀, 1) > 1e-8
    q₀ = query - center
    q⟂ = q₀ - (q₀⋅d₀)d₀ / (d₀⋅d₀)
    q⟂⋅q⟂
end

# True if query point is inside cylinder
# query: DenseVector for which to compute cylindrical norm
# d₀: Direction of the cylinder
# center: Center of cylinder goes through this point
# radius: of the cylinder
function is_in_pipe{T}(query::DenseVector{T}, d₀::DenseVector{T}, center::DenseVector{T}, radius::T)
    cylindrical_norm_squared(query, d₀, center) <= (radius^2)
end

# True if in half-space defined by normal and offset
# query: DenseVector for which to compute condition
# d₀: Normal to the plane defining the half-space
# origin: Origin of the plane defining the half-space
function is_in_half_space{T}(query::DenseVector{T}, d₀::DenseVector{T}, origin::DenseVector{T})
    @assert length(d₀) == length(origin) == length(query)
    @assert norm(d₀, 1) > 1e-8
    (query - origin)⋅d₀ >= 0
end

# True if in sphere defined by center and radius
# query: DenseVector for which to compute condition
function is_in_sphere{T}(query::DenseVector{T}, center::DenseVector{T}, radius::T)
    @assert length(center) == length(query)
    q₀ = center - query
    q₀⋅q₀ <= radius^2
end

end # module geometry


# function black_formula(optType::OptionType, strike::Float64, forward::Float64, stdDev::Float64,
#                       discount::Float64 = 1.0, displacement::Float64 = 0.0)
#   # TODO check parameters
#   if stdDev == 0.0
#     return max((forward - strike) * value(optType), 0.0) * discount
#   end
#
#   forward = forward + displacement
#   strike = strike + displacement
#
#   if strike == 0.0
#     return isa(optType, Call) ? forward * discount : 0.0
#   end
#
#   d1 = log(forward / strike) / stdDev + 0.5 * stdDev
#   d2 = d1 - stdDev
#   phi = Normal()
#   nd1 = phi(value(optionType) * d1)
#   nd2 = phi(value(optionType) * d2)
#   result = discount * optionType * (forward * nd1 - strike * nd2)
#   result >= 0.0 || error("negative value")
#
#   return result
# end

function black_scholes_theta(process::AbstractBlackScholesProcess, val::Float64, delta_::Float64, gamma_::Float64)
  u = state_variable(process).value
  r = zero_rate(process.riskFreeRate, 0.0, ContinuousCompounding()).rate
  q = zero_rate(process.dividendYield, 0.0, ContinuousCompounding()).rate
  v = local_vol(process.localVolatility, 0.0, u)

  return r * val - (r - q) * u * delta_ - 0.5 * v * v * u * u * gamma_
end

function black_formula_std_dev_derivative(strike::Float64, forward::Float64, stdDev::Float64, expiry::Float64, disc::Float64, displacement::Float64)
  # TOOD check params
  stdDev >= 0.0 || error("stdDev must be non-negative")
  disc > 0.0 || error("discount must be positive")

  forward = forward + displacement
  strike = strike + displacement

  if stdDev == 0.0 || strike == 0.0
    return 0.0
  end

  d1 = log(forward/strike) / stdDev + 0.5 * stdDev

  return disc * forward * distribution_derivative(Normal(), d1)
end

black_formula_vol_derivative(strike::Float64, forward::Float64, stdDev::Float64, expiry::Float64, disc::Float64, displacement::Float64) =
                            black_formula_std_dev_derivative(strike, forward, stdDev, expiry, disc, displacement) * sqrt(expiry)


module Elements
import GaussQuadrature

"""
    lglpoints(::Type{T}, N::Integer) where T <: AbstractFloat

returns the points `r` and weights `w` associated with the `N+1`-point
Gauss-Legendre-Lobatto quadrature rule of type `T`

"""
function lglpoints(::Type{T}, N::Integer) where {T <: AbstractFloat}
    @assert N ≥ 1
    GaussQuadrature.legendre(T, N + 1, GaussQuadrature.both)
end

"""
    lgpoints(::Type{T}, N::Integer) where T <: AbstractFloat

returns the points `r` and weights `w` associated with the `N+1`-point
Gauss-Legendre quadrature rule of type `T`
"""
function lgpoints(::Type{T}, N::Integer) where {T <: AbstractFloat}
    @assert N ≥ 1
    GaussQuadrature.legendre(T, N + 1, GaussQuadrature.neither)
end

"""
    baryweights(r)

returns the barycentric weights associated with the array of points `r`

Reference:
  Jean-Paul Berrut & Lloyd N. Trefethen, "Barycentric Lagrange Interpolation",
  SIAM Review 46 (2004), pp. 501-517.
  <https://doi.org/10.1137/S0036144502417715>
"""
function baryweights(r::AbstractVector{T}) where {T}
    Np = length(r)
    wb = ones(T, Np)

    for j in 1:Np
        for i in 1:Np
            if i != j
                wb[j] = wb[j] * (r[j] - r[i])
            end
        end
        wb[j] = T(1) / wb[j]
    end
    wb
end

"""
    spectralderivative(r::AbstractVector{T},
                       wb=baryweights(r)::AbstractVector{T}) where T

returns the spectral differentiation matrix for a polynomial defined on the
points `r` with associated barycentric weights `wb`

Reference:
  Jean-Paul Berrut & Lloyd N. Trefethen, "Barycentric Lagrange Interpolation",
  SIAM Review 46 (2004), pp. 501-517.
  <https://doi.org/10.1137/S0036144502417715>
"""
function spectralderivative(
    r::AbstractVector{T},
    wb = baryweights(r)::AbstractVector{T},
) where {T}
    Np = length(r)
    @assert Np == length(wb)
    D = zeros(T, Np, Np)

    for k in 1:Np
        for j in 1:Np
            if k == j
                for l in 1:Np
                    if l != k
                        D[j, k] = D[j, k] + T(1) / (r[k] - r[l])
                    end
                end
            else
                D[j, k] = (wb[k] / wb[j]) / (r[j] - r[k])
            end
        end
    end
    D
end

"""
    interpolationmatrix(rsrc::AbstractVector{T}, rdst::AbstractVector{T},
                        wbsrc=baryweights(rsrc)::AbstractVector{T}) where T

returns the polynomial interpolation matrix for interpolating between the points
`rsrc` (with associated barycentric weights `wbsrc`) and `rdst`

Reference:
  Jean-Paul Berrut & Lloyd N. Trefethen, "Barycentric Lagrange Interpolation",
  SIAM Review 46 (2004), pp. 501-517.
  <https://doi.org/10.1137/S0036144502417715>
"""
function interpolationmatrix(
    rsrc::AbstractVector{T},
    rdst::AbstractVector{T},
    wbsrc = baryweights(rsrc)::AbstractVector{T},
) where {T}
    Npdst = length(rdst)
    Npsrc = length(rsrc)
    @assert Npsrc == length(wbsrc)
    I = zeros(T, Npdst, Npsrc)
    for k in 1:Npdst
        for j in 1:Npsrc
            I[k, j] = wbsrc[j] / (rdst[k] - rsrc[j])
            if !isfinite(I[k, j])
                I[k, :] .= T(0)
                I[k, j] = T(1)
                break
            end
        end
        d = sum(I[k, :])
        I[k, :] = I[k, :] / d
    end
    I
end

end # module


using LinearAlgebra
using Makie

"""
λmap(p)
Sends a point on a 3-sphere to a point in the plane x₄=0 with the given point
in the form of 2 complex numbers representing a unit quaternion. This is the
stereographic projection of a 3-sphere. S³ ↦ R³
"""
function λmap(p)
    [real(p[1]), imag(p[1]), real(p[2])] ./ (1 - imag(p[2]))
end

"""
λ⁻¹map(p)
Sends a point on the plane back to a point on a unit sphere with the given
point. This is the inverse stereographic projection of a 3-sphere.
"""
function λ⁻¹map(p)
    x₁ = 2p[1] / (1 + p[1]^2 + p[2]^2 + p[3]^2)
    x₂ = 2p[2] / (1 + p[1]^2 + p[2]^2 + p[3]^2)
    x₃ = 2p[3] / (1 + p[1]^2 + p[2]^2 + p[3]^2)
    x₄ = (-1 + p[1]^2 + p[2]^2 + p[3]^2) / (1 + p[1]^2 + p[2]^2 + p[3]^2)
    [Complex(x₁, x₂), Complex(x₃, x₄)]
end

"""
S¹action(α, p)
Performs a group action corresponding to moving along the circumference of a
circle with the given angle and the point in the form of 2 complex numbers
representing a unit quaternion.
"""
function S¹action(α, p)
    [exp(im * α) * p[1], exp(im * α) * p[2]]
end

"""
convert_to_cartesian(p)
Converts a point in the geographic coordinate system to a point in the
cartesian one with the given point in radians.
"""
function convert_to_cartesian(p)
    [cos(p[2]) * cos(p[1]),
     cos(p[2]) * sin(p[1]),
     sin(p[2])]
end

"""
convert_to_geographic(p)
Converts a point in the cartesian coordinate system to a point in the
geographic one with the given point.
"""
function convert_to_geographic(p)
    r = sqrt(p[1]^2 + p[2]^2 + p[3]^2)
    if p[1] > 0
          ϕ = atan(p[2] / p[1])
    elseif p[2] > 0
          ϕ = atan(p[2] / p[1]) + pi
    else
          ϕ = atan(p[2] / p[1]) - pi
    end
    θ = asin(p[3] / r)
    [ϕ, θ]
end

"""
get_center(A, B, C)

Finds the center point of the fiber circle under stereographic projection
with the given 3 points on the circle circumference.
"""
function get_center(A, B, C)
    a = LinearAlgebra.norm(B - C)
    b = LinearAlgebra.norm(A - C)
    c = LinearAlgebra.norm(A - B)
    numerator = a^2 * (b^2 + c^2 - a^2) * A +
                b^2 * (a^2 + c^2 - b^2) * B +
                c^2 * (a^2 + b^2 - c^2) * C
    denominator = a^2 * (b^2 + c^2 - a^2) +
                  b^2 * (a^2 + c^2 - b^2) +
                  c^2 * (a^2 + b^2 - c^2)
    numerator / denominator
end

"""
get_flower(;N=4, A=.5, B=-pi/7, P=pi/2, Q=0, number=300)

Calculates the x, y and z points of a flower in the base space.
with the given number of petals N, the fattness of the petals A,
the height of the petals B, the latitude of the flower P,
the rotation of the flower Q, and the total number of points in the grid.
"""
function get_flower(;N=4, A=.5, B=-pi/7, P=pi/2, Q=0, number=300)
    N = 6
    A = .5
    B = -pi/7
    P = pi/3
    Q = 0
    t = range(0, stop = 2pi, length = number)
    az = 2pi .* t + A .* cos.(N .* 2pi .* t) .+ Q
    po = B .* sin.(N .* 2pi .* t) .+ P
    x = cos.(az).*sin.(po)
    y = sin.(az).*sin.(po)
    z = cos.(po)
    points = Array{Float64}(undef, number, 2)
    for i in 1:number
        points[i, :] = convert_to_geographic([x[i], y[i], z[i]])
    end
    points
end

"""
build_surface(scene, points, color; transparency, shading)

Builds a surface with the given scene, points, color, transparency and shading.
"""
function build_surface(scene,
                       points,
                       color;
                       transparency = false,
                       shading = true)
    surface!(scene,
             @lift($points[:, :, 1]),
             @lift($points[:, :, 2]),
             @lift($points[:, :, 3]),
             color = color,
             transparency = transparency,
             shading = shading)
end

"""
rotate3D_geographic(point, q)

Rotates a point in the 3D space with the given point and the unit quaternion.
"""
function rotate3D_geographic(point, q)
    c = convert_to_cartesian(point)
    p = Quaternion(c[1], c[2], c[3], 0.0)
    R = conj(q) * p * q
    convert_to_geographic([R[1], R[2], R[3]])
end

"""
rotate3D_cartesian(point, q)

Rotates a point in the 3D space with the given point and the unit quaternion.
"""
function rotate3D_cartesian(point, q)
    p = Quaternion(point..., 0.0)
    R = conj(q) * p * q
    [R[1], R[2], R[3]]
end

"""
get_fiber(point, segments, samples; r=0.025)

Calculates a torus of revolution for building a surface in a specific way with
the given point in the base space, the number of segments, the number of
samples and the radius of the smaller circle in the torus of revolution.
"""
function get_fiber(point, segments, samples; r=0.01)
    # Find 3 points on the circle
    b = λ⁻¹map(convert_to_cartesian(point))
    A = λmap(S¹action(pi / 6, b))
    B = λmap(S¹action(pi / 4, b))
    C = λmap(S¹action(pi / 3, b))
    # The fiber circle center
    Q = get_center(A, B, C)
    # The bigger radius
    R = norm(Q - A)
    # Get the normal to the plane containing the points
    n = LinearAlgebra.cross(A - Q, B - Q)
    n = n / LinearAlgebra.norm(n)
    # The initial normal to the circle
    i = [0.0, 0.0, 1.0]
    # The axis of rotation
    u = LinearAlgebra.cross(n, i)
    u = u / LinearAlgebra.norm(u)
    # The angle of rotation
    angle = acos(LinearAlgebra.dot(n, i)) / 2.0
    q = Quaternion(sin(angle)*u[1],
                   sin(angle)*u[2],
                   sin(angle)*u[3],
                   cos(angle))
    # Construct a torus of revolution grid
    manifold = Array{Float64}(undef, segments, samples, 3)
    for i in 1:segments
        for j in 1:samples
            longitude = i * 2pi / (segments - 1)
            latitude = j * 2pi / (samples - 1)
            x₁ = (Q[1] + (R + r * cos(longitude)) * cos(latitude)) / R
            x₂ = (Q[2] + (R + r * cos(longitude)) * sin(latitude)) / R
            x₃ = (Q[3] + r * sin(longitude)) / R
            manifold[i, j, :] = rotate3D_cartesian([x₁, x₂, x₃], q)
        end
    end
    manifold
end

# The scene object that contains other visual objects
universe = Scene(backgroundcolor = :white, show_axis=false, resolution = (360, 360))
#universe = Scene(backgroundcolor = :white, show_axis=false, resolution = (2560, 1440))

# Calculate a unit quaternion as the rotation axis
u = [sqrt(3)/3, sqrt(3)/3, sqrt(3)/3]
ϕ = Node(0.0)
q = @lift(Quaternion(sin($ϕ)*u[1],
                     sin($ϕ)*u[2],
                     sin($ϕ)*u[3],
                     cos($ϕ)))
segments = 36
samples = 72
number = 360
points = get_flower(number = number)
for i in 1:number
    rotated = @lift(rotate3D_geographic(points[i, :], $q))
    fiber = @lift(get_fiber($rotated, segments, samples))
    color = @lift begin
        x₁, x₂, x₃ = convert_to_cartesian($rotated)
        fill(RGBAf0(rand()/3+2x₁/3, rand()/3+2x₂/3, rand()/3+2x₃/3, 0.9),
             segments,
             samples)
    end
    build_surface(universe, fiber, color, shading = false)
end

# update eye position
eye_position, lookat, upvector = Vec3f0(0.01, 0, 5), Vec3f0(0), Vec3f0(0, 0, 1.0)
update_cam!(universe, eye_position, lookat)
universe.center = false # prevent scene from recentering on display
# Makie.save("gallery/porta.jpg", universe)
frames = 90
record(universe, "gallery/flower.gif") do io
    for i in 1:frames
        ϕ[] = i*2pi/frames # animate scene
        recordframe!(io) # record a new frame
    end
end


#
#   Posituve Definite Karcher Mean (Matlab Manopt)
#
using Manopt, Manifolds, ManifoldsBase, LinearAlgebra, Random
Random.seed!(42)
n = 5
m = 5
M = SymmetricPositiveDefinite(n)
x = random_point(M)
A = [random_point(M) for _ in 1:m]
A = [Symmetric(a) for a in A]
F(X::Array{Float64,2}) = sum([distance(M, X, B)^2 for B in A]) / (2 * m)
∇F(X::Array{Float64,2}) = -sum([log(M, X, B) for B in A]) / m

@time quasi_Newton(
    M,
    F,
    ∇F,
    x;
    memory_size=100,
    stopping_criterion=StopWhenGradientNormLess(norm(M, x, ∇F(x)) * 10^(-6)),
    debug=[:Iteration, " ", :Cost, "\n", 1, :Stop],
)

# B1 = quasi_Newton(M,F,∇F,x; memory_size = 100, debug = [:Iteration, " ", :Cost, "\n", 1, :Stop])
# B2 = mean(M,A)


# load required modules
using CategoricalArrays
using CSV
using DataFrames
using Distributions
using Plots
using Random
using StatsBase

# set seed for reproducible sequence of numbers
Random.seed!(42);

# draw 1000 observations from a normal distribuition
# of mean 100 and standard deviation 10
xnorm = rand(Normal(100, 10), 1000)
histogram(xnorm)

# draw 1000 observations from a uniform distribuition
# of minimum 0 and maximum 10
xunif = rand(Uniform(0, 10), 1000)
histogram(xunif)

# draw 1000 observation from bernoulli distribution with
# probability of success 0.3
xbern = rand(Bernoulli(0.3), 1000)
histogram(xbern)

# draw 1000 observation from binomial distribution with
# probability of success 0.3 and 30 number of trials
xbin = rand(Binomial(30, 0.3), 1000)
histogram(xbin)

# draw 1000 observations from a poisson distribuition
# of average rate of occurrence 100
xpoiss = rand(Poisson(100), 1000)
histogram(xpoiss)


function spindown_convergence_plot(dtlist, l2err, maxerr, sdc_order)
    orderstring = latexstring("\\mathcal{O}(\\Delta t^$sdc_order)")
    h = loglog(dtlist, maxerr, ".-", label=LaTeXString("\$||\\tilde u^{($sdc_order)} - u||_\\infty\$"))
    color = h[1][:get_color]()
    loglog(dtlist, l2err, "*-", label=LaTeXString("\$||\\tilde u^{($sdc_order)} - u||_2\$"), color=color)
    matchidx = 4
    C = ( maxerr[matchidx] + l2err[matchidx] )/2
    @show dtlist_ext = vcat(dtlist[1]*2, dtlist, dtlist[end]/2)
    loglog(dtlist_ext, C*(dtlist_ext/dtlist[matchidx]).^sdc_order, "--", label=orderstring, color=color)
    #axis("equal")
    grid("on")
    #xlim(left=3e-5, right=1e-2)
    axis("tight")
    ylim([1e-13, 1e-2])
    xlabel(L"Time step $\Delta t$")
    ylabel("Relative error")
    legend(bbox_to_anchor=(1.0,1.03), ncol=2)
    tight_layout(0.1)
end

close("all")
figure(figsize=(9,3))
spindown_convergence_plot(dtlist1, l2err1, maxerr1, 1)
spindown_convergence_plot(dtlist2, l2err2, maxerr2, 2)
spindown_convergence_plot(dtlist3, l2err3, maxerr3, 3)
spindown_convergence_plot(dtlist4, l2err4, maxerr4, 4)

savefig("../docs/paper/figs/convergence_time.pdf")


### A Pluto.jl notebook ###
# v0.14.6

using Markdown
using InteractiveUtils

# ╔═╡ 0f6cd296-6cb8-4a9d-a21f-02b689dfe4e9
begin
	using PlutoUI
	PlutoUI.TableOfContents(aside=true)
end

# ╔═╡ 646a9fc2-3a73-4f8d-a952-2cdcf0e6287a
begin
	# Packages
	using LightGraphs
	using GraphPlot
	using Clustering
	using SparseArrays
	using LinearAlgebra
	using Plots
end

# ╔═╡ 7cd3af86-f459-4e4b-8e35-c29bc6966eed
begin
	# Define sizes
	using Random
	m=[200,100,100]
	n=[100,200,100]
	density=[0.5,0.7,0.2]
	Pts=Array{Any}(undef,3)
	Random.seed!(421)
	for i=1:3
	    # Generate sparse random
	    Pts[i]=sprand(m[i],n[i],density[i])
	end
	B=blockdiag(Pts[1],Pts[2],Pts[3])
end

# ╔═╡ 9ab98160-3e37-44c3-b219-a3e1bf95fd68
using Arpack

# ╔═╡ 89299cb7-a4eb-48f3-b3f1-ab2b5c6cfa2e
md"""
# Spectral Partitioning of Bipartite Graphs

Typical example of bipartite graph is a graph obtained from a collection of documents presented as a _term $\times$ document_  matrix.

__Prerequisites__

The reader should be familiar with k-means algorithm and spectral graph partitioning theory and algorithms.

__Competences__

The reader should be able to apply spectral partitioning of bipartite graphs to data clustering problems.

__Credits.__ The notebook was initially derived from M.Sc. Thesis of Ivančica Mirošević.
"""

# ╔═╡ a92d0e70-479d-444e-b06d-e47d2e385b10
md"""
## Definitions

__Undirected bipartite graph__ $G$ is a triplet $G=(T,D,E)$, where $T=\{t_{1},\cdots ,t_{m}\}$ and $D=\{d_{1},...,d_{n}\}$ are two sets of vertices and $E=\{(t_{i},d_{j}):t_{i}\in R,d_{j}\in D\}$, is a set of edges.

 $G$ is __weighted__ if there is weight $\omega(e)$ associated with each edge $e\in E$.

For example, $D$ is a set of documents, $T$ is a set of terms (words) and  edge $e=(t_{i},d_{j})$ exists if document $d_{j}$ contains term $t_{i}$. Weight $\omega(e)$ can be number of appearances of the term $t_i$ in the document $d_j$.

A __term-by-document-matrix__ is a matrix $A\in\mathbb{R}^{m\times n}$ with $A_{ij}=\omega((t_i,d_j))$.
"""

# ╔═╡ fddcc2a7-f718-4900-a135-764eb36496a7
md"""
## Facts

1. The weight matrix of $G$ is $W=\begin{bmatrix}0 & A \\ A^{T} & 0 \end{bmatrix}$.

2. The Laplacian matrix of $G$ is

$$L=\begin{bmatrix} \Delta_{1} & -A \\ -A^{T} & \Delta_{2}\end{bmatrix},$$

where $\Delta_1$ and $\Delta_2$ are diagonal matrices with elements

$\Delta_{1,ii}=\sum\limits_{j=1}^n A_{ij},\quad i=1,\ldots,m,$

$\Delta_{2,jj}=\sum\limits_{i=1}^m A_{ij},\quad j=1,\ldots,n.$

3. The normalized Laplacian matrix of $G$ is

$$L_n=\begin{bmatrix}
I & -\Delta_{1}^{-\frac{1}{2}}A\Delta_{2}^{-\frac{1}{2}} \\
-\Delta_{2}^{-\frac{1}{2}}A^T\Delta_{1}^{-\frac{1}{2}} & I
\end{bmatrix} \equiv
\begin{bmatrix} I & -A_n \\ -A_n^T & I \end{bmatrix}.$$
"""

# ╔═╡ b3b44560-203a-11eb-2782-170e4c12627a
md"""
4. Let $\lambda$ be an eigenvalue of $L_n$ with an eigenvector $w=\begin{bmatrix} u \\ v\end{bmatrix}$, where $u\in \mathbb{R}^{m}$ $v\in\mathbb{R}^{n}$. Then $L_n w=\lambda w$ implies $A_n v =(1-\lambda)u$ and $A_n^T u=(1-\lambda)v$. Vice versa, if $(u,\sigma,v)$ is a singular triplet of $A_n$, then $1-\sigma$ is an eigenvalue of $L_n$ with (non-unit) eigenvector $w=\begin{bmatrix} u \\ v\end{bmatrix}$.

5. The second largest singular value of $A_n$ corresponds to the second smallest eigenvalue of $L_n$, and computing the former is numerically more stable.
"""

# ╔═╡ 1be90f30-203b-11eb-34da-c3116523ec3a
md"

## Algorithms

### Bipartitioning algorithm

1. For given $A$ compute $A_{n}$.
2. Compute singular vectors of $A_{n}$, $u^{[2]}$ and $v^{[2]}$, which correspond to the second largest singular value, $\sigma_2(A_n)$.
3. Assign the partitions $T=\{T_1,T_2\}$ and $D=\{D_1,D_2\}$ according to the signs of $u^{[2]}$ and $v^{[2]}$. The pair $(T,D)$ is now partitioned as $\{(T_1,D_1),(T_2,D_2)\}$.
"

# ╔═╡ 5211ce2e-203b-11eb-0b6f-55e809177fb9
md"
### Recursive bipartitioning algorithm

1. Compute the bipartition $\pi=\{(T_1,D_1),(T_2,D_2)\}$ of $(T,D)$. Set the counter $c=2$.

2. While $c<k$ repeat

   - compute bipartitions of each of the subpartitions of $(T,D)$,
   - among all $(c+1)$-subpartitions, choose the one with the smallest $\mathop{\mathrm{pcut}}(\pi_{c+1})$ or $\mathop{\mathrm{ncut}}(\pi_{c+1})$, respectively.
   - Set $c=c+1$

3. Stop
"

# ╔═╡ 38333a20-203c-11eb-2b28-59e07b3e50ee
md"
### Multipartitioning algorithm
1. For given $A$ compute $A_{n}$.
2. Compute $k$ left and right singular vectors, $u^{[1]},\ldots,u^{[k]}$ and $v^{[1]},\ldots,v^{[k]}$, which correspond to $k$ largest singular values $\sigma_1\geq \cdots \geq \sigma_k$ of $A_n$.
3. Partition the rows of matrices $\Delta_{1}^{-\frac{1}{2}}\begin{bmatrix} u^{[1]} & \ldots & u^{[k]}\end{bmatrix}$ and $\Delta_{2}^{-\frac{1}{2}}\begin{bmatrix} v^{[1]} & \ldots & v^{[k]}\end{bmatrix}$ with the k-means algorithm.
"

# ╔═╡ 1bc20329-3a6e-435d-a25c-8cce79fb7d9d
md"
## Examples

### Small term-by- document matrix
"

# ╔═╡ b9d0830a-0a33-40c1-a08a-49538cf7a879
begin
	# Make a nicer spy function
	import Plots.spy
	spy(A)=heatmap(A,yflip=true,color=:RdBu,aspectratio=1,clim=(-1,1.0))
end

# ╔═╡ f6492dfe-faac-407f-9a0e-5dcf321a4f2d
begin
	# Sources, targets, and weights
	dn=[6,6,7,6,7,7]
	tn=[1,2,2,3,4,5]
	wn=[3,1,3,2,2,3]
	[dn tn wn]
end

# ╔═╡ 1a8b6f7b-4098-4cd2-adad-8ee788b47333
mynames=["Term 1";"Term 2";"Term 3";"Term 4";"Term 5";"Doc 1";"Doc 2"]

# ╔═╡ 4079aa8e-2bb5-4d51-bca8-04d0997db2bf
begin
	G=Graph(7)
	for i=1:length(dn)
	    add_edge!(G,tn[i],dn[i])
	end
	gplot(G, nodelabel=mynames, edgelabel=wn)
end

# ╔═╡ b6d41ab1-24f2-417d-b604-3e12384e65f8
begin
	# We define some functions
	function WeightMatrix(src::Array,dst::Array,weights::Array)
	    n=nv(G)
	    sparse([src;dst],[dst;src],[weights;weights],n,n)
	end

	Laplacian(W::AbstractMatrix)=spdiagm(0=>vec(sum(W,dims=2)))-W

	function NormalizedLaplacian(L::AbstractMatrix)
	    D=1.0./sqrt.(diag(L))
	    n=length(D)
	    [L[i,j]*(D[i]*D[j]) for i=1:n, j=1:n]
	end
end

# ╔═╡ d5d8ee5d-9cce-45fc-a31b-b2ed8237cc86
W=WeightMatrix(tn,dn,wn)

# ╔═╡ 86fd4f5a-af57-490d-a17f-2339a9414057
Matrix(W)

# ╔═╡ 1f53b5cf-ab6a-4107-b53b-9ce96afcdb7f
begin
	L=Laplacian(W)
	Matrix(L)
end

# ╔═╡ 63403411-cca4-4d71-8a76-8dea8638f69f
Lₙ=NormalizedLaplacian(L)

# ╔═╡ 3cf471bc-d316-4de9-af63-71db9b035c40
begin
	A=W[1:5,6:7]
	Δ₁=sqrt.(sum(A,dims=2))
	Δ₂=sqrt.(sum(A,dims=1))
	Aₙ=[A[i,j]/(Δ₁[i]*Δ₂[j]) for i=1:size(A,1), j=1:size(A,2)]
end

# ╔═╡ ac0b5eda-299c-4ae2-ae78-bda92b378f19
# The partitioning - explain the results!
U,σ,V=svd(Aₙ)

# ╔═╡ 9bf6bbca-f96c-4b38-a4fd-5188ec29a2a2
U[:,2]

# ╔═╡ aa51688c-21a0-40f4-a6d6-01d441bfd4b7
V[:,2]

# ╔═╡ 5f27f1a2-80de-41d1-b1e2-da3a99493217
md"
### Sets of points
"

# ╔═╡ e1149b05-d594-4d75-ae56-1ce18bd0f60d
spy(Matrix(B))

# ╔═╡ b95d4070-c127-4db0-987b-5fc17421b667
# The structure of singular vectors reflects the block
S,rest₀=svds(B,nsv=3);

# ╔═╡ 57d472c1-548d-4678-b884-85f863771aa7
# S is a structure
S.S

# ╔═╡ 4806378b-63cb-445a-8c71-583458fce846
begin
	# Plot the first three left singular vectors
	k=size(B,1)
	x=collect(1:k)
	scatter(x,S.U[:,1],title="Left Singular Vectors",label="U[:,1]")
	scatter!(x,S.U[:,2],label="U[:,2]",legend=:topleft)
	scatter!(x,S.U[:,3],label="U[:,3]")
end

# ╔═╡ a08c99a3-9c10-47be-9dc6-987d92a85087
begin
	# Plot the first three right singular vectors
	scatter(x,S.Vt[1,:],title="Right Singular Vectors",label="V[:,1]")
	scatter!(x,S.Vt[2,:],label="V[:,2]")
	scatter!(x,S.Vt[3,:],label="V[:,3]")
end

# ╔═╡ 6888e8a9-f306-4490-9147-2460fbda4889
begin
	# Add random noise
	noise=sprand(k,k,0.3)
	C=B+noise
	spy(Matrix(C))
end

# ╔═╡ 3bd5ad15-31cd-4efd-b4e3-7cf7ede3edae
begin
	# Apply random permutation to rows and columns of C
	D=C[randperm(k),randperm(k)]
	spy(Matrix(D))
end

# ╔═╡ e4487f35-5de0-492e-96f1-69318d7d8e8b
md"""
__Question.__ Given D, can we recover C?

__Answer.__ Yes (with spectral partitioning)!
"""

# ╔═╡ 81d99f24-82b9-4216-8d17-6a18317beb77
Sₙ,rest=svds(D,nsv=3);

# ╔═╡ 8f925ff7-90b4-4fe4-a869-86dec693c0f4
# K-means on rows of U
outU=kmeans(Matrix(transpose(Sₙ.U)),3)

# ╔═╡ d446724a-97d1-403f-baaf-d0619203f351
# K-means on Vt
outV=kmeans(Sₙ.Vt,3)

# ╔═╡ 7572c241-dbcc-4094-97e3-ee3172255294
sortperm(outU.assignments)

# ╔═╡ 22e90c11-ba4e-4073-9f78-f3b7beb24e69
begin
	# RECOVERY of B
	E=D[sortperm(outU.assignments),sortperm(outV.assignments)]
	spy(Matrix(E))
end

# ╔═╡ Cell order:
# ╟─0f6cd296-6cb8-4a9d-a21f-02b689dfe4e9
# ╟─89299cb7-a4eb-48f3-b3f1-ab2b5c6cfa2e
# ╟─a92d0e70-479d-444e-b06d-e47d2e385b10
# ╟─fddcc2a7-f718-4900-a135-764eb36496a7
# ╟─b3b44560-203a-11eb-2782-170e4c12627a
# ╟─1be90f30-203b-11eb-34da-c3116523ec3a
# ╟─5211ce2e-203b-11eb-0b6f-55e809177fb9
# ╟─38333a20-203c-11eb-2b28-59e07b3e50ee
# ╟─1bc20329-3a6e-435d-a25c-8cce79fb7d9d
# ╠═646a9fc2-3a73-4f8d-a952-2cdcf0e6287a
# ╠═b9d0830a-0a33-40c1-a08a-49538cf7a879
# ╠═b6d41ab1-24f2-417d-b604-3e12384e65f8
# ╠═f6492dfe-faac-407f-9a0e-5dcf321a4f2d
# ╠═1a8b6f7b-4098-4cd2-adad-8ee788b47333
# ╠═4079aa8e-2bb5-4d51-bca8-04d0997db2bf
# ╠═d5d8ee5d-9cce-45fc-a31b-b2ed8237cc86
# ╠═86fd4f5a-af57-490d-a17f-2339a9414057
# ╠═1f53b5cf-ab6a-4107-b53b-9ce96afcdb7f
# ╠═63403411-cca4-4d71-8a76-8dea8638f69f
# ╠═3cf471bc-d316-4de9-af63-71db9b035c40
# ╠═ac0b5eda-299c-4ae2-ae78-bda92b378f19
# ╠═9bf6bbca-f96c-4b38-a4fd-5188ec29a2a2
# ╠═aa51688c-21a0-40f4-a6d6-01d441bfd4b7
# ╟─5f27f1a2-80de-41d1-b1e2-da3a99493217
# ╠═7cd3af86-f459-4e4b-8e35-c29bc6966eed
# ╠═e1149b05-d594-4d75-ae56-1ce18bd0f60d
# ╠═9ab98160-3e37-44c3-b219-a3e1bf95fd68
# ╠═b95d4070-c127-4db0-987b-5fc17421b667
# ╠═57d472c1-548d-4678-b884-85f863771aa7
# ╠═4806378b-63cb-445a-8c71-583458fce846
# ╠═a08c99a3-9c10-47be-9dc6-987d92a85087
# ╠═6888e8a9-f306-4490-9147-2460fbda4889
# ╠═3bd5ad15-31cd-4efd-b4e3-7cf7ede3edae
# ╟─e4487f35-5de0-492e-96f1-69318d7d8e8b
# ╠═81d99f24-82b9-4216-8d17-6a18317beb77
# ╠═8f925ff7-90b4-4fe4-a869-86dec693c0f4
# ╠═d446724a-97d1-403f-baaf-d0619203f351
# ╠═7572c241-dbcc-4094-97e3-ee3172255294
# ╠═22e90c11-ba4e-4073-9f78-f3b7beb24e69

src/Module C - Applications/L12 Spectral Partitioning of Bipartite Graphs.jl


#############################################################################
# max.jl
# Return the maximum of the two arguments. Operates elementwise over arrays.
# All expressions and atoms are subtpyes of AbstractExpr.
# Please read expressions.jl first.
#############################################################################
import Base.max
export max, pos, hinge_loss

# TODO: This can easily be extended to work
### Max Atom
struct MaxAtom <: AbstractExpr
    head::Symbol
    id_hash::UInt64
    children::Tuple{AbstractExpr, AbstractExpr}
    size::Tuple{Int, Int}

    function MaxAtom(x::AbstractExpr, y::AbstractExpr)
        if sign(x) == ComplexSign() || sign(y) == ComplexSign()
            error("Both the arguments should be real instead they are $(sign(x)) and $(sign(y))")
        else
            if x.size == y.size
                sz = x.size
            elseif x.size == (1, 1)
                sz = y.size
            elseif y.size == (1, 1)
                sz = x.size
            else
                error("Got different sizes for x as $(x.size) and y as $(y.size)")
            end
        end

        children = (x, y)
        return new(:max, hash(children), children, sz)
    end
end

function sign(x::MaxAtom)
    sign_one = sign(x.children[1])
    sign_two = sign(x.children[2])
    if sign_one == Positive() || sign_two == Positive()
        return Positive()
    elseif sign_one == Negative() && sign_two == Negative()
        return Negative()
    else
        return sign_one + sign_two
    end
end

# The monotonicity
function monotonicity(x::MaxAtom)
    return (Nondecreasing(), Nondecreasing())
end

# If we have h(x) = f o g(x), the chain rule says h''(x) = g'(x)^T f''(g(x))g'(x) + f'(g(x))g''(x);
# this represents the first term
function curvature(x::MaxAtom)
    return ConvexVexity()
end

function evaluate(x::MaxAtom)
    return max.(evaluate(x.children[1]), evaluate(x.children[2]))
end

# x <= this and y <= this if max(x, y) = this
function conic_form!(x::MaxAtom, unique_conic_forms::UniqueConicForms=UniqueConicForms())
    if !has_conic_form(unique_conic_forms, x)
        this = Variable(x.size[1], x.size[2])
        objective = conic_form!(this, unique_conic_forms)
        for child in x.children
            conic_form!(this >= child, unique_conic_forms)
        end
        cache_conic_form!(unique_conic_forms, x, objective)
    end
    return get_conic_form(unique_conic_forms, x)
end

max(x::AbstractExpr, y::AbstractExpr) = MaxAtom(x, y)
max(x::AbstractExpr, y::Value) = max(x, Constant(y))
max(x::Value, y::AbstractExpr) = max(Constant(x), y)
pos(x::AbstractExpr) = max(x, Constant(0, Positive()))
hinge_loss(x::AbstractExpr) = pos(1 - x)


function f_2by2!(F, x)
    F[1] = (x[1]+3)*(x[2]^3-7)+18
    F[2] = sin(x[2]*exp(x[1])-1)
end

function g_2by2!(J, x)
    J[1, 1] = x[2]^3-7
    J[1, 2] = 3*x[2]^2*(x[1]+3)
    u = exp(x[1])*cos(x[2]*exp(x[1])-1)
    J[2, 1] = x[2]*u
    J[2, 2] = u
end

function nlsolvers_f(x, F, J)
	if F !== nothing
		f_2by2!(F, x)
	end
	if J !== nothing
		g_2by2!(J, x)
	end
	objective_return(F, J)
end

nlsolve!(OnceDiffed(nlsolvers_f), [ -0.5; 1.4], TrustRegion(NLSolvers.Newton(), Dogleg()))
nlsolve!(OnceDiffed(nlsolvers_f), [ -0.5; 1.4], TrustRegion(NLSolvers.Newton()))
nlsolve!(OnceDiffed(nlsolvers_f), [ -0.5; 1.4], TrustRegion(NLSolvers.Newton(), NWI()))
# initial convergence
nlsolve!(OnceDiffed(nlsolvers_f), [ 0.0; 1.0], TrustRegion(NLSolvers.Newton(), Dogleg()))
nlsolve!(OnceDiffed(nlsolvers_f), [ 0.0; 1.0], TrustRegion(NLSolvers.Newton()))
nlsolve!(OnceDiffed(nlsolvers_f), [ 0.0; 1.0], TrustRegion(NLSolvers.Newton(), NWI()))


mutable struct BSpline
	degree   :: Int
	numKnots :: Int
	lower    :: Float64
	upper    :: Float64
	knots    :: Array{Float64,1}

	# constructor with equally spaced interior knots
	function BSpline(nKnots,deg,lb,ub)
		knots = zeros(nKnots+2*deg)
		if (nKnots < 2*(deg+1)-1)
			throw(ArgumentError("need at least 2*(deg+1) - 1 = $(2*(deg+1) - 1) knots"))
		end
		if deg < 0
			throw(ArgumentError("degree must be non-negative integer"))
		end
		if lb > ub
			throw(ArgumentError("require lb < ub"))
		end

		h = (ub - lb) / (nKnots - 1)
		for i=1:length(knots)
			if i < deg + 1
				knots[i] = lb
			elseif i > nKnots + deg
				knots[i] = ub
			else
				knots[i] = lb + (i-deg-1)*h
			end
		end
		# sp = sparsevec(zeros(num_coefs))
		new(deg,nKnots,lb,ub,knots)
	end

	# constructor with user supplied interior knots
	function BSpline(knots::Array,deg::Integer)
		if !issorted(knots)
			throw(ArgumentError("knots must be sorted"))
		end
		lb = knots[1]
		ub = knots[end]
		# extend knot vector by degree on both sides
		knots_new = [ [lb for i=1:deg]; knots; [ub for i=1:deg] ]
		numKnots = length(knots)	# number of INTERIOR knots
		new(deg,numKnots,lb,ub,knots_new)
	end

end

function show(io::IO, b::BSpline)
	print(io,"BSpline object with\n")
	print(io,"degree: $(b.degree)\n")
	print(io,"number of knots: $(b.numKnots)\n")
	print(io,"[lower,upper]: [$(b.lower),$(b.upper)]\n")
	print(io,"knot vector: $(b.knots)\n")
end

function getNumKnots(b::BSpline) return b.numKnots end
function getNumCoefs(b::BSpline) return b.numKnots + b.degree -1 end
function getCoefs(b::BSpline) return b.knots[ (b.degree+1):(length(b.knots)-b.degree)] end
function getDegree(b::BSpline) return b.degree end

# """
# 	compute value of the ``d + 1`` B-splines at ``x``

# returns

# ```math
# b = \right(B_{\mu-d,d}(x),\dots,B_{\mu,d}(x)\left)^T
# ```
# """
# function getBasis2(x::Float64,b::BSpline)

# 	num_nodes = getNumCoefs(b)
# 	deg       = b.degree

# 	B = zeros(num_nodes)  # Bspline

# end

#
function getBasis(x::Float64,b::BSpline)

	num_nodes = getNumCoefs(b)
	deg       = b.degree

	# tmp
	d = 0.0
	e = 0.0

	# create a basis function
	# return a colvec because currntly only CSC format
	# a 1-row sparse matrix is dense
	# bs = spzeros(num_nodes,1)
	bs = spzeros(num_nodes,1)

	# check x
	if x < b.lower
		@warn("x < lb: x=$x, lb=$(b.lower). be careful!")
	elseif x > b.upper
		@warn("x > ub: x=$x, ub=$(b.upper). be careful!")
	end

	# get mu s.t. knot_mu < knot_{mu+1} and x in [knot_mu, knot_{mu+1})
	# i.e. get the index of the lower knot in the active knot span

	# fix bound behaviour
	if x <= b.lower
		mu = deg+1
	elseif x>=b.upper
		mu = num_nodes
	else
		mu = searchsortedlast(b.knots,x)
	end

	# set 0-degree basis function
	# 0-deg basis is an indicator function
	# that is 1.0 in the active knot span and 0.0 else.
	bs[mu] = 1.0

	# loop over degrees
	for k=1:deg

		# loop over basis functions
		for j in mu-k:mu

			# take care of "division by zero" issue
			# dividing by zero must return 0.0
			if j+k <= deg +1
				d = 0.0
			elseif b.knots[j+k]-b.knots[j] == 0
				d = 0.0
			else
				d = bs[j] * (x - b.knots[j]) / (b.knots[j+k]-b.knots[j])
			end

			if j+1 >= num_nodes+1
				e = 0.0
			elseif b.knots[j+k+1]-b.knots[j+1] == 0
				e = 0.0
			else
				e = bs[j+1] * (b.knots[j+k+1] - x)/(b.knots[j+k+1]-b.knots[j+1])
			end

			bs[j] = d + e

		end

	end
	return bs
end

# vector of points
function getBasis(x::Vector{Float64},b::BSpline)

	n = length(x)
	num_nodes = getNumCoefs(b)
	deg       = b.degree

	# if you evaluate a degree 1 basis at its interior knots ("coefs"), you get identity
	if (deg == 1) && (n == num_nodes) && (sum(abs.(x .- getCoefs(b))) < 1e-8)
		return sparse(I*1.0,n,n)
	else
	# sort x?

		# tmp
		d = 0.0
		e = 0.0

		# create a basis function
		# return a colvec because currntly only CSC format
		# a 1-row sparse matrix is dense
		bs = spzeros(n,num_nodes)

		for xi in 1:n

			# check x
			if x[xi] < b.lower
				@warn("x < lb: x=$(x[xi]), lb=$(b.lower). be careful!")
			elseif x[xi] > b.upper
				@warn("x > ub: x=$(x[xi]), ub=$(b.upper). be careful!")
			end

			# get mu s.t. knot_mu < knot_{mu+1} and x in [knot_mu, knot_{mu+1})
			# i.e. get the index of the lower knot in the active knot span

			# fix bound behaviour
			if x[xi] <= b.lower
				mu = deg+1
			elseif x[xi]>=b.upper
				mu = num_nodes
			else
				mu = searchsortedlast(b.knots,x[xi])
			end

			# set 0-degree basis function
			# 0-deg basis is an indicator function
			# that is 1.0 in the active knot span and 0.0 else.
			bs[xi,mu] = 1.0

			# loop over degrees
			for k=1:deg

				# loop over basis functions
				for j in mu-k:mu

					# take care of "division by zero" issue
					# dividing by zero must return 0.0
					if j+k <= deg +1
						d = 0.0
					elseif b.knots[j+k]-b.knots[j] == 0
						d = 0.0
					else
						d = bs[xi,j] * (x[xi] - b.knots[j]) / (b.knots[j+k]-b.knots[j])
					end

					if j+1 >= num_nodes+1
						e = 0.0
					elseif b.knots[j+k+1]-b.knots[j+1] == 0
						e = 0.0
					else
						e = bs[xi,j+1] * (b.knots[j+k+1] - x[xi])/(b.knots[j+k+1]-b.knots[j+1])
					end
					bs[xi,j] = d + e
				end

			end
		end
		return bs

	end

end


export IFS

"""
`IFS` is an iterated function system.
"""
immutable IFS
  funcs::Vector{AffineMap}
end

function IFS(args...)
  fun_list = [ f for f in args ]
  IFS(fun_list)
end

# Apply F to a single vector
function (F::IFS)(x::Vector)
  y = [ f(x) for f in F.funcs ]
  return Set(y)
end

# Apply F to a list of vectors
function (F::IFS)(xlist::Set)
  A = Set{Vector{Float64}}()
  for x in xlist
    A = union(A,F(x))
  end
  return A
end

function show(io::IO,F::IFS)
  nf = size(F.funcs)[1]
  for k=1:nf
    println(io,"#$k:\t$(F.funcs[k])")
  end
end

"""
`square_check(F)` checks if all functions `f` in the `IFS` pass
`square_check(f)`.
"""
function square_check(F::IFS)
  return all(square_check(f) for f in F.funcs)
end

"""
`Triangle_IFS()` returns an `IFS` that creates a Triangle triangle.
"""
function Triangle_IFS()
  A = 0.5 * eye(2)
  f1 = AffineMap(A,[0,0])
  f2 = AffineMap(A,[0.5,0])
  f3 = AffineMap(A,[0.25,0.5])
  return IFS(f1,f2,f3)
end

"""
`Carpet_IFS()` returns an `IFS` for Triangle's Carpet.
"""
function Carpet_IFS()
  A = (1/3)*eye(2)
  f1 = AffineMap(A,[0,0])
  f2 = AffineMap(A,[0,1/3])
  f3 = AffineMap(A,[0,2/3])
  g1 = AffineMap(A,[1/3,0])
  g3 = AffineMap(A,[1/3,2/3])
  h1 = AffineMap(A,[2/3,0])
  h2 = AffineMap(A,[2/3,1/3])
  h3 = AffineMap(A,[2/3,2/3])
  return IFS(f1,f2,f3,g1,g3,h1,h2,h3)
end

"""
`visualize(F)` for an `IFS` draws the unit square (dotted)
and the image of the functions in `F` (solid frames).
"""
function visualize(F::IFS)
  clf()
  for f in F.funcs
    if !square_check(f)
      warn("$f does not map unit square to itself.")
    end
    visualize(f)
  end
end


struct Pos
    x::Float64
    y::Float64
end

function jarvis_cross(point1::Pos, point2::Pos, point3::Pos)
    vec1 = Pos(point2.x - point1.x, point2.y - point1.y)
    vec2 = Pos(point3.x - point2.x, point3.y - point2.y)
    ret_cross = vec1.x*vec2.y - vec1.y*vec2.x
    return ret_cross*ret_cross
end

function jarvis_march(points::Vector{Pos})
    hull = Vector{Pos}()

    # sorting array based on leftmost point
    sort!(points, by = item -> item.x)
    push!(hull, points[1])

    i = 1
    curr_point = points[2]

    # Find cross product between points
    curr_product = jarvis_cross(Pos(0,0), hull[1], curr_point)
    while (curr_point != hull[1])
        for point in points
                product = 0.0
            if (i == 1)
                if (hull[i] != point)
                    product = jarvis_cross(Pos(0,0), hull[i], point)
                end
            else
                if (hull[i] != point && hull[i-1] != point)
                    product = jarvis_cross(hull[i-1], hull[i], point)
                end
            end
            if (product > curr_product)
                curr_point = point
                curr_product = product
            end
        end
        push!(hull, curr_point)
        curr_product = 0
        i += 1
    end

    return hull
end

function main()

    points = [Pos(2,1.5), Pos(1, 1), Pos(2, 4), Pos(3, 1)]
    hull = jarvis_march(points)
    println(hull)
end

main()


"""
    ElementVectorSampler

A sampler that makes it easy to draw signed element vectors (of a specified dimension)
uniformly at random.

Notice that, w.p. `1/2`, each element vector is multiplied by `-1`. This is useful for
implementing random walks.

An `element vector` `e_i` is a vector in which the `i`th element is a `1`, and all other
elements are `0`.

##### Fields
- `dimension::Int` : the dimension (i.e. length) of element vectors.
"""
struct ElementVectorSampler <: Sampleable{Multivariate, Discrete}
    dimension::Int
end

"""
    EVecSampler

Alias for `ElementVectorSampler`
"""
EVecSampler = ElementVectorSampler

"""
    Base.show(io::IO, elementvectorsampler::ElementVectorSampler) -> nothing

Show an `ElementVectorSampler` in a human-friendly manner.
"""
function Base.show(io::IO, elementvectorsampler::ElementVectorSampler)
    println(io, typeof(elementvectorsampler))
    println(io, " dimension: ", elementvectorsampler.dimension)
    return nothing
end

"""
    Base.rand(rng::AbstractRNG, elementvectorsampler::ElementVectorSampler) -> Vector{Int}

Return a vector with a specified `length` and with exactly one element equal to `1`.
"""
function Base.rand(rng::AbstractRNG, elementvectorsampler::ElementVectorSampler)
    vector = zeros(Int, elementvectorsampler.dimension)
    # pick an index to replace with 1 or -1 (uses multiple dispatch)
    vector[rand(1:elementvectorsampler.dimension)] = 1*rand((-1,1))
    return vector
end