"""Solver interface for optimization and constraint programming.

Exports:
    SolverModel: Build optimization models with variables, objectives, constraints.
    make_name: Construct variable/constraint names.
    all_different: Global constraint for pairwise distinct values.
    implies: Logical implication constraint.
    special_ordered_set_type_2: SOS2 constraint for piecewise-linear modeling.
"""

from __future__ import annotations

from typing import Any

from relationalai.semantics import Integer, TypeVar, Any as CoreAny
from relationalai.semantics.frontend import base as b

from .solver_model import SolverModel, make_name

# Library for solver-specific aggregate relations.
library = b.Library("aggregates")

# Package version.
__version__ = "0.0.0"

__all__ = [
    "SolverModel",
    "make_name",
    "all_different",
    "implies",
    "special_ordered_set_type_2",
]


# =============================================================================
# Constraint Constructors
# =============================================================================


_implies = library.Relation(
    "implies", fields=[b.Field.input("left", CoreAny), b.Field.input("right", CoreAny)]
)


def implies(left: Any, right: Any) -> b.Expression:
    """Create a logical implication constraint (left => right).

    Args:
        left: Antecedent of the implication.
        right: Consequent of the implication.

    Returns:
        Expression representing the implication constraint.
    """
    return _implies(left, right)


_all_different = library.Relation(
    "all_different", fields=[b.Field.input("over", TypeVar)]
)


def all_different(*args: b.Value) -> b.Aggregate:
    """Create an all-different global constraint.

    Requires all variables to take pairwise distinct values.

    Args:
        *args: Decision variables that must all differ.

    Returns:
        Aggregate expression for the all-different constraint.
    """
    return b.Aggregate(_all_different, *args)


_special_ordered_set_type_2 = library.Relation(
    "special_ordered_set_type_2",
    fields=[b.Field.input("index", Integer), b.Field.input("variables", CoreAny)],
)


def special_ordered_set_type_2(index: Any, variables: Any) -> b.Aggregate:
    """Create a special ordered set type 2 (SOS2) constraint.

    At most two variables can be non-zero, and they must be consecutive
    in the given order. Useful for piecewise-linear approximations.

    Args:
        index: Expression specifying variable order.
        variables: Decision variables forming the ordered set.

    Returns:
        Aggregate expression for the SOS2 constraint.
    """
    return b.Aggregate(_special_ordered_set_type_2, index, variables)
