datatypes.jl

################################################################################
#                              LOGICAL VARIABLES
################################################################################
"""
    LogicalVariableIndex

A type for storing the index of a [`LogicalVariable`](@ref).

**Fields**
- `value::Int64`: The index value.
"""
struct LogicalVariableIndex
    value::Int64
end

"""
    LogicalVariableRef{M <: JuMP.AbstractModel}

A type for looking up logical variables.
"""
struct LogicalVariableRef{M <:JuMP.AbstractModel} <: JuMP.AbstractVariableRef
    model::M
    index::LogicalVariableIndex
end

"""
    LogicalVariable <: JuMP.AbstractVariable

A variable type the logical variables associated with disjuncts in a [`Disjunction`](@ref).

**Fields**
- `fix_value::Union{Nothing, Bool}`: A fixed boolean value if there is one.
- `start_value::Union{Nothing, Bool}`: An initial guess if there is one.
- `logical_complement::Union{Nothing, LogicalVariableRef}`: The logical complement of
   this variable if there is one.
"""
struct LogicalVariable <: JuMP.AbstractVariable 
    fix_value::Union{Nothing, Bool}
    start_value::Union{Nothing, Bool}
    logical_complement::Union{Nothing, LogicalVariableRef}
end

# Wrapper variable type for including arbitrary tags that will be used for 
# creating reformulation variables later on
struct _TaggedLogicalVariable{T} <: JuMP.AbstractVariable
    variable::LogicalVariable
    tag_data::T
end

"""
    Logical{T}

Tag for creating logical variables using `@variable`. Most often this will 
be used to enable the syntax:
```julia
@variable(model, var_expr, Logical, [kwargs...])
```
which creates a [`LogicalVariable`](@ref) that will ultimately be 
reformulated into a binary variable of the form:
```julia
@variable(model, var_expr, Bin, [kwargs...])
```

To include a tag that is used to create the reformulated variables, the syntax 
becomes:
```julia
@variable(model, var_expr, Logical(MyTag()), [kwargs...])
```
which creates a [`LogicalVariable`](@ref) that is associated with `MyTag()` such 
that the reformulation binary variables are of the form:
```julia
@variable(model, var_expr, Bin, MyTag(), [kwargs...])
```
"""
struct Logical{T}
    tag_data::T
end

"""
    LogicalVariableData

A type for storing [`LogicalVariable`](@ref)s and any meta-data they 
possess.

**Fields**
- `variable::LogicalVariable`: The logical variable object.
- `name::String`: The name of the variable.
"""
mutable struct LogicalVariableData
    variable::LogicalVariable
    name::String
end

################################################################################
#                        LOGICAL SELECTOR (CARDINALITY) SETS
################################################################################
# TODO check required methods for AbstractVectorSet: 
# All AbstractVectorSets of type S must implement:
# •  dimension, unless the dimension is
#    stored in the set.dimension field
# •  Utilities.set_dot, unless the dot
#    product between two vectors in the set
#    is equivalent to LinearAlgebra.dot.
"""
    AbstractCardinalitySet <: MOI.AbstractVectorSet

An abstract type for cardinality sets [`_MOIAtLeast`](@ref), [`_MOIExactly`](@ref),
and [`_MOIAtMost`](@ref).
"""
abstract type AbstractCardinalitySet <:_MOI.AbstractVectorSet end

"""
    _MOIAtLeast <: AbstractCardinalitySet

MOI level set for AtLeast constraints, see [`AtLeast`](@ref) for recommended syntax.
"""
struct _MOIAtLeast <: AbstractCardinalitySet
    dimension::Int
end

"""
    _MOIAtMost <: AbstractCardinalitySet

MOI level set for AtMost constraints, see [`AtMost`](@ref) for recommended syntax.
"""
struct _MOIAtMost <: AbstractCardinalitySet
    dimension::Int
end

"""
    _MOIExactly <: AbstractCardinalitySet

MOI level set for Exactly constraints, see [`Exactly`](@ref) for recommended syntax.
"""
struct _MOIExactly <: AbstractCardinalitySet
    dimension::Int
end

