const BBox = Rect2D{Float32}

const Optional{T} = Union{Nothing, T}

struct RectSides{T<:Real}
    left::T
    right::T
    bottom::T
    top::T
end

abstract type Side end

struct Left <: Side end
struct Right <: Side end
struct Top <: Side end
struct Bottom <: Side end
# for protrusion content:
struct TopLeft <: Side end
struct TopRight <: Side end
struct BottomLeft <: Side end
struct BottomRight <: Side end

struct Inner <: Side end
struct Outer <: Side end

abstract type GridDir end
struct Col <: GridDir end
struct Row <: GridDir end

struct RowCols{T <: Union{Number, Vector{Float64}}}
    lefts::T
    rights::T
    tops::T
    bottoms::T
end

"""
Used to specify space that is occupied in a grid. Like 1:1|1:1 for the first square,
or 2:3|1:4 for a rect over the 2nd and 3rd row and the first four columns.
"""
struct Span
    rows::UnitRange{Int64}
    cols::UnitRange{Int64}
end

"""
An object that can be aligned that also specifies how much space it occupies in
a grid via its span.
"""
mutable struct GridContent{G, T} # G should be GridLayout but can't be used before definition
    parent::Optional{G}
    al::T
    sp::Span
    side::Side
    needs_update::Node{Bool}
    protrusions_handle::Optional{Function}
    computedsize_handle::Optional{Function}
end

abstract type AlignMode end

struct Inside <: AlignMode end
struct Outside <: AlignMode
    padding::RectSides{Float32}
end
Outside() = Outside(0f0)
Outside(padding::Real) = Outside(RectSides{Float32}(padding, padding, padding, padding))
Outside(left::Real, right::Real, bottom::Real, top::Real) =
    Outside(RectSides{Float32}(left, right, bottom, top))

abstract type ContentSize end
abstract type GapSize <: ContentSize end

struct Auto <: ContentSize
    trydetermine::Bool # false for determinable size content that should be ignored
    ratio::Float64 # float ratio in case it's not determinable

    Auto(trydetermine::Bool = true, ratio::Real = 1.0) = new(trydetermine, ratio)
end
Auto(ratio::Real) = Auto(true, ratio)

struct Fixed <: GapSize
    x::Float64
end
struct Relative <: GapSize
    x::Float64
end
struct Aspect <: ContentSize
    index::Int
    ratio::Float64
end

mutable struct LayoutNodes{T, G} # G again GridLayout
    suggestedbbox::Node{BBox}
    protrusions::Node{RectSides{Float32}}
    computedsize::Node{NTuple{2, Optional{Float32}}}
    autosize::Node{NTuple{2, Optional{Float32}}}
    computedbbox::Node{BBox}
    gridcontent::Optional{GridContent{G, T}} # the connecting link to the gridlayout
end

mutable struct GridLayout
    content::Vector{GridContent}
    nrows::Int
    ncols::Int
    rowsizes::Vector{ContentSize}
    colsizes::Vector{ContentSize}
    addedrowgaps::Vector{GapSize}
    addedcolgaps::Vector{GapSize}
    alignmode::AlignMode
    equalprotrusiongaps::Tuple{Bool, Bool}
    needs_update::Node{Bool}
    block_updates::Bool
    layoutnodes::LayoutNodes
    attributes::Attributes
    _update_func_handle::Optional{Function} # stores a reference to the result of on(obs)

    function GridLayout(
        content, nrows, ncols, rowsizes, colsizes,
        addedrowgaps, addedcolgaps, alignmode, equalprotrusiongaps, needs_update,
        layoutnodes, attributes)

        gl = new(content, nrows, ncols, rowsizes, colsizes,
            addedrowgaps, addedcolgaps, alignmode, equalprotrusiongaps,
            needs_update, false, layoutnodes, attributes, nothing)

        validategridlayout(gl)

        # attach_parent!(gl, parent)

        # on(needs_update) do update
        #     request_update(gl)
        # end

        gl
    end
end

struct AxisAspect
    aspect::Float32
end

struct DataAspect end

abstract type Ticks end

struct AutoLinearTicks <: Ticks
    idealtickdistance::Float32
end

struct ManualTicks <: Ticks
    values::Vector{Float32}
    labels::Vector{String}
end

struct AxisContent{T}
    content::T
    attributes::Attributes
end

mutable struct LineAxis
    parent::Scene
    protrusion::Node{Float32}
    attributes::Attributes
    decorations::Dict{Symbol, Any}
    tickpositions::Node{Vector{Point2f0}}
    tickvalues::Node{Vector{Float32}}
    ticklabels::Node{Vector{String}}
end

abstract type LObject end

mutable struct LAxis <: AbstractPlotting.AbstractScene
    parent::Scene
    scene::Scene
    plots::Vector{AxisContent}
    xaxislinks::Vector{LAxis}
    yaxislinks::Vector{LAxis}
    limits::Node{BBox}
    layoutnodes::LayoutNodes
    needs_update::Node{Bool}
    attributes::Attributes
    block_limit_linking::Node{Bool}
    decorations::Dict{Symbol, Any}
end

mutable struct LColorbar <: LObject
    parent::Scene
    scene::Scene
    layoutnodes::LayoutNodes
    attributes::Attributes
    decorations::Dict{Symbol, Any}
end

mutable struct LText <: LObject
    parent::Scene
    layoutnodes::LayoutNodes
    text::AbstractPlotting.Text
    attributes::Attributes
end

mutable struct LRect <: LObject
    parent::Scene
    layoutnodes::LayoutNodes
    rect::AbstractPlotting.Poly
    attributes::Attributes
end

struct LSlider <: LObject
    scene::Scene
    layoutnodes::LayoutNodes
    attributes::Attributes
    decorations::Dict{Symbol, Any}
end

struct LButton <: LObject
    scene::Scene
    layoutnodes::LayoutNodes
    attributes::Attributes
    decorations::Dict{Symbol, Any}
end

struct LToggle <: LObject
    scene::Scene
    layoutnodes::LayoutNodes
    attributes::Attributes
    decorations::Dict{Symbol, Any}
end

abstract type LegendElement end

struct LineElement <: LegendElement
    attributes::Attributes
end

struct MarkerElement <: LegendElement
    attributes::Attributes
end

struct PolyElement <: LegendElement
    attributes::Attributes
end

struct LegendEntry
    elements::Vector{LegendElement}
    attributes::Attributes
end

struct LLegend <: LObject
    scene::Scene
    entries::Node{Vector{LegendEntry}}
    layoutnodes::LayoutNodes
    attributes::Attributes
    decorations::Dict{Symbol, Any}
    entrytexts::Vector{LText}
    entryplots::Vector{Vector{AbstractPlot}}
end

const Indexables = Union{UnitRange, Int, Colon}


using NNlib: DenseConvDims

# descriptor

mutable struct ConvDesc
    ptr::cudnnConvolutionDescriptor_t
end

unsafe_free!(cd::ConvDesc) = cudnnDestroyConvolutionDescriptor(cd.ptr)

Base.unsafe_convert(::Type{cudnnConvolutionDescriptor_t}, cd::ConvDesc)=cd.ptr

function cdsize(w, nd)
    isa(w, Integer) && return Cint[fill(w,nd)...]
    length(w) == nd && return Cint[reverse(w)...]
    length(w) == 2*nd && return Cint[reverse(w[nd+1:end])...]
    throw(DimensionMismatch())
end

pdsize(w, nd)=Cint[reverse(psize(w,nd))...]
function psize(w, nd)
    isa(w, Integer) && return Cint[fill(w,nd)...]
    length(w) == nd && return w
    length(w) == 2*nd && return w[1:nd]
    throw(DimensionMismatch())
end

Base.cconvert(::Type{cudnnConvolutionMode_t}, x::Bool) = x ? CUDNN_CROSS_CORRELATION : CUDNN_CONVOLUTION

function ConvDesc(T, N, padding, stride, dilation, mode)
    cd = Ref{cudnnConvolutionDescriptor_t}()
    cudnnCreateConvolutionDescriptor(cd)
    if version() >= v"4"
        cudnnSetConvolutionNdDescriptor(cd[],N,cdsize(padding,N),cdsize(stride,N),cdsize(dilation,N),mode,cudnnDataType(T))
    elseif version() >= v"3"
        cudnnSetConvolutionNdDescriptor_v3(cd[],N,cdsize(padding,N),cdsize(stride,N),cdsize(dilation,N),mode,cudnnDataType(T))
    else
        cudnnSetConvolutionNdDescriptor(cd[],N,cdsize(padding,N),cdsize(stride,N),cdsize(dilation,N),mode)
    end
    cudnnSetConvolutionMathType(cd[], math_mode())
    this = ConvDesc(cd[])
    finalizer(unsafe_free!, this)
    return this
end

function ConvDesc(T, cdims::DenseConvDims)
    pd = NNlib.padding(cdims)
    if !all(pd[1:2:end] .== pd[2:2:end])
        @warn("CuDNN does not support asymmetric padding; defaulting to symmetric choice")
    end
    return ConvDesc(T, NNlib.spatial_dims(cdims), pd[1:2:end], NNlib.stride(cdims),
                       NNlib.dilation(cdims), NNlib.flipkernel(cdims))
end

# wrappers

# Forward

function cudnnGetConvolutionForwardAlgorithmMaxCount()
    count=@argout(
        cudnnGetConvolutionForwardAlgorithmMaxCount(
            handle(),
            out(Ref{Cint}()))
        )[]
    return count
end

# will be removed in cuDNN 8
function cudnnGetConvolutionForwardAlgorithm(y::DenseCuArray{T,N}, x::DenseCuArray{T,N}, w::DenseCuArray{T,N},
                                             cdims::DenseConvDims; preference=0, workspacesize=1<<32) where {T,N}
    algo=@argout(
        cudnnGetConvolutionForwardAlgorithm(
            handle(), TensorDesc(x),
            FilterDesc(w), ConvDesc(T, cdims),
            TensorDesc(y),
            cudnnConvolutionFwdPreference_t(preference),
            Csize_t(workspacesize),
            out(Ref{cudnnConvolutionFwdAlgo_t}()))
        )[]
    return algo
end

function cudnnGetConvolutionForwardAlgorithm_v7(y::DenseCuArray{T,N}, x::DenseCuArray{T,N}, w::DenseCuArray{T,N},
                                                cdims::DenseConvDims; count=-1) where {T,N}
    if count < 0
        count = cudnnGetConvolutionForwardAlgorithmMaxCount()
    end
    perfResults = Array{cudnnConvolutionFwdAlgoPerf_t, 1}(undef, count)
    returnedAlgoCount=@argout(
        cudnnGetConvolutionForwardAlgorithm_v7(
            handle(), TensorDesc(x),
            FilterDesc(w), ConvDesc(T, cdims),
            TensorDesc(y),
            Cint(count),
            out(Ref{Cint}()),
            perfResults)
        )[]
    return returnedAlgoCount, perfResults
end

function cudnnFindConvolutionForwardAlgorithm(y::DenseCuArray{T,N}, x::DenseCuArray{T,N}, w::DenseCuArray{T,N},
                                              cdims::DenseConvDims; count=-1) where {T,N}
    if count < 0
        count = cudnnGetConvolutionForwardAlgorithmMaxCount()
    end
    perfResults = Array{cudnnConvolutionFwdAlgoPerf_t, 1}(undef, count)
    returnedAlgoCount=@argout(
        cudnnFindConvolutionForwardAlgorithm(
            handle(), TensorDesc(x),
            FilterDesc(w), ConvDesc(T, cdims),
            TensorDesc(y),
            Cint(count),
            out(Ref{Cint}()),
            perfResults)
        )[]
    return returnedAlgoCount, perfResults
end

function cudnnFindConvolutionForwardAlgorithmEx(y::DenseCuArray{T,N}, x::DenseCuArray{T,N}, w::DenseCuArray{T,N},
                                              cdims::DenseConvDims; count=-1, workspacesize=1<<32) where {T,N}
    if count < 0
        count = cudnnGetConvolutionForwardAlgorithmMaxCount()
    end
    @workspace size=workspacesize workspace->begin
        perfResults = Array{cudnnConvolutionFwdAlgoPerf_t, 1}(undef, count)
        returnedAlgoCount=@argout(
            cudnnFindConvolutionForwardAlgorithmEx(
                handle(), TensorDesc(x), x,
                FilterDesc(w), w, ConvDesc(T, cdims),
                TensorDesc(y), y,
                Cint(count),
                out(Ref{Cint}()),
                perfResults,
                workspace,
                workspacesize)
            )[]
        return returnedAlgoCount, perfResults
    end
