import torch
import torch.nn as nn

from .attention import GQA
from .SwiGLU import SwiGLU


class TransformerBlock(nn.Module):
    def __init__(self, hidden_size: int = 1024, intermediate_size: int = 2816,
                 num_heads: int = 16, num_kv_heads: int = 4, max_seq_len: int = 2048,
                 rope_base: float = 10000.0, norm_eps: float = 1e-6,
                 layer_idx: int = 0):
        super().__init__()
        self.attn_norm = nn.RMSNorm(hidden_size, eps=norm_eps)
        self.attn = GQA(
            hidden_size=hidden_size,
            num_heads=num_heads,
            num_kv_heads=num_kv_heads,
            max_seq_len=max_seq_len,
            rope_base=rope_base,
            layer_idx=layer_idx,
        )
        self.mlp_norm = nn.RMSNorm(hidden_size, eps=norm_eps)
        self.mlp = SwiGLU(
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
        )

    def forward(self, x: torch.Tensor, past_key_values=None, past_len: int = 0):
        x = x + self.attn(self.attn_norm(x),
                          past_key_values=past_key_values, past_len=past_len)
        x = x + self.mlp(self.mlp_norm(x))
        return x