text:
####
#### Types for typed and untyped VarInfo
####

####################
# VarInfo metadata #
####################

"""
The `Metadata` struct stores some metadata about the parameters of the model. This helps
query certain information about a variable, such as its distribution, which samplers
sample this variable, its value and whether this value is transformed to real space or
not.

Let `md` be an instance of `Metadata`:
- `md.vns` is the vector of all `VarName` instances.
- `md.idcs` is the dictionary that maps each `VarName` instance to its index in
 `md.vns`, `md.ranges` `md.dists`, `md.orders` and `md.flags`.
- `md.vns[md.idcs[vn]] == vn`.
- `md.dists[md.idcs[vn]]` is the distribution of `vn`.
- `md.gids[md.idcs[vn]]` is the set of algorithms used to sample `vn`. This is used in
 the Gibbs sampling process.
- `md.orders[md.idcs[vn]]` is the number of `observe` statements before `vn` is sampled.
- `md.ranges[md.idcs[vn]]` is the index range of `vn` in `md.vals`.
- `md.vals[md.ranges[md.idcs[vn]]]` is the vector of values of corresponding to `vn`.
- `md.flags` is a dictionary of true/false flags. `md.flags[flag][md.idcs[vn]]` is the
 value of `flag` corresponding to `vn`.

To make `md::Metadata` type stable, all the `md.vns` must have the same symbol
and distribution type. However, one can have a Julia variable, say `x`, that is a
matrix or a hierarchical array sampled in partitions, e.g.
`x[1][:] ~ MvNormal(zeros(2), I); x[2][:] ~ MvNormal(ones(2), I)`, and is managed by
a single `md::Metadata` so long as all the distributions on the RHS of `~` are of the
same type. Type unstable `Metadata` will still work but will have inferior performance.
When sampling, the first iteration uses a type unstable `Metadata` for all the
variables then a specialized `Metadata` is used for each symbol along with a function
barrier to make the rest of the sampling type stable.
"""
struct Metadata{
    TIdcs<:Dict{<:VarName,Int},
    TDists<:AbstractVector{<:Distribution},
    TVN<:AbstractVector{<:VarName},
    TVal<:AbstractVector{<:Real},
    TGIds<:AbstractVector{Set{Selector}},
}
    # Mapping from the `VarName` to its integer index in `vns`, `ranges` and `dists`
    idcs::TIdcs # Dict{<:VarName,Int}

    # Vector of identifiers for the random variables, where `vns[idcs[vn]] == vn`
    vns::TVN # AbstractVector{<:VarName}

    # Vector of index ranges in `vals` corresponding to `vns`
    # Each `VarName` `vn` has a single index or a set of contiguous indices in `vals`
    ranges::Vector{UnitRange{Int}}

    # Vector of values of all the univariate, multivariate and matrix variables
    # The value(s) of `vn` is/are `vals[ranges[idcs[vn]]]`
    vals::TVal # AbstractVector{<:Real}

    # Vector of distributions correpsonding to `vns`
    dists::TDists # AbstractVector{<:Distribution}

    # Vector of sampler ids corresponding to `vns`
    # Each random variable can be sampled using multiple samplers, e.g. in Gibbs, hence the `Set`
    gids::TGIds # AbstractVector{Set{Selector}}

    # Number of `observe` statements before each random variable is sampled
    orders::Vector{Int}

    # Each `flag` has a `BitVector` `flags[flag]`, where `flags[flag][i]` is the true/false flag value corresonding to `vns[i]`
    flags::Dict{String,BitVector}
end

###########
# VarInfo #
###########

"""
```
struct VarInfo{Tmeta, Tlogp} <: AbstractVarInfo
    metadata::Tmeta
    logp::Base.RefValue{Tlogp}
    num_produce::Base.RefValue{Int}
end
```

A light wrapper over one or more instances of `Metadata`. Let `vi` be an instance of
`VarInfo`. If `vi isa VarInfo{<:Metadata}`, then only one `Metadata` instance is used
for all the sybmols. `VarInfo{<:Metadata}` is aliased `UntypedVarInfo`. If
`vi isa VarInfo{<:NamedTuple}`, then `vi.metadata` is a `NamedTuple` that maps each
symbol used on the LHS of `~` in the model to its `Metadata` instance. The latter allows
for the type specialization of `vi` after the first sampling iteration when all the
symbols have been observed. `VarInfo{<:NamedTuple}` is aliased `TypedVarInfo`.

Note: It is the user's responsibility to ensure that each "symbol" is visited at least
once whenever the model is called, regardless of any stochastic branching. Each symbol
refers to a Julia variable and can be a hierarchical array of many random variables, e.g. `x[1] ~ ...` and `x[2] ~ ...` both have the same symbol `x`.
"""
struct VarInfo{Tmeta,Tlogp} <: AbstractVarInfo
    metadata::Tmeta
    logp::Base.RefValue{Tlogp}
    num_produce::Base.RefValue{Int}
end
const UntypedVarInfo = VarInfo{<:Metadata}
const TypedVarInfo = VarInfo{<:NamedTuple}

function VarInfo(old_vi::UntypedVarInfo, spl, x::AbstractVector)
    new_vi = deepcopy(old_vi)
    new_vi[spl] = x
    return new_vi
end

function VarInfo(old_vi::TypedVarInfo, spl, x::AbstractVector)
    md = newmetadata(old_vi.metadata, Val(getspace(spl)), x)
    return VarInfo(
        md, Base.RefValue{eltype(x)}(getlogp(old_vi)), Ref(get_num_produce(old_vi))
    )
end

function VarInfo(
    rng::Random.AbstractRNG,
    model::Model,
    sampler::AbstractSampler=SampleFromPrior(),
    context::AbstractContext=DefaultContext(),
)
    varinfo = VarInfo()
    model(rng, varinfo, sampler, context)
    return TypedVarInfo(varinfo)
end
VarInfo(model::Model, args...) = VarInfo(Random.GLOBAL_RNG, model, args...)

# without AbstractSampler
function VarInfo(rng::Random.AbstractRNG, model::Model, context::AbstractContext)
    return VarInfo(rng, model, SampleFromPrior(), context)
end

@generated function newmetadata(
    metadata::NamedTuple{names}, ::Val{space}, x
) where {names,space}
    exprs = []
    offset = :(0)
    for f in names
        mdf = :(metadata.$f)
        if inspace(f, space) || length(space) == 0
            len = :(length($mdf.vals))
            push!(
                exprs,
                :(
                    $f = Metadata(
                        $mdf.idcs,
                        $mdf.vns,
                        $mdf.ranges,
                        x[($offset + 1):($offset + $len)],
                        $mdf.dists,
                        $mdf.gids,
                        $mdf.orders,
                        $mdf.flags,
                    )
                ),
            )
            offset = :($offset + $len)
        else
            push!(exprs, :($f = $mdf))
        end
    end
    length(exprs) == 0 && return :(NamedTuple())
    return :($(exprs...),)
end

####
#### Internal functions
####

"""
    Metadata()

Construct an empty type unstable instance of `Metadata`.
"""
function Metadata()
    vals = Vector{Real}()
    flags = Dict{String,BitVector}()
    flags["del"] = BitVector()
    flags["trans"] = BitVector()

    return Metadata(
        Dict{VarName,Int}(),
        Vector{VarName}(),
        Vector{UnitRange{Int}}(),
        vals,
        Vector{Distribution}(),
        Vector{Set{Selector}}(),
        Vector{Int}(),
        flags,
    )
end

"""
    empty!(meta::Metadata)

Empty the fields of `meta`.

This is useful when using a sampling algorithm that assumes an empty `meta`, e.g. `SMC`.
"""
function empty!(meta::Metadata)
    empty!(meta.idcs)
    empty!(meta.vns)
    empty!(meta.ranges)
    empty!(meta.vals)
    empty!(meta.dists)
    empty!(meta.gids)
    empty!(meta.orders)
    for k in keys(meta.flags)
        empty!(meta.flags[k])
    end

    return meta
end

# Removes the first element of a NamedTuple. The pairs in a NamedTuple are ordered, so this is well-defined.
if VERSION < v"1.1"
    _tail(nt::NamedTuple{names}) where {names} = NamedTuple{Base.tail(names)}(nt)
else
    _tail(nt::NamedTuple) = Base.tail(nt)
end

const VarView = Union{Int,UnitRange,Vector{Int}}

"""
    getval(vi::UntypedVarInfo, vview::Union{Int, UnitRange, Vector{Int}})

Return a view `vi.vals[vview]`.
"""
getval(vi::UntypedVarInfo, vview::VarView) = view(vi.metadata.vals, vview)

"""
    setval!(vi::UntypedVarInfo, val, vview::Union{Int, UnitRange, Vector{Int}})

Set the value of `vi.vals[vview]` to `val`.
"""
setval!(vi::UntypedVarInfo, val, vview::VarView) = vi.metadata.vals[vview] = val
function setval!(vi::UntypedVarInfo, val, vview::Vector{UnitRange})
    if length(vview) > 0
        vi.metadata.vals[[i for arr in vview for i in arr]] = val
    end
    return val
end

"""
    getmetadata(vi::VarInfo, vn::VarName)

Return the metadata in `vi` that belongs to `vn`.
"""
getmetadata(vi::VarInfo, vn::VarName) = vi.metadata
getmetadata(vi::TypedVarInfo, vn::VarName) = getfield(vi.metadata, getsym(vn))

"""
    getidx(vi::VarInfo, vn::VarName)

Return the index of `vn` in the metadata of `vi` corresponding to `vn`.
"""
getidx(vi::VarInfo, vn::VarName) = getmetadata(vi, vn).idcs[vn]

"""
    getrange(vi::VarInfo, vn::VarName)

Return the index range of `vn` in the metadata of `vi`.
"""
getrange(vi::VarInfo, vn::VarName) = getmetadata(vi, vn).ranges[getidx(vi, vn)]

"""
    getranges(vi::AbstractVarInfo, vns::Vector{<:VarName})

Return the indices of `vns` in the metadata of `vi` corresponding to `vn`.
"""
function getranges(vi::AbstractVarInfo, vns::Vector{<:VarName})
    return mapreduce(vn -> getrange(vi, vn), vcat, vns; init=Int[])
end

"""
    getdist(vi::VarInfo, vn::VarName)

Return the distribution from which `vn` was sampled in `vi`.
"""
getdist(vi::VarInfo, vn::VarName) = getmetadata(vi, vn).dists[getidx(vi, vn)]

"""
    getval(vi::VarInfo, vn::VarName)

Return the value(s) of `vn`.

The values may or may not be transformed to Euclidean space.
"""
getval(vi::VarInfo, vn::VarName) = view(getmetadata(vi, vn).vals, getrange(vi, vn))

"""
    setval!(vi::VarInfo, val, vn::VarName)

Set the value(s) of `vn` in the metadata of `vi` to `val`.

The values may or may not be transformed to Euclidean space.
"""
setval!(vi::VarInfo, val, vn::VarName) = getmetadata(vi, vn).vals[getrange(vi, vn)] = [val;]

"""
    getval(vi::VarInfo, vns::Vector{<:VarName})

Return the value(s) of `vns`.

The values may or may not be transformed to Euclidean space.
"""
function getval(vi::AbstractVarInfo, vns::Vector{<:VarName})
    return mapreduce(vn -> getval(vi, vn), vcat, vns)
end

"""
    getall(vi::VarInfo)

Return the values of all the variables in `vi`.

The values may or may not be transformed to Euclidean space.
"""
getall(vi::UntypedVarInfo) = vi.metadata.vals
getall(vi::TypedVarInfo) = vcat(_getall(vi.metadata)...)
@generated function _getall(metadata::NamedTuple{names}) where {names}
    exprs = []
    for f in names
        push!(exprs, :(metadata.$f.vals))
    end
    return :($(exprs...),)
end

