text:
ENV["GUROBI_HOME"] = ARGS[9]
ENV["GRB_LICENSE_FILE"] = ARGS[10]

using Gurobi, JuMP, CSV, DataFrames, Random, RLEVectors, StatsBase

Random.seed!(1033)

@time begin

const CURR_PUMA = ARGS[1]
const SYN_HHS_FILE = ARGS[2]
const SYN_HHS_SPATIAL_FILE = ARGS[3]
const SYN_INDVS_FILE = ARGS[4]
const SYN_INDVS_SPATIAL_FILE = ARGS[5]
const PUMA_TRACT_EQUIV_FILE = ARGS[6]
const MARG_DIR = ARGS[7]
const GUROBI_LOGFILE = ARGS[8]
const TRACT_MARG_FILES = ["tract_hhtype.csv", "tract_tenur_hhinc.csv", "tract_nwork.csv"]
const BLKGP_MARG_FILES = ["blkgp_tenur_hhsiz.csv"]
const MARG_COLS = [11, 21, 4, 22]
const INDV_TRACT_MARG_FILES = ["tract_i_sex_i_age.csv"]
const INDV_BLKGP_MARG_FILES = ["blkgp_emply.csv"]
const INDV_MARG_COLS = [19, 20]

## Get tracts and block groups
puma_tract_equiv = CSV.read(PUMA_TRACT_EQUIV_FILE, DataFrame)
curr_geo = filter(row -> row.STATEFP == parse(Float64, CURR_PUMA[1:2]) && row.PUMA5CE == parse(Float64, CURR_PUMA[3:7]), puma_tract_equiv)
const TRACT_GEOIDS = [string(lpad(string(sct[1]), 2, "0"), lpad(string(sct[2]), 3, "0"), lpad(string(sct[3]), 6, "0"))
                        for sct in zip(curr_geo.STATEFP, curr_geo.COUNTYFP, curr_geo.TRACTCE)]

blkgp_marg_df = CSV.read(string(MARG_DIR, BLKGP_MARG_FILES[1]), DataFrame)
blkgp_marg_df = filter(row -> lpad(string(row.geoid), 12, "0")[1:11] in TRACT_GEOIDS, blkgp_marg_df)
const BLKGP_GEOIDS = lpad.(string.(blkgp_marg_df.geoid), 12, "0")

# create tract-block group mapping
const tract_blkgp = [findall(tract -> tract == tract_geoid, SubString.(BLKGP_GEOIDS, 1, 11)) for tract_geoid in TRACT_GEOIDS]

## Process marginals
indv_index = length(vcat(TRACT_MARG_FILES, BLKGP_MARG_FILES))+1
const m = length(BLKGP_GEOIDS)
marginals = []
levelss = []
const MARG_FILES = vcat(TRACT_MARG_FILES, BLKGP_MARG_FILES, INDV_TRACT_MARG_FILES, INDV_BLKGP_MARG_FILES)
for marg_file in MARG_FILES
        marg_df = CSV.read("$MARG_DIR$marg_file", DataFrame)
        marg_df = filter(row -> lpad(string(row.geoid), (marg_file in vcat(TRACT_MARG_FILES, INDV_TRACT_MARG_FILES) ? 11 : 12), "0") in (marg_file in vcat(TRACT_MARG_FILES, INDV_TRACT_MARG_FILES) ? TRACT_GEOIDS : BLKGP_GEOIDS), marg_df)
        select!(marg_df, Not(:geoid))
        # convert df into array of arrays
        marg = [col[:] for col in eachcol(Matrix(marg_df))]

        push!(marginals, marg)
        push!(levelss, names(marg_df))
end

const tract_hh_pops = sum(marginals[1], dims=1)[1]
const blkgp_hh_pops = sum(marginals[indv_index-1], dims=1)[1]
const tract_indv_pops = sum(marginals[indv_index], dims=1)[1]
const blkgp_indv_pops = sum(marginals[length(marginals)], dims=1)[1]

## Process population
syn_hhs = CSV.read(SYN_HHS_FILE, DataFrame)
puma_df = filter(row -> lpad(string(row.puma), 7, "0") == CURR_PUMA, syn_hhs)
pop = [row[:] for row in eachrow(Matrix(puma_df))]

const n = length(pop)

syn_indvs = CSV.read(SYN_INDVS_FILE, DataFrame)
n_indvs = DataFrames.nrow(syn_indvs)
syn_hhs_hhsizs = rle(syn_indvs.HHID)[2]
# precompute individual variable level counts for each household; factors for
#     household-level variables can all be 1 or whatever
indv_factors = Any[]
for lvls in levelss[1:(indv_index-1)]
        push!(indv_factors, [[1 for _ in 1:n] for _ in 1:length(lvls)])
end
for (indv_marg_col, lvls) in zip(INDV_MARG_COLS, levelss[indv_index:length(levelss)])
        with_counts = combine(groupby(syn_indvs, :HHID), indv_marg_col => hh_values -> Tuple(count(val->val==l, hh_values) for l in lvls))
        counts_mat = [[v for v in vals] for vals in with_counts[!,2]]
        counts_mat = [[row[col] for row in counts_mat] for col in 1:size(counts_mat[1])[1]]
        push!(indv_factors, counts_mat)
end

println("Block groups: ", m)
println("Tracts: ", length(TRACT_GEOIDS))
println("Households: ", n)
println("Individuals: ", n_indvs)
println("Marginals: ", MARG_FILES)

#######################################################
## Define optimization model
model = Model(Gurobi.Optimizer)
set_optimizer_attribute(model, "LogFile", GUROBI_LOGFILE)

## Variables
@variable(model, x[1:n*m]) # define variable x
n_cells = sum(sum(length(lvl) for lvl in marg) for marg in marginals)
@variable(model, y[1:n_cells]) # define absolute value helper variable y

## Objective
@objective(model, Min, sum(y))

## Constraints
@constraint(model, con_lb[i=1:n*m], 0 <= x[i]) # lower bound constraint
@constraint(model, con_ub[i=1:n*m], x[i] <= 1) # upper bound constraint
@constraint(model, con_rowsum[i=1:n], sum(x[m*(i-1)+j] for j in 1:m) == 1)
@constraint(model, con_colsum[j=1:m], sum(x[m*(i-1)+j] for i in 1:n) == blkgp_hh_pops[j])

# add marginal cell constraints to enforce absolute value
marg_cell_i = 1
for (marg_i, (marg_col, marg, lvls, indv_facs)) in enumerate(zip(vcat(MARG_COLS, INDV_MARG_COLS), marginals, levelss, indv_factors))
        cell_weight = length(lvls) * length(marg[1])
        for (lvl_i, (lvl, ifs)) in enumerate(zip(lvls, indv_facs))
                if length(marg[1]) == m # block group marginal
                        for j in 1:m
                                @constraint(model, y[marg_cell_i] >= (sum(ifs[i] * x[m*(i-1)+j] for (i,hh) in enumerate(pop) if marg_i>=indv_index || hh[marg_col]==lvl) - marg[lvl_i][j])/(1+marg[lvl_i][j])/cell_weight*blkgp_hh_pops[j])
                                @constraint(model, y[marg_cell_i] >= -(sum(ifs[i] * x[m*(i-1)+j] for (i,hh) in enumerate(pop) if marg_i>=indv_index || hh[marg_col]==lvl) - marg[lvl_i][j])/(1+marg[lvl_i][j])/cell_weight*blkgp_hh_pops[j])
                                global marg_cell_i += 1
                        end
                else # tract marginal
                        for j in 1:length(tract_hh_pops)
                                @constraint(model, y[marg_cell_i] >= (sum(sum(ifs[i] * x[m*(i-1)+k] for k in tract_blkgp[j]) for (i,hh) in enumerate(pop) if marg_i>=indv_index || hh[marg_col]==lvl) - marg[lvl_i][j])/(1+marg[lvl_i][j])/cell_weight*tract_hh_pops[j])
                                @constraint(model, y[marg_cell_i] >= -(sum(sum(ifs[i] * x[m*(i-1)+k] for k in tract_blkgp[j]) for (i,hh) in enumerate(pop) if marg_i>=indv_index || hh[marg_col]==lvl) - marg[lvl_i][j])/(1+marg[lvl_i][j])/cell_weight*tract_hh_pops[j])
                                global marg_cell_i += 1
                        end
                end
        end
end

## Optimize
status=optimize!(model)

println("****************************************************")
println("final objective value =", objective_value(model))

## assign households to optimized tracts
sol_matrix = transpose(reshape(value.(x), (m,n)))
blkgp_assignments = [BLKGP_GEOIDS[argmax(row)] for row in eachrow(sol_matrix)]
blkgp_by_hhid = DataFrame(HHID = [row[1] for row in pop],
                          geoid = blkgp_assignments)

syn_hhs_spatial = leftjoin(syn_hhs, blkgp_by_hhid, on=:HHID)
select!(syn_hhs_spatial, Not(:tenur_hhinc_prox))
select!(syn_hhs_spatial, Not(:tenur_hhsiz_prox))
CSV.write(SYN_HHS_SPATIAL_FILE, syn_hhs_spatial)

syn_indvs_spatial = leftjoin(syn_indvs, blkgp_by_hhid, on=:HHID)
CSV.write(SYN_INDVS_SPATIAL_FILE, syn_indvs_spatial)

end


text:
# Elementary Householder reflection
struct Householder{T,V<:AbstractVector,R<:IndexRange}
    β::T
    v::V
    r::R
end

Base.adjoint(H::Householder) = Householder(conj(H.β), H.v, H.r)

function householder(x::AbstractVector, r::IndexRange = axes(x, 1), k = first(r))
    i = findfirst(isequal(k), r)
    i isa Nothing && error("k = $k should be in the range r = $r")
    β, v, ν = _householder!(x[r], i)
    return Householder(β, v, r), ν
end
# Householder reflector h that zeros the elements A[r,col] (except for A[k,col]) upon lmul!(A,h)
function householder(A::AbstractMatrix, r::IndexRange, col::Int, k = first(r))
    i = findfirst(isequal(k), r)
    i isa Nothing && error("k = $k should be in the range r = $r")
    β, v, ν = _householder!(A[r, col], i)
    return Householder(β, v, r), ν
end
# Householder reflector that zeros the elements A[row,r] (except for A[row,k]) upon rmulc!(A,h)
function householder(A::AbstractMatrix, row::Int, r::IndexRange, k = first(r))
    i = findfirst(isequal(k), r)
    i isa Nothing && error("k = $k should be in the range r = $r")
    β, v, ν = _householder!(conj!(A[row, r]), i)
    return Householder(β, v, r), ν
