import torch.nn as nn

from model.attention import CausalSelfAttention
from model.mlp import FeedForward


class TransformerBlock(nn.Module):

    def __init__(
        self,
        d_model,
        n_heads,
        context_length,
        dropout=0.1
    ):
        super().__init__()

        self.norm1 = nn.LayerNorm(
            d_model
        )

        self.attention = (
            CausalSelfAttention(
                d_model,
                n_heads,
                context_length,
                dropout
            )
        )

        self.norm2 = nn.LayerNorm(
            d_model
        )

        self.feed_forward = (
            FeedForward(
                d_model,
                dropout
            )
        )

    def forward(self, x, kv_cache=None, cache_position=None):

        attn_out = self.attention(
            self.norm1(x),
            kv_cache=kv_cache,
            cache_position=cache_position
        )

        x = x + attn_out

        x = x + self.feed_forward(
            self.norm2(x)
        )

        return x