"""Independently implemented Fourier operator with exact Darcy boundary values."""
from __future__ import annotations

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


class FourierKernel(nn.Module):
    def __init__(self, width: int, modes: int):
        super().__init__()
        self.modes = modes
        self.upper = nn.Parameter(torch.rand(width, width, modes, modes, dtype=torch.complex64) / width**2)
        self.lower = nn.Parameter(torch.rand(width, width, modes, modes, dtype=torch.complex64) / width**2)

    def forward(self, field: torch.Tensor) -> torch.Tensor:
        spectrum = torch.fft.rfft2(field)
        m = self.modes
        if 2 * m > field.shape[-2] or m > spectrum.shape[-1]:
            raise ValueError("Fourier truncation exceeds the input resolution")
        # Contract channels at each retained frequency, with batch broadcasting.
        result = torch.zeros_like(spectrum)
        for band, weights in ((slice(0, m), self.upper), (slice(-m, None), self.lower)):
            coefficients = spectrum[:, :, band, :m].permute(2, 3, 0, 1)
            result[:, :, band, :m] = torch.matmul(
                coefficients, weights.permute(2, 3, 0, 1)
            ).permute(2, 3, 0, 1)
        return torch.fft.irfft2(result, s=field.shape[-2:])


class VINODarcy(nn.Module):
    """Map nodal permeability (B,H,W,1) to hydraulic head on [0,1]^2.

    Architecture details absent from the paper are documented in
    results/reproduction_spec.md. No solution labels enter this module.
    """
    def __init__(self, modes: int = 12, width: int = 32, depth: int = 4,
                 projection_width: int = 128, input_range=(-1.0, 1.0)):
        super().__init__()
        self.input_range = tuple(input_range)
        self.lift = nn.Linear(3, width)
        self.kernels = nn.ModuleList(FourierKernel(width, modes) for _ in range(depth))
        self.spectral_maps = nn.ModuleList(
            nn.Sequential(nn.Conv2d(width, width, 1), nn.GELU(), nn.Conv2d(width, width, 1))
            for _ in range(depth)
        )
        self.local_maps = nn.ModuleList(nn.Conv2d(width, width, 1) for _ in range(depth))
        self.project = nn.Sequential(nn.Linear(width, projection_width), nn.GELU(), nn.Linear(projection_width, 1))
        self.register_buffer("grid", torch.empty(0), persistent=False)

    def forward(self, permeability: torch.Tensor) -> torch.Tensor:
        if permeability.ndim != 4 or permeability.shape[-1] != 1:
            raise ValueError("VINODarcy expects permeability with shape (B,H,W,1)")
        batch, height, width, _ = permeability.shape
        if self.grid.shape != (1, height, width, 2) or self.grid.device != permeability.device:
            x = torch.linspace(0, 1, height, device=permeability.device, dtype=permeability.dtype)
            y = torch.linspace(0, 1, width, device=permeability.device, dtype=permeability.dtype)
            self.grid = torch.stack(torch.meshgrid(x, y, indexing="ij"), dim=-1).unsqueeze(0)
        coordinates = self.grid.expand(batch, -1, -1, -1)
        lo, hi = self.input_range
        features = torch.cat(((2 * permeability - hi - lo) / (hi - lo), coordinates), dim=-1)
        hidden = self.lift(features).movedim(-1, 1)
        for index, (kernel, spectral_map, local_map) in enumerate(zip(self.kernels, self.spectral_maps, self.local_maps)):
            hidden = spectral_map(kernel(hidden)) + local_map(hidden)
            if index + 1 < len(self.kernels):
                hidden = F.gelu(hidden)
        raw = self.project(hidden.movedim(1, -1))
        x, y = coordinates[..., 0:1], coordinates[..., 1:2]
        return raw * (x * (1 - x) * y * (1 - y))