end

function cudnnConvolutionForward(y::DenseCuArray{T,N}, x::DenseCuArray{T,N}, w::DenseCuArray{T,N},
                                 cdims::DenseConvDims; algo=0, alpha=1, beta=0) where {T,N}
    @workspace size=@argout(
            cudnnGetConvolutionForwardWorkspaceSize(
                handle(), TensorDesc(x),
                FilterDesc(w), ConvDesc(T, cdims),
                TensorDesc(y),
                cudnnConvolutionFwdAlgo_t(algo),
                out(Ref{Csize_t}()))
        )[] workspace->begin
            cudnnConvolutionForward(
                handle(), Ref{T}(alpha), TensorDesc(x), x, FilterDesc(w), w,
                ConvDesc(T,cdims), cudnnConvolutionFwdAlgo_t(algo), workspace,
                sizeof(workspace), Ref{T}(beta), TensorDesc(y), y)
        end
    return y
end

function cudnnConvolutionBiasActivationForward(y::DenseCuArray{T,N}, x::DenseCuArray{T,N}, w::DenseCuArray{T,N}, z::DenseCuArray{T,N}, bias::DenseCuArray{T,N},
                                               cdims::DenseConvDims; algo=0, alpha1=1, alpha2=1,
                                               activationMode=CUDNN_ACTIVATION_RELU, activationCoeff=0.0, activationReluNanOpt=CUDNN_NOT_PROPAGATE_NAN) where {T,N}
    @workspace size=@argout(
            cudnnGetConvolutionForwardWorkspaceSize(
                handle(), TensorDesc(x),
                FilterDesc(w), ConvDesc(T, cdims),
                TensorDesc(y),
                cudnnConvolutionFwdAlgo_t(algo),
                out(Ref{Csize_t}()))
        )[] workspace->begin
            cudnnConvolutionBiasActivationForward(
                handle(), Ref(T(alpha1)), TensorDesc(x), x, FilterDesc(w), w,
                ConvDesc(T, cdims), cudnnConvolutionFwdAlgo_t(algo), workspace,
                sizeof(workspace), Ref(T(alpha2)), TensorDesc(z), z, TensorDesc(bias), bias, ActivationDesc(activationMode, activationCoeff, activationReluNanOpt), TensorDesc(y),y)
        end
    return y
end

# Backward data

function cudnnGetConvolutionBackwardDataAlgorithmMaxCount()
    count=@argout(
        cudnnGetConvolutionBackwardDataAlgorithmMaxCount(
            handle(),
            out(Ref{Cint}()))
        )[]
    return count
end

# will be removed in cuDNN 8
function cudnnGetConvolutionBackwardDataAlgorithm(dx::DenseCuArray{T,N}, w::DenseCuArray{T,N}, dy::DenseCuArray{T,N},
                                                  cdims::DenseConvDims; preference=0, workspacesize=1<<32) where {T,N}
    algo=@argout(
        cudnnGetConvolutionBackwardDataAlgorithm(
            handle(), FilterDesc(w), TensorDesc(dy), ConvDesc(T, cdims),
            TensorDesc(dx), cudnnConvolutionBwdDataPreference_t(preference),
            Csize_t(workspacesize), out(Ref{cudnnConvolutionBwdDataAlgo_t}()))
        )[]
    return algo
end

function cudnnGetConvolutionBackwardDataAlgorithm_v7(dx::DenseCuArray{T,N}, w::DenseCuArray{T,N}, dy::DenseCuArray{T,N},
                                                     cdims::DenseConvDims; count=-1) where {T,N}
    if count < 0
        count = cudnnGetConvolutionBackwardDataAlgorithmMaxCount()
    end
    perfResults = Array{cudnnConvolutionBwdDataAlgoPerf_t, 1}(undef, count)
    returnedAlgoCount=@argout(
        cudnnGetConvolutionBackwardDataAlgorithm_v7(
            handle(), FilterDesc(w), TensorDesc(dy),
            ConvDesc(T, cdims), TensorDesc(dx),
            Cint(count),
            out(Ref{Cint}()), perfResults)
        )[]
    return returnedAlgoCount, perfResults
end

function cudnnFindConvolutionBackwardDataAlgorithm(dx::DenseCuArray{T,N}, w::DenseCuArray{T,N}, dy::DenseCuArray{T,N},
                                                   cdims::DenseConvDims; count=-1) where {T,N}
    if count < 0
        count = cudnnGetConvolutionBackwardDataAlgorithmMaxCount()
    end
    perfResults = Array{cudnnConvolutionBwdDataAlgoPerf_t, 1}(undef, count)
    returnedAlgoCount=@argout(
        cudnnFindConvolutionBackwardDataAlgorithm(
            handle(), FilterDesc(w), TensorDesc(dy),
            ConvDesc(T, cdims), TensorDesc(dx),
            Cint(count),
            out(Ref{Cint}()), perfResults)
        )[]
    return returnedAlgoCount, perfResults
end

function cudnnFindConvolutionBackwardDataAlgorithmEx(dx::DenseCuArray{T,N}, w::DenseCuArray{T,N}, dy::DenseCuArray{T,N},
                                                     cdims::DenseConvDims; count=-1, workspacesize=1<<32) where {T,N}
    if count < 0
        count = cudnnGetConvolutionBackwardDataAlgorithmMaxCount()
    end
    @workspace size=workspacesize workspace->begin
        perfResults = Array{cudnnConvolutionBwdDataAlgoPerf_t, 1}(undef, count)
        returnedAlgoCount=@argout(
            cudnnFindConvolutionBackwardDataAlgorithmEx(
                handle(), FilterDesc(w), w, TensorDesc(dy), dy,
                ConvDesc(T, cdims), TensorDesc(dx), dx,
                Cint(count),
                out(Ref{Cint}()),
                perfResults, workspace,
                workspacesize)
            )[]
        return returnedAlgoCount, perfResults
    end
end

function cudnnConvolutionBackwardData(dx::DenseCuArray{T,N}, w::DenseCuArray{T,N}, dy::DenseCuArray{T,N},
                                      cdims::DenseConvDims; algo=0, alpha=1, beta=0) where {T,N}
    @workspace size=@argout(
            cudnnGetConvolutionBackwardDataWorkspaceSize(
                handle(), FilterDesc(w),
                TensorDesc(dy), ConvDesc(T, cdims), TensorDesc(dx),
                cudnnConvolutionBwdDataAlgo_t(algo),
                out(Ref{Csize_t}()))
        )[] workspace->begin
            cudnnConvolutionBackwardData(
                handle(), Ref{T}(alpha), FilterDesc(w), w,
                TensorDesc(dy), dy, ConvDesc(T, cdims),
                cudnnConvolutionBwdDataAlgo_t(algo),
                workspace, sizeof(workspace),
                Ref{T}(beta), TensorDesc(dx), dx)
        end
    return dx
end

# Backward filter

function cudnnGetConvolutionBackwardFilterAlgorithmMaxCount()
    count=@argout(
        cudnnGetConvolutionBackwardFilterAlgorithmMaxCount(
            handle(),
            out(Ref{Cint}()))
        )[]
    return count
end

# will be removed in cuDNN 8
function cudnnGetConvolutionBackwardFilterAlgorithm(dw::DenseCuArray{T,N}, x::DenseCuArray{T,N}, dy::DenseCuArray{T,N},
                                                    cdims::DenseConvDims; preference=0, workspacesize=1<<32) where {T,N}
    algo=@argout(
        cudnnGetConvolutionBackwardFilterAlgorithm(
            handle(), TensorDesc(x), TensorDesc(dy),
            ConvDesc(T, cdims), FilterDesc(dw), cudnnConvolutionBwdFilterPreference_t(preference),
            Csize_t(workspacesize), out(Ref{cudnnConvolutionBwdFilterAlgo_t}()))
        )[]
    return algo
end

function cudnnGetConvolutionBackwardFilterAlgorithm_v7(dw::DenseCuArray{T,N}, x::DenseCuArray{T,N}, dy::DenseCuArray{T,N},
                                                     cdims::DenseConvDims; count=-1) where {T,N}
    if count < 0
        count = cudnnGetConvolutionBackwardFilterAlgorithmMaxCount()
    end
    perfResults = Array{cudnnConvolutionBwdFilterAlgoPerf_t, 1}(undef, count)
    returnedAlgoCount=@argout(
        cudnnGetConvolutionBackwardFilterAlgorithm_v7(
            handle(), TensorDesc(x), TensorDesc(dy),
            ConvDesc(T, cdims), FilterDesc(dw),
            Cint(count),
            out(Ref{Cint}()),
            perfResults)
        )[]
    return returnedAlgoCount, perfResults
end

function cudnnFindConvolutionBackwardFilterAlgorithm(dw::DenseCuArray{T,N}, x::DenseCuArray{T,N}, dy::DenseCuArray{T,N},
                                                     cdims::DenseConvDims; count=-1) where {T,N}
    if count < 0
        count = cudnnGetConvolutionBackwardFilterAlgorithmMaxCount()
    end
    perfResults = Array{cudnnConvolutionBwdFilterAlgoPerf_t, 1}(undef, count)
    returnedAlgoCount=@argout(
        cudnnFindConvolutionBackwardFilterAlgorithm(
            handle(), TensorDesc(x), TensorDesc(dy),
            ConvDesc(T, cdims), FilterDesc(dw),
            Cint(count),
            out(Ref{Cint}()),
            perfResults)
        )[]
    return returnedAlgoCount, perfResults
end

function cudnnFindConvolutionBackwardFilterAlgorithmEx(dw::DenseCuArray{T,N}, x::DenseCuArray{T,N}, dy::DenseCuArray{T,N},
                                                     cdims::DenseConvDims; count=-1, workspacesize=1<<32) where {T,N}
    if count < 0
        count = cudnnGetConvolutionBackwardFilterAlgorithmMaxCount()
    end
    @workspace size=workspacesize workspace->begin
        perfResults = Array{cudnnConvolutionBwdFilterAlgoPerf_t, 1}(undef, count)
        returnedAlgoCount=@argout(
            cudnnFindConvolutionBackwardFilterAlgorithmEx(
                handle(), TensorDesc(x), x, TensorDesc(dy),
                dy, ConvDesc(T, cdims), FilterDesc(dw), dw,
                Cint(count),
                out(Ref{Cint}()),
                perfResults, workspace,
                workspacesize)
            )[]
        return returnedAlgoCount, perfResults
    end
end

function cudnnConvolutionBackwardFilter(dw::DenseCuArray{T,N}, x::DenseCuArray{T,N}, dy::DenseCuArray{T,N},
                                        cdims::DenseConvDims; algo=0, alpha=1, beta=0) where {T,N}
    @workspace size=@argout(
            cudnnGetConvolutionBackwardFilterWorkspaceSize(
                handle(), TensorDesc(x),
                TensorDesc(dy),
                ConvDesc(T, cdims),
                FilterDesc(dw),
                cudnnConvolutionBwdFilterAlgo_t(algo),
                out(Ref{Csize_t}()))
        )[] workspace->begin
            cudnnConvolutionBackwardFilter(
                handle(), Ref{T}(alpha), TensorDesc(x), x,
                TensorDesc(dy), dy, ConvDesc(T, cdims),
                cudnnConvolutionBwdFilterAlgo_t(algo), workspace,
                sizeof(workspace), Ref{T}(beta), FilterDesc(dw), dw)
        end
    return dw
end

# Backward bias

function cudnnConvolutionBackwardBias(db::DenseCuArray{T,N}, dy::DenseCuArray{T,N}; alpha=1, beta=0) where {T,N}
    cudnnConvolutionBackwardBias(handle(),
                                 Ref(T(alpha)), TensorDesc(dy), dy,
                                 Ref(T(beta )), TensorDesc(db), db)
    return db
end


"""
    AbstractFitnessStrategy

Base type for fitness strategies.

See [`fitnessfun`](@ref)
"""
abstract type AbstractFitnessStrategy end

"""
    fitnessfun(s::T, x, y) where T <: AbstractFitnessStrategy

Returns an `AbstractFitness` given the features `x` and labels `y`.

Allows for simple control over e.g. train vs validation split.
"""
fitnessfun(s::T, x, y) where T <: AbstractFitnessStrategy = error("Not implemented for $(T)!")

"""
    AbstractEvolutionStrategy

Base type for strategies on how to perform the evolution.

See [`evostrategy`](@ref)
"""
abstract type AbstractEvolutionStrategy end

