import Base: print, length, push!

type FIRDecim{Th,Tx}
    q::Ptr{Th}
    tx::Type{Tx}
    decimation::Int
end

Base.show{Th,Tx}(io::IO, obj::FIRDecim{Th,Tx}) = print(io::IO, "FIRDecim{$Th,$Tx}")

function Base.showall(io::IO, obj::FIRDecim)
    println(io::IO)
    print(obj)
end

for (sigstr, Ty, Th, Tx) in (rrrf, crcf, cccf)

    # /* create decimator from external coefficients              */  \
    # /*  _M      : decimation factor                             */  \
    # /*  _h      : filter coefficients [size: _h_len x 1]        */  \
    # /*  _h_len  : filter coefficients length                    */  \
    # FIRDECIM() FIRDECIM(_create)(unsigned int _M,                   \
    #                              TC *         _h,                   \
    #                              unsigned int _h_len);              \

    @eval begin
        function FIRDecim(::Type{$Tx}, decimation::Integer, h::Vector{$Th})
            q   = ccall(($"firdecim_$(sigstr)_create", liquiddsp), Ptr{$Th}, (Cuint, Ptr{$Th}, Cuint), decimation, h, length(h))
            obj = FIRDecim(q, $Tx, decimation)
            finalizer(obj, destroy)
            return obj
        end
    end

    # /* create decimator from prototype
    # /*  _M      : decimation factor
    # /*  _m      : filter delay (symbols)
    # /*  _As     : stop-band attenuation [dB]
    # FIRDECIM() FIRDECIM(_create_prototype)(unsigned int _M,
    #                                        unsigned int _m,
    #                                        float        _As);

    # # TODO: this function doesn't appear to be working on the liquid-dsp side of th house
    # @eval begin
    #     function FIRDecim(::Type{$Tx}, decimation::Integer, attenuation::Real, delay::Integer = 11)
    #         delay > 0 || throw(ArgumentError("Delay, (hLen-1)/(2*decimation), must be greater than 0"))
    #         q   = ccall(($"firdecim_$(sigstr)_create_prototype", liquiddsp), Ptr{$Th}, (Cuint, Cuint, Float32), decimation, delay, attenuation)
    #         obj = FIRDecim(q, $Tx, decimation)
    #         finalizer(obj, destroy)
    #         return obj
    #     end
    # end

    # /* print decimator object propreties to stdout              */  \
    # void FIRDECIM(_print)(FIRDECIM() _q);                           \

    @eval begin
        function print( obj::FIRDecim{$Th,$Tx} )
            obj.q == C_NULL && error("`obj` is a NULL pointer")
            ccall(($"firdecim_$(sigstr)_print", liquiddsp), Void, (Ptr{$Th},), obj.q)
        end
    end

    # /* reset decimator object internal state                    */  \
    # void FIRDECIM(_clear)(FIRDECIM() _q);                           \

    @eval begin
        function reset!( obj::FIRDecim{$Th,$Tx} )
            obj.q == C_NULL && error("`obj` is a NULL pointer")
            ccall(($"firdecim_$(sigstr)_clear", liquiddsp), Void, (Ptr{$Th},), obj.q)
        end
    end

    # /* destroy decimator object                                 */  \
    # void FIRDECIM(_destroy)(FIRDECIM() _q);                         \

    @eval begin
        function destroy( obj::FIRDecim{$Th,$Tx} )
            obj.q == C_NULL && return
            ccall(($"firdecim_$(sigstr)_destroy", liquiddsp), Void, (Ptr{$Th},), obj.q)
            obj.q = C_NULL
        end
    end

    # /* execute decimator on block of _n*_M input samples        */  \
    # /*  _q      : decimator object                              */  \
    # /*  _x      : input array [size: _n*_M x 1]                 */  \
    # /*  _n      : number of _output_ samples                    */  \
    # /*  _y      : output array [_size: _n x 1]                  */  \
    # void FIRDECIM(_execute_block)(FIRDECIM()   _q,                  \
    #                               TI *         _x,                  \
    #                               unsigned int _n,                  \
    #                               TO *         _y);                 \

    @eval begin
        function execute( obj::FIRDecim{$Th,$Tx}, x::Vector{$Tx} )
            obj.q == C_NULL && error("`obj` is a NULL pointer")
            xLen = length(x)
            xLen % obj.decimation == 0 || error("input must be an integer multiple of the decimation factor")
            yLen = div(xLen, obj.decimation)
            y = Array($Ty, yLen)
            ccall(($"firdecim_$(sigstr)_execute_block", liquiddsp), Void, (Ptr{$Th}, Ptr{$Tx}, Cuint, Ptr{$Ty}), obj.q, x, yLen, y)
            return y
        end
    end
end

# Handle conversion of taps tp Float32
FIRDecim(Tx::Type, decimation::Integer, h::AbstractVector) = FIRDecim(Tx, decimation, limitprecision(h))

# If passed only taps, assume x will be the same type
function FIRDecim(decimation::Integer, h::AbstractVector)
    h = limitprecision(h)
    FIRDecim(eltype(h), decimation, h)
end

function gettaps{Th,Tx}(obj::FIRDecim{Th,Tx})
    hLenPtr = convert(Ptr{Cuint}, obj.q+sizeof(obj.q))
    hLen    = unsafe_load(hLenPtr)
    hPtr    = convert(Ptr{Ptr{Th}}, obj.q)
    h       = pointer_to_array(unsafe_load(hPtr), hLen)
end

#################################################################
# _ _  _ ___  _    ____ _  _ ____ _  _ ___ ____ ___ _ ____ _  _ #
# | |\/| |__] |    |___ |\/| |___ |\ |  |  |__|  |  | |  | |\ | #
# | |  | |    |___ |___ |  | |___ | \|  |  |  |  |  | |__| | \| #
#                                                               #
#             ___  ____ _  _ ___  _ _  _ ____                   #
#             |__] |___ |\ | |  \ | |\ | | __                   #
#             |    |___ | \| |__/ | | \| |__]                   #
#################################################################

# /* create square-root Nyquist decimator
# /*  _type   : filter type (e.g. LIQUID_FIRFILT_RRC)
# /*  _M      : samples/symbol (decimation factor)
# /*  _m      : filter delay (symbols)
# /*  _beta   : rolloff factor (0 < beta <= 1)
# /*  _dt     : fractional sample delay
# FIRDECIM() FIRDECIM(_create_rnyquist)(int          _type,
#                                       unsigned int _M,
#                                       unsigned int _m,
#                                       float        _beta,
#                                       float        _dt);
#
#
#
# /* execute decimator on _M input samples
# /*  _q      : decimator object
# /*  _x      : input samples [size: _M x 1]
# /*  _y      : output sample pointer
# void FIRDECIM(_execute)(FIRDECIM() _q,
#                         TI *       _x,
#                         TO *       _y);
#


"""
    compare(s1, s2, dist)

return a similarity score between 0 and 1 for the strings `s1` and
`s2` based on the distance `dist`.

### Examples
```julia-repl
julia> compare("martha", "marhta", Levenshtein())
0.6666666666666667
```
"""
compare(s1, s2, dist::StringDistance; min_score = 0.0) = 1 - normalize(dist)(s1, s2, 1 - min_score)

"""
    findnearest(s, itr, dist::StringDistance; min_score = 0.0) -> (x, index)

`findnearest` returns the value and index of the element of `itr` that has the
highest similarity score with `s` according to the distance `dist`.
It returns `(nothing, nothing)` if none of the elements has a similarity score
higher or equal to `min_score` (default to 0.0).

It is particularly optimized for [`Levenshtein`](@ref) and [`DamerauLevenshtein`](@ref) distances
(as well as their modifications via [`Partial`](@ref), [`TokenSort`](@ref), [`TokenSet`](@ref), or [`TokenMax`](@ref)).

### Examples
```julia-repl
julia> using StringDistances
julia> s = "Newark"
julia> iter = ["New York", "Princeton", "San Francisco"]
julia> findnearest(s, iter, Levenshtein())
("NewYork", 1)
julia> findnearest(s, iter, Levenshtein(); min_score = 0.9)
(nothing, nothing)
```
"""
function findnearest(s, itr, dist::StringDistance; min_score = 0.0)
    min_score_atomic = Threads.Atomic{Float64}(min_score)
    scores = [0.0 for _ in 1:Threads.nthreads()]
    is = [0 for _ in 1:Threads.nthreads()]
    s = _helper(s, dist)
    # need collect since @threads requires a length method
    Threads.@threads for i in collect(eachindex(itr))
        score = compare(s, _helper(itr[i], dist), dist; min_score = min_score_atomic[])
        score_old = Threads.atomic_max!(min_score_atomic, score)
        if score >= score_old
            scores[Threads.threadid()] = score
            is[Threads.threadid()] = i
        end
    end
    imax = is[argmax(scores)]
    imax == 0 ? (nothing, nothing) : (itr[imax], imax)
end

function _helper(s, dist::QGramDistance)
    s !== missing ? QGramSortedVector(s, dist.q) : s
end
_helper(s, dist::StringDistance) = s

function Base.findmax(s, itr, dist::StringDistance; min_score = 0.0)
    @warn "findmax(s, itr, dist; min_score) is deprecated. Use findnearest(s, itr, dist; min_score)"
    findnearest(s, itr, dist; min_score = min_score)
end
"""
    findall(s, itr , dist::StringDistance; min_score = 0.8)

`findall` returns the vector of indices for elements of `itr` that have a
similarity score higher or equal than `min_score` according to the distance `dist`.
If there are no such elements, return an empty array.

It is particularly optimized for [`Levenshtein`](@ref) and [`DamerauLevenshtein`](@ref) distances
(as well as their modifications via `Partial`, `TokenSort`, `TokenSet`, or `TokenMax`).

### Examples
```julia-repl
julia> using StringDistances
julia> s = "Newark"
julia> iter = ["Newwark", "Princeton", "San Francisco"]
julia> findall(s, iter, Levenshtein())
1-element Array{Int64,1}:
 1
julia> findall(s, iter, Levenshtein(); min_score = 0.9)
0-element Array{Int64,1}
```
"""
function Base.findall(s, itr, dist::StringDistance; min_score = 0.8)
    out = [Int[] for _ in 1:Threads.nthreads()]
    s = _helper(s, dist)
    # need collect since @threads requires a length method
    Threads.@threads for i in collect(eachindex(itr))
        score = compare(s, _helper(itr[i], dist), dist; min_score = min_score)
        if score >= min_score
            push!(out[Threads.threadid()], i)
        end
    end
    vcat(out...)
end


import AxisArrays: axisdim, axisnames, axisname, axisvalues

# All ArrayTypes must define `axisnames` function
const ArrayTypes{T,M} = Union{AxisArray{T,M}} # JuMP.JuMPArray{T,M}

struct WrappedArray{T,N,M,D,Ax} <: AbstractArray{T,N}
    data::D
    axes::Ax
    # Duck-typing, we assume axisnames is defined for data
    function WrappedArray(data::T, axes::NTuple{N,Axis}) where {T<:Union{Number,Symbol}, N}
        new{T,N,0,T,typeof(axes)}(data, axes)
    end
    function WrappedArray(data::ArrayTypes{T,M}, axes::NTuple{N,Axis}) where {T,N,M}
        @assert(issubset(axisnames(data), axisnames(axes...)),
                "WrappedArray is missing axes: $(setdiff(axisnames(data), axisnames(axes...)))")
        new{T,N,M,typeof(data),typeof(axes)}(data, axes)
    end
end
WrappedArray(data, axes::Axis...) = WrappedArray(data, axes)
# Not all axes are Axis, since we insist, you have to provide a function that
# convert one for your iterator type by adding methods to this recursive
# _to_axis
WrappedArray(data, axes...) = WrappedArray(data, _to_axis(axes...)...)
_to_axis(ax::Axis, axes...) = (ax, _to_axis(axes...)...)
_to_axis() = ()

@inline Base.size(A::WrappedArray) = length.(A.axes)
@inline Base.size(A::WrappedArray, Ax::Axis) = length(A.axes[axisdim(A, Ax)])

# Slightly adapted axisdim, axisnames and axisvalues from AxisArrays

axisdim(A::WrappedArray, ax::Axis) = axisdim(A, typeof(ax))
@generated function axisdim(A::WrappedArray, ax::Type{Ax}) where Ax<:Axis
    dim = axisdim(A, Ax)
    :($dim)
end

# The actual computation is done in the type domain, which is a little tricky
# due to type invariance.
function axisdim(::Type{WrappedArray{T,N,M,D,Ax}}, ::Type{<:Axis{name,S} where S}) where {T,N,M,D,Ax,name}
    isa(name, Int) && return name <= N ? name : error("axis $name greater than array dimensionality $N")
    names = axisnames(Ax)
    idx = findfirst(isequal(name), names)
    isnothing(idx) && error("axis $name not found in array axes $names")
    idx
end

axisnames(::WrappedArray{T,N,M,D,Ax}) where {T,N,M,D,Ax}       = AxisArrays._axisnames(Ax)
axisnames(::Type{WrappedArray{T,N,M,D,Ax}}) where {T,N,M,D,Ax} = AxisArrays._axisnames(Ax)
axisvalues(A::WrappedArray) = axisvalues(A.axes...)

# Probably if there are colons or ranges on any of the virtual axes,
# we would instead want to give back a new WrappedArray, but this works well enough.
@generated function Base.getindex(A::WrappedArray{T,N,M,D,Ax}, idxs...) where {T,N,M,D,Ax}
    meta = Expr(:meta, :inline, :propagate_inbounds)
    if M == 0
        :(A.data)
    else
        names = axisnames(A)
        :($meta; A.data[$((Expr(:ref, :idxs, something(findfirst(isequal(n), names))) for n=axisnames(D))...)])
    end
end

# # RelIterators allow to use + Int, and - Int to get to adjacent indices, BUT
# # they are about a 100x slower, so should be used sparingly

