# This code is derived from the Julia SubArray implementation (MIT license)

# L is true if the view itself supports fast linear indexing
"""
    SSubArray

Statically sized SubArray.
"""
struct SSubArray{S,T,N,P,I,L} <: StaticArray{S, T, N}
    parent::P
    indices::I
    offset1::Int       # for linear indexing and pointer, only valid when L==true
    stride1::Int       # used only for linear indexing
    function SSubArray{S,T,N,P,I,L}(parent, indices, offset1, stride1) where {S,T,N,P,I,L}
        Base.@_inline_meta
        Base.check_parent_index_match(parent, indices)
        new(parent, indices, offset1, stride1)
    end
end
# Compute the linear indexability of the indices, and combine it with the linear indexing of the parent
function SSubArray(parent::Union{StaticArray{S}, HybridArray{S}}, indices::Tuple) where S
    Base.@_inline_meta
    SSubArray(IndexStyle(Base.viewindexing(indices), IndexStyle(parent)), all_dynamic_fixed_val(S, indices...), parent, Base.ensure_indexable(indices), Base.index_dimsum(indices...))
end
function SSubArray(::IndexCartesian, ::Val{:dynamic_fixed_true}, parent::Union{StaticArray{S}, HybridArray{S}}, indices::I, ::NTuple{N,Any}) where {S,I,N}
    Base.@_inline_meta
    newsize = new_out_size(S, indices...)
    SSubArray{newsize, eltype(typeof(parent)), N, typeof(parent), I, false}(parent, indices, 0, 0)
end
function SSubArray(::IndexLinear, ::Val{:dynamic_fixed_true}, parent::Union{StaticArray{S}, HybridArray{S}}, indices::I, ::NTuple{N,Any}) where {S,I,N}
    Base.@_inline_meta
    # Compute the stride and offset
    stride1 = Base.compute_stride1(parent, indices)
    newsize = new_out_size(S, indices...)
    SSubArray{newsize, eltype(typeof(parent)), N, typeof(parent), I, true}(parent, indices, Base.compute_offset1(parent, stride1, indices), stride1)
end
function SSubArray(indexing::Any, ::Val{:dynamic_fixed_false}, parent::Union{StaticArray{S}, HybridArray{S}}, indices::I, idsum::NTuple{N,Any}) where {S,I,N}
    return SubArray(indexing, parent, indices, idsum)
end

if !isdefined(Base, :__subarray_throw_boundserror)
    @noinline __subarray_throw_boundserror(::Type{T}, parent, indices, offset1, stride1, I) where {T} =
    throw(BoundsError(T(parent, indices, offset1, stride1), I))
else
    using Base: __subarray_throw_boundserror
end

# This makes it possible to elide view allocation in cases where the
# view is indexed with a boundscheck but otherwise all its uses
# are inlined
@inline Base.throw_boundserror(A::SSubArray, I) =
    __subarray_throw_boundserror(typeof(A), A.parent, A.indices, A.offset1, A.stride1, I)

# Simple utilities
size(V::SSubArray) = (Base.@_inline_meta; map(n->Int(Base.unsafe_length(n)), axes(V)))

similar(V::SSubArray, T::Type, dims::Dims) = similar(V.parent, T, dims)

Base.sizeof(V::SSubArray) = length(V) * sizeof(eltype(V))

copy(V::SSubArray) = V.parent[V.indices...]

Base.parent(V::SSubArray) = V.parent
Base.parentindices(V::SSubArray) = V.indices

## Aliasing detection
dataids(A::SSubArray) = (dataids(A.parent)..., Base._splatmap(dataids, A.indices)...)
Base.unaliascopy(A::SSubArray) = typeof(A)(Base.unaliascopy(A.parent), map(Base.unaliascopy, A.indices), A.offset1, A.stride1)

# When the parent is an Array we can trim the size down a bit. In the future this
# could possibly be extended to any mutable array.
function Base.unaliascopy(V::SSubArray{S,T,N,A,I,LD}) where {S,T,N,A<:Array,I<:Tuple{Vararg{Union{Real,AbstractRange,Array}}},LD}
    dest = Array{T}(undef, index_lengths(V.indices...))
    copyto!(dest, V)
    SSubArray{S,T,N,A,I,LD}(dest, map(_trimmedindex, V.indices), 0, Int(LD))
end

function Base.unsafe_view(A::Union{StaticArray, HybridArray}, I::Vararg{Base.ViewIndex,<:Any})
    Base.@_inline_meta
    SSubArray(A, I)
end

Base.unsafe_view(V::SSubArray, I::Vararg{Base.ViewIndex,N}) where {N} =
    (Base.@_inline_meta; Base._maybe_reindex(V, I))
Base._maybe_reindex(V::SSubArray, I) = (Base.@_inline_meta; Base._maybe_reindex(V, I, I))
Base._maybe_reindex(V::SSubArray, I, ::Tuple{AbstractArray{<:Base.AbstractCartesianIndex}, Vararg{Any}}) =
    (Base.@_inline_meta; SSubArray(V, I))
# But allow arrays of CartesianIndex{1}; they behave just like arrays of Ints
Base._maybe_reindex(V::SSubArray, I, A::Tuple{AbstractArray{<:Base.AbstractCartesianIndex{1}}, Vararg{Any}}) =
    (Base.@_inline_meta; Base._maybe_reindex(V, I, Base.tail(A)))
Base._maybe_reindex(V::SSubArray, I, A::Tuple{Any, Vararg{Any}}) = (Base.@_inline_meta; Base._maybe_reindex(V, I, Base.tail(A)))
function Base._maybe_reindex(V::SSubArray, I, ::Tuple{})
    Base.@_inline_meta
    @inbounds idxs = to_indices(V.parent, _reindex(V.indices, I))
    return Base.unsafe_view(V.parent, idxs...)
end

# reindex was copied from Base because in pre-1.2 it has an incompatible implementation
_reindex(::Tuple{}, ::Tuple{}) = ()

# Skip dropped scalars, so simply peel them off the parent indices and continue
_reindex(idxs::Tuple{Base.ScalarIndex, Vararg{Any}}, subidxs::Tuple{Vararg{Any}}) =
    (Base.@_propagate_inbounds_meta; (idxs[1], _reindex(Base.tail(idxs), subidxs)...))

# Slices simply pass their subindices straight through
_reindex(idxs::Tuple{Base.Slice, Vararg{Any}}, subidxs::Tuple{Any, Vararg{Any}}) =
    (Base.@_propagate_inbounds_meta; (subidxs[1], _reindex(Base.tail(idxs), Base.tail(subidxs))...))

# Re-index into parent vectors with one subindex
_reindex(idxs::Tuple{AbstractVector, Vararg{Any}}, subidxs::Tuple{Any, Vararg{Any}}) =
    (Base.@_propagate_inbounds_meta; (idxs[1][subidxs[1]], _reindex(Base.tail(idxs), Base.tail(subidxs))...))

# Parent matrices are re-indexed with two sub-indices
_reindex(idxs::Tuple{AbstractMatrix, Vararg{Any}}, subidxs::Tuple{Any, Any, Vararg{Any}}) =
    (Base.@_propagate_inbounds_meta; (idxs[1][subidxs[1], subidxs[2]], _reindex(Base.tail(idxs), Base.tail(Base.tail(subidxs)))...))

# In general, we index N-dimensional parent arrays with N indices
@generated function _reindex(idxs::Tuple{AbstractArray{T,N}, Vararg{Any}}, subidxs::Tuple{Vararg{Any}}) where {T,N}
    if length(subidxs.parameters) >= N
        subs = [:(subidxs[$d]) for d in 1:N]
        tail = [:(subidxs[$d]) for d in N+1:length(subidxs.parameters)]
        :(Base.@_propagate_inbounds_meta; (idxs[1][$(subs...)], _reindex(Base.tail(idxs), ($(tail...),))...))
    else
        :(throw(ArgumentError("cannot re-index SubArray with fewer indices than dimensions\nThis should not occur; please submit a bug report.")))
    end
end

# In general, we simply re-index the parent indices by the provided ones
SlowSSubArray{S,T,N,P,I} = SSubArray{S,T,N,P,I,false}
function getindex(V::SSubArray{S,T,N}, I::Vararg{Int,N}) where {S,T,N}
    Base.@_inline_meta
    @boundscheck checkbounds(V, I...)
    @inbounds r = V.parent[_reindex(V.indices, I)...]
    r
end

# overrides default erroring getindex from `StaticArrays`
function getindex(V::SSubArray, i::Int)
    Base.@_inline_meta
    @boundscheck checkbounds(V, i)
    subind = Base._to_subscript_indices(V, i)
    @inbounds r = V.parent[_reindex(V.indices, subind)...]
    return r
end

# But SSubArrays with fast linear indexing pre-compute a stride and offset
FastSSubArray{S,T,N,P,I} = SSubArray{S,T,N,P,I,true}
function getindex(V::FastSSubArray, i::Int)
    Base.@_inline_meta
    @boundscheck checkbounds(V, i)
    @inbounds r = V.parent[V.offset1 + V.stride1*i]
    r
end

# disambiguation
function getindex(V::FastSSubArray{S,T,1}, i::Int) where {S,T}
    Base.@_inline_meta
    @boundscheck checkbounds(V, i)
    @inbounds r = V.parent[V.offset1 + V.stride1*i]
    r
end

# We can avoid a multiplication if the first parent index is a Colon or AbstractUnitRange,
# or if all the indices are scalars, i.e. the view is for a single value only
FastContiguousSSubArray{S,T,N,P,I<:Union{Tuple{Union{Base.Slice, AbstractUnitRange}, Vararg{Any}},
                                      Tuple{Vararg{Base.ScalarIndex}}}} = SSubArray{S,T,N,P,I,true}
function getindex(V::FastContiguousSSubArray, i::Int)
    Base.@_inline_meta
    @boundscheck checkbounds(V, i)
    @inbounds r = V.parent[V.offset1 + i]
    r
end
# For vector views with linear indexing, we disambiguate to favor the stride/offset
# computation as that'll generally be faster than (or just as fast as) re-indexing into a range.
function getindex(V::FastSSubArray{<:Any, 1}, i::Int)
    Base.@_inline_meta
    @boundscheck checkbounds(V, i)
    @inbounds r = V.parent[V.offset1 + V.stride1*i]
    r
end
function getindex(V::FastContiguousSSubArray{<:Any, 1}, i::Int)
    Base.@_inline_meta
    @boundscheck checkbounds(V, i)
    @inbounds r = V.parent[V.offset1 + i]
    r
end

# Indexed assignment follows the same pattern as `getindex` above
function setindex!(V::SSubArray{T,N}, x, I::Vararg{Int,N}) where {T,N}
    Base.@_inline_meta
    @boundscheck checkbounds(V, I...)
    @inbounds V.parent[_reindex(V.indices, I)...] = x
    V
end
# overrides the default setindex! from `StaticArrays`
function setindex!(V::SSubArray, x, i::Int)
    Base.@_inline_meta
    subind = Base._to_subscript_indices(V, i)
    @inbounds V.parent[_reindex(V.indices, subind)...] = x
    V
end
function setindex!(V::FastSSubArray, x, i::Int)
    Base.@_inline_meta
    @boundscheck checkbounds(V, i)
    @inbounds V.parent[V.offset1 + V.stride1*i] = x
    V
end
function setindex!(V::FastContiguousSSubArray, x, i::Int)
    Base.@_inline_meta
    @boundscheck checkbounds(V, i)
    @inbounds V.parent[V.offset1 + i] = x
    V
end
function setindex!(V::FastSSubArray{<:Any, 1}, x, i::Int)
    Base.@_inline_meta
    @boundscheck checkbounds(V, i)
    @inbounds V.parent[V.offset1 + V.stride1*i] = x
    V
end
function setindex!(V::FastContiguousSSubArray{<:Any, 1}, x, i::Int)
    Base.@_inline_meta
    @boundscheck checkbounds(V, i)
    @inbounds V.parent[V.offset1 + i] = x
    V
end

Base.IndexStyle(::Type{<:FastSSubArray}) = IndexLinear()
Base.IndexStyle(::Type{<:SSubArray}) = IndexCartesian()

if VERSION < v"1.2"
    # substrides copied from base because it's different in Julia < v1.2
    substrides(strds::Tuple{}, ::Tuple{}) = ()
    substrides(strds::NTuple{N,Int}, I::Tuple{Base.ScalarIndex, Vararg{Any}}) where N = (substrides(Base.tail(strds), Base.tail(I))...,)
    substrides(strds::NTuple{N,Int}, I::Tuple{Base.Slice, Vararg{Any}}) where N = (first(strds), substrides(Base.tail(strds), Base.tail(I))...)
    substrides(strds::NTuple{N,Int}, I::Tuple{AbstractRange, Vararg{Any}}) where N = (first(strds)*step(I[1]), substrides(Base.tail(strds), Base.tail(I))...)
    substrides(strds, I::Tuple{Any, Vararg{Any}}) = throw(ArgumentError("strides is invalid for SSubArrays with indices of type $(typeof(I[1]))"))
