from pymatgen.core import Structure
from pymatgen.core.periodic_table import Element, Specie
from pymatgen.analysis.local_env import CrystalNN
from pymatgen.analysis.structure_matcher import StructureMatcher
from pymatgen.io.cif import CifParser
import numpy as np
class crystal:
    def __init__(self, file_path, element_positive='Na', mixed_anions=None):
        # self.parse = CifParser(file_path)
        # self.structure = self.parse.get_structures()[0]
        if mixed_anions is None:
            mixed_anions = {frozenset({'S', 'O'}), frozenset({'Cl', 'Br'}),frozenset({'Cl', 'O'}),frozenset({'Cl', 'Br'}),frozenset({'S', 'Cl'})}
        self.structure = Structure.from_file(file_path)
        self.file_path = file_path
        self.element_positive = element_positive
        self.check_all = False
        self.check_basic_result = False
        self.check_high_cn_and_face_sharing_result = False
        self.check_percolation_radius_result = False
        self.check_practical_result = False
        self.anion = ""
        self.anions = ""
        self.mixed_anions = mixed_anions
        #self.initialize()
        
    def initialize(self):
        print("e")
        # self.check_basic_result=self.check_basic()
        # self.check_high_cn_and_face_sharing_result = self.check_high_cn_and_face_sharing()
        # self.check_percolation_radius_result = self.check_percolation_radius()
        # self.check_all = self.check_basic_result and self.check_high_cn_and_face_sharing_result and self.check_percolation_radius_result
        # print(f"{self.file_path}done")

    def check_practical(self):
        structure = self.structure

        # 检查是否为Li-X-O,Li-P-S
        excluded_X_elements = {'S', 'I', 'Si', 'C', 'P', 'Al', 'Ge', 'Se', 'B', 'Cl'}
        chemical_system_set = structure.chemical_system_set

        try:
            if len(chemical_system_set) == 3:
                if "Li" in chemical_system_set and "O" in chemical_system_set:
                    for element in excluded_X_elements:
                        if element in chemical_system_set:
                            return False
                if "Li" in chemical_system_set and "P" in chemical_system_set and "S" in chemical_system_set:
                    return False
        except Exception as e:
            print(f"Error during Li-X-O check: {e}")
            return False

        # 排除过渡金属元素
        excluded_transition_metals = {'Fe', 'Mn', 'Ni', 'Ti', 'Mo', 'V', 'Co'}
        try:
            if "Li" in chemical_system_set and "O" in chemical_system_set:
                for element in excluded_transition_metals:
                    if element in chemical_system_set:
                        return False
        except Exception as e:
            print(f"Error during transition metal check: {e}")
            return False

        # 检查是否包含N, Re, Ho, Hf, Ru, Eu, Lu
        excluded_elements = {'N', 'Re', 'Ho', 'Hf', 'Ru', 'Lu'}
        try:
            for element in excluded_elements:
                if element in chemical_system_set:
                    return False
        except Exception as e:
            print(f"Error during excluded elements check: {e}")
            return False

        # 检查是否共享位点
        try:
            for site in structure.sites:
                if 'Li' in site.species_string and len(site.species) > 1:
                    return False
        except Exception as e:
            print(f"Error during site sharing check: {e}")
            return False

        self.check_practical_result = True
        return True

    def check_basic(self):
        structure = self.structure
        #判断是否为二元化合物
        if len(structure.types_of_specie) == 2:
            return False
        #判断阴离子是否为多种阴离子
        # anions = {'O', 'S', 'Se', 'Te', 'F', 'Cl', 'Br', 'I'}
        anions = {'O', 'S','Br','Cl'}
        try:
            for site in self.structure.sites:
                try:
                    #if site.specie.symbol in anions:
                    if site.species.chemical_system in anions:
                        self.anion = site.specie.symbol
                        break
                except AttributeError as e:
                    a=1
                try:
                    if site.species.chemical_system in anions:
                        self.anion = site.specie.symbol
                        break
                except AttributeError as e:
                    print(e)
            if self.anion in anions:
                a=1
            else:
                if not self.mixed_anions:
                    print("不是所选阴离子化合物")
                    return False
        except Exception as e:
            print(e)
            return False
        #这里添加对多种阴离子的支持
        try:
            # 创建一个集合来收集所有发现的阴离子
            found_anions = set()

            # 遍历structure以收集所有阴离子
            for site in self.structure.sites:
                try:
                    if site.species.chemical_system in anions:
                        found_anions.add(site.specie.symbol)
                except AttributeError:
                    try:
                        if site.specie.symbol in anions:
                            found_anions.add(site.specie.symbol)
                    except AttributeError:
                        continue

            # 检查找到的阴离子情况
            if len(found_anions) == 0:
                print("未找到任何预定义的阴离子")
                return False
            elif len(found_anions) == 1:
                # 只有一种阴离子
                self.anion = list(found_anions)[0]
                print(f"发现单一阴离子: {self.anion}")
            else:
                # 有多种阴离子，检查是否匹配预定义的混合阴离子组合
                found_anions_frozen = frozenset(found_anions)
                if found_anions_frozen in self.mixed_anions:
                    self.anions = found_anions
                    self.anion = "+".join(sorted(found_anions))  # 例如: "Cl+S"
                    print(f"发现匹配的混合阴离子组合: {self.anion}")
                else:
                    # 如果找到的阴离子组合不在预定义列表中
                    print(f"发现的阴离子组合 {found_anions} 不在预定义的混合阴离子列表中")
                    return False
        except Exception as e:
            print(f"处理阴离子时出错: {e}")
            return False

        #这里还要调试
        # try:
        #     # 初始化总电荷
        #     total_charge = 0
        #
        #     # 检查是否所有元素都有氧化态
        #     for site in structure:
        #         try:
        #             oxi_state = site.species.charge  # 检查是否有氧化态
        #             total_charge += oxi_state  # 累加氧化态
        #         except AttributeError:
        #             print(f"元素 {site.specie.symbol} 缺少氧化态定义")
        #             return False
        #     # 检查是否电荷平衡
        #     if total_charge == 0:
        #         print("所有元素的价态之和为 0，结构电荷平衡")
        #     else:
        #         print(f"所有元素的价态之和为 {total_charge}，结构不平衡")
        #         return False
        # except Exception as e:
        #     print(f"发生错误: {e}")
        #     return False

        #判断原子个数
        try:
            if not self.mixed_anions:
                if structure.num_sites>300:
                    return False
            else:
                if structure.num_sites>900:
                    return False
        except Exception:
            print("原子个数判断失败")
            return False

        #判断有几个阴离子
        # anions = {'O', 'S', 'Se', 'Te', 'F', 'Cl', 'Br', 'I'}
        # try:
        #     anion_elements = {site.species.chemical_system for site in structure if site.species.chemical_system in anions}
        #     if len(anion_elements) > 1:
        #         return False
        # except Exception:
        #     print("阴离子个数判断失败")
        #     return False

        #判断是否有放射性元素
        radioactive_elements = {'U', 'Th', 'Pu', 'Ra', 'Rn', 'Po', 'Np', 'Am', 'Cm', 'Bk', 'Cf', 'Es', 'Fm', 'Md', 'No',
                                'Lr'}

        try:
            # 遍历结构中的元素
            for site in structure:
                if site.species.chemical_system in radioactive_elements:
                    return False  # 存在放射性元素
        except Exception:
            print("放射性元素判断失败")
            return False


        #判断是否存在共占位
        try:
            for site in structure.sites:
                if self.element_positive in [specie.symbol for specie in site.species.keys()] and len(site.species) > 1:
                    return False
        except Exception:
            print("共占位判断失败")
            return False

        #判读是否有无序或部分占位的阴离子
        try:
            for site in structure.sites:
                for specie, occupancy in site.species.items():
                    if specie.symbol in anions and occupancy < 1:
                        return False
        except Exception:
            print("无序或部分占位的阴离子判断失败")
            return False
                
        #判断是否有水分子
        try:
            oxygen_sites = [site for site in structure.sites if site.species.chemical_system == "O"]
            hydrogen_sites = [site for site in structure.sites if site.species.chemical_system == "H"]
            
            for o_site in oxygen_sites:
                nearby_hydrogens = [h_site for h_site in hydrogen_sites if o_site.distance(h_site) < 1.2]
                
                if len(nearby_hydrogens) == 2:
                    return False
        except Exception:
            print("水分子判断失败")
            return False
        #接下来判断是否有标准信息
        try:
            for site in structure.sites:
                for specie in site.species.keys():
                    element = Element(specie.symbol)
                    
                    if not element.common_oxidation_states:
                        return False
                    
                    try:
                        _ = Specie(element.symbol, max(element.common_oxidation_states)).ionic_radius
                    except:
                        return False
        except Exception:
            print("标准信息判断失败")
            return False
        #暂时不判断是否为电中性
        #可能需要通过ovito等库来做判断
        #存在一些文件不提供各元素的电负性
        
        #判断电中性是否存在

        self.check_basic_result = True
        return True
    
    def check_high_cn_and_face_sharing(self,cut_distance = 3.1):
        structure = self.structure
        #基于固角权重的计算
        nn_finder = CrystalNN()
        
        #遍历结构中的所有Na位点，检查配位数
        #是所有Na位点都需要还是只检查高配位数的位点？
        high_cn_ep_sites = []
        try:
            for i,site in enumerate(structure):
                if site.specie == Element(self.element_positive):
                    cn = nn_finder.get_cn(structure,i)
                    if cn>=5:
                        high_cn_ep_sites.append(i)
            if len(high_cn_ep_sites)==0:
                return False
        except Exception:
            print("高配位Na离子判断失败")
            return False
        #检查共面
        try:
            for i in high_cn_ep_sites:
                neighbors = nn_finder.get_nn_info(structure,i)
                x_neighbors = []
                x_neighbors = [
                    neighbor["site_index"]
                    for neighbor in neighbors
                    if structure[neighbor["site_index"]].specie.symbol == self.element_positive
                    and neighbor["site"].distance(structure[i]) <= cut_distance
                ]
                
                if not self._check_face_sharing(i, x_neighbors):
                    print(f"Na site {i} does not share a face with other high-CN Na sites.")
                    return False

            print("All high-CN Na sites are face-sharing.")
        except Exception:
            print("共面判断失败")
        return True
                    
                    
    def _check_face_sharing(self,site_index,neighbor_indices):
        # 获取当前 Na 位点的坐标
        site_coords = self.structure[site_index].coords

        # 遍历邻居
        for neighbor_index in neighbor_indices:
            # 获取邻居的坐标
            neighbor_coords = self.structure[neighbor_index].coords

            # 获取两个原子之间共享的面（使用简单的距离或角度计算）
            # 假设共享面的法向量计算可以从 Voronoi 构造
            shared_face_normal = self._calculate_face_normal(site_coords, neighbor_coords)

            # 判断是否共面（如果法向量的绝对值接近 0，可以认为共面）
            if shared_face_normal is not None:
                return True

        return False
    def _calculate_face_normal(self, coords1, coords2):

        # 示例计算：用两个原子之间的向量生成法向量
        vector = coords2 - coords1
        norm = np.linalg.norm(vector)

        # 如果向量接近零，返回 None
        if norm < 1e-6:
            return None

        # 正则化向量作为法向量
        return vector / norm
    
    def check_percolation_radius(self):
        return True
def group_structures_by_framework(structures):
    matcher = StructureMatcher()
    grouped_structures = []

    for structure in structures:
        matched = False
        for group in grouped_structures:

            if matcher.fit(structure, group[0]):  # 比较结构是否匹配
                group.append(structure)
                matched = True
                break
        if not matched:
            grouped_structures.append([structure])

    return grouped_structures