"""DenseNet121 with an MC-Dropout head for chest X-ray classification.

The backbone is the CheXNet-standard DenseNet121. We replace the classifier
with a Dropout -> Linear head. At inference we keep the Dropout layers active
(see `enable_mc_dropout`) so repeated forward passes give a distribution over
predictions -- the basis for our uncertainty estimate.
"""
import torch
import torch.nn as nn
from torchvision import models


def build_model(num_classes: int = 1, dropout_p: float = 0.3,
                pretrained: bool = True) -> nn.Module:
    """DenseNet121 with an MC-Dropout classifier head.

    num_classes=1 -> binary task, use with BCEWithLogitsLoss (sigmoid output).
    num_classes=N (N>1, multi-label e.g. ChestX-ray14) -> also BCEWithLogitsLoss.
    """
    weights = models.DenseNet121_Weights.IMAGENET1K_V1 if pretrained else None
    net = models.densenet121(weights=weights)
    in_features = net.classifier.in_features
    net.classifier = nn.Sequential(
        nn.Dropout(p=dropout_p),
        nn.Linear(in_features, num_classes),
    )
    return net


def enable_mc_dropout(model: nn.Module) -> None:
    """Put the model in eval mode but re-activate Dropout layers.

    This is what makes MC Dropout work: BatchNorm etc. stay in eval mode
    (using running stats), while Dropout keeps sampling.
    """
    model.eval()
    for m in model.modules():
        if isinstance(m, nn.Dropout):
            m.train()


@torch.no_grad()
def mc_predict(model: nn.Module, x: torch.Tensor, n_passes: int = 30):
    """Run n stochastic forward passes.

    Returns:
        mean_prob:  (B, C) predictive probability (mean over passes)
        uncertainty:(B, C) predictive std over passes (epistemic signal)
    """
    enable_mc_dropout(model)
    probs = []
    for _ in range(n_passes):
        logits = model(x)
        probs.append(torch.sigmoid(logits))
    probs = torch.stack(probs, dim=0)          # (T, B, C)
    mean_prob = probs.mean(dim=0)              # (B, C)
    uncertainty = probs.std(dim=0)            # (B, C)
    return mean_prob, uncertainty