#!/usr/bin/env python
"""Data preparation for the BITS Phishing Detector.

Download -> normalise -> dedupe -> decontaminate -> split -> leakage-check -> manifests.
Procedure and expected counts: docs/research/verify_data.md sections 7 and 8.

Subcommands
  inspect   features + 2 rows of all four datasets (streaming; nothing is fully downloaded)
  build     produce data/processed/{pretrain,finetune_train,finetune_val,test_holdout}.parquet,
            splits_manifest.json and data_report.md
  check     re-run the leakage assertions on the written files and print the manifest summary
  build3    THREE-WAY labels (RUN_LOG entry 6): add label3/label_set to the binary splits, ingest the Champa
            phishing corpora (Nazario, Nigerian_Fraud) [+ bigshp if data.threeway.bigshp.enabled], carve the
            clean-provenance test_holdout3.parquet, rewrite finetune_train/val + test_holdout (2 columns added)
  check3    recompute every binary + three-way leakage assertion from the written files

Leakage contract (hard failures, asserted identically by `build` and `check`): the held-out
TREC-06 sub-source never shares an exact key (k_full / k_pre / k_sb / k_sb_full) or a >=3-shingle
near-duplicate link with the fine-tune train/val splits or the MLM pretraining corpus; train and val
never share a k_full; val never shares a key with the pretraining corpus; the holdout source never
appears in train/val. Info-only counters: train/val near-dup links, train/pretrain key overlap,
val/pretrain near-dup links.

Three-way contract (build3 / check3): label3 in {0 legitimate, 1 phishing, 2 spam, -1 unwanted-but-unsplit},
label_set in {legit, phish, spam, unwanted}; binary `label` unchanged. Binary split membership is unchanged
except for the rows moved from finetune_val into test_holdout3. test_holdout3 is asserted to share zero keys and
zero >=3-shingle links with train, val and the pretraining corpus and zero keys with the TREC-06 holdout.

The assertions are an independent recomputation, not the removal bookkeeping: after every removal
step the shingle index is rebuilt from the surviving rows (and the pretraining rows are re-derived
from their stored `text` exactly as `check` does), so a bug in the removal logic surfaces as a
failed assertion rather than a silent pass.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import logging
import os
import re
import resource
import shutil
import sys
import unicodedata
from collections import Counter, defaultdict
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Iterable

import numpy as np
import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq

sys.path.insert(0, str(Path(__file__).parent))
import common  # noqa: E402

LOG = logging.getLogger("data_prep")

# ----------------------------------------------------------------------------- dials
@dataclass
class Dials:
    """Dedup / pretraining thresholds. Read from config.yaml when present, else these defaults.

    Optional config keys (all under `data`): `dedup.key_prefix_chars`, `dedup.key_prefix_min_len`,
    `dedup.near_dup_shingle_cap`, `pretrain_min_alnum_chars`, `quote_strip_min_alnum`. The body-key
    fallback threshold is always `data.min_alnum_chars` so the two cannot diverge.
    """

    key_prefix_chars: int = 300        # k_pre / k_sb hash the first N normalised chars
    key_prefix_min_len: int = 80       # shorter bases use k_full as k_pre (a prefix would be the whole text)
    body_key_min_alnum: int = 20       # below this the body alone is degenerate; key on subject+body instead
    pretrain_min_alnum: int = 50       # verify_data.md section 7 step 2: corbt bodies with <50 alnum are dropped
    quote_strip_min_alnum: int = 200   # a quote/forward cut is applied only if this much text survives above it
    shingle_cap: int = 2000            # shingles shared by more rows than this are boilerplate and ignored

    @classmethod
    def from_config(cls, cfg: dict) -> "Dials":
        d = cfg["data"]
        dd = d.get("dedup", {})
        return cls(
            key_prefix_chars=int(dd.get("key_prefix_chars", cls.key_prefix_chars)),
            key_prefix_min_len=int(dd.get("key_prefix_min_len", cls.key_prefix_min_len)),
            body_key_min_alnum=int(d["min_alnum_chars"]),
            pretrain_min_alnum=int(d.get("pretrain_min_alnum_chars", cls.pretrain_min_alnum)),
            quote_strip_min_alnum=int(d.get("quote_strip_min_alnum", cls.quote_strip_min_alnum)),
            shingle_cap=int(dd.get("near_dup_shingle_cap", cls.shingle_cap)),
        )


DIALS = Dials()   # replaced by Dials.from_config(cfg) in main(); module-level so row_keys stays a plain function

LABELLED_COLUMNS = ["text", "subject", "label", "dataset_name"]
PRETRAIN_COLUMNS = ["subject", "body"]
FINETUNE_SCHEMA = ["doc_id", "text", "subject", "body", "label", "source", "cluster_id"]
PRETRAIN_SCHEMA = ["doc_id", "text", "n_chars"]

_CRLF = re.compile(r"\r\n?")
_BLANK_RUN = re.compile(r"\n(?:[ \t]*\n){2,}")   # 3+ newlines (2+ blank lines) -> one blank line
_QUOTE_MARKERS = re.compile(
    r"(?P<original>-----\s*Original Message\s*-----)"
    r"|(?P<forwarded>-{3,}\s*Forwarded by\b)"
    r"|(?P<header_block>^From:[^\n]*\n(?:[^\n]*\n){0,2}(?:Sent|To|Subject):)",
    re.MULTILINE,
)


class SpamPolicyError(ValueError):
    """Raised for spam policies that the available corpora cannot implement."""


# ----------------------------------------------------------------------------- text utils
def clean_text(s: Any) -> str:
    """NFKC, normalised line endings, collapsed blank-line runs, stripped.

    Keeps the email readable for a byte-level BPE (no aggressive stripping of punctuation or
    HTML) while removing the formatting noise that would otherwise make identical emails look
    different to the tokenizer and inflate the near-duplicate count.
    """
    if not isinstance(s, str) or not s:
        return ""
    s = unicodedata.normalize("NFKC", s).replace("\x00", "")
    s = _CRLF.sub("\n", s)
    s = _BLANK_RUN.sub("\n\n", s)
    return s.strip()


def model_input(subject: str, body: str) -> str:
    """subject + blank line + body. Never sender/receiver/date (TREC headers are synthetic)."""
    return "\n\n".join(part for part in (subject, body) if part)


def strip_quotes(body: str) -> tuple[str, str | None, bool]:
    """Cut a corbt body at its first quote/forward marker when >=200 alphanumerics survive.

    Enron threads quote the whole history, so without this ~43% of pretraining characters would
    be verbatim repeats of other rows. Returns (body, marker_name, cut_applied).
    """
    m = _QUOTE_MARKERS.search(body)
    if m is None:
        return body, None, False
    head = body[: m.start()].rstrip()
    if len(common.normalize_text(head)) >= DIALS.quote_strip_min_alnum:
        return head, m.lastgroup, True
    return body, m.lastgroup, False


def row_keys(subject: str, body: str) -> tuple[str, str, str, str, str, int]:
    """Dedup keys per verify_data.md section 7.

    Returns (k_full, k_pre, k_sb, k_sb_full, base, n_alnum):
      k_full    sha1 of the normalised body (the content-level identity of the email)
      k_pre     sha1 of its first 300 chars (catches appended footers / truncated copies), or
                k_full when the body is shorter than 80 chars
      k_sb      sha1 of the first 300 chars of normalised subject+body (catches corpora that
                embed the subject in the body, e.g. Enron-Spam, and corbt subject+body)
      k_sb_full sha1 of the whole normalised subject+body (corbt dedup key)
      base      the normalised string used for k_full/k_pre and shingles
      n_alnum   alphanumerics in subject+body (the min_alnum_chars filter)
    When the body holds fewer than `min_alnum_chars` alphanumerics the subject+body string is used
    as `base` instead, otherwise every empty-body email would collide on sha1("").
    """
    norm_body = common.normalize_text(body)
    norm_sb = common.normalize_text(subject) + norm_body
    base = norm_body if len(norm_body) >= DIALS.body_key_min_alnum else norm_sb
    k_full = common.sha1_hex(base)
    k_pre = common.sha1_hex(base[: DIALS.key_prefix_chars]) if len(base) >= DIALS.key_prefix_min_len else k_full
    k_sb = common.sha1_hex(norm_sb[: DIALS.key_prefix_chars])
    k_sb_full = common.sha1_hex(norm_sb)
    return k_full, k_pre, k_sb, k_sb_full, base, len(norm_sb)


def split_text(text: str) -> tuple[str, str]:
    """Recover (subject, body) from a stored `text` = model_input(subject, body).

    Exact whenever the subject was non-empty (subjects never contain a blank line). For an
    empty-subject email with a multi-paragraph body the first paragraph is mistaken for the
    subject; `text_identity` is therefore only a lower bound on the build-time identity, and
    build_pretrain compensates by matching on BOTH identities for empty-subject rows so the
    post-hoc assertion (shared with `check`) can never see a link the removal pass did not.
    """
    head, sep, rest = text.partition("\n\n")
    return (head, rest) if sep else ("", head)


def text_identity(text: str) -> tuple[set[str], str]:
    """(keys, shingle base) of a stored pretraining `text`, recomputed the same way by build and check.

    normalize(subject + "\\n\\n" + body) == normalize(subject) + normalize(body), so k_sb / k_sb_full
    are always exact; k_full / k_pre / base are exact unless the subject was empty (see split_text).
    """
    k_full, k_pre, k_sb, k_sb_full, base, _ = row_keys(*split_text(text))
    return {k_full, k_pre, k_sb, k_sb_full}, base


def add_key_columns(df: pd.DataFrame) -> list[str]:
    """Attach k_full/k_pre/k_sb/k_sb_full/n_alnum columns; return the shingle bases (window-trimmed)."""
    cols: dict[str, list] = defaultdict(list)
    bases: list[str] = []
    for subj, body in zip(df["subject"].tolist(), df["body"].tolist()):
        k_full, k_pre, k_sb, k_sb_full, base, n_alnum = row_keys(subj, body)
        cols["k_full"].append(k_full)
        cols["k_pre"].append(k_pre)
        cols["k_sb"].append(k_sb)
        cols["k_sb_full"].append(k_sb_full)
        cols["n_alnum"].append(n_alnum)
        bases.append(base)
    for name, values in cols.items():
        df[name] = values
    return bases


def key_union(df: pd.DataFrame) -> set[str]:
    out: set[str] = set()
    for col in ("k_full", "k_pre", "k_sb", "k_sb_full"):
        out.update(df[col].tolist())
    return out


def rows_hitting(df: pd.DataFrame, keys: set[str]) -> np.ndarray:
    """Boolean mask: row has any of its four keys inside `keys`."""
    mask = np.zeros(len(df), dtype=bool)
    for col in ("k_full", "k_pre", "k_sb", "k_sb_full"):
        mask |= np.fromiter((k in keys for k in df[col].tolist()), dtype=bool, count=len(df))
    return mask


# ----------------------------------------------------------------------------- near-dup index
@dataclass
class NearDupConfig:
    shingle_chars: int = 50
    stride: int = 25
    window_chars: int = 1500
    min_shared: int = 3
    cap: int = 2000

    @classmethod
    def from_config(cls, cfg: dict, cap: int) -> "NearDupConfig":
        d = cfg["data"]["dedup"]
        return cls(
            shingle_chars=int(d["near_dup_shingle_chars"]),
            stride=int(d["near_dup_stride"]),
            window_chars=int(d["near_dup_window_chars"]),
            min_shared=int(d["near_dup_min_shared"]),
            cap=cap,
        )


def shingle_hashes(base: str, nd: NearDupConfig) -> np.ndarray:
    """Distinct 64-bit hashes of the fixed-stride shingles of base[:window]."""
    s = base[: nd.window_chars]
    n = len(s)
    if n < nd.shingle_chars:
        return np.empty(0, dtype=np.int64)
    digests = b"".join(
        hashlib.blake2b(s[i : i + nd.shingle_chars].encode("ascii"), digest_size=8).digest()
        for i in range(0, n - nd.shingle_chars + 1, nd.stride)
    )
    return np.unique(np.frombuffer(digests, dtype=np.int64))


class ShingleIndex:
    """Inverted shingle index over a list of normalised texts.

    Shingles shared by more than `cap` rows are boilerplate (list footers, pharma templates) and
    are ignored: linking everything that carries the same footer would merge half a corpus into
    one cluster and remove legitimate training rows. This is a known, logged limitation.
    """

    def __init__(self, bases: list[str], nd: NearDupConfig):
        self.nd = nd
        self.n = len(bases)
        per_row = [shingle_hashes(b, nd) for b in bases]
        lens = np.fromiter((len(h) for h in per_row), dtype=np.int64, count=self.n)
        self.offsets = np.concatenate(([0], np.cumsum(lens)))
        self.hashes = np.concatenate(per_row) if self.n else np.empty(0, dtype=np.int64)
        rid = np.repeat(np.arange(self.n, dtype=np.int64), lens)
        order = np.argsort(self.hashes, kind="stable")
        self.sorted_rid = rid[order]
        self.uniq, self.start, self.counts = np.unique(
            self.hashes[order], return_index=True, return_counts=True
        )
        capped = self.counts > nd.cap
        self.n_capped_shingles = int(capped.sum())
        self.n_rows_touching_capped = int(
            np.unique(self.sorted_rid[np.repeat(capped, self.counts)]).size
        )
        LOG.info(
            "shingle index: %d rows, %d shingle occurrences, %d distinct, %d shingles over cap=%d "
            "(ignored; touching %d rows)",
            self.n, self.hashes.size, self.uniq.size, self.n_capped_shingles, nd.cap,
            self.n_rows_touching_capped,
        )

    def links(self) -> tuple[np.ndarray, np.ndarray]:
        """All pairs (a < b) sharing >= min_shared distinct shingles."""
        out_a: list[np.ndarray] = []
        out_b: list[np.ndarray] = []
        t0 = common.timer()
        for r in range(self.n):
            h = self.hashes[self.offsets[r] : self.offsets[r + 1]]
            if h.size == 0:
                continue
            g = np.searchsorted(self.uniq, h)
            cnt = self.counts[g]
            ok = (cnt >= 2) & (cnt <= self.nd.cap)
            if not ok.any():
                continue
            starts = self.start[g[ok]]
            lens = cnt[ok]
            idx = np.repeat(starts - np.concatenate(([0], np.cumsum(lens)[:-1])), lens) + np.arange(lens.sum())
            cand = self.sorted_rid[idx]
            cand = cand[cand > r]
            if cand.size < self.nd.min_shared:
                continue
            u, c = np.unique(cand, return_counts=True)
            hit = u[c >= self.nd.min_shared]
            if hit.size:
                out_a.append(np.full(hit.size, r, dtype=np.int64))
                out_b.append(hit)
            if r and r % 50_000 == 0:
                LOG.info("  near-dup scan %d/%d rows (%.1fs)", r, self.n, common.timer() - t0)
        if not out_a:
            return np.empty(0, dtype=np.int64), np.empty(0, dtype=np.int64)
        a, b = np.concatenate(out_a), np.concatenate(out_b)
        LOG.info("near-dup links: %d pairs in %.1fs", a.size, common.timer() - t0)
        return a, b

    def query(self, bases: list[str]) -> tuple[np.ndarray, np.ndarray]:
        """Pairs (query_idx, indexed_idx) where a foreign text shares >= min_shared shingles with an indexed row.

        Used to screen the (large) pretraining corpus against the (small) holdout+val reference set
        without indexing the corpus itself: cost is one searchsorted per query row, so the millions
        of pretrain-pretrain Enron-thread links are never materialised.
        """
        out_q: list[np.ndarray] = []
        out_r: list[np.ndarray] = []
        t0 = common.timer()
        for qi, base in enumerate(bases):
            h = shingle_hashes(base, self.nd)
            if h.size == 0 or self.uniq.size == 0:
                continue
            g = np.searchsorted(self.uniq, h)
            inb = g < self.uniq.size
            g, h = g[inb], h[inb]
            g = g[self.uniq[g] == h]          # foreign shingles absent from the index drop out here
            if g.size == 0:
                continue
            cnt = self.counts[g]
            ok = cnt <= self.nd.cap
            if not ok.any():
                continue
            starts = self.start[g[ok]]
            lens = cnt[ok]
            idx = np.repeat(starts - np.concatenate(([0], np.cumsum(lens)[:-1])), lens) + np.arange(lens.sum())
            u, c = np.unique(self.sorted_rid[idx], return_counts=True)
            hit = u[c >= self.nd.min_shared]
            if hit.size:
                out_q.append(np.full(hit.size, qi, dtype=np.int64))
                out_r.append(hit)
            if qi and qi % 100_000 == 0:
                LOG.info("  near-dup query %d/%d rows (%.1fs)", qi, len(bases), common.timer() - t0)
        if not out_q:
            return np.empty(0, dtype=np.int64), np.empty(0, dtype=np.int64)
        q, r = np.concatenate(out_q), np.concatenate(out_r)
        LOG.info("near-dup query: %d rows -> %d links in %.1fs", len(bases), q.size, common.timer() - t0)
        return q, r


def recount_shingle_links(
    train_bases: list[str], val_bases: list[str], hold_bases: list[str], pretrain_bases: list[str], nd: NearDupConfig
) -> dict[str, int]:
    """Independent near-dup link counts between the FINAL sets, from a fresh index.

    This is the leakage gate proper: it never reuses the index or masks that drove the removals,
    so it fails if the removal logic is wrong. Same code path in `build` (after writing nothing yet)
    and in `check` (from the parquet files). Pretrain rows are queried against a holdout+val index
    rather than indexed, so pretrain-pretrain links are never enumerated.
    """
    n_tr, n_va = len(train_bases), len(val_bases)
    index = ShingleIndex(train_bases + val_bases + hold_bases, nd)
    la, lb = index.links()
    part_a = np.where(la < n_tr, 0, np.where(la < n_tr + n_va, 1, 2))
    part_b = np.where(lb < n_tr, 0, np.where(lb < n_tr + n_va, 1, 2))
    pair = Counter(zip(np.minimum(part_a, part_b).tolist(), np.maximum(part_a, part_b).tolist()))
    ref = ShingleIndex(hold_bases + val_bases, nd)
    _q, r = ref.query(pretrain_bases)
    out = {
        "holdout_train": int(pair[(0, 2)]),
        "holdout_val": int(pair[(1, 2)]),
        "train_val": int(pair[(0, 1)]),
        "holdout_pretrain": int(np.sum(r < len(hold_bases))),
        "val_pretrain": int(np.sum(r >= len(hold_bases))),
    }
    LOG.info("recount shingle links (final sets): %s", out)
    return out


def union_find(n: int, pairs: Iterable[tuple[int, int]]) -> np.ndarray:
    """Dense cluster ids (0..K-1, in order of first appearance) from an edge list."""
    parent = list(range(n))

    def find(x: int) -> int:
        root = x
        while parent[root] != root:
            root = parent[root]
        while parent[x] != root:
            parent[x], x = root, parent[x]
        return root

    for a, b in pairs:
        ra, rb = find(a), find(b)
        if ra != rb:
            parent[max(ra, rb)] = min(ra, rb)
    roots = np.fromiter((find(i) for i in range(n)), dtype=np.int64, count=n)
    _, dense = np.unique(roots, return_inverse=True)
    return dense.astype(np.int64)


def same_key_pairs(keys: list[str]) -> Iterable[tuple[int, int]]:
    """Consecutive-member edges for every group of rows sharing a key."""
    first: dict[str, int] = {}
    for i, k in enumerate(keys):
        j = first.setdefault(k, i)
        if j != i:
            yield j, i


def same_key_pairs_multi(columns: list[list[str]]) -> Iterable[tuple[int, int]]:
    """Edges for rows sharing a key value within OR across the given key columns.

    k_pre (first 300 chars of the body) and k_sb (first 300 chars of subject+body) can collide across columns
    when one corpus embeds the subject in the body; the leakage gate's key_union treats such a pair as shared,
    so the clustering must too or two near-identical rows could land on different sides of a split.
    """
    first: dict[str, int] = {}
    for col in columns:
        for i, k in enumerate(col):
            j = first.setdefault(k, i)
            if j != i:
                yield j, i


# ----------------------------------------------------------------------------- loading
def hub_sha(repo: str) -> str | None:
    try:
        from huggingface_hub import HfApi

        return HfApi().dataset_info(repo).sha
    except Exception as exc:  # offline or rate-limited: record unknown, keep going
        LOG.warning("could not fetch commit sha for %s: %s", repo, exc)
        return None


def load_labelled(cfg: dict, limit: int | None, rng: np.random.Generator) -> tuple[pd.DataFrame, dict]:
    """All Hub splits of the primary set concatenated; the Hub's own split is ignored (it leaks)."""
    from datasets import concatenate_datasets, load_dataset

    repo = cfg["data"]["labelled_primary"]
    cache_dir = str(Path(cfg["paths"]["data_raw"]) / "hf_cache")
    sha = hub_sha(repo)
    dsd = load_dataset(repo, cache_dir=cache_dir, revision=sha)
    split_names = list(dsd.keys())
    parts = [dsd[s].select_columns(LABELLED_COLUMNS) for s in split_names]
    ds = concatenate_datasets(parts)
    df = ds.to_pandas()
    df = df.rename(columns={"text": "body", "dataset_name": "source"})
    df["row_idx"] = np.arange(len(df), dtype=np.int64)
    meta = {
        "repo": repo,
        "commit_sha": sha,
        "hub_splits_concatenated": {s: int(dsd[s].num_rows) for s in split_names},
        "datasets_fingerprints": {s: dsd[s]._fingerprint for s in split_names},
        "rows_loaded": int(len(df)),
    }
    LOG.info("labelled primary %s @ %s: %d rows from splits %s", repo, sha, len(df), meta["hub_splits_concatenated"])
    if limit:
        keep = []
        for source, grp in df.groupby("source", sort=True):
            n = min(limit, len(grp))
            keep.append(grp.iloc[np.sort(rng.choice(len(grp), n, replace=False))])
        df = pd.concat(keep).sort_values("row_idx").reset_index(drop=True)
        LOG.warning("--limit %d: labelled subsample %d rows (per source <= %d)", limit, len(df), limit)
    return df, meta


