abstract type AbstractAggregateTools end

const ElIndices = Vector{Vector{T}} where {T<:Integer}
const VibIndices = Vector{Vector{Vector{Int64}}} where {T<:Integer}
# TODO: be more specific for T
const Indices = Vector{Vector{Vector{T} where T}}
const IndicesMap = Vector{Vector{T}} where {T<:Integer}
const FCfactorsT = Matrix{T} where {T<:AbstractFloat}
const FCproductT = Matrix{T} where {T<:AbstractFloat}

"""
    vibrationalIndices(aggCore)

Get all electronic indices of the [`Aggregate`](@ref).

# Arguments
* `aggCore`: Instance of [`Aggregate`](@ref).
"""
function electronicIndices(aggCore::AggregateCore)::ElIndices
    # TODO: add typeofs
    vibInds = Array{Array{Int64,1},1}(undef, 0)
    molCount = length(aggCore.molecules)
    currentInds = fill(1, (molCount))
    push!(vibInds, copy(currentInds))
    for indPos = 1:molCount
        currentInds[indPos] += 1
        push!(vibInds, copy(currentInds))
        currentInds[indPos] -= 1
    end
    return vibInds
end

"""
    vibrationalIndices(aggCore)

Get all vibrational indices of the [`Aggregate`](@ref).

"""
function vibrationalIndices(aggCore::AggregateCore)::VibIndices
    # TODO: add typeofs
    NvibMols = getNvib(aggCore)
    molLen = length(aggCore.molecules)
    NvibIndMols = Array{Array{Array{Int64,1},1},1}(undef, 0)
    for mol_i = 1:molLen
        push!(NvibIndMols, vibrationalIndices(NvibMols[mol_i]))
    end
    molNvib = map((vibInds) -> length(vibInds), NvibIndMols)
    molInds = vibrationalIndices(molNvib)
    aggInds = Array{Array{Array{Int64,1},1},1}(undef, 0)
    for molInd in molInds
        aggInd = Array{Array{Int64,1},1}(undef, 0)
        for mol_i = 1:molLen
            push!(aggInd, NvibIndMols[mol_i][molInd[mol_i]])
        end
        push!(aggInds, aggInd)
    end
    return aggInds
end

"""
    getIndices(aggCore)

Get all indices (electronic and vibrational) of the [`Aggregate`](@ref).

# Arguments
* `aggCore`: Instance of [`Aggregate`](@ref).
"""
function getIndices(aggCore::AggregateCore)::Indices
    vibInds = vibrationalIndices(aggCore)
    vibIndsLen = length(vibInds)
    elInds = electronicIndices(aggCore)
    elIndsLen = length(elInds)
    indices = Array{Array{Array{Int64,1} where Int64,1},1}(undef, 0)
    for el_i = 1:elIndsLen
        for vib_i = 1:vibIndsLen
            push!(indices, [elInds[el_i], vibInds[vib_i]])
        end
    end
    return indices
end

"""
getIndicesMap(aggCore)

Get pointers (integers) to the indices of the [`Aggregate`](@ref) separated by
electronic states (e.g. [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]]).

# Arguments
* `agg`: Instance of [`Aggregate`](@ref).
"""
function getIndicesMap(aggCore::AggregateCore, aggIndices::Indices)::IndicesMap
    aggIndLen = length(aggIndices)
    indicesMap = Array{Array{Int64,1},1}(undef, 0)
    elLen = length(aggCore.molecules) + 1
    for el_i = 1:elLen
        push!(indicesMap, Array{Int64,1}(undef, 0))
    end

    for I = 1:aggIndLen
        elind1, vibind1 = aggIndices[I]
        elOrder1 = OpenQuantumSystems.elIndOrder(elind1)
        push!(indicesMap[elOrder1], I)
    end
    return indicesMap
end

"""
    getFranckCondonFactors(aggCore, aggIndices)
    getFranckCondonFactors(aggCore)

Get Frack-Condon factors of the [`Aggregate`](@ref) in a form of matrix.

# Arguments
* `aggCore`: Instance of [`Aggregate`](@ref).
* `aggIndices`: Aggregate indices generated by [`getIndices`](@ref).
"""
function getFranckCondonFactors(aggCore::AggregateCore, aggIndices::Indices)::FCfactorsT
    aggIndLen = length(aggIndices)
    molLen = length(aggCore.molecules)
    FC = zeros(Float64, (aggIndLen, aggIndLen))
    for I = 1:aggIndLen
        elind1, vibind1 = aggIndices[I]
        for J = 1:aggIndLen
            elind2, vibind2 = aggIndices[J]
            if elind1 == elind2
                if vibind1 == vibind2
                    FC[I, J] = 1.0
                end
            else
                fc = 1.0::Float64
                for mi = 1:molLen
                    mol = aggCore.molecules[mi]
                    fc *=
                        getMolStateFC(mol, elind1[mi], vibind1[mi], elind2[mi], vibind2[mi])
                end
                FC[I, J] = fc
            end
        end
    end
    return FC
end

getFranckCondonFactors(aggCore::AggregateCore)::FCfactorsT =
    (aggIndices = getIndices(aggCore);
    getFranckCondonFactors(aggCore, aggIndices))

# TODO: resurrect sparse version
#=
"""
    getFranckCondonFactorsSparse(agg, aggIndices)
    getFranckCondonFactorsSparse(agg)

Sparse version of [`getFranckCondonFactors`](@ref).

"""
function getFranckCondonFactorsSparse(
    agg::C,
    aggIndices::Any
) where {C<:AbstractAggregateCore}
    if aggIndices === nothing
        aggIndices = getIndices(agg)
    end
    aggIndLen = length(aggIndices)
    molLen = length(agg.molecules)
    FC = spzeros(C2, aggIndLen, aggIndLen)
    for I = 1:aggIndLen
        elind1, vibind1 = aggIndices[I]
        for J = 1:aggIndLen
            elind2, vibind2 = aggIndices[J]
            if elind1 == elind2
                if vibind1 == vibind2
                    FC[I, J] = 1.0
                end
            else
                fc = 1.0::C2
                for mi = 1:molLen
                    mol = agg.molecules[mi]
                    fc *=
                        getMolStateFC(mol, elind1[mi], vibind1[mi], elind2[mi], vibind2[mi])
                end
                if fc != 0.0
                    FC[I, J] = fc
                end
            end
        end
    end
    return FC
end

getFranckCondonFactorsSparse(
    agg::C
) where {C<:AbstractAggregateCore} = getFranckCondonFactorsSparse(
    agg::C,
    nothing
)
=#

"""
    getFCProd(agg, FCFact, aggIndices, vibindices)

Get product of Franck-Condon factors. This way the trace over bath will be faster

"""
function getFCproduct(
    aggCore::AggregateCore,
    indices::Indices,
    indicesMap::IndicesMap,
    FCfactors::FCfactorsT,
)::FCproductT
    elLen = length(aggCore.molecules)
    aggIndLen = length(indices)
    vibLen = length(indicesMap[2])
    elLen += 1
    FCproduct = zeros(eltype(FCfactors), aggIndLen, aggIndLen)

    for I = 1:aggIndLen
        elind1, vibind1 = indices[I]
        elOrder1 = OpenQuantumSystems.elIndOrder(elind1)

        for J = 1:aggIndLen
            elind2, vibind2 = indices[J]
            elOrder2 = OpenQuantumSystems.elIndOrder(elind2)
            #=
            K1 = vibindices[elOrder1][1]
            K2 = vibindices[elOrder1][end]
            L1 = vibindices[elOrder2][1]
            L2 = vibindices[elOrder2][end]
            FCProd[I, J] = sum(FCFact[K1:K2, I] .* FCFact[J, L1:L2])
            =#

            for m = 1:vibLen
                # according to quantarhei, trace_over_vibrations()
                K = indicesMap[1][m]
                L = indicesMap[1][m]
                FCproduct[I, J] += FCfactors[K, I] * FCfactors[J, L]
            end
        end
    end
    return FCproduct
end

struct AggregateTools <: AbstractAggregateTools
    elIndices::ElIndices
    vibIndices::VibIndices
    indices::Indices
    indicesMap::IndicesMap
    FCfactors::FCfactorsT
    FCproduct::FCproductT
    basisSystem::GenericBasis
    basisBath::GenericBasis
    basis::GenericBasis
    bSystemSize::Int64
    bBathSize::Int64
    bSize::Int64
    function AggregateTools(
        elIndices::ElIndices,
        vibIndices::VibIndices,
        indices::Indices,
        indicesMap::IndicesMap,
        FCfactors::FCfactorsT,
        FCproduct::FCproductT,
        basisSystem::GenericBasis,
        basisBath::GenericBasis,
        basis::GenericBasis,
        bSystemSize::Int64,
        bBathSize::Int64,
        bSize::Int64,
    )::AggregateTools
        new(
            elIndices,
            vibIndices,
            indices,
            indicesMap,
            FCfactors,
            FCproduct,
            basisSystem,
            basisBath,
            basis,
            bSystemSize,
            bBathSize,
            bSize,
        )
    end
end

function AggregateTools(aggCore::AggregateCore)::AggregateTools
    elIndices = electronicIndices(aggCore)
    vibIndices = vibrationalIndices(aggCore)
    indices = getIndices(aggCore)
    indicesMap = getIndicesMap(aggCore, indices)
    FCfactors = getFranckCondonFactors(aggCore, indices)
    FCproduct = getFCproduct(aggCore, indices, indicesMap, FCfactors)

    bSystemSize = length(elIndices)
    bBathSize = length(vibIndices)
    bSize = length(indices)

    basisSystem = GenericBasis([bSystemSize])
    basisBath = GenericBasis([bBathSize])
    basis = GenericBasis([bSize])

    return AggregateTools(
        elIndices,
        vibIndices,
        indices,
        indicesMap,
        FCfactors,
        FCproduct,
        basisSystem,
        basisBath,
        basis,
        bSystemSize,
        bBathSize,
        bSize,
    )
end

Base.:(==)(x::AggregateTools, y::AggregateTools) =
    x.elIndices == y.elIndices &&
    x.vibIndices == y.vibIndices &&
    x.indices == y.indices &&
    x.FCfactors == y.FCfactors &&
    x.FCproduct == y.FCproduct &&
    x.bSystemSize == y.bSystemSize &&
    x.bBathSize == y.bBathSize &&
    x.bSize == y.bSize &&
    x.basisSystem == y.basisSystem &&
    x.basisBath == y.basisBath &&
    x.basis == y.basis

"""
    take_el_part(A, a, b, indicesMap)

Take electric part specified by electric indices `a` and `b` from the A (type of Array).

"""
function take_el_part(A::Array, a, b, indicesMap)
    a1 = indicesMap[a][1]
    a2 = indicesMap[a][end]
    b1 = indicesMap[b][1]
    b2 = indicesMap[b][end]

    return A[a1:a2, b1:b2]
end


# MIT License

#Copyright (c) 2021 hondoRandale <jules.rasetaharison@tutanota.com> and contributors

#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.

