using DocOpt

doc="""Align two feature vector sequences and create parallel data in batch.
Note that filename of source and target feature is assumed to be same.

Usage:
    align.jl [options] <src_dir> <tgt_dir> <dst_dir>
    align.jl --version
    align.jl -h | --help

Options:
    -h --help       show this message
    --threshold=TH  threshold that is used to remove silence [default: -14.0]
    --max=MAX       Maximum number that will be processed [default: 200]
"""

using VoiceConversion
using HDF5, JLD

using Logging
@Logging.configure(level=DEBUG, output=STDOUT)

searchdir(path, key) = filter(x -> contains(x, key), readdir(path))

function mkdir_if_not_exist(dir)
    if !isdir(dir)
        println("Create $(dir)")
        run(`mkdir -p $dir`)
    end
end

function align!(src::Dict,
                tgt::Dict;
                threshold::AbstractFloat=-14.0,
                )
    src_fm = src["feature_matrix"]
    tgt_fm = tgt["feature_matrix"]

    max_num_frames = max(size(src_fm, 2), size(tgt_fm, 2))

    if src["type"] == "MelCepstrum"
        src_fm, tgt_fm = align_mcep(src_fm, tgt_fm,
                                    src["alpha"],
                                    src["fftlen"];
                                    threshold=threshold)
    else
        src_fm, tgt_fm = _align(src_fm, tgt_fm)
    end

    @assert size(src_fm) == size(tgt_fm)

    @info("The number of aligned frames: $(size(src_fm, 2))/$(max_num_frames)")
    if size(src_fm, 2) ==  0
        @warn("No frame found in aligned data. Probably threshold is too high.")
    end

    @assert !any(isnan(src_fm))
    @assert !any(isnan(tgt_fm))

    src["feature_matrix"] = src_fm
    tgt["feature_matrix"] = tgt_fm

    # type Dict{Union(UTF8String, ASCIIString), Any} is saved as
    # Dict{UTF8String, Any} and cause error in reading JLD file.
    # remove off Union and then save do the trick (but why? bug in HDF5?)
    @assert isa(src, Dict{String, Any})
    @assert isa(tgt, Dict{String, Any})

    Dict{String, Any}(src), Dict{String, Any}(tgt)
end

let
    args = docopt(doc, version=v"0.0.1")

    srcdir = args["<src_dir>"]
    tgtdir = args["<tgt_dir>"]
    dstdir = args["<dst_dir>"]
    mkdir_if_not_exist(dstdir)

    threshold = parse(Float64, args["--threshold"])
    nmax = parse(Int, args["--max"])

    files = searchdir(srcdir, ".jld")
    @info("$(length(files)) data found.")

    count = 0
    for filename in files
        # filename is assumed to be same between src and tgt
        srcpath = joinpath(srcdir, filename)
        tgtpath = joinpath(tgtdir, filename)
        savepath = joinpath(dstdir, string(splitext(basename(srcpath))[1],
                                      "_parallel.jld"))

        @info("Start processing $(srcpath) and $(tgtpath)")
        elapsed = @elapsed begin
            src = load(srcpath)
            tgt = load(tgtpath)

            src, tgt = align!(src, tgt, threshold=threshold)
            save(savepath,
                 "description", "Parallel data",
                 "src", src,
                 "tgt", tgt
                 )
        end
        @info("Elapsed time in alignment is $(elapsed) sec.")
        @info("Dumped to $(savepath)")

        count += 1
        count >= nmax && break
    end

    @info("Finished")
end


function instantiate(model_type::DataType, data, pools, param::Dict{Symbol, <:Any})
    debug(LOGGER, "Instantiating model of type $(model_type) with params $(param)")
    m = model_type(data, pools)
    set_param!(m, param)
    return m
end

function set_param!(model, param::Dict{Symbol, <:Any})
    has_changed = false
    for (k, v) in param
        k in keys(get_model_params(model)) || throw(ArgumentError("Model $(typeof(model)) has no parameter $k. Valid parameters are $(get_model_params(model))."))
        fieldtype = typeof(getfield(model, k))
        is_valid_param_value(model, Val{k}, v) || throw(ArgumentError("Parameter $k has value $v which is outside the valid range for model $(typeof(model))."))
        if getfield(model,k) != v
            has_changed = true
            setfield!(model, k, v)
        end
    end
    has_changed && invalidate_solution!(model)
    return nothing
end

function set_pools!(model, pools::Vector{Symbol})
    islabelenc(pools, learning_pool_enc) || throw(ArgumentError("Pools must be a LabelEncoding of $(learning_pool_enc.label)."))
    set_pools!(model, labelmap(pools))
    return nothing
end

function set_pools!(model, pools::Dict{Symbol, Vector{Int}})
    if !isdefined(model, :pools) ||model.pools != pools
        model.pools = pools
        invalidate_solution!(model)
    end
    return nothing
end

function set_adjust_K!(model, adjust_K::Bool)
    if adjust_K && !model.adjust_K
        model.adjust_K = adjust_K
        update_K!(model)
        invalidate_solution!(model)
    elseif !adjust_K && model.adjust_K
        model.adjust_K = adjust_K
        invalidate_solution!(model)
    end
    return nothing
end

is_K_adjusted(model)::Bool = model.adjust_K

calculate_kernel_matrix(model) = MLKernels.kernelmatrix(Val(:col), model.kernel_fct, model.data)

function calculate_kernel_matrix(model::SubSVDD)
     map(k -> MLKernels.kernelmatrix(Val(:col), model.kernel_fct[k], model.data[model.subspaces[k],:]), eachindex(model.subspaces))
 end

function update_K!(model)
    updated_K = calculate_kernel_matrix(model)
    if !isdefined(model, :K) || updated_K != model.K || (!isdefined(model, :K_adjusted) && model.adjust_K)
        debug(LOGGER, "[UPDATE_K] Updating Kernel matrix.")
        invalidate_solution!(model)
        model.K = updated_K
        if model.adjust_K
            model.K_adjusted = adjust_kernel_matrix(updated_K)
        end
    else
        debug(LOGGER, "[UPDATE_K] Nothing to update, kernel matrix has not changed.")
    end
    return nothing
end

function set_data!(model, data::Array{T, 2}) where T <: Real
    if model.data != data
        model.data = data
        update_K!(model)
        invalidate_solution!(model)
    end
    return nothing
end

function set_kernel!(model::OCClassifier, kernel_fct)
    if !isdefined(model, :kernel_fct) || model.kernel_fct != kernel_fct
        model.kernel_fct = kernel_fct
        update_K!(model)
        invalidate_solution!(model)
    end
    return nothing
end

function initialize!(model::OCClassifier, strategy::InitializationStrategy)
    debug(LOGGER, "[INITIALIZE] Initializing $(typeof(model)) with strategy: $(typeof(strategy))")
    C, kernel_fct = get_parameters(model, strategy)
    debug(LOGGER, "[INITIALIZE] Fitted parameters are: C = $(C) and kernel_fct = $(kernel_fct)")
    set_C!(model, C)
    set_kernel!(model, kernel_fct)
    model.state = model_initialized
    return nothing
end

get_kernel(model::OCClassifier) = model.kernel_fct

apply_update_strategy!(model, new_pools, query_ids, old_idx_remaining, new_idx_remaining) = nothing

function update_with_feedback!(model, new_data, new_pools::Vector{Symbol}, query_ids::Vector{Int},
    old_idx_remaining::Vector{Int},
    new_idx_remaining::Vector{Int})
    @assert size(new_data, 2) == length(new_pools)
    @assert islabelenc(new_pools, SVDD.learning_pool_enc)

    apply_update_strategy!(model, new_pools, query_ids, old_idx_remaining, new_idx_remaining)

    set_pools!(model, new_pools)
    set_data!(model, new_data)
    invalidate_solution!(model)
    return nothing
end


## parameter checks, error structures and messages

#@noinline
function checkrange(value::Int64, min::Int64, max::Int64)
    (min <= value <= max) || throw(BoundsError(min:max),value)
    nothing
end

#@noinline
function checkulimit(value::UInt64,max::UInt64)
    value <= max || throw(BoundsError(0:max,value))
end

checkrange(value::Integer,min::Integer,max::Integer) = checkrange(Int64(value),Int64(min),Int64(max))

checkulimit(value::Unsigned,limit::Unsigned) = checkulimit(value%UInt64,limit%UInt64)

"verify value nonnegative and <= limit"
checkulimit(value::Integer,limit::Integer) = checkrange(value,0,limit)

ckeckbyteofs(ofs::UInt32, t) = ofs<usize(t) || throw(BoundsError(t,ofs))

#Base.@propagate_inbounds function checksize(size::Unsigned, maxsize)

"test if index range is valid for given string "
checkrange(s::AbstractString,first::Integer, last::Integer) = checkbounds(s,first:last)

"test if offset/size range is valid for given s, defaults to offset+size<=usize(d)"
function checkrange(offset::UInt32, size::UInt64, s)
    checkbounds(s,offset+size)
end

"""
    dumpbyte(io:IO, byte::UInt8)

write a byte value ascii-readable: replace control characters and nonAscii values
 a hex notation prefixed with a dollar character. a dollar in text contents is doubled
"""
function dumpbyte(io::IO, byte::UInt8)
    if byte<32 || byte>127 || byte==UInt8('$')
        write(io,'$', Base.bytes2hex[1 + byte>>>4], Base.bytes2hex[1 + byte&0x0f])
    else
        if byte==UInt8('$')
            write(io,byte) # double $ to distinguish from hex notation
        end
        write(io,byte)
    end
end


abstract type AbstractNBodySystem end
abstract type AbstractInteraction end
abstract type AbstractSummationAlgorithm{T} end

struct Pairwise{isparallel} <: AbstractSummationAlgorithm{isparallel} end
struct PairwiseSymmetric{isparallel} <: AbstractSummationAlgorithm{isparallel} end
struct PairList{isparallel} <: AbstractSummationAlgorithm{isparallel} end

function number_of_particles(::AbstractNBodySystem, ::AbstractInteraction) end
function interactions(::AbstractNBodySystem) end
function pairlist(::AbstractNBodySystem, ::AbstractInteraction) end
function initial_conditions(::AbstractNBodySystem) end
function atype(::AbstractNBodySystem) end
function force!(dv, ::AbstractInteraction, ::AbstractNBodySystem) end
function force!(dv, ::AbstractInteraction, ::AbstractNBodySystem, i, j) end
function particle_mass(::AbstractNBodySystem, i) end
function summation_alg(::AbstractInteraction) end
function issparse(::AbstractInteraction) end

boundary_conditions!(s) = CallbackSet()
thermostats!(s) = CallbackSet()
neighbours(s) = CallbackSet()


module SuperLU

    import Base: \, lufact, size, nfilled

    export
        DMat,
        NCMat,
        superlu_options_t,
        SuperLUStat_t,
        SuperMatrix,
        nv,
        sp_ienv

    include("superlu_h.jl")
    include("superlu_types.jl")
    include("SuperLUCode.jl")

end # module


################################################################################
#                   CORE DISPATCHVARIABLEREF METHOD EXTENSIONS
################################################################################
# Extend dispatch_variable_ref
function dispatch_variable_ref(model::InfiniteModel,
                               index::DependentParameterIndex
                               )::DependentParameterRef
    return DependentParameterRef(model, index)
end

# Extend _add_data_object
function _add_data_object(model::InfiniteModel,
                          object::MultiParameterData
                          )::DependentParametersIndex
    index = MOIUC.add_item(model.dependent_params, object)
    push!(model.param_object_indices, index)
    return index
end

# Extend _data_dictionary (type based)
function _data_dictionary(model::InfiniteModel,
    ::Type{DependentParameters})::MOIUC.CleverDict
    return model.dependent_params
end

# Extend _data_dictionary (ref based)
function _data_dictionary(pref::DependentParameterRef)::MOIUC.CleverDict
    return JuMP.owner_model(pref).dependent_params
end

# Extend _data_object
function _data_object(pref::DependentParameterRef)::MultiParameterData
    object = get(_data_dictionary(pref), JuMP.index(pref).object_index, nothing)
    object === nothing && error("Invalid dependent parameter reference, cannot find " *
                           "corresponding parameter in the model. This is likely " *
                           "caused by using the reference of a deleted parameter.")
    return object
end

# Extend _core_variable_object
function _core_variable_object(pref::DependentParameterRef)::DependentParameters
    return _data_object(pref).parameters
end

