#!/usr/bin/env python3
"""HaluEval + FEVER hallucination AUROC for SRT-Adapter v8a.

For each benchmark we score the mean r_hat (BEN reflexivity) over each candidate
text and compute an AUROC where the "positive" class is the hallucination /
refuted claim.

HaluEval QA (pminervini/HaluEval, 'qa'): paired (question + right_answer) vs
(question + hallucinated_answer). AUROC = does mean r_hat distinguish?

FEVER (copenlu/fever_gold_evidence, 'validation'): claim alone, label in
{SUPPORTS, REFUTES, NOT ENOUGH INFO}. AUROC computed on SUPPORTS vs REFUTES.

Usage:
    python scripts/hallucination_auroc.py --bench halueval --max-samples 500
    python scripts/hallucination_auroc.py --bench fever    --max-samples 500
    python scripts/hallucination_auroc.py --bench both     --max-samples 500
"""

from __future__ import annotations

import argparse
import json
import sys
import time
from pathlib import Path

import torch
from sklearn.metrics import roc_auc_score
from transformers import AutoTokenizer

HERE = Path(__file__).resolve().parent
ROOT = HERE.parent
sys.path.insert(0, str((ROOT / "src").resolve()))

from srt.adapter import SRTAdapter  # noqa: E402
from srt.config import (  # noqa: E402
    SRTConfig, MAHConfig, RRMConfig, BENConfig, CommunityConfig, LossConfig,
)


def build_config(p: Path) -> SRTConfig:
    raw = json.loads(p.read_text())
    return SRTConfig(
        backbone_id=raw["backbone_id"],
        backbone_dtype=raw["backbone_dtype"],
        mah_layer_indices=list(raw["mah_layer_indices"]),
        rrm_inject_indices=list(raw["rrm_inject_indices"]),
        community_layer_idx=raw["community_layer_idx"],
        num_mah_layers=raw["num_mah_layers"],
        mah=MAHConfig(**raw["mah"]),
        rrm=RRMConfig(**raw["rrm"]),
        ben=BENConfig(**raw["ben"]),
        community=CommunityConfig(**raw["community"]),
        loss=LossConfig(**{k: v for k, v in raw["loss"].items() if k in LossConfig.__dataclass_fields__}),
    )


def load_model(config_path: Path, adapter_path: Path, device: str):
    config = build_config(config_path)
    model = SRTAdapter(config).to(device)
    if adapter_path.suffix == ".safetensors":
        from safetensors.torch import load_file
        state = load_file(str(adapter_path), device=device)
    else:
        state = torch.load(str(adapter_path), map_location=device)
    model.load_state_dict(state, strict=False)
    model.eval()
    tok = AutoTokenizer.from_pretrained(config.backbone_id)
    return model, tok


@torch.no_grad()
def score_text(model, tok, text: str, device: str, max_len: int = 256) -> dict:
    """Return all candidate hallucination signals for a text."""
    enc = tok(text, return_tensors="pt", truncation=True, max_length=max_len).to(device)
    out = model(input_ids=enc.input_ids, attention_mask=enc.attention_mask)
    if out.ben_output is None:
        return {"r_hat_mean": float("nan"), "p_super_mean": float("nan"),
                "p_super_max": float("nan"), "div2_mean": float("nan")}
    r_hat = out.ben_output.r_hat[0].float()
    p_super = torch.softmax(out.ben_output.regime_logits[0].float(), dim=-1)[:, 1]
    div2 = out.divergences[2][0].norm(dim=-1).float() if len(out.divergences) >= 3 else torch.tensor([0.0])
    return {
        "r_hat_mean": float(r_hat.mean().cpu()),
        "p_super_mean": float(p_super.mean().cpu()),
        "p_super_max": float(p_super.max().cpu()),
        "div2_mean": float(div2.mean().cpu()),
    }


def auroc_for_signal(scores_pos, scores_neg, signal_key: str) -> float:
    pos = [s[signal_key] for s in scores_pos]
    neg = [s[signal_key] for s in scores_neg]
    y = [1] * len(pos) + [0] * len(neg)
    s = pos + neg
    return float(roc_auc_score(y, s))


