"""Causal decoder backbone: pre-norm blocks (GQA attention + SwiGLU), tied embeddings, single final norm."""
from __future__ import annotations

from typing import Optional

import torch
import torch.nn as nn

from .attention import CausalSelfAttention, LayerKVCache, RMSNorm, SwiGLU
from .config import PrismeConfig


class DecoderBlock(nn.Module):
    def __init__(self, cfg: PrismeConfig):
        super().__init__()
        self.attn_norm = RMSNorm(cfg.d_model, cfg.norm_eps)
        self.attn = CausalSelfAttention(cfg.d_model, cfg.n_heads, cfg.n_kv_heads, cfg.rope_theta)
        self.ffn_norm = RMSNorm(cfg.d_model, cfg.norm_eps)
        self.ffn = SwiGLU(cfg.d_model, cfg.ffn_width)

    def forward(self, x, positions, cache: Optional[LayerKVCache] = None, attn_mask=None):
        x = x + self.attn(self.attn_norm(x), positions, cache, attn_mask)
        x = x + self.ffn(self.ffn_norm(x))
        return x


class Backbone(nn.Module):
    """Returns the final residual stream H (pre final-norm). The LM head is `head()`."""

    def __init__(self, cfg: PrismeConfig):
        super().__init__()
        self.cfg = cfg
        self.embed = nn.Embedding(cfg.vocab_size, cfg.d_model)
        self.layers = nn.ModuleList([DecoderBlock(cfg) for _ in range(cfg.n_layers)])
        self.final_norm = RMSNorm(cfg.d_model, cfg.norm_eps)
        if not cfg.tie_embeddings:
            self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
        self.apply(self._init)
        # scaled init for residual output projections (GPT-2 style)
        for n, p in self.named_parameters():
            if n.endswith("wo.weight") or n.endswith("w_down.weight"):
                nn.init.normal_(p, mean=0.0, std=cfg.init_std / (2 * cfg.n_layers) ** 0.5)

    def _init(self, m):
        if isinstance(m, nn.Linear):
            nn.init.normal_(m.weight, mean=0.0, std=self.cfg.init_std)
        elif isinstance(m, nn.Embedding):
            nn.init.normal_(m.weight, mean=0.0, std=self.cfg.init_std)

    def new_cache(self) -> list[LayerKVCache]:
        return [LayerKVCache() for _ in self.layers]

    def forward(self, tokens: torch.Tensor, positions: Optional[torch.Tensor] = None,
                cache: Optional[list[LayerKVCache]] = None, attn_mask=None) -> torch.Tensor:
        B, T = tokens.shape
        if positions is None:
            start = 0 if cache is None else cache[0].length
            positions = torch.arange(start, start + T, device=tokens.device).unsqueeze(0).expand(B, T)
        x = self.embed(tokens)
        for i, layer in enumerate(self.layers):
            x = layer(x, positions, None if cache is None else cache[i], attn_mask)
        return x  # [B, T, d], pre final-norm residual stream

    def head(self, h: torch.Tensor) -> torch.Tensor:
        """logits = E @ RMSNorm(h) (tied) or lm_head(RMSNorm(h))."""
        hn = self.final_norm(h)
        if self.cfg.tie_embeddings:
            return hn @ self.embed.weight.t()
        return self.lm_head(hn)