module NameList
# See Table 2 of Cohen et al, 2020
# namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["entrainment_factor"] = 0.13
# namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["detrainment_factor"] = 0.51
# namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["entrainment_massflux_div_factor"] = 0.0
# namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["turbulent_entrainment_factor"] = 0.015
# namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["entrainment_smin_tke_coeff"] = 0.3
# namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["updraft_mixing_frac"] = 0.25
# namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["entrainment_sigma"] = 10.0
# namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["entrainment_scale"] = 0.004
# namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["sorting_power"] = 2.0
# namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["aspect_ratio"] = 0.2

# See Table 1 of Lopez Gomez et al, 2020
# namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["tke_ed_coeff"] = 0.14
# namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["tke_diss_coeff"] = 0.22
# namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["static_stab_coeff"] = 0.4 # Square of value in the paper

# See Table ? of He et al, 2021
# namelist["turbulence"]["EDMF_PrognosticTKE"]["pressure_normalmode_buoy_coeff1"] = 0.12 ==> alpha_b (scaling constant for virtual mass term)
# namelist["turbulence"]["EDMF_PrognosticTKE"]["pressure_normalmode_buoy_coeff2"] = 0.0 ==> alpha_b (scaling constant for virtual mass term)
# namelist["turbulence"]["EDMF_PrognosticTKE"]["pressure_normalmode_adv_coeff"] = 0.1 alpha_a (scaling constant for advection term)
# namelist["turbulence"]["EDMF_PrognosticTKE"]["pressure_normalmode_drag_coeff"] = 10.0 ==> alpha_d (scaling constant for drag term)
# namelist["turbulence"]["EDMF_PrognosticTKE"]["pressure_plume_spacing"] ==> r_d (horizontal length scale of plume spacing)

#NB: except for Bomex and life_cycle_Tan2018 cases, the parameters listed have not been thoroughly tuned/tested
# and should be regarded as placeholders only. Optimal parameters may also depend on namelist options, such as
# entrainment/detrainment rate formulation, diagnostic vs. prognostic updrafts, and vertical resolution
export default_namelist

using ArgParse
using ArtifactWrappers

import StaticArrays
const SA = StaticArrays

#! format: off
LESDrivenSCM_output_dataset = ArtifactWrapper(
    @__DIR__,
    isempty(get(ENV, "CI", "")),
    "LESDrivenSCM_output_dataset",
    ArtifactFile[
        ArtifactFile(url = "https://caltech.box.com/shared/static/0hnf7nkttueraaqf9tpkqsx38gjqx41p.nc", filename = "Stats.cfsite23_HadGEM2-A_amip_2004-2008.07.nc",),
    ],
)
LESDrivenSCM_output_dataset_path = get_data_folder(LESDrivenSCM_output_dataset)
#! format: on

import JSON

function parse_commandline()
    s = ArgParseSettings(; description = "namelist Generator")

    @add_arg_table! s begin
        "case_name"
        help = "The case name"
        arg_type = String
        required = true
    end

    return parse_args(s)
end

function default_namelist(::Nothing)

    args = parse_commandline()
    case_name = args["case_name"]
    return default_namelist(case_name)
end

function default_namelist(case_name::String; root::String = ".", write::Bool = true)

    namelist_defaults = Dict()
    namelist_defaults["meta"] = Dict()
    namelist_defaults["meta"]["uuid"] = basename(tempname())

    namelist_defaults["turbulence"] = Dict()
    namelist_defaults["turbulence"]["Ri_bulk_crit"] = 0.2
    namelist_defaults["turbulence"]["prandtl_number_0"] = 0.74

    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"] = Dict()
    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["surface_area"] = 0.1
    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["max_area"] = 0.9
    # mixing_length
    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["tke_ed_coeff"] = 0.14
    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["tke_diss_coeff"] = 0.22
    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["static_stab_coeff"] = 0.4
    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["Prandtl_number_scale"] = 53.0 / 13.0
    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["smin_ub"] = 0.1
    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["smin_rm"] = 1.5
    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["l_max"] = 1.0e6
    # entrainment
    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["entrainment_factor"] = 0.13
    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["detrainment_factor"] = 0.51
    # 1-layer nn parameters
    #! format: off
    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["general_ent_params"] =
        SA.SVector(0.3038, 0.719,-0.910,-0.483,
                   0.739, 0.0755, 0.178, 0.521,
                   0.0, 0.0, 0.843,-0.340,
                   0.655, 0.113, 0.0, 0.0)
    #! format: on

    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["entrainment_massflux_div_factor"] = 0.0
    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["turbulent_entrainment_factor"] = 0.015
    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["entrainment_smin_tke_coeff"] = 0.3
    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["updraft_mixing_frac"] = 0.25
    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["entrainment_sigma"] = 10.0
    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["area_limiter_scale"] = 10.0
    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["area_limiter_power"] = 3.0
    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["entrainment_scale"] = 0.004
    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["sorting_power"] = 2.0
    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["min_upd_velocity"] = 0.001
    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["aspect_ratio"] = 0.2
    # pressure
    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["min_updraft_top"] = 500.0
    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["pressure_normalmode_buoy_coeff1"] = 0.12
    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["pressure_normalmode_buoy_coeff2"] = 0.0
    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["pressure_normalmode_adv_coeff"] = 0.1
    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["pressure_normalmode_drag_coeff"] = 10.0

    # stochastic closures
    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["stochastic"] = Dict()
    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["stochastic"]["closure"] = "none"
    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["stochastic"]["entr_lognormal_var"] = 0.0
    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["stochastic"]["detr_lognormal_var"] = 0.0
    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["stochastic"]["sde_entr_theta"] = 1.0
    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["stochastic"]["sde_entr_std"] = 0.0
    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["stochastic"]["sde_detr_theta"] = 1.0
    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["stochastic"]["sde_detr_std"] = 0.0

    # From namelist
    namelist_defaults["grid"] = Dict()
    namelist_defaults["grid"]["dims"] = 1

    namelist_defaults["thermodynamics"] = Dict()
    namelist_defaults["thermodynamics"]["thermal_variable"] = "thetal"
    namelist_defaults["thermodynamics"]["sgs"] = "quadrature"
    namelist_defaults["thermodynamics"]["quadrature_order"] = 3
    namelist_defaults["thermodynamics"]["quadrature_type"] = "log-normal" #"gaussian" or "log-normal"

    namelist_defaults["time_stepping"] = Dict()
    namelist_defaults["time_stepping"]["dt_max"] = 12.0
    namelist_defaults["time_stepping"]["dt_min"] = 1.0
    namelist_defaults["time_stepping"]["adapt_dt"] = true
    namelist_defaults["time_stepping"]["cfl_limit"] = 0.5

    namelist_defaults["microphysics"] = Dict()
    namelist_defaults["microphysics"]["precipitation_model"] = "None"

    namelist_defaults["turbulence"]["scheme"] = "EDMF_PrognosticTKE"

    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["updraft_number"] = 1
    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["entrainment"] = "moisture_deficit" # "moisture_deficit" or "NN"
    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["extrapolate_buoyancy"] = true
    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["use_local_micro"] = true
    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["constant_area"] = false
    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["calculate_tke"] = true
    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["mixing_length"] = "sbtd_eq"
    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["env_buoy_grad"] = "quadratures"

    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["pressure_closure_buoy"] = "normalmode"
    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["pressure_closure_drag"] = "normalmode"
    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["pressure_closure_asp_label"] = "const"

    namelist_defaults["output"] = Dict()
    namelist_defaults["output"]["output_root"] = "./"

    namelist_defaults["stats_io"] = Dict()
    namelist_defaults["stats_io"]["stats_dir"] = "stats"
    namelist_defaults["stats_io"]["frequency"] = 60.0
    namelist_defaults["stats_io"]["skip"] = false

    if case_name == "Soares"
        namelist = Soares(namelist_defaults)
    elseif case_name == "Nieuwstadt"
        namelist = Nieuwstadt(namelist_defaults)
    elseif case_name == "Bomex"
        namelist = Bomex(namelist_defaults)
    elseif case_name == "life_cycle_Tan2018"
        namelist = life_cycle_Tan2018(namelist_defaults)
    elseif case_name == "Rico"
        namelist = Rico(namelist_defaults)
    elseif case_name == "TRMM_LBA"
        namelist = TRMM_LBA(namelist_defaults)
    elseif case_name == "ARM_SGP"
        namelist = ARM_SGP(namelist_defaults)
    elseif case_name == "GATE_III"
        namelist = GATE_III(namelist_defaults)
    elseif case_name == "DYCOMS_RF01"
        namelist = DYCOMS_RF01(namelist_defaults)
    elseif case_name == "GABLS"
        namelist = GABLS(namelist_defaults)
    elseif case_name == "SP"
        namelist = SP(namelist_defaults)
    elseif case_name == "DryBubble"
        namelist = DryBubble(namelist_defaults)
    elseif case_name == "LES_driven_SCM"
        namelist = LES_driven_SCM(namelist_defaults)
    else
        error("Not a valid case name")
    end

    if write
        write_file(namelist, root)
    end
    return namelist
end
function Soares(namelist_defaults)

    namelist = deepcopy(namelist_defaults)
    namelist["meta"]["casename"] = "Soares"

    namelist["grid"]["nz"] = 75
    namelist["grid"]["dz"] = 50.0

    namelist["time_stepping"]["t_max"] = 8 * 3600.0
    namelist["time_stepping"]["dt_min"] = 1.0

    namelist["meta"]["simname"] = "Soares"
    namelist["meta"]["casename"] = "Soares"

    return namelist
end
function Nieuwstadt(namelist_defaults)

    namelist = deepcopy(namelist_defaults)
    namelist["meta"]["casename"] = "Nieuwstadt"

    namelist["grid"]["nz"] = 75
    namelist["grid"]["dz"] = 50.0

    namelist["time_stepping"]["t_max"] = 8 * 3600.0
    namelist["time_stepping"]["dt_min"] = 1.2

    namelist["meta"]["simname"] = "Nieuwstadt"
    namelist["meta"]["casename"] = "Nieuwstadt"

    return namelist
end
function Bomex(namelist_defaults)

    namelist = deepcopy(namelist_defaults)
    namelist["meta"]["casename"] = "Bomex"

    namelist["grid"]["nz"] = 60
    namelist["grid"]["dz"] = 50.0

    namelist["time_stepping"]["t_max"] = 21600.0
    namelist["time_stepping"]["dt_min"] = 6.0

    namelist["meta"]["simname"] = "Bomex"
    namelist["meta"]["casename"] = "Bomex"

    return namelist
end
function life_cycle_Tan2018(namelist_defaults)

    namelist = deepcopy(namelist_defaults)
    namelist["meta"]["casename"] = "life_cycle_Tan2018"

    namelist["grid"]["nz"] = 75
    namelist["grid"]["dz"] = 40.0

    namelist["time_stepping"]["t_max"] = 6 * 3600.0
    namelist["time_stepping"]["dt_min"] = 10.0
    namelist["meta"]["simname"] = "life_cycle_Tan2018"
    namelist["meta"]["casename"] = "life_cycle_Tan2018"

    return namelist
end
function Rico(namelist_defaults)

    namelist = deepcopy(namelist_defaults)

    namelist["meta"]["casename"] = "Rico"

    namelist["grid"]["nz"] = 80
    namelist["grid"]["dz"] = 50.0

    # Rico cannot run with adaptive dt for now
    namelist["time_stepping"]["adapt_dt"] = true
    namelist["time_stepping"]["t_max"] = 86400.0
    namelist["time_stepping"]["dt_max"] = 9.0
    namelist["time_stepping"]["dt_min"] = 1.5

    namelist["microphysics"]["precipitation_model"] = "clima_1m"
    namelist["microphysics"]["τ_acnv_rai"] = 2500.0

    namelist["meta"]["simname"] = "Rico"
    namelist["meta"]["casename"] = "Rico"
    return namelist
end
function TRMM_LBA(namelist_defaults)

    namelist = deepcopy(namelist_defaults)

    namelist["meta"]["casename"] = "TRMM_LBA"

    namelist["grid"]["nz"] = 320
    namelist["grid"]["dz"] = 50.0

    namelist["time_stepping"]["t_max"] = 60 * 60 * 6.0
    namelist["time_stepping"]["dt_min"] = 1.0

    namelist["microphysics"]["precipitation_model"] = "cutoff"
    namelist["microphysics"]["τ_precip"] = 50.0

    namelist["meta"]["simname"] = "TRMM_LBA"
    namelist["meta"]["casename"] = "TRMM_LBA"

    return namelist
