const DimOrDimTuple = Union{Dimension,Tuple{Vararg{<:Dimension}}}
const IntOrIntTuple = Union{Int,Tuple{Vararg{<:Int}}}

"""
    aggregate(method, object, scale)

Aggregate array, or all arrays in a stack or series, by some scale.

# Arguments

- `method`: a function such as `mean` or `sum` that can combine the
  value of multiple cells to generate the aggregated cell, or a [`Locus`]($DDlocusdocs)
  like `Start()` or `Center()` that species where to sample from in the interval.
- `object`: Object to aggregate, like `AbstractGeoSeries`, `AbstractStack`,
  `AbstractGeoArray`, `Dimension`
- `scale`: the aggregation factor, which can be an integer, a tuple of integers
  for each dimension, or any `Dimension`, `Selector` or `Int` combination you can
  usually use in `getindex`. Using a `Selector` will determine the scale by the
  distance from the start of the index.
"""
function aggregate end

"""
    aggregate(method, series::AbstractGeoSeries, scale)

Aggregate an [`AbstractGeoSeries`](@ref) by `scale` using `method`.

Returns a [`GeoSeries`](ref).
"""
function aggregate(
    method, series::AbstractGeoSeries, scale, args...; progress=true, kwargs...
)
    f = i -> aggregate(method, series[i], scale, args...; progress=false, kwargs...)
    data = if progress
        @showprogress "Aggregating series..." map(f, 1:length(series))
    else
        map(f, 1:length(series))
    end
    return rebuild(series, data)
end
"""
    aggregate(method, stack::AbstractGeoStack, scale)

Aggregate an [`AbstractGeoStack`](@ref) by `scale` using `method`.

Returns a [`GeoStack`](ref).
"""
function aggregate(
    method, stack::AbstractGeoStack, scale; keys=keys(stack), progress=true
)
    f = key -> aggregate(method, stack[key], scale)
    keys_nt = NamedTuple{keys}(keys)
    data = if progress
        @showprogress "Aggregating stack..." map(f, keys_nt)
    else
        map(f, keys_nt)
    end
    return GeoStack(stack; data=data)
end

# DimensionalData methods

"""
    aggregate(method, src::AbstractDimArray, scale)

Aggregate an `AbstractDimArray` by `scale` using `method`.

[`DiskGeoArray`](@ref) will be converted to [`GeoArray`](@ref).
"""
aggregate(method, src::AbstractDimArray, scale) =
    aggregate!(method, alloc_ag(method, src, scale), src, scale)
aggregate(method, src::DiskGeoArray, scale) =
    aggregate(method, GeoArray(src), scale)

"""
    aggregate(method, dim::Dimension, scale)

Aggregate a `Dimension` by `scale` using `method`.
"""
function aggregate(method, dim::Dimension, scale)
    intscale = scale2int(dim, scale)
    intscale == 1 && return dim
    start, stop = endpoints(dim, method, intscale)
    return rebuild(dim, val(dim)[start:scale:stop], aggregate(method, mode(dim), intscale))
end

function endpoints(dim, method, scale)
    start = firstindex(dim) + agoffset(method, dim, scale)
    stop = (length(dim) ÷ scale) * scale
    return start, stop
end

"""
    aggregate(method, dim::IndexMode, scale)

Aggregate an `IndexMode` by `scale` using `method`.
"""
aggregate(method, mode::IndexMode, scale) = mode
aggregate(method, mode::AbstractSampled, scale) =
    rebuild(mode; span=aggregate(method, span(mode), scale))

"""
    aggregate(method, dim::Span, scale)

Aggregate a `Span` by `scale` using `method`.
"""
aggregate(method, span::Span, scale) = span
aggregate(method, span::Regular, scale) = Regular(val(span) * scale)

"""
    aggregate!(method, dst::AbstractDimArray, src::AbstractDimArray, scale)

Aggregate array `src` to array `dst` by `scale`, using `method`.

# Arguments

- `method`: a function such as `mean` or `sum` that can combine the
  value of multiple cells to generate the aggregated cell, or a [`Locus`]($DDlocusdocs)
  like `Start()` or `Center()` that species where to sample from in the interval.
- `scale`: the aggregation factor, which can be an integer, a tuple of integers
  for each dimension, or any `Dimension`, `Selector` or `Int` combination you can
  usually use in `getindex`. Using a `Selector` will determine the scale by the
  distance from the start of the index in the `src` array.
"""
aggregate!(locus::Locus, dst::AbstractDimArray, src, scale) =
    aggregate!((locus,), dst, src, scale)
function aggregate!(loci::Tuple{Locus,Vararg}, dst::AbstractDimArray, src, scale)
    intscale = scale2int(dims(src), scale)
    offsets = agoffset.(loci, intscale)
    for I in CartesianIndices(dst)
        dst[I] = src[(upsample.(Tuple(I), intscale) .+ offsets)...]
    end
    return dst
end
# Function/functor methods
function aggregate!(f, dst::AbstractDimArray, src, scale)
    intscale = scale2int(dims(src), scale)
    for I in CartesianIndices(dst)
        topleft = upsample.(Tuple(I), intscale)
        bottomright = topleft .+ intscale .- 1
        dst[I] = f(view(src, map(:, topleft, bottomright)...))
    end
    return dst
end

"""
    alloc_ag(method, A::AbstractDimArray, scale)

Allocate an array of the correct size to aggregate `A` by `scale`
"""
alloc_ag(method, A::AbstractDimArray, scale) = alloc_ag((method,), A, scale)
function alloc_ag(method::Tuple, A::AbstractDimArray, scale)
    intscale = scale2int(dims(A), scale)
    # Aggregate the dimensions
    dims_ = aggregate.(method, dims(A), intscale)
    # Dim aggregation determines the array size
    data_ = similar(data(A), map(length, dims_)...)
    return rebuild(A; data=data_, dims=dims_)
end

# Convert scale or tuple of scale to integer using dims2indices
scale2int(dims, scale::Tuple) = dims2indices(dims, scale)
scale2int(dims, scale::Int) = scale

"""
    disaggregate(method, object, scale)

Disaggregate array, or all arrays in a stack or series, by some scale.

# Arguments

- `method`: a function such as `mean` or `sum` that can combine the
  value of multiple cells to generate the aggregated cell, or a [`Locus`]($DDlocusdocs)
  like `Start()` or `Center()` that species where to sample from in the interval.
- `object`: Object to aggregate, like `AbstractGeoSeries`, `AbstractStack`,
  `AbstractGeoArray`, `Dimension`
- `scale`: the aggregation factor, which can be an integer, a tuple of integers
  for each dimension, or any `Dimension`, `Selector` or `Int` combination you can
  usually use in `getindex`. Using a `Selector` will determine the scale by the
  distance from the start of the index.
"""
function disaggregate end

"""
    disaggregate(method, series::AbstractGeoSeries, scale)

Disaggregate an [`AbstractGeoSeries`](@ref) by `scale` using `method`.
"""
function disaggregate(method, series::AbstractGeoSeries, scale; progress=true, kwargs...)
    f = i -> disaggregate(method, series[i], scale; progress=false, kwargs...)
    return if progress
        @showprogress "Disaggregating series..." map(f, 1:length(series))
    else
        map(f, 1:length(series))
    end
end
"""
    disaggregate(method, stack::AbstractGeoStack, scale)

Disaggregate an [`AbstractGeoStack`](@ref) by `scale` using `method`.
"""
function disaggregate(method, stack::AbstractGeoStack, scale; keys=keys(stack), progress=true)
    f = key -> disaggregate(method, stack[key], scale)
    keys_nt = NamedTuple{keys}(keys)
    data = if progress
        @showprogress "Disaggregating stack..." map(f, keys_nt)
    else
        map(f, keys_nt)
    end
    return GeoStack(stack; data=data)
end

# DimensionalData methods
"""
    disaggregate(method, src::AbstractDimArray, scale)

Disaggregate an `AbstractDimArray` by `scale` using `method`.

[`DiskGeoArray`](@ref) will be converted to [`GeoArray`](@ref).
"""
disaggregate(method, src::AbstractDimArray, scale) =
    disaggregate!(method, alloc_disag(method, src, scale), src, scale)