def iter_pretrain_batches(cfg: dict, limit: int | None, rng: np.random.Generator, batch_size: int = 20_000):
    """Yield (row_indices, subjects, bodies) batches of corbt without materialising the corpus."""
    from datasets import load_dataset

    repo = cfg["data"]["pretrain_corpus"]
    cache_dir = str(Path(cfg["paths"]["data_raw"]) / "hf_cache")
    sha = hub_sha(repo)
    ds = load_dataset(repo, split="train", cache_dir=cache_dir, revision=sha)
    meta = {"repo": repo, "commit_sha": sha, "datasets_fingerprint": ds._fingerprint, "rows_total": int(ds.num_rows)}
    indices = np.arange(ds.num_rows, dtype=np.int64)
    if limit:
        n = min(limit * 4, ds.num_rows)
        indices = np.sort(rng.choice(ds.num_rows, n, replace=False))
        ds = ds.select(indices)
        LOG.warning("--limit %d: corbt subsample %d rows", limit, n)
    ds = ds.select_columns(PRETRAIN_COLUMNS)
    LOG.info("pretrain corpus %s @ %s: %d rows total, processing %d", repo, sha, meta["rows_total"], len(indices))

    def batches():
        pos = 0
        for batch in ds.iter(batch_size=batch_size):
            n = len(batch["body"])
            yield indices[pos : pos + n], batch["subject"], batch["body"]
            pos += n

    return meta, batches()


# ----------------------------------------------------------------------------- policy
def check_spam_policy(policy: str) -> None:
    """Only `positive_class` is implementable; the others are kept so the decision stays visible."""
    if policy == "positive_class":
        return
    if policy in ("map_to_zero", "exclude"):
        raise SpamPolicyError(
            f"spam_policy={policy!r} cannot be implemented: no dataset in scope distinguishes spam "
            "from phishing. Every label-1 row of puyang2025/seven-phishing-email-datasets comes from "
            "a spam-track corpus (TREC-05/06/07, CEAS-08, Enron-Spam, SpamAssassin, Ling-Spam) with a "
            "single binary label, so mapping spam to 0 or excluding spam sources would leave ~0 "
            "positives (verify_data.md section 5). Use positive_class (positive = phishing/spam, "
            "'unwanted, potentially malicious') and report the credential-lure slice separately."
        )
    raise SpamPolicyError(f"unknown spam_policy {policy!r}; expected positive_class|map_to_zero|exclude")


# ----------------------------------------------------------------------------- split
def grouped_stratified_split(
    pool: pd.DataFrame, val_fraction: float, rng: np.random.Generator
) -> tuple[np.ndarray, dict]:
    """Assign whole near-dup clusters to val until each source x label stratum reaches ~val_fraction.

    A cluster never straddles train/val, so a template seen in training can never inflate the
    validation score. Clusters spanning several strata are assigned under their majority stratum;
    a cluster larger than a stratum's whole validation target goes to train.
    """
    strata = pool["source"].astype(str) + "|" + pool["label"].astype(str)
    cluster = pool["cluster_id"].to_numpy()
    members: dict[int, list[int]] = defaultdict(list)
    for pos, cid in enumerate(cluster.tolist()):
        members[int(cid)].append(pos)
    stratum_arr = strata.to_numpy()
    by_stratum: dict[str, list[tuple[int, int]]] = defaultdict(list)
    mixed = 0
    for cid, pos_list in members.items():
        counts = Counter(stratum_arr[p] for p in pos_list)
        if len(counts) > 1:
            mixed += 1
        by_stratum[counts.most_common(1)[0][0]].append((cid, len(pos_list)))

    is_val = np.zeros(len(pool), dtype=bool)
    per_stratum: dict[str, dict[str, Any]] = {}
    for stratum in sorted(by_stratum):
        groups = by_stratum[stratum]
        n_rows = sum(size for _, size in groups)
        target = int(round(val_fraction * n_rows))
        order = rng.permutation(len(groups))
        val_rows = 0
        oversized = 0
        for gi in order:
            cid, size = groups[gi]
            if size > target:
                oversized += 1
                continue
            if val_rows < target and abs(val_rows + size - target) <= abs(val_rows - target):
                for p in members[cid]:
                    is_val[p] = True
                val_rows += size
        per_stratum[stratum] = {
            "rows": n_rows, "clusters": len(groups), "val_target": target, "val_rows": val_rows,
            "val_fraction": round(val_rows / n_rows, 4) if n_rows else 0.0,
            "clusters_too_large_for_val": oversized,
        }
    info = {
        "val_fraction_target": val_fraction,
        "val_fraction_achieved": round(float(is_val.mean()), 4) if len(pool) else 0.0,
        "clusters_total": len(members),
        "clusters_mixed_stratum": mixed,
        "largest_cluster": max((len(v) for v in members.values()), default=0),
        "per_stratum": per_stratum,
    }
    return is_val, info


# ----------------------------------------------------------------------------- bookkeeping
def nested_counts(df: pd.DataFrame) -> dict[str, dict[str, int]]:
    """{source: {label: n}} with string keys for JSON."""
    out: dict[str, dict[str, int]] = {}
    if len(df) == 0:
        return out
    for (source, label), n in df.groupby(["source", "label"], sort=True).size().items():
        out.setdefault(str(source), {})[str(label)] = int(n)
    return out


def split_totals(df: pd.DataFrame) -> dict[str, Any]:
    n = int(len(df))
    pos = int((df["label"] == 1).sum()) if n else 0
    return {"rows": n, "positives": pos, "negatives": n - pos, "base_rate": round(pos / n, 4) if n else None}


def md_table(headers: list[str], rows: list[list[Any]]) -> str:
    lines = ["| " + " | ".join(headers) + " |", "|" + "|".join("---" for _ in headers) + "|"]
    for row in rows:
        lines.append("| " + " | ".join(f"{v:,}" if isinstance(v, int) and not isinstance(v, bool) else str(v) for v in row) + " |")
    return "\n".join(lines)


def source_label_table(counts: dict[str, dict[str, int]]) -> str:
    rows = []
    for source in sorted(counts):
        neg = counts[source].get("0", 0)
        pos = counts[source].get("1", 0)
        tot = neg + pos
        rows.append([source, neg, pos, tot, f"{pos / tot:.1%}" if tot else "-"])
    return md_table(["source", "label 0", "label 1", "rows", "positive rate"], rows)


def peak_rss_gb() -> float:
    return round(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 2**20, 2)


# ----------------------------------------------------------------------------- leakage
@dataclass
class LeakageResults:
    checks: dict[str, dict[str, Any]] = field(default_factory=dict)

    def add(self, name: str, count: int, assert_zero: bool, detail: str = "") -> None:
        self.checks[name] = {"count": int(count), "asserted": assert_zero, "pass": (count == 0) or not assert_zero, "detail": detail}
        level = logging.ERROR if (assert_zero and count) else logging.INFO
        LOG.log(level, "leakage %-44s %8d %s", name, count, "(assert zero)" if assert_zero else "(info)")

    def failures(self) -> list[str]:
        return [k for k, v in self.checks.items() if v["asserted"] and v["count"]]

    def raise_if_failed(self) -> None:
        bad = self.failures()
        if bad:
            raise AssertionError("LEAKAGE ASSERTIONS FAILED: " + ", ".join(f"{k}={self.checks[k]['count']}" for k in bad))


def leakage_checks(
    train: pd.DataFrame,
    val: pd.DataFrame,
    holdout: pd.DataFrame,
    pretrain_keys: set[str],
    n_pretrain: int,
    shingle_links: dict[str, int],
    holdout_source: str,
) -> LeakageResults:
    """The gate. `shingle_links` must come from recount_shingle_links over the final sets and
    `pretrain_keys` from text_identity over the final pretraining texts, in build and check alike."""
    res = LeakageResults()
    hold_keys = key_union(holdout)
    val_keys = key_union(val)
    train_keys = key_union(train)
    res.add("holdout_vs_train_key_rows", int(rows_hitting(train, hold_keys).sum()), True)
    res.add("holdout_vs_val_key_rows", int(rows_hitting(val, hold_keys).sum()), True)
    res.add("holdout_vs_pretrain_key_rows", len(pretrain_keys & hold_keys), True, "pretrain keys re-derived from stored text (text_identity)")
    res.add("holdout_vs_train_shingle_links", shingle_links["holdout_train"], True)
    res.add("holdout_vs_val_shingle_links", shingle_links["holdout_val"], True)
    res.add("holdout_vs_pretrain_shingle_links", shingle_links["holdout_pretrain"], True, "pretrain bases re-derived from stored text (text_identity)")
    res.add("train_vs_val_k_full_rows", len(set(train["k_full"]) & set(val["k_full"])), True)
    res.add("train_vs_val_any_key_rows", len(train_keys & val_keys), False)
    res.add("train_vs_val_shingle_links", shingle_links["train_val"], False)
    res.add("val_vs_pretrain_key_rows", len(pretrain_keys & val_keys), True)
    res.add("val_vs_pretrain_shingle_links", shingle_links["val_pretrain"], False, "near-twins of val emails in MLM text; counted, not removed (val is model selection, the holdout is the gate)")
    res.add("train_vs_pretrain_key_overlap", len(pretrain_keys & train_keys), False, "allowed: TREC-05/Enron ham are Enron mail")
    res.add("holdout_rows_in_train_or_val_by_source", int((train["source"] == holdout_source).sum() + (val["source"] == holdout_source).sum()), True)
    res.add("test_holdout_rows_from_other_sources", int((holdout["source"] != holdout_source).sum()), True)
    res.add("pretrain_rows", n_pretrain, False)
    return res


