module TurbulenceConvectionUtils

import Logging
using JLD2
using JSON
using Random
using DocStringExtensions

using ..ReferenceModels
import ..ReferenceModels: NameList

using ..ReferenceStats
# EKP modules
using EnsembleKalmanProcesses.ParameterDistributions
import EnsembleKalmanProcesses: construct_initial_ensemble
using TurbulenceConvection

using ..HelperFuncs

export ModelEvaluator
export run_SCM, run_SCM_handler, run_reference_SCM
export generate_scm_input, get_gcm_les_uuid, eval_single_ref_model
export save_full_ensemble_data
export precondition

"""
    ModelEvaluator

A structure containing the information required to perform
a forward model evaluation at a given parameter vector.

# Fields

$(TYPEDFIELDS)

# Constructors

$(METHODLIST)
"""
Base.@kwdef struct ModelEvaluator{FT <: Real}
    "Parameter vector in constrained (i.e. physical) space"
    param_cons::Vector{FT}
    "Parameter names associated with parameter vector"
    param_names::Vector{String}
    "Vector of reference models"
    ref_models::Vector{ReferenceModel}
    "Reference statistics for the inverse problem"
    ref_stats::ReferenceStatistics

    function ModelEvaluator(
        param_cons::Vector{FT},
        param_names::Vector{String},
        RM::Vector{ReferenceModel},
        RS::ReferenceStatistics,
    ) where {FT <: Real}
        return new{FT}(param_cons, param_names, RM, RS)
    end
end

"""
    run_SCM(
        u::Vector{FT},
        u_names::Vector{String},
        RM::Vector{ReferenceModel},
        RS::ReferenceStatistics;
        error_check::Bool = false,
        namelist_args = nothing,
        failure_handler = "high_loss",
    ) where {FT <: Real}
    run_SCM(
        ME::ModelEvaluator;
        error_check::Bool = false,
        namelist_args = nothing,
        failure_handler = "high_loss",
    ) where {FT <: Real}

Run the single-column model (SCM) using a set of parameters
and return the value of outputs defined in y_names, possibly
after normalization and projection onto lower dimensional
space using PCA.

The function also outputs a boolean diagnosing whether the
requested SCM simulation failed.

Inputs:

 - `u`               :: Values of parameters to be used in simulations.
 - `u_names`         :: SCM names for parameters `u`.
 - `RM`              :: Vector of `ReferenceModel`s
 - `RS`              :: reference statistics for simulation
 - `error_check`     :: Returns as an additional argument whether the SCM call errored.
 - `namelist_args`   :: Additional arguments passed to the TurbulenceConvection namelist.
 - `failure_handler` :: Method used to handle failed simulations.

Outputs:

 - `sim_dirs`    :: Vector of simulation output directories
 - `g_scm`       :: Vector of model evaluations concatenated for all flow configurations.
 - `g_scm_pca`   :: Projection of `g_scm` onto principal subspace spanned by eigenvectors.
 - `model_error` :: Whether the simulation errored with the requested configuration.
"""
function run_SCM(
    u::Vector{FT},
    u_names::Vector{String},
    RM::Vector{ReferenceModel},
    RS::ReferenceStatistics;
    error_check::Bool = false,
    namelist_args = nothing,
    failure_handler = "high_loss",
) where {FT <: Real}

    mkpath(joinpath(pwd(), "tmp"))
    result_arr = map(x -> eval_single_ref_model(x..., RS, u, u_names, namelist_args), enumerate(RM))
    # Unpack
    sim_dirs, g_scm, g_scm_pca, sim_errors = (getindex.(result_arr, i) for i in 1:4)
    g_scm = vcat(g_scm...)
    g_scm_pca = vcat(g_scm_pca...)

    model_error = any(sim_errors)
    if model_error && failure_handler == "high_loss"
        g_scm .= 1e5
        g_scm_pca .= 1e5
    end

    @info "Length of g_scm (full, pca): ($(length(g_scm)), $(length(g_scm_pca)))"
    if error_check
        return sim_dirs, g_scm, g_scm_pca, model_error
    else
        return sim_dirs, g_scm, g_scm_pca
    end
end
function run_SCM(
    ME::ModelEvaluator;
    error_check::Bool = false,
    namelist_args = nothing,
    failure_handler = "high_loss",
) where {FT <: Real}
    return run_SCM(
        ME.param_cons,
        ME.param_names,
        ME.ref_models,
        ME.ref_stats,
        error_check = error_check,
        namelist_args = namelist_args,
        failure_handler = failure_handler,
    )
end

"""
    eval_single_ref_model(
        m_index::IT,
        m::ReferenceModel,
        RS::ReferenceStatistics,
        u::Vector{FT},
        u_names::Vector{String},
        namelist_args = nothing,
    ) where {FT <: Real, IT <: Int}

Runs the single-column model (SCM) under a single configuration
(i.e., ReferenceModel) using a set of parameters u, and returns
the forward model evaluation in both the original and the latent
PCA space.

Inputs:

 - `m_index`       :: The index of the ReferenceModel within the overarching
                    ref_models vector used to construct the ReferenceStatistics.
 - `m`             :: A ReferenceModel.
 - `RS`            :: reference statistics for simulation
 - `u`             :: Values of parameters to be used in simulations.
 - `u_names`       :: SCM names for parameters `u`.
 - `namelist_args` :: Additional arguments passed to the TurbulenceConvection namelist.

Outputs:

 - `sim_dir`     ::  Simulation output directory.
 - `g_scm`       :: Forward model evaluation in original output space.
 - `g_scm_pca`   :: Projection of `g_scm` onto principal subspace spanned by eigenvectors.
 - `model_error` :: Whether the simulation errored with the requested configuration.
"""
function eval_single_ref_model(
    m_index::IT,
    m::ReferenceModel,
    RS::ReferenceStatistics,
    u::Vector{FT},
    u_names::Vector{String},
    namelist_args = nothing,
) where {FT <: Real, IT <: Int}
    # create temporary directory to store SCM data in
    tmpdir = mktempdir(joinpath(pwd(), "tmp"))
    # run TurbulenceConvection.jl. Get output directory for simulation data
    sim_dir, model_error = run_SCM_handler(m, tmpdir, u, u_names, namelist_args)
    filename = get_stats_path(sim_dir)
    if model_error
        g_scm = fill(NaN, length(get_z_obs(m)) * length(m.y_names))
    else
        g_scm = get_profile(m, filename, z_scm = get_z_obs(m))
        g_scm = normalize_profile(g_scm, length(m.y_names), RS.norm_vec[m_index])
    end
    # perform PCA reduction
    g_scm_pca = RS.pca_vec[m_index]' * g_scm
    return sim_dir, g_scm, g_scm_pca, model_error
end

"""
    run_reference_SCM(m::ReferenceModel; overwrite::Bool = false, run_single_timestep = true)

Run the single-column model (SCM) for a reference model object
using default parameters, and write the output to file.

Inputs:
 - `m`                    :: A `ReferenceModel`.
 - `overwrite`            :: if true, run TC.jl and overwrite existing simulation files.
 - `run_single_timestep`  :: if true, run only one time step.
"""
function run_reference_SCM(m::ReferenceModel; overwrite::Bool = false, run_single_timestep = true)
    output_dir = scm_dir(m)
    if ~isdir(joinpath(output_dir, "stats")) | overwrite
        namelist = get_scm_namelist(m)

        default_t_max = namelist["time_stepping"]["t_max"]
        default_adapt_dt = namelist["time_stepping"]["adapt_dt"]
        if run_single_timestep
            # Run only 1 timestep -- since we don't need output data, only simulation config
            namelist["time_stepping"]["adapt_dt"] = false
            namelist["time_stepping"]["t_max"] = namelist["time_stepping"]["dt_min"]
        end
        # run TurbulenceConvection.jl
        logger = Logging.ConsoleLogger(stderr, Logging.Warn)
        _, ret_code = Logging.with_logger(logger) do
            main1d(namelist; time_run = false)
        end
        if ret_code ≠ :success
            @warn "Default TurbulenceConvection.jl simulation $(basename(m.y_dir)) failed."
        end
        if run_single_timestep
            # reset t_max to default and overwrite stored namelist file
            namelist["time_stepping"]["t_max"] = default_t_max
            namelist["time_stepping"]["adapt_dt"] = default_adapt_dt
            open(namelist_directory(output_dir, m), "w") do io
                JSON.print(io, namelist, 4)
            end
        end
    end
end

"""
    run_SCM_handler(
        m::ReferenceModel,
        tmpdir::String,
        u::Array{FT, 1},
        u_names::Array{String, 1},
        namelist_args = nothing,
    ) where {FT<:AbstractFloat}

Run a case using a set of parameters `u_names` with values `u`,
and return directory pointing to where data is stored for simulation run.

Inputs:

 - m             :: Reference model
 - tmpdir        :: Temporary directory to store simulation results in
 - u             :: Values of parameters to be used in simulations.
 - u_names       :: SCM names for parameters `u`.
 - namelist_args :: Additional arguments passed to the TurbulenceConvection namelist.

Outputs:

 - output_dir   :: directory containing output data from the SCM run.
 - model_error   :: Boolean specifying whether the simulation failed.
"""
function run_SCM_handler(
    m::ReferenceModel,
    tmpdir::String,
    u::Vector{FT},
    u_names::Vector{String},
    namelist_args = nothing,
) where {FT <: AbstractFloat}

    # fetch namelist
    namelist = get_scm_namelist(m)
    # output subset of variables needed for calibration
    namelist["stats_io"]["calibrate_io"] = true

    # run TurbulenceConvection.jl with modified parameters
    return run_SCM_handler(
        m.case_name,
        tmpdir;
        u = u,
        u_names = u_names,
        namelist = namelist,
        namelist_args = namelist_args,
        uuid = basename(tmpdir), # set random uuid
        les = get(namelist["meta"], "lesfile", nothing),
    )
end

