"""
BADAS - Production Ready Inference
==================================
High-performance collision prediction with FP16 + torch.compile optimization.

Features:
    - Zero Lightning dependency for inference
    - Memory-aware model loading (GPU-direct or CPU-first)
    - FP16 inference for maximum speed
    - torch.compile optimization
    - GPU preprocessing with incremental buffer
    - Configurable temporal smoothing (causal/streaming-safe)
    - Debug mode for saving inference windows
    - Two inference paths: streaming (frame-by-frame) and video file

Two inference paths share identical core helpers (_add_frame, _get_tensor,
_build_prediction). Path-specific logic is only in the wrapper that feeds
frames and attaches extra metadata fields.

Usage:
    from badas.inference import BADAS, BADASConfig

    predictor = BADAS(checkpoint_path)

    # Path 1 — streaming (camera / sensor)
    result = predictor.predict_frame(bgr_frame)

    # Path 2 — video file
    for pred in predictor.predict_stream("video.mp4", stride=4):
        print(pred)

    results = predictor.predict_video("video.mp4", stride=4, verbose=True)
"""

import cv2
import torch
import numpy as np
import time
from typing import TYPE_CHECKING, Generator, Dict, Any, List, Optional, Callable, Tuple

if TYPE_CHECKING:
    from badas.inference.callbacks import PredictionCallback
from pathlib import Path

# Optional tqdm import
try:
    from tqdm.auto import tqdm
    TQDM_AVAILABLE = True
except ImportError:
    TQDM_AVAILABLE = False

# Project imports
from badas.inference.config import BADASConfig
from badas.inference.backends import (
    PyTorchBackend,
    ONNXBackend,
    TensorRTBackend,
)
from badas.inference.smoothing import TemporalSmoother
from badas.inference.debug import DebugWindowSaver
from badas.inference.risk import create_risk_assessment
from badas.inference.checkpoint import load_model_from_checkpoint
from badas.core.preprocessing import VideoPreprocessor
from badas.inference.frame_sampling import (
    expected_sampled_frames,
    get_video_info,
    iter_sampled_frames_sequential,
    calculate_timestamp,
)

# Backwards compatibility alias
GPUPreprocessor = VideoPreprocessor

import torch._dynamo
torch._dynamo.config.verbose = False
torch._dynamo.config.suppress_errors = True

import logging
logging.getLogger("torch._inductor").setLevel(logging.WARNING)


