src/compute_descriptors.py

#!/usr/bin/env python3

"""
Descriptor computation module for AbMelt inference pipeline.
Handles extraction of MD descriptors from trajectories and aggregation into ML-ready format.
"""

import os
import sys
import logging
from pathlib import Path
from typing import Dict, List, Optional, Tuple
import numpy as np
import pandas as pd
import glob

try:
    import gromacs
    from order_param import order_s2, avg_s2_blocks, get_lambda, order_lambda
    from res_sasa import core_surface, get_core_surface, get_slope
    from timing import time_step
except ImportError as e:
    logging.error(f"Failed to import required modules: {e}")
    raise

logger = logging.getLogger(__name__)


def compute_descriptors(simulation_result: Dict, config: Dict) -> Dict:
    """
    Main entry point for descriptor computation.
    
    Args:
        simulation_result: Dictionary containing trajectory files and work directory
        config: Configuration dictionary
        
    Returns:
        Dictionary containing descriptors DataFrame and metadata
    """
    logger.info("Starting descriptor computation...")
    
    work_dir = Path(simulation_result["work_dir"])
    trajectory_files = simulation_result["trajectory_files"]
    antibody_name = work_dir.name
    
    # Get descriptor computation parameters
    desc_config = config["descriptors"]
    eq_time = desc_config["equilibration_time"]
    block_lengths = desc_config["block_length"]  # Now expects a list [2.5, 25]
    if not isinstance(block_lengths, list):
        block_lengths = [block_lengths]  # Backwards compatibility
    core_surface_k = desc_config["core_surface_k"]
    compute_lambda = desc_config["compute_lambda"]
    use_dummy_s2 = desc_config.get("use_dummy_s2", False)  # Default to False if not specified
    
    # Extract temperatures from trajectory files
    temps = [str(temp) for temp in trajectory_files.keys()]
    
    # Store original working directory
    original_cwd = os.getcwd()
    
    try:
        # Change to work directory
        os.chdir(work_dir)
        logger.info(f"Changed to work directory: {work_dir}")
        
        # Step 1: Compute GROMACS-based descriptors
        with time_step("GROMACS Descriptors", parent="Descriptor Computation"):
            logger.info("Step 1: Computing GROMACS descriptors...")
            xvg_files = _compute_gromacs_descriptors(work_dir, temps, eq_time)
            logger.info(f"Generated {len(xvg_files)} GROMACS descriptor files")
        
        # Step 2: Compute order parameters
        with time_step("Order Parameters", parent="Descriptor Computation"):
            logger.info("Step 2: Computing order parameters...")
            master_s2_dicts = _compute_order_parameters(work_dir, temps, eq_time, block_lengths, antibody_name, use_dummy_s2)
        
        # Step 3: Compute core/surface SASA
        with time_step("Core/Surface SASA", parent="Descriptor Computation"):
            logger.info("Step 3: Computing core/surface SASA...")
            sasa_dict = _compute_core_surface_sasa(work_dir, temps, eq_time, core_surface_k)
        
        # Step 4: Compute multi-temperature features (lambda)
        with time_step("Lambda Features", parent="Descriptor Computation"):
            if len(temps) >= 2 and compute_lambda:
                logger.info("Step 4: Computing multi-temperature lambda...")
                all_lambda_features = _compute_lambda_features(master_s2_dicts, temps, eq_time, antibody_name)
            else:
                logger.warning(f"Skipping lambda computation: need >=2 temperatures, got {len(temps)}")
                all_lambda_features = None
        
        # Step 5: Aggregate all descriptors into DataFrame
        with time_step("Aggregate to DataFrame", parent="Descriptor Computation"):
            logger.info("Step 5: Aggregating descriptors to DataFrame...")
            descriptors_df = _aggregate_descriptors_to_dataframe(
                work_dir, temps, antibody_name, eq_time, master_s2_dicts, 
                all_lambda_features, sasa_dict, core_surface_k
            )
        
        logger.info(f"Descriptor computation completed. DataFrame shape: {descriptors_df.shape}")
        logger.info(f"Features: {list(descriptors_df.columns)}")
        
        # Save descriptors to file for future use
        try:
            descriptors_csv = work_dir / "descriptors.csv"
            descriptors_pkl = work_dir / "descriptors.pkl"
            
            descriptors_df.to_csv(descriptors_csv, index=False)
            logger.info(f"Saved descriptors to {descriptors_csv}")
            
            import pickle
            with open(descriptors_pkl, 'wb') as f:
                pickle.dump(descriptors_df, f)
            logger.info(f"Saved descriptors to {descriptors_pkl}")
        except Exception as e:
            logger.warning(f"Failed to save descriptors to file: {e}")
            logger.warning("Continuing without saving descriptors")
        
        result = {
            "status": "success",
            "descriptors_df": descriptors_df,
            "xvg_files": xvg_files,
            "work_dir": str(work_dir),
            "message": "Descriptor computation completed successfully"
        }
        
        return result
        
    except Exception as e:
        logger.error(f"Descriptor computation failed: {e}")
        raise
    finally:
        # Restore original working directory
        os.chdir(original_cwd)


