from pathlib import Path
import shutil

# 设置源目录和目标目录模板
src_dir = Path('./newfix')
dest_dir_template = Path('./pdb_test1/runner_{pdb_id}')

# 查找所有符合模式的文件
for file_path in src_dir.glob('*.modellerfix.pdb'):
    # 提取PDB ID
    pdb_id = file_path.stem.split('.')[0]
    # 设置目标目录
    dest_dir = dest_dir_template.with_name(dest_dir_template.name.format(pdb_id=pdb_id))
    # 确保目标目录存在
    dest_dir.mkdir(parents=True, exist_ok=True)
    # 复制文件
    shutil.copy(file_path, dest_dir / file_path.name)

# 检查和统计文件
count = 0
for pdb_id in src_dir.glob('*.modellerfix.pdb'):
    pdb_id = pdb_id.stem.split('.')[0]
    dest_dir = dest_dir_template.with_name(dest_dir_template.name.format(pdb_id=pdb_id))
    if (dest_dir / pdb_id).exists():
        count += 1

print(f"共复制并核对了 {count} 个文件")

# 遍历pdb_test1下的所有runner_*目录
for runner_dir in Path('./pdb_test1').glob('runner_*'):
    pdb_id = runner_dir.name.split('_')[1]
    files = list(runner_dir.glob(f'{pdb_id}.modellerfix.pdb'))

    # 确保每个目录中恰有一个匹配的文件
    assert len(files) == 1, f"目录 {runner_dir} 中找到了 {len(files)} 个匹配的文件，期望是 1 个"

print("所有runner_*目录下的文件检查完毕，均符合一对一的条件。")