# This is part of Calendars.jl. See the copyright note there.
# ========================= ISO dates =======================

# Day number of the start of the ISO calendar.
# ID-0001-01-01 = CE-0001-01-01 ~ DN#1 = EpochIso.
# ID-1111-45-04 = CE-1111-11-11 ~ DN#0405733
# ID-9999-52-05 = CE-9999-12-31 ~ DN#3652059
const EpochIso = 1
const ValidYearsIso = (0, 9999) # by convention

# Day of week in ISO is represented by an integer from 1 through 7,
# beginning with Monday and ending with Sunday, which is
# different from our enumeration, so we have to adjust.
function DayOfWeekIso(dn::DPart)
    day = rem(dn, weeklen)
    return day == 0 ? 7 : day
end

# True if year is an Iso leap year.
function isLeapYearIso(year::DPart)
    f(y) = rem(y + div(y, 4) - div(y, 100) + div(y, 400), 7)
    f(year) == 4 || f(year - 1) == 3
end

# Return the last week of the year for the Iso calendar.
function LastWeekOfYearIso(year::DPart)
    isLeapYearIso(year) ? 53 : 52
end

# Is the date a valid Iso date?
function isValidDateIso(cd::CDate, warn=true)
    cal, year, week, day = cd
    val = ( cal == ID
        && (ValidYearsIso[1] <= year && year <= ValidYearsIso[2])
        && (1 <= week && week <= LastWeekOfYearIso(year))
        && (1 <= day && day <= 7)
    )
    ! val && warn && @warn(Warning(cd))
    return val
end

# Represents the day number a valid Iso date?
function isValidDateIso(dn::DPart, warn=true)
    # val = EpochGregorian - 2 <= dn
    val = -1 <= dn
    ! val && warn && @warn(Warning(dn))
    return val
end

# Return the days this year so far.
# Does not depend on 'year' but keep it in the signature.
function DayOfYearIso(year::DPart, week::DPart, day::DPart)
    day += 7 * (week - 1)  # days in prior weeks this year
    return day - 1
end

XdayOnOrBefore(dn, x) = dn - rem(dn - x, weeklen)

# Return the number of days in prior years.
function DaysInPriorYears(year::DPart)
    # dn = DayNumberGregorian(year, 1, 4)
    # Let's make it self-contained!
    dn = ( 4           # days this year
    + 365 * year       # days in previous years ignoring leap days
    + div(year, 4)     # Julian leap days before this year...
    - div(year, 100)   # ...minus prior century years...
    + div(year, 400)   # ...plus prior years divisible by 400
    )
    return XdayOnOrBefore(dn, 1)
end

# Returns the day number from a valid ISO date.
function DayNumberIso(year::DPart, week::DPart, day::DPart)
    (
    DaysInPriorYears(year - 1) # days in prior years
    + 7 * (week - 1)           # days in prior weeks this year
    + day - 1                  # prior days this week
    )
end

# Computes the day number from a valid ISO date.
DayNumberIso(d::CDate) = DayNumberIso(d[2], d[3], d[4])

# Computes the ISO date from a day number.
function DateIso(dn::DPart)
    ! isValidDateIso(dn) && return InvalidDate

    # Patch in the two allowed exceptions.
    dn == -1 && return (ID, 0, 52, 6)
    dn ==  0 && return (ID, 0, 52, 7)

    # year = DateGregorian(dn - 3)[2]
    # Let's make it self-contained!
    year = div(400 * (dn - 4), 146097) + 1

    if dn >= YearStartIso(year + 1)
        year += 1
    end
    week = 1 + div(dn - YearStartIso(year), weeklen)
    day = DayOfWeekIso(dn)
    return (ID, year, week, day)::CDate
end

# Return the day number of the first day in the given Iso year.
function YearStartIso(year::DPart)
    return DayNumberIso(year, 1, 1)
end

# Return the day number of the last day in the given Iso year.
function YearEndIso(year::DPart)
    return DayNumberIso(year + 1, 1, 1) - 1
end

# Return the number of weeks in the given Iso year.
function WeeksInYearIso(year::DPart)
    return LastWeekOfYearIso(year)
end

# Return the number of days in the given Iso year.
function DaysInYearIso(year::DPart)
    return YearStartIso(year + 1) - YearStartIso(year)
end


if !isdefined(Main, :CustomStencil)
    include("../module/CustomStencil.jl")
    using Main.CustomStencil
end
include("misc.jl")
coefs = Float64.([1 1 1])

thumbstack_stencil = @def_stencil_expression(
            @sum(i,-2,2,(
                @sum(j,-2,2, c[max(abs(i),abs(j))]*D[x+i,y+j,z])
            )) + @sum(i,1,2, c[i]*(D[x,y,z+i]+D[x,y,z-i]))
)

st_inst = NewStencilInstance(thumbstack_stencil, 2, coefs)

radius = 2
nx = 8
ny = 8
nz = 8
dx = 1<<(nx)
dy = 1<<(ny)
dz = 1<<(nz)

t_steps = 5
## Actual time steps
data = CreateData(dx,dy,dz,radius)
dout = ApplyStencil(st_inst, data, t_steps)
#heatmap(dout1[:,:,Int(round(dz/2))])
o = Array(dout)


@doc doc"""
    Grassmann{n,k,F} <: Manifold

The Grassmann manifold $\operatorname{Gr}(n,k)$ consists of all subspaces spanned
by $k$ linear independent vectors $\mathbb F^n$, where
$\mathbb F \in \{\mathbb R, \mathbb C\}$ is either the real- (or complex-) valued vectors.
This yields all $k$-dimensional subspaces of $\mathbb R^n$ for the real-valued case and all
$2k$-dimensional subspaces of $\mathbb C^n$ for the second.

The manifold can be represented as

````math
\operatorname{Gr}(n,k) \coloneqq \bigl\{ \operatorname{span}(x)
: x \in \mathbb F^{n\times k}, \bar{x}^\mathrm{T}x = I_k\},
````

where ${\bar\cdot}^{\mathrm{T}}$ denotes the complex conjugate transpose and
$I_k$ is the $k\times k$ identity matrix. This means, that the columns of $x$
form an orthonormal basis of the subspace, that is a point on
$\operatorname{Gr}(n,k)$, and hence the subspace can actually be represented by
a whole equivalence class of representers.
Another interpretation is, that

````math
\operatorname{Gr}(n,k) = \operatorname{St}(n,k) / \operatorname{O}(k),
````

i.e the Grassmann manifold is the quotient of the [`Stiefel`](@ref) manifold and
the orthogonal group $\operatorname{O}(k)$ of orthogonal $k\times k$ matrices.

The tangent space at a point (subspace) $x$ is given by

````math
T_x\mathrm{Gr}(n,k) = \bigl\{
v \in \mathbb{F}^{n\times k} :
{\bar v}^{\mathrm{T}}x + {\bar x}^{\mathrm{T}}v = 0_{k} \bigr\},
````

where $0_{k}$ denotes the $k\times k$ zero matrix.

Note that a point $x\in \operatorname{Gr}(n,k)$ might be represented by
different matrices (i.e. matrices with orthonormal column vectors that span
the same subspace). Different representations of $x$ also lead to different
representation matrices for the tangent space $T_x\mathrm{Gr}(n,k)$

The manifold is named after
[Hermann G. Graßmann](https://en.wikipedia.org/wiki/Hermann_Grassmann) (1809-1877).

# Constructor

    Grassmann(n,k,F=ℝ)

Generate the Grassmann manifold $\operatorname{Gr}(n,k)$, where the real-valued
case $\mathbb F = \mathbb R$ is the default.
"""
struct Grassmann{n,k,F} <: Manifold end

Grassmann(n::Int, k::Int, F::AbstractNumbers = ℝ) = Grassmann{n,k,F}()

@doc doc"""
    check_manifold_point(M::Grassmann{n,k,F}, x)

Check whether `x` is representing a point on the [`Grassmann`](@ref) `M`, i.e. its
a `n`-by-`k` matrix of unitary column vectors and of correct `eltype` with respect to `F`.
"""
function check_manifold_point(M::Grassmann{n,k,F}, x; kwargs...) where {n,k,F}
    if (F === ℝ) && !(eltype(x) <: Real)
        return DomainError(
            eltype(x),
            "The matrix $(x) is not a real-valued matrix, so it does noe lie on the Grassmann manifold of dimension ($(n),$(k)).",
        )
    end
    if (F === ℂ) && !(eltype(x) <: Real) && !(eltype(x) <: Complex)
        return DomainError(
            eltype(x),
            "The matrix $(x) is neiter real- nor complex-valued matrix, so it does noe lie on the complex Grassmann manifold of dimension ($(n),$(k)).",
        )
    end
    if size(x) != representation_size(M)
        return DomainError(
            size(x),
            "The matrix $(x) is does not lie on the Grassmann manifold of dimension ($(n),$(k)), since its dimensions are wrong.",
        )
    end
    c = x' * x
    if !isapprox(c, one(c); kwargs...)
        return DomainError(
            norm(c - one(c)),
            "The point $(x) does not lie on the Grassmann manifold of dimension ($(n),$(k)), because x'x is not the unit matrix.",
        )
    end
end

@doc doc"""
    check_tangent_vector(M::Grassmann{n,k,F}, x, v)

Check whether `v` is a tangent vector in the tangent space of `x` on the [`Grassmann`](@ref)
`M`, i.e. that `v` is of size and type as well as that

````math
    x^{\mathrm{H}}v + v^{\mathrm{H}}x = 0_k,
````

where $\cdot^{\mathrm{H}}$ denotes the complex conjugate transpose or Hermitian and $0_k$
denotes the $k\times k$ zero natrix.
"""
function check_tangent_vector(G::Grassmann{n,k,F}, x, v; kwargs...) where {n,k,F}
    t = check_manifold_point(G, x)
    t === nothing || return t
    if (F === ℝ) && !(eltype(v) <: Real)
        return DomainError(
            eltype(v),
            "The matrix $(v) is not a real-valued matrix, so it can not be a tangent vector to the Grassmann manifold of dimension ($(n),$(k)).",
        )
    end
    if (F === ℂ) && !(eltype(v) <: Real) && !(eltype(v) <: Complex)
        return DomainError(
            eltype(v),
            "The matrix $(v) is neiter real- nor complex-valued matrix, so it can not bea tangent vector to the complex Grassmann manifold of dimension ($(n),$(k)).",
        )
    end
    if size(v) != representation_size(G)
        return DomainError(
            size(v),
            "The matrix $(v) is does not lie in the tangent space of $(x) on the Grassmann manifold of dimension ($(n),$(k)), since its dimensions are wrong.",
        )
    end
    if !isapprox(x' * v + v' * x, zeros(k, k); kwargs...)
        return DomainError(
            norm(x' * v + v' * x),
            "The matrix $(v) is does not lie in the tangent space of $(x) on the Grassmann manifold of dimension ($(n),$(k)), since x'v + v'x is not the zero matrix.",
        )
    end
end