"""
    run_SCM_handler(
        case_name::String,
        out_dir::String;
        u::Vector{FT},
        u_names::Vector{String},
        namelist::Union{Dict, Nothing} = nothing,
        namelist_args::Union{Tuple, Nothing} = nothing,
        uuid::String = "01",
        les::Union{NamedTuple, String} = nothing,
    )

Run a TurbulenceConvection.jl case and return directory pointing to where data is stored for simulation run.

Inputs:
 - case_name     :: case name
 - out_dir       :: Directory to store simulation results in.
 Optional Inputs:
 - u             :: Values of parameters to be used in simulations.
 - u_names       :: SCM names for parameters `u`.
 - namelist      :: namelist to use for simulation.
 - namelist_args :: Additional arguments passed to the TurbulenceConvection namelist.
 - uuid          :: uuid of SCM run
 - les           :: path to LES stats file, or NamedTuple with keywords {forcing_model, month, experiment, cfsite_number} needed to specify path.
 Outputs:
 - output_dir   :: directory containing output data from the SCM run.
 - model_error   :: Boolean specifying whether the simulation failed.
"""
function run_SCM_handler(
    case_name::String,
    out_dir::String;
    u::Vector{FT},
    u_names::Vector{String},
    namelist::Union{Dict, Nothing} = nothing,
    namelist_args::Union{Vector, Nothing} = nothing,
    uuid::String = "01",
    les::Union{NamedTuple, String, Nothing} = nothing,
) where {FT <: AbstractFloat}
    model_error = false

    # fetch default namelist if not provided
    if isnothing(namelist)
        namelist = NameList.default_namelist(case_name)
    end

    namelist["meta"]["uuid"] = uuid
    # set output dir to `out_dir`
    namelist["output"]["output_root"] = out_dir

    u_names, u = create_parameter_vectors(u_names, u)
    # Set optional namelist args
    if !isnothing(namelist_args)
        for namelist_arg in namelist_args
            change_entry!(namelist, namelist_arg)
        end
    end

    # update learnable parameter values
    @assert length(u_names) == length(u)
    for (pName, pVal) in zip(u_names, u)
        if haskey(namelist["turbulence"]["EDMF_PrognosticTKE"], pName)
            namelist["turbulence"]["EDMF_PrognosticTKE"][pName] = pVal
        elseif haskey(namelist["microphysics"], pName)
            namelist["microphysics"][pName] = pVal
        elseif haskey(namelist["time_stepping"], pName)
            namelist["time_stepping"][pName] = pVal
        else
            throw(
                ArgumentError("Parameter $pName cannot be calibrated. Consider adding namelist dictionary if needed."),
            )
        end
    end

    if case_name == "LES_driven_SCM"
        if isnothing(les)
            error("les path or keywords required for LES_driven_SCM case!")
        elseif isa(les, NamedTuple)
            les = get_stats_path(
                get_cfsite_les_dir(
                    les.cfsite_number;
                    forcing_model = les.forcing_model,
                    month = les.month,
                    experiment = les.experiment,
                ),
            )
        end
        namelist["meta"]["lesfile"] = les
    end

    # run TurbulenceConvection.jl with modified parameters
    logger = Logging.ConsoleLogger(stderr, Logging.Warn)
    _, ret_code = Logging.with_logger(logger) do
        main1d(namelist; time_run = false)
    end
    if ret_code ≠ :success
        model_error = true
        message = ["TurbulenceConvection.jl simulation $out_dir failed with parameters: \n"]
        append!(message, ["$param_name = $param_value \n" for (param_name, param_value) in zip(u_names, u)])
        @warn join(message)
    end

    return data_directory(out_dir, case_name, uuid), model_error
end

"""
    create_parameter_vectors(u_names::Vector{String}, u::Vector{FT}) where {FT <: AbstractFloat}

Given vector of parameter names and corresponding values, combine any vector components
into single parameter vectors for input into SCM.

Inputs:

 - `u_names` :: SCM names for parameters `u`, which may contain vector components.
 - `u` :: Values of parameters to be used in simulations, which may contain vector components.

Outputs:

 -  `u_names_out` :: SCM names for parameters `u`.
 -  `u_out` :: Values of parameters to be used in simulations.
"""
function create_parameter_vectors(u_names::Vector{String}, u::Vector{FT}) where {FT <: AbstractFloat}

    u_names_out = String[]
    u_out = []

    vector_param_inds = occursin.(r"{?}", u_names)
    pv_name_elem = rsplit.(u_names[vector_param_inds], "_", limit = 2) # get param name and element index
    u_vec_names, uvi = ~isempty(pv_name_elem) ? eachrow(reduce(hcat, pv_name_elem)) : ([], [])  # "transpose" `pv_name_elem`
    u_vec_inds = @. parse(Int64, only(split(uvi, (('{', '}'),), keepempty = false)))  # get `i` from "{i}" as Int64

    # collect scalar parameters
    scalar_param_inds = .!vector_param_inds
    append!(u_names_out, u_names[scalar_param_inds])
    append!(u_out, u[scalar_param_inds])

    # collect vector parameters
    for u_name in unique(u_vec_names)
        u_name_inds = u_vec_names .== u_name
        u_vals = u[vector_param_inds][u_name_inds]
        u_vals_sort_inds = u_vec_inds[u_name_inds]
        permute!(u_vals, u_vals_sort_inds)
        push!(u_names_out, u_name)
        push!(u_out, u_vals)
    end

    return u_names_out, u_out
end

"""
    generate_scm_input(
        model_evaluators::Vector{ModelEvaluator{FT}},
        outdir_path::String = pwd(),
    ) where {FT <: AbstractFloat}

Writes to file a set of ModelEvaluator used to initialize SCM
evaluations at different parameter vectors, as well as their
assigned numerical version.
"""
function generate_scm_input(
    model_evaluators::Vector{ModelEvaluator{FT}},
    outdir_path::String = pwd(),
    batch_indices::Union{Vector{Int}, Nothing} = nothing,
) where {FT <: AbstractFloat}
    # Generate versions conditioned on being unique within the batch.
    used_versions = Vector{Int}()
    for model_evaluator in model_evaluators
        version = rand(11111:99999)
        while version in used_versions
            version = rand(11111:99999)
        end
        jldsave(scm_init_path(outdir_path, version); model_evaluator, version, batch_indices)
        push!(used_versions, version)
    end
    return used_versions
end

"""
    get_gcm_les_uuid(
        cfsite_number::Integer;
        forcing_model::String,
        month::Integer,
        experiment::String,)

Generate unique and self-describing uuid given information about a GCM-driven LES
simulation from [Shen2022](@cite).
"""
function get_gcm_les_uuid(
    cfsite_number::Integer;
    forcing_model::String = "HadGEM2-A",
    month::Integer = 7,
    experiment::String = "amip",
)
    cfsite_number = string(cfsite_number)
    month = string(month, pad = 2)
    return join([cfsite_number, forcing_model, month, experiment], '_')
end

""" Save full EDMF data from every ensemble"""
function save_full_ensemble_data(save_path, sim_dirs_arr, ref_models)
    # get a simulation directory `.../Output.SimName.UUID`, and corresponding parameter name
    for (ens_i, sim_dirs) in enumerate(sim_dirs_arr)  # each ensemble returns a list of simulation directories
        ens_i_path = joinpath(save_path, "ens_$ens_i")
        mkpath(ens_i_path)
        for (ref_model, sim_dir) in zip(ref_models, sim_dirs)
            scm_name = ref_model.case_name
            # Copy simulation data to output directory
            dirname = splitpath(sim_dir)[end]
            @assert dirname[1:7] == "Output."  # sanity check
            # Stats file
            tmp_data_path = joinpath(sim_dir, "stats/Stats.$scm_name.nc")
            save_data_path = joinpath(ens_i_path, "Stats.$scm_name.$ens_i.nc")
            cp(tmp_data_path, save_data_path)
            # namefile
            tmp_namefile_path = namelist_directory(sim_dir, scm_name)
            save_namefile_path = namelist_directory(ens_i_path, scm_name)
            cp(tmp_namefile_path, save_namefile_path)
        end
    end
end

"""
    precondition(
        param::Vector{FT},
        priors,
        ref_models::Vector{ReferenceModel},
        ref_stats::ReferenceStatistics,
        namelist_args = nothing;
        counter::Integer = 0,
        max_counter::Integer = 10,
    ) where {FT <: Real}

Substitute parameter vector `param` by a parameter vector drawn
from the same prior, conditioned on the forward model being stable.

Inputs:

 - `param`      :: A parameter vector that may possibly result in unstable
    forward model evaluations (in unconstrained space).
 - `priors`      :: Priors from which the parameters were drawn.
 - `ref_models`  :: Vector of ReferenceModels to check stability for.
 - `ref_stats`   :: ReferenceStatistics of the ReferenceModels.
 - `namelist_args` :: Arguments passed to the TC.jl namelist.
 - `counter` :: Accumulator tracking number of recursive calls to preconditioner.
 - `max_counter` :: Maximum number of recursive calls to the preconditioner.

Outputs:

 - `new_param` :: A new parameter vector drawn from the prior, conditioned on simulations being stable (in unconstrained space).
"""
function precondition(
    param::Vector{FT},
    priors::ParameterDistribution,
    ref_models::Vector{ReferenceModel},
    ref_stats::ReferenceStatistics,
    namelist_args = nothing;
    counter::Integer = 0,
    max_counter::Integer = 10,
) where {FT <: Real}
    param_names = priors.name
    # Wrapper around SCM
    g_(u::Array{Float64, 1}) =
        run_SCM(u, param_names, ref_models, ref_stats, error_check = true, namelist_args = namelist_args)

    param_cons = deepcopy(transform_unconstrained_to_constrained(priors, param))
    _, _, _, model_error = g_(param_cons)
    if model_error && counter < max_counter
        message = ["Unstable parameter vector found: \n"]
        append!(message, ["$param_name = $param \n" for (param_name, param) in zip(param_names, param_cons)])
        @warn join(message)
        @warn "Sampling new parameter vector from prior..."
        return precondition(
            vec(construct_initial_ensemble(priors, 1)),
            priors,
            ref_models,
            ref_stats,
            namelist_args,
            counter = counter + 1,
            max_counter = max_counter,
        )
    elseif model_error
        @error "Number of recursive calls to preconditioner exceeded $(max_counter). Returning last failed parameter."
        return param
    else
        @info "Preconditioning finished."
        return param
    end
end

"""
    precondition(ME::ModelEvaluator, priors)

Substitute the parameter vector of a ModelEvaluator by another
one drawn from the given `priors`, conditioned on the forward
model being stable.

Inputs:
 - ME          :: A ModelEvaluator.
 - priors      :: Priors from which the parameters were drawn.
Outputs:
 - A preconditioned ModelEvaluator.
"""
function precondition(ME::ModelEvaluator, priors; namelist_args = nothing)
    # Precondition in unconstrained space
    u_orig = transform_constrained_to_unconstrained(priors, ME.param_cons)
    u = precondition(u_orig, priors, ME.ref_models, ME.ref_stats, namelist_args)
    # Transform back to constrained space
    param_cons = transform_unconstrained_to_constrained(priors, u)
    return ModelEvaluator(param_cons, ME.param_names, ME.ref_models, ME.ref_stats)
end

end # module


__precompile__()
module Arrowhead

using Base, HierarchicalMatrices

import Base: size, getindex, inv, issymmetric, chol, eig, eigvals, +, -

