using Symbolics
using SatisfiabilityInterface

vars = @variables x, y, z, w

constraints = [
    x ∈ 1:2
    y ∈ 2:5
    z ∈ 3:7
    z == x * y
]

prob = DiscreteCSP(constraints)

solve(prob)

all_solutions(prob)

###

vars = @variables x, y, z, w

constraints = [
    x ∈ 1:2
    y ∈ 2:5
    z ∈ 3:7
    w ∈ 4:6
    z == x + y
    w == z + x
]

prob = DiscreteCSP(constraints)
solve(prob)
all_solutions(prob)

###

constraints = [
    x ∈ 1:2
    y ∈ 2:5
    z ∈ 3:7
    z == 2x
]

all_solutions(DiscreteCSP(constraints))

###

vars = @variables x, y, z

# constraints = [
#     x ∈ 1:4
#     y ∈ 2:5
#     z == x + y
#     z ≤ 4
# ]

### Subset sum Problem

# https://en.wikipedia.org/wiki/Subset_sum_problem

# Given a set of integers like {−7, −3, −2, 5, 8},
# do a subset sum to 0?

s = [-7, -3, -2, 5, 8]
n = length(s)
x = [Num(Variable(:x, i)) for i in 1:n]

constraints = [
    [x[i] ∈ (0:1) for i in 1:n]
    sum(x[i] * s[i] for i in 1:n) == 0
    sum(x) ≠ 0   # at least one is not 0
    ]

prob = DiscreteCSP(constraints)

all_solutions(prob)

#####

@variables x, y, z

constraints = [
    x ∈ 1:2
    y ∈ 2:5
    z ∈ 3:7
    z == 2x
]

prob = DiscreteCSP(constraints)
all_solutions(prob)

## Example from https://github.com/newptcai/BeeEncoder.jl

function make_vector(name::Symbol, range)
    return Num.(Variable.(name, range))
end

@variables x, y, z, w
b = make_vector(:b, 1:4)
# @variables b[1:4]

¬(x::Num) = 1 - x   # assumes a Boolean variable

constraints = [
    x ∈ 0:5
    y ∈ -4:9
    z ∈ -5:10
    w ∈ 0:10

    [bb ∈ 0:1 for bb ∈ b]

    x + y == z

    b[1] == ¬b[2]
    b[2] == 1

    ¬b[1] + b[2] + ¬b[3] + b[4] == w
]

prob = DiscreteCSP(constraints)
solve(prob)

# @time length(all_solutions(prob))   # 1184 solutions in 38 seconds

###

constraints = [
    x ∈ -5:5
    y ∈ -4:9
    x - y^2 == 3
    # sum(b) == w

]

prob = DiscreteCSP(constraints)
all_solutions(prob)

###

@variables x, y

constraints = [
    x ∈ -5:5
    y ∈ -4:9
    10x + y == 9
    # sum(b) == w
]

prob = DiscreteCSP(constraints)
all_solutions(prob)


using Miletus.DayCounts
using Dates

@test 0.497724380567 ≈ yearfraction(ISDAActualActual(), Date(2003,November,1), Date(2004, May, 1))
@test 0.50           ≈ yearfraction(ISMAActualActual(), Date(2003,November,1), Date(2004, May, 1))
@test 0.497267759563 ≈ yearfraction(AFBActualActual(), Date(2003,November,1), Date(2004, May, 1))

@test 0.410958904110 ≈ yearfraction(ISDAActualActual(), Date(1999,February,1), Date(1999, July, 1))
@test 0.410958904110 ≈ yearfraction(ISMAActualActual(), Date(1999,February,1), Date(1999, July, 1), Date(1998, July, 1), Date(1999, July, 1))
@test 0.410958904110 ≈ yearfraction(AFBActualActual(), Date(1999,February,1), Date(1999, July, 1))

@test 1.001377348600 ≈ yearfraction(ISDAActualActual(), Date(1999,July,1), Date(2000, July, 1))
@test 1.0            ≈ yearfraction(ISMAActualActual(), Date(1999,July,1), Date(2000, July, 1), Date(1999, July, 1), Date(2000, July, 1))
@test 1.0            ≈ yearfraction(AFBActualActual(), Date(1999,July,1), Date(2000, July, 1))

@test 0.915068493151 ≈ yearfraction(ISDAActualActual(), Date(2002,August,15), Date(2003, July, 15))
@test 0.915760869565 ≈ yearfraction(ISMAActualActual(), Date(2002,August,15), Date(2003, July, 15), Date(2003, January, 15), Date(2003, July, 15))
@test 0.915068493151 ≈ yearfraction(AFBActualActual(), Date(2002,August,15), Date(2003, July, 15))

@test 0.504004790778 ≈ yearfraction(ISDAActualActual(), Date(2003,July,15), Date(2004, January, 15))
@test 0.5            ≈ yearfraction(ISMAActualActual(), Date(2003,July,15), Date(2004, January, 15), Date(2003, July, 15), Date(2004, January, 15))
@test 0.504109589041 ≈ yearfraction(AFBActualActual(), Date(2003,July,15), Date(2004, January, 15))

@test 0.503892506924 ≈ yearfraction(ISDAActualActual(), Date(1999,July,30), Date(2000, January, 30))
@test 0.5            ≈ yearfraction(ISMAActualActual(), Date(1999,July,30), Date(2000, January, 30), Date(1999,July,30), Date(2000, January, 30))
@test 0.504109589041 ≈ yearfraction(AFBActualActual(), Date(1999,July,30), Date(2000, January, 30))

@test 0.415300546448 ≈ yearfraction(ISDAActualActual(), Date(2000,January,30), Date(2000, June, 30))
@test 0.417582417582 ≈ yearfraction(ISMAActualActual(), Date(2000,January,30), Date(2000, June, 30), Date(2000,January,30), Date(2000, July, 30))
@test 0.415300546448 ≈ yearfraction(AFBActualActual(), Date(2000,January,30), Date(2000, June, 30))

using BusinessDays
c=BusinessDays.UnitedKingdom()
adjust(ModPreceding(), c, Date(2016,5,2)) == Date(2016, 5, 3)


#=
It is possible to show that the square root of two can be expressed as an infinite continued fraction.

√ 2 = 1 + 1/(2 + 1/(2 + 1/(2 + ... ))) = 1.414213...

By expanding this for the first four iterations, we get:

1 + 1/2 = 3/2 = 1.5
1 + 1/(2 + 1/2) = 7/5 = 1.4
1 + 1/(2 + 1/(2 + 1/2)) = 17/12 = 1.41666...
1 + 1/(2 + 1/(2 + 1/(2 + 1/2))) = 41/29 = 1.41379...

The next three expansions are 99/70, 239/169, and 577/408, but the eighth expansion, 1393/985, is the first example where the number of digits in the numerator exceeds the number of digits in the denominator.

In the first one-thousand expansions, how many fractions contain a numerator with more digits than denominator?
=#
function calc()
  p = BigInt(1)
  q = BigInt(1)
  count = 1
  total = 0
  while count <= 1000
    np = p + 2q
    nq = p + q
    p,q = np,nq
    if length(string(p)) > length(string(q))
      total = total + 1
    end
    count = count + 1
  end
  total
end
@time println(calc())


using Printf
using Plots
using Statistics
using LinearAlgebra
using MDBM
using PyPlot
pygui(true);
using Interpolations

