"""Solver model implementation using expression graph format.

Provides SolverModel for defining optimization problems with variables,
objectives, and constraints. Expressions are serialized to a Hash-based AST
format using numeric operator codes.

## Architecture

This module implements a declarative interface for optimization modeling:

1. **Model Building**: Users define decision variables, objectives, and constraints
   using the SolverModel API (solve_for, minimize, maximize, satisfy).

2. **Expression Rewriting**: SymbolifyContext converts DSL expressions into a
   Hash-based Abstract Syntax Tree (AST) suitable for serialization:
   - Each AST node is identified by a Hash value
   - Operators are encoded as numeric codes (see _*_CODES constants)
   - Arguments reference child nodes via Hash IDs or literal values

3. **Serialization Format**: The AST is exported to CSV tables:
   - Variables: properties (name, type, bounds, start)
   - Expressions: root nodes (objectives/constraints)
   - Operators: Hash node -> operator code
   - Arguments: Hash node -> child nodes (ordered or unordered)

4. **Solver Integration**: Serialized model is sent to external solver service,
   results are imported back from Snowflake views.

## Example AST Structure

For an expression like "x + y * 2":
- Node1 (Hash=abc123): operator=MUL(22), args=[y_hash, 2]
- Node2 (Hash=def456): operator=ADD(20), args=[x_hash, abc123]
- Expression root: Hash=def456

## Classes

    ExpressionFormatter: Formats expression trees as human-readable strings.
    SolverModel: Main interface for building optimization models.
    Symbolic: Wrapper marking expressions containing solver variables.
    SymbolifyContext: Manages expression rewriting to symbolic Hash-based AST form.
"""

from __future__ import annotations

from typing import Any, Optional
from collections import defaultdict
import uuid
from pandas import DataFrame

from relationalai.semantics.frontend import base as b
from relationalai.semantics import std, String, Float, Integer, Any as CoreAny
from relationalai.semantics.frontend.core import Hash, cast as core_cast
from relationalai.util.schema import get_provider
from v0.relationalai.semantics.reasoners.optimization import Solver
from v0.relationalai.clients.resources.snowflake import Resources as SnowflakeResources

# Library for custom aggregation relations used in solver expression handling.
library = b.Library("aggregates")


# =============================================================================
# Type Aliases
# =============================================================================

# Type aliases for improved type safety and clarity
DSLExpression = b.Expression | b.Variable | b.Concept | b.Ref | b.Chain
NumericValue = b.Variable | float | int
SymbolicResult = Optional["Symbolic"] | Any


# =============================================================================
# Solver Serialization Format Constants
# =============================================================================
# These codes define the wire format for solver model serialization.
# Codes are grouped by category for clarity and easy external synchronization.

# Variable types (codes 0-9, reserve 3-9 for future)
_VARIABLE_TYPE_CODES: dict[str, int] = {
    "cont": 0,  # Continuous variables
    "int": 1,  # Integer variables
    "bin": 2,  # Binary variables (0 or 1)
}

# Expression types (codes 10-19, reserve 13-19 for future)
_EXPRESSION_TYPE_CODES: dict[str, int] = {
    "min_objective": 10,
    "max_objective": 11,
    "constraint": 12,
}

# Arithmetic operators (codes 20-39, reserve 25-39 for future) - using PyRel operator names
_ARITHMETIC_OPERATOR_CODES: dict[str, int] = {
    "+": 20,  # Addition
    "-": 21,  # Subtraction
    "*": 22,  # Multiplication
    "/": 23,  # Division (real division)
    "^": 24,  # Power/exponentiation
}

# Unsupported arithmetic operators - these are not supported by any MOI solver backend.
# They are reserved for future use if solver support is added.
_UNSUPPORTED_ARITHMETIC_OPS: set[str] = {"%", "//"}

# Math functions (codes 40-69, reserve 48-69 for future) - using PyRel function names
_MATH_FUNCTION_CODES: dict[str, int] = {
    "abs": 40,  # Absolute value
    "exp": 41,  # Exponential (e^x)
    "natural_log": 42,  # Natural logarithm (ln)
    "log": 43,  # Logarithm with base
    "sqrt": 44,  # Square root
    "square": 45,  # Square (x^2)
    "cbrt": 46,  # Cube root
    "power": 47,  # Power function (alternative to ^)
}

# Unsupported math functions - not supported by MOI/MiniZinc solver backends.
_UNSUPPORTED_MATH_OPS: set[str] = {"floor", "ceil"}

# Comparison operators (codes 70-79, reserve 76-79 for future) - using PyRel operator names
_COMPARISON_OPERATOR_CODES: dict[str, int] = {
    "=": 70,  # Equal
    "!=": 71,  # Not equal
    "<": 72,  # Less than
    "<=": 73,  # Less than or equal
    ">": 74,  # Greater than
    ">=": 75,  # Greater than or equal
}

# Logical operators (codes 80-89, reserve 81-89 for future)
_LOGICAL_OPERATOR_CODES: dict[str, int] = {
    "implies": 80,  # Logical implication (a => b)
}

# Aggregate functions (codes 90-109, reserve 94-109 for future) - using PyRel aggregate names
_AGGREGATE_FUNCTION_CODES: dict[str, int] = {
    "sum": 90,  # Sum aggregate
    "min": 91,  # Minimum aggregate
    "max": 92,  # Maximum aggregate
    "count": 93,  # Count aggregate
}

# Global constraints (codes 110-129, reserve 113-129 for future)
_GLOBAL_CONSTRAINT_CODES: dict[str, int] = {
    "all_different": 110,  # All variables must have different values
    "special_ordered_set_type_1": 111,  # SOS1: at most one variable non-zero
    "special_ordered_set_type_2": 112,  # SOS2: at most two consecutive variables non-zero
}

# Combine operator dictionaries for convenience
_FIRST_ORDER_OPERATOR_CODES = _ARITHMETIC_OPERATOR_CODES | _MATH_FUNCTION_CODES
_FIRST_ORDER_COMPARISON_CODES = _COMPARISON_OPERATOR_CODES | _LOGICAL_OPERATOR_CODES
_HIGHER_ORDER_OPERATOR_CODES = _AGGREGATE_FUNCTION_CODES | _GLOBAL_CONSTRAINT_CODES

# Operators that use infix notation in display
_INFIX_OPERATORS: frozenset[str] = frozenset(_ARITHMETIC_OPERATOR_CODES.keys())
_INFIX_OPERATORS_AND_COMPARISONS: frozenset[str] = _INFIX_OPERATORS | frozenset(
    _COMPARISON_OPERATOR_CODES.keys()
)

# Reverse mapping from numeric codes to names (for display/debugging)
_CODE_TO_NAME: dict[int, str] = {
    v: k
    for k, v in (
        _VARIABLE_TYPE_CODES
        | _EXPRESSION_TYPE_CODES
        | _FIRST_ORDER_OPERATOR_CODES
        | _FIRST_ORDER_COMPARISON_CODES
        | _HIGHER_ORDER_OPERATOR_CODES
    ).items()
}


