qdot/perception/ood.py

"""
qdot/perception/ood.py
======================
MahalanobisOOD — out-of-distribution detector for the Inspection Agent.

Uses Mahalanobis distance computed on PCA-projected penultimate-layer
features of TinyCNN model 0 (the reference model in the ensemble).

When a real device measurement is flagged as OOD, it means the scan
topology is genuinely outside the training distribution — not a quality
issue (that's DQC's job), but a device-specific signature the model
hasn't seen before. This triggers the DisorderLearner in Phase 3.

Important asymmetry in Phase 1:
    Training distribution = CIM-generated data
    OOD population        = QFlow real experimental data

This means QFlow scans *will* produce elevated OOD scores at test time
because real devices have charge disorder the CIM wasn't trained with.
That's the correct behaviour — the Phase 3 DisorderLearner is designed
to resolve exactly this gap. The OOD flag is not a failure; it's a
diagnostic that triggers the right module.

Calibration:
    Threshold is set at the 95th percentile of Mahalanobis distances on
    a held-out validation set from the CIM training distribution.
    This means ~5% false-positive rate on in-distribution data, which
    keeps the DisorderLearner from firing on normal CIM variation.
"""

from __future__ import annotations

import pickle
from pathlib import Path
from typing import Optional, Tuple

import numpy as np
from sklearn.decomposition import PCA
from sklearn.preprocessing import StandardScaler

from qdot.core.types import OODResult
from uuid import UUID