#interpolate
function it(A)
    matrdim=size(A,2)-1
    scaleitp=real(A[:,1])
    knots=(scaleitp,)
    ARe=real(A)
    AIm=imag(A)

    imRe=[interpolate(knots, ARe[:,2],Gridded(Linear()))]
    for j = 3:matrdim+1
        imRetemp=[interpolate(knots, real(ARe[:,j]), Gridded(Linear()))]
        imRe=vcat(imRe,imRetemp)
    end
    imIm=[interpolate(knots, AIm[:,2],Gridded(Linear()))]
    for j = 3:matrdim+1
        imImtemp=[interpolate(knots, real(AIm[:,j]), Gridded(Linear()))]
        imIm=vcat(imIm,imImtemp)
    end
    return(hcat(imRe,imIm))
end

B=transpose([0 0.5 1 1.5 2; 3*im 2+4*im -5+1*im 6*im 5*im-4; 7*im -8+3*im -5+2*im 7*im 3*im-2])
real(B)
typeof(real(B))

iterB=it(B)

interpolate((real(B[:,1]),), real(B[:,2]),Gridded(Linear()))
#substitution in interpolating function
function sub(it,t)
    subdim=size(it,1)
    out=zeros(ComplexF64,subdim)
    for j = 1:subdim
        out[j]=it[j,1](t)+it[j,2](t)*im
    end
    return(out)
end

sub(iterB,1.25)


# nuclear Norm (times a constant)

export NuclearNorm

"""
**Nuclear norm**

    NuclearNorm(λ=1.0)

Returns the function
```math
f(X) = \\|X\\|_* = λ ∑_i σ_i(X),
```
where `λ` is a positive parameter and ``σ_i(X)`` is ``i``-th singular value of matrix ``X``.
"""
struct NuclearNorm{R <: Real} <: ProximableFunction
  lambda::R
  function NuclearNorm{R}(lambda::R) where {R <: Real}
    if lambda < 0
      error("parameter λ must be nonnegative")
    else
      new(lambda)
    end
  end
end

is_convex(f::NuclearNorm) = true

NuclearNorm(lambda::R=1.0) where {R <: Real} = NuclearNorm{R}(lambda)

function (f::NuclearNorm{R})(X::AbstractMatrix{T}) where {R <: Real, T <: Union{R, Complex{R}}}
  F = svd(X);
  return f.lambda * sum(F.S);
end

function prox!(Y::AbstractMatrix{T}, f::NuclearNorm{R}, X::AbstractMatrix{T}, gamma::R=one(R)) where {R <: Real, T <: Union{R, Complex{R}}}
  F = svd(X)
  S_thresh = max.(zero(R), F.S .- f.lambda*gamma)
  rankY = findfirst(S_thresh .== zero(R))
  if rankY === nothing
    rankY = minimum(size(X))
  end
  Vt_thresh = view(F.Vt, 1:rankY, :)
  U_thresh = view(F.U, :, 1:rankY)
  # TODO: the order of the following matrix products should depend on the shape of x
  M = S_thresh[1:rankY] .* Vt_thresh
  mul!(Y, U_thresh, M)
  return f.lambda * sum(S_thresh);
end

fun_name(f::NuclearNorm) = "nuclear norm"
fun_dom(f::NuclearNorm) = "AbstractArray{Real,2}, AbstractArray{Complex,2}"
fun_expr(f::NuclearNorm) = "X ↦ λ∑σ_i(X)"
fun_params(f::NuclearNorm) = "λ = $(f.lambda)"

function prox_naive(f::NuclearNorm, X::AbstractMatrix{T}, gamma=1.0) where T
  F = svd(X)
  S = max.(0, F.S .- f.lambda*gamma)
  Y = F.U * (Diagonal(S) * F.Vt)
  return Y, f.lambda * sum(S)
end


using Plots
using Printf

"""
    untildivergence(f, x, maxiter=50)

Iterate `f` on `x` until its module is greater than 2 or `maxiter` iterations have been done.

Returns the number of iterations.
"""
function untildivergence(f, x, maxiter = 50)
    i = 0
    while abs(x) < 2 && (i += 1) <= maxiter
        x = f(x)
    end
    return i
end

"""
    julia_set_wh(f; kwargs)

Approximate the Julia set of `f` in a grid using the escape time algorithm.

The grid is defined by its bounds `xmin`, `xmax`, `ymin` and `ymax` as well as its width `w` and height `h` in pixels. The optional parameter `maxiter` is passed on to `untildivergence`.
"""
function julia_set_wh(f; xmin::Real=-2, xmax::Real=2, ymin::Real=-2, ymax::Real=2, w::Int=1000, h::Int=1000, maxiter::Int=50)
    img = Array{Int,2}(undef, h, w)
    x = LinRange(xmin, xmax, w)
    y = LinRange(ymin, ymax, h)
    # creating labels for use with `heatmap`
    xticks = LinRange(1,w,11)
    yticks = LinRange(1,h,11)
    xlabels = map(z->@sprintf("%.2f",z), LinRange(xmin,xmax,11))
    ylabels = map(z->@sprintf("%.2f",z), LinRange(ymin,ymax,11))

    for i in 1:w, j in 1:h
        img[j,i] = untildivergence(f, x[i]+im*y[j], maxiter)
    end
    return img, xticks, xlabels, yticks, ylabels
end

"""
    julia_set(f, cellsize; kwargs)

Approximate the Julia of `f` set in a grid of cells of length `cellsize`.

The grid is defined by the same parameters as in `julia_set_wh`, except its width and height in pixels are computed with `cellsize`.
"""
function julia_set(f, cellsize::Real; xmin=-2, xmax=2, ymin=-2, ymax=2, maxiter=50)
    w = Integer(round(abs(xmax-xmin)/cellsize))
    h = Integer(round(abs(ymax-ymin)/cellsize))
    return julia_set_wh(f, xmin=xmin, xmax=xmax, ymin=ymin, ymax=ymax, w=w, h=h, maxiter=maxiter)
end

"""
    mandelbrot_set_wh(f; kwargs)

Approximate the Mandelbrot set of `f` in a grid using the escape time algorithm.

The grid is defined by its bounds `xmin`, `xmax`, `ymin` and `ymax` as well as its width `w` and height `h` in pixels. The optional parameter `maxiter` is passed on to `untildivergence`.
"""
function mandelbrot_set_wh(f; xmin::Real=-2, xmax::Real=2, ymin::Real=-2, ymax::Real=2, w::Integer=1000, h::Integer=1000, maxiter::Integer=50)
    img = Array{Int,2}(undef, h, w)
    x = LinRange(xmin, xmax, w)
    y = LinRange(ymin, ymax, h)
    # creating labels for use with `heatmap`
    xticks = LinRange(1,w,11)
    yticks = LinRange(1,h,11)
    xlabels = map(z->@sprintf("%.2f",z), LinRange(xmin,xmax,11))
    ylabels = map(z->@sprintf("%.2f",z), LinRange(ymin,ymax,11))

    for i in 1:w, j in 1:h
        img[j,i] = untildivergence(f(x[i]+im*y[j]), 0, maxiter)
    end
    return img, xticks, xlabels, yticks, ylabels
end

"""
    mandelbrot_set(f, cellsize; kwargs)

Approximate the Mandelbrot of `f` set in a grid of cells of length `cellsize`.

The grid is defined by the same parameters as in `julia_set_wh`, except its width and height in pixels are computed with `cellsize`.
"""
function mandelbrot_set(f, cellsize; xmin=-2, xmax=2, ymin=-2, ymax=2, maxiter=50)
    w = Integer(round(abs(xmax-xmin)/cellsize))
    h = Integer(round(abs(ymax-ymin)/cellsize))
    return mandelbrot_set_wh(f, xmin=xmin, xmax=xmax, ymin=ymin, ymax=ymax, w=w, h=h, maxiter=maxiter)