"""
    evostrategy(s::AbstractEvolutionStrategy, inshape)

Returns an `AbstractEvolution`.

Argument `inshape` is the size of the input feature maps (i.e. how many pixels images are) and may be used to determine which mutation operations are allowed for example to avoid that feature maps accidentally become 0 sized.
"""
evostrategy(s::AbstractEvolutionStrategy, inshape) = evostrategy_internal(s, inshape)

"""
    struct TrainSplitAccuracy{T} <: AbstractFitnessStrategy
        TrainSplitAccuracy(;split, accuracyconfig, accuracyfitness, trainconfig, trainfitness)

Strategy to train model on a subset of the training data and measure fitness as accuracy on the rest.

Size of subset for accuracy fitness is `ceil(Int, split * nobs)` where `nobs` is the size of along the last dimension of the input.

# Arguments
- `accuracyconfig` (default [`BatchedIterConfig()`](@ref)) determine how to iterate over the accuracy subset.
- `accuracyfitness` (default [`AccuracyVsSize`](@ref)) determine how to measure fitness based on `accuracyconfig`.
- `trainconfig` (default [`TrainIterConfig()`](@ref)) determines how to iterate over the training subset.
- `trainfitness` is a function accepting the iterator produced by `trainconfig` and the fitness strategy produced
   by `accuracyfitness` and returns the `AbstractFitness` to be used (default [`TrainThenFitness`](@ref) wrapped in
   [`GpuFitness`](@ref)).
"""
struct TrainSplitAccuracy{S, VC, VF, TC, TF} <: AbstractFitnessStrategy
    split::S
    accconfig::VC
    accfitness::VF
    trainconfig::TC
    trainfitness::TF
end
function TrainSplitAccuracy(;split=0.1,
            accuracyconfig=BatchedIterConfig(),
            accuracyfitness=AccuracyVsSize,
            trainconfig=TrainIterConfig(),
            trainfitness=(iter, accf) -> GpuFitness(TrainThenFitness(StatefulGenerationIter(iter), Flux.Losses.logitcrossentropy, ADAM(), accf, 0.0)))

    return TrainSplitAccuracy(split, accuracyconfig, accuracyfitness, trainconfig, trainfitness)
end
function fitnessfun(s::TrainSplitAccuracy, x, y)
    train_x, acc_x = split_examples(x, ceil(Int, last(size(x)) * s.split))
    train_y, acc_y = split_examples(y, ceil(Int, last(size(y)) * s.split))

    aiter = dataiter(s.accconfig, acc_x, acc_y)
    afitness = s.accfitness(aiter)

    titer = dataiter(s.trainconfig, train_x, train_y)
    return s.trainfitness(titer, afitness)
end

split_examples(a::AbstractArray{T, 1}, splitpoint) where T = a[1:end-splitpoint], a[end-splitpoint+1:end]
split_examples(a::AbstractArray{T, 2}, splitpoint) where T = a[:,1:end-splitpoint], a[:,end-splitpoint+1:end]
split_examples(a::AbstractArray{T, 4}, splitpoint) where T = a[:,:,:,1:end-splitpoint], a[:,:,:,end-splitpoint+1:end]

"""
    AccuracyVsSize(data, accdigits=2, accwrap=identity)

Produces an `AbstractFitness` which measures fitness accuracy on `data` and based on number of parameters.

The two are combined so that a candidate `a` which achieves higher accuracy rounded to the first `accdigits` digits compared to a candidate `b` will always have a better fitness.

If the first `accdigits` of accuracy is the same the candidate with fewer parameters will get higher fitness.

Accuracy part of the fitness is calculated by `accwrap(`[`AccuracyFitness(data)`](@ref)`)`.
"""
AccuracyVsSize(data, accdigits=2; accwrap=identity) = sizevs(accwrap(AccuracyFitness(data)), accdigits)
function sizevs(f::AbstractFitness, accdigits=2)
    truncacc = MapFitness(x -> round(x, digits=accdigits), f)

    size = SizeFitness()
    sizefit = MapFitness(x -> min(10.0^-(accdigits+1), 1 / x), size)

    return AggFitness(+, truncacc, sizefit)
end

"""
    struct TrainAccuracyVsSize <: AbstractFitnessStrategy
    TrainAccuracyVsSize(;trainconfig, trainfitness)

Produces an `AbstractFitness` which measures fitness accuracy on training data and based on number of parameters combined in the same way as is done for [`AccuracyVsSize`](@ref).

# Arguments
- `trainconfig` (default [`TrainIterConfig()`](@ref)) determines how to iterate over the training subset.
- `trainfitness` is a function accepting the iterator produced by `trainconfig` and the fitness strategy produced
   by `accuracyfitness` and returns the `AbstractFitness` to be used (default [`TrainAccuracyFitness`](@ref) wrapped in
   [`GpuFitness`](@ref)).

Beware that fitness as accuracy on training data will make evolution favour overfitted candidates.
"""
struct TrainAccuracyVsSize{TC, TF} <: AbstractFitnessStrategy
    trainconfig::TC
    trainfitness::TF
end
function TrainAccuracyVsSize(;
                        trainconfig=TrainIterConfig(),
                        trainfitness = dataiter -> sizevs(GpuFitness(TrainAccuracyFitness(
                                                                            dataiter=StatefulGenerationIter(dataiter),
                                                                            defaultloss=Flux.Losses.logitcrossentropy, defaultopt = ADAM()))))
        return TrainAccuracyVsSize(trainconfig, trainfitness)
end
function fitnessfun(s::TrainAccuracyVsSize, x, y)
    iter = dataiter(s.trainconfig, x, y)
    return s.trainfitness(iter)
end

"""
    struct BatchedIterConfig{T, V}
    BatchedIterConfig(;batchsize=32, dataaug=identity, iterwrap=GpuIterator)

Configuration for creating batch iterators from array data.

The function `dataiter(s::BatchedIterConfig, x, y)` creates an iterator which returns a tuple of batches from `x` and `y` respectively.

More specifically, the result of `s.iterwrap(zip(s.dataaug(bx), by))` will be returned where `bx` and `by` are `BatchIterator`s.
"""
struct BatchedIterConfig{T, V}
    batchsize::Int
    dataaug::T
    iterwrap::V
end
BatchedIterConfig(;batchsize=32, dataaug=identity, iterwrap=GpuIterator) = BatchedIterConfig(batchsize, dataaug, iterwrap)
dataiter(s::BatchedIterConfig, x, y) = dataiter(x, y, s.batchsize, s.dataaug) |> s.iterwrap

"""
    struct ShuffleIterConfig{T, V}
    ShuffleIterConfig(;batchsize=32, seed=123, dataaug=identity, iterwrap=GpuIterator)

Configuration for creating shuffled batch iterators from array data. Data will be re-shuffled every time the iterator restarts.

The function `dataiter(s::ShuffleIterConfig, x, y)` creates an iterator which returns a tuple of batches from `x` and `y` respectively.

More specifically, the result of `s.iterwrap(Iterators.map(((x,y),) -> (s.dataaug(x), y), iter))` will be returned where `iter` is a `BatchIterator` over `x` and `y` with `shuffle=true`.

Note there there is no upper bound on how many generations are supported as the returned iterator cycles the data indefinitely. Use e.g. `Iterators.take(itr, cld(nepochs * nbatches_per_epoch, nbatches_per_gen))` to limit to `nepochs`
epochs.
"""
struct ShuffleIterConfig{T, V}
    batchsize::Int
    seed::Int
    dataaug::T
    iterwrap::V
end
ShuffleIterConfig(;batchsize=32, seed=123, dataaug=identity, iterwrap=GpuIterator) = ShuffleIterConfig(batchsize, seed, dataaug, iterwrap)
dataiter(s::ShuffleIterConfig, x, y) = dataiter(x, y, s.batchsize, s.seed, s.dataaug) |> s.iterwrap

"""
    struct TrainIterConfig{T}
    TrainIterConfig(nbatches_per_gen, baseconfig)
    TrainIterConfig(;nbatches_per_gen=400, baseconfig=ShuffleIterConfig())

Standard training strategy for creating a batched iterators from data. Data from `baseconfig` is cycled in partitions of `nbatches_per_gen` each generation using a [`RepeatPartitionIterator`](@ref).

"""
struct TrainIterConfig{T}
    nbatches_per_gen::Int
    baseconfig::T
end
TrainIterConfig(;nbatches_per_gen=400, baseconfig=ShuffleIterConfig()) = TrainIterConfig(nbatches_per_gen,baseconfig)
function dataiter(s::TrainIterConfig, x, y)
    baseiter = dataiter(s.baseconfig, x, y)
    return RepeatPartitionIterator(baseiter, s.nbatches_per_gen)
end

function dataiter(x::AbstractArray,y::AbstractArray{T, 1}, args...) where T
    _dataiter(x, Flux.onehotbatch(y, sort(unique(y))), args...)
end
dataiter(x::AbstractArray,y::AbstractArray{T, 2}, args...) where T = _dataiter(x,y, args...)

function _dataiter(x::AbstractArray,y::AbstractArray, bs::Integer, seed::Integer, xwrap, ywrap = identity)
    biter = shuffleiter((x,y), bs, seed)

    # iterates feature/label pairs
    # Cycle is pretty important here. It is what makes it so that each generation sees a new shuffle
    # This is also what makes SeedIterators useful with random data augmentation:
    #   - Even if SeedIterator makes the augmentation itself the same, it will be applied to different images
    # Something ought to be done about the inconsistency vs BatchedIterConfig through...
    citer = Iterators.map(biter) do (x, y)
        xwrap(x), ywrap(y)
    end |> Iterators.cycle
    # Ensures all models see the exact same examples, even when a RepeatStatefulIterator restarts iteration.
    return SeedIterator(citer;rng=biter.rng)
end
shuffleiter(data, batchsize, seed) = BatchIterator(data, batchsize; shuffle=MersenneTwister(seed))

function _dataiter(x::AbstractArray, y::AbstractArray, bs::Integer, xwrap, ywrap = identity)
    biter = BatchIterator((x,y), bs)
    # iterates feature/label pairs
    return Iterators.map(biter) do (x, y)
        xwrap(x), ywrap(y)
    end
end

"""
    struct GlobalOptimizerMutation{S<:AbstractEvolutionStrategy, F} <: AbstractEvolutionStrategy
    GlobalOptimizerMutation(base::AbstractEvolutionStrategy)
    GlobalOptimizerMutation(base::AbstractEvolutionStrategy, optfun)

Maps the optimizer of each candidate in a population through `optfun` (default `randomlrscale()`).

Basically a thin wrapper for `global_optimizer_mutation`.

Useful for applying the same mutation to every candidate, e.g. global learning rate schedules which all models follow.
"""
struct GlobalOptimizerMutation{S<:AbstractEvolutionStrategy, F} <: AbstractEvolutionStrategy
    base::S
    optfun::F
end
GlobalOptimizerMutation(base::AbstractEvolutionStrategy) = GlobalOptimizerMutation(base, NaiveGAflux.randomlrscale())

function evostrategy_internal(s::GlobalOptimizerMutation, inshape)
    base = evostrategy_internal(s.base, inshape)
    return AfterEvolution(base, pop -> NaiveGAflux.global_optimizer_mutation(pop, s.optfun))
end

"""
    struct EliteAndSusSelection <: AbstractEvolutionStrategy
    EliteAndSusSelection(popsize, nelites, evolve)
    EliteAndSusSelection(;popsize=50, nelites=2, evolve = crossovermutate())

Standard evolution strategy.

Selects `nelites` candidates to move on to the next generation without any mutation.

Also selects `popsize - nelites` candidates out of the whole population using [`SusSelection`](@ref) to evolve by applying random mutation.

Mutation operations are both applied to the model itself (change sizes, add/remove vertices/edges) as well as to the optimizer (change learning rate and optimizer algorithm).

Finally, models are renamed so that the name of each vertex of the model of candidate `i` is prefixed with "model`i`".
"""
struct EliteAndSusSelection{F} <: AbstractEvolutionStrategy
    popsize::Int
    nelites::Int
    evolve::F
end
EliteAndSusSelection(;popsize=50, nelites=2, evolve=crossovermutate()) = EliteAndSusSelection(popsize, nelites, evolve)

function evostrategy_internal(s::EliteAndSusSelection, inshape)
    elite = EliteSelection(s.nelites)

    evolve = SusSelection(s.popsize - s.nelites, EvolutionChain(ShuffleCandidates(), s.evolve(inshape)))

    combine = CombinedEvolution(elite, evolve)
    return AfterEvolution(combine, rename_models ∘ clear_redundant_vertices)
