"""Expression nodes for the SRHarness symbolic engine."""

from __future__ import annotations

from copy import deepcopy
from dataclasses import dataclass, field
from collections.abc import Iterator
from typing import TYPE_CHECKING, Any, Mapping

if TYPE_CHECKING:
    from .optimize import FitResult


class Expression:
    """Base class of every symbolic expression node."""

    def __add__(self, other: Any) -> Expression:
        return Binary("+", self, as_expression(other))

    def __radd__(self, other: Any) -> Expression:
        return Binary("+", as_expression(other), self)

    def __sub__(self, other: Any) -> Expression:
        return Binary("-", self, as_expression(other))

    def __rsub__(self, other: Any) -> Expression:
        return Binary("-", as_expression(other), self)

    def __mul__(self, other: Any) -> Expression:
        return Binary("*", self, as_expression(other))

    def __rmul__(self, other: Any) -> Expression:
        return Binary("*", as_expression(other), self)

    def __truediv__(self, other: Any) -> Expression:
        return Binary("/", self, as_expression(other))

    def __rtruediv__(self, other: Any) -> Expression:
        return Binary("/", as_expression(other), self)

    def __pow__(self, other: Any) -> Expression:
        return Binary("**", self, as_expression(other))

    def __rpow__(self, other: Any) -> Expression:
        return Binary("**", as_expression(other), self)

    def __neg__(self) -> Expression:
        return Unary("-", self)

    def __getitem__(self, index: Index | tuple[Index, ...] | str) -> Expression:
        if isinstance(index, tuple):
            indices = tuple(as_index(item) for item in index)
        else:
            indices = (as_index(index),)
        return Indexed(self, indices)

    def evaluate(
        self,
        values: Mapping[str, Any] | None = None,
        *,
        parameters: Mapping[str, Any] | None = None,
        time: Any = None,
        delay_resolver: Any = None,
        num_nodes: int | None = None,
    ) -> Any:
        """Evaluate this expression with NumPy values.

        Args:
            values: Values keyed by symbol name, or a context exposing ``data``
                and optional ``num_nodes`` metadata.
            parameters: Fitted parameter values keyed by parameter name.
            time: Optional sample times.
            delay_resolver: Optional callback that resolves delayed values.
            num_nodes: Explicit node count for indexed expressions.

        Returns:
            The evaluated scalar or array.
        """
        from .context import resolve_context
        from .evaluation import evaluate

        values, _, num_nodes = resolve_context(values, num_nodes=num_nodes)

        return evaluate(
            self,
            values=values,
            parameters=parameters,
            time=time,
            delay_resolver=delay_resolver,
            num_nodes=num_nodes,
        )

    eval = evaluate

    @property
    def operands(self) -> tuple[Expression, ...]:
        """Child expressions, exposed as an immutable tuple.

        Returns:
            The direct child nodes in expression order.
        """
        from .tree import children

        return children(self)

    def iter_preorder(self) -> Iterator[Expression]:
        """Yield this node followed by its descendants.

        Yields:
            Expression nodes in preorder.
        """
        from .tree import iter_preorder

        yield from iter_preorder(self)

    def iter_postorder(self) -> Iterator[Expression]:
        """Yield descendants followed by this node.

        Yields:
            Expression nodes in postorder.
        """
        from .tree import iter_postorder

        yield from iter_postorder(self)

    def copy(self) -> Expression:
        """Return an independent copy.

        Returns:
            A deep copy of this expression.
        """
        return deepcopy(self)

    def replace(self, old: Expression, new: Expression) -> Expression:
        """Return a tree in which the exact *old* node is replaced by *new*.

        Args:
            old: Existing expression node to replace.
            new: Replacement expression node.

        Returns:
            A copied expression tree with matching nodes replaced.
        """
        from .tree import replace

        return replace(self, old, new)

    def to_str(
        self,
        *,
        latex: bool = False,
        number_format: str = "",
    ) -> str:
        """Render the expression as plain text or LaTeX.

        Args:
            latex: Whether to render LaTeX notation.
            number_format: Format specification for numeric literals.

        Returns:
            The rendered expression.
        """
        from .render import render

        return render(self, latex=latex, number_format=number_format)

    def to_tree(self, *, number_format: str = "") -> str:
        """Render a compact preorder tree for diagnostics.

        Args:
            number_format: Format specification for numeric literals.

        Returns:
            A multiline representation of the expression tree.
        """
        from .tree import children

        lines = [self.to_str(number_format=number_format)]

        def visit(node: Expression, prefix: str, connector: str) -> None:
            lines.append(f"{prefix}{connector}{node.to_str(number_format=number_format)}")
            values = children(node)
            continuation = "  " if connector == "└ " else "│ "
            for index, child in enumerate(values):
                last = index == len(values) - 1
                visit(child, prefix + continuation, "└ " if last else "├ ")

        values = children(self)
        for index, child in enumerate(values):
            visit(child, "", "└ " if index == len(values) - 1 else "├ ")
        return "\n".join(lines)

    def __len__(self) -> int:
        return sum(1 for _ in self.iter_preorder())

    def fit(
        self,
        values: Mapping[str, Any],
        target: Any = None,
        *,
        initial: Mapping[str, Any] | None = None,
        method: str = "BFGS",
        options: Mapping[str, Any] | None = None,
        num_nodes: int | None = None,
    ) -> FitResult:
        """Fit named and grouped parameters against a target array.

        Args:
            values: Values keyed by symbol name, or a context exposing ``data``,
                ``target``, and optional ``num_nodes`` metadata.
            target: Target name or target values. Omit when supplied by context.
            initial: Optional initial parameter values.
            method: Optimization method name.
            options: Optional optimizer settings.
            num_nodes: Explicit node count for indexed expressions.

        Returns:
            The fitted expression, parameter values, predictions, and loss.
        """
        from .context import resolve_context
        from .optimize import fit

        values, target, num_nodes = resolve_context(values, target, num_nodes)
        if target is None:
            raise ValueError("A target array or a context target is required for fitting.")

        return fit(
            self, values, target, initial=initial, method=method,
            options=options, num_nodes=num_nodes,
        )

    def fold_constants(self) -> Expression:
        """Return a copy with closed numerical subexpressions evaluated.

        Returns:
            A simplified expression with closed numeric branches folded.
        """
        from .analysis import fold_constants

        return fold_constants(self)

    def count_parameters(
        self,
        values: Mapping[str, Any] | None = None,
        *,
        parameters: Mapping[str, Any] | None = None,
    ) -> int:
        """Count independent fitted values represented by this expression.

        Args:
            values: Values keyed by symbol name.
            parameters: Fitted parameter values keyed by parameter name.

        Returns:
            The number of independent scalar parameter values.
        """
        from .analysis import count_parameters

        return count_parameters(self, values, parameters=parameters)

    def __str__(self) -> str:
        from .render import render

        return render(self)