# ----------------------------------------------------------------------------- build
def cmd_build(args: argparse.Namespace, cfg: dict) -> None:
    t_start = common.timer()
    policy = args.spam_policy or cfg["data"]["spam_policy"]
    check_spam_policy(policy)
    seed = int(cfg["seed"])
    common.set_seed(seed)
    rng = np.random.default_rng(seed)
    strip = not args.no_strip_quotes
    nd = NearDupConfig.from_config(cfg, cap=args.shingle_cap or DIALS.shingle_cap)
    holdout_source = cfg["data"]["holdout_test_source"]
    min_alnum = int(cfg["data"]["min_alnum_chars"])
    val_fraction = float(cfg["data"]["finetune_val_fraction"])
    out_dir = Path(cfg["paths"]["data_processed"])
    out_dir.mkdir(parents=True, exist_ok=True)
    LOG.info("build: spam_policy=%s strip_quotes=%s limit=%s seed=%d holdout=%s", policy, strip, args.limit, seed, holdout_source)

    # 1-2. labelled rows -> model input, clean, keys, min-alnum filter
    df, labelled_meta = load_labelled(cfg, args.limit, rng)
    df["subject"] = df["subject"].map(clean_text)
    df["body"] = df["body"].map(clean_text)
    df["text"] = [model_input(s, b) for s, b in zip(df["subject"].tolist(), df["body"].tolist())]
    df["label"] = df["label"].astype(np.int64)
    df["source"] = df["source"].astype(str)
    bases = add_key_columns(df)
    df["base"] = bases
    LOG.info("labelled: cleaned + keyed %d rows in %.1fs", len(df), common.timer() - t_start)

    short = (df["n_alnum"] < min_alnum).to_numpy()
    drop_min_alnum = nested_counts(df.loc[short])
    LOG.info("drop <%d alnum: %d rows %s", min_alnum, int(short.sum()), drop_min_alnum)
    df = df.loc[~short].reset_index(drop=True)

    # 3-5. holdout, Ling exclusion, exact dedup of the pool
    is_holdout = (df["source"] == holdout_source).to_numpy()
    is_ling = (df["source"] == "Ling").to_numpy()
    holdout = df.loc[is_holdout].reset_index(drop=True)
    ling_excluded = {str(k): int(v) for k, v in df.loc[is_ling].groupby("label").size().items()}
    pool = df.loc[~is_holdout & ~is_ling].reset_index(drop=True)
    if (pool["source"] == holdout_source).any():
        raise AssertionError("holdout source present in the training pool")
    LOG.info("holdout %s: %d rows (%s); Ling excluded: %d rows %s; pool before dedup: %d",
             holdout_source, len(holdout), split_totals(holdout), int(is_ling.sum()), ling_excluded, len(pool))
    holdout_internal_dups = int(holdout["k_full"].duplicated().sum())

    # identical bodies carrying both labels (Enron ham vs TREC-05 spam energy newsletters): keeping one
    # by Hub row order would be arbitrary label noise, so every member of a conflicting group is dropped
    conflict = (pool.groupby("k_full")["label"].transform("nunique") > 1).to_numpy()
    label_conflicts = {"groups": int(pool.loc[conflict, "k_full"].nunique()), "rows_dropped": int(conflict.sum()),
                       "by_source_label": nested_counts(pool.loc[conflict])}
    pool = pool.loc[~conflict].reset_index(drop=True)
    LOG.info("pool label conflicts (same k_full, both labels): dropped %d rows in %d groups %s",
             label_conflicts["rows_dropped"], label_conflicts["groups"], label_conflicts["by_source_label"])

    dup_mask = pool["k_full"].duplicated(keep="first").to_numpy()
    pool_exact_dups = nested_counts(pool.loc[dup_mask])
    pool = pool.loc[~dup_mask].reset_index(drop=True)
    LOG.info("pool exact dedup (k_full keep first): removed %d, kept %d", int(dup_mask.sum()), len(pool))

    # 4. near-dup links over pool + holdout
    n_pool = len(pool)
    index = ShingleIndex(pool["base"].tolist() + holdout["base"].tolist(), nd)
    la, lb = index.links()
    a_hold, b_hold = la >= n_pool, lb >= n_pool
    pool_pool = ~a_hold & ~b_hold
    pool_hold = a_hold ^ b_hold
    hold_hold = a_hold & b_hold
    link_counts = {"pool_pool": int(pool_pool.sum()), "pool_holdout": int(pool_hold.sum()), "holdout_holdout": int(hold_hold.sum())}
    n_capped, n_touching = index.n_capped_shingles, index.n_rows_touching_capped
    LOG.info("near-dup links: %s", link_counts)

    # 5-6. decontaminate the pool against the frozen holdout key/shingle sets
    holdout_keys = frozenset(key_union(holdout))
    key_hit = rows_hitting(pool, holdout_keys)
    shingle_hit = np.zeros(n_pool, dtype=bool)
    shingle_hit[np.where(b_hold[pool_hold], la[pool_hold], lb[pool_hold])] = True   # the pool side of each pool-holdout pair
    contaminated = key_hit | shingle_hit
    decontam = {
        "by_source_label": nested_counts(pool.loc[contaminated]),
        "by_reason": {"key_match_only": int((key_hit & ~shingle_hit).sum()), "shingle_link_only": int((shingle_hit & ~key_hit).sum()), "both": int((key_hit & shingle_hit).sum()), "total": int(contaminated.sum())},
    }
    LOG.info("decontamination vs holdout: removed %d rows %s", int(contaminated.sum()), decontam["by_reason"])

    # 7. clusters: pool-pool shingle links + shared k_pre / k_sb, holdout-holdout links separately
    edges = list(zip(la[pool_pool].tolist(), lb[pool_pool].tolist()))
    edges += list(same_key_pairs(pool["k_pre"].tolist()))
    edges += list(same_key_pairs(pool["k_sb"].tolist()))
    edges += list(zip(la[hold_hold].tolist(), lb[hold_hold].tolist()))
    cluster_all = union_find(n_pool + len(holdout), edges)
    pool["cluster_id"] = cluster_all[:n_pool]
    holdout["cluster_id"] = cluster_all[n_pool:]
    pool = pool.loc[~contaminated].reset_index(drop=True)

    # 9. grouped stratified train/val split
    is_val, split_info = grouped_stratified_split(pool, val_fraction, rng)
    train = pool.loc[~is_val].reset_index(drop=True)
    val = pool.loc[is_val].reset_index(drop=True)
    LOG.info("split: train %s val %s (%s)", split_totals(train), split_totals(val), {k: v for k, v in split_info.items() if k != "per_stratum"})
    del index, la, lb

    # 8. pretraining corpus: keys and near-dups vs holdout removed, keys vs val removed, near-dups vs val counted
    val_keys = frozenset(key_union(val))
    train_keys = frozenset(key_union(train))
    ref_index = ShingleIndex(holdout["base"].tolist() + val["base"].tolist(), nd)
    pretrain_df, pretrain_meta, pretrain_stats = build_pretrain(
        cfg, args.limit, rng, strip, holdout_keys, val_keys, train_keys, ref_index, len(holdout), holdout["label"].to_numpy(),
    )
    del ref_index

    # 11. leakage assertions: an independent recomputation over the FINAL sets, same code path as `check`
    pretrain_keys: set[str] = set()
    pretrain_bases: list[str] = []
    for text in pretrain_df["text"].tolist():
        keys, base = text_identity(text)
        pretrain_keys |= keys
        pretrain_bases.append(base[: nd.window_chars])
    shingle_links = recount_shingle_links(train["base"].tolist(), val["base"].tolist(), holdout["base"].tolist(), pretrain_bases, nd)
    del pretrain_bases
    leak = leakage_checks(train, val, holdout, pretrain_keys, len(pretrain_df), shingle_links, holdout_source)
    leak.raise_if_failed()

    # 10. write outputs
    files = write_outputs(out_dir, train, val, holdout, pretrain_df)
    elapsed = round(common.timer() - t_start, 1)
    manifest = {
        "timestamp": datetime.now(timezone.utc).isoformat(timespec="seconds"),
        "limit": args.limit,
        "dev_subsample": bool(args.limit),
        "spam_policy": policy,
        "strip_quotes": strip,
        "seed": seed,
        "holdout_test_source": holdout_source,
        "elapsed_seconds": elapsed,
        "peak_rss_gb": peak_rss_gb(),
        "datasets": {
            "labelled_primary": labelled_meta,
            "pretrain_corpus": pretrain_meta,
            "excluded": cfg["data"]["labelled_excluded"],
        },
        "counts": {
            "splits": {"finetune_train": nested_counts(train), "finetune_val": nested_counts(val), "test_holdout": nested_counts(holdout)},
            "totals": {"finetune_train": split_totals(train), "finetune_val": split_totals(val), "test_holdout": split_totals(holdout), "pretrain": {"rows": int(len(pretrain_df)), "chars": int(pretrain_df["n_chars"].sum())}},
        },
        "normalization": "common.normalize_text: NFKC -> lowercase -> keep [a-z0-9] only (import it, never re-derive it; NFKC is an extension of the verify_data.md section 7 wording)",
        "dials": dict(DIALS.__dict__),
        "drops": {
            "labelled_min_alnum": {"threshold": min_alnum, "rows": int(short.sum()), "by_source_label": drop_min_alnum},
            "ling_excluded_by_label": ling_excluded,
            "pool_label_conflicts": label_conflicts,
            "pool_exact_dups_by_source_label": pool_exact_dups,
            "decontamination_vs_holdout": decontam,
        },
        "holdout": {"rows": int(len(holdout)), "internal_exact_dups_kept": holdout_internal_dups},
        "near_dup": {**nd.__dict__, "capped_shingles": n_capped, "rows_touching_capped": n_touching, "links_at_removal_stage": link_counts},
        "split": split_info,
        "pretrain": pretrain_stats,
        "leakage": leak.checks,
        "files": files,
        "config_snapshot": {k: v for k, v in cfg.items() if not k.startswith("_")},
    }
    common.write_json(manifest, out_dir / "splits_manifest.json")
    (out_dir / "data_report.md").write_text(render_report(manifest), encoding="utf-8")
    LOG.info("wrote %s and %s in %.1fs (peak RSS %.2f GB)", out_dir / "splits_manifest.json", out_dir / "data_report.md", elapsed, manifest["peak_rss_gb"])
    print_summary(manifest)


def build_pretrain(
    cfg: dict, limit: int | None, rng: np.random.Generator, strip: bool,
    holdout_keys: frozenset[str], val_keys: frozenset[str], train_keys: frozenset[str],
    ref_index: ShingleIndex, n_holdout: int, holdout_labels: np.ndarray,
) -> tuple[pd.DataFrame, dict, dict]:
    """corbt -> deduped, quote-stripped MLM corpus with no exact or near-duplicate of a holdout email.

    Streaming pass: exact dedup, quote strip, length filter, key screens against holdout and val.
    Then one near-dup query of every surviving row against `ref_index` (holdout bases first, then
    val bases): rows linked to a holdout row are removed (verify_data.md section 7 step 8, which
    the key screens alone miss for quoted/forwarded variants and 419 templates), rows linked to a
    val row are counted but kept. Empty-subject rows are additionally screened under their
    text-derived identity (split_text) so the post-hoc text-only assertion shared with `check` can
    never see a key or link this pass did not.
    """
    meta, batches = iter_pretrain_batches(cfg, limit, rng)
    seen_raw: set[str] = set()
    seen_stripped: set[str] = set()
    doc_ids: list[str] = []
    texts: list[str] = []
    n_chars: list[int] = []
    q_bases: list[str] = []      # shingle bases to screen (one per row, two for empty-subject rows)
    q_row: list[int] = []        # kept-row index of each q_base
    window = ref_index.nd.window_chars
    st = Counter()
    marker_hits = Counter()
    chars_before = chars_after = chars_cut_bodies_before = chars_cut_bodies_after = 0
    t0 = common.timer()
    for idx_batch, subjects, bodies in batches:
        for row_idx, subj, body in zip(idx_batch.tolist(), subjects, bodies):
            st["rows_in"] += 1
            subj = clean_text(subj)
            body = clean_text(body)
            raw_key = common.sha1_hex(common.normalize_text(subj) + common.normalize_text(body))
            if raw_key in seen_raw:
                st["exact_dups_raw"] += 1
                continue
            seen_raw.add(raw_key)
            chars_before += len(body)
            if strip:
                new_body, marker, cut = strip_quotes(body)
                if marker:
                    marker_hits[marker] += 1
                if cut:
                    st["bodies_cut"] += 1
                    chars_cut_bodies_before += len(body)
                    chars_cut_bodies_after += len(new_body)
                elif marker:
                    st["marker_but_kept_full"] += 1
                body = new_body
            chars_after += len(body)
            k_full, k_pre, k_sb, k_sb_full, base, _n = row_keys(subj, body)
            if len(common.normalize_text(body)) < DIALS.pretrain_min_alnum:
                st["short_after_strip"] += 1
                continue
            if k_sb_full in seen_stripped:
                st["exact_dups_after_strip"] += 1
                continue
            seen_stripped.add(k_sb_full)
            keys = {k_full, k_pre, k_sb, k_sb_full}
            text = model_input(subj, body)
            alt_base: str | None = None
            if not subj:                      # `check` will see (first paragraph, rest): screen that identity too
                d_keys, d_base = text_identity(text)
                keys |= d_keys
                if d_base != base:
                    alt_base = d_base
            if keys & holdout_keys:
                st["holdout_key_removed"] += 1
                continue
            if keys & val_keys:
                st["val_key_removed"] += 1
                continue
            if keys & train_keys:
                st["train_overlap_kept"] += 1
            q_bases.append(base[:window])
            q_row.append(len(texts))
            if alt_base is not None:
                q_bases.append(alt_base[:window])
                q_row.append(len(texts))
            doc_ids.append(f"corbt:{row_idx:06d}")
            texts.append(text)
            n_chars.append(len(text))
        if st["rows_in"] % 100_000 < 20_000:
            LOG.info("  corbt %d rows in, %d kept (%.1fs)", st["rows_in"], len(texts), common.timer() - t0)

    # near-dup screen against the frozen holdout (removed) and val (counted) shingle sets
    q, r = ref_index.query(q_bases)
    del q_bases
    row_of = np.asarray(q_row, dtype=np.int64)[q] if q.size else np.empty(0, dtype=np.int64)
    hit_hold = r < n_holdout
    drop = np.zeros(len(texts), dtype=bool)
    drop[row_of[hit_hold]] = True
    val_linked_rows = np.unique(row_of[~hit_hold & ~drop[row_of]])
    hold_rows_hit = np.unique(r[hit_hold])
    near_dup = {
        "holdout_links": int(hit_hold.sum()),
        "holdout_rows_removed": int(drop.sum()),
        "holdout_rows_hit_by_label": {str(k): int(v) for k, v in Counter(holdout_labels[hold_rows_hit].tolist()).items()},
        "val_links": int((~hit_hold).sum()),
        "val_rows_hit": int(np.unique(r[~hit_hold] - n_holdout).size),
        "pretrain_rows_near_val_kept": int(val_linked_rows.size),
    }
    LOG.info("pretrain near-dup screen: %s", near_dup)
    keep = ~drop
    df = pd.DataFrame({
        "doc_id": [d for d, k in zip(doc_ids, keep) if k],
        "text": [t for t, k in zip(texts, keep) if k],
        "n_chars": np.asarray(n_chars, dtype=np.int64)[keep],
    })
    stats = {
        **{k: int(st[k]) for k in ("rows_in", "exact_dups_raw", "short_after_strip", "exact_dups_after_strip", "holdout_key_removed", "val_key_removed", "train_overlap_kept")},
        "holdout_near_dup_removed": near_dup["holdout_rows_removed"],
        "near_dup_screen": near_dup,
        "rows_out": int(len(df)),
        "total_chars_out": int(df["n_chars"].sum()) if len(df) else 0,
        "mean_chars_out": round(float(df["n_chars"].mean()), 1) if len(df) else 0.0,
        "min_alnum_after_strip": DIALS.pretrain_min_alnum,
        "quote_strip": {
            "enabled": strip,
            "bodies_cut": int(st["bodies_cut"]),
            "marker_present_but_kept_full": int(st["marker_but_kept_full"]),
            "marker_hits": dict(marker_hits),
            "retained_char_fraction_all_bodies": round(chars_after / chars_before, 4) if chars_before else None,
            "retained_char_fraction_cut_bodies": round(chars_cut_bodies_after / chars_cut_bodies_before, 4) if chars_cut_bodies_before else None,
            "min_alnum_above_cut": DIALS.quote_strip_min_alnum,
        },
    }
    LOG.info("pretrain corpus: %s", {k: v for k, v in stats.items() if k not in ("quote_strip", "near_dup_screen")})
    LOG.info("quote strip: %s", stats["quote_strip"])
    return df, meta, stats