end

"""
    struct EliteAndTournamentSelection <: AbstractEvolutionStrategy
    EliteAndTournamentSelection(popsize, nelites, k, p, evolve)
    EliteAndTournamentSelection(;popsize=50, nelites=2; k=2, p=1.0, evolve = crossovermutate())

Standard evolution strategy.

Selects `nelites` candidates to move on to the next generation without any mutation.

Also selects `popsize - nelites` candidates out of the whole population using [`TournamentSelection`](@ref) to evolve by applying random mutation.

Mutation operations are determined by `evolve` both applied to the model itself (change sizes, add/remove vertices/edges) as well as to the optimizer (change learning rate and optimizer algorithm).

Finally, models are renamed so that the name of each vertex of the model of candidate `i` is prefixed with "model`i`".
"""
struct EliteAndTournamentSelection{R,F} <: AbstractEvolutionStrategy
    popsize::Int
    nelites::Int
    k::Int
    p::R
    evolve::F
end
EliteAndTournamentSelection(;popsize=50, nelites=2, k=2, p=1.0, evolve = crossovermutate()) = EliteAndTournamentSelection(popsize, nelites, k, p, evolve)

function evostrategy_internal(s::EliteAndTournamentSelection, inshape)
    elite = EliteSelection(s.nelites)

    evolve = TournamentSelection(s.popsize - s.nelites, s.k, s.p, s.evolve(inshape))

    combine = CombinedEvolution(elite, evolve)
    return AfterEvolution(combine, rename_models ∘ clear_redundant_vertices)
end

"""
    crossovermutate(;pcrossover=0.3, pmutate=0.9)

Return a function which creates an [`EvolutionChain`](@ref) when called with an inputshape.

Crossover will be applied with a probability of `pcrossover` while mutation will be applied with a probability of `pmutate`. Note that these probabilities apply to models only, not optimizers.

Crossover is done using [`CrossoverSwap`](@ref) for models and [`LearningRateCrossover`](@ref) and [`OptimizerCrossover`](@ref) for optimizers.

Mutation is applied both to the model itself (change sizes, add/remove vertices/edges) as well as to the optimizer (change learning rate and optimizer algorithm).
"""
crossovermutate(;pcrossover=0.3, pmutate=0.9) = function(inshape)
    cross = candidatecrossover(pcrossover)
    crossoverevo = AfterEvolution(PairCandidates(EvolveCandidates(cross)), align_vertex_names)

    mutate = candidatemutation(pmutate, inshape)
    mutationevo = EvolveCandidates(mutate)

    return EvolutionChain(crossoverevo, mutationevo)
end

candidatemutation(p, inshape) = evolvemodel(MutationProbability(graphmutation(inshape), p), optmutation())
candidatecrossover(p) = evolvemodel(MutationProbability(graphcrossover(), p), optcrossover())

function clear_redundant_vertices(pop)
    foreach(cand -> NaiveGAflux.model(check_apply, cand), pop)
    return pop
end

function align_vertex_names(pop)
    for i in eachindex(pop)
        m = match(r"model(\d+).\.*", modelname(pop[i]))
        if !isnothing(m)
            pop[i] = rename_model(m.captures[1], pop[i])
        end
    end
    return pop
end

function rename_models(pop)
    for i in eachindex(pop)
        pop[i] = rename_model(i, pop[i])
    end
    return pop
 end

function rename_model(i, cand)
    # No need to copy layers now
    return fmap(cand; walk=(f, x) -> x isa NaiveNASflux.AbstractMutableComp ? x : Functors._default_walk(f, x)) do x
        x isa String ? replace(x, r"^model\d+\.*" => "model$i.") : x
    end
end

function optcrossover(poptswap=0.3, plrswap=0.4)
    lrc = MutationProbability(LearningRateCrossover(), plrswap) |> OptimizerCrossover
    oc = MutationProbability(OptimizerCrossover(), poptswap) |> OptimizerCrossover
    return MutationChain(lrc, oc)
end

function optmutation(p=0.1)
    lrm = LearningRateMutation()
    om = MutationProbability(OptimizerMutation([Descent, Momentum, Nesterov, ADAM, NADAM, ADAGrad]), p)
    return MutationChain(lrm, om)
end

function graphcrossover()
    vertexswap = LogMutation(((v1,v2)::Tuple) -> "Crossover swap between $(name(v1)) and $(name(v2))", CrossoverSwap(0.05; mergefun=NaiveGAflux.default_mergefun(;layerfun=ActivationContributionLow)))
    crossoverop = VertexCrossover(MutationProbability(vertexswap, 0.05), 0.05)
    return LogMutation(((g1,g2)::Tuple) -> "Crossover between $(modelname(g1)) and $(modelname(g2))", crossoverop)
end

function graphmutation(inshape)
    acts = default_actfuns()

    change_nout = NoutMutation(-0.2, 0.2) # Max 20% change in output size
    decrease_nout = NoutMutation(-0.2, 0)
    add_vertex = add_vertex_mutation(acts)
    add_downsampling = AddVertexMutation(ConditionalArchSpace(candownsample(inshape), downsamplingspace(default_layerconf(); outsizes = 2 .^(4:9), activations=acts)))
    rem_vertex = RemoveVertexMutation()
    # [-2, 2] keeps kernel size odd due to CuArrays issue# 356 (odd kernel size => symmetric padding)
    change_kernel = KernelSizeMutation(ParSpace2D([-2, 2]), maxsize=maxkernelsize(inshape))
    decrease_kernel = KernelSizeMutation(ParSpace2D([-2]))
    mutate_act = ActivationFunctionMutation(acts)

    add_edge = AddEdgeMutation(0.1; mergefun=NaiveGAflux.default_mergefun(;layerfun=ActivationContributionLow))
    rem_edge = RemoveEdgeMutation()

    # Create a shorthand alias for MutationProbability
    mpn(m, p) = VertexMutation(MutationProbability(m, p))
    mph(m, p) = VertexMutation(HighUtilityMutationProbability(m, p))
    mpl(m, p) = VertexMutation(LowUtilityMutationProbability(m, p))

    cnout = mpn(LogMutation(v -> "Mutate output size of vertex $(name(v))", change_nout), 0.05)
    dnout = mpl(LogMutation(v -> "Reduce output size of vertex $(name(v))", decrease_nout), 0.05)
    maddv = mph(LogMutation(v -> "Add vertex after $(name(v))", add_vertex), 0.005)
    maddd = mpn(LogMutation(v -> "Add downsampling after $(name(v))", add_downsampling), 0.01)
    mremv = mpl(LogMutation(v -> "Remove vertex $(name(v))", rem_vertex), 0.01)
    ckern = mpn(MutationFilter(allowkernelmutation, LogMutation(v -> "Mutate kernel size of $(name(v))", change_kernel)), 0.02)
    dkern = mpn(MutationFilter(allowkernelmutation, LogMutation(v -> "Decrease kernel size of $(name(v))", decrease_kernel)), 0.02)
    mactf = mpl(LogMutation(v -> "Mutate activation function of $(name(v))", mutate_act), 0.005)
    madde = mph(LogMutation(v -> "Add edge from $(name(v))", add_edge), 0.02)
    mreme = mpn(MutationFilter(v -> length(outputs(v)) > 1, LogMutation(v -> "Remove edge from $(name(v))", rem_edge)), 0.02)

    mremv = MutationFilter(g -> nvertices(g) > 5, mremv)

    # Create two possible mutations: One which is guaranteed to not increase the size:
    dsize = MutationChain(mremv, dnout, dkern, mreme, maddd)
    # ...and another which is not guaranteed to decrease the size
    csize = MutationChain(mremv, cnout, ckern, mreme, madde, maddd, maddv)
    # Add mutation last as new vertices with neuronutility == 0 screws up outputs selection as per https://github.com/DrChainsaw/NaiveNASlib.jl/issues/39

    mgp = VertexMutation(MutationProbability(LogMutation(v -> "Mutate global pool type for $(name(v))", MutateGlobalPool()), 0.1), SelectGlobalPool())

    # If isbig then perform the mutation operation which is guaranteed to not increase the size
    # This is done mostly to avoid OOM and time outs.
    mall = MutationChain(MutationFilter(isbig, dsize), MutationFilter(!isbig, csize), mactf, mgp)

    return LogMutation(g -> "Mutate $(modelname(g))", mall)
end

isbig(g) = nparams(g) > 20e7

candownsample(inshape) = v -> candownsample(v, inshape)
candownsample(v::AbstractVertex, inshape) = is_convtype(v) && !infork(v) && canshrink(v, inshape)

function canshrink(v, inshape)
    # Note assumes stride = 2!
    # Also assumes single output after a global pool and flatten
    allvs = all_in_graph(v)
    gpv = allvs[findfirst(is_globpool, allvs)] |> inputs |> first
    return all(fshape(shapetrace(gpv) |> squashshapes, inshape) .> 1)
end

function infork(v, inputcnt = Dict{AbstractVertex, Int}(inputs(v) .=> 1), seen = Set())
    v in seen && return any(x -> x < 0, values(inputcnt))
    push!(seen, v)

    # How many times do we expect to find v as input to some subsequent vertex if we did not start inside a fork?
    inputcnt[v] = get(inputcnt, v, 0) + length(outputs(v))

    for vi in inputs(v)
        # And now subtract by 1 each time we find it
        inputcnt[vi] = get(inputcnt, vi, 0) - 1
    end

    foreach(vo -> infork(vo, inputcnt, seen), outputs(v))
    # If we have any "unaccounted" for inputs when we hit the output then we started inside a fork
    return any(x -> x < 0, values(inputcnt))
end

maxkernelsize(inshape) = v -> maxkernelsize(v, inshape)
function maxkernelsize(v::AbstractVertex, inshape)
    ks = fshape(shapetrace(v) |> squashshapes, inshape)
    # Kernel sizes must be odd due to CuArrays issue# 356 (odd kernel size => symmetric padding)
    return @. ks - !isodd(ks)
end

allowkernelmutation(v) = allowkernelmutation(NaiveNASflux.layertype(v), v)
allowkernelmutation(l, v) = false
allowkernelmutation(::NaiveNASflux.FluxConvolutional{N}, v) where N = all(isodd, size(NaiveNASflux.weights(layer(v)))[1:N])

function add_vertex_mutation(acts)

    function outselect(vs)
        rss = randsubseq(vs, 0.5)
        return isempty(rss) ? [rand(vs)] : rss
    end

    wrapitup(as) = AddVertexMutation(rep_fork_res(as, 1,loglevel=Logging.Info), outselect)

    add_conv = wrapitup(convspace(default_layerconf(),2 .^(4:9), 1:2:5, acts,loglevel=Logging.Info))
    add_dense = wrapitup(LoggingArchSpace(VertexSpace(default_layerconf(), NamedLayerSpace("dense", DenseSpace(2 .^(4:9), acts)));level = Logging.Info))

    return MutationChain(MutationFilter(is_convtype, add_conv), MutationFilter(!is_convtype, add_dense))
end

is_convtype(v::AbstractVertex) = any(is_globpool.(outputs(v))) || any(is_convtype.(outputs(v)))
is_globpool(v::AbstractVertex) = is_globpool(layer(v))
is_globpool(f) = f isa NaiveGAflux.GlobalPool

struct SelectGlobalPool <:AbstractVertexSelection
    s::AbstractVertexSelection
end
SelectGlobalPool() = SelectGlobalPool(AllVertices())
NaiveGAflux.select(s::SelectGlobalPool, g::CompGraph, ms...) = filter(is_globpool, NaiveGAflux.select(s.s, g, ms...))

struct MutateGlobalPool <: AbstractMutation{AbstractVertex} end
function (::MutateGlobalPool)(v::AbstractVertex)
    # Remove the old vertex and replace it with a new one with a different global pool type
    iv = inputs(v)[]
    RemoveVertexMutation()(v)
    v in outputs(iv) && return # failed to remove for some reason

    newname = first(splitext(name(v)))
    insert!(iv, vx -> GlobalPoolSpace(newpool(layer(v)))(newname, vx))
end

newpool(::GlobalPool{MaxPool}) = MeanPool
newpool(::GlobalPool{MeanPool}) = MaxPool


