text:
####
#### Implementation for Hamiltonian dynamics trajectories
####
#### Developers' Notes
####
#### Not all functions that use `rng` requires a callback function with `GLOBAL_RNG`
#### as default. In short, only those exported to other libries need such a callback
#### function. Internal uses shall always use the explict `rng` version. (Kai Xu 6/Jul/19)

"""
A transition that contains the phase point and
other statistics of the transition.
"""
struct Transition{P<:PhasePoint, NT<:NamedTuple}
    z       ::  P
    stat    ::  NT
end

stat(t::Transition) = t.stat

"""
Abstract Markov chain Monte Carlo proposal.
"""
abstract type AbstractProposal end

"""
Hamiltonian dynamics numerical simulation trajectories.
"""
abstract type AbstractTrajectory{I<:AbstractIntegrator} <: AbstractProposal end

##
## Sampling methods for trajectories.
##

"""
Sampler carried during the building of the tree.
"""
abstract type AbstractTrajectorySampler end

"""
Slice sampler carried during the building of the tree.
It contains the slice variable and the number of acceptable condidates in the tree.
"""
struct SliceTS{F<:AbstractFloat} <: AbstractTrajectorySampler
    zcand   ::  PhasePoint
    ℓu      ::  F     # slice variable in log space
    n       ::  Int   # number of acceptable candicates, i.e. those with prob larger than slice variable u
end

Base.show(io::IO, s::SliceTS) = print(io, "SliceTS(ℓu=$(s.ℓu), n=$(s.n))")

"""
Multinomial sampler carried during the building of the tree.
It contains the weight of the tree, defined as the total probabilities of the leaves.
"""
struct MultinomialTS{F<:AbstractFloat} <: AbstractTrajectorySampler
    zcand   ::  PhasePoint
    ℓw      ::  F     # total energy for the given tree, i.e. sum of energy of all leaves
end

"""
Slice sampler for the starting single leaf tree.
Slice variable is initialized.
"""
SliceTS(rng::AbstractRNG, z0::PhasePoint) = SliceTS(z0, log(rand(rng)) - energy(z0), 1)

"""
Multinomial sampler for the starting single leaf tree.
(Log) weights for leaf nodes are their (unnormalised) Hamiltonian energies.

Ref: https://github.com/stan-dev/stan/blob/develop/src/stan/mcmc/hmc/nuts/base_nuts.hpp#L226
"""
MultinomialTS(rng::AbstractRNG, z0::PhasePoint) = MultinomialTS(z0, zero(energy(z0)))

"""
Create a slice sampler for a single leaf tree:
- the slice variable is copied from the passed-in sampler `s` and
- the number of acceptable candicates is computed by comparing the slice variable against the current energy.
"""
SliceTS(s::SliceTS, H0::AbstractFloat, zcand::PhasePoint) =
    SliceTS(zcand, s.ℓu, (s.ℓu <= -energy(zcand)) ? 1 : 0)

"""
Multinomial sampler for a trajectory consisting only a leaf node.
- tree weight is the (unnormalised) energy of the leaf.
"""
MultinomialTS(s::MultinomialTS, H0::AbstractFloat, zcand::PhasePoint) =
    MultinomialTS(zcand, H0 - energy(zcand))

function combine(rng::AbstractRNG, s1::SliceTS, s2::SliceTS)
    @assert s1.ℓu == s2.ℓu "Cannot combine two slice sampler with different slice variable"
    n = s1.n + s2.n
    zcand = rand(rng) < s1.n / n ? s1.zcand : s2.zcand
    SliceTS(zcand, s1.ℓu, n)
end

function combine(zcand::PhasePoint, s1::SliceTS, s2::SliceTS)
    @assert s1.ℓu == s2.ℓu "Cannot combine two slice sampler with different slice variable"
    n = s1.n + s2.n
    SliceTS(zcand, s1.ℓu, n)
end

function combine(rng::AbstractRNG, s1::MultinomialTS, s2::MultinomialTS)
    ℓw = logaddexp(s1.ℓw, s2.ℓw)
    zcand = rand(rng) < exp(s1.ℓw - ℓw) ? s1.zcand : s2.zcand
    MultinomialTS(zcand, ℓw)
end

function combine(zcand::PhasePoint, s1::MultinomialTS, s2::MultinomialTS)
    ℓw = logaddexp(s1.ℓw, s2.ℓw)
    MultinomialTS(zcand, ℓw)
end

mh_accept(
    rng::AbstractRNG,
    s::SliceTS,
    s′::SliceTS
) = rand(rng) < min(1, s′.n / s.n)

mh_accept(
    rng::AbstractRNG,
    s::MultinomialTS,
    s′::MultinomialTS
) = rand(rng) < min(1, exp(s′.ℓw - s.ℓw))

"""
    transition(τ::AbstractTrajectory{I}, h::Hamiltonian, z::PhasePoint)

Make a MCMC transition from phase point `z` using the trajectory `τ` under Hamiltonian `h`.

NOTE: This is a RNG-implicit callback function for `transition(GLOBAL_RNG, τ, h, z)`
"""
transition(
    τ::AbstractTrajectory{I},
    h::Hamiltonian,
    z::PhasePoint
) where {I<:AbstractIntegrator} = transition(GLOBAL_RNG, τ, h, z)

###
### Standard HMC implementation with fixed leapfrog step numbers.
###
struct StaticTrajectory{I<:AbstractIntegrator} <: AbstractTrajectory{I}
    integrator  ::  I
    n_steps     ::  Int
end
Base.show(io::IO, τ::StaticTrajectory) = print(io, "StaticTrajectory(integrator=$(τ.integrator), λ=$(τ.n_steps)))")

function transition(
    rng::AbstractRNG,
    τ::StaticTrajectory,
    h::Hamiltonian,
    z::PhasePoint
) where {T<:Real}
    H0 = energy(z)
    integrator = jitter(rng, τ.integrator)
    z′ = step(rng, integrator, h, z, τ.n_steps)
    # Accept via MH criteria
    is_accept, α = mh_accept_ratio(rng, H0, energy(z′))
    if is_accept
        # Reverse momentum variable to preserve reversibility
        z = PhasePoint(z′.θ, -z′.r, z′.ℓπ, z′.ℓκ)
    end
    H = energy(z)
    tstat = merge(
        (
         n_steps=τ.n_steps,
         is_accept=is_accept,
         acceptance_rate=α,
         log_density=z.ℓπ.value,
         hamiltonian_energy=H,
         hamiltonian_energy_error=H - H0
        ),
        stat(integrator)
    )
    return Transition(z, tstat)
end

abstract type DynamicTrajectory{I<:AbstractIntegrator} <: AbstractTrajectory{I} end

###
### Standard HMC implementation with fixed total trajectory length.
###
struct HMCDA{I<:AbstractIntegrator} <: DynamicTrajectory{I}
    integrator  ::  I
    λ           ::  AbstractFloat
end
Base.show(io::IO, τ::HMCDA) = print(io, "HMCDA(integrator=$(τ.integrator), λ=$(τ.λ)))")

function transition(
    rng::AbstractRNG,
    τ::HMCDA,
    h::Hamiltonian,
    z::PhasePoint
) where {T<:Real}
    # Create the corresponding static τ
    n_steps = max(1, floor(Int, τ.λ / nom_step_size(τ.integrator)))
    static_τ = StaticTrajectory(τ.integrator, n_steps)
    return transition(rng, static_τ, h, z)
end

###
### Advanced HMC implementation with (adaptive) dynamic trajectory length.
###

##
## Variants of no-U-turn criteria
##

abstract type AbstractTerminationCriterion end

struct ClassicNoUTurn <: AbstractTerminationCriterion end

ClassicNoUTurn(::PhasePoint) = ClassicNoUTurn()

struct GeneralisedNoUTurn{T<:AbstractVector{<:Real}} <: AbstractTerminationCriterion
    rho::T
end

GeneralisedNoUTurn(z::PhasePoint) = GeneralisedNoUTurn(z.r)

combine(cleft::T, cright::T) where {T<:ClassicNoUTurn} = T()
combine(cleft::T, cright::T) where {T<:GeneralisedNoUTurn} = T(cleft.rho + cright.rho)

##
## NUTS
##

"""
Dynamic trajectory HMC using the no-U-turn termination criteria algorithm.
"""
struct NUTS{
    S<:AbstractTrajectorySampler,
    C<:AbstractTerminationCriterion,
    I<:AbstractIntegrator,
    F<:AbstractFloat
} <: DynamicTrajectory{I}
    integrator      ::  I
    max_depth       ::  Int
    Δ_max           ::  F
end

Base.show(io::IO, τ::NUTS{S,C,I,F}) where {I,F,S<:SliceTS,C<:ClassicNoUTurn} =
    print(io, "NUTS{SliceTS}(integrator=$(τ.integrator), max_depth=$(τ.max_depth)), Δ_max=$(τ.Δ_max))")
Base.show(io::IO, τ::NUTS{S,C,I,F}) where {I,F,S<:SliceTS,C<:GeneralisedNoUTurn} =
    print(io, "NUTS{SliceTS,Generalised}(integrator=$(τ.integrator), max_depth=$(τ.max_depth)), Δ_max=$(τ.Δ_max))")
Base.show(io::IO, τ::NUTS{S,C,I,F}) where {I,F,S<:MultinomialTS,C<:ClassicNoUTurn} =
    print(io, "NUTS{MultinomialTS}(integrator=$(τ.integrator), max_depth=$(τ.max_depth)), Δ_max=$(τ.Δ_max))")
