####
#### 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


# Defining our own `_unsafe_getindex!` means we don't have to rely on Base internals.
# If there is a performance cost, we may want to switch to a copy of the Base version,
# but we would then need to call `to_indices` earlier to avoid passing in `Colon`.
function _unsafe_getindex!(
    dest::AbstractArray,
    src::AbstractArray,
    I::Vararg{Union{Real, AbstractArray, Colon}, N}
) where N  # `where N` used to force specialisation
    @inbounds for (i, j) in zip(eachindex(dest), Iterators.product(to_indices(src, I)...))
        dest[i] = src[j...]
    end
    return dest
end

#=
Optimized version (from Base) of the `_unsafe_getindex!` function above.
Can use this if it improves performance, but ideally we'd avoid an `@generated` function.

# This code is copied from Base/multidimensional.jl
# https://github.com/JuliaLang/julia/blob/f12cde0e8c7f4268838d3201dd9e13b90ea44633/base/multidimensional.jl#L761-L777
# Does not support `Colon`, as it assumes `I == to_indices(src, I)`.
@generated function _unsafe_getindex!(dest::AbstractArray, src::AbstractArray, I::Vararg{Union{Real, AbstractArray}, N}) where N
    quote
        Base.@_inline_meta
        D = eachindex(dest)
        Dy = iterate(D)
        @inbounds Base.Cartesian.@nloops $N j d->I[d] begin
            # This condition is never hit, but at the moment
            # the optimizer is not clever enough to split the union without it
            Dy === nothing && return dest
            (idx, state) = Dy
            dest[idx] = Base.Cartesian.@ncall $N getindex src j
            Dy = iterate(D, state)
        end
        return dest
    end
end
=#


function read(fn::AbstractString, bands = nothing)
    isfile(fn) || error("File not found.")
    dataset = ArchGDAL.unsafe_read(fn)

    if bands === nothing
        nbands = ArchGDAL.nraster(dataset)
        bands = 1:nbands
    end
    A = ArchGDAL.read(dataset, bands)
    am = get_affine_map(dataset)

    # nodata masking
    # A = Array{Union{Missing, eltype(A)}}(A)
    mask = falses(size(A))

    for i = bands
        band = ArchGDAL.getband(dataset, i)
        maskflags = mask_flags(band)

        # All values are valid, skip masking
        if :GMF_ALL_VALID in maskflags
            @debug "No masking"
            continue
        # Mask is valid for all bands
        elseif :GMF_PER_DATASET in maskflags
            @debug "Mask for each band"
            maskband = ArchGDAL.getmaskband(band)
            m = ArchGDAL.read(maskband) .== 0
            mask[:,:,i] = m
        # Alpha layer
        elseif :GMF_ALPHA in maskflags
            @warn "Dataset has band $i with an Alpha band, which is unsupported for now."
            continue
        # Nodata values
        elseif :GMF_NODATA in maskflags
            @debug "Flag NODATA"
            nodata = get_nodata(band)
            mask[:, :, i] = A[:,:,i] .== nodata
        else
            @warn "Unknown/unsupported mask."
        end
    end

    if any(mask)
        A = Array{Union{Missing,eltype(A)}}(A)
        A[mask] .= missing
    end

    wkt = ArchGDAL.getproj(dataset)
    ArchGDAL.destroy(dataset)
    GeoArray(A, am, wkt)
end

"""
- rast(fn::AbstractString, bands = nothing): construct `GeoArray` object from file
- rast(ga::GeoArray; vals::AbstractArray{T}): Reconstruct `GeoArray` object
"""
rast(fn::AbstractString, bands = nothing) = read(fn, bands)
function rast(ga::GeoArray; vals::AbstractArray{T}) where T<: Real
    GeoArray(vals, ga.f, ga.crs)
end

"""
    st_read(file::String, options...)
    st_read(files::Array{String,1}, options)

# Arguments:
- `options`: other parameters to `ArchGDAL.read(dataset, options...)`.
"""
function st_read(file::String, options...)
    ArchGDAL.read(file) do dataset
        ArchGDAL.read(dataset, options...)
    end
end

# read multiple tiff files and cbind
function st_read(files::Array{String,1}, options)
    # bands = collect(bands)
    # bands = collect(Int32, bands)
    res = map(file -> st_read(file, options...), files)
    vcat(res...)
end

export st_read, rast


# This file contains mostly random utilities
# TODO: make the matrix aggrigator only take the output array

@doc """
    Scattered = wrap(layers::stParallel{K, M}, results::Array{T,N}, X) where {K, N, M, T<:Number}

given a stParallel and an array as produced by the thin ST, wrap the
results in the easier to process Scattered type. note that the data is zero
when produced this way.

TODO: the 2nd layer still has a weird size, since both indices are stuck together (see the cwt example)
"""
function rollSt(layers::stParallel{K,1}, results::AbstractArray{T,N}, X;
    percentage = 0.9, outputSubsample = (-1, -1), full = false) where {K,N,T<:Number}
    if full
        wrapped = ScatteredFull(layers, X, outputSubsample = outputSubsample)
    else
        wrapped = ScatteredOut(layers, X, outputSubsample = outputSubsample)
    end
    n, q, dataSizes, outputSizes, resultingSize = calculateSizes(layers,
        outputSubsample,
        size(X),
        percentage = percentage)
    @info "" n' q' dataSizes' outputSizes' resultingSize'
    outerAxes = axes(X)[2:end]
    addedLayers = getListOfScaleDims(layers, n)
    for i = 1:depth(layers)+1
        concatStart = sum([resultingSize[m] * prod(outputSizes[m][2:end-ndims(X)+1]) for m = 1:i-1]) + 1
        for js in indexOverScales(addedLayers, i)
            j = linearFromTuple(js, addedLayers[i])
            accessed = results[concatStart.+j*resultingSize[i].+(0:(resultingSize[i]-1)),
                outerAxes...]
            thingToWrite = reshape(accessed,
                (resultingSize[i],
                    size(X)[2:end]...,))
            if js == 1
                wrapped.output[i][:, outerAxes...] = thingToWrite
            else
                wrapped.output[i][:, js..., outerAxes...] = thingToWrite
            end
        end
    end
    return wrapped
end

