"""Action catalog and ephemeral workspace helpers.

Action catalog (M = 16 slots):
    SINGLE_j  (j = 0..M-1)        -> edit one slot            : actions 0 .. M-1
    PAIR_{2i,2i+1} (i = 0..M/2-1) -> edit two fixed slots     : actions M .. M+M/2-1
    STOP                          -> no edit                  : action M+M/2
Exactly M + M/2 nonterminal actions + STOP (= 25 for M = 16). Arbitrary pairs are NOT available.
"""
from __future__ import annotations

import torch

from .config import PrismeConfig


def action_table(cfg: PrismeConfig, device=None) -> torch.Tensor:
    """[n_actions + 1, 2] slot indices per action; -1 marks an unused query lane. Last row = STOP."""
    M = cfg.n_slots
    rows = [[j, -1] for j in range(M)]
    if getattr(cfg, "pair_scheme", "adjacent") == "offset":
        rows += [[i, i + M // 2] for i in range(M // 2)]
    else:
        rows += [[2 * i, 2 * i + 1] for i in range(M // 2)]
    rows += [[-1, -1]]
    return torch.tensor(rows, dtype=torch.long, device=device)


def action_names(cfg: PrismeConfig) -> list[str]:
    tbl = action_table(cfg).tolist()
    return [f"SINGLE_{a}" if b < 0 else f"PAIR_{a}_{b}" for a, b in tbl[:-1]] + ["STOP"]


def action_group_size(cfg: PrismeConfig, device=None) -> torch.Tensor:
    """k(a): number of edited slots per action (0 for STOP)."""
    M = cfg.n_slots
    return torch.tensor([1] * M + [2] * (M // 2) + [0], dtype=torch.long, device=device)


def init_workspace(slot_embed: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
    """S0[j] = e[j] + z_t.  slot_embed: [M, w]; z: [..., w] -> [..., M, w]."""
    return slot_embed + z.unsqueeze(-2)


def rms(x: torch.Tensor) -> torch.Tensor:
    return x.float().pow(2).mean(-1).sqrt()