"""Periodic one-dimensional CNN for mass-aware data-assimilation correction."""

from __future__ import annotations

import torch
from torch import nn
from torch.nn import functional as F


class PeriodicConv1d(nn.Module):
    """Conv1d with explicit circular padding and unchanged spatial length."""

    def __init__(self, in_channels: int, out_channels: int, kernel_size: int):
        super().__init__()
        if kernel_size % 2 != 1:
            raise ValueError("kernel_size must be odd")
        self.pad = kernel_size // 2
        self.conv = nn.Conv1d(in_channels, out_channels, kernel_size, padding=0)

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        return self.conv(F.pad(inputs, (self.pad, self.pad), mode="circular"))


class MassConservingCNN(nn.Module):
    """Four hidden SELU convolutions followed by the u/h/r output layer."""

    def __init__(self, input_channels: int = 4, hidden_channels: int = 32,
                 hidden_layers: int = 4, kernel_size: int = 3):
        super().__init__()
        if input_channels != 4 or hidden_layers != 4 or kernel_size != 3:
            raise ValueError("paper architecture requires 4 inputs, 4 hidden layers, kernel size 3")
        layers = []
        channels = input_channels
        for _ in range(hidden_layers):
            layers.extend((PeriodicConv1d(channels, hidden_channels, kernel_size), nn.SELU()))
            channels = hidden_channels
        self.hidden = nn.Sequential(*layers)
        self.output = PeriodicConv1d(hidden_channels, 3, kernel_size)

    @property
    def influence_radius(self) -> int:
        return 5

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        if inputs.ndim != 3 or inputs.shape[1] != 4 or inputs.shape[2] != 250:
            raise ValueError(f"expected float tensor [B,4,250], got {tuple(inputs.shape)}")
        if not inputs.is_floating_point():
            raise TypeError("inputs must have a floating-point dtype")
        raw = self.output(self.hidden(inputs))
        return torch.cat((raw[:, :2], F.relu(raw[:, 2:3])), dim=1)