module SimpleFWA

  using InvertedIndices
  using Random
  using Statistics

  export FWA, simpleFWA, metropolisHastingsFWA

  mutable struct FWA
    X::Matrix{Float32};
    fitness_fireworks::Vector{Float32};
    S;
    fitness_sparks;
    x_b;
    y_min::Float32;
    iter::Int;
    err_conv::Float32;
  end

  function dynExplosionAmplitude( y_min::Float32,
                                  y_cf::Float32,
                                  A_cf_old::Float32,
                                  C_a::Float32,
                                  C_r::Float32 )
    A_cf = A_cf_old::Float32
    if y_min - y_cf < 0.0f0
      A_cf *= C_a;
    else
      A_cf *= C_r;
    end
    return A_cf
  end

  function explosionSparks( ;X_i, A_i, s, lower, upper, rng )
    nSparks = size( s, 2 );
    d       = size( s, 1 );
    for j ∈ 1:1:nSparks
      idz = bitrand( rng, d );
      for ii ∈ 1:1:d
        if idz[ii]
          @inbounds ΔX_k    = ( A_i * ( rand( rng, Float32, 1 )[1] * 2.0f0 - 1.0f0 ) )::Float32;
          @inbounds val     = ( X_i[ii] + ΔX_k )::Float32;
          if val >= lower[ii] && val <= upper[ii]
            @inbounds s[ii,j] = val::Float32;
          else
            @inbounds s[ii,j] = ( rand( rng, Float32, 1 )[1] * ( upper[ii] - lower[ii] ) + lower[ii] )::Float32;
          end
        end
      end
    end
  end

  function drawFireworksPositionsBoundary!(;rng,
                                           X::Matrix{Float32},
                                           rand_mat::Matrix{Float32},
                                           upper::Vector{Float32},
                                           lower::Vector{Float32} )
    @assert length( upper ) == length( lower )
    @assert all( size( X )  == size( rand_mat ) )

    nFireworks = size( X, 2 );
    d          = length( lower );
    rand_mat  .= Float32.( rand( rng, d, nFireworks ) );
    for iter ∈ 1:1:nFireworks
      x_sp     = view( X,        :, iter );
      rand_vec = view( rand_mat, :, iter );
      for k ∈ 1:1:d
        lu      = lower[k]::Float32;
        x_sp[k] = ( rand_vec[k] * ( upper[k] - lu ) + lu )::Float32;
      end
    end
  end

  function explosionAmplitudes!(;A::Vector{Float32},
                                 λ_0::Float32,
                                 ϵ_A::Float32,
                                 y_min::Float32,
                                 fitness_fireworks::Vector{Float32} )
    n      = length( fitness_fireworks );
    cacheA = Vector{Float32}( undef, n );
    for ii ∈ 1:1:n
      @inbounds fx         = fitness_fireworks[ii];
      @inbounds cacheA[ii] = ( fx - y_min )::Float32;
    end
    denomA = ( sum( cacheA ) + ϵ_A )::Float32;
    for ii ∈ 1:1:n
      @inbounds A[ii] = ( λ_0 * ( ( cacheA[ii] + ϵ_A ) / denomA ) )::Float32;
    end
  end

  function evaluateCandidates( ;kwargs::Tuple,
                               f::Function,
                               X,
                               fitness )
    @assert size( X, 2 ) == size( fitness, 1 )
    n = size( X, 2 )::Int;
    @assert n > 0
    for ii ∈ 1:1:n
      x            = view( X, :, ii );
      fitness[ii]  = f( x; kwargs=kwargs );
    end
  end

    """
        simpleFWA( nFireworks::Int,
                   nSparks::Int;
                   λ_0::Float32,
                   ϵ_A::Float32,
                   C_a::Float32,
                   C_r::Float32,
                   lower::Vector{Float32},
                   upper::Vector{Float32},
                   objFunction::Function,
                   XPrimary::Vector{ Matrix{Float32} },
                   yPrimary::Vector{ Matrix{Float32} },
                   maxiter::Int,
                   ϵ_conv::Float32=1f-6 )
        minimize objective function objFunction, the solution space is limited by lower and upper bound.
        The optimization algorithm utilized is an simplified version of dynFireWorksAlgorithm. The nFireworks
        parameter governs the number of fireworks being evaluated in parallel in each iteration.  nSparks is
        the number of sparks per firework, in remains constant foreach firework. ϵ_A is the smoothing parameter
        controlling the variance of amplitudes computed foreach fw. C_a ist the upscaling parameter for explosion
        amplitudes. C_r is the downscaling parameter for explosion amplitudes. XPrimary is the feature set of the
        primary algorithm to be tuned. yPrimary is the target set of the primary algorithm. maxiter is the maximum
        number iteraions.  ϵ_conv denotes the convergence parameter.
    """
  @views function simpleFWA( nFireworks::Int,
                             nSparks::Int,
                             kwargs...;
                             λ_0::Float32,
                             ϵ_A::Float32,
                             C_a::Float32,
                             C_r::Float32,
                             lower::Vector{Float32},
                             upper::Vector{Float32},
                             objFunction::Function,
                             maxiter::Int,
                             ϵ_conv::Float32=0.001f0 )

    @assert nFireworks > 0
    @assert nSparks    > 0
    @assert λ_0        > 0.0f0
    @assert ϵ_A        > 0.0f0
    @assert C_a        > 0.0f0
    @assert C_r        > 0.0f0
    @assert length( upper ) == length( lower )
    @assert all( lower .<  upper )
    @assert maxiter   > 0

    d        = length( lower );
    rng      = RandomDevice();

    ## explosion amplitudes of firework roots
    fitness_fireworks = Vector{Float32}( undef, nFireworks );
    fitness_sparks    = Vector{ Vector{Float32} }( undef,  nFireworks );
    A                 = Vector{Float32}( undef, nFireworks );
    A_old             = zeros( Float32, nFireworks );
    ## spark positions for all fireworks
    S                 = Vector{ Matrix{Float32} }( undef, nFireworks );

    for ii in 1:1:nFireworks
      fitness_sparks[ii] = zeros( Float32, nSparks );
      S[ii]              = zeros( Float32, d, nSparks );
    end

    ## firework root explosion positions
    X                 = Matrix{Float32}( undef, d, nFireworks );
    ## random mat needed for inititialization of fw positions
    rand_mat          = Matrix{Float32}( undef, d, nFireworks );
    # generate nFireworks randomly within boundaries
    drawFireworksPositionsBoundary!( ;rng=rng, X=X, rand_mat=rand_mat, upper=upper, lower=lower );
    ## evaluate fireworks
    evaluateCandidates(;kwargs   = kwargs,
                        f        = objFunction,
                        X        = X,
                        fitness  = fitness_fireworks );
    cf        = argmin( fitness_fireworks )::Int;
    y_min     = fitness_fireworks[cf]::Float32;
    y_min_old = fitness_fireworks[Not( cf )][1]::Float32;
    conv_val  = -1.0f0;

    # compute explosion amplitudes
    explosionAmplitudes!( ;A, λ_0, ϵ_A, y_min, fitness_fireworks );
    A[cf]   = ( minimum( A[ Not( cf ) ] ) / 2.0f0 )::Float32;
    iter    = 1;
    err     = 2.0f0 * ϵ_conv;
    x_b     = X[:,cf];
    x_b_old = zeros( Float32, d );
    while iter <= maxiter
      cf              = argmin( fitness_fireworks )::Int;
      @inbounds y_cf  = fitness_fireworks[cf]::Float32;
      @inbounds A_cf  = A[cf]::Float32;
      # compute explosion amplitudes
      explosionAmplitudes!( ;A, λ_0, ϵ_A, y_min, fitness_fireworks );
      @inbounds A[cf]  = A_cf::Float32;
      @inbounds A[cf]  = dynExplosionAmplitude( y_min, y_cf, A_old[cf], C_a, C_r )::Float32;
      A_old .= A;
      # compute explosion sparks
      Threads.@threads for jj ∈ 1:1:nFireworks
        explosionSparks(;X_i             = X[:,jj],
                         A_i             = A[jj],
                         s               = S[jj],
                         lower           = lower,
                         upper           = upper,
                         rng             = rng );

        evaluateCandidates(;kwargs   = kwargs,
                            f        = objFunction,
                            X        = S[jj],
                            fitness  = fitness_sparks[jj] );
      end
      for jj ∈ 1:1:nFireworks
        @inbounds val = minimum( fitness_sparks[jj] );
        @inbounds idx = argmin( fitness_sparks[jj] );
        ## update best globally found solution
        if val < y_min
          y_min = val;
          @inbounds x_b   = S[jj][:,idx];
        end
        ## update fw move to best found local solution
        if val < fitness_fireworks[jj]
          @inbounds X[:,jj]              .= S[jj][:,idx];
          @inbounds fitness_fireworks[jj] = val;
        end
      end
      μFW      = mean( fitness_fireworks );
      conv_val = std( fitness_fireworks;mean=mean( fitness_fireworks ) )^0.10;
      if iter > 10 &&  conv_val < ϵ_conv
        break
      end
      y_min_old = y_min
      iter     += 1;
    end
    # check if best global solution is contained in fw position
    if y_min > minimum( fitness_fireworks )
      idx             = argmin( fitness_fireworks );
      @inbounds y_min = fitness_fireworks[ idx ];
      @inbounds x_b   = X[:,idx];
    end
    return FWA( X, fitness_fireworks, S, fitness_sparks, x_b, y_min, iter-1, conv_val )
  end

  @views function metropolisHastingsFWA( nFireworks::Int,
                                         nSparks::Int;
                                         λ_0::Float32,
                                         ϵ_A::Float32,
                                         C_a::Float32,
                                         C_r::Float32,
                                         lower::Vector{Float32},
                                         upper::Vector{Float32},
                                         objFunction::Function,
                                         XPrimary::Vector{ Matrix{Float32} },
                                         yPrimary::Vector{ Matrix{Float32} },
                                         maxiter::Int,
                                         ϵ_conv::Float32=1f-40 )

    @assert nFireworks > 0
    @assert nSparks    > 0
    @assert λ_0        > 0.0f0
    @assert ϵ_A        > 0.0f0
    @assert C_a        > 0.0f0
    @assert C_r        > 0.0f0
    @assert length( upper ) == length( lower )
    @assert all( lower .<  upper )
    @assert size( XPrimary, 2 ) == length( yPrimary )
    @assert maxiter   > 0

    d        = length( lower );
    rng      = RandomDevice();
    ## explosion amplitudes of firework roots
    fitness_fireworks = Vector{Float32}( undef, nFireworks );
    fitness_sparks    = Vector{ Vector{Float32} }( undef,  nFireworks );
    A                 = Vector{Float32}( undef, nFireworks );
    A_old             = zeros( Float32, nFireworks );
    ## spark positions for all fireworks
    S                 = Vector{ Matrix{Float32} }( undef, nFireworks );

    for ii in 1:1:nFireworks
      fitness_sparks[ii] = zeros( Float32, nSparks );
      S[ii]              = zeros( Float32, d, nSparks );
    end

    ## firework root explosion positions
    X                 = Matrix{Float32}( undef, d, nFireworks );
    ## random mat needed for inititialization of fw positions
    rand_mat          = Matrix{Float32}( undef, d, nFireworks );
    # generate nFireworks randomly within boundaries
    drawFireworksPositionsBoundary!( ;rng=rng, X=X, rand_mat=rand_mat, upper=upper, lower=lower );
    ## evaluate fireworks
    evaluateCandidates(;XPrimary = XPrimary,
                        yPrimary = yPrimary,
                        f        = objFunction,
                        X        = X,
                        fitness  = fitness_fireworks );
    cf    = argmin( fitness_fireworks )::Int;
    y_min = fitness_fireworks[cf]::Float32;

    # compute explosion amplitudes
    explosionAmplitudes!( ;A, λ_0, ϵ_A, y_min, fitness_fireworks );
    A[cf]   = ( minimum( A[ Not( cf ) ] ) / 2.0f0 )::Float32;
    iter    = 0;
    err     = 2.0f0 * ϵ_conv;
    x_b     = X[:,cf];
    x_b_old = zeros( Float32, d );
    while err > ϵ_conv && iter <= maxiter
      cf              = argmin( fitness_fireworks )::Int;
      @inbounds y_cf  = fitness_fireworks[cf]::Float32;
      @inbounds A_cf  = A[cf]::Float32;
      # compute explosion amplitudes
      explosionAmplitudes!( ;A, λ_0, ϵ_A, y_min, fitness_fireworks );
      @inbounds A[cf]  = A_cf::Float32;
      @inbounds A[cf]  = dynExplosionAmplitude( y_min, y_cf, A_old[cf], C_a, C_r )::Float32;
      A_old .= A;
      # compute explosion sparks
      Threads.@threads for jj ∈ 1:1:nFireworks
        explosionSparks(;X_i             = X[:,jj],
                         A_i             = A[jj],
                         s               = S[jj],
                         lower           = lower,
                         upper           = upper,
                         rng             = rng );

        evaluateCandidates(;XPrimary = XPrimary,
                            yPrimary = yPrimary,
                            f        = objFunction,
                            X        = S[jj],
                            fitness  = fitness_sparks[jj] );
      end
      for jj ∈ 1:1:nFireworks
        @inbounds val = minimum( fitness_sparks[jj] );
        @inbounds idx = argmin( fitness_sparks[jj] );
        ## update best globally found solution
        if val < y_min
          y_min = val;
          @inbounds x_b   = S[jj][:,idx];
        end
        ## update fw by metropolis hastings rejection sampling
        α = val / fitness_fireworks[jj];
        if rand(Float32) < α
          @inbounds X[:,jj]              .= S[jj][:,idx];
          @inbounds fitness_fireworks[jj] = val;
        end
      end
      μFW = mean( fitness_fireworks );
      #err = sqrt( sum( std( X, mean=mean( X, dims=2 )[:,1], dims=2 )[:,1] ) );
      iter += 1;
    end
    # check if best global solution is contained in fw position
    if y_min > minimum( fitness_fireworks )
      idx             = argmin( fitness_fireworks );
      @inbounds y_min = fitness_fireworks[ idx ];
      @inbounds x_b   = X[:,idx];
    end
    return FWA( X, fitness_fireworks, S, fitness_sparks, x_b, y_min, iter, 0.0f0 )
  end

end


@testset "io" begin
    fname = "testdata/n3-window-1m-LXe-20mm-1-5.h5"
    pdf   = read_abc(fname)

    @test length(pdf.volume_names ) ==    88
    @test length(pdf.process_names) ==     3
    @test nrow(pdf.sensor_xyz     ) == 47428
    @test ncol(pdf.sensor_xyz     ) ==     4
    @test nrow(pdf.primaries      ) ==     5
    @test ncol(pdf.primaries      ) ==     7
    @test nrow(pdf.vertices       ) ==     5
    @test ncol(pdf.vertices       ) ==    13
    @test nrow(pdf.total_charge   ) ==  1080
    @test ncol(pdf.total_charge   ) ==     3
    @test nrow(pdf.waveform       ) ==  8098
    @test ncol(pdf.waveform       ) ==     3

    evt_range = in(4995000:4995004)
    @test all(evt_range.(pdf.primaries.event_id))
    @test all(evt_range.(pdf.vertices.event_id))
    @test all(evt_range.(pdf.waveform.event_id))

    @test sum(pdf.total_charge.charge) == 8098

    outfile = "testdata/evtPar_example.h5"
    norm, rmin, rmax, evtdf = read_evtpar([outfile])

    @test norm == 5000
    @test isapprox(rmin, 351.5)
    @test isapprox(rmax, 372.9)
    @test nrow(evtdf) == 538
    exp_fields = fieldnames(ATools.EventParameters)
    @test ncol(evtdf) == length(exp_fields)
    @test all(in(exp_fields).(propertynames(evtdf)))
end


### A Pluto.jl notebook ###
# v0.16.1

using Markdown
using InteractiveUtils

# ╔═╡ 912254bc-2532-11ec-3a0e-bb48e4656479
begin
	using MarkovModels
	using HDF5
	using JLD2
	using SparseArrays
	using PlutoUI
end

