test = """939
7,13,x,x,59,x,31,19
"""

function parse_input(input::AbstractString)::Tuple
    (departure, buses) = split(input, r"\s+"; keepempty = false)
    buses = tryparse.(Int, split(buses, ','))
    return (parse(Int, departure), buses)
end

function q1(departure::Integer, buses::AbstractVector)::Int
    buses = filter(!isnothing, buses)
    bus = buses[argmin([bus - departure % bus for bus in buses])]
    return (bus - departure % bus) * bus
end

function q2(buses::AbstractVector)::Int
    return mod_lcm([(bus, rem(1 - i, bus, RoundDown))
                    for (i, bus) in enumerate(buses)
                    if bus !== nothing])
end

function mod_lcm(vs::AbstractVector)::Int
    isempty(vs) && return 0
    ((v1, m1), rest...) = vs
    r = map(rest) do (v, m)
        vp = v ÷ gcd(v, v1)
        k = findfirst(==(m), [(k * v1 + m1) % v for k in 0:vp - 1]) - 1
        return (vp, k)
    end |> mod_lcm
    return r * v1 + m1
end

let
    (departure, buses) = parse_input(test)
    @assert q1(departure, buses) == 295
    @assert q2(buses) == 1068781
end

(departure, buses) = parse_input(read("day13.in", String))
println("Q1: ", q1(departure, buses))
println("Q1: ", q2(buses))


using TuringModels, StatsFuns

Turing.setadbackend(:reversediff)

# outcome and predictor almost perfectly associated

x = repeat([-1], 9); append!(x, repeat([1],11))
y = repeat([0], 10); append!(y, repeat([1],10))

@model m_good_stan(x,y) = begin
    α ~ Normal(0,10)
    β ~ Normal(0,10)

    logits = α .+ β * x

    y .~ BinomialLogit.(1, logits)
end

chns = sample(m_good_stan(x,y), Turing.NUTS(0.65), 1000)

# Stan results

m_10_x,_results = "
    mean   sd   5.5% 94.5% n_eff Rhat
 a -5.09 4.08 -12.62 -0.25   100 1.00
 b  7.86 4.09   2.96 15.75   104 1.01
";

# Look at the proper draws (in corrected chn2)

describe(chns)

# End of 10/m10.xxt.jl


using Plots

V(x,y) = (1/2)*(x^2 + y^2 + 2*x^2*y - (2/3)*y^3)
x = y = range(-1,stop = 1, length = 1000)
V(x,y) = (1/2)*(x^2 + y^2 + 2*x^2*y - (2/3)*y^3)
x = y = range(-1,stop = 1, length = 1000)
contourf(x,y,V)
contour!(x,y,V,levels=0:1/6:1,title = "Henon-Heiles potential: Intensity map",linecolor = :white)
# savefig("potential_intensity.png")

plot(x,y,V,title = "Henon-Heiles potential: Surface map")
# savefig("potential_surface.png")


using DynamicHMCModels

Random.seed!(1)

cd(@__DIR__)
include("simulateGaussian.jl")

struct GaussianProb{TY <: AbstractVector}
   "Observations."
   y::TY
end

function (problem::GaussianProb)(θ)
   @unpack y = problem   # extract the data
   @unpack mu, sigma = θ
   loglikelihood(Normal(mu, sigma), y) + logpdf(Normal(0,1), mu) +
   logpdf(Truncated(Cauchy(0,5),0,Inf), sigma)
end

# Define problem with data and inits.

data = simulateGaussian(;Nd=100)
p = GaussianProb(data.y);
p((mu = 0.0, sigma = 1.0))

# Write a function to return properly dimensioned transformation.

problem_transformation(p::GaussianProb) =
  as((mu  = as(Real, -25, 25), sigma = asℝ₊), )

# Use Flux for the gradient.
P = TransformedLogDensity(problem_transformation(p), p)

∇P = ADgradient(:ForwardDiff, P)

#import Zygote
#∇P = ADgradient(:Zygote, P)

# Sample from the posterior.
chain, NUTS_tuned = NUTS_init_tune_mcmc(∇P, 4000);

# Undo the transformation to obtain the posterior from the chain.
posterior = TransformVariables.transform.(Ref(problem_transformation(p)),
get_position.(chain));

chns = nptochain(posterior, NUTS_tuned)

describe(chns)


include("../src/autodiff_module.jl")
include("../../__lib__/math/opti/src/opti_module.jl")

module optitest
using ..autodiff
using ..opti

function quadratic(x)
    y = zero(eltype(x))
    for ii in eachindex(x)
        y += x[ii]^2
    end
    return y
end

function quadratic(x, x0)
    y = zero(eltype(x))
    for ii in eachindex(x)
        y += (x[ii] - x0[ii])^2
    end
    return y
end

# x1a = autodiff.DualNumber(3.0)
# x2a = autodiff.DualNumber(4.0)

# x_vec = [x1a, x2a]
# g = autodiff.gradient(quadratic, x_vec)

x0 = [3.0, 4.0]

c = [5.0, 6.0]
obj_fcn(x)              = quadratic(x, c)
grad_fcn(objfcn, x)     = autodiff.gradient(autodiff.ForwardMode(), objfcn, x)
dir_diff_fcn(fcn, x, s) = autodiff.directional_derivative(fcn, x, s)

tol_options      = opti.ToleranceOptions(
    step_size_tol = 1e-15,
    f_abs_tol = 1e-15,
    f_rel_tol = 1e-15,
    g_abs_tol = 1e-15,
    max_iter = 100)

# (x_sol, y_sol, iter, stopping_crit) = opti.quasi_newton(
#     obj_fcn,
#     x0;
#     tol_options          = tol_options,
#     gradient_fcn         = grad_fcn,
#     directional_diff_fcn = dir_diff_fcn)

stat = opti.cg_pr(
    obj_fcn,
    x0;
    tol_options          = tol_options,
    gradient_fcn         = grad_fcn,
    directional_diff_fcn = dir_diff_fcn)

# directional diff
# lim[h -> 0] (fcn(x + h * s) - fcn(x)) / h

# x0 = [3.0, 4.0]
# obj_fcn(x) = quadratic(x)

# # s = [-4.0, 3.0]
# # s = -[-3.0, 4.0]
# s = [0.0, 5.0]
# s = [5.0, 0.0]

# dg = autodiff.directional_derivative(obj_fcn, x0, s)

end


"""
        capacity(P::Drip;      # Drip structure
                 unit = "bit"  # optional: unit of capacity (bit or nat).
                 )
Returns the amount of information processes per unit of time in the steady state
of the DRIP `P`.
"""
function capacity(P::Drip;      # Drip structure
                  unit = "bit"  # optional: unit of capacity (bit or nat).
                  )
    if unit == "bit"
        κ = 0.5*log(det(P.ss.Σ_1)/det(P.ss.Σ_p))/log(2); #returns capacity in bits
    elseif unit == "nat"
        κ = 0.5*log(det(P.ss.Σ_1)/det(P.ss.Σ_p));        #returns capacity in nats
    else
        println("Invalid input for unit! Capacity is reported in bits.")
        κ = 0.5*log(det(P.ss.Σ_1)/det(P.ss.Σ_p))/log(2);
    end
    return(κ)
end

"""
    infinitesum(func; tol = 1e-6,maxit = 1000,start=0)
Returns the infinite sum `Σₓfunc(x)` starting from `x = start` up to tolderance
`tol` or max iteration `maxit`.
"""
function infinitesum(func; tol = 1e-6,maxit = 1000,start=0)
    diff  = 1.0
    infsum = func(start)
    it    = start + 1
    while (diff > tol) & (it < maxit)
        func_it = func(it)
        infsum += func_it
        diff = maximum(func_it)
        it += 1
    end
    return(infsum)
end


using IntUtils

function is_abundant(n::Integer)
    d = proper_divisors(n)
    s = sum(d)
    if s > n
        return true
    else
        return false
    end
end

function is_sum_of_pairs_from(n::Integer, s::Set{Integer})
    for elt in s
        d = n - elt
        if d > 0 && d in s
            return true
        end
    end
    return false
end

function main()
    abundants = Set{Integer}()
    for i in 12:28123
        if is_abundant(i)
            push!(abundants, i)
        end
    end
    s = 0
    for i in 1:28123
        if !is_sum_of_pairs_from(i, abundants)
            s += i
        end
    end
    @printf("The sum of all numbers not the sum of two abundant numbers is %d\n", s)
end

main()


#' ---
#' title: N-Queens
#' ---

#' **Originally Contributed by**: Matthew Helm

#' The N-Queens problem involves placing N queens on an N x N chessboard such that none of the queens attacks another. In chess, a
#' queen can move vertically, horizontally, and diagonally so there cannot be more than one queen on any given row, column, or
#' diagonal.

#'

#' *Note that none of the queens above are able to attack any other as a result of their careful placement.*

using GLPK
using JuMP
using LinearAlgebra

# N-Queens
N = 8

model = Model(GLPK.Optimizer);

#' Next, let's create an N x N chessboard of binary values. 0 will represent an empty space on the board and 1 will represent a
#' space occupied by one of our queens:

@variable(model, x[i=1:N, j=1:N], Bin)

#' Now we can add our constraints:

# There must be exactly one queen in a given row/column
for i=1:N
    @constraint(model, sum(x[i, :]) == 1)
    @constraint(model, sum(x[:, i]) == 1)
end

# There can only be one queen on any given diagonal
for i in -(N-1):(N-1)
    @constraint(model, sum(diag(x,i)) <= 1)
    @constraint(model, sum(diag(reverse(x,dims=1), i)) <=1)
end

#' That's it! We are ready to put our model to work and see if it is able to find a feasible solution:

optimize!(model)

#' We can now review the solution that our model found:

solution = convert.(Int,value.(x))

#'


function exkrls{T}(x::Array{T}, y::Array{T}; nu=1., α=.5, β=0.5, q=.1, λ=0.1, kernelfunc=linear_kernel, maxdict=100)

# krls(x, y)
#
# Julia version of the Extended Kernel Recursive Least Squares Algorithm
#
# Input:
# 			  x : d(dimension) x N(samples) array
# 			  y : 1 x N(samples) vector
# 	         nu : approximate linear dependency (ALD) threshold, controls sparsity (default: 1.)
# 	     λ : regularization parameter (default: 0.1)
#        kernel : type of kernel (default : 'linear_kernel')
#       maxdict : maximum dictionary size (default: 100)
#         index : how to cycle through samples, 'lin' goes in order, 'rand' goes randomly, or you can provide your own vector index (default: 'lin')
#
# Output:
#		  alpha : weights on dictionary samples
#		   dict : dictionary samples
#		   Kinv : Inverse of kernel matrix
#	   dict_idx : dictionary sample indicies
#
#
# Citation:
# Y. Engel, S. Mannor, and R. Meir, “The kernel recursive least-squares algorithm,” IEEE Transactions on Signal Processing, vol. 52, no. 8, pp. 2275–2285, 2004.

λ = convert(T, λ)
nu = convert(T, nu)

λ2 = λ.^2

sz = size(x)

dict = zeros(eltype(x), sz[1], maxdict)
dict_idx = zeros(eltype(x), maxdict, 1)

## Initialize
K = kernelfunc(x[:,1], x[:,1])
Kinv = (1./K)'
alpha = (α*y[1] ./ (λ*β+K))'
pr = λ*β / (abs(α)^2)
dict[:,1] = x[:,1]
dict_idx[1] = 1
P = one(eltype(x))
m = 1
m2 = 1