# struct RelativeIndex{Ax}
#     ax::Ax
#     i::Int64
# end

# struct RelativeIterator{Ax}
#     ax::Ax
# end

# Base.start(it::RelativeIterator) = 1
# Base.done(it::RelativeIterator, state) = state == length(it.ax) + 1
# Base.next(it::RelativeIterator, state) = RelativeIndex(it.ax, state), state + 1
# Base.length(it::RelativeIterator) = length(it.ax)
# Base.:+(r::RelativeIndex, d::Int) = RelativeIndex(r.ax, r.i+d)
# Base.:-(r::RelativeIndex, d::Int) = RelativeIndex(r.ax, r.i-d)

# # We intercept the getindex calls of JuMPArray and AxisArray to resolve to the
# # actual index value
# for T = (:(JuMP.JuMPArray), :AxisArray)
#     @eval Base.getindex(a::$T, r::RelativeIndex) = a[r.ax[r.i]]
# end

# @inline unrelative(it) = it
# @inline unrelative(it::RelativeIterator) = it.ax

# _to_axis(it::RelativeIterator, axes...) = (unrelative(it), _to_axis(axes...)...)


import Base: ==, hash

# Define struct not yet exported in POMDPs.jl framework
struct AlphaVec
    alpha::Vector{Float64} # alpha vector
    action::Any # action associated wtih alpha vector
end

AlphaVec() = AlphaVec([0.0], 0)

# define alpha vector equality
Base.hash(a::AlphaVec, h::UInt) = hash(a.alpha, hash(a.action, h))
==(a::AlphaVec, b::AlphaVec) = (a.alpha,a.action) == (b.alpha, b.action)

import BeliefUpdaters: pdf
pdf(b::DiscreteBelief, s) = b.b[typeof(HorizonLength(b.pomdp)) == InfiniteHorizon ?
                                stateindex(b.pomdp, s) : stage_stateindex(b.pomdp, s)]

import BeliefUpdaters: initialize_belief
function initialize_belief(bu::DiscreteUpdater, dist::FiniteHorizonPOMDPs.InStageDistribution)
    state_list = ordered_stage_states(bu.pomdp, stage(dist))
    ns = length(state_list)
    b = zeros(ns)
    belief = DiscreteBelief(bu.pomdp, state_list, b)
    for s in support(dist)
        sidx = stage_stateindex(bu.pomdp, s)
        belief.b[sidx] = pdf(dist, s)
    end
    return belief
end

import BeliefUpdaters: update
function update(bu::DiscreteUpdater, b::DiscreteBelief, a, o)
    pomdp = bu.pomdp
    state_space = b.state_list
    bp = zeros(length(state_space))

    for (si, s) in enumerate(state_space)
        si = stage_stateindex(pomdp, s)

        if pdf(b, s) > 0.0
            td = transition(pomdp, s, a)
            for (sp, tp) in weighted_iterator(td)
                spi = stage_stateindex(pomdp, sp)
                op = obs_weight(pomdp, s, a, sp, o) # shortcut for observation probability from POMDPModelTools
                op2 = pdf(observation(pomdp, a, sp), o)

                bp[spi] += op * tp * b.b[si]
            end
        end
    end

    bp_sum = sum(bp)

    if bp_sum == 0.0
        error("""
              Failed discrete belief update: new probabilities sum to zero.

              b = $b
              a = $a
              o = $o

              Failed discrete belief update: new probabilities sum to zero.
              """)
    end

    # Normalize
    bp ./= bp_sum

    return DiscreteBelief(pomdp, ordered_stage_states(pomdp, stage(pomdp, o)+1), bp)
end

methods(update)

#test
let
    hor = 5
    pomdp = fixhorizon(TigerPOMDP(), hor)

    bu = DiscreteUpdater(pomdp.m)
    b0 = initialize_belief(bu, initialstate_distribution(pomdp.m))

    Γs = []
    # I only know correct alpha vectors from horizon stage
    for i in 1:4
        push!(Γs, [AlphaVec([0., 0.], 0), AlphaVec([0., 0.], 0), AlphaVec([0., 0.], 0)])
    end
    push!(Γs, [AlphaVec([10., 0.], 1), AlphaVec([0., 10.], 2), AlphaVec([-0.5, -0.5], 0)])

    staged_policies = [AlphaVectorPolicy(pomdp.m, length(stage_states(pomdp, t)), [a.alpha for a in Γs[t]],
        convert(Vector{actiontype(pomdp)}, [a.action for a in Γs[t]])) for t in 1:horizon(pomdp)]

    policy = StagedPolicy(pomdp, staged_policies)

    # these values were gotten from FiVI.jl tests
    @test Set(alphapairs(policy, hor)) == Set([[10., 0.]=>1, [0., 10.]=>2, [-0.5, -0.5]=>0])
    @test Set(alphavectors(policy, hor)) == Set([[10., 0.], [0., 10.], [-0.5, -0.5]])

    # initial belief is 50% tiger position confidence
    @test isapprox(value(policy, b0, hor), hor)
    @test isapprox(value(policy, [0.5, 0.5], hor), 5)
    # I am not sure whether the result should be [-0.5, 5. 5.] or [-1., 5. 5.]
    # may check the correctness in the future
    @test_broken isapprox(actionvalues(policy, b0, hor), [-1, 5., 5.])
    @test length(actionvalues(policy, b0, hor)) == length(actions(pomdp))

    # The agent only has one decision, so he has to risk choosing one door
    @test action(policy, b0, hor) == 1 || action(policy, b0, hor) == 2

    # try pushing new vector
    push!(policy, [0.0,0.0], 0, hor)
    @test isapprox(length(policy.staged_policies[hor].alphas), 4)

    # test simulation
    up = updater(policy)
    @test isapprox(simulate(RolloutSimulator(), pomdp, policy, up), 4.435, rtol=.1)

    # test infinite horizon pomdp
    pomdp = BabyPOMDP()
    @test_throws ArgumentError StagedPolicy(pomdp, staged_policies)
end


prune(strategy, m) = mappruneable(strategy, m; exclude = x -> isleaf(x) && x isa AbstractArray)
prune(strategies::Union{<:Tuple, <:AbstractVector}, m::Chain) = Chain(prune.(strategies, m)...)

function iterativeprune(finetune!, stages, m; maxtries = 5)
    for stage in stages
        # prune to current stage
        m̄ = prune(stage, m)
        @show count(iszero, m̄[end].weight) / length(m̄[end].weight)

        # finetune pruned model as needed
        attempts = 0
        while (attempts < maxtries) && !finetune!(m̄)
            attempts += 1
        end
        m = m̄
    end

    return m
end


#
#= = # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
# Description
#
#   This file contains functions to retrieve space indices for many models.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# References
#
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # = =#

export get_space_index, init_space_indices

# Interpolation types used in space indices.
_space_indices_itp_constant{T} =
    Interpolations.GriddedInterpolation{T,1,T,Gridded{Interpolations.Constant{Interpolations.Nearest}},Tuple{Array{Float64,1}}}

_space_indices_itp_linear{T} =
    Interpolations.GriddedInterpolation{T,1,T,Gridded{Linear},Tuple{Array{Float64,1}}}

include("./dtcfile.jl")
include("./fluxtable.jl")
include("./solfsmy.jl")
include("./wdcfiles.jl")
include("./dstfiles.jl")
include("./space_indices_helpers.jl")

################################################################################
#                                Private Macros
################################################################################

macro _check_data(itp, DT)
    quote
        if ($(esc(DT)) < $(esc(itp)).knots[1][1]) ||
           ($(esc(DT)) > $(esc(itp)).knots[1][end])
            error("The data for the requested Julian Day is not available!")
        end
    end
end

################################################################################
#                               Public Functions
################################################################################

"""
    get_space_index(T, DT::Number; ...)

Return the space index `T` at the day `DT` [Julian Day]. `T` can be:

## Daily 10.7-cm solar flux

The daily 10.7-cm solar flux can be obtained using:

* `F10()`: 10.7-cm adjusted solar flux \\[10⁻²² W/(M² Hz)].
* `F10adj()`: 10.7-cm adjusted solar flux \\[10⁻²² W/(M² Hz)].
* `F10obs()`: 10.7-cm observed solar flux \\[10⁻²² W/(M² Hz)].

These indices require `fluxtable` (see `init_space_indices`).

## Daily average 10.7-cm solar flux

The daily average 10.7-cm solar flux, centered at `DT`, can be obtained using:

* `F10M()`: 10.7-cm adjusted solar flux \\[10⁻²² W/(M² Hz)].
* `F10Madj()`: 10.7-cm adjusted solar flux \\[10⁻²² W/(M² Hz)].
* `F10Mobs()`: 10.7-cm observed solar flux \\[10⁻²² W/(M² Hz)].

In this case, the keyword `window::Int` can be passed to select the size of the
window. By default, it is selected as 81.

These indices require `fluxtable` (see `init_space_indices`).

# Daily Kp and Ap

* `Kp()`: Kp index (daily mean).
* `Kp_vect()`: A vector containing the Kp index for the following hours of the
               day: 0-3h, 3-6h, 6-9h, 9-12h, 12-15h, 15-18h, 18-20h, 20-23h.
* `Ap()`: Ap index (daily mean).
* `Ap_vect()`: A vector containing the Ap index for the following hours of the
               day: 0-3h, 3-6h, 6-9h, 9-12h, 12-15h, 15-18h, 18-20h, 20-23h.

These indices require `wdcfiles` (see `init_space_indices`).

# Daily S10, M10, and Y10

* `S10()`: EUV index (26-34 nm) scaled to F10.7.
* `M10()`: MG2 index scaled to F10.7.
* `Y10()`: Solar X-ray & Lya index scaled to F10.7.

These indices require `solfsmy` (see `init_space_indices`).

# 81-day centered average of S10, M10, and Y10.

* `S81a`: EUV 81-day averaged centered index.
* `M81a`: MG2 81-day averaged centered index.
* `Y81a`: Solar X-ray & Lya 81-day averaged centered index.

These indices require `solfsmy` (see `init_space_indices`).

# Exospheric temperature variation due to Dst

* `DstΔTc`: Exospheric temperature variation due to `Dst` [K].

This index requires `dtcfile` (see `init_space_indices`).

"""
@inline get_space_index(::SPI{:F10}, DT::Number) =
    get_space_index(Val(:F10adj), DT)

@inline function get_space_index(::SPI{:F10obs}, DT::Number)
    @_check_data(get(_fluxtable_data).F10obs, DT)
get(_fluxtable_data).F10obs(DT)
end

@inline function get_space_index(::SPI{:F10adj}, DT::Number)
    @_check_data(get(_fluxtable_data).F10adj, DT)
get(_fluxtable_data).F10adj(DT)
end

@inline get_space_index(::SPI{:F10M}, DT::Number; window::Int=81) =
    get_space_index(Val(:F10Madj), DT; window=window)

@inline function get_space_index(::SPI{:F10Mobs}, DT::Number; window::Int=81)
    Δ = floor(Int, (window - 1) / 2)
    @_check_data(get(_fluxtable_data).F10obs, DT - Δ)
    @_check_data(get(_fluxtable_data).F10obs, DT + Δ)
    mean(get(_fluxtable_data).F10obs(DT - Δ:1:DT + Δ))
end

@inline function get_space_index(::SPI{:F10Madj}, DT::Number; window::Int=81)
    Δ = floor(Int, (window - 1) / 2)
    @_check_data(get(_fluxtable_data).F10adj, DT - Δ)
    @_check_data(get(_fluxtable_data).F10adj, DT + Δ)
    mean(get(_fluxtable_data).F10adj(DT - Δ:1:DT + Δ))
end

@inline function get_space_index(::SPI{:Kp_vect}, DT::Number)
    @_check_data(get(_wdc_data).Kp, DT)
    get(_wdc_data).Kp(DT)
end

@inline function get_space_index(::SPI{:Ap_vect}, DT::Number)
    @_check_data(get(_wdc_data).Ap, DT)
    get(_wdc_data).Ap(DT)
end

@inline get_space_index(::SPI{:Kp}, DT::Number) =
    mean(get_space_index(Val(:Kp_vect), DT))

@inline get_space_index(::SPI{:Ap}, DT::Number) =
    mean(get_space_index(Val(:Ap_vect), DT))

@inline function get_space_index(::SPI{:Dst_vect}, DT::Number)
    @_check_data(get(_dst_data).Dst, DT)
    get(_dst_data).Dst(DT)
end

@inline get_space_index(::SPI{:Dst}, DT::Number) =
    mean(get_space_index(Val(:Dst_vect), DT))

@inline function get_space_index(::SPI{:S10}, DT::Number)
    @_check_data(get(_solfsmy_data).S10, DT)
    get(_solfsmy_data).S10(DT)
end

@inline function get_space_index(::SPI{:S81a}, DT::Number)
    @_check_data(get(_solfsmy_data).S81a, DT)
    get(_solfsmy_data).S81a(DT)
end

@inline function get_space_index(::SPI{:M10}, DT::Number)
    @_check_data(get(_solfsmy_data).M10, DT)
    get(_solfsmy_data).M10(DT)
end

@inline function get_space_index(::SPI{:M81a}, DT::Number)
    @_check_data(get(_solfsmy_data).M81a, DT)
    get(_solfsmy_data).M81a(DT)
end

@inline function get_space_index(::SPI{:Y10}, DT::Number)
    @_check_data(get(_solfsmy_data).Y10, DT)
    get(_solfsmy_data).Y10(DT)
end

@inline function get_space_index(::SPI{:Y81a}, DT::Number)
    @_check_data(get(_solfsmy_data).Y81a, DT)
    get(_solfsmy_data).Y81a(DT)
end

@inline function get_space_index(::SPI{:DstΔTc}, DT::Number)
    @_check_data(get(_dtcfile_data).DstΔTc, DT)
    get(_dtcfile_data).DstΔTc(DT)
end