else
    using Base: substrides
end

# Strides are the distance in memory between adjacent elements in a given dimension
# which we determine from the strides of the parent
Base.strides(V::SSubArray) = substrides(strides(V.parent), V.indices)

Base.stride(V::SSubArray, d::Integer) = d <= ndims(V) ? strides(V)[d] : strides(V)[end] * size(V)[end]

Base.elsize(::Type{<:SSubArray{<:Any,<:Any,<:Any,P}}) where {P} = Base.elsize(P)

Base.iscontiguous(A::SSubArray) = Base.iscontiguous(typeof(A))
Base.iscontiguous(::Type{<:SSubArray}) = false
Base.iscontiguous(::Type{<:FastContiguousSSubArray}) = true

Base.first_index(V::FastSSubArray) = V.offset1 + V.stride1 # cached for fast linear SSubArrays
function Base.first_index(V::SSubArray)
    P, I = parent(V), V.indices
    s1 = Base.compute_stride1(P, I)
    s1 + Base.compute_offset1(P, s1, I)
end
Base.unsafe_convert(::Type{Ptr{T}}, V::SSubArray{T,N,P,<:Tuple{Vararg{Base.RangeIndex}}}) where {T,N,P} =
    unsafe_convert(Ptr{T}, V.parent) + (first_index(V)-1)*sizeof(T)

Base.pointer(V::FastSSubArray, i::Int) = pointer(V.parent, V.offset1 + V.stride1*i)
Base.pointer(V::FastContiguousSSubArray, i::Int) = pointer(V.parent, V.offset1 + i)
Base.pointer(V::SSubArray, i::Int) = Base._pointer(V, i)
if isdefined(Base, :_pointer)
    Base._pointer(V::SSubArray{<:Any,1}, i::Int) = pointer(V, (i,))
    Base._pointer(V::SSubArray, i::Int) = pointer(V, Base._ind2sub(axes(V), i))
end

function Base.pointer(V::SSubArray{T,N,<:Array,<:Tuple{Vararg{Base.RangeIndex}}}, is::Tuple{Vararg{Int}}) where {T,N}
    index = Base.first_index(V)
    strds = strides(V)
    for d = 1:length(is)
        index += (is[d]-1)*strds[d]
    end
    return pointer(V.parent, index)
end

if VERSION < v"1.2"
    # copied from Base because in Julia < 1.2 it works in an incompatible way
    _indices_sub(::Real, I...) = (Base.@_inline_meta; _indices_sub(I...))
    _indices_sub() = ()
    function _indices_sub(i1::AbstractArray, I...)
        Base.@_inline_meta
        (Base.unsafe_indices(i1)..., _indices_sub(I...)...)
    end
else
    using Base: _indices_sub
end

axes(S::SSubArray) = (Base.@_inline_meta; _indices_sub(S.indices...))


# This file is a part of RRRMC.jl. License is MIT: http://github.com/carlobaldassi/RRRMC.jl/LICENCE.md

module LE

using ExtractMacro
using ..Interface
using ..Common

using ...RRRMC # this is silly but it's required for correct cross-linking in docstrings, apparently

export GraphLE, GraphLocalEntropy, LEenergies

import ..Interface: energy, delta_energy, neighbors, allΔE,
                    update_cache!, delta_energy_residual,
                    cenergy, distances

mutable struct GraphLE{M,γT} <: DiscrGraph{Float64}
    N::Int
    Nk::Int
    cache::LocalFields{Int}

    function GraphLE{M,γT}(N::Integer) where {M,γT}
        isa(M, Int) || throw(ArgumentError("invalid parameter M, expected Int, given: $(typeof(M))"))
        M > 2 || throw(ArgumentError("M must be greater than 2, given: $M"))
        isa(γT, Float64) || throw(ArgumentError("invalid parameter γT, expected Float64, given: $(typeof(γT))"))
        N % (M + 1) == 0 || throw(ArgumentError("N must be divisible by M+1, given: N=$N M=$M"))
        Nk = N ÷ (M+1)

        cache = LocalFields{Int}(N)
        return new{M,γT}(N, Nk, cache)
    end
end

@doc """
    GraphLE{M,γT}(N::Integer) <: DiscrGraph

An auxiliary `DiscrGraph` used to implement the interactions in the
Local Entropy Ensemble.

It is only useful when implementing other graph types; see [`GraphLocalEntropy`](@ref).
""" -> GraphLE{M,γT}(N::Integer)

GraphLE(X::GraphLE{M,oldγ}, newγ::Float64) where {M,oldγ} = GraphLE{M,newγ}(X.N)

# @generated function ΔElist{M,γT}(::Type{GraphLE{M,γT}})
#     Expr(:tuple, ntuple(d->fk(2*(d - 1 - (M-1) >>> 0x1) - iseven(M), γT), M)...)
# end
# lstind(μ̄::Int, M::Int) = (μ̄ + M-1) >>> 0x1 + 1
#
# function getk{M,γT}(X::GraphLE{M,γT}, μ̄::Int)
#     @inbounds k = ΔElist(GraphLE{M,γT})[lstind(μ̄, M)]
#     return k
# end

function energy(X::GraphLE{M,γT}, C::Config) where {M,γT}
    # @assert X.N == C.N
    @extract X : Nk cache
    @extract cache : lfields lfields_last
    @extract C : s

    n = 0
    j = 0
    for i = 1:Nk
        j += 1
        # @assert j == 1 + (i-1) * (M+1)
        jc = j
        σc = 2s[jc] - 1

        μ = 0
        for k = 2:(M+1)
            j += 1
            # @assert j == k + (i-1) * (M+1)
            σj = 2s[j] - 1
            lfields[j] = σc * σj
            μ += σj
        end
        f = σc * μ
        lfields[jc] = f
        n -= f
    end
    cache.move_last = 0
    fill!(lfields_last, 0.0)
    return n * γT
end

function kinterval(move::Integer, M::Integer)
    j0 = move - ((move-1) % (M+1))
    j1 = j0 + M
    return j0:j1
end

function update_cache!(X::GraphLE{M,γT}, C::Config, move::Int) where {M,γT}
    # @assert X.N == C.N
    # @assert 1 ≤ move ≤ C.N
    @extract C : N s
    @extract X : Nk cache

    k = mod1(move, M+1)
    if k ≠ 1
        @inbounds σx = 2s[move] - 1
        i = (move-1) ÷ (M+1) + 1
        jc = 1 + (i-1) * (M+1)
    else
        Ux = kinterval(move, M)
    end

    @extract cache : lfields lfields_last move_last
    if move_last == move
        @inbounds begin
            #for y in neighbors(X, move)
            if k ≠ 1
                lfields[jc], lfields_last[jc] = lfields_last[jc], lfields[jc]
                lfields[move], lfields_last[move] = lfields_last[move], lfields[move]
            else
                for y in Ux
                    lfields[y], lfields_last[y] = lfields_last[y], lfields[y]
                end
            end
            #lfields[move] = -lfields[move]
            #lfields_last[move] = -lfields_last[move]
        end
        return
    end

    @inbounds begin
        if k ≠ 1
            #jc = 1 + (i-1) * (M+1)
            σc = 2s[jc] - 1
            lfc = lfields[jc]
            lfm = lfields[move]
            lfields_last[jc] = lfc
            lfields_last[move] = lfm
            lfields[jc] = lfc + 2 * (σc * σx)
            lfields[move] = -lfm
        else
            for y in Ux
                lfy = lfields[y]
                lfields_last[y] = lfy
                lfields[y] = -lfy
            end
        end
        #lfm = lfields[move]
        #lfields_last[move] = lfm
        #lfields[move] = -lfm
    end
    cache.move_last = move

    # lfields_bk = copy(lfields)
    # energy(X, C)
    # lfields_bk ≠ lfields && @show move hcat(lfields,lfields_bk) find(lfields .≠ lfields_bk)
    # @assert lfields_bk == lfields

    return
end

@inline function delta_energy(X::GraphLE{M,γT}, C::Config, move::Int) where {M,γT}
    # @assert X.N == C.N
    # @assert 1 ≤ move ≤ C.N
    @extract X : cache
    @extract cache : lfields

    @inbounds Δ = 2γT * lfields[move]
    return Δ
end

@inline function neighbors(X::GraphLE{M}, j::Int) where {M}
    r = (j-1) % (M+1)
    if r == 0
        return (j+1):(j+M)
    else
        j0 = j - r
        return j0:j0
    end
end

@generated function allΔE(::Type{GraphLE{M,γT}}) where {M,γT}
    iseven(M) ? Expr(:tuple, insert!([4*d*abs(γT) for d = 0:(M÷2)], 2, 2abs(γT))...) :
                Expr(:tuple, [2*(2d-1)*abs(γT) for d = 1:((M+1)÷2)]...)
end

# Replicate an existsing graph

mutable struct GraphLocalEntropy{M,γT,G<:AbstractGraph} <: DoubleGraph{DiscrGraph{Float64},Float64}
    N::Int
    Nk::Int
    X0::GraphLE{M,γT}
    Xc::G
    X1::Vector{G}
    Cc::Config
    C1::Vector{Config}
    function GraphLocalEntropy{M,γT,G}(N::Integer, g0::G, Gconstr, args...) where {M,γT,G}
        X0 = GraphLE{M,γT}(N)
        Nk = X0.Nk
        X1 = Array{G}(undef, M)
        Xc = g0
        for k = 1:M
            X1[k] = Gconstr(args...)
        end
        Cc = Config(Nk, init=false)
        C1 = [Config(Nk, init=false) for k = 1:M]
        return new{M,γT,G}(N, Nk, X0, Xc, X1, Cc, C1)
    end
end

"""
    GraphLocalEntropy(N::Integer, M::Integer, γ::Float64, β::Float64, Gconstr, args...) <: DoubleGraph

A `DoubleGraph` which implements a "Local Entropy" model, given any other Ising model previously defined.
This simulates a replicated system in which each replica interacts with an extra "reference"
configuration.
This kind of graph can be simulated efficiently with [`rrrMC`](@ref).

`N` is the number of spins, `M` the number of replicas, `γ` the interaction strength,
`β` the inverse temperature. `GConstr` is the (original) graph constructor, and `args` the
arguments to the contructor.

This is similar to [`GraphRobustEnsemble`](@ref RRRMC.GraphRobustEnsemble), but the reference is simulated explicitly.

See also [`LEenergies`](@ref) and [`cenergy`](@ref).
"""
function GraphLocalEntropy(Nk::Integer, M::Integer, γ::Float64, β::Float64, Gconstr, args...)
    g0 = Gconstr(args...)
    G = typeof(g0)
    return GraphLocalEntropy{M,γ/β,G}(Nk * (M+1), g0, Gconstr, args...)
end

function update_cache!(X::GraphLocalEntropy{M}, C::Config, move::Int) where {M}
    @extract X : X0 Xc X1 Cc C1
    k = mod1(move, M+1)
    i = (move - 1) ÷ (M+1) + 1

    if k == 1
        #spinflip!(Xc, Cc, i)
        spinflip!(Cc, i) # Xc cache not updated!
    else
        spinflip!(X1[k-1], C1[k-1], i)
    end

    update_cache!(X0, C, move)
end

function energy(X::GraphLocalEntropy{M}, C::Config) where {M}
    # @assert X.N == C.N
    @extract X : Nk X0 Xc X1 Cc C1
    @extract C : s

    E = energy(X0, C)

    Cc.s .= s[1:(M+1):(1 + (M+1) * (Nk-1))]
    energy(Xc, Cc)

    for k = 1:M
        C1[k].s .= s[(k+1):(M+1):(k+1 + (M+1) * (Nk-1))]
        E += energy(X1[k], C1[k])
    end

    return E
end

"""
    LEenergies(X::GraphLocalEntropy)

Returns a Vector with the individual energy (as defined by the original model)
of each replica in a [`GraphLocalEntropy`](@ref) graph.
"""
function LEenergies(X::GraphLocalEntropy{M}) where {M}
    @extract X : X1 C1
    return Float64[energy(X1[k], C1[k]) for k = 1:M]
end

function cenergy(X::GraphLocalEntropy{M}) where {M}
    @extract X : Xc Cc
    return energy(Xc, Cc)
end

function delta_energy_residual(X::GraphLocalEntropy{M}, C::Config, move::Int) where {M}
    @extract X : X1 C1

    k = mod1(move, M+1)
    k == 1 && return 0.0

    i = (move - 1) ÷ (M+1) + 1

    return delta_energy(X1[k-1], C1[k-1], i)
end

function delta_energy(X::GraphLocalEntropy, C::Config, move::Int)
    return delta_energy(X.X0, C, move) +
           delta_energy_residual(X, C, move)
end