end
function ARM_SGP(namelist_defaults)

    namelist = deepcopy(namelist_defaults)
    namelist["meta"]["casename"] = "ARM_SGP"

    namelist["grid"]["nz"] = 88
    namelist["grid"]["dz"] = 50.0

    namelist["time_stepping"]["t_max"] = 3600.0 * 14.5
    namelist["time_stepping"]["dt_min"] = 2.0

    namelist["meta"]["simname"] = "ARM_SGP"
    namelist["meta"]["casename"] = "ARM_SGP"

    return namelist
end
function GATE_III(namelist_defaults)

    namelist = deepcopy(namelist_defaults)
    namelist["meta"]["casename"] = "GATE_III"

    namelist["grid"]["nz"] = 1700
    namelist["grid"]["dz"] = 10

    namelist["time_stepping"]["t_max"] = 3600.0 * 24.0
    namelist["time_stepping"]["dt_min"] = 2.0

    namelist["meta"]["simname"] = "GATE_III"
    namelist["meta"]["casename"] = "GATE_III"

    return namelist
end
function DYCOMS_RF01(namelist_defaults)

    namelist = deepcopy(namelist_defaults)

    namelist["meta"]["casename"] = "DYCOMS_RF01"

    namelist["grid"]["nz"] = 30
    namelist["grid"]["dz"] = 50

    namelist["time_stepping"]["t_max"] = 60 * 60 * 16.0
    namelist["time_stepping"]["dt_min"] = 6.0

    namelist["meta"]["simname"] = "DYCOMS_RF01"
    namelist["meta"]["casename"] = "DYCOMS_RF01"

    return namelist
end
function GABLS(namelist_defaults)

    namelist = deepcopy(namelist_defaults)

    namelist["meta"]["casename"] = "GABLS"

    namelist["grid"]["nz"] = 8
    namelist["grid"]["dz"] = 50.0

    namelist["time_stepping"]["t_max"] = 9 * 3600.0
    namelist["time_stepping"]["dt_min"] = 4.0
    namelist["time_stepping"]["dt_max"] = 8.0
    namelist["meta"]["simname"] = "GABLS"
    namelist["meta"]["casename"] = "GABLS"

    return namelist
end
# Not fully implemented yet - Ignacio
function SP(namelist_defaults)

    namelist = deepcopy(namelist_defaults)
    namelist["meta"]["casename"] = "SP"

    # this case is resolution dependent, we should check why
    namelist["grid"]["nz"] = 64
    namelist["grid"]["dz"] = 32

    namelist["time_stepping"]["t_max"] = 7200.0
    namelist["time_stepping"]["dt_min"] = 1.0
    namelist["meta"]["simname"] = "SP"
    namelist["meta"]["casename"] = "SP"

    return namelist
end
function DryBubble(namelist_defaults)
    namelist = deepcopy(namelist_defaults)
    namelist["meta"]["casename"] = "DryBubble"
    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["surface_area"] = 0.0

    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["pressure_normalmode_buoy_coeff1"] = 0.12
    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["pressure_normalmode_adv_coeff"] = 0.25
    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["pressure_normalmode_drag_coeff"] = 0.1

    namelist["grid"]["nz"] = 200
    namelist["grid"]["dz"] = 50.0

    namelist["stats_io"]["frequency"] = 10.0
    namelist["time_stepping"]["t_max"] = 1000.0
    namelist["time_stepping"]["dt_min"] = 0.5

    namelist["meta"]["simname"] = "DryBubble"
    namelist["meta"]["casename"] = "DryBubble"

    namelist["turbulence"]["EDMF_PrognosticTKE"]["entrainment_massflux_div_factor"] = 0.4

    return namelist
end

function LES_driven_SCM(namelist_defaults)
    namelist = deepcopy(namelist_defaults)
    namelist["grid"]["dz"] = 50.0
    namelist["grid"]["nz"] = 80

    namelist["stats_io"]["frequency"] = 10.0
    namelist["time_stepping"]["t_max"] = 3600.0 * 6
    namelist["time_stepping"]["dt_min"] = 1.0

    # use last 6 hours of LES simulation to drive LES
    namelist["t_interval_from_end_s"] = 3600.0 * 6
    # average in 1 hour interval around `t_interval_from_end_s`
    namelist["initial_condition_averaging_window_s"] = 3600.0

    namelist["meta"]["lesfile"] =
        joinpath(LESDrivenSCM_output_dataset_path, "Stats.cfsite23_HadGEM2-A_amip_2004-2008.07.nc")
    namelist["meta"]["simname"] = "LES_driven_SCM"
    namelist["meta"]["casename"] = "LES_driven_SCM"
    namelist["forcing"] = Dict()
    namelist["forcing"]["nudging_timescale"] = 6.0 * 3600.0

    return namelist
end

function write_file(namelist, root::String = ".")
    mkpath(root)

    @assert haskey(namelist, "meta")
    @assert haskey(namelist["meta"], "simname")

    casename = namelist["meta"]["casename"]
    open(joinpath(root, "namelist_$casename.in"), "w") do io
        JSON.print(io, namelist, 4)
    end

    return
end

if abspath(PROGRAM_FILE) == @__FILE__
    default_namelist(nothing)
end

end


"""
    MapFun <: Augmentor.Operation

Description
--------------

Maps the given function over all individual array elements.

This means that the given function is called with an individual
elements and is expected to return a transformed element that
should take the original's place. This further implies that the
function is expected to be unary. It is encouraged that the
function should be consistent with its return type and
type-stable.

Usage
--------------

    MapFun(fun)

Arguments
--------------

- **`fun`** : The unary function that should be mapped over all
    individual array elements.

See also
--------------

[`AggregateThenMapFun`](@ref), [`ConvertEltype`](@ref), [`augment`](@ref)

Examples
--------------

```julia
using Augmentor, ColorTypes
img = testpattern()

# subtract the constant RGBA value from each pixel
augment(img, MapFun(px -> px - RGBA(0.5, 0.3, 0.7, 0.0)))

# separate channels to scale each numeric element by a constant value
pl = SplitChannels() |> MapFun(el -> el * 0.5) |> CombineChannels(RGBA)
augment(img, pl)
```
"""
struct MapFun{T} <: Operation
    fun::T
end

@inline supports_lazy(::Type{<:MapFun}) = true

function applyeager(op::MapFun, img::AbstractArray, param)
    plain_array(map(op.fun, img))
end

function applylazy(op::MapFun, img::AbstractArray, param)
    mappedarray(op.fun, img)
end

function showconstruction(io::IO, op::MapFun)
    print(io, nameof(typeof(op)), '(', op.fun, ')')
end

function Base.show(io::IO, op::MapFun)
    if get(io, :compact, false)
        print(io, "Map function \"", op.fun, "\" over image")
    else
        print(io, "Augmentor.")
        showconstruction(io, op)
    end
end

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

"""
    AggregateThenMapFun <: Augmentor.Operation

Description
--------------

Compute some aggregated value of the current image using the
given function `aggfun`, and map that value over the current
image using the given function `mapfun`.

This is particularly useful for achieving effects such as
per-image normalization.

Usage
--------------

    AggregateThenMapFun(aggfun, mapfun)

Arguments
--------------

- **`aggfun`** : A function that takes the whole current image as
    input and which result will also be passed to `mapfun`. It
    should have a signature of `img -> agg`, where `img` will the
    the current image. What type and value `agg` should be is up
    to the user.

- **`mapfun`** : The binary function that should be mapped over
    all individual array elements. It should have a signature of
    `(px, agg) -> new_px` where `px` is a single element of the
    current image, and `agg` is the output of `aggfun`.

See also
--------------

[`MapFun`](@ref), [`ConvertEltype`](@ref), [`augment`](@ref)

Examples
--------------

```julia
using Augmentor
img = testpattern()

# subtract the average RGB value of the current image
augment(img, AggregateThenMapFun(img -> mean(img), (px, agg) -> px - agg))
```
"""
struct AggregateThenMapFun{A,M} <: Operation
    aggfun::A
    mapfun::M
end

@inline supports_lazy(::Type{<:AggregateThenMapFun}) = true

function applyeager(op::AggregateThenMapFun, img::AbstractArray, param)
    # promote the result to float type because `aggfun` might not widely support `Gray{N0f8}`/`N0f8` types well
    # e.g, type instability in https://github.com/JuliaGraphics/ColorVectorSpace.jl/issues/134
    agg = convert(floattype(eltype(img)), op.aggfun(img))
    plain_array(map(x -> op.mapfun(x, agg), img))
end

function applylazy(op::AggregateThenMapFun, img::AbstractArray, param)
    agg = convert(floattype(eltype(img)), op.aggfun(img))
    mappedarray(x -> op.mapfun(x, agg), img)
end

function showconstruction(io::IO, op::AggregateThenMapFun)
    print(io, nameof(typeof(op)), '(', op.aggfun, ", ", op.mapfun, ')')
end

function Base.show(io::IO, op::AggregateThenMapFun)
    if get(io, :compact, false)
        print(io, "Map result of \"", op.aggfun, "\" using \"", op.mapfun, "\" over image")
    else
        print(io, "Augmentor.")
        showconstruction(io, op)
    end
end


module SortBenchmarks

include(joinpath(dirname(@__FILE__), "..", "utils", "RandUtils.jl"))

using .RandUtils
using BenchmarkTools

const SUITE = BenchmarkGroup()
const LIST_SIZE = 50000
const LISTS = (
    ("ascending", collect(1:LIST_SIZE)),
    ("descending", collect(LIST_SIZE:-1:1)),
    ("ones", ones(LIST_SIZE)),
    ("random", samerand(LIST_SIZE))
)

#####################################
# QuickSort/MergeSort/InsertionSort #
#####################################

for (group, Alg) in (("quicksort", QuickSort), ("mergesort", MergeSort), ("insertionsort", InsertionSort))
    g = addgroup!(SUITE, group)
    for (kind, list) in LISTS
        ix = collect(1:length(list))
        g["sort forwards", kind] = @benchmarkable sort($list; alg = $Alg)
        g["sort reverse", kind] = @benchmarkable sort($list; alg = $Alg, rev = true)
        g["sortperm forwards", kind] = @benchmarkable sortperm($list; alg = $Alg)
        g["sortperm reverse", kind] = @benchmarkable sortperm($list; alg = $Alg, rev = true)
        g["sort! forwards", kind] = @benchmarkable sort!(x; alg = $Alg) setup=(x = copy($list))
        g["sort! reverse", kind] = @benchmarkable sort!(x; alg = $Alg, rev = true) setup=(x = copy($list))
        g["sortperm! forwards", kind] = @benchmarkable sortperm!(x, $list; alg = $Alg) setup=(x = copy($ix))
        g["sortperm! reverse", kind] = @benchmarkable sortperm!(x, $list; alg = $Alg, rev = true) setup=(x = copy($ix))
    end
    for b in values(g)
        b.params.time_tolerance = 0.30
    end
end

############
# issorted #
############

g = addgroup!(SUITE, "issorted")

for (kind, list) in LISTS
    g["forwards", kind] = @benchmarkable issorted($list)
    g["reverse", kind] = @benchmarkable issorted($list; rev = true)
end

for b in values(g)
    b.params.time_tolerance = 0.30
end

end # module


using XInts
using XInts: SLimb, slimbmin, slimbmax, Limb, ClongMax, CulongMax, CdoubleMax, BITS_PER_LIMB,
             add!, sub!, mul!, com!, xor!, lshift!, rshift!, and!, ior!, SLimbW, LimbW, vec, is_short
using BitIntegers
using Random

function validate(x::XInt)
    @test x isa XInt
    if XInts.is_short(x)
        @test x.x != slimbmin
    else
        xv = x.v
        xl = abs(x.x)
        @test length(xv) >= xl >= 1
        @test !iszero(xv[xl])
        if xl == 1
            @test xv[1] >= XInts.limb1min
        end
    end
    x
end

validate(x) = @test x isa XInt

vint(x) = validate(XInt(x))

@testset "constructor" begin
    function test(a)
        x = validate(XInt(a))
        @test x isa XInt
        if a isa XInt
            @test x === a
        else
            @test x == a
        end
        if !(a isa BigInt)
            @test x == BigInt(a)
        end
    end

    @testset "XInt(::XInt)" begin
        test(XInt(2))
        test(XInt(big(2)^200))
    end
    @testset "XInt(::Bool)" begin
        test(true)
        test(false)
    end

    @testset "XInt(::$T)" for T in [Base.BitInteger_types...,
                                    Int256, UInt256, Int512, UInt512, Int1024, UInt1024]
        for x = rand(T, 20)
            test(x)
        end
        for d = T[0, 1, 2]
            test(d)
            if T <: Signed
                test(-d)
            end
            test(typemin(T) + d)
            test(typemax(T) - d)
            test((typemax(T) >>> 1) + d)
            test((typemax(T) >>> 1) - d)
        end
    end

    @testset "XInt(::BigInt)" begin
        for T = (Bool, Int8, Int, Int128)
            for x = BigInt.(rand(T, 5))
                test(x)
            end
        end
    end

    @testset "XInt(::$F)" for F in [Float16, Float32, Float64]
        for T = (Bool, Int8, Int32)
            for x = F.(rand(T, 5))
                isinteger(x) && test(x)
            end
        end
        @test XInt(F(2.0)) === XInt(2)
        @test_throws InexactError XInt(F(2.2))
    end
