#!/usr/bin/env python3
"""Evaluate QEV2 multitask records by language, origin, and question type."""
from __future__ import annotations
import argparse, json, math, time
from collections import Counter, defaultdict
from pathlib import Path

import torch
from qev.model import QevModel
from qev.train import load_records, normalize
from qev.records import load_jsonl


def load_records_with_meta(path):
    """Normalize records while retaining source metadata for language splits."""
    out = []
    for raw in load_jsonl(path):
        row = normalize(raw)
        meta = dict(raw.get("_meta", {})) if isinstance(raw, dict) else {}
        # Some corpora place language at the top level instead of _meta.
        if "language" not in meta and isinstance(raw, dict) and raw.get("language"):
            meta["language"] = raw["language"]
        row["_meta"] = meta
        out.append(row)
    return out


def add(cm, y, p):
    cm["n"] += 1
    cm["correct"] += int(y == p)
    if y and p: cm["tp"] += 1
    elif y and not p: cm["fn"] += 1
    elif not y and p: cm["fp"] += 1
    else: cm["tn"] += 1


def metric(cm):
    n = cm["n"]
    acc = cm["correct"] / n if n else 0.0
    prec = cm["tp"] / (cm["tp"] + cm["fp"]) if cm["tp"] + cm["fp"] else 0.0
    rec = cm["tp"] / (cm["tp"] + cm["fn"]) if cm["tp"] + cm["fn"] else 0.0
    f1 = 2 * prec * rec / (prec + rec) if prec + rec else 0.0
    f2 = 5 * prec * rec / (4 * prec + rec) if 4 * prec + rec else 0.0
    return {"n": n, "accuracy": acc, "precision": prec, "recall": rec, "f1": f1, "f2": f2,
            "tp": cm["tp"], "fp": cm["fp"], "fn": cm["fn"], "tn": cm["tn"]}


def evaluate(args):
    dtype = torch.bfloat16 if args.dtype == "bf16" and args.device.startswith("cuda") else torch.float32
    model = QevModel.load_checkpoint(args.checkpoint, args.device, dtype, trainable=False)
    model.eval()
    records = load_records_with_meta(args.data)
    if args.max_records: records = records[:args.max_records]
    overall = defaultdict(lambda: Counter()); by_lang = defaultdict(lambda: defaultdict(Counter)); by_origin = defaultdict(lambda: defaultdict(Counter))
    type_counts = Counter(); mismatches = Counter(); started = time.time()
    with torch.inference_mode():
        for i, row in enumerate(records, 1):
            logits_many = model.logits_many(row, row["questions"], max_state=args.max_state, max_branch=args.max_branch)
            preds = {}
            for q, logits in zip(row["questions"], logits_many):
                pred = int(torch.argmax(logits).item())
                y = int(q["label"])
                p = bool(pred)
                truth = bool(y)
                qid = q["id"]
                add(overall[qid], truth, p); add(by_lang[row.get("_meta", {}).get("language")][qid], truth, p); add(by_origin[row.get("_meta", {}).get("origin")][qid], truth, p)
                if qid != "is_dlp": type_counts[qid] += 1
                preds[qid] = p
            if "is_dlp" in preds and any(k != "is_dlp" and v for k, v in preds.items()) != preds["is_dlp"]:
                lang = row.get("_meta", {}).get("language") or "unknown"; mismatches[lang] += 1
            if i % args.progress == 0: print(json.dumps({"processed": i, "seconds": round(time.time()-started,1)}), flush=True)
    report = {"checkpoint": args.checkpoint, "data": args.data, "records": len(records), "seconds": round(time.time()-started, 1),
              "overall": {k: metric(v) for k,v in overall.items()},
              "by_language": {lang: {k: metric(v) for k,v in qs.items()} for lang, qs in by_lang.items()},
              "by_origin": {origin: {k: metric(v) for k,v in qs.items()} for origin, qs in by_origin.items()},
              "logical_mismatch": dict(mismatches), "type_question_counts": dict(type_counts)}
    Path(args.out).write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
    print(json.dumps(report, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("--checkpoint", required=True); p.add_argument("--data", required=True); p.add_argument("--out", required=True)
    p.add_argument("--device", default="cuda"); p.add_argument("--dtype", choices=["fp32", "bf16"], default="bf16")
    p.add_argument("--max_records", type=int); p.add_argument("--max_state", type=int, default=12288); p.add_argument("--max_branch", type=int, default=16384); p.add_argument("--progress", type=int, default=250)
    evaluate(p.parse_args())