def _compute_gromacs_descriptors(work_dir: Path, temps: List[str], eq_time: int) -> List[str]:
    """
    Compute GROMACS-based descriptors from trajectories.
    
    Args:
        work_dir: Working directory containing trajectories
        temps: List of temperature strings
        eq_time: Equilibration time in ns
        
    Returns:
        List of generated .xvg file paths
    """
    xvg_files = []
    
    for temp in temps:
        logger.info(f"Computing GROMACS descriptors for {temp}K...")
        
        # Check if required files exist
        final_xtc = f'md_final_{temp}.xtc'
        final_gro = f'md_final_{temp}.gro'
        tpr_file = f'md_{temp}.tpr'
        index_file = 'index.ndx'
        
        if not os.path.exists(final_xtc):
            raise ValueError(f'Trajectory file not found: {final_xtc}')
        if not os.path.exists(final_gro):
            raise ValueError(f'Structure file not found: {final_gro}')
        if not os.path.exists(tpr_file):
            raise ValueError(f'TPR file not found: {tpr_file}')
        if not os.path.exists(index_file):
            logger.warning(f'Index file not found: {index_file}. CDR-specific features may fail.')
            # Try to create index file if it doesn't exist
            # This should have been created during MD simulation, but handle gracefully
            logger.warning('Attempting to create index file...')
            try:
                from preprocess import canonical_index
                annotation = canonical_index(pdb='processed.pdb')
                gromacs.make_ndx(f='processed.gro', o='index.ndx', input=annotation)
                logger.info('Index file created successfully')
            except Exception as e:
                logger.error(f'Failed to create index file: {e}')
                logger.error('CDR-specific features will be skipped')
        
        try:
            # Global features
            logger.debug(f"  Computing global SASA...")
            gromacs.sasa(f=final_xtc, s=final_gro, o=f'sasa_{temp}.xvg', input=['1'])
            xvg_files.append(f'sasa_{temp}.xvg')
            
            logger.debug(f"  Computing hydrogen bonds and contacts...")
            # Use legacy hbond to get both hydrogen bonds AND contacts in one file
            gromacs.hbond_legacy(f=final_xtc, s=tpr_file, num=f'bonds_{temp}.xvg', input=['1', '1'])
            xvg_files.append(f'bonds_{temp}.xvg')
            
            logger.debug(f"  Computing RMSD...")
            gromacs.rms(f=final_xtc, s=final_gro, o=f'rmsd_{temp}.xvg', input=['3', '3'])
            xvg_files.append(f'rmsd_{temp}.xvg')
            
            logger.debug(f"  Computing gyration radius...")
            gromacs.gyrate(f=final_xtc, s=final_gro, o=f'gyr_{temp}.xvg', n=index_file, input=['1'])
            xvg_files.append(f'gyr_{temp}.xvg')
            
            # CDR-specific features
            cdr_regions = {
                'cdrl1': '12',
                'cdrl2': '13',
                'cdrl3': '14',
                'cdrh1': '15',
                'cdrh2': '16',
                'cdrh3': '17',
                'cdrs': '18'
            }
            
            # SASA for each CDR
            logger.debug(f"  Computing CDR SASA...")
            for cdr_name, index_group in cdr_regions.items():
                gromacs.sasa(f=final_xtc, s=final_gro, o=f'sasa_{cdr_name}_{temp}.xvg', 
                           n=index_file, input=[index_group])
                xvg_files.append(f'sasa_{cdr_name}_{temp}.xvg')
            
            # H-bonds between light and heavy chains
            logger.debug(f"  Computing light-heavy bonds...")
            gromacs.hbond(f=final_xtc, s=tpr_file, num=f'bonds_lh_{temp}.xvg', 
                         n=index_file, input=['10', '11'])
            xvg_files.append(f'bonds_lh_{temp}.xvg')
            
            # RMSF for each CDR
            logger.debug(f"  Computing CDR RMSF...")
            eq_time_ps = str(eq_time * 1000)
            for cdr_name, index_group in cdr_regions.items():
                gromacs.rmsf(f=final_xtc, s=final_gro, o=f'rmsf_{cdr_name}_{temp}.xvg',
                           n=index_file, b=eq_time_ps, input=[index_group, index_group])
                xvg_files.append(f'rmsf_{cdr_name}_{temp}.xvg')
            
            # Gyration radius for each CDR
            logger.debug(f"  Computing CDR gyration...")
            for cdr_name, index_group in cdr_regions.items():
                gromacs.gyrate(f=final_xtc, s=final_gro, o=f'gyr_{cdr_name}_{temp}.xvg',
                             n=index_file, input=[index_group])
                xvg_files.append(f'gyr_{cdr_name}_{temp}.xvg')
            
            # Conformational entropy (S_conf)
            logger.debug(f"  Computing conformational entropy...")
            try:
                gromacs.trjconv(f=final_xtc, s=tpr_file, dt='0', fit='rot+trans', 
                             n=index_file, o=f'md_final_covar_{temp}.xtc', input=['1','1'])
                gromacs.covar(f=f'md_final_covar_{temp}.xtc', s=tpr_file, n=index_file,
                            o=f'covar_{temp}.xvg', av=f'avg_covar{temp}.pdb',
                            ascii=f'covar_matrix_{temp}.dat', v=f'covar_{temp}.trr',
                            input=['4', '4'])
                # Note: anaeig output goes to log file via shell redirection
                # The original implementation uses shell redirection in input parameter
                gromacs.anaeig(f=f'md_final_covar_{temp}.xtc', v=f'covar_{temp}.trr',
                             entropy=True, temp=temp, s=tpr_file, nevskip='6',
                             n=index_file, b=eq_time_ps, input=[f'> sconf_{temp}.log'])
            except Exception as e:
                logger.warning(f"Conformational entropy computation failed for {temp}K: {e}")
            
            # Electrostatic potential for each CDR
            logger.debug(f"  Computing CDR electrostatic potential...")
            for cdr_name, index_group in cdr_regions.items():
                try:
                    gromacs.potential(f=final_xtc, s=tpr_file, spherical=True, sl='10',
                                   o=f'potential_{cdr_name}_{temp}.xvg',
                                   oc=f'charge_{cdr_name}_{temp}.xvg',
                                   of=f'field_{cdr_name}_{temp}.xvg',
                                   n=index_file, input=[index_group])
                    xvg_files.append(f'potential_{cdr_name}_{temp}.xvg')
                except Exception as e:
                    logger.warning(f"Potential computation failed for {cdr_name} at {temp}K: {e}")
            
            # Dipole moment
            logger.debug(f"  Computing dipole moment...")
            gromacs.dipoles(f=final_xtc, s=tpr_file, o=f'dipole_{temp}.xvg',
                          n=index_file, input=['1'])
            xvg_files.append(f'dipole_{temp}.xvg')
            
        except Exception as e:
            logger.error(f"Failed to compute GROMACS descriptors for {temp}K: {e}")
            raise
    
    return xvg_files