for ii = 2:length(y)

	m2 = m2 + 1

	kt = kernelfunc(dict[:,1:m], x[:,ii]) + λ2

	ktt = kernelfunc(x[:,ii], x[:,ii]) + λ2

	at = Kinv * kt'

	dt = ktt - kt*at

	kta = kt*alpha

	e = y[ii] - kta[1]

	if abs(dt[1]) > nu && m < maxdict

		m = m + 1

		dict[:,m] = x[:,ii]
		dict_idx[m] = ii

		Kinv = (1 ./ dt) .* [dt.*Kinv + at*at' -at; -at' 1]

		P = [P zeros(eltype(x),size(P,1), 1); zeros(eltype(x),1, size(P,1)) 1]

		alpha = [alpha - ((at ./ dt) * e); (1./dt)*e]

		if mod(m,50)==0
			println("Dictionary Size: $m of $maxdict")
		end

	else

		Pat = P*at
		atPat = 1 + at'*Pat

		qt = Pat / atPat[1]

		atP = (at'*P)

		# P -= ((Pat*(at'*P)) ./ atPat)
		# optimized as
		Base.LinAlg.BLAS.gemm!('N', 'N', -one(eltype(y)), qt,atP, one(eltype(y)), P)

		# alpha +=  Kinv*qt*(y[ii] - kt*alpha)
		# optimized as

		Base.LinAlg.BLAS.gemm!('N', 'N', e, Kinv,qt, one(eltype(y)), alpha)

		if mod(m2,50)==0
			println("On sample: $m2 of $(sz[2])")
		end

	end

end

dict = dict[:,1:m]
dict_idx = dict_idx[1:m]

(alpha, dict, Kinv, dict_idx)

end


mutable struct RegularizedSVD{T} <: MatrixFactorization{T}
    P::Array
    Q::Array
    preference::Persa.Preference{T}
    users::Int
    items::Int
end

const RSVD = RegularizedSVD

function RegularizedSVD(dataset::Persa.Dataset, features::Int)
    (users, items) = size(dataset)

    P = rand(users, features)
    Q = rand(items, features)

    return RegularizedSVD(P, Q, dataset.preference, Persa.users(dataset), Persa.items(dataset))
end

Persa.predict(model::RegularizedSVD, user::Int, item::Int) = model.P[user, :]' * model.Q[item, :]

function objective(model::RegularizedSVD, dataset::Persa.Dataset, λ::Float64)
    total = 0

    for (u, v, r) in dataset
        total += (r - model[u, v])^2
        total += λ * (norm(model.P[u,:])^2 + norm(model.Q[v,:])^2)
    end

    return total
end

function update!(model::RegularizedSVD, dataset::Persa.Dataset, γ::Float64, λ::Float64)

    idx = shuffle(1:length(dataset))

    for i = 1:length(dataset)
        (u, v, r) = dataset[idx[i]]

        e = r - Persa.predict(model, u, v)

        P = model.P[u,:]
        Q = model.Q[v,:]

        model.P[u,:] += γ * (e .* Q .- λ .* P)
        model.Q[v,:] += γ * (e .* P .- λ .* Q)
    end
end


# Plot of various Mathieu functions, inspired by https://dlmf.nist.gov/28.3
using Mathieu, PyPlot
const MPL = PyPlot.matplotlib

# get the same line colors
PyPlot.PyDict(MPL."rcParams")["axes.prop_cycle"] =
    MPL.cycler(color=["#2ca02c","#d62728","#1f77b4","#ff7f0e"])

# plot parameters
z = LinRange(0,π/2,101)
n = 0:3

for q in [1,10]
    f,ax = subplots(2,2,figsize=(10,8));
    f.suptitle("Mathieu Functions for \$q=$q\$");

    ax[1].set_title("Even π-Periodic Solutions");
    ax[1].plot(z, Mathieu.cep(n,q,z));
    ax[1].legend(string.("\$ce_",2n,"\$"));

    ax[2].set_title("Even π-Antiperiodic Solutions");
    ax[2].plot(z, Mathieu.cea(n,q,z));
    ax[2].legend(string.("\$ce_",2n.+1,"\$"));

    ax[3].set_title("Odd π-Antiperiodic Solutions");
    ax[3].plot(z, Mathieu.sea(n,q,z));
    ax[3].legend(string.("\$se_",2n.+1,"\$"));

    ax[4].set_title("Odd π-Periodic Solutions");
    ax[4].plot(z, Mathieu.sep(n,q,z));
    ax[4].legend(string.("\$se_",2n.+2,"\$"));
end


"""
   LaguerreWeight(α)

is a quasi-vector representing `x^α * exp(-x)` on `0..Inf`.
"""
struct LaguerreWeight{T} <: Weight{T}
    α::T
end

LaguerreWeight{T}() where T = LaguerreWeight{T}(zero(T))
LaguerreWeight() = LaguerreWeight{Float64}()
axes(::LaguerreWeight{T}) where T = (Inclusion(ℝ),)
function getindex(w::LaguerreWeight, x::Number)
    x ∈ axes(w,1) || throw(BoundsError())
    x^w.α * exp(-x)
end

sum(L::LaguerreWeight{T}) where T = gamma(L.α + 1)

struct Laguerre{T} <: OrthogonalPolynomial{T}
    α::T
    Laguerre{T}(α) where T = new{T}(convert(T, α))
end
Laguerre{T}() where T = Laguerre{T}(zero(T))
Laguerre() = Laguerre{Float64}()
Laguerre(α::T) where T = Laguerre{float(T)}(α)
orthogonalityweight(L::Laguerre)= LaguerreWeight(L.α)

==(L1::Laguerre, L2::Laguerre) = L1.α == L2.α
axes(::Laguerre{T}) where T = (Inclusion(HalfLine{T}()), oneto(∞))

"""
     laguerrel(n, α, z)

computes the `n`-th generalized Laguerre polynomial, orthogonal with
respec to `x^α * exp(-x)`, at `z`.
"""
laguerrel(n::Integer, α, z::Number) = Base.unsafe_getindex(Laguerre{promote_type(typeof(α), typeof(z))}(α), z, n+1)

"""
     laguerrel(n, z)

computes the `n`-th Laguerre polynomial, orthogonal with
respec to `exp(-x)`, at `z`.
"""
laguerrel(n::Integer, z::Number) = laguerrel(n, 0, z)

# L_{n+1} = (-1/(n+1) x + (2n+α+1)/(n+1)) L_n - (n+α)/(n+1) L_{n-1}
# - (n+α) L_{n-1} + (2n+α+1)* L_n -(n+1) L_{n+1} = x  L_n
# x*[L_0 L_1 L_2 …] = [L_0 L_1 L_2 …] * [(α+1)    -(α+1); -1  (α+3)     -(α+2);0  -2   (α+5) -(α+3); …]
function jacobimatrix(L::Laguerre{T}) where T
    α = L.α
    Tridiagonal(-(1:∞), (α+1):2:∞, -(α+1:∞))
end

recurrencecoefficients(L::Laguerre{T}) where T = ((-one(T)) ./ (1:∞), ((L.α+1):2:∞) ./ (1:∞), (L.α:∞) ./ (1:∞))

##########
# Derivatives
##########

@simplify function *(D::Derivative, L::Laguerre)
    T = promote_type(eltype(D),eltype(L))
    D = _BandedMatrix(Fill(-one(T),1,∞), ∞, -1,1)
    Laguerre(L.α+1)*D
end


using Symbolics, Test

@sym x y z m ω t μ

@testset "   Euler-Lagrange Solver" begin
    function Γ(w)
       t -> UpTuple([t, w(t), D(w)(t)])
    end

    function Lagrange_Equations(L)
        w -> D(∂(3)(L)∘Γ(w)) - ∂(2)(L)∘Γ(w)
    end

    function L_SHO(local_tuple::UpTuple)
        t, q, qdot = local_tuple.data
        (0.5m)*qdot^2 - (0.5m*ω^2)*q^2
    end

    function L_free(local_tuple::UpTuple)
        t, q, qdot = local_tuple.data
        (0.5m)*qdot^2
    end

    function L_pendulum(local_tuple::UpTuple)
        t, q, qdot = local_tuple.data
        m*qdot^2/2+cos(q)
    end

    @test Lagrange_Equations(L_SHO)(x)(t)  == (D(D(x)))(t) * m + (x)(t) * m * ω ^ 2
    @test Lagrange_Equations(L_free)(x)(t) == (D(D(x)))(t) * m
    @test Lagrange_Equations(L_pendulum)(x)(t) == (D(D(x)))(t) * m + sin((x)(t))

    # Lagrangian for the free particle in space (3D)
    function L_free_3d(local_tuple::UpTuple)
        t, q, qdot = local_tuple.data
        (qdot[1]^2+qdot[2]^2+qdot[3]^2)/2
    end

    q = UpTuple([x, y, z])

    lagrange_eqs_fp3d = Lagrange_Equations(L_free_3d)(q)
    lagrange_eqs_fp3d_t = Lagrange_Equations(L_free_3d)(q)(t)
    @test typeof(lagrange_eqs_fp3d_t) == UpTuple
    fp3d_eqs = [(D(D(x)))(t), (D(D(y)))(t), (D(D(z)))(t)] # eqs of motion for the free particle in space
    @test lagrange_eqs_fp3d_t.data == fp3d_eqs
    @test lagrange_eqs_fp3d_t == up(fp3d_eqs)
    @test lagrange_eqs_fp3d_t == UpTuple(fp3d_eqs)
    @test lagrange_eqs_fp3d(t).data == fp3d_eqs
    @test lagrange_eqs_fp3d(t) == up(fp3d_eqs)
    @test lagrange_eqs_fp3d(t) == UpTuple(fp3d_eqs)

    # Lagrangian for the Kepler problem
    function L_kepler(local_tuple::UpTuple)
        t, q, qdot = local_tuple.data
        r = (q[1]^2+q[2]^2)^(1/2)
        (qdot[1]^2+qdot[2]^2)*(m/2)+μ/r
    end

    q = UpTuple([x, y])

    # Equivalent to Eq. (1.48) of SICM
    lagrange_eqs_kepler = Lagrange_Equations(L_kepler)(q)
    lagrange_eqs_kepler_t = Lagrange_Equations(L_kepler)(q)(t)
    @test typeof(lagrange_eqs_kepler_t) == UpTuple
    kepler_eqs = [m * (D(D(x)))(t) + ((x)(t) ^ 2 + (y)(t) ^ 2) ^ -1.5 * (x)(t) * μ, m * (D(D(y)))(t) + ((x)(t) ^ 2 + (y)(t) ^ 2) ^ -1.5 * (y)(t) * μ]
    @test lagrange_eqs_kepler_t.data == kepler_eqs
    @test lagrange_eqs_kepler_t == up(kepler_eqs)
    @test lagrange_eqs_kepler_t == UpTuple(kepler_eqs)
    @test lagrange_eqs_kepler(t).data == kepler_eqs
    @test lagrange_eqs_kepler(t) == up(kepler_eqs)
    @test lagrange_eqs_kepler(t) == UpTuple(kepler_eqs)

end


struct FractionalFactorial <: GSAMethod
end

"""
Code based on the theory presetned in:
Saltelli, A. (2008). Global sensitivity analysis: The primer. Chichester: Wiley, pp. 71-76.
"""

using LinearAlgebra

function _recursive_hadamard(k::Integer)
     """
     Generate a hadamard matrix via recursion.
     """
     # base case
     if k == 2
          return [1 1; 1 -1]
     else
          h = _recursive_hadamard(k ÷ 2)
     end

     return hcat(vcat(h, h), vcat(h, -h))
end

function _expanding_window_hadamard(k::Integer)
     """
     Generate hadamard matrix of size k using expanding window approach.
     """
     @assert ispow2(k)

     # intialize
     h = ones(Int64, k, k)
     h[2,2] = -1

     let
        bot_row = 2
        right_col = 2
        while bot_row < k

             @inbounds cur = h[1:bot_row, 1:right_col]

             new_bot = bot_row+1:bot_row*2
             new_right = right_col+1:right_col*2
             # update right
             @inbounds h[1:bot_row, new_right] = cur

             # update below
             @inbounds h[new_bot, 1:right_col] = cur

             # update diagonal
             @inbounds h[new_bot, new_right] = - cur

             # update window for each to 'copy' from.
             bot_row *= 2
             right_col *= 2
        end
     end
     return h
end

function _generate_hadamard(k::Integer)
     @assert ispow2(k)
     return _expanding_window_hadamard(k)
end

function generate_ff_design_matrix(num_parameters::Integer)
     """
     param: num_parameters
          The number of parameters to be sampled for.

     return:
          A 2-level fractional factorial design matrix of -1s and 1s with a design of resolution IV.

          The number of rows is 2*number of parameters and the number of columns is the number of parameters.

          See equation (2.31) in Saltelli, A. (2008) for details.
     """

     # If k is not a power of 2, get the next larger number that is
     # so the hadamard matrices can be computed.
     k = Integer(round(2^ceil(log2(num_parameters))))

     s = _generate_hadamard(k)
     design_matrix = vcat(s, -s) # design of resolution IV.

     return design_matrix
end

function generate_ff_sample_matrix(design_matrix::Array{Int64,2},  levels_list=nothing)
      """

      Convert the desgin matrix to a matrix whose rows can
      be used as inputs to a model by replacing low and high values for each parameter
      with the the low and high end of each parameters range, as specified in levels_list.

      param: design_matrix

            The matrix outputted by 'ff_design' i.e a 2-level fractional
            factorial design matrix of -1s and 1s with a design of resolution IV.
            The number of rows is 2*number of parameters and the
            number of columns is the number of parameters.

      param: levels_list (optional)
          A list of length equal to the number of parameters, where the elements of
          the list are 2-tuples.

          The ith tuple corresponds to the the low-value and the high-value of the ith parameter.
          i.e. levels_list[2][1] corresponds to the low-value of the 2nd parameter,
          and levels_list[2][2] corresponds to the high-value of the 2nd parameter.

          If no list is provided, the low value for each parameter
          is assumed to be 0 and the high value is assumed to be 1.

      returns:
          A matrix with the same dimensions as the design matrix, but with the '-1's replaced with
          the low value specified for each parameter in level_list, and similar the '1's are replaced
          with the high values.

          Each row corresponds to a single input vector to be used as input to a model/function.
      """
      sample_matrix = copy(design_matrix)
      if !isnothing(levels_list)
           for (col_index, (low_value, high_value)) in enumerate(levels_list)
                design_col = @view design_matrix[:,col_index]
                sample_col = @view sample_matrix[:,col_index]

                sample_col[design_col .== -1] .= low_value
                sample_col[design_col .== 1] .= high_value
            end
      end
      return sample_matrix
end

function run_model(sample_matrix::AbstractArray, model)
      """

      param: sample_matrix
            A matrix where each row corresponds to a input vector for "model"
      param: model
            The function/model to be executed on each row of the sample_matrix.
            Must return a scalar.
      return:
            A vector containing the result of running each row of the
            sample_matrix through model.
      """

      N = size(sample_matrix)[1]
      y_out = zeros(N)
      Threads.@threads for i in 1:N
          y_out[i] = model(@view sample_matrix[i,:])
      end
      return y_out
end

function ff_main_effects(design_matrix::AbstractArray, response_values::AbstractArray)
      """
           Computes the main effect for each parameter

           param: design_matrix

                The matrix outputted by 'ff_design' i.e a 2-level fractional
                factorial design matrix of -1s and 1s with a design of resolution IV.
                The number of rows is 2*number of parameters and the
                number of columns is the number of parameters.

           param: response_values

                An array of values corresponding to the result of running the
                model on each row of a sample matrix generated by the design_matrix

           return: main_effects

               A vector of the main effects of each parameter on the output.

            See Saltelli, A. equation 2.32 for details.
      """

      num_rows, num_cols = size(design_matrix)
      @assert num_rows == length(response_values) "Number of rows in design matrix must match number of responses"

      main_effects = zeros(num_cols)
      for column in 1:num_cols
            contrast_vector = @view design_matrix[:,column]
            main_effects[column] = dot(response_values, contrast_vector)
      end
      return main_effects ./ num_rows
end

function gsa(f, method::FractionalFactorial; num_params, N, p_range = nothing, kwargs...)
    design_matrix = generate_ff_design_matrix(num_params)
    sample_matrix = generate_ff_sample_matrix(design_matrix, p_range)

    response_vec = run_model(sample_matrix, f)

    main_effects = ff_main_effects(design_matrix, response_vec)
    return main_effects, main_effects.^2
end


using GARCH
using Pkg, Test, DelimitedFiles, Statistics

filename = Pkg.dir("GARCH", "test", "data", "price.csv")
price = Array{Float64}(readdlm(filename, ',')[:,2])

ret = diff(log.(price))
ret = ret .- mean(ret)
fit = garchFit(ret)

param = [2.469347e-06, 1.142268e-01, 8.691734e-01] #R fGarch garch(1,1) estimated params

@test fit.params ≈ param atol=1e-3
@test predict(fit) ≈ 0.005617744 atol = 1e-4
@test predict(fit, 2) ≈ [0.005617744, 0.005788309] atol = 1e-4


using LinearAlgebra
using Base: require_one_based_indexing

"""
    systemSolver(A::AbstractMatrix, b::Vector, x_initial::Vector) -> Vector

    Compute the orthogonal projection of a given point to the
solution set of a system of linear equations. Also a direct
method for solving systems of linear equations. The output
is either the projection or inconsistency of the system.

# Arguments
 - `A::AbstractMatrix{Real}`: coefficients of linear equations
 - `b::AbstractVector{Real}`: intended results of equations
 - `x0::AbstractVector{Real}`: initial point for projection
 - `β::Real`: set in order to obtain comparable summands in norm

# References
> Ján Plesník, Finding the orthogonal projection of a point onto
> an affine subspace, Linear Algebra and its Applications, Volume 422,
> Issues 2–3, 2007, Pages 455-470
"""
function systemSolver(A::AbstractMatrix{<:Real},    # in: equations matrix (m x n) rows = coefficients
                    b::AbstractVector{<:Real},      # in: equations results
                    x0::AbstractVector{<:Real},     # in: initial point
                    β::Real=1.0)                    # in: optional normalizer

    require_one_based_indexing(A, b, x0)
    m, n = size(A)
    if m != length(b) || n != length(x0)
        throw(DimensionMismatch("dimensions of `A`, $m * $n, does not match one of the lengths of `x0`, $(length(x0)), or `b`, $(length(b))"))
    end
    for i in 1:m
        iszero(A[i, :]) && throw(ArgumentError("`A` rows must be nonzero"))
    end
    iszero(β) && throw(ArgumentError("`β` must be nonzero"))

    if m == 0
        return x0
    end

    k = 1
    p = 1
    V = hcat(A[1, :])
    x = x0 + (b[1] - dot(A[1, :], x0))/(dot(A[1, :], V[:, 1])) * V[:, 1]

    denominators = zeros(Real, m-1)

    while k != m
        k = k + 1
        p = p + 1

        denominators[k-1] = dot((A[k, :] - A[1, :]), V[:, k-1])
        if denominators[k-1] == 0
            A[k, :] = - A[k, :]
            b[k] = - b[k]
            denominators[k-1] = dot((A[k, :] - A[1, :]), V[:, k-1])
        end

        w = x0 + β * A[k, :]
        for j in 2:k
            w = w + (b[j]-b[1]-dot(A[j, :]-A[1, :], w))/denominators[j-1] * V[:, j-1]
        end

        u = x + (b[k] - dot(A[k, :], x))/denominators[k-1] * V[:, k-1]
        if u != w
            V = hcat(V, w - u)
            x = u + (b[1] - dot(A[1, :], u))/dot(A[1, :], V[:, k]) * V[:, k]
        elseif dot(A[k, :], x) == b[k]
            if k != m
                A = A[setdiff(1:end, k), :]
                b = b[setdiff(1:end, k)]
                k = k - 1
                m = m - 1
            end
        else
            throw(DomainError("the ($p)th equation is inconsistent"))
        end
    end
    return x
end


"""
     l63!(du,u,p,t)

 Steps forward the Lorenz63 equations for time `t`.
 Prognostic variables x, y, z are defined as
 elements 1, 2 and 3 of vector `u` and with parameters
 σ, ρ and β defined as elements of vector `p`.
 Elements of du are set to dxdt, dydt and dzdt respectively.

```
 julia> using DifferentialEquations, Plots
 julia> u0=[1.,0.,0.]
 julia> p=[10.,28.,8/3]
 julia> tspan=(0.,100.)
 julia> prob=ODEProblem(l63!,u0,tspan,p)
 julia> sol=solve(prob)
 julia> plot(sol,vars=(1,2,3))
```

"""
function l63!(du,u,p,t)
        x, y, z=u
        σ, ρ, β=p

        dxdt=σ*(y-x)
        dydt=ρ*x-x*z-y
        dzdt=x*y-β*z

        du[1], du[2], du[3]=
        dxdt , dydt,  dzdt
        return
end


@testset "spectral operators" begin
    function P5(r::AbstractVector{T}) where {T}
        return T(1) / T(8) * (T(15) * r - T(70) * r .^ 3 + T(63) * r .^ 5)
    end
    function P6(r::AbstractVector{T}) where {T}
        return T(1) / T(16) * (-T(5) .+ T(105) * r .^ 2 - T(315) * r .^ 4 + T(231) * r .^ 6)
    end
    function DP6(r::AbstractVector{T}) where {T}
        return T(1) / T(16) * (T(2 * 105) * r - T(4 * 315) * r .^ 3 + T(6 * 231) * r .^ 5)
    end

    IPN(::Type{T}, N) where {T} = T(2) / T(2 * N + 1)

    for T in (Float32, Float64, BigFloat)
        r, w = legendregausslobatto(T, 7)
        x = LinRange{T}(-1, 1, 101)
        D = spectralderivative(r)
        P = spectralinterpolation(r, x)

        @test sum(P5(r) .^ 2 .* w) ≈ IPN(T, 5)
        @test D * P6(r) ≈ DP6(r)
        @test P * P6(r) ≈ P6(x)
    end

    for T in (Float32, Float64, BigFloat)
        r, w = legendregauss(T, 7)
        D = spectralderivative(r)

        @test sum(P5(r) .^ 2 .* w) ≈ IPN(T, 5)
        @test sum(P6(r) .^ 2 .* w) ≈ IPN(T, 6)
        @test D * P6(r) ≈ DP6(r)
    end
end


function nlp_mi_005_010(
    optimizer,
    objective_tol,
    primal_tol,
    dual_tol,
    termination_target = TERMINATION_TARGET_LOCAL,
    primal_target = PRIMAL_TARGET_LOCAL,
)
    # Test Goals:
    # - function /

    model = Model(optimizer)

    @variable(model, x >= 0, Int)
    @variable(model, y >= 0)

    @objective(model, Min, x + y)
    @NLconstraint(model, y >= 1 / (x + 0.1) - 0.5)
    @NLconstraint(model, x >= y^(-2) - 0.5)
    @NLconstraint(model, 4 / (x + y + 0.1) >= 1)

    optimize!(model)

    check_status(model, FEASIBLE_PROBLEM, termination_target, primal_target)
    check_objective(model, 1.8164965727459055, tol = objective_tol)
    return check_solution([x, y], [1, 0.816496581496872], tol = primal_tol)
end


### A Pluto.jl notebook ###
# v0.14.1

using Markdown
using InteractiveUtils

# ╔═╡ db4a6b48-99a7-11eb-097b-676aed417c48
begin

include("MultiArmBandit.jl")

using .MultiArmBandit
using DataFrames
using StatsBase, Statistics
using Distributions
using Gadfly
import Random

end

# ╔═╡ b65dd189-456b-40ac-88fc-d05b5e29b396
md"

# Multi-Arm Bandits

A Multi-Arm Bandit is a concept from Reinforcement Learning which is useful for decision making under uncertainty. Take a situation where you have limited time or resources to test multiple drugs, perhaps the COVID vaccines for example. While you could use the standard approach of a clinical trial with randomization, a balanced design, and performing AB tests, this results in many participants not recieving the best treatment, as some of them will be allocated to inferior choices. Ideally, we would like to do both: 1) Show a given treatment is the best choice and 2) Have most people end up benefitting from it.

