import torch
import torch.nn as nn

class Sadim54MStable(nn.Module):
    def __init__(self, input_dim=5):
        super().__init__()
        # The exact architecture of SADIM 54M
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, 1024), nn.ReLU(), nn.LayerNorm(1024),
            nn.Linear(1024, 2048), nn.ReLU(), nn.LayerNorm(2048),
            nn.Linear(2048, 512)
        )
        self.decoder = nn.Sequential(
            nn.Linear(512, 2048), nn.ReLU(),
            nn.Linear(2048, 1024), nn.ReLU(),
            nn.Linear(1024, input_dim)
        )

    def forward(self, x):
        return self.decoder(self.encoder(x))

    def get_anomaly_score(self, x):
        '''Calculate MSE loss for anomaly detection'''
        with torch.no_grad():
            reconstructed = self.forward(x)
            error = torch.mean((x - reconstructed)**2, dim=1)
        return error