# -*- coding: utf-8 -*-
"""
scenesolver_core.py — Scene Solver v3 : Coherent Reasoning System
=================================================================
Faithful extraction of the nine-stage SceneSolver pipeline from the
authoritative Colab notebook (NewChecker.ipynb), packaged for the
CrimeSceneSolver Hugging Face Space.

Stage 0 : AutoEncoder (AE)           : PRIMARY anomaly signal (batched)
Stage 1 : TimeSformer Binary         : Semantic gate (anomaly/normal)
Stage 2 : TimeSformer Multi-class    : Fine-grained anomaly classification
Stage 3 : YOLO Weapon (EXPERIMENTAL) : Hard safety trigger (conditional)
Stage 4 : Audio Feature Extraction   : Environmental signal
Stage 5 : Tracking / Behavior        : Motion behavior extraction
Stage 6 : Fusion Layer               : xAI-style reasoning core
Stage 7 : RL3033 (EXPERIMENTAL)      : Temporal reasoning agent (FINAL)
Stage 8 : LLaVA (EXPERIMENTAL)       : Context explanation (multi-frame)

LOGIC-NORMAL: When Stage 1 = Normal, status is overridden to SAFE,
Stages 2 and 3 are skipped, and LLaVA runs with neutral prompts.

Model assets (public repositories):
  - samxkevin/scenesolver-assets   (AE / TFB / TFM / RL / YOLO / audio)
  - samxkevin/SceneDescriptionModel (LLaVA VLM)

DEPLOYMENT DIFFERENCES vs the notebook (compatibility only — no pipeline
logic changes):
  1. Checkpoint downloads are lazy (resolved on first use, cached) instead
     of running inside the Config class body at import time.
  2. An ffmpeg binary shim (imageio-ffmpeg) is placed on PATH when the
     system ffmpeg is missing (HF Gradio Spaces have no apt access).
  3. Stage 8 (LLaVA, ~14 GB) is released after each analysis so its
     weights do not stay resident between runs on constrained hosts.
"""


# ── Standard library ──────────────────────────────────────────────────────────
import json
import os
import gc
import math
import time
import logging
import threading
import shutil
from datetime import datetime
from typing import List, Dict, Optional, Tuple, Any

# ── Third-party ───────────────────────────────────────────────────────────────
import cv2
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from PIL import Image, ImageDraw, ImageFont
from tqdm.auto import tqdm
from torchvision import transforms
from torchvision.models import resnet50, ResNet50_Weights
from reportlab.lib.pagesizes import letter
from reportlab.lib import colors
from reportlab.pdfgen import canvas as rl_canvas
from reportlab.platypus import Table, TableStyle

# ── Optional heavy imports ────────────────────────────────────────────────────
try:
    from transformers import (
        TimesformerForVideoClassification,
        LlavaForConditionalGeneration,
        AutoProcessor,
        VideoMAEImageProcessor,
    )
    TRANSFORMERS_AVAILABLE = True
except ImportError:
    TRANSFORMERS_AVAILABLE = False
    print("[WARNING] transformers not available — Stages 1/2/8 will be skipped.")

try:
    from stable_baselines3 import PPO
    RL_AVAILABLE = True
except ImportError:
    RL_AVAILABLE = False
    print("[WARNING] stable-baselines3 not available — Stage 7 (RL) will be skipped.")

try:
    import librosa
    LIBROSA_AVAILABLE = True
except ImportError:
    LIBROSA_AVAILABLE = False
    print("[WARNING] librosa not available — Stage 4 (Audio) will be limited.")

try:
    from ultralytics import YOLO
    YOLO_AVAILABLE = True
except ImportError:
    YOLO_AVAILABLE = False
    print("[INFO] ultralytics not available — Stage 3 (YOLO) will be disabled.")

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%H:%M:%S",
)
log = logging.getLogger("SceneSolver")


# ══════════════════════════════════════════════════════════════════════════════
# CONFIGURATION
# ══════════════════════════════════════════════════════════════════════════════

# ── Deployment fix: guarantee an ffmpeg binary on PATH ────────────────────────
# Gradio SDK Spaces cannot apt-install; imageio-ffmpeg ships a static binary.
def _ensure_ffmpeg_on_path() -> None:
    import tempfile as _tempfile
    if shutil.which("ffmpeg"):
        return
    try:
        import imageio_ffmpeg
        exe = imageio_ffmpeg.get_ffmpeg_exe()
        bindir = os.path.join(_tempfile.gettempdir(), "scenesolver_bin")
        os.makedirs(bindir, exist_ok=True)
        link = os.path.join(bindir, "ffmpeg")
        if not os.path.exists(link):
            try:
                os.symlink(exe, link)
            except OSError:
                shutil.copy2(exe, link)
        os.environ["PATH"] = bindir + os.pathsep + os.environ.get("PATH", "")
        log.info(f"[FFmpeg] Using bundled binary: {exe}")
    except Exception as e:
        log.warning(f"[FFmpeg] Bundled binary unavailable: {e}")


_ensure_ffmpeg_on_path()


from huggingface_hub import hf_hub_download, snapshot_download
from functools import lru_cache

HF_ASSET_REPO = "samxkevin/scenesolver-assets"
LLAVA_REPO    = "samxkevin/SceneDescriptionModel"

# Deployment fix: the notebook resolved every checkpoint with hf_hub_download /
# snapshot_download inside the class body at import time. On a Space that
# blocks startup on multi-GB downloads, so each asset is now resolved lazily
# on first access and cached (identical paths, identical repositories).
_asset_cache: Dict[str, str] = {}
_asset_lock = threading.Lock()


def _hf_asset(filename: str) -> str:
    with _asset_lock:
        if filename not in _asset_cache:
            log.info(f"[HF] Fetching asset: {HF_ASSET_REPO}/{filename}")
            _asset_cache[filename] = hf_hub_download(
                repo_id=HF_ASSET_REPO, filename=filename)
        return _asset_cache[filename]


def _hf_llava_snapshot() -> str:
    with _asset_lock:
        if "__LLAVA__" not in _asset_cache:
            log.info(f"[HF] Fetching VLM snapshot: {LLAVA_REPO} (~14 GB, first use only)")
            _asset_cache["__LLAVA__"] = snapshot_download(repo_id=LLAVA_REPO)
        return _asset_cache["__LLAVA__"]


class _LazyAssets(dict):
    """dict-compatible mapping whose values download on first access."""

    def __init__(self, resolvers: Dict[str, Any]):
        super().__init__()
        self._resolvers = resolvers
        self._resolved: Dict[str, str] = {}

    def __getitem__(self, key: str) -> str:
        if key not in self._resolved:
            self._resolved[key] = self._resolvers[key]()
        return self._resolved[key]

    def get(self, key: str, default: Any = None) -> Any:
        if key not in self._resolvers:
            return default
        try:
            return self[key]
        except Exception as e:
            log.warning(f"[HF] Asset '{key}' unavailable: {e}")
            return default

    def __contains__(self, key: object) -> bool:
        return key in self._resolvers


AUDIOPATHS = _LazyAssets({
    "CALLER_TUNE":      lambda: _hf_asset("audio/phone.mp3"),
    "XTTS_SPEAKER_WAV": lambda: _hf_asset("audio/DoctorVoice_22050.wav"),
})


class Config:
    # NOTE: `import spaces` (done first in app.py) keeps torch.cuda.is_available()
    # True on ZeroGPU via CUDA emulation, so DEVICE resolves to cuda there and
    # to cpu on CPU-only hosts — the notebook's device idiom is preserved.
    DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    PATHS = _LazyAssets({
        "AAE":       lambda: _hf_asset("AAE/best_model.pth"),
        "AE":        lambda: _hf_asset("AE/best_autoencoder.pth"),
        "RL_AGENT":  lambda: _hf_asset("RL/best_model.zip"),
        "TS_BINARY": lambda: _hf_asset("TFB/timesformer_best_binary.pth"),
        "TS_MULTI":  lambda: _hf_asset("TFM/best_model.pth"),
        "YOLO":      lambda: _hf_asset("YOLO/weapon_yolov8n.pt"),
        "LLAVA":     _hf_llava_snapshot,
    })
    DRIVE_ROOT = "/tmp/scenesolver"
    log.info("[HF] SceneSolver asset registry ready (lazy download on first use).")

    SAFE_THRESHOLD_FLOOR    = 0.030
    SAFE_THRESHOLD_LEGACY   = 0.055
    AUDIO_ANOMALY_THRESHOLD = 0.1

    CLASS_NAMES = {
        0: "Aggression",     1: "TheftOrLarceny", 2: "Firebombing",
        3: "LawEnforcement", 4: "Shooting",        5: "Vandalism",
        6: "RoadAccidents",
    }
    WEAPON_CLASS_NAMES = {0: "Melee", 1: "Firearms"}

    # RL3033 observation space dims (must match training)
    RL_AE_WINDOW   = 32
    RL_AE_CHANNELS = 6
    RL_LATENT_MOT  = 1536
    RL_LATENT_STK  = 2560
    RL_SPATIAL_DIM = 196
    RL_CONTEXT_DIM = 5

    AE_BATCH_SIZE   = 16
    AE_FRAME_STRIDE = 5

    EMA_SHORT_ALPHA = 0.3
    EMA_LONG_ALPHA  = 0.05

    FUSION_WEAPON_BOOST       = 0.40
    FUSION_MOTION_BOOST       = 0.15
    FUSION_AUDIO_BOOST        = 0.10
    FUSION_NO_SUPPORT_PENALTY = 0.15

    ALERT_RISK_HIGH   = 0.70
    ALERT_RISK_MEDIUM = 0.40

    AE_WINDOW_PAD_RIGHT = True


# ══════════════════════════════════════════════════════════════════════════════
# UTILITIES
# ══════════════════════════════════════════════════════════════════════════════
def cleanup_memory():
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


def compute_dynamic_threshold(scores: List[float],
                               floor: float = Config.SAFE_THRESHOLD_FLOOR) -> float:
    if not scores:
        return Config.SAFE_THRESHOLD_LEGACY
    arr = np.array(scores, dtype=np.float64)
    return max(float(np.mean(arr) + 2.0 * np.std(arr)), floor)


def extract_frames(video_path: str) -> Tuple[List[np.ndarray], float]:
    cap    = cv2.VideoCapture(video_path)
    frames = []
    fps    = cap.get(cv2.CAP_PROP_FPS) or 25.0
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
    cap.release()
    return frames, fps


def extract_audio(video_path: str) -> Optional[str]:
    output_path = "/tmp/audio.wav"
    os.system(
        f"ffmpeg -y -loglevel quiet -i '{video_path}' "
        f"-ac 1 -ar 22050 '{output_path}' 2>/dev/null"
    )
    return output_path if os.path.exists(output_path) else None


def frame_to_timestamp(frame_idx: int, fps: float) -> str:
    """Convert a frame index to MM:SS.mmm."""
    if fps <= 0:
        fps = 25.0
    total_seconds = frame_idx / fps
    minutes = int(total_seconds // 60)
    seconds = total_seconds % 60
    return f"{minutes:02d}:{seconds:06.3f}"


def calculate_enhanced_score(img: torch.Tensor, recon: torch.Tensor) -> float:
    mse = torch.mean((img - recon) ** 2).item()
    mae = torch.mean(torch.abs(img - recon)).item()
    return 0.6 * mse + 0.4 * mae


def compute_ema(values: np.ndarray, alpha: float) -> np.ndarray:
    ema    = np.zeros_like(values)
    ema[0] = values[0]
    for i in range(1, len(values)):
        ema[i] = alpha * values[i] + (1.0 - alpha) * ema[i - 1]
    return ema


def find_anomaly_extent(scores: List[float], threshold: float) -> Tuple[int, int]:
    above = [i for i, s in enumerate(scores) if s > threshold]
    if not above:
        peak = int(np.argmax(scores))
        half = max(1, len(scores) // 10)
        return max(0, peak - half), min(len(scores) - 1, peak + half)
    return above[0], above[-1]


def z_score_normalize(arr: np.ndarray) -> np.ndarray:
    mean = np.mean(arr)
    std  = np.std(arr)
    return (arr - mean) / max(std, 1e-8)


def check_audio_has_content(audio_path: str) -> bool:
    if not audio_path or not os.path.exists(audio_path) or not LIBROSA_AVAILABLE:
        return False
    try:
        y, _ = librosa.load(audio_path, sr=22050, duration=10)
        return float(np.sqrt(np.mean(y ** 2))) > 1e-4
    except Exception:
        return False


# ══════════════════════════════════════════════════════════════════════════════
# SEGMENT-WISE AE-WEIGHTED TEMPORAL SAMPLING
# ══════════════════════════════════════════════════════════════════════════════
def segmentwise_ae_sampling(
    frames: List[np.ndarray],
    ae_scores: List[float],
    sampled_indices: List[int],
    num_frames: int = 16,
) -> List[np.ndarray]:
    n_total = len(frames)
    blank   = np.zeros((224, 224, 3), dtype=np.uint8)

    if n_total == 0:
        return [blank] * num_frames
    if n_total < num_frames:
        padded = list(frames)
        while len(padded) < num_frames:
            padded.append(padded[-1])
        return padded
    if not ae_scores or not sampled_indices:
        step = n_total / num_frames
        return [frames[min(int(i * step), n_total - 1)] for i in range(num_frames)]

    if len(ae_scores) != len(sampled_indices):
        raise ValueError(
            f"segmentwise_ae_sampling: len(ae_scores)={len(ae_scores)} != "
            f"len(sampled_indices)={len(sampled_indices)}."
        )
    for pos, fi in enumerate(sampled_indices):
        if not (0 <= int(fi) < n_total):
            raise ValueError(
                f"segmentwise_ae_sampling: sampled_indices[{pos}]={fi} out of "
                f"range for frames of length {n_total}."
            )

    segment_size = n_total / num_frames
    buckets: List[List[Tuple[int, float]]] = [[] for _ in range(num_frames)]
    for raw_fi, raw_score in zip(sampled_indices, ae_scores):
        fi  = int(raw_fi)
        seg = min(int(fi / segment_size), num_frames - 1)
        buckets[seg].append((fi, float(raw_score)))

    selected: List[np.ndarray] = []
    for seg_idx in range(num_frames):
        seg_start = int(seg_idx * segment_size)
        seg_end   = min(int((seg_idx + 1) * segment_size), n_total)
        if seg_start >= seg_end:
            best_fi = min(seg_start, n_total - 1)
        elif buckets[seg_idx]:
            best_fi = max(buckets[seg_idx], key=lambda x: x[1])[0]
        else:
            best_fi = (seg_start + seg_end) // 2
        selected.append(frames[min(max(best_fi, 0), n_total - 1)])
    return selected


# ══════════════════════════════════════════════════════════════════════════════
# AUTOENCODER MODEL
# ══════════════════════════════════════════════════════════════════════════════
class ResNetAE(nn.Module):
    def __init__(self, latent_dim: int = 512):
        super().__init__()
        r            = resnet50(weights=ResNet50_Weights.IMAGENET1K_V2)
        self.encoder = nn.Sequential(*list(r.children())[:-2])
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc_enc  = nn.Linear(2048, latent_dim)
        self.fc_dec  = nn.Linear(latent_dim, 2048)
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(2048, 1024, 4, 1, 0), nn.BatchNorm2d(1024), nn.ReLU(),
            nn.ConvTranspose2d(1024, 512,  4, 2, 1), nn.BatchNorm2d(512),  nn.ReLU(),
            nn.ConvTranspose2d(512,  256,  4, 2, 1), nn.BatchNorm2d(256),  nn.ReLU(),
            nn.ConvTranspose2d(256,  128,  4, 2, 1), nn.BatchNorm2d(128),  nn.ReLU(),
            nn.ConvTranspose2d(128,  64,   4, 2, 1), nn.BatchNorm2d(64),   nn.ReLU(),
            nn.ConvTranspose2d(64,   3,    4, 2, 1),
        )
        self.upsample = nn.Upsample((224, 224), mode="bilinear", align_corners=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        f = self.encoder(x)
        l = self.fc_enc(self.avgpool(f).flatten(1))
        r = self.upsample(self.decoder(self.fc_dec(l).view(-1, 2048, 1, 1)))
        return torch.sigmoid(r)


# ══════════════════════════════════════════════════════════════════════════════
# STAGE 0 — AUTOENCODER
# ══════════════════════════════════════════════════════════════════════════════
class Stage0AE:
    def __init__(self, config: Config):
        self.config    = config
        self.model: Optional[ResNetAE] = None
        self.transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
        ])

    def load_model(self):
        log.info("[Stage0-AE] Loading autoencoder...")
        self.model = ResNetAE().to(self.config.DEVICE)
        path = self.config.PATHS["AE"]
        if not os.path.exists(path):
            log.warning(f"[Stage0-AE] Checkpoint not found: {path}")
            self.model.eval()
            return self
        ckpt = torch.load(path, map_location=self.config.DEVICE, weights_only=False)
        if isinstance(ckpt, dict):
            for k in ("model_state", "model_state_dict", "state_dict"):
                if k in ckpt:
                    state = ckpt[k]
                    break
            else:
                raise ValueError(f"No valid model state key found: {list(ckpt.keys())}")
        else:
            state = ckpt
        self.model.load_state_dict(state)
        self.model.eval()
        return self

    def analyze(self, video_path: str) -> dict:
        if self.model is None:
            self.load_model()
        log.info("[Stage0-AE] Extracting frames...")
        frames, fps = extract_frames(video_path)
        log.info(f"[Stage0-AE] {len(frames)} frames @ {fps:.1f} FPS")

        sampled_indices = list(range(0, len(frames), self.config.AE_FRAME_STRIDE))
        scores: List[float] = []
        bs = self.config.AE_BATCH_SIZE
        log.info(f"[Stage0-AE] Scoring {len(sampled_indices)} frames (batch={bs})...")

        for batch_start in range(0, len(sampled_indices), bs):
            batch_indices = sampled_indices[batch_start: batch_start + bs]
            tensors = [self.transform(Image.fromarray(frames[i])) for i in batch_indices]
            batch   = torch.stack(tensors).to(self.config.DEVICE)
            with torch.no_grad():
                recon = self.model(batch)
            for i in range(len(batch_indices)):
                scores.append(calculate_enhanced_score(batch[i:i+1], recon[i:i+1]))
            del batch, recon
            cleanup_memory()

        if not scores:
            return {
                "status": "SAFE", "max_error": 0.0, "avg_error": 0.0,
                "peak_frame_idx": 0, "peak_score_idx": 0,
                "frames": frames, "scores": scores, "fps": fps,
                "threshold": self.config.SAFE_THRESHOLD_LEGACY,
                "anomaly_start_frame": 0, "anomaly_end_frame": 0,
                "sampled_indices": sampled_indices,
            }

        threshold      = compute_dynamic_threshold(scores)
        max_error      = float(np.max(scores))
        avg_error      = float(np.mean(scores))
        peak_score_idx = int(np.argmax(scores))
        peak_frame_idx = sampled_indices[peak_score_idx]
        anom_start_si, anom_end_si = find_anomaly_extent(scores, threshold)
        is_anomaly = max_error > threshold or avg_error > threshold * 0.7
        status     = "CRITICAL" if is_anomaly else "SAFE"
        cleanup_memory()
        log.info(f"[Stage0-AE] {status} | max={max_error:.6f} | "
                 f"avg={avg_error:.6f} | threshold={threshold:.6f}")
        return {
            "status":              status,
            "max_error":           max_error,
            "avg_error":           avg_error,
            "peak_frame_idx":      peak_frame_idx,
            "peak_score_idx":      peak_score_idx,
            "frames":              frames,
            "scores":              scores,
            "sampled_indices":     sampled_indices,
            "fps":                 fps,
            "threshold":           threshold,
            "anomaly_start_frame": sampled_indices[anom_start_si],
            "anomaly_end_frame":   sampled_indices[anom_end_si],
        }

    def build_ae_signals(self, scores: List[float], peak_score_idx: int) -> np.ndarray:
        W   = self.config.RL_AE_WINDOW
        eps = 1e-8
        if not scores:
            return np.zeros((W, 6), dtype=np.float32)

        arr   = np.array(scores, dtype=np.float64)
        n     = len(arr)
        delta = np.zeros(n); accel = np.zeros(n)
        delta[1:]  = arr[1:] - arr[:-1]
        accel[2:]  = delta[2:] - delta[1:-1]
        ema_s = compute_ema(arr, self.config.EMA_SHORT_ALPHA)
        ema_l = compute_ema(arr, self.config.EMA_LONG_ALPHA)
        spike = arr / (ema_l + eps)

        full_signals = np.stack([arr, delta, accel, ema_s, ema_l, spike], axis=1)
        for c in range(6):
            full_signals[:, c] = z_score_normalize(full_signals[:, c])

        half_w    = W // 2
        start_idx = max(0, peak_score_idx - half_w)
        end_idx   = min(n, start_idx + W)
        window    = full_signals[start_idx:end_idx]
        if len(window) < W:
            pad    = np.zeros((W - len(window), 6), dtype=np.float32)
            window = np.concatenate([window, pad], axis=0)
        return window[:W].astype(np.float32)