export SymArrow, UpperArrow

include("SymArrow.jl")
include("UpperArrow.jl")
include("secular.jl")

end # module


# Load required packages
using FishEBM, DataFrames

#fill params
# Specify stock assumptions:
# * s_a.naturalmortality = Age specific mortality
# * s_a.halfmature = Age at 50% maturity
# * s_a.broodsize = Age specific fecundity
# * s_a.fecunditycompensation = Compensatory strength - fecundity
# * s_a.maturitycompensation = Compensatory strength - age at 50% maturity
# * s_a.mortalitycompensation = Compensatory strength - adult natural mortality
# * s_a.catchability = Age specific catchability
adult_a = AdultAssumptions([0.35, 0.40, 0.45, 0.50, 0.55, 0.60, 0.65],
                       5,
                       [2500, 7500, 15000, 20000, 22500, 27500, 32500],
                       2,
                       0.25,
                       1,
                       [0.00001, 0.00002, 0.000025, 0.000025, 0.000025, 0.000025, 0.000025])

# Specify environment assumptions:
# * e_a.spawning = Spawning areas
# * e_a.habitat = Habitat types
# * e_a.risk = Risk areas
# * e_a.harvest = Harvest areas
spawnPath = string(split(Base.source_path(), "FishEBM.jl")[1], "FishEBM.jl/maps/LakeHuron_1km_spawning.csv")
habitatPath = string(split(Base.source_path(), "FishEBM.jl")[1], "FishEBM.jl/maps/LakeHuron_1km_habitat.csv")
riskPath = string(split(Base.source_path(), "FishEBM.jl")[1], "FishEBM.jl/maps/LakeHuron_1km_risk.csv")
harvestPath = string(split(Base.source_path(), "FishEBM.jl")[1], "FishEBM.jl/maps/LakeHuron_1km_harvest.csv")

enviro_a = initEnvironment(spawnPath, habitatPath, riskPath, harvestPath)

#a_db = AgentDB(e_a)
#e_a = generateEnvironment(spawnPath, habitatPath, riskPath)

# Specify agent assumptions:
# * a_a.naturalmortality =  Weekly natural mortality rate (by habitat type in the rows, and stage in the columns)
# * a_a.extramortality = Weekly risk mortality (by stage)
# * a_a.growth = Stage length (in weeks)
# * a_a.movement = Movement weight matrices
# * a_a.autonomy =  Movement autonomy

a_a = AgentAssumptions([[0.80 0.095 0.09 0.05]
                        [0.10 0.095 0.09 0.10]
                        [0.80 0.095 0.09 0.20]
                        [0.80 0.80 0.09 0.30]
                        [0.80 0.80 0.80 0.40]
                        [0.80 0.80 0.80 0.50]],
                       [0.0, 0.0, 0.0, 0.0],
                       [19, 52, 104, 0],
                       Array[[[0. 0. 0.]
                              [0. 1. 0.]
                              [0. 0. 0.]], [[1. 2. 1.]
                                            [1. 2. 1.]
                                            [1. 1. 1.]], [[1. 2. 1.]
                                                          [1. 1. 1.]
                                                          [1. 1. 1.]], [[1. 2. 1.]
                                                                        [1. 1. 1.]
                                                                        [1. 1. 1.]]],
                       [0., 0.5, 0.75, 0.5])

#generate agent db and hash environemt for Testing
a_db = AgentDB(enviro_a)
hashEnvironment!(a_db, enviro_a)

#inject agents into agent db for testing nonempty agent db
addingStock = [5000, 10000, 15000, 20000]
for i = 1:4
  injectAgents!(a_db, enviro_a.spawningHash, addingStock[5-i], -a_a.growth[((7-i)%4)+1])
end

#non-empty a_db
a_db


function process_icr!(;collection, openable_cards, account, rarity, parameters, logging_vector)
    if isempty(openable_cards.by_rarity[rarity])
        @goto vault_increment
    end
    card = rand(openable_cards.by_rarity[rarity])
    if collection[card] < 4
        collection[card] += 1
        push!(logging_vector, card)
        if collection[card] >= 4 && parameters.prevent_duplicates
            remove_fully_collected!(card, rarity; openable_cards = openable_cards)
        end
    else
        @label vault_increment
        account.vault_pct += parameters.duplicate_vpct[rarity]
    end
    collection
end


using Test
using LogicCircuits
using ProbabilisticCircuits

@testset "Circuit saver test" begin
    mktempdir() do tmp

        circuit, vtree = load_struct_prob_circuit(
                            zoo_psdd_file("little_4var.psdd"), zoo_vtree_file("little_4var.vtree"))

        # load, save, and load as .psdd
        save_circuit("$tmp/temp.psdd", circuit, vtree)
        save(vtree, "$tmp/temp.vtree");
        load_struct_prob_circuit("$tmp/temp.psdd", "$tmp/temp.vtree")

        # save and load as .sdd
        save_circuit("$tmp/temp.sdd", circuit, vtree)
        save(vtree, "$tmp/temp.vtree")

    end

    #TODO add some actual @test statements
    @test true
end


using Revise, Parameters, Logging, Test
using CUDA
includet("C:\\GitHub\\GitHub\\NuclearMedEval\\test\\includeAllUseFullForTest.jl")
using ..CUDAGpuUtils ,..IterationUtils,..ReductionUtils , ..MemoryUtils,..CUDAAtomicUtils, ..PrepareArrtoBool
using ..ResultListUtils, ..MetadataAnalyzePass,..MetaDataUtils,..WorkQueueUtils,..ProcessMainDataVerB,..HFUtils, ..ScanForDuplicates

@testset "getBoolKernel" begin

@testset "getIndexOfQueue" begin
        """
    1)   Left FP
    2)   Left FN
    3)   Right FP
    4)   Right FN
    5)   Posterior FP
    6)   Posterior FN
    7)   Anterior FP
    8)   Anterior FN
    9)   Top FP
    10)   Top FN
    11)   Bottom FP
    12)   Bottom FN
    13)   Total block Fp
    14)   Total block Fn

    """
    dataBdim= (5,5,5)
    @test getIndexOfQueue(1,1,1, (5,5,5),false)==1
    @test getIndexOfQueue(1,1,1, (5,5,5),true)==2

    @test getIndexOfQueue(5,1,1, (5,5,5),false)==3
    @test getIndexOfQueue(5,1,1, (5,5,5),true)==4

    @test getIndexOfQueue(2,1,1, (5,5,5),false)==5
    @test getIndexOfQueue(2,1,1, (5,5,5),true)==6

    @test getIndexOfQueue(2,5,1, (5,5,5),false)==7
    @test getIndexOfQueue(2,5,1, (5,5,5),true)==8

    @test getIndexOfQueue(2,2,1, (5,5,5),false)==9
    @test getIndexOfQueue(2,2,1, (5,5,5),true)==10

    @test getIndexOfQueue(2,2,5, (5,5,5),false)==11
    @test getIndexOfQueue(2,2,5, (5,5,5),true)==12

    @test getIndexOfQueue(5,5,5, (5,5,5),false)==3
    @test getIndexOfQueue(5,5,5, (5,5,5),true)==4

    @test getIndexOfQueue(2,2,2, (5,5,5),false)==13
    @test getIndexOfQueue(2,2,2, (5,5,5),true)==14

end

# metadat= zeros(Int64,100,50,40,15);
# metaDataDims= size(metadat)
# for i in 1:length(metadat)
#     metadat[i]=i
# end
# linIdexMetaPrim= (100*50*3)+303
# linIdexMeta= (-1)+ (100*50*3)+303
# xMeta= mod(linIdexMeta,metaDataDims[1])
# zMeta= fld((linIdexMeta),metaDataDims[1]*metaDataDims[2])
# yMeta= fld((linIdexMeta-((zMeta*metaDataDims[1]*metaDataDims[2] ) + xMeta )),metaDataDims[1])
# qn= 7
# metadat[linIdexMetaPrim+(qn-1)*metaDataDims[1]*metaDataDims[2]*metaDataDims[3] ]
# metadat[xMeta+1,yMeta+1,zMeta+1,qn]

##### iter3dOuter

@testset "iter3dOuter" begin

        using Revise, Parameters, Logging, Test
        using CUDA
        includet("C:\\GitHub\\GitHub\\NuclearMedEval\\test\\includeAllUseFullForTest.jl")
        using ..CUDAGpuUtils ,..IterationUtils,..ReductionUtils , ..MemoryUtils,..CUDAAtomicUtils
        using ..ResultListUtils, ..MetadataAnalyzePass,..MetaDataUtils,..WorkQueueUtils,..ProcessMainDataVerB,..HFUtils, ..ScanForDuplicates

    singleVal = CUDA.zeros(14)

    # threads=(32,20)

    threads=(32,8)
    # blocks =2

    # mainArrDims= (7,7,7)
    # dataBdim = (2,2,2)

    blocks =136
    dataBdim = (33,32,32)
    mainArrDims= (512, 512, 500)

    metaDataDims= (cld(mainArrDims[1],dataBdim[1] ),cld(mainArrDims[2],dataBdim[2]),cld(mainArrDims[3],dataBdim[3]))
    #we are iterating here block by block sequentially
    metaDataDims= (16, 16, 16,90)
    metaDataLength= metaDataDims[1]*metaDataDims[2]*metaDataDims[3]
    loopMeta= fld(metaDataLength,blocks )

    # metaDataDims (16, 16, 16, 90)  loopMeta 409 metaDataLength 4096
    function iter3dOuterKernel(mainArrDims,singleVal,metaDataDims,dataBdim, metaDataLength, loopMeta )
        @iter3dOuter(metaDataDims,loopMeta,metaDataLength,
        begin
        @ifXY 1 1 CUDA.@atomic singleVal[]+=1
     @ifXY 1 1    CUDA.@cuprint " linIdexMeta $(linIdexMeta)  zMeta $(zMeta)  yMeta$(yMeta) xMeta$(xMeta)  idX $(blockIdxX()) \n"

    end)

        return
    end
    @cuda threads=threads blocks=blocks iter3dOuterKernel(mainArrDims,singleVal,metaDataDims,dataBdim, metaDataLength, loopMeta )
    @test singleVal[1]==metaDataDims[1]*metaDataDims[2]*metaDataDims[3]

end #"iter3dOuter"

