from __future__ import annotations

"""Grad-CAM visualization for DINOv3 ViT encoders.

Generates attention heatmaps overlaid on input images to show
which regions the model focuses on for drone/satellite matching.

Hook target: last DINOv3Block output (before final LayerNorm).
Uses patch tokens (excluding CLS + registers) to build spatial map.

Usage:
    cam = DINOv3GradCAM(model.drone_encoder, target_layer_idx=-1)
    heatmap = cam.generate(image_tensor, class_idx=None)  # [H, W] numpy
    overlay = cam.overlay(image_tensor, heatmap)  # [3, H, W] torch
    cam.remove_hooks()
"""

import logging
from typing import Any

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

LOGGER = logging.getLogger("caption_test.gradcam")


class DINOv3GradCAM:
    """Grad-CAM for DINOv3 ViT encoder.

    Hooks into a specific transformer block to capture activations
    and gradients, then produces spatial attention heatmaps.

    Args:
        encoder: DINOv3ViT model instance.
        target_layer_idx: Which block to hook (-1 = last block).
        num_registers: Number of register tokens (skipped in spatial map).
    """

    def __init__(
        self,
        encoder: nn.Module,
        target_layer_idx: int = -1,
        num_registers: int = 4,
    ) -> None:
        self.encoder = encoder
        self.num_registers = num_registers
        self._activations: torch.Tensor | None = None
        self._gradients: torch.Tensor | None = None

        # Hook into the target block.
        target_block = encoder.layer[target_layer_idx]
        self._fwd_hook = target_block.register_forward_hook(self._save_activation)
        self._bwd_hook = target_block.register_full_backward_hook(self._save_gradient)

    def _save_activation(self, module: nn.Module, input: Any, output: torch.Tensor) -> None:
        self._activations = output.detach()

    def _save_gradient(self, module: nn.Module, grad_input: Any, grad_output: tuple) -> None:
        self._gradients = grad_output[0].detach()

    @torch.enable_grad()
    def generate(
        self,
        image: torch.Tensor,
        target_embedding: torch.Tensor | None = None,
    ) -> np.ndarray:
        """Generate Grad-CAM heatmap for a single image.

        Args:
            image: Input image [1, 3, H, W].
            target_embedding: Optional target to maximize similarity against.
                If None, uses the CLS token norm as the scalar output.

        Returns:
            Heatmap [H_patches, W_patches] as numpy array, values in [0, 1].
        """
        self.encoder.zero_grad()
        image = image.requires_grad_(True)

        # Forward through the full encoder.
        cls_token = self.encoder(image)  # [1, D]

        # Scalar to backprop from.
        if target_embedding is not None:
            score = (cls_token * target_embedding).sum()
        else:
            score = cls_token.norm(dim=-1).sum()

        score.backward(retain_graph=True)

        if self._activations is None or self._gradients is None:
            LOGGER.warning("No activations/gradients captured")
            return np.zeros((16, 16), dtype=np.float32)

        # Gradients: [1, N_tokens, D] — average over token dim for weights.
        weights = self._gradients.mean(dim=1, keepdim=True)  # [1, 1, D]

        # Weighted activation map.
        # Skip CLS (idx 0) + registers (idx 1..num_registers).
        n_special = 1 + self.num_registers
        patch_activations = self._activations[:, n_special:, :]  # [1, N_patches, D]

        cam = (patch_activations * weights).sum(dim=-1)  # [1, N_patches]
        cam = F.relu(cam)  # Only positive contributions.

        # Reshape to spatial grid.
        n_patches = cam.shape[1]
        h = w = int(n_patches ** 0.5)
        cam = cam.reshape(1, 1, h, w)

        # Normalize to [0, 1].
        cam = cam - cam.min()
        cam_max = cam.max()
        if cam_max > 0:
            cam = cam / cam_max

        return cam.squeeze().cpu().numpy()

    def overlay(
        self,
        image: torch.Tensor,
        heatmap: np.ndarray,
        alpha: float = 0.5,
        colormap: str = "jet",
    ) -> torch.Tensor:
        """Overlay heatmap on image.

        Args:
            image: Original image [1, 3, H, W] or [3, H, W] (ImageNet normalized).
            heatmap: Grad-CAM heatmap [h, w] in [0, 1].
            alpha: Overlay transparency.
            colormap: Matplotlib colormap name.

        Returns:
            Overlaid image [3, H, W] as torch tensor, values in [0, 1].
        """
        import matplotlib.cm as cm

        if image.dim() == 4:
            image = image.squeeze(0)

        _, H, W = image.shape

        # Denormalize from ImageNet stats.
        mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1).to(image.device)
        std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1).to(image.device)
        img = (image * std + mean).clamp(0, 1).cpu().numpy().transpose(1, 2, 0)  # [H, W, 3]

        # Resize heatmap to image size.
        heatmap_resized = torch.tensor(heatmap).unsqueeze(0).unsqueeze(0).float()
        heatmap_resized = F.interpolate(heatmap_resized, size=(H, W), mode="bilinear", align_corners=False)
        heatmap_np = heatmap_resized.squeeze().numpy()

        # Apply colormap.
        cmap = cm.get_cmap(colormap)
        heatmap_colored = cmap(heatmap_np)[:, :, :3]  # [H, W, 3]

        # Blend.
        overlay = alpha * heatmap_colored + (1 - alpha) * img
        overlay = np.clip(overlay, 0, 1)

        return torch.from_numpy(overlay.transpose(2, 0, 1)).float()

    def remove_hooks(self) -> None:
        """Remove forward/backward hooks."""
        self._fwd_hook.remove()
        self._bwd_hook.remove()


