import pymatgen.core as mg
from pymatgen.io.cif import CifParser
from pymatgen.transformations.standard_transformations import SupercellTransformation
import random
import os


def create_ordered_pnma_structure(disordered_structure):
    """
    手动将Pnma相的无序结构（主要为Li的部分占位）转换为有序结构。
    """
    s = disordered_structure.copy()

    # 根据origin.cif, Li位点的占位率为0.75 [5]
    partial_li_indices = [i for i, site in enumerate(s.sites) if "Li" in site.species and not site.is_ordered]

    # 根据0.75的占位率随机选择要保留的Li原子
    num_to_keep = int(round(len(partial_li_indices) * 0.75))
    keep_indices = set(random.sample(partial_li_indices, num_to_keep))

    # 找出需要删除的原子索引
    to_remove_indices = [i for i in partial_li_indices if i not in keep_indices]

    s.remove_sites(sorted(to_remove_indices, reverse=True))

    # 重新创建一个新的、完全有序的结构，避免任何副作用
    ordered_species = []
    ordered_coords = []
    for site in s.sites:
        # 只取每个位点的主要元素
        main_specie = site.species.elements[0]
        ordered_species.append(main_specie)
        ordered_coords.append(site.frac_coords)

    final_structure = mg.Structure(s.lattice, ordered_species, ordered_coords)

    return final_structure


def create_multiple_pnma_supercells(cif_path, num_configs=3, output_path="."):
    """
    读取Pnma相CIF，为不同尺寸的超胞生成多个具有不同反位缺陷位置的构型。
    """
    if not os.path.exists(cif_path):
        print(f"错误: 文件 '{cif_path}' 不存在。")
        return

    print(f"正在从 {cif_path} 读取Pnma结构...")
    parser = CifParser(cif_path)
    disordered_structure = parser.parse_structures(primitive=False)[0]

    structure = create_ordered_pnma_structure(disordered_structure)
    print(f"成功将无序Pnma结构转换为一个包含 {len(structure)} 个原子的有序单胞。")

    os.makedirs(output_path, exist_ok=True)

    target_sizes = [60, 90]
    for size in target_sizes:
        print(f"\n--- 正在为约 {size} 原子的版本生成 {num_configs} 个不同构型 ---")

        # 1. 构建基准超胞
        if size == 60:
            tf = SupercellTransformation([[1, 0, 0], [0, 1, 0], [0, 0, 2]])
            filename_suffix = "60_approx"
        else:  # size == 90
            tf = SupercellTransformation([[1, 0, 0], [0, 1, 0], [0, 0, 3]])
            filename_suffix = "90_approx"

        base_supercell = tf.apply_transformation(structure)
        print(f"已生成基准超胞，实际原子数: {len(base_supercell)}")

        li_indices = [i for i, site in enumerate(base_supercell.sites) if site.species_string == 'Li']
        y_indices = [i for i, site in enumerate(base_supercell.sites) if site.species_string == 'Y']

        if not li_indices or not y_indices:
            print("错误：在超胞中未找到足够的Li或Y原子来引入缺陷。")
            continue

        # 2. 循环生成多个独特的缺陷构型
        used_pairs = set()
        for i in range(num_configs):
            defect_supercell = base_supercell.copy()

            # 确保随机选择的交换对是全新的
            while True:
                li_swap_idx = random.choice(li_indices)
                y_swap_idx = random.choice(y_indices)
                # 使用排序后的元组作为键，确保(a,b)和(b,a)被视为相同
                pair = tuple(sorted((li_swap_idx, y_swap_idx)))
                if pair not in used_pairs:
                    used_pairs.add(pair)
                    break

            # 引入缺陷
            defect_supercell.replace(li_swap_idx, "Y")
            defect_supercell.replace(y_swap_idx, "Li")

            print(f"  配置 {i}: 成功引入一对反位缺陷 (Li at index {li_swap_idx} <-> Y at index {y_swap_idx})。")

            # 3. 保存为带编号的POSCAR文件
            poscar_filename = f"POSCAR_Pnma_{filename_suffix}_antisite_defect_{i}"
            poscar_path = os.path.join(output_path, poscar_filename)
            defect_supercell.to(fmt="poscar", filename=poscar_path)
            print(f"  已保存文件: {poscar_path}")


if __name__ == '__main__':
    # 请将您的Pnma相CIF文件保存，并修改此路径
    # 这里我们使用您提供的参考文件名 'origin.cif'
    cif_file_path = "data/Pnma/origin.cif"
    output_directory = "raw/Pnma/output"

    create_multiple_pnma_supercells(cif_file_path, num_configs=3, output_path=output_directory)

    print("\nPnma相处理完成！")