end

# the traditional julia/mandelbrot function
f(c) = z -> z^2 + c

img, xticks, xlabels, yticks, ylabels = mandelbrot_set(f, 0.00005, xmin=-2.01, xmax = -1.8, ymin = -0.1, ymax = 0.1)

img, xticks, xlabels, yticks, ylabels = julia_set(f(im), 0.0005, xmin=-2, xmax=2, ymin=-2,ymax=2)

heatmap(img,
        xticks = (xticks, xlabels),
        yticks = (yticks, ylabels),
        aspect_ratio=:equal
        )

# alternative, for better looking images
using Images, Colors
g = Gray.(img/51)

g = Gray.(.!(img .>= 50))

mset = img .>= 50

#save("julia_minus2.png", g)

include("FractalDimensions.jl")

B = boundary(mset)
Gray.(B)
points = array_to_points(B)

d, N, ε = box_counting_dimension(points, 10)

C, r = correlation_dimension(points)

loglog_regression(C,r)


### Gali's Book Chapter 3 model
### We use this as a consistency check of gensys implementation
### By Daniel Coutinho and Gilberto Boaretto

#### Parameters from Gali's book example, chapter 3, page 52, first edition

bet = 0.99
sig = 1
phi = 1
alfa = 1/3
epsilon = 6
theta = 2/3
phi_pi = 1.5
phi_y = 0.5/4
rho_v = 0.5

true_pars = (alfa = alfa, bet = bet, epsilon = epsilon, theta = theta, sig = sig, s2 = 1, phi = phi, phi_pi = phi_pi, phi_y = phi_y, rho_v = rho_v)

THETA = (1-alfa)/(1-alfa+alfa*epsilon)
lamb = (1-theta)*(1-bet*theta)/theta*THETA
kappa = lamb*(sig+(phi+alfa)/(1-alfa))

## Gali's give analytical expressions for the effects of a shock in the monetary policy over inflation and output gap,chapter 3, page 51, first edition

LAMBDA_v = ((1-bet*rho_v)*(sig*(1-rho_v)+phi_y)+kappa*(phi_pi-rho_v))^(-1)

y_irf(v) = -(1-bet*rho_v)*LAMBDA_v*v
pi_irf(v) = -kappa*LAMBDA_v*v
r_irf(v) = sig*(1-rho_v)*(1-bet*rho_v)*LAMBDA_v*v
i_irf(v) = (sig*(1-rho_v)*(1-bet*rho_v) - rho_v*kappa)*LAMBDA_v*v
v_irf(v,uu) = rho_v*v + uu

## Test drive: lets see if this gives the same IRFs that we get in Gali's book
## This is just for debugging, so comment it out when using it to check gensys consistency

#using Plots

irfs_true = zeros(15,5)

irfs_true[1,1] = 0

uu = 0.25

for t in 2:15
    irfs_true[t,4] = v_irf(irfs_true[t-1,4],uu)
    irfs_true[t,1] = pi_irf(irfs_true[t,4])
    irfs_true[t,2] = y_irf(irfs_true[t,4])
    irfs_true[t,3] = i_irf(irfs_true[t,4])
    irfs_true[t,5] = r_irf(irfs_true[t,4])
    global uu = 0
end

#plot(irfs[2:13,1], label = "v")
#plot(irfs[2:13,2]*4, label = "pi")
#plot(irfs[2:13,3], label = "y")
#plot(4*irfs[2:13,4], label = "i")
#plot(4*irfs[2:13,5], label = "r")


# AAA algorithm from the paper "The AAA Alorithm for Rational Approximation"
# by Y. Nakatsukasa, O. Sete, and L.N. Trefethen, SIAM Journal on Scientific Computing
# 2018

struct AAAapprox{T <: AbstractArray} <: BRInterp
    x::T
    f::T
    w::T
    errvec::T
end

(a::AAAapprox)(zz) = reval(zz, a.x, a.f, a.w)

# Handle function inputs as well
function aaa(Z::AbstractArray{T,1}, F::S;  tol=1e-13, mmax=100, verbose=false, clean=true) where {T, S<:Function}
    aaa(Z, F.(Z), tol=tol, mmax=mmax, verbose=verbose, clean=clean)
end

"""aaa  rational approximation of data F on set Z
        r = aaa(Z, F; tol, mmax, verbose, clean)

 Input: Z = vector of sample points
        F = vector of data values at the points in Z
        tol = relative tolerance tol, set to 1e-13 if omitted
        mmax: max type is (mmax-1, mmax-1), set to 100 if omitted
        verbose: print info while calculating default = false
        clean: detect and remove Froissart doublets default = true

 Output: r = an AAA approximant as a callable struct with fields
         z, f, w = vectors of support pts, function values, weights
         errvec = vector of errors at each step

 NOTE: changes from matlab version:
         switched order of Z and F in function signature
         added verbose and clean boolean flags
         pol, res, zer = vectors of poles, residues, zeros are now only calculated on demand
         by calling prz(z::AAAapprox)
"""
function aaa(Z::AbstractArray{T,1}, F::AbstractArray{S,1}; tol=1e-13, mmax=100,
             verbose=false, clean=true) where {S, T}
    # filter out any NaN's or Inf's in the input
    keep = isfinite.(F)
    F = F[keep]
    Z = Z[keep]

    M = length(Z)                    # number of sample points
    mmax = min(M, mmax)              # max number of support points

    reltol = tol*norm(F, Inf)
    SF = spdiagm(M, M, 0 => F)       # left scaling matrix

    F, Z = promote(F, Z)
    P = promote_type(S, T)

    J = [1:M;]
    z = P[]                          # support points
    f = P[]                          # function values at support points
    C = P[]
    w = P[]

    errvec = P[]
    R = F .- mean(F)
    @inbounds for m = 1:mmax
        j = argmax(abs.(F .- R))               # select next support point
        push!(z, Z[j])
        push!(f, F[j])
        deleteat!(J, findfirst(isequal(j), J))   # update index vector

        # next column of Cauchy matrix
        C = isempty(C) ? reshape((1 ./ (Z .- Z[j])), (M,1)) : [C (1 ./ (Z .- Z[j]))]

        Sf = diagm(f)                         # right scaling matrix
        A = SF * C - C * Sf                   # Loewner matrix
        G = svd(A[J, :])

        # A[J, :] might not have full rank, so svd can return fewer than m columns
        #w = G.V[:, m]
        w = G.V[:, end]                       # weight vector = min sing vector

        N = C * (w .* f)                      # numerator
        D = C * w
        R .= F
        R[J] .= N[J] ./ D[J]                  # rational approximation

        err = norm(F - R, Inf)
        verbose && println("Iteration: ", m, "  err: ", err)
        errvec = [errvec; err]                # max error at sample points
        err <= reltol && break                # stop if converged
    end
    r = AAAapprox(z, f, w, errvec)

    # remove Frois. doublets if desired.  We do this in place
    if clean
        pol, res, zer = prz(r)            # poles, residues, and zeros
        ii = findall(abs.(res) .< 1e-13)  # find negligible residues
        length(ii) != 0 && cleanup!(r, pol, res, zer, Z, F)
    end
    return r
end

