from __future__ import annotations

import torch
from torch import nn


class TemporalAutoencoder(nn.Module):
    def __init__(self, channels: int = 10) -> None:
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Conv1d(channels, 20, kernel_size=3, stride=2, padding=1),
            nn.GELU(),
            nn.Conv1d(20, 10, kernel_size=3, stride=2, padding=1),
            nn.GELU(),
        )
        self.decoder = nn.Sequential(
            nn.ConvTranspose1d(
                10,
                20,
                kernel_size=4,
                stride=2,
                padding=1,
            ),
            nn.GELU(),
            nn.ConvTranspose1d(
                20,
                channels,
                kernel_size=4,
                stride=2,
                padding=1,
            ),
        )

    def forward(self, windows: torch.Tensor) -> torch.Tensor:
        return self.decoder(self.encoder(windows))


def parameter_count(model: nn.Module) -> int:
    return sum(parameter.numel() for parameter in model.parameters())