import os, json
from collections import defaultdict

save_path = "/mnt/shared-storage-user/caipengxiang/workspace/ChemBOMAS/info_encoder/distance_cluster_results/summary"
json_path = "/mnt/shared-storage-user/caipengxiang/workspace/ChemBOMAS/info_encoder/distance_cluster_results"

os.makedirs(save_path, exist_ok=True)

methods = ['kmeans', 'hierarchical', 'Spectral']

for method in methods:
    # 获取当前目录所有JSON文件
    files = [f for f in os.listdir(json_path) if f.endswith('.json') and method in f]

    # 创建按聚类方法分类的数据结构
    clusters = defaultdict(dict)

    for filename in files:
        filename = os.path.join(json_path, filename)

        # 读取文件内容
        with open(filename, "r") as f:
            data = json.load(f)
        
        # 存储到对应聚类方法
        # import pdb; pdb.set_trace()
        clusters[method] = {**clusters[method], **data}

    # 保存合并后的文件
    for method, data in clusters.items():
        output_file = os.path.join(save_path, f'{method}_clusters.json')
        with open(output_file, 'w') as f:
            json.dump(data, f, indent=2)