import torch
import torch.nn as nn
import numpy as np
import joblib
import time

MODEL_PATH = "/app/models/lstm_trained_L8_finetuned_FINAL_12pct.pt"
NORM_STATS_PATH = "/app/models/ppi_norm_stats_40k.pkl"
THRESHOLD = 0.0155

=class StackedLSTM_AE(nn.Module):
    def __init__(self, n_features=3, seq_len=30, hidden=32, latent=16, dropout_p=0.2):
        super().__init__()
        self.seq_len = seq_len
        self.enc_lstm = nn.LSTM(n_features, hidden, num_layers=2, batch_first=True, dropout=dropout_p)
        self.enc_norm = nn.LayerNorm(hidden)
        self.to_latent = nn.Linear(hidden, latent)
        self.from_latent = nn.Linear(latent, hidden)
        self.dec_lstm = nn.LSTM(hidden, hidden, num_layers=2, batch_first=True, dropout=dropout_p)
        self.dec_norm = nn.LayerNorm(hidden)
        self.out = nn.Linear(hidden, n_features)

    def forward(self, x):
        _, (h, _) = self.enc_lstm(x)
        h_last = self.enc_norm(h[-1])
        z = self.to_latent(h_last)
        d = self.from_latent(z).unsqueeze(1).repeat(1, self.seq_len, 1)
        d, _ = self.dec_lstm(d)
        d = self.dec_norm(d)
        return self.out(d)

=device = torch.device("cpu")

model = StackedLSTM_AE(seq_len=8).to(device)
model.load_state_dict(torch.load(MODEL_PATH, map_location=device))
model.eval()

norm_stats = joblib.load(NORM_STATS_PATH)
MEAN_TIMING, STD_TIMING = norm_stats[0]
MEAN_SIZE, STD_SIZE = norm_stats[2]

print(" Inference Engine: model and scaler loaded successfully")

def preprocess_ppi_matrix(raw_matrix):
    processed = raw_matrix.copy()

    processed[:, 0] = np.log1p(processed[:, 0])
    processed[:, 2] = np.log1p(processed[:, 2])

    
    processed[:, 0] = (processed[:, 0] - MEAN_TIMING) / STD_TIMING
    processed[:, 2] = (processed[:, 2] - MEAN_SIZE) / STD_SIZE

    return processed

def classify_flow(raw_ppi_matrix):
    start_time = time.perf_counter()

    processed = preprocess_ppi_matrix(raw_ppi_matrix)

    input_tensor = torch.tensor(processed, dtype=torch.float32).unsqueeze(0).to(device)

    with torch.no_grad():
        reconstructed = model(input_tensor)
        error = torch.mean((reconstructed - input_tensor) ** 2).item()

    decision = "MALICIOUS" if error > THRESHOLD else "BENIGN"
    latency_ms = (time.perf_counter() - start_time) * 1000

    return decision, error, latency_ms