def write_outputs(out_dir: Path, train: pd.DataFrame, val: pd.DataFrame, holdout: pd.DataFrame, pretrain: pd.DataFrame) -> dict[str, str]:
    files = {}
    for name, df in (("finetune_train", train), ("finetune_val", val), ("test_holdout", holdout)):
        out = df.copy()
        out["doc_id"] = [f"primary:{i:06d}" for i in out["row_idx"].tolist()]
        out = out[FINETUNE_SCHEMA]
        out["label"] = out["label"].astype(np.int64)
        out["cluster_id"] = out["cluster_id"].astype(np.int64)
        path = out_dir / f"{name}.parquet"
        out.to_parquet(path, index=False)
        files[name] = str(path)
        LOG.info("wrote %s (%d rows)", path, len(out))
    path = out_dir / "pretrain.parquet"
    pretrain[PRETRAIN_SCHEMA].to_parquet(path, index=False)
    files["pretrain"] = str(path)
    LOG.info("wrote %s (%d rows)", path, len(pretrain))
    return files


# ----------------------------------------------------------------------------- report
def render_report(m: dict) -> str:
    c = m["counts"]
    tot = c["totals"]
    lines = [f"# Data report — {m['timestamp']}", ""]
    if m["dev_subsample"]:
        lines += [f"**DEV SUBSAMPLE (--limit {m['limit']}): counts below are NOT the real build.**", ""]
    lines += [
        f"spam_policy `{m['spam_policy']}` (positive = phishing/spam, unwanted); quote stripping {'on' if m['strip_quotes'] else 'off'}; "
        f"seed {m['seed']}; holdout source `{m['holdout_test_source']}`; elapsed {m['elapsed_seconds']}s; peak RSS {m['peak_rss_gb']} GB.",
        "",
        "## Sources",
        f"- labelled primary: `{m['datasets']['labelled_primary']['repo']}` @ `{m['datasets']['labelled_primary']['commit_sha']}` "
        f"({m['datasets']['labelled_primary']['rows_loaded']:,} rows, hub splits {m['datasets']['labelled_primary']['hub_splits_concatenated']} concatenated; the hub split leaks and is ignored)",
        f"- pretraining corpus: `{m['datasets']['pretrain_corpus']['repo']}` @ `{m['datasets']['pretrain_corpus']['commit_sha']}` ({m['datasets']['pretrain_corpus']['rows_total']:,} rows)",
        f"- excluded (no body text / 92% subset): {', '.join('`' + r + '`' for r in m['datasets']['excluded'])}",
        "",
        "## Splits",
        md_table(["split", "rows", "positives", "negatives", "base rate"],
                 [[k, v["rows"], v["positives"], v["negatives"], v["base_rate"]] for k, v in tot.items() if k != "pretrain"]
                 + [["pretrain (MLM)", tot["pretrain"]["rows"], "-", "-", f"{tot['pretrain']['chars']:,} chars"]]),
        "",
    ]
    for split in ("finetune_train", "finetune_val", "test_holdout"):
        lines += [f"### {split} by source", source_label_table(c["splits"][split]), ""]
    d = m["drops"]
    lc = d["pool_label_conflicts"]
    lines += [
        "## Normalisation and dials",
        f"- {m['normalization']}",
        f"- dials: {m['dials']}",
        "",
        "## Drops and removals",
        f"- labelled rows with <{d['labelled_min_alnum']['threshold']} alphanumerics (subject+body): {d['labelled_min_alnum']['rows']:,} rows",
        source_label_table(d["labelled_min_alnum"]["by_source_label"]),
        f"- Ling excluded (academic list corpus, {sum(d['ling_excluded_by_label'].values()):,} rows by label {d['ling_excluded_by_label']})",
        f"- pool label conflicts (identical normalised body labelled both 0 and 1; all members dropped): {lc['rows_dropped']:,} rows in {lc['groups']:,} groups",
        source_label_table(lc["by_source_label"]),
        f"- pool exact duplicates (k_full, keep first): {sum(sum(v.values()) for v in d['pool_exact_dups_by_source_label'].values()):,} rows",
        source_label_table(d["pool_exact_dups_by_source_label"]),
        f"- decontamination against the holdout: {d['decontamination_vs_holdout']['by_reason']}",
        source_label_table(d["decontamination_vs_holdout"]["by_source_label"]),
        f"- holdout internal exact duplicates kept (test set is never deduped): {m['holdout']['internal_exact_dups_kept']:,}",
        "",
        "## Near-duplicate clustering",
        f"- {m['near_dup']['shingle_chars']}-char shingles, stride {m['near_dup']['stride']}, first {m['near_dup']['window_chars']} chars, link at >= {m['near_dup']['min_shared']} shared; "
        f"{m['near_dup']['capped_shingles']:,} boilerplate shingles over cap {m['near_dup']['cap']:,} ignored (touching {m['near_dup']['rows_touching_capped']:,} rows)",
        f"- links at the removal stage (pool + holdout index): {m['near_dup']['links_at_removal_stage']}",
        f"- clusters {m['split']['clusters_total']:,} (mixed source x label: {m['split']['clusters_mixed_stratum']:,}; largest {m['split']['largest_cluster']:,}); "
        f"val fraction target {m['split']['val_fraction_target']} achieved {m['split']['val_fraction_achieved']}",
        md_table(["stratum", "rows", "clusters", "val target", "val rows", "val frac", "clusters too large"],
                 [[k, v["rows"], v["clusters"], v["val_target"], v["val_rows"], v["val_fraction"], v["clusters_too_large_for_val"]] for k, v in m["split"]["per_stratum"].items()]),
        "",
        "## Pretraining corpus (corbt/enron-emails)",
    ]
    p = m["pretrain"]
    lines += [md_table(["stage", "rows"], [
        ["rows in", p["rows_in"]], ["exact dups (subject+body, raw)", p["exact_dups_raw"]],
        [f"< {p['min_alnum_after_strip']} alnum after stripping", p["short_after_strip"]],
        ["exact dups after stripping", p["exact_dups_after_strip"]],
        ["matches a holdout key (removed)", p["holdout_key_removed"]], ["matches a fine-tune val key (removed)", p["val_key_removed"]],
        [f">= {m['near_dup']['min_shared']} shingles shared with a holdout row (removed)", p["holdout_near_dup_removed"]],
        [f">= {m['near_dup']['min_shared']} shingles shared with a fine-tune val row (kept, counted)", p["near_dup_screen"]["pretrain_rows_near_val_kept"]],
        ["overlaps a fine-tune train key (kept, allowed)", p["train_overlap_kept"]], ["rows out", p["rows_out"]],
    ])]
    ns = p["near_dup_screen"]
    q = p["quote_strip"]
    lines += [
        f"- near-dup screen: {ns['holdout_links']:,} links to {sum(ns['holdout_rows_hit_by_label'].values()):,} holdout rows (by label {ns['holdout_rows_hit_by_label']}); "
        f"{ns['val_links']:,} links to {ns['val_rows_hit']:,} val rows",
        f"- quote/forward stripping {'on' if q['enabled'] else 'off'}: {q['bodies_cut']:,} bodies cut, {q['marker_present_but_kept_full']:,} had a marker but kept the full body "
        f"(< {q['min_alnum_above_cut']} alnum above it); marker hits {q['marker_hits']}; retained chars {q['retained_char_fraction_all_bodies']} of all bodies, "
        f"{q['retained_char_fraction_cut_bodies']} of cut bodies",
        f"- {p['total_chars_out']:,} chars out, mean {p['mean_chars_out']} chars/doc",
        "",
        "## Leakage assertions",
        md_table(["check", "count", "asserted zero", "pass"], [[k, v["count"], v["asserted"], v["pass"]] for k, v in m["leakage"].items()]),
        "",
        "## Files",
    ] + [f"- `{k}`: `{v}`" for k, v in m["files"].items()]
    return "\n".join(lines) + "\n"


def print_summary(m: dict) -> None:
    print("\n=== data_prep summary ===")
    if m["dev_subsample"]:
        print(f"*** DEV SUBSAMPLE --limit {m['limit']} — not a real build ***")
    for k, v in m["counts"]["totals"].items():
        print(f"{k:>16}: {v}")
    print("leakage:")
    for k, v in m["leakage"].items():
        print(f"  {k:<44} {v['count']:>8}  {'ASSERT' if v['asserted'] else 'info  '}  {'ok' if v['pass'] else 'FAIL'}")
    print(f"pretrain: {m['pretrain']['rows_out']:,} rows, quote strip {m['pretrain']['quote_strip']['bodies_cut']:,} cut")


# ----------------------------------------------------------------------------- check
def cmd_check(args: argparse.Namespace, cfg: dict) -> None:
    out_dir = Path(cfg["paths"]["data_processed"])
    manifest = common.read_json(out_dir / "splits_manifest.json")
    nd = NearDupConfig.from_config(cfg, cap=int(manifest["near_dup"]["cap"]))
    frames = {}
    for name in ("finetune_train", "finetune_val", "test_holdout"):
        df = pd.read_parquet(out_dir / f"{name}.parquet")
        if list(df.columns) not in (FINETUNE_SCHEMA, THREEWAY_SCHEMA):   # build3 appends label3/label_set; the binary gate ignores them
            raise AssertionError(f"{name}: unexpected columns {list(df.columns)}")
        df = df[FINETUNE_SCHEMA]
        df["subject"] = df["subject"].fillna("")
        df["body"] = df["body"].fillna("")
        df["base"] = add_key_columns(df)
        frames[name] = df
        LOG.info("read %s: %d rows", name, len(df))
    pre = pd.read_parquet(out_dir / "pretrain.parquet", columns=["text"])
    if list(pre.columns) != ["text"]:
        raise AssertionError("pretrain.parquet: missing text column")
    pretrain_keys: set[str] = set()
    pretrain_bases: list[str] = []
    for t in pre["text"].tolist():
        keys, base = text_identity(t)
        pretrain_keys |= keys
        pretrain_bases.append(base[: nd.window_chars])
    LOG.info("read pretrain: %d rows", len(pre))

    train, val, hold = frames["finetune_train"], frames["finetune_val"], frames["test_holdout"]
    holdout_source = cfg["data"]["holdout_test_source"]
    shingle_links = recount_shingle_links(train["base"].tolist(), val["base"].tolist(), hold["base"].tolist(), pretrain_bases, nd)
    leak = leakage_checks(train, val, hold, pretrain_keys, len(pre), shingle_links, holdout_source)
    print_summary({**manifest, "leakage": leak.checks})
    print("\nmanifest timestamp:", manifest["timestamp"], "| limit:", manifest["limit"], "| spam_policy:", manifest["spam_policy"])
    print("build-time leakage results:", {k: v["count"] for k, v in manifest["leakage"].items()})
    leak.raise_if_failed()
    print("ALL LEAKAGE ASSERTIONS PASS")


# ----------------------------------------------------------------------------- inspect
def cmd_inspect(args: argparse.Namespace, cfg: dict) -> None:
    from datasets import get_dataset_split_names, load_dataset

    cache_dir = str(Path(cfg["paths"]["data_raw"]) / "hf_cache")
    repos = [cfg["data"]["labelled_primary"], cfg["data"]["pretrain_corpus"], *cfg["data"]["labelled_excluded"]]
    for repo in repos:
        print(f"\n=== {repo} ===")
        sha = hub_sha(repo)
        try:
            splits = get_dataset_split_names(repo, revision=sha)
            ds = load_dataset(repo, split=splits[0], streaming=True, cache_dir=cache_dir, revision=sha)
            print(f"commit: {sha}  splits: {splits}  (showing '{splits[0]}')")
            rows = list(ds.take(2))
            if ds.features is not None:
                print("features:")
                for name, feat in ds.features.items():
                    print(f"  {name}: {feat}")
            else:   # streaming CSV/JSON builders declare no schema; show the python types of the first row
                print("features: not declared by the builder; python types of row 0:")
                for name, value in (rows[0].items() if rows else []):
                    print(f"  {name}: {type(value).__name__}")
            for i, row in enumerate(rows):
                print(f"row {i}:")
                for k, v in row.items():
                    s = repr(v)
                    print(f"  {k}: {s[:200]}{'...' if len(s) > 200 else ''}")
        except Exception as exc:
            print(f"ERROR inspecting {repo}: {type(exc).__name__}: {exc}")
    role = {cfg["data"]["labelled_primary"]: "labelled primary (own split built here)", cfg["data"]["pretrain_corpus"]: "MLM pretraining"}
    print("\nroles:", {r: role.get(r, "EXCLUDED from training/eval") for r in repos})


