function gausschebyshev(n::Int, kind::Int=1)
    # GAUSS-CHEBYSHEV NODES AND WEIGHTS.

    # Use known explicit formulas. Complexity O(n).
    if kind == 1
        # Gauss-ChebyshevT quadrature, i.e., w(x) = 1/sqrt(1-x^2)
        [cos((2 * k - 1) * π / (2 * n)) for k = n:-1:1], fill(π / n, n)
    elseif kind == 2
        # Gauss-ChebyshevU quadrature, i.e., w(x) = sqrt(1-x^2)
        ([cos(k * π / (n + 1)) for k = n:-1:1],
         [π/(n + 1) * sin(k / (n + 1) * π)^2 for k = n:-1:1])
    elseif kind == 3
        # Gauss-ChebyshevV quadrature, i.e., w(x) = sqrt((1+x)/(1-x))
        ([cos((k - .5) * π / (n + .5)) for k = n:-1:1],
         [2π / (n + .5) * cos((k - .5) * π / (2 * (n + .5)))^2 for k = n:-1:1])
    elseif kind == 4
        # Gauss-ChebyshevW quadrature, i.e., w(x) = sqrt((1-x)/(1+x))
        ([cos(k * π / (n + .5)) for k = n:-1:1],
         [2π / (n + .5) * sin(k * π / (2 * (n + .5)))^2 for k = n:-1:1])
    else
        throw(ArgumentError("Chebyshev kind should be 1, 2, 3, or 4"))
    end
end


"""
Shapey value using Harsanyi dividends
"""
function fₛⁱ(G::Game, i::Int)
    ∑(Δₕ(G, T) / length(T) for T in subsets(G.N) if i ∈ T)
end

"""
Shapey value using permutations
"""
function fₛ(G::SimpleGame)
    N = length(G.N)
    Nσ = factorial(N)

    shapley = zeros(N)

    for n in 1:Nσ
        σ = nthperm(G.N, n)

        for (i, s) in enumerate(σ)
            shapley[s] += G.v(σ[1:i]) - G.v(σ[1:i - 1])
        end

    end

    return shapley ./ Nσ
end
fₛ(G::GraphGame) = fₛ(graphtoMyerson(G))

"""
Myerson solution, it is equivalent to the Shapley value of the Myerson game.
"""
μ(G::SimpleGame, L::Array{NTuple{2,Int}}) = fₛ(GraphGame(G.N, G.v, L))


""" Computes divided differences of tabulated data n points x and
    for each of the n points, up to m derivatives  dy
"""
function DividedDifference2(x,dy)  # generate array of divided differences
    n, m = size(dy)                # n data points, m derivatives (0 to m-1)
    dd = zeros(typeof(x[1]),n*m,n*m)               # matrix of divided differences
    z = zeros(typeof(x[1]),n*m)
    k=1;
    for i = 1:n                    # n data points
        for j = 1:m                # m derivatives (0 to m-1) at each point
            k = (i-1)*m + j        # row index
            z[k] = x[i]
            dd[k,1] = dy[i,1]      # 0th divided difference in first column
            for l = 2:k            # column index for the remaining columns
                if dd[k,l-1] == dd[k-1,l-1]  # left and top-left neighbors are repeated
                    if l<=m
                        dd[k,l] = dy[i,l] / factorial(l-1)
                    else dd[k,l] = big(0.0)
                    end
                else
                    dd[k,l] = (dd[k,l-1] - dd[k-1,l-1]) / (z[k]-z[k-l+1])
                end
            end
        end
    end
    return dd
end


using Distributions
using Turing
using Test
using Random

Random.seed!(125)

D = [1.0 1.0 4.0 4.0]

@model MoGtest(D) = begin
  mu1 ~ Normal(1, 1)
  mu2 ~ Normal(4, 1)
  z1 ~ Categorical(2)
  if z1 == 1
    D[1] ~ Normal(mu1, 1)
  else
    D[1] ~ Normal(mu2, 1)
  end
  z2 ~ Categorical(2)
  if z2 == 1
    D[2] ~ Normal(mu1, 1)
  else
    D[2] ~ Normal(mu2, 1)
  end
  z3 ~ Categorical(2)
  if z3 == 1
    D[3] ~ Normal(mu1, 1)
  else
    D[3] ~ Normal(mu2, 1)
  end
  z4 ~ Categorical(2)
  if z4 == 1
    D[4] ~ Normal(mu1, 1)
  else
    D[4] ~ Normal(mu2, 1)
  end
  z1, z2, z3, z4, mu1, mu2
end

pmmh = PMMH(500, SMC(10, :z1, :z2, :z3, :z4), MH(1, :mu1, :mu2))
chain = sample(MoGtest(D), pmmh)

@test mean(chain[:z1]) ≈ 1.0 atol=0.1
@test mean(chain[:z2]) ≈ 1.0 atol=0.1
@test mean(chain[:z3]) ≈ 2.0 atol=0.1
@test mean(chain[:z4]) ≈ 2.0 atol=0.1
@test mean(chain[:mu1]) ≈ 1.0 atol=0.1
@test mean(chain[:mu2]) ≈ 4.0 atol=0.1


"""
OLSFn

Run the (first-stage) time-series regressions for the asset pricing model to obtain the alpha and betas for each test assets

INPUT
`Y': T x N matrix of dependent varaible
`X': T x (K + 1) matrix with K explanatory variables and one intercept term
`cov_method': method for calculating the variance-covariance matrix of the error terms:
            - time iid: error terms are iid across time but cross-sectionally correlated
            - iid: error terms are iiid both across time and assets
            - NW: Newey-West covariance matrix with m lags
            - HH: Hansen-Hodrick covariane matrix with m lags
`m': number of lags for the Newey-West or Hansen-Hodrick covariance matrix

OUTPUT
`alpha': N x 1 vector of alphas of the test assets
`beta': N x K matrix of the betas of the test assets. There are N test assets and K factors
`resid': T x N matrix of the pricing errors
`cov_resid': N x N matrix of covariance matrix of the pricing errors
"""

function OLSFn(Y, X, cov_method = "time iid", m = 0)
    beta = (X' * X) \ (X' * Y); # (K+1) x N
    resid = Y - X * beta;
    alpha = beta[1, :]; # alphas N x 1
    beta = beta[2:end, :]'; # betas N x K

    # Method of estimating the covariance matrix of residuals for OLS
    if cov_method == "time iid" # iid across time but cross-sectionally correlated
        cov_resid = cov(resid);
    elseif cov_method == "iid" # iid
        cov_resid = cov(resid);
        cov_resid = diagm(diag(cov_resid));
    elseif cov_method == "NW" # Newey-West
        cov_resid = cov(resid);
        for i = 1:m
            w = 1 - i / (m + 1);
            cov_resid = cov_resid + 2 * w * (1 / (T - 1 - i)) * resid[1:end - i, :]' * resid[1 + i:end, :];
        end
    elseif cov_method == "HH" # Hansen-Hodrick
        cov_resid = cov(resid)
        for i = 1:m
            cov_resid = cov_resid + (2 / (T - 1 - i)) * resid[1:end - i, :]' * resid[1 + i:end, :];
        end
    end

    return (alpha, beta, resid, cov_resid)
end


function aliquotclassifier{T<:Integer}(n::T)
    a = T[n]
    b = divisorsum(a[end])
    len = 1
    while len < 17 && !(b in a) && 0 < b && b < 2^47+1
        push!(a, b)
        b = divisorsum(a[end])
        len += 1
    end
    if b in a
        1 < len || return ("Perfect", a)
        if b == a[1]
            2 < len || return ("Amicable", a)
            return ("Sociable", a)
        elseif b == a[end]
            return ("Aspiring", a)
        else
            return ("Cyclic", push!(a, b))
        end
    end
    push!(a, b)
    b != 0 || return ("Terminating", a)
    return ("Non-terminating", a)
end


using LinearAlgebra, Statistics, StatsBase, Plots; pyplot()

n, N = 5, 10^6
P = diagm(-1 => fill(1/3,n-1),
           0 => fill(1/3,n),
           1 => fill(1/3,n-1))
P[1,n], P[n,1] = 1/3, 1/3

A = UpperTriangular(ones(n,n))
C = P*A

function f1(x,u)
    for xNew in 1:n
        if u <= C[x+1,xNew]
            return xNew-1
        end
    end
end

f2(x,xi) = mod(x + xi , n)

function countTau(f,rnd)
    t = 0
    visits = fill(false,n)
    state = 0
    while sum(visits) < n
        state = f(state,rnd())
        visits[state+1] |= true
        t += 1
    end
    return t-1
end

data1 = [countTau(f1,rand) for _ in 1:N]
data2 = [countTau(f2,()->rand([-1,0,1]) ) for _ in 1:N]
est1, est2 = mean(data1), mean(data2)
c1, c2 = counts(data1)/N,counts(data2)/N
println("Estimated mean value of tau using f1: ",est1)
println("Estimated mean value of tau using f2: ",est2)
println("\nThe matrix P:", P)
scatter(4:33,c1[1:30],
	c=:blue, ms=5, msw=0,
	label="Transition probability matrix")
scatter!(4:33,c2[1:30],
	c=:red, ms=5, msw=0, shape=:cross,
	label="Stochastic recursive formula", xlabel="Time", ylabel="Probability")


export iCycle, iPath, iGrid, iKnight, iCube, iShift

"""
`iCycle(n::Int, simple::Bool=true)` creates an implicit graph that is
a cycle graph with `n` vertices `1` through `n`. When `simple`
is `true`, the graph is undirected; when `false` it's a directed
cycle `1 → 2 → 3 → ⋯ → n → 1`.
"""
function iCycle(n::Int, simple::Bool = true)::ImplicitGraph{Int}
    if n < 3
        error("Number of vertices must be at least 3")
    end

    function N(v::Int)
        a = mod1(v + 1, n)
        b = mod1(v - 1, n)
        if simple
            return [b, a]
        end
        return [a]
    end

    has_vertex(v::Int) = 1 <= v <= n

    return ImplicitGraph{Int}(has_vertex, N)
end

"""
`iPath(simple::Bool=true)` creates an implicit graph that is
a path graph on the integers. If `simple` is `true` vertex `v` is
has an edge to `v-1` and `v+1`; otherwise, `v` has an edge only to
`v+1`.
"""
function iPath(simple::Bool = true)::ImplicitGraph{Int}
    yes(v::Int)::Bool = true

    function N(v::Int)::Vector{Int}
        if simple
            return [v - 1, v + 1]
        else
            return [v + 1]
        end
    end
    return ImplicitGraph{Int}(yes, N)
end

"""
`iGrid()` returns an infinite two-dimensional grid graph. Vertices are
of type `Tuple{Int,Int}`.
"""
function iGrid()::ImplicitGraph{Tuple{Int,Int}}
    yes(v::Tuple{Int,Int})::Bool = true
    function N(v::Tuple{Int,Int})::Vector{Tuple{Int,Int}}
        a, b = v
        return [(a, b - 1), (a, b + 1), (a - 1, b), (a + 1, b)]
    end
    return ImplicitGraph{Tuple{Int,Int}}(yes, N)
end

"""
`iKnight()` returns the Knight's move graph on an infinite
chessboard.
"""
function iKnight()::ImplicitGraph{Tuple{Int,Int}}
    yes(v::Tuple{Int,Int})::Bool = true
    function N(v::Tuple{Int,Int})::Vector{Tuple{Int,Int}}
        a, b = v
        neigh = [
            (a + 1, b + 2),
            (a + 1, b - 2),
            (a + 2, b + 1),
            (a + 2, b - 1),
            (a - 1, b + 2),
            (a - 1, b - 2),
            (a - 2, b + 1),
            (a - 2, b - 1),
        ]
        return neigh
    end
    return ImplicitGraph{Tuple{Int,Int}}(yes, N)
end

"""
`iCube(d::Int)` creates an (implict) `d`-dimensional cube graph.
"""
function iCube(d::Int)::ImplicitGraph{String}
    if d < 1
        error("Dimension must be positive")
    end

    function dvec_check(s::String)::Bool
        if length(s) != d
            return false
        end
        for i = 1:d
            if s[i] ∉ "01"
                return false
            end
        end
        return true
    end

    function N(v::String)
        result = Vector{String}(undef, d)
        for i = 1:d
            head = v[1:i-1]
            c = v[i] == '0' ? "1" : "0"
            tail = v[i+1:end]
            result[i] = head * c * tail
        end
        return result
    end

    return ImplicitGraph{String}(dvec_check, N)
end

using IterTools
"""
`iShift(alphabet,n)` creates an implicit shift digraph whose vertices
are `n`-tuples of elements of `alphabet`.
"""
function iShift(alphabet, n::Int)::ImplicitGraph
    elts = collect(distinct(alphabet))
    T = eltype(elts)

    function has_vertex(v)::Bool
        for i = 1:n
            if v[i] ∉ elts
                return false
            end
        end
        return true
    end

    function N(v)
        base = v[2:end]
        result = [(base..., j) for j in elts]
        return result
    end

    return ImplicitGraph{NTuple{n,T}}(has_vertex, N)
end


### wiener-time-series.jl
#
# Copyright (C) 2016 Mosè Giordano.
#
# Maintainer: Mosè Giordano <mose AT gnu DOT org>
# Keywords: deconvolution, signal processing
#
# This file is a part of Deconvolution.jl.
#
# License is MIT "Expat".
#
### Commentary:
#
# This file produces an example image used in Deconvolution.jl documentation.
#
### Code:

using Random

Random.seed!(7)

using LombScargle, Deconvolution, Plots, Statistics, FFTW
t = range(0, stop=10, length=1000) # observation times
x = sinpi.(t) .* cos.(5t) - 1.5cospi.(t) .* sin.(2t) # the original signal

# Gaussian blurring kernel
kernel = exp.( - 10 .* (t .- 5).^2)
kernel ./= sum(kernel) # normalize kernel to sum of 1
kernel = ifftshift(kernel) # move center to index pos 1

n = rand(length(x)) # noise to be added
noise = 3 .* (n .- mean(n))
y = x + noise # observed noisy signal
# blurred and noise signal
y_blurred = real(ifft(fft(kernel) .* fft(x))) + noise

# Lomb-Scargle periodogram
p = lombscargle(t, y, maximum_frequency=2, samples_per_peak=10)
plot(freqpower(p)...)

m1 = LombScargle.model(t, y, findmaxfreq(p, [0, 0.5])[1]) # first model
m2 = LombScargle.model(t, y, findmaxfreq(p, [0.5, 1])[1]) # second model
m3 = LombScargle.model(t, y, findmaxfreq(p, [1, 1.5])[1]) # third model

signal = m1 + m2 + m3
polished = wiener(y, signal, noise)
deblurred = wiener(y_blurred, signal, noise, kernel)

