scripts/04_report.py

"""Generate the results tables that the README embeds.

Every number in the model card comes from here, which reads only the JSON
written by scripts 02 and 03. Nothing is typed by hand.

    python scripts/04_report.py
"""

import json
import sys
from pathlib import Path

# RESULTS.md is written as UTF-8 regardless; this only stops a Windows console
# on a legacy codepage (cp1252/cp1254) from crashing on the arrows and dashes.
try:
    sys.stdout.reconfigure(encoding="utf-8", errors="replace")
except (AttributeError, ValueError):  # pragma: no cover
    pass

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

from cti_attack import config, data  # noqa: E402

ORDER = ["frequency", "keyword", "tfidf_lr", "modernbert", "deberta", "securebert"]
PRETTY = {
    "frequency": "Frequency prior",
    "keyword": "ATT&CK keyword match",
    "tfidf_lr": "TF-IDF + one-vs-rest LR",
    "modernbert": "ModernBERT-base (fine-tuned)",
    "deberta": "DeBERTa-v3-base (fine-tuned)",
    "securebert": "SecureBERT (fine-tuned)",
}
TRAINED = {"modernbert", "deberta", "securebert"}


def collect() -> dict:
    out: dict[str, dict] = {}
    for scheme in ("document", "random"):
        p = config.RESULTS_DIR / f"baselines__{scheme}.json"
        if p.exists():
            for name, payload in json.loads(p.read_text(encoding="utf-8")).items():
                out.setdefault(name, {})[scheme] = payload
        for name in TRAINED:
            q = config.RESULTS_DIR / f"{name}__{scheme}.json"
            if q.exists():
                out.setdefault(name, {})[scheme] = json.loads(q.read_text(encoding="utf-8"))
    return out


def _cell(res: dict, scheme: str, metric: str, mode: str) -> str:
    try:
        return f"{res[scheme][mode][metric]:.3f}"
    except (KeyError, TypeError):
        return "—"


def main_table(results: dict, mode: str = "per_class_threshold") -> str:
    lines = [
        "| Model | doc macro-F1 | doc micro-F1 | random macro-F1 | random micro-F1 | inflation |",
        "|---|---|---|---|---|---|",
    ]
    for name in ORDER:
        if name not in results:
            continue
        res = results[name]
        dm = _cell(res, "document", "macro_f1", mode)
        rm = _cell(res, "random", "macro_f1", mode)
        infl = "n/a"
        if dm != "—" and rm != "—" and float(dm) > 0:
            infl = f"{(float(rm) / float(dm) - 1) * 100:+.1f}%"
        label = PRETTY[name]
        if name == "modernbert":
            label = f"**{label}**"
        lines.append(
            f"| {label} | {dm} | {_cell(res, 'document', 'micro_f1', mode)} "
            f"| {rm} | {_cell(res, 'random', 'micro_f1', mode)} | {infl} |")
    return "\n".join(lines)


def per_class_table(results: dict, model: str = "modernbert", top: int = 12) -> str:
    try:
        pc = results[model]["document"]["per_class_threshold"]["per_class"]
    except KeyError:
        return "_(not available — train a model first)_"
    rows = sorted(pc.items(), key=lambda kv: -kv[1]["support"])[:top]

    try:
        from cti_attack import attack_meta
        names = attack_meta.build_technique_names()
        status = attack_meta.build_technique_status()
    except Exception:
        names, status = {}, {}

    lines = ["| Technique | Name | Support | P | R | F1 |", "|---|---|---|---|---|---|"]
    for tid, m in rows:
        nm = names.get(tid, "")
        if len(nm) > 46:
            nm = nm[:43] + "…"
        if status.get(tid) == "revoked":
            nm += " ⚠️*revoked*"
        lines.append(f"| `{tid}` | {nm} | {m['support']} | {m['precision']:.2f} "
                     f"| {m['recall']:.2f} | {m['f1']:.2f} |")
    return "\n".join(lines)


