import torch
import torch.nn as nn

class FaceDetectionMultiBoxCNN(nn.Module):
    def __init__(self, S=20, B=3):
        super().__init__()
        self.S = S
        self.B = B
        self.output_size = S * S * B * 5  # (x, y, w, h, confidence)

        self.features = nn.Sequential(
            nn.Conv2d(3, 32, 3, 1, 1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(2),  # 320->160

            nn.Conv2d(32, 64, 3, 1, 1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(2),  # 160->80

            nn.Conv2d(64, 128, 3, 1, 1),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.MaxPool2d(2),  # 80->40

            nn.Conv2d(128, 256, 3, 1, 1),
            nn.BatchNorm2d(256),
            nn.ReLU(),
            nn.MaxPool2d(2),  # 40->20

            nn.Conv2d(256, 512, 3, 1, 1),
            nn.BatchNorm2d(512),
            nn.ReLU(),
            nn.MaxPool2d(2),  # 20->10
        )

        # Calculate flatten size dynamically
        dummy = torch.zeros(1, 3, 320, 320)
        with torch.no_grad():
            dummy_out = self.features(dummy)
            flatten_size = dummy_out.view(1, -1).shape[1]

        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(flatten_size, 1024),
            nn.ReLU(),
            nn.Linear(1024, self.output_size)
        )

    def forward(self, x):
        x = self.features(x)
        x = self.classifier(x)
        return x.view(-1, self.S, self.S, self.B * 5)