src/cti_attack/data.py

"""Dataset construction: ingest -> clean -> dedupe -> filter -> split.

The output is two parallel views of the same corpus:

  * ``document`` split — no source report appears in more than one split.
  * ``random``   split — naive sentence-level shuffle.

Publishing both is the point of this repo. The corpus has 19k sentences drawn
from only 151 reports, so a sentence-level shuffle scatters near-identical
prose from one report across train and test. Measuring that gap is a result,
not a footnote.
"""

from __future__ import annotations

import json
import random
import re
import urllib.request
from collections import Counter, defaultdict
from dataclasses import dataclass, field
from pathlib import Path

from . import config

_BOILERPLATE = re.compile(config.BOILERPLATE_PREFIX_RE, re.IGNORECASE | re.DOTALL)
_BOILERPLATE_MARKER = re.compile(config.BOILERPLATE_MARKER_RE, re.IGNORECASE)
_NON_ALNUM = re.compile(r"[^a-z0-9 ]")
_WS = re.compile(r"\s+")


# --------------------------------------------------------------------------
# ingest
# --------------------------------------------------------------------------
def download_raw(force: bool = False) -> Path:
    """Fetch the TRAM multi-label corpus (Apache-2.0) if not already cached."""
    if config.TRAM_RAW.exists() and not force:
        return config.TRAM_RAW
    print(f"downloading {config.TRAM_MULTILABEL_URL}")
    urllib.request.urlretrieve(config.TRAM_MULTILABEL_URL, config.TRAM_RAW)
    return config.TRAM_RAW


def load_raw() -> list[dict]:
    with open(download_raw(), encoding="utf-8") as fh:
        return json.load(fh)


def download_single_raw(force: bool = False) -> Path:
    """Fetch the TRAM single-label corpus (Apache-2.0) if not already cached."""
    if config.TRAM_SINGLE_RAW.exists() and not force:
        return config.TRAM_SINGLE_RAW
    print(f"downloading {config.TRAM_SINGLELABEL_URL}")
    urllib.request.urlretrieve(config.TRAM_SINGLELABEL_URL, config.TRAM_SINGLE_RAW)
    return config.TRAM_SINGLE_RAW


def load_single_raw() -> list[dict]:
    """Single-label rows normalised to the multi-label schema.

    Each row is ``{"text", "label", "doc_title"}``; it becomes
    ``{"sentence", "labels": [label], "doc_title"}`` so both corpora flow through
    one cleaning/dedup/split pipeline. Dedup takes the union of labels, so a
    sentence present in both files ends up with every technique either assigned.
    """
    with open(download_single_raw(), encoding="utf-8") as fh:
        rows = json.load(fh)
    return [
        {"sentence": r["text"],
         "labels": [r["label"]] if r.get("label") else [],
         "doc_title": r["doc_title"]}
        for r in rows
    ]


# --------------------------------------------------------------------------
# cleaning
# --------------------------------------------------------------------------
def clean_sentence(text: str) -> str:
    """Strip scraped ``title: … url: …`` headers and normalise whitespace.

    Rows carrying a bare ``title:`` with no URL keep whatever headline follows —
    it is real English prose that carries no technique, which makes it a valid
    negative. Only rows that reduce to nothing are dropped, by the caller.
    """
    text = _BOILERPLATE.sub("", text)
    text = _BOILERPLATE_MARKER.sub("", text)
    return _WS.sub(" ", text).strip()


def dedup_key(text: str) -> str:
    """Aggressive normalisation used only for duplicate detection."""
    return _WS.sub(" ", _NON_ALNUM.sub("", text.lower())).strip()


