process_single_file.py

#!/usr/bin/env python3
"""
Process a single h5ad file using kompot for differential abundance and expression analysis.
Based on prototype_processing.ipynb
"""

import argparse
import os
import sys
import importlib.util
from pathlib import Path
import yaml
import logging
import platform
import time
from datetime import datetime

import anndata as ad
import numpy as np
import pandas as pd
import scanpy as sc
import palantir
import kompot


def setup_logging(log_file):
    """Setup logging to both file and console, including kompot and mellon loggers."""
    # Configure root logger
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(log_file),
            logging.StreamHandler()
        ]
    )
    
    # Ensure kompot and mellon loggers write to our log file
    kompot_logger = logging.getLogger('kompot')
    mellon_logger = logging.getLogger('mellon')
    
    # Add file handler to kompot and mellon loggers if not already present
    file_handler = logging.FileHandler(log_file)
    file_handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s'))
    
    kompot_logger.addHandler(file_handler)
    mellon_logger.addHandler(file_handler)
    
    # Set appropriate log levels
    kompot_logger.setLevel(logging.INFO)
    mellon_logger.setLevel(logging.INFO)


def get_environment_info():
    """Collect environment and package version information."""
    try:
        from importlib.metadata import version, PackageNotFoundError
    except ImportError:
        # Fallback for Python < 3.8
        from importlib_metadata import version, PackageNotFoundError
    
    env_info = {
        'platform': platform.platform(),
        'python_version': platform.python_version(),
        'hostname': platform.node(),
        'timestamp': datetime.now().isoformat(),
        'packages': {}
    }
    
    # Get versions of key packages
    key_packages = ['anndata', 'scanpy', 'numpy', 'pandas', 'kompot', 'mellon']
    for pkg in key_packages:
        try:
            env_info['packages'][pkg] = version(pkg)
        except PackageNotFoundError:
            env_info['packages'][pkg] = 'not found'
    
    return env_info


def load_config(config_path):
    """Load configuration from Python module."""
    spec = importlib.util.spec_from_file_location("config", config_path)
    config = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(config)
    return config


def create_config_object(dataset_config):
    """Convert dataset config dict to object with attributes."""
    class Config:
        pass

    config = Config()

    # Map centralized config keys to expected attribute names
    config.SAMPLE_COLUMN = dataset_config["sample_column"]
    config.GROUPING_COLUMN = dataset_config["grouping_column"]
    config.CONDITIONS = dataset_config["conditions"]
    config.CELL_TYPE_COLUMN = dataset_config["cell_type_column"]
    config.DIMENSIONALITY_REDUCTION = dataset_config["dimensionality_reduction"]
    config.LAYER_FOR_EXPRESSION = dataset_config["layer_for_expression"]
    config.N_TOP_GENES = dataset_config.get("n_top_genes", None)  # HVG filtering before analysis
    config.TOPN_GENES = dataset_config["topn_genes"]
    config.BATCH_SIZE = dataset_config["batch_size"]
    config.STORE_ARRAYS_ON_DISK = dataset_config["store_arrays_on_disk"]

    # PCA configuration - read from centralized config with defaults
    config.PCA_KEY = dataset_config.get("pca_key", "X_pca")
    config.PCA_N_COMPONENTS = dataset_config.get("pca_n_components", 50)
    config.FORCE_RECOMPUTE_PCA = dataset_config.get("force_recompute_pca", False)
    config.SKIP_PCA = dataset_config.get("skip_pca", False)

    # Diffusion maps configuration
    config.DIFFUSION_MAPS_KEY = dataset_config.get("dimensionality_reduction", "DM_EigenVectors")
    config.FORCE_RECOMPUTE_DIFFUSION_MAPS = dataset_config.get("force_recompute_diffusion_maps", False)
    config.diffusion_maps_n_components = dataset_config.get("diffusion_maps_n_components", 10)

    return config


def get_adata_info(adata):
    """Get information about the AnnData object."""
    return {
        'n_obs': adata.n_obs,
        'n_vars': adata.n_vars,
        'obs_columns': list(adata.obs.columns),
        'var_columns': list(adata.var.columns),
        'obsm_keys': list(adata.obsm.keys()) if adata.obsm else [],
        'layers': list(adata.layers.keys()) if adata.layers else [],
        'uns_keys': list(adata.uns.keys()) if adata.uns else []
    }


