text:
################################################################################
#                       MEASURE VARIABLE CREATION METHODS
################################################################################
"""
    make_point_variable_ref(write_model::Union{InfiniteModel, JuMP.Model},
                            ivref::GeneralVariableRef,
                            support::Vector{Float64}
                            )::GeneralVariableRef

Make a point variable for infinite variable/derivative `ivref` at `support`, add it to
the `write_model`, and return the `GeneralVariableRef`. This is an internal method
for point variables produced by expanding measures via [`expand_measure`](@ref).
This is also useful for those writing extension optimizer models and wish to
expand measures without modifiying the `InfiniteModel`. In such cases, `write_model`
should be the optimizer model and [`add_measure_variable`](@ref add_measure_variable(::JuMP.Model, ::Any, ::Any))
should be extended appropriately for point variables. Errors if `write_model` is
an optimizer model and `add_measure_variable` is not properly extended.

Note this is also accomodates infinite parameter functions, in which case the
infinite parameter function is called with the support as input.
"""
function make_point_variable_ref(write_model::InfiniteModel,
                                 ivref::GeneralVariableRef,
                                 support::Vector{Float64})
    return make_point_variable_ref(write_model, ivref, support, _index_type(ivref))
end

# Infinite variable index
function make_point_variable_ref(write_model::InfiniteModel, ivref, support,
    ::Union{Type{InfiniteVariableIndex}, Type{DerivativeIndex}}
    )::GeneralVariableRef
    prefs = parameter_list(ivref)
    for i in eachindex(support)
        support[i] = round(support[i], sigdigits = significant_digits(prefs[i]))
    end
    base_info = JuMP.VariableInfo(false, NaN, false, NaN, false, NaN, false,
                                  NaN, false, false)
    new_info = _update_point_info(base_info, dispatch_variable_ref(ivref), support)
    var = PointVariable(_make_float_info(new_info), ivref, support)
    return JuMP.add_variable(write_model, var; add_support = false)
end

# Infinite parameter function index (this works for both optimizer and infinite models)
function make_point_variable_ref(write_model, fref, support,
    ::Type{ParameterFunctionIndex}
    )::Float64
    prefs = raw_parameter_refs(fref)
    supp_tuple = Tuple(support, prefs)
    func = raw_function(fref)
    return func(supp_tuple...)
end

"""
    add_measure_variable(model::JuMP.Model, var,
                         key::Val{:ext_key_name})::GeneralVariableRef

Add a measure variable `var` to the optimizer model `model` (with `key`) and
return the correct `InfiniteOpt` variable reference. This is an internal method
used by [`make_point_variable_ref`](@ref) and [`make_reduced_variable_ref`](@ref)
to make point variables and reduced variables when the `write_model` is an
optimizer model. This is useful for extensions that wish to expand measures, but
without changing the original `InfiniteModel`. Thus, this should be extended for
adding `PointVariable`s and `ReducedVariable`s for such extensions.
Otherwise, an error is thrown for unextended variable and/or optimizer model types.
Note if this is extended, than [`internal_reduced_variable`](@ref) should also
be extended in order to direct reduced variables references to the underlying
[`ReducedVariable`](@ref).
"""
function add_measure_variable(model::JuMP.Model, var, key)
    error("`add_measure_variable` not defined for variable of type `$(typeof(var))` " *
          "and an optimizer model with key `$(typeof(key).parameters[1])`.")
end

# Store/add the variable to the optimizer model via add_measure_variable
# This avoids changing the InfiniteModel unexpectedly
function make_point_variable_ref(write_model::JuMP.Model, # this should be an optimizer model
                                 ivref::GeneralVariableRef,
                                 support::Vector{Float64})
    return make_point_variable_ref(write_model, ivref, support, _index_type(ivref))
end

# Infinite variable index
function make_point_variable_ref(write_model::JuMP.Model, ivref, support,
    ::Union{Type{InfiniteVariableIndex}, Type{DerivativeIndex}}
    )::GeneralVariableRef
    prefs = parameter_list(ivref)
    for i in eachindex(support)
        support[i] = round(support[i], sigdigits = significant_digits(prefs[i]))
    end
    base_info = JuMP.VariableInfo(false, NaN, false, NaN, false, NaN, false,
                                  NaN, false, false)
    new_info = _update_point_info(base_info, dispatch_variable_ref(ivref), support)
    var = PointVariable(_make_float_info(new_info), ivref, support)
    opt_key = optimizer_model_key(write_model)
    return add_measure_variable(write_model, var, Val(opt_key))
end

"""
    make_reduced_variable_ref(write_model::Union{InfiniteModel, JuMP.Model},
                              ivref::GeneralVariableRef,
                              indices::Vector{Int},
                              values::Vector{Float64}
                              )::GeneralVariableRef

Make a reduced variable for infinite variable/derivative/parameter function `ivref` at `support`, add it to
the `write_model`, and return the `GeneralVariableRef`. This is an internal method
for reduced variables produced by expanding measures via [`expand_measure`](@ref).
This is also useful for those writing extension optimizer models and wish to
expand measures without modifiying the `InfiniteModel`. In such cases, `write_model`
should be the optimizer model and [`add_measure_variable`](@ref add_measure_variable(::JuMP.Model, ::Any, ::Any))
should be extended appropriately for reduced variables. Errors if `write_model`
is an optimizer model and `add_measure_variable` is not properly extended.
Note this is only intended for optimizer models that are currently stored in
`InfiniteModel.optimizer_model`.
"""
function make_reduced_variable_ref(write_model::InfiniteModel,
                                   ivref::GeneralVariableRef,
                                   indices::Vector{Int},
                                   values::Vector{Float64}
                                   )::GeneralVariableRef
    eval_supps = Dict(indices[i] => values[i] for i in eachindex(indices))
    var = JuMP.build_variable(error, ivref, eval_supps, check = false)
    return JuMP.add_variable(write_model, var)
end

# Add reduced infinite variables in the optimizer model without modifying the InfiniteModel
function make_reduced_variable_ref(write_model::JuMP.Model,
                                   ivref::GeneralVariableRef,
                                   indices::Vector{Int},
                                   values::Vector{Float64}
                                   )::GeneralVariableRef
    eval_supps = Dict(indices[i] => values[i] for i in eachindex(indices))
    var = JuMP.build_variable(error, ivref, eval_supps, check = false)
    key = optimizer_model_key(write_model)
    return add_measure_variable(write_model, var, Val(key))
end

"""
    delete_internal_reduced_variable(write_model::Union{InfiniteModel, JuMP.Model},
                                     rvref::ReducedVariableRef)::Nothing

Delete the variable associated with `rvref` from `write_model` if it is purely
an internal variable only used for measure expansion and is no longer needed.
For `write_model`s that are an optimizer model,
[`delete_reduced_variable`](@ref delete_reduced_variable(::JuMP.Model, ::Any,::Any))
will need to be extended for this this to work. Otherwise, a warning will be thrown.
Note that this is intended as an internal method to assist with extensions to
[`expand_measure`](@ref).
"""
function delete_internal_reduced_variable(write_model::InfiniteModel,
                                          rvref::ReducedVariableRef)::Nothing
    if !used_by_measure(rvref) && !used_by_constraint(rvref)
        JuMP.delete(write_model, rvref)
    end
    return
end

"""
    delete_reduced_variable(model::JuMP.Model, vref, key::Val{:ext_key_name})::Nothing

Delete the reduced variable associated with `vref` from the optimizer model
`model` with associated extension key `:ext_key_name`. A warning is thrown if this
is not properly extended. This is intended as a helper function for
[`delete_internal_reduced_variable`](@ref) which is used by [`expand_measure`](@ref).
"""
function delete_reduced_variable(write_model::JuMP.Model, vref, key)
    @warn "'delete_reduced_variable' not extended for reduced variable type " *
          "`$(typeof(vref))` and optimizer model with key `$(typeof(key).parameters[1])`."
    return
end

# Delete reduced infinite variable from optimizer model if it was not made by the InfiniteModel
function delete_internal_reduced_variable(write_model::JuMP.Model,
                                          rvref::ReducedVariableRef)::Nothing
    if !JuMP.is_valid(JuMP.owner_model(rvref), rvref)
        key = optimizer_model_key(write_model)
        delete_reduced_variable(write_model, rvref, Val(key))
    end
    return
end

################################################################################
#                          EXPAND_MEASURE DEFINITIONS
################################################################################
"""
    expand_measure(expr, data::AbstractMeasureData,
                   write_model::JuMP.AbstractModel)::JuMP.AbstractJuMPScalar

Return the finite reformulation of a measure containing a variable/parameter
expression `expr` with measure data `data`. Here `write_model` is the target
model where this expanded expression will be used. Thus, any variables that need
to be created will be added to `write_model`. The methods [`make_point_variable_ref`](@ref)
and [`make_reduced_variable_ref`](@ref) should be used as appropriate to create
these variables. Developers might also choose to use
[`delete_internal_reduced_variable`](@ref) in order to remove reduced variables
once they are no longer needed. Note this is intended as an internal function,
but will need to be extended for unsupported `expr` types and for user-defined
measure data types. Principally, this is leveraged to enable the user methods
[`expand`](@ref) and [`expand_all_measures!`](@ref).
"""
function expand_measure end

# GeneralVariableRef
function expand_measure(vref::GeneralVariableRef,
                        data::DiscreteMeasureData,
                        write_model::JuMP.AbstractModel
                        )::Union{JuMP.AbstractJuMPScalar, Float64}
    return expand_measure(vref, _index_type(vref), data, write_model)
end

# InfiniteVariableRef/DerivativeRef/ParameterFunctionRef (1D DiscreteMeasureData)
function expand_measure(ivref::GeneralVariableRef,
                        index_type::Union{Type{InfiniteVariableIndex}, Type{DerivativeIndex}, Type{ParameterFunctionIndex}},
                        data::DiscreteMeasureData{GeneralVariableRef, 1},
                        write_model::JuMP.AbstractModel)
    # pull in the needed information
    var_prefs = parameter_list(ivref)
    pref = parameter_refs(data)
    supps = supports(data)
    coeffs = coefficients(data)
    w = weight_function(data)
    # treat variable as constant if doesn't have measure parameter
    if !(pref in var_prefs)
        var_coef = sum(coeffs[i] * w(supps[i]) for i in eachindex(coeffs))
        return JuMP.GenericAffExpr{Float64, GeneralVariableRef}(0, ivref => var_coef)
    # make point variables if var_prefs = pref (it is the only dependence)
    elseif length(var_prefs) == 1
        return JuMP.@expression(write_model, sum(coeffs[i] * w(supps[i]) *
                    make_point_variable_ref(write_model, ivref, [supps[i]])
                    for i in eachindex(coeffs)))
    # make reduced variables if the variable contains other parameters
    else
        index = [findfirst(isequal(pref), var_prefs)]
        return JuMP.@expression(write_model, sum(coeffs[i] * w(supps[i]) *
                    make_reduced_variable_ref(write_model, ivref, index, [supps[i]])
                    for i in eachindex(coeffs)))
    end
end

