"""Value controller: scores every nonterminal action from target-free features.

Per state: a shared state representation u = f(z_t, step, remaining) is computed once (one SwiGLU MLP);
each action a then costs only one w x w projection of its chosen-slot summary s_a (mean of the slots in I(a)),
a nonlinearity and a 4-way linear head. Outputs per action: g_hat (predicted loss gain), c_hat (predicted cost
proxy), log_sigma (uncertainty scale), harm_logit (P[loss increases]).
Action choice: Q(s,a) = g_hat - lambda*c(a) - kappa*sigma_hat, Q(STOP) = 0.
The controller never sees the target token and never executes candidate edits.
"""
from __future__ import annotations

import torch
import torch.nn as nn
import torch.nn.functional as F

from .attention import RMSNorm, SwiGLU
from .config import PrismeConfig
from .workspace import action_table, action_group_size


class ValueController(nn.Module):
    N_CTX = 2  # step/R, remaining/R

    def __init__(self, cfg: PrismeConfig):
        super().__init__()
        self.cfg = cfg
        hidden = cfg.controller_hidden or cfg.ffn_w
        w = cfg.w
        self.z_proj = nn.Linear(cfg.d_model if cfg.ctrl_input == "h" else w, w, bias=False)
        self.ent_proj = nn.Linear(2, w, bias=False) if cfg.ctrl_use_entropy else None
        self.ctx_proj = nn.Linear(self.N_CTX, w, bias=True)
        self.norm = RMSNorm(w, cfg.norm_eps)
        self.mlp = SwiGLU(w, hidden)
        self.s_proj = nn.Linear(w, w, bias=False)
        self.k_embed = nn.Parameter(torch.zeros(2, w))  # group size 1 or 2
        self.heads = nn.Linear(w, 4, bias=True)
        self.register_buffer("actions", action_table(cfg), persistent=False)
        self.register_buffer("group_size", action_group_size(cfg), persistent=False)

    def state_repr(self, z: torch.Tensor, step: torch.Tensor, remaining: torch.Tensor, entropy=None) -> torch.Tensor:
        """z: [B,Tq,w] summary (or [B,Tq,d] normalised residual when ctrl_input == "h"); entropy: [B,Tq] optional."""
        ctx = torch.stack([step.to(z.dtype) / self.cfg.r_max, remaining.to(z.dtype) / self.cfg.r_max], -1)
        u = self.z_proj(z) + self.ctx_proj(ctx)
        if self.ent_proj is not None:
            assert entropy is not None, "ctrl_use_entropy requires the backbone entropy feature"
            e = entropy.to(z.dtype).unsqueeze(-1)
            u = u + self.ent_proj(torch.cat([e, torch.log(e + 1e-4)], -1))
        return u + self.mlp(self.norm(u))                                 # [B,Tq,w]

    def slot_summaries(self, S: torch.Tensor) -> torch.Tensor:
        A = self.cfg.n_actions
        acts = self.actions[:A]
        amask = (acts >= 0).to(S.dtype)
        sel = S[..., acts.clamp(min=0), :]                                # [B,Tq,A,2,w]
        return (sel * amask.view(1, 1, A, 2, 1)).sum(-2) / amask.sum(-1).view(1, 1, A, 1)

    def forward(self, z, S, step, remaining, entropy=None):
        """z: controller input ([B,Tq,w] summary or [B,Tq,d] residual); S: [B,Tq,M,w]; step/remaining: [B,Tq]."""
        A = self.cfg.n_actions
        u = self.state_repr(z, step, remaining, entropy).unsqueeze(2)     # [B,Tq,1,w]
        s = self.s_proj(self.slot_summaries(S))                           # [B,Tq,A,w]
        kidx = (self.group_size[:A] - 1).clamp(min=0)
        h = F.silu(u + s + self.k_embed[kidx].view(1, 1, A, -1))
        out = self.heads(h)
        return {"g_hat": out[..., 0], "c_hat": out[..., 1], "log_sigma": out[..., 2], "harm_logit": out[..., 3]}

    def q_values(self, preds: dict, cost: torch.Tensor, lambda_cost: float | None = None, kappa_risk: float | None = None):
        lam = self.cfg.lambda_cost if lambda_cost is None else lambda_cost
        kap = self.cfg.kappa_risk if kappa_risk is None else kappa_risk
        sigma = preds["log_sigma"].clamp(-7, 3).exp()
        return preds["g_hat"] - lam * cost - kap * sigma

    def select(self, preds: dict, cost: torch.Tensor, feasible: torch.Tensor | None = None, **kw) -> torch.Tensor:
        """Greedy action with STOP (index A) when max Q <= 0. feasible: optional [B,Tq,A] bool."""
        q = self.q_values(preds, cost, **kw)
        if feasible is not None:
            q = q.masked_fill(~feasible, float("-inf"))
        best_q, best_a = q.max(-1)
        return torch.where(best_q > 0, best_a, torch.full_like(best_a, self.cfg.stop_action))