"""
    setall!(vi::VarInfo, val)

Set the values of all the variables in `vi` to `val`.

The values may or may not be transformed to Euclidean space.
"""
setall!(vi::UntypedVarInfo, val) = vi.metadata.vals .= val
setall!(vi::TypedVarInfo, val) = _setall!(vi.metadata, val)
@generated function _setall!(metadata::NamedTuple{names}, val, start=0) where {names}
    expr = Expr(:block)
    start = :(1)
    for f in names
        length = :(length(metadata.$f.vals))
        finish = :($start + $length - 1)
        push!(expr.args, :(metadata.$f.vals .= val[($start):($finish)]))
        start = :($start + $length)
    end
    return expr
end

"""
    getgid(vi::VarInfo, vn::VarName)

Return the set of sampler selectors associated with `vn` in `vi`.
"""
getgid(vi::VarInfo, vn::VarName) = getmetadata(vi, vn).gids[getidx(vi, vn)]

"""
    settrans!(vi::VarInfo, trans::Bool, vn::VarName)

Set the `trans` flag value of `vn` in `vi`.
"""
function settrans!(vi::AbstractVarInfo, trans::Bool, vn::VarName)
    return trans ? set_flag!(vi, vn, "trans") : unset_flag!(vi, vn, "trans")
end

"""
    syms(vi::VarInfo)

Returns a tuple of the unique symbols of random variables sampled in `vi`.
"""
syms(vi::UntypedVarInfo) = Tuple(unique!(map(getsym, vi.metadata.vns)))  # get all symbols
syms(vi::TypedVarInfo) = keys(vi.metadata)

# Get all indices of variables belonging to SampleFromPrior:
#   if the gid/selector of a var is an empty Set, then that var is assumed to be assigned to
#   the SampleFromPrior sampler
@inline function _getidcs(vi::UntypedVarInfo, ::SampleFromPrior)
    return filter(i -> isempty(vi.metadata.gids[i]), 1:length(vi.metadata.gids))
end
# Get a NamedTuple of all the indices belonging to SampleFromPrior, one for each symbol
@inline function _getidcs(vi::TypedVarInfo, ::SampleFromPrior)
    return _getidcs(vi.metadata)
end
@generated function _getidcs(metadata::NamedTuple{names}) where {names}
    exprs = []
    for f in names
        push!(exprs, :($f = findinds(metadata.$f)))
    end
    length(exprs) == 0 && return :(NamedTuple())
    return :($(exprs...),)
end

# Get all indices of variables belonging to a given sampler
@inline function _getidcs(vi::AbstractVarInfo, spl::Sampler)
    # NOTE: 0b00 is the sanity flag for
    #         |\____ getidcs   (mask = 0b10)
    #         \_____ getranges (mask = 0b01)
    #if ~haskey(spl.info, :cache_updated) spl.info[:cache_updated] = CACHERESET end
    # Checks if cache is valid, i.e. no new pushes were made, to return the cached idcs
    # Otherwise, it recomputes the idcs and caches it
    #if haskey(spl.info, :idcs) && (spl.info[:cache_updated] & CACHEIDCS) > 0
    #    spl.info[:idcs]
    #else
    #spl.info[:cache_updated] = spl.info[:cache_updated] | CACHEIDCS
    idcs = _getidcs(vi, spl.selector, Val(getspace(spl)))
    #spl.info[:idcs] = idcs
    #end
    return idcs
end
@inline _getidcs(vi::UntypedVarInfo, s::Selector, space) = findinds(vi.metadata, s, space)
@inline _getidcs(vi::TypedVarInfo, s::Selector, space) = _getidcs(vi.metadata, s, space)
# Get a NamedTuple for all the indices belonging to a given selector for each symbol
@generated function _getidcs(
    metadata::NamedTuple{names}, s::Selector, ::Val{space}
) where {names,space}
    exprs = []
    # Iterate through each varname in metadata.
    for f in names
        # If the varname is in the sampler space
        # or the sample space is empty (all variables)
        # then return the indices for that variable.
        if inspace(f, space) || length(space) == 0
            push!(exprs, :($f = findinds(metadata.$f, s, Val($space))))
        end
    end
    length(exprs) == 0 && return :(NamedTuple())
    return :($(exprs...),)
end
@inline function findinds(f_meta, s, ::Val{space}) where {space}
    # Get all the idcs of the vns in `space` and that belong to the selector `s`
    return filter(
        (i) ->
            (s in f_meta.gids[i] || isempty(f_meta.gids[i])) &&
                (isempty(space) || inspace(f_meta.vns[i], space)),
        1:length(f_meta.gids),
    )
end
@inline function findinds(f_meta)
    # Get all the idcs of the vns
    return filter((i) -> isempty(f_meta.gids[i]), 1:length(f_meta.gids))
end

# Get all vns of variables belonging to spl
_getvns(vi::AbstractVarInfo, spl::Sampler) = _getvns(vi, spl.selector, Val(getspace(spl)))
function _getvns(vi::AbstractVarInfo, spl::Union{SampleFromPrior,SampleFromUniform})
    return _getvns(vi, Selector(), Val(()))
end
function _getvns(vi::UntypedVarInfo, s::Selector, space)
    return view(vi.metadata.vns, _getidcs(vi, s, space))
end
function _getvns(vi::TypedVarInfo, s::Selector, space)
    return _getvns(vi.metadata, _getidcs(vi, s, space))
end
# Get a NamedTuple for all the `vns` of indices `idcs`, one entry for each symbol
@generated function _getvns(metadata, idcs::NamedTuple{names}) where {names}
    exprs = []
    for f in names
        push!(exprs, :($f = metadata.$f.vns[idcs.$f]))
    end
    length(exprs) == 0 && return :(NamedTuple())
    return :($(exprs...),)
end

# Get the index (in vals) ranges of all the vns of variables belonging to spl
@inline function _getranges(vi::AbstractVarInfo, spl::Sampler)
    ## Uncomment the spl.info stuff when it is concretely typed, not Dict{Symbol, Any}
    #if ~haskey(spl.info, :cache_updated) spl.info[:cache_updated] = CACHERESET end
    #if haskey(spl.info, :ranges) && (spl.info[:cache_updated] & CACHERANGES) > 0
    #    spl.info[:ranges]
    #else
    #spl.info[:cache_updated] = spl.info[:cache_updated] | CACHERANGES
    ranges = _getranges(vi, spl.selector, Val(getspace(spl)))
    #spl.info[:ranges] = ranges
    return ranges
    #end
end
# Get the index (in vals) ranges of all the vns of variables belonging to selector `s` in `space`
@inline function _getranges(vi::AbstractVarInfo, s::Selector, space)
    return _getranges(vi, _getidcs(vi, s, space))
end
@inline function _getranges(vi::UntypedVarInfo, idcs::Vector{Int})
    return mapreduce(i -> vi.metadata.ranges[i], vcat, idcs; init=Int[])
end
@inline _getranges(vi::TypedVarInfo, idcs::NamedTuple) = _getranges(vi.metadata, idcs)

@generated function _getranges(metadata::NamedTuple, idcs::NamedTuple{names}) where {names}
    exprs = []
    for f in names
        push!(exprs, :($f = findranges(metadata.$f.ranges, idcs.$f)))
    end
    length(exprs) == 0 && return :(NamedTuple())
    return :($(exprs...),)
end

@inline function findranges(f_ranges, f_idcs)
    # Old implementation was using `mapreduce` but turned out
    # to be type-unstable.
    results = Int[]
    for i in f_idcs
        append!(results, f_ranges[i])
    end
    return results
end

"""
    set_flag!(vi::VarInfo, vn::VarName, flag::String)

Set `vn`'s value for `flag` to `true` in `vi`.
"""
function set_flag!(vi::VarInfo, vn::VarName, flag::String)
    return getmetadata(vi, vn).flags[flag][getidx(vi, vn)] = true
end

####
#### APIs for typed and untyped VarInfo
####

# VarInfo

VarInfo(meta=Metadata()) = VarInfo(meta, Ref{Float64}(0.0), Ref(0))

"""
    TypedVarInfo(vi::UntypedVarInfo)

This function finds all the unique `sym`s from the instances of `VarName{sym}` found in
`vi.metadata.vns`. It then extracts the metadata associated with each symbol from the
global `vi.metadata` field. Finally, a new `VarInfo` is created with a new `metadata` as
a `NamedTuple` mapping from symbols to type-stable `Metadata` instances, one for each
symbol.
"""
function TypedVarInfo(vi::UntypedVarInfo)
    meta = vi.metadata
    new_metas = Metadata[]
    # Symbols of all instances of `VarName{sym}` in `vi.vns`
    syms_tuple = Tuple(syms(vi))
    for s in syms_tuple
        # Find all indices in `vns` with symbol `s`
        inds = findall(vn -> getsym(vn) === s, meta.vns)
        n = length(inds)
        # New `vns`
        sym_vns = getindex.((meta.vns,), inds)
        # New idcs
        sym_idcs = Dict(a => i for (i, a) in enumerate(sym_vns))
        # New dists
        sym_dists = getindex.((meta.dists,), inds)
        # New gids, can make a resizeable FillArray
        sym_gids = getindex.((meta.gids,), inds)
        @assert length(sym_gids) <= 1 || all(x -> x == sym_gids[1], @view sym_gids[2:end])
        # New orders
        sym_orders = getindex.((meta.orders,), inds)
        # New flags
        sym_flags = Dict(a => meta.flags[a][inds] for a in keys(meta.flags))

        # Extract new ranges and vals
        _ranges = getindex.((meta.ranges,), inds)
        # `copy.()` is a workaround to reduce the eltype from Real to Int or Float64
        _vals = [copy.(meta.vals[_ranges[i]]) for i in 1:n]
        sym_ranges = Vector{eltype(_ranges)}(undef, n)
        start = 0
        for i in 1:n
            sym_ranges[i] = (start + 1):(start + length(_vals[i]))
            start += length(_vals[i])
        end
        sym_vals = foldl(vcat, _vals)

        push!(
            new_metas,
            Metadata(
                sym_idcs,
                sym_vns,
                sym_ranges,
                sym_vals,
                sym_dists,
                sym_gids,
                sym_orders,
                sym_flags,
            ),
        )
    end
    logp = getlogp(vi)
    num_produce = get_num_produce(vi)
    nt = NamedTuple{syms_tuple}(Tuple(new_metas))
    return VarInfo(nt, Ref(logp), Ref(num_produce))
end
TypedVarInfo(vi::TypedVarInfo) = vi

"""
    empty!!(vi::VarInfo)

Empty the fields of `vi.metadata` and reset `vi.logp[]` and `vi.num_produce[]` to
zeros.

This is useful when using a sampling algorithm that assumes an empty `vi`, e.g. `SMC`.
"""
function BangBang.empty!!(vi::VarInfo)
    _empty!(vi.metadata)
    resetlogp!!(vi)
    reset_num_produce!(vi)
    return vi
end
@inline _empty!(metadata::Metadata) = empty!(metadata)
@generated function _empty!(metadata::NamedTuple{names}) where {names}
    expr = Expr(:block)
    for f in names
        push!(expr.args, :(empty!(metadata.$f)))
    end
    return expr
end

# Functions defined only for UntypedVarInfo
"""
    keys(vi::AbstractVarInfo)

Return an iterator over all `vns` in `vi`.
"""
Base.keys(vi::UntypedVarInfo) = keys(vi.metadata.idcs)

@generated function Base.keys(vi::TypedVarInfo{<:NamedTuple{names}}) where {names}
    expr = Expr(:call)
    push!(expr.args, :vcat)

    for n in names
        push!(expr.args, :(vi.metadata.$n.vns))
    end

    return expr
end

"""
    setgid!(vi::VarInfo, gid::Selector, vn::VarName)

Add `gid` to the set of sampler selectors associated with `vn` in `vi`.
"""
function setgid!(vi::VarInfo, gid::Selector, vn::VarName)
    return push!(getmetadata(vi, vn).gids[getidx(vi, vn)], gid)
