"""
GradCAM spatial attribution for BADAS V-JEPA2 models.

Draws a bounding box on the last frame of a clip showing which spatial
region triggered a high risk prediction.

Key geometry (V-JEPA2 ViT-L, 16 frames, 256x256):
    backbone.encoder.layer  — 24 VJEPA2Layer blocks
    block output            — (1, 2048, 1024)  [present tokens only]
    token layout            — flattened (T_p=8, H_p=16, W_p=16)
    last temporal position  — tokens [1792:2048] → (16, 16) spatial grid
    patch size              — 16px  (each patch = 16x16 pixels)
"""

import numpy as np
import torch
import torch.nn as nn
from dataclasses import dataclass
from scipy.ndimage import gaussian_filter, label as _scipy_label
from typing import List, Optional, Tuple


@dataclass(frozen=True)
class ModelGeometry:
    """Spatial/temporal geometry of a V-JEPA2 model variant.

    Default values match ViT-L/ViT-B with 16 frames @ 256px (16×16 patches).
    Adjust for ViT-G (different patch count) or different input resolutions.
    """
    h_patches: int = 16        # spatial patch grid height
    w_patches: int = 16        # spatial patch grid width
    patch_px: int = 16         # pixels per patch
    n_t_positions: int = 8     # temporal token positions

    @property
    def t_tokens(self) -> int:
        """Number of spatial tokens per temporal position (h_patches * w_patches)."""
        return self.h_patches * self.w_patches

    @property
    def last_t_start(self) -> int:
        """Token index of the start of the last temporal block."""
        return (self.n_t_positions - 1) * self.t_tokens

    @property
    def last_t_end(self) -> int:
        """Token index one past the end of the last temporal block."""
        return self.n_t_positions * self.t_tokens


# Default geometry for V-JEPA2 ViT-L/B, 16 frames, 256px input
_DEFAULT_GEOMETRY = ModelGeometry()
# TODO: SpatialAttribution.__init__ should accept an optional geometry parameter
#       so callers can pass a custom ModelGeometry for ViT-G or non-256px inputs.


class CompiledDeepAttentionForward(nn.Module):
    """
    Compiled full-forward for attention extraction from a fixed encoder layer range.

    Runs the complete backbone forward from pixel_values_videos — no instance-dict
    patching, no forward hooks.  Early blocks (before attn_start) run normally
    without capturing attention.  Late blocks (attn_start to attn_end) are called
    with output_attentions=True as a static literal, slicing to the last-T region
    immediately to avoid materialising the full 2048×2048 matrix.

    Requires the backbone encoder to use attn_implementation='eager', which is how
    attach_attribution() loads the attr_model.

    Compile with torch.compile for a 3–5× speedup over uncompiled execution:
        fwd = torch.compile(CompiledDeepAttentionForward(embeddings, early, late))

    Returns:
        (n_late_blocks, n_heads, last_t_len, last_t_len) float tensor.
    """

    def __init__(
        self,
        embeddings: nn.Module,
        early_blocks: list,
        late_blocks: list,
        last_t_start: int = _DEFAULT_GEOMETRY.last_t_start,
        last_t_end: int = _DEFAULT_GEOMETRY.last_t_end,
    ):
        super().__init__()
        if len(late_blocks) == 0:
            raise ValueError("CompiledDeepAttentionForward requires at least one late block")
        self.embeddings = embeddings
        self.early_blocks = nn.ModuleList(early_blocks)
        self.late_blocks = nn.ModuleList(late_blocks)
        self._lt_start = last_t_start
        self._lt_end = last_t_end

    def forward(self, pixel_values_videos: torch.Tensor) -> torch.Tensor:
        hidden = self.embeddings(pixel_values_videos)
        # Early blocks: forward only, no attention capture
        for block in self.early_blocks:
            hidden = block(hidden)[0]
        # Late blocks: capture attention, slice to last-T immediately
        slices = []
        for block in self.late_blocks:
            out = block(hidden, output_attentions=True)
            hidden = out[0]
            # out[1]: (1, n_heads, N, N) — slice to last-T immediately
            slices.append(
                out[1][0, :, self._lt_start:self._lt_end, self._lt_start:self._lt_end]
            )
        return torch.stack(slices)  # (n_late_blocks, n_heads, 256, 256)