# --------------------------------------------------------------------------
# statistics carried through the build so the dataset card can cite them
# --------------------------------------------------------------------------
@dataclass
class BuildStats:
    raw_sentences: int = 0
    single_label_merged: bool = False
    single_label_rows_added: int = 0
    empty_after_cleaning: int = 0
    duplicate_groups: int = 0
    duplicates_removed: int = 0
    cross_document_duplicates: int = 0
    labels_recovered_by_merge: int = 0
    dropped_techniques: list[str] = field(default_factory=list)
    dropped_label_instances: int = 0
    final_sentences: int = 0
    final_labelled: int = 0
    final_techniques: int = 0
    final_documents: int = 0

    def as_dict(self) -> dict:
        return {k: v for k, v in self.__dict__.items()}


# --------------------------------------------------------------------------
# dedupe
# --------------------------------------------------------------------------
def dedupe(records: list[dict], stats: BuildStats) -> list[dict]:
    """Collapse duplicate sentences, taking the *union* of their labels.

    Two copies of one sentence annotated ``[T1027]`` and ``[T1027, T1140]`` are
    the same sentence annotated inconsistently. Keeping the union recovers the
    label rather than silently discarding it with the duplicate row.

    The surviving row keeps the first document it was seen in, so a sentence
    can never span two documents — which would defeat the document-level split.
    """
    by_key: dict[str, dict] = {}
    seen_docs: dict[str, set[str]] = defaultdict(set)
    counts: Counter = Counter()

    for rec in records:
        key = dedup_key(rec["sentence"])
        if not key:
            continue
        counts[key] += 1
        seen_docs[key].add(rec["doc_title"])
        if key not in by_key:
            by_key[key] = {
                "sentence": rec["sentence"],
                "labels": set(rec["labels"]),
                "doc_title": rec["doc_title"],
            }
        else:
            before = len(by_key[key]["labels"])
            by_key[key]["labels"].update(rec["labels"])
            stats.labels_recovered_by_merge += len(by_key[key]["labels"]) - before

    stats.duplicate_groups = sum(1 for c in counts.values() if c > 1)
    stats.duplicates_removed = sum(c - 1 for c in counts.values() if c > 1)
    stats.cross_document_duplicates = sum(1 for k, d in seen_docs.items() if len(d) > 1)

    return [
        {"sentence": v["sentence"], "labels": sorted(v["labels"]), "doc_title": v["doc_title"]}
        for v in by_key.values()
    ]


# --------------------------------------------------------------------------
# label filtering
# --------------------------------------------------------------------------
def filter_rare_techniques(
    records: list[dict], stats: BuildStats, min_docs: int = config.MIN_DOCS_PER_TECHNIQUE
) -> list[dict]:
    """Drop techniques that occur in fewer than ``min_docs`` distinct documents.

    Such a technique cannot be split leak-free: every one of its examples ends
    up on one side of the split, so it is either untrainable or unevaluable.
    Dropping it and saying so is more honest than reporting an F1 of 0.0 for it.
    """
    tech_docs: dict[str, set[str]] = defaultdict(set)
    for rec in records:
        for lab in rec["labels"]:
            tech_docs[lab].add(rec["doc_title"])

    drop = {t for t, docs in tech_docs.items() if len(docs) < min_docs}
    stats.dropped_techniques = sorted(drop)

    for rec in records:
        kept = [l for l in rec["labels"] if l not in drop]
        stats.dropped_label_instances += len(rec["labels"]) - len(kept)
        rec["labels"] = kept
    return records


