"""Score the hand-written CiteNet at citation level on a LePhantomCite split."""
import argparse, collections, os, sys
import numpy as np, torch
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from citenet import net as cn, text
from lepc import load

ap = argparse.ArgumentParser()
ap.add_argument("split", default="aux_train", nargs="?")
ap.add_argument("--show", type=int, default=25)
a = ap.parse_args()

model, meta = cn.make()
rows = load(a.split)
x = torch.from_numpy(text.encode([r["cite"] for r in rows]))
with torch.no_grad():
    p, ev, _ = model(x)
p = p[:, 0].numpy(); ev = ev.numpy()
lab = np.array([r["label"] for r in rows])
flag = p >= 0.5
tot, nz = cn.param_counts(model)
print(f"params {tot:,} (nonzero {nz:,})")
print(collections.Counter(lab))
tp = int((flag & (lab == "phantom")).sum()); fn = int((~flag & (lab == "phantom")).sum())
fp = int((flag & (lab == "genuine")).sum()); tn = int((~flag & (lab == "genuine")).sum())
oth = int((flag & (lab == "other")).sum())
print(f"phantom recall {tp}/{tp+fn} = {tp/max(1,tp+fn):.3f}")
print(f"genuine false-flag {fp}/{fp+tn} = {fp/max(1,fp+tn):.4f}")
print(f"precision {tp/max(1,tp+fp):.3f}")
print(f"other-hallucination cites flagged {oth}/{int((lab=='other').sum())}")
C = meta["concepts"]
cls = [i for i, n in enumerate(C) if n.startswith("class:")]
with torch.no_grad():
    conc = model(x)[2].numpy()
multi = np.where(conc[:, cls].sum(1) > 1.5)[0]
print(f"cites where >1 reporter class fired: {len(multi)}", [rows[i]["cite"] for i in multi[:5]])
names = meta["evidence"]
def why(i):
    return ", ".join(f"{n}={ev[i,j]:.2f}" for j, n in enumerate(names) if ev[i, j] > 0.01)
print("\n-- missed phantoms"); [print(f"  {p[i]:.3f} {rows[i]['cite']!r} -> {text.normalize(rows[i]['cite'])!r} [{why(i)}]") for i in np.where(~flag & (lab=='phantom'))[0][:a.show]]
print("\n-- false flags"); [print(f"  {p[i]:.3f} {rows[i]['cite']!r} -> {text.normalize(rows[i]['cite'])!r} [{why(i)}]") for i in np.where(flag & (lab=='genuine'))[0][:a.show]]
amber = np.where((p > 0.1) & (p < 0.5) & (lab == "genuine"))[0]
print(f"\n-- genuine but amber (0.1-0.5): {len(amber)}")
for i in amber[:a.show]: print(f"  {p[i]:.3f} {text.normalize(rows[i]['cite'])!r} [{why(i)}]")