################ iter data block
@testset "iter data block" begin

    using Revise, Parameters, Logging, Test
    using CUDA
    includet("C:\\GitHub\\GitHub\\NuclearMedEval\\test\\includeAllUseFullForTest.jl")
    using ..CUDAGpuUtils ,..IterationUtils,..ReductionUtils , ..MemoryUtils,..CUDAAtomicUtils
    using ..ResultListUtils, ..MetadataAnalyzePass,..MetaDataUtils,..WorkQueueUtils,..ProcessMainDataVerB,..HFUtils, ..ScanForDuplicates

    ##### iter data block
    singleVal = CUDA.zeros(Int64,14)
    threads=(32,5)

    blocks =17
    mainArrDims= (67,177,90)
    dataBdim = (43,21,17)
    indices = CUDA.zeros(Bool,mainArrDims)

    metaDataDims= (cld(mainArrDims[1],dataBdim[1] ),cld(mainArrDims[2],dataBdim[2]),cld(mainArrDims[3],dataBdim[3]))

    inBlockLoopX,inBlockLoopY,inBlockLoopZ= (fld(dataBdim[1] ,threads[1]),fld(dataBdim[2] ,threads[2]),dataBdim[3]    )
    #we are iterating here block by block sequentially
    metaDataLength= metaDataDims[1]*metaDataDims[2]*metaDataDims[3]
    loopMeta= fld(metaDataLength,blocks )

    function iterDataBlocksKernel(loopMeta,metaDataLength,mainArrDims,singleVal,metaDataDims,dataBdim, inBlockLoopX,inBlockLoopY,inBlockLoopZ ,indices)
        @iter3dOuter(metaDataDims,loopMeta,metaDataLength,
        begin
            @iterDataBlock(mainArrDims,dataBdim, inBlockLoopX,inBlockLoopY,inBlockLoopZ,xMeta,yMeta,zMeta,
            begin
           CUDA.@atomic singleVal[1]+=1
            indices[x,y,z]=true
            #CUDA.@cuprint "x $(x) y $(y) z $(z) zMeta $(zMeta)  yMeta $(yMeta) xMeta $(xMeta)  idX $(blockIdxX()) \n"

        end)end)

        return
    end
    @cuda threads=threads blocks=blocks iterDataBlocksKernel(loopMeta,metaDataLength,mainArrDims,singleVal,metaDataDims,dataBdim, inBlockLoopX,inBlockLoopY,inBlockLoopZ ,indices)
    @test singleVal[1]==mainArrDims[1]*mainArrDims[2]*mainArrDims[3]
    Int64(singleVal[1])
    7*7*7
    inddd = Array(indices);
    cartt = CartesianIndices(inddd);
    filtered = filter(cc-> !inddd[cc]  ,cartt)
    sum(indices)

end #iter data block

##### @uploadLocalfpFNCounters
@testset "uploadLocalfpFNCounters" begin

    localQuesValues = CUDA.zeros(UInt32,14)
    threads=(32,5)
    blocks =1

    mainArrDims= (99,99,99)
    dataBdim = (17,7,12)

    metaDataDims= (cld(mainArrDims[1],dataBdim[1] ),cld(mainArrDims[2],dataBdim[2]),cld(mainArrDims[3],dataBdim[3]))

    inBlockLoopX,inBlockLoopY,inBlockLoopZ= (fld(dataBdim[1] ,threads[1]),fld(dataBdim[2] ,threads[2]),dataBdim[3]    )
    #we are iterating here block by block sequentially
    metaDataLength= metaDataDims[1]*metaDataDims[2]*metaDataDims[3]
    loopMeta= fld(metaDataLength,blocks )

    function uploadLocalfpFNCountersKernel(mainArrDims,localQuesValues,metaDataDims,loopMeta,metaDataLength,dataBdim, inBlockLoopX,inBlockLoopY,inBlockLoopZ)
        # localQuesValues= @cuStaticSharedMem(Float32, 14)

        # @iter3dOuter(metaDataDims,loopMeta,metaDataLength,
        # begin
            #@ifXY 1 1    CUDA.@cuprint "  xMeta $(xMeta) yMeta $(yMeta)  zMeta $(zMeta) \n"
            xMeta= 2
            zMeta= 2
            yMeta= 2

            @iterDataBlock(mainArrDims,dataBdim, inBlockLoopX,inBlockLoopY,inBlockLoopZ,xMeta,yMeta,zMeta,
            begin
                        boolGold=true
                        boolSegm=false
                        #coord=PrepareArrtoBool.getIndexOfQueue((xdim * blockDimX())+threadIdxX() ,(ydim * blockDimY())+threadIdxY(),(zdim+1),dataBdim,false)
                        PrepareArrtoBool.@uploadLocalfpFNCounters()
                        #CUDA.@cuprint " spot $(coord)  x $((xdim * blockDimX())+threadIdxX()  )) y $((ydim * blockDimY())+threadIdxY()  )) z $((zdim+1)) dataBdim $(dataBdim[1]),$(dataBdim[2]),$(dataBdim[3]) \n"
                        #CUDA.atomic_add!(pointer(localQuesValues, coord),Float32(1))

                    #@inboundsCUDA.@atomic localQuesValues[coord]+=Float32(1)
                    #CUDAAtomicUtils.atomicallyAddToSpot(Float32,localQuesValues,getIndexOfQueue(x,y,z,dataBdim,boolSegm),1)
                    #PrepareArrtoBool.@uploadLocalfpFNCounters()

                    end)#end)

        return
    end

    @cuda threads=threads blocks=blocks uploadLocalfpFNCountersKernel(mainArrDims,localQuesValues,metaDataDims,loopMeta,metaDataLength,dataBdim, inBlockLoopX,inBlockLoopY,inBlockLoopZ)
    @test sum(localQuesValues)>0

    @test Int64(localQuesValues[1])==dataBdim[2]*dataBdim[3]
    @test Int64(localQuesValues[3])==dataBdim[2]*dataBdim[3]

    @test Int64(localQuesValues[5])==dataBdim[1]*dataBdim[3]-2*dataBdim[3]
    @test Int64(localQuesValues[7])==dataBdim[1]*dataBdim[3]-2*dataBdim[3]

    @test Int64(localQuesValues[9])==(dataBdim[1]-2)*(dataBdim[2]-2)
    @test Int64(localQuesValues[11])==(dataBdim[1]-2)*(dataBdim[2]-2)

    @test Int64(localQuesValues[13])==dataBdim[1]*dataBdim[2]*dataBdim[3] - sum(localQuesValues[1:12])

end #"uploadLocalfpFNCounters"

using Revise, Parameters, Logging, Test
using CUDA
includet("C:\\GitHub\\GitHub\\NuclearMedEval\\test\\includeAllUseFullForTest.jl")
using ..CUDAGpuUtils ,..IterationUtils,..ReductionUtils , ..MemoryUtils,..CUDAAtomicUtils
using ..ResultListUtils, ..MetadataAnalyzePass,..MetaDataUtils,..WorkQueueUtils,..ProcessMainDataVerB,..HFUtils, ..ScanForDuplicates

# ########### uploadDataToMetaData
#  @testset "uploadDataToMetaData 1 " begin

#     localQuesValues = CUDA.zeros(Float32,14)
#     threads=(32,5)
#     blocks =8
#     mainArrDims= (634,521,632)
#     dataBdim = (43,21,17)

#     metaDataDims= (cld(mainArrDims[1],dataBdim[1] ),cld(mainArrDims[2],dataBdim[2]),cld(mainArrDims[3],dataBdim[3]))
#     loopAXFixed,loopBXfixed,loopAYFixed,loopBYfixed,loopAZFixed,loopBZfixed,loopdataDimMainX,loopdataDimMainY,loopdataDimMainZ,inBlockLoopX,inBlockLoopY,inBlockLoopZ,metaDataLength,loopMeta,loopWarpMeta,clearIterResShmemLoop,clearIterSourceShmemLoop,clearIterResShmemLoop,clearIterSourceShmemLoop,resShmemTotalLength,sourceShmemTotalLength = calculateLoopsIter(dataBdim,threads[1],threads[2],metaDataDims,blocks)

#     #we are iterating here block by block sequentially
#     metaData = MetaDataUtils.allocateMetadata(mainArrDims,dataBdim);
#     metaDataLength= metaDataDims[1]*metaDataDims[2]*metaDataDims[3]
#     loopMeta= fld(metaDataLength,blocks )
#     function uploadDataToMetaDataKernel(loopMeta,metaDataLength,mainArrDims,metaDataDims,dataBdim, inBlockLoopX,inBlockLoopY,inBlockLoopZ,metaData)
#         localQuesValues= @cuStaticSharedMem(UInt32, 14)
#         localBool=false
#         anyPositive=false

#         @iter3dOuter(metaDataDims,loopMeta,metaDataLength,
#         begin

#             @iterDataBlock(mainArrDims,dataBdim, inBlockLoopX,inBlockLoopY,inBlockLoopZ,
#             begin
#                         boolGold=true
#                         boolSegm=false
#                         localBool=true
#                         if(xMeta==2 && yMeta==2 && zMeta==4)

#                             PrepareArrtoBool.@uploadLocalfpFNCounters()
#                         end

#                     end)
#                     sync_threads()

#                     isAnyPositive= true
#                         PrepareArrtoBool.@uploadDataToMetaData()
#                         sync_threads()

#                         sync_threads()
#                         #threadfence()
#                         @ifY 1 if(threadIdxX()<15)
#                             localQuesValues[threadIdxX()]=0
#                         end
#                 end)

#         return
#     end

#     @cuda threads=threads blocks=blocks uploadDataToMetaDataKernel(loopMeta,metaDataLength,mainArrDims,metaDataDims,loopXMeta,loopYZMeta,yTimesZmeta,dataBdim, inBlockLoopX,inBlockLoopY,inBlockLoopZ,metaData)

#     @test metaData[3,3,5,getBeginingOfFpFNcounts()+1]==dataBdim[2]*dataBdim[3]
#     @test metaData[3,3,5,getBeginingOfFpFNcounts()+3]==dataBdim[2]*dataBdim[3]

#     @test metaData[3,3,5,getBeginingOfFpFNcounts()+5]==dataBdim[1]*dataBdim[3]-2*dataBdim[3]
#     @test metaData[3,3,5,getBeginingOfFpFNcounts()+7]==dataBdim[1]*dataBdim[3]-2*dataBdim[3]

#     @test metaData[3,3,5,getBeginingOfFpFNcounts()+9]==(dataBdim[1]-2)*(dataBdim[2]-2)
#     @test metaData[3,3,5,getBeginingOfFpFNcounts()+11]==(dataBdim[1]-2)*(dataBdim[2]-2)

#     @test metaData[3,3,5,getBeginingOfFpFNcounts()+13]==dataBdim[1]*dataBdim[2]*dataBdim[3] - sum(metaData[3,3,5,getBeginingOfFpFNcounts():getBeginingOfFpFNcounts()+12])

#     @test metaData[3,3,5,getBeginingOfFpFNcounts()+15]==dataBdim[1]*dataBdim[2]*dataBdim[3]

