expressions.jl

################################################################################
#                          PARAMETER FUNCTION METHODS
################################################################################
# Extend dispatch_variable_ref
function dispatch_variable_ref(
    model::InfiniteModel,
    index::ParameterFunctionIndex
    )::ParameterFunctionRef
    return ParameterFunctionRef(model, index)
end

# Extend _add_data_object
function _add_data_object(
    model::InfiniteModel,
    object::ParameterFunctionData
    )::ParameterFunctionIndex
    return MOIUC.add_item(model.param_functions, object)
end

# Extend _data_dictionary (reference based)
function _data_dictionary(
    fref::ParameterFunctionRef
    )::MOIUC.CleverDict{ParameterFunctionIndex, ParameterFunctionData{<:ParameterFunction}}
    return JuMP.owner_model(fref).param_functions
end

# Extend _data_object
function _data_object(fref::ParameterFunctionRef)
    object = get(_data_dictionary(fref), JuMP.index(fref), nothing)
    if isnothing(object) 
        error("Invalid parameter function reference, cannot find ",
              "corresponding object in the model. This is likely ",
              "caused by using the reference of a deleted function.")
    end
    return object
end

"""
    core_object(fref::ParameterFunctionRef)::ParameterFunction

Retrieve the underlying core [`ParameterFucntion`](@ref) object for `fref`. 
This is intended as an advanced method for developers.
"""
function core_object(fref::ParameterFunctionRef)
    return _data_object(fref).func
end

"""
    parameter_group_int_indices(fref::ParameterFunctionRef)::Vector{Int}

Return the list of infinite parameter group integer indices used by `fref`.
"""
function parameter_group_int_indices(fref::ParameterFunctionRef)
    return core_object(fref).group_int_idxs
end

# Extend _parameter_numbers
function _parameter_numbers(fref::ParameterFunctionRef)
    return core_object(fref).parameter_nums
end

"""
    build_parameter_function(
        _error::Function, 
        func::Function, 
        parameter_refs::Union{GeneralVariableRef, AbstractArray{<:GeneralVariableRef}, Tuple}
        )::ParameterFunction

Build an [`ParameterFunction`](@ref) object that employs a parameter function 
`func` that takes instances of the infinite parameter(s) as input. This can 
ultimately by incorporated into expressions to enable nonlinear infinite parameter 
behavior and/or incorporate data over infinite domains.

Here `func` should be of the form:
```
func(paramvals...)::Float64
```
where the formatting of `paramvals` is analagous to point variables (and will be 
based on the tuple of infinite parameter references given in `parameter_refs`). 

Errors if the infinite parameter tuple is formatted incorrectly. The allowed 
format follows that of infinite variables. Also errors if the function doesn't 
accept a support realization of the `parameter_refs` as input.

**Example**
```julia-repl 
julia> f = build_parameter_function(error, sin, t);
```
"""
function build_parameter_function(
    _error::Function, 
    func::Function, 
    parameter_refs::Union{GeneralVariableRef, AbstractArray{<:GeneralVariableRef}, Tuple};
    extra_kwargs...
    )
    # check for unneeded keywords
    for (kwarg, _) in extra_kwargs
        _error("Keyword argument $kwarg is not for use with parameter functions.")
    end
    # process the parameter reference inputs and check
    prefs = Collections.VectorTuple(parameter_refs)
    _check_parameter_tuple(_error, prefs)
    # check the function
    supp_type = typeof(Tuple(Vector{Float64}(undef, length(prefs)), prefs))
    if !hasmethod(func, supp_type)
        _error("Parameter function method `$(func)` is not defined for `$(func)(",
               join(Tuple(prefs), ", ") * ")`. Note that the infinite parameter ",
               "arguments `" * join(Tuple(prefs), ", ") * "` are checked via a ",
               "numeric support (each parameter is a `Float64`) of the same format.")
    end  
    # get the parameter group integer indices
    group_int_idxs = Int[]
    for pref in prefs 
        union!(group_int_idxs, parameter_group_int_index(pref))
    end
    # get the parameter numbers 
    param_nums = [_parameter_number(pref) for pref in prefs]
    # make the variable and return
    return ParameterFunction(func, prefs, group_int_idxs, param_nums)
