"""
utils/gradcam.py

Grad-CAM (Gradient-weighted Class Activation Mapping) implementation.

What is Grad-CAM?
─────────────────
Grad-CAM creates a heatmap that shows WHICH PARTS of the image the model 
focused on when making a prediction. This helps us understand and trust 
the model's decisions.

How it works (simplified):
1. Run image through the model.
2. Capture the output of the last convolutional layer (a feature map).
3. Compute how much each feature map channel contributed to the final prediction 
   by looking at gradients (how much the output changes when feature maps change).
4. Weight the feature maps by their gradient importance.
5. Average them to get a single 2D heatmap.
6. Resize to original image size.
7. Apply a colormap (red = most important, blue = least important).
8. Overlay on original image.

This implementation works with ResNet models.
"""

import torch
import torch.nn.functional as F
import numpy as np
from PIL import Image
import cv2


class GradCAM:
    """
    Grad-CAM implementation that hooks into the last convolutional layer.
    """

    def __init__(self, model: torch.nn.Module):
        """
        Args:
            model: A loaded PyTorch model (ResNet architecture).
        """
        self.model = model
        self.model.eval()

        # These will store the captured activations and gradients
        self.feature_maps = None
        self.gradients = None

        # For ResNet: the last conv layer is `layer4`
        # We attach "hooks" — functions that run automatically during forward/backward pass
        self.target_layer = model.layer4

        # Forward hook: captures the output (feature maps) of layer4
        self.forward_hook = self.target_layer.register_forward_hook(
            self._save_feature_maps
        )

        # Backward hook: captures the gradients flowing back through layer4
        self.backward_hook = self.target_layer.register_full_backward_hook(
            self._save_gradients
        )

    def _save_feature_maps(self, module, input, output):
        """Called automatically during forward pass. Saves feature maps."""
        self.feature_maps = output.detach()

    def _save_gradients(self, module, grad_input, grad_output):
        """Called automatically during backward pass. Saves gradients."""
        self.gradients = grad_output[0].detach()

    def compute(self, tensor: torch.Tensor, class_idx: int = None) -> np.ndarray:
        """
        Compute the Grad-CAM heatmap for the given input tensor.

        Args:
            tensor: Preprocessed image tensor of shape [1, 3, 224, 224].
            class_idx: Which class to explain. If None, uses predicted class.

        Returns:
            A 2D numpy array (normalized 0–1) — the heatmap.
        """
        device = next(self.model.parameters()).device
        tensor = tensor.to(device)
        tensor.requires_grad = True

        # Forward pass
        output = self.model(tensor)  # [1, num_classes]

        # Use the predicted class if not specified
        if class_idx is None:
            class_idx = output.argmax(dim=1).item()

        # Zero existing gradients
        self.model.zero_grad()

        # Select the score for the target class and backpropagate
        # This tells us how each part of the network contributed to THIS class
        class_score = output[0, class_idx]
        class_score.backward()

        # Compute importance weights: global average pooling over spatial dimensions
        # gradients shape: [1, channels, h, w]
        weights = self.gradients.mean(dim=(2, 3), keepdim=True)  # [1, channels, 1, 1]

        # Weighted combination of feature maps
        cam = (weights * self.feature_maps).sum(dim=1, keepdim=True)  # [1, 1, h, w]

        # Apply ReLU — only care about features that positively influence the class
        cam = F.relu(cam)

        # Normalize to [0, 1]
        cam = cam.squeeze().cpu().numpy()  # [h, w]
        cam = (cam - cam.min()) / (cam.max() - cam.min() + 1e-8)

        return cam

    def remove_hooks(self):
        """Clean up hooks to prevent memory leaks."""
        self.forward_hook.remove()
        self.backward_hook.remove()


def generate_gradcam(
    original_image: Image.Image,
    model: torch.nn.Module,
    tensor: torch.Tensor
) -> Image.Image:
    """
    Generate a Grad-CAM heatmap overlay for a given model and image.

    Args:
        original_image: The original PIL Image (before preprocessing).
        model: The trained PyTorch model.
        tensor: The preprocessed tensor [1, 3, 224, 224].

    Returns:
        A PIL Image with the heatmap overlaid on the original image.
    """
    # Create the Grad-CAM object and compute heatmap
    gradcam = GradCAM(model)

    try:
        cam = gradcam.compute(tensor)
    finally:
        # Always remove hooks, even if an error occurs
        gradcam.remove_hooks()

    # Resize the heatmap to match original image size
    original_np = np.array(original_image.resize((224, 224)))
    cam_resized = cv2.resize(cam, (224, 224))

    # Apply colormap: converts grayscale heatmap to colorful visualization
    # COLORMAP_JET: blue (cold/ignored) → green → red (hot/focused)
    heatmap_colored = cv2.applyColorMap(
        np.uint8(255 * cam_resized),
        cv2.COLORMAP_JET
    )

    # Convert from BGR (OpenCV default) to RGB (PIL/Streamlit default)
    heatmap_rgb = cv2.cvtColor(heatmap_colored, cv2.COLOR_BGR2RGB)

    # Blend: 60% original image + 40% heatmap
    overlay = cv2.addWeighted(original_np, 0.6, heatmap_rgb, 0.4, 0)

    # Convert numpy array back to PIL Image for Streamlit display
    return Image.fromarray(overlay)