"""Held-out token-level cross-entropy / perplexity at fixed forced budgets and under the learned controller."""
from __future__ import annotations

import math

import numpy as np
import torch

from prisme.losses import token_nll
from training.train_editor import sample_forced_actions


@torch.no_grad()
def perplexity_by_policy(model, cfg, batches_fn, device, cost, budgets=(0, 1, 2, 4, 8), controller=True, seed=0):
    model.eval()
    res = {}
    for b in budgets:
        rng = np.random.default_rng(seed)
        tot, n = 0.0, 0
        for x, y in batches_fn():
            if b == 0 or not cfg.use_workspace:
                logits = model(x, None)["logits"]
            else:
                acts, _ = sample_forced_actions(cfg, x.shape[0], x.shape[1], [b], [1.0], rng, device)
                logits = model(x, acts)["logits"]
            nll = token_nll(logits, y); v = y != -100
            tot += float(nll[v].sum()); n += int(v.sum())
        res[f"forced_b{b}"] = {"loss": tot / n, "ppl": math.exp(tot / n), "n_tokens": n}
    if controller and cfg.use_workspace:
        tot, n, ed = 0.0, 0, 0.0
        for x, y in batches_fn():
            out = model.forward_policy(x, cfg.r_max, cost)
            nll = token_nll(out["logits"], y); v = y != -100
            tot += float(nll[v].sum()); n += int(v.sum()); ed += float(out["n_edits"][v].float().sum())
        res["controller"] = {"loss": tot / n, "ppl": math.exp(tot / n), "n_tokens": n, "edits_per_token": ed / n}
    return res