"""Independent implementation of FluxNet-LAP (arXiv:2602.01941v2).

The transport equations and epsilon follow Sections 3.2, 3.3 and 3.5;
the backbone dimensions follow Appendix C.2. See reproduction_spec.md
for the explicitly identified gaps filled by read-only official references.
"""
import torch
from torch import nn
from torch.nn import functional as F


class FeatureBlock(nn.Module):
    """Two periodic convolutions, with a learned concatenation skip fusion."""

    def __init__(self, width, kernel_size):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Conv2d(width, width, kernel_size, padding=kernel_size // 2,
                      padding_mode="circular"),
            nn.BatchNorm2d(width), nn.GELU(),
            nn.Conv2d(width, width, kernel_size, padding=kernel_size // 2,
                      padding_mode="circular"),
            nn.BatchNorm2d(width), nn.GELU(),
        )
        self.fuse = nn.Conv2d(2 * width, width, 1)

    def forward(self, features):
        return self.fuse(torch.cat((self.layers(features), features), dim=1))


class FluxNetShallowWater(nn.Module):
    """Map physical states [B,3,H,W] ordered h,mx,my to the next state.

    This operator redistributes cumulative amounts, with no dt multiplier,
    clipping, mass projection, or numerical-solver substeps.
    """

    def __init__(self, base_channels=64, num_blocks=6, kernel_size=5,
                 neighborhood_size=3, epsilon=1e-6):
        super().__init__()
        if kernel_size % 2 != 1 or neighborhood_size % 2 != 1:
            raise ValueError("Convolution and transport sizes must be odd")
        if neighborhood_size < 3 or epsilon <= 0:
            raise ValueError("Invalid transport size or epsilon")
        self.epsilon = float(epsilon)
        radius = neighborhood_size // 2
        self.offsets = tuple((dy, dx) for dy in range(-radius, radius + 1)
                             for dx in range(-radius, radius + 1)
                             if (dy, dx) != (0, 0))
        self.encoder = nn.Sequential(
            nn.Conv2d(3, base_channels, kernel_size, padding=kernel_size // 2,
                      padding_mode="circular"),
            nn.BatchNorm2d(base_channels), nn.GELU(),
            *(FeatureBlock(base_channels, kernel_size) for _ in range(num_blocks)),
        )
        directions = len(self.offsets)
        self.depth_head = nn.Conv2d(base_channels, directions + 1, 1)
        self.pressure_x = nn.Conv2d(base_channels, directions, 1)
        self.pressure_y = nn.Conv2d(base_channels, directions, 1)

    def transport(self, state):
        features = self.encoder(state)
        h = state[:, :1]
        raw = self.depth_head(features)
        outgoing_depth = h * raw[:, :1].sigmoid()
        depth_transfers = outgoing_depth * raw[:, 1:].softmax(dim=1)
        velocity = state[:, 1:] / (h + self.epsilon)
        gate = h.square()
        momentum_x = depth_transfers * velocity[:, :1] + gate * F.softplus(self.pressure_x(features))
        momentum_y = depth_transfers * velocity[:, 1:] + gate * F.softplus(self.pressure_y(features))
        transfers = torch.stack((depth_transfers, momentum_x, momentum_y), dim=1)
        return outgoing_depth, transfers

    def forward(self, state):
        outgoing_depth, transfers = self.transport(state)
        # For h, use the analytically summed alpha*h to retain the lower bound
        # in floating point. Both forms are algebraically the same (Prop. 2).
        outgoing = torch.cat((outgoing_depth, transfers[:, 1:].sum(dim=2)), dim=1)
        updated = state - outgoing
        for index, offset in enumerate(self.offsets):
            # Offset denotes receiver minus sender; rolling moves sender's
            # outgoing field to its receiver on the periodic torus.
            updated = updated + torch.roll(transfers[:, :, index], offset, dims=(-2, -1))
        return updated