"""Ryskina & Knight (2021) next-term eval for byte-level spark-gpt OEIS checkpoints.

Each item is a short sequence. The prompt is every term but the last, written
"t1,t2,...,tk," with a trailing comma. The answer is the last term.

Two scores per item:
  exact   greedy decode until the next comma or EOS; correct iff the text equals
          str(answer). This matches the NextTerm model-card protocol.
  logprob log P(answer is the next term) in nats: the sum of the log-probs of the
          answer's bytes, plus log(P(",") + P(EOS)) at the position after it. A
          continuous score, much less noisy than 57 binary outcomes.

Encoding is BOS + UTF-8 bytes, exactly as spark-gpt trains. The script checks
that the exported HF tokenizer agrees.

Usage: python eval_ryskina.py <hf_dir> <ryskina.jsonl> <out_prefix>
Writes <out_prefix>.items.jsonl and <out_prefix>.summary.json.
"""
import json
import math
import statistics
import sys

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

MAX_NEW_TOKENS = 40  # longest answer is far shorter; a runaway decode just scores wrong


def main():
    hf_dir, data_path, out_prefix = sys.argv[1:4]
    tok = AutoTokenizer.from_pretrained(hf_dir)
    model = AutoModelForCausalLM.from_pretrained(hf_dir, dtype=torch.float32).cuda().eval()
    bos, eos = tok.bos_token_id, tok.eos_token_id
    comma = ord(",")

    def encode(text):
        ids = [bos] + list(text.encode("utf-8"))
        assert tok(text).input_ids == ids, f"tokenizer disagrees on {text!r}"
        return ids

    items = [json.loads(line)["seq"] for line in open(data_path)]
    results = []
    with torch.no_grad():
        for idx, seq in enumerate(items):
            prompt_text = ",".join(str(x) for x in seq[:-1]) + ","
            answer = str(seq[-1])
            prompt = encode(prompt_text)

            # exact match: greedy until comma / EOS
            ids = torch.tensor([prompt], device="cuda")
            generated = []
            for _ in range(MAX_NEW_TOKENS):
                nxt = int(model(ids).logits[0, -1].argmax())
                if nxt in (comma, eos):
                    break
                generated.append(nxt)
                ids = torch.cat([ids, torch.tensor([[nxt]], device="cuda")], dim=1)
            prediction = bytes(b for b in generated if b < 256).decode("utf-8", "replace")

            # log P(answer is the next term)
            ans_ids = list(answer.encode("utf-8"))
            full = torch.tensor([prompt + ans_ids], device="cuda")
            logp = torch.log_softmax(model(full).logits[0].float(), dim=-1)
            start = len(prompt) - 1  # logits at position i predict token i + 1
            lp_answer = sum(float(logp[start + j, t]) for j, t in enumerate(ans_ids))
            end = logp[start + len(ans_ids)]
            lp_stop = float(torch.logaddexp(end[comma], end[eos]))
            lp = lp_answer + lp_stop

            results.append({
                "idx": idx,
                "prompt": prompt_text,
                "answer": answer,
                "prediction": prediction,
                "correct": prediction == answer,
                "logprob": lp,
                "logprob_answer_bytes": lp_answer,
                "logprob_stop": lp_stop,
            })

    lps = [r["logprob"] for r in results]
    n_correct = sum(r["correct"] for r in results)
    summary = {
        "checkpoint": hf_dir,
        "items": len(results),
        "exact_correct": n_correct,
        "exact_acc": n_correct / len(results),
        "mean_logprob": statistics.fmean(lps),
        "median_logprob": statistics.median(lps),
        "mean_answer_prob": statistics.fmean(math.exp(x) for x in lps),
    }
    with open(out_prefix + ".items.jsonl", "w") as f:
        for r in results:
            f.write(json.dumps(r) + "\n")
    with open(out_prefix + ".summary.json", "w") as f:
        json.dump(summary, f, indent=2)
    print(json.dumps(summary, indent=2))


if __name__ == "__main__":
    main()