main.py

from preprocessing import load_data, create_windows, balance_data
from train import prepare_data, flatten_data
from models import build_lstm, build_transformer, build_autoencoder, build_random_forest, build_isolation_forest
from evaluate import evaluate_lstm, evaluate_autoencoder, evaluate_hybrid

from sklearn.metrics import roc_auc_score
import os

os.makedirs("results", exist_ok=True)

# -------------------- LOAD DATA --------------------
file_path = "data/Veremi_final_dataset.csv"
df = load_data(file_path)

windows, labels = create_windows(df)
windows, labels = balance_data(windows, labels)

X_train, X_test, y_train, y_test = prepare_data(windows, labels)

# -------------------- LSTM --------------------
print("\n=== LSTM TRAINING ===")
model = build_lstm()
model.fit(X_train, y_train, epochs=20, batch_size=32, validation_split=0.1)
lstm_probs, lstm_auc = evaluate_lstm(model, X_test, y_test)

# -------------------- TRANSFORMER --------------------
print("\n=== TRANSFORMER TRAINING ===")
transformer = build_transformer()
transformer.fit(X_train, y_train, epochs=20, batch_size=32, validation_split=0.1)
transformer_probs, transformer_auc = evaluate_lstm(transformer, X_test, y_test)
print(f"Transformer ROC-AUC: {transformer_auc:.4f}")

# -------------------- BASELINES --------------------
print("\n=== BASELINES ===")

X_train_flat = flatten_data(X_train)
X_test_flat = flatten_data(X_test)

rf = build_random_forest()
rf.fit(X_train_flat, y_train)
rf_probs = rf.predict_proba(X_test_flat)[:, 1]
rf_auc = roc_auc_score(y_test, rf_probs)
print("Random Forest ROC-AUC:", rf_auc)

iso = build_isolation_forest()
iso.fit(X_train_flat)
iso_scores = -iso.decision_function(X_test_flat)
iso_auc = roc_auc_score(y_test, iso_scores)
print("Isolation Forest ROC-AUC:", iso_auc)

# -------------------- AUTOENCODER --------------------
print("\n=== AUTOENCODER TRAINING ===")
autoencoder = build_autoencoder()
X_normal = X_train[y_train == 0]
autoencoder.fit(X_normal, X_normal, epochs=20, batch_size=32, validation_split=0.1)
mse_scores, ae_auc = evaluate_autoencoder(autoencoder, X_test, y_test)

# -------------------- HYBRID --------------------
print("\n=== HYBRID MODEL ===")
hybrid_auc = evaluate_hybrid(
    model, transformer, autoencoder, X_test, y_test, use_meta_learner=True
)

# -------------------- SUMMARY --------------------
print("\n=== FINAL SUMMARY ===")
print(f"LSTM ROC-AUC:             {lstm_auc:.4f}")
print(f"Transformer ROC-AUC:      {transformer_auc:.4f}")
print(f"Random Forest ROC-AUC:    {rf_auc:.4f}")
print(f"Isolation Forest ROC-AUC: {iso_auc:.4f}")
print(f"Autoencoder ROC-AUC:      {ae_auc:.4f}")
print(f"Hybrid ROC-AUC:           {hybrid_auc:.4f}")