"""Loading, tokenising and batching the synthetic task splits (char-level or BPE tokenizer)."""
from __future__ import annotations

import json
from collections import defaultdict
from pathlib import Path

import numpy as np
import torch

from data.char_tokenizer import CharTokenizer


class AnyTokenizer:
    def __init__(self, spec: str):
        self.kind = "char" if spec == "char" else "bpe"
        if self.kind == "char":
            self.tok = CharTokenizer(); self.eot_id = self.tok.eot_id; self.pad_id = self.tok.pad_id
            self.vocab_size = self.tok.vocab_size
        else:
            from tokenizers import Tokenizer
            self.tok = Tokenizer.from_file(spec); self.eot_id = self.tok.token_to_id("<|endoftext|>")
            self.pad_id = self.tok.token_to_id("<|pad|>"); self.vocab_size = self.tok.get_vocab_size()

    def encode(self, s: str) -> list[int]:
        return self.tok.encode(s) if self.kind == "char" else self.tok.encode(s).ids

    def decode(self, ids) -> str:
        ids = [int(i) for i in ids]
        return self.tok.decode(ids) if self.kind == "char" else self.tok.decode(ids, skip_special_tokens=True)


def load_items(path: str | Path) -> list[dict]:
    return [json.loads(l) for l in open(path)]


def encode_item(tok: AnyTokenizer, it: dict) -> dict:
    p = tok.encode(it["prompt"]); a = tok.encode(it["answer"]) + [tok.eot_id]
    return {"ids": p + a, "prompt_len": len(p), "answer_len": len(a), "item": it}


class SyntheticBatches:
    """Right-padded batches with a loss mask over the answer span (+EOT). Targets elsewhere are -100."""

    def __init__(self, items: list[dict], tok: AnyTokenizer, max_len: int = 256):
        self.enc = [encode_item(tok, it) for it in items]
        self.enc = [e for e in self.enc if len(e["ids"]) <= max_len]
        self.pad = tok.pad_id

    def __len__(self):
        return len(self.enc)

    def collate(self, encs, device, loss_span: str = "answer"):
        L = max(len(e["ids"]) for e in encs)
        x = torch.full((len(encs), L), self.pad, dtype=torch.long)
        y = torch.full((len(encs), L), -100, dtype=torch.long)
        for i, e in enumerate(encs):
            ids = torch.tensor(e["ids"])
            x[i, :len(ids)] = ids
            if loss_span == "answer":
                s = e["prompt_len"] - 1              # position predicting the first answer token
                y[i, s:len(ids) - 1] = ids[s + 1:]
            else:
                y[i, :len(ids) - 1] = ids[1:]
        return x.to(device), y.to(device)

    def random_batches(self, batch_size: int, seed: int, device, loss_span="answer"):
        rng = np.random.default_rng(seed)
        while True:
            idx = rng.integers(0, len(self.enc), size=batch_size)
            yield self.collate([self.enc[i] for i in idx], device, loss_span)

    def sequential_batches(self, batch_size: int, device, loss_span="answer", max_items=None):
        encs = self.subset(max_items)
        for i in range(0, len(encs), batch_size):
            yield self.collate(encs[i:i + batch_size], device, loss_span)

    def subset(self, max_items=None):
        """Stratified deterministic subset: every k-th item so that all families/templates/depths are covered."""
        if not max_items or max_items >= len(self.enc):
            return self.enc
        k = len(self.enc) / max_items
        return [self.enc[int(i * k)] for i in range(max_items)]

    def prompt_buckets(self, batch_size: int, max_items=None):
        """Groups of encoded items with identical prompt length (for cache-based free generation)."""
        encs = self.subset(max_items)
        by_len = defaultdict(list)
        for e in encs:
            by_len[e["prompt_len"]].append(e)
        for L, group in sorted(by_len.items()):
            for i in range(0, len(group), batch_size):
                yield group[i:i + batch_size]


@torch.no_grad()
def exact_match(model, batches: SyntheticBatches, tok: AnyTokenizer, device, max_budget: int, cost=None,
                forced_actions_fn=None, batch_size: int = 64, max_items=None, max_new=12, lambda_cost=None, kappa_risk=None):
    """Free greedy generation of the answer; returns dict with accuracy, per-family accuracy, and per-item records."""
    from prisme.generation import generate
    model.eval()
    recs = []
    for group in batches.prompt_buckets(batch_size, max_items):
        prompt = torch.tensor([e["ids"][:e["prompt_len"]] for e in group], device=device)
        out = generate(model, prompt, max_new, max_budget=max_budget, cost=cost, temperature=0.0, eos_id=tok.eot_id,
                       forced_budget_actions=forced_actions_fn, lambda_cost=lambda_cost, kappa_risk=kappa_risk)
        gen = out["tokens"][:, prompt.shape[1]:]
        ne = out["n_edits"]
        for i, e in enumerate(group):
            ids = gen[i].tolist()
            if tok.eot_id in ids:
                ids = ids[:ids.index(tok.eot_id)]
            pred = tok.decode(ids).strip()
            truth = e["item"]["answer"].strip()
            recs.append({"uid": e["item"]["uid"], "family": e["item"]["family"], "template": e["item"]["template"],
                         "depth": e["item"]["depth"], "pred": pred, "truth": truth, "correct": pred == truth,
                         "edits_per_token": float(ne[i].float().mean()) if ne is not None else 0.0})
    acc = float(np.mean([r["correct"] for r in recs])) if recs else float("nan")
    fam = defaultdict(list)
    for r in recs:
        fam[r["family"]].append(r["correct"])
    return {"accuracy": acc, "n": len(recs), "by_family": {k: float(np.mean(v)) for k, v in fam.items()},
            "edits_per_token": float(np.mean([r["edits_per_token"] for r in recs])) if recs else 0.0, "records": recs}