# InfiniteVariableRef/DerivativeRef/ParameterFunctionRef (Multi DiscreteMeasureData)
function expand_measure(ivref::GeneralVariableRef,
                        index_type::Union{Type{InfiniteVariableIndex}, Type{DerivativeIndex}, Type{ParameterFunctionIndex}},
                        data::DiscreteMeasureData{Vector{GeneralVariableRef}, 2},
                        write_model::JuMP.AbstractModel)
    # pull in the needed information
    var_prefs = parameter_list(ivref)
    prefs = parameter_refs(data)
    supps = supports(data)
    coeffs = coefficients(data)
    w = weight_function(data)
    # var_prefs == prefs so let's make a point variable
    if var_prefs == prefs
        return JuMP.@expression(write_model, sum(coeffs[i] * w(supps[:, i]) *
                    make_point_variable_ref(write_model, ivref, supps[:, i])
                    for i in eachindex(coeffs)))
    # treat variable as constant if doesn't have measure parameter
    elseif !any(pref in var_prefs for pref in prefs)
        var_coef = sum(coeffs[i] * w(supps[:, i]) for i in eachindex(coeffs))
        return JuMP.GenericAffExpr{Float64, GeneralVariableRef}(0, ivref => var_coef)
    # make point variables if all var_prefs are contained in prefs
    elseif all(pref in prefs for pref in var_prefs)
        indices = [findfirst(isequal(pref), prefs) for pref in var_prefs]
        new_supps = supps[indices, :]
        return JuMP.@expression(write_model, sum(coeffs[i] * w(supps[:, i]) *
                    make_point_variable_ref(write_model, ivref, new_supps[:, i])
                    for i in eachindex(coeffs)))
    # make reduced variables if the variable contains other parameters
    else
        # get indices of each pref to map properly
        indices = [findfirst(isequal(pref), var_prefs) for pref in prefs]
        # check that if any of the indices are empty and truncate as needed
        empty = map(i -> i === nothing, indices)
        if any(empty)
            indices = convert(Vector{Int}, deleteat!(indices, empty))
            supps = supps[.!empty, :]
        end
        return JuMP.@expression(write_model, sum(coeffs[i] * w(supps[:, i]) *
                    make_reduced_variable_ref(write_model, ivref, indices, supps[:, i])
                    for i in eachindex(coeffs)))
    end
end

# Write point support given reduced info and the support
function _make_point_support(orig_prefs::Vector{GeneralVariableRef},
                             support_dict::Dict{Int, Float64},
                             index::Int, value::Float64)::Vector{Float64}
    return [i == index ? value : support_dict[i] for i in eachindex(orig_prefs)]
end

# ReducedVariableRef (1D DiscreteMeasureData)
function expand_measure(rvref::GeneralVariableRef,
                        index_type::Type{ReducedVariableIndex},
                        data::DiscreteMeasureData{GeneralVariableRef, 1},
                        write_model::JuMP.AbstractModel)
    # pull in the needed information
    drvref = dispatch_variable_ref(rvref)
    ivref = infinite_variable_ref(drvref)
    orig_prefs = parameter_list(ivref)
    var_prefs = parameter_list(drvref)
    eval_supps = eval_supports(drvref)
    pref = parameter_refs(data)
    supps = supports(data)
    coeffs = coefficients(data)
    w = weight_function(data)
    # treat variable as constant if doesn't have measure parameter
    if !(pref in var_prefs)
        var_coef = sum(coeffs[i] * w(supps[i]) for i in eachindex(coeffs))
        expr = JuMP.GenericAffExpr{Float64, GeneralVariableRef}(0, rvref => var_coef)
    # make point variables if var_prefs = pref (it is the only dependence)
    elseif length(var_prefs) == 1
        index = findfirst(isequal(pref), orig_prefs)
        expr = JuMP.@expression(write_model, sum(coeffs[i] * w(supps[i]) *
                    make_point_variable_ref(write_model, ivref,
                    _make_point_support(orig_prefs, eval_supps, index, supps[i]))
                    for i in eachindex(coeffs)))
        delete_internal_reduced_variable(write_model, drvref) # TODO not sure this is helpful
    # make reduced variables if the variable contains other parameters
    else
        index = findfirst(isequal(pref), orig_prefs)
        collected_indices = collect(keys(eval_supps))
        vals = map(k -> eval_supps[k], collected_indices) # a support will be appended on the fly
        indices = push!(collected_indices, index)
        expr = JuMP.@expression(write_model, sum(coeffs[i] * w(supps[i]) *
                    make_reduced_variable_ref(write_model, ivref, indices, vcat(vals, supps[i]))
                    for i in eachindex(coeffs)))
        delete_internal_reduced_variable(write_model, drvref) # TODO not sure this is helpful
    end
    return expr
end

# Write point support given reduced info and the support
function _make_point_support(orig_prefs::Vector{GeneralVariableRef},
                             support_dict::Dict{Int, Float64},
                             new_indices::Vector{Int},
                             values::Vector{Float64})::Vector{Float64}
    # these might overlap with the old dict, so we favor the old dict
    new_dict = Dict(new_indices[i] => values[i] for i in eachindex(values))
    return [haskey(support_dict, i) ? support_dict[i] : new_dict[i]
            for i in eachindex(orig_prefs)]
end

# ReducedVariableRef (Multi DiscreteMeasureData)
function expand_measure(rvref::GeneralVariableRef,
                        index_type::Type{ReducedVariableIndex},
                        data::DiscreteMeasureData{Vector{GeneralVariableRef}, 2},
                        write_model::JuMP.AbstractModel)
    # pull in the needed information
    drvref = dispatch_variable_ref(rvref)
    ivref = infinite_variable_ref(drvref)
    orig_prefs = parameter_list(ivref)
    var_prefs = parameter_list(drvref)
    eval_supps = eval_supports(drvref)
    prefs = parameter_refs(data)
    supps = supports(data)
    coeffs = coefficients(data)
    w = weight_function(data)
    # treat variable as constant if doesn't have measure parameter
    if !any(pref in var_prefs for pref in prefs)
        var_coef = sum(coeffs[i] * w(supps[:, i]) for i in eachindex(coeffs))
        expr = JuMP.GenericAffExpr{Float64, GeneralVariableRef}(0, rvref => var_coef)
    # make point variables if prefs includes all var_prefs
    elseif all(pref in prefs for pref in var_prefs)
        # get the indices of measure prefs to reorder/truncate as needed
        supp_indices = [findfirst(isequal(pref), prefs) for pref in var_prefs]
        # reorder/truncate if necesary
        if supp_indices != 1:size(supps, 1)
            supps = supps[supp_indices, :]
        end
        # get the parameter indices of the variable parameters to be reduced
        indices = [findfirst(isequal(pref), orig_prefs) for pref in var_prefs]
        # make the expression
        expr = JuMP.@expression(write_model, sum(coeffs[i] * w(supps[:, i]) *
                    make_point_variable_ref(write_model, ivref,
                    _make_point_support(orig_prefs, eval_supps, indices, supps[:, i]))
                    for i in eachindex(coeffs)))
        delete_internal_reduced_variable(write_model, drvref) # TODO not sure this is helpful
    # make reduced variables if the variable contains other parameters
    else
        # get the indices of prefs in terms of the ivref
        new_indices = [findfirst(isequal(pref), orig_prefs) for pref in prefs]
        # check that if any of the indices are empty or already reduced and truncate as needed
        bad_index = map(i -> i === nothing || i in keys(eval_supps), new_indices)
        if any(bad_index)
            deleteat!(new_indices, bad_index)
            supps = supps[.!bad_index, :]
        end
        # prepare the indices and values for reduced variable construction
        collected_indices = collect(keys(eval_supps))
        vals = map(k -> eval_supps[k], collected_indices) # a support will be appended on the fly
        indices = append!(collected_indices, new_indices)
        # make the expression
        expr = JuMP.@expression(write_model, sum(coeffs[i] * w(supps[:, i]) *
                    make_reduced_variable_ref(write_model, ivref, indices, vcat(vals, supps[:, i]))
                    for i in eachindex(coeffs)))
        delete_internal_reduced_variable(write_model, drvref) # TODO not sure this is helpful
    end
    return expr
end

# FiniteVariableRef (1D DiscreteMeasureData)
function expand_measure(vref::GeneralVariableRef,
                        index_type::Type{V},
                        data::DiscreteMeasureData{GeneralVariableRef, 1},
                        write_model::JuMP.AbstractModel
                        )::JuMP.GenericAffExpr where {V <: FiniteVariableIndex}
    # pull in the needed information
    supps = supports(data)
    coeffs = coefficients(data)
    w = weight_function(data)
    # treat the variable as a constant and build the expression
    var_coef = sum(coeffs[i] * w(supps[i]) for i in eachindex(coeffs))
    return JuMP.GenericAffExpr{Float64, GeneralVariableRef}(0, vref => var_coef)
end

# FiniteVariableRef (Multi DiscreteMeasureData)
function expand_measure(vref::GeneralVariableRef,
                        index_type::Type{V},
                        data::DiscreteMeasureData{Vector{GeneralVariableRef}, 2},
                        write_model::JuMP.AbstractModel
                        )::JuMP.GenericAffExpr where {V <: FiniteVariableIndex}
    # pull in the needed information
    supps = supports(data)
    coeffs = coefficients(data)
    w = weight_function(data)
    # treat the variable as a constant and build the expression
    var_coef = sum(coeffs[i] * w(supps[:, i]) for i in eachindex(coeffs))
    return JuMP.GenericAffExpr{Float64, GeneralVariableRef}(0, vref => var_coef)
end

# InfiniteParameterRef (1D DiscreteMeasureData)
function expand_measure(pref::GeneralVariableRef,
                        index_type::Type{P},
                        data::DiscreteMeasureData{GeneralVariableRef, 1},
                        write_model::JuMP.AbstractModel
                        )::Union{JuMP.GenericAffExpr, Float64} where {P <: InfiniteParameterIndex}
    # pull in the needed information
    meas_pref = parameter_refs(data)
    supps = supports(data)
    coeffs = coefficients(data)
    w = weight_function(data)
    # treat the parameter
    if meas_pref != pref
        par_coef = sum(coeffs[i] * w(supps[i]) for i in eachindex(coeffs))
        return JuMP.GenericAffExpr{Float64, GeneralVariableRef}(0, pref => par_coef)
    # replace the parameter with its value
    else
        return sum(coeffs[i] * w(supps[i]) * supps[i] for i in eachindex(coeffs))
    end
end

# InfiniteParameterRef (Multi DiscreteMeasureData)
function expand_measure(pref::GeneralVariableRef,
                        index_type::Type{P},
                        data::DiscreteMeasureData{Vector{GeneralVariableRef}, 2},
                        write_model::JuMP.AbstractModel
                        )::Union{JuMP.GenericAffExpr, Float64} where {P <: InfiniteParameterIndex}
    # pull in the needed information
    prefs = parameter_refs(data)
    supps = supports(data)
    coeffs = coefficients(data)
    w = weight_function(data)
    # find the position of pref if it is in the data
    index = findfirst(isequal(pref), prefs)
    # treat the parameter as a constant
    if index === nothing
        par_coef = sum(coeffs[i] * w(supps[:, i]) for i in eachindex(coeffs))
        return JuMP.GenericAffExpr{Float64, GeneralVariableRef}(0, pref => par_coef)
    # replace the parameter with its value
    else
        return sum(coeffs[i] * w(supps[:, i]) * supps[index, i] for i in eachindex(coeffs))
    end
end

# GenericAffExpr (1D DiscreteMeasureData)
function expand_measure(expr::JuMP.GenericAffExpr{C, GeneralVariableRef},
                        data::DiscreteMeasureData{GeneralVariableRef, 1},
                        write_model::JuMP.AbstractModel
                        )::Union{JuMP.AbstractJuMPScalar, Float64} where {C}
    # pull in the needed information
    supps = supports(data)
    coeffs = coefficients(data)
    w = weight_function(data)
    # expand each variable independently and add all together
    constant_coef = sum(coeffs[i] * w(supps[i]) for i in eachindex(coeffs))
    return JuMP.@expression(write_model, sum(coef * expand_measure(var, data, write_model)
                for (var, coef) in expr.terms) + expr.constant * constant_coef)