# =============================================================================
# Utilities
# =============================================================================


def make_name(*args: Any, sep: Optional[str] = "_") -> Any:
    """Construct a name by concatenating arguments.

    Args:
        *args: Arguments to concatenate (strings, numbers, or lists).
        sep: Separator between arguments. Defaults to "_".

    Returns:
        String or expression representing the concatenated name.

    Raises:
        ValueError: If no arguments provided.
    """
    if not args:
        raise ValueError("No arguments provided to `make_name`")

    if len(args) == 1:
        arg = args[0]
        if isinstance(arg, str):
            return arg
        elif isinstance(arg, list):
            return make_name(*arg, sep=sep)
        else:
            return std.strings.string(arg)

    if sep:
        str_args: list[Any] = []
        for a in args:
            str_args.append(std.strings.string(a))
            str_args.append(sep)
        str_args.pop()  # Remove trailing separator.
    else:
        str_args = [std.strings.string(a) for a in args]

    return std.strings.concat(*str_args)


# Relation for computing hash values from argument lists.
_hash = library.Relation(
    "hash", [b.Field.input("args", CoreAny, is_list=True), b.Field("hash", Hash)]
)


def hash_with_res(args: list[b.Value], res: b.Value) -> b.Expression:
    """Compute a hash of the given arguments and bind to res."""
    return _hash(b.TupleVariable(args), res)


# Mapping of computed attribute names to (concept_attr, type_code) for __getattr__.
_MODEL_INFO_ATTRS: dict[str, tuple[str, Optional[int]]] = {
    "num_variables": ("Variable", None),
    "num_min_objectives": ("Expression", _EXPRESSION_TYPE_CODES["min_objective"]),
    "num_max_objectives": ("Expression", _EXPRESSION_TYPE_CODES["max_objective"]),
    "num_constraints": ("Expression", _EXPRESSION_TYPE_CODES["constraint"]),
}

# Valid result info keys that can be accessed as attributes
_RESULT_INFO_KEYS = frozenset(
    [
        "error",
        "termination_status",
        "solver_version",
        "printed_model",
        "solve_time_sec",
        "objective_value",
        "result_count",
    ]
)


# =============================================================================
# Main Solver Model Class
# =============================================================================