def tail_table(results: dict, model: str = "modernbert") -> str:
    """Head vs tail: the number that macro-F1 exists to expose."""
    try:
        pc = results[model]["document"]["per_class_threshold"]["per_class"]
    except KeyError:
        return "_(not available)_"
    head = [m for m in pc.values() if m["support"] >= 20]
    tail = [m for m in pc.values() if m["support"] < 20]
    lines = ["| Bucket | Techniques | Mean F1 |", "|---|---|---|"]
    for nm, bucket in (("head (>=20 test examples)", head), ("tail (<20 test examples)", tail)):
        if bucket:
            lines.append(f"| {nm} | {len(bucket)} | "
                         f"{sum(m['f1'] for m in bucket) / len(bucket):.3f} |")
    return "\n".join(lines)


def _paired(a: list[float], b: list[float]) -> dict:
    """Paired difference stats for a - b across folds, with a p-value if scipy is present."""
    import math
    d = [x - y for x, y in zip(a, b)]
    n = len(d)
    mean = sum(d) / n
    var = sum((x - mean) ** 2 for x in d) / (n - 1) if n > 1 else 0.0
    std = math.sqrt(var)
    se = std / math.sqrt(n) if std > 0 else 0.0
    t = mean / se if se > 0 else float("inf")
    pos = sum(1 for x in d if x > 0)
    p = None
    try:
        from scipy import stats  # optional
        p = float(stats.ttest_rel(a, b).pvalue)
    except Exception:
        pass
    return {"mean": mean, "std": std, "t": t, "pos": pos, "n": n, "p": p}


def _sig(p, t, n) -> str:
    if p is not None:
        return f"p = {p:.3f}"
    # df = n-1; two-sided 0.05 critical t for df=4 is 2.776, df=3 is 3.182
    crit = {2: 4.303, 3: 3.182, 4: 2.776, 5: 2.571, 6: 2.447}.get(n - 1, 2.776)
    return f"t = {t:.2f} ({'p<0.05' if abs(t) > crit else 'n.s.'})"


def cv_ensemble_section() -> str:
    """Cross-validated ensemble results, read from the CV and ensemble JSON.

    This is the section that turns 'higher on one split' into 'higher across
    independent leak-free folds, with a paired significance test'.
    """
    cvp = config.RESULTS_DIR / "cv_5fold__modernbert.json"
    enp = config.RESULTS_DIR / "ensemble_tfidf_modernbert__document.json"
    if not cvp.exists():
        return ""
    cv = json.loads(cvp.read_text(encoding="utf-8"))
    S, folds = cv["summary"], cv["per_fold"]

    def ms(k):
        return f"{S[k]['mean']:.4f} +/- {S[k]['std']:.4f}"

    lines = [
        f"One 70/15/15 split of a {cv['n_documents']}-document corpus is one "
        f"draw, so the gap between models can be noise. To check it, here is "
        f"{cv['folds']}-fold cross-validation grouped by document. Each report "
        f"lands in exactly one test fold. The blend weight and thresholds are "
        f"tuned on each fold's own dev set, never on its test set.\n",
        "| Model | doc macro-F1 (5-fold CV, mean +/- std) |",
        "|---|---|",
        f"| TF-IDF + LR (per-class) | {ms('tfidf_perclass')} |",
        f"| TF-IDF + LR (global) | {ms('tfidf_global')} |",
        f"| ModernBERT (per-class) | {ms('modernbert_perclass')} |",
        f"| **Ensemble TF-IDF+ModernBERT (per-class)** | **{ms('ensemble_perclass')}** |",
        "",
        "Paired per-fold test, per-class thresholds on both sides:",
    ]
    ep = [f["ensemble_perclass"] for f in folds]
    tests = [
        ("ensemble vs TF-IDF, same per-class regime", ep, [f["tfidf_perclass"] for f in folds]),
        ("ensemble vs ModernBERT, same per-class regime", ep, [f["modernbert_perclass"] for f in folds]),
        ("ensemble vs TF-IDF at its stronger global regime", ep, [f["tfidf_global"] for f in folds]),
    ]
    for label, a, b in tests:
        r = _paired(a, b)
        lines.append(
            f"- {label}: **{r['mean']:+.4f}**, {r['pos']}/{r['n']} folds positive, {_sig(r['p'], r['t'], r['n'])}")

    if enp.exists():
        en = json.loads(enp.read_text(encoding="utf-8"))
        a = en.get("alpha")
        pc = en["per_class_threshold"]["macro_f1"]
        comp = en["components_test_macro_f1"]
        lines.append(
            f"\nOn the original single 70/15/15 document split, the same ensemble "
            f"({a:.0%} TF-IDF, {1-a:.0%} ModernBERT) scores "
            f"**{pc:.4f}** macro-F1, against {comp['tfidf_lr']['per_class']:.4f} "
            f"for TF-IDF and {comp['modernbert']['per_class']:.4f} for ModernBERT.")

    lines.append(
        "\nHow to read this: the ensemble has the highest mean and the lowest "
        "variance of everything tested. Compared against each base model at the "
        "same threshold setting, it wins in all five folds. The one comparison it "
        "does not clearly win is against TF-IDF at its own best threshold setting, "
        "where the lead is small enough to be noise over five folds. That row is "
        "in the table on purpose.")
    return "\n".join(lines)