end

"""
    istrans(vi::VarInfo, vn::VarName)

Return true if `vn`'s values in `vi` are transformed to Euclidean space, and false if
they are in the support of `vn`'s distribution.
"""
istrans(vi::AbstractVarInfo, vn::VarName) = is_flagged(vi, vn, "trans")

"""
    getlogp(vi::VarInfo)

Return the log of the joint probability of the observed data and parameters sampled in
`vi`.
"""
getlogp(vi::AbstractVarInfo) = vi.logp[]

"""
    setlogp!!(vi::VarInfo, logp)

Set the log of the joint probability of the observed data and parameters sampled in
`vi` to `logp`, mutating if it makes sense.
"""
function setlogp!!(vi::VarInfo, logp)
    vi.logp[] = logp
    return vi
end

"""
    acclogp!!(vi::VarInfo, logp)

Add `logp` to the value of the log of the joint probability of the observed data and
parameters sampled in `vi`, mutating if it makes sense.
"""
function acclogp!!(vi::VarInfo, logp)
    vi.logp[] += logp
    return vi
end

"""
    resetlogp!!(vi::AbstractVarInfo)

Reset the value of the log of the joint probability of the observed data and parameters
sampled in `vi` to 0, mutating if it makes sense.
"""
resetlogp!!(vi::AbstractVarInfo) = setlogp!!(vi, zero(getlogp(vi)))

"""
    get_num_produce(vi::VarInfo)

Return the `num_produce` of `vi`.
"""
get_num_produce(vi::VarInfo) = vi.num_produce[]

"""
    set_num_produce!(vi::VarInfo, n::Int)

Set the `num_produce` field of `vi` to `n`.
"""
set_num_produce!(vi::VarInfo, n::Int) = vi.num_produce[] = n

"""
    increment_num_produce!(vi::VarInfo)

Add 1 to `num_produce` in `vi`.
"""
increment_num_produce!(vi::VarInfo) = vi.num_produce[] += 1

"""
    reset_num_produce!(vi::AbstractVarInfo)

Reset the value of `num_produce` the log of the joint probability of the observed data
and parameters sampled in `vi` to 0.
"""
reset_num_produce!(vi::AbstractVarInfo) = set_num_produce!(vi, 0)

"""
    isempty(vi::VarInfo)

Return true if `vi` is empty and false otherwise.
"""
isempty(vi::UntypedVarInfo) = isempty(vi.metadata.idcs)
isempty(vi::TypedVarInfo) = _isempty(vi.metadata)
@generated function _isempty(metadata::NamedTuple{names}) where {names}
    expr = Expr(:&&, :true)
    for f in names
        push!(expr.args, :(isempty(metadata.$f.idcs)))
    end
    return expr
end

# X -> R for all variables associated with given sampler
"""
    link!(vi::VarInfo, spl::Sampler)

Transform the values of the random variables sampled by `spl` in `vi` from the support
of their distributions to the Euclidean space and set their corresponding `"trans"`
flag values to `true`.
"""
function link!(vi::UntypedVarInfo, spl::Sampler)
    # TODO: Change to a lazy iterator over `vns`
    vns = _getvns(vi, spl)
    if ~istrans(vi, vns[1])
        for vn in vns
            @debug "X -> ℝ for $(vn)..."
            dist = getdist(vi, vn)
            # TODO: Use inplace versions to avoid allocations
            setval!(
                vi,
                vectorize(dist, Bijectors.link(dist, reconstruct(dist, getval(vi, vn)))),
                vn,
            )
            settrans!(vi, true, vn)
        end
    else
        @warn("[DynamicPPL] attempt to link a linked vi")
    end
end
function link!(vi::TypedVarInfo, spl::AbstractSampler)
    return link!(vi, spl, Val(getspace(spl)))
end
function link!(vi::TypedVarInfo, spl::AbstractSampler, spaceval::Val)
    vns = _getvns(vi, spl)
    return _link!(vi.metadata, vi, vns, spaceval)
end
@generated function _link!(
    metadata::NamedTuple{names}, vi, vns, ::Val{space}
) where {names,space}
    expr = Expr(:block)
    for f in names
        if inspace(f, space) || length(space) == 0
            push!(
                expr.args,
                quote
                    f_vns = vi.metadata.$f.vns
                    if ~istrans(vi, f_vns[1])
                        # Iterate over all `f_vns` and transform
                        for vn in f_vns
                            @debug "X -> R for $(vn)..."
                            dist = getdist(vi, vn)
                            setval!(
                                vi,
                                vectorize(
                                    dist,
                                    Bijectors.link(dist, reconstruct(dist, getval(vi, vn))),
                                ),
                                vn,
                            )
                            settrans!(vi, true, vn)
                        end
                    else
                        @warn("[DynamicPPL] attempt to link a linked vi")
                    end
                end,
            )
        end
    end
    return expr
end

# R -> X for all variables associated with given sampler
"""
    invlink!(vi::VarInfo, spl::AbstractSampler)

Transform the values of the random variables sampled by `spl` in `vi` from the
Euclidean space back to the support of their distributions and sets their corresponding
`"trans"` flag values to `false`.
"""
function invlink!(vi::UntypedVarInfo, spl::AbstractSampler)
    vns = _getvns(vi, spl)
    if istrans(vi, vns[1])
        for vn in vns
            @debug "ℝ -> X for $(vn)..."
            dist = getdist(vi, vn)
            setval!(
                vi,
                vectorize(dist, Bijectors.invlink(dist, reconstruct(dist, getval(vi, vn)))),
                vn,
            )
            settrans!(vi, false, vn)
        end
    else
        @warn("[DynamicPPL] attempt to invlink an invlinked vi")
    end
end
function invlink!(vi::TypedVarInfo, spl::AbstractSampler)
    return invlink!(vi, spl, Val(getspace(spl)))
end
function invlink!(vi::TypedVarInfo, spl::AbstractSampler, spaceval::Val)
    vns = _getvns(vi, spl)
    return _invlink!(vi.metadata, vi, vns, spaceval)
end
@generated function _invlink!(
    metadata::NamedTuple{names}, vi, vns, ::Val{space}
) where {names,space}
    expr = Expr(:block)
    for f in names
        if inspace(f, space) || length(space) == 0
            push!(
                expr.args,
                quote
                    f_vns = vi.metadata.$f.vns
                    if istrans(vi, f_vns[1])
                        # Iterate over all `f_vns` and transform
                        for vn in f_vns
                            @debug "ℝ -> X for $(vn)..."
                            dist = getdist(vi, vn)
                            setval!(
                                vi,
                                vectorize(
                                    dist,
                                    Bijectors.invlink(
                                        dist, reconstruct(dist, getval(vi, vn))
                                    ),
                                ),
                                vn,
                            )
                            settrans!(vi, false, vn)
                        end
                    else
                        @warn("[DynamicPPL] attempt to invlink an invlinked vi")
                    end
                end,
            )
        end
    end
    return expr
end

"""
    islinked(vi::VarInfo, spl::Union{Sampler, SampleFromPrior})

Check whether `vi` is in the transformed space for a particular sampler `spl`.

Turing's Hamiltonian samplers use the `link` and `invlink` functions from
[Bijectors.jl](https://github.com/TuringLang/Bijectors.jl) to map a constrained variable
(for example, one bounded to the space `[0, 1]`) from its constrained space to the set of
real numbers. `islinked` checks if the number is in the constrained space or the real space.
"""
function islinked(vi::UntypedVarInfo, spl::Union{Sampler,SampleFromPrior})
    vns = _getvns(vi, spl)
    return istrans(vi, vns[1])
end
function islinked(vi::TypedVarInfo, spl::Union{Sampler,SampleFromPrior})
    vns = _getvns(vi, spl)
    return _islinked(vi, vns)
end
@generated function _islinked(vi, vns::NamedTuple{names}) where {names}
    out = []
    for f in names
        push!(out, :(length(vns.$f) == 0 ? false : istrans(vi, vns.$f[1])))
    end
    return Expr(:||, false, out...)
end

# The default getindex & setindex!() for get & set values
# NOTE: vi[vn] will always transform the variable to its original space and Julia type
"""
    getindex(vi::VarInfo, vn::VarName)
    getindex(vi::VarInfo, vns::Vector{<:VarName})

Return the current value(s) of `vn` (`vns`) in `vi` in the support of its (their)
distribution(s).

If the value(s) is (are) transformed to the Euclidean space, it is
(they are) transformed back.
"""
function getindex(vi::AbstractVarInfo, vn::VarName)
    @assert haskey(vi, vn) "[DynamicPPL] attempted to replay unexisting variables in VarInfo"
    dist = getdist(vi, vn)
    return if istrans(vi, vn)
        Bijectors.invlink(dist, reconstruct(dist, getval(vi, vn)))
    else
        reconstruct(dist, getval(vi, vn))
    end
end
function getindex(vi::AbstractVarInfo, vns::Vector{<:VarName})
    @assert haskey(vi, vns[1]) "[DynamicPPL] attempted to replay unexisting variables in VarInfo"
    dist = getdist(vi, vns[1])
    return if istrans(vi, vns[1])
        Bijectors.invlink(dist, reconstruct(dist, getval(vi, vns), length(vns)))
    else
        reconstruct(dist, getval(vi, vns), length(vns))
    end
end

"""
    getindex(vi::VarInfo, spl::Union{SampleFromPrior, Sampler})

Return the current value(s) of the random variables sampled by `spl` in `vi`.

The value(s) may or may not be transformed to Euclidean space.
"""
getindex(vi::AbstractVarInfo, spl::SampleFromPrior) = copy(getall(vi))
getindex(vi::AbstractVarInfo, spl::SampleFromUniform) = copy(getall(vi))
getindex(vi::UntypedVarInfo, spl::Sampler) = copy(getval(vi, _getranges(vi, spl)))
function getindex(vi::TypedVarInfo, spl::Sampler)
    # Gets the ranges as a NamedTuple
    ranges = _getranges(vi, spl)
    # Calling getfield(ranges, f) gives all the indices in `vals` of the `vn`s with symbol `f` sampled by `spl` in `vi`
    return vcat(_getindex(vi.metadata, ranges)...)
end
# Recursively builds a tuple of the `vals` of all the symbols
@generated function _getindex(metadata, ranges::NamedTuple{names}) where {names}
    expr = Expr(:tuple)
    for f in names
        push!(expr.args, :(metadata.$f.vals[ranges.$f]))
    end
    return expr
end

"""
    setindex!(vi::VarInfo, val, vn::VarName)

Set the current value(s) of the random variable `vn` in `vi` to `val`.

The value(s) may or may not be transformed to Euclidean space.
"""
setindex!(vi::AbstractVarInfo, val, vn::VarName) = (setval!(vi, val, vn); return vi)
function BangBang.setindex!!(vi::AbstractVarInfo, val, vn::VarName)
    return (setindex!(vi, val, vn); return vi)
end

"""
    setindex!(vi::VarInfo, val, spl::Union{SampleFromPrior, Sampler})

Set the current value(s) of the random variables sampled by `spl` in `vi` to `val`.

The value(s) may or may not be transformed to Euclidean space.
"""
setindex!(vi::AbstractVarInfo, val, spl::SampleFromPrior) = setall!(vi, val)
setindex!(vi::UntypedVarInfo, val, spl::Sampler) = setval!(vi, val, _getranges(vi, spl))
function setindex!(vi::TypedVarInfo, val, spl::Sampler)
    # Gets a `NamedTuple` mapping each symbol to the indices in the symbol's `vals` field sampled from the sampler `spl`
    ranges = _getranges(vi, spl)
    _setindex!(vi.metadata, val, ranges)
    return nothing
end

