__precompile__()
"""
    EasyEcon

Easy & friendly mathematics/economics in OLG model;
recommend to using it for fast debugging & mass use
"""
module EasyEcon

# ==============================================================================
    """
        CDProdFunc( A::Real, K::Real, L::Real, β::Real; κ::Real = 0.05, GetPrice::Bool = true, TechType::String = "Hicks"  )

    Cobb-Douglas production function, where A is technology, K is capital, L is labour, β is income share of capital;
    allows three types of technology: "Hicks", "Harod" & "Solow", controled by TechType parameter;
    if GetPrice is true, separately returns: output (Y), equilibrium NET interest rate (r-κ) & equilibrium wage level (w̄);
    where the capital depreciation rate is controled by keyword parameter κ;
    if GetPrice is false, only returns: output (Y), a real number
    """
    function CDProdFunc( A::Real, K::Real, L::Real, β::Real; κ::Real = 0.05, GetPrice::Bool = true, TechType::String = "Hicks"  )
        # 1. bounding check
        tmp = [A,K,L,β,κ]
        @assert( all(tmp .> 0), string("bound failure [A,K,L,β,κ], all positive required: ",tmp) )
        # 2. compute
        if TechType == "Hicks"
            Y = A * K^β * L^(1-β)
            if GetPrice
                r = β * A * (K/L)^(β-1) - κ
                w̄ = (1-β) * A * (K/L)^β
                return Y::Real, r::Real, w̄::Real
            else
                return Y::Real
            end
        elseif TechType == "Harod"
            Y = K^β * (A * L)^(1-β)
            if GetPrice
                r = β * (K/A/L)^(β-1) - κ
                w̄ = (1-β) * A * (K/A/L)^β
                return Y::Real, r::Real, w̄::Real
            else
                return Y::Real
            end
        elseif TechType == "Solow"
            Y = (A * K)^β * L^(1-β)
            if GetPrice
                r = β * (A*K/L)^(β-1) - κ
                w̄ = (1-β) * (A*K/L)^β
                return Y::Real, r::Real, w̄::Real
            else
                return Y::Real
            end
        else
            throw( ErrorException("invalid TechType, pls choose from \"Hicks\", \"Harod\", \"Solow\"") )
        end
        return nothing
    end

    # --------------------------------------------
    """
        WageProfile( w̄::Real, ε::Vector, Nst::Vector, Lst::Vector; GetScalingCoef::Bool = false  )

    profile wage among different working cohorts in year t where age ``s=1,...,S_r`` (``S_r`` is retiring age);
    using equation: ``w̄ L_{t} = \\sum^{S_r}_{s=1} w_{s,t} L_{s,t} N_{s,t}``;
    where ``L_{t} = \\sum^{S_r} L_{s,t}N_{s,t}`` is labour supply, ``w_{s,t} = o_{t} \varepsilon_{s}`` is wage profiling,
    ``o_{t}`` is a year-specific scaling coefficient (computed by the equation), ``\varepsilon_{s}`` is age-specific profiling coefficients;
    using Lst for ``L_{s,t}``, and Nst for ``N_{s,t}``;
    if GetSaclingCoef is true, separately returns: wprofiled::Vector (``w_{s,t}``), o::Real (``o_{t}``);
    if false, only returns: wprofiled;

    ## Inputs:
    1. ̄w: average wage level
    2. ε: wage profiling coefficient
    3. Nst: working population vector
    4. Lst: labour vector (1 - leisure)
    5. GetScalingCoef: if to return the wage scaling coefficient
    """
    function WageProfile( w̄::Real, ε::Vector, Nst::Vector, Lst::Vector; GetScalingCoef::Bool = false  )
        # 1. size check
        @assert( length(ε) == length(Nst) == length(Lst), "requires same length: ε, Nst, Lst" )
        # 2. compute
        L = sum( Nst .* Lst )  # get total labour supply
        o = L / sum( Nst .* Lst .* ε )  # get scaling coefficient
        wprofiled = w̄ .* o .* ε  # get profiled wage levels
        # 3. return
        if GetScalingCoef
            return wprofiled::Vector, o::Real
        else
            return wprofiled::Vector
        end
        return nothing
    end

    # --------------------------------------------
    """
        PAYGPension( πCoef::Real, w::Vector, N::Vector, Lab::Vector, Sr::Int )

    Computes an average pension benefit amount for retired popultion:
        `` \\pi_{t} \\sum^{S_r}_{s=1} w_{s,t} N_{s,t} L_{s,t} = \\Lambda_{t} \\sum^{S}_{s>S_r} N_{s,t}``
    where ``\\pi_{t}`` is total contribution rate in year ``t``, ``L_{s,t}`` is labor supply vector, ``\\Lambda_{t}`` is average pension benefit amount in year ``t``, ``N_{s,t}`` is population.

    ##Inputs:
    1. πCoef : ``\\pi_{t}``
    2. w : ``w_{s,t}``, nomial wage level, len = Sr
    3. N : compelete population, from 1 to Sr (retirement age) to S (max age), len = S
    4. Lab : ``L_{s,t} = 1 - l_{s,t}``, where ``l_{s,t}`` is leisure, len = Sr
    5. Sr : retirement age

    returns a Real number, the average pension benefit amount ``\\Lambda_{t}``.
    """
    function PAYGPension( πCoef::Real, w::Vector, N::Vector, Lab::Vector, Sr::Int )
        # 1. size check
        @assert( length(w) == length(N[1:Sr]) == length(Lab), "uncompatible vector size: w, Nwork or Lab"  )
        # 2. accounting
        local Λ::Float64
        Λ = πCoef * sum( w .* N[1:Sr] .* Lab ) / sum( N[Sr+1:end] )
        return Λ
    end

    # --------------------------------------------
    """
        Get𝕡( 𝕓::Real, πMf::Real, w::Vector, L::Vector, N::Vector, Sr::Int )

    Gets transfer payment amount 𝕡 from firm contribution to UE-BMI in year ``t`` to those retired in this year.
    based on a transfer rate 𝕓.

    ## Inputs:
    1. 𝕓: transfer rate/fraction from firm total contribution amount (to UE-BMI) to retired generations
    2. πMf: total contribution rate by firm on nomial wage level ``w_{s,t}``; it is defined as: ``\\pi^{Mf} = \\frac{\\zeta_{t}}{1+z \\eta + \\zeta}``
    3. w: nomial wage level vector, len = Sr
    4. Lab: labor supply, len = Sr
    5. N: population, len = S (env.MAX_AGE)
    6. Sr: retirement age

    returns a Real number 𝕡, the transfer payment amount per capita to retired generations.
    """
    function Get𝕡( 𝕓::Real, πMf::Real, w::Vector, Lab::Vector, N::Vector, Sr::Int )
        # 1. size check
        @assert( length(w) == length(Lab) == length(N[1:Sr]) , "uncompatible vector size: w, Lab or N" )
        # 2. accounting
        local 𝕡::Float64
        𝕡 = 𝕓 * πMf * sum( w .* N[1:Sr] .* Lab ) / sum( N[Sr+1:end] )
        return 𝕡
    end

# ==============================================================================
end  # module ends


"""
# Description
Constructor for the isotropic Squared Exponential kernel (covariance)

k(x,x') = σ²exp(-(x-x')ᵀ(x-x')/2ℓ²)
# Arguments:
* `ll::Float64`: Log of the length scale ℓ
* `lσ::Float64`: Log of the signal standard deviation σ
"""
type SEIso <: Isotropic
    ℓ2::Float64      # Length scale
    σ2::Float64      # Signal std
    SEIso(ll::Float64, lσ::Float64) = new(exp(2*ll), exp(2*lσ))
end

function set_params!(se::SEIso, hyp::Vector{Float64})
    length(hyp) == 2 || throw(ArgumentError("Squared exponential only has two parameters"))
    se.ℓ2, se.σ2 = exp(2.0*hyp)
end

get_params(se::SEIso) = Float64[log(se.ℓ2)/2.0, log(se.σ2)/2.0]
get_param_names(::SEIso) = [:ll, :lσ]
num_params(se::SEIso) = 2

metric(se::SEIso) = SqEuclidean()
cov(se::SEIso, r::Float64) = se.σ2*exp(-0.5*r/se.ℓ2)

function grad_kern(se::SEIso, x::Vector{Float64}, y::Vector{Float64})
    r = distance(se, x, y)
    exp_r = exp(-0.5*r/se.ℓ2)
    g1 = se.σ2*r/se.ℓ2*exp_r # dK_d(log ℓ)
    g2 = 2.0*se.σ2*exp_r     # dK_d(log σ)
    return [g1,g2]
end

function grad_stack!(stack::AbstractArray, se::SEIso, X::Matrix{Float64}, data::IsotropicData)
    nobsv = size(X, 2)
    R = distance(se, X, data)
    exp_R = exp(-0.5*R/se.ℓ2)
    for i in 1:nobsv, j in 1:i
        @inbounds stack[i,j,1] = se.σ2*R[i,j]/se.ℓ2*exp_R[i,j] # dK_dℓ
        @inbounds stack[j,i,1] = stack[i,j,1]
        @inbounds stack[i,j,2] = 2.0*se.σ2*exp_R[i,j]
        @inbounds stack[j,i,2] = stack[i,j,2]
    end

    return stack
end


function equations_to_matrix_form(eqs, vars)
    expand_eqs = expand_derivatives.(eqs)
    A = Symbolics.jacobian(expand_eqs, vars; simplify = true)
    # # Ensure none of x is in A
    # vars_A = Symbolics.get_variables(A)
    # if !isempty(vars_A)
    #     for v in vars
    #         @assert all(vars_A .!= v)
    #     end
    # end

    b = simplify.(-substitute(expand_eqs, Dict(v => 0 for v in vars)))
    A, b
end