disaggregate(method, src::DiskGeoArray, scale) = disaggregate(method, GeoArray(src), scale)

"""
    disaggregate(method, dim::Dimension, scale)

Disaggregate a `Dimension` by `scale` using `method`.
"""
function disaggregate(locus::Locus, dim::Dimension, scale)
    intscale = scale2int(dim, scale)
    intscale == 1 && return dim
    len = length(dim) * intscale
    step_ = step(mode(dim)) / intscale
    start = dim[1] - agoffset(locus, intscale) * step_
    stop = start + (len - 1)  * step_
    index = LinRange(start, stop, len)
    return rebuild(dim, index, disaggregate(locus, mode(dim), intscale))
end
disag_index(locus::Start, dim, scale) =
    LinRange(dim[1], dim[end] + (scale - 1) * step(dim), length(dim) * scale)
disag_index(locus::End, dim, scale) =
    LinRange(dim[1], dim[end] + (scale - agoffset(locus)) * step(dim), length(dim) * scale)

"""
    disaggregate(method, dim::IndexMode, scale)

Disaggregate an `IndexMode` by `scale` using `method`.
"""
disaggregate(method, mode::IndexMode, scale) = mode
disaggregate(method, mode::AbstractSampled, scale) =
    rebuild(mode; span=disaggregate(method, span(mode), scale))

"""
    disaggregate(method, dim::Span, scale)

Disaggregate a `Span` by `scale` using `method`.
"""
disaggregate(method, span::Span, scale) = span
disaggregate(method, span::Regular, scale) = Regular(val(span) / scale)

"""
    disaggregate!(method, dst::AbstractDimArray, src::AbstractDimArray, scale)

Disaggregate array `src` to array `dst` by some scale, using `method`.

- `method`: a function such as `mean` or `sum` that can combine the
  value of multiple cells to generate the aggregated cell, or a [`Locus`]($DDlocusdocs)
  like `Start()` or `Center()` that species where to sample from in the interval.
- `scale`: the aggregation factor, which can be an integer, a tuple of integers
  for each dimension, or any `Dimension`, `Selector` or `Int` combination you can
  usually use in `getindex`. Using a `Selector` will determine the scale by the
  distance from the start of the index in the `src` array.
"""
disaggregate!(locus::Locus, dst::AbstractDimArray, src, scale) =
    disaggregate!((locus,), dst, src, scale)
function disaggregate!(loci::Tuple{Locus,Vararg}, dst::AbstractDimArray, src, scale)
    intscale = scale2int(dims(src), scale)
    for I in CartesianIndices(dst)
        dst[I] = src[(downsample.(Tuple(I), intscale))...]
    end
    return dst
end

"""
    alloc_ag(method, A::AbstractDimArray, scale)

Allocate an array of the correct size to disaggregate `A` by `scale`
"""
alloc_disag(method, A::AbstractDimArray, scale) = alloc_disag((method,), A, scale)
function alloc_disag(method::Tuple, A::AbstractDimArray, scale)
    intscale = scale2int(dims(A), scale)
    dims_ = disaggregate.(method, dims(A), intscale)
    # Dim aggregation determines the array size
    data_ = similar(data(A), map(length, dims_)...)
    return rebuild(A; data=data_, dims=dims_)
end

"""
    upsample(index::Int, scale::Int)

Convert indicies from the aggregated array to the larger original array.
"""
upsample(index::Int, scale::Int) = (index - 1) * scale + 1

"""
    downsample(index::Int, scale::Int)

Convert indicies from the original array to the aggregated array.
"""
downsample(index::Int, scale::Int) = (index - 1) ÷ scale + 1

agoffset(locus::Locus, dim::Dimension, scale) = agoffset(locus, scale)
agoffset(method, dim::Dimension, scale) = agoffset(locus(dim), scale)
agoffset(locus::Start, scale) = 0
agoffset(locus::End, scale) = scale - 1
agoffset(locus::Center, scale) = scale ÷ 2


using Parameters
using DifferentialEquations
using ForwardDiff

abstract type AbstractReactor end
export AbstractReactor

struct Reactor{D<:AbstractDomain} <: AbstractReactor
    domain::D
    ode::ODEProblem
end

function Reactor(domain::T,y0::Array{W,1},tspan::Tuple) where {T<:AbstractDomain,W<:Real}
    dydt(y::Array{T,1},p::V,t::Q) where {T<:Real,Q<:Real,V} = dydtreactor!(y,t,domain)
    ode = ODEProblem(dydt,y0,tspan)
    return Reactor(domain,ode)
end
export Reactor

@inline function getrate(rxn::T,cs::Array{W,1},kfs::Array{Q,1},krevs::Array{Q,1}) where {T<:AbstractReaction,Q,W<:Real}
    Nreact = length(rxn.reactantinds)
    Nprod = length(rxn.productinds)
    R = 0.0
    if Nreact == 1
        @fastmath @inbounds R += kfs[rxn.index]*cs[rxn.reactantinds[1]]
    elseif Nreact == 2
        @fastmath @inbounds R += kfs[rxn.index]*cs[rxn.reactantinds[1]]*cs[rxn.reactantinds[2]]
    elseif Nreact == 3
        @fastmath @inbounds R += kfs[rxn.index]*cs[rxn.reactantinds[1]]*cs[rxn.reactantinds[2]]*cs[rxn.reactantinds[3]]
    end

    if Nprod == 1
        @fastmath @inbounds R -= krevs[rxn.index]*cs[rxn.productinds[1]]
    elseif Nprod == 2
        @fastmath @inbounds R -= krevs[rxn.index]*cs[rxn.productinds[1]]*cs[rxn.productinds[2]]
    elseif Nprod == 3
        @fastmath @inbounds R -= krevs[rxn.index]*cs[rxn.productinds[1]]*cs[rxn.productinds[2]]*cs[rxn.productinds[3]]
    end

    return R
end
export getrate

@inline function addreactionratecontributions!(dydt::Array{Q,1},rarray::Array{UInt16,2},cs::Array{W,1},kfs::Array{Z,1},krevs::Array{Z,1}) where {Q<:Real,Z<:Real,T<:Integer,W<:Real}
    @inbounds @simd for i = 1:size(rarray)[2]
        if @inbounds rarray[2,i] == 0
            @inbounds @fastmath fR = kfs[i]*cs[rarray[1,i]]
        elseif @inbounds rarray[3,i] == 0
            @inbounds @fastmath fR = kfs[i]*cs[rarray[1,i]]*cs[rarray[2,i]]
        else
            @inbounds @fastmath fR = kfs[i]*cs[rarray[1,i]]*cs[rarray[2,i]]*cs[rarray[3,i]]
        end
        if @inbounds rarray[5,i] == 0
            @inbounds @fastmath rR = krevs[i]*cs[rarray[4,i]]
        elseif @inbounds rarray[6,i] == 0
            @inbounds @fastmath rR = krevs[i]*cs[rarray[4,i]]*cs[rarray[5,i]]
        else
            @inbounds @fastmath rR = krevs[i]*cs[rarray[4,i]]*cs[rarray[5,i]]*cs[rarray[6,i]]
        end
        @fastmath R = fR - rR
        @inbounds @fastmath dydt[rarray[1,i]] -= R
        if @inbounds rarray[2,i] != 0
            @inbounds @fastmath dydt[rarray[2,i]] -= R
            if @inbounds rarray[3,i] != 0
                @inbounds @fastmath dydt[rarray[3,i]] -= R
            end
        end
        @inbounds @fastmath dydt[rarray[4,i]] += R
        if @inbounds rarray[5,i] != 0
            @inbounds @fastmath dydt[rarray[5,i]] += R
            if @inbounds rarray[6,i] != 0
                @inbounds @fastmath dydt[rarray[6,i]] += R
            end
        end
    end
end
export addreactionratecontributions!