function BangBang.setindex!!(vi::AbstractVarInfo, val, spl::AbstractSampler)
    setindex!(vi, val, spl)
    return vi
end

# Recursively writes the entries of `val` to the `vals` fields of all the symbols as if they were a contiguous vector.
@generated function _setindex!(metadata, val, ranges::NamedTuple{names}) where {names}
    expr = Expr(:block)
    offset = :(0)
    for f in names
        f_vals = :(metadata.$f.vals)
        f_range = :(ranges.$f)
        start = :($offset + 1)
        len = :(length($f_range))
        finish = :($offset + $len)
        push!(expr.args, :(@views $f_vals[$f_range] .= val[($start):($finish)]))
        offset = :($offset + $len)
    end
    return expr
end

"""
    tonamedtuple(vi::VarInfo)

Convert a `vi` into a `NamedTuple` where each variable symbol maps to the values and
indexing string of the variable.

For example, a model that had a vector of vector-valued
variables `x` would return

```julia
(x = ([1.5, 2.0], [3.0, 1.0], ["x[1]", "x[2]"]), )
```
"""
function tonamedtuple(vi::VarInfo)
    return tonamedtuple(vi.metadata, vi)
end
@generated function tonamedtuple(metadata::NamedTuple{names}, vi::VarInfo) where {names}
    length(names) === 0 && return :(NamedTuple())
    expr = Expr(:tuple)
    map(names) do f
        push!(
            expr.args,
            Expr(:(=), f, :(getindex.(Ref(vi), metadata.$f.vns), string.(metadata.$f.vns))),
        )
    end
    return expr
end

@inline function findvns(vi, f_vns)
    if length(f_vns) == 0
        throw("Unidentified error, please report this error in an issue.")
    end
    return map(vn -> vi[vn], f_vns)
end

function Base.eltype(vi::AbstractVarInfo, spl::Union{AbstractSampler,SampleFromPrior})
    return eltype(Core.Compiler.return_type(getindex, Tuple{typeof(vi),typeof(spl)}))
end

"""
    haskey(vi::VarInfo, vn::VarName)

Check whether `vn` has been sampled in `vi`.
"""
haskey(vi::VarInfo, vn::VarName) = haskey(getmetadata(vi, vn).idcs, vn)
function haskey(vi::TypedVarInfo, vn::VarName)
    metadata = vi.metadata
    Tmeta = typeof(metadata)
    return getsym(vn) in fieldnames(Tmeta) && haskey(getmetadata(vi, vn).idcs, vn)
end

function Base.show(io::IO, ::MIME"text/plain", vi::UntypedVarInfo)
    vi_str = """
    /=======================================================================
    | VarInfo
    |-----------------------------------------------------------------------
    | Varnames  :   $(string(vi.metadata.vns))
    | Range     :   $(vi.metadata.ranges)
    | Vals      :   $(vi.metadata.vals)
    | GIDs      :   $(vi.metadata.gids)
    | Orders    :   $(vi.metadata.orders)
    | Logp      :   $(getlogp(vi))
    | #produce  :   $(get_num_produce(vi))
    | flags     :   $(vi.metadata.flags)
    \\=======================================================================
    """
    return print(io, vi_str)
end

const _MAX_VARS_SHOWN = 4

function _show_varnames(io::IO, vi)
    md = vi.metadata
    vns = md.vns

    vns_by_name = Dict{Symbol,Vector{VarName}}()
    for vn in vns
        group = get!(() -> Vector{VarName}(), vns_by_name, getsym(vn))
        push!(group, vn)
    end

    L = length(vns_by_name)
    if L == 0
        print(io, "0 variables, dimension 0")
    else
        (L == 1) ? print(io, "1 variable (") : print(io, L, " variables (")
        join(io, Iterators.take(keys(vns_by_name), _MAX_VARS_SHOWN), ", ")
        (L > _MAX_VARS_SHOWN) && print(io, ", ...")
        print(io, "), dimension ", length(md.vals))
    end
end

function Base.show(io::IO, vi::UntypedVarInfo)
    print(io, "VarInfo (")
    _show_varnames(io, vi)
    print(io, "; logp: ", round(getlogp(vi); digits=3))
    return print(io, ")")
end

"""
    push!!(vi::VarInfo, vn::VarName, r, dist::Distribution)

Push a new random variable `vn` with a sampled value `r` from a distribution `dist` to
the `VarInfo` `vi`, mutating if it makes sense.
"""
function BangBang.push!!(vi::AbstractVarInfo, vn::VarName, r, dist::Distribution)
    return BangBang.push!!(vi, vn, r, dist, Set{Selector}([]))
end

"""
    push!!(vi::VarInfo, vn::VarName, r, dist::Distribution, spl::AbstractSampler)

Push a new random variable `vn` with a sampled value `r` sampled with a sampler `spl`
from a distribution `dist` to `VarInfo` `vi`, if it makes sense.

The sampler is passed here to invalidate its cache where defined.
"""
function BangBang.push!!(
    vi::AbstractVarInfo, vn::VarName, r, dist::Distribution, spl::Sampler
)
    return BangBang.push!!(vi, vn, r, dist, spl.selector)
end
function BangBang.push!!(
    vi::AbstractVarInfo, vn::VarName, r, dist::Distribution, spl::AbstractSampler
)
    return BangBang.push!!(vi, vn, r, dist)
end

"""
    push!!(vi::VarInfo, vn::VarName, r, dist::Distribution, gid::Selector)

Push a new random variable `vn` with a sampled value `r` sampled with a sampler of
selector `gid` from a distribution `dist` to `VarInfo` `vi`.
"""
function BangBang.push!!(
    vi::AbstractVarInfo, vn::VarName, r, dist::Distribution, gid::Selector
)
    return BangBang.push!!(vi, vn, r, dist, Set([gid]))
end
function BangBang.push!!(
    vi::VarInfo, vn::VarName, r, dist::Distribution, gidset::Set{Selector}
)
    if vi isa UntypedVarInfo
        @assert ~(vn in keys(vi)) "[push!!] attempt to add an exisitng variable $(getsym(vn)) ($(vn)) to VarInfo (keys=$(keys(vi))) with dist=$dist, gid=$gidset"
    elseif vi isa TypedVarInfo
        @assert ~(haskey(vi, vn)) "[push!!] attempt to add an exisitng variable $(getsym(vn)) ($(vn)) to TypedVarInfo of syms $(syms(vi)) with dist=$dist, gid=$gidset"
    end

    val = vectorize(dist, r)

    meta = getmetadata(vi, vn)
    meta.idcs[vn] = length(meta.idcs) + 1
    push!(meta.vns, vn)
    l = length(meta.vals)
    n = length(val)
    push!(meta.ranges, (l + 1):(l + n))
    append!(meta.vals, val)
    push!(meta.dists, dist)
    push!(meta.gids, gidset)
    push!(meta.orders, get_num_produce(vi))
    push!(meta.flags["del"], false)
    push!(meta.flags["trans"], false)

    return vi
end

"""
    setorder!(vi::VarInfo, vn::VarName, index::Int)

Set the `order` of `vn` in `vi` to `index`, where `order` is the number of `observe
statements run before sampling `vn`.
"""
function setorder!(vi::VarInfo, vn::VarName, index::Int)
    metadata = getmetadata(vi, vn)
    if metadata.orders[metadata.idcs[vn]] != index
        metadata.orders[metadata.idcs[vn]] = index
    end
    return vi
end

#######################################
# Rand & replaying method for VarInfo #
#######################################

"""
    is_flagged(vi::VarInfo, vn::VarName, flag::String)

Check whether `vn` has a true value for `flag` in `vi`.
"""
function is_flagged(vi::VarInfo, vn::VarName, flag::String)
    return getmetadata(vi, vn).flags[flag][getidx(vi, vn)]
end

"""
    unset_flag!(vi::VarInfo, vn::VarName, flag::String)

Set `vn`'s value for `flag` to `false` in `vi`.
"""
function unset_flag!(vi::VarInfo, vn::VarName, flag::String)
    getmetadata(vi, vn).flags[flag][getidx(vi, vn)] = false
    return vi
end

"""
    set_retained_vns_del_by_spl!(vi::VarInfo, spl::Sampler)

Set the `"del"` flag of variables in `vi` with `order > vi.num_produce[]` to `true`.
"""
function set_retained_vns_del_by_spl!(vi::UntypedVarInfo, spl::Sampler)
    # Get the indices of `vns` that belong to `spl` as a vector
    gidcs = _getidcs(vi, spl)
    if get_num_produce(vi) == 0
        for i in length(gidcs):-1:1
            vi.metadata.flags["del"][gidcs[i]] = true
        end
    else
        for i in 1:length(vi.orders)
            if i in gidcs && vi.orders[i] > get_num_produce(vi)
                vi.metadata.flags["del"][i] = true
            end
        end
    end
    return nothing
end
function set_retained_vns_del_by_spl!(vi::TypedVarInfo, spl::Sampler)
    # Get the indices of `vns` that belong to `spl` as a NamedTuple, one entry for each symbol
    gidcs = _getidcs(vi, spl)
    return _set_retained_vns_del_by_spl!(vi.metadata, gidcs, get_num_produce(vi))
end
@generated function _set_retained_vns_del_by_spl!(
    metadata, gidcs::NamedTuple{names}, num_produce
) where {names}
    expr = Expr(:block)
    for f in names
        f_gidcs = :(gidcs.$f)
        f_orders = :(metadata.$f.orders)
        f_flags = :(metadata.$f.flags)
        push!(
            expr.args,
            quote
                # Set the flag for variables with symbol `f`
                if num_produce == 0
                    for i in length($f_gidcs):-1:1
                        $f_flags["del"][$f_gidcs[i]] = true
                    end
                else
                    for i in 1:length($f_orders)
                        if i in $f_gidcs && $f_orders[i] > num_produce
                            $f_flags["del"][i] = true
                        end
                    end
                end
            end,
        )
    end
    return expr
end

"""
    updategid!(vi::VarInfo, vn::VarName, spl::Sampler)

Set `vn`'s `gid` to `Set([spl.selector])`, if `vn` does not have a sampler selector linked
and `vn`'s symbol is in the space of `spl`.
"""
function updategid!(vi::AbstractVarInfo, vn::VarName, spl::Sampler)
    if inspace(vn, getspace(spl))
        setgid!(vi, spl.selector, vn)
    end
end

# TODO: Maybe rename or something?
"""
    _apply!(kernel!, vi::AbstractVarInfo, values, keys)

Calls `kernel!(vi, vn, values, keys)` for every `vn` in `vi`.
"""
function _apply!(kernel!, vi::AbstractVarInfo, values, keys)
    keys_strings = map(string, collectmaybe(keys))
    num_indices_seen = 0

    for vn in Base.keys(vi)
        indices_found = kernel!(vi, vn, values, keys_strings)
        if indices_found !== nothing
            num_indices_seen += length(indices_found)
        end
    end

    if length(keys) > num_indices_seen
        # Some keys have not been seen, i.e. attempted to set variables which
        # we were not able to locate in `vi`.
        # Find the ones we missed so we can warn the user.
        unused_keys = _find_missing_keys(vi, keys_strings)
        @warn "the following keys were not found in `vi`, and thus `kernel!` was not applied to these: $(unused_keys)"
    end

    return vi
end

function _apply!(kernel!, vi::TypedVarInfo, values, keys)
    return _typed_apply!(kernel!, vi, vi.metadata, values, collectmaybe(keys))
end

@generated function _typed_apply!(
    kernel!, vi::TypedVarInfo, metadata::NamedTuple{names}, values, keys
) where {names}
    updates = map(names) do n
        quote
            for vn in metadata.$n.vns
                indices_found = kernel!(vi, vn, values, keys_strings)
                if indices_found !== nothing
                    num_indices_seen += length(indices_found)
                end
            end
        end
    end

    return quote
        keys_strings = map(string, keys)
        num_indices_seen = 0

        $(updates...)

        if length(keys) > num_indices_seen
            # Some keys have not been seen, i.e. attempted to set variables which
            # we were not able to locate in `vi`.
            # Find the ones we missed so we can warn the user.
            unused_keys = _find_missing_keys(vi, keys_strings)
            @warn "the following keys were not found in `vi`, and thus `kernel!` was not applied to these: $(unused_keys)"
        end

        return vi
    end