# --------------------------------------------------------------------------
# splitting
# --------------------------------------------------------------------------
def split_by_document(records: list[dict], fractions: dict, seed: int) -> dict[str, str]:
    """Group-aware multi-label stratification: assign whole documents to splits.

    Two phases, because plain greedy stratification quietly fails here.

    *Phase 1 — coverage seeding.* With a 70/15/15 target, the train split always
    shows the largest absolute label deficit, so a naive greedy pass hands it
    every document containing a rare technique and the test set ends up with
    zero examples of it. Phase 1 therefore reserves one document per technique
    for each of train, dev and test before general packing begins, rarest
    technique first. Every retained technique is then guaranteed to be both
    trainable and evaluable.

    *Phase 2 — packing.* Remaining documents go to whichever split is furthest
    below quota, measured as deficit *normalised by split size* so that the
    small dev/test splits can still compete with train for scarce labels.
    """
    doc_labels: dict[str, Counter] = defaultdict(Counter)
    doc_size: Counter = Counter()
    for rec in records:
        doc_size[rec["doc_title"]] += 1
        for lab in rec["labels"]:
            doc_labels[rec["doc_title"]][lab] += 1

    total_labels: Counter = Counter()
    for c in doc_labels.values():
        total_labels.update(c)
    total_size = sum(doc_size.values())

    desired_labels = {
        s: {t: n * f for t, n in total_labels.items()} for s, f in fractions.items()
    }
    desired_size = {s: total_size * f for s, f in fractions.items()}
    have_labels: dict[str, Counter] = {s: Counter() for s in fractions}
    have_size: Counter = Counter({s: 0 for s in fractions})
    assignment: dict[str, str] = {}

    rng = random.Random(seed)

    def place(doc: str, split: str) -> None:
        assignment[doc] = split
        have_labels[split].update(doc_labels.get(doc, Counter()))
        have_size[split] += doc_size[doc]

    # ---- phase 1: guarantee every technique reaches every split ----------
    tech_docs: dict[str, list[str]] = defaultdict(list)
    for doc, labs in doc_labels.items():
        for t in labs:
            tech_docs[t].append(doc)

    # test and dev are seeded before train: they are the splits that starve.
    for tech in sorted(tech_docs, key=lambda t: (len(tech_docs[t]), t)):
        for split in ("test", "dev", "train"):
            if have_labels[split][tech] > 0:
                continue
            candidates = [d for d in tech_docs[tech] if d not in assignment]
            if not candidates:
                continue  # already spent; phase 2 cannot recover it
            rng.shuffle(candidates)
            # spend the cheapest document that satisfies the requirement, and
            # prefer one that also fits the split's remaining size budget
            candidates.sort(
                key=lambda d: (
                    sum(doc_labels[d].values()),
                    abs((desired_size[split] - have_size[split]) - doc_size[d]),
                )
            )
            place(candidates[0], split)

    # ---- phase 2: pack the rest, rarest-first ----------------------------
    remaining = [d for d in doc_size if d not in assignment]
    rng.shuffle(remaining)

    def rarity(doc: str) -> tuple:
        labs = doc_labels.get(doc)
        if not labs:
            return (1, 0, 0)  # unlabelled documents are placed last, on size only
        return (0, min(total_labels[t] for t in labs), -sum(labs.values()))

    remaining.sort(key=rarity)

    for doc in remaining:
        labs = doc_labels.get(doc, Counter())

        def need(s: str, _labs: Counter = labs) -> tuple:
            # normalising by the split fraction converts "train is biggest" into
            # a fair comparison of how starved each split actually is
            deficit = sum(
                max(0.0, desired_labels[s][t] - have_labels[s][t]) for t in _labs
            ) / fractions[s]
            size_room = (desired_size[s] - have_size[s]) / fractions[s]
            return (deficit, size_room)

        place(doc, max(fractions, key=need))

    return assignment


def split_random(records: list[dict], fractions: dict, seed: int) -> list[str]:
    """Naive sentence-level shuffle — the split this repo argues against."""
    idx = list(range(len(records)))
    random.Random(seed).shuffle(idx)
    n = len(idx)
    n_train = int(n * fractions["train"])
    n_dev = int(n * fractions["dev"])
    out = [""] * n
    for rank, i in enumerate(idx):
        if rank < n_train:
            out[i] = "train"
        elif rank < n_train + n_dev:
            out[i] = "dev"
        else:
            out[i] = "test"
    return out


