import pathlib
import subprocess

def convert_cif_to_pdb(input_folder):
    """
    Convert all CIF files in the specified folder to PDB format using OpenBabel.
    The PDB files will be saved in the same folder.
    
    Args:
        input_folder (str): Path to the folder containing CIF files.
    """
    folder_path = pathlib.Path(input_folder)
    
    if not folder_path.exists():
        print(f"Error: Folder '{folder_path}' does not exist.")
        return
    
    # Find all .cif files in the folder
    cif_files = list(folder_path.glob("*.cif"))
    if not cif_files:
        print(f"No CIF files found in '{folder_path}'.")
        return
    
    print(f"Found {len(cif_files)} CIF file(s) to convert.")
    
    successes = []
    failures = []
    
    for cif_file in cif_files:
        pdb_file = cif_file.with_suffix(".pdb")
        
        # OpenBabel command
        cmd = [
            "obabel",
            "-icif", str(cif_file),
            "-opdb", "-O", str(pdb_file),
            "-e"  
        ]
        
        print(f"Converting '{cif_file.name}' to PDB...")
        
        try:
            # Run OpenBabel
            subprocess.run(cmd, check=True, capture_output=True, text=True)
            
            if pdb_file.exists() and pdb_file.stat().st_size > 0:
                successes.append(cif_file.name)
                print(f"Success: Saved as '{pdb_file.name}'")
            else:
                failures.append(cif_file.name)
                print("Warning: No valid PDB file was created.")
                
        except subprocess.CalledProcessError as e:
            failures.append(cif_file.name)
            print(f"Error: Failed to convert '{cif_file.name}': {e.stderr.strip()}")
        except FileNotFoundError:
            print("Error: 'obabel' command not found. Is OpenBabel installed?")
            return
        except Exception as e:
            failures.append(cif_file.name)
            print(f"Unexpected error converting '{cif_file.name}': {str(e)}")
    
    # Summary
    print("\nConversion Summary:")
    print(f"Successful: {len(successes)}")
    print(f"Failed: {len(failures)}")
    if failures:
        print("\nFailed files:")
        for name in failures:
            print(f"- {name}")


INPUT_FOLDER = r"C:\Users\Lenovo\pdb"  #path to the input files (just write the path to folder, can handle all cif files in batch)

if __name__ == "__main__":
    convert_cif_to_pdb(INPUT_FOLDER)
  

import json
from pathlib import Path

def ligandmpnn_to_json_filtered(input_file_path, output_file=None):
    """
    Parse LigandMPNN output file, OMIT the first sequence, and convert to JSON.
    Uses the existing 'id' from the file for labeling.
    """
    # Convert input path to Path object
    input_path = Path(input_file_path)
    
    if not input_path.exists():
        print(f"Error: The file '{input_file_path}' does not exist.")
        return
    
    # create a file or just use the older one
    if output_file is None:
        output_file = input_path.with_suffix('.json')
    
    try:
        with open(input_path, 'r', encoding='utf-8') as f:
            content = f.read()
        
        # Parse the file and extract sequences only with ids specified (to avoid adding on the wt protien and waste valuable quota)
        sequences_with_ids = []  # Store tuples of (sequence, id_from_file)
        current_record = None
        current_sequence = []
        current_id = None
        found_first_record = False
        
        for line in content.split('\n'):
            line = line.strip()
            if line.startswith('>P1;'):
                if current_record is not None and current_sequence and current_id is not None:
                    sequence_str = ''.join(current_sequence).replace(' ', '')
                    sequences_with_ids.append((sequence_str, current_id))
                
                # repeat 
                current_record = line
                current_sequence = []
                
                # Check if this header contains an ID
                if 'id=' in line:
                    # Extract the ID from the header
                    import re
                    id_match = re.search(r'id=(\d+)', line)
                    if id_match:
                        current_id = int(id_match.group(1))
                    else:
                        current_id = None
                else:
                    # defined the first record (the wt) as one not having and id (for mpnn  output this is usually the case)
                    current_id = None
                    found_first_record = True
                
            elif line.startswith(', 0 氨基酸'): #apply to Chinese version of Snapgene fasta output
                # Skip the metadata line
                continue
            elif line and current_record is not None:
                # This is a sequence line - remove spaces and add to current sequence (the space is a formating in snapgene fasta file)
                cleaned_line = line.replace(' ', '')
                if cleaned_line and all(c in 'ACDEFGHIKLMNPQRSTVWY' for c in cleaned_line):
                    current_sequence.append(cleaned_line)
        
        # Don't forget the last record (if it has an ID)
        if current_record is not None and current_sequence and current_id is not None:
            sequence_str = ''.join(current_sequence).replace(' ', '')
            sequences_with_ids.append((sequence_str, current_id))
        
        # Create JSON structure
        json_data = []
        base_name = input_path.stem  # Get filename without extension
        
        for sequence, original_id in sequences_with_ids:
            entry = {
                "name": f"{base_name}_{original_id}",  
                "modelSeeds": [],
                "sequences": [
                    {
                        "proteinChain": {
                            "count": 1,
                            "sequence": sequence
                        }
                    }
                ],
                "dialect": "alphafoldserver",
                "version": 1
            }
            json_data.append(entry)
        
        # Write to JSON file
        with open(output_file, 'w', encoding='utf-8') as out_f:
            json.dump(json_data, out_f, indent=4, ensure_ascii=False)
        
        print(f"Successfully converted filtered LigandMPNN output to JSON")
        print(f"Input: {input_path}")
        print(f"Output: {output_file}")
        print(f"Processed {len(sequences_with_ids)} sequences (first sequence omitted)")
        
        # Show preview
        if sequences_with_ids:
            print(f"\nPreview of first processed entry:")
            print(f"   Name: {json_data[0]['name']}")
            print(f"   Original ID from file: {sequences_with_ids[0][1]}")
            print(f"   Sequence length: {len(sequences_with_ids[0][0])}")
            print(f"   First 50 chars: {sequences_with_ids[0][0][:50]}...")
        
    except Exception as e:
        print(f"Error during conversion: {str(e)}")
        return 0
    
    return len(sequences_with_ids)