function prz(r::AAAapprox)
    z, f, w = r.x, r.f, r.w
    m = length(w)
    B = diagm(ones(m+1))
    B[1, 1] = 0.0
    E = [0.0  transpose(w); ones(m) diagm(z)]
    pol, _ = eigen(E, B)
    pol = pol[isfinite.(pol)]
    dz = 1e-5 * exp.(2im*pi*[1:4;]/4)

    # residues
    res = r(pol .+ transpose(dz)) * dz ./ 4

    E = [0 transpose(w .* f); ones(m) diagm(z)]
    zer, _ = eigen(E, B)
    zer = zer[isfinite.(zer)]
    pol, res, zer
end

function reval(zz, z, f, w)
    # evaluate r at zz
    zv = size(zz) == () ? [zz] : vec(zz)
    CC = 1.0 ./ (zv .- transpose(z))           # Cauchy matrix
    r = (CC * (w .* f)) ./ (CC * w)            # AAA approx as vector
    r[isinf.(zv)] .= sum(f .* w) ./ sum(w)

    ii = findall(isnan.(r))                    # find values NaN = Inf/Inf if any
    @inbounds for j in ii
        if !isnan(zv[j]) && ((v = findfirst(isequal(zv[j]), z)) !== nothing)
            r[j] = f[v]  # force interpolation there
        end
    end
    r = size(zz) == () ? r[1] : reshape(r, size(zz))             # the AAA approximation
end

# Only calculate the updated z, f, and w
function cleanup!(r, pol, res, zer, Z, F)
    z, f, w = r.x, r.f, r.w
    m = length(z)
    M = length(Z)
    ii = findall(abs.(res) .< 1e-13)  # find negligible residues
    ni = length(ii)
    ni == 0 && return
    println("$ni Froissart doublets. Number of residues = ", length(res))

    @inbounds for j = 1:ni
        azp = abs.(z .- pol[ii[j]] )
        jj = findall(isequal(minimum(azp)), azp)
        deleteat!(z, jj)    # remove nearest support points
        deleteat!(f, jj)
    end

    @inbounds for j = 1:length(z)
        jj = findall(isequal(z[j]), Z)
        deleteat!(F, jj)
        deleteat!(Z, jj)
    end
    m = m - length(ii)
    SF = spdiagm(M-m, M-m, 0 => F)
    Sf = diagm(f)
    C = 1 ./ (Z .- transpose(z))
    A = SF*C - C*Sf
    G = svd(A)
    w[:] .= G.V[:, m]
    println("cleanup: ", size(z), "  ", size(f), "  ", size(w))
    return nothing
end


@testset "F-test 1" begin
    x1 = [19.4, 23.9, 21.0, 22.8, 15.8]
    x2 = [34.9, 27.1, 6.9, 13.2, 7.3]
    res_ftest = f_test(x1, x2)
    test_show(res_ftest)
    @test res_ftest.statistic ≈ 15.5838287752675
    @test res_ftest.dofs[1] == 4
    @test res_ftest.dofs[2] == 4
    @test res_ftest.pval ≈ 0.0209393152025605
end

@testset "F-test 2" begin
    df_sleep = dataset("datasets", "sleep")
    res_ftest = f_test(@formula(Extra ~ Group), df_sleep)
    test_show(res_ftest)
    @test res_ftest.statistic ≈ 1.2525950355841
    @test res_ftest.dofs[1] == 9
    @test res_ftest.dofs[2] == 9
    @test res_ftest.pval ≈ 0.742719931726045
end


#############################
##  Numerical experiments associated with section 5.4.1
#############################
## Nonlinear equations are treated by first defining an objective function, then using standard iterative methods
## such as Newton's method via NLsolve.jl. Using more sophisticated methods, linesearches etc. can sometimes result in significant speed boosts.
##

using NLsolve, ApproxFun, MultivariateOrthogonalPolynomials, BandedMatrices, BlockBandedMatrices, Plots
using SparseVolterraExamples

#####################################################
## Problem in Equation (25), u_1(x) = exp(x)
#####################################################
####
## Define g(x), the nonlinearity and the kernel
g(x) = exp(x)+1/3*x*(1-exp(3*x))
    gFun = Fun(g,Jacobi(0,1,0..1))
    f(u) = u^3
    K(x,y) = x
    n = 15 # Desired polynomial order of approximation for accuracy
####
## Define the objective function and compute the solution using a standard Newton method
objective(x) = triNonLinearVolterraObjective(x,f,K,gFun,n)
sol = nlsolve(objective,zeros(n),method=:newton,ftol=1e-8)
####
## Plot function aganinst analytic solution
plot(Fun(Jacobi(0,1,0..1),sol.zero),grid=false,legend=:topleft,xlabel="x",ylabel="u(x)",label="sparse method")
plot!(exp,0,1,grid=false,legend=:topleft,xlabel="x",ylabel="u(x)",label="analytic exp(x)")
####
## Or plot error. A simple loop over polynomial degree can then produce Figure 7(a).
plot(Fun(Jacobi(0,1,0..1),sol.zero)-Fun(x->exp(x),Jacobi(0,1,0..1)),grid=false,legend=:bottomleft,xlabel="x",ylabel="error",label=false)

#####################################################
## Problem in Equation (26), u_2(x) = sin(x)
#####################################################
####
## Define g(x), the nonlinearity and the kernel
g(x) = sin(x)+1/4*sin(x)^2-1/4*x^2
    gFun = Fun(g,Jacobi(0,1,0..1))
    f(u) = u^2
    K(x,y) = x-y
    n = 14 # Desired polynomial order of approximation for accuracy
####
## Define the objective function and compute the solution using a standard Newton method
objective(x) = triNonLinearVolterraObjective(x,f,K,gFun,n)
sol = nlsolve(objective,zeros(n),method=:newton,ftol=1e-15)
####
## Plot function aganinst analytic solution
plot(Fun(Jacobi(0,1,0..1),sol.zero),grid=false,legend=:topleft,xlabel="x",ylabel="u(x)",label="sparse method")
plot!(sin,0,1,grid=false,legend=:topleft,xlabel="x",ylabel="u(x)",label="analytic sin(x)")
####
## Or plot error. A simple loop over polynomial degree can then produce Figure 7(b).
plot(Fun(Jacobi(0,1,0..1),sol.zero)-Fun(x->sin(x),Jacobi(0,1,0..1)),grid=false,legend=:bottomleft,xlabel="x",ylabel="error",label=false)

examples/Section 5-4-1 - Nonlinear Volterra integral equations.jl


using LinearAlgebra

"""
    is_schur(A::AbstractMatrix)

Check if a square matrix is Schur, i.e. the eigenvalues lie within the unit circle.

### Input
- `A` -- square matrix

### Ouput
Boolean value indicating if the matrix `A` is Schur.
"""
is_schur(A::AbstractMatrix) = maximum(abs.(eigvals(A))) < 1.0


export solarposition, zenith, zenith_azimuth, sines_and_cosines

