#!/usr/bin/env python3
"""评测双头模型的六种精度变体：二分类 / 7 类指标 + 与 fp32 预测一致率。

用法:
  python scripts/eval_dual_variants.py --root .work/hf/v2_six \
      --val data/processed_dual/val.jsonl --holdout data/processed_dual/holdout827.jsonl \
      --out release/v2_variant_eval.json
"""
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__).parent))
from dual_head_model import load_dual_head_any  # noqa: E402
from train_dual_head import evaluate, tokenize  # noqa: E402


def run_variant(d, rows, split, dtype):
    tok = AutoTokenizer.from_pretrained(d)
    ids, am = tokenize(rows, tok, 384)
    lengths = [len(x) for x in ids]
    model = load_dual_head_any(d, device="cuda", dtype=dtype)
    m, pb, pk = evaluate(model, rows, ids, am,
                         [r["label_bin"] for r in rows],
                         [r["label_kind"] for r in rows], lengths, 128, 384)
    del model
    torch.cuda.empty_cache()
    return m, pb, pk


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--root", default=".work/hf/v2_six")
    ap.add_argument("--val", default="data/processed_dual/val.jsonl")
    ap.add_argument("--holdout", default="data/processed_dual/holdout827.jsonl")
    ap.add_argument("--out", default="release/v2_variant_eval.json")
    args = ap.parse_args()

    root = Path(args.root)
    splits = {
        "val": [json.loads(l) for l in open(args.val, encoding="utf-8") if l.strip()],
        "holdout827": [json.loads(l) for l in open(args.holdout, encoding="utf-8") if l.strip()],
    }
    variants = ["fp32", "fp16", "bf16", "fp8", "int8", "int4"]
    results = {}
    ref = {}
    print(f"{'variant':<8}{'val bin':>9}{'val kind':>10}{'val agree':>11}"
          f"{'hold bin':>10}{'hold kind':>11}{'hold agree':>12}{'size MB':>10}")
    for v in variants:
        d = root / v
        dtype = torch.float32 if v == "fp32" else torch.bfloat16
        results[v] = {"size_mb": round((d / "model.safetensors").stat().st_size / 2**20, 1)}
        for split, rows in splits.items():
            m, pb, pk = run_variant(d, rows, split, dtype)
            entry = {k: m[k] for k in ("binary_acc", "ad_f1", "kind_acc", "kind_core_f1")}
            if split == "val" and v == "fp32":
                ref["val"] = (pb[:, 0] >= 0.5, pk.argmax(1))
            if split == "holdout827" and v == "fp32":
                ref["holdout827"] = (pb[:, 0] >= 0.5, pk.argmax(1))
            if split in ref:
                rp, rk = ref[split]
                pred = pb[:, 0] >= 0.5
                entry["agree_bin"] = round(float((pred == rp).mean()), 4)
                kk = np.array([i for i, r in enumerate(rows) if r["label_kind"] >= 0])
                entry["agree_kind"] = round(float((pk.argmax(1)[kk] == rk[kk]).mean()), 4)
            results[v][split] = entry
        rv, rh = results[v]["val"], results[v]["holdout827"]
        print(f"{v:<8}{rv['binary_acc']*100:>8.2f}%{rv['kind_acc']*100:>9.2f}%"
              f"{rv.get('agree_bin', 1)*100:>10.2f}%"
              f"{rh['binary_acc']*100:>9.2f}%{rh['kind_acc']*100:>10.2f}%"
              f"{rh.get('agree_bin', 1)*100:>11.2f}%{results[v]['size_mb']:>10.1f}")

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


if __name__ == "__main__":
    main()