#!/usr/bin/env python3
"""四头模型推理：一次前向输出 is_ad / kind / scope / severity（0–3 期望分数）。"""
import argparse
import json
from pathlib import Path

import torch
from transformers import AutoTokenizer

from dual_head_model import load_dual_head_any

KINDS = ["crypto", "gambling", "porn", "porn_bait", "promo", "scam", "spam_flood"]
DEMO = [
    "加微信 abc123 免费领取涨粉秘籍，名额有限！",
    "今晚合唱团排练改到七点半，别走错了",
    "这款面霜绝了，用了一个月肉眼可见变好，链接放评论区",
]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", default="outputs/ad-4head-teacher-heads/best")
    ap.add_argument("--text", default=None)
    ap.add_argument("--file", default=None, help="jsonl({text}) 或 txt，每行一条")
    ap.add_argument("--threshold", type=float, default=0.5)
    ap.add_argument("--kind_threshold", type=float, default=0.0,
                    help=">0 时 kind 置信度低于该值输出 ad_uncertain")
    ap.add_argument("--batch_size", type=int, default=64)
    ap.add_argument("--max_len", type=int, default=384)
    args = ap.parse_args()

    if args.text:
        texts = [args.text]
    elif args.file:
        if args.file.endswith(".jsonl"):
            texts = [json.loads(l)["text"] for l in open(args.file, encoding="utf-8") if l.strip()]
        else:
            texts = [l.rstrip("\n") for l in open(args.file, encoding="utf-8") if l.strip()]
    else:
        texts = DEMO

    cfg_path = Path(args.model) / "dual_config.json"
    dc = json.loads(cfg_path.read_text(encoding="utf-8")) if cfg_path.is_file() else {}
    kind_labels = dc.get("kind_labels") or KINDS
    scope_labels = dc.get("scope_labels") or []
    severity_labels = dc.get("severity_labels") or []

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = load_dual_head_any(args.model, device=device, dtype=torch.bfloat16)
    tok = AutoTokenizer.from_pretrained(args.model)

    with torch.inference_mode():
        for i in range(0, len(texts), args.batch_size):
            chunk = texts[i:i + args.batch_size]
            enc = tok(chunk, truncation=True, max_length=args.max_len,
                      padding=True, return_tensors="pt").to(device)
            with torch.autocast(device_type=device, dtype=torch.bfloat16, enabled=device == "cuda"):
                blog, klog, slog, vlog, _ = model.forward_all(enc["input_ids"], enc["attention_mask"])
            p_ad = torch.softmax(blog.float(), -1)[:, 0].cpu()
            pk = torch.softmax(klog.float(), -1).cpu()
            ps = torch.softmax(slog.float(), -1).cpu() if slog is not None else None
            pv = torch.softmax(vlog.float(), -1).cpu() if vlog is not None else None
            for j, (t, pa) in enumerate(zip(chunk, p_ad)):
                is_ad = bool(pa >= args.threshold)
                k = int(pk[j].argmax())
                kind = kind_labels[k] if is_ad else "normal"
                if is_ad and args.kind_threshold > 0 and float(pk[j][k]) < args.kind_threshold:
                    kind = "ad_uncertain"
                out = {"is_ad": is_ad, "p_ad": round(float(pa), 6),
                       "kind": kind, "kind_conf": round(float(pk[j][k]), 6),
                       "kind_probs": {n: round(float(pk[j][x]), 6) for x, n in enumerate(kind_labels)}}
                if is_ad and ps is not None:
                    s = int(ps[j].argmax())
                    out["scope"] = scope_labels[s] if scope_labels else str(s)
                    out["scope_probs"] = {n: round(float(ps[j][x]), 6) for x, n in enumerate(scope_labels)}
                if is_ad and pv is not None:
                    probs = [float(x) for x in pv[j]]
                    out["severity_score"] = round(sum(x * p for x, p in enumerate(probs)), 4)
                    out["severity_probs"] = {n: round(probs[x], 6) for x, n in enumerate(severity_labels)}
                print(json.dumps({"text": t, **out}, ensure_ascii=False))


if __name__ == "__main__":
    main()