@inline function dydtreactor!(y::Array{U,1},t::Z,domain::Q;sensitivity::Bool=true) where {Z<:Real,U<:Real,J<:Integer,Q<:AbstractDomain}
    dydt = zeros(U,length(y))
    if sensitivity #if sensitivity isn't explicitly set to false set it to domain.sensitivity
        sensitivity = domain.sensitivity
    end
    ns,cs,T,P,V,C,N,mu,kfs,krevs,Hs,Us,Gs,diffs,Cvave = calcthermo(domain,y,t)
    addreactionratecontributions!(dydt,domain.rxnarray,cs,kfs,krevs)
    dydt *= V
    if sensitivity && isa(domain,ConstantVDomain)
        wV = copy(dydt[domain.indexes[1]:domain.indexes[2]])
    else
        wV = Array{U,1}()
    end
    calcdomainderivatives!(domain,dydt;t=t,T=T,P=P,Us=Us,V=V,C=C,ns=ns,N=N,Cvave=Cvave)
    if sensitivity
        Nspcs = length(cs)
        Nrxns = length(domain.phase.reactions)
        jacobian!(domain.jacobian,y,nothing,t,domain)
        dgdk = ratederivative(domain;cs=cs,V=V,T=T,kfs=kfs,Us=Us,Cvave=Cvave,N=N,krevs=krevs,wV=wV,sparse=false)
        for j  = 1:Nspcs+Nrxns #kfs and Gfs
            for i = 1:(domain.indexes[end]-domain.indexes[1]+1) #species
                for z in 1:(domain.indexes[end]-domain.indexes[1]+1)
                    dydt[j*Nspcs+i] += domain.jacobian[i,z]*y[j*Nspcs+z]
                end
                dydt[j*Nspcs+i] += dgdk[i,j]
            end
        end
    end
    return dydt
end
export dydtreactor!

function jacobian!(J::Q,y::U,p::W,t::Z,domain::V) where {Q<:AbstractArray,U<:AbstractArray,W,Z<:Real,V<:AbstractDomain}
    if domain.t[1] == t && domain.jacuptodate == true
        return
    else
        f(y::Array{T,1}) where {T<:Real} = dydtreactor!(y,domain.t[1],domain;sensitivity=false)
        ForwardDiff.jacobian!(domain.jacobian,f,y[1:domain.indexes[end]-domain.indexes[1]+1])
        domain.jacuptodate[1] = true
        return
    end
end
export jacobian!

function ratederivative(d::W; cs::Q,V::Y,T::Y2,Us::Z3,Cvave::Z4,N::Z5,kfs::Z,krevs::X,wV::Q2,sparse::Bool=false) where {W<:Union{ConstantTPDomain,ConstantTVDomain},Z4<:Real,Z5<:Real,Z3<:AbstractArray,Q2<:AbstractArray,Q<:AbstractArray,Y2<:Real,Y<:Real,Z<:AbstractArray,X<:AbstractArray}
    Nspcs = length(cs)
    rxns = d.phase.reactions
    Nrxns = length(rxns)

    if sparse
        ratederiv = spzeros(Nspcs,Nspcs+Nrxns)
    else
        ratederiv = zeros(Nspcs,Nspcs+Nrxns)
    end
    RTinv = 1.0/(R*T)

    for (j,rxn) in enumerate(rxns)
        Nreact = length(rxn.reactantinds)
        Nprod = length(rxn.productinds)

        if Nreact == 1
            rind1 = rxn.reactantinds[1]
            fderiv = cs[rind1]
        elseif Nreact == 2
            rind1,rind2 = rxn.reactantinds
            fderiv = cs[rind1]*cs[rind2]
        else
            rind1,rind2,rind3 = rxn.reactantinds
            fderiv = cs[rind1]*cs[rind2]*cs[rind3]
        end

        if Nprod == 1
            pind1 = rxn.productinds[1]
            rderiv = krevs[j]/kfs[j]*cs[pind1]
        elseif Nprod == 2
            pind1,pind2 = rxn.productinds
            rderiv = krevs[j]/kfs[j]*cs[pind1]*cs[pind2]
        else
            pind1,pind2,pind3 = rxn.productinds
            rderiv = krevs[j]/kfs[j]*cs[pind1]*cs[pind2]*cs[pind3]
        end

        flux = fderiv-rderiv
        gderiv = rderiv*kfs[j]*RTinv

        deriv = zeros(Nspcs)

        deriv[rind1] += gderiv
        if Nreact > 1
            deriv[rind2] += gderiv
            if Nreact > 2
                deriv[rind3] == gderiv
            end
        end

        deriv[pind1] -= gderiv
        if Nprod > 1
            deriv[pind2] -= gderiv
            if Nprod > 2
                deriv[pind3] -= gderiv
            end
        end

        ratederiv[rind1,j] -= flux
        ratederiv[rind1,Nrxns+1:Nrxns+Nspcs] .-= deriv
        if Nreact > 1
            ratederiv[rind2,j] -= flux
            ratederiv[rind2,Nrxns+1:Nrxns+Nspcs] .-= deriv
            if Nreact > 2
                ratederiv[rind3,j] -= flux
                ratederiv[rind3,Nrxns+1:Nrxns+Nspcs] .-= deriv
            end
        end

        ratederiv[pind1,j] += flux
        ratederiv[pind1,Nrxns+1:Nrxns+Nspcs] .+= deriv
        if Nprod > 1
            ratederiv[pind2,j] += flux
            ratederiv[pind2,Nrxns+1:Nrxns+Nspcs] .+= deriv
            if Nprod > 2
                ratederiv[pind3,j] += flux
                ratederiv[pind3,Nrxns+1:Nrxns+Nspcs] .+= deriv
            end
        end
    end
    return V*ratederiv
end

