"""Global latent-recurrence adapter (Huginn/Ouro/ELR-*inspired*, simplified): the whole top representation is
refined r times by ONE shared block operating at workspace width w, with a zero-initialised residual back to d.

Same frozen backbone, same forced mixed-budget training recipe and same budget semantics as PRISME; the only
difference is that every "action" rewrites the entire latent (global, dense) instead of 1-2 of 16 slots, and
there is no slot selection to learn. Disclosed simplifications: no prelude/coda retraining, no truncated
backprop through time, adapter-level recurrence on a frozen backbone rather than from-scratch recurrent pretraining.
Parameter count is matched to the PRISME workspace (ffn width chosen in the experiment config)."""
from __future__ import annotations

import torch
import torch.nn as nn

from prisme.attention import RMSNorm, SwiGLU
from prisme.backbone import Backbone
from prisme.config import PrismeConfig
import torch.nn.functional as F


class RecurrentBlock(nn.Module):
    def __init__(self, w: int, heads: int, ffn: int, eps: float):
        super().__init__()
        self.h, self.hd = heads, w // heads
        self.n1 = RMSNorm(w, eps); self.q = nn.Linear(w, w, bias=False); self.k = nn.Linear(w, w, bias=False)
        self.v = nn.Linear(w, w, bias=False); self.o = nn.Linear(w, w, bias=False)
        self.n2 = RMSNorm(w, eps); self.ffn = SwiGLU(w, ffn)

    def forward(self, x, ctx_mask):
        B, T, w = x.shape
        xn = self.n1(x)
        q = self.q(xn).view(B, T, self.h, self.hd).transpose(1, 2)
        k = self.k(xn).view(B, T, self.h, self.hd).transpose(1, 2)
        v = self.v(xn).view(B, T, self.h, self.hd).transpose(1, 2)
        a = F.scaled_dot_product_attention(q, k, v, attn_mask=ctx_mask)
        x = x + self.o(a.transpose(1, 2).reshape(B, T, w))
        return x + self.ffn(self.n2(x))


class GlobalRecurrenceModel(nn.Module):
    """Interface-compatible with PrismeModel.forward(tokens, actions): any nonterminal action = one global
    refinement step; STOP = no step. n_edits semantics identical."""

    def __init__(self, cfg: PrismeConfig, ffn_w: int | None = None, heads: int | None = None):
        super().__init__()
        self.cfg = cfg
        self.backbone = Backbone(cfg)
        w = cfg.w
        self.bank_norm = RMSNorm(cfg.d_model, cfg.norm_eps)
        self.down = nn.Linear(cfg.d_model, w, bias=False)
        self.block = RecurrentBlock(w, heads or cfg.n_editor_heads, ffn_w or cfg.ffn_w, cfg.norm_eps)
        self.ctx_proj = nn.Linear(2, w, bias=True)
        self.up = nn.Linear(w, cfg.d_model, bias=False)
        nn.init.zeros_(self.ctx_proj.weight); nn.init.zeros_(self.ctx_proj.bias)
        if cfg.readout_zero_init:
            nn.init.zeros_(self.up.weight)
        else:
            nn.init.normal_(self.up.weight, std=cfg.readout_init_std)
        self.controller = nn.Identity()  # placeholder so parameter-group helpers work

    def forward(self, tokens, actions=None, return_state=False):
        cfg = self.cfg
        H = self.backbone(tokens)
        if actions is None or actions.shape[-1] == 0:
            return {"logits": self.backbone.head(H), "H": H, "n_edits": None}
        B, T, R = actions.shape
        x0 = self.down(self.bank_norm(H))
        x = x0
        mask = torch.ones(T, T, dtype=torch.bool, device=H.device).tril().view(1, 1, T, T)
        stopped = torch.zeros(B, T, dtype=torch.bool, device=H.device)
        n_edits = torch.zeros(B, T, dtype=torch.long, device=H.device)
        for r in range(R):
            a = torch.where(stopped, torch.full_like(actions[:, :, r], cfg.stop_action), actions[:, :, r])
            act = (a != cfg.stop_action)
            stopped |= ~act
            ctx = torch.stack([torch.full((B, T), r / cfg.r_max, device=H.device), torch.full((B, T), (R - r - 1) / cfg.r_max, device=H.device)], -1)
            xn = self.block(x + self.ctx_proj(ctx), mask)
            x = torch.where(act.unsqueeze(-1), xn, x)     # positions that stopped keep their state
            n_edits += act.long()
        delta = self.up(x - x0) * self.cfg.readout_scale
        edited = n_edits > 0
        delta = torch.where(edited.unsqueeze(-1), delta, torch.zeros_like(delta))
        return {"logits": self.backbone.head(H + delta), "H": H, "n_edits": n_edits, "delta_out": delta}