"""Benchmark the humanizer model.

Measures:
- corpus BLEU (sacrebleu) between outputs and rule references on AI->human
  held-out rows (meaning preservation),
- AI-cliché removal rate on AI-flavored inputs,
- identity stability on ALL clean-prose (identity) pairs in the corpus,
- protected-span preservation on ALL examples containing protected spans,
- latency (characters per second, MPS/CPU).
"""

import argparse
import json
import re
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from humanize import Humanizer, protect

EN_CLICHES = [
    "it is important to note that", "it is worth noting that", "moreover",
    "furthermore", "in conclusion", "leverage", "seamless", "robust",
    "testament", "delve", "cutting-edge", "state-of-the-art",
]
ZH_CLICHES = [
    "值得注意的是", "综上所述", "赋能", "降本增效", "闭环", "无缝",
    "由此可见", "总而言之",
]


def scan(data_dir: str):
    rows = []
    for split in ("train", "val", "test"):
        with open(Path(data_dir) / f"{split}.jsonl", encoding="utf-8") as f:
            for line in f:
                rows.append(json.loads(line))
    return rows


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", default="checkpoints/humanize-text-model")
    parser.add_argument("--data", default="data")
    parser.add_argument("--ai-limit", type=int, default=500)
    args = parser.parse_args()

    from sacrebleu import corpus_bleu

    humanizer = Humanizer(model_id=args.model)
    rows = scan(args.data)

    ai_rows = [r for r in rows if r["input_text"] != r["output_text"]][: args.ai_limit]
    identity_rows = [r for r in rows if r["input_text"] == r["output_text"]]
    protected_rows = [
        r for r in rows
        if re.search(r"PROTECTED_\d+", r["input_text"]) or re.search(
            r"https?://|PROTECTED_\d+|\d", r["input_text"]
        )
    ]
    protected_rows = [r for r in protected_rows if "PROTECTED_" in r["input_text"] or "https://" in r["input_text"]]

    # Deduplicate by input so samples are unique.
    def dedupe(items):
        seen, out = set(), []
        for r in items:
            if r["input_text"] in seen:
                continue
            seen.add(r["input_text"])
            out.append(r)
        return out

    ai_rows = dedupe(ai_rows)
    identity_rows = dedupe(identity_rows)
    protected_rows = dedupe(protected_rows)

    print(f"AI->human rows: {len(ai_rows)} | identity rows: {len(identity_rows)} | protected rows: {len(protected_rows)}", flush=True)

    # 1. AI -> human: BLEU + cliché removal.
    refs, hyps = [], []
    cliche_in = 0
    cliche_removed = 0
    t0 = time.time()
    for n, row in enumerate(ai_rows, 1):
        out = humanizer.humanize(row["input_text"], num_beams=3)
        hyps.append(out)
        refs.append(row["output_text"])
        lower = out.lower()
        hits_in = [c for c in (ZH_CLICHES if row["lang"] == "zh" else EN_CLICHES)
                   if c in (row["input_text"] if row["lang"] == "zh" else row["input_text"].lower())]
        hits_out = [c for c in (ZH_CLICHES if row["lang"] == "zh" else EN_CLICHES)
                    if c in lower]
        if hits_in:
            cliche_in += 1
            if not hits_out:
                cliche_removed += 1
        if n % 50 == 0:
            print(f"  ...ai {n}/{len(ai_rows)}", flush=True)
    bleu = corpus_bleu(hyps, [refs]).score

    # 2. Identity stability.
    identity_ok = 0
    for n, row in enumerate(identity_rows, 1):
        out = humanizer.humanize(row["input_text"], num_beams=3)
        # Allow light edits, no big drift: use char-level edit ratio.
        ratio = sum(1 for a, b in zip(out, row["input_text"]) if a != b) / max(len(row["input_text"]), 1)
        if ratio <= 0.15:
            identity_ok += 1
        if n % 25 == 0:
            print(f"  ...identity {n}/{len(identity_rows)}", flush=True)

    # 3. Protected spans.
    protected_ok = 0
    protected_total = 0
    failures = []
    for n, row in enumerate(protected_rows, 1):
        masked, spans = protect(row["input_text"])
        out = humanizer.humanize(row["input_text"], num_beams=3)
        for span in spans:
            protected_total += 1
            if span in out:
                protected_ok += 1
            else:
                failures.append((row["lang"], span, out))
        if n % 25 == 0:
            print(f"  ...protected {n}/{len(protected_rows)}", flush=True)

    elapsed = time.time() - t0
    chars = sum(len(r["input_text"]) for r in ai_rows + identity_rows + protected_rows)

    print("=" * 60)
    print(f"model: {args.model}")
    print(f"corpus BLEU (AI->human, n={len(ai_rows)}): {bleu:.1f}")
    print(f"cliche removal rate: {cliche_removed / max(cliche_in, 1):.1%} ({cliche_removed}/{cliche_in})")
    print(f"identity stability (n={len(identity_rows)}): {identity_ok / max(len(identity_rows), 1):.1%} ({identity_ok}/{len(identity_rows)})")
    print(f"protected-span preservation: {protected_ok}/{protected_total} "
          f"({protected_ok / max(protected_total, 1):.1%})")
    print(f"throughput: {chars / max(elapsed, 0.01):.0f} chars/s")
    if failures:
        print(f"\nprotected failures ({len(failures)}):")
        for lang, span, out in failures[:10]:
            print(f"  [{lang}] lost {span!r} in {out[:80]!r}")

    print("\nsample outputs:")
    for row in ai_rows[:5]:
        print("  IN :", row["input_text"][:100])
        print("  OUT:", humanizer.humanize(row["input_text"], num_beams=3)[:100])
        print()


if __name__ == "__main__":
    main()