"""``unet2_5ch``: the backbone of the Euler and SQG models.

Channel parity between the endpoint and the local arms: the endpoint head IS the local head
(``LocalDrift2D``, 5 input channels [x_τ, x_0, τ, t_start, t_end]) with the segment's start and end
time fixed to the forecast horizon (t_start = 0, t_end = T). Both heads are therefore the identical
module class with the identical parameter count (1,515,265 at base_channels 64); the two extra
planes carry no information for an endpoint arm (constants) — the fix removes an *interface*
asymmetry, not a capacity difference.

``t_end`` must equal the forecast horizon: ``bifurcation.registry.build_model`` checks the model
configuration's ``t_start`` / ``t_end`` against ``method.sampling.t_values`` and raises on disagreement.
"""
from __future__ import annotations

import torch

from ..registry import register_model
from .unet2 import K_C, LocalDrift2D, _require_packaged


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

    forward(x, sigma_embedded, x_0) -> (B, D): global SI passes sigma_embedded = tau."""

    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) -> "Conditional2D5ch":
        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", Conditional2D5ch, LocalDrift2D, needs_dataset_norm=True,
               description="unet2 with channel parity: both heads are LocalDrift2D (5 ch); the endpoint head "
                           "fixes (t_start, t_end) = (0, T). The backbone of the Euler and SQG models.")