"""Shared model/Grad-CAM code, extracted from
mit-group8-explainable-defect-detection.ipynb (sections 3 and 6) so the
Space app and the training script can both import it without duplication."""
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as T
import matplotlib

CLASS_NAMES = [
    "crazing",
    "inclusion",
    "patches",
    "pitted_surface",
    "rolled-in_scale",
    "scratches",
]
IMAGE_SIZE = 128
MEAN = (0.5, 0.5, 0.5)
STD = (0.25, 0.25, 0.25)

eval_transform = T.Compose([
    T.Resize((IMAGE_SIZE, IMAGE_SIZE)),
    T.ToTensor(),
    T.Normalize(MEAN, STD),
])

_mean_t = torch.tensor(MEAN).view(3, 1, 1)
_std_t = torch.tensor(STD).view(3, 1, 1)


def to_image(tensor):
    return (tensor.detach().cpu() * _std_t + _mean_t).clamp(0, 1).permute(1, 2, 0).numpy()


class ConvBlock(nn.Module):
    def __init__(self, in_channels, out_channels, stride):
        super().__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride,
                               padding=1, bias=False)
        self.norm = nn.BatchNorm2d(out_channels)
        self.activation = nn.ReLU()

    def forward(self, x):
        return self.activation(self.norm(self.conv(x)))


class DefectCNN(nn.Module):
    def __init__(self, number_of_classes=6):
        super().__init__()
        self.features = nn.Sequential(
            ConvBlock(3, 24, stride=2),
            ConvBlock(24, 48, stride=2),
            ConvBlock(48, 96, stride=2),
            ConvBlock(96, 128, stride=1),
        )
        self.pool = nn.AdaptiveAvgPool2d(1)
        self.classifier = nn.Sequential(
            nn.Dropout(0.20),
            nn.Linear(128, number_of_classes),
        )

    def forward(self, x):
        features = self.features(x)
        return self.classifier(self.pool(features).flatten(1))


class GradCAM:
    def __init__(self, model, target_layer):
        self.model = model
        self.target_layer = target_layer
        self.activations = None
        self.gradients = None
        self.forward_handle = target_layer.register_forward_hook(self._capture_activations)

    def _capture_activations(self, module, inputs, output):
        self.activations = output
        if output.requires_grad:
            output.register_hook(self._capture_gradients)

    def _capture_gradients(self, gradient):
        self.gradients = gradient

    def __call__(self, images, target_classes):
        self.model.eval()
        self.model.zero_grad(set_to_none=True)

        with torch.enable_grad():
            logits = self.model(images)
            target_classes = target_classes.to(logits.device, dtype=torch.long)
            selected_logits = logits.gather(1, target_classes[:, None]).sum()
            selected_logits.backward()

        weights = self.gradients.mean(dim=(2, 3), keepdim=True)
        cam = torch.relu((weights * self.activations).sum(dim=1, keepdim=True))
        cam = F.interpolate(cam, size=images.shape[-2:], mode="bilinear", align_corners=False)

        flat = cam.flatten(start_dim=1)
        minimum = flat.min(dim=1).values[:, None, None, None]
        maximum = flat.max(dim=1).values[:, None, None, None]
        cam = (cam - minimum) / (maximum - minimum + 1e-8)
        return cam[:, 0].detach(), logits.detach()

    def close(self):
        self.forward_handle.remove()


def overlay_heatmap(image_tensor, heatmap, heatmap_weight=0.45):
    image = to_image(image_tensor)
    color = matplotlib.colormaps["inferno"](heatmap.detach().cpu().numpy())[..., :3]
    return np.clip((1 - heatmap_weight) * image + heatmap_weight * color, 0, 1)


def load_model(checkpoint_path, device="cpu"):
    model = DefectCNN(number_of_classes=len(CLASS_NAMES)).to(device)
    state_dict = torch.load(checkpoint_path, map_location=device)
    model.load_state_dict(state_dict)
    model.eval()
    return model