surface_code_in_stem/decoders/ising_predecoder.py

"""Ising pre-decoder adapter for surface-code decoders.

The adapter preserves the shared decoder protocol while executing an optional
pre-decoding stage and then deferring to an existing deterministic fallback
decoder (MWPM). When pre-decoding is unavailable, unsupported, or fails, the
adapter falls back cleanly to MWPM and emits deterministic diagnostics.
"""

from __future__ import annotations

from dataclasses import dataclass, field
import hashlib
import os
from pathlib import Path
import time
from typing import Any, Callable, Mapping

import numpy as np

from .base import BoolArray, DecoderMetadata, DecoderOutput, DecoderProtocol
from .mwpm import MWPMDecoder


def _coerce_non_negative_int(value: Any) -> int | None:
    try:
        parsed = int(value)
    except Exception:
        return None
    if parsed < 0:
        return None
    return parsed


def _coerce_bool_array(detector_events: BoolArray) -> np.ndarray:
    events = np.asarray(detector_events, dtype=np.bool_)
    if events.ndim != 2:
        raise ValueError("detector_events must be a 2D bool array.")
    return events


def _extract_geometry(metadata: DecoderMetadata) -> tuple[int | None, int | None]:
    extra = metadata.extra or {}
    if not isinstance(extra, Mapping):
        return None, None

    distance = _coerce_non_negative_int(extra.get("distance", extra.get("d")))
    rounds = _coerce_non_negative_int(extra.get("rounds", extra.get("num_rounds")))
    if distance is None or rounds is None:
        distance = _coerce_non_negative_int(extra.get("surface_distance"))
        rounds = _coerce_non_negative_int(extra.get("surface_rounds"))
    return distance, rounds


