import torch
import torch.nn as nn

class AnomalyDetector(nn.Module):
    def __init__(self, input_dim=41, hidden_dim=64):
        """
        Autoencoder for Network Anomaly Detection.
        Input dimension 41 corresponds to NSL-KDD dataset features.
        """
        super(AnomalyDetector, self).__init__()
        
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Linear(hidden_dim // 2, hidden_dim // 4)
        )
        
        self.decoder = nn.Sequential(
            nn.Linear(hidden_dim // 4, hidden_dim // 2),
            nn.ReLU(),
            nn.Linear(hidden_dim // 2, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, input_dim)
        )

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

def detect_anomaly(model, data, threshold=0.05):
    """
    Detects anomalies based on reconstruction error.
    """
    model.eval()
    with torch.no_grad():
        reconstructed = model(data)
        mse = torch.mean((data - reconstructed) ** 2, dim=1)
        anomalies = mse > threshold
    return anomalies, mse

if __name__ == "__main__":
    model = AnomalyDetector()
    dummy_traffic = torch.randn(10, 41)
    anomalies, scores = detect_anomaly(model, dummy_traffic)
    print("Anomalies:", anomalies)
    print("Scores:", scores)