"""Sparse 3D backbone with AdaLN-zero TX modulation (AC-4).

Two implementations selected by `BackboneConfig.kind`:

  * `"trellis2"`: real sparse path — imports `ModulatedSparseTransformerBlock`,
    `SparseTransformerBlock`, and `SparseDownsample` from the vendored
    TRELLIS-2 subset and stacks them with `S in {1, 2, 3}` downsample stages.
    TX enters ONLY via AdaLN-zero modulation. RX never enters.

  * `"dense_fallback"`: a CPU-friendly placeholder used for unit tests. It
    keeps the same public API, implements AdaLN-zero on dense features,
    pretends voxel coordinates survive each downsample pass by passing a
    `VoxelLevel` dict through unchanged, and proves that swapping TX changes
    modulation features without changing coordinates. Tests can thus run
    without TRELLIS-2 CUDA extensions.

The model code always talks to the backbone through the same forward signature
`forward(voxel_level, tx_emb) -> scene_tokens`.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any

import torch
from torch import nn


@dataclass(slots=True)
class BackboneConfig:
    kind: str = "dense_fallback"  # or "trellis2"
    channels: int = 256
    tx_emb_channels: int = 128
    num_downsample_stages: int = 2
    blocks_per_stage: int = 2
    num_heads: int = 8
    mlp_ratio: float = 4.0


class AdaLNZeroModulation(nn.Module):
    """AdaLN-zero modulation: produce (shift, scale, gate) from a condition."""

    def __init__(self, channels: int, cond_channels: int) -> None:
        super().__init__()
        self.proj = nn.Linear(cond_channels, channels * 3)
        nn.init.zeros_(self.proj.weight)
        nn.init.zeros_(self.proj.bias)

    def forward(self, cond: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        out = self.proj(cond)  # [..., 3C]
        shift, scale, gate = out.chunk(3, dim=-1)
        return shift, scale, gate


class DenseAdaLNBlock(nn.Module):
    """MSA + FFN block with AdaLN-zero modulation; dense fallback only."""

    def __init__(self, channels: int, cond_channels: int, num_heads: int, mlp_ratio: float) -> None:
        super().__init__()
        self.norm1 = nn.LayerNorm(channels, elementwise_affine=False)
        self.norm2 = nn.LayerNorm(channels, elementwise_affine=False)
        self.attn = nn.MultiheadAttention(channels, num_heads, batch_first=True)
        self.mlp = nn.Sequential(
            nn.Linear(channels, int(channels * mlp_ratio)),
            nn.GELU(),
            nn.Linear(int(channels * mlp_ratio), channels),
        )
        self.mod_msa = AdaLNZeroModulation(channels, cond_channels)
        self.mod_mlp = AdaLNZeroModulation(channels, cond_channels)

    def forward(self, feats: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
        # feats: [B, N, C]; cond: [B, C_cond]
        shift_msa, scale_msa, gate_msa = self.mod_msa(cond)
        shift_mlp, scale_mlp, gate_mlp = self.mod_mlp(cond)

        h = self.norm1(feats) * (1.0 + scale_msa.unsqueeze(1)) + shift_msa.unsqueeze(1)
        h, _ = self.attn(h, h, h, need_weights=False)
        feats = feats + gate_msa.unsqueeze(1) * h

        h = self.norm2(feats) * (1.0 + scale_mlp.unsqueeze(1)) + shift_mlp.unsqueeze(1)
        h = self.mlp(h)
        feats = feats + gate_mlp.unsqueeze(1) * h
        return feats


class DenseDownsample(nn.Module):
    """Stride-2-like downsample over a [B, N, C] feature table.

    Keeps a passthrough list of "coordinates" so we can assert the
    coordinates-invariance property of TX-only modulation in unit tests.
    """

    def forward(self, feats: torch.Tensor) -> torch.Tensor:
        N = feats.shape[1]
        if N <= 1:
            return feats
        if N % 2 == 1:
            feats = feats[:, : N - 1]
        return (feats[:, 0::2] + feats[:, 1::2]) / 2.0


class DensePlainBlock(nn.Module):
    """Plain MSA + FFN block used by the TX-FREE scene_encode pass (no modulation)."""

    def __init__(self, channels: int, num_heads: int, mlp_ratio: float) -> None:
        super().__init__()
        self.norm1 = nn.LayerNorm(channels)
        self.norm2 = nn.LayerNorm(channels)
        self.attn = nn.MultiheadAttention(channels, num_heads, batch_first=True)
        self.mlp = nn.Sequential(
            nn.Linear(channels, int(channels * mlp_ratio)),
            nn.GELU(),
            nn.Linear(int(channels * mlp_ratio), channels),
        )

    def forward(self, feats: torch.Tensor) -> torch.Tensor:
        h = self.norm1(feats)
        h, _ = self.attn(h, h, h, need_weights=False)
        feats = feats + h
        feats = feats + self.mlp(self.norm2(feats))
        return feats


class DenseFallbackBackbone(nn.Module):
    """Dense-tensor fallback backbone with a TRUE scene/TX split.

    Two sub-pipelines run on the same `[B, N, C]` feature grid:

      * `scene_blocks` + `downsamplers` — **TX-independent**. Runs inside
        `scene_encode_forward(voxel_level)` and produces reusable per-scene
        tokens (cached once per unique scene in the batch per AC-5).
      * `modulation_blocks` — **AdaLN-zero conditioned on TX only**. Runs
        inside `modulation_forward(scene_tokens, tx_emb)` once per unique
        `(scene, TX)` pair in the batch.

    `forward(voxel_level, tx_emb)` remains available as a one-shot path for
    unit tests and anything that doesn't want the split.
    """

    def __init__(self, config: BackboneConfig) -> None:
        super().__init__()
        if config.num_downsample_stages not in (1, 2, 3):
            raise ValueError(
                f"num_downsample_stages must be 1, 2, or 3; got {config.num_downsample_stages}"
            )
        self.config = config
        C = config.channels

        # TX-FREE scene pass.
        self.scene_blocks = nn.ModuleList()
        self.downsamplers = nn.ModuleList()
        for _ in range(config.num_downsample_stages):
            for _ in range(config.blocks_per_stage):
                self.scene_blocks.append(DensePlainBlock(C, config.num_heads, config.mlp_ratio))
            self.downsamplers.append(DenseDownsample())

        # TX-conditioned modulation pass (runs on the cached scene tokens).
        self.modulation_blocks = nn.ModuleList([
            DenseAdaLNBlock(C, config.tx_emb_channels, config.num_heads, config.mlp_ratio)
            for _ in range(max(1, config.blocks_per_stage))
        ])

    # ------------------------------------------------------------------
    # Split forward
    # ------------------------------------------------------------------

    def scene_encode_forward(
        self,
        voxel_level: dict[str, Any],
        return_intermediates: bool = False,
    ) -> dict[str, Any]:
        """TX-independent backbone pass; cache this output per unique scene.

        V1.5: when return_intermediates=True, also returns stage-0 output (pre-downsample)
        so hierarchical memory (corridor + global) can be built in radiomap_head.
        """
        feats = voxel_level["feats"]
        coords = voxel_level["coords"]
        block_iter = iter(self.scene_blocks)
        intermediates: list[dict[str, Any]] = []
        for stage in range(self.config.num_downsample_stages):
            for _ in range(self.config.blocks_per_stage):
                feats = next(block_iter)(feats)
            if return_intermediates and stage == 0:
                intermediates.append({
                    "feats": feats, "coords": coords,
                    "voxel_size_m": 0.1 * (2 ** stage),
                })
            feats = self.downsamplers[stage](feats)
            if coords.shape[1] > 1:
                Nc = coords.shape[1]
                if Nc % 2 == 1:
                    coords = coords[:, : Nc - 1]
                coords = (coords[:, 0::2] + coords[:, 1::2]) / 2.0
        final = {
            "feats": feats, "coords": coords,
            "voxel_size_m": 0.1 * (2 ** self.config.num_downsample_stages),
        }
        if return_intermediates:
            return {"levels": intermediates + [final], "final": final, **final}
        return final

    def modulation_forward(self, scene_tokens: dict[str, Any], tx_emb: torch.Tensor) -> dict[str, Any]:
        """TX-conditioned pass on the cached scene tokens."""
        feats = scene_tokens["feats"]
        coords = scene_tokens["coords"]
        B = feats.shape[0]
        if tx_emb.dim() == 1:
            tx_emb = tx_emb.unsqueeze(0).expand(B, -1)
        for blk in self.modulation_blocks:
            feats = blk(feats, tx_emb)
        # Coordinates are a passthrough — the modulation does NOT touch geometry.
        return {"feats": feats, "coords": coords}

    # ------------------------------------------------------------------
    # Legacy one-shot forward (tests + backward compatibility)
    # ------------------------------------------------------------------

    def forward(self, voxel_level: dict[str, Any], tx_emb: torch.Tensor) -> dict[str, Any]:
        feats = voxel_level["feats"]   # [B, N, C]
        coords = voxel_level["coords"]  # [B, N, 3]
        B = feats.shape[0]
        if tx_emb.dim() == 1:
            tx_emb = tx_emb.unsqueeze(0).expand(B, -1)

        block_iter = iter(self.scene_blocks)
        for stage in range(self.config.num_downsample_stages):
            for _ in range(self.config.blocks_per_stage):
                feats = next(block_iter)(feats)
            feats = self.downsamplers[stage](feats)
            if coords.shape[1] > 1:
                Nc = coords.shape[1]
                if Nc % 2 == 1:
                    coords = coords[:, : Nc - 1]
                coords = (coords[:, 0::2] + coords[:, 1::2]) / 2.0
        # Apply the TX-conditioned modulation pass so the one-shot forward
        # remains equivalent to `scene_encode_forward -> modulation_forward`.
        for blk in self.modulation_blocks:
            feats = blk(feats, tx_emb)
        return {"feats": feats, "coords": coords}


class Trellis2ModulatedBackbone(nn.Module):
    """Real sparse backbone using vendored TRELLIS-2 blocks.

    Round-5 wiring:
        Stage i (i = 1 .. S):
            blocks_per_stage × ModulatedSparseTransformerBlock(channels, num_heads, mlp_ratio)
            SparseDownsample(factor=2)  # (AC-4: at least one sparse 3D downsample stage)

    TX enters only through AdaLN-zero modulation (the `mod` argument on each
    ModulatedSparseTransformerBlock). RX does NOT enter the backbone.
    """

    def __init__(self, config: BackboneConfig) -> None:
        super().__init__()
        from wiser.third_party.trellis2_sparse import get_blocks, get_modulated, get_spatial_basic
        self._blocks_mod = get_blocks()
        self._mod_mod = get_modulated()
        self._spatial_mod = get_spatial_basic()
        self.config = config

        # TX embedding → modulation conditioning vector. AdaLN-zero (used inside
        # ModulatedSparseTransformerBlock) takes a `mod: Tensor [B, C]` input
        # and produces the six (shift, scale, gate) × (MSA, MLP) affine knobs.
        self.tx_mod_proj = nn.Sequential(
            nn.Linear(config.tx_emb_channels, config.channels),
            nn.SiLU(),
            nn.Linear(config.channels, config.channels),
        )

        # TX-FREE scene pass (plain SparseTransformerBlocks + SparseDownsample).
        self.scene_stages = nn.ModuleList()
        self.scene_downsamplers = nn.ModuleList()
        for _ in range(config.num_downsample_stages):
            self.scene_stages.append(nn.ModuleList([
                self._blocks_mod.SparseTransformerBlock(
                    channels=config.channels,
                    num_heads=config.num_heads,
                    mlp_ratio=config.mlp_ratio,
                )
                for _ in range(config.blocks_per_stage)
            ]))
            self.scene_downsamplers.append(self._spatial_mod.SparseDownsample(factor=2))

        # TX-conditioned modulation pass (ModulatedSparseTransformerBlocks).
        self.modulation_blocks = nn.ModuleList([
            self._mod_mod.ModulatedSparseTransformerBlock(
                channels=config.channels,
                num_heads=config.num_heads,
                mlp_ratio=config.mlp_ratio,
            )
            for _ in range(max(1, config.blocks_per_stage))
        ])

    def _build_sparse_tensor(self, feats_bnc: torch.Tensor, coords_bn3: torch.Tensor):
        B, N, C = feats_bnc.shape
        batch_idx = torch.arange(B, device=feats_bnc.device).repeat_interleave(N).unsqueeze(-1)
        int_coords = coords_bn3.reshape(B * N, 3).long()
        flat_coords = torch.cat([batch_idx, int_coords], dim=-1)
        # Respect the incoming dtype so CUDA FlashAttention (requires fp16/bf16)
        # works when the caller is in autocast/bf16 mode. Only .contiguous() —
        # no forced fp32 cast.
        flat_feats = feats_bnc.reshape(B * N, C).contiguous()
        SparseTensor = self._blocks_mod.SparseTensor
        return SparseTensor(flat_feats, flat_coords)

    def scene_encode_forward(self, voxel_level, return_intermediates: bool = False):  # pragma: no cover - sparse backend
        """TX-FREE sparse scene pass; cache the output per unique scene.

        V1.5: when return_intermediates=True, also returns stage-0 output (pre-downsample)
        for hierarchical memory in radiomap_head.
        """
        x = self._build_sparse_tensor(voxel_level["feats"], voxel_level["coords"])
        intermediates: list[dict] = []
        for stage_idx, (stage_blocks, ds) in enumerate(zip(self.scene_stages, self.scene_downsamplers, strict=True)):
            for blk in stage_blocks:
                x = blk(x)
            if return_intermediates and stage_idx == 0:
                intermediates.append({
                    "feats": x.feats, "coords": x.coords,
                    "voxel_size_m": 0.1 * (2 ** stage_idx),
                    "_sparse_tensor": x,
                })
            x = ds(x)
        final = {
            "feats": x.feats, "coords": x.coords,
            "_sparse_tensor": x,
            "voxel_size_m": 0.1 * (2 ** self.config.num_downsample_stages),
        }
        if return_intermediates:
            return {"levels": intermediates + [final], "final": final, **final}
        return final

    def modulation_forward(self, scene_tokens, tx_emb):  # pragma: no cover - sparse backend
        """Per-(scene, TX) TRELLIS modulation on the cached scene tokens."""
        if tx_emb.dim() == 1:
            tx_emb = tx_emb.unsqueeze(0)
        mod = self.tx_mod_proj(tx_emb)  # [B, C]
        x = scene_tokens.get("_sparse_tensor")
        if x is None:
            # Lightweight reconstruction from feats/coords when cache was not
            # routed through scene_encode_forward (e.g., off-path unit test).
            feats = scene_tokens["feats"]
            coords = scene_tokens["coords"]
            if feats.dim() == 2:
                feats = feats.unsqueeze(0)
            if coords.dim() == 2:
                coords = coords.unsqueeze(0)
            x = self._build_sparse_tensor(feats, coords)
        for blk in self.modulation_blocks:
            x = blk(x, mod)
        return {"feats": x.feats, "coords": x.coords}

    def forward(self, voxel_level, tx_emb):  # pragma: no cover - needs sparse backend runtime
        scene = self.scene_encode_forward(voxel_level)
        return self.modulation_forward(scene, tx_emb)


def build_backbone(config: BackboneConfig) -> nn.Module:
    if config.kind == "dense_fallback":
        return DenseFallbackBackbone(config)
    if config.kind == "trellis2":
        return Trellis2ModulatedBackbone(config)
    raise ValueError(f"unknown backbone kind: {config.kind}")


__all__ = [
    "BackboneConfig",
    "AdaLNZeroModulation",
    "DenseFallbackBackbone",
    "Trellis2ModulatedBackbone",
    "build_backbone",
]