"""
    function tb_train!(mdl, opti, trn, vld=nothing; epochs=1,
                      lr_decay=nothing, lrd_freq=1, lrd_linear=false,
                      l2=0.0,
                      eval_size=0.2, eval_freq=1,
                      acc_fun=nothing,
                      mb_loss_freq=100,
                      cp_freq=nothing, cp_dir="checkpoints",
                      tb_dir="logs", tb_name="run",
                      tb_text=\"\"\"Description of tb_train!() run.\"\"\",
                      opti_args...)

Train function with TensorBoard integration. TB logs are written with
the TensorBoardLogger.jl package.
The model is updated (in-place) and the trained model is returned.

### Arguments:
+ `mdl`: model; i.e. forward-function for the net
+ `opti`: Knet-stype optimiser type
+ `trn`: training data; iterator to provide (x,y)-tuples with
        minibatches
+ `vld`: validation data; iterator to provide (x,y)-tuples with
        minibatches. Set to `nothing`, if not defined.

### Keyword arguments:
#### Optimiser:
+ `epochs=1`: number of epochs to train
+ `lr_decay=nothing`: do a leraning rate decay if not `nothing`:
        the value given is the final learning rate after `epochs*lrd_freq`
        steps of decay. `lr_decay` is only applied if both start learning rate
        `lr` and final learning rate `lr_decay` are defined explicitly.
        Example: `lr=0.01, lr_decay=0.001` will reduce the lr from
        0.01 to 0.001 during the training.
+ `lrd_freq=1`: frequency of learning rate decay steps. Default is
        to modify the lr after every epoch
+ `lrd_linear=false`: type of learning rate decay;
        If `false`, lr is modified
        by a constant factor (e.g. 0.9) resulting in an exponential decay.
        If `true`, lr is modified by the same step size, i.e. linearly.
+ `l2=0.0`: L2 regularisation; implemented as weight decay per
        parameter
+ `opti_args...`: optional keyword arguments for the optimiser can be specified
        (i.e. `lr`, `gamma`, ...).

#### Model evaluation:
+ `eval_size=0.2`: fraction of validation data to be used for calculating
        loss and accuracy for train and validation data during training.
+ `eval_freq=1`: frequency of evaluation; default=1 means evaluation is
        calculated after each epoch. With eval_freq=10 eveluation is
        calculated 10 times per epoch.
+ `acc_fun=nothing`: function to calculate accuracy. The function
        must implement the following signature: `fun(model; data)` where
        data is an iterator that provides (x,y)-tuples of minibatches.
        For classification tasks, `accuracy` from the Knet package is
        a good choice. For regression a correlation or mean error
        may be used.
+ `mb_loss_freq=100`: frequency of training loss reporting. default=100
        means that 100 loss-values per epoch will be logged to TensorBoard.
        If mb_loss_freq is greater then the number of minibatches,
        loss is logged for each minibatch.
+ `cp_freq=nothing`: frequency of model checkpoints written to disk.
        Default is `nothing`, i.e. no checkpoints are written.
        To write the model after each epoch with
        name `model` use freq=1; to write every 2 epochs freq=0.5.
+ `cp_dir="checkpoints"`: directory for checkpoints

#### TensorBoard:
TensorBoard log-directory is created from 3 parts:
`tb_dir/tb_name/<current date time>`.

+ `tb_dir="logs"`: root directory for tensorborad logs.
+ `tb_name="run"`: name of training run. `tb_name` will be used as
        directory name and should not include whitespace
+ `tb_text`:  description
        to be included in the TensorBoard log as *text* log.
"""
function tb_train!(mdl, opti, trn, vld=nothing; epochs=1,
                  lr_decay=nothing, lrd_freq=1, lrd_linear=false,
                  l2=0.0,
                  eval_size=0.2, eval_freq=1, acc_fun=nothing,
                  mb_loss_freq=100,
                  cp_freq=nothing, cp_dir="checkpoints",
                  tb_dir="logs", tb_name="run",
                  tb_text="""Description of tb_train!() run.""",
                  opti_args...)

    # use every n-th mb for evaluation (based on vld if defined):
    #
    n_trn = length(trn)

    if vld == nothing
        n_vld = 0
        eval_vld = nothing
        n_eval = Int(ceil(n_trn * eval_size))
    else
        n_vld = length(vld)
        n_eval = Int(ceil(n_vld * eval_size))
        nth_vld = cld(n_vld, n_eval)
        eval_vld = takenth(vld, nth_vld)
    end
    nth_trn = cld(n_trn, n_eval)
    eval_trn = takenth(trn, nth_trn)

    # frequencies of actions during training:
    #
    eval_nth = cld(n_trn, eval_freq)
    mb_loss_nth = cld(n_trn, mb_loss_freq)
    lr_nth = cld(n_trn, lrd_freq)
    if cp_freq != nothing
        cp_nth = cld(n_trn, cp_freq)
    end

    # do lr-decay only if lr is explicitly defined:
    #
    if lr_decay != nothing && haskey(opti_args, :lr)
       lr_decay = calc_d_η(opti_args[:lr], lr_decay, lrd_linear, lrd_freq*epochs)
    else
        lr_decay = nothing
    end

    # mk log directory:
    #
    start_time = Dates.now()
    tb_log_dir = joinpath(pwd(), tb_dir, tb_name,
                    Dates.format(start_time, "yyyy-mm-ddTHH-MM-SS"))
    println("Training $epochs epochs with $n_trn minibatches/epoch")
    if vld != nothing
        println("    (and $n_vld validation mbs).")
    end
    println("Evaluation is performed every $eval_nth minibatches (with $n_eval mbs).")
    println("Watch the progress with TensorBoard at: $tb_log_dir")

    # Tensorboard logger:
    #
    tbl = TensorBoardLogger.TBLogger(tb_log_dir,
                    min_level=Logging.Info)
    log_text(tbl, tb_log_dir, tb_name, start_time, tb_text,
             opti, trn, vld, epochs,
             lr_decay, lrd_freq, l2,
             cp_freq, opti_args)
    calc_and_report_loss(mdl, eval_trn, eval_vld, tbl, 0)

    if acc_fun != nothing
        calc_and_report_acc(mdl, acc_fun, eval_trn, eval_vld, tbl, 0)
    end

    # set optimiser - only if not yet set:
    #
    if first(params(mdl)).opt == nothing
        for p in params(mdl)
            p.opt = opti(;opti_args...)
        end
    else
        if haskey(opti_args, :lr)
            set_learning_rate(mdl, opti_args[:lr])
        end
    end

    # Training:
    #
    mb_losses = Float32[]
    @showprogress for (i, (x,y)) in enumerate(ncycle(trn,epochs))

        loss = @diff mdl(x,y)
        mb_loss = value(loss)

        if isnan(mb_loss)
            println("ERROR: training aborted because of loss value NaN!")
            break
        end

        for p in params(loss)
            Δw = grad(loss, p) + p .* l2
            # println("updating $i: $(p.opt.lr), Δw: -")
            Knet.update!(p, Δw)
        end

        # TensorBoard:
        #
        # println("mb_loss: $mb_loss"); flush(stdout)
        push!(mb_losses, mb_loss)
        if (i % eval_nth) == 0
            calc_and_report_loss(mdl, eval_trn, eval_vld, tbl, eval_nth)

            if acc_fun != nothing
                calc_and_report_acc(mdl, acc_fun, eval_trn, eval_vld,
                                    tbl, eval_nth)
            end
        end

        if (i % mb_loss_nth) == 0
            TensorBoardLogger.log_value(tbl,
                    "Minibatch loss (epoch = $n_trn steps)",
                    mean(mb_losses), step=i)
            # println("mb_loss-mean: $(mean(mb_losses))"); flush(stdout)
            mb_losses = Float32[]
        end

        # checkpoints:
        #
        if (cp_freq != nothing) && (i % cp_nth) == 0
            write_cp(mdl, i, tb_log_dir)
        end

        # lr decay:
        #
        if (lr_decay != nothing) && (i % lr_nth == 0)
            lr = first(params(mdl)).opt.lr
            lr = lrd_linear ? lr + lr_decay : lr * lr_decay
            @printf("Setting learning rate to η=%.2e\n", lr)
            set_learning_rate(mdl, lr)
            # for p in params(mdl)
            #     p.opt.lr = lr
            #     # println("adapting lr in $p to $(p.opt.lr)")
            # end
        end
    end

    println("Training finished with:")
    println("Training loss:       $(calc_loss(mdl, data=trn))")
    if acc_fun != nothing
        println("Training accuracy:   $(calc_acc(mdl, acc_fun, data=trn))")
    end

    if vld != nothing
        println("Validation loss:     $(calc_loss(mdl, data=vld))")
        if acc_fun != nothing
            println("Validation accuracy: $(calc_acc(mdl, acc_fun, data=vld))")
        end
    end

    # save final model:
    #
    if (cp_freq != nothing)
        write_cp(mdl, n_trn*epochs+1, tb_log_dir)
    end
    return mdl
end

function set_learning_rate(mdl, lr)

    for p in params(mdl)
        p.opt.lr = lr
    end
end

function log_text(tbl, tb_log_dir, tb_name, start_time, tb_text,
                  opti, trn, vld, epochs,
                  lr_decay, lrd_freq, l2,
                  cp_freq, opti_args)

    if vld == nothing
        vld = []
    end

    tb_log_text =
    "
NNHelferlein.jl tb_train!() log
" *
    "
 " *
    "
dir:  $tb_log_dir
 " *
    "
name: $tb_name
 " *
    "
time: $(Dates.format(start_time, "E, yyyy/mm/dd, HH:MM:SS"))
 " *
    "
 " *
    "
 " *
    "   $tb_text " *
    "
 " *
    "
Training parameters:
" *
    "
" *
    "
Training minibatches: $(length(trn))
" *
    "
Validation minibatches: $(length(vld))
" *
    "
Optimiser: $opti
" *
    "
Epochs: $epochs
"

    if lr_decay != nothing
        tb_log_text *= "
learning rate step size: $lr_decay with frequency $lrd_freq
"
    end
    if l2 > 0
        tb_log_text *= "
L2 regularisation: $l2
"
    end
    if cp_freq != nothing
        tb_log_text *= "
Checkpoints are saved $cp_freq times per epoch
"
    end

    for arg in keys(opti_args)
        tb_log_text *= "
$arg: $(opti_args[arg])
"
    end
    tb_log_text *= "
"

    TensorBoardLogger.log_text(tbl, "Description", tb_log_text, step=0)
    # with_logger(tbl) do
    #         @info "Description" text=tb_log_text log_step_increment=0
    # end
end

function write_cp(model, step, dir)

    dir_name = joinpath(dir, "checkpoints")
    if !isdir(dir_name)
        mkdir(dir_name)
    end
    fname = joinpath(dir_name, "checkpoint_$step.jld2")
    JLD2.@save fname model
end

# Helper to calc loss and acc with only ONE forward run:
#
function calc_loss(mdl; data)

    loss = 0.0
    for (x,y) in data
        loss += mdl(x,y)
    end
    return loss/length(data)
end

function calc_acc(mdl, fun; data)

    # acc = 0.0
    # for (x,y) in data
    #     acc += fun(mdl(x), y)
    # end
    # return acc/length(data)
    return fun(mdl, data=data)
end

# Helper for TensorBoardLogger:
#
function calc_and_report_loss(mdl, trn, vld, tbl, step)

    loss_trn = calc_loss(mdl, data=trn)

    if vld != nothing
        loss_vld = calc_loss(mdl, data=vld)
        with_logger(tbl) do
            @info "Evaluation Loss" train=loss_trn valid=loss_vld log_step_increment=step
        end
    else
        with_logger(tbl) do
            @info "Evaluation Loss" train=loss_trn log_step_increment=0
        end
    end
end

function calc_and_report_acc(mdl, acc_fun, trn, vld, tbl, step)

    acc_trn = calc_acc(mdl, acc_fun, data=trn)

    if vld != nothing
        acc_vld = calc_acc(mdl, acc_fun, data=vld)
        with_logger(tbl) do
            @info "Evaluation Accuracy" train=acc_trn valid=acc_vld log_step_increment=step
        end
    else
        with_logger(tbl) do
            @info "Evaluation Accuracy" train=acc_trn log_step_increment=step
        end
    end
end

# calc step size for lr-decay based on
# start rate, end rate and num steps:
#
function calc_d_η(η_start, η_end, lrd_linear, steps)
    if !lrd_linear
        d = log(η_end/η_start) / steps
        d = exp(d)
    else
        d = (η_end - η_start) / steps
    end
    return d
end

