"""``unet2_5ch_coord``: the backbone for datasets without translation symmetry (the passive scalar).

The passive scalar's fixed cellular flow is not translation invariant, so the equivariant ``unet2`` cannot represent
its conditional-mean map. This variant is ``unet2_5ch`` built on the ``unet2_coord`` core, which adds two coordinate
input planes: both heads are ``LocalDrift2DCoord`` with 7 input channels [x_tau, x_tk, tau, t_start, t_end, X, Y],
and the endpoint head fixes (t_start, t_end) = (0, T) as ``Conditional2D5ch`` does. The parameter count is that of
unet2_5ch plus the two extra 3x3 input filters (1,152 at base_channels 64), identical for both heads, so EDM, Direct
SI and Cascaded SI share one module. ``t_end`` must equal the forecast horizon (checked by
``bifurcation.registry.build_model``).
"""
from __future__ import annotations

import torch

from ..registry import register_model
from .spectral import K_C
from .unet2 import _require_packaged
from .unet2_coord import LocalDrift2DCoord


class Conditional2D5chCoord(LocalDrift2DCoord):
    """Endpoint head with the segment head's exact interface: t_k = t_start, t_{k+1} = t_end."""

    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) -> "Conditional2D5chCoord":
        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)),
                   t_start=float(cfg.get("t_start", 0.0)), t_end=float(cfg.get("t_end", 1.0)))


register_model("unet2_5ch_coord", Conditional2D5chCoord, LocalDrift2DCoord, needs_dataset_norm=True,
               description="unet2_coord with channel parity: both heads are LocalDrift2DCoord (7 ch, not translation-"
                           "equivariant); the endpoint head fixes (t_start, t_end) = (0, T). The backbone for "
                           "datasets without translation symmetry (passive scalar).")