# 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


"""
    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


### 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())


"""
    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


# 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


# %% 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)


#!/usr/bin/env julia
#
# Jerzy Wroczyński (nr. indeksu 250075)
# Obliczenia Naukowe, Laboratorium
# Lista 2, Zadanie 1.
#

# assert ARGS
err_mess = "provide two arguments: [64, 32] [up, down, max_to_min, min_to_max]"
if length(ARGS) != 2
    println(err_mess)
    exit(1)
end

# name the arguments
fl_str = ARGS[1]
mode = ARGS[2]

# parse the arguments
if !in(fl_str, ["64", "32"]) || !in(mode, ["up", "down", "max_to_min", "min_to_max"])
    println(err_mess)
    exit(1)
end

fl = Float64

# apply the first argument
if fl_str== "32"
    global fl = Float32
end
# end assert

# x₄ and x₅ were altered
x = [2.718281828, -3.141592654, 1.414213562, 0.577215664, 0.301029995]
y = [1486.2497, 878366.9879, -22.37492, 4773714.647, 0.000185049]

function calc_scalar_product(x, y)
    S = fl(0)
    # (a), (b)
    if in(mode, ["up", "down"])
        # change the direction of the range depending on the provided program argument
        for i in (mode == "down" ? (5:-1:1) : (1:5))
            S = fl(S + fl(x[i] * y[i]))
        end
    # (c), (d)
    elseif in(mode, ["max_to_min", "min_to_max"])
        # save partial products
        mid_positive::Array{fl} = []
        mid_negative::Array{fl} = []
        # iterate over the vectors
        for i in (1:5)
            out = fl(x[i] * y[i])
            if out < fl(0)
                push!(mid_negative, out)
            else
                push!(mid_positive, out)
            end
        end

        # sort the partial products
        sort!(mid_negative, rev=(mode == "max_to_min"))
        sort!(mid_positive, rev=(mode == "max_to_min"))

        # add the partial products
        S_positive = fl(0)
        S_negative = fl(0)
        for i in (1:length(mid_negative))
            S_negative = fl(S_negative + mid_negative[i])
        end
        for i in (1:length(mid_positive))
            S_positive = fl(S_positive + mid_positive[i])
        end
        # finally, add the partial sums
        S = fl(S_positive + S_negative)
    end
    return S
end

# print out the results
println(calc_scalar_product(x,y))


using Optim

f(x) = -sin(x[1]^2/2 - x[2]^2/4 + 3) * cos(2x[1] + 1 - exp(x[2])) # Same as last section
res = optimize(f, [0, 0.5])

function df(x)
    # Same as last section
    a1 = x[1]^2/2 - x[2]^2/4 + 3
    a2 = 2x[1] + 1 - exp(x[2])
    b1 = cos(a1)*cos(a2)
    b2 = sin(a1)*sin(a2)
    return -[x[1]*b1 - 2b2, -x[2]/2*b1 + exp(x[2])*b2]
end
res = optimize(f, df, [0, 0.5]; inplace=false)

res = optimize(f, df, [0, 0.5], GradientDescent(); inplace=false)

res = optimize(f, [0, 0.5], GradientDescent(); autodiff=:forward)

res = optimize(f, [0, 0.5], Newton(); autodiff=:forward)

res = optimize(f, [0, 0.5], BFGS(); autodiff=:forward)


module XiCor

using Random

export xicor

