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

class VoxelCNNEncoder(nn.Module):
    """
    Enhanced 3D CNN encoder for voxelized obstruction data with multi-channel support.
    Processes environment obstacles, start position, and goal position.
    """
    def __init__(self,
                 input_channels=3,  # obstacles + start + goal
                 filters_1=32,
                 kernel_size_1=(3, 3, 3),
                 pool_size_1=(2, 2, 2),
                 filters_2=64,
                 kernel_size_2=(3, 3, 3),
                 pool_size_2=(2, 2, 2),
                 filters_3=128,
                 kernel_size_3=(3, 3, 3),
                 pool_size_3=(2, 2, 2),
                 dense_units=512,
                 input_voxel_dim=(32, 32, 32),
                 dropout_rate=0.2
                ):
        super(VoxelCNNEncoder, self).__init__()

        self.input_voxel_dim = input_voxel_dim
        self.input_channels = input_channels

        # First 3D Convolutional Block (Conv-BN-ReLU)
        padding_1 = tuple([(k - 1) // 2 for k in kernel_size_1])
        self.conv1 = nn.Conv3d(input_channels, filters_1, kernel_size_1, padding=padding_1)
        self.bn1 = nn.BatchNorm3d(filters_1)
        self.pool1 = nn.MaxPool3d(pool_size_1)
        self.dropout1 = nn.Dropout3d(dropout_rate)

        # Second 3D Convolutional Block (Conv-BN-ReLU)
        padding_2 = tuple([(k - 1) // 2 for k in kernel_size_2])
        self.conv2 = nn.Conv3d(filters_1, filters_2, kernel_size_2, padding=padding_2)
        self.bn2 = nn.BatchNorm3d(filters_2)
        self.pool2 = nn.MaxPool3d(pool_size_2)
        self.dropout2 = nn.Dropout3d(dropout_rate)

        # Third 3D Convolutional Block (Conv-BN-ReLU)
        padding_3 = tuple([(k - 1) // 2 for k in kernel_size_3])
        self.conv3 = nn.Conv3d(filters_2, filters_3, kernel_size_3, padding=padding_3)
        self.bn3 = nn.BatchNorm3d(filters_3)
        self.pool3 = nn.MaxPool3d(pool_size_3)
        self.dropout3 = nn.Dropout3d(dropout_rate)

        # Calculate flattened size
        self._to_linear_input_size = self._get_conv_output_size()

        # Dense layers with residual connection
        self.fc1 = nn.Linear(self._to_linear_input_size, dense_units)
        self.fc2 = nn.Linear(dense_units, dense_units)
        self.dropout_fc = nn.Dropout(dropout_rate)

    def _get_conv_output_size(self):
        with torch.no_grad():
            dummy_input = torch.zeros(1, self.input_channels, *self.input_voxel_dim)
            # Standardized Conv-BN-ReLU order
            x = self.conv1(dummy_input)
            x = self.bn1(x)
            x = F.relu(x)
            x = self.pool1(x)
            x = self.dropout1(x)
            
            x = self.conv2(x)
            x = self.bn2(x)
            x = F.relu(x)
            x = self.pool2(x)
            x = self.dropout2(x)
            
            x = self.conv3(x)
            x = self.bn3(x)
            x = F.relu(x)
            x = self.pool3(x)
            x = self.dropout3(x)
            
            return x.numel()

    def forward(self, x):
        # First conv block (Conv-BN-ReLU)
        x = self.conv1(x)
        x = self.bn1(x)
        x = F.relu(x)
        x = self.pool1(x)
        x = self.dropout1(x)

        # Second conv block (Conv-BN-ReLU)
        x = self.conv2(x)
        x = self.bn2(x)
        x = F.relu(x)
        x = self.pool2(x)
        x = self.dropout2(x)

        # Third conv block (Conv-BN-ReLU)
        x = self.conv3(x)
        x = self.bn3(x)
        x = F.relu(x)
        x = self.pool3(x)
        x = self.dropout3(x)

        # Flatten and dense layers
        x = x.view(x.size(0), -1)
        x1 = F.relu(self.fc1(x))
        x1 = self.dropout_fc(x1)
        x2 = F.relu(self.fc2(x1))
        
        # Residual connection
        return x1 + x2


class PositionEncoder(nn.Module):
    """
    Encodes start and goal positions with learned embeddings.
    """
    def __init__(self, voxel_dim=(32, 32, 32), position_embed_dim=64):
        super(PositionEncoder, self).__init__()
        self.voxel_dim = voxel_dim
        self.position_embed_dim = position_embed_dim
        
        # Calculate dimensions for each axis to sum to position_embed_dim
        dim_per_axis = position_embed_dim // 3
        remainder = position_embed_dim % 3
        
        x_dim = dim_per_axis + (1 if remainder > 0 else 0)
        y_dim = dim_per_axis + (1 if remainder > 1 else 0)
        z_dim = dim_per_axis
        
        # Learned position embeddings for each dimension
        self.x_embed = nn.Embedding(voxel_dim[0], x_dim)
        self.y_embed = nn.Embedding(voxel_dim[1], y_dim)
        self.z_embed = nn.Embedding(voxel_dim[2], z_dim)
        
        # Additional processing - fixed input dimension
        self.fc = nn.Linear(2 * position_embed_dim, position_embed_dim)
        
    def forward(self, positions):
        """
        positions: (batch_size, 2, 3) - [start_pos, goal_pos] with (x, y, z)
        """
        batch_size = positions.size(0)
        
        # Extract coordinates
        # Clamp coordinates defensively to valid index ranges to avoid embedding OOB
        x_coords = positions[:, :, 0].long().clamp_(0, self.voxel_dim[0] - 1)  # (batch_size, 2)
        y_coords = positions[:, :, 1].long().clamp_(0, self.voxel_dim[1] - 1)  # (batch_size, 2)
        z_coords = positions[:, :, 2].long().clamp_(0, self.voxel_dim[2] - 1)  # (batch_size, 2)
        
        # Get embeddings
        x_emb = self.x_embed(x_coords)  # (batch_size, 2, x_dim)
        y_emb = self.y_embed(y_coords)  # (batch_size, 2, y_dim)
        z_emb = self.z_embed(z_coords)  # (batch_size, 2, z_dim)
        
        # Concatenate embeddings
        pos_emb = torch.cat([x_emb, y_emb, z_emb], dim=-1)  # (batch_size, 2, position_embed_dim)
        
        # Flatten start and goal embeddings
        pos_emb = pos_emb.view(batch_size, -1)  # (batch_size, 2 * position_embed_dim)
        
        return F.relu(self.fc(pos_emb))


class PathPlannerTransformer(nn.Module):
    """
    Transformer-based path planner that generates action sequences.
    Fixed token IDs to avoid collision:
    - Actions: 0-5 (Forward, Back, Left, Right, Up, Down)
    - START: 6
    - END: 7
    - PAD: 8 (used only for teacher forcing inputs; targets still use -1 for ignore)
    """
    def __init__(self, 
                 env_feature_dim=512,
                 pos_feature_dim=64,
                 hidden_dim=256,
                 num_heads=8,
                 num_layers=4,
                 max_sequence_length=100,
                 num_actions=6,  # Forward, Back, Left, Right, Up, Down
                 use_end_token=True):
        super(PathPlannerTransformer, self).__init__()
        
        self.hidden_dim = hidden_dim
        self.max_sequence_length = max_sequence_length
        self.num_actions = num_actions
        self.use_end_token = use_end_token
        
        # Fixed token IDs to avoid collision
        self.start_token_id = num_actions  # 6
        self.end_token_id = num_actions + 1 if use_end_token else None  # 7
        # Reserve a PAD token for embedding inputs during teacher forcing
        self.pad_token_id = (num_actions + 2) if use_end_token else (num_actions + 1)
        # Total tokens include PAD
        self.total_tokens = (num_actions + 3) if use_end_token else (num_actions + 2)
        
        # Feature fusion
        self.feature_fusion = nn.Linear(env_feature_dim + pos_feature_dim, hidden_dim)
        
        # Action embeddings
        self.action_embed = nn.Embedding(self.total_tokens, hidden_dim)
        
        # Positional encoding - register as buffer for proper device handling
        self.register_buffer('pos_encoding', self._create_positional_encoding(max_sequence_length, hidden_dim))
        
        # Transformer decoder
        decoder_layer = nn.TransformerDecoderLayer(
            d_model=hidden_dim,
            nhead=num_heads,
            dim_feedforward=hidden_dim * 4,
            dropout=0.1,
            batch_first=True
        )
        self.transformer_decoder = nn.TransformerDecoder(decoder_layer, num_layers)
        
        # Output projection
        self.output_proj = nn.Linear(hidden_dim, self.total_tokens)
        
        # Turn head (now supervised via BCE-with-logits against turn labels)
        self.turn_penalty_head = nn.Linear(hidden_dim, 1)
        
    def _create_positional_encoding(self, max_len, d_model):
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1).float()
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * -(np.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        return pe.unsqueeze(0)
    
    def forward(self, env_features, pos_features, target_actions=None):
        """
        env_features: (batch_size, env_feature_dim)
        pos_features: (batch_size, pos_feature_dim)
        target_actions: (batch_size, seq_len) - for training (contains action IDs 0-5 and END token 7)
        """
        batch_size = env_features.size(0)
        
        # Fuse environment and position features
        fused_features = self.feature_fusion(torch.cat([env_features, pos_features], dim=1))
        
        # Create memory (encoder output) by repeating fused features
        memory = fused_features.unsqueeze(1).repeat(1, self.max_sequence_length, 1)
        
        if target_actions is not None:
            # Training mode: use teacher forcing
            seq_len = target_actions.size(1)
            
            # Create input sequence (START token + target_actions[:-1])
            start_tokens = torch.full((batch_size, 1), self.start_token_id, 
                                    dtype=torch.long, device=target_actions.device)
            input_seq = torch.cat([start_tokens, target_actions[:, :-1]], dim=1)
            # Replace padding (-1) in teacher-forced inputs with PAD token id to avoid OOB in embedding
            input_seq = torch.where(input_seq < 0, torch.full_like(input_seq, self.pad_token_id), input_seq)
            
            # Embed actions and add positional encoding
            embedded = self.action_embed(input_seq)
            embedded = embedded + self.pos_encoding[:, :seq_len, :]
            
            # Generate attention mask (causal mask)
            tgt_mask = self._generate_square_subsequent_mask(seq_len).to(embedded.device)
            
            # Transformer decoder forward pass
            output = self.transformer_decoder(
                tgt=embedded,
                memory=memory[:, :seq_len, :],
                tgt_mask=tgt_mask
            )
            
            # Output projections
            action_logits = self.output_proj(output)
            # Turn logits for supervised turn classification
            turn_logits = self.turn_penalty_head(output)
            
            return action_logits, turn_logits
        else:
            # Inference mode: generate sequence autoregressively
            return self._generate_path(memory, batch_size)
    
    def _generate_square_subsequent_mask(self, sz):
        mask = torch.triu(torch.ones(sz, sz), diagonal=1)
        mask = mask.masked_fill(mask == 1, float('-inf'))
        return mask
    
    def _generate_path(self, memory, batch_size):
        """
        Generate path sequence autoregressively, handling batches correctly.
        Fixes bugs related to premature stopping and inclusion of special tokens.
        """
        device = memory.device
        
        # Start with START token
        input_seq = torch.full((batch_size, 1), self.start_token_id, dtype=torch.long, device=device)
        
        # Keep track of sequences that have generated an END token
        is_finished = torch.zeros(batch_size, dtype=torch.bool, device=device)

        for step in range(self.max_sequence_length):
            # Embed current sequence
            embedded = self.action_embed(input_seq)
            seq_len = embedded.size(1)
            embedded = embedded + self.pos_encoding[:, :seq_len, :]

            # Generate causal mask
            tgt_mask = self._generate_square_subsequent_mask(seq_len).to(device)

            # Forward pass
            output = self.transformer_decoder(
                tgt=embedded,
                memory=memory[:, :seq_len, :],
                tgt_mask=tgt_mask
            )
            
            # Get next action probabilities from the last token in the sequence
            next_action_logits = self.output_proj(output[:, -1, :])
            next_actions = torch.argmax(next_action_logits, dim=-1, keepdim=True)
            
            # Append the predicted actions to the sequence
            input_seq = torch.cat([input_seq, next_actions], dim=1)
            
            # Update the finished mask for any sequence that just produced an END token
            if self.use_end_token:
                is_finished |= (next_actions.squeeze(-1) == self.end_token_id)
            
            # If all sequences in the batch are finished, we can stop early
            if is_finished.all():
                break
        
        # Post-processing to create a clean, dense tensor of valid actions
        # Remove the initial START token from all sequences
        raw_paths = input_seq[:, 1:]
        
        clean_paths_list = []
        max_len = 0
        
        for i in range(batch_size):
            path = []
            for token_id in raw_paths[i]:
                # Stop decoding for this path if an END token is found
                if self.use_end_token and token_id.item() == self.end_token_id:
                    break
                # Only include valid movement actions in the final path
                if token_id.item() < self.num_actions:
                    path.append(token_id.item())
            
            clean_paths_list.append(path)
            if len(path) > max_len:
                max_len = len(path)
        
        # Return an empty tensor if no valid actions were generated
        if max_len == 0:
            return torch.zeros(batch_size, 0, dtype=torch.long, device=device)
        
        # Pad all paths to the length of the longest path in the batch
        # We use the END token ID for padding, as downstream functions like
        # check_collisions are designed to ignore non-action tokens.
        pad_value = self.end_token_id if self.use_end_token else self.num_actions
        padded_paths = torch.full((batch_size, max_len), pad_value, dtype=torch.long, device=device)
        
        for i, path in enumerate(clean_paths_list):
            if len(path) > 0:
                padded_paths[i, :len(path)] = torch.tensor(path, dtype=torch.long, device=device)
                
        return padded_paths


class PathfindingNetwork(nn.Module):
    """
    Complete pathfinding network combining CNN encoder, position encoder, and transformer planner.
    """
    def __init__(self, 
                 voxel_dim=(32, 32, 32),
                 input_channels=3,
                 env_feature_dim=512,
                 pos_feature_dim=64,
                 hidden_dim=256,
                 num_actions=6,
                 use_end_token=True):
        super(PathfindingNetwork, self).__init__()
        
        self.voxel_dim = voxel_dim
        self.num_actions = num_actions
        
        self.voxel_encoder = VoxelCNNEncoder(
            input_channels=input_channels,
            dense_units=env_feature_dim,
            input_voxel_dim=voxel_dim
        )
        
        self.position_encoder = PositionEncoder(
            voxel_dim=voxel_dim,
            position_embed_dim=pos_feature_dim
        )
        
        self.path_planner = PathPlannerTransformer(
            env_feature_dim=env_feature_dim,
            pos_feature_dim=pos_feature_dim,
            hidden_dim=hidden_dim,
            num_actions=num_actions,
            use_end_token=use_end_token
        )
        
    def forward(self, voxel_data, positions, target_actions=None):
        """
        voxel_data: (batch_size, 3, D, H, W) - [obstacles, start_mask, goal_mask]
        positions: (batch_size, 2, 3) - [start_pos, goal_pos]
        target_actions: (batch_size, seq_len) - for training
        """
        # Encode environment
        env_features = self.voxel_encoder(voxel_data)
        
        # Encode positions
        pos_features = self.position_encoder(positions)
        
        # Generate path
        if target_actions is not None:
            action_logits, turn_penalties = self.path_planner(env_features, pos_features, target_actions)
            return action_logits, turn_penalties
        else:
            generated_path = self.path_planner(env_features, pos_features)
            return generated_path
    
    def check_collisions(self, voxel_data, positions, actions):
        """
        Check if actions lead to collisions with obstacles.
        
        voxel_data: (batch_size, 3, D, H, W)
        positions: (batch_size, 2, 3) - start positions
        actions: (batch_size, seq_len) - action sequences
        
        Returns: (batch_size, seq_len) collision mask
        """
        batch_size, seq_len = actions.shape
        device = actions.device
        
        # Extract obstacle channel
        obstacles = voxel_data[:, 0, :, :, :]  # (batch_size, D, H, W)
        
        # Action to direction mapping
        directions = torch.tensor([
            [1, 0, 0],   # Forward (z+)
            [-1, 0, 0],  # Back (z-)
            [0, 1, 0],   # Left (x+)
            [0, -1, 0],  # Right (x-)
            [0, 0, 1],   # Up (y+)
            [0, 0, -1]   # Down (y-)
        ], dtype=torch.long, device=device)
        
        collision_mask = torch.zeros(batch_size, seq_len, dtype=torch.bool, device=device)
        current_pos = positions[:, 0, :].clone()  # Start from start position
        
        for t in range(seq_len):
            # Get actions for this timestep
            action_t = actions[:, t]
            
            # Only process valid actions (0-5), skip special tokens
            valid_actions = action_t < self.num_actions
            
            # Update positions based on actions
            for b in range(batch_size):
                if valid_actions[b]:
                    direction = directions[action_t[b]]
                    new_pos = current_pos[b] + direction
                    
                    # Check bounds
                    if (new_pos >= 0).all() and (new_pos[0] < self.voxel_dim[0]) and \
                       (new_pos[1] < self.voxel_dim[1]) and (new_pos[2] < self.voxel_dim[2]):
                        # Check collision
                        if obstacles[b, new_pos[0], new_pos[1], new_pos[2]] > 0:
                            collision_mask[b, t] = True
                        else:
                            current_pos[b] = new_pos
                    else:
                        # Out of bounds counts as collision
                        collision_mask[b, t] = True
        
        return collision_mask


class PathfindingLoss(nn.Module):
    """
    Custom loss function that balances path correctness and turn minimization.
    Properly handles special tokens (START=6, END=7) and action tokens (0-5).
    Turn loss is supervised: a turn occurs when consecutive valid actions differ.
    """
    def __init__(self, turn_penalty_weight=0.1, collision_penalty_weight=10.0, 
                 num_actions=6, use_end_token=True):
        super(PathfindingLoss, self).__init__()
        self.turn_penalty_weight = turn_penalty_weight
        self.collision_penalty_weight = collision_penalty_weight
        self.num_actions = num_actions
        self.use_end_token = use_end_token
        self.start_token_id = num_actions  # 6
        self.end_token_id = num_actions + 1 if use_end_token else None  # 7
        self.ce_loss = nn.CrossEntropyLoss(ignore_index=-1)  # Ignore padding
        # BCE with logits for supervised turn prediction
        self.turn_bce = nn.BCEWithLogitsLoss(reduction='sum')
        
    def forward(self, action_logits, turn_penalties, target_actions, collision_mask=None):
        """
        action_logits: (batch_size, seq_len, total_tokens) - includes all tokens (0-7)
        turn_penalties: (batch_size, seq_len, 1) - interpreted as turn logits
        target_actions: (batch_size, seq_len) - contains action IDs (0-5) and possibly END (7)
        collision_mask: (batch_size, seq_len) - 1 if collision, 0 if safe
        """
        batch_size, seq_len, total_tokens = action_logits.shape
        
        # Reshape for cross entropy loss
        action_logits_flat = action_logits.view(-1, total_tokens)
        target_actions_flat = target_actions.view(-1)
        
        # Path correctness loss - now properly handles all token IDs
        path_loss = self.ce_loss(action_logits_flat, target_actions_flat)
        
        # Supervised turn loss
        # Compute valid action mask (exclude special tokens)
        valid_actions_mask = (target_actions < self.num_actions)
        # Previous actions (pad first timestep with itself; will be masked out anyway)
        prev_actions = torch.cat([target_actions[:, :1], target_actions[:, :-1]], dim=1)
        prev_valid_mask = torch.cat([torch.zeros_like(valid_actions_mask[:, :1], dtype=torch.bool),
                                     valid_actions_mask[:, :-1]], dim=1)
        # A turn occurs if both current and previous are valid actions and they differ
        both_valid = valid_actions_mask & prev_valid_mask
        is_turn = ((target_actions != prev_actions) & both_valid).float()
        
        # Turn logits predicted by the model
        turn_logits = turn_penalties.squeeze(-1)
        
        # Compute BCE-with-logits only over valid pairs
        num_pairs = both_valid.sum().clamp_min(1).float()
        if num_pairs > 0:
            bce_sum = self.turn_bce(turn_logits[both_valid], is_turn[both_valid])
            turn_loss = bce_sum / num_pairs
        else:
            turn_loss = torch.tensor(0.0, device=action_logits.device)
        
        # Collision penalty - only apply to actual movement actions
        collision_loss = torch.tensor(0.0, device=action_logits.device)
        if collision_mask is not None:
            # Mask collisions to only count for actual movement actions
            masked_collisions = collision_mask.float() * valid_actions_mask.float()
            if valid_actions_mask.sum() > 0:
                collision_loss = (masked_collisions.sum() / valid_actions_mask.sum()) * self.collision_penalty_weight
            
        total_loss = path_loss + self.turn_penalty_weight * turn_loss + collision_loss
        
        return {
            'total_loss': total_loss,
            'path_loss': path_loss,
            'turn_loss': turn_loss,
            'collision_loss': collision_loss
        }


# Utility functions for data preparation
def create_voxel_input(obstacles, start_pos, goal_pos, voxel_dim=(32, 32, 32)):
    """
    Create multi-channel voxel input.
    
    obstacles: (D, H, W) binary array
    start_pos: (x, y, z) tuple
    goal_pos: (x, y, z) tuple
    """
    # Channel 0: obstacles
    obstacle_channel = obstacles.astype(np.float32)
    
    # Channel 1: start position
    start_channel = np.zeros(voxel_dim, dtype=np.float32)
    start_channel[start_pos] = 1.0
    
    # Channel 2: goal position
    goal_channel = np.zeros(voxel_dim, dtype=np.float32)
    goal_channel[goal_pos] = 1.0
    
    # Stack channels
    voxel_input = np.stack([obstacle_channel, start_channel, goal_channel], axis=0)
    
    return voxel_input


def prepare_training_targets(action_sequence, use_end_token=True, num_actions=6):
    """
    Prepare target action sequences for training.
    Ensures action IDs are in range [0, num_actions-1] and adds END token if needed.
    
    action_sequence: list or tensor of action IDs (0-5)
    use_end_token: whether to append END token
    num_actions: number of valid actions
    
    Returns: tensor with proper token IDs
    """
    if isinstance(action_sequence, list):
        action_sequence = torch.tensor(action_sequence)
    
    # Ensure actions are in valid range
    assert (action_sequence >= 0).all() and (action_sequence < num_actions).all(), \
        f"Actions must be in range [0, {num_actions-1}]"
    
    if use_end_token:
        # Append END token (ID = num_actions + 1 = 7)
        end_token = torch.tensor([num_actions + 1])
        target = torch.cat([action_sequence, end_token])
    else:
        target = action_sequence
    
    return target


# Example usage and testing
if __name__ == "__main__":
    # Define problem parameters
    voxel_dim = (32, 32, 32)
    batch_size = 4
    num_actions = 6  # Forward, Back, Left, Right, Up, Down
    
    # Create the complete pathfinding network
    pathfinding_net = PathfindingNetwork(
        voxel_dim=voxel_dim,
        input_channels=3,
        env_feature_dim=512,
        pos_feature_dim=64,
        hidden_dim=256,
        num_actions=num_actions,
        use_end_token=True
    )
    
    print("=== 3D Pathfinding Network Architecture ===")
    print(f"Total parameters: {sum(p.numel() for p in pathfinding_net.parameters()):,}")
    print(f"\nToken ID mapping:")
    print(f"  Actions: 0-5 (Forward, Back, Left, Right, Up, Down)")
    print(f"  START token: {pathfinding_net.path_planner.start_token_id}")
    print(f"  END token: {pathfinding_net.path_planner.end_token_id}")
    
    # Create dummy data
    dummy_voxel_data = torch.randn(batch_size, 3, *voxel_dim)
    dummy_positions = torch.randint(0, 32, (batch_size, 2, 3))  # start and goal positions
    
    # Create proper target actions with END token
    dummy_actions = torch.randint(0, num_actions, (batch_size, 19))  # 19 movement actions
    dummy_target_actions = torch.cat([
        dummy_actions, 
        torch.full((batch_size, 1), pathfinding_net.path_planner.end_token_id)
    ], dim=1)  # Add END token
    
    print(f"\n=== Testing Forward Pass ===")
    print(f"Input voxel shape: {dummy_voxel_data.shape}")
    print(f"Input positions shape: {dummy_positions.shape}")
    print(f"Target actions shape: {dummy_target_actions.shape}")
    print(f"Target action values range: [{dummy_target_actions.min().item()}, {dummy_target_actions.max().item()}]")
    
    # Training forward pass
    pathfinding_net.train()
    action_logits, turn_penalties = pathfinding_net(
        dummy_voxel_data, 
        dummy_positions, 
        dummy_target_actions
    )
    
    print(f"\nTraining mode outputs:")
    print(f"Action logits shape: {action_logits.shape} (should be {(batch_size, 20, 8)})")
    print(f"Turn logits shape: {turn_penalties.shape}")
    
    # Inference forward pass
    pathfinding_net.eval()
    with torch.no_grad():
        generated_paths = pathfinding_net(dummy_voxel_data, dummy_positions)
    
    print(f"\nInference mode outputs:")
    print(f"Generated paths shape: {generated_paths.shape}")
    if generated_paths.shape[1] > 0:
        print(f"Generated action values range: [{generated_paths.min().item()}, {generated_paths.max().item()}]")
    
    # Test collision checking
    test_actions = generated_paths if generated_paths.shape[1] > 0 else dummy_actions
    collision_mask = pathfinding_net.check_collisions(
        dummy_voxel_data, 
        dummy_positions, 
        test_actions
    )
    print(f"Collision mask shape: {collision_mask.shape}")
    
    # Test loss function with proper masking
    loss_fn = PathfindingLoss(
        turn_penalty_weight=0.1, 
        num_actions=num_actions,
        use_end_token=True
    )
    
    # Adjust collision mask to match target sequence length
    if collision_mask.shape[1] >= 20:
        collision_mask_adjusted = collision_mask[:, :20]
    else:
        # Pad with zeros if collision mask is shorter
        padding = torch.zeros(batch_size, 20 - collision_mask.shape[1], 
                             dtype=torch.bool, device=collision_mask.device)
        collision_mask_adjusted = torch.cat([collision_mask, padding], dim=1)
    
    loss_dict = loss_fn(action_logits, turn_penalties, dummy_target_actions, collision_mask_adjusted)
    
    print(f"\n=== Loss Components ===")
    for key, value in loss_dict.items():
        print(f"{key}: {value.item():.4f}")
    
    # Verify that the loss properly masks special tokens
    print(f"\n=== Verification Tests ===")
    
    # Test 1: Verify token ID assignments
    print(f"1. Token IDs are correctly assigned:")
    print(f"   - Movement actions use IDs 0-5: ✓")
    print(f"   - START token uses ID {pathfinding_net.path_planner.start_token_id}: ✓")
    print(f"   - END token uses ID {pathfinding_net.path_planner.end_token_id}: ✓")
    
    # Test 2: Verify Conv-BN-ReLU order
    print(f"2. Conv-BN-ReLU order is standardized: ✓")
    
    # Test 3: Verify supervised turn labels mask
    with torch.no_grad():
        # Create a sequence with mixed actions and END token
        test_sequence = torch.tensor([[0, 1, 2, 3, 4, 5, 7]])  # Actions 0-5 then END
        valid_mask = (test_sequence < num_actions)
        prev_seq = torch.cat([test_sequence[:, :1], test_sequence[:, :-1]], dim=1)
        prev_valid = torch.cat([torch.zeros_like(valid_mask[:, :1], dtype=torch.bool), valid_mask[:, :-1]], dim=1)
        both_valid = valid_mask & prev_valid
        is_turn = ((test_sequence != prev_seq) & both_valid).float()
        print(f"3. Supervised turn labels test:")
        print(f"   - Test sequence: {test_sequence.tolist()}")
        print(f"   - Valid mask: {valid_mask.tolist()}")
        print(f"   - Both valid mask: {both_valid.tolist()}")
        print(f"   - Turn labels: {is_turn.tolist()}")
    
    # Test 4: Verify action generation doesn't output START token
    print(f"4. Generated paths contain only valid action IDs (0-5):")
    if generated_paths.shape[1] > 0:
        contains_only_valid = (generated_paths >= 0).all() and (generated_paths < num_actions).all()
        print(f"   - Generated actions in valid range: {'✓' if contains_only_valid else '✗'}")
    else:
        print(f"   - No actions generated (early END token)")
    
    print(f"\n=== Network Ready for Training ===")