end 

# Fallback for weird macro inputs
function build_parameter_function(_error::Function, func, prefs; kwargs...)
    _error("Invalid syntax causing unexpected parsing of the function and " * 
           "infinite parameter arguments.")
end 

# Used to update parameter-parameter function mappings
function _update_param_var_mapping(
    fref::ParameterFunctionRef,
    prefs::Collections.VectorTuple
    )::Nothing
    for pref in prefs
        dependency_list = _parameter_function_dependencies(pref)
        if !(JuMP.index(fref) in dependency_list)
            push!(dependency_list, JuMP.index(fref))
        end
    end
    return
end

"""
    add_parameter_function(model::InfiniteModel, pfunc::ParameterFunction, 
                           [name::String])::GeneralVariableRef

Add an [`ParameterFunction`](@ref) `pfunc` to the `model` using `name` for 
printing and return a `GeneralVariableRef` such that it can be embedded in 
expressions. Errors if the parameter function `pfunc` points to do not belong to 
`model`. Note that `pfunc` should be created using [`build_parameter_function`](@ref).

**Example**
```julia-repl
julia> f = build_parameter_function(error, sin, t);

julia> fref = add_parameter_function(model, f)
sin(t)
```
"""
function add_parameter_function(
    model::InfiniteModel, 
    pfunc::ParameterFunction, 
    name::String = String(nameof(pfunc.func))
    )::GeneralVariableRef
    _check_parameters_valid(model, pfunc.parameter_refs)
    data_object = ParameterFunctionData(pfunc, name)
    findex = _add_data_object(model, data_object)
    fref = ParameterFunctionRef(model, findex)
    _update_param_var_mapping(fref, pfunc.parameter_refs)
    return GeneralVariableRef(model, findex)
end

"""
    JuMP.name(fref::ParameterFunctionRef)::String

Extend `JuMP.name` to return the base name of `fref`.

**Example**
```julia-repl
julia> name(fref)
"func_name"
```
"""
function JuMP.name(fref::ParameterFunctionRef)::String 
    object = get(_data_dictionary(fref), JuMP.index(fref), nothing)
    return isnothing(object) ? "" : object.name
end

"""
    JuMP.set_name(fref::ParameterFunctionRef, name::String)::Nothing

Extend `JuMP.set_name` to set the name of a parameter function.

**Example**
```julia-repl
julia> set_name(fref, "func_name")

julia> name(fref)
"func_name"
```
"""
function JuMP.set_name(fref::ParameterFunctionRef, name::String)::Nothing
    _data_object(fref).name = name
    return
end

"""
    parameter_function(func::Function, 
                       pref_inputs::Union{GeneralVariableRef, AbstractArray{GeneralVariableRef}, Tuple}; 
                       [name::String = [the name of `func`]]
                       )::GeneralVariableRef

Make a parameter function and return a `GeneralVariableRef` that can be 
embedded in InfiniteOpt expressions. This serves as a convenient wrapper for 
[`build_parameter_function`](@ref) and [`add_parameter_function`](@ref). For an 
even more convenient definition method see [`@parameter_function`](@ref).

Here `func` denotes the function that will take a support of infinite parameters as 
input (formatted like `pref_inputs`) and will return a scalar value. Specifically, 
`func` should be of the form:
```
func(paramvals...)::Float64
```
where the formatting of `paramvals` is analagous to point variables (and will be 
based on the tuple of infinite parameter references given in `parameter_refs`).
Moreover, `func` must be a function that returns a scalar numeric value. 

Errors if `func` will not take a support formatted like `pref_inputs` in 
combination with the `fargs` and `fkwargs` specified. Also errors if `pref_inputs` 
follow an invalid input format.

**Example**
```julia-repl 
julia> p_func = parameter_function(sin, t)
sin(t)

julia> p_func3 = parameter_function((t_supp) -> 2 * sin(2 * t_supp), t, name = "mysin")
mysin(t)

julia> p_func4 = parameter_function(t, name = "mysin") do t_supp
                    if t_supp <= 5
                        return sin(t_supp)
                    else 
                        return 2 * sin(2 * t_supp)
                    end
                 end

mysin(t)
```
"""
function parameter_function(
    func::Function, 
    pref_inputs; 
    name::String = String(nameof(func))
    )::GeneralVariableRef
    f = build_parameter_function(error, func, pref_inputs)
    model = JuMP.owner_model(first(f.parameter_refs))
    return add_parameter_function(model, f, name)