def _compute_order_parameters(work_dir: Path, temps: List[str], eq_time: int, 
                             block_lengths: List[float], antibody_name: str, use_dummy_s2: bool = False) -> Dict[float, Dict[int, Dict]]:
    """
    Compute N-H bond order parameters (S²) for each temperature and block length.
    
    Args:
        work_dir: Working directory
        temps: List of temperature strings
        eq_time: Equilibration time in ns
        block_lengths: List of block lengths for order parameter calculation in ns
        antibody_name: Name of antibody
        use_dummy_s2: If True, generate dummy S2 values instead of computing from trajectory
        
    Returns:
        Dictionary mapping block_length (float) to dictionary mapping temperature (int) to S² values per residue
    """
    all_master_s2_dicts = {}
    
    if use_dummy_s2:
        logger.info("Using dummy S2 values for testing (use_dummy_s2=True)")
    
    for block_length in block_lengths:
        logger.info(f"Computing order parameters for block_length={block_length}ns...")
        master_s2_dict = {int(temp): {} for temp in temps}
        
        for temp in temps:
            logger.info(f"  Computing S² for {temp}K, block={block_length}ns...")
            
            final_xtc = f'md_final_{temp}.xtc'
            final_gro = f'md_final_{temp}.gro'
            
            if not os.path.exists(final_xtc) or not os.path.exists(final_gro):
                logger.warning(f"Trajectory files not found for {temp}K, skipping")
                continue
            
            try:
                s2_blocks_dict = order_s2(mab=antibody_name, temp=temp, 
                                        block_length=block_length, start=eq_time, use_dummy=use_dummy_s2)
                master_s2_dict[int(temp)] = avg_s2_blocks(s2_blocks_dict)
                logger.info(f"  Order parameters computed for {temp}K")
            except Exception as e:
                logger.warning(f"Order parameter computation failed for {temp}K: {e}")
                logger.warning("This is common with short trajectories. Continuing...")
        
        all_master_s2_dicts[block_length] = master_s2_dict
    
    return all_master_s2_dicts