"""
    function predict_top5(mdl, x; top_n=5, classes=nothing)

Run the model `mdl` for data in `x` and print the top 5
predictions as softmax probabilities.

### Arguments:
+ `top_n`: print top *n* hits
+ `classes`: optional list of human readable class labels.
"""
function predict_top5(mdl, x; top_n=5, classes=nothing)

    y = predict(mdl, x, softmax=false)

    if classes == nothing
        classes = repeat(["-"], size(y)[1])
    end
    for (i,o) in enumerate(eachcol(y))
        o = Knet.softmax(vec(Array(o)))
        top = sortperm(vec(Array(o)), rev=true)[1:top_n]
        println("top-$top_n hits for sample $i: $top"); flush(stdout)

        @printf("%6s  %6s   %s\n", "softmax", "#", "class label")
        for t in top
            @printf(" %6.2f  %6i   \"%s\"\n", o[t], t, classes[t])
        end
        println(" ")
    end
    return y
end

"""
    function predict(mdl, x; softmax=false)

Return the prediction for x.

### Arguments:
+ `mdl`: executable network model
+ `x`: iterator providing minibatches
        of input data
+ `softmax`: if true and if model is a `::Classifier` the predicted
        softmax probabilities are returned instead of raw
        activations.
"""
function predict(mdl, x; softmax=false)

    if x isa AbstractArray
        y = mdl(x)
    else
        # ys = [mdl(i) for i in x]
        y = cat((mdl(i) for i in x)..., dims=2)
    end
    y = convert(Array{Float32}, y)

    if softmax && mdl isa Classifier
        return Knet.softmax(y)
    else
        return y
    end
end

"""
    function hamming_dist(p, t; accuracy=false, vocab=nothing, pad=0)

Return the Hamming distance between two sequences or two minibatches
of sequences. Predicted sequences `p` and teaching input sequences `t`
may be of different length but the number of sequences in the minibatch
must be the same.

### Arguments:
+ `p`, `t`: n-dimensional arrays of type `Int` with predictions
        and teaching input for a minibatch of sequences.
        Shape of the arrays must be identical except of the first dimension
        (i.e. the sequence length) that may differ between `p` and `t`.
+ `accuracy=false`: if `false`, the mean Hamming distance in the minibatch
        is returned (i.e. the average number of differences in the sequences).
        If `true`, the accuracy is returned
        for all not padded positions in a range (0.0 - 1.0).
+ `vocab=nothing`: target laguage vocabulary of type `NNHelferlein.WordTokenizer`.
        If defined,
        the padding token of `vocab` is used to mask all control tokens in the
        sequences (i.e. '<start>, <end>, <unknwon>, <pad>').
+ `pad=0`: if `vocab` is undefined, `pad` is used to pad `p`, if the sequence
        length of `p` is smaller than the length of `t`.
"""
function hamming_dist(p, t; accuracy=false, vocab=nothing, pad=0)

    # make 2d matrix of sequences:
    #
    n_seq_t = size(t)[1]
    t = reshape(t, n_seq_t,:)

    n_seq_p = size(p)[1]
    p = reshape(p, n_seq_p,:)

    n_mb = size(t)[2]

    # make all control-tokens the same:
    #
    if vocab != nothing
        START = vocab("<start>")
        END = vocab("<end>")
        UNK = vocab("<unknown>")
        PAD = vocab("<pad>")

        t[t .== START] .= PAD
        t[t .== END] .= PAD
        t[t .== UNK] .= PAD
    else
        PAD = pad
    end

    # make n_seq of p the same as t:
    #
    if n_seq_p > n_seq_t
        p = p[1:n_seq_t,:]
    end
    while size(p)[1] < n_seq_t
        p = vcat(p, repeat([PAD], inner=(1,n_mb)))
    end

    # mask preds same as teaching and count all
    # mask positions:
    #
    p[t .== PAD] .= PAD
    num_pad = length(t[t .== PAD])
    return accuracy ? (sum(p .== t) - num_pad)/(length(t)-num_pad) : sum(p .!= t)/n_mb
end


## Main loop to extract in e2e the DBSCAN parameters.
## The list of votable is selected directly from the directory

using DataFrames, Query
using CSV, Glob, Dates
using Statistics, Random
using Printf

rootdir =  ENV["GAIA_ROOT"]

push!(LOAD_PATH,"$rootdir/run/src")
using GaiaClustering

## directory
wdir    = "$rootdir/e2e_products"
votdir  = "$wdir/votable"
plotdir = "$wdir/plotSelect"
ocdir   = "$wdir/oc"

## load a liist of votable and update the file if done
## add results
##

function readlist_votable(filelist::String)
    vot = DataFrames.copy(CSV.read(filelist))
    return(vot)
end

function getdata(filevot)
    voname = filevot

    data       = read_votable(voname)
    df         = filter_data(data, [0,2100])
    dfcart     = add_cartesian(df)
    blck       = [[1,2,3],[4,5], [6,7,8]]
    wghtblck   = [4.0,5.0,1.0]
    norm       = "identity"

    dfcartnorm , scale8 = normalization_PerBlock(dfcart, blck, wghtblck , norm, false)
    return(df, dfcart , dfcartnorm)
end

## to check if done and record
## check and updt if votname analyzed. If not done return false
function updt_votcompleted(fileres, votname , cycletot=1, flag= 0 , onlycheck=true)
    let
        if  onlycheck
            if !isfile(fileres)
                return(0, false)
            else
                res = DataFrames.copy(CSV.read(fileres, delim=";"))
                if votname in res.votname
                    x = @from i in res begin
                        @where i.votname == votname
                        @select i
                        @collect DataFrame
                    end
                    return(x, true)
                else
                    return(0, false)
                end
            end
            ## UPDTE
        else
            if !isfile(fileres)
                res = DataFrame(votname=votname, cycle=cycletot, flag=flag)
                CSV.write(fileres,res,delim=';')
                println("## $fileres created...")
                return(res,true)
            else
                res = DataFrames.copy(CSV.read(fileres, delim=";"))
                newrow = DataFrame(votname=votname,cycle=cycletot, flag=flag)
                append!(res,newrow)
                CSV.write(fileres,res,delim=';')
                return(res,true)
            end
        end
    end
end
## check for blacklist
##
function read_blacklist(blackname)
    if isfile(blackname)
        df= DataFrames.copy(CSV.read(blackname, delim=";"))
        blacklist= df.votname
    else
        blacklist= [""]
    end

    return(blacklist)
end

#######################
## Main loop
##
function main(filelist, metafile)
    let
        wd= pwd() ; nfile= size(filelist)[1] ; totalTime= 0.
        println("## Starting main loop using optimization with cycles...")
        println("## It can be very long but will be resumed to the last reduced file.")
        println("## Working directory: $wd")
        println("## $nfile files to analyze...")

        m= read_params(metafile, false)
        m.plotdir= plotdir
        m.ocdir= ocdir
        fileres= "$(m.prefile).done.csv"

        # read a possible votname blacklist
        blackname= "blacklist.csv"
        blacklist= read_blacklist(blackname)
        println("## blacklist read...")

        for i in 1:nfile
            votname = filelist[i]
            res, votfound= updt_votcompleted(fileres, votname, 0, 0, true)

            ## test blacklist
            if votname in blacklist
                println("##")
                println("## $votname in blacklist...")
                println("##")
                votfound= true
            end

            if !votfound
                tstart= now()
                println("###########################")
                println("## Starting with $votname")
                println("## Starting at $tstart")

                df , dfcart , dfcartnorm = getdata(votdir*"/"*votname)
                m.votname= votname
                cycle, flag= cycle_extraction(df, dfcart, m)
                res,  votfound= updt_votcompleted(fileres, votname , cycle, flag, false)

                tend= now()
                println("## Ending at $tend")
                println("## number of cycle: $cycle , flag:$flag ")
                println("##########################")
                println("##")

                duration= Dates.value(tend-tstart) / (1000*3600)
                totalTime += duration
                meanTime= totalTime / i
                ETA= meanTime * (nfile-i) / 24
                nleft= nfile-i
                ETAstr= @sprintf("%3.3f", ETA) ; durationstr= @sprintf("%3.3f", duration)
                @printf("## %s \n",specialstr("Duration: $durationstr hours","YELLOW"))
                @printf("## %s \n",specialstr("ETA: $ETAstr days","YELLOW"))
                @printf("## %s \n",specialstr("Votable done: $votname","YELLOW"))
                @printf("## %s \n",specialstr("Files analyzed: $i","YELLOW"))
                @printf("## %s \n",specialstr("Files to go: $nleft","YELLOW"))
                println("##\n##")

            end
        end
    end
    println("## Main loop done.")
end
###############################################################################
header_extract()

cd(votdir)
votlist= glob("*.vot")
cd(wdir)

main(votlist,"configAll.ext")


module CoupledConditionalSMC

using RNGPool
import Statistics.mean
using Random

include("structures.jl")
include("subroutines.jl")
include("algorithms.jl")
include("unbiased.jl")
include("couplingTime.jl")
include("jls.jl")

export CCSMCIO, ccXpf!, initializeCCSMC

end # module


function checksizes(data, layout)
  for rng in layout
    size_range = size(rng)
    el = data[_data_id(rng)]
    size_el = size(el)
    if size_range != size_el
      error(
        """
        size mismatch
        rng: $rng
        el: $el
        size_range: $size_range, size_el: $size_el\n
        """)
    end
  end
  return nothing
end
type Table
  data::Vector
  layout
  styles::Vector{AbstractStyleElement}

  function Table(data, layout, styles)
    checksizes(data, layout)
    return new(data, layout, styles)
  end
end

function size(tbl::Table)
  n, m = 0, 0

  for rng in tbl.layout
    ncand = ymax(rng.box)
    mcand = xmax(rng.box)
    if n < ncand
      n = ncand
    end
    if m < mcand
      m = mcand
    end
  end
  return n, m
end


import Base:
  # Types
  Array,
  CartesianIndex,
  IndexStyle,
  Tuple,
  # Symbols
  +,
  -,
  *,
  # Methods
  checkbounds,
  complex,
  convert,
  conj,
  copy,
  copyto!,
  eachindex,
  eltype,
  fill,
  fill!,
  getindex,
  hash,
  isempty,
  isless,
  iterate,
  length,
  ndims,
  permutedims,
  permutedims!,
  promote_rule,
  randn,
  reshape,
  setindex,
  setindex!,
  show,
  size,
  stride,
  strides,
  summary,
  to_indices,
  unsafe_convert,
  view,
  zeros

import Base.Broadcast: Broadcasted, BroadcastStyle

import LinearAlgebra: exp

import TupleTools: isperm


#= NOTE: If running tests incrementally, import the following packages:
using Revise, Test
=#

import Cytof5.Util: similarity_FM, meanabsdiff, sumabsdiff

@testset "Util similarity_FM" begin
  I = 3
  J = 20
  K1 = 5
  K2 = 4

  W1 = rand(I, K1)
  W1 ./= sum(W1, dims=2)

  Z1 = Int.(rand(J, K1) .> .5)
  Z2 = Int.(rand(J, K2) .> .5)

  @time s = similarity_FM(Z1, W1, Z2)
  @assert s <= J * I

  @test true
end


module Data

using YAML: load_file
using Images
using Images.FileIO
using LMDB
using LMDB: create, open, close, put!, start

struct MainDataset
    config::Dict
    num_frames::Int
    txn_map::Dict
    dbi_map::Dict
    idx_map::Dict
    yaw_map::Dict
    file_map::Dict

    function MainDataset(data_dir::AbstractString, config_path::AbstractString)
        # read the config file
        num_frames = 0
        config = load_file(config_path)
        txn_map = Dict()
        dbi_map = Dict()
        idx_map = Dict()
        yaw_map = Dict()
        file_map = Dict()

        for full_path in readdir(data_dir, join=true)
            if !isdir(full_path)
                continue
            end
            env = create()
            open(env, full_path)
            txn = start(env)

            dbi = open(txn)

            n = parse(Int64, LMDB.get(txn, dbi, "len", String)) - config["num_plan"] # last num_plan labels are missing from the dataset
            if n < config["num_plan"]+1
                print(full_path, " is too small. consider deleting")
                close(env, dbi)
            else
                offset = num_frames
                for i in 0:n-config["num_plan"]-1
                    num_frames+=1
                    for j in 0:length(config["camera_yaws"])-1
                        txn_map[(offset+i)*length(config["camera_yaws"])+j] = txn
                        dbi_map[(offset+i)*length(config["camera_yaws"])+j] = dbi
                        idx_map[(offset+i)*length(config["camera_yaws"])+j] = i
                        yaw_map[(offset+i)*length(config["camera_yaws"])+j] = j
                        file_map[(offset+i)*length(config["camera_yaws"])+j] = full_path
                    end
                end
            end
        end
        print(data_dir,": ", num_frames, " frames ", "x ", length(config["camera_yaws"]))
        new(config, num_frames, txn_map, dbi_map, idx_map, yaw_map, file_map)
    end