"""
    init_space_indices(...)

Initialize all space indices. The files that will be initialized must be
indicated by the array of symbols passed to the keyword argument
`enabled_files`. If this is `nothing`, which is the default, then all files will
be initialized. The symbol related to each file is described next.

Notice that the initialization process can be changed by a set of keywords as
described next.

## DTCFILE

**Symbol**: `:dtcfile`

This file contains the exospheric temperature variation caused by the Dst index.
This is used for the JB2008 atmospheric model.

### Keywords

* `dtcfile_path`: Path for the file `DTCFILE.TXT`. If `nothing`, then it will be
                  downloaded. (**Default** = `nothing`)
* `dtcfile_force_download`: If `true`, then the file will always be downloaded
                            if the path is not specified. (**Default** =
                            `false`).

## fluxtable

**Symbol**: `:fluxtable`

This file contains the F10.7 flux data in different formats.

### Keywords

* `fluxtable_path`: Path for the file `fluxtable.txt`. If `nothing`, then it
                    will be downloaded. (**Default** = `nothing`)
* `fluxtable_force_download`: If `true`, then the file will always be downloaded
                              if the path is not specified.
                              (**Default** = `false`).

## SOLFSMY

**Symbol**: `:solfsmy`

This files contains the indices necessary for the JB2008 atmospheric model.

### Keywords

* `solfsmy_path`: Path for the file `SOLFSMY.TXT`. If `nothing`, then it will be
                  downloaded. (**Default** = `nothing`)
* `solfsmy_force_download`: If `true`, then the file will always be downloaded
                            if the path is not specified. (**Default** =
                            `false`).

## WDC Files

**Symbol**: `:wdcfiles`

This set of files contain the Kp and Ap indices.

### Keywords

* `wdcfiles_path`: Path for the directory with the WDC files. If `nothing`, then
                   they will be downloaded. (**Default** = `nothing`)
* `wdcfiles_force_download`: If `true`, then the files will always be downloaded
                            if the path is not specified. (**Default** =
                            `false`).
* `wdcfiles_oldest_year`: Oldest year in which the WDC file will be obtained.
                          (**Default** = past 3 years).
* `wdcfiles_newest_year`: Newest year in which the WDC file will be obtained. If
                          it is `nothing`, then it defaults to the current year.
                          (**Default** = `nothing`).

"""
function init_space_indices(;
            enabled_files=nothing,
            dtcfile_path=nothing,
            dtcfile_force_download=false,

            fluxtable_path=nothing,
            fluxtable_force_download=false,

            solfsmy_path=nothing,
            solfsmy_force_download=false,

            wdcfiles_dir=nothing,
            wdcfiles_force_download=false,
            wdcfiles_oldest_year=year(now()) - 3,
            wdcfiles_newest_year=nothing,

            dstfiles_dir=nothing,
            dstfiles_force_download=false,
            dstfiles_oldest_year=year(now()) - 3,
            dstfiles_newest_year=nothing
        )
    dtcfile   = (enabled_files == nothing) || (:dtcfile in enabled_files)
    fluxtable = (enabled_files == nothing) || (:fluxtable in enabled_files)
    solfsmy   = (enabled_files == nothing) || (:solfsmy in enabled_files)
    wdcfiles  = (enabled_files == nothing) || (:wdcfiles in enabled_files)
    dstfiles  = (enabled_files == nothing) || (:dstfiles in enabled_files)

    dtcfile && _init_dtcfile(local_path=dtcfile_path,
                            force_download=dtcfile_force_download)

    fluxtable && _init_fluxtable(local_path=fluxtable_path,
                                force_download=fluxtable_force_download)

    solfsmy && _init_solfsmy(local_path=solfsmy_path,
                            force_download=solfsmy_force_download)

    wdcfiles && _init_wdcfiles(local_dir=wdcfiles_dir,
                            force_download=wdcfiles_force_download,
                            wdcfiles_oldest_year=wdcfiles_oldest_year,
                            wdcfiles_newest_year=wdcfiles_newest_year)

    dstfiles && _init_dstfiles(local_dir=dstfiles_dir,
                            force_download=dstfiles_force_download,
                            dstfiles_oldest_year=dstfiles_oldest_year,
                            dstfiles_newest_year=dstfiles_newest_year)

    nothing
end


@testset "mountain_car" begin

    env = MountainCarEnv(; rng = StableRNG(123))
    RLBase.test_interfaces!(env)
    RLBase.test_runnable!(env)

    env = ContinuousMountainCarEnv(; rng = StableRNG(123))
    RLBase.test_interfaces!(env)
    RLBase.test_runnable!(env)

end


"""

$(TYPEDEF)

Structure to contain a list of all the distances smaller than
the cutoff.

$(TYPEDFIELDS)

"""
struct CutoffDistances

  # To store all distances smaller than the cutoff
  nd::Vector{Int} # Number of small distances (vector size = 1), just to be mutable
  d::Vector{Float64} # All distances smaller than the cutoff
  iat::Vector{Int} # atom of the solute
  jat::Vector{Int} # atom of the solvent
  imol::Vector{Int} # molecule of the solute
  jmol::Vector{Int} # molecule of the solvent

  # Size of the arrays
  maxdim::Vector{Int}

end

# Generator
CutoffDistances(natoms::Int) =
  CutoffDistances(zeros(Int,1), # nd
                  zeros(Float64,natoms), # d
                  zeros(Int,natoms), # iat
                  zeros(Int,natoms), # jat
                  zeros(Int,natoms), # imol
                  zeros(Int,natoms), # jmol
                  [ natoms ]) # maxdim

# Function that zeroes all the values in this structure
function reset!(c::CutoffDistances)
  c.nd[1] = 0
  @. c.d = 0.
  @. c.iat = 0
  @. c.jat = 0
  @. c.imol = 0
  @. c.jmol = 0
  return nothing
end

function reset!(c::Vector{CutoffDistances})
  for i in 1:size(c,1)
    reset!(c[i])
  end
  return nothing
end


module FuncKinds

# Many of these are written elaborately to defeat inference

function hasfoo(list)   # a test for anonymous functions
    hf = false
    hf = map(list) do item
        if isa(item, AbstractString)
            (str->occursin("foo", str))(item)
        else
            false
        end
    end
    return any(hf)
end

## A test for keyword functions
const randnums = Any[1, 1.0, UInt(14), 2.0f0]

@noinline returnsrandom() = randnums[rand(1:length(randnums))]

@noinline function haskw(x, y; a="hello", b=1, c=returnsrandom())
    if isa(b, Integer)
        return cos(rand()) + c + x + y
    elseif isa(b, AbstractFloat)
        s = 0.0
        for i = 1:rand(1:10)
            s += log(rand())
        end
        return s + c + x + y
    end
    return "string"
end

function callhaskw()
    ret = Any[]
    for i = 1:5
        push!(ret, haskw(returnsrandom(), returnsrandom()))
    end
    push!(ret, haskw(returnsrandom(), returnsrandom(); b = 2.0))
    return ret
end

@generated function gen(x::T) where T
    Tbigger = T == Float32 ? Float64 : BigFloat
    :(convert($Tbigger, x))
end

function gen2(x::Int, y)
    if @generated
        return y <: Integer ? :(x*y) : :(x+y)
    else
        return 2x+3y
    end
end

function hasinner(x, y)
    inner(z) = 2z

    s = 0
    for i = 1:10
        s += inner(returnsrandom())
    end
    return s
end

# Two kwarg generated functions; one will be called from the no-kw call, the other from a kwcall
@generated function genkw1(; b=2)
    :(string(typeof($b)))
end
@generated function genkw2(; b=2)
    :(string(typeof($b)))
end

# Function styles from JuliaInterpreter
f1(x::Int) = 1
f1(x) = 2
# where signatures
f2(x::T) where T = -1
f2(x::T) where T<:Integer = T
f2(x::T) where Unsigned<:T<:Real = 0
f2(x::V) where V<:SubArray{T} where T = 2
f2(x::V) where V<:Array{T,N} where {T,N} = 3
f2(x::V) where V<:Base.ReshapedArray{T,N} where T where N = 4
# Varargs
f3(x::Int, y...) = 1
f3(x::Int, y::Symbol...) = 2
f3(x::T, y::U...) where {T<:Integer,U} = U
f3(x::Array{Float64,K}, y::Vararg{Symbol,K}) where K = K
# Default args
f4(x, y::Int=0) = 2
f4(x::UInt, y="hello", z::Int=0) = 3
f4(x::Array{Float64,K}, y::Int=0) where K = K
# Keyword args
f5(x::Int8; y=0) = y
f5(x::Int16; y::Int=0) = 2
f5(x::Int32; y="hello", z::Int=0) = 3
f5(x::Int64;) = 4
f5(x::Array{Float64,K}; y::Int=0) where K = K
# Default and keyword args
f6(x, y="hello"; z::Int=0) = 1

end


using ArgParse
using Formatting
using JSON,TOML
using DataStructures

include("runCmd.jl")

s = ArgParseSettings()
@add_arg_table s begin

    "--casename"
        help = "Prefix of the file."
        arg_type = String
        required = true

    "--input-dir"
        help = "The folder where case folders are contained"
        arg_type = String
        required = true

    "--output-dir"
        help = "Output folder"
        arg_type = String
        required = true

    "--year-rng"
        help = "Range of years being diagnosed"
        arg_type = Int64
        nargs = 2
        required = true

    "--domain-file"
        help = "Domain file of ocean model"
        arg_type = String
        required = true

    "--overwrite"
        action = :store_true

    "--remap-file-nn"
        help = "Remap weighting file"
        arg_type = String
        required = true

    "--pop2"
        arg_type = Bool
        default = false
end

parsed = DataStructures.OrderedDict(parse_args(ARGS, s))

JSON.print(parsed, 4)

for y = parsed["year-rng"][1]:parsed["year-rng"][2], m=1:12

    date_str = format("{:04d}-{:02d}", y, m)

    if ! parsed["pop2"]
        old_file  = joinpath(parsed["input-dir"],  "$(parsed["casename"]).EMOM.h0.mon.$(date_str).nc")
        Nx = "Nx"
        Ny = "Ny"
        Nz = "Nz"
    else
        old_file  = joinpath(parsed["input-dir"],  "$(parsed["casename"]).pop.h.$(date_str).nc")
        Nx = "nlon"
        Ny = "nlat"
        Nz = "z_t"
    end

    new_file1 = joinpath(parsed["output-dir"], "$(parsed["casename"]).EMOM_extra1.$(date_str).nc")
    new_file2 = joinpath(parsed["output-dir"], "$(parsed["casename"]).EMOM_extra1_rg.$(date_str).nc")

    if !isfile(new_file1) || !isfile(new_file2) || parsed["overwrite"]

        println("Generating file: $(new_file1)")

#        pleaseRun(`ncap2 -h -O -v
#                    -s "SST=array(0.0,0.0,/\$time,\$$Ny,\$$Nx/);  SST(0,:,:)=TEMP(0, 0, :, :);"
#                    -s '*tmp=HMXL*1.0;'
#                    -s "HMXL=array(0.0,0.0,/\$time,\$$Ny,\$$Nx/); HMXL(0,:,:)=tmp;"
#                    $old_file $new_file1
#        `; igs=false)

        pleaseRun(`ncap2 -h -O -v
                    -s "SST=array(0.0,0.0,/\$time,\$$Ny,\$$Nx/);  SST(0,:,:)=TEMP(0, 0, :, :);"
                    $old_file $new_file1
        `; igs=false)

        if parsed["pop2"]
            pleaseRun(`ncrename -d nlat,Ny -d nlon,Nx -d z_t,Nz $new_file1`; igs=true)
        end
        println("Generating file: $(new_file2)")
        pleaseRun(`ncremap -R '--rgr lat_nm_in=Ny --rgr lon_nm_in=Nx' -m $(parsed["remap-file-nn"]) $new_file1 $new_file2`;igs=false)

    end

    # Special: stratification
    new_file3 = joinpath(parsed["output-dir"], "$(parsed["casename"]).EMOM_extra2.$(date_str).nc")
    new_file4 = joinpath(parsed["output-dir"], "$(parsed["casename"]).EMOM_extra2_rg.$(date_str).nc")
    if !isfile(new_file3) || parsed["overwrite"]

        println("Generating file: $(new_file3)")

        if !parsed["pop2"]
            dz_cmd = "*dz = dz_cT(:, 0, 0);"
        else
            dz_cmd = "*dz=array(0.0,0.0,/\$$Nz/); for (i=0;i<\$z_t.size;i++) { dz(i) = (z_w_bot(i) - z_w_top(i))/100.0; };"
        end
        pleaseRun(`ncap2 -h -O -v
                    -s 'defdim("Nz1", 5); defdim("Nz2", 28);'
                    -s "$dz_cmd"
                    -s '*dz1=dz(0:4).ttl(); *dz2=dz(5:32).ttl();'
                    -s "*TEMPdz[\$time,\$$Nz,\$$Ny,\$$Nx]=TEMP*dz; TEMPdz.set_miss(TEMP@_FillValue);"
                    -s "*_T1=array(0.0,0.0,/\$time,\$Nz1,\$$Ny,\$$Nx/);_T1.set_miss(TEMP@_FillValue);"
                    -s "*_T2=array(0.0,0.0,/\$time,\$Nz2,\$$Ny,\$$Nx/);_T2.set_miss(TEMP@_FillValue);"
                    -s '_T1(:, :, :, :)=TEMPdz(:, 0:4, :, :);'
                    -s '_T2(:, :, :, :)=TEMPdz(:, 5:32, :, :);'
                    -s 'T1=_T1.ttl($Nz1)/dz1;'
                    -s 'T2=_T2.ttl($Nz2)/dz2;'
                    -s 'STRAT=T1-T2;'
                    $old_file $new_file3
        `; igs=false)

        println("Generating file: $(new_file4)")
        pleaseRun(`ncremap -R '--rgr lat_nm_in=Ny --rgr lon_nm_in=Nx' -m $(parsed["remap-file-nn"]) $new_file3 $new_file4`;igs=false)

    end

    # Ocean heat content
    new_file5 = joinpath(parsed["output-dir"], "$(parsed["casename"]).EMOM_extra3.$(date_str).nc")
    if !isfile(new_file5) || parsed["overwrite"]

        println("Generating file: $(new_file5)")

        if !parsed["pop2"]
            pleaseRun(`ncwa -a N1 -v TEMP,SALT,area_sT $old_file $new_file5`)
            pleaseRun(`ncwa -O -w area_sT -a Nx,Ny -v TEMP,SALT $new_file5 $new_file5`)
        else
            pleaseRun(`ncwa -w TAREA -a nlon,nlat -v TEMP,SALT $old_file $new_file5`)
        end

    end

    # Ocean heat transport
    new_file6 = joinpath(parsed["output-dir"], "$(parsed["casename"]).EMOM_extra4_OHT.$(date_str).nc")
    if !isfile(new_file6) || parsed["overwrite"]

        println("Generating file: $(new_file6)")

        if !parsed["pop2"]
            pleaseRun(`ncwa -a N1 -v ADVT,ADVS,WKRSTT,WKRSTS,Q_LOST,TEMP,SALT,dz_cT $old_file $new_file6`)
            pleaseRun(`ncwa -O -y ttl -w dz_cT -a Nz -v ADVT,ADVS,WKRSTT,WKRSTS,Q_LOST,TEMP,SALT $new_file6 $new_file6`)
