"""
model/model.py
──────────────
DenseNet-121 backbone fine-tuned for binary classification of
breast cancer histopathology images (benign vs. malignant).

Architecture
────────────
  DenseNet-121 (pretrained on ImageNet)
    └─ Adaptive Average Pool → flatten
    └─ Classifier head
         ├─ BatchNorm1d(1024)
         ├─ Dropout(p=0.4)
         ├─ Linear(1024 → 256)  + ReLU
         ├─ BatchNorm1d(256)
         ├─ Dropout(p=0.3)
         └─ Linear(256 → 2)   ← raw logits  [benign, malignant]

Outputs
────────
  logits  : Tensor[1, 2]   — raw scores (pre-softmax)
  probs   : Tensor[1, 2]   — calibrated probabilities via softmax
"""

import torch
import torch.nn as nn
from torchvision import models


class BreastCancerClassifier(nn.Module):
    """
    DenseNet-121 backbone with a custom two-class head.

    Parameters
    ----------
    pretrained : bool
        Load ImageNet weights into the DenseNet-121 backbone (default True).
    freeze_backbone : bool
        Freeze all DenseNet layers except the classifier head (default False).
        Set True for pure feature-extraction / fast fine-tuning scenarios.
    dropout_rate : float
        Dropout probability applied in the classifier head (default 0.4).
    """

    def __init__(
        self,
        pretrained: bool = True,
        freeze_backbone: bool = False,
        dropout_rate: float = 0.4,
    ) -> None:
        super().__init__()

        # ── Backbone ────────────────────────────────────────────────────────
        weights = models.DenseNet121_Weights.IMAGENET1K_V1 if pretrained else None
        densenet = models.densenet121(weights=weights)

        # Keep every layer except the original FC classifier
        self.features = densenet.features          # Conv + DenseBlocks + Transitions
        self.pool     = nn.AdaptiveAvgPool2d((1, 1))

        in_features = densenet.classifier.in_features   # 1024 for DenseNet-121

        # ── Classifier head ─────────────────────────────────────────────────
        self.classifier = nn.Sequential(
            nn.BatchNorm1d(in_features),
            nn.Dropout(p=dropout_rate),
            nn.Linear(in_features, 256),
            nn.ReLU(inplace=True),
            nn.BatchNorm1d(256),
            nn.Dropout(p=dropout_rate * 0.75),
            nn.Linear(256, 2),          # 2 logits: [benign, malignant]
        )

        # ── Optional backbone freeze ─────────────────────────────────────────
        if freeze_backbone:
            for param in self.features.parameters():
                param.requires_grad = False

        self._init_classifier_weights()

    # ────────────────────────────────────────────────────────────────────────
    def _init_classifier_weights(self) -> None:
        """Kaiming / Xavier initialisation for the custom head."""
        for module in self.classifier.modules():
            if isinstance(module, nn.Linear):
                nn.init.kaiming_normal_(module.weight, nonlinearity="relu")
                if module.bias is not None:
                    nn.init.zeros_(module.bias)
            elif isinstance(module, nn.BatchNorm1d):
                nn.init.ones_(module.weight)
                nn.init.zeros_(module.bias)

    # ────────────────────────────────────────────────────────────────────────
    def forward(self, x: torch.Tensor) -> dict:
        """
        Forward pass.

        Parameters
        ----------
        x : torch.Tensor
            Normalised image tensor of shape (B, 3, 224, 224).

        Returns
        -------
        dict with keys
            "logits" : Tensor[B, 2]  — raw model outputs
            "probs"  : Tensor[B, 2]  — softmax probabilities
        """
        features = self.features(x)             # (B, 1024, 7, 7)
        pooled   = self.pool(features)           # (B, 1024, 1, 1)
        flat     = torch.flatten(pooled, 1)      # (B, 1024)
        logits   = self.classifier(flat)         # (B, 2)
        probs    = torch.softmax(logits, dim=1)  # (B, 2)

        return {"logits": logits, "probs": probs}

    # ────────────────────────────────────────────────────────────────────────
    def get_feature_maps(self, x: torch.Tensor) -> torch.Tensor:
        """
        Return the final DenseNet feature maps before pooling.
        Used by Grad-CAM and other spatial explainability modules.

        Returns
        -------
        Tensor[B, 1024, 7, 7]
        """
        return self.features(x)