Using a strategy known as $\epsilon$-greedy Q learning, the Multi Arm bandit can help accomplish this task in an online manner where the treatments are allocated on the spot. The arm refers to different treatment groups that the bandit can select from. In the beginning, each treatment arm is assigned the same value of $Q_{tmt}=Q_{0}$. Some small amount of noise can be added to break ties. At each step, $(1-\epsilon)$ of the time the bandit randomly selects from the treatment with the highest $Q_{tmt}$ and updates its corresponding value. $\epsilon$ of the time, the bandit will select a random treatment instead of the one with the highest Q. The update rule is as follows:

$Q_{tmt,t} = Q_{tmt,t-1} + \frac{R-Q_{tmt,t-1}}{N_{tmt,t}}$

where R is the reward, which is usually just the response Y corresponding to the treatment given. $N_{t+1}$ corresponds to the sample size in the given treatment group after the sample from the group is randomly selected. The bandit will be trying to go for the treatment which maximizes the reward, and so care should be taken that the response is transformed or sign-flipped if necessary so that higher is better. The above update rule can be shown to correspond to the calculation of the sample mean in an online manner if $Q_{0} = 0$. However, we can modify  the updates by adding an explicit learning rate parameter $\eta$:

$Q_{tmt,t+1} = Q_{tmt,t} + \eta(R-Q_{tmt,t})$

This can help improve learning. In this notebook, we will simulate data from 5 groups and examine the impact of setting different parameters for $Q_{0}$ and $\eta$.

The source code for the functions used in this demo Pluto notebook can be found in the **MultiArmBandit.jl** module in this repo.

"

# ╔═╡ 38528382-a1ec-4117-88d4-fbcc019c7302
md"

## Simulate Data

Below, we simulate n=10000 points from 5 groups each A, B, C, D, E from a normal distibution with mean $\mu = (-0.3,0.0,0.5,0.8,1.1)$ and $\sigma = (1, 0.5, 0.3, 0.3, 0.5)$ respectively. The bandit will select points from this simulated dataset. Note that in the ideal situation the ground truth in this simulated example is:

$ E > D > C > B > A $

Ideally, the bandit will come to this same conclusion and choose E most of the time. However, the different SDs for the groups can complicate this choice, as the bandit is greedy and thus converging to E in a finite number of iterations (in the real world, this would correspond to resources or number of patients in the trial) can be more complicated. We can track the bandit's history and plot it over time to examine when convergence occurs.

The sample mean/SD of each Arm are shown below. The bandit will be selecting samples from each group in an online manner, and so at each step it does NOT have access to every single data point. Instead it must figure out in the finite number of iterations which group is the best

"

# ╔═╡ 60e50c4e-8120-4306-9c6c-9282fff8eb47
begin

n = 1000
Random.seed!(101)
A= rand(Normal(-0.3,1),n)
B= rand(Normal(0,0.5),n)
C = rand(Normal(0.5,0.3),n)
D = rand(Normal(0.8,0.3),n)
E = rand(Normal(1.1,0.5),n)

simdata = DataFrame(Arm = vcat(fill("A",n),fill("B",n),fill("C",n),fill("D",n),fill("E",n)),
                    Y = vcat(A,B,C,D,E))

summary_simdata = combine(groupby(simdata,:Arm),nrow,:Y=>mean=>:meanY,:Y=>std=>:sdY)
summary_simdata
end

# ╔═╡ 65107f38-7916-43a4-b36f-1ef012c43b51
md"

## Create Bandit: $Q_0=0, \eta=\frac{1}{N_{tmt,t}}, \epsilon = 0.1$

Below we instantiate the first bandit with all default settings, and allow it to explore 10% of the time while 90% of the time it selects the current best group in a greedy manner.

"

# ╔═╡ 8ed00c06-7348-49dc-a9f2-041465beb969
band1=MultiArmBandit.CreateBandit(simdata,:Arm,:Y,track_hist=true)

# ╔═╡ 472fee46-6e48-47a7-9235-ccb7c460f735
begin
Random.seed!(1001)
MultiArmBandit.Train!(band1,999,ϵ = 0.1)
end

# ╔═╡ 1243d0ab-f95c-482d-b81e-3b3eb60907eb
band1_df = MultiArmBandit.GetHistoryDF(band1)

# ╔═╡ 794d15a3-a4b4-4ffa-9afb-bb19d79e79e5
plot(band1_df,x=:steps,y=:Q,color=:Group,Geom.line)

# ╔═╡ 5e8e7408-aa48-4a2d-a427-40e8adf81d3c
band1_summary = summary(band1)

# ╔═╡ 32f83952-c660-4fde-8bd5-ae1a97e7bacd
md"

Based on the above figure, the correct ordering is found after around 250 steps. For the first few steps, the bandit picks Group D which is the 2nd best but eventually finds Group C as a result of the exploration. Lets try to see what happens if we lower the bandit's exploration by setting $\epsilon = 0.03$. That is, the bandit will select a group randomly (instead of the group it deems the best group at the moment) 3% of the time.

## Create Bandit: $Q_0=0, \eta=\frac{1}{N_{tmt,t}}, \epsilon = 0.03$

"

# ╔═╡ 65945e4f-17e9-40ba-870f-88da0c5d13d0
begin
band2=MultiArmBandit.CreateBandit(simdata,:Arm,:Y,track_hist=true)
Random.seed!(1102)
MultiArmBandit.Train!(band2,999,ϵ = 0.03)
band2_df = MultiArmBandit.GetHistoryDF(band2)
band2_summary = summary(band2)
end

# ╔═╡ 3566a579-84dc-4cb6-91ab-833d0d5e9782
plot(band2_df,x=:steps,y=:Q,color=:Group,Geom.line)

# ╔═╡ a15c8f8f-f077-4758-a668-01c8b562ca95
md"

In this situation, it took longer (~100 iterations) for the bandit to arrive at the best group

In the next section, we will change the starting value to be higher at $Q_{0} = 5$. This will tend to encourage a lot more early exploration as this value is much higher than what is possible based on sampling.

## Create Bandit: $Q_0=5, \eta=\frac{1}{N_{tmt,t}}, \epsilon = 0.1$

"

# ╔═╡ 77d6e956-219e-40e9-978d-f7e0906a8d6c
begin
band3=MultiArmBandit.CreateBandit(simdata,:Arm,:Y,Q₀=5,track_hist=true)
Random.seed!(901)
MultiArmBandit.Train!(band3,99,ϵ = 0.1)
band3_df = MultiArmBandit.GetHistoryDF(band3)
band3_summary = summary(band3)
end