# This is type-unstable and inefficient. On the other hand, it is
# basically only written for testing purposes...
function neighbors(X::GraphLocalEntropy{M}, i::Int) where {M}
    @extract X : X0 X1
    jts = neighbors(X0, i)

    k = mod1(i, M+1)

    k == 1 && return jts

    i1 = (i - 1) ÷ (M+1) + 1

    is = neighbors(X1[k-1], i1)

    return tuple(jts..., map(i->((i-1)*(M+1)+k), is)...)
end

function distances(X::GraphLocalEntropy{M}) where {M}
    @extract X : C1
    # @extract C : s
    # @assert Nk == C.N ÷ (M+1)

    # rs = BitVector[s[(k-1)*(M+1) + (2:(M+1):end)] for k = 1:M]

    # return [sum(s1 .⊻ s2) for s1 in rs, s2 in rs]
    return [sum(C1[k1].s .⊻ C1[k2].s) for k1 = 1:M, k2 = 1:M]
end

end # module


module RPOMDPModels

using RPOMDPs
using RPOMDPToolbox
using Distributions, SimpleProbabilitySets
using JuMP, Clp
using StaticArrays
using AutoHashEquals
using StatsBase
using RCall

importall RPOMDPs

import Base.rand!
import Base.rand
import Base.==
import Base.hash

import RPOMDPs: initial_state, generate_s, generate_o, generate_sor
import RPOMDPs: observation

include("Baby2POMDP.jl")
export
    Baby2POMDP,
    Baby2RPOMDP,
    Baby2IPOMDP,
    Baby2RIPOMDP

include("BabyRobust.jl")
export
    BabyPOMDP,
    BabyRPOMDP,
    BabyIPOMDP,
    BabyRIPOMDP

include("BabyRobustInfo.jl")
export
    BabyInfoPOMDP,
    BabyInfoRPOMDP

include("SimpleBaby2Robust.jl")
export
    SimpleBaby2RPOMDP

include("SimpleTigerRobust.jl")
export
    SimpleTigerRPOMDP

include("TigerRobust.jl")
export
    TigerPOMDP,
    TigerRPOMDP,
    TigerIPOMDP,
    TigerRIPOMDP

include("TigerRobustInfo.jl")
export
    TigerInfoPOMDP,
    TigerInfoRPOMDP

include("CyberRobust.jl")
export
    CyberPOMDP,
    CyberRPOMDP,
    CyberIPOMDP,
    CyberRIPOMDP

include("CyberRobustMisspec.jl")
export
    CyberTestIPOMDP

include("SimpleRobust.jl")
export
    SimpleIPOMDP,
    SimpleRIPOMDP

include("RockDiagnosis.jl")
export
    RockIPOMDP,
    RockRIPOMDP

export
    n_states,
    n_actions,
    n_observations,
    state_index,
    action_index,
    observation_index,
    observation,
    dynamics,
    reward,
    rewardalpha,
    transition,
    observation,
    action,
    initial_belief,
    generate_sor,
    generate_sor_worst,
    initial_state_distribution,
    initial_belief_distribution

end # module


"""

Type for storing steady state response (SSR) data.

The minimum required information for SSR recording is stored in the type.
Additional information can be stored in the `processing` field.
To facilitate processing, specific names are used in the processing dictionary.

#### Fields

* `data`: contains the recorded data
* `trigers`: contains information about timing for creation of epochs
* `system_codes`: contains system information
* `samplingrate`: the sampling rate of the data
* `modulationrate`: the modulation rate of the stimulus
* `reference_channel`: the channel the data has been referenced to
* `file_path` and `file_name`: where the file was read in from
* `channel_names`: the names of the channels
* `processing`: dictionary type to store analysis
* `header`: additional information read from the file

#### Additional `processing` fields
The following standard names are used when saving data to the processing dictionary.

* `Name`: The identifier for the participant
* `Side`: Side of stimulation
* `Carrier_Frequency`: Carrier frequency of the stimulus
* `Amplitude`: Amplitude of the stimulus
* `epochs`: The epochs extracted from the recording
* `sweeps`: The extracted sweeps from the recording

#### Example

Put an example here

```julia
s = SSR("filename")
```

"""
mutable struct SSR
    data::Array
    sensors::Array{Sensor}
    triggers::Dict
    system_codes::Dict
    samplingrate::typeof(1.0u"Hz")
    modulationrate::typeof(1.0u"Hz")
    reference_channel::Array{AbstractString,1}
    file_path::AbstractString
    file_name::AbstractString
    processing::Dict
    header::Dict
end

#######################################
#
# SSR info
#
#######################################

"""
Return the sampling rate of a steady state type.
If no type is provided, the sampling rate is returned as a floating point.

#### Example

Return the sampling rate of a recording

```julia
s = read_SSR(filename)
samplingrate(s)
```
"""
samplingrate(t, s::SSR) = convert(t, ustrip(s.samplingrate))
samplingrate(s::SSR) = samplingrate(AbstractFloat, s)

"""
Return the modulation rate of a steady state type.
If no type is provided, the modulation rate is returned as a floating point.

#### Example

Return the modulation rate of a recording

```julia
s = read_SSR(filename)
modulationrate(s)
```
"""
modulationrate(t, s::SSR) = convert(t, ustrip(s.modulationrate))
modulationrate(s::SSR) = modulationrate(AbstractFloat, s)

"""
Return the names of sensors in SSR measurement.

#### Example

```julia
s = read_SSR(filename)
channelnames(s)
```
"""
channelnames(s::SSR) = labels(s.sensors)

"""
Change the names of sensors in SSR measurement.

#### Example

```julia
s = read_SSR(filename)
channelnames(s, 1, "Fp1")
```
"""
function channelnames(s::SSR, i::Int, l::S) where {S<:AbstractString}

    s.sensors[i].label = l
    return s
end
function channelnames(s::SSR, l::AbstractVector{S}) where {S<:AbstractString}
    for li = 1:length(l)
        s = channelnames(s, li, l[li])
    end
    return s
end

#######################################
#
# Show
#
#######################################

import Base.show
function Base.show(io::IO, a::SSR)
    time_length = round.(size(a.data, 1) / samplingrate(a) / 60)
    println(
        io,
        "SSR measurement of $time_length mins with $(size(a.data,2)) channels sampled at $(a.samplingrate)",
    )
    println(io, "  Modulation frequency: $(a.modulationrate )")

    if haskey(a.processing, "Amplitude")
        println(io, "  Stimulation amplitude: $(a.processing["Amplitude"]) dB")
    end
    if haskey(a.processing, "Name")
        println(io, "  Participant name: $(a.processing["Name"] )")
    end
    if haskey(a.processing, "Side")
        println(io, "  Stimulation side: $(a.processing["Side"] )")
    end
    if haskey(a.processing, "Carrier_Frequency")
        println(io, "  Carrier frequency: $(a.processing["Carrier_Frequency"] ) Hz")
    end

end

#######################################
#
# Operate on type
#
#######################################

import Base.hcat
"""
Append one SSR type to another, simulating a longer recording.

#### Example

```julia
hcat(a, b)
```
"""
function hcat(a::SSR, b::SSR)

    if channelnames(a) != channelnames(b)
        throw(
            ArgumentError(
                string("Channels do not match $(channelnames(a)) != $(channelnames(b))"),
            ),
        )
    end

    if haskey(a.processing, "epochs")
        @warn("Epochs have already been extracted and will no longer be valid")
    end
    if haskey(a.processing, "statistics")
        @warn("Statistics have already been calculated and will no longer be valid")
    end

    @debug(
        "Appending two SSRs with $(size(a.data, 2)) .& $(size(b.data, 2)) channels and lengths $(size(a.data, 1)) $(size(b.data, 1))"
    )

    join_triggers(a, b)
    a.data = [a.data; b.data]

    return a
end

"""
Append the trigger information of one SSR type to another.
Places the trigger information at the end of first file

#### Example

```julia
join_triggers(a, b)
```
"""
function join_triggers(a, b; offset = size(a.data, 1))

    a.triggers["Index"] = [a.triggers["Index"]; (b.triggers["Index"] .+ offset)]
    a.triggers["Code"] = [a.triggers["Code"]; b.triggers["Code"]]
    a.triggers["Duration"] = [a.triggers["Duration"]'; b.triggers["Duration"]']'

    a
end

#######################################
#
# Manipulate channels
#
#######################################

"""
Add a channel to the SSR type with specified channel names.

#### Example

Add a channel called `Merged`

```julia
s = read_SSR(filename)
new_channel = mean(s.data, 2)
s = add_channel(s, new_channel, "Merged")
```
"""
function add_channel(a::SSR, data::Vector, chanLabel::AbstractString; kwargs...)

    @info("Adding channel $chanLabel")

    a.data = hcat(a.data, data)
    push!(a.sensors, Electrode(chanLabel, Talairach(NaN, NaN, NaN), Dict()))

    return a
end

"""
Remove specified channels from SSR.

#### Example

Remove channel Cz and those in the set called `EEG_Vanvooren_2014_Right`

```julia
a = read_SSR(filename)
remove_channel!(a, [EEG_Vanvooren_2014_Right, "Cz"])
```
"""
function remove_channel!(
    a::SSR,
    channel_names::Array{S};
    kwargs...,
) where {S<:AbstractString}
    @debug("Removing channels $(join(channel_names, " "))")
    remove_channel!(
        a,
        Int[something(findfirst(isequal(c), channelnames(a)), 0) for c in channel_names],
    )
end

function remove_channel!(a::SSR, channel_name::S; kwargs...) where {S<:AbstractString}
    @debug("Removing channel $(channel_name)")
    remove_channel!(a, [channel_name])
end

remove_channel!(a::SSR, channel_names::Int; kwargs...) =
    remove_channel!(a, [channel_names]; kwargs...)

function remove_channel!(a::SSR, channel_idx::Array{Int}; kwargs...)

    channel_idx = channel_idx[channel_idx.!=0]

    @debug("Removing channel(s) $channel_idx")
    if any(channel_idx .== 0)
        @warn("Failed to remove a channel")
    end

    keep_idx = [1:size(a.data)[end];]
    for c in sort(channel_idx, rev = true)
        try
            splice!(keep_idx, c)
        catch
            # Nothing
        end
    end

    if haskey(a.processing, "epochs")
        if size(a.processing["epochs"], 3) == size(a.data, 2)
            @debug("Removing channel(s) from epoch data")
            a.processing["epochs"] = a.processing["epochs"][:, :, keep_idx]
        end
    end
    if haskey(a.processing, "sweeps")
        if size(a.processing["sweeps"], 3) == size(a.data, 2)
            @debug("Removing channel(s) from sweep data")
            a.processing["sweeps"] = a.processing["sweeps"][:, :, keep_idx]
        end
    end

    a.data = a.data[:, keep_idx]
    a.sensors = a.sensors[keep_idx]

    return a
end

"""
Remove all channels except those requested from SSR.

#### Example

Remove all channels except Cz and those in the set called `EEG_Vanvooren_2014_Right`

```julia
a = read_SSR(filename)
keep_channel!(a, [EEG_Vanvooren_2014_Right, "Cz"])
```
"""
function keep_channel!(a::SSR, channel_names::Array{S}; kwargs...) where {S<:AbstractString}
    @info("Keeping channel(s) $(join(channel_names, " "))")
    keep_channel!(
        a,
        vec(
            round.(
                Int,
                [
                    something(findfirst(isequal(c), channelnames(a)), 0) for
                    c in channel_names
                ],
            ),
        ),
    )
end

function keep_channel!(a::SSR, channel_name::AbstractString; kwargs...)
    keep_channel!(a, [channel_name]; kwargs...)
end

function keep_channel!(a::SSR, channel_idx::AbstractVector{Int}; kwargs...)

    remove_channels = [1:size(a.data, 2);]

    channel_idx = sort(channel_idx, rev = true)
    for c in channel_idx
        splice!(remove_channels, c)
    end

    remove_channel!(a, remove_channels; kwargs...)
end

#######################################
#
# Trim channels
#
#######################################

"""
Trim SSR recording by removing data after `stop` specifed samples.

#### Optional Parameters

* `start` Remove samples before this value

#### Example

Remove the first 8192 samples and everything after 8192*300 samples

```julia
s = trim_channel(s, 8192*300, start=8192)
```
"""
function trim_channel(a::SSR, stop::Int; start::Int = 1, kwargs...)

    @info("Trimming $(size(a.data)[end]) channels between $start and $stop")

    a.data = a.data[start:stop, :]

    to_keep = findall((a.triggers["Index"] .>= start) .& (a.triggers["Index"] .<= stop))
    a.triggers["Index"] = a.triggers["Index"][to_keep]
    a.triggers["Duration"] = a.triggers["Duration"][to_keep]
    a.triggers["Code"] = a.triggers["Code"][to_keep]
    a.triggers["Index"] .-= (start - 1)

    to_keep =
        findall((a.system_codes["Index"] .>= start) .& (a.system_codes["Index"] .<= stop))
    a.system_codes["Index"] = a.system_codes["Index"][to_keep]
    a.system_codes["Duration"] = a.system_codes["Duration"][to_keep]
    a.system_codes["Code"] = a.system_codes["Code"][to_keep]
    a.system_codes["Index"] .-= (start - 1)

    return a