def _expected_surface_detectors(distance: int | None, rounds: int | None) -> int | None:
    if distance is None or rounds is None:
        return None
    if distance < 2 or rounds < 1:
        return None
    return int(2 * rounds * ((distance * distance - 1) // 2))


def _coerce_predecode_output(raw: Any, num_detectors: int, shots: int) -> tuple[np.ndarray, np.ndarray]:
    if isinstance(raw, Mapping):
        if "residual" in raw:
            residual = np.asarray(raw["residual"], dtype=np.bool_)
        elif "syndrome" in raw:
            residual = np.asarray(raw["syndrome"], dtype=np.bool_)
        elif "residual_detectors" in raw:
            residual = np.asarray(raw["residual_detectors"], dtype=np.bool_)
        else:
            pre_l = np.asarray(raw.get("pre_l"), dtype=np.bool_)  # type: ignore[call-arg]
            residual = np.asarray(raw.get("residual"), dtype=np.bool_)
            if pre_l.shape == () or residual.shape == ():
                raise ValueError("predecode mapping-style output missing required fields.")
            pre_l = pre_l.astype(np.bool_).reshape(shots)
            residual = residual.astype(np.bool_)
            if residual.shape != (shots, num_detectors):
                residual = residual[:shots, :num_detectors]
            return pre_l, residual

        pre_l = np.asarray(raw.get("pre_l"), dtype=np.bool_) if "pre_l" in raw else np.zeros(shots, dtype=np.bool_)  # type: ignore[call-arg]
        residual = np.asarray(residual, dtype=np.bool_)
        if residual.ndim != 2 or residual.shape[0] != shots:
            raise ValueError(
                "predecode mapping-style output must be 2D with shape (shots, num_detectors)."
            )
        if residual.shape[1] != num_detectors:
            raise ValueError(
                f"predecode residual width mismatch: expected {num_detectors}, got {residual.shape[1]}"
            )
        if pre_l.ndim == 0:
            pre_l = np.full(shots, bool(pre_l), dtype=np.bool_)
        pre_l = np.asarray(pre_l, dtype=np.bool_).reshape(shots)
        return pre_l, residual

    if isinstance(raw, (tuple, list)):
        if len(raw) == 0:
            raise ValueError("predecode output is empty.")
        pre_l = np.asarray(raw[0], dtype=np.bool_)
        residual = np.asarray(raw[1], dtype=np.bool_) if len(raw) > 1 else np.asarray(raw[0], dtype=np.bool_)
        if pre_l.shape == (shots, num_detectors + 1):
            pre_l, residual = pre_l[:, 0], pre_l[:, 1:]
        if pre_l.ndim == 0:
            pre_l = np.full(shots, bool(pre_l), dtype=np.bool_)
        if residual.ndim != 2:
            residual = np.asarray(residual, dtype=np.bool_).reshape(shots, -1)
        pre_l = np.asarray(pre_l, dtype=np.bool_).reshape(shots)
        residual = np.asarray(residual, dtype=np.bool_)
    else:
        residual = np.asarray(raw, dtype=np.bool_)
        pre_l = np.zeros(shots, dtype=np.bool_)

    if residual.ndim != 2:
        raise ValueError("predecode residual output must be 2D.")
    if residual.shape[0] != shots:
        raise ValueError(
            f"predecode residual shot-axis mismatch: expected {shots}, got {residual.shape[0]}"
        )
    if residual.shape[1] == num_detectors + 1:
        residual = residual[:, :num_detectors]
    elif residual.shape[1] != num_detectors:
        raise ValueError(
            f"predecode residual width mismatch: expected {num_detectors}, got {residual.shape[1]}"
        )
    if pre_l.size == 1:
        pre_l = np.full(shots, bool(pre_l), dtype=np.bool_)
    return pre_l[:shots], residual[:, :num_detectors]


def _identity_predecode(events: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    shots = events.shape[0]
    return np.zeros(shots, dtype=np.bool_), np.asarray(events, dtype=np.bool_)


def _safe_fingerprint(path: Path) -> str:
    try:
        digest = hashlib.sha1()
        with path.open("rb") as handle:
            while True:
                chunk = handle.read(1024 * 1024)
                if not chunk:
                    break
                digest.update(chunk)
        return digest.hexdigest()
    except Exception:
        return ""


def _as_numpy(raw: Any) -> Any:
    if hasattr(raw, "detach"):
        try:
            return raw.detach().cpu().numpy()
        except Exception:
            pass
    if hasattr(raw, "numpy"):
        try:
            return raw.numpy()
        except Exception:
            pass
    return raw


def _predecoder_backend_is_disabled(value: Any) -> bool:
    if not isinstance(value, str):
        return False
    return value.strip().lower() in {"disabled", "none", "off", "false"}


def _candidate_predecode_inputs(events: np.ndarray, metadata: DecoderMetadata) -> list[np.ndarray]:
    distance, rounds = _extract_geometry(metadata)
    variants: list[np.ndarray] = [np.asarray(events, dtype=np.float32)]

    # Flattened detector stream as a simple sequence feature map.
    if events.size > 0:
        variants.append(np.asarray(events, dtype=np.float32).reshape(events.shape[0], -1, 1))

    if distance is not None and rounds is not None and distance >= 1 and rounds >= 1:
        half = (distance * distance - 1) // 2
        if half > 0 and events.shape[1] == 2 * rounds * half:
            xz = np.asarray(events, dtype=np.float32).reshape(events.shape[0], 2, rounds, half)
            variants.append(xz)

            # Optional 4D shape when checkpoints expect spatial layout.
            side = int(distance)
            if side > 0:
                flat = xz.reshape(events.shape[0], -1)
                max_dims = side * side * 2 * rounds
                resized = np.zeros((events.shape[0], max_dims), dtype=np.float32)
                if flat.shape[1] >= max_dims:
                    resized[:, :max_dims] = flat[:, :max_dims]
                else:
                    resized[:, : flat.shape[1]] = flat
                variants.append(resized.reshape(events.shape[0], 2, rounds, side, side))

    unique: list[np.ndarray] = []
    seen: set[tuple[int, ...]] = set()
    for candidate in variants:
        key = candidate.shape
        if key not in seen:
            seen.add(key)
            unique.append(candidate.astype(np.float32))
    return unique


def _invoke_candidate_signatures(
    predecoder: Callable[..., Any],
    events: np.ndarray,
    metadata: DecoderMetadata,
    num_detectors: int,
) -> Any:
    last_error: Exception | None = None
    candidates = _candidate_predecode_inputs(events, metadata)
    signature_args: list[tuple[Any, ...]] = []
    for candidate in candidates:
        signature_args.extend(
            [
                (candidate,),
                (candidate, metadata),
                (candidate, metadata.extra),
                (candidate, num_detectors),
            ]
        )
    if not signature_args:
        signature_args.append((np.asarray(events, dtype=np.float32), metadata))

    for args in signature_args:
        try:
            raw = predecoder(*args)  # type: ignore[misc]
            return _as_numpy(raw)
        except TypeError:
            continue
        except Exception as exc:
            last_error = exc
            continue

    if last_error is not None:
        raise last_error
    raise RuntimeError("predecoder callable rejected all candidate signatures")


@dataclass
class IsingDecoder(DecoderProtocol):
    """Decoder adapter that performs optional Ising pre-decoding then MWPM fallback."""

    name: str = "ising"
    fallback_decoder: MWPMDecoder = field(default_factory=MWPMDecoder)
    predecoder_backend: str = "identity"
    predecoder_callable: Callable[..., Any] | None = None
    predecoder_artifact: str | os.PathLike[str] | None = None
    predecoder_seed: int | None = None

    _backend_resolution_error: str | None = field(default=None, init=False, repr=False)
    _resolved_backend: str | None = field(default=None, init=False, repr=False)
    _predecoder: Callable[..., Any] | None = field(default=None, init=False, repr=False)

    def __post_init__(self) -> None:
        self._resolved_backend = None
        if self.predecoder_callable is not None and callable(self.predecoder_callable):
            self._predecoder = self.predecoder_callable
            self._resolved_backend = "callable"
            return

        if _predecoder_backend_is_disabled(self.predecoder_backend):
            self._resolved_backend = "disabled"
            return

        artifact_path = None
        if self.predecoder_artifact is not None:
            artifact_path = Path(self.predecoder_artifact)
            if not artifact_path.exists():
                self._backend_resolution_error = f"artifact does not exist: {artifact_path}"
                self._resolved_backend = "artifact-missing"
                return
            self._resolved_backend = f"artifact:{artifact_path.name}"
            return

        self._resolved_backend = "identity"

    def _apply_metadata_config(self, metadata: DecoderMetadata) -> None:
        extra = metadata.extra
        if not isinstance(extra, Mapping):
            return

        requested_backend = extra.get("predecoder_backend")
        requested_artifact = extra.get("predecoder_artifact")
        requested_seed = extra.get("predecoder_seed")

        changed = False
        if isinstance(requested_backend, str) and requested_backend != self.predecoder_backend:
            self.predecoder_backend = requested_backend
            changed = True
        if "predecoder_artifact" in extra:
            if requested_artifact is not None and requested_artifact != self.predecoder_artifact:
                self.predecoder_artifact = requested_artifact
                changed = True
            elif requested_artifact is None and self.predecoder_artifact is not None:
                self.predecoder_artifact = None
                changed = True
        if isinstance(requested_seed, int) and requested_seed != self.predecoder_seed:
            self.predecoder_seed = requested_seed
            changed = True

        if changed:
            self._backend_resolution_error = None
            self._resolved_backend = None
            self._predecoder = None
            self.__post_init__()

    def _load_artifact_predecoder(self, num_detectors: int, metadata: DecoderMetadata) -> Callable[[np.ndarray], Any] | None:
        artifact = self.predecoder_artifact
        if artifact is None:
            return None
        artifact_path = Path(artifact)
        if not artifact_path.exists():
            self._backend_resolution_error = f"artifact does not exist: {artifact_path}"
            return None

        requested_backend = (self.predecoder_backend or "auto").strip().lower()
        suffix = artifact_path.suffix.lower()
        if requested_backend in {"auto", "torch", "torchscript", "pytorch", "pt"} or suffix in {".pt", ".pth", ".ckpt"}:
            try:
                import torch
            except Exception:
                self._backend_resolution_error = "torch is required for torch-based predecoder backends"
                return None

            try:
                loaded = torch.jit.load(str(artifact_path), map_location="cpu")
                if callable(loaded):
                    return lambda events: loaded(torch.as_tensor(events))
            except Exception:
                pass

            try:
                loaded = torch.load(str(artifact_path), map_location="cpu")
            except Exception as exc:
                self._backend_resolution_error = f"Unable to load torch artifact '{artifact_path}': {exc}"
                return None

            if callable(loaded):
                return lambda events: loaded(events)
            if isinstance(loaded, Mapping):
                for key in ("predecoder", "model", "module", "decoder"):
                    candidate = loaded.get(key)
                    if callable(candidate):
                        return lambda events: candidate(events)

            self._backend_resolution_error = (
                f"torch artifact '{artifact_path}' does not expose a callable predecoder module."
            )
            return None

        if requested_backend in {"safetensor", "safetensors", "numpy", "np"} or suffix in {".npy", ".npz", ".safetensors"}:
            try:
                if suffix == ".npz":
                    loaded = np.load(str(artifact_path), allow_pickle=False)
                    if isinstance(loaded, np.lib.npyio.NpzFile):
                        first = None
                        for value in loaded.values():
                            candidate = np.asarray(value, dtype=np.float32)
                            if candidate.size > 0:
                                first = candidate
                                break
                        if first is None:
                            self._backend_resolution_error = f"npz artifact '{artifact_path}' is empty"
                            return None
                    else:
                        first = np.asarray(loaded, dtype=np.float32)
                elif suffix == ".safetensors":
                    try:
                        from safetensors.numpy import load as load_safetensor_array
                    except Exception as exc:
                        self._backend_resolution_error = (
                            f"Failed to import safetensors for '{artifact_path}': {exc}"
                        )
                        return None
                    tensors = load_safetensor_array(str(artifact_path))
                    first_values = list(tensors.values())
                    if not first_values:
                        self._backend_resolution_error = f"No tensors in safe tensor artifact '{artifact_path}'"
                        return None
                    first = np.asarray(first_values[0], dtype=np.float32)
                else:
                    first = np.asarray(np.load(str(artifact_path), allow_pickle=False), dtype=np.float32)
            except Exception as exc:
                self._backend_resolution_error = f"Unable to read artifact '{artifact_path}': {exc}"
                return None

            if first.ndim == 1:
                first = first.reshape(1, -1)
            if first.ndim != 2:
                self._backend_resolution_error = (
                    f"Unsupported artifact shape {first.shape}; expected matrix-like transform."
                )
                return None

            weight = first
            if weight.shape[1] != num_detectors:
                self._backend_resolution_error = (
                    f"Artifact width mismatch: expected {num_detectors}, got {weight.shape[1]}"
                )
                return None

            def _matrix_predecode(events: np.ndarray) -> np.ndarray:
                flat = np.asarray(events, dtype=np.float32).reshape(events.shape[0], -1)
                return flat @ weight.T

            return _matrix_predecode

        self._backend_resolution_error = f"Unsupported predecoder backend '{requested_backend}' for artifact '{artifact_path}'"
        return None

    def _run_predecoder(
        self,
        events: np.ndarray,
        metadata: DecoderMetadata,
        num_detectors: int,
    ) -> tuple[np.ndarray, np.ndarray, dict[str, Any], str | None]:
        self._apply_metadata_config(metadata)
        if self._resolved_backend in {None, "disabled", "artifact-missing"}:
            reason = self._backend_resolution_error or "predecoder disabled"
            pre_l, residual = _identity_predecode(events)
            details: dict[str, Any] = {
                "predecoder_backend": self._resolved_backend or "disabled",
                "predecoder_fallback_reason": reason,
            }
            return pre_l, residual, details, reason

        if self._predecoder is None:
            pre_l, residual = _identity_predecode(events)
            if self._resolved_backend == "identity":
                details = {"predecoder_backend": "identity", "predecoder_fallback_reason": None}
                return pre_l, residual, details, None

            if self.predecoder_artifact is None:
                details = {"predecoder_backend": self._resolved_backend, "predecoder_fallback_reason": None}
                return pre_l, residual, details, None

            self._predecoder = self._load_artifact_predecoder(num_detectors, metadata)
            if self._predecoder is None:
                details = {
                    "predecoder_backend": self._resolved_backend,
                    "predecoder_fallback_reason": self._backend_resolution_error,
                    "predecoder_artifact": str(self.predecoder_artifact),
                    "predecoder_artifact_available": True,
                }
                pre_l, residual = _identity_predecode(events)
                return pre_l, residual, details, self._backend_resolution_error

        predecoder_start = time.perf_counter_ns()
        try:
            raw = _invoke_candidate_signatures(
                self._predecoder,  # type: ignore[arg-type]
                events,
                metadata,
                num_detectors=num_detectors,
            )
            pre_l, residual = _coerce_predecode_output(raw, num_detectors=num_detectors, shots=events.shape[0])
            details = {
                "predecoder_backend": self._resolved_backend or self.predecoder_backend,
                "predecoder_fallback_reason": None,
                "predecoder_latency_ms": (time.perf_counter_ns() - predecoder_start) / 1_000_000,
            }
            return pre_l.astype(np.bool_), np.asarray(residual, dtype=np.bool_), details, None
        except Exception as exc:
            details = {
                "predecoder_backend": self._resolved_backend or self.predecoder_backend,
                "predecoder_fallback_reason": str(exc),
            }
            pre_l, residual = _identity_predecode(events)
            return pre_l, residual, details, str(exc)

    def _artifact_info(self) -> dict[str, Any]:
        if self.predecoder_artifact is None:
            return {}
        path = Path(self.predecoder_artifact)
        if not path.exists():
            return {
                "predecoder_artifact": str(path),
                "predecoder_artifact_exists": False,
            }
        return {
            "predecoder_artifact": str(path),
            "predecoder_artifact_exists": True,
            "predecoder_artifact_size": path.stat().st_size,
            "predecoder_artifact_fingerprint": _safe_fingerprint(path),
            "predecoder_artifact_backend": self._resolved_backend or "identity",
        }

    def decode(self, detector_events: BoolArray, metadata: DecoderMetadata) -> DecoderOutput:
        events = _coerce_bool_array(detector_events)
        shots, num_detectors = events.shape
        if metadata.num_observables <= 0:
            raise ValueError("metadata.num_observables must be > 0")

        start_ns = time.perf_counter_ns()
        diagnostics: dict[str, Any] = {
            "backend": self.name,
            "backend_id": self.name,
            "backend_enabled": True,
            "backend_available": True,
            "backend_contract": True,
            "backend_error": None,
            "backend_chain": [f"requested:{self.name}"],
            "fallback_chain": [f"requested:{self.name}"],
            "contract_flags": "backend_enabled,contract_met",
            "degraded": False,
            "num_shots": int(shots),
            "num_detectors": int(num_detectors),
            "num_observables": int(metadata.num_observables),
            "predecoder_seed": self.predecoder_seed,
        }
        diagnostics.update(self._artifact_info())
        diagnostics["predecoder_available"] = self._predecoder is not None or self._resolved_backend == "identity"

        distance, rounds = _extract_geometry(metadata)
        expected = _expected_surface_detectors(distance, rounds)
        if expected is not None and expected != num_detectors:
            reason = (
                f"Detector width mismatch: expected {expected} for distance={distance}, rounds={rounds}, "
                f"got {num_detectors}"
            )
            diagnostics.update(
                {
                    "backend_contract": False,
                    "predecoder_fallback_reason": reason,
                    "backend_error": reason,
                    "contract_flags": "backend_disabled,contract_fallback",
                    "degraded": True,
                    "backend_chain": diagnostics["backend_chain"] + ["predecoder_contract_failed"],
                    "fallback_chain": diagnostics["fallback_chain"] + ["predecoder_contract_failed"],
                    "predecoder_backend": self._resolved_backend or self.predecoder_backend,
                    "predecoder_available": False,
                    "predecoder_latency_ms": 0.0,
                }
            )
        pre_l, residual, predecoder_details, predecode_error = self._run_predecoder(
            events=events,
            metadata=metadata,
            num_detectors=num_detectors,
        )
        diagnostics["predecoder_backend"] = predecoder_details.get("predecoder_backend")
        diagnostics["predecoder_fallback_reason"] = predecoder_details.get("predecoder_fallback_reason")
        if "predecoder_latency_ms" in predecoder_details:
            diagnostics["predecoder_latency_ms"] = predecoder_details["predecoder_latency_ms"]

        if predecode_error is not None:
            diagnostics["backend_chain"].append("predecoder_fallback")
            diagnostics["fallback_chain"].append("predecoder_fallback")
            diagnostics["backend_error"] = predecode_error
            diagnostics["contract_flags"] = "backend_disabled,contract_fallback"
            diagnostics["degraded"] = True
            diagnostics["predecoder_available"] = False
        else:
            diagnostics["backend_chain"].append("selected:ising_predecoder")
            diagnostics["fallback_chain"].append("selected:ising_predecoder")
            diagnostics["predecoder_available"] = True

        if pre_l.shape[0] == shots:
            diagnostics["predecoder_logical_bit"] = int(np.sum(pre_l.astype(np.uint8)))

        if residual.shape != (shots, num_detectors):
            diagnostics.update(
                {
                    "backend_error": f"Residual shape malformed: {residual.shape}",
                    "backend_chain": diagnostics["backend_chain"] + ["predecoder_shape_fallback"],
                    "fallback_chain": diagnostics["fallback_chain"] + ["predecoder_shape_fallback"],
                    "contract_flags": "backend_disabled,contract_fallback",
                    "degraded": True,
                }
            )
            residual = events

        try:
            fallback_output = self.fallback_decoder.decode(residual, metadata)
            logicals = np.asarray(fallback_output.logical_predictions, dtype=np.bool_)
            diagnostics["fallback_decoder"] = fallback_output.decoder_name
            diagnostics["fallback_decoder_diagnostics"] = dict(fallback_output.diagnostics)
            diagnostics["backend_chain"].append(f"selected:{fallback_output.decoder_name}")
            diagnostics["fallback_chain"].append(f"selected:{fallback_output.decoder_name}")
            diagnostics["contract_flags"] = "backend_enabled,contract_met"
        except Exception as exc:
            diagnostics.update(
                {
                    "backend_error": str(exc),
                    "backend_chain": diagnostics["backend_chain"] + ["fallback_decoding_failed"],
                    "fallback_chain": diagnostics["fallback_chain"] + ["fallback_decoding_failed"],
                    "contract_flags": "backend_disabled,contract_fallback",
                    "degraded": True,
                }
            )
            logicals = np.zeros((shots, int(metadata.num_observables)), dtype=np.bool_)

        if logicals.shape != (shots, int(metadata.num_observables)):
            raise ValueError(
                f"Decoder returned predictions with invalid shape {logicals.shape}, expected ({shots}, {metadata.num_observables})"
            )

        latency_ms = (time.perf_counter_ns() - start_ns) / 1_000_000
        diagnostics["latency_ms"] = latency_ms
        diagnostics["sample_us"] = latency_ms * 1000.0
        diagnostics["degraded"] = bool(diagnostics.get("backend_error") is not None)

        if "predecoder_latency_ms" not in diagnostics:
            diagnostics["predecoder_latency_ms"] = 0.0

        return DecoderOutput(
            logical_predictions=logicals,
            decoder_name=self.name,
            diagnostics=diagnostics,
        )