# ╔═╡ d130ae9e-f21b-436e-81d5-19913c29b896
den_fsm = jldopen("/people/ondel/Repositories/SpeechLab/recipes/lfmmi/graphs//wsj/train_alignments_fsms.jld2", "r") do f
	 f["48ic030f"]
end;

# ╔═╡ 74ab1c58-1b5b-48e8-a3bc-23db2585a195
"$den_fsm"

# ╔═╡ 74934d7e-b3b2-44c0-a59f-0c8d22ee2c7a
ω = den_fsm.T[1:end-1, end]

# ╔═╡ 510098ba-5632-4b16-84a6-ac805ff2354a
findnz(den_fsm.T[1:end-1, 1:end-1])

# ╔═╡ cea70023-79c2-4139-bc59-d6cf59c5ddf0
den_fsm.T[1,1]

# ╔═╡ 455b3530-3f02-40df-b890-bc7aabb056e1
@__DIR__

# ╔═╡ 42b894cc-7e28-48b6-8304-13f4b8b166e7
open("num_fsm_wsj.txt", "w") do f
	for (i, v) in zip(findnz(den_fsm.π[1:end-1])...)
		println(f, "0 $i $i $i $(-v.val)")
	end

	for (i, j, v) in zip(findnz(den_fsm.T[1:end-1,1:end-1])...)
		println(f, "$i $j $i $j $(-v.val)")
	end

	for (i, v) in zip(findnz(ω)...)
		println(f, "$i $(-v.val)")
	end
end

# ╔═╡ 00000000-0000-0000-0000-000000000001
PLUTO_PROJECT_TOML_CONTENTS = """
[deps]
HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
MarkovModels = "c2ae6250-d0a1-11ea-0991-234599ce5244"
PlutoUI = "7f904dfe-b85e-4ff6-b463-dae2292396a8"
SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"

[compat]
HDF5 = "~0.15.6"
JLD2 = "~0.4.14"
MarkovModels = "~0.7.0"
PlutoUI = "~0.7.14"
"""

# ╔═╡ 00000000-0000-0000-0000-000000000002
PLUTO_MANIFEST_TOML_CONTENTS = """
# This file is machine-generated - editing it directly is not advised

[[AbstractFFTs]]
deps = ["LinearAlgebra"]
git-tree-sha1 = "485ee0867925449198280d4af84bdb46a2a404d0"
uuid = "621f4979-c628-5d54-868e-fcf4e3e8185c"
version = "1.0.1"

[[Adapt]]
deps = ["LinearAlgebra"]
git-tree-sha1 = "84918055d15b3114ede17ac6a7182f68870c16f7"
uuid = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
version = "3.3.1"

[[ArgTools]]
uuid = "0dad84c5-d112-42e6-8d28-ef12dabb789f"

[[Artifacts]]
uuid = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"

[[BFloat16s]]
deps = ["LinearAlgebra", "Test"]
git-tree-sha1 = "4af69e205efc343068dc8722b8dfec1ade89254a"
uuid = "ab4f0b2a-ad5b-11e8-123f-65d77653426b"
version = "0.1.0"

[[Base64]]
uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"

[[Blosc]]
deps = ["Blosc_jll"]
git-tree-sha1 = "84cf7d0f8fd46ca6f1b3e0305b4b4a37afe50fd6"
uuid = "a74b3585-a348-5f62-a45c-50e91977d574"
version = "0.7.0"

[[Blosc_jll]]
deps = ["Artifacts", "JLLWrappers", "Libdl", "Lz4_jll", "Pkg", "Zlib_jll", "Zstd_jll"]
git-tree-sha1 = "e747dac84f39c62aff6956651ec359686490134e"
uuid = "0b7ba130-8d10-5ba8-a3d6-c5182647fed9"
version = "1.21.0+0"

[[CEnum]]
git-tree-sha1 = "215a9aa4a1f23fbd05b92769fdd62559488d70e9"
uuid = "fa961155-64e5-5f13-b03f-caf6b980ea82"
version = "0.4.1"

[[CUDA]]
deps = ["AbstractFFTs", "Adapt", "BFloat16s", "CEnum", "CompilerSupportLibraries_jll", "ExprTools", "GPUArrays", "GPUCompiler", "LLVM", "LazyArtifacts", "Libdl", "LinearAlgebra", "Logging", "Printf", "Random", "Random123", "RandomNumbers", "Reexport", "Requires", "SparseArrays", "SpecialFunctions", "TimerOutputs"]
git-tree-sha1 = "335b3d2373733919b4972a51215a6840c7a33828"
uuid = "052768ef-5323-5732-b1bb-66c8b64840ba"
version = "3.4.2"

[[ChainRulesCore]]
deps = ["Compat", "LinearAlgebra", "SparseArrays"]
git-tree-sha1 = "a325370b9dd0e6bf5656a6f1a7ae80755f8ccc46"
uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
version = "1.7.2"

[[Compat]]
deps = ["Base64", "Dates", "DelimitedFiles", "Distributed", "InteractiveUtils", "LibGit2", "Libdl", "LinearAlgebra", "Markdown", "Mmap", "Pkg", "Printf", "REPL", "Random", "SHA", "Serialization", "SharedArrays", "Sockets", "SparseArrays", "Statistics", "Test", "UUIDs", "Unicode"]
git-tree-sha1 = "31d0151f5716b655421d9d75b7fa74cc4e744df2"
uuid = "34da2185-b29b-5c13-b0c7-acf172513d20"
version = "3.39.0"

[[CompilerSupportLibraries_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "e66e0078-7015-5450-92f7-15fbd957f2ae"

[[DataStructures]]
deps = ["Compat", "InteractiveUtils", "OrderedCollections"]
git-tree-sha1 = "7d9d316f04214f7efdbb6398d545446e246eff02"
uuid = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
version = "0.18.10"

[[Dates]]
deps = ["Printf"]
uuid = "ade2ca70-3891-5945-98fb-dc099432e06a"

[[DelimitedFiles]]
deps = ["Mmap"]
uuid = "8bb1440f-4735-579b-a4ab-409b98df4dab"

[[Distributed]]
deps = ["Random", "Serialization", "Sockets"]
uuid = "8ba89e20-285c-5b6f-9357-94700520ee1b"

[[DocStringExtensions]]
deps = ["LibGit2"]
git-tree-sha1 = "a32185f5428d3986f47c2ab78b1f216d5e6cc96f"
uuid = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
version = "0.8.5"

[[Downloads]]
deps = ["ArgTools", "LibCURL", "NetworkOptions"]
uuid = "f43a241f-c20a-4ad4-852c-f6b1247861c6"

[[ExprTools]]
git-tree-sha1 = "b7e3d17636b348f005f11040025ae8c6f645fe92"
uuid = "e2ba6199-217a-4e67-a87a-7c52f15ade04"
version = "0.1.6"

[[FileIO]]
deps = ["Pkg", "Requires", "UUIDs"]
git-tree-sha1 = "3c041d2ac0a52a12a27af2782b34900d9c3ee68c"
uuid = "5789e2e9-d7fb-5bc7-8068-2c6fae9b9549"
version = "1.11.1"

[[GPUArrays]]
deps = ["Adapt", "LinearAlgebra", "Printf", "Random", "Serialization", "Statistics"]
git-tree-sha1 = "69faa5f1c5706ca9ca067604acf797ee3a8ec6f6"
uuid = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
version = "8.1.1"

[[GPUCompiler]]
deps = ["ExprTools", "InteractiveUtils", "LLVM", "Libdl", "Logging", "TimerOutputs", "UUIDs"]
git-tree-sha1 = "4ed2616d5e656c8716736b64da86755467f26cf5"
uuid = "61eb1bfa-7361-4325-ad38-22787b887f55"
version = "0.12.9"

[[HDF5]]
deps = ["Blosc", "Compat", "HDF5_jll", "Libdl", "Mmap", "Random", "Requires"]
git-tree-sha1 = "83173193dc242ce4b037f0263a7cc45afb5a0b85"
uuid = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
version = "0.15.6"

[[HDF5_jll]]
deps = ["Artifacts", "JLLWrappers", "LibCURL_jll", "Libdl", "OpenSSL_jll", "Pkg", "Zlib_jll"]
git-tree-sha1 = "fd83fa0bde42e01952757f01149dd968c06c4dba"
uuid = "0234f1f7-429e-5d53-9886-15a909be8d59"
version = "1.12.0+1"

[[HypertextLiteral]]
git-tree-sha1 = "72053798e1be56026b81d4e2682dbe58922e5ec9"
uuid = "ac1192a8-f4b3-4bfe-ba22-af5b92cd3ab2"
version = "0.9.0"

[[IOCapture]]
deps = ["Logging", "Random"]
git-tree-sha1 = "f7be53659ab06ddc986428d3a9dcc95f6fa6705a"
uuid = "b5f81e59-6552-4d32-b1f0-c071b021bf89"
version = "0.2.2"

[[InteractiveUtils]]
deps = ["Markdown"]
uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240"

[[IrrationalConstants]]
git-tree-sha1 = "f76424439413893a832026ca355fe273e93bce94"
uuid = "92d709cd-6900-40b7-9082-c6be49f344b6"
version = "0.1.0"

[[JLD2]]
deps = ["DataStructures", "FileIO", "MacroTools", "Mmap", "Pkg", "Printf", "Reexport", "TranscodingStreams", "UUIDs"]
git-tree-sha1 = "192934b3e2a94e897ce177423fd6cf7bdf464bce"
uuid = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
version = "0.4.14"

[[JLLWrappers]]
deps = ["Preferences"]
git-tree-sha1 = "642a199af8b68253517b80bd3bfd17eb4e84df6e"
uuid = "692b3bcd-3c85-4b1f-b108-f13ce0eb3210"
version = "1.3.0"

[[JSON]]
deps = ["Dates", "Mmap", "Parsers", "Unicode"]
git-tree-sha1 = "8076680b162ada2a031f707ac7b4953e30667a37"
uuid = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
version = "0.21.2"

[[LLVM]]
deps = ["CEnum", "LLVMExtra_jll", "Libdl", "Printf", "Unicode"]
git-tree-sha1 = "46092047ca4edc10720ecab437c42283cd7c44f3"
uuid = "929cbde3-209d-540e-8aea-75f648917ca0"
version = "4.6.0"

[[LLVMExtra_jll]]
deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
git-tree-sha1 = "6a2af408fe809c4f1a54d2b3f188fdd3698549d6"
uuid = "dad2f222-ce93-54a1-a47d-0025e8a3acab"
version = "0.0.11+0"

[[LazyArtifacts]]
deps = ["Artifacts", "Pkg"]
uuid = "4af54fe1-eca0-43a8-85a7-787d91b784e3"

[[LibCURL]]
deps = ["LibCURL_jll", "MozillaCACerts_jll"]
uuid = "b27032c2-a3e7-50c8-80cd-2d36dbcbfd21"

[[LibCURL_jll]]
deps = ["Artifacts", "LibSSH2_jll", "Libdl", "MbedTLS_jll", "Zlib_jll", "nghttp2_jll"]
uuid = "deac9b47-8bc7-5906-a0fe-35ac56dc84c0"

[[LibGit2]]
deps = ["Base64", "NetworkOptions", "Printf", "SHA"]
uuid = "76f85450-5226-5b5a-8eaa-529ad045b433"

[[LibSSH2_jll]]
deps = ["Artifacts", "Libdl", "MbedTLS_jll"]
uuid = "29816b5a-b9ab-546f-933c-edad1886dfa8"

[[Libdl]]
uuid = "8f399da3-3557-5675-b5ff-fb832c97cbdb"

[[LinearAlgebra]]
deps = ["Libdl"]
uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

[[LogExpFunctions]]
deps = ["ChainRulesCore", "DocStringExtensions", "IrrationalConstants", "LinearAlgebra"]
git-tree-sha1 = "34dc30f868e368f8a17b728a1238f3fcda43931a"
uuid = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
version = "0.3.3"

[[Logging]]
uuid = "56ddb016-857b-54e1-b83d-db4d58db5568"

[[Lz4_jll]]
deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
git-tree-sha1 = "5d494bc6e85c4c9b626ee0cab05daa4085486ab1"
uuid = "5ced341a-0733-55b8-9ab6-a4889d929147"
version = "1.9.3+0"

[[MacroTools]]
deps = ["Markdown", "Random"]
git-tree-sha1 = "5a5bc6bf062f0f95e62d0fe0a2d99699fed82dd9"
uuid = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
version = "0.5.8"

[[Markdown]]
deps = ["Base64"]
uuid = "d6f4376e-aef5-505a-96c1-9c027394607a"

[[MarkovModels]]
deps = ["CUDA", "LinearAlgebra", "SparseArrays"]
git-tree-sha1 = "4f01fbb60caea1c9f83770dda7d4d8cc7d63db65"
uuid = "c2ae6250-d0a1-11ea-0991-234599ce5244"
version = "0.7.0"

[[MbedTLS_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "c8ffd9c3-330d-5841-b78e-0817d7145fa1"

[[Mmap]]
uuid = "a63ad114-7e13-5084-954f-fe012c677804"

[[MozillaCACerts_jll]]
uuid = "14a3606d-f60d-562e-9121-12d972cd8159"

[[NetworkOptions]]
uuid = "ca575930-c2e3-43a9-ace4-1e988b2c1908"

[[OpenLibm_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "05823500-19ac-5b8b-9628-191a04bc5112"

[[OpenSSL_jll]]
deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
git-tree-sha1 = "15003dcb7d8db3c6c857fda14891a539a8f2705a"
uuid = "458c3c95-2e84-50aa-8efc-19380b2a3a95"
version = "1.1.10+0"

[[OpenSpecFun_jll]]
deps = ["Artifacts", "CompilerSupportLibraries_jll", "JLLWrappers", "Libdl", "Pkg"]
git-tree-sha1 = "13652491f6856acfd2db29360e1bbcd4565d04f1"
uuid = "efe28fd5-8261-553b-a9e1-b2916fc3738e"
version = "0.5.5+0"

[[OrderedCollections]]
git-tree-sha1 = "85f8e6578bf1f9ee0d11e7bb1b1456435479d47c"
uuid = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
version = "1.4.1"

[[Parsers]]
deps = ["Dates"]
git-tree-sha1 = "a8709b968a1ea6abc2dc1967cb1db6ac9a00dfb6"
uuid = "69de0a69-1ddd-5017-9359-2bf0b02dc9f0"
version = "2.0.5"

[[Pkg]]
deps = ["Artifacts", "Dates", "Downloads", "LibGit2", "Libdl", "Logging", "Markdown", "Printf", "REPL", "Random", "SHA", "Serialization", "TOML", "Tar", "UUIDs", "p7zip_jll"]
uuid = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"

[[PlutoUI]]
deps = ["Base64", "Dates", "HypertextLiteral", "IOCapture", "InteractiveUtils", "JSON", "Logging", "Markdown", "Random", "Reexport", "UUIDs"]
git-tree-sha1 = "d1fb76655a95bf6ea4348d7197b22e889a4375f4"
uuid = "7f904dfe-b85e-4ff6-b463-dae2292396a8"
version = "0.7.14"

[[Preferences]]
deps = ["TOML"]
git-tree-sha1 = "00cfd92944ca9c760982747e9a1d0d5d86ab1e5a"
uuid = "21216c6a-2e73-6563-6e65-726566657250"
version = "1.2.2"

[[Printf]]
deps = ["Unicode"]
uuid = "de0858da-6303-5e67-8744-51eddeeeb8d7"

[[REPL]]
deps = ["InteractiveUtils", "Markdown", "Sockets", "Unicode"]
uuid = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"

[[Random]]
deps = ["Serialization"]
uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"

[[Random123]]
deps = ["Libdl", "Random", "RandomNumbers"]
git-tree-sha1 = "0e8b146557ad1c6deb1367655e052276690e71a3"
uuid = "74087812-796a-5b5d-8853-05524746bad3"
version = "1.4.2"

[[RandomNumbers]]
deps = ["Random", "Requires"]
git-tree-sha1 = "043da614cc7e95c703498a491e2c21f58a2b8111"
uuid = "e6cf234a-135c-5ec9-84dd-332b85af5143"
version = "1.5.3"

[[Reexport]]
git-tree-sha1 = "45e428421666073eab6f2da5c9d310d99bb12f9b"
uuid = "189a3867-3050-52da-a836-e630ba90ab69"
version = "1.2.2"

[[Requires]]
deps = ["UUIDs"]
git-tree-sha1 = "4036a3bd08ac7e968e27c203d45f5fff15020621"
uuid = "ae029012-a4dd-5104-9daa-d747884805df"
version = "1.1.3"

[[SHA]]
uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce"

[[Serialization]]
uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b"

[[SharedArrays]]
deps = ["Distributed", "Mmap", "Random", "Serialization"]
uuid = "1a1011a3-84de-559e-8e89-a11a2f7dc383"

[[Sockets]]
uuid = "6462fe0b-24de-5631-8697-dd941f90decc"

[[SparseArrays]]
deps = ["LinearAlgebra", "Random"]
uuid = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"

[[SpecialFunctions]]
deps = ["ChainRulesCore", "IrrationalConstants", "LogExpFunctions", "OpenLibm_jll", "OpenSpecFun_jll"]
git-tree-sha1 = "793793f1df98e3d7d554b65a107e9c9a6399a6ed"
uuid = "276daf66-3868-5448-9aa4-cd146d93841b"
version = "1.7.0"

[[Statistics]]
deps = ["LinearAlgebra", "SparseArrays"]
uuid = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"

[[TOML]]
deps = ["Dates"]
uuid = "fa267f1f-6049-4f14-aa54-33bafae1ed76"

[[Tar]]
deps = ["ArgTools", "SHA"]
uuid = "a4e569a6-e804-4fa4-b0f3-eef7a1d5b13e"

[[Test]]
deps = ["InteractiveUtils", "Logging", "Random", "Serialization"]
uuid = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

[[TimerOutputs]]
deps = ["ExprTools", "Printf"]
git-tree-sha1 = "7cb456f358e8f9d102a8b25e8dfedf58fa5689bc"
uuid = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
version = "0.5.13"

[[TranscodingStreams]]
deps = ["Random", "Test"]
git-tree-sha1 = "216b95ea110b5972db65aa90f88d8d89dcb8851c"
uuid = "3bb67fe8-82b1-5028-8e26-92a6c54297fa"
version = "0.9.6"

[[UUIDs]]
deps = ["Random", "SHA"]
uuid = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"

[[Unicode]]
uuid = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5"

[[Zlib_jll]]
deps = ["Libdl"]
uuid = "83775a58-1f1d-513f-b197-d71354ab007a"

[[Zstd_jll]]
deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
git-tree-sha1 = "cc4bf3fdde8b7e3e9fa0351bdeedba1cf3b7f6e6"
uuid = "3161d3a3-bdf6-5164-811a-617609db77b4"
version = "1.5.0+0"

[[nghttp2_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "8e850ede-7688-5339-a07c-302acd2aaf8d"

[[p7zip_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "3f19e933-33d8-53b3-aaab-bd5110c3b7a0"
"""

