import os
from pymatgen.core import Structure
from pymatgen.core.periodic_table import Element
# 导入你的CSM计算工具库 (根据 provided context [11])
try:
    from utils.analyze_env_st import extract_sites, export_envs
except ImportError:
    print("Error: 找不到 utils.analyze_env_st 模块，请检查 utils 文件夹。")
    exit()

from tqdm import tqdm

# ================= 配置区域 =================
# 输入目录：使用筛选后的目录，只计算符合要求的材料
INPUT_DIR = "../../solidstate-tools/corner-sharing/data/1209/input"
# 输出目录
OUTPUT_DIR = "../output/CSM"
# 分析参数
TARGET_ELEMENT = 'Na'
ENV_TYPE = 'both'  # 可选 'tet', 'oct', 'both'
# ===========================================

def run_csm_analysis():
    """
    遍历 after_screening 文件夹，计算 CSM 并生成 .dat 文件到 output/CSM
    """
    if not os.path.exists(INPUT_DIR):
        print(f"输入目录不存在: {INPUT_DIR}，请先运行筛选步骤。")
        return

    # 收集所有需要处理的 CIF 文件
    cif_files = []
    for root, dirs, files in os.walk(INPUT_DIR):
        for file in files:
            if file.endswith(".cif"):
                # 保存完整路径
                cif_files.append(os.path.join(root, file))

    print(f"开始进行 CSM 分析，共找到 {len(cif_files)} 个筛选后的材料...")

    for cif_path in tqdm(cif_files, desc="Calculating CSM"):
        try:
            # 1. 确定输出路径，保持目录结构
            # 获取相对路径 (例如: S/195819.cif 或 S+O/S/195819.cif)
            rel_path = os.path.relpath(cif_path, INPUT_DIR)
            # 获取所在文件夹 (例如: S 或 S+O/S)
            rel_dir = os.path.dirname(rel_path)
            # 获取文件名 (例如: 195819)
            file_base = os.path.splitext(os.path.basename(cif_path))[0]

            # 构建目标文件夹: ../output/CSM/S/
            target_dir = os.path.join(OUTPUT_DIR, rel_dir)
            if not os.path.exists(target_dir):
                os.makedirs(target_dir)

            # 构建目标文件路径: ../output/CSM/S/195819.dat
            target_dat_path = os.path.join(target_dir, f"{file_base}.dat")

            # 2. 如果已经存在，跳过 (可选，视需求而定，这里默认覆盖)
            # if os.path.exists(target_dat_path):
            #     continue

            # 3. 读取结构
            struct = Structure.from_file(cif_path)

            # 检查是否包含目标元素 (Li)
            if Element(TARGET_ELEMENT) not in struct.composition.elements:
                # print(f"Skipping {file_base}: No {TARGET_ELEMENT}")
                continue

            # 4. 计算 CSM (引用 utils 中的函数)
            # extract_sites 返回环境列表
            env_list = extract_sites(struct, sp=TARGET_ELEMENT, envtype=ENV_TYPE)

            # 5. 导出结果 (引用 utils 中的函数)
            # export_envs 将结果写入 .dat 文件
            if env_list:
                export_envs(env_list, sp=TARGET_ELEMENT, envtype=ENV_TYPE, fname=target_dat_path)
            else:
                # 如果没有提取到环境（例如没有配位环境），生成一个空文件或记录日志
                with open(target_dat_path, 'w') as f:
                    f.write("No environments found.")

        except Exception as e:
            print(f"处理出错 {cif_path}: {e}")

    print(f"CSM 分析完成，结果已保存至 {OUTPUT_DIR}")

if __name__ == "__main__":
    run_csm_analysis()