#!/usr/bin/env python3
"""评测 v2.1 四头模型的六种精度变体：teacher holdout / earlier val + 与 fp32 一致率。"""
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


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--root", default=".work/hf/v2.1_six")
    ap.add_argument("--teacher", default="data/optimized_4h/teacher_holdout.jsonl")
    ap.add_argument("--earlier", default="data/optimized_4h/earlier_val.jsonl")
    ap.add_argument("--out", default="release/v2.1_variant_eval.json")
    args = ap.parse_args()

    root = Path(args.root)
    sets = {"teacher_holdout": [json.loads(l) for l in open(args.teacher, encoding="utf-8")],
            "earlier_val": [json.loads(l) for l in open(args.earlier, encoding="utf-8")]}
    variants = ["fp32", "fp16", "bf16", "fp8", "int8", "int4"]
    results = {}
    ref = {}
    print(f"{'variant':<8}{'set':<18}{'bin_acc':>9}{'adF1':>8}{'kind':>8}{'core':>8}"
          f"{'scope':>8}{'sevMAE':>8}{'agree':>8}{'MB':>8}")
    for v in variants:
        d = root / v
        tok = AutoTokenizer.from_pretrained(d)
        model = load_dual_head_any(d, device="cuda",
                                   dtype=torch.float32 if v == "fp32" else torch.bfloat16)
        results[v] = {"size_mb": round((d / "model.safetensors").stat().st_size / 2**20, 1)}
        for name, rows in sets.items():
            ids, am = tokenize(rows, tok, 384)
            lens = [len(x) for x in ids]
            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], 64, 384)
            vidx = [i for i, r in enumerate(rows) if r["label_sev"] >= 0]
            em = None
            if vidx:
                y = np.array([rows[i]["label_sev"] for i in vidx])
                exp = (pv[vidx] * np.arange(4)).sum(1)
                em = float(np.abs(exp - y).mean())
            entry = {k: m.get(k) for k in ("binary_acc", "ad_f1", "kind_acc", "kind_core_f1",
                                           "scope_acc", "sev_acc", "sev_mae", "kind_n", "scope_n", "sev_n")}
            entry["sev_expected_mae"] = em
            pred = pb[:, 0] >= 0.5
            kk = np.array([i for i, r in enumerate(rows) if r["label_kind"] >= 0])
            if v == "fp32":
                ref[name] = (pred, pk.argmax(1))
                entry["agree_bin"] = 1.0
                entry["agree_kind"] = 1.0
            else:
                rp, rk = ref[name]
                entry["agree_bin"] = round(float((pred == rp).mean()), 4)
                entry["agree_kind"] = round(float((pk.argmax(1)[kk] == rk[kk]).mean()), 4)
            results[v][name] = entry
            print(f"{v:<8}{name:<18}{m['binary_acc']*100:>8.2f}%{m['ad_f1']:>8.3f}"
                  f"{m.get('kind_acc',0)*100:>7.2f}%{m.get('kind_core_f1',0):>8.3f}"
                  f"{m.get('scope_acc',0)*100:>7.2f}%{em if em is not None else 0:>8.3f}"
                  f"{entry['agree_bin']*100:>7.2f}%{results[v]['size_mb']:>8.1f}")
        del model
        torch.cuda.empty_cache()
    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()