# Return the number of dependent parameters involved
function _num_parameters(pref::DependentParameterRef)::Int
    return length(_data_object(pref).names)
end

# Extend _delete_data_object
function _delete_data_object(vref::DependentParameterRef)::Nothing
    delete!(_data_dictionary(vref), JuMP.index(vref).object_index)
    return
end

################################################################################
#                             PARAMETER DEFINITION
################################################################################
# Store partially processed individual dependent parameters
struct _DependentParameter{S <: AbstractInfiniteSet, M <: AbstractDerivativeMethod}
    set::S
    supports::Vector{Float64}
    name::String
    deriv_method::M
    function _DependentParameter(set::S,
        supports::Union{Vector{<:Real}, Real},
        name::String,
        method::M
        )::_DependentParameter{S, M} where {S <: AbstractInfiniteSet, M <: AbstractDerivativeMethod}
        if supports isa Real
            return new{S, M}(set, [supports], name, method)
        else
            return new{S, M}(set, supports, name, method)
        end
    end
end

## Use type dispatch to efficiently check the set(s)
# All InfiniteScalarSets
function _check_param_sets(_error::Function,
    params::AbstractArray{<:_DependentParameter{<:InfiniteScalarSet}}
    )::Nothing
    return
end

# All MultiDistributionSet with non SparseAxisArray
function _check_param_sets(_error::Function,
    params::AbstractArray{<:_DependentParameter{<:MultiDistributionSet}}
    )::Nothing
    dist = first(params).set.distribution
    set_type = typeof(first(params).set)
    if size(dist) != size(params)
        _error("The dimensions of the parameters and the multi-dimensional " *
               "distribution $dist.")
    end
    return
end

# All MultiDistributionSet with SparseAxisArray
function _check_param_sets(_error::Function,
    params::JuMPC.SparseAxisArray{<:_DependentParameter{<:MultiDistributionSet}}
    )
    _error("Cannot specify multiple-dimensional distribution set with a " *
           "`SparseAxisArray` of dependent infinite parameters.")
end

# All CollectionSet
function _check_param_sets(_error::Function,
    params::AbstractArray{<:_DependentParameter{<:CollectionSet}}
    )::Nothing
    sets = collection_sets(first(params).set)
    set_type = typeof(first(params).set)
    if length(sets) != length(params)
        _error("The dimensions of the parameters and the specified CollectionSet " *
               "do not match.")
    elseif params isa JuMPC.SparseAxisArray
        @warn("CollectionSet order may not match the given `SparseAxisArray` " *
              "of specified dependent infinite parameters, consider instead " *
              "specifying the `InfiniteScalarSet` for each parameter using " *
              "the `set` keyword and the appropriate indices.")
    end
    return
end

# All some InfiniteArraySet (for extensions)
function _check_param_sets(_error::Function,
    params::AbstractArray{<:_DependentParameter{<:InfiniteArraySet}}
    )::Nothing
    return
end

# Mixed sets
function _check_param_sets(_error::Function,
    params::AbstractArray{<:_DependentParameter}
    )::Nothing
    set_type = typeof(first(params).set)
    if !all(param.set isa InfiniteScalarSet for param in params)
        _error("Cannot specify multiple `InfiniteScalarSets` for one container " *
               "of infinite dependent parameters.")
    end
    return
end

# Fallback
function _check_param_sets(_error::Function, params)
    _error("Unrecognized input for infinite set.")
end

## Define methods for checking the the derivative methods
# Expected format
function _check_derivative_methods(_error::Function,
    params::AbstractArray{<:_DependentParameter})::Nothing
    if !all(p.deriv_method isa NonGenerativeDerivativeMethod for p in params)
        _error("Cannot use generative derivative evaluation methods with dependent " *
               "infinite parameters. Only subtypes of `NonGenerativeDerivativeMethod` " *
               "can be used.")
    end
    return
end

# Fallback
function _check_derivative_methods(_error::Function, params)
    _error("Unrecognized input for derivative method.")
end

## Use set type dispatch to make the proper InfiniteArraySet
# InfiniteArraySet
function _make_array_set(params::Vector{<:_DependentParameter{T}}
                         )::T where {T <: InfiniteArraySet}
    return first(params).set
end

# InfiniteScalarSets
function _make_array_set(params::Vector{<:_DependentParameter{T}}
                         )::CollectionSet{T} where {T <: InfiniteScalarSet}
    return CollectionSet([p.set for p in params])
end

# Build a DependentParameters object given an array of _DependentParameters
function _build_parameters(_error::Function,
                           params::AbstractArray{<:_DependentParameter};
                           num_supports::Int = 0,
                           sig_digits::Int = DefaultSigDigits,
                           extra_kw_args...)
    # error with extra keywords
    for (kwarg, _) in extra_kw_args
       _error("Unrecognized keyword argument $kwarg")
    end
    # check the formatting
    _check_param_sets(_error, params)
    _check_derivative_methods(_error, params)
    # vectorize the parameter array
    indices = Collections._get_indices(params)
    ordered_params = Collections._make_ordered(params, indices)
    vector_params = _make_vector(ordered_params)
    # make the set
    set = _make_array_set(vector_params)
    # make the supports and labels
    lens = map(p -> length(p.supports), vector_params)
    _allequal(lens) || _error("Inconsistent support dimensions.")
    # we have supports
    if first(lens) != 0
        # build the support array transpose to fill in column order (leverage locality)
        trans_supps = Array{Float64}(undef, first(lens), length(vector_params))
        for i in 1:size(trans_supps, 2)
            trans_supps[:, i] = vector_params[i].supports
        end
        supps = permutedims(trans_supps)
        supports_in_set(supps, set) || _error("Supports violate infinite set domain.")
        supps = round.(supps, sigdigits = sig_digits)
        label = UserDefined
        supp_dict = Dict{Vector{Float64}, Set{DataType}}(@views supps[:, i] =>
                                         Set([label]) for i in 1:size(supps, 2))
    # we want to generate supports
    elseif num_supports != 0
        supps, label = generate_support_values(set, num_supports = num_supports,
                                               sig_digits = sig_digits)
        supp_dict = Dict{Vector{Float64}, Set{DataType}}(@views supps[:, i] =>
                                         Set([label]) for i in 1:size(supps, 2))
    # no supports are specified
    else
        supp_dict = Dict{Vector{Float64}, Set{DataType}}()
    end
    # make the parameter object
    names = map(p -> p.name, vector_params)
    methods = map(p -> p.deriv_method, vector_params)
    return DependentParameters(set, supp_dict, sig_digits, methods), names, indices
end

"""
    add_parameters(model::InfiniteModel,
                   params::DependentParameters,
                   names::Vector{String} = ["noname", "noname", ...],
                   indices = nothing
                   )::AbstractArray{<:GeneralVariableRef}

Add `params` to `model` and return an appropriate container of the dependent
infinite parameter references. This is intended as an internal method for use
with [`@dependent_parameters`](@ref). However, if desired users can use this
to create a container of infinite dependent parameter without the use of a
macro. `names` denote the name of each parameter and `indices` denote the
indices of the expected container as used by `Containers._make_array`
(implemented by `VectorTuple`s), by default a `Vector` is returned.

**Example**
```julia-repl
julia> using Distributions

julia> dist = MvNormal(ones(3)); # 3 dimensional

julia> set = MultiDistributionSet(dist); # 3 dimensional

julia> params = DependentParameters(set, Dict{Vector{Float64}, Set{DatatType}}(), 10);

julia> prefs = add_parameters(model, params, ["par1", "par2", "par3"])
3-element Array{GeneralVariableRef,1}:
 par1
 par2
 par3
```
"""
function add_parameters(model::InfiniteModel,
                        params::DependentParameters,
                        names::Vector{String} = String[],
                        indices = nothing
                        )::AbstractArray{<:GeneralVariableRef}
    # get the number of parameters
    num_params = length(params.set)
    # process the names
    if isempty(names)
        names = ["noname" for i in 1:num_params]
    end
    # process the indices
    if indices === nothing
        indices = CartesianIndices(1:num_params)
    end
    # make the parameter model object
    obj_num = length(_param_object_indices(model)) + 1
    first_param_num = model.last_param_num + 1
    last_param_num = model.last_param_num += num_params
    param_nums = first_param_num:last_param_num
    data_object = MultiParameterData(params, obj_num, param_nums, names)
    # add the data object to the model and make the references
    obj_index = _add_data_object(model, data_object)
    prefs = [GeneralVariableRef(model, obj_index.value, DependentParameterIndex, i)
             for i in 1:num_params]
    return Collections._make_array(prefs, indices)
end

# Construct an expression to build an infinite set(s) (use with @dependent_parameters)
function _construct_array_set(_error::Function, info::_ParameterInfoExpr)
    if (info.has_lb || info.has_ub) && !(info.has_lb && info.has_ub)
        _error("Must specify both an upper bound and a lower bound")
    elseif info.has_lb
        check = :(isa($(info.lower_bound), Real))
        return :($(check) ? IntervalSet($(info.lower_bound), $(info.upper_bound)) : error("Bounds must be a real number."))
    elseif info.has_dist
        check = :(isa($(info.distribution), Distributions.UnivariateDistribution))
        return :($(check) ? UniDistributionSet($(info.distribution)) : MultiDistributionSet($(info.distribution)))
    elseif info.has_set
        check1 = :(isa($(info.set), AbstractInfiniteSet))
        check2 = :(isa($(info.set), Distributions.UnivariateDistribution))
        return :($(check1) ? $(info.set) : ($(check2) ? UniDistributionSet($(info.set)) : MultiDistributionSet($(info.set))))
    else
        _error("Must specify upper/lower bounds, a distribution, or a set")
    end
end

################################################################################
#                                  NAMING
################################################################################
# Get the parameter index in the DependentParameters object
_param_index(pref::DependentParameterRef)::Int = JuMP.index(pref).param_index

"""
    JuMP.name(pref::DependentParameterRef)::String

Extend [`JuMP.name`](@ref JuMP.name(::JuMP.VariableRef)) to return the names of
infinite dependent parameters.

**Example**
```julia-repl
julia> name(pref)
"par_name"
```
"""
function JuMP.name(pref::DependentParameterRef)::String
    object = get(_data_dictionary(pref), JuMP.index(pref).object_index, nothing)
    return object === nothing ? "" : object.names[_param_index(pref)]
end

"""
    JuMP.set_name(pref::DependentParameterRef, name::String)::Nothing

Extend [`JuMP.set_name`](@ref JuMP.set_name(::JuMP.VariableRef, ::String)) to set
names of dependent infinite parameters.

**Example**
```julia-repl
julia> set_name(vref, "par_name")

julia> name(vref)
"para_name"
```
"""
function JuMP.set_name(pref::DependentParameterRef, name::String)::Nothing
    _data_object(pref).names[_param_index(pref)] = name
    JuMP.owner_model(pref).name_to_param = nothing
    return
end

################################################################################
#                           PARAMETER DEPENDENCIES
################################################################################
# Extend _infinite_variable_dependencies
function _infinite_variable_dependencies(pref::DependentParameterRef
    )::Vector{InfiniteVariableIndex}
    return _data_object(pref).infinite_var_indices
end

# Extend _parameter_function_dependencies
function _parameter_function_dependencies(pref::DependentParameterRef
    )::Vector{ParameterFunctionIndex}
    return _data_object(pref).parameter_func_indices
end

# Extend _measure_dependencies
function _measure_dependencies(pref::DependentParameterRef
    )::Vector{MeasureIndex}
    return _data_object(pref).measure_indices[_param_index(pref)]
end

# Extend _constraint_dependencies
function _constraint_dependencies(pref::DependentParameterRef
    )::Vector{ConstraintIndex}
    return _data_object(pref).constraint_indices[_param_index(pref)]
end

# Extend _derivative_dependencies
function _derivative_dependencies(pref::DependentParameterRef
    )::Vector{DerivativeIndex}
    return _data_object(pref).derivative_indices[_param_index(pref)]
end

"""
    used_by_infinite_variable(pref::DependentParameterRef)::Bool

Return a `Bool` indicating if the dependent infinite parameter `pref` is used by
an infinite variable.

**Example**
```julia-repl
julia> used_by_infinite_variable(pref)
true
```
"""
function used_by_infinite_variable(pref::DependentParameterRef)::Bool
    return !isempty(_infinite_variable_dependencies(pref))
end