Base.show(io::IO, τ::NUTS{S,C,I,F}) where {I,F,S<:MultinomialTS,C<:GeneralisedNoUTurn} =
    print(io, "NUTS{MultinomialTS,Generalised}(integrator=$(τ.integrator), max_depth=$(τ.max_depth)), Δ_max=$(τ.Δ_max))")

const NUTS_DOCSTR = """
    NUTS{S,C}(
        integrator::I,
        max_depth::Int=10,
        Δ_max::F=1000.0
    ) where {I<:AbstractIntegrator,F<:AbstractFloat,S<:AbstractTrajectorySampler,C<:AbstractTerminationCriterion}

Create an instance for the No-U-Turn sampling algorithm.
"""

"$NUTS_DOCSTR"
function NUTS{S,C}(
    integrator::I,
    max_depth::Int=10,
    Δ_max::F=1000.0
) where {I<:AbstractIntegrator,F<:AbstractFloat,S<:AbstractTrajectorySampler,C<:AbstractTerminationCriterion}
    return NUTS{S,C,I,F}(integrator, max_depth, Δ_max)
end

"""
    NUTS(args...) = NUTS{MultinomialTS,GeneralisedNoUTurn}(args...)

Create an instance for the No-U-Turn sampling algorithm
with multinomial sampling and original no U-turn criterion.

Below is the doc for NUTS{S,C}.

$NUTS_DOCSTR
"""
NUTS(args...) = NUTS{MultinomialTS,GeneralisedNoUTurn}(args...)

###
### The doubling tree algorithm for expanding trajectory.
###

"""
Termination reasons
- `dynamic`: due to stoping criteria
- `numerical`: due to large energy deviation from starting (possibly numerical errors)
"""
struct Termination
    dynamic::Bool
    numerical::Bool
end

Base.show(io::IO, d::Termination) = print(io, "Termination(dynamic=$(d.dynamic), numerical=$(d.numerical))")
Base.:*(d1::Termination, d2::Termination) = Termination(d1.dynamic || d2.dynamic, d1.numerical || d2.numerical)
isterminated(d::Termination) = d.dynamic || d.numerical

"""
Check termination of a Hamiltonian trajectory.
"""
Termination(
    s::SliceTS,
    nt::NUTS,
    H0::F,
    H′::F
) where {F<:AbstractFloat} = Termination(false, !(s.ℓu < nt.Δ_max + -H′))

"""
Check termination of a Hamiltonian trajectory.
"""
Termination(
    s::MultinomialTS,
    nt::NUTS,
    H0::F,
    H′::F
) where {F<:AbstractFloat} = Termination(false, !(-H0 < nt.Δ_max + -H′))

"""
A full binary tree trajectory with only necessary leaves and information stored.
"""
struct BinaryTree{C<:AbstractTerminationCriterion}
    zleft   # left most leaf node
    zright  # right most leaf node
    c::C    # termination criterion
    sum_α   # MH stats, i.e. sum of MH accept prob for all leapfrog steps
    nα      # total # of leap frog steps, i.e. phase points in a trajectory
    ΔH_max  # energy in tree with largest absolute different from initial energy
end

"""
    maxabs(a, b)

Return the value with the largest absolute value.
"""
@inline maxabs(a, b) = abs(a) > abs(b) ? a : b

"""
    combine(treeleft::BinaryTree, treeright::BinaryTree)

Merge a left tree `treeleft` and a right tree `treeright` under given Hamiltonian `h`,
then draw a new candidate sample and update related statistics for the resulting tree.
"""
combine(treeleft::BinaryTree, treeright::BinaryTree) =
    BinaryTree(treeleft.zleft, treeright.zright, combine(treeleft.c, treeright.c), treeleft.sum_α + treeright.sum_α, treeleft.nα + treeright.nα, maxabs(treeleft.ΔH_max, treeright.ΔH_max))

"""
Detect U turn for two phase points (`zleft` and `zright`) under given Hamiltonian `h`
using the (original) no-U-turn cirterion.

Ref: https://arxiv.org/abs/1111.4246, https://arxiv.org/abs/1701.02434
"""
function isterminated(h::Hamiltonian, t::BinaryTree{C}) where {C<:ClassicNoUTurn}
    # z0 is starting point and z1 is ending point
    z0, z1 = t.zleft, t.zright
    Δθ = z1.θ - z0.θ
    s = (dot(Δθ, ∂H∂r(h, -z0.r)) >= 0) || (dot(-Δθ, ∂H∂r(h, z1.r)) >= 0)
    return Termination(s, false)
end

"""
Detect U turn for two phase points (`zleft` and `zright`) under given Hamiltonian `h`
using the generalised no-U-turn criterion.

Ref: https://arxiv.org/abs/1701.02434
"""
function isterminated(h::Hamiltonian, t::BinaryTree{C}) where {C<:GeneralisedNoUTurn}
    # z0 is starting point and z1 is ending point
    z0, z1 = t.zleft, t.zright
    rho = t.c.rho
    s = (dot(rho, ∂H∂r(h, -z0.r)) >= 0) || (dot(-rho, ∂H∂r(h, z1.r)) >= 0)
    return Termination(s, false)
end

"""
Recursivly build a tree for a given depth `j`.
"""
function build_tree(
    rng::AbstractRNG,
    nt::NUTS{S,C,I,F},
    h::Hamiltonian,
    z::PhasePoint,
    sampler::AbstractTrajectorySampler,
    v::Int,
    j::Int,
    H0::AbstractFloat
) where {I<:AbstractIntegrator,F<:AbstractFloat,S<:AbstractTrajectorySampler,C<:AbstractTerminationCriterion}
    if j == 0
        # Base case - take one leapfrog step in the direction v.
        z′ = step(rng, nt.integrator, h, z, v)
        H′ = energy(z′)
        ΔH = H′ - H0
        α′ = exp(min(0, -ΔH))
        sampler′ = S(sampler, H0, z′)
        return BinaryTree(z′, z′, C(z′), α′, 1, ΔH), sampler′, Termination(sampler′, nt, H0, H′)
    else
        # Recursion - build the left and right subtrees.
        tree′, sampler′, termination′ = build_tree(rng, nt, h, z, sampler, v, j - 1, H0)
        # Expand tree if not terminated
        if !isterminated(termination′)
            # Expand left
            if v == -1
                tree′′, sampler′′, termination′′ = build_tree(rng, nt, h, tree′.zleft, sampler, v, j - 1, H0) # left tree
                treeleft, treeright = tree′′, tree′
            # Expand right
            else
                tree′′, sampler′′, termination′′ = build_tree(rng, nt, h, tree′.zright, sampler, v, j - 1, H0) # right tree
                treeleft, treeright = tree′, tree′′
            end
            tree′ = combine(treeleft, treeright)
            sampler′ = combine(rng, sampler′, sampler′′)
            termination′ = termination′ * termination′′ * isterminated(h, tree′)
        end
        return tree′, sampler′, termination′
    end
end

function transition(
    rng::AbstractRNG,
    τ::NUTS{S,C,I,F},
    h::Hamiltonian,
    z0::PhasePoint
) where {I<:AbstractIntegrator,F<:AbstractFloat,S<:AbstractTrajectorySampler,C<:AbstractTerminationCriterion}
    H0 = energy(z0)
    tree = BinaryTree(z0, z0, C(z0), zero(F), zero(Int), zero(H0))
    sampler = S(rng, z0)
    termination = Termination(false, false)
    zcand = z0

    integrator = jitter(rng, τ.integrator)
    τ = reconstruct(τ, integrator=integrator)

    j = 0
    while !isterminated(termination) && j < τ.max_depth
        # Sample a direction; `-1` means left and `1` means right
        v = rand(rng, [-1, 1])
        if v == -1
            # Create a tree with depth `j` on the left
            tree′, sampler′, termination′ = build_tree(rng, τ, h, tree.zleft, sampler, v, j, H0)
            treeleft, treeright = tree′, tree
        else
            # Create a tree with depth `j` on the right
            tree′, sampler′, termination′ = build_tree(rng, τ, h, tree.zright, sampler, v, j, H0)
            treeleft, treeright = tree, tree′
        end
        # Perform a MH step and increse depth if not terminated
        if !isterminated(termination′)
            j = j + 1   # increment tree depth
            if mh_accept(rng, sampler, sampler′)
                zcand = sampler′.zcand
            end
        end
        # Combine the proposed tree and the current tree (no matter terminated or not)
        tree = combine(treeleft, treeright)
        # Update sampler
        sampler = combine(zcand, sampler, sampler′)
        # update termination
        termination = termination * termination′ * isterminated(h, tree)
    end

    H = energy(zcand)
    tstat = merge(
        (
         n_steps=tree.nα,
         is_accept=true,
         acceptance_rate=tree.sum_α / tree.nα,
         log_density=zcand.ℓπ.value,
         hamiltonian_energy=H,
         hamiltonian_energy_error=H - H0,
         max_hamiltonian_energy_error=tree.ΔH_max,
         tree_depth=j,
         numerical_error=termination.numerical,
        ),
        stat(τ.integrator)
    )

    return Transition(zcand, tstat)
end

"""
A single Hamiltonian integration step.

NOTE: this function is intended to be used in `find_good_eps` only.
"""
function A(rng, h, z, ϵ)
    z′ = step(rng, Leapfrog(ϵ), h, z)
    H′ = energy(z′)
    return z′, H′