end

"""
    raw_parameter_refs(fref::ParameterFunctionRef)::VectorTuple

Return the raw [`VectorTuple`](@ref InfiniteOpt.Collections.VectorTuple) of the 
parameter references that `fref` depends on. This is primarily an internal method 
where [`parameter_refs`](@ref parameter_refs(::ParameterFunctionRef)) 
is intended as the preferred user function.
"""
function raw_parameter_refs(fref::ParameterFunctionRef)
    return core_object(fref).parameter_refs 
end

"""
    parameter_refs(fref::ParameterFunctionRef)::Tuple

Return the parameter references associated with `fref`. This
is formatted as a Tuple of containing the parameter references as they inputted
to define `fref`.

**Example**
```julia-repl
julia> parameter_refs(p_func)
(t,)
```
"""
function parameter_refs(fref::ParameterFunctionRef)
    return Tuple(raw_parameter_refs(fref))
end

"""
    parameter_list(fref::ParameterFunctionRef)::Vector{GeneralVariableRef}

Return a vector of the parameter references that `fref` depends on. This is
primarily an internal method where [`parameter_refs`](@ref parameter_refs(::ParameterFunctionRef))
is intended as the preferred user function.
"""
function parameter_list(fref::ParameterFunctionRef)::Vector{GeneralVariableRef}
    return raw_parameter_refs(fref).values
end

"""
    raw_function(fref::ParameterFunctionRef)::Function

Returns the raw function behind `fref` that takes a particular support of `fref`'s 
infinite parameters as input. 
"""
function raw_function(fref::ParameterFunctionRef)
    return core_object(fref).func
end

"""
    call_function(fref::ParameterFunctionRef, support...)::Float64

Safely evaluates the [`raw_function`](@ref) of `fref` at a particular support `support`
point that matches the format of the infinite parameter tuple given when the `fref` 
was defined. This is essentially equivalent to `raw_function(fref)(supps...)`. 
"""
function call_function(fref::ParameterFunctionRef, supps...)::Float64
    pfunc = core_object(fref)
    return pfunc.func(supps...)
end

# Extend _semi_infinite_variable_dependencies
function _semi_infinite_variable_dependencies(
    fref::ParameterFunctionRef
     )::Vector{SemiInfiniteVariableIndex}
    return _data_object(fref).semi_infinite_var_indices
end

# Extend _derivative_dependencies
function _derivative_dependencies(
    fref::ParameterFunctionRef
     )::Vector{DerivativeIndex}
    return _data_object(fref).derivative_indices
end

# Extend _measure_dependencies
function _measure_dependencies(fref::ParameterFunctionRef)::Vector{MeasureIndex}
    return _data_object(fref).measure_indices
end

# Extend _constraint_dependencies
function _constraint_dependencies(
    fref::ParameterFunctionRef
    )::Vector{InfOptConstraintIndex}
    return _data_object(fref).constraint_indices
end

"""
    used_by_semi_infinite_variable(fref::ParameterFunctionRef)::Bool

Return a `Bool` indicating if `fref` is used by a semi-infinite infinite variable.

**Example**
```julia-repl
julia> used_by_semi_infinite_variable(fref)
false
```
"""
function used_by_semi_infinite_variable(fref::ParameterFunctionRef)::Bool
    return !isempty(_semi_infinite_variable_dependencies(fref))
