using NearestNeighbors: TreeData, get_leaf_range, isleaf, getleft, getright

const BARNES_HUT_LEAFSIZE = 16
const BARNES_HUT_THETA = 1/4

# could generalize to something like HierarchicalFactorization
# have field that determines algorithm to execute Taylor(p), BarnesHut(), FKT(p)
struct BarnesHutFactorization{T, KT, XT<:AbstractVector, YT<:AbstractVector, TT, DT, RT} <: Factorization{T}
    k::KT # kernel function, for taylor branch, needs to be evaluable as k(r²)
    x::XT
    y::YT
    Tree::TT
    D::DT # diagonal correction (i.e. noise variance for GPs)
    θ::RT # distance threshold for compression
    # expansion_type::ET # IDEA to generalize tree-based algorithms: problem: would be nice to pre-allocate moments
    # w::WT # temporary storage for weight vectors, useful to split positive and negative values
    # i::BT
    # centers_of_mass::CT # pre-allocated, has to be able to work with type of w
end

# NOTE on ball tree indices
# if reorder = true, and T = BallTree(X), we have
# X[:, T.indices[i]] = T.data[i]
# X[:, i] == T.data[invpermute!(collect(1:length(T.data)), T.indices)][i]
function BarnesHutFactorization(k, x, y = x, D = nothing; θ::Real = BARNES_HUT_THETA, leafsize::Int = BARNES_HUT_LEAFSIZE)
    xs = vector_of_static_vectors(x)
    ys = x === y ? xs : vector_of_static_vectors(y)
    Tree = BallTree(ys, leafsize = leafsize)
    m = length(y)
    XT, YT, KT, TT, DT, RT = typeof.((xs, ys, k, Tree, D, θ))
    # w = zeros(length(m))
    # i = zeros(Bool, m)
    # WT, BT = typeof(w), typeof(i)
    T = gramian_eltype(k, xs[1], ys[1])
    BarnesHutFactorization{T, KT, XT, YT, TT, DT, RT}(k, xs, ys, Tree, D, θ) #, w, i)
end
function BarnesHutFactorization(G::Gramian; θ::Real = BARNES_HUT_THETA, leafsize::Int = BARNES_HUT_LEAFSIZE)
    BarnesHutFactorization(G.k, G.x, G.y, θ = θ, leafsize = leafsize)
end
Base.size(F::BarnesHutFactorization) = (length(F.x), length(F.y))
Base.size(F::BarnesHutFactorization, i::Int) = i ≤ 2 ? size(F)[i] : 1
Base.getindex(F::BarnesHutFactorization, i::Int, j::Int) = F.k(F.x[i], F.y[j])
Base.eltype(F::BarnesHutFactorization{T}) where T = T

function LinearAlgebra.mul!(y::AbstractVector, F::BarnesHutFactorization, x::AbstractVector, α::Real = 1, β::Real = 0)
    # barneshut!(y, F, x, α, β)
    if all(≥(0), x) # if all weights are positive, proceed with regular BarnesHut
        barneshut!(y, F, x, α, β, split = false)
    else
        taylor!(y, F, x, α, β)
    end
end
function Base.:*(F::BarnesHutFactorization{<:Number}, x::AbstractVector{<:Number})
    T = promote_type(eltype(F), eltype(x))
    y = zeros(T, size(F, 1))
    mul!(y, F, x)
end

# # use cg! if it's positive definite
# keywords include maxiter = size(F, 2)
# verbose: convergence info
# log: keep track of residual norm
# Pl, Pr: left and right preconditioners (not implemented?)
function LinearAlgebra.ldiv!(x::AbstractVector, F::BarnesHutFactorization, b::AbstractVector; kwargs...)
    # cg!(x, F, b; kwargs...)
    minres!(x, F, b; kwargs...)
end
function Base.:\(F::BarnesHutFactorization, b::AbstractVector; kwargs...)
    T = promote_type(eltype(F), eltype(b))
    x = zeros(T, size(F, 2))
    ldiv!(x, F, b; initially_zero = true, kwargs...)
end

############################ Core Barnes Hut algorithm ##########################
# uses the Barnes-Hut algorithm for an approximate n*log(n) multiply with the kernel matrix G
function barneshut!(b::AbstractVector, F::BarnesHutFactorization{<:Number}, w::AbstractVector,
                    α::Number = 1, β::Number = 0, θ::Real = F.θ; split::Bool = true)
    size(F, 2) == length(w) || throw(DimensionMismatch("length of w does not match second dimension of F: $(length(w)) ≠ $(size(F, 2))"))
    eltype(b) == promote_type(eltype(F), eltype(w)) ||
            throw(TypeError("eltype of target vector b not equal to eltype of matrix-vector product: $(eltype(b)) and $(promote_type(eltype(F), eltype(w)))"))

    @. b = iszero(β) ? 0 : β * b # this avoids trouble if b is initialized with NaN's, e.g. thorugh "similar"
    if split && any(<(0), w) # if split is on, we multiply with positive and negative component of w separately, tends to increase accuracy because of center of mass calculation
        splitting_barneshut!(b, F, w, α, β, θ)
    else
        sums_w = node_sums(w, F.Tree) # IDEA: could pre-allocate
        centers_of_mass = compute_centers_of_mass(F, w)
        @threads for i in eachindex(F.x) # exactly 4 * length(y) allocations?
            Fw_i = bh_recursion(1, F.k, F.x[i], F.y, w, sums_w, θ, F.Tree, centers_of_mass)::eltype(b) # x[i] creates an allocation
            b[i] += α * Fw_i
        end
        if !isnothing(F.D) # if there is a diagonal correction, need to add it
            mul!(b, F.D, w, α, 1)
        end
    end
    return b
end

# splits weight vector a into positive and negative components, tends to have
# much better accuracy for negative a, for only ~2x time penalty
function splitting_barneshut!(b::AbstractVector, F::BarnesHutFactorization, a::AbstractVector,
                    α::Number = 1, β::Number = 0, θ::Real = F.θ)
    c = copy(a) # IDEA: pre-allocate
    i = a .< 0 # first, blend out negative indices (multiply with positive part)
    @. c[i] = 0
    barneshut!(b, F, c, α, β, θ, split = false)
    @. c = -a # multiply with -a[i] to make entries positive, and use negative α to make result correct
    @. i = !i # blend out non-negative indices
    @. c[i] = 0
    barneshut!(b, F, c, -α, 1, θ, split = false) # use β = 1 because β is already taken care of above
    return b
end

# supposing y is a single value now
# k sould be isotropic kernel with euclidean norm
# x is single target data point
# y are all source data points
# sums_a[index] contains sum(a[get_node_indices(T.tree_data, index)])
# θ < 1 is the distance threshold beyond which compression is used
# e.g. when θ = 1/10, then we compress when the current target node is 10 times further
# from the center of the hypersphere than its radius
# when θ = 1, everything is compressed, when θ = 0, nothing is
function bh_recursion(index, k, xi, y::AbstractVector, w::AbstractVector,
                      sums_w::AbstractVector, θ::Real, T::BallTree, centers_of_mass)
    h = T.hyper_spheres[index]
    if isleaf(T.tree_data.n_internal_nodes, index) # do direct computation
        elty = promote_type(eltype(xi), eltype(eltype(y)), eltype(w))
        val = zero(elty)
        @inbounds @simd for i in get_leaf_range(T.tree_data, index) # loop vectorization?
            j = T.indices[i]
            val += k(xi, y[j]) * w[j]
        end
        return val

    elseif h.r < θ * euclidean(xi, centers_of_mass[index]) # compress
        return k(xi, centers_of_mass[index]) * sums_w[index]

    else # recurse NOTE: parallelizing here is not as efficient as parallelizing over target points
        l = bh_recursion(getleft(index), k, xi, y, w, sums_w, θ, T, centers_of_mass)
        r = bh_recursion(getright(index), k, xi, y, w, sums_w, θ, T, centers_of_mass)
        return l + r
    end
end

############################# centers of mass ##################################
# this is a weighted sum, could be generalized to incorporate node_sums
function compute_centers_of_mass(w::AbstractVector, x::AbstractVector, T::BallTree)
    D = eltype(x) <: StaticVector ? length(eltype(x)) : length(x[1]) # if x is static vector
    com = [zero(MVector{D, Float64}) for _ in 1:length(T.hyper_spheres)]
    compute_centers_of_mass!(com, w, x, T)
end

function compute_centers_of_mass(F::BarnesHutFactorization, w::AbstractVector)
    compute_centers_of_mass(w, F.y, F.Tree)
end

function compute_centers_of_mass!(com::AbstractVector, w::AbstractVector, x::AbstractVector, T::BallTree)
    abs_w = abs.(w)
    weighted_node_sums!(com, abs_w, x, T)
    sum_w = node_sums(abs_w, T)
    ε = eps(eltype(w)) # ensuring division by zero it not a problem
    @. com ./= sum_w + ε
end

node_sums(x::AbstractVector, T::BallTree) = weighted_node_sums(Ones(length(x)), x, T)
function node_sums!(sums, x::AbstractVector, T::BallTree)
    weighted_node_sums!(sums, Ones(length(x)), x, T)
end

