hubmap_dir_conversion_v3.6.py

import argparse
import csv
import json
import logging
import os
import shutil
import sys
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import pandas as pd


class XeniumHubMapOrganizer:
    def __init__(self, input_dir: str, output_dir: str, sample_name: str):
        self.input_dir = Path(input_dir)
        self.output_dir = Path(output_dir)
        self.sample_name = sample_name
        self.validation_errors = []
        self.validation_warnings = []

        # Setup logging
        logging.basicConfig(level=logging.INFO)
        self.logger = logging.getLogger(__name__)

    def organize_data(
        self,
        microscope_hardware_path: Optional[str] = None,
        microscope_settings_path: Optional[str] = None,
        overlay_image_path: Optional[str] = None,
        tissue_boundary_geojson_path: Optional[str] = None,
        markers_csv_path: Optional[str] = None,
        additional_panels_csv_path: Optional[str] = None,
        custom_probe_set_csv_path: Optional[str] = None,
        custom_probe_set_bed_path: Optional[str] = None,
        custom_gene_list_csv_path: Optional[str] = None,
        probes_csv_path: Optional[str] = None,
    ) -> bool:
        """
        Main function to organize Xenium data according to HuBMAP schema
        """
        try:
            # Create output directory structure
            self._create_directory_structure()

            # Process extras folder
            self._process_extras(microscope_hardware_path, microscope_settings_path)

            # Process raw folder
            self._process_raw(
                overlay_image_path,
                markers_csv_path,
                additional_panels_csv_path,
                custom_probe_set_csv_path,
                custom_probe_set_bed_path,
                custom_gene_list_csv_path,
                probes_csv_path,
            )

            # Process lab_processed folder
            self._process_lab_processed(tissue_boundary_geojson_path)

            # Validate the output
            is_valid = self._validate_output()

            # Generate validation report
            self._generate_validation_report()

            return is_valid

        except Exception as e:
            self.logger.error(f"Error organizing data: {str(e)}")
            return False

    def _create_directory_structure(self):
        """Create the required directory structure"""
        directories = [
            "extras",
            "raw",
            "lab_processed/images",
            "lab_processed/xenium_bundle",
        ]

        for dir_path in directories:
            (self.output_dir / dir_path).mkdir(parents=True, exist_ok=True)

    def _process_extras(
        self,
        microscope_hardware_path: Optional[str],
        microscope_settings_path: Optional[str],
    ):
        """Process files for extras folder"""
        extras_dir = self.output_dir / "extras"

        # Copy microscope hardware file (required for HuBMAP submission)
        if microscope_hardware_path and Path(microscope_hardware_path).exists():
            shutil.copy2(
                microscope_hardware_path, extras_dir / "microscope_hardware.json"
            )
        else:
            if microscope_hardware_path:
                self.validation_errors.append(
                    f"Microscope hardware file not found: {microscope_hardware_path}"
                )
            else:
                self.validation_warnings.append(
                    "microscope_hardware.json not provided - required for HuBMAP submission. "
                    "Generate using micro-meta app (contact help@hubmapconsortium.org)"
                )

        # Copy microscope settings file (optional)
        if microscope_settings_path and Path(microscope_settings_path).exists():
            shutil.copy2(
                microscope_settings_path, extras_dir / "microscope_settings.json"
            )

    def _process_raw(
        self,
        overlay_image_path: Optional[str],
        markers_csv_path: Optional[str],
        additional_panels_csv_path: Optional[str],
        custom_probe_set_csv_path: Optional[str],
        custom_probe_set_bed_path: Optional[str],
        custom_gene_list_csv_path: Optional[str],
        probes_csv_path: Optional[str],
    ):
        """Process files for raw folder"""
        raw_dir = self.output_dir / "raw"

        # Copy overlay image (optional)
        if overlay_image_path and Path(overlay_image_path).exists():
            overlay_path = Path(overlay_image_path)
            if overlay_path.suffix.lower() in [".jpeg", ".jpg", ".tiff", ".tif"]:
                shutil.copy2(
                    overlay_image_path, raw_dir / f"overlay{overlay_path.suffix}"
                )

        # Copy optional CSV files
        optional_files = {
            "markers.csv": markers_csv_path,
            "additional_panels_used.csv": additional_panels_csv_path,
            "custom_probe_set.csv": custom_probe_set_csv_path,
            "custom_gene_list.csv": custom_gene_list_csv_path,
            "probes.csv": probes_csv_path,
        }

        for filename, file_path in optional_files.items():
            if file_path and Path(file_path).exists():
                shutil.copy2(file_path, raw_dir / filename)

        # Copy custom probe set BED file (optional)
        if custom_probe_set_bed_path and Path(custom_probe_set_bed_path).exists():
            shutil.copy2(custom_probe_set_bed_path, raw_dir / "custom_probe_set.bed")

        # Process transcript locations (required) - convert from transcripts.csv.gz
        self._process_transcript_locations()

        # Copy gene panel JSON (required)
        gene_panel_source = self.input_dir / "gene_panel.json"
        if gene_panel_source.exists():
            shutil.copy2(gene_panel_source, raw_dir / "gene_panel.json")
        else:
            self.validation_errors.append(
                "Required gene_panel.json file not found in input directory"
            )

    def _process_transcript_locations(self):
        """Convert transcripts data to transcript_locations.csv format"""
        raw_dir = self.output_dir / "raw"

        # Try to find transcripts file (prioritize filtered transcripts)
        transcript_files = [
            self.input_dir / "qv20-filtered-transcripts" / "transcripts.parquet",
            self.input_dir / "qv20-filtered-transcripts" / "transcripts.csv.gz",
            self.input_dir / "transcripts.parquet",
            self.input_dir / "transcripts.csv.gz",
        ]

        transcript_file = None
        for file_path in transcript_files:
            if file_path.exists():
                transcript_file = file_path
                break

        if transcript_file:
            try:
                # Read transcript data
                if transcript_file.suffix == ".parquet":
                    df = pd.read_parquet(transcript_file)
                else:
                    df = pd.read_csv(transcript_file)

                # More comprehensive column mapping for different Xenium versions
                column_mapping = {
                    "gene": "gene",
                    "gene_id": "gene",
                    "feature_name": "gene",
                    "x_location": "x",
                    "x": "x",
                    "y_location": "y",
                    "y": "y",
                    "z_location": "z",
                    "z": "z",
                    "qv": "quality_score",
                    "quality_value": "quality_score",
                }

                # Check which columns actually exist
                available_cols = df.columns.tolist()
                self.logger.info(f"Available transcript columns: {available_cols}")

                # Apply mapping only for existing columns
                rename_dict = {
                    k: v for k, v in column_mapping.items() if k in available_cols
                }
                df_renamed = df.rename(columns=rename_dict)

                # Select required columns
                required_cols = ["gene", "x", "y"]
                missing_required = [
                    col for col in required_cols if col not in df_renamed.columns
                ]

                if missing_required:
                    raise ValueError(f"Missing required columns: {missing_required}")

                # Build output columns
                output_cols = ["gene", "x", "y"]
                if "z" in df_renamed.columns:
                    output_cols.append("z")
                if "quality_score" in df_renamed.columns:
                    output_cols.append("quality_score")

                df_output = df_renamed[output_cols]

                # Save as CSV
                df_output.to_csv(raw_dir / "transcript_locations.csv", index=False)
                self.logger.info(
                    f"Created transcript_locations.csv with {len(df_output)} transcripts"
                )

            except Exception as e:
                self.validation_errors.append(
                    f"Error processing transcript locations: {str(e)}"
                )
        else:
            self.validation_errors.append(
                "Required transcript data not found in input directory"
            )

    def _process_lab_processed(self, tissue_boundary_geojson_path: Optional[str]):
        """Process files for lab_processed folder"""
        lab_processed_dir = self.output_dir / "lab_processed"
        images_dir = lab_processed_dir / "images"
        xenium_bundle_dir = lab_processed_dir / "xenium_bundle"

        # Process images
        self._process_images(images_dir, tissue_boundary_geojson_path)

        # Process xenium bundle
        self._process_xenium_bundle(xenium_bundle_dir)

    def _process_images(
        self, images_dir: Path, tissue_boundary_geojson_path: Optional[str]
    ):
        """Process image files"""
        # Copy morphology OME-TIFF (use morphology.ome.tif as main image)
        morphology_source = self.input_dir / "morphology.ome.tif"
        if morphology_source.exists():
            ome_tiff_name = f"{self.sample_name}.ome.tiff"
            shutil.copy2(morphology_source, images_dir / ome_tiff_name)

            # Create channels CSV file
            self._create_channels_csv(images_dir, ome_tiff_name)
        else:
            self.validation_errors.append("Required morphology.ome.tif file not found")

        # Copy tissue boundary GeoJSON (optional)
        if tissue_boundary_geojson_path and Path(tissue_boundary_geojson_path).exists():
            geojson_name = f"{self.sample_name}.tissue-boundary.geojson"
            shutil.copy2(tissue_boundary_geojson_path, images_dir / geojson_name)

    def _create_channels_csv(self, images_dir: Path, ome_tiff_name: str):
        """Create channels CSV file for OME-TIFF"""
        self.logger.info(f"Extracting channel metadata from {ome_tiff_name}")

        try:
            # Try to extract actual channel info from OME-TIFF metadata
            channels_data = self._extract_channel_metadata(images_dir / ome_tiff_name)
            self.logger.info(f"Successfully extracted {len(channels_data)} channels")
        except Exception as e:
            self.logger.warning(f"Could not extract channel metadata: {e}")
            # Fall back to default 4-channel Xenium configuration
            channels_data = self._get_default_xenium_channels(4)
            self.logger.info(f"Using default 4-channel Xenium configuration")

        # Log channel purposes for verification
        for i, channel in enumerate(channels_data):
            nuclei_seg = channel.get("is_channel_used_for_nuclei_segmentation", "No")
            cell_seg = channel.get("is_channel_used_for_cell_segmentation", "No")
            is_antibody = channel.get("is_antibody", "No")
            self.logger.info(
                f"Channel {i}: nuclei_seg={nuclei_seg}, cell_seg={cell_seg}, antibody={is_antibody}"
            )

        channels_df = pd.DataFrame(channels_data)
        channels_csv_name = ome_tiff_name.replace(".ome.tiff", ".ome-tiff.channels.csv")
        channels_df.to_csv(images_dir / channels_csv_name, index=False)

        self.logger.info(
            f"Created {channels_csv_name} with {len(channels_data)} channels"
        )

    def _extract_channel_metadata(self, ome_tiff_path: Path) -> List[Dict]:
        """Extract channel metadata from OME-TIFF file"""
        try:
            import tifffile

            # Read OME-TIFF metadata
            with tifffile.TiffFile(ome_tiff_path) as tif:
                num_pages = len(tif.pages)
                self.logger.info(f"TIFF file info: {num_pages} pages")

                # Get OME-XML metadata
                ome_xml = None
                if hasattr(tif, "ome_metadata") and tif.ome_metadata:
                    ome_xml = tif.ome_metadata
                    self.logger.info("Found OME metadata via ome_metadata attribute")
                elif hasattr(tif, "pages") and len(tif.pages) > 0:
                    # Try to get from first page description
                    page = tif.pages[0]
                    if hasattr(page, "description") and page.description:
                        if page.description.startswith("<?xml"):
                            ome_xml = page.description
                            self.logger.info("Found OME metadata in page description")

                if ome_xml:
                    # Debug: show first 500 chars of OME-XML
                    self.logger.debug(f"OME-XML preview: {ome_xml[:500]}...")
                    channels_from_xml = self._parse_ome_xml_channels(ome_xml)

                    # Special handling for Xenium files that might have incorrect OME metadata
                    if len(channels_from_xml) == 1 and num_pages > 1:
                        self.logger.warning(
                            f"OME-XML reports 1 channel but TIFF has {num_pages} pages"
                        )

                        # For Xenium morphology files, check if this looks like a multi-channel structure
                        if num_pages in [4, 8, 12, 16]:  # Common multi-channel patterns
                            # Assume 4 channels for Xenium (standard configuration)
                            inferred_channels = 4
                            if (
                                num_pages == 12
                            ):  # 12 pages might be 4 channels x 3 Z-planes
                                inferred_channels = 4
                            elif (
                                num_pages == 8
                            ):  # 8 pages might be 4 channels x 2 Z-planes
                                inferred_channels = 4
                            elif (
                                num_pages == 16
                            ):  # 16 pages might be 4 channels x 4 Z-planes
                                inferred_channels = 4
                            else:
                                inferred_channels = min(num_pages, 4)

                            self.logger.info(
                                f"Overriding OME metadata: inferring {inferred_channels} channels from {num_pages} pages"
                            )
                            return self._get_default_xenium_channels(inferred_channels)

                    return channels_from_xml
                else:
                    self.logger.info(
                        "No OME-XML metadata found, inferring from TIFF structure"
                    )
                    # Fallback: try to determine channels from image shape
                    return self._infer_channels_from_tiff(tif)

        except ImportError:
            self.logger.warning(
                "tifffile package not available. Install with: pip install tifffile"
            )
            return self._get_default_xenium_channels()
        except Exception as e:
            self.logger.warning(
                f"Could not extract channel metadata from OME-TIFF: {e}"
            )
            return self._get_default_xenium_channels()

    def _analyze_tiff_structure(self, tif) -> Dict:
        """Analyze TIFF structure to understand channel organization"""
        try:
            analysis = {
                "num_pages": len(tif.pages),
                "page_shapes": [],
                "likely_channels": 1,
            }

            # Analyze first few pages
            for i, page in enumerate(tif.pages[:8]):  # Check first 8 pages
                analysis["page_shapes"].append(page.shape)

            # Look for patterns that suggest multi-channel structure
            num_pages = analysis["num_pages"]

            if num_pages == 12:
                # 12 pages could be: 4 channels × 3 Z-planes, or 3 channels × 4 Z-planes
                analysis["likely_channels"] = 4
                analysis["interpretation"] = "4 channels × 3 Z-planes"
            elif num_pages == 8:
                # 8 pages could be: 4 channels × 2 Z-planes, or 2 channels × 4 Z-planes
                analysis["likely_channels"] = 4
                analysis["interpretation"] = "4 channels × 2 Z-planes"
            elif num_pages == 4:
                # 4 pages likely represents 4 channels
                analysis["likely_channels"] = 4
                analysis["interpretation"] = "4 channels"
            elif num_pages == 16:
                # 16 pages could be 4 channels × 4 Z-planes
                analysis["likely_channels"] = 4
                analysis["interpretation"] = "4 channels × 4 Z-planes"
            else:
                analysis["likely_channels"] = min(num_pages, 4)
                analysis["interpretation"] = (
                    f"Assuming {analysis['likely_channels']} channels"
                )

            return analysis

        except Exception as e:
            self.logger.warning(f"Error analyzing TIFF structure: {e}")
            return {"num_pages": 0, "page_shapes": [], "likely_channels": 4}

    def _parse_ome_xml_channels(self, ome_xml: str) -> List[Dict]:
        """Parse OME-XML to extract channel information"""
        try:
            import xml.etree.ElementTree as ET

            # Parse XML
            root = ET.fromstring(ome_xml)

            # Try multiple namespace variations
            namespaces = [
                {"ome": "http://www.openmicroscopy.org/Schemas/OME/2016-06"},
                {"ome": "http://www.openmicroscopy.org/Schemas/OME/2015-01"},
                {"ome": "http://www.openmicroscopy.org/Schemas/OME/2013-06"},
                {},  # No namespace
            ]

            channels = []
            channel_elements = []

            # Try different namespace combinations
            for ns in namespaces:
                if ns:
                    channel_elements = root.findall(".//ome:Channel", ns)
                    if channel_elements:
                        break
                else:
                    # Try without namespace
                    channel_elements = root.findall(".//Channel")
                    if channel_elements:
                        break

            # Also try to get channel count from Pixels element
            num_channels_from_pixels = None
            for ns in namespaces:
                if ns:
                    pixels = root.find(".//ome:Pixels", ns)
                else:
                    pixels = root.find(".//Pixels")

                if pixels is not None:
                    size_c = pixels.get("SizeC")
                    if size_c:
                        try:
                            num_channels_from_pixels = int(size_c)
                            self.logger.info(
                                f"Found SizeC={num_channels_from_pixels} in Pixels element"
                            )
                            break
                        except:
                            pass

            # Determine actual number of channels
            if channel_elements:
                num_channels = len(channel_elements)
                self.logger.info(f"Found {num_channels} Channel elements in OME-XML")
            elif num_channels_from_pixels:
                num_channels = num_channels_from_pixels
                self.logger.info(
                    f"Using channel count from Pixels SizeC: {num_channels}"
                )
            else:
                self.logger.warning(
                    "Could not determine channel count from OME-XML, using default"
                )
                return self._get_default_xenium_channels()

            # Get default Xenium channel templates
            default_channels = self._get_default_xenium_channels(num_channels)

            # If we have actual channel elements, try to extract their IDs
            if channel_elements:
                for i, channel in enumerate(channel_elements):
                    if i < len(default_channels):
                        channel_info = default_channels[i].copy()

                        # Override with actual OME-TIFF channel ID if available
                        channel_id = channel.get("ID", f"Channel:0:{i}")
                        channel_info["channel_id"] = channel_id

                        channels.append(channel_info)
            else:
                # Use default channels with proper numbering
                channels = default_channels

            self.logger.info(f"Generated {len(channels)} channel entries")
            return channels

        except Exception as e:
            self.logger.warning(f"Error parsing OME-XML: {e}")
            return self._get_default_xenium_channels()

    def _infer_channels_from_tiff(self, tif) -> List[Dict]:
        """Infer channel information from TIFF structure when OME-XML is not available"""
        try:
            # Get image dimensions
            if hasattr(tif, "series") and len(tif.series) > 0:
                series = tif.series[0]
                shape = series.shape
                self.logger.info(f"TIFF series shape: {shape}")

                # Determine number of channels from shape
                num_channels = 4  # Default for Xenium

                if len(shape) >= 3:
                    # Try to identify channel dimension
                    # Common patterns: (C, Y, X), (Y, X, C), (Z, C, Y, X), etc.
                    for i, dim_size in enumerate(shape):
                        if (
                            2 <= dim_size <= 10
                        ):  # Reasonable range for number of channels
                            num_channels = dim_size
                            self.logger.info(
                                f"Found potential channel dimension at axis {i}: {dim_size}"
                            )
                            break

            elif hasattr(tif, "pages") and len(tif.pages) > 0:
                # Check if it's a multi-page TIFF
                num_pages = len(tif.pages)
                page = tif.pages[0]
                shape = page.shape
                self.logger.info(
                    f"TIFF has {num_pages} pages, first page shape: {shape}"
                )

                # For Xenium multi-page TIFF, pages usually represent Z-planes × channels
                # Common Xenium patterns:
                # - 4 pages = 4 channels × 1 Z-plane
                # - 8 pages = 4 channels × 2 Z-planes
                # - 12 pages = 4 channels × 3 Z-planes
                # - 16 pages = 4 channels × 4 Z-planes

                if num_pages == 4:
                    num_channels = 4
                    self.logger.info("Interpreting 4 pages as 4 channels")
                elif num_pages == 8:
                    num_channels = 4  # 4 channels × 2 Z-planes
                    self.logger.info("Interpreting 8 pages as 4 channels × 2 Z-planes")
                elif num_pages == 12:
                    num_channels = 4  # 4 channels × 3 Z-planes
                    self.logger.info("Interpreting 12 pages as 4 channels × 3 Z-planes")
                elif num_pages == 16:
                    num_channels = 4  # 4 channels × 4 Z-planes
                    self.logger.info("Interpreting 16 pages as 4 channels × 4 Z-planes")
                elif num_pages > 1 and num_pages <= 6:
                    # For other small page counts, assume each page is a channel
                    num_channels = num_pages
                    self.logger.info(
                        f"Interpreting {num_pages} pages as {num_channels} channels"
                    )
                else:
                    # Default to 4 channels for Xenium
                    num_channels = 4
                    self.logger.info(f"Using default 4 channels for {num_pages} pages")
            else:
                return self._get_default_xenium_channels()

            self.logger.info(f"Inferred {num_channels} channels from TIFF structure")
            return self._get_default_xenium_channels(num_channels)

        except Exception as e:
            self.logger.warning(f"Error inferring channels from TIFF: {e}")
            return self._get_default_xenium_channels()

    def _get_default_xenium_channels(self, num_channels: int = 4) -> List[Dict]:
        """Get default Xenium channel configuration based on HuBMAP requirements"""

        # Xenium-specific channel configuration based on multimodal cell segmentation
        xenium_channels = [
            {
                # Channel 0: Nuclear Label (DAPI)
                "channel_id": "Channel:0:0",
                "is_channel_used_for_nuclei_segmentation": "Yes",
                "is_channel_used_for_cell_segmentation": "No",
                "is_antibody": "No",  # DAPI is a nucleic acid stain, not an antibody
                "channel_quality": "Good",
                "default_for_display": True,
                "threshold": 65535,
                "minimum_threshold": 0,
            },
            {
                # Channel 1: Membrane Boundary Label (ATP1A1, E-Cadherin, CD45)
                "channel_id": "Channel:0:1",
                "is_channel_used_for_nuclei_segmentation": "No",
                "is_channel_used_for_cell_segmentation": "Yes",
                "is_antibody": "Yes",  # Uses antibodies for membrane proteins
                "channel_quality": "Good",
                "default_for_display": True,
                "threshold": 65535,
                "minimum_threshold": 0,
            },
            {
                # Channel 2: Interior - RNA Label (18S Ribosomal RNA)
                "channel_id": "Channel:0:2",
                "is_channel_used_for_nuclei_segmentation": "No",
                "is_channel_used_for_cell_segmentation": "Yes",
                "is_antibody": "No",  # RNA probe, not antibody
                "channel_quality": "Good",
                "default_for_display": False,
                "threshold": 65535,
                "minimum_threshold": 0,
            },
            {
                # Channel 3: Interior - Protein Label (alphaSMA/Vimentin)
                "channel_id": "Channel:0:3",
                "is_channel_used_for_nuclei_segmentation": "No",
                "is_channel_used_for_cell_segmentation": "Yes",
                "is_antibody": "Yes",  # Uses antibodies for cytoskeletal proteins
                "channel_quality": "Good",
                "default_for_display": False,
                "threshold": 65535,
                "minimum_threshold": 0,
            },
        ]

        # Return only the requested number of channels
        channels = xenium_channels[:num_channels]

        # Ensure at least one channel is marked for nuclei segmentation and one for cell segmentation
        if num_channels > 0:
            # First channel should be nuclei segmentation
            channels[0]["is_channel_used_for_nuclei_segmentation"] = "Yes"
            channels[0]["is_channel_used_for_cell_segmentation"] = "No"

            # Second channel (if exists) should be cell segmentation
            if num_channels > 1:
                channels[1]["is_channel_used_for_nuclei_segmentation"] = "No"
                channels[1]["is_channel_used_for_cell_segmentation"] = "Yes"

        return channels

    def _process_xenium_bundle(self, xenium_bundle_dir: Path):
        """Process xenium bundle files"""
        # Required files mapping
        required_files = {
            "cell_feature_matrix.h5": "cell_feature_matrix.h5",
            "experiment.xenium": "experiment.xenium",
            "nucleus_boundaries.parquet": "nucleus_boundaries.parquet",
            "cell_boundaries.parquet": "cell_boundaries.parquet",
            "transcripts.parquet": "transcripts.parquet",
            "cells.parquet": "cells.parquet",
        }

        # Copy required files
        for source_name, dest_name in required_files.items():
            source_path = self.input_dir / source_name
            if source_path.exists():
                shutil.copy2(source_path, xenium_bundle_dir / dest_name)
            else:
                self.validation_errors.append(
                    f"Required file {source_name} not found in input directory"
                )

        # Copy morphology focus files
        morphology_focus_dir = self.input_dir / "morphology_focus"
        if morphology_focus_dir.exists():
            for focus_file in morphology_focus_dir.glob("morphology_focus_*.ome.tif"):
                shutil.copy2(focus_file, xenium_bundle_dir / focus_file.name)
        else:
            self.validation_errors.append(
                "Required morphology_focus directory not found"
            )

        # Copy morphology MIP if available
        morphology_mip = self.input_dir / "morphology.ome.tif"
        if morphology_mip.exists():
            shutil.copy2(morphology_mip, xenium_bundle_dir / "morphology_mip.ome.tif")

        # Copy additional bundle files (zarr, etc.)
        additional_files = [
            "analysis.zarr.zip",
            "cell_feature_matrix.zarr.zip",
            "cells.zarr.zip",
            "transcripts.zarr.zip",
        ]

        for file_name in additional_files:
            source_path = self.input_dir / file_name
            if source_path.exists():
                shutil.copy2(source_path, xenium_bundle_dir / file_name)

    def update_existing_output(self, **kwargs) -> bool:
        """Update existing output directory with newly provided files"""
        try:
            print("Checking for new files to add...")

            # Only process files that don't already exist in output
            self._update_extras(
                kwargs.get("microscope_hardware_path"),
                kwargs.get("microscope_settings_path"),
            )

            self._update_raw(
                kwargs.get("overlay_image_path"),
                kwargs.get("markers_csv_path"),
                kwargs.get("additional_panels_csv_path"),
                kwargs.get("custom_probe_set_csv_path"),
                kwargs.get("custom_probe_set_bed_path"),
                kwargs.get("custom_gene_list_csv_path"),
                kwargs.get("probes_csv_path"),
            )

            self._update_lab_processed(kwargs.get("tissue_boundary_geojson_path"))

            # Validate the updated output
            is_valid = self._validate_output()

            # Generate validation report
            self._generate_validation_report()

            return is_valid

        except Exception as e:
            self.logger.error(f"Error updating output: {str(e)}")
            return False

    def _update_extras(
        self,
        microscope_hardware_path: Optional[str],
        microscope_settings_path: Optional[str],
    ):
        """Update extras folder with new files only"""
        extras_dir = self.output_dir / "extras"

        # Add microscope hardware if missing
        hardware_dest = extras_dir / "microscope_hardware.json"
        if (
            not hardware_dest.exists()
            and microscope_hardware_path
            and Path(microscope_hardware_path).exists()
        ):
            shutil.copy2(microscope_hardware_path, hardware_dest)
            print(f"Added: {hardware_dest}")

        # Add microscope settings if missing
        settings_dest = extras_dir / "microscope_settings.json"
        if (
            not settings_dest.exists()
            and microscope_settings_path
            and Path(microscope_settings_path).exists()
        ):
            shutil.copy2(microscope_settings_path, settings_dest)
            print(f"Added: {settings_dest}")

    def _update_raw(
        self,
        overlay_image_path: Optional[str],
        markers_csv_path: Optional[str],
        additional_panels_csv_path: Optional[str],
        custom_probe_set_csv_path: Optional[str],
        custom_probe_set_bed_path: Optional[str],
        custom_gene_list_csv_path: Optional[str],
        probes_csv_path: Optional[str],
    ):
        """Update raw folder with new files only"""
        raw_dir = self.output_dir / "raw"

        # Add overlay image if missing
        if overlay_image_path and Path(overlay_image_path).exists():
            overlay_path = Path(overlay_image_path)
            overlay_dest = raw_dir / f"overlay{overlay_path.suffix}"
            if not overlay_dest.exists():
                shutil.copy2(overlay_image_path, overlay_dest)
                print(f"Added: {overlay_dest}")

        # Add optional CSV files if missing
        optional_files = {
            "markers.csv": markers_csv_path,
            "additional_panels_used.csv": additional_panels_csv_path,
            "custom_probe_set.csv": custom_probe_set_csv_path,
            "custom_gene_list.csv": custom_gene_list_csv_path,
            "probes.csv": probes_csv_path,
        }

        for filename, file_path in optional_files.items():
            dest_path = raw_dir / filename
            if not dest_path.exists() and file_path and Path(file_path).exists():
                shutil.copy2(file_path, dest_path)
                print(f"Added: {dest_path}")

        # Add custom probe set BED file if missing
        if custom_probe_set_bed_path and Path(custom_probe_set_bed_path).exists():
            bed_dest = raw_dir / "custom_probe_set.bed"
            if not bed_dest.exists():
                shutil.copy2(custom_probe_set_bed_path, bed_dest)
                print(f"Added: {bed_dest}")

    def _update_lab_processed(self, tissue_boundary_geojson_path: Optional[str]):
        """Update lab_processed folder with new files only"""
        images_dir = self.output_dir / "lab_processed" / "images"

        # Add tissue boundary GeoJSON if missing
        if tissue_boundary_geojson_path and Path(tissue_boundary_geojson_path).exists():
            geojson_name = f"{self.sample_name}.tissue-boundary.geojson"
            geojson_dest = images_dir / geojson_name
            if not geojson_dest.exists():
                shutil.copy2(tissue_boundary_geojson_path, geojson_dest)
                print(f"Added: {geojson_dest}")

    def _validate_output(self) -> bool:
        """Validate the organized output against HuBMAP schema"""
        is_valid = True

        # Required file patterns
        required_patterns = [
            "raw/transcript_locations.csv",
            "raw/gene_panel.json",
            "lab_processed/images/*.ome.tiff",
            "lab_processed/images/*.ome-tiff.channels.csv",
            "lab_processed/xenium_bundle/cell_feature_matrix.h5",
            "lab_processed/xenium_bundle/experiment.xenium",
            "lab_processed/xenium_bundle/nucleus_boundaries.parquet",
            "lab_processed/xenium_bundle/cell_boundaries.parquet",
            "lab_processed/xenium_bundle/transcripts.parquet",
            "lab_processed/xenium_bundle/cells.parquet",
            "lab_processed/xenium_bundle/morphology_focus*.ome.tif",
        ]

        # Check for microscope_hardware.json separately (required for HuBMAP but not for script validation)
        hardware_file = self.output_dir / "extras" / "microscope_hardware.json"
        if not hardware_file.exists():
            self.validation_warnings.append(
                "microscope_hardware.json missing - required for HuBMAP submission. "
                "Generate using micro-meta app (contact help@hubmapconsortium.org)"
            )

        for pattern in required_patterns:
            if "*" in pattern:
                # Handle glob patterns
                matches = list(self.output_dir.glob(pattern))
                if not matches:
                    self.validation_errors.append(
                        f"Required file pattern not found: {pattern}"
                    )
                    is_valid = False
            else:
                # Handle exact paths
                if not (self.output_dir / pattern).exists():
                    self.validation_errors.append(f"Required file not found: {pattern}")
                    is_valid = False

        return is_valid and len(self.validation_errors) == 0

    def _generate_validation_report(self):
        """Generate validation report"""
        report_path = self.output_dir / "validation_report.txt"

        with open(report_path, "w") as f:
            f.write(f"HuBMAP Xenium Data Organization Validation Report\n")
            f.write(f"Sample: {self.sample_name}\n")
            f.write(f"Generated: {pd.Timestamp.now()}\n\n")

            if self.validation_errors:
                f.write("ERRORS:\n")
                for error in self.validation_errors:
                    f.write(f"- {error}\n")
                f.write("\n")

            if self.validation_warnings:
                f.write("WARNINGS:\n")
                for warning in self.validation_warnings:
                    f.write(f"- {warning}\n")
                f.write("\n")

            if not self.validation_errors and not self.validation_warnings:
                f.write("✓ All validation checks passed!\n")
            elif not self.validation_errors:
                f.write("✓ No critical errors found, but check warnings above.\n")

        self.logger.info(f"Validation report saved to: {report_path}")


