"""Pre-norm Transformer encoder (nn.TransformerEncoderLayer, norm_first=True) with key padding masks."""
from __future__ import annotations

import torch
import torch.nn as nn


class Encoder(nn.Module):
    def __init__(self, d_model: int, n_layers: int, n_heads: int, d_ff: int, dropout: float = 0.1):
        super().__init__()
        layer = nn.TransformerEncoderLayer(d_model, n_heads, d_ff, dropout, activation="gelu", batch_first=True,
                                           norm_first=True)
        self.layers = nn.TransformerEncoder(layer, n_layers, enable_nested_tensor=False)
        self.norm = nn.LayerNorm(d_model)

    def forward(self, tokens: torch.Tensor, key_padding_mask: torch.Tensor | None = None) -> torch.Tensor:
        # key_padding_mask: True = ignore. Guard against all-masked rows (would produce NaN attention).
        if key_padding_mask is not None:
            all_masked = key_padding_mask.all(dim=1)
            if all_masked.any():
                key_padding_mask = key_padding_mask.clone()
                key_padding_mask[all_masked, 0] = False
        return self.norm(self.layers(tokens, src_key_padding_mask=key_padding_mask))