## activate package environment
cd(@__DIR__)
using Pkg
Pkg.activate("Project.toml")

## parameters (calculation for perpendicular incident)
λ = 520 # nm
N = [1.0; 3.812 + 0.1903im; 4.1891+0.3633im] # air - window layer - absorber layer
D = 100 # thickness of window layer
# (refractive indices are from https://refractiveindex.info/, data by Aspnes 1986)

## function for fresnel coefficients (for perpendicular incident)
function fresnel_coeff(n,m,pol)
    if pol == "s"
        r = (n-m)/(n+m)
        t = (2*n)/(n+m)
    elseif pol == "p"
        r = (m-n)/(m+n)
        t = (2*n)/(m+n)
    end
    return r,t
end

## function for interface propagation matrix
function interface_prop(n,m,pol)
    r,t = fresnel_coeff(n,m,pol)
    I = 1/t * [1 r;r 1]
    return I
end

## function for propagation through film Matrix
function film_prop(d,n)
    beta = 2*pi*d*n/λ
    L = [exp(-1im*beta) 0; 0 exp(1im*beta)]
    return L
end

## calcuation of scattering matrix
# Emanuele Centurioni: Generalized matrix method for calculation of internal light energy flux
# in mixed coherent and incoherent multilayers
Ss = interface_prop(N[1],N[2],"s") * film_prop(D,N[2]) * interface_prop(N[2],N[3],"s")
Sp = interface_prop(N[1],N[2],"p") * film_prop(D,N[2]) * interface_prop(N[2],N[3],"p")

## backside transmittance
Ts = abs2(1/Ss[1,1])*real(N[3])/real(N[1])
Tp = abs2(1/Ss[1,1])*real(conj(N[3]))/real(N[1])
T  = round((Ts+Tp)/2,sigdigits=4)

## ouput result
println("")
println("transmission into absorber T = $T")


"""
Calculates Bézout coefficients (see `gcdx`)
"""
function bezout(a::R, b::R) where {R}
    rev = a < b
    x, y = rev ? (a,b) : (b,a)

    s0, s1 = oneunit(R), zero(R)
    t0, t1 = zero(R), oneunit(R)

    while y != zero(R)
        q = div(x, y)
        x, y = y, x - y * q
        s0, s1 = s1, s0 - q * s1
        t0, t1 = t1, t0 - q * t1
    end

    s, t = rev ? (s0, t0) : (t0, s0)
    g = x

    if g == a
        s = one(R)
        t = zero(R)
    elseif g == -a
        s = -one(R)
        t = zero(R)
    end

    return s, t, g
end


using Revise
using AdFem
using PyPlot

# forward computation
mmesh = Meh(joinpath(PDATA, "twoholes_large.stl"))
xy = gauss_nodes(mmesh)
κ = @. sin(xy[:,1]) * (1+xy[:,2]^2) + 1.0
f = 1e5 * @. xy[:,1] + xy[:,2]
K = compute_fem_laplace_matrix1(κ, mmesh)
F = compute_fem_source_term1(f, mmesh)
bdnode = bcnode(mmesh)
K, F = impose_Dirichlet_boundary_conditions(K, F, bdnode, zeros(length(bdnode)))
sol = K\F

# inverse modeling
nn_κ = squeeze(fc(xy, [20,20,20,1])) + 1
K = compute_fem_laplace_matrix1(nn_κ, mmesh)
F = compute_fem_source_term1(f, mmesh)
bdnode = bcnode(mmesh)
K, F = impose_Dirichlet_boundary_conditions(K, F, bdnode, zeros(length(bdnode)))
nn_sol = K\F
loss = sum((nn_sol - sol)^2)

sess = Session(); init(sess)
BFGS!(sess, loss)

nn_val = run(sess, nn_κ)
close("all")
figure(figsize=(18,5))
subplot(131)
visualize_scalar_on_gauss_points(κ, mmesh)
title("Reference")
subplot(132)
visualize_scalar_on_gauss_points(nn_val, mmesh)
title("DNN")
subplot(133)
visualize_scalar_on_gauss_points(abs.(κ-nn_val), mmesh)
title("Absolute Difference")
savefig("poisson_kappa.png")
close("all")
visualize_scalar_on_fem_points(sol, mmesh)
savefig("poisson_solution.png")


@testset "FNT 1D" begin
    for t in BigInt.(1:13)
        @show t
        x = [1:2^(t+1);] .|> BigInt
        @show length(x)
        g = 2 |> BigInt
        q = 2^2^t + 1 |> BigInt
        @time  @test ifnt(fnt(x, g, q), g, q) == x
    end
end

@testset "FNT and NTT coherence" begin
    # FNT is special case with faster implementation of general NTT and should
    # give the same results
    t = 2
    F_t = 2^2^t + 1 #a Fermat number
    x = [1:8;]

    @test fnt(x, 2, F_t) == ntt(x, 2, F_t)
end

@testset "FNT 1D convolution" begin
    using DSP

    t = 3
    q = 2^2^t + 1
    g = 2
    # padding inputs with zeros to get non circular convolution
    x = [1:8; zeros(Int, 8);]
    h = [1:8; zeros(Int, 8);]
    X = fnt(x, g, q)
    H = fnt(h, g, q)
    Y = mod.(X .* H, q)
    y = ifnt(Y, g, q)

    @test y[1:15] == DSP.conv(1:8, 1:8)
end

@testset "FNT 2D" begin
    t = 2
    q = 2^2^t + 1
    g = 2
    x = mod.(rand(Int, 8, 8), q)
    @test ifnt(fnt(x, g, q), g, q) == x

    t = 4
    q = 2^2^t + 1
    g = 314
    x = mod.(rand(Int, 512, 512), q)
    @test ifnt(fnt(x, g, q), g, q) == x
end

@testset "FNT 2D convolution" begin
    using DSP

    t = 3
    q = 2^2^t + 1
    g = 2
    x = mod.(rand(Int, 8, 8), 5)
    x_padded = zeros(Int64, 16, 16)
    x_padded[1:8, 1:8] = x
    h = mod.(rand(Int, 8, 8), 3)
    h_padded = zeros(Int64, 16, 16)
    h_padded[1:8, 1:8] = h
    X = fnt(x_padded, g, q)
    H = fnt(h_padded, g, q)
    Y = mod.(X .* H, q)
    y = ifnt(Y, g, q)

    @test y[1:15, 1:15] == DSP.conv(x, h)
end

@testset "FNT in-place" begin
    t = 2
    q = 2^2^t + 1
    g = 2

    x = mod.(rand(Int, 8), q)
    y = fnt(x, g, q)
    @test x != y
    fnt!(x, g, q)
    @test x == y

    x = mod.(rand(Int, 8, 8), q)
    y = fnt(x, g, q)
    @test x != y
    fnt!(x, g, q)
    @test x == y

end

@testset "FNT fails for non fermat modulus" begin
    (g, q, n) = (2, 31, 5)
    x = [1:n;]
    @test_throws AssertionError fnt(x, g, q)

    (g, q, n) = (2, 33, 10)
    x = [1:n;]
    @test_throws AssertionError fnt(x, g, q)
end

@testset "isfermat() tests" begin
    known = map(n->2^2^n+1, [0:4;])
    expected = map(v -> v in known, [1:maximum(known);])
    actual = isfermat.([1:maximum(known);])

    @test expected == actual
end

@testset "modfermat() tests" begin

    for t in 0:3
        q = 2^2^t+1
        x = 0:(q-1)^2
        @test mod.(x, q) == modfermat.(x, q)
    end

    for t in 0:10
        q = BigInt(2)^2^t+1
        limit = (q-1)^2
        x = mod.(rand(0:limit, 1000), (q-1)^2)
        @test mod.(x, q) == modfermat.(x, q)
    end
end


using SparseArrays
using StatsBase
using LinearAlgebra
using Random
using ProgressMeter
using LinearAlgebra
using Test
using ForwardDiff
using Calculus
#include("core.jl")
J=GeoPhyInv

nx=5
nz=5
nznx=nz*nx
ddx=22.5
ddz=10.3

p=randn(nznx+1)
p[end]=0.0
Q=abs.(randn(nz,nx))
σ=abs.(randn(nz,nx))
psi0=randn(nznx)

#paA=Param(nx,nz,σ)
#paB=Param(nx,nz,Q)
#A=paA.A
#B=paB.A

function func_test(x,p)
	global psi0
	x=reshape(x,nz,nx)
	paA=J.Poisson.Param(nx,nz,x)

	psi=vec(J.Poisson.applyinvA(p,paA))
	f= sum(abs2,psi.-psi0)
	return f
end