end

# generate Householder vector based on vector v, such that applying the reflection
# to v yields a vector with single non-zero element on position i, whose value is
# positive and thus equal to norm(v)
function _householder!(v::AbstractVector{T}, i::Int) where {T}
    β::T = zero(T)
    @inbounds begin
        σ = abs2(zero(T))
        @simd for k in 1:i-1
            σ += abs2(v[k])
        end
        @simd for k in i+1:length(v)
            σ += abs2(v[k])
        end
        vi = v[i]
        ν = sqrt(abs2(vi) + σ)

        if iszero(σ) && vi == ν
            β = zero(vi)
        else
            if real(vi) < 0
                vi = vi - ν
            else
                vi = ((vi - conj(vi)) * ν - σ) / (conj(vi) + ν)
            end
            @simd for k in 1:i-1
                v[k] /= vi
            end
            v[i] = 1
            @simd for k in i+1:length(v)
                v[k] /= vi
            end
            β = -conj(vi) / (ν)
        end
    end
    return β, v, ν
end

function LinearAlgebra.lmul!(H::Householder, x::AbstractVector)
    v = H.v
    r = H.r
    β = H.β
    iszero(β) && return x
    @inbounds begin
        μ::eltype(x) = zero(eltype(x))
        i = 1
        @simd for j in r
            μ += conj(v[i]) * x[j]
            i += 1
        end
        μ *= β
        i = 1
        @simd for j in H.r
            x[j] -= μ * v[i]
            i += 1
        end
    end
    return x
end
function LinearAlgebra.lmul!(H::Householder, A::AbstractMatrix, cols = axes(A, 2))
    v = H.v
    r = H.r
    β = H.β
    iszero(β) && return A
    @inbounds begin
        for k in cols
            μ::eltype(A) = zero(eltype(A))
            i = 1
            @simd for j in r
                μ += conj(v[i]) * A[j, k]
                i += 1
            end
            μ *= β
            i = 1
            @simd for j in H.r
                A[j, k] -= μ * v[i]
                i += 1
            end
        end
    end
    return A
end
function LinearAlgebra.rmul!(A::AbstractMatrix, H::Householder, rows = axes(A, 1))
    v = H.v
    r = H.r
    β = H.β
    iszero(β) && return A
    w = similar(A, length(rows))
    fill!(w, zero(eltype(w)))
    @inbounds begin
        l = 1
        for k in r
            j = 1
            vl = v[l]
            @simd for i in rows
                w[j] += A[i, k] * vl
                j += 1
            end
            l += 1
        end
        l = 1
        for k in r
            j = 1
            vl = β * conj(v[l])
            @simd for i in rows
                A[i, k] -= w[j] * vl
                j += 1
            end
            l += 1
        end
    end
    return A
end
function LinearAlgebra.rmul!(b::OrthonormalBasis, H::Householder)
    v = H.v
    r = H.r
    β = H.β
    iszero(β) && return b
    w = similar(b[first(r)])
    @inbounds begin
        unproject!(w, b, v, 1, 0, r)
        rank1update!(b, w, v, -β, 1, r)
    end
    return b
end

ext: jl


text:
#based on the optimizing radius tutorial
@testset "optimize radius" begin
    er = 4.5
    k0 = 2π
    kin = sqrt(er)*k0
    l0 = 2π/k0
    a = 0.2*l0
    ui = LineSource(-3a, 0)
    P = 5
    centers = square_grid(5, a)
    M = size(centers,1)
    # let's impose symmetry wrt x-axis
    centers_abs = centers[:,1] + 1im*abs.(centers[:,2])
    ids, centers_abs = uniqueind(centers_abs)
    J = maximum(ids) #number of optim vars
    φs = zeros(M)

    fmm_options = FMMoptions(true, acc = 6, dx = 2a)
    optim_options =  Optim.Options(x_tol = 1e-4, outer_x_tol = 1e-4,
                            iterations = 10, outer_iterations = 10,
                            store_trace = true, allow_f_increases = true)

    points = [3a 0.0]
    r_max = (0.4*a)*ones(J)
    r_min = (1e-3*a)*ones(J)
    rs0 = (0.25*a)*ones(J)
    @assert verify_min_distance([CircleParams(r_max[i]) for i = 1:J],
            centers, ids, points)

    res = optimize_radius(rs0, r_min, r_max, points, ids, P, ui, k0, kin,
            centers, fmm_options, optim_options, minimize = true)
    rs = res.minimizer
    @test res.x_converged

    optim_options2 =  Optim.Options(f_tol = 1e-7, outer_f_tol = 1e-7,
                            iterations = 10, outer_iterations = 10,
                            store_trace = true, show_trace = false,
                            allow_f_increases = true)
    res2 = optimize_radius(rs0, r_min, r_max, points, ids, P, ui, k0, kin,
            centers, fmm_options, optim_options2, minimize = false,
            method = "LBFGS")
    @test res2.f_converged

    sp1 = ScatteringProblem([CircleParams(rs0[i]) for i = 1:J], ids, centers, φs)
    sp2 = ScatteringProblem([CircleParams(rs[i]) for i = 1:J], ids, centers, φs)

    points_ = [3a 0.0; 0.0 0.0]
    u1_5 = calc_near_field(k0, kin, 5, sp1, points_, ui; opt = fmm_options,
            verbose = false)
    u1_6 = calc_near_field(k0, kin, 6, sp1, points_, ui; opt = fmm_options,
            verbose = false)
    u1_7 = calc_near_field(k0, kin, 7, sp1, points_, ui; opt = fmm_options,
            verbose = false)
    u2_5 = calc_near_field(k0, kin, 5, sp2, points_, ui; opt = fmm_options,
            verbose = false)
    u2_6 = calc_near_field(k0, kin, 6, sp2, points_, ui; opt = fmm_options,
            verbose = false)
    u2_7 = calc_near_field(k0, kin, 7, sp2, points_, ui; opt = fmm_options,
            verbose = false)

    @test (norm(u1_5-u1_6)/norm(u1_6) < 1e-6) && (norm(u1_5-u1_7)/norm(u1_7) < 1e-6)
    @test (norm(u2_5-u2_6)/norm(u2_6) < 1e-6) && (norm(u2_5-u2_7)/norm(u2_7) < 1e-6)
end


text:
module Simhashes

    include("simhash.jl")
    include("simhash_index.jl")

    export Simhash, SimhashIndex
    export distance, get_near_dups, delete_from_index, add_to_index

end # module


text:
#Returns the length of a expression tuple, or 1 if it is not an expression tuple (probably a  Symbol/Numerical).
function tup_leng(ex::ExprValues)
    (typeof(ex)==Expr && ex.head == :tuple) && (return length(ex.args))
    return 1
end

#Gets the i'th element in a expression tuple, or return the input itself if it is not an expression tuple (probably a  Symbol/Numerical).
function get_tup_arg(ex::ExprValues,i::Int)
    (tup_leng(ex) == 1) && (return ex)
    return ex.args[i]
end

#This will be called whenever a function stored in funcdict is called.
function replace_names(expr, old_names, new_names)
    mapping = Dict(zip(old_names, new_names))
    MacroTools.postwalk( x -> x in old_names ? x= mapping[x] : x, expr)
end

# below are currently not used

# function splitplus!(ex)
#     dosplit = ex.head == :(=) && ex.args[2] isa Expr && ex.args[2].head == :call && ex.args[2].args[1] == :(+)
#     if dosplit
#       summands = ex.args[2].args[2:end]
#       ex.args[2] = foldl((x,y)->(:(($x + $y))), summands)
#     end
#     dosplit
# end

# #Creates an expression which can be evaluated to an actual function. Input is an array of expression were each entry is a line in the function. Uses the array of expressions generated in either get_f or get_g.
# function make_func(func_expr::Vector{Expr},reactants::OrderedDict{Symbol,Int}, parameters::OrderedDict{Symbol,Int})
#     system = Expr(:block)
#     for func_line in deepcopy(func_expr)
#         ex = recursive_replace!(func_line, (reactants,:internal_var___u), (parameters, :internal_var___p))
#         splitplus!(ex)
#         push!(system.args,ex)
#     end
#     push!(system.args, :(nothing))
#     return :((internal_var___du,internal_var___u,internal_var___p,t) -> @inbounds $system)
# end

# """
# clean_subtractions(ex::Expr)
# Replace additions of negative terms with subtractions.
# This is a fairly stupid function which is designed for a specific problem
# with reaction networks. It is neither recursive nor very general.
# Return :: cleaned out expression

# From Latexify.jl with permission:
# [see](https://github.com/JuliaDiffEq/Catalyst.jl/issues/89#issuecomment-462147882)
# """
# function clean_subtractions(ex::Expr)
#     ex.args[1] != :+ && return ex

#     term = ex.args[2]

#     ### Sort out the first term
#     if term isa Expr && length(term.args) >= 3 && term.args[1:2] == [:*, -1]
#         result = :(- *($(term.args[3:end]...)))
#     else
#         result = :($term)
#     end

#     ### Sort out the other terms
#     for term in ex.args[3:end]
#         if term isa Expr && length(term.args) >= 3 && term.args[1:2] == [:*, -1]
#             result = :($result - *($(term.args[3:end]...)))
#         else
#             result = :($result + $term)
#         end
#     end
#     return result
# end

