dataset_process/visualize_sample_features.py

#!/usr/bin/env python3
"""
Visualize Sample Features using Open3D

This script visualizes processed training samples with local features.
Features are colorized using PCA (Principal Component Analysis) where the first 3 
principal components are mapped to RGB colors, or using distinct colors for each part.

Features:
- White background as default (better for presentations/papers)
- Bird's Eye View (BEV) as default viewing angle
- Automatic random sequence selection if none specified
- Toggle between PCA-based colors and part index colors using 'C' key
- Toggle between white and black background using 'B' key
- Load next sample using Spacebar (next in sequence, or first of next sequence)
- Estimate normals for point clouds that don't have them
- Interactive 3D visualization with mouse controls

Usage:
    # Basic usage with PCA colors (default) - sequence will be randomly selected
    python ./dataset_process/visualize_sample_features.py --input /path/to/processed --sample_id 123
    
    # Specify a particular sequence
    python ./dataset_process/visualize_sample_features.py --input /path/to/processed --sequence 00 --sample_id 123
    
    # Start with part index colors
    python ./dataset_process/visualize_sample_features.py --input /path/to/processed --sequence 00 --sample_id 123 --color_mode part
    
    # With custom visualization options
    python ./dataset_process/visualize_sample_features.py --input /path/to/processed --sequence 00 --sample_id 123 --point_size 5.0 --color_mode pca
    
    # Visualize raw samples (before feature extraction) with part-based coloring
    python ./dataset_process/visualize_sample_features.py --input /path/to/raw_samples_output --sequence 00 --sample_id 123 --raw_samples
    
    # Estimate normals for point clouds that don't have them
    python ./dataset_process/visualize_sample_features.py --input /path/to/processed --sequence 00 --sample_id 123 --estimate_normals --normal_estimation_radius 0.2
"""

import os
import sys
import numpy as np
import open3d as o3d
import argparse
import logging
from pathlib import Path
from typing import List, Tuple, Optional, Dict, Union
from sklearn.decomposition import PCA
from sklearn.preprocessing import MinMaxScaler
import glob
import threading
from dataset_process.utils.io_utils import CMAP_DEFAULT

# Setup logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


