"""Teacher-forced loss on held-out pet data: torch fp32 vs the exported 4-bit model.

Usage: quant_eval.py CHECKPOINT EXPORT_DIR DATA.jsonl [N]
Uses EXPORT_DIR/eval/model_q4_full.onnx, the published weights with logits for
every position (the published file only computes the last one).
"""
import json
import random
import sys

import numpy as np
import onnxruntime as ort
import torch
from transformers import AutoTokenizer, GPT2LMHeadModel

import gen_data as G

def fresh(row):
    try:
        G.recheck(row)
        return True
    except ValueError:
        return False


ckpt, exp, data = sys.argv[1:4]
N = int(sys.argv[4]) if len(sys.argv) > 4 else 40
tok = AutoTokenizer.from_pretrained(ckpt)
rows = [json.loads(l) for l in open(data)]
rows = [r for r in rows if fresh(r)]  # as train.py --drop-stale; a no-op on clean data
random.seed(0)
random.shuffle(rows)
rows = rows[:N]  # train.py holds out the first rows of the same shuffle


def encode(row):
    enc = tok(row["text"], return_offsets_mapping=True)
    mark = bytearray(len(row["text"]) + 1)
    for a, b in row["spans"]:
        mark[a:b] = b"\x01" * (b - a)
    ids = enc["input_ids"][:1024]
    lab = [any(mark[s:e]) for s, e in enc["offset_mapping"][:1024]]
    return np.array(ids), np.array(lab)


def nll(logits, ids, lab):
    lp = torch.log_softmax(torch.tensor(logits[:-1], dtype=torch.float32), -1)
    tgt = torch.tensor(ids[1:])
    l = -lp[torch.arange(len(tgt)), tgt]
    m = torch.tensor(lab[1:])
    return l[m].sum().item(), int(m.sum())


model = GPT2LMHeadModel.from_pretrained(ckpt).eval()
results = {"torch fp32": [0, 0]}
so = ort.SessionOptions()
so.add_session_config_entry("session.disable_quant_qdq", "1")  # as web/brain.ts: fold to fp32
sessions = {"q4 (published)": ort.InferenceSession(f"{exp}/eval/model_q4_full.onnx", so, providers=["CPUExecutionProvider"])}
heads = model.config.n_head
for f in sessions:
    results[f] = [0, 0]
for row in rows:
    ids, lab = encode(row)
    with torch.no_grad():
        lg = model(torch.tensor(ids)[None]).logits[0].numpy()
    a, b = nll(lg, ids, lab)
    results["torch fp32"][0] += a; results["torch fp32"][1] += b
    for f, s in sessions.items():
        ins = {i.name: i for i in s.get_inputs()}
        feed = {"input_ids": ids[None].astype(np.int64), "attention_mask": np.ones((1, len(ids)), np.int64)}
        if "position_ids" in ins:
            feed["position_ids"] = np.arange(len(ids), dtype=np.int64)[None]
        for k, v in ins.items():
            if k.startswith("past"):
                feed[k] = np.zeros((1, heads, 0, 64), np.float32)
        a, b = nll(s.run(["logits"], feed)[0][0], ids, lab)
        results[f][0] += a; results[f][1] += b
for k, (a, b) in results.items():
    print(f"{k:22} loss {a / b:.4f}  ppl {np.exp(a / b):.3f}  ({b} tokens)")