#            pleaseRun(`ncap2 -O -s 'SHF=ADVT+WKRSTT;' $new_file6 $new_file6`)
        else
            pleaseRun(`ncwa -O -y ttl -w dz -a z_t -v TEMP,SALT $old_file $new_file6`)
            pleaseRun(`ncks -A -v SHF,N_HEAT $old_file $new_file6`)
        end

    end

end


### Guru Interfaces
finufft_plan{T} = Ptr{T} where T <: Union{Float32,Float64}

function finufft_makeplan(type::Integer,
                          dim::Integer,
                          n_modes::T,
                          iflag::Integer,
                          ntrans::Integer,
                          eps::Float64,
                          opts::nufft_opts=finufft_default_opts()) where T <: Union{Array{BIGINT}, NTuple{3,BIGINT}}

    # see https://stackoverflow.com/questions/40140699/the-proper-way-to-declare-c-void-pointers-in-julia for how to declare c-void pointers in julia
    plan_p = Ref{finufft_plan{Float64}}()

    ret = ccall( (:finufft_makeplan, libfinufft),
            Cint,
            (Cint,
            Cint,
            Ref{BIGINT},
            Cint,
            Cint,
            Cdouble,
            Ptr{finufft_plan{Cdouble}},
            Ref{nufft_opts}),
            type,dim,n_modes,iflag,ntrans,eps,plan_p,opts
            )
    check_ret(ret)

    plan = plan_p[]
    return plan
end
function finufft_makeplan(type::Integer,
    dim::Integer,
    n_modes::T,
    iflag::Integer,
    ntrans::Integer,
    eps::Float32,
    opts::nufft_opts=finufft_default_opts()) where T <: Union{Array{BIGINT}, NTuple{3,BIGINT}}

    # see https://stackoverflow.com/questions/40140699/the-proper-way-to-declare-c-void-pointers-in-julia for how to declare c-void pointers in julia
    plan_p = Ref{finufft_plan{Float32}}()

    ret = ccall( (:finufftf_makeplan, libfinufft),
            Cint,
            (Cint,
            Cint,
            Ref{BIGINT},
            Cint,
            Cint,
            Cfloat,
            Ptr{finufft_plan{Cfloat}},
            Ref{nufft_opts}),
            type,dim,n_modes,iflag,ntrans,eps,plan_p,opts
            )
    check_ret(ret)

    plan = plan_p[]
    return plan
end

function finufft_setpts(plan::finufft_plan{Float64},
                        xj::Array{Float64},
                        yj::Array{Float64}=Float64[],
                        zj::Array{Float64}=Float64[],
                        s::Array{Float64}=Float64[],
                        t::Array{Float64}=Float64[],
                        u::Array{Float64}=Float64[])

    nj = length(xj)
    nk = length(s)

    ret = ccall( (:finufft_setpts, libfinufft),
                Cint,
                (finufft_plan{Cdouble},
                BIGINT,
                Ref{Cdouble},
                Ref{Cdouble},
                Ref{Cdouble},
                BIGINT,
                Ref{Cdouble},
                Ref{Cdouble},
                Ref{Cdouble}),
                plan,nj,xj,yj,zj,nk,s,t,u
                )

    check_ret(ret)
end
function finufft_setpts(plan::finufft_plan{Float32},
    xj::Array{Float32},
    yj::Array{Float32}=Float32[],
    zj::Array{Float32}=Float32[],
    s::Array{Float32}=Float32[],
    t::Array{Float32}=Float32[],
    u::Array{Float32}=Float32[])

    nj = length(xj)
    nk = length(s)

    ret = ccall( (:finufftf_setpts, libfinufft),
                Cint,
                (finufft_plan{Cfloat},
                BIGINT,
                Ref{Cfloat},
                Ref{Cfloat},
                Ref{Cfloat},
                BIGINT,
                Ref{Cfloat},
                Ref{Cfloat},
                Ref{Cfloat}),
                plan,nj,xj,yj,zj,nk,s,t,u
                )

    check_ret(ret)
end

function finufft_exec(plan::finufft_plan{Float64}, cj::Array{ComplexF64}, fk::Array{ComplexF64})

    ret = ccall( (:finufft_execute, libfinufft),
                Cint,
                (finufft_plan{Cdouble},
                Ref{ComplexF64},
                Ref{ComplexF64}),
                plan,cj,fk
                )

    check_ret(ret)
end
function finufft_exec(plan::finufft_plan{Float32}, cj::Array{ComplexF32}, fk::Array{ComplexF32})

    ret = ccall( (:finufftf_execute, libfinufft),
                Cint,
                (finufft_plan{Cfloat},
                Ref{ComplexF32},
                Ref{ComplexF32}),
                plan,cj,fk
                )

    check_ret(ret)
end

function finufft_destroy(plan::finufft_plan{Float64})
    ret = ccall( (:finufft_destroy, libfinufft),
                Cint,
                (finufft_plan{Cdouble},),
                plan
                )

    check_ret(ret)
end
function finufft_destroy(plan::finufft_plan{Float32})
    ret = ccall( (:finufftf_destroy, libfinufft),
            Cint,
            (finufft_plan{Cfloat},),
            plan
            )

    check_ret(ret)
end


struct NamedVector{K,L,T,V} <: FieldVector{L,T}
    nt::NamedTuple{K,V}
    NamedVector{K,L,T,V}(nt::NT) where {K,L,T,V<:Tuple,NT<:NamedTuple{K,Tuple{Vararg{T,L}}}} = new{K,L,T,V}(nt)
end
NamedVector{K,L,T,V}(t::Tuple{Vararg{T,L}}) where {K,L,T,V<:Tuple} = NamedVector{K,L,T,V}(NamedTuple{K}(t))
NamedVector{K,L,T,V}(x::T, xs::T...) where {K,L,T,V} = NamedVector{K,L,T,V}((x, xs...))

NamedVector{K,L,T}(t::V) where {K,L,T,V<:Tuple} = NamedVector{K,L,T,V}(t)
NamedVector{K,L,T}(t::V) where {K,L,T,V<:NamedTuple{K,V1}} where V1 = NamedVector{K,L,T,V1}(t)
NamedVector(nt::NamedTuple{K,V}) where {K,V<:Tuple{Vararg{Any,L}}} where L = NamedVector{K,L}(nt)
NamedVector{K}(t::V) where {K,V<:Tuple{Vararg{<:Any,L}}} where L = NamedVector{K,L}(t)
# Type promotion step...
function NamedVector{K,L}(t::V) where {K,L,V<:Tuple{Vararg{<:Any,L}}}
    t_uniform_type = map(x -> convert(promote_type(map(typeof, t)...), x), t)
    T = promote_type(map(typeof, t)...)
    NamedVector{K,L,T,NTuple{L,T}}(NamedTuple{K,NTuple{L,T}}(t_uniform_type))
end
function NamedVector{K,L}(nt::NT) where {K,L,NT<:NamedTuple{K,<:Tuple{Vararg{<:Any,L}}}}
    # TODO: why doesn't this compile away...
    NamedVector{K,L}(Tuple(nt))
end
function NamedVector{K,L}(v::SVector{L,T}) where {K,L,T}
    NamedVector{K,L,T,NTuple{L,T}}(NamedTuple{K,NTuple{L,T}}(Tuple(v)))
end
NamedVector(; kw...) = NamedVector{keys(kw),length(kw)}(values(kw))

Base.parent(a::NamedVector) = getfield(a, :nt)
Base.values(a::NamedVector) = values(parent(a))
Base.getproperty(a::NamedVector, x::Symbol) = getproperty(parent(a), x)
Base.@propagate_inbounds Base.getindex(a::NamedVector, i::Int) = getindex(parent(a), i)
Base.@propagate_inbounds Base.setindex!(a::NamedVector, x, i::Int) = (setindex!(parent(a), i, x); a)
# Base.map(f, as::Union{NamedVector{K},NamedTuple{K}}...) where K = begin
Base.map(f, a1::NamedVector{K,L}, as::NamedVector{K}...) where {K,L} = begin
    NamedVector{K,L}(map(f, map(SVector, (a1, as...))...))
end

_maybeparent(a::NamedVector) = parent(a)
_maybeparent(nt::NamedTuple) = nt

Base.tail(a::NamedVector{K,L,T}) where {K,L,T} = NamedVector(Base.tail(parent(a)))


get_connection_type(s) = getmetadata(unwrap(s), VariableConnectType, Equality)

function with_connector_type(expr)
    @assert expr isa Expr && (expr.head == :function || (expr.head == :(=) &&
                                       expr.args[1] isa Expr &&
                                       expr.args[1].head == :call))

    sig = expr.args[1]
    body = expr.args[2]

    fname = sig.args[1]
    args = sig.args[2:end]

    quote
        function $fname($(args...))
            function f()
                $body
            end
            res = f()
            $isdefined(res, :connector_type) && $getfield(res, :connector_type) === nothing ? $Setfield.@set!(res.connector_type = $connector_type(res)) : res
        end
    end
end

macro connector(expr)
    esc(with_connector_type(expr))
end

abstract type AbstractConnectorType end
struct StreamConnector <: AbstractConnectorType end
struct RegularConnector <: AbstractConnectorType end

function connector_type(sys::AbstractSystem)
    sts = get_states(sys)
    #TODO: check the criteria for stream connectors
    n_stream = 0
    n_flow = 0
    for s in sts
        vtype = get_connection_type(s)
        if vtype === Stream
            isarray(s) && error("Array stream variables are not supported. Got $s.")
            n_stream += 1
        end
        vtype === Flow && (n_flow += 1)
    end
    (n_stream > 0 && n_flow > 1) && error("There are multiple flow variables in $(nameof(sys))!")
    n_stream > 0 ? StreamConnector() : RegularConnector()
end

Base.@kwdef struct Connection
    inners = nothing
    outers = nothing
end

# everything is inner by default until we expand the connections
Connection(syss) = Connection(inners=syss)
get_systems(c::Connection) = c.inners
function Base.in(e::Symbol, c::Connection)
    any(k->nameof(k) === e, c.inners) || any(k->nameof(k) === e, c.outers)
end

const EMPTY_VEC = []

function Base.show(io::IO, ::MIME"text/plain", c::Connection)
    # It is a bit unfortunate that the display of an array of `Equation`s won't
    # call this.
    @unpack outers, inners = c
    if outers === nothing && inners === nothing
        print(io, "<Connection>")
    else
        syss = Iterators.flatten((something(inners, EMPTY_VEC), something(outers, EMPTY_VEC)))
        splitting_idx = length(inners)
        sys_str = join((string(nameof(s)) * (i <= splitting_idx ? ("::inner") : ("::outer")) for (i, s) in enumerate(syss)), ", ")
        print(io, "<", sys_str, ">")
    end
end

# symbolic `connect`
function connect(sys1::AbstractSystem, sys2::AbstractSystem, syss::AbstractSystem...)
    syss = (sys1, sys2, syss...)
    length(unique(nameof, syss)) == length(syss) || error("connect takes distinct systems!")
    Equation(Connection(), Connection(syss)) # the RHS are connected systems
end

# the actual `connect`.
function connect(c::Connection; check=true)
    @unpack inners, outers = c

    flow_eqs = Equation[]
    other_eqs = Equation[]

    ncnts = length(inners) + length(outers)
    cnts = Iterators.flatten((inners, outers))
    fs, ss = Iterators.peel(cnts)
    splitting_idx = length(inners) # anything after the splitting_idx is outer.
    first_sts = get_states(fs)
    first_sts_set = Set(getname.(first_sts))
    for sys in ss
        current_sts = getname.(get_states(sys))
        Set(current_sts) == first_sts_set || error("$(nameof(sys)) ($current_sts) doesn't match the connection type of $(nameof(fs)) ($first_sts).")
    end

    seen = Set()
    ceqs = Equation[]
    for s in first_sts
        name = getname(s)
        fix_val = getproperty(fs, name) # representative
        fix_val in seen && continue
        push!(seen, fix_val)

        vtype = get_connection_type(fix_val)
        vtype === Stream && continue

        isarr = isarray(fix_val)

        if vtype === Flow
            rhs = isarr ? zeros(Int, ncnts) : 0
            for (i, c) in enumerate(cnts)
                isinner = i <= splitting_idx
                # https://specification.modelica.org/v3.4/Ch15.html
                var = scalarize(getproperty(c, name))
                rhs += isinner ? var : -var
            end
            if isarr
                for r in rhs
                    push!(ceqs, 0 ~ r)
                end
            else
                push!(ceqs, 0 ~ rhs)
            end
        else # Equality
            for c in ss
                var = getproperty(c, name)
                if isarr
                    vs = scalarize(var)
                    for (i, v) in enumerate(vs)
                        push!(ceqs, fix_val[i] ~ v)
                    end
                else
                    push!(ceqs, fix_val ~ var)
                end
            end
        end
    end

    return ceqs