def print_usage_examples():
    """Print detailed usage examples"""
    examples = """
USAGE EXAMPLES:

1. Basic usage (minimal required arguments):
   python xenium_hubmap_organizer.py \\
     --input /path/to/xenium/output-SAMPLE123 \\
     --output /path/to/hubmap_submission \\
     --sample SAMPLE123

2. With microscope hardware file:
   python xenium_hubmap_organizer.py \\
     --input /path/to/xenium/output-SAMPLE123 \\
     --output /path/to/hubmap_submission \\
     --sample SAMPLE123 \\
     --microscope-hardware /path/to/microscope_hardware.json

3. Full usage with all optional files:
   python xenium_hubmap_organizer.py \\
     --input /path/to/xenium/output-SAMPLE123 \\
     --output /path/to/hubmap_submission \\
     --sample SAMPLE123 \\
     --microscope-hardware /path/to/microscope_hardware.json \\
     --microscope-settings /path/to/microscope_settings.json \\
     --overlay-image /path/to/overlay.tiff \\
     --tissue-boundary /path/to/tissue_boundary.geojson \\
     --markers /path/to/markers.csv \\
     --additional-panels /path/to/additional_panels.csv \\
     --custom-probes /path/to/custom_probes.csv \\
     --custom-probes-bed /path/to/custom_probes.bed \\
     --custom-genes /path/to/custom_genes.csv \\
     --probes /path/to/probes.csv

4. Update existing output with microscope hardware:
   python xenium_hubmap_organizer.py \\
     --input /path/to/xenium/output-SAMPLE123 \\
     --output /path/to/hubmap_submission \\
     --sample SAMPLE123 \\
     --microscope-hardware /path/to/microscope_hardware.json \\
     --update

5. Validate existing output only:
   python xenium_hubmap_organizer.py \\
     --input /path/to/xenium/output-SAMPLE123 \\
     --output /path/to/hubmap_submission \\
     --sample SAMPLE123 \\
     --validate-only

REQUIRED FILES:
- Input Xenium directory with standard output structure

OPTIONAL FILES (but recommended for HuBMAP submission):
- microscope_hardware.json (generate using micro-meta app - contact help@hubmapconsortium.org)
- microscope_settings.json (from micro-meta app)
- overlay.tiff/jpeg (overlay image used for ROI selection)
- tissue_boundary.geojson (tissue boundary annotations)
- markers.csv (morphology markers description)
- additional_panels.csv (additional probe panels used)
- custom_probes.csv (custom probe sequences)
- custom_probes.bed (BED format custom probes)
- custom_genes.csv (custom gene list)
- probes.csv (probe panel description)

OUTPUT STRUCTURE:
hubmap_submission/
├── extras/
│   ├── microscope_hardware.json (required for HuBMAP submission)
│   └── microscope_settings.json (optional)
├── raw/
│   ├── transcript_locations.csv (auto-generated from Xenium data)
│   ├── gene_panel.json (from Xenium output)
│   └── [optional CSV files]
└── lab_processed/
    ├── images/
    │   ├── SAMPLE.ome.tiff (from morphology.ome.tif)
    │   ├── SAMPLE.ome-tiff.channels.csv (auto-generated)
    │   └── SAMPLE.tissue-boundary.geojson (optional)
    └── xenium_bundle/
        ├── cell_feature_matrix.h5
        ├── experiment.xenium
        ├── nucleus_boundaries.parquet
        ├── cell_boundaries.parquet
        ├── transcripts.parquet
        ├── cells.parquet
        ├── morphology_focus_*.ome.tif
        └── [additional zarr files]

VALIDATION:
- A validation report will be generated at: output_dir/validation_report.txt
- Check this file for any errors or warnings after running the script
- The script will warn if microscope_hardware.json is missing (required for HuBMAP submission)

TROUBLESHOOTING:
- Ensure input directory contains standard Xenium output structure
- Generate microscope_hardware.json using micro-meta app when ready for HuBMAP submission
- Contact HuBMAP Help Desk (help@hubmapconsortium.org) for micro-meta assistance
- Check validation_report.txt for specific missing files or errors
- Install tifffile package for proper channel extraction: pip install tifffile
"""
    print(examples)


