"""Two-component U-Net for nonlinear Klein-Gordon forecasts.

The state contains 196 coarse displacement coordinates and 196 scaled velocity
coordinates. The segment head receives both components of the noisy state and
conditioning state, together with interpolation time and segment boundaries.
It emits both components through the shared periodic U-Net and frozen spectral
normalization wrappers. Internal parameter names are retained for checkpoint loading."""
from __future__ import annotations

import torch

from ..registry import register_model
from .unet2 import K_C, _broadcast_scalar, _PeriodicUNet2Core, _require_packaged, _SpectralFieldWrapper


class _TwoComponent(_PeriodicUNet2Core):
    """Shared plumbing: one frozen wrapper applied to each half of the 2*D_c state vector."""

    def __init__(self, D: int = 392, base_channels: int = 64, n_grid: int = 32, k_c: int = K_C,
                 packaged: str | None = None, dropout: float = 0.0, in_channels: int = 7):
        super().__init__(in_channels=in_channels, base_channels=base_channels, dropout=dropout,
                         out_channels=2)
        self.D = int(D)
        if packaged is None:
            raise ValueError("unet2_2comp needs the training dataset's normalization: pass packaged")
        # one wrapper per component, each with that component's slice of the frozen normalization
        self.wrap = _SpectralFieldWrapper(packaged, n_grid, k_c, component=0)
        self.wrap_im = _SpectralFieldWrapper(packaged, n_grid, k_c, component=1)
        if self.wrap.D * 2 != self.D:
            raise ValueError(f"two-component state needs D = 2 * {self.wrap.D}, got {self.D}")
        self.Dc = self.wrap.D

    def _fields(self, v: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        """(B, 2*Dc) -> two (B, 1, Ng, Ng) field channels."""
        return (self.wrap.inv(v[:, :self.Dc]).unsqueeze(1),
                self.wrap_im.inv(v[:, self.Dc:]).unsqueeze(1))

    def _vec(self, field2: torch.Tensor) -> torch.Tensor:
        """(B, 2, Ng, Ng) -> (B, 2*Dc)."""
        return torch.cat([self.wrap.fwd(field2[:, 0]), self.wrap_im.fwd(field2[:, 1])], dim=1)


class LocalDrift2Comp(_TwoComponent):
    """forward(x_tau, tau, x_tk, t_k, t_k1) -> (B, D); channels [Re/Im x_tau, Re/Im x_tk, tau, t_k, t_k1]."""

    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
        a_re, a_im = self._fields(x_tau)
        b_re, b_im = self._fields(x_tk)
        tau_bc = _broadcast_scalar(tau, B, ng)
        tk_bc = t_k.view(B, 1, 1, 1).expand(B, 1, ng, ng)
        tk1_bc = t_k1.view(B, 1, 1, 1).expand(B, 1, ng, ng)
        out = self._forward_core(torch.cat([a_re, a_im, b_re, b_im, tau_bc, tk_bc, tk1_bc], dim=1))
        return self._vec(out)

    @classmethod
    def from_config(cls, cfg: dict) -> "LocalDrift2Comp":
        return cls(D=int(cfg.get("D", 392)), 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 Conditional2Comp(LocalDrift2Comp):
    """Endpoint head with the segment head's interface (t_k = t_start, t_k1 = t_end), as unet2_5ch."""

    def __init__(self, *args, t_start: float = 0.0, t_end: float = 1.0, **kwargs):
        super().__init__(*args, **kwargs)
        self.t_start, self.t_end = float(t_start), float(t_end)

    def forward(self, x: torch.Tensor, sigma_embedded: torch.Tensor, x_0: torch.Tensor) -> torch.Tensor:  # type: ignore[override]
        B = x.shape[0]
        tk = torch.full((B,), self.t_start, dtype=x.dtype, device=x.device)
        tk1 = torch.full((B,), self.t_end, dtype=x.dtype, device=x.device)
        return super().forward(x, sigma_embedded, x_0, tk, tk1)

    @classmethod
    def from_config(cls, cfg: dict) -> "Conditional2Comp":
        return cls(D=int(cfg.get("D", 392)), 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)),
                   t_start=float(cfg.get("t_start", 0.0)), t_end=float(cfg.get("t_end", 1.0)))


register_model("unet2_2comp", endpoint_cls=Conditional2Comp, segment_cls=LocalDrift2Comp,
               needs_dataset_norm=True,
               description="two-component (Re,Im) coarse complex state; unet2 core, per-component FFT wrapper")