#!/usr/bin/env python3
"""广告二分类错误分析：逐条落盘概率，打印全部错判样本供人工核对。

用法:
  python scripts/analyze_ad_errors.py --model release/laya-ad-v1.1-bf16 \
      --data data/processed_v2/val.jsonl --only_archive
  python scripts/analyze_ad_errors.py --model release/laya-ad-v1.1-bf16 \
      --data data/processed_v2/val.jsonl --out outputs/ad_errors_val.jsonl
"""
import argparse
import json
from collections import Counter

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

LABEL_MAP = {"ad": 0, "normal": 1}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", default="release/laya-ad-v1.1-bf16")
    ap.add_argument("--data", default="data/processed_v2/val.jsonl")
    ap.add_argument("--only_archive", action="store_true",
                    help="只评 archive 净新增子集（827 条无污染 holdout）")
    ap.add_argument("--archive_file",
                    default="data/raw/archive/converted/archive.strong.binary.jsonl")
    ap.add_argument("--out", default="outputs/ad_preds.jsonl")
    ap.add_argument("--max_len", type=int, default=512)
    ap.add_argument("--batch_size", type=int, default=64)
    ap.add_argument("--full_text", action="store_true", help="错例打印全文而非截断")
    args = ap.parse_args()

    rows = [json.loads(l) for l in open(args.data, 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]
    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()
    names = [model.config.id2label[i] for i in range(model.config.num_labels)]

    probs_all = []
    with torch.inference_mode():
        for i in range(0, len(rows), args.batch_size):
            enc = tok([r["text"] for r in rows[i:i + args.batch_size]], truncation=True,
                      max_length=args.max_len, padding=True,
                      return_tensors="pt").to("cuda")
            probs_all.append(torch.softmax(model(**enc).logits.float(), -1).cpu().numpy())
    probs = np.concatenate(probs_all)
    pred = probs.argmax(1)

    print(f"模型 {args.model} | 数据 {args.data} | 样本 {len(rows)} "
          f"(ad {int((y == 0).sum())} / normal {int((y == 1).sum())})")
    print(classification_report(y, pred, target_names=[names[0], names[1]], digits=4))
    print(f"混淆矩阵（行=真实，列=预测，顺序 = {names[0]}, {names[1]}）:")
    print(confusion_matrix(y, pred))
    prec, rec, f1, _ = precision_recall_fscore_support(
        y, pred, labels=[0], average=None, zero_division=0)
    print(f"ad 精确 {prec[0]:.4f} / 召回 {rec[0]:.4f} / F1 {f1[0]:.4f} | "
          f"漏判(ad→normal) {int(((pred == 1) & (y == 0)).sum())} | "
          f"误判(normal→ad) {int(((pred == 0) & (y == 1)).sum())}")

    with open(args.out, "w", encoding="utf-8") as f:
        for r, p in zip(rows, probs):
            f.write(json.dumps({"text": r["text"], "label": r["label"],
                                "p_ad": float(p[0]), "p_normal": float(p[1]),
                                "pred": names[int(p.argmax())]},
                               ensure_ascii=False) + "\n")
    print("逐条概率已存:", args.out)

    err_idx = np.where(pred != y)[0]
    print(f"\n===== 错判明细（{len(err_idx)} 条）=====")
    for i in err_idx:
        t = rows[i]["text"] if args.full_text else rows[i]["text"][:220].replace("\n", "\\n")
        print(f"[{i}] 真={names[y[i]]} 判={names[pred[i]]} "
              f"p_ad={probs[i][0]:.4f} p_normal={probs[i][1]:.4f}")
        print(f"     {t}")
        print()

    # 低置信度（不管对错）——最可能藏着标签噪声或边界样本
    conf = probs.max(1)
    low = np.where(conf < 0.9)[0]
    print(f"===== 低置信度样本（max prob < 0.9，共 {len(low)} 条，按置信度升序前 30）=====")
    for i in low[np.argsort(conf[low])][:30]:
        t = rows[i]["text"] if args.full_text else rows[i]["text"][:160].replace("\n", "\\n")
        mark = "错" if pred[i] != y[i] else "对"
        print(f"[{i}] {mark} 真={names[y[i]]} 判={names[pred[i]]} "
              f"p_ad={probs[i][0]:.4f}  {t}")

    # ad 类按 p_ad 分桶，看召回风险集中在哪
    ad_idx = np.where(y == 0)[0]
    buckets = Counter()
    for i in ad_idx:
        p = probs[i][0]
        buckets["0.00-0.10" if p < .1 else "0.10-0.50" if p < .5 else
                "0.50-0.90" if p < .9 else "0.90-0.99" if p < .99 else "0.99-1.00"] += 1
    print("\n真广告的 p_ad 分布:", dict(buckets))


if __name__ == "__main__":
    main()