end

#######################################
#
# Merge channels
#
#######################################

"""
Merge `SSR` channels listed in `merge_Chans` and label the averaged channel as `new_name`

If multiple channels are listed then the average of those channels will be added.

#### Example

```julia
s = merge_channels(s, ["P6", "P8"], "P68")
```
"""
function merge_channels(
    a::SSR,
    merge_Chans::Array{S},
    new_name::S;
    kwargs...,
) where {S<:AbstractString}

    @debug("Number of original channels: $(length(channelnames(a)))")

    keep_idxs =
        vec([something(findfirst(isequal(i), channelnames(a)), 0) for i in merge_Chans])

    if sum(keep_idxs .== 0) > 0
        @warn(
            "Could not merge as these channels don't exist: $(join(vec(merge_Chans[keep_idxs .== 0]), " "))"
        )
        keep_idxs = keep_idxs[keep_idxs.>0]
    end

    @info("Merging channels $(join(vec(channelnames(a)[keep_idxs,:]), " "))")
    @debug("Merging channels $keep_idxs")

    a = add_channel(a, vec(mean(a.data[:, keep_idxs], dims = 2)), new_name; kwargs...)
end

function merge_channels(
    a::SSR,
    merge_Chans::S,
    new_name::S;
    kwargs...,
) where {S<:AbstractString}
    a = merge_channels(a, [merge_Chans], new_name; kwargs...)
end

#######################################
#
# Helper functions
#
#######################################

function assr_frequency(
    rounded_freq::Number;
    stimulation_samplingrate::Number = 32000,
    stimulation_frames_per_epoch::Number = 32768,
)

    round.(rounded_freq / (stimulation_samplingrate / stimulation_frames_per_epoch)) *
    stimulation_samplingrate / stimulation_frames_per_epoch
end

function assr_frequency(rounded_freq::AbstractVector)

    [assr_frequency(f) for f in rounded_freq]
end


using ..Grids
using ..MPIStateArrays
using ..DGBalanceLawDiscretizations
using ..AtmosDycore.VanillaAtmosDiscretizations

"""
    writevtk(prefix, Q::MPIStateArray, disc::DGBalanceLaw [, fieldnames])

Write a vtk file for all the fields in the state array `Q` using geometry and
connectivity information from `disc.grid`. The filename will start with `prefix`
which may also contain a directory path. The names used for each of the fields
in the vtk file can be specified through the collection of strings `fieldnames`;
if not specified the fields names will be `"Q1"` through `"Qk"` where `k` is the
number of states in `Q`, i.e., `k = size(Q,2)`.

"""
function writevtk(prefix, Q::MPIStateArray, disc::DGBalanceLaw,
                  fieldnames=nothing)
  vgeo = disc.grid.vgeo
  host_array = Array ∈ typeof(Q).parameters
  (h_vgeo, h_Q) = host_array ? (vgeo, Q.Q) : (Array(vgeo), Array(Q))
  writevtk_helper(prefix, h_vgeo, h_Q, disc.grid, fieldnames)
end

"""
    writevtk(prefix, Q::MPIStateArray, disc::DGBalanceLaw, fieldnames,
             auxstate::MPIStateArray, auxfieldnames)

Write a vtk file for all the fields in the state array `Q` and auxiliary state
`auxstate` using geometry and connectivity information from `disc.grid`. The
filename will start with `prefix` which may also contain a directory path. The
names used for each of the fields in the vtk file can be specified through the
collection of strings `fieldnames` and `auxfieldnames`.

If `fieldnames === nothing` then the fields names will be `"Q1"` through `"Qk"`
where `k` is the number of states in `Q`, i.e., `k = size(Q,2)`.

If `auxfieldnames === nothing` then the fields names will be `"aux1"` through
`"auxk"` where `k` is the number of states in `auxstate`, i.e., `k =
size(auxstate,2)`.

"""
function writevtk(prefix, Q::MPIStateArray, disc::DGBalanceLaw,
                  fieldnames, auxstate, auxfieldnames)
  vgeo = disc.grid.vgeo
  host_array = Array ∈ typeof(Q).parameters
  (h_vgeo, h_Q, h_aux) = host_array ? (vgeo, Q.Q, auxstate.Q) :
                                      (Array(vgeo), Array(Q), Array(auxstate))
  writevtk_helper(prefix, h_vgeo, h_Q, disc.grid, fieldnames, h_aux,
                  auxfieldnames)
end

"""
    writevtk_helper(prefix, vgeo::Array, Q::Array, grid, fieldnames)

Internal helper function for `writevtk`
"""
function writevtk_helper(prefix, vgeo::Array, Q::Array, grid,
                         fieldnames, auxstate=nothing, auxfieldnames=nothing)

  dim = dimensionality(grid)
  N = polynomialorder(grid)
  Nq  = N+1

  nelem = size(Q)[end]
  Xid = (grid.xid, grid.yid, grid.zid)
  X = ntuple(j->reshape((@view vgeo[:, Xid[j], :]),
                        ntuple(j->Nq, dim)...,
                        nelem), dim)
  if fieldnames == nothing
    fields = ntuple(i->("Q$i", reshape((@view Q[:, i, :]),
                                       ntuple(j->Nq, dim)..., nelem)),
                    size(Q, 2))
  else
    fields = ntuple(i->(fieldnames[i], reshape((@view Q[:, i, :]),
                                               ntuple(j->Nq, dim)..., nelem)),
                    size(Q, 2))
  end
  if auxstate !== nothing
    if auxfieldnames === nothing
      auxfields = ntuple(i->("aux$i", reshape((@view auxstate[:, i, :]),
                                              ntuple(j->Nq, dim)..., nelem)),
                         size(auxstate, 2))
    else
      auxfields = ntuple(i->(auxfieldnames[i], reshape((@view auxstate[:, i, :]),
                                                       ntuple(j->Nq, dim)...,
                                                       nelem)),
                         size(auxstate, 2))
    end
    fields = (fields..., auxfields...)
  end
  writemesh(prefix, X...; fields=fields, realelems=grid.topology.realelems)
end

function writevtk(prefix, Q::MPIStateArray, disc::VanillaAtmosDiscretization)
  vgeo = disc.grid.vgeo
  host_array = Array ∈ typeof(Q).parameters
  (vgeo, Q) = host_array ? (vgeo, Q.Q) : (Array(vgeo), Array(Q))
  writevtk_VanillaAtmosDiscretization(prefix, vgeo, Q, disc.grid)
end

function writevtk_VanillaAtmosDiscretization(prefix, vgeo::Array, Q::Array,
                                             G::Grids.AbstractGrid{T, dim, N}
                                            ) where {T, dim, N}

  Nq  = N+1

  nelem = size(Q)[end]
  Xid = (G.xid, G.yid, G.zid)
  X = ntuple(j->reshape((@view vgeo[:, Xid[j], :]),
                        ntuple(j->Nq, dim)...,
                        nelem), dim)

  _ρ = VanillaAtmosDiscretizations._ρ
  _U = VanillaAtmosDiscretizations._U
  _V = VanillaAtmosDiscretizations._V
  _W = VanillaAtmosDiscretizations._W
  _E = VanillaAtmosDiscretizations._E

  ρ = reshape((@view Q[:, _ρ, :]), ntuple(j->Nq, dim)..., nelem)
  U = reshape((@view Q[:, _U, :]), ntuple(j->Nq, dim)..., nelem)
  V = reshape((@view Q[:, _V, :]), ntuple(j->Nq, dim)..., nelem)
  W = reshape((@view Q[:, _W, :]), ntuple(j->Nq, dim)..., nelem)
  E = reshape((@view Q[:, _E, :]), ntuple(j->Nq, dim)..., nelem)
  writemesh(prefix, X...;
            fields=(("ρ", ρ), ("U", U), ("V", V), ("W", W), ("E", E)),
            realelems=G.topology.realelems)
end


"""
Extensible Framework for Quantum Algorithm Design for Humans.
"""
module Yao

export 幺

"""
Extensible Framework for Quantum Algorithm Design for Humans.

简单易用可扩展的量子算法设计框架。

幺 means normalized but not orthogonal in Chinese.
"""
const 幺 = Yao

using Reexport
@reexport using YaoBase, YaoArrayRegister, YaoBlocks, YaoSym

include("deprecations.jl")

end # module


module ColorFunctions
	using Colors
	import ..Lattices: intersect_lattice_with_plane, dimension
	using LightGraphs

	is_leaf(g::SimpleDiGraph, s) = !mapreduce(|, vertices(g)) do in; has_edge(g, in, s) end

	mutable struct ColorMapping
		inner::Dict{Any, Color}
		full::Dict{Any, Color}
	end
	ColorMapping(seed::Dict{<:Any, <:Color}) = ColorMapping(seed, copy(seed))
	min_colors(palette) = ColorMapping(Dict([0 =>colorant"white", 1 => palette[1]]))

	default_palette = distinguishable_colors(256, colorant"midnightblue")

	"""
		lighten(C, p)

		Lighten a color `C` by `p` percent. Internally converts to HSV colorspace,
		reduces saturation, increases value (brightness) and converts back to
		the orginal colorspace.
	"""
	function lighten(C,p)
	    T = typeof(C)
	    hsv = HSV(C)
	    convert(T, HSV(hsv.h, min(hsv.s * (1-0*p), 1.0), min(hsv.v * (1+p), 1.0) ))
	end

	function color_phylo!(C::ColorMapping, state; depth=1,
		 palette=default_palette, transition=(x::Color, d)->x)

		P = state.Phylogeny

		## Find all vertices which are at most `depth` removed from the root.
		## Each of which is going to be assigned its own color.
		## They form the 'inner' part.
		if depth == -1
			depth = nv(P)
		end
		ring = neighborhood(P, 1, depth, dir=:in)
		# How many colors are currently in use?
		j = length(unique(values(C.inner))) - 1

		### Prune inner colors ###
		# Assign genotypes to indices.
		inner_gs = map(x->state.meta.genotypes[x], ring)
		# If a genotype is no longer present, remove it from the color-map.
		for k in keys(C.inner)
			if k!=0 && !(k in inner_gs)
				delete!(C.inner, k)
				j -= 1
				# @info "$k pruned"
			end
		end
		# Assign colors to new genotypes within the inner ring.
		for ringv in ring
			g = state.meta.genotypes[ringv]
			if !haskey(C.inner ,g)
				push!(C.inner, g => palette[j%length(palette)+1])
			end
			j += 1
		end

		### Build a new _full_ color-map from the inner colors.
		C.full = copy(C.inner)
		# @info C.inner
		for ringv in ring
			nb = neighborhood(P, ringv, nv(P), dir=:in)
			g_parent = state.meta.genotypes[ringv]
			for (d,n) in enumerate(nb)
				g = state.meta.genotypes[n]
				newcolor = range(C.inner[g_parent], colorant"white", length=round(Int, 1.5*length(nb)))[d]
				#oc = RGBA(C.inner[g_parent])
				#newcolor = RGBA(oc.r, oc.g, oc.b, clamp( oc.alpha*(1-0.3), 0., 1.))
				push!(C.full, g => newcolor)
			end
		end

		map(reshape(state.lattice.data, length(state.lattice.data))) do s
			C.full[s]
		end
	end
	color_phylo(state; depth=1, palette=default_palette, kwargs...) = color_phylo!(min_colors(palette), state; depth=depth, kwargs...)

	function color_phylo_by_genomic_distance!(C::ColorMapping, state; depth=1,
		 palette=reverse(colormap("Blues", 20))
		)
		P = state.Phylogeny

		C.full = copy(C.inner)

		for j in 2:length(state.meta.genotypes)
			l = length(state.meta.snps[j])
			# push!(C.full, state.meta.genotypes[j] => palette[l%length(palette)+1])
			push!(C.full, state.meta.genotypes[j] => palette[min(l, 20)] )
		end

		map(reshape(state.lattice.data, state.lattice.Na^2)) do s
			C.full[s]
		end
	end
	function color_generic(state;palette=rand(RGB,maximum(state.lattice.data)))
		v0 = state.lattice.data
		map(v0) do x
			if x==0
				RGBA{Float32}(1.0,1.0,1.0,0.0)
			else
				RGBA{Float32}(palette[x])
			end
		end |> X->reshape(X, length(v0))
	end

	function color_genotype(state, g; color=colorant"blue")
		v0 = state.lattice.data
		color = RGBA(color)
		map(v0) do x
			if x!=g
				RGBA{Float32}(1.0,1.0,1.0,0.0)
			else
				color
			end
		end |> X->reshape(X, length(v0))
	end

	function color_neighbors(state, I::CartesianIndex)
		nn = neighbors(state.lattice, I)
		C = similar(state.lattice.data, RGBA{Float32})
		fill!(C, colorant"gray")

		C[nn] .= colorant"cornflowerblue"
		return reshape(C, length(C))
	end

	function color_column(state, q)
		C = similar(state.lattice.data, RGBA{Float32})
		fill!(C, colorant"gray")

		C[:,q] .= colorant"cornflowerblue"
		return reshape(C, length(C))
	end

	function color_row(state, r)
		C = similar(state.lattice.data, RGBA{Float32})
		fill!(C, colorant"gray")

		C[r,:] .= colorant"cornflowerblue"
		return reshape(C, length(C))
	end

	function color_rand(state)
		C = similar(state.lattice.data, RGBA{Float32})
		reshape(rand(RGB{Float32}, size(state.lattice.data)), length(state.lattice.data))
	end

	function hide!(v, A::BitArray)
		for i in eachindex(v)
			if !A[i]
				v[i] = RGBA{Float32}(0., 0., 0., 0.)
			end
		end
		v
	end