# #Recursively traverses an expression and removes things like X^1, 1*X. Will not actually have any affect on the expression when used as a function, but will make it much easier to look at it for debugging, as well as if it is transformed to LaTeX code.
# function recursive_clean!(expr::ExprValues)
#     (expr == :no___noise___scaling) && (return 1)
#     (typeof(expr)!=Expr) && (return expr)
#     for i = 1:length(expr.args)
#         expr.args[i] = recursive_clean!(expr.args[i])
#     end
#     (expr.args[1] == :^) && (expr.args[3] == 1) && (return expr.args[2])
#     if expr.args[1] == :*
#         in(0,expr.args) && (return 0)
#         i = 1
#         while (i = i + 1) <= length(expr.args)
#              if (typeof(expr.args[i]) == Expr) && (expr.args[i].head == :call) && (expr.args[i].args[1] == :*)
#                  for arg in expr.args[i].args
#                      (arg != :*) && push!(expr.args, arg)
#                  end
#              end
#         end
#         for i = length(expr.args):-1:2
#             (typeof(expr.args[i]) == Expr) && (expr.args[i].head == :call) && (expr.args[i].args[1] == :*) && deleteat!(expr.args,i)
#             (expr.args[i] == 1) && deleteat!(expr.args,i)
#         end
#         (length(expr.args) == 2) && (return expr.args[2])                   # We have a multiplication of only one thing, return only that thing.
#         (length(expr.args) == 1) && (return 1)                              #We have only * and no real argumenys.
#         (length(expr.args) == 3) && (expr.args[2] == -1) && return :(-$(expr.args[3]))
#         (length(expr.args) == 3) && (expr.args[3] == -1) && return :(-$(expr.args[2]))
#     end
#     if expr.head == :call
#         (expr.args[1] == :/) && (expr.args[3] == 1) && (return expr.args[2])
#         haskey(funcdict, expr.args[1]) && return funcdict[expr.args[1]](expr.args[2:end])
#         in(expr.args[1],hill_name) && return hill(expr)
#         in(expr.args[1],hillR_name) && return hillR(expr)
#         in(expr.args[1],mm_name) && return mm(expr)
#         in(expr.args[1],mmR_name) && return mmR(expr)
#         (expr.args[1] == :binomial) && (expr.args[3] == 1) && return expr.args[2]
#         #@isdefined($(expr.args[1])) || error("Function $(expr.args[1]) not defined.")
#     end
#     return expr
# end

# #Recursively traverses an expression and replace instances of variables and parmaters with things that the DifferentialEquations packakes simulation algorithms can understand. E.g. X --> u[1], kB1 --> p[1] etc.
# function recursive_replace!(expr::ExprValues, replace_requests::Tuple{OrderedDict{Symbol,Int},Symbol}...)
#     if typeof(expr) == Symbol
#         for rr in replace_requests
#             (haskey(rr[1],expr)) && (return :($(rr[2])[$(rr[1][expr])]))
#         end
#     elseif typeof(expr) == Expr
#         for i = 1:length(expr.args)
#             expr.args[i] = recursive_replace!(expr.args[i], replace_requests...)
#         end
#     end
#     return expr
# end

# #Recursively traverses an expression and replaces a symbol with another.
# function recursive_replace!(expr::ExprValues, replace_requests::Dict{Symbol,Symbol})
#     if typeof(expr) == Symbol
#         haskey(replace_requests,expr) && return replace_requests[expr]
#     elseif typeof(expr) == Expr
#         for i = 1:length(expr.args)
#             expr.args[i] = recursive_replace!(expr.args[i], replace_requests)
#         end
#     end
#     return expr
# end

# #Recursive Contains, checks whenever an expression contains a certain symbol.
# function recursive_contains(s,ex)
#     (typeof(ex)!=Expr) && (return s==ex)
#     for arg in ex.args
#         recursive_contains(s,arg) && (return true)
#     end
#     return false
# end

# #Parses an expression, and returns a set with all symbols in the expression, which is also a part of the provided vector with symbols (syms).
# function recursive_content(ex,syms::Vector{Symbol},content::Vector{Symbol})
#     if ex isa Symbol
#         in(ex,syms) && push!(content,ex)
#     elseif ex isa Expr
#         foreach(arg -> recursive_content(arg,syms,content), ex.args)
#     end
#     return content
# end

# function recursive_content(ex,symsmap::OrderedDict{Symbol,Int},content::Vector{Symbol})
#     if ex isa Symbol
#         haskey(symsmap,ex) && push!(content,ex)
#     elseif ex isa Expr
#         foreach(arg -> recursive_content(arg,symsmap,content), ex.args)
#     end
#     return content
# end

# #Turns an array of expressions to a expression block with corresponding expressions.
# function expr_arr_to_block(exprs)
#     block = :(begin end)
#     foreach(expr -> push!(block.args, expr), exprs)
#     return block
#   end


text:
# ---
# title: JuliaRL\_NFSP\_KuhnPoker
# cover: assets/logo.svg
# description: NFSP applied to KuhnPokerEnv
# date: 2021-08-07
# author: "[Peter Chen](https://github.com/peterchen96)"
# ---

#+ tangle=true
using ReinforcementLearning
using StableRNGs
using Flux
using Flux.Losses

mutable struct KuhnNFSPHook <: AbstractHook
    eval_freq::Int
    episode_counter::Int
    episode::Vector{Int}
    results::Vector{Float64}
end

function (hook::KuhnNFSPHook)(::PostEpisodeStage, policy, env)
    hook.episode_counter += 1
    if hook.episode_counter % hook.eval_freq == 0
        push!(hook.episode, hook.episode_counter)
        push!(hook.results, RLZoo.nash_conv(policy, env))
    end
end

function RL.Experiment(
    ::Val{:JuliaRL},
    ::Val{:NFSP},
    ::Val{:KuhnPoker},
    ::Nothing;
    seed = 123,
)
    rng = StableRNG(seed)

    ## Encode the KuhnPokerEnv's states for training.
    env = KuhnPokerEnv()
    wrapped_env = StateTransformedEnv(
        env;
        state_mapping = s -> [findfirst(==(s), state_space(env))],
        state_space_mapping = ss -> [[findfirst(==(s), state_space(env))] for s in state_space(env)]
        )
    player = 1 # or 2
    ns, na = length(state(wrapped_env, player)), length(action_space(wrapped_env, player))

    ## construct rl_agent(use `DQN`) and sl_agent(use `BehaviorCloningPolicy`)
    rl_agent = Agent(
        policy = QBasedPolicy(
            learner = DQNLearner(
                approximator = NeuralNetworkApproximator(
                    model = Chain(
                        Dense(ns, 64, relu; init = glorot_normal(rng)),
                        Dense(64, na; init = glorot_normal(rng))
                    ) |> cpu,
                    optimizer = Descent(0.01),
                ),
                target_approximator = NeuralNetworkApproximator(
                    model = Chain(
                        Dense(ns, 64, relu; init = glorot_normal(rng)),
                        Dense(64, na; init = glorot_normal(rng))
                    ) |> cpu,
                ),
                γ = 1.0f0,
                loss_func = huber_loss,
                batch_size = 128,
                update_freq = 128,
                min_replay_history = 1000,
                target_update_freq = 1000,
                rng = rng,
            ),
            explorer = EpsilonGreedyExplorer(
                kind = :linear,
                ϵ_init = 0.06,
                ϵ_stable = 0.001,
                decay_steps = 1_000_000,
                rng = rng,
            ),
        ),
        trajectory = CircularArraySARTTrajectory(
            capacity = 200_000,
            state = Vector{Int} => (ns, ),
        ),
    )

    sl_agent = Agent(
        policy = BehaviorCloningPolicy(;
            approximator = NeuralNetworkApproximator(
                model = Chain(
                        Dense(ns, 64, relu; init = glorot_normal(rng)),
                        Dense(64, na; init = glorot_normal(rng))
                    ) |> cpu,
                optimizer = Descent(0.01),
            ),
            explorer = WeightedSoftmaxExplorer(),
            batch_size = 128,
            min_reservoir_history = 1000,
            rng = rng,
        ),
        trajectory = ReservoirTrajectory(
            2_000_000;# reservoir capacity
            rng = rng,
            :state => Vector{Int},
            :action => Int,
        ),
    )

    ## set parameters and initial NFSPAgentManager
    η = 0.1 # anticipatory parameter
    nfsp = NFSPAgentManager(
        Dict(
            (player, NFSPAgent(
                deepcopy(rl_agent),
                deepcopy(sl_agent),
                η,
                rng,
                128, # update_freq
                0, # initial update_step
                true, # initial NFSPAgent's training mode
            )) for player in players(wrapped_env) if player != chance_player(wrapped_env)
        )
    )

    stop_condition = StopAfterEpisode(1_200_000, is_show_progress=!haskey(ENV, "CI"))
    hook = KuhnNFSPHook(10_000, 0, [], [])

    Experiment(nfsp, wrapped_env, stop_condition, hook, "# run NFSP on KuhnPokerEnv")
end

#+ tangle=false
using Plots
ex = E`JuliaRL_NFSP_KuhnPoker`
run(ex)
plot(ex.hook.episode, ex.hook.results, xaxis=:log, xlabel="episode", ylabel="nash_conv")

savefig("assets/JuliaRL_NFSP_KuhnPoker.png")#hide

# ![](assets/JuliaRL_NFSP_KuhnPoker.png)


text:
# Definitions of layers
# This stuff should be removed

"""
    DoubleLayer(b::Body/BodyList,g::PhysicalGrid,u::GridData)

Construct a double-layer operator for a body or bodies `b`. When the
resulting operator acts upon scalar point data `f`, it returns scalar
grid data of the form ``D R_F (f\\circ n)``, where ``D`` is the discrete
divergence operator on grid `g`, ``n`` are the normals associated with
body/bodies `b`, and ``R_F`` is the regularization operator to (edge data on)
`g`. When `f` is of vector point type, then it returns tensor grid data
of the form ``D R_T (f\\times n + n\\times f)``, where ``\\times`` is
the Cartesian product, and ``R_T`` is the regularization operator to
edge gradient data on `g`.
"""
mutable struct DoubleLayer{N,D,G,P} <: LayerType{N}
    weight :: Float64
    nds :: VectorData{N,Float64,D}
    H :: RegularizationMatrix{G,P}
    Pbuf :: P
    Qbuf :: P
    Gbuf :: G
end

function DoubleLayer(body::Union{Body,BodyList},H::RegularizationMatrix{G,P};weight::Float64 = 1.0) where {G,P}
  nrm = normals(body)*weight
  Pbuf = P() #_allocate_point_data(P)
  Qbuf = P()
  Gbuf = G() #_allocate_grid_data(G)
  return DoubleLayer(weight,nrm,H,Pbuf,Qbuf,Gbuf)
end

"""
    DoubleLayer!(dl::DoubleLayer,body::Union{Body,BodyList},g::PhysicalGrid)
"""
function DoubleLayer!(dl::DoubleLayer{N,D,G,P},body::Union{Body,BodyList},H::RegularizationMatrix{G,P}) where {N,D,G,P}
  numpts(body) == N || error("Inconsistent number of points in body")
  dl.nds = normals(body)*dl.weight
  dl.H = H
  return dl
end

