module SStools

import Distributions.MvNormal

export kalman_filter, fast_state_smoother, simulate

"""
Perform Kalman filtering on the data y.
Conventions are as in Durbin and Koopman (2012).
Relevant dimensions are:
- Nt: number of time points
- Np: dimension of observation space
- Nm: dimension of state space
- Nr: dimension of state noise covariance

Parameters:
- a0: mean of prior on states: (Nm, 1) or (Nm,)
- P0: variance of prior on states: (Nm, Nm)
- Z: either a (Np, Nm, Nt) array or a (Np, Nm) array
- H: either a (Np, Np, Nt) array or a (Np, Np) array
- T: either a (Nm, Nm, Nt) array or a (Nm, Nm) array
- R: either a (Nm, Nr, Nt) array or a (Nm, Nr) array
- Q: either a (Nr, Nr, Nt) array or a (Nr, Nr) array

Returns:
- v: (Np, Nt) vector of residuals at each time
- K: (Nm, Np, Nt) Kalman gain matrix
- Finv: (Np, Np, Nt) inverse of prediction variance matrix
"""
function kalman_filter(y, a0, P0, Z, H, T, R, Q)
    Np, Nt = size(y)
    Nm = size(P0, 1)

    # preallocate arrays
    a = Array(Float64, Nm, Nt)
    P = Array(Float64, Nm, Nm, Nt)
    v = Array(Float64, Np, Nt)
    Finv = Array(Float64, Np, Np, Nt)
    K = Array(Float64, Nm, Np, Nt)

    # initialize
    a[:, 1] = a0;
    P[:, :, 1] = P0;

    # iterate
    for t in 1:Nt
        local a_t = slice(a, :, t)
        local P_t = slice(P, :, :, t)
        local y_t = slice(y, :, t)
        local Z_t = ndims(Z) < 3 ? Z : slice(Z, :, :, t)
        local H_t = ndims(H) < 3 ? H : slice(H, :, :, t)
        local T_t = ndims(T) < 3 ? T : slice(T, :, :, t)
        local R_t = ndims(R) < 3 ? R : slice(R, :, :, t)
        local Q_t = ndims(Q) < 3 ? Q : slice(Q, :, :, t)

        v_t = y_t - Z_t * a_t
        F_t = Z_t * P_t * Z_t' + H_t
        Finv_t = inv(F_t)
        K_t = T_t * P_t * Z_t' * Finv_t
        L_t = T_t - K_t * Z_t

        if t < Nt
            a[:, t + 1] = T_t * a_t + K_t * v_t
            P[:, :, t + 1] = T_t * P_t * L_t' + R_t * Q_t * R_t'
        end

        v[:, t] = v_t
        K[:, :, t] = K_t
        Finv[:, :, t] = Finv_t
    end

    return v, K, Finv, a, P
end

"""
Runs an efficient smoothing algorithm to get posterior mean of state
vector.

Parameters are as in kalman_filter.
"""
function fast_state_smoother(v, K, Finv, a0, P0, Z, T, R, Q)
    # infer dimensions
    Np, Nt = size(v)
    Nm = size(P0, 1)

    # preallocate
    r0 = 0.
    r = Array(Float64, Nm, Nt)
    r[:, end] = 0
    alpha = Array(Float64, Nm, Nt);

    # iterate backward
    for t in Nt:-1:1
        local Finv_t = slice(Finv, :, :, t)
        local v_t = slice(v, :, t)
        local K_t = slice(K, :, :, t)
        local r_t = slice(r, :, t)
        local Z_t = ndims(Z) < 3 ? Z : slice(Z, :, :, t)
        local T_t = ndims(T) < 3 ? T : slice(T, :, :, t)

        local u = Finv_t * v_t - K_t' * r_t
        local thisr = Z_t' * u + T_t' * r_t

        if t - 1 > 0
            r[:, t - 1] = thisr
        else
            r0 = thisr
        end
    end

    # run model forward
    alpha[:, 1] = a0 + P0 * r0

    for t in 1:(Nt - 1)
        local T_t = ndims(T) < 3 ? T : slice(T, :, :, t)
        local R_t = ndims(R) < 3 ? R : slice(R,  :, :, t)
        local Q_t = ndims(Q) < 3 ? Q : slice(Q,  :, :, t)
        local RQR = R_t * Q_t * R_t'

        alpha[:, t + 1] = T_t * alpha[:, t] + RQR * r[:, t]
    end

    return alpha
end

"""
Draw a sample of the state trajectory from the smoothed posterior.
"""
function simulate(y, a0, P0, Z, H, T, R, Q; interleaved=true)
    # get dimensions
    Np, Nt = size(y)
    Nm = size(a0, 1)
    Nr = size(Q, 1)

    # preallocate
    alpha_plus = Array(Float64, Nm, Nt)
    y_plus = Array(Float64, Np, Nt)

    # simulate data
    alpha_plus[:, 1] = rand(MvNormal(a0, P0))

    for t in 1:Nt
        local Z_t = ndims(Z) < 3 ? Z : slice(Z, :, :, t)
        local T_t = ndims(T) < 3 ? T : slice(T, :, :, t)
        local R_t = ndims(R) < 3 ? R : slice(R, :, :, t)
        # random number generators should take copies, not views?
        local H_t = ndims(H) < 3 ? H : H[:, :, t]
        local Q_t = ndims(Q) < 3 ? Q : Q[:, :, t]

        # draw disturbances
        local ϵ = rand(MvNormal(full(H_t)))
        y_plus[:, t] = Z_t * alpha_plus[:, t] + ϵ

        if t < Nt
            η = rand(MvNormal(full(Q_t)))
            alpha_plus[:, t + 1] = T_t * alpha_plus[:, t] + R_t * η
        end
    end

    # calculate smoothed means:
    if interleaved
        # actual data
        v, K, Finv, a, P = interleaved_kalman_filter(y, a0, P0, Z, H, T, R, Q)
        alpha_hat = interleaved_state_smoother(v, K, Finv, a0, P0, Z, T, R, Q)

        # simulated data
        v, K, Finv, a, P = interleaved_kalman_filter(y_plus, a0, P0, Z, H, T, R, Q)
        alpha_hat_plus = interleaved_state_smoother(v, K, Finv, a0, P0, Z, T, R, Q)

    else
        # actual data
        v, K, Finv, a, P = kalman_filter(y, a0, P0, Z, H, T, R, Q)
        alpha_hat = fast_state_smoother(v, K, Finv, a0, P0, Z, T, R, Q)

        # simulated data
        v, K, Finv, a, P = kalman_filter(y_plus, a0, P0, Z, H, T, R, Q)
        alpha_hat_plus = fast_state_smoother(v, K, Finv, a0, P0, Z, T, R, Q)
    end

    # combine
    alpha_draw = alpha_plus - alpha_hat_plus + alpha_hat

    return alpha_draw
end

