"""``linear-algebra@v1`` π* — matrix-operation canonicalizer.

Domain: ``linear-algebra``. One π* covers the whole linear-algebra
surface via an ``op`` field in the input JSON. Matrix entries are
folded through SymPy's exact-rational arithmetic so
``[[1, 2], [3, 4]]`` and ``[[1.0, 2.0], [3.0, 4.0]]`` collapse.

Why this matters
----------------
Ticket #000030 §2.2 Phase 6. Linear algebra rounds out the math
substrate alongside the calculus quartet. Same ``[math]`` extras
gate; same `srepr`-based canonical-bytes discipline as the rest
of the SymPy substrate.

Input shape (JSON, UTF-8)
-------------------------
::

    {
      "op": "rref" | "det" | "eigenvalues" | "inverse",
      "matrix": [[<cell>, <cell>, ...], ...]
    }

- ``op`` selects the operation. Each operation has its own canonical
  output shape; see below.
- ``matrix`` is a row-major 2D list. Cells go through
  :func:`sympy.sympify` so ``"1/2"``, ``1`` and ``1.0`` all work
  and fold to the same rational (see ``arithmetic@v1`` discipline).

Canonical output by op
----------------------

``op = "rref"`` — Reduced Row Echelon Form. Output:

    rref;rows=<r>;cols=<c>:cell00|cell01|...||cell10|...

Cells are printed as ``num/den`` (rational lowest-terms) for
integer + rational entries; symbolic entries fall through
``sp.srepr``.

``op = "det"`` — determinant. Output:

    det:<num/den-or-srepr>

``op = "eigenvalues"`` — eigenvalues with multiplicity. Output:

    eigenvalues:<value-1>x<mult-1>|<value-2>x<mult-2>|...

Sorted by value's `srepr` for determinism (eigenvalues come back
as a dict from SymPy and would otherwise depend on insertion
order).

``op = "inverse"`` — matrix inverse, same shape as ``rref``:

    inverse;rows=<r>;cols=<c>:cell00|...

Equivalence classes preserved
-----------------------------
- Cell-formatting variations (``1`` ≡ ``1.0`` ≡ ``"1.0"``) collapse.
- Eigenvalue ordering is canonicalized (sorted lexically by
  ``srepr``); two equivalent eigen-decompositions canonicalize
  identically.

Equivalence classes kept distinct
---------------------------------
- Different ``op`` for same matrix → different canonical (the op
  is part of identity).
- Different matrix entries → different canonical.
- Different matrix shape (rows/cols) → different canonical.

Errors
------
- Singular matrix on ``inverse`` → :class:`PiStarError`.
- Non-square matrix on ``det`` / ``inverse`` / ``eigenvalues`` →
  :class:`PiStarError`.
- Unknown ``op`` → :class:`PiStarError`.

Optional dependency
-------------------
Same as ``algebra-symbolic@v1``: SymPy via the ``[math]`` extras.

Source: ticket #000030 §2.2 Phase 6 — landed 2026-05-09.
"""

from __future__ import annotations

import json
from dataclasses import dataclass

from arborist.pi_star.protocol import PiStarError
from arborist.pi_star.registry import register


try:
    import sympy as sp  # type: ignore[import-not-found]
except ImportError:  # pragma: no cover — optional extra
    sp = None  # type: ignore[assignment]


_VALID_OPS = ("rref", "det", "eigenvalues", "inverse")


_SYMPY_REQUIRED_MSG = (
    "linear-algebra@v1 requires sympy; install with "
    "`pip install 'arborist[math]'`"
)


def _format_cell(cell) -> str:
    """Render a SymPy entry as canonical text. Rationals get
    ``num/den``; symbolic values fall through ``srepr``."""
    if isinstance(cell, sp.Rational):
        return f"{cell.p}/{cell.q}"
    return sp.srepr(cell)


def _format_matrix(M) -> str:
    rows = []
    for i in range(M.rows):
        cells = [_format_cell(M[i, j]) for j in range(M.cols)]
        rows.append("|".join(cells))
    return "||".join(rows)