end

"""
Find a good initial leap-frog step-size via heuristic search.
"""
function find_good_eps(
    rng::AbstractRNG,
    h::Hamiltonian,
    θ::AbstractVector{T};
    max_n_iters::Int=100
) where {T<:Real}
    # Initialize searching parameters
    ϵ′ = ϵ = T(0.1)
    a_min, a_cross, a_max = T(0.25), T(0.5), T(0.75) # minimal, crossing, maximal accept ratio
    d = T(2.0)
    # Create starting phase point
    r = rand(rng, h.metric) # sample momentum variable
    z = phasepoint(h, θ, r)
    H = energy(z)

    # Make a proposal phase point to decide direction
    z′, H′ = A(rng, h, z, ϵ)
    ΔH = H - H′ # compute the energy difference; `exp(ΔH)` is the MH accept ratio
    direction = ΔH > log(a_cross) ? 1 : -1

    # Crossing step: increase/decrease ϵ until accept ratio cross a_cross.
    for _ = 1:max_n_iters
        # `direction` being  `1` means MH ratio too high
        #     - this means our step size is too small, thus we increase
        # `direction` being `-1` means MH ratio too small
        #     - this means our step szie is too large, thus we decrease
        ϵ′ = direction == 1 ? d * ϵ : 1 / d * ϵ
        z′, H′ = A(rng, h, z, ϵ)
        ΔH = H - H′
        DEBUG && @debug "Crossing step" direction H′ ϵ "α = $(min(1, exp(ΔH)))"
        if (direction == 1) && !(ΔH > log(a_cross))
            break
        elseif (direction == -1) && !(ΔH < log(a_cross))
            break
        else
            ϵ = ϵ′
        end
    end
    # Note after the for loop,
    # `ϵ` and `ϵ′` are the two neighbour step sizes across `a_cross`.

    # Bisection step: ensure final accept ratio: a_min < a < a_max.
    # See https://en.wikipedia.org/wiki/Bisection_method

    ϵ, ϵ′ = ϵ < ϵ′ ? (ϵ, ϵ′) : (ϵ′, ϵ)  # ensure ϵ < ϵ′;
    # Here we want to use a value between these two given the
    # criteria that this value also gives us a MH ratio between `a_min` and `a_max`.
    # This condition is quite mild and only intended to avoid cases where
    # the middle value of `ϵ` and `ϵ′` is too extreme.
    for _ = 1:max_n_iters
        ϵ_mid = middle(ϵ, ϵ′)
        z′, H′ = A(rng, h, z, ϵ_mid)
        ΔH = H - H′
        DEBUG && @debug "Bisection step" H′ ϵ_mid "α = $(min(1, exp(ΔH)))"
        if (exp(ΔH) > a_max)
            ϵ = ϵ_mid
        elseif (exp(ΔH) < a_min)
            ϵ′ = ϵ_mid
        else
            ϵ = ϵ_mid
            break
        end
    end

    return ϵ
end

find_good_eps(
    h::Hamiltonian,
    θ::AbstractVector{T};
    max_n_iters::Int=100
) where {T<:Real} = find_good_eps(GLOBAL_RNG, h, θ; max_n_iters=max_n_iters)

"""
Perform MH acceptance based on energy, i.e. negative log probability.
"""
function mh_accept_ratio(
    rng::AbstractRNG,
    Horiginal::T,
    Hproposal::T
) where {T<:Real}
    α = min(1.0, exp(Horiginal - Hproposal))
    accept = rand(rng) < α
    return accept, α
end

####
#### Adaption
####

function update(
    h::Hamiltonian,
    τ::AbstractProposal,
    pc::Adaptation.AbstractPreconditioner
)
    metric = renew(h.metric, getM⁻¹(pc))
    h = reconstruct(h, metric=metric)
    return h, τ
end

function update(
    h::Hamiltonian,
    τ::AbstractProposal,
    da::NesterovDualAveraging
)
    integrator = reconstruct(τ.integrator, ϵ=getϵ(da))
    τ = reconstruct(τ, integrator=integrator)
    return h, τ
end

function update(
    h::Hamiltonian,
    τ::AbstractProposal,
    ca::Union{Adaptation.NaiveHMCAdaptor, Adaptation.StanHMCAdaptor}
)
    metric = renew(h.metric, getM⁻¹(ca.pc))
    h = reconstruct(h, metric=metric)
    integrator = reconstruct(τ.integrator, ϵ=getϵ(ca.ssa))
    τ = reconstruct(τ, integrator=integrator)
    return h, τ
end

function update(
    h::Hamiltonian,
    θ::AbstractVector{T}
) where {T<:Real}
    metric = h.metric
    if length(metric) != length(θ)
        metric = typeof(metric)(length(θ))
        h = reconstruct(h, metric=metric)
    end
    return h
end


text:
#=
 Algorithms for triple-word arithmetic
 Nicolas Fabiano, Jean-Michel Muller, Joris Picot
=#

function vec_sum(x0::T, x1::T, x2::T) where {T}
    s1, e2 = two_sum(x1, x2)
    s0, e1 = two_sum(x0, s1)
    return s0,e1,e2
end

function vecsum(x0::T,x1::T,x2::T) where {T}
   x0 === 0.0 && return (two_sum(x1, x2)..., 0.0)
   s1, e2 = two_sum(x1, x2)
   e0, e1 = two_sum(x0, s1)
   return (e0, e1, e2)
end

function vseb(a0::T, a1::T, a2::T, a3::T, a4::T, a5::T) where {T}
    a2 === 0.0 && return (a0, a1, a2)
	r0, e1 = two_sum(a0, a1)
	if e1 != 0.0
	  y0 = r0
	else
	  e1 = r0
	end
	y1, y2  = two_sum(e1, a2)
	return (y0, y1, y2)
end

function vseb(a0::T, a1::T, a2::T) where {T}
    a2 === 0.0 && return (a0, a1, a2)
	r0, e1 = two_sum(a0, a1)
	if e1 != 0.0
	  y0 = r0
	else
	  e1 = r0
	end
	y1, y2  = two_sum(e1, a2)
	return (y0, y1, y2)
end

function vecseb(x0::T, x1::T, x2::T) where {T}
   a0, a1, a2 = vecsum(x0, x1, x2)
   return vseb(a0, a1, a2)
end

function tripleword(a::T, b::T, c::T) where {T}
    d0, d1 = two_sum(a,b)
    # e0, e1, e2 = vecsum(d0, d1, c)
    s1, e2 = two_sum(d1, c)
    e0, e1 = two_sum(d0, s1)
    # s0, s1, s2 = vseb(e0, e1, e2)
    r0, e1 = two_sum(e0, e1)
    if e1 != 0.0
	s0 = r0
    else
	e1 = r0
    end
    s1, s2  = two_sum(e1, e2)
    return (s0, s1, s2)
end

using SetRounding

function Float64(x0::Float64, x1::Float64, x2::Float64)
    s, e = two_hilo_sum(x0, x1+x1)
    if e !== 0.0
        y = x0 + x1
    elseif x2 > 0.0
        setrounding(Float64, RoundUp)
        y = x0 + x1
        setrounding(Float64, RoundNearest)
    elseif x2 < 0.0
        setrounding(Float64, RoundDown)
        y = x0 + x1
        setrounding(Float64, RoundNearest)
    else
        y = x0 + x1
    end
    return y
end

@inline max_min(x,y) = ifelse(abs(y) < abs(x), (x,y) , (y,x))

function maxtomin(a::T, b::T, c::T) where {T}
    b, c = max_min(b, c)
    a, c = max_min(a, c)
    a, b = max_min(a, b)
    return a, b, c
end

function vec_sum(x1::T, x2::T, x3::T) where {T}
    s3 = x3
    s2, e3 = two_sum(x2, s3)
    s1, e2 = two_sum(x1, s2)
    return s1,e2,e3
end

function vecsum_errbranch(x::NTuple{3,T}) where {T}
    y = r = e = zeros(T, 3)
    j = 1
    e[1] = x[1]
    for i = 1:1
        r[i], t = two_sum(e[i], x[i+1])
        if t !== zero(T)
            y[j] = r[i]
            e[i+1] = t
            j += 1
        else
            e[i+1] = r[i]
        end
    end
    y[j], y[j+1] = two_sum(e[3], x[4])
    return y
end

function fast_vecsum_errbranch(x1::T,x2::T,x3::T) where {T}
    y = zeros(T, 3)
    j = 1
    # e[1] = x1
    # i = 1
    r, t = two_sum(x1, x2)
    if t !== zero(T)
       y[j] = r
       e = t
       j += 1
    else
       e = r
    end

    y[j], y[j+1] = two_sum(e, x3)
    return y
end

function tripleword(x1::T, x2::T, x3::T) where {T}
    a1, a2 = two_sum(x1, x2)
    c1, c2 = two_sum(a1, x3)
    e1to3 = vec_sum(c1,c2,a2)
    y = vecsum_errbranch(e1to3)
    return (y...,)
end

@inline function fast_tripleword(x1::T, x2::T, x3::T) where {T}
    a,b,c,d = maxtomin(x1,x2,x3)
    return fast_vecsum_errbranch(a,b,c)
end

ext: jl


text:
################################################################################
#
#  Abstract types
#
################################################################################

# abstract spaces
abstract type AbsSpace{S} end