def run_halueval(model, tok, device, max_samples: int, max_len: int) -> dict:
    from datasets import load_dataset
    ds = load_dataset("pminervini/HaluEval", "qa", split="data")
    n = min(len(ds), max_samples)
    print(f"[halueval] {n} examples")
    pos, neg = [], []  # pos = hallucinated, neg = right
    t0 = time.time()
    for i in range(n):
        ex = ds[i]
        q = ex["question"]
        neg.append(score_text(model, tok, f"Q: {q}\nA: {ex['right_answer']}", device, max_len))
        pos.append(score_text(model, tok, f"Q: {q}\nA: {ex['hallucinated_answer']}", device, max_len))
        if (i + 1) % 100 == 0:
            print(f"  [{i+1}/{n}] elapsed={time.time()-t0:.1f}s")
    aurocs = {k: auroc_for_signal(pos, neg, k)
              for k in ("r_hat_mean", "p_super_mean", "p_super_max", "div2_mean")}
    paired_acc = {k: float(sum(1 for p, n_ in zip(pos, neg) if p[k] > n_[k]) / len(pos))
                  for k in aurocs}
    return {
        "n_pairs": n,
        "auroc_unpaired": aurocs,
        "paired_accuracy_hall_gt_right": paired_acc,
        "wall_time_sec": time.time() - t0,
    }


def run_fever(model, tok, device, max_samples: int, max_len: int) -> dict:
    from datasets import load_dataset
    ds = load_dataset("copenlu/fever_gold_evidence", split="validation")
    keep = [ex for ex in ds if ex["label"] in ("SUPPORTS", "REFUTES")]
    if max_samples > 0 and len(keep) > max_samples:
        sup = [ex for ex in keep if ex["label"] == "SUPPORTS"][:max_samples // 2]
        ref = [ex for ex in keep if ex["label"] == "REFUTES"][:max_samples // 2]
        keep = sup + ref
    print(f"[fever] {len(keep)} examples ({sum(1 for e in keep if e['label']=='REFUTES')} REFUTES)")
    pos, neg = [], []
    t0 = time.time()
    for i, ex in enumerate(keep):
        s = score_text(model, tok, ex["claim"], device, max_len)
        (pos if ex["label"] == "REFUTES" else neg).append(s)
        if (i + 1) % 100 == 0:
            print(f"  [{i+1}/{len(keep)}] elapsed={time.time()-t0:.1f}s")
    aurocs = {k: auroc_for_signal(pos, neg, k)
              for k in ("r_hat_mean", "p_super_mean", "p_super_max", "div2_mean")}
    return {
        "n_samples": len(keep),
        "n_refutes": len(pos),
        "auroc_refutes_vs_supports": aurocs,
        "wall_time_sec": time.time() - t0,
    }


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--config", default=str(ROOT / "config.json"))
    ap.add_argument("--adapter", default=str(ROOT / "adapter.safetensors"))
    ap.add_argument("--bench", choices=("halueval", "fever", "both"), default="both")
    ap.add_argument("--max-samples", type=int, default=500)
    ap.add_argument("--max-len", type=int, default=256)
    args = ap.parse_args()

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model, tok = load_model(Path(args.config), Path(args.adapter), device)
    print(f"loaded adapter on {device}")

    results: dict = {}
    if args.bench in ("halueval", "both"):
        results["halueval_qa"] = run_halueval(model, tok, device, args.max_samples, args.max_len)
        a = results["halueval_qa"]["auroc_unpaired"]
        print(f"[halueval] AUROC: r_hat={a['r_hat_mean']:.3f}  "
              f"p_super_mean={a['p_super_mean']:.3f}  "
              f"p_super_max={a['p_super_max']:.3f}  "
              f"div2={a['div2_mean']:.3f}")
    if args.bench in ("fever", "both"):
        results["fever"] = run_fever(model, tok, device, args.max_samples, args.max_len)
        a = results["fever"]["auroc_refutes_vs_supports"]
        print(f"[fever]    AUROC: r_hat={a['r_hat_mean']:.3f}  "
              f"p_super_mean={a['p_super_mean']:.3f}  "
              f"p_super_max={a['p_super_max']:.3f}  "
              f"div2={a['div2_mean']:.3f}")

    out = ROOT / "benchmarks" / "hallucination_auroc.json"
    out.write_text(json.dumps(results, indent=2))
    print(f"\nwrote {out}")


if __name__ == "__main__":
    main()