def main() -> None:
    results = collect()
    if not results:
        print("no results found — run scripts 02 and 03 first")
        return

    stats = json.loads((config.BUILD_DIR / "build_stats.json").read_text(encoding="utf-8"))
    labels = data.load_labels()
    cv_block = cv_ensemble_section()

    doc = [
        "<!-- generated by scripts/04_report.py - do not edit by hand -->",
        "### Test-set results\n",
        "Per-class thresholds tuned on dev. Macro-F1 is the headline metric. It "
        "weights all 49 techniques equally, so a handful of common ones cannot "
        "cover for a weak long tail.\n",
        main_table(results, "per_class_threshold"),
        "\n### Same table, single global threshold\n",
        main_table(results, "global_threshold"),
        *(["\n### Ensemble + cross-validated results\n", cv_block] if cv_block else []),
        "\n### ModernBERT per-technique, document split (12 most frequent)\n",
        per_class_table(results),
        "\n### Head vs tail, document split\n",
        tail_table(results),
        "\n### Corpus\n",
        f"- {stats['raw_sentences']:,} raw sentences, {stats['final_sentences']:,} after cleaning and dedup",
        f"- {stats['final_labelled']:,} carry at least one technique "
        f"({stats['final_labelled'] / stats['final_sentences']:.1%})",
        f"- {len(labels)} techniques across {stats['final_documents']} source documents",
        f"- {stats['duplicates_removed']:,} duplicate sentences removed "
        f"({stats['cross_document_duplicates']} of them appeared in more than one document)",
        f"- dropped for having too few documents to split: {', '.join(stats['dropped_techniques'])}",
    ]
    out = config.RESULTS_DIR / "RESULTS.md"
    out.write_text("\n".join(doc) + "\n", encoding="utf-8")
    print("\n".join(doc))
    print(f"\n-> {out.relative_to(config.REPO_ROOT)}")

    inject_into_readme("\n".join(doc[1:]))  # drop the "generated by" comment


BEGIN = "<!-- BEGIN GENERATED RESULTS - scripts/04_report.py -->"
END = "<!-- END GENERATED RESULTS -->"


def inject_into_readme(block: str) -> None:
    """Rewrite the results section of README.md in place.

    The tables in the model card are generated, not transcribed. Pasting them
    by hand is how a card drifts out of sync with the data it describes, which
    is the exact failure this repo argues against.
    """
    readme = config.REPO_ROOT / "README.md"
    if not readme.exists():
        return
    text = readme.read_text(encoding="utf-8")
    if BEGIN not in text or END not in text:
        print("! README.md is missing the generated-results markers; skipping injection")
        return
    head, rest = text.split(BEGIN, 1)
    _, tail = rest.split(END, 1)
    readme.write_text(f"{head}{BEGIN}\n{block}\n{END}{tail}", encoding="utf-8")
    print(f"-> README.md results section updated")


if __name__ == "__main__":
    main()