"""
Multi-scale prediction head with cross-attention to decoder features.
Query-based predictor using RoPE-aware attention for position encoding.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import List, Optional, Tuple

# Import RoPE modules for position-aware attention
from src.models.decoder import CrossAttentionWithRoPE, SelfAttentionWithRoPE


class QueryBasedMultiScalePredictor(nn.Module):
    """
    Query-based predictor with cross-attention to multi-scale features.
    Anchors attend to multi-scale decoder features instead of pooling them.
    """
    
    def __init__(
        self,
        state_dim: int = 768,
        query_dim: int = 3,
        hidden_dim: int = 256,
        output_dim: int = 3,
        layer_indices: List[int] = [0, 6, 9, 12],
        num_heads: int = 8,
        dropout: float = 0.1,
        use_ref_pose: bool = False,  # Whether to use reference pose
        norm_type: str = 'layer_norm',  # 'layer_norm' | 'rms_norm'
        qk_norm: bool = False,          # normalize q/k before attention
        fusion_activation: str = 'swiglu',  # 'swiglu' | 'gelu' | 'silu'
        use_rope: bool = True,  # Position-aware attention with RoPE
        rope_dim: int = None,  # Custom RoPE dimension
        dim_per_coord_3d: int = None,  # Dim per coordinate for single-point (3D) RoPE
        dim_per_coord_multi: int = None,  # Dim per coordinate for multi-point RoPE
        layer_index_mapping: dict = None,  # Maps original layer index to saved position
        elementwise_gate: bool = False,  # Elementwise gated attention
        cross_object_attention: bool = False,  # Cross-object anchor attention
        cross_object_kv: bool = False,  # Cross-object key-value: anchors query all objects
        no_global_state: bool = False  # Skip state modules entirely when True
    ):
        super().__init__()
        self.cross_object_attention = cross_object_attention
        self.cross_object_kv = cross_object_kv
        self.no_global_state = no_global_state
        self._cross_object_attention_logged = False  # One-time log flag
        self._cross_object_kv_logged = False  # One-time log flag
        self._p0_mask_logged = False  # One-time padding mask log
        self._p1_no_state_logged = False  # One-time no-state log
        if cross_object_attention:
            print(f"[Predictor] Cross-object anchor attention enabled: anchors from different objects can interact")
        if cross_object_kv:
            print(f"[Predictor] Cross-object key-value enabled: anchors query all objects' features")
        if no_global_state:
            print(f"[Predictor] no_global_state=True: state modules omitted, fusion uses object features only")

        self.layer_indices = layer_indices
        self.num_scales = len(layer_indices)

        # Maps original layer index to its position in the saved feature list.
        # If None, all layers are assumed saved.
        if layer_index_mapping is not None:
            self.layer_index_mapping = layer_index_mapping
        else:
            # Default: assume sequential saving (all layers saved)
            self.layer_index_mapping = {idx: idx for idx in layer_indices}
        self.use_ref_pose = use_ref_pose
        self.use_rope = use_rope
        self.hidden_dim = hidden_dim

        # Scale-specific projections for object features
        self.scale_projections = nn.ModuleList([
            nn.Linear(state_dim, hidden_dim) for _ in range(self.num_scales)
        ])

        # Scale-specific projections for state features (skipped when no_global_state)
        if not no_global_state:
            self.state_scale_projections = nn.ModuleList([
                nn.Linear(state_dim, hidden_dim) for _ in range(self.num_scales)
            ])

        # Query projection; query_dim already includes the ref-pose concatenation
        self.query_proj = nn.Linear(query_dim, hidden_dim)

        # Cross-attention to object features for each scale with RoPE
        if use_rope:
            self.cross_attentions = nn.ModuleList([
                CrossAttentionWithRoPE(
                    d_model=hidden_dim,
                    num_heads=num_heads,
                    dropout=dropout,
                    use_rope=True,
                    rope_dim=rope_dim,
                    dim_per_coord_3d=dim_per_coord_3d,
                    dim_per_coord_multi=dim_per_coord_multi,
                    elementwise_gate=elementwise_gate
                ) for _ in range(self.num_scales)
            ])
            # Cross-attention to state features (skipped when no_global_state)
            if not no_global_state:
                self.state_cross_attentions = nn.ModuleList([
                    CrossAttentionWithRoPE(
                        d_model=hidden_dim,
                        num_heads=num_heads,
                        dropout=dropout,
                        use_rope=True,
                        rope_dim=rope_dim,
                        dim_per_coord_3d=dim_per_coord_3d,
                        dim_per_coord_multi=dim_per_coord_multi,
                        elementwise_gate=elementwise_gate
                    ) for _ in range(self.num_scales)
                ])
            # Self-attention for anchors within each object (with RoPE)
            self.self_attentions = nn.ModuleList([
                SelfAttentionWithRoPE(
                    d_model=hidden_dim,
                    num_heads=num_heads,
                    dropout=dropout,
                    use_rope=True,
                    rope_dim=rope_dim,
                    dim_per_coord_3d=dim_per_coord_3d,
                    dim_per_coord_multi=dim_per_coord_multi,
                    elementwise_gate=elementwise_gate
                ) for _ in range(self.num_scales)
            ])
        else:
            # Standard attention
            self.cross_attentions = nn.ModuleList([
                nn.MultiheadAttention(
                    hidden_dim, num_heads, dropout=dropout, batch_first=True
                ) for _ in range(self.num_scales)
            ])
            if not no_global_state:
                self.state_cross_attentions = nn.ModuleList([
                    nn.MultiheadAttention(
                        hidden_dim, num_heads, dropout=dropout, batch_first=True
                    ) for _ in range(self.num_scales)
                ])
            # Self-attention for anchors (standard attention)
            self.self_attentions = nn.ModuleList([
                nn.MultiheadAttention(
                    hidden_dim, num_heads, dropout=dropout, batch_first=True
                ) for _ in range(self.num_scales)
            ])

        # Norm modules (configurable type) for norm-first and qk-norm
        Norm = nn.LayerNorm if norm_type == 'layer_norm' else getattr(nn, 'RMSNorm', nn.LayerNorm)
        # For object attention
        self.query_in_norms = nn.ModuleList([Norm(hidden_dim) for _ in range(self.num_scales)])
        self.kv_in_norms = nn.ModuleList([Norm(hidden_dim) for _ in range(self.num_scales)])
        self.post_attn_norms = nn.ModuleList([Norm(hidden_dim) for _ in range(self.num_scales)])
        # For state attention (skipped when no_global_state)
        if not no_global_state:
            self.state_kv_in_norms = nn.ModuleList([Norm(hidden_dim) for _ in range(self.num_scales)])
            self.state_post_attn_norms = nn.ModuleList([Norm(hidden_dim) for _ in range(self.num_scales)])
        # For self-attention (within same object)
        self.self_attn_in_norms = nn.ModuleList([Norm(hidden_dim) for _ in range(self.num_scales)])
        self.self_attn_post_norms = nn.ModuleList([Norm(hidden_dim) for _ in range(self.num_scales)])
        self._use_qk_norm = qk_norm

        # Cross-object anchor self-attention
        # Allows anchors from different objects to interact directly
        if cross_object_attention:
            if use_rope:
                self.cross_obj_self_attentions = nn.ModuleList([
                    SelfAttentionWithRoPE(
                        d_model=hidden_dim,
                        num_heads=num_heads,
                        dropout=dropout,
                        use_rope=True,
                        rope_dim=rope_dim,
                        dim_per_coord_3d=dim_per_coord_3d,
                        dim_per_coord_multi=dim_per_coord_multi,
                        elementwise_gate=elementwise_gate
                    ) for _ in range(self.num_scales)
                ])
            else:
                self.cross_obj_self_attentions = nn.ModuleList([
                    nn.MultiheadAttention(
                        hidden_dim, num_heads, dropout=dropout, batch_first=True
                    ) for _ in range(self.num_scales)
                ])
            self.cross_obj_in_norms = nn.ModuleList([Norm(hidden_dim) for _ in range(self.num_scales)])
            self.cross_obj_post_norms = nn.ModuleList([Norm(hidden_dim) for _ in range(self.num_scales)])
        
        # Fusion and prediction
        # With no_global_state, fusion uses object features only (no dead zero channels)
        if no_global_state:
            fusion_input_dim = hidden_dim * self.num_scales  # Only object features
            print(f"[Predictor] fusion_input_dim = {fusion_input_dim} (object features only)")
        else:
            fusion_input_dim = hidden_dim * self.num_scales * 2  # × 2 for object + state
        if use_ref_pose:
            # Add extra projection for reference pose in fusion
            # Reference pose is always 3D (positions)
            self.ref_pose_proj = nn.Linear(3, hidden_dim)
            fusion_input_dim += hidden_dim
            print(f"[Predictor] fusion_input_dim with ref_pose = {fusion_input_dim}")
        
        # Fusion head: norm-first + SwiGLU (default)
        def _fusion_block(in_dim: int, out_dim: int) -> nn.Sequential:
            if fusion_activation == 'swiglu':
                return nn.Sequential(
                    nn.LayerNorm(in_dim) if norm_type == 'layer_norm' else getattr(nn, 'RMSNorm', nn.LayerNorm)(in_dim),
                    nn.Linear(in_dim, out_dim),
                    nn.Linear(in_dim, out_dim),  # gate
                )
            else:
                act = nn.GELU() if fusion_activation == 'gelu' else nn.SiLU()
                return nn.Sequential(
                    nn.LayerNorm(in_dim) if norm_type == 'layer_norm' else getattr(nn, 'RMSNorm', nn.LayerNorm)(in_dim),
                    nn.Linear(in_dim, out_dim),
                    act
                )

        # Build two-stage fusion: [fusion_dim -> hidden_dim*2 -> hidden_dim] then linear to output
        self.fusion_dim1 = hidden_dim * 2
        self.fusion_norm1 = nn.LayerNorm(fusion_input_dim) if norm_type == 'layer_norm' else getattr(nn, 'RMSNorm', nn.LayerNorm)(fusion_input_dim)
        self.fusion_w1 = nn.Linear(fusion_input_dim, self.fusion_dim1)
        self.fusion_g1 = nn.Linear(fusion_input_dim, self.fusion_dim1)
        self.fusion_drop = nn.Dropout(dropout)
        self.fusion_norm2 = nn.LayerNorm(self.fusion_dim1) if norm_type == 'layer_norm' else getattr(nn, 'RMSNorm', nn.LayerNorm)(self.fusion_dim1)
        self.fusion_w2 = nn.Linear(self.fusion_dim1, hidden_dim)
        self.fusion_g2 = nn.Linear(self.fusion_dim1, hidden_dim)
        self.fusion_out = nn.Linear(hidden_dim, output_dim)

        # Lightweight penetration head for vertex-level sign prediction: +1 (outside)
        # vs -1 (inside) for all vertices, computed from the last-layer object
        # feature, a global state feature, and the vertex position.
        # Two 32-d projections + a 16-d position projection -> 1-d tanh score.
        self._pen_obj_proj = nn.Linear(state_dim, 32)
        self._pen_state_proj = nn.Linear(state_dim, 32)
        self._pen_pos_proj = nn.Linear(3, 16)
        self._pen_head_norm = nn.LayerNorm(32 + 32 + 16) if norm_type == 'layer_norm' else getattr(nn, 'RMSNorm', nn.LayerNorm)(32 + 32 + 16)
        self._pen_head = nn.Linear(32 + 32 + 16, 1)
        
    # torch.nn.RMSNorm built without an explicit eps uses torch.finfo(input.dtype).eps: 0.0078125 for the bf16 inputs
    # four module groups receive under autocast in training, 1.19e-7 for fp32 inputs, and the evaluation rollouts run
    # without autocast. pin_rmsnorm_eps() gives every module the eps of the dtype it receives in bf16-autocast
    # training (input dtypes measured in bf16-autocast training), so evaluation computes the trained function;
    # bf16-autocast training itself is unchanged.
    _RMSNORM_BF16_TRAINING_INPUT = ('query_in_norms', 'kv_in_norms', 'post_attn_norms', 'fusion_norm2')
    _RMSNORM_FP32_TRAINING_INPUT = ('self_attn_in_norms', 'self_attn_post_norms', 'cross_obj_in_norms',
                                    'cross_obj_post_norms', 'fusion_norm1', '_pen_head_norm')

    def pin_rmsnorm_eps(self) -> int:
        bf16, fp32 = torch.finfo(torch.bfloat16).eps, torch.finfo(torch.float32).eps
        n = 0
        for name, m in self.named_modules():
            if not isinstance(m, nn.RMSNorm):
                continue
            top = name.split('.')[0]
            if top in self._RMSNORM_BF16_TRAINING_INPUT:
                m.eps = bf16
            elif top in self._RMSNORM_FP32_TRAINING_INPUT:
                m.eps = fp32
            else:
                raise ValueError(f"pin_rmsnorm_eps: training input dtype of RMSNorm '{name}' was not measured")
            n += 1
        return n

    def forward(
        self,
        multi_scale_features: List[torch.Tensor],
        query: torch.Tensor,
        ref_pose: Optional[torch.Tensor] = None,
        query_pos: Optional[torch.Tensor] = None,      # Anchor 3D positions for RoPE
        key_pos: Optional[torch.Tensor] = None,        # Object anchor positions for RoPE
        multi_scale_state_features: Optional[List[torch.Tensor]] = None,  # State features from decoder
        state_pos: Optional[torch.Tensor] = None,      # State 3D grid positions for RoPE
        obj_mask: Optional[torch.Tensor] = None        # (B, Nobj) bool mask for valid objects
    ) -> torch.Tensor:
        """
        Query-based prediction using cross-attention to multi-scale object and state features.

        Supports two input formats:
        1. Grouped: query (B, Nobj, Nq, feature_dim) - preferred for per-object attention
        2. Flat: query (B, M, feature_dim) where M=Nobj*Nq

        Args:
            multi_scale_features: List of object features from decoder layers (B, Nobj, D)
            query: Query features
                - Grouped: (B, Nobj, Nq, 9/12) - dist_vec + velocity + rel_pos
                - Flat: (B, Nobj*Nq, 9/12)
            ref_pose: Reference pose
                - Grouped: (B, Nobj, Nq, 3)
                - Flat: (B, Nobj*Nq, 3)
            query_pos: Anchor 3D positions for RoPE
                - Grouped: (B, Nobj, Nq, 3)
                - Flat: (B, Nobj*Nq, 3)
            key_pos: Object anchor positions (B, Nobj, Nq*3) for multi-point RoPE
                     (RenderFormer-style: flattened multi-point representation)
            multi_scale_state_features: List of state features from decoder layers (B, N_state, D)
            state_pos: State 3D grid positions (B, N_state, 3) for RoPE

        Returns:
            Predictions matching input format:
                - Grouped input -> (B, Nobj, Nq, output_dim)
                - Flat input -> (B, Nobj*Nq, output_dim)
        """
        # Detect input format
        is_grouped = (query.dim() == 4)  # (B, Nobj, Nq, D) vs (B, M, D)

        if is_grouped:
            B, Nobj, Nq, query_dim = query.shape
            M = Nobj * Nq
        else:
            B, M, query_dim = query.shape
            # Infer Nobj from key_pos or from the first available feature layer
            if key_pos is not None:
                Nobj = key_pos.shape[1]
            else:
                # Safely access the first available layer in the saved features
                # (layer 0 might not be saved with memory optimization)
                Nobj = multi_scale_features[0].shape[1] if len(multi_scale_features) > 0 else 1
            Nq = M // Nobj
            assert M == Nobj * Nq, f"M={M} must be divisible by Nobj={Nobj}"

        # Select object features from specified layers using the mapping
        # The mapping converts from original layer index to position in saved list
        selected_features = [multi_scale_features[self.layer_index_mapping[i]] for i in self.layer_indices]

        # Select state features from specified layers (if provided)
        selected_state_features = None
        if multi_scale_state_features is not None:
            selected_state_features = [multi_scale_state_features[self.layer_index_mapping[i]] for i in self.layer_indices]

        # Convert to grouped format for processing
        if is_grouped:
            # Already grouped
            query_grouped = query  # (B, Nobj, Nq, query_dim)
            ref_pose_grouped = ref_pose  # (B, Nobj, Nq, 3)
            query_pos_grouped = query_pos if query_pos is not None else None  # (B, Nobj, Nq, 3)
        else:
            # Reshape flat to grouped
            query_grouped = query.view(B, Nobj, Nq, query_dim)  # (B, Nobj, Nq, query_dim)
            if ref_pose is not None:
                ref_pose_grouped = ref_pose.view(B, Nobj, Nq, 3)
            if query_pos is not None:
                query_pos_grouped = query_pos.view(B, Nobj, Nq, 3)

        # Prepare query with reference pose
        assert ref_pose_grouped is not None
        query_input_grouped = torch.cat([query_grouped, ref_pose_grouped], dim=-1)  # (B, Nobj, Nq, query_dim+3)

        # Reshape for batched per-object attention: (B*Nobj, Nq, query_dim+3)
        query_input = query_input_grouped.view(B * Nobj, Nq, -1)

        # Validate query dimensionality against the projection layer
        import logging
        logger = logging.getLogger(__name__)
        expected_dim = self.query_proj.in_features
        actual_dim = query_input.shape[-1]
        if expected_dim != actual_dim:
            logger.error(f"Dimension mismatch: query_proj expects {expected_dim}D input, got {actual_dim}D")
            logger.error(f"  query shape: {query.shape}, ref_pose shape: {ref_pose.shape if ref_pose is not None else None}")
            raise RuntimeError(f"Predictor input dimension mismatch: expected {expected_dim}, got {actual_dim}")

        # Project query: (B*Nobj, Nq, query_dim+3) -> (B*Nobj, Nq, hidden_dim)
        query_proj = self.query_proj(query_input)  # (B*Nobj, Nq, hidden_dim)

        # Process each scale with cross-attention (with RoPE if enabled)
        # For batched per-object attention, we reshape features to (B*Nobj, 1, hidden_dim)
        # so each object's features can be attended by its Nq anchors
        # When cross_object_kv=True, each anchor can attend to ALL objects' features (B*Nobj, Nobj, hidden_dim)
        attended_features = []
        attended_state_features = []

        # Build attention masks for cross-object attention from obj_mask
        kv_mask_for_cross_obj = None  # For cross-object KV attention
        anchor_mask_for_global = None  # For global anchor self-attention
        if obj_mask is not None and (self.cross_object_kv or self.cross_object_attention):
            # kv mask for cross-object KV: each anchor group may only attend to
            # valid objects' features; obj_mask (B, Nobj) expands to all anchor groups
            kv_mask_expanded = obj_mask.unsqueeze(1).expand(B, Nobj, Nobj)  # (B, Nobj, Nobj)
            kv_mask_expanded = kv_mask_expanded.reshape(B * Nobj, Nobj)  # (B*Nobj, Nobj)
            # SDPA additive format: (B*Nobj, 1, Nq, Nobj); True=attend becomes 0.0
            kv_mask_for_cross_obj = kv_mask_expanded.unsqueeze(1).unsqueeze(1).expand(-1, 1, Nq, -1)  # (B*Nobj, 1, Nq, Nobj)
            kv_mask_for_cross_obj = torch.where(kv_mask_for_cross_obj, 0.0, float('-inf'))  # SDPA additive mask

            # Mask for global anchor self-attention: expand obj_mask to per-anchor
            # (B, Nobj) -> (B, Nobj*Nq)
            anchor_mask_flat = obj_mask.unsqueeze(-1).expand(B, Nobj, Nq).reshape(B, Nobj * Nq)  # (B, Nobj*Nq)
            # SDPA mask: (B, 1, Nobj*Nq, Nobj*Nq); only valid anchors attend to valid anchors
            anchor_mask_2d = anchor_mask_flat.unsqueeze(1) & anchor_mask_flat.unsqueeze(2)  # (B, Nobj*Nq, Nobj*Nq)
            anchor_mask_for_global = torch.where(anchor_mask_2d.unsqueeze(1), 0.0, float('-inf'))  # (B, 1, Nobj*Nq, Nobj*Nq)

            if not self._p0_mask_logged:
                n_valid = obj_mask.sum(dim=1).float().mean().item()
                n_pad = Nobj - n_valid
                print(f"[Predictor] Padding mask active: {n_valid:.0f} valid / {n_pad:.0f} padding objects, "
                      f"{n_pad*Nq:.0f} padded anchors masked in attention")
                self._p0_mask_logged = True

        for i, feat in enumerate(selected_features):
            # Project features: (B, Nobj, hidden_dim)
            feat_proj = self.scale_projections[i](feat)  # (B, Nobj, hidden_dim)

            if self.cross_object_kv:
                # === Cross-Object Key-Value Mode ===
                # Each anchor can attend to ALL objects' features
                # Expand feat_proj for all objects: (B, Nobj, D) -> (B*Nobj, Nobj, D)
                # Each of the B*Nobj anchor-groups sees all Nobj object features
                kv_batched = feat_proj.unsqueeze(1).expand(B, Nobj, Nobj, -1).reshape(B * Nobj, Nobj, -1)
                # kv_batched: (B*Nobj, Nobj, hidden_dim)

                # One-time runtime confirmation log
                if not self._cross_object_kv_logged:
                    print(f"[Predictor Runtime] Cross-object KV active: each anchor queries {Nobj} objects' features")
                    self._cross_object_kv_logged = True
            else:
                # === Standard Per-Object Mode ===
                # Each object has 1 feature vector that Nq anchors will attend to
                kv_batched = feat_proj.view(B * Nobj, 1, -1)  # (B*Nobj, 1, hidden_dim)

            # Norm-first with optional qk-norm
            q_in = self.query_in_norms[i](query_proj)  # (B*Nobj, Nq, hidden_dim)
            kv_in = self.kv_in_norms[i](kv_batched) if self._use_qk_norm else kv_batched

            # Per-object attention: Nq anchors attend to object feature(s)
            if self.use_rope:
                # Use RoPE-aware cross-attention
                if query_pos_grouped is None or key_pos is None:
                    raise ValueError(
                        "QueryBasedMultiScalePredictor with use_rope=True requires "
                        "query_pos (anchor positions) and key_pos (object positions)"
                    )
                # Reshape positions for batched per-object
                # query_pos: individual anchor positions (B, Nobj, Nq, 3) -> (B*Nobj, Nq, 3)
                query_pos_batched = query_pos_grouped.view(B * Nobj, Nq, 3)  # (B*Nobj, Nq, 3)

                if self.cross_object_kv:
                    # Key positions for ALL objects: (B, Nobj, Nq*3) -> (B*Nobj, Nobj, Nq*3)
                    # Each anchor-group sees all objects' anchor positions
                    key_pos_batched = key_pos.unsqueeze(1).expand(B, Nobj, Nobj, -1).reshape(B * Nobj, Nobj, -1)
                    # key_pos_batched: (B*Nobj, Nobj, Nq*3)
                else:
                    # key_pos: flattened multi-point anchors (B, Nobj, Nq*3) -> (B*Nobj, 1, Nq*3)
                    key_pos_batched = key_pos.view(B * Nobj, 1, -1)  # (B*Nobj, 1, Nq*3) for multi-point RoPE

                attn_out = self.cross_attentions[i](
                    query=q_in,                 # (B*Nobj, Nq, hidden_dim)
                    key_value=kv_in,            # (B*Nobj, 1 or Nobj, hidden_dim)
                    pos_q=query_pos_batched,    # (B*Nobj, Nq, 3) - individual anchor positions
                    pos_k=key_pos_batched,      # (B*Nobj, 1 or Nobj, Nq*3) - multi-point anchors
                    attn_mask=kv_mask_for_cross_obj if self.cross_object_kv else None
                )  # (B*Nobj, Nq, hidden_dim)
            else:
                # Standard attention
                attn_out, _ = self.cross_attentions[i](
                    query=q_in,
                    key=kv_in,
                    value=kv_in
                )  # (B*Nobj, Nq, hidden_dim)

            # Residual + post norm
            attn_out = self.post_attn_norms[i](attn_out + query_proj)  # (B*Nobj, Nq, hidden_dim)
            attended_features.append(attn_out)

            # === Cross-attention to state features (skipped when no_global_state) ===
            if not self.no_global_state and selected_state_features is not None:
                state_feat = selected_state_features[i]  # (B, N_state, D)
                state_feat_proj = self.state_scale_projections[i](state_feat)  # (B, N_state, hidden_dim)

                # Expand state features for all anchors across all objects
                # state_feat_proj: (B, N_state, hidden_dim) -> (B*Nobj, N_state, hidden_dim)
                state_feat_expanded = state_feat_proj.unsqueeze(1).expand(B, Nobj, -1, -1).reshape(B * Nobj, -1, self.hidden_dim)

                # Norm
                state_kv_in = self.state_kv_in_norms[i](state_feat_expanded) if self._use_qk_norm else state_feat_expanded

                # Cross-attention: anchors attend to state
                if self.use_rope:
                    if state_pos is None:
                        raise ValueError("state_pos is required when using RoPE for state attention")

                    # Expand state_pos for all objects: (B, N_state, 3) -> (B*Nobj, N_state, 3)
                    state_pos_expanded = state_pos.unsqueeze(1).expand(B, Nobj, -1, -1).reshape(B * Nobj, -1, 3)

                    state_attn_out = self.state_cross_attentions[i](
                        query=q_in,                     # (B*Nobj, Nq, hidden_dim)
                        key_value=state_kv_in,          # (B*Nobj, N_state, hidden_dim)
                        pos_q=query_pos_batched,        # (B*Nobj, Nq, 3) - anchor positions
                        pos_k=state_pos_expanded,       # (B*Nobj, N_state, 3) - state grid positions
                        attn_mask=None
                    )  # (B*Nobj, Nq, hidden_dim)
                else:
                    state_attn_out, _ = self.state_cross_attentions[i](
                        query=q_in,
                        key=state_kv_in,
                        value=state_kv_in
                    )

                # Residual + post norm
                state_attn_out = self.state_post_attn_norms[i](state_attn_out + query_proj)
                attended_state_features.append(state_attn_out)
            else:
                if not self.no_global_state:
                    # Without a global state input: append zeros to keep fusion dimensions consistent
                    attended_state_features.append(torch.zeros_like(attended_features[-1]))
                # With no_global_state, attended_state_features stays empty

            # === Self-attention among anchors within each object ===
            # With no_global_state, only object features feed the self-attention
            if self.no_global_state or len(attended_state_features) <= i:
                sa_input = attended_features[-1]  # (B*Nobj, Nq, hidden_dim), object features only
            else:
                sa_input = attended_features[-1] + attended_state_features[-1]  # (B*Nobj, Nq, hidden_dim)

            # Norm-first
            sa_input_normed = self.self_attn_in_norms[i](sa_input)  # (B*Nobj, Nq, hidden_dim)

            # Self-attention: anchors attend to each other within the same object
            if self.use_rope:
                sa_out = self.self_attentions[i](
                    query=sa_input_normed,          # (B*Nobj, Nq, hidden_dim)
                    key_value=sa_input_normed,      # (B*Nobj, Nq, hidden_dim)
                    pos_q=query_pos_batched,        # (B*Nobj, Nq, 3) - anchor positions
                    pos_k=query_pos_batched,        # (B*Nobj, Nq, 3) - same positions for self-attn
                    attn_mask=None
                )  # (B*Nobj, Nq, hidden_dim)
            else:
                sa_out, _ = self.self_attentions[i](
                    query=sa_input_normed,
                    key=sa_input_normed,
                    value=sa_input_normed
                )  # (B*Nobj, Nq, hidden_dim)

            # Residual + post norm
            sa_out = self.self_attn_post_norms[i](sa_out + sa_input)  # (B*Nobj, Nq, hidden_dim)

            # === Cross-object anchor self-attention ===
            # Allows anchors from different objects to interact directly
            if self.cross_object_attention:
                # One-time runtime confirmation log
                if not self._cross_object_attention_logged:
                    print(f"[Predictor Runtime] Cross-object anchor attention active: B={B}, Nobj={Nobj}, Nq={Nq}, total_anchors={Nobj*Nq}")
                    self._cross_object_attention_logged = True
                # Reshape from per-object to global: (B*Nobj, Nq, D) -> (B, Nobj*Nq, D)
                sa_out_global = sa_out.view(B, Nobj * Nq, -1)  # (B, Nobj*Nq, hidden_dim)

                # Reshape anchor positions: (B, Nobj, Nq, 3) -> (B, Nobj*Nq, 3)
                query_pos_global = query_pos_grouped.view(B, Nobj * Nq, 3)  # (B, Nobj*Nq, 3)

                # Norm-first
                cross_obj_input = self.cross_obj_in_norms[i](sa_out_global)  # (B, Nobj*Nq, hidden_dim)

                # Cross-object self-attention: all anchors attend to all anchors
                if self.use_rope:
                    cross_obj_out = self.cross_obj_self_attentions[i](
                        query=cross_obj_input,          # (B, Nobj*Nq, hidden_dim)
                        key_value=cross_obj_input,      # (B, Nobj*Nq, hidden_dim)
                        pos_q=query_pos_global,         # (B, Nobj*Nq, 3) - all anchor positions
                        pos_k=query_pos_global,         # (B, Nobj*Nq, 3) - same for self-attn
                        attn_mask=anchor_mask_for_global  # Mask padding anchors
                    )  # (B, Nobj*Nq, hidden_dim)
                else:
                    cross_obj_out, _ = self.cross_obj_self_attentions[i](
                        query=cross_obj_input,
                        key=cross_obj_input,
                        value=cross_obj_input
                    )  # (B, Nobj*Nq, hidden_dim)

                # Residual + post norm
                cross_obj_out = self.cross_obj_post_norms[i](cross_obj_out + sa_out_global)  # (B, Nobj*Nq, hidden_dim)

                # Reshape back to per-object: (B, Nobj*Nq, D) -> (B*Nobj, Nq, D)
                sa_out = cross_obj_out.view(B * Nobj, Nq, -1)  # (B*Nobj, Nq, hidden_dim)

            # Replace this scale's object features with the enhanced version
            attended_features[-1] = sa_out
        
        # Concatenate multi-scale features for fusion
        # With no_global_state, only object features are concatenated
        if self.no_global_state:
            multi_scale_output = torch.cat(attended_features, dim=-1)
            # (B*Nobj, Nq, hidden_dim * num_scales)
            if not self._p1_no_state_logged:
                print(f"[Predictor] Fusion uses object features only: {multi_scale_output.shape[-1]}D")
                self._p1_no_state_logged = True
        else:
            multi_scale_output = torch.cat(attended_features + attended_state_features, dim=-1)
            # (B*Nobj, Nq, hidden_dim * num_scales * 2)

        # Add reference pose to fusion if enabled
        if self.use_ref_pose:
            # ref_pose_grouped: (B, Nobj, Nq, 3) -> (B*Nobj, Nq, 3)
            ref_pose_batched = ref_pose_grouped.view(B * Nobj, Nq, 3)
            ref_pose_features = self.ref_pose_proj(ref_pose_batched)  # (B*Nobj, Nq, hidden_dim)
            multi_scale_output = torch.cat([multi_scale_output, ref_pose_features], dim=-1)
            # (B*Nobj, Nq, hidden_dim * num_scales + hidden_dim)
        
        # Final fusion and prediction (SwiGLU blocks)
        # All operations are per-anchor: (B*Nobj, Nq, dim)
        # Block 1
        x = self.fusion_norm1(multi_scale_output)
        if hasattr(torch.nn.functional, 'silu'):
            x = self.fusion_drop(torch.nn.functional.silu(self.fusion_g1(x)) * self.fusion_w1(x))
        else:
            x = self.fusion_drop(torch.sigmoid(self.fusion_g1(x)) * self.fusion_w1(x))
        # Block 2
        x = self.fusion_norm2(x)
        if hasattr(torch.nn.functional, 'silu'):
            x = self.fusion_drop(torch.nn.functional.silu(self.fusion_g2(x)) * self.fusion_w2(x))
        else:
            x = self.fusion_drop(torch.sigmoid(self.fusion_g2(x)) * self.fusion_w2(x))
        if getattr(self, 'fp32_output', False):
            # model.predictor.fp32_output: the acceleration read-out in fp32 also under bf16 autocast (training), as
            # in the fp32 evaluation rollout
            with torch.autocast(device_type=x.device.type, enabled=False):
                predictions = self.fusion_out(x.float())  # (B*Nobj, Nq, output_dim)
        else:
            predictions = self.fusion_out(x)  # (B*Nobj, Nq, output_dim)

        # Reshape output to match input format
        if is_grouped:
            # Return grouped format: (B, Nobj, Nq, output_dim)
            predictions = predictions.view(B, Nobj, Nq, -1)
        else:
            # Return flat format: (B, Nobj*Nq, output_dim)
            predictions = predictions.view(B, M, -1)

        return predictions

    def predict_penetration_vertices(
        self,
        multi_scale_features: List[torch.Tensor],
        multi_scale_state_features: Optional[List[torch.Tensor]],
        pos_full: torch.Tensor  # (B, Nobj, Nv, 3)
    ) -> torch.Tensor:
        """
        Predict penetration sign for ALL vertices (vertex-level supervision).

        Returns penetration score in [-1, 1] where:
        - +1 indicates vertex is outside other objects (no penetration)
        - -1 indicates vertex is inside other objects (penetration)

        Args:
            multi_scale_features: List of object features from decoder layers (B, Nobj, D)
            multi_scale_state_features: List of state features from decoder layers (B, N_state, D) or None
            pos_full: Vertex positions (B, Nobj, Nv, 3) - AFTER Kabsch transform

        Returns:
            Penetration scores (B, Nobj, Nv, 1) with tanh activation, requires_grad=True for training
        """
        # Use layer_index_mapping to get the correct position in the saved features list
        last_layer_idx = self.layer_indices[-1]
        mapped_idx = self.layer_index_mapping[last_layer_idx]

        obj_last = multi_scale_features[mapped_idx]  # (B, Nobj, D)
        B, Nobj, _ = obj_last.shape
        obj_proj = self._pen_obj_proj(obj_last)                   # (B, Nobj, 32)

        if multi_scale_state_features is not None and len(multi_scale_state_features) > 0:
            state_last = multi_scale_state_features[mapped_idx]  # (B, N_state, D)
            state_global = state_last.mean(dim=1)                            # (B, D)
        else:
            state_global = obj_last.mean(dim=1) * 0.0                        # (B, D)
        state_proj = self._pen_state_proj(state_global)                      # (B, 32)
        state_proj = state_proj.unsqueeze(1).expand(-1, Nobj, -1)            # (B, Nobj, 32)

        Nv = pos_full.shape[2]
        pos_proj = self._pen_pos_proj(pos_full)                              # (B, Nobj, Nv, 16)
        obj_proj_v = obj_proj.unsqueeze(2).expand(-1, -1, Nv, -1)            # (B, Nobj, Nv, 32)
        state_proj_v = state_proj.unsqueeze(2).expand(-1, -1, Nv, -1)        # (B, Nobj, Nv, 32)

        pen_feat = torch.cat([obj_proj_v, state_proj_v, pos_proj], dim=-1)   # (B, Nobj, Nv, 80)
        pen_feat = self._pen_head_norm(pen_feat)
        pen = torch.tanh(self._pen_head(pen_feat))                           # (B, Nobj, Nv, 1)
        return pen