end

"""
    used_by_derivative(fref::ParameterFunctionRef)::Bool

Return a `Bool` indicating if `fref` is used by a derivative.

**Example**
```julia-repl
julia> used_by_derivative(vref)
true
```
"""
function used_by_derivative(fref::ParameterFunctionRef)::Bool
    return !isempty(_derivative_dependencies(fref))
end

"""
    used_by_measure(fref::ParameterFunctionRef)::Bool

Return a `Bool` indicating if `fref` is used by a measure.

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

"""
    used_by_constraint(fref::ParameterFunctionRef)::Bool

Return a `Bool` indicating if `fref` is used by a constraint.

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

"""
    is_used(fref::ParameterFunctionRef)::Bool

Return a `Bool` indicating if `fref` is used in the model.

**Example**
```julia-repl
julia> is_used(fref)
true
```
"""
function is_used(fref::ParameterFunctionRef)::Bool
    return used_by_measure(fref) || used_by_constraint(fref) || 
           used_by_semi_infinite_variable(fref) || used_by_derivative(fref)
end

"""
    JuMP.delete(model::InfiniteModel, fref::ParameterFunctionRef)::Nothing

Extend `JuMP.delete` to delete parameter functions and their dependencies. Errors 
if `fref` is invalid, meaning it has already been deleted or it belongs to 
another model.
"""
function JuMP.delete(model::InfiniteModel, fref::ParameterFunctionRef)::Nothing 
    @assert JuMP.is_valid(model, fref) "Parameter function is invalid."
    # update the transformation backend status
    if is_used(fref)
        set_transformation_backend_ready(model, false)
    end
    # update parameter mapping
    all_prefs = parameter_list(fref)
    for pref in all_prefs
        filter!(e -> e != JuMP.index(fref), _parameter_function_dependencies(pref))
    end
    gvref = GeneralVariableRef(model, JuMP.index(fref))
    # delete associated semi-infinite variables and mapping
    for index in _semi_infinite_variable_dependencies(fref)
        JuMP.delete(model, dispatch_variable_ref(model, index))
    end
    # delete associated derivative variables and mapping 
    for index in _derivative_dependencies(fref)
        JuMP.delete(model, dispatch_variable_ref(model, index))
    end
    # remove from measures if used
    for mindex in _measure_dependencies(fref)
        mref = dispatch_variable_ref(model, mindex)
        func = measure_function(mref)
        data = measure_data(mref)
        if func isa GeneralVariableRef
            new_func = zero(JuMP.GenericAffExpr{Float64, GeneralVariableRef})
            new_meas = Measure(new_func, data, Int[], Int[], true)
        else
            _remove_variable(func, gvref)
            new_meas = build_measure(func, data)
        end
        _set_core_object(mref, new_meas)
    end
    # remove from constraints if used
    for cindex in copy(_constraint_dependencies(fref))
        cref = InfOptConstraintRef(model, cindex)
        func = JuMP.jump_function(JuMP.constraint_object(cref))
        if func isa GeneralVariableRef
            set = JuMP.moi_set(JuMP.constraint_object(cref))
            new_func = zero(JuMP.GenericAffExpr{Float64, GeneralVariableRef})
            new_constr = JuMP.ScalarConstraint(new_func, set)
            _set_core_object(cref, new_constr)
            empty!(parameter_group_int_indices(cref))
        elseif func isa AbstractArray && any(isequal(gvref), func)
            JuMP.delete(model, cref)
        else
            _remove_variable(func, gvref)
            _data_object(cref).group_int_idxs = sort(parameter_group_int_indices(func))
        end
    end
    # delete the data object
    _delete_data_object(fref)
    return
end

################################################################################
#                              VARIABLE ITERATION
################################################################################
## Create helper methods to interrogate the variables of an expr w/ a function
# Constant
function _interrogate_variables(interrogator::Function, c)
    return
end

