"""Weight-tied RNN: one 2-layer cell applied K times, ACT learned halting, GRU digit decoder.

Spec pseudocode (design/experiment-spec.md):
    state = embed(input_number)
    for step in range(max_steps):
        state = step_module(state)      # same weights every iteration
        if halt_condition(state): break
    output = project(state)

Initial state: masked mean-pool of (digit embedding + sinusoidal positional encoding)
passed through a small MLP, so digit ORDER reaches the tied cell.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F

from src.config import Config
from src.model_api import PrimeModel


class TiedRNN(PrimeModel):
    def __init__(self, cfg: Config):
        super().__init__()
        self.cfg = cfg
        d = cfg.d_model
        self.embed = nn.Embedding(cfg.vocab + 1, d)      # +1 row = pad
        self.in_proj = nn.Sequential(nn.Linear(d, d), nn.GELU(), nn.Linear(d, d))
        self.ln1 = nn.LayerNorm(d)
        self.cell_ln = nn.LayerNorm(d)
        self.cell_w1 = nn.Linear(d, d)
        self.cell_w2 = nn.Linear(d, d)
        self.halt_head = nn.Linear(d, 1)
        self.decoder = nn.GRUCell(d, d)
        self.out_head = nn.Linear(d, cfg.vocab)

    @staticmethod
    def _sinusoidal(T: int, d: int) -> torch.Tensor:
        pe = torch.zeros(T, d)
        pos = torch.arange(T).float().unsqueeze(1)
        i = torch.arange(d).float().unsqueeze(0)
        pe[:, 0::2] = torch.sin(pos / 10000 ** (2 * i[:, 0::2] / d))
        pe[:, 1::2] = torch.cos(pos / 10000 ** (2 * i[:, 1::2] / d))
        return pe

    def _initial_state(self, x: torch.Tensor) -> torch.Tensor:
        B, T = x.shape
        mask = (x != self.cfg.pad_id).float().unsqueeze(-1)          # (B,T,1)
        pos = self._sinusoidal(T, self.cfg.d_model).to(x.device)     # (T,d)
        e = self.embed(x) + pos.unsqueeze(0)                          # (B,T,d)
        h = (e * mask).sum(1) / mask.sum(1).clamp(min=1)             # (B,d)
        return self.in_proj(h)

    def _run_cell(self, h0: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        """Tied cell x K steps. Returns (final_state (B,d), mean_steps (B,))."""
        cfg = self.cfg
        B = h0.shape[0]
        device = h0.device
        if not cfg.halting:
            h = h0
            for _ in range(cfg.max_steps):
                h = h + self.cell_w2(F.gelu(self.cell_w1(self.cell_ln(h))))
            steps = h0.new_full((B,), float(cfg.max_steps))
            return h, steps
        # ACT: run all K steps, accumulate weighted average (K=20 -> no early break needed)
        h_list, p_list = [], []
        h = h0
        for t in range(cfg.max_steps):
            h = h + self.cell_w2(F.gelu(self.cell_w1(self.cell_ln(h))))
            p = torch.sigmoid(self.halt_head(self.ln1(h))).squeeze(-1)   # (B,)
            if t < cfg.min_steps:
                p = p * 0.0
            h_list.append(h)
            p_list.append(p)
        final = torch.zeros_like(h0)
        steps = torch.zeros(B, device=device)
        remaining = torch.ones(B, device=device)
        for t in range(cfg.max_steps):
            p = p_list[t]
            w = remaining * p
            final = final + w.unsqueeze(-1) * h_list[t]
            steps = steps + (t + 1) * w
            remaining = remaining * (1 - p)
        final = final + remaining.unsqueeze(-1) * h_list[-1]
        steps = steps + remaining * cfg.max_steps
        return final, steps

    def forward(self, x: torch.Tensor, y_in: torch.Tensor) -> dict:
        cfg = self.cfg
        h = self._initial_state(x)
        h, steps = self._run_cell(h)
        # autoregressive digit decoder, teacher-forced during training
        e = self.embed(y_in)                       # (B,T_out,d)
        outs = []
        for t in range(y_in.shape[1]):
            h = self.decoder(e[:, t], h)
            outs.append(self.out_head(h))
        logits = torch.stack(outs, dim=1)          # (B,T_out,vocab)
        return {"logits": logits, "halt_steps": steps}