import os
import shutil
from pathlib import Path
import sys

# 读取参数
import argparse

parser = argparse.ArgumentParser(description="合并 vina 拆分输出并检查转换结果")
parser.add_argument('--n_splits', type=int, default=12, help='拆分数量')
parser.add_argument('--out_dir', type=str, default='./result', help='meeko 转化出 sdf 文件所在目录')
parser.add_argument('--output_prefix', type=str, default='poses', help='每个子任务输出目录前缀')
parser.add_argument('--poses_dir', type=str, default='./result/poses_all', help='合并后SDF输出目录')
args = parser.parse_args()

n_splits = args.n_splits
out_dir = Path(args.out_dir)
poses_dir = Path(args.poses_dir)
poses_dir.mkdir(exist_ok=True)


merged_count = 0
# 合并所有 output{1..n_splits}/*.sdf 文件并复制
for i in range(1, n_splits + 1):
    part_dir = out_dir / f'{args.output_prefix}{i}'
    if part_dir.exists():
        sdfs = list(part_dir.glob('*_converted.sdf'))
        for f in sdfs:
            shutil.copy2(f, poses_dir / f.name)
            merged_count += 1
    else:
        print(f"Warning: {part_dir} does not exist.")

# 检查数量
sdf_stems = set(f.stem.replace('_converted', '') for f in poses_dir.glob('*_converted.sdf'))
pdbqt_stems = set(f.stem for f in out_dir.glob('*_out.pdbqt'))

print(f"pdbqt数量: {len(pdbqt_stems)}, 合并后sdf数量: {len(sdf_stems)}")

if len(sdf_stems) == len(pdbqt_stems) and sdf_stems == pdbqt_stems:
    print("所有配体都已成功转换！")
else:
    missing = tuple(sorted(pdbqt_stems - sdf_stems))
    if missing:
        print(f"缺少{len(missing)}个配体转换结果，缺失文件名如下：")
        for name in missing:
            print(name)
    extra = tuple(sorted(sdf_stems - pdbqt_stems))
    if extra:
        print(f"有{len(extra)}个多余的SDF文件（没有对应的pdbqt），多余文件名如下：")
        for name in extra:
            print(name)

# 合并12
# python vina_merge_and_check.py
# 合并24
# python vina_merge_and_check.py --n_splits 128 --out_dir ./result --output_prefix poses --poses_dir ./result/poses_all