# ══════════════════════════════════════════════════════════════════════════════
# STAGE 1 — TIMESFORMER BINARY
# ══════════════════════════════════════════════════════════════════════════════
class Stage1TSBinary:
    def __init__(self, config: Config):
        self.config = config
        self.model  = None
        self._last_logits: Optional[np.ndarray] = None
        self.transform = transforms.Compose([
            transforms.ToPILImage(),
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
        ])

    def load_model(self):
        if not TRANSFORMERS_AVAILABLE:
            log.warning("[Stage1-TSBinary] transformers unavailable.")
            return self
        log.info("[Stage1-TSBinary] Loading binary TimeSformer...")
        self.model = TimesformerForVideoClassification.from_pretrained(
            "facebook/timesformer-base-finetuned-k400",
            num_labels=2,
            ignore_mismatched_sizes=True,
        ).to(self.config.DEVICE)
        path = self.config.PATHS["TS_BINARY"]
        if os.path.exists(path):
            raw = torch.load(path, map_location=self.config.DEVICE, weights_only=False)
            if isinstance(raw, dict):
                state = raw.get("model") or raw.get("model_state_dict") or raw
            else:
                state = raw
            self.model.load_state_dict(state, strict=False)
            log.info("[Stage1-TSBinary] Weights loaded.")
        else:
            log.warning(f"[Stage1-TSBinary] Checkpoint not found: {path}")
        self.model.eval()
        return self

    def analyze(self, frames: List[np.ndarray], peak_idx: int) -> dict:
        if self.model is None:
            self.load_model()
        if self.model is None or not frames:
            self._last_logits = np.zeros(2, dtype=np.float32)
            return {"classification": "Unknown", "confidence": 0.0,
                    "prediction_idx": -1, "logits": self._last_logits.tolist()}

        log.info("[Stage1-TSBinary] Binary classification (fixed-interval clip)...")
        clip = frames[max(0, peak_idx - 16): peak_idx + 16]
        while len(clip) < 32:
            clip.append(clip[-1] if clip else np.zeros((224, 224, 3), dtype=np.uint8))

        tensors = torch.stack([self.transform(f) for f in clip]).unsqueeze(0).to(self.config.DEVICE)
        with torch.no_grad():
            outputs = self.model(pixel_values=tensors)
            logits  = outputs.logits
            pred    = torch.argmax(logits, 1).item()
            conf    = torch.softmax(logits, 1).max().item()
            self._last_logits = logits.cpu().numpy().flatten().astype(np.float32)

        classification = "Normal" if pred == 0 else "Suspicious"
        cleanup_memory()
        log.info(f"[Stage1-TSBinary] {classification} ({conf:.2%})")
        return {"classification": classification, "confidence": conf,
                "prediction_idx": pred, "logits": self._last_logits.tolist()}

    def get_logits(self) -> np.ndarray:
        return self._last_logits if self._last_logits is not None else np.zeros(2, dtype=np.float32)


# ══════════════════════════════════════════════════════════════════════════════
# STAGE 2 — TIMESFORMER MULTI-CLASS
# ══════════════════════════════════════════════════════════════════════════════
class _TSMultiWrapper(nn.Module):
    """Replica of TimeSformerAnomalyClassifier from training."""
    def __init__(self, n_classes: int = 7):
        super().__init__()
        self.timesformer = TimesformerForVideoClassification.from_pretrained(
            "facebook/timesformer-base-finetuned-k400",
            num_labels=n_classes,
            ignore_mismatched_sizes=True,
        )
        h = self.timesformer.config.hidden_size
        self.classifier = nn.Sequential(
            nn.Linear(h, h // 2), nn.ReLU(), nn.Dropout(0.2), nn.Linear(h // 2, n_classes),
        )
        self.timesformer.classifier = self.classifier

    def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
        return self.timesformer(pixel_values=pixel_values).logits


class Stage2TSMulti:
    def __init__(self, config: Config):
        self.config    = config
        self.model     = None
        self.processor = None
        self._last_logits: Optional[np.ndarray] = None

    def load_model(self):
        if not TRANSFORMERS_AVAILABLE:
            return self
        log.info("[Stage2-TSMulti] Loading multi-class TimeSformer...")
        path = self.config.PATHS["TS_MULTI"]
        if not os.path.exists(path):
            log.warning(f"[Stage2-TSMulti] Not found: {path}")
            return self
        try:
            raw = torch.load(path, map_location=self.config.DEVICE, weights_only=False)
            state_dict = (raw.get("model_state_dict") or raw.get("model") or
                          raw.get("state_dict") or raw) if isinstance(raw, dict) else raw
            model   = _TSMultiWrapper(n_classes=7)
            cleaned = {k: v for k, v in state_dict.items()
                       if not k.startswith("timesformer.classifier.")}
            n_removed = len(state_dict) - len(cleaned)
            if n_removed:
                log.info(f"[Stage2-TSMulti] Stripped {n_removed} duplicate keys.")
            missing, unexpected = model.load_state_dict(cleaned, strict=False)
            if missing:
                log.warning(f"[Stage2-TSMulti] Missing keys ({len(missing)}): {missing[:3]}")
            if unexpected:
                log.warning(f"[Stage2-TSMulti] Unexpected keys ({len(unexpected)}): {unexpected[:3]}")
            self.model     = model.to(self.config.DEVICE).eval()
            self.processor = VideoMAEImageProcessor.from_pretrained(
                "facebook/timesformer-base-finetuned-k400")
            log.info("[Stage2-TSMulti] Loaded successfully.")
        except Exception as e:
            log.error(f"[Stage2-TSMulti] Load failed: {e}")
            import traceback; traceback.print_exc()
            self.model = None
        return self

    def analyze(self, frames: List[np.ndarray], peak_idx: int,
                ae_scores: Optional[List[float]] = None,
                sampled_indices: Optional[List[int]] = None) -> dict:
        if self.model is None or self.processor is None:
            self.load_model()
        null = {"class": "Unknown", "confidence": 0.0, "top3": [],
                "all_probabilities": {}, "logits": [0.0] * 7, "sampling_method": "unavailable"}
        if self.model is None:
            self._last_logits = np.zeros(7, dtype=np.float32)
            return null

        log.info("[Stage2-TSMulti] Multi-class classification...")
        if ae_scores and sampled_indices:
            clip            = segmentwise_ae_sampling(frames, ae_scores, sampled_indices, 16)
            sampling_method = "segmentwise_ae"
        else:
            start = max(0, peak_idx - 8)
            clip  = frames[start: min(len(frames), peak_idx + 8)]
            while len(clip) < 16:
                clip.append(clip[-1] if clip else np.zeros((224, 224, 3), dtype=np.uint8))
            clip            = clip[:16]
            sampling_method = "fixed_window_fallback"

        pil = [Image.fromarray(f) if isinstance(f, np.ndarray) else f for f in clip]
        try:
            inputs = self.processor(pil, return_tensors="pt")
            pv     = inputs["pixel_values"].to(self.config.DEVICE)
            with torch.no_grad():
                logits = self.model(pixel_values=pv)
                probs  = torch.softmax(logits, dim=-1)[0]
                self._last_logits = logits.cpu().numpy().flatten()[:7].astype(np.float32)
            top3_p, top3_i = torch.topk(probs, min(3, len(probs)))
            cls   = self.config.CLASS_NAMES.get(top3_i[0].item(), "Unknown")
            top3  = [(self.config.CLASS_NAMES.get(i.item(), "Unknown"), p.item())
                     for i, p in zip(top3_i, top3_p)]
            all_p = {self.config.CLASS_NAMES.get(i, "Unknown"): probs[i].item()
                     for i in range(len(probs))}
            cleanup_memory()
            log.info(f"[Stage2-TSMulti] {cls} ({top3_p[0].item():.2%}) via {sampling_method}")
            return {"class": cls, "confidence": top3_p[0].item(),
                    "top3": top3, "all_probabilities": all_p,
                    "logits": self._last_logits.tolist(), "sampling_method": sampling_method}
        except Exception as e:
            log.error(f"[Stage2-TSMulti] Inference failed: {e}")
            self._last_logits = np.zeros(7, dtype=np.float32)
            return {"class": "Error", "confidence": 0.0, "top3": [],
                    "all_probabilities": {}, "logits": self._last_logits.tolist(),
                    "sampling_method": sampling_method}

    def get_logits(self) -> np.ndarray:
        return self._last_logits if self._last_logits is not None else np.zeros(7, dtype=np.float32)


# ══════════════════════════════════════════════════════════════════════════════
# STAGE 3 — YOLO WEAPON DETECTOR (EXPERIMENTAL)
# ══════════════════════════════════════════════════════════════════════════════
class Stage3YOLOWeapon:
    _NULL_RESULT = {
        "weapon_flag":          False,
        "weapon_count":         0,
        "dominant_class":       "None",
        "best_detection":       {},
        "all_detections":       [],
        "top_frames_annotated": [],
        "spatial_heatmap":      None,  # filled at runtime
        "summary":              "YOLO weapon detection unavailable",
    }

    def __init__(self, config: Config):
        self.config = config
        self.model  = None

    def _null(self, summary: str = "") -> dict:
        r = dict(self._NULL_RESULT)
        r["spatial_heatmap"] = np.zeros(196, dtype=np.float32)
        if summary:
            r["summary"] = summary
        return r

    def load_model(self):
        if not YOLO_AVAILABLE:
            return self
        path = self.config.PATHS["YOLO"]
        try:
            self.model = YOLO(path if os.path.exists(path) else "yolov8n.pt")
            log.info(f"[Stage3-YOLO] Loaded from: {path if os.path.exists(path) else 'yolov8n.pt'}")
        except Exception as e:
            log.warning(f"[Stage3-YOLO] Load failed: {e}")
            self.model = None
        return self

    def analyze(self, frames: List[np.ndarray], anomaly_start_frame: int,
                anomaly_end_frame: int, max_frames: int = 32) -> dict:
        if self.model is None:
            self.load_model()
        if self.model is None or not frames:
            return self._null()

        a_start = max(0, anomaly_start_frame)
        a_end   = min(len(frames) - 1, anomaly_end_frame)
        if a_end - a_start < 0:
            return self._null()

        step         = max(1, (a_end - a_start + 1) // max_frames)
        scan_indices = list(range(a_start, a_end + 1, step))[:max_frames]
        log.info(f"[Stage3-YOLO] Scanning {len(scan_indices)} frames in [{a_start},{a_end}]...")

        all_detections: List[dict]           = []
        best_detection: dict                 = {}
        best_conf                            = 0.0
        class_counts: Dict[int, int]         = {0: 0, 1: 0}
        spatial_grid                         = np.zeros(196, dtype=np.float32)
        frame_conf_list: List[Tuple]         = []

        for fi in scan_indices:
            try:
                results_yolo    = self.model(frames[fi], verbose=False)[0]
                frame_has_det   = False
                frame_max_conf  = 0.0
                annotated_frame = frames[fi].copy()
                for box in results_yolo.boxes:
                    cls_idx = int(box.cls[0])
                    if cls_idx not in self.config.WEAPON_CLASS_NAMES:
                        continue
                    conf = float(box.conf[0])
                    x1, y1, x2, y2 = [int(v) for v in box.xyxy[0]]
                    det = {"class_idx": cls_idx,
                           "class_name": self.config.WEAPON_CLASS_NAMES[cls_idx],
                           "confidence": conf, "frame_idx": fi}
                    all_detections.append(det)
                    class_counts[cls_idx] = class_counts.get(cls_idx, 0) + 1
                    if conf > best_conf:
                        best_conf      = conf
                        best_detection = det
                    h, w = frames[fi].shape[:2]
                    gx   = min(int((x1 + x2) / 2 / max(w, 1) * 14), 13)
                    gy   = min(int((y1 + y2) / 2 / max(h, 1) * 14), 13)
                    spatial_grid[gy * 14 + gx] += conf
                    cv2.rectangle(annotated_frame, (x1, y1), (x2, y2), (255, 50, 50), 3)
                    cv2.putText(annotated_frame,
                                f"{self.config.WEAPON_CLASS_NAMES[cls_idx]} {conf:.0%}",
                                (x1, max(y1 - 10, 15)), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 50, 50), 2)
                    frame_has_det  = True
                    frame_max_conf = max(frame_max_conf, conf)
                if frame_has_det:
                    frame_conf_list.append((fi, frame_max_conf, annotated_frame))
            except Exception as e:
                log.warning(f"[Stage3-YOLO] Frame {fi} failed: {e}")

        frame_conf_list.sort(key=lambda x: x[1], reverse=True)
        weapon_count   = len(all_detections)
        weapon_flag    = weapon_count > 0
        dom_cls_idx    = max(class_counts, key=class_counts.get)
        dominant_class = self.config.WEAPON_CLASS_NAMES[dom_cls_idx] if weapon_flag else "None"
        s_max          = spatial_grid.max()
        if s_max > 0:
            spatial_grid /= s_max
        summary = (f"Weapons detected, dominant: {dominant_class}, best confidence: {best_conf:.2%}"
                   if weapon_flag else "No weapons detected")
        log.info(f"[Stage3-YOLO] {summary}")
        return {
            "weapon_flag":          weapon_flag,
            "weapon_count":         weapon_count,
            "dominant_class":       dominant_class,
            "best_conf":            best_conf,
            "best_detection":       best_detection,
            "all_detections":       all_detections,
            "top_frames_annotated": [(fi, ann) for fi, _, ann in frame_conf_list[:2]],
            "spatial_heatmap":      spatial_grid,
            "summary":              summary,
        }


# ══════════════════════════════════════════════════════════════════════════════
# STAGE 4 — AUDIO FEATURE EXTRACTION
# ══════════════════════════════════════════════════════════════════════════════
class Stage4Audio:
    _NORM_RANGES = {
        "zcr":               (0.0, 0.5),
        "mfcc_variance":     (0.0, 200.0),
        "spectral_flux":     (0.0, 50.0),
        "rms_energy":        (0.0, 0.5),
        "spectral_centroid": (0.0, 8000.0),
    }

    def __init__(self, config: Config):
        self.config = config

    def _normalize(self, value: float, key: str) -> float:
        lo, hi = self._NORM_RANGES.get(key, (0.0, 1.0))
        return float(np.clip((value - lo) / max(hi - lo, 1e-8), 0.0, 1.0))

    def _null_result(self, status: str) -> dict:
        return {
            "status": status, "audio_present": False,
            "zcr": 0.0, "mfcc_variance": 0.0, "spectral_flux": 0.0,
            "rms_energy": 0.0, "spectral_centroid": 0.0,
            "audio_score": 0.0, "audio_features_vec": np.zeros(5, dtype=np.float32),
        }

    def analyze(self, video_path: str) -> dict:
        log.info("[Stage4-Audio] Extracting audio features...")
        audio_path = extract_audio(video_path)
        if not audio_path:
            return self._null_result("NO_AUDIO")
        if not check_audio_has_content(audio_path):
            log.info("[Stage4-Audio] Silent or absent audio.")
            return self._null_result("SILENT")
        if not LIBROSA_AVAILABLE:
            return self._null_result("LIBROSA_UNAVAILABLE")
        try:
            y, sr = librosa.load(audio_path, sr=22050, duration=30)
            zcr_raw  = float(np.mean(librosa.feature.zero_crossing_rate(y)[0]))
            mfcc_raw = librosa.feature.mfcc(y=y, sr=sr, n_mfcc=13)
            mfcc_var = float(np.mean(np.var(mfcc_raw, axis=1)))
            S        = np.abs(librosa.stft(y))
            flux     = float(np.mean(np.sqrt(np.mean(np.diff(S, axis=1) ** 2, axis=0))))
            rms      = float(np.mean(librosa.feature.rms(y=y)[0]))
            sc       = float(np.mean(librosa.feature.spectral_centroid(y=y, sr=sr)[0]))

            zcr_n  = self._normalize(zcr_raw,  "zcr")
            mfcc_n = self._normalize(mfcc_var,  "mfcc_variance")
            flux_n = self._normalize(flux,       "spectral_flux")
            rms_n  = self._normalize(rms,        "rms_energy")
            sc_n   = self._normalize(sc,         "spectral_centroid")

            audio_score = float(np.clip(
                0.30 * flux_n + 0.25 * mfcc_n + 0.20 * zcr_n +
                0.15 * rms_n  + 0.10 * sc_n, 0.0, 1.0))
            n_anom = sum([flux_n > 0.4, mfcc_n > 0.4, zcr_n > 0.4, rms_n > 0.4])
            status = "ANOMALY" if n_anom >= 2 else "SUSPICIOUS" if n_anom == 1 else "NORMAL"
            log.info(f"[Stage4-Audio] {status} | audio_score={audio_score:.3f}")
            return {
                "status": status, "audio_present": True,
                "zcr": zcr_raw, "mfcc_variance": mfcc_var,
                "spectral_flux": flux, "rms_energy": rms, "spectral_centroid": sc,
                "audio_score": audio_score,
                "audio_features_vec": np.array([zcr_n, mfcc_n, flux_n, rms_n, sc_n], dtype=np.float32),
                "normalized": {"zcr": zcr_n, "mfcc": mfcc_n, "flux": flux_n, "rms": rms_n, "sc": sc_n},
            }
        except Exception as e:
            log.error(f"[Stage4-Audio] Failed: {e}")
            return self._null_result("ERROR")


# ══════════════════════════════════════════════════════════════════════════════
# STAGE 5 — TRACKING / BEHAVIOR EXTRACTION
# ══════════════════════════════════════════════════════════════════════════════
class Stage5Tracking:
    def __init__(self, config: Config):
        self.config = config

    def analyze(self, frames: List[np.ndarray], anomaly_start_frame: int,
                n_frames: int = 32) -> dict:
        null = {"motion_intensity": 0.0, "behaviors": [], "track_count": 0,
                "tracking_features": np.zeros(20, dtype=np.float32), "status": "UNAVAILABLE"}
        if not frames:
            return null
        clip = frames[max(0, anomaly_start_frame): max(0, anomaly_start_frame) + n_frames]
        if len(clip) < 2:
            return null
        if YOLO_AVAILABLE:
            try:
                return self._yolo_track(clip)
            except Exception as e:
                log.warning(f"[Stage5-Tracking] YOLO tracking failed: {e}, falling back to optical flow.")
        return self._optical_flow_track(clip)

    def _yolo_track(self, clip: List[np.ndarray]) -> dict:
        try:
            tracker_model = YOLO("yolov8n.pt")
        except Exception as e:
            raise RuntimeError(f"Cannot load YOLOv8n for tracking: {e}")

        all_tracks: Dict[int, List[Tuple[float, float]]] = {}
        per_frame_counts = []
        for frame in clip:
            try:
                results = tracker_model.track(frame, persist=True, verbose=False)[0]
                if results.boxes.id is None:
                    per_frame_counts.append(0)
                    continue
                per_frame_counts.append(len(results.boxes.id))
                for tid, box in zip(results.boxes.id.int().tolist(),
                                    results.boxes.xyxy.tolist()):
                    cx = (box[0] + box[2]) / 2
                    cy = (box[1] + box[3]) / 2
                    all_tracks.setdefault(int(tid), []).append((cx, cy))
            except Exception:
                per_frame_counts.append(0)

        behaviors, motion_intensity = self._analyze_tracks(all_tracks)
        return {
            "motion_intensity":  motion_intensity,
            "behaviors":         behaviors,
            "track_count":       len(all_tracks),
            "tracking_features": self._build_tracking_features(
                all_tracks, motion_intensity, behaviors, per_frame_counts),
            "status": "YOLO_TRACK",
        }

    def _optical_flow_track(self, clip: List[np.ndarray]) -> dict:
        magnitudes = []
        prev_gray  = cv2.cvtColor(clip[0], cv2.COLOR_RGB2GRAY)
        for frame in clip[1:]:
            gray = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY)
            try:
                flow = cv2.calcOpticalFlowFarneback(
                    prev_gray, gray, None, 0.5, 3, 15, 3, 5, 1.2, 0)
                mag, _ = cv2.cartToPolar(flow[..., 0], flow[..., 1])
                magnitudes.append(float(np.mean(mag)))
            except Exception:
                pass
            prev_gray = gray

        if not magnitudes:
            return {"motion_intensity": 0.0, "behaviors": [], "track_count": 0,
                    "tracking_features": np.zeros(20, dtype=np.float32),
                    "status": "OPTICAL_FLOW_FAILED"}

        mag_arr          = np.array(magnitudes)
        mean_mag         = float(np.mean(mag_arr))
        std_mag          = float(np.std(mag_arr))
        motion_intensity = float(np.clip(mean_mag / 15.0, 0.0, 1.0))
        behaviors        = []
        if mean_mag > 8.0:
            behaviors.append("FAST_MOVEMENT")
        elif mean_mag < 1.5:
            behaviors.append("STATIC")
        if std_mag > mean_mag * 0.5 and mean_mag > 3.0:
            behaviors.append("CLUSTERING")

        feat    = np.zeros(20, dtype=np.float32)
        feat[0] = float(np.clip(mean_mag / 15.0, 0, 1))
        feat[1] = float(np.clip(float(np.max(mag_arr)) / 30.0, 0, 1))
        feat[2] = float(np.clip(std_mag / 10.0, 0, 1))
        feat[3] = float("FAST_MOVEMENT" in behaviors)
        feat[4] = float("STATIC" in behaviors)
        feat[5] = float("CLUSTERING" in behaviors)
        series_len = min(14, len(magnitudes))
        feat[6:6 + series_len] = np.clip(
            np.array(magnitudes[:series_len]) / 15.0, 0, 1).astype(np.float32)
        return {"motion_intensity": motion_intensity, "behaviors": behaviors,
                "track_count": 0, "tracking_features": feat, "status": "OPTICAL_FLOW"}

    @staticmethod
    def _analyze_tracks(tracks: Dict[int, List[Tuple[float, float]]]) -> Tuple[List[str], float]:
        behaviors, all_displacements = [], []
        tids = list(tracks.keys())
        for pts in tracks.values():
            if len(pts) < 2:
                continue
            disp = [math.hypot(pts[i][0]-pts[i-1][0], pts[i][1]-pts[i-1][1])
                    for i in range(1, len(pts))]
            all_displacements.extend(disp)
            avg = sum(disp) / max(len(disp), 1)
            if avg > 8.0:
                behaviors.append("FAST_MOVEMENT")
            if avg < 2.0 and len(pts) >= 10:
                behaviors.append("STATIC")
        for i in range(len(tids)):
            for j in range(i + 1, len(tids)):
                a, b = tracks[tids[i]], tracks[tids[j]]
                n    = min(len(a), len(b))
                if n < 6:
                    continue
                dists = [math.hypot(a[k][0]-b[k][0], a[k][1]-b[k][1]) for k in range(n)]
                dec   = sum(1 for k in range(1, n) if dists[k] < dists[k-1])
                if dec / n > 0.75 and dists[0] - dists[-1] > 50:
                    behaviors.append("CLUSTERING")
                    break
        behaviors        = list(dict.fromkeys(behaviors)) or ["STATIC"]
        motion_intensity = float(np.clip(np.mean(all_displacements) / 15.0, 0.0, 1.0)) if all_displacements else 0.0
        return behaviors, motion_intensity

    @staticmethod
    def _build_tracking_features(tracks, motion_intensity, behaviors, per_frame_counts) -> np.ndarray:
        feat    = np.zeros(20, dtype=np.float32)
        feat[0] = float(np.clip(motion_intensity, 0, 1))
        feat[1] = float(np.clip(len(tracks) / 10.0, 0, 1))
        feat[2] = float("FAST_MOVEMENT" in behaviors)
        feat[3] = float("STATIC" in behaviors)
        feat[4] = float("CLUSTERING" in behaviors)
        feat[5] = float(np.clip(np.mean(per_frame_counts) / 10.0, 0, 1)) if per_frame_counts else 0.0
        all_v   = [math.hypot(pts[i][0]-pts[i-1][0], pts[i][1]-pts[i-1][1])
                   for pts in tracks.values() for i in range(1, len(pts)) if len(pts) >= 2]
        feat[6] = float(np.clip(np.std(all_v) / 10.0, 0, 1)) if all_v else 0.0
        series_len  = min(13, len(per_frame_counts))
        feat[7:7 + series_len] = np.clip(
            np.array(per_frame_counts[:series_len]) / 10.0, 0, 1).astype(np.float32)
        return feat