#     sum(Array(metaData[3,3,5,getBeginingOfFpFNcounts():getBeginingOfFpFNcounts()+13]))

#     metaData[3,3,5,getBeginingOfFpFNcounts()+1]+ metaData[3,3,5,getBeginingOfFpFNcounts()+3]+ metaData[3,3,5,getBeginingOfFpFNcounts()+5]+ metaData[3,3,5,getBeginingOfFpFNcounts()+7]+ metaData[3,3,5,getBeginingOfFpFNcounts()+9]+ metaData[3,3,5,getBeginingOfFpFNcounts()+11]+ metaData[3,3,5,getBeginingOfFpFNcounts()+13]

# end #test set
# ########### uploadDataToMetaData 2
# @testset "uploadDataToMetaData2" begin

#     # localQuesValues = CUDA.zeros(Float32,14)
#     # threads=(32,5)
#     # blocks =8

#     # mainArrDims= (634,521,632)
#     # dataBdim = (43,21,17)

#     # metaDataDims= (cld(mainArrDims[1],dataBdim[1] ),cld(mainArrDims[2],dataBdim[2]),cld(mainArrDims[3],dataBdim[3]))
#     # metaData = MetaDataUtils.allocateMetadata(mainArrDims,dataBdim);

#     # inBlockLoopX,inBlockLoopY,inBlockLoopZ= (fld(dataBdim[1] ,threads[1]),fld(dataBdim[2] ,threads[2]),dataBdim[3]    )
#     # #we are iterating here block by block sequentially
#     # loopXMeta,loopYZMeta= (metaDataDims[1],fld(metaDataDims[2]*metaDataDims[3] ,blocks)  )
#     # yTimesZmeta= metaDataDims[2]*metaDataDims[3]
#     # metaDataLength= metaDataDims[1]*metaDataDims[2]*metaDataDims[3]
#     # loopMeta= fld(metaDataLength,blocks )
#     # function uploadDataToMetaDataKernel(loopMeta,metaDataLength,mainArrDims,metaDataDims,loopXMeta,loopYZMeta,yTimesZmeta,dataBdim, inBlockLoopX,inBlockLoopY,inBlockLoopZ,metaData)

#     #     localQuesValues= @cuStaticSharedMem(UInt32, 14)
#     #     localBool=false
#     #     @iter3dOuter(metaDataDims,loopMeta,metaDataLength,
#     #     begin

#     #         @iterDataBlock(mainArrDims,dataBdim, inBlockLoopX,inBlockLoopY,inBlockLoopZ,
#     #         begin
#     #                     boolGold=false
#     #                     boolSegm=true
#     #                     localBool=true
#     #                     if(xMeta==2 && yMeta==2 && zMeta==4)

#     #                         PrepareArrtoBool.@uploadLocalfpFNCounters()
#     #                     end

#     #                 end)
#     #                 sync_threads()

#     #                 isAnyPositive= true
#     #                     PrepareArrtoBool.@uploadDataToMetaData()
#     #                     sync_threads()

#     #                     sync_threads()
#     #                     #threadfence()
#     #                     @ifY 1 if(threadIdxX()<15)
#     #                         localQuesValues[threadIdxX()]=0
#     #                     end
#     #             end)

#     #     return
#     # end

#     # @cuda threads=threads blocks=blocks uploadDataToMetaDataKernel(loopMeta,metaDataLength,mainArrDims,metaDataDims,loopXMeta,loopYZMeta,yTimesZmeta,dataBdim, inBlockLoopX,inBlockLoopY,inBlockLoopZ,metaData)

#     # @test metaData[3,3,5,getBeginingOfFpFNcounts()+2]==dataBdim[2]*dataBdim[3]
#     # @test metaData[3,3,5,getBeginingOfFpFNcounts()+4]==dataBdim[2]*dataBdim[3]

#     # @test metaData[3,3,5,getBeginingOfFpFNcounts()+6]==dataBdim[1]*dataBdim[3]-2*dataBdim[3]
#     # @test metaData[3,3,5,getBeginingOfFpFNcounts()+8]==dataBdim[1]*dataBdim[3]-2*dataBdim[3]

#     # @test metaData[3,3,5,getBeginingOfFpFNcounts()+10]==(dataBdim[1]-2)*(dataBdim[2]-2)
#     # @test metaData[3,3,5,getBeginingOfFpFNcounts()+12]==(dataBdim[1]-2)*(dataBdim[2]-2)

#     # @test metaData[3,3,5,getBeginingOfFpFNcounts()+14]==dataBdim[1]*dataBdim[2]*dataBdim[3] - sum(metaData[3,3,5,getBeginingOfFpFNcounts():getBeginingOfFpFNcounts()+12])

#     # @test metaData[3,3,5,getBeginingOfFpFNcounts()+16]==dataBdim[1]*dataBdim[2]*dataBdim[3]

#     # sum(Array(metaData[3,3,5,getBeginingOfFpFNcounts():getBeginingOfFpFNcounts()+13]))

# end#testset

#################   uploadMinMaxesToShmem

@testset "uploadMinMaxesToShmem" begin
    localQuesValues = CUDA.zeros(Float32,14)
    threads=(32,5)
    blocks =17

    mainArrDims= (634,521,632)
    dataBdim = (17,7,12)

    metaDataDims= (cld(mainArrDims[1],dataBdim[1] ),cld(mainArrDims[2],dataBdim[2]),cld(mainArrDims[3],dataBdim[3]))

    inBlockLoopX,inBlockLoopY,inBlockLoopZ= (fld(dataBdim[1] ,threads[1]),fld(dataBdim[2] ,threads[2]),dataBdim[3]    )
    #we are iterating here block by block sequentially
    loopXMeta,loopYZMeta= (metaDataDims[1],fld(metaDataDims[2]*metaDataDims[3] ,blocks)  )
    yTimesZmeta= metaDataDims[2]*metaDataDims[3]
    metaData = MetaDataUtils.allocateMetadata(mainArrDims,dataBdim);

    minX= CuArray([Float32(1110.0) ])
    maxX= CuArray([Float32(0.0)])
    minY= CuArray([Float32(1110.0)])
    maxY= CuArray([Float32(0.0)    ])
    minZ= CuArray([Float32(1110.0)])
    maxZ= CuArray([Float32(0.0) ])
    metaDataLength= metaDataDims[1]*metaDataDims[2]*metaDataDims[3]
    loopMeta= fld(metaDataLength,blocks )

    function uploadDataToMetaDataKernel(loopMeta,metaDataLength,minX, maxX, minY,maxY,minZ,maxZ,mainArrDims,metaDataDims,loopXMeta,loopYZMeta,yTimesZmeta,dataBdim, inBlockLoopX,inBlockLoopY,inBlockLoopZ,metaData)
        localQuesValues= @cuStaticSharedMem(UInt32, 14)
        isAnyPositive=true
        @iter3dOuter(metaDataDims,loopMeta,metaDataLength,
        begin

            @iterDataBlock(mainArrDims,dataBdim, inBlockLoopX,inBlockLoopY,inBlockLoopZ,xMeta,yMeta,zMeta,
            begin

                        if(xMeta>=2 && yMeta>=1 && zMeta>=3 && xMeta<=4 && yMeta<=5 && zMeta<=6  )
                            sync_threads()
                            PrepareArrtoBool.@uploadDataToMetaData()

                            PrepareArrtoBool.@uploadMinMaxesToShmem()
                            sync_threads()

                        end

                    end)

                end)

        return
    end

    @cuda threads=threads blocks=blocks uploadDataToMetaDataKernel(loopMeta,metaDataLength,minX, maxX, minY,maxY,minZ,maxZ,mainArrDims,metaDataDims,loopXMeta,loopYZMeta,yTimesZmeta,dataBdim, inBlockLoopX,inBlockLoopY,inBlockLoopZ,metaData)
    @test minX[1]== 2+1
    @test maxX[1]== 4+1
    @test minY[1]== 1+1
    @test maxY[1]== 5+1
    @test minZ[1]== 3+1
    @test maxZ[1]== 6+1
end #test set