# ----------------------------------------------------------------------------- three-way labels (RUN_LOG entry 6)
# Contract shared with finetune.py / evaluate.py / app: label3 in {0 legitimate, 1 phishing, 2 spam, -1 unwanted-but-
# unsplit}, label_set in {legit, phish, spam, unwanted}; the binary `label` stays (1 for phish/spam/unwanted).
LABEL3_OF_SET: dict[str, int] = {"legit": 0, "phish": 1, "spam": 2, "unwanted": -1}
SET_OF_LABEL3: dict[int, str] = {v: k for k, v in LABEL3_OF_SET.items()}
LABEL3_NAMES: dict[int, str] = {0: "Legitimate", 1: "Phishing", 2: "Spam"}
UNWANTED_NAME = "Unwanted (phishing or spam)"
THREEWAY_SCHEMA = FINETUNE_SCHEMA + ["label3", "label_set"]
THREEWAY_SPLITS = ("finetune_train", "finetune_val", "test_holdout", "test_holdout3")
BINARY_SNAPSHOT_DIR = "binary_snapshot"
BIGSHP_SOURCE = "bigshp"
PART_ORDER = ("holdout", "test3", "train", "val")   # canonical order of the shingle-link pair names (a_b with a before b)
MBOX_INTERNAL_SUBJECT = "DON'T DELETE THIS MESSAGE -- FOLDER INTERNAL DATA"   # mbox folder header mis-exported as an email
_WORD = re.compile(r"[a-z]{2,}")
_PUNCT = re.compile(r"[^\w\s]")

THREEWAY_DEFAULTS: dict[str, Any] = {
    # Provenance -> label3 of a binary-positive row. TREC/CEAS spam tracks never separated spam from phishing
    # ("unwanted", -1, trains on P(phish)+P(spam)); SpamAssassin / Enron-Spam / Ling label-1 rows are spam folders;
    # Nazario and Nigerian_Fraud are phishing corpora by construction.
    "positive_label3_by_source": {
        "Assassin": 2, "Enron": 2, "Ling": 2,
        "TREC-05": -1, "TREC-06": -1, "TREC-07": -1, "CEAS-08": -1,
        "Nazario": 1, "Nigerian_Fraud": 1,
    },
    "champa": {
        "repo": "kudzaiprichard/aura-phishing-email-corpus",
        "files": ["Nazario.csv", "Nigerian_Fraud.csv"],
        "local_dir": "data/raw/aura",
        "label3": 1,
    },
    "bigshp": {
        "enabled": False,
        "verdict": "DO_NOT_USE",
        "repo": "locuoco/the-biggest-spam-ham-phish-email-dataset-300000",
        "file": "df.csv",
        "local_dir": "data/raw/bigshp",
        "drop_phish_class": True,
        "min_chars": 160,
        "min_english_fraction": 0.7,
        "near_dup_token_key_len": 30,
    },
    "new_source_test_fraction": 0.15,
    "new_source_val_fraction": 0.10,
    "drop_new_clusters_linked_to_both_train_and_val": True,
    "holdout3_from_val": {
        "legit_sources": ["CEAS-08", "TREC-07", "Assassin"],
        "spam_sources": ["Assassin", "Enron"],
        "legit_rows": 2000,
        "spam_rows": 700,
    },
}


def threeway_config(cfg: dict) -> dict[str, Any]:
    """config `data.threeway` merged over THREEWAY_DEFAULTS one level deep.

    Every decision (which source is spam vs unwanted, whether bigshp is ingested, the holdout3 targets) is
    therefore visible in config.yaml and reversible, while a missing key never silently changes a decision.
    """
    user = (cfg.get("data") or {}).get("threeway") or {}
    out: dict[str, Any] = {}
    for key, default in THREEWAY_DEFAULTS.items():
        val = user.get(key, default)
        if isinstance(default, dict) and isinstance(val, dict):
            merged = dict(default)
            merged.update(val)
            val = merged
        out[key] = val
    return out


def label3_from_provenance(source: np.ndarray, label: np.ndarray, pos_map: dict[str, int]) -> np.ndarray:
    """label3 of every row from (source, binary label): label 0 -> 0, label 1 -> pos_map[source].

    Raises on a positive row whose source has no entry: silently defaulting it to "unwanted" would hide a
    provenance decision that belongs in config.yaml.
    """
    label = np.asarray(label, dtype=np.int64)
    source = np.asarray(source).astype(str)
    pos = label == 1
    unknown = sorted(set(source[pos].tolist()) - set(pos_map))
    if unknown:
        raise ValueError(f"no data.threeway.positive_label3_by_source entry for positive rows of source(s) {unknown}")
    out = np.zeros(len(label), dtype=np.int64)
    for s in set(source[pos].tolist()):
        out[pos & (source == s)] = int(pos_map[s])
    return out


def add_label3_columns(df: pd.DataFrame, label3: np.ndarray) -> None:
    df["label3"] = np.asarray(label3, dtype=np.int64)
    df["label_set"] = [SET_OF_LABEL3[int(v)] for v in df["label3"].tolist()]


def contract_violations(df: pd.DataFrame, pos_map: dict[str, int]) -> dict[str, int]:
    """Per-rule counts of rows breaking the three-way contract (all zero on a correct file).

    label_set must be a known set, label3 must equal LABEL3_OF_SET[label_set], the binary label must be
    (label3 != 0), and for every source except bigshp (which carries its own label) label3 must equal the
    provenance mapping. check3 asserts the sum is zero.
    """
    if len(df) == 0:
        return {"label_set_unknown": 0, "label3_vs_label_set": 0, "label_vs_label3": 0, "label3_vs_provenance": 0}
    label3 = df["label3"].to_numpy().astype(np.int64)
    lset = df["label_set"].astype(str).to_numpy()
    label = df["label"].to_numpy().astype(np.int64)
    expect3 = np.fromiter((LABEL3_OF_SET.get(s, -99) for s in lset.tolist()), dtype=np.int64, count=len(lset))
    out = {
        "label_set_unknown": int((expect3 == -99).sum()),
        "label3_vs_label_set": int((expect3 != label3).sum()),
        "label_vs_label3": int(((label3 != 0).astype(np.int64) != label).sum()),
    }
    own = df["source"].astype(str).to_numpy() == BIGSHP_SOURCE
    try:
        prov = label3_from_provenance(df["source"].to_numpy()[~own], label[~own], pos_map)
        out["label3_vs_provenance"] = int((prov != label3[~own]).sum())
    except ValueError as exc:
        LOG.error("provenance mapping: %s", exc)
        out["label3_vs_provenance"] = int((~own).sum())
    out["bigshp_label3_outside_0_1_2"] = int((~np.isin(label3[own], [0, 1, 2])).sum())
    return out


def nested_counts3(df: pd.DataFrame) -> dict[str, dict[str, int]]:
    """{source: {label3: n}} with string keys for JSON."""
    out: dict[str, dict[str, int]] = {}
    if len(df) == 0:
        return out
    for (source, l3), n in df.groupby(["source", "label3"], sort=True).size().items():
        out.setdefault(str(source), {})[str(int(l3))] = int(n)
    return out


def split_totals3(df: pd.DataFrame) -> dict[str, Any]:
    n = int(len(df))
    by3 = {str(k): int(v) for k, v in df["label3"].value_counts().sort_index().items()} if n else {}
    byset = {str(k): int(v) for k, v in df["label_set"].value_counts().sort_index().items()} if n else {}
    pos = int((df["label"] == 1).sum()) if n else 0
    return {"rows": n, "by_label3": by3, "by_label_set": byset, "binary_positives": pos,
            "binary_base_rate": round(pos / n, 4) if n else None,
            "rows_with_full_label3": int(df["label3"].isin([0, 1, 2]).sum()) if n else 0}


def source_label3_table(counts: dict[str, dict[str, int]]) -> str:
    rows = []
    for source in sorted(counts):
        c = counts[source]
        rows.append([source, c.get("0", 0), c.get("1", 0), c.get("2", 0), c.get("-1", 0), sum(c.values())])
    return md_table(["source", "0 Legitimate", "1 Phishing", "2 Spam", "-1 unwanted (unsplit)", "rows"], rows)


def first_reason(n: int, reasons: list[tuple[str, np.ndarray]]) -> np.ndarray:
    """Per-row first matching drop reason ("" = kept), so every drop is counted exactly once."""
    out = np.full(n, "", dtype=object)
    for name, mask in reasons:
        hit = np.asarray(mask, dtype=bool) & (out == "")
        out[hit] = name
    return out


def by_source_reason(source: np.ndarray, reason: np.ndarray) -> dict[str, dict[str, int]]:
    out: dict[str, dict[str, int]] = {}
    for (s, r), n in Counter(zip(source.tolist(), reason.tolist())).items():
        if r:
            out.setdefault(str(s), {})[str(r)] = int(n)
    return out


def write_table_atomic(table: "pa.Table", path: Path) -> None:
    """tmp + os.replace: other tools (finetune, evaluate, tokenizer stats) read these files while we write."""
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_name(path.name + ".tmp")
    pq.write_table(table, tmp)
    os.replace(tmp, path)