# ══════════════════════════════════════════════════════════════════════════════
# STAGE 6 — FUSION LAYER
# ══════════════════════════════════════════════════════════════════════════════
class Stage6Fusion:
    def __init__(self, config: Config):
        self.config = config

    def fuse(self, ae_result: dict, ts_binary_result: dict, ts_multi_result: dict,
             yolo_result: dict, audio_result: dict, tracking_result: dict) -> dict:
        reasoning_trace: List[str] = []
        threshold     = ae_result.get("threshold", self.config.SAFE_THRESHOLD_LEGACY)
        ae_max        = float(ae_result.get("max_error", 0.0))
        ae_score_norm = float(np.clip(ae_max / max(threshold, 1e-8), 0.0, 1.0))

        ts_binary_susp = (ts_binary_result.get("classification") == "Suspicious")
        ts_binary_conf = float(ts_binary_result.get("confidence", 0.0) if ts_binary_susp else 0.0)
        weapon_flag    = bool(yolo_result.get("weapon_flag", False))
        audio_score    = float(audio_result.get("audio_score", 0.0))
        audio_anomaly  = (audio_result.get("status") == "ANOMALY")
        motion_int     = float(tracking_result.get("motion_intensity", 0.0))
        fast_movement  = "FAST_MOVEMENT" in tracking_result.get("behaviors", [])

        # Rule 0: base AE weight
        fused_risk = ae_score_norm * 0.50
        reasoning_trace.append(
            f"Baseline: AE normalised score is {ae_score_norm:.3f} (50% weight), "
            f"giving initial base risk of {fused_risk:.3f}.")

        # Rule 1: weapon + AE hard escalation
        if weapon_flag and ae_max > threshold:
            fused_risk = max(fused_risk, 0.90)
            reasoning_trace.append(
                "Rule 1 TRIGGERED: Weapon detected AND AE exceeds threshold. "
                "Risk escalated to CRITICAL (minimum 0.90).")
        elif ts_binary_susp:
            ts_boost    = ts_binary_conf * 0.20
            fused_risk += ts_boost
            reasoning_trace.append(
                f"Rule 2: Binary TimeSformer classified as Suspicious at "
                f"{ts_binary_conf:.2f} confidence (20% weight), adding {ts_boost:.3f}.")

        # Rule 3: motion + audio boosts
        if fast_movement and audio_anomaly and ae_max > threshold:
            escalation  = self.config.FUSION_MOTION_BOOST + self.config.FUSION_AUDIO_BOOST
            fused_risk += escalation
            reasoning_trace.append(
                f"Rule 3 FULLY TRIGGERED: Fast movement, audio anomaly, and AE anomaly "
                f"all present. Adding {escalation:.3f}.")
        elif fast_movement:
            boost       = self.config.FUSION_MOTION_BOOST * 0.5
            fused_risk += boost
            reasoning_trace.append(f"Rule 3 (partial): Fast movement detected. Adding {boost:.3f}.")
        elif audio_anomaly:
            boost       = self.config.FUSION_AUDIO_BOOST * 0.5
            fused_risk += boost
            reasoning_trace.append(f"Rule 3 (partial): Audio anomaly detected. Adding {boost:.3f}.")

        # Rule 4: no support penalty
        has_support = ts_binary_susp or weapon_flag or audio_anomaly or fast_movement
        if ae_max > threshold and not has_support:
            fused_risk -= self.config.FUSION_NO_SUPPORT_PENALTY
            reasoning_trace.append(
                f"Rule 4: AE exceeds threshold but no supporting signal. "
                f"Applying deduction of {self.config.FUSION_NO_SUPPORT_PENALTY:.3f}.")

        # Rule 5: weapon floor
        if weapon_flag:
            pre_floor  = fused_risk
            fused_risk = max(fused_risk, self.config.FUSION_WEAPON_BOOST)
            if fused_risk != pre_floor:
                reasoning_trace.append(
                    f"Rule 5: Weapon floor applied at {self.config.FUSION_WEAPON_BOOST:.2f}.")

        fused_risk = float(np.clip(fused_risk, 0.0, 1.0))
        if weapon_flag and ae_max > threshold:
            risk_level = "CRITICAL"
        elif fused_risk >= 0.70:
            risk_level = "HIGH"
        elif fused_risk >= 0.40:
            risk_level = "MODERATE"
        else:
            risk_level = "LOW"
        reasoning_trace.append(
            f"Final: fused risk is {fused_risk:.4f}, classified as {risk_level}.")

        context_vector = np.array([
            float(np.clip(ae_score_norm, 0.0, 1.0)),
            float(ts_binary_conf),
            float(np.clip(audio_score,  0.0, 1.0)),
            float(weapon_flag),
            float(np.clip(motion_int,   0.0, 1.0)),
        ], dtype=np.float32)

        log.info(f"[Stage6-Fusion] {risk_level} | fused_risk={fused_risk:.4f}")
        return {
            "fused_risk":       fused_risk,
            "risk_level":       risk_level,
            "reasoning_trace":  reasoning_trace,
            "context_vector":   context_vector,
            "ae_score_norm":    ae_score_norm,
            "ts_binary_conf":   ts_binary_conf,
            "audio_score":      audio_score,
            "weapon_flag":      weapon_flag,
            "motion_intensity": motion_int,
        }


# ══════════════════════════════════════════════════════════════════════════════
# STAGE 7 — RL3033 (EXPERIMENTAL)
# ══════════════════════════════════════════════════════════════════════════════
class Stage7RL3033:
    def __init__(self, config: Config):
        self.config = config
        self.agent  = None

    def load_agent(self):
        if not RL_AVAILABLE:
            log.warning("[Stage7-RL3033] stable-baselines3 unavailable.")
            return self
        path = self.config.PATHS["RL_AGENT"]
        if not os.path.exists(path):
            log.warning(f"[Stage7-RL3033] Agent not found: {path}")
            return self
        log.info("[Stage7-RL3033] Loading RL3033 agent...")
        try:
            self.agent = PPO.load(path, device="cpu")
            log.info(f"[Stage7-RL3033] Loaded. obs_space={self.agent.observation_space}")
        except Exception as e:
            log.error(f"[Stage7-RL3033] Load failed: {e}")
            self.agent = None
        return self

    def _ch_stats(self, ch: np.ndarray) -> np.ndarray:
        return np.array([
            np.mean(ch), np.std(ch), np.max(ch), np.min(ch),
            np.percentile(ch, 75), np.percentile(ch, 25),
            np.median(ch), float(np.sum(ch > 0)) / max(len(ch), 1),
        ], dtype=np.float32)

    def _build_latent_motion(self, ae_signals: np.ndarray, tracking_result: dict) -> np.ndarray:
        ae_flat    = ae_signals.flatten()
        track_feat = tracking_result.get("tracking_features", np.zeros(20, dtype=np.float32))
        ch_desc    = np.concatenate([self._ch_stats(ae_signals[:, c]) for c in [0, 1, 2, 5]])
        motion_int = float(tracking_result.get("motion_intensity", 0.0))
        behaviors  = tracking_result.get("behaviors", [])
        motion_sum = np.array([
            motion_int,
            float("FAST_MOVEMENT" in behaviors),
            float("STATIC"        in behaviors),
            float("CLUSTERING"    in behaviors),
            float(np.clip(tracking_result.get("track_count", 0) / 10.0, 0, 1)),
        ], dtype=np.float32)
        base = np.concatenate([ae_flat, ch_desc, track_feat, motion_sum])
        return np.tile(base, math.ceil(1536 / len(base)))[:1536].astype(np.float32)

    def _build_latent_stack(self, ae_signals, ts_binary_logits, ts_multi_logits,
                             audio_result, tracking_result, fusion_result) -> np.ndarray:
        fusion_signals = np.array([
            float(fusion_result.get("fused_risk",       0.0)),
            float(fusion_result.get("ae_score_norm",    0.0)),
            float(fusion_result.get("ts_binary_conf",   0.0)),
            float(fusion_result.get("audio_score",      0.0)),
            float(fusion_result.get("weapon_flag",      False)),
            float(fusion_result.get("motion_intensity", 0.0)),
            float(fusion_result.get("risk_level") == "CRITICAL"),
            float(fusion_result.get("risk_level") == "HIGH"),
        ], dtype=np.float32)
        base = np.concatenate([
            ae_signals.flatten().astype(np.float32),
            ts_binary_logits.astype(np.float32),
            ts_multi_logits[:7].astype(np.float32),
            audio_result.get("audio_features_vec", np.zeros(5, dtype=np.float32)),
            tracking_result.get("tracking_features", np.zeros(20, dtype=np.float32)),
            fusion_result.get("context_vector", np.zeros(5, dtype=np.float32)),
            fusion_signals,
        ])
        return np.tile(base, math.ceil(2560 / len(base)))[:2560].astype(np.float32)

    def _build_spatial(self, yolo_result: dict, tracking_result: dict,
                        frames: List[np.ndarray], peak_frame_idx: int) -> np.ndarray:
        spatial      = np.zeros(196, dtype=np.float32)
        spatial     += yolo_result.get("spatial_heatmap", np.zeros(196, dtype=np.float32)) * 0.5
        if frames and peak_frame_idx < len(frames):
            gray = cv2.cvtColor(frames[peak_frame_idx], cv2.COLOR_RGB2GRAY).astype(np.float32)
            H, W = gray.shape
            bh, bw = H // 14, W // 14
            if bh > 0 and bw > 0:
                for gy in range(14):
                    for gx in range(14):
                        spatial[gy*14 + gx] += float(np.var(gray[gy*bh:(gy+1)*bh, gx*bw:(gx+1)*bw])) / 10000.0
        motion_int = float(tracking_result.get("motion_intensity", 0.0))
        for cy in range(4, 10):
            for cx in range(4, 10):
                spatial[cy*14 + cx] += motion_int * 0.3
        s_max = spatial.max()
        if s_max > 0:
            spatial /= s_max
        return spatial

    def _build_observation(self, ae_signals, ts_binary_logits, ts_multi_logits,
                            yolo_result, audio_result, tracking_result,
                            fusion_result, frames, peak_frame_idx) -> dict:
        obs = {
            "ae_signals":    ae_signals.astype(np.float32),
            "latent_motion": self._build_latent_motion(ae_signals, tracking_result),
            "latent_stack":  self._build_latent_stack(ae_signals, ts_binary_logits,
                                                       ts_multi_logits, audio_result,
                                                       tracking_result, fusion_result),
            "spatial":       self._build_spatial(yolo_result, tracking_result, frames, peak_frame_idx),
            "context":       fusion_result.get("context_vector", np.zeros(5, dtype=np.float32)).astype(np.float32),
        }
        expected = {
            "ae_signals":    (self.config.RL_AE_WINDOW, self.config.RL_AE_CHANNELS),
            "latent_motion": (self.config.RL_LATENT_MOT,),
            "latent_stack":  (self.config.RL_LATENT_STK,),
            "spatial":       (self.config.RL_SPATIAL_DIM,),
            "context":       (self.config.RL_CONTEXT_DIM,),
        }
        for key, exp_shape in expected.items():
            if obs[key].shape != exp_shape:
                flat   = obs[key].flatten()
                needed = exp_shape[0] if len(exp_shape) == 1 else exp_shape[0] * exp_shape[1]
                resized = flat[:needed] if len(flat) >= needed else np.pad(flat, (0, needed - len(flat)))
                obs[key] = resized.reshape(exp_shape) if len(exp_shape) == 2 else resized
        return obs

    def _fallback_result(self, fusion_result: dict, label: str = "") -> dict:
        fused_risk  = float(fusion_result.get("fused_risk", 0.0))
        tfb_score   = float(fusion_result.get("ts_binary_conf", 0.0))
        tfm_score   = float(fusion_result.get("ae_score_norm", 0.0))
        final_risk  = float(np.clip(
            0.70 * tfb_score + 0.20 * tfm_score + 0.10 * fused_risk, 0.0, 1.0))
        if fusion_result.get("weapon_flag", False):
            final_risk = max(final_risk, 0.85)
        # Deterministic dummy RL value derived from fused/final risk — never random
        rl_raw = float(np.clip(0.5 * fused_risk + 0.5 * final_risk, 0.0, 1.0))
        status = ("CRITICAL" if final_risk >= 0.85 else "HIGH_RISK" if final_risk >= 0.65
                  else "MODERATE" if final_risk >= 0.40 else "LOW_RISK")
        return {
            "final_risk": final_risk, "fused_risk": fused_risk, "rl_raw": rl_raw, "status": status,
            "interpretation": (
                f"{label}Final risk from weighted formula: "
                f"70% TFB ({tfb_score:.2%}) + 20% TFM ({tfm_score:.2%}) + "
                f"10% Fused ({fused_risk:.2%}) = {final_risk:.2%} [{status}]. "
                f"RL component approximated deterministically as {rl_raw:.2%} "
                f"(average of fused and final risk)."
            ),
        }

    def analyze(self, ae_signals, ts_binary_logits, ts_multi_logits,
                yolo_result, audio_result, tracking_result,
                fusion_result, frames, peak_frame_idx) -> dict:
        if self.agent is None:
            self.load_agent()
        if self.agent is None:
            return self._fallback_result(fusion_result, label="RL agent unavailable. ")

        try:
            obs = self._build_observation(ae_signals, ts_binary_logits, ts_multi_logits,
                                           yolo_result, audio_result, tracking_result,
                                           fusion_result, frames, peak_frame_idx)
            action, _ = self.agent.predict(obs, deterministic=True)
            action_arr = np.asarray(action, dtype=np.float64).flatten()
            rl_raw     = float(np.clip(action_arr[0] if action_arr.size else 0.0, 0.0, 1.0))

            tfb_score   = float(fusion_result.get("ts_binary_conf", 0.0))
            tfm_score   = float(fusion_result.get("ae_score_norm",  0.0))
            fused_score = float(np.clip(float(fusion_result.get("fused_risk", 0.0)), 0.0, 1.0))
            final_risk  = float(np.clip(
                0.70 * tfb_score + 0.20 * tfm_score + 0.05 * fused_score + 0.05 * rl_raw,
                0.0, 1.0))
            if fusion_result.get("weapon_flag", False):
                final_risk = max(final_risk, 0.85)

            status = ("CRITICAL" if final_risk >= 0.85 else "HIGH_RISK" if final_risk >= 0.65
                      else "MODERATE" if final_risk >= 0.40 else "LOW_RISK")
            log.info(f"[Stage7-RL3033] tfb={tfb_score:.4f} tfm={tfm_score:.4f} "
                     f"fused={fused_score:.4f} rl={rl_raw:.4f} final={final_risk:.4f} => {status}")
            return {
                "fused_risk":     float(fusion_result.get("fused_risk", 0.0)),
                "final_risk":     final_risk,
                "status":         status,
                "tfb_score":      tfb_score,
                "tfm_score":      tfm_score,
                "rl_raw":         rl_raw,
                "interpretation": (
                    f"Final risk: 70% TFB ({tfb_score:.2%}) + 20% TFM ({tfm_score:.2%}) + "
                    f"5% Fused ({fused_score:.2%}) + 5% RL ({rl_raw:.2%}) = {final_risk:.2%} [{status}]."
                ),
            }
        except Exception as e:
            import traceback
            log.error(f"[Stage7-RL3033] Inference failed: {e}")
            traceback.print_exc()
            result = self._fallback_result(fusion_result, label="RL inference failed, using deterministic fallback. ")
            result["status"] = "FALLBACK"
            result["error"]  = str(e)
            return result

