from pymatgen.core import Structure
from pymatgen.core.periodic_table import Element, Specie
from pymatgen.io.cif import CifWriter
from crystal_2 import crystal
import crystal_2
import os
import shutil


def read_files_check_basic(folder_path):
    file_contents = []

    if not os.path.exists(folder_path):
        print(f"{folder_path} 文件夹不存在")
        return file_contents

    for filename in os.listdir(folder_path):
        file_path = os.path.join(folder_path, filename)

        if os.path.isfile(file_path):
            try:
                temp = crystal(file_path)
                file_contents.append(temp)
            except Exception as e:
                print(e)
                continue  # 如果出错跳过当前循环，避免temp未定义报错

            print(f"正在处理{filename}")
            temp.check_basic()

            if temp.check_basic_result:
                # 获取不带后缀的文件名，用于创建同名文件夹
                file_base_name = os.path.splitext(filename)[0]

                if not "+" in temp.anion:
                    # 单一阴离子情况
                    # 路径变为: ../data/after_step1/Anion/FileBaseName/
                    base_anion_folder = os.path.join("../data/after_step1", f"{temp.anion}")
                    target_folder = os.path.join(base_anion_folder, file_base_name)

                    if not os.path.exists(target_folder):
                        os.makedirs(target_folder)

                    # 目标文件路径
                    target_file_path = os.path.join(target_folder, filename)
                    # 复制文件到目标文件夹
                    shutil.copy(file_path, target_file_path)
                    print(f"文件 {filename}通过基本筛选,已复制到 {target_folder}")
                else:
                    # 混合阴离子情况
                    anions = temp.anion.split("+")
                    for anion in anions:
                        # 路径变为: ../data/after_step1/AnionCombination/Anion/FileBaseName/
                        base_group_folder = os.path.join("../data/after_step1", f"{temp.anion}")
                        base_anion_folder = os.path.join(base_group_folder, anion)
                        target_folder = os.path.join(base_anion_folder, file_base_name)

                        if not os.path.exists(target_folder):
                            os.makedirs(target_folder)

                        # 目标文件路径
                        target_file_path = os.path.join(target_folder, filename)
                        # 复制文件到目标文件夹
                        shutil.copy(file_path, target_file_path)
                        print(f"文件 {filename}通过基本筛选,已复制到 {target_folder}")


if __name__ == "__main__":
    read_files_check_basic("../data/input")