# SPDX-License-Identifier: Apache-2.0
"""C06 goldens: golden tensor files, a tap registry, gates that are never loosened, JSON compare reports.

- **Golden files** are ``.npz`` (loaded with ``allow_pickle=False``) holding named tensors ("taps") and an optional
  JSON metadata record under the key ``__meta__`` (source file hashes, reference script, sample, date).
- **Taps** are named intermediates recorded by :class:`TapRegistry` from the CPU reference or from an *eager*
  device run (never inside a trace capture: a traced intermediate is freed at the end of the capture).
- **Gates** (PLAN.md section 0.3): thresholds are declared in the test (from the spec targets) and frozen in a JSON
  file next to the test at the first green run (``test_pcc_device.py`` -> ``test_pcc_device.gates.json``). From then
  on a declaration looser than the frozen threshold raises :class:`GateLoosenedError`; tightening is allowed and
  re-frozen. Changing a frozen gate means editing the JSON by hand and disclosing it (VERIFICATION_<date>.md).
- **Reports**: :class:`CompareReport` collects per-tap statistics and gate verdicts and writes JSON.

numpy only (device tensors are converted through :func:`.tensors.to_numpy`, which imports ttnn lazily).
"""
from __future__ import annotations

import contextlib
import datetime as _dt
import fnmatch
import json
import math
import os
import sys
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any, Dict, Iterator, List, Mapping, Optional, Sequence, Union

import numpy as np

from . import metrics as _m
from .weights import file_sha256

__all__ = [
    "save_goldens",
    "load_goldens",
    "Goldens",
    "TapRegistry",
    "Gate",
    "GateResult",
    "GateLoosenedError",
    "GateRegistry",
    "METRICS",
    "DIRECTIONS",
    "TapComparison",
    "CompareReport",
    "compare_taps",
]

PathLike = Union[str, os.PathLike]
META_KEY = "__meta__"


def _now() -> str:
    return _dt.datetime.now(_dt.timezone.utc).isoformat(timespec="seconds")


def _host(value: Any) -> np.ndarray:
    from .tensors import to_numpy

    return np.asarray(to_numpy(value))


def _atomic_write_text(path: Path, text: str) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_name(f".{path.name}.{os.getpid()}.tmp")
    tmp.write_text(text)
    os.replace(tmp, path)


# --------------------------------------------------------------------------------------------- golden files

def save_goldens(path: PathLike, tensors: Mapping[str, Any], meta: Optional[Mapping[str, Any]] = None, *,
                 compress: bool = False) -> Path:
    """Write ``tensors`` (numpy / torch / ttnn) and ``meta`` (JSON-able) to an ``.npz`` file."""
    path = Path(path)
    arrays = {str(k): _host(v) for k, v in tensors.items()}
    if META_KEY in arrays:
        raise ValueError(f"{META_KEY!r} is reserved for metadata")
    arrays[META_KEY] = np.array(json.dumps(dict(meta or {}), sort_keys=True, default=str))
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_name(f".{path.stem}.{os.getpid()}.tmp.npz")
    (np.savez_compressed if compress else np.savez)(tmp, **arrays)
    os.replace(tmp, path)
    return path


class Goldens(Mapping):
    """Read-only, lazily loaded golden tensors of one ``.npz`` file (``with load_goldens(p) as g: g["neck"]``)."""

    def __init__(self, path: PathLike):
        self.path = Path(path)
        self._npz = np.load(self.path, allow_pickle=False)
        self._cache: Dict[str, np.ndarray] = {}
        self.meta: Dict[str, Any] = (json.loads(str(self._npz[META_KEY])) if META_KEY in self._npz.files else {})
        self._sha256: Optional[str] = None

    @property
    def sha256(self) -> str:
        if self._sha256 is None:
            self._sha256 = file_sha256(self.path)
        return self._sha256

    def __getitem__(self, name: str) -> np.ndarray:
        if name == META_KEY or name not in self._npz.files:
            raise KeyError(f"{name!r} not in {self.path.name}; have {sorted(self)[:20]}")
        if name not in self._cache:
            self._cache[name] = self._npz[name]
        return self._cache[name]

    def __iter__(self) -> Iterator[str]:
        return (k for k in self._npz.files if k != META_KEY)

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

    def close(self) -> None:
        self._npz.close()

    def __enter__(self) -> "Goldens":
        return self

    def __exit__(self, *exc) -> None:
        self.close()