class SampleVisualizer:
    """Visualizer for processed training samples with feature colorization."""
    
    def __init__(self, 
                 pca_components: int = 3,
                 point_size: float = 3.0,
                 background_color: List[float] = [1.0, 1.0, 1.0],
                 no_coordinate_frame: bool = False,
                 raw_samples: bool = False,
                 estimate_normals: bool = False,
                 normal_estimation_radius: float = 0.1):
        """
        Initialize the sample visualizer.
        
        Args:
            pca_components: Number of PCA components to use (should be 3 for RGB)
            point_size: Size of points in the visualization
            background_color: Background color as [R, G, B] in range [0, 1] (default: white)
            raw_samples: If True, visualize raw samples (before feature extraction),
                         only showing parts colored by index.
            estimate_normals: If True, estimate normals for point clouds that don't have them
            normal_estimation_radius: Radius for normal estimation (default: 0.1)
        """
        self.pca_components = pca_components
        self.point_size = point_size
        self.background_color = background_color
        self.no_coordinate_frame = no_coordinate_frame
        self.raw_samples = raw_samples
        self.estimate_normals = estimate_normals
        self.normal_estimation_radius = normal_estimation_radius
        self.pca = PCA(n_components=pca_components)
        self.scaler = MinMaxScaler()
        
        # Color mode state
        self.use_pca_colors = True
        self.colored_pcds = []
        self.pca_colors = []
        self.part_colors = []
        self.vis = None
        
        # Background toggle state
        self.use_white_background = True
        self.white_background = [1.0, 1.0, 1.0]
        self.black_background = [0.0, 0.0, 0.0]
        
        # Check if initial background is closer to white or black
        if np.mean(self.background_color) > 0.5:
            self.use_white_background = True
        else:
            self.use_white_background = False
        
        # Sample navigation state
        self.base_dir = None
        self.current_sequence = None
        self.current_sample_id = None
        self.padding_length = 6
        self.sequence_samples_cache = {}  # Cache for sequence sample lists
        
        # Global PCA state for consistent coloring across samples
        self.global_pca = None
        self.global_scaler = None
        self.use_global_pca = False
        
        # Supported directory name prefixes
        self.directory_prefixes = ['sample', 'fracture', 'part']
    
    def _detect_dataset_name(self, base_dir: str, sequence: str) -> Optional[str]:
        """
        Detect dataset name from directory structure.
        
        Args:
            base_dir: Base directory containing processed samples
            sequence: Sequence name
            
        Returns:
            Dataset name if found, None otherwise
        """
        if not base_dir or not os.path.exists(base_dir):
            return None
        
        # Check if base_dir itself is a dataset directory (e.g., named "modelnet")
        base_dir_name = os.path.basename(base_dir.rstrip('/'))
        if base_dir_name.lower().startswith('modelnet'):
            # Check if sequences exist directly under base_dir
            seq_path = os.path.join(base_dir, sequence)
            if os.path.exists(seq_path):
                sample_dirs = self._find_sample_directories(seq_path)
                if sample_dirs:
                    return base_dir_name
        
        # Try direct structure first: base_dir/sequence/
        seq_path = os.path.join(base_dir, sequence)
        if os.path.exists(seq_path):
            # Check if this is a direct sequence (no dataset parent)
            sample_dirs = self._find_sample_directories(seq_path)
            if sample_dirs:
                # Check if base_dir itself is a dataset directory
                # by checking if it looks like a dataset name
                if any(base_dir_name.lower().startswith(prefix) for prefix in ['modelnet', 'partnet', 'ikea']):
                    return base_dir_name
                return None
        
        # Try dataset structure: base_dir/dataset_name/sequence/
        for dataset_dir in os.listdir(base_dir):
            dataset_path = os.path.join(base_dir, dataset_dir)
            if os.path.isdir(dataset_path):
                seq_path = os.path.join(dataset_path, sequence)
                if os.path.exists(seq_path):
                    sample_dirs = self._find_sample_directories(seq_path)
                    if sample_dirs:
                        return dataset_dir
                
                # Also check nested structure: base_dir/dataset_name/scene/seq/
                for subseq_dir in os.listdir(dataset_path):
                    subseq_path = os.path.join(dataset_path, subseq_dir)
                    if os.path.isdir(subseq_path):
                        nested_seq_path = os.path.join(subseq_path, sequence.split('/')[-1] if '/' in sequence else sequence)
                        if os.path.exists(nested_seq_path):
                            sample_dirs = self._find_sample_directories(nested_seq_path)
                            if sample_dirs:
                                return dataset_dir
        
        return None
    
    def _determine_padding_length(self, base_dir: str, sequence: str) -> int:
        """
        Determine appropriate padding length by examining actual directory names.
        
        Args:
            base_dir: Base directory containing processed samples
            sequence: Sequence name
            
        Returns:
            Padding length detected from actual directories (default: 6)
        """
        dataset_name = self._detect_dataset_name(base_dir, sequence)
        if dataset_name and dataset_name.lower().startswith('modelnet'):
            return 1
        
        # Try to detect actual padding length from directory names
        # Find the sequence path
        seq_path = None
        
        # Try direct structure: base_dir/sequence/
        direct_path = os.path.join(base_dir, sequence)
        if os.path.exists(direct_path):
            seq_path = direct_path
        
        # Try dataset structure: base_dir/dataset_name/sequence/
        if seq_path is None:
            for dataset_dir in os.listdir(base_dir):
                dataset_path = os.path.join(base_dir, dataset_dir)
                if os.path.isdir(dataset_path):
                    candidate_path = os.path.join(dataset_path, sequence)
                    if os.path.exists(candidate_path):
                        seq_path = candidate_path
                        break
        
        # If we found the sequence path, check actual directory names
        if seq_path:
            sample_dirs = self._find_sample_directories(seq_path)
            if sample_dirs:
                # Extract sample ID from first directory and check its length
                sample_name = os.path.basename(sample_dirs[0])
                sample_id = self._extract_sample_id(sample_name)
                if sample_id and sample_id.isdigit():
                    # Count leading zeros + digits to determine padding
                    # If sample_id is "00000", padding_length should be 5
                    # If sample_id is "0", padding_length should be 1
                    # But we want the total length including leading zeros
                    padding_length = len(sample_id)
                    logger.debug(f"Detected padding_length={padding_length} from directory '{sample_name}' (sample_id='{sample_id}')")
                    return padding_length
        
        # Default fallback
        logger.debug(f"Could not detect padding length from directories, using default 6")
        return 6
    
    def toggle_background_color(self):
        """Toggle between white and black background colors."""
        self.use_white_background = not self.use_white_background
        
        if self.use_white_background:
            new_background = self.white_background
            logger.info("Switched to white background")
        else:
            new_background = self.black_background
            logger.info("Switched to black background")
        
        # Update background color
        self.background_color = new_background
        
        # Update visualization if visualizer exists
        if self.vis is not None:
            render_option = self.vis.get_render_option()
            render_option.background_color = np.array(new_background)
    
    def _find_sample_directories(self, directory: str) -> List[str]:
        """
        Find all sample directories matching any supported prefix pattern.
        
        Args:
            directory: Directory to search in
            
        Returns:
            List of matching directory paths
        """
        sample_dirs = []
        for prefix in self.directory_prefixes:
            if self.raw_samples:
                pattern = os.path.join(directory, f"{prefix}_*")
            else:
                pattern = os.path.join(directory, f"{prefix}_*_processed")
            matches = glob.glob(pattern)
            sample_dirs.extend(matches)
        return sample_dirs
    
    def _extract_sample_id(self, directory_name: str) -> Optional[str]:
        """
        Extract sample ID from directory name, supporting multiple prefixes.
        
        Args:
            directory_name: Name of the directory (e.g., "sample_123", "fracture_456_processed")
            
        Returns:
            Sample ID string, or None if not matched
        """
        for prefix in self.directory_prefixes:
            if self.raw_samples:
                if directory_name.startswith(f"{prefix}_"):
                    return directory_name[len(f"{prefix}_"):]
            else:
                if directory_name.startswith(f"{prefix}_") and directory_name.endswith("_processed"):
                    return directory_name[len(f"{prefix}_"):-10]  # Remove prefix_ and _processed
        return None
    
    def generate_part_colors(self, num_parts: int) -> List[np.ndarray]:
        """
        Generate distinct colors for each part using the default color map from render.py.
        
        Args:
            num_parts: Number of parts to generate colors for
            
        Returns:
            List of color arrays, one per part
        """
        colors = []
        
        for i in range(num_parts):
            # Use default color map, cycling through colors if we have more parts than colors
            color_idx = i % len(CMAP_DEFAULT)
            color = CMAP_DEFAULT[color_idx]
            colors.append(color)
        
        logger.info(f"Generated {num_parts} distinct part colors using default color map")
        if num_parts > len(CMAP_DEFAULT):
            logger.info(f"Note: Cycling through colors since {num_parts} parts > {len(CMAP_DEFAULT)} available colors")
        
        return colors
    
    def load_sample_data(self, sample_dir: str, center_pcds: bool = True) -> Tuple[List[np.ndarray], List[np.ndarray], List[str], List[Optional[np.ndarray]]]:
        """
        Load point clouds and features from a processed sample directory.
        
        Args:
            sample_dir: Path to processed sample directory
            center_pcds: Whether to center all point clouds at origin
            
        Returns:
            Tuple of (point_clouds, features, part_names, normals) or (point_clouds, part_ids, part_names, normals) if raw_samples is True.
            normals is a list of normal arrays (one per part), or None if normals are not available.
        """
        if not os.path.exists(sample_dir):
            raise FileNotFoundError(f"Sample directory not found: {sample_dir}")
        
        # Find all PLY files
        ply_files = sorted(glob.glob(os.path.join(sample_dir, "*.ply")))
        
        if not ply_files:
            raise FileNotFoundError(f"No PLY files found in {sample_dir}")
        
        point_clouds = []
        features = []
        part_names = []
        part_ids = []
        normals = []

        for i, ply_file in enumerate(ply_files):
            # Extract part name from PLY filename
            part_name = os.path.splitext(os.path.basename(ply_file))[0]
            
            try:
                # Load point cloud
                pcd = o3d.io.read_point_cloud(ply_file)
                points = np.asarray(pcd.points)
                
                if len(points) == 0:
                    logger.warning(f"Empty point cloud in {ply_file}")
                    continue
                
                # Check for normals and estimate if needed
                part_normals = None
                if pcd.has_normals():
                    part_normals = np.asarray(pcd.normals)
                    if len(part_normals) != len(points):
                        logger.warning(f"Mismatch between points ({len(points)}) and normals ({len(part_normals)}) for {part_name}")
                        part_normals = None
                    else:
                        logger.debug(f"Loaded normals for part '{part_name}': {len(part_normals)} normals")
                
                # Estimate normals if they don't exist and estimation is requested
                if part_normals is None and self.estimate_normals:
                    logger.info(f"Estimating normals for part '{part_name}' (radius={self.normal_estimation_radius})")
                    # Estimate normals using Open3D
                    pcd.estimate_normals(
                        search_param=o3d.geometry.KDTreeSearchParamHybrid(
                            radius=self.normal_estimation_radius,
                            max_nn=30
                        )
                    )
                    # Orient normals consistently (optional, but helps with visualization)
                    pcd.orient_normals_consistent_tangent_plane(k=15)
                    part_normals = np.asarray(pcd.normals)
                    logger.info(f"Estimated {len(part_normals)} normals for part '{part_name}'")
                elif part_normals is None:
                    logger.debug(f"No normals found for part '{part_name}'")
                
                if self.raw_samples:
                    point_clouds.append(points)
                    part_names.append(part_name)
                    part_ids.append(np.full(len(points), i, dtype=int)) # Assign unique part ID
                    normals.append(part_normals)
                    logger.debug(f"Loaded raw part '{part_name}': {len(points)} points")
                else:
                    # Find corresponding feature file
                    feature_file = os.path.join(sample_dir, f"features_{part_name}.npy")
                    
                    if not os.path.exists(feature_file):
                        logger.warning(f"Feature file not found for {part_name}: {feature_file}")
                        continue

                    # Load features
                    part_features = np.load(feature_file)
                    
                    if len(points) != len(part_features):
                        logger.warning(f"Mismatch between points ({len(points)}) and features ({len(part_features)}) for {part_name}")
                        continue
                    
                    point_clouds.append(points)
                    features.append(part_features)
                    part_names.append(part_name)
                    normals.append(part_normals)
                    
                    logger.debug(f"Loaded part '{part_name}': {len(points)} points, {part_features.shape[1]} feature dims")
                
            except Exception as e:
                logger.error(f"Failed to load {part_name}: {e}")
                continue
        
        if not point_clouds:
            raise RuntimeError(f"No valid parts loaded from {sample_dir}")
        
        # Log detailed point count information
        total_points = sum(len(points) for points in point_clouds)
        
        logger.info("=" * 60)
        logger.info(f"SAMPLE POINT COUNT SUMMARY:")
        logger.info(f"  Sample Directory: {os.path.basename(sample_dir)}")
        logger.info(f"  Number of Parts: {len(point_clouds)}")

        if self.raw_samples:
            logger.info("  Mode: Raw Samples (no features)")
            for i, (points, part_name) in enumerate(zip(point_clouds, part_names)):
                percentage = (len(points) / total_points) * 100
                logger.info(f"  Part {i+1:2d} ({part_name:<12}): {len(points):6d} points ({percentage:5.1f}%) [Part ID: {i}]")
            
        else:
            feature_dims = features[0].shape[1] if features else 0
            logger.info(f"  Feature Dimensions: {feature_dims}")
            for i, (points, part_name) in enumerate(zip(point_clouds, part_names)):
                percentage = (len(points) / total_points) * 100
                logger.info(f"  Part {i+1:2d} ({part_name:<12}): {len(points):6d} points ({percentage:5.1f}%) ")
        logger.info(f"  TOTAL POINTS: {total_points:6d}")
        
        # Log normal availability
        normals_count = sum(1 for n in normals if n is not None)
        if normals_count > 0:
            logger.info(f"  Normals available: {normals_count}/{len(normals)} parts")
        else:
            logger.info(f"  Normals available: None")
        logger.info("=" * 60)
        
        # Center point clouds if requested
        if center_pcds:
            point_clouds, center_offset = self.center_point_clouds(point_clouds)
            logger.info(f"Point clouds centered at origin with offset: {center_offset}")
            # Note: Normals don't need to be modified when centering (they're direction vectors)
        
        if self.raw_samples:
            all_part_ids = np.concatenate(part_ids, axis=0)
            return point_clouds, all_part_ids, part_names, normals
        else:
            return point_clouds, features, part_names, normals
    
    def compute_pca_colors(self, features: List[np.ndarray]) -> List[np.ndarray]:
        """
        Compute PCA-based colors for features.
        
        Args:
            features: List of feature arrays, one per part
            
        Returns:
            List of RGB color arrays, one per part
        """
        if self.raw_samples:
            logger.warning("Attempted to compute PCA colors in raw samples mode. Returning empty list.")
            return []
        
        # Concatenate all features for global PCA
        all_features = np.concatenate(features, axis=0)
        logger.info(f"Computing PCA on {all_features.shape[0]} points with {all_features.shape[1]} features")
        
        # Fit PCA on all features
        pca_features = self.pca.fit_transform(all_features)
        
        # Normalize PCA components to [0, 1] range
        pca_normalized = self.scaler.fit_transform(pca_features)
        
        # Split back into parts
        colors = []
        start_idx = 0
        
        for feature_array in features:
            end_idx = start_idx + len(feature_array)
            part_colors = pca_normalized[start_idx:end_idx]
            colors.append(part_colors)
            start_idx = end_idx
        
        logger.info(f"PCA explained variance ratio: {self.pca.explained_variance_ratio_}")
        
        return colors
    
    def create_colored_point_clouds(self, 
                                  point_clouds: List[np.ndarray], 
                                  colors_data: Union[List[np.ndarray], np.ndarray],
                                  part_names: List[str],
                                  normals: Optional[List[Optional[np.ndarray]]] = None) -> List[o3d.geometry.PointCloud]:
        """
        Create Open3D point clouds with both PCA and part index colors.
        
        Args:
            point_clouds: List of point arrays
            colors_data: List of PCA-based color arrays (RGB in [0, 1]) OR a single np.ndarray of part_ids
            part_names: List of part names
            normals: Optional list of normal arrays (one per part), or None if not available
            
        Returns:
            List of colored Open3d point clouds
        """
        colored_pcds = []
        
        # Generate part index colors
        part_base_colors = self.generate_part_colors(len(point_clouds))
        
        # Store colors for toggling
        if not self.raw_samples:
            self.pca_colors = colors_data # In non-raw mode, colors_data is pca_colors
        self.part_colors = []
        
        for i, (points, part_name) in enumerate(zip(point_clouds, part_names)):
            pcd = o3d.geometry.PointCloud()
            pcd.points = o3d.utility.Vector3dVector(points)
            
            # Add normals if available
            if normals and i < len(normals) and normals[i] is not None:
                pcd.normals = o3d.utility.Vector3dVector(normals[i])
            
            if self.raw_samples:
                # For raw samples, merge all into one PCD and color by part ID
                # This loop will run for each part, but we only create one merged PCD
                # We need to collect all points and colors first
                pass # Actual merging and coloring will happen after the loop
            else:
                # Generate part index colors (same color for all points in this part)
                part_color = np.tile(part_base_colors[i], (len(points), 1))
                self.part_colors.append(part_color)
                
                # Set initial colors based on current mode
                if self.use_pca_colors:
                    pcd.colors = o3d.utility.Vector3dVector(colors_data[i]) # In non-raw mode, colors_data[i] is pca_part_colors
                else:
                    pcd.colors = o3d.utility.Vector3dVector(part_color)
                
            colored_pcds.append(pcd)
            logger.info(f"Created colored point cloud for '{part_name}' (Part {i})")
        
        if self.raw_samples:
            # Merge all point clouds into a single one for raw samples and color by part ID
            merged_pcd = o3d.geometry.PointCloud()
            all_points = np.concatenate(point_clouds, axis=0)
            all_part_ids_flat = colors_data # In raw mode, colors_data is the concatenated part_ids

            # Generate colors based on part IDs
            unique_part_ids = np.unique(all_part_ids_flat)
            num_unique_parts = len(unique_part_ids)
            part_id_colors = self.generate_part_colors(num_unique_parts)

            # Map part IDs to colors
            colors_array = np.zeros((len(all_points), 3))
            for i, part_id in enumerate(unique_part_ids):
                colors_array[all_part_ids_flat == part_id] = part_id_colors[i]

            merged_pcd.points = o3d.utility.Vector3dVector(all_points)
            merged_pcd.colors = o3d.utility.Vector3dVector(colors_array)
            
            # Merge normals if available
            if normals:
                all_normals = []
                for part_normals in normals:
                    if part_normals is not None:
                        all_normals.append(part_normals)
                if all_normals:
                    merged_normals = np.concatenate(all_normals, axis=0)
                    merged_pcd.normals = o3d.utility.Vector3dVector(merged_normals)
            
            colored_pcds = [merged_pcd] # Only one PCD for raw samples
            logger.info(f"Created single merged point cloud for raw samples, colored by {num_unique_parts} part indices.")

        self.colored_pcds = colored_pcds
        return colored_pcds
    
    def toggle_color_mode(self):
        """Toggle between PCA colors and part index colors."""
        self.use_pca_colors = not self.use_pca_colors
        
        if self.raw_samples:
            logger.info("Color mode toggle ignored: Raw samples mode only supports part index coloring.")
            return
        
        if self.use_pca_colors:
            logger.info("Switched to PCA-based colors")
            colors_to_use = self.pca_colors
        else:
            logger.info("Switched to part index-based colors")
            colors_to_use = self.part_colors
        
        # Update colors for all point clouds
        for pcd, colors in zip(self.colored_pcds, colors_to_use):
            pcd.colors = o3d.utility.Vector3dVector(colors)
        
        # Update visualization if visualizer exists
        if self.vis is not None:
            for pcd in self.colored_pcds:
                self.vis.update_geometry(pcd)
    
    def visualize_sample(self, 
                        sample_dir: str,
                        window_name: Optional[str] = None,
                        show_coordinate_frame: bool = True,
                        save_screenshot: Optional[str] = None,
                        center_pcds: bool = True,
                        raw_samples: bool = False) -> None:
        """
        Visualize a processed sample with PCA-colorized features.
        
        Args:
            sample_dir: Path to processed sample directory
            window_name: Custom window name
            show_coordinate_frame: Whether to show coordinate frame
            save_screenshot: Path to save screenshot (optional)
            center_pcds: Whether to center all point clouds at origin (default: True)
            raw_samples: If True, visualize raw samples (before feature extraction),
                         only showing parts colored by index.
        """
        logger.info(f"Visualizing sample from: {sample_dir}")
        
        # Load sample data
        point_clouds, features, part_names, normals = self.load_sample_data(sample_dir, center_pcds=center_pcds)
        
        # Compute PCA colors (use global PCA if available)
        if self.use_global_pca:
            pca_colors = self.apply_global_pca_colors(features)
        else:
            pca_colors = self.compute_pca_colors(features)
        
        # Create colored point clouds (both PCA and part colors)
        if raw_samples:
            # In raw samples mode, features are actually part_ids
            # And there is no pca_colors needed, so we pass the part_ids as colors_data
            colored_pcds = self.create_colored_point_clouds(point_clouds, features, part_names, normals)
        else:
            colored_pcds = self.create_colored_point_clouds(point_clouds, pca_colors, part_names, normals)
        
        # Try to use VisualizerWithKeyCallback, fallback to regular Visualizer
        try:
            self.vis = o3d.visualization.VisualizerWithKeyCallback()
            has_key_callback = True
        except AttributeError:
            self.vis = o3d.visualization.Visualizer()
            has_key_callback = False
            logger.info("Key callback not supported in this Open3D version, using manual toggle")
        
        self.vis.create_window(window_name="Sample 3D Visualizer")
        
        # Add point clouds to visualizer
        for pcd in colored_pcds:
            self.vis.add_geometry(pcd)
        
        # Add coordinate frame if requested
        if show_coordinate_frame and not self.no_coordinate_frame:
            coordinate_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=1.0)
            self.vis.add_geometry(coordinate_frame)
        
        # Set render options
        render_option = self.vis.get_render_option()
        render_option.background_color = np.array(self.background_color)
        render_option.point_size = self.point_size
        
        # Center and fit view to point cloud
        self._center_and_fit_view()
        
        # Register keyboard callback for color toggling if supported
        if has_key_callback:
            try:
                # Create callback function for color toggling
                def color_callback(vis):
                    self.toggle_color_mode()
                    return False
                
                # Create callback function for background toggling
                def background_callback(vis):
                    self.toggle_background_color()
                    return False
                
                # Create callback function for loading next sample
                def next_sample_callback(vis):
                    success = self.load_next_sample()
                    if success:
                        logger.info("Spacebar pressed: Loaded next sample")
                    else:
                        logger.warning("Spacebar pressed: Failed to load next sample")
                    return False
                
                # Create callback function for loading next sequence
                def next_sequence_callback(vis):
                    success = self.load_next_sequence()
                    if success:
                        logger.info("N pressed: Jumped to next sequence")
                    else:
                        logger.warning("N pressed: Failed to jump to next sequence")
                    return False
                
                self.vis.register_key_callback(ord('C'), color_callback)
                self.vis.register_key_callback(ord('B'), background_callback)
                self.vis.register_key_callback(ord(' '), next_sample_callback)  # Spacebar
                self.vis.register_key_callback(ord('N'), next_sequence_callback)  # N key
                logger.info("Keyboard callbacks registered successfully")
            except Exception as e:
                logger.warning(f"Could not register keyboard callback: {e}")
                has_key_callback = False
        
        logger.info("Visualization ready. Press 'Q' to quit.")
        logger.info("Controls:")
        logger.info("  - Mouse: Rotate view")
        logger.info("  - Scroll: Zoom")
        logger.info("  - Shift+Mouse: Pan")
        logger.info("  - R: Reset view")
        logger.info("  - H: Print help")
        
        if has_key_callback:
            logger.info("  - C: Toggle between PCA colors and part index colors")
            logger.info("  - B: Toggle background color")
            logger.info("  - Spacebar: Load next sample (next in sequence, or first of next sequence)")
            logger.info("  - N: Jump to next sequence (first sample of next sequence)")
        
        if self.raw_samples:
            logger.info("  Note: In raw samples mode, 'C' toggle is disabled as only part index colors are shown.")
        else:
            logger.info(f"  - Color mode: {self.get_color_mode_description()}")
            logger.info("    (To change color mode, close and restart with different initial setting)")
        
        logger.info(f"Current color mode: {'PCA-based' if self.use_pca_colors else 'Part index-based'}")
        logger.info("=" * 60)
        # Display padded sample ID for consistency
        padded_sample_id = self.current_sample_id.zfill(self.padding_length) if self.current_sample_id else "N/A"
        logger.info(f"CURRENT SAMPLE: Sequence {self.current_sequence}, Sample {padded_sample_id}")
        logger.info("=" * 60)
        
        # Run visualizer
        if has_key_callback:
            # Use the callback-based approach
            self.vis.run()
        else:
            # Use manual polling approach for compatibility
            self._run_with_manual_controls()
        
        # Save screenshot if requested
        if save_screenshot:
            self.vis.capture_screen_image(save_screenshot)
            logger.info(f"Screenshot saved to: {save_screenshot}")
        
        self.vis.destroy_window()
        self.vis = None  # Clean up reference
    
    def _run_with_manual_controls(self):
        """Run visualizer with manual control polling for older Open3D versions."""
        logger.info("Manual control mode - press ESC in console and type commands + Enter:")
        logger.info("  'c' = toggle colors, 'b' = toggle background, 'n' = next sample, 's' = next sequence, 'q' = quit")
        logger.info(f"Current: Sequence {self.current_sequence}, Sample {self.current_sample_id}")
        
        def input_thread():
            while True:
                try:
                    command = input().strip().lower()
                    if command == 'c':
                        self.toggle_color_mode()
                    elif command == 'b':
                        self.toggle_background_color()
                    elif command == 'n':
                        success = self.load_next_sample()
                        if success:
                            logger.info("Loaded next sample")
                        else:
                            logger.warning("Failed to load next sample")
                    elif command == 's':
                        success = self.load_next_sequence()
                        if success:
                            logger.info("Jumped to next sequence")
                        else:
                            logger.warning("Failed to jump to next sequence")
                    elif command in ['q', 'quit', 'exit']:
                        break
                except (EOFError, KeyboardInterrupt):
                    break
        
        # Start input thread
        input_thread_obj = threading.Thread(target=input_thread, daemon=True)
        input_thread_obj.start()
        
        # Run the visualizer
        self.vis.run()
    
    def get_color_mode_description(self) -> str:
        """Get description of current color mode."""
        if self.use_pca_colors:
            return "PCA-based colors (features)"
        else:
            return "Part index-based colors (distinct per part)"
    
    def set_navigation_state(self, base_dir: str, sequence: str, sample_id: str, padding_length: Optional[int] = None):
        """
        Set the navigation state for sample browsing.
        
        Args:
            base_dir: Base directory containing processed samples
            sequence: Current sequence name
            sample_id: Current sample ID
            padding_length: Length to pad sample ID to (auto-detected if None)
        """
        self.base_dir = base_dir
        self.current_sequence = sequence
        self.current_sample_id = sample_id
        
        # Auto-detect padding_length if not provided
        if padding_length is None:
            padding_length = self._determine_padding_length(base_dir, sequence)
            dataset_name = self._detect_dataset_name(base_dir, sequence)
            if dataset_name:
                logger.info(f"Auto-detected padding_length={padding_length} for dataset '{dataset_name}'")
        
        self.padding_length = padding_length
        logger.info(f"Navigation state set: {base_dir}, sequence {sequence}, sample {sample_id}, padding_length={padding_length}")
    
    def compute_global_pca(self, max_samples: int = 1000):
        """
        Compute global PCA on a subset of all available samples for consistent coloring.
        
        Args:
            max_samples: Maximum number of samples to use for PCA computation (for performance)
        """
        if not self.base_dir or not os.path.exists(self.base_dir):
            logger.warning("Base directory not set, cannot compute global PCA")
            return False
        
        logger.info("Computing global PCA for consistent coloring across samples...")
        
        # Get all available sequences
        sequences = self.get_available_sequences()
        if not sequences:
            logger.warning("No sequences found for global PCA computation")
            return False
        
        all_features = []
        sample_count = 0
        
        # Collect features from samples across all sequences
        for sequence in sequences:
            if sample_count >= max_samples:
                break
                
            samples = self.get_sequence_samples(sequence)
            for sample_id in samples:
                if sample_count >= max_samples:
                    break
                    
                try:
                    # Find sample directory
                    sample_dir = find_sample_directory(self.base_dir, sequence, sample_id, self.padding_length, raw_samples=self.raw_samples, directory_prefixes=self.directory_prefixes)
                    
                    # Load sample data
                    point_clouds, features, part_names, normals = self.load_sample_data(sample_dir, center_pcds=False)
                    
                    # Concatenate all features from this sample
                    if features:
                        sample_features = np.concatenate(features, axis=0)
                        all_features.append(sample_features)
                        sample_count += 1
                        
                        if sample_count % 100 == 0:
                            logger.info(f"Processed {sample_count} samples for global PCA...")
                    
                except Exception as e:
                    logger.warning(f"Failed to load sample {sequence}/{sample_id} for global PCA: {e}")
                    continue
        
        if not all_features:
            logger.warning("No features collected for global PCA")
            return False
        
        # Concatenate all features
        all_features = np.concatenate(all_features, axis=0)
        logger.info(f"Computing global PCA on {all_features.shape[0]} points with {all_features.shape[1]} features from {sample_count} samples")
        
        # Compute global PCA
        self.global_pca = PCA(n_components=self.pca_components)
        self.global_scaler = MinMaxScaler()
        
        # Fit PCA and scaler
        pca_features = self.global_pca.fit_transform(all_features)
        self.global_scaler.fit(pca_features)
        
        logger.info(f"Global PCA computed successfully. Explained variance ratio: {self.global_pca.explained_variance_ratio_}")
        self.use_global_pca = True
        return True
    
    def apply_global_pca_colors(self, features: List[np.ndarray]) -> List[np.ndarray]:
        """
        Apply global PCA transformation to features for consistent coloring.
        
        Args:
            features: List of feature arrays, one per part
            
        Returns:
            List of RGB color arrays, one per part
        """
        if not self.use_global_pca or self.global_pca is None:
            logger.warning("Global PCA not available, falling back to local PCA")
            return self.compute_pca_colors(features)
        
        # Concatenate all features
        all_features = np.concatenate(features, axis=0)
        
        # Apply global PCA transformation
        pca_features = self.global_pca.transform(all_features)
        
        # Apply global scaling
        pca_normalized = self.global_scaler.transform(pca_features)
        
        # Split back into parts
        colors = []
        start_idx = 0
        
        for feature_array in features:
            end_idx = start_idx + len(feature_array)
            part_colors = pca_normalized[start_idx:end_idx]
            colors.append(part_colors)
            start_idx = end_idx
        
        return colors
    
    def get_available_sequences(self) -> List[str]:
        """Get list of available sequences in the base directory."""
        if not self.base_dir or not os.path.exists(self.base_dir):
            return []
        
        sequences = []
        
        # Check direct sequence directories
        for item in os.listdir(self.base_dir):
            item_path = os.path.join(self.base_dir, item)
            if os.path.isdir(item_path):
                # Check if this directory contains processed samples
                sample_dirs = self._find_sample_directories(item_path)
                if sample_dirs:
                    sequences.append(item)
        
        # Also check for dataset structure: base_dir/dataset_name/sequence/
        for dataset_dir in os.listdir(self.base_dir):
            dataset_path = os.path.join(self.base_dir, dataset_dir)
            if os.path.isdir(dataset_path):
                for seq_dir in os.listdir(dataset_path):
                    seq_path = os.path.join(dataset_path, seq_dir)
                    if os.path.isdir(seq_path):
                        sample_dirs = self._find_sample_directories(seq_path)
                        if sample_dirs and seq_dir not in sequences:
                            sequences.append(seq_dir)
                        
                        # Also check for nested structure: base_dir/dataset_name/scene/seq/
                        # (e.g., for ThreeDMatch: base_dir/dataset_name/bundlefusion-office0/seq-01/)
                        for subseq_dir in os.listdir(seq_path):
                            subseq_path = os.path.join(seq_path, subseq_dir)
                            if os.path.isdir(subseq_path):
                                # Check if subseq_dir itself matches a prefix pattern
                                if any(subseq_dir.startswith(f"{prefix}_") for prefix in self.directory_prefixes):
                                    continue
                                sample_dirs = self._find_sample_directories(subseq_path)
                                if sample_dirs:
                                    nested_seq_name = f"{seq_dir}/{subseq_dir}"
                                    if nested_seq_name not in sequences:
                                        sequences.append(nested_seq_name)
        
        return sorted(sequences)
    
    def get_sequence_samples(self, sequence: str) -> List[str]:
        """
        Get list of available sample IDs for a given sequence.
        
        Args:
            sequence: Sequence name (can be nested like "scene/seq-01")
            
        Returns:
            List of sample IDs (as strings)
        """
        if sequence in self.sequence_samples_cache:
            return self.sequence_samples_cache[sequence]
        
        if not self.base_dir or not os.path.exists(self.base_dir):
            return []
        
        sample_ids = []
        
        # Try direct structure: base_dir/sequence/{prefix}_XXXXXX{_processed}
        seq_path = os.path.join(self.base_dir, sequence)
        if os.path.exists(seq_path):
            sample_dirs = self._find_sample_directories(seq_path)
            for sample_dir in sample_dirs:
                # Extract sample ID from directory name
                sample_name = os.path.basename(sample_dir)
                sample_id = self._extract_sample_id(sample_name)
                if sample_id and sample_id not in sample_ids: # Ensure uniqueness
                    sample_ids.append(sample_id)
        
        # Try dataset structure: base_dir/*/sequence/{prefix}_XXXXXX{_processed}
        if not sample_ids:
            for dataset_dir in os.listdir(self.base_dir):
                dataset_path = os.path.join(self.base_dir, dataset_dir)
                if os.path.isdir(dataset_path):
                    seq_path = os.path.join(dataset_path, sequence)
                    if os.path.exists(seq_path):
                        sample_dirs = self._find_sample_directories(seq_path)
                        for sample_dir in sample_dirs:
                            sample_name = os.path.basename(sample_dir)
                            sample_id = self._extract_sample_id(sample_name)
                            if sample_id and sample_id not in sample_ids: # Ensure uniqueness
                                sample_ids.append(sample_id)
        
        # Sort sample IDs numerically
        sample_ids.sort(key=lambda x: int(x) if x.isdigit() else x)
        
        # Cache the result
        self.sequence_samples_cache[sequence] = sample_ids
        
        return sample_ids
    
    def find_next_sample(self) -> Tuple[str, str]:
        """
        Find the next sample to load.
        
        Returns:
            Tuple of (next_sequence, next_sample_id)
        """
        if not self.current_sequence or not self.current_sample_id:
            logger.warning("Navigation state not set, cannot find next sample")
            return None, None
        
        # Get available sequences
        sequences = self.get_available_sequences()
        if not sequences:
            logger.warning("No sequences found")
            return None, None
        
        # Get samples for current sequence
        current_samples = self.get_sequence_samples(self.current_sequence)
        if not current_samples:
            logger.warning(f"No samples found for sequence {self.current_sequence}")
            return None, None
        
        # Pad current sample ID to 6 digits for proper indexing
        # In raw samples mode, sample IDs are usually not padded, so use the original ID for lookup
        sample_id_for_lookup = self.current_sample_id if self.raw_samples else self.current_sample_id.zfill(self.padding_length)
        
        # Find current sample index
        try:
            current_idx = current_samples.index(sample_id_for_lookup)
        except ValueError:
            # Try with padded version if it's a digit and not raw_samples (for robustness, though should be covered by sample_id_for_lookup)
            if sample_id_for_lookup.isdigit() and not self.raw_samples and self.current_sample_id.zfill(self.padding_length) in current_samples:
                current_idx = current_samples.index(self.current_sample_id.zfill(self.padding_length))
                logger.warning(f"Current sample '{sample_id_for_lookup}' not found, but found padded '{self.current_sample_id.zfill(self.padding_length)}' in sequence '{self.current_sequence}'. Using padded.")
            else:
                logger.warning(f"Current sample '{sample_id_for_lookup}' not found in sequence '{self.current_sequence}")
                return None, None
        
        # Check if there's a next sample in current sequence
        if current_idx + 1 < len(current_samples):
            next_sample_id = current_samples[current_idx + 1]
            logger.info(f"Next sample in sequence {self.current_sequence}: {next_sample_id}")
            return self.current_sequence, next_sample_id
        
        # If we're at the last sample, move to next sequence
        current_seq_idx = sequences.index(self.current_sequence)
        if current_seq_idx + 1 < len(sequences):
            next_sequence = sequences[current_seq_idx + 1]
            next_samples = self.get_sequence_samples(next_sequence)
            if next_samples:
                next_sample_id = next_samples[0]  # First sample of next sequence
                logger.info(f"Moving to next sequence {next_sequence}, first sample: {next_sample_id}")
                return next_sequence, next_sample_id
        
        # If we're at the last sample of the last sequence, loop back to first
        first_sequence = sequences[0]
        first_samples = self.get_sequence_samples(first_sequence)
        if first_samples:
            first_sample_id = first_samples[0]
            logger.info(f"Looping back to first sequence {first_sequence}, first sample: {first_sample_id}")
            return first_sequence, first_sample_id
        
        logger.warning("No next sample found")
        return None, None
    
    def find_next_sequence(self) -> Tuple[str, str]:
        """
        Find the first sample of the next sequence to load.
        
        Returns:
            Tuple of (next_sequence, first_sample_id)
        """
        if not self.current_sequence:
            logger.warning("Navigation state not set, cannot find next sequence")
            return None, None
        
        # Get available sequences
        sequences = self.get_available_sequences()
        if not sequences:
            logger.warning("No sequences found")
            return None, None
        
        # Find current sequence index
        try:
            current_seq_idx = sequences.index(self.current_sequence)
        except ValueError:
            logger.warning(f"Current sequence {self.current_sequence} not found in available sequences")
            return None, None
        
        # Check if there's a next sequence
        if current_seq_idx + 1 < len(sequences):
            next_sequence = sequences[current_seq_idx + 1]
        else:
            # If we're at the last sequence, loop back to first
            next_sequence = sequences[0]
        
        # Get first sample of next sequence
        next_samples = self.get_sequence_samples(next_sequence)
        if next_samples:
            first_sample_id = next_samples[0]
            logger.info(f"Jumping to next sequence {next_sequence}, first sample: {first_sample_id}")
            return next_sequence, first_sample_id
        else:
            logger.warning(f"No samples found in next sequence {next_sequence}")
            return None, None
    
    def load_next_sequence(self):
        """Load and visualize the first sample of the next sequence."""
        next_sequence, first_sample_id = self.find_next_sequence()
        
        if next_sequence is None or first_sample_id is None:
            logger.warning("No next sequence available")
            return False
        
        try:
            # Find the sample directory
            sample_dir = find_sample_directory(self.base_dir, next_sequence, first_sample_id, self.padding_length, raw_samples=self.raw_samples, directory_prefixes=self.directory_prefixes)
            
            # Update navigation state (store unpadded version for consistency)
            self.current_sequence = next_sequence
            self.current_sample_id = first_sample_id
            
            # Clear current visualization
            if self.vis is not None:
                # Clear all existing geometries (including point clouds and coordinate frame)
                self.vis.clear_geometries()
            
            # Load new sample data
            point_clouds, data_for_colors, part_names, normals = self.load_sample_data(sample_dir, center_pcds=True)
            
            # Compute PCA colors or use part_ids for coloring
            if self.raw_samples:
                pca_colors = [] # Not used in raw mode, but argument required for type consistency if we keep the signature
                colored_pcds = self.create_colored_point_clouds(point_clouds, data_for_colors, part_names, normals) # data_for_colors is part_ids here
            else:
                pca_colors = self.apply_global_pca_colors(data_for_colors) if self.use_global_pca else self.compute_pca_colors(data_for_colors)
                colored_pcds = self.create_colored_point_clouds(point_clouds, pca_colors, part_names, normals)
            
            # Add new geometries to visualizer
            if self.vis is not None:
                for pcd in colored_pcds:
                    self.vis.add_geometry(pcd, False) # Keep False for efficiency in loop
                    self.vis.update_geometry(pcd) # Explicitly update geometry
                
                # Add coordinate frame if enabled (it was cleared by clear_geometries)
                if not self.no_coordinate_frame:
                    coordinate_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=1.0)
                    self.vis.add_geometry(coordinate_frame, False)
                    self.vis.update_geometry(coordinate_frame) # Explicitly update geometry
                
                # Update view
                self._center_and_fit_view()
                self.vis.update_renderer()
            
            logger.info("=" * 60)
            # Display padded sample ID for consistency
            if self.raw_samples:
                logged_sample_id = first_sample_id
            else:
                logged_sample_id = first_sample_id.zfill(self.padding_length) if first_sample_id else "N/A"
            logger.info(f"JUMPED TO NEXT SEQUENCE: Sequence {next_sequence}, Sample {logged_sample_id}")
            logger.info("=" * 60)
            return True
            
        except Exception as e:
            logger.error(f"Failed to load next sequence: {e}")
            return False
    
    def load_next_sample(self):
        """Load and visualize the next sample."""
        next_sequence, next_sample_id = self.find_next_sample()
        
        if next_sequence is None or next_sample_id is None:
            logger.warning("No next sample available")
            return False
        
        try:
            # Find the sample directory
            sample_dir = find_sample_directory(self.base_dir, next_sequence, next_sample_id, self.padding_length, raw_samples=self.raw_samples, directory_prefixes=self.directory_prefixes)
            
            # Update navigation state (store unpadded version for consistency)
            self.current_sequence = next_sequence
            self.current_sample_id = next_sample_id
            
            # Clear current visualization
            if self.vis is not None:
                # Clear all existing geometries (including point clouds and coordinate frame)
                self.vis.clear_geometries()
            
            # Load new sample data
            point_clouds, data_for_colors, part_names, normals = self.load_sample_data(sample_dir, center_pcds=True)
            
            # Compute PCA colors or use part_ids for coloring
            if self.raw_samples:
                pca_colors = [] # Not used in raw mode, but argument required for type consistency if we keep the signature
                colored_pcds = self.create_colored_point_clouds(point_clouds, data_for_colors, part_names, normals) # data_for_colors is part_ids here
            else:
                pca_colors = self.apply_global_pca_colors(data_for_colors) if self.use_global_pca else self.compute_pca_colors(data_for_colors)
                colored_pcds = self.create_colored_point_clouds(point_clouds, pca_colors, part_names, normals)
            
            # Add new geometries to visualizer
            if self.vis is not None:
                for pcd in colored_pcds:
                    self.vis.add_geometry(pcd, False) # Keep False for efficiency in loop
                    self.vis.update_geometry(pcd) # Explicitly update geometry
                
                # Add coordinate frame if enabled (it was cleared by clear_geometries)
                if not self.no_coordinate_frame:
                    coordinate_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=1.0)
                    self.vis.add_geometry(coordinate_frame, False)
                    self.vis.update_geometry(coordinate_frame) # Explicitly update geometry
                
                # Update view
                self._center_and_fit_view()
                self.vis.update_renderer()
                
            logger.info("=" * 60)
            # Display padded sample ID for consistency
            if self.raw_samples:
                logged_sample_id = next_sample_id
            else:
                logged_sample_id = next_sample_id.zfill(self.padding_length) if next_sample_id else "N/A"
            logger.info(f"LOADED NEXT SAMPLE: Sequence {next_sequence}, Sample {logged_sample_id}")
            logger.info("=" * 60)
            return True
            
        except Exception as e:
            logger.error(f"Failed to load next sample: {e}")
            return False
    
    def center_point_clouds(self, point_clouds: List[np.ndarray]) -> Tuple[List[np.ndarray], np.ndarray]:
        """
        Center all point clouds at the origin by subtracting their mean.
        
        Args:
            point_clouds: List of point cloud arrays
            
        Returns:
            Tuple of (centered_point_clouds, center_offset)
        """
        if not point_clouds:
            return point_clouds, np.zeros(3)
        
        # Concatenate all point clouds to find the global center
        all_points = np.concatenate(point_clouds, axis=0)
        center = np.mean(all_points, axis=0)
        
        # Center each point cloud
        centered_point_clouds = []
        for points in point_clouds:
            centered_points = points - center
            centered_point_clouds.append(centered_points)
        
        logger.info(f"Centered all point clouds at origin. Global center was: {center}")
        return centered_point_clouds, center
    
    def _center_and_fit_view(self):
        """Center and fit the view to show all point clouds properly."""
        if not self.colored_pcds:
            return
        
        # Get view control
        view_control = self.vis.get_view_control()
        
        # Calculate bounding box of all point clouds
        all_points = []
        for pcd in self.colored_pcds:
            points = np.asarray(pcd.points)
            if len(points) > 0:
                all_points.append(points)
        
        if not all_points:
            return
        
        # Concatenate all points to get overall bounding box
        all_points = np.concatenate(all_points, axis=0)
        
        # Calculate center and extent
        center = np.mean(all_points, axis=0)
        min_coords = np.min(all_points, axis=0)
        max_coords = np.max(all_points, axis=0)
        extent = max_coords - min_coords
        max_extent = np.max(extent)
        
        # Set camera parameters for a good view
        # Position camera at a distance proportional to the point cloud size
        camera_distance = max_extent * (2.0 if self.raw_samples else 1.0)  # Adjust this multiplier as needed
        
        # Set up a bird's eye view (top-down) by default
        # You can modify these angles for different viewing angles
        front = np.array([0.0, 0.0, 1.0])  # Looking down
        up = np.array([0.0, 1.0, 0.0])      # Y-axis up
        lookat = center                      # Look at the center of the point cloud
        
       
        # Set the view
        view_control.set_front(front)
        view_control.set_lookat(lookat)
        view_control.set_up(up)
        view_control.set_zoom(0.8)  # Adjust zoom level
        
        logger.info(f"Centered view on point cloud center: {center}")
        logger.info(f"Point cloud extent: {extent}, max extent: {max_extent}")


