inference.py

"""
Inference engines for SURPRISE.
"""

import logging
from abc import ABC, abstractmethod
from typing import Dict, List, Tuple

import numpy as np

from config import settings

log = logging.getLogger("surprise.inference")


class InferenceEngine(ABC):
    @abstractmethod
    def analyze(self, frames: np.ndarray) -> Dict:
        ...

    def _classify(self, surprise: float, straightness: float) -> Tuple[str, str, float]:
        s_low = settings.surprise_threshold_real
        s_high = settings.surprise_threshold_fake
        st_high = settings.straightness_threshold_real
        st_low = settings.straightness_threshold_fake

        if surprise < s_low and straightness > st_high:
            conf = 80 + (1 - surprise) * 15
            return "real", "AUTHENTIC", min(conf, 96)
        if surprise > s_high or straightness < st_low:
            conf = 60 + min(surprise * 40, 35)
            return "fake", "AI-GENERATED", min(conf, 95)
        conf = 50 + abs(surprise - 0.4) * 30
        return "uncertain", "INCONCLUSIVE", min(conf, 75)


class MockLeWMInference(InferenceEngine):
    def __init__(self):
        log.info("MockLeWMInference initialized (no real model)")

    def analyze(self, frames: np.ndarray) -> Dict:
        N = len(frames)
        diffs = np.abs(frames[1:].astype(np.float32) - frames[:-1].astype(np.float32))
        motion = diffs.mean(axis=(1, 2, 3)) / 255.0
        variance = float(frames.std(axis=(1, 2, 3)).mean()) / 255.0

        seed = int((motion.sum() + variance * 100) * 1000) % (2**32)
        rng = np.random.default_rng(seed)

        base = 0.20 + variance * 0.3
        per_frame = base + rng.normal(0, 0.04, size=N)

        if len(motion) > 0:
            m_mean = motion.mean()
            m_std = motion.std() + 1e-6
            for i in range(1, N):
                if motion[i - 1] > m_mean + 1.5 * m_std:
                    per_frame[i] += 0.25 + rng.uniform(0, 0.15)

        per_frame = np.clip(per_frame, 0.05, 1.0)
        surprise = float(per_frame.mean())
        straightness = float(np.clip(0.95 - per_frame.std() * 1.5, 0.5, 0.99))

        threshold = float(np.quantile(per_frame, 0.95))
        flagged = [int(i) for i in np.where(per_frame >= threshold)[0].tolist()][:8]

        verdict, label, conf = self._classify(surprise, straightness)
        flag_reason = self._reason(verdict, flagged)

        return {
            "verdict": verdict,
            "label": label,
            "confidence": round(float(conf), 1),
            "surprise_score": round(surprise, 3),
            "straightness": round(straightness, 3),
            "per_frame_surprise": [round(float(v), 3) for v in per_frame],
            "flagged_frames": flagged,
            "flag_reason": flag_reason,
            "backend": "mock",
        }

    @staticmethod
    def _reason(verdict, flagged):
        if verdict == "real":
            return "Smooth temporal dynamics; no physical violations detected"
        if verdict == "fake":
            if flagged:
                shown = flagged[:3]
                return f"Surprise spike at frames {', '.join(str(f) for f in shown)}"
            return "Inconsistent latent dynamics"
        return "Mixed signals"


class RealLeWMInference(InferenceEngine):
    def __init__(self):
        import torch
        from models import TinyEncoder, TinyPredictor

        self.device = self._pick_device()
        log.info(f"RealLeWMInference loading from {settings.checkpoint_path} on {self.device}")

        ckpt = torch.load(settings.checkpoint_path, map_location=self.device, weights_only=False)
        cfg = ckpt.get("config", {})
        self.embed_dim = cfg.get("embed_dim", 64)
        self.image_size = cfg.get("image_size", 64)

        self.encoder = TinyEncoder(embed_dim=self.embed_dim).to(self.device)
        self.predictor = TinyPredictor(embed_dim=self.embed_dim).to(self.device)
        self.encoder.load_state_dict(ckpt["encoder_state"])
        self.predictor.load_state_dict(ckpt["predictor_state"])
        self.encoder.eval()
        self.predictor.eval()

        log.info(f"Loaded checkpoint: embed_dim={self.embed_dim}, image_size={self.image_size}")

    def _pick_device(self):
        import torch
        if settings.device == "auto":
            return "cuda" if torch.cuda.is_available() else "cpu"
        return settings.device

    def analyze(self, frames: np.ndarray) -> Dict:
        import torch

        resized = self._resize_batch(frames, self.image_size)
        x = torch.from_numpy(resized).float() / 255.0
        x = x.permute(0, 3, 1, 2).contiguous().to(self.device)

        with torch.no_grad():
            z = self.encoder(x)
            z_pred = self.predictor(z[:-1])
            per_frame = ((z_pred - z[1:]) ** 2).mean(dim=-1).cpu().numpy()
            v = (z[1:] - z[:-1])
            v_norm = v / (v.norm(dim=-1, keepdim=True) + 1e-8)
            if v_norm.shape[0] > 1:
                cos_sim = (v_norm[:-1] * v_norm[1:]).sum(dim=-1)
                straightness = float(cos_sim.mean().cpu().item())
                straightness = (straightness + 1.0) / 2.0
            else:
                straightness = 0.5

        surprise = float(per_frame.mean())

        if len(per_frame) > 0:
            threshold = float(np.quantile(per_frame, 0.95))
            flagged = [int(i) for i in np.where(per_frame >= threshold)[0].tolist()][:8]
        else:
            flagged = []

        rescaled_surprise = min(surprise / 0.10, 1.0)
        verdict, label, conf = self._classify(rescaled_surprise, straightness)
        flag_reason = self._reason(verdict, flagged)

        return {
            "verdict": verdict,
            "label": label,
            "confidence": round(float(conf), 1),
            "surprise_score": round(float(rescaled_surprise), 3),
            "straightness": round(straightness, 3),
            "per_frame_surprise": [round(float(v), 4) for v in per_frame],
            "flagged_frames": flagged,
            "flag_reason": flag_reason,
            "raw_surprise": round(surprise, 5),
            "backend": "lewm",
        }

    @staticmethod
    def _resize_batch(frames, size):
        import cv2
        out = np.zeros((len(frames), size, size, 3), dtype=np.uint8)
        for i, f in enumerate(frames):
            out[i] = cv2.resize(f, (size, size), interpolation=cv2.INTER_AREA)
        return out

    @staticmethod
    def _reason(verdict, flagged):
        if verdict == "real":
            return "Smooth temporal dynamics; no anomalies detected"
        if verdict == "fake":
            if flagged:
                shown = flagged[:3]
                return f"Surprise spike at frames {', '.join(str(f) for f in shown)}"
            return "Inconsistent temporal dynamics"
        return "Mixed signals"


def get_inference_engine() -> InferenceEngine:
    backend = settings.model_backend
    if backend == "mock":
        return MockLeWMInference()
    if backend == "lewm":
        return RealLeWMInference()
    raise ValueError(f"Unknown MODEL_BACKEND: {backend!r}")