"""
    xicor(X, Y[, break_ties_randomly=false[, rng=nothing]])

Computes the correlation ξ between X and Y.

Unlike most coefficients of correlation, ξ ranges from -0.5 to 1.

If there are duplicate values in X, then ties are
broken based on the order in which they are observed.
If the order of X is not random, then you should
set `break_ties_randomly` to `true` to avoid a biased
estimate. You can use the `rng` parameter to
deterministically break ties.

See _A new coefficient of correlation_ by Chatterjee.

[arXiv:1909.10140 [math.ST]](https://arxiv.org/abs/1909.10140)

# Examples
```julia-repl
julia> ξ = xicor(1:100, 1:100)
0.9702970297029703
julia> x = trunc.((1:100) ./ 10);  # create x with lots of duplicates
julia> y = rand(MersenneTwister(0), 100);
julia> ξ = xicor(x, y, true, MersenneTwister(0))
-0.01830183018301823
julia> ξ = xicor(x, y, true, MersenneTwister(42))
0.004500450045004545
```
"""
function xicor(X, Y, break_ties_randomly=false, rng=nothing)
    if break_ties_randomly
        if !isnothing(rng)
            index = randperm(rng, length(X))
        else
            index = randperm(length(X))
        end
        X = X[index]
        Y = Y[index]
    end
    n = length(X)
    Y = Y[sortperm(X)]  # how should offset arrays be handled?
    sorter = sortperm(Y)
    R = zeros(Int, n)  # R[i] is the number of j such that Y[j] ≤ Y[i]
    L = zeros(Int, n)  # L[i] is the number of j such that Y[j] ≥ Y[i]

    i = 1
    while i <= n
        curr = Y[sorter[i]]
        counter = 1
        i += 1  # look ahead for repeated Y values
        while i <= n && Y[sorter[i]] == curr
            counter += 1
            i += 1
        end
        i -= 1
        for j = i-counter+1:i  # fill in R and L values wherever Y == curr
            R[sorter[j]] = i
            L[sorter[j]] = (n-i) + counter
        end
        i += 1
    end

    1 - n * sum(abs.(diff(R))) / (2 * sum(L .* (n .- L)))
end

end # module


######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
@doc """

    kalmanFilter(data, ssmodel)

Description:
Apply Kalman filter to observed data.
Measurement Equation:
    y_{t} = H_{t} β_{t} + A z_{t} + e_{t} .
Transition Equation:
    β_{t} = μ + F β_{t-1} + v_{t};
    e_{t} ~ i.i.d.N(0,R);
    v_{t} ~ i.i.d.N(0,Q);
    z_{t} ~ i.i.d.N(0,Z);
    E(e_t v_s') = 0.

Inputs:
- data      = observed data
- H         = measurement eq. state coef. matrix
- A         = measurement eq. exogenous coef. matrix
- F         = state eq. companion matrix
- μ         = state eq. intercept term
- R         = covariance matrix on measurement disturbance
- Q         = covariance matrix on state disturbance
- Z         = covariance matrix on predetermined var vector
"""
function kalmanFilter(data_y, ssmodel)

    @unpack H, A, F, μ, R, Q, Z = ssmodel

    # Save number of observations
    num_obs = size(data_y)[1]

    # Empty filtered data matrices
    data_filtered_y = similar(data_y)
    data_filtered_β = zeros(num_obs, size(Q)[1])

    # Create empty lists for P_{t}, P_{t|t-1}
    Ptt = Any[]
    Pttlag = Any[]

    # Initialize β_pred and P_pred
    β_pred_laglag = inv(I - F) * μ
    P_pred_laglag = ones(size(Q)[1], size(Q)[1])

    for t = 1:num_obs
        # Save current obs of y
        y = data_y[t, :]

        # Prediction
        β_pred_lag = μ + F * β_pred_laglag
        P_pred_lag = F * P_pred_laglag * transpose(F) + Q
        y_pred_lag = H * β_pred_lag
        η_pred_lag = y - y_pred_lag
        f_pred_lag = H * P_pred_lag * transpose(H) + R

        # Save P_{t|t-1}
        push!(Pttlag, P_pred_lag)

        # Updating
        K = P_pred_lag * transpose(H) * pinv(f_pred_lag)
        β_pred = β_pred_lag + K * η_pred_lag
        P_pred = P_pred_lag - K * H * P_pred_lag

        # Save P_{t|t}
        push!(Ptt, P_pred)

        # Save data
        data_filtered_y[t, :] = y_pred_lag
        data_filtered_β[t, :] = β_pred

        # Lag the predictions
        β_pred_laglag = β_pred
        P_pred_laglag = P_pred
    end

    # Returned filtered series
    # for obs variable and state
    return data_filtered_y, data_filtered_β, Pttlag, Ptt