# ╔═╡ Cell order:
# ╠═912254bc-2532-11ec-3a0e-bb48e4656479
# ╠═d130ae9e-f21b-436e-81d5-19913c29b896
# ╠═74ab1c58-1b5b-48e8-a3bc-23db2585a195
# ╠═74934d7e-b3b2-44c0-a59f-0c8d22ee2c7a
# ╠═510098ba-5632-4b16-84a6-ac805ff2354a
# ╠═cea70023-79c2-4139-bc59-d6cf59c5ddf0
# ╠═455b3530-3f02-40df-b890-bc7aabb056e1
# ╠═42b894cc-7e28-48b6-8304-13f4b8b166e7
# ╟─00000000-0000-0000-0000-000000000001
# ╟─00000000-0000-0000-0000-000000000002


function grad(Q::MPIStateArray, _1)
    Q_x1 = view(Q, (:,:,_1))

    v_∇Q_x1 = MPIStateArray{FT}(mpicomm, ArrayType, ijksize, nrealelem, 3)
    s_∇Q_x1 = MPIStateArray{FT}(mpicomm, ArrayType, ijksize, nrealelem, 3)

    event = launch_volume_gradient!(grid, v_∇Q, Q_x1)
    wait(event)

    event = launch_interface_gradient!(grid, s_∇Q, Q_x1)
    wait(event)

    grad = v_∇Q_x1 + s_∇Q_x1

    return grad
end

function grad(Q::MPIStateArray, _1, _2, _3)
    Q_x1 = view(Q, (:,:,_1))
    Q_x2 = view(Q, (:,:,_2))
    Q_x3 = view(Q, (:,:,_3))

    v_∇Q_x1 = MPIStateArray{FT}(mpicomm, ArrayType, ijksize, nrealelem, 3)
    v_∇Q_x2 = MPIStateArray{FT}(mpicomm, ArrayType, ijksize, nrealelem, 3)
    v_∇Q_x3 = MPIStateArray{FT}(mpicomm, ArrayType, ijksize, nrealelem, 3)

    s_∇Q_x1 = MPIStateArray{FT}(mpicomm, ArrayType, ijksize, nrealelem, 3)
    s_∇Q_x2 = MPIStateArray{FT}(mpicomm, ArrayType, ijksize, nrealelem, 3)
    s_∇Q_x3 = MPIStateArray{FT}(mpicomm, ArrayType, ijksize, nrealelem, 3)

    ## strong form gradients here
    event_x1 = launch_volume_gradient!(grid, v_∇Q_x1, Q_x1)
    event_x2 = launch_volume_gradient!(grid, v_∇Q_x2, Q_x2)
    event_x3 = launch_volume_gradient!(grid, v_∇Q_x3, Q_x3)
    wait(event_x1)
    wait(event_x2)
    wait(event_x3)

    ## strong form gradients here
    event_x1 = launch_interface_gradient!(grid, s_∇Q_x1, Q_x1)
    event_x2 = launch_interface_gradient!(grid, s_∇Q_x2, Q_x2)
    event_x3 = launch_interface_gradient!(grid, s_∇Q_x3, Q_x3)
    wait(event_x1)
    wait(event_x2)
    wait(event_x3)

    ∇Q_x1 = v_∇Q_x1 + s_∇Q_x1
    ∇Q_x2 = v_∇Q_x2 + s_∇Q_x2
    ∇Q_x3 = v_∇Q_x3 + s_∇Q_x3

    return (; ∇Q_x1, ∇Q_x2, ∇Q_x3)
end

function div(Q::MPIStateArray, _1, _2, _3)
    ∇Q_x1, ∇Q_x2, ∇Q_x3 = grad(Q, _1, _2, _3)

    div = MPIStateArray{FT}(mpicomm, ArrayType, ijksize, nrealelem, 1)
    div[:,:,1] = ∇Q_x1[:,:,1] + ∇Q_x2[:,:,2] + ∇Q_x3[:,:,3]

    return div
end

function curl(Q::MPIStateArray, _1, _2, _3)
    ∇Q_x1, ∇Q_x2, ∇Q_x3 = grad(Q, _1, _2, _3)

    curl = MPIStateArray{FT}(mpicomm, ArrayType, ijksize, nrealelem, 3)

    curl[:,:,1] = ∇Q_x3[:,:,2] - ∇Q_x2[:,:,3]
    curl[:,:,2] = ∇Q_x1[:,:,3] - ∇Q_x3[:,:,1]
    curl[:,:,3] = ∇Q_x2[:,:,1] - ∇Q_x1[:,:,2]

    return curl
end


@testset "bots" begin

@testset "MCTSBot" begin
    UCT_C = 2.

    init_bot(game, max_simulations, evaluator) = MCTSBot(game, evaluator, UCT_C, max_simulations, 5, true, 42, false, UCT, 0., 0.)

    @testset "can play tic_tac_toe" begin
        game = load_game("tic_tac_toe")
        max_simulations = 100
        evaluator = RandomRolloutEvaluator(20, 42)
        bot0 = init_bot(game, max_simulations, CxxPtr(evaluator))
        bot1 = init_bot(game, max_simulations, CxxPtr(evaluator))
        results = evaluate_bots(new_initial_state(game), [bot0, bot1], 42)
        @test results[1] + results[2] == 0
    end

    @testset "can play single player" begin
        game = load_game("catch")
        max_simulations = 100
        evaluator = RandomRolloutEvaluator(20, 42)
        bot = init_bot(game, max_simulations, CxxPtr(evaluator))
        results = evaluate_bots(new_initial_state(game), [bot], 42)
        @test results[] > 0
    end

    @testset "play three player stochastic games" begin
        game = load_game("pig(players=3,winscore=20,horizon=30)")
        max_simulations = 1000
        evaluator = RandomRolloutEvaluator(20, 42)
        bot0 = init_bot(game, max_simulations, CxxPtr(evaluator))
        bot1 = init_bot(game, max_simulations, CxxPtr(evaluator))
        bot2 = init_bot(game, max_simulations, CxxPtr(evaluator))
        results = evaluate_bots(new_initial_state(game), [bot0, bot1, bot2], 42)
        @test sum(results) == 0
    end

    function get_action_by_str(state, action_str)
        for action in legal_actions(state)
            if action_str == action_to_string(state, current_player(state), action)
                return action
            end
        end
        @error "Illegal action: $action_str"
    end

    function search_tic_tac_toe_state(initial_actions)
        game = load_game("tic_tac_toe")
        state = new_initial_state(game)
        for action_str in split(initial_actions, " ")
            apply_action(state, get_action_by_str(state, action_str))
        end
        evaluator = RandomRolloutEvaluator(20, 42)
        bot = MCTSBot(game, CxxPtr(evaluator), UCT_C, 10000, 10, true, 42, false, UCT, 0., 0.)
        mcts_search(bot, state), state
    end

    @testset "solve draw" begin
        root, state = search_tic_tac_toe_state("x(1,1) o(0,0) x(2,2)")
        @test to_string(state) == "o..\n.x.\n..x"
        @test get_outcome(root)[get_player(root)+1] == 0
        for c in get_children(root)
            @test get_outcome(c)[get_player(c)+1] <= 0
        end
        best = best_child(root)[]
        @test get_outcome(best)[get_player(best)+1] == 0
    end

    @testset "solve loss" begin
        root, state = search_tic_tac_toe_state("x(1,1) o(0,0) x(2,2) o(1,0) x(2,0)")
        @test to_string(state) == "oox\n.x.\n..x"
        @test get_outcome(root)[get_player(root)+1] == -1
        for c in get_children(root)
            @test get_outcome(c)[get_player(c)+1] == -1
        end
    end

    @testset "solve win" begin
        root, state = search_tic_tac_toe_state("x(1,0) o(2,2)")
        @test to_string(state) == ".x.\n...\n..o"
        @test get_outcome(root)[get_player(root)+1] == 1
        best = best_child(root)[]
        @test get_outcome(best)[get_player(best)+1] == 1
        @test action_to_string(state, get_player(best), get_action(best)) == "x(2,0)"
    end
end

end


module PiCamera

export preview