# GeneralVariableRef
function _interrogate_variables(
    interrogator::Function, 
    v::JuMP.AbstractVariableRef
    )
    interrogator(v)
    return
end

# AffExpr
function _interrogate_variables(
    interrogator::Function, 
    aff::JuMP.GenericAffExpr
    )
    for (v, _) in aff.terms
        interrogator(v)
    end
    return
end

# QuadExpr
function _interrogate_variables(
    interrogator::Function, 
    quad::JuMP.GenericQuadExpr
    )
    for (p, _) in quad.terms
        interrogator(p.a)
        interrogator(p.b)
    end
    _interrogate_variables(interrogator, quad.aff)
    return
end

# NonlinearExpr (avoid recursion to handle deeply nested expressions)
function _interrogate_variables(interrogator::Function, nlp::JuMP.GenericNonlinearExpr)
    stack = Vector{Any}[nlp.args]
    while !isempty(stack)
        args = pop!(stack)
        for arg in args
            if arg isa JuMP.GenericNonlinearExpr
                push!(stack, arg.args)
            else
                _interrogate_variables(interrogator, arg)
            end
        end
    end
    return
end

# AbstractArray
function _interrogate_variables(interrogator::Function, arr::AbstractArray)
    for ex in arr
        _interrogate_variables(interrogator, ex)
    end
    return
end

################################################################################
#                            VARIABLE LIST MAKING
################################################################################
"""
    all_expression_variables(expr::JuMP.AbstractJuMPScalar)::Vector

Returns a vector of all the variable references contained in `expr`.

**Example**
```julia-repl
julia> all_expr_variables(y^2 + z - t)
3-element Array{GeneralVariableRef,1}:
 y(t)
 z
 t
```
"""
function all_expression_variables(f)
    error("`all_expression_variables` not defined for expression of type $(typeof(f)).")
end

# GeneralVariableRef
function all_expression_variables(f::GeneralVariableRef)
    return [f]
end

# GenericAffExpr
function all_expression_variables(f::JuMP.GenericAffExpr)
    return collect(keys(f.terms))
end

# GenericQuadExpr
function all_expression_variables(f::JuMP.GenericQuadExpr)
    vref_set = Set(keys(f.aff.terms))
    for (pair, _) in f.terms
        push!(vref_set, pair.a, pair.b)
    end
    return collect(vref_set)
end

# NonlinearExpr or array of expressions
function all_expression_variables(f::Union{JuMP.GenericNonlinearExpr, AbstractArray})
    vref_set = Set{GeneralVariableRef}()
    _interrogate_variables(v -> push!(vref_set, v), f)
    return collect(vref_set)
end

################################################################################
#                        GROUP/PARAMETER NUMBER METHODS
################################################################################
## Return the unique set of parameter group integer indices in an expression
# Dispatch fallback (--> should be defined for each non-empty variable type)
parameter_group_int_indices(v::DispatchVariableRef) = Int[]

"""
    parameter_group_int_indices(vref::GeneralVariableRef)::Vector{Int}

Return the list of infinite parameter group integer indices used by `vref`.
"""
function parameter_group_int_indices(v::GeneralVariableRef)
    return parameter_group_int_indices(dispatch_variable_ref(v))
end

"""
    parameter_group_int_indices(expr::JuMP.AbstractJuMPScalar)::Vector{Int}

Return the list of infinite parameter group integer indices used by `expr`.
"""
function parameter_group_int_indices(expr)
    group_int_idxs = Set{Int}()
    _interrogate_variables(v -> union!(group_int_idxs, parameter_group_int_indices(v)), expr)
    return collect(group_int_idxs)
end

## Return the unique set of parameter numbers in an expression
# Dispatch fallback (--> should be defined for each non-empty variable type)
_parameter_numbers(v::DispatchVariableRef) = Int[]

# GeneralVariableRef
function _parameter_numbers(v::GeneralVariableRef)
    return _parameter_numbers(dispatch_variable_ref(v))
end

