"""Same CiteNet architecture, trained instead of hand-written.

random  : random init, Adam + class-weighted BCE on aux_train (genuine vs phantom)
handinit: start from the hand-written weights, same training
Scored on the LePhantomCite eval split, and on the planted fakes in the demo brief
(court mismatch / volume-year / state mismatch: kinds the train split's phantoms
never show, since they are all invented series).
"""
import os, sys, copy
import numpy as np, torch, torch.nn.functional as F
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from citenet import net as cn, text
from lepc import load

torch.manual_seed(0); np.random.seed(0)
tr = [r for r in load("aux_train") if r["label"] != "other"]
ev = load("eval")
Xtr = torch.from_numpy(text.encode([r["cite"] for r in tr])); ytr = torch.tensor([r["label"] == "phantom" for r in tr], dtype=torch.float32)
Xev = torch.from_numpy(text.encode([r["cite"] for r in ev])); lev = np.array([r["label"] for r in ev])
brief = open(os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "demo", "sample_brief.txt")).read()
bc = [c for _, _, c in text.extract(brief)]
hand, _ = cn.make()
with torch.no_grad():
    bl = (hand(torch.from_numpy(text.encode(bc)))[0][:, 0] >= 0.5).numpy()   # hand-written labels = planted truth here
Xb = torch.from_numpy(text.encode(bc))


def score(m, tag):
    with torch.no_grad():
        p = m(Xev)[0][:, 0].numpy(); pb = m(Xb)[0][:, 0].numpy()
    f = p >= 0.5
    rec = (f & (lev == "phantom")).sum(), (lev == "phantom").sum()
    ff = (f & (lev == "genuine")).sum(), (lev == "genuine").sum()
    fb = pb >= 0.5
    ood = [c for c, t, y in zip(bc, bl, fb) if t and "th " not in c.split("(")[0]]
    caught = sum(1 for c, t, y in zip(bc, bl, fb) if t and y and "th " not in c.split("(")[0])
    series = sum(1 for c, t, y in zip(bc, bl, fb) if t and y and "th " in c.split("(")[0])
    fpb = sum(1 for t, y in zip(bl, fb) if y and not t)
    print(f"{tag:22s} eval phantoms {rec[0]}/{rec[1]}  eval false flags {ff[0]}/{ff[1]}  | brief: series fakes {series}/2, "
          f"other fakes {caught}/{len(ood)}, real flagged {fpb}/{int((~bl).sum())}")


def train(m, steps=1500, lr=3e-3):
    m.train()
    opt = torch.optim.Adam(m.parameters(), lr=lr)
    pw = torch.tensor((1 - ytr.mean()) / ytr.mean())
    for s in range(steps):
        idx = torch.randint(0, len(Xtr), (256,))
        p = m(Xtr[idx])[0][:, 0].clamp(1e-6, 1 - 1e-6)
        logit = torch.log(p) - torch.log1p(-p)
        loss = F.binary_cross_entropy_with_logits(logit, ytr[idx], pos_weight=pw)
        opt.zero_grad(); loss.backward(); opt.step()
    return m.eval()


score(hand, "hand-written (0 steps)")
for seed in (0, 1, 2):
    torch.manual_seed(seed)
    rnd = cn.CiteNet(cn.build()[0])
    for mod in rnd.modules():
        if isinstance(mod, (torch.nn.Conv2d, torch.nn.Linear)):
            mod.reset_parameters()
    score(train(rnd), f"random init, seed {seed}")
torch.manual_seed(0)
score(train(copy.deepcopy(hand), lr=1e-4), "hand init + fine-tune")