end;
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
@doc """

    kalmanSmoother(data, ssmodel)

Description:
Apply Kalman smoother to observed data.
Measurement Equation:
    y_{t} = H_{t} β_{t} + A z_{t} + e_{t}.
Transition Equation:
    β_{t} = μ + F β_{t-1} + v_{t};
    e_{t} ~ i.i.d.N(0,R);
    v_{t} ~ i.i.d.N(0,Q);
    z_{t} ~ i.i.d.N(0,Z);
    E(e_t v_s') = 0.

Inputs:
- data      = observed data
- H         = measurement eq. state coef. matrix
- A         = measurement eq. exogenous coef. matrix
- F         = state eq. companion matrix
- μ         = state eq. intercept term
- R         = covariance matrix on measurement disturbance
- Q         = covariance matrix on state disturbance
- Z         = covariance matrix on predetermined var vector
"""
function kalmanSmoother(data_y, ssmodel)

    @unpack H, A, F, μ, R, Q, Z = ssmodel

    # Save number of observations
    num_obs = size(data_y)[1]

    # Empty filtered data matrices
    data_smoothed_y = similar(data_y)
    data_smoothed_β = zeros(num_obs, size(Q)[1])

    # Create empty list for P_{t|T}
    PtT = Any[]

    # Run Kalman filter
    data_filtered_y, data_filtered_β, Pttlag, Ptt = kalmanFilter(data_y, ssmodel)

    # Initialize β_{t+1|T} (β_{T|T})
    βtflagT = data_filtered_β[end]
    data_smoothed_β[end] = βtflagT

    # Initialize P_{t+1|T} (P_{T|T})
    Ptflag_T = Ptt[end]
    push!(PtT, Ptflag_T)

    # Initialize y_{t|T} (y_{T|T})
    data_smoothed_y[end] = data_filtered_y[end]

    # Run Kalman smoother
    for i = 1:(num_obs-1)

        # Retrieve β_{t|t}
        βtt = data_filtered_β[end-i]

        # Compute β_{t|T} using β_{t+1|T}, β_{t|t}, P_{t|t}, and P_{t+1|t}
        βtT = βtt +
              Ptt[end-i] * transpose(F) * inv(Pttlag[end-i+1]) *
              (βtflagT - F * βtt - μ)

        # Store β_{t+1|T} in smoothed data
        data_smoothed_β[end-i] = βtT

        # Set β_{t|T} as new β_{t+1|T} for next iteration
        βtflagT = βtT

        # Compute P_{t|T} using P_{t|t}, P_{t+1|t}, and P_{t+1|T}
        Pt_T = Ptt[end-i] +
               Ptt[end-i] * transpose(F) * inv(Pttlag[end-i+1]) *
               (Ptflag_T - Pttlag[end-i+1]) *
               transpose(Ptt[end-i] * transpose(F) * inv(Pttlag[end-i+1]))

        # Store P_{t|T}
        push!(PtT, Pt_T)

        # Set P_{t|T} as new P_{t+1|T} for next iteration
        Ptflag_T = Pt_T

        # Generate y_{t|T} (smoothed obs.)
        ytT = H * βtT

        # Store smoothed obs.
        data_smoothed_y[end-i] = ytT
    end

    # Flip P_{t|T} list
    PtT = reverse(PtT)

    # Returned filtered series
    # for obs variable and state
    return data_smoothed_y, data_smoothed_β, PtT