#_allocate_point_data(P::Type{VectorData{N}}) where {N} = VectorData(N,dtype=eltype(P))
#_allocate_point_data(P::Type{TensorData{N}}) where {N} = TensorData(N,dtype=eltype(P))
#_allocate_grid_data(G::Type{Edges{C,NX,NY}}) where {C,NX,NY} = Edges(C,(NX,NY),dtype=eltype(G))
#_allocate_grid_data(G::Type{EdgeGradient{C,D,NX,NY}}) where {C,D,NX,NY} = EdgeGradient(C,(NX,NY),dtype=eltype(G))

function DoubleLayer(body::Union{Body,BodyList},g::PhysicalGrid,w::ScalarGridData{NX,NY,T};
                      ddftype=CartesianGrids.Yang3) where {NX,NY,T}
  reg = _get_regularization(body,g,ddftype,GridScaling)
  out = _regularization_matrix(reg,VectorData(numpts(body),dtype=T),Edges(celltype(w),w,dtype=T))
  #return DoubleLayer(body,out, weight = 1/cellsize(g))
  return DoubleLayer(body,out, weight = 1.0)
end

function DoubleLayer(body::Union{Body,BodyList},g::PhysicalGrid,w::VectorGridData{NX,NY,T};
                      ddftype=CartesianGrids.Yang3) where {NX,NY,T}
  reg = _get_regularization(body,g,ddftype,GridScaling)
  out = _regularization_matrix(reg,TensorData(numpts(body),dtype=T),EdgeGradient(celltype(w),w,dtype=T))
  #return DoubleLayer(body,out, weight = 1/cellsize(g))
  return DoubleLayer(body,out, weight = 1.0)
end

function DoubleLayer!(dl::DoubleLayer{N,D,G,P},body::Union{Body,BodyList},g::PhysicalGrid;
                      ddftype=CartesianGrids.Yang3) where {N,D,G,P}

  reg = _get_regularization(body,g,ddftype,GridScaling)
  DoubleLayer!(dl,body,_regularization_matrix(reg,dl.Pbuf,dl.Gbuf))
  return dl
end

# in-place
function (μ::DoubleLayer{N,D,<:Edges{C}})(u::Nodes{C},p::ScalarData{N}) where {N,D,C}
    product!(μ.Pbuf,p,μ.nds)
    divergence!(u,mul!(μ.Gbuf,μ.H,μ.Pbuf))
    return u
end

function (μ::DoubleLayer{N,D,<:EdgeGradient{C}})(u::Edges{C},p::VectorData{N}) where {N,D,C}
    pointwise_tensorproduct!(μ.Pbuf,p,μ.nds)
    transpose!(μ.Qbuf,μ.Pbuf)
    μ.Pbuf .+= μ.Qbuf
    divergence!(u,mul!(μ.Gbuf,μ.H,μ.Pbuf))
    return u
end

#(μ::DoubleLayer{N})(p::ScalarData{N}) where {N} = divergence(μ.H*(p∘μ.nds))
# out of place
(μ::DoubleLayer{N,D,G,P})(p::ScalarData{N,DG}) where {N,D,G<:Edges{C,NX,NY},P<:PointData{N,DG1},DG} where {C,NX,NY,DG1} =
          μ(Nodes(C,(NX,NY),dtype=DG1),p)

#(μ::DoubleLayer{N})(p::VectorData{N}) where {N} = divergence(μ.H*(p*μ.nds+μ.nds*p))
(μ::DoubleLayer{N,D,G,P})(p::VectorData{N,DG}) where {N,D,G<:EdgeGradient{C,F,NX,NY},P<:PointData{N,DG1},DG} where {C,F,NX,NY,DG1} =
          μ(Edges(C,(NX,NY),dtype=DG1),p)

function (μ::DoubleLayer{N,D,G,P})(p::Number) where {N,D,G<:GridData,P<:PointData}
  ϕ = ScalarData(N,dtype=eltype(P))
  ϕ .= p
  return μ(ϕ)
end

function Base.show(io::IO, H::DoubleLayer{N,D,G,P}) where {N,D,G<:GridData{NX,NY,DG},P<:PointData{N,DP}} where {NX,NY,DG,DP}
    println(io, "Double-layer operator")
    println(io, "  from $N point data of $P type")
    println(io, "  to a $NX x $NY grid of $G data")
end

"""
    SingleLayer(b::Body/BodyList,g::PhysicalGrid,u::GridData)

Construct a single-layer operator for a body or bodies `b`. When the
resulting operator acts upon scalar point data `f`, it returns
scalar grid data of form ``R_C f``, where ``R_C`` is the
regularization operator to (node data on) `g`. When it acts upon
vector point data `f`, it returns vector (edge) grid data of the
form ``R_F f``, where ``R_F`` is the regularization operato to edge
data on `g`.
"""
mutable struct SingleLayer{N,D,G,P} <: LayerType{N}
    weight :: Float64
    ds :: ScalarData{N,Float64,D}
    H :: RegularizationMatrix{G,P}
    Pbuf :: P
    Gbuf :: G
end

function SingleLayer(body::Union{Body,BodyList},H::RegularizationMatrix{G,P};weight::Float64 = 1.0) where {G,P}
  ds = ScalarData(numpts(body))
  ds .= weight
  Pbuf = P()
  Gbuf = G()
  return SingleLayer(weight,ds,H,Pbuf,Gbuf)
end

function SingleLayer!(sl::SingleLayer{N,D,G,P},body::Union{Body,BodyList},H::RegularizationMatrix{G,P}) where {N,D,G,P}
  numpts(body) == N || error("Inconsistent number of points in body")
  sl.H = H
  return sl
end

function SingleLayer(body::Union{Body,BodyList},g::PhysicalGrid,w::GridData{NX,NY,T};
                      ddftype=CartesianGrids.Yang3) where {NX,NY,T}
  reg = _get_regularization(body,g,ddftype,GridScaling)
  out = _regularization_matrix(reg,ScalarData(numpts(body),dtype=T),Nodes(celltype(w),w,dtype=T))
  return SingleLayer(body,out) #,weight=cellsize(g)^2)
end

function SingleLayer!(sl::SingleLayer{N,D,G,P},body::Union{Body,BodyList},g::PhysicalGrid;
                      ddftype=CartesianGrids.Yang3) where {N,D,G,P}

  reg = _get_regularization(body,g,ddftype,GridScaling)
  SingleLayer!(sl,body,_regularization_matrix(reg,sl.Pbuf,sl.Gbuf))
  return sl
end

# In place
function (μ::SingleLayer{N,D,G})(u::G,p::ScalarData{N}) where {N,D,G}
    product!(μ.Pbuf,p,μ.ds)
    mul!(u,μ.H,μ.Pbuf)
    return u
end

# Out of place
(μ::SingleLayer{N,D,G})(p::ScalarData{N}) where {N,D,G} = μ(G(),p)
#(μ::SingleLayer{N})(p::ScalarData{N}) where {N} = μ.H*(p∘μ.ds)

function (μ::SingleLayer{N,D,G,P})(p::Number) where {N,D,G<:GridData,P<:PointData}
  μ.Pbuf .= p
  return μ(μ.Pbuf)
end

function Base.show(io::IO, H::SingleLayer{N,D,G,P}) where {N,D,G<:GridData{NX,NY,DG},P<:PointData{N,DP}} where {NX,NY,DG,DP}
    println(io, "Single-layer operator")
    println(io, "  from $N point data of $P type")
    println(io, "  to a $NX x $NY grid of $G data")
end

abstract type MaskType end

struct Mask{N,NX,NY,G} <: MaskType
  data :: Nodes{G,NX,NY}
end

struct ComplementaryMask{N,NX,NY,G} <: MaskType
  mask :: Mask{N,NX,NY,G}
end

"""
    Mask(b::Body,g::PhysicalGrid,w::GridData)

Returns a `Nodes` mask that sets every grid point equal to 1 if it lies inside the body
`b` and 0 outside. Returns grid data of the same type as `w`, corresponding to
physical grid `g`.
"""
function Mask(dlayer::DoubleLayer{N,D,G,P},g::PhysicalGrid) where {N,D,G<:GridData{NX,NY,DG},P<:PointData{N,DP}} where {NX,NY,DG,DP}
  L = plan_laplacian(CartesianGrids.node_inds(celltype(G),size(g)),with_inverse=true,dtype=DG)
  #return Mask{N,NX,NY,G}(-cellsize(g)^2*(L\dlayer(1)))
  return Mask{N,NX,NY,celltype(G)}(-cellsize(g)*(L\dlayer(1)))
end

Mask(body::Union{Body,BodyList},g::PhysicalGrid,w::GridData{NX,NY,T}) where {NX,NY,T} =
    Mask(DoubleLayer(body,g,w),g)

(m::Mask{N,NX,NY,G})(w::Nodes{G,NX,NY,T}) where {N,NX,NY,G,T} = m.data ∘ w

(m::Mask{N,NX,NY,G})(out::Nodes{G,NX,NY,T},w::Nodes{G,NX,NY,T}) where {N,NX,NY,G,T} = product!(out,m.data,w)

(m::ComplementaryMask{N,NX,NY,G})(w::Nodes{G,NX,NY,T}) where {N,NX,NY,G,T} = w - m.mask(w)

(m::ComplementaryMask{N,NX,NY,G})(out::Nodes{G,NX,NY,T},w::Nodes{G,NX,NY,T}) where {N,NX,NY,G,T} = (m.mask(out,w); out .= w - out)


text:
# VLSV reader in Julia

include("vlsvvariables.jl")

"Velocity mesh information."
struct VMeshInfo
   "number of velocity blocks"
   vblocks::NTuple{3, Int64}
   vblock_size::NTuple{3, Int64}
   vmin::NTuple{3, Float64}
   vmax::NTuple{3, Float64}
   dv::NTuple{3, Float64}
end

"Variable information from the VLSV footer."
struct VarInfo
   "unit of the variable as a string"
   unit::String
   "unit of the variable as a LaTeX-formatted string"
   unitLaTeX::LaTeXString
   "description of the variable as a LaTeX-formatted string"
   variableLaTeX::LaTeXString
   "conversion factor to SI units as a string"
   unitConversion::String
end

"VLSV meta data."
struct MetaVLSV
   name::String
   dir::String
   fid::IOStream
   footer::EzXML.Node
   variable::Vector{String}
   "unsorted cell IDs"
   cellid::Vector{UInt64}
   "ordered sequence indexes of raw cell IDs"
   cellindex::Vector{Int64}
   time::Float64
   maxamr::Int64
   hasvdf::Bool
   ncells::NTuple{3, Int64}
   block_size::NTuple{3, Int64}
   coordmin::NTuple{3, Float64}
   coordmax::NTuple{3, Float64}
   dcoord::NTuple{3, Float64}
   species::Vector{String}
   meshes::Dict{String, VMeshInfo}
