"""Grad-CAM heatmaps for DenseNet-121, mapped back onto the original image."""
from __future__ import annotations

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

from .preprocess import to_tensor


def gradcam(model, img: Image.Image, class_idx: int) -> np.ndarray:
    """Return a heatmap in [0, 1] with the same size as ``img``."""
    device = next(model.parameters()).device
    x = to_tensor(img).to(device)
    with torch.enable_grad():
        feats = model.features(x)
        feats.retain_grad()
        out = F.relu(feats)
        logits = model.classifier(torch.flatten(F.adaptive_avg_pool2d(out, 1), 1))
        model.zero_grad(set_to_none=True)
        logits[0, class_idx].backward()
        w = feats.grad.mean(dim=(2, 3), keepdim=True)
        cam = F.relu((w * feats).sum(1, keepdim=True))
    cam = F.interpolate(cam, size=(448, 448), mode="bilinear", align_corners=False)[0, 0]
    cam = cam.detach().cpu().numpy()

    # Undo the square padding: the image occupies a centred region of the padded square.
    W, H = img.size
    m = max(W, H)
    left, top = int((m - W) / 2) / m * 448, int((m - H) / 2) / m * 448
    cam = cam[int(round(top)):int(round(448 - top)), int(round(left)):int(round(448 - left))]
    cam = np.asarray(Image.fromarray(cam.astype(np.float32), mode="F").resize((W, H), Image.BILINEAR))
    rng = cam.max() - cam.min()
    return (cam - cam.min()) / rng if rng > 0 else np.zeros_like(cam)


def overlay(img: Image.Image, cam: np.ndarray, alpha: float = 0.45) -> Image.Image:
    """Blend a heatmap (red = most influential) onto the image."""
    base = np.asarray(img.convert("RGB"), dtype=np.float32)
    c = np.clip(cam, 0, 1)[..., None]
    heat = np.concatenate([np.clip(1.5 * c, 0, 1), np.clip(1.5 * c - 0.5, 0, 1) * 0.6, 0.2 * (1 - c)], axis=-1) * 255
    a = alpha * c
    return Image.fromarray(np.clip(base * (1 - a) + heat * a, 0, 255).astype(np.uint8))