# Other
function _parameter_numbers(expr)
    param_nums = Set{Int}()
    _interrogate_variables(v -> union!(param_nums, _parameter_numbers(v)), expr)
    return collect(param_nums)
end

################################################################################
#                            VARIABLE REMOVAL BOUNDS
################################################################################
## Delete variables from an expression
# GenericAffExpr
function _remove_variable(f::JuMP.GenericAffExpr, vref::GeneralVariableRef)
    if haskey(f.terms, vref)
        delete!(f.terms, vref)
    end
    return
end

# GenericQuadExpr
function _remove_variable(f::JuMP.GenericQuadExpr, vref::GeneralVariableRef)
    _remove_variable(f.aff, vref)
    for (pair, _) in f.terms
        if isequal(pair.a, vref) || isequal(pair.b, vref)
            delete!(f.terms, pair)
        end
    end
    return
end

# Helper functions for nonlinear variable deletion
function _remove_variable_from_leaf(c, vref)
    return c
end
function _remove_variable_from_leaf(n_vref::GeneralVariableRef, vref)
    return isequal(n_vref, vref) ? 0.0 : n_vref
end
function _remove_variable_from_leaf(
    ex::Union{JuMP.GenericAffExpr, JuMP.GenericQuadExpr}, 
    vref
    )
    _remove_variable(ex, vref)
    return ex
end

# Nonlinear (avoid recursion to handle deeply nested expressions)
function _remove_variable(f::JuMP.GenericNonlinearExpr, vref::GeneralVariableRef)
    stack = Tuple{Vector{Any}, Int}[]
    for i in eachindex(f.args) # should be reverse, but order doesn't matter
        push!(stack, (f.args, i))
    end
    while !isempty(stack)
        arr, idx = pop!(stack)
        expr = arr[idx]
        if expr isa JuMP.GenericNonlinearExpr
            for i in eachindex(expr.args) # should be reverse, but order doesn't matter
                push!(stack, (expr.args, i))
            end
        else
            arr[idx] = _remove_variable_from_leaf(expr, vref)
        end
    end
    return
end

# AbstractArray 
function _remove_variable(arr::AbstractArray, vref::GeneralVariableRef)
    for f in arr
        _remove_variable(f, vref)
    end
    return
end

################################################################################
#                                MAPPING METHODS
################################################################################
"""
    map_expression(transform::Function, 
                   expr::JuMP.AbstractJuMPScalar)::JuMP.AbstractJuMPScalar

Map and return a new expression of `expr` where each variable is transformed 
via `transform`. This can be helpful for writing user extensions.
"""
function map_expression(transform::Function, v::JuMP.AbstractVariableRef)
    return transform(v)
end

# Constant
function map_expression(transform::Function, c::Number)
    return c
end

# AffExpr
function map_expression(transform::Function, aff::JuMP.GenericAffExpr)
    # flatten needed in case expression becomes nonlinear with transform
    return JuMP.flatten!(@_expr(sum(c * transform(v) 
                        for (c, v) in JuMP.linear_terms(aff)) + 
                        JuMP.constant(aff)))
end

# QuadExpr
function map_expression(transform::Function, quad::JuMP.GenericQuadExpr)
    # flatten needed in case expression becomes nonlinear with transform
    return JuMP.flatten!(@_expr(sum(c * transform(v1) * transform(v2) 
                        for (c, v1, v2) in JuMP.quad_terms(quad)) + 
                        map_expression(transform, quad.aff)))
end

# NonlinearExpr 
function map_expression(transform::Function, nlp::JuMP.GenericNonlinearExpr)
    # TODO: Figure out how to make the recursionless code work 
    # stack = Tuple{Vector{Any}, Vector{Any}}[]
    # new_nlp = JuMP.GenericNonlinearExpr{NewVrefType}(nlp.head, Any[]) # Need to add `NewVrefType` arg throughout pkg
    # push!(stack, (nlp.args, new_nlp.args))
    # while !isempty(stack)
    #     args, cloned = pop!(stack)
    #     for arg in args
    #         if arg isa JuMP.GenericNonlinearExpr
    #             new_expr = JuMP.GenericNonlinearExpr{NewVrefType}(arg.head, Any[])
    #             push!(stack, (arg.args, new_expr.args))
    #         else
    #             new_expr = map_expression(transform, arg)
    #         end
    #         push!(cloned, new_expr)
    #     end
    # end
    # return new_nlp
    return JuMP.GenericNonlinearExpr(nlp.head, Any[map_expression(transform, arg) for arg in nlp.args])