@doc """
    MatrixAggrigator(pardir::String; keep=[],named="output")

    gathers all sheared matrices in a file with the name grouped(keep) and makes a single matrix out of them, with each row being a single sheared transform. Keep determines which depth to keep, with an empty matrix meaning keep all layers. named is the name of the variable in the file, while it is saved to "grouped(keep)"
"""
function MatrixAggrigator(pardir::String; keep = [], named = "output")
    vectorLength = 0
    # First determine the size of a single entry
    for (root, dir, files) in walkdir(pardir)
        @debug "folder $root"
        for name in files
            if name[end-3:end] == ".jld"
                results = load("$root/$name", named)
                if keep == []
                    vectorLength = [prod(size(results))]
                else
                    vectorLength = [prod(size(x)) for x in results.output[keep]]
                end
                break
            end
        end
    end
    # quick walkthrough to determine the number of examples
    NumFiles = 0
    for (root, dir, files) in walkdir(pardir)
        for name in files
            if name[end-3:end] == ".jld"
                NumFiles += 1
            end
        end
    end
    @debug "vectorLength = $vectorLength"
    grouped = zeros(NumFiles, sum(vectorLength))
    if keep == []
        keep = 1:size(vectorLength, 1)
    end
    i = 1
    # Acutally load the data into the matrix
    for (root, dir, files) in walkdir(pardir)
        for name in files
            if name[end-3:end] == ".jld" && name[1:4] == "data"
                results = load("$root/$name", named)
                grouped[i, :] = reshape(results, (1, vectorLength[1]))
                i += 1
            end
        end
    end
    save("$(pardir)/grouped$(keep).jld", "grouped", grouped)
    return grouped
end

function reshapeFlattened(mat::Array{Float64}, sheared::S) where {S<:Scattered}
    return reshape(mat, size(sheared.output[end]))
end

@doc """
    transformFolder(sourceFolder::String, destFolder::String, layers::stParallel; separate::Bool=false, loadThis::Function=defaultLoadFunction, nonlinear::Function=abs, subsam::Function=bspline, stType::String="full", thin::Bool=true, postSubsample::Tuple{Int,Int}=(-1,-1), overwrite::Bool=false)

Given the (relative) name of a folder and a layered transform, it will load and transform the entire folder, and save the output into a similarly structured set of folders in destFolder. If separate is true, it will also save a file containing just the Scattered coefficients from all inputs.

The data is loaded via a user supplied function loadThis. This should have the signature `(data,hasContents) = loadThis(filename::String)`. `filename` is a (relative) path to the file, while `data` should be arranged so that the  last dimension is to be transformed. `hasContents` is a boolean that indicates if there is actually any data in the file.
If overwrite is false, then skip files that already exist. If overwrite is true, for right now, it uses whatever is already there as the last entries
"""
function transformFolder(sourceFolder::String, destFolder::String, layers::stParallel; separate::Bool = false, loadThis::Function = defaultLoadFunction, nonlinear::Function = abs, subsam::Function = bspline, stType::String = "full", thin::Bool = true, postSubsample::Tuple{Int,Int} = (-1, -1), overwrite::Bool = false)
    for (root, dirs, files) in walkdir(sourceFolder)
        for dir in dirs
            mkpath(joinpath(destFolder, dir))
        end
        @debug "" root
        for file in files
            # data is expected as *column* vectors
            @debug "starting file $(joinpath(root, file))"
            savefile = joinpath(destFolder, relpath(root, sourceFolder), "$(file).h5")
            savefileSize = stat(savefile).size
            if overwrite || savefileSize < 10
                (fullMatrix, hasContents) = loadThis(joinpath(root, file))
                if hasContents
                    # breaking up the data into managable chunks
                    nEx = size(fullMatrix, 1)
                    # if the save file already has meaningful content, load this content
                    if savefileSize > 10
                        prevWork = h5read(savefile, "result/result")
                    end
                    # if the savefile doesn't have content, or if there should be more examples than what the save file has, calculate the new coefficients
                    if savefileSize < 10 || (nEx > 2 && size(prevWork, 1) <= 2)
                        innerAxes = axes(fullMatrix)[2:end]
                        usefulSlice = 1:min(2, nEx)
                        if thin
                            tmpResult = thinSt(fullMatrix[usefulSlice, innerAxes...], layers, nonlinear = nonlinear, subsam = subsam, stType = stType, outputSubsample = postSubsample)
                        else
                            tmpResult = st(fullMatrix[usefulSlice, innerAxes...], layers, nonlinear = nonlinear, subsam = subsam, stType = stType)
                        end
                        # if the data is too big, we'll need extra room to store the results
                        result = zeros(nEx, size(tmpResult)[2:end]...)
                        innerResultAxes = axes(tmpResult)[2:end]
                        result[usefulSlice, innerResultAxes...] = tmpResult
                        if nEx > 2
                            if savefileSize > 10
                                usefulSlice = (2*ceil(Int, nEx / 2)-1):nEx
                                result[usefulSlice, innerResultAxes...] = prevWork
                            end
                            maxValue = (savefileSize > 10 && nEx > 2) ? ceil(Int, nEx / 2) - 1 : ceil(Int, nEx / 2)
                            for i = 2:maxValue
                                usefulSlice = (2*i-1):min(nEx, 2 * i)
                                if thin
                                    tmpResult = thinSt(fullMatrix[usefulSlice, innerAxes...], layers, nonlinear = nonlinear, subsam = subsam, stType = stType, outputSubsample = postSubsample)
                                else
                                    tmpResult = st(fullMatrix[usefulSlice, innerAxes...], layers, nonlinear = nonlinear, subsam = subsam, stType = stType)
                                end
                                result[usefulSlice, innerResultAxes...] = tmpResult
                            end
                            # repair a problem in the previous run
                        end
                        @info "saving to $(savefile)"
                        if savefileSize < 10 && isfile(savefile)
                            rm(savefile)
                        end
                        h5write(savefile, "result/result", result)
                    end
                end
            end
        end
    end
    @info "saving settings to $(joinpath(destFolder, "settings.h5"))"
    save(joinpath(destFolder, "settings.jld"), "layers", layers)
    return
end

function defaultLoadFunction(filename)
    return (h5read(filename, "data"), true)