end

function _find_missing_keys(vi::AbstractVarInfo, keys)
    string_vns = map(string, collectmaybe(Base.keys(vi)))
    # If `key` isn't subsumed by any element of `string_vns`, it is not present in `vi`.
    missing_keys = filter(keys) do key
        !any(Base.Fix2(subsumes_string, key), string_vns)
    end

    return missing_keys
end

"""
    setval!(vi::AbstractVarInfo, x)
    setval!(vi::AbstractVarInfo, values, keys)
    setval!(vi::AbstractVarInfo, chains::AbstractChains, sample_idx::Int, chain_idx::Int)

Set the values in `vi` to the provided values and leave those which are not present in
`x` or `chains` unchanged.

## Notes
This is rather limited for two reasons:
1. It uses `subsumes_string(string(vn), map(string, keys))` under the hood,
   and therefore suffers from the same limitations as [`subsumes_string`](@ref).
2. It will set every `vn` present in `keys`. It will NOT however
   set every `k` present in `keys`. This means that if `vn == [m[1], m[2]]`,
   representing some variable `m`, calling `setval!(vi, (m = [1.0, 2.0]))` will
   be a no-op since it will try to find `m[1]` and `m[2]` in `keys((m = [1.0, 2.0]))`.

## Example
```jldoctest
julia> using DynamicPPL, Distributions, StableRNGs

julia> @model function demo(x)
           m ~ Normal()
           for i in eachindex(x)
               x[i] ~ Normal(m, 1)
           end
       end;

julia> rng = StableRNG(42);

julia> m = demo([missing]);

julia> var_info = DynamicPPL.VarInfo(rng, m);

julia> var_info[@varname(m)]
-0.6702516921145671

julia> var_info[@varname(x[1])]
-0.22312984965118443

julia> DynamicPPL.setval!(var_info, (m = 100.0, )); # set `m` and and keep `x[1]`

julia> var_info[@varname(m)] # [✓] changed
100.0

julia> var_info[@varname(x[1])] # [✓] unchanged
-0.22312984965118443

julia> m(rng, var_info); # rerun model

julia> var_info[@varname(m)] # [✓] unchanged
100.0

julia> var_info[@varname(x[1])] # [✓] unchanged
-0.22312984965118443
```
"""
setval!(vi::AbstractVarInfo, x) = setval!(vi, values(x), keys(x))
setval!(vi::AbstractVarInfo, values, keys) = _apply!(_setval_kernel!, vi, values, keys)
function setval!(
    vi::AbstractVarInfo, chains::AbstractChains, sample_idx::Int, chain_idx::Int
)
    return setval!(vi, chains.value[sample_idx, :, chain_idx], keys(chains))
end

function _setval_kernel!(vi::VarInfo, vn::VarName, values, keys)
    indices = findall(Base.Fix1(subsumes_string, string(vn)), keys)
    if !isempty(indices)
        val = reduce(vcat, values[indices])
        setval!(vi, val, vn)
        settrans!(vi, false, vn)
    end

    return indices
end

"""
    setval_and_resample!(vi::AbstractVarInfo, x)
    setval_and_resample!(vi::AbstractVarInfo, values, keys)
    setval_and_resample!(vi::AbstractVarInfo, chains::AbstractChains, sample_idx, chain_idx)

Set the values in `vi` to the provided values and those which are not present
in `x` or `chains` to *be* resampled.

Note that this does *not* resample the values not provided! It will call `setflag!(vi, vn, "del")`
for variables `vn` for which no values are provided, which means that the next time we call `model(vi)` these
variables will be resampled.

## Note
- This suffers from the same limitations as [`setval!`](@ref). See `setval!` for more info.

## Example
```jldoctest
julia> using DynamicPPL, Distributions, StableRNGs

julia> @model function demo(x)
           m ~ Normal()
           for i in eachindex(x)
               x[i] ~ Normal(m, 1)
           end
       end;

julia> rng = StableRNG(42);

julia> m = demo([missing]);

julia> var_info = DynamicPPL.VarInfo(rng, m);

julia> var_info[@varname(m)]
-0.6702516921145671

julia> var_info[@varname(x[1])]
-0.22312984965118443

julia> DynamicPPL.setval_and_resample!(var_info, (m = 100.0, )); # set `m` and ready `x[1]` for resampling

julia> var_info[@varname(m)] # [✓] changed
100.0

julia> var_info[@varname(x[1])] # [✓] unchanged
-0.22312984965118443

julia> m(rng, var_info); # sample `x[1]` conditioned on `m = 100.0`

julia> var_info[@varname(m)] # [✓] unchanged
100.0

julia> var_info[@varname(x[1])] # [✓] changed
101.37363069798343
```

## See also
- [`setval!`](@ref)
"""
function setval_and_resample!(vi::AbstractVarInfo, x)
    return setval_and_resample!(vi, values(x), keys(x))
end
function setval_and_resample!(vi::AbstractVarInfo, values, keys)
    return _apply!(_setval_and_resample_kernel!, vi, values, keys)
end
function setval_and_resample!(
    vi::AbstractVarInfo, chains::AbstractChains, sample_idx::Int, chain_idx::Int
)
    return setval_and_resample!(vi, chains.value[sample_idx, :, chain_idx], keys(chains))
end

function _setval_and_resample_kernel!(vi::VarInfo, vn::VarName, values, keys)
    indices = findall(Base.Fix1(subsumes_string, string(vn)), keys)
    if !isempty(indices)
        val = reduce(vcat, values[indices])
        setval!(vi, val, vn)
        settrans!(vi, false, vn)
    else
        # Ensures that we'll resample the variable corresponding to `vn` if we run
        # the model on `vi` again.
        set_flag!(vi, vn, "del")
    end

    return indices
end

"""
    values_as(vi::AbstractVarInfo)
"""
values_as(vi::VarInfo) = vi.metadata

"""
    values_as(vi::AbstractVarInfo, ::Type{NamedTuple})
    values_as(vi::AbstractVarInfo, ::Type{Dict})

Return values in `vi` as the specified type.
"""
function values_as(vi::UntypedVarInfo, ::Type{NamedTuple})
    iter = values_from_metadata(vi.metadata)
    return NamedTuple(map(p -> Symbol(p.first) => p.second, iter))
end
values_as(vi::UntypedVarInfo, ::Type{Dict}) = Dict(values_from_metadata(vi.metadata))

function values_as(vi::VarInfo{<:NamedTuple{names}}, ::Type{NamedTuple}) where {names}
    iter = Iterators.flatten(values_from_metadata(getfield(vi.metadata, n)) for n in names)
    return NamedTuple(map(p -> Symbol(p.first) => p.second, iter))
end

function values_as(vi::VarInfo{<:NamedTuple{names}}, ::Type{Dict}) where {names}
    iter = Iterators.flatten(values_from_metadata(getfield(vi.metadata, n)) for n in names)
    return Dict(iter)
end

function values_from_metadata(md::Metadata)
    return (
        vn => reconstruct(md.dists[md.idcs[vn]], md.vals[md.ranges[md.idcs[vn]]]) for
        vn in md.vns
    )
end

max_stars_repo_name: TuringLang/MicroPPL

max_issues_repo_name: TuringLang/MicroPPL

max_forks_repo_name: TuringLang/MicroPPL


text:
baremodule Reagents

export ⨟, compose, opcompose

import Base

using CompositionsBase: ⨟, compose, opcompose

abstract type Reagent end

abstract type Ref{T} end

abstract type Failure end
struct Block <: Failure end
struct Retry <: Failure end

# The original implementation by Turon (2012) combines what is called `Reactor`
# and `Reagent` here simply as reagent.  This is structurally identical to how
# Transducers.jl splits transducers (`Reagent` here) and reducing functions
# (`Reactor` here). It makes implemneting the pairting combinator `&` (`Both`)
# easier (which is similar to `Transducers.TeeZip`).
struct Reactor{R<:Reagent,C}
    reagent::R
    continuation::C
end

# TODO: Is it better to provide factory functions like `updating(f, ref)`
# instead of exposing types (constructors) which is rather leaky?

struct Identity <: Reagent end

struct Update{F,R<:Ref} <: Reagent
    f::F
    ref::R
end

struct Read{R<:Ref} <: Reagent
    ref::R
end

struct CAS{T,R<:Ref{T}} <: Reagent
    ref::R
    old::T
    new::T
end

struct Return{T} <: Reagent
    value::T
end

struct Computed{F} <: Reagent
    f::F
end

struct Map{F} <: Reagent
    f::F
end

struct WithNack{F} <: Reagent
    f::F
end

struct Until{F,R<:Reagent} <: Reagent
    f::F
    reagent::R
end

struct PostCommit{F} <: Reagent
    f::F
end

function channel end
function dissolve end
function clear! end

function try! end
function trysync! end
function trysyncing end

module Internal

using Accessors: @set
using CompositionsBase: ⨟

using ..Reagents:
    Block,
    CAS,
    Computed,
    Failure,
    Identity,
    Map,
    PostCommit,
    Reactor,
    Read,
    Reagent,
    Reagents,
    Retry,
    Return,
    Until,
    Update,
    WithNack

include("utils.jl")
include("immutablelists.jl")
include("tracing.jl")
include("anchors.jl")
include("bags.jl")
include("core.jl")
include("reactions.jl")
include("react.jl")
include("combinators.jl")
include("computational.jl")
include("refs.jl")
include("channels.jl")
include("dissolve.jl")

end  # module Internal

# For defining the docstrings
const (|) = Internal.Base.:|
const (&) = Internal.Base.:&

module __Reagents_API__
using ..Reagents
export ⨟, compose, opcompose, Reagents

module DefaultNames end

for n in names(Reagents; all = true)
    try
        getfield(DefaultNames, n)
        continue
    catch
    end
    match(r"^[a-z]"i, string(n)) === nothing && continue
    x = try
        getfield(Reagents, n)
    catch
        continue
    end
    x isa Union{Type,Function} || continue
    @eval begin
        $n = $x
        export $n
    end
end

end  # module __Reagents_API__

"""
    Reagents.*

A module that re-exports a subset of Reagents public API at top-level.

Use `using Reagents.*` to import all public APIs. Do not use this inside a
package. The set of exported names is not the part of stable API. For example,
if a name collistion with `Base` or important packages are recognized, the
corresponding name may be removed in the next realease, without incrementing
the leading non-zero version number.
"""
const * = __Reagents_API__

Internal.define_docstrings()

end  # baremodule Reagents

ext: jl


text:
function label_propagation{V,T<:Real}(graph::AbstractGraph{V};
                                      weights::Vector{T} = Array(Float64, 0),
                                      proximity::Vector{T} = Array(Float64, 0),
                                      initial::Dict{V,Int} = Dict{V,Int}(),
                                      fixed::Set{V} = Set{V}())
    !is_directed(graph) || error("graph must be undirected.")
    if !isempty(proximity)
        @graph_requires graph edge_map incidence_list
        if length(proximity) != num_edges(graph)
            error("Invalid proximity vector length")
        elseif minimum(proximity) < 0 || maximum(proximity) > 1
            error("Proximity must between 0 and 1")
        end
        mix_label_propagation!(graph, _init(graph, initial, fixed)..., proximity)
    else
        if !isempty(weights)
            @graph_requires graph edge_map incidence_list
            if length(weights) != num_edges(graph)
                error("Invalid weight vector length")
            elseif minimum(weights) < 0
                error("Weights must be non-negative")
            end
            label_propagation!(graph, _init(graph, initial, fixed)..., weights)
        else
            label_propagation!(graph, _init(graph, initial, fixed)...)
        end
    end