# Create our own JuMP level sets to infer the dimension using the expression
"""
    AtLeast{T<:Union{Int,LogicalVariableRef}} <: JuMP.AbstractVectorSet

Convenient alias for using [`_MOIAtLeast`](@ref).
"""
struct AtLeast{T<:Union{Int, LogicalVariableRef}} <: JuMP.AbstractVectorSet
    value::T
end

"""
    AtMost{T<:Union{Int,LogicalVariableRef}} <: JuMP.AbstractVectorSet

Convenient alias for using [`_MOIAtMost`](@ref).
"""
struct AtMost{T<:Union{Int, LogicalVariableRef}} <: JuMP.AbstractVectorSet
    value::T
end

"""
    Exactly <: JuMP.AbstractVectorSet

Convenient alias for using [`_MOIExactly`](@ref).
"""
struct Exactly{T<:Union{Int, LogicalVariableRef}} <: JuMP.AbstractVectorSet
    value::T 
end

# Extend JuMP.moi_set as needed
JuMP.moi_set(::AtLeast, dim::Int) = _MOIAtLeast(dim)
JuMP.moi_set(::AtMost, dim::Int) = _MOIAtMost(dim)
JuMP.moi_set(::Exactly, dim::Int) = _MOIExactly(dim)

################################################################################
#                              LOGICAL CONSTRAINTS
################################################################################
const _LogicalExpr{M} = JuMP.GenericNonlinearExpr{LogicalVariableRef{M}}

"""
    ConstraintData{C <: JuMP.AbstractConstraint}

A type for storing constraint objects in [`GDPData`](@ref) and any meta-data 
they possess.

**Fields**
- `constraint::C`: The constraint.
- `name::String`: The name of the proposition.
"""
mutable struct ConstraintData{C <: JuMP.AbstractConstraint}
    constraint::C
    name::String
end

"""
    LogicalConstraintIndex

A type for storing the index of a logical constraint.

**Fields**
- `value::Int64`: The index value.
"""
struct LogicalConstraintIndex
    value::Int64
end

"""
    LogicalConstraintRef{M <: JuMP.AbstractModel}

A type for looking up logical constraints.
"""
struct LogicalConstraintRef{M <: JuMP.AbstractModel}
    model::M
    index::LogicalConstraintIndex
end

################################################################################
#                              DISJUNCT CONSTRAINTS
################################################################################
"""
    Disjunct

Used as a tag for constraints that will be used in disjunctions. This is done via 
the following syntax:
```julia-repl
julia> @constraint(model, [constr_expr], Disjunct)

julia> @constraint(model, [constr_expr], Disjunct(lvref))
```
where `lvref` is a [`LogicalVariableRef`](@ref) that will ultimately be associated 
with the disjunct the constraint is added to. If no `lvref` is given, then one is 
generated when the disjunction is created.
"""
struct Disjunct{M <: JuMP.AbstractModel}
    indicator::LogicalVariableRef{M}
end

# Create internal type for temporarily packaging constraints for disjuncts
struct _DisjunctConstraint{C <: AbstractConstraint, L <: LogicalVariableRef}
    constr::C
    lvref::L
end

"""
    DisjunctConstraintIndex

A type for storing the index of a [`Disjunct`](@ref).

**Fields**
- `value::Int64`: The index value.
"""
struct DisjunctConstraintIndex
    value::Int64
end

"""
    DisjunctConstraintRef{M <: JuMP.AbstractModel}

A type for looking up disjunctive constraints.
"""
struct DisjunctConstraintRef{M <: JuMP.AbstractModel}
    model::M
    index::DisjunctConstraintIndex
end

################################################################################
#                              DISJUNCTIONS
################################################################################
"""
    Disjunction{M <: JuMP.AbstractModel} <: JuMP.AbstractConstraint

A type for a disjunctive constraint that is comprised of a collection of 
disjuncts of indicated by a unique [`LogicalVariableIndex`](@ref).

**Fields**
- `indicators::Vector{LogicalVariableref}`: The references to the logical variables 
(indicators) that uniquely identify each disjunct in the disjunction.
- `nested::Bool`: Is this disjunction nested within another disjunction?
"""
struct Disjunction{M <: JuMP.AbstractModel} <: JuMP.AbstractConstraint
    indicators::Vector{LogicalVariableRef{M}}
    nested::Bool