end

@testset "conversions" begin
    @testset "rem(::XInt, ::Type{Bool})" begin
        for x = Any[rand(Limb), big(rand(Int128)), rand(-big(2)^200:big(2)^200)]
            @test vint(x) % Bool === x % Bool
        end
    end
    @testset "rem(::XInt, ::Type{$T})" for T in Base.BitInteger_types
        xs = T[0, 2, rand(T, 10)..., typemax(T)]
        if T <: Signed
            push!(xs, -1, -2, typemin(T))
        end
        for x = xs
            @test vint(x) % T === x
            @test vint(big(rand(T)) << (8*sizeof(T)) +
                       big(x)) % T === x
            @test vint(big(rand(T)) << (16*sizeof(T)) +
                       big(rand(T)) << (8*sizeof(T)) +
                       big(x)) % T === x
        end
    end

    @testset "Bool(::XInt)" begin
        @test Bool(XInt(0)) == false
        @test Bool(XInt(1)) == true
        for x = Any[rand(Int8, 10); rand(Int128, 10)]
            (x == 0 || x == 1) && continue
            @test_throws InexactError Bool(XInt(x))
        end
    end
    @testset "$T(::XInt)" for T in Base.BitInteger_types
        for x = rand(T, 10)
            @test T(vint(x)) === x
        end
        if sizeof(T) > sizeof(SLimb)
            for x = T.(rand(T <: Signed ? Int32 : UInt32, 10))
                @test T(vint(x)) === x
            end
        end
        @test T(vint(typemax(T))) === typemax(T)
        @test T(vint(typemin(T))) === typemin(T)
        @test_throws InexactError T(XInt(typemax(T))+XInt(1))
        @test_throws InexactError T(XInt(typemax(T))+XInt(+10))
        @test_throws InexactError T(XInt(typemin(T))+XInt(-1))
        @test_throws InexactError T(XInt(typemin(T))+XInt(-10))
        @test_throws InexactError T(XInt(2)^200) # for sizeof(T) < sizeof(Limb)
        # test un-reduced XInt
        @test T(XInts._XInt(1, Limb[1])) === T(1)
    end

    @testset "BigInt(::XInt)" begin
        for T = (Bool, Int8, Int, Int128)
            for x = rand(T, 5)
                z = BigInt(x)
                y = vint(x)
                @test z == BigInt(y) isa BigInt
                @test z == big(y) isa BigInt
                @test z == y % BigInt isa BigInt
            end
        end
    end

    @testset "$F(::XInt, $R)" for
        F in (Float16, Float32, Float64, BigFloat, AbstractFloat, float),
        R in ((F===float) ? ("",) : (RoundToZero, RoundDown, RoundUp, RoundNearest, ""))

        for T = Base.BitInteger_types
            for x = rand(T, 10)
                if R == ""
                    @test F(vint(x)) == F(BigInt(x))
                    # have to use BigInt because of a Julia bug
                else
                    @test F(vint(x), R) == F(BigInt(x), R)
                end
            end
        end
    end
end

