"""Reference routing policies that share PRISME's editor (actions chosen without the learned value controller).
All are target-free and deployable; they differ only in how the action is picked at each step.
  random      uniform over the 24 nonterminal actions, fixed budget b
  uniform     round-robin single-slot schedule SINGLE_0, SINGLE_1, ... (deterministic, fixed budget b)
  fixed_pair  PAIR_0_1 repeated (fixed budget b)
  entropy     continue (random action) while the current predictive entropy exceeds a threshold, else STOP
  norm        continue while the last write moved the slot by more than a threshold (RMS), else STOP
"""
from __future__ import annotations

import torch

from prisme.model import PrismeModel


@torch.no_grad()
def rollout(model: PrismeModel, tokens, policy: str, budget: int, threshold: float = 0.0, seed: int = 0):
    cfg = model.cfg
    ctx = model.encode(tokens)
    S = ctx.S0
    B, T = tokens.shape
    g = torch.Generator(device="cpu").manual_seed(seed)
    stopped = torch.zeros(B, T, dtype=torch.bool, device=tokens.device)
    n_edits = torch.zeros(B, T, dtype=torch.long, device=tokens.device)
    trace = torch.full((B, T, budget), cfg.stop_action, dtype=torch.long, device=tokens.device)
    last_rms = torch.full((B, T), float("inf"), device=tokens.device)
    for r in range(budget):
        if policy == "random":
            a = torch.randint(0, cfg.n_actions, (B, T), generator=g).to(tokens.device)
        elif policy == "uniform":
            a = torch.full((B, T), r % cfg.n_slots, device=tokens.device)
        elif policy == "fixed_pair":
            a = torch.full((B, T), cfg.n_slots + (r % cfg.n_pairs), device=tokens.device)
        elif policy == "entropy":
            logits, _ = model.logits_from_state(S, ctx, n_edits > 0)
            p = torch.softmax(logits.float(), -1)
            ent = -(p * torch.log(p.clamp_min(1e-12))).sum(-1)
            a = torch.randint(0, cfg.n_actions, (B, T), generator=g).to(tokens.device)
            a = torch.where(ent > threshold, a, torch.full_like(a, cfg.stop_action))
        elif policy == "norm":
            a = torch.randint(0, cfg.n_actions, (B, T), generator=g).to(tokens.device)
            a = torch.where(last_rms > threshold, a, torch.full_like(a, cfg.stop_action))
        else:
            raise ValueError(policy)
        a = torch.where(stopped, torch.full_like(a, cfg.stop_action), a)
        stopped |= a == cfg.stop_action
        trace[:, :, r] = a
        S_new, info = model.apply_action(S, ctx, a, r, budget - r - 1)
        moved = (info["eta"].squeeze(-1) * info["delta_rms"] * info["scale"])  # [B,T,k] realised rms step per lane
        last_rms = torch.where(a != cfg.stop_action, moved.max(-1).values, last_rms)
        S = S_new
        n_edits += (a != cfg.stop_action).long()
    logits, _ = model.logits_from_state(S, ctx, n_edits > 0)
    return {"logits": logits, "actions": trace, "n_edits": n_edits}