end

# GenericAffExpr (Multi DiscreteMeasureData)
function expand_measure(expr::JuMP.GenericAffExpr{C, GeneralVariableRef},
                        data::DiscreteMeasureData{Vector{GeneralVariableRef}, 2},
                        write_model::JuMP.AbstractModel
                        )::Union{JuMP.AbstractJuMPScalar, Float64} where {C}
    # pull in the needed information
    supps = supports(data)
    coeffs = coefficients(data)
    w = weight_function(data)
    # expand each variable independently and add all together
    constant_coef = sum(coeffs[i] * w(supps[:, i]) for i in eachindex(coeffs))
    return JuMP.@expression(write_model, sum(coef * expand_measure(var, data, write_model)
                for (var, coef) in expr.terms) + expr.constant * constant_coef)
end

# GenericQuadExpr (1D DiscreteMeasureData)
function expand_measure(expr::JuMP.GenericQuadExpr{C, GeneralVariableRef},
                        data::DiscreteMeasureData{GeneralVariableRef, 1},
                        write_model::JuMP.AbstractModel
                        )::Union{JuMP.AbstractJuMPScalar, Float64} where {C}
    # get needed info
    pref = parameter_refs(data)
    supps = supports(data)
    coeffs = coefficients(data)
    w = weight_function(data)
    label = support_label(data)
    lb = JuMP.lower_bound(data)
    ub = JuMP.upper_bound(data)
    is_expect = _is_expect(data)
    # expand the affine expression
    new_expr = expand_measure(expr.aff, data, write_model)
    # make viable data objects so we can multiply the terms
    coef_data = DiscreteMeasureData(pref, ones(1), ones(1), label, w, lb, ub,
                                    is_expect)
    simple_data = DiscreteMeasureData(pref, ones(1), ones(1), label,
                                      default_weight, lb, ub, is_expect)
    # expand the quadratic terms
    for i in eachindex(coeffs)
        # update the temp data
        coefficients(coef_data)[1] = coeffs[i]
        supports(coef_data)[1] = supps[i]
        supports(simple_data)[1] = supps[i]
        new_expr = JuMP.@expression(write_model, new_expr + sum(coef *
                        expand_measure(pair.a, coef_data, write_model) *
                        expand_measure(pair.b, simple_data, write_model)
                        for (pair, coef) in expr.terms))
    end
    return new_expr
end

# GenericQuadExpr(Multi DiscreteMeasureData)
function expand_measure(expr::JuMP.GenericQuadExpr{C, GeneralVariableRef},
                        data::DiscreteMeasureData{Vector{GeneralVariableRef}, 2},
                        write_model::JuMP.AbstractModel
                        )::Union{JuMP.AbstractJuMPScalar, Float64} where {C}
    # get needed info
    prefs = parameter_refs(data)
    supps = supports(data)
    coeffs = coefficients(data)
    w = weight_function(data)
    label = support_label(data)
    lbs = JuMP.lower_bound(data)
    ubs = JuMP.upper_bound(data)
    is_expect = _is_expect(data)
    # expand the affine expression
    new_expr = expand_measure(expr.aff, data, write_model)
    # make viable data objects so we can multiply the terms
    coef_data = DiscreteMeasureData(prefs, ones(1), ones(length(prefs), 1),
                                    label, w, lbs, ubs, is_expect)
    simple_data = DiscreteMeasureData(prefs, ones(1), ones(length(prefs), 1),
                                      label, default_weight, lbs, ubs, is_expect)
    # expand the quadratic terms
    for i in eachindex(coeffs)
        # update the temp data
        coefficients(coef_data)[1] = coeffs[i]
        supports(coef_data)[:, 1] = @view(supps[:, i])
        supports(simple_data)[:, 1] = @view(supps[:, i])
        new_expr = JuMP.@expression(write_model, new_expr + sum(coef *
                        expand_measure(pair.a, coef_data, write_model) *
                        expand_measure(pair.b, simple_data, write_model)
                        for (pair, coef) in expr.terms))
    end
    return new_expr
end

# MeasureRef
function expand_measure(mref::GeneralVariableRef,
                        index_type::Type{MeasureIndex},
                        data::DiscreteMeasureData,
                        write_model::JuMP.AbstractModel
                        )::Union{JuMP.AbstractJuMPScalar, Float64}
    # determine function and data of the inner measure
    deeper_func = measure_function(mref)
    deeper_data = measure_data(mref)
    # expand the inner measure (note this is recursive for nested measures)
    new_func = expand_measure(deeper_func, deeper_data, write_model)
    # expand current level with the inner measure now expanded
    return expand_measure(new_func, data, write_model)
end

# Call add_generative_supports if needed
function _prep_generative_supps(prefs, info_type::Type{NoGenerativeSupports})::Nothing
    return
end
function _prep_generative_supps(pref, info_type)::Nothing
    add_generative_supports(pref)
    return
end

# FunctionalDiscreteMeasureData
function expand_measure(expr,
    data::FunctionalDiscreteMeasureData{P, B, I},
    write_model::JuMP.AbstractModel
    )::Union{JuMP.AbstractJuMPScalar, Float64} where {P, B, I}
    # get the info
    prefs = parameter_refs(data)
    _prep_generative_supps(prefs, I)
    supps = supports(data)
    coeffs = coefficients(data)
    label = support_label(data)
    w = weight_function(data)
    lbs = JuMP.lower_bound(data)
    ubs = JuMP.upper_bound(data)
    is_expect = _is_expect(data)
    # prepare the DiscreteMeasureData and dispatch
    new_data = DiscreteMeasureData(prefs, coeffs, supps, label, w, lbs, ubs,
                                   is_expect)
    return expand_measure(expr, new_data, write_model)
end

# Catch all method for undefined behavior
function expand_measure(expr, data::AbstractMeasureData,
                        write_model::JuMP.AbstractModel)
    expr_type = typeof(expr)
    data_type = typeof(data)
    error("Undefined behavior to expand expression of type `$expr_type` with " *
          "measure data `$data_type`. If this functionality is needed consider " *
          "extending `expand_measure`.")
end

################################################################################
#                          ANALYTIC EXPANSION METHODS
################################################################################
"""
    analytic_expansion(expr, data::AbstractMeasureData,
                       write_model::JuMP.AbstractModel)::JuMP.AbstractJuMPScalar

Analytically, evaluate measure in the simple case where the measure expression
`expr` doesn't depend on `data` and thus `expr` can be treated as a constant in
conjunction with an analytic result of the `data`. This is intended as an internal
method that is used by [`expand`](@ref) and [`expand_measures`](@ref). For
unrecognized `data` types, `expand_measure` is called instead. User defined
measure data type may choose to extend this method if desired. This is triggered
when `is_analytic(mref) = true`.
"""
function analytic_expansion end

# 1D DiscreteMeasureData/FunctionalDiscreteMeasureData
function analytic_expansion(expr::JuMP.AbstractJuMPScalar,
                            data::Union{DiscreteMeasureData{GeneralVariableRef, 1},
                                        FunctionalDiscreteMeasureData{GeneralVariableRef}},
                            write_model::JuMP.AbstractModel # needed for fallback
                            )::JuMP.AbstractJuMPScalar
    # get the bounds and expect
    lb = JuMP.lower_bound(data)
    ub = JuMP.upper_bound(data)
    is_expect = _is_expect(data)
    w = weight_function(data)
    # build the expression
    if w == default_weight && !isnan(lb) && !isnan(ub)
        return (ub - lb) * expr
    elseif w == default_weight && is_expect
        return expr # this is an expectation and should be equal to 1 * expr
    else
        supps = supports(data)
        coeffs = coefficients(data)
        constant_coef = sum(coeffs[i] * w(supps[i]) for i in eachindex(coeffs))
        return constant_coef * expr
    end
end

# Multi DiscreteMeasureData/FunctionalDiscreteMeasureData
function analytic_expansion(expr::JuMP.AbstractJuMPScalar,
                            data::Union{DiscreteMeasureData{Vector{GeneralVariableRef}, 2},
                                        FunctionalDiscreteMeasureData{Vector{GeneralVariableRef}}},
                            write_model::JuMP.AbstractModel # needed for fallback
                            )::JuMP.AbstractJuMPScalar
    # get the bounds and expect
    lbs = JuMP.lower_bound(data)
    ubs = JuMP.upper_bound(data)
    is_expect = _is_expect(data)
    w = weight_function(data)
    # build the expression
    if w == default_weight && !isnan(first(lbs)) && !isnan(first(ubs))
        return prod(ubs .- lbs) * expr
    elseif w == default_weight && is_expect
        return expr # this is an expectation and should be equal to 1 * expr
    else
        supps = supports(data)
        coeffs = coefficients(data)
        constant_coef = sum(coeffs[i] * w(supps[:, i]) for i in eachindex(coeffs))
        return constant_coef * expr
    end
end

# Fallback
function analytic_expansion(expr, data::AbstractMeasureData,
                            write_model::JuMP.AbstractModel
                            )::JuMP.AbstractJuMPScalar
    return expand_measure(expr, data, write_model)
end

################################################################################
#                              EXPANSION METHODS
################################################################################
"""
    expand(mref::MeasureRef)::JuMP.AbstractJuMPScalar

Return a JuMP scalar function containing the explicit expansion of the measure
`mref`. This expansion is done according to the measure data. Note that
variables are added to the model as necessary to accomodate the expansion (i.e.,
point variables and reduced infinite variables are made as needed). Errors if
expansion is undefined for the measure data and/or the measure expression. If
desired this can be used in combination with [`measure`](@ref) to expand measures
on the fly.

This is useful for extensions that employ a custom optimizer model since it
can be used evaluate measures before expressions are translated to the new model.
This method can also be extended to handle custom measure data types by extending
[`expand_measure`](@ref). Optionally, [`analytic_expansion`](@ref) can also
be extended which is triggered by [`is_analytic`](@ref) for such types if
analytic expansion is possible in certain cases.

**Example**
```julia-repl
julia> tdata = DiscreteMeasureData(t, [0.5, 0.5], [0, 1])

julia> expr = expand(measure(g + z + T - h - 2, tdata))
0.5 g(0) + 0.5 g(1) + z + 0.5 T(0, x) + 0.5 T(1, x) - h(x) - 2
```
"""
function expand(mref::MeasureRef)::Union{JuMP.AbstractJuMPScalar, Float64}
    if is_analytic(mref)
        return analytic_expansion(measure_function(mref), measure_data(mref),
                                  JuMP.owner_model(mref))
    else
        return expand_measure(measure_function(mref), measure_data(mref),
                          JuMP.owner_model(mref))
    end
end

"""
    expand_measures(expr, write_model::JuMP.AbstractModel)::JuMP.AbstractJuMPScalar

Expand all `MeasureRef`s in `expr` in-place via [`expand_measure`](@ref) and
return the expanded expression. This is an internal method used by
[`expand_all_measures!`](@ref) and `TranscriptionOpt` but can be useful for
user-defined optimizer model extensions that add implement
[`add_measure_variable`](@ref) in combination with `expand_measure`. `write_model`
is the model that the measure variables are added to as described in
[`expand_measure`](@ref).
"""
function expand_measures end