end

instream(a) = term(instream, unwrap(a), type=symtype(a))
SymbolicUtils.promote_symtype(::typeof(instream), _) = Real

isconnector(s::AbstractSystem) = has_connector_type(s) && get_connector_type(s) !== nothing
isstreamconnector(s::AbstractSystem) = isconnector(s) && get_connector_type(s) isa StreamConnector
isstreamconnection(c::Connection) = any(isstreamconnector, c.inners) || any(isstreamconnector, c.outers)

function print_with_indent(n, x)
    print(" " ^ n)
    show(stdout, MIME"text/plain"(), x)
    println()
end

function split_sys_var(var)
    var_name = string(getname(var))
    sidx = findlast(isequal('₊'), var_name)
    sidx === nothing && error("$var is not a namespaced variable")
    connector_name = Symbol(var_name[1:prevind(var_name, sidx)])
    streamvar_name = Symbol(var_name[nextind(var_name, sidx):end])
    connector_name, streamvar_name
end

function flowvar(sys::AbstractSystem)
    sts = get_states(sys)
    for s in sts
        vtype = get_connection_type(s)
        vtype === Flow && return s
    end
    error("There in no flow variable in $(nameof(sys))")
end

collect_instream!(set, eq::Equation) = collect_instream!(set, eq.lhs) | collect_instream!(set, eq.rhs)

function collect_instream!(set, expr, occurs=false)
    istree(expr) || return occurs
    op = operation(expr)
    op === instream && (push!(set, expr); occurs = true)
    for a in SymbolicUtils.unsorted_arguments(expr)
        occurs |= collect_instream!(set, a, occurs)
    end
    return occurs
end

positivemax(m, ::Any; tol=nothing)= max(m, something(tol, 1e-8))

function expand_connections(sys::AbstractSystem; debug=false, tol=1e-10)
    subsys = get_systems(sys)
    isempty(subsys) && return sys

    # post order traversal
    @set! sys.systems = map(s->expand_connections(s, debug=debug), subsys)

    outer_connectors = Symbol[]
    for s in subsys
        n = nameof(s)
        isconnector(s) && push!(outer_connectors, n)
    end
    isouter = let outer_connectors=outer_connectors
        function isouter(sys)::Bool
            s = string(nameof(sys))
            isconnector(sys) || error("$s is not a connector!")
            idx = findfirst(isequal('₊'), s)
            parent_name = Symbol(idx === nothing ? s : s[1:prevind(s, idx)])
            parent_name in outer_connectors
        end
    end

    sys = flatten(sys)
    eqs′ = get_eqs(sys)
    eqs = Equation[]
    instream_eqs = Equation[]
    instream_exprs = []
    cts = [] # connections
    for eq in eqs′
        if eq.lhs isa Connection
            push!(cts, get_systems(eq.rhs))
        elseif collect_instream!(instream_exprs, eq)
            push!(instream_eqs, eq)
        else
            push!(eqs, eq) # split connections and equations
        end
    end

    # if there are no connections, we are done
    isempty(cts) && return sys

    sys2idx = Dict{Symbol,Int}() # system (name) to n-th connect statement
    narg_connects = Connection[]
    for (i, syss) in enumerate(cts)
        # find intersecting connections
        exclude = 0 # exclude the intersecting system
        idx = 0     # idx of narg_connects
        for (j, s) in enumerate(syss)
            idx′ = get(sys2idx, nameof(s), nothing)
            idx′ === nothing && continue
            idx = idx′
            exclude = j
        end
        if exclude == 0
            outers = []
            inners = []
            for s in syss
                isouter(s) ? push!(outers, s) : push!(inners, s)
            end
            push!(narg_connects, Connection(outers=outers, inners=inners))
            for s in syss
                sys2idx[nameof(s)] = length(narg_connects)
            end
        else
            # fuse intersecting connections
            for (j, s) in enumerate(syss); j == exclude && continue
                sys2idx[nameof(s)] = idx
                c = narg_connects[idx]
                isouter(s) ? push!(c.outers, s) : push!(c.inners, s)
            end
        end
    end

    isempty(narg_connects) && error("Unreachable reached. Please file an issue.")

    @set! sys.connections = narg_connects

    # Bad things happen when there are more than one intersections
    for c in narg_connects
        @unpack outers, inners = c
        len = length(outers) + length(inners)
        allconnectors = Iterators.flatten((outers, inners))
        dups = find_duplicates(nameof(c) for c in allconnectors)
        length(dups) == 0 || error("Connection([$(join(map(nameof, allconnectors), ", "))]) has duplicated connections: [$(join(collect(dups), ", "))].")
    end

    if debug
        println("============BEGIN================")
        println("Connections for [$(nameof(sys))]:")
        foreach(Base.Fix1(print_with_indent, 4), narg_connects)
    end

    connection_eqs = Equation[]
    for c in narg_connects
        ceqs = connect(c)
        debug && append!(connection_eqs, ceqs)
        append!(eqs, ceqs)
    end

    if debug
        println("Connection equations:")
        foreach(Base.Fix1(print_with_indent, 4), connection_eqs)
        println("=============END=================")
    end

    # stream variables
    stream_connects = filter(isstreamconnection, narg_connects)
    instream_eqs, additional_eqs = expand_instream(instream_eqs, instream_exprs, stream_connects; debug=debug, tol=tol)

    @set! sys.eqs = [eqs; instream_eqs; additional_eqs]
    return sys
end

function expand_instream(instream_eqs, instream_exprs, connects; debug=false, tol)
    sub = Dict()
    seen = Set()
    for ex in instream_exprs
        var = only(arguments(ex))
        connector_name, streamvar_name = split_sys_var(var)

        # find the connect
        cidx = findfirst(c->connector_name in c, connects)
        cidx === nothing && error("$var is not a variable inside stream connectors")
        connect = connects[cidx]
        connectors = Iterators.flatten((connect.inners, connect.outers))
        # stream variable
        sv = getproperty(first(connectors), streamvar_name; namespace=false)
        inner_sc = connect.inners
        outer_sc = connect.outers

        n_outers = length(outer_sc)
        n_inners = length(inner_sc)
        outer_names = (nameof(s) for s in outer_sc)
        inner_names = (nameof(s) for s in inner_sc)
        if debug
            println("Expanding: $ex")
            isempty(inner_names) || println("Inner connectors: $(collect(inner_names))")
            isempty(outer_names) || println("Outer connectors: $(collect(outer_names))")
        end

        # expand `instream`s
        # https://specification.modelica.org/v3.4/Ch15.html
        # Based on the above requirements, the following implementation is
        # recommended:
        if n_inners == 1 && n_outers == 0
            connector_name === only(inner_names) || error("$var is not in any stream connector of $(nameof(ogsys))")
            sub[ex] = var
        elseif n_inners == 2 && n_outers == 0
            connector_name in inner_names || error("$var is not in any stream connector of $(nameof(ogsys))")
            idx = findfirst(c->nameof(c) === connector_name, inner_sc)
            other = idx == 1 ? 2 : 1
            sub[ex] = states(inner_sc[other], sv)
        elseif n_inners == 1 && n_outers == 1
            isinner = connector_name === only(inner_names)
            isouter = connector_name === only(outer_names)
            (isinner || isouter) || error("$var is not in any stream connector of $(nameof(ogsys))")
            if isinner
                outerstream = states(only(outer_sc), sv) # c_1.h_outflow
                sub[ex] = outerstream
            end
        else
            fv = flowvar(first(connectors))
            i = findfirst(c->nameof(c) === connector_name, inner_sc)
            if i !== nothing
                si = isempty(outer_sc) ? 0 : sum(s->max(states(s, fv), 0), outer_sc)
                for j in 1:n_inners; j == i && continue
                    f = states(inner_sc[j], fv)
                    si += max(-f, 0)
                end

                num = 0
                den = 0
                for j in 1:n_inners; j == i && continue
                    f = states(inner_sc[j], fv)
                    tmp = positivemax(-f, si; tol=tol)
                    den += tmp
                    num += tmp * states(inner_sc[j], sv)
                end
                for k in 1:n_outers
                    f = states(outer_sc[k], fv)
                    tmp = positivemax(f, si; tol=tol)
                    den += tmp
                    num += tmp * instream(states(outer_sc[k], sv))
                end
                sub[ex] = num / den
            end

        end
    end

    # additional equations
    additional_eqs = Equation[]
    for c in connects
        outer_sc = c.outers
        isempty(outer_sc) && continue
        inner_sc = c.inners
        n_outers = length(outer_sc)
        n_inners = length(inner_sc)
        connector_representative = first(outer_sc)
        fv = flowvar(connector_representative)
        for sv in get_states(connector_representative)
            vtype = get_connection_type(sv)
            vtype === Stream || continue
            if n_inners == 1 && n_outers == 1
                innerstream = states(only(inner_sc), sv)
                outerstream = states(only(outer_sc), sv)
                push!(additional_eqs, outerstream ~ innerstream)
            elseif n_inners == 0 && n_outers == 2
                # we don't expand `instream` in this case.
                v1 = states(outer_sc[1], sv)
                v2 = states(outer_sc[2], sv)
                push!(additional_eqs, v1 ~ instream(v2))
                push!(additional_eqs, v2 ~ instream(v1))
            else
                sq = 0
                for q in 1:n_outers
                    sq += sum(s->max(-states(s, fv), 0), inner_sc)
                    for k in 1:n_outers; k == q && continue
                        f = states(outer_sc[k], fv)
                        sq += max(f, 0)
                    end

                    num = 0
                    den = 0
                    for j in 1:n_inners
                        f = states(inner_sc[j], fv)
                        tmp = positivemax(-f, sq; tol=tol)
                        den += tmp
                        num += tmp * states(inner_sc[j], sv)
                    end
                    for k in 1:n_outers; k == q && continue
                        f = states(outer_sc[k], fv)
                        tmp = positivemax(f, sq; tol=tol)
                        den += tmp
                        num += tmp * instream(states(outer_sc[k], sv))
                    end
                    push!(additional_eqs, states(outer_sc[q], sv) ~ num / den)
                end
            end
        end
    end

    instream_eqs = map(Base.Fix2(substitute, sub), instream_eqs)
    if debug
        println("===========BEGIN=============")
        println("Expanded equations:")
        for eq in instream_eqs
            print_with_indent(4, eq)
        end
        if !isempty(additional_eqs)
            println("Additional equations:")
            for eq in additional_eqs
                print_with_indent(4, eq)
            end
        end
        println("============END==============")
    end
    return instream_eqs, additional_eqs
end


module FlowWorkspace

    using GigaSOM,DataFrames
    using EzXML,MetaGraphs

    using PolygonOps, StaticArrays
    using Glob: glob
    using HTTP: escapeuri

    include("Groups.jl")
    include("Gating.jl")
    include("Utils.jl")

    export load
end


@testset "Batchnorm" begin
    v = CUDA.rand(Float32, 2)
    m = CUDA.rand(Float32, 2, 5)

    @testset for training in (true, false), track_stats in (true, false)
        kws = (training=training, track_stats=track_stats)

        # Normal
        NNlibCUDA.batchnorm(v, v, m, v, v, 1.0; kws...)
        NNlibCUDA.∇batchnorm(v, v, m, m, v, v, 1.0; kws...)

        # No affine
        NNlibCUDA.batchnorm(nothing, nothing, m, v, v, 1.0; kws...)
        NNlibCUDA.∇batchnorm(nothing, nothing, m, m, v, v, 1.0; kws...)

        # No tracking
        NNlibCUDA.batchnorm(v, v, m, nothing, nothing, 1.0; kws...)
        NNlibCUDA.∇batchnorm(v, v, m, m, nothing, nothing, 1.0; kws...)

        # Both or neither tracked or affine params must be set
        for (α, β) in ((v, nothing), (nothing, v))
            @test_throws MethodError NNlibCUDA.batchnorm(α, β, m, v, v, 1.0; kws...)
            @test_throws MethodError NNlibCUDA.∇batchnorm(α, β, m, m, v, v, 1.0; kws...)
            @test_throws ArgumentError NNlibCUDA.batchnorm(v, v, m, α, β, 1.0; kws...)
        end
    end
end


### Abstract MALA state

abstract type MALAState{F<:VariateForm} <: LMCSamplerState{F} end

### MALA state subtypes

## UnvMALAState holds the internal state ("local variables") of the MALA sampler for univariate parameters

mutable struct UnvMALAState <: MALAState{Univariate}
  pstate::ParameterState{Continuous, Univariate} # Parameter state used internally by MALA
  tune::MCTunerState
  ratio::Real
  μ::Real
  diagnosticindices::Dict{Symbol, Integer}

  function UnvMALAState(
    pstate::ParameterState{Continuous, Univariate},
    tune::MCTunerState,
    ratio::Real,
    μ::Real,
    diagnosticindices::Dict{Symbol, Integer}
  )
    if !isnan(ratio)
      @assert ratio > 0 "Acceptance ratio should be positive"
    end
    new(pstate, tune, ratio, μ, diagnosticindices)
  end
end

UnvMALAState(pstate::ParameterState{Continuous, Univariate}, tune::MCTunerState=BasicMCTune()) =
  UnvMALAState(pstate, tune, NaN, NaN, Dict{Symbol, Integer}())

## MuvMALAState holds the internal state ("local variables") of the MALA sampler for multivariate parameters