# Calculate the position of the sun as (zenith,azimuth) in degrees for a location
# on earth given as (latitude,longitude) in degrees. The DateTime should be given
# as Universal Time, e.g. UTC.
function solarposition_wiki(datetime::DateTime, latitude, longitude)
	# Position of the Sun as seen from Earth.
	# Ported from Matlab, original helpstring below.

	#   This is the most basic algorithm. It is documented in Seinfeld &
	#   Pandis, Duffie & Beckman and Wikipedia.
	#
	# [ANGLES,PROJECTION] = SOLARPOSITION(DATE,TIME,LATITUDE,LONGITUDE,TIME_ZONE)
	# returns ZENITH & AZIMUTH for all DATE & TIME pairs at LATITUDE, LONGITUDE.
	# ANGLES = [ZENITH,AZIMUTH] and PROJECTION = [PHI_X, PHI_Y]
	# PHI_X is projection on x-z plane & PHI_Y is projection on y-z plane.
	# DATETIME can be string, vector [YEAR, MONTH, DAY, HOURS, MINUTES, SECONDS],
	#   cellstring or matrix N x [YEAR, MONTH, DAY, HOURS, MINUTES, SECONDS] for N
	#   times.
	# LATITUDE [degrees] and LONGITUDE [degrees] are the coordinates of the site.
	# TIME_ZONE [hours] of the site.
	# ROTATION [degrees] clockwise rotation of system relative to north.
	# DST [logical] flag for daylight savings time, typ. from March to November
	#   in the northern hemisphere.
	#
	# References:
	# http://en.wikipedia.org/wiki/Solar_azimuth_angle
	# http://en.wikipedia.org/wiki/Solar_elevation_angle
	#
	# Mark A. Mikofski
	# Copyright (c) 2013

	# Equation of Time
	# "Alternative calculation" on wikipedia:
	# https://en.wikipedia.org/wiki/Equation_of_time#Calculating_the_equation_of_time
	# Supposedly accurate to 6 seconds (0.1 minutes)

	year = Dates.year(datetime)
	# date in days starting at zero on 1 January (the subtraction produces a difference in milliseconds)
	d_ms = datetime - DateTime(year,1,1)
	d = d_ms.value/1000/3600/24

	# Earth's mean angular orbital velocity [degrees/day]
	w = 360/365.24
	# angle the earth moves on its orbit at its average speed from the December solstice
	# 10 is the approximate number of days from the December solstice to January 1
	a = w * (d+10)
	# The angle the Earth moves from the solstice to date D, including a first-order correction
	# for the Earth's orbital eccentricity, 0.0167. The number 2 is the number of days from 1 January
	# to the date of the Earth's perihelion.
	b = a + (360/pi*0.0167)*sind(w*(d-2))
	# C is the difference between the angles moved at mean speed, and at the corrected speed projected
	# onto the equatorial plane, and divided by 180 to get the difference in "half turns". The value 23.44°
	# is the obliquity (tilt) of the Earth's axis. The subtraction gives the conventional sign to the equation of time.
	c = (a - atand(tand(b)/cosd(23.44))) / 180
	# For any given value of x, arctan x (sometimes written as tan−1 x) has multiple values, differing from each other
	# by integer numbers of half turns. This may cause C to be wrong by an integer number of half turns. The excess
	# half turns are removed in the next step of the calculation to give the equation of time:
	ET = 720*(c - round(c))

	# approximate solar time [hours]
	solarTime = 24*(d - floor(d)) + longitude*24/360 + ET/60
	t_h = 15*(solarTime - 12) # [degrees] hour angle

	# declination [degrees]
	# accurate to 0.2 degrees, see https://en.wikipedia.org/wiki/Position_of_the_Sun#Calculations
	delta = -asind(sind(23.44)*cosd(b))
	zenith = acosd(sind(latitude)*sind(delta) + cosd(latitude)*cosd(delta)*cosd(t_h)) # [degrees] zenith

	# azimuth [0, 180], absolute value measured from due south, so east = west = 90,
	# south = 0, north = 180
	cos_phi = (cosd(zenith)*sind(latitude) - sind(delta)) ./ (sind(zenith)*cosd(latitude)) # cosine(azimuth)
	phi_south = acosd(clamp(cos_phi,-1,1))
	#phi_south = acosd(cos_phi);

	# azimuth [0, 360], measured clockwise from due north,
	# so east = 90, south = 180, and west = 270 degrees
	azimuth = 180 + sign(t_h)*phi_south # Shift domain to 0-360 deg

	#angles = [theta, phi]; # [degrees] zenith, azimuth
	return zenith, azimuth, solarTime, delta
end

function solarposition_wiki_faster(datetime::DateTime, latitude, longitude)
	year = Dates.year(datetime)
	d_ms = datetime - DateTime(year,1,1)
	d = d_ms.value/1000/3600/24
	w = 360/365.24
	a = w * (d+10)
	b = a + (360/pi*0.0167)*sind(w*(d-2))
	c = (a - atand(tand(b)/cosd(23.44))) / 180
	ET = 720*(c - round(c))
	solarTime = 24*(d - floor(d)) + longitude*24/360 + ET/60
	t_h = 15*(solarTime - 12) # [degrees] hour angle
	delta = -asind(sind(23.44)*cosd(b))
	slat, clat = sind(latitude), cosd(latitude)
	sdel, cdel = sind(delta), cosd(delta)
	czen = slat*sdel + clat*cdel*cosd(t_h)
	zenith = acosd(czen) # [degrees] zenith

	# azimuth [0, 180], absolute value measured from due south, so east = west = 90,
	# south = 0, north = 180
	cos_phi = (czen*slat - sdel) / (sind(zenith)*clat) # cosine(azimuth)
	phi_south = acosd(clamp(cos_phi,-1,1))
	#phi_south = acosd(cos_phi);

	# azimuth [0, 360], measured clockwise from due north,
	# so east = 90, south = 180, and west = 270 degrees
	azimuth = 180 + sign(t_h)*phi_south # Shift domain to 0-360 deg

	#angles = [theta, phi]; # [degrees] zenith, azimuth
	return zenith, azimuth, solarTime, delta
end

# Splits the DateTime into (year, month, day) (as integers) and the hour (as a floating point)
function splitdatetime(datetime::DateTime)
	# date = floor(datetime, Dates.Day)
	y, m, d = Dates.yearmonthday(datetime)
	return y, m, d, (datetime - DateTime(y,m,d)).value/1000/3600
end

# time correction term, difference between Terrestrial Time and Universal Time, see Grena (2012) section 2.1.
# Conversion to TT is necessary since UT follows the rotation of the earth, which is not uniform.
# t [days] days from January 1, 2060 (in UT)
# returns Δτ [seconds]
differenceTTUT(t) = 96.4 + 0.00158*t

# Calculate time variable needed for Grena (2012) algorithms, see section 3.1.
# The time t is the input DateTime expressed as days from January 1, 2060 (in UT).
# returns t [days]
calctime(datetime::DateTime) = (datetime-DateTime(2060,1,1)).value/1000/3600/24

# Same result as above, but this is the function in Grena's paper. The Julia version is faster.
function calctime_grena(datetime::DateTime)
	y, m, d, h = splitdatetime(datetime)
	if m <= 2
		m += 12
		y -= 1
	end
	return trunc(365.25*(y-2000)) + trunc(30.6001*(m+1)) - trunc(0.01*y) + d + 0.0416667*h - 21958
end

