"""Evaluate one arm (or Gemma 3 270M as the reference) and save EVERY per-question result, so error bars can be
computed afterwards by bars.py without re-running anything.

  python eval.py --arm B --ckpt ckpt/B/step_07820.pt [--limit N]
  python eval.py --gemma [--limit N]                      # google/gemma-3-270m, the shape our body copies

Benchmarks (EleutherAI lm-evaluation-harness, standard zero-shot settings): lambada_openai (acc), hellaswag (acc_norm),
arc_easy (acc_norm), piqa (acc_norm). Our arms are wrapped as an lm-eval model: every request is scored by the
log-probability of the continuation given the context, exactly as the harness does for Hugging Face models. Documents
in our training data are separated by <eos> (no <bos>), so every context starts with <eos> for our arms; Gemma
3 270M was trained with <bos>, so its contexts start with <bos>.
Validation loss: mean next-token loss per 1,024-token sequence of the held-out split (data/val.npy), one number per
sequence, so its error bar is over sequences.

Output: reports/eval_<name>.json  = {summary, per_question: {task: [0/1 ...]}, val_loss_per_seq: [...], meta}
"""
import argparse, json, math, os, pathlib, time
import numpy as np
import torch
import torch.nn.functional as F

TASKS = {"lambada_openai": "acc", "hellaswag": "acc_norm", "arc_easy": "acc_norm", "piqa": "acc_norm"}
SEQ = 1024


def load_model(a):
    from transformers import AutoTokenizer
    tok = AutoTokenizer.from_pretrained("google/gemma-3-270m")
    if a.gemma:
        from transformers import AutoModelForCausalLM
        g = AutoModelForCausalLM.from_pretrained("google/gemma-3-270m", torch_dtype=torch.float32).cuda().eval()
        def logits_fn(ids):
            return g(input_ids=ids, use_cache=False).logits
        return "gemma3_270m", tok, logits_fn
    import train
    counts = np.load(pathlib.Path(a.data) / "unigram_counts.npy")
    m = train.Model(a.arm, counts).cuda().eval()
    st = torch.load(a.ckpt, map_location="cuda", weights_only=False)
    m.load_state_dict(st["model"])
    def logits_fn(ids):
        return m.logits(m.hidden(ids))
    return f"arm{a.arm}_step{st.get('step', 0)}", tok, logits_fn


def make_lm(tok, logits_fn, bs, sep):
    from lm_eval.api.model import LM

    class Wrapped(LM):
        def __init__(self):
            super().__init__()

        @torch.no_grad()
        def _score(self, ctx_ids, cont_ids):
            ids = (ctx_ids + cont_ids)[-SEQ:]
            n = len(cont_ids)
            x = torch.tensor([ids], device="cuda")
            with torch.autocast("cuda", dtype=torch.bfloat16):
                lg = logits_fn(x[:, :-1]).float()
            lp = F.log_softmax(lg[0, -n:], -1)
            tgt = x[0, -n:]
            ll = float(lp.gather(-1, tgt[:, None]).sum())
            greedy = bool((lp.argmax(-1) == tgt).all())
            return ll, greedy

        def loglikelihood(self, requests):
            out = []
            for r in requests:
                ctx, cont = r.args
                ctx_ids = [sep] + tok(ctx, add_special_tokens=False).input_ids
                whole = tok(ctx + cont, add_special_tokens=False).input_ids
                cont_ids = whole[len(ctx_ids) - 1:] if whole[:len(ctx_ids) - 1] == ctx_ids[1:] else tok(cont, add_special_tokens=False).input_ids
                out.append(self._score(ctx_ids, cont_ids))
            return out

        def loglikelihood_rolling(self, requests):
            raise NotImplementedError

        def generate_until(self, requests):
            raise NotImplementedError

    return Wrapped()


@torch.no_grad()
def val_losses(logits_fn, val, n_seq):
    x = torch.from_numpy(val[:n_seq * (SEQ + 1)].astype(np.int64).reshape(n_seq, SEQ + 1))
    out = []
    for i in range(n_seq):
        ids = x[i:i + 1].cuda()
        with torch.autocast("cuda", dtype=torch.bfloat16):
            lg = logits_fn(ids[:, :-1]).float()
        out.append(float(F.cross_entropy(lg[0], ids[0, 1:])))
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--arm", choices=["A", "B", "C", "D"])
    ap.add_argument("--ckpt")
    ap.add_argument("--gemma", action="store_true")
    ap.add_argument("--data", default="data")
    ap.add_argument("--limit", type=int, default=0, help="questions per task (0 = all); for smoke tests only")
    ap.add_argument("--val-seqs", type=int, default=512)
    ap.add_argument("--name", default="")
    a = ap.parse_args()
    t0 = time.time()
    name, tok, logits_fn = load_model(a)
    name = a.name or name
    import lm_eval
    res = lm_eval.simple_evaluate(model=make_lm(tok, logits_fn, 1, tok.bos_token_id if a.gemma else tok.eos_token_id), tasks=list(TASKS), num_fewshot=0,
                                  limit=a.limit or None, log_samples=True, bootstrap_iters=0)
    per_q, summary = {}, {}
    for task, metric in TASKS.items():
        samples = sorted(res["samples"][task], key=lambda s: s["doc_id"])
        per_q[task] = [float(s[metric]) for s in samples]
        summary[task] = {"metric": metric, "n": len(per_q[task]), "score": float(np.mean(per_q[task]))}
    val = np.load(pathlib.Path(a.data) / "val.npy", mmap_mode="r")
    vl = val_losses(logits_fn, val, a.val_seqs)
    summary["val_loss"] = {"metric": "nats/token", "n": len(vl), "score": float(np.mean(vl))}
    out = {"name": name, "summary": summary, "per_question": per_q, "val_loss_per_seq": vl,
           "meta": {"ckpt": a.ckpt, "gemma": a.gemma, "limit": a.limit, "seconds": round(time.time() - t0),
                    "lm_eval": lm_eval.__version__, "date": time.strftime("%Y-%m-%d %H:%M")}}
    pathlib.Path("reports").mkdir(exist_ok=True)
    p = pathlib.Path("reports") / f"eval_{name}.json"
    json.dump(out, open(p, "w"))
    print(json.dumps(summary, indent=1))
    print("saved", p)


if __name__ == "__main__":
    main()