"""WiSER CIR end-to-end CIR set-prediction model.

Pipeline (AC-4 + AC-5 split):

  voxel_level  --[scene_encode]-->  scene_tokens (no TX, once per scene)
  scene_tokens --[modulate_scene_tx(tx_emb)]-->  tx_mem (per (scene, TX))
  tx_mem + rx --[head]-->  (exists, delay_ns, peak_db) per query

`SparseCsiDetrModel.encode_tx(...)` remains a thin helper for Fourier + MLP.
`SparseCsiDetrModel.scene_encode(...)` is the TX-independent backbone
forward used by the cache wrapper; `modulate_scene_tx(...)` is the
TX-conditioned pass. The trainer's `CachedSceneTxForward` wraps these so
`C_enc ≤ U` counts real scene-encoder forwards and `C_txmod ≤ P` counts
real TX-modulation forwards.
"""

from __future__ import annotations

from dataclasses import dataclass, field

import torch
from torch import nn

from ..config import ModelConfig as SharedModelConfig
from .backbone import BackboneConfig, build_backbone
from .detr_head import CIRPathSetDETRHead, DetrHeadConfig, ContinuousFourierEmbed


@dataclass(slots=True)
class ModelConfig:
    """Runtime model config. Derive from `config.ModelConfig` via `.from_shared()`."""

    backbone: BackboneConfig = field(default_factory=BackboneConfig)
    head: DetrHeadConfig = field(default_factory=DetrHeadConfig)

    @classmethod
    def from_shared(cls, shared: SharedModelConfig) -> "ModelConfig":
        return cls(
            backbone=BackboneConfig(
                kind=shared.backbone_kind,
                channels=shared.backbone_channels,
                tx_emb_channels=shared.backbone_tx_emb_channels,
                num_downsample_stages=shared.backbone_downsample_stages,
                blocks_per_stage=shared.backbone_blocks_per_stage,
                num_heads=shared.backbone_num_heads,
                mlp_ratio=shared.backbone_mlp_ratio,
            ),
            head=DetrHeadConfig(
                channels=shared.backbone_channels,
                num_queries=shared.query_budget,
                num_decoder_layers=shared.head_num_decoder_layers,
                num_heads=shared.head_num_heads,
                dropout=shared.head_dropout,
                db_low=shared.db_low,
                db_high=shared.db_high,
                peak_db_head_arch=getattr(shared, "peak_db_head_arch", "flat"),
                peak_db_hidden=int(getattr(shared, "peak_db_hidden", 512)),
                peak_db_init_mean_db=float(getattr(shared, "peak_db_init_mean_db", -55.0)),
                delay_head_arch=getattr(shared, "delay_head_arch", "flat"),
                delay_hidden=int(getattr(shared, "delay_hidden", 512)),
                delay_init_mean_ns=float(getattr(shared, "delay_init_mean_ns", 4.0)),
                exists_head_arch=getattr(shared, "exists_head_arch", "flat"),
                exists_hidden=int(getattr(shared, "exists_hidden", 512)),
                exists_init_prob=float(getattr(shared, "exists_init_prob", 0.2)),
            ),
        )


class SparseCsiDetrModel(nn.Module):
    """The importable model class referenced by AC-1's smoke test.

    Exposes `scene_encode` and `modulate_scene_tx` as two separable stages so
    the cache wrapper can enforce AC-5 invariants. The backbone itself
    implements both: for the Round-2 `dense_fallback` kind, `scene_encode`
    runs every sparse block WITHOUT AdaLN-zero modulation (zero condition
    vector), and `modulate_scene_tx` adds the per-TX modulation pass. The
    `trellis2` kind delegates to the vendored `ModulatedSparseTransformerBlock`.
    """

    def __init__(self, config: ModelConfig | None = None) -> None:
        super().__init__()
        self.config = config or ModelConfig()
        self.backbone = build_backbone(self.config.backbone)
        self.tx_embed = ContinuousFourierEmbed(in_dim=3, num_bands=8)
        self.tx_proj = nn.Sequential(
            nn.Linear(self.tx_embed.out_dim, self.config.backbone.tx_emb_channels),
            nn.GELU(),
            nn.Linear(self.config.backbone.tx_emb_channels, self.config.backbone.tx_emb_channels),
        )
        self.head = CIRPathSetDETRHead(self.config.head)

    @classmethod
    def from_shared(cls, shared: SharedModelConfig) -> "SparseCsiDetrModel":
        return cls(ModelConfig.from_shared(shared))

    def encode_tx(self, tx_xyz_norm: torch.Tensor) -> torch.Tensor:
        return self.tx_proj(self.tx_embed(tx_xyz_norm))

    def scene_encode(self, voxel_level: dict) -> dict:
        """TX-independent backbone pass; produces reusable scene tokens.

        Delegates to `backbone.scene_encode_forward` which runs the TX-FREE
        block stack (`DensePlainBlock` × S for dense_fallback, or plain
        `SparseTransformerBlock` × S for trellis2) + the downsample stack.
        The returned dict is safe to cache per unique scene in the batch.
        """
        return self.backbone.scene_encode_forward(voxel_level)

    def modulate_scene_tx(self, scene_tokens: dict, tx_emb: torch.Tensor) -> dict:
        """Per-(scene, TX) modulation pass on the cached scene tokens.

        Delegates to `backbone.modulation_forward`:
          * dense_fallback → `DenseAdaLNBlock × blocks_per_stage` with TX as
            AdaLN-zero condition (`shift, scale, gate` all derived from TX).
          * trellis2 → `ModulatedSparseTransformerBlock × blocks_per_stage`
            with TX projected through `tx_mod_proj` into the vendored
            AdaLN-zero path.

        This is the ONLY code path where TX enters the scene tensor; RX never
        enters here.
        """
        if tx_emb.dim() == 1:
            tx_emb = tx_emb.unsqueeze(0)
        return self.backbone.modulation_forward(scene_tokens, tx_emb)

    def modulate_scene(self, voxel_level: dict, tx_emb: torch.Tensor) -> dict:
        """Legacy one-shot: scene_encode + modulate_scene_tx in one call."""
        scene = self.scene_encode(voxel_level)
        return self.modulate_scene_tx(scene, tx_emb)

    def head_forward(
        self,
        tx_mem: torch.Tensor,
        rx_xyz_norm: torch.Tensor,
        tx_mem_key_padding_mask: torch.Tensor | None = None,
    ) -> dict[str, torch.Tensor]:
        return self.head(tx_mem, rx_xyz_norm, tx_mem_key_padding_mask)


__all__ = ["ModelConfig", "SparseCsiDetrModel"]