def _compute_core_surface_sasa(work_dir: Path, temps: List[str], eq_time: int, k: int) -> Dict:
    """
    Compute core/surface SASA using mdtraj.
    
    Args:
        work_dir: Working directory
        temps: List of temperature strings
        eq_time: Equilibration time in ns
        k: Number of residues for core/surface classification
        
    Returns:
        Dictionary with SASA statistics per temperature
    """
    sasa_dict = {}
    
    for temp in temps:
        logger.info(f"Computing core/surface SASA for {temp}K...")
        
        final_xtc = f'md_final_{temp}.xtc'
        final_gro = f'md_final_{temp}.gro'
        
        if not os.path.exists(final_xtc) or not os.path.exists(final_gro):
            logger.warning(f"Trajectory files not found for {temp}K, skipping SASA")
            continue
        
        try:
            # Compute residue-level SASA
            core_surface(temp)
            
            # Aggregate statistics
            sasa_dict[temp] = {}
            sasa_dict = get_core_surface(sasa_dict, temp, k=k, start=eq_time)
            logger.info(f"Core/surface SASA computed for {temp}K")
        except Exception as e:
            logger.warning(f"Core/surface SASA computation failed for {temp}K: {e}")
    
    return sasa_dict


def _compute_lambda_features(master_s2_dicts: Dict[float, Dict[int, Dict]], temps: List[str],
                            eq_time: int, antibody_name: str) -> Dict[float, Tuple[Dict, Dict]]:
    """
    Compute multi-temperature lambda (order parameter slope) for each block length.
    
    Args:
        master_s2_dicts: Dictionary mapping block_length to dictionary of S² values per temperature
        temps: List of temperature strings
        eq_time: Equilibration time
        antibody_name: Name of antibody
        
    Returns:
        Dictionary mapping block_length to tuple of (lambda_dict, r_dict) - lambda values and correlation coefficients per residue
    """
    temp_ints = [int(t) for t in temps]
    all_lambda_features = {}

    for block_length, master_s2_dict in master_s2_dicts.items():
        logger.info(f"Computing lambda features for block_length={block_length}ns...")
        
        # Filter out temperatures that don't have S² data
        available_temps = [t for t in temp_ints if t in master_s2_dict and len(master_s2_dict[t]) > 0]
        
        if len(available_temps) < 2:
            logger.warning(f"Need at least 2 temperatures with S² data for lambda, got {len(available_temps)}")
            all_lambda_features[block_length] = (None, None)
            continue
        
        try:
            # Use order_lambda function from order_param module (saves CSV)
            # Note: start parameter expects picoseconds
            order_lambda(master_dict=master_s2_dict, mab=antibody_name, 
                        temps=available_temps, block_length=str(block_length), 
                        start=str(eq_time * 1000))
            
            # Compute lambda and r for each residue directly
            lambda_dict, r_dict = get_lambda(master_s2_dict, temps=available_temps)
            
            logger.info(f"Lambda computed for {len(lambda_dict)} residues at block={block_length}ns")
            all_lambda_features[block_length] = (lambda_dict, r_dict)
            
        except Exception as e:
            logger.warning(f"Lambda computation failed for block={block_length}ns: {e}")
            all_lambda_features[block_length] = (None, None)

    return all_lambda_features