#= End module =#
end


#####################
# Aggregate flow circuits
# (like a flow circuit but aggregates flows over several examples and batches)
#####################

abstract type AggregateFlowΔNode{O,A} <: DecoratorΔNode{O} end
abstract type AggregateFlowLeafNode{O,A} <: AggregateFlowΔNode{O,A} end
abstract type AggregateFlowInnerNode{O,A} <: AggregateFlowΔNode{O,A} end

struct AggregateFlowLiteral{O,A} <: AggregateFlowLeafNode{O,A}
    origin::O
end

struct AggregateFlowConstant{O,A} <: AggregateFlowLeafNode{O,A}
    origin::O
end

struct AggregateFlow⋀{O,A} <: AggregateFlowInnerNode{O,A}
    origin::O
    children::Vector{<:AggregateFlowΔNode{<:O,A}}
end

mutable struct AggregateFlow⋁{O,A} <: AggregateFlowInnerNode{O,A}
    origin::O
    children::Vector{<:AggregateFlowΔNode{<:O,A}}
    aggr_flow::A
    aggr_flow_children::Vector{A}
end

const AggregateFlowΔ{O,A} = AbstractVector{<:AggregateFlowΔNode{O,A}}

#####################
# traits
#####################

@inline GateType(::Type{<:AggregateFlowLiteral}) = LiteralGate()
@inline GateType(::Type{<:AggregateFlowConstant}) = ConstantGate()

@inline GateType(::Type{<:AggregateFlow⋀}) = ⋀Gate()
@inline GateType(::Type{<:AggregateFlow⋁}) = ⋁Gate()

#####################
# constructors and conversions
#####################

function AggregateFlowΔ(circuit::Δ, ::Type{A}) where {A}
    cache = Dict{ΔNode, AggregateFlowΔNode}()
    sizehint!(cache, length(circuit)*4÷3)

    O = grapheltype(circuit) # type of node in the origin
    af_node(::LiteralGate, n::ΔNode) = AggregateFlowLiteral{O,A}(n)
    af_node(::ConstantGate, n::ΔNode) = AggregateFlowConstant{O,A}(n)

    af_node(::⋀Gate, n::ΔNode) = begin
        children = map(c -> cache[c], n.children)
        AggregateFlow⋀{O,A}(n, children)
    end

    af_node(::⋁Gate, n::ΔNode) = begin
        children = map(c -> cache[c], n.children)
        AggregateFlow⋁{O,A}(n, children, zero(A), some_vector(A, num_children(n)))
    end

    map(circuit) do node
        afn = af_node(GateType(node), node)
        cache[node] = afn
        afn
    end
end

#####################
# methods
#####################

@inline literal(n::AggregateFlowLiteral)::Lit  = literal(n.origin)
@inline constant(n::AggregateFlowConstant)::Bool = constant(n.origin)
@inline children(n::AggregateFlowInnerNode) = n.children

function collect_aggr_flows(afc::AggregateFlowΔ, batches::XBatches{Bool})
    reset_aggregate_flows(afc)
    accumulate_aggr_flows(afc, batches)
end

function collect_aggr_flows(fc::FlowΔ, batches::XBatches{Bool})
    @assert origin(fc) isa AggregateFlowΔ
    reset_aggregate_flows(origin(fc))
    accumulate_aggr_flows(fc, batches)
end

# set flow counters to zero
reset_aggregate_flows(afc::AggregateFlowΔ) = foreach(n->reset_aggregate_flow(n), afc)
reset_aggregate_flow(::AggregateFlowΔNode) = () # do nothing
reset_aggregate_flow(n::AggregateFlow⋁{O,A}) where {O,A} = (n.aggr_flow = zero(A) ; n.aggr_flow_children .= zero(A))

const opts_accumulate_flows = (flow_opts★..., compact⋁=false) #keep default options but insist on Bool flows

function accumulate_aggr_flows(afc::AggregateFlowΔ, batches::XBatches{Bool})
    fc = FlowΔ(afc, max_batch_size(batches), Bool, opts_accumulate_flows)
    accumulate_aggr_flows(fc, batches)
end

function accumulate_aggr_flows(fc::FlowΔ, batches::XBatches{Bool})
    for batch in batches
        accumulate_aggr_flows_batch(fc, batch)
    end
end

function accumulate_aggr_flows_batch(fc::FlowΔ, batch::XData{Bool})
    # pass a mini-batch through the flow circuit
    pass_up_down(fc, unweighted_data(batch))
    accumulate_aggr_flows_cached(fc, batch)
end

function accumulate_aggr_flows_cached(fc::FlowΔ, batch::XData{Bool})
    aggregate_data(xd::PlainXData, f::AbstractArray) = sum(f)
    aggregate_data(xd::PlainXData, f::AbstractArray{Bool}) = count(f)
    aggregate_data(xd::WXData, f::AbstractArray) = sum(f .* weights(xd))
    aggregate_data(xd::WXData, f::AbstractArray{Bool}) = sum(weights(xd)[f]) # see AggregateKernelBenchmark

    aggregate_data_factorized(xd::PlainXData, x1::BitVector, xs::BitVector...) = count_conjunction(x1, xs...)
    aggregate_data_factorized(xd::WXData, x1::BitVector, xs::BitVector...) = sum_weighted_product(weights(xd), x1, xs...)

    for n in fc
         # collect flows from mini-batch into aggregate statistics
        if n isa DownFlow⋁
            orig = grand_origin(n)::AggregateFlow⋁
            additional_flow = aggregate_data(batch,downflow(n))
            orig.aggr_flow += additional_flow
            # compute aggregate flows for all but the final child
            for i in 1:num_children(n)-1
                child = n.children[i]
                child_aggr_flow = aggregate_data_factorized(batch, downflow(n), pr_factors(origin(child))...)
                orig.aggr_flow_children[i] += child_aggr_flow
                additional_flow -= child_aggr_flow
            end
            # assing the remainder flow to the final child
            orig.aggr_flow_children[end] += additional_flow
            @assert isapprox(sum(orig.aggr_flow_children), orig.aggr_flow, rtol=0.0001) "Flow is leaking between parent and children: $(origin.aggr_flow) should equal $(origin.aggr_flow_children)"
            # no need to normalize flows here: we can leave this to parameter learning code
        end
    end
end


import ForwardDiff
import MacroTools
import LinearAlgebra
import Parameters: @with_kw

#######################
# trace transform DSL #
#######################

# See this math writeup for an understanding of how this code works:
# docs/tex/mcmc.pdf

"""
    TraceTransformDSLProgram

A program compiled from the [Trace Transform DSL](@ref).
"""
mutable struct TraceTransformDSLProgram
    fn!::Function
    inverse::Union{Nothing,TraceTransformDSLProgram}
end

"""
    pair_bijections!(f1::TraceTransformDSLProgram, f2::TraceTransformDSLProgram)

Assert that a pair of bijections contsructed using the [Trace Transform DSL](@ref) are inverses of one another.
"""
function pair_bijections!(f1::TraceTransformDSLProgram, f2::TraceTransformDSLProgram)
    f1.inverse = f2
    f2.inverse = f1
    return nothing
end

"""
    is_involution!(f::TraceTransformDSLProgram)

Assert that a bijection constructed with the [Trace Transform DSL](@ref) is its own inverse.
"""
function is_involution!(f::TraceTransformDSLProgram)
    f.inverse = f
    return nothing
end

"""
    b::TraceTransformDSLProgram = inverse(a::TraceTransformDSLProgram)

Obtain the inverse of a bijection that was constructed with the [Trace Transform DSL](@ref).

The inverse must have been associated with the bijection either via [`pair_bijections!`](@ref) or [`is_involution!`])(@ref).
"""
function inverse(bijection::TraceTransformDSLProgram)
    if isnothing(bijection.inverse)
        error("inverse bijection was not defined")
    end
    return bijection.inverse
end

struct ModelInputTraceToken{T}
    args::T
end

struct AuxInputTraceToken{T}
    args::T
end

struct ModelInputTraceRetValToken
end

struct AuxInputTraceRetValToken
end

struct ModelOutputTraceToken
end

struct AuxOutputTraceToken
end

struct ModelInputAddress{T}
    addr::T
end

struct AuxInputAddress{T}
    addr::T
end

struct ModelOutputAddress{T}
    addr::T
end

struct AuxOutputAddress{T}
    addr::T
end

Base.getindex(::ModelInputTraceToken, addr) = ModelInputAddress(addr) # model_in[addr]
Base.getindex(::ModelOutputTraceToken, addr) = ModelOutputAddress(addr) # model_out[addr]
Base.getindex(::AuxInputTraceToken, addr) = AuxInputAddress(addr) # aux_in[addr]
Base.getindex(::AuxOutputTraceToken, addr) = AuxOutputAddress(addr) # aux_out[addr]
Base.getindex(::ModelInputTraceToken) = ModelInputTraceRetvalToken() # model_in[]
Base.getindex(::AuxInputTraceToken) = AuxInputTraceRetValToken() # aux_in[]
get_args(token::ModelInputTraceToken) = token.args # get_args(model_in)
get_args(token::AuxInputTraceToken) = token.args # get_args(aux_in)

const bij_state = gensym("bij_state")

"""
    @transform f[(params...)] (in1 [,in2]) to (out1 [,out2])
        ..
    end

Write a program in the [Trace Transform DSL](@ref).
"""
macro transform(f_expr, src_expr, to_symbol::Symbol, dest_expr, body)
    syntax_err = "valid syntactic forms:\n@transform f (..) to (..) begin .. end\n@transform f(..) (..) to (..) begin .. end"
    err = false
    if MacroTools.@capture(f_expr, f_(args__))
    elseif MacroTools.@capture(f_expr, f_)
        args = []
    else
        err = true
    end
    err = err || (to_symbol != :to)
    if MacroTools.@capture(src_expr, (model_in_, aux_in_))
    elseif MacroTools.@capture(src_expr, (model_in_))
        aux_in = gensym("dummy_aux")
    else
        err = true
    end
    if MacroTools.@capture(dest_expr, (model_out_, aux_out_))
    elseif MacroTools.@capture(src_expr, (model_out_))
        aux_out = gensym("dummy_aux")
    else
        err = true
    end

    fn! = gensym("$(esc(f))_fn!")

    return quote

        # mutates the state
        function $fn!(
                $(esc(bij_state))::Union{FirstPassState,JacobianPassState},
                $(map(esc, args)...))
            model_args = get_args($(esc(bij_state)).model_trace)
            aux_args = get_args($(esc(bij_state)).aux_trace)
            $(esc(model_in)) = ModelInputTraceToken(model_args)
            $(esc(model_out)) = ModelOutputTraceToken()
            $(esc(aux_in)) = AuxInputTraceToken(aux_args)
            $(esc(aux_out)) = AuxOutputTraceToken()
            $(esc(body))
            return nothing
        end

        Core.@__doc__ $(esc(f)) = TraceTransformDSLProgram($fn!, nothing)

    end
end

macro tcall(ex)
    MacroTools.@capture(ex, f_(args__)) || error("expected syntax: f(..)")
    return quote $(esc(f)).fn!($(esc(bij_state)), $(map(esc, args)...)) end
end

# handlers

struct DiscreteAnn end
struct ContinuousAnn end

const DISCRETE = [:discrete, :disc]
const CONTINUOUS = [:continuous, :cont]

function typed(annotation::Symbol)
    if annotation in DISCRETE
        return DiscreteAnn()
    elseif annotation in CONTINUOUS
        return ContinuousAnn()
    else
        error("error")
    end
end

"""
    @read(<source>, <annotation>)

Macro for reading the value of a random choice from an input trace in the [Trace Transform DSL](@ref).

<source> is of the form <trace>[<addr>] where <trace> is an input trace, and <annotation> is either :discrete or :continuous.
"""
macro read(src, ann::QuoteNode)
    return quote read($(esc(bij_state)), $(esc(src)), $(esc(typed(ann.value)))) end
end