mutable struct MuvMALAState <: MALAState{Multivariate}
  pstate::ParameterState{Continuous, Multivariate} # Parameter state used internally by MALA
  tune::MCTunerState
  ratio::Real
  μ::RealVector
  diagnosticindices::Dict{Symbol, Integer}

  function MuvMALAState(
    pstate::ParameterState{Continuous, Multivariate},
    tune::MCTunerState,
    ratio::Real,
    μ::RealVector,
    diagnosticindices::Dict{Symbol, Integer}
  )
    if !isnan(ratio)
      @assert ratio > 0 "Acceptance ratio should be positive"
    end
    new(pstate, tune, ratio, μ, diagnosticindices)
  end
end

MuvMALAState(pstate::ParameterState{Continuous, Multivariate}, tune::MCTunerState=BasicMCTune()) =
  MuvMALAState(pstate, tune, NaN, Array{eltype(pstate)}(pstate.size), Dict{Symbol, Integer}())

### Metropolis-adjusted Langevin Algorithm (MALA)

struct MALA <: LMCSampler
  driftstep::Real

  function MALA(driftstep::Real)
    @assert driftstep > 0 "Drift step is not positive"
    new(driftstep)
  end
end

MALA() = MALA(1.)

### Initialize MALA sampler

## Initialize parameter state

function initialize!(
  pstate::ParameterState{Continuous, F},
  parameter::Parameter{Continuous, F},
  sampler::MALA,
  outopts::Dict
) where F<:VariateForm
  parameter.uptogradlogtarget!(pstate)
  @assert isfinite(pstate.logtarget) "Log-target not finite: initial value out of support"
  @assert all(isfinite.(pstate.gradlogtarget)) "Gradient of log-target not finite: initial values out of support"

  if !isempty(outopts[:diagnostics])
    pstate.diagnostickeys = copy(outopts[:diagnostics])
    pstate.diagnosticvalues = Array{Any}(length(pstate.diagnostickeys))
  end
end

## Initialize MALA state

# Stopped at this point, fixing the sampler_state functions

function sampler_state(
  parameter::Parameter{Continuous, Univariate},
  sampler::MALA,
  tuner::MCTuner,
  pstate::ParameterState{Continuous, Univariate},
  vstate::VariableStateVector,
  diagnostickeys::Vector{Symbol}
)
  sstate = UnvMALAState(
    generate_empty(pstate, parameter.diffmethods, parameter.diffopts), tuner_state(parameter, sampler, tuner)
  )
  set_diagnosticindices!(sstate, [:accept], diagnostickeys)
  sstate
end

function sampler_state(
  parameter::Parameter{Continuous, Multivariate},
  sampler::MALA,
  tuner::MCTuner,
  pstate::ParameterState{Continuous, Multivariate},
  vstate::VariableStateVector,
  diagnostickeys::Vector{Symbol}
)
  sstate = MuvMALAState(
    generate_empty(pstate, parameter.diffmethods, parameter.diffopts), tuner_state(parameter, sampler, tuner)
  )
  set_diagnosticindices!(sstate, [:accept], diagnostickeys)
  sstate
end

## Reset parameter state

function reset!(
  pstate::ParameterState{Continuous, Univariate},
  x::Real,
  parameter::Parameter{Continuous, Univariate},
  sampler::MALA
)
  pstate.value = x
  parameter.uptogradlogtarget!(pstate)
end

function reset!(
  pstate::ParameterState{Continuous, Multivariate},
  x::RealVector,
  parameter::Parameter{Continuous, Multivariate},
  sampler::MALA
)
  pstate.value = copy(x)
  parameter.uptogradlogtarget!(pstate)
end

show(io::IO, sampler::MALA) = print(io, "MALA sampler: drift step = $(sampler.driftstep)")


# Julia interface when not defined in source files

eltype{T}(x::DistMultiVec{T}) = T

function size(A::ElementalMatrix, i::Integer)
    if i < 1
        error("dimension out of range")
    elseif i == 1
        return Int(height(A))
    elseif i == 2
        return Int(width(A))
    else
        return 1
    end
end

size(A::ElementalMatrix) = (size(A, 1), size(A, 2))

(*){T<:ElementalMatrix}(A::T, B::T)      = A_mul_B!(one(eltype(A)), A, B, zero(eltype(A)), similar(A, (size(A, 1), size(B, 2))))
(*){T}(A::DistSparseMatrix{T}, B::DistMultiVec{T}) = A_mul_B!(one(T), A, B, zero(T), similar(B, (size(A, 1), size(B, 2))))
Ac_mul_B{T<:ElementalMatrix}(A::T, B::T) = Ac_mul_B!(one(eltype(A)), A, B, zero(eltype(A)), similar(A, (size(A, 2), size(B, 2))))
Ac_mul_B{T}(A::DistSparseMatrix{T}, B::DistMultiVec{T}) = Ac_mul_B!(one(T), A, B, zero(T), similar(B, (size(A, 2), size(B, 2))))

# Spectral
svd(A::ElementalMatrix) = svd!(copy(A))
svd(A::ElementalMatrix, ctrl::SVDCtrl) = svd!(copy(A), ctrl)
svdvals(A::ElementalMatrix, ctrl::SVDCtrl) = svdvals!(copy(A), ctrl)

# conversions to and from julia arrays

# function copy!{T}(dest::Matrix{T}, src::Base.VecOrMat{T})
#     m, n = size(src, 1), size(src, 2)
#     resize!(dest, m, n)
#     Base.unsafe_copy!(pointer(dest), pointer(src), m*n)
#     return dest
# end
# function copy!{T}(dest::Base.Matrix{T}, src::Matrix{T})
#     m, n = size(dest)
#     if m != size(src, 1) || n != size(src, 2)
#         throw(DimensionMisMatch("source and destination must have same shape"))
#     end
#     Base.unsafe_copy!(pointer(dest), pointer(src), m*n)
#     return dest
# end

function copy!{T}(dest::DistMatrix{T}, src::Base.VecOrMat)
    m, n = size(src, 1), size(src, 2)
    zeros!(dest, m, n)
    if MPI.commRank(comm(dest)) == 0
        for j = 1:n
            for i = 1:m
                queueUpdate(dest, i, j, src[i,j])
            end
        end
    end
    processQueues(dest)
    return dest
end

function convert{T}(::Type{Matrix{T}}, A::Base.VecOrMat{T})
    m, n = size(A, 1), size(A, 2)
    B = Matrix(T)
    resize!(B, m, n)
    Base.unsafe_copy!(pointer(B), pointer(A), m*n)
    return B
end
function convert{T}(::Type{Base.Matrix{T}}, A::Matrix{T})
    m, n = size(A)
    B = Base.Matrix{T}(m, n)
    Base.unsafe_copy!(pointer(B), pointer(A), m*n)
    return B
end

function convert{T}(::Type{DistMatrix{T}}, A::Base.VecOrMat{T})
    m, n = size(A, 1), size(A, 2)
    B = DistMatrix(T)
    zeros!(B, m, n)
    if MPI.commRank(comm(B)) == 0
        for j = 1:n
            for i = 1:m
                queueUpdate(B, i, j, A[i,j])
            end
        end
    end
    processQueues(B)
    return B
end

function convert{T}(::Type{DistMultiVec{T}}, A::Base.VecOrMat{T})
    m, n = size(A, 1), size(A, 2)
    B = DistMultiVec(T)
    zeros!(B, m, n)
    if MPI.commRank(comm(B)) == 0
        for j = 1:n
            for i = 1:m
                queueUpdate(B, i, j, A[i,j])
            end
        end
    end
    processQueues(B)
    return B
end

function convert{T}(::Type{DistMatrix{T}}, A::DistMultiVec{T})
    m, n = size(A)
    B = DistMatrix(T)
    zeros!(B, m, n)
    for j = 1:n
        for i = 1:localHeight(A)
            xij = getLocal(A, i, j)
            queueUpdate(B, globalRow(A, i), j, xij)
        end
    end
    processQueues(B)
    return B
end

function norm(x::ElementalMatrix)
    if size(x, 2) == 1
        return nrm2(x)
    else
        return twoNorm(x)
    end
end

# Multiplication
(*){T}(A::DistMatrix{T}, B::Base.VecOrMat{T}) = A*convert(DistMatrix{T}, B)
(*){T}(A::DistMultiVec{T}, B::Base.VecOrMat{T}) = convert(DistMatrix{T}, A)*convert(DistMatrix{T}, B)
(*){T}(A::DistSparseMatrix{T}, B::Base.VecOrMat{T}) = A*convert(DistMultiVec{T}, B)
Ac_mul_B{T}(A::DistMatrix{T}, B::Base.VecOrMat{T}) = Ac_mul_B(A, convert(DistMatrix{T}, B))
Ac_mul_B{T}(A::DistMultiVec{T}, B::Base.VecOrMat{T}) = Ac_mul_B(convert(DistMatrix{T}, A), convert(DistMatrix{T}, B))
Ac_mul_B{T}(A::DistSparseMatrix{T}, B::Base.VecOrMat{T}) = Ac_mul_B(A, convert(DistMultiVec{T}, B))


import ArchGDAL; const AG = ArchGDAL

@test AG.OF_ReadOnly | 0x04 == 0x04
@test 0x06 | AG.OF_ReadOnly == 0x06
@test AG.OF_ReadOnly | AG.OF_GNM == GDAL.GDAL_OF_READONLY | GDAL.GDAL_OF_GNM

for dt in (AG.GDT_Byte, AG.GDT_Float64, AG.GDT_Unknown, AG.GDT_Int16,
           AG.GDT_UInt16, AG.GDT_Float32, AG.GDT_Int32, AG.GDT_UInt32)
    print("size: $(AG.typesize(dt)), name: $(AG.typename(dt)), ")
    println("type: $(AG.gettype(AG.typename(dt)))")
end

@test AG.typeunion(AG.GDT_UInt16, AG.GDT_Byte) == AG.GDT_UInt16
@test AG.iscomplex(AG.GDT_Float32) == false

for name in (AG.GARIO_COMPLETE,AG.GARIO_ERROR,AG.GARIO_PENDING,AG.GARIO_UPDATE)
    println("$(AG.getname(name)) $(AG.asyncstatustype(AG.getname(name)))")
end

for color in (:Alpha,:Green,:Palette,:YCbCr,:Black,:Hue,:Red,:Blue,:Gray,
              :LightnessBand,:Saturation,:Yellow,:Cyan,:Magenta,:Undefined,
              :alpha,:black,:blackband,:hueband,:GrayIndex,:saturation)
    println("$color: $(AG.colorinterp(string(color)))")
end

for p in (AG.GPI_Gray, AG.GPI_RGB, AG.GPI_CMYK, AG.GPI_HLS)
    println("$p: $(AG.getname(p))")
end

for ft in (AG.OFTDate,AG.OFTStringList,AG.OFTWideString,AG.OFTBinary,
           AG.OFTString,AG.OFTIntegerList,AG.OFTRealList,AG.OFTInteger64,
           AG.OFTInteger,AG.OFTReal,AG.OFTWideStringList,AG.OFTInteger64List,
           AG.OFTDateTime,AG.OFTTime)
    println("$ft: $(AG.getname(ft))")
    for fst in (AG.OFSTNone, AG.OFSTBoolean, AG.OFSTInt16, AG.OFSTFloat32)
        print("  $fst: $(AG.getname(fst)), ")
        println("compatible: $(AG.arecompatible(ft,fst))")
    end
end


__precompile__(false)
module ElTopo

using Cxx
using Parameters
using GeometryTypes

include("../deps/deps.jl")

"""
    SurfTrack([...])

Constructs a SurfTrack instance which holds parameters for ElTopo surface tracker stabilization.
"""
@with_kw struct SurfTrack
    "Elements closer than this are considered \"near\" (or proximate)"
    proximity_epsilon::Float64 = 1e-4
    friction_coefficient::Float64 = 0.0
    min_triangle_area::Float64 = 1e-7

    " Collision epsilon to use during mesh improvment operations (i.e. if any mesh elements are closer than this, the operation is aborted). NOTE: This should be greater than collision_epsilon, to prevent improvement operations from moving elements into
     a collision configuration."
    improve_collision_epsilon::Float64 = 2e-6

    "Whether to set the min and max edge lengths as fractions of the initial average edge length"
    use_fraction::Bool = false

    # If use_fraction is true, these are taken to be fractions of the average edge length of the new surface. If use_fraction is false, these are absolute.
    min_edge_length::Float64 = 0.05
    max_edge_length::Float64 = 0.2
    max_volume_change::Float64 = 0.1

    "In-triangle minimum angle to enforce"
    min_triangle_angle::Float64 = 0
    "In-triangle maximum angle to enforce"
    max_triangle_angle::Float64 = 180

    use_curvature_when_splitting::Bool = false
    use_curvature_when_collapsing::Bool = false

    "Clamp curvature scaling with minimum bound"
    min_curvature_multiplier::Float64 = 1
    "Clamp curvature scaling with maximum bound"
    max_curvature_multiplier::Float64 = 1

    allow_vertex_movement::Bool = false

    "Minimum edge length improvement in order to flip an edge"
    edge_flip_min_length_change::Float64 = 0.05

    "Elements within this distance will trigger a merge attempt"
    merge_proximity_epsilon::Float64 = 1e-5

    "Whether to enforce collision-free surfaces (including during mesh maintenance operations)"
    collision_safety::Bool = true

    "Whether to allow changes in topology"
    allow_topology_changes::Bool = false

    "Wether to allow non-manifold (edges incident on more than two triangles)"
    allow_non_manifold::Bool = false

    "Whether to allow mesh improvement"
    perform_improvement::Bool = true
end

Libdl.dlopen(eltopo,Libdl.RTLD_GLOBAL)

addHeaderDir(dirname(@__FILE__)*"/../deps/usr/include/",kind=C_System)

cxxinclude("vector")
cxxinclude("subdivisionscheme.h")
cxxinclude("surftrack.h")

