"""Measure the ADPO dilution pathology on our own preference pairs.

arXiv 2609.32445 (Masking Frequent Tokens Sharpens DPO) reports that under canonical Qwen
tokenization on Anthropic HH-RLHF, 69 token types account for 55.1% of all response tokens and
85.9% of *within-pair shared* token mass, and that those tokens sit symmetrically in the chosen and
rejected replies. A sequence-sum objective then spends most of its gradient on tokens that carry no
preference information between the two arms.

This script computes the same two statistics on our own pairs, per category, because our 6C anti-slop
pairs are two replies in the same voice differing mainly in slop markers -- if anything the shared
mass should be higher here, which would make the dilution worse than their headline number.

Definitions (following the paper):
  within-pair shared mass = sum_i min(count_chosen[i], count_rejected[i]) / (len(chosen)+len(rejected))
  token-type share        = count[i] / sum(count)

Caveat, stated rather than hidden: tokenization here covers the assistant reply text only
(`reasoning_content` + `content`), NOT the chat-template rendering of the full conversation, which is
what the trainer tokenizes. The prompt side cancels between the two arms of a pair, so the shared-mass
statistic is unaffected; absolute token counts per reply may differ slightly.

    uv run --with regex --offline python3 persona/scripts/adpo_dilution_check.py
"""
import argparse
import collections
import json
import pathlib
import sys

ROOT = pathlib.Path(__file__).resolve().parents[2]   # repo root: this file is persona/scripts/
sys.path.insert(0, str(ROOT / "vendor" / "strata" / "tools"))
import strata_tokenizer  # noqa: E402


def load_tokenizer(directory):
    """`Tokenizer(tokens, merges)` needs the GGUF-style lists, so rebuild them from the pack directory."""
    d = pathlib.Path(directory)
    vocab = json.loads((d / "vocab.json").read_text(encoding="utf-8"))
    tokens = [""] * (max(vocab.values()) + 1)
    for token, index in vocab.items():
        tokens[index] = token
    merges = [m for m in (d / "merges.txt").read_text(encoding="utf-8").splitlines()
              if m and not m.startswith("#") and " " in m]
    types_path = d / "token_type.json"
    token_types = json.loads(types_path.read_text(encoding="utf-8")) if types_path.exists() else None
    return strata_tokenizer.Tokenizer(tokens, merges, token_types=token_types)


def reply_text(msg):
    parts = [msg.get("reasoning_content") or "", msg.get("content") or ""]
    return "\n".join(p for p in parts if p)


def analyse(tok, path, categories, ks):
    shared = collections.Counter()
    total_ids = collections.Counter()
    per_cat = collections.defaultdict(lambda: {"pairs": 0, "chosen": 0, "rejected": 0,
                                               "shared": 0, "total": 0})
    for line in path.open(encoding="utf-8"):
        pair = json.loads(line)
        cat = pair.get("category")
        if categories and cat not in categories:
            continue
        if pair.get("split") != "train":
            continue
        a = tok.encode(reply_text(pair["chosen_assistant"]))
        b = tok.encode(reply_text(pair["rejected_assistant"]))
        if not a or not b:
            continue
        ca, cb = collections.Counter(a), collections.Counter(b)
        sh = sum(min(v, cb[k]) for k, v in ca.items())
        shared.update({k: min(v, cb[k]) for k, v in ca.items()})
        total_ids.update(a)
        total_ids.update(b)
        st = per_cat[cat]
        st["pairs"] += 1
        st["chosen"] += len(a)
        st["rejected"] += len(b)
        st["shared"] += sh
        st["total"] += len(a) + len(b)
    return shared, total_ids, per_cat


def main(argv=None):
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("--pairs", default=str(ROOT / "persona/data/agent_persona_v1/ftpo_dpo_pairs.jsonl"))
    ap.add_argument("--tokenizer", default=str(ROOT / "models/K512_iq2_xs/tokenizer"))
    ap.add_argument("--categories", default="", help="comma list; default all")
    ap.add_argument("--ks", default="32,69,128,256,512")
    a = ap.parse_args(argv)
    tok = load_tokenizer(a.tokenizer)
    cats = [c for c in a.categories.split(",") if c] or None
    ks = [int(k) for k in a.ks.split(",")]
    shared, total_ids, per_cat = analyse(tok, pathlib.Path(a.pairs), cats, ks)

    print(f"pairs analysed: {sum(v['pairs'] for v in per_cat.values())}"
          f"   distinct token types: {len(total_ids):,}")
    print("\nper category:")
    for cat, st in sorted(per_cat.items()):
        mass = st["shared"] / st["total"] if st["total"] else 0.0
        print(f"  {cat:26s} n={st['pairs']:4d}  reply tok mean "
              f"{(st['chosen']+st['rejected'])/2/st['pairs']:6.1f}  "
              f"shared mass {mass:6.2%}")

    grand = sum(total_ids.values())
    shared_total = sum(shared.values())
    print(f"\ncorpus: {grand:,} reply tokens, {shared_total:,} within-pair shared "
          f"({shared_total/grand:.2%} of all)")
    print(f"\n{'top types':>10} {'share of all tokens':>20} {'share of shared mass':>21}")
    for k in ks:
        types = dict(total_ids.most_common(k))
        all_share = sum(types.values()) / grand
        sh_share = sum(shared[t] for t in types) / shared_total
        print(f"{k:>10} {all_share:>19.2%} {sh_share:>20.2%}")
    return 0


if __name__ == "__main__":
    sys.exit(main())