end

"""
    map_expression_to_ast(var_mapper::Function, [op_mapper::Function,] expr::JuMP.AbstractJuMPScalar)::Expr

Map the expression `expr` to a Julia AST expression where each variable is mapped 
via `var_mapper` and is directly interpolated into the AST expression. Any nonlinear 
operators can be mapped if needed via `op_mapper` and will be inserted into the 
AST expression. This is only intended for developers and advanced users.
"""
function map_expression_to_ast(var_mapper::Function, expr)
    return map_expression_to_ast(var_mapper, identity, expr)
end

# Constant
function map_expression_to_ast(::Function, ::Function, c)
    return c
end

# GeneralVariableRef
function map_expression_to_ast(
    var_mapper::Function, 
    ::Function, 
    vref::GeneralVariableRef
    )
    return var_mapper(vref)
end

# GenericAffExpr
function map_expression_to_ast(
    var_mapper::Function, 
    ::Function, 
    aff::GenericAffExpr
    )
    ex = Expr(:call, :+)
    for (c, v) in JuMP.linear_terms(aff)
        if isone(c)
            push!(ex.args, var_mapper(v))
        else
            push!(ex.args, Expr(:call, :*, c, var_mapper(v)))
        end
    end
    if !iszero(aff.constant) || isempty(ex.args[2:end])
        push!(ex.args, aff.constant)
    end
    return ex
end

# GenericQuadExpr
function map_expression_to_ast(
    var_mapper::Function, 
    op_mapper::Function, 
    quad::GenericQuadExpr
    )
    ex = Expr(:call, :+)
    for (c, v1, v2) in JuMP.quad_terms(quad)
        if isone(c)
            push!(ex.args, Expr(:call, :*, var_mapper(v1), var_mapper(v2)))
        else
            push!(ex.args, Expr(:call, :*, c, var_mapper(v1), var_mapper(v2)))
        end
    end
    aff_ex = map_expression_to_ast(var_mapper, op_mapper, quad.aff)
    if aff_ex.args != [:+, 0.0] || isempty(ex.args[2:end])
        append!(ex.args, aff_ex.args[2:end])
    end
    return ex
end

# GenericNonlinearExpr
function map_expression_to_ast(
    var_mapper::Function, 
    op_mapper::Function, 
    expr::JuMP.GenericNonlinearExpr
    )
    ex = Expr(:call, op_mapper(expr.head))
    append!(ex.args, map_expression_to_ast(var_mapper, op_mapper, arg) for arg in expr.args)
    return ex
end

################################################################################
#                            EXPRESSION RESTICTIONS
################################################################################
"""
    restrict(expr::JuMP.AbstractJuMPScalar, supps...)::JuMP.AbstractJuMPScalar

Restrict an infinite expression `expr` to be enforced over infinite parameter 
supports `supps`. This is limited to expressions only contain infinite variables 
with the same kind of infinite parameter dependencies. Note that more conveniently 
the expression can be treated as a function for the syntax `expr(supps...)`. 

**Example**
```julia-repl
julia> ex = @expression(model, 3y - 2)
3 y(t) - 2

julia> restrict(ex, 0)
3 y(0) - 2

julia> ex(0)
3 y(0) - 2
```
"""
function restrict(expr::JuMP.AbstractJuMPScalar, supps...)
    # check to make sure all variables depend on the same infinite parameters
    pref_tuples = Set()
    _interrogate_variables(expr) do v
        if v.index_type <: InfiniteParameterIndex
            error("Restrictions on expressions with infinite parameters are not supported.")
        elseif !(v.index_type <: Union{FiniteVariableIndex, PointVariableIndex, FiniteParameterIndex})
            push!(pref_tuples, raw_parameter_refs(v))
        end
    end
    if length(pref_tuples) > 1
        error("Unable to restrict expression `$expr` with supports `$supps`. " *
              "Not all the variables use the same infinite parameters in the " *
              "same format.")
    end
    # restrict the expression using supps and return
    return map_expression(expr) do v
        if isempty(parameter_group_int_indices(v))
            return v
        else
            return restrict(v, supps...)
        end
    end
