"""Visibility-aware radiomap decoder operating on gathered voxel context."""

from __future__ import annotations

import torch
import torch.nn as nn
import torch.nn.functional as F


def _resolve_group_count(channels: int, max_groups: int = 8) -> int:
    groups = min(max(int(max_groups), 1), max(int(channels), 1))
    while groups > 1 and channels % groups != 0:
        groups -= 1
    return max(groups, 1)


class ResidualMLPBlock(nn.Module):
    """Residual MLP block used to deepen the decoder head stack."""

    def __init__(self, dim: int, hidden_dim: int, dropout: float = 0.0) -> None:
        super().__init__()
        self.norm = nn.LayerNorm(dim)
        self.fc1 = nn.Linear(dim, hidden_dim)
        self.act = nn.GELU()
        self.drop1 = nn.Dropout(dropout)
        self.fc2 = nn.Linear(hidden_dim, dim)
        self.drop2 = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        residual = x
        x = self.norm(x)
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop1(x)
        x = self.fc2(x)
        x = self.drop2(x)
        return residual + x


class QueryHeadStem(nn.Module):
    """Shared token-space stem before scattering query features back to the RX plane."""

    def __init__(self, input_dim: int, hidden_dim: int, depth: int, dropout: float = 0.0) -> None:
        super().__init__()
        self.input_proj = nn.Sequential(
            nn.LayerNorm(input_dim),
            nn.Linear(input_dim, hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
        )
        self.blocks = nn.ModuleList(
            [ResidualMLPBlock(hidden_dim, hidden_dim * 2, dropout=dropout) for _ in range(max(int(depth), 0))]
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.input_proj(x)
        for block in self.blocks:
            x = block(x)
        return x


class DeepOutputHead(nn.Module):
    """Legacy deeper per-task MLP head kept for compatibility/debugging."""

    def __init__(self, input_dim: int, hidden_dim: int, depth: int, out_dim: int = 1, dropout: float = 0.0) -> None:
        super().__init__()
        self.blocks = nn.ModuleList(
            [ResidualMLPBlock(input_dim, hidden_dim * 2, dropout=dropout) for _ in range(max(int(depth), 0))]
        )
        self.out = nn.Sequential(
            nn.LayerNorm(input_dim),
            nn.Linear(input_dim, hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, out_dim),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        for block in self.blocks:
            x = block(x)
        return self.out(x)


class DenseSelfAttentionBlock(nn.Module):
    """Pre-norm dense transformer block for query refinement."""

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

    def forward(self, x: torch.Tensor, key_padding_mask: torch.Tensor | None = None) -> torch.Tensor:
        attn_out, _ = self.attn(self.norm1(x), self.norm1(x), self.norm1(x), key_padding_mask=key_padding_mask, need_weights=False)
        x = x + attn_out
        return x + self.mlp(self.norm2(x))


class DenseCrossAttentionBlock(nn.Module):
    """Pre-norm cross-attention block for query-to-memory fusion."""

    def __init__(self, dim: int, context_dim: int, num_heads: int, mlp_ratio: float = 4.0, dropout: float = 0.0) -> None:
        super().__init__()
        self.norm_q = nn.LayerNorm(dim)
        self.norm_ctx = nn.LayerNorm(context_dim)
        self.attn = nn.MultiheadAttention(
            dim,
            num_heads,
            dropout=dropout,
            batch_first=True,
            kdim=context_dim,
            vdim=context_dim,
        )
        self.norm2 = nn.LayerNorm(dim)
        self.mlp = nn.Sequential(
            nn.Linear(dim, int(dim * mlp_ratio)),
            nn.GELU(),
            nn.Linear(int(dim * mlp_ratio), dim),
        )

    def forward(self, x: torch.Tensor, context: torch.Tensor, key_padding_mask: torch.Tensor | None = None) -> torch.Tensor:
        attn_out, _ = self.attn(
            self.norm_q(x),
            self.norm_ctx(context),
            self.norm_ctx(context),
            key_padding_mask=key_padding_mask,
            need_weights=False,
        )
        x = x + attn_out
        return x + self.mlp(self.norm2(x))


class SpatialResidualBlock(nn.Module):
    """Conv residual block for spatial radiomap decoding."""

    def __init__(self, channels: int, hidden_channels: int, dropout: float = 0.0) -> None:
        super().__init__()
        groups = _resolve_group_count(channels)
        hidden_groups = _resolve_group_count(hidden_channels)
        self.norm1 = nn.GroupNorm(groups, channels)
        self.conv1 = nn.Conv2d(channels, hidden_channels, kernel_size=3, padding=1)
        self.act = nn.GELU()
        self.drop1 = nn.Dropout2d(dropout)
        self.norm2 = nn.GroupNorm(hidden_groups, hidden_channels)
        self.conv2 = nn.Conv2d(hidden_channels, channels, kernel_size=3, padding=1)
        self.drop2 = nn.Dropout2d(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        residual = x
        x = self.norm1(x)
        x = self.conv1(x)
        x = self.act(x)
        x = self.drop1(x)
        x = self.norm2(x)
        x = self.conv2(x)
        x = self.drop2(x)
        return residual + x


class SpatialHeadStem(nn.Module):
    """Shared spatial trunk inspired by the older conv-based radiomap head."""

    def __init__(self, channels: int, depth: int, dropout: float = 0.0, use_coord_channels: bool = True) -> None:
        super().__init__()
        self.channels = int(channels)
        self.use_coord_channels = bool(use_coord_channels)
        extra_channels = 1 + (2 if self.use_coord_channels else 0)
        self.input_proj = nn.Sequential(
            nn.Conv2d(self.channels + extra_channels, self.channels, kernel_size=3, padding=1),
            nn.GroupNorm(_resolve_group_count(self.channels), self.channels),
            nn.GELU(),
            nn.Dropout2d(dropout),
        )
        self.blocks = nn.ModuleList(
            [
                SpatialResidualBlock(
                    channels=self.channels,
                    hidden_channels=max(self.channels * 2, self.channels + 32),
                    dropout=dropout,
                )
                for _ in range(max(int(depth), 1))
            ]
        )

    def forward(self, x: torch.Tensor, observed_mask: torch.Tensor, extent_mask: torch.Tensor) -> torch.Tensor:
        features = [x, observed_mask]
        if self.use_coord_channels:
            batch_size, _, height, width = x.shape
            yy = torch.linspace(-1.0, 1.0, steps=height, device=x.device, dtype=x.dtype)
            xx = torch.linspace(-1.0, 1.0, steps=width, device=x.device, dtype=x.dtype)
            grid_y, grid_x = torch.meshgrid(yy, xx, indexing="ij")
            coord = torch.stack([grid_x, grid_y], dim=0).unsqueeze(0).expand(batch_size, -1, -1, -1)
            features.append(coord)
        x = self.input_proj(torch.cat(features, dim=1))
        x = x * extent_mask
        for block in self.blocks:
            x = block(x) * extent_mask
        return x


class SpatialOutputHead(nn.Module):
    """Small U-Net-like conv head that predicts dense plane outputs."""

    def __init__(self, channels: int, hidden_channels: int, depth: int, out_channels: int = 1, dropout: float = 0.0) -> None:
        super().__init__()
        self.pre_blocks = nn.ModuleList(
            [
                SpatialResidualBlock(
                    channels=channels,
                    hidden_channels=max(channels * 2, channels + 32),
                    dropout=dropout,
                )
                for _ in range(max(int(depth), 1))
            ]
        )
        self.down = nn.Sequential(
            nn.GroupNorm(_resolve_group_count(channels), channels),
            nn.GELU(),
            nn.Conv2d(channels, hidden_channels, kernel_size=3, stride=2, padding=1),
        )
        self.bottleneck_blocks = nn.ModuleList(
            [
                SpatialResidualBlock(
                    channels=hidden_channels,
                    hidden_channels=max(hidden_channels * 2, hidden_channels + 32),
                    dropout=dropout,
                )
                for _ in range(max(int(depth), 1))
            ]
        )
        self.up_proj = nn.Sequential(
            nn.GroupNorm(_resolve_group_count(hidden_channels), hidden_channels),
            nn.GELU(),
            nn.Conv2d(hidden_channels, channels, kernel_size=3, padding=1),
        )
        self.fuse = nn.Sequential(
            nn.Conv2d(channels * 2, channels, kernel_size=3, padding=1),
            nn.GroupNorm(_resolve_group_count(channels), channels),
            nn.GELU(),
        )
        self.refine_blocks = nn.ModuleList(
            [
                SpatialResidualBlock(
                    channels=channels,
                    hidden_channels=max(channels * 2, channels + 32),
                    dropout=dropout,
                )
                for _ in range(max(int(depth) - 1, 0))
            ]
        )
        self.out = nn.Conv2d(channels, out_channels, kernel_size=1)

    def forward(self, x: torch.Tensor, extent_mask: torch.Tensor) -> torch.Tensor:
        skip = x * extent_mask
        for block in self.pre_blocks:
            skip = block(skip) * extent_mask

        down = self.down(skip)
        down_mask = F.interpolate(extent_mask, size=down.shape[-2:], mode="nearest")
        down = down * down_mask
        for block in self.bottleneck_blocks:
            down = block(down) * down_mask

        up = self.up_proj(down)
        up = F.interpolate(up, size=skip.shape[-2:], mode="bilinear", align_corners=False)
        fused = self.fuse(torch.cat([skip, up], dim=1)) * extent_mask
        for block in self.refine_blocks:
            fused = block(fused) * extent_mask
        return self.out(fused) * extent_mask


class VisibilityAwareQueryDecoder(nn.Module):
    """Fuse per-query voxel memory and predict radiomap outputs."""

    def __init__(
        self,
        *,
        query_dim: int,
        memory_dim: int,
        decoder_dim: int,
        num_heads: int,
        cross_depth: int,
        self_depth: int,
        output_head_hidden_dim: int = 0,
        output_head_shared_depth: int = 2,
        output_head_branch_depth: int = 2,
        output_head_dropout: float = 0.0,
        mlp_ratio: float = 4.0,
        dropout: float = 0.0,
        predict_valid: bool = False,
        predict_visibility: bool = False,
        predict_boundary: bool = False,
        predict_uncertainty: bool = False,
        use_dual_gain_heads: bool = False,
        use_los_residual_heads: bool = False,
    ) -> None:
        super().__init__()
        self.predict_valid = bool(predict_valid)
        self.predict_visibility = bool(predict_visibility)
        self.predict_boundary = bool(predict_boundary)
        self.predict_uncertainty = bool(predict_uncertainty)
        self.use_los_residual_heads = bool(use_los_residual_heads)
        requested_dual_gain = bool(use_dual_gain_heads)
        if self.use_los_residual_heads and requested_dual_gain:
            raise ValueError("use_los_residual_heads and use_dual_gain_heads are mutually exclusive")
        if self.use_los_residual_heads and not self.predict_visibility:
            self.predict_visibility = True
        self.use_dual_gain_heads = bool(requested_dual_gain and self.predict_visibility)

        head_hidden_dim = int(output_head_hidden_dim) if int(output_head_hidden_dim) > 0 else int(decoder_dim * 2)
        spatial_head_dim = min(max(head_hidden_dim // 2, 96), 256)
        aux_hidden_dim = min(max(spatial_head_dim, 96), 256)

        self.query_proj = nn.Sequential(nn.Linear(query_dim, decoder_dim), nn.LayerNorm(decoder_dim))
        self.memory_proj = nn.Sequential(nn.Linear(memory_dim, decoder_dim), nn.LayerNorm(decoder_dim))
        self.cross_blocks = nn.ModuleList(
            [
                DenseCrossAttentionBlock(
                    dim=decoder_dim,
                    context_dim=decoder_dim,
                    num_heads=num_heads,
                    mlp_ratio=mlp_ratio,
                    dropout=dropout,
                )
                for _ in range(max(int(cross_depth), 1))
            ]
        )
        self.self_blocks = nn.ModuleList(
            [
                DenseSelfAttentionBlock(
                    dim=decoder_dim,
                    num_heads=num_heads,
                    mlp_ratio=mlp_ratio,
                    dropout=dropout,
                )
                for _ in range(max(int(self_depth), 0))
            ]
        )
        self.head_stem = QueryHeadStem(
            input_dim=int(decoder_dim),
            hidden_dim=head_hidden_dim,
            depth=int(output_head_shared_depth),
            dropout=float(output_head_dropout),
        )
        self.spatial_token_proj = nn.Sequential(
            nn.LayerNorm(head_hidden_dim),
            nn.Linear(head_hidden_dim, spatial_head_dim),
            nn.GELU(),
            nn.Dropout(output_head_dropout),
        )
        self.spatial_stem = SpatialHeadStem(
            channels=spatial_head_dim,
            depth=int(output_head_shared_depth),
            dropout=float(output_head_dropout),
            use_coord_channels=True,
        )

        self.path_gain_head = (
            None
            if (self.use_dual_gain_heads or self.use_los_residual_heads)
            else SpatialOutputHead(
                channels=spatial_head_dim,
                hidden_channels=aux_hidden_dim,
                depth=int(output_head_branch_depth),
                out_channels=1,
                dropout=float(output_head_dropout),
            )
        )
        self.valid_head = (
            SpatialOutputHead(
                channels=spatial_head_dim,
                hidden_channels=max(aux_hidden_dim // 2, 64),
                depth=int(output_head_branch_depth),
                out_channels=1,
                dropout=float(output_head_dropout),
            )
            if self.predict_valid
            else None
        )
        self.visibility_head = (
            SpatialOutputHead(
                channels=spatial_head_dim,
                hidden_channels=max(aux_hidden_dim // 2, 64),
                depth=int(output_head_branch_depth),
                out_channels=1,
                dropout=float(output_head_dropout),
            )
            if self.predict_visibility
            else None
        )
        self.boundary_head = (
            SpatialOutputHead(
                channels=spatial_head_dim,
                hidden_channels=max(aux_hidden_dim // 2, 64),
                depth=int(output_head_branch_depth),
                out_channels=1,
                dropout=float(output_head_dropout),
            )
            if self.predict_boundary
            else None
        )
        self.clear_gain_head = (
            SpatialOutputHead(
                channels=spatial_head_dim,
                hidden_channels=aux_hidden_dim,
                depth=int(output_head_branch_depth),
                out_channels=1,
                dropout=float(output_head_dropout),
            )
            if self.use_dual_gain_heads
            else None
        )
        self.blocked_gain_head = (
            SpatialOutputHead(
                channels=spatial_head_dim,
                hidden_channels=aux_hidden_dim,
                depth=int(output_head_branch_depth),
                out_channels=1,
                dropout=float(output_head_dropout),
            )
            if self.use_dual_gain_heads
            else None
        )
        self.los_gain_head = (
            SpatialOutputHead(
                channels=spatial_head_dim,
                hidden_channels=aux_hidden_dim,
                depth=int(output_head_branch_depth),
                out_channels=1,
                dropout=float(output_head_dropout),
            )
            if self.use_los_residual_heads
            else None
        )
        self.residual_gain_head = (
            SpatialOutputHead(
                channels=spatial_head_dim,
                hidden_channels=aux_hidden_dim,
                depth=int(output_head_branch_depth),
                out_channels=1,
                dropout=float(output_head_dropout),
            )
            if self.use_los_residual_heads
            else None
        )
        self.uncertainty_head = (
            SpatialOutputHead(
                channels=spatial_head_dim,
                hidden_channels=max(aux_hidden_dim // 2, 64),
                depth=int(output_head_branch_depth),
                out_channels=1,
                dropout=float(output_head_dropout),
            )
            if self.predict_uncertainty
            else None
        )

    @staticmethod
    def _db_to_power(db: torch.Tensor) -> torch.Tensor:
        return torch.pow(10.0, db.clamp(min=-200.0, max=80.0) / 10.0)

    @staticmethod
    def _power_to_db(power: torch.Tensor) -> torch.Tensor:
        return 10.0 * torch.log10(power.clamp_min(1e-20))

    @staticmethod
    def _build_extent_mask(grid_shape: torch.Tensor, max_h: int, max_w: int, dtype: torch.dtype) -> torch.Tensor:
        batch_size = int(grid_shape.shape[0])
        mask = torch.zeros((batch_size, 1, max_h, max_w), dtype=dtype, device=grid_shape.device)
        for batch_idx in range(batch_size):
            grid_h = int(grid_shape[batch_idx, 0].item())
            grid_w = int(grid_shape[batch_idx, 1].item())
            mask[batch_idx, :, :grid_h, :grid_w] = 1.0
        return mask

    @staticmethod
    def _scatter_query_tokens_to_dense(
        token_features: torch.Tensor,
        grid_shape: torch.Tensor,
        rx_flat_index: torch.Tensor,
        query_padding_mask: torch.Tensor | None,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        batch_size, num_queries, channels = token_features.shape
        max_h = int(grid_shape[:, 0].max().item())
        max_w = int(grid_shape[:, 1].max().item())
        dense_flat = token_features.new_zeros((batch_size, channels, max_h * max_w))
        count_flat = token_features.new_zeros((batch_size, 1, max_h * max_w))
        if query_padding_mask is None:
            valid_queries = torch.ones((batch_size, num_queries), dtype=torch.bool, device=token_features.device)
        else:
            valid_queries = query_padding_mask > 0.5
        valid_queries = valid_queries & (rx_flat_index >= 0)

        for batch_idx in range(batch_size):
            sample_valid = valid_queries[batch_idx]
            if not bool(sample_valid.any()):
                continue
            sample_index = rx_flat_index[batch_idx, sample_valid].long()
            sample_tokens = token_features[batch_idx, sample_valid]
            grid_w = int(grid_shape[batch_idx, 1].item())
            rows = torch.div(sample_index, grid_w, rounding_mode="floor")
            cols = sample_index.remainder(grid_w)
            dense_index = rows * max_w + cols
            dense_flat[batch_idx].scatter_add_(
                1,
                dense_index.unsqueeze(0).expand(channels, -1),
                sample_tokens.transpose(0, 1),
            )
            count_flat[batch_idx, 0].scatter_add_(
                0,
                dense_index,
                torch.ones_like(dense_index, dtype=token_features.dtype),
            )

        observed_mask = (count_flat > 0).view(batch_size, 1, max_h, max_w).to(dtype=token_features.dtype)
        dense = (dense_flat / count_flat.clamp_min(1.0)).view(batch_size, channels, max_h, max_w)
        extent_mask = VisibilityAwareQueryDecoder._build_extent_mask(grid_shape, max_h, max_w, token_features.dtype)
        return dense * extent_mask, observed_mask, extent_mask

    @staticmethod
    def _gather_dense_predictions(
        dense_map: torch.Tensor,
        grid_shape: torch.Tensor,
        rx_flat_index: torch.Tensor,
    ) -> torch.Tensor:
        batch_size, channels, _, max_w = dense_map.shape
        flat_map = dense_map.flatten(start_dim=2)
        outputs = []
        for batch_idx in range(batch_size):
            sample_index = rx_flat_index[batch_idx].long()
            sample_out = dense_map.new_zeros((channels, sample_index.numel()))
            valid = sample_index >= 0
            if bool(valid.any()):
                grid_w = int(grid_shape[batch_idx, 1].item())
                rows = torch.div(sample_index[valid], grid_w, rounding_mode="floor")
                cols = sample_index[valid].remainder(grid_w)
                dense_index = rows * max_w + cols
                sample_out[:, valid] = flat_map[batch_idx].index_select(1, dense_index)
            outputs.append(sample_out.transpose(0, 1))
        return torch.stack(outputs, dim=0)

    def _apply_spatial_head(
        self,
        head: SpatialOutputHead | None,
        dense_hidden: torch.Tensor,
        extent_mask: torch.Tensor,
        grid_shape: torch.Tensor,
        rx_flat_index: torch.Tensor,
    ) -> torch.Tensor | None:
        if head is None:
            return None
        dense_map = head(dense_hidden, extent_mask)
        gathered = self._gather_dense_predictions(dense_map, grid_shape=grid_shape, rx_flat_index=rx_flat_index)
        return gathered.squeeze(-1)

    def forward(
        self,
        *,
        query_tokens: torch.Tensor,
        memory_tokens: torch.Tensor,
        memory_valid_mask: torch.Tensor,
        grid_shape: torch.Tensor,
        rx_flat_index: torch.Tensor,
        query_padding_mask: torch.Tensor | None = None,
    ) -> dict[str, torch.Tensor]:
        x = self.query_proj(query_tokens)
        memory = self.memory_proj(memory_tokens)
        if memory.dim() == 4:
            batch_size, num_queries, memory_len, dim = memory.shape
            flat_query = x.view(batch_size * num_queries, 1, dim)
            flat_memory = memory.view(batch_size * num_queries, memory_len, dim)
            flat_memory_invalid = (~memory_valid_mask.bool()).view(batch_size * num_queries, memory_len)
            for block in self.cross_blocks:
                flat_query = block(flat_query, flat_memory, key_padding_mask=flat_memory_invalid)
            x = flat_query.view(batch_size, num_queries, dim)
        elif memory.dim() == 3:
            memory_invalid = ~memory_valid_mask.bool()
            for block in self.cross_blocks:
                x = block(x, memory, key_padding_mask=memory_invalid)
        else:
            raise ValueError(f"Unsupported memory token rank: {memory.dim()}")

        query_invalid = None
        if query_padding_mask is not None:
            query_invalid = ~(query_padding_mask > 0.5)
        for block in self.self_blocks:
            x = block(x, key_padding_mask=query_invalid)

        head_hidden = self.head_stem(x)
        spatial_tokens = self.spatial_token_proj(head_hidden)
        dense_hidden, observed_mask, extent_mask = self._scatter_query_tokens_to_dense(
            spatial_tokens,
            grid_shape=grid_shape,
            rx_flat_index=rx_flat_index,
            query_padding_mask=query_padding_mask,
        )
        dense_hidden = self.spatial_stem(dense_hidden, observed_mask=observed_mask, extent_mask=extent_mask)

        if self.use_los_residual_heads:
            los_gain_db = self._apply_spatial_head(self.los_gain_head, dense_hidden, extent_mask, grid_shape, rx_flat_index)
            residual_gain_db = self._apply_spatial_head(self.residual_gain_head, dense_hidden, extent_mask, grid_shape, rx_flat_index)
            los_logits = self._apply_spatial_head(self.visibility_head, dense_hidden, extent_mask, grid_shape, rx_flat_index)
            assert los_gain_db is not None and residual_gain_db is not None and los_logits is not None
            los_prob = torch.sigmoid(los_logits)
            total_power = los_prob * self._db_to_power(los_gain_db) + self._db_to_power(residual_gain_db)
            path_gain_db = self._power_to_db(total_power)
            clear_path_gain_db = los_gain_db
            blocked_path_gain_db = residual_gain_db
            visibility_logits = los_logits
        elif self.use_dual_gain_heads:
            clear_path_gain_db = self._apply_spatial_head(self.clear_gain_head, dense_hidden, extent_mask, grid_shape, rx_flat_index)
            blocked_path_gain_db = self._apply_spatial_head(self.blocked_gain_head, dense_hidden, extent_mask, grid_shape, rx_flat_index)
            visibility_logits = self._apply_spatial_head(self.visibility_head, dense_hidden, extent_mask, grid_shape, rx_flat_index)
            assert clear_path_gain_db is not None and blocked_path_gain_db is not None and visibility_logits is not None
            visibility_prob = torch.sigmoid(visibility_logits)
            path_gain_db = visibility_prob * clear_path_gain_db + (1.0 - visibility_prob) * blocked_path_gain_db
        else:
            path_gain_db = self._apply_spatial_head(self.path_gain_head, dense_hidden, extent_mask, grid_shape, rx_flat_index)
            assert path_gain_db is not None
            clear_path_gain_db = path_gain_db
            blocked_path_gain_db = path_gain_db
            visibility_logits = self._apply_spatial_head(self.visibility_head, dense_hidden, extent_mask, grid_shape, rx_flat_index)

        outputs = {
            "query_tokens": x,
            "path_gain_db": path_gain_db,
            "base_path_gain_db": clear_path_gain_db,
            "coarse_path_gain_db": clear_path_gain_db,
        }
        if self.predict_valid and self.valid_head is not None:
            valid_logits = self._apply_spatial_head(self.valid_head, dense_hidden, extent_mask, grid_shape, rx_flat_index)
            if valid_logits is not None:
                outputs["valid_logits"] = valid_logits
        if self.use_dual_gain_heads or self.use_los_residual_heads:
            outputs["clear_path_gain_db"] = clear_path_gain_db
            outputs["blocked_path_gain_db"] = blocked_path_gain_db
        if self.predict_visibility and visibility_logits is not None:
            outputs["visibility_logits"] = visibility_logits
            outputs["occlusion_logits"] = -visibility_logits
            if self.use_dual_gain_heads or self.use_los_residual_heads:
                outputs["residual_path_gain_db"] = (
                    blocked_path_gain_db if self.use_los_residual_heads else blocked_path_gain_db - clear_path_gain_db
                )
        if self.use_los_residual_heads and visibility_logits is not None:
            outputs["los_logits"] = visibility_logits
            outputs["los_gain_db"] = clear_path_gain_db
            outputs["residual_multipath_gain_db"] = blocked_path_gain_db
            outputs["los_prob"] = torch.sigmoid(visibility_logits)
        if self.predict_boundary and self.boundary_head is not None:
            boundary_logits = self._apply_spatial_head(self.boundary_head, dense_hidden, extent_mask, grid_shape, rx_flat_index)
            if boundary_logits is not None:
                outputs["boundary_logits"] = boundary_logits
                outputs["boundary_prob"] = torch.sigmoid(boundary_logits)
        if self.predict_uncertainty and self.uncertainty_head is not None:
            uncertainty_logits = self._apply_spatial_head(self.uncertainty_head, dense_hidden, extent_mask, grid_shape, rx_flat_index)
            if uncertainty_logits is not None:
                outputs["uncertainty_logits"] = uncertainty_logits
        return outputs


__all__ = [
    "DeepOutputHead",
    "DenseCrossAttentionBlock",
    "DenseSelfAttentionBlock",
    "QueryHeadStem",
    "ResidualMLPBlock",
    "SpatialHeadStem",
    "SpatialOutputHead",
    "SpatialResidualBlock",
    "VisibilityAwareQueryDecoder",
]