def create_argument_parser():
    """Create and configure argument parser with detailed help"""
    parser = argparse.ArgumentParser(
        description="Organize Xenium data products for HuBMAP submission compliance",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
For detailed usage examples, run: python xenium_hubmap_organizer.py --examples

This script organizes Xenium spatial transcriptomics data according to HuBMAP 
directory schema v3.6. It takes a standard Xenium output directory and reorganizes 
the files into the required HuBMAP structure with proper validation.

Required inputs:
- Xenium output directory (standard structure from Xenium Onboard Analysis)

Optional but recommended for HuBMAP submission:
- microscope_hardware.json file (generated using micro-meta app)

The script will:
1. Create HuBMAP-compliant directory structure
2. Copy and organize files to appropriate locations
3. Convert transcript data to required format
4. Generate channel metadata for OME-TIFF files with proper Xenium channel configuration
5. Validate output against schema requirements
6. Generate detailed validation report

Contact HuBMAP Consortium Help Desk (help@hubmapconsortium.org) for assistance
with generating microscope metadata files using the micro-meta app.
        """,
    )

    # Required arguments
    required = parser.add_argument_group("required arguments")
    required.add_argument(
        "--input",
        "-i",
        required=True,
        help="Path to Xenium output directory (e.g., output-SAMPLE123__timestamp)",
    )
    required.add_argument(
        "--output",
        "-o",
        required=True,
        help="Path to output directory for HuBMAP-compliant structure",
    )
    required.add_argument(
        "--sample", "-s", required=True, help="Sample identifier/name for file naming"
    )

    # Optional arguments
    optional = parser.add_argument_group("optional arguments")
    optional.add_argument(
        "--microscope-hardware",
        help="Path to microscope_hardware.json file (required for HuBMAP submission)",
    )
    optional.add_argument(
        "--microscope-settings",
        help="Path to microscope_settings.json file (from micro-meta app)",
    )
    optional.add_argument(
        "--overlay-image",
        help="Path to overlay image file (JPEG or TIFF) used for ROI selection",
    )
    optional.add_argument(
        "--tissue-boundary", help="Path to tissue boundary GeoJSON file"
    )
    optional.add_argument(
        "--markers", help="Path to markers.csv file (morphology markers description)"
    )
    optional.add_argument(
        "--additional-panels", help="Path to additional_panels_used.csv file"
    )
    optional.add_argument("--custom-probes", help="Path to custom_probe_set.csv file")
    optional.add_argument(
        "--custom-probes-bed", help="Path to custom_probe_set.bed file"
    )
    optional.add_argument("--custom-genes", help="Path to custom_gene_list.csv file")
    optional.add_argument(
        "--probes", help="Path to probes.csv file (probe panel description)"
    )

    # Utility arguments
    parser.add_argument(
        "--examples", action="store_true", help="Show detailed usage examples and exit"
    )
    parser.add_argument(
        "--validate-only",
        action="store_true",
        help="Only validate existing output directory without reorganizing",
    )
    parser.add_argument(
        "--update",
        action="store_true",
        help="Update existing output directory with newly provided files",
    )
    parser.add_argument(
        "--verbose", "-v", action="store_true", help="Enable verbose logging output"
    )

    return parser


def main():
    """Main function with command-line interface"""
    parser = create_argument_parser()

    # Show examples if requested
    if len(sys.argv) == 2 and sys.argv[1] == "--examples":
        print_usage_examples()
        sys.exit(0)

    args = parser.parse_args()

    # Show examples and exit
    if args.examples:
        print_usage_examples()
        sys.exit(0)

    # Configure logging
    if args.verbose:
        logging.basicConfig(level=logging.DEBUG)
    else:
        logging.basicConfig(level=logging.INFO)

    # Validate input directory exists
    if not Path(args.input).exists():
        print(f"ERROR: Input directory does not exist: {args.input}")
        sys.exit(1)

    # Create organizer instance
    organizer = XeniumHubMapOrganizer(args.input, args.output, args.sample)

    if args.validate_only:
        # Only validate existing output
        if not Path(args.output).exists():
            print(
                f"ERROR: Output directory does not exist for validation: {args.output}"
            )
            sys.exit(1)

        print("Validating existing output directory...")
        is_valid = organizer._validate_output()
        organizer._generate_validation_report()

        if is_valid:
            print("✓ Validation successful!")
            if organizer.validation_warnings:
                print("⚠ Check validation_report.txt for warnings.")
            sys.exit(0)
        else:
            print("✗ Validation failed. Check validation_report.txt for details.")
            sys.exit(1)

    elif args.update:
        # Update existing output with new files
        if not Path(args.output).exists():
            print(f"ERROR: Output directory does not exist for update: {args.output}")
            print("Use without --update flag to create new output directory.")
            sys.exit(1)

        print(f"Updating existing output directory: {args.output}")
        print("Only adding newly provided files...")

        success = organizer.update_existing_output(
            microscope_hardware_path=args.microscope_hardware,
            microscope_settings_path=args.microscope_settings,
            overlay_image_path=args.overlay_image,
            tissue_boundary_geojson_path=args.tissue_boundary,
            markers_csv_path=args.markers,
            additional_panels_csv_path=args.additional_panels,
            custom_probe_set_csv_path=args.custom_probes,
            custom_probe_set_bed_path=args.custom_probes_bed,
            custom_gene_list_csv_path=args.custom_genes,
            probes_csv_path=args.probes,
        )

        if success:
            print("✓ Output directory updated successfully!")
            if organizer.validation_warnings:
                print("⚠ Check validation_report.txt for warnings.")
            print(f"Check validation report at: {args.output}/validation_report.txt")
            sys.exit(0)
        else:
            print("✗ Update failed!")
            print(
                f"Check validation report at: {args.output}/validation_report.txt for details"
            )
            sys.exit(1)

    # Regular organization (create new output)
    print(f"Organizing Xenium data from: {args.input}")
    print(f"Output directory: {args.output}")
    print(f"Sample name: {args.sample}")

    success = organizer.organize_data(
        microscope_hardware_path=args.microscope_hardware,
        microscope_settings_path=args.microscope_settings,
        overlay_image_path=args.overlay_image,
        tissue_boundary_geojson_path=args.tissue_boundary,
        markers_csv_path=args.markers,
        additional_panels_csv_path=args.additional_panels,
        custom_probe_set_csv_path=args.custom_probes,
        custom_probe_set_bed_path=args.custom_probes_bed,
        custom_gene_list_csv_path=args.custom_genes,
        probes_csv_path=args.probes,
    )

    if success:
        print("✓ Data organization completed successfully!")
        if organizer.validation_warnings:
            print(
                "⚠ Check validation_report.txt for warnings about HuBMAP submission requirements."
            )
        print(f"Check validation report at: {args.output}/validation_report.txt")
        sys.exit(0)
    else:
        print("✗ Data organization failed!")
        print(
            f"Check validation report at: {args.output}/validation_report.txt for details"
        )
        sys.exit(1)


# Legacy function for backward compatibility
def organize_xenium_data(
    input_dir: str, output_dir: str, sample_name: str, **kwargs
) -> bool:
    """
    Organize Xenium data for HuBMAP submission

    Args:
        input_dir: Path to Xenium output directory
        output_dir: Path to organized output directory
        sample_name: Sample identifier
        **kwargs: Optional file paths for additional files

    Returns:
        bool: True if organization successful and valid
    """
    organizer = XeniumHubMapOrganizer(input_dir, output_dir, sample_name)
    return organizer.organize_data(**kwargs)


if __name__ == "__main__":
    main()