"""Tiny Transformer encoder over channel × time windows.

Linear projection of channels → d_model, learned positional embedding,
two encoder layers, mean-pool over time, linear head. Deliberately
small (d_model=64, 4 heads, 2 layers) — the dataset is small enough
that anything bigger overfits within a few epochs."""
from __future__ import annotations

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


@register("transformer")
class Transformer(_SeqBase):
    def _build_module(self, *, n_channels_in: int, n_timesteps: int,
                      n_classes: int, d_model: int = 64, n_heads: int = 4,
                      n_layers: int = 2, ffn_hidden: int = 128,
                      dropout: float = 0.1):
        return _TransformerClassifier(
            n_channels_in=n_channels_in, n_timesteps=n_timesteps,
            n_classes=n_classes, d_model=d_model, n_heads=n_heads,
            n_layers=n_layers, ffn_hidden=ffn_hidden, dropout=dropout,
        )


import torch                # noqa: E402
from torch import nn        # noqa: E402


class _TransformerClassifier(nn.Module):
    def __init__(self, *, n_channels_in: int, n_timesteps: int, n_classes: int,
                 d_model: int, n_heads: int, n_layers: int, ffn_hidden: int,
                 dropout: float):
        super().__init__()
        self.proj = nn.Linear(n_channels_in, d_model)
        self.pos = nn.Parameter(torch.zeros(1, n_timesteps, d_model))
        nn.init.trunc_normal_(self.pos, std=0.02)
        layer = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=n_heads, dim_feedforward=ffn_hidden,
            dropout=dropout, batch_first=True, activation="gelu",
            norm_first=True,
        )
        self.encoder = nn.TransformerEncoder(layer, num_layers=n_layers)
        self.head = nn.Sequential(nn.LayerNorm(d_model),
                                   nn.Dropout(dropout),
                                   nn.Linear(d_model, n_classes))

    def forward(self, x):                       # (B, C, T) → (B, T, C)
        x = x.transpose(1, 2)
        h = self.proj(x) + self.pos[:, : x.size(1), :]
        h = self.encoder(h)                      # (B, T, d_model)
        h = h.mean(dim=1)                        # mean-pool over time
        return self.head(h)