function grad_test(x,p)
	global psi0
	paA=J.Poisson.Param(nx,nz,x)

	psi=vec(J.Poisson.applyinvA(p,paA))

	adjsrc=-2.0*(psi.-psi0)

	lambda=vec(J.Poisson.applyinvAt(adjsrc,paA))
	#println(size(lambda))

	return vec((lambda*psi'))'*hcat(vec.(paA.dAdx)...)
end

f = x->func_test(x, p);

gcalc = Calculus.gradient(f);
g3=reshape(gcalc(vec(σ)),nz,nx);

g2=reshape(grad_test(σ,p),nz,nx)

@test ≈(g2,g3, rtol=1e-5)


function binomial_coeffs(n::Integer)::Vector{Int}
    if n == 1
        return [1]
    end
    if n == 2
        return [1, 1]
    end
    c = [1, 1]
    for _ in 3:n
        c=[1, (view(c[1:end-1]) .+ view(c[2:end])), 1]
    end
    return c
end

function binomial_coeffs_2(n::Integer)
    if n == 1
        return [1]
    end
    if n == 2
        return [1, 1]
    end
    c = [1 , 1]
    m = 2
    for _ in 3:n
        for i in 1:(m - 1)
            c[i] = c[i] + c[i+1]
            m += 1
        end
        pushfirst!(c, 1)
    end
    return c
end


export nullspacex

"""
`nullspacex(A)` returns an exact basis for the matrix `A`
"""
function nullspacex(A::AbstractMatrix{T}) where {T}
    r, c = size(A)
    B = rrefx(A)

    leads = Int[]
    # in each row, find first 1
    for i = 1:r
        row = B[i, :]
        if all(row .== 0)
            continue
        end
        k = findfirst(row .!= 0)
        append!(leads, k)
    end

    frees = setdiff(collect(1:c), leads)

    result = Matrix{T}(undef, c, 0)
    for f in frees
        v = zeros(T, c)
        v[f] = T(1)
        for i = 1:length(leads)
            l = leads[i]  # (i,l) is a leading 1
            if l < f
                v[l] = -B[i, f]
            end
        end
        result = [result v]
    end
    return result
end

function nullspacex(A::AbstractMatrix{T}) where {T<:IntegerX}
    return nullspacex(big.(A) // 1)
end


const brown_dennis = let res_init=zeros(20), jac_init=zeros(20,4), x_init=[25., 5., -5., -1.]

    function res(x, r)
        n = length(r)
        for i = 1:n
            ti = 0.2e0 * Float64(i)
            r[i] = (x[1] + x[2] * ti - exp(ti)) ^ 2 +
                (x[3] + x[4] * sin(ti) - cos(ti))^2
        end
        return r
    end

    function jac(x, j)
        n = size(j)[1]
        for i = 1:n
            ti = 0.2e0 * Float64(i)
            j[i, 1] = 2.0e0 * (x[1] + x[2]*ti - exp(ti))
            j[i, 2] = ti * j[i, 1]
            t = sin(ti)
            j[i, 3] = 2.0e0 *(x[3] + x[4]*t - cos(ti))
            j[i, 4] = t * j[i, 3]
        end
        return j
    end

    f(;scale=1, verbose=false, print_steps=false) = testone("brown_dennis", res, jac,
                                                            res_init, jac_init, x_init;
                                                            scale=scale, verbose=verbose,
                                                            print_steps=print_steps)
end


function calculate_pi(NBIN::Int= 1000000000)::Float64
    step = 1/NBIN
    sum = 0
    for i=1:NBIN
        x = (i+0.5)*step;
        sum += 4 / (1 + x^2)
    end
    pi = sum*step
    return pi
end

if length(ARGS)>0
    @time @show calculate_pi(parse(Int, ARGS[1]))
else
    @time @show calculate_pi()
end


# //  This is a gear train design problem taken from the GAMS test library
# //
# //  A compound gear train is to be designed to achieve a specific
# //  gear ratio between the driver and driven shafts. The objective
# //  of the gear train design is to find the number of teeth of the
# //  four gears and to obtain a required gear ratio of 1/6.931.
# //
# //  The problem originated from:
# //  Deb, K, and Goyal, M, Optimizing Engineering Designs Using a
# //  Combined Genetic Search. In Back, T, Ed, Proceedings of the
# //  Seventh International Conference on Genetic Algorithms. 1997,
# //  pp. 521-528.

function gear(oct::Bool = false)
    m = JuMP.Model()
    @variable(m, 12 <= i[1:4] <= 60, Int)
    for j = 1:4
        JuMP.set_start_value(i[j], 24)
    end
    @constraint(m, e2, - i[3] + i[4] >= 0)
    @constraint(m, e3, i[1] - i[2] >= 0)
    if !oct
        @NLobjective(m, Min, (6.931 - i[1]*i[2]/(i[3]*i[4]))^2 + 1)
        return m
    else
        @variable(m, obj)
        @objective(m, Min, obj)
        gm = GlobalModel(model = m, name = "gear")
        add_nonlinear_constraint(gm, :((i) -> (6.931 - i[1]*i[2]/(i[3]*i[4]))^2 + 1), dependent_var = obj)
        return gm
    end
end


using Measurements
using Printf

printstyled("\nExperimento 1: Capacidade térmica\n", color =:blue);

# em gramas
Δm = 0.01;
mcopo = 51.82;
maq = 122.07;
maf = 125.02;

# em graus celsius
ΔT = 0.1;
taf = 12.4;
taq = 53.2;
tf = 31.4;

# calor especifico da água
ca = 1;
@printf("%.3f · %.3f ( %.3f - %.3f) / (%.3f - %.3f) - %.3f · %.3f \n", maq,ca,taq,tf,tf,taf,maf,ca);
C = maq * ca * ( (taq - tf) / (tf - taf) ) - maf * ca;

ΔC = ((Δm * taq + ΔT * maq) * (tf - taf) + (2 * ΔT) * (maq * taq) ) / (tf - taf)^2

println("\nCapacidade térmica do calorimetro");
@printf("Ccal = %.2f ± %.2f \n", C, ΔC);

# em grama
mcopo = 51.82 ±  Δm;
cAluminio = 0.218;
Ccopo =  mcopo * cAluminio;

println("\nCapacidade térmica do copo");
@printf("Ccopo = %.3f ± %0.3f \n", Ccopo.val, Ccopo.err);

println("\nRelação de equivalência");
@printf("| x1 - x2 | < 2 · (σ1 + σ2) \n");
@printf("| %.3f - %.3f| < 2 · (%.3f + %.3f) \n", C,Ccopo.val,ΔC,Ccopo.err);
println(abs(C - Ccopo.val) < 2 * (ΔC + Ccopo.err));

printstyled("\nExperimento 2: Calor específico de um metal  ", color =:blue);

# em gramas
mobjeto = 205.75 ± Δm;
ma = 216.07 ± Δm;

# em graus celsius
ta = 13.1 ± ΔT;
to = 97.1 ± ΔT;
tf = 19.2 ± ΔT;

co = ((ma * ca + C ) * (ta - tf)) / (mobjeto * (tf - to));

println("\nCalor específico de um metal");
@printf("Co = %.4f ± %.4f \n", co.val, co.err);

printstyled("\nExperimento 3: Calor latente de condensação da água\n",color=:blue);

ta = 8.8  ± ΔT;
tf = 71.7  ± ΔT;
tc = 100  ± ΔT;

ma = 167.479  ± Δm;
mc = 187.02 - ma  ± Δm;

@printf("Massa condensada = %.3f", mc.val);

L = (((ma + C) * (ta - tf)) / mc) + (tc - tf);

println("\nCalor latente");
@printf("Lc = %.3f ± %.3f",L.val, L.err);


@testset "Type Stability" begin
    function rosenbrock(x::Vector{T}) where T
        o = one(T)
        c = convert(T,100)
        return (o - x[1])^2 + c * (x[2] - x[1]^2)^2
    end

    function rosenbrock_gradient!(storage::Vector{T}, x::Vector{T}) where T
        o = one(T)
        c = convert(T,100)
        storage[1] = (-2*o) * (o - x[1]) - (4*c) * (x[2] - x[1]^2) * x[1]
        storage[2] = (2*c) * (x[2] - x[1]^2)
    end

    function rosenbrock_hessian!(storage::Matrix{T}, x::Vector{T}) where T
        o = one(T)
        c = convert(T,100)
        f = 4*c
        storage[1, 1] = (2*o) - f * x[2] + 3 * f * x[1]^2
        storage[1, 2] = -f * x[1]
        storage[2, 1] = -f * x[1]
        storage[2, 2] = 2*c
    end

    for method in (NelderMead(),
                   SimulatedAnnealing(),
                   BFGS(),
                   ConjugateGradient(),
                   GradientDescent(),
                   MomentumGradientDescent(),
                   AcceleratedGradientDescent(),
                   LBFGS(),
                   Newton())
        for T in (Float32, Float64)
            result = optimize(rosenbrock, rosenbrock_gradient!, rosenbrock_hessian!, fill(zero(T), 2), method)
            @test eltype(Optim.minimizer(result)) == T
        end
    end
end


using Plots;gr()

f(x) = x*x - 2
df(x) = 2x

function newton(x_0, max_iter)
    x = x_0
    error = zeros(max_iter)
    for i in 1:max_iter
        x = x - f(x) / df(x)
        error[i] = abs(x - sqrt(2))

    end

    (x, error)
end

x_0 = 1.0
max_iter = 20

x, error = newton(x_0, max_iter)
print(error)
t = 1:max_iter

plot(t, error, color=:cornflowerblue, label="newton",xlabel="t",ylabel="Error")
savefig("newton.png")


using DrWatson
@quickactivate "StatReth"

# %%
using DataFrames
using CSV
using StatsBase
using Distributions
using StatsPlots

# %% 4.7 - 4.11
d = DataFrame(CSV.File(datadir("exp_raw/Howell_1.csv")))
# precis(d)
d.height

d2 = filter(row -> row.age >= 18, d)    # either
d2 = d[d.age .>= 18, :]                 # or

# %% 4.12, 4.13
plot(p -> pdf(Normal(178, 20), p), 100:250)
plot(p -> pdf(Uniform(0, 50), p), -10:60)

# %% 4.14
N = 10_000
d_μ = Normal(178, 20)
sample_mu = rand(d_μ, N)
d_σ = Uniform(0, 50)
sample_sigma = rand(d_σ, N)

prior_h = rand.(Normal.(sample_mu, sample_sigma))
density(prior_h)

# %% 4.15
d_μ_100 = Normal(178, 100)
sample_μ_100 = rand(d_μ_100, N)

sample_height_100 = rand.(Normal.(sample_μ_100, sample_sigma))
# or is
# mix = MixtureModel(Normal.(sample_μ_100, sample_sigma))
# sample_height_100 = rand(mix, N)
# more correct?
density(sample_height_100)

# %% 4.16, 4.17
μ_grid = range(150, 160, length = 100)
σ_grid = range(7, 9, length = 100)
post = [(; :μ => μ, :σ => σ) for μ in μ_grid, σ in σ_grid]  # this is `post`

ll = [loglikelihood(Normal(μ, σ), d2.height) for (μ, σ) in post]
prior = [logpdf(d_μ, μ) + logpdf(d_σ, σ) for (μ, σ) in post]
lposterior = ll .+ prior
lposterior .-= maximum(lposterior)
posterior = exp.(lposterior)

wireframe(μ_grid, σ_grid, posterior)

# %% 4.19 - 4.22
samples = sample(post, pweights(posterior), N) |> DataFrame

scatter(samples.μ, samples.σ, alpha = 0.05, ms = 5, xlabel = "μ", ylabel = "σ")

density(samples.μ)
density(samples.σ)

quantile(samples.μ, (0.1, 0.9))
quantile(samples.σ, (0.1, 0.9))

# %% 4.23 - 4.25
d3_height = d2.height[1:20]
μ_grid = range(140, 160, length = 200)
σ_grid = range(4, 20, length = 200)
post = [(μ = μ, σ = σ) for μ in μ_grid, σ in σ_grid]

ll = [loglikelihood(Normal(μ, σ), d3_height) for (μ, σ) in post]
prior = [logpdf(d_μ, μ) + logpdf(d_σ, σ) for (μ, σ) in post]
lposterior = ll .+ prior
lposterior .-= maximum(lposterior)
posterior = exp.(lposterior)

wireframe(μ_grid, σ_grid, posterior)

samples2 = sample(post, pweights(posterior), N) |> DataFrame

scatter(samples2.μ, samples2.σ, alpha = 0.05, ms = 5)
density(samples2.μ)
density(samples2.σ)


### A Pluto.jl notebook ###
# v0.12.10

using Markdown
using InteractiveUtils

# ╔═╡ e077ab16-22e5-11eb-15ed-5f0654bf7328
using Pkg, DrWatson

# ╔═╡ e75a9a06-22e5-11eb-2e36-8d4f62b830ed
begin
	@quickactivate "StatisticsWithJuliaPlutoNotebooks"
	using Random, Distributions, KernelDensity, Plots
	Random.seed!(0)
end

# ╔═╡ ca5fbdc8-22e5-11eb-0a60-cb0a127a3ca5
md"## Listing4.16"

# ╔═╡ fa68607e-22e5-11eb-0558-c9a4d9f77426
begin
	mu1, sigma1 = 10, 5
	mu2, sigma2 = 40, 12
	dist1, dist2 = Normal(mu1,sigma1), Normal(mu2,sigma2)
	p = 0.3
	mixRv() = (rand() <= p) ? rand(dist1) : rand(dist2)
	mixPDF(x) = p*pdf(dist1,x) + (1-p)*pdf(dist2,x)

	n = 2000
	data = [mixRv() for _ in 1:n]

	kdeDist = kde(data)
end

# ╔═╡ 45f535f0-250a-11eb-2a6b-03853e2d8eb1
begin
	xGrid = -20:0.1:80
	pdfKDE = pdf(kdeDist,xGrid)

	plot(xGrid, pdfKDE, c=:blue, label="KDE PDF")
	stephist!(data, bins=50, c=:black, normed=:true, label="Histogram")
	p1 = plot!(xGrid, mixPDF.(xGrid), c=:red, label="Underlying PDF",
		xlims=(-20,80), ylims=(0,0.035), legend=:topleft,
		xlabel="X", ylabel = "Density")
end

# ╔═╡ 45f5722a-250a-11eb-2840-a7d3dd59711b
begin
	hVals = [0.5,2,10]
	kdeS = [kde(data,bandwidth=h) for h in hVals]
	plot(xGrid, pdf(kdeS[1],xGrid), c = :green, label= "h=$(hVals[1])")
	plot!(xGrid, pdf(kdeS[2],xGrid), c = :blue, label= "h=$(hVals[2])")
	p2 = plot!(xGrid, pdf(kdeS[3],xGrid), c = :purple, label= "h=$(hVals[3])",
		xlims=(-20,80), ylims=(0,0.035), legend=:topleft,
		xlabel="X", ylabel = "Density")
	plot(p1,p2,size = (800,400))
end

# ╔═╡ 475ff888-22e6-11eb-2354-09f8f40a8e12
md"## End of listing4.16"

# ╔═╡ Cell order:
# ╟─ca5fbdc8-22e5-11eb-0a60-cb0a127a3ca5
# ╠═e077ab16-22e5-11eb-15ed-5f0654bf7328
# ╠═e75a9a06-22e5-11eb-2e36-8d4f62b830ed
# ╠═fa68607e-22e5-11eb-0558-c9a4d9f77426
# ╠═45f535f0-250a-11eb-2a6b-03853e2d8eb1
# ╠═45f5722a-250a-11eb-2840-a7d3dd59711b
# ╟─475ff888-22e6-11eb-2354-09f8f40a8e12


###########################################################################
# Gauss-Kronrod integration-weight computation for arbitrary floating-point
# types and precision, implemented based on the description in:
#
#    Dirk P. Laurie, "Calculation of Gauss-Kronrod quadrature rules,"
#    Mathematics of Computation, vol. 66, no. 219, pp. 1133-1145 (1997).
#
# for the Kronrod rule, and for the Gauss rule from the description in
#
#    Lloyd N. Trefethen and David Bau, Numerical Linear Algebra (SIAM, 1997).
#
# Arbitrary-precision eigenvalue (eignewt & eigpoly) and eigenvector
# (eigvec1) routines are written by SGJ, independent of the above sources.
#
# Since we only implement Gauss-Kronrod rules for the unit weight function,
# the diagonals of the Jacobi matrices are zero and certain things simplify
# compared to the general case of an arbitrary weight function.

# Given a symmetric tridiagonal matrix H with H[i,i] = 0 and
# H[i-1,i] = H[i,i-1] = b[i-1], compute p(z) = det(z I - H) and its
# derivative p'(z), returning (p,p').
function eigpoly(b,z,m=length(b)+1)
    d1 = z
    d1deriv = d2 = one(z)
    d2deriv = zero(z)
    for i = 2:m
        b2 = b[i-1]^2
        d = z * d1 - b2 * d2
        dderiv = d1 + z * d1deriv - b2 * d2deriv
        d2 = d1
        d1 = d
        d2deriv = d1deriv
        d1deriv = dderiv
    end
    return (d1, d1deriv)
end

# compute the n smallest eigenvalues of the symmetric tridiagonal matrix H
# (defined from b as in eigpoly) using a Newton iteration
# on det(H - lambda I).  Unlike eig, handles BigFloat.
function eignewt(b,m,n)
    # get initial guess from eig on Float64 matrix
    H = SymTridiagonal(zeros(m), Float64[ b[i] for i in 1:m-1 ])
    lambda0 = sort(eigvals(H))

    lambda = Array{eltype(b)}(undef, n)
    for i = 1:n
        lambda[i] = lambda0[i]
        for k = 1:1000
            (p,pderiv) = eigpoly(b,lambda[i],m)
            lambda[i] = (lamold = lambda[i]) - p / pderiv
            if abs(lambda[i] - lamold) < 10 * eps(lambda[i]) * abs(lambda[i])
                break
            end
        end
        # do one final Newton iteration for luck and profit:
        (p,pderiv) = eigpoly(b,lambda[i],m)
        lambda[i] = lambda[i] - p / pderiv
    end
    return lambda
end

# given an eigenvalue z and the matrix H(b) from above, return
# the corresponding eigenvector, normalized to 1.
function eigvec1(b,z::Number,m=length(b)+1)
    # "cheat" and use the fact that our eigenvector v must have a
    # nonzero first entries (since it is a quadrature weight), so we
    # can set v[1] = 1 to solve for the rest of the components:.
    v = Array{eltype(b)}(undef, m)
    v[1] = 1
    if m > 1
        s = v[1]
        v[2] = z * v[1] / b[1]
        s += v[2]^2
        for i = 3:m
            v[i] = - (b[i-2]*v[i-2] - z*v[i-1]) / b[i-1]
            s += v[i]^2
        end
        rmul!(v, 1 / sqrt(s))
    end
    return v
end

"""
    gauss([T,] N, a=-1, b=1)

Return a pair `(x, w)` of `N` quadrature points `x[i]` and weights `w[i]` to
integrate functions on the interval `(a, b)`,  i.e. `sum(w .* f.(x))`
approximates the integral.  Uses the method described in Trefethen &
Bau, Numerical Linear Algebra, to find the `N`-point Gaussian quadrature
in O(`N`²) operations.

`T` is an optional parameter specifying the floating-point type, defaulting
to `Float64`. Arbitrary precision (`BigFloat`) is also supported.
"""
function gauss(::Type{T}, N::Integer) where T<:AbstractFloat
    if N < 1
        throw(ArgumentError("Gauss rules require positive order"))
    end
    o = one(T)
    b = T[ n / sqrt(4n^2 - o) for n = 1:N-1 ]
    x = eignewt(b,N,N)
    w = T[ 2*eigvec1(b,x[i])[1]^2 for i = 1:N ]
    return (x, w)
end

gauss(N::Integer) = gauss(Float64, N) # integration on the standard interval (-1,1)

# re-scaled to an arbitrary interval:
gauss(N::Integer, a::Real, b::Real) = gauss(typeof(float(b-a)), N, a, b)
function gauss(::Type{T}, N::Integer, a::Real, b::Real) where T<:AbstractFloat
    x, w = gauss(T, N)
    s = T(b-a)/2
    x .= a .+ (x .+ 1) .* s
    w .*= abs(s)
    return (x, w)
end

"""
    kronrod([T,] n)

Compute `2n+1` Kronrod points `x` and weights `w` based on the description in
Laurie (1997), appendix A, simplified for `a=0`, for integrating on `[-1,1]`.
Since the rule is symmetric, this only returns the `n+1` points with `x <= 0`.
The function Also computes the embedded `n`-point Gauss quadrature weights `gw`
(again for `x <= 0`), corresponding to the points `x[2:2:end]`. Returns `(x,w,wg)`
in O(`n`²) operations.

`T` is an optional parameter specifying the floating-point type, defaulting
to `Float64`. Arbitrary precision (`BigFloat`) is also supported.

Given these points and weights, the estimated integral `I` and error `E` can
be computed for an integrand `f(x)` as follows:

    x, w, wg = kronrod(n)
    fx⁰ = f(x[end])                # f(0)
    x⁻ = x[1:end-1]                # the x < 0 Kronrod points
    fx = f.(x⁻) .+ f.((-).(x⁻))    # f(x < 0) + f(x > 0)
    I = sum(fx .* w[1:end-1]) + fx⁰ * w[end]
    if isodd(n)
        E = abs(sum(fx[2:2:end] .* wg[1:end-1]) + fx⁰*wg[end] - I)
    else
        E = abs(sum(fx[2:2:end] .* wg[1:end])- I)
    end
"""
function kronrod(::Type{T}, n::Integer) where T<:AbstractFloat
    if n < 1
        throw(ArgumentError("Kronrod rules require positive order"))
    end
    o = one(T)
    b = zeros(T, 2n+1)
    b[1] = 2*o
    for j = 1:div(3n+1,2)
        b[j+1] = j^2 / (4j^2 - o)
    end
    s = zeros(T, div(n,2) + 2)
    t = zeros(T, div(n,2) + 2)
    t[2] = b[n+2]
    for m = 0:n-2
        u = zero(T)
        for k = div(m+1,2):-1:0
            l = m - k + 1
            k1 = k + n + 2
            u += b[k1]*s[k+1] - b[l]*s[k+2]
            s[k+2] = u
        end
        s,t = t,s
    end
    for j = div(n,2):-1:0
        s[j+2] = s[j+1]
    end
    for m = n-1:2n-3
        u = zero(T)
        for k = m+1-n:div(m-1,2)
            l = m - k + 1
            j = n - l
            k1 = k + n + 2
            u -= b[k1]*s[j+2] - b[l]*s[j+3]
            s[j+2] = u
        end
        k = div(m+1,2)
        if 2k != m
            j = n - (m - k + 2)
            b[k+n+2] = s[j+2] / s[j+3]
        end
        s,t = t,s
    end
    for j = 1:2n
        b[j] = sqrt(b[j+1])
    end

    # get negative quadrature points x
    x = eignewt(b,2n+1,n+1) # x <= 0

    # get quadrature weights
    w = T[ 2*eigvec1(b,x[i],2n+1)[1]^2 for i in 1:n+1 ]

    # Get embedded Gauss rule from even-indexed points, using
    # the method described in Trefethen and Bau.
    for j = 1:n-1
        b[j] = j / sqrt(4j^2 - o)
    end
    gw = T[ 2*eigvec1(b,x[i],n)[1]^2 for i = 2:2:n+1 ]

    return (x, w, gw)
end

kronrod(N::Integer) = kronrod(Float64, N)

###########################################################################
# Type-stable cache of quadrature rule results, so that we don't
# repeat the kronrod calculation unnecessarily.

# precomputed n=7 rule in double precision (computed in 100-bit arithmetic),
# since this is the common case.
const xd7 = [-9.9145537112081263920685469752598e-01,
             -9.4910791234275852452618968404809e-01,
             -8.6486442335976907278971278864098e-01,
             -7.415311855993944398638647732811e-01,
             -5.8608723546769113029414483825842e-01,
             -4.0584515137739716690660641207707e-01,
             -2.0778495500789846760068940377309e-01,
             0.0]
const wd7 = [2.2935322010529224963732008059913e-02,
             6.3092092629978553290700663189093e-02,
             1.0479001032225018383987632254189e-01,
             1.4065325971552591874518959051021e-01,
             1.6900472663926790282658342659795e-01,
             1.9035057806478540991325640242055e-01,
             2.0443294007529889241416199923466e-01,
             2.0948214108472782801299917489173e-01]
const gwd7 = [1.2948496616886969327061143267787e-01,
              2.797053914892766679014677714229e-01,
              3.8183005050511894495036977548818e-01,
              4.1795918367346938775510204081658e-01]

# cache of T -> n -> (x,w,gw) Kronrod rules, to avoid recomputing them
# unnecessarily for repeated integration.   We initialize it with the
# default n=7 rule for double-precision calculations.  We use a cache
# of caches to allow us to evaluate the cache in a type-stable way with
# a generated function below.
const rulecache = Dict{Type,Dict}(
    Float64 => Dict{Int,NTuple{3,Vector{Float64}}}(7 => (xd7,wd7,gwd7)),
    Float32 => Dict{Int,NTuple{3,Vector{Float32}}}(7 => (xd7,wd7,gwd7)))

# for BigFloat rules, we need a separate cache keyed by (n,precision)
const bigrulecache = Dict{Tuple{Int,Int}, NTuple{3,Vector{BigFloat}}}()

cachedrule(::Type{T}, n::Integer) where {T<:Number} = cachedrule(typeof(float(real(one(T)))), n::Integer)

function cachedrule(::Type{BigFloat}, n::Integer)
    key = (n, precision(BigFloat))
    haskey(bigrulecache, key) ? bigrulecache[key] : (bigrulecache[key] = kronrod(BigFloat, n))
end

# use a generated function to make this type-stable
@generated function cachedrule(::Type{T}, n::Integer) where {T<:AbstractFloat}
    cache = haskey(rulecache, T) ? rulecache[T] : (rulecache[T] = Dict{Int,NTuple{3,Vector{T}}}())
    :(haskey($cache, n) ? $cache[n] : ($cache[n] = kronrod($T, n)))
end


# ------------------------------------------------------------
# Reconstruction of impulse responses, i.e., a single coefficient in the Haar
# basis from equidistant frequency samples

# Define scaling function and mother wavelet

function Haar_scaling(x::Real)
	if x >= 0 && x <= 1.0
		return 1.0
	else
		return 0.0
	end
end

function Haar_wavelet(x::Real)
	if x >= 0 && x < 0.5
		return 1.0
	elseif x >= 0.5 && x <= 1.0
		return -1.0
	else
		return 0.0
	end
end

# ------------------------------------------------------------
# Reconstruction of scaling function

# Input:
# M: The number of frequency samples
# epsislon: The distance between two consecutive samples
# N: The number of wavelet coefficients recovered

M = 32
epsilon = 1
#J = floor(Int, log2(M))
J = 2

# Sampling points
xi = [-M:M;]*epsilon

# Change of basis matrix
Tx, Ty = freq2Haar( xi, xi, J )
T = kron(Ty, Tx)

# Fourier observations
bx = FourHaarScaling(xi, J, 3)
by = FourHaarScaling(xi, J, 3)
b = bx * by.'

# Scaling function coefficients: 1 in the last (16th) entry, 0 elsewhere
y = pinv(T)*b[:] # Brute force
yy = pinv(Tx)*b*pinv(Ty).' # With Kronecker products


# p1.jl - convergence of fourth-order finite differences

# For various N, set up grid in [-pi,pi] and function u(x):
Nvec = 2.^(3:12);
clf(); axes([.1,.4,.8,.5]);
for N = Nvec
    h = 2*pi/N; x = -pi + (1:N)*h;
    u = @. exp(sin(x)^2);
    uprime = @. 2*sin(x)*cos(x)*u;

    # Construct sparse fourth-order differentiation matrix:
    e = ones(N);
    D = sparse(1:N,[2:N;1],2*e/3) - sparse(1:N,[3:N;1:2],e/12);
    D = (D-D')/h;

    # Plot max(abs(D*u-uprime)):
    error = norm(D*u-uprime,Inf);
    loglog(N,error,"k.",markersize=6);
end
grid(true); xlabel("N"); ylabel("error");
title("Convergence of fourth-order finite differences");
loglog(Nvec,1.0./Nvec.^4,"--");
text(105,5e-8,L"N^{-4}",fontsize=18);


function linear_index(IND)
  J = round(Int,floor(-.5 + .5 * sqrt(1 + 8 * (IND - 1))) + 2)
  I = round(Int,J .* (3 - J) / 2 + IND - 1)
  return I, J
end


### A Pluto.jl notebook ###
# v0.12.4

using Markdown
using InteractiveUtils

# ╔═╡ e1a3f480-12e7-11eb-01db-71003ec47f08
md"""
## What *is* an array, actually?

We have seen several examples of arrays.
When you think about an array, you probably think of a chain of boxes with the same shape, stored one after the other, consecutively in memory. And indeed this is the original meaning of the word "array".

However, as a user, we don't actually care about how the data is actually stored! All we really care about is the functionality of an array. Namely, we want the following properties:

- Elements have an integer index (from 1 to `n`);
- An element at position `i` can be accessed using the syntax `A[i]`;
- An element at position `i` can be stored using `A[i] = i` (if the array is mutable)
- We should be able to see the size of an array.
"""

# ╔═╡ 51a660b4-12f3-11eb-2b8a-f5544e650b1c
md"""
Julia provides an **abstraction** of an array called `AbstractArray`. We can define new types of array that "behave like" arrays but are stored differently.
"""

# ╔═╡ 853ac7d0-12f3-11eb-1b1c-993e336cf23a
md"""
## How does indexing work?
"""

# ╔═╡ 94bb1e6c-12f3-11eb-2e18-bb232d2c9cfe
md"""
First let's think about indexing, which is really the key operation on an array.
Let's define a vector:
"""

# ╔═╡ 314ba27a-12ec-11eb-12ce-67650c0020b7
arr = [34,45,67]

# ╔═╡ ab3afe6e-12f3-11eb-196a-79d26b66f681
md"""
We know that we can index into it using square brackets:
"""

# ╔═╡ b32acfda-12f3-11eb-323d-ef91b9ff9f0d
arr[2]

# ╔═╡ b53bc5ba-12f3-11eb-2b09-d38cc266b069
md"""
This is just syntax? What is actually going on? We can see this by looking at the **lowered** form:
"""

# ╔═╡ 3a202684-12ec-11eb-26fc-df7b7484e671
Meta.@lower arr[2]

# ╔═╡ c225d324-12f3-11eb-23d7-45548e7b32b6
md"""
We see that `arr[2]` is just "syntactic sugar" for calling the `getindex` function as `getindex(arr, 2)`; in other words, it's just a nice short and suggestive way to write the same thing.
"""

# ╔═╡ d92c5aa2-12f3-11eb-286e-8769ecdf03cb
md"""
So to implement our own objects that behave like arrays, we will define a new type and extend `getindex` on that type.
"""

# ╔═╡ f10cf87a-12f3-11eb-1104-0186b9e43757
md"""
## Example: A zero-based vector
"""

# ╔═╡ f7cadd08-12f3-11eb-1882-8dd4ce8ee2c2
md"""
There are interminable discussions online about whether zero-based indexing (i.e. array indices starting at 0) or one-based indexing (starting at 1) is superior. The answer is that *neither* is -- each is more suitable in different contexts.

Fortunately Julia is flexible enough to allow you to use *any* type of indexing that suits you -- see e.g. the `OffsetArrays.jl` package.

Let's see how to implement the simplest case, namely a zero-based vector.

So we define a new type and extend `getindex` on it. In order to use the full Julia machinery we need to extend a couple of other functions too:
"""

# ╔═╡ 1ac0095e-12eb-11eb-3df2-e98d0ff66dd5
struct ZeroFirstVector{T} <: AbstractVector{T}
	array::Vector{T}
end

# ╔═╡ a66afe14-12eb-11eb-3afc-fdac2325ff84
Base.size(vec::ZeroFirstVector) = size(vec.array)

# ╔═╡ e86c53fa-12eb-11eb-0ee4-41455a218d84
Base.axes(vec::ZeroFirstVector) = (0:length(vec.array)-1,)

# ╔═╡ 69b10842-12eb-11eb-381a-0bda698a59ed
Base.getindex(vec::ZeroFirstVector, i::Int) = vec.array[i+1]

# ╔═╡ 9c57e4aa-12eb-11eb-16fb-c7a18c238192
zarr = ZeroFirstVector([1, 3, 4])

# ╔═╡ 3cf51826-12f4-11eb-2405-db9616fa3152
md"""
Now we are able to use zero-based indexing!:
"""

# ╔═╡ fed5e2e4-12eb-11eb-169f-8b050f9eea4a
zarr[0]

# ╔═╡ 464c22f2-12f4-11eb-0b3c-872621c385f6
md"""
## Modifying elements: `setindex!`
"""

# ╔═╡ 4dec992e-12f4-11eb-1f20-b96632769bed
md"""
In a similar way to `getindex`, modifying elements of an array uses the mutating function `setindex!`:
"""

# ╔═╡ 5b8f3508-12ec-11eb-0a31-61ddf796ac74
Meta.@lower arr[2] = 3

# ╔═╡ 0c8ebad2-12ec-11eb-03f1-c5b9c79bb4da
Base.setindex!(vec::ZeroFirstVector, val, i::Int) = (vec.array[i+1] = val)

# ╔═╡ 8438424c-12ec-11eb-1de2-0bb3924ec78c
zarr[0] = 55

# ╔═╡ 9cee3206-12ec-11eb-09c7-1f0ebad11b8e
zarr

# ╔═╡ b0ab2844-12ec-11eb-1908-0318510b3292
md"""
### Julia gives you added bonuses here:
"""

# ╔═╡ bc8154b8-12ec-11eb-1fcc-4765fcd0796d
zarr[0:1]

# ╔═╡ 8c0c8188-12f4-11eb-1576-23958903a69a
@which zarr[0:1]

# ╔═╡ 7425ef6e-12f4-11eb-21c3-19c221c1ebdb
md"""
We never defined how to index using a *range*, only an integer. Julia has a generic fallback mechanism that understands how to do this once you have implemented `getindex` for a single element.

Similarly:
"""

# ╔═╡ c67fbbbc-12ec-11eb-32be-2ff720964fe6
zarr[ [2, 1, 0] ]

# ╔═╡ d13a65b6-12ec-11eb-2996-37987cfae782
zarr[0:1] .= [7,8];

# ╔═╡ dbb95f42-12ec-11eb-01d5-df2e21affd86
zarr

# ╔═╡ e6428aa6-12ec-11eb-06d0-95404a8a4039
zarr[ [2, 1, 0] ] .= [1, 2, 3];

# ╔═╡ e21658ec-12ef-11eb-06fd-b1d282edf7b6
zarr

# ╔═╡ b405b9ac-12f4-11eb-12fa-f7575ff1fa73
md"""
Various other functions also automatically work, e.g.:
"""

# ╔═╡ e57329a2-12ef-11eb-165d-51f080c73652
sort!(zarr)

# ╔═╡ e8b34188-12ef-11eb-2d7f-fde100b582d1
md"""
### Some more examples

"""

# ╔═╡ bd6bfd10-12f4-11eb-17fe-6f9ce83f0ce8
md"""
We have already seen some other examples of array types in the course, e.g. ranges:
"""

# ╔═╡ f0d952da-12ef-11eb-2452-471b184bc529
r = 53:83

# ╔═╡ 0fc7e3e6-12f0-11eb-1b49-e37196154ec6
size(r)

# ╔═╡ 03b23fe0-12f0-11eb-0970-33b3b973e902
r[20]

# ╔═╡ 00088cda-12f0-11eb-36b7-9b900fa6d0bf
typeof(r)

# ╔═╡ f6024924-12ef-11eb-1d6c-3ffff2546914
supertypes(typeof(r))

# ╔═╡ cdf3e532-12f4-11eb-1ce0-873bcf0dea58
md"""
Let's also look at what happens if we reshape an array:
"""

# ╔═╡ f9bf5200-12ef-11eb-30ee-9b3cbbccc886
A = reshape(1:12, 3, 4)

# ╔═╡ 07c1cdb6-12f1-11eb-0b31-15b5c88caac6
typeof(A)

# ╔═╡ 09d69238-12f1-11eb-1d9f-f72e8c6c14d9
supertypes(typeof(A))

# ╔═╡ ff598e1c-12f0-11eb-24f6-bdaa74db2ad7
size(A)

# ╔═╡ 2f23dd4e-12f0-11eb-01fc-53f66d2a55e4
Aᵀ = transpose(A)

# ╔═╡ ea8e0794-12f0-11eb-2194-bdc46224dc69
size(Aᵀ)

# ╔═╡ 2a76a87a-12f1-11eb-3e58-e5f919ba81e8
@which Aᵀ[2,3]

# ╔═╡ 892e7c40-12f0-11eb-2871-0d308063f025
md"The `getindex` method for `Transpose` is defined [here](https://github.com/JuliaLang/julia/blob/master/stdlib/LinearAlgebra/src/adjtrans.jl#L190)"

# ╔═╡ da41895c-12f4-11eb-2ee4-f1ffc4344f61
md"""
## Conclusions
"""

# ╔═╡ dfac92ea-12f4-11eb-1253-b74be632de64
md"""
Summarising, we see that we -- and Julia -- can think of an array as "anything that behaves like an array", i.e. which you can index into.

What does it mean for something to behave like an array? In the end it means that we *could* "materialise" it to produce an actual array that did live as data stored linearly in memory, and our program would behave in exactly the same way.

But Julia provides a mechanism by which we can store less data, or store it in a more convenient way, and have indexing be almost as efficient.
"""

# ╔═╡ Cell order:
# ╟─e1a3f480-12e7-11eb-01db-71003ec47f08
# ╟─51a660b4-12f3-11eb-2b8a-f5544e650b1c
# ╟─853ac7d0-12f3-11eb-1b1c-993e336cf23a
# ╟─94bb1e6c-12f3-11eb-2e18-bb232d2c9cfe
# ╠═314ba27a-12ec-11eb-12ce-67650c0020b7
# ╟─ab3afe6e-12f3-11eb-196a-79d26b66f681
# ╠═b32acfda-12f3-11eb-323d-ef91b9ff9f0d
# ╟─b53bc5ba-12f3-11eb-2b09-d38cc266b069
# ╠═3a202684-12ec-11eb-26fc-df7b7484e671
# ╟─c225d324-12f3-11eb-23d7-45548e7b32b6
# ╟─d92c5aa2-12f3-11eb-286e-8769ecdf03cb
# ╟─f10cf87a-12f3-11eb-1104-0186b9e43757
# ╟─f7cadd08-12f3-11eb-1882-8dd4ce8ee2c2
# ╠═1ac0095e-12eb-11eb-3df2-e98d0ff66dd5
# ╠═a66afe14-12eb-11eb-3afc-fdac2325ff84
# ╠═e86c53fa-12eb-11eb-0ee4-41455a218d84
# ╠═69b10842-12eb-11eb-381a-0bda698a59ed
# ╠═9c57e4aa-12eb-11eb-16fb-c7a18c238192
# ╟─3cf51826-12f4-11eb-2405-db9616fa3152
# ╠═fed5e2e4-12eb-11eb-169f-8b050f9eea4a
# ╟─464c22f2-12f4-11eb-0b3c-872621c385f6
# ╟─4dec992e-12f4-11eb-1f20-b96632769bed
# ╠═5b8f3508-12ec-11eb-0a31-61ddf796ac74
# ╠═0c8ebad2-12ec-11eb-03f1-c5b9c79bb4da
# ╠═8438424c-12ec-11eb-1de2-0bb3924ec78c
# ╠═9cee3206-12ec-11eb-09c7-1f0ebad11b8e
# ╟─b0ab2844-12ec-11eb-1908-0318510b3292
# ╠═bc8154b8-12ec-11eb-1fcc-4765fcd0796d
# ╠═8c0c8188-12f4-11eb-1576-23958903a69a
# ╟─7425ef6e-12f4-11eb-21c3-19c221c1ebdb
# ╠═c67fbbbc-12ec-11eb-32be-2ff720964fe6
# ╠═d13a65b6-12ec-11eb-2996-37987cfae782
# ╠═dbb95f42-12ec-11eb-01d5-df2e21affd86
# ╠═e6428aa6-12ec-11eb-06d0-95404a8a4039
# ╠═e21658ec-12ef-11eb-06fd-b1d282edf7b6
# ╟─b405b9ac-12f4-11eb-12fa-f7575ff1fa73
# ╠═e57329a2-12ef-11eb-165d-51f080c73652
# ╟─e8b34188-12ef-11eb-2d7f-fde100b582d1
# ╟─bd6bfd10-12f4-11eb-17fe-6f9ce83f0ce8
# ╠═f0d952da-12ef-11eb-2452-471b184bc529
# ╠═0fc7e3e6-12f0-11eb-1b49-e37196154ec6
# ╠═03b23fe0-12f0-11eb-0970-33b3b973e902
# ╠═00088cda-12f0-11eb-36b7-9b900fa6d0bf
# ╠═f6024924-12ef-11eb-1d6c-3ffff2546914
# ╟─cdf3e532-12f4-11eb-1ce0-873bcf0dea58
# ╠═f9bf5200-12ef-11eb-30ee-9b3cbbccc886
# ╠═07c1cdb6-12f1-11eb-0b31-15b5c88caac6
# ╠═09d69238-12f1-11eb-1d9f-f72e8c6c14d9
# ╠═ff598e1c-12f0-11eb-24f6-bdaa74db2ad7
# ╠═2f23dd4e-12f0-11eb-01fc-53f66d2a55e4
# ╠═ea8e0794-12f0-11eb-2194-bdc46224dc69
# ╠═2a76a87a-12f1-11eb-3e58-e5f919ba81e8
# ╟─892e7c40-12f0-11eb-2871-0d308063f025
# ╟─da41895c-12f4-11eb-2ee4-f1ffc4344f61
# ╟─dfac92ea-12f4-11eb-1253-b74be632de64


# estimate with MCMC (Turing) for conjugate prior
# compare with jlbayes_conjugate_regression.jl

using Turing
using Distributions
using Random
using LinearAlgebra
using StatsPlots
using Optim

# Generate moc data
n = 50
dim = 1
Random.seed!(99)
u = rand(Normal(0, 0.7), n) # true value σ² = 0.49
x = rand(Uniform(-sqrt(3.0), sqrt(3.0)), n, dim) # x generated from Uniform[-√3, √3]
X = [ones(n) x] # dependent variable matrix with ones
β = Vector(1:1:dim+1) # true coefficients
y = X*β + u # true data generating process

# set prior (conjugate prior for camparison with analytic result)
# p.78 eq (3.32)
k = size(X)[2]
β0 = zeros(k)
τ0 = 0.2
A0 = τ0 * I # Identity matrix
ν0 = 5.0
λ0 = 7.0
# H0 = (λ0 / ν0) * inv(A0) # scale matrix for marginal prior of multivariate coefficient (mv Tdist)
# h0 = diag(sqrt(Matrix(H0, k, k))) # scale params for marginal prior of single coefficient (T dist) H0の対角成分
sd0 = diag(sqrt(Matrix(A0, k, k))) # the parameter for conditional prior for β|σ² ~ N_k(β0, σ²*inv(A0)), (inv(A0) part)

@model function linear_regression(X, y)
    # conjugate priors
    # https://turing.ml/dev/tutorials/05-linear-regression/

    # Set variance priors
    σ² ~ InverseGamma(ν0/2, λ0/2)

    # Set the priors on our coefficients.
    nfeatures = size(X, 2)
    β ~ MvNormal(β0, σ² * inv(A0))

    # Write likelihood.
    y ~ MvNormal(X * β, σ² * I)
end

# Draw random number from posterior.
model = linear_regression(X, y)
n_draws = 5_000
n_chains = 4
n_tune = 1_000
# sample(model, sampler, parallel_type, n, n_chains)
# chain = sample(model, NUTS(), MCMCThreads(), n_draws, n_chains, discard_initial=n_tune)
# How to implement burn-in / warmup? --> discard initial value automatically. ; discard_adapt=false to turn off
chn = sample(model, NUTS(), MCMCThreads(), n_draws, n_chains)
chn = sample(model, NUTS(), MCMCThreads(), n_draws, n_chains; discard_adapt=false)

# Chain info.
typeof(chn)
size(chn)
describe(chn)
summarize(chn)
quantile(chn)
plot(chn)

# Indexing a Chains object.
β1_chn = chn["β[1]"]
β2_chn = chn[Symbol("β[2]")]
σ²_chn = chn[:σ²]
all_β_chn = group(chn, :β) # get subset of all parameters include "β"

typeof(β1_chn)
β1_chn[iter = 1000:5000, chain = 1]
mean(β1_chn)
std(β1_chn)

# Monte Carlo Standard Error
mcse(θ_chn) = sqrt((1 / (length(θ_chn) * (length(θ_chn) - 1))) * sum((θ_chn .- mean(θ_chn)).^2))

[β1_chn, β2_chn, σ²_chn]

mcse.([β1_chn, β2_chn, σ²_chn])

# without burn-in
mean(β1_chn[iter = 1:5000])

histogram(β1_chn[:, 1])
histogram!(β1_chn[:, 2])
histogram!(β1_chn[:, 3])
histogram!(β1_chn[:, 4])

# mode (MLE and MAP) estimates (Optimization with Optim.jl) as an option (continuous only)
# Note that loading Optim explicitly is required for mode estimation to function,
# as Turing does not load the opimization suite unless Optim is loaded as well.

# Generate a MLE estimate (default optimizer is LBFGS optimizer).
mle_estimate = optimize(model, MLE())
# Use Newton
# mle_estimate = optimize(model, MLE(), Newton())

# Generate a MAP estimate (default optimizer is LBFGS optimizer).
map_estimate = optimize(model, MAP())

# Analyze mode estimate
# Import StatsBase to use it's statistical methods.
# using StatsBase

# Print out the coefficient table.
# coeftable(mle_estimate)

# Analytical Solution
inv(X' * X) * X' * y


using LinearAlgebra
import LinearAlgebra.Eigen
using SparseArrays
#this lib will be used to estimate the executon time
using BenchmarkTools
using Test
using Distributions

struct SVDCM
	U::Matrix{Float64}
	S::SparseMatrixCSC{Float64, Int64}
	V::Matrix{Float64}

	function SVDCM(A::Matrix{Float64})
		U, S, V = svdCm(A)
		return new( U, S, V)
	end
	function svdCm(A::Matrix{Float64})

		r,c = size(A)
		shorter=min(r,c)
		longer=max(r,c)

		Asq = zeros(shorter,shorter)
		mul!(Asq, A', A)
		a, V = eigen!(Asq)

		S = sparse(zeros(longer,shorter))
		Sin = sparse(zeros(shorter,longer))

		S[diagind(S)] .= (sqrt.(a))
		Sin[diagind(Sin)] .= 1 ./ S[diagind(S)]

		U1 = zeros(r,c)
		U = zeros(longer,longer)

		mul!(U1, A, V)
		mul!(U,U1,Sin)
		return U, S, V
	end
end

R = rand(Uniform(1., 100.), 10000,1000)
MM = [1 0 0 0 2; 0 0 3 0 0; 0 0 0 0 0; 0 2 0 0 0; 1 0 0 0 2; 0 0 3 0 0; 0 0 0 0 0; 0 2 0 0 0; 1 0 0 0 2; 0 0 3 0 0; 0 0 0 0 0; 0 2 0 0 0; 1 0 0 0 2; 0 0 3 0 0; 0 0 0 0 0; 0 2 0 0 0;1 0 0 0 2; 0 0 3 0 0; 0 0 0 0 0; 0 2 0 0 0;1 0 0 0 2; 0 0 3 0 0; 0 0 0 0 0; 0 2 0 0 0 ]

C = R
println("size of matrix: ", size(C) )

println("Julia SVD time:")
SVD = @time svd(C)

println("SVD time:")
mySvd = @time SVDCM(C)

F = mySvd.U * mySvd.S * mySvd.V'
#display(F)
@test F ≈ C atol=1e-5


# Legendre-Gauss quadrature weights and nodes
# This code was originally written for Matlab
#  by Greg von Winckel - 02/25/2004
# posted on Mathworks website and licensed under BSD License
# http://fr.mathworks.com/matlabcentral/fileexchange/4540-legendre-gauss-quadrature-weights-and-nodes
#
#

function lgwt(N,a,b)

	N=N-1;
	N1=N+1; N2=N+2;
	xu= transpose(range(-1,stop = 1, length = N1));

	y = transpose(cos.((2*transpose(0:N).+1)*pi/(2*N+2))+(0.27/N1)*sin.(pi*xu*N/N2));

	L = zeros(N1,N2);

	Lp=zeros(N1,N2);
	y0=2;

	while maximum(abs.(y.-y0))>eps()
		L[:,1] .= 1;
		Lp[:,1] .= 0;
		L[:,2] = y;
		for k=2:N1
			L[:,k+1]=( (2*k-1)*y.*L[:,k]-(k-1)*L[:,k-1] )/k;
		end
		Lp=(N2)*( L[:,N1] - y .* L[:,N2] )./(1 .- y.^2);
		y0=y;
		y=y0-L[:,N2]./Lp;
	end

	x = (a*(1 .- y)+b*(1 .+ y))/2;

	w=(b-a)./((1 .- y.^2).*Lp.^2)*(N2/N1)^2;

	return x,w
end


function house!(h, j, X, k)
    T = eltype(h)
    n = size(X,1)
    σ = zero(T)
    @inbounds @simd for i in j+1:n
        σ += X[i,k]^2
    end
    @inbounds begin
        h[j] = one(T)
        @simd for i in j+1:n
            h[i] = X[i,k]
        end
    end
    if σ == zero(T)
        β = zero(T)
    else
        @inbounds begin
            μ = sqrt(X[j,k]^2 + σ)
            if X[j,k] <= zero(T)
                h[j] = X[j,k] - μ
            else
                h[j] = -σ/(X[j,k] + μ)
            end
            β = 2*(h[j]^2) / (σ + h[j]^2)
            d = h[j]
            @simd for i in j:n
                h[i] /= d
            end
        end
    end
    β
end

"""
accumulate!(Q, A, β)

Q: Matrix to be over-written with Q, gets returned
A: Has the Householder vectors on the lower triangle
β: Has the Householder coefficients
"""
function accumulate!(Q, A, β)
    T = eltype(Q)
    m, n = size(Q)
    @inbounds begin
        Q .= zero(T)
        @simd for i in 1:n
            Q[i,i] = one(T)
        end
    end
    @inbounds for j in n:-1:1
        @simd for t in j:n
            dp = Q[j,t]
            for i in j+1:m
                dp += A[i,j]*Q[i,t]
            end
            Q[j,t] -= β[j]*dp
            for i in j+1:m
                Q[i,t] -= β[j]*dp*A[i,j]
            end
        end
    end
    Q
end

"""
getR!(R, A)

Copies the right triangular part of A into the right triangular part of R.
"""
function getR!(R, A)
    T = eltype(R)
    m, n = size(A)
    @inbounds @simd for j in 1:n
        for i in 1:j
            R[i,j] = A[i,j]
        end
    end
    R
end

"""
qrfact!(A, β, h=zeros(eltype(A), size(A,1)))

A: Matrix to be QR factorized, returned by the function
β: Empty buffer to store the Householder coefficients, returned by the function
h: Empty buffer for in-place operations
"""
function qrfact!(A, β, h=zeros(eltype(A), size(A,1)))
    m, n = size(A)
    @inbounds for j in 1:n
        β[j] = house!(h, j, A, j)
        @simd for k in j:n
            dp = A[j,k]
            for i in j+1:m
                dp += h[i]*A[i,k]
            end
            for l in j:m
                A[l,k] -= (β[j] * dp) * h[l]
            end
        end
        if j < m
            @simd for k in 1:(m-j)
                A[k+j,j] = h[k+j]
            end
        end
    end
    A, β
end


"Evaluate the Laplace 2D single layer potential kernel"
function laplace_slp(x, y)
    T = eltype(x)
    if norm(x-y) < eps(T)
         return zero(Complex{T})
    else
        return complex(1/2/pi*log(norm(x-y)))
    end
end

export Laplace_SLP_2D
"The single layer potential kernel of the Laplace equation."
struct Laplace_SLP_2D{T} <: BoundaryKernel
end

Laplace_SLP_2D() = Laplace_SLP_2D{Float64}()

(kernel::Laplace_SLP_2D)(t, tau, param, x, y) = laplace_slp(x, y)
(kernel::Laplace_SLP_2D)(x, tau, param, y) = laplace_slp(x, y)

is_symmetric(::Laplace_SLP_2D) = true

singularity(::Laplace_SLP_2D) = LogSingularDiagonal()

"Evaluate the Laplace 2D double layer potential kernel"
laplace_dlp(x, y, normal_y, z = norm(x-y)) =
    (normal_y' * (x-y)) / z^2/2/pi

"Evaluate the Laplace 2D double layer potential kernel"

function laplace_dlp_kernel(x, y, tau, param)
    z = norm(x-y)
    if abs(z) > eps(eltype(x))
        return Complex{Float64}(laplace_dlp(x, y, normal(param, tau), z))
    else
        return Complex{Float64}(0)
    end
end

export Laplace_DLP_2D

"The double layer potential kernel of the Laplace equation in 2D."
struct Laplace_DLP_2D{T} <: BoundaryKernel
end

Laplace_DLP_2D() = Laplace_DLP_2D{Float64}()

BasisFunctions.name(kernel::Laplace_DLP_2D) = "2D Laplace double layer potential kernel"

(kernel::Laplace_DLP_2D)(t, tau, param, x, y) =
    laplace_dlp_kernel(x, y, tau, param)
(kernel::Laplace_DLP_2D)(x, tau, param, y) =
    laplace_dlp_kernel(x, y, tau, param)

is_symmetric(::Laplace_DLP_2D) = false

# The kernel is continuous but its derivative is singular
singularity(::Laplace_DLP_2D) = LogSingularDiagonal()

"Adjoint Double layer potential for Laplace equation"

export Laplace_adjDLP_2D
"The adjoint double layer potential kernel of the Laplace equation in 2D."
struct Laplace_adjDLP_2D{T} <: BoundaryKernel
end

laplace_adjdlp(x, y, normal_x, z = norm(x-y)) =
    (normal_x' * (y-x)) / z^2/2/pi

"Evaluate the adjoint of the Laplace 2D double layer potential kernel"

function laplace_adjdlp_kernel(x, y, t, param)
    z = norm(x-y)
    if abs(z) > 10*eps(eltype(x))
        return Complex{Float64}(laplace_adjdlp(x, y, normal(param, t), z))
    else
        return Complex{Float64}(0)
    end
end

Laplace_adjDLP_2D() = Laplace_adjDLP_2D{Float64}()

BasisFunctions.name(kernel::Laplace_adjDLP_2D) = "Adjoint of 2D Laplace double layer potential kernel"

(kernel::Laplace_adjDLP_2D)(t, tau, param, x, y) =
    laplace_adjdlp_kernel(x, y, t, param)

is_symmetric(::Laplace_adjDLP_2D) = false

# The kernel is continuous but its derivative is singular
singularity(::Laplace_adjDLP_2D) = LogSingularDiagonal()
## Some boundary conditions

export laplace_bcond_harmonic, laplace_param_bcond_harmonic, laplace_parboundary_condition_potential_flow

harmonic_poly(z) = z^2-z+2

boundary_condition_harmonic(x) = real(harmonic_poly(x[1]+im*x[2]))

"Return a field function that evaluates to a plane wave."
laplace_bcond_harmonic() = (x,y) -> boundary_condition_harmonic(SVector(x,y))

"Return a parametric boundary function that evaluates to a plane wave."
laplace_param_bcond_harmonic(param) =
    t -> boundary_condition_harmonic(applymap(param, t))

"Return a parametric boundary function corresponding to the unperturbed flow at r=∞"
laplace_parboundary_condition_potential_flow(param, velocity) =
    t -> velocity'*normal(param, t)


using Images
using ImageView
using MAT
using LinearAlgebra
using Statistics
using Plots

vars = matread("../../DATA/allFaces.mat")

n = Int(vars["n"])
m = Int(vars["m"])
faces = vars["faces"]
nfaces = Int.(vars["nfaces"])

# We use the first 36 people for training data
trainingFaces = faces[:,1:sum(nfaces[1:36])];
avgFace = mean(trainingFaces, dims = 2) # size n*m by 1

# Compute eigenfaces on mean-subtracted training data
X = trainingFaces - avgFace * ones(1, size(trainingFaces)[2])
U, S, V = svd(X)

imshow(reshape(avgFace, n, m)) # Show avg face
imshow(reshape(U[:,1], n, m))  # Show first eigenface

# Now show eigenface reconstruction of image that was omitted from test set

testFace = faces[:,1+sum(nfaces[1:36])] # First face of person 37
imshow(reshape(testFace, n, m))

testFaceMS = testFace - avgFace
for r = [25, 50, 100, 200, 400, 800, 1600]
    reconFace = avgFace + (U[:,1:r]*(U[:,1:r]'*testFaceMS))
    imshow(reshape(reconFace, n, m))
end

# Project person 2 and 7 onto PC5 and PC6

P1num = 2; # Person number 2
P2num = 7; # Person number 7

P1 = faces[:,1+sum(nfaces[1:P1num-1]):sum(nfaces[1:P1num])]
P2 = faces[:,1+sum(nfaces[1:P2num-1]):sum(nfaces[1:P2num])]

P1 = P1 - avgFace * ones(1, size(P1)[2])
P2 = P2 - avgFace * ones(1, size(P2)[2])

PCAmodes = [5, 6] # Project onto PCA modes 5 and 6
PCACoordsP1 = U[:,PCAmodes]'*P1
PCACoordsP2 = U[:,PCAmodes]'*P2

p1 = scatter(PCACoordsP1[1,:], PCACoordsP1[2,:], mc = :blue, legend = false)
scatter!(PCACoordsP2[1,:], PCACoordsP2[2,:], mc = :red)


"""
    compute_linear_return_array(priceArray::Array{Float64,1})

Stuff will go here. Awesome stuff, the most beautiful stuff ever.
"""
function compute_linear_return_array(priceArray::Array{Float64,1})::PSResult

    # initialize -
    linear_return_array = Array{Float64,1}()

    # get the size -
    number_of_time_steps = length(priceArray)
    for time_index = 2:number_of_time_steps
        value = (priceArray[time_index] / priceArray[time_index - 1]) - 1
        push!(linear_return_array, value)
    end

    # return -
    return PSResult(linear_return_array)
end

"""
    compute_linear_return_array(dataTable::DataFrame; key::Symbol = Symbol("adjusted_close"))

Stuff will go here. Awesome stuff, the most beautiful stuff ever.
"""
function compute_linear_return_array(dataTable::DataFrame; key::Symbol=Symbol("adjusted_close"))

    # initialize -
    price_array = Float64[]

    # iterate to build the price array -
    (number_of_rows, number_of_cols) = size(dataTable)
    for row_index = 1:number_of_rows
        data_value = dataTable[row_index,key]
        push!(price_array, data_value)
    end

    # return -
    return compute_linear_return_array(price_array)
end

"""
    compute_log_return_array(priceArray::Array{Float64})

Stuff will go here. Awesome stuff, the most beautiful stuff ever.
"""
function compute_log_return_array(priceArray::Array{Float64})::PSResult

    # initialize -
    log_return_array = Array{Float64,1}()

    # get the size -
    number_of_time_steps = length(priceArray)
    for time_index = 2:number_of_time_steps
        value = log((priceArray[time_index] / priceArray[time_index - 1]))
        push!(log_return_array, value)
    end

    # return -
    return PSResult(log_return_array)
end

"""
    compute_log_return_array(dataTable::DataFrame; key::Symbol = Symbol("adjusted_close"))

Stuff will go here. Awesome stuff, the most beautiful stuff ever.
"""
function compute_log_return_array(dataTable::DataFrame; key::Symbol=Symbol("adjusted_close"))::PSResult

    # initialize -
    price_array = Float64[]

    # iterate to build the price array -
    (number_of_rows, number_of_cols) = size(dataTable)
    for row_index = 1:number_of_rows
        data_value = dataTable[row_index,key]
        push!(price_array, data_value)
    end

    # return -
    return compute_log_return_array(price_array)
end

"""
    compute_return_volatility(dataTable::DataFrame; returnCalcFunction::Function=compute_linear_return_array,
        key::Symbol = Symbol("adjusted_close"))

Stuff will go here. Awesome stuff, the most beautiful stuff ever.
"""
function compute_return_volatility(dataTable::DataFrame; returnCalcFunction::Function=compute_linear_return_array,
    key::Symbol=Symbol("adjusted_close"))::PSResult

    # initialize -
    price_array = Float64[]

    # iterate to build the price array -
    (number_of_rows, number_of_cols) = size(dataTable)
    for row_index = 1:number_of_rows
        data_value = dataTable[row_index,key]
        push!(price_array, data_value)
    end

    # we have the difference -
    result = returnCalcFunction(price_array)
    if (isa(result.value, Exception) == true)
        return result
    end
    price_return_array = result.value

    # compute the std -
    volatlity = std(price_return_array)

    # return -
    return PSResult(volatlity)
end

"""
    compute_return_volatility(priceArray::Array{Float64};
        returnCalcFunction::Function=compute_linear_return_array)

Stuff will go here. Awesome stuff, the most beautiful stuff ever.
"""
function compute_return_volatility(priceArray::Array{Float64};
    returnCalcFunction::Function=compute_linear_return_array)::PSResult

    # compute the return array -
    result = returnCalcFunction(priceArray)
    if (isa(result.value, Exception) == true)
        return result
    end
    volatlity = std(result.value)

    # return -
    return PSResult(volatlity)
end


function decompose_vector(v::Vector3, k::Vector3)
    v_parallel = (v ⋅ k) * k
    v_perpendicular = v - v_parallel
    v_parallel, v_perpendicular
end

function rotate_around(v::Vector3, k::Vector3, φ::Real)
    k = normalize(k)
    v_parallel, v_perpendicular = decompose_vector(v, k)
    ω = k × v_perpendicular
    v_parallel + v_perpendicular * cos(φ) + ω * sin(φ)
end

@_with_default_rng function random_3d_direction()
    θ = 2π * rand(rng)
    z = 2 * rand(rng) - 1
    r = √(1 - z ^ 2)
    Vector3(r * cos(θ), r * sin(θ), z)
end

"""
    random_direction_on_plane([rng], normal_vector)
    random_direction_on_plane([rng], u, v)

Generate a random direction on a plane defined by a normal vector or two vectors on the plane.
"""
random_direction_on_plane(rng::AbstractRNG, u::Vector3, v::Vector3) = let φ = 2π * rand(rng)
    w = cos(φ) * u + sin(φ) * v
end
@_with_default_rng function random_direction_on_plane(normal_vector::Vector3)
    a, b, c = normal_vector
    u = normalize(Vector3(b - c, -a + c, a - b))
    v = normalize(normal_vector × u)
    random_direction_on_plane(rng, u, v)
end


function find_closest_element(arr::Array{Int}, k::UInt, x::Int)
    larger = findfirst(v -> v>=x, arr)
    smaller = larger

    #
    len = lastindex(arr)
    larger isa Nothing && begin
        return arr[(len - k + 1):len]
        end
    smaller == 1 && return arr[1:k]

    #
    smaller -= 1
    count = 0

    while true
        smaller < 1 && return arr[1:k]
        larger > len && return arr[(len - k + 1):len]

        count == k && break

        if x - arr[smaller] <= arr[larger] - x
            smaller -= 1
        else
            larger += 1
        end
        count += 1
    end

    arr[smaller+1:larger-1]
end


# using OrdinaryDiffEq
using ModelingToolkit
using DataDrivenDiffEq
using LinearAlgebra, DiffEqSensitivity, Optim
using DiffEqFlux, Flux
using Plots
using DifferentialEquations
gr()

function lotka(du, u, p, t)
    α, β, γ, δ = p
    du[1] = α*u[1] - β*u[2]*u[1]
    du[2] = γ*u[1]*u[2]  - δ*u[2]
end

# Define the experimental parameter
tspan = (0.0f0,3.0f0)
u0 = Float32[0.44249296,4.6280594]
p_ = Float32[1.3, 0.9, 0.8, 1.8]
prob = ODEProblem(lotka, u0,tspan, p_)
solution = solve(prob, Vern7(), abstol=1e-12, reltol=1e-12, saveat = 0.1)

scatter(solution, alpha = 0.25)
plot!(solution, alpha = 0.5)

# Ideal data
X = Array(solution)
# Add noise to the data
println("Generate noisy data")
Xₙ = X + Float32(1e-3)*randn(eltype(X), size(X))

# Define the neueral network which learns L(x, y, y(t-τ))
# Actually, we do not care about overfitting right now, since we want to
# extract the derivative information without numerical differentiation.
L = FastChain(FastDense(2, 32, tanh),FastDense(32, 32, tanh), FastDense(32, 2))
p = initial_params(L)

function dudt_(u, p,t)
    x, y = u
    z = L(u,p)
    [p_[1]*x + z[1],
    -p_[4]*y + z[2]]
end

prob_nn = ODEProblem(dudt_,u0, tspan, p)
sol_nn = concrete_solve(prob_nn, Tsit5(), u0, p, saveat = solution.t)

plot(solution)
plot!(sol_nn)

function predict(θ)
    Array(concrete_solve(prob_nn, Tsit5(), u0, θ, saveat = solution.t,
                         abstol=1e-6, reltol=1e-6,
                         sensealg = InterpolatingAdjoint(autojacvec=ReverseDiffVJP())))
    # @show Array(concrete_solve(prob_nn, Vern7(), u0, θ, saveat = solution.t,
    #                                           abstol=1e-6, reltol=1e-6,
    #                                           sensealg = InterpolatingAdjoint(autojacvec=ReverseDiffVJP())))
  #   tmp_prob = remake(prob_nn,u0=u0,p=θ)
  # tmp_sol =  solve(tmp_prob, Rodas4(), saveat = solution.t,
  #                                           abstol=1e-6, reltol=1e-6,
  #                                           sensealg = InterpolatingAdjoint(autojacvec=ReverseDiffVJP()))
  # # @show size(Array(tmp_sol))
  # pred =  Array(tmp_sol)
end

# No regularisation right now
function loss(θ)
    pred = predict(θ)
    sum(abs2, Xₙ .- pred), pred
end

# Test
loss(p)

const losses = []

callback(θ,l,pred) = begin
    @show l
    push!(losses, l)
    if length(losses)%50==0
        println("Current loss after $(length(losses)) iterations: $(losses[end])")
    end
    false
end

# First train with ADAM for better convergence
res1 = DiffEqFlux.sciml_train(loss, p, ADAM(0.5), cb=callback, maxiters = 100)

# Train with BFGS
res2 = DiffEqFlux.sciml_train(loss, res1.minimizer, BFGS(initial_stepnorm=0.01), cb=callback, maxiters = 10)

println("Final training loss after $(length(losses)) iterations: $(losses[end])")

# Plot the losses
plot(losses, yaxis = :log, xaxis = :log, xlabel = "Iterations", ylabel = "Loss")

# Plot the data and the approximation
NNsolution = predict(res2.minimizer)
# Trained on noisy data vs real solution
plot(solution.t, NNsolution')
plot!(solution.t, X')

# Ideal derivatives
DX = Array(solution(solution.t, Val{1}))

prob_nn2 = ODEProblem(dudt_,u0, tspan, res2.minimizer)
_sol = solve(prob_nn2, Tsit5())
DX_ = Array(_sol(solution.t, Val{1}))

# The learned derivatives
plot(DX')
plot!(DX_')

# Ideal data
L̄ = [-p_[2]*(X[1,:].*X[2,:])';p_[3]*(X[1,:].*X[2,:])']
# Neural network guess
L̂ = L(Xₙ,res2.minimizer)

scatter(L̄')
plot!(L̂')

# Plot the error
scatter(abs.(L̄-L̂)', yaxis = :log)

## Sparse Identification

# Create a Basis
@variables u[1:2]
# Lots of polynomials
polys = Operation[1]

for i ∈ 1:5
    push!(polys, u[1]^i)
    push!(polys, u[2]^i)
    for j ∈ i:5
        if i != j
            push!(polys, (u[1]^i)*(u[2]^j))
            push!(polys, u[2]^i*u[1]^i)
        end
    end
end

# And some other stuff
h = [cos.(u)...; sin.(u)...; polys...]
basis = Basis(h, u)

# Create an optimizer for the SINDy problem
opt = SR3()
# Create the thresholds which should be used in the search process
λ = exp10.(-7:0.1:3)
# Target function to choose the results from; x = L0 of coefficients and L2-Error of the model
g(x) = x[1] < 1 ? Inf : norm(x, 2)
# Test on original data and without further knowledge
println("SINDy on full ideal, unavailable data")
Ψ = SINDy(Xₙ[:, :], DX[:, :], basis, λ, opt, g = g, maxiter = 10000) # Fail
println(Ψ)
print_equations(Ψ)

# Test on ideal derivative data for unknown function ( not available )
println("SINDy on partial ideal, unavailable data")
Ψ = SINDy(Xₙ[:, 1:end], L̄[:, 1:end], basis, λ,opt, g = g, maxiter = 10000) # Succeed
println(Ψ)
print_equations(Ψ)

# Test on uode derivative data
println("SINDy on learned, partial, available data")
Ψ = SINDy(Xₙ[:, 2:end], L̂[:, 2:end], basis, λ,  opt, g = g, maxiter = 10000, normalize = true, denoise = true) # Succeed
println(Ψ)
print_equations(Ψ)

# Extract the parameter
p̂ = parameters(Ψ)
println("First parameter guess : $(p̂)")

# The parameters are a bit off, but the equations are recovered
# Start another SINDy run to get closer to the ground truth
# Create function
unknown_sys = ODESystem(Ψ)
unknown_eq = ODEFunction(unknown_sys)

# Just the equations
b = Basis((u, p, t)->unknown_eq(u, [1.; 1.], t), u)

# Retune for better parameters -> we could also use DiffEqFlux or other parameter estimation tools here.
Ψf = SINDy(Xₙ[:, 2:end], L̂[:, 2:end], b, STRRidge(0.01), maxiter = 100, convergence_error = 1e-18) # Succeed
println(Ψf)
p̂ = parameters(Ψf)
println("Second parameter guess : $(p̂)")

# Create function
recovered_sys = ODESystem(Ψf)
recovered_eq = ODEFunction(recovered_sys)

# Build a ODE for the estimated system
function dudt(du, u, p, t)
    # Add SINDy Term
    α, δ, β, γ = p
    z = recovered_eq(u, [β; γ], t)
    du[1] = α*u[1] + z[1]
    du[2] = -δ*u[2] + z[2]
end

# Create the approximated problem and solution
ps = [p_[[1,4]]; p̂]
approximate_prob = ODEProblem(dudt, u0, tspan, ps)
approximate_solution = solve(approximate_prob, Tsit5(), saveat = 0.01)

# Plot
plot(solution)
plot!(approximate_solution)

## Simulation

# Look at long term prediction
t_long = (0.0, 50.0)
approximate_prob = ODEProblem(dudt, u0, t_long, ps)
approximate_solution_long = solve(approximate_prob, Tsit5(), saveat = 0.1) # Using higher tolerances here results in exit of julia
plot(approximate_solution_long)

true_prob = ODEProblem(lotka, u0, t_long, p_)
true_solution_long = solve(true_prob, Tsit5(), saveat = approximate_solution_long.t)
plot!(true_solution_long)

using JLD2
@save "knowledge_enhanced_NN.jld2" solution recovered_sys approximate_solution_long NNsolution L true_solution_long Xₙ L̄ L̂
@load "knowledge_enhanced_NN.jld2" solution recovered_sys approximate_solution_long NNsolution L true_solution_long Xₙ L̄ L̂

c1 = 3 # RGBA(174/255,192/255,201/255,1) # Maroon
c2 = :orange # RGBA(132/255,159/255,173/255,1) # Red
c3 = :blue # RGBA(255/255,90/255,0,1) # Orange
c4 = :purple # RGBA(153/255,50/255,204/255,1) # Purple

p1 = plot(0.1:0.1:tspan[end],abs.(Array(solution)[:,2:end] .- NNsolution[:,2:end])' .+ eps(Float32),
          lw = 3, yaxis = :log, title = "Timeseries of UODE Error",
          color = [3 :orange], xlabel = "t",
          label = ["x(t)" "y(t)"],
          titlefont = "Helvetica", legendfont = "Helvetica",
          legend = :topright)

# Plot L₂
p2 = plot(X[1,:], X[2,:], L̂[2,:], lw = 3,
     title = "Neural Network Fit of U2(t)", color = c1,
     label = "Neural Network", xaxis = "x", yaxis="y",
     titlefont = "Helvetica", legendfont = "Helvetica",
     legend = :bottomright)
plot!(X[1,:], X[2,:], L̄[2,:], lw = 3, label = "True Missing Term", color=c2)

p3 = scatter(solution, color = [c1 c2], label = ["x data" "y data"],
             title = "Extrapolated Fit From Short Training Data",
             titlefont = "Helvetica", legendfont = "Helvetica",
             markersize = 5)

plot!(p3,true_solution_long, color = [c1 c2], linestyle = :dot, lw=5, label = ["True x(t)" "True y(t)"])
plot!(p3,approximate_solution_long, color = [c3 c4], lw=1, label = ["Estimated x(t)" "Estimated y(t)"])
plot!(p3,[2.99,3.01],[0.0,maximum(hcat(Array(solution),Array(approximate_solution_long)))],lw=2,color=:black)
annotate!([(1.5,9,text("Training \nData", 10, :center, :top, :black, "Helvetica"))])
l = @layout [grid(1,2)
             grid(1,1)]
plot(p1,p2,p3,layout = l)

savefig("sindy_extrapolation.pdf")


using Random
Random.seed!(1)

N = 10^5
prob1 = 0.7
eps0, eps1 = 0.1, 0.05

flipWithProb(bit,prob) = rand() < prob ? xor(bit,1) : bit

TxData = rand(N) .< prob1
RxData = [x == 0 ? flipWithProb(x,eps0) : flipWithProb(x,eps1) for x in TxData]

numTx1 = 0
totalRx1 = 0
for i in 1:N
   if RxData[i] == 1
        global totalRx1 += 1
        global numTx1 += TxData[i]
    end
end

monteCarlo = numTx1/totalRx1
analytic = ((1-eps1)*0.7)/((1-eps1)*0.7+0.3*eps0)

println("Monte Carlo: ", monteCarlo, "\t\tAnalytic: ", analytic)


import Statistics.cor

# Autocorrelation

#TODO: move to another file to maintain a more logical structure

function lagindex(x::Vector{Float64}, n::Int=1)::UnitRange{Int}
    @assert n < size(x,1)
    N = size(x, 1)
    idx = 1:N
    if n == 0
        return idx
    elseif n > 0
        return idx[n+1:N]
    elseif n < 0
        return idx[1:N+n]
    end
end

function lag(x::Vector{Float64}, n::Int=1; pad::Bool=true, padval::Float64=NaN)
    if n==0
        return x
    elseif n>0
        if pad
            return [ones(n)*padval; x[n+1:end]]
        else
            return x[n+1:end]
        end
    elseif n<0
        n *= -1
        if pad
            return [x[1:end-n]; ones(n)*padval]
        else
            return x[1:end-n]
        end
    end
end

function corlag(x::AbstractArray{T,1}, n::Int=1) where {T<:Number}
    if n == 0
        return 1.0
    end
    @assert n > 0
    @assert n < size(x,1) - 2
    idx = collect(lagindex(x, n))
    return cor(x[idx.-n], x[idx])
end

function acf(x::Vector{T}, maxlag::Int=15; lags::AbstractArray{Int,1}=0:maxlag)::Vector{Float64} where {T<:Number}
    @assert all(lags .< size(x,1)-2)
    return map((n) -> corlag(x, n), lags)
end

function acf(x::Matrix{T}, maxlag::Int=15; lags::AbstractVector{Int}=0:maxlag)::Matrix{Float64} where {T<:Number}
    k = size(x,2)
    out = zeros(Float64, (length(lags), k))
    @inbounds for j in 1:k
        out[:,j] = acf(x[:,j], maxlag, lags=lags)
    end
    return out
end

"""
    acf(x::TS,maxlag::Int=15;lags::AbstractArray{Int,1}=0:maxlag)
    acf(x::Vector{T},maxlag::Int=15;lags::AbstractArray{Int,1}=0:maxlag)::Vector{Float64}where{T<:Number}
    acf(x::Matrix{T},maxlag::Int=15;lags::AbstractVector{Int}=0:maxlag)::Matrix{Float64}where{T<:Number}

Compute the autocorrelation function of a time series.

The output will be a matrix with the same number of columns as the input time series `x`
and number of rows equal to the number of lags used as inputs to the autocorrelation function.

...
# Arguments
- `x::TS`: Time series object array containing columns on which to compute autocorrelation.
- `maxlag::Int=15`: Maximum lag of the autocorrelation series.
Optional args:
- `lags::AbstractArray{Int,1}=0:maxlag`: Explicitly specified list of lags to use (overrides use of `maxlag`).
...
"""
acf(x::TS, maxlag::Int=15; lags::AbstractArray{Int,1}=0:maxlag) = acf(x.values, maxlag; lags=lags)


using NumericalMethodsforEngineers, Test

a = [2. -3. 1.; 1. -1. -2.; 3. 1. -1.]
b = [7., -2., 0.]

(lower, upper) = lufac(a)
println("NumericalMethodsforEngineers lufac(a) - lower:")
println(lower)
println()
println("NumericalMethodsforEngineers lufac(a) - upper:")
println(upper)
println()
@show lower * upper
println()

c1 = subfor(lower, b)
c2 = subbac(upper, c1)

subfor!(lower, b)
println("Updated RHS:")
println(b)
println()

@test ≈(c1,  b, atol=1e-10)
println()

c3 = subbac(upper, c1)

subbac!(upper, b)
println("Solution vector:")
println(b)
println()

c = a \ [7., -2., 0.]
c |> display
@test ≈(c2, c, atol=1e-10)
@test ≈(c3, c, atol=1e-10)

@test ≈(a * c, [7., -2., 0.], atol=1e-10)
println()

println("Julia lu(a):")
f = lu(a);
f |> display
println()
f.U |> display


"""
    norm1(v)

Normalize `v` to sum 1.
"""
function norm1(v)
    vs = sum(v)
    if vs ≈ 1.0
        return v
    else
        return v ./ vs
    end
end

"""
    infovalue(p, q)

Compute the symmetric Relative Entropy (Kullback-Liebler Divergence) between `p` and `q`,
two probability vectors or two frequency count vectors.
Value is bounded within [0, ∞].
"""
function infovalue(p, q)
    pn = norm1(p)
    qn = norm1(q)

    kldivergence(pn, qn) + kldivergence(qn, pn)
end
infovalue(f::AbstractMatrix) = infovalue(f[:, 1], f[:, 2])
infovalue(f::NamedArray) = infovalue(f.array)

"""
    cramerv(f::AbstractMatrix{T}) where T <: Real

Compute Cramer's V of a contingency table `f`.\\
ϕ coefficient is √(χ²/n), bounded within [0, √min(row-1,col-1)],
Cramer's V is ϕ/√min(row-1,col-1), bounded within [0,1].
"""
function cramerv(f::AbstractMatrix{T}) where {T<:Real}
    minimum(size(f)) >= 2 || throw(ArgumentError("Matrix needed, not single row or column"))

    p = norm1(f)
    rp = sum(p; dims = 2)               # row marginal
    cp = sum(p; dims = 1)               # column marginal
    ep = rp .* cp                       # expected probabilities
    ϕ² = sum((p .- ep) .^ 2 ./ ep)      # ϕ²

    sqrt(ϕ² / (minimum(size(f)) - 1))   # Cramer's V
end
cramerv(f::NamedArray) = cramerv(f.array)

"""
    cramerv(x::Vector{T} where T<:Integer, y::Vector{T} where T<:Integer)

Compute Cramer's V of two presumed discrete vectors `x` and `y`.\\
"""
cramerv(x::Vector{T} where {T<:Integer}, y::Vector{T} where {T<:Integer}) =
    cramerv(freqtable(x, y).array)

"""
    cramerv(x::Vector, y::Vector; groups = 10)

Compute Cramer's V of two presumed continuous vectors `x` and `y`,
vectors are binned into `groups` first.
"""
function cramerv(x::AbstractVector, y::AbstractVector; groups = 10)
    xr = ranks(x; groups = groups)
    yr = ranks(y, groups = groups)

    cramerv(freqtable(xr, yr).array)
end

"""
    mutualinfo(f::AbstractMatrix{T}) where T <: Real

Compute the mutual information of frequency matrix `f`;
value is bounded within [0, log(min(size(`f`)))].
"""
function mutualinfo(f::AbstractMatrix{T} where {T<:Real})
    minimum(size(f)) >= 2 || throw(ArgumentError("Matrix needed, not single row or column"))

    ps = norm1(f)
    px = sum(ps; dims = 2)
    py = sum(ps; dims = 1)

    entropy(px) + entropy(py) - entropy(ps)
end
mutualinfo(f::NamedArray) = mutualinfo(f.array)

"""
    mutualinfo(x::Vector{T} where T<:Integer, y::Vector{T} where T<:Integer)

Compute mutual information of two presumed discrete vectors `x` and `y`.\\
"""
mutualinfo(x::Vector{T} where {T<:Integer}, y::Vector{T} where {T<:Integer}) =
    mutualinfo(freqtable(x, y).array)

"""
    mutualinfo(x::Vector, y::Vector)

Compute the mutual information of two presumed continuous vectors `x` and `y`.
"""
function mutualinfo(x::AbstractVector, y::AbstractVector; groups = 10)
    xr = ranks(x; groups = groups)
    yr = ranks(y, groups = groups)

    mutualinfo(freqtable(xr, yr).array)
end

"""
    eda(df::AbstractDataFrame, target::Symbol)

Return a dataframe of Mutual Information and Cramer's V between `target` and other
variables in `df`. If `target` is binary, Information Value is also returned.

The dataframe is sorted by descending Mutual Information.
"""
function eda(df::AbstractDataFrame, target::Symbol; groups = 20)::AbstractDataFrame
    t = df[!, target]
    tnlvl = length(unique(t))
    tnlvl <= 1 && throw(ArgumentError("Target is single valued"))

    if tnlvl > groups
        t = ranks(t, groups = groups)
        tnlvl = length(unique(t))
    end

    println("Target: $target   type: $(typeof(t))   Levels: $tnlvl")

    lk = ReentrantLock()

    if tnlvl == 2
        out = DataFrame(
            Variable   = Symbol[],
            Vartype    = String[],
            Varlvls    = Int[],
            MutualInfo = Float64[],
            CramerV    = Float64[],
            InfoValue  = Float64[],
        )

        Threads.@threads for v in propertynames(df)
            v == target && continue

            vtype = eltype(df[!, v])
            if vtype <: Union{Missing,Real}
                vb = ranks(df[!, v], groups = groups)
            else
                vb = df[!, v]
            end

            frq = proptable(vb, t).array
            vnlvl = size(frq, 1)
            if vnlvl == 1
                println("Warning: [$v] is singled valued, skipped.")
                continue
            elseif vnlvl > 50
                println("Note: [$v] has more than 100 levels, consider binning")
            end

            mutin = mutualinfo(frq)
            cramv = cramerv(frq)
            iv = infovalue(frq[:, 1], frq[:, 2])

            lock(lk)
            push!(out, (v, string(vtype), vnlvl, mutin, cramv, iv))
            unlock(lk)
        end
    else
        out = DataFrame(
            Variable   = Symbol[],
            Vartype    = String[],
            Varlvls    = Int[],
            MutualInfo = Float64[],
            CramerV    = Float64[],
        )

        Threads.@threads for v in propertynames(df)
            v == target && continue

            vtype = eltype(df[!, v])
            if vtype <: Union{Missing,Real}
                vb = ranks(df[!, v], groups = groups)
            else
                vb = df[!, v]
            end

            frq = proptable(vb, t).array
            vnlvl = size(frq, 1)
            if vnlvl == 1
                println("Warning: [$v] is singled valued, skipped.")
                continue
            elseif vnlvl > 50
                println("Note: [$v] has more than 100 levels, consider binning")
            end

            mutin = mutualinfo(frq)
            cramv = cramerv(frq)

            lock(lk)
            push!(out, (v, string(vtype), vnlvl, mutin, cramv))
            unlock(lk)
        end
    end

    sort!(out, [:MutualInfo, :CramerV], rev = true)
end


"""
    Multivector2{T <: Real} <: AbstractMultivector{T}

An immutable 4-tuple of real numbers that represents a member
of a 2-dimensional multivector algebra.

Each `Multivector2` has the form
```math
    p+qB
```
where ``p`` and ``q`` are 1-dimensional multivectors, or
```math
    a+bA+cB+dAB
```
where ``a``, ``b``, ``c``, and ``d`` are real (and of the same type),
and ``A * A = 0``, ``B * B = 0``, and ``AB = A * B = -B * A``.
Here ``*`` is the wedge product.
"""
struct Multivector2{T <: Real} <: AbstractMultivector{T}
    l::Multivector1{T}
    r::Multivector1{T}

    Multivector2{U}(l::Multivector1{U}, r::Multivector1{U}) where {U <: Real} = new(l, r)
end

function Multivector2(x::Multivector1{T}, y::Multivector1{T}) where {T <: Real}
    Multivector2{T}(x, y)
end

function Multivector2(z::Multivector1)
    Multivector2(z, zero(z))
end

function Multivector2(a::T, b::T, c::T, d::T) where {T <: Real}
    Multivector2{T}(Multivector1{T}(a, b), Multivector1{T}(c, d))
end

function Multivector2(a::T, b::T, c::T) where {T <: Real}
    Multivector2{T}(Multivector1{T}(a, b), Multivector1{T}(c, zero(T)))
end

function Multivector2(a::T, b::T) where {T <: Real}
    Multivector2{T}(Multivector1{T}(a, b), zero(Multivector1{T}))
end

function Multivector2(a::T) where {T <: Real}
    Multivector2{T}(Multivector1{T}(a, zero(T)), zero(Multivector1{T}))
end

function Multivector2(a::Real, b::Real, c::Real, d::Real)
    Multivector2(promote(a, b, c, d)...)
end

function Multivector2(a::Real, b::Real, c::Real)
    Multivector2(promote(a, b, c)...)
end

function Multivector2(a::Real, b::Real)
    Multivector2(promote(a, b)...)
end

function Multivector2(x::Multivector1, y::Multivector1)
    Multivector2(promote(x.l, x.r, y.l, y.r)...)
end

function show(io::IO, z::Multivector2)
    print(io, "[1: ")
    print(io, z.l.l)
    print(io, ", A: ")
    print(io, z.l.r)
    print(io, ", B: ")
    print(io, z.r.l)
    print(io, ", AB: ")
    print(io, z.r.r)
    print(io, "]")
end

function real(z::Multivector2)
    real(z.l)
end

function unreal(z::Multivector2)
    vcat(unreal(z.l), asarray(z.r))
end

function zero(z::Multivector2{T}) where {T <: Real}
    Multivector2{T}(zero(z.l), zero(z.r))
end

function zero(::Type{Multivector2{T}}) where {T <: Real}
    Multivector2(zero(T), zero(T), zero(T), zero(T))
end

function one(z::Multivector2{T}) where {T <: Real}
    Multivector2{T}(one(z.l), zero(z.r))
end

function one(::Type{Multivector2{T}}) where {T <: Real}
    Multivector2(one(T), zero(T), zero(T), zero(T))
end

"""
    conj{T <: Real}(z::Multivector2{T})

The `Multivector2` conjugate. If ``z=a+bA+cB+dAB``, then `conj(z)` gives
```math
    a-bA-cB-dAB
```
This operation is an involution.
"""
function conj(z::Multivector2{T}) where {T <: Real}
    Multivector2{T}(conj(z.l), -z.r)
end

"""
    cloak{T <: Real}(z::Multivector2{T})

The cloak conjugate changes the sign of even blades. If ``z=a+bA+cB+dAB``, then `cloak(z)` gives
```math
    -a+bA+cB-dAB
```
This operation is equivalent to `-dagger(z)` and thus is also an involution.
"""
function cloak(z::Multivector2{T}) where {T <: Real}
    Multivector2{T}(cloak(z.l), dagger(z.r))
end

"""
    dagger{T <: Real}(z::Multivector2{T})

The dagger conjugate changes the sign of odd blades. If ``z=a+bA+cB+dAB``, then `dagger(z)` gives
```math
    a-bA-cB+dAB
```
This operation is an involution.
"""
function dagger(z::Multivector2{T}) where {T <: Real}
    Multivector2{T}(dagger(z.l), cloak(z.r))
end

"""
    star{T <: Real}(z::Multivector2{T})

Returns the Hodge star conjugate. If ``z=a+bA+cB+dAB``, then `star(z)` gives
```math
    d-cA+bB+aAB
```
This operation is not an involution, but `star(star(z)) = dagger(z)`.
"""
function star(z::Multivector2{T}) where {T <: Real}
    Multivector2{T}(dagger(star(z.r)), star(z.l))
end

"""
    selfstar(z::Multivector2)

The self-star-conjugate part.
This operation is idempotent.
"""
function selfstar(z::Multivector2)
    (z + star(z) + dagger(z) + star(dagger(z))) / 4
end

"""
    antiselfstar(z::Multivector2)

The anti-self-star-conjugate part.
This operation is idempotent.
"""
function antiselfstar(z::Multivector2)
    (z - star(z) + dagger(z) - star(dagger(z))) / 4
end

function (+)(x::Multivector2, y::Multivector2)
    Multivector2(x.l + y.l, x.r + y.r)
end

function (+)(x::Multivector2, y::Multivector1)
    Multivector2(x.l + y, x.r)
end

function (+)(x::Multivector1, y::Multivector2)
    Multivector2(x + y.l, y.r)
end

function (+)(z::Multivector2, a::Real)
    Multivector2(z.l + a, z.r)
end

function (+)(a::Real, z::Multivector2)
    Multivector2(z.l + a, z.r)
end

function (-)(z::Multivector2{T}) where {T <: Real}
    Multivector2{T}(-z.l, -z.r)
end

function (-)(x::Multivector2, y::Multivector2)
    Multivector2(x.l - y.l, x.r - y.r)
end

function (-)(x::Multivector2, y::Multivector1)
    Multivector2(x.l - y, x.r)
end

function (-)(x::Multivector1, y::Multivector2)
    Multivector2(x - y.l, -y.r)
end

function (-)(z::Multivector2, a::Real)
    Multivector2(z.l - a, z.r)
end

function (-)(a::Real, z::Multivector2)
    Multivector2(a - z.l, -z.r)
end

"""
    (*)(x::Multivector2, y::Multivector2)

Wedge product of 2-dimensional multivectors.
This operation is non-commutative but associative.
"""
function (*)(x::Multivector2, y::Multivector2)
    Multivector2(x.l * y.l,  (y.r * x.l) + (x.r * conj(y.l)))
end

function (*)(x::Multivector2, y::Multivector1)
    Multivector2(x.l * y, x.r * conj(y))
end

function (*)(x::Multivector1, y::Multivector2)
    Multivector2(x * y.l,  y.r * x)
end

"""
    (*)(z::Multivector2, a::Real)
    (*)(a::Real, z::Multivector2)

Scaling and/or reflection of a `Multivector2` by a real number.
"""
function (*)(z::Multivector2, a::Real)
    Multivector2(z.l * a, z.r * a)
end

function (*)(a::Real, z::Multivector2)
    Multivector2(a * z.l, a * z.r)
end

function (/)(x::Multivector2, y::Multivector2)
    if iszerodivisor(y)
        error(ZeroDivisorDenominator)
    end

    x * inv(y)
end

function (\)(y::Multivector2, x::Multivector2)
    if iszerodivisor(y)
        error(ZeroDivisorDenominator)
    end

    inv(y) * x
end

function (/)(z::Multivector2, a::Real)
    Multivector2(z.l / a, z.r / a)
end

function (\)(a::Real, z::Multivector2)
    Multivector2(a \ z.l, a \ z.r)
end

function random(::Type{Multivector2{T}}) where T <: Real
    Multivector2{T}(random(Multivector1{T}), random(Multivector1{T}))
end


"""
    eom_henon2(dx, x, p, n)

Equations of motion for a 2D Henon system consisting of two identical
Henon maps with unidirectional forcing ``X \\to Y `` [1].

## Equations of motion

The equations of motion are

```math
\\begin{aligned}
x_1(t+1) &= 1.4 - x_1^2(t) + 0.3x_2(t) \\\\
x_2(t+1) &= x_1(t) \\\\
y_1(t+1) &= 1.4 - [c_{xy} x_1(t) y_1(t) + (1-c_{xy}) y_1^2(t)] + 0.3 y_2(t) \\\\
y_2(t+1) &= y_1(t)
\\end{aligned}
```

## References

1. Krakovská, A., Jakubík, J., Chvosteková, M., Coufal, D., Jajcay, N., & Paluš, M. (2018).
    Comparison of six methods for the detection of causality in a bivariate time series.
    Physical Review E, 97(4), 042207.
"""
function eom_henon2(x, p, n)
    c_xy = p[1]
    x₁, x₂, y₁, y₂ = (x...,)
    dx₁ = 1.4 - x₁^2 + 0.3*x₂
    dx₂ = x₁
    dy₁ = 1.4 - (c_xy * x₁ * y₁  +  (1 - c_xy)*y₁^2) + 0.3*y₂
    dy₂ = y₁
    return SVector{4}(dx₁, dx₂, dy₁, dy₂)
end

function henon2(u₀, c_xy)
    p = [c_xy]
    DiscreteDynamicalSystem(eom_henon2, u₀, p)
end

"""
    henon2(;u₀ = rand(4), c_xy = 2.0) -> DiscreteDynamicalSystem

Initialize a 2D Henon system consisting of two identical Henon maps with
unidirectional forcing ``X \\to Y `` [1].

## Equations of motion

The equations of motion are

```math
\\begin{aligned}
x_1(t+1) &= 1.4 - x_1^2(t) + 0.3x_2(t) \\\\
x_2(t+1) &= x_1(t) \\\\
y_1(t+1) &= 1.4 - [c_{xy} x_1(t) y_1(t) + (1-c_{xy}) y_1^2(t)] + 0.3 y_2(t) \\\\
y_2(t+1) &= y_1(t)
\\end{aligned}
```

## References

1. Krakovská, A., Jakubík, J., Chvosteková, M., Coufal, D., Jajcay, N., & Paluš, M. (2018).
    Comparison of six methods for the detection of causality in a bivariate time series.
    Physical Review E, 97(4), 042207.
"""
henon2(;u₀ = rand(4), c_xy = 2.0) = henon2(u₀, c_xy)


@testset "Test Spline Basis Evaluation" begin

    # constuct specific case that we can compute by hand

    n = 5
    a = 0
    b = 1
    params = SplineParams(range(a,stop=b,length=n), 0, 1)

    x = rand(10000)

    basestr = @inferred BasisMatrices.evalbase(params, x)
    base = Array(basestr)

    # vector of basis functions

    function linBspline(u)
        if u >= 0 && u < 0.25
            return [1-4*u,4*u,0,0,0]

        elseif u >= 0.25 && u < 0.5
            return [0,2*(1 - 2*u),4*u-1,0,0]

        elseif u >= 0.5 && u < 0.75
            return [0,0,3 - 4*u,2*(2*u - 1),0]

        else
            return [0,0,0,4*(1 - u),4*u-3]
        end
    end

    manualbase = Array{Float64}(undef, length(x), n)

    for i in 1:length(x)
        manualbase[i, :] = linBspline(x[i])
    end

    @testset "test evalbase with linear B spline" begin
        @test all(manualbase .== base)
    end

    # test stuff that isn't implemented
    params2 = SplineParams(10, 0, 1, 2)
    @test_throws ErrorException evalbase(SplineSparse, params, nodes(params), -1)
    @test_throws ErrorException evalbase(params, nodes(params), -1)

    @test BasisMatrices.family(SplineParams) == Spline
    @test BasisMatrices.family_name(SplineParams) == "Spline"
    @test issparse(SplineParams)

    @test BasisMatrices.family(params) == Spline
    @test BasisMatrices.family_name(params) == "Spline"
    @test issparse(params)

    @testset "non Float64 eltypes" begin
        p = SplineParams(Float32[0.1, 0.2, 0.3, 0.4, 0.5], 0, 3)
        @test Float32 == eltype(@inferred BasisMatrices.evalbase(p, Float32[0.25, 0.35]))
        @test Float32 == eltype(@inferred BasisMatrices.evalbase(p, Float32[0.25, 0.35], 1))
        @test Float64 == eltype(@inferred BasisMatrices.evalbase(p, [0.25, 0.35]))
    end
end


using LifeContingencies
using MortalityTables
using Dates
using DataFrames

mort_table_name = "2012 IAM Period Table – Male, ANB"
improv_table = "Projection Scale G2 – Male, ANB"
iss_date = Date(2017,09,01)
val_date = Date(2017,12,31)
issue_age = 65
annual_income = 12000
deferral_period = 0
benefit_period = 120
interest_rate = 0.04

# Not used because table name provides
# sex = :M
# projection_scale = :G2
# age_rule = :ANB

mort = MortalityTables.table(mort_table_name)
imp = MortalityTables.table(improv_table)

table_end = omega(mort)

# the projection scale table only goes to age 105, wheras we want  through age 120
function imp_rate(table,age)
    if age > lastindex(table)
        return 0.0
    else
        return table[age]
    end
end

SRA = map(0:(table_end - issue_age)) do time

    att_age = issue_age + time

    if time > 1
        imp_factor = prod(map(age -> 1 - imp_rate(imp,age),issue_age:(att_age-1)))
    else
        imp_factor = 1.0
    end

    # returned values for each time

    (
        att_age = att_age,
        t=time,
        q=mort[att_age],
        imp_factor = imp_factor,
        q_imp = mort[att_age] * imp_factor,

    )
end

### A type to handle life contingent maths via the LifeContingencies package
ins = LifeContingency(
    SingleLife(
        mort = UltimateMortality([x.q_imp for x in SRA],start_age = issue_age), # mort is indexed by attained age, not starting at 1
        issue_age = issue_age
    ),
    InterestRate(interest_rate)
)

ä(ins,55) * annual_income
sum(LifeContingencies.APV.(ins,1:55) * 12000)

ℓ₁ = [43302,42854,42081,41351,40050]
ℓ₂ = [47260,47040,46755,46500,46227]

ps₁ =  ℓ₁ ./ ℓ₁[1]
qs₁ = [1 - ps₁[t] / ps₁[t - 1] for t in 2:5 ]

ps₂ =  ℓ₂ ./ ℓ₂[1]
qs₂ = [1 - ps₂[t] / ps₂[t - 1] for t in 2:5 ]

m1 = UltimateMortality(qs₁, start_age=65)
m2 = UltimateMortality(qs₂, start_age=60)

l1 = SingleLife(mort = m1, issue_age = 65)
l2 = SingleLife(mort = m2, issue_age = 60)

jl = JointLife(lives=(l1, l2), contingency = LastSurvivor(), joint_assumption=Frasier())

ins = LifeContingency(jl,InterestRate(0.05))
ins_l1 = LifeContingency(jl.lives[1],InterestRate(0.05))
ins_l2 = LifeContingency(jl.lives[2],InterestRate(0.05))
# problem 9.1.f


using LinearAlgebra
using DataFrames
using CSV
using Plots

obs = CSV.File("../../DATA/ovariancancer_obs.csv", header = false) |> Tables.matrix
grp = CSV.File("../../DATA/ovariancancer_grp.csv", header = false) |> Tables.matrix

U, S, V = svd(obs)

p1 = plot(S, yaxis = :log, legend = false)
p2 = plot(cumsum(S)./sum(S), legend = false)

cancer = [[],[],[]]
normal = [[],[],[]]

for i = 1:size(obs)[1]
    x = V[:,1]'*obs[i,:]
    y = V[:,2]'*obs[i,:]
    z = V[:,3]'*obs[i,:]
    if (grp[i] == "Cancer")
        append!.(cancer, [x, y, z])
    else
        append!.(normal, [x, y, z])
    end
end

p3 = scatter(cancer..., mc = :red, label = "Cancer")
scatter!(normal..., mc = :blue, label = "Normal")


# Log sum exp.
function lse(arr)
    max = maximum(arr)
    max == -Inf ? -Inf : max + log(sum(exp.(arr .- max)))
end

function lse(x1::Real, x2::Real)
    m = max(x1, x2)
    m == -Inf ? m : m + log(exp(x1 - m) + exp(x2 - m))
end

# Effective sample size.
function ess(lnw::Vector{Float64})
    log_ess = -lse(2. * lnw)
    return exp(log_ess)
end

# Normalize log weights.
function nw(lw::Vector{Float64})
    lt = lse(lw)
    lnw = lw .- lt
    return (lt, lnw)
end


using ArgParse

s=ArgParseSettings()
@add_arg_table s begin
    "--part2"
        help="Do part 2"
        action=:store_true
end
parsed_args = parse_args(ARGS, s)

function calculate_gamma(numbers, bitlen)
    counts = zeros(bitlen)
    for n in numbers
        for idx in 1:bitlen
            mask = 0x1000 >>> idx
            if (n & mask) != 0
                counts[idx] += 1
            end
        end
    end

    l = div(length(numbers), 2)
    gamma::UInt32 = 0
    for idx in 1:bitlen
        gamma <<= 1
        if counts[idx] >= l
            gamma |= 1
        end
    end
    gamma
end

function part1(bits)
    gamma = calculate_gamma(bits, 12)
    epsilon = xor(gamma, 0x0fff)
    epsilon * gamma
end

function filter_values(values, use_first)
    bit = 0x0800
    while bit != 0 && length(values) > 1
        set = count(v -> (v & bit) != 0 , values)
        if use_first
            good = set >= (length(values) - set) ? bit : 0
        else
            good = set < (length(values)-set) ? bit : 0
        end
        values = filter(b -> (b & bit) == good, values)
        bit >>>= 1
    end
    values[1]
end

function part2(values)
    o2_rating = filter_values(values, true)
    co2_rating = filter_values(values, false)
    o2_rating * co2_rating
end

function parseLine(line)
    parse(UInt32, line, base=2)
end

bits = [parseLine(line) for line in eachline("input/day_3.txt")]
if parsed_args["part2"]
    answer = part2(bits)
else
    answer = part1(bits)
end
println("The answer is: $answer")


#=
  From the WebPPL model:https://github.com/probmods/ppaml2016/blob/gh-pages/chapters/5-data.md
  """
  Here, we explore the result of an experiment with 15 trials and binary outcomes (e.g., flipping a coin with an
  uncertain weight, asking people if they'll vote for Hillary Clinton or Donald Trump, ...)
  """
=#

using Turing
include("jl_utils.jl")

@model function bda(k=1,n=15)
    p ~ Uniform(0, 1)

    # Observed k number of successes, assuming a binomial
    k ~ Binomial(n, p)

    # sample from binomial with updated p
    posteriorPredictive ~ Binomial(n,p)

    # sample fresh p
    prior_p ~ Uniform(0,1)
    # sample from binomial with fresh p
    priorPredictive ~ Binomial(n,prior_p)

end

# data
k = 1  # number of successes
n = 15 # number of attempts

model = bda(k,n)

# chns = sample(model, MH(), 10_000)
# chns = sample(model, PG(15), 10_000)
# chns = sample(model, IS(), 10_000)
chns = sample(model, SMC(), 10_000)
# chns = sample(model, SMC(), MCMCThreads(), 10_000, 4)

display(chns)

show_var_dist_pct(chns, :priorPredictive)
show_var_dist_pct(chns, :posteriorPredictive)


"""
A recursive implementation of divided differences.
"""
mutable struct DividedDifference{T}
    memo::Dict{Tuple{Int64,Int64},T}
    DividedDifference{T}() where {T <: Real} = new(Dict{Tuple{Int64,Int64},T}())
end

"""
Calculate divided differences of a function f at points τ. This calculates the divided difference
from τ_i to τ_j. The `dd` is a memoization.
"""
function divided_difference(i::Integer, j::Integer, τ::Vector, f::Vector, dd::DividedDifference)
    if (i, j) ∈ keys(dd.memo)
        return dd.memo[(i, j)]
    end

    val = zero(Float64)
    if i == j
        val = f[i]
    elseif i < j
        val = (divided_difference(i + 1, j, τ, f, dd) - divided_difference(i, j - 1, τ, f, dd)) /
            (τ[j] - τ[i])
    else
        throw(ArgumentError("divided_difference needs i <= j, found i = {i}, j = {j}"))
    end
    dd.memo[(i, j)] = val
    val
end

function divided_difference_explicit(
    i::Integer, k::Integer, τ::Vector, f::Vector{T}) where {T <: Real}

    total = zero(T)
    for j = i:k
        denom = one(T)
        for l = i:k
            if l != j
                denom *= τ[j] -τ[l]
            end
        end
        total += f[j] / denom
    end
    total
end


# using JuMP, VariationalInequality
# using Base.Test

using JuMP, VariationalInequality
# include("../src/model.jl")
# include("../src/algorithms.jl")

# https://cdr.lib.unc.edu/indexablecontent/uuid:778ca632-74ca-4858-8c3c-6dcfc7e6e703
# Example 4.2. This example is adapted from the example in (Fukushima, 1986). Let
m = VIPModel()
@variable(m, x[1:3])

@NLconstraint(m, x[1]^2 + 0.4x[2]^2 + 0.6x[3]^2 <= 1)
@NLconstraint(m, 0.6x[1]^2 + 0.4x[2]^2 + x[3]^2 <= 1)
@NLconstraint(m, x[1] + x[2] + x[3] >= sqrt(3))

@NLexpression(m, F1, 2x[1] + 0.2x[1]^3 - 0.5x[2] + 0.1x[3] - 4)
@NLexpression(m, F2, -0.5x[1] + x[2] + 0.1x[2]^3 + 0.5)
@NLexpression(m, F3, 0.5x[1] - 0.2x[2] + 2x[3] - 0.5)

correspond(m, F1, x[1])
correspond(m, F2, x[2])
correspond(m, F3, x[3])
# or
# correspond(m, [F1, F2, F3], x)

solveVIP(m, algorithm=:fixed_point, max_iter=1000, step_size=0.1)
sol1, Fval1, gap1 = saveSolution(m)
@assert 0<= gap1 < 1e-6

# The above `solveVIP` sets the value of variables at the solution
clearValues(m)

solveVIP(m, algorithm=:extra_gradient, max_iter=1000, step_size=0.1)
sol2, Fval2, gap2 = saveSolution(m)
@assert 0<= gap2 < 1e-6

# sol2, Fval2, gap2 = solveVIP(m, algorithm=:hyperplane, max_iter=1000, step_size=0.1)

@show sol1
@show sol2
@show gap1
@show gap2

# OKAY.
# x∗ = (0.9168, 0.4850, 0.3303)


factorial_ij(i::Integer, j::Integer) = factorial(big(min(i,j))) / factorial(big(max(i,j)))

z(x::Real, p::Real) = sqrt(2.)*(x + p*im)
z(x::Vector{<:Real}, p::Vector{<:Real}) = z.(x, p')

α(m::Integer, n::Integer) = n - m
α(m::Vector{<:Integer}, n::Vector{<:Integer}) = α.(min.(m, n'), max.(m, n'))

gaussian_function(x::Real, p::Real) = exp(-0.5 * abs2(z(x,p))) / π
gaussian_function(x::Vector{<:Real}, p::Vector{<:Real}) = gaussian_function.(x,p')

neg_one_to_power_of(i::Integer) = (i % 2 == 0) ? 1 : -1

function coefficient_of_wave_function(m::Integer, n::Integer)
    if n ≥ m
        return neg_one_to_power_of(m) * sqrt(factorial_ij(m, n))
    else
        return neg_one_to_power_of(n) * sqrt(factorial_ij(m, n))
    end
end

coefficient_of_wave_function(m::Vector{<:Integer}, n::Vector{<:Integer}) = coefficient_of_wave_function.(m, n')

function z_to_power(m::Integer, n::Integer, x::Real, p::Real)
    if n ≥ m
        return conj(z(x, p'))^(n - m)
    else
        return z(x, p')^(m - n)
    end
end

function z_to_power(m::Vector{<:Integer}, n::Vector{<:Integer}, x::Vector{<:Real}, p::Vector{<:Real})
    x = reshape(x, 1, 1, length(x))
    p = reshape(p, 1, 1, 1, length(p))
    z_to_power.(m, n', x, p)
end

function z_to_power(m::Vector{<:Integer}, n::Vector{<:Integer})
    function z_to_power_xp(x::Vector{<:Real}, p::Vector{<:Real})
        x = reshape(x, 1, 1, length(x))
        p = reshape(p, 1, 1, 1, length(p))
        z_to_power.(m, n', x, p)
    end
    return z_to_power_xp
end

function z_to_power(x::Vector{<:Real}, p::Vector{<:Real})
    x = reshape(x, 1, 1, length(x))
    p = reshape(p, 1, 1, 1, length(p))
    function z_to_power_mn(m::Vector{<:Integer}, n::Vector{<:Integer})
        z_to_power.(m, n', x, p)
    end
    return z_to_power_mn
end


#############################################################################
# JuMP
# An algebraic modelling langauge for Julia
# See http://github.com/JuliaOpt/JuMP.jl
#############################################################################
# minellipse.jl
#
# This example is from the Boyd & Vandenberghe book "Convex Optimization".
# If PyPlot is installed, passing the argument `plot` will plot
# the solution, e.g. julia minellipse.jl plot
#
# Given a set of ellipses centered on the origin
# E(A) = { u | u^T inv(A) u <= 1 }
# find a "minimal" ellipse that contains the provided ellipses
#
# We can formulate this as an SDP:
#     minimize  trace(WX)
#   subject to  X >= A_i,    i = 1,...,m
#               X PSD
# where W is a PD matrix of weights to choose between different solutions
#############################################################################

using JuMP
using Mosek
using PyPlot  # Comment out if not installed

# We will use three ellipses
m = 3
# Two "simple" ones
As = Any[ [2.0  0.0;
           0.0  1.0],
          [1.0  0.0;
           0.0  3.0]]
# and a random one
randA = rand(2,2)
push!(As, (randA' * randA) * (rand()*2+1))

# We change the weights to see different solutions, if they exist
W = [1.0 0.0
     0.0 1.0];

mod = Model()
@variable(mod, X[1:2,1:2], SDP)
@objective(mod, Min, trace(W*X))
for i = 1:m
    @SDconstraint(mod, X >= As[i])
end
solve(mod)

X_val = getvalue(X)
println(X_val)

# Plot it, if desired (e.g. julia minellipse.jl plot)
if length(ARGS) > 0 && Pkg.installed("PyPlot") !== nothing
    # Setup the figure
    fig = figure()
    ax = fig[:gca]()
    ax[:set_xticks]([-4:+4])
    ax[:set_yticks]([-4:+4])

    # Draw provided ellipses
    for i = 1:m
        xs = Float64[]
        ys = Float64[]
        for angle in linspace(0, 2*pi, 100)
            u = [cos(angle),sin(angle)]
            x = As[i] * u
            push!(xs, x[1])
            push!(ys, x[2])
        end
        plot(xs, ys, "b", linewidth=2.0)
    end

    # Draw bounding ellipse
    xs = Float64[]
    ys = Float64[]
    for angle in linspace(0, 2*pi, 100)
        u = [cos(angle),sin(angle)]
        x = X_val * u
        push!(xs, x[1])
        push!(ys, x[2])
    end
    plot(xs, ys, "r", linewidth=2.0)

    fig[:canvas][:draw]()
    readline()
end


using Test
using Recommenders: σ, Logloss, grad

@testset "Sigmoid." begin
    @test σ(0.0) == 0.5
    @test σ(100) ≈ 1.0
    @test σ(-100) < 1e-20
    @test σ(1.0) ≈ 1.0 / (1 + exp(-1))
    @test σ(0.1) > 0.5
    @test σ(-0.1) < 0.5
end

@testset "Logloss." begin
    logloss = Logloss()
    @test logloss(100, 1) ≈ 0
    @test logloss(-100, 1) > 20
    @test logloss(3, 0) > 3
    @test logloss(-100, 0) ≈ 0
    @test logloss(0.0, 1) == logloss(0.0, 0)
end

@testset "Grad of Logloss." begin
    logloss = Logloss()

    @test grad(logloss, 0.3243, 0) ≈ σ(0.3243)
    @test grad(logloss, 0.3243, 1) ≈ σ(0.3243) - 1

    # check sign
    logit = 0.4
    label = 1
    loss = logloss(logit, label)
    grad_logloss = grad(logloss, logit, label)
    @test logloss(logit - 0.1 * grad_logloss, label) < loss

    label = 0
    loss = logloss(logit, label)
    grad_logloss = grad(logloss, logit, label)
    @test logloss(logit - 0.1 * grad_logloss, label) < loss

    logit = -0.3
    label = 1
    loss = logloss(logit, label)
    grad_logloss = grad(logloss, logit, label)
    @test logloss(logit - 0.1 * grad_logloss, label) < loss

    label = 0
    loss = logloss(logit, label)
    grad_logloss = grad(logloss, logit, label)
    @test logloss(logit - 0.1 * grad_logloss, label) < loss
end


# Examples borrowed from: http://arxiv.org/pdf/1211.3759.pdf

# 2/3 N(x|0,1) + 1/3 N(x|0,1/10)
function kurtotic(theta)
    if theta[3] <= 0 || theta[6] <= 0
        return -Inf
    end
    llk = 0
    for i in 1:length(y)
      llk += log(theta[1] * pdf(Normal(theta[2], theta[3]), y[i]) +
                 theta[4] * pdf(Normal(theta[5], theta[6]), y[i]))
    end
    return llk
end

function bimodal_density(x)
    lik = (1/2) * pdf(Normal(-1,2/3),x) + (1/2) * pdf(Normal(1,2/3),x)
    return log(lik)
end

# 1/2 N(x|-1,2/3) + 1/2 N(x|1,2/3)
function bimodal_loglikelihood(theta)
    if theta[3] < 0 || theta[6] < 0
        return -Inf
    end
    llk = 0
    for i in 1:length(y)
      llk += log(theta[1] * pdf(Normal(theta[2], theta[3]), y[i]) +
                 theta[4] * pdf(Normal(theta[5], theta[6]), y[i]))
    end
    return -llk
end

# requires y and x
function logistic_regression(theta)
    llk = 0
    eta = [1 / (1 + exp(- x[i,:] * theta))[1] for i in 1:length(y)]
    for i = 1:length(y)
        llk += y[i] * log(eta[i]) + (1-y[i]) * log(1-eta[i])
    end
    return llk
end