@dataclass(frozen=True, slots=True)
class Number(Expression):
    """Fixed numeric literal."""
    value: int | float


@dataclass(frozen=True, slots=True)
class Symbol(Expression):
    """Named input symbol with an optional bound value."""
    name: str
    value: Any = field(default=None, compare=False, repr=False)


@dataclass(frozen=True, slots=True)
class Parameter(Expression):
    """Named scalar parameter optimized during fitting."""
    name: str
    value: float | None = field(default=None, compare=False)


@dataclass(frozen=True, slots=True)
class GroupedParameter(Expression):
    """Parameter with one fitted value per category."""
    by: Expression
    name: str | None = None
    value: Mapping[Any, float] | None = field(default=None, compare=False)
    default: float | None = field(default=None, compare=False)


@dataclass(frozen=True, slots=True)
class Index:
    """Symbolic relation index."""
    name: str

    def __str__(self) -> str:
        return self.name


@dataclass(frozen=True, slots=True)
class Unary(Expression):
    """Unary expression node."""
    operator: str
    operand: Expression


@dataclass(frozen=True, slots=True)
class Binary(Expression):
    """Binary expression node."""
    operator: str
    left: Expression
    right: Expression


@dataclass(frozen=True, slots=True)
class Function(Expression):
    """Named function-call expression node."""
    name: str
    arguments: tuple[Expression, ...]


@dataclass(frozen=True, slots=True)
class Indexed(Expression):
    """Expression annotated with symbolic indices."""
    base: Expression
    indices: tuple[Index, ...]


@dataclass(frozen=True, slots=True)
class Reduction(Expression):
    """Sum reduction with an optional relation binder."""
    indices: tuple[Index, ...]
    operand: Expression
    relation: Expression | None = None


@dataclass(frozen=True, slots=True)
class Gather(Expression):
    """Weighted collection of a structural expression on relation entries."""
    relation: Expression
    operand: Expression


@dataclass(frozen=True, slots=True)
class Aggregate(Expression):
    """Convenience aggregation node lowered to indexed syntax."""
    relation: Expression
    operand: Expression


