src/train_balanced.py

"""
train_balanced.py
==================
Trains the two balanced-training ablation variants on top of the
splits produced by balance_data.py:

  --mode balanced_recon
      Same unsupervised reconstruction-error training as train.py's
      normal-only model -- MSE loss, no labels used -- but fed the
      balanced normal+attack training set instead of normal-only.
      Tests whether reconstruction error still separates the classes
      when the model is trained to reconstruct BOTH well.

  --mode balanced_supervised
      Labels are used directly. The encoder's bottleneck vector feeds
      a small classifier head trained with binary cross-entropy,
      combined with the reconstruction loss (weighted). This is
      structurally closest to the Elsayed et al. (LSTM-AE + OC-SVM
      on the latent representation) precedent already cited in
      Literature.tex, except the classifier head is trained jointly
      rather than as a separate downstream SVM step.

Both variants share the same encoder/decoder architecture, hidden
size, window size, and seed as the canonical normal-only model in
train.py -- only the training data and (for balanced_supervised) the
loss function differ. Run balance_data.py first to generate the
required *_balanced_{run_id}.npy arrays.

Usage
-----
python balance_data.py --dataset csic2010 --window 5
python train_balanced.py --dataset csic2010 --mode balanced_recon
python train_balanced.py --dataset csic2010 --mode balanced_supervised

Author : K.A.D.S.D. Kandanaarachchi (2020/ICT/19)
Project: Detecting Anomalous REST API Traffic -- IT4216 (balanced-
         training ablation)
"""

import argparse
import json
import logging
import random
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset

from model import (
    LSTMAutoencoder,
    build_model_cicids2018,
    build_model_csic2010,
    build_model_unsw,
)

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s  %(levelname)s  %(message)s",
    datefmt="%H:%M:%S",
)
log = logging.getLogger(__name__)


def set_seed(seed=42):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.backends.cudnn.deterministic = True


class LSTMAutoencoderWithHead(nn.Module):
    """
    Wraps the existing LSTMAutoencoder (encoder+decoder unchanged) and
    adds a small linear classifier head on top of the encoder's
    bottleneck vector, for the balanced_supervised variant.

    Kept as a separate wrapper class -- rather than modifying
    model.py's LSTMAutoencoder directly -- so the canonical
    normal-only model and all existing scripts (evaluate.py,
    active_learning.py, tui.py) are completely unaffected by this
    ablation study.
    """

    def __init__(self, base_model: LSTMAutoencoder, hidden_size: int = 64):
        super().__init__()
        self.base = base_model
        self.classifier = nn.Linear(hidden_size, 1)

    def forward(self, x):
        if self.base.input_mode == "embedding":
            x_emb = self.base.embedding(x)
        else:
            x_emb = x
        context = self.base.encoder(x_emb)          # bottleneck vector
        reconstruction = self.base.decoder(context)
        logit = self.classifier(context).squeeze(-1)
        return reconstruction, logit, x_emb

    def reconstruction_error(self, x):
        with torch.no_grad():
            recon, _, x_emb = self.forward(x)
            error = ((recon - x_emb) ** 2).mean(dim=(1, 2))
        return error

    def predict_proba(self, x):
        with torch.no_grad():
            _, logit, _ = self.forward(x)
            return torch.sigmoid(logit)


def _to_tensor(X: np.ndarray, dataset: str) -> torch.Tensor:
    if dataset == "csic2010":
        return torch.tensor(X, dtype=torch.long)
    X = np.nan_to_num(X, nan=0.0, posinf=0.0, neginf=0.0)
    return torch.tensor(X, dtype=torch.float32)