"""
    used_by_parameter_function(pref::DependentParameterRef)::Bool

Return a `Bool` indicating if the dependent infinite parameter `pref` is used by
an infinite parameter function.

**Example**
```julia-repl
julia> used_by_parameter_function(pref)
true
```
"""
function used_by_parameter_function(pref::DependentParameterRef)::Bool
    return !isempty(_parameter_function_dependencies(pref))
end

"""
    used_by_measure(pref::DependentParameterRef)::Bool

Return a `Bool` indicating if the dependent infinite parameter `pref` is used
by a measure.

**Example**
```julia-repl
julia> used_by_measure(pref)
true
```
"""
function used_by_measure(pref::DependentParameterRef)::Bool
    return !isempty(_measure_dependencies(pref))
end

"""
    used_by_constraint(pref::DependentParameterRef)::Bool

Return a `Bool` indicating if the dependent infinite parameter `pref` is used by
a constraint.

**Example**
```julia-repl
julia> used_by_constraint(pref)
false
```
"""
function used_by_constraint(pref::DependentParameterRef)::Bool
    return !isempty(_constraint_dependencies(pref))
end

"""
    used_by_derivative(pref::DependentParameterRef)::Bool

Return a `Bool` indicating if the dependent infinite parameter `pref` is used by
a derivative.

**Example**
```julia-repl
julia> used_by_derivative(pref)
false
```
"""
function used_by_derivative(pref::DependentParameterRef)::Bool
    return !isempty(_derivative_dependencies(pref))
end

# Extend used by objective
used_by_objective(pref::DependentParameterRef)::Bool = false

"""
    is_used(pref::DependentParameterRef)::Bool

Return a `Bool` indicating if the dependent infinite parameter `pref` is used in
the model.

**Example**
```julia-repl
julia> is_used(pref)
true
```
"""
function is_used(pref::DependentParameterRef)::Bool
    return used_by_measure(pref) || used_by_constraint(pref) ||
           used_by_infinite_variable(pref) || used_by_derivative(pref) ||
           used_by_parameter_function(pref)
end

################################################################################
#                          PARAMETER OBJECT METHODS
################################################################################
# Extend _parameter_number
function _parameter_number(pref::DependentParameterRef)::Int
    return _data_object(pref).parameter_nums[_param_index(pref)]
end

# Extend _parameter_numbers
function _parameter_numbers(pref::DependentParameterRef)::Vector{Int}
    return [_parameter_number(pref)]
end

# Extend _object_number
function _object_number(pref::DependentParameterRef)::Int
    return _data_object(pref).object_num
end

# Extend _object_numbers
function _object_numbers(pref::DependentParameterRef)::Vector{Int}
    return [_object_number(pref)]
end

## Set helper methods for adapting data_objects with parametric changes
# No change needed
function _adaptive_data_update(pref::DependentParameterRef, params::P,
    data::MultiParameterData{P})::Nothing where {P <: DependentParameters}
    data.parameters = params
    return
end

# Reconstruction is necessary
function _adaptive_data_update(pref::DependentParameterRef, params::P1,
    data::MultiParameterData{P2})::Nothing  where {P1, P2}
    new_data = MultiParameterData(params, data.object_num, data.parameter_nums,
                                  data.names, data.parameter_func_indices,
                                  data.infinite_var_indices,
                                  data.derivative_indices, data.measure_indices,
                                  data.constraint_indices,
                                  data.has_internal_supports,
                                  data.has_deriv_constrs)
    _data_dictionary(pref)[JuMP.index(pref).object_index] = new_data
    return
end

# Extend _set_core_variable_object
function _set_core_variable_object(pref::DependentParameterRef,
                                   params::DependentParameters)::Nothing
    _adaptive_data_update(pref, params, _data_object(pref))
    return
end

################################################################################
#                        DERIVATIVE METHOD FUNCTIONS
################################################################################
# Extend fallback for dependent parameters
function has_derivative_supports(pref::DependentParameterRef)::Bool
    return false
end

# Extend fallback for dependent parameters
function _set_has_derivative_supports(pref::DependentParameterRef,
                                      status::Bool)::Nothing
    return
end

# Extend has derivative constraints
function has_derivative_constraints(pref::DependentParameterRef)::Bool
    return _data_object(pref).has_deriv_constrs[_param_index(pref)]
end

# Extend setting if has derivative constraints
function _set_has_derivative_constraints(pref::DependentParameterRef,
                                         status::Bool)::Nothing
    _data_object(pref).has_deriv_constrs[_param_index(pref)] = status
    return
end

# Get the raw derivative method vector
function _derivative_methods(pref::DependentParameterRef)
    return _core_variable_object(pref).derivative_methods
end

"""
    derivative_method(pref::DependentParameterRef)::NonGenerativeDerivativeMethod

Returns the numerical derivative evaluation method employed with `pref` when it
is used as an operator parameter in a derivative.

**Example**
```julia-repl
julia> derivative_method(pref)
FiniteDifference
```
"""
function derivative_method(pref::DependentParameterRef)::NonGenerativeDerivativeMethod
    return _derivative_methods(pref)[_param_index(pref)]
end

## Define helper methods for setting the derivative method efficiently
# Compatible with vector type
function _adaptive_method_update(pref,
    p::DependentParameters{S, M1},
    method::M2
    )::Nothing where {S, M1 <: NonGenerativeDerivativeMethod, M2 <: M1}
    p.derivative_methods[_param_index(pref)] = method
    return
end

# Not compatible
function _adaptive_method_update(pref,
    p::DependentParameters{S, M1},
    method::M2
    )::Nothing where {S, M1, M2}
    methods = p.derivative_methods
    new_methods = [i == _param_index(pref) ? method : m
                   for (i, m) in enumerate(methods)]
    new_params = DependentParameters(p.set, p.supports, p.sig_digits, new_methods)
    _set_core_variable_object(pref, new_params)
    return
end

"""
    set_derivative_method(pref::DependentParameterRef,
                          method::NonGenerativeDerivativeMethod)::Nothing

Specfies the desired derivative evaluation method `method` for derivatives that are
taken with respect to `pref`. Errors if `method` is generative (i.e., it requires
the definition of additional supports)

**Example**
```julia-repl
julia> set_derivative_method(d, FiniteDifference())

```
"""
function set_derivative_method(pref::DependentParameterRef,
    method::AbstractDerivativeMethod
    )::Nothing
    if !(method isa NonGenerativeDerivativeMethod)
        error("Must specify a subtype of `NonGenerativeDerivativeMethod` for " *
              "for a dependent parameter.")
    end
    _adaptive_method_update(pref, _core_variable_object(pref), method)
    _reset_derivative_evaluations(pref)
    if is_used(pref)
        set_optimizer_model_ready(JuMP.owner_model(pref), false)
    end
    return
end

"""
    set_all_derivative_methods(model::InfiniteModel,
                               method::AbstractDerivativeMethod)::Nothing

Sets the desired evaluation method `method` for all the derivatives currently added
to `model`. Note that this is done with respect to the infinite parameters. Errors
if a generative method is specified and the model contains dependent parameters.

**Example**
```julia-repl
julia> set_all_derivative_methods(model, OrthogonalCollocation(2))

```
"""
function set_all_derivative_methods(model::InfiniteModel,
    method::AbstractDerivativeMethod
    )::Nothing
    for pref in all_parameters(model, InfiniteParameter)
        set_derivative_method(pref, method)
    end
    return
end

################################################################################
#                             INFINITE SET METHODS
################################################################################
## Get the individual infinite set if possible
# raw_set
function _parameter_set(pref::DependentParameterRef)::InfiniteArraySet
    return _core_variable_object(pref).set
end

# CollectionSet
function _parameter_set(set::CollectionSet{S},
                        pref::DependentParameterRef
                        )::S where {S <: InfiniteScalarSet}
    return collection_sets(set)[_param_index(pref)]
end

# InfiniteArraySet (Fallback)
function _parameter_set(set::InfiniteArraySet, pref::DependentParameterRef)
    error("An individual infinite set is not well-defined for $pref which " *
          "is part of a group of dependent infinite parameters that correspond " *
          "to an multi-dimensional infinite set of type `$(typeof(set))`.")
end

"""
    infinite_set(pref::DependentParameterRef)::InfiniteScalarSet

Return the infinite set associated with the particular infinite dependent
parameter `pref` if valid. Errors if the underlying [`DependentParameters`](@ref)
object does not use a [`CollectionSet`](@ref).

**Example**
```julia-repl
julia> infinite_set(x[1])
[-1, 1]
```
"""
function infinite_set(pref::DependentParameterRef)::InfiniteScalarSet
    return _parameter_set(_parameter_set(pref), pref)
end

# Check that prefs are complete
function _check_complete_param_array(
    prefs::AbstractArray{<:DependentParameterRef}
    )::Nothing
    if length(prefs) != _num_parameters(first(prefs))
        error("Dimensions of parameter container and the infinite set do not " *
              "match, ensure all related dependent parameters are included.")
    end
    return
end

"""
    infinite_set(prefs::AbstractArray{<:DependentParameterRef})::InfiniteArraySet

Return the infinite set associated with the container of infinite dependent
parameters `prefs`. Errors if the container `prefs` is incomplete.

**Example**
```julia-repl
julia> infinite_set(x)
ZeroMeanDiagNormal(
dim: 2
μ: [0.0, 0.0]
Σ: [1.0 0.0; 0.0 1.0]
)
```
"""
function infinite_set(prefs::AbstractArray{<:DependentParameterRef}
                      )::InfiniteArraySet
    _check_complete_param_array(prefs)
    return _parameter_set(first(prefs))
end

# Update the underlying set and delete the supports
function _update_parameter_set(pref::DependentParameterRef,
                               new_set::InfiniteArraySet)::Nothing
    old_params = _core_variable_object(pref)
    new_supports = Dict{Vector{Float64}, Set{DataType}}()
    sig_figs = significant_digits(pref)
    methods = _derivative_methods(pref)
    new_params = DependentParameters(new_set, new_supports, sig_figs, methods)
    _set_core_variable_object(pref, new_params)
    for i in 1:length(new_set)
        idx = DependentParameterIndex(JuMP.index(pref).object_index, i)
        p = DependentParameterRef(JuMP.owner_model(pref), idx)
        _reset_derivative_evaluations(p)
    end
    _set_has_internal_supports(pref, false)
    if is_used(pref)
        set_optimizer_model_ready(JuMP.owner_model(pref), false)
    end
    return
end

"""
    set_infinite_set(pref::DependentParameterRef,
                     set::InfiniteScalarSet)::Nothing

Specify the scalar infinite set of the dependent infinite parameter `pref` to
`set` if `pref` is part of a [`CollectionSet`](@ref), otherwise an error is
thrown. Note this will reset/delete all the supports contained in the
underlying [`DependentParameters`](@ref) object. Also, errors if `pref` is used
by a measure.

**Example**
```julia-repl
julia> set_infinite_set(x[1], IntervalSet(0, 2))

julia> infinite_set(x[1])
[0, 2]
```
"""
function set_infinite_set(pref::DependentParameterRef,
                          set::InfiniteScalarSet)::Nothing
    old_set = _parameter_set(pref)
    if !(old_set isa CollectionSet)
        error("Cannot set the individual infinite set of $pref if the " *
              "underlying set is not a CollectionSet.")
    elseif used_by_measure(pref)
        error("Cannot override the infinite set of $pref since it is used by " *
              "a measure.")
    end
    param_idx = _param_index(pref)
    new_set = CollectionSet([i != param_idx ? collection_sets(old_set)[i] : set
                             for i in eachindex(collection_sets(old_set))])
    _update_parameter_set(pref, new_set)
    return
end

"""
    set_infinite_set(prefs::AbstractArray{<:DependentParameterRef},
                     set::InfiniteArraySet)::Nothing

Specify the multi-dimensional infinite set of the dependent infinite parameters
`prefs` to `set`. Note this will reset/delete all the supports contained in the
underlying [`DependentParameters`](@ref) object. This will error if the not all
of the dependent infinite parameters are included, if any of them are used by
measures.

**Example**
```julia-repl
julia> set_infinite_set(x, CollectionSet([IntervalSet(0, 1), IntervalSet(0, 2)]))
```
"""
function set_infinite_set(prefs::AbstractArray{<:DependentParameterRef},
                          set::InfiniteArraySet)::Nothing
    if any(used_by_measure(pref) for pref in prefs)
        error("Cannot override the infinite set of $prefs since it is used by " *
              "a measure.")
    end
    _check_complete_param_array(prefs)
    _update_parameter_set(first(prefs), set)
    return
end

