src/balance_data.py

"""
balance_data.py
================
Builds a SEPARATE balanced train/test split for the balanced-training
ablation study, without touching the existing normal-only arrays used
by Chapter 4's canonical results.

Why this exists
----------------
preprocessing.py currently puts 100% of attack sessions into the test
set and 0% into train (train = normal-only, by design, for the
unsupervised model). To train a model on a roughly-balanced mix of
normal + attack sessions, some attack sessions must move into a
training split -- which necessarily means this ablation study uses a
DIFFERENT, smaller test set than Table tab:lstm_results. That's fine,
but it must be reported as such, not silently compared 1:1 against
the canonical normal-only numbers.

What this script does
----------------------
1. Loads the existing X_test_{dataset}_w{window}.npy / y_test_*.npy
   (the only place attack sessions currently exist on disk).
2. Splits the ATTACK sessions 80/20 into attack_train / attack_test,
   mirroring how normal traffic is already split in preprocessing.py.
3. Builds a balanced training set: attack_train sessions + an equal
   (or ratio-controlled) number of normal sessions, drawn from the
   existing normal-only X_train (for csic2010/cicids2018/unsw this is
   the *_train_{run_id}.npy already on disk).
4. Builds a new, smaller balanced test set: attack_test + a matching
   slice of held-out normal sessions (drawn from whatever normal
   sessions were NOT used in the balanced training set, so nothing
   leaks between train/test).
5. Saves everything under a "_balanced" tag so the original arrays
   are never modified:
       X_train_balanced_{run_id}.npy
       y_train_balanced_{run_id}.npy   (0=normal, 1=attack -- needed
                                         for the supervised variant;
                                         harmless/unused for the
                                         reconstruction-only variant)
       X_test_balanced_{run_id}.npy
       y_test_balanced_{run_id}.npy

Usage
-----
python balance_data.py --dataset csic2010 --window 5
python balance_data.py --dataset cicids2018 --window 5 --normal-ratio 1.0
python balance_data.py --dataset unsw --window 5

--normal-ratio controls the normal:attack ratio in the TRAINING set
(1.0 = strictly balanced 1:1). CIC-IDS2018 only has 924 attack test
sessions to draw from -- see the printed summary for exactly how many
sessions end up in each split before you commit to a training run.

Author : K.A.D.S.D. Kandanaarachchi (2020/ICT/19)
Project: Detecting Anomalous REST API Traffic -- IT4216 (balanced-
         training ablation, added after supervisor/co-supervisor
         discussion on normal-only training methodology)
"""

import argparse
import logging
from pathlib import Path

import numpy as np

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

SEED = 42