# Full version of algorithm 1 of Grena (2012), see section 3.2.
# Calculate the position of the sun as (zenith,azimuth) in radians for a location
# on earth given as (latitude,longitude) in degrees. The DateTime should be given
# as Universal Time, e.g. UTC.
function solarposition_grena1(datetime::DateTime, latitude, longitude)
	ω = 0.017202786 # [1/day]
	t = calctime(datetime)
	t_e = t + 1.1574e-5 * differenceTTUT(t)
	s1, c1 = sin(ω*t_e), cos(ω*t_e)
	s2, c2 = 2*s1*c1, (c1+s1)*(c1-s1)
	α = mod(-1.38880 + 1.72027920e-2*t_e + 3.199e-2*s1 - 2.65e-3*c1 + 4.050e-2*s2 + 1.525e-2*c2, 2π)
	δ = 6.57e-3 + 7.347e-2*s1 - 3.9919e-1*c1 + 7.3e-4*s2 - 6.60e-3*c2
	H = mod(1.75283 + 6.3003881*t + deg2rad(longitude) - α + π, 2π) - π
	sϕ, cϕ = sind(latitude), cosd(latitude)
	sδ, cδ = sin(δ), cos(δ)
	sH, cH = sin(H), cos(H)
	se0 = sϕ*sδ + cϕ*cδ*cH
	ep = asin(se0) - 4.26e-5*sqrt(1-se0^2)  #cos(e0)
	azimuth = atan(sH, cH*sϕ - sδ*cϕ/cδ)	#azimuth = 0 towards south and positive direction towards west
	zenith = π/2 - ep
	return zenith, azimuth				# radians
	# return zenith, azimuth, δ, H		# radians
	# return rad2deg(zenith), rad2deg(mod(azimuth+π,2π)), rad2deg(H)/15+12, rad2deg(δ) # degrees, same azimuth origin as solarposition0
end

# First half of algorithm 1 of Grena (2012), see section 3.2.
# Calculate the position of the sun as (declination, hour angle) in radians for a location
# on earth given by its longitude in degrees. The DateTime should be given
# as Universal Time, e.g. UTC.
function solarposition(datetime::DateTime, longitude)
	ω = 0.017202786 # [1/day]
	t = calctime(datetime)
	t_e = t + 1.1574e-5 * differenceTTUT(t)
	s1, c1 = sin(ω*t_e), cos(ω*t_e)
	s2, c2 = 2*s1*c1, (c1+s1)*(c1-s1)
	α = mod(-1.38880 + 1.72027920e-2*t_e + 3.199e-2*s1 - 2.65e-3*c1 + 4.050e-2*s2 + 1.525e-2*c2, 2π)
	δ = 6.57e-3 + 7.347e-2*s1 - 3.9919e-1*c1 + 7.3e-4*s2 - 6.60e-3*c2
	H = mod(1.75283 + 6.3003881*t + deg2rad(longitude) - α + π, 2π) - π
	return δ, H
end

# Full version of algorithm 1 of Grena (2012), see section 3.2.
# Calculate the position of the sun as (zenith, azimuth) in degrees for a location
# on earth given as (latitude,longitude) in degrees. The DateTime should be given
# as Universal Time, e.g. UTC.
function solarposition(datetime::DateTime, latitude, longitude)
	δ, H = solarposition(datetime, longitude)
	zenith, azimuth = zenith_azimuth(latitude, sines_and_cosines(δ, H)...)
	return rad2deg(zenith), rad2deg(azimuth)
end

sines_and_cosines(δ, H) = sin(δ), cos(δ), sin(H), cos(H)

function zenith_azimuth(latitude, sδ, cδ, sH, cH)
	sϕ, cϕ = sind(latitude), cosd(latitude)
	se0 = sϕ*sδ + cϕ*cδ*cH
	azimuth = atan(sH, cH*sϕ - sδ*cϕ/cδ)	# azimuth = 0 towards south and positive direction towards west
	zenith = π/2 - asin(se0) + 4.26e-5*sqrt(1-se0^2)  # sqrt(1-se0^2) = cos(e0), but slightly faster
	return zenith, azimuth		# radians
end

function zenith(latitude, sδ, cδ, sH, cH)
	sϕ, cϕ = sind(latitude), cosd(latitude)
	se0 = sϕ*sδ + cϕ*cδ*cH
	return π/2 - asin(se0) + 4.26e-5*sqrt(1-se0^2)  # sqrt(1-se0^2) = cos(e0), but slightly faster
end

const solarconstant = 1361  # W/m2

# Annual variation of solar insolation outside the atmosphere (W/m2):
# https://www.itacanet.org/the-sun-as-a-source-of-energy/part-2-solar-energy-reaching-the-earths-surface/#2.1.-The-Solar-Constant
solarinsolation(dt::DateTime) = solarconstant * (1 + 0.034*cos(2*π*(dt - DateTime(year(dt))).value/1000/3600/24/365.25))

# dt = DateTime(2019,4,15,13,46)
# δ, H = solarposition(dt, 12)	# return rad2deg(zenith), rad2deg(mod(azimuth+π,2π)), rad2deg(H)/15+12, rad2deg(δ)

# pos = sines_and_cosines(δ, H)
# zen = zenith(58, pos...)


function nlp_cvx_109_012(
    optimizer,
    objective_tol,
    primal_tol,
    dual_tol,
    termination_target = TERMINATION_TARGET_LOCAL,
    primal_target = PRIMAL_TARGET_LOCAL,
)
    model = Model(optimizer)

    @variable(model, x >= 0.00001, start = 0.1)
    @variable(model, y >= 0.00001, start = 0.1)

    @NLobjective(model, Max, log(x + y))
    @NLconstraint(model, (y - 2)^2 <= -x + 2)

    optimize!(model)

    check_status(model, FEASIBLE_PROBLEM, termination_target, primal_target)
    check_objective(model, log(7 / 4 + 5 / 2), tol = objective_tol)
    return check_solution([x, y], [7 / 4, 5 / 2], tol = primal_tol)
end


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

actualAlpha, actualLambda = 2,3
gammaDist = Gamma(actualAlpha,1/actualLambda)
n = 10^2
sample = rand(gammaDist, n)

alphaGrid = 1:0.02:3
lambdaGrid = 2:0.02:5

likelihood = [prod([pdf.(Gamma(a,1/l),v) for v in sample])
                        for l in lambdaGrid, a in alphaGrid]

surface(alphaGrid, lambdaGrid, likelihood, lw=0.1,
	c=cgrad([:blue, :red]), legend=:none, camera = (135,20),
	xlabel=L"\alpha", ylabel=L"\lambda", zlabel="Likelihood")


using NewtonsMethod
using Test