# Plots
plot(t, x, size=(900, 600), label="Original signal", linewidth=2)
plot!(t, y, label="Observed signal")
savefig("time-series-observed.png")

plot(t, x, size=(900, 600), label="Original signal", linewidth=2)
plot!(t, y_blurred, label="Blurred signal")
savefig("time-series-observed-blurred.png")

plot(t, x, size=(900, 600), label="Original signal", linewidth=2)
plot!(t, polished, label="Recovered with Wiener")
plot!(t, signal, label="Lomb–Scargle model")
savefig("time-series-recovered.png")

plot(t, x, size=(900, 600), label="Original signal", linewidth=2)
plot!(t, deblurred, label="Deblurred with Wiener")
plot!(t, signal, label="Lomb–Scargle model")
savefig("time-series-deblurred.png")


#== # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# Description
#
#   Functions to convert anomalies related to the orbit.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# References
#
#   [1] Vallado, D. A (2013). Fundamentals of Astrodynamics and Applications.
#       Microcosm Press, Hawthorn, CA, USA.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # ==#

export M_to_E, M_to_f
export E_to_f, E_to_M
export f_to_E, f_to_M

################################################################################
#                              From Mean Anomaly
################################################################################

"""
    M_to_E(e::T1, M::T2; kwargs...) where {T1, T2}

Compute the eccentric anomaly (0,2π) \\[rad] given the eccentricity `e` and the
mean anomaly `M` [rad].

This function uses the Newton-Raphson algorithm to solve the Kepler's equation.

# Keywords

- `tol::Union{Nothing, Number}`: Tolerance to accept the solution from
    Newton-Raphson algorithm. If `tol` is `nothing`, then it will be
    `eps(T)`, where `T` is a floating-point type obtained from the promotion of
    `T1` and `T2`. (**Default** = `nothing`)
- `max_iterations::Number`: Maximum number of iterations allowed for the
    Newton-Raphson algorithm. If it is lower than 1, then it is set to 10.
    (**Default** = 10)
"""
@inline function M_to_E(
    e::T1,
    M::T2;
    max_iterations::Integer = 10,
    tol::Union{Nothing, Number} = nothing
) where {T1, T2}
    T = float(promote_type(T1, T2))

    # Compute the eccentric anomaly using the Newton-Raphson method.
    # ==============================================================

    # Make sure that M is in the interval [0,2π].
    M = mod(M, T(2π))

    # Initial guess.
    #
    # See [1, p. 75].
    E = (M > π) ? M - e : M + e

    sin_E, cos_E = sincos(E)

    # Check the tolerance.
    if tol === nothing
        δ = eps(T)
    else
        δ = T(tol)
    end

    # Check the maximum number of iterations.
    if max_iterations < 1
        max_iterations = 10
    end

    # Newton-Raphson iterations.
    for i in 1:max_iterations
        abs(E - e*sin_E - M) ≤ δ && break
        E = E - (E - e * sin_E - M) / (1 - e * cos_E)
        sin_E, cos_E = sincos(E)
    end

    # Return the eccentric anomaly in the interval [0, 2π].
    return mod(E, T(2π))
end

"""
    M_to_f(e::T1, M::T2; kwargs...) where {T1, T2}

Compute the true anomaly (0,2π) \\[rad] given the eccentricity `e` and the mean
anomaly `M` [rad].

This function uses the Newton-Raphson algorithm to solve the Kepler's equation.

# Keywords

- `tol::Union{Nothing, Number}`: Tolerance to accept the solution from
    Newton-Raphson algorithm. If `tol` is `nothing`, then it will be
    `eps(T)`, where `T` is a floating-point type obtained from the promotion of
    `T1` and `T2`. (**Default** = `nothing`)
- `max_iterations::Number`: Maximum number of iterations allowed for the
    Newton-Raphson algorithm. If it is lower than 1, then it is set to 10.
    (**Default** = 10)
"""
@inline function M_to_f(e::T1, M::T2; kwargs...) where {T1, T2}
    # Compute the eccentric anomaly.
    E = M_to_E(e, M; kwargs...)

    # Compute the true anomaly in the interval [0,2π].
    return E_to_f(e, E)
end

################################################################################
#                            From Eccentric Anomaly
################################################################################

"""
    E_to_f(e::T1, E::T2) where {T1, T2}

Compute the true anomaly (0,2π) \\[rad] given the eccentricity `e` and the
eccentric anomaly `E` [rad].
"""
@inline function E_to_f(e::T1, E::T2) where {T1, T2}
    T = float(promote_type(T1, T2))

    sin_Eo2, cos_Eo2 = sincos(E / 2)

    # Compute the true anomaly in the interval [0, 2*π].
    return mod(2atan(sqrt(1 + e) * sin_Eo2, sqrt(1 - e) * cos_Eo2), T(2π))
end

"""
    E_to_M(e::T1, E::T2) where {T1, T2}

Compute the mean anomaly (0,2π) \\[rad] given the eccentricity `e` and the
eccentric anomaly `E` [rad].
"""
@inline function E_to_M(e::T1, E::T2) where {T1, T2}
    T = float(promote_type(T1, T2))
    return mod(E - e * sin(E), T(2π))
end

################################################################################
#                              From True Anomaly
################################################################################

"""
    f_to_E(e::T1, f::T2) where {T1, T2}

Compute the eccentric anomaly (0,2π) \\[rad] given the eccentricity `e` and
the true anomaly `f` [rad].
"""
@inline function f_to_E(e::T1, f::T2) where {T1, T2}
    T = float(promote_type(T1, T2))
    sin_fo2, cos_fo2 = sincos(f / 2)
    return mod(2atan(sqrt(1 - e) * sin_fo2, sqrt(1 + e) * cos_fo2), T(2π) )
end

"""
    f_to_M(e::T1, f::T2) where {T1, T2}

Compute the mean anomaly (0,2π) \\[rad] given the eccentricity `e` and the
true anomaly `f` [rad].
"""
@inline function f_to_M(e::T1, f::T2) where {T1, T2}
    # Compute the eccentric anomaly.
    E = f_to_E(e, f)

    # Compute the true anomaly in the interval [0, 2π].
    return E_to_M(e, E)
end


mutable struct Model{T <: Real}
    y::Vector{Int}
    X::Matrix{T}
    threshold::T
    num_obs::Int
    beta_hat::Vector{T}
    pi_hat::Vector{T}
    y_hat::Vector{Int}
    dof_log::Int
    dof_resid::Int
    dof_total::Int
    llk::T
    aic::T
    bic::T
    # r2::T
    dev_residuals::Vector{T}
    dev_residuals_var::Float64
    sigma::Matrix{T}
    std_error::Vector{T}
    # z_value::Vector{T}
    # z_test_p_value::Vector{T}
end


using Plots

function circle(x)
    θ = x[1];
    r = 1/2 + 1/2 * (2*x[2]/(1+x[2]^2));
    return [1-r*cos(θ), 1-r*sin(θ)];
end


# One dimensional wave equation
using ConservationLawsDiffEq
using OrdinaryDiffEq
const CL = ConservationLawsDiffEq

const CFL = 0.45
const Tend = 1.0
const cc = 1.0

Jf(u::AbstractVector) = [0.0 cc;cc 0.0]
f(u::AbstractVector) = [0.0 cc;cc 0.0]*u
f0(x) = [sin(4*π*x),0.0]

Nflux(ϕl::AbstractVector, ϕr::AbstractVector) = 0.5*(f(ϕl)+f(ϕr))
exact_sol(x, t::Float64) = [0.5*(sin(4*π*(-t+x))+sin(4*π*(t+x))),
0.5*(sin(4*π*(-t+x))-sin(4*π*(t+x)))]

# Now discretizate the domain
mesh = Uniform1DFVMesh(50, [-1.0, 1.0])

function get_problem(alg, mesh)
  #Compute discrete initial data
  u0 = getInitialState(mesh,f0,use_threads = true)

  # Now get a explicit semidiscretization (discrete in space) du_h(t)/dt = f_h(u_h(t))
  f_h = getSemiDiscretization(f,alg,mesh,[Periodic()]; Df = Jf, use_threads = false,numvars = 2)

  #Setup ODE problem for a time interval = [0.0,1.0]
  ode_prob = ODEProblem(f_h,u0,(0.0,Tend))

  #Setup callback in order to fix CFL constant value
  cb = getCFLCallback(f_h, CFL)

  #Estimate an initial dt
  dt = update_dt!(u0, f_h, CFL)
  return ode_prob, cb, dt
end

ode_prob, cb, dt = get_problem(FVSKTScheme(), mesh)
sol = solve(ode_prob,SSPRK22(); dt = dt, callback = cb)

ode_prob, cb, dt = get_problem(FVTecnoScheme(Nflux;order=3), mesh)
sol2 = solve(ode_prob,SSPRK22(); dt = dt, callback = cb)

ode_prob, cb, dt = get_problem(FVCompWENOScheme(), mesh)
sol3 = solve(ode_prob,SSPRK22(); dt = dt, callback = cb)

ode_prob, cb, dt = get_problem(FVCompMWENOScheme(), mesh)
sol4 = solve(ode_prob,SSPRK22(); dt = dt, callback = cb)

ode_prob, cb, dt = get_problem(FVSpecMWENOScheme(), mesh)
sol5 = solve(ode_prob,SSPRK22(); dt = dt, callback = cb)

#Estimate errors
u1_h = fv_solution(sol, mesh; vars = 2)
u2_h = fv_solution(sol2, mesh; vars = 2)
u3_h = fv_solution(sol3, mesh; vars = 2)
u4_h = fv_solution(sol4, mesh; vars = 2)
u5_h = fv_solution(sol5, mesh; vars = 2)
get_L1_error(exact_sol,u1_h)
get_L1_error(exact_sol,u2_h)
get_L1_error(exact_sol,u3_h)
get_L1_error(exact_sol,u4_h)
get_L1_error(exact_sol,u5_h)

#Plot
using Plots
plot(u1_h, vars=1, lab="KT y",line = (:dot,2))
plot!(u2_h, vars=1,lab="Tecno y",line=(:dot,3))
plot!(u3_h, vars=1,lab="Comp WENO y",line=(:dot,3))
plot!(u4_h, vars=1,lab="Comp MWENO y",line=(:dot,3))
plot!(u5_h, vars=1,lab="Spec MWENO y",line=(:dot,3))
plot!(CL.cell_centers(mesh), [exact_sol(x,Tend)[1] for x in CL.cell_centers(mesh)],lab="Ref y")


"""
Features extraction from returns proposed by Betina and al.
"""
#KLT stands for the long-term period
#KST stands for the short-term period
#ϱ is the number of standard deviations above the means
function feature_indexes_sgns(asset_rtn_sgn, kst, klt, kmom, ϱ=2)
    #current time
    t = size(asset_rtn_sgn, 1)
    # sig_1 and sig_2 are
    # short-term and long-term simple moving averages
    sig_1 = (1 / kst) * sum(asset_rtn_sgn[t - d] for d in 0:(kst - 1))
    sig_2 = (1 / klt) * sum(asset_rtn_sgn[t - d] for d in 0:(klt - 1))

    # sig_3 and sig_4 are exponential moving average (short and long)
    sig_3 =
        (1.0 / sum(exp(-d / kst) for d in 1:kst)) *
        sum(exp(-d / kst) * asset_rtn_sgn[t - d] for d in 0:(kst - 1))
    sig_4 =
        (1.0 / sum(exp(-d / klt) for d in 1:klt)) *
        sum(exp(-d / klt) * asset_rtn_sgn[t - d] for d in 0:(klt - 1))
    # RSI
    R_pos = 0
    R_neg = 0
    num_pos = 0
    num_neg = 0
    for i in (t - kmom + 1):t
        if (asset_rtn_sgn[i] > 0)
            R_pos = R_pos + asset_rtn_sgn[i]
            num_pos = num_pos + 1
        else
            R_neg = R_neg + asset_rtn_sgn[i]
            num_neg = num_neg + 1
        end
    end
    med_pos = R_pos / num_pos
    if num_neg == 0
        med_neg = 1
    else
        med_neg = abs(R_neg / num_neg)
    end
    sig_5 = med_pos / med_neg
    # estimated shor-t and long-t standard deviations
    σ_st_6 = sqrt((1.0 / kst) * sum((asset_rtn_sgn[t - d] - sig_1)^2 for d in 0:(kst - 1)))
    σ_lt_7 = sqrt((1.0 / klt) * sum((asset_rtn_sgn[t - d] - sig_2)^2 for d in 0:(klt - 1)))
    # sig_6 and sig_7 are

    if sig_1 <= 0
        sig_6 = sig_1 - ϱ * σ_st_6
    else
        sig_6 = sig_1 + ϱ * σ_st_6
    end
    if sig_2 <= 0
        sig_7 = sig_2 - ϱ * σ_lt_7
    else
        sig_7 = sig_2 + ϱ * σ_lt_7
    end
    σ_st_8 = sqrt((1.0 / kst) * sum((asset_rtn_sgn[t - d] - sig_3)^2 for d in 0:(kst - 1)))
    σ_lt_9 = sqrt((1.0 / klt) * sum((asset_rtn_sgn[t - d] - sig_4)^2 for d in 0:(klt - 1)))
    if sig_3 <= 0
        sig_8 = sig_3 - ϱ * σ_st_8
    else
        sig_8 = sig_3 + ϱ * σ_st_8
    end
    if sig_4 <= 0
        sig_9 = sig_4 - ϱ * σ_lt_9
    else
        sig_9 = sig_4 + ϱ * σ_lt_9
    end

    return [sig_1; sig_2; sig_3; sig_4; sig_5; sig_6; sig_7; sig_8; sig_9]
end