def build_balanced_split(
    dataset: str,
    window: int,
    data_dir: Path,
    attack_train_ratio: float = 0.8,
    normal_ratio: float = 1.0,
) -> dict:
    """
    Parameters
    ----------
    attack_train_ratio : fraction of attack sessions moved into the
        balanced training set (rest held out as balanced test attacks).
        0.8 mirrors the normal-traffic train/test split already used
        elsewhere in this pipeline.
    normal_ratio : ratio of normal:attack sessions in the TRAINING set.
        1.0 = perfectly balanced (equal counts). Set >1.0 if you decide
        with your co-supervisor that strict 1:1 throws away too much
        signal for a given dataset (e.g. CIC-IDS2018).

    Returns a dict summary (also printed) so you can sanity-check
    counts before committing to a full training run.
    """
    rng = np.random.default_rng(SEED)
    run_id = f"{dataset}_w{window}"

    X_test_all = np.load(data_dir / f"X_test_{run_id}.npy")
    y_test_all = np.load(data_dir / f"y_test_{run_id}.npy")
    X_train_normal_pool = np.load(data_dir / f"X_train_{run_id}.npy")

    attack_idx = np.where(y_test_all == 1)[0]
    normal_test_idx = np.where(y_test_all == 0)[0]

    n_attack_total = len(attack_idx)
    log.info("%s: %d attack sessions available in existing test set", dataset, n_attack_total)

    # Shuffle attacks, split 80/20 train/test (mirrors normal split elsewhere)
    shuffled_attack_idx = rng.permutation(attack_idx)
    n_attack_train = int(round(n_attack_total * attack_train_ratio))
    attack_train_idx = shuffled_attack_idx[:n_attack_train]
    attack_test_idx = shuffled_attack_idx[n_attack_train:]

    X_attack_train = X_test_all[attack_train_idx]
    X_attack_test = X_test_all[attack_test_idx]

    # How many normal sessions needed for the training set, given normal_ratio
    n_normal_train_needed = int(round(len(attack_train_idx) * normal_ratio))
    if n_normal_train_needed > len(X_train_normal_pool):
        log.warning(
            "  Requested %d normal training sessions but only %d available "
            "in the normal-only pool -- using all of them (ratio will be "
            "lower than requested).",
            n_normal_train_needed, len(X_train_normal_pool),
        )
        n_normal_train_needed = len(X_train_normal_pool)

    normal_train_sel = rng.choice(
        len(X_train_normal_pool), n_normal_train_needed, replace=False
    )
    X_normal_train = X_train_normal_pool[normal_train_sel]

    # Remaining normal sessions (not used in training) are available for
    # the balanced test set, so nothing leaks between splits.
    remaining_normal_mask = np.ones(len(X_train_normal_pool), dtype=bool)
    remaining_normal_mask[normal_train_sel] = False
    remaining_normal_pool = X_train_normal_pool[remaining_normal_mask]

    # Also fold in the ORIGINAL held-out normal test sessions -- these
    # were never used for training under any variant, so they're safe
    # to reuse here too.
    normal_test_pool = np.concatenate(
        [remaining_normal_pool, X_test_all[normal_test_idx]], axis=0
    )

    n_normal_test_needed = min(len(X_attack_test) * 4, len(normal_test_pool))
    # ^ keep test set realistically imbalanced-ish (4:1) rather than
    # forcing 1:1 at test time too -- balancing is a TRAINING decision;
    # evaluation should still reflect a plausible operating distribution.
    # Adjust this multiplier with your co-supervisor if you want the
    # balanced-test evaluation itself to also be strictly 1:1.
    normal_test_sel = rng.choice(
        len(normal_test_pool), n_normal_test_needed, replace=False
    )
    X_normal_test = normal_test_pool[normal_test_sel]

    # Assemble final arrays
    X_train_balanced = np.concatenate([X_normal_train, X_attack_train], axis=0)
    y_train_balanced = np.concatenate(
        [np.zeros(len(X_normal_train), dtype=np.int32),
         np.ones(len(X_attack_train), dtype=np.int32)]
    )
    X_test_balanced = np.concatenate([X_normal_test, X_attack_test], axis=0)
    y_test_balanced = np.concatenate(
        [np.zeros(len(X_normal_test), dtype=np.int32),
         np.ones(len(X_attack_test), dtype=np.int32)]
    )

    # Shuffle train (test order doesn't matter for evaluation)
    perm = rng.permutation(len(X_train_balanced))
    X_train_balanced = X_train_balanced[perm]
    y_train_balanced = y_train_balanced[perm]

    # Save under a distinct tag -- originals are never touched
    np.save(data_dir / f"X_train_balanced_{run_id}.npy", X_train_balanced)
    np.save(data_dir / f"y_train_balanced_{run_id}.npy", y_train_balanced)
    np.save(data_dir / f"X_test_balanced_{run_id}.npy", X_test_balanced)
    np.save(data_dir / f"y_test_balanced_{run_id}.npy", y_test_balanced)

    summary = {
        "dataset": dataset,
        "window": window,
        "train_normal": len(X_normal_train),
        "train_attack": len(X_attack_train),
        "train_ratio_normal_to_attack": round(len(X_normal_train) / max(len(X_attack_train), 1), 3),
        "test_normal": len(X_normal_test),
        "test_attack": len(X_attack_test),
    }
    log.info("=" * 55)
    log.info("BALANCED SPLIT SUMMARY -- %s", dataset.upper())
    log.info("=" * 55)
    for k, v in summary.items():
        log.info("  %-30s %s", k, v)
    log.info(
        "  NOTE: this test set (%d normal / %d attack) is DIFFERENT from "
        "the canonical Chapter 4 test set (%d normal / %d attack). "
        "Report balanced-training results against THIS test set only -- "
        "do not compare its metrics directly to Table tab:lstm_results "
        "without re-stating that the test populations differ.",
        summary["test_normal"], summary["test_attack"],
        len(normal_test_idx), n_attack_total,
    )

    if dataset == "cicids2018":
        log.info(
            "  SCOPE NOTE: attack sessions here are drawn only from the "
            "928 web/brute-force attack flows in the two capture days "
            "(02-22/23-2018) this study uses -- NOT the full CIC-IDS2018 "
            "dataset, which also includes DDoS, Botnet, Infiltration, and "
            "other volumetric/network-layer attack days outside this "
            "study's stated scope of sequence-dependent API-logic attacks "
            "(Section 1.5). This is why balancing this dataset requires "
            "aggressive normal undersampling (or SMOTE) rather than more "
            "attack volume -- pulling in the other capture days would "
            "ease balancing but would silently broaden the thesis's scope."
        )

    return summary


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset", required=True, choices=["csic2010", "cicids2018", "unsw"])
    parser.add_argument("--window", type=int, default=5)
    parser.add_argument("--attack-train-ratio", type=float, default=0.8)
    parser.add_argument(
        "--normal-ratio", type=float, default=1.0,
        help="normal:attack ratio in the TRAINING set (1.0 = balanced)",
    )
    parser.add_argument("--data-dir", default="data/processed")
    args = parser.parse_args()

    build_balanced_split(
        dataset=args.dataset,
        window=args.window,
        data_dir=Path(args.data_dir),
        attack_train_ratio=args.attack_train_ratio,
        normal_ratio=args.normal_ratio,
    )


if __name__ == "__main__":
    main()