"""
    @write(<destination>, <value>, <annotation>)

Macro for writing the value of a random choice to an output trace in the [Trace Transform DSL](@ref).

<destination> is of the form <trace>[<addr>] where <trace> is an input trace, and <annotation> is either :discrete or :continuous.
"""
macro write(dest, val, ann::QuoteNode)
    return quote write($(esc(bij_state)), $(esc(dest)), $(esc(val)), $(esc(typed(ann.value)))) end
end

"""
    @copy(<source>, <destination>)

Macro for copying the value of a random choice (or a whole namespace of random choices) from an input trace to an output trace in the [Trace Transform DSL](@ref).

<destination> is of the form <trace>[<addr>] where <trace> is an input trace, and <annotation> is either :discrete or :continuous.
"""
macro copy(src, dest)
    return quote copy($(esc(bij_state)), $(esc(src)), $(esc(dest))) end
end

# TODO make more consistent by allowing us to read any hierarchical address,
# including return values of intermediate calls, not just the top-level call.

# TODO add haskey(model_in, addr), and haskey(aux_in, addr)

################################
# first pass through transform #
################################

struct FirstPassResults

    "subset of the output model choice map ``t'``"
    constraints::ChoiceMap

    "output proposal choice map ``u'``"
    u_back::ChoiceMap

    t_cont_reads::Dict
    u_cont_reads::Dict
    t_cont_writes::Dict
    u_cont_writes::Dict
    t_copy_reads::DynamicSelection
    u_copy_reads::DynamicSelection
end

function FirstPassResults()
    return FirstPassResults(
        choicemap(), choicemap(),
        Dict(), Dict(), Dict(), Dict(),
        DynamicSelection(), DynamicSelection())
end

struct FirstPassState

    "trace containing the input model choice map ``t``"
    model_trace::Trace
    model_choices::ChoiceMap

    "the input proposal choice map ``u``"
    aux_trace::Trace
    aux_choices::ChoiceMap

    results::FirstPassResults
end

function FirstPassState(model_trace, aux_trace)
    return FirstPassState(
        model_trace, get_choices(model_trace),
        aux_trace, get_choices(aux_trace), FirstPassResults())
end

function run_first_pass(transform::TraceTransformDSLProgram, model_trace, aux_trace)
    state = FirstPassState(model_trace, aux_trace)
    transform.fn!(state) # TODO allow for other args to top-level transform function
    return state.results
end

function read(state::FirstPassState, src::ModelInputTraceRetValToken, ::DiscreteAnn)
    return get_retval(state.model_trace)
end

function read(state::FirstPassState, src::AuxInputTraceRetValToken, ::DiscreteAnn)
    return get_retval(state.aux_trace)
end

function read(state::FirstPassState, src::ModelInputAddress, ::DiscreteAnn)
    addr = src.addr
    return state.model_trace[addr]
end

function read(state::FirstPassState, src::ModelInputAddress, ::ContinuousAnn)
    addr = src.addr
    state.results.t_cont_reads[addr] = state.model_choices[addr]
    return state.model_trace[addr]
end

function read(state::FirstPassState, src::AuxInputAddress, ::DiscreteAnn)
    addr = src.addr
    return state.aux_trace[addr]
end

function read(state::FirstPassState, src::AuxInputAddress, ::ContinuousAnn)
    addr = src.addr
    state.results.u_cont_reads[addr] = state.aux_choices[addr]
    return state.aux_trace[addr]
end

function write(state::FirstPassState, dest::ModelOutputAddress, value, ::DiscreteAnn)
    addr = dest.addr
    state.results.constraints[addr] = value
    return value
end

function write(state::FirstPassState, dest::ModelOutputAddress, value, ::ContinuousAnn)
    addr = dest.addr
    has_value(state.results.constraints, addr) && error("Model address $addr already written to")
    state.results.constraints[addr] = value
    state.results.t_cont_writes[addr] = value
    return value
end

function write(state::FirstPassState, dest::AuxOutputAddress, value, ::DiscreteAnn)
    addr = dest.addr
    state.results.u_back[addr] = value
    return value
end

function write(state::FirstPassState, dest::AuxOutputAddress, value, ::ContinuousAnn)
    addr = dest.addr
    has_value(state.results.u_back, addr) && error("Proposal address $addr already written to")
    state.results.u_back[addr] = value
    state.results.u_cont_writes[addr] = value
    return value
end

function copy(state::FirstPassState, src::ModelInputAddress, dest::ModelOutputAddress)
    from_addr, to_addr = src.addr, dest.addr
    model_choices = get_choices(state.model_trace)
    push!(state.results.t_copy_reads, from_addr)
    if has_value(model_choices, from_addr)
        state.results.constraints[to_addr] = model_choices[from_addr]
    else
        set_submap!(state.results.constraints, to_addr, get_submap(model_choices, from_addr))
    end
    return nothing
end

function copy(state::FirstPassState, src::ModelInputAddress, dest::AuxOutputAddress)
    from_addr, to_addr = src.addr, dest.addr
    model_choices = get_choices(state.model_trace)
    push!(state.results.t_copy_reads, from_addr)
    if has_value(model_choices, from_addr)
        state.results.u_back[to_addr] = model_choices[from_addr]
    else
        set_submap!(state.results.u_back, to_addr, get_submap(model_choices, from_addr))
    end
    return nothing
end

function copy(state::FirstPassState, src::AuxInputAddress, dest::AuxOutputAddress)
    from_addr, to_addr = src.addr, dest.addr
    push!(state.results.u_copy_reads, from_addr)
    aux_choices = get_choices(state.aux_trace)
    if has_value(aux_choices, from_addr)
        state.results.u_back[to_addr] = aux_choices[from_addr]
    else
        set_submap!(state.results.u_back, to_addr, get_submap(aux_choices, from_addr))
    end
    return nothing
end

function copy(state::FirstPassState, src::AuxInputAddress, dest::ModelOutputAddress)
    from_addr, to_addr = src.addr, dest.addr
    push!(state.results.u_copy_reads, from_addr)
    aux_choices = get_choices(state.aux_trace)
    if has_value(aux_choices, from_addr)
        state.results.constraints[to_addr] = aux_choices[from_addr]
    else
        set_submap!(state.results.constraints, to_addr, get_submap(aux_choices, from_addr))
    end
    return nothing
end

#####################################################################
# second pass through transform (gets automatically differentiated) #
#####################################################################

struct JacobianPassState{T<:Real}
    model_trace
    aux_trace
    input_arr::AbstractArray{T}
    output_arr::Array{T}
    t_key_to_index::Dict
    u_key_to_index::Dict
    cont_constraints_key_to_index::Dict
    cont_u_back_key_to_index::Dict
end

function read(state::JacobianPassState, src::ModelInputTraceRetValToken, ::DiscreteAnn)
    return get_retval(state.model_trace)
end

function read(state::JacobianPassState, src::AuxInputTraceRetValToken, ::DiscreteAnn)
    return get_retval(state.aux_trace)
end

function read(state::JacobianPassState, src::ModelInputAddress, ::DiscreteAnn)
    addr = src.addr
    return state.model_trace[addr]
end

function read(state::JacobianPassState, src::AuxInputAddress, ::DiscreteAnn)
    addr = src.addr
    return state.aux_trace[addr]
end

function _read_continuous(input_arr, addr_info::Int)
    return input_arr[addr_info]
end

function _read_continuous(input_arr, addr_info::Tuple{Int,Int})
    # TODO to handle things other than vectors, store shape in addr info and reshape?
    (start_idx, len) = addr_info
    return input_arr[start_idx:start_idx+len-1]
 end

function read(state::JacobianPassState, src::ModelInputAddress, ::ContinuousAnn)
    addr = src.addr
    if haskey(state.t_key_to_index, addr)
        return _read_continuous(state.input_arr, state.t_key_to_index[addr])
    else
        return state.model_trace[addr]
    end
end

function read(state::JacobianPassState, src::AuxInputAddress, ::ContinuousAnn)
    addr = src.addr
    if haskey(state.u_key_to_index, addr)
        return _read_continuous(state.input_arr, state.u_key_to_index[addr])
    else
        return state.aux_trace[addr]
    end
end

function write(state::JacobianPassState, dest::ModelOutputAddress, value, ::DiscreteAnn)
    return value
end

function write(state::JacobianPassState, dest::AuxOutputAddress, value, ::DiscreteAnn)
    return value
end

function _write_continuous(output_arr, addr_info::Int, value)
    return output_arr[addr_info] = value
end

function _write_continuous(output_arr, addr_info::Tuple{Int,Int}, value)
    (start_idx, len) = addr_info
    return output_arr[start_idx:start_idx+len-1] = value
end

function write(state::JacobianPassState, dest::AuxOutputAddress, value, ::ContinuousAnn)
    addr = dest.addr
    return _write_continuous(state.output_arr, state.cont_u_back_key_to_index[addr], value)
end

function write(state::JacobianPassState, dest::ModelOutputAddress, value, ::ContinuousAnn)
    addr = dest.addr
    return _write_continuous(state.output_arr, state.cont_constraints_key_to_index[addr], value)
end

function copy(state::JacobianPassState, src, dest)
    return nothing
end

#################################
# computing jacobian correction #
#################################

discard_skip_read_addr(addr, discard::ChoiceMap) = !has_value(discard, addr)
discard_skip_read_addr(addr, discard::Nothing) = false

function store_addr_info!(dict::Dict, addr, value::Real, next_index::Int)
    dict[addr] = next_index
    return 1 # number of elements of array
end

function store_addr_info!(dict::Dict, addr, value::AbstractArray{<:Real}, next_index::Int)
    len = length(value)
    dict[addr] = (next_index, len)
    return len # number of elements of array
end

function assemble_input_array_and_maps(
        t_cont_reads, t_copy_reads, u_cont_reads, u_copy_reads, discard::Union{ChoiceMap,Nothing})

    input_arr = Vector{Float64}()
    next_input_index = 1

    t_key_to_index = Dict()
    for (addr, v) in t_cont_reads
        if addr in t_copy_reads
            continue
        end
        if discard_skip_read_addr(addr, discard)
            # note: only happens when the model is unchanged
            continue
        end
        next_input_index += store_addr_info!(t_key_to_index, addr, v, next_input_index)
        append!(input_arr, v)
    end

    u_key_to_index = Dict()
    for (addr, v) in u_cont_reads
        if addr in u_copy_reads
            continue
        end
        next_input_index += store_addr_info!(u_key_to_index, addr, v, next_input_index)
        append!(input_arr, v)
    end

    return (t_key_to_index, u_key_to_index, input_arr)
end

function assemble_output_maps(t_cont_writes, u_cont_writes)
    next_output_index = 1

    cont_constraints_key_to_index = Dict()
    for (addr, v) in t_cont_writes
        next_output_index += store_addr_info!(cont_constraints_key_to_index, addr, v, next_output_index)
    end

    cont_u_back_key_to_index = Dict()
    for (addr, v) in u_cont_writes
        next_output_index += store_addr_info!(cont_u_back_key_to_index, addr, v, next_output_index)
    end

    return (cont_constraints_key_to_index, cont_u_back_key_to_index, next_output_index-1)
end

function jacobian_correction(transform::TraceTransformDSLProgram, prev_model_trace, proposal_trace, first_pass_results, discard)

    # create input array and mappings input addresses that are needed for Jacobian
    # exclude addresses that were copied explicitly to another address
    (t_key_to_index, u_key_to_index, input_arr) = assemble_input_array_and_maps(
        first_pass_results.t_cont_reads,
        first_pass_results.t_copy_reads,
        first_pass_results.u_cont_reads,
        first_pass_results.u_copy_reads, discard)

    # create mappings for output addresses that are needed for Jacobian
    (cont_constraints_key_to_index, cont_u_back_key_to_index, n_output) = assemble_output_maps(
        first_pass_results.t_cont_writes,
        first_pass_results.u_cont_writes)

    # this function is the partial application of the continuous part of the
    # transform, with inputs corresponding to a particular superset of the
    # columns of the reduced Jacobian matrix
    function f_array(input_arr::AbstractArray{T}) where {T <: Real}

        # closing over:
        # - trace, u
        # - u_key_to_index, t_key_to_index, cont_constraints_key_to_index, cont_u_back_key_to_index
        # - proposal_args, proposal_retval

        output_arr = Vector{T}(undef, n_output)

        jacobian_pass_state = JacobianPassState(
            prev_model_trace, proposal_trace, input_arr, output_arr,
            t_key_to_index, u_key_to_index,
            cont_constraints_key_to_index,
            cont_u_back_key_to_index)

        # mutates the state
        transform.fn!(jacobian_pass_state)

        # return the output array
        output_arr
    end

    # compute Jacobian matrix of f_array, where columns are inputs, rows are outputs
    J = ForwardDiff.jacobian(f_array, input_arr)
    @assert size(J)[2] == length(input_arr)
    num_outputs = size(J)[1]
    if size(J) != (num_outputs, num_outputs)
        error("Jacobian was not square (size was $(size(J)); the function may not be an bijection")
    end

    # log absolute value of Jacobian determinant
    correction = LinearAlgebra.logabsdet(J)[1]
    if isinf(correction)
        @error "Weight correction is infinite; the function may not be an bijection"
    end

    return correction
