"""Zero-shot multiple-choice evaluation (lm-evaluation-harness conventions) that fits next to a training job."""
from __future__ import annotations

import json
import random
import re
from pathlib import Path

import torch

EVAL_DIR = Path(__file__).resolve().parent.parent / "data" / "evals"
# chance level, for the dashboard
CHANCE = {"hellaswag": 25.0, "arc_easy": 25.0, "arc_challenge": 25.0, "piqa": 50.0, "sciq": 25.0, "openbookqa": 25.0, "winogrande": 50.0}
METRIC = {"hellaswag": "acc_norm", "arc_easy": "acc", "arc_challenge": "acc_norm", "piqa": "acc_norm", "sciq": "acc_norm", "openbookqa": "acc_norm", "winogrande": "acc"}


def load_tasks(limit: int | None = None, seed: int = 0) -> dict[str, list[dict]]:
    tasks = {}
    for f in sorted(EVAL_DIR.glob("*.json")):
        items = json.loads(f.read_text())
        if limit and len(items) > limit:
            items = random.Random(seed).sample(items, limit)
        tasks[f.stem] = items
    return tasks


@torch.no_grad()
def loglik_batch(model, pairs, device, pad_id=1, max_ctx=480, bs=48):
    """pairs: list of (ctx_ids, cont_ids). Returns (sum logprob, n_tokens) per pair."""
    res = [None] * len(pairs)
    order = sorted(range(len(pairs)), key=lambda i: len(pairs[i][0]) + len(pairs[i][1]))
    w = model.out_weight
    for s in range(0, len(order), bs):
        idxs = order[s : s + bs]
        seqs = []
        for i in idxs:
            c, k = pairs[i]
            c = c[-max_ctx:]
            seqs.append((c + k, len(c), len(k)))
        L = max(len(x[0]) for x in seqs)
        inp = torch.full((len(seqs), L), pad_id, dtype=torch.long)
        for r, (ids, _, _) in enumerate(seqs):
            inp[r, : len(ids)] = torch.tensor(ids)
        inp = inp.to(device)
        with torch.autocast(device.type, torch.bfloat16):
            h = model.hidden(inp)
        rows, cols, tgts = [], [], []
        for r, (ids, nc, nk) in enumerate(seqs):
            for j in range(nk):
                rows.append(r); cols.append(nc + j - 1); tgts.append(ids[nc + j])
        hs = h[torch.tensor(rows, device=device), torch.tensor(cols, device=device)]
        logits = (hs @ w.to(hs.dtype).t()).float()
        lp = torch.log_softmax(logits, -1).gather(1, torch.tensor(tgts, device=device)[:, None]).squeeze(1).cpu()
        p = 0
        for r, i in enumerate(idxs):
            nk = seqs[r][2]
            res[i] = (float(lp[p : p + nk].sum()), nk); p += nk
    return res


def evaluate(model, tok, tasks: dict[str, list[dict]], device, pad_id=1, bs=48) -> dict[str, dict]:
    model.eval()
    out = {}
    for name, items in tasks.items():
        pairs, spans, labels, chs = [], [], [], []
        for it in items:
            start = len(pairs)
            if "alt_ctx" in it:  # WinoGrande partial scoring: same continuation, two candidate contexts
                tail = tok.encode(it["choices"][0]).ids
                pairs += [(tok.encode(it["ctx"]).ids, tail), (tok.encode(it["alt_ctx"]).ids, tail)]
                labels.append(it["wino_label"]); chs.append([it["choices"][0]] * 2)
            else:
                cids = tok.encode(it["ctx"]).ids
                for ch in it["choices"]:
                    pairs.append((cids, tok.encode(ch).ids or [0]))
                labels.append(it["label"]); chs.append(it["choices"])
            spans.append((start, len(pairs) - start))
        ll = loglik_batch(model, pairs, device, pad_id=pad_id, bs=bs)
        acc = accn = 0
        for (st, n), lab, ch in zip(spans, labels, chs):
            raw = [ll[st + k][0] for k in range(n)]
            norm = [raw[k] / max(1, len(ch[k].encode())) for k in range(n)]
            acc += int(max(range(n), key=raw.__getitem__) == lab)
            accn += int(max(range(n), key=norm.__getitem__) == lab)
        out[name] = {"acc": 100 * acc / len(items), "acc_norm": 100 * accn / len(items), "n": len(items)}
        out[name]["score"] = out[name][METRIC.get(name, "acc")]
    model.train()
    return out