scripts/02_run_baselines.py

"""Run the non-neural baselines on both split schemes.

    python scripts/02_run_baselines.py
"""

import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

import numpy as np  # noqa: E402

from cti_attack import baselines, data, evaluate  # noqa: E402


def run_scheme(scheme: str, labels: list[str]) -> dict:
    print(f"\n{'=' * 62}\n  split scheme: {scheme}\n{'=' * 62}")
    tr = data.load_split(scheme, "train")
    dv = data.load_split(scheme, "dev")
    te = data.load_split(scheme, "test")

    Ytr, Ydv, Yte = (evaluate.to_matrix(x, labels) for x in (tr, dv, te))
    txt_tr = [r["sentence"] for r in tr]
    txt_dv = [r["sentence"] for r in dv]
    txt_te = [r["sentence"] for r in te]

    results = {}

    # ---- frequency prior -------------------------------------------------
    dev_s = baselines.frequency_scores(Ytr, len(dv))
    test_s = baselines.frequency_scores(Ytr, len(te))
    results["frequency"] = _score("frequency", Ydv, dev_s, Yte, test_s, labels)

    # ---- ATT&CK keyword match -------------------------------------------
    dev_s = baselines.keyword_scores(txt_dv, labels)
    test_s = baselines.keyword_scores(txt_te, labels)
    results["keyword"] = _score("keyword", Ydv, dev_s, Yte, test_s, labels)

    # ---- TF-IDF + one-vs-rest logistic regression ------------------------
    print("  fitting tfidf_lr (49 one-vs-rest classifiers) …")
    s = baselines.tfidf_lr_scores(txt_tr, Ytr, {"dev": txt_dv, "test": txt_te})
    results["tfidf_lr"] = _score("tfidf_lr", Ydv, s["dev"], Yte, s["test"], labels)

    return results


def _score(name, Ydv, dev_scores, Yte, test_scores, labels) -> dict:
    gt, _ = evaluate.tune_global_threshold(Ydv, dev_scores)
    pct = evaluate.tune_per_class_thresholds(Ydv, dev_scores)

    rep_g = evaluate.evaluate(Yte, evaluate.apply_thresholds(test_scores, gt), labels)
    rep_p = evaluate.evaluate(Yte, evaluate.apply_thresholds(test_scores, pct), labels)

    print(f"  {name:11} global t={gt:<5} macro-F1={rep_g.macro_f1:.4f} "
          f"micro-F1={rep_g.micro_f1:.4f}   |  per-class macro-F1={rep_p.macro_f1:.4f} "
          f"micro-F1={rep_p.micro_f1:.4f}")

    return {
        "global_threshold": {"threshold": gt, **rep_g.as_dict()},
        "per_class_threshold": {
            "thresholds": {l: float(t) for l, t in zip(labels, pct)},
            **rep_p.as_dict(),
        },
    }


def main() -> None:
    labels = data.load_labels()
    all_results = {}
    for scheme in ("document", "random"):
        all_results[scheme] = run_scheme(scheme, labels)
        evaluate.save_report("baselines", scheme, all_results[scheme])

    print(f"\n{'=' * 62}\n  leakage gap (test macro-F1, per-class thresholds)\n{'=' * 62}")
    for name in ("frequency", "keyword", "tfidf_lr"):
        d = all_results["document"][name]["per_class_threshold"]["macro_f1"]
        r = all_results["random"][name]["per_class_threshold"]["macro_f1"]
        print(f"  {name:11} document={d:.4f}  random={r:.4f}  "
              f"inflation={r - d:+.4f} ({(r / d - 1) * 100 if d else float('nan'):+.1f}%)")


if __name__ == "__main__":
    main()