using AbstractPlotting, GLMakie, Makie, VideoIO, CoordinateTransformations, Interact, Observables

# cam_rotation = slider(range(0, stop = 2π, length = 100))
# cam_position = slider(range(0, stop = 2π, length = 100))
# cam_transform = PerspectiveMap() ∘ inv(AffineMap(cam_rotation, cam_position))
# screen_points = map(cam_transform, points)
function preview()
    VideoIO.viewcam()
end

end # module


include("./1-indrnn.jl")
include("./2-indlstm.jl")
include("./3-rnn.jl")

export RNN
export RNNs
export IndRNN
export IndRNNs
export IndLSTM
export IndLSTMs

export resethidden
export PadSeqPackBatch
export PackedSeqPredict
export PackedSeqForward

global RNNLIST = [IndRNN,IndRNNs,
                  IndLSTM,IndLSTMs,
                  RNN,RNNs];
# global RNNLIST = [RNN, rin, lstm, IndRNN, IndLSTM, RNNs, RIN, LSTM, IndLSTMs];

"""
    PadSeqPackBatch(inputs::Vector; epsilon::Real=0.0) -> output
+ `inputs` <: AbstractArray{Real,2}
+ `output` <: AbstractArray{Real,3}
pad epsilon to align raw input features probably with different length
# Examples
    julia> PadSeqPackBatch([ones(2,1), 2ones(2,2), 3ones(2,3)])
    2×3×3 Array{Float64,3}:
    [:, :, 1] =
     1.0  0.0  0.0
     1.0  0.0  0.0

    [:, :, 2] =
     2.0  2.0  0.0
     2.0  2.0  0.0

    [:, :, 3] =
     3.0  3.0  3.0
     3.0  3.0  3.0
"""
function PadSeqPackBatch(inputs::Vector; epsilon::Real=0.0)
    # all Array of inputs shall have the same size in dim-1
    batchSize = length(inputs)
    lengths   = [size(inputs[i], 2) for i in 1:batchSize]
    featDims  = size(inputs[1], 1)
    maxSteps  = maximum(lengths)
    RNNBatch  = zeros(eltype(inputs[1]), featDims, maxSteps, batchSize)
    fill!(RNNBatch, epsilon)

    for i = 1:batchSize
        Tᵢ = lengths[i]
        RNNBatch[:,1:Tᵢ,i] .= inputs[i]
    end
    return RNNBatch
end

function PackedSeqForward(chain::Block, x::Variable{S}) where S
    T = size(x, 2)
    y = Vector{Variable{S}}(undef, T)
    resethidden(chain)
    for t = 1:T
        y[t] = forward(chain, x[:,t,:])
    end

    timeSteps = T
    featsDims = size(y[1], 1)
    batchSize = size(y[1], 2)
    RNNBatch  = Variable{S}(S(undef, featsDims, timeSteps, batchSize), x.backprop)

    for t = 1:timeSteps
        RNNBatch.value[:,t,:] .= y[t].value
    end

    if RNNBatch.backprop
        function PackSeqSlicesBackward()
            for t = 1:timeSteps
                if need2computeδ!(y[t])
                    y[t].delta .+= RNNBatch.delta[:,t,:]
                end
            end
            ifNotKeepδThenFreeδ!(RNNBatch)
        end
        push!(graph.backward, PackSeqSlicesBackward)
    end
    return RNNBatch
end

function PackedSeqPredict(chain::Block, x::AbstractArray{S}) where S
    T = size(x,2)
    y = Vector{AbstractArray{S}}(undef,T)
    resethidden(chain)
    for t = 1:T
        y[t] = predict(chain, x[:,t,:])
    end

    timeSteps = T
    featsDims = size(y[1], 1)
    batchSize = size(y[1], 2)
    RNNBatch  = typeof(x)(undef, featsDims, timeSteps, batchSize)
    for t = 1:timeSteps
        RNNBatch[:,t,:] .= y[t]
    end
    return RNNBatch
end


# --------------------------------------------------------------------------- #
# Abstract type of a processing Model

abstract type Model{I <: Image} end

# --------------------------------------------------------------------------- #
# Get the image type that the model is constructed for

imgtype{I <: Image}(:: Model{I}) = I

# --------------------------------------------------------------------------- #
# Reference to the variables of a model trained by backpropagation

weights(m :: Model) = error("Not implemented")

# --------------------------------------------------------------------------- #
# Reference to other internals of a model

state(m :: Model) = error("Not implemented")

# --------------------------------------------------------------------------- #
# Forward evaluation of the model

# Low-level function needed to support automated differentiation
density(w, s, data, ::Type{Model}) = error("Not implemented")

# Convenience function built on top
function density{I <: Image}(m :: Model{I}, img :: I)
  x = reshape(img.data, (imgsize(img)..., imgchannels(I), 1))
  return density(weights(m), state(m), x, typeof(m))[:,:,1]
end

# A full image may be too large to go through the model in a single run.
# So divide the image into patches, apply the model to the patches, and
# stich the density-patches together again.
function density_patched{I <: Image}(model :: Model{I}, img :: I;
                                     patchsize = 256, overlap = 24,
                                     at = Array{Float32},
                                     callback :: Function = (i, n) -> nothing)

  s = (n, m) = imgsize(img)
  d = overlap

  # TODO: The first two can be relaxed
  @assert patchsize >= 48 "patchsize must be larger than or equal to 48"
  @assert min(n, m) > patchsize "image sizes must be larger than patchsize"
  @assert d >= 8 "overlap must be larger than or equal to 8"

  # Find parameters for suitable patch regions
  # k = number of patches in y/x direction
  # l = good length of the patches in y/x direction
  k = ceil.( Int, (s .- patchsize) ./ (patchsize .- d) ) .+ 1
  l = ceil.( Int, (s .+ (k .- 1) .* d) ./ (8 .* k) ) .* 8

  mirrorinds(a, l, n) = (a+l-1 <= n) ? (a:(a+l-1)) : ([a:n; n:-1:(n-(a+l-2-n))])

  # Pack the model, i.e. bring it to graphics memory if suitable at is given
  # TODO: These functions are only introduced in the file training.jl.
  # Should be moved?
  w = packweights(weights(model), at = at)
  s = packstate(state(model), at = at)

  # Create patch data
  patchdata = zeros(Float32, l..., imgchannels(img), prod(k))
  for i in 1:k[1], j in 1:k[2]
    y = (i == 1) ? 1 : (i-1)*(l[1] - d)
    x = (j == 1) ? 1 : (j-1)*(l[2] - d)

    # Get indices for this patch. If the patch would protrude from the
    # image, continue the patch on a mirror version after the image border
    yind = mirrorinds(y, l[1], n)
    xind = mirrorinds(x, l[2], m)

    patchdata[:,:,:,sub2ind((k[2], k[1]), j, i)] = imgdata(img)[yind,xind,:]
  end

  # Load the patch data to gpu if suitable at is given
  pd = convert(at, patchdata)

  # Since Knet does not yet support slicing in 4-d arrays, convert pd to a 2d
  # array for slicing
  sz = size(pd)
  pd = reshape(pd, (prod(sz[1:end-1]), sz[end]))

  # Bring the density patches in rectangular layout
  densities = Array{Array{Float32, 2}}(k...)
  for i in 1:k[1], j in 1:k[2]
    ind = sub2ind((k[2], k[1]), j, i)
    data = reshape(pd[:, ind:ind], (sz[1:end-1]..., 1))
    densities[i, j] = convert(at, density(w, s, data, typeof(model)))[:,:,1]

    # Give some feedback about the status via callback
    callback(ind, prod(k))
  end

  # Merge the density patches with smooth interpolation profiles
  ftrans = [ 1 - sin(r*(pi/2) / (d-1))^2 for r in 1:d ]
  rtrans = reverse(ftrans)

  sprofile = [ones(patchsize - d); ]
  eprofile = sprofile[end:-1:1]
  mprofile = sprofile .* eprofile

  dens = zeros(Float32, (k .* l .- (k.-1) .* d)...)

  for i in 1:k[1], j in 1:k[2]

    # If we are at the first or last patches in the x or y direction,
    # use start or end profiles; use middle profiles otherwise
    if i == 1
      y = 1
      yprofile = [ones(l[1] - d); ftrans]
    elseif i == k[1]
      y = (i-1)*(l[1] - d)
      yprofile = [rtrans; ones(l[1] - d)]
    else
      y = (i-1)*(l[1] - d)
      yprofile = [rtrans; ones(l[1] - 2d); ftrans]
    end

    if j == 1
      x = 1
      xprofile = [ones(l[2] - d); ftrans]
    elseif j == k[2]
      x = (j-1)*(l[2] - d)
      xprofile = [rtrans; ones(l[2] - d)]
    else
      x = (j-1)*(l[2] - d)
      xprofile = [rtrans; ones(l[2] - 2d); ftrans]
    end

    profile = yprofile .* xprofile'
    dens[y:y+l[1]-1, x:x+l[2]-1] += profile .* densities[i, j]

  end

  return dens[1:n, 1:m]
end

# --------------------------------------------------------------------------- #
# Estimate labels
# Densities are the output of the application of a neural network regressor

function label(dens :: Array{Float32, 2};
               candidates = false, level = 25, clevel = 10)

  # Find local maxima that are sufficiently high

  label = Array{Int}[]
  cands = Array{Int}[]

  m, n = size(dens)

  for i in 1:m, j in 1:n
    xsel = max(i-1, 1):min(i+1, m)
    ysel = max(j-1, 1):min(j+1, n)

    if dens[i, j] < maximum(dens[xsel, ysel])
      continue
    end

    if dens[i, j] >= level
      push!(label, [j, i])
    elseif dens[i, j] >= clevel
      push!(cands, [j, i])
    end
  end

  label = isempty(label) ? zeros(Int, 2, 0) : hcat(label...)
  cands = isempty(cands) ? zeros(Int, 2, 0) : hcat(cands...)

  if candidates
    return Label(label), Label(cands)
  else
    return Label(label)
  end
end

function label(dens :: Array{Float32, 3}; kwargs...)
  return label.([dens[:,:,i] for i in 1:size(dens, 3)]; kwargs...)
end

function label(args...; kwargs...)
  dmap = density(args...; kwargs...)
  return label(convert(Array{Float32}, dmap))
end

# --------------------------------------------------------------------------- #
# Some auxiliary functions for building neural network models

include("models/primitives.jl")

# --------------------------------------------------------------------------- #
# Class of models from (Pen, Yang, Li, et al.; 2018)
# and (Xie, Noble, Zisserman; 2015)

abstract type FCModel{I, BN} <: Model{I} end

weights(m :: FCModel) = m.weights
state(m :: FCModel) = m.bnmoments
hasbatchnorm{I, BN}(:: FCModel{I, BN}) = BN

# Different architectures of this model class

include("models/unetlike.jl")
include("models/multiscale3.jl")
include("models/fcrna.jl")


# *****************************************************************************
# Written by Ritchie Lee, ritchie.lee@sv.cmu.edu
# *****************************************************************************
# Copyright ã ``2015, United States Government, as represented by the
# Administrator of the National Aeronautics and Space Administration. All
# rights reserved.  The Reinforcement Learning Encounter Simulator (RLES)
# platform is licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License. You
# may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0. Unless required by applicable
# law or agreed to in writing, software distributed under the License is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the specific language
# governing permissions and limitations under the License.
# _____________________________________________________________________________
# Reinforcement Learning Encounter Simulator (RLES) includes the following
# third party software. The SISLES.jl package is licensed under the MIT Expat
# License: Copyright (c) 2014: Youngjun Kim.
# 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.
# *****************************************************************************

"""
Probabilistic Context-Free Grammar
Wraps a CFG and adds a parallel data structure for the probabilities
"""
module PCFGs

export PCFG, fit_mle!, weighted_sum!

using GrammaticalEvolution
using LinearDerivTrees
using AbstractTrees
using RLESUtils, StatUtils
import Base: rand, rand!, getindex, setindex!, normalize!, fill!, copy

typealias CFG Grammar
typealias ProbDict Dict{Symbol,Vector{Float64}}

type PCFG
    cfg::CFG
    probs::ProbDict #probability vector for each Or rule
end

"""
PCFG constructor that initializes all probabilities to uniform
"""
function PCFG(cfg::CFG)
    probs = ProbDict()
    for (sym,rule) in cfg.rules
        if isa(rule, OrRule) || isa(rule, RangeRule)
            probs[sym] = uniform_probs(length(rule))
        end
    end
    PCFG(cfg, probs)
end

getindex(pcfg::PCFG, index) = pcfg.probs[index]
setindex!(pcfg::PCFG, index, val) = pcfg.probs[index] = val

"""
Draw a random sample from pcfg.  The output is in-place
into tree.
"""
function rand!(tree::LinearDerivTree, pcfg::PCFG)
    rand!(Base.GLOBAL_RNG, tree, pcfg)
end
function rand!(rng::AbstractRNG, tree::LinearDerivTree, pcfg::PCFG)
    probs = pcfg.probs
    initialize!(tree)
    while !isdone(tree)
        as = actionspace(tree)
        sym = get_sym(tree)
        step!(tree, weighted_rand(rng, as, probs[sym]))
    end
    iscomplete(tree)
end

"""
Draw random samples according to pcfg and output them in-place into samples.
"""
function rand!(samples::Vector{LinearDerivTree}, pcfg::PCFG)
    rand!(Base.GLOBAL_RNG, samples, pcfg)
end
function rand!(rng::AbstractRNG, samples::Vector{LinearDerivTree}, pcfg::PCFG)
    for i = 1:length(samples)
        rand!(rng, samples[i], pcfg)
    end
