src/visualise.py

"""
visualise.py
============
Fully data-driven thesis figure generator.

Reads from (filenames are window-suffixed, e.g. "_w5", "_w7"):
  models/history_{dataset}_w{window}.json     — per-epoch training curves
  models/threshold_{dataset}_w{window}.json   — detection thresholds
  results/evaluation_{dataset}_w{window}.json — model comparison metrics
  results/active_learning_{dataset}_w{window}.json — AL iteration results
  results/errors_{dataset}_w{window}.npy      — real per-sample reconstruction errors
  data/processed/y_test_{dataset}_w{window}.npy   — label counts

Which window to load per dataset is set via WINDOW_BY_DATASET below.
Change it and re-run to compare results across different window sizes
without re-running the whole pipeline.

Dataset record counts (normal/attack totals) are fixed constants
from the original dataset papers — not model outputs.

Usage
-----
cd ~/api-anomaly-detection
python src/visualise.py

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

import json
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np

# ── Style ─────────────────────────────────────────────────────────────────────
plt.rcParams.update(
    {
        "font.family": "DejaVu Sans",
        "font.size": 17,
        "axes.titlesize": 20,
        "axes.titleweight": "bold",
        "axes.labelsize": 17,
        "xtick.labelsize": 15,
        "ytick.labelsize": 15,
        "legend.fontsize": 15,
        "axes.spines.top": False,
        "axes.spines.right": False,
        "axes.grid": True,
        "grid.alpha": 0.3,
        "grid.linestyle": "--",
        "figure.dpi": 150,
        "savefig.dpi": 300,
        "savefig.bbox": "tight",
        "savefig.facecolor": "white",
    }
)

TEAL = "#028090"
NAVY = "#0D1B2A"
MINT = "#02C39A"
CORAL = "#D85A30"
AMBER = "#F59E0B"
PURPLE = "#7C3AED"
GRAY = "#94A3B8"

BASE = Path(__file__).parent.parent
RES_DIR = BASE / "results"
DATA_DIR = BASE / "data" / "processed"
MDL_DIR = BASE / "models"
FIG_DIR = RES_DIR / "figures"
FIG_DIR.mkdir(parents=True, exist_ok=True)

# ── Fixed dataset constants (from original dataset papers) ────────────────────
DATASET_COUNTS = {
    "csic2010": {"normal": 72000, "attack": 25065, "label": "CSIC 2010"},
    "cicids2018": {
        "normal": 2096222,
        "attack": 928,
        "label": "CIC-IDS2018\n(used in this thesis)",
    },
    "unsw": {"normal": 93000, "attack": 164673, "label": "UNSW-NB15"},
}

# CIC-IDS2018 scope note: this thesis uses only the two web-attack capture
# days (02-22-2018.csv, 02-23-2018.csv) out of the full ten-day CIC-IDS2018
# release. Row counts below come from the full per-file scan documented in
# 02_cicids2018_exploration.md (Label-column counts per day), including the
# large DDoS/Bot/Infiltration days this thesis deliberately excludes per
# its stated scope of sequence-dependent, API-logic attacks (Section 1.5).
_CICIDS2018_DAILY = {
    # day: (normal, attack)
    "02-14": (667_626, 193_360 + 187_589),          # FTP/SSH Brute Force
    "02-15": (996_077, 41_508 + 10_990),             # DoS GoldenEye/Slowloris
    "02-16": (446_772, 461_912 + 139_890),           # DoS Hulk/SlowHTTPTest
    "02-20": (7_372_557, 576_191),                   # DDoS LOIC-HTTP
    "02-21": (360_833, 686_012 + 1_730),             # DDoS HOIC/LOIC-UDP
    "02-22": (1_048_213, 249 + 79 + 34),              # Web BF/XSS/SQLi <- used
    "02-23": (1_048_009, 362 + 151 + 53),             # Web BF/XSS/SQLi <- used
    "02-28": (544_200, 68_871),                       # Infiltration
    "03-01": (238_037, 93_063),                       # Infiltration
    "03-02": (762_384, 286_191),                      # Bot
}
CICIDS2018_FULL_NORMAL = sum(n for n, a in _CICIDS2018_DAILY.values())
CICIDS2018_FULL_ATTACK = sum(a for n, a in _CICIDS2018_DAILY.values())
CICIDS2018_FULL_RELEASE_ROWS = CICIDS2018_FULL_NORMAL + CICIDS2018_FULL_ATTACK
CICIDS2018_USED_ROWS = (
    DATASET_COUNTS["cicids2018"]["normal"] + DATASET_COUNTS["cicids2018"]["attack"]
)
CICIDS2018_USED_PCT = CICIDS2018_USED_ROWS / CICIDS2018_FULL_RELEASE_ROWS * 100

# A separate, fourth pie-chart entry for the full ten-day release, shown
# alongside the two-day subset actually used, so Figure 12 makes the scope
# reduction visually explicit rather than only stating it in a caption.
DATASET_COUNTS["cicids2018_full"] = {
    "normal": CICIDS2018_FULL_NORMAL,
    "attack": CICIDS2018_FULL_ATTACK,
    "label": "CIC-IDS2018\n(full 10-day release)",
}

DATASET_COLORS = {
    "csic2010": TEAL,
    "cicids2018": AMBER,
    "cicids2018_full": GRAY,
    "unsw": PURPLE,
}

# ── Window size to load per dataset ────────────────────────────────────────────
# preprocessing.py/train.py/evaluate.py now save every artifact with a
# "_w{window}" filename suffix, so multiple window sizes can coexist on
# disk. Set which window's results this notebook/script should plot for
# each dataset. Change these and re-run the figure functions to compare
# different window sizes without re-running the whole pipeline.
WINDOW_BY_DATASET = {
    "csic2010": 5,
    "cicids2018": 5,
    "unsw": 5,
}


def _run_id(dataset: str) -> str:
    """Build the window-suffixed filename tag for a dataset, e.g. 'csic2010_w5'."""
    return f"{dataset}_w{WINDOW_BY_DATASET[dataset]}"


def discover_windows(dataset: str) -> list[int]:
    """
    Find every window size that has a threshold_{dataset}_w{N}.json on
    disk, by scanning models/. Used by the window-comparison figure so
    it doesn't need WINDOW_BY_DATASET edited by hand for every window
    you've run — it just picks up whatever preprocessing/train/evaluate
    have actually produced.
    """
    pattern = f"threshold_{dataset}_w*.json"
    windows = []
    for p in MDL_DIR.glob(pattern):
        # filename: threshold_{dataset}_w{N}.json
        stem = p.stem  # threshold_{dataset}_w{N}
        tail = stem.rsplit("_w", 1)[-1]
        if tail.isdigit():
            windows.append(int(tail))
    return sorted(windows)


def load_threshold_full_for(dataset: str, window: int) -> dict:
    """Same as load_threshold_full but for an explicit window, not WINDOW_BY_DATASET."""
    path = MDL_DIR / f"threshold_{dataset}_w{window}.json"
    if not path.exists():
        return {}
    return json.load(open(path))


def load_eval_for(dataset: str, window: int) -> dict:
    """Same as load_eval but for an explicit window, not WINDOW_BY_DATASET."""
    path = RES_DIR / f"evaluation_{dataset}_w{window}.json"
    if not path.exists():
        return {}
    data = json.load(open(path))
    return {r["model"]: r for r in data["results"]}


def _grid_2x2(figsize=(14.5, 11.5), n_used=3):
    """
    Create a 2x2 grid of axes for figures that previously used a single
    wide row (1x3 for three-dataset panels, or 1x4 for four-panel
    figures). Reflowing to 2 columns x 2 rows gives each panel roughly
    double the linear size for the same total figure area, which is
    the main lever for print legibility (font/marker/line sizes are
    fixed in points, so a panel with twice the area renders everything
    proportionally larger on the page).

    n_used=3: fills axes 0,1,2 (row-major) and hides axes[3] (blank
    bottom-right cell) — for the common three-dataset case.
    n_used=4: fills all four axes — for four-panel figures.
    Returns (fig, axes_flat) where axes_flat has length 4 regardless,
    so callers can always index axes_flat[0..n_used-1].
    """
    fig, axes = plt.subplots(2, 2, figsize=figsize)
    axes_flat = axes.flatten()
    if n_used < 4:
        for ax in axes_flat[n_used:]:
            ax.axis("off")
    return fig, axes_flat


def _fig_suffix() -> str:
    """
    Filename suffix for saved figures, reflecting WINDOW_BY_DATASET.
    If all three datasets use the same window, use "_w{N}" (clean and
    readable). If they differ (comparing mixed windows across
    datasets), spell out each one so outputs never silently overwrite
    a previous run with different settings.
    """
    windows = set(WINDOW_BY_DATASET.values())
    if len(windows) == 1:
        return f"_w{windows.pop()}"
    return "_" + "-".join(f"{k}w{v}" for k, v in WINDOW_BY_DATASET.items())


# ── Disk readers ──────────────────────────────────────────────────────────────


def load_history(dataset: str) -> dict:
    path = MDL_DIR / f"history_{_run_id(dataset)}.json"
    if not path.exists():
        print(f"  WARNING: {path.name} not found")
        return {}
    return json.load(open(path))


def load_threshold(dataset: str) -> float | None:
    path = MDL_DIR / f"threshold_{_run_id(dataset)}.json"
    if not path.exists():
        return None
    return json.load(open(path))["threshold"]


def load_threshold_full(dataset: str) -> dict:
    path = MDL_DIR / f"threshold_{_run_id(dataset)}.json"
    if not path.exists():
        return {}
    return json.load(open(path))


def load_eval(dataset: str) -> dict:
    path = RES_DIR / f"evaluation_{_run_id(dataset)}.json"
    if not path.exists():
        print(f"  WARNING: {path.name} not found")
        return {}
    data = json.load(open(path))
    return {r["model"]: r for r in data["results"]}


def load_al(dataset: str) -> list:
    path = RES_DIR / f"active_learning_{_run_id(dataset)}.json"
    if not path.exists():
        return []
    return json.load(open(path))["history"]


def load_y_test(dataset: str) -> np.ndarray | None:
    path = DATA_DIR / f"y_test_{_run_id(dataset)}.npy"
    if not path.exists():
        return None
    return np.load(path)


def load_real_errors(dataset: str) -> tuple[np.ndarray, np.ndarray] | None:
    """
    Load real per-sample reconstruction errors + labels saved by
    evaluate.py (evaluate_lstm + save_errors). Returns (errors, labels)
    or None if evaluate.py hasn't been (re-)run since this feature was added.
    """
    run_id = _run_id(dataset)
    err_path = RES_DIR / f"errors_{run_id}.npy"
    lbl_path = RES_DIR / f"errors_labels_{run_id}.npy"
    if not (err_path.exists() and lbl_path.exists()):
        return None
    return np.load(err_path), np.load(lbl_path)


# ═══════════════════════════════════════════════════════════════════════════════
# FIGURE 1 — Learning curves  (from models/history_*.json)
# ═══════════════════════════════════════════════════════════════════════════════
def fig_learning_curves():
    datasets = [
        ("csic2010", "CSIC 2010", TEAL),
        ("cicids2018", "CIC-IDS2018", AMBER),
        ("unsw", "UNSW-NB15", PURPLE),
    ]

    fig, axes = _grid_2x2(n_used=3)

    for ax, (key, label, color) in zip(axes, datasets):
        h = load_history(key)
        if not h:
            ax.text(
                0.5,
                0.5,
                "history not found",
                ha="center",
                va="center",
                transform=ax.transAxes,
            )
            ax.set_title(label)
            continue

        ep = h["epochs"]
        train_loss = h["train_loss"]
        val_loss = h["val_loss"]

        ax.plot(ep, train_loss, color=color, lw=2, label="Train loss", zorder=3)
        ax.plot(
            ep,
            val_loss,
            color=color,
            lw=2,
            linestyle="--",
            alpha=0.7,
            label="Val loss",
            zorder=3,
        )

        # Mark best val epoch
        best_ep = ep[int(np.argmin(val_loss))]
        ax.axvline(
            best_ep,
            color=CORAL,
            linestyle=":",
            lw=1.5,
            label=f"Best val (ep {best_ep})",
        )

        ax.set_title(label)
        ax.set_xlabel("Epoch")
        ax.set_ylabel("MSE Loss")
        ax.legend(fontsize=14)
        ax.set_xlim(1, max(ep))

    plt.tight_layout()
    out = FIG_DIR / f"01_learning_curves{_fig_suffix()}.png"
    plt.savefig(out)
    plt.close()
    print(f"  Saved: {out.name}")


# ═══════════════════════════════════════════════════════════════════════════════
# FIGURE 2 — Error distributions
# (threshold from models/threshold_*.json, stats from evaluation JSON)
# ═══════════════════════════════════════════════════════════════════════════════
def fig_error_distributions():
    configs = [
        ("csic2010", "CSIC 2010", TEAL),
        ("cicids2018", "CIC-IDS2018", AMBER),
        ("unsw", "UNSW-NB15", PURPLE),
    ]

    fig, axes = _grid_2x2(n_used=3)

    rng = np.random.default_rng(42)

    for ax, (key, name, color) in zip(axes, configs):
        threshold = load_threshold(key)
        if threshold is None:
            ax.text(
                0.5, 0.5, "threshold not found",
                ha="center", va="center", transform=ax.transAxes,
            )
            ax.set_title(name)
            continue

        real = load_real_errors(key)

        if real is not None:
            errors, labels = real
            normal_errors = errors[labels == 0]
            attack_errors = errors[labels == 1]
            is_simulated = False
        else:
            # Fallback: evaluate.py hasn't been re-run to save raw
            # errors yet. Simulate an approximation from mean/std so
            # the figure still renders, but say so on the plot —
            # never present this silently as measured data.
            thresh_data = load_threshold_full(key)
            mean_err = thresh_data.get("mean_error", threshold * 0.3)
            raw_std = thresh_data.get("std_error", threshold * 0.2)
            std_err = min(raw_std, threshold * 1.5)

            y = load_y_test(key)
            n_normal = int((y == 0).sum()) if y is not None else 5000
            n_attack = int((y == 1).sum()) if y is not None else 2000

            normal_errors = np.abs(rng.normal(mean_err * 0.5, std_err * 0.3, min(n_normal, 5000)))
            normal_errors = np.clip(normal_errors, 0, threshold * 0.95)
            attack_errors = np.abs(rng.normal(threshold * 4, threshold * 2, min(n_attack, 2000)))
            attack_errors = np.clip(attack_errors, threshold * 0.5, threshold * 15)
            is_simulated = True

        # Normal and attack reconstruction errors typically differ by
        # orders of magnitude (e.g. CSIC 2010: normal ~6e-4, attack up
        # to ~2.5) — a shared narrow LINEAR window cannot show both
        # without one of them falling entirely outside it (which
        # previously produced empty, all-zero-count bins and a
        # divide-by-zero warning from density normalisation). A
        # log-scale x-axis with log-spaced bins keeps both the
        # near-zero normal cluster and the heavy attack tail visible
        # on the same panel.
        eps = max(threshold * 1e-3, 1e-12)
        normal_clipped = np.clip(normal_errors, eps, None)
        attack_clipped = np.clip(attack_errors, eps, None)
        combined_max = max(normal_clipped.max(), attack_clipped.max(), threshold) * 1.1
        combined_min = min(normal_clipped.min(), attack_clipped.min(), threshold) * 0.9
        combined_min = max(combined_min, eps)
        log_bins = np.logspace(
            np.log10(combined_min), np.log10(combined_max), 60
        )

        ax.hist(
            normal_clipped, bins=log_bins, alpha=0.65, color=color,
            label="Normal", density=True, zorder=3,
        )
        ax.hist(
            attack_clipped, bins=log_bins, alpha=0.55, color=CORAL,
            label="Attack", density=True, zorder=2,
        )
        ax.axvline(
            threshold, color=NAVY, lw=2, linestyle="-.",
            label=f"τ = {threshold:.5f}", zorder=4,
        )

        ax.set_xscale("log")
        # Normal errors are typically far more tightly clustered than
        # attack errors (e.g. CSIC 2010: normal peak density ~3000+ vs
        # attack peak density <1, a >1000x difference), even though
        # both are now correctly positioned on the x-axis. On a shared
        # LINEAR y-axis the shorter distribution becomes visually flat
        # and disappears. A log y-axis keeps both shapes visible
        # simultaneously; zero-count bins are simply omitted, which is
        # standard behaviour for log-scale histograms.
        ax.set_yscale("log")
        ax.set_title(name)
        ax.set_xlabel("Reconstruction Error (MSE, log scale)")
        ax.set_ylabel("Density (log scale)")
        ax.legend(fontsize=14)

        if is_simulated:
            ax.text(
                0.98, 0.98, "SIMULATED\n(re-run evaluate.py to plot real errors)",
                transform=ax.transAxes, ha="right", va="top",
                fontsize=12.5, color=CORAL, fontweight="bold",
                bbox=dict(boxstyle="round", facecolor="white", edgecolor=CORAL, alpha=0.85),
            )

    plt.tight_layout()
    out = FIG_DIR / f"02_error_distributions{_fig_suffix()}.png"
    plt.savefig(out)
    plt.close()
    print(f"  Saved: {out.name}")


# ═══════════════════════════════════════════════════════════════════════════════
# FIGURE 3 — Model comparison CSIC 2010  (from results/evaluation_csic2010.json)
# ═══════════════════════════════════════════════════════════════════════════════
def fig_model_comparison():
    ev = load_eval("csic2010")
    if not ev:
        print("  SKIP fig 3 — evaluation_csic2010.json not found")
        return

    model_order = [
        "LSTM-Autoencoder",
        "Isolation Forest",
        "Random Forest",
        "WAF Simulation",
    ]
    short_labels = ["LSTM-AE", "Isolation\nForest", "Random\nForest", "WAF\nSim"]
    colors = [TEAL, AMBER, CORAL, GRAY]

    f1 = [ev[m]["f1"] for m in model_order]
    recall = [ev[m]["recall"] for m in model_order]
    fpr = [ev[m]["fpr"] for m in model_order]
    latency = [ev[m]["latency_ms"] for m in model_order]

    fig, axes = _grid_2x2(n_used=4)

    metrics = [
        (axes[0], f1, "F1-Score", [0, 1.05]),
        (axes[1], recall, "Recall", [0, 1.05]),
        (axes[2], fpr, "False Positive Rate", [0, max(fpr) * 1.4]),
        (axes[3], latency, "Latency (ms)", None),
    ]

    for ax, vals, label, ylim in metrics:
        bars = ax.bar(
            short_labels,
            vals,
            color=colors,
            width=0.55,
            edgecolor="white",
            linewidth=1.5,
            zorder=3,
        )
        bars[0].set_edgecolor(NAVY)
        bars[0].set_linewidth(2.5)
        for bar, val in zip(bars, vals):
            h = bar.get_height()
            ax.text(
                bar.get_x() + bar.get_width() / 2,
                h + max(vals) * 0.02,
                f"{val:.3f}",
                ha="center",
                va="bottom",
                fontsize=14,
                fontweight="bold",
            )
        ax.set_title(label)
        ax.set_ylabel(label)
        if ylim:
            ax.set_ylim(ylim)
        ax.tick_params(axis="x", labelsize=13)

    plt.tight_layout()
    out = FIG_DIR / f"03_model_comparison_csic{_fig_suffix()}.png"
    plt.savefig(out)
    plt.close()
    print(f"  Saved: {out.name}")


# ═══════════════════════════════════════════════════════════════════════════════
# FIGURE 4 — Cross-dataset LSTM  (from results/evaluation_*.json)
# ═══════════════════════════════════════════════════════════════════════════════
def fig_cross_dataset():
    keys = ["csic2010", "cicids2018", "unsw"]
    labels = ["CSIC 2010", "CIC-IDS2018", "UNSW-NB15"]

    names = []
    precision = []
    recall = []
    f1 = []
    fpr = []

    for key, label in zip(keys, labels):
        ev = load_eval(key)
        if not ev or "LSTM-Autoencoder" not in ev:
            print(f"  WARNING: LSTM-Autoencoder missing from {key} evaluation")
            continue
        r = ev["LSTM-Autoencoder"]
        names.append(label)
        precision.append(r["precision"])
        recall.append(r["recall"])
        f1.append(r["f1"])
        fpr.append(r["fpr"])

    if not names:
        print("  SKIP fig 4 — no evaluation data found")
        return

    x = np.arange(len(names))
    width = 0.2

    fig, ax = plt.subplots(figsize=(13.0, 6.5))
    
    ax.bar(
        x - 1.5 * width,
        precision,
        width,
        label="Precision",
        color=TEAL,
        alpha=0.85,
        edgecolor="white",
        zorder=3,
    )
    ax.bar(
        x - 0.5 * width,
        recall,
        width,
        label="Recall",
        color=MINT,
        alpha=0.85,
        edgecolor="white",
        zorder=3,
    )
    ax.bar(
        x + 0.5 * width,
        f1,
        width,
        label="F1-Score",
        color=NAVY,
        alpha=0.85,
        edgecolor="white",
        zorder=3,
    )
    ax.bar(
        x + 1.5 * width,
        fpr,
        width,
        label="FPR",
        color=CORAL,
        alpha=0.85,
        edgecolor="white",
        zorder=3,
    )

    ax.set_xticks(x)
    ax.set_xticklabels(names, fontsize=16)
    ax.set_ylabel("Score")
    ax.set_ylim(0, 1.15)
    ax.legend(loc="upper right", fontsize=15)

    # Annotate CIC precision if available
    if "CIC-IDS2018" in names:
        cic_idx = names.index("CIC-IDS2018")
        cic_prec = precision[cic_idx]
        ax.annotate(
            "* 453:1 imbalance",
            xy=(cic_idx - 1.5 * width, cic_prec + 0.02),
            xytext=(cic_idx - 0.5, 0.3),
            fontsize=13,
            color=CORAL,
            arrowprops=dict(arrowstyle="->", color=CORAL, lw=1),
        )

    plt.tight_layout()
    out = FIG_DIR / f"04_cross_dataset_performance{_fig_suffix()}.png"
    plt.savefig(out)
    plt.close()
    print(f"  Saved: {out.name}")


# ═══════════════════════════════════════════════════════════════════════════════
# FIGURE 5 — Latency comparison  (from results/evaluation_*.json)
# ═══════════════════════════════════════════════════════════════════════════════
def fig_latency():
    keys = ["csic2010", "cicids2018", "unsw"]
    dataset_labels = ["CSIC 2010", "CIC-IDS2018", "UNSW-NB15"]
    model_keys = ["LSTM-Autoencoder", "Isolation Forest", "Random Forest"]
    model_labels = ["LSTM-AE", "Isolation Forest", "Random Forest"]
    colors = [TEAL, AMBER, CORAL]

    latencies = {mk: [] for mk in model_keys}

    for key in keys:
        ev = load_eval(key)
        for mk in model_keys:
            val = ev.get(mk, {}).get("latency_ms", 0.0)
            latencies[mk].append(val)

    x = np.arange(len(dataset_labels))
    width = 0.25

    fig, ax = plt.subplots(figsize=(13.0, 6.5))
    
    ymax = max(max(v) for v in latencies.values())

    for i, (mk, ml) in enumerate(zip(model_keys, model_labels)):
        offset = (i - 1) * width
        bars = ax.bar(
            x + offset,
            latencies[mk],
            width,
            label=ml,
            color=colors[i],
            alpha=0.85,
            edgecolor="white",
            zorder=3,
        )
        for bar, val in zip(bars, latencies[mk]):
            if val > 0:
                ax.text(
                    bar.get_x() + bar.get_width() / 2,
                    bar.get_height() + ymax * 0.015,
                    f"{val:.3f}",
                    ha="center",
                    va="bottom",
                    fontsize=14,
                )

    # Speedup annotations — placed well above the value labels (which sit
    # just above each bar) so the two text layers never collide, and given
    # enough headroom via ax.set_ylim below.
    lstm_v = latencies["LSTM-Autoencoder"]
    if_v = latencies["Isolation Forest"]
    rf_v = latencies["Random Forest"]
    for xi in range(len(keys)):
        if lstm_v[xi] > 0:
            ax.annotate(
                f"{if_v[xi] / lstm_v[xi]:.1f}\u00d7",
                xy=(xi - width / 2, if_v[xi] + ymax * 0.10),
                ha="center",
                fontsize=15,
                color=AMBER,
                fontweight="bold",
            )
            ax.annotate(
                f"{rf_v[xi] / lstm_v[xi]:.1f}\u00d7",
                xy=(xi + width / 2, rf_v[xi] + ymax * 0.10),
                ha="center",
                fontsize=15,
                color=CORAL,
                fontweight="bold",
            )

    ax.set_ylim(0, ymax * 1.28)
    ax.set_xticks(x)
    ax.set_xticklabels(dataset_labels)
    ax.set_ylabel("Latency (ms)")
    # Placed OUTSIDE the axes (above the plot, horizontal) rather than
    # inside a corner: an in-plot "upper left" legend box previously sat
    # directly on top of the CSIC 2010 Random Forest speedup annotation
    # (22.2x), making it render faint/illegible underneath the legend's
    # semi-transparent panel. Outside placement guarantees no overlap
    # with any bar, value label, or speedup annotation regardless of
    # which dataset group they belong to.
    ax.legend(
        fontsize=15, loc="lower center", bbox_to_anchor=(0.5, 1.01),
        ncol=3, frameon=False,
    )
    ax.text(
        0.02,
        0.62,
        "\u00d7 annotations show speedup\nrelative to LSTM-AE",
        transform=ax.transAxes,
        fontsize=14,
        color=GRAY,
        va="top",
    )

    plt.tight_layout()
    out = FIG_DIR / f"05_latency_comparison{_fig_suffix()}.png"
    plt.savefig(out)
    plt.close()
    print(f"  Saved: {out.name}")


# ═══════════════════════════════════════════════════════════════════════════════
# FIGURE 6 — Active Learning  (from results/active_learning_*.json)
# ═══════════════════════════════════════════════════════════════════════════════
def fig_active_learning():
    configs = [
        ("csic2010", "CSIC 2010", TEAL),
        ("cicids2018", "CIC-IDS2018", AMBER),
        ("unsw", "UNSW-NB15", PURPLE),
    ]

    fig, axes = _grid_2x2(n_used=3)
    
    for ax, (key, label, color) in zip(axes, configs):
        history = load_al(key)
        if not history:
            ax.text(
                0.5,
                0.5,
                "Data not found",
                ha="center",
                va="center",
                transform=ax.transAxes,
            )
            ax.set_title(label)
            continue

        iters = [h["iteration"] for h in history]
        thresh = [h["threshold"] for h in history]
        fpr_v = [h["fpr"] for h in history]
        f1_v = [h["f1"] for h in history]

        ax2 = ax.twinx()

        (l1,) = ax.plot(
            iters,
            thresh,
            color=color,
            lw=2.5,
            marker="o",
            ms=6,
            label="Threshold",
            zorder=3,
        )
        (l2,) = ax2.plot(
            iters,
            fpr_v,
            color=CORAL,
            lw=2,
            marker="s",
            ms=5,
            linestyle="--",
            label="FPR",
            zorder=3,
        )
        (l3,) = ax2.plot(
            iters,
            f1_v,
            color=MINT,
            lw=2,
            marker="^",
            ms=5,
            linestyle=":",
            label="F1",
            zorder=3,
        )

        ax.set_title(label)
        ax.set_xlabel("AL Iteration")
        ax.set_ylabel("Threshold", color=color)
        ax2.set_ylabel("FPR / F1")
        ax.tick_params(axis="y", labelcolor=color)
        ax.legend(
            [l1, l2, l3], ["Threshold", "FPR", "F1"], fontsize=14, loc="best"
        )

    plt.tight_layout()
    out = FIG_DIR / f"06_active_learning{_fig_suffix()}.png"
    plt.savefig(out)
    plt.close()
    print(f"  Saved: {out.name}")


# ═══════════════════════════════════════════════════════════════════════════════
# FIGURE 7 — Error separation box plots
# (threshold from disk; distributions approximated from threshold stats)
# ═══════════════════════════════════════════════════════════════════════════════
def fig_error_separation():
    configs = [
        ("csic2010", "CSIC 2010", TEAL),
        ("cicids2018", "CIC-IDS2018", AMBER),
        ("unsw", "UNSW-NB15", PURPLE),
    ]

    fig, axes = _grid_2x2(n_used=3)
    
    rng = np.random.default_rng(42)

    for ax, (key, name, color) in zip(axes, configs):
        threshold = load_threshold(key)
        thresh_data = load_threshold_full(key)
        if threshold is None:
            ax.set_title(name)
            continue

        real = load_real_errors(key)
        if real is not None:
            errors, labels = real
            normal_e = errors[labels == 0]
            attack_e = errors[labels == 1]
            is_simulated = False
        else:
            mean_err = thresh_data.get("mean_error", threshold * 0.3)
            std_err = min(thresh_data.get("std_error", threshold * 0.2), threshold * 1.5)

            # Normal errors: clustered well below threshold
            normal_e = np.abs(rng.normal(mean_err * 0.4, std_err * 0.3, 2000))
            normal_e = np.clip(normal_e, 0, threshold * 0.9)

            # Attack errors: spread above threshold
            attack_e = np.abs(rng.normal(threshold * 4, threshold * 2, 1000))
            attack_e = np.clip(attack_e, threshold * 0.3, threshold * 15)
            is_simulated = True

        # Same root cause as Figure 2: a handful of extreme outliers
        # (attack errors reaching into the hundreds or thousands) blow
        # out a LINEAR y-axis so far that the actual 25th-75th
        # percentile box — sitting down near 0 — gets squashed into an
        # invisible hairline. A log y-axis keeps the box itself visible
        # while still showing the outlier whiskers/fliers. Clip to a
        # tiny positive floor first since log-scale can't plot exact
        # zeros (MSE is non-negative but can occasionally be ~0).
        eps = max(threshold * 1e-3, 1e-12)
        normal_e = np.clip(normal_e, eps, None)
        attack_e = np.clip(attack_e, eps, None)

        bp = ax.boxplot(
            [normal_e, attack_e],
            tick_labels=["Normal", "Attack"],
            patch_artist=True,
            medianprops=dict(color=NAVY, lw=2.5),
            whiskerprops=dict(lw=1.5),
            capprops=dict(lw=1.5),
            flierprops=dict(marker="o", markersize=2, alpha=0.3),
            widths=0.5,
        )
        bp["boxes"][0].set_facecolor(color)
        bp["boxes"][0].set_alpha(0.6)
        bp["boxes"][1].set_facecolor(CORAL)
        bp["boxes"][1].set_alpha(0.6)

        ax.axhline(
            threshold,
            color=NAVY,
            lw=2,
            linestyle="-.",
            label=f"τ = {threshold:.5f}",
            zorder=4,
        )
        ax.set_yscale("log")
        ax.set_title(name)
        ax.set_ylabel("Reconstruction Error (MSE, log scale)")
        ax.legend(fontsize=14)

        if is_simulated:
            ax.text(
                0.98, 0.02, "SIMULATED\n(re-run evaluate.py to plot real errors)",
                transform=ax.transAxes, ha="right", va="bottom",
                fontsize=12.5, color=CORAL, fontweight="bold",
                bbox=dict(boxstyle="round", facecolor="white", edgecolor=CORAL, alpha=0.85),
            )

    plt.tight_layout()
    out = FIG_DIR / f"07_error_separation{_fig_suffix()}.png"
    plt.savefig(out)
    plt.close()
    print(f"  Saved: {out.name}")


# ═══════════════════════════════════════════════════════════════════════════════
# FIGURE 8 — Confusion matrices, all models × all datasets
# (from results/evaluation_*.json — tp/fp/tn/fn were computed but never plotted)
# ═══════════════════════════════════════════════════════════════════════════════
def fig_confusion_matrices():
    datasets = [
        ("csic2010", "CSIC 2010"),
        ("cicids2018", "CIC-IDS2018"),
        ("unsw", "UNSW-NB15"),
    ]
    model_order = ["LSTM-Autoencoder", "Isolation Forest", "Random Forest", "WAF Simulation"]
    short_labels = {
        "LSTM-Autoencoder": "LSTM-AE",
        "Isolation Forest": "Isolation Forest",
        "Random Forest": "Random Forest",
        "WAF Simulation": "WAF Sim",
    }

    # Split into one 2x2-grid image per dataset (LSTM-AE/IF on top,
    # RF/WAF on bottom) instead of one dense 3x4 combined image — each
    # panel gets roughly 3x the linear size for the same page width.
    any_saved = False
    for key, label in datasets:
        ev = load_eval(key)
        if not ev:
            print(f"  SKIP fig 8 ({key}) — evaluation_{key} not found")
            continue

        fig, axes = _grid_2x2(figsize=(13.5, 12.0), n_used=4)

        for ax, model in zip(axes, model_order):
            r = ev.get(model)
            if not r:
                ax.text(0.5, 0.5, "no data", ha="center", va="center", transform=ax.transAxes)
                ax.set_xticks([]); ax.set_yticks([])
                ax.set_title(short_labels[model], fontsize=17)
                continue

            any_saved = True
            cm = np.array([[r["tn"], r["fp"]], [r["fn"], r["tp"]]])
            cm_norm = cm / cm.sum(axis=1, keepdims=True).clip(min=1)

            ax.imshow(cm_norm, cmap="Blues", vmin=0, vmax=1, aspect="auto")
            for i in range(2):
                for j in range(2):
                    val = cm[i, j]
                    pct = cm_norm[i, j] * 100
                    txt_color = "white" if cm_norm[i, j] > 0.5 else NAVY
                    ax.text(
                        j, i, f"{val:,}\n({pct:.1f}%)",
                        ha="center", va="center", fontsize=16, color=txt_color,
                    )
            ax.set_xticks([0, 1]); ax.set_xticklabels(["Normal", "Attack"], fontsize=15)
            ax.set_yticks([0, 1]); ax.set_yticklabels(["Normal", "Attack"], fontsize=15)
            ax.set_title(short_labels[model], fontsize=17)
            ax.set_ylabel("True label", fontsize=15)
            ax.set_xlabel("Predicted label", fontsize=15)

        plt.tight_layout()
        out = FIG_DIR / f"08_confusion_matrices_{key}{_fig_suffix()}.png"
        plt.savefig(out)
        plt.close()
        print(f"  Saved: {out.name}")

    if not any_saved:
        print("  SKIP fig 8 — no evaluation data found for any dataset")


# ═══════════════════════════════════════════════════════════════════════════════
# FIGURE 9 — Accuracy and throughput
# (from results/evaluation_*.json — both computed, neither previously plotted)
# ═══════════════════════════════════════════════════════════════════════════════
def fig_accuracy_throughput():
    keys = ["csic2010", "cicids2018", "unsw"]
    labels = ["CSIC 2010", "CIC-IDS2018", "UNSW-NB15"]
    model_order = ["LSTM-Autoencoder", "Isolation Forest", "Random Forest", "WAF Simulation"]
    short_labels = ["LSTM-AE", "Isolation\nForest", "Random\nForest", "WAF\nSim"]
    colors = [TEAL, AMBER, CORAL, GRAY]

    fig, axes = plt.subplots(1, 2, figsize=(18.2, 6.5))
    
    # Left: grouped accuracy bars, one group per dataset
    ax = axes[0]
    x = np.arange(len(keys))
    width = 0.2
    any_acc = False
    for i, model in enumerate(model_order):
        vals = []
        for key in keys:
            ev = load_eval(key)
            vals.append(ev.get(model, {}).get("accuracy", np.nan))
        if not all(np.isnan(v) for v in vals):
            any_acc = True
        ax.bar(x + (i - 1.5) * width, vals, width, label=short_labels[i].replace("\n", " "), color=colors[i])
    ax.set_xticks(x); ax.set_xticklabels(labels)
    ax.set_ylabel("Accuracy")
    ax.set_ylim(0, 1.05)
    ax.set_title("Accuracy by model and dataset")
    ax.legend(fontsize=13)
    if not any_acc:
        ax.text(0.5, 0.5, "no data", ha="center", va="center", transform=ax.transAxes)

    # Right: throughput (log scale — LSTM/RF/IF/WAF can differ by orders of magnitude)
    ax = axes[1]
    any_thr = False
    for i, model in enumerate(model_order):
        vals = []
        for key in keys:
            ev = load_eval(key)
            vals.append(ev.get(model, {}).get("throughput", np.nan))
        if not all(np.isnan(v) for v in vals):
            any_thr = True
        ax.bar(x + (i - 1.5) * width, vals, width, label=short_labels[i].replace("\n", " "), color=colors[i])
    ax.set_xticks(x); ax.set_xticklabels(labels)
    ax.set_ylabel("Throughput (sessions/sec, log scale)")
    ax.set_yscale("log")
    ax.set_title("Inference throughput by model and dataset")
    ax.legend(fontsize=13)
    if not any_thr:
        ax.text(0.5, 0.5, "no data", ha="center", va="center", transform=ax.transAxes)

    if not (any_acc or any_thr):
        print("  SKIP fig 9 — no evaluation data found")
        plt.close()
        return

    plt.tight_layout()
    out = FIG_DIR / f"09_accuracy_throughput{_fig_suffix()}.png"
    plt.savefig(out)
    plt.close()
    print(f"  Saved: {out.name}")


# ═══════════════════════════════════════════════════════════════════════════════
# FIGURE 10 — Active Learning: precision/recall and labeling effort
# (from results/active_learning_*.json — precision, recall, n_labeled
#  were all recorded per iteration but Figure 6 only showed threshold/F1/FPR)
# ═══════════════════════════════════════════════════════════════════════════════
def fig_active_learning_effort():
    configs = [
        ("csic2010", "CSIC 2010", TEAL),
        ("cicids2018", "CIC-IDS2018", AMBER),
        ("unsw", "UNSW-NB15", PURPLE),
    ]

    any_saved = False
    for key, label, color in configs:
        history = load_al(key)
        if not history:
            print(f"  SKIP fig 10 ({key}) — active_learning_{key} not found")
            continue

        any_saved = True
        fig, (ax_top, ax_bot) = plt.subplots(2, 1, figsize=(9.5, 12.5))

        iters = [h["iteration"] for h in history]
        prec = [h["precision"] for h in history]
        rec = [h["recall"] for h in history]
        n_labeled = [h["n_labeled"] for h in history]
        n_uncertain = [h.get("n_uncertain", 0) for h in history]

        # Top: precision & recall over iterations
        ax_top.plot(iters, prec, "o-", color=color, lw=2, ms=6, label="Precision")
        ax_top.plot(iters, rec, "s--", color=CORAL, lw=2, ms=5, label="Recall")
        ax_top.set_title(label)
        ax_top.set_xlabel("AL Iteration")
        ax_top.set_ylabel("Score")
        ax_top.set_ylim(0, 1.05)
        ax_top.legend(fontsize=15)

        # Bottom: cumulative labeling effort vs uncertain pool size
        ax_bot.bar(iters, n_uncertain, alpha=0.35, color=GRAY, label="Uncertain pool size", zorder=2)
        ax2 = ax_bot.twinx()
        ax2.plot(iters, n_labeled, "^-", color=NAVY, lw=2, ms=6, label="Cumulative labeled", zorder=3)
        ax_bot.set_xlabel("AL Iteration")
        ax_bot.set_ylabel("Uncertain pool size", color=GRAY)
        ax2.set_ylabel("Cumulative sessions labeled", color=NAVY)

        lines1, labels1 = ax_bot.get_legend_handles_labels()
        lines2, labels2 = ax2.get_legend_handles_labels()
        ax_bot.legend(lines1 + lines2, labels1 + labels2, fontsize=13, loc="upper left")

        plt.tight_layout()
        out = FIG_DIR / f"10_active_learning_effort_{key}{_fig_suffix()}.png"
        plt.savefig(out)
        plt.close()
        print(f"  Saved: {out.name}")

    if not any_saved:
        print("  SKIP fig 10 — no active learning data found")


# ═══════════════════════════════════════════════════════════════════════════════
# FIGURE 11 — Window size comparison (LSTM-Autoencoder only)
# Scans disk for every window size that's been trained+evaluated per
# dataset (via discover_windows) rather than relying on
# WINDOW_BY_DATASET, so this figure always reflects everything you've
# actually run regardless of what the rest of the module is set to plot.
# ═══════════════════════════════════════════════════════════════════════════════
def fig_window_comparison():
    configs = [
        ("csic2010", "CSIC 2010", TEAL),
        ("cicids2018", "CIC-IDS2018", AMBER),
        ("unsw", "UNSW-NB15", PURPLE),
    ]

    fig, axes = _grid_2x2(n_used=3)
    
    any_data = False
    for ax, (key, label, color) in zip(axes, configs):
        windows = discover_windows(key)
        if not windows:
            ax.text(0.5, 0.5, "No trained windows found", ha="center", va="center", transform=ax.transAxes)
            ax.set_title(label)
            continue

        f1s, recalls, fprs, threshes = [], [], [], []
        valid_windows = []
        for w in windows:
            ev = load_eval_for(key, w)
            lstm = ev.get("LSTM-Autoencoder")
            if lstm is None:
                continue
            valid_windows.append(w)
            f1s.append(lstm["f1"])
            recalls.append(lstm["recall"])
            fprs.append(lstm["fpr"])

        if not valid_windows:
            ax.text(
                0.5, 0.5, "Trained but not\nevaluated yet",
                ha="center", va="center", transform=ax.transAxes,
            )
            ax.set_title(label)
            continue

        any_data = True
        ax.plot(valid_windows, f1s, "o-", color=color, lw=2.5, ms=7, label="F1", zorder=3)
        ax.plot(valid_windows, recalls, "^--", color=MINT, lw=2, ms=6, label="Recall", zorder=3)
        ax.plot(valid_windows, fprs, "s:", color=CORAL, lw=2, ms=6, label="FPR", zorder=3)

        ax.set_title(label)
        ax.set_xlabel("Window size (sessions)")
        ax.set_ylabel("Score")
        ax.set_xticks(valid_windows)
        ax.set_ylim(-0.02, 1.05)
        ax.legend(fontsize=14)

    if not any_data:
        print("  SKIP fig 11 — no multi-window evaluation data found "
              "(train + evaluate at more than one --window to populate this)")
        plt.close()
        return

    plt.tight_layout()
    out = FIG_DIR / "11_window_comparison.png"  # not suffixed — this IS the cross-window comparison
    plt.savefig(out)
    plt.close()
    print(f"  Saved: {out.name}")


# ═══════════════════════════════════════════════════════════════════════════════
# FIGURE 12 — Dataset summary pie charts
# (fixed constants — total record counts from dataset papers)
# ═══════════════════════════════════════════════════════════════════════════════
def fig_dataset_summary():
    panel_order = ["csic2010", "cicids2018_full", "cicids2018", "unsw"]
    fig, axes = _grid_2x2(figsize=(15.0, 14.0), n_used=4)

    def _autopct_fmt(pct):
        # A plain "%1.1f%%" rounds anything under 0.05% down to "0.0%"
        # (and its complement up to "100.0%"), which for CIC-IDS2018's
        # 928-in-2,097,150 attack share reads as "no attacks at all" —
        # misleading given the whole point of this panel is to show
        # that tiny-but-nonzero attack fraction. Use more decimal
        # places whenever the value is close to either extreme.
        if pct < 1 or pct > 99:
            return f"{pct:.2f}%"
        return f"{pct:.1f}%"

    for ax, key in zip(axes, panel_order):
        info = DATASET_COUNTS[key]
        n = info["normal"]
        a = info["attack"]
        total = n + a
        color = DATASET_COLORS[key]

        _, _, autotexts = ax.pie(
            [n / total * 100, a / total * 100],
            explode=(0, 0.05),
            labels=["Normal", "Attack"],
            colors=[color, CORAL],
            autopct=_autopct_fmt,
            startangle=90,
            pctdistance=0.75,
            wedgeprops=dict(edgecolor="white", linewidth=2),
        )
        for at in autotexts:
            at.set_fontsize(15)
            at.set_fontweight("bold")

        ax.set_title(f"{info['label']}\n({total:,} total records)")
        ax.text(
            0,
            -1.35,
            f"Normal: {n:,}\nAttack: {a:,}",
            ha="center",
            fontsize=14,
            color="black",
        )

        if key == "cicids2018":
            # This panel is the two-day subset (02-22/02-23-2018) actually
            # used for training/evaluation in this thesis; compare directly
            # against the full-release panel immediately to its left.
            ax.text(
                0,
                -1.68,
                f"Used subset: 2 of 10 capture days\n"
                f"({CICIDS2018_USED_ROWS:,} / "
                f"{CICIDS2018_FULL_RELEASE_ROWS:,} rows, "
                f"{CICIDS2018_USED_PCT:.1f}% of full release)",
                ha="center",
                fontsize=13,
                color=CORAL,
                fontweight="bold",
                style="italic",
            )
        elif key == "cicids2018_full":
            ax.text(
                0,
                -1.68,
                "Includes DDoS/Bot/Infiltration days\n"
                "outside this thesis's stated scope\n"
                "(sequence-dependent API-logic attacks)",
                ha="center",
                fontsize=13,
                color="black",
                style="italic",
            )

    plt.tight_layout()
    plt.subplots_adjust(bottom=0.1, hspace=0.55)
    out = FIG_DIR / f"12_dataset_summary{_fig_suffix()}.png"
    plt.savefig(out)
    plt.close()
    print(f"  Saved: {out.name}")


# ═══════════════════════════════════════════════════════════════════════════════
# MAIN — print a summary of what was read from disk before generating
# ═══════════════════════════════════════════════════════════════════════════════
def _parse_args():
    import argparse

    parser = argparse.ArgumentParser(
        description="Generate thesis figures from saved pipeline results."
    )
    parser.add_argument(
        "--window",
        type=int,
        default=None,
        help="Window size to use for ALL three datasets (e.g. 5, 8, 10). "
        "Overrides WINDOW_BY_DATASET so you can re-run this script once "
        "per window size you've trained/evaluated.",
    )
    parser.add_argument(
        "--window-csic2010",
        type=int,
        default=None,
        help="Override window size for csic2010 only (takes precedence over --window).",
    )
    parser.add_argument(
        "--window-cicids2018",
        type=int,
        default=None,
        help="Override window size for cicids2018 only (takes precedence over --window).",
    )
    parser.add_argument(
        "--window-unsw",
        type=int,
        default=None,
        help="Override window size for unsw only (takes precedence over --window).",
    )
    return parser.parse_args()


if __name__ == "__main__":
    args = _parse_args()

    # Apply CLI overrides on top of the WINDOW_BY_DATASET defaults above.
    # --window sets all three; the per-dataset flags win if also given.
    if args.window is not None:
        for key in WINDOW_BY_DATASET:
            WINDOW_BY_DATASET[key] = args.window
    if args.window_csic2010 is not None:
        WINDOW_BY_DATASET["csic2010"] = args.window_csic2010
    if args.window_cicids2018 is not None:
        WINDOW_BY_DATASET["cicids2018"] = args.window_cicids2018
    if args.window_unsw is not None:
        WINDOW_BY_DATASET["unsw"] = args.window_unsw

    print(f"Output directory: {FIG_DIR}\n")
    print("Window sizes in use:", WINDOW_BY_DATASET, "\n")

    print("── Data loaded from disk ─────────────────────────────────────────")
    for key in ["csic2010", "cicids2018", "unsw"]:
        label = DATASET_COUNTS[key]["label"].replace("\n", " ")
        run_id = _run_id(key)
        h = load_history(key)
        t = load_threshold_full(key)
        ev = load_eval(key)
        al = load_al(key)

        epochs = len(h.get("epochs", []))
        threshold = t.get("threshold", "N/A")
        lstm = ev.get("LSTM-Autoencoder", {})
        f1 = lstm.get("f1", "N/A")
        recall = lstm.get("recall", "N/A")
        latency = lstm.get("latency_ms", "N/A")
        al_iters = len(al)

        thresh_str = f"{threshold:.6f}" if isinstance(threshold, (int, float)) else str(threshold)
        print(
            f"  {label:<14} ({run_id})  "
            f"epochs={epochs}  "
            f"τ={thresh_str}  "
            f"F1={f1}  "
            f"Rec={recall}  "
            f"Lat={latency}ms  "
            f"AL_iters={al_iters}"
        )
    print()

    print("── Generating figures ────────────────────────────────────────────")
    fig_learning_curves()
    fig_error_distributions()
    fig_model_comparison()
    fig_cross_dataset()
    fig_latency()
    fig_active_learning()
    fig_error_separation()
    fig_confusion_matrices()
    fig_accuracy_throughput()
    fig_active_learning_effort()
    fig_window_comparison()
    fig_dataset_summary()

    print(f"\nDone — up to 12 figures saved to {FIG_DIR}/ (some may be skipped if data is missing)")