using EconPDEs

Base.@kwdef  struct CampbellCochraneModel
    # consumption process parameters
    μ::Float64 = 0.0189
    σ::Float64 = 0.015
    # utility
    γ::Float64 = 2.0
    ρ::Float64 = 0.116
    # habit
    κs::Float64 = 0.138
    b::Float64 = 0.0
end
# I choose persistence so that monthly simulation of the model matches processes in CC (1999)
# ρ = 12 * (1 - 0.89^(1/12))
# κs = 12 * (1 - 0.87^(1/12))

function initialize_stategrid(m::CampbellCochraneModel; sn = 1000)
    μ = m.μ ; σ = m.σ ; γ = m.γ ; ρ = m.ρ ; κs = m.κs ; b = m.b
    Sbar = σ * sqrt(γ / (κs - b / γ))
    sbar = log.(Sbar)
    smax =  sbar + 0.5 * (1 - Sbar^2)
    # corresponds to Grid 3 in Wachter (2005)
    shigh = log.(range(0.0, exp(smax), length = div(sn, 10)))
    slow = range(-300.0, shigh[2], length = sn - div(sn, 10))
    OrderedDict(:s => vcat(slow[1:(end-1)], shigh[2:end]))
end

function (m::CampbellCochraneModel)(state::NamedTuple, y::NamedTuple)
    (; μ, σ, γ, ρ, κs, b) = m
    (; s) = state
    (; p, ps_up, ps_down, pss) = y

    # drift and volatility of  s and p
    Sbar = σ * sqrt(γ / (κs - b / γ))
    sbar = log(Sbar)
    λ = 1 / Sbar * sqrt(1 - 2 * (s - sbar)) - 1
    μs = - κs * (s - sbar)
    ps = (μs >= 0) ? ps_up : ps_down
    σs = λ * σ
    σp = ps / p * σs
    μp = ps / p * μs + 0.5 * pss / p * σs^2

    # market price of risk κ
    κ = γ * (σ + σs)

    # risk free rate  r
    r = ρ + γ * μ - (γ * κs - b) / 2 + b * (sbar - s)

    # PDE
    pt = - p * (1 / p + μ + μp + σp * σ - r - κ * (σ + σp))
    return (; pt)
end

# Campbell Cochrane (1999)
m = CampbellCochraneModel()
stategrid = initialize_stategrid(m)
yend = OrderedDict(:p => ones(length(stategrid[:s])))
result = pdesolve(m, stategrid, yend)
@assert result.residual_norm <= 1e-5

# Wachter (2005) calibration
# m = CampbellCochraneModel(μ = 0.022, σ = 0.0086, γ = 2.0, ρ = 0.073, κs = 0.116, b = 0.011)
# stategrid = initialize_stategrid(m)
# yend = OrderedDict(:p => ones(length(stategrid[:s])))
# y, result, distance = pdesolve(m, stategrid, yend)


using Test

@testset "BIC"  begin

    using StatsModelComparisons

    v1 = bic(-200.0, 10, 100)
    v2 = bic(-200.0, 2, 100)
    v3 = bic(-200.0, 2, 10)
    v4 = bic(-100.0, 2, 10)
    @test v2 < v1
    @test v3 < v2
    @test v4 < v3
end

@testset "AIC"  begin

    using StatsModelComparisons

    v1 = aic(-200.0, 10)
    v2 = aic(-200.0, 2)
    v3 = aic(-100.0, 2)
    @test v2 < v1
    @test v3 < v2
end

@testset "DIC"  begin

    using StatsModelComparisons
    x = randn(100)
    LLs = x .- mean(x)
    v1 = dic(LLs)
    v2 = dic(LLs*2)
    v3 = dic(LLs .+ 1)
    @test v1 < v2
    @test v3 < v1
end

@testset "WAIC and LOO" begin

    using StatsModelComparisons, StanSample
    using StatsFuns, RDatasets, Random, CSV
    import StanSample: read_csv_files

    cd(@__DIR__)
    Random.seed!(89905)
    # df = RDatasets.dataset("datasets", "cars")

    # stream = open("cars_model.stan", "r")
    # cars_stan = read(stream, String)
    # close(stream)

    # cars_stan_model = SampleModel("cars.model", cars_stan)
    # data = (N = size(df, 1), speed = df.Speed, dist = df.Dist)
    # rc = stan_sample(cars_stan_model; data)

    # stan_summary(cars_stan_model, true)
    # nt_cars = read_samples(cars_stan_model);
    # log_lik = nt_cars.log_lik'

    # files = readdir("test_chains")
    # files = map(f->"test_chains/"*f, files)
    # samples = read_csv_files(files, 2000, :namedtuple; include_internals=false)
    # df = vcat(dfs...)

    files = readdir("test_chains/")

    dfs = map(f->CSV.read("test_chains/"*f, DataFrame), files)
    dfs = map(x->x[1001:end,:], dfs)
    df = vcat(dfs...)

    col_idx = findall(x->occursin("log_lik", x), names(df))
    pw_log_liks = Array(df[:,col_idx])
    WAIC = waic(pw_log_liks)
    loo,_,_ = psisloo(pw_log_liks)

    @test WAIC.WAIC ≈ 421.0 rtol = .01
    @test WAIC.std_err ≈ 16.3 rtol = .01
    @test loo ≈ -210.5 rtol = .05
end


module StatisticTools

export weightedmean

"""
weightedmean(values::Array{Float64,1}, errors::Array{Float64,1})

computes the weighted mean of values, using as weights w = 1/ sigma^2
"""
function weightedmean(x::Array{Float64,1}, errors::Array{Float64,1})
    @assert length(x) == length(errors)
    weights = 1./ (errors.^2)
    sigma = sqrt( 1./ sum( weights ) )
    mean = sigma^2 * sum(x .* weights)

    mean, sigma
end

end # module


module Epidemiology

using AlgebraicPetri
using Petri
using Catlab
using Catlab.Theories

export InfectiousDiseases, FunctorGenerators, F_epi, S, E, I, R, D, transmission, exposure, illness, recovery, death

ob = PetriCospanOb(1)
spontaneous_petri = PetriCospan([1], Petri.Model(1:2, [(Dict(1=>1), Dict(2=>1))]), [2])
transmission_petri = PetriCospan([1], Petri.Model(1:2, [(Dict(1=>1, 2=>1), Dict(2=>2))]), [2])
exposure_petri = PetriCospan([1, 2], Petri.Model(1:3, [(Dict(1=>1, 2=>1), Dict(3=>1, 2=>1))]), [3, 2])

""" InfectiousDiseases
"""
@present InfectiousDiseases(FreeBiproductCategory) begin
    S::Ob
    E::Ob
    I::Ob
    R::Ob
    D::Ob
    transmission::Hom(S⊗I, I)
    exposure::Hom(S⊗I, E⊗I)
    illness::Hom(E,I)
    recovery::Hom(I,R)
    death::Hom(I,D)
end

""" generators
"""
S,E,I,R,D,transmission,exposure,illness,recovery,death = generators(InfectiousDiseases);

""" FunctorGenerators
"""
const FunctorGenerators = Dict(S=>ob, E=>ob, I=>ob, R=>ob, D=>ob,
        transmission=>transmission_petri, exposure=>exposure_petri,
        illness=>spontaneous_petri, recovery=>spontaneous_petri, death=>spontaneous_petri)

""" F_epi
"""
F_epi(ex) = functor((PetriCospanOb, PetriCospan), ex, generators=FunctorGenerators)

end