def write_text_atomic(text: str, path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_name(path.name + ".tmp")
    tmp.write_text(text, encoding="utf-8")
    os.replace(tmp, path)


def copy_atomic(src: Path, dst: Path) -> None:
    dst.parent.mkdir(parents=True, exist_ok=True)
    tmp = dst.with_name(dst.name + ".tmp")
    shutil.copy2(src, tmp)
    os.replace(tmp, dst)


# ----------------------------------------------------------------------------- three-way: inputs
def binary_base_dir(processed_dir: Path, dev: bool) -> tuple[Path, str]:
    """Directory holding the BINARY (pre-build3) splits that build3 derives from.

    build3 rewrites finetune_{train,val}.parquet and test_holdout.parquet in place, so a second run could not
    start from the same inputs. The first non-dev run therefore snapshots the binary files into
    <data_processed>/binary_snapshot/ (atomic copies) and every later run derives from that snapshot: build3 is
    idempotent and the binary split membership stays byte-for-byte recoverable. If the live files already carry
    label3 and there is no snapshot, the binary originals are gone and `build` must be re-run first.
    """
    names = ("finetune_train", "finetune_val", "test_holdout")
    snap = processed_dir / BINARY_SNAPSHOT_DIR
    if all((snap / f"{n}.parquet").exists() for n in names):
        return snap, "snapshot"
    for n in names:
        live = processed_dir / f"{n}.parquet"
        if not live.exists():
            raise FileNotFoundError(f"{live} missing: run `data_prep.py build` first")
        if "label3" in pq.read_schema(live).names:
            raise RuntimeError(f"{live} already carries label3 but {snap} has no snapshot of the binary files; "
                               "re-run `data_prep.py build` to recreate the binary splits before build3")
    if dev:
        return processed_dir, "live (dev run: no snapshot taken)"
    for n in names:
        copy_atomic(processed_dir / f"{n}.parquet", snap / f"{n}.parquet")
    LOG.info("snapshotted the binary splits to %s", snap)
    return snap, "snapshot (created by this run)"


def per_source_sample(source: np.ndarray, limit: int, rng: np.random.Generator) -> np.ndarray:
    """Sorted row indices: at most `limit` rows per source (dev subsample)."""
    keep = []
    for s in np.unique(source):
        idx = np.flatnonzero(source == s)
        keep.append(rng.choice(idx, min(limit, idx.size), replace=False))
    return np.sort(np.concatenate(keep)) if keep else np.empty(0, dtype=np.int64)


def read_binary_split(path: Path, limit: int | None, rng: np.random.Generator) -> tuple["pa.Table", pd.DataFrame]:
    """(arrow table, keyed frame) of a binary split; the two stay row-aligned so the table can be rewritten in place."""
    table = pq.read_table(path)
    if list(table.schema.names) != FINETUNE_SCHEMA:
        raise AssertionError(f"{path}: expected binary columns {FINETUNE_SCHEMA}, got {list(table.schema.names)}")
    if limit:
        idx = per_source_sample(np.asarray(table.column("source").to_pylist(), dtype=object), limit, rng)
        table = table.take(pa.array(idx))
    df = table.select(["doc_id", "subject", "body", "label", "source", "cluster_id"]).to_pandas()
    df["subject"] = df["subject"].fillna("")
    df["body"] = df["body"].fillna("")
    df["label"] = df["label"].astype(np.int64)
    df["source"] = df["source"].astype(str)
    df["base"] = add_key_columns(df)
    LOG.info("read %s: %d rows%s", path.name, len(df), f" (dev subsample, <= {limit}/source)" if limit else "")
    return table, df


def pretrain_identity(path: Path, nd: NearDupConfig, limit: int | None, rng: np.random.Generator
                      ) -> tuple[set[str], list[str], int, "pa.Table | None"]:
    """(keys, window-trimmed shingle bases, n_rows, dev table) of the pretraining corpus.

    Re-derived from the stored `text` with text_identity exactly as `check` does, in 20k-row batches so the
    408M-char corpus is never fully resident. In dev mode a 4*limit subsample is used AND returned so the dev
    directory gets its own pretrain.parquet: check3 must screen against the same corpus build3 screened against.
    """
    keys: set[str] = set()
    bases: list[str] = []
    dev_table = None
    if limit:
        full = pq.read_table(path)
        n = min(limit * 4, full.num_rows)
        idx = np.sort(rng.choice(full.num_rows, n, replace=False))
        dev_table = full.take(pa.array(idx))
        del full
        batches = dev_table.select(["text"]).to_batches(max_chunksize=20_000)
        LOG.warning("--limit %d: pretrain subsample %d rows", limit, n)
    else:
        batches = pq.ParquetFile(path).iter_batches(batch_size=20_000, columns=["text"])
    n_rows = 0
    t0 = common.timer()
    for batch in batches:
        for text in batch.column(0).to_pylist():
            k, base = text_identity(text or "")
            keys |= k
            bases.append(base[: nd.window_chars])
            n_rows += 1
    LOG.info("pretrain identity: %d rows, %d keys (%.1fs)", n_rows, len(keys), common.timer() - t0)
    return keys, bases, n_rows, dev_table


def champa_commit_sha(local_dir: Path, fname: str, repo: str) -> str | None:
    """Commit sha recorded by hf_hub_download in <local_dir>/.cache/huggingface/download/<file>.metadata (offline), else the Hub."""
    meta = local_dir / ".cache" / "huggingface" / "download" / f"{fname}.metadata"
    if meta.exists():
        first = meta.read_text(encoding="utf-8").splitlines()[0].strip()
        if re.fullmatch(r"[0-9a-f]{40}", first):
            return first
    return hub_sha(repo)


def load_champa(tw: dict, limit: int | None, rng: np.random.Generator) -> tuple[pd.DataFrame, dict]:
    """Nazario.csv + Nigerian_Fraud.csv from the Champa et al. mirror -> raw frame (subject, body, label, source, label3).

    Both files are phishing corpora by construction (all label 1 in the file, label3 = 1 by provenance). The
    sender/receiver/date/urls columns are never used: model input is subject + body only, like the primary set.
    Downloaded with hf_hub_download when missing; the commit sha is recorded per file.
    """
    c = tw["champa"]
    local_dir = Path(c["local_dir"])
    if not local_dir.is_absolute():
        local_dir = common.project_root() / local_dir
    meta: dict[str, Any] = {"repo": c["repo"], "local_dir": str(local_dir), "label3": int(c["label3"]), "files": {}}
    frames = []
    for fname in c["files"]:
        path = local_dir / fname
        if not path.exists():
            from huggingface_hub import hf_hub_download

            LOG.info("downloading %s/%s -> %s", c["repo"], fname, local_dir)
            hf_hub_download(c["repo"], fname, repo_type="dataset", local_dir=str(local_dir))
        sha = champa_commit_sha(local_dir, fname, c["repo"])
        raw = pd.read_csv(path, dtype=str, keep_default_na=False)
        missing = {"subject", "body", "label"} - set(raw.columns)
        if missing:
            raise ValueError(f"{path}: missing columns {sorted(missing)}")
        labels = raw["label"].astype(np.int64).to_numpy()
        if (labels != 1).any():
            raise ValueError(f"{path}: expected every label == 1 (phishing corpus), found {Counter(labels.tolist())}")
        stem = Path(fname).stem
        df = pd.DataFrame({
            "doc_id": [f"aura:{stem}:{i:06d}" for i in range(len(raw))],
            "subject": raw["subject"].astype(str), "body": raw["body"].astype(str),
            "label": np.ones(len(raw), dtype=np.int64), "source": stem,
            "label3": np.full(len(raw), int(c["label3"]), dtype=np.int64),
        })
        if limit:
            df = df.iloc[np.sort(rng.choice(len(df), min(limit, len(df)), replace=False))].reset_index(drop=True)
        meta["files"][fname] = {"rows": int(len(raw)), "rows_used": int(len(df)), "commit_sha": sha, "path": str(path),
                                "bytes": int(path.stat().st_size)}
        LOG.info("champa %s @ %s: %d rows (%d used), %d null subjects", fname, sha, len(raw), len(df), int((raw["subject"] == "").sum()))
        frames.append(df)
    return pd.concat(frames, ignore_index=True), meta


def english_vocab(train: pd.DataFrame, rng: np.random.Generator, sample: int = 20_000, min_count: int = 3) -> set[str]:
    """Lower-case word types seen >= min_count times in a sample of legitimate training bodies (bigshp filter)."""
    legit = np.flatnonzero((train["label"] == 0).to_numpy())
    idx = rng.choice(legit, min(sample, legit.size), replace=False)
    counts: Counter = Counter()
    for body in train["body"].to_numpy()[idx].tolist():
        counts.update(_WORD.findall(body.lower()))
    return {w for w, n in counts.items() if n >= min_count}


def load_bigshp(tw: dict, limit: int | None, rng: np.random.Generator, vocab: set[str]) -> tuple[pd.DataFrame, dict]:
    """locuoco/...-300000 df.csv behind data.threeway.bigshp.enabled, with the research verdict's filter chain.

    Research verdict (docs/research/bigshp_research.md): DO_NOT_USE. The file is a union of lossy re-packagings of
    corpora we already hold; its phish label is source-derived (87% of resolvable phish rows are spam-folder or
    spam-track mail) and the whole phish class went through a stopword-stripping pipeline that would act as a label
    shortcut. This loader exists so the decision is reversible from config.yaml. Filters, in order:
    drop the label-1 class; drop rows shorter than min_chars; label-2 rows with any non-ASCII byte or < min_english_fraction
    of word tokens in our legitimate-mail vocabulary; label-0 rows that look like Ling-Spam list mail (first token
    "subject", no punctuation); within-file dedup on k_full, k_pre and the first-N-token key. The cross-set screens
    (keys and shingles against train/val/holdout/pretrain) then apply exactly as for the Champa rows; a surviving
    row whose key matches one of our rows is dropped there, with label contradictions counted separately.
    NOTE: no digit-folding adapter ("escapenumber") is applied to our keys, so exact-key recall on bigshp text is a
    lower bound and the shingle screen carries the near-duplicate removal.
    """
    import pyarrow.csv as pacsv

    b = tw["bigshp"]
    local_dir = Path(b["local_dir"])
    if not local_dir.is_absolute():
        local_dir = common.project_root() / local_dir
    path = local_dir / b["file"]
    if not path.exists():
        from huggingface_hub import hf_hub_download

        LOG.info("downloading %s/%s -> %s", b["repo"], b["file"], local_dir)
        hf_hub_download(b["repo"], b["file"], repo_type="dataset", local_dir=str(local_dir))
    sha = champa_commit_sha(local_dir, b["file"], b["repo"])
    table = pacsv.read_csv(
        path,
        parse_options=pacsv.ParseOptions(newlines_in_values=True),
        read_options=pacsv.ReadOptions(block_size=64 << 20),
        convert_options=pacsv.ConvertOptions(column_types={"label": pa.int64(), "text": pa.large_string()}),
    )
    n_total = table.num_rows
    if limit:
        idx = np.sort(rng.choice(n_total, min(limit * 4, n_total), replace=False))
        table = table.take(pa.array(idx))
        LOG.warning("--limit %d: bigshp subsample %d rows", limit, table.num_rows)
    row_idx = idx if limit else np.arange(n_total)
    label3 = np.asarray(table.column("label").to_pylist(), dtype=np.int64)
    texts = [t or "" for t in table.column("text").to_pylist()]
    del table
    n = len(texts)
    lengths = np.fromiter((len(t) for t in texts), dtype=np.int64, count=n)
    non_ascii = np.fromiter((not t.isascii() for t in texts), dtype=bool, count=n)
    eng = np.ones(n, dtype=np.float64)
    ling_like = np.zeros(n, dtype=bool)
    tok_keys: list[str] = []
    k = int(b["near_dup_token_key_len"])
    for i, t in enumerate(texts):
        words = _WORD.findall(t.lower())
        if words:
            eng[i] = sum(1 for w in words if w in vocab) / len(words)
        tok_keys.append(common.sha1_hex(" ".join(words[:k])))
        if label3[i] == 0:
            first = t.strip().split(None, 1)[0].lower() if t.strip() else ""
            ling_like[i] = first == "subject" and _PUNCT.search(t) is None
    keyed = pd.DataFrame({"subject": [""] * n, "body": texts})
    add_key_columns(keyed)
    reasons = [
        ("bigshp_phish_class_source_derived", (label3 == 1) if b["drop_phish_class"] else np.zeros(n, dtype=bool)),
        ("bigshp_label_outside_0_1_2", ~np.isin(label3, [0, 1, 2])),
        (f"bigshp_shorter_than_{int(b['min_chars'])}_chars", lengths < int(b["min_chars"])),
        ("bigshp_spam_non_ascii", (label3 == 2) & non_ascii),
        (f"bigshp_spam_english_fraction_lt_{b['min_english_fraction']}", (label3 == 2) & (eng < float(b["min_english_fraction"]))),
        ("bigshp_ham_ling_list_mail", (label3 == 0) & ling_like),
        ("bigshp_within_file_k_full_dup", keyed["k_full"].duplicated(keep="first").to_numpy()),
        ("bigshp_within_file_k_pre_dup", keyed["k_pre"].duplicated(keep="first").to_numpy()),
        (f"bigshp_within_file_first{k}_token_dup", pd.Series(tok_keys).duplicated(keep="first").to_numpy()),
    ]
    reason = first_reason(n, reasons)
    keep = reason == ""
    df = pd.DataFrame({
        "doc_id": [f"bigshp:{int(i):06d}" for i in row_idx[keep].tolist()],
        "subject": "", "body": [t for t, kp in zip(texts, keep) if kp],
        "label": (label3[keep] != 0).astype(np.int64), "source": BIGSHP_SOURCE, "label3": label3[keep],
    })
    meta = {
        "repo": b["repo"], "file": b["file"], "commit_sha": sha, "path": str(path), "bytes": int(path.stat().st_size),
        "rows_total": int(n_total), "rows_read": int(n), "rows_after_file_filters": int(len(df)),
        "file_filter_drops": {name: int(v) for name, v in Counter(reason[~keep].tolist()).items()},
        "filters": {key: b[key] for key in ("drop_phish_class", "min_chars", "min_english_fraction", "near_dup_token_key_len")},
        "label_counts_read": {str(int(c)): int(v) for c, v in Counter(label3.tolist()).items()},
    }
    LOG.info("bigshp @ %s: %d rows read, %d survive the file filters %s", sha, n, len(df), meta["file_filter_drops"])
    return df, meta


def prepare_new_rows(df: pd.DataFrame) -> None:
    """Clean exactly like the primary set: clean_text, model_input, keys/base, label3 -> label_set."""
    df["subject"] = df["subject"].map(clean_text)
    df["body"] = df["body"].map(clean_text)
    df["text"] = [model_input(s, b) for s, b in zip(df["subject"].tolist(), df["body"].tolist())]
    df["label"] = df["label"].astype(np.int64)
    df["source"] = df["source"].astype(str)
    df["base"] = add_key_columns(df)
    add_label3_columns(df, df["label3"].to_numpy())


# ----------------------------------------------------------------------------- three-way: leakage
def recount_links_multi(parts: dict[str, list[str]], pretrain_bases: list[str], pretrain_ref: tuple[str, ...], nd: NearDupConfig) -> dict[str, int]:
    """Independent near-dup link counts between every pair of final sets, from a fresh index.

    Generalises recount_shingle_links to four parts (holdout, test3, train, val): one index over all parts
    enumerates every cross-part link; the pretraining corpus is queried against an index of the small reference
    parts so pretrain-pretrain links are never materialised. Pair names follow PART_ORDER (holdout_train,
    test3_val, ...), so the binary gate's keys (holdout_train, holdout_val, train_val, holdout_pretrain,
    val_pretrain) come out unchanged.
    """
    names = [p for p in PART_ORDER if p in parts]
    bounds = np.cumsum([0] + [len(parts[n]) for n in names])
    index = ShingleIndex([b for n in names for b in parts[n]], nd)
    la, lb = index.links()
    pa_ = np.searchsorted(bounds, la, side="right") - 1
    pb_ = np.searchsorted(bounds, lb, side="right") - 1
    pair = Counter(zip(np.minimum(pa_, pb_).tolist(), np.maximum(pa_, pb_).tolist()))
    out = {f"{names[i]}_{names[j]}": int(pair[(i, j)]) for i in range(len(names)) for j in range(i + 1, len(names))}
    del index, la, lb
    ref_names = [n for n in pretrain_ref if n in parts]
    ref = ShingleIndex([b for n in ref_names for b in parts[n]], nd)
    _q, r = ref.query(pretrain_bases)
    rb = np.cumsum([0] + [len(parts[n]) for n in ref_names])
    part_of = np.searchsorted(rb, r, side="right") - 1
    for i, n in enumerate(ref_names):
        out[f"{n}_pretrain"] = int((part_of == i).sum())
    LOG.info("recount shingle links (final sets): %s", out)
    return out


def leakage_checks3(train: pd.DataFrame, val: pd.DataFrame, holdout: pd.DataFrame, test3: pd.DataFrame,
                    pretrain_keys: set[str], n_pretrain: int, links: dict[str, int], holdout_source: str,
                    pos_map: dict[str, int]) -> LeakageResults:
    """The binary gate (unchanged) plus the test_holdout3 assertions.

    test_holdout3 must be text the model never saw in any form: zero exact keys and zero >=3-shingle links against
    train, val and the MLM corpus, zero keys shared with the TREC-06 holdout (its shingle links to TREC-06 are
    counted only: the two test sets may legitimately contain the same 1998-2008 spam templates), every row fully
    labelled (label3 in {0,1,2}), no TREC-06 row, no doc_id in two splits, no contract violation anywhere.
    """
    res = leakage_checks(train, val, holdout, pretrain_keys, n_pretrain, links, holdout_source)
    t3_keys = key_union(test3)
    hold_keys = key_union(holdout)
    res.add("holdout3_vs_train_key_rows", int(rows_hitting(train, t3_keys).sum()), True)
    res.add("holdout3_vs_val_key_rows", int(rows_hitting(val, t3_keys).sum()), True)
    res.add("holdout3_vs_pretrain_key_rows", len(pretrain_keys & t3_keys), True, "pretrain keys re-derived from stored text (text_identity)")
    res.add("holdout3_vs_holdout_key_rows", int(rows_hitting(test3, hold_keys).sum()), True)
    res.add("holdout3_vs_train_shingle_links", links["test3_train"], True)
    res.add("holdout3_vs_val_shingle_links", links["test3_val"], True)
    res.add("holdout3_vs_pretrain_shingle_links", links["test3_pretrain"], True, "pretrain bases re-derived from stored text (text_identity)")
    res.add("holdout3_vs_holdout_shingle_links", links["holdout_test3"], False, "two test sets; shared 1998-2008 spam templates are allowed")
    res.add("holdout3_rows_without_full_label3", int((~test3["label3"].isin([0, 1, 2])).sum()), True)
    res.add("holdout3_rows_from_holdout_source", int((test3["source"] == holdout_source).sum()), True)
    ids = [set(df["doc_id"].tolist()) for df in (train, val, holdout, test3)]
    res.add("doc_id_shared_between_splits", sum(len(ids[i] & ids[j]) for i in range(4) for j in range(i + 1, 4)), True)
    viol = {name: contract_violations(df, pos_map) for name, df in zip(THREEWAY_SPLITS, (train, val, holdout, test3))}
    res.add("label3_contract_violations", sum(sum(v.values()) for v in viol.values()), True, str({k: v for k, v in viol.items() if sum(v.values())}))
    res.add("holdout3_rows", len(test3), False)
    return res


# ----------------------------------------------------------------------------- three-way: build
def cmd_build3(args: argparse.Namespace, cfg: dict) -> None:
    """Derive the three-way files from the binary splits + the Champa phishing corpora (+ bigshp if enabled).

    Binary split membership is preserved except for the rows moved from finetune_val into test_holdout3; every
    other existing row only gains label3/label_set. New rows are cleaned, keyed, deduplicated within themselves
    and against train/val/holdout/pretrain (exact keys), screened for >=3-shingle links against the TREC-06
    holdout (dropped), pinned to train/val when they near-duplicate an existing train/pretrain/val row (so no
    template straddles a split), and the free clusters are split test3 / val / train by grouped_stratified_split.
    """
    t_start = common.timer()
    seed = int(cfg["seed"])
    common.set_seed(seed)
    rng = np.random.default_rng(seed + 3)       # independent stream from `build` so the binary split is never re-drawn
    tw = threeway_config(cfg)
    if args.bigshp:
        tw["bigshp"]["enabled"] = args.bigshp == "on"
    pos_map = {str(k): int(v) for k, v in tw["positive_label3_by_source"].items()}
    nd = NearDupConfig.from_config(cfg, cap=args.shingle_cap or DIALS.shingle_cap)
    holdout_source = cfg["data"]["holdout_test_source"]
    min_alnum = int(cfg["data"]["min_alnum_chars"])
    processed_dir = Path(cfg["paths"]["data_processed"])
    dev = bool(args.limit)
    out_dir = Path(args.out_dir) if args.out_dir else (processed_dir / f"dev_threeway_limit{args.limit}" if dev else processed_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    LOG.info("build3: limit=%s seed=%d out_dir=%s bigshp.enabled=%s holdout=%s", args.limit, seed, out_dir, tw["bigshp"]["enabled"], holdout_source)

    # 1. binary inputs (snapshot), keyed; label3 from provenance
    base_dir, base_origin = binary_base_dir(processed_dir, dev)
    tables: dict[str, pa.Table] = {}
    frames: dict[str, pd.DataFrame] = {}
    for name in ("finetune_train", "finetune_val", "test_holdout"):
        tables[name], frames[name] = read_binary_split(base_dir / f"{name}.parquet", args.limit, rng)
    train, val, holdout = frames["finetune_train"], frames["finetune_val"], frames["test_holdout"]
    for df in (train, val, holdout):
        add_label3_columns(df, label3_from_provenance(df["source"].to_numpy(), df["label"].to_numpy(), pos_map))
    base_counts = {n: nested_counts3(frames[n]) for n in frames}
    pretrain_keys, pretrain_bases, n_pretrain, pretrain_dev_table = pretrain_identity(processed_dir / "pretrain.parquet", nd, args.limit, rng)
    LOG.info("binary base from %s: train %d val %d holdout %d pretrain %d (%.1fs)", base_origin, len(train), len(val), len(holdout), n_pretrain, common.timer() - t_start)

    # 2. new sources, cleaned exactly like the primary set
    new, champa_meta = load_champa(tw, args.limit, rng)
    bigshp_meta: dict[str, Any] = {"enabled": bool(tw["bigshp"]["enabled"]), "verdict": tw["bigshp"]["verdict"],
                                   "research": "docs/research/bigshp_research.md"}
    if tw["bigshp"]["enabled"]:
        bdf, bmeta = load_bigshp(tw, args.limit, rng, english_vocab(train, rng))
        bigshp_meta.update(bmeta)
        new = pd.concat([new, bdf], ignore_index=True)
    prepare_new_rows(new)
    loaded_by_source = {str(k): int(v) for k, v in new["source"].value_counts().items()}
    LOG.info("new rows: %s", loaded_by_source)

    # 3. exact screens: min-alnum, mbox artefacts, within-new dups, keys against holdout / train / val / pretrain
    hold_keys = frozenset(key_union(holdout))
    train_keys = frozenset(key_union(train))
    val_keys = frozenset(key_union(val))
    hit_hold, hit_train, hit_val, hit_pre = (rows_hitting(new, ks) for ks in (hold_keys, train_keys, val_keys, pretrain_keys))
    reasons = [
        (f"lt_{min_alnum}_alnum", (new["n_alnum"] < min_alnum).to_numpy()),
        ("mbox_internal_data", (new["subject"] == MBOX_INTERNAL_SUBJECT).to_numpy()),
        ("within_new_k_full_dup", new["k_full"].duplicated(keep="first").to_numpy()),
        ("key_match_test_holdout", hit_hold),
        ("key_match_finetune_train", hit_train),
        ("key_match_finetune_val", hit_val),
        ("key_match_pretrain", hit_pre),
    ]
    reason = first_reason(len(new), reasons)
    contradictions = 0
    if tw["bigshp"]["enabled"]:
        # a bigshp row that exact-matches one of our rows is dropped anyway; count how often its label would have
        # disagreed with our provenance label (the research verdict's "label contradicts provenance" filter)
        own_label: dict[str, int] = {}
        for df in (train, val, holdout):
            for col in ("k_full", "k_pre", "k_sb", "k_sb_full"):
                own_label.update(zip(df[col].tolist(), df["label3"].tolist()))
        for i in np.flatnonzero((new["source"] == BIGSHP_SOURCE).to_numpy() & (hit_train | hit_val | hit_hold)):
            match = [own_label[k] for k in (new.at[i, "k_full"], new.at[i, "k_pre"], new.at[i, "k_sb"], new.at[i, "k_sb_full"]) if k in own_label]
            if match and any(m != int(new.at[i, "label3"]) and m != -1 for m in match):
                contradictions += 1
    survivors = new.loc[reason == ""].reset_index(drop=True)
    LOG.info("exact screens dropped %d of %d new rows: %s", int((reason != "").sum()), len(new), by_source_reason(new["source"].to_numpy(), reason))

    # 4. one shingle index over (new survivors + val): holdout -> drop, train/pretrain -> pin train, val -> pin val
    n_new = len(survivors)
    ref = ShingleIndex(survivors["base"].tolist() + val["base"].tolist(), nd)
    _q, r_hold = ref.query(holdout["base"].tolist())
    new_hold_linked = np.zeros(n_new, dtype=bool)
    new_hold_linked[r_hold[r_hold < n_new]] = True
    val_hold_linked = int(np.unique(r_hold[r_hold >= n_new]).size)
    _q, r_tr = ref.query(train["base"].tolist())
    _q, r_pre = ref.query(pretrain_bases)
    new_train_linked = np.zeros(n_new, dtype=bool)
    new_train_linked[r_tr[r_tr < n_new]] = True
    new_pre_linked = np.zeros(n_new, dtype=bool)
    new_pre_linked[r_pre[r_pre < n_new]] = True
    val_train_linked = np.zeros(len(val), dtype=bool)
    val_train_linked[r_tr[r_tr >= n_new] - n_new] = True
    val_pre_linked = np.zeros(len(val), dtype=bool)
    val_pre_linked[r_pre[r_pre >= n_new] - n_new] = True
    la, lb = ref.links()
    a_new, b_new = la < n_new, lb < n_new
    new_new = a_new & b_new
    new_val = a_new ^ b_new
    new_side = np.where(a_new[new_val], la[new_val], lb[new_val])
    val_side = np.where(a_new[new_val], lb[new_val], la[new_val]) - n_new
    new_val_linked = np.zeros(n_new, dtype=bool)
    new_val_linked[new_side] = True
    val_new_linked = np.zeros(len(val), dtype=bool)
    val_new_linked[val_side] = True
    del ref

    # 5. drop holdout-linked new rows; cluster the rest (shingle links + shared k_pre / k_sb); pin clusters
    reason_s = first_reason(n_new, [("shingle_link_test_holdout", new_hold_linked)])
    keep_s = reason_s == ""
    edges = [(int(a), int(b)) for a, b in zip(la[new_new].tolist(), lb[new_new].tolist()) if keep_s[a] and keep_s[b]]
    edges += [(a, b) for a, b in same_key_pairs_multi([survivors["k_pre"].tolist(), survivors["k_sb"].tolist()]) if keep_s[a] and keep_s[b]]
    cl = union_find(n_new, edges)
    # A cluster that near-duplicates rows on BOTH sides of the binary train/val split (or the MLM corpus and val)
    # is a template family the primary set already holds on both sides; placing it anywhere would create
    # train/val near-dup links the binary build had ruled out by construction, so it is dropped instead.
    train_side_cl = set(cl[(new_train_linked | new_pre_linked) & keep_s].tolist())
    val_side_cl = set(cl[new_val_linked & keep_s].tolist())
    both_cl = train_side_cl & val_side_cl if tw["drop_new_clusters_linked_to_both_train_and_val"] else set()
    both_rows = keep_s & np.isin(cl, list(both_cl))
    reason_s[both_rows] = "near_dup_of_both_train_and_val"
    keep_s = reason_s == ""
    pin_train_cl = train_side_cl - both_cl
    pin_val_cl = val_side_cl - both_cl
    hold_touched_cl = set(cl[new_hold_linked].tolist())
    dest = np.full(n_new, "", dtype=object)
    dest[keep_s & np.isin(cl, list(pin_train_cl))] = "train"
    dest[keep_s & np.isin(cl, list(pin_val_cl))] = "val"
    free = keep_s & (dest == "")
    # free clusters: ~15% to test3 (cluster-grouped, stratified by source x label3), ~10% of the rest to val
    tmp = pd.DataFrame({"source": survivors["source"], "label": survivors["label3"], "cluster_id": cl})
    is_test, split_test_info = grouped_stratified_split(tmp.loc[free].reset_index(drop=True), float(tw["new_source_test_fraction"]), rng)
    free_idx = np.flatnonzero(free)
    dest[free_idx[is_test]] = "test3"
    rest_idx = free_idx[~is_test]
    is_val_new, split_val_info = grouped_stratified_split(tmp.iloc[rest_idx].reset_index(drop=True), float(tw["new_source_val_fraction"]), rng)
    dest[rest_idx[is_val_new]] = "val"
    dest[rest_idx[~is_val_new]] = "train"
    assert ((dest != "") == keep_s).all()
    cluster_offset = int(max(train["cluster_id"].max(), val["cluster_id"].max(), holdout["cluster_id"].max())) + 1
    survivors["cluster_id"] = (cl + cluster_offset).astype(np.int64)
    survivors["dest"] = dest
    new_pins = {
        "clusters_total": int(np.unique(cl[keep_s]).size), "clusters_pinned_train_by_train_or_pretrain_link": len(pin_train_cl),
        "clusters_pinned_val_by_val_link": len(pin_val_cl), "clusters_dropped_linked_to_both_train_and_val": len(both_cl),
        "rows_dropped_linked_to_both_train_and_val": int(both_rows.sum()), "clusters_with_a_holdout_linked_member": len(hold_touched_cl),
        "rows_linked_to_train": int((new_train_linked & keep_s).sum()), "rows_linked_to_pretrain": int((new_pre_linked & keep_s).sum()),
        "rows_linked_to_val": int((new_val_linked & keep_s).sum()), "rows_dropped_holdout_shingle_link": int(new_hold_linked.sum()),
        "largest_cluster": int(np.bincount(cl[keep_s]).max()) if keep_s.any() else 0,
    }
    LOG.info("new rows: pins %s; destinations %s", new_pins, Counter(dest[keep_s].tolist()))
    all_reason = reason.copy()
    all_reason[np.flatnonzero(reason == "")] = reason_s
    drops_by_source_reason = by_source_reason(new["source"].to_numpy(), all_reason)
    new_kept = survivors.loc[keep_s].reset_index(drop=True)
    dest_by_source = {s: {d: int(n) for d, n in grp["dest"].value_counts().items()} for s, grp in new_kept.groupby("source")}

    # 6. test_holdout3 part (b): clean-provenance legit / spam rows of the CURRENT val, whole clusters only
    h3 = tw["holdout3_from_val"]
    l3 = val["label3"].to_numpy()
    src = val["source"].to_numpy()
    cand = ((l3 == 0) & np.isin(src, list(h3["legit_sources"]))) | ((l3 == 2) & np.isin(src, list(h3["spam_sources"])))
    val_pre_key = rows_hitting(val, pretrain_keys)
    val_hold_key = rows_hitting(val, hold_keys)
    train_clusters = set(train["cluster_id"].tolist())
    in_train_cl = val["cluster_id"].isin(train_clusters).to_numpy()
    excl = first_reason(len(val), [
        ("not_a_candidate_source_or_label", ~cand),
        ("key_match_pretrain", val_pre_key), ("shingle_link_pretrain", val_pre_linked),
        ("key_or_shingle_link_test_holdout", val_hold_key | (np.isin(np.arange(len(val)), r_hold[r_hold >= n_new] - n_new))),
        ("shingle_link_train", val_train_linked), ("cluster_also_in_train", in_train_cl),
        ("shingle_link_to_a_new_row", val_new_linked),
    ])
    row_ok = excl == ""
    cluster_ok = pd.Series(row_ok).groupby(val["cluster_id"].to_numpy()).transform("all").to_numpy()
    excl[row_ok & ~cluster_ok] = "cluster_has_an_ineligible_member"
    eligible = excl == ""
    moved = np.zeros(len(val), dtype=bool)
    h3_split_info = {}
    for cls, target_rows in ((0, int(h3["legit_rows"])), (2, int(h3["spam_rows"]))):
        pool_mask = eligible & (l3 == cls)
        pool_idx = np.flatnonzero(pool_mask)
        if pool_idx.size == 0:
            LOG.warning("holdout3 from val: no eligible rows for label3=%d", cls)
            continue
        frac = min(1.0, target_rows / pool_idx.size)
        tmp = pd.DataFrame({"source": src[pool_idx], "label": l3[pool_idx], "cluster_id": val["cluster_id"].to_numpy()[pool_idx]})
        sel, info = grouped_stratified_split(tmp, frac, rng)
        moved[pool_idx[sel]] = True
        h3_split_info[LABEL3_NAMES[cls]] = {"target_rows": target_rows, "eligible_rows": int(pool_idx.size), "fraction": round(frac, 4), "moved_rows": int(sel.sum()), "split": info}
    excl_counts = by_source_reason(src, np.where(cand, excl, ""))
    LOG.info("holdout3 from val: candidates %d, eligible %d, moved %d; exclusions %s", int(cand.sum()), int(eligible.sum()), int(moved.sum()), excl_counts)

    # 7. final frames (keyed) for the independent recount + gate
    cols = THREEWAY_SCHEMA[:1] + ["subject", "body", "label", "source", "cluster_id", "label3", "label_set", "base", "k_full", "k_pre", "k_sb", "k_sb_full", "n_alnum"]
    new_train = new_kept.loc[new_kept["dest"] == "train"]
    new_val = new_kept.loc[new_kept["dest"] == "val"]
    new_test = new_kept.loc[new_kept["dest"] == "test3"]
    train_f = pd.concat([train[cols], new_train[cols]], ignore_index=True)
    val_f = pd.concat([val.loc[~moved, cols], new_val[cols]], ignore_index=True)
    test3 = pd.concat([val.loc[moved, cols], new_test[cols]], ignore_index=True)
    links = recount_links_multi(
        {"train": train_f["base"].tolist(), "val": val_f["base"].tolist(), "holdout": holdout["base"].tolist(), "test3": test3["base"].tolist()},
        pretrain_bases, ("holdout", "val", "test3"), nd,
    )
    del pretrain_bases
    leak = leakage_checks3(train_f, val_f, holdout, test3, pretrain_keys, n_pretrain, links, holdout_source, pos_map)
    leak.raise_if_failed()

    # 8. write (atomic): existing rows keep their text byte-for-byte (arrow tables), new rows appended
    def out_table(base: pa.Table, keep: np.ndarray | None, frame: pd.DataFrame, extra: pd.DataFrame | None, tag: str) -> pa.Table:
        t = base.select(FINETUNE_SCHEMA)
        if keep is not None:
            t = t.filter(pa.array(keep))
        if t.num_rows != len(frame) or (t.num_rows and t.column("doc_id")[0].as_py() != frame["doc_id"].iloc[0]):
            raise AssertionError(f"{tag}: arrow table / frame misaligned")
        t = t.append_column("label3", pa.array(frame["label3"].to_numpy(), pa.int64()))
        t = t.append_column("label_set", pa.array(frame["label_set"].tolist(), pa.large_string()))
        if extra is not None and len(extra):
            ex = pa.Table.from_pandas(extra[THREEWAY_SCHEMA].reset_index(drop=True), preserve_index=False).cast(t.schema)
            t = pa.concat_tables([t, ex]).combine_chunks()
        meta = {b"threeway": json.dumps({"label3": SET_OF_LABEL3, "display_names": LABEL3_NAMES, "binary_rollup": UNWANTED_NAME, "built_by": "data_prep.py build3"}).encode()}
        return t.replace_schema_metadata(meta)

    out_tables = {
        "finetune_train": out_table(tables["finetune_train"], None, train, new_train, "train"),
        "finetune_val": out_table(tables["finetune_val"], ~moved, val.loc[~moved], new_val, "val"),
        "test_holdout": out_table(tables["test_holdout"], None, holdout, None, "holdout"),
        "test_holdout3": out_table(tables["finetune_val"], moved, val.loc[moved], new_test, "test3"),
    }
    files: dict[str, str] = {}
    for name, t in out_tables.items():
        path = out_dir / f"{name}.parquet"
        write_table_atomic(t, path)
        files[name] = str(path)
        LOG.info("wrote %s (%d rows, %s)", path, t.num_rows, Counter(t.column("label3").to_pylist()))
    if pretrain_dev_table is not None:
        write_table_atomic(pretrain_dev_table, out_dir / "pretrain.parquet")
        files["pretrain"] = str(out_dir / "pretrain.parquet")
    del tables, out_tables

    # 9. manifest section + report section
    finals = {"finetune_train": train_f, "finetune_val": val_f, "test_holdout": holdout, "test_holdout3": test3}
    elapsed = round(common.timer() - t_start, 1)
    section = {
        "timestamp": datetime.now(timezone.utc).isoformat(timespec="seconds"),
        "limit": args.limit, "dev_subsample": dev, "seed": seed, "rng_stream": "default_rng(seed + 3)",
        "elapsed_seconds": elapsed, "peak_rss_gb": peak_rss_gb(),
        "contract": {"label3": {str(k): v for k, v in SET_OF_LABEL3.items()}, "display_names": {str(k): v for k, v in LABEL3_NAMES.items()},
                     "binary_rollup_name": UNWANTED_NAME, "positive_label3_by_source": pos_map,
                     "binary_label": "unchanged: 1 for phish / spam / unwanted, 0 for legit"},
        "binary_base": {"from": base_origin, "dir": str(base_dir), "counts_by_source_label3": base_counts,
                        "rows": {n: int(len(frames[n])) for n in frames}},
        "sources": {"champa": champa_meta, "bigshp": bigshp_meta},
        "new_rows": {"loaded_by_source": loaded_by_source, "drops_by_source_reason": drops_by_source_reason,
                     "bigshp_label_contradictions_among_key_matches": int(contradictions),
                     "kept_by_source": {str(k): int(v) for k, v in new_kept["source"].value_counts().items()},
                     "destination_by_source": dest_by_source, "pins": new_pins,
                     "test_fraction": tw["new_source_test_fraction"], "val_fraction_of_remainder": tw["new_source_val_fraction"],
                     "split_test3": split_test_info, "split_val": split_val_info, "cluster_id_offset": cluster_offset},
        "holdout3_from_val": {"legit_sources": h3["legit_sources"], "spam_sources": h3["spam_sources"],
                              "candidates_by_source_label3": nested_counts3(val.loc[cand]),
                              "excluded_candidates_by_source_reason": excl_counts,
                              "eligible_by_source_label3": nested_counts3(val.loc[eligible]),
                              "moved_by_source_label3": nested_counts3(val.loc[moved]), "per_class": h3_split_info,
                              "val_rows_linked_to_holdout_info": val_hold_linked},
        "counts": {n: nested_counts3(df) for n, df in finals.items()},
        "totals": {n: split_totals3(df) for n, df in finals.items()},
        "shingle_links": links, "near_dup": nd.__dict__, "leakage": leak.checks, "files": files,
        "config_snapshot": tw,
    }
    manifest_src = processed_dir / "splits_manifest.json"
    manifest = common.read_json(manifest_src) if manifest_src.exists() else {}
    manifest["threeway"] = section
    tmp_path = out_dir / "splits_manifest.json.tmp"
    common.write_json(manifest, tmp_path)
    os.replace(tmp_path, out_dir / "splits_manifest.json")
    report_src = processed_dir / "data_report.md"
    report = report_src.read_text(encoding="utf-8") if report_src.exists() else "# Data report\n"
    marker = "\n## Three-way labels"
    if marker in report:
        report = report[: report.index(marker)].rstrip("\n") + "\n"
    write_text_atomic(report.rstrip("\n") + "\n" + render_threeway_report(section), out_dir / "data_report.md")
    LOG.info("wrote %s and %s in %.1fs (peak RSS %.2f GB)", out_dir / "splits_manifest.json", out_dir / "data_report.md", elapsed, section["peak_rss_gb"])
    print_summary3(section)


def render_threeway_report(s: dict) -> str:
    lines = [f"\n## Three-way labels — build3 {s['timestamp']}", ""]
    if s["dev_subsample"]:
        lines += [f"**DEV SUBSAMPLE (--limit {s['limit']}): counts below are NOT the real build.**", ""]
    c = s["contract"]
    lines += [
        f"label3 {c['label3']}; display names {c['display_names']}; binary roll-up \"{c['binary_rollup_name']}\" = P(phishing)+P(spam). "
        f"Positive rows by provenance: {c['positive_label3_by_source']}. {c['binary_label']}. "
        f"Elapsed {s['elapsed_seconds']}s; peak RSS {s['peak_rss_gb']} GB; binary base from {s['binary_base']['from']} (`{s['binary_base']['dir']}`).",
        "",
        "### New sources",
    ]
    ch = s["sources"]["champa"]
    for fname, f in ch["files"].items():
        lines.append(f"- `{ch['repo']}` / `{fname}` @ `{f['commit_sha']}`: {f['rows']:,} rows ({f['rows_used']:,} used), label3 {ch['label3']} (phishing by provenance)")
    bg = s["sources"]["bigshp"]
    if bg["enabled"]:
        lines.append(f"- bigshp `{bg['repo']}` @ `{bg.get('commit_sha')}` ENABLED despite verdict {bg['verdict']}: {bg['rows_read']:,} rows read, "
                     f"{bg['rows_after_file_filters']:,} after the file filters {bg['file_filter_drops']}; filters {bg['filters']}")
    else:
        lines.append(f"- bigshp (`{bg.get('repo', 'locuoco/the-biggest-spam-ham-phish-email-dataset-300000')}`): research verdict **{bg['verdict']}**, not ingested (`{bg['research']}`)")
    nr = s["new_rows"]
    rows = []
    for src in sorted(nr["loaded_by_source"]):
        d = nr["drops_by_source_reason"].get(src, {})
        dest = nr["destination_by_source"].get(src, {})
        rows.append([src, nr["loaded_by_source"][src], sum(d.values()), nr["kept_by_source"].get(src, 0), dest.get("train", 0), dest.get("val", 0), dest.get("test3", 0)])
    lines += ["", md_table(["source", "loaded", "dropped", "kept", "-> train", "-> val", "-> test_holdout3"], rows),
              "", f"- drops by reason: {nr['drops_by_source_reason']}", f"- near-dup pins: {nr['pins']}",
              f"- bigshp label contradictions among exact matches: {nr['bigshp_label_contradictions_among_key_matches']}",
              "", "### test_holdout3 rows taken from finetune_val",
              f"- legitimate from {s['holdout3_from_val']['legit_sources']}, spam from {s['holdout3_from_val']['spam_sources']}; whole near-dup clusters only; "
              f"excluded candidates by reason: {s['holdout3_from_val']['excluded_candidates_by_source_reason']}",
              f"- per class: { {k: {kk: vv for kk, vv in v.items() if kk != 'split'} for k, v in s['holdout3_from_val']['per_class'].items()} }",
              source_label3_table(s["holdout3_from_val"]["moved_by_source_label3"]), "", "### Splits by label3"]
    for name in THREEWAY_SPLITS:
        t = s["totals"][name]
        lines += [f"#### {name}: {t['rows']:,} rows, by label3 {t['by_label3']}, binary positives {t['binary_positives']:,} ({t['binary_base_rate']})",
                  source_label3_table(s["counts"][name]), ""]
    lines += ["### Leakage assertions (binary gate + three-way)",
              md_table(["check", "count", "asserted zero", "pass"], [[k, v["count"], v["asserted"], v["pass"]] for k, v in s["leakage"].items()]),
              "", "### Files"] + [f"- `{k}`: `{v}`" for k, v in s["files"].items()]
    return "\n".join(lines) + "\n"


def print_summary3(s: dict) -> None:
    print("\n=== data_prep build3 summary (three-way) ===")
    if s["dev_subsample"]:
        print(f"*** DEV SUBSAMPLE --limit {s['limit']} — not a real build ***")
    for name in THREEWAY_SPLITS:
        t = s["totals"][name]
        print(f"{name:>16}: rows {t['rows']:>7,}  by label3 {t['by_label3']}  binary positives {t['binary_positives']:,}")
        for src, cnt in sorted(s["counts"][name].items()):
            print(f"{'':>18}{src:<16} {cnt}")
    print("leakage:")
    for k, v in s["leakage"].items():
        print(f"  {k:<44} {v['count']:>8}  {'ASSERT' if v['asserted'] else 'info  '}  {'ok' if v['pass'] else 'FAIL'}")


# ----------------------------------------------------------------------------- three-way: check
def cmd_check3(args: argparse.Namespace, cfg: dict) -> None:
    """Recompute every three-way assertion from the written files (no reuse of build3 bookkeeping)."""
    processed_dir = Path(cfg["paths"]["data_processed"])
    out_dir = Path(args.out_dir) if args.out_dir else processed_dir
    manifest = common.read_json(out_dir / "splits_manifest.json")
    if "threeway" not in manifest:
        raise AssertionError(f"{out_dir / 'splits_manifest.json'} has no 'threeway' section: run build3 first")
    section = manifest["threeway"]
    tw = threeway_config(cfg)
    pos_map = {str(k): int(v) for k, v in tw["positive_label3_by_source"].items()}
    nd = NearDupConfig.from_config(cfg, cap=int(section["near_dup"]["cap"]))
    holdout_source = cfg["data"]["holdout_test_source"]
    frames: dict[str, pd.DataFrame] = {}
    for name in THREEWAY_SPLITS:
        df = pd.read_parquet(out_dir / f"{name}.parquet")
        if list(df.columns) != THREEWAY_SCHEMA:
            raise AssertionError(f"{name}: expected columns {THREEWAY_SCHEMA}, got {list(df.columns)}")
        df["subject"] = df["subject"].fillna("")
        df["body"] = df["body"].fillna("")
        df["label"] = df["label"].astype(np.int64)
        df["label3"] = df["label3"].astype(np.int64)
        df["source"] = df["source"].astype(str)
        df["base"] = add_key_columns(df)
        frames[name] = df
        LOG.info("read %s: %d rows, label3 %s", name, len(df), dict(Counter(df["label3"].tolist())))
    pre_path = out_dir / "pretrain.parquet" if (out_dir / "pretrain.parquet").exists() else processed_dir / "pretrain.parquet"
    pretrain_keys, pretrain_bases, n_pretrain, _ = pretrain_identity(pre_path, nd, None, np.random.default_rng(0))
    train, val, hold, test3 = (frames[n] for n in THREEWAY_SPLITS)
    links = recount_links_multi(
        {"train": train["base"].tolist(), "val": val["base"].tolist(), "holdout": hold["base"].tolist(), "test3": test3["base"].tolist()},
        pretrain_bases, ("holdout", "val", "test3"), nd,
    )
    del pretrain_bases
    leak = leakage_checks3(train, val, hold, test3, pretrain_keys, n_pretrain, links, holdout_source, pos_map)

    # the TREC-06 holdout must be the binary file with two columns added: same rows, same order, same labels
    snap = processed_dir / BINARY_SNAPSHOT_DIR / "test_holdout.parquet"
    if snap.exists() and not section["dev_subsample"]:
        ref = pq.read_table(snap, columns=["doc_id", "label"])
        same = ref.column("doc_id").to_pylist() == hold["doc_id"].tolist() and ref.column("label").to_pylist() == hold["label"].tolist()
        leak.add("test_holdout_rows_changed_vs_binary_snapshot", 0 if same else 1, True, str(snap))
        ref_tr = pq.read_table(processed_dir / BINARY_SNAPSHOT_DIR / "finetune_train.parquet", columns=["doc_id"]).column("doc_id").to_pylist()
        ref_va = pq.read_table(processed_dir / BINARY_SNAPSHOT_DIR / "finetune_val.parquet", columns=["doc_id"]).column("doc_id").to_pylist()
        binary_ids = set(ref_tr) | set(ref_va)
        now_ids = set(train["doc_id"]) | set(val["doc_id"]) | set(test3["doc_id"])
        leak.add("binary_train_val_rows_missing_from_train_val_holdout3", len(binary_ids - now_ids), True, "every binary train/val row must still exist in train, val or test_holdout3")
        leak.add("binary_train_rows_not_in_train", len(set(ref_tr) - set(train["doc_id"])), True, "train membership is only ever extended")
    for name, df in frames.items():
        want = section["totals"][name]["rows"]
        if want != len(df):
            leak.add(f"{name}_rows_vs_manifest", abs(want - len(df)), True, f"manifest {want} vs file {len(df)}")
    syn = processed_dir / "synthetic.parquet"
    if syn.exists():
        st = pq.read_table(syn)
        if {"label3", "label_set", "label"} <= set(st.schema.names):
            sdf = st.select(["label", "label3", "label_set"]).to_pandas()
            sdf["source"] = "synthetic"
            v = contract_violations(sdf, {"synthetic": 1})
            leak.add("synthetic_label3_contract_violations", sum(v.values()), True, str(v))
        else:
            leak.add("synthetic_missing_label3_columns", 1, False, "synthetic.parquet predates the three-way contract; finetune --synthetic on will refuse it")
    print_summary3({**section, "leakage": leak.checks})
    print("\nmanifest threeway timestamp:", section["timestamp"], "| limit:", section["limit"], "| bigshp enabled:", section["sources"]["bigshp"]["enabled"])
    print("build-time results:", {k: v["count"] for k, v in section["leakage"].items() if v["asserted"]})
    leak.raise_if_failed()
    print("ALL THREE-WAY LEAKAGE ASSERTIONS PASS")


# ----------------------------------------------------------------------------- main
def main(argv: list[str] | None = None) -> None:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--config", default="config.yaml")
    sub = ap.add_subparsers(dest="cmd", required=True)
    sub.add_parser("inspect", help="features + 2 rows of all four datasets (streaming)")
    b = sub.add_parser("build", help="build processed parquet files + manifest + report")
    b.add_argument("--limit", type=int, default=None, help="DEV ONLY: N rows per labelled source, 4N corbt rows")
    b.add_argument("--spam-policy", choices=["positive_class", "map_to_zero", "exclude"], default=None)
    b.add_argument("--no-strip-quotes", action="store_true", help="keep quoted/forwarded history in corbt bodies")
    b.add_argument("--shingle-cap", type=int, default=None, help="ignore shingles shared by more rows than this (default: config data.dedup.near_dup_shingle_cap, else 2000)")
    sub.add_parser("check", help="re-run leakage assertions on the written files")
    b3 = sub.add_parser("build3", help="three-way labels: label3/label_set + Champa phishing corpora + test_holdout3 (reads the processed binary files)")
    b3.add_argument("--limit", type=int, default=None, help="DEV ONLY: N rows per source of every split, N per new file, 4N pretrain rows; writes to data/processed/dev_threeway_limitN/")
    b3.add_argument("--out-dir", default=None, help="output directory (default: data.processed, or the dev dir with --limit)")
    b3.add_argument("--bigshp", choices=["on", "off"], default=None, help="override data.threeway.bigshp.enabled (research verdict: DO_NOT_USE)")
    b3.add_argument("--shingle-cap", type=int, default=None, help="as for build")
    c3 = sub.add_parser("check3", help="recompute the binary + three-way leakage assertions on the written files")
    c3.add_argument("--out-dir", default=None, help="directory written by build3 (default: data.processed)")
    args = ap.parse_args(argv)

    cfg = common.load_config(args.config)
    global LOG, DIALS
    LOG = common.setup_logging("data_prep", cfg["paths"]["logs"])
    DIALS = Dials.from_config(cfg)
    LOG.info("dials: %s", DIALS)
    {"inspect": cmd_inspect, "build": cmd_build, "check": cmd_check, "build3": cmd_build3, "check3": cmd_check3}[args.cmd](args, cfg)


if __name__ == "__main__":
    main()