# ╔═╡ 043f9d99-e8b5-4b33-82ac-3870a2238f85
plot(band3_df,x=:steps,y=:Q,color=:Group,Geom.line)

# ╔═╡ 50196d6d-8a52-4043-aa9f-e3f43a9ed2a1
md"

In this case, starting at a value much higher than possible leads to earlier exploration, resulting in the Q value for each group falling rapidly. However, the bandit arrives at the best option E much quicker since it is pulled down less and is essentially guaranteed to be selected since the bandit desperately switches options more in the beginning to try to maintain a high Q value.

In this next part, we keep $Q_0 = 5$ but set the learning rate $\eta=0.01$

## Create Bandit: $Q_0=5, \eta=0.01, \epsilon = 0.1$

"

# ╔═╡ 9f3d4773-ac0a-4dd2-992d-b549f58f4021
begin
band4=MultiArmBandit.CreateBandit(simdata,:Arm,:Y,Q₀=5,track_hist=true)
Random.seed!(555)
MultiArmBandit.Train!(band4,2999,ϵ = 0.1,η=0.01)
band4_df = MultiArmBandit.GetHistoryDF(band4)
band4_summary = summary(band4)
end

# ╔═╡ 3789a09f-73fb-4ca8-b267-bbe2d7b32848
plot(band4_df,x=:steps,y=:Q,color=:Group,Geom.line)

# ╔═╡ 7825e432-9f3b-4099-bf37-0b1b3ccf38c1
md"

In this case, setting the learning rate to 0.01 leads the bandit to take longer to reach the optimal solution E. Until around 1500 steps, the Q value for all the groups is even.

What happens if we increase the learning rate to 1 (but keep $Q_0 = 5$) ?

## Create Bandit: $Q_0=5, \eta=1, \epsilon = 0.1$

"

# ╔═╡ 2342907c-75e0-4126-b2df-6ce5a5901bf1
begin
band5=MultiArmBandit.CreateBandit(simdata,:Arm,:Y,Q₀=5,track_hist=true)
Random.seed!(1556)
MultiArmBandit.Train!(band5,2999,ϵ = 0.1,η=1)
band5_df = MultiArmBandit.GetHistoryDF(band5)
band5_summary = summary(band5)
end

# ╔═╡ 7c511272-01f8-4270-bdf7-f4f66e78df07
plot(band5_df,x=:steps,y=:Q,color=:Group,Geom.line)

# ╔═╡ c1deba65-a93f-4992-958e-38a987346691
md"

Now the plot is extremely noisy, and the bandit cannot seem to differentiate between Group D and Group E very well. We can see in the summary that the bandit tended to select D the most but in the end the Q values of both D/E are about the same.

In this final parts, we now set $Q_0=0$ again and examine the impact of the two learning rates.

## Create Bandit: $Q_0=0, \eta=0.01, \epsilon = 0.1$

"

# ╔═╡ 8bbfae0c-6648-4f42-a40c-ada03d17bb3a
begin
band6=MultiArmBandit.CreateBandit(simdata,:Arm,:Y,Q₀=0,track_hist=true)
Random.seed!(121)
MultiArmBandit.Train!(band6,2999,ϵ = 0.1,η=0.01)
band6_df = MultiArmBandit.GetHistoryDF(band6)
band6_summary = summary(band6)
end

# ╔═╡ 613e3215-4402-4ba2-8402-2ab5bf1f1247
plot(band6_df,x=:steps,y=:Q,color=:Group,Geom.line)

# ╔═╡ e7772e4c-b798-4bdd-85c1-6ee291fce282
md"

In this case, starting at $Q_0=0$ and setting the learning rate $\eta = 0.01$ leads to the bandit favoring Group C, a suboptimal group, for a long time. It is only after ~2700 iterations that it picks up on Group E.

What happens if the learning rate is increased? Will we see the same noisy behavior as before when starting from $Q_0=0$?

## Create Bandit: $Q_0=0, \eta=1, \epsilon = 0.1$

"

# ╔═╡ bf545060-f55c-4a91-99f8-c41543c25ebb
begin
band7=MultiArmBandit.CreateBandit(simdata,:Arm,:Y,Q₀=0,track_hist=true)
Random.seed!(13)
MultiArmBandit.Train!(band7,2999,ϵ = 0.1,η=1)
band7_df = MultiArmBandit.GetHistoryDF(band7)
band7_summary = summary(band7)
end

# ╔═╡ 5af03195-5199-48b9-8792-f02da75dc320
plot(band7_df,x=:steps,y=:Q,color=:Group,Geom.line)

# ╔═╡ a94f1ce9-b874-4da2-af2b-ee8c6840a2d2
md"

The higher learning rate again leads to noise, and in the end, the bandit has trouble differentiating Group D and E just as before.

In summary, lowering the learning rate appears to increase the chance of arriving at a suboptimal solution. Increasing the learning rate too much confuses the Bandit and the Q value curves become extremely noisy.

Starting at a higher Q (with no set learning rate) encourages quicker exploration between the groups

"

# ╔═╡ Cell order:
# ╟─b65dd189-456b-40ac-88fc-d05b5e29b396
# ╠═db4a6b48-99a7-11eb-097b-676aed417c48
# ╟─38528382-a1ec-4117-88d4-fbcc019c7302
# ╠═60e50c4e-8120-4306-9c6c-9282fff8eb47
# ╟─65107f38-7916-43a4-b36f-1ef012c43b51
# ╠═8ed00c06-7348-49dc-a9f2-041465beb969
# ╠═472fee46-6e48-47a7-9235-ccb7c460f735
# ╠═1243d0ab-f95c-482d-b81e-3b3eb60907eb
# ╠═794d15a3-a4b4-4ffa-9afb-bb19d79e79e5
# ╠═5e8e7408-aa48-4a2d-a427-40e8adf81d3c
# ╟─32f83952-c660-4fde-8bd5-ae1a97e7bacd
# ╠═65945e4f-17e9-40ba-870f-88da0c5d13d0
# ╠═3566a579-84dc-4cb6-91ab-833d0d5e9782
# ╟─a15c8f8f-f077-4758-a668-01c8b562ca95
# ╠═77d6e956-219e-40e9-978d-f7e0906a8d6c
# ╠═043f9d99-e8b5-4b33-82ac-3870a2238f85
# ╟─50196d6d-8a52-4043-aa9f-e3f43a9ed2a1
# ╠═9f3d4773-ac0a-4dd2-992d-b549f58f4021
# ╠═3789a09f-73fb-4ca8-b267-bbe2d7b32848
# ╟─7825e432-9f3b-4099-bf37-0b1b3ccf38c1
# ╠═2342907c-75e0-4126-b2df-6ce5a5901bf1
# ╠═7c511272-01f8-4270-bdf7-f4f66e78df07
# ╟─c1deba65-a93f-4992-958e-38a987346691
# ╠═8bbfae0c-6648-4f42-a40c-ada03d17bb3a
# ╠═613e3215-4402-4ba2-8402-2ab5bf1f1247
# ╟─e7772e4c-b798-4bdd-85c1-6ee291fce282
# ╠═bf545060-f55c-4a91-99f8-c41543c25ebb
# ╠═5af03195-5199-48b9-8792-f02da75dc320
# ╟─a94f1ce9-b874-4da2-af2b-ee8c6840a2d2


function Threshold(input, threshold, value)
    output = map(x -> x > threshold ? x : value, input)
    return output
end

function Threshold!(input, threshold, value)
    input = map(x -> x > threshold ? x : value, input)
    return input
end

function ReLU(input)
    return max(0, input)
end

function ReLU!(input)
    input = max(0, input)
    return input
end

function Hardtanh(input, min_value, max_value)
    return map(x -> x > max_value ? 1 : x < min_value ? -1 : x, input)
end

function Hardtanh(input, min_value, max_value)
    input = map(x -> x > max_value ? 1 : x < min_value ? -1 : x, input)
    return input
end

function ReLU6(input)
    return min(max(0, x), 6)
end

function ReLU6!(input)
    input = min(max(0, x), 6)
    return input
end

Sigmoid(x) = 1 ./ (1 + exp(-x))

Tanh(x) = (exp(x) - exp(-x)) ./ (exp(x) + exp(-x))

ELU(x) = max(0, x) + min(0, alpha * (exp(x) - 1))

function hardshrink(input, lambda)
    return map(x -> x > lambda || x < lambda ? x : 0, input)
end

function hardshrink!(input, lambda)
    input = map(x -> x > lambda || x < lambda ? x : 0, input)
end

LeakyReLU(x, negative_slope) = max(0, x) + negative_slope(0, x)

LogSigmoid(x) = log(Sigmoid(x))

Softplus(x, beta) = 1 ./ beta * log(1 + exp(beta * x))

function Softshrink(input, lambda)
    output = map(input) do x
        if x > lambda
            return x - lambda
        elseif x < -lambda
            return x + lambda
        else
            return 0
        end
    end
    return output
end

function Softshrink!(input, lambda)
    input = map(input) do x
        if x > lambda
            return x - lambda
        elseif x < - lambda
            return x + lambda
        else
            return 0
        end
    end
    return input
end

PReLU(x, a) = max(0, x) + a * min(0, x)

Softsign(x) = x ./ (1 + abs(x))

Tanhshrink(x) = x - Tanh(x)

function Softmin(input)
    shift(x) = max(input) - x
    return map(x -> exp(-x - shift(x)) ./ sum(exp(map(y -> y - shift(x), input))), input)
end

function Softmin!(input)
    shift(x) = max(input) - x
    input = map(x -> exp(-x - shift(x)) ./ sum(map(y -> exp(y - shift(x)), input)), input)
end

function Softmax(input)
    return map(x -> exp(x - max(x)) ./ sum(map(y -> exp(y - max(x)), input)), input)
end

function Softmax!(input)
    input = map(x -> exp(x - max(x)) ./ sum(map(y -> exp(y - max(x)), input)), input)
end

LogSoftmax(x) = log(Softmax(x))


"""Computes the value of the function ackerman(n,m)"""
function ack(m :: Int64, n :: Int64) :: Int64
    if m == 0
        return n + 1
    elseif n == 0
        return ack(m - 1, 1)
    else
        return ack(m - 1, ack(m, n - 1))
    end
end


function bernoulli(n)
    A = Vector{Rational{BigInt}}(undef, n + 1)
    for m = 0 : n
        A[m + 1] = 1 // (m + 1)
        for j = m : -1 : 1
            A[j] = j * (A[j] - A[j + 1])
        end
    end
    return A[1]
end

function display(n)
    B = map(bernoulli, 0 : n)
    pad = mapreduce(x -> ndigits(numerator(x)) + Int(x < 0), max, B)
    argdigits = ndigits(n)
    for i = 0 : n
        if numerator(B[i + 1]) & 1 == 1
            println(
                "B(", lpad(i, argdigits), ") = ",
                lpad(numerator(B[i + 1]), pad), " / ", denominator(B[i + 1])
            )
        end
    end
end

display(60)


"""
    black_body(T, K₀, σ)
    black_body(T)

Thermal infrared, *i.e.* longwave radiation emitted from a black body at temperature T.

- `T`: temperature of the object in Celsius degree
- `K₀`: absolute zero (°C)
- `σ` (``W\\ m^{-2}\\ K^{-4}``) [Stefan-Boltzmann constant](https://en.wikipedia.org/wiki/Stefan%E2%80%93Boltzmann_law)

# Note

`K₀` and `σ` are taken from [`Constants`](@ref) if not provided.

"""
function black_body(T, K₀, σ)
  Tₖ = T - K₀
  σ * (Tₖ^4.0)
end

function black_body(T)
    constants = Constants()
    black_body(T, constants.K₀, constants.σ)
end

"""
Thermal infrared, *i.e.* longwave radiation emitted from an object at temperature T.

- `T`: temperature of the object in Celsius degree
- `ε` object [emissivity](https://en.wikipedia.org/wiki/Emissivity) (not to confuse with ε the
ratio of molecular weights from [`Constants`](@ref)). A typical value for a leaf is 0.955.
- `K₀`: absolute zero (°C)
- `σ` (``W\\ m^{-2}\\ K^{-4}``) [Stefan-Boltzmann constant](https://en.wikipedia.org/wiki/Stefan%E2%80%93Boltzmann_law)

# Note

`K₀` and `σ` are taken from [`Constants`](@ref) if not provided.

# Examples

```julia
# Thermal infrared radiation of water at 25 °C:
grey_body(25.0, 0.96)
```
"""
function grey_body(T, ε, K₀, σ)
  ε * black_body(T, K₀, σ)
end

function grey_body(T, ε)
    constants = Constants()
    grey_body(T, ε, constants.K₀, constants.σ)
end