@dataclass(frozen=True, slots=True)
class RelationLift(Expression):
    """Convenience source or target projection used inside an aggregation."""
    role: str
    operand: Expression
    relation: Expression | None = None


def as_expression(value: Any) -> Expression:
    """Convert a numeric literal or expression into an expression node.

    Args:
        value: Existing expression or numeric literal.

    Returns:
        The corresponding expression node.
    """
    if isinstance(value, Expression):
        return value
    if isinstance(value, bool) or not isinstance(value, (int, float)):
        raise TypeError(f"Expected an expression or number, got {type(value).__name__}.")
    return Number(value)


Variable = Symbol


def as_index(value: Index | str) -> Index:
    """Convert an index name into an index node.

    Args:
        value: Existing index or valid Python identifier.

    Returns:
        The corresponding symbolic index.
    """
    if isinstance(value, Index):
        return value
    if isinstance(value, str) and value.isidentifier():
        return Index(value)
    raise TypeError(f"Expected an index name, got {value!r}.")


def param(name: str, value: float | None = None) -> Parameter:
    """Create a named scalar parameter.

    Args:
        name: Parameter name shared by all matching occurrences.
        value: Optional initial or fixed value.

    Returns:
        A symbolic scalar parameter.
    """
    return Parameter(name, value)


def grouped_param(
    by: Expression,
    *,
    name: str | None = None,
    value: Mapping[Any, float] | None = None,
    default: float | None = None,
) -> GroupedParameter:
    """Create a parameter with one fitted value per category.

    Args:
        by: Symbol or expression containing category labels.
        name: Optional parameter-map name.
        value: Optional initial values keyed by category.
        default: Value used for categories absent from ``value``.

    Returns:
        A category-dependent parameter expression.
    """
    return GroupedParameter(by, name=name, value=value, default=default)


def function(name: str, *arguments: Any) -> Function:
    """Create a supported symbolic function call.

    Args:
        name: Function name recognized by the evaluator.
        *arguments: Function operands.

    Returns:
        A symbolic function node.
    """
    return Function(name, tuple(as_expression(argument) for argument in arguments))


def reduction(
    indices: Index | tuple[Index, ...],
    operand: Any,
    relation: Any = None,
) -> Reduction:
    """Create an indexed sum reduction.

    Args:
        indices: Symbolic indices.
        operand: Expression being reduced or transformed.
        relation: Relation expression that binds symbolic indices.

    Returns:
        A symbolic reduction node.
    """
    if not isinstance(indices, tuple):
        indices = (indices,)
    return Reduction(
        tuple(as_index(index) for index in indices),
        as_expression(operand),
        None if relation is None else as_expression(relation),
    )


def gather(relation: Any, operand: Any) -> Gather:
    """Collect structural values at the nonzero entries of a relation.

    Args:
        relation: Indexed relation or relation-valued expression.
        operand: Structural expression evaluated at relation coordinates.

    Returns:
        A relation-aligned symbolic field.
    """
    return Gather(as_expression(relation), as_expression(operand))


def aggr(relation: Any, operand: Any = None) -> Expression:
    """Build and lower target-wise graph aggregation syntax.

    Args:
        relation: Edge relation, or a product containing it when ``operand`` is omitted.
        operand: Message expression to aggregate by target node.

    Returns:
        The equivalent canonical indexed reduction.
    """
    from .desugar import desugar
    from .parser import _split_aggregation

    if operand is None:
        relation, operand = _split_aggregation(as_expression(relation))
    return desugar(Aggregate(as_expression(relation), as_expression(operand)))


def targ(*arguments: Any) -> RelationLift:
    """Project node values onto relation targets inside ``aggr``.

    Args:
        *arguments: Either ``value`` or ``relation, value``.

    Returns:
        A target projection used by aggregation desugaring.
    """
    if len(arguments) == 1:
        return RelationLift("target", as_expression(arguments[0]))
    if len(arguments) == 2:
        return RelationLift("target", as_expression(arguments[1]), as_expression(arguments[0]))
    raise TypeError("targ expects targ(value) or targ(relation, value).")


def sour(*arguments: Any) -> RelationLift:
    """Project node values onto relation sources inside ``aggr``.

    Args:
        *arguments: Either ``value`` or ``relation, value``.

    Returns:
        A source projection used by aggregation desugaring.
    """
    if len(arguments) == 1:
        return RelationLift("source", as_expression(arguments[0]))
    if len(arguments) == 2:
        return RelationLift("source", as_expression(arguments[1]), as_expression(arguments[0]))
    raise TypeError("sour expects sour(value) or sour(relation, value).")