"""Exact-match evaluation of a toy model on a Split under greedy decoding, with per-item outputs."""
from __future__ import annotations

from dataclasses import dataclass

import numpy as np
import torch

from wp_ai_0001.datasets.common import Split, Vocab, encode_split
from wp_ai_0001.models.decoding import greedy_decode
from wp_ai_0001.models.toy_transformer import ToyTransformer


@dataclass
class EvalResult:
    split: str
    split_hash: str
    n: int
    accuracy: float  # final-answer exact match
    trace_accuracy: float  # all scratchpad tokens exact match
    correct: np.ndarray  # [N] bool (final answer)
    trace_correct: np.ndarray  # [N] bool
    clusters: list[str]
    predictions: list[list[str]]
    answers: list[list[str]]

    def to_summary(self) -> dict:
        return {"split": self.split, "split_hash": self.split_hash, "n": self.n, "accuracy": self.accuracy, "trace_accuracy": self.trace_accuracy}


@torch.no_grad()
def evaluate(model: ToyTransformer, split: Split, vocab: Vocab, device: torch.device, bs: int = 1024) -> EvalResult:
    enc = encode_split(split, vocab, model.cfg.max_seq_len)
    K = enc["answers"].shape[1]
    preds = []
    for i in range(0, len(enc["ids"]), bs):
        ids = torch.from_numpy(enc["ids"][i : i + bs]).to(device)
        pl = torch.from_numpy(enc["prompt_len"][i : i + bs]).to(device)
        na = torch.from_numpy(enc["n_answer"][i : i + bs]).to(device)
        preds.append(greedy_decode(model, ids, pl, na, K).cpu().numpy())
    pred = np.concatenate(preds)
    na = enc["n_answer"]
    ar = np.arange(len(na))
    final_ok = pred[ar, na - 1] == enc["answers"][ar, na - 1]
    valid = np.arange(K)[None] < na[:, None]
    trace_ok = ((pred == enc["answers"]) | ~valid).all(-1)
    return EvalResult(split.name, split.hash, len(na), float(final_ok.mean()), float(trace_ok.mean()), final_ok, trace_ok,
                      [e.cluster for e in split.examples], [vocab.decode(pred[i, : na[i]]) for i in ar],
                      [list(e.answers) for e in split.examples])


def chance_level(split: Split) -> float:
    """Majority-class rate of the answer distribution (the relevant chance baseline)."""
    return max(split.answer_distribution().values())