using Polynomials
@testset "NewtonsMethod.jl" begin
    @testset "root_of_a_known_function" begin
        f(x) = (x-1)^3
        f_prime(x) = 3*(x-1)^2
        root = 1
        @test root ≈ fixedpointmap(f, f_prime, 2).value #test with first method (f, f', x_0)
        @test root ≈ fixedpointmap(f, 2).value #test with second method (f, x_0)

        f(x) = (x-4)/x
        f_prime(x) = 4/(x^2)
        root = 4
        @test root ≈ fixedpointmap(f, f_prime, 2).value #test with first method (f, f', x_0)
        @test root ≈ fixedpointmap(f, 2).value #test with second method (f, x_0)

    end
    @testset "comparing with Polynomials.roots" begin
        #check the one of the roots of 2 - 5*x + 2*x^2 and comparing it with roots()
        p = Poly([2, -5, 2], :x)
        p′ = polyder(p)
        root = roots(p)
        @test any(root .≈ fixedpointmap(p, p′, 2).value) #test with first method (f, f', x_0)
        @test any(root .≈ fixedpointmap(p, 2).value) #test with second method (f, x_0)

        p = Poly([2, -5, 2, 10], :x)
        p′ = polyder(p)
        root = roots(p)
        @test any(root .≈ fixedpointmap(p, p′, 2).value) #test with first method (f, f', x_0)
        @test any(root .≈ fixedpointmap(p, 2).value) #test with second method (f, x_0)
    end

    @testset "testing BigFloat" begin
        #check the one of the roots of 2 - 5*x + 2*x^2 and comparing it with roots()
        f(x) = (x-1)^3
        f_prime(x) = 3*(x-1)^2
        root = 1.0
        @test root ≈ fixedpointmap(f, f_prime, BigFloat(2)).value #test with first method (f, f', x_0)
        @test root ≈ fixedpointmap(f, BigFloat(2)).value #test with second method (f, x_0)
    end

    @testset "testing non-convergence" begin
        #check the one of the roots of 2 - 5*x + 2*x^2 and comparing it with roots()
        f(x) = x^2 + 1
        f_prime(x) = 2*x
        @test nothing == fixedpointmap(f, f_prime, 2) #test with first method (f, f', x_0)
        @test nothing == fixedpointmap(f, 2) #test with second method (f, x_0)
    end

    @testset "maxiter is working" begin
        #check the one of the roots of 2 - 5*x + 2*x^2 and comparing it with roots()
        f(x) = (x-1)^3
        f_prime(x) = 3*(x-1)^2
        root = 1
        #if with maxiter = 2, we doesn't get the root, then the test will be passed: it means maxiter is working.
        @test_broken root ≈ fixedpointmap(f, f_prime, 2, maxiter = 2).value #test with first method (f, f', x_0)
        @test_broken root ≈ fixedpointmap(f, 2, maxiter = 2).value #test with second method (f, x_0)
    end
    @testset "tolerance is working" begin
        #check the one of the roots of 2 - 5*x + 2*x^2 and comparing it with roots()
        f(x) = (x-1)^3
        f_prime(x) = 3*(x-1)^2
        root = 1
        #if with tolerance = 1E-3, we doesn't get the root, then the test will be passed: it means maxiter is working.
        @test_broken root ≈ fixedpointmap(f, f_prime, 2, tolerance = 1E-3).value #test with first method (f, f', x_0)
        @test_broken root ≈ fixedpointmap(f, 2, tolerance = 1E-3).value #test with second method (f, x_0)
    end

end


using StatsBase, Distributions, Plots; pyplot()

function prn(lambda)
    k, p = 0, 1
    while p > MathConstants.e^(-lambda)
        k += 1
        p *= rand()
    end
    return k-1
end

xGrid, lambda, N = 0:16, 5.5, 10^6

pDist = Poisson(lambda)
bPmf = pdf.(pDist,xGrid)
data = counts([prn(lambda) for _ in 1:N],xGrid)/N

plot( xGrid, data,
	line=:stem, marker=:circle,
	c=:blue, ms=10, msw=0, lw=4, label="MC estimate")
plot!( xGrid, bPmf, line=:stem,
	marker=:xcross, c=:red, ms=6, msw=0, lw=2, label="PMF",
	ylims=(0,0.2), xlabel="x", ylabel="Probability of x events")


function identical_matrix(n::Integer)
    ret = zeros(n, n)
    for itr1 in 1 : n
        ret[itr1, itr1] = 1
    end
    return ret
end

function exchange_row!(mat::Matrix, a::Integer, b::Integer)
    # if $a$ or $b$ are out of bound, do nothing
    if a > size(mat, 1) || b > size(mat, 1)
        return mat
    end
    # exchange row $a$ and row $b$
    for itr1 in 1 : size(mat, 2)
        mat[a, itr1], mat[b, itr1] = mat[b, itr1], mat[a, itr1]
    end
    return mat
end

function similarity_transformation!(
        mat::Matrix, a::Integer, b::Integer)
    # if $a$ or $b$ are out of bound, do nothing
    if a > size(mat, 1) || b > size(mat, 1)
        return mat
    end
    # taking advantage of the matrix, the transformation
    # is just exchanging part of the row $a$ and row $b$
    for itr1 in 1 : (min(a, b) - 1)
        mat[a, itr1], mat[b, itr1] = mat[b, itr1], mat[a, itr1]
    end
    return mat
end

function plu(mat::Matrix)
    # make a copy of the original matrix, which will finally become $U$
    U = float(copy(mat))
    # # U = copy(mat)
    # it's assumed that the matrix has $m$ rows and $n$ columns
    m, n = size(U)
    # initialize two identical matrices as $P$ and $L$
    P = identical_matrix(m)
    L = identical_matrix(m)
    # $i$ and $j$ are used to represent iterators of row and column,
    # respectively
    i = j = 1
    while i <= m && j <= n
        # find the entry with largest absolute value
        imax, vmax = 0.0, 0.0
        for itr1 in i : m
            if vmax < abs(U[itr1, j])
                imax, vmax = itr1, abs(U[itr1, j])
            end
        end
        (vmax > 0) ? (k = imax) : (j += 1; continue)
        # left multiply $U$ and $P$ with a permutation matrix
        if k != i
            exchange_row!(U, i, k)
            exchange_row!(P, i, k)
        end
        # left and right multiply $L$ with the same matrix
        similarity_transformation!(L, i, k)
        # `forward elimination'
        for itr1 in (i + 1) : m
            scaler = U[itr1, j] / U[i, j]
            for itr2 in j : n
                U[itr1, itr2] -= U[i, itr2] * scaler
            end
            # right multiply $L$ with an elimination matrix
            L[itr1, i] = scaler
        end
        i += 1
        j += 1
    end
    return P, L, U
end


using ThreeBodyProblem
using DifferentialEquations
using Plots
using LinearAlgebra
# using ForwardDiff

# The first step is to define the system. ThreeBodyProblem.jl has built-in functions for common systems like Sun/Jupiter.
sys = sun_jupiter()

## There's even a recipe to plot the system for us!
plot(sys)

## We can see the Lagrange Points, but where are the Sun and Jupiter? Turns out that the distance between them is so much larger than the radius of either one that it makes it hard to see them on a truly scaled picture. Let's scale the Sun and Jupiter to make them visible.
plot(sys, scaled = true)

## That's better! Now let's compute the L1 and L2 Lagrange points of our system

L1, L2 = computeLpts(sys)
# Declare state vectors for L1 and L2 (with zero velocity)
rv1 = [L1; zeros(3)] # state at Lagrange point L1 (zero velocity added on)
rv2 = [L2; zeros(3)] # state at Lagrange point L2

## The Lagrange points are equilibrium points in our dynamics. This means that an object placed there perfectly will stay there forever. But, Lagrange points L1, L2, and L3 are unstable, meaning if you perturb the object slightly, it may fall away. This is how we will generate the invariant manifolds.
## We first need to linearize the dynamics about these Lagrange points to determine which directions are stable and unstable

Φ₀ = I(6) # The initial State Transition Matrix (STM)
w₀ = vcat(rv1,reshape(Φ₀,36,1)) # Augmented state with state and STM
# Set up the ODE problem
tspan = (0.,1.)
prob = ODEProblem(CR3BPstm!,w₀,tspan,sys)
sol = solve(prob, reltol=1e-6)
Φₜ = Matrix(reshape(sol.u[end][7:42],6,6))

# last time step should match very closely with rv1
rvₜ = sol.u[end][1:6]
err = norm(rv1 - rvₜ)

