"""Solver model implementation using expression graph format.

Provides SolverModel for defining optimization problems with variables,
objectives, and constraints. Uses numeric codes for serialization.

Classes:
    SolverModel: Main interface for building optimization models.
    Symbolic: Wrapper for symbolified solver expressions.
    SymbolifyContext: Manages expression rewriting to symbolic form.
"""

from __future__ import annotations

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

from relationalai.semantics.frontend import base as b
from relationalai.semantics import std, count, String, Float, Integer, Any as CoreAny
from relationalai.semantics.frontend.core import Hash, cast as core_cast

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


# =============================================================================
# Solver Serialization Format Constants and Helpers
# =============================================================================


def _invert_dict(d: dict[str, int]) -> dict[int, str]:
    """Invert a dictionary, swapping keys and values."""
    return {value: key for key, value in d.items()}


# Variable type codes.
# - cont: continuous (real-valued), int: integer, bin: binary (0/1).
_VARIABLE_TYPE_CODES: dict[str, int] = {
    "cont": 0,
    "int": 1,
    "bin": 2,
}

# Expression type codes.
# - min_objective: minimization, max_objective: maximization, constraint: constraint.
_EXPRESSION_TYPE_CODES: dict[str, int] = {
    "min_objective": 10,
    "max_objective": 11,
    "constraint": 12,
}

# First-order operators: arithmetic and mathematical functions.
_FIRST_ORDER_OPERATOR_CODES: dict[str, int] = {
    "+": 20,
    "-": 21,
    "*": 22,
    "/": 23,
    "^": 24,
    "abs": 25,
    "exp": 26,
    "log": 27,
}

# Operators that use infix notation (e.g., "x + y" instead of "+(x, y)").
_INFIX_OPERATORS: frozenset[str] = frozenset(["+", "-", "*", "/", "^"])

# First-order comparison operators: relational constraints.
_FIRST_ORDER_COMPARISON_CODES: dict[str, int] = {
    "=": 50,
    "!=": 51,
    "<=": 52,
    ">=": 53,
    "<": 54,
    ">": 55,
    "implies": 56,
}
# Combined set of all infix operators and comparisons for precedence handling.
_INFIX_OPERATORS_AND_COMPARISONS: frozenset[str] = (
    _INFIX_OPERATORS | frozenset(_FIRST_ORDER_COMPARISON_CODES.keys())
)

# Higher-order operators: aggregation and global constraints.
_HIGHER_ORDER_OPERATOR_CODES: dict[str, int] = {
    "sum": 70,
    "min": 71,
    "max": 72,
    "count": 73,
    "all_different": 80,
    "special_ordered_set_type_1": 81,
    "special_ordered_set_type_2": 82,
}

# Reverse mapping from numeric codes back to operator/type names.
_CODE_MAP: dict[int, str] = _invert_dict(
    _VARIABLE_TYPE_CODES
    | _EXPRESSION_TYPE_CODES
    | _FIRST_ORDER_OPERATOR_CODES
    | _FIRST_ORDER_COMPARISON_CODES
    | _HIGHER_ORDER_OPERATOR_CODES
)


# =============================================================================
# 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"]),
}


# =============================================================================
# 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 solver model.

        Args:
            model: The RelationalAI model to attach solver constructs to.
            num_type: Numeric data type for variable values. Must be Float or Integer.

        Raises:
            ValueError: If num_type is not Float or 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
        # 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}")

    # =========================================================================
    # 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.

        Args:
            expr: Relationship, Chain, or Expression defining variable structure.
            where: Grounding conditions to filter variable instances.
            name: Variable name or components to construct a name.
            type: 'cont', 'int', or 'bin'. Defaults based on num_type.
            lower: Lower bound.
            upper: Upper bound.
            start: Starting value hint for solver.
        """
        # Extract relationship and args from expression.
        relationship, args = self._extract_relationship(expr)
        key = _get_key(relationship)

        if key in self._variable_relationships:
            raise ValueError(f"Variables already defined for {relationship}.")

        # 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.")
            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__}.")
            if len(expr._args) != len(rel._fields):
                raise ValueError(f"Arg count {len(expr._args)} != field count {len(rel._fields)}.")
            return rel, expr._args
        raise TypeError(f"Invalid expr type: {type(expr).__name__}. Expected Relationship, Chain, or Expression.")

    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 str, got {type(var_type).__name__}.")
        if var_type not in _VARIABLE_TYPE_CODES:
            raise ValueError(f"Invalid type '{var_type}'. Valid: {', '.join(_VARIABLE_TYPE_CODES)}.")
        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"'{name}' must be numeric, got {type(val).__name__}.")

    # =========================================================================
    # 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 var_df.empty:
            raise ValueError("No variables found in the model.")
        var_types = var_df["type"].astype(int)

        expr_df = ctx.select(Expr.type, Expr.name, Expr.root).to_df()
        expr_types = expr_df["type"].astype(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.
        var_df["type"] = var_types.map(lambda x: _CODE_MAP.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()

        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."""
        # 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.
        expr_df["expression"] = [
            _expr_strings_rec(root, variables, operators, args)
            for root in expr_df["root"]
        ]

    # =========================================================================
    # Solving and Result Handling
    # =========================================================================

    def __getattr__(self, name: str) -> Any:
        """Get computed model attributes (num_variables, num_constraints, etc.)."""
        if name not in _MODEL_INFO_ATTRS:
            raise AttributeError(f"'{type(self).__name__}' has no attribute '{name}'.")

        concept_attr, type_code = _MODEL_INFO_ATTRS[name]
        concept = getattr(self, concept_attr)
        if type_code is None:
            result_df = self._model.select(count(concept) | 0).to_df()
        else:
            result_df = self._model.select(
                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]