end

function check_round_trip(trace, trace_rt)
    choices = get_choices(trace)
    choices_rt = get_choices(trace_rt)
    if !isapprox(choices, choices_rt)
        @error("choices: $(sprint(show, "text/plain", choices))")
        @error("choices after round trip: $(sprint(show, "text/plain", choices_rt))")
        error("transform round trip check failed")
    end
    return nothing
end

function check_round_trip(
            prev_model_trace, prev_model_trace_rt,
            forward_proposal_trace, forward_proposal_trace_rt)
    check_round_trip(prev_model_trace, prev_model_trace_rt)
    check_round_trip(forward_proposal_trace, forward_proposal_trace_rt)
    return nothing
end

################################
# DeterministicTraceTranslator #
################################

"""
    translator = DeterministicTraceTranslator(;
        p_new::GenerativeFunction, p_args::Tuple=();
        new_observations::ChoiceMap=EmptyChoiceMap()
        f::TraceTransformDSLProgram)

Constructor for a deterministic trace translator.

Run the translator with:

    (output_trace, log_weight) = translator(input_trace)
"""
@with_kw struct DeterministicTraceTranslator
    p_new::GenerativeFunction
    p_args::Tuple = ()
    new_observations::ChoiceMap = EmptyChoiceMap()
    f::TraceTransformDSLProgram # a bijection
end

function deterministic_trace_translator_run_transform(
        f::TraceTransformDSLProgram, new_observations::ChoiceMap,
        prev_model_trace::Trace, p_new::GenerativeFunction, p_new_args::Tuple)
    first_pass_results = run_first_pass(f, prev_model_trace, nothing)
    log_abs_determinant = jacobian_correction(
        f, prev_model_trace, nothing, first_pass_results, nothing)
    constraints = merge(first_pass_results.constraints, new_observations)
    (new_model_trace, _) = generate(p_new, p_new_args, constraints)
    return (new_model_trace, log_abs_determinant)
end

"""
    (new_trace, log_weight) = (translator::DeterministicTraceTranslator)(trace)

Apply a trace translator.
"""
function (translator::DeterministicTraceTranslator)(
        prev_model_trace::Trace; check=false, prev_observations=EmptyChoiceMap())

    # apply trace transform
    (new_model_trace, log_abs_determinant) = deterministic_trace_translator_run_transform(
        translator.f, translator.new_observations, prev_model_trace, translator.p_new, translator.p_new_args)

    # compute log weight
    prev_model_score = get_score(prev_model_trace)
    new_model_score = get_score(new_model_trace)
    log_weight = new_model_score - prev_model_score + log_abs_determinant

    if check
        check_observations(get_choices(new_model_trace), observations)
        (prev_model_trace_rt, _) = deterministic_trace_translator_run_transform(
            inverse(translator.f), prev_observations, new_model_trace,
            get_gen_fn(prev_model_trace), get_args(prev_model_trace))
        check_round_trip(prev_model_trace, prev_model_trace_rt)
    end

    return (new_model_trace, log_weight)
end

##########################
# GeneralTraceTranslator #
##########################

"""
    translator = GeneralTraceTranslator(;
        p_new::GenerativeFunction,
        p_new_args::Tuple = (),
        new_observations::ChoiceMap = EmptyChoiceMap(),
        q_forward::GenerativeFunction,
        q_forward_args::Tuple  = (),
        q_backward::GenerativeFunction,
        q_backward_args::Tuple  = (),
        f::TraceTransformDSLProgram)

Constructor for a general trace translator.

Run the translator with:

    (output_trace, log_weight) = translator(input_trace; check=false, prev_observations=EmptyChoiceMap())

Use `check` to enable a bijection check (this requires that the transform `f` has been paired with its inverse using [`pair_bijections!](@ref) or [`is_involution`](@ref)).

If `check` is enabled, then `prev_observations` is a choice map containing the observed random choices in the previous trace.
"""
@with_kw struct GeneralTraceTranslator
    p_new::GenerativeFunction
    p_new_args::Tuple = ()
    new_observations::ChoiceMap = EmptyChoiceMap()
    q_forward::GenerativeFunction
    q_forward_args::Tuple  = ()
    q_backward::GenerativeFunction
    q_backward_args::Tuple  = ()
    f::TraceTransformDSLProgram # a bijection
end

function general_trace_translator_run_transform(
        f::TraceTransformDSLProgram, new_observations::ChoiceMap,
        prev_model_trace::Trace, forward_proposal_trace::Trace,
        p_new::GenerativeFunction, p_new_args::Tuple,
        q_backard::GenerativeFunction, q_backward_args::Tuple)
    first_pass_results = run_first_pass(f, prev_model_trace, forward_proposal_trace)
    log_abs_determinant = jacobian_correction(
        f, prev_model_trace, forward_proposal_trace, first_pass_results, nothing)
    constraints = merge(first_pass_results.constraints, new_observations)
    (new_model_trace, _) = generate(p_new, p_new_args, constraints)
    backward_proposal_trace, = generate(
        q_backward, (new_model_trace, q_backward_args...), first_pass_results.u_back)
    return (new_model_trace, backward_proposal_trace, log_abs_determinant)
end

function (translator::GeneralTraceTranslator)(
        prev_model_trace::Trace; check=false, prev_observations=EmptyChoiceMap())

    # sample auxiliary trace
    forward_proposal_trace = simulate(proposal, (prev_model_trace, translator.q_forward_args...,))

    # apply trace transform
    (new_model_trace, backward_proposal_trace, log_abs_determinant) = general_trace_translator_run_transform(
        translator.f, prev_model_trace, forward_proposal_trace, translator.new_observations,
        translator.p_new, translator.p_new_args, translator.q_backward, translator.q_backward_args)

    # compute log weight
    prev_model_score = get_score(prev_model_trace)
    new_model_score = get_score(new_model_trace)
    forward_proposal_score = get_score(forward_proposal_trace)
    backward_proposal_score = get_score(backward_proposal_trace)
    log_weight = new_model_score - prev_model_score + backward_proposal_score + forward_proposal_score + log_abs_determinant

    if check
        forward_proposal_choices = get_choices(forward_proposal_trace)
        f_inv = inverse(translator.f)
        (prev_model_trace_rt, forward_proposal_trace_rt, _) = general_trace_translator_run_transform(
            inverse(translator.f), prev_observations,
            new_model_trace, backward_proposal_trace,
            get_gen_fn(prev_model_trace), get_args(prev_model_trace),
            translator.q_forward, translator.q_forward_args)
        check_round_trip(
            prev_model_trace, prev_model_trace_rt,
            forward_proposal_trace, forward_proposal_trace_rt)
    end

    return (new_model_trace, log_weight)
end

##################################
# SimpleExtendingTraceTranslator #
##################################

"""
    translator = SimpleExtendingTraceTranslator(;
        p_new_args::Tuple = (),
        argdiffs::Tuple = (),
        new_obs::ChoiceMap = EmptyChoiceMap(),
        q_fwd::GenerativeFunction,
        q_fwd_args::Tuple  = ())

Constructor for a simple extending trace translator.

Run the translator with:

    (output_trace, log_weight) = translator(input_trace)
"""
@with_kw struct SimpleExtendingTraceTranslator
    p_new_args::Tuple = ()
    argdiffs::Tuple = ()
    new_obs::ChoiceMap = EmptyChoiceMap()
    q_fwd::GenerativeFunction
    q_fwd_args::Tuple  = ()
end

function (translator::SimpleExtendingTraceTranslator)(prev_model_trace::Trace)

    # simulate from auxiliary program
    forward_proposal_trace = simulate(translator.q_fwd, (prev_model_trace, translator.q_fwd_args...,))
    forward_proposal_score = get_score(forward_proposal_trace)

    # computing the new trace via update
    constraints = merge(get_choices(forward_proposal_trace), translator.new_obs)
    (new_model_trace, log_model_weight, _, discard) = update(
        prev_model_trace, translator.p_new_args,
        translator.argdiffs, constraints)

    if !isempty(discard)
        @error("can only extend the trace with random choices, cannot remove random choices")
        error("Invalid SimpleExtendingTraceTranslator")
    end

    log_weight = log_model_weight - forward_proposal_score
    return (new_model_trace, log_weight)
end

############################
# SymmetricTraceTranslator #
############################

"""
    translator = SymmetricTraceTranslator(;
        q::GenerativeFunction,
        q_args::Tuple = (),
        involution::Union{TraceTransformDSLProgram,Function})

Constructor for a symmetric trace translator.

The involution is either constructed via the [`@transform`](@ref) macro (recommended), or can be provided as a Julia function.

Run the translator with:

    (output_trace, log_weight) = translator(input_trace; check=false, observations=EmptyChoiceMap())

Use `check` to enable the involution check (this requires that the transform `f` has been marked with [`is_involution`](@ref)).

If `check` is enabled, then `observations` is a choice map containing the observed random choices, and the check will additionally ensure they are not mutated by the involution.
"""
@with_kw struct SymmetricTraceTranslator{T <: Union{TraceTransformDSLProgram,Function}}
    q::GenerativeFunction
    q_args::Tuple = ()
    involution::T # an involution
end

function symmetric_trace_translator_run_transform(
        involution::TraceTransformDSLProgram,
        prev_model_trace::Trace, forward_proposal_trace::Trace,
        q::GenerativeFunction, q_args::Tuple)
    first_pass_results = run_first_pass(involution, prev_model_trace, forward_proposal_trace)
    (new_model_trace, log_model_weight, _, discard) = update(
        prev_model_trace, get_args(prev_model_trace),
        map((_) -> NoChange(), get_args(prev_model_trace)),
        first_pass_results.constraints)
    log_abs_determinant = jacobian_correction(
        involution, prev_model_trace, forward_proposal_trace, first_pass_results, discard)
    backward_proposal_trace, = generate(
        q, (new_model_trace, q_args...), first_pass_results.u_back)
    return (new_model_trace, backward_proposal_trace, log_abs_determinant)
end

function (translator::SymmetricTraceTranslator{TraceTransformDSLProgram})(
        prev_model_trace::Trace; check=false, observations=EmptyChoiceMap())

    # simulate from auxiliary program
    forward_proposal_trace = simulate(translator.q, (prev_model_trace, translator.q_args...,))

    # apply trace transform
    (new_model_trace, backward_proposal_trace, log_abs_determinant) = symmetric_trace_translator_run_transform(
        translator.involution, prev_model_trace, forward_proposal_trace, translator.q, translator.q_args)

    # compute log weight
    prev_model_score = get_score(prev_model_trace)
    new_model_score = get_score(new_model_trace)
    forward_proposal_score = get_score(forward_proposal_trace)
    backward_proposal_score = get_score(backward_proposal_trace)
    log_weight = new_model_score - prev_model_score + backward_proposal_score - forward_proposal_score + log_abs_determinant

    if check
        check_observations(get_choices(new_model_trace), observations)
        forward_proposal_choices = get_choices(forward_proposal_trace)
        (prev_model_trace_rt, forward_proposal_trace_rt, _) = symmetric_trace_translator_run_transform(
            translator.involution, new_model_trace, backward_proposal_trace, translator.q, translator.q_args)
        check_round_trip(
            prev_model_trace, prev_model_trace_rt,
            forward_proposal_trace, forward_proposal_trace_rt)
    end

    return (new_model_trace, log_weight)
end

function (translator::SymmetricTraceTranslator{<:Function})(
        prev_model_trace::Trace; check=false, observations=EmptyChoiceMap())

    forward_trace = simulate(translator.q, (prev_model_trace, translator.q_args...,))
    forward_score = get_score(forward_trace)
    forward_choices = get_choices(forward_trace)
    forward_retval = get_retval(forward_trace)
    (new_model_trace, backward_choices, log_weight) = translator.involution(
        prev_model_trace, forward_choices, forward_retval, translator.q_args)
    (backward_score, backward_retval) = assess(translator.q, (new_model_trace, translator.q_args...), backward_choices)

    log_weight += (backward_score - forward_score)

    if check
        check_observations(get_choices(new_model_trace), observations)
        (prev_model_trace_rt, forward_choices_rt, _) = translator.involution(
            new_model_trace, backward_choices, backward_retval, translator.q_args)
        (forward_trace_rt, _) = generate(translator.q, (prev_model_trace, translator.q_args...), forward_choices_rt)
        check_round_trip(
            prev_model_trace, prev_model_trace_rt,
            forward_trace, forward_trace_rt)
    end

    return (new_model_trace, log_weight)
end

export @transform
export @read, @write, @copy, @tcall
export TraceTransformDSLProgram, pair_bijections!, is_involution!, inverse
export DeterministicTraceTranslator, SymmetricTraceTranslator, SimpleExtendingTraceTranslator, GeneralTraceTranslator