end

"""
Draw N random samples from pcfg.  Derivation trees are limited to maxsteps.
Outputs a vector of LinearDerivTrees.
"""
function rand(pcfg::PCFG, N::Int64, maxsteps::Int64)
    rand(Base.GLOBAL_RNG, pcfg, N, maxsteps)
end
function rand(rng::AbstractRNG, pcfg::PCFG, N::Int64, maxsteps::Int64)
    params = LDTParams(pcfg.cfg, maxsteps)
    samples = Array(LinearDerivTree, N)
    for i = 1:N
        s = LinearDerivTree(params)
        rand!(rng, s, pcfg)
        samples[i] = s
    end
    samples
end

"""
Compute the maximum likelihood estimator from all the transitions in the
trees in samples and store the output in-lce into pcfg.
"""
function fit_mle!(pcfg::PCFG, samples::Vector{LinearDerivTree})
    fill!(pcfg, 0.0)
    for s in samples
        if iscomplete(s) #don't consider incomplete elite samples
            for node in PreOrderDFS(s.derivtree.root)
                if isa(node.rule, DecisionRule)
                    sym = Symbol(node.rule.name)
                    a = node.action
                    pcfg[sym][a] += 1.0
                end
            end
        end
    end
    normalize!(pcfg)
end

"""
Computes the weighted sum of two pcfgs and returns the output in-place
pcfg = w1 * pcfg + w2 * pcfg2
"""
function weighted_sum!(pcfg::PCFG, w1::Float64, pcfg2::PCFG, w2::Float64)
    for (k,v1) in pcfg.probs
        v2 = pcfg2.probs[k]
        v1[:] = w1*v1 + w2*v2
    end
end

"""
Fill all probabilities in pcfg to value x
"""
function fill!(pcfg::PCFG, x)
    for v in values(pcfg.probs)
        fill!(v, x)
    end
end

"""
Normalize the sum of each probability vector in pcfg
"""
function normalize!(pcfg::PCFG)
    for v in values(pcfg.probs)
        if sum(v) > 0.0 #make sure valid denom
            normalize!(v, 1)
        else #set to uniform, catches NaNs and 0.0
            len = length(v)
            fill!(v, 1.0/len)
        end
    end
end

function copy(pcfg::PCFG)
    PCFG(pcfg.cfg, deepcopy(pcfg.probs))
end

end #module


"""
Helper functions

RT_to_scans: Converts retention time to indexes in spectrum
filter_XIC: Returns spectrum filtered at chosen mz values
batch_import: Imports spectra present in given folder
retrieve_sample_name: Retrieves sample names for the spectra from the agilent .D folder

"""

"""Returns index range of given retention time range"""
function RT_to_scans(spectrum, RT_range)
	# Only possible if Rt is sorted
	if !issorted(spectrum["Rt"])
		error("Error: Retention time array is not in order")
	end

	# Convert Integer to tuple
	if RT_range isa(Number)
		RT_range = (RT_range, RT_range)
	end

	return searchsortedfirst(spectrum["Rt"], RT_range[1]), searchsortedlast(spectrum["Rt"], RT_range[2])
end

"""Returns filtered spectrum based on given mass (mz) values"""
function filter_XIC(spectrum, mass_values)
	# Return complete spectrum
	if mass_values == 0
		return reduce(+, spectrum["Mz_intensity"], dims=2)
	# Convert to tuple
	elseif typeof(mass_values) == Int
		mass_values = (mass_values,)
	end

	indices = Array{CartesianIndex{2}, 1}()

	# Add indices that correspond to given mz values
	for mass in mass_values
		mass_range = [mass - MAX_MZ_DEVIATION, mass + MAX_MZ_DEVIATION]

		# Store which indices are between mass range
		append!(indices, findall(mz -> (mz .>= mass_range[1]) .& (mz .<= mass_range[2]), spectrum["Mz_values"]))
	end

	# Create XIC spectrum
	spectrum_XIC = zeros(Int, size(spectrum["Mz_values"], 1))

	# Filter spectrum using stored indices
	for I in indices
		spectrum_XIC[I[1]] += spectrum["Mz_intensity"][I]
	end

	return spectrum_XIC
end

"""
	parse_data(row)
Parses mz values and RT overlap values from string to array
return mz_values, RT_overlap_values
"""
function parse_data(row)
	# Parse mz values
	mz_values_string = split(string(row.mz), ";")
	mz_values = Array{Any,1}(undef, length(mz_values_string))
	for (i, mz) in enumerate(mz_values_string)

		# Tuple of mz values
		if startswith(mz, "(")
			mz_tuple_string = split(mz[2:end - 1], ",")
			mz_values[i] = Tuple(([parse(Float32, sub_mass) for sub_mass in mz_tuple_string]))

		# Separate mz value
		else
			mz_values[i] = parse(Float32, mz)
		end
	end

	# Parse overlap values
	RT_overlap_values_string = split(string(row.overlap), ",")
	RT_overlap_values = zeros(Float16, 2)
	for RT_overlap in RT_overlap_values_string
		RT_overlap = RT_overlap == "missing" ? 0 : parse(Float16, RT_overlap)
		index = RT_overlap < row.RT ? 1 : 2
		RT_overlap_values[index] = RT_overlap
	end

	return mz_values, RT_overlap_values

end

"""
	batch_import(pathin, settings_json; use_filters=true)
Returns list containing all imported spectra from folder and metadata headers
Set use_filters to false to ignore filters and import all spectra,
"""
function batch_import(pathin, settings_json; use_filters=true)

	files = readdir(pathin, sort=false)
	files_supported = zeros(Bool, length(files),1)

	# Determine supported files
	for i=1:length(files)
		if length(files[i]) > length(".mzxml") && lowercase(files[i][end-4:end]) == "mzxml"
			files_supported[i] = true
		end
	end

	# Save indices of supported files
	supported_indices=findall(x -> x == true, files_supported)
	num_of_spectra = length(supported_indices)

	# Retrieve which metadata to save
	info_containing = settings_json[Symbol("output_metadata")]
	sample_info_metadata = Dict(i["element_name"] => i["header"] for i in info_containing if i["type"] == "sample_info_xml")
	metadata_headers = append!([], values(sample_info_metadata), ["Filename", "Folder Name"])
	# Retrieve filters
	filters = settings_json[Symbol("filters")]
	filters_must_contain = Dict(i["element_name"] => i["value"] for i in filters if i["filter"] == "must_contain")
	empty_field_allowed = Dict(i["element_name"] => i["empty_field_allowed"] for i in filters)

	spectra = []
	# Load files into spectra array
	for (i, file_index) in enumerate(supported_indices)
		filename = files[file_index]

		# Retrieve sample info and check filter requirements
		element_names = append!([], keys(sample_info_metadata), keys(filters_must_contain))
		sample_info = retrieve_sample_info(pathin, filename, element_names)
		filters_accepted = true
		for (element_name, value) in filters_must_contain
			if !(contains(sample_info[element_name], value) || (empty_field_allowed[element_name] && isempty(sample_info[element_name])))
				filters_accepted = false
				break
			end
		end

		if filters_accepted == false && use_filters == true
			@info "Spectrum $(i[1]) of $num_of_spectra did not meet filter requirements"
			continue
		end

		# Read spectrum
		spectrum = import_files_light(pathin,[filename])
		push!(spectra, spectrum)

		# Add metadata
		spectrum["MS1"]["Filename"] = filename
		merge!(spectrum["MS1"], Dict(header => sample_info[element_name] for (element_name, header) in sample_info_metadata))
		spectrum["MS1"]["Folder Name"] = split(pathin, "\\")[end]

		@info "Read spectra $(i[1]) of $num_of_spectra"
	end

	return spectra, metadata_headers
end

function retrieve_sample_info(pathin, filename, element_names)
	# Look for file in same folder (created by .bat file)
	filename_without_ext = filename[begin:end-6]
	sample_info_file = joinpath(pathin, filename_without_ext*"_sample_info.xml")

	# If it does not exist, look for file in .D folder
	if !isfile(sample_info_file)
		folder_loc = joinpath(pathin, filename_without_ext*".D")
		sample_info_file = joinpath(folder_loc, "AcqData\\sample_info.xml")
	end

	# File does not exist at all
	if !isfile(sample_info_file)
		@warn "Sample info file for $filename_without_ext does not exist, metadata will not be saved"
		return Dict(element_name => "" for element_name in element_names)
	end

	# Load xml file into array
	xdoc = parse_file(sample_info_file)
	xroot = root(xdoc)
	xarray = xroot["Field"]

	# Find index where element is stored and retrieve its value
	elements = Dict{String, String}()
	for element_name in element_names
		element_index = findfirst(i -> content(find_element(i, "Name")) == element_name, xarray)
		element_value = content(find_element(xarray[element_index], "Value"))
		elements[element_name] = element_value
	end

	# Free memory
	free(xdoc)

	return elements
end


export PPOPolicy, PPOTrajectory, MaskedPPOTrajectory

const PPOTrajectory = Trajectory{
    <:NamedTuple{
        (:action_log_prob, SART...),
        <:Tuple{
            <:CircularArrayBuffer,
            <:CircularArrayBuffer,
            <:CircularArrayBuffer,
            <:CircularArrayBuffer,
            <:CircularArrayBuffer,
        },
    },
}

function PPOTrajectory(; capacity, action_log_prob, kwargs...)
    merge(
        CircularArrayTrajectory(;
            capacity = capacity + 1,
            action_log_prob = action_log_prob,
        ),
        CircularArraySARTTrajectory(; capacity = capacity, kwargs...),
    )
end

const MaskedPPOTrajectory = Trajectory{
    <:NamedTuple{
        (:action_log_prob, SLART...),
        <:Tuple{
            <:CircularArrayBuffer,
            <:CircularArrayBuffer,
            <:CircularArrayBuffer,
            <:CircularArrayBuffer,
            <:CircularArrayBuffer,
            <:CircularArrayBuffer,
        },
    },
}

function MaskedPPOTrajectory(; capacity, action_log_prob, kwargs...)
    merge(
        CircularArrayTrajectory(;
            capacity = capacity + 1,
            action_log_prob = action_log_prob,
        ),
        CircularArraySLARTTrajectory(; capacity = capacity, kwargs...),
    )
end

function Base.length(t::Union{PPOTrajectory,MaskedPPOTrajectory})
    x = t[:terminal]
    size(x, ndims(x))
end

"""
    PPOPolicy(;kwargs)

# Keyword arguments

- `approximator`,
- `γ = 0.99f0`,
- `λ = 0.95f0`,
- `clip_range = 0.2f0`,
- `max_grad_norm = 0.5f0`,
- `n_microbatches = 4`,
- `n_epochs = 4`,
- `actor_loss_weight = 1.0f0`,
- `critic_loss_weight = 0.5f0`,
- `entropy_loss_weight = 0.01f0`,
- `dist = Categorical`,
- `rng = Random.GLOBAL_RNG`,

By default, `dist` is set to `Categorical`, which means it will only works
on environments of discrete actions. To work with environments of continuous
actions `dist` should be set to `Normal` and the `actor` in the `approximator`
should be a `GaussianNetwork`. Using it with a `GaussianNetwork` supports
multi-dimensional action spaces, though it only supports it under the assumption
that the dimensions are independent since the `GaussianNetwork` outputs a single
`μ` and `σ` for each dimension which is used to simplify the calculations.
"""
mutable struct PPOPolicy{A<:ActorCritic,D,R} <: AbstractPolicy
    approximator::A
    γ::Float32
    λ::Float32
    clip_range::Float32
    max_grad_norm::Float32
    n_microbatches::Int
    n_epochs::Int
    actor_loss_weight::Float32
    critic_loss_weight::Float32
    entropy_loss_weight::Float32
    rng::R
    n_random_start::Int
    update_freq::Int
    update_step::Int
    # for logging
    norm::Matrix{Float32}
    actor_loss::Matrix{Float32}
    critic_loss::Matrix{Float32}
    entropy_loss::Matrix{Float32}
    loss::Matrix{Float32}
end

function PPOPolicy(;
    approximator,
    update_freq,
    n_random_start = 0,
    update_step = 0,
    γ = 0.99f0,
    λ = 0.95f0,
    clip_range = 0.2f0,
    max_grad_norm = 0.5f0,
    n_microbatches = 4,
    n_epochs = 4,
    actor_loss_weight = 1.0f0,
    critic_loss_weight = 0.5f0,
    entropy_loss_weight = 0.01f0,
    dist = Categorical,
    rng = Random.GLOBAL_RNG,
)
    PPOPolicy{typeof(approximator),dist,typeof(rng)}(
        approximator,
        γ,
        λ,
        clip_range,
        max_grad_norm,
        n_microbatches,
        n_epochs,
        actor_loss_weight,
        critic_loss_weight,
        entropy_loss_weight,
        rng,
        n_random_start,
        update_freq,
        update_step,
        zeros(Float32, n_microbatches, n_epochs),
        zeros(Float32, n_microbatches, n_epochs),
        zeros(Float32, n_microbatches, n_epochs),
        zeros(Float32, n_microbatches, n_epochs),
        zeros(Float32, n_microbatches, n_epochs),
    )
end

function RLBase.prob(
    p::PPOPolicy{<:ActorCritic{<:GaussianNetwork},Normal},
    state::AbstractArray,
    mask,
)
    if p.update_step < p.n_random_start
        @error "todo"
    else
        μ, logσ = p.approximator.actor(send_to_device(device(p.approximator), state)) |> send_to_host
        StructArray{Normal}((μ, exp.(logσ)))
    end