input_file_path = r"C:\Users\Lenovo\file" #directly set the path to YOUR FILE, not folder

if __name__ == "__main__":
    try:
        num_sequences = ligandmpnn_to_json_filtered(input_file_path)
        if num_sequences > 0:
            print(f"\n Success! Your JSON file has been created with {num_sequences} sequences")
            print("   The first sequence (without an 'id') has been omitted.")
        else:
            print("No sequences were processed")
    except Exception as e:
        print(f"Unexpected error: {str(e)}")
                
              

import os
import glob
import argparse
import numpy as np
import subprocess
import sys
import tempfile
import time
import logging
from concurrent.futures import ProcessPoolExecutor
import multiprocessing
import pandas as pd # Import pandas for CSV integration

# Import Schrödinger specific modules (these require Schrödinger environment)
try:
    from schrodinger import structure
    from schrodinger.structutils import analyze
    from schrodinger.structutils import rmsd
except ImportError:
    # Handle cases where Schrödinger modules are not available (e.g., for testing parsing)
    # Full script functionality will require a properly configured Schrödinger environment.
    logging.warning("Schrödinger modules (structure, analyze, rmsd) not found. Script functionality will be limited.")
    structure = None
    analyze = None
    rmsd = None

# Configure logging to file and console
logging.basicConfig(
    level=logging.INFO, # Set to logging.DEBUG to see detailed argument types/values
    format='%(asctime)s - %(processName)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('dock_schrodinger.log'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

def parse_arguments():
    """Parse command-line arguments."""
    parser = argparse.ArgumentParser(description="Dock native ligands into structures using Schrödinger API with parallel processing.")
    parser.add_argument("structures_directory", help="Directory containing structure files (.pdb or .mae) with native ligands")
    parser.add_argument("--max-workers", type=int, default=multiprocessing.cpu_count(),
                        help="Number of concurrent processes (default: number of CPU cores)")
    parser.add_argument("--integrate-results", action="store_true",
                        help="If set, combine all individual *_glide.csv results into a single CSV file, excluding *_glide_skip.csv files.")
    parser.add_argument("--reference-structure", type=str,
                        help="Optional: Path to a reference structure (.pdb or .mae) containing the ligand for RMSD calculation relative to docked poses.")
    return parser.parse_args()

def load_structure(file_path):
    """Load a structure from a file, handling .pdb or .mae formats."""
    if structure is None:
        raise ImportError("Schrödinger 'structure' module not available. Cannot load structures.")
    try:
        reader = structure.StructureReader(file_path)
        st = next(reader)
        logger.info(f"Loaded structure from {file_path} with {len(st.atom)} atoms")
        return st
    except Exception as e:
        logger.error(f"Failed to load structure from {file_path}: {e}")
        raise

def extract_receptor_and_ligand(target_structure):
    """Extract the receptor and the first ligand from the structure."""
    if analyze is None:
        raise ImportError("Schrödinger 'analyze' module not available. Cannot extract receptor/ligand.")
    try:
        # This function finds a ligand based on common definitions.
        ligands = analyze.find_ligands(target_structure)
        if not ligands:
            logger.warning("No ligand found in structure")
            return None, None
        
        # Take the first ligand found
        ligand = ligands[0]
        ligand_indexes = set(ligand.atom_indexes)
        
        # Get all atom indexes in the structure
        all_indexes = set(range(1, target_structure.atom_total + 1))
        
        # Receptor is all atoms not part of the first ligand
        receptor_indexes = list(all_indexes - ligand_indexes)
        
        receptor_structure = target_structure.extract(receptor_indexes)
        ligand_structure = ligand.st # The 'st' attribute is the structure.Structure object for the ligand
        return receptor_structure, ligand_structure
    except Exception as e:
        logger.error(f"Failed to extract receptor and ligand: {e}")
        raise

def prepare_receptor_with_prewizard(receptor_structure, prepwizard_script, options, output_file, base_name):
    """
    Prepares the receptor using prepwizard_driver.py.
    A temporary .mae file is created for input.
    """
    temp_dir = tempfile.gettempdir()
    # Create a unique temporary .mae file path for the receptor input
    temp_input = os.path.join(temp_dir, f"{base_name}_temp_{os.getpid()}.mae")
    
    try:
        # Write the receptor structure directly to the temporary .mae file path.
        # This resolves the issue where the file object 'f' was being passed
        # to Schrodinger's structure.write, causing it to misinterpret the input.
        receptor_structure.write(temp_input) 
        logger.info(f"Wrote temporary receptor file: {temp_input}")
        
        # Add debugging for the types and values of arguments passed to subprocess
        logger.debug(f"DEBUG: prepwizard_script type: {type(prepwizard_script)}, value: {prepwizard_script}")
        logger.debug(f"DEBUG: temp_input type: {type(temp_input)}, value: {temp_input}")
        logger.debug(f"DEBUG: output_file type: {type(output_file)}, value: {output_file}")

        if not os.path.exists(temp_input):
            raise FileNotFoundError(f"Temporary file not created: {temp_input}")
        
        # Construct the command to run prepwizard_driver.py
        # sys.executable ensures the script is run with the same Python interpreter
        cmd = [sys.executable, prepwizard_script, temp_input, output_file] + options
        logger.info(f"Running prepwizard_driver.py with command: {' '.join(cmd)}")
        
        # Execute the subprocess
        result = subprocess.run(cmd, check=True, capture_output=True, text=True)
        logger.info(f"Protein preparation completed for {output_file}")
        if result.stderr:
            logger.warning(f"prepwizard_driver.py stderr: {result.stderr}")
        
        # Verify if the output file was created
        if not os.path.exists(output_file):
            raise RuntimeError(f"Protein preparation failed: output file {output_file} not found")
    except subprocess.CalledProcessError as e:
        logger.error(f"Protein preparation failed with exit code {e.returncode}: {e.stderr}")
        raise
    finally:
        # Attempt to clean up the temporary input file
        max_attempts = 5
        for attempt in range(max_attempts):
            try:
                if os.path.exists(temp_input):
                    os.remove(temp_input)
                    logger.info(f"Cleaned up temporary file: {temp_input}")
                break
            except PermissionError as e:
                # Handle cases where the file might still be in use by another process
                if attempt < max_attempts - 1:
                    logger.warning(f"Failed to delete {temp_input} (attempt {attempt + 1}/{max_attempts}): {e}")
                    time.sleep(1)
                else:
                    logger.error(f"Failed to delete {temp_input} after {max_attempts} attempts: {e}")

def calculate_grid_center_and_size(ligand_structure):
    """Calculate grid center and box size based on the ligand."""
    try:
        coords = np.array([atom.xyz for atom in ligand_structure.atom])
        centroid = coords.mean(axis=0)
        distances = np.linalg.norm(coords - centroid, axis=1)
        max_distance = distances.max()
        # Add a buffer of 10 Å around the ligand
        outer_box_size = 2 * max_distance + 10
        return centroid, outer_box_size
    except Exception as e:
        logger.error(f"Failed to calculate grid parameters: {e}")
        raise

def write_combined_inp(receptor_file, ligand_file, job_name, centroid, outer_box_size):
    """Write a combined .inp file for grid generation and docking."""
    try:
        inp_file = f"{job_name}.inp"
        content = f"""
[ GRIDGEN ]
    RECEP_FILE {receptor_file}
    RECEP_VSCALE 1.0
    GRID_CENTER {centroid[0]:.6f}, {centroid[1]:.6f}, {centroid[2]:.6f}
    INNERBOX 3, 3, 3
    OUTERBOX 39, 39, 39

[ DOCKING ]
    LIGANDFILE {ligand_file}
    PRECISION XP
    EXPANDED_SAMPLING   True
    POSES_PER_LIG 10
    LIG_VSCALE 0.8
    WRITE_XP_DESC False
    MAX_ITERATIONS   150
    MAXKEEP   500000
    MAXREF   4000
"""
        with open(inp_file, 'w') as f:
            f.write(content.strip())
        logger.info(f"Generated combined input file: {inp_file}")
        return inp_file
    except Exception as e:
        logger.error(f"Failed to write .inp file: {e}")
        raise

def run_glide(inp_file, schrodinger_path, job_name):
    """Run Glide using the provided .inp file."""
    try:
        glide_cmd = os.path.join(schrodinger_path, 'glide')
        # Added -JOBNAME flag with the job_name
        args = [glide_cmd, inp_file, '-HOST', 'localhost:1', '-NJOBS', '1', '-JOBNAME', job_name, '-TMPLAUNCHDIR']
        logger.info(f"Running Glide with command: {' '.join(args)}")
        result = subprocess.run(args, text=True, check=True, capture_output=True)
        logger.info(f"Glide completed for {inp_file}")
        if result.stderr:
            logger.warning(f"Glide stderr: {result.stderr}")
    except subprocess.CalledProcessError as e:
        logger.error(f"Glide failed for {inp_file} with exit code {e.returncode}: {e.stderr}")
        raise

def integrate_glide_results(directory, rmsd_data):
    """
    Finds all *_glide.csv files (excluding *_glide_skip.csv) in the specified directory
    and combines them into a single CSV file, including RMSD data.
    """
    logger.info(f"Integrating Glide results from directory: {directory}")
    all_csv_files = glob.glob(os.path.join(directory, "*.csv"))
    
    result_files = []
    for csv_file in all_csv_files:
        if csv_file.endswith("_glide.csv") and not csv_file.endswith("_glide_skip.csv"):
            result_files.append(csv_file)

    if not result_files:
        logger.warning(f"No *_glide.csv files (excluding *_skip.csv) found for integration in {directory}.")
        return

    combined_df = pd.DataFrame()
    for f in result_files:
        try:
            df = pd.read_csv(f)
            combined_df = pd.concat([combined_df, df], ignore_index=True)
            logger.info(f"Added {os.path.basename(f)} to combined results.")
        except Exception as e:
            logger.warning(f"Could not read {f} for integration: {e}")

    if not combined_df.empty:
        # Create a DataFrame from RMSD data
        rmsd_df = pd.DataFrame(rmsd_data)
        # Assuming 's_l_Title' in Glide CSV for joining, which should match the base_name
        rmsd_df.rename(columns={'base_name': 's_l_Title'}, inplace=True) 
        
        # Merge combined_df with rmsd_df based on the ligand title
        combined_df = pd.merge(combined_df, rmsd_df, on='s_l_Title', how='left')

        output_path = os.path.join(directory, "combined_docking_results.csv")
        combined_df.to_csv(output_path, index=False)
        logger.info(f"Combined results saved to: {output_path}")
    else:
        logger.info("No data to combine. Combined results file not created.")


def process_structure(structure_file, prepwizard_script, options, schrodinger_path, reference_structure_path=None):
    """
    Process a single structure: load, prepare, generate grid, and dock.
    Returns a dictionary with base_name and calculated RMSD (or None).
    """
    # Extract base name without extension
    base_name = os.path.splitext(os.path.basename(structure_file))[0]
    
    # Define a sub-directory for prepared structures relative to the input file
    output_directory = os.path.join(os.path.dirname(structure_file), "prepared_structures")
    os.makedirs(output_directory, exist_ok=True) # Ensure the directory exists

    # Define absolute paths for output files
    receptor_file = os.path.join(output_directory, f"{base_name}_receptor_prepared.maegz")
    ligand_file = os.path.join(output_directory, f"{base_name}_ligand.maegz")
    job_name = f"{base_name}_glide" # Glide job name, used for -JOBNAME flag

    rmsd_value = None # Initialize RMSD value

    try:
        # Load input structure (original complex)
        input_complex_st = load_structure(structure_file)
        
        # Extract receptor and ligand from the input complex
        receptor_st_input, ligand_st_input = extract_receptor_and_ligand(input_complex_st)
        if receptor_st_input is None:
            logger.warning(f"No ligand found in {structure_file}, skipping processing for this file.")
            return {"base_name": base_name, "RMSD": None} # Return for integration even if skipped
        
        # Prepare receptor using prepwizard
        prepare_receptor_with_prewizard(receptor_st_input, prepwizard_script, options, receptor_file, base_name)
        
        # Save the extracted ligand to an MAEGZ file
        ligand_st_input.write(ligand_file)
        logger.info(f"Saved ligand to {ligand_file}")
        
        # Calculate grid center and size based on the ligand's coordinates
        centroid, outer_box_size = calculate_grid_center_and_size(ligand_st_input)
        
        # Write the combined .inp file for Glide
        inp_file = write_combined_inp(receptor_file, ligand_file, job_name, centroid, outer_box_size)
        
        # Run Glide, passing the job_name for the -JOBNAME flag
        run_glide(inp_file, schrodinger_path, job_name)

        # --- Optional RMSD Calculation ---
        if reference_structure_path:
            # Check if Schrödinger's RMSD module is available
            if rmsd is None or analyze is None or structure is None:
                logger.error("Schrödinger 'rmsd', 'analyze', or 'structure' module not available. Cannot calculate RMSD.")
            elif not os.path.exists(reference_structure_path):
                logger.error(f"Reference structure not found at '{reference_structure_path}'. Skipping RMSD calculation for {base_name}.")
            else:
                glide_output_maegz = f"{job_name}-out.maegz" # Standard Glide output name
                if os.path.exists(glide_output_maegz):
                    try:
                        # Load the full docked complex (receptor + ligand) from Glide output
                        docked_complex_st = load_structure(glide_output_maegz)
                        
                        # Load the full reference structure (receptor + ligand)
                        ref_overall_st = load_structure(reference_structure_path)

                        # Extract receptor and ligand from both structures
                        docked_receptor_st, docked_ligand_st = extract_receptor_and_ligand(docked_complex_st)
                        reference_receptor_st, reference_ligand_st = extract_receptor_and_ligand(ref_overall_st)

                        if docked_ligand_st and reference_ligand_st and docked_receptor_st and reference_receptor_st:
                            # 1. Superimpose the docked complex onto the reference complex based on protein atoms
                            # The transformation is calculated using the receptor structures
                            # We use atom_match_prop='element' for robust atom matching during superposition
                            
                            # Ensure both receptors have atoms for superposition
                            if docked_receptor_st.atom_total > 0 and reference_receptor_st.atom_total > 0:
                                # Get transformation based on receptor alignment
                                # The returned xform applies to the 'moving_st' (docked_receptor_st here)
                                _, xform = rmsd.get_rmsd_and_superimpose_transformation(
                                    docked_receptor_st, reference_receptor_st, atom_match_prop='element'
                                )
                                
                                # Apply the transformation to the *entire* docked complex
                                # This creates a new structure where the protein is aligned to reference
                                docked_complex_st_aligned = xform.transform_structure(docked_complex_st)
                                
                                # Extract the ligand from the newly aligned complex
                                _, docked_ligand_st_aligned = extract_receptor_and_ligand(docked_complex_st_aligned)
                                
                                if docked_ligand_st_aligned:
                                    # 2. Calculate RMSD of the docked ligand relative to the reference ligand
                                    # after the protein-based alignment
                                    rmsd_value = rmsd.get_rmsd(
                                        docked_ligand_st_aligned, reference_ligand_st, atom_match_prop='element'
                                    )
                                    logger.info(f"RMSD of docked ligand relative to reference (protein-aligned) for {base_name}: {rmsd_value:.3f} Å")
                                else:
                                    logger.warning(f"Could not re-extract ligand from aligned docked complex for {base_name}. Skipping RMSD calculation.")
                            else:
                                logger.warning(f"Receptor structures are empty for {base_name}. Skipping protein-aligned RMSD calculation.")
                        else:
                            logger.warning(f"Could not extract both receptor and ligand from either docked or reference structure for {base_name}. Skipping RMSD calculation.")
                    except Exception as e:
                        logger.error(f"Error encountered during protein-aligned RMSD calculation for {base_name}: {e}", exc_info=True)
                else:
                    logger.warning(f"Glide output MAEGZ file not found at '{glide_output_maegz}' for RMSD calculation for {base_name}.")
        # --- End Optional RMSD Calculation ---
        
        logger.info(f"Completed processing for {structure_file}")
        return {"base_name": base_name, "RMSD": rmsd_value} # Return the RMSD value for integration
    except Exception as e:
        logger.error(f"Error processing {structure_file}: {e}", exc_info=True)
        # Return a dictionary with None for RMSD if an error occurs during processing
        return {"base_name": base_name, "RMSD": None} 

def main():
    args = parse_arguments()
    structures_directory = os.path.abspath(args.structures_directory)
    max_workers = args.max_workers
    reference_structure_path = args.reference_structure # Get the reference structure path

    # Retrieve Schrödinger installation path from environment variable
    schrodinger_path = os.environ.get('SCHRODINGER', '')
    if not schrodinger_path:
        logger.error("SCHRODINGER environment variable not set. Please set it to your Schrödinger installation directory.")
        sys.exit(1)
    
    # Construct the path to prepwizard_driver.py
    # Assumes a standard Schrödinger installation structure
    mmshare_dir = os.path.join(schrodinger_path, 'mmshare') 
    prepwizard_script = os.path.join(mmshare_dir, 'python', 'scripts', 'prepwizard_driver.py')
    
    if not os.path.exists(prepwizard_script):
        # Fallback to a common older path if the first one doesn't exist
        mmshare_dir = os.path.join(schrodinger_path, 'mmshare-v6.9') 
        prepwizard_script = os.path.join(mmshare_dir, 'python', 'scripts', 'prepwizard_driver.py')
        if not os.path.exists(prepwizard_script):
            logger.error(f"prepwizard_driver.py not found at expected paths: {os.path.join(schrodinger_path, 'mmshare', 'python', 'scripts', 'prepwizard_driver.py')} or {os.path.join(schrodinger_path, 'mmshare-v6.9', 'python', 'scripts', 'prepwizard_driver.py')}. Please verify your SCHRODINGER environment variable and Schrödinger installation.")
            sys.exit(1)
    
    # Command-line options for prepwizard_driver.py
    # These options configure how the protein is prepared (e.g., adding hydrogens, optimizing, etc.)
    options = [
        "-fillsidechains",
        "-disulfides",
        "-assign_all_residues",
        "-rehtreat",
        "-max_states", "1",
        "-epik_pH", "7.4",
        "-epik_pHt", "2.0",
        "-antibody_cdr_scheme", "Kabat",
        "-samplewater",
        "-include_epik_states",
        "-propka_pH", "7.4",
        "-f", "S-OPLS",
        "-rmsd", "0.3",
        "-watdist", "5.0"
    ]
    
    # Get all .pdb and .mae files from the specified directory
    structure_files = glob.glob(os.path.join(structures_directory, "*.pdb")) + \
                      glob.glob(os.path.join(structures_directory, "*.mae"))
    
    if not structure_files:
        logger.warning(f"No .pdb or .mae files found in directory: {structures_directory}. Exiting.")
        sys.exit(1)
    
    # Process structures in parallel using a ProcessPoolExecutor
    logger.info(f"Starting parallel processing for {len(structure_files)} structures with {max_workers} workers.")
    
    # List to store RMSD results for integration
    all_rmsd_results = []

    with ProcessPoolExecutor(max_workers=max_workers) as executor:
        futures = [
            # Pass the reference_structure_path to process_structure
            executor.submit(process_structure, structure_file, prepwizard_script, options, schrodinger_path, reference_structure_path)
            for structure_file in structure_files
        ]
        # Iterate over futures to catch and log any exceptions from worker processes
        for future in futures:
            try:
                # Store the result from each process (which includes base_name and RMSD)
                result_data = future.result()
                if result_data: # Ensure result_data is not None (e.g., if a process totally failed)
                    all_rmsd_results.append(result_data)
            except Exception as e:
                logger.error(f"A task failed during parallel execution: {e}", exc_info=True) # exc_info=True logs traceback

    # Integrate results if the --integrate-results flag is set
    if args.integrate_results:
        # Changed to os.getcwd() to look for CSVs in the terminal's working directory
        integrate_glide_results(os.getcwd(), all_rmsd_results)

if __name__ == "__main__":
    main()

                

# superpose_pymol.py

import os
import glob
import argparse
import numpy as np 
import subprocess
import time
import sys

# Define a timeout
FILE_WAIT_TIMEOUT = 30
FILE_WAIT_INTERVAL = 1

def parse_arguments():
    """Parses command-line arguments for directory, reference structure."""
    parser = argparse.ArgumentParser(description="Superpose PDB files onto a reference structure using PyMOL.")
    parser.add_argument("input_directory", help="Directory containing PDB files to superpose")
    parser.add_argument("ref_structure", help="Path to the reference structure file (e.g., .pdb or .mae)")
    parser.add_argument("--output_directory", default=os.getcwd(), 
                        help="Directory to save the superposed structure files. Defaults to current working directory.")
    return parser.parse_args()

def load_structure_simple(file_path):
    """
    Load a structure from a file for simple PyMOL use.
    This is a basic placeholder; PyMOL handles the actual PDB parsing.
    We just need to confirm file existence for input.
    """
    if not os.path.exists(file_path):
        raise FileNotFoundError(f"Structure file not found: {file_path}")
    print(f"Confirmed existence of structure file: {file_path}")
    # We don't actually parse it here, PyMOL will.
    return file_path # Return path as confirmation

def _get_pymol_executable_path():
    """
    Determines the best PyMOL executable path.
    Prefers 'pymol.exe' (console) over 'PyMOLWin.exe' (GUI).
    """
    base_pymol_path = r"C:\ProgramData\pymol" # Adjust this path if your PyMOL is installed elsewhere
    pymol_console_exe = os.path.join(base_pymol_path, "pymol.exe")
    pymol_win_exe = os.path.join(base_pymol_path, "PyMOLWin.exe")

    if os.path.exists(pymol_console_exe):
        print(f"Using console PyMOL executable: {pymol_console_exe}")
        return pymol_console_exe
    elif os.path.exists(pymol_win_exe):
        print(f"Using GUI PyMOL executable (fallback): {pymol_win_exe}")
        return pymol_win_exe
    else:
        raise FileNotFoundError(f"Neither '{pymol_console_exe}' nor '{pymol_win_exe}' found. Please verify PyMOL installation path in _get_pymol_executable_path().")

def superpose_structures_pymol(mobile_pdb_path, fixed_complex_pdb_path, output_aligned_pdb_path):
    """
    Superpose a mobile protein structure onto the protein part of a fixed complex (protein+ligand)
    using PyMOL's 'super' command, and then save the aligned mobile protein along with the ligand
    from the fixed complex.
    """
    fixed_complex_obj_name = "ref_complex"
    mobile_protein_obj_name = "input_protein"
    output_combined_obj_name = "aligned_complex_with_ligand"

    # Convert all Windows backslashes to forward slashes for PyMOL commands
    mobile_pdb_path_fwd = mobile_pdb_path.replace('\\', '/')
    fixed_complex_pdb_path_fwd = fixed_complex_pdb_path.replace('\\', '/')
    output_aligned_pdb_path_fwd = output_aligned_pdb_path.replace('\\', '/')

    # PyMOL commands to be written into the .pml script
    pymol_commands = [
        f"load \"{fixed_complex_pdb_path_fwd}\", {fixed_complex_obj_name}",
        f"load \"{mobile_pdb_path_fwd}\", {mobile_protein_obj_name}",
        f"super {mobile_protein_obj_name}, {fixed_complex_obj_name} and polymer.protein",
        f"create {output_combined_obj_name}, {mobile_protein_obj_name} or ({fixed_complex_obj_name} and not polymer.protein)",
        f"save {str(output_aligned_pdb_path_fwd)}, {output_combined_obj_name}", 
        "quit" # Ensure PyMOL exits cleanly
    ]

    command_string = "; ".join(pymol_commands)
    
    # Get the appropriate PyMOL executable path
    pymol_cmd_path = _get_pymol_executable_path()
    
    # Create a temporary PyMOL script file to execute commands
    temp_script_path = os.path.join(os.getcwd(), "temp_pymol_script.pml")
    temp_script_path_fwd = temp_script_path.replace('\\', '/') 

    with open(temp_script_path, 'w') as f:
        f.write(command_string)
    
    # print("\n--- Diagnostic Information for PyMOL Execution ---")
    # print(f"Python interpreter running script: {sys.executable}")
    # print(f"PyMOL executable path used: {pymol_cmd_path}")
    # print(f"Temporary PyMOL script path: {temp_script_path}")
    # print(f"Mobile PDB (input) path: {os.path.abspath(mobile_pdb_path)}")
    # print(f"Fixed Complex PDB (reference) path: {os.path.abspath(fixed_complex_pdb_path)}")
    # print(f"Output Aligned PDB expected path: {os.path.abspath(output_aligned_pdb_path)}")
    # print(f"Current Working Directory (cwd): {os.getcwd()}")
    # print("--- Key Environment Variables for PyMOL Subprocess ---")
    
    pymol_env = os.environ.copy()
    pymol_base_dir = os.path.dirname(pymol_cmd_path) 
    pymol_env['PATH'] = pymol_base_dir + os.pathsep + pymol_env.get('PATH', '')
    # print(f"  PATH (for PyMOL subprocess) will start with: {pymol_base_dir}{os.pathsep}...")
    # print("-------------------------------------------\n")
    # print(f"Contents of temp_pymol_script.pml:\n{command_string}")
    # print("-------------------------------------------\n")

    full_command_args = [pymol_cmd_path, '-c', temp_script_path_fwd]
    
    print(f"Attempting to execute PyMOL with command arguments: {full_command_args}")
    try:
        result = subprocess.run(full_command_args, text=True, check=True, cwd=os.getcwd(), capture_output=True, env=pymol_env) 
        
        print(f"PyMOL process finished with exit code {result.returncode}.")
        if result.stdout:
            print(f"PyMOL stdout:\n{result.stdout}")
        if result.stderr:
            print(f"PyMOL stderr:\n{result.stderr}")
        
        print(f"Waiting for output file to appear: {output_aligned_pdb_path}")
        file_found = False
        start_time = time.time()
        while time.time() - start_time < FILE_WAIT_TIMEOUT:
            if os.path.exists(output_aligned_pdb_path) and os.path.getsize(output_aligned_pdb_path) > 0:
                print(f"Output file found and is non-empty: {output_aligned_pdb_path}")
                file_found = True
                break
            time.sleep(FILE_WAIT_INTERVAL)
        
        if not file_found:
            raise RuntimeError(f"PyMOL output file '{output_aligned_pdb_path}' was not found or remained empty after {FILE_WAIT_TIMEOUT} seconds.")

        print(f"Successfully superposed {mobile_pdb_path} onto {fixed_complex_pdb_path} and saved to {output_aligned_pdb_path}")

        return output_aligned_pdb_path
    except FileNotFoundError:
        raise RuntimeError(f"PyMOL executable not found at '{pymol_cmd_path}'. Please ensure the path is correct or 'pymol' is in your system PATH.")
    except subprocess.CalledProcessError as e:
        print(f"PyMOL process exited with error code {e.returncode}")
        print(f"Command executed by subprocess: {' '.join(e.cmd)}") 
        if e.stdout:
            print(f"PyMOL stdout (from error):\n{e.stdout}")
        if e.stderr:
            print(f"PyMOL stderr (from error):\n{e.stderr}") 
        raise RuntimeError(f"PyMOL superposition failed. See above output for details.")
    except Exception as e:
        raise RuntimeError(f"An unexpected error occurred during PyMOL execution: {e}")
    finally:
        if os.path.exists(temp_script_path):
            os.remove(temp_script_path)


def main():
    args = parse_arguments()
    input_directory = os.path.abspath(args.input_directory)
    ref_structure_path = os.path.abspath(args.ref_structure)
    output_directory = os.path.abspath(args.output_directory)

    os.makedirs(output_directory, exist_ok=True)
    
    # Load reference structure (just check existence, PyMOL will do actual parsing)
    fixed_pymol_ref = load_structure_simple(ref_structure_path)

    pdb_files = glob.glob(os.path.join(input_directory, "*.pdb"))
    if not pdb_files:
        print(f"No .pdb files found in input directory: {input_directory}")
        return

    for pdb_file in pdb_files:
        base_name = os.path.splitext(os.path.basename(pdb_file))[0]
        output_aligned_pdb_path = os.path.join(output_directory, f"{base_name}_aligned_complex.pdb")
        print(f"Processing {pdb_file} for superposition...")
        
        try:
            superpose_structures_pymol(
                mobile_pdb_path=pdb_file, 
                fixed_complex_pdb_path=fixed_pymol_ref, 
                output_aligned_pdb_path=output_aligned_pdb_path
            )
            print(f"Finished superposition for {pdb_file}. Output saved to {output_aligned_pdb_path}")
        except Exception as e:
            print(f"Error superposing {pdb_file}: {e}")
            continue 

if __name__ == "__main__":
    main()
                

<ROSETTASCRIPTS>
    <SCOREFXNS>
        <ScoreFunction name="r15_cart" weights="ref2015"/>
    </SCOREFXNS>
    <RESIDUE_SELECTORS>
    </RESIDUE_SELECTORS>
    <PACKER_PALETTES>
    </PACKER_PALETTES>
    <TASKOPERATIONS>
    </TASKOPERATIONS>
    <MOVE_MAP_FACTORIES>
    </MOVE_MAP_FACTORIES>
    <SIMPLE_METRICS>
        <SasaMetric name="TotalSASA" sasa_metric_mode="all_sasa" />
        <RMSDMetric name="rmsd" rmsd_type="rmsd_protein_bb_ca" use_native="1" super="true"/>
    </SIMPLE_METRICS>
    <FILTERS>
        <ExposedHydrophobics name="Exposed" />
    </FILTERS>
    <MOVERS>
        <RunSimpleMetrics name="run_metrics1" metrics="rmsd"/>
    </MOVERS>
    <PROTOCOLS>
        <Add metrics="TotalSASA" filter_name="Exposed" />
        <Add mover_name="run_metrics1"/>
    </PROTOCOLS>
    <OUTPUT scorefxn="r15_cart" />
</ROSETTASCRIPTS>