from pathlib import Path
import pyrosetta
from multiprocessing import Pool
from shutil import copyfile
import logging

# fix = ['6bga', '5ksb', '4z7v', '4p2o', '4ozi', '3qiu', '2ypl', '1zgl', '7z50', '6u3n']
fix = ['1zgl']

def fix_optimize(file: Path, out_file: Path):
    # 设置日志
    log_file = out_file.with_suffix('.log')
    logger = logging.getLogger(log_file.name)
    logger.setLevel(logging.INFO)
    formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
    file_handler = logging.FileHandler(log_file)
    file_handler.setFormatter(formatter)
    logger.addHandler(file_handler)

    logger.info(f'Processing {file.name}')

    pyrosetta.init()
    pose = pyrosetta.pose_from_pdb(file.as_posix())
    scorefxn = pyrosetta.create_score_function('ref2015')
    relax = pyrosetta.rosetta.protocols.relax.FastRelax(scorefxn)
    relax.apply(pose)
    pose.dump_pdb(out_file.as_posix())

    logger.info(f'Finished processing {file.name}')

if __name__ == '__main__':
    dir1 = Path('../pdb_test8')
    dir2 = Path('../pdb_test7')
    dirs = [dir1, dir2]
    files = []
    for dir in dirs:
        files.extend(list(dir.rglob('*.modellerfix.pdb')))
    for file in files:
        if file.name.split('.')[0] in fix:
            print(file)
            target = Path('/mnt/mydrive/analysis_pdb-dev/manualfix') / file.name
            copyfile(file.as_posix(), target.as_posix())
    pyrosetta_fix = list(Path('/mnt/mydrive/analysis_pdb-dev/manualfix').rglob('*.modellerfix.pdb'))
    with Pool(16) as p:
        p.starmap(fix_optimize, [(file, file.with_stem(file.stem + '.fastrelax')) for file in pyrosetta_fix])
    print('fastrelax done')