"""unet2 with coordinate channels — a NON-equivariant variant of the FFT-wrapped 2-D U-Net.

The passive scalar's fixed cellular flow
psi_K is not translation-invariant, so its conditional-mean map x_0 -> M_T x_0 is not
translation-equivariant; ``unet2`` (exactly equivariant to shifts by its pooling stride)
provably cannot represent it. This variant keeps unet2's field reconstruction, circular
convolutions and heads but appends two constant coordinate planes (x, y in [-1, 1)) to the
input channels, which breaks translation equivariance and lets position-dependent filters
form.

  - ``Conditional2DCoord``  forward(x, sigma_embedded, x_0) -> (B, D)       (5 input channels)
  - ``LocalDrift2DCoord``   forward(x_tau, tau, x_tk, t_k, t_k1) -> (B, D)   (7 input channels)

"""

from __future__ import annotations

import torch

from .spectral import K_C
from .unet2 import _PeriodicUNet2Core, _SpectralFieldWrapper, _broadcast_scalar, _require_packaged


def _coord_planes(B: int, ng: int, like: torch.Tensor) -> torch.Tensor:
    """(B, 2, ng, ng) constant planes with x and y coordinates in [-1, 1)."""
    c = torch.arange(ng, device=like.device, dtype=like.dtype) * (2.0 / ng) - 1.0
    X = c.view(1, 1, ng, 1).expand(B, 1, ng, ng)
    Y = c.view(1, 1, 1, ng).expand(B, 1, ng, ng)
    return torch.cat([X, Y], dim=1)


class Conditional2DCoord(_PeriodicUNet2Core):
    def __init__(self, D: int = 196, base_channels: int = 64, n_grid: int = 32,
                 k_c: int = K_C, packaged: str | None = None, dropout: float = 0.0):
        super().__init__(in_channels=5, base_channels=base_channels, dropout=dropout)
        self.D = D
        if packaged is None:
            raise ValueError("unet2_coord needs the training dataset's normalization: pass packaged")
        self.wrap = _SpectralFieldWrapper(packaged, n_grid, k_c)
        assert self.wrap.D == D

    def forward(self, x: torch.Tensor, sigma_embedded: torch.Tensor, x_0: torch.Tensor) -> torch.Tensor:
        B, ng = x.shape[0], self.wrap.n_grid
        x_f = self.wrap.inv(x).unsqueeze(1)
        field = self._forward_core(torch.cat([x_f, self.wrap.inv(x_0).unsqueeze(1),
                                              _broadcast_scalar(sigma_embedded, B, ng).to(x_f),
                                              _coord_planes(B, ng, x_f)], dim=1))
        return self.wrap.fwd(field)

    @classmethod
    def from_config(cls, cfg: dict) -> "Conditional2DCoord":
        return cls(D=int(cfg.get("D", 196)), base_channels=int(cfg.get("base_channels", 64)),
                   n_grid=int(cfg.get("n_grid", 32)), k_c=int(cfg.get("k_c", K_C)),
                   packaged=_require_packaged(cfg), dropout=float(cfg.get("dropout", 0.0)))


class LocalDrift2DCoord(_PeriodicUNet2Core):
    def __init__(self, D: int = 196, base_channels: int = 64, n_grid: int = 32,
                 k_c: int = K_C, packaged: str | None = None, dropout: float = 0.0):
        super().__init__(in_channels=7, base_channels=base_channels, dropout=dropout)
        self.D = D
        if packaged is None:
            raise ValueError("unet2_coord needs the training dataset's normalization: pass packaged")
        self.wrap = _SpectralFieldWrapper(packaged, n_grid, k_c)
        assert self.wrap.D == D

    def forward(self, x_tau: torch.Tensor, tau: torch.Tensor, x_tk: torch.Tensor,
                t_k: torch.Tensor, t_k1: torch.Tensor) -> torch.Tensor:
        B, ng = x_tau.shape[0], self.wrap.n_grid
        xt = self.wrap.inv(x_tau).unsqueeze(1)
        field = self._forward_core(torch.cat([
            xt, self.wrap.inv(x_tk).unsqueeze(1), _broadcast_scalar(tau, B, ng).to(xt),
            t_k.view(B, 1, 1, 1).expand(B, 1, ng, ng).to(xt), t_k1.view(B, 1, 1, 1).expand(B, 1, ng, ng).to(xt),
            _coord_planes(B, ng, xt)], dim=1))
        return self.wrap.fwd(field)

    @classmethod
    def from_config(cls, cfg: dict) -> "LocalDrift2DCoord":
        return cls(D=int(cfg.get("D", 196)), base_channels=int(cfg.get("base_channels", 64)),
                   n_grid=int(cfg.get("n_grid", 32)), k_c=int(cfg.get("k_c", K_C)),
                   packaged=_require_packaged(cfg), dropout=float(cfg.get("dropout", 0.0)))


from ..registry import register_model  # noqa: E402

register_model("unet2_coord", Conditional2DCoord, LocalDrift2DCoord, needs_dataset_norm=True,
               description="unet2 + two coordinate input planes (not translation-equivariant; passive scalar)")