def get_kompot_params(config, exclude_for='none'):
    """
    Extract kompot-relevant parameters from config using negative selection.

    This extracts ALL lowercase parameters from config EXCEPT those we know are
    NOT kompot parameters. This way, new kompot parameters are automatically
    picked up without code changes.

    Args:
        config: Configuration object
        exclude_for: 'da' to exclude DE-only params, 'de' to exclude DA-only params,
                     'none' to include all kompot params

    Returns:
        dict: Parameters with lowercase keys for kompot function calls
    """
    # Parameters that are our config-specific, NOT kompot parameters
    NON_KOMPOT_PARAMS = {
        'SAMPLE_COLUMN', 'GROUPING_COLUMN', 'CONDITIONS', 'CELL_TYPE_COLUMN',
        'DIMENSIONALITY_REDUCTION', 'LAYER_FOR_EXPRESSION', 'N_TOP_GENES',
        'TOPN_GENES', 'PCA_KEY', 'PCA_N_COMPONENTS', 'FORCE_RECOMPUTE_PCA',
        'SKIP_PCA', 'DIFFUSION_MAPS_KEY', 'FORCE_RECOMPUTE_DIFFUSION_MAPS',
        'diffusion_maps_n_components', 'n_components'
    }

    # Parameters exclusive to specific functions (for filtering)
    DE_ONLY_PARAMS = {'sigma', 'batch_size', 'store_arrays_on_disk'}
    DA_ONLY_PARAMS = set()  # Currently no DA-exclusive params

    # Extract all parameters from config
    params = {}
    for attr in dir(config):
        if attr.startswith('_'):
            continue
        if attr in NON_KOMPOT_PARAMS:
            continue

        value = getattr(config, attr, None)
        # Skip methods and None values
        if callable(value) or value is None:
            continue

        # Convert to lowercase for kompot function parameters
        param_key = attr.lower()

        # Apply exclusions based on function type
        if exclude_for == 'da' and param_key in DE_ONLY_PARAMS:
            continue
        if exclude_for == 'de' and param_key in DA_ONLY_PARAMS:
            continue

        params[param_key] = value

    return params


