"""Position-wise feed-forward sub-layer used in transformer blocks."""

from __future__ import annotations

import torch
import torch.nn as nn
from torch import Tensor

__all__ = ["FeedForwardLayer"]


class FeedForwardLayer(nn.Module):
    """
    Position-wise feed-forward layer used in Transformer blocks.

    Architecture:
        fc1: Linear(d_model -> d_ff)
        activation: ReLU
        dropout: nn.Dropout(dropout_rate)
        fc2: Linear(d_ff -> d_model)

    Args:
        d_model (int): Dimensionality of model embeddings.
        d_ff (int): Hidden dimensionality of feed-forward layer.
        dropout_rate (float): Dropout probability between 0 and 1 (exclusive).

    Shape:
        Input:  (B, S, D) where D == d_model
        Output: (B, S, D)
    """

    def __init__(self, d_model: int, d_ff: int, dropout_rate: float = 0.1):
        super().__init__()

        if not isinstance(d_ff, int):
            raise TypeError(f"d_ff must be an int, got {type(d_ff)}")
        if not isinstance(d_model, int):
            raise TypeError(f"d_model must be an int, got {type(d_model)}")
        if not isinstance(dropout_rate, float):
            raise TypeError(f"dropout_rate must be a float, got {type(dropout_rate)}")

        if d_ff <= 0:
            raise ValueError(f"d_ff must be strictly greater than 0, got {d_ff}")
        if d_model <= 0:
            raise ValueError(f"d_model must be strictly greater than 0, got {d_model}")
        if not (0.0 <= dropout_rate < 1.0):
            raise ValueError(f"dropout_rate must be in [0,1), got {dropout_rate}")

        self.d_model = d_model
        self.d_ff = d_ff
        self.dropout_rate = dropout_rate

        self.fc1 = nn.Linear(d_model, d_ff)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(dropout_rate)
        self.fc2 = nn.Linear(d_ff, d_model)

    def forward(self, x: Tensor) -> Tensor:
        if not isinstance(x, torch.Tensor):
            raise TypeError(f"x must be a torch.Tensor, got {type(x)}")
        if x.ndim != 3:
            raise ValueError(f"x must be 3D of shape (B,S,D); got shape {tuple(x.shape)}")
        if x.shape[-1] != self.d_model:
            raise ValueError(f"Last dim {x.shape[-1]} must match d_model {self.d_model}")

        return self.fc2(self.dropout(self.relu(self.fc1(x))))