import torch
import torch.nn as nn

class FastSCNNWithNorm(nn.Module):
    def __init__(self, num_classes, mean=(123.675,116.28,103.53), std=(58.395,57.12,57.375), to_rgb=True):
        super().__init__()
        self.backbone = FastSCNN(num_classes=num_classes)
        # registra constantes como buffers (vão pro ONNX)
        m = torch.tensor(mean).view(1,3,1,1)
        s = torch.tensor(std).view(1,3,1,1)
        self.register_buffer("mean", m, persistent=False)
        self.register_buffer("std",  s, persistent=False)
        self.to_rgb = to_rgb  # se tua câmera entregar BGR, podemos inverter canais aqui

    def forward(self, x):
        # x chega como FP16 (convertido pelo NN node), mas garantimos float32 pra estabilidade das consts
        x = x.float()
        if self.to_rgb:
            # se a ColorCamera estiver em BGR, inverte canais aqui (BGR->RGB)
            x = x[:, [2,1,0], :, :]
        x = (x - self.mean) / self.std
        return self.backbone(x)

class FastSCNN(nn.Module):
    def __init__(self, num_classes):
        super().__init__()
        self.down1 = nn.Sequential(
            nn.Conv2d(3, 32, 3, stride=2, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True)
        )
        self.down2 = nn.Sequential(
            nn.Conv2d(32, 48, 3, stride=2, padding=1),
            nn.BatchNorm2d(48),
            nn.ReLU(inplace=True)
        )
        self.down3 = nn.Sequential(
            nn.Conv2d(48, 64, 3, stride=2, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True)
        )
        self.classifier = nn.Sequential(
            nn.Conv2d(64, num_classes, 1)
        )
        # Upsampling substituído por ConvTranspose2d
        self.up1 = nn.ConvTranspose2d(num_classes, num_classes, kernel_size=2, stride=2)
        self.up2 = nn.ConvTranspose2d(num_classes, num_classes, kernel_size=2, stride=2)
        self.up3 = nn.ConvTranspose2d(num_classes, num_classes, kernel_size=2, stride=2)

    def forward(self, x):
        x = self.down1(x)
        x = self.down2(x)
        x = self.down3(x)
        x = self.classifier(x)
        x = self.up1(x)
        x = self.up2(x)
        x = self.up3(x)
        return x