# =============================================================================
# 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 SymbolifyContext:
    """Rewrites DSL expressions into solver AST form using Hash nodes.

    State:
    - variable_map: Expression keys -> Hash-based variable refs
    - defines: Property definitions (operator codes, arguments)
    - wheres: Grounding conditions
    - subcontexts: Child contexts for nested scopes
    """

    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. Returns Symbolic, rewritten expr, or None."""
        if expr is None or isinstance(expr, (int, float, str, b.Literal)):
            return None

        # Dispatch to type-specific handlers.
        if isinstance(expr, (b.Concept, b.Ref)):
            return self._rewrite_variable_ref(expr)
        if isinstance(expr, b.Alias):
            return self.rewrite(expr._source_item)
        if isinstance(expr, b.Chain):
            return self._rewrite_chain(expr)
        if isinstance(expr, b.Relationship):
            return self._rewrite_relationship(expr)
        if isinstance(expr, b.Match):
            return self._rewrite_match(expr)
        if isinstance(expr, b.FilterBy):
            return self._rewrite_filter_by(expr)
        if isinstance(expr, b.Expression):
            return self._rewrite_expression(expr)
        if isinstance(expr, b.Aggregate):
            return self._rewrite_aggregate(expr)
        if isinstance(expr, b.Fragment):
            return self._rewrite_fragment(expr)

        raise NotImplementedError(f"Cannot rewrite {type(expr).__name__}.")

    # -------------------------------------------------------------------------
    # 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_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_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:
            raise NotImplementedError(f"Unsupported operator '{op_name}'.")

        # 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 _rewrite_aggregate(self, expr: b.Aggregate) -> Optional[Symbolic | Any]:
        """Rewrite aggregate (sum/min/max/count/etc). Creates subcontext for inner scope."""
        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 symbolic arg and other args (position varies by operator).
        args = list(expr._args)
        proj_args = list(expr._projection_args)

        if op_name == "count":
            sym_arg, args, proj_args = 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.")
            sym_arg, args, proj_args = None, args[0]._items, []
        elif op_name in ("special_ordered_set_type_1", "special_ordered_set_type_2"):
            sym_arg, args, proj_args = args[1], args[:1], proj_args
        else:
            sym_arg, args, proj_args = args[0], args[1:], proj_args

        # Rewrite 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)

        if rw_sym is None:
            orig_where = expr._where._where
            nothing_changed = (
                rw_proj == list(proj_args)
                and rw_args == list(args)
                and rw_where == list(orig_where)
            )
            if nothing_changed:
                return None

        # Non-symbolic: rebuild aggregate with rewritten components.
        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]
            agg = b.Aggregate(op, *new_args, distinct=expr._distinct)
            return agg.per(*rw_group).where(*rw_where, *subctx.wheres)

        # Symbolic: build AST node.
        if op_name not in _HIGHER_ORDER_OPERATOR_CODES:
            raise NotImplementedError(f"Unsupported aggregate '{op_name}'.")

        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":
            # SOS2 needs ordered args: (rank_value, variable) pairs.
            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:
            # Other aggregates use unordered args with a combined hash key.
            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


def _expr_strings_rec(
    node: Any,
    variables: dict[Any, str],
    operators: dict[Any, int],
    args: dict[Any, list[Any]],
) -> str:
    """Recursively build human-readable string for expression tree."""
    if node in variables:
        return variables[node]

    op_code = operators.get(node)
    if op_code is None:
        return str(node)

    op = _CODE_MAP[op_code]
    children = args.get(node, [])

    # Format children, parenthesizing nested infix operators for precedence.
    is_infix = op in _INFIX_OPERATORS
    parts: list[str] = []
    for child in children:
        child_str = _expr_strings_rec(child, variables, operators, args)
        if is_infix:
            child_op_code = operators.get(child)
            needs_parens = (
                child_op_code is not None
                and _CODE_MAP[child_op_code] in _INFIX_OPERATORS_AND_COMPARISONS
            )
            parts.append(f"({child_str})" if needs_parens else child_str)
        else:
            parts.append(child_str)

    if op in _HIGHER_ORDER_OPERATOR_CODES:
        parts.sort()

    if op in _INFIX_OPERATORS_AND_COMPARISONS:
        return f"{parts[0]} {op} {parts[1]}"
    return f"{op}({', '.join(parts)})"