end

function Base.show(io::IO, meta::MetaVLSV)
   println(io, "File: ", meta.name)
   println(io, "Time: ", round(meta.time, digits=2))
   println(io, "Dimension: ", ndims(meta))
   println(io, "Maximum AMR level: ", meta.maxamr)
   println(io, "Contains VDF: ", meta.hasvdf)
   println(io, "Variables: ", meta.variable)
end

function Base.show(io::IO, s::VarInfo)
   println(io, "Variable in LaTeX: ", s.variableLaTeX)
   println(io, "Unit: ", s.unit)
   println(io, "Unit in LaTeX: ", s.unitLaTeX)
   println(io, "Unit conversion: ", s.unitConversion)
end

function Base.show(io::IO, vmesh::VMeshInfo)
   println(io, "vblocks: ", vmesh.vblocks)
   println(io, "vblock size: ", vmesh.vblock_size)
   foreach((vmin,dv,vmax,comp) -> println(io, "v$comp range: ", vmin, ":", dv, ":", vmax),
      vmesh.vmin, vmesh.dv, vmesh.vmax, 'x':'z')
end

"Return the XML footer of opened VLSV file."
function getfooter(fid::IOStream)
   # First 8 bytes indicate big-endian or else
   endian_offset = 8
   seek(fid, endian_offset)
   # Obtain the offset of the XML footer
   offset = read(fid, UInt64)
   seek(fid, offset)
   footer = read(fid, String) |> parsexml |> root
end

"Return size and type information for the inquired parameter/data with `name`."
function getObjInfo(footer, name, tag, attr)
   local arraysize, datasize, datatype, vectorsize, variable_offset
   isFound = false

   for var in findall("//$tag", footer)
      if var[attr] == String(name)
         arraysize = parse(Int, var["arraysize"])
         datasize = parse(Int, var["datasize"])
         datatype = var["datatype"]
         vectorsize = parse(Int, var["vectorsize"])
         variable_offset = parse(Int, nodecontent(var))
         isFound = true
         break
      end
   end

   isFound || throw(ArgumentError("unknown variable $name"))

   T::Type =
      if datatype == "float"
         datasize == 4 ? Float32 : Float64
      elseif datatype == "int"
         datasize == 4 ? Int32 : Int64
      elseif datatype == "uint"
         datasize == 4 ? UInt32 : UInt64
      end

   T, variable_offset, arraysize, datasize, vectorsize
end

"Return vector of `name` from the VLSV file with `footer` associated with stream `fid`."
function readvector(fid::IOStream, footer, name, tag)
   T, offset, asize, dsize, vsize = getObjInfo(footer, name, tag, "name")

   if Sys.total_memory() > 8*asize*vsize*dsize
      w = vsize == 1 ?
         Vector{T}(undef, asize) :
         Array{T,2}(undef, vsize, asize)
      seek(fid, offset)
      read!(fid, w)
   else
      @warn "Large array detected. Using memory-mapped I/O!" maxlog=1
      a = mmap(fid, Vector{UInt8}, dsize*vsize*asize, offset)
      w = vsize == 1 ?
         reinterpret(T, a) :
         reshape(reinterpret(T, a), vsize, asize)
   end

   w
end

"""
    load(file) -> MetaVLSV

Generate a MetaVLSV object from `file` of VLSV format.
"""
function load(file::AbstractString)
   isfile(file) || throw(ArgumentError("Cannot open \'$file\': not a file"))
   fid = open(file, "r")

   footer = getfooter(fid)

   local cellid
   let
      T, offset, asize, dsize, vsize = getObjInfo(footer, "CellID", "VARIABLE", "name")
      a = mmap(fid, Vector{UInt8}, dsize*vsize*asize, offset)
      cellid = reinterpret(T, a)
   end

   cellindex = sortperm(cellid)

   bbox = readmesh(fid, footer, "SpatialGrid", "MESH_BBOX")::Vector{UInt}

   nodeCoordsX = readmesh(fid, footer, "SpatialGrid", "MESH_NODE_CRDS_X")::Vector{Float64}
   nodeCoordsY = readmesh(fid, footer, "SpatialGrid", "MESH_NODE_CRDS_Y")::Vector{Float64}
   nodeCoordsZ = readmesh(fid, footer, "SpatialGrid", "MESH_NODE_CRDS_Z")::Vector{Float64}

   @inbounds ncells = (bbox[1], bbox[2], bbox[3])
   @inbounds block_size = (bbox[4], bbox[5], bbox[6])
   @inbounds coordmin = (nodeCoordsX[begin], nodeCoordsY[begin], nodeCoordsZ[begin])
   @inbounds coordmax = (nodeCoordsX[end], nodeCoordsY[end], nodeCoordsZ[end])

   dcoord = ntuple(i -> (coordmax[i] - coordmin[i]) / ncells[i], Val(3))

   meshes = Dict{String, VMeshInfo}()

   # Find all species by the BLOCKIDS tag
   species = String[]

   for varinfo in findall("//BLOCKIDS", footer)

      if haskey(varinfo, "name")
         # VLSV 5.0 file with bounding box
         popname = varinfo["name"]

         bbox = readmesh(fid, footer, popname, "MESH_BBOX")::Vector{UInt}

         nodeCoordsX = readmesh(fid, footer, popname, "MESH_NODE_CRDS_X")::Vector{Float64}
         nodeCoordsY = readmesh(fid, footer, popname, "MESH_NODE_CRDS_Y")::Vector{Float64}
         nodeCoordsZ = readmesh(fid, footer, popname, "MESH_NODE_CRDS_Z")::Vector{Float64}
         vblocks = (bbox[1], bbox[2], bbox[3])
         vblock_size = (bbox[4], bbox[5], bbox[6])
         vmin = (nodeCoordsX[begin], nodeCoordsY[begin], nodeCoordsZ[begin])
         vmax = (nodeCoordsX[end], nodeCoordsY[end], nodeCoordsZ[end])
         dv = ntuple(i -> (vmax[i] - vmin[i]) / vblocks[i] / vblock_size[i], Val(3))
      else
         popname = "avgs"

         if "vxblocks_ini" in getindex.(findall("//PARAMETER", footer), "name")
            # In VLSV before 5.0 the mesh is defined with parameters.
            vblocks_str = ("vxblocks_ini", "vyblocks_ini", "vzblocks_ini")
            vmin_str = ("vxmin", "vymin", "vzmin")
            vmax_str = ("vxmax", "vymax", "vzmax")
            vblocks = ntuple(i -> readparameter(fid, footer, vblocks_str[i]), Val(3))
            vblock_size = (4, 4, 4)
            vmin = ntuple(i -> readparameter(fid, footer, vmin_str[i]), Val(3))
            vmax = ntuple(i -> readparameter(fid, footer, vmax_str[i]), Val(3))
            dv = ntuple(i -> (vmax[i] - vmin[i]) / vblocks[i] / vblock_size[i], Val(3))
         else
            # No velocity space info, e.g., file not written by Vlasiator
            vblocks = (0, 0, 0)
            vblock_size = (4, 4, 4)
            vmin = (0, 0, 0)
            vmax = (0, 0, 0)
            dv = (1, 1, 1)
         end
      end

      # Update list of active species
      if popname ∉ species
         push!(species, popname)
      end

      # Create a new object for this population
      popVMesh = VMeshInfo(vblocks, vblock_size, vmin, vmax, dv)

      meshes[popname] = popVMesh
   end

   if hasname(footer, "PARAMETER", "time") # Vlasiator 5.0+
      timesim = readparameter(fid, footer, "time")::Float64
   elseif hasname(footer, "PARAMETER", "t")
      timesim = readparameter(fid, footer, "t")::Float64
   else
      timesim = Inf
   end

   # Obtain maximum refinement level
   maxamr = getmaxrefinement(cellid, cellindex, ncells)

   varinfo = findall("//VARIABLE", footer)
   vars = [info["name"] for info in varinfo]

   hasvdf = let
      vcells = readmesh(fid, footer, "SpatialGrid", "CELLSWITHBLOCKS")::Vector{UInt}
      !isempty(vcells)
   end

   # File IOstream is not closed for sake of data processing later.

   meta = MetaVLSV(basename(file), dirname(file), fid, footer, vars, cellid,
      cellindex, timesim, maxamr, hasvdf, ncells, block_size, coordmin, coordmax,
      dcoord, species, meshes)
end

# Allow do ... end syntax.
function load(f::Function, file::AbstractString)
   meta = load(file)
   try
      f(meta)
   finally
      close(meta.fid)
   end
end

function getmaxrefinement(cellid, cellindex, ncells)
   ncell = prod(ncells)
   maxamr, cid = 0, ncell
   while @inbounds cid < cellid[cellindex[end]]
      maxamr += 1
      cid += ncell*8^maxamr
   end
   maxamr
end

"""
    readvariablemeta(meta, var) -> VarInfo

Return VarInfo about `var` in the VLSV file associated with `meta`.
"""
function readvariablemeta(meta::MetaVLSV, var)

   varSym = isa(var, AbstractString) ? Symbol(var) : var

   unit, unitLaTeX, variableLaTeX, unitConversion = "", "", "", ""

   if varSym in keys(units_predefined)
      unit, variableLaTeX, unitLaTeX = units_predefined[varSym]
   elseif hasvariable(meta, var) # For Vlasiator 5 files, MetaVLSV is included
      for varinfo in findall("//VARIABLE", meta.footer)
         if varinfo["name"] == var
            haskey(varinfo, "unit") || break
            unit = varinfo["unit"]
            unitLaTeX = varinfo["unitLaTeX"]
            variableLaTeX = varinfo["variableLaTeX"]
            unitConversion = varinfo["unitConversion"]
         end
      end
   end

   VarInfo(unit, unitLaTeX, variableLaTeX, unitConversion)
end

"Return mesh related variable."
function readmesh(fid::IOStream, footer, typeMesh, varMesh)
   T, offset, arraysize, _, _ = getObjInfo(footer, typeMesh, varMesh, "mesh")

   w = Vector{T}(undef, arraysize)
   seek(fid, offset)
   read!(fid, w)

   w
end