def load_goldens(path: PathLike) -> Goldens:
    return Goldens(path)


# ------------------------------------------------------------------------------------------------- taps

class TapRegistry:
    """Records named intermediates. ``tap(name, value)`` returns ``value`` unchanged, so it can be inlined::

        taps = TapRegistry(include=["backbone.*", "neck"])
        x = taps.tap("backbone.layer1", layer1(x))        # stored as numpy (host read for device tensors)
        with taps.scope("head"):
            y = taps.tap("cls", cls_head(x))                # recorded as "head.cls"
        taps.save("golden/sample0_taps.npz", meta={...})

    A disabled registry (``TapRegistry(enabled=False)``, or :data:`NULL_TAPS`) does nothing, so production code can
    keep its ``tap`` calls. Device tensors are read immediately (they may be deallocated later by the model);
    calling ``tap`` on a device tensor while a trace capture is active raises."""

    def __init__(self, enabled: bool = True, *, include: Sequence[str] = ("*",), exclude: Sequence[str] = ()):
        self.enabled = enabled
        self.include = tuple(include)
        self.exclude = tuple(exclude)
        self._prefix: List[str] = []
        self._values: Dict[str, np.ndarray] = {}

    def _full(self, name: str) -> str:
        return ".".join(self._prefix + [name])

    def wants(self, name: str) -> bool:
        """Whether ``name`` (full dotted name) would be recorded."""
        return (self.enabled and any(fnmatch.fnmatchcase(name, p) for p in self.include)
                and not any(fnmatch.fnmatchcase(name, p) for p in self.exclude))

    def tap(self, name: str, value: Any) -> Any:
        if not self.enabled:
            return value
        full = self._full(name)
        if not self.wants(full):
            return value
        ttnn = sys.modules.get("ttnn")
        if ttnn is not None and isinstance(value, getattr(ttnn, "Tensor", ())):
            if value.storage_type() == ttnn.StorageType.DEVICE and ttnn.is_trace_capture_active(value.device()):
                raise RuntimeError(f"tap {full!r} inside a trace capture: return it as a trace output instead")
        self._values[full] = _host(value)
        return value

    @contextlib.contextmanager
    def scope(self, name: str) -> Iterator["TapRegistry"]:
        self._prefix.append(name)
        try:
            yield self
        finally:
            self._prefix.pop()

    def names(self) -> List[str]:
        return list(self._values)

    def __contains__(self, name: str) -> bool:
        return name in self._values

    def __getitem__(self, name: str) -> np.ndarray:
        return self._values[name]

    def items(self):
        return self._values.items()

    def to_dict(self) -> Dict[str, np.ndarray]:
        return dict(self._values)

    def clear(self) -> None:
        self._values.clear()

    def save(self, path: PathLike, meta: Optional[Mapping[str, Any]] = None, *, compress: bool = False) -> Path:
        return save_goldens(path, self._values, meta, compress=compress)


NULL_TAPS = TapRegistry(enabled=False)

# ------------------------------------------------------------------------------------------------ gates

# metric name -> (function(test, ref, **kw) -> float, direction): "min" = higher is better (value >= threshold)
METRICS: Dict[str, Any] = {
    "pcc": (_m.pcc, "min"),
    "masked_pcc": (_m.masked_pcc, "min"),
    "argmax_agreement": (_m.argmax_agreement, "min"),
    "label_agreement": (_m.label_agreement, "min"),
    "mask_iou": (_m.mask_iou, "min"),
    "max_abs": (lambda t, r: _m.error_stats(t, r)["max_abs"], "max"),
    "mean_abs": (lambda t, r: _m.error_stats(t, r)["mean_abs"], "max"),
    "rel_l2": (lambda t, r: _m.error_stats(t, r)["rel_l2"], "max"),
}

