"""
Cell Memory: per-layer attention over cell embeddings.

Each cell gets a learned embedding that serves as both key and value.
Forward pass queries over cells, adds result residually.

New params per layer:
  - cell_emb: [max_cells, d_model]
  - q_proj, out_proj: [d_model, d_model]
  For 28 layers, max_cells=256, d_model=1536:
    cell_emb: 28 × 256 × 1536 = 11.0 M
    q_proj + out_proj: 28 × 2 × 1536² = 132 M
    Total: ~143 M new params (14% of R7)

Checkpoint compatibility:
  - Zero-init out_proj → memory contributes 0.0 at start
  - Checkpoint 3200 loads fine; new params learned from scratch
"""
import torch
import torch.nn as nn
import torch.nn.functional as F


class CellMemory(nn.Module):
    def __init__(self, d_model, max_cells=256, n_heads=8):
        super().__init__()
        assert d_model % n_heads == 0
        self.d_model = d_model
        self.max_cells = max_cells
        self.n_heads = n_heads
        self.head_dim = d_model // n_heads

        self.cell_emb = nn.Embedding(max_cells, d_model)
        self.q_proj = nn.Linear(d_model, d_model, bias=False)
        self.out_proj = nn.Linear(d_model, d_model, bias=False)

        # Init: small cell emb, xavier q, ZERO out (identity at start)
        nn.init.normal_(self.cell_emb.weight, std=0.02)
        nn.init.xavier_uniform_(self.q_proj.weight)
        nn.init.zeros_(self.out_proj.weight)

    def forward(self, hidden, slot_ids, slot_mask):
        """
        hidden: [B, T, d]
        slot_ids: [N] long tensor in [0, max_cells)
        slot_mask: [N] 1 for active, 0 for padding
        """
        device = hidden.device
        if slot_ids.device != device:
            slot_ids = slot_ids.to(device)
        if slot_mask.device != device:
            slot_mask = slot_mask.to(device)

        B, T, D = hidden.shape
        N = slot_ids.shape[0]
        H, hd = self.n_heads, self.head_dim

        # Multi-head query
        q = self.q_proj(hidden).view(B, T, H, hd).transpose(1, 2)  # [B,H,T,hd]

        # Cell keys = values
        k = self.cell_emb(slot_ids)              # [N, D]
        k = k.view(N, H, hd).transpose(0, 1)     # [H, N, hd]
        v = k

        # Attention
        attn = torch.matmul(q, k.transpose(-2, -1)) * (hd ** -0.5)  # [B,H,T,N]
        attn = attn.masked_fill(slot_mask.view(1, 1, 1, N) == 0, float("-inf"))
        attn = F.softmax(attn, dim=-1)
        attn = torch.nan_to_num(attn, nan=0.0)

        out = torch.matmul(attn, v)              # [B,H,T,hd]
        out = out.transpose(1, 2).contiguous().view(B, T, D)
        return self.out_proj(out)


if __name__ == "__main__":
    import torch
    cm = CellMemory(d_model=64, max_cells=16, n_heads=4)
    hidden = torch.randn(1, 8, 64)
    slot_ids = torch.arange(4)
    slot_mask = torch.ones(4)
    out = cm(hidden, slot_ids, slot_mask)
    print(f"out shape: {tuple(out.shape)}")
    print(f"all zeros (out_proj=0): {torch.allclose(out, torch.zeros_like(out))}")