def _aggregate_descriptors_to_dataframe(work_dir: Path, temps: List[str], 
                                       antibody_name: str, eq_time: int,
                                       master_s2_dicts: Dict[float, Dict[int, Dict]], 
                                       all_lambda_features: Dict[float, Tuple[Dict, Dict]],
                                       sasa_dict: Dict, core_surface_k: int) -> pd.DataFrame:
    """
    Aggregate all computed descriptors into a single-row DataFrame.
    
    Args:
        work_dir: Working directory
        temps: List of temperature strings
        antibody_name: Name of antibody
        eq_time: Equilibration time in ns
        master_s2_dicts: Dictionary mapping block_length to order parameter dictionary per temperature
        all_lambda_features: Dictionary mapping block_length to tuple of (lambda_dict, r_dict)
        sasa_dict: Core/surface SASA dictionary
        core_surface_k: Number of residues for core/surface classification
        
    Returns:
        Single-row DataFrame with all descriptors
    """
    descriptor_dict = {}
    
    # Parse all .xvg files
    xvg_files = glob.glob('*.xvg')
    
    for xvg_file in xvg_files:
        try:
            metric_name = Path(xvg_file).stem
            
            # Extract temperature from filename
            temp = None
            for t in temps:
                if t in metric_name:
                    temp = t
                    break
            
            if temp is None:
                continue
            
            # Parse the xvg file
            data = _parse_xvg_file(xvg_file)
            
            if data is None or len(data) == 0:
                continue
            
            # Compute equilibrated statistics
            # Note: RMSF files contain per-residue data (not time-series), 
            # and equilibration is already handled by the -b flag in GROMACS
            if 'rmsf' in metric_name:
                # RMSF: per-residue data, no time-based equilibration needed
                equilibrated_data = data
            else:
                # Time-series data: apply equilibration slicing
                eq_time_ps = eq_time * 1000  # Convert to ps
                eq_start_idx = int(eq_time_ps / 10)  # Assuming 10 ps per frame (adjust if needed)
                if len(data) <= eq_start_idx:
                    continue
                equilibrated_data = data[eq_start_idx:]
            
            if len(equilibrated_data) > 0:
                
                # Handle different data shapes
                if equilibrated_data.ndim == 1:
                    # Single column data
                    mu = np.mean(equilibrated_data)
                    std = np.std(equilibrated_data)
                    
                    # Create feature names based on metric type
                    # Match exact naming conventions from training data
                    if 'bonds' in metric_name:
                        if 'lh' in metric_name:
                            descriptor_dict[f'bonds_lh_mu_{temp}'] = mu
                            descriptor_dict[f'bonds_lh_std_{temp}'] = std
                        else:
                            # bonds file has hbonds and contacts - handled in 2D case
                            descriptor_dict[f'bonds_hbonds_mu_{temp}'] = mu
                            descriptor_dict[f'bonds_hbonds_std_{temp}'] = std
                    elif 'sasa' in metric_name:
                        region = metric_name.replace('sasa_', '').replace(f'_{temp}', '')
                        descriptor_dict[f'sasa_{region}_mu_{temp}'] = mu
                        descriptor_dict[f'sasa_{region}_std_{temp}'] = std
                    elif 'rmsd' in metric_name:
                        descriptor_dict[f'rmsd_mu_{temp}'] = mu
                        descriptor_dict[f'rmsd_std_{temp}'] = std
                    elif 'rmsf' in metric_name:
                        region = metric_name.replace('rmsf_', '').replace(f'_{temp}', '')
                        # Some models use mu, some use std - include both
                        descriptor_dict[f'rmsf_{region}_mu_{temp}'] = mu
                        descriptor_dict[f'rmsf_{region}_std_{temp}'] = std
                    elif 'gyr' in metric_name:
                        region = metric_name.replace('gyr_', '').replace(f'_{temp}', '')
                        # Training data shows: gyr_cdrs_Rg_std_350, gyr_cdrs_Rg_std_400
                        descriptor_dict[f'gyr_{region}_Rg_mu_{temp}'] = mu
                        descriptor_dict[f'gyr_{region}_Rg_std_{temp}'] = std
                    elif 'potential' in metric_name:
                        # Potential features should NOT use time-series mean
                        # They need specific radius index extraction (handled in 2D case below)
                        # This 1D case should not happen for potential files
                        pass
                    elif 'dipole' in metric_name:
                        # Dipole files should have 4 columns, handled in 2D case
                        # This 1D case should not happen for dipole files
                        # But if it does, use the mean
                        descriptor_dict[f'dipole_mu_{temp}'] = mu
                        descriptor_dict[f'dipole_std_{temp}'] = std
                
                elif equilibrated_data.ndim == 2:
                    # Multi-column data (e.g., gyration with Rg, Rx, Ry, Rz, potential with multiple radii)
                    
                    # Handle potential files (multiple radii/slices)
                    if 'potential' in metric_name:
                        region = metric_name.replace('potential_', '').replace(f'_{temp}', '')
                        # Potential is measured at specific radius indices
                        # Original code uses radius index based on region type
                        if region in ['cdrl1', 'cdrl2', 'cdrl3', 'cdrh1', 'cdrh2', 'cdrh3']:
                            # Individual CDRs use radius index 2
                            radius_idx = 2
                        elif region == 'cdrs':
                            # Combined CDRs use radius index 5
                            radius_idx = 5
                        else:
                            # Default to radius index 2
                            radius_idx = 2
                        
                        # Extract value at specific radius (column 1 is potential value, column 0 is radius)
                        if equilibrated_data.shape[0] > radius_idx:
                            # Potential files have 2 columns: [radius, potential]
                            # We need the potential value (column 1) at the specific radius index (row)
                            descriptor_dict[f'potential_{region}_mu_{temp}'] = equilibrated_data[radius_idx, 1]
                            # Original code sets std to 0 for potential
                            descriptor_dict[f'potential_{region}_std_{temp}'] = 0
                        else:
                            logger.warning(f"Not enough radii in potential file for {region} (need idx {radius_idx}, have {equilibrated_data.shape[0]} rows)")
                    
                    elif equilibrated_data.shape[1] >= 4:
                        # Gyration radius components
                        if 'gyr' in metric_name:
                            region = metric_name.replace('gyr_', '').replace(f'_{temp}', '')
                            r_values = equilibrated_data[:, 0]  # Rg
                            x_values = equilibrated_data[:, 1]  # Rx
                            y_values = equilibrated_data[:, 2]  # Ry
                            z_values = equilibrated_data[:, 3]  # Rz
                            
                            # Match training data format: gyr_cdrs_Rg_std_350
                            descriptor_dict[f'gyr_{region}_Rg_mu_{temp}'] = np.mean(r_values)
                            descriptor_dict[f'gyr_{region}_Rg_std_{temp}'] = np.std(r_values)
                            descriptor_dict[f'gyr_{region}_Rx_mu_{temp}'] = np.mean(x_values)
                            descriptor_dict[f'gyr_{region}_Rx_std_{temp}'] = np.std(x_values)
                            descriptor_dict[f'gyr_{region}_Ry_mu_{temp}'] = np.mean(y_values)
                            descriptor_dict[f'gyr_{region}_Ry_std_{temp}'] = np.std(y_values)
                            descriptor_dict[f'gyr_{region}_Rz_mu_{temp}'] = np.mean(z_values)
                            descriptor_dict[f'gyr_{region}_Rz_std_{temp}'] = np.std(z_values)
                    
                    elif equilibrated_data.shape[1] == 2:
                        # Two-column data (e.g., bonds with hbonds and contacts)
                        if 'bonds' in metric_name:
                            hbonds = equilibrated_data[:, 0]
                            contacts = equilibrated_data[:, 1]
                            
                            # Match training data format: bonds_contacts_std_350
                            descriptor_dict[f'bonds_hbonds_mu_{temp}'] = np.mean(hbonds)
                            descriptor_dict[f'bonds_hbonds_std_{temp}'] = np.std(hbonds)
                            descriptor_dict[f'bonds_contacts_mu_{temp}'] = np.mean(contacts)
                            descriptor_dict[f'bonds_contacts_std_{temp}'] = np.std(contacts)
                    
                    elif equilibrated_data.shape[1] == 3:
                        # Three-column data (e.g., dipole with Mx, My, Mz)
                        if 'dipole' in metric_name:
                            # Dipole files have columns: Mx, My, Mz(magnitude), [|Mtot|]
                            # Original code uses Z (Mz, the magnitude) = column index 2
                            dipole_z = equilibrated_data[:, 2]  # Mz column
                            descriptor_dict[f'dipole_mu_{temp}'] = np.mean(dipole_z)
                            descriptor_dict[f'dipole_std_{temp}'] = np.std(dipole_z)
                    
                    elif equilibrated_data.shape[1] == 4:
                        # Four-column data (e.g., dipole with Mx, My, Mz, |Mtot|)
                        if 'dipole' in metric_name:
                            # Dipole files have 4 columns: Mx, My, Mz(magnitude), |Mtot|
                            # Original code uses Z (Mz) = column index 2
                            dipole_z = equilibrated_data[:, 2]  # Mz column
                            descriptor_dict[f'dipole_mu_{temp}'] = np.mean(dipole_z)
                            descriptor_dict[f'dipole_std_{temp}'] = np.std(dipole_z)
        
        except Exception as e:
            logger.warning(f"Failed to parse {xvg_file}: {e}")
            continue
    
    # Add order parameter features for each block length
    for block_length, master_s2_dict in master_s2_dicts.items():
        for temp_int, s2_values in master_s2_dict.items():
            if s2_values and len(s2_values) > 0:
                temp_str = str(temp_int)
                s2_mean = np.mean(list(s2_values.values()))
                s2_std = np.std(list(s2_values.values()))
                # Include block length in feature name for clarity (optional)
                descriptor_dict[f'order_s2_{temp_str}_b={block_length}_mu'] = s2_mean
                descriptor_dict[f'order_s2_{temp_str}_b={block_length}_std'] = s2_std
    
    # Add lambda features for each block length
    if all_lambda_features:
        for block_length, (lambda_dict, r_dict) in all_lambda_features.items():
            if lambda_dict and r_dict:
                lambda_mean = np.mean(list(lambda_dict.values()))
                r_mean = np.mean(list(r_dict.values()))
                
                # Generate features with correct values
                descriptor_dict[f'all-temp_lamda_b={block_length}_eq={eq_time}'] = lambda_mean
                descriptor_dict[f'r-lamda_b={block_length}_eq={eq_time}'] = r_mean  # FIX: was lambda_mean
                descriptor_dict[f'all-temp_lamda_r_b={block_length}_eq={eq_time}'] = r_mean
    
    # Add core/surface SASA features
    if sasa_dict:
        # Per-temperature SASA features
        for temp, sasa_data in sasa_dict.items():
            if isinstance(sasa_data, dict):
                for key, value in sasa_data.items():
                    descriptor_dict[f'sasa_{key}_{temp}'] = value
        
        # Cross-temperature SASA slopes
        if len(temps) >= 2:
            temp_ints = sorted([int(t) for t in temps])
            sasa_slopes = {}
            
            for key in ['total_mean', 'core_mean', 'surface_mean', 'total_std', 'core_std', 'surface_std']:
                data_points = [(int(t), sasa_dict[t][key]) for t in temps if t in sasa_dict and key in sasa_dict[t]]
                if len(data_points) >= 2:
                    slope = get_slope(data_points)
                    sasa_slopes[key] = slope
            
            for key, slope in sasa_slopes.items():
                descriptor_dict[f'all-temp-sasa_{key}_k={core_surface_k}_eq={eq_time}'] = slope
    
    # Parse conformational entropy from log files
    for temp in temps:
        log_file = f'sconf_{temp}.log'
        if os.path.exists(log_file):
            try:
                with open(log_file, 'r') as f:
                    for line in f:
                        if 'Entropy' in line and 'J/mol K' in line:
                            if 'Schlitter' in line:
                                parts = line.split()
                                if len(parts) > 8:
                                    entropy = float(parts[8])
                                    descriptor_dict[f'sconf_schlitter_{temp}'] = entropy
                            elif 'Quasiharmonic' in line:
                                parts = line.split()
                                if len(parts) > 8:
                                    entropy = float(parts[8])
                                    descriptor_dict[f'sconf_quasiharmonic_{temp}'] = entropy
            except Exception as e:
                logger.warning(f"Failed to parse entropy log {log_file}: {e}")
    
    # Create DataFrame
    df = pd.DataFrame([descriptor_dict])
    
    return df