end

"""
    DisjunctionIndex

A type for storing the index of a [`Disjunction`](@ref).

**Fields**
- `value::Int64`: The index value.
"""
struct DisjunctionIndex
    value::Int64
end

"""
    DisjunctionRef{M <: JuMP.AbstractModel}

A type for looking up disjunctive constraints.
"""
struct DisjunctionRef{M <: JuMP.AbstractModel}
    model::M
    index::DisjunctionIndex
end

################################################################################
#                              CLEVER DICTS
################################################################################
## Extend the CleverDicts key access methods
# index_to_key
function _MOIUC.index_to_key(::Type{LogicalVariableIndex}, index::Int64)
    return LogicalVariableIndex(index)
end
function _MOIUC.index_to_key(::Type{DisjunctConstraintIndex}, index::Int64)
    return DisjunctConstraintIndex(index)
end
function _MOIUC.index_to_key(::Type{DisjunctionIndex}, index::Int64)
    return DisjunctionIndex(index)
end
function _MOIUC.index_to_key(::Type{LogicalConstraintIndex}, index::Int64)
    return LogicalConstraintIndex(index)
end

# key_to_index
function _MOIUC.key_to_index(key::LogicalVariableIndex)
    return key.value
end
function _MOIUC.key_to_index(key::DisjunctConstraintIndex)
    return key.value
end
function _MOIUC.key_to_index(key::DisjunctionIndex)
    return key.value
end
function _MOIUC.key_to_index(key::LogicalConstraintIndex)
    return key.value
end

################################################################################
#                              SOLUTION METHODS
################################################################################
"""
    AbstractSolutionMethod

An abstract type for solution methods used to solve `GDPModel`s.
"""
abstract type AbstractSolutionMethod end

"""
    AbstractReformulationMethod <: AbstractSolutionMethod

An abstract type for reformulation approaches used to solve `GDPModel`s.
"""
abstract type AbstractReformulationMethod <: AbstractSolutionMethod end

"""
    BigM{T} <: AbstractReformulationMethod

A type for using the big-M reformulation approach for disjunctive constraints.

**Fields**
- `value::T`: Big-M value (default = `1e9`).
- `tight::Bool`: Attempt to tighten the Big-M value (default = `true`)?
"""
struct BigM{T} <: AbstractReformulationMethod
    value::T
    tighten::Bool
    function BigM(val::T = 1e9, tight = true) where {T}
        new{T}(val, tight)
    end
end

"""
    MBM{O, T, L <: LogicalVariableRef} <: AbstractReformulationMethod

A type for using the multiple big-M reformulation approach for disjunctive constraints.

**Fields**
- `optimizer::O`: Optimizer to use when solving mini-models (required).
- `default_M::T`: Default big-M value to use if no big-M is specified for a logical variable (1e9).
"""
mutable struct MBM{O, T} <: AbstractReformulationMethod
    optimizer::O
    default_M::T
    
    # Constructor with optimizer (required) and optional default_M
    function MBM(optimizer::O, default_M::T = 1e9) where {O, T}
        new{O, T}(optimizer, default_M)
    end
end

mutable struct _MBM{O, T, M <: JuMP.AbstractModel} <: AbstractReformulationMethod
    optimizer::O
    M::Dict{LogicalVariableRef{M}, Any}
    default_M::T
    subproblem_indicators::Vector{LogicalVariableRef{M}}
    # Cached submodels: indicator => GDPSubmodel.
    # Typed Any so extensions can store different types.
    model_cache::Dict{LogicalVariableRef{M}, Any}

    function _MBM(method::MBM{O, T}, model::M) where {O, T, M <: JuMP.AbstractModel}
        new{O, T, M}(
            method.optimizer,
            Dict{LogicalVariableRef{M}, Any}(),
            method.default_M,
            Vector{LogicalVariableRef{M}}(),
            Dict{LogicalVariableRef{M}, Any}()
        )
    end
end