end

function access(tag::AbstractString, txn::Transaction, dbi::DBI, index::Int, T::Int, dtype::Type; preprocess = x->x)
    try
        reduce(hcat, preprocess.([LMDB.get(txn, dbi, "$(tag)_$(lpad(t, 5, '0'))", dtype) for t in index:index+(T-1)]))
    catch
        @show tag
    end
end

function get_item(d::MainDataset, state::Int)
    T = d.config["num_plan"]
    idx = state-1
    if !d.config["multi_cam"]
        idx *= length(d.config["camera_yaws"])
    end

    lmdb_txn = d.txn_map[idx]
    lmdb_dbi = d.dbi_map[idx]
    index = d.idx_map[idx]
    cam_index = d.yaw_map[idx]

    locs = access("loc", lmdb_txn, lmdb_dbi, index, T+1, Vector{Float32})
    rots = access("rot", lmdb_txn, lmdb_dbi, index, T, Vector{Float32})
    spds = access("spd", lmdb_txn, lmdb_dbi, index, T, Vector{Float32})

    decode = (x -> FileIO.load(IOBuffer(x)))
    lbls = [access("lbl_$(lpad(l, 2, '0'))", lmdb_txn, lmdb_dbi, index+1, T, Vector{UInt8}; preprocess=decode) |> channelview for l in 0:11]
    wide_rgb = permuteddimsview(access("wide_$(cam_index)", lmdb_txn, lmdb_dbi, index, 1, Vector{UInt8}; preprocess=decode) |> channelview, (2,3,1))
    wide_sem = access("wide_sem_$(cam_index)", lmdb_txn, lmdb_dbi, index, 1, Vector{UInt8}; preprocess=decode) |> channelview
    narr_rgb = access("narr_$(cam_index)", lmdb_txn, lmdb_dbi, index, 1, Vector{UInt8}; preprocess=decode) |> channelview

    cmd = access("cmd", lmdb_txn, lmdb_dbi, index, 1, Float32)
    wide_rgb, wide_sem, narr_rgb, lbls, locs, rots, spds, trunc.(Int, cmd)
end

Base.length(d::MainDataset) = (if (d.config["multi_cam"]) d.num_frames*length(d.config["camera_yaws"]) else d.num_frames end)
Base.iterate(d::MainDataset, state=1) = state > length(d) ? nothing : (get_item(d, state), state+1)

end


@testset "getindex" begin
    @testset "getindex(::Matrix{<:Number},...)" begin
        x = [1.0 2.0 3.0; 10.0 20.0 30.0]
        x̄ = [1.4 2.5 3.7; 10.5 20.1 30.2]
        full_ȳ = [7.4 5.5 2.7; 8.5 11.1 4.2]

        @testset "single element" begin
            rrule_test(getindex, 2.3, (x, x̄), (2, nothing))
            rrule_test(getindex, 2.3, (x, x̄), (2, nothing), (1, nothing))
            rrule_test(getindex, 2.3, (x, x̄), (2, nothing), (2, nothing))

            rrule_test(getindex, 2.3, (x, x̄), (CartesianIndex(2, 3), nothing))
        end

        @testset "slice/index postions" begin
            rrule_test(getindex, [2.3, 3.1], (x, x̄), (2:3, nothing))
            rrule_test(getindex, [2.3, 3.1], (x, x̄), (3:-1:2, nothing))
            rrule_test(getindex, [2.3, 3.1], (x, x̄), ([3,2], nothing))
            rrule_test(getindex, [2.3, 3.1], (x, x̄), ([2,3], nothing))

            rrule_test(getindex, [2.3 3.1; 4.1 5.1], (x, x̄), (1:2, nothing), (2:3, nothing))
            rrule_test(getindex, [2.3 3.1; 4.1 5.1], (x, x̄), (:, nothing), (2:3, nothing))

            rrule_test(getindex, [2.3, 3.1], (x, x̄), (1:2, nothing), (1, nothing))
            rrule_test(getindex, [2.3, 3.1], (x, x̄), (1, nothing), (1:2, nothing))

            rrule_test(getindex, [2.3 3.1; 4.1 5.1], (x, x̄), (1:2, nothing), (2:3, nothing))
            rrule_test(getindex, [2.3 3.1; 4.1 5.1], (x, x̄), (:, nothing), (2:3, nothing))

            rrule_test(getindex, full_ȳ, (x, x̄), (:, nothing), (:, nothing))
            rrule_test(getindex, full_ȳ[:], (x, x̄), (:, nothing))
        end

        @testset "masking" begin
            rrule_test(getindex, full_ȳ, (x, x̄), (trues(size(x)), nothing))
            rrule_test(getindex, full_ȳ[:], (x, x̄), (trues(length(x)), nothing))

            mask = falses(size(x))
            mask[2,3] = true
            mask[1,2] = true
            rrule_test(getindex, [2.3, 3.1], (x, x̄), (mask, nothing))

            rrule_test(
                getindex, full_ȳ[1,:], (x, x̄), ([true, false], nothing), (:, nothing)
            )
        end

        @testset "By position with repeated elements" begin
            rrule_test(getindex, [2.3, 3.1], (x, x̄), ([2, 2], nothing))
            rrule_test(getindex, [2.3, 3.1, 4.1], (x, x̄), ([2, 2, 2], nothing))
            rrule_test(
                getindex, [2.3 3.1; 4.1 5.1], (x, x̄), ([2,2], nothing), ([3,3], nothing)
            )
        end
    end
end


using TestImages

real_imgs = [
    first(splitext(img_name)) => testimage(img_name)
    for img_name
    in TestImages.remotefiles
    if endswith(img_name, ".png")
]

@testset "TestImages.jl" begin
    for (case, image) in real_imgs
        @testset "$(case)" begin
            expected = collect(PNGFiles._prepare_buffer(image))
            fpath = joinpath(PNG_TEST_PATH, "test_img_$(case).png")
            @testset "write" begin
                @test PNGFiles.save(fpath, image) == 0
            end
            @testset "read" begin
                global read_in_pngf = PNGFiles.load(fpath)
                @test read_in_pngf isa Matrix
            end
            @testset "compare" begin
                @test all(expected .≈ read_in_pngf)
            end
            global read_in_immag = _standardize_grayness(ImageMagick.load(fpath))
            @testset "$(case): ImageMagick read type equality" begin
                # The lena image is Grayscale saved as RGB...
                @test eltype(_standardize_grayness(read_in_pngf)) == eltype(read_in_immag)
            end
            @testset "$(case): ImageMagick read values equality" begin
                # The lena image is Grayscale saved as RGB...
                imdiff_val = imdiff(_standardize_grayness(read_in_pngf), read_in_immag)
                onfail(@test imdiff_val < 0.01) do
                    PNGFiles._inspect_png_read(fpath)
                    _add_debugging_entry(fpath, case, imdiff_val)
                end
            end
            path, ext = splitext(fpath)
            newpath = path * "_new" * ext
            PNGFiles.save(newpath, read_in_pngf)
            @testset "$(case): IO is idempotent" begin
                @test all(read_in_pngf .≈ PNGFiles.load(newpath))
            end
        end
    end
end


using Interpolations: Flag, Quadratic, InPlace

export ItpSharedEV, set_up_dcdp_workers

struct ItpSharedEV{T,A1<:Interpolations.AbstractInterpolation{T},A2<:Interpolations.AbstractInterpolation{T},A3<:Interpolations.AbstractInterpolation{T},TT<:Tuple}
  EV::A1
  dEV::A2
  dEVσ::A3
  itypes::TT
end

function ItpSharedEV(σ::Real, ψflag::Flag=Quadratic(InPlace()) )
  global g_SharedEV
  global g_dcdp_primitives
  return ItpSharedEV(g_SharedEV, g_dcdp_primitives, σ, ψflag)
end

@GenGlobal g_ItpSharedEV

function ItpSharedEV(sev::SharedEV{T,N,N2,TT}, p::dcdp_primitives, σ::Real=1.0; ψflag::Flag=Quadratic(InPlace())) where {T,N,N2,TT}

    # dimensions of things
    ntyp = length(sev.itypes)
    nθ = _nθt(p)
    nS = _nS(p)
    nSexp1 = _nSexp(p)

    # form spline specifications
    splinetype(r::Union{StepRange{<:AbstractFloat},StepRangeLen{<:AbstractFloat}}) = BSpline(Quadratic(InPlace()))
    splinetype(r::AbstractVector) = NoInterp()
    splinetype(r::LinRange{<:AbstractFloat}) = BSpline(Constant())

    it_z     = splinetype.(p.zspace)
    it_itype = splinetype.(sev.itypes)

    it_EV   = (it_z..., BSpline(ψflag), NoInterp(),             it_itype...)
    it_dEV  = (it_z..., BSpline(ψflag), NoInterp(), NoInterp(), it_itype...)
    it_dEVσ = (it_z..., BSpline(ψflag), NoInterp(),             it_itype...)

    # form UNFILTERED spline objects.
    # These specify the SharedArray as the coef field and add type information about spline interpolation
    # if quadtratic, we will have to solve a system of linear equations later to interpolate
    itp_EV    = BSplineInterpolation(tweight(sev.EV  ), sev.EV  , it_EV  , OnCell(), Val{0}())
    itp_dEV   = BSplineInterpolation(tweight(sev.dEV ), sev.dEV , it_dEV , OnCell(), Val{0}())
    itp_dEVσ  = BSplineInterpolation(tweight(sev.dEVσ), sev.dEVσ, it_dEVσ, OnCell(), Val{0}())

    # information for how to scale the interpolation object
    scalegrid(x::AbstractRange{S}) where {S<:AbstractFloat} = StepRangeLen(x)
    scalegrid(x::Integer) = Base.OneTo(x)
    scalegrid(x::AbstractVector) = Base.OneTo(length(x))

    # scaled interpolation
    scl_EV   = scalegrid.((p.zspace..., _ψspace(p),         nS, sev.itypes...))
    scl_dEV  = scalegrid.((p.zspace..., _ψspace(p), nθ,     nS, sev.itypes...))
    scl_dEVσ = scalegrid.((p.zspace..., _ψspace(p),     nSexp1, sev.itypes...))

    sit_EV   = scale(itp_EV   , scl_EV...)
    sit_dEV  = scale(itp_dEV  , scl_dEV...)
    sit_dEVσ = scale(itp_dEVσ , scl_dEVσ...)
    # sit_dEVψ = Interpolations.scale(it_dEVψ , scl_dEVσ...)

    return ItpSharedEV{T,typeof(sit_EV),typeof(sit_dEV),typeof(sit_dEVσ),TT}(sit_EV, sit_dEV, sit_dEVσ, sev.itypes)
end

for typ in (:parallel, :serial)
    fun = Symbol("$(typ)_prefilterByView!")
    @eval begin
        function ($fun)(sev::SharedEV{T,N,N2}, isev::ItpSharedEV, dograd::Bool=true) where {T,N,N2}
            Ntyps = (length(sev.itypes[end]), )
            ($fun)(sev.EV, isev.EV.itp,         Ntyps...)
            if dograd
                ($fun)(sev.dEV, isev.dEV.itp,   Ntyps...)
                ($fun)(sev.dEVσ, isev.dEVσ.itp, Ntyps...)
            end
        end
    end
end

# function parallel_prefilterByView!(sev::SharedEV{T,N,N2}, isev::ItpSharedEV)
#     Ntyps = length.(sev.itypes)
#     parallel_prefilterByView!(sev.EV, isev.EV,     N- Ntyps+1:N...)
#     parallel_prefilterByView!(sev.dEV, isev.dEV,   N2-Ntyps+1:N...)
#     parallel_prefilterByView!(sev.dEVσ, isev.dEVσ, N- Ntyps+1:N...)
# end
#
# function parallel_prefilterByView!(sev::SharedEV{T,N,N2}, isev::ItpSharedEV)
#     Ntyps = length.(sev.itypes)
#     parallel_prefilterByView!(sev.EV, isev.EV,     N- Ntyps+1:N...)
#     parallel_prefilterByView!(sev.dEV, isev.dEV,   N2-Ntyps+1:N...)
#     parallel_prefilterByView!(sev.dEVσ, isev.dEVσ, N- Ntyps+1:N...)
# end