def generate_gradcam_samples(
    model: nn.Module,
    dataloader: torch.utils.data.DataLoader,
    device: str,
    output_dir: str,
    n_samples: int = 8,
    epoch: int = 0,
) -> list[torch.Tensor]:
    """Generate Grad-CAM visualizations for a batch of samples.

    Creates overlaid heatmaps for both drone and satellite encoders.
    Saves to output_dir/gradcam/epoch_{N}/.

    Args:
        model: AsymmetricEncoder instance.
        dataloader: Validation dataloader.
        device: Torch device.
        output_dir: Base output directory.
        n_samples: Number of samples to visualize.
        epoch: Current epoch (for naming).

    Returns:
        List of overlay tensors [3, H, W] for logging.
    """
    from pathlib import Path
    import matplotlib.pyplot as plt

    save_dir = Path(output_dir) / "gradcam" / f"epoch_{epoch:03d}"
    save_dir.mkdir(parents=True, exist_ok=True)

    model.eval()
    cam_drone = DINOv3GradCAM(model.drone_encoder, target_layer_idx=-1)
    cam_sat = DINOv3GradCAM(model.sat_encoder, target_layer_idx=-1)

    overlays = []
    batch = next(iter(dataloader))
    drone_imgs = batch["drone_img"][:n_samples].to(device)
    sat_imgs = batch["sat_img"][:n_samples].to(device)

    for i in range(min(n_samples, drone_imgs.shape[0])):
        drone_img = drone_imgs[i:i+1]
        sat_img = sat_imgs[i:i+1]

        # Drone Grad-CAM.
        heatmap_d = cam_drone.generate(drone_img)
        overlay_d = cam_drone.overlay(drone_img, heatmap_d)
        overlays.append(overlay_d)

        # Satellite Grad-CAM.
        heatmap_s = cam_sat.generate(sat_img)
        overlay_s = cam_sat.overlay(sat_img, heatmap_s)
        overlays.append(overlay_s)

        # Save as side-by-side figure.
        fig, axes = plt.subplots(1, 2, figsize=(10, 5))
        axes[0].imshow(overlay_d.permute(1, 2, 0).numpy())
        axes[0].set_title(f"Drone #{i}")
        axes[0].axis("off")
        axes[1].imshow(overlay_s.permute(1, 2, 0).numpy())
        axes[1].set_title(f"Satellite #{i}")
        axes[1].axis("off")
        plt.tight_layout()
        fig.savefig(save_dir / f"sample_{i:03d}.png", dpi=150, bbox_inches="tight")
        plt.close(fig)

    cam_drone.remove_hooks()
    cam_sat.remove_hooks()

    LOGGER.info("Grad-CAM: saved %d samples to %s", len(overlays) // 2, save_dir)
    return overlays