"""
Like Kalman filter, but works by transforming a vector-valued
observation into a sequence of scalar observations.
"""
function interleaved_kalman_filter(y, a0, P0, Z, H, T, R, Q)
    # test for diagonality of H
    # H (minus its diagonals) should be the 0 matrix
    if ndims(H) == 2
        assert(H == diagm(diag(H)))
    end

    Np, Nt = size(y)
    Nm = size(P0, 1)

    # preallocate
    a = Array(Float64, Nm, Np + 1, Nt)
    P = Array(Float64, Nm, Nm, Np + 1, Nt)
    v = Array(Float64, Np, Nt)
    Finv = Array(Float64, Np, Nt)
    K = Array(Float64, Nm, Np, Nt)

    # initialize
    a[:, 1, 1] = a0
    P[:, :, 1, 1] = P0

    for t in 1:Nt, i in 1:Np
        local Z_t = ndims(Z) < 3 ? Z : slice(Z, :, :, t)
        local H_t = ndims(H) < 3 ? H : slice(H, :, :, t)
        local T_t = ndims(T) < 3 ? T : slice(T, :, :, t)
        local R_t = ndims(R) < 3 ? R : slice(R, :, :, t)
        local Q_t = ndims(Q) < 3 ? Q : slice(Q, :, :, t)
        z = vec(full(Z_t[i, :]))  # now a column vector

        v[i, t] = y[i, t] - dot(z, a[:, i, t])
        F = dot(z, P[:, :, i, t] * z) + H_t[i, i]
        Finv[i, t] = 1. / F
        if F != 0
            K[:, i, t] = P[:, :, i, t] * z * Finv[i, t]
        else
            K[:, i, t] = 0
        end

        a[:, i + 1, t] = a[:, i, t] + K[:, i, t] * v[i, t]
        P[:, :, i + 1, t] = P[:, :, i, t] - F * (K[:, i, t] * K[:, i, t]')

        if i == Np && t < Nt
            a[:, 1, t + 1] = T_t * a[:, Np + 1, t]
            P[:, :, 1, t + 1] = T_t * P[:, :, Np + 1, t] * T_t' + R_t * Q_t * R_t'
        end
    end

    return v, K, Finv, a, P
end

function interleaved_state_smoother(v, K, Finv, a0, P0, Z, T, R, Q)
    # infer dimensions
    Np, Nt = size(v)
    Nm = size(P0, 1)

    # preallocate
    r = Array(Float64, Nm, Np + 1, Nt)

    # initialize
    r[:, end, end] = 0

    # iterate backward
    for t in Nt:-1:1, i in (Np + 1):-1:2
        ii = i - 1  # handles offset between r and Z/K/F indices

        local Z_t = ndims(Z) < 3 ? Z : slice(Z, :, :, t)
        L = eye(Nm) - K[:, ii, t] * Z_t[ii, :]
        z = vec(full(Z_t[ii, :]))  # now a column vector

        r[:, i - 1, t] = z * (v[ii, t] * Finv[ii, t]) + L' * r[:, i, t]

        if i == 2 && t > 1
            # NOTE: t - 1 BELOW !!!
            local T_t = ndims(T) < 3 ? T : slice(T, :, :, t - 1)
            r[:, Np + 1, t - 1] = T_t' * r[:, 1, t]
        end
    end

    # initialize again
    α = Array(Float64, Nm, Nt)
    α[:, 1] = a0 + P0 * r[:, 1, 1]

    for t in 1:(Nt - 1)
        local T_t = ndims(T) < 3 ? T : slice(T, :, :, t)
        local R_t = ndims(R) < 3 ? R : slice(R, :, :, t)
        local Q_t = ndims(Q) < 3 ? Q : slice(Q, :, :, t)
        local RQR = R_t * Q_t * R_t'

        α[:, t + 1] = T_t * α[:, t] + RQR * r[:, 1, t + 1]
    end

    return α
end

end  # module


import Distributions: logpdf, pdf

struct SDT{T1,T2} <: ContinuousUnivariateDistribution
    d::T1
    c::T2
end

logpdf(d::SDT,data::Vector{Int64}) = logpdf(d,data...)

logpdf(d::SDT,data::Tuple{Vararg{Int64}}) = logpdf(d,data...)

function logpdf(d::SDT,hits,fas,Nd)
    @unpack d,c=d
    θhit=cdf(Normal(0,1),d/2-c)
    θfa=cdf(Normal(0,1),-d/2-c)
    loghits = logpdf(Binomial(Nd,θhit),hits)
    logfas = logpdf(Binomial(Nd,θfa),fas)
    return loghits+logfas
end

pdf(d::SDT,data::Vector{Int64}) = exp(logpdf(d,data...))


using Revise
using AdFem
using LinearAlgebra
using PyPlot
using DelimitedFiles

function calc_residual_and_jacobian(θ, u)
    ψ, J = neo_hookean(u, μ, λ, mmesh)
    B = compute_fem_source_term(zeros(get_ngauss(mmesh)), -0.5*ones(get_ngauss(mmesh)), mmesh)
    ψ = ψ - B
    J, ψ = impose_Dirichlet_boundary_conditions(J, ψ, bdnode, zeros(length(bdnode)))
    ψ, J
end

mmesh = Mesh(8, 8 , 1/8)

E, nu = 10.0, 0.3
μ, λ = E/(2*(1 + nu)) * ones(get_ngauss(mmesh)), E*nu/((1 + nu)*(1 - 2*nu)) * ones(get_ngauss(mmesh))
bdnode = bcnode((x,y)->(x<1e-5) || (x>1-1e-5), mmesh)
bdnode = [bdnode; bdnode .+ mmesh.ndof]
nr = newton_raphson(calc_residual_and_jacobian, zeros(2mmesh.ndof), missing)

ADCME.options.newton_raphson.verbose = true
sess = Session(); init(sess)
nr = run(sess, nr)

close("all")
# visualize_displacement(nr.x, mmesh)
u = nr.x[1:mmesh.nnode]
v = nr.x[mmesh.ndof+1:mmesh.ndof + mmesh.nnode]
visualize_scalar_on_fem_points(sqrt.(u.^2+v.^2), mmesh)
savefig("test_AdFem.png")


include("utils.jl")

struct CST{T<:AbstractFloat, N} <: AbstractAirfoil{T}
    A₀::T
    Aᵤ::Vector{T}
    Aₗ::Vector{T}
    dz::T
    C₁::T
    C₂::T
end

function CST(A::AbstractVector{T}; dz::T=0, c1::T=0.5, c2::T=1.0) where {T<:AbstractFloat}
    n, r = divrem(length(A), 2)
    if (r!=1)
        throw(ArgumentError("length of parameters must be `2N-1`"))
    end
    CST(A[1], A[2:n+1], A[n+2:end]; dz=dz, c1=c1, c2=c2)
end

function CST(A₀::T, Aᵤ::AbstractVector{T}, Aₗ::AbstractVector{T};
             dz::T=0, c1::T=0.5, c2::T=1.0) where {T<:AbstractFloat}
    if length(Aₗ) != length(Aᵤ)
        throw(ArgumentError("length(Aᵤ) != length(Aₗ)"))
    end
    n = length(Aᵤ)
    CST{T,n}(A₀, copy(Aᵤ), copy(Aₗ), dz, c1, c2)
end

function CST(n::Int,
             xᵤ::AbstractVector{T}, yᵤ::AbstractVector{T},
             xₗ::AbstractVector{T}, yₗ::AbstractVector{T},
             opt::OptimizerFactory;
             c1::T=0.5, c2::T=1.0) where {T<:AbstractFloat}
    if length(xᵤ)!=length(yᵤ)
        throw(ArgumentError("imbalanced upper-surface data: length(xᵤ)!=length(yᵤ)"))
    end
    if length(xₗ)!=length(yₗ)
        throw(ArgumentError("imbalanced lower-surface data: length(xₗ)!=length(yₗ)"))
    end
    if !(xᵤ[1] == xₗ[1] == 0)
        throw(ArgumentError("`xᵤ[1]` and `xₗ[1]` must be 0"))
    end
    if !(xᵤ[end] == xₗ[end] == 1)
        throw(ArgumentError("`xᵤ[end]` and `xₗ[end]` must be 0"))
    end
    if !isapprox(yᵤ[1], yₗ[1])
        throw(ArgumentError("`yᵤ[1]` must approximate `yₗ[1]`"))
    end
    # move LE to (0,0)
    oy = (yᵤ[1]+yₗ[1] )/2
    yₗ .-= oy
    yᵤ .-= oy
    #
    dz = (yᵤ[end] - yₗ[end])/2
    if dz < -eps(T)
        throw(ArgumentError("`yᵤ[end]` must be greater than `yₗ[end]`"))
    elseif dz < eps(T)
        dz = zero(T) # set to zero
    end
    A₀, Aᵤ, Aₗ, status = fit_cst(n, dz, xᵤ, yᵤ, xₗ, yₗ, opt; c1=c1, c2=c2)
    CST{T,n}(A₀, Aᵤ, Aₗ, dz, c1, c2)
end

function y_upper(cst::CST{T,N}, x::T)::T where{T<:AbstractFloat,N}
    y = cst.A₀*cst_S(N,0,x)
    for i = 1:N
        y += cst.Aᵤ[i]*cst_S(N,i,x)
    end
    y*cst_C(cst.C₁, cst.C₂, x) + cst.dz*x
end

function y_lower(cst::CST{T,N}, x::T)::T where{T<:AbstractFloat,N}
    y = -cst.A₀*cst_S(N,0,x)
    for i = 1:N
        y += cst.Aₗ[i]*cst_S(N,i,x)
    end
    y*cst_C(cst.C₁, cst.C₂, x)-cst.dz*x
end

function dy_upper(cst::CST{T,N}, x::T)::T where{T<:AbstractFloat,N}
    dy = cst.A₀*cst_dS(N,0,x)
    y = cst.A₀*cst_S(N,0,x)
    for i = 1:N
        y += cst.Aᵤ[i]*cst_S(N,i,x)
        dy += cst.Aᵤ[i]*cst_dS(N,i,x)
    end
    y*cst_dC(cst.C₁, cst.C₂, x)+dy*cst_C(cst.C₁, cst.C₂,x)+cst.dz
end

function dy_lower(cst::CST{T,N}, x::T)::T where{T<:AbstractFloat,N}
    dy = -cst.A₀*cst_dS(N,0,x)
    y = -cst.A₀*cst_S(N,0,x)
    for i = 1:N
        dy += cst.Aₗ[i]*cst_dS(N,i,x)
        y += cst.Aₗ[i]*cst_S(N,i,x)
    end
    y*cst_dC(cst.C₁, cst.C₂, x)+dy*cst_C(cst.C₁, cst.C₂,x)-cst.dz
end

function fy_upper(cst::CST{T,N}, x::T)::Tuple{T,T} where{T<:AbstractFloat,N}
    dy = cst.A₀*cst_dS(N,0,x)
    y = cst.A₀*cst_S(N,0,x)
    for i = 1:N
        y += cst.Aᵤ[i]*cst_S(N,i,x)
        dy += cst.Aᵤ[i]*cst_dS(N,i,x)
    end
    C = cst_C(cst.C₁, cst.C₂, x)
    dC = cst_dC(cst.C₁, cst.C₂, x)
    y*C+cst.dz*x, y*dC+dy*C+cst.dz
end

function fy_lower(cst::CST{T,N}, x::T)::Tuple{T,T} where{T<:AbstractFloat,N}
    dy = -cst.A₀*cst_dS(N,0,x)
    y = -cst.A₀*cst_S(N,0,x)
    for i = 1:N
        dy += cst.Aₗ[i]*cst_dS(N,i,x)
        y += cst.Aₗ[i]*cst_S(N,i,x)
    end
    C = cst_C(cst.C₁, cst.C₂, x)
    dC = cst_dC(cst.C₁, cst.C₂, x)
    y*C-cst.dz*x, y*dC+dy*C-cst.dz
end


# This module represents an explicit, direct sparse 2D finite-difference Poisson solver for heterogeneous media,
# i.e. media having spatially varying (space-dependent) medium parameters.
# The following functionality is currently available in this module:
# * Apply operator ``A=∇⋅(σ(x,z)∇)`` on a field ``ψ`` to get ``p``.
# * Apply ``A^{-1}`` in order to solve for ``ψ`` in ``Aψ=p``, given ``p``.
# Current implementation assumes Neumann boundary conditions at all the boundaries.

#-

# As a demo, start with loading some packages.

using SparseArrays
using StatsBase
using LinearAlgebra
using Random
using ProgressMeter
using LinearAlgebra
using Test
using ForwardDiff
using Calculus
#src # include("core.jl")
#src # include("expt.jl")

# From here on, consider the following Poisson experiment:
# ```math
# ∇⋅(σ(x,z)∇) ψ(t) = ∇⋅(Q(x,z)∇) p(t),
# ```
# ```math
# Q = k * Qv / η.
# ```
# Dimensions and spatial grids are allocated as follows.

nx=21
nz=21
nt=4
nznx=nz*nx
mgrid=[range(-div(nz,2), step=1.0, length=nz), range(-div(nx,2), step=1.0, length=nx)]
tgrid=range(0.0,step=0.5, length=nt)

# Now lets allocate the inputs for a toy experiment.

Qv=abs.(randn(nz,nx))
η=abs.(randn(nz,nx))
k=abs.(randn(nz,nx))
σ=abs.(randn(nz,nx))
p=randn(nz,nx,nt)

for it in 1:nt
	for ix in 1:nx
		for iz in 1:nz
			p[iz,ix,it]=inv(abs(iz-11)+1)*inv(abs(ix-11)+1)
		end
	end
end

# These medium parameters are used to generate the *observed* field ``ψ``.

σobs=abs.(randn(nz,nx))
Qobs=abs.(randn(nz,nx))

fill!(σobs, 1.0)
fill!(Qobs, 1.0)
fill!(Qv, 1.0)
fill!(η, 1.0)
fill!(k, 1.0)
σobs[11,11]=20.0
fill!(σ, 1.0)

# Generate the configuration for the Poisson experiment, using the arrays above.

acqgeom=J.Acquisition.Geom_circ(nss=1,nr=30,rad=[5.,5.])
#ACQ=sprandn(200,441,0.9)
ACQ=J.Acquisition.ACQmat(acqgeom,mgrid)
paE=J.Poisson.ParamExpt(p, tgrid, mgrid, Qv, k, η, σ, ACQ, σobs=σobs, Qobs=Qobs,)
#paE=J.Poisson.ParamExpt(p, tgrid, mgrid, Qv, k, η, σ, σobs=σobs, Qobs=Qobs,)

#src # J.Poisson.updateLP!(paE, paE.Q)

# Calculate the data misfit.

@time f=J.Poisson.func(σ,paE)

# Compute the gradient with finite-difference approximation.

f = x->J.Poisson.func(x, paE);
gcalc = Calculus.gradient(f);
g_fd=reshape(gcalc(vec(σ)),nz,nx);

# Compute the gradient using adjoint state method.

@time J.Poisson.mod!(paE,σ,J.Poisson.FGσ())
g=paE.g;

# Check gradient accuracy.

@test ≈(g,g_fd, rtol=1e-5)


filename() = "day12.txt"

function parse_line(line)
    instr = line[1]
    val = parse(Int, line[2:end])
    return (; instr, val)
end

read_data() = parse_line.(readlines(filename()))

function part1()
    pos = (x = 0, y = 0)
    rad = 0
    for d in read_data()
        if d.instr == 'N'
            pos = (x = pos.x, y = pos.y - d.val)
        elseif d.instr == 'S'
            pos = (x = pos.x, y = pos.y + d.val)
        elseif d.instr == 'E'
            pos = (x = pos.x + d.val, y = pos.y)
        elseif d.instr == 'W'
            pos = (x = pos.x - d.val, y = pos.y)
        elseif d.instr == 'F'
            pos = (x = pos.x + cos(rad) * d.val, y = pos.y + sin(rad) * d.val)
        elseif d.instr == 'L'
            rad -= deg2rad(d.val)
        elseif d.instr == 'R'
            rad += deg2rad(d.val)
        end
        @info "pos=$pos rad=$rad"
    end
    return abs(pos.x) + abs(pos.y)
end

function part2()
    pos = (x = 0, y = 0)
    way = (x = 10, y = -1)
    for d in read_data()
        dx, dy = way.x - pos.x, way.y - pos.y
        if d.instr == 'N'
            way = (x = way.x, y = way.y - d.val)
        elseif d.instr == 'S'
            way = (x = way.x, y = way.y + d.val)
        elseif d.instr == 'E'
            way = (x = way.x + d.val, y = way.y)
        elseif d.instr == 'W'
            way = (x = way.x - d.val, y = way.y)
        elseif d.instr == 'F'
            pos = (x = pos.x + d.val * dx, y = pos.y + d.val * dy)
            way = (x = way.x + d.val * dx, y = way.y + d.val * dy)
        elseif d.instr == 'L'
            r = sqrt(dx^2 + dy^2)
            θ = atan(dy, dx)
            ϕ = θ - deg2rad(d.val)
            way = (x = pos.x + r * cos(ϕ), y = pos.y + r * sin(ϕ))
        elseif d.instr == 'R'
            r = sqrt(dx^2 + dy^2)
            θ = atan(dy, dx)
            ϕ = θ + deg2rad(d.val)
            way = (x = pos.x + r * cos(ϕ), y = pos.y + r * sin(ϕ))
        else
            error("wat")
        end
    end
    return abs(pos.x) + abs(pos.y)
end

# Animation

using Plots
using FileIO

function part2_animated()
    pos = (x = 0, y = 0)
    way = (x = 10, y = 1)
    anim = Animation()
    # img = FileIO.load("/Users/tomkwong/Downloads/iu.jpeg")
    # plot(img)
    # frame(anim)
    @animate for (i,d) in enumerate(read_data())
        dx, dy = way.x - pos.x, way.y - pos.y
        if d.instr == 'N'
            way = (x = way.x, y = way.y + d.val)
        elseif d.instr == 'S'
            way = (x = way.x, y = way.y - d.val)
        elseif d.instr == 'E'
            way = (x = way.x + d.val, y = way.y)
        elseif d.instr == 'W'
            way = (x = way.x - d.val, y = way.y)
        elseif d.instr == 'F'
            pos = (x = pos.x + d.val * dx, y = pos.y + d.val * dy)
            way = (x = way.x + d.val * dx, y = way.y + d.val * dy)
        elseif d.instr == 'L'
            r = sqrt(dx^2 + dy^2)
            θ = atan(dy, dx)
            ϕ = θ + deg2rad(d.val)
            way = (x = pos.x + r * cos(ϕ), y = pos.y + r * sin(ϕ))
        elseif d.instr == 'R'
            r = sqrt(dx^2 + dy^2)
            θ = atan(dy, dx)
            ϕ = θ - deg2rad(d.val)
            way = (x = pos.x + r * cos(ϕ), y = pos.y + r * sin(ϕ))
        else
            error("wat")
        end
        if i % 2 == 0
            s, adx, ady = 3, abs(dx), abs(dy)
            ms = 15
            xlims = (min(pos.x, way.x) - s * adx, max(pos.x, way.x) + s * adx)
            ylims = (min(pos.y, way.y) - s * ady, max(pos.y, way.y) + s * ady)
            scatter([pos.x], [pos.y], xlims = xlims, ylims = ylims, markersize = ms,
                legend = nothing, ticks = nothing, showaxis = false,
                background = "black", markershape = :octagon,
                title = "AoC Day 12 - Iteration $i")
            scatter!([way.x], [way.y],
                xlims = xlims, ylims = ylims, markersize = ms + 5,
                markershape = :star)
            frame(anim)
        end
    end
    gif(anim, "day12_anim.gif", fps = 15)
end


module AssigmentProblem
using JuMP
using GLPK

"""
Problema de asignación. (Vanderbei LP, pg. 227). Dado un conjunto S de m personas, un conjunto D de m tareas, y para cada iS,jD un costo cij de asignar la tarea j a la persona i, el problema de asignación consiste en asignar a cada persona una y solo una tarea, de manera tal que se minimize el costo total. Formule este problema como un problema de programación lineal con enteros. Escriba una función para resolverlo, que tome como entrada una matriz de costos cij arbitraria.
"""

#https://github.com/josePereiro/AssigmentProblem.jl/

include("BruteForce.jl")
include("LinearProgramming.jl")
include("Common.jl");

function hello()
	println("hello");
end

end # module


function drelu_dZ(NN::NeuralNetwork, l::Int)
    return @. NN.cache.Z[l] > 0
end

function dlogistic_dZ(NN::NeuralNetwork, l::Int)
    return @. NN.cache.A[l] * (1 - NN.cache.A[l])
end

function dtanh_dZ(NN::NeuralNetwork, l::Int)
    return @. 1 - NN.cache.A[l]^2
end

function dsoftmax_dZ(NN::NeuralNetwork, l::Int)
    return @. NN.cache.A[l] * (1 - NN.cache.A[l])
end

function propagate_back!(NN::NeuralNetwork, mini_batch::MiniBatch)
    for l = NN.L:-1:1
        compute_dJ_dA!(NN, l, mini_batch)
        compute_dJ_dZ!(NN, l)
        compute_dJ_dW!(NN, l, mini_batch)
        compute_dJ_db!(NN, l, mini_batch)
    end
end

function compute_dJ_dA!(NN::NeuralNetwork, l::Int, mini_batch::MiniBatch)
    if l ≠ NN.L
        NN.cache.dA[l] = NN.params.W[l+1]' * NN.cache.dZ[l+1]
    else
        if NN.params.loss_function == cross_entropy_binary
            NN.cache.dA[l] = @. -mini_batch.Y / NN.cache.A[NN.L] +
               (1 - mini_batch.Y) / (1 - NN.cache.A[NN.L])
        elseif NN.params.loss_function == cross_entropy_multi
            NN.cache.dA[l] = @. -mini_batch.Y / NN.cache.A[NN.L]
        elseif NN.params.loss_function == mean_squared_error
            NN.cache.dA[l] = @. NN.cache.A[NN.L] - mini_batch.Y
        else
            error("unknown loss function provided in params")
        end
    end
end

function compute_dJ_dZ!(NN::NeuralNetwork, l::Int)
    if NN.params.activation_functions[l] == tanh
        NN.cache.dZ[l] = NN.cache.dA[l] .* drelu_dZ(NN, l)
    elseif NN.params.activation_functions[l] == relu
        NN.cache.dZ[l] = NN.cache.dA[l] .* dtanh_dZ(NN, l)
    elseif NN.params.activation_functions[l] == logistic
        NN.cache.dZ[l] = NN.cache.dA[l] .* dlogistic_dZ(NN, l)
    elseif NN.params.activation_functions[l] == softmax
        NN.cache.dZ[l] = NN.cache.dA[l] .* dsoftmax_dZ(NN, l)
    else
        error("unknown activation function(s) provided in params")
    end
end

function compute_dJ_dW!(NN::NeuralNetwork, l::Int, mini_batch::MiniBatch)
    if l ≠ 1
        NN.cache.dW[l] = (NN.cache.dZ[l] ./ size(mini_batch.X, 2)) * NN.cache.A[l-1]'
    else
        NN.cache.dW[l] = (NN.cache.dZ[l] ./ size(mini_batch.X, 2)) * mini_batch.X'
    end

    # Add in regularization if using gradient descent
    if NN.hparams.optimization == :gd
        NN.cache.dW[l] += (NN.hparams.λ / size(mini_batch.X, 2)) .* NN.params.W[l]
    end
end

function compute_dJ_db!(NN::NeuralNetwork, l::Int, mini_batch::MiniBatch)
    NN.cache.db[l] = sum(NN.cache.dZ[l], dims = 2) ./ size(mini_batch.X, 2)
end


using DiffEqBase, MATLABDiffEq, ParameterizedFunctions, Test

f = @ode_def_bare LotkaVolterra begin
  dx = a*x - b*x*y
  dy = -c*y + d*x*y
end a b c d
p = [1.5,1,3,1]
tspan = (0.0,10.0)
u0 = [1.0,1.0]
prob = ODEProblem(f,u0,tspan,p)
sol = solve(prob,MATLABDiffEq.ode45())

function lorenz(du,u,p,t)
 du[1] = 10.0(u[2]-u[1])
 du[2] = u[1]*(28.0-u[3]) - u[2]
 du[3] = u[1]*u[2] - (8/3)*u[3]
end
u0 = [1.0;0.0;0.0]
tspan = (0.0,100.0)
prob = ODEProblem(lorenz,u0,tspan)
sol = solve(prob,MATLABDiffEq.ode45())

algs = [MATLABDiffEq.ode23
        MATLABDiffEq.ode45
        MATLABDiffEq.ode113
        MATLABDiffEq.ode23s
        MATLABDiffEq.ode23t
        MATLABDiffEq.ode23tb
        MATLABDiffEq.ode15s]

for alg in algs
  sol = solve(prob,alg())
end


"""
    JoeCopula{d,T}

Fields:
  - θ::Real - parameter

Constructor

    JoeCopula(d, θ)

The [Joe](https://en.wikipedia.org/wiki/Copula_(probability_theory)#Most_important_Archimedean_copulas) copula in dimension ``d`` is parameterized by ``\\theta \\in [1,\\infty)``. It is an Archimedean copula with generator :

```math
\\phi(t) = 1 - \\left(1 - e^{-t}\\right)^{\\frac{1}{\\theta}}
```
"""
struct JoeCopula{d,T} <: ArchimedeanCopula{d}
    θ::T
end
JoeCopula(d,θ) = θ >= 1 ? JoeCopula{d,typeof(θ)}(θ) : @error "Theta must be greater than one."
ϕ(  C::JoeCopula,          t) = 1-(1-exp(-t))^(1/C.θ)
ϕ⁻¹(C::JoeCopula,          t) = -log(1-(1-t)^C.θ)
τ(C::JoeCopula) = 1 - 4sum(1/(k*(2+k*C.θ)*(C.θ*(k-1)+2)) for k in 1:1000) # 446 in R copula.
radial_dist(C::JoeCopula) = Sibuya(1/C.θ)


function solve()
    k = parse(Int, readline())
    if k % 2 == 0 || k % 5 == 0
        -1
    else
        ret = 1
        x = 7
        while x % k ≠ 0
            x = x % k * 10 + 7
            ret += 1
        end
        ret
    end
end

println(solve())


""" Cálculo de IMC
    Leia o peso e a altura do usuário
    imprima seu Índice de Massa Corporal.
        * IMC = peso / (altura^2)
    OBS: use ponto final para informar o peso
"""

println("informe seu peso")
weight = parse(Float32, readline())

println("informe sua altura")
height = parse(Float32, readline())

imc = (weight / (height^2))
println("IMC = ", imc)


# Cazzolato, Ben. "Derivations of the Dynamics of the Ball and Beam System." School of mechanical engineering, The University of Adelaide 11 (2007): 2010-05.
# https://www.researchgate.net/profile/Ben_Cazzolato/publication/235898231_Derivation_of_the_Dynamics_of_the_Ball_and_Beam_System/links/569338cf08aec14fa55db5a8/Derivation-of-the-Dynamics-of-the-Ball-and-Beam-System.pdf

# fixed parameters from the paper
# L  (m)  length of the beam
# d  (m)  distance from pivot to the plane of the ball contact on the beam
# D  (m)  distance from the pivot to the center of mass of the beam
# R0 (m)  radius of the ball
# R1 (m)  distance between the axis of rotation of the ball (center of gravity) and the point of contact of the ball with the beam
# m  (kg) mass of the ball
# M  (kg) mass of the beam (located at d (m) from the pivot)
# C1 (N/(m/s)) is the viscous friction coefficient between the ball and the beam (and accounts for the ball rotational viscous losses too since these enter at the same dynamic order)
# C2 is the damping term for the beams velocity
# Jb (kg.m^2)  is the moment of inertia of the beam, including all rotational components such as rotor include beam mass offset (parallel axis theorem)
# g=9.81 (m/s^2) is the gravitational acceleration

# independent variables
# r  (m)   displacement of the ball along the beam. Positive r is when the ball is traveling to the right and r=0 represents the center of the beam
# θ  (rad) is the angular rotation of the beam, where a counter-clockwise direction is Positive

# dependent variables
# τ  (Nm) is the torque the servo-motor applies to the beam, where a counter-closewise directin is Positive

# servo params (not from paper)
# kp  proportional constant
# servo model: produces a torque corresponding to a reference beam angle u
# τ = kp * (u - θ)
# The beam's friction constant acts as the kd term in a PD controller

# uses the linearized system dynamics from the paper.
function ballbeam_sim!(state, u, params, dt, sim_steps::Int)
	r, ṙ, θ, θ̇ = state[1], state[2], state[3], state[4]

    L, d, D, R0, R1, m, M, C1, C2, Jb, g, kp = params

    d² = d^2
    R0² = R0^2
    R1² = R1^2

    Δt = dt / sim_steps
    K = 2*R0²*m*d² + 5*R1²*Jb + 2*R0²*Jb
    Q = 5*R1* + 2*R0² + 5 * R1²
    Q2 = (5*R1² + 2 * R0²)
    A21 = -m*g*R1 * Q / K
    A22 = -R1²*C1 * (2*R0² + 5*d² + 10*R1*d + 5*R1² + 5*Jb / m) / K
    A23 = g*R1*(2*R0²*M*D - 5*R1*Jb + 2*R0²*m*d + 5*R1²*M*D + 5*R1*d*M*D) / K
    A24 = -C2*R1*Q / K

    A41 = -m*g*Q2 / K
    A42 = -C1*R1*Q/K
    A43 = g * (2*R0²*m*d + 5*R1²*M*D + 2*R0²*M*D) / K
    A44 = -C2*Q2/K

    B21 = R1 * Q / K
    B41 = Q2 / K

    for i in 1:sim_steps
        τ = kp * (u - θ) #- C2 * θ̇
        dr = ṙ
        dθ = θ̇

        dṙ = A21 * r + A22 * ṙ + A23 * θ + A24 * θ̇ + B21 * τ

        # println(dr)

        dθ̇ = A41 * r + A42 * ṙ + A43 * θ + A44 * θ̇ + B41 * τ

        r += Δt * dr
        ṙ += Δt * dṙ
        θ += Δt * dθ
        θ̇ += Δt * dθ̇
	end
    # if ball hits the side
    Lhalf = L/2 - R0
    if r > Lhalf
        ṙ = 0.0#0.25 * (Lhalf - r)  # if ball hits the side bounce back with half the velocity (damping by 0.25 so it cannot have infinite energy)
        r = Lhalf
    elseif r < -Lhalf
        ṙ = 0.0#0.25 * (-Lhalf - r)  # if ball hits the side bounce back with half the velocity (damping by 0.25 so it cannot have infinite energy)
        r = -Lhalf
    end

    # if beam tilts past 5 degress it halts (all energy is absorbed)
    θlim = deg2rad(5)
    if θ > θlim
        θ = θlim
        θ̇ = 0.0
    elseif θ < -θlim
        θ = -θlim
        θ̇ = 0.0
    end
	state[1:4] .= r, ṙ, θ, θ̇
	return nothing
end

function ballbeam_params(randomize=false)
    L = 0.25 # 0.25  # range [0.25, 0.5]
    D = 0.0 # torque applied to center of beam
    d = 0.01
    R0 = 0.008
    R1 = R0  # ball rest ontop of beam
    m = 0.05
    M = 0.25
    C1 = 0.2
    C2 = 3.0  # stable constant perhaps don't change.
    Jb = (1.0/12.0) * M * L^2 # moment of inertia
    g = 9.81
    kp = 10.0 #8.0  # range [8, 20] # 5 is a pretty slow response

    if randomize
        L = rand() * 0.25 + 0.25    # [0.250,  0.50]
        m = rand() * 0.975 + 0.025  # [0.025,  0.10]
        M = rand() * 0.75 + 0.25    # [0.250,  1.00]
        C1 = rand() * 0.15 + 0.1    # [0.100,  0.25]
        kp = rand() * 12 + 8.0      # [8.000, 20.00]
    end

    return L, d, D, R0, R1, m, M, C1, C2, Jb, g, kp
end

function ballandbeam_fixedgoal(;randomize=false, maxT=20.0, dt=0.05, droptime=true, stochastic_start=true, action_type=:Discrete)
    params = ballbeam_params(randomize)

	return create_ballandbeam_finitetime(params, maxT=maxT, dt=dt, droptime=droptime, stochastic_start=stochastic_start, action_type=action_type, goal_mode=:fixed)
end

function ballandbeam_randomgoal(;randomize=false, maxT=20.0, dt=0.05, droptime=true, stochastic_start=true, action_type=:Discrete)
    params = ballbeam_params(randomize)

	return create_ballandbeam_finitetime(params, maxT=maxT, dt=dt, droptime=droptime, stochastic_start=stochastic_start, action_type=action_type, goal_mode=:random)
end

function ballandbeam_tacking(;randomize=false, maxT=20.0, dt=0.05, droptime=true, stochastic_start=true, action_type=:Discrete)
    params = ballbeam_params(randomize)

	return create_ballandbeam_finitetime(params, maxT=maxT, dt=dt, droptime=droptime, stochastic_start=stochastic_start, action_type=action_type, goal_mode=:track)
end

function create_ballandbeam_finitetime(params; maxT=20.0, dt=0.05, droptime=true, stochastic_start=true, action_type=:Discrete, goal_mode=:fixed)
    X = zeros((5,2))
    L, R0 = params[1], params[4]
    Lhalf = L/2
	X[1,:] .= [-Lhalf+R0, Lhalf-R0]   # ball position range
    X[2,:] .= [-0.7, 0.7]             # ball velocity range (seems about right for fastest ball acceleration)
    θlim = deg2rad(5)
    X[3,:] .= [-θlim, θlim]          # beam is limited to ± 5 degrees
    X[4,:] .= [-2*θlim, 2*θlim]      # this is too high and is probably closer to ±0.1 for default parameters
    X[5,:] .= [-(L - 2*R0), L - 2*R0] # delta to goal position
	S = ([0.0, maxT], X[1:4, :], X[1, :])
    if action_type == :Discrete
        A = 1:3
    else
        A = [-1.0, 1.0]
    end
	sim_steps = 100  # cannot reduce much below here for all configurations of ball and beam. Long beams are easy to simulate.
    if goal_mode == :track
        d = (L/2 - R0) * 0.9
        y = t->cos(t/2.0+π)*d
        step_fn = (s,a)->ballandbeam_step_track!(s,a,params,y,dt,sim_steps,maxT)
        intfn = ()->bb_sample_initial_track(params, y)
    elseif goal_mode == :fixed
        step_fn = (s,a)->ballandbeam_step!(s,a,params,dt,sim_steps,maxT)
        intfn = ()->bb_sample_initial_fixed(params, stochastic_start)
    elseif goal_mode == :random
        step_fn = (s,a)->ballandbeam_step!(s,a,params,dt,sim_steps,maxT)
        intfn = ()->bb_sample_initial_random(params)
    end
    if droptime
		function get_outcome1(s,a,step_fn)
			t, x, g, r, γ = step_fn(s,a)
			s = (t,x,g)
			return s,x,r,γ
		end
        p = (s,a)->get_outcome1(s,a,step_fn)

        function bbd0_obs()
            t,x,g = intfn()
            return (t,x,g), x
        end
        d0 = bbd0_obs
	else
		X = (S[1], X)
		function get_outcome2(s,a,step_fn)
			t, x, g, r, γ = step_fn(s,a)
            s = (t,x,g)
            obs = (t,x)
			return s,obs,r,γ
		end
        p = (s,a)->get_outcome2(s,a,step_fn)
        function bbd0()
            t,x,g = intfn()
            s = (t,x,g)
            obs = (t,x)
            return s,obs
        end
		d0 = bbd0
    end

    numT = ceil(Int, maxT / dt)
	meta = Dict{Symbol,Any}()
    meta[:minreward] = -22.21  # max speed and distance from goal on longest beam
    meta[:maxreward] = 0
    meta[:minreturn] = -2000.0 # chosen experimentally as a lower bound (more likely to be around 1100)
    meta[:maxreturn] = -10.0 # chosen experimentall as an upper bound. Highest I've seen is low -20s
    meta[:stochastic] = false
    meta[:minhorizon] = numT
    meta[:maxhorizon] = numT
	meta[:discounted] = false
	meta[:episodes] = 10000
    meta[:threshold] = -50  # somewhat generous threshold. Should quickly balance the ball around the goal location

	render = (state,clearplot=false)->ballbeamplot(state, params)

	m = SequentialProblem(S,X,A,p,d0,meta,render)
end

function bbget_refpos(action::Int)
    θlim = deg2rad(5)
    if action == 1
        return -θlim
    elseif action == 2
        return 0.0
    else
        return θlim
    end
end

function bbget_refpos(action::Float64)
    return clamp(action, -1.0, 1.0) * deg2rad(5.0)
end

function ballandbeam_step!(state, action, params, dt, sim_steps, maxT)
	u = bbget_refpos(action)
    t,x,g = state
    dcoef = 10.0
    vcoef = 2.0
    olddist = (x[5]*dcoef)^2
    oldv = (x[2]*vcoef)^2
    ballbeam_sim!(x, u, params, dt, sim_steps)
    t += dt
    x[5] = g - x[1]
    curdist = (x[5]*dcoef)^2
    curv = (x[2]*vcoef)^2

    # reward = (olddist - curdist) + (oldv - curv) # shaping reward
    reward = -(curdist + curv)  # negative quadratic cost
	γ = 1.0
	done = t ≥ maxT

	if done
        γ = 0.0
	end

	return t, x, g, reward, γ
end

function ballandbeam_step_track!(state, action, params, y, dt, sim_steps, maxT)
	u = bbget_refpos(action)
    t,x,g = state
    dcoef = 10.0
    olddist = (x[5]*dcoef)^2
    ballbeam_sim!(x, u, params, dt, sim_steps)
    t += dt
    g = y(t)
    x[5] = g - x[1]
    curdist = (x[5]*dcoef)^2

    reward = -curdist
	γ = 1.0
	done = t ≥ maxT

	if done
        γ = 0.0
	end

	return t, x, g, reward, γ
end

function bb_sample_initial_fixed(params, stochastic_start=true)
    x = zeros(5)
    L, R = params[1], params[4]
    x[1] = -L/2 + R

    if stochastic_start
        x[1] += rand() * 0.2 * L
    end
    g = L/4
    x[5] = g - x[1]

	return 0.0, x, g
end

function bb_sample_initial_random(params)
    x = zeros(5)
    L, R = params[1], params[4]
    Lhalf = L/2 - R

    g = rand() * 2*Lhalf - Lhalf
    if g < 0
        x[1] = rand() * 0.5 * Lhalf
    else
        x[1] = -rand() * 0.5 * Lhalf
    end

    x[5] = g - x[1]

	return 0.0, x, g
end

function bb_sample_initial_track(params, y)
    x = zeros(5)
    L, R = params[1], params[4]
    x[1] = -L/2 + R
    t = 0.0
    g = y(t)
    x[5] = g - x[1]

	return t, x, g
end

# Position of the ball
function bb_ballpos(r, θ, d)
    x1, y1 = r * cos(θ), r * sin(θ)
    x2, y2 = d * cos(θ+π/2), d * sin(θ + π/2)
    return x1+x2, y1+y2
end

@userplot BallBeamPlot
@recipe function f(ap::BallBeamPlot)
    state, params = ap.args
    t,x,g = state
    r, θ = x[1], x[3]
    L, d, D, R0= params[1:4]
    # println("R0 $R0")
	maxlen = (L/2)*1.2
    dlow = -d

    ballx, bally = bb_ballpos(r, θ, d + R0)
    goalx, goaly = bb_ballpos(g, θ, d + R0)

    # println(ballx, " ", bally)
	legend := false
	xlims := (-maxlen, maxlen)
	ylims := (-maxlen*sin(deg2rad(5))-(2*R0+d), maxlen*sin(deg2rad(5))+(2*R0+d))
	grid := false
	ticks := nothing
	foreground_color := :white
	aspect_ratio := 1.

    # beam
	@series begin
        seriescolor := :orange
        linecolor := nothing
        seriestype := :shape
		beamx11, beamy11 = bb_ballpos(-L/2, θ, d)
        beamx21, beamy21 = bb_ballpos(L/2, θ, d)
        beamx22, beamy22 = bb_ballpos(L/2, θ, dlow)
        beamx12, beamy12 = bb_ballpos(-L/2, θ, dlow)
		[beamx11, beamx21, beamx22, beamx12], [beamy11, beamy21, beamy22, beamy12]
    end

    # left end
	@series begin
        seriescolor := :orange
        linecolor := nothing
        seriestype := :shape
		beamx11, beamy11 = bb_ballpos(-L/2 - d, θ, d+2*R0)
        beamx21, beamy21 = bb_ballpos(-L/2, θ, d+2*R0)
        beamx22, beamy22 = bb_ballpos(-L/2, θ, dlow)
        beamx12, beamy12 = bb_ballpos(-L/2 -d, θ, dlow)
		[beamx11, beamx21, beamx22, beamx12], [beamy11, beamy21, beamy22, beamy12]
    end

    # right end
	@series begin
        seriescolor := :orange
        linecolor := nothing
        seriestype := :shape
		beamx11, beamy11 = bb_ballpos(L/2, θ, d+2*R0)
        beamx21, beamy21 = bb_ballpos(L/2 + d, θ, d+2*R0)
        beamx22, beamy22 = bb_ballpos(L/2 + d, θ, dlow)
        beamx12, beamy12 = bb_ballpos(L/2, θ, dlow)
		[beamx11, beamx21, beamx22, beamx12], [beamy11, beamy21, beamy22, beamy12]
    end

	# goal location
	@series begin
		seriestype := :shape
		linecolor := nothing
        seriescolor := :green
        seriescolor := :green
		aspect_ratio := 1.
		fillalpha := 0.4

		circle_shape(goalx, goaly, R0)
	end

    # ball location
	@series begin
		seriestype := :shape
		linecolor := nothing
		seriescolor := :red
        aspect_ratio := 1.

		circle_shape(ballx, bally, R0)
    end

end


@doc raw"""
    StateSpaceSystem

Abstract type that unifies the definition of state space models matrices such as ``y, Z, d, T, c, R, H, Q`` for
linear models.
"""
abstract type StateSpaceSystem end

@doc raw"""
    LinearUnivariateTimeInvariant{Fl}(
        y::Vector{Fl},
        Z::Vector{Fl},
        T::Matrix{Fl},
        R::Matrix{Fl},
        d::Fl,
        c::Vector{Fl},
        H::Fl,
        Q::Matrix{Fl},
    ) where Fl <: AbstractFloat

Definition of the system matrices ``y, Z, d, T, c, R, H, Q`` for linear univariate time invariant state space models.

```math
\begin{gather*}
    \begin{aligned}
        y_{t} &=  Z\alpha_{t} + d + \varepsilon_{t} \quad &\varepsilon_{t} \sim \mathcal{N}(0, H)\\
        \alpha_{t+1} &= T\alpha_{t} + c + R\eta_{t} \quad &\eta_{t} \sim \mathcal{N}(0, Q)\\
    \end{aligned}
\end{gather*}
```

where:

* ``y_{t}`` is a scalar
* ``Z`` is a ``m \times 1`` vector
* ``d`` is a scalar
* ``T`` is a ``m \times m`` matrix
* ``c`` is a ``m \times 1`` vector
* ``R`` is a ``m \times r`` matrix
* ``H`` is a scalar
* ``Q`` is a ``r \times r`` matrix
"""
mutable struct LinearUnivariateTimeInvariant{Fl<:AbstractFloat} <: StateSpaceSystem
    y::Vector{Fl}
    Z::Vector{Fl}
    T::Matrix{Fl}
    R::Matrix{Fl}
    d::Fl
    c::Vector{Fl}
    H::Fl
    Q::Matrix{Fl}

    function LinearUnivariateTimeInvariant{Fl}(
        y::Vector{Fl},
        Z::Vector{Fl},
        T::Matrix{Fl},
        R::Matrix{Fl},
        d::Fl,
        c::Vector{Fl},
        H::Fl,
        Q::Matrix{Fl},
    ) where Fl <: AbstractFloat
        mz = length(Z)
        mt1, mt2 = size(T)
        mr, rr = size(R)
        mc = length(c)
        rq1, rq2 = size(Q)

        dim_str =
            "Z is $(mz)x1, T is $(mt1)x$(mt2), R is $(mr)x$(rr), " *
            "d is a scalar, c is $(mc)x1, H is a scalar."

        !(mz == mt1 == mt2 == mr == mc) && throw(DimensionMismatch(dim_str))
        !(rr == rq1 == rq2) && throw(DimensionMismatch(dim_str))

        return new{Fl}(y, Z, T, R, d, c, H, Q)
    end
end

num_states(system::LinearUnivariateTimeInvariant) = size(system.T, 1)

@doc raw"""
    LinearUnivariateTimeVariant{Fl}(
        y::Vector{Fl},
        Z::Vector{Vector{Fl}},
        T::Vector{Matrix{Fl}},
        R::Vector{Matrix{Fl}},
        d::Vector{Fl},
        c::Vector{Vector{Fl}},
        H::Vector{Fl},
        Q::Vector{Matrix{Fl}},
    ) where Fl <: AbstractFloat

Definition of the system matrices ``y, Z, d, T, c, R, H, Q`` for linear univariate time variant state space models.

```math
\begin{gather*}
    \begin{aligned}
        y_{t} &=  Z_{t}\alpha_{t} + d_{t} + \varepsilon_{t} \quad &\varepsilon_{t} \sim \mathcal{N}(0, H_{t})\\
        \alpha_{t+1} &= T_{t}\alpha_{t} + c_{t} + R_{t}\eta_{t} \quad &\eta_{t} \sim \mathcal{N}(0, Q_{t})\\
    \end{aligned}
\end{gather*}
```

where:

* ``y_{t}`` is a scalar
* ``Z_{t}`` is a ``m \times 1`` vector
* ``d_{t}`` is a scalar
* ``T_{t}`` is a ``m \times m`` matrix
* ``c_{t}`` is a ``m \times 1`` vector
* ``R_{t}`` is a ``m \times r`` matrix
* ``H_{t}`` is a scalar
* ``Q_{t}`` is a ``r \times r`` matrix
"""
mutable struct LinearUnivariateTimeVariant{Fl<:AbstractFloat} <: StateSpaceSystem
    y::Vector{Fl}
    Z::Vector{Vector{Fl}}
    T::Vector{Matrix{Fl}}
    R::Vector{Matrix{Fl}}
    d::Vector{Fl}
    c::Vector{Vector{Fl}}
    H::Vector{Fl}
    Q::Vector{Matrix{Fl}}

    function LinearUnivariateTimeVariant{Fl}(
        y::Vector{Fl},
        Z::Vector{Vector{Fl}},
        T::Vector{Matrix{Fl}},
        R::Vector{Matrix{Fl}},
        d::Vector{Fl},
        c::Vector{Vector{Fl}},
        H::Vector{Fl},
        Q::Vector{Matrix{Fl}},
    ) where Fl <: AbstractFloat

        # TODO assert dimensions

        return new{Fl}(y, Z, T, R, d, c, H, Q)
    end
end

num_states(system::LinearUnivariateTimeVariant) = size(system.T[1], 1)

"""
TODO
"""
mutable struct LinearMultivariateTimeInvariant{Fl<:Real} <: StateSpaceSystem
    y::Matrix{Fl}
    Z::Matrix{Fl}
    T::Matrix{Fl}
    R::Matrix{Fl}
    d::Vector{Fl}
    c::Vector{Fl}
    H::Matrix{Fl}
    Q::Matrix{Fl}

    function LinearMultivariateTimeInvariant{Fl}(
        y::Matrix{Fl},
        Z::Matrix{Fl},
        T::Matrix{Fl},
        R::Matrix{Fl},
        d::Vector{Fl},
        c::Vector{Fl},
        H::Matrix{Fl},
        Q::Matrix{Fl},
    ) where Fl

        # TODO assert dimensions

        return new{Fl}(y, Z, T, R, d, c, H, Q)
    end
end

hasnan(y::Vector{Fl}) where Fl = findfirst(isnan, y) === nothing ? false : true
num_states(system::LinearMultivariateTimeInvariant) = size(system.T, 1)

"""
TODO
"""
mutable struct LinearMultivariateTimeVariant{Fl<:Real} <: StateSpaceSystem
    y::Matrix{Fl}
    Z::Vector{Matrix{Fl}}
    T::Vector{Matrix{Fl}}
    R::Vector{Matrix{Fl}}
    d::Vector{Vector{Fl}}
    c::Vector{Vector{Fl}}
    H::Vector{Matrix{Fl}}
    Q::Vector{Matrix{Fl}}

    function LinearMultivariateTimeVariant{Fl}(
        y::Matrix{Fl},
        Z::Vector{Matrix{Fl}},
        T::Vector{Matrix{Fl}},
        R::Vector{Matrix{Fl}},
        d::Vector{Vector{Fl}},
        c::Vector{Vector{Fl}},
        H::Vector{Matrix{Fl}},
        Q::Vector{Matrix{Fl}},
    ) where Fl

        # TODO assert dimensions

        return new{Fl}(y, Z, T, R, d, c, H, Q)
    end
end

# Some useful functions for dealing with these systems
function fill_system_matrice_with_value_in_time(vector::Vector{Fl}, value::Fl) where Fl
    for t in axes(vector, 1)
        vector[t] = value
    end
end

function repeat_system_matrice_n_times(value::T, n::Int) where T
    vector_of_matrice = Vector{T}(undef, n)
    for i in 1:n
        vector_of_matrice[i] = value
    end
    return vector_of_matrice
end

function to_multivariate_Z(Z::Vector)
    return permutedims(Z)
end
function to_multivariate_Z(Z::Vector{Vector{Fl}}) where Fl
    multivariate_Z = Vector{Matrix{Fl}}(undef, length(Z))
    for i in 1:length(Z)
        multivariate_Z[i] = permutedims(Z[i])
    end
    return multivariate_Z
end
function to_multivariate_H(H::Fl) where Fl
    return [H][:, :]
end
function to_multivariate_H(H::Vector{Fl}) where Fl
    multivariate_H = Vector{Matrix{Fl}}(undef, length(H))
    for i in 1:length(H)
        multivariate_H[i] = permutedims([H[i]])
    end
    return multivariate_H
end
function to_multivariate_d(d::Fl) where Fl
    return [d]
end
function to_multivariate_d(d::Vector{Fl}) where Fl
    multivariate_d = Vector{Vector{Fl}}(undef, length(d))
    for i in 1:length(d)
        multivariate_d[i] = [d[i]]
    end
    return multivariate_d
end

# Bridges between Linear systems

# LinearUnivariateTimeInvariant to (LinearUnivariateTimeVariant,
#                                   LinearMultivariateTimeInvariant, TODO
#                                   LinearMultivariateTimeVariant)
function to_univariate_time_variant(system::LinearUnivariateTimeInvariant{Fl}) where Fl
    n = length(system.y)
    y = system.y
    Z = repeat_system_matrice_n_times(system.Z, n)
    T = repeat_system_matrice_n_times(system.T, n)
    R = repeat_system_matrice_n_times(system.R, n)
    d = repeat_system_matrice_n_times(system.d, n)
    c = repeat_system_matrice_n_times(system.c, n)
    H = repeat_system_matrice_n_times(system.H, n)
    Q = repeat_system_matrice_n_times(system.Q, n)
    return LinearUnivariateTimeVariant{Fl}(y, Z, T, R, d, c, H, Q)
end
function to_multivariate_time_variant(system::LinearUnivariateTimeInvariant)
    univariate_time_variant = to_univariate_time_variant(system)
    return to_multivariate_time_variant(univariate_time_variant)
end

# LinearUnivariateTimeVariant to (LinearMultivariateTimeVariant)
function to_multivariate_time_variant(system::LinearUnivariateTimeVariant{Fl}) where Fl
    y = system.y[:, :]
    Z = to_multivariate_Z(system.Z)
    T = system.T
    R = system.R
    d = to_multivariate_d(system.d)
    c = system.c
    H = to_multivariate_H(system.H)
    Q = system.Q
    return LinearMultivariateTimeVariant{Fl}(y, Z, T, R, d, c, H, Q)
end

to_multivariate_time_variant(system::LinearMultivariateTimeVariant) = system

# Functions for simulations
function simulate(
    sys::LinearUnivariateTimeInvariant{Fl},
    initial_state::Vector{Fl},
    n::Int;
    return_simulated_states::Bool=false,
) where Fl
    m = size(sys.T, 1)
    y = Vector{Fl}(undef, n)
    alpha = Matrix{Fl}(undef, n + 1, m)
    # Sampling errors
    chol_H = sqrt(sys.H)
    chol_Q = cholesky(sys.Q)
    standard_ε = randn(n)
    standard_η = randn(n + 1, size(sys.Q, 1))

    # The first state of the simulation is the update of a_0
    alpha[1, :] = sys.T * initial_state + sys.c + sys.R * chol_Q.L * standard_η[1, :]
    # Simulate scenarios
    for t in 1:n
        y[t] = dot(sys.Z, alpha[t, :]) + sys.d + chol_H * standard_ε[t]
        alpha[t + 1, :] = sys.T * alpha[t, :] + sys.c + sys.R * chol_Q.L * standard_η[t, :]
    end

    if return_simulated_states
        return y, alpha[1:n, :]
    end
    return y
end


#==========================================================================
 Some functions from Julia package MLKernels.jl (v0.1), that I've borrowed and modified.
 See https://github.com/trthatcher/MLKernels.jl/blob/master/LICENSE.md:
"The MIT License (MIT)
Copyright (c) 2015 Tim Thatcher
Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE."
==========================================================================#

export Kernel
export kernelmatrix, SquaredDistanceKernel, PolynomialKernel, LinearKernel
export GaussianKernel

abstract type Kernel{T} end
abstract type StandardKernel{T<:AbstractFloat} <: Kernel{T} end
abstract type BaseKernel{T<:AbstractFloat} <: StandardKernel{T} end
abstract type CompositeKernel{T<:AbstractFloat} <: StandardKernel{T} end

#==========================================================================
  Additive Kernel: k(x,y) = sum(k(x_i,y_i))    x ∈ ℝⁿ, y ∈ ℝⁿ
  Separable Kernel: k(x,y) = k(x)k(y)    x ∈ ℝ, y ∈ ℝ
==========================================================================#

abstract type AdditiveKernel{T<:AbstractFloat} <: BaseKernel{T} end
abstract type SeparableKernel{T<:AbstractFloat} <: AdditiveKernel{T} end

phi(κ::SeparableKernel{T}, x::T, y::T) where {T<:AbstractFloat} = phi(κ,x) * phi(κ,y)
isnonnegative(κ::Kernel) = kernelrange(κ) == :Rp

#==========================================================================
  Squared Distance Kernel
  k(x,y) = (x-y)²ᵗ    x ∈ ℝ, y ∈ ℝ, t ∈ (0,1]
==========================================================================#

struct SquaredDistanceKernel{T<:AbstractFloat,CASE} <: AdditiveKernel{T}
    t::T
    function (::Type{SquaredDistanceKernel{T,CASE}})(t::T) where {T,CASE}
        0 < t <= 1 || error("Parameter t = $(t) must be in range (0,1]")
        new{T,CASE}(t)
    end
end

function SquaredDistanceKernel(t::T = 1.0) where {T<:AbstractFloat}
    CASE =  if t == 1
                :t1
            elseif t == 0.5
                :t0p5
            else
                :∅
            end
    SquaredDistanceKernel{T,CASE}(t)
end

isnegdef(::SquaredDistanceKernel) = true
kernelrange(::SquaredDistanceKernel) = :Rp

phi(κ::SquaredDistanceKernel{T,:t1}, x::Vector{T}, y::Vector{T}) where {T<:AbstractFloat} = sum(abs2, x-y)
phi(κ::SquaredDistanceKernel{T,:t0p5}, x::Vector{T}, y::Vector{T}) where {T<:AbstractFloat} = sum(abs, x-y)
phi(κ::SquaredDistanceKernel{T}, x::Vector{T}, y::Vector{T}) where {T<:AbstractFloat} = sum(abs2, x-y)^κ.t

#==========================================================================
  Scalar Product Kernel
  (Code preserved for future reasons)
==========================================================================#

struct ScalarProductKernel{T<:AbstractFloat} <: SeparableKernel{T} end
ScalarProductKernel() = ScalarProductKernel{Float64}()

ismercer(::ScalarProductKernel) = true

convert(::Type{ScalarProductKernel{T}}, κ::ScalarProductKernel) where {T<:AbstractFloat} = ScalarProductKernel{T}()

phi(κ::ScalarProductKernel{T}, x::T) where {T<:AbstractFloat} = x

convert(::Type{Kernel{T}}, κ::ScalarProductKernel) where {T<:AbstractFloat} = convert(ScalarProductKernel{T}, κ)

#==========================================================================
  Polynomial Kernel
==========================================================================#

struct PolynomialKernel{T<:AbstractFloat,CASE} <: CompositeKernel{T}
    k::BaseKernel{T}
    alpha::T
    c::T
    d::T
    function (::Type{PolynomialKernel{T,CASE}})(κ::BaseKernel{T}, α::T, c::T, d::T) where {T,CASE}
        ismercer(κ) == true || error("Composed kernel must be a Mercer kernel.")
        α > 0 || error("α = $(α) must be greater than zero.")
        c >= 0 || error("c = $(c) must be non-negative.")
        (d > 0 && trunc(d) == d) || error("d = $(d) must be an integer greater than zero.")
        if CASE == :d1 && d != 1
            error("Special case d = 1 flagged but d = $(convert(Int64,d))")
        end
        new{T,CASE}(κ, α, c, d)
    end
end

PolynomialKernel(κ::BaseKernel{T}, α::T = one(T), c::T = one(T), d::T = convert(T, 2)) where {T<:AbstractFloat} = PolynomialKernel{T, d == 1 ? :d1 : :Ø}(κ, α, c, d)
PolynomialKernel(α::T = 1.0, c::T = one(T), d::T = convert(T, 2)) where {T<:AbstractFloat} = PolynomialKernel(convert(Kernel{T},ScalarProductKernel()), α, c, d)
LinearKernel(α::T = 1.0, c::T = one(T)) where {T<:AbstractFloat} = PolynomialKernel(ScalarProductKernel(), α, c, one(T))

phi(κ::PolynomialKernel{T}, x::Vector{T}, y::Vector{T}) where {T<:AbstractFloat} = (κ.alpha*dot(x,y) + κ.c)^κ.d
phi(κ::PolynomialKernel{T,:d1}, x::Vector{T}, y::Vector{T}) where {T<:AbstractFloat} = κ.alpha*dot(x,y) + κ.c

#==========================================================================
  Exponential Kernel
==========================================================================#

struct ExponentialKernel{T<:AbstractFloat,CASE} <: CompositeKernel{T}
    k::BaseKernel{T}
    alpha::T
    gamma::T
    function (::Type{ExponentialKernel{T,CASE}})(κ::BaseKernel{T}, α::T, γ::T) where {T,CASE}
        isnegdef(κ) == true || error("Composed kernel must be negative definite.")
        isnonnegative(κ) || error("Composed kernel must attain only non-negative values.")
        α > 0 || error("α = $(α) must be greater than zero.")
        0 < γ <= 1 || error("γ = $(γ) must be in the interval (0,1].")
        if CASE == :γ1 &&  γ != 1
            error("Special case γ = 1 flagged but γ = $(γ)")
        end
        new{T,CASE}(κ, α, γ)
    end
end
ExponentialKernel(κ::BaseKernel{T}, α::T = one(T), γ::T = one(T)) where {T<:AbstractFloat} = ExponentialKernel{T, γ == 1 ? :γ1 : :Ø}(κ, α, γ)
ExponentialKernel(α::T = 1.0, γ::T = one(T)) where {T<:AbstractFloat} = ExponentialKernel(convert(Kernel{T}, SquaredDistanceKernel()), α, γ)

GaussianKernel(α::T = 1.0) where {T<:AbstractFloat} = ExponentialKernel(SquaredDistanceKernel(one(T)), α)
RadialBasisKernel(α::T = 1.0) where {T<:AbstractFloat} = ExponentialKernel(SquaredDistanceKernel(one(T)),α)
LaplacianKernel(α::T = 1.0) where {T<:AbstractFloat} = ExponentialKernel(SquaredDistanceKernel(one(T)),α, convert(T, 0.5))

function convert(::Type{ExponentialKernel{T}}, κ::ExponentialKernel) where {T<:AbstractFloat}
    ExponentialKernel(convert(Kernel{T}, κ.k), convert(T, κ.alpha), convert(T, κ.gamma))
end

phi(κ::ExponentialKernel{T}, x::Vector{T}, y::Vector{T}) where {T<:AbstractFloat} = exp.(-κ.alpha * sum(abs2, x-y)^κ.gamma)
phi(κ::ExponentialKernel{T,:γ1}, x::Vector{T}, y::Vector{T}) where {T<:AbstractFloat} = exp.(-κ.alpha * sum(abs2, x-y))

#==========================================================================
  Generic Kernel Matrix Functions
==========================================================================#

function init_pairwise(X::Matrix{T}, Y::Matrix{T}, is_trans::Bool) where {T<:AbstractFloat}
    n_dim = is_trans ? 2 : 1
    n = size(X, n_dim)
    m = size(Y, n_dim)
    return Array{T}(undef, n, m)
end

function kernelmatrix(κ::Kernel{T}, X::Matrix{T}, Y::Matrix{T}, is_trans::Bool = false) where {T<:AbstractFloat}
    kernelmatrix!(init_pairwise(X, Y, is_trans), κ, X, Y, is_trans)
end

#==========================================================================
  Base and Composite Kernel Matrix Functions
==========================================================================#

function kernelmatrix!(K::Matrix{T}, κ::BaseKernel{T}, X::Matrix{T}, Y::Matrix{T}, is_trans::Bool) where {T<:AbstractFloat}
    pairwise!(K, κ, X, Y, is_trans)
end

function kernelmatrix!(K::Matrix{T}, κ::CompositeKernel{T}, X::Matrix{T}, Y::Matrix{T}, is_trans::Bool) where {T<:AbstractFloat}
    pairwise!(K, κ, X, Y, is_trans)
end

#===================================================================================================
  Default Pairwise Computation
===================================================================================================#

# Initiate pairwise matrices

function init_pairwise(X::Matrix{T}, is_trans::Bool) where {T<:AbstractFloat}
    n = size(X, is_trans ? 2 : 1)
    Array(T, n, n)
end

# Pairwise definition

function pairwise!(K::Matrix{T}, κ::Kernel{T}, X::Matrix{T}, Y::Matrix{T}, is_trans::Bool) where {T<:AbstractFloat}
    if is_trans
        pairwise_XtYt!(K, κ, X, Y)
    else
        pairwise_XY!(K, κ, X, Y)
    end
end

function pairwise_XY!(K::Matrix{T}, κ::Kernel{T}, X::Matrix{T}, Y::Matrix{T}) where {T<:AbstractFloat}
    (n = size(X,1)) == size(K,1) || throw(DimensionMismatch("Dimension 1 of X must match dimension 1 of K."))
    (m = size(Y,1)) == size(K,2) || throw(DimensionMismatch("Dimension 1 of Y must match dimension 2 of K."))
    size(X,2) == size(Y,2) || throw(DimensionMismatch("Dimension 2 of Y must match dimension 2 of X."))
    for j = 1:m
        y = Y[j,:]
        for i = 1:n
            K[i,j] = phi(κ, X[i,:], y)
        end
    end
    K
end

function pairwise_XtYt!(K::Matrix{T}, κ::Kernel{T}, X::Matrix{T}, Y::Matrix{T}) where {T<:AbstractFloat}
    (n = size(X,2)) == size(K,1) || throw(DimensionMismatch("Dimension 2 of X must match dimension 1 of K."))
    (m = size(Y,2)) == size(K,2) || throw(DimensionMismatch("Dimension 2 of Y must match dimension 2 of K."))
    size(X,1) == size(Y,1) || throw(DimensionMismatch("Dimension 1 of Y must match dimension 1 of X."))
    for j = 1:m
        y = Y[:,j]
        for i = 1:n
            K[i,j] = phi(κ, X[:,i], y)
        end
    end
    K
end


### A Pluto.jl notebook ###
# v0.17.4

using Markdown
using InteractiveUtils

# ╔═╡ 31e93532-d915-4638-ac42-6d5a8b51472e
using PlutoUI

# ╔═╡ 9ba84189-80cd-4760-9e48-feca71f313c5
using DualNumbers

# ╔═╡ a37e3c58-bd22-4bc6-b003-bb8d350d49ed
using Unitful

# ╔═╡ b4b177dd-97db-4a6d-987c-64426b90bd51
bigbreak = HTML("<br>"^4);

# ╔═╡ 47e2cc32-8290-11ec-2253-e528a7b8198e
md"""
# Multiple dispatch ⭐
"""

# ╔═╡ ab103ac0-dba1-4f4f-91d8-3872f234436f
md"""### Starting with the familiar: Generic functions

If we write generic functions, Julia will figure out the types and try to make sense of it (if possible):
"""

# ╔═╡ 808b1e07-2912-4417-aef7-1b5de503e244
f(x) = x.^2

# ╔═╡ 849bd1c6-969b-479e-b519-f5d444bcb840
md"Hence, all of these inputs work:"

# ╔═╡ 67437820-a50b-49f2-af6b-7ba80277103e
(
	f_int    = f(1),
	f_array  = f([1,2]),
	f_string = f("hi")
)

# ╔═╡ e0fd7d97-187e-4937-9ca1-6473e167de4a
bigbreak

# ╔═╡ 19e8f83e-ca03-4677-9c6a-44f392d92a12
md"Indeed, under the hood, Julia does compile for all of these inputs specialised functions:"

# ╔═╡ 3e6b096d-58d3-46a6-8936-93186042c9b7
with_terminal() do
	@code_llvm f( (1.0, 1.0,1.0, 1.0) )
end

# ╔═╡ a213f381-52bf-4008-815e-b3c007b08b50
with_terminal() do
	@code_llvm f( [1.0, 1.0, 1.0] )
end

# ╔═╡ e6fc063f-19ef-4af4-b476-2c28ba161a93
md"""
### Short interlude

If you want to deep dive into how Julia works, it is often very interesting to check how Julia actually converts simple commands.
It needs some time to get used to the output, but it is worth it!

The `@code_*` commands are
```julia
@code_lowered f(x)   # Julia's internal representation after removing syntax sugar
@code_warntype f(x)  # highlighted lines mark type instabilities
@code_llvm f(x)      # the final product of the Julia JIT compiler
@code_native f(x)    # after optimisation with the LLVM for your machine
```

"""

# ╔═╡ 08714061-f90d-45f7-83ad-191b8d309e74

# ╔═╡ 2a255075-0bf6-4347-bf2a-1ded087d2514

# ╔═╡ 08d4c62e-1265-4925-8465-9ff1904057a3

# ╔═╡ 93051f8a-7802-4a2b-891b-46037a7f97b3
md"## Specifying the types

Sometimes the same function needs different implementation depending on the type:
"

# ╔═╡ 1c2b3808-ce65-4313-9380-58a42844eba9
begin
	# this function picks randomly between a and b:
	random(a, b) = rand([a,b])

	# integers between a and b
	random(a::Int64, b::Int64) = rand(a:b)

	# floats between a and b
	random(a::Float64, b::Float64) = a + (b-a)*rand()

	random() = 1
end

# ╔═╡ 1ff517d6-2145-4ab5-93e2-faeb0882dcdc
methods(random)

# ╔═╡ 4292041a-db6d-4235-809e-3bac423a7342
random(1.0, 2.0)

# ╔═╡ c0834dbb-0128-4219-a119-6df81bf876e4
random(1, 3)

# ╔═╡ a92b084e-ede0-408b-8bee-f8508fc2a470
random(1, 6.32)

# ╔═╡ 6682b551-f8da-4bc9-ad77-d41444ba6022

# ╔═╡ a9c019a5-a08e-467a-b635-fab0debd158d
md"#### Definition:

- **A function** (say `random`) is identified by the function name.
- The **different implementations** (say `random(::Int64,::Int64), ...`) of the function **are called methods**.
Hence, each function can have several methods.
"

# ╔═╡ 97036832-f4ca-4b34-b2da-6857439d529b

# ╔═╡ 88d8106c-2008-455b-a42a-fa3dd582c831
md"### → Julia always applies the most specific method for the given inputs!

These methods could overlap in crazy ways! _(It (usually) just works!™)_"

# ╔═╡ e53b709a-dd90-4a8c-bc59-781c685f4954

# ╔═╡ 2f381972-5972-4369-bb2d-67d01199eaf2

# ╔═╡ 2161f348-7efe-4cd0-b998-2d301fac22e4
md"## The key point of multiple dispatch

We want maximal control about how functions look for different types.

Multiple dispatch is
- **dynamic**
  - deciding at runtime what to do
- **modular**
  - reducing duplicate code
- **extensible**
  - allowing to add something to existing code
- **~ fast**
  - trying to keep amount of recomplilation minimal

(see: [Expression problem](https://en.wikipedia.org/wiki/Expression_problem))
"

# ╔═╡ 54eac40f-31aa-4df5-a1e6-3b863c12198f

# ╔═╡ 433fca14-845a-4267-b6be-b4ea1d9100c4
bigbreak

# ╔═╡ dd260866-8e10-4b36-a62f-2f930617e579
md"""
## The type hiearachy

How does Julia decide which method to call?

#### With a partial ordering `<:` on all types! 🧮 😎
"""

# ╔═╡ ba329815-e6bd-410f-82b4-bc7f89655661
Float64 <: Real <: Number

# ╔═╡ b7751fd9-6d65-4d5f-a570-4d71af818132
Int64 <: Integer <: Real <: Number

# ╔═╡ 72476e34-5ff3-48bf-848b-3e9be57faca3
Float64 <: Integer

# ╔═╡ 2431f98b-4e98-4528-acf6-65e27abb08c7
String <: AbstractString <: Any

# ╔═╡ f38fccdb-167c-44bd-8d9d-c3ae2fb79db7
Vector{Float64} <: Vector <: AbstractArray

# ╔═╡ ba7efcc8-479f-409b-846f-2696001d7a6b

# ╔═╡ 1ca91764-5f29-403e-8417-7d7cbb55042f
md"For each input with types `T1, T2, T3, ...`, Julia calls basically the method with input types

$\inf \{ \mathbf{S} \mid \text{method with argument types } \mathbf{S} \text{ exists and } \mathbf{T} <: \mathbf{S} \}$

with respect to the type ordering `<:`."

# ╔═╡ de4767b1-e882-4c55-9d22-0b91cc0d49d9
begin
	g(x::Number) = x
	g(x::Real) = x^2
	g(x::Float64) = x^3
end

# ╔═╡ 4d2f8eaa-fd3f-4b65-a804-7c1658a319c8
g( 1 + 2im )

# ╔═╡ 8d6623d2-60cf-43e7-b040-6787e8b7edad
g(3.0)

# ╔═╡ 352f54bf-3520-4f1c-9124-8d2871466e47
bigbreak

# ╔═╡ ce1a54c5-1ff0-4458-82f5-a1984b5be4fc
md"**Quick exercise:**

Find a case in which Julia cannot determine which method to call!"

# ╔═╡ 684ca6d9-ba49-446c-b394-3d4ca985e1d2
begin
	fnc(x::Int64,  y::Int64) = 1
	fnc(x::Number, y::Number) = 2
end

# ╔═╡ c76d5344-ebae-4bef-9d7c-09940be5c492
fnc(1,1.0)

# ╔═╡ b31fdae8-76f1-4998-945b-ca37f6173c0f

# ╔═╡ 82a2733a-da60-485f-8482-6189101933aa

# ╔═╡ 32a1a877-2146-4e8e-90c1-3b021b6d9348
md"**Note:** A case which is important to keep in mind is this one:"

# ╔═╡ 0a9e5987-869d-4b10-987d-4408769894c1
Vector{Int64} <: Vector

# ╔═╡ 9eb5aa13-7df1-4f2d-8f9b-d2b8b7e13a84
md"However, with so called _parametric types_ we can handle that:"

# ╔═╡ 57e75a99-1af9-452e-af11-81157551c12e
Vector{Int64} <: Vector{T} where T <: Integer

# ╔═╡ 61b1d73d-f7c3-4136-b92b-3a95b5143d2c
bigbreak

# ╔═╡ e97156d1-bea8-491f-b190-5c95d9a5b536
md"### How is that related to Object Oriented Programming?

> Yes and No: Multiple dispatch is more flexible 😃, but less ordered 😢.

In other words: OOP ensures that the programmer does not too stupid things. In Julia you can literally be a **type pirate** ☠.
"

# ╔═╡ c000aa68-b1ab-40ed-9376-32be9c51dc1a
md"
#### Short interlude: Let's be a type pirate ☠
_Run for example: `Base.:+(x::Float64, y::Float64) = 0.0` in a Julia REPL and compute `sin(1)`. (Afterwards, do something like that never again!)_
"

# ╔═╡ 4737b51a-539e-45ec-bf32-e3798d5ae5cd

# ╔═╡ c7a2a4bd-3e1d-438f-9270-725c3c9616f8
md" → You should usually only define methods for your custom types. Don't mess too much with other peoples types or it will backfire!"

# ╔═╡ 19412b07-a5ce-4c7a-b193-2b5e60c521f7
bigbreak

# ╔═╡ e4b689a9-f03a-4e28-8739-e018041d3c79
md"""
## Putting it all together: Being fast and generic is cool!

#### Example: (Derivative of solutions with respect to inital data)
"""

# ╔═╡ 3e537363-610e-40e0-8b8e-44e4713a17cd
md"Let's consider this simple implementation of a (scalar) explicit Euler method."

# ╔═╡ 4e965956-3633-4f4c-80ad-ff3f29958b54
function euler(f, x0::T, dt, n) where T
	x = Array{T}(undef, n)
	x[1] = x0

	for i in 2:n
		x[i] = x[i-1] + dt * f(x[i-1])
	end
	return x
end

# ╔═╡ ae2eedbd-6cb8-4356-924f-eed0926f388b
sol = euler( x -> -sin(x), 1.0, 0.1, 100)

# ╔═╡ 9228a5bc-84d5-4e5e-a971-211ac9b47fe3

# ╔═╡ 65c828b1-ff10-4428-afd9-ec1748054a5b
md"Now, let's see what happens if we plug in another type of `Numbers`!"

# ╔═╡ 71449f89-4c20-4eb6-b41b-5ac39f2ab829

# ╔═╡ c8639b54-fef0-485a-8e8d-10d90ab6ef95
md"### Automatic differentiation"

# ╔═╡ aba97bf0-2519-4ebb-8ead-d7abb4042adb
dual_sol = euler( x -> -sin(x), Dual(1.0, 1.0), 0.1, 100)

# ╔═╡ b3a96ab1-2285-4db3-b67e-040278da2312

# ╔═╡ 6890676a-37eb-4808-8967-5ac17ec286f4
md"##### It just works!™

We have implemented a method to compute the gradient of the numerical solution with respect to initial data. (aka Automatic Differentiation)"

# ╔═╡ 39072267-77e1-4b7e-9efc-f0a904ebf587
md"""
The solution $x(t, x_0)$ of

$\dot x = -\sin(x), \quad x(0) = x_0$

has approximately the derivative
$\frac{\partial x}{\partial x_0}(10, 1) ≈$ $(dualpart(dual_sol[end])).

_Notice that in order to do automatic differentiation, not a single line of the `euler` function had to change!_

"""

# ╔═╡ 6bf56af9-ad4f-417b-9356-9b582f9655e9

# ╔═╡ 91c8ea46-9fc4-4d94-94ec-0824555d4a2d
md"#### Units"

# ╔═╡ f5113d03-754f-4ace-91f3-27bf015cce9c
unit_sol = euler( x -> -x/u"hr", 1.0u"μm", 0.1u"s", 100)

# ╔═╡ 384f9598-abca-4727-8f26-037bc5f377bb

# ╔═╡ 23c8d96c-314d-40d8-a14e-a317378ad705

# ╔═╡ 59ca17c8-98bc-467b-9abc-495c978298d2

# ╔═╡ 5683bafe-96a6-404d-9ef3-4b493ac7e5e4

# ╔═╡ 7b92479f-2638-4bcc-822f-ce638a073d3c

# ╔═╡ 0a97a970-199e-4853-895e-fefc6d7753e6
md"### Unlimited possibilities to hijack this Euler method!"

# ╔═╡ fe5ad85b-f1d6-4d3e-a382-3d050c60c4cf
md"_Reminder:_ All of that compiled to almost C-like machine code!"

# ╔═╡ 8beec153-5a22-44b2-ae98-4cff82ec263c
bigbreak

# ╔═╡ e101a76d-d2ca-43c7-ab01-953c8397b7c9
md"""
# Exercises

7.1

A famous example for multiple dispatch is the
`Base.show(io, x)` function.

It is responsible for displaying Julia objects.
_(It works for text, numbers, images, html, ... you name it, then you can implement it)._

Overwrite the show function for `Tuple{Float64,Float64}` such that it displays
> 1 ❤ 2
when you call `print( (1,2) )`.

Hint: Use `print(io, ...)` to write something directly to input-output `io`.
"""

# ╔═╡ c1ed7bc2-2936-452c-b5b7-06de12adbe40
Base.show(io::IO, x::Tuple{Int64, Int64}) = nothing

# ╔═╡ 35f7e19f-bca1-4d5a-8f8d-dc4c4c77592e
with_terminal() do
	print( (1, 2+5) )
end

# ╔═╡ c79e90d9-d5e0-4f37-8abe-f49731a5d830
bigbreak

# ╔═╡ 716a277e-525b-4fed-81be-c8ceb50f9aec
md"""
7.2

Implement the product for `Tuple{Number,Number}` as
$(x,y) * (a,b) = xb - ya$

_Hint: The multiplication is called `Base.:*`_
"""

# ╔═╡ 8d48eb42-bbca-45c1-8c32-be9da3712835
Tup2f = Tuple{Number,Number} # if you want, you can use this type

# ╔═╡ 7e9c3310-e098-4e98-a02e-fdd905292f64
# enter your code here

# ╔═╡ 1fe523b2-f9b2-4a90-b45f-54dd9cb9dcd4

# ╔═╡ 214287e1-ac41-458a-aedf-16150f0b1da3
# (1.0,2) * (4,5) == -3.0  # test

# ╔═╡ 9b3fc16a-bee4-4917-92cf-41e64ab6ce39
bigbreak

# ╔═╡ e1befdce-9555-4a14-ad8e-48c64b4a74a6
md"""
7.3

Define a function `h` which returns the sum of integer arrays, but the product for all non-integer arrays.
"""

# ╔═╡ 838b3265-4724-4e56-95f9-56566862ffff
# enter your code here

# ╔═╡ 35e901dc-93f1-4927-a6c7-c01a6934ca83
# h([1, 2]) == 3
# h([1.0, 2.0]) == 2.0

# ╔═╡ 00000000-0000-0000-0000-000000000001
PLUTO_PROJECT_TOML_CONTENTS = """
[deps]
DualNumbers = "fa6b7ba4-c1ee-5f82-b5fc-ecf0adba8f74"
PlutoUI = "7f904dfe-b85e-4ff6-b463-dae2292396a8"
Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"

[compat]
DualNumbers = "~0.6.6"
PlutoUI = "~0.7.32"
Unitful = "~1.10.1"
"""

# ╔═╡ 00000000-0000-0000-0000-000000000002
PLUTO_MANIFEST_TOML_CONTENTS = """
# This file is machine-generated - editing it directly is not advised

julia_version = "1.7.1"
manifest_format = "2.0"

[[deps.AbstractPlutoDingetjes]]
deps = ["Pkg"]
git-tree-sha1 = "8eaf9f1b4921132a4cff3f36a1d9ba923b14a481"
uuid = "6e696c72-6542-2067-7265-42206c756150"
version = "1.1.4"

[[deps.ArgTools]]
uuid = "0dad84c5-d112-42e6-8d28-ef12dabb789f"

[[deps.Artifacts]]
uuid = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"

[[deps.Base64]]
uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"

[[deps.Calculus]]
deps = ["LinearAlgebra"]
git-tree-sha1 = "f641eb0a4f00c343bbc32346e1217b86f3ce9dad"
uuid = "49dc2e85-a5d0-5ad3-a950-438e2897f1b9"
version = "0.5.1"

[[deps.ChainRulesCore]]
deps = ["Compat", "LinearAlgebra", "SparseArrays"]
git-tree-sha1 = "f9982ef575e19b0e5c7a98c6e75ee496c0f73a93"
uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
version = "1.12.0"

[[deps.ChangesOfVariables]]
deps = ["ChainRulesCore", "LinearAlgebra", "Test"]
git-tree-sha1 = "bf98fa45a0a4cee295de98d4c1462be26345b9a1"
uuid = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0"
version = "0.1.2"

[[deps.ColorTypes]]
deps = ["FixedPointNumbers", "Random"]
git-tree-sha1 = "024fe24d83e4a5bf5fc80501a314ce0d1aa35597"
uuid = "3da002f7-5984-5a60-b8a6-cbb66c0b333f"
version = "0.11.0"

[[deps.Compat]]
deps = ["Base64", "Dates", "DelimitedFiles", "Distributed", "InteractiveUtils", "LibGit2", "Libdl", "LinearAlgebra", "Markdown", "Mmap", "Pkg", "Printf", "REPL", "Random", "SHA", "Serialization", "SharedArrays", "Sockets", "SparseArrays", "Statistics", "Test", "UUIDs", "Unicode"]
git-tree-sha1 = "44c37b4636bc54afac5c574d2d02b625349d6582"
uuid = "34da2185-b29b-5c13-b0c7-acf172513d20"
version = "3.41.0"

[[deps.CompilerSupportLibraries_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "e66e0078-7015-5450-92f7-15fbd957f2ae"

[[deps.ConstructionBase]]
deps = ["LinearAlgebra"]
git-tree-sha1 = "f74e9d5388b8620b4cee35d4c5a618dd4dc547f4"
uuid = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
version = "1.3.0"

[[deps.Dates]]
deps = ["Printf"]
uuid = "ade2ca70-3891-5945-98fb-dc099432e06a"

[[deps.DelimitedFiles]]
deps = ["Mmap"]
uuid = "8bb1440f-4735-579b-a4ab-409b98df4dab"

[[deps.Distributed]]
deps = ["Random", "Serialization", "Sockets"]
uuid = "8ba89e20-285c-5b6f-9357-94700520ee1b"

[[deps.DocStringExtensions]]
deps = ["LibGit2"]
git-tree-sha1 = "b19534d1895d702889b219c382a6e18010797f0b"
uuid = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
version = "0.8.6"

[[deps.Downloads]]
deps = ["ArgTools", "LibCURL", "NetworkOptions"]
uuid = "f43a241f-c20a-4ad4-852c-f6b1247861c6"

[[deps.DualNumbers]]
deps = ["Calculus", "NaNMath", "SpecialFunctions"]
git-tree-sha1 = "84f04fe68a3176a583b864e492578b9466d87f1e"
uuid = "fa6b7ba4-c1ee-5f82-b5fc-ecf0adba8f74"
version = "0.6.6"

[[deps.FixedPointNumbers]]
deps = ["Statistics"]
git-tree-sha1 = "335bfdceacc84c5cdf16aadc768aa5ddfc5383cc"
uuid = "53c48c17-4a7d-5ca2-90c5-79b7896eea93"
version = "0.8.4"

[[deps.Hyperscript]]
deps = ["Test"]
git-tree-sha1 = "8d511d5b81240fc8e6802386302675bdf47737b9"
uuid = "47d2ed2b-36de-50cf-bf87-49c2cf4b8b91"
version = "0.0.4"

[[deps.HypertextLiteral]]
git-tree-sha1 = "2b078b5a615c6c0396c77810d92ee8c6f470d238"
uuid = "ac1192a8-f4b3-4bfe-ba22-af5b92cd3ab2"
version = "0.9.3"

[[deps.IOCapture]]
deps = ["Logging", "Random"]
git-tree-sha1 = "f7be53659ab06ddc986428d3a9dcc95f6fa6705a"
uuid = "b5f81e59-6552-4d32-b1f0-c071b021bf89"
version = "0.2.2"

[[deps.InteractiveUtils]]
deps = ["Markdown"]
uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240"

[[deps.InverseFunctions]]
deps = ["Test"]
git-tree-sha1 = "a7254c0acd8e62f1ac75ad24d5db43f5f19f3c65"
uuid = "3587e190-3f89-42d0-90ee-14403ec27112"
version = "0.1.2"

[[deps.IrrationalConstants]]
git-tree-sha1 = "7fd44fd4ff43fc60815f8e764c0f352b83c49151"
uuid = "92d709cd-6900-40b7-9082-c6be49f344b6"
version = "0.1.1"

[[deps.JLLWrappers]]
deps = ["Preferences"]
git-tree-sha1 = "abc9885a7ca2052a736a600f7fa66209f96506e1"
uuid = "692b3bcd-3c85-4b1f-b108-f13ce0eb3210"
version = "1.4.1"

[[deps.JSON]]
deps = ["Dates", "Mmap", "Parsers", "Unicode"]
git-tree-sha1 = "8076680b162ada2a031f707ac7b4953e30667a37"
uuid = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
version = "0.21.2"

[[deps.LibCURL]]
deps = ["LibCURL_jll", "MozillaCACerts_jll"]
uuid = "b27032c2-a3e7-50c8-80cd-2d36dbcbfd21"

[[deps.LibCURL_jll]]
deps = ["Artifacts", "LibSSH2_jll", "Libdl", "MbedTLS_jll", "Zlib_jll", "nghttp2_jll"]
uuid = "deac9b47-8bc7-5906-a0fe-35ac56dc84c0"

[[deps.LibGit2]]
deps = ["Base64", "NetworkOptions", "Printf", "SHA"]
uuid = "76f85450-5226-5b5a-8eaa-529ad045b433"

[[deps.LibSSH2_jll]]
deps = ["Artifacts", "Libdl", "MbedTLS_jll"]
uuid = "29816b5a-b9ab-546f-933c-edad1886dfa8"

[[deps.Libdl]]
uuid = "8f399da3-3557-5675-b5ff-fb832c97cbdb"

[[deps.LinearAlgebra]]
deps = ["Libdl", "libblastrampoline_jll"]
uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

[[deps.LogExpFunctions]]
deps = ["ChainRulesCore", "ChangesOfVariables", "DocStringExtensions", "InverseFunctions", "IrrationalConstants", "LinearAlgebra"]
git-tree-sha1 = "e5718a00af0ab9756305a0392832c8952c7426c1"
uuid = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
version = "0.3.6"

[[deps.Logging]]
uuid = "56ddb016-857b-54e1-b83d-db4d58db5568"

[[deps.Markdown]]
deps = ["Base64"]
uuid = "d6f4376e-aef5-505a-96c1-9c027394607a"

[[deps.MbedTLS_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "c8ffd9c3-330d-5841-b78e-0817d7145fa1"

[[deps.Mmap]]
uuid = "a63ad114-7e13-5084-954f-fe012c677804"

[[deps.MozillaCACerts_jll]]
uuid = "14a3606d-f60d-562e-9121-12d972cd8159"

[[deps.NaNMath]]
git-tree-sha1 = "b086b7ea07f8e38cf122f5016af580881ac914fe"
uuid = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
version = "0.3.7"

[[deps.NetworkOptions]]
uuid = "ca575930-c2e3-43a9-ace4-1e988b2c1908"

[[deps.OpenBLAS_jll]]
deps = ["Artifacts", "CompilerSupportLibraries_jll", "Libdl"]
uuid = "4536629a-c528-5b80-bd46-f80d51c5b363"

[[deps.OpenLibm_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "05823500-19ac-5b8b-9628-191a04bc5112"

[[deps.OpenSpecFun_jll]]
deps = ["Artifacts", "CompilerSupportLibraries_jll", "JLLWrappers", "Libdl", "Pkg"]
git-tree-sha1 = "13652491f6856acfd2db29360e1bbcd4565d04f1"
uuid = "efe28fd5-8261-553b-a9e1-b2916fc3738e"
version = "0.5.5+0"

[[deps.Parsers]]
deps = ["Dates"]
git-tree-sha1 = "0b5cfbb704034b5b4c1869e36634438a047df065"
uuid = "69de0a69-1ddd-5017-9359-2bf0b02dc9f0"
version = "2.2.1"

[[deps.Pkg]]
deps = ["Artifacts", "Dates", "Downloads", "LibGit2", "Libdl", "Logging", "Markdown", "Printf", "REPL", "Random", "SHA", "Serialization", "TOML", "Tar", "UUIDs", "p7zip_jll"]
uuid = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"

[[deps.PlutoUI]]
deps = ["AbstractPlutoDingetjes", "Base64", "ColorTypes", "Dates", "Hyperscript", "HypertextLiteral", "IOCapture", "InteractiveUtils", "JSON", "Logging", "Markdown", "Random", "Reexport", "UUIDs"]
git-tree-sha1 = "ae6145ca68947569058866e443df69587acc1806"
uuid = "7f904dfe-b85e-4ff6-b463-dae2292396a8"
version = "0.7.32"

[[deps.Preferences]]
deps = ["TOML"]
git-tree-sha1 = "2cf929d64681236a2e074ffafb8d568733d2e6af"
uuid = "21216c6a-2e73-6563-6e65-726566657250"
version = "1.2.3"

[[deps.Printf]]
deps = ["Unicode"]
uuid = "de0858da-6303-5e67-8744-51eddeeeb8d7"

[[deps.REPL]]
deps = ["InteractiveUtils", "Markdown", "Sockets", "Unicode"]
uuid = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"

[[deps.Random]]
deps = ["SHA", "Serialization"]
uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"

[[deps.Reexport]]
git-tree-sha1 = "45e428421666073eab6f2da5c9d310d99bb12f9b"
uuid = "189a3867-3050-52da-a836-e630ba90ab69"
version = "1.2.2"

[[deps.SHA]]
uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce"

[[deps.Serialization]]
uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b"

[[deps.SharedArrays]]
deps = ["Distributed", "Mmap", "Random", "Serialization"]
uuid = "1a1011a3-84de-559e-8e89-a11a2f7dc383"

[[deps.Sockets]]
uuid = "6462fe0b-24de-5631-8697-dd941f90decc"

[[deps.SparseArrays]]
deps = ["LinearAlgebra", "Random"]
uuid = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"

[[deps.SpecialFunctions]]
deps = ["ChainRulesCore", "IrrationalConstants", "LogExpFunctions", "OpenLibm_jll", "OpenSpecFun_jll"]
git-tree-sha1 = "e6bf188613555c78062842777b116905a9f9dd49"
uuid = "276daf66-3868-5448-9aa4-cd146d93841b"
version = "2.1.0"

[[deps.Statistics]]
deps = ["LinearAlgebra", "SparseArrays"]
uuid = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"

[[deps.TOML]]
deps = ["Dates"]
uuid = "fa267f1f-6049-4f14-aa54-33bafae1ed76"

[[deps.Tar]]
deps = ["ArgTools", "SHA"]
uuid = "a4e569a6-e804-4fa4-b0f3-eef7a1d5b13e"

[[deps.Test]]
deps = ["InteractiveUtils", "Logging", "Random", "Serialization"]
uuid = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

[[deps.UUIDs]]
deps = ["Random", "SHA"]
uuid = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"

[[deps.Unicode]]
uuid = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5"

[[deps.Unitful]]
deps = ["ConstructionBase", "Dates", "LinearAlgebra", "Random"]
git-tree-sha1 = "b95e0b8a8d1b6a6c3e0b3ca393a7a285af47c264"
uuid = "1986cc42-f94f-5a68-af5c-568840ba703d"
version = "1.10.1"

[[deps.Zlib_jll]]
deps = ["Libdl"]
uuid = "83775a58-1f1d-513f-b197-d71354ab007a"

[[deps.libblastrampoline_jll]]
deps = ["Artifacts", "Libdl", "OpenBLAS_jll"]
uuid = "8e850b90-86db-534c-a0d3-1478176c7d93"

[[deps.nghttp2_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "8e850ede-7688-5339-a07c-302acd2aaf8d"

[[deps.p7zip_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "3f19e933-33d8-53b3-aaab-bd5110c3b7a0"
"""

# ╔═╡ Cell order:
# ╟─31e93532-d915-4638-ac42-6d5a8b51472e
# ╟─b4b177dd-97db-4a6d-987c-64426b90bd51
# ╟─47e2cc32-8290-11ec-2253-e528a7b8198e
# ╟─ab103ac0-dba1-4f4f-91d8-3872f234436f
# ╠═808b1e07-2912-4417-aef7-1b5de503e244
# ╟─849bd1c6-969b-479e-b519-f5d444bcb840
# ╠═67437820-a50b-49f2-af6b-7ba80277103e
# ╟─e0fd7d97-187e-4937-9ca1-6473e167de4a
# ╟─19e8f83e-ca03-4677-9c6a-44f392d92a12
# ╠═3e6b096d-58d3-46a6-8936-93186042c9b7
# ╠═a213f381-52bf-4008-815e-b3c007b08b50
# ╟─e6fc063f-19ef-4af4-b476-2c28ba161a93
# ╟─08714061-f90d-45f7-83ad-191b8d309e74
# ╟─2a255075-0bf6-4347-bf2a-1ded087d2514
# ╟─08d4c62e-1265-4925-8465-9ff1904057a3
# ╟─93051f8a-7802-4a2b-891b-46037a7f97b3
# ╠═1c2b3808-ce65-4313-9380-58a42844eba9
# ╟─1ff517d6-2145-4ab5-93e2-faeb0882dcdc
# ╠═4292041a-db6d-4235-809e-3bac423a7342
# ╠═c0834dbb-0128-4219-a119-6df81bf876e4
# ╠═a92b084e-ede0-408b-8bee-f8508fc2a470
# ╟─6682b551-f8da-4bc9-ad77-d41444ba6022
# ╟─a9c019a5-a08e-467a-b635-fab0debd158d
# ╟─97036832-f4ca-4b34-b2da-6857439d529b
# ╟─88d8106c-2008-455b-a42a-fa3dd582c831
# ╟─e53b709a-dd90-4a8c-bc59-781c685f4954
# ╟─2f381972-5972-4369-bb2d-67d01199eaf2
# ╟─2161f348-7efe-4cd0-b998-2d301fac22e4
# ╟─54eac40f-31aa-4df5-a1e6-3b863c12198f
# ╟─433fca14-845a-4267-b6be-b4ea1d9100c4
# ╟─dd260866-8e10-4b36-a62f-2f930617e579
# ╠═ba329815-e6bd-410f-82b4-bc7f89655661
# ╠═b7751fd9-6d65-4d5f-a570-4d71af818132
# ╠═72476e34-5ff3-48bf-848b-3e9be57faca3
# ╠═2431f98b-4e98-4528-acf6-65e27abb08c7
# ╠═f38fccdb-167c-44bd-8d9d-c3ae2fb79db7
# ╟─ba7efcc8-479f-409b-846f-2696001d7a6b
# ╟─1ca91764-5f29-403e-8417-7d7cbb55042f
# ╠═de4767b1-e882-4c55-9d22-0b91cc0d49d9
# ╠═4d2f8eaa-fd3f-4b65-a804-7c1658a319c8
# ╠═8d6623d2-60cf-43e7-b040-6787e8b7edad
# ╟─352f54bf-3520-4f1c-9124-8d2871466e47
# ╟─ce1a54c5-1ff0-4458-82f5-a1984b5be4fc
# ╠═684ca6d9-ba49-446c-b394-3d4ca985e1d2
# ╠═c76d5344-ebae-4bef-9d7c-09940be5c492
# ╟─b31fdae8-76f1-4998-945b-ca37f6173c0f
# ╟─82a2733a-da60-485f-8482-6189101933aa
# ╟─32a1a877-2146-4e8e-90c1-3b021b6d9348
# ╠═0a9e5987-869d-4b10-987d-4408769894c1
# ╟─9eb5aa13-7df1-4f2d-8f9b-d2b8b7e13a84
# ╠═57e75a99-1af9-452e-af11-81157551c12e
# ╟─61b1d73d-f7c3-4136-b92b-3a95b5143d2c
# ╟─e97156d1-bea8-491f-b190-5c95d9a5b536
# ╟─c000aa68-b1ab-40ed-9376-32be9c51dc1a
# ╟─4737b51a-539e-45ec-bf32-e3798d5ae5cd
# ╟─c7a2a4bd-3e1d-438f-9270-725c3c9616f8
# ╟─19412b07-a5ce-4c7a-b193-2b5e60c521f7
# ╟─e4b689a9-f03a-4e28-8739-e018041d3c79
# ╟─3e537363-610e-40e0-8b8e-44e4713a17cd
# ╠═4e965956-3633-4f4c-80ad-ff3f29958b54
# ╠═ae2eedbd-6cb8-4356-924f-eed0926f388b
# ╟─9228a5bc-84d5-4e5e-a971-211ac9b47fe3
# ╟─65c828b1-ff10-4428-afd9-ec1748054a5b
# ╟─71449f89-4c20-4eb6-b41b-5ac39f2ab829
# ╟─c8639b54-fef0-485a-8e8d-10d90ab6ef95
# ╠═9ba84189-80cd-4760-9e48-feca71f313c5
# ╠═aba97bf0-2519-4ebb-8ead-d7abb4042adb
# ╟─b3a96ab1-2285-4db3-b67e-040278da2312
# ╟─6890676a-37eb-4808-8967-5ac17ec286f4
# ╟─39072267-77e1-4b7e-9efc-f0a904ebf587
# ╟─6bf56af9-ad4f-417b-9356-9b582f9655e9
# ╟─91c8ea46-9fc4-4d94-94ec-0824555d4a2d
# ╠═a37e3c58-bd22-4bc6-b003-bb8d350d49ed
# ╠═f5113d03-754f-4ace-91f3-27bf015cce9c
# ╟─384f9598-abca-4727-8f26-037bc5f377bb
# ╠═23c8d96c-314d-40d8-a14e-a317378ad705
# ╠═59ca17c8-98bc-467b-9abc-495c978298d2
# ╟─5683bafe-96a6-404d-9ef3-4b493ac7e5e4
# ╟─7b92479f-2638-4bcc-822f-ce638a073d3c
# ╟─0a97a970-199e-4853-895e-fefc6d7753e6
# ╟─fe5ad85b-f1d6-4d3e-a382-3d050c60c4cf
# ╟─8beec153-5a22-44b2-ae98-4cff82ec263c
# ╟─e101a76d-d2ca-43c7-ab01-953c8397b7c9
# ╠═c1ed7bc2-2936-452c-b5b7-06de12adbe40
# ╠═35f7e19f-bca1-4d5a-8f8d-dc4c4c77592e
# ╟─c79e90d9-d5e0-4f37-8abe-f49731a5d830
# ╟─716a277e-525b-4fed-81be-c8ceb50f9aec
# ╠═8d48eb42-bbca-45c1-8c32-be9da3712835
# ╠═7e9c3310-e098-4e98-a02e-fdd905292f64
# ╟─1fe523b2-f9b2-4a90-b45f-54dd9cb9dcd4
# ╠═214287e1-ac41-458a-aedf-16150f0b1da3
# ╟─9b3fc16a-bee4-4917-92cf-41e64ab6ce39
# ╟─e1befdce-9555-4a14-ad8e-48c64b4a74a6
# ╠═838b3265-4724-4e56-95f9-56566862ffff
# ╠═35e901dc-93f1-4927-a6c7-c01a6934ca83
# ╟─00000000-0000-0000-0000-000000000001
# ╟─00000000-0000-0000-0000-000000000002


using ParamLevelSet
using jInv.Mesh
using SparseArrays
using LinearAlgebra
using Test

plotting = false;
if plotting
	using jInvVis
	using PyPlot
end

println("~~~~~~~~~~~~~~~~Testing Mesh Free RBFs.~~~~~~~~~~~~~~~~~~")

n = [32,32,32];
Mesh = getRegularMesh([0.0;3.0;0.0;3.0;0.0;3.0],n);
alpha = [1.5;2.5;-2.0;-1.0];
beta = [2.5;2.0;-1.5;2.5];
Xs = [0.5 0.5 0.5 ; 2.0 2.0 2.0; 1.2 2.3 1.5; 2.2 1.5 2.0];
m = wrapTheta(alpha,beta,Xs);

u, = MeshFreeParamLevelSetModelFunc(Mesh,m;computeJacobian=0);
if plotting
	figure()
	u = reshape(u,tuple(Mesh.n...));
	plotModel(u)
end

sigmaH = getDefaultHeaviside();

us, = MeshFreeParamLevelSetModelFunc(Mesh,m;computeJacobian=0,sigma = sigmaH);
if plotting
	figure()
	plotModel(us)
end

bf = 1;
dm = 0.01*randn(length(alpha)*5);

hh = 1.0;
u0,JBuilder = MeshFreeParamLevelSetModelFunc(Mesh,m;computeJacobian=1,bf = bf);
J0 = getSparseMatrix(JBuilder);
for k=0:10
	hhh = (0.5^k)*hh;
	ut = MeshFreeParamLevelSetModelFunc(Mesh,m+ hhh*dm;computeJacobian=0,bf = bf)[1];
	println("norm(ut-u0): ",norm(ut[:]-u0[:]),", norm(ut - u0 - J0*dm): ",norm(ut[:] - u0[:] - J0*hhh*dm));
end

println("With Heaviside func");

hh = 1.0;

u0,JBuilder = MeshFreeParamLevelSetModelFunc(Mesh,m;computeJacobian=1,sigma=sigmaH,bf = bf);
J0 = getSparseMatrix(JBuilder);

for k=0:10
	hhh = (0.5^k)*hh;
    ut = MeshFreeParamLevelSetModelFunc(Mesh,m + hhh*dm;computeJacobian=0,sigma=sigmaH,bf = bf)[1];
	println("norm(ut-u0): ",norm(ut[:]-u0[:]),", norm(ut - u0 - J0*dm): ",norm(ut[:] - u0[:] - J0*hhh*dm));
end


include("utils.jl")

@testset "Test Notation" begin
    M = Sphere(2)
    p1 = [1.0, 0.0, 0.0]
    @test (p1 ∈ M) == is_manifold_point(M, p1)
    @test (2 * p1 ∈ M) == is_manifold_point(M, 2 * p1)
    X1 = [0.0, 1.0, 0.0]
    X2 = [1.0, 0.0, 0.0]
    TpM = TangentSpaceAtPoint(M, p1)
    @test (X1 ∈ TpM) == is_tangent_vector(M, p1, X1)
    @test (X2 ∈ TpM) == is_tangent_vector(M, p1, X2)
end


include("distribution.jl")

export MI, PI, CMI

ϵ = 0.0000001

# predictive information
PI(data::Vector{Int64}; base=2, mode="ML") = MI(hcat(data[1:end-1], data[2:end]), base=base, mode=mode)

function MI_ML(data::Matrix{Int64}; base=2)
  max = maximum(data)
  px  = fe1p(data[:,1])
  py  = fe1p(data[:,2])
  pxy = fe2p(hcat(data[:,[1:2]]))

  r = 0
  for x=1:length(px)
    for y=1:length(py)
      if px[x] != 0.0 && py[y] != 0.0 && pxy[x,y] > 0.0
        r = r + pxy[x,y] * (log(base, pxy[x,y]) - log(base, px[x]*py[y]))
      end
    end
  end
  r
end

# mutual information
function MI(data::Matrix{Int64}; base=2, mode="ML", pseudocount=0)

  modes      = ["ML", "Maximum Likelihood"]
  umodes     = map(x->uppercase(x), modes)
  known_mode = uppercase(mode) in umodes
  pmodes     = map(x->", $x",modes)
  pmodes     = foldl(*, pmodes[1][3:end], pmodes[2:end])
  @assert known_mode "Mode may be any of the following: [$pmodes]."

  if uppercase(mode) in umodes[1:2]
    MI_ML(data, base=base)
  end

end

# MI(X;Y) = ∑ p(x,y) * log( p(x,y) / (p(x) * p(y)))
function MI(pxy::Matrix{Float64}; base=2)
  px = sum(pxy,1)
  py = sum(pxy,2)
  r  = 0
  for x = 1:size(pxy)[1]
    for y = 1:size(pxy)[2]
      if abs(px[x]) > ϵ && abs(py[y]) > ϵ && abs(pxy[x,y]) > ϵ
        r = r + pxy[x,y] * (log(base, pxy[x,y]) - log(base, px[x] * py[y]))
      end
    end
  end
  return r
end

# MI(X;Y|Z) = ∑ p(x,y,z) * log( p(x,y|z) / (p(x|z) * p(y|z)))
function CMI(pxyz::Array{Float64,3}; base=2)
  r = 0

  pz  = sum(pxyz, (1,2))
  pxz = sum(pxyz, 2)
  pyz = sum(pxyz, 1)

  for x = 1:size(pxyz)[1]
    for y = 1:size(pxyz)[2]
      for z = 1:size(pxyz)[3]
        if abs(pxyz[x,y,z]) > ϵ && abs(pz[1,1,z]) > ϵ && abs(pxz[x,1,z]) > ϵ && abs(pyz[1,y,z]) > ϵ
           r = r + pxyz[x,y,z] * ( log(base, pxyz[x,y,z] / pz[1,1,z]) - log(base, pxz[x,1,z] / pz[1,1,z] * pyz[1,y,z] / pz[1,1,z]))
         end
      end
    end
  end
  return r
end


module MCMC_LR_Tests

using ArgCheck
using Distributions
using MCMCDiagnostics

export mean_LR_pvalue, cov_LR_pvalue

"""
    column_ρs(X)

For observations in the rows of `X`, return the estimated `ρ=ESS/n` for each column (variable).

Provides the default value to [`mean_LR_statistic`](@ref), otherwise only useful for debugging.
"""
column_ρs(X::AbstractMatrix) = vec(mapslices(first ∘ ess_factor_estimate, X, 1))

"""
    mean_LR_pvalue(μ, X, [ρ])

P-value for a likelihood ratio statistic for testing that the columnwise mean of `X` are equal to the vector `μ`, if the rows of `X` are normally distributed (with unknown variance, estimated from the sample).

Assume that some form of a CLT holds, so

``n(x̄(n)-x̄) → N(0, Σ)`` in distribution, where `x̄(n)` is the running mean.

The test statistic is

``(x̄(n)-μ)' Σ⁻¹ (x̄(n)-μ) ∼ T²(k, n-1)``

where `T²` is [Hotelling's T²-distribution](https://en.wikipedia.org/wiki/Hotelling%27s_T-squared_distribution), if `x̄ = μ` (Dootika et al 2017).

For IID samples, `Σ` can be estimated from the sample. However, MCMC draws are autocorrelated, so this is more difficult. Let ρᵢ = ESSᵢ/n denote the scaling factor to obtain the effective sample size. For a crude approximation, we rescale the variance by ρᵢ, and covariances by √(ρᵢρⱼ).
"""
function mean_LR_pvalue(μ::AbstractVector, X::AbstractMatrix, ρ = column_ρs(X))
    N, K = size(X)
    @argcheck K == length(μ) == length(ρ)
    R = Diagonal(.√ρ)           # correction by ρ = ESS/n
    S = cov(X, 1, false)        # sample covariance
    δ = vec(mean(X, 1)) - μ     # deviations from μ
    h = (δ' * (R*inv(S)*R) * δ) * N    # has Hotelling T²(K, N-1) distribution
    ccdf(FDist(K, N-K), h*((N-K)/((N-1)*K))) # convert to F(K, N-K) distribution
end

"""
    cov_LR_pvalue(Σ, X)

P-value for likelihood ratio statistic for comparing the covariance of `X` (with observations in rows) to `Σ`.

Note that this is an *asymptotic* statistic, so it has better coverage when `N ≫ K`. Coverage may be very bad when `N` is small, see Bai et al (2009).
"""
function cov_LR_pvalue(Σ::AbstractMatrix, X::AbstractMatrix)
    N, K = size(X)
    @argcheck size(Σ) == (K, K)
    U = chol(Symmetric(Σ))
    S = cov(X / U, 1)
    L = trace(S) - logdet(S) - K
    ccdf(Chisq(K*(K+1)/2), N*L)
end

end # module


module DEACfrequency
using Random;
using JLD;
const functionNames = ["exponential_frequency_bins",
                       "fixed_frequency_bins",
                       "load_frequency_bins"]

function exponential_frequency_bins(start::Float64,stop::Float64,length::Int64;epsilon::Float64=1e-5)
    frequency_bins = Array{Float64,1}(undef,length + 1);
    exponential_frequency_bins!(frequency_bins,start,stop,epsilon=epsilon);
    frequency_bins
end

function exponential_frequency_bins!(frequency_bins::Array{Float64,1},start::Float64,stop::Float64;epsilon::Float64=1e-5)
    length = size(frequency_bins,1) - 1;
    frequency_bins[1] = start;
    _logspace = range(log10(start+epsilon), stop=log10(stop), length=length)
    @inbounds for i in 2:length+1
        frequency_bins[i] = exp10(_logspace[i-1]);
    end
    nothing
end

function fixed_frequency_bins(start::Float64,stop::Float64,length::Int64;epsilon::Float64=1e-5)
    frequency_bins = Array{Float64,1}(undef,length+1);
    fixed_frequency_bins!(frequency_bins,start,stop);
    frequency_bins
end

function fixed_frequency_bins!(frequency_bins::Array{Float64,1},start::Float64,stop::Float64;epsilon::Float64=1e-5)
    length = size(frequency_bins,1);
    _linspace = range(start,stop=stop,length=length);
    for i in 1:length
        frequency_bins[i]=_linspace[i];
    end
    nothing
end

function load_frequency_bins(start::Float64,stop::Float64,length::Int64;filename::String="./frequency.jld")
    frequency_bins = load_frequency_bins(filename);
    frequency_bins
end

function load_frequency_bins(filename::String="./frequency.jld")
    frequency_bins = JLD.load(filename,"frequency_bins");
    frequency_bins
end

end


#==============================================================================
This script reconstructs the binary matrix and frequency vector from experimental
measurement data stored in MAT file realdata.mat.

In addition, the dependence between the frequency vector and the reconstruction
error is examined with synthetic data which is generated with the given binary
matrix, different frequency vectors, and additive Gaussian noise.

Parallelization can be enabled, e.g., by starting julia with "julia -p 16".
Additional packages may be needed to install for the visualization.

This is the set of experiments presented (with nT=200 and higher resolution) in
Section 4.4 of the article:

L. Mustonen, X. Gao, A. Santana, R. Mitchell, Y. Vigfusson, and L. Ruthotto,
A Bayesian framework for molecular strain identification from mixed diagnostic samples
https://arxiv.org/abs/1803.02916
==============================================================================#

@everywhere using StrainRecon
using PyPlot
using jInv.Utils
using jInv.Mesh
using jInvVis
using MAT
srand(2) # for reproducibility

# load data from MAT file
data = matread("realdata.mat")
M_true = data["M_true"]
w_true = vec(data["w_true"])
d_meas = vec(data["d_meas"])

m = length(d_meas) # number of measurements
n = 3 # number of strains in the reconstruction
R = 10 # variance of the measurement noise is assumed to be 0.01
param = getStrainReconParam(d_meas,n,R=R)
nT = 200 # number of trials in BCD

# compute sample standard deviation of the error
errstd = sqrt(var(M_true*w_true-d_meas))
println("Sample standard deviation of the measurement error is $errstd")

# compute MAP estimates
MrecoBCD,wrecoBCD = solveMAP(param,maxTry=nT)
MrecoMIQP,wrecoMIQP = solveConvexMAP(param)

# compute posterior means and variances
Mcm,wcm,dummy,Mvar,wvar = wrapperIntegratePosterior(param)

# show results
imcolor = "gray"
imlimits = [.5, m+.5, n+.5, .5]

subplot(5,1,1)
imshow(M_true',cmap=imcolor,extent=imlimits,vmin=0,vmax=1)
title("true M")
println("true w:")
println(w_true)
println()

subplot(5,1,2)
imshow(MrecoBCD',cmap=imcolor,extent=imlimits,vmin=0,vmax=1)
title("M from BCD")
println("w from BCD:")
println(wrecoBCD)
println()

subplot(5,1,3)
imshow(MrecoMIQP',cmap=imcolor,extent=imlimits,vmin=0,vmax=1)
title("M from MIQP")
println("w from MIQP:")
println(wrecoMIQP)
println()

subplot(5,1,4)
imshow(Mcm',cmap=imcolor,extent=imlimits,vmin=0,vmax=1)
title("posterior mean of M")
println("posterior mean of w:")
println(wcm)
println()

subplot(5,1,5)
imshow(sqrt.(Mvar)',cmap=imcolor,extent=imlimits,vmin=0,vmax=1)
title("posterior standard deviation of M")
println("posterior standard deviation of w:")
println(sqrt.(wvar))
println()

# keep M fixed, vary w and reconstruct from synthetic, noisy measurements
noiseVar_range = [1e-2, 1e-4]
nT = 20 # decrease the number of trials for efficiency
resolution = 64
Ma = getRegularMesh([1/3 1 0 1/2 ],[resolution resolution])
as = getCellCenteredGrid(Ma)
recoerrs = SharedArray{Float64}(Ma.nc,2)
recoerrs[:] = NaN # white background

for noiseVar_idx = 1:2
    @sync @parallel for k=1:size(as,1)
        at = vec([ vec(as[k,:]); 1-sum(as[k,:])])
        if all(at.>=0) && sum(at)<=1 && all(diff(at).<=0)
            d = M_true * at + sqrt(noiseVar_range[noiseVar_idx])*randn(m)
            param = getStrainReconParam(d,n)
            Mreco,wreco = solveMAP(param,maxTry=nT)
            recoerrs[k,noiseVar_idx] = srDistance(BitArray(M_true),at,BitArray(Mreco),wreco,pn=1.0)
        end
    end
end

figure()
subplot(1,2,1)
viewImage2D(recoerrs[:,1],Ma)
xlabel("w1")
ylabel("w2")
axis("image")
title("reconstruction error, noise $(sqrt(noiseVar_range[1]))")
colorbar()
set_cmap("Spectral")
clim(0,17)

subplot(1,2,2)
viewImage2D(recoerrs[:,2],Ma)
xlabel("w1")
ylabel("w2")
axis("image")
title("reconstruction error, noise $(sqrt(noiseVar_range[2]))")
colorbar()
set_cmap("Spectral")
clim(0,17)


# function invC(x::Vector{Float64}, Q::AbstractSparseMatrix{T})
#     v = x
#     for j in 1:Q.n
#         tmpx = 0.0
#         for z in Q.colptr[j]:(Q.colptr[j+1]-1)
#             i = Q.rowval[z]
#             if i == j
#                 tmpx += v[j]
#                 v[j] = tmpx / (-Q.nzval[z])
#                 break
#             else
#                 tmpx -= (-Q.nzval[z]) * v[i]
#             end
#         end
#     end
#     v
# end

# function ex(x::Vector{Float64}, Q::SparseMatrixCSC;
#     maxiter=5000, tol=1.0e-16)::Vector{Float64}
#     b = zeros(Q.n)
#     v = x
#     y = invC(v, Q)
#     while maximum(v) > tol
#         gsstep!(v, Q, b)
#         y += v
#     end
#     y
# end

# function ex2(Q::SparseMatrixCSC, b::Vector{Float64};
#     x0::Vector{Float64}=fill(1/Q.n, Q.n),
#     maxiter=5000, steps=50, rtol=1.0e-6)::Vector{Float64}
#     x = x0
#     iter = 0
#     conv = false
#     rerror = 0.0
#     while true
#         prevx = x
#         for i in 1:steps
#             gsstep!(x, Q, b)
#         end
#         rerror = maximum(abs.((x - prevx) ./ x))
#         iter += steps
#         if rerror < rtol
#             conv = true
#             break
#         end
#         if iter >= maxiter
#             break
#         end
#     end
#     @printf "convergence   : %s\n" conv ? "true" : "false"
#     @printf "iteration     : %d / %d\n" iter maxiter
#     @printf "relative error: %e < %e\n" rerror rtol
#     x
# end

# function test2(Q, x)
#     C = triu(Q, 0)
#     D = tril(Q, -1)
#     y = zeros(size(x))
#     x = (-C)' \ x
#     y += x
#     for i in 1:100
#         x = D' * x
#         x = (-C)' \ x
#         y += x
#     end
#     y
# end

# function test(Q::SparseMatrixCSC;
#     x0::Vector{Float64}=fill(1/Q.n, Q.n),
#     maxiter=5000, steps=50, rtol=1.0e-6)::Vector{Float64}
#     x = x0
#     iter = 0
#     conv = false
#     rerror = 0.0
#     while true
#         prevx = x
#         for i in 1:steps
#             b = x
#             gsstep!(x, Q, b)
#             x /= sum(x)
#         end
#         rerror = maximum(abs.((x - prevx) ./ x))
#         iter += steps
#         if rerror < rtol
#             conv = true
#             break
#         end
#         if iter >= maxiter
#             break
#         end
#     end
#     @printf "convergence   : %s\n" conv ? "true" : "false"
#     @printf "iteration     : %d / %d\n" iter maxiter
#     @printf "relative error: %e < %e\n" rerror rtol
#     x
# end


# Tests for Revenue DEA Models
@testset "RevenueDEAModel" begin

    ## Test Revenue DEA Model with Cooper et al. (2007)
    # Test agains book results
    X = [3 2; 1 3; 4 6]
    Y = [3; 5; 6]
    P = [6; 6; 6]

    dearevenuecooper = dearevenue(X, Y, P, rts = :CRS)
    @test efficiency(dearevenuecooper, :Economic)   ≈ [0.9; 1; 0.6] atol = 1e-3
    @test efficiency(dearevenuecooper, :Technical)  ≈ [0.9; 1; 0.6  ] atol = 1e-3
    @test efficiency(dearevenuecooper, :Allocative) ≈ [1; 1; 1] atol = 1e-3

    ## Test Revenue DEA Model with Zofío and Prieto (2006) data.
    # Test agains results with R
    X = [5 3; 2 4; 4 2; 4 8; 7 9]
    Y = [7 4; 10 8; 8 10; 5 4; 3 6]
    P = [3 2; 3 2; 3 2; 3 2; 3 2.0]

    # Revnue CRS
    dearevenuecrs = dearevenue(X, Y, P, rts = :CRS)
    @test efficiency(dearevenuecrs) ≈ [0.4915254237;
                                    1.000;
                                    1.000;
                                    0.250;
                                    0.1735537190]
    @test efficiency(dearevenuecrs, :Technical) ≈ [0.6364;
                                1.000;
                                1.000;
                                0.250;
                                0.2609]  atol = 1e-3
    @test efficiency(dearevenuecrs, :Allocative) ≈ [0.7723970944;
                                1.000;
                                1.000;
                                1.000;
                                0.6652892562]

    # Revenue VRS
    dearevenuevrs = dearevenue(X, Y, P, rts = :VRS)
    @test efficiency(dearevenuevrs) ≈ [0.6444444444 ;
                                    1.000;
                                    1.000;
                                    0.500;
                                    0.4565217391]
    @test efficiency(dearevenuevrs, :Technical) ≈ [0.7777777778 ;
                                1.000;
                                1.000;
                                0.500 ;
                                0.6000000000]
    @test efficiency(dearevenuevrs, :Allocative) ≈ [0.8285714286;
                                1.000;
                                1.000;
                                1.000;
                                0.7608695652]

    # Check defaults
    @test efficiency(dearevenue(X, Y, P)) == efficiency(dearevenuevrs)
    @test efficiency(dearevenuevrs, :Economic) == efficiency(dearevenuevrs)

    # Print
    show(IOBuffer(), dearevenuecooper)

    # Test errors
    @test_throws ErrorException dearevenue([1; 2 ; 3], [4 ; 5], [1; 1; 1]) #  Different number of observations
    @test_throws ErrorException dearevenue([1; 2; 3], [4; 5; 6], [1; 2; 3], rts = :Error) # Invalid returns to scale
    @test_throws ErrorException dearevenue([1; 2; 3], [4; 5; 6], [1; 2; 3; 4]) # Different number of observation in prices
    @test_throws ErrorException dearevenue([1; 2; 3], [4 4; 5 5; 6 6], [4 4 4; 5 5 5; 6 6 6]) # Different number of output prices and outputs

end