"""
NoduleClassifier - Step 3 of the thyroid analysis pipeline.

Wraps ResNet50 (fine-tuned for binary benign/malignant classification) and
Grad-CAM for saliency visualisation.

Crop extraction logic:
  - Keep only the largest connected component of the UNet++ mask
  - Derive bounding box from mask pixel coordinates
  - Apply 5% margin on each side  (matches training-time padding)
  - Resize ROI to 224 x 224 with ImageNet normalisation

Classification threshold:
  - Loaded from checkpoint key 'clinical_threshold'
  - Optimised for sensitivity/specificity balance on the validation set

Grad-CAM target layer: resnet_model.layer4[-1]  (last conv block)
"""

import cv2
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.models as models
import torchvision.transforms as transforms
from PIL import Image
from pytorch_grad_cam import GradCAM
from pytorch_grad_cam.utils.image import show_cam_on_image
from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget

from config import (
    RESNET_W, RESNET_IMG_SIZE, IMAGENET_MEAN, IMAGENET_STD, DEVICE,
)
from pipeline.segmentor import NoduleSegmentor


class NoduleClassifier:
    """
    Classifies a nodule ROI as BENIGN or MALIGNANT using ResNet50.

    Architecture:
        backbone : ResNet50 (torchvision, pretrained on ImageNet)
        head     : Dropout(0.5) -> Linear(2048->512) -> ReLU
                   -> Dropout(0.3) -> Linear(512->2)

    Attributes
    model     : nn.Module   - ResNet50
    cam       : GradCAM     - Grad-CAM instance attached to layer4[-1]
    threshold : float       - clinical malignancy probability threshold
    device    : torch.device
    """

    def __init__(self):
        print("[NoduleClassifier] Loading ResNet50...")
        self.device = DEVICE

        self.model = self._build_resnet50().to(self.device)

        ckpt = torch.load(RESNET_W, map_location=self.device, weights_only=False)
        self.model.load_state_dict(ckpt["model_state_dict"])
        self.model.eval()

        self.threshold = float(ckpt.get("clinical_threshold", 0.48))
        print(f"[NoduleClassifier] ResNet50 loaded OK - threshold = {self.threshold:.2f}")

        self.cam = GradCAM(
            model        = self.model,
            target_layers= [self.model.layer4[-1]],
        )
        print("[NoduleClassifier] Grad-CAM initialised on layer4[-1]")


    def classify(
        self,
        image_bgr : np.ndarray,
        mask      : np.ndarray,
    ) -> tuple[float | None, str | None, np.ndarray | None, np.ndarray | None]:
        """
        Extract nodule crop, classify, and compute Grad-CAM.

        Parameters
        image_bgr : np.ndarray (H, W, 3) uint8, BGR
        mask      : np.ndarray (H, W) uint8, 0/255  - UNet++ binary mask

        Returns
            (prob_malign, label, crop_rgb, cam_img)
            prob_malign : float  - P(malignant), in [0, 1]
            label       : str    - "MALIGNANT" or "BENIGN"
            crop_rgb    : np.ndarray (224, 224, 3) uint8 - denormalised crop
            cam_img     : np.ndarray (224, 224, 3) uint8 - Grad-CAM overlay

        Returns (None, None, None, None) if:
            - The clean mask is below the minimum area threshold (0.5% of image)
            - The extracted ROI is empty
        """
        # Keep only largest connected component
        clean_mask   = NoduleSegmentor.get_largest_component(mask)
        h_img, w_img = image_bgr.shape[:2]
        min_area     = int(h_img * w_img * 0.005)

        if int((clean_mask > 0).sum()) < min_area:
            return None, None, None, None

        # Derive bounding box from mask
        coords       = np.where(clean_mask > 0)
        ymin, ymax   = int(coords[0].min()), int(coords[0].max())
        xmin, xmax   = int(coords[1].min()), int(coords[1].max())

        # 5% margin - identical to training-time crop augmentation
        mx  = int((xmax - xmin) * 0.05)
        my  = int((ymax - ymin) * 0.05)
        rx1 = max(0,     xmin - mx)
        ry1 = max(0,     ymin - my)
        rx2 = min(w_img, xmax + mx)
        ry2 = min(h_img, ymax + my)

        roi_bgr = image_bgr[ry1:ry2, rx1:rx2]
        if roi_bgr.size == 0:
            return None, None, None, None

        # Preprocessing transform
        roi_pil = Image.fromarray(cv2.cvtColor(roi_bgr, cv2.COLOR_BGR2RGB))
        eval_tf = transforms.Compose([
            transforms.Resize((RESNET_IMG_SIZE, RESNET_IMG_SIZE)),
            transforms.ToTensor(),
            transforms.Normalize(IMAGENET_MEAN, IMAGENET_STD),
        ])
        img_tensor = eval_tf(roi_pil).unsqueeze(0).to(next(self.model.parameters()).device)

        # Forward pass - classification probability
        with torch.no_grad():
            logits      = self.model(img_tensor)
            prob_malign = float(F.softmax(logits, dim=1)[0, 1].item())

        label      = "MALIGNANT" if prob_malign >= self.threshold else "BENIGN"
        pred_class = 1 if label == "MALIGNANT" else 0

        # Grad-CAM on the predicted class
        self.model.zero_grad()
        cam_map = self.cam(
            input_tensor = img_tensor,
            targets      = [ClassifierOutputTarget(pred_class)],
        )

        # Denormalise crop for RGB visualisation
        mean_t   = torch.tensor(IMAGENET_MEAN).view(3, 1, 1)
        std_t    = torch.tensor(IMAGENET_STD).view(3, 1, 1)
        crop_vis = img_tensor.squeeze().cpu() * std_t + mean_t
        crop_vis = (crop_vis.permute(1, 2, 0).numpy() * 255).clip(0, 255).astype(np.uint8)

        cam_img = show_cam_on_image(
            crop_vis.astype(np.float32) / 255.0,
            cam_map[0],
            use_rgb=True,
        )

        return prob_malign, label, crop_vis, cam_img


    @staticmethod
    def _build_resnet50() -> nn.Module:
        """Construct ResNet50 with the custom classification head."""
        m = models.resnet50(weights=None)
        m.fc = nn.Sequential(
            nn.Dropout(0.5),
            nn.Linear(m.fc.in_features, 512),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(512, 2),
        )
        return m