###########################   getBoolCubeKernel
@testset "getBoolCubeKernel" begin

    using Revise, Parameters, Logging, Test
    using CUDA
    includet("C:\\GitHub\\GitHub\\NuclearMedEval\\test\\includeAllUseFullForTest.jl")
    using ..CUDAGpuUtils ,..IterationUtils,..ReductionUtils , ..MemoryUtils,..CUDAAtomicUtils, ..PrepareArrtoBool, ..BitWiseUtils
    using ..ResultListUtils, ..MetadataAnalyzePass,..MetaDataUtils,..WorkQueueUtils,..ProcessMainDataVerB,..HFUtils, ..ScanForDuplicates, ..Housdorff

    threads=(32,5)
    blocks =8
    mainArrDims= (512,512,155)
    dataBdim = (2,2,2)
    metaData = MetaDataUtils.allocateMetadata(mainArrDims,dataBdim);
    metaDataDims= size(metaData)
    arrGold = zeros(Int32,mainArrDims);
    arrSegm = zeros(Int32,mainArrDims);

        # so we should get all mins as 2 and all maxes as 5

        arrGold[9,2,5]= 2
        arrGold[3,6,4]= 2
        arrGold[3,3,4]= 2

        arrGold[4,4,4]= 2
        arrSegm[4,4,4]= 2

        arrSegm[3,18,4]= 2
        arrSegm[5,2,11]= 2

        gold3d,segm3d = CuArray(arrGold), CuArray(arrSegm)
        numberToLooFor =Int32(2)
        robustnessPercent = 0.9

        boolKernelArgs, mainKernelArgs,threadsBoolKern,blocksBoolKern ,threadsMainKern,blocksMainKern ,shmemSizeBool,shmemSizeMain=    preparehousedorfKernel(gold3d,segm3d,robustnessPercent,numberToLooFor)
        mainArrDims,dataBdim,metaData,metaDataDims,reducedGoldA,reducedSegmA,loopXinPlane,loopYinPlane,minxRes,maxxRes,minyRes,maxyRes,minzRes,maxzRes,fn,fp ,numberToLooFor,inBlockLoopXZIterWithPadding,shmemblockDataLoop,shmemblockDataLenght,loopAXFixed,loopBXfixed,loopAYFixed,loopBYfixed,loopAZFixed,loopBZfixed,loopdataDimMainX,loopdataDimMainY,loopdataDimMainZ,inBlockLoopX,inBlockLoopY,inBlockLoopZ,metaDataLength,loopMeta,loopWarpMeta,clearIterResShmemLoop,clearIterSourceShmemLoop,resShmemTotalLength,sourceShmemTotalLength=boolKernelArgs

        function locForKernel(goldGPU,segmGPU,mainArrDims,dataBdim,metaData,metaDataDims,reducedGoldA,reducedSegmA,loopXinPlane,loopYinPlane,minxRes,maxxRes,minyRes,maxyRes,minzRes,maxzRes,fn,fp ,numberToLooFor,inBlockLoopXZIterWithPadding,shmemblockDataLoop,shmemblockDataLenght,loopAXFixed,loopBXfixed,loopAYFixed,loopBYfixed,loopAZFixed,loopBZfixed,loopdataDimMainX,loopdataDimMainY,loopdataDimMainZ,inBlockLoopX,inBlockLoopY,inBlockLoopZ,metaDataLength,loopMeta,loopWarpMeta,clearIterResShmemLoop,clearIterSourceShmemLoop,resShmemTotalLength,sourceShmemTotalLength)

        @getBoolCubeKernel()

        return
    end

        @cuda threads=threadsBoolKern blocks=blocksBoolKern locForKernel(gold3d,segm3d,boolKernelArgs...)
            @test  Int64( maxzRes[])==6

        @test Int64(fn[])==3
        @test Int64(fp[])==2

        @test Int64(minxRes[])==2
        @test Int64( maxxRes[])==5
        @test Int64(minyRes[])==1
        @test  Int64(maxyRes[])==9
        @test  Int64( minzRes[])==2
        @test  Int64( maxzRes[])==6

        @test  reducedGoldA[9,2,5]==true
        @test   reducedGoldA[3,6,4]==true
        @test   reducedGoldA[3,3,4]==true
        @test   reducedGoldA[4,4,4]==true
        @test   reducedGoldA[5,5,5]==false
        @test   reducedGoldA[5,2,5]==false
        @test   reducedGoldA[5,2,90]==false

        @test  reducedGoldB[9,2,5]==true
        @test   reducedGoldB[3,6,4]==true
        @test   reducedGoldB[3,3,4]==true
        @test   reducedGoldB[4,4,4]==true
        @test   reducedGoldB[5,5,5]==false
        @test   reducedGoldB[5,2,5]==false
        @test   reducedGoldB[5,2,90]==false

        @test  reducedSegmA[3,18,4]==true
        @test   reducedSegmA[5,2,11]==true
        @test   reducedSegmA[4,4,4]==true
        @test   reducedSegmA[5,5,5]==false
        @test   reducedSegmA[5,2,5]==false
        @test   reducedSegmA[5,2,90]==false

        @test  reducedSegmB[3,18,4]==true
        @test   reducedSegmB[5,2,11]==true
        @test   reducedSegmB[4,4,4]==true
        @test   reducedSegmB[5,5,5]==false
        @test   reducedSegmB[5,2,5]==false
        @test   reducedSegmB[5,2,90]==false
end#test set

###########################   getBoolCubeKernel  2

using Revise, Parameters, Logging, Test
using CUDA
includet("C:\\GitHub\\GitHub\\NuclearMedEval\\test\\includeAllUseFullForTest.jl")
using ..CUDAGpuUtils ,..IterationUtils,..ReductionUtils , ..MemoryUtils,..CUDAAtomicUtils, ..PrepareArrtoBool
using ..ResultListUtils, ..MetadataAnalyzePass,..MetaDataUtils,..WorkQueueUtils,..ProcessMainDataVerB,..HFUtils, ..ScanForDuplicates

 @testset "getBoolCubeKernel 2 " begin
    localQuesValues = CUDA.zeros(Float32,14)
    threads=(32,5)
    blocks =7

    mainArrDims= (512,512,400)
    dataBdim = (10,10,10)
    metaData = MetaDataUtils.allocateMetadata(mainArrDims,dataBdim);
    metaDataDims= size(metaData)
    arrGold = zeros(Int32,mainArrDims);
    arrSegm = zeros(Int32,mainArrDims);
    loopAXFixed,loopBXfixed,loopAYFixed,loopBYfixed,loopAZFixed,loopBZfixed,loopdataDimMainX,loopdataDimMainY,loopdataDimMainZ,inBlockLoopX,inBlockLoopY,inBlockLoopZ,metaDataLength,loopMeta,loopWarpMeta,clearIterResShmemLoop,clearIterSourceShmemLoop,clearIterResShmemLoop,clearIterSourceShmemLoop,resShmemTotalLength,sourceShmemTotalLength = calculateLoopsIter(dataBdim,threads[1],threads[2],metaDataDims,blocks)
    minxRes,maxxRes,minyRes,maxyRes,minzRes,maxzRes,fn,fp  =getSmallForBoolKernel();
    reducedGoldA,reducedSegmA,reducedGoldB,reducedSegmB=  getLargeForBoolKernel(mainArrDims,dataBdim);

        # so we should get all mins as 2 and all maxes as 5

        # 1)   Left FP
        # 2)   Left FN
        arrGold[11,13,13]= 2
        arrSegm[11,13,15]= 2

    # 3)   Right FP
    # 4)   Right FN
        arrGold[20,13,13]= 2
        arrSegm[20,13,15]= 2

        arrGold[20,14,13]= 2
        arrSegm[20,14,15]= 2

    # 5)   Posterior FP

        # 6)   Posterior FN

        arrGold[13,11,14]= 2
        arrSegm[14,11,13]= 2

        arrGold[15,11,13]= 2
        arrSegm[13,11,15]= 2

        arrGold[16,11,13]= 2
        arrSegm[13,11,16]= 2

        # 7)   Anterior FP
        # 8)   Anterior FN
        arrGold[14,20,13]= 2
        arrSegm[13,20,14]= 2

        arrGold[15,20,13]= 2
        arrSegm[14,20,16]= 2

        arrGold[17,20,13]= 2
        arrSegm[15,20,17]= 2

        arrGold[18,20,17]= 2
        arrSegm[16,20,18]= 2

        # 9)   Top FP
        # 10)   Top FN
        arrGold[14,15,11]= 2
        arrSegm[14,16,11]= 2

        arrGold[17,15,11]= 2
        arrSegm[17,16,11]= 2

        arrGold[15,15,11]= 2
        arrSegm[15,16,11]= 2

        arrGold[16,15,11]= 2
        arrSegm[16,16,11]= 2

        arrGold[18,15,11]= 2
        arrSegm[18,16,11]= 2

        # 11)   Bottom FP
        # 12)   Bottom FN

        arrGold[14,15,20]= 2
        arrSegm[14,16,20]= 2

        arrGold[17,15,20]= 2
        arrSegm[17,16,20]= 2

        arrGold[15,15,20]= 2
        arrSegm[15,16,20]= 2

        arrGold[16,15,20]= 2
        arrSegm[16,16,20]= 2

        arrGold[18,15,20]= 2
        arrSegm[18,16,20]= 2

        arrGold[19,15,20]= 2
        arrSegm[19,16,20]= 2

        # 13)   main block Fp
        # 14)   main block Fn

        arrGold[14,15,14]= 2
        arrSegm[14,16,16]= 2

        arrGold[17,15,17]= 2
        arrSegm[17,16,13]= 2

        arrGold[15,15,13]= 2
        arrSegm[15,16,13]= 2

        arrGold[16,15,13]= 2
        arrSegm[16,16,13]= 2

        arrGold[18,15,13]= 2
        arrSegm[18,16,13]= 2

        arrGold[19,15,13]= 2
        arrSegm[19,16,13]= 2

        arrGold[19,12,13]= 2
        arrSegm[19,12,13]= 2

        arrGold[12,12,17]= 2
        arrSegm[12,12,16]= 2

    ########### some noise around

    arrGold[22,15,13]= 2
    arrSegm[18,22,13]= 2

    arrGold[19,22,13]= 2
    arrSegm[22,16,13]= 2

    arrGold[19,22,13]= 2
    arrSegm[19,12,22]= 2

    arrGold[12,12,17]= 2
    arrSegm[12,12,22]= 2

    arrGold[9,15,13]= 2
    arrSegm[18,9,13]= 2

    arrGold[19,22,9]= 2
    arrSegm[9,16,13]= 2

    arrGold[19,9,13]= 2
    arrSegm[19,12,9]= 2

    arrGold[9,12,17]= 2
    arrSegm[12,9,22]= 2

    gold3d,segm3d = CuArray(arrGold), CuArray(arrSegm)
    numberToLooFor =Int32(2)

    function locForKernelB(mainArrDims,dataBdim,metaData,metaDataDims,reducedGoldA,reducedSegmA,reducedGoldB,reducedSegmB,minxRes,maxxRes,minyRes,maxyRes,minzRes,maxzRes,fn,fp ,gold3d,segm3d,numberToLooFor,loopAXFixed,loopBXfixed,loopAYFixed,loopBYfixed,loopAZFixed,loopBZfixed,loopdataDimMainX,loopdataDimMainY,loopdataDimMainZ,inBlockLoopX,inBlockLoopY,inBlockLoopZ,metaDataLength,loopMeta,loopWarpMeta,clearIterResShmemLoop,clearIterSourceShmemLoop,clearIterResShmemLoop,clearIterSourceShmemLoop,resShmemTotalLength,sourceShmemTotalLength)

    @getBoolCubeKernel()

    return
    end

    @cuda threads=threads blocks=blocks locForKernelB(mainArrDims,dataBdim,metaData,metaDataDims,reducedGoldA,reducedSegmA,reducedGoldB,reducedSegmB,minxRes,maxxRes,minyRes,maxyRes,minzRes,maxzRes,fn,fp ,gold3d,segm3d,numberToLooFor,loopAXFixed,loopBXfixed,loopAYFixed,loopBYfixed,loopAZFixed,loopBZfixed,loopdataDimMainX,loopdataDimMainY,loopdataDimMainZ,inBlockLoopX,inBlockLoopY,inBlockLoopZ,metaDataLength,loopMeta,loopWarpMeta,clearIterResShmemLoop,clearIterSourceShmemLoop,clearIterResShmemLoop,clearIterSourceShmemLoop,resShmemTotalLength,sourceShmemTotalLength)

            #   @test   metaData[2,2,2,getBeginingOfFpFNcounts()+1]==2
            #   @test   metaData[2,3,2,getBeginingOfFpFNcounts()+2]==3
            #   @test   metaData[2,2,4,getBeginingOfFpFNcounts()+3]==4

            @test Int64(metaData[2,2,2,getBeginingOfFpFNcounts()+1])==1
            @test Int64(metaData[2,2,2,getBeginingOfFpFNcounts()+2])==1
            @test Int64(metaData[2,2,2,getBeginingOfFpFNcounts()+3])==2
            @test Int64(metaData[2,2,2,getBeginingOfFpFNcounts()+4])==2
            @test Int64(metaData[2,2,2,getBeginingOfFpFNcounts()+5])==3
            @test Int64(metaData[2,2,2,getBeginingOfFpFNcounts()+6])==3
            @test Int64(metaData[2,2,2,getBeginingOfFpFNcounts()+7])==4
            @test Int64(metaData[2,2,2,getBeginingOfFpFNcounts()+8])==4
            @test Int64(metaData[2,2,2,getBeginingOfFpFNcounts()+9])==5
            @test Int64(metaData[2,2,2,getBeginingOfFpFNcounts()+10])==5
            @test Int64(metaData[2,2,2,getBeginingOfFpFNcounts()+11])==6
            @test Int64(metaData[2,2,2,getBeginingOfFpFNcounts()+12])==6
            @test Int64(metaData[2,2,2,getBeginingOfFpFNcounts()+13])==7
            @test Int64(metaData[2,2,2,getBeginingOfFpFNcounts()+14])==7
            @test Int64(metaData[2,2,2,getBeginingOfFpFNcounts()+15])== 1+2+3+4+5+6+7
            @test Int64(metaData[2,2,2,getBeginingOfFpFNcounts()+16])== 1+2+3+4+5+6+7