# MeasureRef
function expand_measures(mref::GeneralVariableRef,
                         index_type::Type{MeasureIndex},
                         write_model::JuMP.AbstractModel
                         )::Union{JuMP.AbstractJuMPScalar, Float64}
    if is_analytic(mref)
        return analytic_expansion(measure_function(mref), measure_data(mref),
                                  write_model)
    else
        return expand_measure(measure_function(mref), measure_data(mref),
                              write_model)
    end
end

# NonMeasureRef
function expand_measures(vref::GeneralVariableRef,
                         index_type::Type{V},
                         write_model::JuMP.AbstractModel
                         )::GeneralVariableRef where {V <: AbstractInfOptIndex}
    return vref
end

# GeneralVariableRef
function expand_measures(vref::GeneralVariableRef,
                         write_model::JuMP.AbstractModel
                         )::Union{JuMP.AbstractJuMPScalar, Float64}
    return expand_measures(vref, _index_type(vref), write_model)
end

# GenericAffExpr
function expand_measures(expr::JuMP.GenericAffExpr{C, GeneralVariableRef},
                         write_model::JuMP.AbstractModel
                         )::JuMP.AbstractJuMPScalar where {C}
    return JuMP.@expression(write_model, sum(coef * expand_measures(var, write_model)
                            for (var, coef) in expr.terms) + expr.constant)
end

# GenericQuadExpr
function expand_measures(expr::JuMP.GenericQuadExpr{C, GeneralVariableRef},
                         write_model::JuMP.AbstractModel
                         )::JuMP.AbstractJuMPScalar where {C}
    return JuMP.@expression(write_model, sum(coef * expand_measures(pair.a, write_model) *
                expand_measures(pair.b, write_model) for (pair, coef) in expr.terms) +
                expand_measures(expr.aff, write_model))
end

# Fallback
function expand_measures(expr, write_model::JuMP.AbstractModel)
    error("`expand_measures` not defined for expressions of type `$(typeof(expr))`.")
end

"""
    expand_all_measures!(model::InfiniteModel)::Nothing

Expand all of the measures used in the objective and/or constraints of `model`.
The objective and constraints are updated accordingly. Note that
variables are added to the model as necessary to accomodate the expansion (i.e.,
point variables and reduced infinite variables are made as needed). Errors if
expansion is undefined for the measure data and/or the measure expression.

This is useful for extensions that employ a custom optimizer model since it
can be used evaluate measures before `model` is translated into the new model.
This method can also be extended to handle custom measure data types by extending
[`expand_measure`](@ref). Note that this method leverages `expand_measure` via
[`expand_measures`](@ref). Optionally, [`analytic_expansion`](@ref) can also
be extended which is triggered by [`is_analytic`](@ref) for such types if
analytic expansion is possible in certain cases.

**Example**
```julia-repl
julia> print(model)
Min integral{t ∈ [0, 6]}[g(t)*t] + z
Subject to
 T(t, x) ≥ 0.0, ∀ t ∈ [0, 6], xi ∈ [-1, 1]
 z ≥ 0.0
 g(t) + z ≥ 42.0, ∀ t ∈ [0, 6]
 integral{t ∈ [0, 6]}[T(t, x)] ≥ 0.0, ∀ x ∈ [-1, 1]

julia> expand_all_measures!(model)

julia> print(model)
Min 3 g(6) + z
Subject to
 T(t, x) ≥ 0.0, ∀ t ∈ [0, 6], xi ∈ [-1, 1]
 z ≥ 0.0
 g(t) + z ≥ 42.0, ∀ t ∈ [0, 6]
 0.5 T(0, x) + 0.5 T(6, xi) ≥ 0.0, ∀ x ∈ [-1, 1]
```
"""
function expand_all_measures!(model::InfiniteModel)::Nothing
    # expand the objective if it contains measures
    if objective_has_measures(model)
        new_obj = expand_measures(JuMP.objective_function(model), model)
        JuMP.set_objective_function(model, new_obj)
    end
    # expand all of the constraints that contain measures
    for (cindex, object) in model.constraints
        if !isempty(object.measure_indices)
            old_constr = object.constraint
            # clear the old dependencies
            old_vrefs = _all_function_variables(JuMP.jump_function(old_constr))
            for vref in old_vrefs
                filter!(e -> e != cindex, _constraint_dependencies(vref))
            end
            # expand the expression
            new_func = expand_measures(JuMP.jump_function(old_constr), model)
            offset = JuMP.constant(new_func)
            JuMP.add_to_expression!(new_func, -offset)
            new_set = MOIU.shift_constant(JuMP.moi_set(old_constr), -offset)
            vrefs = _all_function_variables(new_func)
            # make the new constraint object
            if old_constr isa BoundedScalarConstraint
                orig_bounds = original_parameter_bounds(old_constr)
                new_constr = BoundedScalarConstraint(new_func, new_set,
                                                     copy(orig_bounds),
                                                     orig_bounds)
            else
                new_constr = JuMP.ScalarConstraint(new_func, new_set)
            end
            # update the bounds if there are bounded hold variables in the model
            if model.has_hold_bounds
                new_constr = _check_and_update_bounds(model, new_constr, vrefs)
            end
            # update the constraint data
            cref = _temp_constraint_ref(model, cindex)
            _set_core_constraint_object(cref, new_constr)
            empty!(object.measure_indices)
            _update_var_constr_mapping(vrefs, cref)
        end
    end
    return
end


text:
@testset "Catenary" begin
    offset = SVector(1.0, 0.0, 5.0)
    Θ = pi / 4
    trans = inv(Translation(offset) ∘ LinearMap(RotZ(Θ)))
    lmin = -10.0
    lmax = 10.0
    a = 5.0

    cat = Catenary(trans, lmin, lmax, a)

    @test Catenary(Θ, offset[1], offset[2], offset[3], lmin, lmax, a) === cat

    @test cat[0.0] ≈ SVector(1.0, 0.0, 5.0)
    @test cat[sqrt(2.0)] ≈ SVector(2.0, 1.0, 5.2013368939727722)

    @test all(map(isapprox, cat[[0.0, sqrt(2.0)]], [SVector(1.0, 0.0, 5.0), SVector(2.0, 1.0, 5.2013368939727722)]))

    @test Quadratic(cat)[0.0] ≈ cat[-10.0]
    @test Quadratic(cat)[10.0] ≈ cat[0.0]
    @test Quadratic(cat)[20.0] ≈ cat[10.0]

    # Test these are invertible (up to numerical error)
    @test Catenary(Quadratic(cat))[-10.0] ≈ cat[-10.0]
    @test Catenary(Quadratic(cat))[0.0] ≈ cat[0.0]
    @test Catenary(Quadratic(cat))[10.0] ≈ cat[10.0]

    @test Quadratic(Catenary(Quadratic(cat)))[0.0] ≈ Quadratic(cat)[0.0]
    @test Quadratic(Catenary(Quadratic(cat)))[10.0] ≈ Quadratic(cat)[10.0]
    @test Quadratic(Catenary(Quadratic(cat)))[20.0] ≈ Quadratic(cat)[20.0]

    # Test stability for tight, steep wires like stays
    points = [SVector(0.0, 0.0, 10.0), SVector(5.0, 0.0, 5.0 - 0.000001), SVector(10.0, 0.0, 0.0)]
    cat = Catenary(points...)
    @test cat[cat.lmin] ≈ points[1]
    @test cat[(cat.lmin + cat.lmax)/2] ≈ points[2]
    @test cat[cat.lmax] ≈ points[3]

    # Exactly linear
    points = [SVector(0.0, 0.0, 10.0), SVector(5.0, 0.0, 5.0), SVector(10.0, 0.0, 0.0)]
    cat = Catenary(points...)
    @test isapprox(cat[cat.lmin], points[1]; atol = 1e-5)
    @test isapprox(cat[(cat.lmin + cat.lmax)/2], points[2]; atol = 1e-5)
    @test isapprox(cat[cat.lmax], points[3]; atol = 1e-5)

    # Degeneracy - flat
    points = [SVector(0.0, 0.0, 10.0), SVector(10.0, 0.0, 10.0), SVector(10.0, 0.0, 10.0)]
    cat = Catenary(points...)
    @test cat[cat.lmin] ≈ points[1]
    @test cat[cat.lmax] ≈ points[3]

    # Degeneracy - sloped
    points = [SVector(0.0, 0.0, 10.0), SVector(10.0, 0.0, 0.0), SVector(10.0, 0.0, 0.0)]
    cat = Catenary(points...)
    @test cat[cat.lmin] ≈ points[1]
    @test cat[cat.lmax] ≈ points[3]

    # Degeneracy - almost flat
    points = [SVector(0.0, 0.0, 10.0), SVector(10.0, 0.0, nextfloat(10.0)), SVector(10.0, 0.0, nextfloat(10.0))]
    cat = Catenary(points...)
    @test cat[cat.lmin] ≈ points[1]
    @test cat[cat.lmax] ≈ points[3]

    # Test stability for flat wires
    cat = Catenary(AffineMap(RotZ(0.374587), SVector(1.67667,-6.34842,-4.604)),-23.078470200626043,-18.554467003815063,2.067954545961777e11)
    p1 = cat[cat.lmin]
    p2 = cat[(cat.lmin + cat.lmax)/2]
    p3 = cat[cat.lmax]
    cat2 = Catenary(p1, p2, p3)
    @test cat2[cat2.lmin] ≈ p1
    @test cat2[(cat2.lmin + cat2.lmax)/2] ≈ p2
    @test cat2[cat2.lmax] ≈ p3

    cat = Catenary(AffineMap(RotZ(0.296116), SVector(1.17404,-6.46032,-4.395)),159.39653663015895,163.80649017276102,6.345e9)
    p1 = cat[cat.lmin]
    p2 = cat[(cat.lmin + cat.lmax)/2]
    p3 = cat[cat.lmax]
    cat2 = Catenary(p1, p2, p3)
    @test cat2[cat2.lmin] ≈ p1
    @test cat2[(cat2.lmin + cat2.lmax)/2] ≈ p2
    @test cat2[cat2.lmax] ≈ p3

    # TODO test transformations
end

ext: jl


text:
using OrdinaryDiffEq, StochasticDiffEq, Flux, DiffEqFlux

x = Float32[2.; 0.]
tspan = (0.0f0,25.0f0)
dudt = Chain(Dense(2,50,tanh),Dense(50,2))

neural_ode(dudt,x,tspan,Tsit5(),save_everystep=false,save_start=false)
neural_ode(dudt,x,tspan,Tsit5(),saveat=0.1)

neural_msde(dudt,x,[0.1,0.1],tspan,SOSRI(),saveat=0.1)


text:
using Test
import SharedVertexMesh

const STL = joinpath("data", "nist.stl")

mesh = SharedVertexMesh.from(STL)
@test length(mesh.triangles) == 7392
@test length(mesh.vertices) == 3688


text:
using Flux, Test
using Flux.CUDA
using Flux: cpu, gpu
using Statistics: mean
using LinearAlgebra: I, cholesky, Cholesky
using SparseArrays: sparse, SparseMatrixCSC, AbstractSparseArray