function generate_equations_of_motion()
    @variables t
    @variables u(t) w(t) θ′(t)

    @variables θ_t(t) θ_k(t) γ_OA(t)

    LT = symbolic_kinetic_energy(u, w, θ′, t, θ_t = θ_t, θ_k = θ_k, γ_OA = γ_OA)

    D = Differential(t)
    # Need to generate A and b for Ax = b where x is D(u), D(w), D(θ′)

    # To me A = jacobian of LT w.r.t x assuming LT is linear in x
    x = [D(u), D(w), D(θ′)]
    A, b = equations_to_matrix_form(LT, x)
    println("Variables in A and b after equations_to_matrix_form")
    println(Symbolics.get_variables.(A))
    println(Symbolics.get_variables.(b))

    generated_dir = joinpath(@__DIR__, "..", "generated")
    funA, _ = build_function(A, (θ_t, θ_k, γ_OA))
    write(joinpath(generated_dir, "inertia_lhs.jl"), string(funA))

    funb, _ = build_function(b, (u, w, θ′, #=t,=# θ_t, D(θ_t), D(D(θ_t)), θ_k, D(θ_k), D(D(θ_k)), γ_OA, D(γ_OA), D(D(γ_OA))))
    write(joinpath(generated_dir, "inertia_rhs.jl"), string(funb))
    (funA, funb)
end


"""
```
normals{VT,FD,FT,FO}(vertices::Vector{Point{3, VT}},
                     faces::Vector{Face{FD,FT,FO}},
                     NT = Normal{3, VT})
```

Compute all vertex normals.
"""
function normals(
        vertices::AbstractVector{Point{3, VT}},
        faces::AbstractVector{F},
        NT = Normal{3, VT}
    ) where {VT, F <: Face}
    normals_result = zeros(Point{3, VT}, length(vertices)) # initilize with same type as verts but with 0
    for face in faces
        v = vertices[face]
        # we can get away with two edges since faces are planar.
        a = v[2] - v[1]
        b = v[3] - v[1]
        n = cross(a, b)
        for i =1:length(F)
            fi = face[i]
            normals_result[fi] = normals_result[fi] + n
        end
    end
    normals_result .= NT.(normalize.(normals_result))
    normals_result
end

"""
Slice an AbstractMesh at the specified Z axis value.
Returns a Vector of LineSegments generated from the faces at the specified
heights. Note: This will not slice in-plane faces.
"""
function slice(mesh::AbstractMesh{Point{3,VT}, Face{3, FT}}, height::Number) where {VT<:AbstractFloat,FT<:Integer}

    height_ct = length(height)
    # intialize the LineSegment array
    slice = Simplex{2,Point{2,VT}}[]

    for face in mesh.faces
        v1,v2,v3 = mesh.vertices[face]
        zmax = max(v1[3], v2[3], v3[3])
        zmin = min(v1[3], v2[3], v3[3])
        if height > zmax
            continue
        elseif zmin <= height
            if v1[3] < height && v2[3] >= height && v3[3] >= height
                p1 = v1
                p2 = v3
                p3 = v2
            elseif v1[3] > height && v2[3] < height && v3[3] < height
                p1 = v1
                p2 = v2
                p3 = v3
            elseif v2[3] < height && v1[3] >= height && v3[3] >= height
                p1 = v2
                p2 = v1
                p3 = v3
            elseif v2[3] > height && v1[3] < height && v3[3] < height
                p1 = v2
                p2 = v3
                p3 = v1
            elseif v3[3] < height && v2[3] >= height && v1[3] >= height
                p1 = v3
                p2 = v2
                p3 = v1
            elseif v3[3] > height && v2[3] < height && v1[3] < height
                p1 = v3
                p2 = v1
                p3 = v2
            else
                continue
            end

            start = Point{2,VT}(p1[1] + (p2[1] - p1[1]) * (height - p1[3]) / (p2[3] - p1[3]),
                                p1[2] + (p2[2] - p1[2]) * (height - p1[3]) / (p2[3] - p1[3]))
            finish = Point{2,VT}(p1[1] + (p3[1] - p1[1]) * (height - p1[3]) / (p3[3] - p1[3]),
                                 p1[2] + (p3[2] - p1[2]) * (height - p1[3]) / (p3[3] - p1[3]))

            push!(slice, Simplex{2,Point{2, VT}}(start, finish))
        end
    end

    return slice
end

# TODO this should be checkbounds(Bool, ...)
"""
```
checkbounds{VT,FT,FD,FO}(m::AbstractMesh{VT,Face{FD,FT,FO}})
```

Check the `Face` indices to ensure they are in the bounds of the vertex
array of the `AbstractMesh`.
"""
function Base.checkbounds(m::AbstractMesh{VT, Face{FD, FT}}) where {VT, FD, FT}
    isempty(faces(m)) && return true # nothing to worry about I guess
    flat_inds = reinterpret(FT, faces(m))
    checkbounds(Bool, vertices(m), flat_inds)
end


# by Lazaro Alonso
using GLMakie
GLMakie.activate!()
let
    n = 50
    x = y = z = LinRange(-5,5,n)
    vals = zeros(n,n,n)
    for i in 1:n, j in 1:n, k in 1:n
        vals[i,j,k] = 0.5*i^2  + j^2 + 2*k^2
    end
    cmap = :Spectral_11
    fig, _ = contour(x, y, z, vals, colormap = cmap, levels=[100,2500,7000],
         algorithm=:mip, alpha = 0.05, ambient = Vec3f0(0.65, 0.65, 0.65))
    fig
end


module ImageReconstruction

export radon, iradon

"""
radon transform

https://en.wikipedia.org/wiki/Radon_transform
"""
radon

"""
inverse radon transform

https://en.wikipedia.org/wiki/Radon_transform
"""
iradon

"""
Radon transform of an image using a pixel-driven algorithm.
"""
function radon(image::AbstractMatrix, θ::AbstractRange, t::AbstractRange)
    P = zeros(eltype(image), length(t), length(θ))
    nr, nc = size(image)

    for i = 1:nr, j = 1:nc
        x = j - nc / 2 + 0.5
        y = i - nr / 2 + 0.5
        @inbounds for (k, θₖ) in enumerate(θ)
            t′ = x * cos(θₖ) + y * sin(θₖ)

            a = convert(Int, round((t′ - t.start) / step(t) + 1))

            (a < 1 || a > length(t)) && continue
            α = abs(t′ - t[a])
            P[a, k] += (1 - α) * image[i, j]

            (a > length(t) + 1) && continue
            P[a+1, k] += α * image[i, j]
        end
    end

    P
end

end # module


# for now only univariate problems and with discretization

struct NormalConvolutionProblem
    prior::Distribution
    marginal::Vector{Float64} # measure after truncation & binning
    marginal_grid::Vector{Float64}
    marginal_h::Float64
end

struct DiscretizedNormalConvolutionProblem
    prior::Vector{Float64}
    prior_grid::Vector{Float64}
    prior_h::Float64
    marginal::Vector{Float64} # measure after truncation & binning
    marginal_grid::Vector{Float64}
    marginal_h::Float64
end

function marginal_grid_l(d::T) where T<:Union{NormalConvolutionProblem, DiscretizedNormalConvolutionProblem}
    d.marginal_grid .- d.marginal_h./2
end

function marginal_grid_r(d::T) where T<:Union{NormalConvolutionProblem, DiscretizedNormalConvolutionProblem}
    d.marginal_grid .+ d.marginal_h./2
end

# this is kind of a bottleneck right now (i.e. slower than it should be)
function NormalConvolutionProblem(prior, marginal_grid)
    marginal_grid = collect(marginal_grid)
    marginal_h = marginal_grid[2] - marginal_grid[1]
    d = NormalConvolutionProblem(prior, zeros(Float64, length(marginal_grid)), marginal_grid, marginal_h)
    marginal_l = marginal_grid_l(d)
    marginal_r = marginal_grid_r(d)

    Z = Normal()

    f(x) = pdf(prior,x[1])*pdf(Z, x[2]-x[1]) #x[1] is the mean

    for i=1:length(marginal_l)
        # TODO: Fix hardcoded 20
        # Idea: Speed this up for normal distributions..
       d.marginal[i] = hcubature(f, [-20, marginal_l[i]], [+20, marginal_r[i]];
                                 rtol=1e-13)[1]
    end
    d.marginal ./= sum(d.marginal)
    d
end

function NormalConvolutionProblem(prior::Normal, marginal_grid)
    marginal_grid = collect(marginal_grid)
    marginal_h = marginal_grid[2] - marginal_grid[1]
    d = NormalConvolutionProblem(prior, zeros(Float64, length(marginal_grid)), marginal_grid, marginal_h)
    marginal_l = marginal_grid_l(d)
    marginal_r = marginal_grid_r(d)

    μ_marginal = prior.μ
    σ_marginal = sqrt(1 + prior.σ^2)
    Z_marginal = Normal(μ_marginal, σ_marginal)

    for i=1:length(marginal_l)
        #TODO: I really don't like how boundary is being handled right now
        # Idea: Speed this up for normal distributions..
       d.marginal[i] = quadgk(β ->  pdf(Z_marginal, β) , marginal_l[i],marginal_r[i])[1]
    end
    d.marginal ./= sum(d.marginal)
    d
end

function convolution_matrix(::Type{DiscretizedNormalConvolutionProblem}, prior_grid, marginal_grid; σ=1.0)
    #marginal_h = marginal_grid[2] - marginal_grid[1]
    A = reshape([pdf(Normal(0,σ), x-μ) for x in prior_grid for μ in marginal_grid],
                length(marginal_grid), length(prior_grid))#.*marginal_h
    A
end

function DiscretizedNormalConvolutionProblem(prior, prior_grid, marginal_grid)
    prior_grid = collect(prior_grid)
    marginal_grid = collect(marginal_grid)

    prior_h = prior_grid[2]-prior_grid[1]

    marginal_h = marginal_grid[2]-marginal_grid[1]
    A = convolution_matrix(DiscretizedNormalConvolutionProblem,
                                prior_grid, marginal_grid) # should this be stored maybe?

    f_marginal = A*prior
    f_marginal ./= sum(f_marginal)

    DiscretizedNormalConvolutionProblem(prior, prior_grid, prior_h,
                             f_marginal, marginal_grid, marginal_h)
end

function DiscretizedNormalConvolutionProblem(prior_distr::Distribution, prior_grid, marginal_grid)
    prior_h = prior_grid[2]-prior_grid[1]
    prior = pdf.(prior_distr, prior_grid).*prior_h # makes this an approximate prob. measure
    DiscretizedNormalConvolutionProblem(prior, prior_grid, marginal_grid)
end

#---- TODO: maybe deprecate these two functions eventually
function posterior_stats(d::NormalConvolutionProblem, f, x)
    prior = d.prior
    Z = Normal()
    post_num =  quadgk(β ->  f(β)*pdf(prior, β)*pdf(Z, x-β) , -10,10)[1]
    post_denom = quadgk(β ->  pdf(prior, β)*pdf(Z, x-β) , -10,10)[1]
    (post_num, post_denom, post_num/post_denom)
end

function posterior_stats(d::DiscretizedNormalConvolutionProblem, ψ, x)
    prior = d.prior
    prior_grid = d.prior_grid
    post_num = d.prior'*(ψ.(prior_grid).*pdf.(Normal(), prior_grid .- x))
    post_denom = d.prior'* pdf.(Normal(), prior_grid .- x)
    (post_num, post_denom, post_num/post_denom)
end
#----------------------------------------------------------------------
# main fun -- rename eventually

function posterior_stats(prior::ContinuousUnivariateDistribution, t::LinearInferenceTarget)
    quadgk(β ->  pdf(prior, β)*riesz_representer(t, β) , -10,10)[1]
end

function posterior_stats(d::NormalConvolutionProblem, t::LinearInferenceTarget)
    posterior_stats(d.prior, t)
end

function posterior_stats(d::NormalConvolutionProblem, t::PosteriorTarget)
    num = posterior_stats(d, t.num)
    denom = posterior_stats(d, t.denom)
    (num, denom, num/denom)
end

function StatsBase.rand(d::NormalConvolutionProblem, n::Integer)
     μs = rand(d.prior, n)
     μs .+ randn(n)
end

# standardize if we mean marginal or prior pdf?
function pdf(d::NormalConvolutionProblem, x)
    posterior_stats(d, MarginalDensityTarget(x))
end

# mmaybe lazy eval this and get marginal_map only once has been set
struct MixingNormalConvolutionProblem{T<:Distribution}
    priors::Vector{T}
    prior_mixture_coef::Vector{Float64}
    marginal_map::Matrix{Float64} #map pi_s -> marginal dbn
    marginal_grid::Vector{Float64}
    marginal_h::Float64
end

function MixingNormalConvolutionProblem(priors::Vector{T},
            marginal_grid,
            prior_mixture_coef=ones(length(priors))./length(priors)) where T<:Distribution

    marginal_grid = collect(marginal_grid)

    #marginal as function of specified prior..
    n_priors = length(priors)
    n_marginal_grid = length(marginal_grid)
    marginal_h = marginal_grid[2] - marginal_grid[1]

    marginal_map = zeros(Float64, n_marginal_grid, n_priors)
    for i=1:n_priors
        tmp = NormalConvolutionProblem(priors[i], marginal_grid)
        marginal_map[:,i] = tmp.marginal
    end
    MixingNormalConvolutionProblem(priors, ones(n_priors)./n_priors,
                    marginal_map, marginal_grid, marginal_h)
end

function MixingNormalConvolutionProblem(::Type{Normal}, σ, prior_grid, marginal_grid)
    priors = [Normal(μ, σ) for μ in prior_grid]
    MixingNormalConvolutionProblem(priors, marginal_grid)
end

# TODO: Implement using iterator specification instead..
function index_mixing(d::MixingNormalConvolutionProblem, i)
    NormalConvolutionProblem(d.priors[i], d.marginal_grid)
end

function posterior_stats(d::MixingNormalConvolutionProblem, t)
    [posterior_stats(index_mixing(d,i), t) for i=1:length(d.priors)]
end

function Base.length(ds::MixingNormalConvolutionProblem)
    length(ds.priors)
end

function Distributions.pdf(ds::MixingNormalConvolutionProblem, mixing_coeff, grid)
     d = MixtureModel(ds.priors, vec(mixing_coeff))
     pdf.(d, grid)
end

function Distributions.pdf(ds::MixingNormalConvolutionProblem, grid)
     Distributions.pdf(ds, ds.prior_mixture_coef, grid)
end

function NormalConvolutionProblem(ds::MixingNormalConvolutionProblem, mixing_coeff)
    prior =  MixtureModel(ds.priors, vec(mixing_coeff))
    NormalConvolutionProblem(prior, ds.marginal_grid)
end


function mass(u, q, w, v, x, t, mu, eta)
    m = [1.0, 1.0, 1.0, 1.0];
    return m;
end
function flux(u, q, w, v, x, t, mu, eta)
    gam = mu[1];
    gam1 = gam - 1.0;
    r = u[1];
    ru = u[2];
    rv = u[3];
    rE = u[4];
    r1 = 1/r;
    uv = ru*r1;
    vv = rv*r1;
    E = rE*r1;
    p = gam1*(rE-r*0.5*(uv*uv+vv*vv));
    h = E+p*r1;
    f = [ru, ru*uv+p, rv*uv, ru*h, rv, ru*vv, rv*vv+p, rv*h];
    f = reshape(f,(4,2));
    return f;
end
function source(u, q, w, v, x, t, mu, eta)
    s = [0.0, 0.0, 0.0, 0.0];
    return s;
end
function ubou(u, q, w, v, x, t, mu, eta, uhat, n, tau)
    ub = zeros((4,2));
    return ub;
end
function fbou(u, q, w, v, x, t, mu, eta, uhat, n, tau)
    f = flux(u, q, w, v, x, t, mu, eta);

    gam = mu[1];
    gam1 = gam - 1.0;
    r = u[1];
    ru = u[2];
    rv = u[3];
    rE = u[4];
    nx = n[1];
    ny = n[2];

    r1 = 1/r;
    uv = ru*r1;
    vv = rv*r1;
    E = rE*r1;
    p = gam1*(rE-r*0.5*(uv*uv+vv*vv));
    h = E+p*r1;
    a = sqrt(gam*p*r1);

    run = ru*nx + rv*ny;
    rut = -ru*ny + rv*nx;
    un = run/r;
    ut = rut/r;

    K = [ 1  1  0  1 ; un-a  un  0  un+a ; ut  ut  1  ut ; h - un*a  (1/2)*(un^2 + ut^2)  ut  h+un*a];
    Kinv = (gam1/(2*a^2))*[ h + (a/gam1)*(un-a)  -(un+a/gam1)  -ut  1; -2*h + (4/gam1)*a^2  2*un  2*ut  -2 ; -2*(ut*a^2)/gam1  0  2*(a^2)/gam1  0 ; h - a*(un+a)/gam1  -un+a/gam1  -ut  1];
    T = [ 1 0 0 0; 0 nx ny 0; 0 -ny nx 0; 0 0 0 1];
    Tinv = [ 1 0 0 0; 0 nx -ny 0; 0 ny nx 0; 0 0 0 1];
    Lambda = [ tanh(1e3*(un-a)) 0 0 0 ; 0 tanh(1e3*(un)) 0 0 ; 0 0 tanh(1e3*(un)) 0 ; 0 0 0 tanh(1e3*(un+a)) ];
    # E = simplify(K * Lambda * Kinv);
    # An = simplify(Tinv * E * T);
    E = (K * Lambda * Kinv);
    An = (Tinv * E * T);

    # freestream boundary condition
    uinf = mu[3:6]; # freestream flow
    uinf = uinf[:];
    u = u[:];          # state variables
    ui = 0.5*((u.+uinf) + An*(u.-uinf));  # Riemann solution
    fi = f[:,1]*n[1] + f[:,2]*n[2] + tau[1]*(u-ui);

    uw = u[:];
    uw[2] = ru - run*nx;
    uw[3] = rv - run*ny;
    uw = 0.5*((u.+uw) + An*(u.-uw));  # Riemann solution
    fw = f[:,1]*n[1] + f[:,2]*n[2] + tau[1]*(u-uw); # numerical flux at freestream boundary
    fw[1] = 0;
    fw[end] = 0;
    fb = hcat(fw[:], fi[:]); # wall and freestream boundary conditions

    return fb;
end
function initu(x, mu, eta)

    u0 = mu[3:6]; # freestream flow
    return u0;
end


#---------------------------------------------------------------------#
#This code computes the element mass and differentiation matrices
#Written by F.X. Giraldo on April 24, 2019
#           Department of Applied Mathematics
#           Naval Postgraduate School
#           Monterey; CA 93943-5216
#---------------------------------------------------------------------#
function element_matrices(ψ,dψ,ξ_x,ξ_y,η_x,η_y,jac,ωq,Ne,Np,Nq,DFloat)

    #Initialize Matrices
    M=zeros(DFloat,Np,Np,Np,Np,Ne)
    D_x=zeros(DFloat,Np,Np,Np,Np,Ne)
    D_y=zeros(DFloat,Np,Np,Np,Np,Ne)

    #Construct Mass and Differentiation Matrices
    for e=1:Ne
        for l=1:Nq, k=1:Nq
            wq=ωq[k]*ωq[l]*jac[k,l,e]
            for j=1:Np, i=1:Np
                Ψ_JK=ψ[i,k]*ψ[j,l] #h_ik*h_jl
                dΨdx_JK=dψ[i,k]*ψ[j,l]*ξ_x[k,l,e] + ψ[i,k]*dψ[j,l]*η_x[k,l,e]
                dΨdy_JK=dψ[i,k]*ψ[j,l]*ξ_y[k,l,e] + ψ[i,k]*dψ[j,l]*η_y[k,l,e]
                for n=1:Np, m=1:Np
                    Ψ_IK=ψ[m,k]*ψ[n,l] #h_ik*h_jl
                    dΨdx_IK=dψ[m,k]*ψ[n,l]*ξ_x[k,l,e] + ψ[m,k]*dψ[n,l]*η_x[k,l,e]
                    dΨdy_IK=dψ[m,k]*ψ[n,l]*ξ_y[k,l,e] + ψ[m,k]*dψ[n,l]*η_y[k,l,e]
                    M[m,n,i,j,e]+=wq*Ψ_IK*Ψ_JK
                    D_x[m,n,i,j,e]+=wq*dΨdx_IK*Ψ_JK
                    D_y[m,n,i,j,e]+=wq*dΨdy_IK*Ψ_JK
                end #m,n
            end #j,i
        end #k,l
    end #e

    return (M,D_x,D_y)
end


# Indefinite integral of pdist.pdf*f is clenshaw_asin(cdf(x),clenshaw_coefs)

function clenshaw_coefs{N}(pdist::Distribution, f::Function, degreetype::Val{N} = Val{50}())
    f1 = Fun(x-> 1/2*cos(x)*f(quantile(pdist, (sin(x)+1)/2)), -π/2..π/2, N-1)
    g1 = cumsum(f1)
    len = min(length(g1.coefficients),N)
    (len < length(g1.coefficients)) && warn("I'm truncating the coefficients!")
    u =  SVector{N, Float64}(vcat(g1.coefficients[1:len],zeros(N-len)))
    return u
end

#clenshaw_asin(x,c) = ApproxFun.clenshaw(2*(asin(2x-1))/π, c)
clenshaw_asin(x,c) = clenshaw_halved(4*(asin(2x-1))/π, c)
@generated function clenshaw_halved{N, R<:Real}(x::R, c::SVector{N})
    Expr(:block,
    :((a,b) = (zero(R), zero(R))),
    (:((a,b) = (muladd(x,a,c[$k]-b),a)) for k = N:-1:2)...,
    :(muladd(x/2,a,c[1]-b)))
end


# epsxc only version
function XC_c_pbe_spin_E( rho, zeta, grho )

    ga = 0.031091
    be = 0.06672455060314922
    third = 1.0/3.0
    pi34 = 0.6203504908994

    xkf = 1.919158292677513
    xks = 1.128379167095513

    rs = pi34/rho^third

    ec = XC_c_pw_spin_E( rho, zeta )

    kf = xkf / rs
    ks = xks * sqrt(kf)
    fz = 0.5 * ( (1.0 + zeta)^(2.0 / 3.0) + (1.0 - zeta)^(2.0 / 3.0) )

    fz2 = fz * fz
    fz3 = fz2 * fz

    t = sqrt(grho) / (2.0 * fz * ks * rho)
    expe = exp( - ec / (fz3 * ga) )
    af = be / ga * (1.0 / (expe - 1.0) )

    y = af * t * t

    xy = (1.0 + y) / (1.0 + y + y * y)
    qy = y * y * (2.0 + y) / (1.0 + y + y * y)^2

    s1 = 1.0 + be / ga * t * t * xy
    h0 = fz3 * ga * log(s1)

    sc = rho * h0

    return sc

end


"""
  vulnerability(N::UnipartiteNetwork)

This function is mesuring the vulnerability of a given network. The vulnerability of a network represents the ink density (L/S) divided by the number of preys. The number of preys: (prop. of basal + prop. of top) x S
"""

function vulnerability(N::UnipartiteNetwork)
    out_degree = collect(values(degree(N; dims=1)))
    in_degree = collect(values(degree(N; dims=2)))
    degrees = DataFrame(out=out_degree, in=in_degree)
    B = nrow(degrees[degrees.out .== 0, :]) / nrow(degrees)
    T = nrow(degrees[degrees.in .== 0, :]) / nrow(degrees)
    I = nrow(degrees[(degrees.out .!= 0) .& (degrees.in .!= 0), :]) / nrow(degrees)
    L = links(N)
    Vulnerability = links_1 ./ (S.^2 .* (B .+ I))
    return Vulnerability
end


@defcomp emissions begin
	SIG0	= Variable()				#Carbon intensity 2010-2015 (kgCO2 per output 2010 USD)
	GSIG	= Variable(index=[time])	#Change in sigma (cumulative improvement of energy efficiency)
	SIGMA	= Variable(index=[time])	#CO2-equivalent-emissions output ratio
    EIND    = Variable(index=[time])    #Industrial emissions (GtCO2 per year)
	ETREE	= Variable(index=[time])	#Emissions from deforestation
    E       = Variable(index=[time])    #Total CO2 emissions (GtCO2 per year)
	CUMETREE= Variable(index=[time])	#Cumulative from land
    CCA     = Variable(index=[time])    #Cumulative industrial emissions
	CCATOT	= Variable(index=[time])	#Cumulative total carbon emissions

	gsigma1 = Parameter()				#Initial growth of sigma (per year)
	dsig	= Parameter()				#Decline rate of decarbonization (per period)
	eland0	= Parameter()				#Carbon emissions from land 2015 (GtCO2 per year)
	deland	= Parameter()				#Decline rate of land emissions (per period)
	e0		= Parameter()				#Industrial emissions 2015 (GtCO2 per year)
    MIU     = Parameter(index=[time])   #Emission control rate GHGs
    YGROSS  = Parameter(index=[time])   #Gross world product GROSS of abatement and damages (trillions 2010 USD per year)
    cca0    = Parameter()               #Initial cumulative industrial emissions
	cumetree0=Parameter()				#Initial emissions from deforestation (see GAMS code)

    function run_timestep(p, v, d, t)
		#Define SIG0
			v.SIG0 = p.e0/(p.YGROSS[TimestepIndex(1)] * (1 - p.MIU[TimestepIndex(1)]))

		#Define function for GSIG
		if is_first(t)
			v.GSIG[t] = p.gsigma1
		else
			v.GSIG[t] = v.GSIG[t-1] * ((1 + p.dsig)^5)
		end

		#Define function for SIGMA
		if is_first(t)
			v.SIGMA[t] = v.SIG0
		else
			v.SIGMA[t] = v.SIGMA[t-1] * exp(v.GSIG[t-1] * 5)
		end

        #Define function for EIND
		v.EIND[t] = v.SIGMA[t] * p.YGROSS[t] * (1 - p.MIU[t])

		#Define function for ETREE
		if is_first(t)
			v.ETREE[t] = p.eland0
		else
			v.ETREE[t] = v.ETREE[t - 1] * (1 - p.deland)
		end

        #Define function for E
        v.E[t] = v.EIND[t] + v.ETREE[t]

		#Define function for CUMETREE
		if is_first(t)
			v.CUMETREE[t] = p.cumetree0
		else
			v.CUMETREE[t] = v.CUMETREE[t-1] + v.ETREE[t-1]
		end

        #Define function for CCA
        if is_first(t)
            v.CCA[t] = p.cca0
        else
            v.CCA[t] = v.CCA[t-1] + v.EIND[t-1] * 5/3.666
        end

		#Define function for CCATOT
			v.CCATOT[t] = v.CCA[t] + v.CUMETREE[t]

    end
end


# Implementing Forward Mode
# =========================
#
# In the [intro notebook](./intro.ipynb) we covered forward-mode differentiation
# thoroughly, but we don't yet have a real implementation that can work on our
# programs. Implementing AD effectively and efficiently is a field of its own,
# and we'll need to learn a few more tricks to get off the ground.

include("utils.jl");

# Up to now, we have differentiated things by creating a new Wengert list which
# contains parts of the original expression.

y = Wengert(:(5sin(log(x))))
derive(y, :x)

# We're now going to explicitly split our lists into two pieces: the original
# expression, and a new one which only calculates derivatives (but might refer
# back to values from the first). For example:

y = Wengert(:(5sin(log(x))))
#-
dy = derive(y, :x, out = Wengert(variable = :dy))
#-
Expr(dy)

# If we want to distinguish them, we can call `y` the *primal code* and `dy`
# the *tangent code*. Nothing fundamental has changed here, but it's useful
# to organise things this way.
#
# Almost all of the subtlety in differentiating programs comes from a
# mathematically trivial question: in what order do we evaluate the statements
# of the Wengert list? We have discussed the
# [forward/reverse](./backandforth.ipynb) distinction, but even once that choice
# is made, we have plenty of flexibility, and those choices can affect efficiency.
#
# For example, imagine if we straightforwardly evaluate `y` followed by `dy`. If
# we only cared about the final output of `y`, this would be no problem at all,
# but in general `dy` also needs to re-use variables like `y1` (or possibly any
# $y_i$). If our primal Wengert list has, say, a billion instructions, we end up
# having to store a billion intermediate $y_i$ before we run our tangent code.
#
# Alternatively, one can imagine running each instruction of the tangent code as
# early as possible; as soon as we run `y1 = log(x)`, for example, we know we
# can run `dy2 = cos(y1)` also. Then our final, combined program would look
# something like this:

# ```julia
# y0 = x
# dy = 1
# y1 = log(y0)
# dy = dy/y0
# y2 = cos(y1)
# dy = dy*sin(y1)
#   ...
# ```

# Now we can throw out `y1` soon after creating it, and we no longer have to
# store those billion intermediate results.
#
# The ability to do this is a very general property of forward differentiation;
# once we run $a = f(b)$, we can then run $\frac{da}{dx} = \frac{da}{db}
# \frac{db}{dx}$ using only `a` and `b`. It's really just a case of replacing
# basic instructions like `cos` with versions that calculate both the primal and
# tangent at once.

# Dual Numbers
# ------------
#
# Finally, the trick that we've been building up to: making our programming
# language do this all for us! Almost all common languages – with the notable
# exception of C – provide good support for *operator overloading*, which allows
# us to do exactly this replacement.
#
# To start with, we'll make a container that holds both a $y$ and a
# $\frac{dy}{dx}$, called a *dual number*.

struct Dual{T<:Real} <: Real
  x::T
  ϵ::T
end

Dual(1, 2)
#-
Dual(1.0,2.0)

# Let's print it nicely.

Base.show(io::IO, d::Dual) = print(io, d.x, " + ", d.ϵ, "ϵ")

Dual(1, 2)

# And add some of our rules for differentiation. The rules have the same basic
# pattern-matching structure as the ones we originally applied to our Wengert
# list, just with different notation.

import Base: +, -, *, /
a::Dual + b::Dual = Dual(a.x + b.x, a.ϵ + b.ϵ)
a::Dual - b::Dual = Dual(a.x - b.x, a.ϵ - b.ϵ)
a::Dual * b::Dual = Dual(a.x * b.x, b.x * a.ϵ + a.x * b.ϵ)
a::Dual / b::Dual = Dual(a.x * b.x, b.x * a.ϵ - a.x * b.ϵ)

Base.sin(d::Dual) = Dual(sin(d.x), d.ϵ * cos(d.x))
Base.cos(d::Dual) = Dual(cos(d.x), - d.ϵ * sin(d.x))

Dual(2, 2) * Dual(3, 4)

# Finally, we'll hook into Julia's number promotion system; this isn't essential
# to understand, but just makes it easier to work with Duals since we can now
# freely mix them with other number types.

Base.convert(::Type{Dual{T}}, x::Dual) where T = Dual(convert(T, x.x), convert(T, x.ϵ))
Base.convert(::Type{Dual{T}}, x::Real) where T = Dual(convert(T, x), zero(T))
Base.promote_rule(::Type{Dual{T}}, ::Type{R}) where {T,R} = Dual{promote_type(T,R)}

Dual(1, 2) * 3

# We already have enough to start taking derivatives of some simple functions.
# If we pass a dual number into a function, the $\epsilon$ component represents
# the derivative.

f(x) = x / (1 + x*x)

f(Dual(5., 1.))

# We can make a utility which allows us to differentiate any function.

D(f, x) = f(Dual(x, one(x))).ϵ

D(f, 5.)

# Dual numbers seem pretty scarcely related to all the Wengert list stuff we
# were talking about earlier. But we need take a closer look at how this is
# working. To start with, look at Julia's internal representation of `f`.

@code_typed f(1.0)

# This is just a Wengert list! Though the naming is a little different – `mul_float`
# rather than the more general `*` and so on – it's still essentially the same
# data structure we were working with earlier. Moreover, you'll recognise
# the code for the derivative, too!

@code_typed D(f, 1.0)

# This code is again the same as the Wengert list derivative we worked out at
# the very beginning of this handbook. The order of operations is just a little
# different, and there's the odd missing or new instruction due to the different
# set of optimisations that Julia applies. Still, we have not escaped our fundamental
# symbolic differentiation algorithm, just tricked the compiler into doing most
# of the work for us.

derive(Wengert(:(sin(cos(x)))), :x)
#-
@code_typed D(x -> sin(cos(x)), 0.5)

# What of data structures, control flow, function calls? Although these things
# are all present in Julia's internal "Wengert list", they end up being the same
# in the tangent program as in the primal; so an operator overloading approach
# need not deal with them explicitly to do the right thing. This won't be true
# when we come to talk more about reverse mode, which demands a more complex
# approach.

# Perturbation Confusion
# ----------------------
#
# Actually, that's not quite true. Operator-overloading-based forward mode
# *almost always* does the right thing, but it is not flawless. This more
# advanced section will talk about nested differentiation and the nasty bug that
# can come with it.
#
# We can differentiate any function we want, as long as we have the right
# primitive definitions for it. For example, the derivative of $\sin(x)$ is
# $\cos(x)$.

D(sin, 0.5), cos(0.5)

# We can also differentiate the differentiation operator itself. We'll find that
# the second derivative of $\sin(x)$ is $-\sin(x)$.

D(x -> D(sin, x), 0.5), -sin(0.5)

# This worked because we ended up nesting dual numbers. If we create a dual number
# whose $\epsilon$ component is another dual number, then we end up tracking the
# derivative of the derivative.

# The issue comes about when we close over a variable that *is itself* being
# differentiated.

D(x -> x*D(y -> x+y, 1), 1) # == 1

# The derivative $\frac{d}{dy} (x + y) = 1$, so this is equivalent to
# $\frac{d}{dx}x$, which should also be $1$. So where did this go wrong? The
# problem is that when we closed over $x$, we didn't just get a numeric value
# but a dual number with $\epsilon = 1$. When we then calculated $x + y$, both
# epsilons were added as if $\frac{dx}{dy} = 1$ (effectively $x = y$). If we had
# written this down, the answer would be correct.

D(x -> x*D(y -> y+y, 1), 1)

# I leave this second example as an excercise to the reader. Needless to say,
# this has caught out many an AD implementor.

D(x -> x*D(y -> x*y, 1), 4) # == 8

# More on Dual Numbers
# --------------------
#
# The above discussion presented dual numbers as essentially being a trick for
# applying the chain rule. I wanted to take the opportunity to present an
# alternative viewpoint, which might be appealing if, like me, you have any
# training in physics.
#
# Complex arithmetic involves a new number, $i$, which behaves like no other:
# specifically, because $i^2 = -1$. We'll introduce a number called $\epsilon$,
# which is a bit like $i$ except that $\epsilon^2 = 0$; this is effectively a
# way of saying the $\epsilon$ is a very small number. The relevance of this
# comes from the original definition of differentiation, which also requires
# $\epsilon$ to be very small.
#
# $$
# \frac{d}{dx} f(x) = \lim_{\epsilon \to 0} \frac{f(x+\epsilon)-f(x)}{\epsilon}
# $$
#
# We can see how our definition of $\epsilon$ works out by applying it to
# $f(x+\epsilon)$; let's say that $f(x) = sin(x^2)$.
#
# \begin{align}
# f(x + \epsilon) &= \sin((x + \epsilon)^2) \\
#                 &= \sin(x^2 + 2x\epsilon + \epsilon^2) \\
#                 &= \sin(x^2 + 2x\epsilon) \\
#                 &= \sin(x^2)\cos(2x\epsilon) + \cos(x^2)\sin(2x\epsilon) \\
#                 &= \sin(x^2) + 2x\cos(x^2)\epsilon \\
# \end{align}
#
# A few things have happened here. Firstly, we directly expand $(x+\epsilon)^2$
# and remove the $\epsilon^2$ term. We expand $sin(a+b)$ and then apply a *small
# angle approximation*: for small $\theta$, $\sin(\theta) \approx \theta$ and
# $\cos(\theta) \approx 1$. (This sounds pretty hand-wavy, but does follow from
# our original definition of $\epsilon$ if we look at the Taylor expansion of
# both functions). Finally we can plug this into our derivative rule.
#
# \begin{align}
# \frac{d}{dx} f(x) &= \frac{f(x+\epsilon)-f(x)}{\epsilon} \\
#                   &= 2x\cos(x^2)
# \end{align}
#
# This is, in my opinion, a rather nice way to derive functions by hand.
#
# This also leads to another nice trick, and a third way to look at forward-mode
# AD; if we replace $x + \epsilon$ with $x + \epsilon i$ then we still have
# $(\epsilon i)^2 = 0$. If $\epsilon$ is a small real number (say
# $1\times10^{-10}$), this is still true within floating point error, so our
# derivative still works out.

ϵ = 1e-10im
x = 0.5

f(x) = sin(x^2)

(f(x+ϵ) - f(x)) / ϵ
#-
2x*cos(x^2)

# So complex numbers can be used to get exact derivatives! This is very efficient
# and can be written using only one call to `f`.

imag(f(x+ϵ)) / imag(ϵ)

# Another way of looking at this is that we are doing bog-standard numerical
# differentiation, but the use of complex numbers avoids the typical problem
# with that technique (i.e. that a small perturbation ends up being overwhelmed
# by floating point error). The dual number is then a slight variation which
# makes the limit $\epsilon \rightarrow 0$ exact, rather than approximate.
# Forward mode AD can be described as "just" a clever implementation of
# numerical differentiation. Both numerical and forward derivatives propagate
# a perturbation $\epsilon$ using the same basic rules, and they have the
# same algorithmic properties.


n = rand(1: 10)
@test matrixdepot("moler", n) == matrixdepot("moler", n, -1)

M = matrixdepot("moler", n)
for i = 1:n, j = 1:n
    i != j ? (@test_approx_eq M[i,j] min(i,j) - 2) : (@test_approx_eq M[i,i] i)
end

println("'moler' passed test...")


### A Pluto.jl notebook ###
# v0.14.4

using Markdown
using InteractiveUtils

# ╔═╡ f11023e5-8f7b-4f40-86d3-3407b61863d9
using PlutoUI, Viznet, Compose, Plots

# ╔═╡ ce44f8bd-692e-4eab-9ba4-055b25e40c81
using ForwardDiff: Dual

# ╔═╡ a1ef579e-4b66-4042-944e-7e27c660095e
md"""
```math
\newcommand{\comment}[1]{{\bf  \color{blue}{\text{◂~ #1}}}}
```
"""

# ╔═╡ 9d11e058-a7d0-11eb-1d78-6592ff7a1b43
md"# An introduction to automatic differentiation

-- GiggleLiu"

# ╔═╡ b73157bf-1a77-47b8-8a06-8d6ec2045023
html"<button onclick='present()'>present</button>"

# ╔═╡ ec13e0a9-64ff-4f66-a5a6-5fef53428fa1
md"""
* What is automatic differentiation (AD)?
    * A true history of AD
    * Forward mode AD
    * Reverse mode AD
        * primitves on tensors (including tensorflow, pytorch et al.)
        * primitves on elementary instructions (usually source code transformation based)
        * defined on a reversible program
* Some applications in **scientific computing**
    * solving the graph embedding problem
    * inverse engineering a hamiltonian
    * obtaining maximum independent set (MIS) configurations
    * towards differentiating `expmv` ``\comment{will be used in our emulator}``
"""

# ╔═╡ f8b0d1ce-99f7-4729-b46e-126da540cbbe
md"""
## The true history of automatic differentiation
"""

# ╔═╡ 435ac19e-1c0c-4ee5-942d-f2a97c8c4d80
md"""
* 1964 ~ Robert Edwin Wengert, A simple automatic derivative evaluation program. ``\comment{first forward mode AD}``
* 1970 ~ Seppo Linnainmaa, Taylor expansion of the accumulated rounding error. ``\comment{first backward mode AD}``
* 1986 ~ Rumelhart, D. E., Hinton, G. E., and Williams, R. J., Learning representations by back-propagating errors.
* 1992 ~ Andreas Griewank, Achieving logarithmic growth of temporal and spatial complexity in reverse automatic differentiation. ``\comment{foundation of source code transformation based AD.}``
* 2000s ~ The boom of tensor based AD frameworks for machine learning.
* 2018 ~ People re-invented AD as differential programming ([wiki](https://en.wikipedia.org/wiki/Differentiable_programming) and this [quora answer](https://www.quora.com/What-is-Differentiable-Programming).)
![](https://qph.fs.quoracdn.net/main-qimg-fb2f8470f2120eb49c8142b08d9c4132)
* 2020 ~ Me, Differentiate everything with a reversible embeded domain-specific language ``\comment{AD based on reversible programming}``.
"""

# ╔═╡ 48ecd619-d01d-43ff-8b52-7c2566c3fa2b
md"## Forward mode automatic differentiation"

# ╔═╡ 4878ce45-40ff-4fae-98e7-1be41e930e4d
md"""
Forward mode AD attaches a infitesimal number $\epsilon$ to a variable, when applying a function $f$, it does the following transformation
```math
\begin{align}
    f(x+g \epsilon) = f(x) + f'(x) g\epsilon + \mathcal{O}(\epsilon^2)
\end{align}
```

The higher order infinitesimal is ignored.

**In the program**, we can define a *dual number* with two fields, just like a complex number
```
f((x, g)) = (f(x), f'(x)*g)
```
"""

# ╔═╡ b2c1936c-2c27-4fbb-8183-e38c5e858483
res = sin(Dual(π/4, 2.0))

# ╔═╡ 8be1b812-fcac-404f-98aa-0571cb990f34
res === Dual(sin(π/4), cos(π/4)*2.0)

# ╔═╡ 33e0c762-c75e-44aa-bfe2-bff92dd1ace8
md"
We can apply this transformation consecutively, it reflects the chain rule.
```math
\begin{align}
\frac{\partial \vec y_{i+1}}{\partial x} &= \boxed{\frac{\partial \vec y_{i+1}}{\partial \vec y_i}}\frac{\partial \vec y_i}{\partial x}\\
&\text{local Jacobian}
\end{align}
```
"

# ╔═╡ c59c35ee-1907-4736-9893-e22c052150ca
let
	lb = textstyle(:math, fontsize(8), width=0.5, height=0.5)
	tb = textstyle(:default, fontsize(10), Compose.font("monospace"))
	tb_big = textstyle(:default, fontsize(3.5), fill("white"), Compose.font("monospace"))
	nb = nodestyle(:circle, fill("white"), Compose.stroke("black"); r=0.08)
	tri = nodestyle(:triangle, Compose.stroke("transparent"), fill("black"); r=0.02)
	eb = bondstyle(:default, linewidth(0.5mm))
	ebr = bondstyle(:default, Compose.stroke("red"), linewidth(0.5mm))
	ebd = bondstyle(:default, linewidth(0.5mm), dashed=true)
	eba = bondstyle(:default, linewidth(0.5mm), Compose.arrow(), Compose.stroke("red"), Compose.fill("red"))

	function arrow(x, y)
		mid = (x .+ y) ./ 2
		t = nodestyle(:triangle, fill("red"), θ=π/2-atan((y .- x)...)-1π/6)
		ebr >> (x, y)
		t >> mid
	end

	Compose.set_default_graphic_size(15cm, 5cm)
	x = (0.1, 0.5)
	fi0 = (0.35, 0.5)
	fi1 = (0.7, 0.5)
	fi2 = (1.0, 0.5)
	img = canvas() do
		nb >> fi0
		nb >> fi1
		lb >> (fi0 .- (0.05, 0.1), "f_{i-1}")
		lb >> (fi1 .- (0.02, 0.1), "f_{i}")
		lb >> (x, "x")
		lb >> ((fi1 .+ fi0) ./ 2 .- (0.02, 0.0), raw"\vec{y}_{i}")
		lb >> ((fi1 .+ fi2) ./ 2 .- (0.05, 0.0), raw"\vec{y}_{i+1}")
		lb >> ((fi1 .+ fi2) ./ 2 .- (0.05, 0.0), "\\vec{y}_{i+1}")
		lb >> (x .- (0.00, 0.25), raw"\color{red}{1}")
		lb >> ((fi1 .+ fi0) ./ 2 .- (0.05, 0.45), raw"\color{red}{\frac{\partial \vec{y}_{i}}{\partial x}}")
		lb >> ((fi1 .+ fi2) ./ 2 .- (0.08, 0.45), raw"\color{red}{\frac{\partial \vec{y}_{i+1}}{\partial x}}")
		ebd >> (x, fi0)
		eb >> (fi0, fi1)
		eb >> (fi1, fi2)
		#arrow((fi1 .+ fi0) ./ 2 .+ (0.08, -0.3), (fi1 .+ fi2) ./ 2 .+ (-0.08, -0.3))
		arrow((fi1 .+ fi0) ./ 2 .+ (0.08, -0.3), (fi1 .+ fi2) ./ 2 .+ (-0.08, -0.3))
	end
	img
end

# ╔═╡ f12b25d8-7c78-4686-b46d-00b34e565605
let
	x = Dual(π/4, 1.0)
	for i=1:10
		x = sin(x)
	end
	x
end

# ╔═╡ d90c3cc9-084d-4cf7-9db7-42cea043030b
md"""
**Example:** Computing two gradients $\frac{\partial z\sin x}{\partial x}$ and $\frac{\partial \sin^2x}{\partial x}$ at one sweep
"""

# ╔═╡ 93c98cb2-18af-47df-afb3-8c5a34b4723c
let
	lb = textstyle(:math, fontsize(8), width=1.0, height=0.5)
	tb = textstyle(:default, fontsize(3.5), Compose.font("monospace"))
	tb_big = textstyle(:default, fontsize(4.5), fill("white"), Compose.font("monospace"))
	nb = nodestyle(:circle, fill("black"), Compose.stroke("transparent"); r=0.05)
	tri = nodestyle(:triangle, Compose.stroke("transparent"), fill("black"); r=0.02)
	eb = bondstyle(:default, linewidth(0.5mm))

	x_x = (0.1, 0.25)
	x_y = (0.9, 0.5)
	x_y2 = (0.9, 0.25)
	x_z = (0.3, 0.5)
	x_sin = (0.3, 0.25)
	x_mul = (0.5, 0.5)
	x_square = (0.5, 0.25)

	function arrow(x, y)
		mid = (x .+ y) ./ 2
		t = nodestyle(:triangle, θ=π/2-atan((y .- x)...)-1π/6)
		eb >> (x, y)
		t >> mid
	end

	img = canvas() do
		nb >> x_sin
		nb >> x_mul
		nb >> x_square
		tb_big >> (x_sin, "sin")
		tb_big >> (x_mul .+ (0, 0.01), "*")
		tb_big >> (x_square, "^2")
		arrow(x_sin, x_mul)
		arrow(x_x, x_sin)
		arrow(x_mul, x_y)
		arrow(x_square, x_y2)
		arrow(x_z, x_mul)
		arrow(x_sin, x_square)
		tb >> ((x_x .+ x_sin) ./ 2 .- (0.02, 0.04), "x+ϵˣ")
		tb >> ((x_sin .+ x_mul) ./ 2 .- (0.08, 0.04), "sin(x)+cos(x)*ϵˣ")
		tb >> ((x_y .+ x_mul) ./ 2 .- (-0.04, 0.055), "z*sin(x)\n+z*cos(x)*ϵˣ")
		tb >> ((x_y2 .+ x_square) ./ 2 .- (-0.04, 0.055), "sin(x)^2\n+2*sin(x)*cos(x)*ϵˣ")
		tb >> ((x_z .+ x_mul) ./ 2 .- (0.05, 0.02), "z")
	end

	Compose.set_default_graphic_size(100mm, 100mm/2)
	Compose.compose(context(0, -0.15, 1, 2), img)
end

# ╔═╡ 2dc74e15-e2ea-4961-b43f-0ada1a73d80a
md"so the gradients are $z\cos x$ and $2\sin x\cos x$"

# ╔═╡ 7ee75a15-eaea-462a-92b6-293813d2d4d7
md"""
**What if we want to compute gradients for multiple inputs?**

The computing time grows **linearly** as the number of variables that we want to differentiate. But does not grow significantly with the number of outputs.
"""

# ╔═╡ 02a25b73-7353-43b1-8738-e7ca472d0cc7
md"""
## Reverse mode automatic differentiation

"""

# ╔═╡ 2afb984f-624e-4381-903f-ccc1d8a66a17
md"On the other side, the back-propagation can differentiate **many inputs** with respect to a **single output** efficiently"

# ╔═╡ 7e5d5e69-90f2-4106-8edf-223c150a8168
md"""
```math
\begin{align}
    \frac{\partial \mathcal{L}}{\partial \vec y_i} = \frac{\partial \mathcal{L}}{\partial \vec y_{i+1}}&\boxed{\frac{\partial \vec y_{i+1}}{\partial \vec y_i}}\\
&\text{local jacobian?}
\end{align}
```
"""

# ╔═╡ 92d7a938-9463-4eee-8839-0b8c5f762c79
let
	lb = textstyle(:math, fontsize(8), width=0.5, height=0.5)
	tb = textstyle(:default, fontsize(10), Compose.font("monospace"))
	tb_big = textstyle(:default, fontsize(3.5), fill("white"), Compose.font("monospace"))
	nb = nodestyle(:circle, fill("white"), Compose.stroke("black"); r=0.08)
	tri = nodestyle(:triangle, Compose.stroke("transparent"), fill("black"); r=0.02)
	eb = bondstyle(:default, linewidth(0.5mm))
	ebr = bondstyle(:default, Compose.stroke("red"), linewidth(0.5mm))
	ebd = bondstyle(:default, linewidth(0.5mm), dashed=true)
	eba = bondstyle(:default, linewidth(0.5mm), Compose.arrow(), Compose.stroke("red"), Compose.fill("red"))

	function arrow(x, y)
		mid = (x .+ y) ./ 2
		t = nodestyle(:triangle, fill("red"), θ=π/2-atan((y .- x)...)-1π/6)
		ebr >> (x, y)
		t >> mid
	end

	Compose.set_default_graphic_size(15cm, 5cm)
	x = (0.1, 0.5)
	fi0 = (0.35, 0.5)
	fi1 = (0.7, 0.5)
	fi2 = (0.9, 0.5)
	img = canvas() do
		nb >> fi0
		nb >> fi1
		lb >> (fi0 .- (0.02, 0.1), "f_{i}")
		lb >> (fi1 .- (0.05, 0.1), "f_{i+1}")
		lb >> (fi2 .- (0.05, 0.0), raw"\mathcal{L}")
		lb >> ((fi0 .+ x) ./ 2 .- (0.05, 0.0), raw"\vec{y}_{i}")
		lb >> ((fi0 .+ fi1) ./ 2 .- (0.05, 0.0), raw"\vec{y}_{i+1}")
		lb >> ((fi0 .+ fi1) ./ 2 .- (0.05, 0.0), "\\vec{y}_{i+1}")
		lb >> (fi2 .- (0.05, 0.25), raw"\color{red}{1}")
		lb >> ((fi0 .+ x) ./ 2 .- (0.08, 0.45), raw"\color{red}{\frac{\partial \vec{y}_{i}}{\partial x}}")
		lb >> ((fi0 .+ fi1) ./ 2 .- (0.08, 0.45), raw"\color{red}{\frac{\partial \vec{y}_{i+1}}{\partial x}}")
		ebd >> (fi1, fi2)
		eb >> (fi0, fi1)
		eb >> (x, fi0)
		#arrow((fi1 .+ fi0) ./ 2 .+ (0.08, -0.3), (fi1 .+ fi2) ./ 2 .+ (-0.08, -0.3))
		arrow( (fi0 .+ fi1) ./ 2 .+ (-0.08, -0.3), (fi0 .+ x) ./ 2 .+ (0.05, -0.3),)
	end
	img
end

# ╔═╡ 4b1a0b59-ddc6-4b2d-b5f5-d92084c31e46
md"### How to visite local Jacobians in the reversed order? "

# ╔═╡ a7fc71a2-6d45-4162-8073-4ddb85ded2e8
md"
**Design Decision**

1. Compute forward pass and caching inetermediate results into a global stack $\Sigma$ （packages except NiLang），
2. reversible programming."

# ╔═╡ fb6c3a48-550a-4d2e-a00b-a1e40d86b535
md"""
**Example:** Computing the gradient $\frac{\partial z\sin x}{\partial x}$ and $\frac{\partial z\sin x}{\partial z}$ by back propagating cached local information.
"""

# ╔═╡ ab6fa4ac-29ed-4722-88ed-fa1caf2072f3
let
	lb = textstyle(:math, fontsize(10), width=1.0, height=0.5)
	tb = textstyle(:default, fontsize(3.5), Compose.font("monospace"))
	tbc = textstyle(:default, fontsize(3.5), fill("red"), Compose.font("monospace"))
	tb_big = textstyle(:default, fontsize(4), fill("white"), Compose.font("monospace"))
	nb = nodestyle(:circle, fill("black"), Compose.stroke("transparent"); r=0.05)
	tri = nodestyle(:triangle, Compose.stroke("transparent"), fill("black"); r=0.02)
	eb = bondstyle(:default, linewidth(0.5mm))

	x_x = (0.1, 0.2)
	x_y = (0.9, 0.5)
	x_z = (0.1, 0.7)
	x_sin = (0.3, 0.3)
	x_mul = (0.5, 0.5)

	function arrow(x, y)
		mid = (x .+ y) ./ 2
		t = nodestyle(:triangle, θ=π/2-atan((y .- x)...)-1π/6)
		eb >> (x, y)
		t >> mid
	end
	img1 = canvas() do
		nb >> x_sin
		nb >> x_mul
		tb_big >> (x_sin, "sin")
		tb_big >> (x_mul .+ (0, 0.01), "*")
		arrow(x_sin, x_mul)
		arrow(x_x, x_sin)
		arrow(x_mul, x_y)
		arrow(x_z, x_mul)
		tb >> ((x_x .+ x_sin) ./ 2 .- (0.0, 0.1), "x \n push(Σ,x)")
		tb >> ((x_sin .+ x_mul) ./ 2 .- (-0.15, 0.04), "s = sin(x) \n push(Σ,s)")
		tb >> ((x_y .+ x_mul) ./ 2 .- (-0.05, 0.04), "y = z*sin(x)")
		tb >> ((x_z .+ x_mul) ./ 2 .- (0.05, 0.07), "z\n push(Σ,z)")
	end
	img2 = canvas() do
		nb >> x_sin
		nb >> x_mul
		tb_big >> (x_sin, "sin")
		tb_big >> (x_mul .+ (0, 0.01), "*")
		arrow(x_mul, x_sin)
		arrow(x_sin, x_x)
		arrow(x_y, x_mul)
		arrow(x_mul, x_z)
		tb >> ((x_x .+ x_sin) ./ 2 .- (0.0, 0.1), "x = pop(Σ)\nx̄ = cos(x)*s̄")
		tb >> ((x_sin .+ x_mul) ./ 2 .- (-0.12, 0.04), "z = pop(Σ)\ns̄ = z*ȳ")
		tb >> ((x_y .+ x_mul) ./ 2 .- (-0.05, 0.06), "y\nȳ=1")
		tb >> ((x_z .+ x_mul) ./ 2 .- (0.05, 0.07), "s = pop(Σ)\nz̄ = s*ȳ")
	end

	Compose.set_default_graphic_size(150mm, 75mm/1.4)
	Compose.compose(context(),
	(context(0, -0.1, 0.5, 1.4), img1),
	(context(0.5, -0.1, 0.5, 1.4), img2)
	)
end

# ╔═╡ 8e72d934-e307-4505-ac82-c06734415df6
md"Here, we use $\overline y$ for $\frac{\partial \mathcal{L}}{\partial y}$, which is also called the adjoint."

# ╔═╡ e6ff86a9-9f54-474b-8111-a59a25eda506
md"### Primitives on different scales"

# ╔═╡ 9c1d9607-a634-4350-aacd-2d40984d647d
md"We call the leaf nodes defining AD rules \"**primitives**\""

# ╔═╡ 63db2fa2-50b2-4940-b8ee-0dc6e3966a57
md"
**Design Decision**

* A: If we define primitives on **arrays**, we need tons of manually defined backward rules. (Jax, Pytorch, Zygote.jl, ReverseDiff.jl et al.)
* B: If we define primitives on **scalar instructions**, we will have worse tensor performance. (Tapenade, Adept, NiLang et al.)

*Note*: Here, implementing AD on scalars means specifically the **optimal checkpointing** approach, rather than a package like Jax, Zygote and ReverseDiff that having scalar support.
"

# ╔═╡ 693167e7-e80c-401d-af89-55b5fae30848
let
	w, h = 0.22, 0.1
	lb = Compose.compose(context(), polygon([(-w, -h), (-w, h), (w, h), (w, -h)]), Compose.stroke("transparent"))
	lb2 = Compose.compose(context(), polygon([(-w, -h), (-w, h), (w, h), (w, -h)]), Compose.stroke("transparent"), fill("red"))
	tb = Compose.compose(context(), Compose.text(0.0, 0.0, ""), fontsize(3), Compose.font("monospace"))
	tb_big = textstyle(:default, fontsize(3), fill("white"), Compose.font("monospace"))
	eb = bondstyle(:default, linewidth(0.5mm))
	ar = bondstyle(:default, linewidth(0.3mm), Compose.arrow())
	xprog = (0.25, 0.15)
	xtensors = (0.25, 0.5)
	t1 = (0.5, 0.15)
	t2 = (0.5, 0.5)
	t3 = (0.5, 0.85)
	xscalars2 = (0.25, 0.85)

	function box(loc, text; color="black")
		(color=="black" ? lb : lb2) >> loc
		tb_big >> (loc, text)
	end
	Compose.set_default_graphic_size(10cm, 5cm)
	canvas() do
		box(xprog, "Program")
		ar >> (xprog, xtensors .+ (0, -h-0.03))
		#ar >> (xprog, xscalars .+ (-w/2, -h-0.03))
		ar >> (xtensors, xscalars2 .+ (0, -h-0.05))
		box(xtensors, "Functions on arrays")
		#box(xscalars, "Functions on Scalars")
		box(xscalars2, "Finite instructions"; color="red")
		tb >> (t1, "Neural networks")
		tb >> (t2, "matrix multiplication")
		tb >> (t3, "+, -, *")
	end
end

# ╔═╡ 4cd70901-2142-4868-9a33-c46ca0d064ec
html"""

on tensors
on finite instructions

meaning
defining backward rules manully for functions on tensors
defining backward rules on a limited set of basic scalar operations, and generate gradient code using source code transformation

pros and cons

Good tensor performance

Mature machine learning ecosystem

Need to define backward rules manually

Reasonalbe scalar performance

hard to utilize GPU kernels (except NiLang.jl) and BLAS

packages
Jax
PyTorch
Tapenade

Adept

NiLang.jl

"""

# ╔═╡ 4ff09f7c-aeac-48bd-9d58-8446137c3acd
md"""
## The AD ecosystem in Julia

Please check JuliaDiff: [https://juliadiff.org/](https://juliadiff.org/)

A short list:
* Forward mode AD: ForwardDiff.jl
* Reverse mode AD (tensor): ReverseDiff.jl/Zygote.jl
* Reverse mode AD (scalar): NiLang.jl

Warnings
* The main authors of `Tracker`, `ReverseDiff` and `Zygote` are not maintaining them anymore.
"""
#=
|       |   Rules | Favors Tensor? | Type |
| ---- | ---- | --- | --- |
|  Zygote   |  C  |  ✓   |   R     |
|  ReverseDiff  |  D    | ✓    | R |
|  Nabla   |  D→C  |   ✓  |   R     |
|  Tracker  |  D    | ✓    | R |
|  Yota   |  C  |  ✓   |     R   |
|  NiLang   |  -  |  ×   |  R      |
|  Enzyme   |  -  |  ×   |  R      |
|  ForwardDiff   |  -  |  ×   |    F    |
|  Diffractor   |  ?  |  ?   |  ?      |

* R: reverse mode
* F: forward mode
* C: ChainRules
* D: DiffRules
"""
=#

# ╔═╡ ea44037b-9359-4fbd-990f-529d88d54351
md"# Quick summary
1. The history of AD is longer than many people have thought. People are most familar with *reverse mode AD with primitives implemented on tensors* that brings the boom of machine learning. There are also AD frameworks that can differentiate a general program directly, which does not require users defining AD rules manually.
2. **Forward mode AD** propagate gradients forward, it has a computational overhead propotional to the number of input parameters.
2. **Backward mode AD** propagate gradients backward, it has a computational overhead propotional to the number of output parameters.
    * primitives on **tensors** v.s. **scalars**
    * reverse the program tape by **caching/checkpointing** v.s. **reversible programming**
4. Julia has one of the most active AD community!

#### Forward v.s. Backward
when is forward mode AD more useful?

* It is often combined with backward mode AD for obtaining Hessians (forward over backward).
* Having <20 input parameters.

when is backward mode AD more useful?
* In most variational optimizations, especially when we are training a neural network with ~ 100M parameters.
"

# ╔═╡ aa1547f2-5edd-4b7e-b93e-bdfc4e4fc6d5
md"""# The true reverse mode "automatic" differentiation"""

# ╔═╡ 71f4b476-027d-4c8f-b561-1ee418bc9e61
md"""
![](https://raw.githubusercontent.com/GiggleLiu/WuLiXueBao/master/paper/bennett_treeverse_pebbles.svg)
"""

# ╔═╡ ae096ad2-3ae9-4440-a959-0d7d9a174f1d
md"## How to differentiate sparse matrix multiplication"

# ╔═╡ 11557d6b-3a1e-416d-874f-b8d217976f76
md"## How to differentiate QR"

# ╔═╡ 6e7bf956-67b1-430c-9cf5-1746282727b2
md"## Time and space"

# ╔═╡ Cell order:
# ╟─a1ef579e-4b66-4042-944e-7e27c660095e
# ╠═f11023e5-8f7b-4f40-86d3-3407b61863d9
# ╟─9d11e058-a7d0-11eb-1d78-6592ff7a1b43
# ╟─b73157bf-1a77-47b8-8a06-8d6ec2045023
# ╟─ec13e0a9-64ff-4f66-a5a6-5fef53428fa1
# ╟─f8b0d1ce-99f7-4729-b46e-126da540cbbe
# ╟─435ac19e-1c0c-4ee5-942d-f2a97c8c4d80
# ╟─48ecd619-d01d-43ff-8b52-7c2566c3fa2b
# ╟─4878ce45-40ff-4fae-98e7-1be41e930e4d
# ╠═ce44f8bd-692e-4eab-9ba4-055b25e40c81
# ╠═b2c1936c-2c27-4fbb-8183-e38c5e858483
# ╠═8be1b812-fcac-404f-98aa-0571cb990f34
# ╟─33e0c762-c75e-44aa-bfe2-bff92dd1ace8
# ╟─c59c35ee-1907-4736-9893-e22c052150ca
# ╠═f12b25d8-7c78-4686-b46d-00b34e565605
# ╟─d90c3cc9-084d-4cf7-9db7-42cea043030b
# ╟─93c98cb2-18af-47df-afb3-8c5a34b4723c
# ╟─2dc74e15-e2ea-4961-b43f-0ada1a73d80a
# ╟─7ee75a15-eaea-462a-92b6-293813d2d4d7
# ╟─02a25b73-7353-43b1-8738-e7ca472d0cc7
# ╟─2afb984f-624e-4381-903f-ccc1d8a66a17
# ╟─7e5d5e69-90f2-4106-8edf-223c150a8168
# ╟─92d7a938-9463-4eee-8839-0b8c5f762c79
# ╟─4b1a0b59-ddc6-4b2d-b5f5-d92084c31e46
# ╟─a7fc71a2-6d45-4162-8073-4ddb85ded2e8
# ╟─fb6c3a48-550a-4d2e-a00b-a1e40d86b535
# ╟─ab6fa4ac-29ed-4722-88ed-fa1caf2072f3
# ╟─8e72d934-e307-4505-ac82-c06734415df6
# ╟─e6ff86a9-9f54-474b-8111-a59a25eda506
# ╟─9c1d9607-a634-4350-aacd-2d40984d647d
# ╟─63db2fa2-50b2-4940-b8ee-0dc6e3966a57
# ╟─693167e7-e80c-401d-af89-55b5fae30848
# ╟─4cd70901-2142-4868-9a33-c46ca0d064ec
# ╟─4ff09f7c-aeac-48bd-9d58-8446137c3acd
# ╟─ea44037b-9359-4fbd-990f-529d88d54351
# ╟─aa1547f2-5edd-4b7e-b93e-bdfc4e4fc6d5
# ╟─ae096ad2-3ae9-4440-a959-0d7d9a174f1d
# ╟─11557d6b-3a1e-416d-874f-b8d217976f76
# ╟─6e7bf956-67b1-430c-9cf5-1746282727b2
# ╟─71f4b476-027d-4c8f-b561-1ee418bc9e61


module TestGradient

using Printf

export test_gradient

struct GradientComponent
  i::Int
  g::Float64
  gbest::Float64
  error::Float64
  stepbest::Float64
end
Base.show(io::IO, c::GradientComponent) =
  @printf("%9i %12.5e %12.5e %12.5e %12.5e", c.i, c.g, c.gbest, c.error, c.stepbest)

title() = println("Component   Analytical     Discrete        Error    Best_Step")
function Base.show(io::IO, ::MIME"text/plain", c::AbstractVector{GradientComponent} )
  title()
  for i in 1:min(length(c),3)
    println(c[i])
  end
  if length(c) > 7
    @printf("%30s\n","⋮ ")
  end
  for i in max(4,length(c)-2):length(c)
    println(c[i])
  end
  max_error = findmax(map(x->x.error,c))
  print("""

  Maximum error = $(max_error[1]) at component: $(max_error[2])
  """)
  title()
  println(c[max_error[2]])
end

function discret(icomp,x,step,f)
  save = x[icomp]
  x[icomp] = save + step
  fplus = f(x)
  x[icomp] = save - step
  fminus = f(x)
  x[icomp] = save
  return (fplus - fminus) / (2 * step)
end

"""

```
test_gradient(f::Function,g!::Function,x::AbstractVector{T}) where T<:Real
```

Function that performs finite difference and analytical gradient
comparision. Used only for test purpouses. The output of `test_gradient`
is a vector of type `GradientComponent`, containing the index of the
component, the analytical and discrete gradientes, and the error associated
to the best step found for that component.

### Example

```julia-repl
julia> f(x) = sum(x.^2)
f (generic function with 1 method)

julia> g!(g,x) = 2*x
g! (generic function with 1 method)

julia> x = rand(20);

julia> test_gradient(f,g!,x)

Comparing analytical and finite-difference gradients... may take a while.

Function value on test point: 8.671516638852811

Component   Analytical     Discrete        Error    Best_Step
        1  1.38356e+00  1.38356e+00  8.66637e-15  1.00000e-02
        2  1.97347e+00  1.97347e+00  4.20806e-14  1.00000e-02
        3  1.53930e-01  1.53930e-01  3.17351e-14  1.00000e-02
                            ⋮
       18  1.89896e-01  1.89896e-01  1.12252e-13  1.00000e-02
       19  1.37622e+00  1.37622e+00  5.84064e-14  1.00000e-02
       20  1.62172e+00  1.62172e+00  2.19071e-15  1.00000e-02

Maximum error = 1.1225244035991744e-13 at component: 18
Component   Analytical     Discrete        Error    Best_Step
       18  1.89896e-01  1.89896e-01  1.12252e-13  1.00000e-02

```

"""
function test_gradient(f::Function,g!::Function,x::AbstractVector{T}) where T<:Real
  n = length(x)

  println("""

  Comparing analytical and finite-difference gradients... may take a while.
  """)

  fx = f(x)
  println("Function value on test point: ", fx)
  println()

  g = similar(x)
  g = g!(g,x)
  output = Vector{GradientComponent}(undef,0)
  t0 = time()
  eworst = 0.
  iworst = 1
  for i in 1:n
    step = 1.e-2
    stepbest = step
    error = +Inf
    gbest = 0.
    if mod(time() - t0,10) == 0
      println("Computing the $(i)th of $n components. Worst error: $eworst")
    end
    while error > 1.e-6 && step >= 1.e-20
      gcomp = discret(i,x,max(abs(g[i]*step),1e-20),f)
      if g[i] ≈ 0.
        steperror = abs(gcomp - g[i])
      else
        steperror = abs( ( gcomp - g[i] ) / g[i] )
      end
      if steperror < error
        error = steperror
        gbest = gcomp
        stepbest = step
      end
      step = step / 10
    end
    push!(output,GradientComponent(i, g[i], gbest, error, stepbest))
    if error > eworst
      iworst = i
      eworst = error
    end
  end
  return output
end

"""

```
test_gradient(f::Function,g!::Function,x::AbstractVector{T}) where T
```

Try to reinterpret the type of input into an array of `Float64`, and then
call `test_gradient(f,g,x::AbstractVector{<:Real})`

### Example

```julia-repl
julia> x = rand(SVector{2,Float64},10);

julia> f(x) = sum(v[1]^2 + v[2]^2 for v in x)
f (generic function with 1 method)

julia> g!(g,x) = @. g = (SVector{2,Float64}(2*v[1],2*v[2]) for v in x)
g! (generic function with 1 method)

julia> test_gradient(f,g!,x)

Comparing analytical and finite-difference gradients... may take a while.

Function value on test point: 6.5919481328739655

Computing the 1th of 20 components. Worst error: 0.0
Component   Analytical     Discrete        Error    Best_Step
        1  1.57869e+00  1.57869e+00  4.78215e-15  1.00000e-02
        2  1.04203e+00  1.04203e+00  1.91779e-14  1.00000e-02
        3  1.72248e+00  1.72248e+00  1.03128e-15  1.00000e-02
                            ⋮
       18  3.32460e-01  3.32460e-01  2.67153e-15  1.00000e-02
       19  1.03341e+00  1.03341e+00  8.16492e-15  1.00000e-02
       20  9.99455e-01  9.99455e-01  8.88663e-16  1.00000e-02

Maximum error = 2.2306785702197458e-12 at component: 7
Component   Analytical     Discrete        Error    Best_Step
        7  1.21440e-02  1.21440e-02  2.23068e-12  1.00000e-02

```

"""
function test_gradient(f::Function, g!::Function, x::AbstractVector{T}) where T
  test_gradient(
    x -> f(reinterpret(T,x)),
    (g,x) -> reinterpret(
      Float64,
      g!(
        reinterpret(T,g),
        reinterpret(T,x)
      ),
    ),
    reinterpret(Float64,x)
  )
end

end # module


using RungeKutta
using Base.Test

# TEST 1 - Correct answer.
x_truth_test1 = [1.0 0.5416666666666; 0.0 -0.8333333333333333];
t_truth_test1 = [0.0 1.0];

f_test1 = [(t,x) -> x[2], (t,x) -> -x[1]];
x0_test1 = [1.0, 0.0];
t0_test1 = 0.0;
h_test1 = 1.0;
n_test1 = 1;

t_answer_test1, x_answer_test1 =
    rk4f(f_test1, x0_test1, t0_test1, h_test1, n_test1);

@test_approx_eq t_answer_test1 t_truth_test1
@test_approx_eq x_answer_test1 x_truth_test1

# TEST 2 - Invalid argument dimension.
f_test2 = [(t,x) -> x[2], (t,x) -> -x[1]];
x0_test2 = [1.0, 0.0, -1.0];
t0_test2 = 0.0;
h_test2 = 1.0;
n_test2 = 1;

@test_throws ArgumentError rk4f(
    f_test2, x0_test2, t0_test2, h_test2, n_test2);

# TEST 3 - Invalid number of iterations.
f_test3 = f_test2;
x0_test3 = x0_test2;
t0_test3 = t0_test2;
h_test3 = h_test2;
n_test3 = -4;

@test_throws ArgumentError rk4f(
    f_test3, x0_test3, t0_test3, h_test3, n_test3);

# TEST 4 - Invalid step size.
f_test4 = f_test3;
x0_test4 = x0_test3;
t0_test4 = t0_test3;
h_test4 = -0.5;
n_test4 = n_test2;

@test_throws ArgumentError rk4f(
    f_test4, x0_test4, t0_test4, h_test4, n_test4);


# The functions in this module are a direct translation of the original
# Matlab functions by E. Doland and J. J. Moré.
# See http://www.mcs.anl.gov/~more/cops.
#
# D. Orban, 2015, 2016.

"""Compute performance ratios used to produce a performance profile.

There is normally no need to call this function directly.
See the documentation of `performance_profile()` for more information.
"""
function performance_ratios(T :: Array{Float64,2}; logscale :: Bool=true)

  (np, ns) = size(T);       # Number of problems and number of solvers.

  T[isinf.(T)] = NaN;
  T[T .< 0] = NaN;
  minperf = mapslices(NaNMath.minimum, T, 2) # Minimal (i.e., best) performance per solver

  # Compute ratios and divide by smallest element in each row.
  r = zeros(np, ns);
  for p = 1 : np
    r[p, :] = T[p, :] / minperf[p];
  end

  logscale && (r = log2.(r));
  max_ratio = NaNMath.maximum(r)

  # Replace failures with twice the max_ratio and sort each column of r.
  failures = isnan.(r);
  r[failures] = 2 * max_ratio;
  r = sort(r, 1);
  return (r, max_ratio)
end

performance_ratios{Td <: Number}(T :: Array{Td,2}; logscale :: Bool=true) = performance_ratios(convert(Array{Float64,2}, T), logscale=logscale)

"""Produce a performance profile.

Each column of the matrix `T` defines the performance data for a solver
(smaller is better). Failures on a given problem are represented by a
negative value, an infinite value, or `NaN`.
The optional argument `logscale` is used to produce a logarithmic (base 2)
performance plot.
"""
function performance_profile(T :: Array{Float64,2}, labels :: Vector{AbstractString};
                             logscale :: Bool=true,
                             title :: AbstractString="",
                             kwargs...)

  (ratios, max_ratio) = performance_ratios(T, logscale=logscale)
  (np, ns) = size(ratios)

  ratios = [ratios; 2.0 * max_ratio * ones(1, ns)]
  xs = [1:np+1;] / np
  length(labels) == 0 && (labels = [@sprintf("column %d", col) for col = 1 : ns])
  profile = Plots.plot(; kwargs...)  # initial empty plot
  for s = 1 : ns
    Plots.plot!(ratios[:, s], xs, t=:steppost, label=labels[s])
  end
  Plots.xlims!(logscale ? 0.0 : 1.0, 1.1 * max_ratio)
  Plots.ylims!(0, 1.1)
  Plots.xlabel!("Within this factor of the best" * (logscale ? " (log scale)" : ""))
  Plots.ylabel!("Proportion of problems")
  Plots.title!(title)
  return profile
end

performance_profile{Tn <: Number, S <: AbstractString}(T :: Array{Tn,2}, labels :: Vector{S}; kwargs...) =
  performance_profile(convert(Array{Float64,2}, T), convert(Vector{AbstractString}, labels); kwargs...)


# utilities

function elMax!{T<:Real}(x::Array{T},y::T)
  x[:] = max.(x,y)
end

function elMin!{T<:Real}(x::Array{T},y::T)
  x[:] = min.(x,y)
end

function resetArray!{T<:Real}(A::Array{T})
  A[:] = zero(eltype(A))
end

function logSumExpWeightsNorm(log_w::Vector{T} where T <: AbstractFloat)
  # function uses log-sum-exp trick to compute normalized weights of a
  # vector of numbers stored as logs, avoiding underflow
  max_entry = maximum(log_w)

  shifted_weights = log_w - max_entry
  return normalized_weights = exp.(shifted_weights)./sum(exp.(shifted_weights))

end

function logSumExpWeightsNorm!(p::Vector{T},log_w::Vector{T}) where T <: AbstractFloat
  # function uses log-sum-exp trick to compute normalized weights of a
  # vector of numbers stored as logs, avoiding underflow
  max_entry = maximum(log_w)

  shifted_weights = log_w - max_entry
  p[:] = exp.(shifted_weights)./sum(exp.(shifted_weights))

end

function logSumExpWeights(log_w::Vector{T} where T <: AbstractFloat)
  # function uses log-sum-exp trick to compute UNnormalized weights of a
  # vector of numbers stored as logs, avoiding underflow
  max_entry = maximum(log_w)

  shifted_weights = exp.(log_w .- max_entry)

end

function logSumExpWeights!(lse_w::Vector{T},log_w::Vector{T}) where T <: AbstractFloat
  # function uses log-sum-exp trick to compute UNnormalized weights of a
  # vector of numbers stored as logs, avoiding underflow
  max_entry = maximum(log_w)

  lse_w[:] = exp.(log_w .- max_entry)

end

function OneHot(n::Int64,k::Int64)

  onehotvector = spzeros(n)
  onehotvector[k] = 1
  return onehotvector

end

function stratifiedResample(weights::Vector{Float64},n_resample::Int64)::Vector{Int64}
  sample_idx = zeros(Int64,n_resample)
  stratifiedResample!(sample_idx,weights,n_resample)
  return sample_idx
end

function stratifiedResample!(sample_idx::Vector{Int64},weights::Vector{Float64},n_resample::Int64)
"""
 Implements stratified resampling as described in Appendix B of
 Fearnhead and Clifford (2003), JRSSB 65(4) 887--899

 weights is a vector of weights; n_resample is the number of resampled
 indices to return
"""
  norm_weights = cumsum(weights./sum(weights))::Vector{Float64} # normalize
  buckets = collect(0.0:(1.0/n_resample):1.0)[1:n_resample]
  locs = rand(n_resample)./n_resample .+ buckets

  sample_idx[:] = [ findfirst( locs[i] .<= norm_weights ) for i in 1:n_resample ]
  # return sample_idx
end

function systematicResample(weights::Vector{Float64},n_resample::Int64)::Vector{Int64}
  sample_idx = zeros(Int64,n_resample)
  systematicResample!(sample_idx,weights,n_resample)
  return sample_idx
end

function systematicResample!(sample_idx::Vector{Int64},weights::Vector{Float64},n_resample::Int64)

  norm_weights = cumsum(weights./sum(weights))
  buckets = collect(0.0:(1.0/n_resample):1.0)[1:n_resample]
  U = rand()*1.0/n_resample
  locs = (buckets .+ U)

  sample_idx[:] = [ findfirst( locs[i] .<= norm_weights ) for i in 1:n_resample ]
  # return sample_idx
end

function tally_ints(Z::Vector{Int},K::Int)
    """
    counts occurrences in `Z` of integers 1 to `K`
    - `Z`: vector of integers
    - `K`: maximum value to count occurences in `Z`
    """
    ret = zeros(Int,K)
    n = size(Z,1)
    idx_all = 1:n
    idx_j = trues(n)
    for j in 1:K
      for i in idx_all[idx_j]
        if Z[i]==j
          ret[j] += 1
          idx_j[i] = false
        end
      end
    end
    return ret
end

function adj2edgelist(A::T where T<:Union{Array{Int64,2},SparseMatrixCSC{Int64,Int64},SparseMatrixCSC{Float64,Int64}})::Array{Int64,2}
  adj2edgelist(A,false)
end

function adj2edgelist(A::T where T<:Union{Array{Int64,2},SparseMatrixCSC{Int64,Int64},SparseMatrixCSC{Float64,Int64}},multigraph::Bool)::Array{Int64,2}
"""
  Convert binary adjacency matrix A to edge list
"""

  M,N,V = findnz(tril(A))
  if multigraph
    return hcat(M,N),V
  else
    return hcat(M,N)
  end
end

function edgelist2adj(edgelist::Array{Int64,2})::SparseMatrixCSC{Int64,Int64}
"""
  Convert edge list to binary adjacency matrix
"""

  edgelist2adj(edgelist,Int64)
end

function edgelist2adj(edgelist::Array{Int64,2},value_type::DataType)
"""
  Convert edge list to binary adjacency matrix, with non-zero values of type value_type
"""
  nv = maximum(edgelist)
  A = sparse([edgelist[:,1];edgelist[:,2]],[edgelist[:,2];edgelist[:,1]],ones(value_type,2*size(edgelist,1)),nv,nv)
end

function edgelist2adj!(A::Union{Array{Float64,2},SparseMatrixCSC{Int64,Int64},SparseMatrixCSC{Float64,Int64}},edge_list::Array{Int64,2})
"""
  Convert edge list to binary adjacency matrix, with non-zero values of type value_type
"""
  et = eltype(A)
  for i=1:size(edge_list,1)
    A[edge_list[i,1],edge_list[i,2]] = one(et)
    A[edge_list[i,2],edge_list[i,1]] = one(et)
  end

end

function normalizedLaplacian(edgelist::Array{Int64,2})::SparseMatrixCSC{Int64,Int64}
"""
  Construct (sparse) normalized Laplacian matrix
"""
  degrees = getDegrees(edgelist)
  nv = maximum(edgelist)
  L = sparse(1*I,nv,nv) - sparse(Diagonal( degrees.^(-1/2) )) * edgelist2adj(edgelist,Float64) * sparse(Diagonal( degrees.^(-1/2) ))
end

function denseNormalizedLaplacian!(L::Array{Float64,2},adj::SparseMatrixCSC{Float64,Int64},degrees::Array{Float64,1},nv::Int64)
"""
  Construct a symmetric (dense) normalized Laplacian matrix
"""
  # degrees = getDegrees(edgelist)
  # nv = maximum(edgelist)
  L[1:nv,1:nv] = Array(sparse(1*I,nv,nv) - sparse(Diagonal( degrees.^(-1/2) )) * adj * sparse(Diagonal( degrees.^(-1/2) )))
end

function denseNormalizedLaplacian!(L::Array{Float64,2},adj::SparseMatrixCSC{Int,Int64},degrees::Array{Float64,1},nv::Int64)
"""
  Construct a symmetric (dense) normalized Laplacian matrix
"""
  # degrees = getDegrees(edgelist)
  # nv = maximum(edgelist)
  L[1:nv,1:nv] = Array(sparse(1*I,nv,nv) - sparse(Diagonal( degrees.^(-1/2) )) * adj * sparse(Diagonal( degrees.^(-1/2) )))
end

function denseNormalizedLaplacian!(L::Array{Float64,2},edge_list::Array{Int64,2},degrees::Array{Float64,1},nv::Int64)
"""
  Construct a symmetric (dense) normalized Laplacian matrix. L should be all zeros.
"""
  # degrees = getDegrees(edgelist)
  # nv = maximum(edgelist)

  # set diagonal
  for n = 1:nv
    L[n,n] = 1.0
  end

  rt_deg = degrees.^(-1/2)
  for m = 1:size(edge_list,1)
    L[edge_list[m,1],edge_list[m,2]] = -(rt_deg[edge_list[m,1]]*rt_deg[edge_list[m,2]])
    L[edge_list[m,2],edge_list[m,1]] = -(rt_deg[edge_list[m,1]]*rt_deg[edge_list[m,2]])
  end

  # L[1:nv,1:nv] = Array(sparse(1*I,nv,nv) - sparse(Diagonal( degrees.^(-1/2) )) * adj * sparse(Diagonal( degrees.^(-1/2) )))
end

function getDegrees(edgelist::Array{Int64,2})::Array{Float64,1}
"""
  Compute degrees of edge list
"""
  return convert(Array{Float64,1},tally_ints(edgelist[:],maximum(edgelist)))
end

function generateEigenSystem!(L::Array{Float64,2},nv::Int64)::Tuple{Array{Float64,1},Array{Float64,2}}
  LL = Symmetric(L[1:nv,1:nv])
  eigvec_tr = zeros(Float64,nv,nv)
  # THIS OVERWRITES L!
  esys_val,esys_vec = LAPACK.syevr!('V','A',LL.uplo,LL.data,-0.0,0.0,0,0,-1.0)
  transpose!(eigvec_tr,esys_vec) # transposed to optimize computation

  return esys_val,eigvec_tr
end

function generateEigenSystem(L::Array{Float64,2},nv::Int64)::Tuple{Array{Float64,1},Array{Float64,2}}
  Lcopy = copy(L)
  LL = Symmetric(Lcopy[1:nv,1:nv])
  eigvec_tr = zeros(Float64,nv,nv)
  # THIS OVERWRITES Lcopy!
  esys_val,esys_vec = LAPACK.syevr!('V','A',LL.uplo,LL.data,-0.0,0.0,0,0,-1.0)
  transpose!(eigvec_tr,esys_vec) # transposed to optimize computation

  return esys_val,eigvec_tr
end

function updateEigenSystem!(L::Array{Float64,2},nv::Int64,p_state::ParticleState)

  denseNormalizedLaplacian!(L,p_state.vertex_unmap[p_state.edge_list],p_state.degrees,nv)
  LL = Symmetric(L[1:nv,1:nv])
  # THIS OVERWRITES L!
  esys_val,esys_vec = LAPACK.syevr!('V','A',LL.uplo,LL.data,-0.0,0.0,0,0,-1.0)
  # esys = eigfact!(Symmetric(L[1:nv,1:nv]))
  esys_val[1] = zero(Float64) # for numerical stability
  esys_val[end] = min(esys_val[end],Float64(2.0)) # for numerical stability

  for i=1:nv
    p_state.eig_vals[i] = esys_val[i]
    for j=1:nv
      # THIS TRANSPOSES THE EIGENVECTOR MATRIX TO OPTIMIZE COMPUTATION
      # ROWS OF p_state.eig_vecs CORRESPOND TO EIGENVECTORS OF L
      p_state.eig_vecs[j,i] = esys_vec[i,j]
    end
  end
  p_state.has_eigensystem[:] = true

end

function saveSamples!(alpha_samples::Array{Float64,1},lambda_samples::Array{Float64,1},
                      particle_trajectory_samples::Array{Int64,2},edge_sequence_samples::Array{Int64,2},
                      s_state::SamplerState,particle_container::Array{Array{ParticleState,1},1},
                      p_idx::Int64,n_sample::Int64)

  alpha_samples[n_sample] = s_state.α[1]
  lambda_samples[n_sample] = s_state.λ[1]
  particle_trajectory_samples[n_sample,:] = s_state.particle_path[:]
  edge_sequence_samples[n_sample,:] = particle_container[p_idx][end].edge_idx_list[:]

end


# # SEIR
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/models/SEIR.ipynb)
#
#md # !!! note "Overview"
#md #     System type: Nonlinear system\
#md #     State dimension: 7\
#md #     Application domain: Epidemiology
#
# ## Model description
#
# The SEIR Model is an Compartmental model, these try to predict things such as
# how a disease spreads, or the total number infected, or the duration of an
# epidemic, and to estimate various epidemiological parameters such as the
# reproductive number.
# The dynamics are described as follows:
#
# ```math
#   \left\{ \begin{array}{lcl}
#   \frac{dS}{dt} & = & \beta I S \\
#   \frac{dE}{dt} & = & \beta I S - \alpha E \\
#   \frac{dI}{dt} & = & -\gamma I - \alpha E \\
#   \frac{dR}{dt} & = & \gamma I
#   \end{array} \right.
# ```
#
# where ``S`` is the stock of susceptible population, ``E`` is the stock of
# exposed population``I`` is the stock of infected, ``R`` is the stock of
# removed population (either by death or recovery), with ``S + E + I + R = N``

using ReachabilityAnalysis, Plots

@taylorize function seir2!(du,u,p,t)
  S, E, I, R, α, β, γ = u

  βIS = β * (I * S)
  αE = α*E
  γI = γ*I

  du[1] = -βIS      # dS
  du[2] = βIS - αE  # dE
  du[3] = -γI + αE  # dI
  du[4] = γI        # dR

  #uncertain parameters
  local zerou = zero(u[1])
  du[5] = zerou
  du[6] = zerou
  du[7] = zerou
end

E₀ = 1e-4
u₀ = [1-E₀, E₀, 0, 0]
param_error = 0.01
α = 0.2 ± param_error
β = 1.0 ± 0.0
γ = 0.5 ± param_error
p = [α, β, γ]
X0 = IntervalBox(vcat(u₀, p));
prob = @ivp(x' = seir2!(x), dim=7, x(0) ∈ X0);

# ## Reachability settings
#
# The initial values used were ``E₀ = 1e-4, u₀ = [1-E₀, E₀, 0, 0]``,
# ``α = 0.2 ± 0.01, β = 1.0 ± 0.0, γ = 0.5 ± 0.01`` and ``p = [α, β, γ]``, for a
# time span of 200.
# The algorithm used was `TMJets` with ``n_T=7`` and ``n_Q=1``

# ## Results

sol = solve(prob, tspan=(0.0, 200.0), alg=TMJets21a(orderT=7, orderQ=1));
solz = overapproximate(sol, Zonotope);

LazySets.set_ztol(Float64, 1e-13)
fig = plot(legend=:outerright)
plot!(fig, solz, vars=(0, 1), color=:blue, lw=0.0, lab="S")
plot!(fig, solz, vars=(0, 2), color=:green, lw=0.0, lab="E")
plot!(fig, solz, vars=(0, 3), color=:red, lw=0.0, lab="I")
plot!(solz, vars=(0, 4), color=:grey, lw=0.0, lab="R")
fig


# --------------
# Initialization
# --------------
using LLLplus
using Test
using LinearAlgebra
using Documenter

# --------------
# tests with small matrices
# --------------
println("tests with small matrices...")

# Matrix from http://home.ie.cuhk.edu.hk/~wkshum/wordpress/?p=442
H =[1   9   1   2;
    1   8   8   3;
    7   4   5   1;
    2   6   7   1];
B1=[2   3  -2  -4
    3  -1   2   1
    1   1   6  -4
    1   3  -1   3];
T1=[0   0   1  -1
    0   1  -1   0
    0   0   0   1
    1  -3   3  -2];
(B,T) = lll(H);
@test sum(abs.(B-B1))==0
@test sum(abs.(T-T1))==0
println("...done")

# --------------
# Run tests with random matrices
# --------------

println("\nIn all the following tests, the first time includes the "*
        "JIT compilation; \nfor the second execution the "*
        "compilation is already done and the time\n"*
        "should be faster.\n")

# Time LLL decomposition of a 1000x1000 real matrix with randn entries
N = 1000;
H = randn(N,N);
println("Testing LLL on $(N)x$(N) real matrix...")
@time (B,T,Q,R) = lll(H);
@time (B,T,Q,R) = lll(H);

# Time LLL, Seysen, VBLAST decompositions of a 10x10 complex matrix with
# randn entries
N = 10;
println("\nTesting LLL on $(N)x$(N) complex matrix...")
H = randn(N,N) + im*randn(N,N);
@time (B,T,Q,R) = lll(H);
@time (B,T,Q,R) = lll(H);
println("Testing Seysen on same $(N)x$(N) complex matrix...")
@time (B,T) = seysen(H);
@time (B,T) = seysen(H);
println("Testing VBLAST on same $(N)x$(N) complex matrix...")
@time (W,P,B) = vblast(H);
@time (W,P,B) = vblast(H);
println("Testing Brun on real part of same $(N)x$(N) matrix...")
@time (B,T) = brun(real.(H));
@time (B,T) = brun(real.(H));

# Test CVP
N=500;
println("\nTesting cvp on $(N)x$(N) BPSK system...")
H=randn(N,N);
Q,R=qr(H);
u=Int.(rand([-1,1],N));
y=H*u+rand(N)/100;
uhat=cvp(Q'*y,R,Val(false),-1,1);
errRate = sum(abs.(u-uhat))
println("Error Rate is $(errRate). It should be zero or very small.\n")

# --------------
# doctests
# --------------
if VERSION >= v"1.1.0"
    println("\nRunning doctests...")
    DocMeta.setdocmeta!(LLLplus, :DocTestSetup,
                        :(using LLLplus, LinearAlgebra,Random); recursive=true)
    doctest(LLLplus)
end


## squared_exponential.jl : implementation of squared exponential covariance function

## SquaredExponential ##
struct SquaredExponential{T} <: IsotropicCovarianceStructure{T}
    λ::T
    σ::T
    p::T

    function SquaredExponential{T}(λ::T, σ::T, p::T) where T
        λ > 0 || throw(DomainError(λ, "correlation length λ of squared exponential covariance cannot be negative or zero"))
        σ > 0 || throw(DomainError(σ, "marginal standard deviation σ of squared exponential covariance cannot be negative or zero"))
        p >= 1 || throw(DomainError(p, "in p-norm, p must be greater than or equal to 1"))
        isinf(p) && throw(DomainError(p, "in p-norm, p cannot be infinity"))

        new{T}(λ, σ, p)
    end
end

"""
    SquaredExponential(λ, σ=1, p=2)

Create a squared exponential (Gaussian) covariance structure with correlation length `λ`, (optional) marginal standard deviation `σ` and (optional) `p`-norm.

# Examples
```jldoctest
julia> s1 = SquaredExponential(0.1)
Gaussian (λ=0.1, σ=1.0, p=2.0)

julia> s2 = SquaredExponential(0.1, σ=2.)
Gaussian (λ=0.1, σ=2.0, p=2.0)

```
"""
SquaredExponential(λ::Real; σ::Real=1.0, p::Real=2) =
    SquaredExponential{promote_type(typeof(λ),typeof(σ),typeof(p))}(promote(λ, σ, p)...)

"""
    Gaussian(λ, σ=1, p=2)

Create a Gaussian (squared exponential) covariance structure with correlation length `λ`, (optional) marginal standard deviation `σ` and (optional) `p`-norm.

# Examples
```jldoctest
julia> g1 = Gaussian(0.1)
Gaussian (λ=0.1, σ=1.0, p=2.0)

julia> g2 = Gaussian(0.1, σ=2.)
Gaussian (λ=0.1, σ=2.0, p=2.0)

```
"""
const Gaussian = SquaredExponential

# evaluate Gaussian covariance
apply(s::SquaredExponential, x::Real) = exp(-(x / s.λ)^2)

# short name
shortname(::SquaredExponential) = "Gaussian"