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


class DenseNet121(nn.Module):
    """DenseNet121 classifier head for the demo app."""

    def __init__(self, num_classes: int, dropout_rate: float = 0.25):
        super().__init__()
        self.densenet_model = models.densenet121(weights=None)
        self.dropout = nn.Dropout(p=dropout_rate)
        self.classifier = nn.Linear(1024, num_classes)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if x.dim() == 3:
            x = x.unsqueeze(0)

        features = self.densenet_model.features(x)
        features = F.gelu(features)
        pooled = F.adaptive_avg_pool2d(features, (1, 1))
        flattened = torch.flatten(pooled, 1)

        if self.training:
            flattened = self.dropout(flattened)

        return self.classifier(flattened)