end

function loadSyntheticMatFile(datafile::String)
    hasContents = (datafile[end-3:end] == ".mat")
    if hasContents
        wef = matread(datafile)
        return (wef["resp"]', hasContents)
    else
        return ([0.0 0.0; 0.0 0.0], hasContents)
    end
end

function loadhdf5(datafile::String)
    hasContents = (datafile[end-3:end] == ".mat")
    if hasContents
        wef = matread(datafile)
        return (wef["resp"]', hasContents)
    else
        return ([0.0 0.0; 0.0 0.0], hasContents)
    end
end
rollSt(results::AbstractArray{T,N}, layers::stParallel{K,1}, X;
    percentage = 0.9, outputSubsample = (-1, -1), full = false) where {K,N,T<:Number} = rollSt(layers, results, X;
    percentage = percentage, outputSubsample = outputSubsample, full = full)


# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/JuliaFEM.jl/blob/master/LICENSE.md

using JuliaFEM
using JuliaFEM.Preprocess
using JuliaFEM.Testing

@testset "read inp file" begin
    model = open(parse_abaqus, Pkg.dir("JuliaFEM")*"/geometry/3d_beam/palkki.inp")
    @test length(model["nodes"]) == 298
    @test length(model["elements"]) == 120
    @test length(model["elsets"]["Body1"]) == 120
    @test length(model["nsets"]["SUPPORT"]) == 9
    @test length(model["nsets"]["LOAD"]) == 9
    @test length(model["nsets"]["TOP"]) == 83
end

@testset "test read element section" begin
    data = """*ELEMENT, TYPE=C3D10, ELSET=BEAM
    1,       243,       240,       191,       117,       245,       242,       244,
    1,         2,       196
    2,       204,       199,       175,       130,       207,       208,       209,
    3,         4,       176
    """
    data = split(data, "\n")
    model = Dict{AbstractString, Any}()
    model["nsets"] = Dict{AbstractString, Vector{Int}}()
    model["elsets"] = Dict{AbstractString, Vector{Int}}()
    model["elements"] = Dict{Integer, Any}()
    parse_section(model, data, :ELEMENT, 1, 5, Val{:ELEMENT})
    @test length(model["elements"]) == 2
    @test model["elements"][1]["connectivity"] == [243, 240, 191, 117, 245, 242, 244, 1, 2, 196]
    @test model["elements"][2]["connectivity"] == [204, 199, 175, 130, 207, 208, 209, 3, 4, 176]
    @test model["elsets"]["BEAM"] == [1, 2]
end

@testset "parse nodes from abaqus .inp file to Mesh" begin
    fn = Pkg.dir("JuliaFEM") * "/test/testdata/cube_tet4.inp"
    mesh = abaqus_read_mesh(fn)
    info(mesh.surfaces)
    @test length(mesh.nodes) == 10
    @test length(mesh.elements) == 17
    @test haskey(mesh.elements, 1)
    @test mesh.elements[1] == [8, 10, 1, 2]
    @test mesh.element_types[1] == :Tet4
    @test haskey(mesh.node_sets, :SYM12)
    @test haskey(mesh.element_sets, :CUBE)
    @test haskey(mesh.surfaces, :LOAD)
    @test length(mesh.surfaces[:LOAD]) == 2
    @test mesh.surfaces[:LOAD][1] == (16, :S1)
    @test mesh.surface_types[:LOAD] == :ELEMENT
end


module SciMLBase

using ConstructionBase
using RecipesBase, RecursiveArrayTools, Tables, TreeViews
using DocStringExtensions
using LinearAlgebra
using Statistics
using Distributed
using StaticArrays

import Logging, ArrayInterface
import IteratorInterfaceExtensions
import CommonSolve: solve, init, solve!

function __solve end
function __init end

abstract type SciMLProblem end

# Problems
"""
$(TYPEDEF)

Base type for all DifferentialEquations.jl problems. Concrete subtypes of
`DEProblem` contain the necessary information to fully define a differential
equation of the corresponding type.
"""
abstract type DEProblem <: SciMLProblem end

"""
$(TYPEDEF)
"""
abstract type DEElement end

"""
$(TYPEDEF)
"""
abstract type DESensitivity end

"""
$(TYPEDEF)

Base for types which define linear systems.
"""
abstract type AbstractLinearProblem{bType,isinplace} <: SciMLProblem end

"""
$(TYPEDEF)

Base for types which define integrals suitable for quadrature.
"""
abstract type AbstractQuadratureProblem{isinplace} <: SciMLProblem end

"""
$(TYPEDEF)

Base for types which define equations for optimization.
"""
abstract type AbstractOptimizationProblem{isinplace} <: SciMLProblem end

"""
$(TYPEDEF)

Base for types which define nonlinear solve problems (f(u)=0).
"""
abstract type AbstractNonlinearProblem{uType,isinplace} <: DEProblem end
const AbstractSteadyStateProblem{uType,isinplace} = AbstractNonlinearProblem{uType,isinplace}

"""
$(TYPEDEF)
"""
abstract type AbstractNoiseProblem <: DEProblem end

"""
$(TYPEDEF)

Base for types which define ODE problems.
"""
abstract type AbstractODEProblem{uType,tType,isinplace} <: DEProblem end

"""
$(TYPEDEF)

Base for types which define discrete problems.
"""
abstract type AbstractDiscreteProblem{uType,tType,isinplace} <:
                      AbstractODEProblem{uType,tType,isinplace} end

"""
$(TYPEDEF)
"""
abstract type AbstractAnalyticalProblem{uType,tType,isinplace} <:
                      AbstractODEProblem{uType,tType,isinplace} end

"""
$(TYPEDEF)

Base for types which define RODE problems.
"""
abstract type AbstractRODEProblem{uType,tType,isinplace,ND} <: DEProblem end

"""
$(TYPEDEF)

Base for types which define SDE problems.
"""
abstract type AbstractSDEProblem{uType,tType,isinplace,ND} <:
                      AbstractRODEProblem{uType,tType,isinplace,ND} end

"""
$(TYPEDEF)

Base for types which define DAE problems.
"""
abstract type AbstractDAEProblem{uType,duType,tType,isinplace} <: DEProblem end

"""
$(TYPEDEF)

Base for types which define DDE problems.
"""
abstract type AbstractDDEProblem{uType,tType,lType,isinplace} <: DEProblem end

"""
$(TYPEDEF)
"""
abstract type AbstractConstantLagDDEProblem{uType,tType,lType,isinplace} <:
                      AbstractDDEProblem{uType,tType,lType,isinplace} end

"""
$(TYPEDEF)
"""
abstract type AbstractSecondOrderODEProblem{uType,tType,isinplace} <: AbstractODEProblem{uType,tType,isinplace} end

"""
$(TYPEDEF)

Base for types which define BVP problems.
"""
abstract type AbstractBVProblem{uType,tType,isinplace} <: AbstractODEProblem{uType,tType,isinplace} end

"""
$(TYPEDEF)

Base for types which define jump problems.
"""
abstract type AbstractJumpProblem{P,J} <: DEProblem end

"""
$(TYPEDEF)

Base for types which define SDDE problems.
"""
abstract type AbstractSDDEProblem{uType,tType,lType,isinplace,ND} <: DEProblem end

"""
$(TYPEDEF)
"""
abstract type AbstractConstantLagSDDEProblem{uType,tType,lType,isinplace,ND} <:
                      AbstractSDDEProblem{uType,tType,lType,isinplace,ND} end

"""
$(TYPEDEF)

Base for types which define PDE problems.
"""
abstract type AbstractPDEProblem <: DEProblem end

# Algorithms
"""
$(TYPEDEF)
"""
abstract type SciMLAlgorithm end

"""
$(TYPEDEF)
"""
abstract type DEAlgorithm <: SciMLAlgorithm end

"""
$(TYPEDEF)
"""
abstract type AbstractLinearAlgorithm <: SciMLAlgorithm end

"""
$(TYPEDEF)
"""
abstract type AbstractNonlinearAlgorithm <: SciMLAlgorithm end

"""
$(TYPEDEF)
"""
abstract type AbstractQuadratureAlgorithm <: SciMLAlgorithm end

"""
$(TYPEDEF)
"""
abstract type AbstractSteadyStateAlgorithm <: DEAlgorithm end

"""
$(TYPEDEF)
"""
abstract type AbstractODEAlgorithm <: DEAlgorithm end

"""
$(TYPEDEF)
"""
abstract type AbstractSecondOrderODEAlgorithm <: DEAlgorithm end

"""
$(TYPEDEF)
"""
abstract type AbstractRODEAlgorithm <: DEAlgorithm end

"""
$(TYPEDEF)
"""
abstract type AbstractSDEAlgorithm <: DEAlgorithm end

"""
$(TYPEDEF)
"""
abstract type AbstractDAEAlgorithm <: DEAlgorithm end

"""
$(TYPEDEF)
"""
abstract type AbstractDDEAlgorithm <: DEAlgorithm end

"""
$(TYPEDEF)
"""
abstract type AbstractSDDEAlgorithm <: DEAlgorithm end
"""
$(TYPEDEF)
"""
abstract type EnsembleAlgorithm <: SciMLAlgorithm end

"""
$(TYPEDEF)
"""
abstract type AbstractSensitivityAlgorithm{CS,AD,FDT} <: SciMLAlgorithm end

"""
$(TYPEDEF)
"""
abstract type DAEInitializationAlgorithm <: SciMLAlgorithm end

# PDE Discretizations

"""
$(TYPEDEF)
"""
abstract type AbstractDiscretization <: SciMLAlgorithm end

# Monte Carlo Simulations
"""
$(TYPEDEF)
"""
abstract type AbstractEnsembleProblem <: SciMLProblem end

"""
$(TYPEDEF)
"""
abstract type AbstractEnsembleEstimator <: SciMLAlgorithm end

export EnsembleProblem
export EnsembleSolution, EnsembleTestSolution, EnsembleSummary

"""
$(TYPEDEF)
"""
abstract type AbstractDiffEqInterpolation end

"""
$(TYPEDEF)
"""
abstract type AbstractDEOptions end

"""
$(TYPEDEF)
"""
abstract type DECache end

"""
$(TYPEDEF)
"""
abstract type DECallback end

"""
$(TYPEDEF)
"""
abstract type AbstractContinuousCallback <: DECallback end

"""
$(TYPEDEF)
"""
abstract type AbstractDiscreteCallback <: DECallback end

# Integrators
"""
$(TYPEDEF)
"""
abstract type DEIntegrator{Alg, IIP, U, T} end

"""
$(TYPEDEF)
"""
abstract type AbstractSteadyStateIntegrator{Alg, IIP, U} <: DEIntegrator{Alg, IIP, U, Nothing} end

"""
$(TYPEDEF)
"""
abstract type AbstractODEIntegrator{Alg, IIP, U, T} <: DEIntegrator{Alg, IIP, U, T} end

"""
$(TYPEDEF)
"""
abstract type AbstractSecondOrderODEIntegrator{Alg, IIP, U, T} <: DEIntegrator{Alg, IIP, U, T} end

"""
$(TYPEDEF)
"""
abstract type AbstractRODEIntegrator{Alg, IIP, U, T} <: DEIntegrator{Alg, IIP, U, T} end

"""
$(TYPEDEF)
"""
abstract type AbstractSDEIntegrator{Alg, IIP, U, T} <: DEIntegrator{Alg, IIP, U, T} end

"""
$(TYPEDEF)
"""
abstract type AbstractDDEIntegrator{Alg, IIP, U, T} <: DEIntegrator{Alg, IIP, U, T} end

"""
$(TYPEDEF)
"""
abstract type AbstractDAEIntegrator{Alg, IIP, U, T} <: DEIntegrator{Alg, IIP, U, T} end

"""
$(TYPEDEF)
"""
abstract type AbstractSDDEIntegrator{Alg, IIP, U, T} <: DEIntegrator{Alg, IIP, U, T} end

# Solutions
"""
$(TYPEDEF)
"""
abstract type AbstractNoTimeSolution{T,N} <: AbstractArray{T,N} end

"""
$(TYPEDEF)
"""
abstract type AbstractTimeseriesSolution{T,N,A} <: AbstractDiffEqArray{T,N,A} end

"""
$(TYPEDEF)
"""
abstract type AbstractEnsembleSolution{T,N,A} <: AbstractVectorOfArray{T,N,A} end

"""
$(TYPEDEF)
"""
abstract type AbstractNoiseProcess{T,N,A,isinplace} <: AbstractDiffEqArray{T,N,A} end

"""
Union of all base solution types.
"""
const SciMLSolution = Union{AbstractTimeseriesSolution,
                         AbstractNoTimeSolution,
                         AbstractEnsembleSolution,
                         AbstractNoiseProcess}
const DESolution = SciMLSolution
export SciMLSolution, DESolution

export AllObserved

"""
$(TYPEDEF)
"""
abstract type AbstractLinearSolution{T,N} <: AbstractNoTimeSolution{T,N} end

"""
$(TYPEDEF)
"""
abstract type AbstractNonlinearSolution{T,N} <: AbstractNoTimeSolution{T,N} end

"""
$(TYPEDEF)
"""
abstract type AbstractQuadratureSolution{T,N} <: AbstractNoTimeSolution{T,N} end

"""
$(TYPEDEF)
"""
const AbstractSteadyStateSolution{T,N} = AbstractNonlinearSolution{T,N}

"""
$(TYPEDEF)
"""
abstract type AbstractAnalyticalSolution{T,N,S} <: AbstractTimeseriesSolution{T,N,S} end

"""
$(TYPEDEF)
"""
abstract type AbstractODESolution{T,N,S} <: AbstractTimeseriesSolution{T,N,S} end

# Needed for plot recipes
"""
$(TYPEDEF)
"""
abstract type AbstractDDESolution{T,N,S} <: AbstractODESolution{T,N,S} end

"""
$(TYPEDEF)
"""
abstract type AbstractRODESolution{T,N,S} <: AbstractODESolution{T,N,S} end

"""
$(TYPEDEF)
"""
abstract type AbstractDAESolution{T,N,S} <: AbstractODESolution{T,N,S} end

"""
$(TYPEDEF)
"""
abstract type AbstractSensitivitySolution{T,N,S} <: AbstractTimeseriesSolution{T,N,S} end

# Misc
"""
$(TYPEDEF)
"""
abstract type AbstractSciMLOperator{T} end

"""
$(TYPEDEF)
"""
abstract type AbstractDiffEqOperator{T} <: AbstractSciMLOperator{T} end

"""
$(TYPEDEF)
"""
abstract type AbstractDiffEqLinearOperator{T} <: AbstractDiffEqOperator{T} end

"""
$(TYPEDEF)
"""
abstract type AbstractDiffEqCompositeOperator{T} <: AbstractDiffEqLinearOperator{T} end

"""
$(TYPEDEF)

Base for types defining SciML functions.
"""
abstract type AbstractSciMLFunction{iip} <: Function end

"""
$(TYPEDEF)

Base for types defining differential equation functions.
"""
abstract type AbstractDiffEqFunction{iip} <: AbstractSciMLFunction{iip} end

"""
$(TYPEDEF)

Base for types defining optimization functions.
"""
abstract type AbstractOptimizationFunction{iip} <: AbstractSciMLFunction{iip} end

"""
$(TYPEDEF)

Base for types which define the history of a delay differential equation.
"""
abstract type AbstractHistoryFunction end

"""
$(TYPEDEF)
"""
abstract type AbstractReactionNetwork <: Function end

"""
$(TYPEDEF)

Base type for AD choices.
"""
abstract type AbstractADType end

include("utils.jl")
include("function_wrappers.jl")
include("scimlfunctions.jl")
include("alg_traits.jl")

include("operators/operators.jl")
include("operators/basic_operators.jl")
include("operators/diffeq_operator.jl")
include("operators/common_defaults.jl")

include("problems/problem_utils.jl")
include("problems/discrete_problems.jl")
include("problems/steady_state_problems.jl")
include("problems/analytical_problems.jl")
include("problems/basic_problems.jl")
include("problems/ode_problems.jl")
include("problems/rode_problems.jl")
include("problems/sde_problems.jl")
include("problems/noise_problems.jl")
include("problems/bvp_problems.jl")
include("problems/dae_problems.jl")
include("problems/dde_problems.jl")
include("problems/sdde_problems.jl")
include("problems/pde_problems.jl")
include("problems/problem_traits.jl")

include("solutions/basic_solutions.jl")
include("solutions/nonlinear_solutions.jl")
include("solutions/ode_solutions.jl")
include("solutions/rode_solutions.jl")
include("solutions/optimization_solutions.jl")
include("solutions/dae_solutions.jl")
include("solutions/solution_interface.jl")

include("ensemble/ensemble_solutions.jl")
include("ensemble/ensemble_problems.jl")
include("ensemble/basic_ensemble_solve.jl")
include("ensemble/ensemble_analysis.jl")

include("solve.jl")
include("interpolation.jl")
include("integrator_interface.jl")
include("tabletraits.jl")
include("remake.jl")

function discretize end
function symbolic_discretize end

isfunctionwrapper(x) = false
function wrapfun_oop end
function wrapfun_iip end
function unwrap_fw end

export isinplace

export solve, solve!, init, discretize, symbolic_discretize

export LinearProblem, NonlinearProblem, QuadratureProblem, OptimizationProblem

export DiscreteProblem
export SteadyStateProblem, SteadyStateSolution
export NoiseProblem
export ODEProblem, ODESolution
export DynamicalODEFunction, DynamicalODEProblem,
       SecondOrderODEProblem, SplitFunction, SplitODEProblem
export SplitSDEProblem
export DynamicalSDEFunction, DynamicalSDEProblem
export RODEProblem, RODESolution, SDEProblem
export DAEProblem, DAESolution
export DDEProblem
export DynamicalDDEFunction, DynamicalDDEProblem,
       SecondOrderDDEProblem
export SDDEProblem
export PDEProblem
export IncrementingODEProblem

export BVProblem, TwoPointBVProblem

export remake

export ODEFunction, DiscreteFunction, SplitFunction, DAEFunction, DDEFunction,
       SDEFunction, SplitSDEFunction, RODEFunction, SDDEFunction, IncrementingODEFunction, NonlinearFunction

export OptimizationFunction

export EnsembleThreads, EnsembleDistributed, EnsembleSplitThreads, EnsembleSerial

export EnsembleAnalysis, EnsembleSummary

export tuples, intervals, TimeChoiceIterator

export AffineDiffEqOperator, DiffEqScaledOperator

export DiffEqScalar, DiffEqArrayOperator, DiffEqIdentity

export step!, deleteat!,addat!,get_tmp_cache,
       full_cache,user_cache,u_cache,du_cache,
       rand_cache,ratenoise_cache,
       resize_non_user_cache!,deleteat_non_user_cache!,addat_non_user_cache!,
       terminate!,
       add_tstop!,add_saveat!,set_abstol!,
       set_reltol!,get_du, get_du!, get_dt,get_proposed_dt,set_proposed_dt!,
       u_modified!, savevalues!,reinit!, auto_dt_reset!, set_t!,
       set_u!, check_error, change_t_via_interpolation!, addsteps!,
       isdiscrete, reeval_internals_due_to_modification!

export update_coefficients!, update_coefficients,
       has_expmv!, has_expmv, has_exp, has_mul, has_mul!, has_ldiv, has_ldiv!

end


module Types
    abstract type AbstractSpec end
    abstract type SpecModifier <: AbstractSpec  end

    struct SingleComponent <: AbstractSpec end
    struct OneMol <: AbstractSpec end

    abstract type MassBasisModifier <: SpecModifier end
    struct MOLAR <: MassBasisModifier end
    struct MASS <: MassBasisModifier end
    struct TOTAL <: MassBasisModifier end

    abstract type CompoundModifier <: SpecModifier end
    struct FRACTION <: CompoundModifier end
    struct TOTAL_AMOUNT <: CompoundModifier end

    abstract type VolumeModifier <: SpecModifier end
    struct DENSITY <: VolumeModifier end
    struct VOLUME <: VolumeModifier end

    abstract type AbstractIntensiveSpec{T1<:MassBasisModifier} <: AbstractSpec  end
    abstract type AbstractTotalSpec <: AbstractSpec  end
    abstract type AbstractFractionSpec <: AbstractSpec  end
    abstract type CategoricalSpec <: AbstractSpec  end

    abstract type AbstractEnergySpec{T1} <: AbstractIntensiveSpec{T1} end

    struct Enthalpy{T} <: AbstractEnergySpec{T} end
    struct InternalEnergy{T} <: AbstractEnergySpec{T} end
    struct Gibbs{T} <: AbstractEnergySpec{T} end
    struct Helmholtz{T} <: AbstractEnergySpec{T} end

    struct Entropy{T} <: AbstractIntensiveSpec{T} end

    struct VolumeAmount{T1,T2<:VolumeModifier} <: AbstractIntensiveSpec{T1} end

    struct Pressure <: AbstractSpec end
    struct Temperature <: AbstractSpec end

    # those are vectors,

    struct MaterialCompounds{T1<:MassBasisModifier,T2<:CompoundModifier} <: AbstractTotalSpec end

    struct MaterialAmount{T1<:MassBasisModifier} <: AbstractTotalSpec end

    #humidity, to support wet gas
    abstract type HumidityModifier <: AbstractSpec end
    struct WetBulbTemperature <:HumidityModifier end
    struct RelativeHumidity <:HumidityModifier end
    struct HumidityRatio <:HumidityModifier end
    struct MassHumidity  <:HumidityModifier end
    struct MolarHumidity <:HumidityModifier end
    struct HumidityDewPoint <:HumidityModifier end

    struct HumiditySpec{T<:HumidityModifier} <: AbstractSpec end

    struct PhaseFractions <: AbstractFractionSpec end
    struct VolumeFraction <: AbstractFractionSpec end
    struct VaporFraction{T<:MassBasisModifier} <: AbstractFractionSpec end

    struct PhaseTag <: CategoricalSpec end
    struct TwoPhaseEquilibrium <: CategoricalSpec end

    struct Options <: CategoricalSpec end

    #not defined for now
    struct MolecularWeight <: AbstractSpec end

    #variable spec, signal to make a variable specs
    struct VariableSpec  end

    #base type to define models
    abstract type ThermoModel end

    export AbstractSpec
    export SpecModifier
    export AbstractIntensiveSpec
    export AbstractTotalSpec
    export AbstractFractionSpec
    export CategoricalSpec
    export SingleComponent
    export OneMol
    export MOLAR
    export MASS
    export TOTAL
    export FRACTION
    export TOTAL_AMOUNT
    export DENSITY
    export VOLUME
    export AbstractEnergySpec
    export Enthalpy
    export InternalEnergy
    export Gibbs
    export Helmholtz
    export Entropy
    export VolumeAmount
    export Pressure
    export Temperature
    export MaterialCompounds
    export MaterialAmount
    export PhaseFractions
    export VaporFraction
    export PhaseTag
    export TwoPhaseEquilibrium
    export Options
    export MolecularWeight
    export VariableSpec
    export ThermoModel

    export HumidityModifier
    export WetBulbTemperature
    export RelativeHumidity
    export HumidityRatio
    export MassHumidity
    export MolarHumidity
    export HumidityDewPoint

    export HumiditySpec

end

using .Types

const KW_TO_SPEC = Dict{Symbol,Any}(
:h =>  Enthalpy{MOLAR}()
,:g =>  Gibbs{MOLAR}()
,:a =>  Helmholtz{MOLAR}()
,:u =>  InternalEnergy{MOLAR}()

,:mol_h =>  Enthalpy{MOLAR}()
,:mol_g =>  Gibbs{MOLAR}()
,:mol_a =>  Helmholtz{MOLAR}()
,:mol_u =>  InternalEnergy{MOLAR}()

,:mass_h =>  Enthalpy{MASS}()
,:mass_g =>  Gibbs{MASS}()
,:mass_a =>  Helmholtz{MASS}()
,:mass_u =>  InternalEnergy{MASS}()

,:total_h =>  Enthalpy{TOTAL}()
,:total_g =>  Gibbs{TOTAL}()
,:total_a =>  Helmholtz{TOTAL}()
,:total_u =>  InternalEnergy{TOTAL}()

,:s =>  Entropy{MOLAR}()
,:mol_s =>  Entropy{MOLAR}()
,:mass_s =>  Entropy{MASS}()
,:total_s =>  Entropy{TOTAL}()

,:p =>  Pressure()
,:P =>  Pressure()
,:t =>  Temperature()
,:T => Temperature()

,:v =>  VolumeAmount{MOLAR,VOLUME}()
,:mol_v =>  VolumeAmount{MOLAR,VOLUME}()
,:mass_v =>  VolumeAmount{MASS,VOLUME}()
,:total_v =>  VolumeAmount{TOTAL,VOLUME}()
,:V =>  VolumeAmount{TOTAL,VOLUME}()

,:rho =>  VolumeAmount{MOLAR,DENSITY}()
,:mol_rho =>  VolumeAmount{MOLAR,DENSITY}()
,:mass_rho =>  VolumeAmount{MASS,DENSITY}()

,:ρ =>  VolumeAmount{MOLAR,DENSITY}()
,:mol_ρ =>  VolumeAmount{MOLAR,DENSITY}()
,:mass_ρ =>  VolumeAmount{MASS,DENSITY}()
,:mass =>  MaterialAmount{MASS}()
,:moles =>  MaterialAmount{MOLAR}()
,:xn =>  MaterialCompounds{MOLAR,FRACTION}()
,:xm =>   MaterialCompounds{MASS,FRACTION}()
,:n =>  MaterialCompounds{MOLAR,TOTAL_AMOUNT}()
,:m =>   MaterialCompounds{MASS,TOTAL_AMOUNT}()

,:mw =>  MolecularWeight()

,:quality =>  VaporFraction{MASS}() #looking for better name
,:vfrac =>  VaporFraction{MOLAR}() #looking for better name

,:phase_fracs =>  PhaseFractions() #looking for better name

,:phase =>PhaseTag()

,:sat => TwoPhaseEquilibrium()

,:single_component => SingleComponent()
,:one_mol => OneMol()
,:options => Options()

,:hum_wetbulb => HumiditySpec{WetBulbTemperature}()
,:hum_ratio => HumiditySpec{HumidityRatio}()
,:hum_molfrac => HumiditySpec{MolarHumidity}()
,:hum_massfrac => HumiditySpec{MassHumidity}()
,:rel_hum => HumiditySpec{RelativeHumidity}()
,:hum_dewpoint => HumiditySpec{HumidityDewPoint}()

)

const SPEC_TO_KW = Dict{Any,Symbol}(
Enthalpy{MOLAR}() => :mol_h
,Gibbs{MOLAR}() => :mol_g
,Helmholtz{MOLAR}() => :mol_a
,InternalEnergy{MOLAR}() => :mol_u

,Enthalpy{MASS}() => :mass_h
,Gibbs{MASS}()  => :mass_g
,Helmholtz{MASS}() => :mass_a
,InternalEnergy{MASS}() => :mass_u

,Enthalpy{TOTAL}() => :total_h
,Gibbs{TOTAL}() => :total_g
,Helmholtz{TOTAL}() => :total_a
,InternalEnergy{TOTAL}() => :total_u

,Entropy{MOLAR}() => :mol_s
,Entropy{MASS}() => :mass_s
,Entropy{TOTAL}() => :total_s

,Pressure() => :p
,Temperature() => :t

,VolumeAmount{MOLAR,VOLUME}() => :mol_v
,VolumeAmount{MASS,VOLUME}() => :mass_v
,VolumeAmount{TOTAL,VOLUME}() => :total_v

,VolumeAmount{MOLAR,DENSITY}()  => :rho
,VolumeAmount{MOLAR,DENSITY}()  => :mol_rho
,VolumeAmount{MASS,DENSITY}()  => :mass_rho

,MaterialAmount{MASS}() => :mass
,MaterialAmount{MOLAR}() => :moles
,MaterialCompounds{MOLAR,FRACTION}() => :xn
,MaterialCompounds{MASS,FRACTION}() => :xm
,MaterialCompounds{MOLAR,TOTAL_AMOUNT}() => :n
,MaterialCompounds{MASS,TOTAL_AMOUNT}() =>  :n

,MolecularWeight() => :mw

,VaporFraction{MASS}() => :quality#looking for better name
,VaporFraction{MOLAR}() => :vfrac#looking for better name

,PhaseFractions() => :phase_fracs #looking for better name

,PhaseTag() => :phase

,TwoPhaseEquilibrium() => :sat

,SingleComponent() => :single_component
,OneMol() => :one_mol
,Options() => :options

,HumiditySpec{WetBulbTemperature}()  => :hum_wetbulb
,HumiditySpec{HumidityRatio}() => :hum_ratio
,HumiditySpec{MolarHumidity}() => :hum_molfrac
,HumiditySpec{MassHumidity}() => :hum_massfrac
,HumiditySpec{RelativeHumidity}() => :rel_hum
,HumiditySpec{HumidityDewPoint}() => :hum_dewpoint

)

module QuickStates
using ..Types
const SinglePT = Tuple{Pressure,Temperature,SingleComponent}
const MultiPT = Tuple{Pressure,Temperature,MaterialCompounds}

pt() = (Pressure(),Temperature(),SingleComponent())
ptx() = (Pressure(),Temperature(),MaterialCompounds{MOLAR,FRACTION}())
ptn() = (Pressure(),Temperature(),MaterialCompounds{MOLAR,TOTAL_AMOUNT}())

const SinglePV = Tuple{Pressure,VolumeAmount,SingleComponent}
const MultiPV = Tuple{Pressure,VolumeAmount,MaterialCompounds}

pv() = (Pressure(),VolumeAmount{MOLAR,VOLUME}(),SingleComponent())
pvx() = (Pressure(),VolumeAmount{MOLAR,VOLUME}(),MaterialCompounds{MOLAR,FRACTION}())
pvn() = (Pressure(),VolumeAmount{MOLAR,VOLUME}(),MaterialCompounds{MOLAR,TOTAL_AMOUNT}())

const SingleVT = Tuple{VolumeAmount,Temperature,SingleComponent}
const MultiVT = Tuple{VolumeAmount,Temperature,MaterialCompounds}

vt() = (VolumeAmount{MOLAR,VOLUME}(),Temperature(),SingleComponent())
vtx() = (VolumeAmount{MOLAR,VOLUME}(),Temperature(),MaterialCompounds{MOLAR,FRACTION}())
vtn() = (VolumeAmount{MOLAR,VOLUME}(),Temperature(),MaterialCompounds{MOLAR,TOTAL_AMOUNT}())

ρt() = (VolumeAmount{MOLAR,DENSITY}(),Temperature(),SingleComponent())
ρtx() = (VolumeAmount{MOLAR,DENSITY}(),Temperature(),MaterialCompounds{MOLAR,FRACTION}())
ρtn() = (VolumeAmount{MOLAR,VOLUME}(),Temperature(),MaterialCompounds{MOLAR,TOTAL_AMOUNT}())

const SinglePS = Tuple{Pressure,Entropy,SingleComponent}
const MultiPS = Tuple{Pressure,Entropy,MaterialCompounds}

st() = (Pressure(),Entropy{MOLAR}(),SingleComponent())
stx() = (Pressure(),Entropy{MOLAR}(),MaterialCompounds{MOLAR,FRACTION}())
stn() = (Pressure(),Entropy{MOLAR}(),MaterialCompounds{MOLAR,TOTAL_AMOUNT}())

const SinglePH = Tuple{Pressure,Enthalpy,SingleComponent}
const MultiPH = Tuple{Pressure,Enthalpy,MaterialCompounds}

ph() = (Pressure(),Enthalpy{MOLAR}(),SingleComponent())
phx() = (Pressure(),Enthalpy{MOLAR}(),MaterialCompounds{MOLAR,FRACTION}())
phn() = (Pressure(),Enthalpy{MOLAR}(),MaterialCompounds{MOLAR,TOTAL_AMOUNT}())

const SingleSatT = Tuple{TwoPhaseEquilibrium,Temperature,SingleComponent}
const MultiSatT = Tuple{TwoPhaseEquilibrium,Temperature,MaterialCompounds}

sat_t() = (TwoPhaseEquilibrium(),Temperature(),SingleComponent())
sat_tx() = (TwoPhaseEquilibrium(),Temperature(),MaterialCompounds{MOLAR,FRACTION}())
sat_tn() = (TwoPhaseEquilibrium(),Temperature(),MaterialCompounds{MOLAR,TOTAL_AMOUNT}())

const SingleSatP = Tuple{TwoPhaseEquilibrium,Pressure,SingleComponent}
const MultiSatP = Tuple{TwoPhaseEquilibrium,Pressure,MaterialCompounds}

sat_p() = (TwoPhaseEquilibrium(),Pressure(),SingleComponent())
sat_px() = (TwoPhaseEquilibrium(),Pressure(),MaterialCompounds{MOLAR,FRACTION}())
sat_pn() = (TwoPhaseEquilibrium(),Pressure(),MaterialCompounds{MOLAR,TOTAL_AMOUNT}())

molar_ϕt() = (VaporFraction{MOLAR}(),Temperature(),SingleComponent())
molar_ϕntx() = (VaporFraction{MOLAR}(),Temperature(),MaterialCompounds{MOLAR,FRACTION}())
molar_ϕtn() = (VaporFraction{MOLAR}(),Temperature(),MaterialCompounds{MOLAR,TOTAL_AMOUNT}())

mass_ϕt() = (VaporFraction{MASS}(),Temperature(),SingleComponent())
mass_ϕntx() = (VaporFraction{MASS}(),Temperature(),MaterialCompounds{MOLAR,FRACTION}())
mass_ϕtn() = (VaporFraction{MASS}(),Temperature(),MaterialCompounds{MOLAR,TOTAL_AMOUNT}())

const SingleΦP = Tuple{VaporFraction,Pressure,SingleComponent}
const MultiΦP = Tuple{VaporFraction,Pressure,MaterialCompounds}
const SingleΦT = Tuple{VaporFraction,Temperature,SingleComponent}
const MultiΦT = Tuple{VaporFraction,Temperature,MaterialCompounds}

const SingleΦnP = Tuple{VaporFraction{MOLAR},Pressure,SingleComponent}
const MultiΦnP = Tuple{VaporFraction{MOLAR},Pressure,MaterialCompounds}
const SingleΦnT = Tuple{VaporFraction{MOLAR},Temperature,SingleComponent}
const MultiΦnT = Tuple{VaporFraction{MOLAR},Temperature,MaterialCompounds}

const SingleΦmP = Tuple{VaporFraction{MASS},Pressure,SingleComponent}
const MultiΦmP = Tuple{VaporFraction{MASS},Pressure,MaterialCompounds}
const SingleΦmT = Tuple{VaporFraction{MASS},Temperature,SingleComponent}
const MultiΦmT = Tuple{VaporFraction{MASS},Temperature,MaterialCompounds}

ϕp() = (VaporQuality(),Pressure(),SingleComponent())
ϕpx() = (VaporQuality(),Pressure(),MaterialCompounds{MOLAR,FRACTION}())
ϕpn() = (VaporQuality(),Pressure(),MaterialCompounds{MOLAR,TOTAL_AMOUNT}())

const MultiT = Union{Tuple{Temperature,AbstractSpec,MaterialCompounds},Tuple{AbstractSpec,Temperature,MaterialCompounds}}
const SingleT = Union{Tuple{Temperature,AbstractSpec,SingleComponent},Tuple{AbstractSpec,Temperature,SingleComponent}}

export SinglePT,MultiPT
export SingleVT,MultiVT
export SinglePS,MultiPS
export SinglePH,MultiPH
export SingleSatT,MultiSatT
export SingleSatP,MultiSatP
export SingleΦT,MultiΦT
export SingleΦP,MultiΦP
export SingleΦmT,MultiΦmT
export SingleΦmP,MultiΦmP
export SingleΦnT,MultiΦnT
export SingleΦnP,MultiΦnP
export SinglePV,MultiPV
export SingleT,MultiT

end


using Test
using AutoViz
using NGSIM
using NBInclude

@nbinclude(joinpath(dirname(pathof(NGSIM)), "..", "jnotebooks", "Demo.ipynb"))

# the following test requires downloading the dataset
# td1 = load_trajdata(1)
# scene = get!(Scene(500), td1, 1000)
# render(scene, ROADWAY_101)