function tvlqr(A, B, Q, R)
    T = length(Q)

    P = [zero(A[1]) for t = 1:T]
    K = [zero(B[1]') for t = 1:T-1]
    P[T] = Q[T]

    for t = T-1:-1:1
        K[t] = (R[t] + B[t]' * P[t+1] *  B[t]) \ (B[t]' * P[t+1] * A[t])
        P[t] = (Q[t] + K[t]' * R[t] * K[t]
                + (A[t] - B[t] * K[t])' * P[t+1] * (A[t] - B[t] * K[t]))
    end

    return K, P
end

function reference_gains(s::Simulation, traj::ContactTraj, obj::TrackingVelocityObjective;
		N::Int=10, κ=2e-4, U_scaling=100, V_scaling=100)

	model = s.model
	H = traj.H
	z = [[deepcopy(traj.z) for i = 1:N]...;]
	θ = [[deepcopy(traj.θ) for i = 1:N]...;]

	rz = [LinearizedStep(s, z[t], θ[t], κ).rz for t = 1:N*H]
	rθ = [LinearizedStep(s, z[t], θ[t], κ).rθ for t = 1:N*H]

	idx_z = indices_z(s)
	idx_θ = indices_θ(model, nf=1)

	∂z∂θ = [-rz[t] \ rθ[t] for t = 1:N*H]

	∂q2∂q1 = [zeros(model.nq, model.nq) for t = 1:N*H]
	∂q2∂q2 = [I(model.nq) for t = 1:N*H]
	∂q3∂q1 = [∂z∂θ[t][idx_z.q, idx_θ.q1] for t = 1:N*H]
	∂q3∂q2 = [∂z∂θ[t][idx_z.q, idx_θ.q2] for t = 1:N*H]

	∂q2∂u = [zeros(model.nq, model.nu) for t = 1:N*H]
	∂q3∂u = [∂z∂θ[t][idx_z.q, idx_θ.u] for t = 1:N*H]

	A = [[∂q2∂q1[t] ∂q2∂q2[t]; ∂q3∂q1[t] ∂q3∂q2[t]] for t = 1:N*H]
	B = [[∂q2∂u[t]; ∂q3∂u[t]] for t = 1:N*H]
	Q = [
		[obj.q[1]+V_scaling*obj.v[1]    -V_scaling*obj.v[1];
		 -V_scaling*obj.v[1]    obj.q[1]+V_scaling*obj.v[1]
		 ] for i = 1:N*H+1]
	R = [U_scaling * obj.u[1] for i = 1:N*H]

	K, P = tvlqr(A, B, Q, R)
	return K[1:H]
end


module Geometry
export volume, surfacearea, lateralarea

abstract type Solid end

# *******************************************
mutable struct Cone <: Solid
    height::Float32
    # one element implies start from one point
    # two elements for truncated cone volume
    radii::Array{Float32, 1}
    function Cone(depth, diameters)
        if length(diameters) > 2
            error("Too many diameter arguments in list")
        end
        if length(diameters) == 1
            push!(diameters,0.0)
        end
        new(depth, diameters / 2)
    end
end
mutable struct Cylinder <: Solid
    height::Float32
    radius::Float32
    function Cylinder(depth, diameter)
        new(depth, (diameter / 2.0))
    end
end
struct Cube <: Solid
    side::Float32
end
struct RectangularPrism <: Solid
    height::Float32
    width::Float32
    depth::Float32
end

function volume(c::Solid)
    tc = typeof(c)
    if tc == Cone
        rads = (c.radii[1]^2 + c.radii[1]*c.radii[2] + c.radii[2]^2)
        return (1/3)*(pi*rads*c.height)
    elseif tc == Cylinder
        return (1/1)*(pi*(c.radius^2)*c.height)
    elseif tc == Cube
        return (1/1)*(c.side^3)
    elseif tc == RectangularPrism
        return (1/1) * (c.height * c.width * c.depth)
    end
    return nothing
end

function lateralarea(c::Solid)
    tc = typeof(c)
    if tc == Cone
        return pi * (c.radii[1] + c.radii[2]) * (((c.radii[1] - c.radii[2])^2 + c.height^2)^(1/2))
    elseif tc == Cylinder
        return 2*pi*c.radius*c.height
    elseif tc == Cube
        return (1/1)*4*(c.side^2)
    elseif tc == RectangularPrism
        return (1/1) * (2 * (c.height * c.width) + 2 * (c.height * c.depth))
    end
    return nothing
end

function surfacearea(c::Solid)
    tc = typeof(c)
    if tc == Cone
        return lateralarea(c) + pi * (c.radii[1]^2 + c.radii[2]^2)
    elseif tc == Cylinder
        return lateralarea(c) + 2*pi*(c.radius^2)
    elseif tc == Cube
        return (1/1)*6*(c.side^2)
    elseif tc == RectangularPrism
        return lateralarea(c) + (1/1) * (2 * (c.width * c.depth))
    end
    return nothing

end

end


using FFTW
using Plots

include("naiveDFT.jl")
include("simpleDFT.jl")
include("benchmarkFunction.jl")

function main()
    #
    # Create test set
    #
    N = 2^8 # 2^8=256
    tmax = 30.0
    Δ = tmax / N

    t = collect(range(0.0, tmax, length = N))
    ω = collect(range(0, stop = N - 1))
    hk = sin.(t) .+ 0.5 .* sin.(5.0 .* t .+ 1.0) .+ sin.(4.0 .* t .- 1.5)

    #
    # Calculations
    #
    resNaive = naiveDFT(hk)                 #naive approach
    resNaiveIter = naiveDFTIterative(hk)    #naive approach iterative
    resSimple = simpleDFT(hk)               #simple recursive dft
    resFFTW = fft(hk)                       #fftw
    resRealFFTW = rfft(hk)                  #real fftw

    println(sum(abs.(resFFTW .- resNaive))) #differences to naive<->fftw
    println(sum(abs.(resFFTW .- resNaiveIter))) #differences to naive<->fftw
    println(sum(abs.(resFFTW .- resSimple))) #differences to simple<->fftw

    #
    # Benchmarking
    #
    benchmarkDFTFunction(fft, 1.0, 5, "FFTW.csv")
    benchmarkDFTFunction(simpleDFT, 1.0, 5, "simpleDFT.csv")
    benchmarkDFTFunction(naiveDFT, 1.0, 5, "naiveDFT.csv")
    benchmarkDFTFunction(naiveDFTIterative, 1.0, 5, "naiveDFTIter.csv")

    #
    # plotting
    #
    gr()
    p1re = plot(ω, real.(resFFTW), title = "FFTW Real", label = "");
    p1im = plot(ω, imag.(resFFTW), title = "FFTW Imag", label = "");
    p1abs = plot(ω, abs.(resFFTW) .^ 2 , title = "FFTW Abs^2", label = "");

    p2re = plot(ω, real.(resNaive), title = "Naive Real", label = "");
    p2im = plot(ω, imag.(resNaive), title = "Naive Imag", label = "");
    p2abs = plot(ω, abs.(resNaive) .^ 2 , title = "Naive Abs^2", label = "");

    p3re = plot(ω, real.(resSimple), title = "Simple Real", label = "");
    p3im = plot(ω, imag.(resSimple), title = "Simple Imag", label = "");
    p3abs = plot(ω, abs.(resSimple) .^ 2 , title = "Simple Abs^2", label = "");

    p = plot(p1re, p1im, p1abs, p2re, p2im, p2abs, p3re, p3im, p3abs, layout = (3, 3), size = (1024, 800))
    savefig(p, "fftPlot.png")

end

main()


"""
```
transition(m::PoolModel{T}) where {T<:AbstractFloat}
```

Assign transition equation

```
x_t = (1-ρ) μ + ρ x_{t-1} + sqrt{1 - ρ^2} σ ϵ_t
ϵ_t ∼ iid N(0,1), x_0 ∼ N(μ,σ^2)
λ_t = Φ(x_t)
```
where Φ(⋅) is the cdf of a N(0,1) random variable, F_ϵ is the distribution of ϵ_t,
and F_λ is the distribution of λ(x_0).
"""
function transition(m::PoolModel{T}) where {T<:AbstractFloat}
    @inline Φ(x::Vector{Float64}, ϵ::Vector{Float64}) = abs.([0;1] .-
                                                             (cdf.(Normal(), (1 - m[:ρ].value) .*
                                                                   m[:μ].value .+ m[:ρ].value .*
                                                                   quantile(Normal(),x[1]) .+
                                                                   sqrt(1 - m[:ρ].value^2) .*
                                                                   m[:σ].value .* ϵ)))
    F_ϵ = Normal(0.,1.)
    F_λ = Uniform(0.,1.)
    return Φ, F_ϵ, F_λ
end


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

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 0507f8ed-8b64-48af-a6c4-ff7c4211b9e9
begin
	using Pkg
	Pkg.activate(".")
end

# ╔═╡ 3310d2f7-450a-4cd1-9c3a-46d02d23a7c6
using Revise

# ╔═╡ d27b2d72-d264-11eb-0be5-13dcacfd2adc
using DeconvOptim, TestImages, ImageShow, Plots, LinearAlgebra, IndexFunArrays, Noise, FourierTools, SpecialFunctions, FFTW, LaTeXStrings, PlutoUI, Images, Tullio

# ╔═╡ 952b251c-207b-4412-b6e1-268fce1647d9
begin
	img = Float32.(testimage("fabio_gray"));
	img_1D = img[:, 200]
end;

# ╔═╡ fa8cd9c9-a2fd-495c-8d22-ada7bb9c39f6
otf(x, Δx=1) = begin
	x = abs(x)
	if x <= Δx
		SpecialFunctions.jinc(x * Δx *(1-x/Δx)) .* 2 / π * (acos(x/Δx) - x/Δx * sqrt(1-(x/Δx)^2))
	else
		zero(x)
	end
end

# ╔═╡ 9a07bc88-be76-4531-bc91-df0d20c3221c
begin
	x = range(-1.5, 1.5, length=size(img, 1))
	freqs = fftshift(fftfreq(size(img_1D, 1), 1))
	psf = Float32.(DeconvOptim.generate_psf(size(img), 20))
	psf_1D = psf[1, :]
	psf_1D ./= sum(psf_1D)
	otf_1D = abs.(ffts(psf_1D))
end;

# ╔═╡ 24fc86da-0b7c-493c-8977-2a19ef6dc133
img_n = Float32.(poisson(DeconvOptim.conv(img, psf), 1000));

# ╔═╡ 3bda922d-552b-42ec-9055-33a141b5841a
blur(x, otf=otf_1D) = iffts(ffts(x) .* otf)

# ╔═╡ 308e2562-5a20-4082-8a6d-9fb135738c49
md"
Mathematically:

$(S * \text{PSF})(\mathbf r) = \int_{-\infty}^{\infty} S(\mathbf r - \mathbf x) \cdot \text{PSF}(\mathbf x) \, \mathrm d \mathbf x$
"

# ╔═╡ b542187a-3ae7-4430-a81d-f968d9000427
img_blurry = DeconvOptim.conv(img, psf);

# ╔═╡ 88d83629-f1b2-4d69-928f-d4acfbc76b70
reg_1D = TV(num_dims=1);

# ╔═╡ d6d0f436-48fd-4e8a-863c-3e9d3850cc91
begin
	reg_tik = Tikhonov()
	reg_TV = TV()
	reg_GR = DeconvOptim.GR()
end

# ╔═╡ 061faf49-662c-4508-83b4-ddcf0970ed0d
md"### DeconvOptim.jl: Microscopy Image Deconvolution
"

# ╔═╡ 23829201-9756-4ef8-90c9-3917b761fe4b
load("../docs/src/assets/logo.png")

# ╔═╡ 30f21bb8-6d09-4fce-9d2a-568bfaf3ff7a
md"
* **Felix Wechsler:** Master Student at the Leibniz Institute of Photonic Technology in  Jena, Germany
* https://github.com/roflmaostc/DeconvOptim.jl
* `]add DeconvOptim`
"

# ╔═╡ 7a0a44c9-fb07-44e6-9a8b-8f720b84e6f6
md"""### Image Convolution

* Typical description of isotropic blur of an image
* The blurring kernel describes blur
    * In optics/microscopy a finite sized dot called Point Spread Function (**PSF**)
    * often a Gaussian function used in image processing
    * cigarre shaped object for motion blur

Discrete version:

$(S * \text{PSF})[i] = \sum_{m} S[i-m] \cdot \text{PSF}[m]$
"""

# ╔═╡ 49686d9a-1683-428f-83ba-a9131c2ad432
[Gray.(img) Gray.(DeconvOptim.conv(img, psf))]

# ╔═╡ 491ec7cb-1663-4acf-b81f-9acafba2b63d
md"## Convolution Theorem

$(S * \text{PSF})(\mathbf r) = \mathcal{F}^{-1}\bigg[ \mathcal{F}[S] \cdot \mathcal{F}[\text{PSF}]  \bigg]$

* we can express the convolution with a Fast Fourier Transform (FFT) which only takes $\mathcal O(N \log(N))$ operations

* For large kernels (especially in 3D), sliding kernels are slower

*  $\mathcal{F}[\text{PSF}]$ is called the $\text{OTF}$

"

# ╔═╡ b6f9e42e-5a23-40ad-9e73-8f02da48f69c
md"### Optical System act as low pass filter

*  $\text{OTF}$ shows the frequency throughput

"

# ╔═╡ e31f4704-5bce-4c8d-b3a7-873a3460d9f8
plot(x, otf.(x), xlabel="frequency / maximum frequency", ylabel="contrast")

# ╔═╡ dc8cf4e2-a87c-47ea-8247-3ae772852241
md"## Frequency spectrum of blurred sample $Y(\mathbf r)$

Blurred sample:
$Y(\mathbf r) = (S * \text{PSF})(\mathbf r)$
"

# ╔═╡ f8033d13-faee-41f5-bdd6-ae8721e8b8a8
begin
	plot(freqs, abs.(ffts(img_blurry)[:, 128]), yaxis=:log, ylabel="real part of FFT output in AU", xlabel="frequency in 1/px", ylims=(1e-4, 1e2), label="blurred")
	plot!(freqs, abs.(ffts(img)[:, 128]), ylabel="abs of FFT output in AU",
				xlabel="frequency in 1/px", yaxis=:log,
				ylims=(1e-4, 1e4), label="ground truth")
	#plot!(freqs, abs.(ffts(DeconvOptim.conv(img_1D, psf_1D))))
end

# ╔═╡ fe7e0292-31f6-43b5-83a5-a38698a87563
md"## Deconvolution Pipeline

* based on:
    * Zygote.jl
    * Optim.jl
    * Tullio.jl
    * CUDA.jl
"

# ╔═╡ e9ef5ba4-56c0-4595-bc28-e04882f44a9a
load("../docs/src/assets/tex/pipeline.png")

# ╔═╡ 90e5708c-05a4-46e4-b1e4-9a61c96dae32
TV_by_hand(x) = @tullio r = sqrt(1f-8 + abs2(x[i, j] - x[i+1, j]) +
						 	     abs2(x[i, j] - x[i, j+1]))

# ╔═╡ 03139ac5-3525-4fad-abf1-84421492b763
DeconvOptim.generate_TV(4, [1,2, 3], [1,1, 1], 1, 0)[1]

# ╔═╡ 4e84e739-9c59-4939-8b04-aec7dc069d67
md"
### Deconvolve with DeconvOptim.jl

"

# ╔═╡ 9d9a5da2-14df-46e6-b7e6-5a33aade1754
@bind reg_list2 Select(["1" => ("Tikhonov"), "2" => ("Total Variation TV"), "3" => ("Good's Roughness GR")])

# ╔═╡ 5ae1a6a3-4505-4123-9f1e-8a1d4ac0b4e1
reg = [reg_tik, reg_TV, reg_GR][parse(Int, reg_list2)]

# ╔═╡ 15d54e5e-64f4-4a1d-8cc8-9334b2e3784f
md"
iterations =
$(@bind iter Slider(0:50, show_value=true))

λ = $(@bind λ Slider(0:0.001:0.3, show_value=true))

regularizer = $(@bind reg_bool CheckBox())"

# ╔═╡ 803368e6-53fd-4413-b3f5-ffe46ee8983e
img_deconv, res_img = deconvolution(img_blurry, psf, regularizer=reg_bool ? reg : nothing, iterations=iter, λ=λ);

# ╔═╡ e58e1f63-2c81-48fb-866a-4bb70bd428a6
Gray.(img_deconv)

# ╔═╡ 438a6639-bd35-464f-a81d-d98eb65e006e
res_1D, o = deconvolution(real(blur(img_1D)), psf_1D, iterations=iter, regularizer=reg_1D, λ=0.01);

# ╔═╡ 7a16df73-95ad-47f5-907c-6fd23c6000cf
[Gray.(img) Gray.(img_blurry) Gray.(img_deconv)]

# ╔═╡ 5c8030b7-8805-4771-9329-23abb2744544
begin
	plot(freqs, abs.(ffts(img_blurry)[:, 128]), yaxis=:log, ylabel="abs of FFT output in AU", xlabel="frequency in 1/px", ylims=(1e-4, 1e2), label="blurred")
	plot!(freqs, abs.(ffts(img)[:, 128]),  yaxis=:log,
				ylims=(1e-4, 1e4), label="ground truth")
	plot!(freqs, abs.(ffts(img_deconv)[:, 128]), label="deconvolved image")
end

# ╔═╡ ff9af79c-f06c-42ac-9c8d-6f09d2ff4056
[Gray.(img_1D); Gray.(res_1D)];

# ╔═╡ d8aee845-e922-41da-a17e-37ffa3e692f0
md"### Real Microscopy Data"

# ╔═╡ fc93fa3d-8599-463b-b9e6-8043b90e9d63
load("figures/real_data_large.png")

# ╔═╡ 97c45bfd-d1ef-49ad-908d-7360c03b0170
md"Image taken from [DeconvolutionLab2](http://bigwww.epfl.ch/deconvolution/deconvolutionlab2/)."

# ╔═╡ bf37664a-f726-46b6-9592-da419165af91
md"## Conclusion - DeconvOptim.jl
"

# ╔═╡ f8117250-bac8-43a1-aae4-5c8bab3a522d
[Gray.(ones(130, 012)) load("../docs/src/assets/logo.png")]

# ╔═╡ b5a70276-e9b9-46e8-8c67-0cad2cfa19da
md"* Flexible Image Deconvolution Software
* N-dimensional signal deconvolution
* Works both on CPU and GPUs
    * GPUs usually 5-15x speed improvement
"

# ╔═╡ Cell order:
# ╠═3310d2f7-450a-4cd1-9c3a-46d02d23a7c6
# ╠═0507f8ed-8b64-48af-a6c4-ff7c4211b9e9
# ╠═d27b2d72-d264-11eb-0be5-13dcacfd2adc
# ╠═952b251c-207b-4412-b6e1-268fce1647d9
# ╠═24fc86da-0b7c-493c-8977-2a19ef6dc133
# ╠═3bda922d-552b-42ec-9055-33a141b5841a
# ╟─fa8cd9c9-a2fd-495c-8d22-ada7bb9c39f6
# ╠═9a07bc88-be76-4531-bc91-df0d20c3221c
# ╟─308e2562-5a20-4082-8a6d-9fb135738c49
# ╠═b542187a-3ae7-4430-a81d-f968d9000427
# ╠═88d83629-f1b2-4d69-928f-d4acfbc76b70
# ╠═e58e1f63-2c81-48fb-866a-4bb70bd428a6
# ╠═d6d0f436-48fd-4e8a-863c-3e9d3850cc91
# ╠═803368e6-53fd-4413-b3f5-ffe46ee8983e
# ╠═5ae1a6a3-4505-4123-9f1e-8a1d4ac0b4e1
# ╠═438a6639-bd35-464f-a81d-d98eb65e006e
# ╟─061faf49-662c-4508-83b4-ddcf0970ed0d
# ╟─23829201-9756-4ef8-90c9-3917b761fe4b
# ╟─30f21bb8-6d09-4fce-9d2a-568bfaf3ff7a
# ╟─7a0a44c9-fb07-44e6-9a8b-8f720b84e6f6
# ╟─49686d9a-1683-428f-83ba-a9131c2ad432
# ╟─491ec7cb-1663-4acf-b81f-9acafba2b63d
# ╟─b6f9e42e-5a23-40ad-9e73-8f02da48f69c
# ╟─e31f4704-5bce-4c8d-b3a7-873a3460d9f8
# ╟─dc8cf4e2-a87c-47ea-8247-3ae772852241
# ╟─f8033d13-faee-41f5-bdd6-ae8721e8b8a8
# ╟─fe7e0292-31f6-43b5-83a5-a38698a87563
# ╟─e9ef5ba4-56c0-4595-bc28-e04882f44a9a
# ╠═90e5708c-05a4-46e4-b1e4-9a61c96dae32
# ╠═03139ac5-3525-4fad-abf1-84421492b763
# ╟─4e84e739-9c59-4939-8b04-aec7dc069d67
# ╟─9d9a5da2-14df-46e6-b7e6-5a33aade1754
# ╟─15d54e5e-64f4-4a1d-8cc8-9334b2e3784f
# ╟─7a16df73-95ad-47f5-907c-6fd23c6000cf
# ╟─5c8030b7-8805-4771-9329-23abb2744544
# ╟─ff9af79c-f06c-42ac-9c8d-6f09d2ff4056
# ╟─d8aee845-e922-41da-a17e-37ffa3e692f0
# ╟─fc93fa3d-8599-463b-b9e6-8043b90e9d63
# ╟─97c45bfd-d1ef-49ad-908d-7360c03b0170
# ╟─bf37664a-f726-46b6-9592-da419165af91
# ╟─f8117250-bac8-43a1-aae4-5c8bab3a522d
# ╟─b5a70276-e9b9-46e8-8c67-0cad2cfa19da


# Test utililties.jl

using Test, Statistics

println("testing utilites")

tol = 1e-2   # Fairly sloppy, needed for the derivative tests

# nonmaxsuppts - tested in test_cornerfeatures.jl

# derivative 3 5 7
# Generate an image of constant gradient
img = [c for r=1:50, c=1:50]
(gx,gy) = derivative3(img, ("x", "y"))

# Test interior of derivatives matrix matches expected values
@test all(abs.(gx[5:end-5,5:end-5] .- 1) .< tol)
@test all(abs.(gy[5:end-5,5:end-5]) .< tol)

(gx,gy, gxx, gyy, gxy) = derivative5(img, ("x", "y", "xx", "yy", "xy"))
# Test interior of derivatives matrix matches expected values
@test all(abs.(gx[5:end-5,5:end-5] .- 1) .< tol)
@test all(abs.(gy[5:end-5,5:end-5]) .< tol)
@test all(abs.(gxx[5:end-5,5:end-5]) .< tol)
@test all(abs.(gyy[5:end-5,5:end-5]) .< tol)
@test all(abs.(gxy[5:end-5,5:end-5]) .< tol)

(gx,gy, gxx, gyy, gxy) = derivative7(img, ("x", "y", "xx", "yy", "xy"))
# Test interior of derivatives matrix matches expected values
@test all(abs.(gx[5:end-5,5:end-5] .- 1) .< tol)
@test all(abs.(gy[5:end-5,5:end-5]) .< tol)
@test all(abs.(gxx[5:end-5,5:end-5]) .< tol)
@test all(abs.(gyy[5:end-5,5:end-5]) .< tol)
@test all(abs.(gxy[5:end-5,5:end-5]) .< tol)

# gaussfilt
img = zeros(100,100)
img[50,50] = 1
simg = gaussfilt(img,2)
@test abs(sum(simg) - 1) < tol

# dilate1d erode1d
# Generate a 1D signal with a single spike
v = zeros(20)
v[10] = 1
k = 5
vd = dilate1d(v,k)
@test all(vd[8:12] .== 1) && all(vd[1:7] .== 0) && all(vd[13:end] .== 0)

ve = erode1d(vd,k)   # Erode back and check we reconstruct v
@test all(ve .== v)

# imdilate imerode
img = falses(100,100)
img[50,50] = true
dimg = imdilate(img, "rect", [5, 3])
@test all(dimg[47:47,48:52] .== [0 0 0 0 0])
@test all(dimg[50:50,48:52] .== [0 1 1 1 0])
@test all(dimg[53:53,48:52] .== [0 0 0 0 0])
@test all(dimg[47:53,50] .== [0, 1, 1, 1, 1, 1, 0])

eimg = imerode(dimg, "rect", [5, 3])
@test all(eimg[49:51,49:51] .== [0 0 0; 0 1 0; 0 0 0])

# circularstruct

# floatxy

# matchbycorrelation

# grey2census

# imgnormalise
a = rand(5,3)
an = imgnormalise(a)
@test isapprox(minimum(an), 0)
@test isapprox(maximum(an), 1)

reqmean = 5
reqvar = 7
an = imgnormalise(a, reqmean, reqvar)
@test isapprox(mean(an), reqmean)
@test isapprox(var(an), reqvar)

# histtruncate
rows = 100;
cols = 200;
img = rand(rows, cols)
lHistCut = 2
uHistCut = 4
htimg = histtruncate(img, lHistCut, uHistCut)
minval = minimum(htimg)
maxval = maximum(htimg)
# Check that the number of saturated pixels at each extreme is
# approximately lHistCut and uHistCut. Note some integer rounding has
# to occur.
v = findall(abs.(htimg .- minval).<eps())
@test abs(length(v)/(rows*cols) * 100 - lHistCut) < 1

v = findall(abs.(htimg .- maxval).<eps())
@test abs(length(v)/(rows*cols) * 100 - uHistCut) < 1


function simsys(sys,nt,seed)
	###
	# simulated systems
	###

	# random seed for the simulations
	srand(seed)

	if sys==1
		#
		# system 1

		# ref: "Subspace Identification for Linear Systems"
		# P. Van Overschee and B. De Moor, page 81
		#
		A = [0.6 0.6 0.;-0.6 0.6 0.;0. 0. 0.4]
		K = [0.1706 -0.1507 0.2772]'
		C = [0.7831 0.5351 0.9701]
		R = 6.3663
		Q = K * R * K'

		et = sqrt(R)*randn(1,nt)
		w = K*et

		ns = 3
 		x = zeros(ns,nt)
		x[:,1] = w[:,1]
		for t in 1:nt-1
			x[:,t+1] = A * x[:,t] + w[:,t]
		end
		y = C * x + et

		salida = y, A, K, C, Q, R
		#
	elseif sys==2
		###
		# ref: "Subspace Identification for Linear Systems"
		# P. Van Overschee and B. De Moor, page 115
		###
		A = [0.6 0.6 0.;-0.6 0.6 0.;0. 0. 0.7]
		B = [1.6161 -0.3481 2.6319]'
		K = [-1.1472 -1.5204 -3.1993]'
		C = [-0.4373 -0.5046 -0.0936]
		D = -0.7759
		R = 0.0432
		Q = K * R * K'

		et = sqrt(R)*randn(1,nt)
		w = K*et

		# input = AR(1)
		u = zeros(1,nt)
		a = randn(nt)
		for t in 1:nt-1
			u[t+1] = 0.3*u[t] + a[t]
		end

		# states
		ns = size(A,1)
		x = zeros(ns,nt)
		x[:,1] = w[:,1]
		for t in 1:nt-1
			x[:,t+1:t+1] = A * x[:,t:t] + B * u[t] + w[:,t:t]
		end
		# output
		y = C * x + D * u + et

		salida = y, u, A, B, K, C, D, Q, R
	end

	return salida

end


using Alistair
using Base.Test

# TEST OPTIONS
const TEST_SIZE_1 = 10
const PRINT_RESULTS = true
RESULTS = Dict()

beta = [0.7; 0.45]
X = hcat(ones(TEST_SIZE_1,1), collect(1.0:TEST_SIZE_1))
Y = X * beta + 0.1*randn(TEST_SIZE_1, 1)
# Y = [1.0 2.0 1.3 3.75 2.25]'

#=
if false
@__test issingletonarray(X) == false
@__test isrowarray(X) == false
@__test iscolumnarray(X) == true
@__test hasintercept(X) == false

@__test aprod(X, 1.0) == X
@__test aprod(1, X) == X
@__test aprod(X, [1]) == X
@__test aprod([1], X) == X
@__test aprod(X', Y) == X' * Y
end
=#

if true
# Test for all types
RESULTS["TIME_OLS_INIT"] = @elapsed OLSRegression = OLS(intercept=true, robust=false)

# OLS test:
RESULTS["TIME_OLS_SOLVE "] = @elapsed res = linregress(X, Y, OLSRegression)
#@__test squeeze(res.beta,2) == beta
#println("OLSRegression")
PRINT_RESULTS ? println(res) : 0

OLSRegression = OLS(intercept=true, robust=HCEVariance())
RESULTS["TIME_OLS_SOLVE_HCE"] = @elapsed res = solve(X, Y, OLSRegression)
#@__test res.beta == [0.785; 0.425]
#println("OLSRegression")
PRINT_RESULTS ? println(res) : 0

# FGLS test
RESULTS["TIME_FGLS_INIT"] = @elapsed FGLSRegression = FGLS(intercept=true, robust=HCEVariance())
RESULTS["TIME_FGLS_SOLVE"] = @elapsed res = linregress(X, Y, FGLSRegression)
#println("FGLSRegression")
PRINT_RESULTS ? println(res) : 0

# IteratedFGLS test
RESULTS["TIME_IFGLS_INIT"] = @elapsed IteratedFGLSRegression = IteratedFGLS(10, intercept=true, robust=HCEVariance())
RESULTS["TIME_IFGLS_SOLVE"] = @elapsed res = linregress(X, Y, IteratedFGLSRegression)
#println("IteratedFGLS")
PRINT_RESULTS ? println(res) : 0
end

# Non-Linear test
# Optimize test

f(X_, beta_) = exp.(beta_[1]*X_[:,1] + beta_[2]*log.(X_[:,2]))

Ynlin = f(X, beta) + 0.1*randn(TEST_SIZE_1, 1)

RESULTS["TIME_OPTIM_INIT"] = @elapsed OptimizeRegression = Optimize(f, [1.0, 1.0])
RESULTS["TIME_OPTIM_SOLVE"] = @elapsed res = nlinfit(X, Ynlin, OptimizeRegression)
PRINT_RESULTS ? println(res) : 0

# Print timing results:
println("#-- TIMING RESULTS --#")
present(RESULTS)


# See Floater and Hormann 2007 Eqn. (18) (Actually, an unnumbered equation
# that is 3 equations later)
# UPDATE: used formula (18) from Hormann's "Barycentric Interpolation"
# part is that the interpolation seems to work OK for either. (We do know
# that the sign of the weights must alternate)
"""
    floater_weights(x, d=0)

Calculate the Floater-Hormann weights for grid x and mixing degree d.  For
many applications d = 3 or 4 works well.

    if we had a regular grid the values would be...
    x: points at which we have function values
    d: mixing degree        weights
    d = 0                 1,1,...,1,1
    d = 1               1,2,2,...,2,2,1
    d = 2             1,3,4,4,...,4,4,3,1
    d = 3           1,4,7,8,8,...,8,8,7,4,1
    d = 4     1,5,11,15,16,16,...,16,16,15,11,5,1

But note, if you do have a regularly spaced grid, it is much better to call
the floater_weights(n, order)
"""
function floater_weights(x::Array{T, 1}, d=0) where {T}
    n = length(x)
    d > n && error("Mixing coefficient must be less than node set size")
    w = zeros(eltype(x), n)
    @inbounds for k in 1:n
        ws = 0.0
        for i in max(1, k-d):min(n-d, k)
            lwp = 0.0
            for j in i:min(n, i+d)
                j == k && continue
                lwp += log(abs(x[k] - x[j]))
            end
            ws += 1.0 / exp(lwp)
        end
        w[k] = k % 2 == 0 ? ws : -ws
    end
    return w ./ norm(w, Inf)
end

# Calculate the FH weights for a uniform grid of size n.  We don't
# worry about the actual grid size, since these weights are intended
# for use in the barycentric (2nd kind) interpolation formula.
# This is from the last equation (unnumbered) in Section 4 of the FH paper
function floater_weights(n::Int, d=0)
    w = Vector{Float64}(undef, n)
    @inbounds for k in 1:n
        s = 0.0
        for i in max(1, k-d):min(n-d, k)
            s += binomial(d, k-i)
        end
        w[k] = (k-d) % 2 == 0 ? s : -s
    end
    return w
end

# Recognize a step range and handle it efficiently
function floater_weights(x::T, d=0) where {T <:StepRangeLen}
    return floater_weights(length(x), d)
end

"""Return the n + 1 chebychev points on the [-1, 1] interval

  A shifted sin is used to make sure the points are symmetric
"""
function chebpts(n)
    m = n - 1
    x = sinpi.([-m:2:m;] ./ 2m)
    return x
end

function chebwts(n)
    w = ones(n)
    w[2:2:end] .= -1.0
    w[1] *= 0.5
    w[end] *= 0.5
    return w
end

# This is faster.  The exp(sum(log.(abs.(...))) version is ~30X slower do n = 10000
"""The weights for Barycentric Lagrange interpolation given the x grid of evaluations"""
function lagrange_weights(x::Vector{T}) where {T}
    n = length(x)
    w = ones(T, n)
    t = copy(w)
    @inbounds for i ∈ 1:n
        t .= x[i] .- x
        t[i] = one(eltype(x))
        w[i] = one(eltype(x)) / prod(t)
        # Here is the "safe" version
        # w[i] = one(eltype(x)) / exp(sum(log.(abs.(t))))
        # neg = count(t .< zero(eltype(x))) % 2 == 1
        # w[i] = neg ? -w[i] : w[i]
    end
    # Normalize weights
    return  w ./ norm(w, Inf)
end


using DReal
using Base.Test

# Ackley's: bounds -5 <= x,y <= 5, min: f(0,0) = 0
begin
println("Trying Ackleys")
ackleys(x::Array) = -20*exp(-0.2*sqrt(0.5(x[1]*x[1] + x[2]*x[2]))) - exp(0.5(cos(2*π*x[1])+cos(2*π*x[2]))) + e + 20
x = Var(Float64,"x2",-5.12,5.12)
y = Var(Float64,"y2",-5.12,5.12)
f = Var(Float64,"f2",-10.,10.)
add!(f == ackleys([x,y]))
cost, optimal_model = minimize(f,x,y; lb=-10.,ub = 10.)

reset_ctx!(default_global_context)

# Commented out by Soonho Kong, 2016/06/09
# The following optimization query returns UNSAT, which is wrong.
# Until I fix the problem, we don't run this test
#
# # # Beales function
# println("Trying Beales")
# beale(x::Array) =  (1.5 - x[1] + x[1]x[2])^2 +
#                 (2.25 - x[1] + x[1]x[2]^2)^2 +
#                 (2.625 - x[1] + x[1]x[2]^3)^2
# x = Var(Float64,"x3",-5.12,5.12)
# y = Var(Float64,"y3",-5.12,5.12)
# f = Var(Float64,"f3",-10.,10.)
# add!(f == beale([x,y]))
# cost, optimal_model = minimize(f,x,y; lb=-10.,ub = 10.)
# @show cost, optimal_model
#
# reset_global_ctx!()

# # rastrigin
println("Trying rastrigin")

rastrigin(x::Array) = 10 * length(x) + sum([xi*xi - 10*cos(2pi*xi) for xi in x])
x = Var(Float64,"x4",-5.12,5.12)
y = Var(Float64,"y4",-5.12,5.12)
f = Var(Float64,"f4",-10.,10.)
add!(f == rastrigin([x,y]))
cost, optimal_model = minimize(f,x,y; lb=-10.,ub = 10.)
@show cost, optimal_model

reset_global_ctx!()

end


"""
    Uniform(a,b)

The *continuous uniform distribution* over an interval ``[a, b]`` has probability density function

```math
f(x; a, b) = \\frac{1}{b - a}, \\quad a \\le x \\le b
```

```julia
Uniform()        # Uniform distribution over [0, 1]
Uniform(a, b)    # Uniform distribution over [a, b]

params(d)        # Get the parameters, i.e. (a, b)
minimum(d)       # Get the lower bound, i.e. a
maximum(d)       # Get the upper bound, i.e. b
location(d)      # Get the location parameter, i.e. a
scale(d)         # Get the scale parameter, i.e. b - a
```

External links

* [Uniform distribution (continuous) on Wikipedia](http://en.wikipedia.org/wiki/Uniform_distribution_(continuous))

"""
struct Uniform{T<:Real} <: ContinuousUnivariateDistribution
    a::T
    b::T
    Uniform{T}(a::T, b::T) where {T <: Real} = new{T}(a, b)
end

function Uniform(a::T, b::T; check_args=true) where {T <: Real}
    check_args && @check_args(Uniform, a < b)
    return Uniform{T}(a, b)
end

Uniform(a::Real, b::Real) = Uniform(promote(a, b)...)
Uniform(a::Integer, b::Integer) = Uniform(float(a), float(b))
Uniform() = Uniform(0.0, 1.0, check_args=false)

function Uniform(::Type{T}, a, b) where {T <: Real}
    return Uniform(T(a), T(b))
end

@distr_support Uniform d.a d.b

#### Conversions
convert(::Type{Uniform{T}}, a::Real, b::Real) where {T<:Real} = Uniform(T(a), T(b))
convert(::Type{Uniform{T}}, d::Uniform{S}) where {T<:Real, S<:Real} = Uniform(T(d.a), T(d.b), check_args=false)

#### Parameters

params(d::Uniform) = (d.a, d.b)
partype(::Uniform{T}) where {T<:Real} = T

location(d::Uniform) = d.a
scale(d::Uniform) = d.b - d.a

#### Statistics

mean(d::Uniform) = middle(d.a, d.b)
median(d::Uniform) = mean(d)
mode(d::Uniform) = mean(d)
modes(d::Uniform) = Float64[]

var(d::Uniform) = (w = d.b - d.a; w^2 / 12)

skewness(d::Uniform{T}) where {T<:Real} = zero(T)
kurtosis(d::Uniform{T}) where {T<:Real} = -6/5*one(T)

entropy(d::Uniform) = log(d.b - d.a)

#### Evaluation

pdf(d::Uniform{T}, x::Real) where {T<:Real} = insupport(d, x) ? 1 / (d.b - d.a) : zero(T)
logpdf(d::Uniform{T}, x::Real) where {T<:Real} = insupport(d, x) ? -log(d.b - d.a) : -T(Inf)

function cdf(d::Uniform{T}, x::Real) where T<:Real
    (a, b) = params(d)
    x <= a ? zero(T) :
    x >= d.b ? one(T) : (x - a) / (b - a)
end

function ccdf(d::Uniform{T}, x::Real) where T<:Real
    (a, b) = params(d)
    x <= a ? one(T) :
    x >= d.b ? zero(T) : (b - x) / (b - a)
end

quantile(d::Uniform, p::Real) = d.a + p * (d.b - d.a)
cquantile(d::Uniform, p::Real) = d.b + p * (d.a - d.b)

function mgf(d::Uniform, t::Real)
    (a, b) = params(d)
    u = (b - a) * t / 2
    u == zero(u) && return one(u)
    v = (a + b) * t / 2
    exp(v) * (sinh(u) / u)
end

function cf(d::Uniform, t::Real)
    (a, b) = params(d)
    u = (b - a) * t / 2
    u == zero(u) && return complex(one(u))
    v = (a + b) * t / 2
    cis(v) * (sin(u) / u)
end

#### Sampling

rand(rng::AbstractRNG, d::Uniform) = d.a + (d.b - d.a) * rand(rng)

#### Fitting

function fit_mle(::Type{<:Uniform}, x::AbstractArray{T}) where T<:Real
    if isempty(x)
        throw(ArgumentError("x cannot be empty."))
    end

    xmin = xmax = x[1]
    for i = 2:length(x)
        xi = x[i]
        if xi < xmin
            xmin = xi
        elseif xi > xmax
            xmax = xi
        end
    end

    Uniform(xmin, xmax)
end


__precompile__()
# ----------------------------------------------------------------------------
module Knapsack
#=
    see also:
        - the walkthrough at http://vladium.com/tutorials/study_julia_with_me/knapsack_benchmark/
        - https://en.wikipedia.org/wiki/Knapsack_problem#0/1_knapsack_problem

    running this benchmark from command line:

    >julia Knapsack.jl
=#
# ............................................................................
"""
    we need a couple of integers to describe a knapsack item, its value and weight.
    a tuple would work, but a struct makes for more readable code
"""
struct Item
    value   ::Int64; # actually, being so type-specific is not necessary in Julia
    weight  ::Int64;
end
# ............................................................................
"""
    first version of the knapsack solver; uses a full matrix for V and will
    not scale to large problems
"""
function opt_value_BAD(W ::Int64, items ::Array{Item}) ::Int64
    n = length(items)

    # V[w,j] stores opt value achievable with capacity 'w' and using items '1..j':

    V = Array{Int64}(undef, W, n) # Wxn matrix with uninitialized storage

    # initialize first column v[:, 1] to trivial single-item solutions:
    # (note "broadcast assignment" syntax)

    V[:, 1] .= 0
    V[items[1].weight:end, 1] .= items[1].value

    # do a pass through remaining columns:

    for j in 2 : n
        itemⱼ = items[j]
        for w in 1 : W
            V_without_itemⱼ = V[w, j - 1]
            V_allow_itemⱼ = (w < itemⱼ.weight
                ? V_without_itemⱼ
                : (itemⱼ.value + (w ≠ itemⱼ.weight ? V[w - itemⱼ.weight, j - 1] : 0)))
            V[w, j] = max(V_allow_itemⱼ, V_without_itemⱼ)
        end
    end

    return V[W, n]
end
# ............................................................................
"""
    second version of the knapsack solver; uses two column buffers for V
"""
function opt_value(W ::Int64, items ::Array{Item}) ::Int64
    n = length(items)

    # V[w] stores opt value achievable with capacity 'w' and using items '1..j':

    V = zeros(Int64, W) # single column of size W, zero-initialized
    V_prev = Array{Int64}(undef, W) # single column of size W, uninitialized storage

    # initialize 'v' to trivial single-item solutions:
    # (note "broadcast assignment" syntax)

    V[items[1].weight:end] .= items[1].value

    # do a pass through remaining columns:

    for j in 2 : n
        V, V_prev = V_prev, V
        itemⱼ = items[j]
        for w in 1 : W
            V_without_itemⱼ = V_prev[w]
            V_allow_itemⱼ = (w < itemⱼ.weight
                ? V_without_itemⱼ
                : (itemⱼ.value + (w ≠ itemⱼ.weight ? V_prev[w - itemⱼ.weight] : 0)))
            V[w] = max(V_allow_itemⱼ, V_without_itemⱼ)
        end
    end

    return V[W]
end
# ............................................................................
"""
    a simple but good quality [xorshift RNG](https://en.wikipedia.org/wiki/Xorshift)
    that can be ported to all benchmark languages to ensure equivalent test problem
    instances across all ports
"""
function xorshift_rand(seed)
    @assert seed ≠ 0
    x ::UInt64 = seed # want unsigned bit arithmetic

    return function _next()
        x ⊻= (x << 13) # note that python would need 'nonlocal' here
        x ⊻= (x >>> 7)
        x ⊻= (x << 17)
        return x
    end
end

Problem = Tuple{Int64, Array{Item}} # use a type alias for less typing

"""
    generate a randomized knapsack problem definition of scale 'W'
"""
function make_random_data(W ::Int64, seed ::Int64) ::Problem
    @assert W > 1000 # don't test on very small problems
    n = W ÷ 100
    rng = xorshift_rand(seed)

    items = Array{Item}(undef, n)
    for i in 1 : n
        v = rng() % 1000
        w = 1 + rng() % 2W
        items[i] = Item(v, w)
    end

    return W, items
end
# ............................................................................
"""
    you can use this function to load any knapsack problem defition from a file
    instead of using make_random_data();
    the format is:

     1. first line: <W>
     2. n subsequent lines: <value> <weight> (separated by spaces)
"""
function load_data(in_file ::String) ::Problem
    W = -1
    items = Item[]
    open(in_file) do fd
        for line in eachline(fd)
            line = strip(line)
            if isempty(line)
                break
            end
            if W < 0 # reading first line
                W = parse(Int, line)
            else # reading items
                tokens = split(line)
                @assert length(tokens) == 2
                item = Item(parse(Int, tokens[1]), parse(Int, tokens[2]))

                push!(items, item)
            end
        end
    end

    return W, items
end
# ............................................................................

function run(repeats = 5)
    @assert repeats > 1

    times = zeros(Float64, repeats)
    seed = 12345

    for W in [5_000, 10_000, 20_000, 40_000, 80_000]
        for repeat in 1 : repeats
            W, items = make_random_data(W, seed += 1)
            times[repeat] = @elapsed V = opt_value(W, items)
            # println("V = $(V), time = $(times[repeat])")
        end
        sort!(times)
        println("julia, ", W, ", ", times[(repeats + 1) ÷ 2]) # report W and median time
    end
end

# comment this out if you don't want to run the benchmark on module load:

run()

end # of module
# ----------------------------------------------------------------------------


# Jonas August 1/2022

#=

This problem was asked by Google.

The power set of a set is the set of all its subsets. Write a function that, given a set, generates its power set.

For example, given the set {1, 2, 3}, it should return {{}, {1}, {2}, {3}, {1, 2}, {1, 3}, {2, 3}, {1, 2, 3}}.

You may also use a list or array to represent a set.

=#

#=

The possible subsets are equivalent to all possible binary indicator vectors, where each bit represents presence (1) or absence (0) of a given element.

Thus 2^n possibilities entries in the power set.

What about repeated elements, and thus indistinguishable subsets?

I will assume that we are to compute the power set _of a set_ (which has distinct elements, by definition), that the array representing the set has unique elements (no repeated elements).  I can test this.

=#

#=

Another approach (more complex, not implemented):

- Let n be length of v.
- Let p be the power set of v, initially empty.
- Each entry in p will be a subset of the elements of v.
- Generate all n-bit patterns, and for each bit pattern select the components of v based where there are 1's in the pattern, and add it to p.

How to generate bit pattern?
Count from 0 through 2^n-1 in an (unsigned?) int of sufficient bits, or create a length n bit array, and "count".  If we use int representation, then counting is built-in.  If we make binary array, then we have to implement increment (with carry).

=#

"Construct array of all subsets of v"
function powerset(v)
   # TODO: assert that there are no repeated elements in v
   p = [[]]
   for item in v
       n = []   # new sets with item appended
       for set in p
           push!(n, push!(copy(set), item))
       end
       push!(p, n...)   # combine p and n
   end
   p
end

using Test
@test powerset([]) == [[]]
@test powerset([1]) == [[], [1]]
@test powerset([1,2]) == [[],[1],[2],[1,2]]
# Set needed because ordering in computed powerset is different than test case.
@test Set(powerset([1,2,3])) == Set([[],[1],[2],[3],[1,2],[1,3],[2,3],[1,2,3]])


### A Pluto.jl notebook ###
# v0.12.16

using Markdown
using InteractiveUtils

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

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

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

# ╔═╡ fa68607e-22e5-11eb-0558-c9a4d9f77426
n, N = 30, 10^6

# ╔═╡ d599e1d4-3a99-11eb-1ab2-637caef1df0d
begin
	dist1 = Uniform(1-sqrt(3),1+sqrt(3))
	dist2 = Exponential(1)
	dist3 = Normal(1,1)

	data1 = [mean(rand(dist1,n)) for _ in 1:N]
	data2 = [mean(rand(dist2,n)) for _ in 1:N]
	data3 = [mean(rand(dist3,n)) for _ in 1:N]
end

# ╔═╡ d59a153c-3a99-11eb-1cfe-f54834ffbc5c
begin
	stephist([data1 data2 data3], bins=100,
		c=[:blue :red :green], xlabel = "x", ylabel = "Density",
		label=["Average of Uniforms" "Average of Exponentials" "Average of Normals"],
		normed=true, xlims=(0,2), ylims=(0,2.5))
end

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

# ╔═╡ Cell order:
# ╟─ca5fbdc8-22e5-11eb-0a60-cb0a127a3ca5
# ╠═e077ab16-22e5-11eb-15ed-5f0654bf7328
# ╠═e75a9a06-22e5-11eb-2e36-8d4f62b830ed
# ╠═fa68607e-22e5-11eb-0558-c9a4d9f77426
# ╠═d599e1d4-3a99-11eb-1ab2-637caef1df0d
# ╠═d59a153c-3a99-11eb-1cfe-f54834ffbc5c
# ╟─475ff888-22e6-11eb-2354-09f8f40a8e12


"""
    +   +---+---+
    | 1   2   3 |
    +---+   +   +
    | 4   5 | 6
    +---+---+---+

    julia> const graph = [
            0 1 0 0 0 0;
            1 0 1 0 1 0;
            0 1 0 0 0 1;
            0 0 0 0 1 0;
            0 1 0 1 0 0;
            0 0 1 0 0 0]

    julia> dist, path = dijkstra(graph, 1)
    (Dict(4=>3,2=>1,3=>2,5=>2,6=>3,1=>0), Dict(4=>5,2=>1,3=>2,5=>2,6=>3,1=>0))

    julia> printpath(path, 6) # Display solution of the maze
    1 -> 2 -> 3 -> 6

"""
function dijkstra(graph, source::Int=1)
    # ensure that the adjacency matrix is squared
    @assert size(graph, 1) == size(graph, 2)
    inf = typemax(Int64)
    n   = size(graph, 1)

    Q    = IntSet(1:n)                  # Set of unvisited nodes
    dist = Dict(n => inf for n in Q)    # Unknown distance function from source to v
    prev = Dict(n => 0   for n in Q)    # Previous node in optimal path from source
    dist[source] = 0                    # Distance from source to source

    function _minimumdist(nodes) # Find the less distant node among nodes
        kmin, vmin = nothing, inf
        for (k, v) in dist
            if k ∈ nodes && v ≤ vmin
                kmin, vmin = k, v
            end
        end
        return kmin
    end
    # Until all nodes are visited...
    while !isempty(Q)
        u = _minimumdist(Q)         # Vertex in Q with smallest dist[]
        pop!(Q, u)
        if dist[u] == inf break end # All remaining vertices are inaccessible from source
        for v in 1:n                # Each neighbor v of u
            if graph[u, v] != 0 && v ∈ Q # where v has not yet been visited
                alt = dist[u] + graph[u, v]
                if alt < dist[v]    # Relax (u, v, a)
                    dist[v] = alt
                    prev[v] = u
                end
            end
        end
    end

    return dist, prev
end

function printpath(prev::Dict, target::Int)
    path = "$target"
    while prev[target] != 0
        target = prev[target]
        path = "$target -> " * path
    end
    println(path)
end

const graph = [
    0 1 0 0 0 0;
    1 0 1 0 1 0;
    0 1 0 0 0 1;
    0 0 0 0 1 0;
    0 1 0 1 0 0;
    0 0 1 0 0 0]

dist, path = dijkstra(graph)
printpath(path, 6)


struct Polynomial{R}
    coeffs::Vector{R}
end

function (p::Polynomial)(x)
    v = p.coeffs[end]
    for i = (length(p.coeffs)-1):-1:1
        v = v*x + p.coeffs[i]
    end
    return v
end

(p::Polynomial)() = p(9)

p = Polynomial([1,9,250])

println(p(3))

println(p())


using MoreBasisFunctions: vandermonde_matrix, vandermonde_matrix_inverse

n = 11
x = collect(range(0., stop=1., length=n))
y = rand(n)

V = vandermonde_matrix(x)
W = vandermonde_matrix_inverse(x)

x = V\y
a = W*y

@test maximum(abs.((a .- x) ./ a)) < eps(Float32)


"""
    crandnToep(t)

Generate `x ~ N_C(0,T)` with Toeplitz cov. `T` that has `t` as
its first row.
"""
function crandnToep(t)
    C = toeplitzHe(t)
    y = sqrtm(C)*crandn(length(t))
end

"""
    toeplitz(t::Vector)

Generate an `n` x `n` Toeplitz matrix `C` with generating vector `t` of length `2n-1`.
`C[1,1] = t[n]`
`C[1,n] = t[2n-1]`
`C[n,1] = t[1]`
"""
function toeplitz(t::Vector)
    n = Int(floor(length(t)/2)+1)

    C = [t[j-i+n] for i=1:n,j=1:n]
end

function applyToep(t,x)
    x2 = [x;zeros(x)]
    t2 = [conj(t);0;t[end:-1:2]]
    y = ifft(fft(t2,1).*fft(x2,1))
    y = y[1:length(x)]
end

"""
    toeplitzHe(t::Vector)

Generate an `n` x `n` Hermitian Toeplitz matrix `C` with generating vector `t` of length `n`.
The vector `t` defines the first row of `C` and `t[1]` must be real.
"""
function toeplitzHe(t::Vector)
    t_full = [ conj(t[end:-1:2]); t[1:end] ]
    C = toeplitz(t_full)
end

"""
    best_circulant_approximation(t::Vector)

If `t` is the length `n` generating vector of a Hermitian Toeplitz matrix `T`, then
`c` is the length `n` generating vector of a Hermitian circulant matrix `C`, which
is the best approximation of `T`.
"""
function best_circulant_approximation(t::Vector)
    n = length(t);
    lins = ((n:-2:-n+2)./n);
    t1 = t;
    t2 = cat(1,[0], conj(reverse(t[2:end])));
    c = 0.5*n*real(ifft(lins .*(t1-t2) + t1 + t2));
end

"""
    best_circulant_approximation(t::Matrix)

If each `t[:,m]` is the length `n` generating vector of a Hermitian Toeplitz matrix `T_m`, then
each `c[:,m]` is the length `n` generating vector of a Hermitian circulant matrix `C_m`, which
is the best approximation of `T_m`.
"""
function best_circulant_approximation(t::Matrix)
    n = size(t,1);
    m = size(t,2);
    lins = ((n:-2:-n+2)./n);
    t1 = t;
    t2 = cat(1,zeros(1,m), conj(flipdim(t[2:end,:],1)));
    c = 0.5*n*real(ifft(lins .*(t1-t2) + t1 + t2,1));
end


using Pkg
Pkg.status()
using Plots
using DifferentialEquations

#----------------------------------------------
# Set Parameters:
# Membrane capacitance and potential
C_m = 1.0
V_m = -65.0

# Max sodium conductance and membrane potential
g_Na = 120.0
V_Na = 56.0

# Max potassium conductance and membrane potential
g_K = 36.0
V_K = -77.0

# Max leak conductance and membrane potential
g_l = 0.3
V_l = -60.0

# Time to integrate over
time = 0.0:0.1:1000

# ion variables
m = 0.05
h = 0.6
n = 0.32

#----------------------------------------------
# Potassium Rate Functoins
a_n(V_m) =  0.01 * (V_m + 55) / (1 - exp(-(V_m + 55.0) / 10))

B_n(V_m) = 0.125 * exp(-(V_m + 65.0) / 80)

#----------------------------------------------
# Sodium Rate Functions
a_m(V_m) = 0.1 * (40 + V_m) / (1.0 - exp(-(V_m + 40.0) / 10.0))

a_h(V_m) = 0.07 * exp(-(V_m + 65.0) / 20.0)

B_m(V_m) = 4 * exp(-(V_m + 65.0) / 18.0)

B_h(V_m) = 1 / (exp(-(35 + V_m) / 10) + 1)

#----------------------------------------------
# Ion Currents
IK(g_K, n, V_m, V_K) = g_K * n^(4) * (V_m - V_K)

INa(g_Na, m, h, V_m, V_Na) =  g_Na * m^(3) * h * (V_m - V_Na)

LI(g_l, V_m, V_l) = g_l * (V_m - V_l)

#----------------------------------------------
# Input current
I_Inj = 20

input(t::AbstractFloat) = ifelse(0 < t < 200, zero(t), ifelse(200<t<300, one(t)*I_Inj, zero(t)))

#----------------------------------------------
# Hodking Huxley model
function hodgkin_huxley!(du, u, p, t)
    # dVdt, dndt, dmdt, dhdt = u
    #
    #
    V_m, n, m, h = u
    g_K, V_K, g_Na, V_Na, g_l, C_m = p

    # Total current through the membrane
    du[1] = input(t) - IK(g_K,n,V_m,V_K) - INa(g_Na, m, h, V_m, V_Na) - LI(g_l, V_m, V_l) / C_m

    # Derivative of n, potassium channel activation, w.r.t. time
    du[2] = a_n(V_m) * (1 - n) - B_n(V_m) * n

    # Derivative of m, sodium channel activion, w.r.t. time
    du[3] = a_m(V_m) * (1 - m) - B_m(V_m) * m

    # Derivative of h, leaky channel in-activion, w.r.t. time

    du[4] = a_h(V_m) * (1 - h) - B_h(V_m) * h

end

#----------------------------------------------
# Define and solve ODE problem
p = [g_K, V_K, g_Na, V_Na, g_l, C_m]
u0 = [V_m, n, m, h]
tspan = (0.0,500.0)

prob = ODEProblem(hodgkin_huxley!,u0,tspan,p)
sol = solve(prob)

#----------------------------------------------
# Plotting Solutions
plot(sol, vars=1)
#animate(sol, vars=(1), every=20)

plot(sol, vars=(2,3))

plot(sol, vars=(2,3,4))


function lin_inter(grid,vals)
    function func(x::Real)
        if x <= grid[1]
            return (vals[2] - vals[1])/(grid[2] - grid[1])*(x - grid[1]) + vals[1]
        elseif x >= grid[end]
            return (vals[end] - vals[end-1])/(grid[end] - grid[end-1])*(x - grid[end]) + vals[end]
        else
            index = searchsortedlast(grid,x)
            return (vals[index+1] - vals[index])/(grid[index+1] - grid[index])*(x - grid[index]) + vals[index]
        end
    end

    function func{T<:Real}(x::AbstractVector{T})
        m = length(x)
        out = Array(Any,m)
        for i in 1:m
            out[i] = func(x[i])
        end
        return out
    end
    return func

end


"""
Obtain coefficients (A, B) of the function Se = B(x - A)^2 for
Se(1.2) = B(1.2 - A)^2 and Se(1.0) = B(1.0 - A)^2 as:
Se(1.0) = Se(1.2)/(1.2 - A)^2 * (1.0 - A)^2 that yields
(1.2 - A)^2 Se(1.0) = Se(1.2) * (1.0 - A)^2 or expanding:
(Se(1.2) - Se(1.0)) A^2 + (2.4 Se(1.0) - 2 Se(1.2)) A + (Se(1.2) - 1.44 Se(1.0)) = 0
and uses the negative solution of the quadratic equation
"""
function get_quadratic_saturation(Se::Tuple{Float64, Float64})
    if Se[1] == 0.0 || Se[2] == 0.0
        @warn "$(Se[1]) or $(Se[2]) equals to zero. Ignoring saturation."
        return (0.0, 0.0)
    end
    if Se[2] <= Se[1]
        throw(IS.ConflictingInputsError("Se(1.2) <= Se(1.0). Saturation data is inconsistent."))
    end
    E1 = 1.0
    E12 = 1.2
    Sat_Se1 = Se[1]
    Sat_Se12 = Se[2]
    Sat_a = sqrt(E1 * Sat_Se1 / (E12 * Sat_Se12)) * ((Sat_Se12 > 0) + (Sat_Se12 < 0))
    Sat_A = E12 - (E1 - E12) / (Sat_a - 1)
    Sat_B = E12 * Sat_Se12 * (Sat_a - 1)^2 * ((Sat_a > 0) + (Sat_a < 0)) / (E1 - E12)^2

    return (Sat_A, Sat_B)
end

"""
Obtain coefficients (A, B) of the function Se = Bx^A for
Se(1.2) = B(1.2)^A and Se(1.0) = B(1.0)^A as:
B = Se(1.0) and hence
(1.2)^A = Se(1.2)/B -> A = log(Se(1.2)/B) / log(1.2)
"""
function get_exponential_saturation(Se::Tuple{Float64, Float64})
    if Se[1] == 0.0 || Se[2] == 0.0
        @warn "$(Se[1]) or $(Se[2]) equals to zero. Ignoring saturation."
        return (0.0, 0.0)
    end
    if Se[2] <= Se[1]
        throw(IS.ConflictingInputsError("Se(1.2) <= Se(1.0). Saturation data is inconsistent."))
    end
    B = Se[1]
    A = log(Se[2] / B) / log(1.2)

    return (A, B)
end


"""
    borehole(param_values::AbstractArray{<:Number, N})

Return an `N` by 1 array of results of the Borehole function (http://www.sfu.ca/~ssurjano/borehole.html)
using the `N` by 8 array of `param_values`.
"""
function borehole(param_values::AbstractArray{<:Number, N}) where N

    numruns = size(param_values, 1)
    Y = zeros(numruns, 1)

    for i in 1:numruns

        # parameters
        rw =    param_values[i, 1] # radius of borehole (m)
        r =     param_values[i, 2] # radius of influence (m)
        Tu =    param_values[i, 3] # transmissivity of upper aquifer (m2/yr)
        Hu =    param_values[i, 4] # potentiometric head of upper aquifer (m)
        Tl =    param_values[i, 5] # transmissivity of lower aquifer (m2/yr)
        Hl =    param_values[i, 6] # potentiometric head of lower aquifer (m)
        L =     param_values[i, 7] # length of borehole (m)
        Kw =    param_values[i, 8] # hydraulic conductivity of borehole (m/yr)

        # calculation
        l = log(r/rw)
        num = 2 * pi * Tu * (Hu - Hl)
        denom = l * (1 + ((2 * L * Tu) / (l * rw^2 * Kw)) + (Tu/Tl))
        Y[i] = num / denom
    end
    return Y
end


using CSV
using LsqFit
using Dates

function calc(dates::Array{Date}, data::Array{Float64}, n::Int64)

    function j_m_exp(t,p)
        J = Array{Float64}(undef, length(t), length(p))
        J[:,1] = exp.(p[2] .* t)       #df/dp[1]
        J[:,2] = t .* p[1] .* J[:,1]   #df/dp[2]
        J
    end

    function j_m_pow2(t,p)
        J = Array{Float64}(undef, length(t), length(p))
        J[:,1] = 2.0 .^ (t ./ p[2])       #df/dp[1]

        c = -log(2.0)*p[1]/(p[2] * p[2])
        J[:,2] =  t .* c .* J[:,1]        #df/dp[2]
        J
    end

    datesWindow = dates[end-n+1 : end]
    date1 = dates[end-n+1]

    x = map(d -> (d-date1).value, datesWindow)
    y = data[end-n+1 : end] .- data[end-n+1];

    m_exp(t, p) = p[1] * exp.(p[2] * t)
    p0_exp = [x[1], 0.5]
    fit_exp = curve_fit(m_exp, j_m_exp, x, y, p0_exp)

    m_pow2(t, p) = p[1] * 2.0 .^ (t/p[2])
    p0_pow2 = [x[1], 0.5]
    fit_pow2 = curve_fit(m_pow2, j_m_pow2, x, y, p0_pow2)

    confidence_intervals = confidence_interval(fit_pow2, 0.1)

    p = fit_pow2.param

    # y_fitted = m_pow2(x,p)
    # diff = 100.0*(abs.(y.-y_fitted))./y
    # println(diff)

    print(p)
    print("\t")
    println(confidence_intervals)
    println("---")

end

df = CSV.File("CoronaData2.csv")

dates = df.Date

dataDE = convert(Array{Float64}, df.Germany)
dataCH = convert(Array{Float64}, df.Switzerland)

print("Germany: ")
# calc(dates[1:end], dataDE, 15)
calc(dates[1:end], dataDE, 10)
# calc(dates[1:end], dataDE, 5)

print("Switzerland: ")
# calc(dates[1:end], dataCH, 15)
calc(dates[1:end], dataCH, 10)
# calc(dates[1:end], dataCH, 5)


function RTBPEnergy(u, mu)
"""
     RTBP-Nbody problem
"""

 @inbounds begin
    x = u[1]   # x
    y = u[2]   # y
    px = u[3]  # px
    py = u[4]  # py

    umu=1-mu

    r1 = ((x+mu)^2 + y^2)^(1/2)
    r2 = ((x-umu)^2 + y^2)^(1/2)

    Energy=(px*px+py*py)/2+px*y-py*x-umu/r1-mu/r2-mu*umu/2

    return(Energy)
    end

end

f = (du,u,p,t) -> begin

  μ=p[1]
  uμ=1-μ
  @inbounds begin
  # 1 = y₁
  # 2 = y₂
  # 3 = y₁'
  # 4 = y₂'
  D₁ = ((u[1]+μ)^2 + u[2]^2)^(3/2)
  D₂ = ((u[1]-uμ)^2 + u[2]^2)^(3/2)
  du[1] = u[3]+u[2]
  du[2] = u[4]-u[1]
  du[3] = u[4] - uμ*(u[1]+μ)/D₁ - μ*(u[1]-uμ)/D₂
  du[4] = -u[3] - uμ*u[2]/D₁ - μ*u[2]/D₂
  end
end


#=
The power set of a set is the set of all its subsets. Write a function that, given a set, generates its power set.

For example, given the set {1, 2, 3}, it should return {{}, {1}, {2}, {3}, {1, 2}, {1, 3}, {2, 3}, {1, 2, 3}}.

You may also use a list or array to represent a set.
=#

using Test
include("Solutions/problem37_generate_powerset.jl")

@test get_powerset([1,2,3]) == [[], [1], [2], [3], [1, 2], [1, 3], [2, 3], [1, 2, 3]]


# WARNING:
# ========
# This blog post https://julialang.org/blog/2019/01/fluxdiffeq/ gives a high level
# description of this file.

# Many function calls are fully qualified. This is not necessary. The package name is
# included to get a better sense of the role of the different packages.

# The Lotka-Volterra differential equations model populations of predators
# and preys over time.

# This example starts with an ODE setup with a given set of parameters.
# Then, using the same equations, optimises starting with a different set of parameters

## Import the required packages
using DifferentialEquations, DiffEqSensitivity
using Flux, DiffEqFlux
using Plots

##
## Setup ODE to optimize
##

# The populations are representated as a vector of 2 numbers:
#    U = [numbers of rabbits, number of wolves]
#
# ODE problems expect the equations to be presented with standard parameters:
#   dU: the derivatives of U. This vector is modified in place for speed purposes
#    U: the functions
#    p: a vector of paramters
#    t: the functions' variable (here time)
function lotka_volterra!(du, u, p, t)
  x, y = u
  α, β, δ, γ = p
  du[1] = dx = α * x - β * x * y
  du[2] = dy = -δ * y + γ * x * y
end

# Initial conditions
u0 = [1.0, 1.0]

# Time over which to simulate
t_span = (0.0, 10.0)

# Initial parameters
p0 = [1.5, 1.0, 3.0, 1.0]

# Setup the ODE
prob0 = DiffEqBase.ODEProblem(lotka_volterra!, u0, t_span, p0)

# Solve the ODE
sol0 = DiffEqBase.solve(prob0, Tsit5())

# Plot it
plot(sol0)

##
## Generate data that will be used to train the network
## The data is the number of rabbits/wolves at set times
##

# Solve the ODE (again) and collect solutions at fixed intervals
target_data = DiffEqBase.solve(prob0, Tsit5(), saveat = 0.1)
rabbits = target_data[1, :]                     # vector of 101 data points
wolves  = target_data[2, :]                    # vector of 101 data points

# Plot the data on top of the full solution
t_steps = 0:0.1:10.0
scatter!(t_steps, rabbits, color = [1], label = "rabbits")
scatter!(t_steps, wolves,  color = [2], label = "wolves")

##
## Parameter optimisation
##

# Loss function is the total squared error
# (the number of points is constant - mean is not necessary)
loss_function = function()
    prediction = DiffEqBase.concrete_solve(prob0, Tsit5(), u0, p;
    sensealg = TrackerAdjoint(), saveat = 0.0:0.1:10.0)

    # Calculate squared error
  return sum(abs2, prediction - target_data)
end

# Callback function to observe training
iter = 0
callback = function ()
  global iter += 1
  if iter % 10 == 1
    # Plot the training data
    scatter(t_steps,  rabbits, color = [1], label = "rabbit data")
    scatter!(t_steps, wolves,  color = [2], label = "wolves data")

    # Use `remake` to re-create the original `prob0` with different current parameters `p`
    remade_sol = DiffEqBase.solve(remake(prob0, p = p), Tsit5(), saveat = 0.1)
    remade_rabbits = remade_sol[1, :]
    remade_wolves  = remade_sol[2, :]
    plot!(t_steps, remade_rabbits, ylim = (0, 6), labels = "rabbit model", color = 1)
    display(plot!(t_steps, remade_wolves, ylim = (0, 6), labels = "wolf model", color = 2))
  end
end

# Vector of new parameters different from p0
p = [4.0, 1.0, 2.0, 0.4]
ODEparams = Flux.params(p)

# Display the ODE with the initial parameter values.
callback()

# The parameter `loss_data` in `train!` (below) contains a list of data that are provided
# to the loss function. `train!`` works by iterating over that list. In our case, the
# calculation of the loss function does not require any additional data.
# Therefore `loss_data` is filled with a whole bunch of nothing.
loss_data = Iterators.repeated((), 1000)

# ADAM is a great default optimiser
optimiser = Flux.ADAM(0.1)

# train! is a function that hides some of the complexity of the library
Flux.train!(loss_function, ODEparams, loss_data, optimiser; cb = callback)


### A Pluto.jl notebook ###
# v0.12.21

using Markdown
using InteractiveUtils

# ╔═╡ 166bbbb2-8f7a-11eb-19c6-19dfe8ee5123
using ChainRules

# ╔═╡ a8439cf0-8eab-11eb-3d6b-a3d65e46a635
md"""
# Automatic Differentiation (AD)

Automatic differentiation is a method to evaluate the derivatives of a given program, you may heard this in context of deep learning as back-propagation. However, in more general scientific machine learning, the situation can be more complicated, thus it'd be necessary to understand better about the mechanism of automatic differentiation.

In this section, we will introduce the automatic differentiation technique in a more systematic aspect
"""

# ╔═╡ 12cb96f8-8f48-11eb-212d-956d24e7a892
md"""
# Dual Numbers

Dual number is the simplest way of implementing AD. And thanks to Julia's powerful generic programming features, we can easily implement one ourselves!

**definition** define ``ϵ`` to satisfy ``ϵ^2 = 0``, then similar to a complex number we define dual number as ``x = a + b ⋅ ϵ``

then we can intermediately get the rule of multiplication

```math
(a + b⋅ϵ)(c + d⋅ϵ) = ac + (ad + bc)ϵ
```

thus we can implement our own dual number as following
"""

# ╔═╡ af43976a-8f48-11eb-1485-0b76c7ee1e79
begin
	struct Dual{T<:Real} <: Real
		a::T
		b::T
	end

	Dual(a) = Dual(a, one(a))
end

# ╔═╡ d1522092-8f48-11eb-30f2-f9c3183ae364
Base.show(io::IO, x::Dual) = print(io, x.a, "+", x.b, "ϵ")

# ╔═╡ e0a6fe78-8f48-11eb-152f-5b448f08fc49
Dual(2)

# ╔═╡ e4fbb1f8-8f48-11eb-220e-f33784f6e617
Base.:(*)(x::Dual, y::Dual) = Dual(x.a * y.a, x.a * y.b + x.b * y.a)

# ╔═╡ 25669170-8f4b-11eb-1cae-a78fc79be709
Base.:(+)(x::Dual, y::Dual) = Dual(x.a + y.a, x.b + y.b)

# ╔═╡ 0c7cdab8-8f49-11eb-3bac-2de9b80bb7d0
Dual(3, 1) * Dual(3, 1)

# ╔═╡ 82b578de-8f49-11eb-13dd-aba04a2c196d
md"now let's define a simply function see what happens if we instead of feeding a normal real number, we feed in a dual number"

# ╔═╡ 28b6ea34-8f49-11eb-2625-d58fbfba667f
f(x) = x^3

# ╔═╡ 9f7849bc-8f49-11eb-0f59-1bdf987ecdfe
f(Dual(2))

# ╔═╡ a7d73bb6-8f49-11eb-2c45-99fbd3a2c248
md"""
the prefactor of ϵ gives the derivative! let's see why

given a real polynomail

```math
P(x) = p_0 + p_1 x + p_2 x^2 + ⋯ + p_n x^n
```

if we put a dual number ``x = a + b ϵ``, we have

```math
\begin{aligned}
P(a+bϵ)
&= p_0 + p_1 a + p_2 a^2 + ⋯ + p_n a^n + p_1 bϵ + 2p_2 abϵ + ⋯ + n p_n a^{n-1}b ϵ\\
&= P(a) + b P'(a)ϵ
\end{aligned}
```

more generally, using Taylor series, this is also true for any real analytical functions.

However, if we plugin our dual number into functions like `sin`, it doesn't work
"""

# ╔═╡ 8a4b5ec8-8f4a-11eb-1155-f358a173daa2
sin(Dual(2))

# ╔═╡ 94bc3b16-8f4a-11eb-0871-59f5469f4c0e
md"""
this is because in practical numerical program, for efficiency, a lot mathematical functions are implemented specifically for floating point type, thus we don't have a generic function for these functions to make things "just work"

thus we will need to define the derivative of `sin` manually here
"""

# ╔═╡ cfdb25b6-8f4a-11eb-2248-3f7783956f35
Base.sin(x::Dual) = Dual(sin(x.a), cos(x.a) * x.b)

# ╔═╡ de6df808-8f4a-11eb-0f70-43400519fd6f
sin(Dual(2))

# ╔═╡ f2b1a960-8f4a-11eb-1276-b333b086605c
md"however if we happen to have a function definition on generic types, e.g the matrix multiplication, this will just work"

# ╔═╡ 0906c78e-8f4b-11eb-1d3b-dd8d6d891175
A = [Dual(1.0) Dual(2.0); Dual(1.0) Dual(2.0)]

# ╔═╡ 149734d0-8f4b-11eb-0af4-6368d8bb34c0
B = [Dual(1.0) Dual(2.0); Dual(1.0) Dual(3.0)]

# ╔═╡ 1a483212-8f4b-11eb-0fa5-5994548ea6b9
A * B

# ╔═╡ f4228cb2-8eab-11eb-1b98-fbb30f8eeac3
md"""
# Forward Mode AD

The forward mode AD is the simplest way of implementing an AD system. The above dual number is one way to implement forward mode AD with a generic type system.

More formally, we say forward mode AD is to apply the chain rule on primal derivatives
when we evaluate the value of an expression, e.g if we are going to evaluate the following function
"""

# ╔═╡ c6d79168-8f4f-11eb-1c4f-4bf4abbb2c72
function foo(x)
	x1 = sin(x)
	x2 = sin(x1)
	x3 = cos(x2) + x2
end

# ╔═╡ d2f7c68e-8f4f-11eb-0d90-4f2f22bd87b6
md"""
when we put in a dual number we actually do the following things:

1. evaluate `sin(x)` and its derivative, store them in `x1`
2. evaluate `sin(x1)` and its derivative, multiply the derivative of `sin(x1)` to previous derivative since we have ``P'(x)b ϵ`` term in the dual number
3. evaluate the derivative of `cos(x2)` and do the same thing as above, but let's call the intermediate value of `cos(x2)` to be `y`
4. evaluate the value of `y+x2` and derivative then multiply and store them in `x3`

thus we see every time we evaluate the derivative, we can throw away the variables we calculated before, but for every single scalar number we need to calculate the entire function, as a result, this gives us the time complexity $O(mn)$ where $n$ is the number of parameters and $m$ is the complexity of the original code.

the memory complexity is $O(m+n)$ where $m$ is the memory complexity of the original code
"""

# ╔═╡ f75c7e22-8f52-11eb-0b32-f9ed6267fb18
md"""
to formalize the above process in terms of chain rule, the forward mode means we evaluate the derivatives in the following accumulation:

```math
\begin{aligned}
\frac{\partial y_n}{\partial x}
&= \frac{\partial y_n}{\partial y_{n-1}} \frac{\partial y_{n-1}}{\partial x}\\
&= \frac{\partial y_n}{\partial y_{n-1}} (\frac{\partial y_{n-1}}{\partial y_{n-2}} \frac{\partial y_{n-2}}{\partial x})\\
&= \cdots\\
&= \frac{\partial y_n}{\partial y_{n-1}} (\frac{\partial y_{n-1}}{\partial y_{n-2}} \cdots (\frac{\partial y_2}{\partial y_1} \frac{\partial y_1}{\partial x}))
\end{aligned}
```

"""

# ╔═╡ 9ef95826-8f53-11eb-11ed-a35e16d20083
md"""
we can further generalize this to multi-variable case using Jacobians

```math
\begin{aligned}
\frac{\partial \vec{y}_n}{\partial x}
&= \frac{\partial \vec{y}_n}{\partial \vec{y}_{n-1}} \frac{\partial \vec{y}_{n-1}}{\partial x}\\
&= \frac{\partial \vec{y}_n}{\partial \vec{y}_{n-1}} (\frac{\partial \vec{y}_{n-1}}{\partial \vec{y}_{n-2}} \frac{\partial \vec{y}_{n-2}}{\partial x})\\
&= \cdots\\
&= \frac{\partial \vec{y}_n}{\partial \vec{y}_{n-1}} (\frac{\partial \vec{y}_{n-1}}{\partial \vec{y}_{n-2}} \cdots (\frac{\partial \vec{y}_2}{\partial \vec{y}_1} \frac{\partial \vec{y}_1}{\partial x}))
\end{aligned}
```
"""

# ╔═╡ 0e07a38e-8f55-11eb-0c46-ef5d11df07de
md"""
now we see to calculate the derivative of ``\frac{\partial \vec{y}_n}{\partial x}`` it is actually about calculatnig a chain of **Jacobian-vector product**
"""

# ╔═╡ 13c539e8-8eac-11eb-35b0-61108d5f3c5f
md"""
# Reverse Mode AD

For large number of parameters, reverse mode is more efficient. This method has been re-discovered many times in history in different fields, you may hear it called back-propagation in context of deep learning.
"""

# ╔═╡ 45826916-8f55-11eb-3f3c-95d6b601191b
md"""
similar to forward mode AD, the reverse mode AD is actually about calculating the chain rule in the following accumulation:

```math
\begin{aligned}
\frac{\partial \vec{y}_n}{\partial x}
&= (\frac{\partial \vec{y}_n}{\partial \vec{y}_{n-1}}) \frac{\partial \vec{y}_{n-1}}{\partial x}\\
&= ((\frac{\partial \vec{y}_n}{\partial \vec{y}_{n-1}}) \frac{\partial \vec{y}_{n-1}}{\partial \vec{y}_{n-2}}) \frac{\partial \vec{y}_{n-2}}{\partial x}\\
&= ⋯\\
&= (\cdots ((\frac{\partial \vec{y}_n}{\partial \vec{y}_{n-1}}) \frac{\partial \vec{y}_{n-1}}{\partial \vec{y}_{n-2}}\cdots \frac{\partial \vec{y}_2}{\partial \vec{y}_1}) \frac{\partial \vec{y}_1}{\partial x}\\
(\frac{\partial \vec{y}_n}{\partial x})^T &= (\frac{\partial \vec{y}_1}{\partial x})^T (\frac{\partial \vec{y}_2}{\partial \vec{y}_1})^{T} \cdots (\frac{\partial \vec{y}_{n-1}}{\partial \vec{y}_{n-2}})^T (\frac{\partial \vec{y}_n}{\partial \vec{y}_{n-1}})^T
\end{aligned}
```

thus the reverse mode automatic differentiation is acutally about **Jacobian-transpose-vector product**.

We can see an intermediate advantage of this accumulation is that if ``x`` is a vector, we can directly calculate the derivative of this vector by a chain of matrix multiplication.
"""

# ╔═╡ e026ffa6-8f5f-11eb-0651-dfb9e28fcc4f
md"""
this process will be easier to understand if we use a graphical language to describe it, this the **computational graph**, to demonstrate this better, I'll use a more complicated function

```math
y = \mathbf{x}^T\mathbf{A}\mathbf{x} + \mathbf{b}\cdot\mathbf{x} + c
```
"""

# ╔═╡ 10ba0d9c-8f5f-11eb-0744-6be46cfacebe
md"""
We will need to call several functions in Julia to get the result $y$, which is

1. ``\mathbf{z_1} = \mathbf{x}^T``: `transpose` function.
2. ``\mathbf{z_2} = \mathbf{z_1} A`` matrix-vector multiplication, which can be `gemv` in `LinearAlgebra.BLAS`, or just `*`.
3. ``y_1 = \mathbf{z_2} \mathbf{x}`` vector dot operation, which is `LinearAlgebra.dot` or the UTF-8 operator `x ⋅ y`
4. ``y_2 = \mathbf{b} \cdot \mathbf{x}`` another vector dot
5. ``y_1 + y_2 + c`` a scalar add function, one can calculate it by simply calling `+` operator in Julia.

In fact, we can draw a graph of this expression, which illustrates the relationship between each variable in this expression.
Each node in the graph with an output arrow represents a variable and each node with an input arrow represents a function/operator.

![](https://blog.rogerluo.me/images/comput-graph-forward.gif)

The evaluation of the math equation above can then be expressed as a process called **forward evaluation**, it starts from the leaf nodes, which represents the inputs of the whole expression, e.g they are $\mathbf{x}, \mathbf{A}, \mathbf{b}, c$ in our expression. Each time, we receive the value of a node in the graph, we mark the node with **green**.
"""

# ╔═╡ 4921f4f8-8f60-11eb-0e45-3f3034dbea83
md"""
the derivative calculation can be then visualized as following

![](https://blog.rogerluo.me/images/comput-graph-backward.gif)
"""

# ╔═╡ 29fa38b6-8f61-11eb-118a-616fc647715a
md"""
as we demonstrate above, although the reverse mode can calculate the derivatives for all the parameters in single evaluation. it will need to store the intermediate values.

thus we need a specific data structure to store these values， this data structure is often called the computational graph, tape or Wengert list.

Now the problem of implementing reverse mode AD becomes how to create the tape.
"""

# ╔═╡ 0580566c-8ff4-11eb-35a3-bf0505fe5a1c
md"""
# Relation with Differential Geometry

Since we have represented our AD algorithm as Jacobian or Jacobian transpose product, it is natural to have a geometric interpretation of them.

## Definitions

Let ``p`` be a point in the space ``\mathbf{M}`` defined by ``(x_1, x_2, ⋯, x_m)``, and we have a function ``f: \mathbf{M} → \mathbf{K}`` that is a smooth map (differentiable everywhere), then we can define the **tangent space** ``\mathbf{T}_p \mathbf{M}`` at point ``p`` as the space span by the derivatives ``\frac{d}{dx}``, the tangent space acts linearly on the space of functions.

Now we define **the collection of the tangent spaces** ``\mathbf{T}_p\mathbf{M}`` for ``p ∈ \mathbf{M}`` as the **tangent bundle** of ``\mathbf{M}``

![tangent space](https://upload.wikimedia.org/wikipedia/commons/thumb/e/e7/Tangentialvektor.svg/200px-Tangentialvektor.svg.png)

let ``df`` denotes the first order information of ``f`` at each point, which is called the **differential** of ``f``. If the derivative of ``f`` and ``g`` agree at ``p``, we say that ``df`` and ``dg`` represents the same cotangent at ``p``. The covectors ``dx_1, dx_2, ⋯, dx_m`` form the basis of the cotangent space ``\mathbf{T}^*_p \mathbf{M}`` for ``p∈\mathbf{M}``

**push-forward**

![](https://upload.wikimedia.org/wikipedia/commons/thumb/3/37/Pushforward.svg/330px-Pushforward.svg.png)

**pullback**

![](https://www.mathphysicsbook.com/wp-content/uploads/2013/10/39.pullback-v2.png)
"""

# ╔═╡ eb4df5e4-8f79-11eb-2429-a31bd0009c4b
md"""
## the manual AD

so we can do reverse mode "automatic" differentiation manually first, to have a feeling about it.

We will use the package `ChainRules` here.
"""

# ╔═╡ 72b98a42-8ff8-11eb-17f3-6d14da3964be
x = 3.0

# ╔═╡ 1a78b2be-8f7a-11eb-0470-e7febaac6272
y, sin_pullback = rrule(sin, x)

# ╔═╡ 627f170a-8ff8-11eb-21eb-47796905e10b
sin_pullback(1)

# ╔═╡ 67c7f3d0-8ff8-11eb-04dc-4f796cfb940e
cos(x)

# ╔═╡ 7cffb62a-8ff8-11eb-0288-f55b0d333571
md"""
now we can manually generate the pullback function for a given function defined on some primal functions, in AD, we usually call these functions **adjoint**s.

let's use the previous example `foo` here
"""

# ╔═╡ 98d147c4-8ff8-11eb-3ab8-21e26217e02a
function adjoint_foo(x)
	x1, x1_pullback = rrule(sin, x)
	x2, x2_pullback = rrule(sin, x1)
	y, y_pullback = rrule(cos, x2)
	x3, x3_pullback = rrule(+, y, x2)

	return x3, function pullback(Δ)
		_, ∂y, ∂x2_1 = x3_pullback(Δ)
		_, ∂x2_2 = y_pullback(∂y)
		_, ∂x1 = x2_pullback(∂x2_1 + ∂x2_2)
		_, ∂x = x1_pullback(∂x1)
		return Zero(), ∂x
	end
end

# ╔═╡ 8e593406-8ff9-11eb-101a-695b2b026a20
_, foo_pullback = adjoint_foo(x)

# ╔═╡ 9e31a492-8ff9-11eb-108d-1b7bbd152170
foo_pullback(1)

# ╔═╡ a21aa0c2-8ff9-11eb-1cc8-7b63a20d0791
(foo(x+1e-10)-foo(x-1e-10))/2e-10

# ╔═╡ c5f14d48-8ff9-11eb-1350-790162aa8b34
md"""
thus the simplest reverse mode AD is about how to create this tiny `pullback` function automatically, and to know how to create this `pullback` function we need to know what primal functions are called by the given function, so that we can simply reverse the order of calls and replace them with the `pullback` functions.
"""

# ╔═╡ 3fc036f8-8eac-11eb-15dc-017c0d5237bc
md"""
## Operator Overloading AD

The easiest way of creating a tape is using operator overloading. By operator overloading, it means the exactly same way we implemented our toy symbolic program in the previous section.

for example, we can define our own type to dispatch the functions to a `track` function so that we can store the function call into a tape.

```julia
mutable struct Variable{T} <: ADExpr
    value::T
    grad::T

    Variable(val::T) where T = new{T}(val)
    Variable(val::T, grad::T) where T = new{T}(val)
end

struct Node{FT <: Function, ArgsT <: Tuple, KwargsT <: NamedTuple} <: ADExpr
    f::FT
    args::ArgsT
    kwargs::KwargsT
end
```

then overload some primal functions to track the call into tape

```
Base.sin(x::ADExpr) = register(Base.sin, x)
```
"""

# ╔═╡ 35cfef70-8ffa-11eb-1a48-e906596eebe6
md"""
I will not demonstrate the details of implementing such AD engine, but once you understand the mechanism, you should be able to write one yourself in mins!

For the curious, checkout my blog post **Implementing Your Own Automatic Differentiation Engine in ONE day**: [https://blog.rogerluo.me/2018/10/23/write-an-ad-in-one-day/](https://blog.rogerluo.me/2018/10/23/write-an-ad-in-one-day/)
"""

# ╔═╡ 73ffb3fe-8ffd-11eb-2b71-038543ea5093
md"""

## What packages/software uses this method?

#### Python

- PyTorch
- autograd
- ...

#### Julia

- Tracker
- YAAD
- Nabla
- ...

"""

# ╔═╡ bd69872c-8f78-11eb-2197-2f760c33b197
md"""

## Limitations

however, the above method **won't work for general programs that contains control flows**. This is because in the representation of Wengert list, or the tape, we don't have the semantic of control flows, and we don't have a way to record control flows using operator overloading, such as `for` loop or `if else` statements.

This is because control flows are not functions in most of the programming languages, we cannot change the control flow behaviour by overloading them on different types, and it may not make sense to do so.

On the other hand, it **won't work for generic programming**, since we have used a specific type to dispatch functions, which kinda "abused" the type system.
"""

# ╔═╡ 46ad0fa2-8eac-11eb-28df-4f2c23f8d3ff
md"""
## Source Code Transformation AD

As you may noticed, in both forward and reverse mode AD, in order to represent intermediate derivatives, we were forced to use a new temperory variable for every function call.

This leads to a special intermediate representation for general programs called the **Static Single Assignment** intermediate representation (SSA IR), by **Static Single Assignment** we mean:

> every variable in the program is only assigned once statically

the control flows in a SSA IR are usually canonicalized to `goto` or `goto ifnot` if it has conditions, since every variable is only assigned once we use `%<name>` or `%<int>` to represent a variable.
"""

# ╔═╡ cbe1949e-8ffc-11eb-0747-e576545133b4
md"""
the SSA IR is a widely used representation in your daily used compilers, for example, your C/C++ compiler clang, or some of your new fancy langauges like Julia, rust and more. This is because SSA IR is a very convenient representation to optimize and analysis.

![](https://slideplayer.com/slide/13900843/85/images/7/C-like+IR+In+SSA+form+x1+%3D+f%28x0%29%3B+x+%3D+f%28x%29%3B+if+%28x1+%3E+y0%29.jpg)
"""

# ╔═╡ 6616ee24-8ffd-11eb-0d13-7f7ea01d3f73
md"""
this is a natural extension of the Wengert list (the tape) with control flows. Thus instead of doing operator overloading, we can directly write a compiler extension to transform the source encoded in SSA IR to the corresponding differential function in SSA IR.

And in Julia, this is extremely easy (in terms of implementing it), because Julia is a dynamic compiled language

![](https://blog.rogerluo.me/images/julia-compile-diagram.png)
"""

# ╔═╡ df1d9e58-90de-11eb-341a-bb942e20e713
md"""
For the curious, checkout my blog post **Build your own source to source automatic differentiation in ONE day!** :[https://blog.rogerluo.me/2019/07/27/yassad/](https://blog.rogerluo.me/2019/07/27/yassad/)
"""

# ╔═╡ c5f18fe2-8ffe-11eb-0a96-1f3ce397e943
md"""
# Software/Package uses this method

- **Julia**: Zygote, Enzyme, Yota, Diffractor (work-in-progress), ...
- **Python**: JAX, ...
"""

# ╔═╡ 420d72f4-8ffe-11eb-1b0f-43bdc9edf955
md"""
## Limitations

No limitations in theory.

We have the information about control flows, which means it **will work for control flows**.

We don't need to reply on type system to track the function calls anymore which means it works for **generic programming**

But covering the entire general purpose programming language with most of the corner cases may take years to accomplish.
"""

# ╔═╡ 84e5fe6e-8eac-11eb-04f8-cf7499b2c218
md"""
## Other type of AD engines

There are other type of AD engines that works under some context, such as polynomials, reversible programs, etc.

You can find a full list of AD engines in Julia at [https://juliadiff.org/](https://juliadiff.org/)
"""

# ╔═╡ 2e3fc494-90de-11eb-1c90-f9cc1357aab3
md"""

### Is there an algebra gives us reverse mode AD?

- reversible programming

![](https://giggleliu.github.io/NiLang.jl/dev/asset/revcomp.jpg)

"""

# ╔═╡ 8d13dbe0-90de-11eb-1bcf-dfbd8118e59d
md"""
## NiLang: reversible domain-specific programming language in Julia

![](https://github.com/GiggleLiu/NiLang.jl/raw/master/docs/src/asset/logo3.png)
"""

# ╔═╡ 4c23d3bc-8eac-11eb-2870-a37dfc123bc5
md"""
## Differentiable Programming in Condensed Matter Physics: Tensor Network

One of the most poplular method in numerical condensed matter physics is the tensor network methods.

Tensor network is also way to approximate the solution, in previous lectures, you have learned how to use neural network states to solve ground state problem using variational Monte Carlo.

More traditionally physicists have been using tensor network as the physical ansatz. However calculating the gradient of tensor networks can be complicated. Thus, by making use of automatic differentiation, one will be able to do gradient based optimization on tensor network easily.

![](https://github.com/Roger-luo/PIML/blob/master/notebooks/assets/tensor-network.png?raw=true)

[https://arxiv.org/pdf/1903.09650.pdf](https://arxiv.org/pdf/1903.09650.pdf)
"""

# ╔═╡ 619d6b8e-90d4-11eb-17fa-dd41790aa420
md"""
## Further Readings

- the ChainRules package documentation: [https://juliadiff.org/ChainRulesCore.jl/stable/](https://juliadiff.org/ChainRulesCore.jl/stable/)
- NiLang tutorial for reversible programming: [https://giggleliu.github.io/NiLang.jl/dev/why/](https://giggleliu.github.io/NiLang.jl/dev/why/)
- the differential and pullback: [https://www.mathphysicsbook.com/mathematics/manifolds/mapping-manifolds/the-differential-and-pullback/](https://www.mathphysicsbook.com/mathematics/manifolds/mapping-manifolds/the-differential-and-pullback/)
"""

# ╔═╡ Cell order:
# ╟─a8439cf0-8eab-11eb-3d6b-a3d65e46a635
# ╟─12cb96f8-8f48-11eb-212d-956d24e7a892
# ╠═af43976a-8f48-11eb-1485-0b76c7ee1e79
# ╠═d1522092-8f48-11eb-30f2-f9c3183ae364
# ╠═e0a6fe78-8f48-11eb-152f-5b448f08fc49
# ╠═e4fbb1f8-8f48-11eb-220e-f33784f6e617
# ╠═25669170-8f4b-11eb-1cae-a78fc79be709
# ╠═0c7cdab8-8f49-11eb-3bac-2de9b80bb7d0
# ╟─82b578de-8f49-11eb-13dd-aba04a2c196d
# ╠═28b6ea34-8f49-11eb-2625-d58fbfba667f
# ╠═9f7849bc-8f49-11eb-0f59-1bdf987ecdfe
# ╟─a7d73bb6-8f49-11eb-2c45-99fbd3a2c248
# ╠═8a4b5ec8-8f4a-11eb-1155-f358a173daa2
# ╟─94bc3b16-8f4a-11eb-0871-59f5469f4c0e
# ╠═cfdb25b6-8f4a-11eb-2248-3f7783956f35
# ╠═de6df808-8f4a-11eb-0f70-43400519fd6f
# ╟─f2b1a960-8f4a-11eb-1276-b333b086605c
# ╠═0906c78e-8f4b-11eb-1d3b-dd8d6d891175
# ╠═149734d0-8f4b-11eb-0af4-6368d8bb34c0
# ╠═1a483212-8f4b-11eb-0fa5-5994548ea6b9
# ╟─f4228cb2-8eab-11eb-1b98-fbb30f8eeac3
# ╠═c6d79168-8f4f-11eb-1c4f-4bf4abbb2c72
# ╟─d2f7c68e-8f4f-11eb-0d90-4f2f22bd87b6
# ╟─f75c7e22-8f52-11eb-0b32-f9ed6267fb18
# ╟─9ef95826-8f53-11eb-11ed-a35e16d20083
# ╟─0e07a38e-8f55-11eb-0c46-ef5d11df07de
# ╟─13c539e8-8eac-11eb-35b0-61108d5f3c5f
# ╟─45826916-8f55-11eb-3f3c-95d6b601191b
# ╟─e026ffa6-8f5f-11eb-0651-dfb9e28fcc4f
# ╟─10ba0d9c-8f5f-11eb-0744-6be46cfacebe
# ╟─4921f4f8-8f60-11eb-0e45-3f3034dbea83
# ╟─29fa38b6-8f61-11eb-118a-616fc647715a
# ╟─0580566c-8ff4-11eb-35a3-bf0505fe5a1c
# ╟─eb4df5e4-8f79-11eb-2429-a31bd0009c4b
# ╠═166bbbb2-8f7a-11eb-19c6-19dfe8ee5123
# ╠═72b98a42-8ff8-11eb-17f3-6d14da3964be
# ╠═1a78b2be-8f7a-11eb-0470-e7febaac6272
# ╠═627f170a-8ff8-11eb-21eb-47796905e10b
# ╠═67c7f3d0-8ff8-11eb-04dc-4f796cfb940e
# ╟─7cffb62a-8ff8-11eb-0288-f55b0d333571
# ╠═98d147c4-8ff8-11eb-3ab8-21e26217e02a
# ╠═8e593406-8ff9-11eb-101a-695b2b026a20
# ╠═9e31a492-8ff9-11eb-108d-1b7bbd152170
# ╠═a21aa0c2-8ff9-11eb-1cc8-7b63a20d0791
# ╟─c5f14d48-8ff9-11eb-1350-790162aa8b34
# ╟─3fc036f8-8eac-11eb-15dc-017c0d5237bc
# ╟─35cfef70-8ffa-11eb-1a48-e906596eebe6
# ╟─73ffb3fe-8ffd-11eb-2b71-038543ea5093
# ╟─bd69872c-8f78-11eb-2197-2f760c33b197
# ╟─46ad0fa2-8eac-11eb-28df-4f2c23f8d3ff
# ╟─cbe1949e-8ffc-11eb-0747-e576545133b4
# ╟─6616ee24-8ffd-11eb-0d13-7f7ea01d3f73
# ╟─df1d9e58-90de-11eb-341a-bb942e20e713
# ╟─c5f18fe2-8ffe-11eb-0a96-1f3ce397e943
# ╟─420d72f4-8ffe-11eb-1b0f-43bdc9edf955
# ╟─84e5fe6e-8eac-11eb-04f8-cf7499b2c218
# ╟─2e3fc494-90de-11eb-1c90-f9cc1357aab3
# ╟─8d13dbe0-90de-11eb-1bcf-dfbd8118e59d
# ╟─4c23d3bc-8eac-11eb-2870-a37dfc123bc5
# ╟─619d6b8e-90d4-11eb-17fa-dd41790aa420


# In this tutorial, we are exploring the application of Ridge and Lasso
# regression to the Hitters R dataset.
#
# ## Getting started

using MLJ
import RDatasets: dataset
using PrettyPrinting
MLJ.color_off() # hide
import Distributions
const D = Distributions

@load LinearRegressor pkg=MLJLinearModels
@load RidgeRegressor pkg=MLJLinearModels
@load LassoRegressor pkg=MLJLinearModels

# We load the dataset using the `dataset` function, which takes the Package and
# dataset names as arguments.

hitters = dataset("ISLR", "Hitters")
@show size(hitters)
names(hitters) |> pprint

# Let's unpack the dataset with the `unpack` function.
# In this case, the target is `Salary` (`==(:Salary)`) and all other columns are features (`col->true`).

y, X = unpack(hitters, ==(:Salary), col->true);

# The target has missing values which we will just ignore.
# We extract the row indices corresponding to non-missing values of the target.
# Note the use of the element-wise operator `.`.
no_miss = .!ismissing.(y);

# We collect the non missing values of the target in an Array.
# And keep only the corresponding features values.
y = collect(skipmissing(y))
X = X[no_miss, :]

# Let's now split our dataset into a train and test sets.
train, test = partition(eachindex(y), 0.5, shuffle=true, rng=424);

# Let's have a look at the target.

using PyPlot
ioff() # hide

figure(figsize=(8,6))
plot(y, ls="none", marker="o")

xticks(fontsize=12); yticks(fontsize=12)
xlabel("Index", fontsize=14), ylabel("Salary", fontsize=14)

savefig(joinpath(@OUTPUT, "ISL-lab-6-g1.svg")) # hide

# \figalt{Salary}{ISL-lab-6-g1.svg}

# That looks quite skewed, let's have a look at a histogram:

figure(figsize=(8,6))
hist(y, bins=50, density=true)

xticks(fontsize=12); yticks(fontsize=12)
xlabel("Salary", fontsize=14); ylabel("Density", fontsize=14)

edfit = D.fit_mle(D.Exponential, y)
xx = range(minimum(y), 2500, length=100)
yy = pdf.(edfit, xx)
plot(xx, yy, lw=3, label="Exponential distribution fit")

legend(fontsize=12)

savefig(joinpath(@OUTPUT, "ISL-lab-6-g2.svg")) # hide

# \figalt{Distribution of salary}{ISL-lab-6-g2.svg}
#
# ### Data preparation
#
# Most features are currently encoded as integers but we will consider them as continuous.
# To coerce `int` features to `Float`, we nest the `autotype` function in the `coerce` function.
# The `autotype` function returns a dictionary containing scientific types, which is then passed to the `coerce` function.
# For more details on the use of `autotype`, see the [Scientific Types](https://alan-turing-institute.github.io/DataScienceTutorials.jl/data/scitype/index.html#autotype)

Xc = coerce(X, autotype(X, rules=(:discrete_to_continuous,)))
scitype(Xc)

# There're a few features that are categorical which we'll one-hot-encode.

# ## Ridge pipeline
# ### Baseline
#
# Let's first fit a simple pipeline with a standardizer, a one-hot-encoder and a basic linear regression:

model = @pipeline(Standardizer(),
                     OneHotEncoder(),
                     LinearRegressor())

pipe  = machine(model, Xc, y)
fit!(pipe, rows=train)
ŷ = predict(pipe, rows=test)
round(rms(ŷ, y[test])^2, sigdigits=4)

# Let's get a feel for how we're doing

figure(figsize=(8,6))

res = ŷ .- y[test]
stem(res)

xticks(fontsize=12); yticks(fontsize=12)
xlabel("Index", fontsize=14); ylabel("Residual (ŷ - y)", fontsize=14)

ylim([-1300, 1000])

savefig(joinpath(@OUTPUT, "ISL-lab-6-g3.svg")) # hide

# \figalt{Residuals}{ISL-lab-6-g3.svg}

figure(figsize=(8,6))
hist(res, bins=30, density=true, color="green")

xx = range(-1100, 1100, length=100)
ndfit = D.fit_mle(D.Normal, res)
lfit  = D.fit_mle(D.Laplace, res)

plot(xx, pdf.(ndfit, xx), lw=3, color="orange", label="Normal fit")
plot(xx, pdf.(lfit, xx), lw=3, color="magenta", label="Laplace fit")

legend(fontsize=12)

xticks(fontsize=12); yticks(fontsize=12)
xlabel("Residual (ŷ - y)", fontsize=14); ylabel("Density", fontsize=14)
xlim([-1100, 1100])

savefig(joinpath(@OUTPUT, "ISL-lab-6-g4.svg")) # hide

# \figalt{Distribution of residuals}{ISL-lab-6-g4.svg}

# ### Basic Ridge
#
# Let's now swap the linear regressor for a Ridge one without specifying the penalty (`1` by default):
# We modify the supervised model in the pipeline directly.

pipe.model.linear_regressor = RidgeRegressor()
fit!(pipe, rows=train)
ŷ = predict(pipe, rows=test)
round(rms(ŷ, y[test])^2, sigdigits=4)

# Ok that's a bit better but surely we can do better with an appropriate selection of the hyperparameter.

# ### Cross validating

# What penalty should you use? Let's do a simple CV to try to find out:

r  = range(model, :(linear_regressor.lambda), lower=1e-2, upper=100_000, scale=:log10)
tm = TunedModel(model=model, ranges=r, tuning=Grid(resolution=50),
                resampling=CV(nfolds=3, rng=4141), measure=rms)
mtm = machine(tm, Xc, y)
fit!(mtm, rows=train)

best_mdl = fitted_params(mtm).best_model
round(best_mdl.linear_regressor.lambda, sigdigits=4)

# right, and  with that we get:

ŷ = predict(mtm, rows=test)
round(rms(ŷ, y[test])^2, sigdigits=4)

# Let's see:

figure(figsize=(8,6))

res = ŷ .- y[test]
stem(res)

xticks(fontsize=12); yticks(fontsize=12)
xlabel("Index", fontsize=14);
ylabel("Residual (ŷ - y)", fontsize=14)
xlim(1, length(res))

ylim([-1300, 1000])

savefig(joinpath(@OUTPUT, "ISL-lab-6-g5.svg")) # hide

# \figalt{Ridge residuals}{ISL-lab-6-g5.svg}
#
# You can compare that with the residuals obtained earlier.

# ## Lasso pipeline
#
# Let's do the same as above but using a Lasso model and adjusting the range a bit:

mtm.model.model.linear_regressor = LassoRegressor()
mtm.model.range = range(model, :(linear_regressor.lambda), lower=500, upper=100_000, scale=:log10)
fit!(mtm, rows=train)

best_mdl = fitted_params(mtm).best_model
round(best_mdl.linear_regressor.lambda, sigdigits=4)

# Ok and let's see how that does:

ŷ = predict(mtm, rows=test)
round(rms(ŷ, y[test])^2, sigdigits=4)

# Pretty good! and the parameters are reasonably sparse as expected:

coefs, intercept = fitted_params(mtm.fitresult).linear_regressor
@show coefs
@show intercept

# with around 50% sparsity:

coef_vals = [c[2] for c in coefs]
sum(coef_vals .≈ 0) / length(coefs)

# Let's visualise this:

figure(figsize=(8,6))
stem(coef_vals)

## name of the features including one-hot-encoded ones
all_names = [:AtBat, :Hits, :HmRun, :Runs, :RBI, :Walks, :Years,
             :CAtBat, :CHits, :CHmRun, :CRuns, :CRBI, :CWalks,
             :League__A, :League__N, :Div_E, :Div_W,
             :PutOuts, :Assists, :Errors, :NewLeague_A, :NewLeague_N]

idxshow = collect(1:length(coef_vals))[abs.(coef_vals) .> 10]
xticks(idxshow .- 1, all_names[idxshow], rotation=45, fontsize=12)
yticks(fontsize=12)
ylabel("Amplitude", fontsize=14)

savefig(joinpath(@OUTPUT, "ISL-lab-6-g6.svg")) # hide

# \figalt{Lasso coefficients}{ISL-lab-6-g6.svg}

# ## Elastic net pipeline

@load ElasticNetRegressor pkg=MLJLinearModels

mtm.model.model.linear_regressor = ElasticNetRegressor()
mtm.model.range = [range(model, :(linear_regressor.lambda), lower=0.1, upper=100, scale=:log10),
                    range(model, :(linear_regressor.gamma),  lower=500, upper=10_000, scale=:log10)]
mtm.model.tuning = Grid(resolution=10)
fit!(mtm, rows=train)

best_mdl = fitted_params(mtm).best_model
@show round(best_mdl.linear_regressor.lambda, sigdigits=4)
@show round(best_mdl.linear_regressor.gamma, sigdigits=4)

# And it's not too bad in terms of accuracy either

ŷ = predict(mtm, rows=test)
round(rms(ŷ, y[test])^2, sigdigits=4)

# But the simple ridge regression seems to work best here.
PyPlot.close_figs() # hide


using Base.Test
using Einsum

## Test that vars in Main aren't overwritten by einsum
let
  i = -1
  y = randn(10)
  @einsum x[i] := y[i]
  @test i == -1
end

## Test that B is overwritten by := operator
let
  B = randn(10,10)
  A = randn(5,10)
  @einsum B[i,j] := A[i,j] # this should run without a problem
  @test size(B) == size(A)
end

## CP decomposition test case ##
let

  # preallocated test case
  A = zeros(5,6,7);
  B = zeros(5,6,7);
  X = randn(5,2);
  Y = randn(6,2);
  Z = randn(7,2);

  @einsum A[i,j,k] = X[i,r]*Y[j,r]*Z[k,r]
  @einsimd B[i,j,k] = X[i,r]*Y[j,r]*Z[k,r]

  for i = 1:5
    for j = 1:6
      for k = 1:7
        s = 0.0
        for r = 1:2
          s += X[i,r]*Y[j,r]*Z[k,r]
        end
        @test isapprox(A[i,j,k],s)
        @test isapprox(B[i,j,k],s)
      end
    end
  end

  # without preallocation
  @einsum A2[i,j,k] := X[i,r]*Y[j,r]*Z[k,r]
  @test isapprox(A,A2)

end

# Interesting test case, can throw an error that
# local vars are declared twice.
let
  A = zeros(5,6,7);
  X = randn(5,2);
  Y = randn(6,2);
  Z = randn(7,2);
  if true
    @einsum A[i,j,k] = X[i,r]*Y[j,r]*Z[k,r]
  else
    @einsum A[i,j,k] = X[i,r]*Y[j,r]*Z[k,r]
  end
end

# At one point this threw an error because the lhs
# had no indices/arguments
let
  x = randn(10)
  y = randn(10)
  @einsum k := x[i]*y[i]
  @test isapprox(k,dot(x,y))
end

# Elementwise multiplication (this should create nested loops with no
# no summation.)
let
  x = randn(10)
  y = randn(10)
  @einsum k[i] := x[i]*y[i]
  @test isapprox(k,x.*y)
end

# Transpose a block matrix
let
  z = Any[ rand(2,2) for i=1:2, j=1:2]
  @einsum t[i,j] := transpose(z[j,i])
  @test isapprox(z[1,1], t[1,1]')
  @test isapprox(z[2,2], t[2,2]')
  @test isapprox(z[1,2], t[2,1]')
  @test isapprox(z[2,1], t[1,2]')
end

# Mapping functions
let
  A = randn(10,10)
  @einsum B[i,j] := exp(A[i,j])
  @test isapprox(exp(A),B)
end

# Example from numpy
let
  A = reshape(collect(1:25),5,5)
  @einsum B[i] := A[i,i]
  @test all(B .== [1,7,13,19,25])
end

# TODO: consider adding support for this:
# @einsum A[i,j] = A[i,j] + 50

## Test in-place operations ##
let
  A = randn(5,6,7);
  B = randn(5,6,7);
  A1 = copy(A);
  B1 = copy(B);

  X = randn(5,2);
  Y = randn(6,2);
  Z = randn(7,2);

  @einsum A[i,j,k] += X[i,r]*Y[j,r]*Z[k,r]
  @einsimd B[i,j,k] += X[i,r]*Y[j,r]*Z[k,r]

  for i = 1:5
      for j = 1:6
          for k = 1:7
              s = 0.0
              for r = 1:2
                  s += X[i,r]*Y[j,r]*Z[k,r]
              end
              @test isapprox(A[i,j,k],A1[i,j,k]+s)
              @test isapprox(B[i,j,k],B1[i,j,k]+s)
          end
      end
  end

  x = randn(10)
  y = randn(10)
  k0 = randn()
  k = k0
  @einsum k += x[i]*y[i]
  @test isapprox(k,k0+dot(x,y))

  # test multiplication

  A1[:] = A[:]
  B1[:] = B[:]

  @einsum A[i,j,k] *= X[i,r]*Y[j,r]*Z[k,r]
  @einsimd B[i,j,k] *= X[i,r]*Y[j,r]*Z[k,r]

  for i = 1:5
      for j = 1:6
          for k = 1:7
              s = 0.0
              for r = 1:2
                  s += X[i,r]*Y[j,r]*Z[k,r]
              end
              @test isapprox(A[i,j,k],A1[i,j,k]*s)
              @test isapprox(B[i,j,k],B1[i,j,k]*s)
          end
      end
  end

  x = randn(10)
  y = randn(10)
  k0 = randn()
  k = k0
  @einsum k *= x[i]*y[i]
  @test isapprox(k,k0*dot(x,y))
end

# Test offsets
let
  X = randn(10)

  # without preallocation
  @einsum A[i] := X[i+5]
  @test size(A) == (5,)
  @test all(A .== X[6:end])

  # with preallocation
  B = zeros(10)
  @einsum B[i] = X[i+5]
  @test size(B) == (10,)
  @test all(B[1:5] .== X[6:end])
end

# Test symbolic offsets
let
  offset = 5
  X = randn(10)

  # without preallocation
  @einsum A[i] := X[i+:offset]
  @test size(A) == (5,)
  @test all(A .== X[6:end])

  # with preallocation
  B = zeros(10)
  @einsum B[i] = X[i+:offset]
  @test size(B) == (10,)
  @test all(B[1:5] .== X[6:end])
end

# Test adding/subtracting constants
let
  k = 5
  X = randn(10)

  # without preallocation
  @einsum A[i] := X[i] + k
  @einsum B[i] := X[i] - k
  @test isapprox(A,X+k)
  @test isapprox(B,X-k)

  # with preallocation
  C,D = zeros(10),zeros(10)
  @einsum C[i] = X[i] + k
  @einsum D[i] = X[i] - k
  @test isapprox(C,X+k)
  @test isapprox(D,X-k)
end

# Test multiplying/dividing constants
let
  k = 5
  X = randn(10)

  # without preallocation
  @einsum A[i] := X[i]*k
  @einsum B[i] := X[i]/k
  @test isapprox(A,X.*k)
  @test isapprox(B,X./k)

  # with preallocation
  C,D = zeros(10),zeros(10)
  @einsum C[i] = X[i]*k
  @einsum D[i] = X[i]/k
  @test isapprox(C,X.*k)
  @test isapprox(D,X./k)
end

# Test indexing with a constant
let
  A = randn(10,2)
  j = 2
  @einsum B[i] := A[i,:j]
  @test all(B .== A[:,j])
  @einsum C[i] := A[i,1]
  @test all(C .== A[:,1])

  D = zeros(10,3)
  @einsum D[i,1] = A[i,:j]
  @test isapprox(D[:,1],A[:,j])
  @einsum D[i,:j] = A[i,:j]
  @test isapprox(D[:,j],A[:,j])
end

# Better type inference on allocating arrays
let
  B1 = ones(Int,5)
  B2 = ones(Float32,5)
  B3 = ones(5)
  C = randn(5)
  @einsum A1[i,j] := B1[i]*C[j]
  @einsum A2[i,j] := B2[i]*C[j]
  @einsum A3[i,j] := B3[i]*C[j]

  @test eltype(A1) == Float64
  @test eltype(A2) == Float64
  @test eltype(A3) == Float64
  @test isapprox(A1,A3)
  @test isapprox(A2,A3)
end


# %% Demonstration of delay embeddings
using DrWatson
@quickactivate "NonlinearDynamicsTextbook"
include(srcdir("style.jl"))
using DynamicalSystems, PyPlot, Random

lo = Systems.lorenz([0, 10, 0.0])
tr = trajectory(lo, 1000; Ttr=10, reltol = 1e-12, abstol = 1e-12)
x, y, z = columns(tr)
w = x
τ = estimate_delay(w, "mi_min")
R = embed(w, 3, τ)

close("all")
N = 5000
fig = figure()
ax1 = subplot(131, projection = "3d")
ax1.plot3D(x[1:N], y[1:N], z[1:N], lw = 1.0, color = "C0")
ax1.set_title("original set \$A\$", pad = 0)
ax = ax1
for a in (ax.xaxis, ax.yaxis, ax.zaxis); a.set_ticklabels([]); end
# ax1.text(-20, -20, 0, "\$D_0(A)\$ = $(round(D_A;digits=3))", color = "C1")

ax2 = subplot(132)
wx = 2400:3200
ax2.plot(wx, w[wx])
midpoint = (wx[1] + wx[end])/2
tau = midpoint:(midpoint+τ)
ax2.plot(tau, fill(10, length(tau)), color = "C3")
ax2.text(midpoint, 7.5, "\$\\tau\$", color = "C3", size = 32)
ax2.set_title("measurement \$w=x\$")
# make fancy axis
ax2.spines["bottom"].set_position("center")
ax2.spines["right"].set_color("none")
ax2.spines["top"].set_color("none")
ax2.set_xticks([])
ax2.set_yticks([])
ax2.set_xlabel("\$t\$")
ax2.set_ylabel("\$w\$", rotation = 0)
ax2.xaxis.set_label_coords(1.0, 0.5)
ax2.yaxis.set_label_coords(-0.05, 0.9)

ax3 = subplot(133, projection = "3d")
ax3.plot(R[1:N, 1], R[1:N, 2], R[1:N, 3], color = "C1", lw = 1.0)
ax = ax3
for a in (ax.xaxis, ax.yaxis, ax.zaxis); a.set_ticklabels([]); end
ax3.set_title("reconstruction \$R\$", pad = 0)
# ax3.set_title("reconstruction \$R\$\n\$(\\gamma=$(γ),\\tau=$(τ))\$")
# ax3.text(-5,-25,-10, "\$D_0(R)\$ = $(round(D_R;digits=3))", color = "C2")
ax1.dist = 9
ax3.dist = 9

fig.tight_layout(pad = 0.3)
wsave(plotsdir("6", "delayembedding"), fig)