class SolverModel:
    """Solver model interface for optimization problems.

    Provides a declarative interface for defining decision variables,
    objective functions, and constraints. Serializes to expression graph format.

    Attributes:
        Variable: Concept for decision variables.
        Expression: Concept for objectives and constraints.
    """

    def __init__(self, model: b.Model, num_type: b.Concept) -> None:
        """Initialize a solver model for mathematical optimization problems.

        Creates Variable and Expression concepts in the provided model, plus internal
        relations for storing the problem structure (operators, arguments, etc.). All
        solver variables will use num_type for their numeric values.

        The num_type determines the default numeric type for all variables:
        - Float: Use for continuous variables (real numbers)
        - Integer: Use for discrete variables (whole numbers)

        Individual variables can override the type using solve_for(type="bin")
        for binary (0/1) variables.

        Args:
            model: The RelationalAI model to attach solver constructs to. Multiple
                   SolverModels can share the same model if needed.
            num_type: Numeric data type for variable values. Must be Float or Integer.
                     Choose Float for continuous optimization, Integer for discrete.

        Raises:
            ValueError: If num_type is neither Float nor Integer.

        Example:
            >>> from relationalai.semantics import Model, Float, Integer
            >>>
            >>> # For continuous optimization (LP, QP, NLP)
            >>> model = Model("continuous_problem")
            >>> s = SolverModel(model, Float)
            >>>
            >>> # For discrete optimization (IP, MIP)
            >>> model2 = Model("discrete_problem")
            >>> s2 = SolverModel(model2, Integer)
        """
        if not (num_type is Float or num_type is Integer):
            raise ValueError(
                f"Invalid numerical type '{num_type}'. Must be Float or Integer."
            )
        self._model = model
        self._num_type = num_type
        self._model_id = str(uuid.uuid4()).upper().replace("-", "_")

        # Maps relationships to their corresponding variable concepts.
        self._variable_relationships: dict[int, b.Concept] = {}

        Concept = model.Concept
        Property = model.Property

        # Variable concept and its properties.
        Variable = self.Variable = Concept("Variable")
        Variable.name = Property(f"{Variable} has {String:name}")
        Variable.type = Property(f"{Variable} has {Integer:type}")
        Variable.lower = Property(f"{Variable} has {num_type:lower}")
        Variable.upper = Property(f"{Variable} has {num_type:upper}")
        Variable.start = Property(f"{Variable} has {num_type:start}")

        # Expression concept and its properties.
        Expression = self.Expression = Concept("Expression")
        Expression.name = Property(f"{Expression} has {String:name}")
        Expression.type = Property(f"{Expression} has {Integer:type}")
        Expression.root = Property(f"{Expression} has {Hash:root}")

        # AST structure properties for expression nodes.
        self._operator = Property(f"{Hash:node} has {Integer:op}")
        self._ordered_args_hash = Property(f"{Hash:node} arg {Integer:i} is {Hash:arg}")
        self._ordered_args_data = Property(
            f"{Hash:node} arg {Integer:i} is {num_type:arg}"
        )
        self._unordered_args_hash = Property(f"{Hash:node} arg {Hash:i} is {Hash:arg}")

        # Result relations for solver outputs.
        self._result_info = model.Relationship(f"{String:name} has {String:value}")
        self._points = Property(
            f"{Integer:sol_index} has {Hash:var_hash} with {num_type:value}"
        )
        self._objective_values = Property(f"{Integer:sol_index} has {num_type:value}")

    # =========================================================================
    # Variable Handling
    # =========================================================================

    def solve_for(
        self,
        expr: b.Relationship | b.Chain | b.Expression,
        where: Optional[list[Any]] = None,
        name: Optional[Any | list[Any]] = None,
        type: Optional[str] = None,
        lower: Optional[b.Variable | float | int] = None,
        upper: Optional[b.Variable | float | int] = None,
        start: Optional[b.Variable | float | int] = None,
    ) -> None:
        """Define decision variables for the optimization problem.

        Declares which relationships represent decision variables that the solver
        should determine. Each relationship can be used at most once.

        For scalar variables, pass a simple Relationship. For indexed variables
        (e.g., x[i] for i in 1..10), pass a Chain or Expression that will create
        one variable per grounded instance.

        Args:
            expr: Variable definition. Can be:
                  - Relationship: Simple scalar variable (e.g., model.Relationship(f"{Float:x}"))
                  - Chain: Indexed by concept (e.g., Item.value)
                  - Expression: Parametrized (e.g., distance(city1, city2))
            where: Optional grounding conditions to filter which variable instances
                   to create. For example, where=[i >= 0, i < 10] creates variables
                   only for valid indices.
            name: Variable name(s) for display. Can be:
                  - String: Scalar name (e.g., "x")
                  - List: Template for indexed vars (e.g., ["x", Item.id])
                  - None: Auto-generated from relationship
            type: Variable type override. Options:
                  - 'cont': Continuous real-valued (uses Float)
                  - 'int': Integer-valued (uses Integer)
                  - 'bin': Binary 0/1 (uses Integer, adds bounds 0-1)
                  Default: 'cont' if num_type is Float, 'int' if Integer
            lower: Lower bound on variable values. Can be constant or expression.
            upper: Upper bound on variable values. Can be constant or expression.
            start: Initial value hint for solver (warm start).

        Raises:
            ValueError: If the relationship is already registered as a variable.
            TypeError: If expr is not a valid variable definition.

        Examples:
            >>> # Scalar variables
            >>> x = model.Relationship(f"{Float:x}")
            >>> s.solve_for(x, name="x", lower=0, upper=10)
            >>>
            >>> # Indexed variables with bounds
            >>> Item = model.Concept("Item", identify_by={"i": Integer})
            >>> Item.quantity = model.Property(f"{Item} has {Float:quantity}")
            >>> s.solve_for(Item.quantity, name=["q", Item.i], lower=0)
            >>>
            >>> # Binary variables with conditions
            >>> Task = model.Concept("Task", identify_by={"id": Integer})
            >>> Task.assigned = model.Property(f"{Task} has {Integer:assigned}")
            >>> s.solve_for(Task.assigned, type="bin", where=[Task.id > 0])
        """
        # Extract relationship and args from expression.
        relationship, args = self._extract_relationship(expr)
        key = _get_key(relationship)

        if key in self._variable_relationships:
            raise ValueError(
                "Variables already defined for this relationship. "
                "Each relationship can only be used once in solve_for()."
            )

        # Create specialized Variable concept for this relationship.
        Var = self._model.Concept(
            f"Variable_{relationship._id}", extends=[self.Variable]
        )
        self._variable_relationships[key] = Var

        # Build variable with field bindings (excluding value field).
        fields = {f.name: args[i] for i, f in enumerate(relationship._fields[:-1])}
        var = Var.new(var_id=relationship._id, **fields)

        # Build definitions list.
        defs: list[Any] = [
            var,
            var.name(make_name(name if name is not None else "_")),
            var.type(_VARIABLE_TYPE_CODES[self._resolve_var_type(type)]),
        ]

        # Add optional properties.
        for attr, val in [("lower", lower), ("upper", upper), ("start", start)]:
            if val is not None:
                self._check_numeric(attr, val)
                defs.append(getattr(var, attr)(val))

        self._model.define(*defs).where(*(where or []))

    def _extract_relationship(
        self, expr: b.Relationship | b.Chain | b.Expression
    ) -> tuple[b.Relationship, list[Any]]:
        """Extract relationship and args from variable expression."""
        if isinstance(expr, b.Relationship):
            return expr, [f.type for f in expr._fields]
        elif isinstance(expr, b.Chain):
            if not isinstance(expr._start, b.Concept) or not isinstance(
                expr._next, b.Relationship
            ):
                raise TypeError(
                    "Chain must have Concept start and Relationship next. "
                    "Example: solve_for(Entity.value)"
                )
            return expr._next, [f.type for f in expr._next._fields]
        elif isinstance(expr, b.Expression):
            rel = expr._op
            if not isinstance(rel, b.Relationship):
                raise TypeError(
                    f"Expression operator must be Relationship, got {type(rel).__name__}. "
                    f"Example: solve_for(my_relation(x, y))"
                )
            if len(expr._args) != len(rel._fields):
                field_names = [f.name for f in rel._fields]
                raise ValueError(
                    f"Argument count mismatch: {len(expr._args)} arguments provided, "
                    f"but relationship '{rel._name}' expects {len(rel._fields)} fields: {', '.join(field_names)}. "
                    f"Example: solve_for({rel._name}({', '.join('...' for _ in rel._fields)}))"
                )
            return rel, expr._args
        raise TypeError(
            f"Invalid expression type: {type(expr).__name__}. "
            f"Expected Relationship, Chain, or Expression. "
            f"Example: solve_for(my_relation) or solve_for(Entity.value)"
        )

    def _resolve_var_type(self, var_type: Optional[str]) -> str:
        """Resolve variable type, defaulting based on model's numeric type."""
        if var_type is None:
            return "cont" if self._num_type is Float else "int"
        if not isinstance(var_type, str):
            raise TypeError(
                f"Variable type must be a string, got {type(var_type).__name__}. "
                f"Example: solve_for(x, type='int')"
            )
        if var_type not in _VARIABLE_TYPE_CODES:
            valid_types = ", ".join(f"'{t}'" for t in _VARIABLE_TYPE_CODES.keys())
            raise ValueError(
                f"Invalid variable type '{var_type}'. Valid types are: {valid_types}. "
                f"Example: solve_for(x, type='cont', lower=0, upper=10)"
            )
        return var_type

    def _check_numeric(self, name: str, val: Any) -> None:
        """Validate that a value is numeric (Variable, float, or int)."""
        if not isinstance(val, (b.Variable, float, int)):
            raise TypeError(
                f"Parameter '{name}' must be numeric (Variable, float, or int), "
                f"got {type(val).__name__}. "
                f"Example: solve_for(x, {name}=0.0)"
            )

    # =========================================================================
    # Expression Handling
    # =========================================================================

    def minimize(
        self,
        expr: b.Variable | float | int | b.Fragment,
        name: Optional[Any | list[Any]] = None,
    ) -> None:
        """Add a minimization objective."""
        self._add_expression("min_objective", expr, name)

    def maximize(
        self,
        expr: b.Variable | float | int | b.Fragment,
        name: Optional[Any | list[Any]] = None,
    ) -> None:
        """Add a maximization objective."""
        self._add_expression("max_objective", expr, name)

    def satisfy(
        self,
        expr: b.Fragment,
        check: bool = False,
        name: Optional[Any | list[Any]] = None,
    ) -> None:
        """Add constraints from the require clause of a Fragment.

        Args:
            expr: Fragment with require clause containing constraints.
            check: If True, also validate as integrity constraint.
            name: Optional constraint name.
        """
        if not isinstance(expr, b.Fragment):
            raise TypeError(f"satisfy expects a Fragment, got {type(expr).__name__}.")
        if not expr._require:
            raise ValueError("Fragment must have a require clause.")
        if expr._select or expr._define:
            raise ValueError("Fragment must not have select or define clauses.")

        # Remove from model roots unless checking as integrity constraint.
        if not check and expr in self._model.requires:
            self._model.requires.remove(expr)
            self._model._bump_version()

        for req in expr._require:
            self._add_expression("constraint", req, name, expr._where)

    def _add_expression(
        self,
        expr_type: str,
        expr: Any,
        name: Optional[Any | list[Any]],
        where: Optional[list[Any]] = None,
    ) -> None:
        """Add an objective or constraint expression to the model."""
        ctx = SymbolifyContext(self)
        if where:
            ctx.wheres.extend(ctx.rewrite_where(*where))

        result = ctx.rewrite(expr)
        if not isinstance(result, Symbolic):
            raise ValueError(
                f"Cannot symbolify {expr}. Ensure it contains solver variables."
            )

        expr_node = self.Expression.new(root=result.expr)
        ctx.defines.extend(
            [
                expr_node,
                expr_node.type(_EXPRESSION_TYPE_CODES[expr_type]),
            ]
        )
        if name is not None:
            ctx.defines.append(expr_node.name(make_name(name)))

        self._apply_context(ctx)

    def _apply_context(self, ctx: SymbolifyContext) -> None:
        """Recursively apply defines and wheres from context tree."""
        self._model.define(*ctx.defines).where(*ctx.wheres)
        for subctx in ctx.subcontexts:
            self._apply_context(subctx)

    # =========================================================================
    # Model Inspection and Display
    # =========================================================================

    def display(self) -> None:
        """Print model summary and variable/expression tables."""
        Var, Expr = self.Variable, self.Expression

        # Query variable, expression, and expression AST data.
        n0 = Hash.ref().alias("node")
        n1 = Hash.ref().alias("arg")
        i = Integer.ref().alias("index")
        h = Hash.ref().alias("hash")
        reachable = self._model.Relationship(f"{Hash}")
        self._model.define(reachable(Expr.root))
        ctx = self._model.where(reachable(n0))
        ctx.define(reachable(n1)).where(self._ordered_args_hash(n0, i, n1))
        ctx.define(reachable(n1)).where(self._unordered_args_hash(n0, h, n1))

        # TODO(coey): In future, issue multiple queries at the same time to reduce overhead.
        # Currently causes failures when running the pytest suite or the run_all.py script.
        # See https://relationalai.slack.com/archives/C08ATF88BC6/p1770146491179649?thread_ts=1768955416.245459&cid=C08ATF88BC6

        var_df = ctx.select(Var.type, Var, Var.name, Var.lower, Var.upper).to_df()
        if not var_df.empty:
            var_types = var_df["type"].astype(int)
        else:
            from pandas import Series

            var_types = Series([], dtype=int)

        expr_df = ctx.select(Expr.type, Expr.name, Expr.root).to_df()
        if not expr_df.empty:
            expr_types = expr_df["type"].astype(int)
        else:
            from pandas import Series

            expr_types = Series([], dtype=int)

        op = Integer.ref().alias("op")
        operators_df = ctx.select(n0, op).where(self._operator(n0, op)).to_df()
        ordered_hash_df = (
            ctx.select(n0, i, n1).where(self._ordered_args_hash(n0, i, n1)).to_df()
        )
        v = self._num_type.ref().alias("value")
        ordered_data_df = (
            ctx.select(n0, i, v).where(self._ordered_args_data(n0, i, v)).to_df()
        )
        unordered_df = (
            ctx.select(n0, n1).where(self._unordered_args_hash(n0, h, n1)).to_df()
        )

        # Print summary.
        var_counts = var_types.value_counts()
        expr_counts = expr_types.value_counts()

        def get_count(counts: Any, codes: dict[str, int], key: str) -> int:
            return int(counts.get(codes[key], 0))

        n_cont = get_count(var_counts, _VARIABLE_TYPE_CODES, "cont")
        n_int = get_count(var_counts, _VARIABLE_TYPE_CODES, "int")
        n_bin = get_count(var_counts, _VARIABLE_TYPE_CODES, "bin")
        n_min = get_count(expr_counts, _EXPRESSION_TYPE_CODES, "min_objective")
        n_max = get_count(expr_counts, _EXPRESSION_TYPE_CODES, "max_objective")
        n_con = get_count(expr_counts, _EXPRESSION_TYPE_CODES, "constraint")

        print("Solver model has:")
        print(
            f"• {len(var_df)} variables: {n_cont} continuous, {n_int} integer, {n_bin} binary"
        )
        print(
            f"• {n_min} minimization objectives, {n_max} maximization objectives, {n_con} constraints"
        )

        # Print variables sorted by name.
        if not var_df.empty:
            var_df["type"] = var_types.map(lambda x: _CODE_TO_NAME.get(x, str(x)))
            cols = ["name", "type"] + [
                c for c in ("lower", "upper") if var_df[c].notna().any()
            ]
            print("\nVariables:")
            print(var_df.sort_values(by="name")[cols].to_string(index=False))  # type: ignore[call-overload]

        # Build and print expressions by type.
        self._make_expr_strings(
            expr_df,
            var_df,
            operators_df,
            unordered_df,
            ordered_data_df,
            ordered_hash_df,
        )
        has_names = expr_df["name"].notna().any() if not expr_df.empty else False

        expr_type_labels = [
            ("min_objective", "Minimization objectives"),
            ("max_objective", "Maximization objectives"),
            ("constraint", "Constraints"),
        ]
        for type_key, label in expr_type_labels:
            type_code = _EXPRESSION_TYPE_CODES[type_key]
            subset = expr_df[expr_types == type_code]
            if subset.empty:
                continue
            print(f"\n{label}:")
            if has_names:
                print(subset.sort_values(by="name")[["name", "expression"]].to_string(index=False))  # type: ignore[call-overload]
            else:
                print("\n".join(sorted(subset["expression"])))

    def _make_expr_strings(
        self,
        expr_df: DataFrame,
        var_df: DataFrame,
        operators_df: DataFrame,
        unordered_df: DataFrame,
        ordered_data_df: DataFrame,
        ordered_hash_df: DataFrame,
    ) -> None:
        """Add 'expression' column with human-readable strings to expr_df."""
        # Early return if no expressions to format
        if expr_df.empty:
            return

        # Build lookup dictionaries from query results.
        variables: dict[Any, str] = (
            dict(zip(var_df["variable"], var_df["name"])) if not var_df.empty else {}
        )
        operators: dict[Any, int] = (
            dict(zip(operators_df["node"], operators_df["op"]))
            if not operators_df.empty
            else {}
        )

        # Build args dict: node -> [child values in order].
        # Unordered args are added first, then ordered args sorted by index.
        args: dict[Any, list[Any]] = defaultdict(list)
        for node, child in unordered_df.itertuples(index=False):
            args[node].append(child)

        # Collect ordered args from both data and hash sources.
        ordered: dict[Any, list[tuple[Any, Any]]] = defaultdict(list)
        for df in (ordered_data_df, ordered_hash_df):
            for node, idx, child in df.itertuples(index=False):
                ordered[node].append((idx, child))

        # Sort by index and append to args.
        for node, indexed_children in ordered.items():
            indexed_children.sort(key=lambda x: x[0])
            args[node].extend(child for _, child in indexed_children)

        # Generate expression strings using the formatter.
        formatter = ExpressionFormatter(variables, operators, args)
        expr_df["expression"] = [formatter.format(root) for root in expr_df["root"]]

    # =========================================================================
    # Solve
    # =========================================================================

    def _export_model_to_csv(self) -> dict[str, str]:
        """Export model relations to CSV files for solver service."""
        model = self._model
        export_to_csv = model._get_executor().export_to_csv  # type: ignore[attr-defined]
        Var, Expr = self.Variable, self.Expression
        inputs = {}

        # Export variables
        inputs["variables"] = export_to_csv(
            model.select(
                Var.alias("VARIABLE"),
                Var.name.alias("NAME"),
                Var.type.alias("TYPE"),
                Var.lower.alias("LOWER"),
                Var.upper.alias("UPPER"),
                Var.start.alias("START"),
            )
        )

        # Export expressions
        inputs["expressions"] = export_to_csv(
            model.select(
                Expr.root.alias("NODE"),
                Expr.name.alias("NAME"),
                Expr.type.alias("TYPE"),
            )
        )

        # Export AST structure
        n0, n1, i, h, op, v = (
            Hash.ref().alias("NODE"),
            Hash.ref().alias("ARG"),
            Integer.ref().alias("INDEX"),
            Hash.ref().alias("HASH"),
            Integer.ref().alias("OP"),
            self._num_type.ref().alias("VALUE"),
        )

        inputs["operators"] = export_to_csv(
            model.select(n0, op).where(self._operator(n0, op))
        )
        inputs["ordered_args_hash"] = export_to_csv(
            model.select(n0, i, n1).where(self._ordered_args_hash(n0, i, n1))
        )
        inputs["ordered_args_data"] = export_to_csv(
            model.select(n0, i, v).where(self._ordered_args_data(n0, i, v))
        )
        inputs["unordered_args_hash"] = export_to_csv(
            model.select(n0, n1).where(self._unordered_args_hash(n0, h, n1))
        )

        return inputs

    def _import_solver_results_from_views(self):
        model = self._model
        resources = get_provider(config=self._model.config).resources
        assert isinstance(resources, SnowflakeResources)
        app_name = resources.get_app_name()

        # # Option 1: Load all tables into the GI database via api.load_data in parallel.
        # # Get DB name: to prevent sql injection, pass the variable as a separate argument to the executor
        # index_db_name = resources._exec(
        #     f"select {resources.get_app_name()}.api.get_index_db_name(?)",
        #     [model.name]
        # )[0][0]
        # logic_engine_name = resources.get_default_engine_name()
        # table_schemas = {
        #     "ancillary": {"name": String, "value": String},
        #     "points": {"sol_index": Integer, "var_hash": String, "value": self._num_type},
        #     "objective_values": {"sol_index": Integer, "value": self._num_type},
        # }
        # tables = {}
        # load_statements = []
        # for name, schema in table_schemas.items():
        #     table = model.Table(f"{app_name}.RESULTS.SOLVERS_{self._model_id}_{name.upper()}", schema)
        #     tables[name] = table
        #     table._skip_cdc = True  # Skip CDC for result tables.
        #     rel_name = table._name.replace(".", "_").lower()
        #     load_statements.append(f"call {app_name}.api.load_data('{index_db_name}', '{logic_engine_name}', '{rel_name}', '{table._name}', '')")
        # with ThreadPoolExecutor(max_workers=len(load_statements)) as ex:
        #     futures = [ex.submit(resources._exec, stmt) for stmt in load_statements]
        #     for future in futures:
        #         future.result()

        # Option 2: Execute SQL to download dataframes to the client (not ideal as it means data leaves SF perimeter).
        def load_table(name: str):
            sql = f"SELECT * FROM {app_name}.RESULTS.SOLVERS_{self._model_id}_{name.upper()}"
            df = DataFrame(resources._exec(sql))
            return None if df.empty else model.data(df)

        ancillary = load_table("ancillary")
        if ancillary is not None:
            model.define(self._result_info(ancillary.name, ancillary.value))
        objective_values = load_table("objective_values")
        if objective_values is not None:
            model.define(
                self._objective_values(
                    objective_values.sol_index, objective_values.value
                )
            )
        points = load_table("points")
        if points is not None:
            model.define(
                self._points(
                    points.sol_index,
                    std.common.parse_uuid(points.var_hash),
                    points.value,
                )
            )

    def solve(
        self, solver: Solver, log_to_console: bool = False, **kwargs: Any
    ) -> None:
        """Solve the optimization problem using the specified solver."""
        options = {**kwargs, "version": 1}
        for option_key, option_value in options.items():
            if not isinstance(option_key, str):
                raise TypeError(
                    f"Solver option keys must be strings, got {type(option_key).__name__}."
                )
            if not isinstance(option_value, (int, float, str, bool)):
                raise TypeError(
                    f"Solver option values must be int, float, str, or bool, got {type(option_value).__name__}."
                )

        # Export model to CSV format
        print("Exporting solver model")
        inputs = self._export_model_to_csv()

        # Execute solver job
        print("Executing solver job")
        payload: dict[str, Any] = {
            "solver": solver.solver_name.lower(),
            "options": options,
            "inputs": inputs,
            "model_id": self._model_id,
        }
        print(f"Solver payload: {payload}")
        solver._exec_job(payload, log_to_console=log_to_console)

        # Import results
        print("Loading solver results")
        self._import_solver_results_from_views()

        print("Finished solve")
        print()

    # =========================================================================
    # Result Handling
    # =========================================================================

    def __getattr__(self, name: str) -> Any:
        """Get model or result attribute (e.g., num_variables, termination_status)."""
        if name in _MODEL_INFO_ATTRS:
            concept_attr, type_code = _MODEL_INFO_ATTRS[name]
            concept = getattr(self, concept_attr)
            if type_code is None:
                result_df = self._model.select(
                    std.aggregates.count(concept) | 0
                ).to_df()
            else:
                result_df = self._model.select(
                    std.aggregates.count(concept).where(concept.type == type_code) | 0
                ).to_df()

            if result_df.shape != (1, 1):
                raise ValueError(
                    f"Expected exactly one value for '{name}', got shape {result_df.shape}."
                )
            return result_df.iloc[0, 0]

        if name not in _RESULT_INFO_KEYS:
            raise AttributeError(f"SolverModel object has no attribute '{name}'")

        # Query the result info relation for the given name.
        value_ref = String.ref()
        result_df = (
            self._model.select(value_ref)
            .where(self._result_info(name, value_ref))
            .to_df()
        )
        if result_df.empty:
            raise AttributeError(
                f"Result info '{name}' not available. Make sure solve() has been called."
            )
        if result_df.shape != (1, 1):
            raise ValueError(
                f"Expected exactly one value for result info '{name}', got shape {result_df.shape}."
            )
        result_value = result_df.iloc[0, 0]
        assert isinstance(result_value, str)
        if name == "solve_time_sec":
            return float(result_value)
        if name == "objective_value":
            return (
                int(result_value) if self._num_type == Integer else float(result_value)
            )
        if name == "result_count":
            return int(result_value)
        return result_value