# ══════════════════════════════════════════════════════════════════════════════
# STAGE 8 — LLAVA (EXPERIMENTAL)
# Always runs when enable_llava=True. Neutral prompts regardless of Stage 1 result.
# ══════════════════════════════════════════════════════════════════════════════
class Stage8LLaVA:
    _TEMPORAL_LABELS = {
        "baseline":   "Frame 1 of 4 (beginning of clip)",
        "anom_start": "Frame 2 of 4 (area of elevated AE error onset)",
        "peak":       "Frame 3 of 4 (highest AE reconstruction error)",
        "anom_end":   "Frame 4 of 4 (end of elevated-error region)",
    }
    _ROLE_INSTRUCTIONS = {
        "baseline":   "Describe the scene at the start of the clip: who is present and what is happening.",
        "anom_start": "Describe what is visible at the onset of the region with elevated reconstruction error.",
        "peak":       "Describe what is most notable in this frame at the reconstruction error peak.",
        "anom_end":   "Describe how the scene looks at the end of the elevated-error region.",
    }

    def __init__(self, config: Config):
        self.config    = config
        self.model     = None
        self.processor = None

    def load_model(self):
        if not TRANSFORMERS_AVAILABLE:
            return self
        path = self.config.PATHS["LLAVA"]
        if not os.path.exists(path):
            log.warning(f"[Stage8-LLaVA] Model not found: {path}")
            return self
        log.info("[Stage8-LLaVA] Loading Fire-LLaVA...")
        try:
            self.model = LlavaForConditionalGeneration.from_pretrained(
                path, torch_dtype=torch.float16, device_map="auto", low_cpu_mem_usage=True)
            self.processor = AutoProcessor.from_pretrained("llava-hf/llava-1.5-7b-hf")
            self.processor.patch_size                     = 14
            self.processor.vision_feature_select_strategy = "default"
            self.model.eval()
            log.info("[Stage8-LLaVA] Loaded.")
        except Exception as e:
            log.error(f"[Stage8-LLaVA] Load failed: {e}")
            self.model = None
        return self

    def unload(self):
        """Deployment fix: release the ~14 GB VLM after each analysis so its
        weights do not remain resident between runs on Space hardware."""
        if self.model is not None:
            del self.model
            self.model = None
        self.processor = None
        cleanup_memory()
        log.info("[Stage8-LLaVA] Model released.")

    def _build_prompt(self, role: str, frame_number: int, yolo_summary: str,
                      motion_behaviors: List[str], audio_status: str,
                      audio_score: float, binary_result: str) -> str:
        t_label       = self._TEMPORAL_LABELS.get(role, f"Frame {frame_number} of 4")
        behaviors_str = ", ".join(motion_behaviors) if motion_behaviors else "STATIC"
        weapon_info   = yolo_summary if "weapon" in yolo_summary.lower() else "No weapons detected"
        instruction   = self._ROLE_INSTRUCTIONS.get(role, "Describe what you see.")
        return (
            f"USER: <image>\n"
            f"You are a forensic security analyst reviewing surveillance footage.\n\n"
            f"=== TEMPORAL CONTEXT ===\n"
            f"Position: {t_label}\n"
            f"Classifier note: The binary classifier result is: {binary_result}. "
            f"Use this as context, but form your own description from the image.\n\n"
            f"=== DETECTED INFORMATION ===\n"
            f"Objects/Weapons: {weapon_info}\n"
            f"Motion Pattern: {behaviors_str}\n"
            f"Audio Status: {audio_status} (score={audio_score:.2f})\n\n"
            f"=== YOUR TASK ===\n"
            f"{instruction}\n"
            f"Be concise (2-3 sentences). Focus on safety-relevant details.\n"
            f"ASSISTANT:"
        )

    def _describe_frame(self, image: Image.Image, role: str, frame_number: int,
                         context: dict, binary_result: str) -> str:
        prompt = self._build_prompt(
            role=role, frame_number=frame_number,
            yolo_summary=context.get("yolo_summary", "Unknown"),
            motion_behaviors=context.get("behaviors", []),
            audio_status=context.get("audio_status", "Unknown"),
            audio_score=context.get("audio_score", 0.0),
            binary_result=binary_result,
        )
        inputs = self.processor(text=prompt, images=image.convert("RGB"), return_tensors="pt")
        inputs = {k: v.to(self.config.DEVICE) for k, v in inputs.items()}
        with torch.no_grad():
            with torch.amp.autocast("cuda"):
                output_ids = self.model.generate(
                    **inputs, max_new_tokens=130, do_sample=True, temperature=0.7,
                    top_p=0.9, top_k=50, repetition_penalty=1.4,
                    no_repeat_ngram_size=3, early_stopping=True)
        text = self.processor.batch_decode(output_ids, skip_special_tokens=True)[0]
        if "ASSISTANT:" in text:
            text = text.split("ASSISTANT:")[-1].strip()
        return self._clean_text(text)

    @staticmethod
    def _clean_text(text: str) -> str:
        if not text or len(text) < 20:
            return text
        seen, out = set(), []
        for s in text.replace("!", ".").replace("?", ".").split("."):
            s   = s.strip()
            key = s.lower()
            if s and len(s) > 10 and key not in seen:
                seen.add(key)
                out.append(s)
        result = ". ".join(out)
        return result + ("." if result and not result.endswith(".") else "")

    @staticmethod
    def synthesise_narrative(frame_descriptions: dict, context: dict,
                              yolo_result: dict, audio_result: dict,
                              tracking_result: dict, binary_result: str = "Unknown") -> str:
        roles    = ("baseline", "anom_start", "peak", "anom_end")
        raw_desc = [frame_descriptions.get(r, "").strip() for r in roles]

        # Deduplicate sentences across all four descriptions
        seen_keys: set         = set()
        unique_parts: List[str] = []
        for raw in raw_desc:
            for sent in raw.replace("!", ".").replace("?", ".").split("."):
                sent = sent.strip()
                key  = " ".join(sent.lower().split())
                if len(sent) > 15 and key not in seen_keys:
                    seen_keys.add(key)
                    unique_parts.append(sent)

        # Build narrative: opening -> onset -> peak -> resolution
        narrative_parts: List[str] = []
        existing_keys = lambda: {" ".join(p.lower().split()) for p in narrative_parts}
        for raw in raw_desc:
            if not raw:
                continue
            key = " ".join(raw.lower().split())
            if key not in existing_keys():
                narrative_parts.append(raw)

        if not narrative_parts:
            return "Visual analysis did not produce a usable scene description for this clip."

        # Append system-derived context notes
        context_notes = []
        weapon_flag  = bool(yolo_result.get("weapon_flag", False))
        weapon_class = yolo_result.get("dominant_class", "None")
        weapon_conf  = float(yolo_result.get("best_conf", 0.0))
        audio_status = audio_result.get("status", "UNKNOWN")
        audio_score  = float(audio_result.get("audio_score", 0.0))
        behaviors    = tracking_result.get("behaviors", ["STATIC"])
        motion_int   = float(tracking_result.get("motion_intensity", 0.0))

        if weapon_flag:
            context_notes.append(
                f"Instrument detection flagged a {weapon_class} with {weapon_conf:.0%} confidence.")
        if audio_status in ("ANOMALY", "SUSPICIOUS"):
            context_notes.append(
                f"The audio track registered as {audio_status} (composite score {audio_score:.2f}), "
                f"suggesting acoustically unusual activity.")
        if "FAST_MOVEMENT" in behaviors:
            context_notes.append(
                f"Motion analysis recorded rapid subject movement (intensity {motion_int:.2f}), "
                f"consistent with an escalating incident.")
        elif "CLUSTERING" in behaviors:
            context_notes.append(
                f"Motion analysis detected converging trajectories (intensity {motion_int:.2f}), "
                f"a pattern associated with confrontational proximity.")
        elif motion_int < 0.15:
            context_notes.append(
                f"Motion analysis recorded minimal movement (intensity {motion_int:.2f}), "
                f"indicating a static scene or post-event lull.")

        combined = " ".join(narrative_parts)
        if context_notes:
            combined = combined.rstrip(".") + ". " + " ".join(context_notes)
        combined = combined.strip()
        if combined and not combined.endswith("."):
            combined += "."
        return combined

    def analyze(self, frames: List[np.ndarray], peak_idx: int,
                anomaly_start_frame: int, anomaly_end_frame: int,
                yolo_result: dict, audio_result: dict, tracking_result: dict,
                binary_result: str = "Unknown") -> dict:
        if self.model is None or self.processor is None:
            self.load_model()
        if self.model is None or not frames:
            return {"description": "LLaVA not available", "analysis": "",
                    "frame_descriptions": {}, "unified_narrative": "LLaVA not available"}

        if anomaly_end_frame < 0:
            anomaly_end_frame = min(peak_idx + 30, len(frames) - 1)
        context = {
            "yolo_summary": yolo_result.get("summary",  "No weapon detection"),
            "behaviors":    tracking_result.get("behaviors", ["STATIC"]),
            "audio_status": audio_result.get("status",   "UNKNOWN"),
            "audio_score":  float(audio_result.get("audio_score", 0.0)),
        }
        keyframes = {
            "baseline":   (frames[0],                                       1),
            "anom_start": (frames[min(anomaly_start_frame, len(frames)-1)], 2),
            "peak":       (frames[min(peak_idx,            len(frames)-1)], 3),
            "anom_end":   (frames[min(anomaly_end_frame,   len(frames)-1)], 4),
        }
        log.info(f"[Stage8-LLaVA] 4-frame analysis (binary={binary_result})...")
        descriptions = {}
        for role, (frame_arr, fn) in keyframes.items():
            log.info(f"[Stage8-LLaVA]   Describing {role} ({fn}/4)...")
            try:
                img               = Image.fromarray(frame_arr) if isinstance(frame_arr, np.ndarray) else frame_arr
                descriptions[role] = self._describe_frame(img, role, fn, context, binary_result)
            except Exception as e:
                log.warning(f"[Stage8-LLaVA] Frame '{role}' failed: {e}")
                descriptions[role] = f"[{role} analysis unavailable]"

        unified_narrative = self.synthesise_narrative(
            frame_descriptions=descriptions, context=context,
            yolo_result=yolo_result, audio_result=audio_result,
            tracking_result=tracking_result, binary_result=binary_result,
        )
        weapon_str = (f"WEAPON DETECTED: {yolo_result.get('dominant_class','?')}"
                      if yolo_result.get("weapon_flag") else "No weapons detected")
        motion_str = ", ".join(tracking_result.get("behaviors", ["STATIC"]))
        audio_str  = f"{context['audio_status']} (score={context['audio_score']:.2f})"
        unified = (
            f"[CONTEXT] {weapon_str}, Motion: {motion_str}, Audio: {audio_str}\n"
            f"[FRAME 1 (BEGINNING)] {descriptions.get('baseline','')}\n"
            f"[FRAME 2 (AE ONSET)]  {descriptions.get('anom_start','')}\n"
            f"[FRAME 3 (AE PEAK)]   {descriptions.get('peak','')}\n"
            f"[FRAME 4 (AE END)]    {descriptions.get('anom_end','')}"
        ).strip()
        cleanup_memory()
        return {"description": unified, "analysis": unified,
                "frame_descriptions": descriptions, "unified_narrative": unified_narrative,
                "context_used": context, "binary_result": binary_result}


# ══════════════════════════════════════════════════════════════════════════════
# ALERT SYSTEM
# ══════════════════════════════════════════════════════════════════════════════
class AlertSystem:
    def __init__(self, config: Config):
        self.config    = config
        self._handlers: List[Any] = []
        self.register(self._default_log_handler)

    def register(self, handler) -> "AlertSystem":
        self._handlers.append(handler)
        return self

    def evaluate_and_fire(self, results: dict) -> Optional[dict]:
        stage7     = results.get("stage7") or {}
        stage6     = results.get("stage6") or {}
        stage3     = results.get("stage3") or {}
        stage4     = results.get("stage4") or {}
        final_risk = float(stage7.get("final_risk", stage6.get("fused_risk", 0.0)))
        level = ("HIGH"   if final_risk >= self.config.ALERT_RISK_HIGH   else
                 "MEDIUM" if final_risk >= self.config.ALERT_RISK_MEDIUM  else
                 "LOW")
        if level == "LOW":
            return None
        alert = {
            "timestamp":       datetime.utcnow().isoformat() + "Z",
            "level":           level,
            "final_risk":      round(final_risk, 4),
            "fused_risk":      round(float(stage6.get("fused_risk",    0.0)), 4),
            "weapon_flag":     bool(stage3.get("weapon_flag", False)),
            "weapon_summary":  stage3.get("summary", ""),
            "audio_score":     round(float(stage4.get("audio_score", 0.0)), 4),
            "risk_level":      stage6.get("risk_level", "UNKNOWN"),
            "reasoning_trace": stage6.get("reasoning_trace", []),
            "rl_status":       stage7.get("status", "UNKNOWN"),
            "video_path":      results.get("video_path", ""),
        }
        for handler in self._handlers:
            try:
                handler(alert)
            except Exception as e:
                log.error(f"[AlertSystem] Handler failed: {e}")
        return alert

    @staticmethod
    def _default_log_handler(alert: dict):
        log.warning(
            f"[ALERT] {alert['level']}, "
            f"final_risk={alert['final_risk']:.2%}, "
            f"fused_risk={alert['fused_risk']:.2%}, "
            f"weapon={'YES' if alert['weapon_flag'] else 'NO'}, "
            f"rl_status={alert['rl_status']}"
        )

    @staticmethod
    def file_alert_handler(log_path: str = "/tmp/scenesolver_alerts.jsonl"):
        def _handler(alert: dict):
            with open(log_path, "a") as f:
                f.write(json.dumps(alert) + "\n")
        return _handler


