import torch.nn as nn

from src.models import register

# Named presets map cnn_preset config values to channel lists
CNN_PRESETS = {
    "micro":  [8, 16],
    "small":  [8, 16, 32],
    "medium": [16, 32, 64, 64],
    "large":  [32, 64, 128, 256],
}


# Each entry in channels builds a Conv -> BN -> ReLU -> Pool block
# the last block pools to 1×1 so the head is resolution-independent
class SimpleCNN(nn.Module):
    def __init__(self, channels=None, in_channels=3, dropout=0.0):
        super().__init__()
        if channels is None:
            channels = CNN_PRESETS["medium"]

        layers = []
        prev = in_channels
        for i, ch in enumerate(channels):
            layers += [nn.Conv2d(prev, ch, 3, padding=1), nn.BatchNorm2d(ch), nn.ReLU()]
            if i < len(channels) - 1:
                layers.append(nn.MaxPool2d(2))
            else:
                layers.append(nn.AdaptiveAvgPool2d(1))
            prev = ch
        self.features = nn.Sequential(*layers)

        head = []
        if dropout > 0:
            head.append(nn.Dropout(dropout))
        head.append(nn.Linear(channels[-1], 1))
        self.classifier = nn.Sequential(*head)

    def forward(self, x):
        return self.classifier(self.features(x).flatten(1))


# Resolves cnn_channels > cnn_preset > "medium" fallback
def build(cfg: dict) -> nn.Module:
    channels = cfg.get("cnn_channels") or CNN_PRESETS.get(cfg.get("cnn_preset", "medium"), CNN_PRESETS["medium"])
    return SimpleCNN(channels=channels, dropout=cfg.get("dropout", 0.0))


register("simple_cnn", build)
