#!/usr/bin/env python3
"""v2.2 vs v2.1 四头模型对比评测：teacher_holdout / earlier_val / gmeow_val128。

gmeow_val128 是 v2.1 训练时没见过的 gmeow 留出切片（128 条，来自 ads_final.jsonl
与 data/cleaned/gmeow/val_split_texts.jsonl 的交集），专测 gmeow 语料带来的增益。
"""
import argparse
import json
import sys
from pathlib import Path

import numpy as np
import torch
from transformers import AutoTokenizer

sys.path.insert(0, str(Path(__file__).resolve().parent))
from dual_head_model import load_dual_head_any  # noqa: E402
from train_multi_head import evaluate, tokenize  # noqa: E402

KINDS = ["crypto", "gambling", "porn", "porn_bait", "promo", "scam", "spam_flood"]


def load(p):
    return [json.loads(l) for l in open(p, encoding="utf-8") if l.strip()]


def gmeow_val_rows(ads_path, val_texts_path):
    hold = {r["text"] for r in load(val_texts_path)}
    rows = []
    for r in load(ads_path):
        if r.get("text") not in hold or r.get("kind") not in KINDS:
            continue
        sev = r.get("severity")
        rows.append({
            "text": r["text"], "label_bin": 0,
            "label_kind": KINDS.index(r["kind"]),
            "label_scope": 0 if r.get("scope") == "account" else 1,
            "label_sev": int(round(sev)) if isinstance(sev, (int, float)) and 0 <= sev <= 3 else -100,
        })
    return rows


def eval_model(model_dir, sets, max_len=384, bs=64):
    tok = AutoTokenizer.from_pretrained(model_dir)
    model = load_dual_head_any(model_dir, device="cuda", dtype=torch.bfloat16)
    out = {}
    for name, rows in sets.items():
        ids, am = tokenize(rows, tok, max_len)
        m, pb, pk, ps, pv = evaluate(
            model, rows, ids, am,
            [r["label_bin"] for r in rows], [r["label_kind"] for r in rows],
            [r["label_scope"] for r in rows], [r["label_sev"] for r in rows], bs, max_len)
        vidx = [i for i, r in enumerate(rows) if r["label_sev"] >= 0]
        if vidx:
            y = np.array([rows[i]["label_sev"] for i in vidx])
            exp = (pv[vidx] * np.arange(4)).sum(1)
            m["sev_expected_mae"] = round(float(np.abs(exp - y).mean()), 4)
        out[name] = m
        print(f"  {name:<16} acc={m['binary_acc']*100:6.2f}% adF1={m['ad_f1']:.3f} "
              f"kind={m.get('kind_acc',0)*100:6.2f}% core={m.get('kind_core_f1',0):.3f} "
              f"scope={m.get('scope_acc',0)*100:6.2f}% sevMAE={m.get('sev_expected_mae')}", flush=True)
    del model
    torch.cuda.empty_cache()
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--v22", default="outputs/ad-4head-v22/best")
    ap.add_argument("--v21", default="outputs/ad-4head-teacher-heads/best")
    ap.add_argument("--teacher", default="data/optimized_4h_v22/teacher_holdout.jsonl")
    ap.add_argument("--earlier", default="data/optimized_4h_v22/earlier_val.jsonl")
    ap.add_argument("--ads", default="data/raw/gmeow/server_export_20261031/ads_final.jsonl")
    ap.add_argument("--val_texts", default="data/cleaned/gmeow/val_split_texts.jsonl")
    ap.add_argument("--out", default="outputs/v22_vs_v21_eval.json")
    args = ap.parse_args()

    sets = {
        "teacher_holdout": load(args.teacher),
        "earlier_val": load(args.earlier),
        "gmeow_val128": gmeow_val_rows(args.ads, args.val_texts),
    }
    print(f"gmeow_val128 n={len(sets['gmeow_val128'])} "
          f"kind={sum(1 for r in sets['gmeow_val128'] if r['label_kind']>=0)} "
          f"sev={sum(1 for r in sets['gmeow_val128'] if r['label_sev']>=0)}")

    results = {}
    for tag, d in (("v2.2", args.v22), ("v2.1", args.v21)):
        print(f"== {tag} ({d}) ==", flush=True)
        results[tag] = eval_model(d, sets)

    out = Path(args.out)
    out.write_text(json.dumps(results, ensure_ascii=False, indent=2), encoding="utf-8")
    print("->", out.resolve())


if __name__ == "__main__":
    main()