from pymatgen.core import Structure
from pymatgen.io.vasp import Poscar

def create_supercell_poscar(cif_path, supercell_matrix, output_filename="POSCAR_supercell"):
    """
    从CIF文件读取晶体结构，根据指定的矩阵进行扩胞，并生成VASP POSCAR文件。

    Args:
        cif_path (str): 输入的CIF文件路径。
        supercell_matrix (list or tuple): 3x3的扩胞矩阵。
            - 对于简单的对角扩胞 (例如 2x2x4)，使用: [[2, 0, 0], [0, 2, 0], [0, 0, 4]]
            - 对于非对角扩胞 (例如 a_s=3a, b_s=2a+4b, c_s=6c)，使用: [[3, 0, 0], [2, 4, 0], [0, 0, 6]]
        output_filename (str): 输出的POSCAR文件名。默认为 "POSCAR_supercell"。

    Returns:
        bool: 如果成功生成文件则返回 True，否则返回 False。
    """
    try:
        # 1. 从CIF文件加载结构
        # 使用 from_file 静态方法直接读取
        # primitive=False 确保我们使用CIF中定义的晶胞，而不是其原胞
        original_structure = Structure.from_file(cif_path, primitive=False)

        print("--- 原始晶胞信息 ---")
        print(f"  原子数: {original_structure.num_sites}")
        print(f"  化学式: {original_structure.composition.reduced_formula}")
        print(f"  晶格参数 (a, b, c, α, β, γ):")
        lat = original_structure.lattice
        print(f"  {lat.a:.4f}, {lat.b:.4f}, {lat.c:.4f}, {lat.alpha:.2f}, {lat.beta:.2f}, {lat.gamma:.2f}")

        # 2. 进行扩胞操作
        # 注意：pymatgen 会自动处理原子坐标的映射
        supercell_structure = original_structure * supercell_matrix

        print("\n--- 扩胞后信息 ---")
        print(f"  扩胞矩阵: {supercell_matrix}")
        print(f"  新原子数: {supercell_structure.num_sites}")
        print(f"  新化学式: {supercell_structure.composition.reduced_formula}")
        print(f"  新晶格参数 (a, b, c, α, β, γ):")
        super_lat = supercell_structure.lattice
        print(f"  {super_lat.a:.4f}, {super_lat.b:.4f}, {super_lat.c:.4f}, {super_lat.alpha:.2f}, {super_lat.beta:.2f}, {super_lat.gamma:.2f}")

        # 3. 创建Poscar对象并写入文件
        # comment 参数可以设置POSCAR文件的第一行注释
        poscar = Poscar(supercell_structure, comment=f"Supercell from {cif_path}")
        poscar.write_file(output_filename)

        print(f"\n成功！已将扩胞结构写入文件: {output_filename}")
        return True

    except Exception as e:
        print(f"发生错误: {e}")
        return False

# --- 使用示例 ---

# 假设您的CIF文件名为 "origin.cif"，并且与此脚本在同一目录下。
# 如果您在复现 Wang Shuo 的工作，他们可能使用了不同的扩胞方案。
# 例如，用于MD模拟的大超胞是 2x2x4 [source_id: 3]。
# 而用于DP-GEN探索的小超胞是 1x1x2 [source_id: 3]。

# 示例1：生成用于DP-GEN探索的 1x1x2 小超胞 (60个原子)
print("="*40)
print("正在生成 1x1x2 超胞 (用于 DP-GEN 探索)...")
matrix_1x1x2 = [[1, 0, 0], [0, 1, 0], [0, 0, 2]]
create_supercell_poscar("data/P3ma/origin.cif", matrix_1x1x2, "data/P3ma/output/POSCAR_1x1x2_60atoms")

# 示例2：生成用于LAMMPS MD模拟的 2x2x4 大超胞 (480个原子)
# print("\n" + "="*40)
# print("正在生成 2x2x4 超胞 (用于 LAMMPS MD 模拟)...")
# matrix_2x2x4 = [[2, 0, 0], [0, 2, 0], [0, 0, 4]]
# create_supercell_poscar("origin.cif", matrix_2x2x4, "POSCAR_2x2x4_480atoms")
#
# # 示例3：生成 Geng 等人研究中使用的大超胞 (2160个原子) [source_id: 1]
# # 这个扩胞矩阵 a_s = 3a_0, b_s = 2a_0 + 4b_0, c_s = 6c_0 [source_id: 1]
# print("\n" + "="*40)
# print("正在生成非对角扩胞超胞 (Geng et al.)...")
# matrix_geng = [[3, 0, 0], [2, 4, 0], [0, 0, 6]]
# create_supercell_poscar("origin.cif", matrix_geng, "POSCAR_Geng_2160atoms")