#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@file               :mainfix.py
@Description:       :批量调用modelbuilder.py进行结构修复
@Date               :2024/01/12 13:38:53
@Author             :lyzeng
@Email              :pylyzeng@gmail.com
@version            :1.0
'''


import argparse
from pathlib import Path
import subprocess
import shutil
import concurrent.futures
import os

def process_pdb_file(pdb_file: Path, script_path: Path, working_dir: Path):
    pdb_id = pdb_file.stem
    runner_dir = working_dir / f"runner_{pdb_id}"
    runner_dir.mkdir(exist_ok=True)

    copied_pdb_file = runner_dir / pdb_file.name
    shutil.copy(pdb_file, copied_pdb_file)

    # 构建并运行modelbuilder.py脚本的命令
    command = ["/home/era/.conda/envs/modeller/bin/python", script_path.absolute().as_posix(), copied_pdb_file.absolute().as_posix()]
    subprocess.run(command, cwd=runner_dir)


def main(pdb_dir_path, script_path):
    pdb_dir = Path(pdb_dir_path)
    pdb_files = list(pdb_dir.glob("*.pdb"))

    script_path = Path(script_path).absolute() # 确保script_path是绝对路径

    max_workers = max(1, os.cpu_count() - 1)  # 保留一个核心给系统

    with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
        futures = [executor.submit(process_pdb_file, pdb_file, script_path, pdb_dir) for pdb_file in pdb_files]
        for future in concurrent.futures.as_completed(futures):
            future.result()

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Process multiple PDB files using modelbuilder.py")
    parser.add_argument("pdb_dir_path", help="Directory containing PDB files")
    parser.add_argument("script_path", help="Path to the modelbuilder.py script")
    args = parser.parse_args()

    main(args.pdb_dir_path, args.script_path)

# python mainfix.py pdb_test1 modelbuilder.py