import os
import shutil
from pymatgen.core import Structure


def get_anion_type(structure):
    """
    简单判断阴离子类型。
    返回: 'O', 'S', 'S+O' 等字符串
    """
    # 定义常见的阴离子列表
    valid_anions = {'O', 'S', 'Se', 'Te', 'F', 'Cl', 'Br', 'I', 'N', 'P'}

    # 获取结构中的所有元素符号
    elements = set([e.symbol for e in structure.composition.elements])

    # 取交集找到当前结构包含的阴离子
    found_anions = elements.intersection(valid_anions)

    if not found_anions:
        return "Unknown"

    # 如果有多个阴离子，按字母顺序排序并用 '+' 连接 (模拟 step1 的逻辑)
    sorted_anions = sorted(list(found_anions))
    return "+".join(sorted_anions)


def organize_files_direct(input_folder, output_base):
    if not os.path.exists(input_folder):
        print(f"输入文件夹不存在: {input_folder}")
        return

    cif_files = [f for f in os.listdir(input_folder) if f.endswith(".cif")]
    print(f"发现 {len(cif_files)} 个 CIF 文件，开始直接整理...")

    for filename in cif_files:
        file_path = os.path.join(input_folder, filename)

        try:
            # 仅读取结构用于分类
            struct = Structure.from_file(file_path)
            anion_type = get_anion_type(struct)

            # 获取不带后缀的文件名 (ID)
            file_base_name = os.path.splitext(filename)[0]

            # --- 构建目标路径逻辑 ---
            # 逻辑：../data/after_step1 / 阴离子类别 / ID / ID.cif

            # 处理混合阴离子情况 (如 S+O)
            if "+" in anion_type:
                # 按照之前的逻辑，如果是混合阴离子，通常会有多层
                # 但为了统一后续处理，我们这里将其放入组合名的文件夹下
                # 比如: after_step1/S+O/S/123/123.cif (复杂)
                # 或者简化为 after_step1/S+O/123/123.cif (简单)
                # 根据你之前的 make_sh.py 和 extract_data.py，
                # 只要是 Folder/ID/ID.cif 结构即可。
                # 为了兼容 analyze_cs.py 的逻辑 (group_name, anion_name)，
                # 这里我们采用 simplified 逻辑：
                # 如果是混合，我们在第一层建 S+O，第二层建具体的 anion 文件夹(比如首字母排序第一个)
                # 或者直接: after_step1/S+O/ID/ID.cif -> 这样 group=S+O, anion=ID (不对)

                # 兼容旧代码的最佳实践：
                # 对于混合 S+O，我们建立 S+O/S/ID/ID.cif 和 S+O/O/ID/ID.cif ?
                # 不，原 Step1 是把一个文件复制了两份到不同文件夹。
                # 这里为了简化，我们只复制一份到主阴离子文件夹，或者直接按组合命名。

                # 让我们采用最稳妥的方式：如果是 S+O，放入 S+O/Mix/ID/ID.cif
                # 这样 group=S+O, anion=Mix。
                # 但为了让 CS_calc 正常工作，最好还是放入具体的元素文件夹。
                # 这里我们简单处理：直接放入 S+O/Combined/ID/
                # 或者根据你的 extract_data.py 逻辑：
                # 它会遍历 top_dir (S+O) -> sub_anion (S, O)

                # 策略：拆分放入。
                sub_anions = anion_type.split("+")
                for sub in sub_anions:
                    target_folder = os.path.join(output_base, anion_type, sub, file_base_name)
                    if not os.path.exists(target_folder):
                        os.makedirs(target_folder)
                    shutil.copy(file_path, os.path.join(target_folder, filename))

                print(f"整理: {filename} -> {anion_type} (已复制到各子类)")

            else:
                # 单一阴离子: after_step1/S/ID/ID.cif
                target_folder = os.path.join(output_base, anion_type, file_base_name)
                if not os.path.exists(target_folder):
                    os.makedirs(target_folder)

                shutil.copy(file_path, os.path.join(target_folder, filename))
                print(f"整理: {filename} -> {anion_type}")

        except Exception as e:
            print(f"处理 {filename} 失败: {e}")


if __name__ == "__main__":
    organize_files_direct("../data/input", "../data/after_step1")