function stabilize(points,faces,p::SurfTrack)
    pointsout = Point{3,Float64}[]
    facesout = Face{3,UInt64}[]

    icxx"""
         std::vector<Vec3d> vs;
         std::vector<Vec3st> ts;
         std::vector<double> ms;

    $:(
    for i in 1:length(points)
        p1,p2,p3 = points[i]
        icxx"vs.push_back(Vec3d($p1,$p2,$p3));"
        icxx"ms.push_back(0.5);"
    end
    );

    $:(
    for i in 1:length(faces)
        v1,v2,v3 = faces[i] .- 1
        icxx"ts.push_back(Vec3st($v1,$v2,$v3));"
    end
    );

      SurfTrackInitializationParameters parameters;
      parameters.m_proximity_epsilon = $(p.proximity_epsilon);
      parameters.m_friction_coefficient = $(p.friction_coefficient);
      parameters.m_min_triangle_area = $(p.min_triangle_area);
      parameters.m_improve_collision_epsilon = $(p.improve_collision_epsilon);
      parameters.m_use_fraction = $(p.use_fraction);
      parameters.m_min_edge_length = $(p.min_edge_length);
      parameters.m_max_edge_length = $(p.max_edge_length);
      parameters.m_max_volume_change = $(p.max_volume_change);
      parameters.m_min_triangle_angle = $(p.min_triangle_angle);
      parameters.m_max_triangle_angle = $(p.max_triangle_angle);
      parameters.m_use_curvature_when_splitting = $(p.use_curvature_when_splitting);
      parameters.m_use_curvature_when_collapsing = $(p.use_curvature_when_collapsing);
      parameters.m_min_curvature_multiplier = $(p.min_curvature_multiplier);
      parameters.m_max_curvature_multiplier = $(p.max_curvature_multiplier);
      parameters.m_allow_vertex_movement = $(p.allow_vertex_movement);
      parameters.m_edge_flip_min_length_change = $(p.edge_flip_min_length_change);
      parameters.m_merge_proximity_epsilon = $(p.merge_proximity_epsilon);
      parameters.m_collision_safety = $(p.collision_safety);
      parameters.m_allow_topology_changes = $(p.allow_topology_changes);
      parameters.m_allow_non_manifold = $(p.allow_non_manifold);
      parameters.m_perform_improvement = $(p.perform_improvement);
      parameters.m_subdivision_scheme = new ButterflyScheme();

      SurfTrack surface_tracker(vs,ts,ms,parameters);

    surface_tracker.m_verbose = false;

    surface_tracker.improve_mesh();
    surface_tracker.topology_changes();
    surface_tracker.defrag_mesh();

    for ( int i = 0; i < surface_tracker.get_num_vertices(); ++i )
        $:(
           p1 = icxx"return surface_tracker.get_position(i)[0];";
           p2 = icxx"return surface_tracker.get_position(i)[1];";
           p3 = icxx"return surface_tracker.get_position(i)[2];";
           push!(pointsout,Point(p1,p2,p3));
        );

    for ( int i = 0; i < surface_tracker.m_mesh.num_triangles(); ++i )
        $:(
           v1 = icxx"return surface_tracker.m_mesh.get_triangle(i)[0];";
           v2 = icxx"return surface_tracker.m_mesh.get_triangle(i)[1];";
           v3 = icxx"return surface_tracker.m_mesh.get_triangle(i)[2];";
           push!(facesout,Face(v1+1,v2+1,v3+1));
        );
    """

    return pointsout,facesout
end

"""
    stabilize(msh::HomogenousMesh,p::SurfTrack)

Computes stabilized mesh for a given initial triangulation `msh` with stabilization parameters `p`.
"""
function stabilize(msh::HomogenousMesh,parameters::SurfTrack)
    newvertices, newfaces = stabilize(msh.vertices,msh.faces,parameters)
    return HomogenousMesh(newvertices,newfaces)
end

export stabilize, SurfTrack

end


module Missile_Types

using Aerospace

export
    Controls,
    SoftwareMsg,
    Aerodynamics,
    Missile

type Controls
    throttle::Float64
    alpha::Float64
    alphadot::Float64
    phi::Float64
    phidot::Float64
end

type SoftwareMsg
    time::Float64
    dt::Float64
    SWmode::UInt
    controls::Controls
    navstates::NavStates
    wypt_num::UInt
    axcmd::Float64
    aycmd::Float64
    azcmd::Float64
    eta_dive_cmd::Float64
    eta_turn_cmd::Float64
    tgo::Float64
    miss::Float64
    missvec_ned::Array{Float64,1}
    vel_err_int::Float64;
end

type Missile
    time::Float64
    dt::Float64
    airframe::Aero5DOF
    control::Controls
    software::SoftwareMsg
end

type Aerodynamics
    sref::Float64
    cl::Float64
    cd::Float64
    faero_body::Array{Float64,1}
end

end ## End of Module


export Segmenter

using Merlin

type Segmenter
    xdict::Dict{String,Int}
    ydict::Dict{String,Int}
    nn
end

function Segmenter()
    Segmenter(Dict{String,Int}(), Dict{String,Int}(), Model(4))
end

function (seg::Segmenter)(str::String)
    ids = map(collect(str)) do c
        id = seg.dict_x[string(c)]
        #id < 0 && throw("Invalid character: $(c).")
        id
    end
    x = reshape(ids, 1, length(ids))
    z = seg(Var(x))
    map(id -> getkey(seg.dict_y,id), z)
end

type Model
    g::Graph
end

function Model(outsize::Int)
    T = Float32
    x = Var([1,2])
    h = Lookup(T,300,10)(x)
    h = window(h, (70,), pads=(30,), strides=(10,))
    h = Linear(T,70,70)(h)
    h = tanh(h)
    h = Linear(T,70,outsize)(h)
    g = Graph([x], [h])
    Model(g)
end

function (m::Model)(x::Var, y=nothing)
    z = m.g(x)
    y == nothing ? argmax(z,1) : crossentropy(y,z)
end

function train(seg::Segmenter, trainpath::String, nepochs::Int)
    train_x, train_y = read(trainpath, seg.xdict, seg.ydict)
    n = floor(Int, length(train_x)*0.9)
    train_x, test_x = train_x[1:n], train_x[n+1:end]
    train_y, test_y = train_y[1:n], train_y[n+1:end]
    info("# training samples: $(length(train_x))")
    info("# test samples: $(length(test_x))")
    info("# x: $(length(seg.xdict))")
    info("# y: $(length(seg.ydict))")

    opt = SGD(0.0005)
    for epoch = 1:nepochs
        println("epoch: $epoch")
        loss = fit(train_x, train_y, seg.nn, opt)
        println("loss: $loss")

        ys = cat(1, map(y -> vec(y.data), test_y)...)
        zs = cat(1, map(x -> vec(seg.nn(x).data), test_x)...)
        acc = mean(i -> ys[i] == zs[i] ? 1.0 : 0.0, 1:length(ys))
        println("test accuracy: $acc")
        println()
    end
end

function read(path::String, xdict, ydict)
    data_x, data_y = Var[], Var[]
    buf_x, buf_y = Int[], Int[]
    lines = open(readlines, path)
    for i = 1:length(lines)
        line = chomp(lines[i])
        if isempty(line) || i == length(lines)
            push!(data_x, Var(reshape(buf_x,1,length(buf_x))))
            push!(data_y, Var(reshape(buf_y,1,length(buf_y))))
            buf_x = Int[]
            buf_y = Int[]
        else
            items = split(line, "\t")
            id = get!(xdict, items[1], length(xdict)+1)
            push!(buf_x, id)
            id = get!(ydict, items[2], length(ydict)+1)
            push!(buf_y, id)
        end
    end
    data_x, data_y
end


import Pkg
Pkg.activate(".")

using GLMakie, VideoIO, Makie, Dates, Printf, Colors, ImageTransformations, FFMPEG

mss = [Millisecond(25), Millisecond(125)] # intervals between frames for the 2 different FPS
fs = Int[] # number of frames per FPS
ts = [Millisecond(0)] # the timestamps
for (i, ms) in enumerate(mss)
    n = 1
    while last(ts) < Second(5i)
        push!(ts, last(ts) + ms)
        n += 1
    end
    push!(fs, n)
end
fs[1] -= 1 # a correction for how `mp4fpsmod` works.
dar = 2 # or PAR, the physical geometry of the pixels.
w = length(ts) # width of the image
h = dar*w # scaling the height
frame = Node(1)
point = lift(frame) do fn
    [(fn, dar*fn)]
end
msg = lift(frame, point) do fn, p
    ti = @sprintf "%03i" fn
    t = @sprintf "%-012s" Time(0) + ts[fn]
    x, y = p[]
    """
    Frame #: $ti
    Time: $t
    x: $x
    y: $y"""
end
limits = FRect3D(Vec3f0(1, 2, 0), Vec3f0(w - 1, 2w - 2, 0))
scene = Scene(resolution = (w, h), limits = limits,  scale_plot = false, show_axis = false)#, padding=(0,0), backgroundcolor=:red)
text!(scene, msg, position = (2, w), align = (:left,  :bottom))#, color = :white)
scatter!(scene, point, markersize = 10, marker = :+)#, color = :white)
imgstack = []
imgstack = map(1:w) do i
    frame[] = i
    AbstractPlotting.update_limits!(scene)
    AbstractPlotting.update!(scene)
    imresize(Gray.(GLMakie.scene2image(scene)), h, h)
end
props = [:color_range=>2, :priv_data => ("crf"=>"0","preset"=>"medium")]
y = "cfr.mp4"
encodevideo(y,imgstack,framerate=30,AVCodecContextProperties=props)
f1, f2 = fs
fps1, fps2 = Millisecond(Second(1)) .÷ mss
run(`mp4fpsmod -r $f1:$fps1  -r $f2:$fps2 -o vfr.mp4 cfr.mp4`)
ffmpeg_exe(` -y -i vfr.mp4 -c copy -bsf:v h264_metadata=sample_aspect_ratio=0.5 testvideo.mp4`)

f = VideoIO.openvideo("testvideo.mp4")
pixelaspectratio = VideoIO.aspect_ratio(f)
@assert pixelaspectratio == 0.5

f = VideoIO.openvideo("testvideo.mp4")
ts2 = Millisecond[]
while !eof(f)
    read(f)
    l = gettime(f)
    push!(ts2, Millisecond(round(Int, 1000l)))
    # push!(ts2, Millisecond(1))
end
@assert ts2 == ts

rm("cfr.mp4")
rm("vfr.mp4")


using SIMDAlignment
using Bio.Seq
using Bio.Align
using Base.Test

function make_submat{score_t}(::Type{score_t})
    submat = Matrix{score_t}(4, 4)
    fill!(submat, -6)
    submat[diagind(submat)] = 0
    return submat
end

function test_same_seqs{score_t}(::Type{score_t})
    submat = make_submat(score_t)
    gap_open = 5
    gap_extend = 3

    seq = seq_t([0x00, 0x01, 0x02, 0x03])
    for n in 1:100
        refs = seq_t[]
        for _ in 1:n
            push!(refs, [0x00, 0x01, 0x02, 0x03])
        end
        alns = paralign_score(submat, gap_open, gap_extend, seq, refs)
        for aln in alns
            @test score(aln) == 0
        end
    end
end

function test_empty_seq{score_t}(::Type{score_t})
    seq = dna""
    refs = [
        dna"A",
        dna"AC",
        dna""
    ]
    submat = make_submat(score_t)
    model = AffineGapScoreModel(submat, gap_open_penalty=5, gap_extend_penalty=3)

    refs′ = [convert(seq_t, ref) for ref in refs]
    alns = paralign_score(submat, model.gap_open_penalty, model.gap_extend_penalty, seq, refs′)
    for i in 1:length(refs)
        ref = refs[i]
        aln = alns[i]
        aln′ = pairalign(GlobalAlignment(), seq, ref, model)
        @test score(aln) == score(aln′)
    end
end

function test_various_seqs{score_t}(::Type{score_t})
    seq = dna"ACGT"
    refs = [
        dna"ACGT",
        dna"ACGG",
        dna"AACGT",
        dna"ACGTT",
        dna"ACATG",
        dna"AACGT",
        dna"AACGTG",
        dna"AACCTG",
        dna"AACCTGA",
        dna"ACG",
        dna"ACT",
        dna"AGT",
        dna"AGG",
        dna"CGT",
        dna"CCT",
        dna"CT",
        dna"C",
        dna"",
        dna"AACCTGA"[2:5],
        dna"AACCTGA"[2:6],
    ]

    submat = make_submat(score_t)
    model = AffineGapScoreModel(submat, gap_open_penalty=5, gap_extend_penalty=3)

    refs′ = [convert(seq_t, ref) for ref in refs]
    alns = paralign_score(submat, model.gap_open_penalty, model.gap_extend_penalty, seq, refs′)

    for i in 1:length(refs)
        ref = refs[i]
        aln = alns[i]
        aln′ = pairalign(GlobalAlignment(), seq, ref, model)
        @test score(aln) == score(aln′)
    end
end

function test_reversed_seqs{score_t}(::Type{score_t})
    seq = dna"ACGTAT"
    refs = [
        dna"TATGCA",
        dna"ACGTAT",
        dna"ATTGA",
        dna"",
    ]

    submat = make_submat(score_t)
    model = AffineGapScoreModel(submat, gap_open_penalty=5, gap_extend_penalty=3)
    refs′ = [seq_t(ref, true) for ref in refs]
    alns = paralign_score(submat, model.gap_open_penalty, model.gap_extend_penalty, seq, refs′)

    for i in 1:length(refs)
        ref = reverse(refs[i])
        aln = alns[i]
        aln′ = pairalign(GlobalAlignment(), seq, ref, model)
        @test score(aln) == score(aln′)
    end
end

# run tests
for score_t in (Int8, Int16, Int32)
    test_same_seqs(score_t)
    test_empty_seq(score_t)
    test_various_seqs(score_t)
end


# Reverse over forward - the compiler should get this itself, but this saves
# some work in the hot path.

struct ∂⃖composeOdd{C, N}
    a
    b
