import os
import csv
from pymatgen.core import Structure
from tqdm import tqdm  # 引入tqdm来显示进度条，如果未安装请运行 pip install tqdm

# --- 导入您自定义的分析函数 ---
# 假设您的函数存放在 'utils/CS_analyse.py' 文件中
# 并且您已经将它们重命名
# from calculate_polyhedra_sharing import calculate_polyhedra_sharing as CS_catulate
# from check_only_corner_sharing import check_only_corner_sharing
# 注意：根据您的描述，您会将函数放在 utils 文件夹中，因此导入方式如下：
from utils.CS_analyse import CS_catulate, check_only_corner_sharing


def process_cif_folder(cif_folder_path: str, output_csv_path: str):
    """
    遍历指定文件夹中的所有CIF文件，计算其角共享特性，并将结果输出到CSV文件。

    参数:
        cif_folder_path (str): 存放CIF文件的文件夹路径。
        output_csv_path (str): 输出的CSV文件的路径。
    """
    # 检查输入文件夹是否存在
    if not os.path.isdir(cif_folder_path):
        print(f"错误: 文件夹 '{cif_folder_path}' 不存在。")
        return

    # 准备存储结果的列表
    results = []

    # 获取所有CIF文件的列表
    try:
        cif_files = [f for f in os.listdir(cif_folder_path) if f.endswith('.cif')]
        if not cif_files:
            print(f"警告: 在文件夹 '{cif_folder_path}' 中没有找到任何 .cif 文件。")
            return
    except FileNotFoundError:
        print(f"错误: 无法访问文件夹 '{cif_folder_path}'。")
        return

    print(f"开始处理 {len(cif_files)} 个CIF文件...")

    # 使用tqdm创建进度条，遍历所有CIF文件
    for filename in tqdm(cif_files, desc="Processing CIFs"):
        file_path = os.path.join(cif_folder_path, filename)

        try:
            # 1. 从CIF文件加载结构
            struct = Structure.from_file(file_path)

            # 2. 调用您的 CS_catulate 函数计算详细的共享关系
            #    这里使用默认参数 sp='Li', anion=['O']
            sharing_details = CS_catulate(struct, sp='Li', anion=['O','S','Cl','F','Br'])

            # 3. 调用 check_only_corner_sharing 函数进行最终判断
            is_only_corner = check_only_corner_sharing(sharing_details)

            # 4. 将文件名和结果存入列表
            results.append([filename, is_only_corner])

        except Exception as e:
            # 如果处理某个文件时出错，打印错误信息并继续处理下一个文件
            print(f"\n处理文件 '{filename}' 时发生错误: {e}")
            results.append([filename, 'Error'])  # 在CSV中标记错误

    # 5. 将结果写入CSV文件
    try:
        with open(output_csv_path, 'w', newline='', encoding='utf-8') as csvfile:
            writer = csv.writer(csvfile)
            # 写入表头
            writer.writerow(['CIF_File', 'Is_Only_Corner_Sharing'])
            # 写入所有结果
            writer.writerows(results)

        print(f"\n处理完成！结果已保存到 '{output_csv_path}'")

    except IOError as e:
        print(f"\n错误: 无法写入CSV文件 '{output_csv_path}': {e}")


# --- 主程序入口 ---
if __name__ == "__main__":
    # ----- 参数配置 -----
    # 请将此路径修改为您存放CIF文件的文件夹的实际路径
    CIF_DIRECTORY = "data/0921"

    # 输出的CSV文件名
    OUTPUT_CSV = "corner_sharing_results.csv"
    # -------------------

    # 调用主函数开始处理
    process_cif_folder(CIF_DIRECTORY, OUTPUT_CSV)