"""Слои трансформера для Hub-экспорта (относительные импорты для trust_remote_code)."""

import torch
import torch.nn as nn


class FeedForward(nn.Module):
    def __init__(self, emb_dim: int, dropout: float = 0.1):
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(emb_dim, 4 * emb_dim),
            nn.GELU(),
            nn.Linear(4 * emb_dim, emb_dim),
            nn.Dropout(dropout),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.network(x)


class MultiHeadAttention(nn.Module):
    def __init__(
        self,
        emb_dim: int,
        n_heads: int,
        context_length: int,
        dropout: float = 0.1,
        qvk_bias: bool = False,
    ):
        super().__init__()
        if emb_dim % n_heads != 0:
            raise ValueError('emb_dim должно быть кратно n_heads')

        self.model_dim = emb_dim
        self.n_heads = n_heads
        self.head_dim = self.model_dim // n_heads

        self.qkv_proj = nn.Linear(
            self.model_dim, 3 * self.model_dim, bias=qvk_bias,
        )
        self.out_proj = nn.Linear(
            self.model_dim, self.model_dim, bias=qvk_bias,
        )
        self.dropout = nn.Dropout(dropout)
        self.register_buffer(
            'mask',
            torch.tril(torch.ones(1, 1, context_length, context_length)),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        batch_size, n_tokens, emb_dim = x.size()

        qkv = self.qkv_proj(x)
        q, k, v = qkv.split(self.model_dim, dim=-1)

        queries = q.view(
            batch_size, n_tokens, self.n_heads, self.head_dim,
        ).transpose(1, 2)
        keys = k.view(
            batch_size, n_tokens, self.n_heads, self.head_dim,
        ).transpose(1, 2)
        values = v.view(
            batch_size, n_tokens, self.n_heads, self.head_dim,
        ).transpose(1, 2)

        attn_weights = (
            queries @ keys.transpose(-2, -1)
        ) / (self.head_dim ** 0.5)
        attn_weights = attn_weights.masked_fill(
            self.mask[:, :, :n_tokens, :n_tokens] == 0, float('-inf'),
        )
        attn_weights = self.dropout(torch.softmax(attn_weights, dim=-1))
        out_per_head = attn_weights @ values
        out = (
            out_per_head.transpose(1, 2)
            .contiguous()
            .view(batch_size, n_tokens, emb_dim)
        )
        return self.out_proj(out)


class TransformerBlock(nn.Module):
    def __init__(
        self,
        emb_dim: int,
        n_heads: int,
        context_length: int,
        dropout: float = 0.1,
        qvk_bias: bool = False,
    ):
        super().__init__()
        self.attn = MultiHeadAttention(
            emb_dim=emb_dim,
            n_heads=n_heads,
            context_length=context_length,
            dropout=dropout,
            qvk_bias=qvk_bias,
        )
        self.ffn = FeedForward(emb_dim=emb_dim, dropout=dropout)
        self.ln_1 = nn.LayerNorm(emb_dim)
        self.ln_2 = nn.LayerNorm(emb_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.attn(self.ln_1(x))
        x = x + self.ffn(self.ln_2(x))
        return x