def train_balanced_recon(
    model, X_train, dataset, run_id, epochs, batch_size, lr, device
):
    """
    Reconstruction-only training on the balanced set -- identical
    training loop to train.py's train(), just fed balanced data
    instead of normal-only data. No labels used.
    """
    model = model.to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    criterion = nn.MSELoss()

    split = int(len(X_train) * 0.9)
    X_tr, X_val = X_train[:split], X_train[split:]
    tr_tensor, val_tensor = _to_tensor(X_tr, dataset), _to_tensor(X_val, dataset)

    tr_loader = DataLoader(TensorDataset(tr_tensor), batch_size=batch_size, shuffle=True)
    val_loader = DataLoader(TensorDataset(val_tensor), batch_size=batch_size)

    train_losses, val_losses = [], []
    best_val = float("inf")

    for epoch in range(1, epochs + 1):
        model.train()
        epoch_loss = 0.0
        for (batch,) in tr_loader:
            batch = batch.to(device)
            optimizer.zero_grad()
            recon = model(batch)
            target = model.embedding(batch).detach() if dataset == "csic2010" else batch
            loss = criterion(recon, target)
            loss.backward()
            nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            optimizer.step()
            epoch_loss += loss.item() * len(batch)
        epoch_loss /= len(X_tr)

        model.eval()
        val_loss = 0.0
        with torch.no_grad():
            for (batch,) in val_loader:
                batch = batch.to(device)
                recon = model(batch)
                target = model.embedding(batch).detach() if dataset == "csic2010" else batch
                val_loss += criterion(recon, target).item() * len(batch)
        val_loss /= len(X_val)

        train_losses.append(epoch_loss)
        val_losses.append(val_loss)
        log.info("Epoch %02d/%02d  train=%.6f  val=%.6f", epoch, epochs, epoch_loss, val_loss)

        if val_loss < best_val - 1e-6:
            best_val = val_loss
            torch.save(model.state_dict(), f"models/best_{run_id}_balanced_recon.pt")

    return train_losses, val_losses


def train_balanced_supervised(
    model, X_train, y_train, dataset, run_id, epochs, batch_size, lr, device,
    recon_weight: float = 0.5,
):
    """
    Joint reconstruction + classification training.

    Loss = recon_weight * MSE(reconstruction, input)
         + (1 - recon_weight) * BCE(classifier_logit, label)

    recon_weight=0.5 is a starting point, not a tuned value -- treat
    it as a hyperparameter to sweep if time allows. At recon_weight=0
    this degenerates to a pure sequence classifier (no autoencoding
    objective at all); at 1.0 it's identical to balanced_recon and the
    classifier head is trained but never influences the shared
    encoder weights via backprop on the classification loss.
    """
    model = model.to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    mse = nn.MSELoss()
    bce = nn.BCEWithLogitsLoss()

    split = int(len(X_train) * 0.9)
    X_tr, X_val = X_train[:split], X_train[split:]
    y_tr, y_val = y_train[:split], y_train[split:]

    tr_tensor, val_tensor = _to_tensor(X_tr, dataset), _to_tensor(X_val, dataset)
    y_tr_t = torch.tensor(y_tr, dtype=torch.float32)
    y_val_t = torch.tensor(y_val, dtype=torch.float32)

    tr_loader = DataLoader(TensorDataset(tr_tensor, y_tr_t), batch_size=batch_size, shuffle=True)
    val_loader = DataLoader(TensorDataset(val_tensor, y_val_t), batch_size=batch_size)

    train_losses, val_losses = [], []
    best_val = float("inf")

    for epoch in range(1, epochs + 1):
        model.train()
        epoch_loss = 0.0
        for batch, labels in tr_loader:
            batch, labels = batch.to(device), labels.to(device)
            optimizer.zero_grad()
            recon, logit, x_emb = model(batch)
            loss = recon_weight * mse(recon, x_emb) + (1 - recon_weight) * bce(logit, labels)
            loss.backward()
            nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            optimizer.step()
            epoch_loss += loss.item() * len(batch)
        epoch_loss /= len(X_tr)

        model.eval()
        val_loss = 0.0
        correct = 0
        with torch.no_grad():
            for batch, labels in val_loader:
                batch, labels = batch.to(device), labels.to(device)
                recon, logit, x_emb = model(batch)
                loss = recon_weight * mse(recon, x_emb) + (1 - recon_weight) * bce(logit, labels)
                val_loss += loss.item() * len(batch)
                pred = (torch.sigmoid(logit) > 0.5).float()
                correct += (pred == labels).sum().item()
        val_loss /= len(X_val)
        val_acc = correct / len(X_val)

        train_losses.append(epoch_loss)
        val_losses.append(val_loss)
        log.info(
            "Epoch %02d/%02d  train=%.6f  val=%.6f  val_acc=%.4f",
            epoch, epochs, epoch_loss, val_loss, val_acc,
        )

        if val_loss < best_val - 1e-6:
            best_val = val_loss
            torch.save(model.state_dict(), f"models/best_{run_id}_balanced_supervised.pt")

    return train_losses, val_losses