end

struct ∂⃖composeEven{C, N}
    a
    b
end

function (x::∂⃖composeOdd)(Δ)
    b, ∂b = x.b(Δ)
    a, ∂a = x.a(b[end])
    a, ∂⃖composeEven{N, plus1(N)}(∂a, ∂b)
end

function (x::∂⃖composeEven)(args...)
    a, ∂a = x.a(args...)
    b, ∂b = x.b(a)
    b, ∂⃖composeOdd{N, plus1(N)}(∂a, ∂b)
end

function (x::∂⃖composeOdd{N,N})(Δ) where {N}
    return x.a(x.b(Δ)[end])
end

function (this::∂⃖{N})(::∂☆internal{1}, args::AbstractTangentBundle{1}...) where {N}
    r, ∂r = this(my_frule, args...)
    if r === nothing
        # Forwards directly to the recursion, no need to ever call the
        # backwards for my_frule.
        return this(∂☆recurse{1}(), args...)
    else
        z, ∂z = this(shuffle_base, r)
        return z, ∂⃖composeOdd{1, c_order(N)}(∂r, ∂z)
    end
end

function shuffle_down_frule(∂☆p, my_frule, args...)
    ∂☆p(my_frule, map(shuffle_down, args)...)
end

function (this::∂⃖{N})(::∂☆internal{M}, args::AbstractTangentBundle{1}...) where {N, M}
    r = this(∂☆shuffle{N}(), args...)
    if primal(r) === nothing
        return this(∂☆recurse{N}(), args...)
    else
        z, ∂z = this(shuffle_up, r)
        return z, ∂⃖composeOdd{1, c_order(N)}(∂r, ∂z)
    end
end

#=
struct FwdIterate{N, T<:AbstractTangentBundle{N}}
    f::T
end
function (f::FwdIterate)(arg::ATB{N}) where {N}
    r = ∂☆{N}()(f.f, arg)
    primal(r) === nothing && return nothing
    (∂☆{N}()(ZeroBundle{N}(getindex), r, ZeroBundle{N}(1)),
     primal(∂☆{N}()(ZeroBundle{N}(getindex), r, ZeroBundle{N}(2))))
end
function (f::FwdIterate)(arg::ATB{N}, st) where {N}
    r = ∂☆{N}()(f.f, arg, ZeroBundle{N}(st))
    primal(r) === nothing && return nothing
    (∂☆{N}()(ZeroBundle{N}(getindex), r, ZeroBundle{N}(1)),
     primal(∂☆{N}()(ZeroBundle{N}(getindex), r, ZeroBundle{N}(2))))
end

function (this::∂☆{N})(::ZeroBundle{N, typeof(Core._apply_iterate)}, iterate::ATB{N}, f::ATB{N}, args::ATB{N}...) where {N}
    Core._apply_iterate(FwdIterate(iterate), this, (f,), args...)
end
=#

function (this::∂⃖{N})(that::∂☆{M}, ::ZeroBundle{M, typeof(Core._apply_iterate)},
        iterate, f, args::ATB{M, <:Tuple}...) where {N, M}
    @assert primal(iterate) === Base.iterate
    x, ∂⃖f = Core._apply_iterate(FwdIterate(iterate), this, (that, f), args...)
    return x, ApplyOdd{1, c_order(N)}(UnApply{map(x->length(primal(x)), args)}(), ∂⃖f)
end

function ChainRules.rrule(∂::∂☆{N}, m::ZeroBundle{N, typeof(map)}, p::ZeroBundle{N, typeof(+)}, A::ATB{N}, B::ATB{N}) where {N}
    ∂(m, p, A, B), Δ->(NoTangent(), NoTangent(), NoTangent(), Δ, Δ)
end

mapev_unbundled(_, js, a) = rebundle(mapev(js, unbundle(a)))
function (∂⃖ₙ::∂⃖{N})(∂☆ₘ::∂☆{M}, ::ZeroBundle{M, typeof(map)},
                    f::ZeroBundle{M}, a::ATB{M, <:Array}) where {N, M}
    @assert Base.issingletontype(typeof(primal(f)))
    js = map(primal(a)) do x
        ∂f = ∂☆{N+M}()(ZeroBundle{N+M}(primal(f)),
                     TaylorBundle{N+M}(x,
                       (one(x), (zero(x) for i = 1:(N+M-1))...,)))
        @assert isa(∂f, TaylorBundle) || isa(∂f, TangentBundle{1})
        Jet{typeof(x), N+M}(x, ∂f.primal,
            isa(∂f, TangentBundle) ? ∂f.partials : ∂f.coeffs)
    end
    ∂⃖ₙ(mapev_unbundled, ∂☆ₘ, js, a)
end


"""
Data structure used by open-loop layer to update current high-level plan.

Attributes:
    - `curr_graph::SimpleVListGraph{OpenLoopVertex{D,C,AD}}` A list of vertices generated so far by the bridge sampler
    - `curr_start_idx::Int64` The vertex index of the current start vertex in the graph, from which planning will happen. Updated each time open_loop_plan is called
    - `curr_goal_idx::Int64` The vertex index of current goal vertex. Updated each time a goal state vertex is popped.
    - `curr_soln_path_idxs::Vector{Int64}` The in-order list of indices from current start to current goal
    - `num_samples::Int64` The N parameter for the number of bridge samples to generate per mode switch
    - `bookkeeping::B` A data structure for tracking information pertaining to the graph search
"""
mutable struct GraphTracker{D,C,AD,M,B,RNG <: AbstractRNG}
    curr_graph::SimpleVListGraph{OpenLoopVertex{D,C,AD,M}}
    curr_start_idx::Int64
    curr_goal_idx::Int64
    curr_soln_path_idxs::Vector{Int64}
    num_samples::Int64
    bookkeeping::B
    rng::RNG
end

metadatatype(::GraphTracker{D,C,AD,M,B,RNG}) where {D,C,AD,M,B,RNG <: AbstractRNG} = M
bookkeepingtype(::GraphTracker{D,C,AD,M,B,RNG}) where {D,C,AD,M,B,RNG <: AbstractRNG} = B

function GraphTracker{D,C,AD,M,B,RNG}(N::Int64, bookkeeping::B, rng::RNG=Random.GLOBAL_RNG) where {D,C,AD,M,B,RNG <: AbstractRNG}
    return GraphTracker{D,C,AD,M,B,RNG}(SimpleVListGraph{OpenLoopVertex{D,C,AD,M}}(),
                        0, 0, Vector{Int64}(undef, 0), N, bookkeeping, rng)
end

"""
Open loop plan from current state from scratch. Return a graph tracker with solution

Arguments:
    - `cmssp::CMSSP{D,C,AD,AC}` The CMSSP instance
    - `s_t::CMSSPState{D,C}` The current state
    - `curr_time::Int64` The current time step value
    - `edge_weight::Function` A generic edge weight function (handed down by hhpc solver)
    - `heuristic::Function` A generic heuristic function (handed down by hhpc solver)
    - `goal_modes::Vector{D}` A set of modes that are goal conditions
    - `graph_tracker::GraphTracker{D,C}` The instance of the object that maintains the current open-loop graph
    - `start_metadata::M` Default metadata for a new start vertex

Returns:
    Updates the `graph_tracker` object in place.
"""
function open_loop_plan!(cmssp::CMSSP, s_t::CMSSPState,
                         curr_time::Int64,
                         edge_weight::Function,
                         heuristic::Function,
                         goal_modes::Vector{D},
                         graph_tracker::GraphTracker{D,C,AD,M,B,RNG},
                         start_metadata::M) where {D,C,AD,M,B,RNG <: AbstractRNG}

    # First update the context
    update_vertices_with_context!(cmssp, graph_tracker, curr_time)

    # Create start vertex and insert in graph
    # Don't need explicit goal - visitor will handle
    add_vertex!(graph_tracker.curr_graph, OpenLoopVertex{D,C,AD,M}(s_t, cmssp.mode_actions[1], TPDistribution([curr_time], [1.0]), start_metadata))
    graph_tracker.curr_start_idx = num_vertices(graph_tracker.curr_graph)
    graph_tracker.curr_goal_idx = 0

    # Clear out current solution
    empty!(graph_tracker.curr_soln_path_idxs)

    # Obtain path and current cost with A*
    a_star_path_soln = a_star_implicit_shortest_path!(graph_tracker.curr_graph, edge_weight,
                                                      graph_tracker.curr_start_idx,
                                                      GoalVisitorImplicit(graph_tracker, cmssp, goal_modes))

    # If unreachable, report warning
    if graph_tracker.curr_goal_idx == 0
        @warn "A goal state is currently unreachable"
        return
    end

    # Check that goal state is indeed terminal
    # And that A* has non-Inf cost path to it
    @assert a_star_path_soln.dists[graph_tracker.curr_goal_idx] < Inf "Path to goal is Inf cost!"

    ## Walk path back to goal
    # Insert goal in soln vector
    pushfirst!(graph_tracker.curr_soln_path_idxs, graph_tracker.curr_goal_idx)
    curr_vertex_idx = graph_tracker.curr_goal_idx

    while curr_vertex_idx != graph_tracker.curr_start_idx
        prev_vertex_idx = a_star_path_soln.parent_indices[curr_vertex_idx]
        pushfirst!(graph_tracker.curr_soln_path_idxs, prev_vertex_idx)
        curr_vertex_idx = prev_vertex_idx
    end
    # @show graph_tracker.curr_soln_path_idxs
    # for idx in graph_tracker.curr_soln_path_idxs
    #     println(graph_tracker.curr_graph.vertices[idx].state," at ",graph_tracker.curr_graph.vertices[idx].tp)
    #     hor_val = graph_tracker.curr_graph.vertices[idx].tp.vals[1]
    #     if hor_val < typemax(Int64)
    #         display_context_future(context_set,hor_val)
    #     end
    # end
end

"""
The visitor object for the implicit A* search. Attributes are obvious.
"""
struct GoalVisitorImplicit <: AbstractDijkstraVisitor
    graph_tracker::GraphTracker
    cmssp::CMSSP
    goal_modes::Vector
end

function Graphs.include_vertex!(vis::GoalVisitorImplicit,
                                u::OpenLoopVertex, v::OpenLoopVertex, d::Float64, nbrs::Vector{Int64})

    # @show v
    # If popped vertex is terminal, then stop
    if isterminal(vis.cmssp, v.state) == true
        vis.graph_tracker.curr_goal_idx = vertex_index(vis.graph_tracker.curr_graph, v)
        return false
    end

    popped_mode = v.state.mode
    popped_cont = v.state.continuous

    # If goal mode but NOT goal state, add samples from goal
    if popped_mode in vis.goal_modes

        # Generate goal sample set
        (vertices_to_add, nbrs_to_add) = generate_goal_vertex_set!(vis.cmssp, v, vis.graph_tracker, vis.graph_tracker.rng)

        for nbr_idx in nbrs_to_add
            push!(nbrs, nbr_idx)
        end

        num_new_vertices = length(vertices_to_add)

        if num_new_vertices > 0

            # Add vertices to graph and to nbrs
            for vtx in vertices_to_add
                add_vertex!(vis.graph_tracker.curr_graph, vtx)
                push!(nbrs, num_vertices(vis.graph_tracker.curr_graph))
            end
        end
    end

    # Now add for next modes
    next_valid_modes = generate_next_valid_modes(vis.cmssp, popped_mode)

    # @show next_valid_modes

    for (action, nvm) in next_valid_modes

        # Generate bridge samples
        (vertices_to_add, nbrs_to_add) = generate_bridge_vertex_set!(vis.cmssp, v,
                                                                    (popped_mode, nvm),
                                                                    vis.graph_tracker,
                                                                    action,
                                                                    vis.graph_tracker.rng)
        for nbr_idx in nbrs_to_add
            push!(nbrs, nbr_idx)
        end

        num_new_vertices = length(vertices_to_add)
        # @show num_new_vertices

        if num_new_vertices > 0

            # Add vertices to graph and to nbrs
            for vtx in vertices_to_add
                add_vertex!(vis.graph_tracker.curr_graph, vtx)
                push!(nbrs, num_vertices(vis.graph_tracker.curr_graph))
            end
        end
    end

    # @show v
    # @show nbrs
    # readline()
    return true
end


#!/usr/bin/env julia

using DataFrames

import Celeste.AccuracyBenchmark
import Celeste.ArgumentParse
import Celeste.Model
import Celeste.Synthetic

const PSF_SIGMA_PX = 2.29 # similar to SDSS
const BAND_SKY_LEVEL_NMGY = [0.2696, 0.3425, 0.7748, 1.6903, 4.9176]
const BAND_NELEC_PER_NMGY = [146.9, 838.1, 829.8, 597.2, 129.8]
const OUTPUT_DIRECTORY = joinpath(splitdir(Base.source_path())[1], "output")

parser = Celeste.ArgumentParse.ArgumentParser()
ArgumentParse.add_argument(
    parser,
    "catalog_csv",
    help="Ground truth CSV catalog",
)
parsed_args = ArgumentParse.parse_args(parser, ARGS)

srand(12345)

catalog_data = AccuracyBenchmark.read_catalog(parsed_args["catalog_csv"])
catalog_entries = [
    AccuracyBenchmark.make_catalog_entry(row)
    for row in eachrow(catalog_data)
]
template_images = AccuracyBenchmark.make_template_images(
    catalog_data,
    PSF_SIGMA_PX,
    BAND_SKY_LEVEL_NMGY,
    BAND_NELEC_PER_NMGY,
)
generated_images = Synthetic.gen_blob(template_images, catalog_entries)

catalog_label = splitext(basename(parsed_args["catalog_csv"]))[1]
output_filename = joinpath(OUTPUT_DIRECTORY, @sprintf("%s_synthetic.fits", catalog_label))
AccuracyBenchmark.save_images_to_fits(output_filename, generated_images)
AccuracyBenchmark.append_hash_to_file(output_filename)