# TODO: MAKE SURE to check if data is within the indices

unsafegetityp(grid::AbstractArray, idx::Real) = idx
unsafegetityp(grid::Vector, idx::Integer) = grid[idx]

getitype(grid::Vector, idx::Integer) = unsafegetityp(grid, idx)

function getitype(grid::AbstractArray, idx::Real)
    minimum(grid) <= idx <= maximum(grid) || throw(DomainError(idx))
    unsafegetityp(grid, idx)
end

function unsafegetityp(sev::Union{SharedEV,ItpSharedEV}, itypidx::Tuple)
    length(sev.itypes) == length(itypidx) || throw(DimensionMismatch())
    return unsafegetityp.(sev.itypes, itypidx)
end

# -------------------------------------------------------------

function set_up_dcdp_workers(pids::AbstractVector{<:Integer}, prim::dcdp_primitives, typegrids::AbstractVector...; σ0::Real=1.0, ψflag::Interpolations.Flag=Quadratic(InPlace()), kwargs...)

    println("initialize shared arrays")
    sev = SharedEV(pids, prim, typegrids...)

    println("setting up tmpvars")
    @eval @everywhere begin
        # Inner loop setup
        set_g_dcdp_primitives($prim)
        set_g_dcdp_tmpvars( dcdp_tmpvars(get_g_dcdp_primitives()) )
        set_g_SharedEV($sev)

        set_g_ItpSharedEV( ItpSharedEV(get_g_SharedEV(), get_g_dcdp_primitives(), $σ0; ψflag=$ψflag) )
    end

    println("worker dcdp problems setup")

    # add a few checks
    fetch(@spawn size(get_g_SharedEV().EV))   == size(sev.EV)               ||  throw(error("remote sharedEV not created (?)"))
    fetch(@spawn get_g_SharedEV().EV  === get_g_ItpSharedEV().EV.itp.coefs) ||  throw(error("remote sharedEV not linked to ItpSharedEV"))
    fetch(@spawn get_g_SharedEV().EV) === sev.EV                            ||  throw(error("remote sharedEV not same object as local"))

    return sev, get_g_ItpSharedEV()
end


############################################################################################
"""
$(SIGNATURES)
Check if 'θ' is of finite value and return Bool. Not exported.

# Examples
```julia
```

"""
function _checkfinite(θ::T, max_val::R=max_val) where {T<:Real,R<:Real}
    return isfinite(θ) && θ < max_val
end
function _checkfinite(θ::Array{T}, max_val::R=max_val) where {T<:Real,R<:Real}
    @inbounds @simd for iter in eachindex(θ)
        if !_checkfinite(θ[iter], max_val)
            return false
        end
    end
    return true
end
function _checkfinite(θ::AbstractArray, max_val::R=max_val) where {R<:Real}
    @inbounds @simd for iter in eachindex(θ)
        if !_checkfinite(θ[iter], max_val)
            return false
        end
    end
    return true
end
function _checkfinite(θ::Tuple, max_val::R=max_val) where {R<:Real}
    return all(map(_checkfinite, θ))
end
function _checkfinite(θ::N, max_val::R=max_val) where {N<:NamedTuple} where {R<:Real}
    return all(map(_checkfinite, θ))
end

############################################################################################
"""
$(SIGNATURES)
Check if 'prior' is a valid density and return Bool. Not exported.

# Examples
```julia
```

"""
function _checkprior(prior)
    return false
end
function _checkprior(prior::S) where {S<:Distributions.Distribution}
    return true
end
function _checkprior(prior::AbstractVector{T}) where {T}
    @inbounds @simd for iter in eachindex(prior)
        if !_checkprior(prior[iter])
            return false
        end
    end
    return true
end
function _checkprior(prior::AbstractArray)
    @inbounds @simd for iter in eachindex(prior)
        if !_checkprior(prior[iter])
            return false
        end
    end
    return true
end
function _checkprior(prior::N) where {N<:NamedTuple}
    return all(map(_checkprior, prior))
end

############################################################################################
"""
$(SIGNATURES)
Check if argument is not fixed. Returns NamedTuple with true/false. Needed in addition to _checkprior for nested NamedTuples. Not exported.

# Examples
```julia
```

"""
function _checksampleable(constraint)
    return false
end
function _checksampleable(constraint::S) where {S<:Distributions.Distribution}
    return true
end
function _checksampleable(constraint::AbstractVector{T}) where {T}
    @inbounds @simd for iter in eachindex(constraint)
        if !_checksampleable(constraint[iter])
            return false
        end
    end
    return true
end
function _checksampleable(constraint::AbstractArray)
    @inbounds @simd for iter in eachindex(constraint)
        if !_checksampleable(constraint[iter])
            return false
        end
    end
    return true
end
function _checksampleable(constraint::NamedTuple{names}) where {names}
    return NamedTuple{names}(map(_checksampleable, constraint))
end

############################################################################################
#=
"""
$(SIGNATURES)
Check if provided val-constraint combination is valid for Param struct. Not exported.

# Examples
```julia
```

"""
function _checkparam(
    _rng::Random.AbstractRNG,
    constraint::AbstractArray,
    val::AbstractArray,
)
    return all(
        map(constraint, val) do constraintᵢ, valᵢ
            _checkparam(_rng, constraintᵢ, valᵢ)
        end,
    )
end
_checkparam(constraint, val) = _checkparam(Random.GLOBAL_RNG, constraint, val)
=#
############################################################################################
"""
$(SIGNATURES)
Check if all values in (Nested) NamedTuple are a 'Param' struct and return Bool. Not exported.

# Examples
```julia
```

"""
function _checkparams(param)
    return false
end
function _checkparams(param::Param)
    return true
end

function _checkparams(param::N) where {N<:NamedTuple}
    return all(map(_checkparams, param))
end

############################################################################################
# Export


using Pkg
Pkg.activate(".")

using ITensors
using LinearAlgebra
using DelimitedFiles

include("options.jl")

println()
print("Using $blas_num_threads BLAS thread")
blas_num_threads > 1 ? println("s") : println()

println()
println("Benchmarking $benchmarks")
println()

println("Bond dimensions set to:")
display(maxdims)
println()

warning_seperator = "X"^70

println()
println(warning_seperator)
if write_results
  println("XXX WARNING: benchmark results are set to be written to disk, may overwrite previous results")
else
  println("XXX WARNING: benchmark results are not set to be written to disk.")
end
println(warning_seperator)
println()

seperator = "#"^70

outputlevel = 1

for benchmark in benchmarks
  println("Running benchmark $benchmark for bond dimensions $(maxdims[benchmark])")
  println()
  for maxdim in maxdims[benchmark]
    if isnothing(which_version) || which_version == "julia"
      julia_dir = joinpath(@__DIR__, "bench_$benchmark", "julia")

      println(seperator)
      println("Run Julia benchmark $benchmark.")
      println()
      println("Benchmark located in path $julia_dir")
      println()
      println("Maximum bond dimension set to $maxdim.")
      println()

      BLAS.set_num_threads(blas_num_threads)

      include(joinpath(julia_dir, "run.jl"))

      # Trigger compilation
      run(maxdim = 10, nsweeps = 2, outputlevel = 0)

      time = @elapsed maxdim_ = run(maxdim = maxdim,
                                    outputlevel = outputlevel)

      println()
      println("Maximum dimension = $maxdim_")
      println("Total runtime = $time seconds")
      println()

      if write_results
        # TODO: add version number to data file name
        # v = Pkg.dependencies()[Base.UUID("9136182c-28ba-11e9-034c-db9fb085ebd5")].version
        # "$(v.major).$(v.minor).$(v.patch)"

        filename = "data_blas_num_threads_$(blas_num_threads)"
        filename *= "_maxdim_$(maxdim_).txt"
        filepath = joinpath(julia_dir, "data", filename)
        println("Writing results to $filepath")
        println()
        mkpath(dirname(filepath))
        writedlm(filepath, time)
      end
    end
    if isnothing(which_version) || which_version == "c++"
      cpp_dir = joinpath(@__DIR__, "bench_$benchmark", "c++")

      println(seperator)
      println("Run C++ benchmark $benchmark.")
      println()
      println("Benchmark located in path $cpp_dir")
      println()
      println("Maximum bond dimension set to $maxdim.")
      println()

      open("run.h", "w") do io
        write(io, "#include \"$(joinpath(cpp_dir, "run.h"))\"")
      end
      # Trigger rebuild
      touch("run.cc")
      println("Compile the benchmark")
      Base.run(`make`)
      println()
      rm("run.h")
      open("run.sh", "w") do io
        write(io, """#!/bin/bash
                     export MKL_NUM_THREADS=$blas_num_threads
                     export OPENBLAS_NUM_THREADS=$blas_num_threads
                     export OMP_NUM_THREADS=$omp_num_threads
                     ./run $maxdim $outputlevel""")
      end
      chmod("run.sh", 0o777)
      time = @elapsed Base.run(`./run.sh`)
      rm("run.sh")

      println()
      println("Maximum dimension = $maxdim")
      println("Total runtime = $time seconds")
      println()

      if write_results
        # TODO: add version number to data file name

        filename = "data_blas_num_threads_$(blas_num_threads)"
        filename *= "_maxdim_$(maxdim).txt"
        filepath = joinpath(cpp_dir, "data", filename)
        println("Writing results to $filepath")
        println()
        mkpath(dirname(filepath))
        writedlm(filepath, time)
      end
    end
  end
end


using Test
using PyCall
import Ghost: play!, Tape, Input
import ONNX: ONNXCtx, push_call!, julia2onnx, onnx2julia

function ort_run(path, args...)
    onnxruntime = pyimport("onnxruntime")
    ort_session = onnxruntime.InferenceSession(path)
    ort_inputs = PyDict(Dict([ort_session.get_inputs()[i].name => args[i] for i=1:length(args)]))
    return ort_session.run(nothing, ort_inputs)
end

function ort_test(tape::Tape, args...)
    mktemp() do path, _
        r1 = play!(tape, args...)
        ONNX.save(path, tape)
        r2 = ort_run(path, julia2onnx.(args)...)[1] |> onnx2julia
        tape2 = ONNX.load(path, args...; exec=true)
        r3 = tape2[tape2.result].val
        @test isapprox(r1, r2)
        @test isapprox(r1, r3)
    end
end

function ort_test(fn::Function, args...)
    tape = Tape(ONNXCtx())
    inp = [push!(tape, Input(arg)) for arg in args]
    res = push_call!(tape, fn, inp...)
    tape.result = res
    ort_test(tape, args...)
end


import Base: Array, convert
import AutoGrad: back
import CUDA: CuArray
using Knet.KnetArrays: DevArray, KnetArray
using AutoGrad: forw, Value, Arg, value

Array(x::Value{K}) where {K<:DevArray}=convert(Array,x)
KnetArray(x::Value{A}) where {A<:AbstractArray}=convert(KnetArray,x)
CuArray(x::Value{A}) where {A<:AbstractArray}=convert(CuArray,x)

convert(::Type{A},x::Value{K}) where {A<:AbstractArray,K<:DevArray}=forw(convert,A,x)
convert(::Type{K},x::Value{A}) where {A<:AbstractArray,K<:DevArray}=forw(convert,K,x)
back(::typeof(convert),::Type{Arg{2}},dy,y,T,x) = convert(typeof(value(x)),dy)

# Array/KnetArray Transfer

# So we will define gradients for convert, KnetArray, Array manually:

# This works but unnecessarily defines new functions:
# cpu2gpu(x::Array)=KnetArray(x)
# @primitive cpu2gpu(x),dy,y (gpu2cpu(dy))
# gpu2cpu(x::KnetArray)=Array(x)
# @primitive gpu2cpu(x),dy,y (cpu2gpu(dy))

# This does not work because !isa(Array,Function)
# @primitive  KnetArray(x::Array),dy  Array(dy)
# @primitive  Array(x::KnetArray),dy  KnetArray(dy)

# This does not work, parametric methods not yet supported, also unnecessary first arg gradient.
# @primitive convert{A<:AbstractArray,K<:KnetArray}(T::Type{K}, x::Value{A}),dy 0 Array(dy)
# @primitive convert{A<:AbstractArray,K<:KnetArray}(T::Type{A}, x::Value{K}),dy 0 KnetArray(dy)

# This gives ambiguity errors:
# @primitive convert(t::Type,x::KnetArray),dy  nothing  convert(KnetArray,dy)
# @primitive convert(t::Type{KnetArray},x::AbstractArray),dy  nothing  convert(Array,dy)