end# test set

end#total test set

using Revise, Parameters, Logging, Test
using CUDA
includet("C:\\GitHub\\GitHub\\NuclearMedEval\\test\\includeAllUseFullForTest.jl")
using ..CUDAGpuUtils ,..IterationUtils,..ReductionUtils , ..MemoryUtils,..CUDAAtomicUtils
using ..ResultListUtils, ..MetadataAnalyzePass,..MetaDataUtils,..WorkQueueUtils,..ProcessMainDataVerB,..HFUtils, ..ScanForDuplicates

##### iter data block
singleVal = CUDA.zeros(Int64,14)
threads=(32,5)

blocks =17
mainArrDims= (67,177,90)
dataBdim = (43,21,17)
indices = CUDA.zeros(Bool,mainArrDims)

metaDataDims= (cld(mainArrDims[1],dataBdim[1] ),cld(mainArrDims[2],dataBdim[2]),cld(mainArrDims[3],dataBdim[3]))

inBlockLoopX,inBlockLoopY,inBlockLoopZ= (fld(dataBdim[1] ,threads[1]),fld(dataBdim[2] ,threads[2]),dataBdim[3]    )
#we are iterating here block by block sequentially
metaDataLength= metaDataDims[1]*metaDataDims[2]*metaDataDims[3]
loopMeta= fld(metaDataLength,blocks )

function iterDataBlocksKernel(loopMeta,metaDataLength,mainArrDims,singleVal,metaDataDims,dataBdim, inBlockLoopX,inBlockLoopY,inBlockLoopZ ,indices)
    @iter3dOuter(metaDataDims,loopMeta,metaDataLength,
    begin
        @iterDataBlockZdeepest(mainArrDims,dataBdim, inBlockLoopX,inBlockLoopY,inBlockLoopZ,xMeta,yMeta,zMeta,
        begin
       CUDA.@atomic singleVal[1]+=1
        indices[x,y,z]=true
        #CUDA.@cuprint "x $(x) y $(y) z $(z) zMeta $(zMeta)  yMeta $(yMeta) xMeta $(xMeta)  idX $(blockIdxX()) \n"

    end,begin indices[1]=true
                 end )end)

    return
end
@cuda threads=threads blocks=blocks iterDataBlocksKernel(loopMeta,metaDataLength,mainArrDims,singleVal,metaDataDims,dataBdim, inBlockLoopX,inBlockLoopY,inBlockLoopZ ,indices)
@test singleVal[1]==mainArrDims[1]*mainArrDims[2]*mainArrDims[3]
Int64(singleVal[1])
7*7*7
inddd = Array(indices);
cartt = CartesianIndices(inddd);
filtered = filter(cc-> !inddd[cc]  ,cartt)
sum(indices)


# OEWS.jl
#= This script pulls the data from
- U.S. Department of Labor
  - Bureau of Labor Statistics
    - Occupational Employment and Wage Statistics
	  - National industry-specific and by ownership
	    - OCC code
		- NAICS
		- Average Annual Salary
		- Number of Employees

For the private sector.

## Source: https://www.bls.gov/oes/tables.htm
=#
using InfoZIP: InfoZIP, unzip
using Downloads: Downloads, download
using HTTP: HTTP, URI, request
using ExcelFiles: ExcelFiles, load, openxl
using XLSX: XLSX, readxlsx, readtable, sheetnames
using CSV: CSV
using DataFrames: DataFrames, DataFrame, combine, groupby, ByRow, dropmissing!, rename!,
				  transform!, select!, select, subset!, leftjoin, Not

#=
We first download the Occupational Employment and Wage Statistics data files for 2008-2019.
=#
for yr in 08:19
	yr = lpad(yr, 2, '0')
	isfile(joinpath("data", "oews_$yr.zip")) ||
		download(string(URI(scheme = "https",
                            host = "www.bls.gov",
                            path = "/oes/special.requests/oesm$(yr)in4.zip")),
				 joinpath("data", "oews_$yr.zip"))
end
#=
Unzip 2008-2011
=#
for yr in 08:11
	yr = lpad(yr, 2, '0')
	isdir(joinpath("data", "oesm$(yr)in4")) || mkdir(joinpath("data", "oesm$(yr)in4"))
    isempty(readdir(joinpath("data", "oesm$(yr)in4"))) &&
        unzip(joinpath("data", "oews_$(yr).zip"), joinpath("data", "oesm$(yr)in4"))
end
#=
Unzip 2012-2019
=#
for yr in 12:19
    isdir(joinpath("data", "oews_$(yr).zip")) ||
		unzip(joinpath("data", "oews_$(yr).zip"), "data")
end
#=
Download the crosswalk for OCC codes from 2010 - 2019.
=#
if !isfile(joinpath("data", "oes_2019_hybrid_structure.xlsx"))
	download(string(URI(scheme = "https",
						host = "www.bls.gov",
						path = "/oes/oes_2019_hybrid_structure.xlsx")),
			 joinpath("data", "oes_2019_hybrid_structure.xlsx"))
end
#=
Download the crosswalk for OCC codes from 2000 - 2010.
=#
if !isfile(joinpath("data", "may_2010_occs.xls"))
	download(string(URI(scheme = "https",
						host = "www.bls.gov",
						path = "/oes/may_2010_occs.xls")),
			 joinpath("data", "may_2010_occs.xls"))
end
#=
Values might be available or not. This helper function handles that logic.
=#
# get_the_number(x) = get_the_number(string(x))
get_the_number(x::Number) = x
get_the_number(x::AbstractString) =
    occursin(r"^\d+\.?\d+$", x) ? parse(Int, x) : missing

"""
	parse_oews(file::AbstractString)::DataFrame

Returns the parsed data from an OEWS file.
"""
function parse_oews(file::AbstractString)
	if endswith(file, "xls")
		data = DataFrame(load(file, only(openxl(file).workbook.sheet_names())))
	elseif endswith(file, "xlsx")
		data = readxlsx(file)
        data = DataFrame(readtable(file, first(sheetnames(data)))...)
	end
	rename!(data, names(data) .=> lowercase.(names(data)))
	data[!,:year] .= parse(Int, match(r"\d{4}", file).match)
	select!(data, [:year, :occ_code, :occ_title, :naics, :tot_emp, :a_mean])
	transform!(data,
			   :a_mean => ByRow(get_the_number),
			   :tot_emp => ByRow(get_the_number),
			   renamecols = false)
end

data = DataFrame(year = Int[],
				 occ_code = String[],
				 occ_title = String[],
				 naics = String[],
				 tot_emp = Union{Missing,Int}[],
				 a_mean = Union{Missing,Int}[])
# empty!(data)
for dir in filter!(dir -> occursin(r"^oesm\d{2}in4$", dir), readdir("data"))
	for (root, dirs, files) in walkdir(joinpath("data", dir))
		yr = match(r"\d{2}", dir).match
		for file in filter!(file -> occursin("natsector_M20$(yr)_dl", file), files)
            occursin(r"\$", file) && continue # Skip opened Windows file
			println(joinpath(root, file))
			append!(data, parse_oews(joinpath(root, file)))
		end
	end
end
subset!(data, :naics => ByRow(!isequal("99")))
unique!(data)

x = subset(data, :occ_code => ByRow(isequal("15-1256")))
transform!(x, [:tot_emp, :a_mean] => ByRow(*) => :s_w)
sum(x[!,:s_w]) / sum(x[!,:tot_emp])
#=
We care about the Software Developers and Software Quality Assurance Analysts and Testers occupation.
This is a hybrid OCC code that started being used in 2019 (15-1256).
We will reconstruct the values for it using the crosswalk.
=#
cw19_18_10 = readxlsx(joinpath("data", "oes_2019_hybrid_structure.xlsx"))
cw19_18_10 = cw19_18_10["OES2019 Hybrid"]
cw19_18_10 = DataFrame(cw19_18_10["A7:I874"],
					   strip.(vec(cw19_18_10["A6:I6"])))
select!(cw19_18_10, [1, 2, 5, 7])
foreach(println, names(cw19_18_10))
rename!(cw19_18_10, [:oes19, :title, :oes18, :soc10])
#=
For 2008 - 2009 we need to use the 2000-2010 SOC crosswalk.
=#
cw10_00 = DataFrame(load(joinpath("data", "may_2010_occs.xls"),
						  "OES use of combined SOC data.!A11:F838"))
select!(cw10_00, [1, 3, 5])
rename!(cw10_00, [:oes10, :soc10, :soc00])
#=
Combining both crosswalks we get the proper one for 2000 - 2019
=#
cw = leftjoin(cw19_18_10, cw10_00, on = :soc10)
subset!(cw, :oes19 => ByRow(isequal("15-1256")))
cw = sort!(unique!(reduce(vcat, eachcol(select(cw, Not(:title))))))
#=
Since it is a single OCC code we can filter the records and aggregate easily.
=#
subset!(data, :occ_code => ByRow(∈(cw)))
#=
In case there are missing values, we can omit those and compute the mean based on the available data.
=#
dropmissing!(data)
#=
We must compute the tot_emp and a_mean.
=#
transform!(data, [:tot_emp, :a_mean] => ByRow(*) => :salary_and_wages)
select!(data, [:year, :naics, :tot_emp, :salary_and_wages])
data = combine(groupby(data, [:year, :naics]),
               [:tot_emp, :salary_and_wages] .=> sum,
               renamecols = false)
transform!(data, [:salary_and_wages, :tot_emp] => ByRow(/) => :a_mean)
select!(data, [:year, :naics, :tot_emp, :a_mean])
sort!(data)
CSV.write(joinpath("data", "oews_15-1256.csv"), data)


