using DiffEqBase
# DDE Example Problems
# examples with constant delays
export prob_dde_1delay, prob_dde_1delay_notinplace, prob_dde_1delay_scalar_notinplace,
       prob_dde_2delays, prob_dde_2delays_notinplace, prob_dde_2delays_scalar_notinplace,
       prob_dde_1delay_long, prob_dde_1delay_long_notinplace,
       prob_dde_1delay_long_scalar_notinplace, prob_dde_2delays_long,
       prob_dde_2delays_long_notinplace, prob_dde_2delays_long_scalar_notinplace,
       prob_dde_mackey, prob_dde_wheldon, prob_dde_qs,
# examples with vanishing time dependent delays
       prob_ddde_neves1, prob_dde_neves_thompson,
# examples with state dependent delays
       prob_dde_paul1, prob_dde_paul2, prob_dde_mahaffy1, prob_dde_mahaffy2,
# examples with vanishing state dependent delays
       prob_neves2, prob_dde_gatica

# DDE examples with analytical solution

## Single constant delay

### In-place function

function f_1delay(du, u, h, p, t)
    du[1] = - h(p, t - oneunit(t))[1] / oneunit(t)
end

function f_1delay_analytic(u₀, p, t)
    z = t/oneunit(t)

    if z < 0
        return zero(u₀)
    elseif z < 1
        return u₀
    elseif z ≤ 10
        if z < 2
            c = @evalpoly(z, 2, -1)
        elseif z < 3
            c = @evalpoly(z, 4, -3, 1//2)
        elseif z < 4
            c = @evalpoly(z, 17//2, -15//2, 2, -1//6)
        elseif z < 5
            c = @evalpoly(z, 115//6, -109//6, 6, -5//6, 1//24)
        elseif z < 6
            c = @evalpoly(z, 1085//24, -1061//24, 197//12, -35//12, 1//4, -1//120)
        elseif z < 7
            c = @evalpoly(z, 13201//120, -13081//120, 521//12, -107//12, 1, -7//120, 1//720)
        elseif z < 8
            c = @evalpoly(z, 39371//144, -39227//144, 27227//240, -3685//144, 487//144,
                          -21//80, 1//90, -1//5040)
        elseif z < 9
            c = @evalpoly(z, 1158379//1680, -1156699//1680, 212753//720, -51193//720,
                          1511//144, -701//720, 1//18, -1//560, 1//40320)
        else
            c = @evalpoly(z, 23615939//13440, -23602499//13440, 7761511//10080,
                          -279533//1440, 89269//2880, -1873//576, 323//1440, -11//1120,
                          1//4032, -1//362880)
        end

        return c .* u₀
    else
        error("This analytical solution is only valid on (-∞,10]")
    end
end

ff_1delay = DDEFunction(f_1delay,analytic=f_1delay_analytic)

build_prob_dde_1delay(u₀, ::T=u₀) where {T} =
    DDEProblem(ff_1delay, [u₀], (p, t)->[zero(u₀)], (zero(T), T(10)),
               constant_lags = [oneunit(T)])

"""
    prob_dde_1delay

Model problem of finding a solution ``u(t)`` in the time span ``t \\in [0,10]`` to the delay
differential equation

```math
\\frac{du}{dt} = -u(t-1)
```

with history function

```math
u(t) = \\begin{cases}
0 & \\text{if } t < 0, \\
u₀ & \\text{if } t = 0,
\\end{cases}
```

for ``t \\leq 0`` with ``u₀ = 1``. Note that the problem is discontinuous at ``t = 0`` for
all ``u₀ \\neq 0``.

An analytical solution of this problem is provided for ``t \\in (-\\infty,10]``.
"""
prob_dde_1delay = build_prob_dde_1delay(1.0)

### Not in-place function

function f_1delay_notinplace(u, h, p, t)
    - h(p, t - oneunit(t)) ./ oneunit(t)
end

ff_1delay_notinplace = DDEFunction(f_1delay_notinplace,analytic=f_1delay_analytic)

#### Vectorized history function

build_prob_dde_1delay_notinplace(u₀, ::T=u₀) where {T} =
    DDEProblem(ff_1delay_notinplace, [u₀], (p, t)->[zero(u₀)], (zero(T), T(10)),
               constant_lags = [oneunit(T)])

"""
    prob_dde_1delay_notinplace

Same as [`prob_dde_1delay`](@ref), but purposefully implemented with a not in-place
function.
"""
prob_dde_1delay_notinplace = build_prob_dde_1delay_notinplace(1.0)

#### Scalar history function

build_prob_dde_1delay_scalar_notinplace(u₀, ::T=u₀) where {T} =
    DDEProblem(ff_1delay_notinplace, u₀, (p, t) -> zero(u₀), (zero(T), T(10)),
               constant_lags = [oneunit(T)])

"""
    prob_dde_1delay_scalar_notinplace

Same as [`prob_dde_1delay_notinplace`](@ref), but purposefully implemented with a scalar
history function.
"""
prob_dde_1delay_scalar_notinplace = build_prob_dde_1delay_scalar_notinplace(1.0)

## Two constant delays

### In-place function

function f_2delays(du, u, h, p, t::T) where T
    du[1] = (- h(p, t - T(1//3))[1] - h(p, t - T(1//5))[1]) / oneunit(t)
end

function f_2delays_analytic(u₀, p, t)
    z = t/oneunit(t)

    if z < 0
        return zero(u₀)
    elseif z < 1/5
        return u₀
    elseif z ≤ 1
        if z < 1/3
            c = @evalpoly(z, 6//5, -1)
        elseif z < 2/5
            c = @evalpoly(z, 23//15, -2)
        elseif z < 8/15
            c = @evalpoly(z, 121//75, -12//5, 1//2)
        elseif z < 3/5
            c = @evalpoly(z, 427//225, -52//15, 3//2)
        elseif z < 2/3
            c = @evalpoly(z, 4351//2250, -547//150, 9//5, -1//6)
        elseif z < 11/15
            c = @evalpoly(z, 539//250, -647//150, 23//10, -1//6)
        elseif z < 4/5
            c = @evalpoly(z, 7942//3375, -128//25, 17//5, -2//3)
        elseif z < 13/15
            c = @evalpoly(z, 39998//16875, -1952//375, 89//25, -4//5, 1//24)
        elseif z < 14/15
            c = @evalpoly(z, 10109//3750, -1583//250, 243//50, -13//10, 1//24)
        else
            c = @evalpoly(z, 171449//60750, -139199//20250, 2579//450, -173//90, 5//24)
        end

        return c .* u₀
    else
        error("This analytical solution is only valid on (-∞,1]")
    end
end

ff_2delays = DDEFunction(f_2delays,analytic=f_2delays_analytic)

build_prob_dde_2delays(u₀, ::T=u₀) where {T} =
    DDEProblem(ff_2delays, [u₀], (p, t) -> [zero(u₀)], (zero(T), oneunit(T)),
    constant_lags = [T(1//3), T(1//5)])

"""
    prob_dde_2delays

Model problem of finding a solution ``u(t)`` in the time span ``t \\in [0,1]``, where
``t`` is of type `T`, to the delay differential equation

```math
\\frac{du}{dt} = -u(t-1/3)-u(t-1/5)
```

with history function

```math
u(t) = \\begin{cases}
0 & \\text{if } t < 0 \\
u₀ & \\text{if } t = 0
\\end{cases}
```

for ``t \\leq 0`` with ``u₀ = 1``. Note that the problem is discontinuous at ``t = 0`` for
all ``u₀ \\neq 0``.

An analytical solution of this problem is provided for ``t \\in (-\\infty,1]``.
"""
prob_dde_2delays = build_prob_dde_2delays(1.0)

### Not in-place function

function f_2delays_notinplace(u, h, p, t::T) where T
    (- h(p, t - T(1//3)) .- h(p, t - T(1//5))) ./ oneunit(t)
end

ff_2delays_notinplace = DDEFunction(f_2delays_notinplace,analytic=f_2delays_analytic)

#### Vectorized history function

build_prob_dde_2delays_notinplace(u₀, ::T=u₀) where {T} =
    DDEProblem(ff_2delays_notinplace, [u₀], (p, t) -> [zero(u₀)], (zero(T), oneunit(T)),
               constant_lags =  [T(1//3), T(1//5)])

"""
    prob_dde_2delays_notinplace

Same as [`prob_dde_2delays`](@ref), but purposefully implemented with a not in-place
function.
"""
prob_dde_2delays_notinplace = build_prob_dde_2delays_notinplace(1.0)

#### Scalar history function

build_prob_dde_2delays_scalar_notinplace(u₀, ::T=u₀) where {T} =
    DDEProblem(ff_2delays_notinplace, u₀, (p, t) -> zero(u₀), (zero(T), oneunit(T)),
               constant_lags = [T(1//3), T(1//5)])

"""
    prob_dde_2delays_scalar_notinplace

Same as [`prob_dde_2delays_notinplace`](@ref), but purposefully implemented with a
scalar history function.
"""
prob_dde_2delays_scalar_notinplace = build_prob_dde_2delays_scalar_notinplace(1.0)

# DDE examples without analytical solution

## Single constant delay

### In-place function

function f_1delay_long(du, u, h, p, t::T) where T
    du[1] = (- h(p, t - T(1//5))[1] + u[1]) / oneunit(t)
end

build_prob_dde_1delay_long(u₀, ::T=u₀) where {T} =
    DDEProblem(f_1delay_long, [u₀], (p, t) -> [zero(u₀)], (zero(T), T(100)),
    constant_lags = [T(1//5)])

"""
    prob_dde_1delay_long

Model problem of finding a solution ``u(t)`` in the time span ``t \\in [0,100]`` to the
delay differential equation

```math
\\frac{du}{dt} = -u(t-0.2) + u(t)
```

with history function

```math
u(t) = \\begin{cases}
0 & \\text{if } t < 0,\\
u₀ & \\text{if } t = 0,
\\end{cases}
```

for ``t \\leq 0`` with ``u₀ = 1``. Note that the problem is discontinuous at ``t = 0`` for
all ``u₀ \\neq 0``.
"""
prob_dde_1delay_long = build_prob_dde_1delay_long(1.0)

### Not in-place function

function f_1delay_long_notinplace(u, h, p, t::T) where T
    (- h(p, t - T(1//5)) .+ u ) ./ oneunit(t)
end

build_prob_dde_1delay_long_notinplace(u₀, ::T=u₀) where {T} =
    DDEProblem(f_1delay_long_notinplace, [u₀], (p, t) -> [zero(u₀)], (zero(T), T(100)),
               constant_lags = [T(1//5)])

"""
    prob_dde_1delay_long_notinplace

Same as [`prob_dde_1delay_long`](@ref), but purposefully implemented with a not
in-place function.
"""
prob_dde_1delay_long_notinplace = build_prob_dde_1delay_long_notinplace(1.0)

build_prob_dde_1delay_long_scalar_notinplace(u₀, ::T=u₀) where {T} =
    DDEProblem(f_1delay_long_notinplace, u₀, (p, t) -> zero(u₀), (zero(T), T(100)),
               constant_lags = [T(1//5)])

"""
    prob_dde_1delay_long_scalar_notinplace

Same as [`prob_dde_1delay_long_notinplace`](@ref), but purposefully implemented with a
scalar history function.
"""
prob_dde_1delay_long_scalar_notinplace = build_prob_dde_1delay_long_scalar_notinplace(1.0)

## Two constant delays

### In-place function

function f_2delays_long(du, u, h, p, t::T) where T
    du[1] = (- h(p, t - T(1//3))[1] - h(p, t - T(1//5))[1]) / oneunit(t)
end

build_prob_dde_2delays_long(u₀, ::T=u₀) where {T} =
    DDEProblem(f_2delays_long, [u₀], (p, t) -> [zero(u₀)], (zero(T), T(100)),
               constant_lags = [T(1//3), T(1//5)])

"""
    prob_dde_2delays_long

Model problem of finding a solution ``u(t)`` in the time span ``t \\in [0,100]`` to the
delay differential equation

```math
\\frac{du}{dt} = -u(t-1/3) - u(1-1/5)
```

with history function

```math
u(t) = \\begin{cases}
0 & \\text{if } t < 0,\\
u₀ & \\text{if } t = 0,
\\end{cases}
```

for ``t < 0`` with ``u₀ = 1``. Note that the problem is discontinuous at ``t = 0`` for all
``u₀ \\neq 0``.
"""
prob_dde_2delays_long = build_prob_dde_2delays_long(1.0)

### Not in-place function

function f_2delays_long_notinplace(u, h, p, t::T) where T
    (- h(p, t - T(1//3)) .- h(p, t - T(1//5))) ./ oneunit(t)
end

#### Vectorized history function

build_prob_dde_2delays_long_notinplace(u₀, ::T=u₀) where {T} =
     DDEProblem(f_2delays_long_notinplace, [u₀], (p, t) -> [zero(u₀)], (zero(T), T(100)),
                constant_lags = [T(1//3), T(1//5)])

"""
    prob_dde_2delays_long_notinplace

Same as [`prob_dde_2delays_long`](@ref), but purposefully implemented with a not
in-place function.
"""
prob_dde_2delays_long_notinplace = build_prob_dde_2delays_long_notinplace(1.0)

#### Scalar history function

build_prob_dde_2delays_long_scalar_notinplace(u₀, ::T=u₀) where {T} =
     DDEProblem(f_2delays_long_notinplace, u₀, (p, t) -> zero(u₀), (zero(T), T(100)),
                constant_lags = [T(1//3), T(1//5)])

"""
    prob_dde_2delays_long_scalar_notinplace

Same as [`prob_dde_2delays_long_notinplace`](@ref), but purposefully implemented with a
scalar history function.
"""
prob_dde_2delays_long_scalar_notinplace = build_prob_dde_2delays_long_scalar_notinplace(1.0)

#=
The following examples are taken from:
W.H. Enright and H. Hayashi, The evaluation of numerical software for delay
differential equations.
=#

function f_dde_mackey(du, u, h, p, t)
    du[1] = 0.2*h(p, t-14)[1]/(1 + h(p, t-14)[1]^10) - 0.1*u[1]
end

"""
    prob_dde_mackey

Model of blood production with constant delay (M. C. Mackey and L. Glass, Oscillation and
chaos in physiological control systems, 1977).
"""
prob_dde_mackey = DDEProblem(f_dde_mackey, [0.5], (p, t) -> [0.5], (0.0, 500.0),
                             constant_lags = [14])

function f_dde_wheldon(du, u, h, p, t)
    du[1] = 1.1/(1 + sqrt(10)*(h(p, t-20)[1])^(5/4)) - 10*u[1]/(1 + 40*u[2])
    du[2] = 100*u[1]/(1 + 40*u[2]) - 2.43*u[2]
end
u0_wheldon = [1.05767027/3; 1.030713491/3]

"""
    prob_dde_wheldon

Model of chronic granulocytic leukemia with constant delay (T. Wheldon, J. Kirk and
H. Finlay, Cyclical granulopoiesis in chronic granulocytic leukemia: A simulation study, 1974).
"""
prob_dde_wheldon = DDEProblem(f_dde_wheldon, u0_wheldon, (p, t) -> u0_wheldon, (0., 100.),
                              constant_lags = [20])

function f_dde_neves1(du, u, h, p, t)
    du[1] = 1 - h(p, exp(1-1/t))[1]
end
# only valid for specific history function
function f_dde_neves1_analytic(u₀, p, t )
    0 < t ≤ 10 && u₀ == [log(0.1)] && return [log(t)]
    error("This analytical solution is only valid on (0, 10] and for history function ϕ(t) = ln(t) for 0 < t ≤ 0.1")
end

ff_dde_neves1 = DDEFunction(f_dde_neves1,analytic=f_dde_neves1_analytic)

"""
    prob_dde_neves1

DDE with vanishing time dependent delay at ``t = 1`` (K. W. Neves, Automatic integratorion
of functional differential equations: An approach, 1975).
"""
prob_dde_neves_1 = DDEProblem(ff_dde_neves1, [log(0.1)], (p, t) -> [log(t)], (0.1, 10.),
                              dependent_lags = [(u,p,t) -> t - exp(1 - 1/t)])

function f_dde_neves_thompson(du, u, h, p, t)
    if h(p, t/2)[1] < 0
        du[1] = 1 - u[1]
    else
        du[1] = -1 - u[1]
    end
end
# only valid for specific history function
function f_dde_neves_thompson_analytic(u₀, p, t )
    if u₀ == [1]
        if 0 ≤ t ≤ 2*log(2)
            return [2*exp(-t) - 1]
        elseif t ≤ 2*log(6)
            return [1 - 6*exp(-t)]
        elseif t ≤ 2*log(66)
            return [66*exp(-t) - 1]
        end
    end
    error("This analytical solution is only valid on [0, 2ln(66)] and for history function ϕ(t) = 1 for t ≤ 0")
end

ff_dde_neves_thompson = DDEFunction(f_dde_neves_thompson,
                                    analytic=f_dde_neves_thompson_analytic)

"""
    prob_dde_neves_thompson

DDE with vanishing time dependent delay at ``t = 0`` (K. W. Neves and S. Thompson, Solution
of systems of functional differential equations with state dependent delays, 1992), given by

```math
u'(t) = \\begin{cases}
-1 - u(t) & \\text{if } u(t/2) \\geq 0 \\
1 - u(t) & \\text{if } u(t/2) < 0
\\end{case}
```

for ``t \\in [0, 2*\\log 66]`` with history function

```math
u(t) = 1
```

for ``t \\leq 0``.
"""
prob_dde_neves_thompson = DDEProblem(ff_dde_neves_thompson, [1.],
                                     (p, t) -> [1.], (0., 2*log(66)),
                                     constant_lags = [(u,p,t) -> t/2])

function f_dde_paul1(du, u, h, p, t)
    du[1] = - 2*h(p, t - 1 - abs(u[1]))[1]*(1 - u[1]^2)
end

"""
    prob_dde_paul1

DDE with state dependent delay (C. A. H. Paul, A test set of functional differential
equations, 1994), given by

```math
u'(t) = - 2u(t - 1 - |u(t)|)(1 - u(t)^2)
```

for ``t \\in [0, 30]`` with history function

```math
u(t) = 1/2
```

for ``t \\leq 0``.
"""
prob_dde_paul1 = DDEProblem(f_dde_paul1, [0.5], (p, t) -> [0.5], (0., 30.),
                            dependent_lags = [(u,p,t) -> 1 + abs(u[1])])

function f_dde_paul2(du, u, h, p, t)
    h1 = h(p, t - u[2])[1]
    du[1] = -2*h1
    du[2] = (abs(h1) - abs(u[1]))/(1 + abs(h1))
end

"""
    prob_dde_paul2

DDE with state dependent delay (C. A. H. Paul, A test set of functional differential equations, 1994).
"""
prob_dde_paul2 = DDEProblem(f_dde_paul2, [1; 0.5], (p, t) -> [1; 0.5], (0., 30.),
                            dependent_lags = [(u,p,t) -> u[2]])

function build_prob_dde_mahaffy(tspan, h, σ₀, T₁, γ, Q, k, a, K, r)
    function f(du, u, h, p, t)
        du[1] = σ₀*h(p, t-T₁)[2] - γ*u[1] - Q
        du[2] = a/(1 + K*u[1]^r) - k*u[2]
        du[3] = 1 - Q*exp(γ*u[3])/(σ₀*h(p, t-T₁-u[3])[2])
    end

    DDEProblem(f, h(nothing, 0), h, tspan,
                constant_lags = [T₁],
                dependent_lags = [(u,p,t) -> T₁ + u[3]])
end

function h_mahaffy1(p, t)
    if t < -6
        return [3.325; 9.5; 120]
    else
        return [3.325; 10; 120]
    end
end
prob_dde_mahaffy1 = build_prob_dde_mahaffy((0., 300.), h_mahaffy1, 0.0031, 6, 0.001, 0.0275,
                                           2.8, 6570, 0.0382, 6.96)

prob_dde_mahaffy2 = build_prob_dde_mahaffy((0., 100.), (p, t) -> [3.5; 10; 50], 0.00372, 3, 0.1,
                                           0.00178, 6.65, 15600, 0.0382, 6.96)

"""
Two models of hematopoiesis with state dependent delay (J. M. Mahaffy, J. Bélair and
M. C. Mackey, Hematopoietic model with moving boundary condition and state dependent delay, 1996).
"""
prob_dde_mahaffy1, prob_dde_mahaffy2

function f_dde_neves2(du, u, h, p, t)
    t2 = exp(1 - u[2])
    du[1] = u[2]
    du[2] = -h(p, t2)[2]*u[2]^2*t2
end
# only valid for specific history function
function f_dde_neves2(::Type{Val{:analytic}}, u₀, p, t )
    0 < t ≤ 5 && u₀ == [log(0.1); 10] && return [log(t); 1/t]
    error("This analytical solution is only valid on (0, 5) and for history function ϕ(t) = [ln(t); 1/t] for 0 < t ≤ 0.1")
end

"""
    prob_dde_neves2

DDE with vanishing state dependent delay at ``t = 1`` (K. W. Neves, Automatic integration
of functional differential equations: An approach, 1975).
"""
prob_dde_neves2 = DDEProblem(f_dde_neves2, [log(0.1); 10], (p, t) -> [log(t); 1/t], (0.1, 5.),
                             dependent_lags = [(u,p,t) -> t - exp(1 - u[2])])

function build_f_dde_gatica(r₁, r₂, α, δ)
    function f_dde_gatica(du, u, h, p, t)
        u₁u₂ = u[1]*u[2]
        r₁u₁u₂ = r₁*u₁u₂
        r₂u₃ = r₂*u[3]
        uₕ = h(p, t - u[4])
        uₕ₁₂ = uₕ[1]*uₕ[2]

        du[1] = -r₁u₁u₂ + r₂u₃
        du[2] = -r₁u₁u₂ + α*r₁*uₕ₁₂
        du[3] = r₁u₁u₂ - r₂u₃
        du[4] = 1 + (3*δ - u₁u₂ - u[3])/(uₕ₁₂ + uₕ[3])*exp(δ*u[4])
    end
end

"""
    prob_dde_gatica

Model of antigen antibody dynamics with fading memory, with vanishing state dependent delay
at ``t = 0`` (J. Gatica and P. Waltman, A threshold model of antigen antibody dynamics with fading memory, 1982).
"""
prob_dde_gatica = DDEProblem(build_f_dde_gatica(0.02, 0.005, 3, 0.01),
                             [5; 0.1; 0; 0], (p, t) -> [5; 0.1; 0; 0],
                             (0., 40.), dependent_lags = [(u,p,t) -> u[4]])

#=
Quorum Sensing model
=#
function build_prob_dde_qs(u₀, tspan, τ, D, γₛ, Kₘ, nₛ, a, αₐ, βₐ, C₁, n₁, γₐ, αC, R, γ₃, Kₑ, αₗ, C₂, n₂, γₗ)
    function hill_growth(x, K, n)
        if x > K
            return 1/(1 + (K / x)^n)
        end

        if K == 0
            return one(x)
        end

        y = (x / K)^n
        return y / (1 + y)
    end

    S₀ = u₀[1]

    function f_dde_qs(du, u, h, p, t)
        if u[1] < 0
            # eq. 1 and 2 not defined for u[1] < 0
            du[1] = 0
            du[2] = 0
        else
            tmp = u[2] * hill_growth(u[1], Kₘ, nₛ)
            du[1] = D * (S₀ - u[1]) - γₛ * tmp
            du[2] = a * tmp - D * u[2]
        end

        if u[4] < 0
            # eq. 3 not defined for u[4] < 0
            du[3] = 0
            du[4] = max(0, αC * (R - u[4]) * u[3] - γ₃ * u[4])
        else
            du[3] = (αₐ + βₐ * hill_growth(u[4], C₁, n₁)) * u[2] -
                (γₐ + D) * u[3] - αC * (R - u[4]) * u[3] + γ₃ * u[4] -
                Kₑ * u[3] * u[5]
            du[4] = αC * (R - u[4]) * u[3] - γ₃ * u[4]
        end

        Cₜ = h(p, t - τ)[4]
        if Cₜ < 0
            # eq. 5 not defined for Cₜ < 0
            du[5] = 0
        else
            du[5] = αₗ * hill_growth(Cₜ, C₂, n₂) * u[2] - (γₗ + D) * u[5]
        end
    end

    DDEProblem(f_dde_qs, u₀, (p, t) -> u₀, tspan, constant_lags = [τ])
end

"""
    prob_dde_qs

Model of Quorum Sensing (QS) of Pseudomonas putida IsoF in continuous cultures with constant
delay (Buddrus-Schiemann et al., Analysis of N-Acylhomoserine Lactone Dynamics in Continuous
Cultures of Pseudomonas Putida IsoF By Use of ELISA and UHPLC/qTOF-MS-derived Measurements
and Mathematical Models, Analytical and Bioanalytical Chemistry, 2014).
"""
prob_dde_qs = build_prob_dde_qs([1; 8.4e8; 2.5e-9; 7.6e-8; 5e-15], # initial values
                                (0., 45.), # time span
                                2, # delay
                                0.1, 1.3e-12, 0.38, 1.3, 0.66, 2.3e-19, 2.3e-18, 70e-9,
                                2.3, 0.05, 4e4, 5e-7, 0.080, 1.5e-4, 1.1e-8, 70e-9, 2.5,
                                0.005)


### A Pluto.jl notebook ###
# v0.14.2

using Markdown
using InteractiveUtils

# ╔═╡ 1bc961f0-66f7-11eb-17c8-2b54d4c4328d
begin
	push!(LOAD_PATH, ".")
	using CellularBase
end

# ╔═╡ f8149ae0-66f6-11eb-0335-39baafdb24c8
begin
	using Plots
	using Images
	using PlutoUI
	using Statistics
end

# ╔═╡ 9a74ef90-67ad-11eb-1fc3-e7080fbfc7e1
html"""
Term Paper 1

This work is a submission for the course Modelling Complex Systems(IDC-631) for the year 2020-21 Spring (Even) Semester by Dhruva Sambrani, MS18163
"""

# ╔═╡ eb2aad32-687f-11eb-0c6b-91b443cdb3d3
md"""# Abstract
This paper looks at the effect of a Lockdown on the spread of a pathogen. This setup is modelled via a Probabilisitic/Non Deterministic 2D cellular automaton.

This paper first [introduces](#7cb22830-67c9-11eb-0e0e-41f2316ed468) some working definitions along with the same definitions in code.
Then features are added, and parallelly introduced to the reader, [layer](#c0318a00-67ac-11eb-39b7-e3c13a5aff57) upon [layer](#1cf29850-66fa-11eb-3f76-db755c0a4c50) with logical reasoning for their implementations.
Following this, the actual topic of the paper is introduced with some [results](#5b1e6e50-6877-11eb-1ca4-3df17b85575e) that show the efficiency of [Lockdowns](#9e71a7c0-67f1-11eb-0282-672d358d2292).
Finally, [future scope of this work](#c7ea4fa0-687b-11eb-0dbc-7f316bd4aefd) is also discussed.

!!! note "Note to Evaluator"
    The author believes that papers which involve a lot of simulation and code should also be accompanied with the code they use so that readers can gain deeper insight into the processes described in the paper. Hence, the author has attempted to create a paper which interlaces the information with the corresponding code.

The coding language used is `Julia`[^julia] and following document is made with `Pluto.jl` [^pluto].

"""

# ╔═╡ 6f7bce82-67ad-11eb-1ec1-79ce2620fe9b
PlutoUI.TableOfContents(title="Term Paper 1 - Index")

# ╔═╡ 3459c040-6858-11eb-09c1-95c7bdb02722
html""

# ╔═╡ 7cb22830-67c9-11eb-0e0e-41f2316ed468
md"# Code Setup"

# ╔═╡ 65ff1d80-66f9-11eb-11c2-69c47b8458e4
md"## `Include packages`"

# ╔═╡ 330799c0-6884-11eb-230d-71c512e4a447
md"`CellularBase`[^cellularbase] is a package written by the author themselves to abstract out certain repetitive funcitions for all Cellular Automata."

# ╔═╡ db4fd9e0-6883-11eb-384e-6df2c1d7a52b
md"The `Plots.jl`[^plots] and `Images.jl`[^images] packages are used for plotting and the `Statistics`[^stats] Package for some statistical analysis. `PlutoUI.jl`[^plutoui] is used for some UI elements such as the ToC"

# ╔═╡ a95c8c70-67ad-11eb-124b-7b66d5547321
md"## Define `Structures` and `Functions`"

# ╔═╡ 73ea0ae0-66f9-11eb-19cf-b7ab6151ef63
md"### Define `LifeGrid`"

# ╔═╡ 038cbdf2-6880-11eb-1b2b-2b4c4974be49
md"""This is the basic structure that defines all the parameters of a particular "grid". The name `LifeGrid` was given by mistake but was not changed as it would entail too many other changes"""

# ╔═╡ 91afcec0-66f9-11eb-37b1-f78de5e0ba8a
md"### Define method of evolution"

# ╔═╡ fcc23790-686e-11eb-03d7-1f0946446f49
md"The method of evolution follows the following algorithm -
1. If the agent is susceptible, infect the agent depending on neighborhood and return. The exact dependacy is discussed later in [Non Deterministic runs](#2563e210-67d5-11eb-089e-f376da973765)
2. Otherwise, the state is increased by 1
3. If the state is recovered, set it to 0

Mathematically this can be written as

$s_{i+1} = \begin{cases}
    1, & s_i=0 \text{ and } \rho(\text{neighbours}) \\
    0, & s_i=0 \text{ and } 1-\rho(\text{neighbours}) \\
	(s_i+1) \mod (\tau_i+\tau_r+1) & \text{otherwise}
\end{cases}$

"

# ╔═╡ b8e7c920-66f9-11eb-397f-3760662a648c
md"### Define States"

# ╔═╡ 8b6833e0-6870-11eb-2593-69edaec0a7b8
md"States are defined in the following way
1. Susceptible $\implies 0$
2. Infected $\implies i \in [1, \tau_i]$
3. Refactory $\implies i \in [\tau_i + 1, \tau_i+\tau_r]$
4. Recovered $\implies i = (\tau_i+\tau_r + 1)$
"

# ╔═╡ de865ed2-66f9-11eb-31dc-c782d4338af3
md"### Define helper `functions`"

# ╔═╡ c3e6cc10-66f7-11eb-039f-df82addf1a5e
begin
	CellularBase.possible_states(l::LifeGrid) = 0:(l.τi + l.τr)
	CellularBase.newstate(grid::LifeGrid) = grid.otherstate;
	function CellularBase.state!(grid::LifeGrid, newstate)
	    grid.otherstate = grid.state
	    grid.state = newstate
	end
end

# ╔═╡ 4d4cdcc0-67b0-11eb-2278-8de64128d6f1
function Base.show(io::IO, ::MIME"text/html", grid::LifeGrid)
	final = Markdown.parse(
		"---" *
		"\n\n**Grid Type:** " * repr(typeof(grid), context=io) *
		"\n\n**\$\\tau_i:\$** " * string(grid.τi) *
		"\n\n**\$\\tau_r:\$** " * string(grid.τr) *
		"\n\n**Max neighbors for probabilistic spread:** " * string(grid.max) *
		"\n\n**Possible States:** " * string(possible_states(grid)) *
		"\n\n**Boundary Condition:** " * string(boundary_conditions(grid)) *
		"\n\n**Size:** " * string(size(grid)) *
		"\n\n**Neighborhood:** "
	)
	show(io, MIME("text/html"), final)
	show(io, MIME("text/html"), grid.neighborhood)
	show(io, MIME("text/html"), md"---")
end

# ╔═╡ 7cdebc60-684b-11eb-38a2-45873e4b51d0
function makegraph(ldgrid, ldresults, ldradii)
    p = plot(1:length(ldresults), ldresults .|> frame -> count(frame .|> p -> susceptible(p, ldgrid)) / length(state(ldgrid)), color=:green, label="Sus")
    plot!(1:length(ldresults), ldresults .|> frame -> count(frame .|> p -> infectious(p, ldgrid)) / length(state(ldgrid)), color=:red, label="Inf")
    plot!(1:length(ldresults), ldresults .|> frame -> count(frame .|> p -> refactory(p, ldgrid)) / length(state(ldgrid)), color=:blue, label="Ref")
    q = plot(ldradii, label="Neighbourhood Radius", legend=:bottomright)
    plot(p, q, size=(600, 500), layout=(2, 1))
end;

# ╔═╡ 5031733e-66f9-11eb-107f-6d69089d9384
function color(i, grid)
    if susceptible(i, grid)
        RGB(0.84, 0.95, 0.80)
    elseif infectious(i, grid)
        RGB(0.83, 0.11, 0.31)
    else
        RGB(0.16, 0.11, 0.42)
    end
end;

# ╔═╡ c0318a00-67ac-11eb-39b7-e3c13a5aff57
md"# Deterministic Case - A First Run"

# ╔═╡ 47b81340-67cb-11eb-372e-972e4d4ff746
md"## Introduction
As a first trial, we can try to run the simulation for the deterministic case, where the infection can spread even if 1 neighbor is infected. We also verify that the results observed are very similar to results obtained in [^Moitra2019]
"

# ╔═╡ 02c24a70-67ae-11eb-2944-61a7963eb947
LifeGrid((100, 100), 4, 9, max=1, nh=VonNeumannNeighborhood(1, 2))

# ╔═╡ fbdba862-67cf-11eb-39bc-f37eb3662785
md"## `state(center patient) = 1`"

# ╔═╡ fe0dda70-66fa-11eb-0e02-73c05666aace
dgrid, dresults = let
	grid = LifeGrid(
		(100, 100),
		4, 9,
		max=1,
		nh=VonNeumannNeighborhood(1, 2));
	results = simulate!(grid, 150)
	grid, results
end;

# ╔═╡ 8315d810-677a-11eb-3953-8faad42ada40
@gif for frame ∈ dresults[5:end]
	plot(frame .|> p -> color(p, dgrid))
end

# ╔═╡ 7d24e300-67d0-11eb-0864-41c6166f13ac
md"## Random states from 0 to $\tau_i$"

# ╔═╡ b8523010-67ce-11eb-2396-4780284dfa01
randgrid, randresults = let
	arr = rand(0:4, 100, 100)
	grid = LifeGrid(
		arr,
		4, 5,
		max=1,
		nh=VonNeumannNeighborhood(1, 2));
	results = simulate!(grid, 50)
	grid, results
end;

# ╔═╡ cf22a450-67ce-11eb-2141-fd6f3efc689a
@gif for frame ∈ randresults
	plot(frame .|> p -> color(p, randgrid))
end

# ╔═╡ 07c440a0-67d1-11eb-3a8b-d559cac6276e
md"## `state(center agent) = τᵢ; τᵢ=τᵣ`"

# ╔═╡ 30ce7a10-67bd-11eb-027d-8dc583201bc1
ϕgrid, ϕresults = let
	arr = zeros(Int, 100, 100)
	arr[50,50] = 4
	grid = LifeGrid(
		arr,
		4, 4,
		max=1,
		nh=VonNeumannNeighborhood(1, 2));
	results = simulate!(grid, 50)
	grid, results
end;

# ╔═╡ 4ca7fe50-67bd-11eb-382c-1304f3c1c3c4
@gif for frame ∈ ϕresults
	plot(frame .|> p -> color(p, ϕgrid))
end

# ╔═╡ 643ffb90-67ad-11eb-20c6-7744050e5448
md"""## Observations in the deterministic case

We can see that the infection spreads outward and then dies off.

The refactory agents act as a buffer between the leading edge of infected agents and the recovered(susceptible) agents inside the wave. This is similar to how a batch of burnt material doesn't allow the fire to move back in.

!!! tip "Spread Back"
    The infection can spread back into the wave only if either -
    -
      1. the radius of the neighborhood is larger than the refactory period
      2. if $\tau_i \ge \tau_r$ and $\text{state(patient zero)}\ge\tau_r$

"""

# ╔═╡ 1cf29850-66fa-11eb-3f76-db755c0a4c50
md"# Non Deterministic Case"

# ╔═╡ 393a5bc0-67cb-11eb-3be1-a9ec7943b5f3
md"""## Why?
Mainly because determinism is so 1900s.

Seriously though, the reason why we don't see too many interesting things in the deterministic case is because of the refactory buffer separating the infected leading edge and the recovered members.

Allowing some agents to escape the infection wave allows breaks in the refactory buffer allowing the infection to come back into the recovered agents. It seems very sadistic to play for the pathogen, but what's the fun otherwise.
"""

# ╔═╡ 2563e210-67d5-11eb-089e-f376da973765
md"""## How is it implemented?
In code, obviously. Oh you mean what's the logic behind it?

If the spread is probabilistic, one way to do it would be to set a non zero probability if the number of infected neighbors is non-negative. While this does make the spread probabilistic, it doesn't model _real life_ very well. In real life, one's chances of infection increases with the increase in infected neighbors. While the exact distribution is unknown, we can model the spread by assuming that the chance of spreading follows

$\rho_\text{infect}{}(\text{count}) = \begin{cases}
  \frac{\text{count}}{\text{max}} & 0\leq \text{count}\leq \text{max} \\
  1 & \text{max}\leq \text{count}
\end{cases}$

where $\text{max}$ is a parameter of the simulation. This basically means that for every infected neighbour, the chance of infection is linear, until a limit, after which the agent will always get infected.

This definition has the added benefit of being independent on the definition of the neighbourhood, and can be attributed only to the pathogen in question.

Every `LifeGrid` has a critical limit of infected neighbors, denoted by `max`. For every agent, $\rho$ is calculated and infection is done by a Bernoulli trial.

!!! info "High values of max"

    Hence, it also makes sense to set `max` higher than the size of the neighbourhood itself, for that would allow us to set the highest possible probability to be lesser than 1.

    After all, what if the neighborhood changes during the run? How? Keep reading
"""

# ╔═╡ dafac300-67e9-11eb-0972-f9070193399e
md"## Max = 4"

# ╔═╡ 46055e90-67e9-11eb-2b6f-dffa58138005
ρgrid, ρresults = let
	grid = LifeGrid(
		(100, 100),
		4, 5,
		max=4,
		nh=VonNeumannNeighborhood(1, 2));
	results = simulate!(grid, 400)
	grid, results
end;

# ╔═╡ 5626db00-67e9-11eb-0e12-a77b6b9c58be
@gif for frame ∈ ρresults
	plot(frame .|> p -> color(p, ρgrid))
end every 2

# ╔═╡ d4c2cff0-67e9-11eb-1f5b-03c7bab91782
md"## Max = 9"

# ╔═╡ 3056cdae-67ec-11eb-0b35-b3d6cde9b0fa
ρ7grid, ρ7results = let
	grid = LifeGrid(
		(100, 100),
		4, 5,
		max=9,
		init=5,
		nh=VonNeumannNeighborhood(1, 2));
	results = simulate!(grid, 100)
	grid, results
end;

# ╔═╡ 3b0024a0-67ec-11eb-3398-eb07250758a9
@gif for frame ∈ ρ7results
	plot(frame .|> p -> color(p, ρ7grid))
end

# ╔═╡ 6d822d52-67ed-11eb-1b71-679d491bf133
md"## Max = 2"

# ╔═╡ 76e43730-67ed-11eb-201f-53bf8d3e3bfb
ρ2grid, ρ2results = let
	grid = LifeGrid(
		(100, 100),
		4, 5,
		max=2,
		nh=VonNeumannNeighborhood(1, 2));
	results = simulate!(grid, 150)
	grid, results
end;

# ╔═╡ 8445d050-67ed-11eb-2411-f776a43c7f40
@gif for frame ∈ ρ2results
	plot(frame .|> p -> color(p, ρ2grid))
end every 2

# ╔═╡ 0f5ba100-67ef-11eb-15f5-3fbaad70dde2
md"""## Observations

Evidently, higher max values imply smaller probability of spreading. This is evident in the run with `max = 9`, where very few infected cases are seen at any point in time, and agents often recover before they spread it on, leading to the pathogen dying out.

On the flip side, if `max` is too close to 1, then there it acts very similar to the deterministic case _most of the time_. This is also evident in the fact that the case where `max` is 1, it is exactly the deterministic case which we've seen also dies out.

!!! tip "Relevant Quote"
    It is possible to have too much of a good thing

    \- **Aesop**
"""

# ╔═╡ 9e71a7c0-67f1-11eb-0282-672d358d2292
md"""# Lockdown
## Introduction
The word _lockdown_ seems to have taken over our lives and completely changed the way that we go about our daily lives. But why are we doing this? And how effective are these lockdowns? We explore this question by looking at how we can simulate a lockdown period for our agents and see how this affects the infection spread.
"""

# ╔═╡ ef666b40-67f4-11eb-2f31-05145a6dfcba
md"""## The logic behind lockdowns

!!! info "Definition"
    A lockdown is a restriction policy for people or community to stay where they are, usually due to specific risks to themselves or to others if they can move and interact freely - [Wikipedia](https://en.wikipedia.org/wiki/Lockdown)

From the definition, it is clear that the intent is to prevent people from interacting and hence hopefully reducing the spread of the disease. If an agent interacts with fewer people, they have a lower chance of getting infected.
"""

# ╔═╡ 5586a260-67f4-11eb-11e2-b706fd0d4dfe
md"""## Neighborhoods, Virulence and Infection Rate
Neighbourhood size and intrinsic virulence of a pathogen together decide the pathogen's spreading capabilities.

Neighbourhood size is a representation of the amount of inter-agent interactions. Due to the way the model is defined, it also shows how _far_ the agent can influence. It defines the number of agents one agent can meet.

Intrinsic virulence of a pathogen is how well it spreads given that two agents meet. There is sufficient discussion regarding this in the above section.

If both these parameters define, in some sense, the same thing, why not just club them together into a single parameter which gives the net spread?

The reason is 3 fold-

1. The **neighborhood is a function of the agents' behavior**, the **virulence is a function of the pathogen's behaviour**. It makes sense to keep them separate to deal with different cases of spread changes.
2. The model is such that there are **geometric differences** that arise by changing the neighborhood radius. The _buffering_ action of the refactory line is reduced by a larger neighborhood. It is not possible to increase the infection rate without also reducing this phenomenon if both were clubbed together.
3. The code would need too many changes.

!!! tips "Relevant quote"
    Any sufficiently advanced bug is indistinguishable from a feature.

    \- **Rich Kulawiec**

"""

# ╔═╡ 21b80dc0-67f9-11eb-0e87-6db59c88c296
md"""## Lockdowns and Neighborhoods

From the above discussion, it must be clear that a lockdown can be simulated pretty easily by varying the size of the neighborhood. The lockdown must also start at some predetermined point during the simulation. The size of the neighbourhood is changed by changing `grid.neighborhood` and the in-simulation change of the neighborhood can be executed by using the `postrunhook` of the `simulate!` function

"""

# ╔═╡ ead37ab2-66f9-11eb-2e56-07f105753991
function lockdown(grid, step; step_dict, radii)
    push!(radii, radius(grid.neighborhood))
    if step in keys(step_dict)
        grid.neighborhood = VonNeumannNeighborhood(step_dict[step], 2);
    end
end

# ╔═╡ 5641a270-684b-11eb-29b4-ef8c59ef9679
md"## 200 → 1"

# ╔═╡ 0f20b710-6840-11eb-1712-9f4614e9ed87
ld1 = let
	step_dict = Dict(200 => 1) # Step -> radius
	radii = []
	grid = LifeGrid(
		(150, 150),
		4, 5,
		max=9,
		init=5,
		nh=VonNeumannNeighborhood(3, 2));
	results = simulate!(grid, 400;
		postrunhook=lockdown, step_dict=step_dict, radii=radii);
	grid, results, radii
end;

# ╔═╡ 6f9dc770-6847-11eb-14ff-2b2e8562076f
@gif for frame ∈ ld1[2]
	plot(frame .|> p -> color(p, ld1[1]))
end every 2

# ╔═╡ 2bf1c112-6848-11eb-3b4a-a3ce569da696
makegraph(ld1...)

# ╔═╡ a9588a00-684b-11eb-2585-09a2cdf9ad80
md"## 50 → 2; 150 → 1; 300 → 3"

# ╔═╡ 4f4fdfd0-684c-11eb-3769-49d1677a1d52
ld2 = let
	step_dict = Dict(50 => 2, 150 => 1, 300 => 3) # Step -> radius
	radii = []
	grid = LifeGrid(
		(150, 150),
		4, 5,
		max=9,
		init=5,
		nh=VonNeumannNeighborhood(3, 2));
	results = simulate!(grid, 400;
		postrunhook=lockdown, step_dict=step_dict, radii=radii);
	grid, results, radii
end;

# ╔═╡ d97dddfe-684c-11eb-1d11-5734830507a2
@gif for frame ∈ ld2[2]
	plot(frame .|> p -> color(p, ld2[1]))
end every 2

# ╔═╡ e1baa1c0-684c-11eb-292c-3336d6d715b1
makegraph(ld2...)

# ╔═╡ 52123eb0-684d-11eb-0179-895599ae7a2f
md"## 100 → 2; 300 → 3"

# ╔═╡ 4e182db0-684d-11eb-0bab-b7c22d3f3ff5
ld3 = let
	step_dict = Dict(100 => 2, 300 => 3) # Step -> radius
	radii = []
	grid = LifeGrid(
		(150, 150),
		4, 5,
		max=9,
		init=5,
		nh=VonNeumannNeighborhood(3, 2));
	results = simulate!(grid, 400;
		postrunhook=lockdown, step_dict=step_dict, radii=radii);
	grid, results, radii
end;

# ╔═╡ 53bb3e50-684e-11eb-1bd3-d582200be8b5
@gif for frame ∈ ld3[2]
	plot(frame .|> p -> color(p, ld3[1]))
end every 2

# ╔═╡ 660d5660-684e-11eb-27ad-cbde3f616410
makegraph(ld3...)

# ╔═╡ f0c45ab2-684e-11eb-24ca-c5942b62f396
md"""## Observations

Apart from being very flashy, we see an interesting phenomenon. When the radius of the neighbourhood is changed, it takes about 25-30 steps for the grid to find a new equilibrium. During this transient phase, the distribution of states can vary wildly.

Further, the equilibrium fractions of the states averaged over a few cycles is nearly equal, but the exact variation may vary on the basis of initial conditions ([See 100 → 2; 300 → 3](#52123eb0-684d-11eb-0179-895599ae7a2f)). This can be attributed to the geometry of the neighborhood, the refactory buffer and the distribution of infected agents.

Because most large scale structures don't last for long times, we can assume that the final state distribution is memoryless or has only short-term memory. Thus our study of the effectiveness of a lockdown can be reduced to a simpler study of equilibrium infection as a function of neighborhood radius and [`grid.max`](#2563e210-67d5-11eb-089e-f376da973765), which, as you probably guessed it, is our next section.
"""

# ╔═╡ 61e28250-6857-11eb-39f9-2ff5565519ab
md"""
# Equilibrium infection distribution
**As a function of Neighbourhood Radius and Virulence of pathogen**
"""

# ╔═╡ 2734ec90-6859-11eb-2430-e3bd97302fd8
md"""## Introduction
This section is mainly computational. With the functionalities we have built up in the previous sections, we simply simulate a random 200X200 grid for 200 steps for different values of neighbourhood radius and virulence, find the average(over 5 cycle-lenghts) infected fraction then plot a scatter plot to see the distribution.
"""

# ╔═╡ f7bf5830-6865-11eb-29d7-2be3d482ab20
function ss(grid, step; kwargs...)
	if count(i -> infectious(i, grid), state(grid)) == 0
		return :shortcurcuit
	end
end

# ╔═╡ 2e5c48b0-6859-11eb-2b69-0b78043bfc13
function equilibrium_infection(nhradius, p_max)
	arr = rand(1:9, 100, 100)
	grid = LifeGrid(arr, 4, 5, max=p_max, nh=VonNeumannNeighborhood(nhradius, 2))
	results = simulate!(grid, 200, postrunhook=ss)
	mean(count.(p -> infectious(p, grid), results[end - 45:end])) / 10000
end

# ╔═╡ 16c3dea0-685b-11eb-1b89-1da8ccf00941
md" ### Verification
We verify that the function is working correctly by checking for the case `nhradius = 3, p_max = 3` against the run [100 → 2; 300 → 3](#52123eb0-684d-11eb-0179-895599ae7a2f)"

# ╔═╡ ba6cd7a0-685b-11eb-3375-1b00d96dc710
begin
	res = equilibrium_infection(3, 9)
	if 0.35 < res < 0.4
		Markdown.parse("""
!!! tips "Success!"
    \$0.35< \\texttt{equilibrium\\_infection(3, 9)} = $(round(res, sigdigits=2)) <0.4\$
""")
	else
		Markdown.parse(md"""
!!! warning "Oh No!"
    \$0.35 \\nless \\texttt{equilibrium\\_infection(3, 9)} = $(round(res, sigdigits=2)) \\nless 0.4\$
    **This should never HAPPEN!!!**
""")
	end
end

# ╔═╡ 977a4062-6861-11eb-17eb-03874dfb3c79
md"""
!!! danger "Long running Code!"
    The following cell takes a long time (≈ 400 secs) to run. Be careful!
    You may want to reduce `nhs` and `maxs`, or skip some elements.
"""

# ╔═╡ 5b1e6e50-6877-11eb-1ca4-3df17b85575e
md"## Results"

# ╔═╡ 104a9570-685e-11eb-06c2-f9df4e63d1ba
begin
	nhs = 1:4
	maxs = 1:2:25
	points = [equilibrium_infection(nhradius, p_max) for nhradius in nhs, p_max in maxs]
end;

# ╔═╡ c6754200-685e-11eb-319f-f1d06650a620

# ╔═╡ 3f6db120-685e-11eb-3fa6-fd5ec5dfe981
let
	plotlyjs()
	p = plot(maxs, nhs, points, st=:surface, ylabel="NH Radius", xlabel="Virulence(max)", zlabel="Frac infected", title="Surface Plot")
	gr()
	p
end

# ╔═╡ 8af3d922-686e-11eb-03bb-d96c3ac24780
plot(maxs, nhs, points, ylabel="NH Radius", xlabel="Virulence(max)", title="Contour Plot", width=3)

# ╔═╡ 97bff6c2-6873-11eb-2f3d-c108ce4accde
begin
	plotlyjs()
	p = plot(maxs, nhs, points, ylabel="NH Radius", xlabel="Virulence(max)", title="HeatMap", st=:heatmap, hover = string.(reshape(points, :)))
	gr()
	p
end

# ╔═╡ 44881ce0-6877-11eb-0cef-f35b60baf357
md"
## Observations
1. As expected, runs with low virulence and low radius (bottom-right) have very low equilibrium infection levels.
2. For very virulent pathogens, even very strong lockdowns do not have much of an effect.
3. Even for mid-level virulence pathogens, weak lockdowns are only able to reduce the infection probability, but are unable to completely kill off the pathogen.
4. As we have seen before, even very low initial viral load can quickly explode in high interaction neighborhoods, so when the lockdown is removed, the infection quickly increases to pre-lockdown levels

## Discussion on the effectiveness of lockdowns

Given that lockdowns don't help in killing off the pathogen, was the whole last year a waste? Was it all just a government hoax to install 5G towers?

While the author can't say much about government plans to install 5G towers, it is clear that lockdowns have a dampening effect in the number of active cases. This decrease in cases allow the society to focus on combating the disease by spending more efforts into research for vaccines, treatments and preventive steps, along with building infrastructure and materials which are necessary to combat the recurrence of the pathogen. It also gives epidemologists and people working in Complex Systems time to make pretty plots and animations.
"

# ╔═╡ c7ea4fa0-687b-11eb-0dbc-7f316bd4aefd
md"# Further scope of this work
While this work has largely focused on the case of pathogen spread, the same work can be extended within the same topic and to other related topics.

## Things within the same topic
1. Fitting a mathematical model of the effect of lockdown and pathogen spread to data obtained by this work, then fitting it for other real life pathogens
2. Verification of model reflecting real pathogens via case studies
3. More extended statistical analysis to quantify equilibrium
4. Search for multi-equilibrium states
5. Contact tracing and quarantine simulations and comparing effectiveness of blanket lockdowns with more surgical methods.

## Other topics which can gain from this work
While we have called agents as people and the spread as a pathogen, any system which relaxes after some time can be modelled with a similar setup. Accordingly the meaning of lockdown can change.
Some examples include -
1. Spread of radioactivity with lockdown as shielding.
2. Complex web of processors which heat up after a while. There, the lockdown would be the tradeoff between a less complex web but a less efficient system of workload balancing.
3. Spread of algae and lockdown being the speed of flow of water.
"

# ╔═╡ 01866aa0-6854-11eb-3369-7deb6c18c3fe
md"""# References
[^Moitra2019]: Moitra, P., Sinha, S. Localized spatial distributions of disease phases yield long-term persistence of infection. Sci Rep 9, 20309 (2019). [https://doi.org/10.1038/s41598-019-56616-3](https://doi.org/10.1038/s41598-019-56616-3)
[^julia]: Julia: A Fresh Approach to Numerical Computing. Jeff Bezanson, Alan Edelman, Stefan Karpinski, Viral B. Shah. (2017) SIAM Review, 59: 65–98. doi: [10.1137/141000671](https://dx.doi.org/10.1137/141000671)
[^plots]: Plots - powerful convenience for visualization in Julia. Thomas Breloff. [docs.juliaplots.org](docs.juliaplots.org)
[^images]: JuliaImages: image processing and machine vision for Julia. [https://juliaimages.org/stable/](https://juliaimages.org/stable/)
[^pluto]: Pluto - Reactive Notebooks for Julia. Fons van der Plas and Mikołaj Bochenski. [https://github.com/fonsp/Pluto.jl](https://github.com/fonsp/Pluto.jl)
[^plutoui]: PlutoUI. Fons van der Plas. [https://github.com/fonsp/PlutoUI.jl](https://github.com/fonsp/PlutoUI.jl)
[^stats]: Statistics - Julia stdlib for Statistics. [https://docs.julialang.org/en/v1/stdlib/Statistics/](https://docs.julialang.org/en/v1/stdlib/Statistics/)
[^cellularbase]: CellularBase.jl: Abstract definitions for Cellular Automaton. Dhruva Sambrani. [https://dhruvasambrani.github.io/idc621/CellularBase](https://dhruvasambrani.github.io/idc621/CellularBase)
"""

# ╔═╡ 263db440-66f9-11eb-03b5-1fce784fad2d
function (_g::LifeGrid)(neighbors::Array{Int}; kwargs...)
	current = neighbors[1]
	if susceptible(current, _g)
		if infect(count(i -> infectious(i, _g), neighbors), _g)
			return 1
		else
			return 0
		end
	else
		return (current + 1) % (_g.τi + _g.τr)
	end
end

# ╔═╡ 35d0be40-66f7-11eb-2f60-01982265ccb0
mutable struct LifeGrid <: AbstractGrid{Int}
	state::Array{Int}
	otherstate::Array{Int}
	bc::BoundaryCondition
	neighborhood::AbstractNeighborhood
	τi::Int64
	τr::Int64
	max::Int64
	function LifeGrid(
			size::Tuple{Int,Int}, τi::Int, τr::Int;
			max::Int64, init=nothing,
			nh::AbstractNeighborhood=VonNeumannNeighborhood(1, 2))
		z = zeros(size)
		if isnothing(init) z[(size .÷ 2)...] = 1
		else z[rand(CartesianIndices(z), init)] .= 1
		end
		new(z, copy(z), FixedMax, nh, τi, τr, max)
	end
	function LifeGrid(
			arr::Array{Int}, τi::Int, τr::Int;
			max::Int64,
			nh::AbstractNeighborhood=VonNeumannNeighborhood(1, 2)
		)
		new(arr, copy(arr), FixedMax, nh, τi, τr, max)
	end
end

# ╔═╡ 21d1eb20-66f8-11eb-26d3-4fb244ca8646
begin
	infect(_count::Int64, g::LifeGrid) = if (_count == 0) false else rand() < 1 / g.max * _count end
	susceptible(i::Int, g::LifeGrid) = i == 0
	infectious(i::Int, g::LifeGrid) = i ∈ 1:g.τi
	refactory(i::Int, g::LifeGrid) = i ∈ (g.τi + 1):(g.τr + g.τi)
	recovered(i::Int, g::LifeGrid) = i == (g.τr + g.τi) + 1
end

# ╔═╡ Cell order:
# ╟─9a74ef90-67ad-11eb-1fc3-e7080fbfc7e1
# ╟─eb2aad32-687f-11eb-0c6b-91b443cdb3d3
# ╟─6f7bce82-67ad-11eb-1ec1-79ce2620fe9b
# ╟─3459c040-6858-11eb-09c1-95c7bdb02722
# ╟─7cb22830-67c9-11eb-0e0e-41f2316ed468
# ╟─65ff1d80-66f9-11eb-11c2-69c47b8458e4
# ╟─330799c0-6884-11eb-230d-71c512e4a447
# ╠═1bc961f0-66f7-11eb-17c8-2b54d4c4328d
# ╟─db4fd9e0-6883-11eb-384e-6df2c1d7a52b
# ╠═f8149ae0-66f6-11eb-0335-39baafdb24c8
# ╟─a95c8c70-67ad-11eb-124b-7b66d5547321
# ╟─73ea0ae0-66f9-11eb-19cf-b7ab6151ef63
# ╟─038cbdf2-6880-11eb-1b2b-2b4c4974be49
# ╠═35d0be40-66f7-11eb-2f60-01982265ccb0
# ╟─91afcec0-66f9-11eb-37b1-f78de5e0ba8a
# ╟─fcc23790-686e-11eb-03d7-1f0946446f49
# ╠═263db440-66f9-11eb-03b5-1fce784fad2d
# ╟─b8e7c920-66f9-11eb-397f-3760662a648c
# ╟─8b6833e0-6870-11eb-2593-69edaec0a7b8
# ╠═21d1eb20-66f8-11eb-26d3-4fb244ca8646
# ╟─de865ed2-66f9-11eb-31dc-c782d4338af3
# ╠═c3e6cc10-66f7-11eb-039f-df82addf1a5e
# ╠═4d4cdcc0-67b0-11eb-2278-8de64128d6f1
# ╠═7cdebc60-684b-11eb-38a2-45873e4b51d0
# ╠═5031733e-66f9-11eb-107f-6d69089d9384
# ╟─c0318a00-67ac-11eb-39b7-e3c13a5aff57
# ╟─47b81340-67cb-11eb-372e-972e4d4ff746
# ╠═02c24a70-67ae-11eb-2944-61a7963eb947
# ╟─fbdba862-67cf-11eb-39bc-f37eb3662785
# ╠═fe0dda70-66fa-11eb-0e02-73c05666aace
# ╠═8315d810-677a-11eb-3953-8faad42ada40
# ╟─7d24e300-67d0-11eb-0864-41c6166f13ac
# ╠═b8523010-67ce-11eb-2396-4780284dfa01
# ╠═cf22a450-67ce-11eb-2141-fd6f3efc689a
# ╟─07c440a0-67d1-11eb-3a8b-d559cac6276e
# ╠═30ce7a10-67bd-11eb-027d-8dc583201bc1
# ╠═4ca7fe50-67bd-11eb-382c-1304f3c1c3c4
# ╟─643ffb90-67ad-11eb-20c6-7744050e5448
# ╟─1cf29850-66fa-11eb-3f76-db755c0a4c50
# ╟─393a5bc0-67cb-11eb-3be1-a9ec7943b5f3
# ╟─2563e210-67d5-11eb-089e-f376da973765
# ╟─dafac300-67e9-11eb-0972-f9070193399e
# ╠═46055e90-67e9-11eb-2b6f-dffa58138005
# ╠═5626db00-67e9-11eb-0e12-a77b6b9c58be
# ╟─d4c2cff0-67e9-11eb-1f5b-03c7bab91782
# ╠═3056cdae-67ec-11eb-0b35-b3d6cde9b0fa
# ╠═3b0024a0-67ec-11eb-3398-eb07250758a9
# ╟─6d822d52-67ed-11eb-1b71-679d491bf133
# ╠═76e43730-67ed-11eb-201f-53bf8d3e3bfb
# ╠═8445d050-67ed-11eb-2411-f776a43c7f40
# ╟─0f5ba100-67ef-11eb-15f5-3fbaad70dde2
# ╟─9e71a7c0-67f1-11eb-0282-672d358d2292
# ╟─ef666b40-67f4-11eb-2f31-05145a6dfcba
# ╟─5586a260-67f4-11eb-11e2-b706fd0d4dfe
# ╟─21b80dc0-67f9-11eb-0e87-6db59c88c296
# ╠═ead37ab2-66f9-11eb-2e56-07f105753991
# ╟─5641a270-684b-11eb-29b4-ef8c59ef9679
# ╠═0f20b710-6840-11eb-1712-9f4614e9ed87
# ╠═6f9dc770-6847-11eb-14ff-2b2e8562076f
# ╠═2bf1c112-6848-11eb-3b4a-a3ce569da696
# ╟─a9588a00-684b-11eb-2585-09a2cdf9ad80
# ╠═4f4fdfd0-684c-11eb-3769-49d1677a1d52
# ╠═d97dddfe-684c-11eb-1d11-5734830507a2
# ╠═e1baa1c0-684c-11eb-292c-3336d6d715b1
# ╟─52123eb0-684d-11eb-0179-895599ae7a2f
# ╠═4e182db0-684d-11eb-0bab-b7c22d3f3ff5
# ╠═53bb3e50-684e-11eb-1bd3-d582200be8b5
# ╠═660d5660-684e-11eb-27ad-cbde3f616410
# ╟─f0c45ab2-684e-11eb-24ca-c5942b62f396
# ╟─61e28250-6857-11eb-39f9-2ff5565519ab
# ╟─2734ec90-6859-11eb-2430-e3bd97302fd8
# ╠═f7bf5830-6865-11eb-29d7-2be3d482ab20
# ╠═2e5c48b0-6859-11eb-2b69-0b78043bfc13
# ╟─16c3dea0-685b-11eb-1b89-1da8ccf00941
# ╟─ba6cd7a0-685b-11eb-3375-1b00d96dc710
# ╟─977a4062-6861-11eb-17eb-03874dfb3c79
# ╟─5b1e6e50-6877-11eb-1ca4-3df17b85575e
# ╠═104a9570-685e-11eb-06c2-f9df4e63d1ba
# ╟─c6754200-685e-11eb-319f-f1d06650a620
# ╠═3f6db120-685e-11eb-3fa6-fd5ec5dfe981
# ╠═8af3d922-686e-11eb-03bb-d96c3ac24780
# ╠═97bff6c2-6873-11eb-2f3d-c108ce4accde
# ╟─44881ce0-6877-11eb-0cef-f35b60baf357
# ╟─c7ea4fa0-687b-11eb-0dbc-7f316bd4aefd
# ╟─01866aa0-6854-11eb-3369-7deb6c18c3fe


"""
Tableau of one-stage, 1st order explicit (forward) Euler method

```julia
TableauExplicitEuler(::Type{T}=Float64) where {T}
```

The constructor takes one optional argument, that is the element type of the tableau.

Reference:

    Leonhard Euler.
    Institutiones calculi differentialis cum eius vsu in analysi finitorum ac doctrina serierum.
    Imp. Acad. Imper. Scient. Petropolitanae, Opera Omnia, Vol.X, [I.6], 1755.
    In: Opera Omnia, 1st Series, Volume 11, Institutiones Calculi Integralis. Teubner, Leipzig, Pages 424-434, 1913.
    Sectio secunda. Caput VII. De integratione aequationum differentialium per approximationem. Problema 85.

"""
function TableauExplicitEuler(::Type{T}=Float64) where {T}
    a = zeros(BigFloat, 1, 1)
    b = ones(BigFloat, 1)
    c = zeros(BigFloat, 1)
    o = 1

    Tableau{T}(:explicit_euler, o, a, b, c)
end

"Alias for [`TableauExplicitEuler`](@ref)"
TableauForwardEuler = TableauExplicitEuler

"""
Tableau of explicit two-stage, 2nd order midpoint method

```julia
TableauExplicitMidpoint(::Type{T}=Float64) where {T}
```

The constructor takes one optional argument, that is the element type of the tableau.

Reference:

    Carl Runge.
    Über die numerische Auflösung von Differentialgleichungen.
    Mathematische Annalen, Volume 46, Pages 167-178, 1895.
    doi: 10.1007/BF01446807.
    Equation (2)

"""
function TableauExplicitMidpoint(::Type{T}=Float64) where {T}
    a = @big [[ 0     0    ]
              [ 1//2  0    ]]
    b = @big  [ 0,    1    ]
    c = @big  [ 0,    1//2 ]
    o = 2

    Tableau{T}(:explicit_midpoint, o, a, b, c)
end

"""
Tableau of Heun's two-stage, 2nd order method

```julia
TableauHeun2(::Type{T}=Float64) where {T}
```

The constructor takes one optional argument, that is the element type of the tableau.

Reference:

    Karl Heun.
    Neue Methoden zur approximativen Integration der Differentialgleichungen einer unabhängigen Veränderlichen.
    Zeitschrift für Mathematik und Physik, Volume 45, Pages 23-38, 1900.
    Algorithm II.

"""
function TableauHeun2(::Type{T}=Float64) where {T}
    a = @big [[ 0      0    ]
              [ 1      0    ]]
    b = @big  [ 1//2,  1//2 ]
    c = @big  [ 0,     1    ]
    o = 2

    Tableau{T}(:heun2, o, a, b, c)
end

"""
Tableau of Heun's three-stage, 3rd order method

```julia
TableauHeun3(::Type{T}=Float64) where {T}
```

The constructor takes one optional argument, that is the element type of the tableau.

Reference:

    Karl Heun.
    Neue Methoden zur approximativen Integration der Differentialgleichungen einer unabhängigen Veränderlichen.
    Zeitschrift für Mathematik und Physik, Volume 45, Pages 23-38, 1900.
    Algorithm VI.

"""
function TableauHeun3(::Type{T}=Float64) where {T}
    a = @big [[ 0      0      0    ]
              [ 1//3   0      0    ]
              [ 0      2//3   0    ]]
    b = @big  [ 1//4,  0,     3//4 ]
    c = @big  [ 0,     1//3,  2//3 ]
    o = 3

    Tableau{T}(:heun3, o, a, b, c)
end

"""
Tableau of Ralston's two-stage, 2nd order method

```julia
TableauRalston2(::Type{T}=Float64) where {T}
```

The constructor takes one optional argument, that is the element type of the tableau.

Reference:

    Anthony Ralston.
    Runge-Kutta Methods with Minimum Error Bounds.
    Mathematics of Computation, Volume 16, Pages 431-437, 1962.
    doi: 10.1090/S0025-5718-1962-0150954-0.
    Equation (3.5)

"""
function TableauRalston2(::Type{T}=Float64) where {T}
    a = @big [[ 0      0    ]
              [ 2//3   0    ]]
    b = @big  [ 1//4,  3//4 ]
    c = @big  [ 0,     2//3 ]
    o = 2

    Tableau{T}(:ralston2, o, a, b, c)
end

"""
Tableau of Ralston's three-stage, 3rd order method

```julia
TableauRalston3(::Type{T}=Float64) where {T}
```

The constructor takes one optional argument, that is the element type of the tableau.

Reference:

    Anthony Ralston.
    Runge-Kutta Methods with Minimum Error Bounds.
    Mathematics of Computation, Volume 16, Pages 431-437, 1962.
    doi: 10.1090/S0025-5718-1962-0150954-0.
    Equation (4.10)

"""
function TableauRalston3(::Type{T}=Float64) where {T}
    a = @big [[ 0      0      0    ]
              [ 1//2   0      0    ]
              [ 0      3//4   0    ]]
    b = @big  [ 2//9,  3//9,  4//9 ]
    c = @big  [ 0,     1//2,  3//4 ]
    o = 3

    Tableau{T}(:ralston3, o, a, b, c)
end

"""
Tableau of Runge's two-stage, 2nd order method

```julia
TableauRunge(::Type{T}=Float64) where {T}
```

The constructor takes one optional argument, that is the element type of the tableau.

Reference:

    Carl Runge
    Über die numerische Auflösung von Differentialgleichungen.
    Mathematische Annalen, Volume 46, Pages 167-178, 1895.
    doi: 10.1007/BF01446807.
    Equation (3)

"""
function TableauRunge(::Type{T}=Float64) where {T}
    a = @big [[ 0      0    ]
              [ 1      0    ]]
    b = @big  [ 1//2,  1//2 ]
    c = @big  [ 0,     1    ]
    o = 2

    Tableau{T}(:runge, o, a, b, c)
end

"Alias for [`TableauRunge`](@ref)"
TableauRunge2 = TableauRunge

"""
Tableau of Kutta's three-stage, 3rd order method

```julia
TableauKutta(::Type{T}=Float64) where {T}
```

The constructor takes one optional argument, that is the element type of the tableau.

Reference:

    Wilhelm Kutta
    Beitrag zur Näherungsweisen Integration totaler Differentialgleichungen
    Zeitschrift für Mathematik und Physik, Volume 46, Pages 435–453, 1901.
    Page 440

"""
function TableauKutta(::Type{T}=Float64) where {T}
    a = @big [[ 0      0      0    ]
              [ 1//2   0      0    ]
              [-1      2      0    ]]
    b = @big  [ 1//6,  4//6,  1//6 ]
    c = @big  [ 0,     1//2,  1    ]
    o = 3

    Tableau{T}(:kutta, o, a, b, c)
end

"Alias for [`TableauKutta`](@ref)"
TableauKutta3 = TableauKutta

"""
Tableau of explicit Runge-Kutta method of order four (1/6 rule)

```julia
TableauRK416(::Type{T}=Float64) where {T}
```

The constructor takes one optional argument, that is the element type of the tableau.

Reference:

    Wilhelm Kutta
    Beitrag zur Näherungsweisen Integration totaler Differentialgleichungen
    Zeitschrift für Mathematik und Physik, Volume 46, Pages 435–453, 1901.
    Page 443

"""
function TableauRK416(::Type{T}=Float64) where {T}
    a = @big [[ 0      0      0      0    ]
              [ 1//2   0      0      0    ]
              [ 0      1//2   0      0    ]
              [ 0      0      1      0    ]]
    b = @big  [ 1//6,  1//3,  1//3,  1//6 ]
    c = @big  [ 0,     1//2,  1//2,  1    ]
    o = 4

    Tableau{T}(:erk4, o, a, b, c)
end

"Alias for [`TableauRK416`](@ref)"
TableauRK4 = TableauRK416

"""
Tableau of explicit Runge-Kutta method of order four (3/8 rule)

```julia
TableauRK438(::Type{T}=Float64) where {T}
```

The constructor takes one optional argument, that is the element type of the tableau.

Reference:

    Wilhelm Kutta
    Beitrag zur Näherungsweisen Integration totaler Differentialgleichungen
    Zeitschrift für Mathematik und Physik, Volume 46, Pages 435–453, 1901.
    Page 441

"""
function TableauRK438(::Type{T}=Float64) where {T}
    a = @big [[ 0      0      0      0    ]
              [ 1//3   0      0      0    ]
              [-1//3   1      0      0    ]
              [ 1     -1      1      0    ]]
    b = @big  [ 1//8,  3//8,  3//8,  1//8 ]
    c = @big  [ 0,     1//3,  2//3,  1    ]
    o = 4

    Tableau{T}(:erk438, o, a, b, c)
end

"""
Tableau of 2rd order Strong Stability Preserving method with two stages and CFL ≤ 1

Alias for [`TableauHeun2`](@ref)

```julia
TableauSSPRK2(::Type{T}=Float64) where {T}
```

The constructor takes one optional argument, that is the element type of the tableau.

Reference:

    Chi-Wang Shu, Stanley Osher.
    Efficient implementation of essentially non-oscillatory shock-capturing schemes.
    Journal of Computational Physics, Volume 77, Issue 2, Pages 439-471, 1988.
    doi: 10.1016/0021-9991(88)90177-5.
    Equation (2.16)

"""
TableauSSPRK2 = TableauHeun2

"""
Tableau of 3rd order Strong Stability Preserving method with three stages and CFL ≤ 1

```julia
TableauSSPRK3(::Type{T}=Float64) where {T}
```

The constructor takes one optional argument, that is the element type of the tableau.

Reference:

    Chi-Wang Shu, Stanley Osher.
    Efficient implementation of essentially non-oscillatory shock-capturing schemes.
    Journal of Computational Physics, Volume 77, Issue 2, Pages 439-471, 1988.
    doi: 10.1016/0021-9991(88)90177-5.
    Equation (2.18)

"""
function TableauSSPRK3(::Type{T}=Float64) where {T}
    a = @big [[ 0      0      0    ]
              [ 1      0      0    ]
              [ 1//4   1//4   0    ]]
    b = @big  [ 1//6,  1//6,  4//6 ]
    c = @big  [ 0,     1,     1//2 ]
    o = 3

    Tableau{T}(:ssprk3, o, a, b, c)
end


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

using Markdown
using InteractiveUtils

# ╔═╡ a450c12c-753d-11eb-06fc-29c5d907cb06
begin
	using Pkg
	Pkg.activate("MLJ_env", shared=true)

	# using MLJ
	using CSV
	using DataFrames
	using PlutoUI
	using Random
	using Test
	using Printf
	# using Plots
end

# ╔═╡ a4ce9686-7549-11eb-2172-2d47391d1260
include("./utils.jl")

# ╔═╡ 5b0c59ae-753d-11eb-1a8a-6fa0bf99a357
md"""
## Week 5: Lasso (Coordinate Descent)

In this notebook, you will implement your very own LASSO solver via coordinate descent. We will:

  - Write a function to normalize features
  - Implement coordinate descent for Lasso
  - Explore effects of L1 penalty

"""

# ╔═╡ 7e12df34-7549-11eb-30a5-b3099a8f77f6
md"""
### Load in house sales data
"""

# ╔═╡ 9e5e1060-7549-11eb-0418-55bcd0b8c757
sales = CSV.File("../../ML_UW_Spec/C02/data/kc_house_test_data.csv";
	header=true) |> DataFrame;

# ╔═╡ 2fa16676-754a-11eb-00cc-896d0e28b246
md"""
### Normalize features

In the house dataset, features vary wildly in their relative magnitude: sqft_living is very large overall compared to bedrooms, for instance. As a result, weight for sqft_living would be much smaller than weight for bedrooms. This is problematic because "small" weights are dropped first as l1_penalty goes up.

To give equal considerations for all features, we need to normalize features as discussed in the lectures: we divide each feature by its 2-norm so that the transformed feature has norm 1.
"""

# ╔═╡ 4e5d02be-754a-11eb-061f-41ab2f747c54
md"""
moved to `utils.jl`

```julia
using LinearAlgebra

function normalize_features(X::Matrix)
  n = size(X)[2]
  norms = zeros(eltype(X), n)'
  for ix ∈ 1:n
    norms[ix] = norm(X[:, ix])
  end
  (X ./ norms, norms)
end
```
"""

# ╔═╡ dc99317e-754a-11eb-0a8e-61ec5dff0507
begin
	M1 = [3. 6. 9.; 4. 8. 12.]
	features1, norms1 = normalize_features(M1)

	@test features1 == [0.6 0.6 0.6; 0.8 0.8 0.8]
	@test norms1 == [5.  10.  15.]  ## or [5.,  10.,  15.]
end

# ╔═╡ e9128af8-7555-11eb-1f4d-d1e17d70f4b3
typeof(features1), typeof(norms1)

# ╔═╡ 0d7f1fa0-754c-11eb-330b-d31a86f5566c
begin
	# 3 rows x 2 cols
	M2 = [3. 4.; 6. 5.; 9. 10.]
	features2, norms2 = normalize_features(M2)

	@test features2 ≈ [0.2672612419124244 0.3368607684266076;
		0.5345224838248488 0.42107596053325946;
		0.8017837257372732 0.8421519210665189]
	@test norms2 ≈ [11.224972160321824 11.874342087037917]
end

# ╔═╡ 984c778a-7557-11eb-21f8-dd2502fe7b73
md"""
### Implementing Coordinate Descent with normalized features

We seek to obtain a sparse set of weights by minimizing the LASSO cost function

$$\Sigma (prediction - output)^2 + \lambda \times \Sigma_{j=1}^{k} | w_j |$$

(By convention, we do not include $w_0$ in the L1 penalty term. We never want to push the intercept to zero.)

The absolute value sign makes the cost function non-differentiable, so simple gradient descent is not viable (you would need to implement a method called subgradient descent). Instead, we will use **coordinate descent**: at each iteration, we will fix all weights but weight `i` and find the value of weight `i` that minimizes the objective. That is, we look for

$$argmin_{w_i}[\Sigma (prediction - output)^2 + \lambda \times \Sigma_{j=1}^{k} | w_j |]$$

where all weights other than $w_i$ are held to be constant. We will optimize one $w_i$ at a time, circling through the weights multiple times.
  1. Pick a coordinate `i`
  2. Compute $w_i$ that minimizes the cost function $\Sigma[ (prediction - output)^2 ] + \lambda \times \Sigma_{j=1}^{k} | w_j |$
  3. Repeat Steps 1 and 2 for all coordinates, multiple times

For this notebook, we use **cyclical coordinate descent with normalized features**, where we cycle through coordinates 0 to (d-1) in order, and assume the features were normalized as discussed above. The formula for optimizing each coordinate is as follows:

$$w_i = \left\{
    \begin{array}{lr}
      \rho_i + \frac{\lambda}{2} & if \rho_i < -\frac{\lambda}{2} \\
       0  &                        if -\frac{\lambda}{2} \le \rho_i \le \frac{\lambda}{2} \\
      \rho_i - \frac{\lambda}{2} & if \rho_i > \frac{\lambda}{2} \\
    \end{array}
\right.$$

where $\rho_i = \Sigma [ [feature_i] \times (output - prediction + w\_i \times [feature_i]) ]$

Note that we do not regularize the weight of the constant feature (intercept) `w[0]`, so, for this weight, the update is simply:
$w_0 = \rho_i$

"""

# ╔═╡ eb96bfc8-7556-11eb-2045-0f0e7adabe76
md"""
#### Effect of L1 penalty
"""

# ╔═╡ 9c8d9610-754f-11eb-05c3-edd271661b84
begin
	simple_features = [:sqft_living, :bedrooms]
	output_ = :price
	(simple_feature_matrix, output) = get_data(sales, simple_features, output_);
end

# ╔═╡ eb872178-754f-11eb-177d-7bb5ee1046f5
begin
	simple_feature_matrix_n, norms_n = normalize_features(simple_feature_matrix);
end

# ╔═╡ 20799d0c-7550-11eb-16ee-676be88a80af
typeof(simple_feature_matrix), typeof(simple_feature_matrix_n)

# ╔═╡ 06be0f04-7557-11eb-3f7c-bfc5e5a172d4
md"""
We assign some random set of initial weights and inspect the values of ρ[i].

we use `predict_output()` to make predictions on this data.
"""

# ╔═╡ 06a0f408-7557-11eb-182e-8312809145d1
begin
	weights₀ = [1., 4., 1.]

	prediction = predict_output(simple_feature_matrix_n, weights₀);
	prediction, size(prediction)
end

# ╔═╡ 06814074-7557-11eb-2359-8fa03900e3f8
md"""
Compute the values of `ρ[i]` for each feature in this simple model, using the formula given above, using the formula:

$\rho_i = \Sigma [ [feature_i] \times (output - prediction + w\_i \times [feature_i]) ]$

"""

# ╔═╡ fcc38dfc-7557-11eb-048c-93e95ce9bbbc
simple_feature_matrix_n[:, 2]

# ╔═╡ 04b4ddf8-7557-11eb-0d69-f7a89e8936f8
ρ = [
	sum(simple_feature_matrix_n[:, i] .*
		(output - prediction + weights₀[i] * simple_feature_matrix[:, i]))
	for i in 1:length(weights₀)
]

# ╔═╡ 0498db3c-7557-11eb-1af6-9de231a4256f
md"""
**Quiz Question**

Recall that, whenever `ρ[i]` falls between `-l1_penalty/2` and `l1_penalty/2`, the corresponding weight `w[i]` is sent to zero.

Now suppose we were to take one step of coordinate descent on either feature 1 or feature 2.

What range of values of `l1_penalty` **would not** set `w[1]` to zero, but **would** set `w[2]` to zero, if we were to take a step in that coordinate?

"""

# ╔═╡ 0480646e-7557-11eb-3aa9-af64a4502ff9
md"""
we know that by definition:

by definition:
  - for w[2]: $|\rho_2| \le \frac{l_1}{2} \equiv  l_1 \ge 2 * \rho_2$
  - for w[1]: $|\rho_1| \le \frac{l_1}{2} \equiv  l_1 \ge 2 * \rho_1$

"""

# ╔═╡ 046a8ad4-7557-11eb-2aa4-e37c7f37f711
## l1 range that would set w[2] to zero (but not w[1])
min(ρ[2] * 2, ρ[1] * 2)

# ╔═╡ 044fc410-7557-11eb-3a80-eb215ebc3f5f
md"""
**Quiz Question**

What range of values of l1_penalty would set both w[1] and w[2] to zero, if we were to take a step in that coordinate?
"""

# ╔═╡ 0434c2be-7557-11eb-0592-1b733cae76ca
## Any value strictly greater than
max(ρ[2] * 2, ρ[1] * 2)

# ╔═╡ 041c6f52-7557-11eb-050e-c97dee4f3220
md"""
So we can say that `ρ[i]` quantifies the significance of the i-th feature: the larger `ρ[i]` is, the more likely it is for the i-th feature to be retained.
"""

# ╔═╡ 4f9ee984-755a-11eb-06cb-8da12ef68907
md"""
### Single Coordinate Descent Step

Using the formula above, implement coordinate descent that minimizes the cost function over a single feature i. Note that the intercept (weight 0) is not regularized. The function should accept feature matrix, output, current weights, l1 penalty, and index of feature to optimize over. The function should return new weight for feature i.
"""

# ╔═╡ 53782228-755a-11eb-1554-394795823745
function lasso_coordinate_descent_step(ix, f_matrix, output, weights, l₁)
	pred = predict_output(f_matrix, weights)

	ρ_ix = sum(f_matrix[:, ix] .* (output - pred + weights[ix] * f_matrix[:, ix]))

	new_weight_ix = if ix == 1   ## intercept -- do not regularize
		 ρ_ix
	elseif ρ_ix < -l₁ / 2.
		ρ_ix + l₁ / 2.
	elseif ρ_ix >  l₁ / 2.
		ρ_ix - l₁ / 2.
	else
		0.
	end

	new_weight_ix
end

# ╔═╡ 535bf7b0-755a-11eb-06ea-8b0e0baec03c
begin
	ϵ = 1e-7
	exp = -1.7706583789422325
	act = lasso_coordinate_descent_step(
		1,
		[3. / √(13)  1. / √(10); 2. / √(13)  3. / √(10)],
		[1., 1.],
		[1., 4.],
	0.1)

	@test abs(act - exp) ≤ ϵ
end

# ╔═╡ 0403eec8-7557-11eb-18f7-13b0403e0a5b
md"""
#### Cyclical coordinate descent

Now that we have a function that optimizes the cost function over a single coordinate, let us implement cyclical coordinate descent where we optimize coordinates 0, 1, ..., (d-1) in order and repeat.

When do we know to stop? Each time we scan all the coordinates (features) once, we measure the change in weight for each coordinate. If no coordinate changes by more than a specified threshold, we stop.

For each iteration:
1. As you loop over features in order and perform coordinate descent, measure how much each coordinate changes.
2. After the loop, if the maximum change across all coordinates falls below the tolerance, stop. Otherwise, go back to step 1.

Return weights

*IMPORTANT: when computing a new weight for coordinate i, make sure to incorporate the new weights for coordinates 0, 1, ..., i-1. One good way is to update your weights variable in-place.*

"""

# ╔═╡ 03ec653c-7557-11eb-2240-8b060cbb17af
function lasso_cyclical_coordinate_descent(f_matrix, output, init_weights,
		l1_penalty, tolerance)

	weights = copy(init_weights)
	while true
		coord = []

		for ix in 1:length(weights)
			prev_weight_ix = weights[ix]
			weights[ix] = lasso_coordinate_descent_step(ix, f_matrix, output,
				weights, l1_penalty)
			push!(coord, abs(prev_weight_ix - weights[ix]))
		end

		maximum(coord) < tolerance && break
	end
	weights
end

# ╔═╡ dcf8bdc2-755b-11eb-04da-39f89dd5709c
begin
	## utility functions, for the coming questions...

	function features_nzw(features, weights)
  		"""features with non-zero weigths"""
  		return filter(
			((_n, w)=t) -> w > 0.,
			collect(zip([:intercept, features...], weights))
		)

	end

	function features_zw(features, weights, ϵ=1e-7)
  		"""features with zero weigths"""
  		return filter(
			((_n, w)=t) -> w ≤ ϵ,
			collect(zip([:intercept, features...], weights))
			# intercept is constant in the quiz
		)
	end
end

# ╔═╡ 03d065a8-7557-11eb-1cf0-add64e6d7596
begin
	## Using the following parameters, learn the weights on the sales dataset.
	s_features = [:sqft_living, :bedrooms]
	s_outputn = :price
	s_init_weights = zeros(Float64, length(s_features) + 1)  # +1 for intercept
	s_l1_penalty = 1e7
	s_tolerance = 1.0

	s_f_matrix, s_output = get_data(sales, s_features, s_outputn)
	(norm_s_f_matrix, s_norms) = normalize_features(s_f_matrix)

	s_weights = lasso_cyclical_coordinate_descent(norm_s_f_matrix, s_output,
		s_init_weights, s_l1_penalty, s_tolerance)

	s_weights, size(s_weights)
end

# ╔═╡ cf49001e-755e-11eb-33d9-9182f04696f8
md"""
**Quiz Questions**

  1. What is the RSS of the learned model on the normalized dataset? (Hint: use the normalized feature matrix when you make predictions.)
  1. Which features had weight zero at convergence?
"""

# ╔═╡ d35515f6-755e-11eb-3da0-0bc64c7ed3b1
with_terminal() do
	## 1. RSS
	rss = calc_rss(norm_s_f_matrix, s_output, s_weights)
	@printf("rss: %.3f / in scientifc notation: %1.5e\n", rss, rss)
end

# ╔═╡ acd4c536-755e-11eb-1b0c-9573130238dc
## 2. features with weight zero
features_zw(s_features, s_weights)

# ╔═╡ a941a03a-755f-11eb-38f7-455aaa090d7f
md"""
### Evaluating LASSO fit with more features
"""

# ╔═╡ c26ef2fe-755f-11eb-28e4-ef66e958b894
begin
	train_data, test_data = train_test_split(sales; split=0.8, seed=42)
	size(train_data), size(test_data)
end

# ╔═╡ c251d9bc-755f-11eb-2fe2-13c496a0abbd
begin
	a_features = [
		:bedrooms, :bathrooms, :sqft_living, :sqft_lot, :floors,
		:waterfront, :view, :condition, :grade, :sqft_above, :sqft_basement,
		:yr_built, :yr_renovated
	]

	tr_outputn = :price
	(tr_f_matrix, tr_output) = get_data(train_data, a_features, tr_outputn)
	(norm_tr_f_matrix, norms_tr) = normalize_features(tr_f_matrix)

	size(norm_tr_f_matrix), size(norms_tr)
end

# ╔═╡ c0ff7904-755f-11eb-2ad7-572acb587b6f
md"""
##### Learn with l₁ of 1e7
"""

# ╔═╡ c0e5f506-755f-11eb-1ad1-6906e5ccfd85
begin
	l1₁=1e7
	init_weights = zeros(length(a_features) + 1)
	tol = 1.0

	weights1e7 = lasso_cyclical_coordinate_descent(norm_tr_f_matrix, tr_output,
		init_weights, l1₁, tol)
	(size(weights1e7), typeof(weights1e7))
end

# ╔═╡ c0ca4482-755f-11eb-156b-a7011072410a
md"""
**Quiz Question**

What features had non-zero weight in this case?

```julia
features_nzw(a_features, weights1e7)
```
"""

# ╔═╡ 230c74ce-7562-11eb-12d7-43e526f2ec26
with_terminal() do
	for t ∈ features_nzw(a_features, weights1e7)
		@printf("%-20s => %2.5e / %15.5f\n", string(t[1]), t[2], t[2])
	end
end

# ╔═╡ c0b05b74-755f-11eb-319e-9f5eff250eeb
md"""
##### Learn with l₁ of 1e8
"""

# ╔═╡ c095ab50-755f-11eb-2347-79fca41a8555
begin
	l1₂=1e8
	# init_weights2 = zeros(length(a_features) + 1)
	# tol₂ = 1.0

	weights1e8 = lasso_cyclical_coordinate_descent(norm_tr_f_matrix, tr_output,
		init_weights, l1₂, tol)
	(size(weights1e8), typeof(weights1e8))
end

# ╔═╡ c07b70a0-755f-11eb-30ed-cba73d960ff5
md"""
**Quiz Question**

What features had non-zero weight in this case?

```julia
features_nzw(a_features, weights1e8)
```
"""

# ╔═╡ 81964a24-7562-11eb-3ff3-4370b4056f74
with_terminal() do
	for t ∈ features_nzw(a_features, weights1e8)
		@printf("%-20s => %2.5e / %15.5f\n", string(t[1]), t[2], t[2])
	end
end

# ╔═╡ c05c69c6-755f-11eb-0442-ef31019d90ff
md"""
##### Learn with l₁ of 1e4
"""

# ╔═╡ 98a45748-7561-11eb-154a-35b7dd014d41
begin
	l1₃=1e4
	# init_weights₃ = zeros(length(a_features) + 1)
	# tol = 1.0

	weights1e4 = lasso_cyclical_coordinate_descent(norm_tr_f_matrix, tr_output,
		init_weights, l1₃, tol)
	(size(weights1e4), typeof(weights1e4))
end

# ╔═╡ 9856bf2e-7561-11eb-0887-ff4f836e557e
md"""
**Quiz Question**

What features had non-zero weight in this case?

```julia
features_nzw(a_features, weights1e4)
```
"""

# ╔═╡ 08a08158-7563-11eb-0141-15cbef1ee9a6
with_terminal() do
	for t ∈ features_nzw(a_features, weights1e4)
		@printf("%-20s => %2.5e / %15.5f\n", string(t[1]), t[2], t[2])
	end
end

# ╔═╡ a7b5fba4-7563-11eb-2a24-f34b6296b447
md"""
#### Rescaling learned weights

Recall that we normalized our feature matrix, before learning the weights.  To use these weights on a test set, we must normalize the test data in the same way.

Alternatively, we can rescale the learned weights to include the normalization, so we never have to worry about normalizing the test data:

In this case, we must scale the resulting weights so that we can make predictions with *original* features:
 1. Store the norms of the original features to a vector called `norms`:
```
features, norms = normalize_features(features)
```
 2. Run Lasso on the normalized features and obtain a `weights` vector
 3. Compute the weights for the original features by performing element-wise division, i.e.
```
weights_normalized = weights / norms
```
Now, we can apply `weights_normalized` to the test data, without normalizing it!
"""

# ╔═╡ c088e8d0-7563-11eb-0d5c-79ed7f100599
## Create a normalized version of each of the weights learned above. (weights1e4, weights1e7, weights1e8).

begin
all_weights = [
		w ./ vec(norms_tr) for w ∈ (weights1e4, weights1e7, weights1e8)
];

size(all_weights)
end

# ╔═╡ c048f716-7563-11eb-214f-35c773b47c1d
md"""
#### Evaluating each of the learned models on the test data

Let's now evaluate the three models on the test data
"""

# ╔═╡ c02e5654-7563-11eb-390b-ef4352bcf1fa
(test_f_matrix, test_output) = get_data(test_data, a_features, :price);

# ╔═╡ 61567f18-7564-11eb-2301-9dbec94cf157
md"""
Compute the RSS of each of the three normalized weights on the (unnormalized) `test_feature_matrix`
"""

# ╔═╡ 9298e410-7566-11eb-0b86-5376fcf9eabc
md"""
**Quiz Question**

Which model performed best on the test data?
"""

# ╔═╡ 611cc566-7564-11eb-168f-052aa0209973
with_terminal() do
	best_rss, best_weights = (nothing, nothing)
	w_names = (:weights1e4, :weights1e7, :weights1e8)

	for (weights, label) ∈ zip(all_weights, w_names)
  		rss = calc_rss(test_f_matrix, test_output, weights)
  		@printf("for %-10s => rss: %1.5e\n", label, rss)

  		if isnothing(best_rss) || rss < best_rss
    		best_rss, best_weights = rss, label
		end
	end

	@printf("\nbest model/weights: for %-10s => rss: %1.5e\n", best_weights, best_rss)
end

# ╔═╡ 6e03f5ea-7564-11eb-3309-9f6131c552cf

# ╔═╡ Cell order:
# ╟─5b0c59ae-753d-11eb-1a8a-6fa0bf99a357
# ╠═a450c12c-753d-11eb-06fc-29c5d907cb06
# ╟─7e12df34-7549-11eb-30a5-b3099a8f77f6
# ╠═9e5e1060-7549-11eb-0418-55bcd0b8c757
# ╠═a4ce9686-7549-11eb-2172-2d47391d1260
# ╟─2fa16676-754a-11eb-00cc-896d0e28b246
# ╟─4e5d02be-754a-11eb-061f-41ab2f747c54
# ╠═dc99317e-754a-11eb-0a8e-61ec5dff0507
# ╠═e9128af8-7555-11eb-1f4d-d1e17d70f4b3
# ╠═0d7f1fa0-754c-11eb-330b-d31a86f5566c
# ╟─984c778a-7557-11eb-21f8-dd2502fe7b73
# ╟─eb96bfc8-7556-11eb-2045-0f0e7adabe76
# ╠═9c8d9610-754f-11eb-05c3-edd271661b84
# ╠═eb872178-754f-11eb-177d-7bb5ee1046f5
# ╠═20799d0c-7550-11eb-16ee-676be88a80af
# ╟─06be0f04-7557-11eb-3f7c-bfc5e5a172d4
# ╠═06a0f408-7557-11eb-182e-8312809145d1
# ╟─06814074-7557-11eb-2359-8fa03900e3f8
# ╠═fcc38dfc-7557-11eb-048c-93e95ce9bbbc
# ╠═04b4ddf8-7557-11eb-0d69-f7a89e8936f8
# ╟─0498db3c-7557-11eb-1af6-9de231a4256f
# ╟─0480646e-7557-11eb-3aa9-af64a4502ff9
# ╠═046a8ad4-7557-11eb-2aa4-e37c7f37f711
# ╟─044fc410-7557-11eb-3a80-eb215ebc3f5f
# ╠═0434c2be-7557-11eb-0592-1b733cae76ca
# ╟─041c6f52-7557-11eb-050e-c97dee4f3220
# ╟─4f9ee984-755a-11eb-06cb-8da12ef68907
# ╠═53782228-755a-11eb-1554-394795823745
# ╠═535bf7b0-755a-11eb-06ea-8b0e0baec03c
# ╟─0403eec8-7557-11eb-18f7-13b0403e0a5b
# ╠═03ec653c-7557-11eb-2240-8b060cbb17af
# ╠═dcf8bdc2-755b-11eb-04da-39f89dd5709c
# ╠═03d065a8-7557-11eb-1cf0-add64e6d7596
# ╟─cf49001e-755e-11eb-33d9-9182f04696f8
# ╠═d35515f6-755e-11eb-3da0-0bc64c7ed3b1
# ╠═acd4c536-755e-11eb-1b0c-9573130238dc
# ╟─a941a03a-755f-11eb-38f7-455aaa090d7f
# ╠═c26ef2fe-755f-11eb-28e4-ef66e958b894
# ╠═c251d9bc-755f-11eb-2fe2-13c496a0abbd
# ╟─c0ff7904-755f-11eb-2ad7-572acb587b6f
# ╠═c0e5f506-755f-11eb-1ad1-6906e5ccfd85
# ╟─c0ca4482-755f-11eb-156b-a7011072410a
# ╠═230c74ce-7562-11eb-12d7-43e526f2ec26
# ╟─c0b05b74-755f-11eb-319e-9f5eff250eeb
# ╠═c095ab50-755f-11eb-2347-79fca41a8555
# ╟─c07b70a0-755f-11eb-30ed-cba73d960ff5
# ╠═81964a24-7562-11eb-3ff3-4370b4056f74
# ╟─c05c69c6-755f-11eb-0442-ef31019d90ff
# ╠═98a45748-7561-11eb-154a-35b7dd014d41
# ╟─9856bf2e-7561-11eb-0887-ff4f836e557e
# ╠═08a08158-7563-11eb-0141-15cbef1ee9a6
# ╟─a7b5fba4-7563-11eb-2a24-f34b6296b447
# ╠═c088e8d0-7563-11eb-0d5c-79ed7f100599
# ╟─c048f716-7563-11eb-214f-35c773b47c1d
# ╠═c02e5654-7563-11eb-390b-ef4352bcf1fa
# ╟─61567f18-7564-11eb-2301-9dbec94cf157
# ╟─9298e410-7566-11eb-0b86-5376fcf9eabc
# ╠═611cc566-7564-11eb-168f-052aa0209973
# ╠═6e03f5ea-7564-11eb-3309-9f6131c552cf


@testset verbose=true "map2salm" begin
    # NOTE: Float16 irfft returns Float32, which leads to type conflicts, so we just don't support
    # map2salm on Float16

    function sYlm(s::Int, ell::Int, m::Int, theta::T, phi::T) where {T<:Real}
        # Eqs. (II.7) and (II.8) of https://arxiv.org/abs/0709.0093v3
        # Note their weird definition w.r.t. `-s`
        k_min = max(0, m + s)
        k_max = min(ell + m, ell + s)
        sin_half_theta, cos_half_theta = sincos(theta / 2)
        return (-1)^(-s) * sqrt((2 * ell + 1) / (4 * T(π))) *
            T(sum(
                (-1) ^ (k)
                * sqrt(factorial(big(ell + m)) * factorial(big(ell - m)) * factorial(big(ell - s)) * factorial(big(ell + s)))
                * (cos_half_theta ^ (2 * ell + m + s - 2 * k))
                * (sin_half_theta ^ (2 * k - s - m))
                / (factorial(big(ell + m - k)) * factorial(big(ell + s - k)) * factorial(big(k)) * factorial(big(k - s - m)))
                for k in k_min:k_max
            )) *
            cis(m * phi)
    end

    m_m2Y2m = [
        (2, NINJA.m2Y22),
        (1, NINJA.m2Y21),
        (0, NINJA.m2Y20),
        (-1, NINJA.m2Y2m1),
        (-2, NINJA.m2Y2m2)
    ]

    @testset "Input expressions $T" for T in [BigFloat, Float64, Float32]
        # These are just internal consistency tests of the sYlm function
        # above, against the explicit expressions `mY2.`
        s = -2
        ℓ = 2
        Nϑ = 17
        Nφ = 18
        for (m, m2Y2m) in m_m2Y2m
            f1 = mapslices(ϕθ -> sYlm(s, ℓ, m, ϕθ[2], ϕθ[1]), phi_theta(Nφ, Nϑ, T), dims=[3])
            f2 = mapslices(ϕθ -> m2Y2m(ϕθ[2], ϕθ[1]), phi_theta(Nφ, Nϑ, T), dims=[3])
            @test f1 ≈ f2 atol=10eps(T) rtol=10eps(T)
        end
    end

    @testset "map2salm $T" for T in [BigFloat, Float64, Float32]
        # These test the ability of map2salm to precisely decompose the results of `sYlm`.
        ℓmax = 7
        Nϑ = 2ℓmax + 1
        Nφ = 2ℓmax + 2
        for s in -2:2
            #for ℓmin in 0:abs(s)
            let ℓmin = 0
                for ℓ in abs(s):ℓmax
                    for m in -ℓ:ℓ
                        f = mapslices(
                            ϕθ -> sYlm(s, ℓ, m, ϕθ[2], ϕθ[1]),
                            phi_theta(Nφ, Nϑ, T),
                            dims=[3]
                        )
                        computed = map2salm(f, s, ℓmax)
                        expected = zeros(Complex{T}, size(computed))
                        expected[SphericalFunctions.Yindex(ℓ, m, ℓmin)] = one(T)
                        if ≉(computed, expected, atol=30eps(T), rtol=30eps(T))
                            @show T
                            @show ℓmax
                            @show Nϑ
                            @show Nφ
                            @show s
                            @show ℓmin
                            @show ℓ
                            @show m
                            println("computed = $computed")
                            println("expected = $expected")
                            println("max_diff = ", maximum(abs, computed .- expected), ";")
                            println()
                        end
                        @test computed ≈ expected atol=30eps(T) rtol=30eps(T)

                        plan = plan_map2salm(f, s, ℓmax)
                        computed2 = map2salm(f, plan)
                        @test array_equal(computed, computed2)
                    end
                end
            end
        end
    end

end


# Import libraries.
using Turing, Turing: Variational
using Distributions
using AbstractGPs, KernelFunctions
import Random
import LinearAlgebra

# NOTE: Import other libraries ...

# NOTE: Read data ...

# Define a kernel.
sqexpkernel(alpha::Real, rho::Real) =
    alpha^2 * transform(SqExponentialKernel(), 1/(rho*sqrt(2)))

# Define model.
@model GPRegression(y, X) = begin
    # Priors.
    alpha ~ LogNormal(0.0, 0.1)
    rho ~ LogNormal(0.0, 1.0)
    sigma ~ LogNormal(0.0, 1.0)

    # Covariance function.
    kernel = sqexpkernel(alpha, rho)

    # GP (implicit zero-mean).
    gp = GP(kernel)

    # Sampling Distribution (MvNormal likelihood).
    y ~ gp(X, sigma^2 + 1e-6)  # add 1e-6 for numerical stability.
end;

# Set random number generator seed.
Random.seed!(0)

# Model creation.
m = GPRegression(y, X)

# Fit via ADVI.
q0 = Variational.meanfield(m)  # initialize variational distribution (optional)
num_elbo_samples, max_iters = (1, 2000)
# Run optimizer.
@time q = vi(m, ADVI(num_elbo_samples, max_iters), q0,
             optimizer=Flux.ADAM(1e-1));

# Fit via HMC.
burn = 1000
nsamples = 1000
@time chain = sample(m, HMC(0.01, 100), burn + nsamples)  # start sampling.

# Fit via NUTS.
@time chain = begin
    nsamples = 1000  # number of MCMC samples
    nadapt = 1000  # number of iterations to adapt tuning parameters in NUTS
    iterations = nsamples + nadapt
    target_accept_ratio = 0.8

    sample(m, NUTS(nadapt, target_accept_ratio, max_depth=10), iterations);
end


using SparseArrays

n = 2 .^ [6 6 7 8 9 10 11 12 13 14 15 16 17 18]

K = 8

MaxGB = 2
MaxGF = 4 * 1.6

m = K .* n
sparse_gbytes = 9 * 8 * m ./ 1e9
sparse_flops = zeros(1,length(n))
sparse_gflops = zeros(1,length(n))
sparse_time = zeros(1,length(n))

for i = 1:length(n)

    global sparse_gflops, sparse_time

    ii = round.(Int,floor.(rand(m[i]) .* n[i]) .+ 1)
    jj = round.(Int,floor.(rand(m[i]) .* n[i]) .+ 1)
    A  = sparse(ii , jj , 1.0, n[i], n[i])
    ii = round.(Int,floor.(rand(m[i]) .* n[i]) .+ 1)
    jj = round.(Int,floor.(rand(m[i]) .* n[i]) .+ 1)
    B  = sparse(ii,jj,1.0,n[i],n[i])

    sparse_time[i] = @elapsed C = A*B

    sparse_flops[i] = 2*sum(C)
    sparse_gflops[i] = sparse_flops[i] ./ sparse_time[i] ./ 1e9

    print("Time: ", string(sparse_time[i]))
    print(", GFlops: ", string(sparse_gflops[i]))
    println(", GBytes: ", string(sparse_gbytes[i]))

end


@testset "Sigma points" begin
    @testset "Weighting parameters" begin
        num_states = 5

        weight_params = WanMerweWeightingParameters(0.5, 2, 0)
        @test @inferred(KalmanFilters.lambda(weight_params, num_states)) == -3.75
        @test @inferred(KalmanFilters.calc_mean_weights(weight_params, num_states)) == (-3, 0.4)
        @test @inferred(KalmanFilters.calc_cov_weights(weight_params, num_states)) == (-0.25, 0.4)
        @test @inferred(KalmanFilters.calc_cholesky_weight(weight_params, num_states)) == 1.25

        weight_params = MeanSetWeightingParameters(0.5)
        @test @inferred(KalmanFilters.calc_mean_weights(weight_params, num_states)) == (0.5, 0.05)
        @test @inferred(KalmanFilters.calc_cov_weights(weight_params, num_states)) == (0.5, 0.05)
        @test @inferred(KalmanFilters.calc_cholesky_weight(weight_params, num_states)) == 10

        weight_params = GaussSetWeightingParameters(3)
        @test all(@inferred(KalmanFilters.calc_mean_weights(weight_params, num_states)) .≈ (-2/3, 1/6))
        @test all(@inferred(KalmanFilters.calc_cov_weights(weight_params, num_states)) .≈ (-2/3, 1/6))
        @test @inferred(KalmanFilters.calc_cholesky_weight(weight_params, num_states)) == 3

        weight_params = ScaledSetWeightingParameters(0.5, 2, 1)
        @test @inferred(KalmanFilters.calc_mean_weights(weight_params, num_states)) == (-19, 2)
        @test @inferred(KalmanFilters.calc_cov_weights(weight_params, num_states)) == (-16.25, 2)
        @test @inferred(KalmanFilters.calc_cholesky_weight(weight_params, num_states)) == 0.25
    end

    @testset "Create sigma points" begin
        weight_params = ScaledSetWeightingParameters(0.5, 2, 1)
        A = [4.0 0.0; 0.0 4.0]
        x = randn(2)
        χ = @inferred KalmanFilters.calc_sigma_points(x, A, weight_params)
        @test χ.x0 == x
        @test χ.P_chol == cholesky(A .* 0.25).L

        P_chol_temp = zeros(2,2)
        χ = @inferred KalmanFilters.calc_sigma_points!(P_chol_temp, x, A, weight_params)
        @test χ.x0 == x
        @test χ.P_chol == cholesky(A .* 0.25).L
    end

    @testset "Weighted means" begin
        weight_params = ScaledSetWeightingParameters(0.5, 2, 1)
        χ = KalmanFilters.TransformedSigmaPoints(ones(5), hcat(ones(5,5) .* 4, ones(5,5) .* 2), weight_params)
        x = @inferred KalmanFilters.mean(χ)
        @test x == ones(5) .* -19 .+ ones(5) .* 20 .* 2 .+ ones(5) .* 10 .* 2

        χ = KalmanFilters.TransformedSigmaPoints(ones(5), hcat(ones(5,5) .* 4, ones(5,5) .* 2), weight_params)
        x = zeros(5)
        x = @inferred KalmanFilters.mean!(x, χ)
        @test x == ones(5) .* -19 .+ ones(5) .* 20 .* 2 .+ ones(5) .* 10 .* 2
    end

    @testset "substract mean" begin
        weight_params = ScaledSetWeightingParameters(0.5, 2, 1)
        χ = KalmanFilters.TransformedSigmaPoints(ones(5), hcat(ones(5,5) .* 4, ones(5,5) .* 2), weight_params)
        x = ones(5) .* 2
        unbiased_χ = @inferred KalmanFilters.substract_mean(χ, x)
        @test unbiased_χ == [-ones(5) hcat(ones(5,5) .* 2, zeros(5,5))]
        @test unbiased_χ.weight_params == weight_params

        other_weight_params = MeanSetWeightingParameters(0.0)
        unbiased_χ_temp = KalmanFilters.TransformedSigmaPoints(randn(5), randn(5,10), other_weight_params)
        unbiased_χ = @inferred KalmanFilters.substract_mean!(unbiased_χ_temp, χ, x)
        @test unbiased_χ == [-ones(5) hcat(ones(5,5) .* 2, zeros(5,5))]
        @test unbiased_χ.weight_params == weight_params
    end

    @testset "Covariance" begin
        weight_params = ScaledSetWeightingParameters(0.5, 2, 1)
        𝓨 = KalmanFilters.TransformedSigmaPoints(ones(5), hcat(ones(5,5) .* 4, ones(5,5) .* 2), weight_params)
        noise = Diagonal(ones(5))
        P = @inferred KalmanFilters.cov(𝓨, noise)
        @test P == -16.25 .* ones(5) * ones(5)' .+ 2 .* ones(5,5) * ones(5,5)' .* 16 .+
            2 .* ones(5,5) * ones(5,5)' .* 4 .+ noise

        P = zeros(5,5)
        P = @inferred KalmanFilters.cov!(P, 𝓨, noise)
        @test P == -16.25 .* ones(5) * ones(5)' .+ 2 .* ones(5,5) * ones(5,5)' .* 16 .+
            2 .* ones(5,5) * ones(5,5)' .* 4 .+ noise

        P_chol = LowerTriangular(ones(5,5))
        x = ones(5) .* 4
        χ = KalmanFilters.SigmaPoints(x, P_chol, weight_params)
        P = @inferred KalmanFilters.cov(χ, 𝓨)
        @test P == 2 .* LowerTriangular(ones(5,5)) * ones(5,5)' .* 4 .+
            2 .* LowerTriangular(-ones(5,5)) * ones(5,5)' .* 2

        P = zeros(5,5)
        P = @inferred KalmanFilters.cov!(P, χ, 𝓨)
        @test P == 2 .* LowerTriangular(ones(5,5)) * ones(5,5)' .* 4 .+
            2 .* LowerTriangular(-ones(5,5)) * ones(5,5)' .* 2
    end

    @testset "Transform sigma points" begin
        P = [2.0 0.0; 0.0 2.0]
        x = [1.0, 1.0]
        weight_params = ScaledSetWeightingParameters(0.5, 2, 1)
        χ = KalmanFilters.SigmaPoints(x, LowerTriangular(P), weight_params)
        F(x) = x .* 2
        𝓨 = @inferred KalmanFilters.transform(F, χ)
        @test 𝓨 == [ones(2) .* 2 [6 2; 2 6] [-2 2; 2 -2]]
        @test 𝓨.weight_params == weight_params

        other_weight_params = MeanSetWeightingParameters(0.5)
        𝓨_temp = KalmanFilters.TransformedSigmaPoints(zeros(2), zeros(2,4), other_weight_params)
        F!(y, x) = y .= x .* 2
        xi_temp = zeros(length(x))
        𝓨 = @inferred KalmanFilters.transform!(𝓨_temp, xi_temp, F!, χ)
        @test 𝓨 == [ones(2) .* 2 [6 2; 2 6] [-2 2; 2 -2]]
        @test 𝓨.weight_params == weight_params
    end
end


struct Dielectric <: Material
    refractive_index::AbstractFloat
end

function schlick(cosine::AbstractFloat, refractive_index::AbstractFloat)
    r0::AbstractFloat = ((1.0 - refractive_index) / (1.0 + refractive_index))^2
    r0 + (1.0 - r0) * (1.0 - cosine)^5
end

function scatter(dielectric::Dielectric, ray::Ray, hp::HitPoint)::Tuple{Bool, Union{Nothing, AbstractVec}, Union{Nothing, Ray}}
    attenuation::AbstractVec = Vec(1.0, 1.0, 1.0)
    local scattered::Ray

    succeeded::Bool, refracted::Union{Nothing, AbstractVec} = refract(ray.direction, hp.normal, 1.0 / dielectric.refractive_index)
    if succeeded
        cosine::AbstractFloat = -dot(ray.direction, hp.normal) / norm(ray.direction)
        reflect_prob::AbstractFloat = schlick(cosine, refractive_index)
        if rand() >= reflect_prob
            scattered = Ray(hp.point, refracted)
            return true, attenuation, scattered
        end
    end

    reflected::AbstractVec = reflect(ray.direction, hp.normal)
    scattered = Ray(hp.point, reflected)
    true, attenuation, scattered
end