"""Evaluation utilities: held-out loss -> bits-per-byte and word perplexity.

`estimate_loss` is used inside the training loop (cheap, fixed #micro-batches).
`evaluate_split` produces a final report dict for a whole split.
"""
from __future__ import annotations

from typing import Dict

import torch

from .config import ByteLMConfig
from .data import ByteDataset
from .metrics import bpb_from_nats, word_ppl_from_bpb


@torch.no_grad()
def estimate_loss(model, dataset: ByteDataset, cfg: ByteLMConfig,
                  device, iters: int, autocast_ctx) -> Dict[str, float]:
    """Mean CE loss (nats/byte) over `iters` random micro-batches -> BPB/ppl."""
    was_training = model.training
    model.eval()
    losses = torch.zeros(iters)
    for i in range(iters):
        batch = dataset.get_batch(device)
        with autocast_ctx():
            _, _, parts = model(batch["x"], batch["y"],
                                batch.get("seg_ids"), batch.get("pos_ids"))
        losses[i] = parts["ce"]
    if was_training:
        model.train()
    ce = losses.mean().item()
    bpb = bpb_from_nats(ce)
    return {"ce": ce, "bpb": bpb}


def attach_word_ppl(report: Dict[str, float], meta: dict) -> Dict[str, float]:
    bpw = meta.get("bytes_per_word")
    if bpw and bpw == bpw:  # not NaN
        report["word_ppl"] = word_ppl_from_bpb(report["bpb"], bpw)
    return report


@torch.no_grad()
def evaluate_split(model, cfg: ByteLMConfig, split: str, device,
                   autocast_ctx, iters: int, meta: dict) -> Dict[str, float]:
    ds = ByteDataset(cfg, split)
    rep = estimate_loss(model, ds, cfg, device, iters, autocast_ctx)
    return attach_word_ppl(rep, meta)