# abstract lattices
abstract type AbsLat{S} end

################################################################################
#
#  Quadratic spaces
#
################################################################################

mutable struct DictWrapper{T}
  x::T
end

function Base.:(==)(x::DictWrapper{T}, y::DictWrapper{S}) where {S, T}
  return S === T && x.x == y.x
end

function Base.hash(x::DictWrapper, h::UInt)
  Base.hash(x.x, h)
end

const QuadSpaceID = AbstractAlgebra.CacheDictType{Any, Any}()

@attributes mutable struct QuadSpace{S, T} <: AbsSpace{S}
  K::S
  gram::T

  function QuadSpace(K::S, G::T, cached::Bool) where {S, T}
    return AbstractAlgebra.get_cached!(QuadSpaceID, DictWrapper(G), cached) do
      new{S, T}(K, G)
    end::QuadSpace{S, T}
  end
end

################################################################################
#
#  Hermitian spaces
#
################################################################################

const HermSpaceID = AbstractAlgebra.CacheDictType{Any, Any}()

@attributes mutable struct HermSpace{S, T, U, W} <: AbsSpace{S}
  E::S
  K::T
  gram::U
  involution::W

  function HermSpace(E::S, K::T, gram::U, involution::W, cached::Bool) where {S, T, U, W}
    return AbstractAlgebra.get_cached!(HermSpaceID, DictWrapper(gram), cached) do
      new{S, T, U, W}(E, K, gram, involution)
    end::HermSpace{S, T, U, W}
  end
end


text:
# Skip the DiffEqBase handling
function solve(prob::OptimizationProblem, opt, args...;kwargs...)
    __solve(prob, opt, args...; kwargs...)
end


text:
module MultivariateTriangle

using ApproxFun, Makie, MultivariateOrthogonalPolynomials
import AbstractPlotting: mesh, mesh!
import ApproxFunOrthogonalPolynomials: _padua_length

export contourf, contourf

contourf(f::Fun; kwds...) = _mesh(meshdata(f)...; shading=false, kwds...)
contourf!(s, f::Fun; kwds...) = _mesh!(s,  meshdata(f)...; shading=false, kwds...)

function _mesh(p, T, v; resolution=(400,400), kwds...)
    T_mat = Array{Int}(undef, length(T), 3)
    for k = 1:length(T)
        T_mat[k,:] .= T[k]
    end
    s = Scene(resolution=resolution)
    mesh!(s, [first.(p) last.(p)], T_mat; color=v, kwds...)
end

function _surface(p, T, v; resolution=(400,400), kwds...)
    T_mat = Array{Int}(undef, length(T), 3)
    for k = 1:length(T)
        T_mat[k,:] .= T[k]
    end
    # s = Scene(resolution=resolution)
    mesh(first.(p), last.(p), vec(v), T_mat; kwds...)
end

function _mesh!(s, p, T, v; kwds...)
    T_mat = Array{Int}(undef, length(T), 3)
    for k = 1:length(T)
        T_mat[k,:] .= T[k]
    end
    mesh!(s, [first.(p) last.(p)], T_mat; color=v, kwds...)
end

function meshdata(f::Fun{<:PiecewiseSpace})
    pTv = MultivariateTriangle.meshdata.(components(f))
    p = vcat(first.(pTv)...)
    T = pTv[1][2]
    cs = length(pTv[1][1])
    for k = 2:length(pTv)
        append!(T, (t -> (cs.+t)).(pTv[k][2]))
        cs += length(pTv[k][1])
    end

    v = vcat(last.(pTv)...)

    p, T, v
end

function meshdata(f::Fun{<:TensorSpace{<:Tuple{<:Chebyshev,<:Chebyshev}}})
    p = points(f)
    v = values(f)
    n = length(p)
    T = Vector{NTuple{3,Int}}()
    d_x,d_y = factors(domain(f))
    a_x,b_x = endpoints(d_x)
    a_y,b_y = endpoints(d_y)
    if iseven(_padua_length(n))
        l = floor(Int, (1+sqrt(1+8n))/4)

        push!(p, Vec(b_x,b_y))
        push!(p, Vec(a_x,b_y))

        push!(v, f(b_x,b_y))
        push!(v, f(a_x,b_y))

        for p = 0:l-2
            for k = (2p*l)+1:(2p*l)+l-1
                push!(T, (k+1, k, l+k+1))
            end
            for k = (2p*l)+1:(2p*l)+l-1
                push!(T, (k, l+k, l+k+1))
            end
            for k = (2p*l)+l+1:(2p*l)+2l-1
                push!(T, (k+1, k, l+k))
            end
            for k = (2p*l)+l+2:(2p*l)+2l
                push!(T, (k, k+l-1, l+k))
            end
        end
        for p=0:l-3
            push!(T, ((2p+1)*l+1, (2p+2)*l+1, (2p+3)*l+1))
        end
        for p =0:l-2
            push!(T, ((2p+1)*l, (2p+2)*l, (2p+3)*l))
        end
        push!(T, (1, n+1, l+1))
        push!(T, (n-2l+1, n+2, n-l+1))
    else
        l = floor(Int, (3+sqrt(1+8n))/4)

        push!(p, Vec(a_x,b_y))
        push!(p, Vec(a_x,a_y))

        push!(v, f(a_x,b_y))
        push!(v, f(a_x,a_y))

        for p = 0:l-2
            for k = p*(2l-1)+1:p*(2l-1)+l-1
                push!(T, (k+1, k, l+k))
            end
            for k = p*(2l-1)+1:p*(2l-1)+l-2
                push!(T, (k+1, l+k, l+k+1))
            end
        end
        for p = 0:l-3
            for k = p*(2l-1)+l+1:p*(2l-1)+2l-2
                push!(T, (k+1, k, l+k))
            end
            for k = p*(2l-1)+l+1:p*(2l-1)+2l-1
                push!(T, (k, k+l-1, l+k))
            end
        end

        for p=0:l-3
            push!(T, (p*(2l-1) + 1, p*(2l-1) + l+1, p*(2l-1) + 2l))
        end

        for p=0:l-3
            push!(T, (p*(2l-1) + l, p*(2l-1) + 2l-1, p*(2l-1) + 3l-1))
        end

        push!(T, (n-2l+2, n+1, n-l+2))
        push!(T, (n-l+1, n+2, n))
    end

    p, T, v
end
meshdata(f::Fun{<:Space{<:Triangle}}) =
    triangle_meshdata(points(f), values(f), (domain(f).a, domain(f).b, domain(f).c),
                                            f.((domain(f).a, domain(f).b, domain(f).c)))

function triangle_meshdata(p, v, (a, b, c), (fa, fb, fc))
    n = length(p)
    T = Vector{NTuple{3,Int}}()

    if iseven(_padua_length(n))
        l = floor(Int, (1+sqrt(1+8n))/4)

        push!(p, b)
        push!(p, c)

        push!(v, fb)
        push!(v, fc)

        for p = 0:l-2
            for k = (2p*l)+1:(2p*l)+l-1
                push!(T, (k+1, k, l+k+1))
            end
            for k = (2p*l)+1:(2p*l)+l-1
                push!(T, (k, l+k, l+k+1))
            end
            for k = (2p*l)+l+1:(2p*l)+2l-1
                push!(T, (k+1, k, l+k))
            end
            for k = (2p*l)+l+2:(2p*l)+2l
                push!(T, (k, k+l-1, l+k))
            end
        end
        for p=0:l-3
            push!(T, ((2p+1)*l+1, (2p+2)*l+1, (2p+3)*l+1))
        end
        for p =0:l-2
            push!(T, ((2p+1)*l, (2p+2)*l, (2p+3)*l))
        end
        push!(T, (1, n+1, l+1))
        push!(T, (n-2l+1, n+2, n-l+1))
    else
        l = floor(Int, (3+sqrt(1+8n))/4)

        push!(p, Vec(c))
        push!(p, Vec(a))

        push!(v, fc)
        push!(v, fa)

        for p = 0:l-2
            for k = p*(2l-1)+1:p*(2l-1)+l-1
                push!(T, (k+1, k, l+k))
            end
            for k = p*(2l-1)+1:p*(2l-1)+l-2
                push!(T, (k+1, l+k, l+k+1))
            end
        end
        for p = 0:l-3
            for k = p*(2l-1)+l+1:p*(2l-1)+2l-2
                push!(T, (k+1, k, l+k))
            end
            for k = p*(2l-1)+l+1:p*(2l-1)+2l-1
                push!(T, (k, k+l-1, l+k))
            end
        end

        for p=0:l-3
            push!(T, (p*(2l-1) + 1, p*(2l-1) + l+1, p*(2l-1) + 2l))
        end

        for p=0:l-3
            push!(T, (p*(2l-1) + l, p*(2l-1) + 2l-1, p*(2l-1) + 3l-1))
        end

        push!(T, (n-2l+2, n+1, n-l+2))
        push!(T, (n-l+1, n+2, n))
    end

    p, T, v
end

end # module


text:
using LinearAlgebra
using Distributions
using SpecialFunctions
using PDMats
using CSV
include("kernel.jl")
include("integration.jl")
include("btgDerivatives.jl")
include("examples.jl")

