"""``function-sampled@v1`` π* — symbolic function sampler.

Domain: ``function-sampled``. Bridges symbolic SymPy expressions
to the existing ``time-series-quantized@v1`` integer-vector
canonical form. Two functions that render identically (within
sample-grid tolerance) collapse to the same canonical bytes.

Why this matters
----------------
Ticket #000030 §2.2 Phase 7. This is what plotting CAN become in
π* terms: a function is identified by its quantized samples on a
declared grid. Different libraries / DPIs / palettes don't matter
— the canonical bytes pin the value-at-each-x. The PNG render is
a downstream view of the same canonical evidence.

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

    {
      "f": "<expression-text>",
      "x": "<variable-name>",
      "x_min": <number>,
      "x_max": <number>,
      "n_samples": <positive int>,
      "dv": <positive number>
    }

- ``f`` — expression to sample (SymPy parsable).
- ``x`` — independent-variable name.
- ``x_min``, ``x_max`` — sample-grid endpoints (inclusive).
- ``n_samples`` — number of points on the grid (linspace).
- ``dv`` — value-quantization step. Same role as
  ``time-series-quantized@v1`` ``dv``: each y-sample is rounded to
  ``round(y / dv)`` (banker's rounding).

Output canonical bytes
----------------------
Reuses ``time-series-quantized@v1`` exactly::

    dt=<dt>;dv=<dv>;n=<n>;t0=<t0>:v0|v1|v2|...

Where ``dt = (x_max - x_min) / (n_samples - 1)``, ``t0 = 0`` (the
linspace index), and each ``v_i = round(f(x_i) / dv)``. This
makes the sampled signature byte-compatible with
``time-series-quantized@v1`` storage and inspection.

Equivalence classes preserved
-----------------------------
- Two expressions that evaluate to the same numeric values on the
  same grid → same bytes. ``sin(x)`` and ``2*sin(x)/2`` collapse
  trivially; the harder cases (``sin(x)**2 + cos(x)**2`` ≡ ``1``)
  don't go through SymPy `simplify` here, so they only collapse
  if the sampler observes the same numeric values within the dv
  tolerance — which they will.
- Different but equivalent JSON whitespace / formatting: same.

Equivalence classes kept distinct
---------------------------------
- Different grid (x_min, x_max, n_samples) → different canonical.
  The sample-grid IS part of identity.
- Different ``dv`` (quantization tightness) → different canonical.
- Different expression that diverges at any sample → different.

Round-trip property
-------------------
Projective. The canonical text isn't valid input JSON; rerunning
on the canonical bytes raises :class:`PiStarError`.

Optional dependency
-------------------
SymPy via the ``[math]`` extras (lambdify needs it). NumPy is
required transitively but ships standalone via SymPy's lambdify
backend; tests use ``pytest.importorskip("sympy")``.

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

from __future__ import annotations

import json
import math
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]


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


def _format_num(x) -> str:
    """Mirror time_series._num: int → 'N'; integer-valued float → 'N';
    otherwise repr (shortest round-trip)."""
    if isinstance(x, int) or (isinstance(x, float) and x.is_integer()):
        return str(int(x))
    return repr(x)


@dataclass
class FunctionSampledV1:
    name: str = "function-sampled"
    version: str = "v1"
    domain: str = "function-sampled"

    def canonicalize(self, raw: bytes) -> bytes:
        if sp is None:
            raise PiStarError(_SYMPY_REQUIRED_MSG)
        if not isinstance(raw, (bytes, bytearray)):
            raise PiStarError(
                "function-sampled@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"function-sampled@v1 input is not valid JSON: {exc}"
            ) from exc
        if not isinstance(obj, dict):
            raise PiStarError(
                "function-sampled@v1 input must be a JSON object"
            )
        for required in ("f", "x", "x_min", "x_max", "n_samples", "dv"):
            if required not in obj:
                raise PiStarError(
                    f"function-sampled@v1 missing required field "
                    f"{required!r}"
                )

        f_text = obj["f"]
        x_name = obj["x"]
        x_min = obj["x_min"]
        x_max = obj["x_max"]
        n_samples = obj["n_samples"]
        dv = obj["dv"]

        if not isinstance(f_text, str) or not f_text.strip():
            raise PiStarError("function-sampled@v1 f must be non-empty string")
        if not isinstance(x_name, str) or not x_name.strip():
            raise PiStarError("function-sampled@v1 x must be non-empty string")
        if not isinstance(x_min, (int, float)) or isinstance(x_min, bool):
            raise PiStarError("function-sampled@v1 x_min must be number")
        if not isinstance(x_max, (int, float)) or isinstance(x_max, bool):
            raise PiStarError("function-sampled@v1 x_max must be number")
        if x_max <= x_min:
            raise PiStarError(
                f"function-sampled@v1 x_max must be > x_min; got "
                f"x_min={x_min!r}, x_max={x_max!r}"
            )
        if isinstance(n_samples, bool) or not isinstance(n_samples, int) or n_samples < 2:
            raise PiStarError(
                "function-sampled@v1 n_samples must be int >= 2"
            )
        if isinstance(dv, bool) or not isinstance(dv, (int, float)) or dv <= 0:
            raise PiStarError(
                f"function-sampled@v1 dv must be positive number; got {dv!r}"
            )

        try:
            x_sym = sp.Symbol(x_name)
            f_expr = sp.sympify(f_text, locals={x_name: x_sym})
        except (sp.SympifyError, SyntaxError, TypeError) as exc:
            raise PiStarError(
                f"function-sampled@v1 cannot parse f: {exc}"
            ) from exc

        # Sample without numpy — keep dependency surface small. The
        # output is byte-identical to time-series-quantized@v1's
        # rounding (Python's round() is banker's-rounding, matching
        # numpy's default).
        try:
            f_callable = sp.lambdify(x_sym, f_expr, "math")
        except (TypeError, ValueError) as exc:
            raise PiStarError(
                f"function-sampled@v1 cannot lambdify f: {exc}"
            ) from exc

        dt_value = (x_max - x_min) / (n_samples - 1)
        values: list[int] = []
        for i in range(n_samples):
            xi = x_min + i * dt_value
            try:
                yi = f_callable(xi)
            except (ValueError, ZeroDivisionError, OverflowError) as exc:
                raise PiStarError(
                    f"function-sampled@v1 evaluation failed at x={xi!r}: "
                    f"{exc}"
                ) from exc
            if isinstance(yi, complex):
                if abs(yi.imag) > 1e-12:
                    raise PiStarError(
                        f"function-sampled@v1 got complex value at "
                        f"x={xi!r}: {yi!r}"
                    )
                yi = yi.real
            if not isinstance(yi, (int, float)):
                raise PiStarError(
                    f"function-sampled@v1 non-numeric sample at x={xi!r}: "
                    f"{yi!r} (type {type(yi).__name__})"
                )
            if math.isnan(yi) or math.isinf(yi):
                raise PiStarError(
                    f"function-sampled@v1 non-finite sample at x={xi!r}: "
                    f"{yi!r}"
                )
            values.append(int(round(yi / dv)))

        # Serialize in time-series-quantized@v1 format so the bytes
        # compose identically (function-sampled output IS a quantized
        # series, just one derived from a closed-form expression
        # rather than an array of samples).
        dt_s = _format_num(dt_value)
        dv_s = _format_num(dv)
        n = len(values)
        body = "|".join(str(v) for v in values)
        out = f"dt={dt_s};dv={dv_s};n={n};t0=0:{body}"
        return out.encode("utf-8")


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