from pathlib import Path
import subprocess
import os

def generate_gmx_top_files(directory, num_cores, forcefield="amber99sb-ildn", water_model="tip3p"):
    """
    Generate GROMACS topology (.top) files for each PDB file in the specified directory
    by running the 'gmx_mpi pdb2gmx' command in batch, with the working directory set to the specified directory.

    Args:
    directory (str): Directory containing PDB files.
    num_cores (int): Number of cores to use for mpirun.
    forcefield (str): Forcefield parameter for GROMACS. Default is "amber99sb-ildn".
    water_model (str): Water model parameter for GROMACS. Default is "tip3p".
    """
    # Ensure the GMXRC path is sourced if set
    gmxrc_path = os.getenv("GMXRC_PATH")
    if gmxrc_path:
        subprocess.run(f"source {gmxrc_path}", shell=True, executable='/bin/bash')

    # Find all PDB files in the specified directory
    pdb_files = Path(directory).glob("*.pdb")

    # Iterate over each PDB file and execute the GROMACS command
    for pdb_file in pdb_files:
        name = pdb_file.stem  # Extract the name without the '.pdb' extension
        output_gro = f"{name}.gro"
        topol_file = "topol.top"

        # Construct and run the command with the working directory set to 'directory'
        cmd = [
            "mpirun", "-np", str(num_cores),
            "gmx_mpi", "pdb2gmx", 
            "-f", str(pdb_file), 
            "-o", output_gro, 
            "-ff", forcefield, 
            "-water", water_model, 
            "-ignh", 
            "-p", topol_file
        ]
        subprocess.run(cmd, check=True, cwd=directory)

# Example usage
directory = "pdb_top"  # Replace with your actual directory path
num_cores = max(os.cpu_count() - 2, 1)  # Determine the number of cores to use
generate_gmx_top_files(directory, num_cores)  # Run the function