end

function _init{V}(graph::AbstractGraph{V},
                  initial::Dict{V,Int} = Dict{V,Int}(), # initial vertex label map
                  fixed::Set{V} = Set{V}())             # initial label fixed nodes

    @graph_requires graph vertex_list vertex_map

    N = num_vertices(graph)
    membership = zeros(Int, N)
    unfixed = V[]

    # Do some initial checks
    if !isempty(fixed) && isempty(initial)
        warn("Ignoring fixed vertices as no initial labeling given")
    end

    if !isempty(initial)
        for v in vertices(graph)
            v_idx = vertex_index(v, graph)
            if haskey(initial, v)
                membership[v_idx] = initial[v]
            else
                membership[v_idx] = 0
            end
        end

        for v in vertices(graph)
            if in(v, fixed)
                if membership[vertex_index(v, graph)] < 1
                    warn("Fixed nodes cannot be unlabeled, ignoring them")
                    push!(unfixed, v)
                end
            else
                push!(unfixed, v)
            end
        end

        i = maximum(membership)
        i <= N || error("elements of the initial labeling vector must be between 1 and |V|")
        i > 0 || error("at least one vertex must be labeled in the initial labeling")
    else
        membership = [1:N]
        unfixed = collect(vertices(graph))
    end

    membership, unfixed
end

function label_propagation!{V,T<:Real}(graph::AbstractGraph{V},
                                       membership::Vector{Int},
                                       unfixed::Vector{V},
                                       weights::Vector{T})

    @graph_requires graph vertex_map edge_map incidence_list

    label_counters = zeros(T, num_vertices(graph))

    running = true
    while running
        running = false

        # Shuffle the node ordering vector
        X = shuffle(unfixed)

        # In the prescribed order, loop over the vertices and reassign labels
        for v in X
            v_idx = vertex_index(v, graph)

            # Clear dominant_labels and nonzero_labels
            dominant_labels = Int[]
            nonzero_labels = Int[]

            # recount
            max_count = zero(T)

            for e in out_edges(v, graph)
                k = membership[vertex_index(target(e, graph), graph)]

                # skip if it has no label yet
                k != 0 || continue

                was_zero = label_counters[k] == zero(T)
                label_counters[k] += weights[edge_index(e, graph)]
                if was_zero && label_counters[k] != zero(T)
                    # counter just became nonzero
                    push!(nonzero_labels, k)
                end
                if max_count < label_counters[k]
                    max_count = label_counters[k]
                    resize!(dominant_labels, 1)
                    dominant_labels[1] = k
                elseif max_count == label_counters[k]
                    push!(dominant_labels, k)
                end
            end

            if !isempty(dominant_labels)

                # Select randomly from the dominant labels
                k = dominant_labels[rand(1:length(dominant_labels))]

                # Check if the current label of the node is also dominant
                if label_counters[membership[v_idx]] != max_count
                    # Nope, we need at least one more iteration
                    running = true
                end

                # Update label of the current node
                membership[v_idx] = k
            end
            # Clear the nonzero elements in label_counters
            for i in nonzero_labels
                label_counters[i] = zero(T)
            end
        end
    end
    permute_labels!(membership)
    membership
end

function label_propagation!{V}(graph::AbstractGraph{V},
                               membership::Vector{Int},
                               unfixed::Vector{V})

    @graph_requires graph vertex_map adjacency_list

    label_counters = zeros(Int, num_vertices(graph))

    running = true
    while running
        running = false

        # Shuffle the node ordering vector
        X = shuffle(unfixed)

        # In the prescribed order, loop over the vertices and reassign labels
        for u in X
            u_idx = vertex_index(u, graph)

            # Clear dominant_labels and nonzero_labels
            dominant_labels = Int[]
            nonzero_labels = Int[]

            # recount
            max_count = 0

            for v in out_neighbors(u, graph)
                k = membership[vertex_index(v, graph)]

                # skip if it has no label yet
                k != 0 || continue

                label_counters[k] += 1

                # counter just became nonzero
                label_counters[k] != 1 || push!(nonzero_labels, k)

                if max_count < label_counters[k]
                    max_count = label_counters[k]
                    resize!(dominant_labels, 1)
                    dominant_labels[1] = k
                elseif max_count == label_counters[k]
                    push!(dominant_labels, k)
                end
            end

            if !isempty(dominant_labels)

                # Select randomly from the dominant labels
                k = dominant_labels[rand(1:length(dominant_labels))]

                # Check if the current label of the node is also dominant
                if label_counters[membership[u_idx]] != max_count
                    # Nope, we need at least one more iteration
                    running = true
                end

                # Update label of the current node
                membership[u_idx] = k
            end
            # Clear the nonzero elements in label_counters
            for i in nonzero_labels
                label_counters[i] = 0
            end
        end
    end
    permute_labels!(membership)
    membership
end

function mix_label_propagation!{V,T<:Real}(graph::AbstractGraph{V},
                                           membership::Vector{Int},
                                           unfixed::Vector{V},
                                           proximity::Vector{T})
    @graph_requires graph edge_map incidence_list vertex_map

    N = num_vertices(graph)
    label_counters = zeros(T, N)
    proximity_sum = zeros(T, N)
    similarity = Array(T, 0)

    running = true
    while running
        running = false

        # Shuffle the node ordering vector
        X = shuffle(unfixed)

        # In the prescribed order, loop over the vertices and reassign labels
        for v in X
            v_idx = vertex_index(v, graph)

            # Clear dominant_labels and nonzero_labels
            dominant_labels = Int[]
            nonzero_labels = Int[]

            # recount
            max_count = zero(T)

            for e in out_edges(v, graph)
                e_idx = edge_index(e, graph)
                k = membership[vertex_index(target(e, graph), graph)]

                # skip if it has no label yet
                k != 0 || continue

                was_zero = label_counters[k] == zero(T)
                label_counters[k] += 1
                proximity_sum[k] += proximity[e_idx]
                if was_zero && label_counters[k] != zero(T)
                    # counter just became nonzero
                    push!(nonzero_labels, k)
                end
                if max_count < label_counters[k]
                    max_count = label_counters[k]
                    resize!(dominant_labels, 1)
                    resize!(similarity, 1)
                    dominant_labels[1] = k
                    similarity[1] = proximity_sum[k]
                elseif max_count == label_counters[k]
                    push!(dominant_labels, k)
                    push!(similarity, proximity_sum[k])
                end
            end

            if !isempty(dominant_labels)

                max_similarity = maximum(similarity)
                candidate_labels = dominant_labels[similarity.>=max_similarity]

                # Select randomly from the dominant labels

                k = candidate_labels[rand(1:length(candidate_labels))]

                # Check if the current label of the node is also dominant
                if label_counters[membership[v_idx]] != max_count
                    # Nope, we need at least one more iteration
                    running = true
                end

                # Update label of the current node
                membership[v_idx] = k
            end
            # Clear the nonzero elements in label_counters
            for i in nonzero_labels
                label_counters[i] = zero(T)
                proximity_sum[i] = zero(T)
            end
        end
    end
    permute_labels!(membership)
    membership
end

function permute_labels!(membership::Vector{Int})
    N = length(membership)
    if maximum(membership) > N || minimum(membership) < 1
        error("Label must between 1 and |V|")
    end
    label_counters = zeros(Int, N)
    j = 1
    for i=1:length(membership)
        k = membership[i]
        if k >= 1
            if label_counters[k] == 0
                # We have seen this label for the first time
                label_counters[k] = j
                k = j
                j += 1
            else
                k = label_counters[k]
            end
        end
        membership[i] = k
    end
end


text:
using Revise

import KRLS

using Plots

filter_param = KRLS.FilterParam(6, 10.0, 15)

spat_data = KRLS.load_sample("Data/Gaussian diff/Gaussian_diff_spat.mat")
(r, c) = size(spat_data)
raw_data = spat_data[:, 1]
vx = KRLS.W * raw_data
vy = spat_data[:, 2]
predicted_data = zeros((r, c))
predicted_data[:, 1] = raw_data

krlsfilter = KRLS.KRLSFilter(vx, vy, '3', filter_param)

println("Something")

vx_arr = KRLS.W * spat_data

@time for k = 1:c - 1
    KRLS.update(vx_arr[:, k], spat_data[:, k + 1], krlsfilter)
    predicted_data[:, k + 1] = KRLS.predict(vx_arr[:, k + 1], krlsfilter)
end