"""
    Hull{T} <: AbstractReformulationMethod

A type for using the convex hull reformulation approach for disjunctive
constraints.

**Fields**
- `value::T`: epsilon value for nonlinear hull reformulations (default = `1e-6`).
- `quadratic::Symbol`: reformulation used for quadratic disjunct
  constraints (default = `:epsilon`). Options are:
  - `:epsilon`: ε-approximated perspective (Furman, Sawaya & Grossmann
    2020).
  - `:exact`: exact hull (Gusev & Bernal Neira 2025); routes each
    constraint to CEHR when its quadratic part is convex and to GEHR
    otherwise (equality constraints always use GEHR since they are
    nonconvex).
  - `:gehr`: always use the General Exact Hull Reformulation.
  - `:cehr`: always use the Conic Exact Hull Reformulation (errors on
    nonconvex quadratic constraints).
  - `:cehr_conic`: CEHR with the cone written out explicitly as a
    rotated second-order cone via a spectral factorization of the
    quadratic part, for solvers that consume cones natively (errors on
    nonconvex quadratic constraints).
"""
struct Hull{T} <: AbstractReformulationMethod
    value::T
    quadratic::Symbol
    function Hull(ϵ::T = 1e-6; quadratic::Symbol = :epsilon) where {T}
        if !(quadratic in (:epsilon, :exact, :gehr, :cehr, :cehr_conic))
            error("Invalid `quadratic` option `:$(quadratic)`. Choose " *
                  "from `:epsilon`, `:exact`, `:gehr`, `:cehr`, or " *
                  "`:cehr_conic`.")
        end
        new{T}(ϵ, quadratic)
    end
end

# temp struct to store variable disaggregations (reset for each disjunction)
mutable struct _Hull{V <: JuMP.AbstractVariableRef, T} <: AbstractReformulationMethod
    value::T
    quadratic::Symbol
    disjunction_variables::Dict{V, Vector{V}}
    disjunct_variables::Dict{Tuple{V, Union{V, JuMP.GenericAffExpr{T, V}}}, V}
    function _Hull(method::Hull{T}, vrefs::Set{V}) where {T, V <: JuMP.AbstractVariableRef}
        new{V, T}(
            method.value,
            method.quadratic,
            Dict{V, Vector{V}}(vref => V[] for vref in vrefs),
            Dict{Tuple{V, Union{V, JuMP.GenericAffExpr{T, V}}}, V}()
        )
    end
end

"""
    CuttingPlanes{O,T} <: AbstractReformulationMethod

A type for using the cutting planes approach for disjunctive constraints.

**Fields**
- `optimizer::O`: Optimizer to use when solving mini-models (required).
- `max_iter::Int`: Number of iterations (default = `3`).
- `seperation_tolerance::T`: Tolerance for the separation problem (default = `1e-6`).
- `final_reform_method::AbstractReformulationMethod`: Final reformulation 
method to use after cutting planes (default = `BigM()`).
- `M_value::T`: Big-M value to use in the final reformulation (default = `1e9`).
"""
struct CuttingPlanes{O, T} <: AbstractReformulationMethod
    optimizer::O;
    max_iter::Int
    seperation_tolerance::T
    final_reform_method::AbstractReformulationMethod
    M_value::T
    function CuttingPlanes(
        optimizer::O;
        max_iter::Int = 3,
        seperation_tolerance::T = 1e-6,
        final_reform_method = BigM(),
        M_value::T = 1e9
    ) where {O, T}
        new{O, T}(optimizer, max_iter, seperation_tolerance, final_reform_method, M_value)
    end
end

################################################################################
#                              GDP SUBMODEL
################################################################################

"""
    GDPSubmodel{M, V, W}

A unified submodel wrapper used by MBM and cutting plane
reformulations. It encapsulates a flat JuMP optimization
submodel built from a single disjunct's feasible region,
along with mappings back to the original model's variables.

## Fields
- `model::M`: The JuMP submodel representing a disjunct's
   feasible region (constraints and variable bounds).
- `decision_vars::Vector{V}`: Ordered decision variables in
   the submodel, matching the original model's ordering.
- `fwd_map::Dict{V, Vector{W}}`: Forward map from original
   model variables to their submodel counterparts.
"""
struct GDPSubmodel{M <: JuMP.AbstractModel,
                   V <: JuMP.AbstractVariableRef,
                   W <: JuMP.AbstractVariableRef}
    model::M
    decision_vars::Vector{V}
    fwd_map::Dict{V, Vector{W}}
