from pymatgen.core.structure import Structure
from pymatgen.analysis.local_env import VoronoiNN
import numpy as np

def check_real(nearest):
    real_nearest = []
    for site in nearest:
        if np.all((site.frac_coords >= 0) & (site.frac_coords <= 1)):
             real_nearest.append(site)

    return real_nearest

def special_check_for_3(site, nearest):
    real_nearest = []
    distances = []
    for site2 in nearest:
        distance = np.linalg.norm(np.array(site.frac_coords) - np.array(site2.frac_coords))
        distances.append(distance)

    sorted_indices = np.argsort(distances)
    for index in sorted_indices[:3]:
        real_nearest.append(nearest[index])

    return real_nearest


def CS_catulate(struct, sp='Li', anion=['O'], tol=0, cutoff=3.0,notice=False,ID=None):
    """
    计算结构中目标元素与最近阴离子的共享关系。

    参数:
        struct (Structure): 输入结构。
        sp (str): 目标元素符号，默认为 'Li'。
        anion (list): 阴离子列表，默认为 ['O']。
        tol (float): VoronoiNN 的容差，默认为 0。
        cutoff (float): VoronoiNN 的截断距离，默认为 3.0。

    返回:
        list: 包含每个目标位点及其最近阴离子索引的列表。
    """
    # 初始化 VoronoiNN 对象
    if sp=='Li':
        tol = 0
        cutoff = 3.0
    voro_nn = VoronoiNN(tol=tol, cutoff=cutoff)
    # 初始化字典，用于统计共享关系
    shared_count = {"2": 0, "3": 0,"4":0,"5":0,"6":0}
    # 存储结果的列表
    atom_dice = []

    # 遍历结构中的每个位点
    for index,site in enumerate(struct.sites):
        # 跳过阴离子位点
        if site.specie.symbol in anion:
            continue
        # 跳过Li原子
        if site.specie.symbol == sp:
            continue
        # 获取 Voronoi 多面体信息
        voro_info = voro_nn.get_voronoi_polyhedra(struct, index)

        # 找到最近的阴离子位点
        nearest_anions = [
            nn_info["site"] for nn_info in voro_info.values()
            if nn_info["site"].specie.symbol in anion
        ]

        # 如果没有找到最近的阴离子，跳过
        if not nearest_anions:
            print(f"No nearest anions found for {ID} site {index}.")
            continue
        if site.specie.symbol == 'B' or site.specie.symbol == 'N':
            nearest_anions = special_check_for_3(site,nearest_anions)
        nearest_anions = check_real(nearest_anions)
        # 将结果添加到 atom_dice 列表中
        atom_dice.append({
            'index': index,
            'nearest_index': [nn.index for nn in nearest_anions]
        })





        # 枚举 atom_dice 中的所有原子对
    for i, atom_i in enumerate(atom_dice):
        for j, atom_j in enumerate(atom_dice[i + 1:], start=i + 1):
            # 获取两个原子的最近阴离子索引
            nearest_i = set(atom_i['nearest_index'])
            nearest_j = set(atom_j['nearest_index'])

            # 比较最近阴离子的交集大小
            shared_count_key = str(len(nearest_i & nearest_j))

            # 更新字典中的计数
            if shared_count_key in shared_count:
                shared_count[shared_count_key] += 1
                if notice:
                    if shared_count_key=='2':
                        print(f"{atom_j['index']}与{atom_i['index']}之间存在共线")
                        print(f"共线的阴离子为{nearest_i & nearest_j}")
                    if shared_count_key=='3':
                        print(f"{atom_j['index']}与{atom_i['index']}之间存在共面")
                        print(f"共面的阴离子为{nearest_i & nearest_j}")

    # # 最后将字典中的值除以 2，因为每个共享关系被计算了两次
    # for key in shared_count.keys():
    #     shared_count[key] //= 2

    return shared_count


def CS_count(struct, shared_count, sp='Li'):
    count = 0
    for site in struct.sites:
        if site.specie.symbol == sp:
            count += 1  # 累加符合条件的原子数量

    CS_count = 0
    for i in range(2, 7):  # 遍历范围 [2, 3, 4, 5]
        if str(i) in shared_count:  # 检查键是否存在
            CS_count += shared_count[str(i)] * i  # 累加计算结果

    if count > 0:  # 防止除以零
        CS_count /= count  # 平均化结果
    else:
        CS_count = 0  # 如果 count 为 0，直接返回 0

    return CS_count

structure = Structure.from_file("data/CS_Table1/Li3Al(BO3)2_mp-6097_computed.cif")
a = CS_catulate(structure,notice=True)
b = CS_count(structure,a)
print(f"{a}\n{b}")