# Compute the eigenvectors of the Monodromy Matrix
Λ,V = eigen(Φₜ,sortby=isreal) # Λ is vector of eigenvalues, V is matrix of eigenvectors
Yw = real(V[:,findall(isreal, Λ)]) # Eigenvectors corresponding to real eigenvalues
Λ = Λ[findall(isreal, Λ)] # Purely real eigenvalues (have 0.0 imaginary component)
Yws = Yw[:,findmin(real(Λ))[2]] # Eigenvector associated with stable eigenvalue λ < 1
Ywu = Yw[:,findmax(real(Λ))[2]] # Eigenvector associated with unstable eigenvalue λ > 1

## We perturb our states in the stable and unstable directions
α = 1e-6

# L1 points
rv1up = rv1 + α*Ywu # unstable manifold + side
rv1un = rv1 - α*Ywu # unstable manifold - side
rv1sp = rv1 + α*Yws # stable manifold + side
rv1sn = rv1 - α*Yws # stable manifold - side

# Integrate trajectories for 2 Jupiter periods
tf = 2*2π
tspan_forward = (0.,tf) # make sure to add decimal points so tspan contains floating point values, not integers
tspan_back = (0.,-tf)
myreltol = 1e-12

# Set up the ODE Problems
prob1up = ODEProblem(CR3BPdynamics!,rv1up,tspan_forward,sys) # Unstable positive
prob1un = ODEProblem(CR3BPdynamics!,rv1un,tspan_forward,sys) # Unstable negative
prob1sp = ODEProblem(CR3BPdynamics!,rv1sp,tspan_back,sys) # Stable positive
prob1sn = ODEProblem(CR3BPdynamics!,rv1sn,tspan_back,sys) # Stable negative

# Solutions to the ODEs
sol1up = solve(prob1up, reltol=myreltol)
sol1un = solve(prob1un, reltol=myreltol)
sol1sp = solve(prob1sp, reltol=myreltol)
sol1sn = solve(prob1sn, reltol=myreltol)

plot(sys,scaled=true)
plot!(sol1up,vars=(1,2),label="Wu+",linecolor=:red)
plot!(sol1un,vars=(1,2),label="Wu-",linecolor=:magenta)
plot!(sol1sp,vars=(1,2),label="Ws+",linecolor=:blue)
plot!(sol1sn,vars=(1,2),label="Ws-",linecolor=:cyan)
plot!(aspect_ratio=1,ylims=[-1,1],xlims=[-0.6,1.3],legend=:topleft,flip=false)

## We can see two strands moving in towards the sun, but the rest is a jumbled mess. Let's zoom in to see what's going on. Also, I need to make sure things are scaled correctly, or Jupiter will look way too big.
plot(sys)
plot!(sol1up,vars=(1,2),label="Wu+",linecolor=:red)
plot!(sol1un,vars=(1,2),label="Wu-",linecolor=:magenta)
plot!(sol1sp,vars=(1,2),label="Ws+",linecolor=:blue)
plot!(sol1sn,vars=(1,2),label="Ws-",linecolor=:cyan)
plot!(aspect_ratio=1,ylims=[-0.04,0.04],xlims=[0.93,1.06],legend=:outerright,flip=false)

# L2 points
rv2up = rv2 + α*Ywu # unstable manifold + side
rv2un = rv2 - α*Ywu # unstable manifold - side
rv2sp = rv2 + α*Yws # stable manifold + side
rv2sn = rv2 - α*Yws # stable manifold - side

## Now we look at the L2 point manifolds
prob2up = ODEProblem(CR3BPdynamics!,rv2up,tspan_forward,sys) # Unstable positive
prob2un = ODEProblem(CR3BPdynamics!,rv2un,tspan_forward,sys) # Unstable negative
prob2sp = ODEProblem(CR3BPdynamics!,rv2sp,tspan_back,sys) # Stable positive
prob2sn = ODEProblem(CR3BPdynamics!,rv2sn,tspan_back,sys) # Stable negative

sol2up = solve(prob2up, reltol=1e-6)
sol2un = solve(prob2un, reltol=1e-6)
sol2sp = solve(prob2sp, reltol=1e-6)
sol2sn = solve(prob2sn, reltol=1e-6)

plot(sys, scaled=true)
plot!(sol2up,vars=(1,2),label="Wu+",linecolor=:red)
plot!(sol2un,vars=(1,2),label="Wu-",linecolor=:magenta)
plot!(sol2sp,vars=(1,2),label="Ws+",linecolor=:blue)
plot!(sol2sn,vars=(1,2),label="Ws-",linecolor=:cyan)
plot!(aspect_ratio=1,ylims=[-1.0,1.0],xlims=[-0.3,1.5],legend=:topright,flip=false)
## This time there are two strands moving out away from the Sun, along with another jumble around Jupiter. Let's zoom in again to take a look.
plot(sys)
plot!(sol2up,vars=(1,2),label="Wu+",linecolor=:red)
plot!(sol2un,vars=(1,2),label="Wu-",linecolor=:magenta)
plot!(sol2sp,vars=(1,2),label="Ws+",linecolor=:blue)
plot!(sol2sn,vars=(1,2),label="Ws-",linecolor=:cyan)
plot!(aspect_ratio=1,ylims=[-0.04,0.04],xlims=[0.93,1.07],legend=:outerright,flip=false)


################################################################################
#
# The LLL-Reduction Algorithm
#
# Implementration of algorithm described in :
#  Lenstra, A.K., Lenstra, H.W. & Lovász, L. Math. Ann. (1982)
#            261: 515. https://doi.org/10.1007/BF01457454
#
################################################################################

"""
  lll(b₀::AbstractMatrix{Float64})::AbstractMatrix{Float64}

Performs Lenstra-Lenstra-Lovasv reduction on the n vectors in b₀.

Parameters
----------
- b₀ : vector of n vectors of n Float64

Returns
-------
- The LLL reduction of the vectors in `b₀`
"""
function lll(b₀::Vector{Vector{Float64}})::Vector{Vector{Float64}}
    b = deepcopy(b₀) # preserve argument by working on a copy
    n = length(b)

    # initialize mu and B with zeros
    μ = zeros(n,n)
    B = zeros(n)

    # function (*) of page 521
    function ☼(l)
        abs(μ[k,l]) < 0.5 && return
        r = round(μ[k,l])
        b[k] .-= r * b[l]
        μ[k,1:l-1] .-= r * μ[l,1:l-1]
        μ[k,l] -= r
    end

    b₂ = similar(b)
    for t in 1:n
        b₂[t] = copy(b[t])
        for j in 1:t-1
            μ[t,j] = dot(b[t], b₂[j]) / B[j]
            b₂[t] .-= μ[t,j] * b₂[j]
        end
        B[t] = dot(b₂[t], b₂[t])
    end

    k = 2
    while k <= n
        ☼(k-1)
        if B[k] >= (0.75 - μ[k,k-1]^2) * B[k-1]
            foreach(☼, k-2:-1:1)
            k += 1
        else
            μₒ = μ[k,k-1]
            Bₒ = B[k] + μₒ^2 * B[k-1]
            μ[k,k-1] = μₒ * B[k-1] / Bₒ
            B[k] = B[k] * B[k-1] / Bₒ
            B[k-1] = Bₒ

            bₒ = b[k] ; b[k] = b[k-1] ; b[k-1] = bₒ
            for j in 1:k-2
                tmp = μ[k,j] ; μ[k,j] = μ[k-1,j] ; μ[k-1,j] = tmp
            end
            for t in k+1:n
                tmp = μ[k,k-1] - μₒ * μ[t,k]
                μ[t,k-1] = μ[t,k] + μ[k,k-1] * tmp
                μ[t,k]   = tmp
            end

            k = max(2, k-1)
        end
    end
    b
end