"""
    readvariable(meta::MetaVLSV, var, sorted::Bool=true) -> Array

Return variable value of `var` from the VLSV file associated with `meta`. By default
`sorted=true`, which means that for DCCRG grid the variables are sorted by cell ID.
"""
function readvariable(meta::MetaVLSV, var, sorted::Bool=true)
   (;fid, footer, cellindex) = meta
   if (local symvar = Symbol(var)) in keys(variables_predefined)
      data = variables_predefined[symvar](meta)
      return data
   end

   raw = readvector(fid, footer, var, "VARIABLE")

   if startswith(var, "fg_") # fsgrid
      bbox = readmesh(fid, footer, "fsgrid", "MESH_BBOX")::Vector{Int}
      # Determine fsgrid domain decomposition
      nIORanks = readparameter(meta, "numWritingRanks")::Int32

      dataOrdered =
         if ndims(raw) > 1
            @inbounds zeros(Float32, size(raw,1), bbox[1], bbox[2], bbox[3])
         else
            @inbounds zeros(Float32, bbox[1], bbox[2], bbox[3])
         end

      @inbounds fgDecomposition = @views getDomainDecomposition(bbox[1:3], nIORanks)

      _fillFGordered!(dataOrdered, raw, fgDecomposition, nIORanks, bbox)

      data = dropdims(dataOrdered, dims=(findall(size(dataOrdered) .== 1)...,))
   elseif sorted # dccrg grid
      @inbounds d = ndims(raw) == 1 ? raw[cellindex] : raw[:,cellindex]
      data = eltype(raw) == Float64 ? Float32.(d) : d
   else
      data = raw
   end
   return data
end

"""
    readvariable(meta::MetaVLSV, var, ids) -> Array

Read variable `var` in a collection of cells `ids` associated with `meta`.
"""
function readvariable(meta::MetaVLSV, var, ids)
   startswith(var, "fg_") && error("Currently does not support reading fsgrid!")
   (;fid, footer, cellid) = meta

   if (local symvar = Symbol(var)) in keys(variables_predefined)
      data = variables_predefined[symvar](meta, ids)
      return data
   end

   T, offset, asize, dsize, vsize = getObjInfo(footer, var, "VARIABLE", "name")

   v = Array{T}(undef, vsize, length(ids))

   w = let
      a = mmap(fid, Vector{UInt8}, dsize*vsize*asize, offset)
      reshape(reinterpret(T, a), vsize, asize)
   end

   id_ = length(ids) < 1000 ?
      [findfirst(==(id), cellid) for id in ids] :
      indexin(ids, cellid)

   _fillv!(v, w, id_, vsize)

   if T == Float64
      v = Float32.(v)
   end
   return v
end

@inline function _fillv!(v, w, id_, vsize)
   for i in eachindex(id_), iv = 1:vsize
      @inbounds v[iv,i] = w[iv,id_[i]]
   end
end

function _fillFGordered!(dataOrdered, raw, fgDecomposition, nIORanks, bbox)

   offsetnow = 1

   @inbounds @views for i = 1:nIORanks
      xyz = (
         (i - 1) ÷ fgDecomposition[3] ÷ fgDecomposition[2],
         (i - 1) ÷ fgDecomposition[3] % fgDecomposition[2],
         (i - 1) % fgDecomposition[3] )

      lsize = ntuple(i -> calcLocalSize(bbox[i], fgDecomposition[i], xyz[i]), Val(3))
      lstart = ntuple(i -> calcLocalStart(bbox[i], fgDecomposition[i], xyz[i]), Val(3))
      offsetnext = offsetnow + prod(lsize)
      lend = @. lstart + lsize - 1
      lrange = map((x,y)->x:y, lstart, lend)
      # Reorder data
      if ndims(raw) > 1
         ldata = reshape(raw[:,offsetnow:offsetnext-1], size(raw,1), lsize...)

         dataOrdered[:,lrange...] = ldata
      else
         ldata = reshape(raw[offsetnow:offsetnext-1], lsize...)

         dataOrdered[lrange...] = ldata
      end
      offsetnow = offsetnext
   end
   return
end

@inline Base.getindex(meta::MetaVLSV, key::AbstractString) = readvariable(meta, key)

"Return 2d scalar/vector data. Nonpublic since it won't work with DCCRG AMR."
function getdata2d(meta::MetaVLSV, var)
   ndims(meta) == 2 || @error "2D outputs required."
   sizes = filter(!=(1), meta.ncells)
   data = readvariable(meta, var)
   data = ndims(data) == 1 || size(data, 1) == 1 ?
      reshape(data, sizes[1], sizes[2]) :
      reshape(data, 3, sizes[1], sizes[2]) # assumes 3-vector, may not work in general
end

"File size in bytes."
@inline Base.size(meta::MetaVLSV) = filesize(joinpath(meta.dir, meta.name))

# Optimize decomposition of this grid over the given number of processors.
# Reference: fsgrid.hpp
function getDomainDecomposition(globalsize, nprocs)
   domainDecomp = (1, 1, 1)
   minValue = typemax(Int)

   @inbounds for i = 1:min(nprocs, globalsize[1])
      iBox = max(globalsize[1]/i, 1)

      for j = 1:min(nprocs, globalsize[2])
         i * j > nprocs && break

         jBox = max(globalsize[2]/j, 1)

         for k = 1:min(nprocs, globalsize[2])
            i * j * k > nprocs && continue

            kBox = max(globalsize[3]/k, 1)

            nyz = ifelse(i > 1, jBox * kBox, 0)
            nzx = ifelse(j > 1, kBox * iBox, 0)
            nxy = ifelse(k > 1, iBox * jBox, 0)

            v = 10*iBox*jBox*kBox + nyz + nzx + nxy

            if i * j * k == nprocs && v < minValue
               minValue = v
               domainDecomp = (i, j, k)
            end
         end
      end
   end

   domainDecomp
end

function calcLocalStart(globalCells, nprocs, lcells)
   ncells = globalCells ÷ nprocs
   remainder = globalCells % nprocs
   lstart = ifelse(lcells < remainder, lcells*(ncells+1) + 1, lcells*ncells + remainder + 1)
end

function calcLocalSize(globalCells, nprocs, lcells)
   ncells = globalCells ÷ nprocs
   remainder = globalCells % nprocs
   lsize = ifelse(lcells < remainder, ncells + 1, ncells)
end

"""
    hasvariable(meta, var) -> Bool

Check if the VLSV file associated with `meta` contains a variable `var`.
"""
hasvariable(meta::MetaVLSV, var) = hasname(meta.footer, "VARIABLE", var)

"""
    readparameter(meta, param)

Return the parameter value from the VLSV file associated with `meta`.
"""
readparameter(meta::MetaVLSV, param) = readparameter(meta.fid, meta.footer, param)

function readparameter(fid::IOStream, footer, param)
   T, offset, _, _, _ = getObjInfo(footer, param, "PARAMETER", "name")
   seek(fid, offset)
   p = read(fid, T)
end

"""
    hasparameter(meta, param) -> Bool

Check if the VLSV file contains a certain parameter `param`.
"""
hasparameter(meta::MetaVLSV, param) = hasname(meta.footer, "PARAMETER", param)

"Check if the XML `element` contains a `tag` with `name`."
function hasname(element, tag, name)
   isFound = false

   for var in findall("//$tag", element)
      var["name"] == name && (isFound = true)
      isFound && break
   end

   isFound
end

"""
    ndims(meta) -> Int

Return the simulation spatial dimension of VLSV data.
"""
Base.ndims(meta::MetaVLSV) = count(>(1), meta.ncells)

"""
    readvcells(meta, cid; species="proton") -> vcellids, vcellf

Read velocity cells of `species` from a spatial cell of ID `cid` associated with `meta`, and
return a map of velocity cell ids `vcellids` and corresponding value `vcellf`.
"""
function readvcells(meta::MetaVLSV, cid; species="proton")
   (;fid, footer) = meta
   (;vblock_size) = meta.meshes[species]
   bsize = prod(vblock_size)

   local offset::Int, nblocks::Int
   let cellsWithVDF = readvector(fid, footer, species, "CELLSWITHBLOCKS"),
      nblock_C = readvector(fid, footer, species, "BLOCKSPERCELL")::Vector{UInt32}

      cellWithVDFIndex = findfirst(==(cid), cellsWithVDF)
      if !isnothing(cellWithVDFIndex)
         @inbounds nblocks = Int(nblock_C[cellWithVDFIndex])
         if nblocks == 0
            throw(ArgumentError("Cell ID $cid does not store velocity distribution!"))
         end
      else
         throw(ArgumentError("Cell ID $cid does not store velocity distribution!"))
      end
      # Offset position to vcell storage
      offset = sum(@view nblock_C[1:cellWithVDFIndex-1])
   end

   local datasize, datatype, vectorsize, variable_offset
   # Read in avgs
   let varinfo = findfirst("//BLOCKVARIABLE", footer)
      datasize = parse(Int, varinfo["datasize"])
      datatype = varinfo["datatype"]
      datatype == "float" || @error "VDFs must be floating numbers!"
      vectorsize = parse(Int, varinfo["vectorsize"])
      variable_offset = parse(Int, nodecontent(varinfo))
   end

   data = let
      Tavg = datasize == 4 ? Float32 : Float64
      a = mmap(fid, Vector{UInt8}, datasize*vectorsize*nblocks,
         offset*vectorsize*datasize + variable_offset)
      reshape(reinterpret(Tavg, a), vectorsize, nblocks)
   end

   # Read in block IDs
   let varinfo = findfirst("//BLOCKIDS", footer)
      datasize = parse(Int, varinfo["datasize"])
      datatype = varinfo["datatype"]
      datatype == "uint" || @error "Block ID must be unsigned integer!"
      variable_offset = parse(Int, nodecontent(varinfo))
   end

   blockIDs = let T = datasize == 4 ? UInt32 : UInt64
      Vector{T}(undef, nblocks)
   end
   seek(fid, offset * datasize + variable_offset)
   read!(fid, blockIDs)

   # Velocity cell IDs and distributions (ordered by blocks)
   vcellids = zeros(UInt32, bsize*nblocks)
   vcellf = zeros(Float32, bsize*nblocks)

   vcellid_local = @inbounds [i + vblock_size[1]*j + vblock_size[1]*vblock_size[2]*k
      for i in 0:vblock_size[1]-1, j in 0:vblock_size[2]-1, k in 0:vblock_size[3]-1]

   _fillvcell!(vcellids, vcellf, vcellid_local, data, blockIDs, bsize)

   vcellids, vcellf
end

