"""Small neural-network layers used by RiboSphere."""

from __future__ import annotations

from torch import Tensor, nn


class FeedForward(nn.Module):
    """Two-layer MLP equivalent to the subset of timm.layers.Mlp used here."""

    def __init__(
        self,
        input_features: int,
        hidden_features: int,
        output_features: int,
        *,
        activation: type[nn.Module] = nn.GELU,
        dropout: float = 0.0,
    ) -> None:
        super().__init__()
        if min(input_features, hidden_features, output_features) <= 0:
            raise ValueError("All feature dimensions must be positive.")
        if not 0.0 <= dropout < 1.0:
            raise ValueError("dropout must be in [0, 1).")

        self.fc1 = nn.Linear(input_features, hidden_features)
        self.activation = activation()
        self.dropout1 = nn.Dropout(dropout)
        self.fc2 = nn.Linear(hidden_features, output_features)
        self.dropout2 = nn.Dropout(dropout)

    def forward(self, inputs: Tensor) -> Tensor:
        """Apply the MLP without changing leading dimensions."""
        outputs = self.fc1(inputs)
        outputs = self.activation(outputs)
        outputs = self.dropout1(outputs)
        outputs = self.fc2(outputs)
        outputs = self.dropout2(outputs)
        return outputs