"""Score a predictions file against RAGTruth response-level labels.

Each line: {"id", "task", "supported": bool (gold), "noul": P(supported)}.
A response is flagged as hallucinated when noul < 1 - THRESHOLD_HALLUCINATED.

    python evaluate.py test_predictions.jsonl            # this model, threshold 0.44 on P(hallucinated)
    python evaluate.py test_predictions_untrained.jsonl 0.5
"""
import json
import sys

import numpy as np


def auroc(score, y):
    r = np.argsort(np.argsort(score)) + 1.0
    return (r[y].sum() - y.sum() * (y.sum() + 1) / 2) / (y.sum() * (~y).sum())


def prf(ph, h, t):
    pred = ph >= t
    tp = (pred & h).sum()
    p, r = tp / max(1, pred.sum()), tp / max(1, h.sum())
    return p, r, 2 * p * r / max(1e-9, p + r)


def decided_at(ph, h, t, accuracy):
    """Largest share of responses, most confident first, whose verdicts reach `accuracy`."""
    conf = np.abs(ph - t)
    ok = (ph >= t) == h
    acc = np.cumsum(ok[np.argsort(-conf)]) / np.arange(1, len(ok) + 1)
    good = np.where(acc >= accuracy)[0]
    return (good.max() + 1) / len(ok) if len(good) else 0.0


def ece(ph, h, bins=10):
    idx = np.minimum((ph * bins).astype(int), bins - 1)
    return sum((idx == b).mean() * abs(ph[idx == b].mean() - h[idx == b].mean()) for b in range(bins) if (idx == b).any())


rows = [json.loads(line) for line in open(sys.argv[1])]
t = float(sys.argv[2]) if len(sys.argv) > 2 else 0.44
ph = np.array([1 - r["noul"] for r in rows])
h = np.array([not r["supported"] for r in rows])
p, r, f = prf(ph, h, t)
print(f"n={len(rows)}  threshold P(hallucinated) >= {t}")
print(f"overall    P {p * 100:.1f}  R {r * 100:.1f}  F1 {f * 100:.1f}  AUROC {auroc(ph, h):.3f}  "
      f"accuracy {((ph >= t) == h).mean() * 100:.1f}  ECE {ece(ph, h):.3f}")
print(f"decided at 95% accuracy {decided_at(ph, h, 0.5, .95) * 100:.0f}%   at 99% {decided_at(ph, h, 0.5, .99) * 100:.0f}%")
for task in ("QA", "Data2txt", "Summary"):
    m = np.array([x["task"] == task for x in rows])
    p, r, f = prf(ph[m], h[m], t)
    print(f"{task:<10} P {p * 100:.1f}  R {r * 100:.1f}  F1 {f * 100:.1f}  AUROC {auroc(ph[m], h[m]):.3f}")