@testset "math" begin
    @testset "isodd/iseven" begin
        for x = Any[rand(Int8, 5)...,
                    rand((-1, 1)) .* rand(0:big(2)^200, 5)...]
            y = vint(x)
            @test iseven(y) == iseven(x) != isodd(y)
        end
    end

    @testset "iszero/isone" begin
        @test iszero(XInt(0))
        @test zero(XInt) === XInt(0)
        @test isone(XInt(1))
        @test one(XInt) === XInt(1)
        x = rand(Int8)
        @test iszero(x) == iszero(XInt(x))
        @test isone(x) == isone(XInt(x))
    end

    @testset "sign" begin
        @test sign(zero(XInt)) === XInt(0)
        @test signbit(zero(XInt)) === false
        @test signbit(one(XInt)) === false
        @test signbit(-one(XInt)) === true

        for x in rand(Int8, 10)
            y = XInt(x)
            @test sign(y) isa XInt
            @test sign(y) == sign(x)
            @test signbit(y) == signbit(x)
        end
    end

    @testset "digits" begin
        for r = big(2).^[5, 8, 64, 200]
            for b = rand(-r:r, 5)
                pad = rand(0:2)
                for base = [rand(8:12, 3); 100] # 100 > 62, cf. implementation of digits!
                    @test ndigits(b, pad=pad, base=base) == ndigits(XInt(b), pad=pad, base=base)
                    @test digits(b, pad=pad, base=base) == digits(XInt(b), pad=pad, base=base)
                end
            end
        end
    end

    @testset "prevpow & nextpow" begin
        # TODO: deduplicate this xs definition (from binary ops)
        xs = Any[1, 2, rand(3:1000, 10)..., slimbmax, slimbmax-1, slimbmax-2,
                 rand(Int128(2)^65:Int128(2)^100, 2)..., rand(big(3):big(2)^200, 2)...,
                 1024, big(2)^100, big(2)^200] # popcount == 1

        for x = xs, b = Int[2, 3, 4, 5, 9, 10, rand(11:100, 6)...]
            y = validate(nextpow(XInt(b), XInt(x)))
            @test y isa XInt
            @test y == nextpow(big(b), big(x))
            if b == 2
                # accepts other integer types for base
                @test y == nextpow(2, XInt(x))
            end
            y = validate(prevpow(XInt(b), XInt(x)))
            @test y isa XInt
            @test y == prevpow(big(b), big(x))
            if b == 2
                # accepts other integer types for base
                @test y == validate(prevpow(2, XInt(x)))
            end
        end
        for b = XInt[1, 0, -1, rand(-100:-2, 5)...]
            @test_throws DomainError prevpow(b, XInt(rand(1:100)))
            @test_throws DomainError nextpow(b, XInt(rand(1:100)))
        end
        for x = [0, -1, rand(-100:-2, 5)...]
            @test_throws DomainError prevpow(XInt(rand(2:20)), vint(x))
            @test_throws DomainError nextpow(XInt(rand(2:20)), vint(x))
        end

        # add some x which lead to increased size (in limbs)
        append!(xs, big.(rand(UInt, 4)) .<< (BITS_PER_LIMB .* [0, 1, 2, 3])')
        xs = [xs; .-(xs)]
        push!(xs, slimbmin)
        for x=xs
            y = validate(Base._nextpow2(XInt(x)))
            @test y isa XInt
            @test y == Base._nextpow2(big(x))
            y = validate(Base._prevpow2(XInt(x)))
            @test y isa XInt
            @test y == Base._prevpow2(big(x))
        end
    end

    @testset "sum" begin
        function testsum(bs)
            bs0 = deepcopy(bs)
            xs = XInt.(bs)
            xs0 = deepcopy(xs)
            sumbs = sum(bs)
            sumxs = sum(xs)
            # we test that sum didn't mess up with internal mutation
            @test is_short(sumxs) || all(x -> x !== sumxs, xs)
            @test xs == xs0 == bs0
            @test sumbs == sumxs
        end
        xs = BigInt.(rand(Int8, rand(1:10)))
        testsum(xs)
        append!(xs, rand(SLimb, rand(1:10)))
        testsum(shuffle!(xs))
        append!(xs, rand(SLimbW, rand(1:10)))
        testsum(shuffle!(xs))
        append!(xs, rand(-big(2)^200:big(2)^200, rand(1:10)))
        testsum(shuffle!(xs))
        # the first number x is a non-short, so when added with the initial
        # 0 in sum, it will return x itself: we must not mutate it later on
        xs = rand(-big(2)^65:big(2)^65, 10)
        testsum(xs)
        # test only big positive or negative numbers
        axs = abs.(xs)
        testsum(axs)
        testsum(.-axs)
        # with some shorts
        append!(axs, rand(slimbmin+1:slimbmax, 10))
        testsum(axs)
        testsum(.-axs)

        # tuples
        z1, z2, z3, z4 = zs = XInt.(xs[1:4])
        @test sum(zs) == sum(Tuple(zs))
        @test sum(zs) == z1 + z2 + z3 + z4
    end
end

@testset "misc" begin
    @testset "[deep]copy, signed, Signed, widen, hastypemax" begin
        # automatic implementation from XInt <: Signed
        for x = XInt[2, big(2)^70]
            @test copy(x) === x
            let y = validate(deepcopy(x))
                @test y == x
                @test x == 2 ? x === y : y !== x
            end
            @test signed(x) === x
            @test Signed(x) === x
            @test widen(x) === x
        end
        @test widen(XInt) == XInt
        @test !Base.hastypemax(XInt)
        if VERSION >= v"1.5"
            @test signed(XInt) == XInt
        end
    end
    @testset "parse, string, show" begin
        @test validate(parse(XInt, "123")) === XInt(123)

        for x = BigInt[-10:10;
                       rand(Int8, 10);
                       rand(Int128, 10);
                       rand(-big(2)^200:big(2)^200, 10)]
            @test string(x) == string(XInt(x))
            @test sprint(show, x) == sprint(show, XInt(x))
            for base = [-62:-2; 2:62], pad = rand(0:220, 30)
                @test string(x, base=base, pad=pad) == string(XInt(x), base=base, pad=pad)
            end
        end
    end
    @testset "float(XInt)" begin
        @test float(XInt) === BigFloat
    end
end

@testset "trunc" begin
    for T in [#=Float16,=# Float32, Float64] # TODO: add Float16, and also for BigInt
        @test validate(trunc(XInt, T(2.2))) === XInt(2)
        @test validate(trunc(XInt, T(-2.2))) === XInt(-2)
        @test_throws InexactError trunc(XInt, T(NaN))
        @test_throws InexactError trunc(XInt, T(Inf))
        @test validate(unsafe_trunc(XInt, T(2.2))) === XInt(2)
        @test validate(unsafe_trunc(XInt, T(-2.2))) === XInt(-2)
    end
end

@testset "comparisons" begin
    @testset "== (integers)" begin
        xs = BigInt[slimbmin, 0, 1, 2, slimbmax-1, slimbmax, slimbmax+1, big(2)^100]
        xs = [xs..., (-).(xs)...]
        for (a, b) in Iterators.product(xs, xs)
            x, y = XInt(a), XInt(b)
            @test x == x
            @test (a == b) == (x == y)
            @test a == x && x == a
            @test b == y && y == b
            for T in Base.BitInteger_types
                if typemin(T) <= b <= typemax(T)
                    z = T(b)
                    @test (x == z) == (a == z) == (z == x)
                end
            end
        end
    end
    @testset "XInt with floats" begin
        for T = Base.uniontypes(CdoubleMax)
            b = Int(maxintfloat(T))
            for x = rand(0:b, 10)
                @test vint(x) == T(x)
                @test T(x) == vint(x)
                @test vint(x-1) < T(x)
                @test T(x) < vint(x+1)
                @test T(x) <= vint(x) <= T(x)

            end
        end
    end
end

opmap(x) = x
opmap(::typeof(add!)) = +
opmap(::typeof(sub!)) = -
opmap(::typeof(mul!)) = *
opmap(::typeof(com!)) = ~
opmap(::typeof(lshift!)) = <<
opmap(::typeof(rshift!)) = >>
opmap(::typeof(and!)) = &
opmap(::typeof(ior!)) = |
opmap(::typeof(xor!)) = xor

function testop2(op, x, y)
    for a = (-x, x), b = (-y, y)
        s = if op === invmod
            try
                op(big(a), big(b))
            catch
                @test_throws DomainError op(XInt(a), XInt(b))
                continue
            end
        else
            opmap(op)(big(a), big(b))
        end
        @assert !(a isa XInt) # so that if op is mutating, it doesn't mutate a
        r = op(vint(a), vint(b))
        if op === divrem
            @test r isa Tuple{XInt,XInt}
            validate.(r)
        elseif op === gcdx
            @test r isa Tuple{XInt,XInt,XInt}
            validate.(r)
        elseif !(s isa BigInt)
            @test typeof(s) == typeof(r)
        else
            validate(r)
        end
        @test s == r
        if op === div
            for R = (RoundToZero, RoundDown, RoundUp)
                @test div(big(a), big(b), R) ==
                    validate(div(XInt(a), XInt(b), R)) isa XInt
            end
        end
        if opmap(op) != op
            # mutating
            tmp = XInts._XInt(0, Limb[])
            aa, bb = vint(a), vint(b)
            rtmp = validate(op(tmp, aa, bb))
            @test rtmp == r
            if !is_short(r)
                @test vec(rtmp) === vec(tmp) || !is_short(aa) && vec(rtmp) === vec(aa) ||
                    vec(rtmp) === vec(bb)
                aa2 = XInts._copy(aa)
                # test aliasing works between the source(s) and destination
                if aa != bb
                    rtmp = validate(op(aa2, aa2, bb))
                else
                    rtmp = validate(op(aa2, aa2, aa2))
                end
                @test rtmp == r
            end
            tmp = XInts._XInt(0)
            @test validate(op(tmp, vint(a), vint(b))) == r
        end
    end
end

function test_alloc(op, x, y, r=nothing)
    if VERSION >= v"1.5"
        @test 0 == @allocated op(x, y)
    end
    z = op(x, y)
    @test validate(z) == r
end

@testset "operations" begin

    xs = BigInt[0, 1, 2, rand(0:1000, 5)..., slimbmax, slimbmax-1, slimbmax-2,
                rand(SLimb, 5)..., rand(Limb, 5)...,
                Int128(slimbmax)+1, Int128(slimbmax)+2,
                (LimbW.(rand(Limb, 3)) .<< BITS_PER_LIMB)...,
                typemax(LimbW), LimbW(typemax(Limb)) << BITS_PER_LIMB, # and!
                (LimbW(typemax(Limb)) << BITS_PER_LIMB) + 1, # and!
                (let h = rand(Limb) # ior!
                     [big(h) << 2BITS_PER_LIMB | big(rand(LimbW)),
                      big(~h) << 2BITS_PER_LIMB | big(rand(LimbW)),
                      ]
                 end)...,
                (Int128(slimbmax).+rand(UInt8, 5))...,
                rand(Int128(2)^65:Int128(2)^100, 2)..., rand(big(0):big(2)^200, 2)...,
                slimbmin] # TODO: slimbmin and 0 are used twice when negated in test(...)

    # add!
    let
        x4 = rand(1:typemax(Limb))
        y4 = x4-1
        z = Limb(1)<<63
        y1 = rand(z : typemax(Limb))
        x1 = rand(Limb(0) : y1-z) # y1-z >= 0
        # y1-x1 >= z, so x1-y1 doesn't have top bit set
        x2 = rand(1:typemax(Limb))
        y2 = x2-1 # 2nd limb must absorb carry, and become 0
        x = XInt(x4) << 3BITS_PER_LIMB | XInt(x2) << BITS_PER_LIMB | x1
        y = XInt(y4) << 3BITS_PER_LIMB | XInt(y2) << BITS_PER_LIMB | y1
        u2 = rand(0:x2-1)
        u = XInt(y4) << 3BITS_PER_LIMB | XInt(u2) << BITS_PER_LIMB | y1
        # w, v: similar x and y, but with a 3rd limb, to have carry propagate
        w3 = rand(1:typemax(Limb))
        w = x | XInt(w3) << 2BITS_PER_LIMB
        v2 = rand(x2:typemax(Limb))
        v3 = rand(0:w3-1)
        v = XInt(y4) << 3BITS_PER_LIMB | XInt(v3) << 2BITS_PER_LIMB | XInt(v2) << BITS_PER_LIMB | y1
        push!(xs, x, y, u, v, w)
    end
    let # xor!
        # x < 0 with 3 limbs, y > 0 with 2 limbs, result with null 3rd limb
        # formula: x ⊻ y = -((-x-1) ⊻ y + 1), so find such result via
        # propagating the carry of +1 till 3rd limb
        x1 = rand(1:typemax(Limb))
        y1 = ~(x1-1)
        y2 = rand(1:typemax(Limb)) # y2 must not be null
        x2 = ~y2
        x3 = typemax(Limb)
        x = XInt(x3) << 2BITS_PER_LIMB | XInt(x2) << BITS_PER_LIMB | XInt(x1)
        y = XInt(y2) << BITS_PER_LIMB | XInt(y1)
        push!(xs, x, y)
    end

    @testset "$op(::XInt, ::XInt)" for op = (+, -, *, mod, rem, gcd, gcdx, lcm, &, |, xor,
                                             /, div, divrem, fld, cld, invmod,
                                             cmp, <, <=, >, >=, ==, flipsign, binomial,
                                             add!, sub!, mul!, and!, ior!, xor!)
        for x=xs, y=xs
            iszero(y) && op ∈ (/, mod, rem, div, divrem, fld, cld, invmod) &&
                continue
            op === binomial && !all(0 .<= [x, y] .<= 1000) && continue
            testop2(op, x, y)
        end
        if op === +
            x = XInt(2)^128
            y = validate(-x+1)
            @test validate(x+y) == big(x) + big(y)
        end

        # special cases
        if op === &
            # special case where the result takes one more limb than the operands
            x = -XInt(XInt(typemax(UInt))<<64 + 1)
            y = -XInt(2)^64
            z = validate(x & y)
            @test z == big(x) & big(y)
        elseif op === |
            x = -(XInt(2)^128)
            y = XInt(2)^128-1
            test_alloc(|, x, y, -1)
        end
    end

    function make_values(T, n=5)
        S = T === BigInt ? Int128 : T
        as = T[0, 1, 2, 3, rand(S, 2n)..., typemax(S), typemax(S)-1, typemax(S)-2]
        xs = BigInt.(filter(isinteger, as))
        if T == Limb # xor!
            x1 = rand(1:typemax(T))
            y = ~(x1-1)
            x = BigInt(typemax(T)) << BITS_PER_LIMB | x1
            push!(as, y)
            push!(xs, x)
        end
        push!(xs, BigInt(2)^BITS_PER_LIMB) # also for xor!
        if T <: Signed
            append!(as, (-).(as))
            push!(as, typemin(S))
        end
        if T <: Integer
            push!(xs, typemin(S))
            append!(xs, rand(S, 2n))
        end
        append!(xs, rand(big(2)^65:big(2)^100, n))
        append!(xs, rand(big(0):big(2)^200, n))
        as, xs
    end

    @testset "$op(::XInt, ::$T) / $op(::$T, ::XInt)" for
        op in (+, -, *, /, cmp, <, <=, >, >=, ==, !=, flipsign, gcd, gcdx, binomial,
               &, |, xor),
        T in [Base.BitInteger_types...; BigInt; Int256; UInt256;
              if op ∈ (gcd, gcdx, binomial, &, |, xor)
                  []
              else
                  [Base.uniontypes(CdoubleMax); BigFloat]
              end]
        as, xs = make_values(T)
        for a = as, y = xs, z = (-y, y), x = (vint(z),)
            op === binomial && !all(-1000 .<= [a, z] .<= 1000) && continue
            for (r, s) = ((op == binomial) ?
                          [(op(x, a), op(z, a))] : # TODO: remove special case
                          Any[(op(a, x), op(a, z)),
                              (op(x, a), op(z, a))])

                if op === flipsign
                    if r isa XInt
                        @test s isa BigInt
                    else
                        @test typeof(s) == typeof(r)
                    end
                elseif s isa BigInt
                    @test validate(r) isa XInt
                elseif op === gcdx
                    @test r isa Tuple{XInt,XInt,XInt}
                    validate.(r)
                else
                    @test typeof(r) == typeof(s)
                end
                @test isequal(r, s) || iszero(r) && iszero(s) && T == BigFloat
                # not `==` for NaN
                # TODO: fix the 0.0 vs -0.0 discrepancy with BigFloat?
            end
        end
    end

    @testset "powermod(::$T/XInt, ::$T/XInt, ::XInt)" for
            T in [Base.uniontypes(CulongMax); Base.uniontypes(ClongMax); [BigInt]]
        # TODO: add more combinations for first two parameters
        as1, xs1 = make_values(T, 2)
        as2, xs2 = make_values(T, 2)
        for a=as1, b=filter(>=(0), as2), x=filter(!iszero, xs1)
            s = powermod(a, b, x)
            r = powermod(a, b, vint(x))
            @test s == r
            @test validate(r) isa XInt
        end
        for a=as1, x=filter(>=(0), xs1), y=filter(!iszero, xs2)
            s = powermod(a, x, y)
            r = powermod(a, vint(x), vint(y))
            @test s == r
            @test validate(r) isa XInt
            r = powermod(vint(a), XInt(x), XInt(y))
            @test s == r
            @test validate(r) isa XInt
        end
    end

    @testset "checked operations" begin
        xs = XInt.([rand(1:1000, 4); rand(1:big(2)^65, 5)])
        xs = [xs; .-(xs)]
        fns = [getfield(Base, f[1]) => f[2] for f in
               [:checked_add => +, :checked_sub => -, :checked_mul => *,
                :checked_div => div, :checked_rem => rem, :checked_fld => fld,
                :checked_mod => mod, :checked_cld => cld]]

        for x=xs
            @test Base.checked_abs(x) == abs(x)
            @test Base.checked_neg(x) == -x
            for y=xs, (c, f) = fns
                @test f(x, y) == c(x, y)
            end
            for y=xs
                @test Base.add_with_overflow(x, y) == (x+y, false)
                @test Base.sub_with_overflow(x, y) == (x-y, false)
                @test Base.mul_with_overflow(x, y) == (x*y, false)
            end
        end
    end
end

@testset "bit and unary ops, etc." begin
    xs = BigInt[1, 2, 3, rand(UInt8, 5)..., rand(UInt, 5)..., rand(SLimb, 10)...,
                typemax(UInt), typemax(UInt)-1, typemax(UInt)-2,
                typemax(SLimb), typemax(SLimb)-1, typemax(SLimb)-2, typemax(SLimb)-3,
                big(2)^192-1] # specifically for rshift!

    append!(xs, rand(big(2)^65:big(2)^100, 5))
    append!(xs, rand(big(0):big(2)^200, 5))
    append!(xs, (-).(xs))
    push!(xs, 0, typemin(SLimb), typemin(SLimb) >> 1, big(typemin(SLimb))-1)

    @testset "$op(::XInt)" for op = (-, ~, isqrt, trailing_zeros, trailing_ones, count_ones,
                                     abs, factorial, com!)

        @assert xs isa Vector{BigInt} # so that mutating ops don't mutate xs themselves
        for x = xs
            op === factorial && !(0 <= x <= 1000) && continue
            s = try
                opmap(op)(x)
            catch
                @test_throws Union{InexactError,DomainError,DivideError} op(XInt(x))
                continue
            end
            r = op(vint(x))
            @test r == s
            if s isa BigInt
                @test validate(r) isa XInt
            else
                @assert s isa Int
                @test r isa Int
            end
        end
    end

    cs = [0, 1, 2, 3, 4, BITS_PER_LIMB, 2*BITS_PER_LIMB, rand(5:4*BITS_PER_LIMB, 20)...]
    cs2 = [cs; (-).(cs)]
    @testset "$op(::XInt, c)" for op = (<<, >>, >>>, (^), lshift!, rshift!)
        for T = [Base.BitInteger_types..., BigInt, Bool, XInt]
            for c = (T <: Signed && op !== (^) ? cs2 :
                     T === Bool ?                [true, false] :
                                                 cs)
                !Base.hastypemax(T) || typemin(T) <= c <= typemax(T) || continue
                for x = xs
                    s = opmap(op)(x, T(c))
                    @assert !(x isa XInt) # for mutation
                    r = op(vint(x), T(c))
                    @test s == r
                    validate(r)
                    if opmap(op) != op
                        # mutating
                        tmp = XInts._XInt(0, Limb[])
                        xx = vint(x)
                        rtmp = validate(op(tmp, xx, T(c)))
                        @test rtmp == r
                        if !is_short(r)
                            @test vec(rtmp) === vec(tmp) || vec(rtmp) === vec(xx)
                        end
                        tmp = XInts._XInt(0)
                        @test validate(op(tmp, vint(x), T(c))) == r
                    end
                end
            end
            if op === (^) && !Base.hastypemax(T) || sizeof(T) > sizeof(Culong)
                c = T(typemax(Culong)) + rand(1:1000)
                z = XInt(0)
                @test z^c === z
                @test_throws OverflowError XInt(2)^c
            end
        end
    end
end

@testset "hashing" begin
    xs = BigInt[rand(UInt, 10); rand(UInt128, 10); rand(1:big(2)^300, 5)]
    xs = [xs; .-(xs)]
    for x=xs
        h = rand(UInt)
        a = XInt(x)
        b = big(x)

        @test hash(b) == hash(a)
        @test hash(b, h) == hash(a, h)
        @test Base.hash_integer(b, h) == Base.hash_integer(a, h)

        for y=xs
            iszero(y) && continue
            a2 = XInt(y)
            b2 = big(y)
            @test hash(b//b2) == hash(a//a2)
            @test hash(b//b2, h) == hash(a//a2, h)
        end
    end
end

@testset "rand" begin
    xs = BigInt[[0, 1, 2]; rand(Int8, 5); rand(SLimb, 10); rand(Int128, 5);
                rand(big(1):big(2)^300, 5)]
    as = [xs; .-(xs)]
    rng = Random.MersenneTwister()

    for a=as, x=xs, b = (x, a+x)
        # tests (a big, b small) and vice-versa, and length(r) small even if a, b bigs
        r = XInt(a):XInt(b)
        if isempty(r)
            @test_throws ArgumentError rand(r)
        else
            @test validate(rand(r)) ∈ r
            sp = Random.Sampler(Random.MersenneTwister, r)
            z = XInt(0)
            @test validate(Random.rand!(rng, z, sp)) ∈ r
        end
    end
end


module CellNumbersTests

using Test
using Gridap

using ..CellValuesMocks

l = 10
a = 3.5

cn = TestIterCellValue(a,l)
@test isa(cn,CellNumber)
@test isa(cn,IterCellNumber)
r = [a for i in 1:l]
test_iter_cell_number(cn,r)
@test !isa(cn,CellArray)
@test !isa(cn,IterCellArray)
@test !isa(cn,CellMap)
@test !isa(cn,IterCellMap)

cn = TestIndexCellValue(a,l)
@test isa(cn,CellNumber)
@test isa(cn,IndexCellNumber)
test_index_cell_number(cn,r)
@test !isa(cn,CellArray)
@test !isa(cn,IndexCellArray)
@test !isa(cn,CellMap)
@test !isa(cn,IndexCellMap)

end # module


# Creation macros #
# ========================== #
"""
    @scenario(def)

Define a scenario type compatible with StochasticPrograms using the syntax
```julia
@scenario name = begin
    ...structdef...

    [@zero begin
        ...
        return zero(scenario)
    end]

    [@expectation begin
        ...
        return expected(scenarios)
     end]
end
```
The generated type is referenced through `name` and a default constructor is always generated. This constructor accepts the keyword `probability` to set the probability of the scenario occuring. Otherwise, any internal variables and specialized constructors are defined in the @scenario block as they would be in any Julia struct.

If possible, a `zero` method and an `expected` method will be generated for the defined type. Otherwise, or if the default implementation is not desired, these can be user provided through [`@zero`](@ref) and [`@expectation`](@ref).

The defined scenario type will be available on all Julia processes.

## Examples

The following defines a simple scenario ``ξ`` with a single value.

```jldoctest
@scenario ExampleScenario = begin
    ξ::Float64
end

ExampleScenario(1.0, probability = 0.5)

# output

ExampleScenario with probability 0.5
  ξ: 1.0

```

See also: [`@zero`](@ref), [`@expectation`](@ref), [`@sampler`](@ref)
"""
macro scenario(arg)
    @capture(arg, scenarioname_Symbol = scenariodef_) || error("Invalid syntax. Expected: scenarioname = begin scenariodef end")
    vars = Vector{Symbol}()
    vartypes = Vector{Union{Expr,Symbol}}()
    vardefs = Vector{Expr}()
    zerodefs = Vector{Expr}()
    expectdefs = Vector{Expr}()
    def = postwalk(prettify(scenariodef)) do x
        if @capture(x, var_Symbol::vartype_)
            push!(vars, var)
            push!(vartypes, vartype)
            push!(vardefs, x)
            return x
        end
        if @capture(x, @zero zerodef_)
            push!(zerodefs, zerodef)
            return @q begin end
        end
        if @capture(x, @expectation expectdef_)
            push!(expectdefs, expectdef)
            return @q begin end
        end
        return x
    end
    # Handle zero definition
    if length(zerodefs) > 1
        error("Only provide one zero implementation")
    end
    provided_zerodef, zerodef = if length(zerodefs) == 1
        zerodef = zerodefs[1]
        true, zerodef
    else
        false, Expr(:block)
    end
    # Handle expectation definition
    if length(expectdefs) > 1
        error("Only provide one expectation implementation")
    end
    provided_expectdef, expectdef = if length(expectdefs) == 1
        expectdef = expectdefs[1]
        expectdef = postwalk(prettify(expectdef)) do x
            @capture(x, constructor_Symbol(args__)) && constructor == scenarioname && return :(StochasticPrograms.ExpectedScenario($x))
            return x
        end
        true, expectdef
    else
        false, Expr(:block)
    end
    # Prepare automatic expectation definition
    combine = Expr(:call)
    push!(combine.args, scenarioname)
    push!(combine.args, :($(Expr(:parameters, :(probability = 1.0)))))
    for var in vars
        push!(combine.args, :(probability(s1)*s1.$var + probability(s2)*s2.$var))
    end
    # Prepare automatic scenariotext definition
    textdef = Expr(:block)
    for var in vars
        key = Meta.quot(var)
        push!(textdef.args, :(print(io, "\n  $($key): $(scenario.$var)")))
    end
    # Define scenario type
    code = @q begin
        if StochasticPrograms.supports_expected([$(vartypes...)], $provided_expectdef) || $provided_expectdef
            if StochasticPrograms.supports_zero([$(vartypes...)], $provided_zerodef) || $provided_zerodef
                struct $scenarioname <: AbstractScenario
                    probability::Probability
                    $def
                    function (::Type{$scenarioname})($(vardefs...); probability::AbstractFloat = 1.0)
                        return new(Probability(probability), $(vars...))
                    end
                end
                if $provided_zerodef
                    function Base.zero(::Type{$scenarioname})
                        $zerodef
                    end
                else
                    function Base.zero(::Type{$scenarioname})
                        return $scenarioname(zero.([$(vartypes...)])...; probability = 1.0)
                    end
                end
                if $provided_expectdef
                    function StochasticPrograms.expected(scenarios::Vector{$scenarioname})
                        isempty(scenarios) && return StochasticPrograms.ExpectedScenario(zero($scenarioname))
                        $expectdef
                    end
                else
                    function StochasticPrograms.expected(scenarios::Vector{$scenarioname})
                        isempty(scenarios) && return StochasticPrograms.ExpectedScenario(zero($scenarioname))
                        return StochasticPrograms.ExpectedScenario(reduce(scenarios) do s1, s2
                                                                   $combine
                                                                   end)
                    end
                end
                function StochasticPrograms.scenariotext(io::IO, scenario::$scenarioname)
                    $textdef
                    return io
                end
            else
                @warn "The scenario type $(string($(Meta.quot(scenarioname)))) was not defined. A user-provided implementation \n\n    function zero(::Type{{$(string($(Meta.quot(scenarioname))))})\n        ...\n    end\n\nis required."
            end
        else
            @warn "The scenario type $(string($(Meta.quot(scenarioname)))) was not defined. A user-provided implementation \n\n    function expected(scenarios::Vector{$(string($(Meta.quot(scenarioname))))})\n        ...\n    end\n\nis required."
        end
    end
    code = prettify(code)
    code = @q begin
        if (@__MODULE__) == $(esc(:Main))
            @everywhere begin
                $code
            end
        else
            $(esc(code))
        end
    end
    return prettify(code)
end
"""
    @zero(def)

Define the additive zero scenario inside a @scenario block using the syntax:
```julia
@zero begin
    ...
    return zero_scenario
end
```

## Examples

The following defines a zero scenario for the example scenario defined in [`@scenario`](@ref)

```julia
@zero begin
    return ExampleScenario(0.0)
end
```

See also [`@scenario`](@ref)
"""
macro zero(def) @warn "@zero should be used inside a @scenario block." end
"""
    @expectation(def)

Define how to form the expected scenario inside a @scenario block. The scenario collection is accessed through the reserved keyword `scenarios`.

```julia
@zero begin
    ...
    return zero_scenario
end
```

## Examples

The following defines expectation for the example scenario defined in [`@scenario`](@ref)

```julia
@expectation begin
    return ExampleScenario(sum([probability(s)*s.ξ for s in scenarios]))
end
```

See also [`@scenario`](@ref)
"""
macro expectation(def) @warn "@expectation should be used inside a @scenario block." end
"""
    @sampler(def)

Define a sampler type for some `scenariotype` compatible with StochasticPrograms using the syntax
```julia
@sampler samplername = begin
    ...internals...

    @sample scenariotype begin
        ...
        return scenario
    end
end
```
Any internal state required by the sampler, as well as any specialized constructor, are defined in the @sampler block as they would be in any Julia struct. Define the sample operation inside the [`@sample`](@ref) block and specify the `scenariotype` that the sampler returns. The defined sampler will be available on all Julia processes.

## Examples

The following defines a simple dummy sampler, with some internal weight value, for the scenario defined in [`@scenario`](@ref), and samples one scenario.

```jldoctest; setup = :(@scenario ExampleScenario = ξ::Float64), filter = r".*"
@sampler ExampleSampler = begin
    w::Float64

    ExampleSampler(w::AbstractFloat) = new(w)

    @sample ExampleScenario begin
        @parameters w
        return ExampleScenario(w*randn(), probability = rand())
    end
end
s = ExampleSampler(2.0)
s()

# output

ExampleScenario with probability 0.29
  ξ: 1.48

```

See also: [`@sample`](@ref), [`@scenario`](@ref)
"""
macro sampler(arg)
    @capture(prettify(arg), samplername_Symbol = samplerdef_) || error("Invalid syntax. Expected: sampler = begin samplerdef end")
    scenarioname = :undefined
    sampledefs = Vector{Expr}()
    def = postwalk(prettify(samplerdef)) do x
        if @capture(x, @parameters args__)
            code = Expr(:block)
            for param in args
                push!(code.args, :($param = sampler.$param))
            end
            return code
        elseif @capture(x, @sample sname_ sampledef_)
            scenarioname = sname
            push!(sampledefs, sampledef)
            return @q begin end
        else
            return x
        end
    end
    # Handle sample definition
    length(sampledefs) == 1 || error("Must provide exactly one @sample operation.")
    sampledef = sampledefs[1]
    # Define sampler type
    code = @q begin
        struct $samplername <: AbstractSampler{$scenarioname}
            $def
        end
        function (sampler::$samplername)()
            $sampledef
        end
    end
    code = prettify(code)
    code = @q begin
        if (@__MODULE__) == $(esc(:Main))
            @everywhere begin
                $code
            end
        else
            $(esc(code))
        end
    end
    return prettify(code)
end
"""
    @sample(def)

Define the sample operaton inside a @sampler block, using the syntax
```julia
@sample begin
    ...
    return sampled_scenario
end
```
The sampler object is referenced through the reserved keyword `sampler`, from which any internals can be accessed.
"""
macro sample(def) @warn "@sample should be used inside a @sampler block." end
"""
    @first_stage(def)

Add a first stage model generation recipe to `stochasticprogram` using the syntax
```julia
@first_stage stochasticprogram::StochasticProgram = begin
    ...
end [defer]
```
where JuMP syntax is used inside the block to define the first stage model. During definition, the first stage model is referenced through the reserved keyword `model`.

Optionally, give the keyword `defer` after the  to delay generation of the first stage model.

## Examples

The following defines the first stage model given by:

```math
  minimize 100x₁ + 150x₂
    s.t  x₁ + x₂ ≤ 120
         x₁ ≥ 40
         x₂ ≥ 20
```

```julia
@first_stage sp = begin
    @variable(model, x₁ >= 40)
    @variable(model, x₂ >= 20)
    @objective(model, Min, 100*x₁ + 150*x₂)
    @constraint(model, x₁ + x₂ <= 120)
end
```

See also: [`@second_stage`](@ref)
"""
macro first_stage(arg) esc(:(@first_stage $arg generate)) end
macro first_stage(arg, defer)
    return esc(:(@stage 1 $arg $defer))
end
"""
    @second_stage(def)

Add a second stage model generation recipe to `stochasticprogram` using the syntax
```julia
@second_stage stochasticprogram::StochasticProgram = begin
    @decision var1 var2 ...
    ...
end [defer]
```
where JuMP syntax is used inside the block to define the second stage model. During definition, the second stage model is referenced through the reserved keyword `model`.

Optionally, give the keyword `defer` after the  to delay generation of the first stage model.

## Examples

The following defines the second stage model given by:
```math
  minimize q₁(ξ)y₁ + q₂(ξ)y₂
    s.t  6y₁ + 10y₂ ≤ 60x₁
         8y₁ + 5y₂ ≤ 60x₂
         0 ≤ y₁ ≤ d₁(ξ)
         0 ≤ y₂ ≤ d₂(ξ)
```
where ``q₁(ξ), q₂(ξ), d₁(ξ), d₂(ξ)`` depend on the scenario ``ξ`` and ``x₁, x₂`` are first stage variables. Two scenarios are added so that two second stage models are generated.

```julia
@second_stage sp = begin
    @decision x₁ x₂
    @uncertain q₁ q₂ d₁ d₂
    @variable(model, 0 <= y₁ <= d₁)
    @variable(model, 0 <= y₂ <= d₂)
    @objective(model, Min, q₁*y₁ + q₂*y₂)
    @constraint(model, 6*y₁ + 10*y₂ <= 60*x₁)
    @constraint(model, 8*y₁ + 5*y₂ <= 80*x₂)
end
```

See also: [`@first_stage`](@ref)
"""
macro second_stage(arg) esc(:(@second_stage $arg generate)) end
macro second_stage(arg, defer)
    return esc(:(@stage 2 $arg $defer))
end
"""
    @parameters(def)

Define the problem parameters in a @stage block
```julia
@parameters param1, param2, ...
```
possibly with default values. Any defined parameter without a default value must be supplied as a keyword argument to [`instantiate`](@ref) or [`SAA`](@ref) when creating models.

## Examples

```julia
@parameters d

@parameters begin
    Crops = [:wheat, :corn, :beets]
    Cost = Dict(:wheat=>150, :corn=>230, :beets=>260)
    Budget = 500
end
```

See also [`@decision`](@ref), [`@uncertain`](@ref), [`@stage`](@ref)
"""
macro parameters(def) @warn "@parameters should be used inside a @stage block." end
"""
    @decision(def)

In a @stage block, annotate each decision taken in the previous stage using the syntax
```julia
@decision var1, var2, ...
```

## Examples

```julia
@decision x₁, x₂
```

See also [`@parameters`](@ref), [`@uncertain`](@ref), [`@stage`](@ref)
"""
macro decision(def) @warn "@decision should be used inside a @stage block." end
"""
    @uncertain(def)

In a @stage block, annotate each uncertain variable using the syntax
```julia
@uncertain var1, var2, ...
```
This assumes that the [`Scenario`] type is used. Alternatively, user-defined scenarios can be specified by annotating the type. Also, inside a @stochastic_model block, user-defined scenarios can be created during the @uncertain annotation, following [`@scenario`](@ref).

## Examples

```julia
@uncertain q₁ q₂ d₁ d₂

@scenario Simple = begin
    q₁::Float64
    q₂::Float64
    d₁::Float64
    d₂::Float64
end
@uncertain ξ::SimpleScenario

@stochastic_model begin
    ...
    @uncertain ξ::SimpleScenario = begin
        q₁::Float64
        q₂::Float64
        d₁::Float64
        d₂::Float64
    end
    ...
end
```

See also [`@scenario`](@ref), [`@parameters`](@ref), [`@decision`](@ref), [`@stage`](@ref)
"""
macro uncertain(def) @warn "@uncertain should be used inside a @stage block." end
"""
    @stage(def)

Add a stage model generation recipe to `stochasticprogram` using the syntax
```julia
@stage stage stochasticprogram::StochasticProgram = begin
    @parameters param1 param2 ...
    @decision var1 var2 ...
    @uncertain ξ
    ... JuMPdef ...
    ...
end [defer]
```
where JuMP syntax is used inside the block to define the stage model. During definition, the second stage model is referenced through the reserved keyword `model`.

Optionally, give the keyword `defer` after the  to delay generation of the stage model.

## Examples

The following defines the first stage model given by:
```math
  minimize 100x₁ + 150x₂
    s.t  x₁ + x₂ ≤ 120
         x₁ ≥ 40
         x₂ ≥ 20
```
and the second-stage model given by:
```math
  minimize q₁(ξ)y₁ + q₂(ξ)y₂
    s.t  6y₁ + 10y₂ ≤ 60x₁
         8y₁ + 5y₂ ≤ 60x₂
         0 ≤ y₁ ≤ d₁(ξ)
         0 ≤ y₂ ≤ d₂(ξ)
```
where ``q₁(ξ), q₂(ξ), d₁(ξ), d₂(ξ)`` depend on the scenario ``ξ`` and ``x₁, x₂`` are first stage variables. Two scenarios are added so that two second stage models are generated.

```jldoctest
ξ₁ = Scenario(q₁ = -24.0, q₂ = -28.0, d₁ = 500.0, d₂ = 100.0, probability = 0.4)
ξ₂ = Scenario(q₁ = -28.0, q₂ = -32.0, d₁ = 300.0, d₂ = 300.0, probability = 0.6)

sp = StochasticProgram([ξ₁, ξ₂])

@stage 1 sp = begin
    @variable(model, x₁ >= 40)
    @variable(model, x₂ >= 20)
    @objective(model, Min, 100*x₁ + 150*x₂)
    @constraint(model, x₁ + x₂ <= 120)
end

@stage 2 sp = begin
    @decision x₁ x₂
    @uncertain q₁ q₂ d₁ d₂
    @variable(model, 0 <= y₁ <= d₁)
    @variable(model, 0 <= y₂ <= d₂)
    @objective(model, Min, q₁*y₁ + q₂*y₂)
    @constraint(model, 6*y₁ + 10*y₂ <= 60*x₁)
    @constraint(model, 8*y₁ + 5*y₂ <= 80*x₂)
end

# output

Stochastic program with:
 * 2 decision variables
 * 2 recourse variables
 * 2 scenarios of type Scenario
Solver is default solver

```

See also: [`@parameters`](@ref), [`@decision`](@ref), [`@uncertain`](@ref)
"""
macro stage(stage, arg) esc(:(@stage $stage $arg generate)) end
macro stage(stage, args, defer)
    generate = if defer == :defer
        :false
    elseif defer == :generate
        :true
    else
        error("Unknown option $defer")
    end
    @capture(args, sp_Symbol = def_) || error("Invalid syntax. Expected stage, multistage = begin JuMPdef end")
    # Save variable definitions separately
    vardefs = Expr(:block)
    for line in block(def).args
        if @capture(line, @constraint(m_Symbol, constdef__)) || @capture(line, @objective(m_Symbol, objdef__)) || @capture(line, @parameters args__) || @capture(line, @decision args__) || @capture(line, @uncertain args__)
            continue
        else
            push!(vardefs.args, line)
        end
    end
    # Handle parameters
    def = postwalk(def) do x
        if @capture(x, @parameters arg_)
            code = Expr(:block)
            for paramdef in prettify(arg).args
                if @capture(paramdef, key_Symbol = val_) || @capture(paramdef, key_Symbol)
                    push!(code.args, :($key = stage.$key))
                else
                    error("Incorrect @parameters syntax. Specify parameter names, possibly with a default value.")
                end
            end
            pushfirst!(vardefs.args, code)
            return code
        elseif @capture(x, @parameters args__)
            code = Expr(:block)
            for param in args
                push!(code.args, :($param = stage.$param))
            end
            pushfirst!(vardefs.args, code)
            return code
        elseif @capture(x, @decision args__)
            stage == 1 && error("@decision links cannot be used in the first stage.")
            code = Expr(:block)
            for var in args
                varkey = Meta.quot(var)
                decisiondef = @q begin
                    haskey(parent.objDict, $varkey) || error("Decision $($varkey) not present in parent model.")
                    $var = parent.objDict[$varkey]
                end
                push!(code.args, decisiondef)
            end
            return code
        elseif @capture(x, @uncertain var_Symbol::t_Symbol)
            stage == 1 && error("@uncertain declarations cannot be used in the first stage.")
            return :($var::$t = scenario)
        elseif @capture(x, @uncertain vars__ from scenvar_Symbol::t_Symbol)
            stage == 1 && error("@uncertain declarations cannot be used in the first stage.")
            code = @q begin
                $scenvar::$t = scenario
            end
            for var in vars
                varkey = Meta.quot(var)
                push!(code.args, :($var = getproperty($scenvar, $varkey)))
            end
            return code
        elseif @capture(x, @uncertain vars__ from t_Symbol)
            stage == 1 && error("@uncertain declarations cannot be used in the first stage.")
            code = @q begin
                scenvar::$t = scenario
            end
            for var in vars
                varkey = Meta.quot(var)
                push!(code.args, :($var = getproperty(scenvar, $varkey)))
            end
            return code
        elseif @capture(x, @uncertain var_Symbol::t_Symbol = scenariodef_)
            stage == 1 && error("@uncertain declarations cannot be used in the first stage.")
            return @q begin
                $var::$t = scenario
            end
        elseif @capture(x, @uncertain vars__ from scenvar_Symbol::t_Symbol = scenariodef_)
            stage == 1 && error("@uncertain declarations cannot be used in the first stage.")
            code = @q begin
                $scenvar::$t = scenario
            end
            for var in vars
                varkey = Meta.quot(var)
                push!(code.args, :($var = getproperty($scenvar, $varkey)))
            end
            return code
        elseif @capture(x, @uncertain vars__ from t_Symbol = scenariodef_)
            stage == 1 && error("@uncertain declarations cannot be used in the first stage.")
            code = @q begin
                scenvar::$t = scenario
            end
            for var in vars
                varkey = Meta.quot(var)
                push!(code.args, :($var = getproperty(scenvar, $varkey)))
            end
            return code
        elseif @capture(x, @uncertain args__)
            stage == 1 && error("@uncertain declarations cannot be used in the first stage.")
            code = Expr(:block)
            for var in args
                varkey = Meta.quot(var)
                push!(code.args, :($var = scenario[$varkey]))
            end
            return code
        else
            return x
        end
    end
    # The first stage is treated different from the subsequent stages
    generatordefs = if stage == 1
        @q begin
            $(esc(sp)).generator[:stage_1] = ($(esc(:model))::JuMP.Model, $(esc(:stage))) -> begin
                $(esc(def))
	        return $(esc(:model))
            end
        end
    else
        @q begin
            $(esc(sp)).generator[Symbol(:stage_,$stage)] = ($(esc(:model))::JuMP.Model, $(esc(:stage)), $(esc(:scenario))::AbstractScenario, $(esc(:parent))::JuMP.Model) -> begin
                $(esc(def))
	        return $(esc(:model))
            end
        end
    end
    # Create definition code
    code = @q begin
        isa($(esc(sp)), StochasticProgram) || error("Given object is not a stochastic program.")
        $(esc(sp)).generator[Symbol(:stage_,$stage,:_vars)] = ($(esc(:model))::JuMP.Model, $(esc(:stage))) -> begin
            $(esc(vardefs))
	    return $(esc(:model))
        end
        $generatordefs
        if $generate
            generate_stage!($(esc(sp)), $stage)
        end
        $(esc(sp))
    end
    # Return code
    return prettify(code)
end
# ========================== #


using SQLQuery

type NewNode{T} <: SQLQuery.QueryNode
    input::T
    args
end

SQLQuery.QUERYNODE[:newnode] = NewNode

SQLQuery.translatesql(nn::NewNode, offset::Int) = "newnode, offset=$offset"

@sqlquery source |>
newnode(are, you.serious) |>
select(*, some, columns)

@sqlquery source |>
filter(name == 3, bar == "whee")

@sqlquery source |>
filter(name == 3, bar == "whee") |>
select(name = foo * 3, col)

@sqlquery source |>
filter(name == 3, bar == "whee") |>
select(distinct(name = foo * 3, col))

@sqlquery source |>
filter(name == 3, bar == "whee") |>
select(all(name = foo * 3, col))

@sqlquery source |>
filter(name == 3, bar == "whee") |>
select(all(*))

@sqlquery source |>
filter(name == 3, bar == "whee") |>
select(distinct(*))

@sqlquery source |>
filter(name == 3, bar == "whee") |>
select(distinct(col))

@sqlquery source |>
filter(name == 3, bar == "whee") |>
select(name = foo * 3, col) |>
orderby(name)

@sqlquery source |>
filter(name == 3, bar == "whee") |>
select(name = foo * 3, col) |>
orderby(name, col)

@sqlquery source |>
filter(name == 3, bar == "whee") |>
select(name = foo * 3, col) |>
orderby(desc(name))

@sqlquery source |>
filter(name == 3, bar == "whee") |>
select(name = foo * 3, col) |>
orderby(name, desc(col))

@sqlquery source |>
filter(name == 3, bar == "whee") |>
select(name = foo * 3, col) |>
orderby(desc(name), desc(col))

@sqlquery source |>
filter(name == 3, bar == "whee") |>
select(name = foo * 3, col) |>
orderby(desc(name), asc(col))

@sqlquery source |>
filter(name == 3, bar == "whee") |>
select(name = foo * 3, col) |>
orderby(asc(name), asc(col))

@sqlquery source |>
filter(name == 3, bar == "whee") |>
select(name = foo * 3, col) |>
orderby(desc(name))

@sqlquery source |>
filter(name == 3, bar == "whee") |>
select(name = foo * 3, col) |>
orderby(desc(name)) |>
limit(10)

@sqlquery source |>
filter(name == 3, bar == "whee") |>
select(name = foo * 3, col) |>
orderby(desc(name)) |>
offset(7) |>
filter(name == 3, bar == "whee") |>
select(name = foo * 3, col) |>
orderby(desc(name)) |>
offset(7) |>
filter(name == 3, bar == "whee") |>
select(name = foo * 3, col) |>
orderby(desc(name)) |>
offset(7) |>
filter(name == 3, bar == "whee") |>
select(name = foo * 3, col) |>
orderby(desc(name)) |>
offset(7) |>
filter(name == 3, bar == "whee") |>
select(name = foo * 3, col) |>
orderby(desc(name)) |>
offset(7) |>
filter(name == 3, bar == "whee") |>
select(name = foo * 3, col) |>
orderby(desc(name)) |>
offset(7) |>
filter(name == 3, bar == "whee") |>
select(name = foo * 3, col) |>
orderby(desc(name)) |>
offset(7) |>
filter(name == 3, bar == "whee") |>
select(name = foo * 3, col) |>
orderby(desc(name)) |>
offset(7)

@sqlquery Artists |>
leftjoin(Songs) |>
leftjoin(Albums) |>
select( song_id = Songs._id,
song_name = Songs.name,
Songs.length,
artist_id = Songs.artist_id,
artist_name = Artists.name,
album_id = Songs.album_id,
album_name = Albums.name)

@sqlquery Artists |>
leftjoin(Songs) |>
leftjoin(Albums) |>
select( song_id = Songs._id,
song_name = Songs.name,
Songs.length,
artist_id = Songs.artist_id,
artist_name = Artists.name,
album_id = Songs.album_id,
album_name = Albums.name) |>
groupby(Songs.length, artist_id)

@sqlquery Artists |>
leftjoin(Songs) |>
leftjoin(Albums) |>
select( song_id = Songs._id,
song_name = Songs.name,
Songs.length,
artist_id = Songs.artist_id,
artist_name = Artists.name,
album_id = Songs.album_id,
album_name = Albums.name) |>
groupby(Songs.length, artist_id) |>
orderby(asc(song_id), desc(Songs.length), desc(artist_id))


function to_pari(io::IO, f::Generic.Poly{nf_elem})
  Zy, y = PolynomialRing(FlintQQ, "y", cached = false)
  k = base_ring(f)
  kt = parent(f)
  s_k = k.S
  s_kt = kt.S
  #important to pari: the top poly needs to have x, so the field should be y...

  k.S = :y
  kt.S = :x

  print(io, "k = nfinit($(k.pol(y)));\n")
  print(io, "f = $f;\n")

  k.S = s_k
  kt.S = s_kt
  nothing
end

function to_pari(s::String, f::Generic.Poly{nf_elem})
  g = open(s, "w")
  to_pari(g, f)
  close(g)
end

function to_pari(f::IO, m::fmpz_mat)
  print(f, "A = [")
  for i=1:nrows(A)
    if i>1
      print(f, "; ")
    end
    for j=1:ncols(A)-1
      print(f, A[i,j], ",")
    end; print(f, A[i, ncols(A)])
  end
  print(f, "];")
end

function to_pari(s::String, m::fmpz_mat)
  g = open(s, "w")
  to_pari(g, m)
  close(g)
end

# gp> #
# gp> read("bla");
# gp> nffactor(k, f);
#=

from
  https://www.math.u-bordeaux.fr/~kbelabas/research/vanhoeij.pdf
  p 26, polys h_1, ..., h_10 (so far)

Zx, x = PolynomialRing(FlintZZ)
l = [x-1,
x^2-19,
x^3-x^2+6,
x^4 - 2*x^3 - 7*x^2 - 7*x + 3,
x^5 - 3*x^3 - 2*x^2+8*x-10,
x^6 - 5*x^4 - x^3 + 10*x^2 - 11*x + 5,
x^7-x^6-32*x^5-100*x^4-130*x^3-70*x^2+ 4*x+ 23,
x^8-8*x^6-x^5+x^3+ 3*x^2-7*x+ 2,
x^9-x^8-42*x^7-180*x^6-375*x^5-452*x^4-327*x^3-133*x^2-31*x-11,
x^10-2*x^9+ 5*x^8+ 8*x^7+ 5*x^6+ 8*x^5-9*x^4-6*x^3+ 7*x^2+ 4*x+ 8,
]

g1 = swinnerton_dyer(7, x);
g2 = swinnerton_dyer(8, x);

for f = l
  k, a = number_field(f)
  kt, t = k["t"]

  g = (g1*g2)(t-a)
  @time length(factor(g))
end

needs prec_scale of 40 rather than 20. Why I don't know (yet)

set_verbose_level(:PolyFactor, 1)

=#


cd("$(homedir())/Documents/GitHub")

# Push Documents/GitHub on the search path
push!(LOAD_PATH, "$(homedir())/Documents/GitHub")

using FileIO, ExcelFiles, DataFrames
#using ExcelReaders
#using MAT
using NEPA

#mat = matread("C:/Users/u0093191/Documents/MATLAB/imputeLCM.mat")
#println(mat)

X = DataFrame(load("C:/Users/u0093191/Documents/MATLAB/WACM/Atkinson-RSanalyis-11May2011.xls","Data!F4:H195"))
W = DataFrame(load("C:/Users/u0093191/Documents/MATLAB/WACM/Atkinson-RSanalyis-11May2011.xls","Data!J4:L195"))
Y = DataFrame(load("C:/Users/u0093191/Documents/MATLAB/WACM/Atkinson-RSanalyis-11May2011.xls","Data!E4:E195"))

#X = readxl("C:/Users/u0093191/Documents/MATLAB/WACM/Atkinson-RSanalyis-11May2011.xls","Data!F4:H195")
#W = readxl("C:/Users/u0093191/Documents/MATLAB/WACM/Atkinson-RSanalyis-11May2011.xls","Data!J4:L195")
#Y = readxl("C:/Users/u0093191/Documents/MATLAB/WACM/Atkinson-RSanalyis-11May2011.xls","Data!E4:E195")

# Convert DataFrame or DataArray to Array so that the  function WACM doesn't crash
X = convert(Array{Float64},X)
W = convert(Array{Float64},W)
Y = convert(Array{Float64},Y)

# Input-oriented efficiencies computed using Maple. We use these to check our computed results
#MapleEff = readxl("C:/Users/u0093191/Documents/MATLAB/WACM/Atkinson-RSanalyis-11May2011.xls","MapleEfficiencies!B2:O193")
MapleEff = DataFrame(load("C:/Users/u0093191/Documents/MATLAB/WACM/Atkinson-RSanalyis-11May2011.xls","MapleEfficiencies!B2:O193"))
MapleEff = convert(Array{Float64},MapleEff)

#eff = WACM(W,X,Y)
#Time execution using @time
#@time eff = WACM(w,x,y)
#@code_warntype WACM(w,x,y)
#eff = WACM(w,x,y)
#print(eff)

#Input- or output-oriented?
input = true

D = FDH_VRS(X,Y,input)
println(maximum(abs(geteff(D()) - MapleEff[:,8])))
D = FDH_CRS(X,Y,input)
println(maximum(abs(geteff(D()) - MapleEff[:,9])))
D = FDH_NIRS(X,Y,input)
println(maximum(abs(geteff(D()) - MapleEff[:,10])))
D = FDH_NDRS(X,Y,input)
println(maximum(abs(geteff(D()) - MapleEff[:,11])))

#M = convert(Int,round(3*size(D,1)/4))
#thetaM = orderm(D,M,100)
#println(thetaM - D())

# This should be equal to the FDH_VRS efficiency!
D = DDF{FreeDisposal,VRS}(X,Y,X,zeros(size(Y)))
println(maximum(abs((1 - geteff(D())) - MapleEff[:,8])))
#println(hcat((1 - geteff(D())), MapleEff[:,8]))

# Hyperbolic equals square root of input oriented FDH under CRS
GDR = Hyperbolic{FreeDisposal,CRS}(X,Y)
D = FDH_CRS(X,Y,input)
println(maximum(abs(geteff(GDR()) - sqrt(geteff(D())))))

D = DEA_VRS(X,Y,input)
println(maximum(abs(geteff(D()) - MapleEff[:,1])))
D = DEA_CRS(X,Y,input)
println(maximum(abs(geteff(D()) - MapleEff[:,2])))
D = DEA_NIRS(X,Y,input)
println(maximum(abs(geteff(D()) - MapleEff[:,3])))
D = DEA_NDRS(X,Y,input)
println(maximum(abs(geteff(D()) - MapleEff[:,4])))
#@time theta = D() #4.389059 seconds
#@time theta = D() #0.123012 seconds

D = DDF{Convex,VRS}(X,Y,X,zeros(size(Y)))
println(maximum(abs((1 - geteff(D())) - MapleEff[:,1])))

# Test indexing of DDF object
println(D[1])
println(X[1,:],Y[1,:],X[1,:],zeros(1,size(Y,2)))

# Do order-m efficiency
#M = convert(Int,round(0.9*size(D,1)))
#thetaM = orderm(D,M,100)
#println(thetaM - geteff(D()))

#data = readdlm("./NEPA/data/GriffellTatjéLovell.txt")
#X0 = data[:,1]
#Y0 = data[:,2]
#X1 = data[:,3]
#Y1 = data[:,4]

#LTFP = Luenberger{Convex,VRS}(X,Y,X,Y,X,Y,X,Y)
#println(LTFP())
#println(TEI(LTFP))
#println(TC(LTFP))

# HMBTFP = HMB{Convex,VRS}(X,Y,X,Y)
# println(HMBTFP())
# println(TEI(HMBTFP))
# println(TC(HMBTFP))
# println(SEC(HMBTFP))


using DynamicGrids, Test, Adapt
using DynamicGrids: SimData

output = ArrayOutput(BitArray(rand(Bool, 10, 10));
    tspan=1:10,
    mask=BitArray(rand(Bool, 10, 10)),
    aux=(aux1=BitArray(rand(Bool, 10, 10)),),
)

rs1 = Ruleset(Life(); opt=NoOpt())
rs2 = Ruleset(Life(); opt=SparseOpt())
sd1 = SimData(output, rs1)
sd2 = SimData(output, rs2)
Adapt.adapt(Array, rs1)
Adapt.adapt(Array, rs2)
b_sd1 = Adapt.adapt(Array, sd1)
b_sd2 = Adapt.adapt(Array, sd2)
@test b_sd2.extent.init._default_ isa Array
@test b_sd2.extent.mask isa Array
@test b_sd2.extent.aux.aux1 isa Array
@test parent(b_sd2.grids[:_default_].source) isa Array
@test parent(b_sd2.grids[:_default_].dest) isa Array
@test parent(b_sd2.grids[:_default_].optdata.sourcestatus) isa Array
@test Adapt.adapt(Array, output)[1] isa Array


abstract type AbstractVar end
abstract type AbstractDomain end

abstract type AbstractSetVar <: AbstractVar end

"""
    abstract type AbstractIntVar end

Abstract type for integer variables.
"""
abstract type AbstractIntVar <: AbstractVar end

"""
    abstract type AbstractBoolVar <: AbstractVar end

Abstract type for all boolean variables.
"""
abstract type AbstractBoolVar <: AbstractVar end

"""
    function id(x::AbstractVar)

Return the `string` identifier of `x`. Every variable must be assigned a unique identifier upon creation, that
will be used as a key to identify the variable in the `CPModel` object.
"""
id(x::AbstractVar) = x.id
parentId(x::AbstractVar) = rootVariable(x).id

include("IntDomain.jl")
include("IntVar.jl")

include("BoolDomain.jl")
include("BoolVar.jl")

abstract type IntVarView <: AbstractIntVar end
abstract type IntDomainView <: AbstractIntDomain end

include("IntVarView.jl")

abstract type BoolVarView <: AbstractBoolVar end
abstract type BoolDomainView <: AbstractBoolDomain end

include("BoolVarView.jl")

include("IntSetDomain.jl")
include("IntSetVar.jl")


include("exp_set_list_all.jl");

println("Printing simulation");

rho=mono.params[:rho];
println("Domain: $(rho), m=",m);
for (i,state)=enumerate(state_list)
    if (!isnothing(state))
        err=showerr(state,output=false,n=state.params[:target_n]*2)
        graphname=state.params[:graphname];
        println("$graphname $(Float64(err))");
    else
        println("x");
    end
end


module Day12

using ..Utils

const CaveSystem = Dict{String, Vector{String}}

function parse_cave_system(input::String)
    system = Dict{String, Vector{String}}()
    for line in split(input, "\n")
        if !isempty(line)
            mergeline!(system, line)
        end
    end
    system
end

function mergeline!(system::CaveSystem, line::AbstractString)
    a, b = split(line, "-")
    if a ∉ keys(system)
        push!(system, a => String[])
    end
    if b ∉ keys(system)
        push!(system, b => String[])
    end
    push!(system[a], b)
    push!(system[b], a)
    system
end

part1(system::CaveSystem) = paths_starting_at("start", system, ["start"]; used_double_visit=true)
part2(system::CaveSystem) = paths_starting_at("start", system, ["start"], used_double_visit=false)

function paths_starting_at(cave::String, system::CaveSystem, path_so_far::Vector{String}; used_double_visit=false)
    result = Vector{String}[]
    neighbors = system[cave]
    for neighbor in neighbors
        path = copy(path_so_far)
        push!(path, neighbor)
        if neighbor == "end"
            push!(result, path)
        elseif neighbor == "start"
            # No-op
        elseif isuppercase(neighbor[1]) # Don't worry about deduplicating big caves
            append!(result, paths_starting_at(neighbor, system, path, used_double_visit=used_double_visit))
        else
            # Lowercase cave
            visit_counts = count(s -> s == neighbor, path)

            if visit_counts == 1
                append!(result, paths_starting_at(neighbor, system, path, used_double_visit=used_double_visit))
            elseif !used_double_visit && visit_counts <= 2
                append!(result, paths_starting_at(neighbor, system, path, used_double_visit=true))
            end
        end
    end
    result
end

function run()
    test1 = """start-A
start-b
A-c
A-b
b-d
A-end
b-end
"""

    test2 = """dc-end
HN-start
start-kj
dc-start
dc-HN
LN-dc
HN-end
kj-sa
kj-HN
kj-dc"""

    test3 = """fs-end
he-DX
fs-he
start-DX
pj-DX
end-zg
zg-sl
zg-pj
pj-he
RW-he
fs-DX
pj-RW
zg-RW
start-pj
he-WI
zg-he
pj-fs
start-RW"""

    @assert test1 |> parse_cave_system |> part1 |> length == 10
    @assert test2 |> parse_cave_system |> part1 |> length == 19
    @assert test3 |> parse_cave_system |> part1 |> length == 226

    @assert test1 |> parse_cave_system |> part2 |> length == 36
    @assert test2 |> parse_cave_system |> part2 |> length == 103
    @assert test3 |> parse_cave_system |> part2 |> length == 3509

    input = read("resources/day12_input.txt", String)
    @show input |> parse_cave_system |> part1 |> length
    @show input |> parse_cave_system |> part2 |> length

end

function test_part2(test)
    part2_result = """
start,A,b,A,b,A,c,A,end
start,A,b,A,b,A,end
start,A,b,A,b,end
start,A,b,A,c,A,b,A,end
start,A,b,A,c,A,b,end
start,A,b,A,c,A,c,A,end
start,A,b,A,c,A,end
start,A,b,A,end
start,A,b,d,b,A,c,A,end
start,A,b,d,b,A,end
start,A,b,d,b,end
start,A,b,end
start,A,c,A,b,A,b,A,end
start,A,c,A,b,A,b,end
start,A,c,A,b,A,c,A,end
start,A,c,A,b,A,end
start,A,c,A,b,d,b,A,end
start,A,c,A,b,d,b,end
start,A,c,A,b,end
start,A,c,A,c,A,b,A,end
start,A,c,A,c,A,b,end
start,A,c,A,c,A,end
start,A,c,A,end
start,A,end
start,b,A,b,A,c,A,end
start,b,A,b,A,end
start,b,A,b,end
start,b,A,c,A,b,A,end
start,b,A,c,A,b,end
start,b,A,c,A,c,A,end
start,b,A,c,A,end
start,b,A,end
start,b,d,b,A,c,A,end
start,b,d,b,A,end
start,b,d,b,end
start,b,end"""

    part2_expected_paths = filter(!isempty, split(part2_result, "\n"))
    part2_actual_paths = [join(l, ",") for l in part2(parse_cave_system(test))]

    display(sort(part2_actual_paths))
    display(sort(part2_expected_paths))
end

end


# ------------ Call sites ------------ #

@inline function (ctx::ScoreContext)(c::typeof(markov),
                                     addr::Address,
                                     call::Function,
                                     len::Int,
                                     args...)
    visit!(ctx, addr => 1)
    ss = get_subselection(ctx, addr)
    ret, w = score(get_sub(ss, 1), call, args...)
    v_ret = Vector{typeof(ret)}(undef, len)
    v_ret[1] = ret
    increment!(ctx, w)
    for i in 2:len
        visit!(ctx, addr => i)
        ret, w = score(get_sub(ss, i), call, v_ret[i-1]...)
        v_ret[i] = ret
        increment!(ctx, w)
    end
    return v_ret
end


## CATEGORICAL ARRAY DECODER UTILITY

"""
    CategoricalDecoder(C::CategoricalArray; eltype=nothing)

Construct a decoder for transforming a `CategoricalArray{T}` object
into an ordinary array, and for re-encoding similar arrays back into a
`CategoricalArray{T}` object having the same `pool` (and, in
particular, the same levels) as `C`. If `eltype` is not specified then
the element type of the transformed array is `T`. Otherwise, the
element type is `eltype` and the elements are promotions of the
internal (integer) `ref`s of the `CategoricalArray`. One
must have `R <: eltype <: Real` where `R` is the reference type of the
`CategoricalArray` (usually `UInt32`).

    transform(decoder::CategoricalDecoder, C::CategoricalArray)

Transform `C` into an ordinary `Array`.

    inverse_transform(decoder::CategoricalDecoder, A::Array)

Transform an array `A` suitably compatible with `decoder` into a
`CategoricalArray` having the same `pool` as `C`.

### Example

````
julia> using CategoricalArrays
julia> C = categorical(["a" "b"; "a" "c"])
2×2 CategoricalArray{String,2,UInt32}:
 "a"  "b"
 "a"  "c"

julia> decoder = MLJBase.CategoricalDecoder(C, eltype=Float64);
julia> A = transform(decoder, C)
2×2 Array{Float64,2}:
 1.0  2.0
 1.0  3.0

julia> inverse_transform(decoder, A[1:1,:])
1×2 CategoricalArray{String,2,UInt32}:
 "a"  "b"

julia> levels(ans)
3-element Array{String,1}:
 "a"
 "b"
 "c"

"""
struct CategoricalDecoder{I<:Real,T,N,R<:Integer}  # I the output eltype
    pool::CategoricalPool{T,R} # abstract type, not optimal
    use_original_type::Bool
    CategoricalDecoder{I,T,N,R}(X::CategoricalArray{T,N,R}, use_original_type) where {I,T,N,R}  =
        new(X.pool, use_original_type)
end

function CategoricalDecoder(X::CategoricalArray{T,N,R}; eltype=nothing) where {T,N,R}
    if eltype ==  nothing
        eltype = R # any integer type will do here; not used
        use_original_type = true
    else
        use_original_type = false
    end
    return CategoricalDecoder{eltype,T,N,R}(X, use_original_type)
end

function transform(decoder::CategoricalDecoder{I,T,N,R}, C::CategoricalArray) where {I,T,N,R}
    if decoder.use_original_type
        return collect(C)
    else
        return broadcast(C.refs) do element
            ref = convert(I, element)
        end
    end
end

function inverse_transform(decoder::CategoricalDecoder{I,T,N,R}, A::Array{J}) where {I,T,N,R,J<:Union{I,T}}
    if decoder.use_original_type
        refs = broadcast(A) do element
            decoder.pool.invindex[element]
        end
    else
        refs = broadcast(A) do element
            round(R, element)
        end
    end
    return CategoricalArray{T,N}(refs, decoder.pool)

end

## UTILITY FOR CONVERTING TABULAR DATA INTO MATRIX

""""
    MLJBase.matrix(X)

Convert an iteratable table source `X` into an `Matrix`; or, if `X` is
a `Matrix`, return `X`.

"""
function matrix(X)
    TableTraits.isiterabletable(X) || error("Argument is not an iterable table.")

    df = @from row in X begin
        @select row
        @collect DataFrames.DataFrame
    end
    return convert(Matrix, df)

end

matrix(X::Matrix) = X

## TOOLS FOR INDEXING QUERYVERSE ITERABLE TABLES

struct Rows end
struct Cols end
struct Schema
    nrows::Int
    ncols::Int
    names::Vector{Symbol}
    eltypes::Vector{DataType}
end

# for accessing tabular data:
retrieve(X, args...) = retrieve(Val(TableTraits.isiterabletable(X)), X, args...)
retrieve(::Val{false}, X, args...) = error("Argument is not an iterable table.")

# Note: to `retrieve` from matrices and other non-iterabletable data,
# we overload the second method above.

# project named tuple onto a tuple with only specified `labels` or indices:
project(t::NamedTuple, labels::AbstractArray{Symbol}) =
    NamedTuple{tuple(labels...)}(tuple([getproperty(t, n) for n in labels]...))
project(t::NamedTuple, label::Symbol) = project(t, [label,])
project(t::NamedTuple, indices::AbstractArray{<:Integer}) =
    NamedTuple{tuple(keys(t)[indices]...)}(tuple([t[i] for i in indices]...))
project(t::NamedTuple, i::Integer) = project(t, [i,])

# to select columns `c` of any tabular data `X` with `retrieve(X, Cols, c)`:
# CURRENTLY RETURNS A DATAFRAME AND NOT THE ORIGINAL TYPE
function retrieve(::Val{true}, X::T, ::Type{Cols}, c::AbstractArray{I}) where {T,I<:Union{Symbol,Integer}}

    row_iterator = @from row in X begin
        @select project(row, c)
        @collect DataFrames.DataFrame # `T` does not work here unless typeof(X) has no type-parameters
    end

end

# to select a single column `c` of any tabular data `X` with
# `retrieve(X, Cols, c)`; note this will try to return a
# CategoricalArray, and does so for X a DataFrame and TypedTable.Table:
function retrieve(::Val{true}, X::T, ::Type{Cols}, c::I) where {T, I<:Union{Symbol,Integer}}

    row_iterator = @from row in X begin
        @select project(row, c)
        @collect DataFrames.DataFrame
    end

    return row_iterator[1]

end

# to select rows `r` of any tabular data `X` with `retrieve(X, Rows, c)`:
# CURRENTLY RETURNS A DATAFRAME AND NOT THE ORIGINAL TYPE
function retrieve(::Val{true}, X::T, ::Type{Rows}, r) where T

    row_iterator = @from row in X begin
        @select row
        @collect
    end

    return @from row in row_iterator[r] begin
        @select row
        @collect DataFrames.DataFrame # `T` does not work here unless typeof(X) has no type-parameters
    end

end

# to get the number of nrows, ncols, feature names and eltypes of an
# tabular data:
function retrieve(::Val{true}, X, ::Type{Schema})

    TableTraits.isiterabletable(X) || error("Argument is not an iterable table.")

    row_iterator = @from row in X begin
        @select row
        @collect
    end

    nrows = length(row_iterator)
    if nrows == 0
        ncols = 0
        _names = Symbol[]
        _eltypes = DataType[]
    else
        row = first(row_iterator) # a named tuple
        ncols = length(row)
        _names = keys(row) |> collect
        _eltypes = DataType[eltype(x) for x in row]
    end

    return Schema(nrows, ncols, _names, _eltypes)

end

# retrieve(df::JuliaDB.IndexedTable, ::Type{Rows}, r) = df[r]
# retrieve(df::JuliaDB.NextTable, ::Type{Cols}, c) = select(df, c)
# retrieve(df::JuliaDB.NextTable, ::Type{Names}) = getfields(typeof(df.columns.columns))
# retrieve(df::JuliaDB.NextTable, ::Type{NRows}) = length(df)

retrieve(::Val{false}, A::Matrix, ::Type{Rows}, r) = A[r,:]
retrieve(::Val{false}, A::Matrix, ::Type{Cols}, c) = A[:,c]
function retrieve(::Val{false}, A::Matrix{T}, ::Type{Schema}) where T
    nrows = size(A, 1)
    ncols = size(A, 2)
    _names = [Symbol(:x, j) for j in 1:ncols]
    _eltypes = fill(T, ncols)
    return Schema(nrows, ncols, _names, _eltypes)
end

retrieve(::Val{false}, v::Vector, ::Type{Rows}, r) = v[r]
retrieve(::Val{false}, v::Vector, ::Type{Cols}, c) = error("Abstract vectors are not column-indexable.")
retrieve(::Val{false}, v::Vector{T}, ::Type{Schema}) where T = Schema(length(v), 1, [:x], [T,])
retrieve(::Val{false}, v::CategoricalArray{T,1,S} where {T,S}, ::Type{Rows}, r) = @inbounds v[r]
retrieve(::Val{false}, v::CategoricalArray{T,1,S} where {T,S}, ::Type{Cols}, r) =
    error("Categorical vectors are not column-indexable.")
retrieve(::Val{false}, v::CategoricalArray{T,1,S} where {T,S}, ::Type{Schema}) =
    retrieve(Val(false), collect(v), Schema)


using ReinforcementLearningEnvironmentGym, Test, ReinforcementLearningBase

test_envinterface(GymEnv("CartPole-v0"))