"""
    JuMP.has_lower_bound(pref::DependentParameterRef)::Bool

Extend the `JuMP.has_lower_bound` function to accomodate a single dependent
infinite parameter.
Return true if the set associated with `pref` has a defined lower bound or if a
lower bound can be found. Extensions with user-defined scalar infinite set types
should extend `JuMP.has_lower_bound(set::NewType)`.

**Example**
```julia-repl
julia> has_lower_bound(x[1])
true
```
"""
function JuMP.has_lower_bound(pref::DependentParameterRef)::Bool
    set = _parameter_set(pref)
    if set isa CollectionSet
        return JuMP.has_lower_bound(collection_sets(set)[_param_index(pref)])
    else
        return false
    end
end

"""
    JuMP.lower_bound(pref::DependentParameterRef)::Number

Extend the `JuMP.lower_bound` function to accomodate a single dependent infinite
parameter. Returns the lower bound associated with the infinite set. Errors if
such a bound is not well-defined.

**Example**
```julia-repl
julia> lower_bound(x[1])
0.0
```
"""
function JuMP.lower_bound(pref::DependentParameterRef)::Number
    if !JuMP.has_lower_bound(pref)
        error("Parameter $(pref) does not have a lower bound.")
    end
    return JuMP.lower_bound(infinite_set(pref))
end

"""
    JuMP.set_lower_bound(pref::DependentParameterRef, lower::Real)::Nothing

Extend the `JuMP.set_lower_bound` function to accomodate a single dependent
infinite parameter. Updates the infinite set lower bound if such an operation
is supported. Infinite scalar set extensions that seek to employ this should extend
`JuMP.set_lower_bound(set::NewType, lower::Number)`. This will call
[`set_infinite_set`](@ref) and will error if this is not well-defined. Note
that existing supports will be deleted.

**Example**
```julia-repl
julia> set_lower_bound(t, -1)

julia> lower_bound(t)
-1.0
```
"""
function JuMP.set_lower_bound(pref::DependentParameterRef, lower::Real)::Nothing
    set = infinite_set(pref)
    new_set = JuMP.set_lower_bound(set, lower)
    set_infinite_set(pref, new_set)
    return
end

"""
    JuMP.has_upper_bound(pref::DependentParameterRef)::Bool

Extend the `JuMP.has_upper_bound` function to accomodate a single dependent
infinite parameter.
Return true if the set associated with `pref` has a defined upper bound or if a
upper bound can be found. Extensions with user-defined scalar infinite set types
should extend `JuMP.has_upper_bound(set::NewType)`.

**Example**
```julia-repl
julia> has_upper_bound(x[1])
true
```
"""
function JuMP.has_upper_bound(pref::DependentParameterRef)::Bool
    set = _core_variable_object(pref).set
    if set isa CollectionSet
        return JuMP.has_upper_bound(collection_sets(set)[_param_index(pref)])
    else
        return false
    end
end

"""
    JuMP.upper_bound(pref::DependentParameterRef)::Number

Extend the `JuMP.upper_bound` function to accomodate a single dependent infinite
parameter. Returns the upper bound associated with the infinite set. Errors if
such a bound is not well-defined.

**Example**
```julia-repl
julia> upper_bound(x[1])
0.0
```
"""
function JuMP.upper_bound(pref::DependentParameterRef)::Number
    if !JuMP.has_upper_bound(pref)
        error("Parameter $(pref) does not have a upper bound.")
    end
    return JuMP.upper_bound(infinite_set(pref))
end

"""
    JuMP.set_upper_bound(pref::DependentParameterRef, upper::Real)::Nothing

Extend the `JuMP.set_upper_bound` function to accomodate a single dependent
infinite parameter. Updates the infinite set upper bound if such an operation
is supported. Infinite scalar set extensions that seek to employ this should extend
`JuMP.set_upper_bound(set::NewType, upper::Number)`. This will call
[`set_infinite_set`](@ref) and will error if this is not well-defined. Note
that existing supports will be deleted.

**Example**
```julia-repl
julia> set_upper_bound(t, -1)

julia> upper_bound(t)
-1.0
```
"""
function JuMP.set_upper_bound(pref::DependentParameterRef, upper::Real)::Nothing
    set = infinite_set(pref)
    new_set = JuMP.set_upper_bound(set, upper)
    set_infinite_set(pref, new_set)
    return
end

################################################################################
#                              SUPPORT METHODS
################################################################################
# Get the raw supports
function _parameter_supports(pref::DependentParameterRef
                             )::Dict{Vector{Float64}, Set{DataType}}
    return _core_variable_object(pref).supports
end

"""
    significant_digits(pref::DependentParameterRef)::Int

Return the number of significant digits enforced on the supports of `pref`.

**Example**
```julia-repl
julia> significant_digits(x[1])
12
```
"""
function significant_digits(pref::DependentParameterRef)::Int
    return _core_variable_object(pref).sig_digits
end

"""
    num_supports(pref::DependentParameterRef;
                 [label::Type{<:AbstractSupportLabel} = PublicLabel])::Int

Return the number of support points associated with a single dependent infinite
parameter `pref`. Specify a subset of supports via `label` to only count the
supports with `label`. By default only the amount of public supports are given, but
the full amount is obtained via `label == All`.

**Example**
```julia-repl
julia> num_supports(x[1])
2

julia> num_supports(x[1], label = MCSample)
0
```
"""
function num_supports(pref::DependentParameterRef;
                      label::Type{<:AbstractSupportLabel} = PublicLabel)::Int
    supp_dict = _parameter_supports(pref)
    if label == All || (!has_internal_supports(pref) && label == PublicLabel)
        return length(supp_dict)
    else
        return count(p -> any(v -> v <: label, p[2]), supp_dict)
    end
end

"""
    num_supports(prefs::AbstractArray{<:DependentParameterRef};
                 [label::Type{<:AbstractSupportLabel} = PublicLabel])::Int

Return the number of support points associated with dependent infinite
parameters `prefs`. Errors if not all from the same underlying object.
Specify a subset of supports via `label` to only count the supports with `label`.
By default only the amount of public supports are given, but the full amount is
obtained via `label == All`.

**Example**
```julia-repl
julia> num_supports(x)
2
```
"""
function num_supports(prefs::AbstractArray{<:DependentParameterRef};
    label::Type{<:AbstractSupportLabel} = PublicLabel
    )::Int
    _check_complete_param_array(prefs)
    return num_supports(first(prefs), label = label)
end

"""
    has_supports(pref::DependentParameterRef)::Bool

Return true if `pref` has supports or false otherwise.

**Example**
```julia-repl
julia> has_supports(x[1])
true
```
"""
has_supports(pref::DependentParameterRef)::Bool = !isempty(_parameter_supports(pref))

"""
    has_supports(prefs::AbstractArray{<:DependentParameterRef})::Bool

Return true if `prefs` have supports or false otherwise. Errors if not all of the
infinite dependent parameters are from the same object.

**Example**
```julia-repl
julia> has_supports(x)
true
```
"""
function has_supports(prefs::AbstractArray{<:DependentParameterRef})::Bool
    _check_complete_param_array(prefs)
    return has_supports(first(prefs))
end

"""
    supports(pref::DependentParameterRef;
             [label::Type{<:AbstractSupportLabel} = PublicLabel])::Vector{Float64}

Return the support points associated with `pref`. A subset of supports can be
returned via `label` to return just the supports associated with `label`. By
default only the public supports are given, but the full set is
obtained via `label == All`.

**Example**
```julia-repl
julia> supports(x[1])
2-element Array{Float64,1}:
 0.0
 1.0
```
"""
function supports(pref::DependentParameterRef;
                  label::Type{<:AbstractSupportLabel} = PublicLabel)::Vector{Float64}
    supp_dict = _parameter_supports(pref)
    pindex = _param_index(pref)
    if label == All || (!has_internal_supports(pref) && label == PublicLabel)
        return Float64[supp[pindex] for supp in keys(supp_dict)]
    else
        reduced_supps = findall(e -> any(v -> v <: label, e), supp_dict)
        return Float64[supp[pindex] for supp in reduced_supps]
    end
end

"""
    supports(prefs::AbstractArray{<:DependentParameterRef};
             [label::Type{<:AbstractSupportLabel} = PublicLabel]
             )::Union{AbstractArray{<:Vector{<:Float64}}, Array{Float64, 2}}

Return the support points associated with `prefs`. Errors if not all of the
infinite dependent parameters are from the same object. This will return a
matrix if `prefs` is `Vector`, otherwise an array of vectors is returned by
calling `supports.(prefs)`. A subset of supports can be
returned via `label` to return just the supports associated with `label`. By
default only the public supports are given, but the full set is  obtained via
`label == All`.

**Example**
```julia-repl
julia> supports(x) # columns are supports
2×2 Array{Float64,2}:
 0.0  1.0
 0.0  1.0
```
"""
function supports(prefs::AbstractArray{<:DependentParameterRef};
                  label::Type{<:AbstractSupportLabel} = PublicLabel
                  )::AbstractArray{<:Vector{<:Float64}}
    _check_complete_param_array(prefs)
    return supports.(prefs, label = label) # TODO make more efficient
end

# More efficient dispatch for Vectors
function supports(prefs::Vector{DependentParameterRef};
                  label::Type{<:AbstractSupportLabel} = PublicLabel
                  )::Array{Float64, 2}
    if !has_supports(prefs)
        return zeros(Float64, _num_parameters(first(prefs)), 0)
    elseif label == All || (!has_internal_supports(first(prefs)) && label == PublicLabel)
        raw_supps = keys(_parameter_supports(first(prefs)))
        if length(raw_supps) == 1
            return reduce(hcat, collect(raw_supps))
        else
            return reduce(hcat, raw_supps)
        end
    else
        raw_supps = findall(e -> any(v -> v <: label, e),
                            _parameter_supports(first(prefs)))
        if isempty(raw_supps)
            return zeros(Float64, _num_parameters(first(prefs)), 0)
        else
            return reduce(hcat, raw_supps)
        end
    end
end

# Define method for overriding the current supports
function _update_parameter_supports(prefs::AbstractArray{<:DependentParameterRef},
                                    supports::Array{<:Real, 2},
                                    label::Type{<:AbstractSupportLabel})::Nothing
    set = _parameter_set(first(prefs))
    new_supps = Dict{Vector{Float64}, Set{DataType}}(@views supports[:, i] =>
                                      Set([label]) for i in 1:size(supports, 2))
    sig_figs = significant_digits(first(prefs))
    methods = _derivative_methods(first(prefs))
    new_params = DependentParameters(set, new_supps, sig_figs, methods)
    _set_core_variable_object(first(prefs), new_params)
    _set_has_internal_supports(first(prefs), label <: InternalLabel)
    for pref in prefs
        _reset_derivative_evaluations(pref)
    end
    if any(is_used(pref) for pref in prefs)
        set_optimizer_model_ready(JuMP.owner_model(first(prefs)), false)
    end
    return
end

# Process an array of vectors into a support matrix
function _make_support_matrix(prefs::AbstractArray{<:DependentParameterRef},
                              supports::AbstractArray{<:Vector{<:Real}}
                              )::Array{<:Real, 2}
    _keys(supports) == _keys(prefs) || error("Inconsistent support indices")
    lens = [length(supp) for supp in supports]
    _allequal(lens) || error("Inconsistent support dimensions.")
    trans_supps = Array{Float64}(undef, first(lens), length(prefs))
    for k in eachindex(prefs)
        trans_supps[:, _param_index(prefs[k])] = supports[k]
    end
    return permutedims(trans_supps)
end

"""
    set_supports(prefs::AbstractArray{<:DependentParameterRef},
                 supports::AbstractArray{<:Vector{<:Real}};
                 [force::Bool = false,
                 label::Type{<:AbstractSupportLabel} = UserDefined])::Nothing

Specify the support points for `prefs`. Errors if the supports violate the domain
of the infinite set, if the dimensions don't match up properly,
if `prefs` and `supports` have different indices, not all of the `prefs` are
from the same dependent infinite parameter container, there are existing
supports and `force = false`. Note that it is strongly preferred to use
`add_supports` if possible to avoid destroying measure dependencies.

```julia
    set_supports(prefs::Vector{DependentParameterRef},
                 supports::Array{<:Real, 2};
                 [force::Bool = false,
                 label::Type{<:AbstractSupportLabel} = UserDefined])::Nothing
```
Specify the supports for a vector `prefs` of dependent infinite parameters.
Here rows of `supports` correspond to `prefs` and the columns correspond to the
supports. This is more efficient than the above method and will error for the
same reasons.

**Example**
```julia-repl
julia> set_supports(y, [[0, 1], [0, 1]])

julia> set_supports(x, [0 1; 0 1])

julia> supports(x)
2×2 Array{Float64,2}:
 0.0  1.0
 0.0  1.0
```
"""
function set_supports(prefs::AbstractArray{<:DependentParameterRef},
                      supports::AbstractArray{<:Vector{<:Real}};
                      force::Bool = false,
                      label::Type{<:AbstractSupportLabel} = UserDefined
                      )::Nothing
    supps = _make_support_matrix(prefs, supports)
    set_supports(_make_vector(prefs), supps, force = force, label = label)
    return