@inline function _fillvcell!(vcellids, vcellf, vcellid_local, data, blockIDs, bsize)
   @inbounds for i in eachindex(blockIDs), j = 1:bsize
      index_ = (i-1)*bsize+j
      vcellids[index_] = vcellid_local[j] + bsize*blockIDs[i]
      vcellf[index_] = data[j,i]
   end
end


text: println("Hello, Julia World!")


text:
using Test
using MPI

MPI.Init()

comm = MPI.COMM_WORLD
@test MPI.Comm_compare(comm, comm) == MPI.IDENT
@test MPI.Comm_compare(comm, MPI.COMM_SELF) == (MPI.Comm_size(comm) == 1 ? MPI.CONGRUENT : MPI.UNEQUAL)
MPI.Barrier(comm)
comm2 = MPI.Comm_dup(comm)
@test MPI.Comm_compare(comm, comm2) == MPI.CONGRUENT
MPI.Barrier(comm2)
comm3 = MPI.Comm_dup(comm2)
@test MPI.Comm_compare(comm, comm3) == MPI.CONGRUENT
MPI.Barrier(comm3)
MPI.free(comm2)
MPI.Barrier(comm3)
# Don't free comm2
MPI.free(comm3)

MPI.Finalize()
@test MPI.Finalized()


text:
module SampleImages

using Colors
using Base64

struct BeadsImageSVG <: Main.SVG
    buf::IOBuffer
end

# The following code is ad-hoc and highly depends on "beads.svg".
# Be careful when modifying the svg file or its code.
function BeadsImageSVG(caption::String; filter=nothing, width="64mm", height="36mm")
    io = IOBuffer()
    id = string(hash(caption), base=16)

    open(joinpath("assets", "figures", "beads.svg"), "r") do file
        for line in eachline(file, keep=true)
            occursin("<?xml", line) && continue
            if occursin("<svg", line)
                line = replace(line, ">"=>
                    """ width="$width" height="$height" style="display:inline; margin-left:1em; margin-bottom:1em">""")
            elseif occursin("filter_beads_g", line)
                line = replace(line, "filter_beads_g"=>"filter_"*id)
            elseif occursin("</svg>", line)
                text = """
                       <text x="16" y="1000" style="fill:black;opacity:0.9;font-size:80px;">$caption</text>
                       </svg>"""
                line = replace(line, "</svg>"=>text)
            end

            m = match(r"data:image/png;base64,([^\"]+)", line)
            if filter === nothing || m === nothing
                write(io, line)
                continue
            end

            head = m.offsets[1]
            write(io, SubString(line, 1, head - 1))
            src = IOBuffer(m.captures[1])
            b64dec = Base64DecodePipe(src) # decode all for simplicity
            b64enc = Base64EncodePipe(io)
            write(b64enc, read(b64dec, 33)) # before the length of "PLTE"
            replace_palette(b64enc, b64dec, filter)
            n = write(b64enc, read(b64dec))
            close(b64enc)
            close(b64dec)
            write(io, SubString(line, head + length(m.captures[1]), length(line)))
        end
    end
    BeadsImageSVG(io)
end

function replace_palette(dest::IO, src::IO, filter)
    buf = IOBuffer() # to calculate chunk CRCs

    u8(x) = write(buf, UInt8(x & 0xFF))
    u16(x) = (u8((x & 0xFFFF)>>8); u8(x))
    u32(x) = (u16((x & 0xFFFFFFFF)>>16); u16(x))
    function read_palette()
        uint32 = (UInt32(read(src, UInt8)) << 16) |
                 (UInt32(read(src, UInt8)) <<  8) | read(src, UInt8)
        reinterpret(RGB24, uint32)
    end
    function write_palette(c::Color)
        rgb24 = convert(RGB24,c)
        u8(rgb24.color>>16); u8(rgb24.color>>8); u8(rgb24.color)
    end
    crct(x) = (for i = 1:8; x = x & 1==1 ? 0xEDB88320 ⊻ (x>>1) : x>>1 end; x)
    table = UInt32[crct(i) for i = 0x00:0xFF]
    function crc32()
        seekstart(buf)
        crc = 0xFFFFFFFF
        while !eof(buf)
            crc = (crc>>8) ⊻ table[(crc&0xFF) ⊻ read(buf, UInt8) + 1]
        end
        u32(crc ⊻ 0xFFFFFFFF)
    end
    lenbytes = read(src, 4)
    len = (UInt32(lenbytes[3]) << 8) | lenbytes[4]
    write(dest, lenbytes)
    write(buf, read(src, 4)) # "PLTE"
    for i = 1:(len÷3)
        write_palette(filter(read_palette()))
    end
    read(src, 4) # CRC
    crc32()
    write(dest, take!(seekstart(buf)))
end

end


text:
"""
    rp_split(data, indices) -> lefts, rights, hyperplane, offset

Splits the data into two sets depending on which side of a random hyperplane
each point lies.
Returns the left indices, right indices, hyperplane, and offset.
"""
function rp_split(data, indices)

    # select two random points and set the hyperplane between them
    # TODO: threading
    # TODO: remove StatsBase dep
    leftidx, rightidx = sample(indices, 2; replace=false)
    hyperplane = data[leftidx] - data[rightidx]

    offset = -adjoint(hyperplane) * (data[leftidx] + data[rightidx]) / 2

    lefts = sizehint!(Int32[], div(length(data), 2))
    rights = sizehint!(Int32[], div(length(data), 2))
    # TODO: threading
    for (i, idx) in enumerate(indices)
        # for each data vector, compute which side of the hyperplane
        select_side(hyperplane, offset, data[idx]) ? append!(lefts, idx) : append!(rights, idx)
    end

    return lefts, rights, hyperplane, offset
end

"""
Project `point` onto `hyperplane` and return `true` if it lies to the "left" of
`offset`, otherwise return `false`.
If the point is less than `√eps`, pick a side randomly.
"""
function select_side(hyperplane, offset, point)
    margin = adjoint(hyperplane) * point + offset
    if abs(margin) <= sqrt(eps(margin))
        # TODO: threading
        return rand() < .5
    end
    return margin < 0
end


text:
using Test, FinancialMonteCarlo;
@show "Black Scholes Model"
S0 = 100.0;
K = 100.0;
r = 0.02;
T = 1.0;
d = 0.01;
D = 90.0;

Nsim = 10000;
Nstep = 30;
sigma = 0.2;
mc = MonteCarloConfiguration(Nsim, Nstep);
mc1 = MonteCarloConfiguration(Nsim, Nstep, FinancialMonteCarlo.AntitheticMC());
toll = 0.8

rfCurve = ZeroRate(r);

FwdData = Forward(T)
@test FinancialMonteCarlo.get_parameters(FwdData)[1] == T
EUData = EuropeanOption(T, K)
EUBin = BinaryEuropeanOption(T, K)
AMData = AmericanOption(T, K)
AmBin = BinaryEuropeanOption(T, K)
BarrierData = BarrierOptionDownOut(T, K, D)
AsianFloatingStrikeData = AsianFloatingStrikeOption(T)
AsianFixedStrikeData = AsianFixedStrikeOption(T, K)
Model = BlackScholesProcess(0.1, Underlying(S0, d));

param_ = FinancialMonteCarlo.get_parameters(Model)

@test param_[1] == 0.1

FinancialMonteCarlo.set_parameters!(Model, [sigma])

display(Model)
@show spot = pricer(Model, rfCurve, mc, Spot());
@show FwdPrice = pricer(Model, rfCurve, mc, FwdData);
@show EuPrice = pricer(Model, rfCurve, mc, EUData);
@show EuBinPrice = pricer(Model, rfCurve, mc, EUBin);
@show AmPrice = pricer(Model, rfCurve, mc, AMData);
@show AmBinPrice = pricer(Model, rfCurve, mc, AmBin);
@show BarrierPrice = pricer(Model, rfCurve, mc, BarrierData);
@show AsianPrice1 = pricer(Model, rfCurve, mc, AsianFloatingStrikeData);
@show AsianPrice2 = pricer(Model, rfCurve, mc, AsianFixedStrikeData);

@test abs(FwdPrice - 99.1078451563562) < toll
@test abs(EuPrice - 8.43005524824866) < toll
@test abs(AmPrice - 8.450489415187354) < toll
@test abs(BarrierPrice - 7.5008664470880735) < toll
@test abs(AsianPrice1 - 4.774451704549382) < toll

@show FwdPrice = pricer(Model, rfCurve, mc1, FwdData);
@show EuPrice = pricer(Model, rfCurve, mc1, EUData);
@show AmPrice = pricer(Model, rfCurve, mc1, AMData);
@show BarrierPrice = pricer(Model, rfCurve, mc1, BarrierData);
@show AsianPrice1 = pricer(Model, rfCurve, mc1, AsianFloatingStrikeData);
@show AsianPrice2 = pricer(Model, rfCurve, mc1, AsianFixedStrikeData);
tollanti = 0.6;
@test abs(FwdPrice - 99.1078451563562) < tollanti
@test abs(EuPrice - 8.43005524824866) < tollanti
@test abs(AmPrice - 8.450489415187354) < tollanti
@test abs(BarrierPrice - 7.5008664470880735) < tollanti
@test abs(AsianPrice1 - 4.774451704549382) < tollanti

EUDataPut = EuropeanOption(T, K, false)
AMDataPut = AmericanOption(T, K, false)
BarrierDataPut = BarrierOptionDownOut(T, K, D, false)
AsianFloatingStrikeDataPut = AsianFloatingStrikeOption(T, false)
AsianFixedStrikeDataPut = AsianFixedStrikeOption(T, K, false)

@show EuPrice = pricer(Model, rfCurve, mc, EUDataPut);
@show AmPrice = pricer(Model, rfCurve, mc, AMDataPut);
@show BarrierPrice = pricer(Model, rfCurve, mc, BarrierDataPut);
@show AsianPrice1 = pricer(Model, rfCurve, mc, AsianFloatingStrikeDataPut);
@show AsianPrice2 = pricer(Model, rfCurve, mc, AsianFixedStrikeDataPut);
tollPut = 0.6;
@test abs(FwdPrice - 99.1078451563562) < tollPut
@test abs(EuPrice - 7.342077422567968) < tollPut
@test abs(AmPrice - 7.467065889603365) < tollPut
@test abs(BarrierPrice - 0.29071929104261723) < tollPut
@test abs(AsianPrice1 - 4.230547012372306) < tollPut
@test abs(AsianPrice2 - 4.236220218194027) < tollPut

@test_throws(ErrorException, MonteCarloConfiguration(Nsim + 1, Nstep, FinancialMonteCarlo.AntitheticMC()));