"""
Mixed signals predictor proposed by Betina and al. Univariate
"""
function mixed_signals_predict_return(asset_rtn_sgn, num_t, kst_a, klt_a, kmom, ϱ=2)
    numD = size(asset_rtn_sgn, 1)
    num_train = min(numD - 1, num_t)

    ## optimize features weights (least squares based on past returns)
    # signal features extraction train (from klt+1 to klt+num_train)
    #kst = max(min(numD-num_train-10,kst_a),0)
    #klt = max(min(numD-num_train-10,klt_a),0)
    sig_features = zeros(num_train, 9) #each line is the i-element's featurs
    for i in 1:num_train
        #kst = max(min(numD-i-2,kst_a),0)
        klt = max(min(numD - i - 2, klt_a), 0)
        kst = Int64(floor(klt / 2.0))
        sig_features[num_train + 1 - i, :] = feature_indexes_sgns(
            asset_rtn_sgn[1:(numD - i)], kst, klt, kmom
        )
    end
    # optimize weights (least squares - pseudo inverse)
    try
        weights_opt = \(
            sig_features'sig_features,
            sig_features'asset_rtn_sgn[(numD - num_train + 1):numD],
        )
    catch
        weights_opt = pinv(sig_features, 1E-07) * asset_rtn_sgn[(numD - num_train + 1):numD]
    end
    #  normalize features weights
    weights_opt = weights_opt ./ sum(weights_opt)
    ## extract features current return
    kst = min(numD, kst_a)
    klt = min(numD, klt_a)
    features = feature_indexes_sgns(asset_rtn_sgn[1:numD], kst, klt, kmom)
    ## extrapolate return

    return (features'weights_opt)[1]
end


@testset "curve_length" begin
    f(x) = sin(x)

    sol = 3.8201977890277120179047620821714432919099676146472747210804966567
    @test isapprox(curve_length(f,0,pi), sol; atol = 1e-10)
    @test !isapprox(curve_length(f,0,5*pi;rtol=1e-4),5*sol,atol=1e-7)
end

@testset "curve_fraction" begin
    f(x) = sin(x)

    @test isapprox(curve_fraction(f,0,pi,0.5), pi/2; atol = 1e-10)
    @test !isapprox(curve_fraction(f,0,pi,0.5;rtol=1e-2), pi/2; atol = 1e-9)
end


using Test

include("../src/weightedgraph.jl")
include("../test/utils.jl")
using .WeightedGraphAlgo
using .WeightedGraphAlgo: WeightedUndirectedEdge

const tiny_ewg_p614 = (
    nv = 8,
    ne = 16,
    edges = [
        (5, 6, 0.35),
        (5, 8, 0.37),
        (6, 8, 0.28),
        (1, 8, 0.16),
        (2, 6, 0.32),
        (1, 5, 0.38),
        (3, 4, 0.17),
        (2, 8, 0.19),
        (1, 3, 0.26),
        (2, 3, 0.36),
        (2, 4, 0.29),
        (3, 8, 0.34),
        (7, 3, 0.40),
        (4, 7, 0.52),
        (7, 1, 0.58),
        (7, 4, 0.93)
    ]
)

@testset "Weighted undirected graph" begin

    @testset "Lazy Prim algorithm" begin
        g = WeightedUndirectedGraph(tiny_ewg_p614)
        mst = MST(g, LazyPrimMST)
        edgeset = Set(mst.result.mstedges)
        @test sum(mst.result.mstvertices) == g.numvertices
        @test sumweights(mst) == 1.81
        @test WeightedUndirectedEdge((1, 8, 0.16)) in edgeset
        @test WeightedUndirectedEdge((6, 8, 0.28)) in edgeset
    end

    @testset "Eager Prim algorithm" begin
        g = WeightedUndirectedGraph(tiny_ewg_p614)
        mst = MST(g, EagerPrimMST)
        edgeset = Set(mst.result.mstedges)
        @test sum(mst.result.mstvertices) == g.numvertices
        @test sumweights(mst) ≈ 1.81
        @test WeightedUndirectedEdge((1, 8, 0.16)) in edgeset
        @test WeightedUndirectedEdge((6, 8, 0.28)) in edgeset
    end

    @testset "Kruskal's algorithm" begin
        g = WeightedUndirectedGraph(tiny_ewg_p614)
        mst = MST(g, KruskalMST)
        edgeset = Set(mst.result.mstedges)
        @test sumweights(mst) == 1.81
        @test WeightedUndirectedEdge((1, 8, 0.16)) in edgeset
        @test WeightedUndirectedEdge((6, 8, 0.28)) in edgeset
    end

end

@testset "Weighted directed graph" begin
    CURRENT_PATH = cd(pwd, ".")
    tinyewd = loadewg(joinpath(CURRENT_PATH, "data/tinyEWD.txt"))

    @testset "Dijkstra shortest path" begin
        g = WeightedDirectedGraph(tinyewd)
        sp = SP(g, 1, DijkstraSP)
        @test distanceto(sp, 2) ≈ 1.05
        @test distanceto(sp, 3) ≈ 0.26
        @test distanceto(sp, 4) ≈ 0.99
        @test distanceto(sp, 5) ≈ 0.38
        @test distanceto(sp, 6) ≈ 0.73
        @test distanceto(sp, 7) ≈ 1.51
        @test distanceto(sp, 8) ≈ 0.60
    end

end


using LinearAlgebra

using GLMakie; GLMakie.activate!()
using AbstractPlotting

using Delaunay
using DEC2D

# convert vector of vectors into a matrix
vv2m(vv) = reduce(vcat, transpose.(vv))

# Simples case
begin
    vertices = vcat([[0, 0]], [[cos(2π * i / 6), sin(2π * i / 6)] for i = 1:6])
    edges = [
        [1, 2],
        [1, 3],
        [1, 4],
        [1, 5],
        [1, 6],
        [1, 7],
        [2, 3],
        [3, 4],
        [4, 5],
        [5, 6],
        [6, 7],
        [7, 2],
    ]
    triangles = [[1, 2, 3], [1, 3, 4], [1, 4, 5], [1, 5, 6], [1, 6, 7], [1, 7, 2]]

    # return the mesh and the indices of the interior vertices
    vint, mesh = DECMesh(vertices, edges, triangles)
    vext = setdiff(1:length(vertices), vint)

    # d₀ = boundary1(mesh)
    # d₁ = boundary2d(mesh)
    # *₁ = hodgestar1(mesh)
    # *₂ = hodgestar2(mesh)
    # Δ = *₂ * d₁ * *₁ * d₀
    Δ = laplacian(mesh) # calculate the laplacian operator

    b = [-1 - 10 * sum(Δ[i, vext]) for i = 1:length(vint)]

    # numerical solution
    zf = Δ[vint, vint] \ b

    x = vv2m(vertices)[:, 1]
    y = vv2m(vertices)[:, 2]
    zn = similar(vertices, Float64)
    zn[vint] = zf
    zn[vext] .= 10

    # exact solution
    f(x, y) = (1 - x^2 - y^2) / 4 + 10
    z = [f(v...) for v in vertices]

    println(
        "La norma de la diferencia entre la solución numérica y la exacta es ",
        norm(zn - z),
    )

    #   ____  _       _   _   _
    #  |  _ \| | ___ | |_| |_(_)_ __   __ _
    #  | |_) | |/ _ \| __| __| | '_ \ / _` |
    #  |  __/| | (_) | |_| |_| | | | | (_| |
    #  |_|   |_|\___/ \__|\__|_|_| |_|\__, |
    #                                 |___/

    # plot numerical solution
    scene = Scene()
    mesh!(
        scene,
        hcat(x, y, zn),
        vv2m(mesh.tri.triangles),
        color = zn,
        scale_plot = false,
    )
    xlims!(scene, -1, 1)
    ylims!(scene, -1, 1)
    zlims!(scene, 10, 10.25)
    wireframe!(scene[end][1], color = (:black, 0.6), linewidth = 2)
    display(scene)
    sleep(5.0)
end

# Caso medio
begin
    # we get the vertices of a more fine mesh
    vertices = [
        [cos(2π * i / (50 * r + 1)), sin(2π * i / (50 * r + 1))] .* r for r = 0:0.1:1
        for i = 1:round(Int, 50 * r + 1, RoundDown)
    ]

    # get thetriangulation using Delaunay algorithm
    meshd = delaunay(vv2m(vertices))
    triangles = [meshd.simplices[i, :] for i = 1:size(meshd.simplices, 1)]
    edges = reduce(
        union,
        [[sort([t[1], t[2]]), sort([t[2], t[3]]), sort([t[3], t[1]])] for t in triangles],
    )

    vint, mesh = DECMesh(vertices, edges, triangles)
    vext = setdiff(1:length(vertices), vint)

    Δ = laplacian(mesh)

    b = [-1 - 10 * sum(Δ[i, vext]) for i = 1:length(vint)]
    zf = Δ[vint, vint] \ b

    x = vv2m(vertices)[:, 1]
    y = vv2m(vertices)[:, 2]
    zn = similar(vertices, Float64)
    zn[vint] = zf
    zn[vext] .= 10

    f(x, y) = (1 - x^2 - y^2) / 4 + 10
    z = [f(v...) for v in vertices]

    println(
        "La norma de la diferencia entre la solución numérica y la exacta es ",
        norm(zn - z),
    )

    scene = Scene()
    mesh!(
        scene,
        hcat(x, y, zn),
        vv2m(mesh.tri.triangles),
        color = zn,
        scale_plot = false,
    )
    xlims!(scene, -1, 1)
    ylims!(scene, -1, 1)
    zlims!(scene, 10, 10.25)
    wireframe!(scene[end][1], color = (:black, 0.6), linewidth = 2)
    display(scene)
    sleep(5.0)
end

# Caso denso
begin
    # we get the vertices of a more fine mesh
    vertices = [
        [cos(2π * i / (200 * r + 1)), sin(2π * i / (200 * r + 1))] .* r for r = 0:0.025:1 for i = 1:round(Int, 200 * r + 1, RoundDown)
    ]

    # get thetriangulation using Delaunay algorithm
    meshd = delaunay(vv2m(vertices))
    triangles = [meshd.simplices[i, :] for i = 1:size(meshd.simplices, 1)]
    edges = reduce(
        union,
        [[sort([t[1], t[2]]), sort([t[2], t[3]]), sort([t[3], t[1]])] for t in triangles],
    )

    vint, mesh = DECMesh(vertices, edges, triangles)
    vext = setdiff(1:length(vertices), vint)

    Δ = laplacian(mesh)

    b = [-1 - 10 * sum(Δ[i, vext]) for i = 1:length(vint)]
    zf = Δ[vint, vint] \ b

    x = vv2m(vertices)[:, 1]
    y = vv2m(vertices)[:, 2]
    zn = similar(vertices, Float64)
    zn[vint] = zf
    zn[vext] .= 10

    f(x, y) = (1 - x^2 - y^2) / 4 + 10
    z = [f(v...) for v in vertices]

    println(
        "La norma de la diferencia entre la solución numérica y la exacta es ",
        norm(zn - z),
    )

    scene = Scene()
    mesh!(
        scene,
        hcat(x, y, zn),
        vv2m(mesh.tri.triangles),
        color = zn,
        scale_plot = false,
    )
    xlims!(scene, -1, 1)
    ylims!(scene, -1, 1)
    zlims!(scene, 10, 10.25)
    wireframe!(scene[end][1], color = (:black, 0.6), linewidth = 2)
    display(scene)
end


@add_problem socp function socp_norm_2_atom(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    x = Variable(2, 1)
    A = [1 2; 2 1; 3 4]
    b = [2; 3; 4]
    p = minimize(norm2(A * x + b))
    if test
        @test vexity(p) == ConvexVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 0.64888 atol=atol rtol=rtol
        @test evaluate(norm2(A * x + b)) ≈ 0.64888 atol=atol rtol=rtol
    end

    x = Variable(2, 1)
    A = [1 2; 2 1; 3 4]
    b = [2; 3; 4]
    lambda = 1
    p = minimize(norm2(A * x + b) + lambda * norm2(x), x >= 1)
    if test
        @test vexity(p) == ConvexVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 14.9049 atol=atol rtol=rtol
        @test evaluate(norm2(A * x + b) + lambda * norm2(x)) ≈ 14.9049 atol=atol rtol=rtol
    end

    x = Variable(2)

    p = minimize(norm2([x[1] + 2x[2] + 2; 2x[1] + x[2] + 3; 3x[1]+4x[2] + 4]) + lambda * norm2(x), x >= 1)
    if test
        @test vexity(p) == ConvexVexity()
    end

    handle_problem!(p)
    if test
        @test p.optval ≈ 14.9049 atol=atol rtol=rtol
        @test evaluate(norm2(A * x + b) + lambda * norm2(x)) ≈ 14.9049 atol=atol rtol=rtol
    end

    x = Variable(2, 1)
    A = [1 2; 2 1; 3 4]
    b = [2; 3; 4]
    lambda = 1
    p = minimize(norm2(A * x + b) + lambda * norm_1(x), x >= 1)
    if test
        @test vexity(p) == ConvexVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 15.4907 atol=atol rtol=rtol
        @test evaluate(norm2(A * x + b) + lambda * norm_1(x)) ≈ 15.4907 atol=atol rtol=rtol
    end
end

@add_problem socp function socp_frobenius_norm_atom(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    m = Variable(4, 5)
    c = [m[3, 3] == 4, m >= 1]
    p = minimize(norm(vec(m), 2), c)
    if test
        @test vexity(p) == ConvexVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ sqrt(35) atol=atol rtol=rtol
        @test evaluate(norm(vec(m), 2)) ≈ sqrt(35) atol=atol rtol=rtol
    end
end

@add_problem socp function socp_quad_over_lin_atom(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    x = Variable(3, 1)
    A = [2 -3 5; -2 9 -3; 5 -8 3]
    b = [-3; 9; 5]
    c = [3 2 4]
    d = -3
    p = minimize(quadoverlin(A*x + b, c*x + d))
    if test
        @test vexity(p) == ConvexVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 17.7831 atol=atol rtol=rtol
        @test (evaluate(quadoverlin(A * x + b, c * x + d)))[1] ≈ 17.7831 atol=atol rtol=rtol
    end
end

@add_problem socp function socp_sum_squares_atom(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    x = Variable(2, 1)
    A = [1 2; 2 1; 3 4]
    b = [2; 3; 4]
    p = minimize(sumsquares(A*x + b))
    if test
        @test vexity(p) == ConvexVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 0.42105 atol=atol rtol=rtol
        @test (evaluate(sumsquares(A * x + b)))[1] ≈ 0.42105 atol=atol rtol=rtol
    end
end

@add_problem socp function socp_square_atom(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    x = Variable(2, 1)
    A = [1 2; 2 1; 3 4]
    b = [2; 3; 4]
    p = minimize(sum(square(A*x + b)))
    if test
        @test vexity(p) == ConvexVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 0.42105 atol=atol rtol=rtol
        @test evaluate(sum(square(A * x + b))) ≈ 0.42105 atol=atol rtol=rtol
    end

    x = Variable(2, 1)
    A = [1 2; 2 1; 3 4]
    b = [2; 3; 4]
    expr = A * x + b
    p = minimize(sum(dot(^)(expr,2))) # elementwise ^
    if test
        @test vexity(p) == ConvexVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 0.42105 atol=atol rtol=rtol
        @test evaluate(sum(broadcast(^, expr, 2))) ≈ 0.42105 atol=atol rtol=rtol
    end

    p = minimize(sum(dot(*)(expr, expr))) # elementwise *
    if test
        @test vexity(p) == ConvexVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 0.42105 atol=atol rtol=rtol
        @test evaluate(sum((dot(*))(expr, expr))) ≈ 0.42105 atol=atol rtol=rtol
    end
end

@add_problem socp function socp_inv_pos_atom(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    x = Variable(4)
    p = minimize(sum(invpos(x)), invpos(x) < 2, x > 1, x == 2, 2 == x)
    if test
        @test vexity(p) == ConvexVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 2 atol=atol rtol=rtol
        @test evaluate(sum(invpos(x))) ≈ 2 atol=atol rtol=rtol
    end

    x = Variable(3)
    p = minimize(sum(dot(/)([3,6,9], x)), x<=3)
    handle_problem!(p)
    if test
        @test x.value ≈ fill(3.0, (3, 1)) atol=atol rtol=rtol
        @test p.optval ≈ 6 atol=atol rtol=rtol
        @test evaluate(sum((dot(/))([3, 6, 9], x))) ≈ 6 atol=atol rtol=rtol
    end

    x = Variable()
    p = minimize(sum([3,6,9]/x), x<=3)
    handle_problem!(p)
    if test
        @test x.value ≈ 3 atol=atol rtol=rtol
        @test p.optval ≈ 6 atol=atol rtol=rtol
        @test evaluate(sum([3, 6, 9] / x)) ≈ 6 atol=atol rtol=rtol
    end
end

@add_problem socp function socp_geo_mean_atom(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    x = Variable(2)
    y = Variable(2)
    p = minimize(geomean(x, y), x >= 1, y >= 2)
    # not DCP compliant
    if test
        @test vexity(p) == ConcaveVexity()
    end
    p = maximize(geomean(x, y), 1 < x, x < 2, y < 2)
    # Just gave it a vector as an objective, not okay
    if test
        @test_throws Exception handle_problem!(p)
    end

    p = maximize(sum(geomean(x, y)), 1 < x, x < 2, y < 2)
    handle_problem!(p)
    if test
        @test p.optval ≈ 4 atol=atol rtol=rtol
        @test evaluate(sum(geomean(x, y))) ≈ 4 atol=atol rtol=rtol
    end
end

@add_problem socp function socp_sqrt_atom(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    x = Variable()
    p = maximize(sqrt(x), 1 >= x)
end

@add_problem socp function socp_quad_form_atom(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    x = Variable(3, 1)
    A = [0.8608 0.3131 0.5458; 0.3131 0.8584 0.5836; 0.5458 0.5836 1.5422]
    p = minimize(quadform(x, A), [x >= 1])
    if test
        @test vexity(p) == ConvexVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 6.1464 atol=atol rtol=rtol
        @test (evaluate(quadform(x, A)))[1] ≈ 6.1464 atol=atol rtol=rtol
    end

    x = Variable(3, 1)
    A = -1.0*[0.8608 0.3131 0.5458; 0.3131 0.8584 0.5836; 0.5458 0.5836 1.5422]
    c = [3 2 4]
    p = maximize(c*x , [quadform(x, A) >= -1])
    if test
        @test vexity(p) == ConvexVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 3.7713 atol=atol rtol=rtol
        @test (evaluate(quadform(x, A)))[1] ≈ -1 atol=atol rtol=rtol
    end
end

@add_problem socp function socp_huber_atom(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    x = Variable(3)
    p = minimize(sum(huber(x, 1)), x >= 2)
    if test
        @test vexity(p) == ConvexVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 9 atol=atol rtol=rtol
        @test evaluate(sum(huber(x, 1))) ≈ 9 atol=atol rtol=rtol
    end
end

@add_problem socp function socp_rational_norm_atom(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    A = [1 2 3; -1 2 3]
    b = A * ones(3)
    x = Variable(3)
    p = minimize(norm(x, 4.5), [A * x == b])
    if test
        @test vexity(p) == ConvexVexity()
    end
    # Solution is approximately x = [1, .93138, 1.04575]
    handle_problem!(p)
    if test
        @test p.optval ≈ 1.2717 atol=atol rtol=rtol
        @test evaluate(norm(x, 4.5)) ≈ 1.2717 atol=atol rtol=rtol
    end
end

@add_problem socp function socp_rational_norm_dual_norm(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    v = [0.463339, 0.0216084, -2.07914, 0.99581, 0.889391]
    x = Variable(5)
    q = 1.379;  # q norm constraint that generates many inequalities
    qs = q / (q - 1);  # Conjugate to q
    p = minimize(x' * v)
    p.constraints += (norm(x, q) <= 1)
    if test
        @test vexity(p) == ConvexVexity()
    end
    handle_problem!(p) # Solution is -norm(v, q / (q - 1))
    if test
        @test p.optval ≈ -2.144087 atol=atol rtol=rtol
        @test sum(evaluate(x' * v)) ≈ -2.144087 atol=atol rtol=rtol
        @test evaluate(norm(x, q)) ≈ 1 atol=atol rtol=rtol
        @test sum(evaluate(x' * v)) ≈ -(sum(abs.(v) .^ qs) ^ (1 / qs)) atol=atol rtol=rtol
    end
end

@add_problem socp function socp_rational_norm_atom_sum(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    A = [-0.719255  -0.229089
        -1.33632   -1.37121
        0.703447  -1.4482]
    b = [-1.82041, -1.67516, -0.866884]
    q = 1.5
    xvar = Variable(2)
    p = minimize(.5 * sumsquares(xvar) + norm(A * xvar - b, q))
    if test
        @test vexity(p) == ConvexVexity()
    end
    handle_problem!(p)

    if test
        # Compute gradient, check it is zero(ish)
        x_opt = xvar.value
        margins = A * x_opt - b
        qs = q / (q - 1);  # Conjugate
        denom = sum(abs.(margins).^q)^(1/qs)
        g = x_opt + A' * (abs.(margins).^(q-1) .* sign.(margins)) / denom
        @test p.optval ≈ 1.7227 atol=atol rtol=rtol
        @test norm(g, 2) ^ 2 ≈ 0 atol=atol rtol=rtol
    end
end

@add_problem socp function socp_norm_consistent_with_Base_for_matrix_variables(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    A = randn(4, 4)
    x = Variable(4, 4)
    x.value = A
    # Matrix norm
    if test
        @test evaluate(opnorm(x)) ≈ opnorm(A) atol=atol rtol=rtol
        @test evaluate(opnorm(x, 1)) ≈ opnorm(A, 1) atol=atol rtol=rtol
        @test evaluate(opnorm(x, 2)) ≈ opnorm(A, 2) atol=atol rtol=rtol
        @test evaluate(opnorm(x, Inf)) ≈ opnorm(A, Inf) atol=atol rtol=rtol
    end
    # Vector norm
    # TODO: Once the deprecation for norm on matrices is removed, remove the `vec` calls
    if test
        @test evaluate(norm(vec(x), 1)) ≈ norm(vec(A), 1) atol=atol rtol=rtol
        @test evaluate(norm(vec(x), 2)) ≈ norm(vec(A), 2) atol=atol rtol=rtol
        @test evaluate(norm(vec(x), 7)) ≈ norm(vec(A), 7) atol=atol rtol=rtol
        @test evaluate(norm(vec(x), Inf)) ≈ norm(vec(A), Inf) atol=atol rtol=rtol
    end
end

@add_problem socp function socp_fix_and_free_addition(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    x = Variable()
    y = Variable()

    p = minimize(x+y, x>=0, y>=0)
    handle_problem!(p)
    if test
        @test p.optval ≈ 0 atol=atol rtol=rtol
    end

    y.value = 4
    fix!(y)
    handle_problem!(p)
    if test
        @test p.optval ≈ 4 atol=atol rtol=rtol
    end

    free!(y)
    handle_problem!(p)
    if test
        @test p.optval ≈ 0 atol=atol rtol=rtol
    end
end

@add_problem socp function socp_fix_multiplication(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    a = [1,2,3,2,1]
    x = Variable(length(a))
    gamma = Variable(Positive())
    fix!(gamma, 0.7)

    p = minimize(norm(x-a) + gamma*norm(x[1:end-1] - x[2:end]))
    handle_problem!(p)
    if test
        o1 = p.optval
        # x should be very close to a
        @test o1 ≈ 0.7 * norm(a[1:end - 1] - a[2:end]) atol=atol rtol=rtol
    end
    # increase regularization
    fix!(gamma, 1.0)
    handle_problem!(p)

    if test
        o2 = p.optval
        # x should be very close to mean(a)
        @test o2 ≈ norm(a .- mean(a)) atol=atol rtol=rtol
    end

    if test
        @test o1 <= o2
    end
end

@add_problem socp function socp_dual_minimal_norm_solutions(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    x = Variable(2)
    A = [1 2; 2 4];
    b = [3, 6];
    p = minimize(norm(x, 1), A*x==b)

    if test
        @test vexity(p) == ConvexVexity()
    end

    handle_problem!(p)
    if test
        @test p.optval ≈ 1.5 atol=atol rtol=rtol
        @test evaluate(x) ≈ [0, 1.5] atol=atol rtol=rtol
        @test evaluate(norm(x, 1)) ≈ p.optval atol=atol rtol=rtol
        @test dot(b, p.constraints[1].dual) ≈ p.optval atol=atol rtol=rtol
    end

    x = Variable(2)
    A = [1 2; 2 4];
    b = [3, 6];
    p = minimize(norm(x, 2), A*x==b)

    test && @test vexity(p) == ConvexVexity()

    handle_problem!(p)

    if test
        @test p.optval ≈ 3/sqrt(5) atol=atol rtol=rtol
        @test evaluate(x) ≈ [3/5, 6/5] atol=atol rtol=rtol
        @test evaluate(norm(x, 2)) ≈ p.optval atol=atol rtol=rtol
        @test dot(b, p.constraints[1].dual) ≈ p.optval atol=atol rtol=rtol
    end

    x = Variable(2)
    A = [1 2; 2 4];
    b = [3, 6];
    p = minimize(norm(x, Inf), A*x==b)

    test && @test vexity(p) == ConvexVexity()

    handle_problem!(p)

    if test
        @test p.optval ≈ 1.0 atol=atol rtol=rtol
        @test evaluate(x) ≈ [1, 1] atol=atol rtol=rtol
        @test evaluate(norm(x, Inf)) ≈ p.optval atol=atol rtol=rtol
        @test dot(b, p.constraints[1].dual) ≈ p.optval atol=atol rtol=rtol
    end
end


using PyCall
@pyimport matplotlib.pyplot as plt

using Gen

#import Gen: Distribution, logpdf

# Example from Section 4 of Reversible jump Markov chain Monte Carlo
# computation and Bayesian model determination

########################
# custom distributions #
########################

# minimum of k draws from uniform_continuous(lower, upper)

# we can sequentially sample the order statistics of a collection of K uniform
# continuous samples on the interval [a, b], by:
# x1 ~ min_uniform_continuous(a, b, K)
# x2 | x1 ~ min_uniform_continuous(x1, b, K-1)
# ..
# xK | x1 .. x_{K-1} ~ min_uniform_continuous(x_{K-1}, b, 1)

struct MinUniformContinuous <: Distribution{Float64} end
const min_uniform_continuous = MinUniformContinuous()

function Gen.logpdf(::MinUniformContinuous, x::Float64, lower::T, upper::U, k::Int) where {T<:Real,U<:Real}
    if x > lower && x < upper
        (k-1) * log(upper - x) + log(k) - k * log(upper - lower)
    else
        -Inf
    end
end

function Gen.random(::MinUniformContinuous, lower::T, upper::U, k::Int) where {T<:Real,U<:Real}
    # inverse CDF method
    p = rand()
    upper - (upper - lower) * (1. - p)^(1. / k)
end

# piecewise homogenous Poisson process

# n intervals - n + 1 bounds
# (b_1, b_2]
# (b_2, b_3]
# ..
# (b_n, b_{n+1}]

function compute_total(bounds, rates)
    num_intervals = length(rates)
    if length(bounds) != num_intervals + 1
        error("Number of bounds does not match number of rates")
    end
    total = 0.
    bounds_ascending = true
    for i=1:num_intervals
        lower = bounds[i]
        upper = bounds[i+1]
        rate = rates[i]
        len = upper - lower
        if len <= 0
            bounds_ascending = false
        end
        total += len * rate
    end
    (total, bounds_ascending)
end

struct PiecewiseHomogenousPoissonProcess <: Distribution{Vector{Float64}} end
const piecewise_poisson_process = PiecewiseHomogenousPoissonProcess()

function Gen.logpdf(::PiecewiseHomogenousPoissonProcess, x::Vector{Float64}, bounds::Vector{Float64}, rates::Vector{Float64})
    cur = 1
    upper = bounds[cur+1]
    lpdf = 0.
    for xi in sort(x)
        if xi < bounds[1] || xi > bounds[end]
            error("x ($xi) lies outside of interval")
        end
        while xi > upper
            cur += 1
            upper = bounds[cur+1]
        end
        lpdf += log(rates[cur])
    end
    (total, bounds_ascending) = compute_total(bounds, rates)
    if bounds_ascending
        lpdf - total
    else
        -Inf
    end
end

function Gen.random(::PiecewiseHomogenousPoissonProcess, bounds::Vector{Float64}, rates::Vector{Float64})
    x = Vector{Float64}()
    num_intervals = length(rates)
    for i=1:num_intervals
        lower = bounds[i]
        upper = bounds[i+1]
        rate = (upper - lower) * rates[i]
        n = random(poisson, rate)
        for j=1:n
            push!(x, random(uniform_continuous, lower, upper))
        end
    end
    x
end

#########
# model #
#########

@gen function model(T::Float64)

    # prior on number of change points
    k = @addr(poisson(3.), "k")

    # prior on the location of (sorted) change points
    change_pts = Vector{Float64}(undef, k)
    lower = 0.
    for i=1:k
        cp = @addr(min_uniform_continuous(lower, T, k-i+1), "cp$i")
        change_pts[i] = cp
        lower = cp
    end

    # k + 1 rate values
    # h$i is the rate for cp$(i-1) to cp$i where cp0 := 0 and where cp$(k+1) := T
    alpha = 1.
    beta = 200.
    rates = Float64[@addr(Gen.gamma(alpha, 1. / beta), "h$i") for i=1:k+1]

    # poisson process
    bounds = vcat([0.], change_pts, [T])
    @addr(piecewise_poisson_process(bounds, rates), "points")
end

function render(trace; ymax=0.02)
    T = get_call_record(trace).args[1]
    assignment = get_assignment(trace)
    k = assignment["k"]
    bounds = vcat([0.], sort([assignment["cp$i"] for i=1:k]), [T])
    rates = [assignment["h$i"] for i=1:k+1]
    for i=1:length(rates)
        lower = bounds[i]
        upper = bounds[i+1]
        rate = rates[i]
        plt.plot([lower, upper], [rate, rate], color="black", linewidth=2)
    end
    points = assignment["points"]
    plt.scatter(points, -rand(length(points)) * (ymax/5.), color="black", s=5)
    ax = plt.gca()
    xlim = [0., T]
    plt.plot(xlim, [0., 0.], "--")
    ax[:set_xlim](xlim)
    ax[:set_ylim](-ymax/5., ymax)
end

function show_prior_samples()
    plt.figure(figsize=(16,16))
    T = 40000.
    for i=1:16
        plt.subplot(4, 4, i)
        trace = simulate(model, (T,))
        render(trace; ymax=0.015)
    end
    plt.tight_layout(pad=0)
    plt.savefig("prior_samples.pdf")
end

#############################
# height and position moves #
#############################

@gen function height_proposal(prev, i::Int)
    prev_assignment = get_assignment(prev)
    height = prev_assignment["h$i"]
    @addr(uniform_continuous(height/2., height*2.), "h$i")
end

@gen function position_proposal(prev, i::Int)
    prev_assignment = get_assignment(prev)
    k = prev_assignment["k"]
    lower = (i == 1) ? 0. : prev_assignment["cp$(i-1)"]
    upper = (i == k) ? T : prev_assignment["cp$(i+1)"]
    @addr(uniform_continuous(lower, upper), "cp$i")
end

function height_move(trace)
    k = get_assignment(trace)["k"]
    i = random(uniform_discrete, 1, k+1)
    mh(model, height_proposal, (i,), trace)
end

function position_move(trace)
    k = get_assignment(trace)["k"]
    i = random(uniform_discrete, 1, k)
    mh(model, position_proposal, (i,), trace)
end

######################
# birth / death move #
######################

# insert a new change point at i, where 1 <= i <= k+1
# the current change point at i, and all after, will be shifted right
# the new change point will be placed between cp$(i-1) and the current cp$i
@gen function birth_proposal(prev, T, i::Int)
    prev_assignment = get_assignment(prev)
    k = prev_assignment["k"]
    lower = (i == 1) ? 0. : prev_assignment["cp$(i-1)"]
    upper = (i == k+1) ? T : prev_assignment["cp$i"]
    @addr(uniform_continuous(lower, upper), "new-cp")
    @addr(uniform_continuous(0., 1.), "u")
end

@gen function death_proposal(prev) end

function birth_move_new_heights(cur_height, new_cp, prev_cp, next_cp, u)
    d_prev = new_cp - prev_cp
    d_next = next_cp - new_cp
    @assert d_prev > 0
    @assert d_next > 0
    d_total = d_prev + d_next
    log_cur_height = log(cur_height)
    log_ratio = log(1 - u) - log(u)
    new_h_prev = exp(log_cur_height - (d_prev / d_total) * log_ratio)
    new_h_next = exp(log_cur_height + (d_next / d_total) * log_ratio)
    @assert new_h_prev > 0.
    @assert new_h_next > 0.
    (new_h_prev, new_h_next)
end

const MODEL = :model
const PROPOSAL = :proposal

@inj function birth_injection(T, i::Int)

    # increment k
    k = @read(MODEL => "k")
    @write(k+1, MODEL => "k")

    # changepoints
    for j=1:i-1
        @copy(MODEL => "cp$j", MODEL => "cp$j")
    end
    @copy(PROPOSAL => "new-cp", MODEL => "cp$i")
    for j=i:k
        @copy(MODEL => "cp$j", MODEL => "cp$(j+1)")
    end

    # compute new heights
    cur_height = @read(MODEL => "h$i")
    prev_cp = (i == 1) ? 0. : @read(MODEL => "cp$(i-1)")
    next_cp = (i == k+1) ? T : @read(MODEL => "cp$i")
    new_cp = @read(PROPOSAL => "new-cp")
    u = @read(PROPOSAL => "u")
    (new_h_prev, new_h_next) = birth_move_new_heights(cur_height, new_cp, prev_cp, next_cp, u)

    # heights
    for j=1:i-1
        @copy(MODEL => "h$j", MODEL => "h$j")
    end
    @write(new_h_prev, MODEL => "h$i")
    @write(new_h_next, MODEL => "h$(i+1)")
    for j=i+1:k+1
        @copy(MODEL => "h$j", MODEL => "h$(j+1)")
    end

    @copy(MODEL => "points", MODEL => "points")
end

function death_move_u_new_height(prev_height, next_height, cur_cp, prev_cp, next_cp)
    d_prev = cur_cp - prev_cp
    d_next = next_cp - cur_cp
    @assert d_prev > 0
    @assert d_next > 0
    d_total = d_prev + d_next
    log_prev_height = log(prev_height)
    log_next_height = log(next_height)
    new_height = exp((d_prev / d_total) * log_prev_height + (d_next / d_total) * log_next_height)
    u = prev_height / (prev_height + next_height)
    @assert new_height > 0.
    (new_height, u)
end

@inj function death_injection(T, i::Int)

    # decrement k
    k = @read(MODEL => "k")
    @assert k > 0
    @write(k-1, MODEL => "k")

    # change points
    for j=1:i-1
        @copy(MODEL => "cp$j", MODEL => "cp$j")
    end
    for j=i+1:k
        @copy(MODEL => "cp$j", MODEL => "cp$(j-1)")
    end

    # compute new height
    cur_cp = @read(MODEL => "cp$i")
    prev_cp = (i == 1) ? 0. : @read(MODEL => "cp$(i-1)")
    next_cp = (i == k) ? T : @read(MODEL => "cp$(i+1)")
    prev_height = @read(MODEL => "h$i")
    next_height = @read(MODEL => "h$(i+1)")
    (new_height, u) = death_move_u_new_height(prev_height, next_height, cur_cp, prev_cp, next_cp)

    # heights
    for j=1:i-1
        @copy(MODEL => "h$j", MODEL => "h$j")
    end
    @write(new_height, MODEL => "h$i")
    for j=i+2:k+1
        @copy(MODEL => "h$j", MODEL => "h$(j-1)")
    end

    @copy(MODEL => "points", MODEL => "points")

    @copy(MODEL => "cp$i", PROPOSAL => "new-cp")
    @write(u, PROPOSAL => "u")
end

function birth_move(trace)
    # if k > 0, then prob_b = 0.25 and prob_d = 0.25
    # the probability that the one we introduce will get deleted, given that
    # we choose a death move, is: 1/(k+1); which is also the probability that
    # we choose the i that we choose here
    k = get_assignment(trace)["k"]
    i = random(uniform_discrete, 1, k+1)
    # prob_b = 1, but prob_d = 0.25, so we correct by -log(4)
    correction = (new_trace) -> (k == 0 ? -log(4) : 0)
    T = get_call_record(trace).args[1]
    rjmcmc(model,
        birth_proposal, (T, i),
        death_proposal, (),
        birth_injection, (T, i),
        trace, correction)
end

function death_move(trace)
    k = get_assignment(trace)["k"]
    @assert k > 0
    i = random(uniform_discrete, 1, k)
    correction = (new_trace) -> (k == 0 ? log(4) : 0)
    T = get_call_record(trace).args[1]
    rjmcmc(model,
        death_proposal, (),
        birth_proposal, (T, i),
        death_injection, (T, i),
        trace, correction)
end

##########################
# Generic MCMC inference #
##########################

function resimulation_mh(selection, trace)
    model_args = get_call_record(trace).args
    (new_trace, weight) = regenerate(model, model_args, NoChange(), trace, selection)
    if log(rand()) < weight
        # accept
        return new_trace
    else
        # reject
        return trace
    end
end

k_selection = DynamicAddressSet()
push_leaf_node!(k_selection, "k")

function generic_mcmc_step(trace)
    k = get_assignment(trace)["k"]
    if k > 0
        prob_h = 1./3
        prob_p = 1./3
        prob_change_k = 1./3
    else
        prob_h = 0.
        prob_p = 0.
        prob_change_k = 1
    end
    move_type = random(categorical, [prob_h, prob_p, prob_change_k])
    if move_type == 1
        height_move(trace)
    elseif move_type == 2
        position_move(trace)
    else
        resimulation_mh(k_selection, trace)
    end
end

#########################
# RJMCMC MCMC inference #
#########################0

function mcmc_step(trace)
    k = get_assignment(trace)["k"]
    if k > 0
        prob_h = 0.25
        prob_p = 0.25
        prob_b = 0.25
        prob_d = 0.25
    else
        prob_h = 0.
        prob_p = 0.
        prob_b = 1.
        prob_d = 0.
    end
    move_type = random(categorical, [prob_h, prob_p, prob_b, prob_d])
    if move_type == 1
        height_move(trace)
    elseif move_type == 2
        position_move(trace)
    elseif move_type == 3
        birth_move(trace)
    elseif move_type == 4
        death_move(trace)
    else
        error("Unknown move type $move_type")
    end
end

function do_mcmc(T, num_steps::Int)
    (trace, _) = generate(model, (T,), observations)
    for iter=1:num_steps
        if iter % 1000 == 0
            println("iter $iter of $num_steps, k: $(get_assignment(trace)["k"])")
        end
        #trace = mcmc_step(trace)
        trace = generic_mcmc_step(trace)
    end
    trace
end

########################
# inference experiment #
########################

Gen.load_generated_functions()

import Random
Random.seed!(1)

# load data set
import CSV
function load_data_set()
    df = CSV.read("coal.csv")
    dates = df[1]
    dates = dates .- minimum(dates)
    dates * 365.25 # convert years to days
end

const points = load_data_set()
const T = maximum(points)
const observations = DynamicAssignment()
observations["points"] = points

function show_posterior_samples()
    plt.figure(figsize=(16,16))
    for i=1:16
        println("replicate $i")
        tic()
        plt.subplot(4, 4, i)
        trace = do_mcmc(T, 5000)#10000)
        toc()
        render(trace; ymax=0.015)
    end
    plt.tight_layout(pad=0)
    plt.savefig("posterior_samples.pdf")
end

function get_rate_vector(trace, test_points)
    assignment = get_assignment(trace)
    k = assignment["k"]
    cps = [assignment["cp$i"] for i=1:k]
    hs = [assignment["h$i"] for i=1:k+1]
    rate = Vector{Float64}()
    cur_h_idx = 1
    cur_h = hs[cur_h_idx]
    next_cp_idx = 1
    upper = (next_cp_idx == k + 1) ? T : cps[next_cp_idx]
    for x in test_points
        while x > upper
            next_cp_idx += 1
            upper = (next_cp_idx == k + 1) ? T : cps[next_cp_idx]
            cur_h_idx += 1
            cur_h = hs[cur_h_idx]
        end
        push!(rate, cur_h)
    end
    rate
end

# compute posterior mean rate curve

function plot_posterior_mean_rate()
    test_points = collect(1.0:10.0:T)
    rates = Vector{Vector{Float64}}()
    num_samples = 0
    num_steps = 5000 # 20000
    for reps=1:10
        (trace, _) = generate(model, (T,), observations)
        for iter=1:num_steps
            if iter % 1000 == 0
                println("iter $iter of $num_steps, k: $(get_assignment(trace)["k"])")
            end
            trace = mcmc_step(trace)
            if iter > 4000
                num_samples += 1
                rate_vector = get_rate_vector(trace, test_points)
                @assert length(rate_vector) == length(test_points)
                push!(rates, rate_vector)
            end
        end
    end
    posterior_mean_rate = zeros(length(test_points))
    for rate in rates
        posterior_mean_rate += rate / Float64(num_samples)
    end
    ymax = 0.010
    plt.figure()
    plt.plot(test_points, posterior_mean_rate, color="black")
    plt.scatter(points, -rand(length(points)) * (ymax/6.), color="black", s=5)
    ax = plt.gca()
    xlim = [0., T]
    plt.plot(xlim, [0., 0.], "--")
    ax[:set_xlim](xlim)
    ax[:set_ylim](-ymax/5., ymax)
    plt.savefig("posterior_mean_rate.pdf")
end

function plot_trace_plot()
    # show the number of clusters
    (trace, _) = generate(model, (T,), observations)
    num_clusters_vec = Int[]
    burn_in = 20000
    for iter=1:burn_in + 5000
        (trace, accept) = mcmc_step(trace)
        if iter > burn_in
            push!(num_clusters_vec, get_assignment(trace)["k"])
        end
    end
    plt.figure()
    plt.plot(num_clusters_vec)
    ax = plt.gca()
    plt.savefig("trace_plot_rjmcmc.pdf")
end

function plot_trace_plot()
    plt.figure(figsize=(8, 4))

    # generic
    (trace, _) = generate(model, (T,), observations)
    num_clusters_vec = Int[]
    burn_in = 20000
    for iter=1:burn_in + 5000
        trace = generic_mcmc_step(trace)
        if iter > burn_in
            push!(num_clusters_vec, get_assignment(trace)["k"])
        end
    end
    plt.subplot(2, 1, 1)
    plt.plot(num_clusters_vec, "r")

    # reversible jump
    (trace, _) = generate(model, (T,), observations)
    height1 = Float64[]
    num_clusters_vec = Int[]
    burn_in = 20000
    for iter=1:burn_in + 5000
        trace = mcmc_step(trace)
        if iter > burn_in
            push!(num_clusters_vec, get_assignment(trace)["k"])
        end
    end
    plt.subplot(2, 1, 2)
    plt.plot(num_clusters_vec, "b")

    ax = plt.gca()
    plt.savefig("trace_plot.pdf")
end

println("showing prior samples...")
show_prior_samples()

println("showing posterior samples...")
show_posterior_samples()

println("estimating posterior mean rate...")
plot_posterior_mean_rate()

println("making trace plot...")
plot_trace_plot()


using BenchmarkTools, FinancialMonteCarlo, ForwardDiff

S0 = 100.0;
K = 100.0;
r = 0.02;
T = 1.0;
d = 0.01;
D = 90.0;

Nsim = 100000;
Nstep = 30;
sigma = 0.2
mc = MonteCarloConfiguration(Nsim, Nstep);
toll = 1e-3;

rfCurve = ZeroRate(r);

FwdData = Forward(T)
EUData = EuropeanOption(T, K)
AMData = AmericanOption(T, K)
BarrierData = BarrierOptionDownOut(T, K, D)
AsianFloatingStrikeData = AsianFloatingStrikeOption(T)
AsianFixedStrikeData = AsianFixedStrikeOption(T, K)
Model = BlackScholesProcess(sigma, Underlying(S0, d));

f__1(x::Vector) = pricer(BlackScholesProcess(x[1], Underlying(x[2], x[4])), ZeroRate(x[3]), mc, EuropeanOption(x[5], K));
x = Float64[sigma, S0, r, d, T]
g__1 = x -> ForwardDiff.gradient(f__1, x);
y0 = g__1(x);
@btime f__1(x);
@btime g__1(x);

f_(x::Vector) = pricer(BlackScholesProcess(x[1], Underlying(x[2], x[4])), ZeroRate(x[3]), mc, AmericanOption(x[5], K));
g_ = x -> ForwardDiff.gradient(f_, x);
y0 = g_(x);
@btime f_(x);
@btime g_(x);


"""
	MMCK(μ, λ, c, k)

Tạo mô hình M/M/c/k
"""
struct MMCK{T} <:AbstractMMCK
	λ::T
	μ::T
	c::Union{T, Integer}
	k::Union{T, Integer}
	r::Union{T, Real}
	ρ::Union{T, Real}
	function MMCK(λ, μ, c, k)
		T = Union{typeof(μ), typeof(λ)}
		new{T}(λ, μ, c,	k, λ/μ,	λ/μ/c)
	end
end

function pn(m::MMCK, n::Integer)
	r = m.r
	k = m.k
	c = m.c
	if n == 0
		1 / (
		     sum(r^i / factorial(i) for i = 0:c-1) +
		     sum(r^i / (c^(i-c) * factorial(c)) for i = c:k)
		     )
	elseif n < c
		pn(m, 0) * r^n / factorial(n)
	elseif n <= k
		pn(m, 0) * r^n / (c^(n-c) * factorial(c))
	else 0
	end
end

function L(m::MMCK)
	sum(i * pn(m, i) for i=0:m.k)
end

function Lq(m::MMCK)
	sum((i - m.c) * pn(m, i) for i=m.c:m.k)
end

function W(m::MMCK)
	L(m) / λe(m)
end

function Wq(m::MMCK)
	Lq(m) / λe(m)
end


# This file is a part of Julia. License is MIT: http://julialang.org/license

# https://github.com/JuliaLang/julia/issues/4707

#----------------- Get the A matrix
function getDivGrad(n1,n2,n3)
        D1 = kron(speye(n3),kron(speye(n2),ddx(n1)))
        D2 = kron(speye(n3),kron(ddx(n2),speye(n1)))
        D3 = kron(ddx(n3),kron(speye(n2),speye(n1)))

        # DIV from faces to cell-centers
        Div = [D1 D2 D3]
        return Div*Div';
end

#----------------- 1D finite difference on staggered grid
function ddx(n)
    return spdiags(ones(n)*[-1 1],[0,1],n,n+1)
end

#------------- Build a diagonal matrix
function spdiags(B,d,m,n)
    d = d[:]
    p = length(d)

    len = zeros(Int, p+1, 1)
    for k = 1:p
        len[k+1] = len[k] + length(max(1,1-d[k]):min(m,n-d[k]))
    end

    a = zeros(Int, len[p+1], 3)
    for k = 1:p
        # Append new d[k]-th diagonal to compact form
        i = max(1,1-d[k]):min(m,n-d[k])
        a[(len[k]+1):len[k+1],:] = [i i+d[k] B[i+(m>=n)*d[k],k]]
    end

    return sparse(a[:,1],a[:,2],a[:,3],m,n)
end


export Zernike,
       zernike_norm,
       zernike_nm_to_fringe,
       zernike_nm_to_ansi_j,
       zernike_ansi_j_to_nm,
       zernike_noll_to_nm,
       zernike_fringe_to_nm,
       zernike_zero_separation

"""
    Zernike([T, ]n, m, norm=true)

Zernike polynomial of orders `n` and `m`. If `norm` is true the output will be normalized. The output type can be specified with `T`, which will default to `Float64`.

# Examples

```jldoctest
julia> z = Zernike(10, 3)
Zernike{Float64}(n=10, m=3)

julia> z(0.8, 0.2)
-1.6860955965619917

julia> Zernike(BigFloat, 10, 3)(0.8, 0.2)
-1.68609559656199170518675600760616362094879150390625
```
"""
struct Zernike{T<:AbstractFloat,F1,F2}
    n::Int
    m::Int
    basis::F1
    trig_func::F2
    Zernike{T}(n, m, basis, trig_func) where {T} = new{T,typeof(basis),typeof(trig_func)}(n, m, basis, trig_func)
end

function Zernike(T, n, m, norm::Bool=true)
    norm_val = norm ? zernike_norm(n, m) : 1.0
    trig_func = m < 0 ? sin : cos
    basis = x -> norm_val * jacobi((n - m) ÷ 2, 0, abs(m), x)
    Zernike{T}(n, m, basis, trig_func)
end

Zernike(n, m, norm::Bool=true) = Zernike(Float64, n, m, norm)

function Base.show(io::IO, z::Zernike{T}) where T
    print(io, "Zernike{$T}(n=$(z.n), m=$(z.m))")
end

function (z::Zernike{T})(ρ, θ) where T
    x = ρ^2 - 1
    return T(ρ^abs(z.m) * z.basis(x) * z.trig_func(z.m * θ))
end

"""
    kronecker(i,j)

1 if i==j, else 0; mathematical kronecker function
"""
kronecker(i, j) = Int(i == j)

"""
    zernike_norm(n, m)

Norm of Zernike polynomial of radial order n, azimuthal order m.

The norm is the average squared distance to zero.  By multiplying a zernike
value by the norm, the term is given unit stdev or RMS.
"""
function zernike_norm(n, m)
    num = √(2 * (n+1)) / (1 + kronecker(m, 0))
end

"""
    zernike_nm_to_fringe(n, m)

Map (n,m) ANSI indices to a single fringe index.
"""
function zernike_nm_to_fringe(n, m)
    term1 = (1 + (n + abs(m))/2)^2
    term2 = 2*abs(m)
    term3 = (1 + sign(m)) / 2
    return int(term2 - term2 - term3) + 1
end

"""
    zernike_nm_to_ansi_j(n, m)

Map (n,m) ANSI indices to a single ANSI j index.

See also:
    - [`zernike_ansi_j_to_nm`](@ref) (reciprocal of this function)
"""
function zernike_nm_to_ansi_j(n, m)
    return int((n * (n + 2) + m) / 2)
end

"""
    zernike_ansi_to_ansi_j(n, m)

Map (n,m) ANSI indices to a single ANSI j index.

See also:
    - [`zernike_nm_to_ansi_j`](@ref) (reciprocal of this function)
"""
function zernike_ansi_j_to_nm(j)
    n = int(ceil((-3 + √(9 + 8j))/2))
    m = 2j - n * (n + 2)
    return n, m
end

"""
    zernike_noll_to_nm(j)

Map j Noll index to ANSI (n,m) indices.
"""
function zernike_noll_to_nm(j)
    n = int(ceil((-1 + √(1 + 8j))/2) - 1)
    if n == 0
        m = 0
    else
        nseries = int((n+1) * (n+2) / 2)
        residual = j - nseries - 1

        if isodd(j)
            sign = -1
        else
            sign = 1
        end

        if isodd(n)
            ms = [1,1]
        else
            ms = [0]
        end

        for i=0:n÷2
            push!(ms, ms[end]+2)
            push!(ms, ms[end])
        end

        m = ms[residual] * sign
    end
    return n, m
end

"""
    zernike_fringe_to_nm(j)

Map j Fringe index to ANSI (n,m) indices.
"""
function zernike_fringe_to_nm(j)
    m_n = 2 * ceil(√j - 1)
    g_s = (m_n / 2)^2 + 1
    n = m_n / 2 + floor((j-g_s)/2)
    m = m_n - n * (1 - mod(j-g_s, 2) * 2)
    return int(n), int(m)
end

"""
    zernike_zero_separation(n)

Minimum zero separation of Zernike polynomial of radial order n.  Useful for
computing sample count requirements.
"""
function zernike_zero_separation(n)
    return 1 / n^2
end

"""
    zernike(n, m, ρ, θ[; norm])

Zernike polynomial of radial order n and azimuthal order m, evaluated at the
point (ρ, θ).  No normalization is required of (ρ, θ), though the polynomials
are orthogonal only over the unit disk.

norm is a boolean flag indicating whether the result should be orthonormalized
(scaled to unit RMS) or not.
"""
function zernike(n, m, ρ, θ; norm::Bool=true)
    x = ρ^2 - 1
    n_j = (n - m) / 2
    am = abs(m)
    # α=0, β=|m|
    # there is a second syntax where you have x reversed, 1 - ρ^2,
    # in which ase you swap α and β.  It makes absolutely no difference
    out = jacobi(n_j, 0, am, x)
    if m != 0
        if sign(m) == -1
            f = sin
        else
            f = cos
        end
        out *= (ρ^am * f(m*θ))
    end
	if norm
		out *= zernike_norm(n,m)
	end
    return out
end


using DataFrames, GLM, Distributions, LinearAlgebra, Random
using Plots, LaTeXStrings;pyplot()
Random.seed!(0)

beta0, beta1 = 2.0, 1.5
sigma = 2.5
n, N = 10, 10^4
alpha = 0.05

function coefEst()
    xVals = collect(1:n)
    yVals = beta0 .+ beta1*xVals + rand(Normal(0,sigma),n)
    data = DataFrame([xVals,yVals],[:X,:Y])
    model = lm(@formula(Y ~ X), data)
    coef(model)
end

ests = [coefEst() for _ in 1:N]

xBar = mean(1:n)
sXX = sum([(x - xBar)^2 for x in 1:n])
sx2 = sum([x^2 for x in 1:n])
var0 = sigma^2 * sx2/(n*sXX)
var1 = sigma^2/sXX
cv = -sigma^2*xBar/sXX

mu = [beta0, beta1]
Sigma = [var0 cv; cv var1]

A = cholesky(Sigma).L
Ai = inv(A)

r = quantile(Rayleigh(),1-alpha)
isInEllipse(x) = norm(Ai*(x-mu)) <= r
estIn = isInEllipse.(ests)

println("Proportion of points inside ellipse: ", sum(estIn)/N)

scatter(first.(ests[estIn]),last.(ests[estIn]),c=:green, ms=2, msw=0)
scatter!(first.(ests[.!estIn]),last.(ests[.!estIn]),c=:blue, ms=2, msw=0)

ellipsePts = [r*A*[cos(t),sin(t)] + mu  for t in 0:0.01:2pi]
scatter!([beta0],[beta1],c=:red, ms=5, msw = 0)
plot!(first.(ellipsePts),last.(ellipsePts),
	c=:red, lw=2, legend=:none,
	xlabel=L"\hat{\beta}_0", ylabel=L"\hat{\beta}_1")


using BenchmarkTools
using StaticArrays
using Test
using LoopVectorization

function pot(x,y,eps4,sig6,sig12)
  r2 = (x[1]-y[1])^2 + (x[2]-y[2])^2
  r6 = r2^3
  r12 = r6^2
  return eps4*(sig12/r12 - sig6/r6)
end

function total_pot(p,eps,sig)
  eps4 = 4*eps
  sig6 = sig^6
  sig12 = sig6^2
  total_pot = 0.
  for i in 1:length(p)-1
    @avx for j in i+1:length(p)
      total_pot = total_pot + pot(p[i],p[j],eps4,sig6,sig12)
    end
  end
  return total_pot
end

eps = 4.
sig = 1.
p = [ 10*rand(2) for i in 1:10 ]
r1 = @btime total_pot(p,eps,sig)

p = [ SVector(p[i]...) for i in 1:10 ]
r2 = @btime total_pot(p,eps,sig)

@test r1 ≈ r2


#Runs a simple Metropolis-Hastings (ie MCMC) algoritm to simulate an
#exponential distribution, which has the probability density
#p(x)=exp(-x/m^2), where m>0.
#
#Author: H. Paul Keeler, 2019.
# Website: hpaulkeeler.com
# Repository: github.com/hpaulkeeler/posts

#Note: Need the .+ for adding a scalar to an array
#Also need . for sqrt, exp, cos, sin etc and assinging scalars to arrays

#clearconsole(); #for clearing Julia REPL console

using Distributions #for random simulations
using PyPlot #for plotting
using StatsBase #for histograms etc
using Random
using LinearAlgebra
PyPlot.close("all");  # close all PyPlot figures

#set random seed for reproducibility
#Random.seed!(1234)

numbSim=10^4; #number of random variables simulated
numbSteps=25; #number of steps for the Markov process
numbBins=50; #number of bins for histogram

sigma=1; #standard deviation for normal random steps
m=2; #parameter (ie mean) for distribution to be simulated

function fun_p(x)
    return ((exp.(-x./m)./m).*(x.>0));
end

xRand=rand(numbSim); #random intial values
probCurrent=fun_p(xRand); #current transition probabilities

for jj=1:numbSteps
    zRand= xRand.+sigma.*rand(Normal(),numbSim);#take a (normally distributed) random step
    #zRand= xRand +2*sigma*(rand(size(xRand))-0.5);#take a (uniformly distributed) random step
    probProposal=fun_p(zRand); #proposed probability

    #acceptance rejection step
    booleAccept=rand(numbSim) .< probProposal./probCurrent;
    #update state of random walk/Markov chain
    xRand[booleAccept]=zRand[booleAccept];
    #update transition probabilities
    probCurrent[booleAccept]=probProposal[booleAccept];

end

#histogram section: empirical probability density
histX=fit(Histogram, xRand,nbins=numbBins); #find histogram data
histX=normalize(histX,mode=:pdf); #normalize histogram
binEdges=histX.edges; #retrieve bin edges
xValues=(binEdges[1][2:end]+binEdges[1][1:end-1])./2; #mid-points of bins
pdfEmp=(histX.weights)
#analytic solution of probability density
pdfExact=fun_p(xValues);

# Plotting
PyPlot.plot(xValues, pdfExact)
PyPlot.scatter(xValues, pdfEmp, marker="x", c="r");
PyPlot.grid(true);
PyPlot.xlabel("x");
PyPlot.ylabel("p(x)");
PyPlot.show();


function update_derivs_pos!(state::Vector{Float64}, derivs::Vector{Float64} )
 # Input: state = [x,y,vx,vy], a vector of two 2-d positions and velocities for a test particle
 # Output: The derivatives of the position are updated in the preallocated array derivs.
  @assert length(state) == 4
  @assert length(derivs) == 4
  v_x = state[3]
  v_y = state[4]
  derivs[1] = v_x
  derivs[2] = v_y
  return derivs;
end

function update_derivs_vel!(state::Vector{Float64}, derivs::Vector{Float64} )
  # Input: state = [x,y,vx,vy], a vector of two 2-d positions and velocities for a test particle
  # Output: The derivatives of the velocity are updated in the preallocated array derivs.
  @assert length(state) == 4
  @assert length(derivs) == 4
  GM = 1.0
  r_x = state[1]
  r_y = state[2]
  r2 = r_x*r_x+r_y*r_y
  a = -GM/r2
  r = sqrt(r2)
  a_x = a * r_x/r
  a_y = a * r_y/r
  derivs[3] = a_x
  derivs[4] = a_y
  return derivs;
end

function update_derivs!(state::Vector{Float64}, derivs::Vector{Float64} )
  # Input: state = [x,y,vx,vy], a vector of two 2-d positions and velocities for a test particle
  # Output: The derivatives are updated in the preallocated array derivs.
  update_derivs_vel!(state,derivs)
  update_derivs_pos!(state,derivs)
  return derivs
end

function advance_leapfrog!(state::Vector{Float64},derivs::Vector{Float64}, dt::Float64; derivs_current::Bool = false)
  # Input/Output: state = array of two 2-d positions and velocities for a test particle
  # Temporary space: The derivatives are updated in the preallocated array derivs.
  # Input: dt is the fixed time step
  # Optional param: derivs_current: whether need to calculate derivatives at beginning
  @assert length(state) == length(derivs)

  if !derivs_current
    update_derivs_pos!(state,derivs);
  end
  state[1] += 0.5*dt*derivs[1]
  state[2] += 0.5*dt*derivs[2]
  update_derivs_vel!(state,derivs);
  state[3] += dt*derivs[3]
  state[4] += dt*derivs[4]
  update_derivs_pos!(state,derivs);
  state[1] += 0.5*dt*derivs[1]
  state[2] += 0.5*dt*derivs[2]
end

# Input/Output: state = [x,y,vx,vy], an array of two 2-d positions and velocities for a test particle
# Input: dt is the fixed time step
# Input: duration is the total
function integrate_leapfrog!(state::Vector{Float64}, dt::Float64, duration::Float64; max_num_log::Integer = 100000)
  @assert(length(state)==4)
  @assert(dt>0.0)
  @assert(duration>0.0)

  # Preallocate array to hold data log (including  initial state)
  nsteps = iceil(duration/dt);
  nskip = (nsteps<max_num_log) ? 1 : iceil(nsteps/(max_num_log-1))
  num_log = iceil(nsteps/nskip)+1
  log = Array(Float64,(num_log,length(state)));

  # Pre-allocate and pre-compute derivaties
  derivs = Array(Float64,4);
  update_derivs!(state,derivs);

  # Log initial state
  log_pos = 1
  log[log_pos,:] = deepcopy(state)

  n = 0
  t = 0.0
  while t<duration
    # ensure don't integrate for more than duration
    dt_tmp = (t+dt<=duration) ? dt : duration-t;

	# advance system by one time step
    advance_leapfrog!(state,derivs,dt_tmp, derivs_current=true)
    t = t + dt_tmp
    n = n + 1

    if (n%nskip==0) # Log data
	   log_pos += 1
   	   @assert( log_pos<=length(log) )
	   @assert( length(log[log_pos,:])==length(state) )
	   log[log_pos,:] = deepcopy(state)
	end
  end
  return log
end

function calc_error_leapfrog_old(dur::Float64, dt::Float64 = 2pi/200.0)
  state = [1.,0.,0.,1.];
  integrate_leapfrog!(state,dt,dur*2pi);
  dist = state[1]^2+state[2]^2
  phase = atan2(state[2],state[1])
  offset = sum((state[1:2].-[1.0,0.0]).^2)
  return (dist-1.0,phase,offset)
end

function calc_end_distance_leapfrog(dur::Integer, dt::Float64 = 2pi/200.0, state::Vector{Float64} = [1., 0., 0., 1.] )
  #state = [1.,0.,0.,1.];
  dist_init = sqrt(state[1]^2+state[2]^2)
  phase_init = atan2(state[2],state[1])
  integrate_leapfrog!(state,dt,dur*2pi);
  # Calculate three metrics of the accuracy of the integration
  dist = sqrt(state[1]^2+state[2]^2)
  phase = atan2(state[2],state[1])
  offset = sqrt(sum((state[1:2].-[1.0,0.0]).^2))
  return (dist-dist_init,phase-phase_init,offset)
end

using Base.Test
function test_leapfrog(dur::Integer, dt::Float64 = 2pi/200.0)
  err = calc_end_distance_leapfrog(dur,dt)
  @test_approx_eq_eps(err[1], 0., 1e-6)
  @test_approx_eq_eps(err[2], 0., 1e-1)
  @test_approx_eq_eps(err[3], 0., 1e-1)
end


### A Pluto.jl notebook ###
# v0.14.5

using Markdown
using InteractiveUtils

# ╔═╡ f6d978ea-b349-11eb-309e-f5a517efc907
begin
	using LinearAlgebra
	using Plots
	using QuadGK
end

# ╔═╡ 99268461-9e3f-4685-b530-e24febfb539e
#=
IDEA OF MULTIGRID :

	G is a mesh of a 1D domain.
	G1 is another, but coarser.
	G1's operator has higher frequency eigenvalues as opposed to  G.
	Use G1 to weed out these values.
	Use smoothers alternatively.
=#

# ╔═╡ b422b21d-b9d3-4aad-b29b-d0669331abc8

begin
	e = 5; #  Tune e for different basis functions
	N = 10;
	Ni = 100;
	xmin = 0;
	xmax = 1;
	x = LinRange(xmin,xmax,N);
	xi = LinRange(x[e-1],x[e+1],Ni);
	h = x[2] - x[1];
end

# ╔═╡ b11596d4-ee5d-41ec-af8f-96734ba2f579
# This shows a heirarchy of grids
grid1 = plot(x,zeros(length(x)),marker="*")

# ╔═╡ c1463b68-12e6-4e8c-9902-c2cc7e6caa91
begin
	# Plot the basis functions
	y = zeros(length(xi));
	for i = 1:length(xi);
		if(xi[i]>=xi[1] && xi[i]<=x[e])
			y[i] = 1/h * (xi[i]-x[e-1]);
		elseif (xi[i]>x[e] && xi[i]<x[e+1])
			y[i] = 1/h * (x[e+1]-xi[i]);
		end
	end
	plot!(xi,y,title="Linear Basis Functions");
end

# ╔═╡ d5da54e6-0a4f-4f43-a7ab-8791574b32c6
# Now we use the linear basis functions and obtain an approximation of the Laplacian operator.

begin
#>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
	function θ(x,i,N)
		n = length(x);
		y = zeros(n);
		if i==1
			for k=1:N
				if x[k]<= x[i+1]
					y[k] = (N-1)*(x[i+1]-x[k]);
				end
			end
		elseif i==N
			for l = 1:N
				if x[l] >= x[i-1]
					y[l] = (N-1) * (x[l]-x[i]);
				end
			end
		else
			for j=2:n-1
				if(x[j]>=x[i-1] && x[j]<=x[i])
					y[j] = (N-1) * (x[j]-x[i-1]);
				elseif(x[j]>=x[i] && x[j]<=x[i+1])
					y[j] = (N-1) * (x[i+1]-x[j]);
				end
			end
		end
		return y;
	end
#>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
	function θ_scalar(x,i,xi,N)
			if(x>=xi[i-1] && x<=xi[i])
				y = (N-1) * (x-xi[i-1]);
			elseif(x>=xi[i] && x<=xi[i+1])
				y = (N-1) * (xi[i+1]-x);
			else
				y=0;
			end
		return y;
	end
#>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
	function stiffness_matrix(n)
		md = 2.0*ones(n-2);
		ld = -1.0*ones(n-3);
		ud = -1.0*ones(n-3);
		return (n-1)*Tridiagonal(ld,md,ud);
	end
#>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
	function intergrate(x,i,N)
		domain = LinRange(0,1,N+2);
		return sin(pi*x) * θ_scalar(x,i,domain,N);
	end

	function compute_rhs(n)
		n = n-2;
		integral = zeros(n);
		for i=1:length(integral)
		integral[i] = 1.5*quadgk(q->intergrate(q,i+1,n),0,1,rtol=1e-8)[1];
		end
		return integral;
	end
#>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
end

# ╔═╡ a5f1a911-3f71-4630-a872-99514d52e5f3
begin
	nk = 12;
	a = LinRange(xmin,xmax,nk);
	z = θ(a,3,nk);
	plot(a,z)
end

# ╔═╡ feba23c9-1b66-400e-9ffb-0dc5fc65896f
begin
	# Get the stiffness matrix and the RHS : The RHS is evaluated and plotted.
	n = 7;
	test_A = stiffness_matrix(n)
	test_b = compute_rhs(n);
	plot(test_b)
	#test_A
	#test_b
end

# ╔═╡ b8b7b670-6039-4bd9-9f95-8652f6e03528
md"""

	#NESTED MULTIGRID

	Start on a coarse grid h_0

	Iterate using a smoother and remove one range of high frequencies.

	Interpolate solution to a finer grid - h_1

	Start step 1 with a smaller time step h_1.
"""

# ╔═╡ a217bef6-3c9f-4d12-8b3f-eb937987c709
# Description

md"""

To implement a nested multigrid method.

We would need access to a smoother. The kind of smoother to choose depends on the nature of the matrix A that we solve.

For the finite element approximation of the poisson equation we have a matrix A that is not diagonally dominant.

So solve the problem we need a method for which convergence could be proven.

In comes the Gauss Seidel method.

"""

# ╔═╡ f2a7eb59-3617-47fe-be19-f6873d0f3c56
function gauss_seidel_relaxation(A,b,x0,max_iterations)
	x = x0;
	res = b - A*x;
	for iter = 1:max_iterations
		for i = 1:length(x)
			x[i] = (1/A[i,i])*(b[i] - dot(A[i,:],x) + A[i,i]*x[i]);
		end
		res = b - A*x;
	end
	return x
end

# ╔═╡ abf72f94-406b-41a6-bbdd-fa3ac2919531
# To perform a nested iteration
begin
	n_nodes = 7;
	A = stiffness_matrix(n_nodes);
	b = compute_rhs(n_nodes);
	x0 = zeros(n_nodes-2);

	y_appx = gauss_seidel_relaxation(A,b,x0,10);
	y_appx = vcat(0,y_appx,0);
	println(y_appx)

	x_nodes = LinRange(0,1,100);
	y_actual = zeros(length(x_nodes));
	for i=1:length(x_nodes)
		y_actual[i] = 1/pi^2 * sin(pi*x_nodes[i]);
	end
	plot(x_nodes,y_actual,label="Actual Solution");
	plot!(LinRange(0,1,length(y_appx)),y_appx,label="Approximate solution");
end

# ╔═╡ Cell order:
# ╠═f6d978ea-b349-11eb-309e-f5a517efc907
# ╠═99268461-9e3f-4685-b530-e24febfb539e
# ╠═b422b21d-b9d3-4aad-b29b-d0669331abc8
# ╠═b11596d4-ee5d-41ec-af8f-96734ba2f579
# ╠═c1463b68-12e6-4e8c-9902-c2cc7e6caa91
# ╠═d5da54e6-0a4f-4f43-a7ab-8791574b32c6
# ╠═a5f1a911-3f71-4630-a872-99514d52e5f3
# ╠═feba23c9-1b66-400e-9ffb-0dc5fc65896f
# ╟─b8b7b670-6039-4bd9-9f95-8652f6e03528
# ╟─a217bef6-3c9f-4d12-8b3f-eb937987c709
# ╠═f2a7eb59-3617-47fe-be19-f6873d0f3c56
# ╠═abf72f94-406b-41a6-bbdd-fa3ac2919531


# Some built-in functions
println("println is a function")
println("parse(Int64, \"123\"): $(parse(Int64, "123"))")
println("parse(Float64, \"123.456\"): $(parse(Float64, "123.456"))")
println("trunc(Int64, 3.67123): $(trunc(Int64, 3.67123))")
println("trunc(Int64, -3.67123): $(trunc(Int64, -3.67123))")
println("Other examples: ", float(32), "\t", string(32), "\t", string(123.456))
println("Natural log and base, log(exp(1)): $(log(exp(1)))")
println("Other log bases, log2(10), log10(10): $(log2(10)), $(log10(10))")
println("Trigonometry, sin(pi), cos(pi), tan(pi): $(sin(pi)), $(cos(pi)), $(tan(pi))")
println("sqrt(2) = $(sqrt(2))")

# Defning functions
function printlines()
    println("This is the first line.")
    println("This is the second line.")
end

# printlines()

function printmorelines()
    printlines()
    printlines()
end

# printmorelines()

function printargument(arg)
    println(arg)
end

printargument(123)
printargument("123456")
printargument(cos(pi/2))

# Locality of variables and parameters
function sumthree(a, b, c)
    a = a + b
    total = a
    total += c
    return total
end
d, e, f = 1, 2, 3
println("sumthree(d, e, f): $(sumthree(d, e, f))")
println("d, e, f: $d, $e, $f")
# These lines will cause an error:
# println(a, b, c)
# println(total)

# Example of void function that returns nothing:
println(printlines())

# Example of a docstring
"""
polyline(t, n, len, angle)

Draws n line segments with the given length and
angle (in degrees) between them.  t is a turtle.
"""
function polyline(t, n, len, angle)
end

# Example of including errors
function fact(n)
    if !(n isa Int64)
        error("Factorial is only defined for integers.")
    elseif n < 0
        error("Factorial is not defined for negative integers.")
    elseif n == 0
        return 1
    else
        return n * fact(n-1)
    end
end

# println(fact(1.5))
# println(fact(-1))

# Shorthand notation for function definition
function func_ex(a=1, b=2)
    return a + 2b
end

# The above function definition is equivalent to the following:
func(a, b) = a + 2b
func(a) = func(a, 2)
func() = func(1, 2)
println("func(7, 5): $(func(7, 5))")
println("func(2): $(func(2))")
println("func(): $(func())")


# for plotif. This identifies a vector of colors
function identify_colors(g, xs, colors=(:red, :blue, :black))
    F = (a,b) -> begin
        ga,gb=g(a),g(b)
        ga * gb < 0 && return nothing
        ga >= 0 && return true
        return false
    end
    find_colors(F, xs, colors)
end

# F(a,b) returns true, false, or nothing
function find_colors(F, xs, colors=(:red, :blue, :black))
    n = length(xs)
    cols = repeat([colors[1]], n)
    for i in 1:n-1
        a,b = xs[i], xs[i+1]
        val = F(a,b)
        if val == nothing
            cols[i] = colors[3]
        elseif val
            cols[i] = colors[1]
        else
            cols[i] = colors[2]
        end
    end
    cols[end] = cols[end-1]
    cols
end

# some plotting utilities
"""
    rangeclamp(f, hi=20, lo=-hi; replacement=NaN)

Modify `f` so that values of `f(x)` outside of `[lo,hi]` are replaced by `replacement`.

Examples
```
f(x) = 1/x
plot(rangeclamp(f), -1, 1)
plot(rangeclamp(f, 10), -1, 1) # no `abs(y)` values exceeding 10
```
"""
rangeclamp(f, hi=20, lo=-hi; replacement=NaN) = x -> lo < f(x) < hi ? f(x) : replacement


# p13.jl - solve linear BVP u_xx = exp(4x), u(-1)=u(1)=0

N = 16;
(D,x) = cheb(N);
D2 = D^2;
D2 = D2[2:N,2:N];                   # boundary conditions
f = @. exp(4*x[2:N]);
u = D2\f;                           # Poisson eq. solved here
u = [0;u;0];
clf(); axes([.1,.4,.8,.5]);
plot(x,u,".",markersize=6);
xx = -1:.01:1;
uu = polyval(polyfit(x,u),xx);      # interpolate grid data
plot(xx,uu);
grid(true);
exact = @. ( exp(4*xx) - sinh(4)*xx - cosh(4) )/16;
title("max err = $(signif(norm(uu-exact,Inf),4))",fontsize=12);


# This file calculates scattering from the gravity Helmholtz equation via the adaptive spectral method.
# Δu + (E+x_2)u = 0,
# u|Γ = 0,
# u^i = Φ(⋅,⋅,E)|Γ
# u = u^i + u^s.
# The normal derivative ∂u/∂n of the entire wave is calculated on the sound-soft boundaries.
# The scattered wave is calculated by convolving ∂u/∂n with the fundamental solution.
# Then, the total wave is obtained by summing the incident and the scattered waves.

using ApproxFun, SingularIntegralEquations
include("Scatteraux.jl")

E = 20.
ω = 2π
ui = (x,y) ->  lhelmfs(complex(x,y),-5.0im,E)

# The gravity Helmholtz Green's function.
g3 = (x,y) ->  lhelmfs(x,y,E)
r = (x,y) ->  lhelm_riemann(x,y,E)

dom = ∪(Segment.([-10.0-3.0im,5.0,-2+5im],[-5.0+0.0im,10.0-3im,2+5im]))
sp = Space(dom)
cwsp = CauchyWeight(sp⊗sp,0)
uiΓ,⨍ = Fun(t->ui(real(t),imag(t)),sp),DefiniteLineIntegral(dom)

@time G = GreensFun(g3,cwsp;method=:unsplit)

@time ∂u∂n = ⨍[G]\uiΓ
println("The length of ∂u∂n is: ",ncoefficients(∂u∂n))
us = (x,y) ->  -linesum(g3,∂u∂n,complex(x,y))


# Local Tangent Space Alignment (LTSA)
# ---------------------------
# Principal Manifolds and Nonlinear Dimension Reduction via Local Tangent Space Alignment,
# Zhang, Zhenyue; Hongyuan Zha (2004),  SIAM Journal on Scientific Computing 26 (1): 313–338.
# doi:10.1137/s1064827502419154.

"""
    LTSA{NN <: AbstractNearestNeighbors, T <: Real} <: AbstractDimensionalityReduction

The `LTSA` type represents a local tangent space alignment model constructed for `T` type data with a help of the `NN` nearest neighbor algorithm.
"""
struct LTSA{NN <: AbstractNearestNeighbors, T <: Real} <: AbstractDimensionalityReduction
    λ::AbstractVector{T}
    proj::Projection{T}
    nearestneighbors::NN
    component::AbstractVector{Int}
end

## properties
outdim(R::LTSA) = size(R.proj, 1)
eigvals(R::LTSA) = R.λ
neighbors(R::LTSA) = R.nearestneighbors.k
vertices(R::LTSA) = R.component

## show
summary(io::IO, R::LTSA) = print(io, "LTSA(outdim = $(outdim(R)), neighbors = $(neighbors(R)))")

## interface functions
"""
    fit(LTSA, data; k=12, maxoutdim=2, nntype=BruteForce)

Fit a local tangent space alignment model to `data`.

# Arguments
* `data`: a matrix of observations. Each column of `data` is an observation.

# Keyword arguments
* `k`: a number of nearest neighbors for construction of local subspace representation
* `maxoutdim`: a dimension of the reduced space.
* `nntype`: a nearest neighbor construction class (derived from `AbstractNearestNeighbors`)

# Examples
```julia
M = fit(LTSA, rand(3,100)) # construct LTSA model
R = transform(M)           # perform dimensionality reduction
```
"""
function fit(::Type{LTSA}, X::AbstractMatrix{T};
        k::Int=12, maxoutdim::Int=2, ɛ::Real=1.0, nntype=BruteForce) where {T<:Real}
    # Construct NN graph
    NN = fit(nntype, X, k)
    D, E = knn(NN, X)
    G, C = largest_component(SimpleWeightedGraph(adjmat(D,E)))
    XX = @view X[:, C]
    n = length(C)

    S = ones(k)./sqrt(k)
    B = spzeros(T, n,n)
    for i=1:n
        II = @view E[:,i]
        VX = view(XX, :, II)

        # re-center points in neighborhood
        μ = mean(VX, dims=2)
        δ_x = VX .- μ

        # Compute orthogonal basis H of θ'
        θ_t = svd(δ_x).V[:,1:maxoutdim]

        # Construct alignment matrix
        G = hcat(S, θ_t)
        B[II, II] .+= Diagonal(fill(one(T), k)) .- G*transpose(G)
    end

    # Align global coordinates
    λ, V = decompose(B, maxoutdim)
    return LTSA{nntype, T}(λ, transpose(V), NN, C)
end

"""
    transform(R::LTSA)

Transforms the data fitted to the local tangent space alignment model `R` into a reduced space representation.
"""
transform(R::LTSA) = R.proj


"""
    swapping_even_odd_bits(n::Int64))

For a given integer value, all odd bits will be swapped with even bits.

# Arguments
- `n::Int64`: Integer value of the number, which has to be swapped.

# Examples
```julia-repl
julia> import ClassicAlgorithmsCollections
julia> ClassicAlgorithmsCollections.function swapping_even_odd_bits(23)
43
```

# Notes
---
The definition of the even and odd bits:
```julia-repl
even_bits = n & 0xAAAAAAAA
odd_bits = n & 0x55555555
```
"""
function swapping_even_odd_bits(n::Int64)
    # Defining even and odd bits
    even_bits = n & 0xAAAAAAAA
    odd_bits = n & 0x55555555

    # Shifting the even and odd bits
    even_bits >>= 1
    odd_bits <<= 1
    return (even_bits | odd_bits)
end


module Percolation

export clustgyration, hklabel, hkmaxclust!, genclust, clustfractal

using DataStructures

"""
Finds the radius of gyration for the largest closed cluster in the percolation grid

Percolation is assumed to occur from top to bottom. Uses the Hoshen-Kopelman algorithm.

# Argument

`grid::Matrix{Bool}`: Represents the grid cells; True if porous (part of the path;
    can pass through) and False if blocked (can't pass through; no path passes through)

# Returns

The gyration radius of the largest closed cluster (the largest cluster that doesn't connect
the top and bottom of the grid)
"""
function clustgyration(grid::BitMatrix)
    labels, list, clusters = hklabel(grid)
    if length(clusters) == 0
        return 0
    end

    maxlabel = hkmaxclust!(labels, list, clusters)
    if maxlabel == 0
        return 0
    end

    maxclust = Tuple.(findall(x -> x != 0 && hkfind(list, x) == maxlabel, labels))
    mass = length(maxclust)
    if mass == 0
        return 0
    end

    centerofmass = reduce(.+, maxclust) ./ mass
    return √(sum(pos -> reduce(.+, (pos .- centerofmass).^2),
        maxclust) / mass)
end

"""
Uses the Hoshen-Kopelman algorithm to label the grid

returns the labels, label mapping list, and cluster lengths
"""
@inline function hklabel(grid::BitMatrix)
    height, width = size(grid)
    labels = zeros(Integer, height, width)
    list = Vector{Integer}()
    clusters = Vector{Integer}()

    # first row, special case because of top boundary
    hkfirstrow!(grid, labels, list, clusters)

    for row in 2:height
        # first column, special case because of left boundary
        if grid[row, 1]
            hkfirstcol!(labels, list, clusters, row)
        end

        # "middle" points, with appropriate neighbors
        for col in 2:width
            if grid[row, col]
                hkiter!(labels, list, clusters, row, col)
            end
        end
    end

    return labels, list, clusters
end

"""
Helper function for `hkpercolate`; Adds new label.
"""
@inline function hknew!(labels::Matrix{<:Integer}, list::Vector{<:Integer},
        clusters::Vector{<:Integer}, row::Integer, col::Integer)
    labels[row, col] = length(list) + 1
    push!(list, labels[row, col])
    push!(clusters, 0)
end

"""
Helper function for `hkpercolate`; Applies the Hoshen-Kopelman algorithm to the first row.
"""
@inline function hkfirstrow!(grid::BitMatrix, labels::Matrix{<:Integer},
        list::Vector{<:Integer}, clusters::Vector{<:Integer})
    width = size(labels)[2]
    if grid[1, 1]
        hknew!(labels, list, clusters, 1, 1)
        clusters[labels[1, 1]] += 1
    end
    for col in 2:width
        if grid[1, col]
            if labels[1, col - 1] != 0
                labels[1, col] = hkfind!(list, labels[1, col - 1])
            else
                hknew!(labels, list, clusters, 1, col)
            end
            clusters[labels[1, col]] += 1
        end
    end
end

"""
Helper function for `hkpercolate`; Applies the Hoshen-Kopelman algorithm to the first column.
"""
@inline function hkfirstcol!(labels::Matrix{<:Integer}, list::Vector{<:Integer},
        clusters::Vector{<:Integer}, row::Integer)
    if labels[row - 1, 1] != 0
        labels[row, 1] = hkfind!(list, labels[row - 1, 1])
    else
        hknew!(labels, list, clusters, row, 1)
    end
    clusters[labels[row, 1]] += 1
end

"""
Helper function for `hkpercolate`; Applies one iteration of the Hoshen-Kopelman algorithm.
"""
@inline function hkiter!(labels::Matrix{<:Integer}, list::Vector{<:Integer},
        clusters::Vector{<:Integer}, row::Integer, col::Integer)
    above, left = labels[row - 1, col], labels[row, col - 1]
    if above == 0
        if left == 0
            hknew!(labels, list, clusters, row, col)
        else
            labels[row, col] = hkfind!(list, left)
        end
    elseif left == 0
        labels[row, col] = hkfind!(list, above)

    else
        labels[row, col] = hkunion!(list, clusters, above, left)
    end
    clusters[labels[row, col]] += 1
end

"""
Union function of the union-find algorithm used by the Hoshen-Kopelman algorithm
"""
@inline function hkunion!(list::Vector{<:Integer}, clusters::Vector{<:Integer},
        x::Integer, y::Integer)
    root = hkfind!(list, x)
    dest = hkfind!(list, y)
    if dest != root
        list[dest] = root
        clusters[root] += clusters[dest]
        clusters[dest] = 0
    end
    return root
end

"""
Find function of the union-find algorithm used by the Hoshen-Kopelman algorithm
"""
@inline function hkfind!(list::Vector{<:Integer}, x::Integer)
    y = x;
    while list[y] != y
        y = list[y];
    end
    while list[x] != x
        z = list[x];
        list[x] = y;
        x = z;
    end
    return y;
end

"""
Same as hkfind!, but doesn't update labels
"""
@inline function hkfind(list::Vector{<:Integer}, x::Integer)
    while x != list[x]
        x = list[x];
    end
    return x
end

"""
Finds the label of the largest closed cluster (a.k.a. finite cluster).

Also, sets the size of open clusters (a.k.a. infinite cluster) to zero, hence the !
"""
@inline function hkmaxclust!(labels::Matrix{Integer}, list::Vector{Integer},
        clusters::Vector{Integer})
    for root in labels[1, :]
        if root > 0
            for dest in labels[end, :]
                if dest > 0
                    toplabel = hkfind(list, root)
                    bottomlabel = hkfind(list, dest)
                    if toplabel == bottomlabel
                        clusters[toplabel] = 0
                    end
                end
            end
        end
    end

    if maximum(clusters) > 0
        return argmax(clusters)
    else
        return 0
    end
end

"""
Uses a depth-first search algorithm to generate a cluster with the given site probability

# Arguments

- `probability`: the probability a site is generated when growing the cluster
- `maxsize::Integer` (optional, defaults to 10000): the maximum size of the cluster
    (breaks the loop after the cluster is grown to this size)

# Returns

an integer matrix (8-bit) with 0 where no site exists, 1 where a cell is generated, and 2
where the growth is stopped
"""
function genclust(probability::Float64; maxsize::Integer=10000, crop::Bool=true)
    cluster = zeros(Int8, maxsize, maxsize)
    # create "fences" to stop the cluster form reaching the bounds of the matrix
    cluster[1, :] .= 2
    cluster[:, 1] .= 2
    cluster[end, :] .= 2
    cluster[:, end] .= 2

    stack = Stack{Tuple{Integer, Integer}}()
    seed = (maxsize÷2, maxsize÷2)
    push!(stack, seed)
    cluster[seed...] = 1

    clustdfs!(stack, cluster, probability, maxsize)

    if crop
        return cropclust!(cluster)
    else
        return cluster
    end
end

"""
The depth-first search algorithm used for `genclust`
"""
@inline function clustdfs!(stack::Stack, cluster::Matrix{<:Integer},
        probability::Float64, maxsize::Integer)
    size = 1
    while !isempty(stack) && size < maxsize
        site = pop!(stack)
        for step in ((-1, 0), (0, -1), (0, 1), (1, 0))
            neighbor = site .+ step
            if cluster[neighbor...] == 0
                if rand() <= probability
                    cluster[neighbor...] = 1
                    push!(stack, neighbor)
                    size += 1
                else
                    cluster[neighbor...] = 2
                end
            end
        end
    end
end

"""
finds the mass and radius of gyration of the cluster from the output of `genclust`
"""
function clustfractal(cluster::Matrix{Int8})
    clust = Tuple.(findall(x -> x == 1, cluster))
    mass = length(clust)
    if mass == 0
        return 0
    end

    centerofmass = reduce(.+, clust) ./ mass
    return mass, √(sum(pos -> reduce(.+, (pos .- centerofmass).^2), clust) / mass)
end

"""
crop the `genclust` function output to remove excess zeros
"""
@inline function cropclust!(cluster::Matrix{<:Integer})
    height, width = size(cluster)
    top = bottom = height ÷ 2
    left = right = width ÷ 2

    while 1 in cluster[top + 1, 2:end-1] || 2 in cluster[top + 1, 2:end-1]
        top += 1
    end
    while 1 in cluster[bottom - 1, 2:end-1] || 2 in cluster[bottom - 1, 2:end-1]
        bottom -= 1
    end
    while 1 in cluster[2:end-1, left - 1] || 2 in cluster[2:end-1, left - 1]
        left -= 1
    end
    while 1 in cluster[2:end-1, right + 1] || 2 in cluster[2:end-1, right + 1]
        right += 1
    end

    return cluster[bottom:top, left:right]
end

end