import torch
import torch.nn as nn

from .config import ARConfig
from .rmsnorm import RMSNorm
from .block import DecoderBlock
from .rope import build_rope_cache


class ARModel(nn.Module):
    """Decoder-only transformer: token IDs in, next-token scores out."""

    def __init__(self, cfg: ARConfig):
        super().__init__()
        self.cfg = cfg

        # token embedding: maps each token ID to a dim-wide vector
        self.embed = nn.Embedding(cfg.vocab_size, cfg.dim)

        # the stack of identical decoder blocks
        self.blocks = nn.ModuleList([
            DecoderBlock(cfg.dim, cfg.n_heads, cfg.n_kv_heads,
                         cfg.head_dim, cfg.hidden_dim, cfg.rms_norm_eps)
            for _ in range(cfg.n_layers)
        ])

        self.final_norm = RMSNorm(cfg.dim, cfg.rms_norm_eps)
        # output head: dim-wide vector -> score for every token in the vocab
        self.lm_head = nn.Linear(cfg.dim, cfg.vocab_size, bias=False)

        # tie weights: embedding table and output head share one matrix
        if cfg.tie_embeddings:
            self.lm_head.weight = self.embed.weight

        # precompute RoPE tables once; reused by every block (not a parameter)
        cos, sin = build_rope_cache(cfg.head_dim, cfg.max_seq_len, cfg.rope_theta)
        self.register_buffer("rope_cos", cos, persistent=False)
        self.register_buffer("rope_sin", sin, persistent=False)

        # initialize all weights small; without this, 24 stacked layers make
        # the output logits explode (std ~32 instead of ~1)
        self.apply(self._init_weights)
        # extra: scale down the two residual-path projections per block by
        # 1/sqrt(2*n_layers) so the signal doesn't compound across depth
        scale = (2 * cfg.n_layers) ** -0.5
        for block in self.blocks:
            torch.nn.init.normal_(block.attn.o_proj.weight, mean=0.0, std=0.02 * scale)
            torch.nn.init.normal_(block.mlp.down_proj.weight, mean=0.0, std=0.02 * scale)

    def _init_weights(self, module):
        # small normal init for all linear and embedding weights
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, token_ids: torch.Tensor, block_mask) -> torch.Tensor:
        B, T = token_ids.shape

        # look up embeddings, then slice the RoPE tables to this sequence length
        x = self.embed(token_ids)
        cos, sin = self.rope_cos[:T], self.rope_sin[:T]

        # pass through every decoder block in turn
        for block in self.blocks:
            x = block(x, cos, sin, block_mask)

        # final clean-up, then project to vocab-sized scores (logits)
        x = self.final_norm(x)
        return self.lm_head(x)



    # def forward(self, token_ids: torch.Tensor) -> torch.Tensor:
    #     B, T = token_ids.shape

    #     # look up embeddings, then slice the RoPE tables to this sequence length
    #     x = self.embed(token_ids)
    #     cos, sin = self.rope_cos[:T], self.rope_sin[:T]

    #     # pass through every decoder block in turn
    #     for block in self.blocks:
    #         x = block(x, cos, sin)

    #     # final clean-up, then project to vocab-sized scores (logits)
    #     x = self.final_norm(x)
    #     return self.lm_head(x)