#!/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
import sys

def process_pdb_file(pdb_file: Path, script_path: Path, working_dir: Path, python_interpreter: str, typestr="refine.very_fast", use_pymol_align_merge=False):
    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 = [python_interpreter, script_path.absolute().as_posix(), copied_pdb_file.absolute().as_posix()]

    # 添加额外的参数
    command += ["--typestr", typestr]
    if use_pymol_align_merge:
        command += ["--use_pymol_align_merge"]

    # 运行modelbuilder.py脚本
    subprocess.run(command, cwd=runner_dir)

def 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")
    parser.add_argument("--python_interpreter", default=sys.executable, help="Path to the Python interpreter")
    args = parser.parse_args()

    pdb_dir_path = args.pdb_dir_path
    script_path = args.script_path
    python_interpreter = args.python_interpreter

    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, python_interpreter) for pdb_file in pdb_files]
        for future in concurrent.futures.as_completed(futures):
            future.result()

if __name__ == "__main__":
    main()


# python mainfix.py pdb_test1 modelbuilder.py --python_interpreter /root/micromamba/envs/pyfastx/bin/python
# 复制修复完成文件到指定文件夹
# import os
# import shutil
# from pathlib import Path

# # 设置源目录和目标目录
# source_dir = Path('pdb_test1')
# target_dir = Path('newfix')

# # 确保目标目录存在
# target_dir.mkdir(exist_ok=True)

# # 遍历源目录
# for subdir in source_dir.glob('runner_*'):
#     if subdir.is_dir():
#         # 查找所有匹配的文件
#         for file in subdir.glob('*.modellerfix.pdb'):
#             # 移动文件到目标目录
#             shutil.move(str(file), target_dir / file.name)

# print("文件移动完成。")