using Flux, Gym, Printf
using Flux.Tracker: data, grad
using Flux.Optimise: _update_params!
using Statistics: mean
using DataStructures: CircularBuffer
using Distributions: sample

using CuArrays

#Load game environment

env = PendulumEnv()

# ----------------------------- Parameters -------------------------------------

STATE_SIZE = length(reset!(env))
ACTION_SIZE = 1#length(env.actions)
ACTION_BOUND = 2#env.action_space.hi
MAX_EP = 50_000
MAX_EP_LENGTH = 1000

BATCH_SIZE = 64
MEM_SIZE = 100_000
MIN_EXP_SIZE = 50_000

γ = 99f-2     # discount rate

τ = 1f-3 # for running average while updating target networks
η_act = 1f-4   # Learning rate
η_crit = 1f-3
L2_DECAY = 1f-2

# Ornstein-Uhlenbeck Noise params
μ = 0f0
θ = 15f-2
σ = 2f-1

# --------------------------------- Memory ------------------------------------

memory = CircularBuffer{Any}(MEM_SIZE)

function getData(batch_size = BATCH_SIZE)
  # Getting data in shape
  minibatch = sample(memory, batch_size)
  x = hcat(minibatch...)

  s      =   hcat(x[1, :]...) |> gpu
  a      =   hcat(x[2, :]...) |> gpu
  r      =   hcat(x[3, :]...) |> gpu
  s′     =   hcat(x[4, :]...) |> gpu
  s_mask = .!hcat(x[5, :]...) |> gpu

  return s, a, r, s′, s_mask
end

# -------------------------------- Action Noise --------------------------------

struct OUNoise
  μ
  θ
  σ
  X
end

ou = OUNoise(μ, θ, σ, zeros(Float32, ACTION_SIZE) |> gpu)

function sample_noise(ou::OUNoise)
  dx     = ou.θ * (ou.μ .- ou.X)
  dx   .+= ou.σ * randn(Float32, length(ou.X)) |> gpu
  ou.X .+= dx
end

# Noise scale
τ_ = 25
ϵ  = exp(-1f0 / τ_)
noise_scale = 1f0 / ACTION_BOUND

# ----------------------------- Model Architecture -----------------------------

w_init(dims...) = 6f-3rand(Float32, dims...) .- 3f-3

actor = Chain(Dense(STATE_SIZE, 400, relu),
	      Dense(400, 300, relu),
              Dense(300, ACTION_SIZE, tanh, initW=w_init),
              x -> x * ACTION_BOUND) |> gpu
actor_target = deepcopy(actor)

struct crit
  state_crit
  act_crit
  sa_crit
end

Flux.@treelike crit

function (c::crit)(state, action)
  s = c.state_crit(state)
  a = c.act_crit(action)
  c.sa_crit(relu.(s .+ a))
end

Base.deepcopy(c::crit) = crit(deepcopy(c.state_crit),
                              deepcopy(c.act_crit),
			      deepcopy(c.sa_crit))

critic = crit(Chain(Dense(STATE_SIZE, 400, relu), Dense(400, 300)) |> gpu,
              Dense(ACTION_SIZE, 300) |> gpu,
	      Dense(300, 1, initW=w_init) |> gpu)
critic_target = deepcopy(critic)

# ------------------------------- Param Update Functions---------------------------------

function update_target!(target, model; τ = 1f0)
  for (p_t, p_m) in zip(params(target), params(model))
    p_t.data .= (1f0 - τ) * p_t.data .+ τ * p_m.data
  end
end

nullify_grad!(p) = p
nullify_grad!(p::TrackedArray) = (p.grad .= 0f0)

zero_grad!(model) = (model = mapleaves(nullify_grad!, model))

# ---------------------------------- Training ----------------------------------

function L2_loss(model)
  l2_loss = 0
  for p in params(model)
    l2_loss += sum(p .^ 2)
  end

  return L2_DECAY * l2_loss
end

opt_crit = ADAM(η_crit)
opt_act  = ADAM(η_act)

function replay()
  s, a, r, s′, s_mask = getData()

  # Update Critic
  a′ = data(actor_target(s′))
  v′ = data(critic_target(s′, a′))
  y = r .+ γ * v′ .* s_mask	# set v′ to 0 where s_ is terminal state

  v = critic(s, a)
  loss_crit = Flux.mse(y, v) + L2_loss(critic)

  zero_grad!(critic)
  Flux.back!(loss_crit)
  _update_params!(opt_crit, params(critic))

  # Update Actor
  actions = actor(s)
  crit_out = critic(s, actions)

  zero_grad!(actor)
  Flux.back!(-sum(crit_out))
  _update_params!(opt_act, params(actor))

  # Update Target models
  update_target!(actor_target, actor; τ = τ)
  update_target!(critic_target, critic; τ = τ)
end

# --------------------------- Helper Functions --------------------------------

# Stores tuple of state, action, reward, next_state, and done
remember(state, action, reward, next_state, done) =
  push!(memory, [data.((state, action, reward[1], next_state))..., done])

# Choose action according to policy PendulumPolicy
function action(state, train=true)
  state = reshape(data(state), size(state)..., 1)
  act_pred = data(actor(state |> gpu))
  if train
    act_pred .+= noise_scale * sample_noise(ou)
  end
  clamp.(act_pred, -ACTION_BOUND, ACTION_BOUND) # returns action
end

function episode!(env, train=true)
  total_reward = 0f0
  s = reset!(env)
  for ep=1:MAX_EP_LENGTH
    a = action(s, train)
    s′, r, done, _ = step!(env, a)
    total_reward += data(r)[1]
    if train
      remember(s, a, r, s′, done)
      replay()
    end
    s = s′
  end
  total_reward
end

# -------------------------------- Testing -------------------------------------

# Returns average score over 100 episodes
function test()
  score_mean = 0f0
  for _=1:100
    total_reward = episode!(env, false)
    score_mean += total_reward / 100
  end
  return score_mean
end

# ------------------------------ Training --------------------------------------

# Populate memory with random actions

s = reset!(env)
for e=1:MIN_EXP_SIZE
  global s
  a = 2rand(Float32) * ACTION_BOUND - ACTION_BOUND
  s′, r, done, _ = step!(env, a)
  remember(s, a, r, s′, done)
  s = s′
end

for e=1:MAX_EP
  global noise_scale
  total_reward = episode!(env)
  total_reward = @sprintf "%9.3f" total_reward
  print("Episode: $e | Score: $total_reward | ")
  score_mean = test()
  score_mean = @sprintf "%9.3f" score_mean
  println("Mean score over 100 test episodes: $score_mean")
  noise_scale *= ϵ
end


module WeyandDB

include("./CSV.jl")
include("./FileOps.jl")
using .CSV, .FileOps

export eachpair

const DATEDICT = Dict{Vector{String}, Int}(
    ["ja", "jan", "january"]   => 01,
    ["fe", "feb", "february"]  => 02,
    ["ma", "mar", "march"]     => 03,
    ["ap", "apr", "april"]     => 04,
    ["may"]                    => 05,
    ["ju", "jun", "june"]      => 06,
    ["jl", "jul", "july"]      => 07,
    ["au", "aug", "august"]    => 08,
    ["se", "sep", "september"] => 09,
    ["oc", "oct", "october"]   => 10,
    ["no", "nov", "november"]  => 11,
    ["de", "dec", "december"]  => 12
)
# ============================================================================ #
struct DB
    db::Vector{Pair{Int,String}}
end
get_ids(db::DB) = first.(db.db)
Base.length(db::DB) = length(db.db)
Base.getindex(db::DB, k::Integer) = db.db[k]
function Base.getindex(db::DB; id::Integer=-1)
    k = findfirst(isequal(id), get_ids(db))
    k == nothing && error("No entry with id $(id) exists in the database")
    return db.db[k]
end
# ============================================================================ #
struct DBIterator
    db::DB
end

eachpair(db::DB) = DBIterator(db)

Base.IteratorSize() = HasLength()
Base.IteratorEltype() = HasEltype()
Base.length(db::DBIterator) = length(db.db)
Base.eltype(db::DBIterator) = Tuple{Vector{Float64}, Vector{Float64}}
Base.iterate(db::DBIterator) = Base.iterate(db, 1)
function Base.iterate(db::DBIterator, state::Integer)
    state > length(db) && return nothing
    return get_data(db.db, state), state+1
end
# ============================================================================ #
function get_month_id(str::AbstractString)
    for (k, v) in DATEDICT
        if str in k
            return v
        end
    end
    return -1
end
# ============================================================================ #
function get_id(file::AbstractString)
    idstr = replace(basename(file), "S_SPK.csv"=>"")
    m = match(r"(\d{4})([A-Za-z]{2,3})(\d+)", idstr)
    m == nothing && error("Failed to parse file name into id")

    yr = parse(Int, m[1])
    mo = get_month_id(lowercase(m[2]))
    id = parse(Int, m[3])

    return yr * Int(1e5) + mo * Int(1e3) + id
end
# ============================================================================ #
function get_database(::String, f::Function=x->true)
    files = find_files(joinpath(@__DIR__, "..", "data"), r".*\.csv")
    filter!(files) do file
        f(get_id(file))
    end
    ids = get_id.(files)
    ks = sortperm(ids)
    return DB(map((x,y)->x=>y, ids[ks], files[ks]))
end
# ============================================================================ #
function get_data(db::DB, idx::Integer=0; id::Integer=-1)
    MS2SEC = 0.001

    ((idx < 1 && id < 0) || idx > length(db)) && error("Invalid index / id!")
    if idx > 0
        file = db[idx].second
    elseif id > -1
        file = db[id=id].second
    end

    d = CSV.parse(file, ',', [Float64, Int])

    lgn = d["timestamps_ms"][findall(isequal(0), d["ids"])] .* MS2SEC
    ret = d["timestamps_ms"][findall(>(0), d["ids"])] .* MS2SEC

    # some records have negative timestamps, so shift both accordingly (plus
    # 100ms so we don't have any timestamps at exactly 0)
    t0 = min(minimum(ret), minimum(lgn))
    t0 = t0 < 0 ? abs(t0) + 0.1 : 0.1

    # retina recordings are S-potentials, so they do NOT include conduction
    # delays, thus we shift the retinal timestamps backwards by 2.4 ms (relative
    # to the LGN) to avoid any weirdness that the short delay may introduce
    # NOTE: 2.4 ms is the shifte required to match the median
    # S-potential -> LGN delay (0.4 ms) to the median RGC->LGN delay from
    # PairsDB (2.8 ms)
    return ret .+ (t0 - 0.0024), lgn .+ t0
end
# ============================================================================ #
end