end;
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
@doc """

    KNFactorSampler(data_y, ssmodel)

Description:
Draw a sample series of dynamic factor from conditional distribution in Ch 8, Kim & Nelson (1999).
Measurement Equation:
    y_{t} = H_{t} β_{t} + A z_{t} + e_{t}.
Transition Equation:
    β_{t} = μ + F β_{t-1} + v_{t};
    e_{t} ~ i.i.d.N(0,R);
    v_{t} ~ i.i.d.N(0,Q);
    z_{t} ~ i.i.d.N(0,Z);
    E(e_t v_s') = 0.

Inputs:
- data      = observed data
- H         = measurement eq. state coef. matrix
- A         = measurement eq. exogenous coef. matrix
- F         = state eq. companion matrix
- μ         = state eq. intercept term
- R         = covariance matrix on measurement disturbance
- Q         = covariance matrix on state disturbance
- Z         = covariance matrix on predetermined var vector
"""
function KNFactorSampler(data_y, ssmodel)

    @unpack H, A, F, μ, R, Q, Z = ssmodel

    # Run Kalman filter
    data_filtered_y, data_filtered_β, Pttlag, Ptt = kalmanFilter(data_y, ssmodel)

    # Create placeholders for factor distr.
    # mean vector and covariance matrix for all t
    β_t_mean = Any[]
    β_t_var = Any[]

    # Record number of time periods
    T = size(data_y)[1]

    # Create empty vector for factor realizations
    β_realized = similar(data_filtered_β)

    # Initialize β_realized
    push!(β_t_mean, data_filtered_β[T, :])
    push!(β_t_var, Ptt[T])
    β_realized[T, :] = sim_MvNormal_alt(β_t_mean[1], β_t_var[1])

    # Generate `β_t_mean` and `β_t_var`
    # for all time periods
    if isposdef(Q) == false
        ## IF Q IS SINGULAR

        # Determine number of rows
        # of state factor used
        num_use_rows = 0
        for i = 1:size(Q)[1]
            if Q[i, i] != 0
                num_use_rows += 1
            end
        end

        # Create modified F and Q matrices
        F_star = F[1:num_use_rows, :]
        Q_star = Q[1:num_use_rows, 1:num_use_rows]

        # Iteratively generate conditional draws
        # of state vector
        for j = 1:(T-1)

            # β_{t|t,β*_{t+1}}
            β_t_mean_temp = data_filtered_β[T-j, :]
            +Ptt[T-j] * transpose(F_star) * inv(F_star * Ptt[T-j] * transpose(F_star) + Q_star) * (β_realized[T+1-j, :][1:num_use_rows] - μ[1:num_use_rows] - F_star * data_filtered_β[T-j, :])
            push!(β_t_mean, β_t_mean_temp)

            # P_{t|t,β*_{t+1}}
            β_t_var_temp = Ptt[T-j] - Ptt[T-j] * transpose(F_star) * inv(F_star * Ptt[T-j] * transpose(F_star) + Q_star) * F_star * Ptt[T-j]
            push!(β_t_var, β_t_var_temp)

            # Draw new β_t
            β_realized[T-j, :] = sim_MvNormal_alt(β_t_mean[1+j], β_t_var[1+j])
        end
    else
        ## IF Q IS NOT SINGULAR (redundant, but potentially faster)
        for j = 1:(T-1)

            # β_{t|t,β_{t+1}}
            β_t_mean_temp = data_filtered_β[T-j, :]
            +Ptt[T-j] * transpose(F) * inv(F * Ptt[T-j] * transpose(F) + Q) * (β_realized[T+1-j, :] - μ - F * data_filtered_β[T-j, :])
            push!(β_t_mean, β_t_mean_temp)

            # P_{t|t,β_{t+1}}
            β_t_var_temp = Ptt[T-j] - Ptt[T-j] * transpose(F) * inv(F * Ptt[T-j] * transpose(F) + Q) * F * Ptt[T-j]
            push!(β_t_var, β_t_var_temp)

            # Draw new β_t
            β_realized[T-j, :] = sim_MvNormal_alt(β_t_mean[1+j], β_t_var[1+j])
        end
    end

    # Return sampled factor series
    # fot t = 1,...,T
    return β_realized
end;
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
@doc """

    linearRegressionSampler(Y,X)

Description:
Estimate β and σ^2 in Y = Xβ + e, e ~ N(0,σ^2 I_T).
Generate samples of β and σ^2.

Inputs:
- Y     = Dependent data matrix
- X     = Independent data matrix
"""
function linearRegressionSampler(Y, X)

    # Save number of obs
    T = size(X)[1]

    # Initialize σ2
    σ2 = 1.0

    ##################################
    ##################################
    # Generate new β

    ## Prior parameters in N(β0,Σ0)
    β0 = zeros(size(X)[2])
    Σ0 = Matrix(I, size(β0)[1], size(β0)[1]) .* 1000.0

    ## Posterior parameters in N(β1,Σ1)
    β1 = inv(inv(Σ0) + inv(σ2) * transpose(X) * X) * (inv(Σ0) * β0 + inv(σ2) * transpose(X) * Y)
    β1 = vec(β1)
    Σ1 = inv(inv(Σ0) + inv(σ2) * transpose(X) * X)
    Σ1 = Hermitian(Σ1)

    ## Generate new β
    β = sim_MvNormal(β1, Σ1)

    ##################################
    ##################################
    # Update σ2

    ## Prior parameters in IG(ν0/2, δ0/2)
    ν0 = 0.002
    δ0 = 0.002

    ## Posterior parameters in IG(ν1/2, δ1/2)
    ν1 = ν0 + T
    δ1 = δ0 + norm(Y - X * β)^2

    ## Generate new σ2
    σ2 = rand(InverseGamma(ν1 / 2, δ1 / 2))

    ##################################
    ##################################
    # Return parameters
    return β, σ2
