import torch
import torch.nn as nn

from .config import IvmeConfig
from .rmsnorm import RMSNorm
from .rope import precompute_rope_freqs
from .attention import CausalSelfAttention
from .feedforward import SwiGLU


class TransformerBlock(nn.Module):
    """One dense transformer layer (Section 3): pre-norm attention + pre-norm SwiGLU,
    with residual connections around each. Identical shape repeated n_layers times --
    no loops, no weight sharing (Section 3.1, distinguishing this from the shelved
    Ivmetron design).
    """

    def __init__(self, cfg: IvmeConfig):
        super().__init__()
        self.attn_norm = RMSNorm(cfg.hidden_dim, eps=cfg.norm_eps)
        self.attn = CausalSelfAttention(cfg.hidden_dim, cfg.n_heads, cfg.dropout)
        self.ffn_norm = RMSNorm(cfg.hidden_dim, eps=cfg.norm_eps)
        self.ffn = SwiGLU(cfg.hidden_dim, cfg.ffn_mult)

    def forward(self, x: torch.Tensor, rope_freqs: torch.Tensor) -> torch.Tensor:
        x = x + self.attn(self.attn_norm(x), rope_freqs)
        x = x + self.ffn(self.ffn_norm(x))
        return x


class IvmeConversateV2(nn.Module):
    """Ivme-Conversate-v2 (Dense) -- the full model described in Section 4.

    ~20M parameters, 10 layers, hidden_dim 384, 6 heads, RoPE, SwiGLU, RMSNorm,
    tied embeddings, 16k vocab, 1024 context. See config.py for the exact spec.
    """

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

        self.tok_embed = nn.Embedding(cfg.vocab_size, cfg.hidden_dim)
        self.blocks = nn.ModuleList([TransformerBlock(cfg) for _ in range(cfg.n_layers)])
        self.final_norm = RMSNorm(cfg.hidden_dim, eps=cfg.norm_eps)

        # Section 4.8: tied embeddings -- output head reuses the input embedding
        # table instead of learning a separate one.
        self.lm_head = nn.Linear(cfg.hidden_dim, cfg.vocab_size, bias=False)
        if cfg.tie_embeddings:
            self.lm_head.weight = self.tok_embed.weight

        rope_freqs = precompute_rope_freqs(cfg.head_dim, cfg.context_len, cfg.rope_theta)
        self.register_buffer("rope_freqs", rope_freqs, persistent=False)

        self.apply(self._init_weights)

    def _init_weights(self, module: nn.Module):
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, idx: torch.Tensor, targets: torch.Tensor | None = None):
        B, T = idx.shape
        assert T <= self.cfg.context_len, (
            f"sequence length {T} exceeds context_len {self.cfg.context_len}"
        )

        x = self.tok_embed(idx)
        for block in self.blocks:
            x = block(x, self.rope_freqs)
        x = self.final_norm(x)
        logits = self.lm_head(x)

        loss = None
        if targets is not None:
            loss = nn.functional.cross_entropy(
                logits.view(-1, logits.size(-1)),
                targets.view(-1),
                ignore_index=-1,
            )
        return logits, loss

    def num_params(self, non_embedding: bool = False) -> int:
        n = sum(p.numel() for p in self.parameters())
        if non_embedding:
            n -= self.tok_embed.weight.numel()
        return n