@doc doc"""
    distance(M::Grassmann, x, y)

Compute the Riemannian distance on [`Grassmann`](@ref) manifold `M`$= \mathrm{Gr}(n,k)$.

Let $USV = x^\mathrm{H}y$ denote the SVD decomposition of
$x^\mathrm{H}y$, where $\cdot^{\mathrm{H}}$ denotes the complex
conjugate transposed or Hermitian. Then the distance is given by
````math
d_{\mathrm{GR}(n,k)}(x,y) = \operatorname{norm}(\operatorname{Re}(b)).
````
where

$b_{i}=\begin{cases} 0 & \text{if} \; S_i \geq 1\\ \operatorname{acos}(S_i) & \, \text{if} \; S_i<1 \end{cases}.$
"""
function distance(M::Grassmann, x, y)
    x ≈ y && return zero(real(eltype(x)))
    a = svd(x' * y).S
    a[a.>1] .= 1
    return sqrt(sum((acos.(a)) .^ 2))
end

@doc doc"""
    exp(M::Grassmann, x, v)

Compute the exponential map on the [`Grassmann`](@ref) `M`$= \mathrm{Gr}(n,k)$ starting in
`x` with tangent vector (direction) `v`. Let $v = USV$ denote the SVD decomposition of $v$.
Then the exponential map is written using

````math
z = x V\cos(S)V^\mathrm{H} + U\sin(S)V^\mathrm{H},
````

where $\cdot^{\mathrm{H}}$ denotes the complex conjugate transposed or Hermitian.
The cosine and sine are applied element wise to the diagonal entries of $S$.
A final QR decomposition $z=QR$ is performed for numerical stability reasons,
yielding the result as
````math
\exp_x v = Q.
````
"""
exp(::Grassmann, ::Any...)

function exp!(M::Grassmann, y, x, v)
    norm(M, x, v) ≈ 0 && return copyto!(y, x)
    d = svd(v)
    z = x * d.V * Diagonal(cos.(d.S)) * d.Vt + d.U * Diagonal(sin.(d.S)) * d.Vt
    # reorthonormalize
    return copyto!(y, Array(qr(z).Q))
end

injectivity_radius(::Grassmann) = π / 2

@doc doc"""
    inner(M::Grassmann, x, v, w)

Compute the inner product for two tangent vectors `v`, `w` from the
tangent space of `x` on the [`Grassmann`](@ref) manifold `M`.
The formula reads

````math
g_x(v,w) = \operatorname{trace}(v^{\mathrm{H}}w),
````

where $\cdot^{\mathrm{H}}$ denotes the complex conjugate transposed or Hermitian.
"""
inner(::Grassmann, x, v, w) = dot(v, w)

@doc doc"""
    inverse_retract(M::Grassmann, x, y, ::PolarInverseRetraction)

Compute the inverse retraction for the [`PolarRetraction`](@ref), on the
[`Grassmann`](@ref), i.e.,

````math
\operatorname{retr}_x^{-1}y = y*(x^\mathrm{H}y)^{-1} - x,
````

where $\cdot^{\mathrm{H}}$ denotes the complex conjugate transposed or Hermitian.
"""
inverse_retract(M::Grassmann, ::Any, ::Any, ::PolarInverseRetraction)

function inverse_retract!(::Grassmann, v, x, y, ::PolarInverseRetraction)
    return copyto!(v, y / (x' * y) - x)
end

@doc doc"""
    inverse_retract(M, x, y, ::QRInverseRetraction)

Compute the inverse retraction valid of the [`QRRetraction`](@ref)

````math
\operatorname{retr}_x^{-1}y = y*(x^\mathrm{H}y)^{-1} - x,
````
where $\cdot^{\mathrm{H}}$ denotes the complex conjugate transposed or Hermitian.
"""
inverse_retract(::Grassmann, ::Any, ::Any, ::QRInverseRetraction)

inverse_retract!(::Grassmann, v, x, y, ::QRInverseRetraction) = copyto!(v, y / (x' * y) - x)

function isapprox(M::Grassmann, x, v, w; kwargs...)
    return isapprox(sqrt(inner(M, x, zero_tangent_vector(M, x), v - w)), 0; kwargs...)
end
isapprox(M::Grassmann, x, y; kwargs...) = isapprox(distance(M, x, y), 0.0; kwargs...)

@doc doc"""
    log(M::Grassmann, x, y)

Compute the logarithmic map on the [`Grassmann`](@ref) `M`$ = \mathcal M=\mathrm{Gr}(n,k)$,
i.e. the tangent vector `v` whose corresponding [`geodesic`](@ref) starting from `x`
reaches `y` after time 1 on `M`. The formula reads

````math
\log_xy = V\cdot \operatorname{atan}(S) \cdot U^\mathrm{H},
````

where $\cdot^{\mathrm{H}}$ denotes the complex conjugate transposed or Hermitian.
$U$ and $V$ are the unitary matrices, and $S$ is a diagonal matrix containing
the singular values of the SVD-decomposition of
````math
USV = (y^\mathrm{H}x)^{-1} ( y^\mathrm{H} - y^\mathrm{H}xx^\mathrm{H} ).
````
In this formula the $\operatorname{atan}$ is meant elementwise.
"""
log(::Grassmann, ::Any...)

function log!(M::Grassmann, v, x, y)
    z = y' * x
    At = y' - z * x'
    Bt = z \ At
    d = svd(Bt')
    return copyto!(v, d.U * Diagonal(atan.(d.S)) * d.Vt)
end

@doc doc"""
    manifold_dimension(M::Grassmann)

Return the dimension of the [`Grassmann(n,k,𝔽)`](@ref) manifold `M`, i.e.

````math
\dim \operatorname{Gr}(n,k) = k(n-k) \dim_ℝ 𝔽,
````

where $\dim_ℝ 𝔽$ is the [`real_dimension`](@ref) of `𝔽`.
"""
manifold_dimension(M::Grassmann{n,k,𝔽}) where {n,k,𝔽} = k * (n - k) * real_dimension(𝔽)

"""
    mean(
        M::Grassmann,
        x::AbstractVector,
        [w::AbstractWeights,]
        method = GeodesicInterpolationWithinRadius(π/4);
        kwargs...,
    )

Compute the Riemannian [`mean`](@ref mean(M::Manifold, args...)) of `x` using
[`GeodesicInterpolationWithinRadius`](@ref).
"""
mean(::Grassmann{n,k,ℝ} where {n,k}, ::Any...)

function mean!(
    M::Grassmann{n,k,ℝ},
    y,
    x::AbstractVector,
    w::AbstractVector;
    kwargs...,
) where {n,k}
    return mean!(M, y, x, w, GeodesicInterpolationWithinRadius(π / 4); kwargs...)
end

@doc doc"""
    project_tangent(M::Grassmann, x, w)

Project the `n`-by-`k` `w` onto the tangent space of `x` on the [`Grassmann`](@ref) `M`,
which is computed by

````math
\operatorname{proj_x}(w) = w - xx^{\mathrm{H}}w,
````

where $\cdot^{\mathrm{H}}$ denotes the complex conjugate transposed or Hermitian.
"""
project_tangent(::Grassmann, ::Any...)

project_tangent!(M::Grassmann, v, x, w) = copyto!(v, w - x * x' * w)

@doc doc"""
    representation_size(M::Grassmann{n,k,F})

Return the represenation size or matrix dimension of a point on the [`Grassmann`](@ref)
`M`, i.e. $(n,k)$ for both the real-valued and the complex value case.
"""
@generated representation_size(::Grassmann{n,k}) where {n,k} = (n, k)

@doc doc"""
    retract(M::Grassmann, x, v, ::PolarRetraction)

Compute the SVD-based retraction [`PolarRetraction`](@ref) on the
[`Grassmann`](@ref) `M`. With $USV = x + v$ the retraction reads
````math
\operatorname{retr}_x v = UV^\mathrm{H},
````

where $\cdot^{\mathrm{H}}$ denotes the complex conjugate transposed or Hermitian.

    retract(M::Grassmann, x, v, ::QRRetraction )

Compute the QR-based retraction [`QRRetraction`](@ref) on the
[`Grassmann`](@ref) `M`. With $QR = x + v$ the retraction reads
````math
\operatorname{retr}_xv = QD,
````
where D is a $m\times n$ matrix with
````math
D = \operatorname{diag}( \operatorname{sgn}(R_{ii}+0,5)_{i=1}^n ).
````
"""
retract(::Grassmann, ::Any...)

function retract!(::Grassmann, y, x, v, ::PolarRetraction)
    s = svd(x + v)
    return mul!(y, s.U, s.Vt)
end
function retract!(::Grassmann{N,K}, y, x, v, ::QRRetraction) where {N,K}
    qrfac = qr(x + v)
    d = diag(qrfac.R)
    D = Diagonal(sign.(sign.(d .+ 0.5)))
    y .= zeros(N, K)
    y[1:K, 1:K] .= D
    return copyto!(y, Array(qrfac.Q) * D)
end

@doc doc"""
    zero_tangent_vector(M::Grassmann, x)

Return the zero tangent vector from the tangent space at `x` on the [`Grassmann`](@ref) `M`,
which is given by a zero matrix the same size as `x`.
"""
zero_tangent_vector(::Grassmann, ::Any...)

zero_tangent_vector!(::Grassmann, v, x) = fill!(v, 0)


# 3D Toy problem to test the solver

using luteos

P = P1() # Polynomial order of solution

mesh   = Mesh3D( "cube", P, N = 30 )#17)
master = Master3D( P )

mat = Material( E = 1, ν = 0.0, dim = mesh.dim )

# source function
function funcS( p::Array{Float64,2} )
  return fill(1.0, size(p,1), 3)
end

function funcB( p::Array{Float64,2} )
  return fill(0.0, size(p,1), 3)
end
function funcBNH( p::Array{Float64,2} )
  return fill(1.0, size(p,1), 3)
end

# prob = Problem( funcS, [funcB, funcB, funcB, funcB, funcB, funcB], [1,1,1,1,1,1],
#                 name="Example 1", bcnorm=false )
prob = Elas( "Example 1", mat, funcS,  [1,1,1,1,1,1], false,[funcB, funcB, funcB, funcB, funcB, funcBNH] )

@time (uhath, uh, σh ) = hdgSolve( master, mesh, prob, method = :GMRES )

# write solution
writeTecplot( "elas3D.dat", prob, mesh, uh, σh )

@printf("No. elements: %i", size(σh,3))


using Compat.Test
using Compat.LinearAlgebra
using MathProgBase

function mixintprogtest(solver)
    @testset "Testing mixintprog with $solver" begin
        # integer knapsack problem
        sol = mixintprog(-[5,3,2,7,4],Float64[2 8 4 2 5],'<',10,:Int,0,1,solver)
        @test sol.status == :Optimal
        @test isapprox(sol.objval, -16.0, atol=1e-6)
        @test isapprox(norm(sol.sol - [1.0, 0.0, 0.0, 1.0, 1.0]), 0.0, atol=1e-4)
    end
end


function get_num_dens_spectrum(CR_N::Vector{T}, CR_S::Vector{T},
    CR_Cut::T, ρ::T;
    par::CRMomentumDistributionConfig,
    mode::Integer = 3) where {T}

    # transform the norm dependent on IO mode
    norm = transform_norm(CR_N, mode)

    bin_width = log10(par.pmax / par.pmin) / par.Nbins

    bounds = 10.0 .^ collect(log10(par.pmin):par.bin_width:log10(par.pmax))
    bound_up = Vector{Float64}(undef, par.Nbins)

    for i = 1:par.Nbins
        bound_up[i] = (bounds[i+1] > CR_Cut) ? CR_Cut : bounds[i+1]
    end

    bound_up[end] = (bounds[end] > CR_Cut) ? CR_Cut : bounds[end]

    bin_centers = Vector{Float64}(undef, par.Nbins)
    density = Vector{Float64}(undef, par.Nbins)

    @inbounds for i = 1:par.Nbins
        bin_centers[i] = 10.0^(0.5 * (log10(bounds[i]) + log10(bound_up[i])))
        density[i] = density_integral(bounds[i], bound_up[i], norm[i], CR_S[i], ρ)
    end

    return CR_NumSpectrum(bin_centers, density)
end


using MatterPower
using PyCall
using Plots, LaTeXStrings

# %% Specify a redshift
redshift = 0

# %% Call the python wrapper for CLASS, `classy`, via PyCall
classy = pyimport("classy")

# Create an instance of the CLASS wrapper
cosmo = classy.Class()

# Create a dictionary of the cosmological parameters
params = Dict(
    "output" => "mPk",
    "P_k_max_h/Mpc" => 30,
    "z_pk" => 5,
    "non linear" => "halofit",
    "A_s" => 2.097e-9,
    "n_s" => 0.9652,
    "k_pivot" => 0.05,
    "h" => 0.674,
    "Omega_b" => 0.049,
    "Omega_cdm" => 0.266,
)

# Set the parameters and run the CLASS code
cosmo.set(params)
cosmo.compute()

# Define functions to return linear and non-linear power spectra
# Note: The CLASS code takes wavenumbers in units of 1/Mpc (no h) and
# return power spectra in units of Mpc^3 (no 1/h^3).
pklin_class(kovh) = cosmo.pk_lin(kovh * params["h"], redshift) * params["h"]^3
pknl_class(kovh) = cosmo.pk(kovh * params["h"], redshift) * params["h"]^3

# %% Plot results and save to "classpk.pdf"
lnk = log(3e-4):0.3:log(30)
kovh = exp.(lnk)
p = plot(
    kovh,
    pknl_class.(kovh),
    xaxis = :log,
    yaxis = :log,
    m = 2,
    lab = "Non-linear P(k)",
    ylab = L"P(k)~[h^{-3}~Mpc^3]",
    xlab = L"k~[h~Mpc^{-1}]",
)
p = plot!(kovh, pklin_class.(kovh), ls = :dash, lab = "Linear P(k)")
savefig("classpk.pdf")
display(p)


"""
    baselevel(d, lags)

Computes baselevel activation according exact equation.

# Arguments

- `d`: decay parameter
- `lags`: a vector of time lags between current time and each usage time
"""
function baselevel(d, lags)
    act = 0.0
    for t in lags
        act += t^-d
    end
    return log(act)
end

"""
    baselevel(d, chunk)

Computes baselevel activation according exact equation.

# Arguments

- `actr`: an `ACTR` model object
- `chunk`: a chunk
"""
function baselevel!(actr, chunk)
    @unpack N,L,k,lags = chunk
    d = actr.parms.d
    exact = baselevel(d, lags)
    approx = 0.0
    if N > k
        tk = lags[k]
        x1 = (N - k) * (L^(1 - d) - tk^(1 - d))
        x2 = (1 - d) * (L - tk)
        approx = x1 / x2
    end
    chunk.act_bll = log(exp(exact) + approx)
    return nothing
end

"""
    baselevel!(actr)

Computes baselevel activation for all chunks according exact equation.

# Arguments

- `actr`: an `ACTR` model object
"""
baselevel!(actr) = activation!.(actr, memory.memory)

"""
    compute_activation!(actr::AbstractACTR, chunks::Vector{<:Chunk}; request...)

Computes the activation of a vector of chunks. By default, current time is computed with `get_time`.

# Arguments

- `actr::AbstractACTR`: an `ACTR` object
- `chunks::Vector{<:Chunk}`: a vector of chunks.

# Keywords

- `request...`: optional keywords for the retrieval request
"""
function compute_activation!(actr::AbstractACTR, chunks::Vector{<:Chunk}; request...)
    return compute_activation!(actr, chunks, get_time(actr); request...)
end

"""
    compute_activation!(actr::AbstractACTR, chunks::Vector{<:Chunk}, cur_time::Float64=0.0; request...)

Computes the activation of a vector of chunks

# Arguments

- `actr::AbstractACTR`: an `ACTR` object
- `chunks::Vector{<:Chunk}`: a vector of chunks.
- `cur_time::Float64`: current simulated time in seconds

# Keywords

- `request...`: optional keywords for the retrieval request
"""
function compute_activation!(actr::AbstractACTR, chunks::Vector{<:Chunk}, cur_time::Float64; request...)
    @unpack sa = actr.parms
    sa ? cache_denomitors(actr) : nothing
    # compute activation for each chunk
    for chunk in chunks
        activation!(actr, chunk, cur_time; request...)
    end
    return nothing
end

"""
    compute_activation!(actr, chunk::Chunk; request...)

Computes the activation of a chunk. By default, current time is computed
with `get_time`.

# Arguments

- `actr`: actr object
- `chunk::Chunk`: a chunk.

# Keywords

- `request...`: optional keywords for the retrieval request
"""
function compute_activation!(actr, chunk::Chunk; request...)
    return compute_activation!(actr, chunk, get_time(actr); request...)
end

"""
    compute_activation!(actr, chunk::Chunk, cur_time; request...)

Computes the activation of a chunk

# Arguments

- `actr`: actr object
- `chunk::Chunk`: a chunk.
- `cur_time`: current simulated time in seconds

# Keywords

- `request...`: optional keywords for the retrieval request
"""
compute_activation!(actr, chunk::Chunk, cur_time; request...) = compute_activation!(actr, [chunk], cur_time; request...)

"""
    compute_activation!(actr::AbstractACTR; request...)

Computes the activation of all chunks in declarative memory. By default, current time is computed
with `get_time`.

# Arguments

- `actr::AbstractACTR`: an `ACTR` object

# Keywords

- `request...`: optional keywords for the retrieval request
"""
compute_activation!(actr::AbstractACTR; request...) = compute_activation!(actr, actr.declarative.memory, get_time(actr); request...)

"""
    compute_activation!(actr::AbstractACTR, cur_time::Float64; request...)

Computes the activation of all chunks in declarative memory

# Arguments

- `actr::AbstractACTR`: an `ACTR` object
- `cur_time`: current simulated time in seconds

# Keywords

- `request...`: optional keywords for the retrieval request
"""
compute_activation!(actr::AbstractACTR, cur_time::Float64; request...) = compute_activation!(actr, actr.declarative.memory, cur_time; request...)

"""
    activation!(actr, chunk::Chunk, cur_time; request...)

Computes the activation of a chunk

# Arguments

- `actr`: an `ACTR` object
- `chunk::Chunk`: a chunk.
- `cur_time`: current simulated time in seconds

# Keywords

- `request...`: optional keywords for the retrieval request
"""
function activation!(actr, chunk::Chunk, cur_time=0.0; request...)
    memory = actr.declarative
    @unpack sa_fun,bll,mmp,sa,noise,blc,τ = actr.parms
    reset_activation!(chunk)
    chunk.act_blc = blc + chunk.bl
    if bll
        update_lags!(chunk, cur_time)
        baselevel!(actr, chunk)
    end
    if mmp
        partial_matching!(actr, chunk; request...)
    end
    if sa
        sa_fun(actr, chunk)
    end
    if noise
        add_noise!(actr, chunk)
    end
    total_activation!(chunk)
    return nothing
end

function reset_activation!(chunk)
    a = chunk.act
    chunk.act_mean = zero(a)
    chunk.act_blc = zero(a)
    chunk.act_bll = zero(a)
    chunk.act_pm = zero(a)
    chunk.act_sa = zero(a)
    chunk.act_noise = zero(a)
    chunk.act = zero(a)
end

"""
    total_activation!(chunk)

Assigns sum of all components to of activation to `chunk.act`.

# Arguments

- `chunk`: a chunk object
"""
function total_activation!(chunk)
    chunk.act_mean = chunk.act_blc + chunk.act_bll - chunk.act_pm +
        chunk.act_sa
    chunk.act = chunk.act_mean + chunk.act_noise
    return nothing
end

function add_noise!(actr, chunk)
    @unpack τ,s = actr.parms
    σ = s * pi / sqrt(3)
    chunk.act_noise = rand(Normal(0, σ))
    return nothing
end

function add_noise!(actr)
    @unpack τ,s = actr.parms
    σ = s * pi / sqrt(3)
    actr.parms.τ′ = rand(Normal(τ, σ))
    nothing
end

"""
    partial_matching!(actr, chunk; request...)

Computes activation for partial matching component

# Arguments

- `actr`: an ACTR object
- `chunk`: a chunk

# Keywords

- `request...`: optional keyword arguments for retrieval request
"""
function partial_matching!(actr, chunk; request...)
    p = actr.parms.mmp_fun(actr, chunk; request...)
    chunk.act_pm = p
    return nothing
end

"""
    set_noise!(actr::AbstractACTR, b)

Sets noise true or false.

# Arguments

- `actr`: ACTR object
- `b`: boolean value
"""
function set_noise!(actr, b)
    actr.parms.noise = b
end

"""
    spreading_activation!(actr, chunk)

Computes the spreading activation for a given chunk

# Arguments

- `actr`: an `ACTR` oject
- `chunk`: the chunk for which spreading activation is computed
"""
function spreading_activation!(actr, chunk)
    imaginal = actr.imaginal
    isempty(imaginal.buffer) ? (return nothing) : nothing
    w = compute_weights(imaginal)
    γ = actr.parms.γ; r = zero(γ); sa = zero(γ)
    slots = imaginal.buffer[1].slots
    denoms = imaginal.denoms
    for (v,d) in zip(slots, denoms)
        num = count_values(chunk, v)
        fan = num / (d + 1)
        r = fan == 0 ? zero(γ) : γ + log(fan)
        sa += w * r
    end
    chunk.act_sa = sa
    return nothing
end

"""
    cache_denomitors(actr)

Caches denominator of spreading activation calculations

# Arguments

- `actr`: an `ACTR` oject
"""
function cache_denomitors(actr)
    @unpack imaginal,declarative = actr
    isempty(imaginal.buffer) ? (return nothing) : nothing
    slots = imaginal.buffer[1].slots
    denoms = fill(0, length(slots))
    for (i,v) in enumerate(slots)
        denoms[i] = compute_denom(declarative, v)
    end
    imaginal.denoms = denoms
    return nothing
end

function compute_weights(mod)
    return mod.ω / length(mod.buffer[1].slots)
end

function compute_denom(memory, value)
    denom = 0
    for c in memory.memory
        denom += count_values(c, value)
    end
    return denom
end

function count_values(chunk, value)
    return count(x -> x == value, values(chunk.slots))
end

"""
    update_recent!(actr, chunk)

Adds a new timestamp to chunk and removes oldest timestamp if
length equals k. By default, current time is computed with `get_time`.

# Arguments
- `actr`: an ACT-R model object
- `chunk`: memory chunk object
"""
function update_recent!(actr::AbstractACTR, chunk)
    return update_recent!(chunk, get_time(actr))
end

"""
    update_recent!(chunk, cur_time)

Adds a new timestamp to chunk and removes oldest timestamp if
length equals k.

# Arguments

* `chunk`: memory chunk object
* `cur_time`: current simulated time in seconds
"""
function update_recent!(chunk, cur_time)
    k = chunk.k; recent = chunk.recent
    if length(recent) == k
        pop!(recent)
    end
    pushfirst!(recent, cur_time)
    return nothing
end

"""
    retrieval_prob(actr::AbstractACTR, target::Array{<:Chunk,1}; request...)

Uses the softmax approximation to compute the retrieval probability of retrieving a chunk.
By default, current time is computed from `get_time`.

# Arguments

- `actr::AbstractACTR`: an ACT-R object
- `target::Array{<:Chunk,1}`: a vector chunks in the numerator of the softmax function

# Keywords

- `request...`: optional keyword pairs representing a retrieval request
"""
function retrieval_prob(actr::AbstractACTR, target::Array{<:Chunk,1}; request...)
    return retrieval_prob(actr, target, get_time(actr); request...)
end

"""
    retrieval_prob(actr::AbstractACTR, target::Array{<:Chunk,1}, cur_time; request...)

Computes the retrieval probability of one chunk from a set of chunks defined in `target`.
Retrieval probability is computed with the softmax approximation.

# Arguments

- `actr::AbstractACTR`: an actr object
- `target::Array{<:Chunk,1}`: a vector chunks in the numerator of the softmax function
- `cur_time`: current time in seconds

# Keywords

- `request...`: optional keywords for the retrieval request
"""
function retrieval_prob(actr::AbstractACTR, target::Array{<:Chunk,1}, cur_time; request...)
    @unpack τ,s,noise = actr.parms
    σ = s * sqrt(2)
    chunks = retrieval_request(actr; request...)
    filter!(x -> (x ∈ chunks), target)
    isempty(target) ? (return (0.0,1.0)) : nothing
    set_noise!(actr, false)
    compute_activation!(actr, chunks, cur_time; request...)
    set_noise!(actr, noise)
    denom = fill(target[1].act_mean, length(chunks) + 1)
    map!(x -> exp(x.act_mean / σ), denom, chunks)
    denom[end] = exp(τ / σ)
    num = map(x -> exp(x.act_mean / σ), target)
    prob = sum(num) / sum(denom)
    fail = denom[end] / sum(denom)
    return prob,fail
end

"""
    retrieval_prob(actr::AbstractACTR, chunk::Chunk; request...)

Uses the softmax approximation to compute the retrieval probability of retrieving a chunk.
By default, current time is computed from `get_time`.

# Arguments

- `actr::AbstractACTR`: an ACT-R object
- `chunk::Chunk`: a chunk

# Keywords

- `request...`: optional keyword pairs representing a retrieval request
"""
function retrieval_prob(actr::AbstractACTR, chunk::Chunk; request...)
    return retrieval_prob(actr, chunk, get_time(actr); request...)
end

"""
    retrieval_prob(actr::AbstractACTR, chunk::Chunk, cur_time=0.0; request...)

Uses the softmax approximation to compute the retrieval probability of retrieving a chunk.

# Arguments

- `actr::AbstractACTR`: an ACT-R object
- `chunk::Chunk`: a chunk
- `cur_time`: current simulated time in seconds

# Keywords

- `request...`: optional keyword pairs representing a retrieval request
"""
function retrieval_prob(actr::AbstractACTR, chunk::Chunk, cur_time; request...)
    @unpack τ,s,noise = actr.parms
    σ = s * sqrt(2)
    chunks = retrieval_request(actr; request...)
    !(chunk ∈ chunks) ? (return (0.0,1.0)) : nothing
    set_noise!(actr, false)
    compute_activation!(actr, chunks, cur_time; request...)
    set_noise!(actr, noise)
    v = fill(chunk.act_mean, length(chunks) + 1)
    map!(x -> exp(x.act_mean / σ), v, chunks)
    v[end] = exp(τ / σ)
    prob = exp(chunk.act_mean / σ) / sum(v)
    fail = v[end] / sum(v)
    return prob,fail
end

"""
    retrieval_probs(actr::AbstractACTR; request...)

Computes the retrieval probability for each chunk matching the retrieval request. By default,
current time is computed from `get_time`.

# Arguments

- `actr::AbstractACTR`: an actr object

# Keywords

- `request...`: optional keyword pairs representing a retrieval request
"""

function retrieval_probs(actr::AbstractACTR; request...)
    return retrieval_probs(actr, get_time(actr); request...)
end

"""
    retrieval_probs(actr::AbstractACTR, cur_time; request...)

Computes the retrieval probability for each chunk matching the retrieval request.

# Arguments

- `actr::AbstractACTR`: an actr object
- `cur_time`: current simulated time in seconds

# Keywords

- `request...`: optional keyword pairs representing a retrieval request
"""
function retrieval_probs(actr::AbstractACTR, cur_time; request...)
    @unpack τ,s,noise = actr.parms
    σ = s * sqrt(2)
    chunks = retrieval_request(actr; request...)
    isempty(chunks) ? (return ([0.0],chunks)) : nothing
    set_noise!(actr, false)
    compute_activation!(actr, chunks, cur_time; request...)
    set_noise!(actr, noise)
    v = Array{typeof(chunks[1].act),1}(undef, length(chunks) + 1)
    map!(x -> exp(x.act_mean / σ), v, chunks)
    v[end] = exp(τ / σ)
    p = v ./ sum(v)
    return p,chunks
end

"""
    update_lags!(chunk::Chunk, cur_time)

Compute lags for each use of a chunk.

# Arguments

- `chunk::Chunks`: a chunk
- `cur_time`: current simulated time in seconds.

"""
function update_lags!(chunk::Chunk, cur_time)
    chunk.L = cur_time - chunk.time_created
    chunk.lags = cur_time .- chunk.recent
    return nothing
end

"""
    update_lags!(actr::AbstractACTR, cur_time)

Compute lags for each use of a chunk. Applies to all chunks in declarative memory.

# Arguments

- `actr::AbstractACTR`: an ACTR object
- `cur_time`: current simulated time in seconds
"""
update_lags!(actr::AbstractACTR, cur_time) = update_lags!(actr.declarative, cur_time)

"""
    update_lags!(memory::Declarative, cur_time)

Compute lags for each use of a chunk. Applies to all chunks in declarative memory.

# Arguments

- `memory::Declarative`: a declarative memory object object
- `cur_time`: current simulated time in seconds
"""
update_lags!(memory::Declarative, cur_time) = update_lags!.(memory.memory, cur_time)

"""
    update_chunk!(chunk, cur_time)

Increments number of uses and adds `cur_time` as the most recent time of use.

# Arguments

- `chunk`: a chunk object
- `cur_time': current simulated time in seconds
"""
function update_chunk!(chunk, cur_time)
    update_recent!(chunk, cur_time)
    chunk.N += 1
    return nothing
end

"""
    add_chunk!(actr::AbstractACTR; slots...)

Adds new chunk to declarative memory or updates existing chunk with new use. The default time is
computed from `get_time`.

# Arguments

- `actr::AbstractACTR`: an ACT-R model object

# Keywords

- `slots...`: optional keyword arguments corresponding to slot-value pairs, e.g. name=:Bob
"""
add_chunk!(actr::AbstractACTR; slots...) = add_chunk!(actr.declarative, get_time(actr); slots...)

"""
    add_chunk!(memory::Declarative, cur_time=0.0; act=0.0, slots...)

Adds new chunk to declarative memory or updates existing chunk with new use

# Arguments

- `memory::Declarative`: declarative memory object
- `cur_time=0.0`: current simulated time in seconds
- `bl=0.0`: baselevel constant for new/updated chunk

# Keywords

- `slots...`: optional keyword arguments corresponding to slot-value pairs, e.g. name=:Bob
"""
function add_chunk!(memory::Declarative, cur_time=0.0; bl::T=0.0, slots...) where {T}
    chunk = get_chunks_exact(memory; slots...)
    if isempty(chunk)
        c = Chunk(;act=zero(T), bl, time_created=cur_time, recent=[cur_time], slots...)
        push!(memory.memory, c)
    else
        update_chunk!(chunk[1], cur_time)
    end
    return nothing
end

"""
    add_chunk!(actr::AbstractACTR, cur_time=0.0; slots...)

Adds new chunk to declarative memory or updates existing chunk with new use

# Arguments

- `actr::AbstractACTR`: an ACT-R model object
- `cur_time=0.0`: current simulated time in seconds

# Keywords

- `slots...`: optional keyword arguments corresponding to slot-value pairs, e.g. name=:Bob
"""
add_chunk!(actr::AbstractACTR, cur_time; slots...) = add_chunk!(actr.declarative, cur_time; slots...)

"""
    get_chunks_exact(memory::Vector{<:Chunk}; criteria...)

Returns all chunks that matches a set criteria and has the same number of slots.

# Arguments

- `memory::Vector{<:Chunk}`: vector of chunk objects

# Keywords

-`criteria...`: optional keyword arguments corresponding to critiria for matching chunk
"""
function get_chunks_exact(memory::Vector{<:Chunk}; criteria...)
    c = filter(x -> match_exact(x, criteria), memory)
    return c
end

get_chunks_exact(d::Declarative; criteria...) = get_chunks_exact(d.memory; criteria...)

function match_exact(chunk::Chunk, request)
    slots = chunk.slots
    length(slots) ≠ length(request) ? (return false) : nothing
    for (k,v) in request
        if !(k ∈ keys(slots)) || (slots[k] != v)
            return false
        end
    end
    return true
end

"""
    get_chunks(memory::Vector{<:Chunk}; criteria...)

Returns all chunks that matches a set criteria.

# Arguments

- `memory::Vector{<:Chunk}`: vector of chunk objects

# Keywords

 - `check_value=true`: check slot value

-`criteria...`: optional keyword arguments corresponding to critiria for matching chunk
"""
function get_chunks(memory::Vector{<:Chunk}; check_value=true, criteria...)
    c = filter(x -> _match(x, criteria; check_value), memory)
    return c
end

"""
    get_chunks(memory::Vector{<:Chunk}, funs...; criteria...)

Returns all chunks that matches a set `criteria` which are evaluted according to the list of functions in `funs`.

# Arguments

- `memory::Vector{<:Chunk}`: vector of chunk objects
- `funs...`: a list of functions

# Keywords

 - `check_value=true`: check slot value

- `criteria...`: optional keyword arguments corresponding to critiria for matching chunk
"""
function get_chunks(memory::Vector{<:Chunk}, funs...; check_value=true, criteria...)
    c = filter(x -> _match(x, funs, criteria; check_value), memory)
    return c
end

"""
    get_chunks(d::Declarative; criteria...)

Returns all chunks that matches a set criteria.

# Arguments

- `d::Declarative`: declarative memory object

# Keywords

 - `check_value=true`: check slot value

- `criteria...`: optional keyword arguments corresponding to critiria for matching chunk
"""
get_chunks(d::Declarative; check_value=true, criteria...) = get_chunks(d.memory; check_value, criteria...)

"""
    get_chunks(actr::AbstractACTR; )

Returns all chunks that matches a set criteria

# Arguments

- `actr::AbstractACTR`: an ACTR Object

#Keywords

 - `check_value=true`: check slot value

- `criteria...`: optional keyword arguments corresponding to critiria for matching chunk
"""
get_chunks(actr::AbstractACTR; check_value=true, criteria...) = get_chunks(actr.declarative.memory; check_value, criteria...)

"""
    get_chunks(d::Declarative, funs...; criteria...)

Returns all chunks that matches a set criteria using `funs...` as matching functions

# Arguments

- `d::Declarative`: declarative memory object
- `funs...`: a list of functions

# Keywords

 - `check_value=true`: check slot value

- `criteria...`: optional keyword arguments corresponding to critiria for matching chunk
"""
get_chunks(d::Declarative, funs...; check_value=true, criteria...) = get_chunks(d.memory, funs...; check_value, criteria...)

"""
    get_chunks(actr::AbstractACTR, funs...; criteria...)

Returns all chunks that matches a set criteria using `funs...` as matching functions.

# Arguments

- `actr::AbstractACTR`: an ACT-R model object
- `funs...`: a list of functions

# Keywords

 - `check_value=true`: check slot value

- `criteria...`: optional keyword arguments corresponding to critiria for matching chunk
"""
get_chunks(actr::AbstractACTR, funs...; check_value=true, criteria...) = get_chunks(actr.declarative.memory, funs...; check_value, criteria...)

"""
    first_chunk(memory::Vector{<:Chunk}; criteria...)

Returns the first chunk in memory that matches a set of criteria

# Arguments

- `memory::Vector{<:Chunk}`: a vector of chunks

# Keywords

 - `check_value=true`: check slot value

- `criteria...`: optional keyword arguments corresponding to critiria for matching chunk
"""
function first_chunk(memory::Vector{<:Chunk}; check_value=true, criteria...)
    chunk = Array{eltype(memory),1}()
    for m in memory
        if _match(m, criteria; check_value)
            push!(chunk, m)
            return chunk
        end
    end
    return chunk
end

"""
    first_chunk(d::Declarative; criteria...)

Returns the first chunk in memory that matches a set of criteria

# Arguments

- `d::Declarative`: a declarative memory object

# Keywords

 - `check_value=true`: check slot value

- `criteria...`: optional keyword arguments corresponding to critiria for matching chunk
"""
first_chunk(d::Declarative; check_value=true, criteria...) = first_chunk(d.memory; check_value, criteria...)

"""
    first_chunk(a::AbstractACTR; criteria...)

Returns the first chunk in memory that matches a set of criteria

# Arguments

- `actr::AbstractACTR`: an ACT-R model object

# Keywords

 - `check_value=true`: check slot value

- `criteria...`: optional keyword arguments corresponding to critiria for matching chunk
"""
first_chunk(actr::AbstractACTR; check_value=true, criteria...) = first_chunk(actr.declarative.memory; check_value, criteria...)

"""
    match(chunk::Chunk, request)

Returns a boolean indicating whether a request matches a chunk.
False is returned if the slot does not exist in the chunk or the value
of the slot does not match the request value.

# Arguments

- `chunk::Chunk`: chunk object
- `request`: a NamedTuple of slot value pairs

# Keywords

 - `check_value=true`: check slot value
"""
function _match(chunk::Chunk, request; check_value=true)
    slots = chunk.slots
    for (k,v) in request
        if !(k ∈ keys(slots))
            return false
        elseif check_value
            if (slots[k] != v)
                return false
            end
        end
    end
    return true
end

"""
    match(chunk::Chunk, f, request)

Returns a boolean indicating whether a request matches a chunk.
False is returned if the slot does not exist in the chunk or the value
of the slot does not match the request value.

# Arguments

- `chunk`: a chunk object
- `f`: a list of functions such as `!=, ==`
- `request`: a NamedTuple of slot value pairs

# Keywords

 - `check_value=true`: check slot value
"""
function _match(chunk::Chunk, f, request; check_value=true)
    slots = chunk.slots
    i = 0
    for (k,v) in request
        i += 1
        if !(k ∈ keys(slots))
            return false
        elseif check_value
            if !(f[i](slots[k], v))
                return false
            end
        end
    end
    return true
end

"""
    match(chunk::Chunk; request...)

Returns a boolean indicating whether a request matches a chunk.
False is returned if the slot does not exist in the chunk or the value
of the slot does not match the request value.

# Arguments

- `chunk::Chunk`: a chunk object

# Keywords

- `request...`: optional keyword arguments corresponding to critiria for matching chunk
"""
function match(chunk::Chunk; check_value=true, request...)
    return _match(chunk, request; check_value)
end
"""
    match(chunk::Chunk, funs...; request...)

Returns a boolean indicating whether a request matches a chunk.
False is returned if the slot does not exist in the chunk or the value
of the slot does not match the request value.

- `chunk`: chunk object
- `funs...`: a list of functions such as `!=, ==`
- `request...`: a NamedTuple of slot value pairs
"""
function match(chunk::Chunk, funs...; check_value=true, request...)
    return _match(chunk, funs, request; check_value)
end

"""
    get_subset(actr; request...)

Returns a filtered subset of the retrieval request when partial matching is on.
By default, slot values for isa and retrieved must match exactly.

# Arguments

- `actr`: an ACTR object

# Keywords

- `request...`: an option set of keyword arguments corresponding to a retrieval request.
"""
function get_subset(actr; request...)
    return Iterators.filter(x -> any(s -> s == x[1], actr.declarative.filtered),
    request)
end

"""
    retrieval_request(actr::AbstractACTR; request...)

Returns chunks matching a retrieval request.

# Arguments

- `memory`: declarative memory object

# Keywords

- `request...`: optional keyword arguments corresponding to retrieval request e.g. dog = :fiddo
"""
function retrieval_request(actr::AbstractACTR; request...)
    @unpack mmp = actr.parms
    !mmp ? (return get_chunks(actr; request...)) : nothing
    chunks = get_chunks(actr; check_value=false, request...)
    c = get_subset(actr; request...)
    return get_chunks(chunks; check_value=true, c...)
end

"""
    modify!(c; args...)

Modifies fields of an object

# Arguments

- `c`: an object

# Keywords

- `args...`: optional keywords for field and values
"""
function modify!(c; args...)
    for (k,v) in args
        setfield!(c, k, v)
    end
    return nothing
end

"""
    modify!(c::NamedTuple; args...)

Modifies fields of NamedTuple

# Arguments

- `c`: a NamedTuple

# Keywords

- `args...`: optional keywords for field and values
"""
function modify!(c::NamedTuple; args...)
    for (k,v) in args
        c[k][1] = v
    end
    return nothing
end

"""
    retrieve(actr::AbstractACTR; request...)

Retrieves a chunk given a retrieval request. By default, current time is
computed with `get_time`.

# Arguments

- `actr`: an ACT-R object

# Keywords

- `request...`: optional keyword arguments representing a retrieval request, e.g. person=:bob
"""
function retrieve(actr::AbstractACTR; request...)
    return retrieve(actr, get_time(actr); request...)
end

"""
    retrieve(actr::AbstractACTR, cur_time; request...)

Retrieves a chunk given a retrieval request

# Arguments

- `actr`: an ACT-R object
- `cur_time`: current simulated time in seconds

# Keywords

- `request...`: optional keyword arguments representing a retrieval request, e.g. person=:bob
"""
function retrieve(actr::AbstractACTR, cur_time; request...)
    @unpack declarative,parms = actr
    arr = Array{eltype(declarative.memory),1}()
    chunks = retrieval_request(actr; request...)
    # add noise to threshold even if result of request is empty
    actr.parms.noise ? add_noise!(actr) : (parms.τ′ = parms.τ)
    isempty(chunks) ? (return arr) : nothing
    compute_activation!(actr, chunks, cur_time; request...)
    best = get_max_active(chunks)
    if best[1].act >= parms.τ′
        return best
    end
    return arr
end

"""
    get_max_active(chunks)

Returns the chunk with maximum activation

# Arguments

- `chunks`: a vector of chunk objects
"""
function get_max_active(chunks)
    a = -Inf
    mx = chunks[1]
    for c in chunks
        if c.act > a
            a = c.act
            mx = c
        end
    end
    return [mx]
end

"""
    compute_RT(actr, chunk)

Generates a reaction time for retrieving a chunk based
on the current activation levels of a chunk. If the vector is empty, time
for a retrieval failure is returned.

# Arguments

- `actr`: ACTR object
- `chunk`: a vector that is empty or contains one chunk
"""
function compute_RT(actr, chunk)
    @unpack τ′,lf = actr.parms
    if isempty(chunk)
        return lf * exp(-τ′)
    end
    return lf * exp(-chunk[1].act)
end

"""
    compute_RT(actr, chunk)

Generates a reaction time for retrieving a chunk based
on the current activation levels of a chunk.

# Arguments

- `actr`: ACTR object
- `chunk`: a chunk
"""
function compute_RT(actr, chunk::Chunk)
    @unpack lf = actr.parms
    return lf * exp(-chunk.act)
end

"""
    get_parm(actr, p)

Returns the value of a parameter

# Arguments

- `actr`: ACTR object
- ` p`: parameter name
"""
function get_parm(actr, p)
    misc = actr.parms.misc
    if p in keys(misc)
        return misc[p]
    end
    return getfield(actr.parms, p)
end

"""
    blend_chunks(actr; request...)

Computes blended value over chunks given a retrieval request. By default,
values are blended over the set of slots formed by the set difference between all
slots of a chunk and the slots specified in the retrieval request. The default time used
in activation calculations is taken from `get_time`. Currently, blended
is only supported for numeric slot-values.

# Arguments

- `actr::AbstractACTR`: an `ACTR` model object

# Keywords

- `request...`: optional keywords for the retrieval request
"""
function blend_chunks(actr; request...)
    return blend_chunks(actr, get_time(actr); request...)
end

"""
    blend_chunks(actr::AbstractACTR, cur_time::Float64; request...)

Computes blended value over chunks given a retrieval request. By default,
values are blended over the set of slots formed by the set difference between all
slots of a chunk and the slots specified in the retrieval request. Currently, blended
is only supported for numeric slot-values.

# Arguments

- `actr::AbstractACTR`: an `ACTR` model object
- `cur_time::Float64`: current simulated time

# Keywords

- `request...`: optional keywords for the retrieval request
"""
function blend_chunks(actr::AbstractACTR, cur_time::Float64; request...)
    blended_slots = setdiff(keys(actr.declarative.memory[1].slots), keys(request))
    return blend_chunks(actr, blended_slots, cur_time; request...)
end

"""
    blend_chunks(actr::AbstractACTR; request...)

Computes blended value over chunks given a retrieval request. By default,
values are blended over the set of slots formed by the set difference between all
slots of a chunk and the slots specified in the retrieval request. The default time used
in activation calculations is taken from `get_time`. Currently, blended
is only supported for numeric slot-values.

# Arguments

- `actr::AbstractACTR`: an `ACTR` model object

# Keywords

- `request...`: optional keywords for the retrieval request
"""
function blend_chunks(actr::AbstractACTR; request...)
    return blend_chunks(actr, get_time(actr); request...)
end

"""
    blend_chunks(actr::AbstractACTR, blended_slots; request...)

Computes blended value over chunks given a retrieval request. By default,
values are blended over the set of slots formed by the set difference between all
slots of a chunk and the slots specified in the retrieval request. The default time used
in activation calculations is taken from `get_time(actr`). Currently, blended
is only supported for numeric slot-values.

# Arguments

- `actr::AbstractACTR`: an `ACTR` model object
- `blended_slots`: a set of slots over which slot-values are blended

# Keywords

- `request...`: optional keywords for the retrieval request
"""
function blend_chunks(actr::AbstractACTR, blended_slots; request...)
    return blend_chunks(actr, blended_slots, get_time(actr); request...)
end

"""
    blend_chunks(actr, blended_slots, cur_time; request...)

Computes blended value over chunks given a retrieval request. Values are blended
over the slots specified in `blended_slots`. Currently, blended is only supported
for numeric slot-values.

# Arguments

- `actr`: an `ACTR` model object
- `blended_slots`: a set of slots over which slot-values are blended
- `cur_time`: current simulated time

# Keywords

- `request...`: optional keywords for the retrieval request
"""
function blend_chunks(actr, blended_slots, cur_time; request...)
    chunks = retrieval_request(actr; request...)
    compute_activation!(actr, chunks, cur_time; request...)
    probs = soft_max(actr, chunks)
    v = blend_slots(chunks, probs, blended_slots)
    return v
end

function blend_slots(chunks, probs, slots)
    T = typeof(probs)
    v::T = map(s -> blend_slots(chunks, probs, s), slots)
    return v
end

"""
    blend_slots(chunks, probs, slot::Symbol)

Computes an expected value over chunks for a specified slot.

# Arguments

- `chunks`: a set of chunks over which slot-values are blended
- `probs`: a vector of retrieval probabilities
- `slot::Symbol`: a slot over which slot-values are blended
"""
function blend_slots(chunks, probs, slot::Symbol)
    T = eltype(probs)
    v::T = 0.0
    for (c,p) in zip(chunks, probs)
        v += p * c.slots[slot]
    end
    return v
end

# non numeric
# blend_chunk(actr, chunks, probs, chunk) = blend_chunk(actr, chunks, probs; chunk.slots...)

# function blend_chunk(actr, chunks, probs; slots...)
#     penalties = actr.parms.mmp_fun.(actr, chunks; slots...)
#     return probs' * penalties
# end

function soft_max(actr, chunks)
    σ = actr.parms.s * sqrt(2)
    v = map(x -> exp(x.act / σ), chunks)
    return v ./ sum(v)
end

"""
    blended_activation(chunks)

Computes a blended activation value by exponentiating, summing and taking the
log of activations across a set of chunks.

# Arguments

- `chunks`: a set of chunks over which slot-values are blended
"""
function blended_activation(chunks)
    exp_act = map(x -> exp(x.act_mean), chunks)
    return log(sum(exp_act))
end

"""
    compute_RT(blended_act)

Computes retrieval time for a given blended activation value.

# Arguments

- `blended_act`: a blended activation value
"""
compute_RT(blended_act) = exp(-blended_act)


my_f(x,y) = 2x+3y
derivative_of_my_f(x,y) = ForwardDiff.derivative(x -> my_f(x,y), x)


## Make MNIST Graph
using MLDatasets
train_x, train_y = MNIST.traindata()
test_x,  test_y  = MNIST.testdata()
##
m = size(train_x,1)*size(train_x,2)
n = size(train_x,3)
X = Float64.(reshape(train_x, m, n))
## Build the 5-NN graph
using NearestNeighbors
@time T = KDTree(X)
##
k = 4
idxs = knn(T, X, k+1)[1]
##
ei = Int[]
ej = Int[]
for i=1:n
  for j=idxs[i]
    if i != j
      push!(ei,i)
      push!(ej,j)
    end
  end
end
##
using SparseArrays
A = sparse(ei,ej,1,n,n)
A = max.(A,A')
##
using DelimitedFiles
writedlm("mnist-train-$k.edges", hcat(findnz(A)[1:2]...))


function utility(consumption::Float64,new_debt::Bool,reserve::Bool,tozero_l::Bool)

    #Put here Your Favorite Utility Function
    check1 = ((new_debt) | tozero_l)
    check2 = reserve
    check3 = consumption > 0
    ( check1 & check2 & check3) ? ( return( (consumption ^ (1 - σ)) / (1 - σ) ) ) : ( return(-Inf) )

end


"""
    `individual_summary(df,xvar,yvar; summary = mean, err = :MOUSE)`
    computes a summary statistic (default = mean) for each individual
    (default = :MOUSE) used to later calculate the error across individual
"""

function individual_summary(df,xvar,yvar; summary = mean, err = :MOUSE)
    #evaluate if it's going over multiple grouping columns
    multiple = typeof(xvar) <: AbstractVector && sizeof(xvar) > 1
    if multiple
        gdc = groupby(df,vcat(xvar,err))
    else
        gdc = groupby(df,[xvar,err])
    end
    df1 = combine(gdc,yvar => summary => yvar)
    sort!(df1,xvar)
    if multiple
        label_df = unique(df1[:,xvar])
    else
        label_df = unique(df1[:,[xvar]])
    end
    label_df[!,:xpos] = 1:nrow(label_df)
    leftjoin(df1,label_df; on = xvar)
    try
        firstval = union(df1[:,xvar])[1]
        df1[!,:xpos] = [v == firstval ? 1 : 2  for v in df1[:,xvar]]
    catch
        println("Can't define x position")
    end
    return df1
end

"""
    `group_summary(df1,xvar,yvar; normality = true)`
    using bootstrap computes central measure plus 95% CI
    if normality is true it uses mean as central measure
    else it uses median
"""

function group_summary(df1,xvar,yvar; normality = true)
    if normality
        central = mean
    else
        central = median
    end
    df2 = combine(groupby(df1,xvar)) do dd
        # using bootstrap to calculate the central point and the 95% CI
        b = bootstrap(central, dd[:,yvar], BasicSampling(100000))
        #bootstrap confint returns a tuple containing a tuple with the statistic
        #estimates and the CI lower and uppper bounds
        m, lower, upper = confint(b,PercentileConfInt(0.95))[1]
        ci1 = m - lower
        ci2 = upper - m
        (Central = m, ERR = (ci1,ci2))
    end
    return df2
end

function survivalrate_algorythm(var; step = 0.5, xaxis = nothing)
    isnothing(xaxis) && (xaxis = range(extrema(var)..., step = step))
    survival = 1 .- ecdf(var).(xaxis)
    return (Xaxis = collect(xaxis), fy = survival)
end

function cumulative_algorythm(var; step = 0.5, xaxis = nothing)
    isnothing(xaxis) && (xaxis = range(extrema(var)..., step = step))
    cum = ecdf(var).(xaxis)
    return (Xaxis = collect(xaxis), fy = cum)
end

function hazardrate_algorythm(var; step = 0.5, xaxis = nothing)
    isnothing(xaxis) && (xaxis = range(extrema(var)..., step = step))
    survival = 1 .- ecdf(var).(xaxis)
    hazard = -pushfirst!(diff(survival),0)./survival
    return (Xaxis = collect(xaxis), fy = hazard)
end

"""
    `function_analysis(df,variable, f; grouping = nothing, step =0.05, calc = :basic,
            color = [:auto], linestyle = [:auto])`
    Apply the function f over the vaariable var per each value of grouping and
    plots the result over the variable var
"""

function function_analysis(df,var, f; grouping = nothing, step = 0.5, calc = :basic,
        xaxis = nothing,
        color = [:auto], linestyle = [:auto])
    subgroups = isnothing(grouping) ? [:MOUSE] : vcat(:MOUSE,grouping)
    isnothing(xaxis) && (xaxis = range(extrema(df[:,var])..., step = step))
    dd1 = combine(groupby(df,subgroups), var => (t-> f(t,xaxis = xaxis, step = step)) => AsTable)
    rename!(dd1, Dict(:Xaxis => var))
    sort!(dd1,[:MOUSE,var])
    if calc == :bootstrapping
        dd2 = combine(groupby(dd1,grouping)) do dd3
            group_summary(dd3,var,:fy; normality = false)
        end
        dd2[!,:low] = [x[1] for x in dd2.ERR]
        dd2[!,:up] = [x[2] for x in dd2.ERR]
    elseif calc == :quantiles
        dd2 = combine(groupby(dd1,[grouping,var]), :fy =>(t-> (Central = mean(t),
        low= abs(mean(t) - quantile(t,0.25),
        up = abs(quantile(t,0.975)-mean(t))),
        # ERR = (abs(mean(t) - quantile(t,0.25)) + abs(quantile(t,0.975)-mean(t)))/2,
        SEM = sem(t))) => AsTable)
    elseif calc == :basic
        dd2 = combine(groupby(dd1,[grouping,var]), :fy =>(t-> (Central = mean(t),up = sem(t), low = sem(t))) => AsTable)
    end
    sort!(dd2,var)

    plt = @df dd2 plot(cols(var),:Central, ribbon = (:low, :up), group = cols(grouping), linecolor = :auto, color = color, linestyle = linestyle)
    return plt, dd2
end


using Pkg
Pkg.activate(@__DIR__)

using SigmaRidgeRegression
using LinearAlgebra
using StatsBase
using Plots
using MLJ
using LaTeXStrings
using Random

using PGFPlotsX
pgfplotsx()
push!(PGFPlotsX.CUSTOM_PREAMBLE, raw"\usepackage{amsmath}")
push!(PGFPlotsX.CUSTOM_PREAMBLE, raw"\usepackage{amssymb}")
push!(PGFPlotsX.CUSTOM_PREAMBLE, raw"\usepackage{bm}")
push!(PGFPlotsX.CUSTOM_PREAMBLE, raw"\usepackage[bbgreekl]{mathbbol}")
push!(PGFPlotsX.CUSTOM_PREAMBLE, raw"\newcommand{\sigmacv}{\bbsigma}")
push!(PGFPlotsX.CUSTOM_PREAMBLE, raw"\newcommand{\loo}{\operatorname{CV}^*}")
push!(PGFPlotsX.CUSTOM_PREAMBLE, raw"\newcommand{\blambda}{\bm{\lambda}}")

main_cols = [:purple :green :grey]

Random.seed!(1)
σ = 4.0
grp = GroupedFeatures([25; 25; 25])
n = 400
p = grp.p

X = randn(n, p)

αs = sqrt.([4.0; 8.0; 12.0])
β = random_betas(grp, αs)
#group_summary(grp, β, norm)
#sum(abs2, β)
Y = X * β .+ σ .* randn(n)

loo_sigma_ridge = LooSigmaRidgeRegressor(; groups=grp, center=false, scale=false)

loo_sigmaridge_machine = machine(loo_sigma_ridge, MLJ.table(X), Y)

fit!(loo_sigmaridge_machine)

σs = loo_sigmaridge_machine.report.params

#  \sigmacv <= 6.2
param_subset = loo_sigmaridge_machine.report.params .<= 6.2
λs = Matrix(hcat(loo_sigmaridge_machine.report.λs...)')
λs_subset = λs[param_subset, :]

βs_list = [
    fit!(
        machine(
            SigmaRidgeRegressor(; groups=grp, σ=σ, center=false, scale=false),
            MLJ.table(X),
            Y,
        ),
    ).fitresult.coef for σ in σs
]
βs = Matrix(hcat(βs_list...)')

pl_left = plot(
    loo_sigmaridge_machine.report.params[param_subset],
    λs_subset;
    legend=:topleft,
    linecolor=main_cols,
    linestyle=[:dot :dashdot :dash],
    xlab=L"\sigmacv",
    ylab=L"\widehat{\lambda}_g(\sigmacv)",
    background_color_legend=:transparent,
    foreground_color_legend=:transparent,
    grid=nothing,
    frame=:box,
    label=["Group 1" "Group 2" "Group 3"],
    thickness_scaling=1.8,
    size=(550, 400),
)
savefig(pl_left, "intro_figure_left.tikz")

pl_right = plot(
    loo_sigmaridge_machine.report.params[param_subset],
    loo_sigmaridge_machine.report.loos[param_subset];
    color=:darkblue,
    grid=nothing,
    frame=:box,
    background_color_legend=:transparent,
    foreground_color_legend=:transparent,
    xlab=L"\sigmacv",
    ylab=L"\loo(\sigmacv)",
    label=nothing,
    thickness_scaling=1.8,
    size=(550, 400),
    ylim=(18.2, 19.2),
    yticks=[18.4; 18.6; 18.8; 19.0],
)

savefig(pl_right, "intro_figure_right.tikz")

pl_left_tree = plot(
    loo_sigmaridge_machine.report.params,
    λs;
    ylim=(0, 15),
    legend=:topleft,
    color=main_cols,
    linestyle=[:dot :dashdot :dash],
    background_color_legend=:transparent,
    foreground_color_legend=:transparent,
    grid=nothing,
    frame=:box,
    xlab=L"\sigmacv",
    ylab=L"\widehat{\lambda}_g(\sigmacv)",
    label=["Group 1" "Group 2" "Group 3"],
    thickness_scaling=1.8,
    size=(550, 400),
)

savefig(pl_left_tree, "christmas_tree_left.tikz")

cols_rep = hcat([fill(col, 1, 25) for col in [:purple; :green; :grey]]...)
ltype_rep = hcat([fill(lt, 1, 25) for lt in [:solid; :dash; :dot]]...)

pl_right_tree = plot(
    σs,
    βs;
    alpha=0.6,
    linewidth=0.8,
    label="",
    grid=nothing,
    frame=:box,
    ylab=L"\widehat{w}_j(\widehat{\blambda}(\sigmacv))",
    xlab=L"\sigmacv",
    color=cols_rep,
    ylim=(-2.1, 2.1),
    thickness_scaling=1.8,
    size=(550, 400),
)

savefig(pl_right_tree, "christmas_tree_right.tikz")


struct ReLU{T, O} <: AbstractBlock
    input::O
    ReLU{T, O}() where {T, O} = new{T, O}()
end

ReLU(::Type{T}, ::Type{O}) where {T, O} = ReLU{T, O}()
ReLU(::Type{T};nbatch=1) where T = ReLU(T, nbatch > 1 ? Matrix{T} : Vector{T})
ReLU(;nbatch=1) = ReLU(Float64; nbatch=nbatch)

_relu(x::T) where {T <: Real} = x > 0 ? x : 0
_relu_grad(x::T, grad::T) where {T <: Real} = x > 0 ? grad : 0

function _relu(x::T) where {T <: Complex}
    if real(x) > 0 && imag(x) > 0
        return x
    elseif real(x) > 0 && imag(x) < 0
        return real(x)
    elseif real(x) < 0 && imag(x) > 0
        return imag(x)
    else
        return 0
    end
end

function _relu_grad(x::T, grad::T) where {T <: Complex}
    if real(x) > 0 && imag(x) > 0
        return grad
    elseif real(x) > 0 && imag(x) < 0
        return real(grad)
    elseif real(x) < 0 && imag(x) > 0
        return -im * real(grad)
    else
        return 0
    end
end

function forward(op::ReLU{T, O}, input::O) where {T, O}
    return _relu.(input)
end

function backward(op::ReLU{T, O}, grad::O) where {T, O}
    return _relu_grad.(op.input, grad)
end


@testset "utils/utils.jl" begin
    @test AGP.jitt isa AGP.Jittering
    @test Float64(AGP.jitt) ≈ 1e-4
    @test Float32(AGP.jitt) ≈ 1e-3
    @test Float16(AGP.jitt) ≈ 1e-2

    @test AGP.δ(0, 1) == 0.0
    @test AGP.δ(1, 1) == 1.0

    A = rand(2, 2)
    B = rand(2, 2)
    x = rand(2)
    @test AGP.hadamard(A, B) == A .* B

    X = copy(A)
    AGP.add_transpose!(X)
    @test X == A + A'

    D = A * A' + I
    C = cholesky(D)
    @test AGP.invquad(C, x) ≈ dot(x, D \ x)

    @test AGP.trace_ABt(A, B) ≈ tr(A * B')
    @test AGP.diag_ABt(A, B) ≈ diag(A * B')
    @test AGP.diagv_B(x, B) ≈ Diagonal(x) * B
    @test AGP.κdiagθκ(A, x) ≈ A' * Diagonal(x) * A
    @test AGP.ρκdiagθκ(2.0, A, x) ≈ 2.0 * A' * Diagonal(x) * A
    @test AGP.opt_add_diag_mat(x, A) ≈ A + Diagonal(x)
    @test AGP.safe_expcosh(2.0, 1.0) ≈ exp(2.0) / cosh(1.0)
    @test AGP.logcosh(2.0) ≈ log(cosh(2.0))

    #TODO test
    #=
    function symcat(S::Symmetric, v::AbstractVector, vv::Real)
        S = vcat(S, v')
        S = hcat(S, vcat(v, vv))
        return Symmetric(S)
    end

    function make_grid(range1, range2)
        return hcat(
            [i for i in range1, j in range2][:],
            [j for i in range1, j in range2][:],
        )
    end
    =#
    @test (2.0 * C).L ≈ cholesky(2.0 * D).L
    @test C * x ≈ D * x
    @test (C + 2.0 * I).L ≈ cholesky(D + 2.0 * I).L
end


############## radius.jl  ###############

import RiemannTheta.radius

≈(a,b) = isapprox(a, b, rtol=1e-1)

@testset "radius" begin

    @test radius(ϵ1, T1) ≈ 3.4
    @test radius(ϵ2, T1) ≈ 5.0
    @test radius(ϵ1, T2) ≈ 3.4
    @test radius(ϵ2, T2) ≈ 5.0
    @test radius(ϵ1, T3) ≈ 4.7
    @test radius(ϵ2, T3) ≈ 6.7

    @test radius(ϵ1, T1, derivs1) ≈ 4.1
    @test radius(ϵ1, T3, derivs2) ≈ 9.1

end

############## innerpoints  ###################

import RiemannTheta.innerpoints

@testset "innerpoints" begin

    R = radius(ϵ1, T1)
    pts = innerpoints(T1, R)
    @test length(pts) > 17
    @test length(pts) < 30
    @test all(length.(pts) .== 2)

    R = radius(ϵ2, T3)
    pts = innerpoints(T3, R)
    @test length(pts) > 100
    @test length(pts) < 300
    @test all(length.(pts) .== 10)

end

########### finite_sum.jl ##############

########### lll.jl ##############


function pancakesort!(arr::Vector{<:Real})
    len = length(arr)
    if len < 2 return arr end
    for i in len:-1:2
        j = indmax(arr[1:i])
        if i == j continue end
        arr[1:j] = reverse(arr[1:j])
        arr[1:i] = reverse(arr[1:i])
    end
    return arr
end

v = rand(-10:10, 10)
println("# unordered: $v\n -> ordered: ", pancakesort!(v))


function johnson_shortest_paths(g::AbstractGraph{U},
distmx::AbstractMatrix{T}=weights(g)) where T <: Real where U <: Integer

    nvg = nv(g)
    type_distmx = typeof(distmx)
#Change when parallel implementation of Bellman Ford available
    wt_transform = bellman_ford_shortest_paths(g, vertices(g), distmx).dists

    if !type_distmx.mutable && type_distmx !=  LightGraphs.DefaultDistance
        distmx = sparse(distmx) #Change reference, not value
    end

#Weight transform not needed if all weights are positive.
    if type_distmx !=  LightGraphs.DefaultDistance
        for e in edges(g)
            distmx[src(e), dst(e)] += wt_transform[src(e)] - wt_transform[dst(e)]
        end
    end

    dijk_state = Parallel.dijkstra_shortest_paths(g, vertices(g), distmx)
    dists = dijk_state.dists
    parents = dijk_state.parents

    broadcast!(-, dists, dists, wt_transform)
    for v in vertices(g)
        dists[:, v] .+= wt_transform[v] #Vertical traversal prefered
    end

    if type_distmx.mutable
        for e in edges(g)
            distmx[src(e), dst(e)] += wt_transform[dst(e)] - wt_transform[src(e)]
        end
    end

    return JohnsonState(dists, parents)
end

function enumerate_paths(s::JohnsonState{T,U}, v::Integer) where T <: Real where U <: Integer
    pathinfo = s.parents[v, :]
    paths = Vector{Vector{U}}()
    for i in 1:length(pathinfo)
        if (i == v) || (s.dists[v, i] == typemax(T))
            push!(paths, Vector{U}())
        else
            path = Vector{U}()
            currpathindex = i
            while currpathindex != 0
                push!(path, currpathindex)
                currpathindex = pathinfo[currpathindex]
            end
            push!(paths, reverse(path))
        end
    end
    return paths
end

enumerate_paths(s::JohnsonState) = [enumerate_paths(s, v) for v in 1:size(s.parents, 1)]
enumerate_paths(st::JohnsonState, s::Integer, d::Integer) = enumerate_paths(st, s)[d]


# Advent of Code 2020 - Day 15

function spoken(list, turn)
    spoken_init = parse.(Int,(split(list, ',')))
    spoken_list = fill(0, turn -1)
    for j=1:length(spoken_init)-1
        spoken_list[spoken_init[j]+1] += j
    end
    last_spoken = spoken_init[end]
    for i=length(spoken_init):turn-1
        if spoken_list[last_spoken+1] == 0
            spoken_list[last_spoken+1] = i
            last_spoken = 0
        else
            difference = i - spoken_list[last_spoken+1]
            spoken_list[last_spoken+1] = i
            last_spoken = difference
        end
    end
    return last_spoken
end

# The next Int32 version reduces allocations by half and time
function spoken32(list, turn)
    spoken_init = parse.(Int32,(split(list, ',')))
    spoken_list = fill(Int32(0), turn -1)
    for j=1:length(spoken_init)-1
        spoken_list[spoken_init[j]+1] += Int32(j)
    end
    last_spoken = spoken_init[end]
    for i=length(spoken_init):turn-1
        if spoken_list[last_spoken+1] == 0
            spoken_list[last_spoken+1] = Int32(i)
            last_spoken = 0
        else
            difference = i - spoken_list[last_spoken+1]
            spoken_list[last_spoken+1] = Int32(i)
            last_spoken = difference
        end
    end
    return last_spoken
end

@show spoken("0,3,6", 2020) == 436
@show spoken("1,3,2", 2020) == 1
@show spoken("2,1,3", 2020) == 10
@show spoken("1,2,3", 2020) == 27
@show spoken("2,3,1", 2020) == 78
@show spoken("3,2,1", 2020) == 438
@show spoken("3,1,2", 2020) == 1836
@show spoken("10,16,6,0,1,17", 2020) == 412

@show spoken("0,3,6", 30000000) == 175594
@show spoken("1,3,2", 30000000) == 2578
@show spoken("2,1,3", 30000000) == 3544142
@show spoken("1,2,3", 30000000) == 261214
@show spoken("2,3,1", 30000000) == 6895259
@show spoken("3,2,1", 30000000) == 18
@show spoken("3,1,2", 30000000) == 362
@show spoken("10,16,6,0,1,17", 30000000) == 243
@show spoken32("10,16,6,0,1,17", 30000000) == 243

#=
```julia
julia> @btime spoken("10,16,6,0,1,17", 30000000)
  1.568 s (6 allocations: 228.88 MiB)
243

julia> @btime spoken32("10,16,6,0,1,17", 30000000)
  1.328 s (6 allocations: 114.44 MiB)
243
```
 =#


# Example adapted from Daniel Shiffman - The Coding Train
# https://youtu.be/0jjeOYMjmDU?list=PLRqwX-V7Uu6ZiZxtDDRCi6uhfTH4FilpH

using Alexya

@init "Fractal Tree" 800 600 # create a Canvas
@layout aside(:v, 200)

slider = @create Slider(0:π/12:2π; init=π/4, @margin(20))
φ = π/4

@use function update()
    global φ = value(slider)

    background("#515151")
    origin((@width)/2, @height)
    sethue("white")
    branch(150)
end

function branch(len)
    line(O, Point(0, -len), :stroke)
    translate(0, -len)

    if len > 4
        gsave()
        rotate(φ)
        branch(.67len)
        grestore()

        gsave()
        rotate(-φ)
        branch(.67len)
        grestore()
    end
end

start()


using DelaunayMeshes
using Base.Test

# write your own tests here
include("TestHelpers.jl")
#include("QuadEdgeTests.jl")
#include("TriangulationTests.jl")
#include("DelaunayMeshTests.jl")
include("PostProcessTests.jl")


module IntervalArithmetic

import Arblib

export Interval,
    IComplex, IComplexF64, mid, diam, rad, mig, mag, hull, isinterior, isdisjoint

import Base: *, /, +, -, ^, inv

@static if VERSION ≥ v"1.5.0-"
    import Base: isdisjoint
end

import Printf

# Implementation of interval arithmetic following "Interval Analysis" - Mayer.

struct Interval{T<:Real}
    lo::T
    hi::T
end

Interval{T}(a::Real) where {T<:Real} = Interval(convert(T, a))
Interval{T}(a::Rational) where {T<:Real} =
    Interval(convert(T, numerator(a))) / Interval(convert(T, denominator(a)))
Interval{T}(a::Interval{S}) where {T<:Real,S<:Real} =
    Interval(convert(T, a.lo), convert(T, a.hi))
Interval(a::T) where {T<:Real} = Interval(a, a)
Interval(a::Rational) = Interval{Float64}(a)
Interval(a::T, b::S) where {T<:Real,S<:Real} = Interval(promote(a, b)...)
Interval(a::T, b::T) where {T<:Integer} = Interval(float(a), float(b))

Base.convert(::Type{Interval{T}}, a::Interval{S}) where {T,S} =
    Interval(convert(T, a.lo), convert(T, a.hi))
Base.convert(::Type{Interval{T}}, a::Real) where {T} = Interval{T}(a)

function interval(a::Real, b::Real)
    is_valid_interval(a, b) || invalid_interval_error(a, b)
    Interval(a, b)
end
is_valid_interval(a::Real, b::Real) = isfinite(a) && isfinite(b) && a ≤ b
@noinline function invalid_interval_error(a, b)
    throw(
        ArgumentError(
            "`[$a, $b]` is not a valid interval. Need `a ≤ b` to construct `interval(a, b)`.",
        ),
    )
end

Base.hash(x::Interval, h::UInt) = hash(x.hi, hash(x.lo, u))
Base.:(==)(a::Interval, b::Interval) = a.lo == b.lo && a.hi == b.hi
Base.eltype(x::Interval{T}) where {T} = T

Base.promote_rule(::Type{Interval{T}}, ::Type{Interval{S}}) where {T<:Real,S<:Real} =
    Interval{promote_type(T, S)}
Base.promote_rule(::Type{Interval{T}}, ::Type{S}) where {T<:Real,S<:Real} =
    Interval{promote_type(T, S)}

round_up(a::AbstractFloat) = nextfloat(a)
round_down(a::AbstractFloat) = prevfloat(a)
macro round(a, b)
    :(Interval(round_down($(esc(a))), round_up($(esc(b)))))
end

mid(a::Interval) = (a.lo + a.hi) / 2
diam(a::Interval) = round_up(a.hi - a.lo)
function rad(a::Interval)
    m = mid(a)
    return round_up(max(m - a.lo, a.hi - m))
end
mag(a::Interval) = max(abs(a.lo), abs(a.hi))
function mig(a::Interval{T}) where {T}
    zero(T) ∈ a && return zero(T)
    min(abs(a.lo), abs(a.hi))
end
hull(a::Interval, b::Interval) = Interval(min(a.lo, b.lo), max(a.hi, b.hi))

Base.issubset(a::Interval, b::Interval) = (a.lo ≥ b.lo) && (a.hi ≤ b.hi)
isinterior(a::Interval, b::Interval) = (a.lo > b.lo) && (a.hi < b.hi)
isdisjoint(a::Interval, b::Interval) = (b.hi < a.lo) || (a.hi < b.lo)
function Base.intersect(a::Interval, b::Interval)
    c = Interval(max(a.lo, b.lo), min(a.hi, b.hi))
    c.lo ≤ c.hi || return Interval(convert(eltype(c), NaN))
    c
end
Base.isempty(a::Interval) = isnan(a.lo) || isnan(a.hi)
Base.in(x::Number, a::Interval) = a.lo ≤ x ≤ a.hi

function Base.show(io::IO, a::Interval)
    print(io, mid(a), " ± ")
    Printf.@printf(io, "%.5g", rad(a))
end
Base.print(io::IO, a::Interval) = print(io, mid(a), " ± ", rad(a))

# arithmetic

# zero, one, typemin, typemax
Base.zero(a::Interval{T}) where {T} = Interval(zero(T))
Base.zero(::Type{Interval{T}}) where {T} = Interval(zero(T))
Base.one(a::Interval{T}) where {T} = Interval(one(T))
Base.one(::Type{Interval{T}}) where {T} = Interval(one(T))
# The arithmetic of Interval is based on the IntervalArithmetic.jl package
# https://github.com/JuliaIntervals/IntervalArithmetic.jl/
# IntervalArithmetic.jl is licensed under the MIT "Expat" License

## Addition and subtraction
+(a::Interval) = a
+(a::Interval, b::Real) = @round(a.lo + b, a.hi + b)
+(b::Real, a::Interval) = a + b
function +(a::Interval, b::Interval)
    @round(a.lo + b.lo, a.hi + b.hi)
end

-(a::Interval) = Interval(-a.hi, -a.lo)
-(a::Interval, b::Real) = @round(a.lo - b, a.hi - b)
-(b::Real, a::Interval) = @round(b - a.hi, b - a.lo)
-(a::Interval, b::Interval) = @round(a.lo - b.hi, a.hi - b.lo)

## Multiplication
function *(x::Real, a::Interval)
    (iszero(a) || iszero(x)) && return zero(a.lo * x)
    if x ≥ 0.0
        return @round(a.lo * x, a.hi * x)
    else
        return @round(a.hi * x, a.lo * x)
    end
end
*(a::Interval, x::Real) = x * a
function *(a::Interval, b::Interval)
    if b.lo >= zero(b.lo)
        a.lo >= zero(a.lo) && return @round(*(a.lo, b.lo), *(a.hi, b.hi))
        a.hi <= zero(a.hi) && return @round(*(a.lo, b.hi), *(a.hi, b.lo))
        return @round(a.lo * b.hi, a.hi * b.hi)   # zero(T) ∈ a
    elseif b.hi <= zero(b.hi)
        a.lo >= zero(a.lo) && return @round(*(a.hi, b.lo), *(a.lo, b.hi))
        a.hi <= zero(a.hi) && return @round(*(a.hi, b.hi), *(a.lo, b.lo))
        return @round(a.hi * b.lo, a.lo * b.lo)   # zero(T) ∈ a
    else
        a.lo > zero(a.lo) && return @round(*(a.hi, b.lo), *(a.hi, b.hi))
        a.hi < zero(a.hi) && return @round(*(a.lo, b.hi), *(a.lo, b.lo))
        return @round(min(*(a.lo, b.hi), *(a.hi, b.lo)), max(*(a.lo, b.lo), *(a.hi, b.hi)))
    end
end

function Base.muladd(a::Interval, b::Interval, c::Interval)
    lo = let
        lo1 = muladd(a.lo, b.lo, c.lo)
        lo2 = muladd(a.lo, b.hi, c.lo)
        lo3 = muladd(a.hi, b.lo, c.lo)
        lo4 = muladd(a.hi, b.hi, c.lo)
        round_down(Base.FastMath.min_fast(lo1, lo2, lo3, lo4))
    end

    hi = let
        hi1 = muladd(a.lo, b.lo, c.hi)
        hi2 = muladd(a.lo, b.hi, c.hi)
        hi3 = muladd(a.hi, b.lo, c.hi)
        hi4 = muladd(a.hi, b.hi, c.hi)
        round_up(Base.FastMath.max_fast(hi1, hi2, hi3, hi4))
    end

    Interval(lo, hi)
end
function Base.muladd(a::Interval{T}, b::T, c::Interval) where {T}
    lo = let
        lo1 = muladd(a.lo, b, c.lo)
        lo2 = muladd(a.hi, b, c.lo)
        round_down(Base.FastMath.min_fast(lo1, lo2))
    end

    hi = let
        hi1 = muladd(a.lo, b, c.hi)
        hi2 = muladd(a.hi, b, c.hi)
        round_up(Base.FastMath.max_fast(hi1, hi2))
    end

    Interval(lo, hi)
end
function Base.muladd(a::T, b::Interval{T}, c::Interval) where {T}
    lo = let
        lo1 = muladd(a, b.lo, c.lo)
        lo2 = muladd(a, b.hi, c.lo)
        round_down(Base.FastMath.min_fast(lo1, lo2))
    end

    hi = let
        hi1 = muladd(a, b.lo, c.hi)
        hi2 = muladd(a, b.hi, c.hi)
        round_up(Base.FastMath.max_fast(hi1, hi2))
    end

    Interval(lo, hi)
end

function inv(a::Interval{T}) where {T<:Real}
    if zero(T) ∈ a
        return Interval(convert(T, NaN))
    end
    @round(inv(a.hi), inv(a.lo))
end

function /(a::Interval{T}, x::T) where {T<:Real}
    iszero(a) && return zero(Interval{T})

    if x ≥ 0.0
        return @round(a.lo / x, a.hi / x)
    else
        return @round(a.hi / x, a.lo / x)
    end
end
function /(a::Interval{T}, b::Interval{T}) where {T<:Real}
    if iszero(a)
        S = typeof(a.lo / b.lo)
        return zero(Interval{S})
    end

    # b contains zero
    if b.lo > zero(T) # b strictly positive
        a.lo >= zero(T) && return @round(a.lo / b.hi, a.hi / b.lo)
        a.hi <= zero(T) && return @round(a.lo / b.lo, a.hi / b.hi)
        return @round(a.lo / b.lo, a.hi / b.lo)  # zero(T) ∈ a

    elseif b.hi < zero(T) # b strictly negative

        a.lo >= zero(T) && return @round(a.hi / b.hi, a.lo / b.lo)
        a.hi <= zero(T) && return @round(a.hi / b.lo, a.lo / b.hi)
        return @round(a.hi / b.hi, a.lo / b.hi)  # zero(T) ∈ a
    else   # b contains zero, but is not zero(b)
        S = typeof(a.lo / b.lo)
        return Interval(convert(S, NaN))
    end
end

function sqr(a::Interval{T}) where {T<:Real}
    if a.lo ≥ zero(T)
        return @round(a.lo^2, a.hi^2)

    elseif a.hi ≤ zero(T)
        return @round(a.hi^2, a.lo^2)
    end

    return @round(mig(a)^2, mag(a)^2)
end
Base.literal_pow(::typeof(^), a::Interval, ::Val{2}) = sqr(a)

function ^(x::Interval, n::Integer)  # fast integer power
    if n < 0
        return inv(pow(x, -n))
    end
    isempty(x) && return x
    if iseven(n) && 0 ∈ x
        return hull(
            zero(x),
            hull(
                Base.power_by_squaring(Interval(mig(x)), n),
                Base.power_by_squaring(Interval(mag(x)), n),
            ),
        )

    else
        return hull(
            Base.power_by_squaring(Interval(x.lo), n),
            Base.power_by_squaring(Interval(x.hi), n),
        )
    end
end

struct IComplex{T} <: Number
    re::Interval{T}
    im::Interval{T}
end
IComplex(x::Real, y::Interval) = IComplex(promote(x, y)...)
IComplex(x::Interval, y::Real) = IComplex(promote(x, y)...)
function IComplex(x::Real, y::Real)
    ix, iy = promote(x, y)
    IComplex(Interval(ix), Interval(iy))
end
IComplex(re::Interval{T}, im::Interval{S}) where {T<:Real,S<:Real} =
    IComplex{promote_type(T, S)}(promote(re, im)...)
IComplex(x::Union{Interval,Real}) = IComplex(x, zero(x))
IComplex(z::Complex) = IComplex(real(z), imag(z))
IComplex(z::IComplex) = z
Base.complex(x::Interval, y::Real) = IComplex(x, y)
Base.complex(x::Real, y::Interval) = IComplex(x, y)
Base.complex(x::Interval, y::Interval) = IComplex(x, y)
Base.complex(x::IComplex) = x

const IComplexF64 = IComplex{Float64}
IComplex{T}(x::Complex{S}) where {T<:Real,S<:Real} =
    IComplex(Interval{T}(real(x)), Interval{T}(imag(x)))
IComplex{T}(x::IComplex{S}) where {T<:Real,S<:Real} =
    IComplex(Interval{T}(real(x)), Interval{T}(imag(x)))
IComplex{T}(x) where {T<:Real} = IComplex(Interval{T}(x))
# (::Type{IComplex{T}})(z::IComplex{T}) where {T<:Real} = z

Base.zero(a::IComplex{T}) where {T} = IComplex(zero(Interval{T}))
Base.zero(::Type{IComplex{T}}) where {T} = IComplex(zero(Interval{T}))
Base.one(a::IComplex{T}) where {T} = IComplex(one(Interval{T}))
Base.one(::Type{IComplex{T}}) where {T} = IComplex(one(Interval{T}))

function Base.show(io::IO, c::IComplex)
    print(io, "(")
    show(io, c.re)
    print(io, ") + (")
    show(io, c.im)
    print(io, ")im")
end
Base.print(io::IO, c::IComplex) = print(io, "(", c.re, ") + (", c.im, ")im")
Base.Base.broadcastable(z::IComplex) = z
Base.promote_rule(::Type{IComplex{T}}, ::Type{S}) where {T,S<:Real} =
    IComplex{promote_type(T, S)}
Base.promote_rule(::Type{IComplex{T}}, ::Type{Complex{S}}) where {T,S} =
    IComplex{promote_type(T, S)}
Base.promote_rule(::Type{IComplex{T}}, ::Type{Interval{S}}) where {T,S<:Real} =
    IComplex{promote_type(T, S)}
Base.promote_rule(::Type{IComplex{T}}, ::Type{IComplex{S}}) where {T,S} =
    IComplex{promote_type(T, S)}
Base.convert(::Type{IComplex{T}}, x::IComplex) where {T} = IComplex{T}(x)
Base.convert(::Type{IComplex{T}}, x::Interval) where {T} = IComplex{T}(x)

Base.widen(::Type{IComplex{T}}) where {T} = IComplex{widen(T)}

Base.real(z::IComplex) = z.re
Base.real(C::Type{IComplex{T}}) where {T} = Interval{T}

Base.imag(z::IComplex) = z.im
Base.imag(C::Type{IComplex{T}}) where {T} = Interval{T}
Base.reim(z::IComplex) = (real(z), imag(z))

Base.conj(z::IComplex) = IComplex(real(z), -imag(z))

+(z::IComplex) = z
+(x::Union{Interval,Real}, z::IComplex) = IComplex(x + real(z), imag(z))
+(z::IComplex, x::Union{Interval,Real}) = IComplex(x + real(z), imag(z))

+(z::IComplex, w::IComplex) = IComplex(real(z) + real(w), imag(z) + imag(w))
-(z::IComplex) = IComplex(-real(z), -imag(z))
-(z::IComplex, w::IComplex) = IComplex(real(z) - real(w), imag(z) - imag(w))
-(x::Union{Interval,Real}, z::IComplex) = IComplex(x - real(z), -imag(z))
-(z::IComplex, x::Union{Interval,Real}) = IComplex(real(z) - x, imag(z))

*(z::IComplex, w::IComplex) =
    IComplex(real(z) * real(w) - imag(z) * imag(w), real(z) * imag(w) + imag(z) * real(w))
*(x::Union{Interval,Real}, z::IComplex) = IComplex(x * real(z), x * imag(z))
*(z::IComplex, x::Union{Interval,Real}) = IComplex(x * real(z), x * imag(z))

Base.muladd(z::Union{Complex,IComplex}, w::Union{Complex,IComplex}, x::IComplex) = IComplex(
    muladd(real(z), real(w), real(x)) - imag(z) * imag(w),
    muladd(real(z), imag(w), muladd(imag(z), real(w), imag(x))),
)

function inv(b::IComplex{T}) where {T}
    bre, bim = reim(b)
    denom = bre^2 + bim^2
    IComplex(bre / denom, (-bim) / denom)
end
/(a::R, z::S) where {R<:Real,S<:IComplex} = (T = promote_type(R, S); a * inv(T(z)))
/(z::IComplex, x::Union{Interval,Real}) = IComplex(real(z) / x, imag(z) / x)
function /(a::IComplex{T}, b::IComplex{T}) where {T}
    are, aim = reim(a)
    bre, bim = reim(b)
    denom = bre^2 + bim^2
    IComplex((are * bre + aim * bim) / denom, (aim * bre - are * bim) / denom)
end

mid(z::IComplex) = Complex(mid(real(z)), mid(imag(z)))
diam(z::IComplex) = max(diam(real(z)), diam(imag(z)))
rad(z::IComplex) = max(rad(real(z)), rad(imag(z)))
mag(z::IComplex) = max(mag(real(z)), mag(imag(z)))
isinterior(a::IComplex, b::IComplex) =
    isinterior(real(a), real(b)) && isinterior(imag(a), imag(b))
Base.issubset(a::IComplex, b::IComplex) = real(a) ⊆ real(b) && imag(a) ⊆ imag(b)
isdisjoint(a::IComplex, b::IComplex) =
    isdisjoint(real(a), real(b)) || isdisjoint(imag(a), imag(b))
Base.intersect(a::IComplex, b::IComplex) =
    IComplex(intersect(real(a), real(b)), intersect(imag(a), imag(b)))
Base.in(x::Number, a::IComplex) = in(real(x), real(a)) && in(imag(x), imag(a))

function inf_norm_bound(A::AbstractMatrix{IComplex{T}}) where {T}
    # ||A|| = maxᵢ ∑ⱼ|Aᵢⱼ| ≤ √2 * maxᵢ ∑ⱼmax(|real(Aᵢⱼ)|,|imag(Aᵢⱼ)|)
    bound = zero(T)
    for i = 1:size(A, 1)
        # use again an interval arithmetic to account for accumulation error
        bᵢ = Interval(zero(T))
        for j = 1:size(A, 2)
            bᵢ += mag(A[i, j])
        end
        bound = max(bound, mag(bᵢ))
    end
    # Upper bound √2 by 1.41422, this also accounts
    # for the multiplcation error in √2 * bound
    1.41422 * bound
end

## Arb
function IComplexF64(
    z::Arblib.AcbLike,
    a = Arblib.Arf(prec = 53),
    b = Arblib.Arf(prec = 53),
)
    Arblib.get_interval!(a, b, Arblib.realref(z); prec = 53)
    re = Interval(Arblib.get_d(a, RoundDown), Arblib.get_d(b, RoundUp))
    Arblib.get_interval!(a, b, Arblib.imagref(z); prec = 53)
    im = Interval(Arblib.get_d(a, RoundDown), Arblib.get_d(b, RoundUp))
    IComplexF64(re, im)
end

function Base.convert(
    ::Type{T},
    x::AbstractVector{IComplexF64},
) where {
    T<:Union{Arblib.AcbVector,Arblib.AcbRefVector,Arblib.AcbMatrix,Arblib.AcbRefMatrix},
}
    y = T(mid.(x), prec = 53)
    m = Arblib.Mag()
    for (i, xᵢ) in enumerate(x)
        m[] = rad(xᵢ)
        Arblib.add_error!(y[i], m)
    end
    y
end

function Base.setindex!(z::Union{Arblib.Acb,Arblib.AcbRef}, x::IComplexF64)
    rz = Arblib.realref(z)
    iz = Arblib.imagref(z)

    Arblib.midref(rz)[] = mid(real(x))
    Arblib.midref(iz)[] = mid(imag(x))

    Arblib.radref(rz)[] = rad(real(x))
    Arblib.radref(iz)[] = rad(imag(x))
    z
end
function Base.setindex!(
    A::Union{Arblib.AcbMatrix,Arblib.AcbRefMatrix},
    x::IComplexF64,
    i::Integer,
    j::Integer,
)
    Arblib.ref(A, i, j)[] = x
end

end #module


using UnivariateDifferentiableOptimizationAlgorithms
using Base.Test
# write your own tests here

# include("testcompare-1D-6-variant.jl")


using Bridge, StaticArrays, Bridge.Models
const R = ℝ
Random.seed!(2)
Random.seed!(3)
πH = 2000. # prior

t = 1.0
T = 5.00
n = 50001 # total imputed length
m = 100 # number of segments
M = div(n-1,m)
skippoints = 2
dt = (T-t)/(n-1)
tt = t:dt:T
si = 3.
# 10, 20, 8/3 srand(2)
P = Bridge.Models.Lorenz(ℝ{3}(10, 28, 8/3), ℝ{3}(si,si,si))
P2 = Psmooth = Bridge.Models.Lorenz(ℝ{3}(10, 28, 8/3), ℝ{3}(0,0,0))

x0 = Models.x0(P)
#x0 =  ℝ{3}(6, 0, 2)

W = sample(tt, Wiener{ℝ{3}}())
X = SamplePath(tt, zeros(ℝ{3}, length(tt)))
Bridge.solve!(Euler(), X, x0, W, P)
W = sample(tt, Wiener{ℝ{3}}())
X2 = SamplePath(tt, zeros(ℝ{3}, length(tt)))
Bridge.solve!(Euler(), X2, x0, W, P2)

Xtrue = copy(X)

# Observation scheme and subsample
_pairs(collection) = Base.Generator(=>, keys(collection), values(collection))
SV = ℝ{3}
SM = typeof(one(Bridge.outer(zero(SV))))
if isdefined(:partial)
    if !partial
        L = I
        Σ = SDiagonal(1., 1., 1.)
        lΣ = chol(Σ)'
        RV = SV
        RM = SM

        V = SamplePath(collect(_pairs(Xtrue))[1:M:end])
        map!(y -> L*y + lΣ*randn(RV), V.yy, V.yy)
    else
        L = @SMatrix [0.0 1.0 0.0; 0.0 0.0 1.0]
        Σ = SDiagonal(1., 1.)
        lΣ = chol(Σ)'
        RV = ℝ{2}
        RM = typeof(one(Bridge.outer(zero(RV))))

        V_ = SamplePath(collect(_pairs(Xtrue))[1:M:end])
        V = SamplePath(V_.tt, map(y -> L*y + lΣ*randn(RV), V_.yy))
    end
else
    println("Info:  No observation scheme defined. Set `partial in [true, false]`")
end


### A Pluto.jl notebook ###
# v0.15.1

using Markdown
using InteractiveUtils

# ╔═╡ f041ef20-f968-11eb-1013-ef1b063dc4cc
begin
	using Pkg
	Pkg.activate(".")
	using Plots
	using SSpline
end

# ╔═╡ 0d9037b8-b0ff-4a4a-8d7f-80b29338ac19
html"""<style> main { max-width: 900px;} """

# ╔═╡ 470b7ee9-b8f7-4c8d-abfd-7fa871abc36f
x = collect(1:10);y = rand(10);

# ╔═╡ d0253bf6-4564-4c3f-9c80-3f59bf3f5370
md"""

(1)

    s[i](x)    = a[i] * (x-x[i])^3 + b[i] * (x-x[i])^2 + c[i] * (x-x[i]) + d[i]

----

    s'[i](x)   = 3 * a[i] * (x-x[i])^2 + 2 * b[i] * (x-x[i]) + c[i]

    s''[i](x)  = 6 * a[i] * (x-x[i]) + 2 * b[i]

    s'''[i](x) = 6 * a[i]

(2)

    s_i(x_i)     = y_i, i=0..n-1

																s_n-1 = y_n

(3)

    s_i(x_i+1)   = s_i+1(x_i+1),      i = 0..n-2

(4)

    s'_i(x_i+1)  = s'_i+1(x_i+1),     i = 0..n-2

(5)

    s''_i(x_i+1) = s''_i+1(x_i+1),    i = 0..n-2

Natural

    s''_0(x_0) = 0
    s''_n-1(x_n) = 0

Clamped

    s'_0(x_0) = y'_0
    s'_n-1(x_n) = y'_n

Periodic

    s'_0(x_0) = s'_n-1(x_n)
    s''_0(x_0) = s''_n-1(x_n)

NAK

    s'''_0(x_1) = s'''_1(x_1)
    s'''_n-2(x_n-1) = s'''_n-1(x_n-1)

----

	h = x_i+1 - x_i

from (1) `s[i](x) = a[i] * (x-x[i])^3 + b[i] * (x-x[i])^2 + c[i] * (x-x[i]) + d[i]`

    d_i = y_i

    σ_i = s''(x_i) = 6 * a_i * (x_i - x_i) + 2 * b_i = 2 * b_i,                          i = 0..n

    σ_i = 2 * b_i

    b_i = σ_i / 2

----

    σ_i+1 = 6 * a_i * h + 2 * b_i

    σ_i+1 = 6 * a_i * h + 2 * σ_i / 2

from (5) `s''_i(x_i+1) = s''_i+1(x_i+1),    i = 0..n-2`

	a_i = (σ_i+1 - σ_i) / (6 * h)

from (3) `s_i(x_i+1) = s_i+1(x_i+1),      i = 0..n-2`

    y_i+1 = a_i * h^3 + b_i * h^2 + c_i * h + y_i

    y_i+1 = (σ_i+1 - σ_i) / (6 * h) * h^3 + σ_i / 2 * h^2 + c_i * h + y_i

    y_i+1 - y_1 = (σ_i+1 - σ_1) / (6 * h) * h^3 + σ_i / 2 * h^2) + c_i * h

    (y_i+1 - y_1) / h = (σ_i+1 - σ_1) / 6 * h + σ_i / 2 * h + c_i

    (y_i+1 - y_1) / h - h * ( (σ_i+1 - σ_1) / 6 + σ_i / 2) = c_i

	c_i = (y_i+1 - y_1) / h - h * ( (σ_i+1 - σ_1) / 6 + 3σ_i / 6 )

    c_i = (y_i+1 - y_1) / h - h * (σ_i+1 - σ_1 + 3 * σ_i) / 6

    c_i = (y_i+1 - y_1) / h - h * (2 * σ_i + σ_i+1) / 6

from (4) `s'_i(x_i+1)  = s'_i+1(x_i+1),     i = 0..n-2`

	c_i+1 = 3 * a_i * h^2 + 2 * b_i * h + c_i

Insert a_i, b_i, c_i, and use

from (2) `s_i(x_i) = y_i, i=0..n-1)`

    c_i+1 = 3 * (σ_i+1 - σ_i) / (6 * h) * h^2 + 2 * σ_i / 2 * h + (y_i+1 - y_1) / h - h * (2 * σ_i + σ_i+1) / 6

"""

# ╔═╡ c82ec4aa-827e-40f0-ab10-543380e5d5d7
begin
	range = collect(1:0.1:10)

	nak = nakspline(x,y)
	kx, ky = interp(nak,range)

	nat = spline3(x,y,:natural)
	nx, ny = interp(nat,range)
end

# ╔═╡ 88c79600-cc20-4ad9-9a98-6d05e2f7512e
nakspline(x,y)

# ╔═╡ 10458fe7-350f-481f-99ca-29e5c185d3b9

# ╔═╡ Cell order:
# ╠═f041ef20-f968-11eb-1013-ef1b063dc4cc
# ╠═0d9037b8-b0ff-4a4a-8d7f-80b29338ac19
# ╠═470b7ee9-b8f7-4c8d-abfd-7fa871abc36f
# ╟─d0253bf6-4564-4c3f-9c80-3f59bf3f5370
# ╟─c82ec4aa-827e-40f0-ab10-543380e5d5d7
# ╠═88c79600-cc20-4ad9-9a98-6d05e2f7512e
# ╠═10458fe7-350f-481f-99ca-29e5c185d3b9


#=
   Gender height.

   Identify a person's sex by height.

   Cf ~/cplint/gender_height.pl
      ~/blog/gender_height.blog
      ~/webppl/gender_height.wppl

   Note: The observations of height and gender seems to be much better
         using Turings standard "parameter observation" than using
         Dirac etc.

=#

# using Memoization
using Turing, StatsPlots, DataFrames
# using ReverseDiff, Zygote, Tracker
# Turing.setadbackend(:reversediff)
# Turing.setadbackend(:zygote)
# Turing.setadbackend(:tracker)
include("jl_utils.jl")

@model function gender_height(gender=missing,height=missing)
# @model function gender_height()

    male = 1
    female = 2
    gender ~ Categorical([0.5,0.5])

    # From https://en.wikipedia.org/wiki/List_of_average_human_height_worldwide
    # Here are the values of Sweden. I'm not sure about the variance of these heights...
    height ~ gender == male ? Normal(181.5,sqrt(50)) : Normal(166.8,sqrt(50))

    # height ~ gender == male ? Normal(178,17.7) : Normal(163,17.3)

    # Note: Using this way of observing is not good in this model.
    # height == 190.0 || begin Turing.@addlogprob! -Inf; return end
    # obs!(_sampler,_varinfo,height == 190.0)
    # true ~ Dirac(height == 190.0)

    ret = []
    if !ismissing(height)
        push!(ret,height)
    end
    if !ismissing(gender)
        push!(ret,gender)
    end

    return height,gender # ret
    # return ret
end

function run_model(gender=missing,height=missing)

    model = gender_height(gender,height)

    num_chains = 4

    # chains = sample(model, Prior(), MCMCThreads(), 10_000, num_chains)

    chains = sample(model, MH(), 10_000)
    # chains = sample(model, MH(), MCMCThreads(), 40_000, num_chains)

    # chains = sample(model, PG(20), MCMCThreads(), 10_000, num_chains)
    # chains = sample(model, PG(20), 1_000)

    # chains = sample(model, IS(), MCMCThreads(), 10_000, num_chains)
    # chains = sample(model, IS(), 10_000)

    # chains = sample(model, SMC(1000), MCMCThreads(), 10_000, num_chains)
    # chains = sample(model, SMC(1000), 10_000)

    # chains = sample(model, NUTS(1000,0.65), 1_000)
    # chains = sample(model, Gibbs(MH(:gender),NUTS(1000,0.65,:height)), 1_000)
    # chains = sample(model, Gibbs(MH(:gender),NUTS(10,0.65,:height)), 1_000)
    # chains = sample(model, Gibbs(MH(:gender),HMC(0.1,5,:height)), 1_000)
    # chains = sample(model, Gibbs(PG(10,:gender),HMC(0.1,5,:height)), 1_000)
    # chains = sample(model, Gibbs(MH(:gender),NUTS(1_000,0.65,:height)), 1_000)

    if ismissing(height)
        display(chains)
        # display(plot(chains))
    end
    show_var_dist_pct(chains, :gender)
    show_var_dist_pct(chains, :height, 20)

    gen = generated_quantities(model, chains)

    if gen[1] != nothing
        show_var_dist_pct(gen, 20)
    end
end

male = 1
female = 2
for height in [150.0,160.0,170.0,175.0,180.0, 185.0,190.0,200.0]
    println("\ngender:missing height:$height")
    run_model(missing,height)
    println()
end

println("\nmale, missing")
run_model(male,missing)

println("\nfemale, missing")
run_model(female,missing)