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
from pymatgen.io.cif import CifWriter

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)
            print(f"正在处理{filename}")
            temp.check_basic()
            if temp.check_basic_result:
                if not "+" in temp.anion:
                    target_folder = os.path.join("../data/after_step1",f"{temp.anion}")
                    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:
                        target_folder = os.path.join("../data/after_step1", f"{temp.anion}")
                        target_folder = os.path.join(target_folder, anion)
                        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}")
read_files_check_basic("../data/input")