#!/usr/bin/env python3
"""按标记形态分列的识别率实测（推理端不做任何预处理）。

同一批文本改写成 6 种标记形态，分别给出 acc / 广告精确率 / 广告召回 / F1 / 漏判 / 误判。
用于验证"模型原生对标记写法不敏感"——不依赖任何归一化步骤。

用法:
  python scripts/eval_marker_metrics.py --model outputs/ad-clf-v2b/best
  python scripts/eval_marker_metrics.py --model outputs/ad-clf-v2b/best --only_archive
"""
import argparse
import json
import re

import numpy as np
import torch
from sklearn.metrics import precision_recall_fscore_support
from transformers import AutoModelForSequenceClassification, AutoTokenizer

LABEL_MAP = {"ad": 0, "normal": 1}
MARKER = re.compile(r"［([^］]{1,16})］")
EN = {"引用": "quoted", "图片": "image", "贴纸": "sticker", "转发自": "forwarded_from",
      "预览": "preview", "隐藏链接": "hidden_link", "按钮": "button", "文件": "file",
      "联系人卡片": "contact_card", "访客bot": "guest_bot", "音频": "audio",
      "故事": "story", "清单": "list", "经由": "via", "入群资料检查": "join_check",
      "前置号复核": "prewarm_check", "富消息": "rich_message", "视频": "video"}


def to_xml_open(t: str) -> str:
    return MARKER.sub(lambda m: f"<{EN.get(m.group(1), 'ref')}>", t)


def to_xml_closed(t: str) -> str:
    if "［引用］" in t:
        head, _, tail = t.partition("［引用］")
        t = head + "<quoted>" + tail + "</quoted>"
    return MARKER.sub(lambda m: f"<{EN.get(m.group(1), 'ref')}>", t)


VARIANTS = {
    "规范 ［中文］": lambda t: t,
    "XML <en> 开标签": to_xml_open,
    "XML <en>…</en> 闭合": to_xml_closed,
    "半角 [en]": lambda t: MARKER.sub(lambda m: f"[{EN.get(m.group(1), m.group(1))}]", t),
    "方头 【中文】": lambda t: MARKER.sub(lambda m: f"【{m.group(1)}】", t),
    "未知标签 <marker_x>": lambda t: MARKER.sub(
        lambda m: f"<marker_{abs(hash(m.group(1))) % 997}>", t),
}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True)
    ap.add_argument("--val", default="data/processed_v2/val.jsonl")
    ap.add_argument("--only_archive", action="store_true",
                    help="只评 archive 净新增子集（对 v1/v2/v2b 都干净的无污染 holdout）")
    ap.add_argument("--archive_file",
                    default="data/raw/archive/converted/archive.strong.binary.jsonl")
    ap.add_argument("--max_len", type=int, default=512)
    ap.add_argument("--batch_size", type=int, default=64)
    args = ap.parse_args()

    rows = [json.loads(l) for l in open(args.val, encoding="utf-8") if l.strip()]
    if args.only_archive:
        clean = {json.loads(l)["text"] for l in open(args.archive_file, encoding="utf-8")}
        rows = [r for r in rows if r["text"] in clean]
        print(f"[干净 holdout] archive 净新增子集 {len(rows)} 条（模型训练时未见过）")
    y = np.array([LABEL_MAP[r["label"]] for r in rows])
    tok = AutoTokenizer.from_pretrained(args.model)
    model = AutoModelForSequenceClassification.from_pretrained(
        args.model, dtype=torch.bfloat16, attn_implementation="sdpa").cuda().eval()

    def run(texts):
        out = []
        with torch.inference_mode():
            for i in range(0, len(texts), args.batch_size):
                enc = tok(texts[i:i + args.batch_size], truncation=True,
                          max_length=args.max_len, padding=True,
                          return_tensors="pt").to("cuda")
                out.append(torch.softmax(model(**enc).logits.float(), -1).cpu().numpy())
        return np.concatenate(out)

    print(f"模型 {args.model} | 推理端不做预处理 | 样本 {len(rows)} "
          f"(ad {int((y==0).sum())} / normal {int((y==1).sum())})\n")
    print(f"{'标记形态':<22}{'acc':>8}{'ad精确':>9}{'ad召回':>9}{'ad-F1':>8}{'漏判':>7}{'误判':>7}")
    print("-" * 70)
    for name, fn in VARIANTS.items():
        pred = run([fn(r["text"]) for r in rows]).argmax(1)
        prec, rec, f1, _ = precision_recall_fscore_support(
            y, pred, labels=[0], average=None, zero_division=0)
        fn_ = int(((pred == 1) & (y == 0)).sum())
        fp_ = int(((pred == 0) & (y == 1)).sum())
        print(f"{name:<22}{(pred==y).mean():>8.4f}{prec[0]:>9.4f}{rec[0]:>9.4f}{f1[0]:>8.4f}"
              f"{fn_:>7}{fp_:>7}")
    print("-" * 70)
    print("漏判 = 真广告判成正常；误判 = 真正常判成广告")


if __name__ == "__main__":
    main()