@testset "CUDA" begin
  x = randn(5, 5)
  cx = gpu(x)
  @test cx isa CuArray

  @test Flux.onecold(gpu([1.0, 2.0, 3.0])) == 3

  x = Flux.onehotbatch([1, 2, 3], 1:3)
  cx = gpu(x)
  @test cx isa Flux.OneHotMatrix && cx.indices isa CuArray
  @test (cx .+ 1) isa CuArray

  m = Chain(Dense(10, 5, tanh), Dense(5, 2), softmax)
  cm = gpu(m)

  @test all(p isa CuArray for p in params(cm))
  @test cm(gpu(rand(10, 10))) isa CuArray{Float32,2}

  xs = rand(5, 5)
  ys = Flux.onehotbatch(1:5,1:5)
  @test collect(cu(xs) .+ cu(ys)) ≈ collect(xs .+ ys)

  c = gpu(Conv((2,2),3=>4))
  x = gpu(rand(10, 10, 3, 2))
  l = c(gpu(rand(10,10,3,2)))
  @test gradient(x -> sum(c(x)), x)[1] isa CuArray

  c = gpu(CrossCor((2,2),3=>4))
  x = gpu(rand(10, 10, 3, 2))
  l = c(gpu(rand(10,10,3,2)))
  @test gradient(x -> sum(c(x)), x)[1] isa CuArray

end

@testset "onehot gpu" begin
  y = Flux.onehotbatch(ones(3), 1:2) |> gpu;
  @test (repr("text/plain", y); true)

  gA = rand(3, 2) |> gpu;
  @test gradient(A -> sum(A * y), gA)[1] isa CuArray

  # construct from CuArray
  x = [1, 3, 2]
  y = Flux.onehotbatch(x, 0:3)
  y2 = Flux.onehotbatch(x |> gpu, 0:3)
  @test y2.indices isa CuArray
  @test y2 |> cpu == y
end

@testset "onecold gpu" begin
  y = Flux.onehotbatch(ones(3), 1:10) |> gpu;
  l = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j']
  @test Flux.onecold(y) isa CuArray
  @test y[3,:] isa CuArray
  @test Flux.onecold(y, l) == ['a', 'a', 'a']
end

@testset "onehot forward map to broadcast" begin
  oa = OneHotArray(rand(1:10, 5, 5), 10) |> gpu
  @test all(map(identity, oa) .== oa)
  @test all(map(x -> 2 * x, oa) .== 2 .* oa)
end

@testset "restructure gpu" begin
  dudt = Dense(1,1) |> gpu
  p,re = Flux.destructure(dudt)
  foo(x) = sum(re(p)(x))
  @test gradient(foo, cu(rand(1)))[1] isa CuArray
end

@testset "GPU functors" begin
  @testset "Cholesky" begin
    M = 2.0*I(10) |> collect
    Q = cholesky(M)
    Q_gpu = Q |> gpu
    @test Q_gpu isa Cholesky{<:Any,<:CuArray}
    Q_cpu = Q_gpu |> cpu
    @test Q_cpu == cholesky(eltype(Q_gpu).(M))
  end

  @testset "isbits array types" begin
    struct SimpleBits
      field::Int32
    end

    @test gpu((;a=ones(1))).a isa CuVector{Float32}
    @test gpu((;a=['a', 'b', 'c'])).a isa CuVector{Char}
    @test gpu((;a=[SimpleBits(1)])).a isa CuVector{SimpleBits}
  end
end