"""
    net_longwave_radiation(T₁,T₂,ε₁,ε₂,F₁,K₀,σ)
    net_longwave_radiation(T₁,T₂,ε₁,ε₂,F₁)

Net longwave radiation fluxes (*i.e.* thermal radiation, W m-2) between an object and another.
The object of interest is at temperature T₁ and has an emissivity ε₁, and the object with
wich it exchanges energy is at temperature T₂ and has an emissivity ε₂.

If the result is positive, then the object of interest gain energy.

# Arguments

- `T₁` (Celsius degree): temperature of the target object (object 1)
- `T₂` (Celsius degree): temperature of the object with which there is potential exchange (object 2)
- `ε₁`: object 1 emissivity
- `ε₂`: object 2 emissivity
- `F₁`: view factor (0-1), *i.e.* visible fraction of object 2 from object 1 (see note)
- `K₀`: absolute zero (°C)
- `σ` (``W\\ m^{-2}\\ K^{-4}``) [Stefan-Boltzmann constant](https://en.wikipedia.org/wiki/Stefan%E2%80%93Boltzmann_law)

# Note

`F₁`, the view factor (also called shape factor) is a coefficient applied to the semi-hemisphere
field of view of object 1 that "sees" object 2. E.g. a leaf can be viewed as a plane. If one side
of the leaf sees only object 2 in its field of view (e.g. the sky), then `F₁ = 1`.
Then the net longwave radiation flux for this part of the leaf is multiplied by its actual
surface to get the exchange. Note that we apply reciprocity between the two objects for
the view factor (they have the same value), *i.e.*: A₁F₁₂ = A₂F₂₁.

Then, if we take a leaf as object 1, and the sky as object 2, the visible fraction of
sky viewed by the leaf would be:

- `0.5` if the leaf is on top of the canopy, *i.e.* the upper side of the leaf sees the sky,
the side bellow sees other leaves and the soil.
- between 0 and 0.5 if it is within the canopy and partly shaded by other objects.

Note that `A₁` for a leaf is twice its common used leaf area, because `A₁` is the **total**
leaf area of the object that exchange energy.

```julia
# Net thermal radiation fluxes between a leaf and the sky considering the leaf at the top of
# the canopy:
Tₗ = 25.0 ; Tₐ = 20.0
ε₁ = 0.955 ; ε₂ = 1.0
Rₗₗ = net_longwave_radiation(Tₗ,Tₐ,ε₁,ε₂,1.0)
Rₗₗ

# Rₗₗ is the net longwave radiation flux between the leaf and the atmosphere per surface area.
# To get the actual net longwave radiation flux we need to multiply by the surface of the
# leaf, e.g. for a leaf of 2cm²:
leaf_area = 2e-4 # in m²
Rₗₗ * leaf_area

# The leaf lose ~0.0055 W towards the atmosphere.
```

# References

Cengel, Y, et Transfer Mass Heat. 2003. A practical approach. New York, NY, USA: McGraw-Hill.
"""
function net_longwave_radiation(T₁,T₂,ε₁,ε₂,F₁,K₀,σ)
    (black_body(T₂,K₀,σ) - black_body(T₁,K₀,σ)) / (1.0 / ε₁ + 1.0 / ε₂ - 1.0) * F₁
end

function net_longwave_radiation(T₁,T₂,ε₁,ε₂,F₁)
    constants = Constants()
    net_longwave_radiation(T₁,T₂,ε₁,ε₂,F₁,constants.K₀,constants.σ)
end

"""
    atmosphere_emissivity(Tₐ,eₐ)

Emissivity of the atmoshpere at a given temperature and vapor pressure.

# Arguments

- `Tₐ` (°C): air temperature
- `eₐ` (kPa): air vapor pressure
- `K₀` (°C): absolute zero

# Examples

```julia
Tₐ = 20.0
VPD = 1.5
atmosphere_emissivity(Tₐ, vapor_pressure(Tₐ,VPD))
```

# References

Leuning, R., F. M. Kelliher, DGG de Pury, et E.-D. SCHULZE. 1995. Leaf nitrogen,
photosynthesis, conductance and transpiration: scaling from leaves to canopies ». Plant,
Cell & Environment 18 (10): 1183‑1200.
"""
function atmosphere_emissivity(Tₐ,eₐ,K₀)
    0.642 * (eₐ * 100 / (Tₐ - K₀))^(1 / 7)
end

function atmosphere_emissivity(Tₐ,eₐ)
    atmosphere_emissivity(Tₐ,eₐ,Constants().K₀)
end


# This code presents an example to illustrate how SmolyakApprox can be used

using SmolyakApprox