end;
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
@doc """

    linearRegressionSamplerRestrictedVariance(Y, X, σ2)

Description:
Estimate β and σ^2 in Y = Xβ + e, e ~ N(0,σ^2 I_T),
where σ^2 is restricted to some chosen value.
Generate samples of β and σ^2.

Inputs:
- Y     = Dependent data matrix
- X     = Independent data matrix
- σ2    = Restricted error variance
"""
function linearRegressionSamplerRestrictedVariance(Y, X, σ2)

    # Save number of obs
    T = size(X)[1]

    ##################################
    ##################################
    # Generate new β

    ## Prior parameters in N(β0,Σ0)
    β0 = zeros(size(X)[2])
    Σ0 = Matrix(I, size(β0)[1], size(β0)[1]) .* 1000.0

    ## Posterior parameters in N(β1,Σ1)
    β1 = inv(inv(Σ0) + inv(σ2) * transpose(X) * X) * (inv(Σ0) * β0 + inv(σ2) * transpose(X) * Y)
    β1 = vec(β1)
    Σ1 = inv(inv(Σ0) + inv(σ2) * transpose(X) * X)
    Σ1 = Hermitian(Σ1)

    ## Generate new β
    β = sim_MvNormal(β1, Σ1)

    ##################################
    ##################################
    # Return parameters
    return β
end;
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
@doc """

    autocorrErrorLinearRegressionSampler(Y, X, error_lag_num)

Description:
Estimate β, σ^2, and ϕ in Y = Xβ + e, e = Eϕ + ν, ν_t ~ i.i.d.N(0,σ^2).
Generate samples of β, σ^2, and ϕ.
Procedure descried in Section 7.4.2 in Kim & Nelson.

Inputs:
- Y             = Dependent data matrix
- X             = Independent data matrix
- error_lag_num = Number of lags in the disturbance DGP
"""

function autocorrErrorLinearRegressionSampler(Y, X, ϕold, σ2old, error_lag_num)

    # Save number of obs
    T = size(X)[1]

    # Initialize σ2
    ϕ = ϕold
    σ2 = σ2old

    ##################################
    ##################################
    # Generate β

    ## Prior parameters in N(b0,A0)
    b0 = zeros(size(X)[2])
    A0 = Matrix(I, size(b0)[1], size(b0)[1]) .* 1000.0

    ## Generate X^⋆
    X_star = similar(X[(1+length(ϕ)):end, :])
    for i = 1:size(X)[2] # iterate over variables in X
        x_temp = X[:, i]
        for p = 1:length(ϕ) # iterate over lag params in ϕ
            x_temp[(1+length(ϕ)):end, :] = x_temp[(1+length(ϕ)):end, :] - ϕ[p] .* lag(x_temp, p)[(1+length(ϕ)):end, :]
        end
        x_temp = x_temp[(1+length(ϕ)):end, :]
        X_star[:, i] = x_temp
    end

    ## Generate Y^⋆
    Y_star = similar(Y[(1+length(ϕ)):end, :])
    y_temp = Y
    for p = 1:length(ϕ)
        y_temp = Y
        y_temp = y_temp[(1+length(ϕ)):end, :] - ϕ[p] .* lag(y_temp, p)[(1+length(ϕ)):end, :]
    end
    Y_star = y_temp

    ## Posterior parameters in N(b1,A1)
    b1 = inv(inv(A0) + inv(σ2) * transpose(X_star) * X_star) * (inv(A0) * b0 + inv(σ2) * transpose(X_star) * Y_star)
    b1 = vec(b1)
    A1 = inv(inv(A0) + inv(σ2) * transpose(X_star) * X_star)
    A1 = Hermitian(A1)

    ## Draw new β
    β = sim_MvNormal(b1, A1)

    ##################################
    ##################################
    # Generate ϕ

    ## Prior parameters in N(c0,B0)
    c0 = zeros(size(ϕ)[1])
    B0 = Matrix(I, size(c0)[1], size(c0)[1]) .* 1000.0

    ## Generate e^⋆
    e_star = Y - X * β

    ## Generate E^⋆
    E_star = zeros(T, length(ϕ))
    E_star = E_star[(1+length(ϕ)):end, :]
    for i = 1:(size(E_star)[2]) # iterate over variables in X
        e_temp = lag(e_star, i)
        E_star[:, i] = e_temp[(1+length(ϕ)):end]
    end
    e_star = e_star[(1+length(ϕ)):end]

    ## Posterior parameters in N(c1,B1)
    c1 = inv(inv(B0) + inv(σ2) * transpose(E_star) * E_star) * (inv(B0) * c0 + inv(σ2) * transpose(E_star) * e_star)
    c1 = vec(c1)
    B1 = inv(inv(B0) + inv(σ2) * transpose(E_star) * E_star)
    B1 = Hermitian(B1)

    ## Generate new ϕ
    ###=
    ind = 0
    accept = 0
    ϕ = similar(ϕold)
    while accept == 0

        ind += 1

        ## Draw ψ
        ϕ = sim_MvNormal(c1, B1)

        ## Check for stationarity
        coef = [-reverse(vec(ϕ), dims = 1); 1]                      # check stationarity
        root = roots(Polynomial(reverse(coef)))
        rootmod = abs.(root)
        accept = min(rootmod...) >= 1.01

        ## If while loop goes on for too long
        if ind > 100
            ϕ = ϕold
            coef = [-reverse(vec(ϕ), dims = 1); 1]                      # check stationarity
            root = roots(Polynomial(reverse(coef)))
            rootmod = abs.(root)
            accept = min(rootmod...) >= 1.01
        end
    end
    ##=#

    #ϕ = sim_MvNormal(c1, B1)

    ##################################
    ##################################
    # Generate σ2

    ## Prior parameters in IG(ν0/2, δ0/2)
    ν0 = 0.002
    δ0 = 0.002

    ## Posterior parameters in IG(ν1/2, δ1/2)
    ν1 = ν0 + T
    δ1 = δ0 .+ norm(Y_star - X_star * β)^2
    δ1 = δ1[1, 1]

    ## Generate new σ2
    σ2 = rand(InverseGamma(ν1 / 2, δ1 / 2))

    ##################################
    ##################################
    # Return parameters
    return β, σ2, ϕ
