"""1D-CNN over channel × time windows.

Three conv blocks + global average pooling. Small enough to fit on the
Pi for live inference, expressive enough to learn cross-channel patterns
the GBT baseline can't see.
"""
from __future__ import annotations

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


@register("cnn")
class CNN(_SeqBase):
    def _build_module(self, *, n_channels_in: int, n_timesteps: int,
                      n_classes: int, ch1: int = 64, ch2: int = 128,
                      ch3: int = 128, dropout: float = 0.1):
        from torch import nn
        return nn.Sequential(
            nn.Conv1d(n_channels_in, ch1, kernel_size=5, padding=2),
            nn.BatchNorm1d(ch1), nn.GELU(),
            nn.MaxPool1d(2),                                # T/2
            nn.Conv1d(ch1, ch2, kernel_size=5, padding=2),
            nn.BatchNorm1d(ch2), nn.GELU(),
            nn.MaxPool1d(2),                                # T/4
            nn.Conv1d(ch2, ch3, kernel_size=3, padding=1),
            nn.BatchNorm1d(ch3), nn.GELU(),
            nn.AdaptiveAvgPool1d(1),                         # → (B, ch3, 1)
            nn.Flatten(),
            nn.Dropout(dropout),
            nn.Linear(ch3, n_classes),
        )