end

"""
    PSplit <: AbstractReformulationMethod

A type for using the P-split reformulation approach for disjunctive constraints.
This method partitions variables into groups and handles each group separately.

# Constructors
- `PSplit(partition::Vector{Vector{V}})`: Create a PSplit with the given 
partition of variables
- `PSplit(n_parts::Int, model::JuMP.AbstractModel)`: Automatically partition 
model variables into `n_parts` groups

# Fields
- `partition::Vector{Vector{V}}`: The partition of variables, where each inner 
vector represents a group of variables that will be handled together
"""
struct PSplit{V <: JuMP.AbstractVariableRef} <: AbstractReformulationMethod
    partition::Vector{Vector{V}}

    function PSplit(partition::Vector{Vector{V}}) where 
        {V <: JuMP.AbstractVariableRef}
        new{V}(partition)
    end

    function PSplit(n_parts::Int, model::JuMP.AbstractModel)
        n_parts > 0 || error("Number of partitions must be 
        positive, got $n_parts")
        variables = collect_all_vars(model)
        n_vars = length(variables)
        
        n_parts = min(n_parts, n_vars)
        n_parts > 0 || error("No variables found in the model")
        
        base_size = n_vars ÷ n_parts
        remaining = n_vars % n_parts
        
        partition = Vector{Vector{eltype(variables)}}()
        start_idx = 1
        
        for i in 1:n_parts
            part_size = i <= remaining ? base_size + 1 : base_size
            end_idx = start_idx + part_size - 1
            push!(partition, variables[start_idx:end_idx])
            start_idx = end_idx + 1
        end
        
        return PSplit(partition)
    end
end

# temp struct to store variable disaggregations (reset for each disjunction)
mutable struct _PSplit{V <: JuMP.AbstractVariableRef, M <: JuMP.AbstractModel, T} <: AbstractReformulationMethod
    partition::Vector{Vector{V}}
    sum_constraints::Dict{LogicalVariableRef{M}, Vector{<:AbstractConstraint}}
    hull::_Hull{V, T}
    function _PSplit(method::PSplit{V}, model::M) where 
        {V <: JuMP.AbstractVariableRef, M <: JuMP.AbstractModel}
        T = JuMP.value_type(M)
        new{V, M, T}(
            method.partition, 
            Dict{LogicalVariableRef{M}, Vector{<:AbstractConstraint}}(), 
            _Hull(Hull(), Set{V}())
        )
    end
end

"""
    Indicator <: AbstractReformulationMethod

A type for using indicator constraint approach for linear disjunctive constraints.
"""
struct Indicator <: AbstractReformulationMethod end

################################################################################
#                              GDP Data
################################################################################
"""
    GDPData{M <: JuMP.AbstractModel, V <: JuMP.AbstractVariableRef, CrefType, ValueType}

The core type for storing information in a [`GDPModel`](@ref).
"""
mutable struct GDPData{M <: JuMP.AbstractModel, V <: JuMP.AbstractVariableRef, C, T}
    # Objects
    logical_variables::_MOIUC.CleverDict{LogicalVariableIndex, LogicalVariableData}
    logical_constraints::_MOIUC.CleverDict{LogicalConstraintIndex, ConstraintData}
    disjunct_constraints::_MOIUC.CleverDict{DisjunctConstraintIndex, ConstraintData}
    disjunctions::_MOIUC.CleverDict{DisjunctionIndex, ConstraintData{Disjunction{M}}}

    # Exactly one constraint mappings
    exactly1_constraints::Dict{DisjunctionRef{M}, LogicalConstraintRef{M}}

    # Indicator variable mappings
    indicator_to_binary::Dict{LogicalVariableRef{M}, Union{V, JuMP.GenericAffExpr{T, V}}}
    indicator_to_constraints::Dict{LogicalVariableRef{M}, Vector{Union{DisjunctConstraintRef{M}, DisjunctionRef{M}}}}
    constraint_to_indicator::Dict{Union{DisjunctConstraintRef{M}, DisjunctionRef{M}}, LogicalVariableRef{M}} # needed for deletion

    # Helpful metadata for most reformulations (not just one of them)
    variable_bounds::Dict{V, Tuple{T, T}}

    # Reformulation variables and constraints
    reformulation_variables::Vector{V}
    reformulation_constraints::Vector{C}

    # Solution data
    solution_method::Union{Nothing, AbstractSolutionMethod}
    ready_to_optimize::Bool

    # Default constructor
    function GDPData{M, V, C}() where {M <: JuMP.AbstractModel, V <: JuMP.AbstractVariableRef, C}
        T = JuMP.value_type(M)
        new{M, V, C, T}(_MOIUC.CleverDict{LogicalVariableIndex, LogicalVariableData}(),
            _MOIUC.CleverDict{LogicalConstraintIndex, ConstraintData}(),
            _MOIUC.CleverDict{DisjunctConstraintIndex, ConstraintData}(),
            _MOIUC.CleverDict{DisjunctionIndex, ConstraintData{Disjunction{M}}}(),
            Dict{DisjunctionRef{M}, LogicalConstraintRef{M}}(),
            Dict{LogicalVariableRef{M}, Union{V, JuMP.GenericAffExpr{T, V}}}(),
            Dict{LogicalVariableRef{M}, Vector{Union{DisjunctConstraintRef{M}, DisjunctionRef{M}}}}(),
            Dict{Union{DisjunctConstraintRef{M}, DisjunctionRef{M}}, LogicalVariableRef{M}}(),
            Dict{V, Tuple{T, T}}(),
            Vector{V}(),
            Vector{C}(),
            nothing,
            false,
        )
    end
end

################################################################################
#                              VARIABLE INFO
################################################################################
"""
    VariableProperties{L, U, F, S, SET, T}

A type for storing variable properties and attributes that can be applied to JuMP variables.
This is used to capture and transfer variable information between models during reformulation.

**Fields**
- `info::JuMP.VariableInfo{L, U, F, S}`: JuMP's VariableInfo struct containing bounds, fixed values, start values, and binary/integer constraints.
- `name::String`: The variable name.
- `set::SET`: The constraint set the variable belongs to (if any), obtained via `JuMP.moi_set`.
- `variable_type::T`: The variable type information, critical for extensions.

**Type Parameters**
- `L, U, F, S`: Type parameters from JuMP.VariableInfo for lower bound, upper bound, fixed value, and start value types.
- `SET`: Type of the constraint set the variable belongs to.
- `T`: Type of the variable type information.

**Constructor**
`VariableProperties(vref::JuMP.GenericVariableRef{T})` creates a VariableProperties instance
from a JuMP variable reference, automatically extracting all relevant properties.
"""
mutable struct VariableProperties{L, U, F, S, SET, T}
    info::JuMP.VariableInfo{L, U, F, S}
    name::String
    set::SET
    variable_type::T
end 

function VariableProperties(vref::JuMP.GenericVariableRef{T}) where T
    info = get_variable_info(vref)
    name = JuMP.name(vref)
    set = JuMP.is_variable_in_set(vref) ? JuMP.moi_set(JuMP.constraint_object(JuMP.VariableInSetRef(vref))) : nothing
    return VariableProperties(info, name, set, nothing)
end

function VariableProperties(vref::JuMP.AbstractVariableRef)
    info = get_variable_info(vref)
    name = JuMP.name(vref)
    return VariableProperties(info, name, nothing, nothing)
end

"""
    VariableProperties(expr)::VariableProperties

Creates a `VariableProperties` object with blank variable info (no bounds, not fixed,
not binary/integer) from an expression. The `expr` argument is provided for
extensions to infer additional properties.

## Arguments
- `expr`: Expression for extensions to extract metadata from

## Returns
A `VariableProperties` object with blank info.
"""
function VariableProperties(expr)
    info = _free_variable_info()
    return VariableProperties(info, "", nothing, nothing)
end