end

function RLBase.prob(p::PPOPolicy{<:ActorCritic,Categorical}, state::AbstractArray, mask)
    logits = p.approximator.actor(send_to_device(device(p.approximator), state))
    if !isnothing(mask)
        logits .+= ifelse.(mask, 0f0, typemin(Float32))
    end
    logits = logits |> softmax |> send_to_host
    if p.update_step < p.n_random_start
        [
            Categorical(fill(1 / length(x), length(x)); check_args = false) for
            x in eachcol(logits)
        ]
    else
        [Categorical(x; check_args = false) for x in eachcol(logits)]
    end
end

function RLBase.prob(p::PPOPolicy, env::MultiThreadEnv)
    mask =  ActionStyle(env) === FULL_ACTION_SET ? legal_action_space_mask(env) : nothing
    prob(p, state(env), mask)
end

function RLBase.prob(p::PPOPolicy, env::AbstractEnv)
    s = state(env)
    s = Flux.unsqueeze(s, ndims(s) + 1)
    mask =  ActionStyle(env) === FULL_ACTION_SET ? legal_action_space_mask(env) : nothing
    prob(p, s, mask)
end

(p::PPOPolicy)(env::MultiThreadEnv) = rand.(p.rng, prob(p, env))

# !!! https://github.com/JuliaReinforcementLearning/ReinforcementLearning.jl/pull/533/files#r728920324
(p::PPOPolicy)(env::AbstractEnv) = rand.(p.rng, prob(p, env))

function (agent::Agent{<:PPOPolicy})(env::MultiThreadEnv)
    dist = prob(agent.policy, env)
    action = rand.(agent.policy.rng, dist)
    if ndims(action) == 2
        action_log_prob = sum(logpdf.(dist, action), dims = 1)
    else
        action_log_prob = logpdf.(dist, action)
    end
    EnrichedAction(action; action_log_prob = vec(action_log_prob))
end

function RLBase.update!(
    p::PPOPolicy,
    t::Union{PPOTrajectory,MaskedPPOTrajectory},
    ::AbstractEnv,
    ::PreActStage,
)
    length(t) == 0 && return  # in the first update, only state & action are inserted into trajectory
    p.update_step += 1
    if p.update_step % p.update_freq == 0
        _update!(p, t)
    end
end

function _update!(p::PPOPolicy, t::AbstractTrajectory)
    rng = p.rng
    AC = p.approximator
    γ = p.γ
    λ = p.λ
    n_epochs = p.n_epochs
    n_microbatches = p.n_microbatches
    clip_range = p.clip_range
    w₁ = p.actor_loss_weight
    w₂ = p.critic_loss_weight
    w₃ = p.entropy_loss_weight
    D = device(AC)

    n_envs, n_rollout = size(t[:terminal])
    @assert n_envs * n_rollout % n_microbatches == 0 "size mismatch"
    microbatch_size = n_envs * n_rollout ÷ n_microbatches

    n = length(t)
    states_plus = send_to_device(D, t[:state])

    states_flatten = flatten_batch(select_last_dim(states_plus, 1:n))
    states_plus_values =
        reshape(send_to_host(AC.critic(flatten_batch(states_plus))), n_envs, :)
    advantages = generalized_advantage_estimation(
        t[:reward],
        states_plus_values,
        γ,
        λ;
        dims = 2,
        terminal = t[:terminal],
    )
    returns = advantages .+ select_last_dim(states_plus_values, 1:n_rollout)

    actions_flatten = flatten_batch(select_last_dim(t[:action], 1:n))
    action_log_probs = select_last_dim(t[:action_log_prob], 1:n)

    # TODO: normalize advantage
    for epoch in 1:n_epochs
        rand_inds = shuffle!(rng, Vector(1:n_envs*n_rollout))
        for i in 1:n_microbatches
            inds = rand_inds[(i-1)*microbatch_size+1:i*microbatch_size]
            if t isa MaskedPPOTrajectory
                lam = send_to_device(
                    D,
                    select_last_dim(
                        flatten_batch(select_last_dim(t[:legal_actions_mask], 2:n+1)),
                        inds,
                    ),
                )
            else
                lam = nothing
            end
            s = send_to_device(D, select_last_dim(states_flatten, inds))  # !!! performance critical
            a = send_to_device(D, select_last_dim(actions_flatten, inds))

            if eltype(a) === Int
                a = CartesianIndex.(a, 1:length(a))
            end

            r = send_to_device(D, vec(returns)[inds])
            log_p = send_to_device(D, vec(action_log_probs)[inds])
            adv = send_to_device(D, vec(advantages)[inds])

            ps = Flux.params(AC)
            gs = gradient(ps) do
                v′ = AC.critic(s) |> vec
                if AC.actor isa GaussianNetwork
                    μ, logσ = AC.actor(s)
                    if ndims(a) == 2
                        log_p′ₐ = vec(sum(normlogpdf(μ, exp.(logσ), a), dims = 1))
                    else
                        log_p′ₐ = normlogpdf(μ, exp.(logσ), a)
                    end
                    entropy_loss = mean(size(logσ, 1) * (log(2.0f0π) + 1) .+ sum(logσ; dims = 1)) / 2
                else
                    # actor is assumed to return discrete logits
                    raw_logit′ = AC.actor(s)
                    if isnothing(lam)
                        logit′ = raw_logit′
                    else
                        logit′ = raw_logit′ .+ ifelse.(lam, 0.0f0, typemin(Float32))
                    end
                    p′ = softmax(logit′)
                    log_p′ = logsoftmax(logit′)
                    log_p′ₐ = log_p′[a]
                    entropy_loss = -sum(p′ .* log_p′) * 1 // size(p′, 2)
                end
                ratio = exp.(log_p′ₐ .- log_p)
                surr1 = ratio .* adv
                surr2 = clamp.(ratio, 1.0f0 - clip_range, 1.0f0 + clip_range) .* adv

                actor_loss = -mean(min.(surr1, surr2))
                critic_loss = mean((r .- v′) .^ 2)
                loss = w₁ * actor_loss + w₂ * critic_loss - w₃ * entropy_loss

                ignore() do
                    p.actor_loss[i, epoch] = actor_loss
                    p.critic_loss[i, epoch] = critic_loss
                    p.entropy_loss[i, epoch] = entropy_loss
                    p.loss[i, epoch] = loss
                end

                loss
            end

            p.norm[i, epoch] = clip_by_global_norm!(gs, ps, p.max_grad_norm)
            update!(AC, gs)
        end
    end
end

function RLBase.update!(
    trajectory::Union{PPOTrajectory,MaskedPPOTrajectory},
    ::PPOPolicy,
    env::MultiThreadEnv,
    ::PreActStage,
    action::EnrichedAction,
)
    push!(
        trajectory;
        state = state(env),
        action = action.action,
        action_log_prob = action.meta.action_log_prob,
    )

    if trajectory isa MaskedPPOTrajectory
        push!(trajectory; legal_actions_mask = legal_action_space_mask(env))
    end
end


using QuantumOptics: NLevelBasis, CompositeBasis, FockBasis
using Test, IonSim
using IonSim.PhysicalConstants: ħ, ca40_qubit_transition_frequency, c, m_ca40
using Suppressor

@suppress_err begin

# setup system
C = Ca40()
L1 = Laser(); L2 = Laser()
chain = LinearChain(
            ions=[C, C], com_frequencies=(x=3e6,y=3e6,z=1e6), vibrational_modes=(;z=[1])
        )

@testset "traps -- Trap" begin
    T = Trap(configuration=chain, lasers=[L1, L1])

    # test construction of CompositeBasis
    @test T.basis.bases[1] ≡ T.configuration.ions[1]
    @test T.basis.bases[2] ≡ T.configuration.ions[2]
    @test T.basis.bases[3] ≡ T.configuration.vibrational_modes.z[1]

    # test construction of T.δB
    t = 0:1:100
    ones = [1 for _ in t]
    T.δB = 1
    @test T.δB.(t) == ones
    @test T._cnst_δB
    T.δB = sin
    @test T.δB.(t) == sin.(t)
    @test !T._cnst_δB

    # test for warning when lasers=[L, L, L, ...] where L point to the same thing
    warn = "Some lasers point to the same thing. Making copies."
    @test_logs (:warn, warn) Trap(configuration=chain, lasers=[L1, L1, L1])
    # and test that, in this case, copies are made
    T1 = Trap(configuration=chain, lasers=[L1, L1, L1, L1], δB=t->t)
    for i in 1:4, j in i+1:4
        @test !(T1.lasers[i] ≡ T1.lasers[j])
    end

    # test setproperty for Trap
    # setting T.δB to a constant
    T.δB = 100
    @test T.δB(0) == 100
    # setting the :basis directly should have no effect
    basis = T.basis
    T.basis = 100
    @test T.basis == basis
    # changing the configuration should also update the basis
    chain1 = LinearChain(
            ions=[C, C, C], com_frequencies=(x=3e6,y=3e6,z=1e6), vibrational_modes=(;z=[1])
        )
    @test length(T.basis.bases) ≡ 3
    T.configuration = chain1
    @test T.configuration ≡ chain1
    @test length(T.basis.bases) ≡ 4

    # make sure print/show doesn't throw errors
    print(T); show(T)
end

@testset "traps -- general functions" begin
    Bhat = (x̂ + ŷ + ẑ)/√3
    T = Trap(configuration=chain, lasers=[L1], Bhat=Bhat)

    # get_basis
    @test T.basis == get_basis(T)

    # global_beam!
    L = Laser()
    global_beam!(T, L)
    @test L.pointing == [(1, 1.0), (2, 1.0)]

    # Efield_from_pi_time
    ion_index = 1
    laser_index = 1
    ion = T.configuration.ions[ion_index]
    transition = ("S-1/2", "D-1/2")
    # compare to specific pre-computed value for both methods
    E1 = Efield_from_pi_time(1e-6, Bhat, L1, ion, transition)
    E2 = Efield_from_pi_time(1e-6, T, laser_index, ion_index, transition)
    @test E1 ≈ 153739.56 rtol=1e-2
    @test E1 == E2
    # confirm in-place versions work
    Efield_from_pi_time!(1e-6, Bhat, L1, ion, transition)
    @test L1.E(0) == E1
    Efield_from_pi_time!(1e-6, T, laser_index, ion_index, transition)
    @test L1.E(0) == E1
    # shouldn't be able to have a laser argument where laser.pointing = []
    L = Laser()
    @test_throws AssertionError Efield_from_pi_time(1e-6, Bhat, L, ion, transition)
    L.pointing = [(1, 1.)]
    @test isinf(Efield_from_pi_time(1e-6, x̂, L, ion, transition))

    # Efield_from_rabi_frequency
    ion_index = 1
    laser_index = 1
    ion = T.configuration.ions[ion_index]
    transition = ("S-1/2", "D-1/2")
    # compare to specific pre-computed value for both methods
    E1 = Efield_from_rabi_frequency(5e5, Bhat, L1, ion, transition)
    E2 = Efield_from_rabi_frequency(5e5, T, laser_index, ion_index, transition)
    @test E1 ≈ 153739.56 rtol=1e-2
    @test E1 == E2
    # confirm in-place versions work
    Efield_from_rabi_frequency!(5e5, Bhat, L1, ion, transition)
    @test L1.E(0) == E1
    Efield_from_rabi_frequency!(5e5, T, laser_index, ion_index, transition)
    @test L1.E(0) == E1
    # shouldn't be able to have a laser argument where laser.pointing = []
    L = Laser()
    @test_throws AssertionError Efield_from_rabi_frequency(5e5, Bhat, L, ion, transition)

    # transition_frequency (test against pre-computed values)
    T.B = 4e-4
    f = transition_frequency(T, C, transition)
    @test f ≈ 2.2393e6 rtol=1e-4
    @test transition_frequency(T.B, C, transition) == f
    @test transition_frequency(T, 1, transition) == f

    # set_gradient
    set_gradient!(T, (1, 2), transition, 1e6)
    f1 = transition_frequency(T, T.configuration.ions[1], transition)
    f2 = transition_frequency(T, T.configuration.ions[2], transition)
    @test abs(f1 - f2) ≈ 1e6

    # test :(==)
    chain1 = LinearChain(
            ions=[C, C], com_frequencies=(x=3e6,y=3e6,z=1e6), vibrational_modes=(x=[1], z=[1])
        )
    T = Trap(configuration=chain1, lasers=[L1, L2])
    xmode = chain1.vibrational_modes.x[1]
    zmode = chain1.vibrational_modes.z[1]
    cb = get_basis(T)
    @test cb == C ⊗ C ⊗ xmode ⊗ zmode
    @test cb != C ⊗ C ⊗ zmode ⊗ xmode
    @test cb != C ⊗ C ⊗ C ⊗ xmode ⊗ zmode

    # test get_η
    f = ca40_qubit_transition_frequency
    λ = c/f
    # η = |k|cos(θ) * √(ħ / (2M ⋅ N ⋅ 2πν)); cos(θ) ≡ k̂ ⋅ mode_axis; N ≡ number of ions
    η(ν) = (2π/λ) * sqrt(ħ / (2 * m_ca40 * 2 * 2π * ν))
    @test abs(get_η(zmode, L, C)) ≈ η(zmode.ν)
    L.k = (x̂ + ẑ) / √2
    @test abs(get_η(xmode, L, C)) ≈ η(xmode.ν) / √2
end
end  # end suppress


include("OSMPBF.jl")

using ProtoBuf, CodecZlib

const BLOCK = Dict(
    "OSMHeader" => OSMPBF.HeaderBlock(),
    "OSMData" => OSMPBF.PrimitiveBlock(),
)