# =============================================================================
# Symbolic Expression Classes
# =============================================================================


class Symbolic:
    """Wrapper marking an expression as containing solver variables."""

    def __init__(self, expr: Any) -> None:
        if isinstance(expr, Symbolic):
            raise TypeError("Cannot double-wrap Symbolic.")
        self.expr = expr


class ExpressionFormatter:
    """Formats expression trees as human-readable strings."""

    def __init__(
        self,
        variables: dict[Any, str],
        operators: dict[Any, int],
        args: dict[Any, list[Any]],
    ) -> None:
        self.variables = variables
        self.operators = operators
        self.args = args

    def format(self, node: Any) -> str:
        """Format an expression tree rooted at the given node."""
        return self._format_recursive(node)

    def _format_recursive(self, node: Any) -> str:
        """Recursively build human-readable string for expression tree."""
        # Base case: variable name
        if node in self.variables:
            return self.variables[node]

        # Base case: not an operator node (e.g., numeric literal)
        op_code = self.operators.get(node)
        if op_code is None:
            return str(node)

        # Recursive case: format operator and children
        op = _CODE_TO_NAME[op_code]
        children = self.args.get(node, [])

        # Format children, adding parentheses for precedence if needed
        is_infix = op in _INFIX_OPERATORS
        parts: list[str] = []
        for child in children:
            child_str = self._format_recursive(child)
            if is_infix:
                # Parenthesize child if it's also an infix/comparison operator
                child_op_code = self.operators.get(child)
                needs_parens = (
                    child_op_code is not None
                    and _CODE_TO_NAME[child_op_code] in _INFIX_OPERATORS_AND_COMPARISONS
                )
                parts.append(f"({child_str})" if needs_parens else child_str)
            else:
                parts.append(child_str)

        # Sort arguments for higher-order operators (for consistent display)
        if op in _HIGHER_ORDER_OPERATOR_CODES:
            parts.sort()

        # Format based on operator type
        if op in _INFIX_OPERATORS_AND_COMPARISONS:
            return f"{parts[0]} {op} {parts[1]}"
        return f"{op}({', '.join(parts)})"