@dataclass
class LinearAlgebraV1:
    name: str = "linear-algebra"
    version: str = "v1"
    domain: str = "linear-algebra"

    def canonicalize(self, raw: bytes) -> bytes:
        if sp is None:
            raise PiStarError(_SYMPY_REQUIRED_MSG)
        if not isinstance(raw, (bytes, bytearray)):
            raise PiStarError(
                "linear-algebra@v1 expects bytes; got "
                f"{type(raw).__name__}"
            )
        try:
            text = raw.decode("utf-8", errors="surrogatepass")
        except UnicodeDecodeError as exc:  # pragma: no cover
            raise PiStarError(f"input not valid UTF-8: {exc}") from exc

        try:
            obj = json.loads(text)
        except json.JSONDecodeError as exc:
            raise PiStarError(
                f"linear-algebra@v1 input is not valid JSON: {exc}"
            ) from exc
        if not isinstance(obj, dict):
            raise PiStarError(
                "linear-algebra@v1 input must be a JSON object"
            )
        for required in ("op", "matrix"):
            if required not in obj:
                raise PiStarError(
                    f"linear-algebra@v1 missing required field "
                    f"{required!r}"
                )

        op = obj["op"]
        matrix_raw = obj["matrix"]
        if op not in _VALID_OPS:
            raise PiStarError(
                f"linear-algebra@v1 op must be one of {_VALID_OPS}; "
                f"got {op!r}"
            )
        if not isinstance(matrix_raw, list) or not matrix_raw:
            raise PiStarError(
                "linear-algebra@v1 matrix must be a non-empty 2D array"
            )

        # Sympify each cell. Floats are folded to exact Rational via
        # Decimal(str(...)) so 1.0 ≡ 1 ≡ "1.0" all collapse to the
        # same Rational(1, 1) — matching arithmetic@v1's discipline.
        # Without this, sympy keeps floats as Float (separate type),
        # and downstream det()/inverse() return Float-shaped bytes
        # that don't match the int-input bytes.
        from decimal import Decimal
        from fractions import Fraction
        sym_rows: list[list] = []
        n_cols = None
        for ri, row in enumerate(matrix_raw):
            if not isinstance(row, list):
                raise PiStarError(
                    f"matrix row[{ri}] must be a JSON array"
                )
            if n_cols is None:
                n_cols = len(row)
            elif len(row) != n_cols:
                raise PiStarError(
                    f"matrix row[{ri}] has {len(row)} cells; "
                    f"row[0] has {n_cols} (jagged matrix)"
                )
            sym_row = []
            for ci, c in enumerate(row):
                try:
                    if isinstance(c, bool):
                        raise PiStarError(
                            f"matrix[{ri}][{ci}] is bool; not a number"
                        )
                    if isinstance(c, float):
                        # Float → exact Fraction → SymPy Rational. Avoids
                        # sympy.Rational(Decimal) signature drift across
                        # versions and matches arithmetic@v1's discipline.
                        f = Fraction(Decimal(str(c)))
                        sym_row.append(sp.Rational(f.numerator, f.denominator))
                    elif isinstance(c, int):
                        sym_row.append(sp.Integer(c))
                    else:
                        sym_row.append(sp.sympify(c, rational=True))
                except (sp.SympifyError, SyntaxError, TypeError, ValueError) as exc:
                    raise PiStarError(
                        f"matrix[{ri}][{ci}] cannot sympify: {exc}"
                    ) from exc
            sym_rows.append(sym_row)

        try:
            M = sp.Matrix(sym_rows)
        except (TypeError, ValueError) as exc:
            raise PiStarError(
                f"linear-algebra@v1 cannot build Matrix: {exc}"
            ) from exc

        # Dispatch by op.
        if op == "rref":
            R, _pivots = M.rref()
            body = (
                f"rref;rows={R.rows};cols={R.cols}:"
                f"{_format_matrix(R)}"
            )
            return body.encode("utf-8")

        if op == "det":
            if M.rows != M.cols:
                raise PiStarError(
                    f"linear-algebra@v1 det requires square matrix; "
                    f"got {M.rows}x{M.cols}"
                )
            d = M.det()
            return f"det:{_format_cell(d)}".encode("utf-8")

        if op == "inverse":
            if M.rows != M.cols:
                raise PiStarError(
                    f"linear-algebra@v1 inverse requires square matrix; "
                    f"got {M.rows}x{M.cols}"
                )
            try:
                Inv = M.inv()
            except Exception as exc:
                # SymPy raises NonInvertibleMatrixError on singular
                # matrices; the import path drifts across versions
                # (sympy.matrices.common in 1.13, sympy.matrices.exceptions
                # in 1.14). Catch broadly + re-raise as PiStarError so
                # we don't pin the import path.
                raise PiStarError(
                    f"linear-algebra@v1 inverse failed: {exc}"
                ) from exc
            body = (
                f"inverse;rows={Inv.rows};cols={Inv.cols}:"
                f"{_format_matrix(Inv)}"
            )
            return body.encode("utf-8")

        if op == "eigenvalues":
            if M.rows != M.cols:
                raise PiStarError(
                    f"linear-algebra@v1 eigenvalues requires square matrix; "
                    f"got {M.rows}x{M.cols}"
                )
            evs = M.eigenvals()  # dict {value: multiplicity}
            # Sort by srepr of the value for determinism.
            ordered = sorted(evs.items(), key=lambda kv: sp.srepr(kv[0]))
            parts = [f"{_format_cell(v)}x{m}" for v, m in ordered]
            return f"eigenvalues:{'|'.join(parts)}".encode("utf-8")

        raise PiStarError(  # pragma: no cover — _VALID_OPS exhausts
            f"linear-algebra@v1 unknown op {op!r}"
        )


if sp is not None:
    register(LinearAlgebraV1())