if false
    (main, dmain, d2main) = partial_theta(10, 2, getExample(1, 1000, 3, 3, 10))

    f = x -> main(x)[1]
    df = x -> dmain(x)[1]
    d2f = x -> d2main(x)[1]

    @time begin
    println(main([.4]))
    end
    @time begin
    println(dmain([.4]))
    end
    @time begin
    println(d2main([.4]))
    end

    #x = .01:.01:2
    #y = f.(x)
    #dy = df.(x)
    #d2y = d2f.(x)
    #plt1 = plot(x, y, title = "f")
    #plt2 = plot(x, dy, title = "df")
    #plt3 = plot(x, d2y, title = "d2f")
    #display(plot(plt1, plt2, plt3))
end
if true
    h = 50:50:500
    times = zeros(3, length(h))
    for i = h
        s1 = rand(10*i, 3)
        Z = K(s1, s1, 10)
        times[1, Int(i/50)] = @elapsed cz = cholesky(Z)
        times[2, Int(i/50)] = @elapsed Z\rand(size(Z, 1), 1)
        times[3, Int(i/50)] = @elapsed cz\rand(size(Z, 1), 1)
    end
    plot(repeat(h, 1,3), times')
end

max_stars_repo_name: dbindel/btg

max_issues_repo_name: dbindel/btg

max_forks_repo_name: dbindel/btg


text:
# Profile the running times of kernels.
#
# Steps that need to be performed
#
# 1. Get the ops in the executable. Clone each of the ops into their own nGraph
#   function.
#
# 2. For each of the ops, we also have to capture the input and output tensor layouts
#   for the op in order to get accurate timing.
#
#   As far as I can tell, if a CPU op is annotated as "use_mkldnn_kernel", then the
#   input and output layout will always be the same. Thus, to do the correct
#   input/output layout conversion, we just have to check if the node we are copying is
#   annotated as mkldnn and make sure we annotate the new node as such.

disable_passes() = ENV["NGRAPH_PASS_ENABLES"] = join((
    "AlgebraicSimplification:0",
    "CoreFusion:0",
    "CPUFusion:0",
    "CPUHorizontalFusion:0",
    "CommonSubexpressionElimination:0",
    "ReshapeElimination:0",
   ), ";"
)

enable_passes() = delete!(ENV, "NGRAPH_PASS_ENABLES")

"""
    profile(args...)

Profile all of the operations in `fex`.

Keyword Arguments
-----------------
* `cache`: A cache to serve running times if a kernel has already been profiled. The cache
    must implement the function `save`.
"""
function profile(fex::nGraph.FluxExecutable; kw...)
    return profile(fex.ex.ngraph_function, fex.ex.backend; kw...)
end

function profile(
        f::nGraph.NFunction,
        backend::nGraph.Backend{T};
        cache = nothing,
        # Force re-profiling
        recache = false,
        show_progress = true,
        cb = nothing
    ) where {T}

    isnothing(cache) && error("Need to define a cache")

    # Create a "FunctionData" object for this function, which allows us to record the graph
    # structures and metadata on the Julia level.
    data = FunctionData(f, T)

    # Do any post-processing on the `data` object.
    isnothing(cb) || cb(data)

    # Get all the configurations we are interested in for this run.
    # Need to make a MOVE node in order to control IO configurations.
    all_configs = possible_configs(backend, data)

    # Convert the configs to a dictionary mapping node name to configs for easier
    # management
    config_dict = Dict{XNode, Vector{IOConfig}}()
    for config in all_configs
        v = get!(config_dict, first(config), IOConfig[])
        push!(v, last(config))
    end

    # Clean up cached configs if passed the `recache` option
    if recache
        @info "Removing Cached Configs"
        for node in nodes(data)
            hasprofile(node) || continue
            configs = config_dict[node]
            kernel_params = params(backend, node)
            for config in configs
                key = (kernel_params, config)
                delete!(cache, key)
            end
        end
    end

    handle_algo_selection!(cache, backend, data)

    # Determine the number of unique configs
    all_configs = map(filter(hasprofile, nodes(data))) do node
        return [(params(backend, node), config) for config in config_dict[node]]
    end |> Iterators.flatten
    println("Total Configs: $(length(collect(all_configs)))")
    num_configs = length(unique(all_configs))

    #num_configs = sum(length(config_dict[node]) for node in nodes(data) if hasprofile(node))
    progress_bar = Progress(num_configs, 0.2)

    # Setup a little update function for all configurations
    # This gives a more fine-grained information than updating for each op
    serviced = Ref(0)
    function _update!(p, op, config, ncached)
        if show_progress
            serviced[] += 1
            ProgressMeter.next!(
                p;
                valuecolor = :white,
                showvalues = [
                    (:iter, serviced[]),
                    (:total, num_configs),
                    (:op, nGraph.name(op)),
                    (:config, config),
                    (:ncached, ncached),
                ]
            )
        end
    end

    # Disable some passes that get rid of the node we actually want
    disable_passes()

    ncached = 0
    for (index, node) in enumerate(nodes(data))
        # Skip unneeded ops
        hasprofile(node) || continue

        # Get the configs to run for this node
        configs = config_dict[node]

        # Before we build a sub-function, get all of the cached ops.
        cached_configs = IOConfig[]
        kernel_params = params(backend, node)
        for config in configs
            key = (kernel_params, config)
            if haskey(cache, key)
                # Update the number of timings serviced from cached ops
                ncached += 1
                #_update!(progress_bar, node, config, ncached)

                settime!(node, config, cache[key])
                push!(cached_configs, config)
            end
        end

        # Abort if everything is cached
        length(cached_configs) == length(configs) && continue

        # Extract a subgraph with just this op
        for config in Iterators.filter(!in(cached_configs), configs)
            _update!(progress_bar, node, config, ncached)

            # Drop the actual profiling into an inner loop to help the GC run inbetween
            # kernels to keep memory consumption down.
            GC.gc()
            _profile!(node, backend, config)

            # Record the saved time into the cache and save the cache to persist the data
            # across transient or intentional (i.e. ctrl+c) failures.
            cache[(kernel_params, config)] = gettime(node, config)
            save(cache)
        end
    end

    # Renable optimization passes
    enable_passes()

    return data
end

@noinline function _profile!(node, backend, config)
    # Extract this particular node with the given configuration
    # Place GC barriers around the instantiation of the nGraph.Executable to try
    # to reduce interference with other extractions.
    ex, inputs, outputs, copied_ops = extract(
        node,
        backend,
        config;
    )

    # Run the function several times. The first time is a warm up.
    # Timing for the second runs should be pretty consistent.
    GC.@preserve inputs outputs begin
        # Due to limitations with the current nGraph interface, we have to cast the
        # input and output tensor arrays to Vector{Any} where the contents are
        # pointers to the runtime Tensors
        input_ptrs = convert(Vector{Any}, nGraph.getpointer.(inputs))
        output_ptrs = convert(Vector{Any}, nGraph.getpointer.(outputs))

        for _ in 1:3
            ex(input_ptrs, output_ptrs)
        end
    end

    # Save the run time into the node
    record_time!(node, ex, copied_ops, config)
    return nothing
end

#####
##### Algorithm selection
#####

# Right now, no algorithm selection happens in the CPU case - so just special case this
# function to save some CPU cycles
handle_algo_selection!(cache, backend::nGraph.Backend{nGraph.CPU}, data::FunctionData) = nothing
function handle_algo_selection!(cache, backend::nGraph.Backend{nGraph.GPU}, data::FunctionData)
    for (index, node) in enumerate(nodes(data))
        hasprofile(node) || continue

        # Get the lookup key for this node
        kernel_params = params(backend, node)

        # If we can select an algorithm for this node, use the built-in timings
        # and skip the actual profiling.
        if nGraph.Lib.can_select_algo(nGraph.getpointer(unx(node)))
            # Hack for the moment - make sure that we only have a single configuration
            # in the GPU case.
            #
            # If we need more flexibility - we'll deal with it in the future
            configs = possible_configs(backend, node)
            @assert length(configs) == 1
            config = first(configs)

            if !haskey(cache, (kernel_params, config))
                # Cleanup
                @info "Getting CUDNN timings for $(nGraph.name(node))"

                enums = UInt32[]
                times = Float32[]
                bytes = UInt64[]
                GC.gc()
                alloc_failed = nGraph.Lib.get_algo_options(
                    nGraph.getpointer(unx(node)),
                    enums,
                    times,
                    bytes
                )

                # cuDNN returns the results of its time in milliseconds.
                #
                # Convert to microseconds here to make it uniform with the rest of
                # the timings.
                algo_list = [
                    CUDNNAlgorithm(e, 1000 * t, b) for (e,t,b) in zip(enums, times, bytes)
                ]

                cache[(kernel_params, config)] = algo_list
                save(cache)
            end
        end
    end
    return nothing
end

#####
##### Record the profiled runtime of a node
#####

function record_time!(
            node::XNode,
            ex::nGraph.Executable,
            ops::Vector{nGraph.Node},
            @nospecialize(config::IOConfig),
        )

    # First, make sure that all of the copied ops given for recording data actually have
    # the same IOConfig as what is expected
    configs = getconfig.(ops)
    inds = findall(isequal(config), configs)

    if isempty(inds)
        @show configs
        @show config
        error()
    end

    # Just pull out the ops that match this configuration
    ops = ops[inds]

    # Get a dictionary mapping node name to runtime in microseconds
    timing_dict = nGraph.get_performance(ex)

    # Get all of the runtimes for the ops being profiled
    times = [timing_dict[nGraph.name(op)] for op in ops]

    # Set the runtime for this node as the average of the runtimes of the individual
    # profiled ops.
    settime!(node, config, mean(times))
    return nothing
end

"""
    extract(node::nGraph.Node, backend::nGraph.Backend{nGraph.CPU}, config)

Create a `nGraph.Executable` with a copy of `node` using the same input and output
parameters. `Move` nodes will be inserted at all inputs and outputs of the copied node
if the inputs come from executable arguments and if the outputs go directly to results.
"""
function extract(
        xnode::XNode,
        backend::nGraph.Backend,
        @nospecialize(config::IOConfig);
    )

    # Manually dispatch some cases.
    if isresult(unx(xnode))
        parameters, output_nodes = _extract_result(xnode)
    else
        parameters, output_nodes = _extract_general(xnode)
    end

    paramvector = nGraph.ParameterVector(parameters)
    nodevector = nGraph.NodeVector(output_nodes)

    # Get the input and output XTensors for this XNode - do a quick sanity check on the
    # inputs and outputs to make sure everything is still okay.
    input_xtensors = filter(!isconstant, inputs(xnode))
    output_xtensors = outputs(xnode)

    for (x,p) in zip(input_xtensors, parameters)
        @assert size(unx(x)) == size(p)
    end
    for (x,o) in zip(output_xtensors, output_nodes)
        @assert size(unx(x)) == size(o)
    end

    # Create a pass callback to setup the function to the given configuration.
    translated_nodes_ref = Ref(nGraph.Node[])
    function configuration_callback(fn)
        # We have to inspect the graph, find the nodes that do not have converted
        # inputs, and insert out synchronous move nodes so we can control the input/output
        # state of the node under test.
        #
        # But first, we have to find what happened to the original node and find it in the
        # new graph.
        translated_nodes = nGraph.Node[]
        node = unx(xnode)
        for op in fn
            # Line it up by description and input/output sizes.
            if params(backend, op) == params(backend, node)
                push!(translated_nodes, op)
            end

            # TODO: I used to allow an option to specify the number of copies.
            # That has been removed, but `translated_nodes`, is still expected to be an
            # iterable of some kind. So, just abort when we find a suitable node.
            #
            # This is mainly needed because `ConvertLayout` has a tendancy to go overboard
            # and segfault.
            length(translated_nodes) == 1 && break
        end

        # Throw an error if we didn't find a node matching what we expected.
        if isempty(translated_nodes)
            for n in translated_nodes
                println("Copied Node: $n")
                println("    name:   $(nGraph.name(n))")
                println("    config: $(config)")
            end
            error("Something done gone wrong!")
        end

        # Finally, configure the translated_nodes to the form that we want.
        #
        # This will handle both internal nGraph tensors as well as top level IO
        for translated_node in translated_nodes
            # Splice in move nodes
            for (index, input) in enumerate(nGraph.get_inputs(translated_node))
                if config[index] == PMEM
                    nGraph.splice(input, 1, translated_node, index, nGraph.move(input))
                end
            end
            for index in 1:nGraph.get_output_size(translated_node)
                for output in nGraph.get_output(translated_node, index)
                    if config.outputs[index] == PMEM
                        nGraph.splice(
                            translated_node,
                            index,
                            output,
                            1,
                            nGraph.move(translated_node, index)
                        )
                    end
                end
            end
            _setup!(translated_node, config)
        end

        # Save the translated nodes to the external translated_nodes_ref
        translated_nodes_ref[] = translated_nodes
        return nothing
    end

    # Setup the above function as a callback
    callbacks = CallbackChain()
    callback!(callbacks, configuration_callback)

    # Compile the function
    ex = nGraph.compile(
        backend,
        paramvector,
        nodevector;
        callback = callbacks,
        emit_timing = true,
    )

    # Make these any to make them compatible with the inner call for nGraph.Executable
    #
    # If we come across any `inplace` annotations, we need to make sure we duplicate the
    # pointer to the underlying runtime tensor
    inplace = Dict{Int, nGraph.TensorView}()

    if length(input_xtensors) == 0
        input_tensors = nGraph.TensorView[]
    else
        input_tensors = map(input_xtensors) do x
            # If this is an inplace node and we've already generated a tensor for it, just
            # return that tensor
            if isinplace(x) && haskey(inplace, x.group)
                return inplace[x.group]
            end

            # If this tensor is persistent - create a PersistentArray to create it from.
            # Otherwise, just use a normal array.
            unxx = unx(x)
            #A = zeros(eltype(unxx), size(unxx))
            A = aligned_alloc(eltype(unxx), size(unxx))
            t = nGraph.TensorView(backend, A)
            isinplace(x) && (inplace[x.group] = t)
            return t
        end
    end
    @assert isa(input_tensors, Vector{nGraph.TensorView})

    output_tensors = map(zip(output_xtensors, output_nodes)) do tup
        x = tup[1]
        out_node = tup[2]

        # Assume we've already seen inplace tensors and simply return.
        # Will error if a match is not found - which is kind of what we want.
        isinplace(x) && return inplace[x.group]
        unxx = unx(x)
        #A = zeros(eltype(unxx), size(unxx))
        A = aligned_alloc(eltype(unxx), size(unxx))
        return nGraph.TensorView(backend, A)
    end
    @assert isa(output_tensors, Vector{nGraph.TensorView})

    return ex, input_tensors, output_tensors, translated_nodes_ref[]
end

function aligned_alloc(::Type{T}, sz; alignment = 64) where {T}
    # Julia guarentees 16 byte alignment for small arrays and 64-byte allignment for
    # larger arrays. However, we want 64-byte alignment for ALL arrays because of the use
    # of AvX instructions in nGraph.
    #
    # Also - make sure to populate the entries with zeros, otherwise data values inside
    # ngraph may go flying around all over the place causing mayhem.
    # While the alignment to the first element is not 64 bytes aligned - try again.
    #
    # If the size if greater than 2KB, allocate it normally.
    # Otherwise, allocate an array bigger than normal and resize it down.
    #
    # For discussion, see: https://discourse.julialang.org/t/aligned-julia-arrays/10993/27
    bytes = sizeof(T) * prod(sz)
    if bytes > 2048
        A = zeros(T, sz)
        return A
    else
        A = zeros(T, div(4096, sizeof(T)))
        resize!(A, prod(sz))
        B = reshape(A, sz)
        return B
    end
end

# We have to extract "result" nodes as well since they will copy if input and output
# tensors live in different memory pools.
#
# We break this up into a general extraction method, and a special case for "results" since
# that case is much simpler
function _extract_general(xnode)
    node = nGraph.Node(unx(xnode))
    # Create parameters for the inputs
    parameters = nGraph.Node[]
    for i in 1:nGraph.get_input_size(node)
        # Check if this input is a constant. If so - copy over the constant.
        this_input = nGraph.get_input(node, i)
        if isconstant(this_input)
            push!(parameters, nGraph.copy(this_input, nGraph.NodeVector()))
        else
            P = nGraph.parameter(
                nGraph.get_input_element_type(node, i),
                nGraph.get_input_shape(node, i),
            )
            push!(parameters, P)
        end
    end

    # Insert layout conversion to match the mkldnn layouts in the original graph.
    links = nGraph.Node[]
    for i in 1:nGraph.get_input_size(node)
        if nGraph.input_needs_conversion(node, i)
            push!(links, nGraph.convert_layout_to(parameters[i], node, i))
        else
            push!(links, parameters[i])
        end
    end

    # Copy the node with the newly created parameters
    copied_node = copy(node, links)

    # Make sure we're using the same version of the node - will always return `false` if
    # compiling for the GPU backend.
    nGraph.is_mkldnn(node) && nGraph.set_mkldnn(copied_node)

    # Compile the new function
    #
    # Now that we've created a copy of the new node, we need to filter out all of the input
    # constants we provided so it compiles correctly.
    filter!(!isconstant, parameters)

    outputs = nGraph.Node[]
    if nGraph.get_output_size(copied_node) > 1
        for i in 1:nGraph.get_output_size(copied_node)
            push!(outputs, nGraph.get_output_element(copied_node, i))
        end
    else
        push!(outputs, copied_node)
    end

    return parameters, outputs
end

_extract_result(node::XNode) = _extract_result(nGraph.Node(unx(node)))
function _extract_result(node)
    @assert isresult(node)
    @assert nGraph.get_input_size(node) == 1

    # We insert a unary negative of the input as the result - this ensure that the result
    # generated by nGraph is not coming directly from a parameter - which means inplace
    # assignment of results should happen the intermediate tensor and result are in the
    # same memory pool.
    P = nGraph.parameter(nGraph.get_input_element_type(node, 1), nGraph.get_input_shape(node, 1))
    parameters = [P]
    outputs = [-P]
    return parameters, outputs
end

max_stars_repo_name: darchr/AutoTM

max_issues_repo_name: darchr/AutoTM

max_forks_repo_name: darchr/AutoTM


text:
const AGGR2STR = Dict{Symbol,String}(:add => "∑", :sub => "-∑", :mul => "∏", :div => "1/∏",
                                     :max => "max", :min => "min", :mean => "𝔼[]")

"""
    GCNConv(graph, in=>out)
    GCNConv(graph, in=>out, σ)

Graph convolutional layer.

# Arguments
- `graph`: should be a adjacency matrix, `SimpleGraph`, `SimpleDiGraph` (from LightGraphs) or `SimpleWeightedGraph`, `SimpleWeightedDiGraph` (from SimpleWeightedGraphs).
- `in`: the dimension of input features.
- `out`: the dimension of output features.
- `bias::Bool=true`: keyword argument, whether to learn the additive bias.

Data should be stored in (# features, # nodes) order.
For example, a 1000-node graph each node of which poses 100 feautres is constructed.
The input data would be a `1000×100` array.
"""
struct GCNConv{T,F}
    weight::AbstractMatrix{T}
    bias::AbstractMatrix{T}
    norm::AbstractMatrix{T}
    σ::F
end

function GCNConv(adj::AbstractMatrix, ch::Pair{<:Integer,<:Integer}, σ = identity;
                 init = glorot_uniform, T::DataType=Float32, bias::Bool=true)
    N = size(adj, 1)
    b = bias ? init(ch[2], N) : zeros(T, ch[2], N)
    GCNConv(init(ch[2], ch[1]), b, normalized_laplacian(adj+I, T), σ)
end

@functor GCNConv

(g::GCNConv)(X::AbstractMatrix) = g.σ.(g.weight * X * g.norm + g.bias)

function Base.show(io::IO, l::GCNConv)
    in_channel = size(l.weight, ndims(l.weight))
    out_channel = size(l.weight, ndims(l.weight)-1)
    print(io, "GCNConv(G(V=", size(l.norm, 1))
    print(io, ", E), ", in_channel, "=>", out_channel)
    l.σ == identity || print(io, ", ", l.σ)
    print(io, ")")
end

"""
    ChebConv(graph, in=>out, k)

Chebyshev spectral graph convolutional layer.

# Arguments
- `graph`: should be a adjacency matrix, `SimpleGraph`, `SimpleDiGraph` (from LightGraphs) or `SimpleWeightedGraph`, `SimpleWeightedDiGraph` (from SimpleWeightedGraphs).
- `in`: the dimension of input features.
- `out`: the dimension of output features.
- `k`: the order of Chebyshev polynomial.
- `bias::Bool=true`: keyword argument, whether to learn the additive bias.
"""
struct ChebConv{T}
    weight::AbstractArray{T,3}
    bias::AbstractMatrix{T}
    L̃::AbstractMatrix{T}
    k::Integer
    in_channel::Integer
    out_channel::Integer
end

function ChebConv(adj::AbstractMatrix, ch::Pair{<:Integer,<:Integer}, k::Integer;
                  init = glorot_uniform, T::DataType=Float32, bias::Bool=true)
    N = size(adj, 1)
    b = bias ? init(ch[2], N) : zeros(T, ch[2], N)
    L̃ = T(2. / eigmax(adj)) * normalized_laplacian(adj, T) - I
    ChebConv(init(ch[2], ch[1], k), b, L̃, k, ch[1], ch[2])
end

@functor ChebConv

function (c::ChebConv)(X::AbstractMatrix{T}) where {T<:Real}
    fin = c.in_channel
    @assert size(X, 1) == fin "Input feature size must match input channel size."
    N = size(c.L̃, 1)
    @assert size(X, 2) == N "Input vertex number must match Laplacian matrix size."
    fout = c.out_channel

    Z = similar(X, fin, N, c.k)
    Z[:,:,1] = X
    Z[:,:,2] = X * c.L̃
    for k = 3:c.k
        Z[:,:,k] = 2*view(Z, :, :, k-1)*c.L̃ - view(Z, :, :, k-2)
    end

    Y = view(c.weight, :, :, 1) * view(Z, :, :, 1)
    for k = 2:c.k
        Y += view(c.weight, :, :, k) * view(Z, :, :, k)
    end
    Y += c.bias
    return Y
end

function Base.show(io::IO, l::ChebConv)
    print(io, "ChebConv(G(V=", size(l.L̃, 1))
    print(io, ", E), ", l.in_channel, "=>", l.out_channel)
    print(io, ", k=", l.k)
    print(io, ")")
end

"""
    GraphConv(graph, in=>out)
    GraphConv(graph, in=>out, aggr)

Graph neural network layer.

# Arguments
- `graph`: should be a adjacency matrix, `SimpleGraph`, `SimpleDiGraph` (from LightGraphs) or `SimpleWeightedGraph`, `SimpleWeightedDiGraph` (from SimpleWeightedGraphs).
- `in`: the dimension of input features.
- `out`: the dimension of output features.
- `bias::Bool=true`: keyword argument, whether to learn the additive bias.
- `aggr::Symbol=:add`: an aggregate function applied to the result of message function. `:add`, `:max` and `:mean` are available.
"""
struct GraphConv{V,T} <: MessagePassing
    adjlist::V
    weight1::AbstractMatrix{T}
    weight2::AbstractMatrix{T}
    bias::AbstractMatrix{T}
    aggr::Symbol
end

function GraphConv(el::AbstractVector{<:AbstractVector{<:Integer}},
                   ch::Pair{<:Integer,<:Integer}, aggr=:add;
                   init = glorot_uniform, bias::Bool=true)
    N = size(el, 1)
    b = bias ? init(ch[2], N) : zeros(T, ch[2], N)
    GraphConv(el, init(ch[2], ch[1]), init(ch[2], ch[1]), b, aggr)
end

function GraphConv(adj::AbstractMatrix, ch::Pair{<:Integer,<:Integer}, aggr=:add;
                   init = glorot_uniform, bias::Bool=true, T::DataType=Float32)
    N = size(adj, 1)
    b = bias ? init(ch[2], N) : zeros(T, ch[2], N)
    GraphConv(neighbors(adj), init(ch[2], ch[1]), init(ch[2], ch[1]), b, aggr)
end

@functor GraphConv

message(g::GraphConv; x_i=zeros(0), x_j=zeros(0)) = g.weight2 * x_j
update(g::GraphConv; X=zeros(0), M=zeros(0)) = g.weight1*X + M + g.bias
(g::GraphConv)(X::AbstractMatrix) = propagate(g, X=X, aggr=:add)

function Base.show(io::IO, l::GraphConv)
    in_channel = size(l.weight1, ndims(l.weight1))
    out_channel = size(l.weight1, ndims(l.weight1)-1)
    print(io, "GraphConv(G(V=", length(l.adjlist), ", E=", sum(length, l.adjlist)÷2)
    print(io, "), ", in_channel, "=>", out_channel)
    print(io, ", aggr=", AGGR2STR[l.aggr])
    print(io, ")")
end

"""
    GATConv(graph, in=>out)

Graph attentional layer.

# Arguments
- `graph`: should be a adjacency matrix, `SimpleGraph`, `SimpleDiGraph` (from LightGraphs) or `SimpleWeightedGraph`, `SimpleWeightedDiGraph` (from SimpleWeightedGraphs).
- `in`: the dimension of input features.
- `out`: the dimension of output features.
- `bias::Bool=true`: keyword argument, whether to learn the additive bias.
- `negative_slope::Real=0.2`: keyword argument, the parameter of LeakyReLU.
"""
struct GATConv{V,T} <: MessagePassing
    adjlist::V
    weight::AbstractMatrix{T}
    bias::AbstractMatrix{T}
    a::AbstractArray{T,3}
    negative_slope::Real
    channel::Pair{<:Integer,<:Integer}
    heads::Integer
    concat::Bool
end

function GATConv(adj::AbstractMatrix, ch::Pair{<:Integer,<:Integer}; heads::Integer=1,
                 concat::Bool=true, negative_slope::Real=0.2, init=glorot_uniform,
                 bias::Bool=true, T::DataType=Float32)
    N = size(adj, 1)
    w = init(ch[2]*heads, ch[1])
    b = bias ? init(ch[2]*heads, N) : zeros(T, ch[2]*heads, N)
    a = init(2*ch[2], heads, 1)
    GATConv(neighbors(adj), w, b, a, negative_slope, ch, heads, concat)
end

@functor GATConv

function message(g::GATConv; x_i=zeros(0), x_j=zeros(0))
    x_i = reshape(x_i, g.channel[2], g.heads, :)
    x_j = reshape(x_j, g.channel[2], g.heads, :)
    n = size(x_j, 3)
    α = cat(repeat(x_i, outer=(1,1,n)), x_j+zero(x_j), dims=1) .* g.a
    α = reshape(sum(α, dims=1), g.heads, n)
    α = leakyrelu.(α, g.negative_slope)
    α = _softmax(α)
    x_j .*= reshape(α, 1, g.heads, n)
    reshape(x_j, g.channel[2]*g.heads, :)
end

function update(g::GATConv; X=zeros(0), M=zeros(0))
    if !g.concat
        M = mean(M, dims=2)
    end
    return M .+ g.bias
end

(g::GATConv)(X::AbstractMatrix) = propagate(g, X=g.weight*X, aggr=:add)

function _softmax(xs)
    xs = exp.(xs)
    s = sum(xs, dims=2)
    return xs ./ s
end

function Base.show(io::IO, l::GATConv)
    in_channel = size(l.weight, ndims(l.weight))
    out_channel = size(l.weight, ndims(l.weight)-1)
    print(io, "GATConv(G(V=", length(l.adjlist), ", E=", sum(length, l.adjlist)÷2)
    print(io, "), ", in_channel, "=>", out_channel)
    print(io, ", LeakyReLU(λ=", l.negative_slope)
    print(io, "))")
end

"""
    GatedGraphConv(graph, out, num_layers)

Gated graph convolution layer.

# Arguments
- `graph`: should be a adjacency matrix, `SimpleGraph`, `SimpleDiGraph` (from LightGraphs) or `SimpleWeightedGraph`, `SimpleWeightedDiGraph` (from SimpleWeightedGraphs).
- `out`: the dimension of output features.
- `num_layers` specifies the number of gated recurrent unit.
- `aggr::Symbol=:add`: an aggregate function applied to the result of message function. `:add`, `:max` and `:mean` are available.
"""
struct GatedGraphConv{V,T,R} <: MessagePassing
    adjlist::V
    weight::AbstractArray{T}
    gru::R
    out_ch::Integer
    num_layers::Integer
    aggr::Symbol
end

function GatedGraphConv(adj::AbstractMatrix, out_ch::Integer, num_layers::Integer;
                        aggr=:add, init=glorot_uniform)
    N = size(adj, 1)
    w = init(out_ch, out_ch, num_layers)
    gru = GRUCell(out_ch, out_ch)
    GatedGraphConv(neighbors(adj), w, gru, out_ch, num_layers, aggr)
end

@functor GatedGraphConv

message(g::GatedGraphConv; x_i=zeros(0), x_j=zeros(0)) = x_j
update(g::GatedGraphConv; X=zeros(0), M=zeros(0)) = M
function (g::GatedGraphConv)(X::AbstractMatrix{T}) where {T<:Real}
    H = X
    m, n = size(H)
    @assert (m <= g.out_ch) "number of input features must less or equals to output features."
    (m < g.out_ch) && (H = vcat(H, zeros(T, g.out_ch - m, n)))

    for i = 1:g.num_layers
        M = view(g.weight, :, :, i) * H
        M = propagate(g, X=M, aggr=g.aggr)
        H, _ = g.gru(H, M)
    end
    H
end

function Base.show(io::IO, l::GatedGraphConv)
    print(io, "GatedGraphConv(G(V=", length(l.adjlist), ", E=", sum(length, l.adjlist)÷2)
    print(io, "), (=>", l.out_ch)
    print(io, ")^", l.num_layers)
    print(io, ", aggr=", AGGR2STR[l.aggr])
    print(io, ")")
end

"""
    EdgeConv(graph, nn)
    EdgeConv(graph, nn, aggr)

Edge convolutional layer.

# Arguments
- `graph`: should be a adjacency matrix, `SimpleGraph`, `SimpleDiGraph` (from LightGraphs) or `SimpleWeightedGraph`, `SimpleWeightedDiGraph` (from SimpleWeightedGraphs).
- `nn`: a neural network
- `aggr::Symbol=:max`: an aggregate function applied to the result of message function. `:add`, `:max` and `:mean` are available.
"""
struct EdgeConv{V} <: MessagePassing
    adjlist::V
    nn
    aggr::Symbol
end

function EdgeConv(adj::AbstractMatrix, nn; aggr::Symbol=:max)
    EdgeConv(neighbors(adj), nn, aggr)
end

@functor EdgeConv

function message(e::EdgeConv; x_i=zeros(0), x_j=zeros(0))
    n = size(x_j, 2)
    e.nn(vcat(repeat(x_i, outer=(1,n)), x_j .- x_i))
end
update(e::EdgeConv; X=zeros(0), M=zeros(0)) = M
(e::EdgeConv)(X::AbstractMatrix) = propagate(e, X=X, aggr=e.aggr)

function Base.show(io::IO, l::EdgeConv)
    print(io, "EdgeConv(G(V=", length(l.adjlist), ", E=", sum(length, l.adjlist)÷2)
    print(io, "), ", l.nn)
    print(io, ", aggr=", AGGR2STR[l.aggr])
    print(io, ")")
end


text:
# This file is part of the StateMachines.jl Julia package
#
# Paul Bayer
# MIT License
# Github: https://github.com/pbayer/StateMachines.jl
#

"""
    step!(A::StateMachine, q::AbstractState, σ::AbstractEvent)

transition function δ causing a state machine A in state q₁ at event σ
to take on a new state q₂.

For all specified transitions of A a `step!(A, q, σ)`-function has to
be implemented. The fallback step function does nothing and is called for all
other calls to `step!(A, q, σ)`.

# Arguments
- `A::StateMachine`: a state machine
- `q::AbstractState`: any state ∈ Q
- `σ::AbstractEvent`: any discrete event ∈ Σ
"""
function step!(A::StateMachine, q::AbstractState, σ::AbstractEvent)
    trace ? println(stderr, "undefined transition: $(A.id), $(typeof(q)), $(typeof(σ))") : nothing
end


text:
# ---
# title: 1020. Number of Enclaves
# id: problem1020
# author: zhwang
# date: 2022-02-12
# difficulty: Medium
# categories: Depth-first Search
# link: <https://leetcode.com/problems/number-of-enclaves/description/>
# hidden: true
# ---
#
# Given a 2D array `A`, each cell is 0 (representing sea) or 1 (representing
# land)
#
# A move consists of walking from one land square 4-directionally to another
# land square, or off the boundary of the grid.
#
# Return the number of land squares in the grid for which we **cannot** walk off
# the boundary of the grid in any number of moves.
#
#
#
# **Example 1:**
#
#
#
#     Input: [[0,0,0,0],[1,0,1,0],[0,1,1,0],[0,0,0,0]]
#     Output: 3
#     Explanation:
#     There are three 1s that are enclosed by 0s, and one 1 that isn't enclosed because its on the boundary.
#
# **Example 2:**
#
#
#
#     Input: [[0,1,1,0],[0,0,1,0],[0,0,1,0],[0,0,0,0]]
#     Output: 0
#     Explanation:
#     All 1s are either on the boundary or can reach the boundary.
#
#
#
#
# **Note:**
#
#   1. `1 <= A.length <= 500`
#   2. `1 <= A[i].length <= 500`
#   3. `0 <= A[i][j] <= 1`
#   4. All rows have the same size.
#
#
## @lc code=start
using LeetCode

## BFS algorithm
function num_of_enclaves_bfs(grid::Vector{Vector{Int}})::Int
    m, n, queue = length(grid), length(grid[1]), Tuple{Int,Int}[]
    (m <= 2 || n <= 2) && return 0
    for i in 1:m ## left and right border
        grid[i][1] == 1 && push!(queue, (i, 1))
        grid[i][n] == 1 && push!(queue, (i, n))
    end
    for i in 2:n-1
        grid[1][i] == 1 && push!(queue, (1, i))
        grid[m][i] == 1 && push!(queue, (m, i))
    end
    walkoff = copy(queue)
    while length(queue) != 0
        i, j = popfirst!(queue)
        for (a, b) in [(i-1, j), (i+1, j), (i,j-1), (i,j+1)]
            if a >= 1 && b >= 1 && a <= m && b <= n && grid[a][b] != 0 && !((a, b) in walkoff)
                push!(walkoff, (a, b))
                push!(queue, (a, b))
            end
        end
    end
    sum(sum(line) for line in grid) - length(walkoff)
end
## @lc code=end

## @lc test=start
@testset "1020.number-of-enclaves.jl" begin
    @test num_of_enclaves_bfs([[0,0,0,0],[1,0,1,0],[0,1,1,0],[0,0,0,0]]) == 3
    @test num_of_enclaves_bfs([[0,1,1,0],[0,0,1,0],[0,0,1,0],[0,0,0,0]]) == 0
    @test num_of_enclaves_bfs([[0],[1],[1],[0],[0]]) == 0
    @test num_of_enclaves_bfs(
        [[0,0,0,1,1,1,0,1,0,0],
         [1,1,0,0,0,1,0,1,1,1],
         [0,0,0,1,1,1,0,1,0,0],
         [0,1,1,0,0,0,1,0,1,0],
         [0,1,1,1,1,1,0,0,1,0],
         [0,0,1,0,1,1,1,1,0,1],
         [0,1,1,0,0,0,1,1,1,1],
         [0,0,1,0,0,1,0,1,0,1],
         [1,0,1,0,1,1,0,0,0,0],
         [0,0,0,0,1,1,0,0,0,1]])==3
end
## @lc test=end

## @lc info=start
# link: [solution to 2006](https://leetcode-cn.com/problems/number-of-enclaves/solution/1020-fei-di-de-shu-liang-by-rex-rs-sav5/)
## @lc info=end


text:
#=

Plot the price-dividend ratio in the risk neutral case, for the Markov asset pricing
lecture.

=#
using QuantEcon
using Plots
using LaTeXStrings

pyplot()

n = 25  # size of state space
beta = 0.9
mc = tauchen(n, 0.96, 0.02)

K = mc.p .* reshape(exp(mc.state_values), 1, n)

I = eye(n)
v = (I - beta * K) \  (beta * K * ones(n, 1))

plot(mc.state_values,
    v,
    lw=2,
    ylabel="price-dividend ratio",
    xlabel="state",
    alpha=0.7,
    label=L"$v$")

max_stars_repo_name: mwhchen/quantecon

max_issues_repo_name: mwhchen/quantecon

max_forks_repo_name: mwhchen/quantecon


text:
#=
Description
    Module for POMDP type objects with state, action, and
    observation space, etc. defined.
=#

module POMDPs

export POMDP

type POMDP

    states          :: Matrix
    actions         :: Array
    observations    :: Matrix
    term_state      :: Vector{Float64}

    function POMDP(states::Matrix, actions::Array, observations::Matrix,
                   term_state::Vector{Float64})
        return new(states, actions, observations, term_state)
    end # function POMDP

end # type POMDP

end # module POMDPs