end

# Efficient method for vector prefs and matrix of supports
function set_supports(prefs::Vector{DependentParameterRef},
                      supports::Array{<:Real, 2};
                      force::Bool = false,
                      label::Type{<:AbstractSupportLabel} = UserDefined
                      )::Nothing
    set = infinite_set(prefs) # this does a check on prefs
    if has_supports(prefs) && !force
        error("Unable set supports for $prefs since they already have supports." *
              " Consider using `add_supports` or use set `force = true` to " *
              "overwrite the existing supports.")
    elseif !supports_in_set(supports, set)
        error("Supports violate the domain of the infinite set.")
    end
    supports = round.(supports, sigdigits = significant_digits(first(prefs)))
    _update_parameter_supports(prefs, supports, label)
    return
end

# Error for single dependent parameters
function set_supports(pref::DependentParameterRef, supports; kwargs...)
    error("Cannot modify the supports of a single dependent infinite parameter.")
end

"""
    add_supports(prefs::AbstractArray{<:DependentParameterRef},
                 supports::AbstractArray{<:Vector{<:Real}};
                 [label::Type{<:AbstractSupportLabel} = UserDefined])::Nothing

Add additional support points for `prefs`. Errors if the supports violate the domain
of the infinite set, if the dimensions don't match up properly,
if `prefs` and `supports` have different indices, or not all of the `prefs` are
from the same dependent infinite parameter container.

```julia
    add_supports(prefs::Vector{DependentParameterRef},
                 supports::Array{<:Real, 2};
                 [label::Type{<:AbstractSupportLabel} = UserDefined])::Nothing
```
Specify the supports for a vector `prefs` of dependent infinite parameters.
Here rows of `supports` correspond to `prefs` and the columns correspond to the
supports. This is more efficient than the above method and will error for the
same reasons.

**Example**
```julia-repl
julia> add_supports(x, [[1], [1]])

julia> supports(x)
2×2 Array{Float64,2}:
 0.0  1.0
 0.0  1.0

julia> add_supports(x, ones(2, 1) * 0.5)

julia> supports(t)
2×3 Array{Float64,2}:
 0.0  1.0  0.5
 0.0  1.0  0.5
```
"""
function add_supports(prefs::AbstractArray{<:DependentParameterRef},
                      supports::AbstractArray{<:Vector{<:Real}};
                      label::Type{<:AbstractSupportLabel} = UserDefined, # interal keyword args
                      check::Bool = true)::Nothing
    supps = _make_support_matrix(prefs, supports)
    add_supports(_make_vector(prefs), supps, label = label, check = check)
    return
end

# More efficient version for supports in the correct format
function add_supports(prefs::Vector{DependentParameterRef},
                      supports::Array{<:Real, 2};
                      label::Type{<:AbstractSupportLabel} = UserDefined, # internal keyword args
                      check::Bool = true)::Nothing
    set = infinite_set(prefs) # this does a check on prefs
    if check && !supports_in_set(supports, set)
        error("Supports violate the domain of the infinite set.")
    end
    supports = round.(supports, sigdigits = significant_digits(first(prefs)))
    current_supports = _parameter_supports(first(prefs))
    for i in 1:size(supports, 2)
        s = @view(supports[:, i])
        if haskey(current_supports, s)
            push!(current_supports[s], label)
        else
            current_supports[s] = Set([label])
        end
    end
    if label <: InternalLabel
        _set_has_internal_supports(first(prefs), true)
    end
    for pref in prefs
        _reset_derivative_evaluations(pref)
    end
    if any(is_used(pref) for pref in prefs)
        set_optimizer_model_ready(JuMP.owner_model(first(prefs)), false)
    end
    return
end

# Error for single dependent parameters
function add_supports(pref::DependentParameterRef, supports; kwargs...)
    error("Cannot modify the supports of a single dependent infinite parameter.")
end

"""
    delete_supports(prefs::AbstractArray{<:DependentParameterRef};
                    [label::Type{<:AbstractSupportLabel} = All])::Nothing

Delete the support points for `prefs`. Errors if any of the parameters are
used by a measure or if not all belong to the same set of dependent parameters.
If `label != All` then that label is removed along with any supports that solely
contain that label.

**Example**
```julia-repl
julia> delete_supports(w)

```
"""
function delete_supports(prefs::AbstractArray{<:DependentParameterRef};
                         label::Type{<:AbstractSupportLabel} = All)::Nothing
    _check_complete_param_array(prefs)
    supp_dict = _parameter_supports(first(prefs))
    for pref in prefs
        _reset_derivative_evaluations(pref)
    end
    if label == All
        if any(used_by_measure(pref) for pref in prefs)
            error("Cannot delete supports with measure dependencies.")
        end
        empty!(supp_dict)
        _set_has_internal_supports(first(prefs), false)
    else
        filter!(p -> !all(v -> v <: label, p[2]), supp_dict)
        for (k, v) in supp_dict
            filter!(l -> !(l <: label), v)
        end
        pref1 = first(prefs)
        if has_internal_supports(pref1) && num_supports(pref1, label = InternalLabel) == 0
            _set_has_internal_supports(pref1, false)
        end
    end
    if any(is_used(pref) for pref in prefs)
        set_optimizer_model_ready(JuMP.owner_model(first(prefs)), false)
    end
    return
end

# Error for single dependent parameters
function delete_supports(pref::DependentParameterRef)
    error("Cannot delete the supports of a single dependent infinite parameter.")
end

# TODO resolve case that there are existing UniformGrid supports
"""
    generate_and_add_supports!(prefs::AbstractArray{<:DependentParameterRef},
                               set::InfiniteArraySet,
                               [method::Type{<:AbstractSupportLabel}];
                               [num_supports::Int = DefaultNumSupports])::Nothing

Generate supports for `prefs` via [`generate_support_values`](@ref) and add them
to `pref`. This is intended as an extendable internal method for
[`fill_in_supports!`](@ref fill_in_supports!(::AbstractArray{<:DependentParameterRef})).
Most extensions that employ user-defined infinite sets can typically enable this
by extending [`generate_support_values`](@ref). However, in some cases it may be
necessary to extend this when more complex operations need to take place then just
adding supports to a set of infinite parameters. Errors if the
infinite set type is not recognized.
"""
function generate_and_add_supports!(prefs::AbstractArray{<:DependentParameterRef},
                                    set::InfiniteArraySet;
                                    num_supports::Int = DefaultNumSupports)::Nothing
    new_supps, label = generate_supports(set,
                                         num_supports = num_supports,
                                    sig_digits = significant_digits(first(prefs)))
    add_supports(_make_vector(prefs), new_supps, check = false, label = label)
    return
end

# Method dispatch
function generate_and_add_supports!(prefs::AbstractArray{<:DependentParameterRef},
                                    set::InfiniteArraySet,
                                    method::Type{<:AbstractSupportLabel};
                                    num_supports::Int = DefaultNumSupports)::Nothing
    new_supps, label = generate_supports(set, method,
                                         num_supports = num_supports,
                                    sig_digits = significant_digits(first(prefs)))
    add_supports(_make_vector(prefs), new_supps, check = false, label = label)
    return
end

"""
    fill_in_supports!(prefs::AbstractArray{<:DependentParameterRef};
                      [num_supports::Int = DefaultNumSupports,
                       modify::Bool = true])::Nothing

Automatically generate support points for a container of dependent infinite
parameters `prefs`. Generating up to `num_supports` for the parameters in accordance
with `generate_and_add_supports!`. Will add nothing if there are supports and
`modify = false`. Extensions that use user defined
set types should extend [`generate_and_add_supports!`](@ref) and/or
[`generate_support_values`](@ref) as needed. Errors if the infinite set type is
not recognized.

**Example**
```julia-repl
julia> fill_in_supports!(x, num_supports = 4)

julia> supports(x)
2×4 Array{Float64,2}:
 0.0  0.333  0.667  1.0
 0.0  0.333  0.667  1.0
```
"""
function fill_in_supports!(prefs::AbstractArray{<:DependentParameterRef};
                           num_supports::Int = DefaultNumSupports,
                           modify::Bool = true)::Nothing
    set = infinite_set(prefs) # does check for bad container
    current_amount = InfiniteOpt.num_supports(first(prefs))
    if (modify || current_amount == 0) && (current_amount < num_supports)
        generate_and_add_supports!(prefs, set,
                                   num_supports = num_supports - current_amount)
    end
    return
end

# Error for single dependent parameters
function fill_in_supports!(pref::DependentParameterRef; kwargs...)
    error("Cannot modify the supports of a single dependent infinite parameter.")
end

"""
    fill_in_supports!(model::InfiniteModel; [num_supports::Int = DefaultNumSupports,
                      modify::Bool = true])::Nothing

Automatically generate support points for all infinite parameters in model.
This calls `fill_in_supports!` for each parameter in the model.
See [`fill_in_supports!`](@ref)
for more information. Errors if one of the infinite set types is unrecognized.
Note that no supports will be added to a particular parameter if it already has
some and `modify = false`.

**Example**
```julia-repl
julia> fill_in_supports!(model, num_supports = 4)

julia> supports(t)
4-element Array{Float64,1}:
 0.0
 0.333
 0.667
 1.0
```
"""
function fill_in_supports!(model::InfiniteModel; num_supports::Int = DefaultNumSupports,
                           modify::Bool = true)::Nothing
    # fill in the the supports of each independent parameter
    for (key, data_object) in model.independent_params
        pref = dispatch_variable_ref(model, key)
        fill_in_supports!(pref, num_supports = num_supports, modify = modify)
    end
    # fill in the supports of each dependent parameter set
    for (key, data_object) in model.dependent_params
        prefs = [dispatch_variable_ref(model, DependentParameterIndex(key, i))
                 for i in 1:length(data_object.names)]
        fill_in_supports!(prefs, num_supports = num_supports, modify = modify)
    end
    return
end

################################################################################
#                          MODEL PARAMETER QUERIES
################################################################################
"""
    num_parameters(model::InfiniteModel,
                   [type::Type{InfOptParameter} = InfOptParameter])::Int

Return the number of `InfiniteOpt` parameters assigned to `model`. By default,
the total number of infinite and finite parameters is returned. The amount
of a particular type is obtained by specifying the concrete parameter type
of [`InfOptParameter`](@ref) via `type`. Type options include:
 - `InfOptParameter`: all parameters
 - `ScalarParameter`: all scalar parameters
 - `InfiniteParameter`: all infinite parameters
 - `FiniteParameter`: all finite parameters
 - `IndependentParameter`: all independent infinite parameters
 - `DependentParameters`: all dependent infinite parameters

**Example**
```julia-repl
julia> num_parameters(model)
3

julia> num_parameters(model, IndependentParameter)
2
```
"""
function num_parameters(model::InfiniteModel,
                        type::Type{InfOptParameter} = InfOptParameter
                        )::Int
    num_pars = num_parameters(model, IndependentParameter)
    num_pars += num_parameters(model, FiniteParameter)
    num_pars += num_parameters(model, DependentParameters)
    return num_pars
end

# Particular scalar parameter types
function num_parameters(model::InfiniteModel,
                        type::Type{C})::Int where {C <: ScalarParameter}
    return length(_data_dictionary(model, type))
end

# ScalarParameter
function num_parameters(model::InfiniteModel,
                        type::Type{ScalarParameter})::Int
    num_pars = num_parameters(model, FiniteParameter)
    num_pars += num_parameters(model, IndependentParameter)
    return num_pars
end

# DependentParameters
function num_parameters(model::InfiniteModel,
                        type::Type{DependentParameters})::Int
    num_pars = 0
    for (_, object) in _data_dictionary(model, type)
        num_pars += length(object.names)
    end
    return num_pars
end

# InfiniteParameter
function num_parameters(model::InfiniteModel,
                        type::Type{InfiniteParameter})::Int
    num_pars = num_parameters(model, IndependentParameter)
    num_pars += num_parameters(model, DependentParameters)
    return num_pars