class SymbolifyContext:
    """Rewrites DSL expressions into solver AST form using Hash nodes."""

    _REWRITE_DISPATCH: dict[type, str] = {
        b.Concept: "_rewrite_variable_ref",
        b.Ref: "_rewrite_variable_ref",
        b.Alias: "_rewrite_alias",
        b.Chain: "_rewrite_chain",
        b.Relationship: "_rewrite_relationship",
        b.Property: "_rewrite_relationship",  # Property is subclass of Relationship
        b.Match: "_rewrite_match",
        b.Union: "_rewrite_match",
        b.FilterBy: "_rewrite_filter_by",
        b.Expression: "_rewrite_expression",
        b.Aggregate: "_rewrite_aggregate",
        b.Fragment: "_rewrite_fragment",
        b.Table: "_rewrite_data_source",
        b.Data: "_rewrite_data_source",
    }

    def __init__(self, solver_model: SolverModel) -> None:
        self.model = solver_model._model
        self.solver_model = solver_model
        self.variable_map: dict[int, Any] = {}
        self.defines: list[Any] = []
        self.wheres: list[Any] = []
        self.subcontexts: list[SymbolifyContext] = []

    def _create_subcontext(self) -> SymbolifyContext:
        """Create child context for nested scope (e.g., aggregates)."""
        subctx = SymbolifyContext(self.solver_model)
        subctx.variable_map = self.variable_map.copy()
        subctx.defines = self.defines.copy()
        subctx.wheres = self.wheres.copy()
        self.subcontexts.append(subctx)
        return subctx

    # =========================================================================
    # Public Rewriting Methods
    # =========================================================================

    def rewrite_where(self, *exprs: Any) -> list[Any]:
        """Rewrite where clauses in two passes: variable relationships first, then rest."""
        results: list[Any] = []

        # Pass 1: Variable relationships populate variable_map.
        for expr in exprs:
            if isinstance(expr, b.Expression) and isinstance(expr._op, b.Relationship):
                if _get_key(expr._op) in self.solver_model._variable_relationships:
                    results.append(self._handle_variable_relationship(expr))
                    continue
            results.append(None)

        # Pass 2: Rewrite remaining expressions (may reference pass 1 variables).
        for i, expr in enumerate(exprs):
            if results[i] is None:
                key = _get_key(expr)
                # Skip if already in variable_map (handled in pass 1).
                results[i] = (
                    expr
                    if key in self.variable_map
                    else self._rewrite_nonsymbolic(expr)
                )

        return results

    def rewrite(self, expr: Any) -> Optional[Symbolic | Any]:
        """Rewrite expression using dispatch table.

        Returns:
            - Symbolic: If expression contains solver variables
            - Rewritten expression: If expression was modified but not symbolic
            - None: If no rewriting needed
        """
        # Base cases: literals and None
        if expr is None or isinstance(expr, (int, float, str, b.Literal)):
            return None

        # Dispatch to type-specific handler
        expr_type = type(expr)
        handler_name = self._REWRITE_DISPATCH.get(expr_type)

        if handler_name is not None:
            handler = getattr(self, handler_name)
            return handler(expr)

        # Provide helpful error messages for known unsupported types
        if isinstance(expr, b.New):
            raise ValueError(
                "Cannot use 'New' expressions in solver constraints. "
                "Solver models work with existing variables and relationships."
            )
        if isinstance(expr, b.FieldRef):
            raise ValueError(
                "Cannot use field indexing in solver constraints. "
                "Try accessing the field directly using dot notation."
            )
        if isinstance(expr, b.MetaRef):
            raise ValueError(
                "Cannot use meta-references in solver constraints. "
                "Solver models require concrete value expressions."
            )
        if isinstance(expr, b.AsBool):
            raise ValueError(
                "Cannot use 'as_bool()' in solver constraints. "
                "Use comparison operators (>, <, =, etc.) directly instead."
            )
        if isinstance(expr, (b.Not, b.Distinct, b.Group)):
            raise ValueError(
                f"'{expr_type.__name__}' is not supported in solver constraints. "
                f"Try restructuring your constraint to avoid using this construct."
            )
        if isinstance(expr, (b.DerivedTable, b.DerivedColumn)):
            raise ValueError(
                f"Cannot use {expr_type.__name__} directly in solver constraints. "
                f"Use variables or relationships instead."
            )

        raise NotImplementedError(
            f"Expression type '{expr_type.__name__}' is not supported in solver constraints. "
            f"Check the solver documentation for supported expression types."
        )

    # -------------------------------------------------------------------------
    # Type-specific rewrite handlers
    # -------------------------------------------------------------------------

    def _rewrite_variable_ref(self, expr: b.Concept | b.Ref) -> Optional[Symbolic]:
        """Rewrite Concept or Ref that may be a solver variable."""
        key = _get_key(expr)
        return Symbolic(self.variable_map[key]) if key in self.variable_map else None

    def _rewrite_alias(self, expr: b.Alias) -> Optional[Symbolic | Any]:
        """Rewrite Alias by delegating to its source item."""
        return self.rewrite(expr._source_item)

    def _rewrite_chain(self, expr: b.Chain) -> Optional[Symbolic]:
        """Rewrite Chain (entity.relationship). Returns Symbolic if relationship is solver var."""
        start, next_expr = expr._start, expr._next
        if not isinstance(start, (b.Concept, b.Ref)):
            raise TypeError(
                f"Chain start must be Concept or Ref, got {type(start).__name__}."
            )
        if self.rewrite(start) is not None:
            raise ValueError("Chain start must not be symbolic.")

        # Check if the relationship part is a solver variable.
        if isinstance(next_expr, b.Relationship):
            var_ref = self._get_variable_ref(next_expr, start)
            if var_ref is not None:
                self.variable_map[_get_key(next_expr)] = var_ref
                return Symbolic(var_ref)
        return None

    def _rewrite_relationship(self, expr: b.Relationship) -> Optional[Symbolic]:
        """Rewrite Relationship that may be a solver variable."""
        key = _get_key(expr)
        if key in self.variable_map:
            return Symbolic(self.variable_map[key])
        var_ref = self._get_variable_ref(expr, None)
        if var_ref is not None:
            self.variable_map[key] = var_ref
            return Symbolic(var_ref)
        return None

    def _rewrite_match(self, expr: b.Match) -> Optional[Symbolic | Any]:
        """Rewrite Match/Union expression."""
        items = expr._items
        rewritten: list[Any] = []
        has_symbolic = False

        for item in items:
            result = self.rewrite(item)
            if isinstance(result, Symbolic):
                has_symbolic = True
                rewritten.append(result.expr)
            else:
                rewritten.append(result if result is not None else item)

        if not has_symbolic and rewritten == list(items):
            return None

        # Build the appropriate Match or Union.
        if isinstance(expr, b.Union):
            combined = self.model.union(*rewritten)
        else:
            combined = b.Match(self.model, *rewritten)

        return Symbolic(combined) if has_symbolic else combined

    def _rewrite_filter_by(self, expr: b.FilterBy) -> Optional[Any]:
        """Rewrite FilterBy expression (symbolic values not allowed)."""
        new_kwargs, modified = {}, False
        for k, v in expr._kwargs.items():
            r = self.rewrite(v)
            if isinstance(r, Symbolic):
                raise ValueError(
                    f"Cannot use symbolic value in FilterBy argument '{k}'."
                )
            new_kwargs[k] = r if r is not None else v
            modified = modified or r is not None

        return expr._params[0].filter_by(**new_kwargs) if modified else None

    def _rewrite_data_source(self, expr: b.Table | b.Data) -> None:
        """Data sources (Table, Data) don't need rewriting."""
        return None

    def _rewrite_expression(self, expr: b.Expression) -> Optional[Symbolic | Any]:
        """Rewrite operator expression. Builds AST node if any arg is symbolic."""
        op = self.rewrite(expr._op)
        if isinstance(op, Symbolic):
            raise ValueError("Expression operator cannot be symbolic.")
        op = op if op is not None else expr._op

        args = expr._args
        rewritten: list[Any] = []
        has_symbolic = False
        for arg in args:
            result = self.rewrite(arg)
            if result is None:
                rewritten.append(arg)
            elif isinstance(result, Symbolic):
                has_symbolic = True
                rewritten.append(result)
            else:
                rewritten.append(result)

        if not has_symbolic and rewritten == list(args):
            return None
        if not has_symbolic:
            return b.Expression(op, rewritten)

        # Build symbolic AST node.
        op_name = op._short_name
        if not isinstance(op_name, str):
            raise TypeError(f"Operator name must be str, got {type(op_name).__name__}.")

        if op_name in _FIRST_ORDER_OPERATOR_CODES:
            op_code = _FIRST_ORDER_OPERATOR_CODES[op_name]
            rewritten = rewritten[:-1]  # Drop auto-filled output arg.
        elif op_name in _FIRST_ORDER_COMPARISON_CODES:
            op_code = _FIRST_ORDER_COMPARISON_CODES[op_name]
        else:
            # Provide helpful error for unsupported operators
            if op_name in _UNSUPPORTED_ARITHMETIC_OPS or op_name in _UNSUPPORTED_MATH_OPS:
                msg = (
                    f"Operator '{op_name}' is not supported by any solver backend "
                    f"(not available in MOI/MiniZinc/Ipopt/HiGHS)."
                )
            else:
                supported = sorted(
                    list(_FIRST_ORDER_OPERATOR_CODES.keys())
                    + list(_FIRST_ORDER_COMPARISON_CODES.keys())
                )
                msg = (
                    f"Operator '{op_name}' is not supported by the solver engine. "
                    f"Supported operators: {', '.join(supported)}"
                )
            raise NotImplementedError(msg)

        # Create Hash node unique per (expr_id, args).
        hash_args = [a.expr if isinstance(a, Symbolic) else a for a in rewritten]
        node = Hash.ref()
        self.wheres.append(hash_with_res([expr._id, *hash_args], node))
        self.defines.append(self.solver_model._operator(node, op_code))

        for i, arg in enumerate(rewritten):
            if isinstance(arg, Symbolic):
                self.defines.append(
                    self.solver_model._ordered_args_hash(node, i, arg.expr)
                )
            else:
                self.defines.append(self.solver_model._ordered_args_data(node, i, arg))

        return Symbolic(node)

    def _extract_aggregate_args(
        self, expr: b.Aggregate, op_name: str
    ) -> tuple[Any, list[Any], list[Any]]:
        """Extract symbolic arg and other args based on operator type."""
        args, proj_args = list(expr._args), list(expr._projection_args)
        if op_name == "count":
            return proj_args.pop(), args, proj_args
        elif op_name == "rank":
            if not isinstance(args[0], b.TupleVariable):
                raise TypeError("Rank aggregate expects TupleVariable as first arg.")
            return None, args[0]._items, []
        elif op_name in ("special_ordered_set_type_1", "special_ordered_set_type_2"):
            return args[1], args[:1], proj_args
        else:
            return args[0], args[1:], proj_args

    def _rewrite_aggregate(self, expr: b.Aggregate) -> Optional[Symbolic | Any]:
        """Rewrite aggregate (sum/min/max/count/etc)."""
        if expr._distinct:
            raise NotImplementedError("Distinct aggregates not yet supported.")

        op = self.rewrite(expr._op)
        if isinstance(op, Symbolic):
            raise ValueError("Aggregate operator cannot be symbolic.")
        op = op if op is not None else expr._op
        op_name = op._short_name

        # Extract arguments based on operator type
        sym_arg, args, proj_args = self._extract_aggregate_args(expr, op_name)

        # Rewrite components in subcontext
        subctx = self._create_subcontext()
        rw_where = subctx.rewrite_where(*expr._where._where)
        rw_proj = [subctx._rewrite_nonsymbolic(a) for a in proj_args]
        rw_args = [subctx._rewrite_nonsymbolic(a) for a in args]
        rw_group = [subctx._rewrite_nonsymbolic(a) for a in expr._group._args]
        rw_sym = None if op_name == "rank" else subctx.rewrite(sym_arg)

        # Early return if nothing changed
        if (
            rw_sym is None
            and rw_proj == list(proj_args)
            and rw_args == list(args)
            and rw_where == list(expr._where._where)
        ):
            return None

        # Non-symbolic case: rebuild aggregate
        if not isinstance(rw_sym, Symbolic):
            sym_arg_final = rw_sym if rw_sym is not None else sym_arg
            if op_name == "rank":
                new_args = [b.TupleVariable(rw_args), expr._args[1], expr._args[2]]
            elif op_name in (
                "special_ordered_set_type_1",
                "special_ordered_set_type_2",
            ):
                new_args = [*rw_proj, *rw_args, sym_arg_final]
            else:
                new_args = [*rw_proj, sym_arg_final, *rw_args]
            return (
                b.Aggregate(op, *new_args, distinct=expr._distinct)
                .per(*rw_group)
                .where(*rw_where, *subctx.wheres)
            )

        # Symbolic case: build Hash-based AST node
        if op_name not in _HIGHER_ORDER_OPERATOR_CODES:
            supported_aggs = list(_HIGHER_ORDER_OPERATOR_CODES.keys())
            msg = f"Aggregate '{op_name}' is not supported by the solver engine."
            if op_name in ("avg", "string_join", "rank", "limit", "top", "bottom"):
                msg += f" Only these aggregates are supported: {', '.join(sorted(supported_aggs))}"
            else:
                msg += f" Supported aggregates: {', '.join(sorted(supported_aggs))}"
            raise NotImplementedError(msg)

        sym_expr = rw_sym.expr
        node = Hash.ref()
        subctx.wheres.append(hash_with_res([expr._id, *rw_group], node))
        subctx.wheres.extend(rw_where)
        subctx.defines.append(
            self.solver_model._operator(node, _HIGHER_ORDER_OPERATOR_CODES[op_name])
        )

        if op_name == "special_ordered_set_type_2":
            if len(rw_proj) != 0 or len(rw_args) != 1:
                raise ValueError("SOS2 expects exactly 2 arguments (rank, variable).")
            subctx.defines.append(
                self.solver_model._ordered_args_hash(node, rw_args[0], sym_expr)
            )
        else:
            arg_hash = std.common.hash(*rw_proj, sym_expr)
            subctx.defines.append(
                self.solver_model._unordered_args_hash(node, arg_hash, sym_expr)
            )

        return Symbolic(self.model.select(node).where(*subctx.wheres))

    def _rewrite_fragment(self, expr: b.Fragment) -> Optional[Symbolic | Any]:
        """Rewrite inline select/where fragment."""
        if expr._define or expr._require:
            raise ValueError("Fragments with define/require not supported.")
        if len(expr._select) != 1:
            raise ValueError(
                f"Fragment must have exactly 1 select, got {len(expr._select)}."
            )

        subctx = self._create_subcontext()
        rw_where = subctx.rewrite_where(*expr._where)
        rw_select = subctx.rewrite(expr._select[0])
        wheres = self.model.where(*rw_where, *subctx.wheres)

        if isinstance(rw_select, Symbolic):
            return Symbolic(wheres.select(rw_select.expr))
        if rw_select is not None:
            return wheres.select(rw_select)
        return None

    # =========================================================================
    # Private Helper Methods
    # =========================================================================

    def _handle_variable_relationship(self, expr: b.Expression) -> Any:
        """Handle variable relationship: store ref in variable_map, return field bindings."""
        relationship = expr._op
        if not isinstance(relationship, b.Relationship):
            raise TypeError(
                f"Expected Relationship operator, got {type(relationship).__name__}."
            )

        params = expr._args
        if len(params) != len(relationship._fields):
            raise ValueError(
                f"Parameter count {len(params)} != field count {len(relationship._fields)}."
            )

        # The last parameter is the output variable (the symbolic value).
        last_param = params[-1]
        if isinstance(last_param, b.Alias):
            last_param = last_param._source_item
        if not isinstance(last_param, (b.Concept, b.Ref)):
            raise TypeError(
                f"Last parameter must be Concept or Ref, got {type(last_param).__name__}."
            )

        # Build field bindings (grounding conditions).
        fields: dict[str, Any] = {}
        for i, param in enumerate(params[:-1]):
            rewritten = self.rewrite(param)
            if isinstance(rewritten, Symbolic):
                raise ValueError(
                    f"Variable relationship parameter {i} cannot be symbolic."
                )
            fields[relationship._fields[i].name] = (
                rewritten if rewritten is not None else param
            )

        key = _get_key(relationship)
        var_ref = self.solver_model._variable_relationships[key].ref()
        self.variable_map[_get_key(last_param)] = core_cast(b.MetaRef(Hash), var_ref)

        return self.solver_model._model.where(
            *[getattr(var_ref, name) == val for name, val in fields.items()]
        )

    def _rewrite_nonsymbolic(self, expr: Any) -> Any:
        """Rewrite expression, raising if result is unexpectedly symbolic."""
        result = self.rewrite(expr)
        if isinstance(result, Symbolic):
            raise ValueError(f"Unexpected Symbolic result for {expr}.")
        return result if result is not None else expr

    def _get_variable_ref(
        self, relationship: b.Relationship, parent: b.Concept | b.Ref | None
    ) -> Optional[Any]:
        """Get Hash-cast var ref if relationship is a solver variable, else None."""
        VarConcept = self.solver_model._variable_relationships.get(
            _get_key(relationship)
        )
        if VarConcept is None:
            return None

        var_ref = VarConcept.ref()
        if parent is not None:
            self.wheres.append(getattr(var_ref, relationship._fields[0].name)(parent))
        return core_cast(b.MetaRef(Hash), var_ref)


def _get_key(expr: Any) -> int:
    """Get unique integer key for an expression."""
    key = expr.to_dict_key()
    if not isinstance(key, int):
        raise TypeError(f"Expression key must be int, got {type(key).__name__}.")
    return key