def main():
    parser = argparse.ArgumentParser(description='Process single h5ad file with kompot')
    parser.add_argument('input_file', help='Input h5ad file path')
    parser.add_argument('--dataset', required=True, help='Dataset name (e.g., aging)')
    parser.add_argument('--config-override', action='append', default=[],
                       help='Override config values (format: key=value)')
    parser.add_argument('--output-dir', help='Output directory (default: data/processed/{dataset})')
    parser.add_argument('--log-dir', help='Log directory (default: same as output)')
    parser.add_argument('--skip-sample-variance', action='store_true',
                       help='Skip sample variance computation (DA and DE with sample_col)')

    args = parser.parse_args()
    
    # Setup paths
    input_path = Path(args.input_file)
    if not input_path.exists():
        print(f"Error: Input file {input_path} does not exist")
        sys.exit(1)
    
    # Default output directory
    if args.output_dir:
        output_dir = Path(args.output_dir)
    else:
        output_dir = Path(f"data/processed/{args.dataset}")
    
    output_dir.mkdir(parents=True, exist_ok=True)
    
    # Output file path
    output_file = output_dir / input_path.name
    
    # Log directory
    if args.log_dir:
        log_dir = Path(args.log_dir)
    else:
        log_dir = output_dir
    
    log_dir.mkdir(parents=True, exist_ok=True)
    
    # Load configuration from centralized processing_config.py
    try:
        from processing_config import resolve_dataset_config
        _, dataset_config = resolve_dataset_config(args.dataset)
        config = create_config_object(dataset_config)
    except Exception as e:
        print(f"Failed to load centralized config for dataset '{args.dataset}': {e}")
        sys.exit(1)
    
    # Handle parameter stability using SLURM_ARRAY_TASK_ID
    array_task_id = os.environ.get('SLURM_ARRAY_TASK_ID')
    parameter_stability = os.environ.get('PARAMETER_STABILITY')
    parameter_type = os.environ.get('PARAMETER_TYPE')
    
    # Collect parameter variation information
    parameter_variation = {}
    if parameter_stability == 'true' and parameter_type:
        # Get parameter metadata from environment
        parameter_variation = {
            'is_parameter_stability': True,
            'parameter_type': parameter_type,
            'default_value': os.environ.get('PARAM_DEFAULT_VALUE'),
            'description': os.environ.get('PARAM_DESCRIPTION'),
            'config_type': os.environ.get('PARAM_CONFIG_TYPE'),
            'sequence_type': os.environ.get('PARAM_SEQUENCE_TYPE'),
            'array_task_id': array_task_id
        }
        
        # Add sequence-specific information
        if parameter_variation['sequence_type'] == 'integer_range':
            parameter_variation.update({
                'min_value': os.environ.get('PARAM_MIN_VALUE'),
                'max_value': os.environ.get('PARAM_MAX_VALUE'),
                'total_values': os.environ.get('PARAM_TOTAL_VALUES')
            })
        elif parameter_variation['sequence_type'] == 'exponential':
            parameter_variation.update({
                'baseline': os.environ.get('PARAM_BASELINE'),
                'decay_rate': os.environ.get('PARAM_DECAY_RATE'),
                'growth_rate': os.environ.get('PARAM_GROWTH_RATE'),
                'n_steps_down': os.environ.get('PARAM_N_STEPS_DOWN'),
                'n_steps_up': os.environ.get('PARAM_N_STEPS_UP')
            })
    else:
        parameter_variation['is_parameter_stability'] = False
    
    # Create log file name that includes parameter information to avoid overwrites
    log_base_name = input_path.stem
    if array_task_id is not None and parameter_stability == 'true' and parameter_type:
        log_base_name = f"{log_base_name}_{parameter_type}_task_{array_task_id}"
    
    log_file = log_dir / f"{log_base_name}.log"
    
    # Setup logging
    setup_logging(log_file)
    logger = logging.getLogger(__name__)
    
    logger.info(f"Processing {input_path} -> {output_file}")
    logger.info(f"Loaded centralized config for dataset: {args.dataset}")
    
    if array_task_id is not None and parameter_stability == 'true' and parameter_type:
        try:
            task_id = int(array_task_id)
            logger.info(f"SLURM array task ID: {task_id}")
            logger.info(f"Parameter stability run for: {parameter_type}")
            
            # Apply parameter stability value based on task ID
            from processing_config import STABILITY_CONFIGS, generate_parameter_sequence, get_subsampling_series
            
            if parameter_type == "subsampling":
                # Special case: task 1 = original, task 2+ = subsampled files
                # Use consolidated directory structure: {dataset}_subsampling/
                output_dir = output_dir.parent / f"{args.dataset}_subsampling"
                output_dir.mkdir(parents=True, exist_ok=True)
                
                if task_id == 1:
                    logger.info("Using original dataset (task 1)")
                    # Create filename with original indicator
                    stem = input_path.stem
                    suffix = input_path.suffix
                    output_file = output_dir / f"{stem}_subsampling_original{suffix}"
                    parameter_variation['current_value'] = "original"
                else:
                    # Use subsampled file
                    subsampled_files = get_subsampling_series(args.dataset)
                    if task_id - 2 < len(subsampled_files):  # -2 because task 1 is original
                        subsampled_file = subsampled_files[task_id - 2]
                        round_num = task_id - 1  # task 2 -> round 1, task 3 -> round 2, etc.
                        logger.info(f"Using subsampled file: {subsampled_file} (round {round_num})")
                        
                        # Update input path and create filename with round indicator
                        input_path = Path(subsampled_file)
                        stem = input_path.stem
                        suffix = input_path.suffix
                        output_file = output_dir / f"{stem}_subsampling_subsampled_round_{round_num:03d}{suffix}"
                        parameter_variation['current_value'] = f"subsampled_round_{round_num:03d}"
            else:
                # Regular parameter sweep
                # Use consolidated directory structure: {dataset}_{parameter_type}/
                output_dir = output_dir.parent / f"{args.dataset}_{parameter_type}"
                output_dir.mkdir(parents=True, exist_ok=True)

                stability_config = STABILITY_CONFIGS[parameter_type].copy()

                # For dataset-specific parameters (like topn_genes), resolve max_value
                if stability_config.get('max_value') is None and parameter_type == 'topn_genes':
                    from processing_config import get_dataset_gene_count
                    stability_config['max_value'] = get_dataset_gene_count(args.dataset)
                    logger.info(f"Resolved topn_genes max_value to {stability_config['max_value']} for dataset {args.dataset}")

                parameter_values = generate_parameter_sequence(stability_config)
                
                if 1 <= task_id <= len(parameter_values):
                    param_value = parameter_values[task_id - 1]  # -1 for 0-based indexing
                    logger.info(f"Using {parameter_type} = {param_value}")
                    
                    # Format parameter value for filename
                    if isinstance(param_value, float):
                        # Use scientific notation for very small/large numbers
                        if abs(param_value) < 0.001 or abs(param_value) > 1000:
                            value_str = f"{param_value:.2e}"
                        else:
                            value_str = f"{param_value:.3f}".rstrip('0').rstrip('.')
                    else:
                        value_str = str(param_value)
                    
                    # Create filename with parameter value
                    stem = input_path.stem
                    suffix = input_path.suffix
                    output_file = output_dir / f"{stem}_{parameter_type}_{value_str}{suffix}"
                    
                    # Apply parameter value and record current value
                    if parameter_type == "ls_factor":
                        setattr(config, "LS_FACTOR", param_value)
                        parameter_variation['current_value'] = param_value
                    elif parameter_type == "sigma":
                        setattr(config, "SIGMA", param_value)
                        parameter_variation['current_value'] = param_value
                    elif parameter_type == "n_landmarks":
                        setattr(config, "N_LANDMARKS", int(param_value))
                        parameter_variation['current_value'] = int(param_value)
                    elif parameter_type == "n_components":
                        # This will need special handling in diffusion maps computation
                        setattr(config, "n_components", int(param_value))
                        parameter_variation['current_value'] = int(param_value)
                    elif parameter_type == "topn_genes":
                        setattr(config, "TOPN_GENES", int(param_value))
                        parameter_variation['current_value'] = int(param_value)
                            
        except Exception as e:
            logger.error(f"Error handling parameter stability: {e}")
            # Continue with normal processing
    
    # Apply config overrides
    for override in args.config_override:
        key, value = override.split('=', 1)
        # Try to evaluate as Python literal, fallback to string
        try:
            value = eval(value)
        except:
            pass
        setattr(config, key, value)
        logger.info(f"Config override: {key} = {value}")
    
    # Load data
    logger.info("Loading AnnData object...")
    adata = ad.read_h5ad(input_path)
    logger.info(f"Loaded data: {adata.n_obs} cells × {adata.n_vars} genes")

    # Apply HVG filtering if configured
    n_top_genes = getattr(config, 'N_TOP_GENES', None)
    if n_top_genes is not None and n_top_genes > 0:
        logger.info(f"Filtering to top {n_top_genes} highly variable genes...")

        # Compute HVGs using scanpy
        if 'highly_variable' not in adata.var.columns:
            logger.info("Computing highly variable genes...")
            sc.pp.highly_variable_genes(
                adata,
                n_top_genes=n_top_genes,
                subset=False,  # Don't subset yet, just mark
                flavor='seurat_v3'
            )

        # Subset to HVGs
        original_n_vars = adata.n_vars
        adata = adata[:, adata.var['highly_variable']].copy()
        logger.info(f"Subsetted from {original_n_vars} to {adata.n_vars} highly variable genes")

        # Clear PCA and diffusion maps as they need recomputation after subsetting
        if 'X_pca' in adata.obsm:
            del adata.obsm['X_pca']
            logger.info("Removed existing PCA (needs recomputation after HVG filtering)")
        if 'DM_EigenVectors' in adata.obsm:
            del adata.obsm['DM_EigenVectors']
            logger.info("Removed existing diffusion maps (needs recomputation after HVG filtering)")

    # Get data info for runinfo
    adata_info = get_adata_info(adata)
    
    # Prepare config dict for saving
    config_dict = {k: v for k, v in config.__dict__.items() if not k.startswith('_')}
    
    # Collect environment info
    env_info = get_environment_info()
    
    # Resolve SLURM log paths if running in SLURM
    slurm_info = {}
    if 'SLURM_JOB_ID' in os.environ:
        # Basic job identification
        slurm_info['job_id'] = os.environ.get('SLURM_JOB_ID')
        slurm_info['array_job_id'] = os.environ.get('SLURM_ARRAY_JOB_ID')
        slurm_info['array_task_id'] = os.environ.get('SLURM_ARRAY_TASK_ID')
        slurm_info['job_name'] = os.environ.get('SLURM_JOB_NAME')
        slurm_info['submit_dir'] = os.environ.get('SLURM_SUBMIT_DIR')
        slurm_info['submit_host'] = os.environ.get('SLURM_SUBMIT_HOST')
        
        # Resource allocation
        slurm_info['cpus_per_task'] = os.environ.get('SLURM_CPUS_PER_TASK')
        slurm_info['mem_per_node'] = os.environ.get('SLURM_MEM_PER_NODE')
        slurm_info['mem_per_cpu'] = os.environ.get('SLURM_MEM_PER_CPU')
        slurm_info['ntasks'] = os.environ.get('SLURM_NTASKS')
        slurm_info['nodes'] = os.environ.get('SLURM_JOB_NUM_NODES')
        
        # Node and partition info
        slurm_info['nodelist'] = os.environ.get('SLURM_JOB_NODELIST')
        slurm_info['partition'] = os.environ.get('SLURM_JOB_PARTITION')
        slurm_info['qos'] = os.environ.get('SLURM_JOB_QOS')
        slurm_info['account'] = os.environ.get('SLURM_JOB_ACCOUNT')
        
        # Timing and limits
        slurm_info['time_limit'] = os.environ.get('SLURM_JOB_TIME_LIMIT')
        slurm_info['start_time'] = os.environ.get('SLURM_JOB_START_TIME')
        
        # Array job specifics
        if slurm_info['array_job_id']:
            slurm_info['array_task_count'] = os.environ.get('SLURM_ARRAY_TASK_COUNT')
            slurm_info['array_task_min'] = os.environ.get('SLURM_ARRAY_TASK_MIN')
            slurm_info['array_task_max'] = os.environ.get('SLURM_ARRAY_TASK_MAX')
            slurm_info['array_task_step'] = os.environ.get('SLURM_ARRAY_TASK_STEP')
        
        # Working directory and user
        slurm_info['work_dir'] = os.environ.get('SLURM_WORKING_DIR')
        slurm_info['user'] = os.environ.get('SLURM_JOB_USER')
        slurm_info['uid'] = os.environ.get('SLURM_JOB_UID')
        
        # Resolve SLURM log paths using the patterns passed via environment variables
        # These are set by processing_config.py and passed through --export
        log_out_pattern = os.environ.get('SLURM_LOG_OUT_PATTERN')
        log_err_pattern = os.environ.get('SLURM_LOG_ERR_PATTERN')
        
        if log_out_pattern and log_err_pattern:
            # Expand SLURM variables %A (array job ID) and %a (array task ID)
            if slurm_info['array_job_id'] and slurm_info['array_task_id']:
                expanded_out = log_out_pattern.replace('%A', slurm_info['array_job_id']).replace('%a', slurm_info['array_task_id'])
                expanded_err = log_err_pattern.replace('%A', slurm_info['array_job_id']).replace('%a', slurm_info['array_task_id'])
            elif slurm_info['job_id']:
                # Non-array job
                expanded_out = log_out_pattern.replace('%A', slurm_info['job_id']).replace('_%a', '')
                expanded_err = log_err_pattern.replace('%A', slurm_info['job_id']).replace('_%a', '')
            else:
                expanded_out = log_out_pattern
                expanded_err = log_err_pattern
            
            # Resolve to absolute paths
            slurm_info['log_out'] = str(Path(expanded_out).resolve())
            slurm_info['log_err'] = str(Path(expanded_err).resolve())
            slurm_info['log_out_pattern'] = log_out_pattern
            slurm_info['log_err_pattern'] = log_err_pattern
        
        # Add job timestamp if available
        job_timestamp = os.environ.get('JOB_TIMESTAMP')
        if job_timestamp:
            slurm_info['job_timestamp'] = job_timestamp
        
        logger.info(f"SLURM job info captured: Job {slurm_info['job_id']}, Node {slurm_info['nodelist']}, Partition {slurm_info['partition']}")
    
    # Initialize timing info
    timing_info = {}
    start_time = time.time()
    
    # Compute PCA if needed
    skip_pca = getattr(config, 'SKIP_PCA', False)
    pca_key = getattr(config, 'PCA_KEY', 'X_pca')
    force_recompute_pca = getattr(config, 'FORCE_RECOMPUTE_PCA', False)
    
    if skip_pca:
        logger.info("Skipping PCA computation (SKIP_PCA=True). Using expression matrix directly.")
        # Store expression matrix in PCA key for downstream compatibility
        if config.LAYER_FOR_EXPRESSION:
            adata.obsm[pca_key] = adata.layers[config.LAYER_FOR_EXPRESSION].copy()
        else:
            adata.obsm[pca_key] = adata.X.copy()
        
        # Ensure it's dense array
        if hasattr(adata.obsm[pca_key], 'toarray'):
            adata.obsm[pca_key] = adata.obsm[pca_key].toarray()
            
        timing_info['pca_seconds'] = 0
        logger.info(f"Using expression matrix as '{pca_key}' with shape {adata.obsm[pca_key].shape}")
    else:
        # Check if we need to recompute PCA due to component count mismatch
        requested_pca_components = getattr(config, 'PCA_N_COMPONENTS', 50)
        need_pca_recompute = force_recompute_pca or pca_key not in adata.obsm
        
        if pca_key in adata.obsm and not force_recompute_pca:
            existing_pca_components = adata.obsm[pca_key].shape[1]
            if existing_pca_components != requested_pca_components:
                logger.info(f"PCA component count mismatch: existing={existing_pca_components}, requested={requested_pca_components}. Recomputing PCA.")
                need_pca_recompute = True
        
        if need_pca_recompute:
            logger.info("Computing PCA...")
            step_start = time.time()
            
            # Get number of components (adaptive or fixed)
            if hasattr(config, 'get_n_pca_components') and callable(config.get_n_pca_components):
                n_components = config.get_n_pca_components(adata.n_obs, adata.n_vars)
            else:
                n_components = requested_pca_components
            
            # Validate PCA components don't exceed features
            max_components = min(adata.n_obs - 1, adata.n_vars)
            if n_components > max_components:
                logger.warning(f"Requested PCA components ({n_components}) > max possible ({max_components}). Using {max_components} components.")
                n_components = max_components
            
            logger.info(f"Computing PCA with {n_components} components")
            sc.tl.pca(adata, n_comps=n_components)
            
            # Store in specified key if different from default
            if pca_key != 'X_pca':
                adata.obsm[pca_key] = adata.obsm['X_pca'].copy()
                
            timing_info['pca_seconds'] = time.time() - step_start
            logger.info(f"PCA completed in {timing_info['pca_seconds']:.2f}s")
        else:
            logger.info(f"PCA already present in {pca_key}")
            timing_info['pca_seconds'] = 0
    
    # Compute diffusion maps if needed
    dm_key = getattr(config, 'DIFFUSION_MAPS_KEY', 'DM_EigenVectors')
    force_recompute_dm = getattr(config, 'FORCE_RECOMPUTE_DIFFUSION_MAPS', False)
    
    # Check if we need to recompute diffusion maps due to:
    # 1. Component count mismatch with requested n_components (unless None - preserve existing)
    # 2. PCA was recomputed (diffusion maps depend on PCA)
    requested_dm_components = getattr(config, 'n_components', getattr(config, 'diffusion_maps_n_components', 10))
    need_dm_recompute = force_recompute_dm or dm_key not in adata.obsm or need_pca_recompute
    
    if dm_key in adata.obsm and not force_recompute_dm and not need_pca_recompute:
        existing_dm_components = adata.obsm[dm_key].shape[1]
        
        # If requested_dm_components is None, preserve existing diffusion maps regardless of size
        if requested_dm_components is not None and existing_dm_components != requested_dm_components:
            logger.info(f"Diffusion maps component count mismatch: existing={existing_dm_components}, requested={requested_dm_components}. Recomputing diffusion maps.")
            need_dm_recompute = True
        elif requested_dm_components is None:
            logger.info(f"Preserving existing diffusion maps with {existing_dm_components} components (no specific count requested)")
            # Update requested_dm_components to match existing for downstream consistency
            requested_dm_components = existing_dm_components
    
    if need_dm_recompute:
        logger.info("Computing diffusion maps...")
        step_start = time.time()
        
        # Get parameters (adaptive or defaults)
        if hasattr(config, 'get_diffusion_maps_params') and callable(config.get_diffusion_maps_params):
            dm_params = config.get_diffusion_maps_params(adata.n_obs, adata.n_vars)
        else:
            dm_params = {}
        
        # Add n_components parameter for diffusion maps
        dm_params['n_components'] = requested_dm_components

        # If PCA was skipped, use expression matrix directly for diffusion maps
        if skip_pca:
            logger.info("Using expression matrix directly for diffusion maps (PCA skipped)")
            if config.LAYER_FOR_EXPRESSION:
                dm_input = adata.layers[config.LAYER_FOR_EXPRESSION]
            else:
                dm_input = adata.X

            # Ensure it's dense for palantir
            if hasattr(dm_input, 'toarray'):
                dm_input = dm_input.toarray()

            # Store in temporary obsm key for palantir
            temp_key = "X_diffusion_input"
            adata.obsm[temp_key] = dm_input
            dm_params['pca_key'] = temp_key
        else:
            # Use configured PCA key for diffusion maps
            dm_params['pca_key'] = pca_key

        logger.info(f"Computing diffusion maps with {requested_dm_components} components using PCA key '{dm_params.get('pca_key', pca_key)}'")
        palantir.utils.run_diffusion_maps(adata, **dm_params)
        
        # Clean up temporary key if used
        if skip_pca and "X_diffusion_input" in adata.obsm:
            del adata.obsm["X_diffusion_input"]
        
        # Store in specified key if different from default
        if dm_key != 'DM_EigenVectors' and 'DM_EigenVectors' in adata.obsm:
            adata.obsm[dm_key] = adata.obsm['DM_EigenVectors'].copy()
            
        timing_info['diffusion_maps_seconds'] = time.time() - step_start
        logger.info(f"Diffusion maps completed in {timing_info['diffusion_maps_seconds']:.2f}s")
    else:
        logger.info(f"Diffusion maps already present in {dm_key}")
        timing_info['diffusion_maps_seconds'] = 0
    
    # Processing steps based on notebook - with error handling
    timing_info['errors'] = []

    # Check if DA computation should be skipped (e.g., for topn_genes parameter stability)
    skip_da = os.environ.get('SKIP_DA_COMPUTATION') == 'true'

    if skip_da:
        logger.info("Skipping DA computation (SKIP_DA_COMPUTATION=true) - DA is invariant to gene selection")
        timing_info['da_no_sample_seconds'] = 0
        timing_info['da_no_sample_success'] = None
        timing_info['da_no_sample_skipped'] = True
        timing_info['da_with_sample_seconds'] = 0
        timing_info['da_with_sample_success'] = None
        timing_info['da_with_sample_skipped'] = True
    else:
        # DA without sample variance
        logger.info("Computing differential abundance without sample variance...")
        step_start = time.time()
        try:
            da_results_no_sample = kompot.compute_differential_abundance(
                adata,
                groupby=config.GROUPING_COLUMN,
                condition1=config.CONDITIONS[0],
                condition2=config.CONDITIONS[1],
                obsm_key=config.DIMENSIONALITY_REDUCTION,
                **get_kompot_params(config, exclude_for='da')
            )
            timing_info['da_no_sample_seconds'] = time.time() - step_start
            timing_info['da_no_sample_success'] = True
            logger.info(f"DA without sample variance completed in {timing_info['da_no_sample_seconds']:.2f}s")
        except Exception as e:
            timing_info['da_no_sample_seconds'] = time.time() - step_start
            timing_info['da_no_sample_success'] = False
            error_msg = f"DA without sample variance failed: {str(e)}"
            timing_info['errors'].append(error_msg)
            logger.error(error_msg, exc_info=True)

        # DA with sample variance (optional)
        if args.skip_sample_variance:
            logger.info("Skipping DA with sample variance (--skip-sample-variance flag set)")
            timing_info['da_with_sample_seconds'] = 0
            timing_info['da_with_sample_success'] = None
            timing_info['da_with_sample_skipped'] = True
        else:
            logger.info("Computing differential abundance with sample variance...")
            step_start = time.time()
            try:
                da_results_with_sample = kompot.compute_differential_abundance(
                    adata,
                    groupby=config.GROUPING_COLUMN,
                    condition1=config.CONDITIONS[0],
                    condition2=config.CONDITIONS[1],
                    obsm_key=config.DIMENSIONALITY_REDUCTION,
                    sample_col=config.SAMPLE_COLUMN,
                    **get_kompot_params(config, exclude_for='da')
                )
                timing_info['da_with_sample_seconds'] = time.time() - step_start
                timing_info['da_with_sample_success'] = True
                logger.info(f"DA with sample variance completed in {timing_info['da_with_sample_seconds']:.2f}s")
            except Exception as e:
                timing_info['da_with_sample_seconds'] = time.time() - step_start
                timing_info['da_with_sample_success'] = False
                error_msg = f"DA with sample variance failed: {str(e)}"
                timing_info['errors'].append(error_msg)
                logger.error(error_msg, exc_info=True)
    
    # Select genes for DE analysis
    genes_for_sample_de = None  # Used for DE with sample variance (old behavior)
    genes_for_no_sample_de = None  # Used for DE without sample variance (only for topn_genes parameter stability)
    topn = adata.n_vars

    if hasattr(config, 'TOPN_GENES') and config.TOPN_GENES is not None:
        n_genes_available = adata.n_vars
        topn = min(config.TOPN_GENES, n_genes_available)
        logger.info(f"Computing gene variance to select top {topn} genes out of {n_genes_available} available")

        # Compute gene variance for selection
        layer_data = adata.layers[config.LAYER_FOR_EXPRESSION] if config.LAYER_FOR_EXPRESSION else adata.X
        if hasattr(layer_data, 'toarray'):
            layer_data = layer_data.toarray()

        gene_var = np.var(layer_data, axis=0)
        top_var_genes = np.argsort(gene_var)[-topn:][::-1]  # Top variance genes, descending
        selected_genes = adata.var.index[top_var_genes]

        # For topn_genes parameter stability (when SKIP_DA_COMPUTATION is set),
        # apply gene selection to DE without sample variance
        # Otherwise, only apply to DE with sample variance (old behavior)
        if skip_da:
            genes_for_no_sample_de = selected_genes
            logger.info(f"Selected top {topn} most variable genes for DE without sample variance (topn_genes parameter stability)")
        else:
            genes_for_sample_de = selected_genes
            logger.info(f"Selected top {topn} most variable genes for DE with sample variance (standard mode)")
    else:
        logger.info("Using all genes for differential expression")

    # DE without sample variance
    logger.info("Computing differential expression without sample variance...")
    step_start = time.time()
    try:
        # Only pass genes parameter for topn_genes parameter stability testing
        de_kwargs = {}
        if genes_for_no_sample_de is not None:
            de_kwargs['genes'] = genes_for_no_sample_de
            # Use inplace=False to prevent modifying adata layers, which avoids
            # the expensive CSR→LIL→CSR sparse matrix conversion (6000+ seconds for 422K cells)
            # when genes parameter triggers "update existing layers" logic
            de_kwargs['inplace'] = False
            de_kwargs['return_full_results'] = True
            logger.info(f"Using inplace=False to avoid sparse matrix update overhead for gene subset")

        de_results_no_sample = kompot.compute_differential_expression(
            adata,
            groupby=config.GROUPING_COLUMN,
            condition1=config.CONDITIONS[0],
            condition2=config.CONDITIONS[1],
            layer=config.LAYER_FOR_EXPRESSION,
            obsm_key=config.DIMENSIONALITY_REDUCTION,
            **de_kwargs,
            **get_kompot_params(config, exclude_for='de')
        )

        # If we used inplace=False, manually add results to adata
        if genes_for_no_sample_de is not None and isinstance(de_results_no_sample, dict):
            # Store var results
            for key, values in de_results_no_sample.items():
                if key.startswith('var_'):
                    col_name = key[4:]  # Remove 'var_' prefix
                    adata.var[col_name] = values
            logger.info(f"Manually stored DE results in adata.var (bypassed sparse layer update)")
        timing_info['de_no_sample_seconds'] = time.time() - step_start
        timing_info['de_no_sample_success'] = True
        logger.info(f"DE without sample variance completed in {timing_info['de_no_sample_seconds']:.2f}s")
    except Exception as e:
        timing_info['de_no_sample_seconds'] = time.time() - step_start
        timing_info['de_no_sample_success'] = False
        error_msg = f"DE without sample variance failed: {str(e)}"
        timing_info['errors'].append(error_msg)
        logger.error(error_msg, exc_info=True)
    
    # DE with sample variance (optional, use selected genes if specified, with null_genes=0 for speed)
    # Also skip if SKIP_DA_COMPUTATION is set (e.g., for topn_genes parameter stability)
    if args.skip_sample_variance or skip_da:
        if skip_da:
            logger.info("Skipping DE with sample variance (SKIP_DA_COMPUTATION=true) - focus on DE without sample variance")
        else:
            logger.info("Skipping DE with sample variance (--skip-sample-variance flag set)")
        timing_info['de_with_sample_seconds'] = 0
        timing_info['de_with_sample_success'] = None
        timing_info['de_with_sample_skipped'] = True
    else:
        logger.info("Computing differential expression with sample variance...")
        step_start = time.time()
        try:
            de_results_with_sample = kompot.compute_differential_expression(
                adata,
                groupby=config.GROUPING_COLUMN,
                condition1=config.CONDITIONS[0],
                condition2=config.CONDITIONS[1],
                layer=config.LAYER_FOR_EXPRESSION,
                obsm_key=config.DIMENSIONALITY_REDUCTION,
                sample_col=config.SAMPLE_COLUMN,
                genes=genes_for_sample_de,
                null_genes=0,  # Skip null distribution when using sample variance on subset
                **get_kompot_params(config, exclude_for='de')
            )
            timing_info['de_with_sample_seconds'] = time.time() - step_start
            timing_info['de_with_sample_success'] = True
            logger.info(f"DE with sample variance completed in {timing_info['de_with_sample_seconds']:.2f}s")
        except Exception as e:
            timing_info['de_with_sample_seconds'] = time.time() - step_start
            timing_info['de_with_sample_success'] = False
            error_msg = f"DE with sample variance failed: {str(e)}"
            timing_info['errors'].append(error_msg)
            logger.error(error_msg, exc_info=True)
    
    timing_info['total_processing_seconds'] = time.time() - start_time
    timing_info['genes_processed'] = topn
    
    # Save runinfo with timing
    runinfo = {
        'input_file': str(input_path.absolute()),
        'output_file': str(output_file.absolute()),
        'dataset': args.dataset,
        'config': config_dict,
        'config_overrides': dict(override.split('=', 1) for override in args.config_override),
        'parameter_variation': parameter_variation,
        'adata_info': adata_info,
        'environment': env_info,
        'slurm': slurm_info,
        'timing': timing_info,
        'log_file': str(log_file.absolute())
    }
    
    runinfo_file = str(output_file) + ".runinfo.yml"
    with open(runinfo_file, 'w') as f:
        yaml.dump(runinfo, f, default_flow_style=False)
    logger.info(f"Saved run info to {runinfo_file}")
    
    # Save results (CSV files only - no AnnData to save disk space)
    logger.info(f"Saving results to {output_file}")
    adata.obs.to_csv(str(output_file) + ".obs.csv")
    adata.var.to_csv(str(output_file) + ".var.csv")
    # adata.write(output_file)  # DISABLED: AnnData objects consume massive disk space (~54TB for all results)
    
    logger.info(f"Processing completed successfully in {timing_info['total_processing_seconds']:.2f}s")


if __name__ == "__main__":
    main()