def main():
    set_seed(42)
    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset", required=True, choices=["csic2010", "cicids2018", "unsw"])
    parser.add_argument("--mode", required=True, choices=["balanced_recon", "balanced_supervised"])
    parser.add_argument("--window", type=int, default=5)
    parser.add_argument("--epochs", type=int, default=30)
    parser.add_argument("--batch_size", type=int, default=256)
    parser.add_argument("--lr", type=float, default=1e-3)
    parser.add_argument("--recon_weight", type=float, default=0.5)
    args = parser.parse_args()

    device = "cuda" if torch.cuda.is_available() else "cpu"
    data_dir = Path("data/processed")
    model_dir = Path("models")
    model_dir.mkdir(exist_ok=True)
    run_id = f"{args.dataset}_w{args.window}"

    X_train = np.load(data_dir / f"X_train_balanced_{run_id}.npy")
    y_train = np.load(data_dir / f"y_train_balanced_{run_id}.npy")
    log.info(
        "Loaded balanced training set: %s  (normal=%d attack=%d)",
        X_train.shape, int((y_train == 0).sum()), int((y_train == 1).sum()),
    )

    if args.dataset == "csic2010":
        vocab_data = json.load(open(data_dir / f"vocab_{run_id}.json"))
        vocab = vocab_data.get("vocab", vocab_data)
        base_model = build_model_csic2010(vocab_size=len(vocab), seq_len=args.window)
    elif args.dataset == "cicids2018":
        base_model = build_model_cicids2018(n_features=X_train.shape[2], seq_len=args.window)
    else:
        base_model = build_model_unsw(n_features=X_train.shape[2], seq_len=args.window)

    if args.mode == "balanced_recon":
        train_losses, val_losses = train_balanced_recon(
            base_model, X_train, args.dataset, run_id,
            args.epochs, args.batch_size, args.lr, device,
        )
        suffix = "balanced_recon"
    else:
        model = LSTMAutoencoderWithHead(base_model, hidden_size=64)
        train_losses, val_losses = train_balanced_supervised(
            model, X_train, y_train, args.dataset, run_id,
            args.epochs, args.batch_size, args.lr, device,
            recon_weight=args.recon_weight,
        )
        suffix = "balanced_supervised"

    history = {
        "dataset": args.dataset,
        "window": args.window,
        "mode": args.mode,
        "epochs": list(range(1, len(train_losses) + 1)),
        "train_loss": train_losses,
        "val_loss": val_losses,
    }
    history_path = model_dir / f"history_{run_id}_{suffix}.json"
    with open(history_path, "w") as f:
        json.dump(history, f, indent=2)
    log.info("History saved -> %s", history_path)
    log.info("Best model saved -> models/best_%s_%s.pt", run_id, suffix)
    log.info(
        "NEXT: evaluate against data/processed/X_test_balanced_%s.npy "
        "(NOT the canonical X_test_%s.npy -- different, smaller test "
        "population; see balance_data.py's printed summary).",
        run_id, run_id,
    )


if __name__ == "__main__":
    main()