end

"""
    all_parameters(model::InfiniteModel,
                   type::Type{InfOptParameter} = InfOptParameter
                   )::Vector{GeneralVariableRef}

Return a list of all the `InfiniteOpt` parameters assigned to `model`. By default,
all of the infinite and finite parameters is returned. The search is reduced to
a particular type is obtained by specifying the concrete parameter type
of [`InfOptParameter`](@ref) via `type`. Type options include:
- `InfOptParameter`: all parameters
- `ScalarParameter`: all scalar parameters
- `InfiniteParameter`: all infinite parameters
- `FiniteParameter`: all finite parameters
- `IndependentParameter`: all independent infinite parameters
- `DependentParameters`: all dependent infinite parameters

**Examples**
```julia-repl
julia> all_parameters(model)
4-element Array{GeneralVariableRef,1}:
 t
 x[1]
 x[2]
 alpha

julia> all_parameters(model, FiniteParameter)
1-element Array{GeneralVariableRef,1}:
 alpha
```
"""
function all_parameters(model::InfiniteModel,
                        type::Type{InfOptParameter} = InfOptParameter
                        )::Vector{GeneralVariableRef}
    prefs_list = all_parameters(model, IndependentParameter)
    append!(prefs_list, all_parameters(model, DependentParameters))
    append!(prefs_list, all_parameters(model, FiniteParameter))
    return prefs_list
end

# Particular scalar parameter types
function all_parameters(model::InfiniteModel,
                        type::Type{C}
                        )::Vector{GeneralVariableRef} where {C <: InfOptParameter}
    prefs_list = Vector{GeneralVariableRef}(undef, num_parameters(model, type))
    for (i, (index, _)) in enumerate(_data_dictionary(model, type))
        prefs_list[i] = _make_parameter_ref(model, index)
    end
    return prefs_list
end

# ScalarParameter
function all_parameters(model::InfiniteModel,
                        type::Type{ScalarParameter})::Vector{GeneralVariableRef}
    prefs_list = all_parameters(model, IndependentParameter)
    append!(prefs_list, all_parameters(model, FiniteParameter))
    return prefs_list
end

# DependentParameters
function all_parameters(model::InfiniteModel,
                        type::Type{DependentParameters}
                        )::Vector{GeneralVariableRef}
    prefs_list = Vector{GeneralVariableRef}(undef, num_parameters(model, type))
    counter = 1
    for (index, object) in _data_dictionary(model, type)
        for i in eachindex(object.names)
            dep_idx = DependentParameterIndex(index, i)
            prefs_list[counter] = _make_parameter_ref(model, dep_idx)
            counter += 1
        end
    end
    return prefs_list
end

# InfiniteParameter
function all_parameters(model::InfiniteModel,
                        type::Type{InfiniteParameter}
                        )::Vector{GeneralVariableRef}
    prefs_list = all_parameters(model, DependentParameters)
    append!(prefs_list, all_parameters(model, IndependentParameter))
    return prefs_list
end

################################################################################
#                                 DELETION
################################################################################
"""
    JuMP.delete(model::InfiniteModel,
                prefs::AbstractArray{<:DependentParameterRef})::Nothing

Extend `JuMP.delete` to delete
dependent infinite parameters and their dependencies. All variables, constraints, and
measure functions that depend on `prefs` are updated to exclude them. Errors if the
parameters are contained in an `AbstractMeasureData` datatype that is employed by
a measure since the measure becomes invalid otherwise. Thus, measures that
contain this dependency must be deleted first. Note that
[`parameter_refs`](@ref parameter_refs(::AbstractMeasureData)) needs to be
extended to allow deletion of parameters when custom `AbstractMeasureData`
datatypes are used. Note that any dependent infinite variables will have their
start values reset to the default via [`reset_start_value_function`](@ref).

**Example**
```julia-repl
julia> print(model)
Min measure(g(t, x)*t + x) + z
Subject to
 z ≥ 0.0
 g(t, x) + z ≥ 42.0, ∀ t ∈ [0, 6], x[1] ∈ [-1, 1], x[2] ∈ [-1, 1]
 g(0.5, x) = 0, x[1] ∈ [-1, 1], x[2] ∈ [-1, 1]

julia> delete(model, x)

julia> print(model)
Min measure(g(t)*t) + z
Subject to
 g(t) + z ≥ 42.0, ∀ t ∈ [0, 6]
 g(0.5) = 0
```
"""
function JuMP.delete(model::InfiniteModel,
                     prefs::AbstractArray{<:DependentParameterRef})::Nothing
    @assert JuMP.is_valid(model, first(prefs)) "Parameter references are invalid."
    _check_complete_param_array(prefs)
    gvrefs = [_make_parameter_ref(model, JuMP.index(pref)) for pref in prefs]
    # ensure deletion is okay (prefs are not used by measure data)
    for pref in gvrefs
        for mindex in _measure_dependencies(pref)
            data = measure_data(dispatch_variable_ref(model, mindex))
            _check_param_in_data(pref, data)
        end
    end
    # make sure it isn't used by parameter function
    if used_by_parameter_function(first(prefs))
        error("Cannot delete `$prefs` since they are used by an infinite " *
              "parameter function.")
    end
    # update optimizer model status
    if any(is_used(pref) for pref in prefs)
        set_optimizer_model_ready(model, false)
    end
    # delete dependence of measures and constraints on prefs
    for pref in gvrefs
        _update_measures(model, pref)
        _update_constraints(model, pref)
    end
    # get the object and parameter numbers
    obj_num = _object_number(first(prefs))
    param_nums = collect(_data_object(first(prefs)).parameter_nums)
    # update infinite variables that depend on pref
    for vindex in _infinite_variable_dependencies(first(prefs))
        # remove the parameter dependences
        vref = InfiniteVariableRef(model, vindex)
        vprefs = raw_parameter_refs(vref)
        tup_index = findfirst(isequal(obj_num), _object_number.(vprefs[:, 1]))
        delete_indices = vprefs.ranges[tup_index]
        deleteat!(vprefs, tup_index, tuple_index = true)
        reset_start_value_function(vref)
        # update any point variables that depend on vref accordingly
        for pindex in _point_variable_dependencies(vref)
            pvref = PointVariableRef(model, pindex)
            deleteat!(raw_parameter_values(pvref), delete_indices)
        end
        # update any reduced variables that depend on vref accordingly
        for rindex in _reduced_variable_dependencies(vref)
            rvref = ReducedVariableRef(model, rindex)
            _update_reduced_variable(rvref, delete_indices)
        end
    end
    # delete derivatives that depend on any of these parameters
    for pref in gvrefs
        for index in _derivative_dependencies(pref)
            JuMP.delete(model, dispatch_variable_ref(model, index))
        end
    end
    # delete parameter information stored in model
    _delete_data_object(first(prefs))
    # update the object numbers and parameter numbers
    _update_model_numbers(model, obj_num, param_nums)
    return
end


#=
import Base: hash, string, show, tryparse

withtypes = [(F,T1,T2) for F in (:hash, :string) for T1 in  (:Float64, :XFloat16) for T2 in (:XFloat16, :Float16)]

for (TF,TX) in ((:Float64, :XFloat32)(:Float64, :XFloat32))
  for F in (:hash, :abs)
     $F(x::$TX) = $F(reinterpret($TF, $TX))#      @eval a = string($T," ",$F); println(a);
           end; end