# The direction of every metric name a gate may use without spelling it out (the METRICS above plus the scalars
# recorded with CompareReport.add_value). A bare-number gate on any other metric is refused rather than guessed:
# a guessed "min" on an error metric (ADE, max |error|) would pass exactly the results it is meant to fail.
DIRECTIONS: Dict[str, str] = {
    **{name: direction for name, (_, direction) in METRICS.items()},
    **dict.fromkeys(("valid_row_pcc", "agreement", "iou", "miou", "mean_iou", "recall", "precision", "topk_overlap",
                     "topk_set_overlap"), "min"),
    **dict.fromkeys(("max_abs_err", "ade", "fde", "center_err", "center_err_max", "center_err_mean", "score_err",
                     "score_err_max", "score_err_mean"), "max"),
}


@dataclass(frozen=True)
class Gate:
    """``direction="min"``: pass when ``value >= threshold`` (PCC, agreement, IoU, recall); ``"max"``: pass when
    ``value <= threshold`` (errors, ADE/FDE). NaN never passes."""

    threshold: float
    direction: str = "min"
    metric: str = "pcc"

    def __post_init__(self) -> None:
        if self.direction not in ("min", "max"):
            raise ValueError(f"gate direction {self.direction!r}: expected 'min' or 'max'")
        object.__setattr__(self, "threshold", float(self.threshold))

    @classmethod
    def of(cls, spec: Union["Gate", float, int, Sequence, Mapping], metric: Optional[str] = None) -> "Gate":
        """``0.999`` | ``(0.999, "min")`` | ``(0.05, "max", "max_abs")`` | ``{"threshold": ..., ...}`` | Gate.
        A bare number takes the direction of ``metric`` (default ``pcc``) from :data:`DIRECTIONS` and raises
        ``ValueError`` for a metric whose direction is unknown; a direction that contradicts a known metric's
        (e.g. a ``"max"`` PCC gate) raises too."""
        if isinstance(spec, Gate):
            gate = spec
        elif isinstance(spec, Mapping):
            gate = cls(**spec)
        elif isinstance(spec, (int, float)):
            name = metric or "pcc"
            if name not in DIRECTIONS:
                raise ValueError(f"gate {spec!r} on metric {name!r}: its direction is unknown; give "
                                 f"(threshold, 'min' | 'max', {name!r})")
            gate = cls(float(spec), DIRECTIONS[name], name)
        else:
            parts = list(spec)
            gate = cls(*parts) if len(parts) == 3 else cls(parts[0], parts[1], metric or "pcc")
        known = DIRECTIONS.get(gate.metric)
        if known is not None and gate.direction != known:
            better = "higher" if known == "min" else "lower"
            raise ValueError(f"gate {gate}: {gate.metric!r} is a {known!r} metric ({better} is better)")
        return gate

    def passes(self, value: float) -> bool:
        if value is None or (isinstance(value, float) and math.isnan(value)):
            return False
        return value >= self.threshold if self.direction == "min" else value <= self.threshold

    def looser_than(self, other: "Gate") -> bool:
        """True if this gate accepts something ``other`` rejects (same direction), or the direction differs."""
        if self.direction != other.direction or self.metric != other.metric:
            return True
        return self.threshold < other.threshold if self.direction == "min" else self.threshold > other.threshold


@dataclass
class GateResult:
    name: str
    value: Optional[float]
    gate: Gate
    passed: bool
    newly_frozen: bool = False

    def __str__(self) -> str:
        op = ">=" if self.gate.direction == "min" else "<="
        v = "nan" if self.value is None else f"{self.value:.6g}"
        return f"{'PASS' if self.passed else 'FAIL'} {self.name}: {self.gate.metric} {v} {op} {self.gate.threshold:g}"


class GateLoosenedError(AssertionError):
    """A declared gate is looser than the threshold frozen at the first green run."""