class MahalanobisOOD:
    """
    OOD detector based on Mahalanobis distance in PCA feature space.

    The feature space is the 32-dimensional penultimate layer of TinyCNN.
    PCA reduces this to `n_components` dimensions where the covariance
    structure is better conditioned.

    Mahalanobis distance:
        d = sqrt( (x - μ)ᵀ Σ⁻¹ (x - μ) )
    where μ, Σ are the mean and covariance of the training distribution
    in PCA space.

    Usage:
        ood = MahalanobisOOD(n_components=16)

        # Fit on training features
        features = ensemble.extract_features_batch(X_train)   # (N, 32)
        ood.fit(features)

        # At test time
        feat = ensemble.extract_features(array)               # (32,)
        result = ood.score(measurement_id, feat)
        if result.flag:
            # trigger DisorderLearner
            ...
    """

    def __init__(
        self,
        n_components: int = 16,
        calibration_percentile: float = 95.0,
    ) -> None:
        """
        Args:
            n_components: PCA dimensionality. 16 retains > 90% variance
                          for typical TinyCNN penultimate features.
            calibration_percentile: Threshold = this percentile of
                          training distances. Default 95 → ~5% FPR.
        """
        self.n_components = n_components
        self.calibration_percentile = calibration_percentile

        # Fitted parameters (None until fit() is called)
        self._scaler: Optional[StandardScaler] = None
        self._pca: Optional[PCA] = None
        self._mu: Optional[np.ndarray] = None          # (n_components,)
        self._precision: Optional[np.ndarray] = None   # (n_components, n_components)
        self._threshold: Optional[float] = None
        self._fitted = False

    # -----------------------------------------------------------------------
    # Fitting
    # -----------------------------------------------------------------------

    def fit(self, features: np.ndarray) -> None:
        """
        Fit the OOD detector on penultimate-layer features from the
        training distribution (CIM-generated data).

        Args:
            features: float32/float64 array of shape (N, 32).
                      Extract with ensemble.extract_features_batch(X_train).
        """
        features = np.asarray(features, dtype=np.float64)
        if features.ndim != 2:
            raise ValueError(f"Expected 2D feature array, got shape {features.shape}")

        n_samples, n_feat = features.shape
        n_comp = min(self.n_components, n_feat, n_samples - 1)

        # Step 1: Standardise features (zero mean, unit variance per dim)
        self._scaler = StandardScaler()
        scaled = self._scaler.fit_transform(features)

        # Step 2: PCA projection
        self._pca = PCA(n_components=n_comp, random_state=42)
        projected = self._pca.fit_transform(scaled)   # (N, n_comp)

        # Step 3: Compute mean + precision matrix of projected features
        self._mu = projected.mean(axis=0)
        cov = np.cov(projected.T) + np.eye(n_comp) * 1e-6   # regularise
        self._precision = np.linalg.inv(cov)

        # Step 4: Calibrate threshold on training distances
        train_distances = self._compute_distances(projected)
        self._threshold = float(
            np.percentile(train_distances, self.calibration_percentile)
        )

        explained = self._pca.explained_variance_ratio_.sum() if n_comp > 1 else 1.0
        self._fitted = True

        print(
            f"OOD detector fitted: n={n_samples}, "
            f"n_components={n_comp} ({explained:.1%} variance), "
            f"threshold={self._threshold:.3f} "
            f"({self.calibration_percentile:.0f}th percentile)"
        )

    # -----------------------------------------------------------------------
    # Scoring
    # -----------------------------------------------------------------------

    def score(self, measurement_id: UUID, features: np.ndarray) -> OODResult:
        """
        Compute OOD score for a single sample.

        Args:
            measurement_id: UUID to attach to the OODResult.
            features: float32/float64 array of shape (32,) or (1, 32).

        Returns:
            OODResult with score, threshold, and flag.
        """
        if not self._fitted:
            raise RuntimeError(
                "OOD detector has not been fitted. Call fit() first."
            )

        feat = np.asarray(features, dtype=np.float64).flatten()
        if feat.ndim == 0 or feat.shape[0] == 0:
            raise ValueError("Empty feature vector.")

        # Project to PCA space
        projected = self._project(feat.reshape(1, -1))   # (1, n_comp)

        # Mahalanobis distance
        diff = projected[0] - self._mu
        dist = float(np.sqrt(diff @ self._precision @ diff))

        flag = dist > self._threshold

        return OODResult(
            measurement_id=measurement_id,
            score=dist,
            threshold=self._threshold,
            flag=flag,
        )

    def score_batch(
        self, features: np.ndarray
    ) -> Tuple[np.ndarray, np.ndarray]:
        """
        Compute OOD scores for a batch of samples.

        Args:
            features: (N, 32) feature matrix

        Returns:
            (scores, flags) — float64 (N,) and bool (N,)
        """
        if not self._fitted:
            raise RuntimeError("OOD detector not fitted.")

        features = np.asarray(features, dtype=np.float64)
        projected = self._project(features)   # (N, n_comp)
        distances = self._compute_distances(projected)
        flags = distances > self._threshold
        return distances, flags

    # -----------------------------------------------------------------------
    # Persistence
    # -----------------------------------------------------------------------

    def save(self, path: str) -> None:
        """Save fitted OOD detector to disk."""
        state = {
            "n_components": self.n_components,
            "calibration_percentile": self.calibration_percentile,
            "scaler": self._scaler,
            "pca": self._pca,
            "mu": self._mu,
            "precision": self._precision,
            "threshold": self._threshold,
            "fitted": self._fitted,
        }
        with open(path, "wb") as f:
            pickle.dump(state, f)

    @classmethod
    def load(cls, path: str) -> "MahalanobisOOD":
        """Load a previously fitted OOD detector."""
        with open(path, "rb") as f:
            state = pickle.load(f)
        obj = cls(
            n_components=state["n_components"],
            calibration_percentile=state["calibration_percentile"],
        )
        obj._scaler = state["scaler"]
        obj._pca = state["pca"]
        obj._mu = state["mu"]
        obj._precision = state["precision"]
        obj._threshold = state["threshold"]
        obj._fitted = state["fitted"]
        return obj

    # -----------------------------------------------------------------------
    # Internal
    # -----------------------------------------------------------------------

    def _project(self, features: np.ndarray) -> np.ndarray:
        """Scale then PCA-project features."""
        scaled = self._scaler.transform(features)
        return self._pca.transform(scaled)

    def _compute_distances(self, projected: np.ndarray) -> np.ndarray:
        """
        Vectorised Mahalanobis distance computation.
        projected: (N, n_components)
        Returns: (N,) distances
        """
        diff = projected - self._mu   # (N, n_comp)
        # d_i = sqrt( diff_i @ precision @ diff_i )
        # vectorised: right_term = diff @ precision, then sum
        right = diff @ self._precision   # (N, n_comp)
        distances_sq = (right * diff).sum(axis=1)   # (N,)
        distances_sq = np.maximum(distances_sq, 0.0)   # numerical safety
        return np.sqrt(distances_sq)


# ---------------------------------------------------------------------------
# Batch feature extraction helper (used during fitting)
# ---------------------------------------------------------------------------

def extract_features_batch(
    ensemble_or_model,
    X: np.ndarray,
    batch_size: int = 256,
    device: str = "cpu",
) -> np.ndarray:
    """
    Extract penultimate-layer features for a batch of preprocessed arrays.

    Args:
        ensemble_or_model: EnsembleCNN or TinyCNN with .extract_features()
        X: float32 array of shape (N, 1, 64, 64) — already preprocessed
        batch_size: batch size for inference

    Returns:
        float32 array of shape (N, 32)
    """
    import torch
    from torch.utils.data import DataLoader, TensorDataset

    dev = torch.device(device)

    # If EnsembleCNN, use model 0 as reference
    model = ensemble_or_model
    if hasattr(model, "models"):
        model = model.models[0]
    model.to(dev).eval()

    dataset = TensorDataset(torch.from_numpy(X).float())
    loader = DataLoader(dataset, batch_size=batch_size, shuffle=False)

    all_features = []
    with torch.no_grad():
        for (batch,) in loader:
            batch = batch.to(dev)
            feat = model.extract_features(batch)   # (B, 32)
            all_features.append(feat.cpu().numpy())

    return np.concatenate(all_features, axis=0).astype(np.float32)