withtypes = [(F,T) for F in (:hash, :string) for T in  (:Float32, :XFloat16)]
  @eval begin
    $F(x::$T) = reinterpret(

=#


function read_legacy(io::IO, ver::Float32)
  Z = getfield(BUF, :buf)
  L = getfield(BUF, :int64_buf)

  # read begins ------------------------------------------------------
  N     = fastread(io, Int64)
  checkbuf_strict!(L, 2*N)
  fast_readbytes!(io, Z, 3*N)
  c1    = copy(Z[1:N])
  c2    = copy(Z[N+1:2*N])
  y     = code2typ.(getindex(Z, 2*N+1:3*N))
  cmp   = fastread(io) % Bool
  read!(io, L)
  nx    = getindex(L, N+1:2*N)
  cmp && checkbuf_8!(Z, maximum(nx))

  ver < 0.5f0 && error("No legacy support for SeisIO file format version < 0.5")

  # Get file creation time
  fname = split(io.name)[2][1:end-1]
  isfile(fname) || error(string("Can't stat file ", fname))
  st = stat(fname)
  t0 = st.ctime
  if ver == 0.5f0
    S = SeisData(N)
    setfield!(S, :id, read_string_vec(io, Z))
    setfield!(S, :name, read_string_vec(io, Z))
    setfield!(S, :loc, InstrumentPosition[read(io, code2loctyp(getindex(c1, i))) for i = 1:N])
    read!(io, S.fs)
    read!(io, S.gain)

    # Here is the hard part. For older files, we read a degenerate InstResp
    R = InstrumentResponse[]
    for i = 1:N
      T = code2resptyp(getindex(c2, i))
      (T == GenResp) && (push!(R, read(io, GenResp)); continue)
      et = Complex{T == PZResp ? Float32 : Float64}

      # skip :c, it was never used
      fastskip(io, T == PZResp ? 4 : 8)

      # read poles
      np = fastread(io, Int64)
      p = zeros(et, np)
      read!(io, p)

      # read zeros
      nz = fastread(io, Int64)
      z = zeros(et, nz)
      read!(io, z)

      # push to R
      push!(R, T(p = p, z = z))
    end
    S.resp .= R
    setfield!(S, :units, read_string_vec(io, Z))
    setfield!(S, :src, read_string_vec(io, Z))
    setfield!(S, :misc, [read_misc(io, Z) for i = 1:N])
    setfield!(S, :notes, [read_string_vec(io, Z) for i = 1:N])
    setfield!(S, :t, [read!(io, Array{Int64, 2}(undef, getindex(L, i), 2)) for i = 1:N])
    setfield!(S, :x,
    FloatArray[cmp ?
    (fast_readbytes!(io, Z, getindex(nx, i)); Blosc.decompress(getindex(y,i), Z)) :
    read!(io, Array{getindex(y,i), 1}(undef, getindex(nx, i))) for i = 1:N])
  elseif ver < 0.53
    #=  Read process for 0.51 and 0.52 is identical; 0.52 added two Types <:
        InstrumentResponse that v0.51 doesn't know about, so this is a safe
        procedure =#
    S = SeisData(N)
    setfield!(S, :id, read_string_vec(io, Z))
    setfield!(S, :name, read_string_vec(io, Z))
    setfield!(S, :loc, InstrumentPosition[read(io, code2loctyp(getindex(c1, i))) for i = 1:N])
    read!(io, S.fs)
    read!(io, S.gain)

    # Here is the hard part. For older files, we read a degenerate InstResp
    R = InstrumentResponse[]
    for i = 1:N
      T = code2resptyp(getindex(c2, i))
      if T in (PZResp, PZResp64, GenResp)
        push!(R, read(io, T))
      elseif T == MultiStageResp
        K = fastread(io, Int64)
        codes = fastread(io, K)
        M = MultiStageResp(K)
        A = Array{RespStage,1}(undef, 0)
        for j = 1:K
          c = codes[j]
          if c == 0x03
            units_out = String(fastread(io, fastread(io, Int64)))
            units_in = String(fastread(io, fastread(io, Int64)))
            M.i[j] = units_in
            M.o[j] = units_out
            if j == 2
              M.i[1] = units_out
            end
            CR = CoeffResp(
                       read!(io, Array{Float64,1}(undef, fastread(io, Int64))),
                       read!(io, Array{Float64,1}(undef, fastread(io, Int64)))
                       )
            push!(A, CR)
          else
            push!(A, read(io, code2resptyp(c)))
          end
        end
        M.stage .= A
        read!(io, M.fs)
        read!(io, M.gain)
        read!(io, M.fg)
        read!(io, M.delay)
        read!(io, M.corr)
        read!(io, M.fac)
        read!(io, M.os)
        push!(R, M)
      end
    end
    S.resp .= R
    setfield!(S, :units, read_string_vec(io, Z))
    for i = 1:N
      if typeof(S.resp[i]) == MultiStageResp
        K = length(S.resp[i].fs)
        if K > 0
          S.resp[i].o[1] = S.units[i]
        end
      end
    end
    setfield!(S, :src, read_string_vec(io, Z))
    setfield!(S, :misc, [read_misc(io, Z) for i = 1:N])
    setfield!(S, :notes, [read_string_vec(io, Z) for i = 1:N])
    setfield!(S, :t, [read!(io, Array{Int64, 2}(undef, getindex(L, i), 2)) for i = 1:N])
    setfield!(S, :x,
    FloatArray[cmp ?
    (fast_readbytes!(io, Z, getindex(nx, i)); Blosc.decompress(getindex(y,i), Z)) :
    read!(io, Array{getindex(y,i), 1}(undef, getindex(nx, i))) for i = 1:N])
  end
  return S
end


const DenseNativeArray = DenseArray{<:NativeTypes}

"""
`vstorent!` (non-temporal store) requires data to be aligned.
`alignstores!` will align `y` in preparation for the non-temporal maps.
"""
function alignstores!(
    f::F, y::DenseArray{T},
    args::Vararg{DenseNativeArray,A}
) where {F, T <: Base.HWReal, A}
    N = length(y)
    ptry = VectorizationBase.zero_offsets(stridedpointer(y))
    ptrargs = VectorizationBase.zero_offsets.(stridedpointer.(args))
    V = VectorizationBase.pick_vector_width_val(T)
    W = unwrap(V)
    zero_index = MM{W}(Static(0))
    uintptry = reinterpret(UInt, pointer(ptry))
    @assert iszero(uintptry & (sizeof(T) - 1)) "The destination vector (`dest`) must be aligned at least to `sizeof(eltype(dest))`."
    alignment = uintptry & (VectorizationBase.REGISTER_SIZE - 1)
    if alignment > 0
        i = reinterpret(Int, W - (alignment >>> VectorizationBase.intlog2(sizeof(T))))
        m = mask(T, i)
        if N < i
            m &= mask(T, N & (W - 1))
        end
        vnoaliasstore!(ptry, f(vload.(ptrargs, ((zero_index,),), m)...), (zero_index,), m)
        gesp(ptry, (i,)), gesp.(ptrargs, ((i,),)), N - i
    else
        ptry, ptrargs, N
    end
end

function vmap_singlethread!(
    f::F, y::DenseArray{T},
    ::Val{NonTemporal},
    args::Vararg{DenseNativeArray,A}
) where {F,T <: Base.HWReal, A, NonTemporal}
    if NonTemporal # if stores into `y` aren't aligned, we'll get a crash
        ptry, ptrargs, N = alignstores!(f, y, args...)
    else
        N = length(y)
        ptry = VectorizationBase.zero_offsets(stridedpointer(y))
        ptrargs = VectorizationBase.zero_offsets.(stridedpointer.(args))
    end
    i = 0
    V = VectorizationBase.pick_vector_width_val(T)
    W = unwrap(V)
    st = VectorizationBase.static_sizeof(T)
    zero_index = MM{W}(Static(0), st)
    while i < N - ((W << 2) - 1)

        # vstore!(stridedpointer(B), VectorizationBase.VecUnroll((v1,v2,v3)), VectorizationBase.Unroll{AU,1,3,AV,W64,zero(UInt)}((i, j, k)))
        # vload(stridedpointer(B), VectorizationBase.Unroll{1,1,4,1,W,0x0000000000000000}((i,)))

        index = VectorizationBase.Unroll{1,1,4,1,W,0x0000000000000000}((i,))
        v = f(vload.(ptrargs, index)...)
        if NonTemporal
            vstorent!(ptry, v, index)
        else
            vnoaliasstore!(ptry, v, index)
        end
        i = vadd(i, 4W)
    end
    while i < N - (W - 1) # stops at 16 when
        vᵣ = f(vload.(ptrargs, ((MM{W}(i),),))...)
        if NonTemporal
            vstorent!(ptry, vᵣ, (MM{W}(i),))
        else
            vnoaliasstore!(ptry, vᵣ, (MM{W}(i),))
        end
        i = vadd(i, W)
    end
    if i < N
        m = mask(T, N & (W - 1))
        vnoaliasstore!(ptry, f(vload.(ptrargs, ((MM{W}(i),),), m)...), (MM{W}(i,),), m)
    end
    y
end

function vmap_multithreaded!(
    f::F,
    y::DenseArray{T},
    ::Val{true},
    args::Vararg{DenseNativeArray,A}
) where {F,T,A}
    ptry, ptrargs, N = alignstores!(f, y, args...)
    N > 0 || return y
    W, Wshift = VectorizationBase.pick_vector_width_shift(T)
    V = VectorizationBase.pick_vector_width_val(T)
    Wsh = Wshift + 2
    Niter = N >>> Wsh
    Base.Threads.@threads for j ∈ 0:Niter-1
        index = VectorizationBase.Unroll{1,1,4,1,W,0x0000000000000000}((j << Wsh,))
        vstorent!(ptry, f(vload.(ptrargs, index)...), index)
    end
    ii = Niter << Wsh
    while ii < N - (W - 1) # stops at 16 when
        vstorent!(ptry, f(vload.(ptrargs, ((MM{W}(ii),),))...), (MM{W}(ii),))
        ii = vadd(ii, W)
    end
    if ii < N
        m = mask(T, N & (W - 1))
        vnoaliasstore!(ptry, f(vload.(ptrargs, ((MM{W}(ii),),), m)...), (MM{W}(ii),), m)
    end
    y
end
function vmap_multithreaded!(
    f::F,
    y::DenseArray{T},
    ::Val{false},
    args::Vararg{DenseNativeArray,A}
) where {F,T,A}
    N = length(y)
    ptry = VectorizationBase.zero_offsets(stridedpointer(y))
    ptrargs = VectorizationBase.zero_offsets.(stridedpointer.(args))
    N > 0 || return y
    W, Wshift = VectorizationBase.pick_vector_width_shift(T)
    V = VectorizationBase.pick_vector_width_val(T)
    Wsh = Wshift + 2
    Niter = N >>> Wsh
    Base.Threads.@threads for j ∈ 0:Niter-1
        index = VectorizationBase.Unroll{1,1,4,1,W,0x0000000000000000}((j << Wsh,))
        vnoaliasstore!(ptry, f(vload.(ptrargs, index)...), index)
    end
    ii = Niter << Wsh
    while ii < N - (W - 1) # stops at 16 when
        vnoaliasstore!(ptry, f(vload.(ptrargs, ((MM{W}(ii),),))...), (MM{W}(ii),))
        ii = vadd(ii, W)
    end
    if ii < N
        m = mask(T, N & (W - 1))
        vnoaliasstore!(ptry, f(vload.(ptrargs, ((MM{W}(ii),),), m)...), (MM{W}(ii),), m)
    end
    y
end

"""
    vmap!(f, destination, a::AbstractArray)
    vmap!(f, destination, a::AbstractArray, b::AbstractArray, ...)

Vectorized-`map!`, applying `f` to each element of `a` (or paired elements of `a`, `b`, ...)
and storing the result in `destination`.
"""
function vmap!(
    f::F, y::DenseArray{T}, args::Vararg{DenseNativeArray,A}
) where {F,T<:Base.HWReal,A}
    vmap_singlethread!(f, y, Val{false}(), args...)
end

"""
    vmapt!(::Function, dest, args...)

Like `vmap!` (see `vmap!`), but uses `Threads.@threads` for parallel execution.
"""
function vmapt!(
    f::F, y::DenseArray{T}, args::Vararg{DenseNativeArray,A}
) where {F,T<:Base.HWReal,A}
    vmap_multithreaded!(f, y, Val{false}(), args...)
end

"""
    vmapnt!(::Function, dest, args...)

This is a vectorized map implementation using nontemporal store operations. This means that the write operations to the destination will not go to the CPU's cache.
If you will not immediately be reading from these values, this can improve performance because the writes won't pollute your cache. This can especially be the case if your arguments are very long.

```julia
julia> using LoopVectorization, BenchmarkTools

julia> x = rand(10^8); y = rand(10^8); z = similar(x);

julia> f(x,y) = exp(-0.5abs2(x - y))
f (generic function with 1 method)

julia> @benchmark map!(f, \$z, \$x, \$y)
BenchmarkTools.Trial:
  memory estimate:  0 bytes
  allocs estimate:  0
  --------------
  minimum time:     439.613 ms (0.00% GC)
  median time:      440.729 ms (0.00% GC)
  mean time:        440.695 ms (0.00% GC)
  maximum time:     441.665 ms (0.00% GC)
  --------------
  samples:          12
  evals/sample:     1

julia> @benchmark vmap!(f, \$z, \$x, \$y)
BenchmarkTools.Trial:
  memory estimate:  0 bytes
  allocs estimate:  0
  --------------
  minimum time:     178.147 ms (0.00% GC)
  median time:      178.381 ms (0.00% GC)
  mean time:        178.430 ms (0.00% GC)
  maximum time:     179.054 ms (0.00% GC)
  --------------
  samples:          29
  evals/sample:     1

julia> @benchmark vmapnt!(f, \$z, \$x, \$y)
BenchmarkTools.Trial:
  memory estimate:  0 bytes
  allocs estimate:  0
  --------------
  minimum time:     144.183 ms (0.00% GC)
  median time:      144.338 ms (0.00% GC)
  mean time:        144.349 ms (0.00% GC)
  maximum time:     144.641 ms (0.00% GC)
  --------------
  samples:          35
  evals/sample:     1
```
"""
function vmapnt!(
    f::F, y::DenseArray{T}, args::Vararg{DenseNativeArray,A}
) where {F,T<:Base.HWReal,A}
    vmap_singlethread!(f, y, Val{true}(), args...)
end

"""
    vmapntt!(::Function, dest, args...)

Like `vmapnt!` (see `vmapnt!`), but uses `Threads.@threads` for parallel execution.
"""
function vmapntt!(
    f::F, y::DenseArray{T}, args::Vararg{DenseNativeArray,A}
) where {F,T<:Base.HWReal,A}
    vmap_multithreaded!(f, y, Val{true}(), args...)
end

# generic fallbacks
@inline vmap!(f, args...) = map!(f, args...)
@inline vmapt!(f, args...) = map!(f, args...)
@inline vmapnt!(f, args...) = map!(f, args...)
@inline vmapntt!(f, args...) = map!(f, args...)

function vmap_call(f::F, vm!::V, args::Vararg{Any,N}) where {V,F,N}
    T = Base._return_type(f, Base.Broadcast.eltypes(args))
    dest = similar(first(args), T)
    vm!(f, dest, args...)
end

"""
    vmap(f, a::AbstractArray)
    vmap(f, a::AbstractArray, b::AbstractArray, ...)

SIMD-vectorized `map`, applying `f` to each element of `a` (or paired elements of `a`, `b`, ...)
and returning a new array.
"""
vmap(f::F, args::Vararg{Any,N}) where {F,N} = vmap_call(f, vmap!, args...)

"""
    vmapt(f, a::AbstractArray)
    vmapt(f, a::AbstractArray, b::AbstractArray, ...)

A threaded variant of [`vmap`](@ref).
"""
vmapt(f::F, args::Vararg{Any,N}) where {F,N} = vmap_call(f, vmapt!, args...)

"""
    vmapnt(f, a::AbstractArray)
    vmapnt(f, a::AbstractArray, b::AbstractArray, ...)

A "non-temporal" variant of [`vmap`](@ref). This can improve performance in cases where
`destination` will not be needed soon.
"""
vmapnt(f::F, args::Vararg{Any,N}) where {F,N} = vmap_call(f, vmapnt!, args...)

"""
    vmapntt(f, a::AbstractArray)
    vmapntt(f, a::AbstractArray, b::AbstractArray, ...)

A threaded variant of [`vmapnt`](@ref).
"""
vmapntt(f::F, args::Vararg{Any,N}) where {F,N} = vmap_call(f, vmapntt!, args...)

# @inline vmap!(f, y, x...) = @avx y .= f.(x...)
# @inline vmap(f, x...) = @avx f.(x...)


export DQNLearner

using Random
using Flux

"""
    DQNLearner(;kwargs...)
See paper: [Human-level control through deep reinforcement learning](https://www.nature.com/articles/nature14236)
# Keywords
- `approximator`::[`AbstractApproximator`](@ref): used to get Q-values of a state.
- `target_approximator`::[`AbstractApproximator`](@ref): similar to `approximator`, but used to estimate the target (the next state).
- `loss_func`: the loss function.
- `γ::Float32=0.99f0`: discount rate.
- `batch_size::Int=32`
- `update_horizon::Int=1`: length of update ('n' in n-step update).
- `min_replay_history::Int=32`: number of transitions that should be experienced before updating the `approximator`.
- `update_freq::Int=4`: the frequency of updating the `approximator`.
- `target_update_freq::Int=100`: the frequency of syncing `target_approximator`.
- `stack_size::Union{Int, Nothing}=4`: use the recent `stack_size` frames to form a stacked state.
- `rng = Random.GLOBAL_RNG`
"""
mutable struct DQNLearner{
    Tq<:AbstractApproximator,
    Tt<:AbstractApproximator,
    Tf,
    S<:Union{Int,Nothing},
    R<:AbstractRNG,
} <: AbstractLearner
    approximator::Tq
    target_approximator::Tt
    loss_func::Tf
    stack_size::S
    γ::Float32
    batch_size::Int
    update_horizon::Int
    min_replay_history::Int
    update_freq::Int
    target_update_freq::Int
    update_step::Int
    rng::R
    loss::Float32
end

function DQNLearner(;
    approximator::Tq,
    target_approximator::Tt,
    loss_func::Tf,
    stack_size::Union{Int,Nothing} = 4,
    γ::Float32 = 0.99f0,
    batch_size::Int = 32,
    update_horizon::Int = 1,
    min_replay_history::Int = 32,
    update_freq::Int = 1,
    target_update_freq::Int = 100,
    update_step::Int = 0,
    rng = Random.GLOBAL_RNG,
) where {Tq,Tt,Tf}
    copyto!(approximator, target_approximator)
    DQNLearner(
        approximator,
        target_approximator,
        loss_func,
        stack_size,
        γ,
        batch_size,
        update_horizon,
        min_replay_history,
        update_freq,
        target_update_freq,
        update_step,
        rng,
        0.0f0,
    )
end

Flux.functor(x::DQNLearner) = (Q = x.approximator, Qₜ = x.target_approximator),
y -> begin
    x = @set x.approximator = y.Q
    x = @set x.target_approximator = y.Qₜ
    x
end

"""

!!! note
    The state of the observation is assumed to have been stacked,
    if `!isnothing(stack_size)`.
"""
function (learner::DQNLearner)(env)
    env |>
    get_state |>
    x ->
        Flux.unsqueeze(x, ndims(x) + 1) |>
        x ->
            send_to_device(device(learner.approximator), x) |>
            learner.approximator |>
            vec |>
            send_to_host
end

function RLBase.update!(learner::DQNLearner, t::AbstractTrajectory)
    length(t[:terminal]) < learner.min_replay_history && return

    learner.update_step += 1

    if learner.update_step % learner.target_update_freq == 0
        copyto!(learner.target_approximator, learner.approximator)
    end

    learner.update_step % learner.update_freq == 0 || return

    experience = extract_experience(t, learner)

    Q = learner.approximator
    Qₜ = learner.target_approximator
    γ = learner.γ
    loss_func = learner.loss_func
    update_horizon = learner.update_horizon
    batch_size = learner.batch_size
    D = device(Q)
    states = send_to_device(D, experience.states)
    actions = CartesianIndex.(experience.actions, 1:batch_size)
    rewards = send_to_device(D, experience.rewards)
    terminals = send_to_device(D, experience.terminals)
    next_states = send_to_device(D, experience.next_states)

    target_q = Qₜ(next_states)
    if haskey(t, :next_legal_actions_mask)
        masked_value = fill(typemin(Float32), size(experience.next_legal_actions_mask))
        masked_value[experience.next_legal_actions_mask] .= 0
        target_q .+= send_to_device(D, masked_value)
    end

    q′ = dropdims(maximum(target_q; dims = 1), dims = 1)
    G = rewards .+ γ^update_horizon .* (1 .- terminals) .* q′

    gs = gradient(params(Q)) do
        q = Q(states)[actions]
        loss = loss_func(G, q)
        ignore() do
            learner.loss = loss
        end
        loss
    end

    update!(Q, gs)
end

function extract_experience(t::AbstractTrajectory, learner::DQNLearner)
    s = learner.stack_size
    h = learner.update_horizon
    n = learner.batch_size
    γ = learner.γ

    valid_ind_range =
        isnothing(s) ? (1:(length(t[:terminal])-h)) : (s:(length(t[:terminal])-h))
    inds = rand(learner.rng, valid_ind_range, n)
    next_inds = inds .+ h

    states = consecutive_view(t[:state], inds; n_stack = s)
    actions = consecutive_view(t[:action], inds)
    next_states = consecutive_view(t[:state], next_inds; n_stack = s)

    if haskey(t, :legal_actions_mask)
        legal_actions_mask = consecutive_view(t[:legal_actions_mask], inds)
        next_legal_actions_mask = consecutive_view(t[:next_legal_actions_mask], next_inds)
    else
        legal_actions_mask = nothing
        next_legal_actions_mask = nothing
    end

    consecutive_rewards = consecutive_view(t[:reward], inds; n_horizon = h)
    consecutive_terminals = consecutive_view(t[:terminal], inds; n_horizon = h)
    rewards, terminals = zeros(Float32, n), fill(false, n)

    # make sure that we only consider experiences in current episode
    for i in 1:n
        m = findfirst(view(consecutive_terminals, :, i))
        if isnothing(m)
            terminals[i] = false
            rewards[i] = discount_rewards_reduced(view(consecutive_rewards, :, i), γ)
        else
            terminals[i] = true
            rewards[i] = discount_rewards_reduced(view(consecutive_rewards, 1:m, i), γ)
        end
    end
    (
        states = states,
        legal_actions_mask = legal_actions_mask,
        actions = actions,
        rewards = rewards,
        terminals = terminals,
        next_states = next_states,
        next_legal_actions_mask = next_legal_actions_mask,
    )
end


function RLCore.Experiment(
    ::Val{:rlpyt},
    ::Val{:PPO},
    ::Val{:Atari},
    name::AbstractString;
    save_dir = nothing,
    seed = 123,
)
    @warn "Currently setting the `seed` will not guarantee the reproducibility. The instability seems to be caused by the `CrossCor` layer when calculating gradient."
    rng = StableRNG(seed)
    if isnothing(save_dir)
        t = Dates.format(now(), "yyyy_mm_dd_HH_MM_SS")
        save_dir = joinpath(pwd(), "checkpoints", "rlpyt_PPO_Atari_$(name)_$(t)")
    end

    lg = TBLogger(joinpath(save_dir, "tb_log"), min_level = Logging.Info)

    N_ENV = 32
    UPDATE_FREQ = 64
    N_FRAMES = 4
    STATE_SIZE = (80, 104)
    env = MultiThreadEnv([
        atari_env_factory(
            name,
            STATE_SIZE,
            N_FRAMES;
            repeat_action_probability = 0,
            seed = seed + i,
        ) for i in 1:N_ENV
    ])
    N_ACTIONS = length(action_space(env[1]))
    INIT_CLIP_RANGE = 0.1f0
    INIT_LEARNING_RATE = 1e-3

    init = orthogonal(rng)

    # share model
    model = Chain(
        x -> x ./ 255,
        CrossCor((8, 8), N_FRAMES => 32, relu; stride = 4, pad = 0, init = init),
        CrossCor((4, 4), 32 => 64, relu; stride = 2, pad = 1, init = init),
        x -> reshape(x, :, size(x)[end]),
        Dense(6912, 512, relu; initW = init),
    )

    agent = Agent(
        policy = RandomStartPolicy(
            num_rand_start = 1000,
            random_policy = RandomPolicy(action_space(env); rng = rng),
            policy = PPOPolicy(
                approximator = ActorCritic(
                    actor = Chain(model, Dense(512, N_ACTIONS; initW = init)),
                    critic = Chain(model, Dense(512, 1; initW = init)),
                    optimizer = ADAM(INIT_LEARNING_RATE),  # decrease learning rate with a hook
                ) |> gpu,
                γ = 0.99f0,
                λ = 0.98f0,
                clip_range = INIT_CLIP_RANGE,  # decrease with a hook
                max_grad_norm = 1.0f0,
                n_microbatches = 4,
                n_epochs = 4,
                actor_loss_weight = 1.0f0,
                critic_loss_weight = 0.5f0,
                entropy_loss_weight = 0.01f0,
                rng = rng,
                update_freq = UPDATE_FREQ,
            ),
        ),
        trajectory = PPOTrajectory(;
            capacity = UPDATE_FREQ,
            state = Array{Float32,4} => (STATE_SIZE..., N_FRAMES, N_ENV),
            action = Vector{Int} => (N_ENV,),
            reward = Vector{Float32} => (N_ENV,),
            terminal = Vector{Bool} => (N_ENV,),
            action_log_prob = Vector{Float32} => (N_ENV,),
        ),
    )

    N_TRAINING_STEPS = 50_000_000 ÷ N_ENV
    EVALUATION_FREQ = N_TRAINING_STEPS ÷ 100
    MAX_EPISODE_STEPS_EVAL = 27_000
    N_CHECKPOINTS = 3
    stop_condition = StopAfterStep(N_TRAINING_STEPS)

    total_batch_reward_per_episode = TotalBatchRewardPerEpisode(N_ENV)
    batch_steps_per_episode = BatchStepsPerEpisode(N_ENV)
    evaluation_result = []

    hook = ComposedHook(
        total_batch_reward_per_episode,
        batch_steps_per_episode,
        DoEveryNStep(UPDATE_FREQ) do t, agent, env
            p = agent.policy.policy
            with_logger(lg) do
                @info "training" loss = mean(p.loss) actor_loss = mean(p.actor_loss) critic_loss =
                    mean(p.critic_loss) entropy_loss = mean(p.entropy_loss) norm =
                    mean(p.norm) log_step_increment = UPDATE_FREQ
            end
        end,
        DoEveryNStep(UPDATE_FREQ) do t, agent, env
            decay = (N_TRAINING_STEPS - t) / N_TRAINING_STEPS
            agent.policy.policy.approximator.optimizer.eta = INIT_LEARNING_RATE * decay
            agent.policy.policy.clip_range = INIT_CLIP_RANGE * Float32(decay)
        end,
        DoEveryNStep() do t, agent, env
            with_logger(lg) do
                rewards = [
                    total_batch_reward_per_episode.rewards[i][end]
                    for i in 1:length(env) if is_terminated(env[i])
                ]
                if length(rewards) > 0
                    @info "training" rewards = mean(rewards) log_step_increment = 0
                end
                steps = [
                    batch_steps_per_episode.steps[i][end]
                    for i in 1:length(env) if is_terminated(env[i])
                ]
                if length(steps) > 0
                    @info "training" steps = mean(steps) log_step_increment = 0
                end
            end
        end,
        DoEveryNStep(EVALUATION_FREQ) do t, agent, env
            @info "evaluating agent at $t step..."
            # switch to GreedyExplorer?
            h = TotalBatchRewardPerEpisode(N_ENV)
            s = @elapsed run(
                agent.policy,
                MultiThreadEnv([
                    atari_env_factory(
                        name,
                        STATE_SIZE,
                        N_FRAMES,
                        MAX_EPISODE_STEPS_EVAL;
                        repeat_action_probability = 0,
                        seed = seed + t + i,
                    ) for i in 1:4
                ]),
                StopAfterStep(27_000; is_show_progress = false),
                h,
            )
            res = (avg_score = mean(Iterators.flatten(h.rewards)),)
            push!(evaluation_result, res)

            @info "finished evaluating agent in $s seconds" avg_score = res.avg_score
            with_logger(lg) do
                @info "evaluating" avg_score = res.avg_score log_step_increment = 0
            end

            policy = cpu(agent.policy)
            mkdir(joinpath(save_dir, string(t)))
            BSON.@save joinpath(save_dir, string(t), "policy.bson") policy
            BSON.@save joinpath(save_dir, string(t), "stats.bson") total_batch_reward_per_episode evaluation_result

            # only keep recent 3 checkpoints
            old_checkpoint_folder =
                joinpath(save_dir, string(t - EVALUATION_FREQ * N_CHECKPOINTS))
            if isdir(old_checkpoint_folder)
                rm(old_checkpoint_folder; force = true, recursive = true)
            end
        end,
    )

    description = """
    # Play Atari($name) with PPO
    The configurations use here should be aligned with [atari_ff_ppo.py](https://github.com/astooke/rlpyt/blob/master/rlpyt/experiments/configs/atari/pg/atari_ff_ppo.py) in rlpyt.
    You can also view the tensorboard logs with `tensorboard --logdir $(joinpath(save_dir, "tb_log"))`
    """

    Experiment(agent, env, stop_condition, hook, description)
end