def find_sample_directory(base_dir: str, sequence: str, sample_id: str, padding_length: int = 6, raw_samples: bool = False, directory_prefixes: List[str] = None) -> str:
    """
    Find the processed sample directory.
    
    Args:
        base_dir: Base directory containing processed samples
        sequence: Sequence name (e.g., "00" or "bundlefusion-office0/seq-01")
        sample_id: Sample ID (e.g., "123" or "000123")
        padding_length: Length to pad sample ID to (default: 6)
        raw_samples: Whether to look for raw samples (without _processed suffix)
        directory_prefixes: List of directory name prefixes to try (default: ['sample', 'fracture', 'part'])
        
    Returns:
        Path to the sample directory
    """
    if directory_prefixes is None:
        directory_prefixes = ['sample', 'fracture', 'part']
    
    # Ensure sample_id is properly padded
    if sample_id.isdigit() and not raw_samples:
        # Only pad if the sample_id has fewer digits than padding_length
        # If it already has the correct number of digits (or more), use it as-is
        if len(sample_id) < padding_length:
            padded_sample_id = sample_id.zfill(padding_length)
        else:
            # Sample ID already has correct padding, use as-is
            padded_sample_id = sample_id
    else:
        padded_sample_id = sample_id
    
    suffix = '_processed' if not raw_samples else ''
    
    # Try different possible directory structures with all prefixes
    possible_paths = []
    for prefix in directory_prefixes:
        # Direct structure: base_dir/sequence/{prefix}_XXXXXX{_processed}
        possible_paths.append(
            os.path.join(base_dir, sequence, f"{prefix}_{padded_sample_id}{suffix}")
        )
        # Dataset structure: base_dir/dataset_name/sequence/{prefix}_XXXXXX{_processed}
        possible_paths.append(
            os.path.join(base_dir, "*", sequence, f"{prefix}_{padded_sample_id}{suffix}")
        )
    
    # Log attempted paths for debugging
    logger.debug(f"Searching for sample directory with sample_id='{sample_id}' -> padded_sample_id='{padded_sample_id}' (padding_length={padding_length})")
    
    for pattern in possible_paths:
        if "*" in pattern:
            # Use glob for wildcard patterns
            matches = glob.glob(pattern)
            if matches:
                logger.debug(f"Found sample directory via glob pattern: {matches[0]}")
                return matches[0]  # Return first match
        else:
            if os.path.exists(pattern):
                logger.debug(f"Found sample directory: {pattern}")
                return pattern
    
    # Log all attempted paths for better error messages
    logger.error(f"Sample directory not found. Attempted paths:")
    for path in possible_paths:
        logger.error(f"  - {path}")
    raise FileNotFoundError(f"Sample directory not found for sequence {sequence}, sample {padded_sample_id} (original: {sample_id}) in {base_dir}. Tried prefixes: {directory_prefixes}")