class GateRegistry:
    """Declared gates + the JSON record of frozen gates next to the test file.

    Usage in a device test::

        GATES = GateRegistry.for_test(__file__, {"backbone": 0.999, "head.reg": 0.99, "dets.recall": (0.95, "min")})

        @pytest.mark.parametrize("stage", GATES.names())
        def test_stage(stage, stage_outputs):
            GATES.require(stage, pcc(*stage_outputs[stage]))      # AssertionError on failure; frozen on first pass

    ``write=False`` (or ``TTAW_GATES_READONLY=1``) never writes the JSON (CI / verification runs)."""

    VERSION = 1

    def __init__(self, path: PathLike, gates: Optional[Mapping[str, Any]] = None, *, write: Optional[bool] = None):
        self.path = Path(path)
        if write is None:
            write = os.environ.get("TTAW_GATES_READONLY", "0").strip() in ("", "0", "false", "no")
        self.write = bool(write)
        self.declared: Dict[str, Gate] = {k: Gate.of(v) for k, v in (gates or {}).items()}
        self.frozen: Dict[str, Dict[str, Any]] = {}
        if self.path.is_file():
            data = json.loads(self.path.read_text())
            self.frozen = dict(data.get("gates", {}))
        self.results: List[GateResult] = []
        loosened = []
        for name, g in self.declared.items():
            if name in self.frozen and g.looser_than(self._frozen_gate(name)):
                f = self._frozen_gate(name)
                loosened.append(f"{name}: declared {g.direction} {g.threshold:g} ({g.metric}) vs frozen "
                                f"{f.direction} {f.threshold:g} ({f.metric})")
        if loosened:
            raise GateLoosenedError(f"gates in {self.path.name} may never be loosened once recorded: "
                                    + "; ".join(loosened))

    @classmethod
    def for_test(cls, test_file: PathLike, gates: Optional[Mapping[str, Any]] = None, **kwargs: Any) -> "GateRegistry":
        """Registry stored as ``<test file stem>.gates.json`` next to ``test_file`` (pass ``__file__``)."""
        p = Path(test_file)
        return cls(p.with_name(p.stem + ".gates.json"), gates, **kwargs)

    def _frozen_gate(self, name: str) -> Gate:
        rec = self.frozen[name]
        return Gate(rec["threshold"], rec["direction"], rec["metric"])

    def names(self) -> List[str]:
        return sorted(set(self.declared) | set(self.frozen))

    def gate(self, name: str) -> Gate:
        """The effective gate: the declared one (never looser than the frozen one), else the frozen one."""
        if name in self.declared:
            return self.declared[name]
        if name in self.frozen:
            return self._frozen_gate(name)
        raise KeyError(f"no gate {name!r} declared or frozen in {self.path.name}")

    def check(self, name: str, value: Optional[float], gate: Optional[Any] = None) -> GateResult:
        """Evaluate ``value``; on a pass, freeze the gate if it is new or tighter than the frozen one."""
        g = Gate.of(gate) if gate is not None else self.gate(name)
        if name in self.frozen and g.looser_than(self._frozen_gate(name)):
            raise GateLoosenedError(f"gate {name!r} {g} is looser than the frozen {self.frozen[name]}")
        v = None if value is None else float(value)
        passed = g.passes(v)
        newly = False
        if passed and (name not in self.frozen or self._frozen_gate(name) != g):
            self.frozen[name] = {**asdict(g), "value_at_freeze": v, "frozen_at": _now()}
            newly = True
            if self.write:
                self.save()
        result = GateResult(name, v, g, passed, newly)
        self.results.append(result)
        return result

    def require(self, name: str, value: Optional[float], gate: Optional[Any] = None) -> GateResult:
        """:meth:`check`, raising ``AssertionError`` when the gate fails."""
        result = self.check(name, value, gate)
        if not result.passed:
            raise AssertionError(str(result))
        return result

    def save(self) -> None:
        body = {"version": self.VERSION, "note": "Frozen at the first green run; never loosen (PLAN.md 0.3).",
                "gates": dict(sorted(self.frozen.items()))}
        _atomic_write_text(self.path, json.dumps(body, indent=1, sort_keys=False) + "\n")


# ----------------------------------------------------------------------------------------------- reports

@dataclass
class TapComparison:
    name: str
    metric: str
    value: Optional[float]
    gate: Optional[Dict[str, Any]]
    passed: Optional[bool]
    stats: Dict[str, Any] = field(default_factory=dict)
    note: str = ""