# ══════════════════════════════════════════════════════════════════════════════
# REAL-TIME STREAM PROCESSOR
# ══════════════════════════════════════════════════════════════════════════════
class RealTimeStream:
    def __init__(self, pipeline: "SceneSolverPipeline"):
        self.pipeline  = pipeline
        self._thread   = None
        self._stop_evt = threading.Event()

    def start(self, source: str, window_frames: int = 50, on_alert: Any = None):
        self._stop_evt.clear()
        self._thread = threading.Thread(
            target=self._run, args=(source, window_frames, on_alert),
            daemon=True, name="RealTimeStream")
        self._thread.start()
        log.info(f"[RealTimeStream] Started: {source}")

    def stop(self):
        self._stop_evt.set()
        if self._thread:
            self._thread.join(timeout=10)
        log.info("[RealTimeStream] Stopped.")

    def _run(self, source: str, window_frames: int, on_alert: Any):
        cfg      = self.pipeline.config
        ae_stage = self.pipeline.stage0
        ae_stage.load_model()
        alert_sys = self.pipeline.alert_system
        transform = ae_stage.transform
        cap       = cv2.VideoCapture(int(source) if source.isdigit() else source)
        if not cap.isOpened():
            log.error(f"[RealTimeStream] Cannot open: {source}")
            return

        buffer: List[np.ndarray] = []
        scores: List[float]      = []
        frame_count              = 0

        while not self._stop_evt.is_set():
            ret, frame = cap.read()
            if not ret:
                break
            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            buffer.append(frame_rgb)
            frame_count += 1

            if frame_count % cfg.AE_FRAME_STRIDE == 0 and ae_stage.model:
                tensor = transform(Image.fromarray(frame_rgb)).unsqueeze(0).to(cfg.DEVICE)
                with torch.no_grad():
                    recon = ae_stage.model(tensor)
                scores.append(calculate_enhanced_score(tensor, recon))
                del tensor, recon

            if len(buffer) >= window_frames:
                threshold = compute_dynamic_threshold(scores)
                max_err   = max(scores) if scores else 0.0
                if max_err > threshold:
                    risk = float(np.clip(max_err / max(threshold, 1e-8), 0, 1))
                    mini_results = {
                        "stage0": {"max_error": max_err},
                        "stage6": {"fused_risk": risk},
                        "stage7": {"final_risk": risk, "status": "STREAM"},
                        "stage3": {}, "stage4": {},
                        "threshold": threshold, "video_path": source, "output_directory": "",
                    }
                    alert = alert_sys.evaluate_and_fire(mini_results)
                    if alert and on_alert:
                        try:
                            on_alert(alert)
                        except Exception as e:
                            log.error(f"[RealTimeStream] on_alert failed: {e}")
                half   = window_frames // 2
                buffer = buffer[-half:]
                scores = scores[-(half // cfg.AE_FRAME_STRIDE):]

        cap.release()
        log.info("[RealTimeStream] Stream ended.")


# ══════════════════════════════════════════════════════════════════════════════
# PDF REPORT GENERATOR
# ══════════════════════════════════════════════════════════════════════════════
class EnhancedReportGenerator:
    """
    PDF Layout (up to 11 pages):
      Page  1 : Cover + Executive Summary
      Page  2 : Stage 0 (AE) temporal signal graph
      Page  3 : Stage 1 (TF Binary)
      Page  4 : Stage 2 (TF Multi-class)
      Page  5 : Stage 3 (Weapon Detection, EXPERIMENTAL)
      Page  6 : Stage 4 (Audio)
      Page  7 : Stage 5 (Tracking)
      Page  8 : Stage 6 (Fusion)
      Page  9 : Stage 7 (RL3033, EXPERIMENTAL)
      Page 10 : Top anomalous frames
      Page 11 : Stage 8 (LLaVA, EXPERIMENTAL)
    """

    DISCLAIMER = (
        "This report is produced by machine learning models and should be treated "
        "as an investigative aid, not a definitive conclusion. Human expert review "
        "is strongly recommended before any operational action is taken."
    )

    # Color palette: pitch black background, white text, orange/saffron accents
    C_BG     = "#000000"
    C_DARK   = "#1a252f"
    C_HEADER = "#2c3e50"
    C_RED    = "#c0392b"
    C_GREEN  = "#1e8449"
    C_BLUE   = "#2471a3"
    C_ORANGE = "#d35400"
    C_PURPLE = "#6c3483"
    C_TEAL   = "#0e6655"
    C_GREY   = "#5d6d7e"
    C_LIGHT  = "#f2f3f4"

    HEADER_H = 52
    MARGIN_X = 48

    @staticmethod
    def _body_top(H: float) -> float:
        return H - EnhancedReportGenerator.HEADER_H - 16

    @staticmethod
    def _page_header(c, title: str, bg_hex: str, W: float = letter[0], H: float = letter[1]):
        c.setFillColor(colors.HexColor(bg_hex))
        c.rect(0, H - EnhancedReportGenerator.HEADER_H, W,
               EnhancedReportGenerator.HEADER_H, fill=True, stroke=False)
        c.setFillColor(colors.white)
        c.setFont("Helvetica-Bold", 15)
        c.drawCentredString(W / 2, H - EnhancedReportGenerator.HEADER_H + 16, title)

    @staticmethod
    def _footer(c, ts: str, W: float = letter[0]):
        c.setStrokeColor(colors.HexColor("#aab7b8"))
        c.setLineWidth(0.5)
        c.line(40, 28, W - 40, 28)
        c.setFont("Helvetica-Oblique", 7)
        c.setFillColor(colors.HexColor("#7f8c8d"))
        c.drawCentredString(W / 2, 14,
            f"Scene Solver v3  |  {ts}  |  Confidential, Investigative Use Only")

    @staticmethod
    def _section_title(c, text: str, y: float, x: float = None,
                       color_hex: str = "#2c3e50", W: float = letter[0]) -> float:
        if x is None:
            x = EnhancedReportGenerator.MARGIN_X
        c.setFillColor(colors.HexColor(color_hex))
        c.rect(x, y - 2, 4, 14, fill=True, stroke=False)
        c.setFont("Helvetica-Bold", 11)
        c.drawString(x + 8, y, text)
        c.setFillColor(colors.black)
        return y - 20

    @staticmethod
    def _kv_row(c, key: str, value: str, x: float, y: float,
                key_w: float = 170, line_h: float = 14, font_size: int = 9) -> float:
        c.setFont("Helvetica-Bold", font_size)
        c.setFillColor(colors.HexColor("#2c3e50"))
        c.drawString(x, y, key + ":")
        c.setFont("Helvetica", font_size)
        c.setFillColor(colors.black)
        c.drawString(x + key_w, y, value)
        return y - line_h

    @staticmethod
    def _wrap_text(c, text: str, x: float, y: float, max_width: float,
                   font: str = "Helvetica", size: int = 9,
                   line_h: int = 12, min_y: float = 50) -> float:
        c.setFont(font, size)
        c.setFillColor(colors.black)
        words, line = text.split(), ""
        for w in words:
            test = line + w + " "
            if c.stringWidth(test, font, size) < max_width:
                line = test
            else:
                if y < min_y:
                    break
                c.drawString(x, y, line.rstrip())
                y   -= line_h
                line = w + " "
        if line.strip() and y >= min_y:
            c.drawString(x, y, line.rstrip())
            y -= line_h
        return y

    @staticmethod
    def _disclaimer_box(c, y: float, W: float = letter[0]) -> float:
        MX    = EnhancedReportGenerator.MARGIN_X
        box_w = W - 2 * MX
        c.setFillColor(colors.HexColor("#fdfefe"))
        c.setStrokeColor(colors.HexColor("#aab7b8"))
        c.setLineWidth(0.6)
        c.roundRect(MX, y - 28, box_w, 32, 4, fill=True, stroke=True)
        c.setFont("Helvetica-Oblique", 7.5)
        c.setFillColor(colors.HexColor("#626567"))
        EnhancedReportGenerator._wrap_text(
            c, EnhancedReportGenerator.DISCLAIMER,
            MX + 6, y - 9, box_w - 12, "Helvetica-Oblique", 7.5, 10)
        return y - 40

    @staticmethod
    def _status_badge(c, text: str, y: float, color_hex: str,
                      W: float = letter[0], MX: float = None) -> float:
        if MX is None:
            MX = EnhancedReportGenerator.MARGIN_X
        bw, bh, br = W - 2 * MX, 38, 8
        c.setFillColor(colors.HexColor(color_hex))
        c.roundRect(MX, y - bh, bw, bh, br, fill=True, stroke=False)
        c.setFillColor(colors.white)
        c.setFont("Helvetica-Bold", 14)
        c.drawCentredString(W / 2, y - bh + 12, text)
        return y - bh - 10

    @staticmethod
    def _experimental_warning_box(c, text: str, y: float,
                                   W: float = letter[0], MX: float = None) -> float:
        """Dynamic-height experimental disclaimer box. Text flows as one continuous paragraph."""
        if MX is None:
            MX = EnhancedReportGenerator.MARGIN_X
        box_w = W - 2 * MX
        font, size, line_h = "Helvetica-Bold", 8.5, 12
        normalised_text = " ".join(text.split())

        # Pre-measure line count
        import io
        from reportlab.pdfgen.canvas import Canvas as _TmpCanvas
        tmp_c = _TmpCanvas(io.BytesIO())
        tmp_c.setFont(font, size)
        words, line_buf, lines = normalised_text.split(), "", []
        for w in words:
            test = line_buf + w + " "
            if tmp_c.stringWidth(test, font, size) < box_w - 16:
                line_buf = test
            else:
                lines.append(line_buf.rstrip())
                line_buf = w + " "
        if line_buf.strip():
            lines.append(line_buf.rstrip())
        box_h = max(len(lines), 1) * line_h + 16

        c.setFillColor(colors.HexColor("#fef9e7"))
        c.setStrokeColor(colors.HexColor("#d35400"))
        c.setLineWidth(1.2)
        c.roundRect(MX, y - box_h, box_w, box_h, 5, fill=True, stroke=True)
        c.setFont(font, size)
        c.setFillColor(colors.HexColor("#d35400"))
        text_y = y - 8 - size
        for ln in lines:
            c.drawString(MX + 8, text_y, ln)
            text_y -= line_h
        c.setFillColor(colors.black)
        return y - box_h - 8

    # ── Graph generators ───────────────────────────────────────────────────
    @staticmethod
    def generate_deviation_graph(scores, threshold, peak_score_indices,
                                  sampled_indices, fps,
                                  output_path="/tmp/deviation_graph.png"):
        ts_labels = [frame_to_timestamp(
                         sampled_indices[i] if i < len(sampled_indices) else i * 5, fps)
                     for i in range(len(scores))]
        x = list(range(len(scores)))
        fig, ax = plt.subplots(figsize=(13, 4.2))
        ax.plot(x, scores, linewidth=2.2, color="#c0392b", label="Reconstruction Error", zorder=2)
        ax.axhline(y=threshold, color="#1e8449", linestyle="--",
                   label=f"Dynamic Threshold ({threshold:.5f})", linewidth=1.8, zorder=1)
        valid_peaks = [i for i in peak_score_indices if i < len(scores)]
        if valid_peaks:
            ax.scatter(valid_peaks, [scores[i] for i in valid_peaks],
                       color="#f39c12", s=220, marker="*",
                       edgecolors="#333", linewidths=1.2, label="Peak Anomaly", zorder=3)
        ax.fill_between(x, scores, threshold,
                        where=[s > threshold for s in scores], color="#c0392b", alpha=0.15)
        tick_step = max(1, len(scores) // 8)
        tick_pos  = list(range(0, len(scores), tick_step))
        ax.set_xticks(tick_pos)
        ax.set_xticklabels([ts_labels[i] for i in tick_pos], rotation=25, fontsize=7.5, ha="right")
        ax.set_xlabel("Timestamp (MM:SS.mmm)", fontsize=10, fontweight="bold")
        ax.set_ylabel("Reconstruction Error",  fontsize=10, fontweight="bold")
        ax.set_title("Autoencoder Temporal Reconstruction Error", fontsize=12, fontweight="bold", pad=10)
        ax.legend(loc="upper right", fontsize=9, framealpha=0.92)
        ax.grid(True, alpha=0.25, linestyle=":", linewidth=0.8)
        fig.tight_layout()
        fig.savefig(output_path, dpi=150, bbox_inches="tight", facecolor="white")
        plt.close(fig)
        return output_path

    @staticmethod
    def generate_classification_pie(probabilities, output_path="/tmp/class_pie.png"):
        """Pie chart with numbered wedges only. Percentages appear in the legend panel."""
        items  = sorted(probabilities.items(), key=lambda x: x[1], reverse=True)
        labels = [i[0] for i in items]
        sizes  = [max(i[1], 1e-6) for i in items]
        clist  = ["#c0392b","#2471a3","#d35400","#6c3483","#117a65","#1a5276","#7d6608"]
        fig, (ax_pie, ax_leg) = plt.subplots(1, 2, figsize=(10, 5),
                                              gridspec_kw={"width_ratios": [1.3, 1]})
        ax_pie.pie(
            sizes,
            labels=[str(i + 1) for i in range(len(sizes))],
            colors=clist[:len(sizes)], startangle=90,
            textprops={"fontsize": 10, "weight": "bold"},
            explode=[0.04 if i == 0 else 0 for i in range(len(sizes))],
            labeldistance=1.10,
        )
        ax_pie.set_title("Anomaly Class Distribution\n(see legend for percentages)",
                         fontsize=10, fontweight="bold", pad=8)
        ax_leg.axis("off")
        for i, (cls, prob) in enumerate(zip(labels, sizes)):
            ax_leg.add_patch(plt.Rectangle((0, 0.85 - i * 0.12), 0.12, 0.09,
                                            color=clist[i % len(clist)]))
            ax_leg.text(0.16, 0.855 - i * 0.12 + 0.01,
                        f"{i+1}. {cls}  ({prob*100:.1f}%)",
                        va="center", fontsize=8.5, fontweight="bold",
                        transform=ax_leg.transAxes)
        ax_leg.set_xlim(0, 1); ax_leg.set_ylim(0, 1)
        fig.tight_layout()
        fig.savefig(output_path, dpi=150, bbox_inches="tight", facecolor="white")
        plt.close(fig)
        return output_path

    @staticmethod
    def generate_confidence_bars(top3, output_path="/tmp/conf_bars.png"):
        if not top3:
            return None
        fig, ax = plt.subplots(figsize=(7, 3.2))
        classes = [t[0] for t in top3]
        confs   = [t[1] * 100 for t in top3]
        bars    = ax.barh(classes, confs,
                          color=["#c0392b","#d35400","#5d6d7e"][:len(top3)],
                          edgecolor="#2c3e50", linewidth=0.8, height=0.5)
        for bar, cv in zip(bars, confs):
            ax.text(bar.get_width() + 0.8, bar.get_y() + bar.get_height() / 2,
                    f"{cv:.1f}%", va="center", fontsize=9, fontweight="bold", color="#2c3e50")
        ax.set_xlabel("Confidence (%)", fontsize=9, fontweight="bold")
        ax.set_title("Top-3 Anomaly Predictions", fontsize=10, fontweight="bold", pad=8)
        ax.set_xlim(0, 118)
        ax.grid(axis="x", alpha=0.25, linestyle="--", linewidth=0.7)
        ax.invert_yaxis()
        fig.tight_layout()
        fig.savefig(output_path, dpi=150, bbox_inches="tight", facecolor="white")
        plt.close(fig)
        return output_path

    @staticmethod
    def extract_top_frames(frames, scores, sampled_indices, n=3):
        if not scores or not frames:
            return []
        top = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)
        out = []
        for si in top[:n]:
            fi = sampled_indices[si] if si < len(sampled_indices) else si * 5
            if fi < len(frames):
                out.append({"frame": frames[fi], "index": fi, "score": scores[si]})
        return out

    # ── Main PDF builder ────────────────────────────────────────────────────
    @staticmethod
    def generate_enhanced_pdf(results: dict, output_path: str = "/tmp/forensic_report.pdf") -> str:
        c    = rl_canvas.Canvas(output_path, pagesize=letter)
        W, H = letter
        ts   = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        G    = EnhancedReportGenerator
        MX   = G.MARGIN_X
        BW   = W - 2 * MX  # body width

        s0 = results.get("stage0") or {}
        s1 = results.get("stage1") or {}
        s2 = results.get("stage2") or {}
        s3 = results.get("stage3") or {}
        s4 = results.get("stage4") or {}
        s5 = results.get("stage5") or {}
        s6 = results.get("stage6") or {}
        s7 = results.get("stage7") or {}
        s8 = results.get("stage8") or {}
        al = results.get("alert")

        fps       = float(s0.get("fps", 25.0))
        threshold = results.get("threshold", Config.SAFE_THRESHOLD_LEGACY)
        overall   = results.get("status", "SAFE")
        is_crit   = (overall == "CRITICAL")
        stage1_normal = (s1.get("classification", "Unknown") == "Normal")

        def ts_of(fi):
            return frame_to_timestamp(int(fi), fps)

        # ── Page 1: Cover + Executive Summary ─────────────────────────────
        G._page_header(c, "FORENSIC ANALYSIS DOSSIER, Scene Solver v3", G.C_HEADER, W, H)
        y = G._body_top(H)

        c.setFont("Helvetica", 8); c.setFillColor(colors.HexColor(G.C_GREY))
        c.drawCentredString(W / 2, y, f"Report generated: {ts}")
        y -= 14

        badge_text = "ANOMALY DETECTED (CRITICAL)" if is_crit else "NO ANOMALY DETECTED (SAFE)"
        y = G._status_badge(c, badge_text, y, G.C_RED if is_crit else G.C_GREEN, W, MX)
        y -= 6
        y = G._section_title(c, "EXECUTIVE SUMMARY", y, MX, G.C_HEADER)

        def _field(label, value, key_w=185):
            nonlocal y
            if y < 60: return
            c.setFont("Helvetica-Bold", 9); c.setFillColor(colors.HexColor(G.C_DARK))
            c.drawString(MX + 10, y, label + ":")
            c.setFont("Helvetica", 9); c.setFillColor(colors.black)
            c.drawString(MX + 10 + key_w, y, str(value))
            y -= 13

        def _spacer(h=5):
            nonlocal y; y -= h

        _field("Overall Status",           overall)
        _field("Dynamic Threshold",        f"{threshold:.6f}  (mean plus 2 standard deviations)")
        _field("AE Max Reconstruction Error", f"{s0.get('max_error', 0):.6f}")
        _field("AE Avg Reconstruction Error", f"{s0.get('avg_error', 0):.6f}")
        _field("Peak Event Timestamp",     ts_of(s0.get("peak_frame_idx", 0)))
        _field("Anomaly Window",
               f"{ts_of(s0.get('anomaly_start_frame', 0))}  to  {ts_of(s0.get('anomaly_end_frame', 0))}")
        _spacer()

        if s1:
            _field("Stage 1 (TF Binary)", f"{s1.get('classification','?')}  ({s1.get('confidence',0):.1%})")
        if s2:
            if s2.get("class") == "Skipped":
                _field("Stage 2 (TF Multi)", "Skipped (Stage 1 = Normal)")
            else:
                samp_note = "  [AE-weighted sampling]" if s2.get("sampling_method") == "segmentwise_ae" else ""
                _field("Stage 2 (TF Multi)",
                       f"{s2.get('class','?')}  ({s2.get('confidence',0):.1%}){samp_note}")
        _spacer()

        if s3:
            wf = bool(s3.get("weapon_flag", False))
            _field("Stage 3 (Weapon Detection)",
                   "Weapons detected" if wf else "No weapons detected", key_w=210)
            if wf:
                _field("  Dominant Class",  s3.get("dominant_class", "?"), key_w=210)
                _field("  Best Confidence", f"{s3.get('best_conf', 0):.2%}", key_w=210)
        _spacer()

        if s4:
            audio_ok = s4.get("audio_present", True)
            audio_st = s4.get("status", "?")
            _field("Stage 4 (Audio)",
                   "No audio or silent track, section skipped"
                   if not audio_ok or audio_st in ("NO_AUDIO", "SILENT")
                   else f"{audio_st}  (score = {s4.get('audio_score', 0):.4f})")
        if s5:
            _field("Stage 5 (Tracking)",
                   f"Motion = {s5.get('motion_intensity',0):.3f},  "
                   f"Behaviors: {', '.join(s5.get('behaviors',[])) or 'None'}")
        _spacer()

        if s6:
            _field("Stage 6 (Fusion)",
                   f"{s6.get('fused_risk',0):.1%}  [{s6.get('risk_level','?')}]")
        if s7:
            _field("Stage 7 (RL3033, Experimental)",
                   f"Final risk = {s7.get('final_risk',0):.1%}  [{s7.get('status','?')}]")
        _spacer()

        if al:
            c.setFont("Helvetica-Bold", 9); c.setFillColor(colors.HexColor(G.C_RED))
            c.drawString(MX + 10, y,
                f"ALERT FIRED: {al['level']}  |  Final risk = {al['final_risk']:.1%}")
            c.setFillColor(colors.black); y -= 13

        # Pipeline interpretation paragraph
        _spacer(8)
        y = G._section_title(c, "PIPELINE INTERPRETATION", y, MX, G.C_HEADER)
        risk_val   = float(s7.get("final_risk", s6.get("fused_risk", 0.0)))
        risk_label = s6.get("risk_level", "LOW")
        cls1_label = s1.get("classification", "Unknown")
        cls1_conf  = float(s1.get("confidence", 0.0))
        cls2_label = s2.get("class", "Unknown")
        cls2_conf  = float(s2.get("confidence", 0.0))
        weapon_flag = bool(s3.get("weapon_flag", False))
        audio_ok    = s4.get("audio_present", True)
        audio_st    = s4.get("status", "NO_AUDIO")
        motion_val  = float(s5.get("motion_intensity", 0.0))
        behaviors   = s5.get("behaviors", [])

        interp = [f"The pipeline classified this video as {overall} with a final risk score of {risk_val:.1%} ({risk_label})."]
        if cls1_label == "Suspicious":
            interp.append(f"The semantic binary classifier (Stage 1) identified the clip as Suspicious at {cls1_conf:.0%} confidence.")
            if cls2_label not in ("Unknown", "Skipped", "Error"):
                interp.append(f"Stage 2 classified the event type as {cls2_label} ({cls2_conf:.0%} confidence) using AE-weighted temporal sampling.")
        else:
            interp.append(f"The semantic binary classifier returned Normal ({cls1_conf:.0%} confidence), overriding the AE signal.")
        if weapon_flag:
            interp.append(f"A weapon ({s3.get('dominant_class','?')}) was detected at {s3.get('best_conf', 0.0):.0%} confidence, triggering hard risk escalation.")
        if audio_ok and audio_st not in ("NO_AUDIO", "SILENT", "NORMAL"):
            interp.append(f"Audio analysis reported {audio_st} status (score {s4.get('audio_score',0):.3f}).")
        if "FAST_MOVEMENT" in behaviors:
            interp.append(f"Motion tracking recorded fast subject movement (intensity {motion_val:.2f}).")
        elif "CLUSTERING" in behaviors:
            interp.append(f"Motion tracking detected converging subject trajectories (intensity {motion_val:.2f}).")

        y = G._wrap_text(c, " ".join(interp), MX + 6, y, BW - 12, "Helvetica", 8.5, 12)
        _spacer(8)
        y = G._disclaimer_box(c, y, W)
        G._footer(c, ts, W)

        # ── Page 2: Stage 0 (AE) ──────────────────────────────────────────
        c.showPage()
        G._page_header(c, "STAGE 0 (AUTOENCODER): TEMPORAL ANOMALY DETECTION", G.C_RED, W, H)
        y = G._body_top(H)
        c.setFont("Helvetica-Oblique", 8.5); c.setFillColor(colors.HexColor(G.C_GREY))
        y = G._wrap_text(c,
            "The Autoencoder is trained exclusively on normal footage. High reconstruction error "
            "indicates unusual content. The dynamic threshold is computed as mean plus two standard "
            "deviations of the per-frame error.",
            MX, y, BW, "Helvetica-Oblique", 8.5, 12)
        y -= 8

        gp = results.get("graph_path")
        if gp and os.path.exists(gp):
            img_h = 200
            c.drawImage(gp, MX, y - img_h, width=BW, height=img_h, preserveAspectRatio=True)
            y -= img_h + 10

        if stage1_normal:
            y = G._section_title(c, "AE Score Summary", y, MX, G.C_GREEN)
            y = G._wrap_text(c,
                "Stage 1 classified this video as NORMAL, overriding the AE signal. "
                "Downstream anomaly stages were skipped. Pipeline status is SAFE.",
                MX + 6, y, BW - 12, "Helvetica", 9, 13)
            y -= 8
            for k, v in [
                ("Max Reconstruction Error",  f"{s0.get('max_error',0):.6f}"),
                ("Avg Reconstruction Error",  f"{s0.get('avg_error',0):.6f}"),
                ("Dynamic Threshold",         f"{threshold:.6f}"),
                ("Threshold Exceeded by AE",  "YES" if s0.get('max_error',0) > threshold else "NO"),
                ("Note", "AE signal overridden to SAFE by Stage 1 semantic classifier"),
                ("Total Sampled Frames",      str(len(s0.get("scores", [])))),
                ("Video FPS",                 f"{fps:.1f}"),
                ("Pipeline Status",           "SAFE (Stage 1 override)"),
            ]:
                y = G._kv_row(c, k, v, MX + 6, y, key_w=220)
                if y < 60: break
        else:
            y = G._section_title(c, "AE Score Summary", y, MX, G.C_RED)
            ae_over = s0.get('max_error', 0) > threshold
            for k, v in [
                ("Max Reconstruction Error",  f"{s0.get('max_error',0):.6f}"),
                ("Avg Reconstruction Error",  f"{s0.get('avg_error',0):.6f}"),
                ("Dynamic Threshold",         f"{threshold:.6f}  (mean + 2 std dev, floor = {Config.SAFE_THRESHOLD_FLOOR})"),
                ("Threshold Exceeded",        "YES" if ae_over else "NO"),
                ("Excess Above Threshold",    f"{s0.get('max_error',0) - threshold:.6f}" if ae_over else "N/A"),
                ("Peak Event Timestamp",      ts_of(s0.get("peak_frame_idx", 0))),
                ("Anomaly Start Timestamp",   ts_of(s0.get("anomaly_start_frame", 0))),
                ("Anomaly End Timestamp",     ts_of(s0.get("anomaly_end_frame", 0))),
                ("Anomaly Duration (frames)", str(s0.get("anomaly_end_frame",0) - s0.get("anomaly_start_frame",0))),
                ("Total Sampled Frames",      str(len(s0.get("scores", [])))),
                ("Sampling Stride",           f"Every {Config.AE_FRAME_STRIDE} frames"),
                ("Video FPS",                 f"{fps:.1f}"),
            ]:
                y = G._kv_row(c, k, v, MX + 6, y, key_w=220)
                if y < 60: break
            if y > 100:
                y -= 6
                y = G._section_title(c, "Signal Interpretation", y, MX, G.C_RED)
                ratio = s0.get('max_error', 0) / max(threshold, 1e-8)
                if ratio >= 2.0:
                    ae_interp = f"Peak reconstruction error is {ratio:.1f}x the threshold — strong, unambiguous anomaly signal."
                elif ratio >= 1.0:
                    ae_interp = f"Peak exceeds threshold by {ratio:.2f}x. Meaningful signal; benefits from corroboration."
                else:
                    ae_interp = f"Error remains below threshold (ratio {ratio:.3f}). AE alone does not flag anomaly."
                y = G._wrap_text(c, ae_interp, MX + 6, y, BW - 12, "Helvetica", 9, 13)

        y -= 6; y = G._disclaimer_box(c, y, W); G._footer(c, ts, W)

        # ── Page 3: Stage 1 (TF Binary) ───────────────────────────────────
        c.showPage()
        G._page_header(c, "STAGE 1 (TIMESFORMER BINARY): SEMANTIC CLASSIFICATION", G.C_BLUE, W, H)
        y = G._body_top(H)
        c.setFont("Helvetica-Oblique", 8.5); c.setFillColor(colors.HexColor(G.C_GREY))
        y = G._wrap_text(c,
            "The binary TimeSformer classifies a 32-frame clip as Normal or Suspicious. "
            "A Normal result overrides the AE signal to SAFE and skips Stages 2 and 3. "
            "LLaVA (Stage 8) still runs regardless, using neutral investigative prompts.",
            MX, y, BW, "Helvetica-Oblique", 8.5, 12)
        y -= 12

        cls1, conf1 = s1.get("classification","?"), s1.get("confidence",0.0)
        y = G._status_badge(c, f"{cls1.upper()}  ({conf1:.1%})", y,
                            G.C_RED if cls1 == "Suspicious" else G.C_GREEN, W, MX)
        y -= 10
        y = G._section_title(c, "Classification Details", y, MX, G.C_BLUE)
        for k, v in [
            ("Result",             cls1),
            ("Confidence",         f"{conf1:.2%}"),
            ("Clip Centre",        ts_of(s0.get("peak_frame_idx", 0))),
            ("Clip Window",        "Peak timestamp, plus or minus 16 frames (32 frames total)"),
            ("Gates YOLO Stage 3", "YES" if cls1 == "Suspicious" else "NO, skipped"),
            ("Gates Stage 2",      "YES" if cls1 == "Suspicious" else "NO, skipped (Stage 1 = Normal)"),
            ("LLaVA Stage 8",      "Runs with full context" if cls1 == "Suspicious"
                                   else "Runs with neutral investigative prompts"),
            ("Overall Status Effect", "Follows AE signal" if cls1 == "Suspicious"
                                      else "Overrides AE — pipeline forced to SAFE"),
        ]:
            y = G._kv_row(c, k, v, MX + 6, y, key_w=210)
        y -= 8
        y = G._section_title(c, "Interpretation and Operational Significance", y, MX, G.C_BLUE)
        if cls1 == "Suspicious":
            interp = (f"Clip at {ts_of(s0.get('peak_frame_idx',0))} classified SUSPICIOUS at {conf1:.1%} confidence. "
                      "Weapon detection and multi-class classification are now active.")
        else:
            interp = (f"Clip classified NORMAL at {conf1:.1%} confidence. "
                      "The AE signal is overridden. Multi-class classification is skipped. "
                      "LLaVA runs with neutral prompts.")
        y = G._wrap_text(c, interp, MX + 6, y, BW - 12, "Helvetica", 9, 13)
        y -= 6; y = G._disclaimer_box(c, y, W); G._footer(c, ts, W)

        # ── Page 4: Stage 2 (TF Multi-class) ──────────────────────────────
        c.showPage()
        G._page_header(c, "STAGE 2 (TIMESFORMER MULTI-CLASS): ANOMALY TYPE", G.C_PURPLE, W, H)
        y = G._body_top(H)
        cls2, conf2 = s2.get("class","Unknown"), s2.get("confidence",0.0)
        samp_meth   = s2.get("sampling_method","?")
        stage2_skipped = (cls2 == "Skipped" or not s2.get("all_probabilities"))

        c.setFont("Helvetica-Oblique", 8.5); c.setFillColor(colors.HexColor(G.C_GREY))
        y = G._wrap_text(c,
            "Multi-class TimeSformer classifies the anomaly into one of seven crime categories. "
            "Only runs when Stage 1 returns Suspicious. Uses segment-wise AE-weighted temporal "
            "sampling: the video is split into 16 equal segments; the highest-error frame per "
            "segment is selected.",
            MX, y, BW, "Helvetica-Oblique", 8.5, 12)
        y -= 8

        if stage2_skipped:
            y = G._status_badge(c, "STAGE 2 SKIPPED (STAGE 1 = NORMAL)", y, G.C_GREY, W, MX)
            y -= 8
            y = G._section_title(c, "Why This Stage Was Skipped", y, MX, G.C_GREY)
            y = G._wrap_text(c,
                "Stage 1 returned Normal, which suppresses all anomaly-specific downstream "
                "processing. Running multi-class classification on a Normal-labelled clip "
                "would produce unreliable category assignments.",
                MX + 6, y, BW - 12, "Helvetica", 9, 13)
        else:
            y = G._status_badge(c, f"{cls2.upper()}  ({conf2:.1%})", y, G.C_PURPLE, W, MX)
            y -= 4
            c.setFont("Helvetica-Oblique", 8); c.setFillColor(colors.HexColor(G.C_GREY))
            samp_label = "Segment-wise AE-weighted (16 segments)" if samp_meth == "segmentwise_ae" else samp_meth
            c.drawCentredString(W / 2, y, f"Sampling: {samp_label}")
            y -= 10
            c.setFont("Helvetica-Oblique", 8); c.setFillColor(colors.HexColor(G.C_PURPLE))
            c.drawCentredString(W / 2, y, "Chart shows class numbers only. Refer to numbered legend for names and percentages.")
            y -= 10

            all_probs = s2.get("all_probabilities", {})
            top3      = s2.get("top3", [])
            if all_probs:
                pie_path = "/tmp/cls_pie_v3.png"
                G.generate_classification_pie(all_probs, pie_path)
                if os.path.exists(pie_path):
                    c.drawImage(pie_path, MX, y - 200, width=310, height=200, preserveAspectRatio=True)
            if top3:
                bar_path = "/tmp/conf_bars_v3.png"
                G.generate_confidence_bars(top3, bar_path)
                if os.path.exists(bar_path):
                    c.drawImage(bar_path, MX + 318, y - 150, width=220, height=150, preserveAspectRatio=True)
            y -= 212
            y = G._section_title(c, "Top-3 Predictions", y, MX, G.C_PURPLE)
            for rank, (cls_name, prob) in enumerate(top3[:3], 1):
                c.setFont("Helvetica", 9); c.setFillColor(colors.black)
                c.drawString(MX + 12, y, f"{rank}.  {cls_name}  ({prob:.2%})")
                y -= 13
            y -= 4
            y = G._section_title(c, "All Class Probabilities", y, MX, G.C_PURPLE)
            for cls_name, prob in sorted(all_probs.items(), key=lambda x: x[1], reverse=True):
                c.setFont("Helvetica", 8.5); c.setFillColor(colors.black)
                c.drawString(MX + 12, y, f"{cls_name:<22s}  {prob:.4f}  ({prob * 100:.2f}%)")
                y -= 12
                if y < 70: break
            if y > 100 and top3:
                y -= 6
                y = G._section_title(c, "Classification Insight", y, MX, G.C_PURPLE)
                top_cls, top_conf = top3[0]
                if top_conf >= 0.60:
                    insight = f"Model assigned highest probability to '{top_cls}' ({top_conf:.1%}) — high-confidence prediction."
                elif top_conf >= 0.35:
                    insight = f"Model assigned highest probability to '{top_cls}' ({top_conf:.1%}) — moderate confidence. Consider second and third candidates."
                else:
                    insight = f"Model assigned highest probability to '{top_cls}' ({top_conf:.1%}) — low confidence. Anomaly pattern may not cleanly match any trained category."
                y = G._wrap_text(c, insight, MX + 6, y, BW - 12, "Helvetica", 9, 13)

        y -= 4; y = G._disclaimer_box(c, y, W); G._footer(c, ts, W)

        # ── Page 5: Stage 3 (Weapon Detection) ────────────────────────────
        c.showPage()
        weapon_flag = bool(s3.get("weapon_flag", False))
        hdr_col3    = G.C_RED if weapon_flag else G.C_TEAL
        G._page_header(c, "STAGE 3 (WEAPON DETECTION, EXPERIMENTAL)", hdr_col3, W, H)
        y = G._body_top(H)
        y = G._experimental_warning_box(c,
            "EXPERIMENTAL: This module is under active development. "
            "Detection results should be independently verified before any action is taken.",
            y, W, MX)
        c.setFont("Helvetica-Oblique", 8.5); c.setFillColor(colors.HexColor(G.C_GREY))
        y = G._wrap_text(c,
            "A custom-trained YOLOv8 model scans frames within the anomaly window for Melee "
            "weapons and Firearms. Only activates when Stage 1 returns Suspicious. A confirmed "
            "weapon detection triggers a hard escalation in Stage 6 (minimum risk 0.90 when "
            "combined with an AE anomaly). Up to 32 frames are scanned.",
            MX, y, BW, "Helvetica-Oblique", 8.5, 12)
        y -= 8

        badge_txt = f"WEAPONS DETECTED ({s3.get('dominant_class','?').upper()})" if weapon_flag else "NO WEAPONS DETECTED"
        y = G._status_badge(c, badge_txt, y, G.C_RED if weapon_flag else G.C_TEAL, W, MX)
        y -= 8

        skipped = str(s3.get("summary", "")).lower().startswith("skipped")
        if skipped:
            c.setFont("Helvetica-Oblique", 9); c.setFillColor(colors.HexColor(G.C_GREY))
            c.drawString(MX + 6, y, "Stage skipped — Stage 1 classified the video as Normal.")
            y -= 20
            y = G._section_title(c, "Operational Note", y, MX, G.C_GREY)
            y = G._wrap_text(c,
                "Weapon scanning is gated behind Stage 1 to reduce false positives. "
                "When Stage 1 returns Suspicious, weapon scanning activates for the anomaly window.",
                MX + 6, y, BW - 12, "Helvetica", 9, 13)
        else:
            y = G._section_title(c, "Detection Summary", y, MX, hdr_col3)
            for k, v in [
                ("Weapon Detected",          "YES" if weapon_flag else "NO"),
                ("Dominant Weapon Class",    s3.get("dominant_class", "None")),
                ("Total Detections",         str(s3.get("weapon_count", 0))),
                ("Best Detection Confidence",
                 f"{s3.get('best_conf', 0.0):.2%}" if weapon_flag else "N/A"),
                ("Frames Scanned",           "Up to 32 frames within the anomaly window"),
                ("Anomaly Window",
                 f"{ts_of(s0.get('anomaly_start_frame',0))} to {ts_of(s0.get('anomaly_end_frame',0))}"),
                ("Fusion Impact",
                 "Risk forced to minimum 0.90 (weapon + AE threshold exceeded)"
                 if weapon_flag and s0.get("max_error",0) > threshold
                 else "Weapon floor applied (0.40 minimum)" if weapon_flag
                 else "No fusion impact from weapon stage"),
            ]:
                y = G._kv_row(c, k, v, MX + 6, y, key_w=210)
            y -= 6
            if weapon_flag:
                y = G._section_title(c, "Individual Detections (top 10)", y, MX, hdr_col3)
                for di, det in enumerate(s3.get("all_detections", [])[:10], 1):
                    c.setFont("Helvetica", 8.5); c.setFillColor(colors.black)
                    c.drawString(MX + 8, y,
                        f"  {di:2d}.  {det.get('class_name','?'):12s}  "
                        f"Confidence = {det.get('confidence',0):.2%},  at {ts_of(det.get('frame_idx',0))}")
                    y -= 12
                    if y < 130: break
                top_frames = s3.get("top_frames_annotated", [])
                if top_frames and y > 80:
                    y -= 4
                    y = G._section_title(c, "Top Detection Frames (with bounding boxes)", y, MX, hdr_col3)
                    n_show = min(2, len(top_frames))
                    img_w  = (BW - 10) / max(n_show, 1)
                    img_h  = min(140, max(y - 60, 60), img_w * 0.6)
                    for fi_idx, (fi, ann_frame) in enumerate(top_frames[:n_show]):
                        tmp_path = f"/tmp/yolo_ann_{fi_idx}.jpg"
                        Image.fromarray(ann_frame).save(tmp_path)
                        cx = MX + fi_idx * (img_w + 10)
                        c.drawImage(tmp_path, cx, y - img_h, width=img_w - 4,
                                    height=img_h, preserveAspectRatio=True)
                        c.setFont("Helvetica", 7.5); c.setFillColor(colors.HexColor(G.C_DARK))
                        c.drawString(cx, y - img_h - 10, f"Timestamp: {ts_of(fi)}")
                    y -= img_h + 22

        y -= 4; y = G._disclaimer_box(c, y, W); G._footer(c, ts, W)

        # ── Page 6: Stage 4 (Audio) ────────────────────────────────────────
        audio_present = s4.get("audio_present", True)
        audio_status  = s4.get("status", "NO_AUDIO")
        audio_silent  = (not audio_present or audio_status in ("NO_AUDIO", "SILENT"))
        c.showPage()
        aud_col = (G.C_RED    if audio_status == "ANOMALY"    else
                   G.C_ORANGE if audio_status == "SUSPICIOUS" else
                   G.C_GREEN  if audio_status == "NORMAL"     else G.C_GREY)
        G._page_header(c, "STAGE 4 (AUDIO FEATURE EXTRACTION)", aud_col, W, H)
        y = G._body_top(H)
        c.setFont("Helvetica-Oblique", 8.5); c.setFillColor(colors.HexColor(G.C_GREY))
        y = G._wrap_text(c,
            "Five acoustic features are extracted and normalised to 0-1 using global range constants. "
            "Classification rule: 2+ features above 0.4 = ANOMALY; 1 = SUSPICIOUS; 0 = NORMAL. "
            "The composite score feeds into Stages 6 and 7.",
            MX, y, BW, "Helvetica-Oblique", 8.5, 12)
        y -= 6

        y = G._section_title(c, "What Each Feature Measures", y, MX, aud_col)
        for fname, fdesc in [
            ("Zero Crossing Rate (ZCR)",
             "How often the audio waveform crosses zero per second. High values indicate noisy or "
             "chaotic sounds (shouting, breaking glass, gunshots)."),
            ("MFCC Variance",
             "Captures tonal texture change over time. High variance means rapidly shifting audio "
             "content — screaming, rapid voice changes, or sudden silence."),
            ("Spectral Flux",
             "Measures how quickly frequency content changes. Sharp spikes indicate explosive impacts, "
             "gunshots, or sudden loud events."),
            ("RMS Energy",
             "Overall loudness. A sudden surge indicates a loud, potentially alarming sound event."),
            ("Spectral Centroid",
             "Centre of gravity of the frequency spectrum. Dramatic shifts can signal alarms, "
             "screaming, or high-pitched metal impacts."),
        ]:
            if y < 90: break
            c.setFont("Helvetica-Bold", 8.5); c.setFillColor(colors.HexColor(G.C_DARK))
            c.drawString(MX + 8, y, fname)
            y -= 12
            y = G._wrap_text(c, fdesc, MX + 16, y, BW - 24, "Helvetica", 8, 11)
            y -= 4

        y -= 4
        if audio_silent:
            c.setFillColor(colors.HexColor("#eaeded"))
            c.roundRect(MX, y - 44, BW, 48, 6, fill=True, stroke=False)
            c.setFillColor(colors.HexColor(G.C_GREY))
            c.setFont("Helvetica-Bold", 12)
            c.drawCentredString(W / 2, y - 20, "NO AUDIO TRACK OR SILENT TRACK")
            c.setFont("Helvetica", 9)
            c.drawCentredString(W / 2, y - 34, "The video has no audio or an entirely silent track.")
            y -= 60
        else:
            y = G._status_badge(c, f"AUDIO STATUS: {audio_status}", y, aud_col, W, MX)
            y -= 8
            y = G._section_title(c, "Raw Feature Values", y, MX, aud_col)
            for k, v in [
                ("Audio Composite Score",    f"{s4.get('audio_score',0):.4f}"),
                ("Zero Crossing Rate (ZCR)", f"{s4.get('zcr',0):.6f}"),
                ("MFCC Variance",            f"{s4.get('mfcc_variance',0):.4f}"),
                ("Spectral Flux",            f"{s4.get('spectral_flux',0):.4f}"),
                ("RMS Energy",               f"{s4.get('rms_energy',0):.6f}"),
                ("Spectral Centroid",        f"{s4.get('spectral_centroid',0):.2f} Hz"),
            ]:
                y = G._kv_row(c, k, v, MX + 6, y, key_w=200)
            norm = s4.get("normalized", {})
            if norm:
                y -= 4
                y = G._section_title(c, "Globally Normalised Values (threshold = 0.4)", y, MX, aud_col)
                flagged = []
                for k, v, key in [
                    ("ZCR (normalised)",               f"{norm.get('zcr',0):.4f}",  "zcr"),
                    ("MFCC Variance (normalised)",     f"{norm.get('mfcc',0):.4f}", "mfcc"),
                    ("Spectral Flux (normalised)",     f"{norm.get('flux',0):.4f}", "flux"),
                    ("RMS Energy (normalised)",        f"{norm.get('rms',0):.4f}",  "rms"),
                    ("Spectral Centroid (normalised)", f"{norm.get('sc',0):.4f}",   "sc"),
                ]:
                    flag = "  [ABOVE THRESHOLD]" if norm.get(key, 0) > 0.4 else ""
                    if flag: flagged.append(k.split(" (")[0])
                    y = G._kv_row(c, k, v + flag, MX + 6, y, key_w=230)
                if flagged and y > 90:
                    y -= 4
                    y = G._wrap_text(c,
                        f"Features above threshold: {', '.join(flagged)}. "
                        "Rule: 2+ features above 0.4 = ANOMALY, 1 = SUSPICIOUS, 0 = NORMAL.",
                        MX + 6, y, BW - 12, "Helvetica-Oblique", 8, 11)

        y -= 6; y = G._disclaimer_box(c, y, W); G._footer(c, ts, W)

        # ── Page 7: Stage 5 (Tracking) ────────────────────────────────────
        c.showPage()
        G._page_header(c, "STAGE 5 (TRACKING AND MOTION BEHAVIOR)", G.C_TEAL, W, H)
        y = G._body_top(H)
        c.setFont("Helvetica-Oblique", 8.5); c.setFillColor(colors.HexColor(G.C_GREY))
        y = G._wrap_text(c,
            "Multi-object tracking records pixel trajectories in the anomaly window. "
            "Three patterns are detected: FAST_MOVEMENT, STATIC, and CLUSTERING. "
            "YOLOv8 tracking is used when available; Farneback optical flow is the fallback.",
            MX, y, BW, "Helvetica-Oblique", 8.5, 12)
        y -= 10

        motion_int   = float(s5.get("motion_intensity", 0.0))
        behaviors    = s5.get("behaviors", [])
        track_count  = s5.get("track_count", 0)
        track_status = s5.get("status", "?")

        bar_col = G.C_RED if motion_int > 0.6 else G.C_ORANGE if motion_int > 0.3 else G.C_GREEN
        c.setFillColor(colors.HexColor("#d5d8dc"))
        c.rect(MX, y - 16, BW, 18, fill=True, stroke=False)
        c.setFillColor(colors.HexColor(bar_col))
        c.rect(MX, y - 16, int(BW * min(motion_int, 1.0)), 18, fill=True, stroke=False)
        c.setFillColor(colors.white if motion_int > 0.15 else colors.HexColor(G.C_DARK))
        c.setFont("Helvetica-Bold", 9)
        c.drawString(MX + 6, y - 9, f"Motion Intensity: {motion_int:.3f}  ({motion_int*100:.1f}%)")
        y -= 28

        y = G._section_title(c, "Tracking Summary", y, MX, G.C_TEAL)
        for k, v in [
            ("Tracking Method",     track_status),
            ("Tracked Objects",     str(track_count) if track_count else "N/A (optical flow)"),
            ("Motion Intensity",    f"{motion_int:.4f}  (0 = static, 1 = maximum displacement)"),
            ("Detected Behaviours", ", ".join(behaviors) if behaviors else "None"),
            ("Fusion Contribution",
             "FAST_MOVEMENT or CLUSTERING applies a motion boost to fused risk"
             if any(b in behaviors for b in ("FAST_MOVEMENT","CLUSTERING"))
             else "No behavioural boost applied"),
        ]:
            y = G._kv_row(c, k, v, MX + 6, y, key_w=210)
        y -= 8

        y = G._section_title(c, "Behaviour Definitions", y, MX, G.C_TEAL)
        for bname, bdesc in [
            ("FAST_MOVEMENT",
             "Average per-frame displacement exceeds 8 pixels. Subjects are running, lunging, or moving rapidly."),
            ("STATIC",
             "Average displacement below 2 pixels over at least 10 frames. Subjects are essentially stationary."),
            ("CLUSTERING",
             "Two subjects converge: distance decreases in over 75% of frames by at least 50 pixels. "
             "Associated with confrontational proximity or pursuit."),
        ]:
            c.setFont("Helvetica-Bold", 9); c.setFillColor(colors.HexColor(G.C_DARK))
            c.drawString(MX + 8, y, bname); y -= 12
            y = G._wrap_text(c, bdesc, MX + 16, y, BW - 24, "Helvetica", 8.5, 12)
            y -= 4

        y -= 4; y = G._disclaimer_box(c, y, W); G._footer(c, ts, W)

        # ── Page 8: Stage 6 (Fusion) ───────────────────────────────────────
        c.showPage()
        fused_risk  = float(s6.get("fused_risk", 0.0))
        risk_level  = s6.get("risk_level", "?")
        fus_col     = (G.C_RED    if risk_level == "CRITICAL" else
                       G.C_ORANGE if risk_level in ("HIGH","MODERATE") else G.C_GREEN)
        G._page_header(c, "STAGE 6 (FUSION LAYER): MULTI-MODAL REASONING", fus_col, W, H)
        y = G._body_top(H)
        c.setFont("Helvetica-Oblique", 8.5); c.setFillColor(colors.HexColor(G.C_GREY))
        y = G._wrap_text(c,
            "The Fusion layer combines all stage outputs into a single risk score via explicit "
            "rule-based reasoning. The base weight for the AE score is 50%. Semantic classification, "
            "motion, audio, and weapon signals apply additive boosts or penalties. "
            "Each step is recorded below, making this an explainable AI decision layer.",
            MX, y, BW, "Helvetica-Oblique", 8.5, 12)
        y -= 8
        y = G._status_badge(c, f"{risk_level}, Fused Risk = {fused_risk:.2%}", y, fus_col, W, MX)
        y -= 10
        y = G._section_title(c, "Input Signal Values", y, MX, fus_col)
        for k, v in [
            ("AE Score (normalised)",    f"{s6.get('ae_score_norm',0):.4f}  (50% base weight)"),
            ("Binary TF Confidence",     f"{s6.get('ts_binary_conf',0):.4f}  (20% weight if Suspicious)"),
            ("Audio Score",              f"{s6.get('audio_score',0):.4f}  (10% boost if ANOMALY)"),
            ("Weapon Detected",
             "YES  (minimum 0.40 floor; 0.90 floor if AE also anomalous)"
             if s6.get("weapon_flag") else "NO"),
            ("Motion Intensity",         f"{s6.get('motion_intensity',0):.4f}  (15% boost if FAST_MOVEMENT + audio)"),
        ]:
            y = G._kv_row(c, k, v, MX + 6, y, key_w=230)
        y -= 8
        y = G._section_title(c, "Step-by-Step Reasoning", y, MX, fus_col)
        for i, line in enumerate(s6.get("reasoning_trace", []), 1):
            y = G._wrap_text(c, f"Step {i}:  {line}", MX + 8, y, BW - 16, "Helvetica", 8.5, 12)
            y -= 4
            if y < 80: break
        y -= 4; y = G._disclaimer_box(c, y, W); G._footer(c, ts, W)

        # ── Page 9: Stage 7 (RL3033) ──────────────────────────────────────
        c.showPage()
        final_risk = float(s7.get("final_risk", 0.0))
        rl_status  = s7.get("status", "?")
        rl_col     = (G.C_RED    if rl_status in ("CRITICAL","ERROR") else
                      G.C_ORANGE if rl_status in ("HIGH_RISK","MODERATE") else G.C_GREEN)
        G._page_header(c, "STAGE 7 (RL3033, EXPERIMENTAL): TEMPORAL REASONING", rl_col, W, H)
        y = G._body_top(H)
        y = G._experimental_warning_box(c,
            "EXPERIMENTAL: The RL3033 agent is actively being trained and refined. "
            "Its output is blended with the deterministic fusion score rather than used in isolation.",
            y, W, MX)
        c.setFont("Helvetica-Oblique", 8.5); c.setFillColor(colors.HexColor(G.C_GREY))
        y = G._wrap_text(c,
            "RL3033 is a PPO reinforcement learning agent trained by trial and reward to predict risk "
            "from multi-modal observation sequences. Its score is blended with the rule-based fusion "
            "score to produce the final risk value: 70% TFB + 20% TFM + 5% Fused + 5% RL.",
            MX, y, BW, "Helvetica-Oblique", 8.5, 12)
        y -= 8
        y = G._status_badge(c, f"FINAL RISK: {final_risk:.2%}  [{rl_status}]", y, rl_col, W, MX)
        y -= 10
        y = G._section_title(c, "Risk Components", y, MX, rl_col)
        for k, v in [
            ("Fusion Risk Input (Stage 6)", f"{float(s7.get('fused_risk', s6.get('fused_risk',0))):.4f}"),
            ("Final Blended Risk",         f"{final_risk:.4f}  (70% TFB + 20% TFM + 5% Fused + 5% RL)"),
            ("Weapon Override Applied",
             "YES, final risk floored at 0.85" if s6.get("weapon_flag") else "NO"),
            ("RL Agent Status",            rl_status),
            ("Risk Classification",
             "CRITICAL (>= 0.85)" if final_risk >= 0.85 else
             "HIGH_RISK (>= 0.65)" if final_risk >= 0.65 else
             "MODERATE (>= 0.40)"  if final_risk >= 0.40 else "LOW_RISK (< 0.40)"),
        ]:
            y = G._kv_row(c, k, v, MX + 6, y, key_w=250)
        y -= 8

        y = G._section_title(c, "What the RL Agent Sees (Observation Space)", y, MX, rl_col)
        c.setFont("Helvetica-Oblique", 8); c.setFillColor(colors.HexColor(G.C_GREY))
        y = G._wrap_text(c,
            "At each decision step the agent receives five named inputs described below.",
            MX + 6, y, BW - 12, "Helvetica-Oblique", 8, 11)
        y -= 6

        for entry_title, entry_desc in [
            ("ae_signals  (32 x 6)",
             "Sliding window of AE reconstruction errors: raw error, frame delta, acceleration, "
             "short EMA, long EMA, and spike ratio. Allows the agent to track anomaly trends over time."),
            ("latent_motion  (1536 values)",
             "AE channel statistics combined with motion tracking data. Correlates movement patterns "
             "with reconstruction error to distinguish genuine events from artefacts."),
            ("latent_stack  (2560 values)",
             "All pipeline signals combined: AE window, TimeSformer logits, audio features, "
             "tracking features, context vector, and Fusion outputs."),
            ("spatial  (196 values, 14x14 grid)",
             "Spatial heatmap of the frame. Records weapon detection locations and pixel-level "
             "variance per cell. Lets the agent weight spatially localised evidence."),
            ("context  (5 values)",
             "High-level snapshot: normalised AE score, binary classifier confidence, audio score, "
             "weapon flag, and motion intensity."),
        ]:
            if y < 80: break
            c.setFont("Helvetica-Bold", 8.5); c.setFillColor(colors.HexColor(G.C_DARK))
            c.drawString(MX + 8, y, entry_title); y -= 12
            y = G._wrap_text(c, entry_desc, MX + 16, y, BW - 28, "Helvetica", 8.5, 11)
            y -= 5

        if s7.get("interpretation"):
            y -= 4
            y = G._section_title(c, "Interpretation", y, MX, rl_col)
            y = G._wrap_text(c, s7["interpretation"], MX + 8, y, BW - 16, "Helvetica", 9, 13)

        y -= 4; y = G._disclaimer_box(c, y, W); G._footer(c, ts, W)

        # ── Page 10: Top Anomalous Frames ─────────────────────────────────
        frames_data  = s0.get("frames", [])
        scores_list  = s0.get("scores", [])
        samp_indices = s0.get("sampled_indices", [])
        if frames_data and scores_list:
            c.showPage()
            G._page_header(c, "TOP ANOMALOUS FRAMES", G.C_RED, W, H)
            y = G._body_top(H)
            c.setFont("Helvetica-Oblique", 8.5); c.setFillColor(colors.HexColor(G.C_GREY))
            y = G._wrap_text(c,
                "The three frames with the highest AE reconstruction error. High error means the AE "
                "could not reproduce the frame accurately from its learned representation of normal footage.",
                MX, y, BW, "Helvetica-Oblique", 8.5, 12)
            y -= 8
            for i, fd in enumerate(G.extract_top_frames(frames_data, scores_list,
                                    samp_indices or list(range(len(scores_list))), n=3), 1):
                tp    = f"/tmp/af_v3_{i}.jpg"
                Image.fromarray(fd["frame"]).save(tp)
                ratio    = fd["score"] / max(threshold, 1e-8)
                severity = "CRITICAL" if ratio >= 2.0 else "HIGH" if ratio >= 1.0 else "BELOW THRESHOLD"
                c.setFillColor(colors.HexColor(G.C_DARK)); c.setFont("Helvetica-Bold", 9)
                c.drawString(MX, y,
                    f"Frame {i}  |  {ts_of(fd['index'])}  |  "
                    f"Error: {fd['score']:.6f}  |  {ratio:.2f}x threshold  |  Severity: {severity}")
                y -= 14
                img_h = 155
                c.drawImage(tp, MX, y - img_h, width=BW, height=img_h, preserveAspectRatio=True)
                y -= img_h + 12
                if y < 80: break
        G._footer(c, ts, W)

        # ── Page 11: Stage 8 (LLaVA) ──────────────────────────────────────
        if s8 and (s8.get("unified_narrative") or s8.get("frame_descriptions")):
            c.showPage()
            G._page_header(c, "STAGE 8 (LLaVA, EXPERIMENTAL): AI VISUAL ANALYSIS", G.C_PURPLE, W, H)
            y = G._body_top(H)
            y = G._experimental_warning_box(c,
                "EXPERIMENTAL: LLaVA output is a language model generation and may contain "
                "hallucinations. Always verify descriptions against the actual frames.",
                y, W, MX)
            c.setFont("Helvetica-Oblique", 8.5); c.setFillColor(colors.HexColor(G.C_GREY))
            binary_note = s8.get("binary_result", "Unknown")
            y = G._wrap_text(c,
                f"LLaVA receives four key frames with contextual metadata (weapon detections, motion, "
                f"audio) and generates a natural-language description for each. Binary classifier "
                f"result was '{binary_note}'. All four descriptions are synthesised below into a "
                f"single coherent scene narrative. Raw per-frame outputs are in results.json.",
                MX, y, BW, "Helvetica-Oblique", 8.5, 12)
            y -= 10

            ctx = s8.get("context_used", {})
            if ctx:
                y = G._section_title(c, "Context Provided to LLaVA", y, MX, G.C_PURPLE)
                for k, v in [
                    ("Weapon Information",  ctx.get("yolo_summary","?")),
                    ("Motion Behaviours",   ", ".join(ctx.get("behaviors",[])) or "None"),
                    ("Audio Status",        ctx.get("audio_status","?")),
                    ("Audio Score",         f"{ctx.get('audio_score',0):.3f}"),
                    ("Binary Classifier",   binary_note),
                    ("Frames Analysed",     "4 keyframes: clip start, AE onset, AE peak, AE end"),
                ]:
                    y = G._kv_row(c, k, v, MX + 6, y, key_w=185)
                y -= 8

            y = G._section_title(c, "Scene Narrative (Synthesised)", y, MX, G.C_PURPLE)
            c.setFont("Helvetica-Oblique", 8); c.setFillColor(colors.HexColor(G.C_GREY))
            c.drawString(MX + 6, y,
                "Coherent summary derived from all four frame analyses, with redundancy removed.")
            y -= 14

            narrative = s8.get("unified_narrative", "") or " ".join(
                s8.get("frame_descriptions", {}).get(r, "")
                for r in ("baseline","anom_start","peak","anom_end")).strip() or "No description available."

            # Measure narrative height
            import io
            from reportlab.pdfgen.canvas import Canvas as _TC2
            _tc = _TC2(io.BytesIO()); _tc.setFont("Helvetica", 9)
            _words, _line, _nlines = narrative.split(), "", 0
            for _w in _words:
                _test = _line + _w + " "
                if _tc.stringWidth(_test, "Helvetica", 9) < BW - 24: _line = _test
                else: _nlines += 1; _line = _w + " "
            if _line.strip(): _nlines += 1
            narr_box_h = min(_nlines * 13 + 24, y - 90)

            c.setFillColor(colors.HexColor("#f4ecf7"))
            c.setStrokeColor(colors.HexColor("#6c3483")); c.setLineWidth(0.8)
            c.roundRect(MX, y - narr_box_h, BW, narr_box_h, 5, fill=True, stroke=True)
            G._wrap_text(c, narrative, MX + 10, y - 11, BW - 24,
                         "Helvetica", 9, 13, min_y=y - narr_box_h + 10)
            y -= narr_box_h + 14

            # Keyframes panel
            fd_map          = s8.get("frame_descriptions", {})
            frames_d_llava  = s0.get("frames", [])
            peak_idx_llava  = int(s0.get("peak_frame_idx", 0))
            anom_s_llava    = int(s0.get("anomaly_start_frame", 0))
            anom_e_llava    = int(s0.get("anomaly_end_frame", peak_idx_llava))
            n_f_llava       = len(frames_d_llava)
            frame_idx_map   = {
                "baseline":   0,
                "anom_start": min(anom_s_llava, n_f_llava - 1) if n_f_llava else 0,
                "peak":       min(peak_idx_llava, n_f_llava - 1) if n_f_llava else 0,
                "anom_end":   min(anom_e_llava, n_f_llava - 1) if n_f_llava else 0,
            }
            frame_roles = [("baseline","Frame 1: Beginning"),("anom_start","Frame 2: AE Onset"),
                           ("peak","Frame 3: AE Peak"),("anom_end","Frame 4: AE End")]
            if frames_d_llava and y > 130:
                y = G._section_title(c, "Keyframes Used in Scene Narrative", y, MX, G.C_PURPLE)
                c.setFont("Helvetica-Oblique", 8); c.setFillColor(colors.HexColor(G.C_GREY))
                c.drawString(MX + 6, y, "Four frames provided to LLaVA as visual input.")
                y -= 14
                n_kf = len(frame_roles)
                kf_w = (BW - (n_kf - 1) * 6) / n_kf
                kf_h = min(90, kf_w * 0.65)
                if y - kf_h - 36 > 40:
                    for kf_i, (role, label) in enumerate(frame_roles):
                        try:
                            tmp_kf = f"/tmp/llava_kf_{kf_i}.jpg"
                            Image.fromarray(frames_d_llava[frame_idx_map[role]]).save(tmp_kf)
                            cx = MX + kf_i * (kf_w + 6)
                            c.drawImage(tmp_kf, cx, y - kf_h, width=kf_w - 4,
                                        height=kf_h, preserveAspectRatio=True)
                            c.setFont("Helvetica-Bold", 7); c.setFillColor(colors.HexColor(G.C_DARK))
                            c.drawString(cx, y - kf_h - 11, label)
                            raw_desc = fd_map.get(role, "")
                            if raw_desc:
                                max_chars = max(1, int((kf_w - 8) / 3.3))
                                snippet   = raw_desc[:max_chars] + ("…" if len(raw_desc) > max_chars else "")
                                c.setFont("Helvetica-Oblique", 6)
                                c.setFillColor(colors.HexColor(G.C_GREY))
                                c.drawString(cx, y - kf_h - 22, snippet)
                        except Exception:
                            pass
                    y -= kf_h + 36
                y -= 8

            # Signal cross-reference
            if y > 130:
                y = G._section_title(c, "Signal Cross-Reference", y, MX, G.C_PURPLE)
                c.setFont("Helvetica-Oblique", 8); c.setFillColor(colors.HexColor(G.C_GREY))
                y = G._wrap_text(c,
                    "Deterministic system signals used as ground truth context for validation.",
                    MX + 6, y, BW - 12, "Helvetica-Oblique", 8, 11)
                y -= 6
                for k, v in [
                    ("AE Peak Timestamp",       ts_of(s0.get("peak_frame_idx",0))),
                    ("AE Max Error vs Threshold",
                     f"{s0.get('max_error',0):.6f} vs {threshold:.6f} "
                     f"({'exceeds' if s0.get('max_error',0) > threshold else 'below'})"),
                    ("Binary Classification",   f"{s1.get('classification','?')} ({s1.get('confidence',0):.1%})"),
                    ("Anomaly Type",
                     f"{s2.get('class','N/A')} ({s2.get('confidence',0):.1%})"
                     if s2.get("class") not in ("Skipped","Unknown",None) else "N/A"),
                    ("Weapon Detected",         "YES" if s3.get("weapon_flag") else "NO"),
                    ("Audio Status",            s4.get("status","N/A")),
                    ("Motion Behaviours",       ", ".join(s5.get("behaviors",[])) or "None"),
                    ("Final Risk",              f"{s7.get('final_risk', s6.get('fused_risk',0)):.1%}"),
                ]:
                    y = G._kv_row(c, k, v, MX + 6, y, key_w=210)
                    if y < 70: break

            y -= 4; y = G._disclaimer_box(c, y, W); G._footer(c, ts, W)

        c.save()
        return output_path

    @staticmethod
    def generate_all(results: dict) -> dict:
        s0        = results.get("stage0") or {}
        scores    = s0.get("scores", [])
        samp      = s0.get("sampled_indices", [])
        threshold = results.get("threshold", Config.SAFE_THRESHOLD_LEGACY)
        peak_si   = s0.get("peak_score_idx", 0)
        fps       = float(s0.get("fps", 25.0))
        if scores:
            results["graph_path"] = EnhancedReportGenerator.generate_deviation_graph(
                scores, threshold, [peak_si], samp, fps, "/tmp/deviation_graph_v3.png")
        results["pdf_path"] = EnhancedReportGenerator.generate_enhanced_pdf(
            results, "/tmp/forensic_report_v3.pdf")
        return results


# ══════════════════════════════════════════════════════════════════════════════
# MAIN PIPELINE — SCENE SOLVER v3
# ══════════════════════════════════════════════════════════════════════════════
class SceneSolverPipeline:
    """
    Stage ordering (enforced):
      0: AE              — primary anomaly signal
      1: TF Binary       — semantic gate (LOGIC-NORMAL: overrides status to SAFE if Normal)
      2: TF Multi-class  — skipped when Stage 1 = Normal
      3: YOLO Weapon     — skipped when Stage 1 = Normal
      4: Audio           — always runs
      5: Tracking        — always runs
      6: Fusion          — always runs
      7: RL3033          — always runs, FINAL stage
      8: LLaVA           — always runs when enable_llava=True, neutral prompts if Stage 1 = Normal
    """

    def __init__(self, config: Optional[Config] = None):
        self.config       = config or Config()
        self.stage0       = Stage0AE(self.config)
        self.stage1       = Stage1TSBinary(self.config)
        self.stage2       = Stage2TSMulti(self.config)
        self.stage3       = Stage3YOLOWeapon(self.config)
        self.stage4       = Stage4Audio(self.config)
        self.stage5       = Stage5Tracking(self.config)
        self.stage6       = Stage6Fusion(self.config)
        self.stage7       = Stage7RL3033(self.config)
        self.stage8       = Stage8LLaVA(self.config)
        self.alert_system = AlertSystem(self.config)
        self.reporter     = EnhancedReportGenerator()

    def release_stage8(self):
        """Free Stage 8 (LLaVA) weights after an analysis (HF Spaces fix)."""
        try:
            self.stage8.unload()
        except Exception:
            pass

    def _make_output_dir(self, video_path: str, base: str) -> str:
        name = os.path.splitext(os.path.basename(video_path))[0]
        ts   = datetime.now().strftime("%Y%m%d_%H%M%S")
        out  = os.path.join(base, "analysis_results", name, ts)
        os.makedirs(out, exist_ok=True)
        return out

    @staticmethod
    def _skipped_stage2() -> dict:
        return {"class": "Skipped", "confidence": 0.0, "top3": [],
                "all_probabilities": {}, "logits": [0.0] * 7, "sampling_method": "skipped"}

    @staticmethod
    def _skipped_stage3() -> dict:
        return {"weapon_flag": False, "weapon_count": 0, "dominant_class": "None",
                "best_detection": {}, "all_detections": [], "top_frames_annotated": [],
                "spatial_heatmap": np.zeros(196, dtype=np.float32), "summary": "Skipped, Stage 1 = Normal"}

    @staticmethod
    def _disabled_stage5() -> dict:
        return {"motion_intensity": 0.0, "behaviors": ["STATIC"], "track_count": 0,
                "tracking_features": np.zeros(20, dtype=np.float32), "status": "DISABLED"}

    @staticmethod
    def _disabled_stage7(fused_risk: float) -> dict:
        return {"final_risk": fused_risk, "status": "DISABLED",
                "interpretation": "RL agent disabled. Using fused risk from Stage 6."}

    def analyze(self, video_path: str, base_output_dir: Optional[str] = None,
                enable_llava: bool = True, enable_rl: bool = True,
                enable_tracking: bool = True) -> dict:
        base = base_output_dir or self.config.DRIVE_ROOT
        out  = self._make_output_dir(video_path, base)
        log.info("=" * 65)
        log.info(f"  SCENE SOLVER v3  |  {os.path.basename(video_path)}")
        log.info(f"  Output: {out}")
        log.info("=" * 65)

        results: dict = {"video_path": video_path, "output_directory": out}

        # Stage 0
        log.info("\n== STAGE 0: AE ==")
        stage0               = self.stage0.analyze(video_path)
        results["stage0"]    = stage0
        results["threshold"] = stage0["threshold"]
        frames               = stage0["frames"]
        scores               = stage0["scores"]
        sampled_indices      = stage0["sampled_indices"]
        peak_frame_idx       = stage0["peak_frame_idx"]
        peak_score_idx       = stage0.get("peak_score_idx", 0)
        threshold            = stage0["threshold"]
        anomaly_start_frame  = stage0.get("anomaly_start_frame", 0)
        anomaly_end_frame    = stage0.get("anomaly_end_frame", peak_frame_idx)
        ae_signals           = self.stage0.build_ae_signals(scores, peak_score_idx)

        # Stage 1
        log.info("\n== STAGE 1: TimeSformer Binary ==")
        stage1           = self.stage1.analyze(frames, peak_frame_idx)
        results["stage1"]     = stage1
        is_suspicious        = (stage1.get("classification") == "Suspicious")
        ts_binary_logits     = self.stage1.get_logits()
        results["status"]    = stage0["status"] if is_suspicious else "SAFE"
        if not is_suspicious:
            log.info("[Pipeline] Stage 1 = Normal => overall status overridden to SAFE")

        # Stage 2
        log.info("\n== STAGE 2: TimeSformer Multi-class ==")
        stage2 = (self.stage2.analyze(frames, peak_frame_idx, scores, sampled_indices)
                  if is_suspicious else self._skipped_stage2())
        if not is_suspicious:
            log.info("[Stage2-TSMulti] Skipped (Stage 1 = Normal)")
        results["stage2"] = stage2
        ts_multi_logits   = self.stage2.get_logits()

        # Stage 3
        log.info("\n== STAGE 3: YOLO Weapon ==")
        stage3 = (self.stage3.analyze(frames, anomaly_start_frame, anomaly_end_frame)
                  if is_suspicious else self._skipped_stage3())
        if not is_suspicious:
            log.info("[Stage3-YOLO] Skipped (Stage 1 = Normal)")
        results["stage3"] = stage3

        # Stage 4
        log.info("\n== STAGE 4: Audio ==")
        stage4            = self.stage4.analyze(video_path)
        results["stage4"] = stage4

        # Stage 5
        if enable_tracking:
            log.info("\n== STAGE 5: Tracking ==")
            stage5 = self.stage5.analyze(frames, anomaly_start_frame, n_frames=32)
        else:
            log.info("\n== STAGE 5: Tracking DISABLED ==")
            stage5 = self._disabled_stage5()
        results["stage5"] = stage5

        # Stage 6
        log.info("\n== STAGE 6: Fusion ==")
        stage6            = self.stage6.fuse(stage0, stage1, stage2, stage3, stage4, stage5)
        results["stage6"] = stage6

        # Stage 7
        if enable_rl:
            log.info("\n== STAGE 7: RL3033 ==")
            stage7 = self.stage7.analyze(ae_signals, ts_binary_logits, ts_multi_logits,
                                          stage3, stage4, stage5, stage6, frames, peak_frame_idx)
        else:
            log.info("\n== STAGE 7: RL3033 DISABLED ==")
            stage7 = self._disabled_stage7(float(stage6.get("fused_risk", 0.0)))
        results["stage7"] = stage7

        # Sync risk level from Stage 7
        _fr = float(stage7.get("final_risk", stage6.get("fused_risk", 0.0)))
        stage6["risk_level"] = ("CRITICAL" if _fr >= 0.85 else "HIGH" if _fr >= 0.70
                                else "MODERATE" if _fr >= 0.40 else "LOW")
        results["risk_label_final"] = stage6["risk_level"]
        # Keep "status" in sync with the final fused/RL risk, not just the AE/Stage-1 gate
        results["status"] = "CRITICAL" if stage6["risk_level"] in ("CRITICAL", "HIGH") else "SAFE"

        # Stage 8
        if enable_llava:
            log.info("\n== STAGE 8: LLaVA ==")
            stage8 = self.stage8.analyze(frames, peak_frame_idx, anomaly_start_frame,
                                          anomaly_end_frame, stage3, stage4, stage5,
                                          binary_result=stage1.get("classification", "Unknown"))
        else:
            stage8 = None
        results["stage8"] = stage8

        # Alert
        alert            = self.alert_system.evaluate_and_fire(results)
        results["alert"] = alert

        # Visualisations + PDF
        results = EnhancedReportGenerator.generate_all(results)

        self._save_outputs(results, out)
        self._print_summary(results)
        if enable_llava:
            self.release_stage8()
        return results

    def _save_outputs(self, results: dict, out_dir: str):
        def _default(o):
            if isinstance(o, np.ndarray): return o.tolist()
            return str(o)

        # Sanitise for JSON serialisation
        save = dict(results)
        if save.get("stage0"):
            save["stage0"] = {k: v for k, v in save["stage0"].items() if k != "frames"}
        for skey in ("stage3", "stage4", "stage5", "stage6"):
            if save.get(skey):
                sx = dict(save[skey])
                for arr_key in ("spatial_heatmap", "audio_features_vec",
                                "tracking_features", "context_vector"):
                    if arr_key in sx and isinstance(sx[arr_key], np.ndarray):
                        sx[arr_key] = sx[arr_key].tolist()
                if skey == "stage3":
                    sx.pop("all_detections", None)
                    sx.pop("top_frames_annotated", None)
                save[skey] = sx

        with open(os.path.join(out_dir, "results.json"), "w") as f:
            json.dump(save, f, indent=2, default=_default)

        # Summary text
        s0  = results.get("stage0") or {}
        fps = float(s0.get("fps", 25.0))
        ts  = lambda fi: frame_to_timestamp(int(fi), fps)
        s1  = results.get("stage1") or {}
        s2  = results.get("stage2") or {}
        s3  = results.get("stage3") or {}
        s4  = results.get("stage4") or {}
        s5  = results.get("stage5") or {}
        s6  = results.get("stage6") or {}
        s7  = results.get("stage7") or {}
        s8  = results.get("stage8") or {}

        with open(os.path.join(out_dir, "summary.txt"), "w") as f:
            f.write(f"SCENE SOLVER v3  |  {os.path.basename(results.get('video_path',''))}\n")
            f.write("=" * 65 + "\n\n")
            f.write(f"Status              : {results.get('status')}\n")
            f.write(f"Threshold           : {results.get('threshold',0):.6f}\n")
            f.write(f"AE Max Error        : {s0.get('max_error',0):.6f}\n")
            f.write(f"Peak Timestamp      : {ts(s0.get('peak_frame_idx',0))}\n")
            if s1: f.write(f"TF Binary           : {s1.get('classification')} ({s1.get('confidence',0):.1%})\n")
            if s2:
                note = " [Skipped]" if s2.get("class") == "Skipped" else f" ({s2.get('confidence',0):.1%}) [{s2.get('sampling_method','?')}]"
                f.write(f"Anomaly Type        : {s2.get('class','?')}{note}\n")
            if s3: f.write(f"Weapon Detection    : {'Detected' if s3.get('weapon_flag') else 'None'}\n")
            if s4:
                ok = s4.get("audio_present", True)
                f.write(f"Audio               : {'No audio' if not ok else s4.get('status','?')}"
                        + (f"  score={s4.get('audio_score',0):.3f}" if ok else "") + "\n")
            if s5: f.write(f"Tracking            : motion={s5.get('motion_intensity',0):.3f} behaviors={s5.get('behaviors',[])} [{s5.get('status')}]\n")
            if s6:
                f.write(f"Fusion Risk         : {s6.get('fused_risk',0):.1%} [{s6.get('risk_level')}]\n")
                for line in s6.get("reasoning_trace", []):
                    f.write(f"  {line}\n")
            if s7: f.write(f"RL3033 Final        : {s7.get('final_risk',0):.1%} [{s7.get('status')}]\n")
            if results.get("alert"):
                al = results["alert"]
                f.write(f"ALERT               : {al['level']} risk={al['final_risk']:.1%}\n")
            if s8 and s8.get("unified_narrative"):
                f.write(f"\nAI Visual Narrative:\n{s8['unified_narrative']}\n")
            elif s8 and s8.get("description"):
                f.write(f"\nAI Visual Description:\n{s8['description']}\n")

        # PDF copy
        pdf_src = results.get("pdf_path")
        if pdf_src and os.path.exists(pdf_src):
            dest = os.path.join(out_dir, "forensic_report.pdf")
            shutil.copy2(pdf_src, dest)
            results["pdf_path_local"] = dest
            colab_dest = "/content/scene_solver_report.pdf"
            try:
                shutil.copy2(pdf_src, colab_dest)
                log.info(f"[Output] PDF saved to {colab_dest}")
                try:
                    from google.colab import files as colab_files
                    colab_files.download(colab_dest)
                except Exception:
                    pass
            except Exception:
                pass
        log.info(f"[Output] Saved to {out_dir}")

    def _print_summary(self, results: dict):
        s0  = results.get("stage0") or {}
        fps = float(s0.get("fps", 25.0))
        ts  = lambda fi: frame_to_timestamp(int(fi), fps)
        s1  = results.get("stage1") or {}
        s2  = results.get("stage2") or {}
        s3  = results.get("stage3") or {}
        s4  = results.get("stage4") or {}
        s5  = results.get("stage5") or {}
        s6  = results.get("stage6") or {}
        s7  = results.get("stage7") or {}
        s8  = results.get("stage8") or {}
        al  = results.get("alert")

        print("\n" + "=" * 65)
        print("  SCENE SOLVER v3, SUMMARY")
        print("=" * 65)
        print(f"  Overall Status        : {results.get('status')}")
        print(f"  Dynamic Threshold     : {results.get('threshold',0):.6f}")
        print(f"\n  [Stage 0 (AE)]  max={s0.get('max_error',0):.6f}  avg={s0.get('avg_error',0):.6f}  peak={ts(s0.get('peak_frame_idx',0))}")
        print(f"  [Stage 1 (TF Binary)] {s1.get('classification','?')} ({s1.get('confidence',0):.2%})")
        if s2: print(f"  [Stage 2 (TF Multi)]  {s2.get('class','?')} ({s2.get('confidence',0):.2%}) [{s2.get('sampling_method','?')}]")
        if s3: print(f"  [Stage 3 (YOLO)]      {'Detected' if s3.get('weapon_flag') else 'No weapons detected'}")
        if s4:
            ok = s4.get("audio_present", True)
            print(f"  [Stage 4 (Audio)]     {'No audio' if not ok else s4.get('status','?')}" +
                  (f"  score={s4.get('audio_score',0):.3f}" if ok else ""))
        if s5: print(f"  [Stage 5 (Tracking)]  motion={s5.get('motion_intensity',0):.3f}  {s5.get('behaviors',[])}  [{s5.get('status')}]")
        if s6: print(f"  [Stage 6 (Fusion)]    {s6.get('fused_risk',0):.2%}  [{s6.get('risk_level','?')}]")
        if s7: print(f"  [Stage 7 (RL3033)]    final={s7.get('final_risk',0):.2%}  [{s7.get('status','?')}]")
        if al: print(f"\n  ALERT: {al['level']}  |  final_risk={al['final_risk']:.2%}")
        if s8:
            narr = s8.get("unified_narrative") or s8.get("description","")
            if narr: print(f"\n  [Stage 8 (LLaVA)]\n    {narr[:300]}{'...' if len(narr) > 300 else ''}")
        print("=" * 65 + "\n")

    def start_stream(self, source: str, window_frames: int = 50, on_alert: Any = None) -> RealTimeStream:
        stream = RealTimeStream(self)
        stream.start(source, window_frames=window_frames, on_alert=on_alert)
        return stream



# ══════════════════════════════════════════════════════════════════════════════
# READY MESSAGE
# ══════════════════════════════════════════════════════════════════════════════
print("=" * 65)
print("[STATUS] SCENE SOLVER v3, REASONING SYSTEM READY")
print("=" * 65)
print(f"Device             : {Config.DEVICE}")
print(f"RL3033 Available   : {RL_AVAILABLE}")
print(f"YOLO Available     : {YOLO_AVAILABLE}")
print(f"Librosa Available  : {LIBROSA_AVAILABLE}")
print()
print("Usage:")
print("  pipeline = SceneSolverPipeline()")
print("  results  = pipeline.analyze('/path/to/video.mp4')")
print("=" * 65)


# ══════════════════════════════════════════════════════════════════════════════
# SHARED PIPELINE SINGLETON (models are cached across analyses)
# ══════════════════════════════════════════════════════════════════════════════
@lru_cache(maxsize=1)
def get_pipeline() -> SceneSolverPipeline:
    return SceneSolverPipeline()