function weighted_node_sums(w::AbstractVector, x::AbstractVector, T::BallTree, index::Int = 1)
    length(x) == 0 && return zero(eltype(x))
    sums = fill(zero(w[1]'x[1]), length(T.hyper_spheres))
    weighted_node_sums!(sums, w, x, T)
end

# NOTE: x should either be vector of numbers or vector of static arrays
function weighted_node_sums!(sums::AbstractVector, w::AbstractVector,
                            x::AbstractVector, T::BallTree, index::Int = 1)
    if isleaf(T.tree_data.n_internal_nodes, index)
        i = get_leaf_range(T.tree_data, index)
        wi, xi = @views w[T.indices[i]], x[T.indices[i]]
        sums[index] = wi'xi
    else
        task = @spawn weighted_node_sums!(sums, w, x, T, getleft(index))
        weighted_node_sums!(sums, w, x, T, getright(index))
        wait(task)
        sums[index] = sums[getleft(index)] + sums[getright(index)]
    end
    return sums
end

# get centers used for compression in taylor factorization
function get_hyper_centers(F::BarnesHutFactorization)
    [hyp.center for hyp in F.Tree.hyper_spheres] # for uniform weights, doesn't loose accuracy in doing this
end


# utility functions for the CUSPARSE wrappers

"""
check that the dimensions of matrix `X` and vector `Y` make sense for a multiplication
"""
function chkmvdims(X, n, Y, m)
    if length(X) != n
        throw(DimensionMismatch("X must have length $n, but has length $(length(X))"))
    elseif length(Y) != m
        throw(DimensionMismatch("Y must have length $m, but has length $(length(Y))"))
    end
end

"""
check that the dimensions of matrices `X` and `Y` make sense for a multiplication
"""
function chkmmdims( B, C, k, l, m, n )
    if size(B) != (k,l)
        throw(DimensionMismatch("B has dimensions $(size(B)) but needs ($k,$l)"))
    elseif size(C) != (m,n)
        throw(DimensionMismatch("C has dimensions $(size(C)) but needs ($m,$n)"))
    end
end


# MIT License
# Copyright (c) 2019: Antonio Saragga Seabra

# Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.

#THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

##################
# Settlement Cycle
##################

"""
    settle(transactiondate::Date, settlementperiod, :HolydayCalendar)

Computes the settlement day using:
- Transaction date
- Settlement cycle period measured in business days (e.g. T+3, or more simply just 3)
- Business days calendars (e.g. :USNYSE, :TARGET,...)
"""
function liqd(transdate::Dates.Date, settcycle::Int, calendar)
    return advancebdays(convert(HolidayCalendar, calendar), transdate, settcycle)
end

function liqd(transdate::Dates.Date, settcycle::String, calendar)
    return advancebdays(convert(HolidayCalendar, calendar), transdate, parse(Int,settcycle[3]))
end

###################################################################
# Additional Days Count Convention
# Note: in DayCounts.jl 'const ActualActualISMA = ActualActualICMA'
###################################################################

"""
    ActualActualISMA()

**Actual/Actual ISMA** day count convention (e.g. London Stock Exchange, NASDAQ).

The year fraction is computed as:
```math
\\frac{\\text{# days between previous coupon and settlement date}}{\\text{# days in coupon period} \\times \\text{# coupon periods in a year}}
```

# Reference
- Information for private investors, January 2015, published by the London Stock Exchange Group
  https://www.londonstockexchange.com/traders-and-brokers/security-types/retail-bonds/accrued-interest-corp-supra.pdf
- Guide to Calculation Methods for the FTSE Fixed Income Indexes, v1.8
  https://www.ftse.com/products/downloads/FTSE_Fixed_Income_Index_Guide_to_Calculation_new.pdf
"""
struct ActualActualISMA <: DayCount end
function frac(startdate::Dates.Date, enddate::Dates.Date, nextcoupondate::Dates.Date, ::ActualActualISMA)
    daysbetweencoupons = Dates.value(nextcoupondate-startdate)
    if startdate > enddate || enddate > nextcoupondate
        error("start date > end date, or enddate > nextcoupondate")
    end
    return Dates.value(enddate-startdate)/(daysbetweencoupons*round(365/daysbetweencoupons))
end

###########################
# Business Days Conventions
###########################

abstract type BusinessDayConvention end
"""
    AllBD

Cashflows can land on weekends or holidays.
"""
struct AllBD <: BusinessDayConvention end
function utild(transdate::Dates.Date, ::Type{AllBD})
    return transdate
end
function utild(transdate::Dates.Date, ::Type{AllBD}, calendar)
    return transdate
end

# (A.1) if "Following" is specified, that date will be the first following day that is a Business Day;
# With following, cashflows are adjusted to land on the next good business day.
"""
    Following

Timings of cashflows that fall on a non-business day are adjusted to be on the the first **following** day that is a business day
"""
struct Following <: BusinessDayConvention end
function utild(transdate::Dates.Date,::Type{Following}, calendar)
    return tobday(convert(HolidayCalendar, calendar), transdate; forward = true)
end

# (A.2) if "Modified Following" or "Modified" is specified, that date will be the first following day that is a Business Day unless that day falls in the next calendar month, in which case that date will be the first preceding day that is a Business Day
# This is probably the most common business day convention used.  Cash flows that fall on weekends or holidays are assumed to be distributed on the following business day. However if the following business day is in a different month, the previous business day is adopted instead.
"""
    ModifiedFollowing

Timings of cashflows that fall on a non-business day are adjusted to be on the first **following** day that is a business day unless that day falls in the next calendar month, in which case that date will be the first **preceding** day that is a business day.
"""
struct ModifiedFollowing <: BusinessDayConvention end
function utild(transdate::Dates.Date,::Type{ModifiedFollowing}, calendar)
    maybe = tobday(convert(HolidayCalendar, calendar), transdate; forward = true)
    if month(maybe) == month(transdate)
        return maybe
    else
        return tobday(convert(HolidayCalendar, calendar), transdate; forward = false)
    end
end

# (A.3) if "Preceding" is specified, that date will be the first preceding day that is a Business Day.
# Cash flows that fall on a non-business day are assumed to be distributed on the previous business day.
"""
    Preceding

Timings of cashflows that fall on a non-business day are adjusted to be on the first **preceding** day that is a business day.
"""
struct Preceding <: BusinessDayConvention end
function utild(transdate::Dates.Date,::Type{Preceding}, calendar)
    return tobday(convert(HolidayCalendar, calendar), transdate; forward = false)
end

# (5) Modified Preceding
# Cash flows that fall on a non-business day are assumed to be distributed on the previous business day. However if the previous business day is in a different month, the following business day is adopted instead.
struct ModifiedPreceding <: BusinessDayConvention end
function utild(transdate::Dates.Date,::Type{ModifiedPreceding}, calendar)
    maybe = tobday(convert(HolidayCalendar, calendar), transdate; forward = false)
    if month(maybe) == month(transdate)
        return maybe
    else
        return tobday(convert(HolidayCalendar, calendar), transdate; forward = true)
    end
end

# End of Month - unadjusted
# All cash flows are assumed to be distributed on the final day of the month, even if it is a non-business day.
"""
    EndOfMonth

All cash flows are assumed to be distributed on the **final day of the month**, even if it is a non-business day.
"""
struct EndOfMonth <: BusinessDayConvention end
function utild(transdate::Dates.Date, ::Type{EndOfMonth})
    return Dates.lastdayofmonth(transdate)
end

# End of Month - Following
# All cash flows are assumed to be distributed on the final day of the month. If the final day of the month is a non-business day, the following business day is adopted.

# End of Month - Preceding
# All cash flows are assumed to be distributed on the final day of the month. If the final day of the month is a non-business day, the previous business day is adopted.

"""
    bdarray(A, convention, calendar)
    bdarray(A, convention, calendar, subsetcolumns)
        A: nxk Array
        convention: BusinessDayConvention
        calender: HolydaysCalender
        subsetcolumns: number of the columns to be processed e.g. [1,4] - first and fouth columns processed, [1] - only first column processed. If omited, all columns will be processed.

Replace holidays and weekends in array A with business days according with a choosen business day convention and calander. The array A may have elements whose type is not Date (which will be skipped).

"""
function bdarray(A, convention, calendar)
    nlines = size(A,1)
    ncolumns = size(A,2)
    for j = 1:ncolumns
        for i = 1:nlines
            if typeof(A[i,j]) == Date
                A[i,j] = utild(A[i,j], convention, calendar)
            end
        end
    end
    return A
end

function bdarray(A, convention, calendar, subsetcolumns)
    nlines = size(A,1)
    for j in subsetcolumns
        for i = 1:nlines
            if typeof(A[i,j]) == Date
                A[i,j] = utild(A[i,j], convention, calendar)
            end
        end
    end
    return A
end


@testset "single op" begin
    @test_throws MethodError augment!(rand(2,2), (rect,), Rotate90())
    @test_throws MethodError augment!((rand(2,2),), rect, Rotate90())
    @test_throws MethodError augment!((rand(2,2),rand(2,2)), (rect,), Rotate90())
    @test_throws BoundsError augment!(rand(2,2), rect, Rotate90())
    for pl in (Augmentor.ImmutablePipeline(Rotate90()), (Rotate90(),))
        img = @inferred Augmentor._augment(rect, pl)
        @test typeof(img) <: Array
        @test img == @inferred(augment(rect, Rotate90()))
        @test img == @inferred(augment(rect, (Rotate90(),)))
        @test eltype(img) <: eltype(rect)
        @test img == rotl90(rect)
        img1, img2 = @inferred augment((square2, rect), pl)
        @test img1 == rotl90(square2)
        @test img2 == img
        @test typeof(img2) == typeof(img)
        img3 = @inferred augment(rect, Rotate90())
        @test img2 == img3
        @test typeof(img3) == typeof(img2)
        out = similar(img)
        @test @inferred(augment!(out, rect, pl)) == img
        outs = (similar(square2), similar(img))
        res = @inferred(augment!(outs, (square2, rect), pl))
        @test res === outs
        @test outs[1] == img1
        @test outs[2] == img2
        out = similar(img)
        @test @inferred(augment!(out, rect, Rotate90())) == img
        @test_throws BoundsError augment!(rand(2,2), rect, pl)
    end
end

ops = Augmentor.ImmutablePipeline(Rotate(90),Rotate(-90)) # forces affine
@testset "$(str_showcompact(ops))" begin
    wv = @inferred Augmentor._augment(camera, ops)
    @test typeof(wv) === typeof(invwarpedview(camera, Augmentor.toaffinemap(NoOp(),rect), Flat()))
    @test wv == camera
    wv1, wv2 = @inferred Augmentor._augment((rgb_rect,camera), ops)
    @test typeof(wv1) === typeof(invwarpedview(rgb_rect, Augmentor.toaffinemap(NoOp(),rect), Flat()))
    @test typeof(wv2) === typeof(invwarpedview(camera, Augmentor.toaffinemap(NoOp(),rect), Flat()))
    @test wv1 == rgb_rect
    @test wv2 == camera
end

ops = (CacheImage(),CacheImage()) # forces eager
@testset "$(str_showcompact(ops))" begin
    img = @inferred Augmentor._augment(camera, ops)
    @test img === camera
    img1, img2 = @inferred Augmentor._augment((camera, square2), ops)
    @test img1 === camera
    @test img2 === square2
end

ops = (ElasticDistortion(4,4),CacheImage()) # forces lazy then eager
@testset "$(str_showcompact(ops))" begin
    img = @inferred Augmentor._augment(camera, ops)
    @test size(img) == size(camera)
    @test typeof(img) == typeof(camera)
    img1, img2 = @inferred Augmentor._augment((camera, N0f8.(camera)), ops)
    @test img1 isa Array{Gray{N0f8}}
    @test img2 isa Array{N0f8}
    @test img1 == img2
end

ops = (ShearX(45),ShearX(-45)) # forces affine
@testset "$(str_showcompact(ops))" begin
    wv = @inferred Augmentor._augment(camera, ops)
    @test typeof(wv) === typeof(invwarpedview(rect, Augmentor.toaffinemap(NoOp(),rect), Flat()))
    @test view(wv,1:512,1:512) == camera
end

ops = (ShearY(45),ShearY(-45)) # forces affine
@testset "$(str_showcompact(ops))" begin
    wv = @inferred Augmentor._augment(camera, ops)
    @test typeof(wv) === typeof(invwarpedview(rect, Augmentor.toaffinemap(NoOp(),rect), Flat()))
    @test view(wv,1:512,1:512) == camera
end

ops = (ShearY(45),ShearX(-2),CacheImage()) # forces affine then eager
@testset "$(str_showcompact(ops))" begin
    img = @inferred Augmentor._augment(camera, ops)
    @test typeof(img) <: OffsetArray
    @test indices(img) == (-255:768, 0:512)
end

ops = (Resize(2,2),Rotate90()) # forces affine
@testset "$(str_showcompact(ops))" begin
    wv = @inferred Augmentor._augment(rect, ops)
    @test typeof(wv) <: SubArray
    @test typeof(wv.indexes) <: Tuple{Vararg{IdentityRange}}
    @test typeof(parent(wv)) <: InvWarpedView
    @test parent(parent(wv)).itp.coefs === rect
    @test round.(Float64.(wv),1) == round.(Float64.(rotl90(imresize(rect,2,2))),1)
end

ops = (Resize(2,2),Rotate90(),CacheImage()) # forces affine then eager
@testset "$(str_showcompact(ops))" begin
    img = @inferred Augmentor._augment(rect, ops)
    @test typeof(img) <: OffsetArray
    @test round.(Float64.(img),1) == round.(Float64.(rotl90(imresize(rect,2,2))),1)
end

buf = rand(Gray{N0f8}, 2, 2)
ops = (Resize(2,2),Rotate90(),CacheImage(buf)) # forces affine then eager
@testset "$(str_showcompact(ops))" begin
    img = @inferred Augmentor._augment(rect, ops)
    @test typeof(img) <: OffsetArray
    @test round.(Float64.(img),1) == round.(Float64.(rotl90(imresize(rect,2,2))),1)
    @test img == ops[3].buffer
    @test parent(img) === ops[3].buffer
end

ops = (Rotate180(),Crop(5:200,200:500),Rotate90(1),Crop(1:250, 1:150))
@testset "$(str_showcompact(ops))" begin
    wv = @inferred Augmentor._augment(camera, ops)
    @test typeof(wv) <: SubArray
    @test typeof(wv.indexes) <: Tuple{Vararg{IdentityRange}}
    @test typeof(parent(wv)) <: InvWarpedView
    @test parent(parent(wv)).itp.coefs === camera
    @test_reference "reference/rot_crop_either_crop.txt" wv
    img = @inferred augment(camera, ops)
    @test img == parent(copy(wv))
    @test typeof(img) <: Array
    @test eltype(img) <: eltype(camera)
    @test_reference "reference/rot_crop_either_crop.txt" img
    out = similar(img)
    @test @inferred(augment!(out, camera, ops)) === out
    @test_reference "reference/rot_crop_either_crop.txt" out
    @test @allocated(augment!(out, camera, ops)) < @allocated(augment(camera, ops))
end

ops = Augmentor.ImmutablePipeline(Rotate180(),Crop(5:200,200:500),Rotate90(),Crop(50:300, 50:195))
@testset "$(str_showcompact(ops))" begin
    wv = @inferred Augmentor._augment(camera, ops)
    @test typeof(wv) <: SubArray
    @test typeof(parent(wv)) <: Base.PermutedDimsArrays.PermutedDimsArray
    @test_reference "reference/rot_crop_rot_crop.txt" wv
    img = @inferred augment(camera, ops)
    @test img == parent(copy(wv))
    @test typeof(img) <: Array
    @test eltype(img) <: eltype(camera)
    @test_reference "reference/rot_crop_rot_crop.txt" img
    out = similar(img)
    @test @inferred(augment!(out, camera, ops)) === out
    @test_reference "reference/rot_crop_rot_crop.txt" out
    @test @allocated(augment!(out, camera, ops)) < @allocated(augment(camera, ops))
end

ops = (Rotate180(),Crop(5:200,200:500),Rotate90(),Crop(50:300, 50:195),Resize(25,15))
@testset "$(str_showcompact(ops))" begin
    wv = @inferred Augmentor._augment(camera, ops)
    @test typeof(wv) <: SubArray
    @test eltype(wv) <: eltype(camera)
    @test typeof(wv.indexes) <: Tuple{Vararg{IdentityRange}}
    @test typeof(parent(wv)) <: InvWarpedView
    @test parent(parent(wv)).itp.coefs === camera
    @test_reference "reference/rot_crop_rot_crop_resize.txt" wv
    img = @inferred Augmentor.augment(camera, ops)
    @test img == parent(copy(wv))
    @test typeof(img) <: Array
    @test eltype(img) <: eltype(camera)
    @test_reference "reference/rot_crop_rot_crop_resize.txt" img
    out = similar(img)
    @test @inferred(augment!(out, camera, ops)) === out
    @test_reference "reference/rot_crop_rot_crop_resize.txt" out
    @test @allocated(augment!(out, camera, ops)) < @allocated(augment(camera, ops))
end

ops = (Rotate(45),CropNative(1:512,1:512))
@testset "$(str_showcompact(ops))" begin
    wv = @inferred Augmentor._augment(camera, ops)
    @test typeof(wv) <: SubArray
    @test typeof(wv.indexes) <: Tuple{Vararg{IdentityRange}}
    @test typeof(parent(wv)) <: InvWarpedView
    @test parent(parent(wv)).itp.coefs === camera
    @test_reference "reference/rot45_crop.txt" wv
    img = @inferred augment(camera, ops)
    @test img == parent(copy(wv))
    @test typeof(img) <: Array
    @test eltype(img) <: eltype(camera)
    @test_reference "reference/rot45_crop.txt" img
end

ops = (Rotate(45),CropSize(512,512))
@testset "$(str_showcompact(ops))" begin
    wv = @inferred Augmentor._augment(camera, ops)
    @test typeof(wv) <: SubArray
    @test typeof(wv.indexes) <: Tuple{Vararg{IdentityRange}}
    @test typeof(parent(wv)) <: InvWarpedView
    @test parent(parent(wv)).itp.coefs === camera
    @test_reference "reference/rot45_crop.txt" wv
    img = @inferred augment(camera, ops)
    @test img == parent(copy(wv))
    @test typeof(img) <: Array
    @test eltype(img) <: eltype(camera)
    @test_reference "reference/rot45_crop.txt" img
end

ops = (Rotate(-45),CropSize(256,256))
@testset "$(str_showcompact(ops))" begin
    wv = @inferred Augmentor._augment(camera, ops)
    @test typeof(wv) <: SubArray
    @test typeof(wv.indexes) <: Tuple{Vararg{IdentityRange}}
    @test typeof(parent(wv)) <: InvWarpedView
    @test parent(parent(wv)).itp.coefs === camera
    @test_reference "reference/rotr45_cropsize.txt" wv
    img = @inferred augment(camera, ops)
    @test img == parent(copy(wv))
    @test typeof(img) <: Array
    @test eltype(img) <: eltype(camera)
    @test_reference "reference/rotr45_cropsize.txt" img
end

ops = (Scale(.1,.2),NoOp())
@testset "$(str_showcompact(ops))" begin
    wv = @inferred Augmentor._augment(camera, ops)
    @test typeof(wv) <: InvWarpedView
    @test parent(wv).itp.coefs === camera
    @test_reference "reference/scale_noop.txt" wv
    img = @inferred augment(camera, ops)
    @test img == parent(copy(wv))
    @test typeof(img) <: Array
    @test eltype(img) <: eltype(camera)
    @test_reference "reference/scale_noop.txt" img
end

ops = (Scale(.1,.2),CropRatio())
@testset "$(str_showcompact(ops))" begin
    wv = @inferred Augmentor._augment(camera, ops)
    @test typeof(wv) <: SubArray
    @test typeof(wv.indexes) <: Tuple{Vararg{IdentityRange}}
    @test typeof(parent(wv)) <: InvWarpedView
    @test parent(parent(wv)).itp.coefs === camera
    @test_reference "reference/scale_cropratio.txt" wv
    img = @inferred augment(camera, ops)
    @test img == parent(copy(wv))
    @test typeof(img) <: Array
    @test eltype(img) <: eltype(camera)
    @test_reference "reference/scale_cropratio.txt" img
end

ops = (ShearX(45),NoOp())
@testset "$(str_showcompact(ops))" begin
    wv = @inferred Augmentor._augment(camera, ops)
    @test typeof(wv) <: InvWarpedView
    @test parent(wv).itp.coefs === camera
    @test_reference "reference/shearx_noop.txt" wv
    img = @inferred augment(camera, ops)
    @test img == parent(copy(wv))
    @test typeof(img) <: Array
    @test eltype(img) <: eltype(camera)
    @test_reference "reference/shearx_noop.txt" img
end

ops = (ShearY(45),CropNative(1:512,1:512))
@testset "$(str_showcompact(ops))" begin
    wv = @inferred Augmentor._augment(camera, ops)
    @test typeof(wv) <: SubArray
    @test typeof(wv.indexes) <: Tuple{Vararg{IdentityRange}}
    @test typeof(parent(wv)) <: InvWarpedView
    @test parent(parent(wv)).itp.coefs === camera
    @test_reference "reference/sheary_crop.txt" wv
    img = @inferred augment(camera, ops)
    @test img == parent(copy(wv))
    @test typeof(img) <: Array
    @test eltype(img) <: eltype(camera)
    @test_reference "reference/sheary_crop.txt" img
end

ops = (Crop(101:200,201:350),Scale(.2,.4))
@testset "$(str_showcompact(ops))" begin
    wv = @inferred Augmentor._augment(camera, ops)
    @test typeof(wv) <: SubArray
    @test typeof(wv.indexes) <: Tuple{Vararg{IdentityRange}}
    @test typeof(parent(wv)) <: InvWarpedView
    @test parent(parent(wv)).itp.coefs === camera
    @test_reference "reference/crop_scale.txt" wv
    img = @inferred augment(camera, ops)
    @test img == parent(copy(wv))
    @test typeof(img) <: Array
    @test eltype(img) <: eltype(camera)
    @test_reference "reference/crop_scale.txt" img
end

ops = (Crop(101:200,201:350),Zoom(1.3))
@testset "$(str_showcompact(ops))" begin
    wv = @inferred Augmentor._augment(camera, ops)
    @test typeof(wv) <: SubArray
    @test typeof(wv.indexes) <: Tuple{Vararg{IdentityRange}}
    @test typeof(parent(wv)) <: InvWarpedView
    @test parent(parent(wv)).itp.coefs === camera
    @test_reference "reference/crop_zoom.txt" wv
    img = @inferred augment(camera, ops)
    @test img == parent(copy(wv))
    @test typeof(img) <: Array
    @test eltype(img) <: eltype(camera)
    @test_reference "reference/crop_zoom.txt" img
end

ops = (Rotate(45),Zoom(2))
@testset "$(str_showcompact(ops))" begin
    wv = @inferred Augmentor._augment(camera, ops)
    @test typeof(wv) <: SubArray
    @test typeof(wv.indexes) <: Tuple{Vararg{IdentityRange}}
    @test typeof(parent(wv)) <: InvWarpedView
    @test parent(parent(wv)).itp.coefs === camera
    @test_reference "reference/rot45_zoom.txt" wv
    wv2 = Augmentor._augment(camera, (ops..., NoOp()))
    @test_reference "reference/rot45_zoom.txt" wv2
    img = @inferred augment(camera, ops)
    @test img == parent(copy(wv))
    @test typeof(img) <: Array
    @test eltype(img) <: eltype(camera)
    @test_reference "reference/rot45_zoom.txt" img
    img1, img2 = @inferred augment((N0f8.(camera),camera), ops)
    @test img1 == img2
    @test img2 == img
    @test img1 isa Array{N0f8}
    @test img2 isa Array{Gray{N0f8}}
    @test_reference "reference/rot45_zoom.txt" img2
end

ops = (Rotate(45),CropSize(200,200),Zoom(1.1),ConvertEltype(RGB{Float64}),SplitChannels())
@testset "$(str_showcompact(ops))" begin
    wv1 = @inferred Augmentor._augment(camera, ops[1:3])
    wv2 = @inferred Augmentor._augment(camera, ops[1:4])
    wv3 = @inferred Augmentor._augment(camera, ops)
    _, wv4 = @inferred Augmentor._augment((rotl90(camera), camera), ops)
    @test wv3 == wv4
    @test typeof(wv3) == typeof(wv4)
    img = colorview(RGB{Float64}, wv3)
    @test RGB{Float64}.(copy(wv1)) ≈ wv2
    @test collect(wv1) ≈ img
    @test_reference "reference/rot45_crop_zoom_convert.txt" wv2
end

# just for code coverage
@test typeof(@inferred(Augmentor.augment_impl(Rotate90()))) <: Expr


using UnicodePlots
import Pkg
using SpikingNeuralNetworks
SNN = SpikingNeuralNetworks
using SpikeSynchrony
using Statistics
using JLD
using Distributed
using SharedArrays
using Plots
using UnicodePlots
using Evolutionary
using Distributions
using LightGraphs
using Metaheuristics

##
# Override to function to include a state.
##
SNN.@load_units
#unicodeplots()

#outputs = simulate!(pop, T; cb = cb, inputs=input)

global Ne = 200;
global Ni = 50
global σee = 1.0
global pee = 0.5
global σei = 1.0
global pei = 0.5

function make_net_from_graph_structure(xx)#;

    xx = Int(round(xx))
    #@show(xx)
    #h = turan_graph(xx, xx)#, seed=1,cutoff=0.3)

    h = circular_ladder_graph(xx)#, xx)#, seed=1,cutoff=0.3)
    #hi = circular_ladder_graph(xx)#, seed=1,cutoff=0.3)
    E = SNN.IZ(; N = Ne, param = SNN.IZParameter(; a = 0.02, b = 0.2, c = -65, d = 8))
    I = SNN.IZ(; N = Ni, param = SNN.IZParameter(; a = 0.1, b = 0.2, c = -65, d = 2))
    #EE = SNN.SpikingSynapse(E, E, :v; σ = σee, p = 1.0)
    EI = SNN.SpikingSynapse(E, I, :v; σ = σei, p = 1.0)
    IE = SNN.SpikingSynapse(I, E, :v; σ = -1.0, p = 1.0)
    II = SNN.SpikingSynapse(I, I, :v; σ = -1.0, p = 1.0)
    # PINningSynapse
    P = [E, I]#, EEA]
    C = [EI, IE, II]#, EEA]
    #EE = SNN.PINningSynapse(E, E, :v; σ=0.5, p=0.8)
    #for n in 1:(N - 1)
    #    SNN.connect!(EE, n, n + 1, 50)
    #end
    #for (i,j) in enumerate(h.fadjlist) println(i,j) end
    EE = SNN.SpikingSynapse(E, E, :v; σ = 0.5, p = 0.8)

    @inbounds for (i, j) in enumerate(h.fadjlist)
        @inbounds for k in j
            SNN.connect!(EE, i, k, 10)
        end
    end

    @inbounds for (i, j) in enumerate(hi.fadjlist)
        @inbounds for k in j
            if i < Ni && k < Ni

                SNN.connect!(EI, i, k, 10)
                SNN.connect!(IE, i, k, 10)
                SNN.connect!(II, i, k, 10)
            end
        end
    end

    #for (i,j) in enumerate(h.fadjlist)
    #    for k in j
    #        SNN.connect!(EI, i, k, 50)
    #    end
    #end
    P = [E, I]#, EEA]
    C = [EE, EI, IE, II]#, EEA]
    return P, C

end

function make_net(Ne, Ni; σee = 1.0, pee = 0.5, σei = 1.0, pei = 0.5)
    Ne = 200
    Ni = 50

    E = SNN.IZ(; N = Ne, param = SNN.IZParameter(; a = 0.02, b = 0.2, c = -65, d = 8))
    I = SNN.IZ(; N = Ni, param = SNN.IZParameter(; a = 0.1, b = 0.2, c = -65, d = 2))
    EE = SNN.SpikingSynapse(E, E, :v; σ = σee, p = pee)
    EI = SNN.SpikingSynapse(E, I, :v; σ = σei, p = pei)
    IE = SNN.SpikingSynapse(I, E, :v; σ = -1.0, p = 0.5)
    II = SNN.SpikingSynapse(I, I, :v; σ = -1.0, p = 0.5)
    P = [E, I]#, EEA]
    C = [EE, EI, IE, II]#, EEA]
    return P, C
end
function get_trains(p)
    fire = p.records[:fire]
    x, y = Float32[], Float32[]
    for time in eachindex(fire)
        for neuron_id in findall(fire[time])
            push!(x, time)
            push!(y, neuron_id)
        end
    end
    cellsa = Array{Union{Missing,Any}}(undef, 1, Int(findmax(y)[1]))
    nac = Int(findmax(y)[1])
    for (inx, cell_id) in enumerate(1:nac)
        cellsa[inx] = []
    end
    @inbounds for cell_id in unique(y)
        @inbounds for (time, cell) in collect(zip(x, y))
            if Int(cell_id) == cell
                append!(cellsa[Int(cell_id)], time)

            end

        end
    end

    cellsa

end
global E
global spkd_ground

#P, C = make_net(Ne, Ni, σee = 0.5, pee = 0.8, σei = 0.5, pei = 0.8, a = 0.02)
#sggcu =[ CuArray(convert(Array{Float32,1},sg)) for sg in spkd_ground ]

#Flux.SGD
#Flux.gpu

function rmse(spkd)
    error = Losses(mean(spkd), spkd; agg = mean)
end

function rmse_depr(spkd)
    total = 0.0
    @inbounds for i = 1:size(spkd, 1)
        total += (spkd[i] - mean(spkd[i]))^2.0
    end
    return sqrt(total / size(spkd, 1))
end

function raster_difference(spkd0, spkd_found)
    maxi0 = size(spkd0)[2]
    maxi1 = size(spkd_found)[2]
    mini = findmin([maxi0, maxi1])[1]
    spkd = ones(mini)
    maxi = findmax([maxi0, maxi1])[1]

    if maxi > 0
        if maxi0 != maxi1
            return sum(ones(maxi))

        end
        if isempty(spkd_found[1, :])
            return sum(ones(maxi))
        end
    end
    spkd = ones(mini)
    @inbounds for i in eachindex(spkd)
        if !isempty(spkd0[i]) && !isempty(spkd_found[i])
            maxt1 = findmax(spkd0[i])[1]
            maxt2 = findmax(spkd_found[i])[1]
            maxt = findmax([maxt1, maxt2])[1]
            if maxt1 > 0.0 && maxt2 > 0.0
                t, S = SpikeSynchrony.SPIKE_distance_profile(
                    unique(sort(spkd0[i])),
                    unique(sort(spkd_found[i]));
                    t0 = 0.0,
                    tf = maxt,
                )
                spkd[i] = SpikeSynchrony.trapezoid_integral(t, S) / (t[end] - t[1])
            end
        end
    end
    spkd
end

#=
using Base.Threads

function raster_difference_threads(spkd0, spkd_found)
    maxi0 = size(spkd0)[2]
    maxi1 = size(spkd_found)[2]
    mini = findmin([maxi0, maxi1])[1]
    spkd = ones(mini)
    maxi = findmax([maxi0, maxi1])[1]
    if maxi > 0
        if maxi0 != maxi1
            return sum(ones(maxi))

        end
        if isempty(spkd_found[1, :])
            return sum(ones(maxi))
        end
    end
    spkd = ones(mini)
    N = length(spkd)
    Threads.foreach(
        Channel(nthreads()) do chnl
            for i in 1:N
                val = 1.0
                if !isempty(spkd0[i]) && !isempty(spkd_found[i])
                    maxt1 = findmax(spkd0[i])[1]
                    maxt2 = findmax(spkd_found[i])[1]
                    maxt = findmax([maxt1, maxt2])[1]
                    if maxt1 > 0.0 && maxt2 > 0.0
                        t, S = SpikeSynchrony.SPIKE_distance_profile(
                            unique(sort(spkd0[i])),
                            unique(sort(spkd_found[i]));
                            t0 = 0.0,
                            tf = maxt,
                        )
                        val = SpikeSynchrony.trapezoid_integral(t, S) / (t[end] - t[1])
                    end
                end
                put!(chnl, val)
                println("thread ", threadid(), ": ", "put ", val)
            end
        end
    end
    spkd
end
=#

function loss(model)
    #σee = model[1]
    #pee = model[2]
    #σei = model[3]
    #pei = model[4]
    #P1, C1 = make_net(Ne, Ni, σee = σee, pee = pee, σei = σei, pei = pei)#,a=a)
    #@show(model)
    println("best candidate ", 26)
    println(" ")
    #println("found ", model[1])
    P1, C1 = make_net_SNN(model[1])#,a=a)

    E1, I1 = P1
    SNN.monitor([E1, I1], [:fire])
    sim_length = 500
    @inbounds for t = 1:sim_length*ms
        E1.I = vec([11.5 for i = 1:sim_length])#vec(E_stim[t,:])#[i]#3randn(Ne)
        SNN.sim!(P1, C1, 1ms)
    end

    spkd_found = get_trains(P1[1])
    #println("Ground Truth \n")
    #SNN.raster([E]) |> display
    #println("Best Candidate \n")

    #SNN.raster([E1]) |> display

    error = raster_difference(spkd_ground, spkd_found)
    error = sum(error)
    #@show(error)
    #println("Broke here")

    error

end

function eval_best(params)
    xx = Int(round(params[1]))
    @show(xx)
    P1, C1 = make_net_SNN(xx)#,a=a)
    println("found ", xx)

    #σee = params[1]
    #pee = params[2]
    #σei = params[3]
    #pei = params[4]
    #P1, C1 = make_net(Ne, Ni, σee = σee, pee = pee, σei = σei, pei = pei)#,a=a)
    E1, I1 = P1
    SNN.monitor([E1, I1], [:fire])
    sim_length = 500
    @inbounds for t = 1:sim_length*ms
        E1.I = vec([11.5 for i = 1:sim_length])#vec(E_stim[t,:])#[i]#3randn(Ne)
        SNN.sim!(P1, C1, 1ms)
    end

    spkd_found = get_trains(P1[1])
    println("Ground Truth: \n")
    SNN.raster([E]) |> display
    println("candidate: \n")

    SNN.raster([E1]) |> display
    #error = raster_difference(spkd_ground,spkd_found)
    E1, spkd_found

end

function init_b(lower, upper)
    gene = []

    for (i, (l, u)) in enumerate(zip(lower, upper))
        p1 = rand(l:u, 1)
        append!(gene, p1)
    end
    gene
end

function initf(n)
    genesb = []
    for i = 1:n
        genes = init_b(lower, upper)
        append!(genesb, [genes])
    end
    genesb
end

function initd()
    population = initf(10)
    garray = zeros((length(population)[1], length(population[1])))
    for (i, p) in enumerate(population)
        garray[i, :] = p
    end
    garray[1, :]
end
function Evolutionary.trace!(
    record::Dict{String,Any},
    objfun,
    state,
    population,
    method::GA,
    options,
)
    idx = sortperm(state.fitpop)
    record["fitpop"] = state.fitpop[:]#idx[1:last(idx)]]
    record["pop"] = population[:]
    #record["σ"] = state.
end

#lower = Float32[0.0 0.0 0.0 0.0]# 0.03 4.0]
#upper = Float32[1.0 1.0 1.0 1.0]# 0.2 20.0]

#lower = Int32[3]# 0.0 0.0 0.0]# 0.03 4.0]
#upper = Int32[40]# 1.0 1.0 1.0]# 0.2 20.0]

#lower = Float32[0.0 0.0 0.0 0.0]# 0.03 4.0]
#upper = Float32[1.0 1.0 1.0 1.0]# 0.2 20.0]
#lower = vec(lower)
#upper = vec(upper)
#using Evolutionary, MultivariateStats
#range = Any[lower,upper]

# mutation = domainrange(fill(0.5,4))
#
#function Evolutionary.trace!(record::Dict{String,Any}, objfun, state, population, method::GA, options)
#    idx = sortperm(state.fitpop)
#record["population"] = population
#    state.fitpop[idx[1:5]]
#    record["fitpop"] = state.fitpop[idx[1:5]]
#end
#using Evolutionary

#=
function Evolutionary.value!(::Val{:serial}, fitness, objfun, population::AbstractVector{IT}) where {IT}
    @show(typeof(fitness))
    fitness = SharedArrays.SharedArray{Float32}(fitness)
    @time @sync @distributed for i in 1:length(population)

    #pmap(f, [::AbstractWorkerPool], c...; distributed=true, batch_size=) -> collection
        fitness[i] = value(objfun, population[i])
        #println("I'm worker $(myid()), working on i=$i")
    end
    @show(typeof(fitness))

    fitness
    #fitness = Array(fitness)
    #@show(fitness)
end
=#


module TestGeom

include("preamble.jl")

@testset "SPoint3D" begin
    @test_throws ArgumentError SPoint3D(rand(4, 10), 3)
    @test_throws BoundsError SPoint3D(rand(3, 10), 11)
    A = rand(Float32, 3, 10)
    @test @inferred(SPoint3D(A, 5)) == A[:, 5]
    @test eltype(A) === Float32
    @test eltype(@inferred(SPoint3D{Float64}(A, 5))) === Float64
end

@testset "MPoint3D" begin
    @test_throws ArgumentError MPoint3D(rand(4, 10), 3)
    @test_throws BoundsError MPoint3D(rand(3, 10), 11)
    A = rand(Float32, 3, 10)
    @test @inferred(MPoint3D(A, 5)) == A[:, 5]
    @test eltype(A) === Float32
    @test eltype(@inferred(MPoint3D{Float64}(A, 5))) === Float64
end

end # module


export FewShotDataLoader, MetaDataSample

using Serialization
using Random: shuffle!
import StatsBase
import Base

struct FewShotDataLoader
    labels::AbstractArray{L,1} where {L}
    samples::AbstractArray
    label2Idices::Dict

    function FewShotDataLoader(path::String)
        data = deserialize(path)
        @assert typeof(data) <: Dict "$path should containing a Dict, get $(typeof(data))"
        labels = data["labels"]
        samples = data["data"]
        @assert ndims(samples) >= 2 "ndims of samples should be at least 2, get $(ndims(samples))"
        N = size(labels)[end]
        label2Idices = Dict{eltype(labels),Array{Integer,1}}()
        for i = 1:N
            label = labels[i]
            if !haskey(label2Idices, label)
                label2Idices[label] = []
            end
            push!(label2Idices[label], i)
        end
        new(labels, samples, label2Idices)
    end
end

function Base.show(io::IO, dloader::FewShotDataLoader)
    print(
        io,
        "FewShotDataLoader(sample:$(size(dloader.samples)), labels:$(typeof(dloader.labels)))",
    )
end

"""
Fields
======

`support_samples`: array of support samples with size (`feature_size`..., `num_support`, `num_tasks`)
`support_labels`: array of labels of support samples with size (`num_support`, `num_tasks`)
`query_samples`: array of query samples with size (`feature_size`..., `num_query`, `num_tasks`)
`query_labels`: array of labels of query samples with size (`num_query`, `num_tasks`)
"""
struct MetaDataSample{T,L}
    support_n_ways::Integer
    support_k_shots::Integer
    query_n_ways::Integer
    query_k_shots::Integer
    support_samples::AbstractArray{T}
    support_labels::AbstractArray{L,2}
    query_samples::AbstractArray{T}
    query_labels::AbstractArray{L,2}
end

function Base.getproperty(m::MetaDataSample, name::Symbol)
    if name == :n_support
        return m.support_k_shots * m.support_n_ways
    elseif name == :n_query
        return m.query_k_shots * m.query_n_ways
    end
    getfield(m, name)
end

function Base.show(io::IO, meta_sample::MetaDataSample)
    print(
        io,
        "MetaDataSample(",
        "support_n_ways: $(meta_sample.support_n_ways), ",
        "support_k_shots: $(meta_sample.support_k_shots), ",
        "query_n_ways: $(meta_sample.query_n_ways), ",
        "query_k_shots: $(meta_sample.query_k_shots), ",
        "support_samples: [$(eltype(meta_sample.support_samples))]$(size(meta_sample.support_samples)), ",
        "support_labels: [$(eltype(meta_sample.support_labels))]$(size(meta_sample.support_labels)), ",
        "query_samples: [$(eltype(meta_sample.query_samples))]$(size(meta_sample.query_samples)), ",
        "query_labels: [$(eltype(meta_sample.query_labels))]$(size(meta_sample.query_labels)))",
    )
end

function MetaDataSample(;
    support_n_ways,
    support_k_shots,
    query_n_ways,
    query_k_shots,
    support_samples,
    support_labels,
    query_samples,
    query_labels,
)
    @assert support_n_ways > 0 "support_n_ways must be positive: $support_n_ways"
    @assert support_k_shots > 0 "support_k_shots must be positive: $support_k_shots"
    @assert query_n_ways > 0 "query_n_ways must be positive: $query_n_ways"
    @assert query_k_shots > 0 "query_k_shots must be positive: $query_k_shots"
    @assert ndims(support_samples) >= 3 "support_samples must be at least 3-dim"
    @assert ndims(query_samples) >= 3 "query_samples must be at least 3-dim"
    @assert support_n_ways * support_k_shots == size(support_samples)[end-1]
    @assert query_n_ways * query_k_shots == size(query_samples)[end-1]
    @assert size(support_samples)[1:end-2] == size(query_samples)[1:end-2] "size of support samples and query samples should be the same"
    MetaDataSample(
        support_n_ways,
        support_k_shots,
        query_n_ways,
        query_k_shots,
        support_samples,
        support_labels,
        query_samples,
        query_labels,
    )
end

function Base.size(meta_data_sample::MetaDataSample)
    return size(meta_data_sample.support_samples)[end]
end

function _sample(
    dloader::FewShotDataLoader;
    support_n_ways = 2,
    support_k_shots = 5,
    query_k_shots = support_k_shots,
)
    # find exclusive train/test samples
    uniq_labels = unique(dloader.labels)
    target_labels = StatsBase.sample(uniq_labels, support_n_ways, replace = false)
    support_idxs = []
    query_idxs = []
    for label in target_labels
        idx_map = Dict([(v, i) for (i, v) in enumerate(dloader.label2Idices[label])])
        selected_idxs = StatsBase.sample(dloader.label2Idices[label], support_k_shots)
        support_idxs = vcat(
            support_idxs,
            selected_idxs,
        )
        weights = ones(length(dloader.label2Idices[label]))
        for idx in selected_idxs
            weights[idx_map[idx]] = 0
        end
        weights = StatsBase.Weights(weights)
        selected_idxs =
            StatsBase.sample(dloader.label2Idices[label], weights, query_k_shots)
        query_idxs = vcat(
            query_idxs,
            selected_idxs
        )
    end
    @assert intersect(support_idxs, query_idxs) |> isempty
    shuffle!(support_idxs)
    shuffle!(query_idxs)
    support_samples = Utils.add_dim(
        dloader.samples[repeat([:], ndims(dloader.samples)-1)..., support_idxs]
    )
    support_labels = Utils.add_dim(dloader.labels[support_idxs])
    query_samples = Utils.add_dim(
        dloader.samples[repeat([:], ndims(dloader.samples)-1)..., query_idxs]
    )
    query_labels = Utils.add_dim(dloader.labels[query_idxs])
    return support_samples, support_labels, query_samples, query_labels
end

function StatsBase.sample(
    dloader::FewShotDataLoader;
    support_n_ways = 2,
    support_k_shots = 5,
    query_k_shots = support_k_shots,
)
    support_samples, support_labels, query_samples, query_labels = _sample(
        dloader,
        support_n_ways = support_n_ways,
        support_k_shots = support_k_shots,
        query_k_shots = query_k_shots,
    )
    MetaDataSample(
        support_n_ways = support_n_ways,
        support_k_shots = support_k_shots,
        query_n_ways = support_n_ways,
        query_k_shots = query_k_shots,
        support_samples = support_samples,
        support_labels = support_labels,
        query_samples = query_samples,
        query_labels = query_labels,
    )
end

function StatsBase.sample(
    dloader::FewShotDataLoader,
    n_tasks::Integer;
    support_n_ways = 2,
    support_k_shots = 5,
    query_k_shots = support_k_shots,
)
    support_sets = []
    support_labels = []
    query_sets = []
    query_labels = []
    for i = 1:n_tasks
        support, support_label, query, query_label = _sample(
            dloader,
            support_n_ways = support_n_ways,
            support_k_shots = support_k_shots,
            query_k_shots = query_k_shots,
        )
        push!(support_sets, support)
        push!(support_labels, support_label)
        push!(query_sets, query)
        push!(query_labels, query_label)
    end
    ndim_sample = ndims(support_sets[1])
    MetaDataSample(
        support_n_ways = support_n_ways,
        support_k_shots = support_k_shots,
        query_n_ways = support_n_ways,
        query_k_shots = query_k_shots,
        support_samples = cat(support_sets..., dims = Val(ndim_sample)),
        support_labels = cat(support_labels..., dims = Val(2)),
        query_samples = cat(query_sets..., dims = Val(ndim_sample)),
        query_labels = cat(query_labels..., dims = Val(2)),
    )
end


# The format of simpleweightedgraph files is as follows: for each graph,
# a one line header: "LightGraphs.SimpleWeightedGraph", <num_vertices>, <num_edges>, {"d" | "u"}, <name>[, <ver>, <vdatatype>, <wdatatype>, <graphcode>]
#   - "LightGraphs.SimpleWeightedGraph" is a fixed string
#   - num_vertices is an integer
#   - num_edges is an integer
#   - "d" for directed graph, "u" for undirected. Note that this
#       option does not perform any additional edge construction; it's
#       merely used to return the correct type of graph.
#   - name is a string
#   - ver is an int
#   - vdatatype is a string ("UInt8", etc.)
#   - wdatatype is a string describing the data type of the weights
#   - graphcode is a string.
# header followed by a list of (comma-delimited) edges - src,dst,weight
# Multiple graphs may be present in one file.

const FIXEDSTR = "LightGraphs.SimpleWeightedGraph"

struct SWGFormat <: AbstractGraphFormat end

struct SWGHeader
    nv::Int
    ne::Int
    is_directed::Bool
    name::String
    ver::Int
    vdtype::DataType    # vertex data type
    wdtype::DataType    # weight data type
    code::String
end
function show(io::IO, h::SWGHeader)
    isdir = h.is_directed ? "d" : "u"
    print(io, "$FIXEDSTR,$(h.nv),$(h.ne),$isdir,$(h.name),$(h.ver),$(h.vdtype),$(h.wdtype),$(h.code)")
end

function _swg_read_one_graph(f::IO, header::SWGHeader)
    T = header.vdtype
    U = header.wdtype
    if header.is_directed
        g = SimpleWeightedDiGraph{T, U}(header.nv)
    else
        g = SimpleWeightedGraph{T, U}(header.nv)
    end
    for i = 1:header.ne
        line = chomp(readline(f))
        if length(line) > 0
            src_s, dst_s, weight_s = split(line, r"\s*,\s*")
            src = parse(T, src_s)
            dst = parse(T, dst_s)
            weight = parse(U, weight_s)
            add_edge!(g, src, dst, weight)
        end
    end
    return g
end

function _swg_skip_one_graph(f::IO, n_e::Integer)
    for _ in 1:n_e
        readline(f)
    end
end

function _parse_header(s::AbstractString)
    addl_info = false
    fixedstr, nvstr, nestr, dirundir, graphname, _ver, _vdtype, _wdtype, graphcode  = split(s, r"s*,s*", limit=9)
    fixedstr != FIXEDSTR && error("Error parsing header.")

    n_v = parse(Int, nvstr)
    n_e = parse(Int, nestr)
    dirundir = strip(dirundir)
    directed = !(dirundir == "u")
    graphname = strip(graphname)
    ver = parse(Int, _ver)
    vdtype = eval(Symbol(_vdtype))
    wdtype = eval(Symbol(_wdtype))

    return SWGHeader(n_v, n_e, directed, graphname, ver, vdtype, wdtype, graphcode)
end

"""
    loadswg_mult(io)

Return a dictionary of (name=>graph) loaded from IO stream `io`.
"""
function loadswg_mult(io::IO)
    graphs = Dict{String,AbstractGraph}()
    while !eof(io)
        line = strip(chomp(readline(io)))
        if startswith(line, "#") || line == ""
            next
        else
            header = _parse_header(line)
            g = _swg_read_one_graph(io, header)
            graphs[header.name] = g
        end
    end
    return graphs
end

function loadswg(io::IO, gname::String)
    while !eof(io)
        line = strip(chomp(readline(io)))
        (startswith(line, "#") || line == "") && continue
        header = _parse_header(line)
        if gname == header.name
            return _swg_read_one_graph(io, header)
        else
            _swg_skip_one_graph(io, header.ne)
        end
    end
    error("Graph $gname not found")
end

"""
    saveswg(io, g, gname)

Write a graph `g` with name `gname` in a proprietary format
to the IO stream designated by `io`. Return 1 (number of graphs written).
"""
function saveswg(io::IO, g::AbstractGraph, gname::String)
    header = SWGHeader(nv(g), ne(g), is_directed(g), gname, 1, eltype(g), weighttype(g), "simpleweightedgraph")
    # write header line
    line = string(header)
    write(io, "$line\n")
    # write edges
    for e in edges(g)
        write(io, "$(src(e)),$(dst(e)),$(weight(e))\n")
    end
    return 1
end

"""
    saveswg_mult(io, graphs)

Write a dictionary of (name=>graph) to an IO stream `io`,
with default `GZip` compression. Return number of graphs written.
"""
function saveswg_mult(io::IO, graphs::Dict)
    ng = 0
    for (gname, g) in graphs
        ng += saveswg(io, g, gname)
    end
    return ng
end

loadgraph(io::IO, gname::String, ::SWGFormat) = loadswg(io, gname)
loadgraphs(io::IO, ::SWGFormat) = loadswg_mult(io)
savegraph(io::IO, g::AbstractGraph, gname::String, ::SWGFormat) = saveswg(io, g, gname)
savegraph(io::IO, g::AbstractGraph, ::SWGFormat) = saveswg(io, g, "graph")
savegraph(io::IO, d::Dict, ::SWGFormat) = saveswg_mult(io, d)


############################################################################################
"""
$(SIGNATURES)
Callable struct to make initializing SMC sampler easier in sampling library.

# Examples
```julia
```

"""
struct SMCConstructor{J<:AbstractConstructor,D<:SMCDefault} <: AbstractConstructor
    "Jitter kernel"
    kernel::J
    "SMC Default Arguments"
    default::D
    function SMCConstructor(
        kernel::J, default::D
    ) where {J<:AbstractConstructor,D<:SMCDefault}
        return new{J,D}(kernel, default)
    end
end
function (constructor::SMCConstructor)(
    _rng::Random.AbstractRNG,
    model::ModelWrapper,
    data::D,
    temperature::F,
    info::BaytesCore.SampleDefault
) where {D, F<:AbstractFloat}
    return SMC(
        _rng,
        constructor.kernel,
        Objective(model, data, Tagged(model), temperature),
        constructor.default,
        info
    )
end
function SMC(
    kernel::J; kwargs...
) where {J<:AbstractConstructor}
    return SMCConstructor(kernel, SMCDefault(; kwargs...))
end

function get_sym(constructor::SMCConstructor)
    return get_sym(constructor.kernel)
end

############################################################################################
"""
$(SIGNATURES)
Print diagnostics result of SMC sampler.

# Examples
```julia
```

"""
function results(
    diagnosticsᵛ::AbstractVector{M}, smc::SMC, Ndigits::Int64, quantiles::Vector{T}
) where {T<:Real,M<:SMCDiagnostics}
    ## Print Trace
    println("SMC algorithm with jitter kernel ", Base.nameof(typeof(smc.particles.kernel[begin])), " Diagnostics:")
    println(
        "Sampler finished after ",
        size(diagnosticsᵛ, 1),
        " iterations with average acceptance ratio of ",
        round(
            Statistics.mean(
                !diagnosticsᵛ[iter].resampled for iter in eachindex(diagnosticsᵛ)
            ) * 100;
            digits=Ndigits,
        ),
        "%.",
    )
    println(
        "Initial average ℓlikelihood per particle: ",
        round(Statistics.mean(diagnosticsᵛ[begin].ℓweights); digits=Ndigits),
        ", variance: ",
        round(Statistics.var(diagnosticsᵛ[begin].ℓweights); digits=Ndigits),
    )
    println(
        "Final average ℓlikelihood per particle: ",
        round(Statistics.mean(diagnosticsᵛ[end].ℓweights); digits=Ndigits),
        ", variance: ",
        round(Statistics.var(diagnosticsᵛ[end].ℓweights); digits=Ndigits),
    )
    ## Number of jittering steps
    println("Total number of jittering steps: ",
        Statistics.sum(diagnosticsᵛ[iter].jittersteps for iter in eachindex(diagnosticsᵛ)),
        ".")
    ## Rejuvenation steps
    t = convert(Int64, floor(length(diagnosticsᵛ) / 2))
    if t > 1
        println(
            "Rejuvenations on average every ",
            round(
                (
                    1 / (
                        1 - Statistics.mean(
                            !diagnosticsᵛ[iter].resampled for iter in eachindex(diagnosticsᵛ)
                        )
                    )
                );
                digits=Ndigits,
            ),
            " step, every ",
            round(
                1 / (1 - Statistics.mean(!diagnosticsᵛ[iter].resampled for iter in 1:t));
                digits=Ndigits,
            ),
            " steps in first half, every ",
            round(
                1 / (
                    1 - Statistics.mean(
                        !diagnosticsᵛ[iter].resampled for iter in (t + 1):length(diagnosticsᵛ)
                    )
                );
                digits=Ndigits,
            ),
            " steps in second half.",
        )
    end
    ##Rejuvenation Correlation
    rejuvenations = [diagnosticsᵛ[iter].resampled for iter in eachindex(diagnosticsᵛ)]
    if sum(rejuvenations) > 0
        println(
            "Quantiles for average rejuvenation correlation of parameter: ",
            round.(
                Statistics.quantile(
                    Statistics.mean.(
                        #!NOTE: If pairs[i] are constant (i.e., all resampled initial parameter come from same index), ρ = NaN for index
                        #filter(!isnan, diagnosticsᵛ[rejuvenations][iter].ρ) for
                        diagnosticsᵛ[rejuvenations][iter].ρ for iter in Base.OneTo(sum(rejuvenations))
                    ),
                    quantiles,
                );
                digits=Ndigits,
            ),
            ".",
        )
        println(
            "Quantiles for average number rejuvenation steps: ",
            round.(
                Statistics.quantile(
                    Statistics.mean.(
                        diagnosticsᵛ[rejuvenations][iter].jittersteps for
                        iter in Base.OneTo(sum(rejuvenations))
                    ),
                    quantiles,
                );
                digits=Ndigits,
            ),
            ".",
        )
    end
    ## ESS quantiles
    println(
        "Quantiles for ESS: ",
        round.(
            Statistics.quantile(
                [diagnosticsᵛ[iter].ESS for iter in eachindex(diagnosticsᵛ)], quantiles
            );
            digits=Ndigits,
        ),
        " for ",
        smc.tune.chains.Nchains,
        " chains.",
    )
    ## Return
    return nothing
end

############################################################################################
"""
$(SIGNATURES)
Infer MCMC diagnostics type.

# Examples
```julia
```

"""
function infer(
    _rng::Random.AbstractRNG,
    diagnostics::Type{AbstractDiagnostics},
    smc::SMC,
    model::ModelWrapper,
    data::D,
) where {D}
    #Pick one of the kernels to obtain type for prediction and jittering diagnostics.
    TPrediction = BaytesCore.infer(_rng, smc.particles.kernel[begin], model, data)
    TJitter = BaytesCore.infer(_rng, AbstractDiagnostics, smc.particles.kernel[begin], model, data)
    TTemperature = model.info.reconstruct.default.output
    TGenerated = infer_generated(_rng, smc, model, data)
    return SMCDiagnostics{TPrediction,TJitter,TTemperature,TGenerated}
end

############################################################################################
"""
$(SIGNATURES)
Infer type of generated quantities of PF sampler.

# Examples
```julia
```

"""
function infer_generated(
    _rng::Random.AbstractRNG, smc::SMC, model::ModelWrapper, data::D
) where {D}
    objective = Objective(model, data, smc.tune.tagged)
    return typeof(generate(_rng, objective, smc.tune.generated))
end

############################################################################################
#export
export SMCConstructor, infer


ENV["PYCALL_JL_RUNTIME_PYTHON"] = Sys.which("python")
include(joinpath(@__DIR__, "stats_evaluation_manipulation_with_unknown_parameters.jl"))

function evaluate()
    data_dir = normpath(joinpath(@__DIR__, "../data",
                                 "manipulation_with_unknown_parameters"));
    data = load(joinpath(data_dir, "statistical_data_manipulation.jld2"))
    safe_rngnums = data["safe_rngnums"];

    cost_tlqg = Array{Float64}[]
    tcalc_tlqg = Array{Float64}[]
    tcalc_tlqg_offline = Array{Float64}[]

    time_idx_array = [];

    num_successful_cases = 0;

    for rngnum in safe_rngnums
        println("------- RNG Seed: $(rngnum) & #Cases Tested: $(num_successful_cases) -------")
        println("4. Simulation with T-LQG Controller")
        y_history,tcalc_true_history,b_history,x_history,u_history,U_pool,Cs,Cu =
        simulate_main("tlqg",rng_num=rngnum, animate=false, plotfig=false, verbose=false);
        push!(cost_tlqg,getMetricHistory(x_history));
        push!(tcalc_tlqg,tcalc_true_history);
        if isempty(time_idx_array)
            push!(time_idx_array, unique([x.t for x in x_history]))
        end

        println("5. Simulation with T-LQG Controller (Offline, for Computation Time Statistics Only)")
        ~,tcalc_true_history,~,~,~,~,~,~ =
        simulate_main("tlqg_offline",rng_num=rngnum, animate=false, plotfig=false, verbose=false);
        push!(tcalc_tlqg_offline,tcalc_true_history);
        if isempty(time_idx_array)
            push!(time_idx_array, unique([x.t for x in x_history]))
        end

        num_successful_cases += 1
    end

    tcalc_tlqg = vcat(tcalc_tlqg...);
    tcalc_tlqg_offline = vcat(tcalc_tlqg_offline...);

    FileIO.save(joinpath(data_dir, "statistical_data_tlqg_manipulation.jld2"),
            "cost_tlqg", cost_tlqg, "tcalc_tlqg", tcalc_tlqg,
            "tcalc_tlqg_offline", tcalc_tlqg_offline,
            "safe_rngnums", safe_rngnums, "time_idx_array", time_idx_array);

    cost_sacbp = data["cost_sacbp"]
    cost_mcts = data["cost_mcts"]
    cost_ilqg = data["cost_ilqg"]
    tcalc_sacbp = data["tcalc_sacbp"]
    tcalc_mcts = data["tcalc_mcts"]
    tcalc_ilqg = data["tcalc_ilqg"];

    # Plot Stats Data
    cost_std_sacbp = [std([cost[ii] for cost in cost_sacbp]) for ii = 1:length(cost_sacbp[1])]
    cost_std_mcts = [std([cost[ii] for cost in cost_mcts]) for ii = 1:length(cost_mcts[1])]
    cost_std_ilqg = [std([cost[ii] for cost in cost_ilqg]) for ii = 1:length(cost_ilqg[1])];
    cost_std_tlqg = [std([cost[ii] for cost in cost_tlqg]) for ii = 1:length(cost_tlqg[1])];

    CostHistorySACBP = mean(cost_sacbp)
    CostHistoryMCTS = mean(cost_mcts)
    plot(time_idx_array[1][1:end-1],CostHistoryMCTS,label="MCTS-DPW",color=:lightcoral,ribbon=cost_std_mcts,fillalpha=0.2,linewidth=2.,linestyle=:dash);
    CostHistoryiLQG = mean(cost_ilqg)
    CostHistoryTLQGOnline = mean(cost_tlqg);
    plot!(time_idx_array[1][1:end-1],CostHistoryTLQGOnline,label="T-LQG",color=:darkgreen,ribbon=cost_std_tlqg,fillalpha=0.1,linewidth=2.,linestyle=:dot)
    plot!(time_idx_array[1][1:end-1],CostHistoryiLQG,label="Belief iLQG",color=:darkorange,ribbon=cost_std_ilqg,fillalpha=0.2,linewidth=2.,linestyle=:dashdot)
    plot!(time_idx_array[1][1:end-1],CostHistorySACBP,label="SACBP (Ours)",color=:darkblue,ribbon=cost_std_sacbp,fillalpha=0.2,linewidth=2.,
          xlabel="Time [s]",ylabel="Residual Norm",size=(400,400),ylim=(0.,5.))
    savefig(joinpath(data_dir, "manipulation_costs_with_tlqg.pdf"))

    tcalcTrueArray = [tcalc_mcts,tcalc_tlqg,tcalc_sacbp]
    meanTcalcTrueArray = mean.(tcalcTrueArray);
    stdTcalcTrueArray = std.(tcalcTrueArray);

    bar([1,2,3],meanTcalcTrueArray,yerror=stdTcalcTrueArray,linecolor=:black,marker=stroke(1.5,:black, :dash),c=[:lightcoral,:palegreen3,:royalblue3],label="");
    xticks!([1,2,3],["MCTS-DPW","T-LQG","SACBP (Ours)"])
    plot!([0.5,3.5],[3/4*0.2,3/4*0.2],width=2.,label="Targeted",size=(300,400),color=:darkred,ylim=(-0.03,0.8))
    ylabel!("Online Replanning Time [s]")
    savefig(joinpath(data_dir, "manipulation_time_with_tlqg.pdf"))
end

if abspath(PROGRAM_FILE) == @__FILE__
    evaluate();
end


"""

```julia
map_naive!(f,output,x,box)
```

Internal function or structure - interface may change.

# Extended help

Function that uses the naive pairwise mapping algorithm, for testing.

"""
function map_naive!(f, output, x, box::Box)
    @unpack unit_cell, cutoff_sq = box
    for i in 1:length(x) - 1
        xᵢ = x[i]
        for j in i + 1:length(x)
            xⱼ = wrap_relative_to(x[j], xᵢ, box)
            d2 = norm_sqr(xᵢ - xⱼ)
            if d2 <= cutoff_sq
                output = f(xᵢ, xⱼ, i, j, d2, output)
            end
        end
    end
    return output
end

function map_naive!(f, output, x, y, box::Box)
    @unpack unit_cell, cutoff_sq = box
    for i in 1:length(x)
        xᵢ = x[i]
        for j in 1:length(y)
            yⱼ = wrap_relative_to(y[j], xᵢ, box)
            d2 = norm_sqr(xᵢ - yⱼ)
            if d2 <= cutoff_sq
                output = f(xᵢ, yⱼ, i, j, d2, output)
            end
        end
    end
    return output
end

"""

```
view_celllist_particles(cl::CellList)
```

Internal function or structure - interface may change.

# Extended help

Auxiliary function to view the particles of a computing box, including images created
for computing purposes.

## Example

```julia-repl
julia> box = Box([ 100 50; 50 100 ],10);

julia> x = [ box.unit_cell_max .* rand(SVector{2,Float64}) for i in 1:1000 ];

julia> cl = CellList(x,box);

julia> p = CellListMap.view_celllist_particles(cl);

julia> using Plots

julia> scatter(Tuple.(p),label=nothing,xlims=(-10,180),ylims=(-10,180))

```

"""
function view_celllist_particles(cl::CellList{N,T}) where {N,T}
    x = SVector{N,T}[]
    for icell in 1:cl.n_cells_with_particles
        cell = cl.cells[icell]
        for ip in 1:cell.n_particles
            p = cell.particles[ip]
            push!(x, p.coordinates)
        end
    end
    return x
end

test_map2(box,cl;parallel=true) =
    map_pairwise!(
        (x, y, i, j, d2, s) -> begin
            s += 1
        end,
        0, box, cl, parallel=parallel
    )

function missing_pair(x1,x2)
    for x in eachrow(x2)
        found = false
        for y in eachrow(x1)
            (x[1] == y[1]) && (x[2] == y[2]) && (found = true)
            (x[1] == y[2]) && (x[2] == y[1]) && (found = true)
            if found
                break
            end
        end
        if !found
            return x
        end
    end
end

test_map(box,cl;parallel=true) = map_pairwise!((x, y, i, j, d2, s) -> s += 1/d2, 0., box, cl, parallel=parallel)
test_naive(x,box) = CellListMap.map_naive!((x, y, i, j, d2, s) -> s += 1/d2, 0., x, box)
function simple_test(x,box;parallel=true)
    cl = CellList(x,box,parallel=parallel)
    r_naive = test_naive(x,box)
    r_map = test_map(box,cl,parallel=parallel)
    println("naive = $r_naive")
    println("map   = $r_map")
    return "Test passed: $(r_naive ≈ r_map)"
end

function check_random_cells(
    N=3, M=10;
    show_progress=true,
    lcell=1,
    parallel=false,
    UnitCellType=OrthorhombicCell
)
    local x, box
    ntrial = 0
    show_progress && (p = Progress(10000, 0.5))
    while ntrial < 2000
        unit_cell_matrix = zeros(SMatrix{N,N,Float64})
        if UnitCellType == OrthorhombicCell
            for i in 1:N
                @set! unit_cell_matrix[i,i] = 1 + 10*rand()
            end
        else
            for i in 1:N, j in 1:N
                @set! unit_cell_matrix[i,j] = 1 + 10*rand()
            end
        end
        cutoff = 1 + rand()
        try
            if !check_unit_cell(unit_cell_matrix, cutoff, printerr=false)
                continue
            end
        catch
            @show unit_cell_matrix, cutoff
            return false
        end
        if UnitCellType == OrthorhombicCell
            box = Box([unit_cell_matrix[i,i] for i in 1:N], cutoff, lcell=lcell)
        else
            box = Box(unit_cell_matrix, cutoff, lcell=lcell)
        end
        if prod(box.nc) > 100000
            continue
        end
        show_progress && (next!(p))
        x = 10 * rand(SVector{N,Float64}, M)
        for i in eachindex(x)
            x[i] = x[i] .- 50
        end
        cl = CellList(x, box, parallel=parallel)
        test = test_map(box, cl, parallel=parallel)
        if test ≈ 0
            continue
        end
        ntrial += 1
        if !(test ≈ test_naive(x, box))
            show_progress && println("FOUND PROBLEMATIC SETUP.")
            return false, x, box
        else

        end
        if parallel
            aux = AuxThreaded(cl)
            cl = UpdateCellList!(x, box, cl, aux, parallel=parallel)
        else
            cl = UpdateCellList!(x, box, cl, parallel=parallel)
        end
        if !(test_map(box, cl, parallel=parallel) ≈ test_naive(x, box))
            show_progress && println("FOUND PROBLEMATIC SETUP.")
            return false, x, box
        end
    end
    show_progress && println(" ALL PASSED! ")
    return true, x, box
end

function drawbox(box::Box{UnitCellType,2}) where UnitCellType
    S = SVector{2,Float64}
    m = box.unit_cell.matrix
    x = [
        S(0., 0.),
        S(m[:,1]),
        S(m[:,1] + m[:,2]),
        S(m[:,2]),
        S(0., 0.)
    ]
    return x
end

function drawbox(box::Box{UnitCellType,3}) where UnitCellType
    S = SVector{3,Float64}
    m = box.unit_cell.matrix
    x = [
        S(0., 0., 0.),
        S(m[:,1]),
        S(m[:,1] + m[:,2]),
        S(m[:,2]),
        S(0., 0., 0.),
        S(m[:,1]),
        S(m[:,1] + m[:,3]),
        S(m[:,3]),
        S(0., 0., 0.),
        S(m[:,2]),
        S(m[:,2] + m[:,3]),
        S(m[:,2]),
        S(0., 0., 0.),
        S(m[:,1]),
        S(m[:,1] + m[:,2]),
        S(m[:,1] + m[:,2]) + m[:,3],
        S(m[:,1] + m[:,3]),
        S(m[:,3]),
        S(m[:,3]) + m[:,2],
        S(m[:,1] + m[:,2]) + m[:,3]
    ]
    return x
end

"""

```
draw_computing_cell(x,box::Box{UnitCellType,2}) where UnitCellType
```

```
draw_computing_cell(cl::CellList,box::Box{UnitCellType,2},x) where UnitCellType
```

This function creates a plot of the computing cell, in two dimensions.

"""
function draw_computing_cell(x, box::Box{UnitCellType,2};
    parallel=true,
    xticks=nothing,
    yticks=nothing,
) where UnitCellType
    cl = CellList(x, box, parallel=parallel)
    return draw_computing_cell(cl,box,x,xticks=xticks,yticks=yticks)
end
function draw_computing_cell(
    cl::CellList, box::Box{UnitCellType,2}, x;
    xticks = nothing,
    yticks = nothing,
) where UnitCellType
    box_points = drawbox(box)
    p = view_celllist_particles(cl)
    plt = Main.plot()
    Main.plot!(plt, Tuple.(box_points), label=:none)
    Main.scatter!(plt, Tuple.(p), label=:none, markeralpha=0.3)
    Main.scatter!(plt, Tuple.(wrap_to_first.(x, Ref(box))), label=:none)
    xmin = minimum(el[1] for el in p) - 3 * box.cell_size[1]
    xmax = maximum(el[1] for el in p) + 3 * box.cell_size[1]
    ymin = minimum(el[2] for el in p) - 3 * box.cell_size[2]
    ymax = maximum(el[2] for el in p) + 3 * box.cell_size[2]
    isnothing(xticks) && (xticks=(round.(digits=3, xmin:box.cell_size[1]:xmax)))
    isnothing(yticks) && (yticks=(round.(digits=3, ymin:box.cell_size[2]:ymax)))
    Main.plot!(plt,
        aspect_ratio=1,framestyle=:box,xrotation=60,
        xlims=(xmin, xmax),
        ylims=(ymin, ymax),
        xticks=xticks,
        yticks=yticks,
    )
    return plt
end

"""

```
draw_computing_cell(x,box::Box{UnitCellType,3}) where UnitCellType
```

This function creates a plot of the computing cell, in three dimensions.

"""
function draw_computing_cell(x, box::Box{UnitCellType,3}; parallel=true) where UnitCellType
    cl = CellList(x, box, parallel=parallel)
    box_points = drawbox(box)
    p = view_celllist_particles(cl)
    plt = Main.plot()
    Main.plot!(plt, Tuple.(box_points), label=:none)
    Main.scatter!(plt, Tuple.(p), label=:none, markeralpha=0.3)
    Main.scatter!(plt, Tuple.(wrap_to_first.(x, Ref(box))), label=:none, markeralpha=0.3)
    lims = Vector{Float64}[]
    for i in 1:3
        push!(lims, [ -2 * box.cell_size[i], box.nc[i] + 2 * box.cell_size[i] ])
    end
    Main.plot!(plt,
        aspect_ratio=1,framestyle=:box,xrotation=60,yrotation=-70,zrotation=0,
        xlims=lims[1],
        ylims=lims[2],
        zlims=lims[3],
        xticks=(round.(digits=3, lims[1][1]:box.cell_size[1]:lims[1][2])),
        yticks=(round.(digits=3, lims[2][1]:box.cell_size[2]:lims[2][2])),
        zticks=(round.(digits=3, lims[3][1]:box.cell_size[3]:lims[3][2])),
    )
    return plt
end

function compare_cells(cl1::CellList{N,T}, cl2::CellList) where {N,T}

    if cl1.n_cells_with_real_particles != cl2.n_cells_with_real_particles
        println("Error: n_cells_with_real_particles differ")
    end
    if cl1.n_particles != cl2.n_particles
        println("Error: n_particles differ")
    end
    for i in 1:cl1.n_cells_with_particles
        linear_index =
        j = findfirst(
            isequal(cl1.cell_indices[i]),
            @view(cl2.cell_indices[1:cl2.n_cells_with_particles])
        )
        if isnothing(j)
            println(" Warning: cells_with_particles: Could not find index $(cl1.cell_indices[i])")
            break
        end
    end
    for i in 1:cl1.n_cells_with_real_particles
        j = findfirst(
            isequal(cl1.cell_indices_real[i]),
            @view(cl2.cell_indices_real[1:cl2.n_cells_with_real_particles])
        )
        if isnothing(j)
            println(" Warning: cells_with_real_particles: Could not find index $(cl1.cell_indices_real[i])")
            break
        end
    end
    if length(cl1.projected_particles) != length(cl2.projected_particles)
        println("Error: length of project_particles differ.")
    end
    ncheck = 0
    nparticle_check = 0
    if length(cl1.cells) != length(cl2.cells)
        println("Error: lengths of cells lists differ")
    else
        differ = false
        for i in 1:cl1.n_cells_with_particles
            if differ
                break
            end
            ncheck += 1
            ci = cl1.cells[i]
            j = findfirst(c -> c.linear_index == ci.linear_index, cl2.cells)
            if isnothing(j)
                println("Error: linear index of $(ci.linear_index) not found.")
                differ = true
            else
                cj = cl2.cells[j]
                if ci.cartesian_index != cj.cartesian_index
                    println("Error: cartesian_index differ")
                    differ = true
                end
                if !(ci.center ≈ cj.center)
                    println("Error: centers differ")
                    differ = true
                end
                if ci.contains_real != cj.contains_real
                    println("Error: contains_real differ")
                    differ = true
                end
                if ci.n_particles != cj.n_particles
                    println("Error: n_particles differ")
                    differ = true
                end
                for p1 in ci.particles
                    j = findfirst(p -> p.index == p1.index, cj.particles)
                    if isnothing(j)
                        println("Error: could not find particle in cl2.")
                        @show p1
                        differ = true
                    else
                        nparticle_check += 1
                        p2 = cj.particles[j]
                        if p1 != p2
                            println("Error: particles differ. ")
                            @show p1
                            @show p2
                            differ = true
                            break
                        end
                    end
                    differ && break
                end
                if sum(p -> p.index, ci.particles[1:ci.n_particles]) !=
                   sum(p -> p.index, cj.particles[1:cj.n_particles])
                    println("Error: particles of cells differ")
                    differ = true
                end
                if differ
                    println("Error: on cell with linear_index: $(ci.linear_index)")
                end
            end
        end
    end
    println(" Checked $ncheck of $(cl1.n_cells_with_particles) cells.")
    println(" Checked $nparticle_check of $(cl1.n_particles) particles.")
end

function check_cl(cl,box;mark)
    for i in 1:cl.n_cells_with_particles-1
        cellᵢ = cl.cells[i]
        for j in i+1:cl.n_cells_with_particles
            cellⱼ = cl.cells[j]
            for ip in 1:cellᵢ.n_particles
                xᵢ = cellᵢ.particles[ip]
                for jp in 1:cellⱼ.n_particles
                    xⱼ = cellⱼ.particles[jp]
                    if norm(xᵢ.coordinates - xⱼ.coordinates) ≈ 0
                        println("duplicate found:")
                        @show mark
                        @show Threads.threadid()
                        @show i, j, cellᵢ.linear_index, cellⱼ.linear_index
                        @show cellᵢ.n_particles
                        @show cellⱼ.n_particles
                        @show xᵢ
                        @show xⱼ
                        @show particle_cell(xᵢ.coordinates,box)
                        @show particle_cell(xⱼ.coordinates,box)
                        @show cell_linear_index(box.nc,particle_cell(xᵢ.coordinates,box))
                        @show cell_linear_index(box.nc,particle_cell(xⱼ.coordinates,box))
                        error()
                    end
                end
            end
        end
    end
end

# Create a set of random particles with an atomic density simiilar to that
# of water at room temperature and pressure.
function xatomic(n=100_000;cutoff=12.,lcell=1,type=:vec)
    atomic_density_of_water = 0.1 # atoms/ Angstrom^3
    vol = n/atomic_density_of_water
    l = vol^(1/3)
    if type == :mat
        x = l*rand(3,n)
    elseif type == :vec
        x = [ l*rand(SVector{3,Float64}) for _ in 1:n ]
    end
    box = Box([l,l,l],cutoff,lcell=lcell)
    return x, box
end

# Create a set of  random particles with the typical density of galaxies
function xgalactic(n=8_000_000;cutoff=5.,lcell=1,type=:vec)
    galactic_density = 12 # galaxyes per cubic mega-parsecs (https://doi.org/10.1093/mnras/106.2.121)
    vol = n/galactic_density
    l = vol^(1/3) # mega-parsecs
    if type == :mat
        x = l*rand(3,n)
    elseif type == :vec
        x = [ l*rand(SVector{3,Float64}) for _ in 1:n ]
    end
    box = Box([l,l,l],cutoff,lcell=lcell)
    return x, box
end


module TestProductOfSets

using Test

import MathOptInterface
const MOI = MathOptInterface

function runtests()
    for name in names(@__MODULE__; all = true)
        if startswith("$(name)", "test_")
            @testset "$(name)" begin
                getfield(@__MODULE__, name)()
            end
        end
    end
    return
end

MOI.Utilities.@mix_of_scalar_sets(
    _ScalarSets,
    MOI.EqualTo{T},
    MOI.GreaterThan{T},
    MOI.LessThan{T},
    MOI.Interval{T},
)

function test_scalar_set_index()
    sets = _ScalarSets{Float64}()
    @test MOI.Utilities.set_index(sets, MOI.EqualTo{Float64}) == 1
    @test MOI.Utilities.set_index(sets, MOI.GreaterThan{Float64}) == 2
    @test MOI.Utilities.set_index(sets, MOI.Interval{Float64}) == 4
    @test MOI.Utilities.set_index(sets, MOI.LessThan{Float64}) == 3
    @test MOI.Utilities.set_index(sets, MOI.ZeroOne) === nothing
end

function test_scalar_set_types()
    sets = _ScalarSets{Float64}()
    @test MOI.Utilities.set_types(sets) == [
        MOI.EqualTo{Float64},
        MOI.GreaterThan{Float64},
        MOI.LessThan{Float64},
        MOI.Interval{Float64},
    ]
end

function test_scalar_basic()
    sets = _ScalarSets{Float64}()
    ci = MOI.ConstraintIndex{
        MOI.ScalarAffineFunction{Float64},
        MOI.EqualTo{Float64},
    }(
        12345,
    )
    @test !MOI.is_valid(sets, ci)
    i = MOI.Utilities.set_index(sets, MOI.EqualTo{Float64})
    ci_value = MOI.Utilities.add_set(sets, i)
    ci = MOI.ConstraintIndex{
        MOI.ScalarAffineFunction{Float64},
        MOI.EqualTo{Float64},
    }(
        ci_value,
    )
    @test MOI.is_valid(sets, ci)
    @test MOI.Utilities.rows(sets, ci) == ci.value
end

function test_scalar_dimension()
    sets = _ScalarSets{Float64}()
    @test MOI.dimension(sets) == 0
    MOI.Utilities.add_set(sets, 1)
    @test MOI.dimension(sets) == 1
    MOI.Utilities.add_set(sets, 1)
    @test MOI.dimension(sets) == 2
    MOI.Utilities.add_set(sets, 2)
    @test MOI.dimension(sets) == 3
end

function test_scalar_empty()
    sets = _ScalarSets{Float64}()
    @test MOI.is_empty(sets)
    MOI.Utilities.add_set(sets, 1)
    @test !MOI.is_empty(sets)
    MOI.empty!(sets)
    @test MOI.is_empty(sets)
end

function test_scalar_ConstraintTypesPresent()
    sets = _ScalarSets{Float64}()
    @test MOI.get(sets, MOI.ListOfConstraintTypesPresent()) == []
    MOI.Utilities.add_set(sets, 1)
    @test MOI.get(sets, MOI.ListOfConstraintTypesPresent()) ==
          [(MOI.ScalarAffineFunction{Float64}, MOI.EqualTo{Float64})]
end

function test_scalar_NumberOfConstraints()
    sets = _ScalarSets{Float64}()
    MOI.Utilities.add_set(sets, 1)
    MOI.Utilities.add_set(sets, 1)
    MOI.Utilities.add_set(sets, 2)
    MOI.Utilities.final_touch(sets)
    @test MOI.get(
        sets,
        MOI.NumberOfConstraints{MOI.VariableIndex,MOI.ZeroOne}(),
    ) == 0
    for (x, S) in zip([2, 1, 0, 0], MOI.Utilities.set_types(sets))
        @test MOI.get(
            sets,
            MOI.NumberOfConstraints{MOI.ScalarAffineFunction{Float64},S}(),
        ) == x
    end
    return
end

function test_scalar_ListOfConstraintIndices()
    sets = _ScalarSets{Float64}()
    MOI.Utilities.add_set(sets, 2)
    MOI.Utilities.add_set(sets, 4)
    MOI.Utilities.add_set(sets, 1)
    for (x, S) in zip([[3], [1], [], [2]], MOI.Utilities.set_types(sets))
        ci = MOI.get(
            sets,
            MOI.ListOfConstraintIndices{MOI.ScalarAffineFunction{Float64},S}(),
        )
        @test ci == MOI.ConstraintIndex{MOI.ScalarAffineFunction{Float64},S}.(x)
    end
    return
end

MOI.Utilities.@product_of_sets(
    _VectorSets,
    MOI.Nonpositives,
    MOI.Nonnegatives,
    MOI.EqualTo{T},
)

function test_vector_set_index()
    sets = _VectorSets{Float64}()
    @test MOI.Utilities.set_index(sets, MOI.Nonpositives) == 1
    @test MOI.Utilities.set_index(sets, MOI.Nonnegatives) == 2
    @test MOI.Utilities.set_index(sets, MOI.EqualTo{Float64}) == 3
    @test MOI.Utilities.set_index(sets, MOI.LessThan{Float64}) === nothing
end

function test_vector_set_types()
    sets = _VectorSets{Float64}()
    @test MOI.Utilities.set_types(sets) ==
          [MOI.Nonpositives, MOI.Nonnegatives, MOI.EqualTo{Float64}]
end

function test_vector_basic()
    sets = _VectorSets{Float64}()
    nonneg_i = MOI.Utilities.add_set(
        sets,
        MOI.Utilities.set_index(sets, MOI.Nonnegatives),
        2,
    )
    equalto_i = MOI.Utilities.add_set(
        sets,
        MOI.Utilities.set_index(sets, MOI.EqualTo{Float64}),
    )
    MOI.Utilities.final_touch(sets)
    VAF = MOI.VectorAffineFunction{Float64}
    @test !MOI.is_valid(sets, MOI.ConstraintIndex{VAF,MOI.Nonnegatives}(12345))
    nonneg_ci = MOI.ConstraintIndex{VAF,MOI.Nonnegatives}(nonneg_i)
    @test MOI.is_valid(sets, nonneg_ci)
    @test MOI.Utilities.rows(sets, nonneg_ci) == 1:2
    SAF = MOI.ScalarAffineFunction{Float64}
    equalto_ci = MOI.ConstraintIndex{SAF,MOI.EqualTo{Float64}}(equalto_i)
    @test MOI.is_valid(sets, equalto_ci)
    @test MOI.Utilities.rows(sets, equalto_ci) == 3
end

function test_vector_dimension()
    sets = _VectorSets{Float64}()
    MOI.Utilities.final_touch(sets)
    @test MOI.dimension(sets) == 0

    sets = _VectorSets{Float64}()
    MOI.Utilities.add_set(sets, 1, 1)
    MOI.Utilities.final_touch(sets)
    @test MOI.dimension(sets) == 1

    sets = _VectorSets{Float64}()
    MOI.Utilities.add_set(sets, 1, 1)
    MOI.Utilities.add_set(sets, 1, 2)
    MOI.Utilities.final_touch(sets)
    @test MOI.dimension(sets) == 3

    sets = _VectorSets{Float64}()
    MOI.Utilities.add_set(sets, 1, 1)
    MOI.Utilities.add_set(sets, 1, 2)
    MOI.Utilities.add_set(sets, 2, 3)
    MOI.Utilities.final_touch(sets)
    @test MOI.dimension(sets) == 6

    sets = _VectorSets{Float64}()
    MOI.Utilities.add_set(sets, 1, 1)
    MOI.Utilities.add_set(sets, 1, 2)
    MOI.Utilities.add_set(sets, 2, 3)
    MOI.Utilities.add_set(sets, 3)
    MOI.Utilities.final_touch(sets)
    @test MOI.dimension(sets) == 7
end

function test_vector_empty()
    sets = _VectorSets{Float64}()
    @test MOI.is_empty(sets)
    MOI.Utilities.add_set(sets, 1, 1)
    @test !MOI.is_empty(sets)
    MOI.empty!(sets)
    @test MOI.is_empty(sets)
end

function test_vector_ConstraintTypesPresent()
    sets = _VectorSets{Float64}()
    MOI.Utilities.final_touch(sets)
    @test MOI.get(sets, MOI.ListOfConstraintTypesPresent()) == []
    sets = _VectorSets{Float64}()
    MOI.Utilities.add_set(sets, 1, 1)
    MOI.Utilities.add_set(sets, 3)
    MOI.Utilities.final_touch(sets)
    @test MOI.get(sets, MOI.ListOfConstraintTypesPresent()) == [
        (MOI.VectorAffineFunction{Float64}, MOI.Nonpositives),
        (MOI.ScalarAffineFunction{Float64}, MOI.EqualTo{Float64}),
    ]
end

function test_vector_NumberOfConstraints()
    sets = _VectorSets{Float64}()
    MOI.Utilities.add_set(sets, 1, 2)
    MOI.Utilities.add_set(sets, 1, 2)
    MOI.Utilities.add_set(sets, 2, 2)
    MOI.Utilities.add_set(sets, 3)
    MOI.Utilities.final_touch(sets)
    @test MOI.get(
        sets,
        MOI.NumberOfConstraints{MOI.VectorAffineFunction{Float64},MOI.Zeros}(),
    ) == 0
    for (x, S) in zip([2, 1], MOI.Utilities.set_types(sets)[1:2])
        @test MOI.get(
            sets,
            MOI.NumberOfConstraints{MOI.VectorAffineFunction{Float64},S}(),
        ) == x
    end
    @test MOI.get(
        sets,
        MOI.NumberOfConstraints{
            MOI.ScalarAffineFunction{Float64},
            MOI.EqualTo{Float64},
        }(),
    ) == 1
    return
end

function test_vector_ListOfConstraintIndices()
    sets = _VectorSets{Float64}()
    MOI.Utilities.add_set(sets, 2, 2)
    MOI.Utilities.add_set(sets, 1, 4)
    MOI.Utilities.add_set(sets, 2, 3)
    MOI.Utilities.add_set(sets, 3)
    MOI.Utilities.final_touch(sets)
    VAF = MOI.VectorAffineFunction{Float64}
    @test MOI.get(sets, MOI.ListOfConstraintIndices{VAF,MOI.Zeros}()) ==
          MOI.ConstraintIndex{VAF,MOI.Zeros}[]
    for (x, S) in zip([[1], [1, 3]], MOI.Utilities.set_types(sets)[1:2])
        ci = MOI.get(sets, MOI.ListOfConstraintIndices{VAF,S}())
        @test ci == MOI.ConstraintIndex{VAF,S}.(x)
    end
    F, S = MOI.ScalarAffineFunction{Float64}, MOI.EqualTo{Float64}
    @test MOI.get(sets, MOI.ListOfConstraintIndices{F,S}()) ==
          [MOI.ConstraintIndex{F,S}(1)]
    return
end

"""
    test_vector_ListOfConstraintIndices2()

Test a more complicated sequence of dimensions to check the `_UnevenIterator`
works appropriately.
"""
function test_vector_ListOfConstraintIndices2()
    sets = _VectorSets{Float64}()
    MOI.Utilities.add_set(sets, 2, 2)
    MOI.Utilities.add_set(sets, 2, 3)
    MOI.Utilities.add_set(sets, 2, 2)
    MOI.Utilities.add_set(sets, 2, 4)
    MOI.Utilities.final_touch(sets)
    S = MOI.Utilities.set_types(sets)[2]
    VAF = MOI.VectorAffineFunction{Float64}
    indices = MOI.get(sets, MOI.ListOfConstraintIndices{VAF,S}())
    @test indices == MOI.ConstraintIndex{VAF,S}.([1, 3, 6, 8])
end

end

TestProductOfSets.runtests()


const gridfieldstosave = [:nx, :ny, :Lx, :Ly, :X, :Y]

"""
    Output(prob, filename, fieldtuples...)

Define output for the Problem prob with fields and functions that calculate
the output in (fldname, func) tuples.
"""
mutable struct Output
  prob::Problem
  filename::String
  fields::Dict{Symbol,Function}
  init::Bool
end

function Output(prob::Problem, filename::String, fields::Dict{Symbol,Function})
  if filename[end-4:end] == ".jld2" # Initialize output: remove trailing ".jld2" to form 'basefilename'
    basefilename = filename[1:end-5]
  else
    basefilename = filename
  end

  filename = basefilename*".jld2"
  n = 0
  while isfile(filename) # append numbers until unique name is found
    n += 1
    filename = basefilename * "_$n.jld2"
  end

  saveproblem(prob, filename)
  Output(prob, filename, fields, true)
end

function Output(prob::Problem, filename::String, fieldtuples...)
  Output(prob, filename, Dict{Symbol,Function}(
    [ (symfld[1], symfld[2]) for symfld in fieldtuples ]
  ))
end

getindex(out::Output, key) = out.fields[key](out.prob)

"""
    saveoutput(out)

Save current output fields for file in out.filename.
"""
function saveoutput(out::Output)
  groupname = "timeseries"
  jldopen(out.filename, "a+") do file
    file["$groupname/t/$(out.prob.step)"] = out.prob.t
    for fieldname in keys(out.fields)
      file["$groupname/$fieldname/$(out.prob.step)"] = out[fieldname]
    end
  end
  nothing
end

"""
    saveproblem(prob, filename)

Save certain aspects of a problem timestepper, grid, and params. Functions
that are fields in params are not saved.
"""
function saveproblem(prob::AbstractProblem, filename::String)
  jldopen(filename, "a+") do file
    file["timestepper/dt"] = prob.ts.dt   # Timestepper

    for field in gridfieldstosave         # Grid
      file["grid/$field"] = getfield(prob.grid, field)
    end

    for name in fieldnames(typeof(prob.params))   # Params
      field = getfield(prob.params, name)
      if !(typeof(field) <: Function)
        file["params/$name"] = field
      end
    end
  end
  nothing
end

saveproblem(out::Output) = saveproblem(out.prob, out.filename)

"""
    savediagnostic(diag, diagname)

Save diagnostics to file, labeled by the string diagname.
"""
function savediagnostic(diag::AbstractDiagnostic, diagname::String, filename::String)
  jldopen(filename, "a+") do file
    file["diags/$diagname/time"] = diag.time
    file["diags/$diagname/data"] = diag.data
  end
  nothing
end


export DynaAgent

"""
    DynaAgent(;kwargs...)

`DynaAgent` is first introduced in: *Sutton, Richard S. "Dyna, an integrated architecture for learning, planning, and reacting." ACM Sigart Bulletin 2.4 (1991): 160-163.*

# Keywords & Fields

- `π`::[`AbstractPolicy`](@ref): the policy to use
- `model`::[`AbstractEnvironmentModel`](@ref): describe the environment to interact with
- `buffer`::[`AbstractTurnBuffer`](@ref): used to store transitions between agent and environment
- `role=:DEFAULT`: used to distinguish different agents
- `plan_step::Int=10`: the count of planning steps

The main difference between [`DynaAgent`](@ref) and [`Agent`](@ref) is that an environment model is involved. It is best described in the book: *Sutton, Richard S., and Andrew G. Barto. Reinforcement learning: An introduction. MIT press, 2018.*

![](../assets/img/RL_book_fig_8_1.png)

![](../assets/img/RL_book_fig_8_2.png)
"""
Base.@kwdef struct DynaAgent{P<:AbstractPolicy, B<:AbstractTurnBuffer, M<:AbstractEnvironmentModel, R} <: AbstractAgent
    π::P
    model::M
    buffer::B
    role::R = :DEFAULT
    plan_step::Int = 10
end

(agent::DynaAgent)(obs::Observation) = agent.π(obs)

function update!(agent::DynaAgent, experience::Pair)
    π, model, buffer = agent.π, agent.model, agent.buffer

    push!(buffer, experience)
    update!(model, buffer, π)  # model learning
    update!(π, buffer, model; plan_step = agent.plan_step)
end

update!(model::AbstractEnvironmentModel, buffer::AbstractTurnBuffer, π::AbstractPolicy) =
    update!(model, buffer)

function update!(model::AbstractEnvironmentModel, buffer::AbstractTurnBuffer)
    transitions = extract_transitions(buffer, model)
    if !isnothing(transitions)
        update!(model, transitions)
    end
end

function update!(
    model::PrioritizedSweepingSampleModel,
    buffer::AbstractTurnBuffer,
    π::AbstractPolicy,
)
    transition = extract_transitions(buffer, model)
    if !isnothing(transition)
        update!(model, transition, priority(transition, π))
    end
end

function update!(
    π::AbstractPolicy,
    buffer::AbstractTurnBuffer,
    model::AbstractEnvironmentModel;
    plan_step = 1,
)
    update!(π, buffer)  # direct RL
    update!(π, model; plan_step = plan_step)  # planning
end