class CompareReport:
    """Per-tap comparison records -> JSON (``logs/<model>/compare_<tag>.json``) and a text table."""

    def __init__(self, title: str = "", meta: Optional[Mapping[str, Any]] = None,
                 gates: Optional[GateRegistry] = None):
        self.title = title
        self.meta = dict(meta or {})
        self.gates = gates
        self.entries: List[TapComparison] = []

    def _gate_for(self, name: str, gate: Any, metric: str) -> Optional[Gate]:
        if gate is not None:
            return Gate.of(gate, metric)
        if self.gates is not None and name in self.gates.names():
            return self.gates.gate(name)
        return None

    def add(self, name: str, test: Any, ref: Any, *, metric: str = "pcc", gate: Any = None,
            **metric_kwargs: Any) -> TapComparison:
        """Compare two tensors with ``metric`` (a :data:`METRICS` name); always records the error statistics."""
        t, r = _host(test), _host(ref)
        stats: Dict[str, Any] = {"shape_test": list(t.shape), "shape_ref": list(r.shape),
                                 "dtype_test": str(t.dtype), "dtype_ref": str(r.dtype)}
        note = ""
        if t.shape == r.shape:
            stats.update(_m.error_stats(t, r))
            fn = METRICS[metric][0]
            value: Optional[float] = float(fn(t, r, **metric_kwargs))
        else:
            value, note = None, "shape mismatch"
        return self._record(name, metric, value, gate, stats, note)

    def add_value(self, name: str, value: Optional[float], *, metric: str, gate: Any = None,
                  stats: Optional[Mapping[str, Any]] = None) -> TapComparison:
        """Record an already computed scalar (detection recall, ADE, argmax agreement ...)."""
        return self._record(name, metric, None if value is None else float(value), gate, dict(stats or {}), "")

    def _record(self, name: str, metric: str, value: Optional[float], gate: Any, stats: Dict[str, Any],
                note: str) -> TapComparison:
        g = self._gate_for(name, gate, metric)
        passed: Optional[bool] = None
        if g is not None:
            passed = (self.gates.check(name, value, g).passed if self.gates is not None else g.passes(value))
        entry = TapComparison(name, metric, value, None if g is None else asdict(g), passed, stats, note)
        self.entries.append(entry)
        return entry

    @property
    def passed(self) -> bool:
        """All gated entries passed (ungated entries are informational)."""
        return all(e.passed for e in self.entries if e.passed is not None)

    def to_dict(self) -> Dict[str, Any]:
        return {"title": self.title, "created": _now(), "meta": self.meta, "passed": self.passed,
                "entries": [asdict(e) for e in self.entries]}

    def save_json(self, path: PathLike) -> Path:
        path = Path(path)
        _atomic_write_text(path, json.dumps(self.to_dict(), indent=1, default=str) + "\n")
        return path

    def table(self) -> str:
        rows = [f"{self.title}".strip(), f"{'name':40s} {'metric':16s} {'value':>12s} {'gate':>14s}  verdict"]
        for e in self.entries:
            gate = ""
            if e.gate is not None:
                gate = (">=" if e.gate["direction"] == "min" else "<=") + f"{e.gate['threshold']:g}"
            val = "nan" if e.value is None else f"{e.value:.6g}"
            verdict = {True: "PASS", False: "FAIL", None: "-"}[e.passed]
            rows.append(f"{e.name:40s} {e.metric:16s} {val:>12s} {gate:>14s}  {verdict} {e.note}".rstrip())
        return "\n".join(r for r in rows if r)


def compare_taps(test: Mapping[str, Any], golden: Mapping[str, Any], *, gates: Optional[Any] = None,
                 metric: str = "pcc", names: Optional[Sequence[str]] = None, title: str = "",
                 meta: Optional[Mapping[str, Any]] = None) -> CompareReport:
    """Compare every tap present in both mappings (or ``names``). ``gates``: a :class:`GateRegistry` or a
    ``{name: gate spec}`` mapping. Taps missing on either side are recorded as failures when gated."""
    registry = gates if isinstance(gates, GateRegistry) else None
    plain = {} if registry is not None or gates is None else {k: Gate.of(v, metric) for k, v in gates.items()}
    report = CompareReport(title, meta, registry)
    for name in (names if names is not None else [n for n in golden if n in test]):
        if name not in test or name not in golden:
            report.add_value(name, None, metric=metric, gate=plain.get(name),
                             stats={"missing": "test" if name not in test else "golden"})
            continue
        report.add(name, test[name], golden[name], metric=metric, gate=plain.get(name))
    return report