"""
model.py – Model definitions for Chest X-Ray classification.

Two architectures:
  - SimpleCNN       : lightweight 3-block CNN trained from scratch (baseline)
  - MobileNetV2     : fine-tuned ImageNet pretrained network (strong model)

Both expose the same forward(x) → logits interface.
"""

import logging
from typing import Literal

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

logger = logging.getLogger(__name__)

ModelName = Literal["simple_cnn", "mobilenet_v2"]


# ─── Baseline: SimpleCNN ──────────────────────────────────────────────────────

class SimpleCNN(nn.Module):
    """
    3-block convolutional network trained from scratch.

    Architecture:
        Conv(3→32) + BN + ReLU + MaxPool  → 112×112
        Conv(32→64) + BN + ReLU + MaxPool → 56×56
        Conv(64→128) + BN + ReLU + MaxPool → 28×28
        AdaptiveAvgPool → 4×4
        FC(2048→256) → Dropout → FC(256→num_classes)

    Chosen as a baseline because it:
      - Has no pretrained dependency
      - Trains quickly on CPU
      - Establishes a meaningful lower bound for comparison
    """

    def __init__(self, num_classes: int = 2, dropout: float = 0.4):
        super().__init__()
        self.features = nn.Sequential(
            # Block 1
            nn.Conv2d(3, 32, kernel_size=3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2),           # → 112×112

            # Block 2
            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2),           # → 56×56

            # Block 3
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2),           # → 28×28
        )
        self.classifier = nn.Sequential(
            nn.AdaptiveAvgPool2d((4, 4)),  # fixed spatial size regardless of input
            nn.Flatten(),
            nn.Linear(128 * 4 * 4, 256),
            nn.ReLU(inplace=True),
            nn.Dropout(dropout),
            nn.Linear(256, num_classes),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.classifier(self.features(x))


# ─── Strong model: MobileNetV2 ────────────────────────────────────────────────

class MobileNetV2Classifier(nn.Module):
    """
    Fine-tuned MobileNetV2 with a custom classification head.

    Transfer learning strategy:
      1. Load ImageNet pretrained weights.
      2. Freeze the backbone – only the new classifier head is trained.
      3. After `unfreeze_after_epoch` epochs, unfreeze the backbone and
         fine-tune end-to-end with a reduced LR (10× lower).

    Why MobileNetV2:
      - Proven accuracy on medical imaging tasks
      - Designed for constrained compute (depthwise separable convolutions)
      - ~3.4M params vs ResNet-50's ~25M – suitable for CPU / laptop GPU
    """

    def __init__(
        self,
        num_classes: int = 2,
        pretrained: bool = True,
        freeze_backbone: bool = True,
    ):
        super().__init__()
        weights = models.MobileNet_V2_Weights.IMAGENET1K_V1 if pretrained else None
        backbone = models.mobilenet_v2(weights=weights)

        # Swap out the classifier head
        in_features = backbone.classifier[1].in_features
        backbone.classifier = nn.Sequential(
            nn.Dropout(0.3),
            nn.Linear(in_features, 256),
            nn.ReLU(inplace=True),
            nn.Dropout(0.2),
            nn.Linear(256, num_classes),
        )
        self.backbone = backbone

        if freeze_backbone:
            self._freeze_backbone()

    def _freeze_backbone(self) -> None:
        """Freeze all feature extraction layers; only the head trains."""
        for param in self.backbone.features.parameters():
            param.requires_grad = False
        logger.info("Backbone frozen – training classifier head only.")

    def unfreeze_backbone(self) -> None:
        """Unfreeze all layers for full end-to-end fine-tuning."""
        for param in self.backbone.features.parameters():
            param.requires_grad = True
        logger.info("Backbone unfrozen – full fine-tuning enabled.")

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.backbone(x)


# ─── Factory ──────────────────────────────────────────────────────────────────

def build_model(cfg: dict) -> nn.Module:
    """
    Instantiate the model described in cfg['model'].

    Args:
        cfg: Full config dict (loaded from configs/config.yaml).

    Returns:
        Initialised nn.Module (not yet on device).
    """
    model_cfg  = cfg["model"]
    name: str  = model_cfg["name"]
    num_classes: int = model_cfg.get("num_classes", 2)

    if name == "simple_cnn":
        model = SimpleCNN(num_classes=num_classes)
        logger.info("Built SimpleCNN baseline.")
    elif name == "mobilenet_v2":
        model = MobileNetV2Classifier(
            num_classes=num_classes,
            pretrained=model_cfg.get("pretrained", True),
            freeze_backbone=model_cfg.get("freeze_backbone", True),
        )
        logger.info("Built MobileNetV2 classifier.")
    else:
        raise ValueError(
            f"Unknown model '{name}'. Valid choices: 'simple_cnn', 'mobilenet_v2'."
        )

    total     = sum(p.numel() for p in model.parameters())
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    logger.info(f"Parameters – total: {total:,} | trainable: {trainable:,}")
    return model