function parseblob(io::IO)
    len = ntoh(read(io, Int32))
    h = read(io, len)
    bh = readproto(IOBuffer(h), OSMPBF.BlobHeader())
    blob_buf = IOBuffer(read(io, bh.datasize))
    b = readproto(blob_buf, OSMPBF.Blob())
    ok = transcode(ZlibDecompressor, b.zlib_data)
    if length(ok) != b.raw_size
        @warn("Uncompressed size $(length(ok)) bytes does not match $(b.raw_size).")
    end
    readproto(IOBuffer(ok), BLOCK[bh._type])
end

function process_block(osm::OSMData, block::OSMPBF.HeaderBlock)
    # We don't currently use that information in `OSMData`
    box = block.bbox
    osm.bounds = Bounds(box.bottom / 1e9, box.top / 1e9, box.left / 1e9, box.right / 1e9)
    return
end

function geo(offset, granularity, int)
    return 0.000000001 * (offset + (granularity * int))
end

function tag(osm, element, table, key_id, value_id)
    tag(osm, element, table[key_id + 1], table[value_id + 1])
end

function process_elements(osm::OSMData, nodes::OSMPBF.DenseNodes, table, lat_offset, lon_offset, granularity)
    ids = nodes.id
    cumsum!(ids, ids)
    lat = nodes.lat
    cumsum!(lat, lat)
    lon = nodes.lon
    cumsum!(lon, lon)
    for i in eachindex(ids)
        @assert !haskey(osm.nodes, ids[i])
        osm.nodes[ids[i]] = LLA(
            geo(lat_offset, granularity, lat[i]),
            geo(lon_offset, granularity, lon[i]),
        )
    end
    i = 1
    j = 1
    keys_vals = nodes.keys_vals
    while j <= length(keys_vals)
        if iszero(keys_vals[j])
            i += 1
            j += 1
        else
            tag(osm, ids[i], table, keys_vals[j], keys_vals[j+1])
            j += 2
        end
    end
end

function process_element(osm, pbf_node::OSMPBF.Node, table, lat_offset, lon_offset, granularity)
    id = pbf_node.id
    osm.nodes[id] = LLA(
        geo(lat_offset, granularity, pbf_node.lat),
        geo(lon_offset, granularity, pbf_node.lon),
    )
    keys = pbf_node.keys
    vals = pbf_node.vals
    for i in eachindex(keys)
        tag(osm, id, table, keys[i], vals[i])
    end
end

function process_element(osm, pbf_way::OSMPBF.Way, table, lat_offset, lon_offset, granularity)
    way = Way(pbf_way.id)
    # `refs` is Δ coded
    refs = pbf_way.refs
    cumsum!(refs, refs)
    way.nodes = refs
    if isdefined(pbf_way, :keys)
        keys = pbf_way.keys
        vals = pbf_way.vals
        for i in eachindex(keys)
            tag(osm, way, table, keys[i], vals[i])
        end
    end
    push!(osm.ways, way)
    return
end

function process_element(osm, pbf_relation::OSMPBF.Relation, table, lat_offset, lon_offset, granularity)
    relation = Relation(pbf_relation.id)
    memids = pbf_relation.memids
    cumsum!(memids, memids)
    types = pbf_relation.types
    for i in eachindex(types)
        push!(relation.members, Dict(
            "type" => OSMPBF.Relation_MemberType[types[i] + 1],
            "ref" => string(memids[i]),
        ))
    end
    if isdefined(pbf_relation, :keys)
        keys = pbf_relation.keys
        vals = pbf_relation.vals
        for i in eachindex(keys)
            tag(osm, relation, table, keys[i], vals[i])
        end
    end
    push!(osm.relations, relation)
    return
end

function process_elements(osm, elements::Vector, table, lat_offset, lon_offset, granularity)
    for element in elements
        process_element(osm, element, table, lat_offset, lon_offset, granularity)
    end
end

function process_block(osm, block::OSMPBF.PrimitiveBlock)
    table = String.(block.stringtable.s)
    granularity = block.granularity
    lat_offset = block.lat_offset
    lon_offset = block.lon_offset
    for group in block.primitivegroup
        for key in fieldnames(OSMPBF.PrimitiveGroup)
            if isdefined(group, key)
                process_elements(osm, getproperty(group, key), table, lat_offset, lon_offset, granularity)
            end
        end
    end
end

function parsePBF(filename::AbstractString)
    osm = OSMData()
    open(filename) do io
        while !eof(io)
            block = parseblob(io)
            process_block(osm, block)
        end
    end
    return osm
end


using SparseArrays
using Distributions
using StatsBase
using GVFHordes
import GVFHordes: AbstractHorde
using MinimalRLCore

ExperienceReplayCuriosity(size, obs_size, obs_type, num_act) =
    DynaExperienceReplay(size,
                         (obs_type, Int, obs_type, Int, Float64, Float64, Float64, Float64, Bool),
                         (obs_size, 1, obs_size, 1, num_act, num_act, 1, 1, 1),
                         (:s, :a, :ns, :na, :actprob, :actprob_prime, :env_r, :rint, :t))

mutable struct AgentER{O, Φ, SC, R} <: AbstractAgent

    behaviour_learner::Learner
    behaviour_demons::Union{AbstractHorde,Nothing} # Round Robin learners have behaviour demons set to nothing
    behaviour_gamma::Float64

    demons::AbstractHorde
    demon_learner::Learner

    last_obs::O
    last_state::Φ

    last_action::Int
    num_actions::Int

    intrinsic_reward::IntrinsicReward
    state_constructor::SC
    use_external_reward::Bool
    exploration::ExplorationStrategy

    random_first_action::Bool

    batch_size::Int
    replay::R

end

function AgentER(horde,
                 behaviour_feature_size::Int,
                 behaviour_learner,
                 behaviour_horde,
                 behaviour_gamma,
                 demon_learner,
                 observation_size::Int,
                 num_actions::Int,
                 intrinsic_reward_type,
                 state_constructor,
                 use_external_reward,
                 exploration_strategy,
                 random_first_action,
                 batch_size,
                 replay_size)

    behaviour_weight_dims = (num_actions, behaviour_feature_size)

    intrinsic_reward = if intrinsic_reward_type == "weight_change"
        #TODO: The intrinsic reward is defined by how the components of the agent are put together. For example, an intrinsic reward
        # could be the model error, which would then require different components that are assembled in the agent
        # Not sure how the construction of intrinisic reward could be abstracted out of this constructor
        WeightChange(demon_learner)
    elseif intrinsic_reward_type == "no_reward"
        NoReward()
    else
        throw(ArgumentError("Not a valid intrinsic reward"))
    end

    replay = ExperienceReplayCuriosity(replay_size, observation_size, Float64, num_actions)

    AgentER(behaviour_learner,
            behaviour_horde,
            behaviour_gamma,

            horde,
            demon_learner,

            zeros(observation_size),
            spzeros(behaviour_feature_size),

            0, # last_action
            num_actions,

            intrinsic_reward,
            state_constructor,

            use_external_reward,
            exploration_strategy,
            random_first_action,
            batch_size,
            replay)

end

function proc_input(agent::AgentER, obs)
    return agent.state_constructor(obs)
end

function get_action(agent::AgentER, state, obs)
    action_probs = if agent.behaviour_learner isa OneDTMazeUtils.RoundRobinPolicy || agent.behaviour_learner isa TwoDGridWorldUtils.RoundRobinPolicy
        get_action_probs(agent.behaviour_learner, state, obs)
    elseif agent.behaviour_learner.update isa TabularRoundRobin
        get_action_probs(agent.behaviour_learner.update, state, obs)
    elseif agent.behaviour_learner isa BaselineUtils.FollowDemon
        agent.behaviour_learner(obs)
    elseif agent.behaviour_learner isa BaselineUtils.RandomDemons
        qs = agent.behaviour_learner(obs)
        agent.exploration(qs)
    else
        qs = agent.behaviour_learner(state)
        agent.exploration(qs)
    end
    action = StatsBase.sample(1:agent.num_actions, Weights(action_probs))
    return action, action_probs
end

function assign_horde!(agent::AgentER, horde)
    agent.demons = horde
end

MinimalRLCore.end!(agent::AgentER, obs, reward, is_terminal) =
    MinimalRLCore.step!(agent, obs, reward, is_terminal)

function MinimalRLCore.start!(agent::AgentER, obs, args...)
    next_state = proc_input(agent, obs)
    #Always exploring starts
    step!(agent.exploration)
    next_action = if agent.random_first_action
        sample(1:agent.num_actions, Weights(ones(agent.num_actions)))
    else
        a, probs = get_action(agent,next_state, obs)
        a
    end

    # NOTE: IS THIS RIGHT, seems wierd!?
    _, discounts, _ = get(agent.demons; state_t = obs, action_t = next_action, state_tp1 = obs, action_tp1 = next_action)

    agent.last_state = next_state
    agent.last_action = next_action
    agent.last_obs = obs
    zero_eligibility_traces!(agent.demon_learner)

    return next_action
end

function MinimalRLCore.step!(agent::AgentER, obs, r, is_terminal, args...)
    next_state = proc_input(agent, obs)
    step!(agent.exploration)
    next_action, next_action_probs = get_action(agent, next_state, obs)

    #=
    Dealing with replay
    =#
    let # (:s, :a, :ns, :na, :actprob, :r, :t)
        ap = get_behaviour_pis(agent, agent.last_state, agent.last_obs)
        ap_prime = get_behaviour_pis(agent, next_state, obs)
        r_int = update_reward!(agent.intrinsic_reward, agent)
        total_reward = agent.use_external_reward ? r_int + r : r_int

        exp = (agent.last_obs, agent.last_action, obs, next_action, ap, ap_prime, r, total_reward, is_terminal)
        # @show typeof(exp)
        # @show typeof(agent.replay)
        push!(agent.replay, exp)
    end

    update_demons!(agent.demon_learner,
                   agent,
                   agent.last_obs,
                   obs,
                   agent.last_state,
                   agent.last_action,
                   next_state,
                   next_action,
                   is_terminal,
                   r)

    #get intrinssic reward

    update_behaviour!(agent)
                      # agent.last_obs,
                      # obs,
                      # agent.last_state,
                      # agent.last_action,
                      # next_state,
                      # next_action,
                      # is_terminal,
                      # total_reward)

    agent.last_state = next_state
    agent.last_action = next_action
    agent.last_obs = obs
    return next_action
end

get_behaviour_pis(agent::AgentER, state, obs) =
    get_action(agent, state, obs)[2]

function μ_π(agent::AgentER, obs)
    state = proc_input(agent, obs)
    get_behaviour_pis(agent,state,obs)
end

function update_demons!(::QLearner, agent::AgentER, args...)#, obs, next_obs, state, action, next_state, next_action, is_terminal, env_reward)

    exp = StatsBase.sample(agent.replay, agent.batch_size)
    # (:s, :a, :ns, :na, :actprob, :actprob_prime, :env_r, :rint, :t))
    for traj in exp
        state = proc_input(agent, traj.s)
        next_state = proc_input(agent, traj.ns)
        update!(agent.demon_learner,
                agent.demons,
                traj.s,
                traj.ns,
                state,
                traj.a,
                next_state,
                traj.na,
                is_terminal,
                (s, obs) -> obs === traj.s ? traj.actprob : traj.actprob_prime, #get_behaviour_pis(agent, state, obs),
                traj.env_r)
    end

end

function update_demons!(::SRLearner, agent::AgentER, obs, next_obs, state, action, next_state, next_action, is_terminal, env_reward)

    # update reward part

    update_rew_part!(agent.demon_learner.update,
                     agent.demon_learner,
                     agent.demons,
                     obs,
                     next_obs,
                     state,
                     action,
                     next_state,
                     next_action,
                     is_terminal,
                     (state, obs) -> get_behaviour_pis(agent, state, obs),
                     env_reward)

    # update SR part
    exp = StatsBase.sample(agent.replay, agent.batch_size)
    # (:s, :a, :ns, :na, :actprob, :actprob_prime, :env_r, :rint, :t))
    for traj in exp
        state = proc_input(agent, traj.s)
        next_state = proc_input(agent, traj.ns)
        update_sr_part!(agent.demon_learner.update,
                        agent.demon_learner,
                        agent.demons,
                        traj.s,
                        traj.ns,
                        state,
                        traj.a,
                        next_state,
                        traj.na,
                        is_terminal,
                        (s, obs) -> obs === traj.s ? traj.actprob : traj.actprob_prime,
                        traj.env_r)
    end

end

function update_behaviour!(agent::AgentER) #, obs, next_obs, state, action, next_state, next_action, is_terminal, reward)

    # behaviour_pis = get_behaviour_pis(agent, state, obs)
    # next_behaviour_pis = get_behaviour_pis(agent, next_state, next_obs)

    #NOTE: Different call than demon updates as the reward and environment pseudotermination function
    #
    exp = StatsBase.sample(agent.replay, agent.batch_size)

    for traj in exp
        state = proc_input(agent, traj.s)
        next_state = proc_input(agent, traj.ns)
        update!(agent.behaviour_learner,
                agent.behaviour_demons,
                traj.s,
                traj.ns,
                state,
                traj.a,
                next_state,
                traj.na,
                is_terminal,
                (s, obs) -> obs === traj.s ? traj.actprob : traj.actprob_prime, #get_behaviour_pis(agent, state, obs),
                traj.rint)

    end
end

function get_demon_prediction(agent::AgentER, obs, action)
    ϕ = agent.state_constructor(obs)
    agent.demon_learner(ϕ, action)
end