class SpatialAttribution:
    """
    Spatial attribution for BADAS V-JEPA2 models.

    Three methods, in order of production relevance:

    compute_deep_attention(x)  — fastest; no backward pass; works on FP16 model;
                                  used by BADAS.attach_attribution() in production.
    compute_fast(x)            — single forward+backward; FP32 uncompiled only;
                                  good quality reference for the analysis notebook.
    compute(x)                 — full SmoothGrad+GradCAM; slow (~5s); highest quality;
                                  analysis notebook only.

    Args:
        model:                Uncompiled VJEPA2Classifier in eval mode (FP32 or FP16).
                              torch.compile is not supported — attention extraction
                              requires instance-dict patching of submodules.
        block_indices:        Encoder layers used by compute() and compute_fast().
                              Default (18, 21) balances locality and semantics.
        block_range:          (start, end) encoder layer range for compute_deep_attention,
                              e.g. (20, 24) for the last 4 layers (~4× faster than all 24).
                              None = use all layers.
        n_smooth:             SmoothGrad passes for compute(). More = less noise, slower.
        smooth_std:           Noise std for SmoothGrad.
        blur_sigma:           Gaussian blur applied to the final 16×16 map.
        threshold_pct:        Zero out patches below this percentile.
        temporal_temperature: None = last temporal position only.  Float = exponentially
                              weighted sum over all 8 positions (opt-in, mixed results).
        use_attention_weighting: Softly amplify GradCAM by attention weights from the
                              last block. Opt-in; helps on some clips.
        attn_top_k_heads:     For attention weighting, keep only the k sharpest
                              (lowest-entropy) heads. None = all heads.
        suppress_ellipses:    List of ``(cy, cx, ry, rx)`` in normalized [0,1]
                              patch-grid coordinates.  Patches whose center falls
                              inside ``((r-cy)/ry)²+((c-cx)/rx)²<=1`` are zeroed
                              before the gaussian blur.  Useful for suppressing
                              fixed-position overlays (logos, watermarks) that
                              bleed into the attribution map.
                              Example — top-left logo:
                              ``suppress_ellipses=[(0.0, 0.0, 0.25, 0.30)]``
    """

    def __init__(
        self,
        model: nn.Module,
        block_indices: Tuple[int, ...] = (18, 21),
        block_range: Optional[Tuple[int, int]] = None,
        n_smooth: int = 5,
        smooth_std: float = 0.02,
        blur_sigma: float = 1.5,
        threshold_pct: float = 30.0,
        temporal_temperature: Optional[float] = None,
        use_attention_weighting: bool = False,
        attn_top_k_heads: Optional[int] = 4,
        suppress_ellipses: Optional[List[Tuple[float, float, float, float]]] = (
            [(0.0, 0.0, 0.20, 0.25)]  # suppress top-left logo (Nexar watermark)
        ),
    ):
        self.model = model
        self.block_indices = block_indices
        self.block_range = block_range
        self.n_smooth = n_smooth
        self.smooth_std = smooth_std
        self.blur_sigma = blur_sigma
        self.threshold_pct = threshold_pct
        self.use_attention_weighting = use_attention_weighting
        self.attn_top_k_heads = attn_top_k_heads
        self.temporal_temperature = temporal_temperature
        self._suppression_mask = self._build_suppression_mask(suppress_ellipses)

        if n_smooth < 1:
            raise ValueError(f"n_smooth must be >= 1, got {n_smooth}")

        self._blocks = model.backbone.encoder.layer
        if len(self._blocks) < max(block_indices) + 1:
            raise ValueError(
                f"Block index {max(block_indices)} out of range "
                f"(model has {len(self._blocks)} layers)"
            )

        # Fast path: compiled full forward for block_range.
        # Runs embeddings + all 24 blocks, capturing attention only from [start:end].
        # No hooks or caching needed — runs from pixel_values_videos directly.
        self._compiled_fwd = None
        if block_range is not None:
            start, end = block_range
            embeddings = model.backbone.encoder.embeddings
            early_blocks = list(self._blocks)[:start]
            late_blocks = list(self._blocks)[start:end]
            compiled_module = CompiledDeepAttentionForward(embeddings, early_blocks, late_blocks)
            try:
                self._compiled_fwd = torch.compile(compiled_module)
            except Exception:
                self._compiled_fwd = compiled_module  # uncompiled fallback, still correct

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def compute(self, x: torch.Tensor) -> np.ndarray:
        """
        Compute spatial attribution map for the last temporal position.

        Args:
            x: Preprocessed clip tensor (1, T, C, H, W), FP32.

        Returns:
            attr_map: (H_patches, W_patches) float ndarray in [0, 1].
                      (16, 16) for the default 256px / 16px-patch model.
        """
        x = x.float()  # ensure FP32 even if input was FP16
        if x.shape[0] != 1:
            raise ValueError(f"SpatialAttribution requires batch size 1, got {x.shape[0]}")
        accumulated = np.zeros((_DEFAULT_GEOMETRY.h_patches, _DEFAULT_GEOMETRY.w_patches), dtype=np.float32)

        for _ in range(self.n_smooth):
            noise = torch.randn_like(x) * self.smooth_std
            x_noisy = (x + noise).detach().requires_grad_(False)

            cam_sum = np.zeros((_DEFAULT_GEOMETRY.h_patches, _DEFAULT_GEOMETRY.w_patches), dtype=np.float32)

            for block_idx in self.block_indices:
                cam = self._gradcam_for_block(x_noisy, block_idx)  # (H_p, W_p)
                cam_sum += cam

            accumulated += cam_sum / len(self.block_indices)

        attr_map = accumulated / self.n_smooth

        if self.use_attention_weighting:
            attn_map = self._get_attention_map(x.detach())   # (H_p, W_p) in [0,1]
            attr_map = attr_map * (0.5 + 0.5 * attn_map)

        attr_map = self._suppress(attr_map)
        attr_map = gaussian_filter(attr_map, sigma=self.blur_sigma)
        attr_map = self._threshold_and_normalize(attr_map)
        return attr_map

    def compute_fast(self, x: torch.Tensor) -> Tuple[np.ndarray, float]:
        """
        Fast GradCAM for real-time use — single forward+backward pass.

        Hooks all blocks in block_indices simultaneously so the entire
        attribution completes in one forward+backward pass (no SmoothGrad,
        no attention weighting).  ~10× faster than compute() with defaults.

        Also returns the risk score so no separate inference call is needed:

            attr_map, risk = attribution.compute_fast(x)
            if risk > 0.4:
                bbox = attribution.get_bbox(attr_map)

        Args:
            x: Preprocessed clip tensor (1, T, C, H, W), FP32.

        Returns:
            attr_map: (H_patches, W_patches) float ndarray in [0, 1].
            risk:     Sigmoid probability in [0, 1].
        """
        x = x.float()
        if x.shape[0] != 1:
            raise ValueError(
                f"SpatialAttribution requires batch size 1, got {x.shape[0]}"
            )

        # Register hooks on ALL blocks simultaneously
        activation_stores = {idx: [None] for idx in self.block_indices}
        handles = []
        for block_idx in self.block_indices:
            block = self._blocks[block_idx]

            def _make_hook(store):
                def _hook(module, inp, out):
                    hidden = out[0] if isinstance(out, (tuple, list)) else out
                    store[0] = hidden
                    hidden.retain_grad()
                return _hook

            handles.append(block.register_forward_hook(_make_hook(activation_stores[block_idx])))

        cfg = self.model.config.model
        orig_use_fp = getattr(cfg, 'use_future_prediction', False)
        cfg.use_future_prediction = False

        try:
            self.model.zero_grad()
            with torch.enable_grad():
                x_leaf = x.detach().requires_grad_(True)
                logit  = self.model(x_leaf)
                logit.squeeze(0).backward()
        finally:
            for h in handles:
                h.remove()
            cfg.use_future_prediction = orig_use_fp

        risk = float(torch.sigmoid(logit.detach().squeeze()))

        # Compute and average GradCAM across all hooked blocks
        cam_sum = np.zeros((_DEFAULT_GEOMETRY.h_patches, _DEFAULT_GEOMETRY.w_patches), dtype=np.float32)
        for block_idx in self.block_indices:
            act  = activation_stores[block_idx][0]
            grad = act.grad if act is not None else None
            if grad is None:
                continue

            if self.temporal_temperature is None:
                act_t  = act[:, _DEFAULT_GEOMETRY.last_t_start:_DEFAULT_GEOMETRY.last_t_end, :]
                grad_t = grad[:, _DEFAULT_GEOMETRY.last_t_start:_DEFAULT_GEOMETRY.last_t_end, :]
                cam    = torch.relu((grad_t * act_t).sum(dim=-1))
                cam_2d = cam[0].detach().cpu().reshape(_DEFAULT_GEOMETRY.h_patches, _DEFAULT_GEOMETRY.w_patches).numpy()
            else:
                T, cam_agg, weight_sum = self.temporal_temperature, None, 0.0
                for t in range(_DEFAULT_GEOMETRY.n_t_positions):
                    s = t * _DEFAULT_GEOMETRY.t_tokens
                    cam_t = torch.relu(
                        (grad[:, s:s+_DEFAULT_GEOMETRY.t_tokens, :] * act[:, s:s+_DEFAULT_GEOMETRY.t_tokens, :]).sum(dim=-1)
                    )[0]
                    w = float(np.exp((t - (_DEFAULT_GEOMETRY.n_t_positions - 1)) / T))
                    cam_agg = w * cam_t if cam_agg is None else cam_agg + w * cam_t
                    weight_sum += w
                cam_2d = (cam_agg / weight_sum).detach().cpu().reshape(_DEFAULT_GEOMETRY.h_patches, _DEFAULT_GEOMETRY.w_patches).numpy()

            cam_sum += cam_2d

        attr_map = cam_sum / len(self.block_indices)
        attr_map = self._suppress(attr_map)
        attr_map = gaussian_filter(attr_map, sigma=self.blur_sigma)
        attr_map = self._threshold_and_normalize(attr_map)
        return attr_map, risk

    def compute_deep_attention(
        self,
        x: torch.Tensor,
        block_range: Optional[Tuple[int, int]] = None,
        _return_raw: bool = False,
    ) -> np.ndarray:
        """
        Multi-layer attention attribution — no backward pass required.

        Collects self-attention within the last temporal position from encoder
        blocks and averages across layers.  More stable than single-layer
        attention (compute_attention_only) because noisy individual-layer patterns
        cancel out.

        Works only on the uncompiled model (self.model) — instance-dict patching
        is used to force output_attentions=True.  Each block's full (N×N)
        attention matrix is sliced immediately to the last-T region to keep GPU
        memory footprint small (~4 MB per layer vs ~268 MB full).

        Args:
            x:           Preprocessed clip tensor (1, T, C, H, W).
            block_range: (start, end) layer indices to use, e.g. (16, 24) for
                         the last 8 layers.  None = use all layers.  Fewer
                         layers = proportionally faster (~55ms for 4 layers
                         vs ~330ms for 24 layers on an uncompiled FP16 model).
                         Overrides the block_range set at construction if provided.
            _return_raw: If True, return the blurred map before
                         _threshold_and_normalize().  Used by the warmup gate
                         to compare slow and fast paths without amplifying the
                         numerical noise introduced by FP16 compilation.

        Returns:
            attr_map: (H_patches, W_patches) float ndarray in [0, 1].
        """
        if x.shape[0] != 1:
            raise ValueError(
                f"SpatialAttribution requires batch size 1, got {x.shape[0]}"
            )

        # Match input dtype to model — FP16 models work fine here (no backward)
        model_dtype = next(self.model.parameters()).dtype
        x = x.to(dtype=model_dtype)

        all_blocks = self.model.backbone.encoder.layer
        effective_range = block_range if block_range is not None else self.block_range
        if effective_range is not None:
            blocks = list(all_blocks)[effective_range[0]:effective_range[1]]
        else:
            blocks = list(all_blocks)

        layer_maps: list = []

        def _make_patch(store_list, orig_fwd):
            def _patched(*args, **kwargs):
                kwargs['output_attentions'] = True
                result = orig_fwd(*args, **kwargs)
                if isinstance(result, (tuple, list)) and len(result) > 1:
                    A = result[1]   # (1, n_heads, N, N)
                    # Slice immediately — avoids holding N × 268 MB on GPU
                    A_slice = A[0, :, _DEFAULT_GEOMETRY.last_t_start:_DEFAULT_GEOMETRY.last_t_end,
                                      _DEFAULT_GEOMETRY.last_t_start:_DEFAULT_GEOMETRY.last_t_end].detach().cpu()
                    store_list.append(A_slice)  # (n_heads, 256, 256)
                return result
            return _patched

        for block in blocks:
            if not hasattr(block, 'attention'):
                continue
            block.attention.__dict__['forward'] = _make_patch(
                layer_maps, block.attention.forward
            )

        cfg = self.model.config.model
        orig_use_fp = getattr(cfg, 'use_future_prediction', False)
        cfg.use_future_prediction = False

        try:
            with torch.no_grad():
                self.model(x)
        finally:
            for block in blocks:
                block.attention.__dict__.pop('forward', None)
            cfg.use_future_prediction = orig_use_fp

        if not layer_maps:
            return np.zeros((_DEFAULT_GEOMETRY.h_patches, _DEFAULT_GEOMETRY.w_patches), dtype=np.float32)

        # Per layer: mean over heads → sum over query dim → (256,) received attention
        maps = []
        for A_slice in layer_maps:   # each: (n_heads, 256, 256)
            received = A_slice.float().mean(0).sum(0).numpy()   # (256,)
            maps.append(received)

        attn_map = np.mean(maps, axis=0).reshape(_DEFAULT_GEOMETRY.h_patches, _DEFAULT_GEOMETRY.w_patches).astype(np.float32)
        attn_map = self._suppress(attn_map)
        attn_map = gaussian_filter(attn_map, sigma=self.blur_sigma)
        if _return_raw:
            return attn_map
        attn_map = self._threshold_and_normalize(attn_map)
        return attn_map

    def compute_compiled_attention(
        self,
        x: torch.Tensor,
        _return_raw: bool = False,
    ) -> np.ndarray:
        """
        Fast attribution using the compiled full forward — no forward hooks needed.

        Runs embeddings + all encoder blocks from pixel_values_videos, capturing
        attention only from the block_range layers.  Output is numerically
        equivalent to compute_deep_attention() for the same input (same checkpoint
        weights, same code path up to FP16 rounding).

        Requires block_range to be set at __init__ and the backbone to use
        attn_implementation='eager' (which attach_attribution() ensures).

        Args:
            x:           Preprocessed clip tensor (1, T, C, H, W).
            _return_raw: If True, return the blurred map before
                         _threshold_and_normalize().  Used by the warmup gate
                         to compare against compute_deep_attention(_return_raw=True).

        Returns:
            attr_map: (H_patches, W_patches) float ndarray in [0, 1].

        Raises:
            RuntimeError: If block_range was not provided at init.
        """
        if self._compiled_fwd is None:
            raise RuntimeError(
                "compute_compiled_attention requires block_range to be set at __init__. "
                "Pass block_range=(start, end) when constructing SpatialAttribution."
            )

        model_dtype = next(self.model.parameters()).dtype
        x = x.to(dtype=model_dtype)

        with torch.no_grad():
            attn = self._compiled_fwd(x)  # (n_blocks, n_heads, 256, 256) GPU

        # Same aggregation as compute_deep_attention:
        # per block: mean over heads → sum over query dim → (256,) received attention
        attn_cpu = attn.float().cpu()
        maps = []
        for b in range(attn_cpu.shape[0]):
            received = attn_cpu[b].mean(dim=0).sum(dim=0).numpy()  # (256,)
            maps.append(received)

        attn_map = (
            np.mean(maps, axis=0).reshape(_DEFAULT_GEOMETRY.h_patches, _DEFAULT_GEOMETRY.w_patches).astype(np.float32)
        )
        attn_map = self._suppress(attn_map)
        attn_map = gaussian_filter(attn_map, sigma=self.blur_sigma)
        if _return_raw:
            return attn_map
        attn_map = self._threshold_and_normalize(attn_map)
        return attn_map

    def get_bboxes(
        self,
        attr_map: np.ndarray,
        energy_fraction: float = 0.5,
        min_component_fraction: float = 0.10,
        min_dominant_fraction: float = 0.15,
        patch_px: int = _DEFAULT_GEOMETRY.patch_px,
        padding_px: int = 4,
    ) -> List[Tuple[int, int, int, int, float]]:
        """Convert a normalized attribution map to ranked pixel bounding boxes.

        Uses energy-based thresholding and connected-component analysis to find
        spatially distinct focal regions, filter noise, and detect when attention
        is too scattered to produce a meaningful box.

        Args:
            attr_map:               Normalized (0–1) attribution map, shape (H_p, W_p).
            energy_fraction:        Threshold 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 connected region must
                                    carry to get its own box.  Filters spurious warm patches.
            min_dominant_fraction:  If the strongest region carries less than this fraction
                                    of the total signal, attention is too scattered — return [].
            patch_px:               Pixels per patch in the model's input space (default 16).
            padding_px:             Extra pixels added around each box on all sides.

        Returns:
            List of ``(x1, y1, x2, y2, importance)`` tuples sorted by importance
            descending, where ``importance`` is the component's share of total signal.
            Returns ``[]`` when attention is too scattered to identify a focal region.
        """
        total_signal = float(attr_map.sum())
        if total_signal < 1e-6:
            return []

        # ── 1. Energy-based threshold ─────────────────────────────────────────
        # Find the activation value such that the patches at-or-above it together
        # account for energy_fraction of the total signal.
        flat_desc = np.sort(attr_map.flatten())[::-1]
        cumsum = np.cumsum(flat_desc)
        idx = int(np.searchsorted(cumsum, energy_fraction * total_signal))
        threshold = float(flat_desc[min(idx, len(flat_desc) - 1)])
        mask = attr_map >= threshold

        # ── 2. Connected components ───────────────────────────────────────────
        labeled, n_components = _scipy_label(mask)
        if n_components == 0:
            return []

        # ── 3. Rank components by energy ──────────────────────────────────────
        components = []
        for comp_id in range(1, n_components + 1):
            comp_mask = labeled == comp_id
            importance = float(attr_map[comp_mask].sum()) / total_signal
            components.append((comp_mask, importance))
        components.sort(key=lambda x: x[1], reverse=True)

        # ── 4. Scattered-map guard ────────────────────────────────────────────
        if components[0][1] < min_dominant_fraction:
            return []

        # ── 5. Filter and convert to pixel bboxes ────────────────────────────
        H_p, W_p = attr_map.shape
        results = []
        for comp_mask, importance in components:
            if importance < min_component_fraction:
                break  # already sorted — everything after is smaller
            rows = np.where(comp_mask.any(axis=1))[0]
            cols = np.where(comp_mask.any(axis=0))[0]
            r1, r2 = int(rows[0]),  int(rows[-1])
            c1, c2 = int(cols[0]), int(cols[-1])
            x1 = max(0,        c1 * patch_px - padding_px)
            y1 = max(0,        r1 * patch_px - padding_px)
            x2 = min(W_p * patch_px, (c2 + 1) * patch_px + padding_px)
            y2 = min(H_p * patch_px, (r2 + 1) * patch_px + padding_px)
            results.append((x1, y1, x2, y2, round(importance, 3)))
        return results

    def get_bbox(
        self,
        attr_map: np.ndarray,
        **kwargs,
    ) -> Optional[Tuple[int, int, int, int]]:
        """Return the single most important bbox, or None if attention is scattered.

        Thin wrapper around :meth:`get_bboxes` for backward compatibility.
        ``energy_fraction``, ``min_component_fraction``, ``min_dominant_fraction``,
        ``patch_px``, and ``padding_px`` are forwarded if provided.
        """
        bboxes = self.get_bboxes(attr_map, **kwargs)
        if not bboxes:
            return None
        x1, y1, x2, y2, _ = bboxes[0]
        return (x1, y1, x2, y2)

    def annotate_frame(
        self,
        frame_bgr: np.ndarray,
        bbox: Optional[Tuple[int, int, int, int]],
        risk: float,
        label: str = "",
    ) -> np.ndarray:
        """
        Draw bbox and risk score on a copy of frame_bgr (OpenCV BGR format).

        Bbox colour: red if risk >= 0.5, yellow otherwise.
        Returns annotated copy; original is unchanged.
        """
        import cv2
        out = frame_bgr.copy()

        if bbox is not None:
            x1, y1, x2, y2 = bbox
            color = (0, 0, 220) if risk >= 0.5 else (0, 255, 255)   # BGR: red / yellow
            cv2.rectangle(out, (x1, y1), (x2, y2), color, thickness=2)

        text = f"risk={risk:.2f}"
        if label:
            text = f"{label} | {text}"
        cv2.putText(out, text, (6, 22), cv2.FONT_HERSHEY_SIMPLEX,
                    0.6, (255, 255, 255), 2, cv2.LINE_AA)
        cv2.putText(out, text, (6, 22), cv2.FONT_HERSHEY_SIMPLEX,
                    0.6, (0, 0, 0), 1, cv2.LINE_AA)
        return out

    # ------------------------------------------------------------------
    # Private helpers
    # ------------------------------------------------------------------

    def _gradcam_for_block(
        self, x: torch.Tensor, block_idx: int
    ) -> np.ndarray:
        """
        Single GradCAM pass for one transformer block.
        Returns (H_patches, W_patches) float ndarray (un-normalised).
        """
        block = self._blocks[block_idx]
        activation_store: list = [None]

        def fwd_hook(module, inp, out):
            # VJEPA2Layer returns a tuple; the first element is the hidden states
            # tensor with shape (B, N_tokens, D).
            hidden = out[0] if isinstance(out, (tuple, list)) else out
            activation_store[0] = hidden
            hidden.retain_grad()   # keep grad for non-leaf tensor

        handle = block.register_forward_hook(fwd_hook)

        # Temporarily disable future prediction so gradients flow only through
        # the direct encoder→classifier path.  The predictor is a 12-layer
        # transformer that mixes all 2048 present tokens together; its gradient
        # back to the encoder is dense and drowns out per-patch spatial signal.
        cfg = self.model.config.model
        orig_use_fp = getattr(cfg, 'use_future_prediction', False)
        cfg.use_future_prediction = False

        try:
            self.model.zero_grad()
            with torch.enable_grad():
                x_leaf = x.detach().requires_grad_(True)
                logit = self.model(x_leaf)          # scalar logit (binary)
                logit = logit.squeeze(0)   # collapse batch dim only
                logit.backward()
        finally:
            handle.remove()
            cfg.use_future_prediction = orig_use_fp  # always restore

        act = activation_store[0]               # (1, N_tokens, D)
        grad = act.grad                         # (1, N_tokens, D)

        if grad is None:
            return np.zeros((_DEFAULT_GEOMETRY.h_patches, _DEFAULT_GEOMETRY.w_patches), dtype=np.float32)

        if self.temporal_temperature is None:
            # Default: attribute only the last temporal position.
            act_t  = act[:, _DEFAULT_GEOMETRY.last_t_start:_DEFAULT_GEOMETRY.last_t_end, :]   # (1, 256, D)
            grad_t = grad[:, _DEFAULT_GEOMETRY.last_t_start:_DEFAULT_GEOMETRY.last_t_end, :]  # (1, 256, D)
            cam = torch.relu((grad_t * act_t).sum(dim=-1))   # (1, 256)
            cam_2d = cam[0].detach().cpu().reshape(_DEFAULT_GEOMETRY.h_patches, _DEFAULT_GEOMETRY.w_patches).numpy()
        else:
            # Temporal aggregation: accumulate weighted GradCAM over all 8 T positions.
            # w_t = exp((t - (T_p-1)) / temperature) → last position gets weight 1.0.
            T = self.temporal_temperature
            cam_agg: Optional[torch.Tensor] = None
            weight_sum = 0.0

            for t in range(_DEFAULT_GEOMETRY.n_t_positions):
                start = t * _DEFAULT_GEOMETRY.t_tokens
                end   = start + _DEFAULT_GEOMETRY.t_tokens
                act_t  = act[:, start:end, :]    # (1, 256, D)
                grad_t = grad[:, start:end, :]   # (1, 256, D)
                cam_t = torch.relu((grad_t * act_t).sum(dim=-1))[0]   # (256,)
                w = float(np.exp((t - (_DEFAULT_GEOMETRY.n_t_positions - 1)) / T))
                cam_agg = w * cam_t if cam_agg is None else cam_agg + w * cam_t
                weight_sum += w

            cam_2d = (
                (cam_agg / weight_sum).detach().cpu()
                .reshape(_DEFAULT_GEOMETRY.h_patches, _DEFAULT_GEOMETRY.w_patches)
                .numpy()
            )

        return cam_2d

    def _get_attention_map(self, x: torch.Tensor) -> np.ndarray:
        """
        Capture attention weights from the last block in block_indices via one
        no-grad forward pass.

        Patches the attention module's forward in the instance __dict__ to force
        output_attentions=True, then restores by removing the instance attribute.

        Returns (H_p, W_p) float32 in [0, 1]: mean attention received at each
        last-T token from all other tokens. Returns all-ones on failure.
        """
        block = self._blocks[max(self.block_indices)]
        attn_store: list = [None]
        orig_attn_fwd = block.attention.forward

        def _patched_attn(*args, **kwargs):
            kwargs['output_attentions'] = True
            result = orig_attn_fwd(*args, **kwargs)
            if isinstance(result, (tuple, list)) and len(result) > 1:
                attn_store[0] = result[1].detach()  # (B, n_heads, N, N)
            return result

        block.attention.__dict__['forward'] = _patched_attn

        cfg = self.model.config.model
        orig_use_fp = getattr(cfg, 'use_future_prediction', False)
        cfg.use_future_prediction = False

        try:
            with torch.no_grad():
                self.model(x.float())
        finally:
            block.attention.__dict__.pop('forward', None)
            cfg.use_future_prediction = orig_use_fp

        if attn_store[0] is None:
            return np.ones((_DEFAULT_GEOMETRY.h_patches, _DEFAULT_GEOMETRY.w_patches), dtype=np.float32)

        # A[b, h, i, j] = attention FROM token i TO token j
        # Self-attention within last-T: only last-T tokens attending to last-T tokens.
        # Avoids dilution from earlier temporal positions attending to background/road.
        A = attn_store[0]                                                              # (1, n_heads, N, N)
        A_last = A[0, :, _DEFAULT_GEOMETRY.last_t_start:_DEFAULT_GEOMETRY.last_t_end,
                         _DEFAULT_GEOMETRY.last_t_start:_DEFAULT_GEOMETRY.last_t_end]  # (n_heads, 256, 256)

        # Head selection by entropy: prefer sharp (low-entropy) heads.
        # entropy_h = mean over last-T tokens of -sum(p * log(p)) for each head.
        # Keep top-k sharpest heads (k = attn_top_k_heads, default 4).
        eps = 1e-9
        entropy_per_head = -(A_last * (A_last + eps).log()).sum(dim=-1).mean(dim=-1)  # (n_heads,)
        k = self.attn_top_k_heads
        if k is not None and k < A_last.shape[0]:
            top_heads = entropy_per_head.topk(k, largest=False).indices               # k lowest-entropy heads
            A_last = A_last[top_heads]                                                 # (k, 256, 256)

        attn_map = A_last.mean(dim=0).sum(dim=0).cpu().numpy()                        # (256,)

        attn_map = attn_map.reshape(_DEFAULT_GEOMETRY.h_patches, _DEFAULT_GEOMETRY.w_patches)
        if attn_map.max() > 0:
            attn_map = attn_map / attn_map.max()
        return attn_map.astype(np.float32)

    def _threshold_and_normalize(self, arr: np.ndarray) -> np.ndarray:
        if arr.max() == 0:
            return arr
        arr = arr / arr.max()
        thresh = np.percentile(arr, self.threshold_pct)
        arr = np.clip(arr - thresh, 0, None)
        if arr.max() > 0:
            arr = arr / arr.max()
        return arr

    def _build_suppression_mask(
        self,
        ellipses: Optional[List[Tuple[float, float, float, float]]],
    ) -> np.ndarray:
        """Build a (H_p, W_p) multiplicative mask with 0 inside each ellipse.

        Each ellipse is ``(cy, cx, ry, rx)`` in normalized [0, 1] patch-grid
        coordinates.  A patch at grid position (r, c) is suppressed when
        ``((r_norm - cy) / ry)² + ((c_norm - cx) / rx)² <= 1``.
        """
        mask = np.ones((_DEFAULT_GEOMETRY.h_patches, _DEFAULT_GEOMETRY.w_patches), dtype=np.float32)
        if not ellipses:
            return mask
        r_norm = np.linspace(0.0, 1.0, _DEFAULT_GEOMETRY.h_patches)
        c_norm = np.linspace(0.0, 1.0, _DEFAULT_GEOMETRY.w_patches)
        R, C = np.meshgrid(r_norm, c_norm, indexing='ij')  # (H_p, W_p)
        for cy, cx, ry, rx in ellipses:
            if ry <= 0 or rx <= 0:
                continue
            inside = ((R - cy) / ry) ** 2 + ((C - cx) / rx) ** 2 <= 1.0
            mask[inside] = 0.0
        return mask

    def _suppress(self, cam: np.ndarray) -> np.ndarray:
        """Zero out patches covered by the suppression ellipses."""
        return cam * self._suppression_mask
