src/models/gradcam.py

"""
gradcam.py — Grad-CAM implementation on the ResNet-50 CNN branch.

Computes gradient-weighted class activation maps for the wildfire risk
output. Highlights spatial regions that most influenced the model's
fire risk prediction.

Output: Semi-transparent heatmap overlay on the original satellite image
using a jet colormap. Red/yellow = high fire risk regions.
"""

import logging
from typing import Optional, Tuple

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

from src.training.config import PATCH_SIZE, DEVICE

logger = logging.getLogger(__name__)


class GradCAM:
    """
    Grad-CAM visualization for the CNN branch of the MTL model.

    Hooks into the last convolutional layer of ResNet-50 (layer4)
    to compute gradient-weighted class activation maps.
    """

    def __init__(self, model, target_layer=None):
        """
        Args:
            model: MultiTaskFusionModel instance.
            target_layer: Specific layer to hook into. Defaults to
                         the last conv layer of ResNet-50 (layer4).
        """
        self.model = model
        self.feature_maps = None
        self.gradients = None

        # Hook into the target layer
        if target_layer is None:
            target_layer = model.cnn_branch.layer4

        self._register_hooks(target_layer)
        logger.info("GradCAM initialized — hooked into target layer.")

    def _register_hooks(self, target_layer):
        """Register forward and backward hooks on the target layer."""

        def forward_hook(module, input, output):
            self.feature_maps = output.detach()

        def backward_hook(module, grad_input, grad_output):
            self.gradients = grad_output[0].detach()

        target_layer.register_forward_hook(forward_hook)
        target_layer.register_full_backward_hook(backward_hook)

    def generate(
        self,
        image: torch.Tensor,
        timeseries: torch.Tensor,
        target_size: Tuple[int, int] = (PATCH_SIZE, PATCH_SIZE),
    ) -> np.ndarray:
        """
        Generate Grad-CAM heatmap for wildfire risk prediction.

        Args:
            image:       (1, 4, 128, 128) single image tensor.
            timeseries:  (1, 7, 6) single time-series tensor.
            target_size: Output heatmap size (H, W).

        Returns:
            Heatmap array of shape (H, W) with values in [0, 1].
        """
        self.model.eval()

        # Need gradients for Grad-CAM
        image_input = image.clone().requires_grad_(True)

        # Forward pass
        output = self.model(image_input, timeseries)
        risk_map = output["risk_map"]  # (1, 128, 128)

        # Backward pass — use the mean risk score as the target
        self.model.zero_grad()
        target = risk_map.mean()
        target.backward(retain_graph=True)

        if self.gradients is None or self.feature_maps is None:
            logger.warning("Grad-CAM: No gradients captured. Returning uniform map.")
            return np.ones(target_size, dtype=np.float32) * 0.5

        # Compute channel-wise weights: global average pooling of gradients
        weights = torch.mean(self.gradients, dim=(2, 3))  # (1, C)

        # Weighted combination of feature maps
        cam = torch.zeros(
            self.feature_maps.shape[2:], device=self.feature_maps.device
        )

        for i, w in enumerate(weights[0]):
            cam += w * self.feature_maps[0, i]

        # ReLU: keep only positive contributions
        cam = F.relu(cam)

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

        # Resize to target size
        cam_np = cam.cpu().numpy()
        cam_resized = cv2.resize(cam_np, target_size, interpolation=cv2.INTER_LINEAR)

        return cam_resized

    def generate_overlay(
        self,
        image: torch.Tensor,
        timeseries: torch.Tensor,
        alpha: float = 0.5,
    ) -> np.ndarray:
        """
        Generate Grad-CAM overlay on the original satellite image.

        Creates a composite image with:
        - Original RGB satellite image as background
        - Grad-CAM heatmap overlaid with jet colormap
        - Alpha blending at 50% opacity

        Args:
            image:      (1, 4, 128, 128) image tensor.
            timeseries: (1, 7, 6) time-series tensor.
            alpha:      Blending factor.

        Returns:
            BGR overlay image of shape (128, 128, 3), uint8.
        """
        # Generate Grad-CAM heatmap
        heatmap = self.generate(image, timeseries)

        # Convert heatmap to color using jet colormap
        heatmap_uint8 = np.uint8(255 * heatmap)
        heatmap_color = cv2.applyColorMap(heatmap_uint8, cv2.COLORMAP_JET)

        # Extract RGB from original image and convert to uint8
        img_np = image[0, :3].detach().cpu().numpy()  # (3, H, W)
        img_np = np.transpose(img_np, (1, 2, 0))  # (H, W, 3) — RGB

        # Normalize to [0, 255]
        img_min, img_max = img_np.min(), img_np.max()
        if img_max - img_min > 0:
            img_np = (img_np - img_min) / (img_max - img_min)
        img_uint8 = np.uint8(255 * img_np)

        # Convert RGB to BGR for OpenCV
        img_bgr = cv2.cvtColor(img_uint8, cv2.COLOR_RGB2BGR)

        # Alpha-blend overlay
        overlay = cv2.addWeighted(img_bgr, 1 - alpha, heatmap_color, alpha, 0)

        return overlay

    def to_base64(
        self,
        image: torch.Tensor,
        timeseries: torch.Tensor,
    ) -> str:
        """
        Generate Grad-CAM overlay and encode as base64 PNG.

        Args:
            image:      (1, 4, 128, 128) image tensor.
            timeseries: (1, 7, 6) time-series tensor.

        Returns:
            Base64-encoded PNG string.
        """
        import base64

        overlay = self.generate_overlay(image, timeseries)

        # Encode as PNG
        success, buffer = cv2.imencode(".png", overlay)
        if not success:
            logger.error("Failed to encode Grad-CAM overlay as PNG.")
            return ""

        return base64.b64encode(buffer).decode("utf-8")


def generate_gradcam_for_prediction(
    model,
    image: torch.Tensor,
    timeseries: torch.Tensor,
) -> Tuple[np.ndarray, np.ndarray, str]:
    """
    Convenience function to generate all Grad-CAM outputs.

    Returns:
        Tuple of (raw_heatmap, overlay_image, base64_string).
    """
    gradcam = GradCAM(model)

    heatmap = gradcam.generate(image, timeseries)
    overlay = gradcam.generate_overlay(image, timeseries)
    b64 = gradcam.to_base64(image, timeseries)

    return heatmap, overlay, b64


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    from src.models.fusion_model import MultiTaskFusionModel

    model = MultiTaskFusionModel(pretrained_cnn=False)
    model.eval()

    img = torch.randn(1, 4, 128, 128)
    ts = torch.randn(1, 7, 6)

    heatmap, overlay, b64 = generate_gradcam_for_prediction(model, img, ts)
    print(f"Heatmap shape: {heatmap.shape}, range: [{heatmap.min():.3f}, {heatmap.max():.3f}]")
    print(f"Overlay shape: {overlay.shape}")
    print(f"Base64 length: {len(b64)} chars")