using Rasters, Test, Statistics, Dates, Plots
using Rasters.LookupArrays, Rasters.Dimensions
import ArchGDAL, NCDatasets, HDF5, CFTime
using Rasters: layerkeys, SMAPfile, FileArray

testpath = joinpath(dirname(pathof(Rasters)), "../test/")
include(joinpath(testpath, "test_utils.jl"))

# TODO example files without a login requirement
path1 = joinpath(testpath, "data/SMAP_L4_SM_gph_20160101T223000_Vv4011_001.h5")
path2 = joinpath(testpath, "data/SMAP_L4_SM_gph_20160102T223000_Vv4011_001.h5")

smapkeys = (
    :baseflow_flux, :heat_flux_ground, :heat_flux_latent, :heat_flux_sensible,
    :height_lowatmmodlay, :land_evapotranspiration_flux, :land_fraction_saturated,
    :land_fraction_snow_covered, :land_fraction_unsaturated, :land_fraction_wilting,
    :leaf_area_index, :net_downward_longwave_flux, :net_downward_shortwave_flux,
    :overland_runoff_flux, :precipitation_total_surface_flux, :radiation_longwave_absorbed_flux,
    :radiation_shortwave_downward_flux, :sm_profile, :sm_profile_pctl, :sm_profile_wetness,
    :sm_rootzone, :sm_rootzone_pctl, :sm_rootzone_wetness, :sm_surface, :sm_surface_wetness,
    :snow_depth, :snow_mass, :snow_melt_flux, :snowfall_surface_flux, :soil_temp_layer1,
    :soil_temp_layer2, :soil_temp_layer3, :soil_temp_layer4, :soil_temp_layer5, :soil_temp_layer6,
    :soil_water_infiltration_flux, :specific_humidity_lowatmmodlay, :surface_pressure, :surface_temp,
    :temp_lowatmmodlay, :vegetation_greenness_fraction, :windspeed_lowatmmodlay
)

if isfile(path1) && isfile(path2)
    # We need to wrap HDF5 for SMAP, as h5 file may not be SMAP files
    @testset "SMAPhdf5 wrapper" begin
        HDF5.h5open(path1) do f
            ds= Rasters.SMAPhdf5(f)
            @test keys(ds) == layerkeys(ds) == smapkeys
            @test dims(ds) isa Tuple{<:X,<:Y}
        end
    end

    @testset "Raster" begin
        @time smaparray = Raster(path1)

        @testset "lazyness" begin
            @time read(Raster(path1));
            @time lazyarray = Raster(path1; lazy=true);
            @time eagerarray = Raster(path1; lazy=false);
            # Lazy is the default
            @test parent(smaparray) isa FileArray
            @test parent(lazyarray) isa FileArray
            @test parent(eagerarray) isa Array
        end

        @testset "open" begin
            @test all(open(A -> A[Y=1], smaparray) .=== smaparray[:, 1])
        end

        @testset "read" begin
            @time A = read(smaparray);
            @test A isa Raster
            @test parent(A) isa Array
            A2 = zero(A)
            read!(smaparray, A2)
            A3 = zero(A)
            @time read!(path1, A3);
            @test A == A2 == A3
        end

        @testset "array properties" begin
            @test smaparray isa Raster
        end

        @testset "dimensions" begin
            @test ndims(smaparray) == 2
            HDF5.h5open(path1) do ds
                @test size(smaparray) == length.(dims(smaparray)) == size(ds["Geophysical_Data/baseflow_flux"])
            end
            @test dims(smaparray) isa Tuple{<:X,<:Y}
            @test refdims(smaparray) == ()
            @test sampling(smaparray) == (Intervals(Center()), Intervals(Center()))
            @test span(smaparray) == (Irregular((-180.0f0, 180.0f0)), Irregular((-85.04456f0, 85.04456f0)))
            @test bounds(smaparray) == ((-180.0f0, 180.0f0), (-85.04456f0, 85.04456f0))
        end

        @testset "other fields" begin @test missingval(smaparray) == -9999.0
            @test metadata(smaparray) isa Metadata{SMAPfile}
            @test name(smaparray) == :baseflow_flux
        end

        @testset "indexing" begin
            @test smaparray[Ti(1)] isa Raster{<:Any,2}
            @test smaparray[Y(1), Ti(1)] isa Raster{<:Any,1}
            @test smaparray[X(1), Ti(1)] isa Raster{<:Any,1}
            @test smaparray[X(1), Y(1)] == -9999.0
            @test smaparray[X(30), Y(30)] isa Float32
            # Russia
            @test smaparray[X(50), Y(100)] == -9999.0
            # Alaska
            @test smaparray[Y(Near(64.2008)), X(Near(149.4937))] == 0.0f0
        end

        @testset "selectors" begin
            a = smaparray[X(Near(21.0)), Y(Between(50, 52))]
            @test bounds(a) == ((50.08451f0, 51.977905f0),)
            x = smaparray[X(Near(150)), Y(Near(30))]
            @test x isa Float32
            dimz = X(Between(-180.0, 180)), Y(Between(-90, 90))
            @test size(smaparray[dimz...]) == (3856, 1624)
            @test index(smaparray[dimz...]) == index(smaparray)
            nca = smaparray[Y(Between(-80, -25)), X(Between(0, 180))]
        end

        @testset "methods" begin
            @test all(mean(smaparray; dims=Y) .=== mean(parent(smaparray); dims=2))
            @testset "trim, crop, extend" begin
                a = read(smaparray)
                a[X(1:20)] .= missingval(a)
                trimmed = trim(a)
                @test size(trimmed) == (3836, 1512)
                cropped = crop(a; to=trimmed)
                dims(trimmed, X)
                dims(cropped, X)
                @test size(cropped) == (3836, 1512)
                @test all(collect(cropped .=== trimmed))
                dims(cropped)
                extended = extend(cropped; to=a)
                @test all(collect(extended .=== a))
            end
            @testset "slice" begin
                ser = Rasters.slice(smaparray, X)
                @test ser isa RasterSeries
                @test size(ser) == (3856,)
                @test index(ser, X) == index(smaparray, X)
                @test bounds(ser, X) == (-180.0f0, 180.0f0)
                A = ser[1]
                @test bounds(A, Y) == (-85.04456f0, 85.04456f0)
            end
        end

        @testset "save" begin
            @testset "to grd" begin
                # TODO save and load subset
                geoA = read(smaparray)
                @test size(geoA) == size(smaparray)
                filename = tempname() * ".grd"
                write(filename, geoA)
                saved = read(Raster(filename; mappedcrs=EPSG(4326))[Band(1)])
                dims(saved)
                @test size(saved) == size(geoA)
                @test missingval(saved) === missingval(geoA)
                @test map(metadata.(dims(saved)), metadata.(dims(Raster))) do s, g
                    all(s .== g)
                end |> all
                @test Rasters.name(saved) == Rasters.name(geoA)
                @test all(lookup.(dims(saved)) .!= lookup.(dims(geoA)))
                @test all(
                          mappedbounds(saved)[1]
                          .≈
                          bounds(geoA)[1])
                @test all(mappedbounds(saved)[2] .≈ mappedbounds(geoA)[2])
                bounds(saved)
                @test all(parent(saved) .=== parent(geoA))
            end
            @testset "to gdal" begin
                gdalfilename = tempname() * ".tif"
                @time write(gdalfilename, read(smaparray))
                gdalarray = Raster(gdalfilename; mappedcrs=EPSG(4326))
                # These come out with slightly different format
                # @test convert(ProjString, crs(gdalarray)) == crs(smaparray)
                @test all(map((a, b) -> all(a .≈ b),
                              mappedbounds(dims(gdalarray))
                              , (bounds(smaparray)..., (1, 1))))
                # Tiff locus = Start, SMAP locus = Center
                @test mappedindex(DimensionalData.shiftlocus(Center(), dims(gdalarray, Y))) ≈ index(smaparray, Y)
                @test mappedindex(DimensionalData.shiftlocus(Center(), dims(gdalarray, X))) ≈ index(smaparray, X)
                @test all(gdalarray .== read(smaparray))
            end
            @testset "to netcdf" begin
                ncdfilename = tempname() * ".nc"
                @time write(ncdfilename, smaparray)
                reorder(smaparray, ForwardOrdered())
                saved = Raster(ncdfilename)
                @test_broken bounds(saved) == bounds(smaparray)
                @test index(saved, Y) == index(smaparray, Y)
                @test index(saved, X) == index(smaparray, X)
            end
        end

        @testset "show" begin
            sh = sprint(show, MIME("text/plain"), smaparray)
            @test occursin("Raster", sh)
            @test occursin("Y", sh)
            @test occursin("X", sh)
        end

        @testset "plot" begin
            smaparray[Ti(1:3:12)] |> plot # FIXME plot is weird
            smaparray[Ti(1)] |> plot
            smaparray[Y(100), Ti(1)] |> plot
        end

    end

    @testset "stack" begin
        @time smapstack = RasterStack(path1)

        @testset "lazyness" begin
            @time read(RasterStack(path1));
            @time lazystack = RasterStack(path1; lazy=true)
            @time eagerstack = RasterStack(path1; lazy=false);
            # Lazy is the default
            @test parent(smapstack[:snow_mass]) isa FileArray
            @test parent(lazystack[:snow_mass]) isa FileArray
            @test parent(eagerstack[:snow_mass]) isa Array
        end

        @testset "read" begin
            @time st = read(smapstack);
            @test st isa RasterStack
            @test st.data isa NamedTuple
            @test first(st.data) isa Array
            st2 = map(a -> a .* 0, st)
            @time read!(smapstack, st2);
            st3 = map(a -> a .* 0, st)
            @time st = read!(path1, st3);
            @test all(map((a, b, c) -> all(a .== b .== c), st, st2, st3))
        end

        @testset "conversion to Raster" begin
            smaparray = smapstack[:soil_temp_layer1];
            @test smaparray isa Raster{Float32,2}
            @test dims(smaparray) isa Tuple{<:X{<:Mapped{Float32}}, <:Y{<:Mapped{Float32}}}
            @test span(smaparray) isa Tuple{Irregular{Tuple{Float32,Float32}},Irregular{Tuple{Float32,Float32}}}
            @test span(smaparray) == (Irregular((-180.0f0, 180.0f0)), Irregular((-85.04456f0, 85.04456f0)))
            @test bounds(smaparray) == ((-180.0f0, 180.0f0), (-85.04456f0, 85.04456f0))
            @test index(smaparray) isa Tuple{Vector{Float32},Vector{Float32}}
            @test missingval(smaparray) == -9999.0
            @test smaparray[1, 1] == -9999.0
            @test name(smaparray) == :soil_temp_layer1
            dt = DateTime(2016, 1, 1, 22, 30)
            step_ = Hour(3)
            # Currently empty
            @test metadata(smaparray) isa Metadata{SMAPfile}
        end

        @testset "conversion to regular RasterStack" begin
            # Stack Constructors
            geostack = RasterStack(smapstack)
            @test parent(geostack) isa NamedTuple;
            @test Symbol.(Tuple(keys(smapstack))) == keys(geostack)
            geostack = RasterStack(smapstack; keys=(:baseflow_flux, :snow_mass, :soil_temp_layer1))
            @test keys(geostack) == (:baseflow_flux, :snow_mass, :soil_temp_layer1)
        end

        if VERSION > v"1.1-"
            @testset "copy" begin
                geoA = zero(smapstack[:soil_temp_layer1])
                @test geoA isa Raster
                @test geoA != smapstack[:soil_temp_layer1]
                copy!(geoA, smapstack, :soil_temp_layer1)
                @test geoA == smapstack[:soil_temp_layer1]
            end
        end

        @testset "show" begin
            sh1 = sprint(show, MIME("text/plain"), smapstack[:soil_temp_layer1])
            @test occursin("Raster", sh1)
            @test occursin("Y", sh1)
            @test occursin("X", sh1)
        end

    end

    @testset "series" begin
        ser = smapseries([path1, path2])
        val.(dims(ser))
        @test ser[1] isa RasterStack
        @test first(bounds(ser, Ti)) == DateTime(2016, 1, 1, 22, 30)
        @test last(bounds(ser, Ti)) == DateTime(2016, 1, 3, 1, 30)
        @time modified_series = modify(Array, ser)
        stackkeys = keys(modified_series[1])
        @test typeof(modified_series) <: RasterSeries{<:RasterStack{<:NamedTuple{stackkeys,<:Tuple{<:Array{Float32,2,},Vararg}}}}
        @testset "read" begin
            # FIXME: uses too much memory
            # Seems to be a HDF5.jl memory leak?
            # @time geoseries = read(ser)
            # @test geoseries isa RasterSeries{<:RasterStack}
            # @test geoseries.data isa Vector{<:RasterStack}
            # @test first(geoseries.data[1].data) isa Array
        end

        @testset "show" begin
            sh = sprint(show, MIME("text/plain"), ser)
            # Test but don't lock this down too much
            @test occursin("RasterSeries", sh)
        end
    end
end