plot(spat_data[:, :]', layout = (3,2), xlabel="Time", ylabel="Temp", size=(900, 700), legend = :none)
plot!(predicted_data[:, :]', layout=(3,2))


text:
using QPSParser
using QPSParser: MPSDescription
using Test, LinearAlgebra, SparseArrays

import InteractiveUtils: subtypes

eye(n,m) = Matrix{Float64}(I, n, m)
eye(n) = eye(n,n)

speye(n,m) = SparseMatrixCSC{Float64,Int}(I,n,m)
speye(n) = speye(n,n)

include("mpsdesc.jl")
include("parseqps.jl")


text:
"""
## fkdiag!

This subroutine returns the elastic dee matrix for ih=3 (plane strain),
ih=4 (axisymmetry or plane strain elastoplasticity) or ih=6
(three dimensions).

### Method
```julia
fkdiag!(kdiag, g)
```

### Arguments
```julia
* kdiag::Vector{Int}      : Bandwidth vector (Updated)
* g::Vector{Int}          : Element steering vector
```
"""
function fkdiag!(kdiag::Vector{Int}, g::Vector{Int})
  idof = size(g, 1)
  for i in 1:idof
    imp1 = 1
    if g[i] !== 0
      for j in 1:idof
        if g[j] !== 0
          im = g[i] - g[j] + 1
          if im > imp1
            imp1 = im
          end
        end
        k = g[i]
        if imp1 > kdiag[k]
          kdiag[k] = imp1
        end
      end
    end
  end
end


text:
"""
Caputo sense fractional derivative algorithms, please refer to [Caputo derivative](https://en.wikipedia.org/wiki/Fractional_calculus#Caputo_fractional_derivative) for more details.
"""
abstract type Caputo <: FracDiffAlg end

"""
# Caputo sense fractional derivative.

    fracdiff(f::Function, α, start_point, end_point, step_size, ::Caputo)

### Example:

```julia-repl
julia> fracdiff(x->x^5, 0.5, 0, 2.5, 0.0001, CaputoDirect())
```

Returns a tuple (**result**, **error**), which means the value of this derivative is 141.59714979764541, and the error estimate is 1.1532243848672914e-6.

Refer to [Caputo derivative](https://en.wikipedia.org/wiki/Fractional_calculus#Caputo_fractional_derivative)

!!! note
    ```Caputo Direct``` algorithms belong to direct computing, precise are ensured, but maybe cause more memory allocation and take more compilation time.

Using the direct mathematic expression:

```math
^CD_t^α=\\frac{1}{\\Gamma(n-α)}\\int_0^t\\frac{f^{(n)}(τ)}{(t-τ)^{α+1-n}}
```

As for the derivative inside the integral, we use the **Complex Step Differentiation** to obtain a more accurate value.
"""
struct CaputoDirect <: Caputo end

"""
# Caputo sense fractional derivative with first derivative known.

    fracdiff(fd, α, start_point, end_point, CaputoDirectFirstDiffKnown())

If the first order derivative of a function is already known, we can use this method to compute the fractional order derivative more precisely.

The inout function should be the first order derivative of a function

### Example

```julia-repl
julia> fracdiff(x->5*x^4, 0.5, 0, 2.5, CaputoDirectFirstDiffKnown())
```
Return the semi-derivative of ``f(x)=x^5`` at ``x=2.5``.

Compared with **CaputoDirect** method, this method don't need to specify step size, more precision are guaranteed.
"""
struct CaputoDirectFirstDiffKnown <: Caputo end

"""
# Caputo sense fractional derivative with first and second derivative known.

    fracdiff(fd1, fd2, α, start_point, end_point, Caputo_Direct_First_Second_Diff_Known)

If the first and second order derivative of a function is already known, we can use this method to compute the fractional order derivative more precisely.

### Example

```julia-repl
julia> fracdiff(x->5*x^4, x->20*x^3, 0.5, 0, 2.5, Caputo_Direct_First_Second_Diff_known())
```
Return the semi-derivative of ``f(x)=x^5`` at ``x=2.5``.

Compared with **CaputoDirect** method, this method don't need to specify step size, more precision are guaranteed.
"""
struct Caputo_Direct_First_Second_Diff_Known <: Caputo end

"""
# Caputo sense Piecewise algorithm

    fracdiff(f, α, end_point, h, CaputoPiecewise())

Using piecewise linear interpolation function to approximate input function and combining Caputo derivative then implement summation.

### Example

```julia-repl
julia> fracdiff(x->x^5, 0.5, 2.5, 0.001, CaputoPiecewise())
```

Return the fractional derivative of ``f(x)=x^5`` at point ``x=2.5``.

```tex
@article{LI20113352,
title = {Numerical approaches to fractional calculus and fractional ordinary differential equation},
author = {Changpin Li and An Chen and Junjie Ye},
}
```

"""
struct CaputoPiecewise <: Caputo end

"""
# Caputo sense Diethelm computation

    fracdiff(f, α, end_point, h, CaputoDiethelm())

Using quadrature weights(derived from product trapezoidal rule) to approximate the derivative.

### Example

```julia-repl
julia> fracdiff(x->x, 0.5, 1, 0.007, CaputoDiethelm())
1.128378318687192
```

!!! info "Scope"
    0 < α < 1

Diethelm's method for computingn Caputo sense fractional derivative using [product trapezoidal rule](https://en.wikipedia.org/wiki/Trapezoidal_rule) and [Richardsin extrapolation](https://en.wikipedia.org/wiki/Richardson_extrapolation)

```tex
@article{10.1093/imanum/17.3.479,
author = {DIETTELM, KAI},
title = "{Generalized compound quadrature formulae for finite-part integrals}",
journal = {IMA Journal of Numerical Analysis},
doi = {10.1093/imanum/17.3.479},
}
```

"""
struct CaputoDiethelm <: Caputo end

"""
# Caputo sense fractioal derivative with p-th order precision.

    fracdiff(f, α, t, p, CaputoHighPrecision())

Use the high precision algorithm to compute the Caputo sense fractional derivative.

The **p** here is the grade of precision.
"""
struct CaputoHighPrecision <: Caputo end

"""
# Caputo sense fractional derivative high order approximation.

!!! info
    By using ```CaputoHighOrder``` method, we can set ``0<\\alpha<2``

### Example

```julia-repl
julia> fracdiff(x->x, 0.5, 1, 0.01, 2, CaputoHighOrder())
1.1283791670955123
```

```tex
@article{Li2017HighOrderAT,
  title={High-Order Approximation to Caputo Derivatives and Caputo-type Advection–Diffusion Equations: Revisited},
  author={Changpin Li and Minhao Cai},
  journal={Numerical Functional Analysis and Optimization},
  year={2017},
  volume={38},
  pages={861 - 890}
}
```
"""
struct CaputoHighOrder <: Caputo end
################################################################
###                    Type definition done                  ###
################################################################

function fracdiff(f::Union{Function, Number}, α::Float64, start_point, end_point, h::Float64, ::CaputoDirect)
    #checks(f, α, start_point, end_point)
    typeof(f) <: Number ? (end_point == 0 ? 0 : f/sqrt(pi*end_point)) : nothing
    #Using complex step differentiation to calculate the first order differentiation
    g(τ) = imag(f(τ+1*im*h) ./ h) ./ ((end_point-τ) .^α)
    result = quadgk(g, start_point, end_point) ./gamma(1-α)
    return result
end

function fracdiff(f::FunctionAndNumber, α::Float64, start_point, end_point::AbstractArray, h::Float64, ::CaputoDirect)::Vector
    #=
    ResultArray = Float64[]
    for (_, value) in enumerate(end_point)
        append!(ResultArray, fracdiff(f, α, start_point, value, h, CaputoDirect()))
    end
    return ResultArray
    =#
    result = map(x->fracdiff(f, α, start_point, x, h, CaputoDirect())[1], end_point)
    return result
end

function fracdiff(f::FunctionAndNumber, α::Float64, end_point, h::Float64, ::CaputoPiecewise)
    typeof(f) <: Number ? (end_point == 0 ? (return 0) : (return f/sqrt(pi*end_point))) : nothing
    end_point == 0 ? (return 0) : nothing
    m = floor(Int, α)+1

    summation = zero(Float64)
    n = floor(Int, end_point/h)

    @fastmath @inbounds @simd for i ∈ 0:n
        summation += W₅(i, n, m, α)*first_order(f, i*h, h)
    end

    result = summation*h^(m-α)/gamma(m-α+2)
    return result
end
function W₅(i, n, m, α)
    if i == 0
        return n^(m-α)*(m-α+1-n) + (n-1)^(m-α+1)
    elseif i == n
        return 1
    else
        return (n-i-1)^(m-α+1) + (n-i+1)^(m-α+1) - 2*(n-i)^(m-α+1)
    end
end
# Deploy Complex Step Differentiation to compute the first order derivative.
first_order(f, point, h::Float64) = imag(f(point + im*h))/h

function fracdiff(f::FunctionAndNumber, α::Float64, end_point::AbstractArray, h::Float64, ::CaputoPiecewise)::Vector
    result = map(x->fracdiff(f, α, x, h, CaputoPiecewise()), end_point)
    return result
end

#=
Caputo Diethelm algorithm
=#
function fracdiff(f::FunctionAndNumber, α::Float64, end_point::Number, h::Float64, ::CaputoDiethelm)
    #checks(f, α, 0, end_point)
    typeof(f) <: Number ? (end_point == 0 ? 0 : f/sqrt(pi*end_point)) : nothing
    N = round(Int, end_point/h)
    result = zero(Float64)

    @fastmath @inbounds @simd for i ∈ 0:N
        result += quadweights(i, N, α)*(f(end_point-i*h) - f(0))
    end

    return result*h^(-α)/gamma(2-α)
end
function quadweights(n, N, α)
    if n == 0
        return 1
    elseif  0 < n < N
        return (n+1)^(1-α) - 2*n^(1-α) + (n-1)^(1-α)
    elseif n == N
        return (1-α)*N^(-α) - N^(1-α) + (N-1)^(1-α)
    end
end

function fracdiff(f, α::Float64, end_point::AbstractArray, h::Float64, ::CaputoDiethelm)::Vector
    result = map(x->fracdiff(f, α, x, h, CaputoDiethelm()), end_point)
    return result
end

#=
Caputo sense high precision algorithm
=#
function fracdiff(y::FunctionAndNumber, α, t, p::Integer, ::CaputoHighPrecision)
    h = t[2]-t[1]
    t = t[:]
    n = length(t)
    y = y.(t)
    q = ceil(Int, α)
    r::Int = max(p, q)
    R = reverse(Vandermonde(collect(0:(r-1)).*h))
    c = inv(R)*y[1:r]
    u = 0
    du = 0
    for i = 1:r
        u = u.+c[i]*t.^(i-1)
    end
    if q < r
        for i = (q+1):p
            du = du.+c[i]*t.^(i-1-α)*gamma(1)/gamma(i-α)
        end
    end

    v = y[:] - u[:]
    dv = fracdiff(v, α, t, p, GL_High_Precision())
    dy = dv[:] + du[:]
    return dy
end

function fracdiff(f::FunctionAndNumber, α, T, h, p::Integer, ::CaputoHighOrder)
    n::Int64 = round(Int, T/h)
    A = wj(n, p, α)
    B = fj(n, T, f)
    C = ones(1, n)
    hh = h^(-α)/(gamma(1-α))
    D = C*A*B*hh
    return D[1]
end

function wj(n::Int64, r::Int64, α)
    A=zeros(n, n+1)
    for iw=1:r-2
        for jw=1:iw+1
            A[iw, jw]=w(iw+1-jw, iw+1, iw, n, α)
        end
    end
    for iw=r-1:n
        for jw=iw-r+2:iw+1
            A[iw, jw]=w(iw+1-jw, r, iw, n, α)
        end
    end
    return A
end

function w(i::Int64, r, j, n, a)
    ar = ones(r-1)
    br = ones(r-1)
    for lj = 1:r-2
        jj = r-lj-1
        kj = r-2
        tj = i-1
        aj = collect(Int64, 0:tj)
        bj = collect(Int64, tj+2:kj+1)
        cj = [aj; bj]
        dj = collect(Int64, -1:tj-1)
        ej = collect(Int64, tj+1:kj)
        fj = [dj; ej]
        yj = binomial.(cj, jj)
        pj = binomial.(fj, jj)
        sj = 1
        tj = 1
        for m = 1:jj
            sj = sj.*yj[:, m]
            tj = tj.*pj[:, m]
            ar[lj] = sum(sj)
            br[lj] = sum(tj)
        end
    end
    s=0
    for l=1:r-1
        k=1
        for m1=1:l
            k=k/(m1-a)
        end
        k=k*factorial(l)
        k=k*(ar[l]*((n-j)^(l-a))-br[l]*((n-j+1)^(l-a)))
        s=s+k
    end
    s=s*((-1)^(i+1))/(factorial(i)*factorial(r-1-i))
    return s
end

function fj(n::Int64, T, fy)
    B = zeros(n+1)
    for i2 = 1:n+1
        B[i2] = fy(T*(i2-1)/n)
    end
    return B
end


text:
#=
    Simple MCMC:
    Utilizes a Metropolis-Hastings Algorithm with a
    standard Gibbs sampler.
=#

#=
    Basic linear equation function
=#

function lnfnc(x::Array{Float64}, a::Float64, b::Float64)
    y0 = Float64[]
    for xi in x
        y0i = (a * xi) + b
	push!(y0, y0i)
    end
    return y0
end

#=
    Basic chi squared function
    Note: data set error is variance not stddev
=#

function chisqr(yo::Array{Float64}, ye::Array{Float64}, syo::Array{Float64})
    csq = Float64[]
    for (yoi, yei, syoi) in zip(yo, ye, syo)
        csqi = (yoi - yei)^2 / syoi
        push!(csq, csqi)
    end
    return sum(csq)
end

#=
    Basic Metropolis-Hastings Algorithm with a
    standard Gibbs sampler.
=#

function mcmc(a0, sa, b0, sb, x, y, dy)
    y0 = lnfnc(x, a0, b0)
    chi0 = chisqr(y0, y, dy)
    atrace = Float64[]
    btrace = Float64[]
    i = 0
    j = 0
    mcn = 500000
    brn = 1000
    while i < mcn
        at = (randn() * sa) + a0
        bt = (randn() * sb) + b0
        yt = lnfnc(x, at, bt)
        chit = chisqr(yt, y, dy)
        acal = e^((chi0 - chit) / 2.0)
        aexp = min(1.0, acal)
        u = rand()
        if u <= aexp
            a0 = at
            b0 = bt
            y0 = yt
            chi0 = chit
            push!(atrace, at)
            push!(btrace, bt)
            i += 1
        end
        j += 1
    end
    return i, j, sum(atrace[brn:end]) / length(atrace[brn:end]), sum(btrace[brn:end]) / length(btrace[brn:end])
end

#=
    Main
=#

x = Float64[203, 58, 210, 202,
	    198, 158, 165, 201,
	    157, 131, 166, 160,
	    186, 125, 218, 146]

y = Float64[495, 173, 479, 504,
	    510, 416, 393, 442,
	    317, 311, 400, 337,
	    423, 334, 533, 344]

dy = Float64[21, 15, 27, 14,
	     30, 16, 14, 25,
	     52, 16, 34, 31,
	     42, 26, 16, 22]

a0 = 2.5
sa = 0.025
b0 = 28.82
sb = 2.5

# timing start
@time mcmc_out = mcmc(a0, sa, b0, sb, x, y, dy)
# timing end
println(mcmc_out)


text:
# The images for this simulation are fairly small compared to the human eye.
# The eye has roughly 126 million sensors (rods and cones).
# This sim handles only images that are 256x256 pixels.
# The pupil (32x32 pixels) will continuosly scan the image trying figure out
# The image.


text:
using TrajectoryOptimization
using StaticArrays
using LinearAlgebra
const TO = TrajectoryOptimization

struct DoubleIntegrator{T} <: AbstractModel
    mass::T
end

function TO.dynamics(model::DoubleIntegrator, x, u)
    SA[x[2], u[1] / model.mass]
end

Base.size(::DoubleIntegrator) = 2,1

# Model and discretization
model = DoubleIntegrator(1.0)
n,m = size(model)
tf = 3.0  # sec
N = 21    # number of knot points

# Objective
x0 = SA[0,0.]  # initial state
xf = SA[1,0.]  # final state

Q = Diagonal(@SVector ones(n))
R = Diagonal(@SVector ones(m))
obj = LQRObjective(Q, R, N*Q, xf, N)

# Constraints
cons = TO.ConstraintSet(n,m,N)
add_constraint!(cons, GoalConstraint(xf), N:N)
add_constraint!(cons, BoundConstraint(n,m, u_min=-10, u_max=10), 1:N-1)

# Create and solve problem
prob = Problem(model, obj, xf, tf, x0=x0, constraints=cons)
solver = ALTROSolver(prob)
cost(solver)           # initial cost
solve!(solver)         # solve with ALTRO
max_violation(solver)  # max constraint violation
cost(solver)           # final cost
iterations(solver)     # total number of iterations

# Get the state and control trajectories
X = states(solver)
U = controls(solver)


text:
abstract type FilterMethod
end
struct JosephForm <: FilterMethod
end
struct SimpleKalman <: FilterMethod
end

function correct(SSM::LinearStateSpaceModel, u::T, y) where T
    x, Ppred = meancov(u)
    H, R = SSM.obs.H, SSM.obs.R
    yres = y - H*x # innovation residual

    S = H*Ppred*H' + R # innovation covariance

    K = Ppred*H'/S # Kalman gain
    x = x + K*yres
    P = (I - K*H)*Ppred*(I - K*H)' + K*R*K'
    T(x, P), yres, S
end

correct(O::LinearObservation, u, y) = correct(JosephForm(), u, (Gaussian(y, O.R), O.H))

function correct(method::JosephForm, u::T, (v, H)::Tuple{<:Gaussian, <:Any}) where T
    x, Ppred = meancov(u)
    y, R = meancov(v)
    yres = y - H*x # innovation residual

    S = H*Ppred*H' + R # innovation covariance

    K = Ppred*H'/S # Kalman gain
    x = x + K*yres
    P = (I - K*H)*Ppred*(I - K*H)' + K*R*K'
    T(x, P), yres, S
end

"""

Single Kalman filter with control being the observation, consisting of a prediction step
(evolution of the underlying system) and a correction step `correct`.
Return filtered state `t => (x, P)` and predicted `(xpred, Ppred)`.

Computes and returns as well the log likelihood of the residual.
"""
function dyniterate(M::StateSpaceModel, ((s, u), ll)::Condition{<:Pair}, v)
    t, y = v
    t, upred = evolve(M.sys, s => u, t)
    u, yres, S = correct(M, upred, y)
    llᵒ = llikelihood(M, yres, S)
    (t => (u, upred, ll + llᵒ)), Condition(t => u, ll + llᵒ)
end

function dyniterate(M::StateSpaceModel, ((value,), dummy)::Condition{<:Start}, v)
    s, u = value
    t, y = v
    t, upred = evolve(M.sys, s => u, t)
    u, yres, S = correct(M, upred, y)
    ll = llikelihood(M, yres, S)
    (t => (u, upred, ll)), Condition(t => u, ll)
end

function dyniterate(O::LinearObservation{<:LinearEvolution}, ((s, u), ll)::Condition, v)
    t, y = v
    t, upred = evolve(O.P, s => u, t)
    u, yres, S = correct(O, upred, y)
    llᵒ = llikelihood(O, yres, S)
    (t => u), Condition(t => u, ll + llᵒ)
end

function dyniterate(O::LinearObservation{<:LinearEvolution}, ((s, u), ll)::Filter, v)
    t, y = v
    t, upred = evolve(O.P, s => u, t)
    u, yres, S = correct(O, upred, y)
    llᵒ = llikelihood(O, yres, S)
    (t => (u, upred, ll + llᵒ)), Filter(t => u, ll + llᵒ)
end

function dyniterate(O::LinearObservation{<:LinearEvolution}, start::Start{<:Filter}, v)
    ((s, u), ll) = start.value
    t, y = v
    t, upred = evolve(O.P, s => u, t)
    u, yres, S = correct(O, upred, y)
    llᵒ = llikelihood(O, yres, S)
    (t => (u, upred, ll + llᵒ)), Filter(t => u, ll + llᵒ)
end

llikelihood(yres, S) = logpdf(Gaussian(zero(yres), symmetrize!(S)), yres)
llikelihood(_, yres, S) = logpdf(Gaussian(zero(yres), symmetrize!(S)), yres)

#=
function dyniterate(M::StateSpaceModel, (s, u)::Pair, (v,)::Observe)
    t, (v, H) = v
    t, upred = evolve(M.sys, s => u, t)
    u, yres, S, K = correct(JosephForm(), upred, (v, H))
    ll = llikelihood(M, yres, S)
    t => (u, Ppred, ll, K), t => u
end
=#

"""
    kalmanfilter(M, t => x0) -> kf

`kf(iter)` is an iterator over `Gaussian`s or `Distributions` representing
the filtered distribution of `x` where `y` iterates over
(enumerated) signal values.

# Example

```
kf = kalmanfilter(M, 0 => prior) #
est1 = collect(kf(Y1))
est2 = collect(kf(Y2))
```
"""
kalmanfilter(M, prior) = iter -> filter(iter, from(M, prior))
kalmanfilter(M) = iter -> filter(iter, M)

function kalmanfilter(M, prior, Y)
    P = filter(Y, M)

    ϕ = dyniterate(P, Start(prior))
    ϕ === nothing && error("no observations")
    (t, u), state = ϕ

    X = trajectory((t => u[1],))
    while true
        ϕ = dyniterate(P, state)
        ϕ === nothing && break
        (t, u), state = ϕ
        push!(X, t => u[1])
    end
    ll = u[3]
    X, ll
end

function kalmanfilter(O::Observation, prior, Y)
    P = Bind(Y, O)

    ϕ = dyniterate(P, Start(Filter(prior, 0.0)))
    ϕ === nothing && error("no observations")
    (t, u), state = ϕ

    X = trajectory((t => u[1],))
    while true
        ϕ = dyniterate(P, state)
        ϕ === nothing && break
        (t, u), state = ϕ
        push!(X, t => u[1])
    end
    ll = u[3]
    X, ll
end


text:
import AutoGrad: Sparse, zeroslike
import LinearAlgebra: axpy!
import Base: +, -, copyto!
using Knet.KnetArrays: DevArray
using Base.Broadcast: Broadcasted

Sparse(a::DevArray{T,N},v::Vector{Any},i::Vector{Any}) where {T,N} = Sparse{T,N}(a,v,i)
zeroslike(a::DevArray)=zero(a)
axpy!(a, x::Sparse, y::DevArray) = addto!(y, a*x)
+(a::DevArray, s::Sparse) = addto!(copy(a), s)
+(s::Sparse, a::DevArray) = addto!(copy(a), s)
-(a::DevArray, s::Sparse) = addto!(copy(a), -s)
-(s::Sparse, a::DevArray) = addto!(-a, s)

function copyto!(a::DevArray, bc::Broadcasted{S,A,F,X}) where
    {S, A, F <: Union{typeof(+),typeof(-)}, X <: Tuple{Any,Sparse}}
    (b,c) = bc.args
    if !(size(a) == size(b) == size(c.container))
        a .= bc.f.(b, full(c))
        return a
    end
    a === b || copyto!(a, b)
    F <: typeof(-) && (c = -c)
    addto!(a, c)
    return a
end


text:
export Sweeps,
       nsweep,
       maxdim,
       mindim,
       cutoff,
       maxdim!,
       mindim!,
       cutoff!,
       noise!,
       sweepnext

mutable struct Sweeps
  nsweep::Int
  maxdim::Vector{Int}
  cutoff::Vector{Float64}
  mindim::Vector{Int}
  noise::Vector{Float64}

  function Sweeps(nsw::Int)
    return new(nsw,fill(1,nsw),zeros(nsw),fill(1,nsw), zeros(nsw))
  end

end

nsweep(sw::Sweeps)::Int = sw.nsweep
length(sw::Sweeps)::Int = sw.nsweep

maxdim(sw::Sweeps,n::Int)::Int = sw.maxdim[n]
mindim(sw::Sweeps,n::Int)::Int = sw.mindim[n]
cutoff(sw::Sweeps,n::Int)::Float64 = sw.cutoff[n]
noise(sw::Sweeps,n::Int)::Float64  = sw.noise[n]

function maxdim!(sw::Sweeps,maxdims::Int...)::Nothing
  Nm = length(maxdims)
  N = min(nsweep(sw),Nm)
  for i=1:N
    sw.maxdim[i] = maxdims[i]
  end
  for i=Nm+1:nsweep(sw)
    sw.maxdim[i] = maxdims[Nm]
  end
end

function mindim!(sw::Sweeps,mindims::Int...)::Nothing
  Nm = length(mindims)
  N = min(nsweep(sw),Nm)
  for i=1:N
    sw.mindim[i] = mindims[i]
  end
  for i=Nm+1:nsweep(sw)
    sw.mindim[i] = mindims[Nm]
  end
end

function cutoff!(sw::Sweeps,cutoffs::Float64...)::Nothing
  Nm = length(cutoffs)
  N = min(nsweep(sw),Nm)
  for i=1:N
    sw.cutoff[i] = cutoffs[i]
  end
  for i=Nm+1:nsweep(sw)
    sw.cutoff[i] = cutoffs[Nm]
  end
end

function noise!(sw::Sweeps,noises::Float64...)::Nothing
  Nm = length(noises)
  N = min(nsweep(sw),Nm)
  for i=1:N
    sw.noise[i] = noises[i]
  end
  for i=Nm+1:nsweep(sw)
    # this is different behavior from cutoff!, min/maxdim!, etc.,
    # but I (CDW) think it's the right thing:
    # if the noise isn't specified, it should be 0
    sw.noise[i] = 0
  end
end

struct SweepNext
  N::Int
end

sweepnext(N::Int)::SweepNext = SweepNext(N)

function iterate(sn::SweepNext,state=(0,1))
  b,ha = state
  if ha==1
    inc = 1
    bstop = sn.N
  else
    inc = -1
    bstop = 0
  end
  new_b = b+inc
  new_ha = ha
  done = false
  if new_b==bstop
    new_b -= inc
    new_ha += 1
    if ha==2
      return nothing
    end
  end
  return ((new_b,new_ha),(new_b,new_ha))
end

function Base.show(io::IO,
                   sw::Sweeps)
  println(io,"Sweeps")
  for n=1:nsweep(sw)
    @printf(io,"%d cutoff=%.1E, maxdim=%d, mindim=%d\n",n,cutoff(sw,n),maxdim(sw,n),mindim(sw,n))
  end
end