"""Long Short-Term Memory over channel × time windows.

Same input/output as GRU, swap the cell. ~30% more parameters than the
GRU at the same hidden size; included so the comparison report can
speak to the cell-choice question."""
from __future__ import annotations

from training.models import register
from training.models._torch_seq import _SeqBase


@register("lstm")
class LSTM(_SeqBase):
    def _build_module(self, *, n_channels_in: int, n_timesteps: int,
                      n_classes: int, hidden: int = 128, n_layers: int = 2,
                      dropout: float = 0.1, bidirectional: bool = False):
        return _LSTMClassifier(
            n_channels_in=n_channels_in, n_classes=n_classes,
            hidden=hidden, n_layers=n_layers,
            dropout=dropout, bidirectional=bidirectional,
        )


from torch import nn  # noqa: E402


class _LSTMClassifier(nn.Module):
    def __init__(self, *, n_channels_in: int, n_classes: int, hidden: int,
                 n_layers: int, dropout: float, bidirectional: bool):
        super().__init__()
        self.lstm = nn.LSTM(
            input_size=n_channels_in, hidden_size=hidden,
            num_layers=n_layers, dropout=dropout if n_layers > 1 else 0.0,
            batch_first=True, bidirectional=bidirectional,
        )
        d_out = hidden * (2 if bidirectional else 1)
        self.head = nn.Sequential(nn.Dropout(dropout), nn.Linear(d_out, n_classes))

    def forward(self, x):                       # (B, C, T) → (B, T, C)
        x = x.transpose(1, 2)
        out, _ = self.lstm(x)
        return self.head(out[:, -1, :])
