from pymatgen.core import Structure
from pymatgen.analysis.structure_matcher import StructureMatcher
import os
import shutil

# 定义函数，用于将 CIF 文件进行分类并输出

def structure_classify(input_folder, output_folder):
    """
    分类 input_folder 中的 CIF 文件，并根据框架分组后存储到 output_folder 中。

    :param input_folder: 输入文件夹路径，包含 .cif 文件。
    :param output_folder: 输出文件夹路径，用于存储分类后的 .cif 文件。
    """
    # 检查输入文件夹是否存在
    if not os.path.exists(input_folder):
        print(f"输入文件夹 {input_folder} 不存在")
        return

    # 创建输出文件夹（如果不存在）
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    # 读取输入文件夹中的所有 .cif 文件
    structures = []
    file_map = {}  # 记录文件路径与结构的对应关系
    for filename in os.listdir(input_folder):
        if filename.endswith(".cif"):
            file_path = os.path.join(input_folder, filename)
            try:
                structure = Structure.from_file(file_path)
                structures.append(structure)
                file_map[id(structure)] = filename  # 使用结构对象的 id 作为键
            except Exception as e:
                print(f"无法读取文件 {filename}：{e}")

    # 检查是否成功加载了任何结构
    if not structures:
        print("未找到有效的 CIF 文件")
        return

    # 分组结构
    matcher = StructureMatcher()
    grouped_structures = []

    for structure in structures:
        matched = False
        for group in grouped_structures:
            if matcher.fit(structure, group[0]):  # 比较结构是否匹配
                group.append(structure)
                matched = True
                break
        if not matched:
            grouped_structures.append([structure])

    # 保存分组后的结构到输出文件夹
    for group_index, group in enumerate(grouped_structures):
        group_folder = os.path.join(output_folder, f"group_{group_index}")
        os.makedirs(group_folder, exist_ok=True)

        for structure_index, structure in enumerate(group):
            output_file = os.path.join(group_folder, f"structure_{structure_index}.cif")
            structure.to(filename=output_file)

            # 获取原始文件名并复制到分组文件夹
            original_filename = file_map[id(structure)]  # 使用结构对象的 id 获取文件名
            original_file_path = os.path.join(input_folder, original_filename)
            shutil.copy(original_file_path, group_folder)

    print(f"处理完成，分类后的结构已保存到 {output_folder}")

if __name__ == "__main__":
    # 示例调用
    structure_classify("../data/after_step5/S", "../data/after_step6/S")
    structure_classify("../data/after_step5/O", "../data/after_step6/O")