import torch
import torch.nn as nn

from .rmsnorm import RMSNorm
from .attention import GroupedQueryAttention
from .mlp import SwiGLU


class DecoderBlock(nn.Module):
    """One transformer block: attention sub-layer, then MLP sub-layer.

    Each sub-layer is pre-normed and wrapped in a residual (add-back) connection.
    """

    def __init__(self, dim, n_heads, n_kv_heads, head_dim, hidden_dim, eps=1e-6):
        super().__init__()
        # a norm before each sub-layer
        self.attn_norm = RMSNorm(dim, eps)
        self.mlp_norm = RMSNorm(dim, eps)
        self.attn = GroupedQueryAttention(dim, n_heads, n_kv_heads, head_dim)
        self.mlp = SwiGLU(dim, hidden_dim)

    def forward(self, x, cos, sin, block_mask):
        # attention sub-layer: normalize, attend, add the result back
        x = x + self.attn(self.attn_norm(x), cos, sin, block_mask)
        # mlp sub-layer: normalize, process per-token, add the result back
        x = x + self.mlp(self.mlp_norm(x))
        return x