function ratederivative(d::ConstantVDomain; cs::Q,V::Y,T::Y2,Us::Z3,Cvave::Y3,N::Y2,kfs::Z,krevs::X,wV::Q2,sparse::Bool=false) where {Z3<:AbstractArray,Q<:AbstractArray,Q2<:AbstractArray,Y3<:Real,Y2<:Real,Y<:Real,Z<:AbstractArray,X<:AbstractArray}
    Nspcs = length(cs)
    rxns = d.phase.reactions
    Nrxns = length(rxns)

    if sparse
        ratederiv = spzeros(Nspcs+1,Nspcs+Nrxns)
    else
        ratederiv = zeros(Nspcs+1,Nspcs+Nrxns)
    end
    RTinv = 1.0/(R*T)

    for (j,rxn) in enumerate(rxns)
        Nreact = length(rxn.reactantinds)
        Nprod = length(rxn.productinds)

        if Nreact == 1
            rind1 = rxn.reactantinds[1]
            fderiv = cs[rind1]
        elseif Nreact == 2
            rind1,rind2 = rxn.reactantinds
            fderiv = cs[rind1]*cs[rind2]
        else
            rind1,rind2,rind3 = rxn.reactantinds
            fderiv = cs[rind1]*cs[rind2]*cs[rind3]
        end

        if Nprod == 1
            pind1 = rxn.productinds[1]
            rderiv = krevs[j]/kfs[j]*cs[pind1]
        elseif Nprod == 2
            pind1,pind2 = rxn.productinds
            rderiv = krevs[j]/kfs[j]*cs[pind1]*cs[pind2]
        else
            pind1,pind2,pind3 = rxn.productinds
            rderiv = krevs[j]/kfs[j]*cs[pind1]*cs[pind2]*cs[pind3]
        end

        flux = fderiv-rderiv
        gderiv = rderiv*kfs[j]*RTinv

        deriv = zeros(Nspcs)

        deriv[rind1] += gderiv
        if Nreact > 1
            deriv[rind2] += gderiv
            if Nreact > 2
                deriv[rind3] == gderiv
            end
        end

        deriv[pind1] -= gderiv
        if Nprod > 1
            deriv[pind2] -= gderiv
            if Nprod > 2
                deriv[pind3] -= gderiv
            end
        end

        ratederiv[rind1,j] -= flux
        ratederiv[rind1,Nrxns+1:Nrxns+Nspcs] .-= deriv
        if Nreact > 1
            ratederiv[rind2,j] -= flux
            ratederiv[rind2,Nrxns+1:Nrxns+Nspcs] .-= deriv
            if Nreact > 2
                ratederiv[rind3,j] -= flux
                ratederiv[rind3,Nrxns+1:Nrxns+Nspcs] .-= deriv
            end
        end

        ratederiv[pind1,j] += flux
        ratederiv[pind1,Nrxns+1:Nrxns+Nspcs] .+= deriv
        if Nprod > 1
            ratederiv[pind2,j] += flux
            ratederiv[pind2,Nrxns+1:Nrxns+Nspcs] .+= deriv
            if Nprod > 2
                ratederiv[pind3,j] += flux
                ratederiv[pind3,Nrxns+1:Nrxns+Nspcs] .+= deriv
            end
        end
    end
    ratederiv *= V
    #Temperature stuff
    @views ratederiv[end,:] += (Us'*ratederiv[1:end-1,:])[1,:]
    ratederiv[end,1:Nspcs] += wV
    ratederiv[end,:] /= (N*Cvave)
    return ratederiv
end
export ratederivative


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


struct WeekdayHolidayObsModel{T,M,A,P,E,W,H,V} <: CovidSurvey.ObservationsModel
    ps::T
    μ::M
    α::A
    ϕ::P
	expecteds::E
    weekdayeffect::W
    holidayeffect::H
	vectorize::V
end

struct SimpleObsModel{T,M,A,P,E,V} <: CovidSurvey.ObservationsModel
    ps::T
    μ::M
    α::A
    ϕ::P
	expecteds::E
	vectorize::V
end

struct SimpleObsModel2{T,M,S,A,P,E,V} <: CovidSurvey.ObservationsModel
    ps::T
    μ::M
	σ::S
    α::A
    ϕ::P
	expecteds::E
	vectorize::V
end

rescale(xs, σ, μs) = @. xs * σ + μs'
# ============================================================================
# random walk model
function random_walks!(Rts, θ, predict, latent_Rts::AbstractVector, R0s, σ_rt)
	@unpack num_regions, rt_step_indices, lockdown_indices, num_observations, link, include_early_dynamic = θ #include_early_dynamic

	for m in 1:num_regions
		rt_step_index = rt_step_indices[m]
		Rt            = Rts[m]
		lockdown      = lockdown_indices[m]
		num_obs       = num_observations[m]
		latent_Rt     = latent_Rts[m]#latent_Rts[:,m] #include_early_dynamic ? latent_Rts[m] : latent_Rts[:,m]
		R0            = R0s[m]

		if include_early_dynamic
			Rt[1:num_obs] = link.(latent_Rt[rt_step_index])
		else
			Rt[1:lockdown] .= R0
			Rt[lockdown+1:num_obs] = link.(latent_Rt[rt_step_index])
		end
		if predict
			Rt[num_obs+1:end] .= Rt[num_obs]
		end
	end
end

function random_walks!(Rts, θ, predict, latent_Rts::AbstractMatrix, R0s, σ_rt)
	@unpack num_regions, rt_step_indices, lockdown_indices, num_observations, link, include_early_dynamic = θ

	for m in 1:num_regions
		rt_step_index = rt_step_indices[m]
		Rt            = Rts[m]
		lockdown      = lockdown_indices[m]
		num_obs       = num_observations[m]
		latent_Rt     = latent_Rts[:,m] #include_early_dynamic ? latent_Rts[m] : latent_Rts[:,m]
		R0            = R0s[m]

		if include_early_dynamic
			Rt[1:num_obs] = link.(latent_Rt[rt_step_index])
		else
			Rt[1:lockdown] .= R0
			Rt[lockdown+1:num_obs] = link.(latent_Rt[rt_step_index])
		end
		if predict
			Rt[num_obs+1:end] .= Rt[num_obs]
		end
	end
end

# function random_walks!(Rts, θ, predict, latent_Rts, R0s)
# 	@unpack num_regions, rt_step_indices, lockdown_indices, num_observations, link = θ
#
# 	for m in 1:num_regions
# 		rt_step_index = rt_step_indices[m]
# 		Rt            = Rts[m]
# 		lockdown      = lockdown_indices[m]
# 		num_obs       = num_observations[m]
# 		latent_Rt     = latent_Rts[:,m]
# 		R0            = R0s[m]
#
# 		Rt[1:lockdown] .= R0
# 		Rt[lockdown+1:num_obs] = link.(latent_Rt[rt_step_index])
# 		if predict
# 			Rt[num_obs+1:end] .= Rt[num_obs]
# 		end
# 	end
# end

# ============================================================================
# semi-parametric model
function semiparametric!(Rts, θ, predict, latent_Rts, R0s, σ_rt, effects)
	@unpack num_regions, lockdown_indices, num_observations, link, rt_step_indices, covariates, covariates_starts, num_total_days, include_early_dynamic= θ #rt_step_indices

	for m in 1:num_regions
		Rt            = Rts[m]
		lockdown      = lockdown_indices[m]
		rt_step_index = rt_step_indices[m]
		num_obs       = num_observations[m]
		latent_Rt     = include_early_dynamic ? latent_Rts[m] : latent_Rts[:,m] #latent_Rts[m]
		R0            = R0s[m]
		num_total     = num_total_days[m]
		covariate     = covariates[m]
		effect        = effects[m]
		covariates_start = covariates_starts[m]

		if include_early_dynamic
			Rt[1:covariates_start-1] = link.(latent_Rt[rt_step_index])
		else
			Rt[1:lockdown] .= R0
			Rt[lockdown+1:covariates_start-1] = link.(latent_Rt[rt_step_index])
		end
		Rt[covariates_start:num_obs] = link.(
			latent_Rt[end] .+ covariate * effect
		)
		if predict
			Rt[num_obs+1:num_total] .= Rt[num_obs]
			# @unpack num_total_days, invlink = θ
			# n  = num_total_days - num_observations
			# n < 1 && return nothing
			# rw = RandomWalk(n, σ_rt, invlink(Rt[num_observations]))
			# Rt[num_observations+1:end] = link.( rand( rw ) )
		end
	end
	return nothing
end

# ============================================================================
# semi-parametric model
function mixed!(Rts, θ, predict, latent_Rts, R0s, σ_rt, effects)
	@unpack (
		num_regions,
		lockdown_indices,
		num_observations,
		link,
		rt_step_indices,
		covariates,
		covariates_starts,
		num_total_days,
		include_early_dynamic
	) = θ #rt_step_indices

	for m in 1:num_regions
		Rt               = Rts[m]
		lockdown_index   = lockdown_indices[m]
		rt_step_index    = rt_step_indices[m]
		num_obs          = num_observations[m]
		latent_Rt        = include_early_dynamic ? latent_Rts[m] : latent_Rts[:,m]#latent_Rts[m]
		R0               = R0s[m]
		num_total        = num_total_days[m]
		covariate        = covariates[m]
		effect           = effects[m]
		covariates_start = covariates_starts[m]
		predictor        = covariate * effect

		if include_early_dynamic
			for i in 1:covariates_start-1
				Rt[i] = link(latent_Rt[rt_step_index[i]])
			end
			for (j,i) in enumerate(covariates_start:num_obs)
				Rt[i] = link(latent_Rt[rt_step_index[i]] + predictor[j])
			end
		else
			Rt[1:lockdown_index] .= R0
			for (j,i) in enumerate(lockdown_index+1:covariates_start-1)
				Rt[i] = link(latent_Rt[rt_step_index[j]])
			end
			n = covariates_start - lockdown_index - 1
			for (j,i) in enumerate(covariates_start:num_obs)
				Rt[i] = link(latent_Rt[rt_step_index[j+n]] + predictor[j])
			end
		end
		if predict
			Rt[num_obs+1:num_total] .= Rt[num_obs]
			# @unpack num_total_days, invlink = θ
			# n  = num_total_days - num_observations
			# n < 1 && return nothing
			# rw = RandomWalk(n, σ_rt, invlink(Rt[num_observations]))
			# Rt[num_observations+1:end] = link.( rand( rw ) )
		end
	end
	return nothing
end

# ============================================================================
# infection model
function infections!(newly_infecteds, cumulative_infecteds, θ, ys, Rts)
	@unpack num_regions, populations = θ
	for m in 1:num_regions
		regionaldata = (
			Rt                     = Rts[m],
			population             = populations[m],
			y                      = ys[m]
		)

		initepidemic!(newly_infecteds[m], cumulative_infecteds[m], θ, regionaldata)
		runepidemic!(newly_infecteds[m], cumulative_infecteds[m], θ, regionaldata)
	end
end

function initepidemic!(newly_infected, cumulative_infected, effective_Rt, θ, regionaldata)
	@unpack num_impute = θ
	@unpack y, Rt, population = regionaldata
	newly_infected[1]      = y
	cumulative_infected[1] = zero( eltype(cumulative_infected) )
	effective_Rt[1]        = Rt[1]
	for t in 2:num_impute
		newly_infected[t] = y
		cumulative_infected[t] = cumulative_infected[t-1] + y
		St = max(population - cumulative_infected[t], 0) / population
		effective_Rt[t] = St * Rt[t]
	end
end

function initepidemic!(newly_infected, cumulative_infected, θ, regionaldata)
	@unpack num_impute = θ
	@unpack y, population = regionaldata
	newly_infected[1]      = y
	cumulative_infected[1] = zero( eltype(cumulative_infected) )
	for t in 2:num_impute
		newly_infected[t] = y
		cumulative_infected[t] = cumulative_infected[t-1] + y
		# St = max(population - cumulative_infected[t], 0) / population
	end
end

function runepidemic!(newly_infected, cumulative_infected, effective_Rts, θ, regionaldata)
	@unpack num_impute, serial_interval, num_si = θ
	@unpack population, Rt = regionaldata
	num_time_step = length(newly_infected)

	for t = (num_impute + 1):num_time_step
		# Update cumulatively infected
		C = cumulative_infected[t] = cumulative_infected[t-1] + newly_infected[t - 1]
		susceptible = max(population - C, 0)

		# Adjusts for portion of population that is susceptible
		#St = max(population - cumulative_infected[t], 0) / population
		# effective number of infectious individuals
		effectively_infectious = sum(newly_infected[τ] * serial_interval[t - τ] for τ = (t - 1):-1:max(t-num_si,1))

		unadjusted_new_infections = effectively_infectious * Rt[t]

		# number of new infections (unobserved)
		# newly_infected[t] = susceptible * (1 - exp(- unadjusted_new_infections / population))
		newly_infected[t] = susceptible * (1 - exp(- unadjusted_new_infections / population))
		# newly_infected[t] = effective_Rt[t] * effectively_infectious

		effective_Rt[t] = susceptible / population * Rt[t]
	end
end

function runepidemic!(newly_infected, cumulative_infected, θ, regionaldata)
	@unpack num_impute, serial_interval, num_si = θ
	@unpack population, Rt = regionaldata
	num_time_step = length(newly_infected)

	for t = (num_impute + 1):num_time_step
		# Update cumulatively infected
		C = cumulative_infected[t] = cumulative_infected[t-1] + newly_infected[t - 1]
		susceptible = max(population - C, 0)

		# Adjusts for portion of population that is susceptible
		effectively_infectious = sum(newly_infected[τ] * serial_interval[t - τ] for τ = (t - 1):-1:max(t-num_si,1))

		unadjusted_new_infections = effectively_infectious * Rt[t]

		# number of new infections (unobserved)
		# newly_infected[t] = susceptible * (1 - exp(- unadjusted_new_infections / population))
		newly_infected[t] = susceptible * (1 - exp(- unadjusted_new_infections / population))

	end
	return nothing
end

# ============================================================================
# unbound infection model
function infections!(newly_infecteds, θ, ys, Rts)
	@unpack num_regions = θ
	for m in 1:num_regions
		regionaldata = (
			Rt                     = Rts[m],
			y                      = ys[m]
		)

		initepidemic!(newly_infecteds[m], θ, regionaldata)
		runepidemic!(newly_infecteds[m], θ, regionaldata)
	end
end

function initepidemic!(newly_infected, θ, regionaldata)
	@unpack num_impute = θ
	@unpack y, Rt = regionaldata
	newly_infected[1] = y
	for t in 2:num_impute
		newly_infected[t] = y
	end
end

function runepidemic!(newly_infected, θ, regionaldata)
	@unpack num_impute, serial_interval, num_si = θ
	@unpack Rt = regionaldata
	num_time_step = length(newly_infected)

	for t = (num_impute + 1):num_time_step
		effectively_infectious = sum(newly_infected[τ] * serial_interval[t - τ] for τ = (t - 1):-1:max(t-num_si,1))
		newly_infected[t] = effectively_infectious * Rt[t]
	end
	return nothing
end

# ============================================================================
# observations

function infection2observation(obsmodel::SimpleObsModel)
	@unpack ps, μ = obsmodel
	@unpack delay_dispersion, delay_length, delay_dist = ps
	i2o = pdf.( Ref(delay_dist(μ, delay_dispersion)), 1:delay_length )
	i2o /= sum(i2o)
	return i2o
end

function infection2observation(obsmodel::SimpleObsModel2)
	@unpack ps, μ, σ = obsmodel
	@unpack  delay_length, delay_dist = ps
	i2o = pdf.( Ref(delay_dist(μ, σ)), 1:delay_length )
	i2o /= sum(i2o)
	return i2o
end

function expected!(obsmodel::Union{SimpleObsModel, SimpleObsModel2}, newly_infecteds)
	@unpack ps, α, expecteds = obsmodel
	@unpack num_regions, delay_length = ps

	i2o = infection2observation(obsmodel)

	for m in 1:num_regions

		expected       = expecteds[m]
		newly_infected = newly_infecteds[m]
		num_time_step  = length(expected)

		expected[1] = 1e-15*newly_infected[1]
		for t = 2:num_time_step
			expected[t] = α * sum(newly_infected[τ] * i2o[t - τ] for τ = (t - 1):-1:max(t-delay_length,1))
			# expected[t] = α * sum(newly_infected[τ] * i2o[t - τ] for τ = (t - 1):-1:max(t-delay_length,1))
		end
	end
	return nothing
end

function expected!(obsmodel::WeekdayHolidayObsModel, newly_infecteds)
	@unpack ps, α, expecteds, holidayeffect, weekdayeffect = obsmodel
	@unpack num_regions, delay_length, holidays, weekdays = ps

	i2o = infection2observation(obsmodel)

	for m in 1:num_regions

		expected       = expecteds[m]
		newly_infected = newly_infecteds[m]
		num_time_step  = length(expected)
		holiday        = holidays[m]
		weekday        = weekdays[m]

		expected[1] = 1e-15*newly_infected[1]
		for t = 2:num_time_step
			weekday_holiday_effect = (1 - holidayeffect * holiday[t]) * weekdayeffect[weekday[t]]
			expected[t] = α * weekday_holiday_effect * sum(newly_infected[τ] * i2o[t - τ] for τ = (t - 1):-1:max(t-delay_length,1))
		end
	end
	return nothing
end

function _fastinference(ys, μs, ϕ)
	dist = arraydist(NegativeBinomial2.(μs, Ref(ϕ)))
	return logpdf(dist, ys)
end

function _slowinference(ys, μs, ϕ)
	ℓ = zero(typeof(ϕ))
	for (μ,y) in zip(μs,ys)
		dist = NegativeBinomial2(μ, ϕ)
		ℓ   += logpdf(dist, y)
	end
	return ℓ
end

function _prediction(ys, μs, ϕ)
	ℓ = zero(typeof(ϕ))
	for μ in μs
		dist = NegativeBinomial2(μ, ϕ)
		y    = rand(dist)
		ℓ   += logpdf(dist, y)
	end
	return ℓ
end

function Turing.logpdf(obsmodel::CovidSurvey.ObservationsModel, observed)
	@unpack ps, ϕ, expecteds, vectorize = obsmodel
	@unpack populations, starts, stops, num_regions = ps

	# _logpdf = if vectorize
	# 	_fastinference
	# else
	# 	if eltype(ϕ) <: Real
	# 		_prediction
	# 	else
	# 		_slowinference
	# 	end
	# end
	_logpdf = _fastinference

	T = typeof(ϕ)
	ℓ = zero(T)
	for m in 1:num_regions
		ts    = starts[m]:stops[m]
		μs = expecteds[m][ts]
		ys = observed[m][ts]

		#_out_of_bounds(μs, populations[m]) && (@warn "out of bounds"; return T(Inf))
		ℓ += _logpdf(ys, μs, ϕ)
	end
	return ℓ
end

_out_of_bounds(x, upper_lim) = !all(@. 0. < x < upper_lim )

# ============================================================================
# pointwise_loglikelihoods
function Turing.pointwise_loglikelihoods(obsmodel::CovidSurvey.ObservationsModel, observed)
	@unpack ps, ϕ, expecteds, vectorize = obsmodel
	@unpack populations, starts, stops, num_regions = ps

	logliks = Vector{Float64}[]
	for m in 1:num_regions
		ts    = starts[m]:stops[m]
		μs = expecteds[m][ts]
		ys = observed[m][ts]

		loglik = [ logpdf(NegativeBinomial2(μ, ϕ), y) for (μ,y) in zip(μs,ys) ]
		push!(logliks, loglik)
	end
	return logliks
end

# ============================================================================
# predictions
function prediction(obsmodel::CovidSurvey.ObservationsModel)
	@unpack ps, ϕ, expecteds = obsmodel
	@unpack starts, stops, num_regions = ps

	ys = Vector{Float64}[]
	for m in 1:num_regions
		μs = expecteds[m]
		μs[μs .< 1e-2] .= 1e-2
		y = [ rand(NegativeBinomial2(μ, ϕ)) for μ in μs ]
		push!(ys, y)
	end
	return ys
end

# ============================================================================
# weekday effect

function weekdayeffect!(weekdayeffect, weekdayeffect_simplex)
	num_days = 7
	for i in eachindex(weekdayeffect)
		@inbounds weekdayeffect[i] = weekdayeffect_simplex[i] * num_days
	end
end

# ============================================================================
# observe sero positive

struct SimpleSeroObsModel{T,E}# <: CovidSurvey.ObservationsModel
    ps::T
	expecteds::E
end

_unrealistic(ŷ,y) = any( ŷ .> 20*y )
# _unrealistic(y,ŷ) = any( ŷ .> 1.5*y )
function Turing.logpdf(obsmodel::SimpleSeroObsModel, observeds)
	@unpack ps, expecteds = obsmodel
	@unpack indexs, delay, populations, std, num_regions = ps

	T = eltype(expecteds[1])
	ℓ = zero(T)
	for m in 1:num_regions
		index      = indexs[m]
		expected   = expecteds[m]
		population = populations[m]
		observed   = observeds[m]

		μ    = expected[index .- delay] / population
		dist = MvNormal(μ, std) # select only early observations
		#_unrealistic(μ, observed) && (return T(Inf))
		ℓ += logpdf(dist, observed)
	end
	return ℓ
end


# Needleman-Wunsch Algorithm
# ==========================
#
# The Needleman-Wunsch algorithm.
#
# This file is a part of BioJulia.
# License is MIT: https://github.com/BioJulia/BioAlignments.jl/blob/master/LICENSE.md

mutable struct NeedlemanWunsch{T<:Union{Signed,AbstractFloat}}
    trace::Matrix{Trace}
    H::Vector{T}
    E::Vector{T}

    function NeedlemanWunsch{T}(m::Integer, n::Integer) where T
        trace = Matrix{Trace}(undef, (m + 1, n + 1))
        fill!(trace, 0xff)
        H = Vector{T}(undef, m + 1)
        E = Vector{T}(undef, m)
        return new{T}(trace, H, E)
    end

    function NeedlemanWunsch{T}() where T
        return NeedlemanWunsch{T}(0, 0)
    end
end

function ensureroom!(nw::NeedlemanWunsch, m, n)
    if size(nw.trace, 1) < m + 1 || size(nw.trace, 2) < n + 1
        # TODO: resize the trace matrix if possible
        nw.trace = Matrix{Trace}(undef, (m + 1, n + 1))
        nw.H = resize!(nw.H, m + 1)
        nw.E = resize!(nw.E, m)
    end
    return nw
end

# Generate a code block that updates the current score and trace.
macro update()
    esc(quote
        e = E[i]
        g = h_diag + submat[a[i],b_j]
        h = max(e, f, g)
        h_diag = H[i+1]
        H[i+1] = h
        t = trace[i+1,j+1] | ft
        e == h && (t |= TRACE_DELETE)
        f == h && (t |= TRACE_INSERT)
        g == h && (t |= TRACE_MATCH)
        trace[i+1,j+1] = t
    end)
end

macro update_nextcol(gap_init, gap_extend)
    esc(quote
        e′ = e + $gap_extend
        e = max(e′, h + $gap_init)
        E[i] = e
        et = TRACE_NONE
        e == e′ && (et |= TRACE_EXTDEL)
        trace[i+1,j+2] = et
    end)
end

macro update_nextrow(gap_init, gap_extend)
    esc(quote
        f′ = f + $gap_extend
        f = max(f′, h + $gap_init)
        ft = TRACE_NONE
        f == f′ && (ft |= TRACE_EXTINS)
    end)
end

#    start   middle  end gaps of a
#    |       |       |
# a: ----ACCA---ATGTG---
# b: AAATACGATTGATGAGGGT

function run!(
        nw::NeedlemanWunsch{T},
        a, b,
        submat::AbstractSubstitutionMatrix{T},
        # a
        start_gap_open_a::T,
        start_gap_extend_a::T,
        middle_gap_open_a::T,
        middle_gap_extend_a::T,
        end_gap_open_a::T,
        end_gap_extend_a::T,
        # b
        start_gap_open_b::T,
        start_gap_extend_b::T,
        middle_gap_open_b::T,
        middle_gap_extend_b::T,
        end_gap_open_b::T,
        end_gap_extend_b::T) where T

    m = length(a)
    n = length(b)
    ensureroom!(nw, m, n)
    H = nw.H
    E = nw.E
    trace = nw.trace

    start_gap_init_a = start_gap_open_a + start_gap_extend_a
    middle_gap_init_a = middle_gap_open_a + middle_gap_extend_a
    end_gap_init_a = end_gap_open_a + end_gap_extend_a

    start_gap_init_b = start_gap_open_b + start_gap_extend_b
    middle_gap_init_b = middle_gap_open_b + middle_gap_extend_b
    end_gap_init_b = end_gap_open_b + end_gap_extend_b

    # ↓: insertion
    # →: deletion
    # ↘: match/mismatch
    H[1] = T(0)
    trace[1,1] = TRACE_NONE
    if m == 0 && n == 0
        return H[1]
    elseif m == 0
        for j in 1:n
            trace[1,j+1] = TRACE_DELETE
        end
        return start_gap_open_a + start_gap_extend_a * n
    elseif n == 0
        for i in 1:m
            trace[i+1,1] = TRACE_INSERT
        end
        return start_gap_open_b + start_gap_extend_b * m
    else
        # initialize the first column of H, E, and trace
        @inbounds for i in 1:m
            H[i+1] = start_gap_open_b + start_gap_extend_b * i
            E[i] = H[i+1] + (i == m ? end_gap_init_a : middle_gap_init_a)
            trace[i+1,1] = TRACE_INSERT
            trace[i+1,2] = TRACE_NONE
        end

        # run dynamic programming column by column (except the last column)
        @inbounds for j in 1:n-1
            b_j = b[j]
            h_diag = H[1]
            H[1] = start_gap_open_a + start_gap_extend_a * j
            f = H[1] + middle_gap_init_b
            ft = TRACE_NONE
            trace[1,j+1] = TRACE_DELETE
            # inner loop along the sequence a
            for i in 1:m-1
                @update
                @update_nextcol middle_gap_init_a middle_gap_extend_a
                @update_nextrow middle_gap_init_b middle_gap_extend_b
            end
            # fill the final row
            let i = m
                @update
                @update_nextcol end_gap_init_a end_gap_extend_a
            end
        end

        # fill the last column
        let j = n
            b_j = b[j]
            h_diag = H[1]
            H[1] = start_gap_open_a + start_gap_extend_a * j
            f = H[1] + end_gap_init_b
            ft = TRACE_NONE
            trace[1,j+1] = TRACE_DELETE
            # inner loop along the sequence a
            @inbounds for i in 1:m-1
                @update
                @update_nextrow end_gap_init_b end_gap_extend_b
            end
            # finally, fill the last element of the DP matrices
            let i = m
                @update
            end
        end

        return H[m+1]
    end
end

function run!(nw::NeedlemanWunsch, a, b, submat, gap_open, gap_extend)
    return run!(nw, a, b, submat,
        # a
        gap_open,
        gap_extend,
        gap_open,
        gap_extend,
        gap_open,
        gap_extend,
        # b
        gap_open,
        gap_extend,
        gap_open,
        gap_extend,
        gap_open,
        gap_extend
    )
end

function traceback(nw::NeedlemanWunsch, a, b, endpos)
    anchors = Vector{AlignmentAnchor}()
    i, j = endpos
    was_extins = false
    was_extdel = false
    @start_traceback
    while i ≥ 1 || j ≥ 1
        t = nw.trace[i+1,j+1]
        @assert !(was_extins && was_extdel)
        if was_extins
            was_extins = t & TRACE_EXTINS > 0
            @anchor OP_INSERT
        elseif was_extdel
            was_extdel = t & TRACE_EXTDEL > 0
            @anchor OP_DELETE
        elseif t & TRACE_INSERT > 0
            was_extins = t & TRACE_EXTINS > 0
            @anchor OP_INSERT
        elseif t & TRACE_DELETE > 0
            was_extdel = t & TRACE_EXTDEL > 0
            @anchor OP_DELETE
        elseif t & TRACE_MATCH > 0
            if a[i] == b[j]
                @anchor OP_SEQ_MATCH
            else
                @anchor OP_SEQ_MISMATCH
            end
        else
            error("failed to trace back")
        end
    end
    @finish_traceback
    return AlignedSequence(a, anchors)
end


# Author: AJ Tyler
# Date: 24/01/22
module Comparison # This module contains functions which compare configurations using their CNA profiles.
include("io.jl") # This file contains the IO module.
using .IO # This module contains functions that read in and write out files.
using Statistics # This module allows easy mean and standard deviation.
export compareConfigs,groupConfigs,configStats # These functions ared used to compare configurations.

"""
Finds the most similar configuration to every configuration over all rCut values based on their CNA profile.

Inputs:
rCutRange: (Vector{Float64}) A range of cut-off radii to test, in units of equilibriumBondLength.
energies: (Vector{Float64}) The energy of each configuration in each the file.
totalProfiles: Array{Dict{String,Int}} 2D array containg the totalCNA profiles of each configuration, for each rCut value.
atomicProfiles: Array{Dict{String,Int}} 3D array containing the CNA profile of each configuration, for each rCut value, for each atom.
configDir: (String) The path to (or name of if in same location as this program) the directory containing the .xyz files.
filename: (String) The prefix of the file.
similarityMeasure: (String) The method of comparing configurations, either "atomic" or "total" CNA.
M: (Int64) The number of rCut values to test.
L: (Int64) The number of configurations in the file.
N: (Int64) The number of atoms in each configuration.

Outputs:
similarConfigs: Vector{Dict{Int64,Vector{Float64}}} A vector of dictionaries which map the most similar configurations to the rCut values
													for which it is maximised, for each configuraiton.
maxSims: Vector{Float64} A Vector of the maximum similarity score for each configuration when compared to every other configuration
						 over all rCut values.
"""
function compareConfigs(rCutRange,energies,totalProfiles,atomicProfiles,configDir,filename,similarityMeasure,M,L,N)
	# Initialise vector to contain most similar configurations for each configuration
	similarConfigs = Vector{Dict{Int64,Vector{Float64}}}(undef,L)
	maxSims = zeros(Float64,1,L) # Initialise the maximum similarity score for each configuration
	for i in 1:M # For all rCut values
		writeProfiles(rCutRange,configDir,filename,energies,totalProfiles,i,L) # Write out the total CNA profiles
		if L == 1
			continue
		end
		atomicProfilesDict = convertAtomicProfiles(atomicProfiles,i,L) # Convert atomic CNA profile format
		# For each pair of configurations
		for configA in 1:L-1
			for configB in configA+1:L
				# Compute the similarity for the cluster pair
				similarity = similarityScore(configA, configB, totalProfiles,atomicProfilesDict,similarityMeasure,N,i)
				# Update the most similar configurations to configA
				if (similarity > maxSims[configA]) # If new similarity score is greater than previous highest
					similarConfigs[configA] = Dict{Int64,Vector{Float64}}(configB=>Vector{Float64}([rCutRange[i]])) # Create new vector, containing configB
					maxSims[configA] = similarity # Update maximum similarity score
				elseif (similarity == maxSims[configA]) # If the new similarity is the same as the previous highest
					if (haskey(similarConfigs[configA],configB)) # If the same configuration pair is already most similar
						# Add the current rCut value to the dictionary
						push!(similarConfigs[configA], configB => push!(get!(similarConfigs[configA],configB,0),rCutRange[i]))
					else
						push!(similarConfigs[configA], configB => Vector{Float64}([rCutRange[i]])) # Add the configuration to the dictionary
					end
				end

				# Same for configB
				if (similarity > maxSims[configB]) # If new similarity score is greater than previous highest
					similarConfigs[configB] = Dict{Int64,Vector{Float64}}(configA=>Vector{Float64}([rCutRange[i]])) # Create new vector, containing configB
					maxSims[configB] = similarity # Update maximum similarity score
				elseif (similarity == maxSims[configB]) # If the new similarity is the same as the previous highest
					if (haskey(similarConfigs[configB],configA))
						push!(similarConfigs[configB], configA => push!(get!(similarConfigs[configB],configA,0),rCutRange[i]))
					else
						push!(similarConfigs[configB], configA => Vector{Float64}([rCutRange[i]]))
					end
				end
			end
		end
	end
	return similarConfigs, maxSims # Return outputs
end

"""
Converts a 2D array of containing the CNA profile of each configuration, for each atom, to a vector of dictionaries,
mapping atom CNA profiles to their frequency, for each configuration.

Inputs:
atomicProfiles: Array{Dict{String,Int}} 3D array containing the CNA profile of each configuration, for each rCut value, for each atom.
i: (Int64) The index of the current rCut value.
L: (Int64) The number of configurations in the file.

Outputs:
atomicProfileDict: Vector{Dict{String,Int64}} ector of dictionaries, mapping atom CNA profiles to their frequency, for each configuration.
"""
function convertAtomicProfiles(atomicProfiles,i,L)
	# Convert atomicProfiles into format better for computing similarity
	atomicProfilesDict = [Dict{String,Int64}() for j in 1:L] # Initialise new format
	for j in 1:L # For all configurations
		atomicProfile = atomicProfiles[i,j,:] # Get atom profile
		for atomProfile in atomicProfile # For each atom
			# Increments atomic CNA profile frequency by 1
			push!(atomicProfilesDict[j], "$atomProfile" => get!(atomicProfilesDict[j],"$atomProfile",0)+1)
		end
	end

	return atomicProfilesDict # Return output
end

"""
This function returns the similarity score of two configurations based on their CNA profiles.

Inputs:
configA: (Int64) The configuration number of one of the configurations to compare.
configB: (Int64) The configuration number of the other of the configurations to compare.
totalProfiles: Array{Dict{String,Int}} 2D array containg the totalCNA profiles of each configuration, for each rCut value.
atomicProfilesDict: Vector{Dict{String,Int64}} ector of dictionaries, mapping atom CNA profiles to their frequency, for each configuration.
similarityMeasure: (String) The method of comparing configurations, either "atomic" or "total" CNA.
N: (Int64) The number of atoms in each configuration.
i: (Int64) The index of the current rCut value.

Outputs:
similarity: (Float64) The similarity score between the two configurations.
"""
function similarityScore(configA, configB, totalProfiles,atomicProfilesDict,similarityMeasure,N,i)
	similarity = 0.0 # Initialise similarity score
	if (similarityMeasure == "atomic") # If want to compare using atomic CNA profile
		# Compute similarity of atomic CNA profiles
		for key in keys(atomicProfilesDict[configA]) # For each atom CNA profile
			# Add the number of atoms with the same atom CNA profile
			similarity += min(get(atomicProfilesDict[configA],key,0),get(atomicProfilesDict[configB],key,0))
		end
		similarity /= N # Normalise the similarity score

	elseif (similarityMeasure == "total") # If want to compare using total CNA profile
		# Compute Jacard index of their total CNA profiles
		similarity = JacardIndex(totalProfiles[i,configA],totalProfiles[i,configB])
	else # If none of the valid similarity measures were used
		throw(DomainError(similarityMeasure, "Invalid similarity measure")) # Throw an error
	end

	return round(similarity; digits = 3) # Return rounded similarity score
end

"""
This function calculates the Jacard Index for two dictionaries,
with their keys being the unique elements of a multiset and the values are their frequencies

Inputs:
A: (Dictionary) Multiset A
B: (Dictionary) Multiset B

Outputs:
similarity: (Float64) The Jacard Index of multsets A & B
"""
function JacardIndex(A,B)
	intersection = 0.0 # Initialise the number of elements in the intersection of A & B
	union = 0.0 # Initialise the number of elements in the union of A & B
	for key in keys(A) # For all the unique elements of A
		intersection += min(get(A,key,0),get(B,key,0)) # Add the minimum number of occurances of the key across both A & B
		union += get(A,key,0) # Add all the elements of A
	end

	for key in keys(B) # For all the unique elements of B
		union += get(B,key,0) # Add all the elements of B
	end

	union -= intersection # union = A + B - intersection
	similarity = intersection/union # Comput the Jacard index
	return similarity # Return the Jacard Index
end

"""
This function groups configurations into distinguisable sets based on their similarity scores.

Inputs:
energies: (Vector{Float64}) The energy of each configuration in each the file.
sortingArray: Vector{Int64} A vector of configuration numbers, sorted by maximum similarity score.
maxSims: Vector{Float64} A Vector of the maximum similarity score for each configuration when compared to every other configuration
						 over all rCut values, in decreasing order.
similarConfigs: Vector{Dict{Int64,Vector{Float64}}} A vector of dictionaries which map the most similar configurations to the rCut values
						 							for which it is maximised, for each configuraiton, sorted by similarity.
similarityThreshold: (Float64) Similarity threshold above which two configurations are considered identical.
M: (Int64) The number of rCut values tested.

Keyword Arguments:
rCutThreshold: (Float64) Proportion of rCut values that have to meet the similarity threshold for two configs to be considered identical.

Outputs:
uniqueConfigs: Vector{Vector{Int64}} Vector of configuration numbers of the distinguisable sets of configurations.
uniqueEnergies: Vector{Vector{Float64}} Vector of configuration energies of the distinguisable sets of configurations.
uniqueSimilarities: Vector{Float64} Vector of similarity scores of each of the distinguiable sets.
"""
function groupConfigs(energies,sortingArray,maxSims,similarConfigs,similarityThreshold,M;rCutThreshold=0.9)
	# Initialise arrays
	uniqueConfigs = Vector{Vector{Int64}}()
	uniqueEnergies = Vector{Vector{Float64}}()
	uniqueSimilarities = Vector{Float64}()
	classified = Vector{Int64}() # Vector of configuration numbers that have already been grouped
	for (j,config) in enumerate(sortingArray) # For all configurations
		if (!(config in classified)) # If the configuration has not already been grouped
			same = Vector{Int64}([config])
			if (maxSims[j] > similarityThreshold) # If the maximum similarity of the configuration exceeds a threshold
				for key in keys(similarConfigs[j])
					if (length(get!(similarConfigs[j],key,0)) > rCutThreshold*M && !(key in classified))
						# Get vector of configuration numbers of configurations that are similar
						push!(same,key)
					end
				end
			end
			push!(uniqueConfigs,same) # Add the set of configurations to the list
			push!(uniqueEnergies,energies[same]) # Add the energies of the configurations to the list
			push!(uniqueSimilarities,maxSims[j]) # Add the similarity score to the list
			classified = vcat(classified,same) # Add the configurations to the list of already classified
		end
	end

	return uniqueConfigs, uniqueEnergies, uniqueSimilarities # Return outputs
end

"""
This function returns some basic statistics of the energies of each of the distinguisbale configuration sets.

Inputs:
uniqueEnergies: (Vector{Vector{Float64}}) Vector of configuration energies of the distinguisable sets of configurations.
numUnique (Int64) Number of distinguiable sets of configurations.

Outputs:
uniqueEnergyStats: Array{Float64} 2D array of energy statistics. First row is the mean, second row is standard deviation.
"""
function configStats(uniqueEnergies,numUnique)
	uniqueEnergyStats = Array{Float64}(undef,2,numUnique) # Initialse array
	for i in 1:numUnique # For each grouped set of configurations
		uniqueEnergyStats[1,i] = Statistics.mean(uniqueEnergies[i]) # Compute mean
		uniqueEnergyStats[2,i] = Statistics.std(uniqueEnergies[i]) # Compute standard deviation
	end

	return uniqueEnergyStats # Return statistics
end

end # End of Module


# Quelques fonctions utilitaires pour les manipulations de solvers PL ou MIP

# Classe abstraite non utilisé
# abstract type AbstractLpSolver end

"""
Création d'un modèle de solver externe (:cplex, :glpk, ...)

Args:
    solver : symbole :cplex, :glpk, :clp, :cbc

        - :cbc : permet de faire du MIP
        - :clp : resteint aux problème LP mais plus efficace de cbc pour les relaxations.

    mode : (:lp ou :mip) utile pour choisir entre :clp ou :cbc
    log_level (expérimental) : entier 0 pour silence, sinon 1, 2, ...

Autre solver à supporter plus tard Gurobi (concurrent de cplex)

Attention:
    DEPENDANCE EXTERNE : module Args (pour Args.get("external_mip_solver"))
"""
function new_model(;solver = :auto,
                     mode = :lp,
                     log_level = 0, # EXPLOITATION PARTIELLE
                     )
    if !(mode in [:mip, :lp])
        error("ERROR: unknown mode $(mode). Should be :cplex, :glpk, ...")
    end
    if solver == :auto
        solver = Args.get("external_mip_solver")
    end
    if solver in [:clp, :cbc]
        # solver = (mode == :mip ? :cbc : :clp )
        if mode==:mip
            solver = :cbc
        else
            solver = :clp
        end
    end
    if ! (solver in [:cplex, :clp, :cbc, :glpk])
        error("ERROR: unknown solver $(solver). Should be :cplex, :glpk, ...")
    end

    if solver == :glpk
        # voir https://github.com/JuliaOpt/GLPK.jl
        # model = Model(GLPK.Optimizer)  # PAS ENCORE SUPPORTE (05/12/2019)
        model = Model(with_optimizer(GLPK.Optimizer))
    elseif solver == :cplex
        model = Model(with_optimizer(CPLEX.Optimizer))
        # model = Model(CPLEX.Optimizer) # PAS ENCORE SUPPORTE (05/12/2019)
    elseif solver == :clp
        # voir https://github.com/JuliaOpt/Clp.jl
        # model = Model(Clp.Optimizer) # PAS ENCORE SUPPORTE (05/12/2019)
        model = Model(with_optimizer(Clp.Optimizer))
    elseif solver == :cbc
        # Voir https://github.com/JuliaOpt/Cbc.jl pour les options
        # model = Model(Cbc.Optimizer) # PAS ENCORE SUPPORTE (05/12/2019)
        model = Model(with_optimizer(Cbc.Optimizer))
    else
        error("LpTimingSolver: unknown external_mip_solver inconnu : $solver")
    end
    # MOI.set(model, MOI.Silent(), true) # Variante de bas niveau
    if log_level <= 3
        set_silent(model) # unset_silent(model) pour réautoriser affichage
    end
    return model
end