@testset "gpu(cpu(x)) inside gradient" begin
  a = randn(Float32, 4, 4)
  ca = cu(a)

  # Trivial functions
  @test gradient(x -> sum(abs, gpu(x)), a)[1] isa Matrix
  @test gradient(x -> sum(gpu(x)), a)[1] isa Matrix
  @test_skip gradient(x -> sum(gpu(x)), a')[1] isa Matrix  # sum(::Adjoint{T,CuArray}) makes a Fill
  @test gradient(x -> sum(abs, cpu(x)), ca)[1] isa CuArray
  # This test should really not go through indirections and pull out Fills for efficiency
  # but we forcefully materialise. TODO: remove materialising CuArray here
  @test gradient(x -> sum(cpu(x)), ca)[1] isa CuArray # This involves FillArray, which should be GPU compatible
  @test gradient(x -> sum(cpu(x)), ca')[1] isa LinearAlgebra.Adjoint

  # Even more trivial: no movement
  @test gradient(x -> sum(abs, cpu(x)), a)[1] isa Matrix
  @test gradient(x -> sum(abs, cpu(x)), a')[1] isa LinearAlgebra.Adjoint
  @test gradient(x -> sum(cpu(x)), a)[1] isa typeof(gradient(sum, a)[1]) # FillArray
  @test gradient(x -> sum(abs, gpu(x)), ca)[1] isa CuArray
  @test_skip gradient(x -> sum(abs, gpu(x)), ca')[1] isa CuArray # KernelError: passing and using non-bitstype argument

  # More complicated, Array * CuArray is an error
  g0 = gradient(x -> sum(abs, (a * (a * x))), a)[1]
  @test g0 ≈ gradient(x -> sum(abs, cpu(ca * gpu(a * x))), a)[1]
  @test cu(g0) ≈ gradient(x -> sum(abs, gpu(a * cpu(ca * x))), ca)[1]
  @test gradient(x -> sum(gpu(cpu(x))), a)[1] isa Matrix
  @test gradient(x -> sum(gpu(cpu(x))), ca)[1] isa CuArray

  g4 = gradient(x -> sum(a * (a' * x)), a)[1]  # no abs, one adjoint
  @test g4 ≈ gradient(x -> sum(cpu(ca * gpu(a' * x))), a)[1]
  @test cu(g4) ≈ gradient(x -> sum(gpu(a * cpu(ca' * x))), ca)[1]

  # Scalar indexing of an array, needs OneElement to transfer to GPU
  # https://github.com/FluxML/Zygote.jl/issues/1005
  @test gradient(x -> cpu(2 .* gpu(x))[1], Float32[1,2,3]) == ([2,0,0],)
  @test gradient(x -> cpu(gpu(x) * gpu(x))[1,2], Float32[1 2 3; 4 5 6; 7 8 9]) == ([2 6 8; 0 2 0; 0 3 0],)
end

@testset "gpu(x) and cpu(x) on structured arrays" begin
  @test cpu(1:3) isa UnitRange
  @test cpu(range(1, 3, length = 4)) isa AbstractRange

  # OneElement isn't GPU compatible
  g1 = Zygote.OneElement(1, (2,3), axes(ones(4,5)))
  @test cpu(g1) isa Zygote.OneElement

  g2 = Zygote.Fill(1f0, 2)
  @test cpu(g2) isa Zygote.FillArrays.AbstractFill

  g3 = transpose(Float32[1 2; 3 4])
  @test parent(cpu(g3)) isa Matrix{Float32}

  @testset "Sparse Arrays" begin
    @test cpu(sparse(rand(3,3))) isa SparseMatrixCSC
    a = sparse(rand(3,3))
    @test cpu(a) === a
    @test gpu(sparse(rand(3,3))) isa CUDA.CUSPARSE.CuSparseMatrixCSC
  end

  # Check that gpu() converts these to CuArrays. This a side-effect of using the same functions
  # in gpu() as in the gradient of cpu(). A different design could avoid having gpu() used alone
  # move these, if that turns out to be desirable.
  @test gpu(g1) isa CuArray
  @test gpu(g1) ≈ cu(Matrix(g1))
  @test gpu(g2) isa CuArray
  @test gpu(g2) ≈ cu(Vector(g2))
  @test parent(gpu(g3)) isa CuArray
end


text:
"""
    cvt_2_plk_n_rm_id()
---
Convert Norwegian final reports to plink format. Also remove ID who were genotyped more than once.
The ID duplicates who were
- of low density genotype
- genotyped earlier

were removed in above precedure order.
"""
function cvt_2_plk_n_rm_id()
    title("Norwegian raw data")
    #############################

    tmp = joinpath(work_dir, "tmp")
    fra = joinpath(work_dir, "data/genotypes/norge")
    til = joinpath(work_dir, "data/genotypes/step-0.plk") # convert to plink and rm ID first
    isdir(til) || mkdir(til)
    empty_dir(til)

    dat = joinpath(fra, "2020-05-11") # where are the 86 files
    ver = "ver"                       # v7, v2, v1, v0

    dic = Dict()

    open(joinpath(fra, "file.desc"), "r") do desc
        _ = readline(desc)      # I defined a header to remember contents.
        for line in eachline(desc)
            nn, vv, _, ff = split(line) # number, version, date, file-name
            subtitle(ff)
            #-------------

            nc = 0                       # number of columns in ff

            if vv ≠ ver                  # Read map dictionary
                ver = vv
                empty!(dic)
                item("Read map $ver")
                dic = ref_map_dict(joinpath(work_dir, "data/maps/updated", "$vv.map"))
            end

            item("Create a GenofileId to GenoId shortlist dictionary")
            sdc = joinpath(til, "$nn.dic")
            open(sdc, "w") do io # output shortlist.
                for line in eachline(joinpath(fra, "idmap_rediverse"))
                    f, x, y = [split(line)[i] for i in [2, 3, 4]]
                    if f == ff
                        write(io, "$x\t$y\n")
                    end
                end
            end

            open(joinpath(fra, "2020-05-11", ff), "r") do io
                for _ in 1:20   # to get into the body
                    line = readline(io)
                end
                nc = length(split(line))
            end                 # determin whether to split v or h.

            empty_dir(tmp)
            file = joinpath(dat, ff)
            cd(tmp)
            if nc >10           # split vertically
                item("Split file $ff vertically")
                _ = read(pipeline(`cat $file`,
                                  `$work_dir/bin/split-v $sdc`),
                         String)
            else                # split horizonly
                item("Split file $ff horizonly")
                _ = read(pipeline(`cat $file`,
                                  `$work_dir/bin/split-h $sdc`),
                         String)
            end
            cd(work_dir)
            list = readdir(tmp)         # the newly created files
            fr2ped(tmp, list, joinpath(tmp, "plink.ped"), "Top")
            sample = list[1]
            create_plink_map(joinpath(tmp, sample), dic, joinpath(tmp, "plink.map"))
            ped_n_map_to_bed(joinpath(tmp, "plink.ped"), joinpath(tmp, "plink.map"), joinpath(til, nn));
        end
    end
end

"""
    merge_raw_norge_gt()
---
Merge the data converted by `cvt_2_plk_n_rm_id` into v0, 1, 2, 7 four files.
"""
function merge_raw_norge_gt()
    title("Merge Norwegian data")
    fra = joinpath(work_dir, "data/genotypes/step-0.plk")
    mid = joinpath(work_dir, "tmp")
    til = joinpath(work_dir, "data/genotypes/step-1.plk")
    isdir(til) || mkdir(til)

    v7 = "0".*[string(i) for i in 1:9]
    v2 = [string(i) for i in 10:84]
    v1 = ["85"]
    v0 = ["86"]
    tt = ["v7", "v2", "v1"]
    it = 1

    ID = Set()                  # ID already dealed
    for vv in [v7, v2, v1]      # v0 was removed as they were all genotyped again later.
        target = tt[it]
        subtitle("Dealing with platform $target")
        empty_dir(mid)
        for s in vv
            open(joinpath(mid, "list"), "w") do list
                for line in eachline(joinpath(fra, "$s.fam"))
                    id = split(line)[2]
                    if id ∉ ID
                        push!(ID, id)
                        write(list, "dummy $id\n")
                    end
                end
            end                 # created keep ID list
            plink_keep_id(joinpath(fra, s),
                          joinpath(mid, "list"),
                          joinpath(mid, s)
                          )
        end
        open(joinpath(mid, "merge.lst"), "w") do list
            for s in vv
                write(list, joinpath(mid, s), '\n')
            end
        end
        merge_beds(joinpath(mid, "merge.lst"),
                   joinpath(til, "norge-$target")
                   )
        it += 1
    end
end

"""
    orgNorgeGT()
---
Organize raw Norwegian data into step-1.plk.  Namely v0, v1, v2, v7.
"""
function orgNorgeGT()
    # Step-0: convert final report to plink, and translate GenofileId to GenoId.
    @time cvt_2_plk_n_rm_id()
    # Step-1: merge to v7, 2, 1, 0 four files, remove dup ID plink by plink
    # ID in v0 were all genotyped again later, so not included.
    @time merge_raw_norge_gt()
end


text:
using Sundials

## Define the system residual function.
function resrob(tres, y, yp, r)
    r[1] = -0.04 * y[1] + 1.0e4 * y[2] * y[3]
    r[2] = -r[1] - 3.0e7 * y[2] * y[2] - yp[2]
    r[1] -= yp[1]
    r[3] = y[1] + y[2] + y[3] - 1.0
end

t = [0.0; 4 * exp10.(range(-1.0, stop = 5.0, length = 7))]
yout, ypout = Sundials.idasol(resrob, [1.0, 0, 0], [-0.04, 0.04, 0.0], t)


text:
function indexing end
macro indexing(ProxyType)
  ProxyType = esc(ProxyType)
  quote
    Base.getindex(p::$ProxyType, args...)	= Base.getindex(ProxyInterfaces.indexing(p), args...)
    Base.firstindex(p::$ProxyType) = Base.firstindex(ProxyInterfaces.indexing(p))
    Base.lastindex(p::$ProxyType)	= Base.lastindex(ProxyInterfaces.indexing(p))
  end
end
macro indexing_mutable(ProxyType)
  ProxyType = esc(ProxyType)
  quote
    @indexing($ProxyType)
    Base.setindex!(p::$ProxyType, args...) = Base.setindex!(ProxyInterfaces.indexing(p), args...)
  end
end


text:
"""
    function grm(gt; nz=0.)
---
This is an instance of Van Raden method I.
Genotypes are of value 0, 1 or 2, and of ID column majored.
They should be stored per ID per column.
This function uses vanRaden method I.
I do no genotype and frequency check here,
As it should be done somewhere else.
"""
function grm(g; nz=0.)
    twop = mean(g, dims = 2)
    s2pq = (1 .- .5 .* twop)'twop # one element array
    r2pq = 1. / s2pq[1]
    Z = g .- twop
    G = Z'Z .* r2pq
    !iszero(nz) && (G += nz .* I)
    G
end


text:
using BenchmarkTools

C_code = raw"""
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

int compare(const void* l, const void* r){
    char* cl = *(char**)l+8;
    char* cr = *(char**)r+8;
    return strcmp(cl, cr);
}

void str_qsort(char **strings, size_t len) {  /* you need to pass len here */

    /* qsort(strings, len, sizeof(char*), (int (*)(void*, void*))strcmp); */
    qsort(strings, len, sizeof(char*), compare);

}
"""

env_path = ""
str_qsort!(::Any) = warn("""
    Either you LD_LIBRARY_PATH environment variable is not defined or gcc is not installed on your computer.
    * Google "how to set environment variable" and set 'LD_LIBRARY_PATH' to a folder/path you have access to
    * On Windows install MinGW. For example from [this link](http://mingw-w64.org/doku.php/download) and
    * add the path of the 'bin' folder to the environment variable PATH
    """)

const Clib = "str_qsort"

try
    env_path = ENV["LD_LIBRARY_PATH"]
    Clib = joinpath(env_path,"strqsort") # tempname()   # make a temporary file
    # compile to a shared library by piping C_code to gcc
    # (works only if you have gcc installed):
    try
        open(`gcc -fPIC -O3 -msse3 -xc -shared -o $(Clib * "." * Libdl.dlext) -`, "w") do f
            print(f, C_code)
        end
        # define a Julia function that calls the C function:
        str_qsort!(X::Array{String}) = ccall(("str_qsort", Clib), Void, (Ptr{UInt64}, Cint), reinterpret(Ptr{UInt64}, pointer(X)), length(X))
    catch e
        warn("""
            Either you LD_LIBRARY_PATH environment variable is not defined or gcc is not installed on your computer.
            * Google "how to set environment variable" and set 'LD_LIBRARY_PATH' to a folder/path you have access to
            * On Windows install MinGW. For example from [this link](http://mingw-w64.org/doku.php/download) and
            * add the path of the 'bin' folder to the environment variable PATH
            """)
    end
catch e
    try
        open(`gcc -fPIC -O3 -msse3 -xc -shared -o $(Clib * "." * Libdl.dlext) -`, "w") do f
            print(f, C_code)
        end
        # define a Julia function that calls the C function:
        str_qsort!(X::Array{String}) = ccall(("str_qsort", Clib), Void, (Ptr{UInt64}, Cint), reinterpret(Ptr{UInt64}, pointer(X)), length(X))
    catch e
        warn("""
            gcc is not installed on your computer:
            * On Windows install MinGW. For example from [this link](http://mingw-w64.org/doku.php/download) and
            * add the path of the 'bin' folder to the environment variable PATH
            """)
    end
    warn("""
    There is no environment variable 'LD_LIBRARY_PATH' defined. We require it to be able to compile the C sort for string sort.
    You have no access to the `str_qsort!` function
    """)
    str_qsort!(::Any) = warn("""
    Either you LD_LIBRARY_PATH environment variable is not defined or gcc is not installed on your computer.
    * Google "how to set environment variable" and set 'LD_LIBRARY_PATH' to a folder/path you have access to
    * On Windows install MinGW. For example from [this link](http://mingw-w64.org/doku.php/download) and
    * add the path of the 'bin' folder to the environment variable PATH
    """)
end


text:
# include("core.jl")

abstract type AbstractAggregate end

mutable struct Aggregate <: AbstractAggregate
    core::Union{AggregateCore,Nothing}
    tools::Union{AggregateTools,Nothing}
    operators::Union{AggregateOperators,Nothing}
    function Aggregate(
        core::Union{AggregateCore,Nothing},
        tools::Union{AggregateTools,Nothing},
        operators::Union{AggregateOperators,Nothing},
    )::Aggregate
        new(core, tools, operators)
    end
end

Aggregate(aggCore::AggregateCore)::Aggregate =
    Aggregate(aggCore, nothing, nothing)::Aggregate

AggregateTools(agg::Aggregate)::AggregateTools = AggregateTools(agg.core)

AggregateOperators(agg::Aggregate)::AggregateOperators =
    AggregateOperators(agg.core, agg.tools)

"""
    setupAggregate(agg; groundEnergy=true, verbose=false)

Generate all basic data from the [`Aggregate`](@ref). Returns
`aggInds, vibindices, aggIndLen, basis, FCFact, FCProd, Ham, Ham_0, Ham_I`.

"""
function setupAggregate(aggCore::AggregateCore; groundEnergy::Bool = true, vib_basis::Symbol = :ground_excited)::Aggregate
    aggTools = AggregateTools(aggCore)
    aggOperators = AggregateOperators(aggCore, aggTools; groundEnergy = groundEnergy, vib_basis = vib_basis)
    return Aggregate(aggCore, aggTools, aggOperators)
end

function setupAggregate!(agg::Aggregate; groundEnergy::Bool = true, vib_basis::Symbol = :ground_excited)::Aggregate
    agg.tools = AggregateTools(agg.core)
    agg.operators = AggregateOperators(agg.core, agg.tools; groundEnergy = groundEnergy, vib_basis = vib_basis)
    return agg
end

Base.:(==)(x::Aggregate, y::Aggregate) =
    x.core == y.core && x.tools == y.tools && x.operators == y.operators

max_stars_repo_name: detrin/OpenQuantumSystems

max_issues_repo_name: detrin/OpenQuantumSystems

max_forks_repo_name: detrin/OpenQuantumSystems


text:
using QuasiArrays, Test

@testset "reducedim" begin
    A = QuasiArray(randn(2,3), (0:0.5:0.5, 1:0.5:2))
    @test sum(A) ≈ sum(A.parent)
    @test sum(A; dims=1) ≈ QuasiArray(sum(A.parent; dims=1), (1:1, 1:0.5:2))
    @test sum(A; dims=2) ≈ QuasiArray(sum(A.parent; dims=2), (0:0.5:0.5, 1:1))

    @test prod(A) ≈ prod(A.parent)
    @test prod(A; dims=1) ≈ QuasiArray(prod(A.parent; dims=1), (1:1, 1:0.5:2))
    @test prod(A; dims=2) ≈ QuasiArray(prod(A.parent; dims=2), (0:0.5:0.5, 1:1))

    @test any(isreal,A) ≈ any(isreal,A.parent)
    @test any(isreal,A; dims=1) ≈ QuasiArray(any(isreal,A.parent; dims=1), (1:1, 1:0.5:2))
    @test any(isreal,A; dims=2) ≈ QuasiArray(any(isreal,A.parent; dims=2), (0:0.5:0.5, 1:1))

    @test all(isreal,A) ≈ all(isreal,A.parent)
    @test all(isreal,A; dims=1) ≈ QuasiArray(all(isreal,A.parent; dims=1), (1:1, 1:0.5:2))
    @test all(isreal,A; dims=2) ≈ QuasiArray(all(isreal,A.parent; dims=2), (0:0.5:0.5, 1:1))

    @test minimum(A) ≈ minimum(A.parent)
    @test_broken minimum(A; dims=1) ≈ QuasiArray(minimum(A.parent; dims=1), (1:1, 1:0.5:2))
    @test_broken minimum(A; dims=2) ≈ QuasiArray(minimum(A.parent; dims=2), (0:0.5:0.5, 1:1))

    @testset "* sum" begin
        b = QuasiVector(randn(3), 1:0.5:2)
        B = QuasiArray(randn(3,2), (1:0.5:2,0:0.5:0.5))

        @test sum(ApplyQuasiArray(*, A, b)) ≈ sum(A*b)
        @test sum(ApplyQuasiArray(*, A, B)) ≈ sum(A*B)
        @test sum(ApplyQuasiArray(*, A, B); dims=1) ≈ sum(A*B; dims=1)
        @test sum(ApplyQuasiArray(*, A, B); dims=2) ≈ sum(A*B; dims=2)

        @test sum(b) ≈ last(cumsum(b)) ≈ cumsum(b)[2]
        @test cumsum(B; dims=1)[2:2,:] ≈ sum(B; dims=1)
        @test cumsum(B; dims=2)[:,0.5:0.5] ≈ sum(B; dims=2)

        @test cumsum(ApplyQuasiArray(*, A, b)) ≈ cumsum(A*b)
        @test cumsum(ApplyQuasiArray(*, A, B); dims=1) ≈ cumsum(A*B; dims=1)
        @test cumsum(ApplyQuasiArray(*, A, B); dims=2) ≈ cumsum(A*B; dims=2)
    end
end


text:
# (A D⁻¹ Aᵀ + δI) Δy = A D⁻¹ ξ₁ + ξ₂
# where D = s_l (x - lvar)⁻¹ + s_u (uvar - x)⁻¹ + ρI,
# and the right hand side of K2 is rhs = [ξ₁]
#                                        [ξ₂]
export K1CholDenseParams

"""
Type to use the K1 formulation with a dense Cholesky factorization.
The input QuadraticModel should have `lcon .== ucon`.
The outer constructor

    sp = K1CholDenseParams(; ρ0 = sqrt(eps()) * 1e5, δ0 = sqrt(eps()) * 1e5)

creates a [`RipQP.SolverParams`](@ref) that should be used to create a [`RipQP.InputConfig`](@ref).
"""
mutable struct K1CholDenseParams <: NormalParams
  uplo::Symbol
  ρ0::Float64
  δ0::Float64
end

function K1CholDenseParams(; ρ0::Float64 = sqrt(eps()) * 1e5, δ0::Float64 = sqrt(eps()) * 1e5)
  uplo = :L # mandatory for LDL fact
  return K1CholDenseParams(uplo, ρ0, δ0)
end

mutable struct PreallocatedDataK1CholDense{T <: Real, S, M <: AbstractMatrix{T}} <:
               PreallocatedDataNormalChol{T, S}
  D::S # temporary top-left diagonal
  invD::Diagonal{T, S}
  AinvD::M
  rhs::S
  regu::Regularization{T}
  K::M # augmented matrix
  diagindK::StepRange{Int, Int}
  tmpldiv::S
end

function convertpad(
  ::Type{<:PreallocatedData{T}},
  pad::PreallocatedDataK1CholDense{T0, S0, M0},
  T02::DataType,
) where {T <: Real, T0 <: Real, S0, M0}
  S = change_vector_eltype(S0, T)
  pad = PreallocatedDataK1CholDense(
    convert(S, pad.D),
    Diagonal(convert(S, pad.invD.diag)),
    convert_mat(pad.AinvD, T),
    convert(S, pad.rhs),
    convert(Regularization{T}, pad.regu),
    convert_mat(pad.K, T),
    pad.diagindK,
    convert(S, pad.tmpldiv),
  )

  if T == Float64 && T0 == Float64
    pad.regu.ρ_min, pad.regu.δ_min = T(sqrt(eps()) * 1e0), T(sqrt(eps()) * 1e0)
  else
    pad.regu.ρ_min, pad.regu.δ_min = T(sqrt(eps(T)) * 1e1), T(sqrt(eps(T)) * 1e1)
  end

  return pad
end

# outer constructor
function PreallocatedData(
  sp::K1CholDenseParams,
  fd::QM_FloatData{T},
  id::QM_IntData,
  itd::IterData{T},
  pt::Point{T},
  iconf::InputConfig{Tconf},
) where {T <: Real, Tconf <: Real}

  # init Regularization values
  D = similar(fd.c, id.nvar)
  if iconf.mode == :mono
    regu = Regularization(T(sp.ρ0), T(sp.δ0), 1e-5 * sqrt(eps(T)), 1e0 * sqrt(eps(T)), :classic)
    D .= T(1.0e0) / 2
  else
    regu =
      Regularization(T(sp.ρ0), T(sp.δ0), T(sqrt(eps(T)) * 1e0), T(sqrt(eps(T)) * 1e0), :classic)
    D .= T(1.0e-2)
  end

  invD = Diagonal(one(T) ./ D)
  AinvD = fd.A * invD
  K = AinvD * fd.A'
  diagindK = diagind(K)
  K[diagindK] .+= regu.δ
  cholesky!(Symmetric(K))

  return PreallocatedDataK1CholDense(
    D,
    invD,
    AinvD,
    similar(D, id.ncon),
    regu,
    K, #K
    diagindK,
    similar(D, id.ncon),
  )
end

# function used to solve problems
# solver LDLFactorization
function solver!(
  dd::AbstractVector{T},
  pad::PreallocatedDataK1CholDense{T},
  dda::DescentDirectionAllocs{T},
  pt::Point{T},
  itd::IterData{T},
  fd::Abstract_QM_FloatData{T},
  id::QM_IntData,
  res::AbstractResiduals{T},
  cnts::Counters,
  T0::DataType,
  step::Symbol,
) where {T <: Real}
  pad.rhs .= @views dd[(id.nvar + 1):end]
  @views mul!(pad.rhs, pad.AinvD, dd[1:(id.nvar)], one(T), one(T))

  ldiv!(pad.tmpldiv, UpperTriangular(pad.K)', pad.rhs)
  ldiv!(pad.rhs, UpperTriangular(pad.K), pad.tmpldiv)

  @views mul!(dd[1:(id.nvar)], fd.A', pad.rhs, one(T), -one(T))
  dd[1:(id.nvar)] ./= pad.D
  dd[(id.nvar + 1):end] .= pad.rhs
  return 0
end

function update_pad!(
  pad::PreallocatedDataK1CholDense{T},
  dda::DescentDirectionAllocs{T},
  pt::Point{T},
  itd::IterData{T},
  fd::Abstract_QM_FloatData{T},
  id::QM_IntData,
  res::AbstractResiduals{T},
  cnts::Counters,
  T0::DataType,
) where {T <: Real}
  if cnts.k != 0
    update_regu!(pad.regu)
  end

  pad.D .= pad.regu.ρ
  pad.D[id.ilow] .+= pt.s_l ./ itd.x_m_lvar
  pad.D[id.iupp] .+= pt.s_u ./ itd.uvar_m_x
  pad.invD.diag .= one(T) ./ pad.D
  mul!(pad.AinvD, fd.A, pad.invD)
  mul!(pad.K, pad.AinvD, fd.A')
  pad.K[pad.diagindK] .+= pad.regu.δ
  cholesky!(Symmetric(pad.K))

  return 0
end


text:
using HDF5

fid = h5open("random_phase.hdf5", "w")
fid["t0"] = θ_A
fid["t5"] = θ_F
fid["ensemble"] = θ̅_F
close(fid)

#=
newfid = h5open("random_phase.hdf5", "r")
array1 = read(newfid["t0"])
=#

max_stars_repo_name: sandreza/FourierCore

max_issues_repo_name: sandreza/FourierCore

max_forks_repo_name: sandreza/FourierCore


text:
module JSONTypeProvider

import JSON3

@static if Base.VERSION < v"1.2"
    function hasfield(::Type{T}, name::Symbol) where T
        return name in fieldnames(T)
    end
    fieldtypes(::Type{T}) where {T} = Tuple(fieldtype(T, i) for i = 1:fieldcount(T))
end

include("build_type.jl")
include("write_type.jl")

end


text:
import JSON

const SUPPORTED_TYPES = Union{EventValues, EventStats}

JSON.lower(es::Event) = event_to_name(es)
function JSON.lower(x::Union{EventValues, EventStats})
    d = Dict{String,Any}()
    T = typeof(x)
    for i = 1:nfields(x)
        name = String(fieldname(T, i))
        value = getfield(x, i)
        ft = typeof(value)
        d[name] = value
    end

    [string(T.name.name), d]
end

function recover(x::Vector)
    length(x) == 2 || throw(ArgumentError("Expecting a vector of length 2"))
    typename = x[1]::String
    fields = x[2]::Dict

    T = if typename == "EventValues"
        EventValues
    elseif typename == "EventStats"
        EventStats
    else
        throw(ArgumentError("Unknown typename $typename"))
    end

    fc = fieldcount(T)
    xs = map(1:fc) do i
        ft = fieldtype(T, i)
        fn = String(fieldname(T, i))

        if ft == Vector{Event}
            map(name_to_event, fields[fn])
        elseif ft <: AbstractMatrix
            C = eltype(ft)
            hcat(Vector{C}.(fields[fn])...)
        else
            convert(ft, fields[fn])
        end
    end

    T(xs...)
end

function save(filename::AbstractString, args::SUPPORTED_TYPES...)
    endswith(filename, ".json") || throw(ArgumentError("Only JSON serialization is supported."))
    open(filename, "w") do io
        save(io, args...)
    end
end

function save(io::IO, args::SUPPORTED_TYPES...)
    isempty(args) && throw(ArgumentError("Nothing to save"))
    JSON.print(io, args)
end

function load(filename::AbstractString)
    endswith(filename, ".json") || throw(ArgumentError("Only JSON serialization is supported."))
    open(filename, "r") do io
        load(io)
    end
end

function load(io::IO)
    parsed = JSON.parse(io)
    map(recover, parsed)
end


text:
using Random, Plots
include_ddp()

# Model
include_model("satellite")

# Time
T = 40
h = 0.1
tf = h * (T - 1)

# Reference trajectory
# https://en.wikipedia.org/wiki/Viviani%27s_curve
# https://mathworld.wolfram.com/VivianisCurve.html
t = range(-2.0 * π, stop = 2.0 * π, length = T)
a = 0.05
xf = a * (1.0 .+ cos.(t))
yf = a * sin.(t)
zf = 2.0 * a * sin.(0.5 .* t)

plot(xf, zf, aspect_ratio = :equal)
plot(xf, yf, aspect_ratio = :equal)
plot(yf, zf, aspect_ratio = :equal)
plot(xf, yf, zf, aspect_ratio = :equal)
pf = [RotZ(0.0 * π) * [xf[t]; yf[t]; zf[t]] for t = 1:T]

ref1 = [a * (1.0 .+ cos.(-π));
		a * sin.(-π);
		2.0 * a * sin.(0.5 .* -π)]
ref2 = [a * (1.0 .+ cos.(0.0));
		a * sin.(0.0);
		2.0 * a * sin.(0.5 .* 0.0)]
ref3 = [a * (1.0 .+ cos.(π));
		a * sin.(0.0);
		2.0 * a * sin.(0.5 .* π)]
ref4 = [a * (1.0 .+ cos.(2.0 * π));
		a * sin.(2.0 * π);
		2.0 * a * sin.(0.5 .* 2.0 * π)]

# Initial conditions, controls, disturbances
mrp1 = MRP(RotXYZ(0.0, 0.0, 0.0))
mrpT = MRP(RotXYZ(0.0, 0.0, 0.0))

x1 = [mrp1.x, mrp1.y, mrp1.z, 0.0, 0.0, 0.0]
xT = [mrpT.x, mrpT.y, mrpT.z, 0.0, 0.0, 0.0]

ū = [1.0e-3 * rand(model.m) for t = 1:T-1]
w = [zeros(model.d) for t = 1:T-1]

# Rollout
x̄ = rollout(model, x1, ū, w, h, T)

obj = StageCosts([NonlinearCost() for t = 1:T], T)

function g(obj::StageCosts, x, u, t)
	T = obj.T
	ω = view(x, 4:6)

	J = 0.0

	# tracking
	p = kinematics(model, x)
	J += transpose(p - pf[t]) * Diagonal(1000.0 * [1.0; 1.0; 1.0]) * (p - pf[t])

	# control limits
	if t < T
		# energy
		J += 0.5 * transpose(ω) * model.J * ω
		J += 1.0e-2 * transpose(u[1:3]) * u[1:3]
	end

	return J / T
end

# Constraints
p_con = [t == T ? model.n : ((t == 10 || t == 20 || t == 30) ? 3 : 0) for t = 1:T]

info_t = Dict(:x1 => ref1, :x2 => ref2, :x3 => ref3)
info_T = Dict(:xT => xT)

con_set = [StageConstraint(p_con[t], t < T ? info_t : info_T) for t = 1:T]

function c!(c, cons::StageConstraints, x, u, t)
	T = cons.T
	n = model.n

	if t == 10
		c[1:3] .= kinematics(model, x) - cons.con[t].info[:x1]
	end

	if t == 20
		c[1:3] .= kinematics(model, x) - cons.con[t].info[:x2]
	end

	if t == 30
		c[1:3] .= kinematics(model, x) - cons.con[t].info[:x3]
	end

	# if t == 40
	# 	c[1:3] .= kinematics(model, x) - cons.con[t].info[:x4]
	# end

	if t == T
		c[1:n] .= x - cons.con[T].info[:xT]
	end
end

prob = problem_data(model, obj, con_set, copy(x̄), copy(ū), w, h, T)

# Solve
@time constrained_ddp_solve!(prob,
    max_iter = 1000, max_al_iter = 10,
	ρ_init = 1.0, ρ_scale = 10.0,
	con_tol = 1.0e-3,
	cache = false)

prob.s_data.obj
x, u = current_trajectory(prob)
x̄, ū = nominal_trajectory(prob)

# Visualize
include(joinpath(pwd(), "models/visualize.jl"))
vis = Visualizer()
render(vis)
visualize!(vis, model, x, Δt = h)

p = [kinematics(model, view(xt, 1:3)) for xt in x]
px = [pt[1] for pt in p]
py = [pt[2] for pt in p]
pz = [pt[3] for pt in p]

plot(yf, zf)
plot!([pt[2] for pt in p], [pt[3] for pt in p])

plot(xf, zf)
plot(xf, yf)
plot(xf)
plot(yf)
plot(zf)

plot(xf, yf, zf, aspect_ratio = :equal)
plot!(px, py, pz, aspect_ratio = :equal)

max_issues_repo_name: thowell/DirectMotionPlanning


text:
################
# shader utils #
################

# from https://github.com/JuliaGL/ModernGL.jl/blob/d56e4ad51f4459c97deeea7666361600a1e6065e/test/util.jl

function validateShader(shader)
	success = GLint[0]
	glGetShaderiv(shader, GL_COMPILE_STATUS, success)
	success[] == GL_TRUE
end

function glErrorMessage()
# Return a string representing the current OpenGL error flag, or the empty string if there's no error.
	err = glGetError()
	err == GL_NO_ERROR ? "" :
	err == GL_INVALID_ENUM ? "GL_INVALID_ENUM: An unacceptable value is specified for an enumerated argument. The offending command is ignored and has no other side effect than to set the error flag." :
	err == GL_INVALID_VALUE ? "GL_INVALID_VALUE: A numeric argument is out of range. The offending command is ignored and has no other side effect than to set the error flag." :
	err == GL_INVALID_OPERATION ? "GL_INVALID_OPERATION: The specified operation is not allowed in the current state. The offending command is ignored and has no other side effect than to set the error flag." :
	err == GL_INVALID_FRAMEBUFFER_OPERATION ? "GL_INVALID_FRAMEBUFFER_OPERATION: The framebuffer object is not complete. The offending command is ignored and has no other side effect than to set the error flag." :
	err == GL_OUT_OF_MEMORY ? "GL_OUT_OF_MEMORY: There is not enough memory left to execute the command. The state of the GL is undefined, except for the state of the error flags, after this error is recorded." : "Unknown OpenGL error with error code $err."
end

function getInfoLog(obj::GLuint)
	# Return the info log for obj, whether it be a shader or a program.
	isShader = glIsShader(obj)
	getiv = isShader == GL_TRUE ? glGetShaderiv : glGetProgramiv
	getInfo = isShader == GL_TRUE ? glGetShaderInfoLog : glGetProgramInfoLog
	# Get the maximum possible length for the descriptive error message
	len = GLint[0]
	getiv(obj, GL_INFO_LOG_LENGTH, len)
	maxlength = len[]
	# TODO: Create a macro that turns the following into the above:
	# maxlength = @glPointer getiv(obj, GL_INFO_LOG_LENGTH, GLint)
	# Return the text of the message if there is any
	if maxlength > 0
		buffer = zeros(GLchar, maxlength)
		sizei = GLsizei[0]
		getInfo(obj, maxlength, sizei, buffer)
		len = sizei[]
		unsafe_string(pointer(buffer), len)
	else
		""
	end
end

function createShader(source, typ)

    # Create the shader
	shader = glCreateShader(typ)::GLuint
	if shader == 0
		error("Error creating shader: ", glErrorMessage())
	end

	# Compile the shader
	glShaderSource(
        shader, 1, convert(Ptr{UInt8},
        pointer([convert(Ptr{GLchar}, pointer(source))])), C_NULL)
	glCompileShader(shader)

	# Check for errors
	!validateShader(shader) && error("Shader creation error: ", getInfoLog(shader))
	shader
end

##################
# OpenGL shaders #
##################

# vertex shader for computing depth image
vertex_source_depth(s) = """
#version $(s) core
uniform mat4 P;
uniform mat4 V;
uniform mat4 pose_mat;
in vec3 position;
void main() {
    gl_Position = P * V * pose_mat * vec4(position, 1);
}
"""

# fragment shader for sillhouette
fragment_source_depth(s) = """
#version $(s) core
out vec4 outColor;
void main()
{
    outColor = vec4(1.0, 1.0, 1.0, 1.0);
}
"""

vertexShader_rgb_basic(s) = """
#version $(s) core
uniform mat4 P;
uniform mat4 V;
uniform mat4 pose_mat;
uniform vec4 color;
in vec3 position;
out vec4 instance_color;
void main() {
    gl_Position = P * V * pose_mat * vec4(position, 1);
    instance_color = color;
}
"""

fragmentShader_rgb_basic(s) = """
#version $(s) core
in vec4 instance_color;
layout(location = 0) out vec4 color;
void main() {
    color = instance_color;
}
"""

vertexShader_rgb(s) = """
#version $(s) core
uniform mat4 P;
uniform mat4 V;
uniform mat4 pose_mat;
uniform mat4 pose_rot;
uniform vec4 color;

in vec3 position;
in vec3 normal;

out vec4 instance_color;
out vec3 normal_out;
out mat4 V_out;
out mat4 pose_mat_out;
out mat4 pose_rot_out;
out vec3 fragVert;

void main() {
    gl_Position = P * V * pose_mat * vec4(position, 1);
    instance_color = color;
    normal_out = normal;
    V_out = V;
    pose_mat_out = pose_mat;
    pose_rot_out = pose_rot;
    fragVert = position;
}
"""

fragmentShader_rgb(s) = """
#version $(s) core
in vec4 instance_color;
in vec3 normal_out;
in mat4 V_out;
in mat4 pose_mat_out;
in mat4 pose_rot_out;
in vec3 fragVert;

layout(location = 0) out vec4 color;
void main() {
    mat3 normalMatrix = transpose(inverse(mat3(pose_mat_out)));
    vec3 normal = normalize(normalMatrix * normal_out);

    //calculate the location of this fragment (pixel) in world coordinates
    vec3 fragPosition = vec3(pose_mat_out * vec4(fragVert, 1));

    //calculate the vector from this pixels surface to the light source
    vec3 surfaceToLight = vec3(0.0, 0.0, 0.0) - fragPosition;

    //calculate the cosine of the angle of incidence
    float brightness = dot(normal, surfaceToLight) / (length(surfaceToLight) * length(normal));
    brightness = clamp(brightness, 0, 1);

    //calculate final color of the pixel, based on:
    // 1. The angle of incidence: brightness
    // 2. The color/intensities of the light: light.intensities
    // 3. The texture and texture coord: texture(tex, fragTexCoord)
    vec4 surfaceColor = instance_color;
    color = vec4(brightness * vec3(1.0, 1.0, 1.0) * surfaceColor.rgb, surfaceColor.a);
}
"""

vertexShader_texture(s) = """
#version $(s) core
uniform mat4 P;
uniform mat4 V;
uniform mat4 pose_mat;
uniform mat4 pose_rot;

in vec3 position;
in vec3 normal;
in vec2 vertTexCoord;

out vec3 normal_out;
out mat4 V_out;
out mat4 pose_mat_out;
out mat4 pose_rot_out;
out vec3 fragVert;
out vec2 fragTexCoord;

void main() {
    gl_Position = P * V * pose_mat * vec4(position, 1);
    normal_out = normal;
    V_out = V;
    pose_mat_out = pose_mat;
    pose_rot_out = pose_rot;
    fragVert = position;
    fragTexCoord = vertTexCoord;
}
"""

fragmentShader_texture(s) = """
#version $(s) core

uniform sampler2D tex;
in vec3 normal_out;
in mat4 V_out;
in mat4 pose_mat_out;
in mat4 pose_rot_out;
in vec3 fragVert;
in vec2 fragTexCoord;

layout(location = 0) out vec4 color;
void main() {
    mat3 normalMatrix = transpose(inverse(mat3(pose_mat_out)));
    vec3 normal = normalize(normalMatrix * normal_out);

    //calculate the location of this fragment (pixel) in world coordinates
    vec3 fragPosition = vec3(pose_mat_out * vec4(fragVert, 1));

    //calculate the vector from this pixels surface to the light source
    vec3 surfaceToLight = vec3(0.0, 0.0, 0.0) - fragPosition;

    //calculate the cosine of the angle of incidence
    float brightness = dot(normal, surfaceToLight) / (length(surfaceToLight) * length(normal));
    brightness = clamp(brightness, 0, 1);

    //calculate final color of the pixel, based on:
    // 1. The angle of incidence: brightness
    // 2. The color/intensities of the light: light.intensities
    // 3. The texture and texture coord: texture(tex, fragTexCoord)
	vec4 surfaceColor = texture(tex, fragTexCoord);
    color = vec4(brightness * vec3(1.0, 1.0, 1.0) * surfaceColor.rgb, surfaceColor.a);
}
"""

vertexShader_texture_mixed(s) = """
#version $(s) core
uniform mat4 P;
uniform mat4 V;
uniform mat4 pose_mat;
uniform mat4 pose_rot;
uniform float textured;
uniform vec4 color;

in vec3 position;
in vec3 normal;
in vec2 vertTexCoord;

out vec3 normal_out;
out mat4 V_out;
out mat4 pose_mat_out;
out mat4 pose_rot_out;
out vec3 fragVert;
out vec2 fragTexCoord;
out float in_textured;
out vec4 in_color;

void main() {
    gl_Position = P * V * pose_mat * vec4(position, 1);
    normal_out = normal;
    V_out = V;
    pose_mat_out = pose_mat;
    pose_rot_out = pose_rot;
    fragVert = position;
    fragTexCoord = vertTexCoord;
    in_textured = textured;
    in_color = color;
}
"""

fragmentShader_texture_mixed(s) = """
#version $(s) core

uniform sampler2D tex;
in vec3 normal_out;
in mat4 V_out;
in mat4 pose_mat_out;
in mat4 pose_rot_out;
in vec3 fragVert;
in vec2 fragTexCoord;
in float in_textured;
in vec4 in_color;

layout(location = 0) out vec4 color;
void main() {
    mat3 normalMatrix = transpose(inverse(mat3(pose_mat_out)));
    vec3 normal = normalize(normalMatrix * normal_out);

    //calculate the location of this fragment (pixel) in world coordinates
    vec3 fragPosition = vec3(pose_mat_out * vec4(fragVert, 1));

    //calculate the vector from this pixels surface to the light source
    vec3 surfaceToLight = vec3(0.0, 0.0, 0.0) - fragPosition;

    //calculate the cosine of the angle of incidence
    float brightness = dot(normal, surfaceToLight) / (length(surfaceToLight) * length(normal));
    brightness = clamp(brightness, 0, 1);

    if (in_textured > 0.5)
    {
		//calculate final color of the pixel, based on:
		// 1. The angle of incidence: brightness
		// 2. The color/intensities of the light: light.intensities
		// 3. The texture and texture coord: texture(tex, fragTexCoord)
        vec4 surfaceColor = texture(tex, fragTexCoord);
        color = vec4(brightness * vec3(1.0, 1.0, 1.0) * surfaceColor.rgb, surfaceColor.a);
	}
	else
	{
        vec4 surfaceColor = in_color;
        color = vec4(brightness * vec3(1.0, 1.0, 1.0) * surfaceColor.rgb, surfaceColor.a);
	}

}
"""