end;
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################
######################


include("FractalDimensions.jl")

#S_0 = [[0.0,0.0], [1.0,0.0], [0.5,1.0]]
S_0 = [[0.0,0.0], [1.0,0.0], [0.5,2.0]]
#S_0 = [[0.0,0.0], [1.0,0.0], [0.0,1.0]]

f1(x) = (S_0[1]+x)/2
f2(x) = (S_0[2]+x)/2
f3(x) = (S_0[3]+x)/2

# WARNING This will use up your memory and slow down your computer for k>10
function sierpinsky_endpoints(k)
        s::Vector{Vector{Float64}} = []
        s1 = S_0
        for i in 1:k
            s2 = setdiff([f1.(s1); f2.(s1); f3.(s1)], s1)
            append!(s, s1)
            s1 = s2
        end
    return s
end

s = sierpinsky_endpoints(10)

using Plots
x = [u[1] for u in s]
y = [u[2] for u in s]

plotlyjs()
scatter(x,y, markersize=1, markerstrokealpha=0.0, legend=false)
savefig("sierpinsky.eps")

d, N, ε = box_counting_dimension(s, 9, 0.0)

loglog_regression(N,ε)

log(3)/log(2)

plot(-log.(ε), log.(N), xlabel="-ln ε", ylabel="ln N", marker=:d, legend=false)
savefig("boxdim_sierpinsky.eps")

C, r = correlation_dimension(s)

loglog_regression(C,r)

plot(log.(r), log.(C), xlabel="ln r", ylabel="ln C", marker=:d, legend=false)
savefig("cordim_sierpinsky.eps")


function acqr_ksmoother(a,xtt,Ptt,xtt1,Ptt1)
	#
	# Kalman smoother for model
	#
	# x_{t+1} = ax_{t} + w_{t}
	# y_{t}   = cx_{t} + v_{t}
	#
	# cov(w_{t},v_{t}) = [q 0;0 r]
	#
	# javier.cara@upm.es, 2020-05
	#

	nt = length(xtt)

	# allocation
	xtN = zeros(nt+1)
	PtN = zeros(nt+1)
	Pt1tN = zeros(nt)

	# values for t = nt+1
	xtN[nt+1] = xtt1[nt+1]
	PtN[nt+1] = Ptt1[nt+1]

	# values for t=nt
	xtN[nt] = xtt[nt]
	PtN[nt] = Ptt[nt]
	Pt1tN[nt] = a*Ptt[nt]

	# smother
	for t = nt-1:-1:1
		# Kalman Smoother matrix J
		Jt = (Ptt[t]*a) / Ptt1[t+1]

		xtN[t] = xtt[t] + Jt*( xtN[t+1] - xtt1[t+1] )
		PtN[t] = Ptt[t] + Jt^2*( PtN[t+1] - Ptt1[t+1] )

		Pt1tN[t] = PtN[t+1]*Jt
	end

	return xtN,PtN,Pt1tN

end