class BADAS:
    """
    Production-ready BADAS inference with two optimized paths.

    Path 1 — Streaming (predict_frame):
        Frame-by-frame input (camera / sensor). Synchronous from caller's
        perspective. Uses the same async backend API as path 2 for uniformity.

    Path 2 — Video file (predict_stream / predict_video):
        Reads a video file with a sliding window. Uses a pending/async pattern:
        forward_async() is called at iteration N and collect_result() at N+1,
        so frame decode overlaps with GPU inference.

    Both paths share _add_frame, _get_tensor, and _build_prediction — changing
    preprocessing, smoothing, risk thresholds, or result dict fields in one
    place automatically applies to both.

    Performance (TensorRT FP16, RTX 5090):
        ~24ms inference, ~0.3ms preprocessing per prediction

    Performance (PyTorch FP16 + torch.compile):
        ~38ms inference, ~0.3ms preprocessing per prediction

    Example:
        >>> predictor = BADAS("model.ckpt")
        >>> result = predictor.predict_frame(bgr_frame)
        >>> results = predictor.predict_video("video.mp4", verbose=True)
    """

    @classmethod
    def from_pretrained(
        cls,
        repo_id: str,
        filename: str = "model.ckpt",
        config: Optional[BADASConfig] = None,
        **hf_kwargs,
    ) -> "BADAS":
        """Load a BADAS model from a HuggingFace Hub repository.

        Args:
            repo_id: HuggingFace repo ID, e.g. "nexar-ai/badas-1.5-flash"
            filename: Checkpoint filename in the repo (default: "model.ckpt")
            config: Optional BADASConfig to override defaults
            **hf_kwargs: Extra kwargs forwarded to hf_hub_download
                         (e.g. token=..., cache_dir=...)

        Example:
            >>> predictor = BADAS.from_pretrained("nexar-ai/badas-1.5-flash")
            >>> results = predictor.predict_video("dashcam.mp4")
        """
        from huggingface_hub import hf_hub_download, list_repo_files
        _CKPT_EXTS = ('.ckpt', '.pth', '.pt', '.bin', '.onnx', '.trt', '.engine')
        try:
            ckpt_path = hf_hub_download(repo_id=repo_id, filename=filename, **hf_kwargs)
            return cls(ckpt_path, config=config)
        except Exception:
            pass
        # model.ckpt not found — scan the repo for any checkpoint file.
        try:
            repo_files = list(list_repo_files(repo_id, **{k: v for k, v in hf_kwargs.items() if k == 'token'}))
            found = next((f for f in repo_files if Path(f).suffix.lower() in _CKPT_EXTS), None)
            if found:
                ckpt_path = hf_hub_download(repo_id=repo_id, filename=found, **hf_kwargs)
                return cls(ckpt_path, config=config)
        except Exception:
            pass
        # No checkpoint found — treat as a VLM repo (e.g. Cosmos safetensors).
        token = hf_kwargs.get('token')
        if token:
            import os
            os.environ.setdefault('HF_TOKEN', token)
        return cls(repo_id, config=config)

    def __init__(
        self,
        checkpoint_path: str,
        config: Optional[BADASConfig] = None
    ):
        self.config = config or BADASConfig()
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        self._log("Initializing BADAS")
        self._log(f"   Device: {self.device}")

        if self.config.debug.enabled:
            self._log("   Debug mode: ENABLED")

        self._preprocessor: Optional[VideoPreprocessor] = None

        self._load_model(checkpoint_path)
        self._apply_optimizations()
        self._setup_preprocessor()
        self._warmup()

        self._attribution = None
        self._attr_threshold = 0.4
        self._attr_triggered = False   # True while risk is above threshold
        self._attr_always_compute = False  # True → compute heatmap on every window
        self._frame_count = 0
        self._predictions_made = 0
        self._use_fast_path = False    # True after warmup verifies compiled fwd is fast
        self._hazard_descriptor = None
        self._describe_if_no_focus = True
        self._bbox_energy_fraction = 0.5
        self._bbox_min_component_fraction = 0.10
        self._bbox_min_dominant_fraction = 0.15

        self.smoother = TemporalSmoother(self.config.smoothing)
        self.debug_saver = DebugWindowSaver(
            self.config.debug,
            self.img_size,
            self.target_fps,
            quiet=self.config.quiet
        )

        self._log("Predictor ready")

    # =========================================================================
    # INITIALIZATION (unchanged)
    # =========================================================================

    def _log(self, msg: str):
        if not self.config.quiet:
            print(msg)

    def _detect_backend(self, model_path: str) -> str:
        # Check Gemini prefixes first — these are not file paths
        if model_path.startswith('gemini-') or model_path.startswith('projects/'):
            return 'gemini'
        ext = Path(model_path).suffix.lower()
        if ext == '.onnx':
            return 'onnx'
        elif ext in ('.trt', '.engine'):
            return 'tensorrt'
        elif not ext:
            # No extension → HuggingFace model ID or safetensors directory → VLM
            return 'vlm'
        else:
            return 'pytorch'

    @property
    def _is_vlm_backend(self) -> bool:
        return getattr(getattr(self, '_backend', None), 'needs_raw_frames', False)

    # Extensions that unambiguously identify a local model file
    _REAL_EXTENSIONS = frozenset({'.ckpt', '.pt', '.pth', '.onnx', '.trt', '.engine', '.bin'})

    def _resolve_hf_ckpt(self, model_path: str) -> str:
        """If model_path looks like a HuggingFace repo ID that contains a model.ckpt,
        download it and return the local path. Otherwise return model_path unchanged.
        This lets BADAS("nexar-ai/badas-1.5") work identically to from_pretrained().

        Note: pathlib.Path.suffix is NOT used for HF detection because names like
        "nexar-ai/badas-1.5-flash-lite" contain dots (in "1.5") that parse as
        file extensions, causing false positives."""
        path_obj = Path(model_path)
        # Actual existing file — use it directly
        if path_obj.exists():
            return model_path
        # Unambiguous file extension (even if file doesn't exist yet) — use as-is;
        # a missing local file will fail loudly in load_model_from_checkpoint
        if path_obj.suffix.lower() in self._REAL_EXTENSIONS:
            return model_path
        # Gemini model IDs (vanilla or Vertex AI) — not HF repos, return as-is
        if model_path.startswith('gemini-') or model_path.startswith('projects/'):
            return model_path
        # Anything else is treated as a HF repo ID — try downloading model.ckpt
        try:
            from huggingface_hub import hf_hub_download
            resolved = hf_hub_download(repo_id=model_path, filename="model.ckpt")
            self._log(f"   Resolved HF repo → {resolved}")
            return resolved
        except Exception:
            return model_path   # no model.ckpt in repo → VLM ID, use as-is

    def _load_model(self, model_path: str):
        model_path = self._resolve_hf_ckpt(model_path)
        self._log(f"   Loading model from: {model_path}")
        self._backend_type = self.config.backend or self._detect_backend(model_path)
        self._log(f"   Backend: {self._backend_type}")

        if self._backend_type == 'onnx':
            self._load_onnx_backend(model_path)
        elif self._backend_type == 'tensorrt':
            self._load_tensorrt_backend(model_path)
        elif self._backend_type == 'vlm':
            self._load_vlm_backend(model_path)
        elif self._backend_type == 'gemini':
            self._load_gemini_backend(model_path)
        else:
            self._load_pytorch_backend(model_path)

        self._log(f"   Model config: {self.window_size} frames @ {self.target_fps} FPS, {self.img_size}x{self.img_size}")

    def _load_pytorch_backend(self, checkpoint_path: str):
        self._checkpoint_path = checkpoint_path
        use_direct_gpu = False
        if self.device.type == 'cuda':
            free_memory = torch.cuda.get_device_properties(0).total_memory
            free_memory -= torch.cuda.memory_allocated()
            free_memory_gb = free_memory / 1024**3
            use_direct_gpu = free_memory_gb >= 5.0
            self._log(f"   Free GPU memory: {free_memory_gb:.1f} GB")

        map_location = 'cuda' if use_direct_gpu else 'cpu'
        if use_direct_gpu:
            self._log("   Loading directly to GPU...")
        else:
            self._log("   Loading via CPU (memory-safe mode)...")

        self.model, model_config = load_model_from_checkpoint(
            checkpoint_path,
            map_location=map_location,
            model_name_override=self.config.model_name_override,
        )
        self.model = self.model.eval().to(self.device)

        self.window_size = model_config.data.frame_count
        self.target_fps = model_config.data.dataset_fps
        self.img_size = model_config.data.img_size

        if self.config.num_frames is not None:
            self._log(f"   Frame count override: {self.window_size} → {self.config.num_frames}")
            self.window_size = self.config.num_frames

        if torch.cuda.is_available():
            torch.cuda.empty_cache()

    def _load_onnx_backend(self, onnx_path: str):
        use_cuda = torch.cuda.is_available()
        self._backend = ONNXBackend(onnx_path, use_cuda=use_cuda, quiet=self.config.quiet)
        self.window_size = self.config.num_frames or 16
        self.target_fps = self.config.model_fps
        self.img_size = self.config.model_img_size
        self.model = None

    def _load_tensorrt_backend(self, engine_path: str):
        self._backend = TensorRTBackend(engine_path, quiet=self.config.quiet)
        self.window_size = self.config.num_frames or 16
        self.target_fps = self.config.model_fps
        self.img_size = self.config.model_img_size
        self.model = None

    def _load_vlm_backend(self, model_path: str):
        from collections import deque
        from badas.inference.backends import CosmosVLMBackend
        self._backend = CosmosVLMBackend(
            model_path,
            fps=self.config.model_fps,
            max_pixels=self.config.vlm_max_pixels,
            gpu_memory_utilization=self.config.vlm_gpu_memory_utilization,
            quiet=self.config.quiet,
        )
        self.window_size = self.config.num_frames or 8
        self.target_fps = float(self.config.model_fps)
        self.img_size = self.config.model_img_size
        self.model = None
        self._raw_frame_buffer = deque(maxlen=self.window_size)

    def _load_gemini_backend(self, model_path: str):
        from collections import deque
        from badas.inference.backends import GeminiVLMBackend
        self._backend = GeminiVLMBackend(
            model_path,
            fps=self.config.model_fps,
            quiet=self.config.quiet,
            temperatures=list(self.config.gemini_temperatures),
            max_workers=self.config.gemini_max_workers,
        )
        self.window_size = self.config.num_frames or 8
        self.target_fps = float(self.config.model_fps)
        self.img_size = self.config.model_img_size
        self.model = None
        self._raw_frame_buffer = deque(maxlen=self.window_size)

    def _setup_preprocessor(self):
        if self._is_vlm_backend:
            self._preprocessor = None
            return
        device = 'cuda' if torch.cuda.is_available() else 'cpu'
        self._preprocessor = VideoPreprocessor(
            img_size=self.img_size,
            num_frames=self.window_size,
            device=device,
            dtype=torch.float32,
        )
        device_label = "GPU" if device == 'cuda' else "CPU"
        self._log(f"   {device_label} preprocessing ready (direct resize to {self.img_size}x{self.img_size})")

    def _apply_optimizations(self):
        if self._backend_type not in ('pytorch',):
            return

        if not torch.cuda.is_available():
            self._log("   CUDA not available, running on CPU")
            self.model.eval()
            self._backend = PyTorchBackend(self.model, self.device, self.config.use_autocast)
            return

        if not self.config.use_autocast:
            self.model = self.model.half()
            self._log("   FP16 applied")
        else:
            self._log("   Autocast mode enabled")

        if self.config.use_compile and hasattr(torch, 'compile'):
            self._log("   Compiling model (this takes ~30-60s on first run)...")
            try:
                self.model = torch.compile(self.model, mode="max-autotune")
                self._log("   torch.compile applied")
            except Exception as e:
                self._log(f"   torch.compile failed: {e}")

        self.model.eval()
        self._backend = PyTorchBackend(self.model, self.device, self.config.use_autocast)

    def _warmup(self):
        if self._is_vlm_backend:
            return  # vLLM warms up internally (CUDA graphs captured on load)
        self._log("   Warming up model...")
        dummy_frames = [
            np.zeros((self.img_size, self.img_size, 3), dtype=np.uint8)
            for _ in range(self.window_size)
        ]

        warmup_iterations = 3
        times = []
        for i in range(warmup_iterations):
            start = time.perf_counter()
            self._run_inference(dummy_frames)
            elapsed = (time.perf_counter() - start) * 1000
            times.append(elapsed)
            self._log(f"      Iteration {i+1}: {elapsed:.1f}ms")

        if self.config.startup_ramp:
            self._log("   Warming up startup shapes...")
            dummy_frame = np.zeros((self.img_size, self.img_size, 3), dtype=np.uint8)
            # Fill the main preprocessor progressively, mirroring actual inference exactly:
            # same buffer, same get_partial_tensor path, same forward_async path.
            self._preprocessor.reset()
            for n in range(2, self.window_size, 2):
                self._preprocessor.add_frame(dummy_frame)
                self._preprocessor.add_frame(dummy_frame)
                shape_times = []
                for _ in range(warmup_iterations):
                    t = time.perf_counter()
                    tensor = self._preprocessor.get_partial_tensor(n).clone()
                    handle = self._backend.forward_async(tensor)
                    self._backend.collect_result(handle)
                    shape_times.append((time.perf_counter() - t) * 1000)
                self._log(f"      {n}-frame: {shape_times[0]:.0f}ms → {shape_times[-1]:.0f}ms")
            self._preprocessor.reset()  # leave clean for actual inference

        if torch.cuda.is_available():
            torch.cuda.empty_cache()

        avg_time = np.mean(times[1:]) if len(times) > 1 else times[0]
        self._log(f"   Warmup complete (avg: {avg_time:.1f}ms)")

    def _warmup_attribution(self) -> None:
        """Warm up attribution, verify compiled forward matches slow path, assert speedup."""
        self._log("   Warming up attribution model...")
        dummy = torch.zeros(
            1, self.window_size, 3, self.img_size, self.img_size,
            dtype=torch.float16, device=self.device,
        )

        # Warm up slow path (uncompiled 24-layer forward with instance-dict patching)
        for i in range(2):
            t = time.perf_counter()
            self._attribution.compute_deep_attention(dummy)
            self._log(f"      Slow path warmup {i+1}: {(time.perf_counter()-t)*1000:.1f}ms")

        if torch.cuda.is_available():
            torch.cuda.synchronize()

        # No block_range → no compiled forward available; stay on slow path.
        if self._attribution._compiled_fwd is None:
            self._log("   Attribution: no block_range set — using slow path (~87ms/event)")
            return

        # Warm up compiled forward (triggers torch.compile JIT on first call)
        for i in range(2):
            t = time.perf_counter()
            self._attribution.compute_compiled_attention(dummy)
            self._log(f"      Compiled path warmup {i+1}: {(time.perf_counter()-t)*1000:.1f}ms")

        if torch.cuda.is_available():
            torch.cuda.synchronize()

        # --- Gate 1: Output identity (compare blurred maps before threshold+normalize) ---
        # The threshold+normalize step is hypersensitive on flat (all-zero) dummy input:
        # tiny FP16 compile differences (~0.002/element) can shift the 30th-percentile
        # threshold enough to produce large absolute differences in the final map.
        # Comparing the blurred-but-not-normalized maps avoids this while still
        # catching any real computational divergence.  atol=0.05 is ~10× the observed
        # post-blur compile noise (~0.004 max on dummy input).
        slow_raw = self._attribution.compute_deep_attention(dummy, _return_raw=True)
        fast_raw = self._attribution.compute_compiled_attention(dummy, _return_raw=True)
        if not np.allclose(slow_raw, fast_raw, atol=0.05, rtol=0.05):
            max_diff = float(np.abs(slow_raw - fast_raw).max())
            self._log(
                f"   WARNING: Compiled/slow attribution maps diverge (max diff={max_diff:.4f}) "
                f"— using slow path fallback"
            )
            return

        # --- Gate 2: Speedup measurement ---
        N = 5
        slow_times, fast_times = [], []
        for _ in range(N):
            if torch.cuda.is_available():
                torch.cuda.synchronize()
            t = time.perf_counter()
            self._attribution.compute_deep_attention(dummy)
            if torch.cuda.is_available():
                torch.cuda.synchronize()
            slow_times.append((time.perf_counter() - t) * 1000)

        for _ in range(N):
            if torch.cuda.is_available():
                torch.cuda.synchronize()
            t = time.perf_counter()
            self._attribution.compute_compiled_attention(dummy)
            if torch.cuda.is_available():
                torch.cuda.synchronize()
            fast_times.append((time.perf_counter() - t) * 1000)

        slow_avg = float(np.mean(slow_times))
        fast_avg = float(np.mean(fast_times))
        speedup = slow_avg / fast_avg if fast_avg > 0 else 0.0

        self._log(
            f"   Attribution compiled path: {fast_avg:.1f}ms "
            f"(vs slow={slow_avg:.1f}ms, speedup={speedup:.1f}\u00d7) \u2713"
        )

        if speedup < 2.0:
            self._log(
                f"   WARNING: Speedup {speedup:.1f}\u00d7 < 2\u00d7 expected "
                f"\u2014 enabling compiled path anyway (output is correct)"
            )
        # Enable fast path — output is verified correct
        self._use_fast_path = True

    def _check_attribution_trigger(self, raw_prob: float) -> bool:
        """First-crossing trigger: fires once when risk rises above threshold,
        resets when risk drops back below it."""
        if raw_prob < self._attr_threshold:
            self._attr_triggered = False
            return False
        if self._attr_triggered:
            return False
        self._attr_triggered = True
        return True

    def _attribution_step(
        self,
        prediction: Dict[str, Any],
        tensor: torch.Tensor,
        frame_bgr=None,
    ) -> None:
        """Run attribution and merge results into *prediction* in-place.

        Two independent features:

        1. Threshold trigger (always active when attribution is attached):
           fires on first threshold crossing, runs full attribution (hazard_map +
           bboxes + description if configured).  Result stored in prediction as a
           hazard event (bboxes key present → overlay draws corners + TTS fires).

        2. always_compute=True: additionally compute attention_map on every non-triggered
           window.  Only attention_map and focus_confidence are stored — no bboxes, so
           the overlay does not draw hazard markers on non-hazard frames.
           Use this to supply attention_map for show_heatmap blending across the whole video.
        """
        if self._attribution is None or tensor is None:
            return

        trigger_fired = self._check_attribution_trigger(prediction['probability'])

        if trigger_fired:
            # Full hazard attribution: hazard_map + bboxes + optional description
            attr = self._run_attribution(tensor, frame_bgr=frame_bgr, probability=prediction['probability'])
            prediction.update(attr)
            # Retry makes sense only when a description was attempted but focus was missing.
            # Without a descriptor there is nothing to retry — keep the trigger locked.
            if (self._retry_on_scatter
                    and self._hazard_descriptor is not None
                    and attr['focus_confidence'] == 0.0):
                self._attr_triggered = False  # retry description on next window

        elif self._attr_always_compute:
            # Background attention map — stored under 'attention_map' (not 'hazard_map')
            # so that visualize_predictions and other tools that look for 'hazard_map'
            # only see triggered attribution events, not every background frame.
            attr = self._run_attribution(tensor, frame_bgr=None, probability=prediction['probability'])
            prediction['attention_map'] = attr['hazard_map']
            prediction['attribution_ms'] = attr['attribution_ms']

    def _run_attribution(
        self,
        tensor: torch.Tensor,
        frame_bgr: Optional[np.ndarray] = None,
        probability: float = 0.0,
    ) -> Dict[str, Any]:
        """Run attribution on tensor. Uses compiled fast path when available."""
        t = time.perf_counter()
        if self._use_fast_path:
            attr_map = self._attribution.compute_compiled_attention(tensor)
        else:
            attr_map = self._attribution.compute_deep_attention(tensor)
        bboxes = self._attribution.get_bboxes(
            attr_map,
            energy_fraction=self._bbox_energy_fraction,
            min_component_fraction=self._bbox_min_component_fraction,
            min_dominant_fraction=self._bbox_min_dominant_fraction,
        )
        result = {
            'hazard_map': attr_map,
            'bboxes': bboxes,
            # backward-compat: primary box or None
            'bbox': (bboxes[0][:4] if bboxes else None),
            # Fraction of total attention signal in the primary region [0, 1].
            # 0.0 = scattered (no focal region found); higher = more focused.
            'focus_confidence': bboxes[0][4] if bboxes else 0.0,
            'attribution_ms': round((time.perf_counter() - t) * 1000, 1),
        }
        if self._hazard_descriptor is not None and frame_bgr is not None:
            t_desc = time.perf_counter()
            desc = self._hazard_descriptor.describe(
                frame_bgr,
                bboxes=bboxes,
                risk=probability,
                img_size=self.img_size,
                describe_if_no_focus=self._describe_if_no_focus,
            )
            if desc is not None:
                result['description'] = desc
                result['description_ms'] = round((time.perf_counter() - t_desc) * 1000, 1)
        return result

    def _run_inference(self, frames: List[np.ndarray]) -> float:
        """Synchronous inference on a pre-assembled window. Used by warmup only."""
        video_tensor = self._preprocessor.preprocess_window(frames)
        return self._backend.forward(video_tensor)

    # =========================================================================
    # SHARED CORE HELPERS — used identically by both inference paths
    # =========================================================================

    def _add_frame(self, frame: np.ndarray) -> Tuple[bool, float]:
        """Preprocess one BGR frame into the rolling buffer.

        Returns:
            (buffer_ready, add_frame_ms)
        """
        t = time.perf_counter()
        if self._is_vlm_backend:
            self._raw_frame_buffer.append(frame)
            ready = len(self._raw_frame_buffer) == self.window_size
        else:
            ready = self._preprocessor.add_frame(frame)
        return ready, (time.perf_counter() - t) * 1000

    def _get_tensor(self) -> Tuple[torch.Tensor, float]:
        """Clone the current buffer view for inference.

        The clone is necessary in the video path: forward_async() submits to
        the GPU and returns immediately, so add_frame() for the next window
        can overwrite the circular buffer before the GPU finishes reading.
        The clone gives the backend its own copy.

        Returns:
            (tensor, get_tensor_ms) — tensor shape: (1, T, C, H, W)
        """
        t = time.perf_counter()
        tensor = self._preprocessor.get_tensor().clone()
        return tensor, (time.perf_counter() - t) * 1000

    def _get_partial_tensor(self, n_frames: int) -> Tuple[torch.Tensor, float]:
        """Clone a partial buffer view for startup-ramp inference.

        Returns:
            (tensor, get_tensor_ms) — tensor shape: (1, n_frames, C, H, W)
        """
        t = time.perf_counter()
        tensor = self._preprocessor.get_partial_tensor(n_frames).clone()
        return tensor, (time.perf_counter() - t) * 1000

    def _build_prediction(
        self,
        raw_prob: float,
        inference_ms: float,
        preprocess_ms: float,
        center_idx: int,
        stride: int,
        smoothing_enabled: bool,
        **extra_fields,
    ) -> Dict[str, Any]:
        """Assemble a complete prediction dict.

        Applies temporal smoothing and risk assessment, then builds the result
        dict with all common fields. Path-specific fields arrive via extra_fields.

        Args:
            raw_prob: Raw model output probability (0-1)
            inference_ms: Time from forward_async() call to collect_result() return
            preprocess_ms: Time for add_frame() + get_tensor()
            center_idx: Frame index used as the smoothing timeline position
            stride: Window stride (controls smoothing window)
            smoothing_enabled: Whether to apply temporal smoothing
            **extra_fields: Path-specific fields merged into the result dict

        Returns:
            Complete prediction dict
        """
        if smoothing_enabled:
            prob = self.smoother.smooth(center_idx, raw_prob)
        else:
            prob = raw_prob
            self.smoother._update_history(center_idx, raw_prob, raw_prob)

        risk = create_risk_assessment(prob, self.config.risk)

        return {
            'type': 'prediction',
            'probability': round(prob, 4),
            'raw_probability': round(raw_prob, 4),
            'risk_level': risk['level'].lower(),
            'risk_emoji': risk['emoji'],
            'inference_ms': round(inference_ms, 1),
            'preprocess_ms': round(preprocess_ms, 1),
            'latency_ms': round(inference_ms + preprocess_ms, 1),
            **extra_fields,
        }

    def _flush_pending(
        self,
        pending: Dict[str, Any],
        stride: int,
        smoothing_enabled: bool,
    ) -> Dict[str, Any]:
        """Collect a previously submitted async inference result and build the prediction dict.

        Called once inside the predict_stream loop (to collect before the next window)
        and once after the loop ends (to flush the last pending window).
        """
        raw_prob = self._backend.collect_result(pending['handle'])
        inference_ms = (time.perf_counter() - pending['t_inf']) * 1000

        prediction = self._build_prediction(
            raw_prob, inference_ms, pending['preprocess_ms'],
            center_idx=pending['center_idx'],
            stride=stride,
            smoothing_enabled=smoothing_enabled,
            **pending['extra_fields'],
        )

        if not self._is_vlm_backend:
            self._attribution_step(prediction, pending['tensor'], frame_bgr=pending.get('frame'))

        if self.config.debug.enabled and pending['frames_snapshot'] is not None:
            center_idx = pending['center_idx']
            timestamp = calculate_timestamp(center_idx, self.target_fps)
            rgb_frames = [
                cv2.cvtColor(f, cv2.COLOR_BGR2RGB)
                for f in pending['frames_snapshot']
            ]
            self.debug_saver.save_window(
                frames=rgb_frames,
                window_index=prediction['window_index'],
                frame_index=center_idx,
                timestamp=timestamp,
                prediction=prediction,
            )

        return prediction

    # =========================================================================
    # PATH 1 — STREAMING: frame-by-frame real-time inference
    # =========================================================================

    def predict_frame(
        self,
        frame: np.ndarray,
        apply_smoothing: bool = True
    ) -> Dict[str, Any]:
        """Add one BGR frame and return a prediction when the buffer is full.

        Primary method for real-time streaming inference (camera / sensor).
        Processes exactly one frame per call using the incremental GPU buffer.
        Uses forward_async + collect_result for backend uniformity with path 2,
        though collection is immediate (no next frame to overlap with).

        Args:
            frame: BGR uint8 image from OpenCV or camera
            apply_smoothing: Apply temporal smoothing to the prediction

        Returns:
            If buffer not yet full:
                {'status': 'buffering', 'frame_count': N, 'buffer_fill': N, 'buffer_size': 16}
            If prediction made:
                {'type': 'prediction', 'probability': 0.xxx, 'raw_probability': 0.xxx,
                 'risk_level': str, 'risk_emoji': str,
                 'inference_ms': N, 'preprocess_ms': N, 'latency_ms': N,
                 'frame_count': N, 'predictions_made': N}
            If error:
                {'error': 'message'}
        """
        try:
            self._frame_count += 1

            buffer_ready, add_frame_ms = self._add_frame(frame)

            if not buffer_ready:
                buf_fill = (len(self._raw_frame_buffer) if self._is_vlm_backend
                            else self._preprocessor.buffer_count)
                return {
                    'status': 'buffering',
                    'frame_count': self._frame_count,
                    'buffer_fill': buf_fill,
                    'buffer_size': self.window_size,
                }

            t_inf = time.perf_counter()
            if self._is_vlm_backend:
                raw_prob = self._backend.forward(list(self._raw_frame_buffer))
                tensor = None
                preprocess_ms = add_frame_ms
            else:
                tensor, get_tensor_ms = self._get_tensor()
                preprocess_ms = add_frame_ms + get_tensor_ms
                handle = self._backend.forward_async(tensor)
                raw_prob = self._backend.collect_result(handle)
            inference_ms = (time.perf_counter() - t_inf) * 1000

            self._predictions_made += 1

            prediction = self._build_prediction(
                raw_prob, inference_ms, preprocess_ms,
                center_idx=self._frame_count,
                stride=1,
                smoothing_enabled=apply_smoothing and self.config.smoothing.enabled,
                frame_count=self._frame_count,
                predictions_made=self._predictions_made,
            )

            if not self._is_vlm_backend:
                self._attribution_step(prediction, tensor, frame_bgr=frame)

            return prediction

        except Exception as e:
            return {'error': f'Prediction error: {str(e)}'}

    def reset_streaming_state(self):
        """Reset all streaming state for a new session.

        Call before starting a new streaming session to clear the frame buffer,
        frame counter, prediction counter, and smoothing history.
        """
        if self._is_vlm_backend:
            self._raw_frame_buffer.clear()
        else:
            self._preprocessor.reset()
        self._frame_count = 0
        self._predictions_made = 0
        self._attr_triggered = False
        self.smoother.reset()

    # =========================================================================
    # PATH 2 — VIDEO FILE: sliding window with async pipeline
    # =========================================================================

    def predict_stream(
        self,
        video_path: str,
        stride: int = 4,
        enable_smoothing: Optional[bool] = None,
    ) -> Generator[Dict[str, Any], None, None]:
        """Stream predictions from a video file using a rolling window.

        Uses a pending/async pattern: forward_async() is called at iteration N
        and collect_result() at iteration N+1, so frame decode (~0.5ms CPU)
        overlaps with GPU inference (~24ms on TRT). Results are yielded in order.

        Args:
            video_path: Path to video file
            stride: Sampled frames between predictions (1 = every frame, 4 = every 4th)
            enable_smoothing: Override config smoothing setting

        Yields:
            Prediction dicts with keys:
                type, probability, raw_probability, risk_level, risk_emoji,
                inference_ms, preprocess_ms, latency_ms,
                frame_index, window_start_frame, window_end_frame,
                timestamp, window_index, window_start, window_end, progress
        """
        self.smoother.reset()
        self.debug_saver.reset()
        self.debug_saver.set_video_source(video_path)
        if self._is_vlm_backend:
            self._raw_frame_buffer.clear()
        else:
            self._preprocessor.reset()

        smoothing_enabled = (
            enable_smoothing if enable_smoothing is not None
            else self.config.smoothing.enabled
        )

        video_info = get_video_info(video_path)
        source_fps = video_info['fps']
        num_sampled = expected_sampled_frames(
            video_info['total_frames'], video_info['fps'], self.target_fps
        )
        total_windows = max(1, (num_sampled - self.window_size) // stride + 1)
        # Round stride up to the nearest even number for startup predictions
        # (architecture constraint: model only accepts even frame counts).
        startup_ramp = self.config.startup_ramp
        startup_step = stride + (stride % 2)
        n_startup_windows = (
            (self.window_size - startup_step) // startup_step
            if startup_ramp and not self._is_vlm_backend else 0
        )
        total_windows_adjusted = n_startup_windows + total_windows

        frame_deque = None
        if self.config.debug.enabled:
            from collections import deque
            frame_deque = deque(maxlen=self.window_size)

        # Pending: async inference submitted at iteration N, collected at N+1.
        # Stores the backend handle plus all metadata needed to build the result.
        pending = None

        for orig_idx, sampled_idx, frame in iter_sampled_frames_sequential(
            video_path, target_fps=self.target_fps, convert_rgb=False
        ):
            if frame_deque is not None:
                frame_deque.append(frame)

            # Collect result from the previous async inference.
            # For TRT: GPU has had at least one frame-decode worth of time
            # (~0.5ms) to progress on its ~24ms execution before we wait here.
            if pending is not None:
                yield self._flush_pending(pending, stride, smoothing_enabled)
                pending = None

            buffer_ready, add_frame_ms = self._add_frame(frame)
            n_buffered = min(sampled_idx + 1, self.window_size)

            is_startup = (
                startup_ramp
                and not self._is_vlm_backend
                and not buffer_ready
                and n_buffered >= startup_step
                and n_buffered % startup_step == 0
            )

            if not buffer_ready and not is_startup:
                continue

            # Normal window stride gate (skip for startup windows)
            if not is_startup:
                frames_past_first_window = sampled_idx - (self.window_size - 1)
                if frames_past_first_window < 0 or frames_past_first_window % stride != 0:
                    continue

            # Compute per-window indices
            if is_startup:
                startup_step_num = n_buffered // startup_step              # 1-indexed
                window_idx = startup_step_num - (self.window_size // startup_step)  # negative
                win_start = 0
                center_idx = sampled_idx
                window_end = n_buffered
                progress_num = startup_step_num
            else:
                frames_past_first_window = sampled_idx - (self.window_size - 1)
                window_idx = frames_past_first_window // stride
                win_start = window_idx * stride
                center_idx = win_start + self.window_size // 2
                window_end = win_start + self.window_size
                progress_num = n_startup_windows + window_idx + 1

            # Compute video-specific metadata before submitting (uses center/end indices)
            end_frame_idx = window_end - 1
            end_timestamp = calculate_timestamp(end_frame_idx, self.target_fps)
            window_start_frame = int(win_start * source_fps / self.target_fps)
            window_end_frame = int(end_frame_idx * source_fps / self.target_fps)

            # Submit inference
            t_inf = time.perf_counter()
            if is_startup:
                tensor, get_tensor_ms = self._get_partial_tensor(n_buffered)
                handle = self._backend.forward_async(tensor)
                preprocess_ms = add_frame_ms + get_tensor_ms
            elif self._is_vlm_backend:
                handle = self._backend.forward_async(list(self._raw_frame_buffer))
                preprocess_ms = add_frame_ms
                tensor = None
            else:
                tensor, get_tensor_ms = self._get_tensor()
                preprocess_ms = add_frame_ms + get_tensor_ms
                handle = self._backend.forward_async(tensor)

            pending = {
                'handle': handle,
                't_inf': t_inf,
                'preprocess_ms': preprocess_ms,
                'center_idx': center_idx,
                'tensor': tensor if (
                    self._attribution is not None
                    and not self._is_vlm_backend
                    and not is_startup
                ) else None,
                'frame': frame if (self._attribution is not None and not self._is_vlm_backend) else None,
                'frames_snapshot': list(frame_deque) if frame_deque is not None else None,
                'extra_fields': {
                    'frame_index': window_end_frame,
                    'window_start_frame': window_start_frame,
                    'window_end_frame': window_end_frame,
                    'timestamp': round(end_timestamp, 3),
                    'window_index': window_idx,
                    'window_start': win_start,
                    'window_end': window_end,
                    'progress': progress_num / max(1, total_windows_adjusted),
                },
            }

        # Flush the last pending result (loop exits before it can be collected)
        if pending is not None:
            yield self._flush_pending(pending, stride, smoothing_enabled)

        if self.config.debug.enabled:
            self.debug_saver.save_json()

    def predict_frames(self, frames: List[np.ndarray]) -> float:
        """Fast inference on a pre-extracted list of BGR frames.

        Useful for benchmarking when frames are already loaded (avoids
        re-opening the video file).  Returns raw probability without smoothing.

        NOTE: This is a benchmarking shortcut. Unlike all other inference paths
        (predict_frame, predict_stream, predict_video, StreamingSession,
        RealtimeVideoPlayer), it bypasses the rolling circular buffer and calls
        preprocess_window() directly on a pre-assembled frame list.
        Parity with the rolling-buffer path is verified in
        tests/test_preprocessing_parity.py.  If the preprocessing logic ever
        changes, both add_frame() and preprocess_window() must be updated.

        Args:
            frames: Exactly window_size BGR uint8 numpy arrays (H, W, 3)

        Returns:
            Probability in [0, 1]
        """
        video_tensor = self._preprocessor.preprocess_window(frames)
        return float(self._backend.forward(video_tensor))

    def predict_batch(self, frame_lists: List[List[np.ndarray]]) -> List[float]:
        """Batch inference across N pre-extracted frame windows.

        Preprocesses each window with the shared preprocessor, stacks into a
        single (N, T, C, H, W) tensor, and runs one batched forward pass.
        Falls back to sequential if the backend does not support batching.

        Args:
            frame_lists: N lists, each of exactly window_size BGR uint8 frames

        Returns:
            N raw probabilities in [0, 1] (unsmoothed)
        """
        if not frame_lists:
            return []
        if self._is_vlm_backend:
            return [self._backend.forward(frames) for frames in frame_lists]
        tensors = [self._preprocessor.preprocess_window(f) for f in frame_lists]
        batch = torch.cat(tensors, dim=0)  # (N, T, C, H, W)
        return self._backend.forward_batch(batch)

    def predict_video(
        self,
        video_path: str,
        stride: int = 1,
        verbose: bool = False,
        quiet: Optional[bool] = None,
        on_prediction: Optional[Callable[[Dict[str, Any]], None]] = None,
        callbacks: Optional[List["PredictionCallback"]] = None,
    ) -> List[Dict[str, Any]]:
        """Run prediction on an entire video and return all results.

        Args:
            video_path: Path to video file
            stride: Sampled frames between predictions (default: 4)
            verbose: Print each prediction and a summary to stdout
            quiet: Suppress progress bar (None = use config.show_progress)
            on_prediction: Optional real-time callback called per prediction
            callbacks: Optional list of :class:`~badas.inference.callbacks.PredictionCallback`
                instances.  Each callback is invoked on every prediction in
                order, and may enrich the prediction dict in-place.  Use this
                to attach TTS synthesis, logging, alerting, etc. so that
                downstream visualisation functions (``generate_overlay_video``)
                can consume the enriched data without needing extra arguments.

                Example::

                    from badas.inference.callbacks import TTSCallback
                    predictions = predictor.predict_video(
                        video, callbacks=[TTSCallback(tts_provider)])

        Returns:
            List of prediction dicts (same schema as predict_stream)
        """
        results = []
        stream = self.predict_stream(video_path, stride=stride)

        show_progress = self.config.show_progress if quiet is None else not quiet
        if show_progress and TQDM_AVAILABLE:
            stream = tqdm(stream, desc="Predicting", unit="pred")
        elif show_progress and not TQDM_AVAILABLE:
            self._log("   tqdm not installed. Install with: pip install tqdm")

        should_print = verbose and not self.config.quiet

        t_loop_start = time.perf_counter()
        first_pred_received = False

        for pred in stream:
            results.append(pred)

            if on_prediction is not None:
                on_prediction(pred)

            for cb in (callbacks or []):
                cb(pred)

            if should_print:
                if not first_pred_received:
                    startup_s = time.perf_counter() - t_loop_start
                    n_first = pred['window_end']
                    print(
                        f"[startup] initial {n_first}-frame buffer filled in {startup_s:.2f}s",
                        flush=True,
                    )
                    first_pred_received = True

                prob = pred['probability']
                raw = pred['raw_probability']
                risk = pred['risk_level']
                emoji = pred.get('risk_emoji', '')
                risk_label = f"{emoji} {risk:<10s}" if emoji else f"{risk:<12s}"
                inf_ms = pred['inference_ms']
                pre_ms = pred['preprocess_ms']
                prob_str = (
                    f"P={prob:.3f} (raw: {raw:.3f})"
                    if abs(prob - raw) > 0.005
                    else f"P={prob:.3f}"
                )
                attr_str = (
                    f" | heatmap: {pred['attribution_ms']:.0f}ms"
                    if 'attribution_ms' in pred else ""
                )
                startup_str = (
                    f" [{pred['window_end']}/{self.window_size}f]"
                    if pred['window_index'] < 0 else ""
                )
                print(
                    f"[{pred['timestamp']:6.2f}s] {risk_label} | "
                    f"{prob_str} | inf: {inf_ms:.1f}ms pre: {pre_ms:.1f}ms{attr_str}{startup_str}",
                    flush=True,
                )
                if 'focus_confidence' in pred:
                    fc = pred['focus_confidence']
                    if fc == 0.0:
                        # "retrying" only when trigger fired (bboxes key present = hazard event)
                        suffix = " — retrying" if 'bboxes' in pred else ""
                        print(f"         → scattered (focus_confidence: 0.00){suffix}", flush=True)
                    elif pred.get('description'):
                        desc_ms_str = (
                            f"  (vlm: {pred['description_ms']:.0f}ms)"
                            if 'description_ms' in pred else ""
                        )
                        print(
                            f"         → [conf: {fc:.2f}] {pred['description']}{desc_ms_str}",
                            flush=True,
                        )
                elif pred.get('description'):
                    desc_ms_str = (
                        f"  (vlm: {pred['description_ms']:.0f}ms)"
                        if 'description_ms' in pred else ""
                    )
                    print(f"         → {pred['description']}{desc_ms_str}", flush=True)

        if should_print and results:
            avg_inf = np.mean([r['inference_ms'] for r in results])
            avg_pre = np.mean([r['preprocess_ms'] for r in results])
            max_prob = max(r['probability'] for r in results)
            high_risk = sum(
                1 for r in results
                if r['probability'] >= self.config.risk.high_threshold
            )
            print(f"\n{'='*60}")
            print(
                f"{len(results)} predictions | "
                f"Avg inf: {avg_inf:.1f}ms  Avg pre: {avg_pre:.1f}ms"
            )
            print(f"   Max probability: {max_prob:.3f}")
            print(f"   High risk windows (>={self.config.risk.high_threshold:.0%}): {high_risk}")

            if self.config.debug.enabled:
                print(f"   Debug windows saved to: {self.config.debug.output_dir}/")

        return results

    # =========================================================================
    # UTILITIES
    # =========================================================================

    def attach_attribution(
        self,
        threshold: float = 0.4,
        block_range: Optional[Tuple[int, int]] = None,
        block_indices: Tuple[int, ...] = (18, 21),
        describe: bool = False,
        describe_model: Optional[str] = None,
        describe_if_no_focus: bool = True,
        energy_fraction: float = 0.5,
        min_component_fraction: float = 0.10,
        min_dominant_fraction: float = 0.15,
        retry_on_scatter: bool = True,
        suppress_ellipses=[(0.0, 0.0, 0.20, 0.25)],
        temporal_temperature: Optional[float] = None,
        always_compute: bool = False,
    ) -> None:
        """Load a second model instance for spatial attribution.

        The attribution model is uncompiled FP16 so it can capture attention
        maps via forward hooks (torch.compile fuses these away in the main
        model).  It is only invoked when risk exceeds `threshold`.
        Warmup takes approximately 1–2 seconds when block_range is set and the
        fast path verification runs.

        Args:
            threshold:              Risk score above which attribution is computed.
            block_range:            (start, end) encoder layer indices for deep
                                    attention, e.g. (16, 24) to use only the last
                                    8 layers (~4× faster than all 24).  None = all.
            block_indices:          GradCAM block indices (kept for API consistency).
            describe:               If True, load a VLM and attach a HazardDescriptor
                                    that generates a one-sentence driver warning per
                                    attribution event (~300-500ms, ~4 GB VRAM).
            describe_model:         HuggingFace model ID to use for hazard descriptions.
                                    Defaults to ``"nexar-ai/badas-reason-qwen3-vl-4b"``.
                                    Only used when ``describe=True``.
            describe_if_no_focus:   When attention is too scattered to find a focal
                                    region (empty bboxes), still generate a description
                                    using the full frame and a prompt that acknowledges
                                    the uncertainty.  Requires ``describe=True``.
            energy_fraction:        Bbox tightness: keep the fewest hottest patches
                                    whose cumulative activation reaches this fraction
                                    of the total signal.  Lower → tighter boxes.
            min_component_fraction: Minimum share of total signal a distinct region
                                    must carry to get its own bounding box.
            min_dominant_fraction:  If the strongest region carries less than this
                                    fraction, attention is considered scattered and
                                    no boxes are returned.
            retry_on_scatter:       If True (default), a scattered heatmap resets the
                                    trigger so attribution is retried on the next frame
                                    above threshold.  Set to False to revert to the
                                    original one-shot behaviour (fire once per
                                    threshold crossing, regardless of focus quality).
            suppress_ellipses:      List of ``(cy, cx, ry, rx)`` ellipses in normalized
                                    [0,1] patch-grid coordinates zeroed out of the
                                    attention map post-hoc.  Use to suppress fixed
                                    overlays (logos, watermarks) that attract attention.
                                    Example: ``[(0.0, 0.0, 0.25, 0.30)]`` for a
                                    top-left logo covering ~25% height × 30% width.
            temporal_temperature:   If set, exponentially weight all 8 temporal
                                    positions instead of using only the last frame.
                                    w_t = exp((t-7) / temperature) so the last
                                    position is always weighted most.  2.0 is a good
                                    starting value.  None (default) = last frame only.
            always_compute:         If True, compute and store a heatmap for every
                                    inference window regardless of the threshold.
                                    Use this to collect heatmaps for the full video
                                    (e.g. for generate_heatmap_video).  When False
                                    (default) the original first-crossing threshold
                                    trigger is used.
        """
        if self._backend_type != 'pytorch':
            self._log("   Attribution requires PyTorch backend — skipped")
            return

        # Reset fast path if attach_attribution is called a second time.
        self._use_fast_path = False
        self._attr_triggered = False  # reset trigger so first event after re-attach fires

        from badas.inference.checkpoint import load_model_from_checkpoint
        from badas.analysis.attribution import SpatialAttribution

        self._log("   Loading attribution model (uncompiled FP16, eager attn)...")
        attr_model, _ = load_model_from_checkpoint(
            self._checkpoint_path,
            map_location='cpu',
            attn_implementation='eager',
        )
        attr_model = attr_model.half().eval().to(self.device)
        self._share_weights_to_attr_model(attr_model)

        self._attribution = SpatialAttribution(
            attr_model, block_indices=block_indices, block_range=block_range,
            suppress_ellipses=suppress_ellipses,
            temporal_temperature=temporal_temperature,
        )
        self._attr_threshold = threshold
        self._attr_always_compute = always_compute
        self._describe_if_no_focus = describe_if_no_focus
        self._bbox_energy_fraction = energy_fraction
        self._bbox_min_component_fraction = min_component_fraction
        self._bbox_min_dominant_fraction = min_dominant_fraction
        self._retry_on_scatter = retry_on_scatter

        self._warmup_attribution()

        if describe:
            from badas.analysis.hazard_description import HazardDescriptor
            kwargs = {"device": str(self.device)}
            if describe_model is not None:
                kwargs["model_id"] = describe_model
            model_label = describe_model or "badas-reason-qwen3-vl-4b"
            self._log(f"   Loading hazard descriptor ({model_label})...")
            self._hazard_descriptor = HazardDescriptor(**kwargs)
            self._hazard_descriptor.load()
            self._log("   Hazard descriptor ready")

        layers = (block_range[1] - block_range[0]) if block_range else len(attr_model.backbone.encoder.layer)
        self._log(f"   Attribution ready — threshold={threshold}, layers={layers}")

    def _share_weights_to_attr_model(self, attr_model: "torch.nn.Module") -> None:
        """Point attr_model's parameter tensors at the main model's tensors.

        Both models share the same checkpoint and differ only in
        attn_implementation (sdpa vs eager), so all weight names and shapes
        are identical.  After this call attr_model holds no independent weight
        storage; the ~600 MB previously allocated for it is freed on the next
        empty_cache call.
        """
        if self.model is None:
            return

        # Unwrap torch.compile wrapper (_orig_mod) if present
        source = getattr(self.model, '_orig_mod', self.model)
        source_params = dict(source.named_parameters())

        shared, skipped = 0, []
        for name, attr_param in attr_model.named_parameters():
            if name not in source_params:
                skipped.append(f"missing: {name}")
                continue
            src = source_params[name]
            if attr_param.shape != src.shape:
                skipped.append(f"shape mismatch {name}: {attr_param.shape} vs {src.shape}")
                continue
            attr_param.data = src.data
            shared += 1

        if skipped:
            self._log(f"   Weight sharing: {shared} shared, {len(skipped)} skipped")
            for s in skipped[:5]:
                self._log(f"      {s}")
        else:
            self._log(f"   Weight sharing: all {shared} params shared (no duplicate weights)")

        if torch.cuda.is_available():
            torch.cuda.empty_cache()

    def get_memory_info(self) -> Dict[str, float]:
        """Get current GPU memory usage information."""
        if not torch.cuda.is_available():
            return {"device": "cpu"}
        return {
            "device": torch.cuda.get_device_name(0),
            "allocated_mb": torch.cuda.memory_allocated() / 1024**2,
            "reserved_mb": torch.cuda.memory_reserved() / 1024**2,
            "total_mb": torch.cuda.get_device_properties(0).total_memory / 1024**2,
        }


def quick_predict(
    checkpoint_path: str,
    video_path: str,
    stride: int = 4,
    verbose: bool = True,
    quiet: bool = False,
    show_progress: bool = False,
) -> List[Dict[str, Any]]:
    """Quick prediction without keeping the predictor in memory."""
    config = BADASConfig(quiet=quiet, show_progress=show_progress)
    predictor = BADAS(checkpoint_path, config=config)
    return predictor.predict_video(video_path, stride=stride, verbose=verbose)