mc1 = MonteCarloConfiguration(Nsim, Nstep, FinancialMonteCarlo.AntitheticMC());
@test variance(Model, rfCurve, mc, EUDataPut) > variance(Model, rfCurve, mc1, EUDataPut);
@test prod(variance(Model, rfCurve, mc, [EUDataPut, AMDataPut]) .>= variance(Model, rfCurve, mc1, [EUDataPut, AMDataPut]));

IC1 = confinter(Model, rfCurve, mc, EUDataPut);
IC2 = confinter(Model, rfCurve, mc1, EUDataPut, 0.99);
L1 = IC1[2] - IC1[1]
L2 = IC2[2] - IC2[1]
@test L1 > L2

IC1 = confinter(Model, rfCurve, mc, [EUDataPut, AMDataPut]);
IC2 = confinter(Model, rfCurve, mc1, [EUDataPut, AMDataPut], 0.99);
L1 = IC1[2][2] - IC1[2][1]
L2 = IC2[2][2] - IC2[2][1]
@test L1 > L2
############## Complete Options

EUDataBin = BinaryEuropeanOption(T, K)
BinAMData = BinaryAmericanOption(T, K)
BarrierDataDI = BarrierOptionDownIn(T, K, D)
BarrierDataUI = BarrierOptionUpIn(T, K, D)
BarrierDataUO = BarrierOptionUpOut(T, K, D)
BermData = BermudanOption(collect(0.2:0.1:T), K);
doubleBarrierOptionDownOut = DoubleBarrierOption(T, K, K / 10.0, 1.2 * K)

@show BarrierPrice = pricer(Model, rfCurve, mc, BarrierDataDI);
@show BarrierPrice = pricer(Model, rfCurve, mc, BarrierDataUI);
@show BarrierPrice = pricer(Model, rfCurve, mc, BarrierDataUO);
@show BermPrice = pricer(Model, rfCurve, mc, BermData);
@show AmBinPrice = pricer(Model, rfCurve, mc, BinAMData);
@show EuBinPrice = pricer(Model, rfCurve, mc, EUDataBin);
@show doubleBarrier = pricer(Model, rfCurve, mc, doubleBarrierOptionDownOut);

@show "Test Black Scholes Parameters"

@test_throws(ErrorException, BlackScholesProcess(-sigma, Underlying(S0, d)))


text:
# Represents a lazy iterator over all solution in a time interval
"""
It represents the solution of a ODE at a given time interval. It is a lazy implementation,
i.e., the object is an iterator that computes the solution at each time step
when accessing the solution at each time step.
"""
abstract type ODESolution <: GridapType end

# First time step
function iterate(u::ODESolution) # (u0,t0)-> (uf,tf) or nothing
  @abstractmethod
end

# Following time steps
function iterate(u::ODESolution,state) # (u0,t0)-> (uf,tf) or nothing
  @abstractmethod
end

# tester

function test_ode_solution(sol::ODESolution)
  for (u_n,t_n) in sol
    @test isa(t_n,Real)
    @test isa(u_n,AbstractVector)
  end
  true
end

# Specialization

struct GenericODESolution <: ODESolution
  solver::ODESolver
  op::ODEOperator
  u0::AbstractVector
  t0::Real
  tF::Real
end

function Base.iterate(sol::GenericODESolution)

  uf = copy(sol.u0)
  u0 = copy(sol.u0)
  t0 = sol.t0

  # Solve step
  uf, tf, cache = solve_step!(uf,sol.solver,sol.op,u0,t0)

  # Update
  u0 .= uf
  state = (uf,u0,tf,cache)

  return (uf, tf), state
end

function Base.iterate(sol::GenericODESolution, state)

  uf,u0,t0,cache = state

  if t0 >= sol.tF - ϵ
    return nothing
  end

  # Solve step
  uf, tf, cache = solve_step!(uf,sol.solver,sol.op,u0,t0,cache)

  # Update
  u0 .= uf
  state = (uf,u0,tf,cache)

  return (uf, tf), state
end


text:
using NNlib

# Activation functions
@cufunc softplus(x::Real) = ifelse(x > 0, x + log1p(exp(-x)), log1p(exp(x)))

@cufunc logσ(x::Real) = -softplus(-x)

@cufunc function gelu(x::Real)
    p = oftype(x / 1, π)
    λ = oftype(x / 1, √(2 / p))
    α = oftype(x / 1, 0.044715)
    h = oftype(x / 1, 0.5)
    h * x * (one(x) + tanh(λ * (x + α * x^3)))
end

@cufunc lisht(x::Real) = x * tanh(x)

@cufunc logcosh(x::Real) = x + softplus(-2x) - log(oftype(x, 2))

@cufunc mish(x::Real) = x * tanh(softplus(x))

@cufunc tanhshrink(x::Real) = x - tanh(x)

# Batched matrix multiplication
# 1st argument is produced by NNlib.storage_type(A)
NNlib._batched_gemm!(::Type{<:CuArray}, transA::Char, transB::Char, α::Number, A, B, β::Number, C) =
     CUBLAS.gemm_strided_batched!(transA, transB, α, A, B, β, C)

Base.unsafe_convert(::Type{CuPtr{T}}, A::NNlib.BatchedAdjOrTrans{T}) where {T} =
    Base.unsafe_convert(CuPtr{T}, parent(A))


text:
function compile(fstr::String)
    nposargs = 0
    argparams = Dict{Argument, Symbol}()
    function arg2param(arg)
        arg isa Argument || return arg  # static argument
        haskey(argparams, arg) && return argparams[arg]
        arg isa Positional && (nposargs = max(arg.position, nposargs))
        param = arg isa Keyword ? arg.name : gensym()
        argparams[arg] = param
        return param
    end

    code_info = Expr(:block)
    code_data = Expr(:block)
    for (i, f) in enumerate(parse(fstr))
        if f isa String
            n = ncodeunits(f)
            info = :(size += $n)
            data = n < 8 ?
                :(@inbounds $(genstrcopy(f)); p += $n) :
                :(copyto!(buf, p, $(codeunits(f)), 1, $n); p += $n)
        else
            @assert f isa Field
            value = esc(arg2param(f.argument))
            fill = esc(arg2param(f.spec.fill))
            width = esc(arg2param(f.spec.width))
            precision = esc(arg2param(f.spec.precision))
            spec = Symbol(:spec, i)
            arg = Symbol(:arg, i)
            meta = Symbol(:meta, i)
            convfun = conv2func(f.conv)
            info = quote
                $spec = Spec($(f.spec), fill = $fill, width = $width, precision = $precision)
                $arg = $(convfun)($value)
                s, $meta = formatinfo($spec, $arg)
                size += s
            end
            data = :(p = formatfield(buf, p, $spec, $arg, $meta))
        end
        push!(code_info.args, info)
        push!(code_data.args, data)
    end

    # function parameters and body
    poparams = [get(argparams, Positional(n), :_) for n in 1:nposargs]
    kwparams = [param for (arg, param) in argparams if arg isa Union{Keyword, Expr}]
    # the parameters are like: function (buf, pos, a, b, c, ...; x, y, z, ...).
    params = Expr(:tuple, Expr(:parameters, esc.(kwparams)...), :buf, :pos, esc.(poparams)...)
    body = quote
        size = 0
        $(code_info)
        if buf === DUMMY_BUFFER
            buf = Base.StringVector(size)
        end
        p::Int = pos
        $(code_data)
        return buf, p - Int(pos)
    end
    func = Expr(:function, params, body)

    makekw((arg, param)) = Expr(:kw, param, esc(arg isa Keyword ? arg.name : arg))
    if any(x -> x isa Expr, keys(argparams))
        @assert nposargs == 0
        return Expr(:call, :stringify, Expr(:call, func, :DUMMY_BUFFER, 1, makekw.(collect(argparams))...))
    else
        return Expr(:call, :Format, fstr, func)
    end
end

function stringify((buf, n))
    n < length(buf) && resize!(buf, n)
    return String(buf)
end

function genstrcopy(s::String)
    code = Expr(:block)
    for i in 1:ncodeunits(s)
        push!(code.args, :(buf[p+$i-1] = $(codeunit(s, i))))
    end
    return code
end

conv2func(conv::Conversion) =
    conv == CONV_REPR   ? repr   :
    conv == CONV_STRING ? string : identity


text:
"""
    DissipativeLambdaChain(N_atoms::Int; γd = 1.0, Γ = 0.0)

Creates a chain of `N_atoms` three-level atoms with spontaneous decay rate `Γ` and decay rate
`γd` into the second ground state.
The coupling rate is set to 1, but can be modified by scaling the Hamiltonian and coupling operator.
Uses `NLevelBasis(3)` as the single atom basis.
"""
function DissipativeLambdaChain(N_atoms::Int; γd = 0.0, Γ = 0.0)
	singleAtomBasis = NLevelBasis(3)
	σ_WG_i, σ_GW_i  = transition(singleAtomBasis, 2, 1), transition(singleAtomBasis, 1, 2)
	σ_DW_i = transition(singleAtomBasis, 3, 2)

	if N_atoms == 1
		groundstateAtoms = nlevelstate(singleAtomBasis, 1)
		H_Atom = 0.0 * one(singleAtomBasis)
		Dissipators = iszero(Γ) ? [√(γd) * σ_DW_i] : [√(γd) * σ_DW_i, √(Γ) * σ_GW_i]

		return H_Atom, Dissipators, σ_GW_i, groundstateAtoms
	end

	multiAtomBasis = reduce(⊗, repeat([singleAtomBasis], N_atoms))

	groundstateAtoms = reduce(⊗, repeat([nlevelstate(singleAtomBasis, 1)], N_atoms))

	σ_col = sum(i -> embed(multiAtomBasis, i, σ_GW_i), [1:N_atoms;])

	Dissipators = [embed(multiAtomBasis, i, √(γd) * σ_DW_i) for i ∈ 1:N_atoms]
	if !iszero(Γ)
		push!(Dissipators, [embed(multiAtomBasis, i, √(Γ) * σ_GW_i) for i ∈ 1:N_atoms]...)
	end

	H_Atoms = -0.5im * sum(i -> sum(
			j -> embed(multiAtomBasis, i, σ_WG_i) * embed(multiAtomBasis, j, σ_GW_i),
			[1:i-1;]
		), [2:N_atoms;]
	)
	H_Atoms = (H_Atoms + dagger(H_Atoms))

	return H_Atoms, Dissipators, σ_col, groundstateAtoms
end