import torch.nn as nn


class BackboneWithFC(nn.Module):
    """Frozen ResNet18 backbone + trainable FC classifier head."""

    def __init__(self, backbone: nn.Module, num_classes: int, dropout: float = 0.4, fc_dim: int = 256):
        super().__init__()
        self.backbone = backbone
        self.classifier = nn.Sequential(
            nn.Dropout(dropout),
            nn.Linear(512, fc_dim),
            nn.ReLU(inplace=True),
            nn.Dropout(dropout),
            nn.Linear(fc_dim, num_classes),
        )

    def forward(self, x):
        return self.classifier(self.backbone(x))


class SimpleCNN(nn.Module):
    def __init__(
        self,
        num_classes: int,
        num_conv_blocks: int = 3,
        base_filters: int = 32,
        kernel_size: int = 3,
        use_batchnorm: bool = True,
        dropout: float = 0.4,
        fc_dim: int = 256,
    ):
        super().__init__()

        padding = kernel_size // 2
        layers = []
        in_channels = 3

        for i in range(num_conv_blocks):
            out_channels = min(base_filters * (2 ** i), 512)
            layers.append(nn.Conv2d(in_channels, out_channels, kernel_size, padding=padding))
            if use_batchnorm:
                layers.append(nn.BatchNorm2d(out_channels))
            layers.append(nn.ReLU(inplace=True))
            layers.append(nn.MaxPool2d(2, 2))
            in_channels = out_channels

        self.features = nn.Sequential(*layers)
        self.pool = nn.AdaptiveAvgPool2d(1)

        self.classifier = nn.Sequential(
            nn.Dropout(dropout),
            nn.Linear(in_channels, fc_dim),
            nn.ReLU(inplace=True),
            nn.Dropout(dropout),
            nn.Linear(fc_dim, num_classes),
        )

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