# --------------------------------------------------------------------------
# build
# --------------------------------------------------------------------------
def build(verbose: bool = True, include_single: bool = False) -> tuple[list[dict], list[str], BuildStats]:
    stats = BuildStats()
    raw = load_raw()
    if include_single:
        single = load_single_raw()
        raw = raw + single
        stats.single_label_merged = True
        stats.single_label_rows_added = len(single)
    stats.raw_sentences = len(raw)

    cleaned = []
    for rec in raw:
        text = clean_sentence(rec["sentence"])
        if not text:
            stats.empty_after_cleaning += 1
            continue
        cleaned.append({"sentence": text, "labels": rec["labels"], "doc_title": rec["doc_title"]})

    records = dedupe(cleaned, stats)
    records = filter_rare_techniques(records, stats)

    labels = sorted({l for r in records for l in r["labels"]})
    doc_assign = split_by_document(records, config.SPLIT_FRACTIONS, config.SPLIT_SEED)
    rand_assign = split_random(records, config.SPLIT_FRACTIONS, config.SPLIT_SEED)

    for rec, rnd in zip(records, rand_assign):
        rec["split_document"] = doc_assign[rec["doc_title"]]
        rec["split_random"] = rnd

    stats.final_sentences = len(records)
    stats.final_labelled = sum(1 for r in records if r["labels"])
    stats.final_techniques = len(labels)
    stats.final_documents = len({r["doc_title"] for r in records})

    if verbose:
        _print_report(records, labels, stats)
    return records, labels, stats


def _print_report(records: list[dict], labels: list[str], stats: BuildStats) -> None:
    print("\n=== build report ===")
    for k, v in stats.as_dict().items():
        print(f"  {k:30} {v}")
    print("\n=== split sizes ===")
    for scheme in ("split_document", "split_random"):
        c = Counter(r[scheme] for r in records)
        lab = Counter(r[scheme] for r in records if r["labels"])
        print(f"  {scheme}")
        for s in ("train", "dev", "test"):
            print(f"    {s:6} {c[s]:6} sentences  {lab[s]:5} labelled")
    print("\n=== technique coverage under document split ===")
    missing = []
    for scheme in ("split_document", "split_random"):
        seen = {s: set() for s in ("train", "dev", "test")}
        for r in records:
            for l in r["labels"]:
                seen[r[scheme]].add(l)
        gaps = {s: len(set(labels) - seen[s]) for s in seen}
        print(f"  {scheme}: techniques absent from train/dev/test = "
              f"{gaps['train']}/{gaps['dev']}/{gaps['test']}")
        if scheme == "split_document":
            missing = sorted(set(labels) - seen["test"])
    if missing:
        print(f"    absent from document-split test set: {missing}")


def write(records: list[dict], labels: list[str], stats: BuildStats) -> None:
    for scheme in ("document", "random"):
        for split in ("train", "dev", "test"):
            path = config.BUILD_DIR / f"{scheme}_{split}.jsonl"
            rows = [r for r in records if r[f"split_{scheme}"] == split]
            with open(path, "w", encoding="utf-8") as fh:
                for r in rows:
                    fh.write(json.dumps(
                        {"sentence": r["sentence"], "labels": r["labels"],
                         "doc_title": r["doc_title"]}, ensure_ascii=False) + "\n")
            print(f"wrote {path.name:26} {len(rows):6} rows")

    (config.BUILD_DIR / "labels.json").write_text(
        json.dumps(labels, indent=2), encoding="utf-8")
    (config.BUILD_DIR / "build_stats.json").write_text(
        json.dumps(stats.as_dict(), indent=2), encoding="utf-8")


def load_split(scheme: str, split: str) -> list[dict]:
    path = config.BUILD_DIR / f"{scheme}_{split}.jsonl"
    with open(path, encoding="utf-8") as fh:
        return [json.loads(line) for line in fh]


def load_labels() -> list[str]:
    return json.loads((config.BUILD_DIR / "labels.json").read_text(encoding="utf-8"))