def load_existing_descriptors(simulation_result: Dict, config: Dict) -> Dict:
    """
    Load existing descriptor computation results.
    
    Args:
        simulation_result: Dictionary containing simulation results
        config: Configuration dictionary
        
    Returns:
        Dictionary matching format from compute_descriptors
        
    Raises:
        FileNotFoundError: If descriptor file not found
    """
    logger.info("Loading existing descriptor computation results...")
    
    work_dir = Path(simulation_result["work_dir"]).resolve()
    
    # Try CSV first, then pickle
    descriptors_csv = work_dir / "descriptors.csv"
    descriptors_pkl = work_dir / "descriptors.pkl"
    
    descriptors_df = None
    
    if descriptors_csv.exists():
        try:
            descriptors_df = pd.read_csv(descriptors_csv)
            logger.info(f"Loaded descriptors from {descriptors_csv}")
        except Exception as e:
            logger.warning(f"Failed to load descriptors from CSV: {e}")
    
    if descriptors_df is None and descriptors_pkl.exists():
        try:
            import pickle
            with open(descriptors_pkl, 'rb') as f:
                descriptors_df = pickle.load(f)
            logger.info(f"Loaded descriptors from {descriptors_pkl}")
        except Exception as e:
            logger.warning(f"Failed to load descriptors from pickle: {e}")
    
    if descriptors_df is None:
        error_msg = f"Descriptor file not found when skipping descriptor computation.\n"
        error_msg += f"Expected one of:\n"
        error_msg += f"  - {descriptors_csv}\n"
        error_msg += f"  - {descriptors_pkl}\n"
        error_msg += f"\nWork directory: {work_dir}"
        raise FileNotFoundError(error_msg)
    
    # Get list of XVG files in work directory (if they exist)
    xvg_files = []
    try:
        xvg_files = [str(f.name) for f in work_dir.glob("*.xvg")]
        logger.info(f"Found {len(xvg_files)} XVG files in work directory")
    except Exception as e:
        logger.warning(f"Could not enumerate XVG files: {e}")
    
    result = {
        "status": "success",
        "descriptors_df": descriptors_df,
        "xvg_files": xvg_files,
        "work_dir": str(work_dir),
        "message": "Descriptor computation results loaded successfully"
    }
    
    logger.info(f"Successfully loaded descriptors. DataFrame shape: {descriptors_df.shape}")
    logger.info(f"Features: {len(descriptors_df.columns)}")
    
    return result


