import torch
import torch.nn as nn
import torch.nn.functional as F


class ConvBlock(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, dropout_rate: float = 0.25):
        super().__init__()

        self.block = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),

            nn.Conv2d(out_channels, out_channels,kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),

            nn.MaxPool2d(kernel_size=2, stride=2),

            nn.Dropout2d(p=dropout_rate),
        )

    def forward(self, x):
        return self.block(x)


class SaraCNN(nn.Module):
    def __init__(self, num_classes: int = 6):
        super().__init__()

        self.features = nn.Sequential(
            ConvBlock(3,   32,  dropout_rate=0.25),   
            ConvBlock(32,  64,  dropout_rate=0.25), 
            ConvBlock(64,  128, dropout_rate=0.25),   
        )

        self.classifier = nn.Sequential(
            nn.Flatten(),

            nn.Linear(128 * 18 * 18, 256),
            nn.BatchNorm1d(256),
            nn.ReLU(inplace=True),
            nn.Dropout(p=0.5),

            nn.Linear(256, num_classes),  
        )

       
        self._init_weights()

    def _init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
            elif isinstance(m, nn.Linear):
                nn.init.xavier_uniform_(m.weight)
                nn.init.zeros_(m.bias)
            elif isinstance(m, (nn.BatchNorm2d, nn.BatchNorm1d)):
                nn.init.ones_(m.weight)
                nn.init.zeros_(m.bias)

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



if __name__ == "__main__":
    model = SaraCNN(num_classes=6)
    dummy = torch.randn(4, 3, 150, 150)  
    out   = model(dummy)
    print("SaraCNN output shape:", out.shape)  

    total = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"Trainable parameters: {total:,}")