function test_smolyak_approx()

  d  = 5  # Set the number of dimensions
  mu = 3  # Set the level of approximation

  g2, m2 =  smolyak_grid(clenshaw_curtis_equidistant,d,mu)
  grid, multi_ind = smolyak_grid(chebyshev_gauss_lobatto,d,mu)  # Construct the Smolyak grid and the multi index

  # An arbitrary test function

  function test(grid)

    y_value = (grid[:,1].+1).^0.1.*exp.(grid[:,2]).*log.(grid[:,3].+2).^0.2.*(grid[:,4].+2).^0.8.*(grid[:,5].+7).^0.1

    return y_value

  end

  y = test(grid)  # Evaluate the test function on the Smolyak grid
  y_pl = test(g2)

  point = [0.75, 0.45, 0.82, -0.15, -0.95]  # Choose a point to evaluate the approximated function

  # One way of computing the weights and evaluating the approximated function

  weights = smolyak_weights(y,grid,multi_ind)        # Compute the Smolyak weights
  y_hat = smolyak_evaluate(weights,point,multi_ind)  # Evaluate the approximated function

  #= A second way of computing the weights and evaluating the approximated function that
  computes the interpolation matrix just once. =#

  interp_mat = smolyak_inverse_interpolation_matrix(grid,multi_ind)  # Compute the interpolation matrix
  w = smolyak_weights(y,interp_mat)             # Compute the Smolyak weights
  y_hatt = smolyak_evaluate(w,point,multi_ind)  # Evaluate the approximated function

  # Piecewise linear

  w_pl = smolyak_pl_weights(y_pl,g2,m2)
  y_pl_hat = smolyak_pl_evaluate(w_pl,point,g2,m2)

  # Evaluate the exact function at point

  y_actual = test(point')

  # Now consider the ansiotropic case

  mu = [3, 2, 2, 2, 3]
  grid, multi_ind = smolyak_grid(chebyshev_gauss_lobatto,d,mu)  # Construct the Smolyak grid and the multi index
  y = test(grid)
  weights = smolyak_weights(y,grid,multi_ind)                   # Compute the Smolyak weights
  y_hat_ansio = smolyak_evaluate(weights,point,multi_ind)       # Evaluate the approximated function
  weights_th = smolyak_weights_threaded(y,grid,multi_ind)                   # Compute the Smolyak weights

  g3, m3 =  smolyak_grid(clenshaw_curtis_equidistant,d,mu)
  y_pl = test(g3)
  w_pl_ansio = smolyak_pl_weights(y_pl,g3,m3)
  y_pl_hat_ansio = smolyak_pl_evaluate(w_pl,point,g3,m3)
  w_pl_ansio_th = smolyak_pl_weights_threaded(y_pl,g3,m3)

  # Now test the full grid results

  mu = 3
  grid_full, multi_ind_full = smolyak_grid_full(chebyshev_gauss_lobatto,d,mu) # Construct the Smolyak grid and the multi index
  y_full = test(grid_full)
  weights_full = smolyak_weights_full(y_full,grid_full,multi_ind_full)        # Compute the Smolyak weights
  y_hat_full = smolyak_evaluate_full(weights_full,point,multi_ind_full)       # Evaluate the approximated function

  return y_hat, y_hatt, y_pl_hat, y_actual, y_hat_ansio, y_pl_hat_ansio, y_hat_full

end

test_smolyak_approx()


using LinearAlgebra
using PyPlot
using Statistics
using Distributions
using ProgressMeter

σ = 1
τ = sqrt(10)
μ = 5
n = 5

X = [9.37;10.18;9.16;11.60;10.33]
μn = mean(X)*n/σ^2/(n/σ^2 + 1/τ^2) + μ*(1/τ^2)/(n/σ^2 + 1/τ^2)
τn2 = 1/(n/σ^2+1/τ^2)

# Hyperparameters
# proposal step size
δ = 3/100
burnin = 2000
N = 10000
θ0 = 0.0

function logf(θ)
    -sum((X.-θ).^2)/2σ^2 - (θ-μ)^2/2τ^2
end

function proposal(x)
    x + (rand()-0.5)*2 * δ
end

sim = zeros(N)
sim[1] = θ0
@showprogress for i = 1:N-1
    x = sim[i]
    x_star = proposal(x)
    Δ =  logf(x_star) - logf(x)
    if log(rand())<Δ
        sim[i+1] = x_star
    else
        sim[i+1] = x
    end
end
L = logf.(sim)

# sim = sim[burnin+1:end]
figure(figsize=(12,4))
subplot(131)
plot(sim)
plot(1:length(sim), ones(length(sim))*μn, "--")
title("\$\\theta\$ Value")
xlabel("Iteration")
ylabel("\$\\theta\$")
subplot(132)
plot(L)
title("Log likelihood")
xlabel("Iteration")
ylabel("Log likelihood")
subplot(133)
hist(sim[burnin+1:end], density=true, bins=50)
y = pdf.(Normal(μn, sqrt(τn2)), LinRange(8.5,11.5,100))
plot(LinRange(8.5,11.5,100), y, label="Exact")
title("Distribution")
xlabel("\$\\theta\$")
ylabel("Density")
legend()
tight_layout()


include("center.jl")
include("cross_validation.jl")
include("mse.jl")

function ridge( X, y, lambda::Number )

    colmeans_X = center!( X )
    mean_y = center!( y )

    #A = X'*X + diagm( fill( lambda, size(X,2) ) )
    #w = A\(X'*y)

    A = [ X ; fill( sqrt( lambda ), ( 1, size(X,2) ) ) ]
    B = [ y ; fill( 0, ( 1, size(y,2) ) ) ]
    w = A\B

    broadcast!( +, X, X, colmeans_X )
    broadcast!( +, y, y, mean_y )

    bias = mean_y - colmeans_X * w

    return w, bias

end

function ridge( X, y, ls, gen = Kfold( size(X,1), 5 ) )
    if Base.length( ls ) == 1
        return ridge( X, y, ls[1] )
    end
    l_errs = Array( Float64, size(ls,1) )
    for i in 1:size(ls,1)
        l = ls[i]
        train_f(is) = ridge( X[is,:], y[is,:], l )
        error_f(model,is) = mse( predict( model, X[is,:] ), y[is,:] )
        l_errs[i] = mean( cross_validate( train_f, error_f, gen ) )
    end
    l = ls[ indmin( l_errs ) ]
    return ridge( X, y, l ), l, l_errs
end

function predict( w, bias, X )
    return broadcast( +, X*w, bias )
end

function ridgeKernel( K, y, lambda )
    m = size( K, 1 )

    colmeans_K = center!( K )
    mean_y = center!( y )

    for i = 1:m
        K[i,i] += lambda
    end
    w = K \ y
    for i = 1:m
        K[i,i] -= lambda
    end

    broadcast!( +, y, y, mean_y )
    broadcast!( +, K, K, colmeans_K )

    bias = mean_y - colmeans_K * w

    return w, bias
end

function ridgeKernel!( K, y, lambda )
    m = size( K, 1 )

    colmeans_K = center!( K )
    mean_y = center!( y )

    for i = 1:m
        K[i,i] += lambda
    end
    w = K \ y

    broadcast!( +, y, y, mean_y )

    bias = mean_y - colmeans_K * w

    return w, bias
end

function predictKernel( w, bias, K )
    return broadcast( +, K*w, bias )
end

function ridge( kernel::Function, X, y, lambda::Number )

    m = size( X, 1 )
    K = Array( eltype(y), m, m )
    for i = 1:m
        for j = i:m
            K[i,j] = kernel( X[i,:], X[j,:] )
            K[j,i] = K[i,j]
        end
    end

    w,bias = ridgeKernel!( K, y, lambda )

    return w, bias, kernel, X

end

function ridge( kernel::Function, X, y, ls, gen = Kfold( size(X,1), 5 ) )

    if Base.length( ls ) == 1
        return ridge( X, y, ls[1], kernel )
    end

    m = size( X, 1 )
    K = Array( eltype(y), m, m )
    for i = 1:m
        for j = i:m
            K[i,j] = kernel( X[i,:], X[j,:] )
            K[j,i] = K[i,j]
        end
    end

    l_errs = Array( Float64, size(ls,1) )
    for i in 1:size(ls,1)
        l = ls[i]
        train_f(is) = ridgeKernel( K[is,is], y[is,:], l ), is

        function error_f(v,is)
            model, train_is = v
            yh = predictKernel( model..., K[is,train_is] )
            return mse( yh, y[is,:] )
        end

        l_errs[i] = mean( cross_validate( train_f, error_f, gen ) )
    end
    l = ls[ indmin( l_errs ) ]
    w,bias = ridgeKernel!( K, y, l )
    return (w, bias, kernel, X), l, l_errs
end

function predict( w, bias, kernel, X_model, X )
    m = size( X, 1 )
    n = size( X_model, 1 )
    K = Array( eltype(w), m, n )
    for i = 1:m
        for j = 1:n
            K[i,j] = kernel( X[i,:], X_model[j,:] )
        end
    end
    return predictKernel( w, bias, K )
end

function predict( model, X )
    predict( model..., X )
end


export getNodalMassMatrix

function getNodalMassMatrix(mesh::TetraMesh, sigma::Array{Float64,1})

	nc = mesh.nc
  nn = mesh.nn

	M = zeros(4,4,nc)

	for i=1:nc
		Ti = mesh.Tetras[i,:]
	  	Xi = mesh.Points[Ti,:]

		M[:,:,i] = getLocalNodalMassMatrix(sigma[i],Xi)

	end
	nodalInd = mesh.Tetras

	ii = zeros(Int64,16*nc)
	jj = zeros(Int64,16*nc)
	# mm = zeros(Float64,16*nc)

	cnt = 1
	for i=1:nc
		for i2=1:4
			for i1=1:4
				ii[cnt] = nodalInd[i,i1]
				jj[cnt] = nodalInd[i,i2]
				# mm[cnt] = M[i1,i2,i]
				cnt    += 1
			end
		end
	end
  mm = vec(M)

	return sparse(ii,jj,mm,nn,nn)

end

# Nodal mass matrix is independent of geometry. Store mass matrix for unit coefficient and unit volume.
const M = (1/20) * [
  2 1 1 1
  1 2 1 1
  1 1 2 1
  1 1 1 2]

function getLocalNodalMassMatrix(sigma::Float64,X::Array{Float64})
	v =  getCellVolume(X[1,:], X[2,:], X[3,:], X[4,:])
	return (v * sigma) * M
end


"""
Compute the principal eigenvector and eigenvalue of a linear operator 𝕋, where 𝕋 is a Metzler matrix (i.e. off-diagonal components are nonnegative)

Denote a = -minimum(Diagonal(V)), which implies 𝕋 + a * I has all positive entries. Applying Perron Frobenus, there a unique largest eigenvalue for aI + 𝕋, which is real, and the correspondind eigenctor is strictly positive.
Note that, in particular, it is the eigenvalue with largest real part, and so this also correspoinds to the eigenvalue with largest real part of 𝕋, which happens to be real.
Denote η(𝕋) the eigenvalue with largest real part of a matrix and ρ(𝕋) the eigenvalue with largest modulus. We have

    η(𝕋) = ρ(𝕋 + a * I) - a

Moreover, the associated eigenvector is real and strictly positive.

Note that, when 𝕋 is generator, its rows sum to zero. This means that eigenvalue with largest real part is
    η(𝕋) = 0
In other words, all eigenvalues of 𝕋 have real part <= 0. This means that 𝕋 is a singular M matrix.
(another proof is to say that, for any s, sI - 𝕋 is a non-singular M-Matrix for any s> 0, since there exists x = e such that (sI - 𝕋) * x > 0).

"""
function principal_eigenvalue(𝕋; r0 = ones(size(𝕋, 1)))
    η, r = 0.0, r0
    a = - minimum(diag(𝕋))
    try
        vals, vecs = Arpack.eigs(𝕋 + a * I; v0 = collect(r0), nev = 1, which = :LM)
        η = vals[1]
        r = vecs[:, 1]
    catch
        vals, vecs = KrylovKit.eigsolve(𝕋 + a * I, collect(r0), 1, :LM; maxiter = size(𝕋, 1))
        η = vals[1]
        r = vecs[1]
    end
    abs(imag(η)) <= eps() || @warn "Principal Eigenvalue has an imaginary part"
    maximum(abs.(imag.(r))) <= eps() || @warn "Principal Eigenvector has an imaginary part"
    real(η) - a, abs.(r)
end


x = [1, 3, -5]
y = [4, -2, -1]
z = dot(x, y)
z = x'*y


function logsumexp(logx::Array{Float64,1})
    maxlog = maximum(logx)
    maxlog + log(sum(exp.(logx - maxlog)))
end

function logsumexp(logx::Array{Float64,2}, dim::Int)
    maxlog = maximum(logx, dim)
    maxlog + log.(sum(exp.(logx .- maxlog), dim))
end


function check(d, M, as)
    # as is sorted
    # check if all the cows can have a distance longer than d
    N = length(as)

    last = 1
    for i in 1:(M-1)
        crt = last + 1
        while crt ≤ N && as[crt] - as[last] < d
            crt += 1
        end

        if crt ≥ N
            return false
        else
            last = crt
        end
    end

    return true
end

function main()
    N, M = parse.(Int, split(readline()))
    as = [parse(Int, readline()) for _ in 1:N]
    sort!(as)

    d̲, d̄ = 0, maximum(as)
    while d̄ - d̲ > 1
        d = (d̄ + d̲) ÷ 2
        if check(d, M, as)
            d̲ = d
        else
            d̄ = d
        end
    end

    if check(d̄, M, as)
        println(d̄)
    else
        println(d̲)
    end
end

main()


## Show the effect of seeded PageRank on MNIST

## Load the MNIST Graph and draw it
##
using DelimitedFiles
using SparseArrays
using LinearAlgebra
ei,ej = copy.(eachcol(Int.(readdlm("mnist-train-4.edges"))))
xy = readdlm("mnist-train-4.xy")
n = maximum(ei)
A = sparse(ei,ej,1,n,n)
@assert issymmetric(A)

##
using Random
function simple_spectral_eigenvector(A,nsteps=500)
  @assert issymmetric(A) # expensive, but useful...
  Random.seed!(1)
  n = size(A,1)
  d = vec(sum(A,dims=1))
  nd2 = d'*d
  x = randn(n)
  # project x orthogonal to d
  x .-= ((d'*x)/(d'*d)).*d
  x ./= x'*(d.*x) # normalize
  for i=1:nsteps
    x = (A*x)./d          # matrix-vector
    x .-= ((d'*x)/nd2).*d
    x ./= sqrt(x'*(d.*x)) # normalize
    # make sure the first component is positive
    x .*= sign(x[1])
  end

  return x
end
f = simple_spectral_eigenvector(A)
##
using MatrixNetworks
f2 = fiedler_vector(A)[1]
# make sure the first component is positive
f2 .*= sign(f2[1])
##
norm(f-f2)
## why different? convergence is slow :(
f = simple_spectral_eigenvector(A,2500)
norm(f-f2)
## why different? convergence is slow :(
f = simple_spectral_eigenvector(A,4000)
norm(f-f2)
##
using Plots
x = f
p = sortperm(x) # this gets things to draw in the right z-order
scatter(xy[p,1],xy[p,2], marker_z=x[p],colorbar=false,alpha=0.5,
  color=:magma,
  size=(1200,1200),
  markerstrokecolor=:white,markerstrokewidth=0.0,label="",framestyle=:none,
  background=nothing,dpi=220)
savefig("spectral-mnist-500.png")
##
using Plots
x = f2
p = sortperm(x) # this gets things to draw in the right z-order
scatter(xy[p,1],xy[p,2], marker_z=x[p],colorbar=false,alpha=0.5,
  color=:magma,
  size=(1200,1200),
  markerstrokecolor=:white,markerstrokewidth=0.0,label="",framestyle=:none,
  background=nothing,dpi=220)
savefig("spectral-mnist-converged.png")
## Let's watch spectra converge
using Plots
using GraphRecipes
Random.seed!(1)
function spectral_convergence_animation(A,xy,nsteps_per_frame,nframes)
  @assert issymmetric(A) # expensive, but useful...
  Random.seed!(1)
  n = size(A,1)
  d = vec(sum(A,dims=1))
  sumd = sum(d)
  x = randn(n)
  # project x orthogonal to d
  x .-= (d'*x)/sum(d)
  x ./= x'*(d.*x) # normalize

  anim = @animate for i=1:nframes

    f = zeros(n)
    for j=1:nsteps_per_frame
      x = (A*x)./d          # matrix-vector
      x .-= (d'*x)/sum(d)   # project out
      x ./= sqrt(x'*(d.*x)) # normalize
      # make sure the first component is positive
      #x .*= sign(x[1])
      f .+= x./nsteps_per_frame
    end

    p = sortperm(f) # this gets things to draw in the right z-order
    scatter(xy[p,1],xy[p,2], marker_z=f[p],colorbar=false,alpha=0.5,
      color=:magma,
      size=(1200,1200),dpi=100,
      markerstrokecolor=nothing,markerstrokewidth=0,label="",framestyle=:none,background=:black)
  end
  return anim
end
anim = spectral_convergence_animation(A,xy,10,500)
gif(anim,"spectral-convergence-mnist.gif",fps=30)
##
anim = spectral_convergence_animation(A,xy,1,500)
gif(anim,"spectral-convergence-mnist-1-500.gif",fps=30)

## Now show the entire graph
using GraphRecipes
using LinearAlgebra
using Plots

ei, ej = findnz(A)
graphplot(ei, ej, x =xy[:,1], y=xy[:,2],
  markercolor=:black, markerstrokecolor=:white,
  size=(1200,1200),dpi=300,
  markersize=0, linecolor=1, linealpha=0.2, linewidth=0.5,
  markeralpha=0.2,colorbar=false,
  axis_buffer=0.02, background=nothing)
p=sortperm(pr)
scatter!(xy[p,1],xy[p,2],marker_z=-abs.(log10.(pr[p])).^0.5,markerstrokewidth=0,alpha=0.5)
savefig("pagerank-mnist.png")


# ============================================================
# Thermodynamics
# ============================================================

"""
$(SIGNATURES)

Calculate heat capacity ratio (monatomic gas)
"""
function heat_capacity_ratio(K, D::T) where {T<:Integer}
    γ = begin
        if D == 1
            (K + 3.0) / (K + 1.0)
        elseif D == 2
            (K + 4.0) / (K + 2.0)
        elseif D == 3
            (K + 5.0) / (K + 3.0)
        end
    end

    return γ
end

"""
$(SIGNATURES)

Calculate heat capacity ratio (diatomic gas)
"""
function heat_capacity_ratio(K, Nr, D::T) where {T<:Integer}
    γ = begin
        if D == 1
            (K + 3.0 + Nr) / (K + 1.0 + Nr)
        elseif D == 2
            (K + 4.0 + Nr) / (K + 2.0 + Nr)
        elseif D == 3
            (K + 5.0 + Nr) / (K + 3.0 + Nr)
        end
    end

    return γ
end

"""
$(SIGNATURES)

Calculate speed of sound
"""
sound_speed(λ::Real, γ::Real) = (0.5 * γ / λ)^0.5

"""
$(SIGNATURES)
"""
sound_speed(prim::AV, γ) = sound_speed(prim[end], γ)

"""
$(TYPEDSIGNATURES)

Calculate sound speed in mixture
"""
function sound_speed(prim::AM, γ)
    c = similar(prim, axes(prim, 2))
    for j in eachindex(c)
        c[j] = sound_speed(prim[end, j], γ)
    end

    return maximum(c)
end


using GeneralizedSampling
import IntervalWavelets: unit
using Base.Test

println("Testing least squares solver...")

#=
Test the least squares solvers:
- Solution with Freq2Wave and the collected version should be identical.
- When sampling from a scaling function used in the Freq2Wave type, the solution should be a unit vector.
=#

# ------------------------------------------------------------
# Reconstruct unit vector

J = 3
M = 2^(J+1)
N = 2^J
xi = grid(M, 0.5)
wavename = "db2"
vm = van_moment(wavename)

T = Freq2Wave(xi, wavename, J)

for k = 1:N
	if k <= vm
		F = FourScalingFunc(xi, wavename, 'L', J)
		f = F[:,k]
	elseif k > N-vm
		F = FourScalingFunc(xi, wavename, 'R', J)
		f = F[:,k-N+vm]
	else
		trans = k - div(N,2) - 1
		f = FourDaubScaling(xi, vm, J, trans)
	end
	y = T \ f

	u = unit(Complex{Float64}, length(y), k)
	#@show k, norm( y - u )
	@test_approx_eq_eps y u 1e-4
end

# ------------------------------------------------------------

begin
	J = 10
	M = 2^(J+2)
	N = 2^J

	# Uniform samples
	xi = grid(M, 0.5)
	TU = Freq2Wave(xi, "Haar", J)
	AU = collect(TU)

	b = rand(M)
	x0 = rand(N)

	x1 = AU \ b
	x2 = TU \ b

	#@show norm(x1 - x2, Inf)
	@test_approx_eq_eps x1 x2 1e-5

	# Non-uniform samples: Requires high M/N ratio
	K = N/2
	xi = 2*K*rand(M) - K
	sort!(xi)
	TNU = Freq2Wave(xi, "Haar", J, K)
	ANU = collect(TNU)
	b = rand(M) + rand(M)*im

	bb = b .* get(TNU.weights)
	z1 = ANU \ bb
	z2 = TNU \ b

	#@show norm(z1 - z2, Inf)
	@test_approx_eq_eps z1 z2 1e-4
end

# 2D
begin
	J = 5
	M = 2^(J+1)
	N = 2^J

	# Uniform samples
	xi = grid((M,M), 0.5)
	TU = Freq2Wave(xi, "Haar", J)
	AU = collect(TU)
	b = rand(size(TU,1))

	x1 = AU \ b
	x2 = TU \ b

	#@show norm(x1 - vec(x2), Inf)
	@test_approx_eq_eps x1 x2 1e-5

	# Non-uniform samples: Requires high M/N ratio
	K = N/2
	xi = 2*K*rand(M^2,2) - K
	TNU = Freq2Wave(xi, "Haar", J, K)
	ANU = collect(TNU)
	b = rand(M^2) + rand(M^2)*im

	bb = b .* get(TNU.weights)
	z1 = ANU \ bb
	z2 = TNU \ b

	#@show norm(z1 - z2, Inf)
	@test_approx_eq_eps z1 z2 1e-4
end


using ParetoFront
using Test

# Suppose we are in two dimensions and want min, min to dominate.
# We know then, that the point that has a minimum value of y will dominate
# all points to the right of it, ie  at x > x(at ymin) will be dominated
# and so for all x on the Pareto front we must have x < x(at ymin)
# and also the point with the minimum value of x will dominate all
# points that are above it, ie that y > y(at xmin) will be dominated
# and so for all y on the Pareto front we must have y < y(at xmin)
function test_min_blob(n)
    v = randn(n, 2)
    ((xmin, ymin), (ix, iy)) = findmin(v, dims=1)
    ymax = v[ix[1], 2]
    xmax = v[iy[1], 1]
    pareto = get_pareto(v, min_idxs=[1,2], max_idxs=[])
    front = Matrix(reduce(hcat, collect(pareto))')
    pxmax, pymax = maximum(front, dims=1)
    return pxmax <= xmax && pymax <= ymax
end

# the test for max, max as optimal
function test_max_blob(n)
    v = randn(n, 2)
    ((xmax, ymax), (ix, iy)) = findmax(v, dims=1)
    ymin = v[ix[1], 2]
    xmin = v[iy[1], 1]
    pareto = get_pareto(v, min_idxs=[], max_idxs=[1,2])
    front = Matrix(reduce(hcat, collect(pareto))')
    pxmin, pymin = minimum(front, dims=1)
    return pxmin >= xmin && pymin >= ymin
end

@testset "ParetoFront.jl" begin
    for i in 10 .^ (1:6)
        @test test_min_blob(i)
        @test test_max_blob(i)
    end
end


# Test that the minimum image distance is correctly calculated

@testset "Minimum image distance" begin
@testset "fcc" begin
    @test MkCell.score_image_distance(fcc1*cellfcc)   ≈ 0.7071067811865476
    @test MkCell.score_image_distance(fcc4*cellfcc)   ≈ 1.000
    @test MkCell.score_image_distance(fcc7*cellfcc)   ≈ 1.224744871391589
    @test MkCell.score_image_distance(fcc13*cellfcc)  ≈ 1.5811388300841898
    @test MkCell.score_image_distance(fcc23*cellfcc)  ≈ 1.7320508075688772
    @test MkCell.score_image_distance(fcc26*cellfcc)  ≈ 1.8708286933869707
    @test MkCell.score_image_distance(fcc32*cellfcc)  ≈ 2.000
    @test MkCell.score_image_distance(fcc38*cellfcc)  ≈ 2.1213203435596424
    @test MkCell.score_image_distance(fcc50*cellfcc)  ≈ 2.345207879911715
    @test MkCell.score_image_distance(fcc70*cellfcc)  ≈ 2.5495097567963922
    @test MkCell.score_image_distance(fcc92*cellfcc)  ≈ 2.7386127875258306
    @test MkCell.score_image_distance(fcc96*cellfcc)  ≈ 2.8284271247461903
    @test MkCell.score_image_distance(fcc100*cellfcc) ≈ 2.9154759474226504
    @test MkCell.score_image_distance(fcc108*cellfcc) ≈ 3.000
end;
@testset "cubic" begin
    @test MkCell.score_image_distance(cubic1*cellcubic)   ≈ 1.000
    @test MkCell.score_image_distance(cubic3*cellcubic)   ≈ 1.4142135623730951
    @test MkCell.score_image_distance(cubic7*cellcubic)   ≈ 1.7320508075688772
    @test MkCell.score_image_distance(cubic8*cellcubic)   ≈ 2.000
    @test MkCell.score_image_distance(cubic13*cellcubic)  ≈ 2.23606797749979
    @test MkCell.score_image_distance(cubic14*cellcubic)  ≈ 2.449489742783178
    @test MkCell.score_image_distance(cubic24*cellcubic)  ≈ 2.8284271247461903
    @test MkCell.score_image_distance(cubic27*cellcubic)  ≈ 3.000
    @test MkCell.score_image_distance(cubic33*cellcubic)  ≈ 3.1622776601683795
    @test MkCell.score_image_distance(cubic38*cellcubic)  ≈ 3.3166247903554
    @test MkCell.score_image_distance(cubic52*cellcubic)  ≈ 3.605551275463989
    @test MkCell.score_image_distance(cubic56*cellcubic)  ≈ 3.7416573867739413
    @test MkCell.score_image_distance(cubic64*cellcubic)  ≈ 4.000
    @test MkCell.score_image_distance(cubic73*cellcubic)  ≈ 4.123105625617661
    @test MkCell.score_image_distance(cubic76*cellcubic)  ≈ 4.242640687119285
    @test MkCell.score_image_distance(cubic93*cellcubic)  ≈ 4.358898943540674
    @test MkCell.score_image_distance(cubic100*cellcubic) ≈ 4.47213595499958
end;
@testset "bcc" begin
    @test MkCell.score_image_distance(bcc1*cellbcc)   ≈ 0.8660254037844386
    @test MkCell.score_image_distance(bcc2*cellbcc)   ≈ 1.000
    @test MkCell.score_image_distance(bcc6*cellbcc)   ≈ 1.4142135623730951
    @test MkCell.score_image_distance(bcc9*cellbcc)   ≈ 1.6583123951777
    @test MkCell.score_image_distance(bcc16*cellbcc)  ≈ 2.000
    @test MkCell.score_image_distance(bcc21*cellbcc)  ≈ 2.179449471770337
    @test MkCell.score_image_distance(bcc25*cellbcc)  ≈ 2.23606797749979
    @test MkCell.score_image_distance(bcc28*cellbcc)  ≈ 2.449489742783178
    @test MkCell.score_image_distance(bcc35*cellbcc)  ≈ 2.598076211353316
    @test MkCell.score_image_distance(bcc48*cellbcc)  ≈ 2.8284271247461903
    @test MkCell.score_image_distance(bcc53*cellbcc)  ≈ 2.958039891549808
    @test MkCell.score_image_distance(bcc54*cellbcc)  ≈ 3.000
    @test MkCell.score_image_distance(bcc65*cellbcc)  ≈ 3.278719262151
    @test MkCell.score_image_distance(bcc84*cellbcc)  ≈ 3.3166247903554
    @test MkCell.score_image_distance(bcc91*cellbcc)  ≈ 3.570714214271425
    @test MkCell.score_image_distance(bcc103*cellbcc) ≈ 3.605551275463989
    @test MkCell.score_image_distance(bcc107*cellbcc) ≈ 3.7416573867739413
    @test MkCell.score_image_distance(bcc112*cellbcc) ≈ 3.840572873934304
    @test MkCell.score_image_distance(bcc128*cellbcc) ≈ 4.000
end;
@testset "hcp" begin
    @test MkCell.score_image_distance(hcp1*cellhcp)   ≈ 0.5612
    @test MkCell.score_image_distance(hcp3*cellhcp)   ≈ 0.9164357624332797
    @test MkCell.score_image_distance(hcp6*cellhcp)   ≈ 0.972026913207654
    @test MkCell.score_image_distance(hcp8*cellhcp)   ≈ 1.1224
    @test MkCell.score_image_distance(hcp11*cellhcp)  ≈ 1.3359232113660824
    @test MkCell.score_image_distance(hcp14*cellhcp)  ≈ 1.4490121692610682
    @test MkCell.score_image_distance(hcp16*cellhcp)  ≈ 1.4847956357694483
    @test MkCell.score_image_distance(hcp21*cellhcp)  ≈ 1.7448417081978143
    @test MkCell.score_image_distance(hcp27*cellhcp)  ≈ 1.9168629232855088
    @test MkCell.score_image_distance(hcp37*cellhcp)  ≈ 2.023435375790391
    @test MkCell.score_image_distance(hcp39*cellhcp)  ≈ 2.0746696957989883
    @test MkCell.score_image_distance(hcp42*cellhcp)  ≈ 2.149232371491428
    @test MkCell.score_image_distance(hcp53*cellhcp)  ≈ 2.221293593081893
    @test MkCell.score_image_distance(hcp62*cellhcp)  ≈ 2.358820914496619
    @test MkCell.score_image_distance(hcp65*cellhcp)  ≈ 2.4246611199643273
    @test MkCell.score_image_distance(hcp69*cellhcp)  ≈ 2.4462140871150257
    @test MkCell.score_image_distance(hcp76*cellhcp)  ≈ 2.6122438375210435
    @test MkCell.score_image_distance(hcp80*cellhcp)  ≈ 2.6718464227321648
    @test MkCell.score_image_distance(hcp88*cellhcp)  ≈ 2.7301481180819964
    @test MkCell.score_image_distance(hcp90*cellhcp)  ≈ 2.749307287299839
    @test MkCell.score_image_distance(hcp100*cellhcp) ≈ 2.806
end;
end;


"Brownian motion process."
immutable BrownianMotion{Y} <: AItoProcess
  y0::Y
end

BrownianMotion() = BrownianMotion(0.0)

convert(::Type{ItoProcess}, bm::BrownianMotion) =
    ItoProcess((t,dt,b,db,y)->db, bm.y0)

function distribution(bm::BrownianMotion, t)
  if t == 0
    Constant(bm.y0)
  elseif ndims(bm.y0) == 0
    Normal(bm.y0, sqrt(t))
  else
    MvNormal(bm.y0, eye(length(bm.y0)) * t)
  end
end

# BrownianMotionWithDrift

immutable BrownianMotionWithDrift{Y,S} <: AItoProcess
  mu::Y
  sigma::S
  y0::Y
end

BrownianMotionWithDrift(mu, sigma) = BrownianMotionWithDrift(mu, sigma, 0.0)

convert(::Type{ItoProcess}, bm::BrownianMotionWithDrift) =
    ItoProcess((t,dt,b,db,y)-> (bm.mu * dt) .+ (bm.sigma * db), bm.y0)

function distribution(bm::BrownianMotionWithDrift, t)
    if t == 0
      Constant(bm.y0)
    elseif ndims(bm.y0) == 0
      Normal(bm.y0 + bm.mu * t, bm.sigma * sqrt(t))
    else
      MvNormal(bm.y0 + bm.mu * t, bm.sigma *bm.sigma' * t)
    end
end

solution(p::BrownianMotionWithDrift, t, b) =
  p.y0 + t * p.mu + p.sigma * b

# Geometric Brownian Motion

type GeometricBrownianMotion <: AItoProcess
    mu::Float64
    sigma::Float64
    y0::Float64
end

convert(::Type{ItoProcess}, bm::GeometricBrownianMotion) =
    ItoProcess((t,dt,b,db,y)-> (bm.mu * dt) * y + bm.sigma * (y .* db), bm.y0)

distribution(bm::GeometricBrownianMotion, t) =
    t == 0 ? Constant(bm.y0) :
    LogNormal(log(bm.y0) + t * (bm.mu - bm.sigma^2 / 2), bm.sigma * sqrt(t))

solution(p::GeometricBrownianMotion, t, b) =
    p.y0 * exp(t * (p.mu - p.sigma^2 / 2) + p.sigma * b)

"Ito integral of f(t, B)"
immutable ItoIntegral{F} <: AItoProcess
  f::F
end

convert{F}(::Type{ItoProcess}, ii::ItoIntegral{F}) =
    ItoProcess((t,dt,b,db,y)->ii.f(t, b).*db, 0.)

"Process corresponding to solutions of stochastic differential equation
 dy = f(t, y) dt + g(t,y) db."
immutable SDE{F,G} <: AItoProcess
  f::F
  g::G
  y0::Float64
end

convert{F,G}(::Type{ItoProcess}, sde::SDE{F,G}) =
    ItoProcess((t,dt,b,db,y)->sde.f(t, y).*dt + sde.g(t, y).*db, sde.y0)


mutable struct ListNode{T<:Number}
    val :: T
    next:: Union{ListNode, Nothing}

    ListNode(data::Array{T,1}) where T<:Number = begin
        (head, rest) = Iterators.peel(data)
        if isempty(rest)
            new{T}(head::T, nothing)
        else
            new{T}(head, ListNode(collect(rest) :: Array{T,1}))
        end
    end
end

function node2list(x::Union{ListNode{T}, Nothing}) :: Array{T} where {T<:Number}
    result = []
    while x.next != nothing
        push!(result, x.val)
        x = x.next
    end
    push!(result, x.val)
    result
end

function addTwoNumbers(a::ListNode{T}, b::ListNode{T}) :: ListNode{T} where T<:Number
    this_a::Union{ListNode, Nothing}, this_b::Union{ListNode, Nothing} = a, b
    flag = false
    result::Array{T} = []
    while true
        (this_a == nothing && this_b == nothing) && begin
            flag && push!(result,1)
            return ListNode(result)
        end

        this_a == nothing && begin
            append!(result, begin
                    cache = node2list(this_b)
                    if flag cache[1] += 1 end
                    cache
                    end)
            return ListNode(result)
        end

        this_b == nothing && begin
            append!(result, begin
                    cache = node2list(this_a)
                    if flag cache[1] += 1 end
                    cache
                    end)
            return ListNode(result)
        end

        push!(result, this_a.val + this_b.val + (flag ? 1 : 0))
        if last(result) >= 10
            result[lastindex(result)] -= 10
            flag = true
        else
            flag = false
        end

        this_a, this_b = this_a.next, this_b.next
    end
    ListNode(result)
end


using LinearAlgebra
function bernsteinBasis(u, deg)
    """
    Function returns the values of  Bernstein basis of degree deg at points u
    Algorithm A1.3 in Piegl & Tiller
    xi is a 1D array
    """
    lenU = length(u)
    B = zeros(lenU, deg+1)
    B[:, 1] = ones(lenU)
    u1 = ones(lenU) - u
    u2 = ones(lenU) + u

    for j in 1:deg
        saved = zeros(lenU)
        for k in 1:j
            temp = B[:,k]
            B[:,k] = saved + u1.*temp
            saved  = u2.*temp
        end
        B[:,j+1] = saved
    end
    B = B./(2^deg)

    #calculate the 1st derivative of Bernstein polynomials
    dB = zeros(lenU, deg)
    dB[:,1] = ones(lenU)
    for j in 1:deg-1
        saved = zeros(lenU)
        for k in 0:j-1
            temp = dB[:,k+1]
            dB[:,k+1] = saved + u1.*temp
            saved = u2.*temp
        end
        dB[:,j+1] = saved
    end
    dB = dB./(2^deg)
    dB = hcat(zeros(lenU,1), dB, zeros(lenU, 1))
    dB = (dB[:,1:end-1]-dB[:,2:end])*deg

    #calculate the 2nd derivative of Bernstein polynomials
    if deg>1
        ddB = zeros(lenU, deg-1)
        ddB[:,1] = ones(lenU)
        for j=1:deg-2
            saved = zeros(lenU)
            for k in 0:j-1
                temp = ddB[:,k+1]
                ddB[:,k+1] = saved + u1.*temp
                saved = u2.*temp
            end
            ddB[:,j+1] = saved
        end
        ddB = ddB./(2^deg)
        ddB = hcat(zeros(lenU,2), ddB, zeros(lenU,2))
        ddB = (ddB[:,1:end-2]-2*ddB[:,2:end-1]+ddB[:,3:end])*deg*(deg-1)
    else
        ddB = zeros(lenU,deg+1)
    end
    return B, dB, ddB

end

function bezierExtraction(knot, deg)
    """
    Bezier extraction
    Based on Algroithm 1, from Borden - Isogeometric finite element data
    structures based on Bezier extraction
    """
    m = length(knot)-deg-1;
    a = deg + 1;
    b = a + 1;
    nb = 1;
    nb_final = length(unique(knot))-1
    C = zeros(deg+1,deg+1,nb_final)
    C[:,:,1] = Matrix{Float64}(I, deg+1, deg+1)

    while b<=m
        C[:,:,nb+1] = Matrix{Float64}(I, deg+1, deg+1)
        i = b
        while (b<=m) && (knot[b+1] == knot[b])
            b = b + 1
        end
        multiplicity = b - i + 1
        alphas = zeros(deg-multiplicity)
        if multiplicity < deg
            numerator = knot[b] - knot[a]
            for j in deg:-1:multiplicity+1
                alphas[j-multiplicity] = numerator/(knot[a+j]-knot[a])
            end
            r = deg - multiplicity
            for j in 1:r
                save = r - j + 1
                s = multiplicity + j
                for k in deg+1:-1:s+1
                    alpha = alphas[k-s]
                    C[:,k,nb] = alpha*C[:,k,nb]+(1-alpha)*C[:,k-1,nb]
                end
                if b<=m
                    C[save:save+j,save,nb+1] = C[deg-j+1:deg+1,deg+1,nb]
                end
            end
            nb = nb + 1
            if b <= m
                a = b
                b = b + 1
            end
        elseif multiplicity==deg
            if b <= m
                nb = nb + 1
                a = b
                b = b + 1
            end
        end
    end
    @assert nb_final == nb
    return C, nb
end

"""
Create the extended knot vector (Subsection 4.3.2 in Scott - Isogeometric
data structures based on the Bézier extraction of T-Splines)
"""
function formExtendedKnot(localKnot, p)
    # Repeat the first knot (if needed) so that it appears p+1 times
    firstKnot = localKnot[1]
    indexFirst = findall(localKnot.==firstKnot)
    numRep = length(indexFirst)
    numNewRepFirst = p+1-numRep

    #repeat the last knot (if needed) so that it appears p+1 times
    lastKnot = localKnot[end]
    indexLast = findall(localKnot.==lastKnot)
    numRep = length(indexLast)
    numNewRepLast = p+1-numRep

    #form the extended knot vector
    extendedKnot = vcat(firstKnot*ones(numNewRepFirst), localKnot, lastKnot*ones(numNewRepLast))
    indexFun = numNewRepFirst + 1
    return extendedKnot, indexFun
end

"""
Compute the Bézier extraction operator corresponding to the basis functions of
degree p with local knot vector localKnot
"""
function bezierExtractionLocal(localKnot, p)
    extendedKnot, indexFun = formExtendedKnot(localKnot, p)
    #perform Bézier extraction and return the basis with index numNewRepFirst+1
    C_temp, nb = bezierExtraction(extendedKnot, p)
    IEN, _ = makeIEN(extendedKnot, nb, p)
    C = zeros(p+1, nb)
    for indexSpan = 1:nb
        C_index = findfirst(IEN[indexSpan,:].==indexFun)
        C[:, indexSpan] = C_temp[C_index,:,indexSpan]
    end
    return C
end


module problem107

description = """
Minimal network
Problem 107

The following undirected network consists of seven vertices and twelve edges with a total weight of 243.

The same network can be represented by the matrix below.
    	A	B	C	D	E	F	G
A	-	16	12	21	-	-	-
B	16	-	-	17	20	-	-
C	12	-	-	28	-	31	-
D	21	17	28	-	18	19	23
E	-	20	-	18	-	-	11
F	-	-	31	19	-	-	27
G	-	-	-	23	11	27	-

However, it is possible to optimise the network by removing some edges and still ensure that all points on the network remain connected. The network which achieves the maximum saving is shown below. It has a weight of 93, representing a saving of 243 − 93 = 150 from the original network.

Using network.txt (right click and 'Save Link/Target As...'), a 6K text file containing a network with forty vertices, and given in matrix form, find the maximum saving which can be achieved by removing redundant edges whilst ensuring that the network remains connected.
"""

immutable Edge
  from :: Int
  to :: Int
  cost :: Int
end

immutable Tree
  edges :: Vector{Edge}
end

immutable Forest
  trees :: Vector{Tree}
end

# Find the index of the tree containing node in the forest, or -1 if not found
function tree_index(node :: Int, f :: Forest)
  for i = 1:length(f.trees)
    for j = 1:length(f.trees[i].edges)
      if (f.trees[i].edges[j].from == node || f.trees[i].edges[j].to == node)
        return i
      end
    end
  end
  return -1
end

# Tests if adding Edge e to Forest would connect two previously disconnected trees
function makes_new_connection(e :: Edge, f :: Forest)
  t1 = tree_index(e.from, f)
  t2 = tree_index(e.to, f)
  return t1 < 0 || t2 < 0 || t1 != t2
end

# Creates Add Edge to Forest, joining any trees that become linked
function add_edge(e :: Edge, f :: Forest)
  t1 = tree_index(e.from, f)
  t2 = tree_index(e.to, f)

  # Simple case, neither from/to existed previously, add a new tree
  if t1 < 0 && t2 < 0
    return Forest([Tree([e]); f.trees])
  end

  # One of the trees exists - make a copy and update
  newtrees = copy(f.trees)
  if t1 > 0 && t2 < 0
    newtrees[t1] = Tree([e; f.trees[t1].edges])
  elseif t1 < 0 && t2 > 0
    newtrees[t2] = Tree([e; f.trees[t2].edges])
  else
    newtrees[t1] = Tree([e; f.trees[t1].edges; f.trees[t2].edges])
    splice!(newtrees, t2, Tree[])
  end
  return Forest(newtrees)
end

# Convert from matrix representation to Edge vector
function matrix_to_edges(connections :: Matrix{Int})
  edges = Edge[]
  for i = 1:size(connections, 1)
    for j = i+1:size(connections, 2)
      if connections[i, j] > 0
        push!(edges, Edge(i, j, connections[i,j]))
      end
    end
  end
  return edges
end

# Calculate the minimum spanning tree by adding the minimum cost edge to a forest if it links up previously disconnected trees
function minimum_spanning_tree(connections :: Matrix{Int})
  edges = matrix_to_edges(connections)
  sort!(edges, by = x -> x.cost)
  f = Forest(Tree[])
  for e in edges
    if makes_new_connection(e, f)
      f = add_edge(e, f)
    end
  end
  return f.trees[1]
end

function read_network()
  m = readcsv("network.txt")
  m[m .== "-"] = -1
  connections = int(m)
  return connections
end

function solve107()
  network = read_network()
  t = minimum_spanning_tree(network)
  initial_cost = sum([e.cost for e in matrix_to_edges(network)])
  final_cost = sum([e.cost for e in t.edges])
  println("$initial_cost - $final_cost = $(initial_cost - final_cost)")
end

using Base.Test

testcase = [
-1	16	12	21	-1	-1	-1
16	-1	-1	17	20	-1	-1
12	-1	-1	28	-1	31	-1
21	17	28	-1	18	19	23
-1	20	-1	18	-1	-1	11
-1	-1	31	19	-1	-1	27
-1	-1	-1	23	11	27	-1]

@test sum([e.cost for e in minimum_spanning_tree(testcase).edges]) == 93

end


# # Kernel Ridge Regression
#
# Building on linear regression, we can fit non-linear data sets by introducing a feature space. In a higher-dimensional feature space, we can overfit the data; ridge regression introduces regularization to avoid this. In this notebook we show how we can use KernelFunctions.jl for *kernel* ridge regression.

## Loading and setup of required packages
using KernelFunctions
using LinearAlgebra
using Distributions

## Plotting
using Plots;
default(; lw=2.0, legendfontsize=11.0, ylims=(-150, 500));

using Random: seed!
seed!(42);

##
# ## Toy data
# Here we use a one-dimensional toy problem. We generate data using the fourth-order polynomial $f(x) = (x+4)(x+1)(x-1)(x-3)$:

f_truth(x) = (x + 4) * (x + 1) * (x - 1) * (x - 3)

x_train = -5:0.5:5
x_test = -7:0.1:7

noise = rand(Uniform(-20, 20), length(x_train))
y_train = f_truth.(x_train) + noise
y_test = f_truth.(x_test)

plot(x_test, y_test; label=raw"$f(x)$")
scatter!(x_train, y_train; seriescolor=1, label="observations")

##
# ## Linear regression
# For training inputs $\mathrm{X}=(\mathbf{x}_n)_{n=1}^N$ and observations $\mathbf{y}=(y_n)_{n=1}^N$, the linear regression weights $\mathbf{w}$ using the least-squares estimator are given by
# ```math
# \mathbf{w} = (\mathrm{X}^\top \mathrm{X})^{-1} \mathrm{X}^\top \mathbf{y}
# ```
# We predict at test inputs $\mathbf{x}_*$ using
# ```math
# \hat{y}_* = \mathbf{x}_*^\top \mathbf{w}
# ```
# This is implemented by `linear_regression`:

function linear_regression(X, y, Xstar)
    weights = (X' * X) \ (X' * y)
    return Xstar * weights
end;

# A linear regression fit to the above data set:

y_pred = linear_regression(x_train, y_train, x_test)
scatter(x_train, y_train; label="observations")
plot!(x_test, y_pred; label="linear fit")

##
# ## Featurization
# We can improve the fit by including additional features, i.e. generalizing to $\tilde{\mathrm{X}} = (\phi(x_n))_{n=1}^N$, where $\phi(x)$ constructs a feature vector for each input $x$. Here we include powers of the input, $\phi(x) = (1, x, x^2, \dots, x^d)$:

function featurize_poly(x; degree=1)
    return repeat(x, 1, degree + 1) .^ (0:degree)'
end

function featurized_fit_and_plot(degree)
    X = featurize_poly(x_train; degree=degree)
    Xstar = featurize_poly(x_test; degree=degree)
    y_pred = linear_regression(X, y_train, Xstar)
    scatter(x_train, y_train; legend=false, title="fit of order $degree")
    return plot!(x_test, y_pred)
end

plot((featurized_fit_and_plot(degree) for degree in 1:4)...)

##
# Note that the fit becomes perfect when we include exactly as many orders in the features as we have in the underlying polynomial (4).
#
# However, when increasing the number of features, we can quickly overfit to noise in the data set:

featurized_fit_and_plot(20)

##
# ## Ridge regression
# To counteract this unwanted behaviour, we can introduce regularization. This leads to *ridge regression* with $L_2$ regularization of the weights ([Tikhonov regularization](https://en.wikipedia.org/wiki/Tikhonov_regularization)).
# Instead of the weights in linear regression,
# ```math
# \mathbf{w} = (\mathrm{X}^\top \mathrm{X})^{-1} \mathrm{X}^\top \mathbf{y}
# ```
# we introduce the ridge parameter $\lambda$:
# ```math
# \mathbf{w} = (\mathrm{X}^\top \mathrm{X} + \lambda \mathbb{1})^{-1} \mathrm{X}^\top \mathbf{y}
# ```
# As before, we predict at test inputs $\mathbf{x}_*$ using
# ```math
# \hat{y}_* = \mathbf{x}_*^\top \mathbf{w}
# ```
# This is implemented by `ridge_regression`:

function ridge_regression(X, y, Xstar, lambda)
    weights = (X' * X + lambda * I) \ (X' * y)
    return Xstar * weights
end

function regularized_fit_and_plot(degree, lambda)
    X = featurize_poly(x_train; degree=degree)
    Xstar = featurize_poly(x_test; degree=degree)
    y_pred = ridge_regression(X, y_train, Xstar, lambda)
    scatter(x_train, y_train; legend=false, title="\$\\lambda=$lambda\$")
    return plot!(x_test, y_pred)
end

plot((regularized_fit_and_plot(20, lambda) for lambda in (1e-3, 1e-2, 1e-1, 1))...)

##
# ## Kernel ridge regression
# Instead of constructing the feature matrix explicitly, we can use *kernels* to replace inner products of feature vectors with a kernel evaluation: $\langle \phi(x), \phi(x') \rangle = k(x, x')$ or $\tilde{\mathrm{X}} \tilde{\mathrm{X}}^\top = \mathrm{K}$, where $\mathrm{K}_{ij} = k(x_i, x_j)$.
#
# To apply this "kernel trick" to ridge regression, we can rewrite the ridge estimate for the weights
# ```math
# \mathbf{w} = (\mathrm{X}^\top \mathrm{X} + \lambda \mathbb{1})^{-1} \mathrm{X}^\top \mathbf{y}
# ```
# using the [matrix inversion lemma](https://tlienart.github.io/pub/csml/mtheory/matinvlem.html#basic_lemmas)
# as
# ```math
# \mathbf{w} = \mathrm{X}^\top (\mathrm{X} \mathrm{X}^\top + \lambda \mathbb{1})^{-1} \mathbf{y}
# ```
# where we can now replace the inner product with the kernel matrix,
# ```math
# \mathbf{w} = \mathrm{X}^\top (\mathrm{K} + \lambda \mathbb{1})^{-1} \mathbf{y}
# ```
# And the prediction yields another inner product,
# ```math
# \hat{y}_* = \mathbf{x}_*^\top \mathbf{w} = \langle \mathbf{x}_*, \mathbf{w} \rangle = \mathbf{k}_* (\mathrm{K} + \lambda \mathbb{1})^{-1} \mathbf{y}
# ```
# where $(\mathbf{k}_*)_n = k(x_*, x_n)$.
#
# This is implemented by `kernel_ridge_regression`:

function kernel_ridge_regression(k, X, y, Xstar, lambda)
    K = kernelmatrix(k, X)
    kstar = kernelmatrix(k, Xstar, X)
    return kstar * ((K + lambda * I) \ y)
end;

# Now, instead of explicitly constructing features, we can simply pass in a `PolynomialKernel` object:

function kernelized_fit_and_plot(kernel, lambda=1e-4)
    y_pred = kernel_ridge_regression(kernel, x_train, y_train, x_test, lambda)
    if kernel isa PolynomialKernel
        title = string("order ", kernel.degree)
    else
        title = string(nameof(typeof(kernel)))
    end
    scatter(x_train, y_train; label=nothing)
    return plot!(x_test, y_pred; label=nothing, title=title)
end

plot((kernelized_fit_and_plot(PolynomialKernel(; degree=degree, c=1)) for degree in 1:4)...)

##
# However, we can now also use kernels that would have an infinite-dimensional feature expansion, such as the squared exponential kernel:

kernelized_fit_and_plot(SqExponentialKernel())


#=
def linear(a, b, x):
    return b + a*x
=#

linear(a, b, x) = b + a * x

#=
# a linear demand function is generated for every
# pair of coefficients in vectors a_vec and b_vec
def demand_hypotheses(a_vec, b_vec):
    for a, b in itertools.product(a_vec, b_vec):
        yield {
            'd': functools.partial(linear, a, b),
            'p_opt': -b/(2*a)
        }
=#

function demand_hypothesis(f, a, b)
	f1(x) = f(a, b, x)
	return DataFrame(
		:a => a,
		:b => b,
		:d => f1,
        :d_opt => f1(-b / (2a)),
		:p_opt => -b / (2a)
	)
end

function generate_demand_hypothesis(a_range, b_range)
    h_vec = DataFrame()
    for a in a_range
    	for b in b_range
    		df1 = demand_hypothesis(linear, a, b)
    		push!(h_vec, df1[1, :])
    	end
    end
    h_vec
end