def _parse_xvg_file(xvg_file: str) -> Optional[np.ndarray]:
    """
    Parse GROMACS .xvg file and return data as numpy array.
    
    Args:
        xvg_file: Path to .xvg file
        
    Returns:
        Numpy array with data (time in first column, data in subsequent columns)
    """
    try:
        t, x, y, z, r = [], [], [], [], []
        
        with open(xvg_file, 'r') as f:
            for line in f:
                # Skip comments and metadata
                if line.startswith('#') or line.startswith('@'):
                    continue
                
                cols = line.split()
                
                if len(cols) == 0:
                    continue
                elif len(cols) == 2:
                    t.append(float(cols[0]))
                    x.append(float(cols[1]))
                elif len(cols) == 3:
                    t.append(float(cols[0]))
                    x.append(float(cols[1]))
                    y.append(float(cols[2]))
                elif len(cols) == 4:
                    t.append(float(cols[0]))
                    x.append(float(cols[1]))
                    y.append(float(cols[2]))
                    z.append(float(cols[3]))
                elif len(cols) == 5:
                    t.append(float(cols[0]))
                    r.append(float(cols[1]))
                    x.append(float(cols[2]))
                    y.append(float(cols[3]))
                    z.append(float(cols[4]))
        
        # Return appropriate array based on what was collected
        if len(r) > 0:
            return np.column_stack([r, x, y, z])
        elif len(z) > 0:
            return np.column_stack([x, y, z])
        elif len(y) > 0:
            return np.column_stack([x, y])
        elif len(x) > 0:
            return np.array(x)
        else:
            return None
            
    except Exception as e:
        logger.warning(f"Failed to parse {xvg_file}: {e}")
        return None