def main():
    parser = argparse.ArgumentParser(description='Visualize processed training samples with PCA-colorized features')
    
    # Input arguments
    parser.add_argument('--input', '-i', type=str, required=True,
                        help='Input directory containing processed samples')
    parser.add_argument('--sequence', '-s', type=str, default=None,
                        help='Sequence name (e.g., "00"). If not specified, a random sequence will be selected.')
    parser.add_argument('--sample_id', '--sample', type=str, default="0",
                        help='Sample ID (e.g., "123" or "000123") - will be padded automatically')
    
    # Formatting options
    parser.add_argument('--padding_length', type=int, default=None,
                        help='Length to pad sample ID to (default: auto-detect based on dataset name, 1 for modelnet, 6 for others)')
    
    # Visualization options
    parser.add_argument('--point_size', '-p', type=float, default=4.0,
                        help='Point size in visualization (default: 4.0)')
    parser.add_argument('--background_color', type=float, nargs=3, default=[1.0, 1.0, 1.0],
                        help='Background color as R G B values in [0,1] (default: 1.0 1.0 1.0 - white)')
    parser.add_argument('--no_coordinate_frame', action='store_true', default=True, 
                        help='Hide coordinate frame')
    parser.add_argument('--save_screenshot', type=str, default=None,
                        help='Save screenshot to specified path')
    parser.add_argument('--window_name', type=str, default=None,
                        help='Custom window name')
    parser.add_argument('--color_mode', type=str, default='pca', choices=['pca', 'part'],
                        help='Initial color mode: pca (feature-based) or part (index-based) (default: pca)')
    parser.add_argument('--no_center', action='store_true',
                        help='Do not center point clouds at origin (keep original positions)')
    parser.add_argument('--raw_samples', '-r', action='store_true', default=False,
                        help='Visualize raw samples (before feature extraction), only showing parts colored by index')
    parser.add_argument('--estimate_normals', '-n', action='store_true', default=False,
                        help='Estimate normals for point clouds that don\'t have them')
    parser.add_argument('--normal_estimation_radius', type=float, default=0.5,
                        help='Radius for normal estimation (default: 0.1)')
    
    # PCA options
    parser.add_argument('--pca_components', type=int, default=3,
                        help='Number of PCA components (default: 3 for RGB)')
    parser.add_argument('--global_pca', action='store_true', default=True,
                        help='Compute global PCA on all samples for consistent coloring across samples')
    parser.add_argument('--max_pca_samples', type=int, default=50,
                        help='Maximum number of samples to use for global PCA computation (default: 50)')
    
    # Utility arguments
    parser.add_argument('--log_level', type=str, default='INFO', 
                        choices=['DEBUG', 'INFO', 'WARNING', 'ERROR'],
                        help='Logging level (default: INFO)')
    
    args = parser.parse_args()
    
    # Setup logging
    logging.basicConfig(
        level=getattr(logging, args.log_level),
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )
    
    # Validate input
    if not os.path.exists(args.input):
        logger.error(f"Input directory does not exist: {args.input}")
        return
    
    try:
        # Create visualizer first to access sequence discovery methods
        visualizer = SampleVisualizer(
            pca_components=args.pca_components,
            point_size=args.point_size,
            background_color=args.background_color,
            no_coordinate_frame=args.no_coordinate_frame,
            raw_samples=args.raw_samples,
            estimate_normals=args.estimate_normals,
            normal_estimation_radius=args.normal_estimation_radius
        )
        
        # Set navigation state temporarily to access sequence discovery
        # Use provided padding_length or None for auto-detection
        visualizer.set_navigation_state(args.input, "temp", "0", args.padding_length if args.padding_length is not None else 6)
        
        # If sequence not specified, select a random one
        if args.sequence is None:
            available_sequences = visualizer.get_available_sequences()
            if not available_sequences:
                logger.error(f"No sequences found in input directory: {args.input}")
                return
            
            import random
            args.sequence = random.choice(available_sequences)
            logger.info(f"No sequence specified, randomly selected: {args.sequence}")
        
        # Auto-detect padding_length if not explicitly provided (after sequence is determined)
        if args.padding_length is None:
            args.padding_length = visualizer._determine_padding_length(args.input, args.sequence)
            dataset_name = visualizer._detect_dataset_name(args.input, args.sequence)
            if dataset_name:
                logger.info(f"Auto-detected padding_length={args.padding_length} for dataset '{dataset_name}'")
            else:
                logger.info(f"Auto-detected padding_length={args.padding_length} (default)")
        
        # Get available samples for the selected sequence
        available_samples = visualizer.get_sequence_samples(args.sequence)
        if not available_samples:
            logger.error(f"No samples found in sequence {args.sequence} in directory: {args.input}")
            return

        # If sample_id is default "0" or not found, use the first available sample
        if args.sample_id == "0" or args.sample_id not in available_samples:
            if args.sample_id == "0":
                logger.info(f"No sample ID specified (default '0'), using first sample in sequence: {available_samples[0]}")
            else:
                logger.warning(f"Specified sample ID '{args.sample_id}' not found in sequence '{args.sequence}', using first sample: {available_samples[0]}")
            args.sample_id = available_samples[0]

        # Find sample directory (with automatic padding)
        sample_dir = find_sample_directory(args.input, args.sequence, args.sample_id, args.padding_length, raw_samples=args.raw_samples, directory_prefixes=visualizer.directory_prefixes)
        logger.info(f"Found sample directory: {sample_dir}")
        
        # Log the padded sample ID for user reference
        if args.sample_id.isdigit() and not args.raw_samples:
            padded_id = args.sample_id.zfill(args.padding_length)
            logger.info(f"Using padded sample ID: {args.sample_id} -> {padded_id}")
        elif args.raw_samples:
            logger.info(f"Using raw sample ID: {args.sample_id} (not padded in raw mode)")
        
        # Log initial sample information
        logger.info("=" * 60)
        logger.info(f"STARTING VISUALIZATION:")
        logger.info(f"  Input Directory: {args.input}")
        logger.info(f"  Sequence: {args.sequence}")
        logger.info(f"  Sample ID: {args.sample_id} (padded: {args.sample_id.zfill(args.padding_length)})")
        logger.info(f"  Sample Directory: {os.path.basename(sample_dir)}")
        logger.info("=" * 60)
        
        # Set initial color mode based on command line argument
        if args.color_mode == 'pca':
            visualizer.use_pca_colors = True
            logger.info("Initial color mode: PCA-based (feature-based)")
        else:
            visualizer.use_pca_colors = False
            logger.info("Initial color mode: Part index-based (distinct per part)")
        
        # If raw samples, force part index color mode
        if args.raw_samples:
            visualizer.use_pca_colors = False
            logger.info("Raw samples mode: Forcing initial color mode to part index-based.")

        # Set navigation state
        visualizer.set_navigation_state(args.input, args.sequence, args.sample_id, args.padding_length)
        
        # Compute global PCA if requested
        if args.global_pca and not args.raw_samples:
            logger.info("Global PCA requested - computing on all available samples...")
            success = visualizer.compute_global_pca(max_samples=args.max_pca_samples)
            if success:
                logger.info("Global PCA computed successfully - colors will be consistent across samples")
            else:
                logger.warning("Failed to compute global PCA - falling back to local PCA per sample")
        elif args.global_pca and args.raw_samples:
            logger.info("Global PCA skipped: Raw samples mode does not require feature-based coloring.")

        # Visualize sample
        visualizer.visualize_sample(
            sample_dir=sample_dir,
            window_name=args.window_name,
            show_coordinate_frame=not args.no_coordinate_frame,
            save_screenshot=args.save_screenshot,
            center_pcds=not args.no_center,
            raw_samples=args.raw_samples
        )
        
        logger.info("Visualization complete!")
        
    except Exception as e:
        logger.error(f"Visualization failed: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()