end

## Make expressions callable for restrictions
# AffExprs
function (aff::JuMP.GenericAffExpr{Float64, GeneralVariableRef})(supps...) 
    return restrict(aff, supps...)
end

# QuadExprs
function (quad::JuMP.GenericQuadExpr{Float64, GeneralVariableRef})(supps...) 
    return restrict(quad, supps...)
end

# NonlinearExprs
function (nl::JuMP.GenericNonlinearExpr{GeneralVariableRef})(supps...) 
    return restrict(nl, supps...)
end

################################################################################
#                             COEFFICIENT METHODS
################################################################################
## Modify linear coefficient of variable in expression
# GeneralVariableRef
function _set_variable_coefficient!(expr::GeneralVariableRef,
    var::GeneralVariableRef,
    coeff::Real
    )
    # Determine if variable is that of the expression and change accordingly
    if isequal(expr, var)
        return Float64(coeff) * var
    else
        return expr + Float64(coeff) * var
    end
end

# GenericAffExpr
function _set_variable_coefficient!(expr::JuMP.GenericAffExpr,
    var::GeneralVariableRef,
    coeff::Real
    )
    # Determine if variable is in the expression and change accordingly
    if haskey(expr.terms, var)
        expr.terms[var] = coeff
        return expr
    else
        return expr + coeff * var
    end
end

# GenericQuadExpr
function _set_variable_coefficient!(expr::JuMP.GenericQuadExpr,
    var::GeneralVariableRef,
    coeff::Real
    )
    # Determine if variable is in the expression and change accordingly
    if haskey(expr.aff.terms, var)
        expr.aff.terms[var] = coeff
        return expr
    else
        return expr + coeff * var
    end
end

# Fallback
function _set_variable_coefficient!(expr, var::GeneralVariableRef, coeff::Real)
    error("Unsupported function type for coefficient modification.")
end

## Query the coefficient of a variable 
# GeneralVariableRef 
function _affine_coefficient(
    func::GeneralVariableRef, 
    var::GeneralVariableRef
    )
    return isequal(func, var) ? 1.0 : 0.0
end

# GenericAffExpr
function _affine_coefficient(
    func::GenericAffExpr, 
    var::GeneralVariableRef
    )
    return get(func.terms, var, 0.0)
end

# GenericQuadExpr
function _affine_coefficient(
    func::GenericQuadExpr, 
    var::GeneralVariableRef
    )
    return get(func.aff.terms, var, 0.0)
end

# Fallback 
function _affine_coefficient(func, var::GeneralVariableRef)
    error("Unsupported function type for coefficient queries.")
end

################################################################################
#                         PARAMETER REFERENCE METHODS
################################################################################
"""
    parameter_refs(expr)::Tuple

Return the tuple of parameter references that determine the infinite
dependencies of `expr`.

**Example**
```julia-repl
julia> parameter_refs(my_expr)
(t,)
```
"""
function parameter_refs(
    expr::Union{JuMP.GenericAffExpr, JuMP.GenericQuadExpr, JuMP.GenericNonlinearExpr}
    )
    model = JuMP.owner_model(expr)
    isnothing(model) && return ()
    prefs = parameter_refs(model)
    group_int_idxs = parameter_group_int_indices(expr)
    length(prefs) == length(group_int_idxs) && return prefs
    return Tuple(prefs[i] for i in eachindex(prefs) if i in group_int_idxs)
end