# Path: pyDFTutils/ase_utils/symbol.py
# def get_symdict(filename='POSCAR',atoms=None):
#     """
#     get a symbol_number: index dict.
#     """
#     if atoms is not None:
#         syms=atoms.get_chemical_symbols()
#     elif filename is not None and atoms is None:
#         syms=read(filename).get_chemical_symbols()
# 
#     symdict=symbol_number(syms)
#     return symdict
# 
# def symbol_number(symbols):
#     """
#     symbols can be also atoms. Thus the chemical symbols will be used.
#     Fe Fe Fe O -> {Fe1:0 Fe2:1 Fe3:2 O1:3}
#     """
#     try:
#         symbs=symbols.copy().get_chemical_symbols()
#     except Exception:
#         symbs=symbols
#     symdict={}
#     result=OrderedDict()
#     for i,sym in enumerate(symbs):
#         if sym in symdict:
#             symdict[sym]=symdict[sym]+1
#         else:
#             symdict[sym]=1
#         result[sym+str(symdict[sym])]=i
#     return result
# 
# def symnum_to_sym(symbol_number):
#     """
#     symnum-> sym. eg: Fe1-> Fe
#     """
#     try:
#         a=re.search('[A-Za-z]+',symbol_number).group()
#         return a
#     except AttributeError:
#         raise AttributeError('%s is not a good symbol_number'%symbol_number)


from ase.utils.geometry import rotate
from ase.io import read
from ase.geometry.cell import cellpar_to_cell, cell_to_cellpar
from pyDFTutils.ase_utils.symbol import get_symdict,symbol_number,symnum_to_sym
from pyDFTutils.ase_utils.ase_utils import scaled_pos_to_pos,force_near_0,pos_to_scaled_pos
from collections import OrderedDict,Iterable
from ase.calculators.neighborlist import NeighborList
    from scipy.sparse import csr_matrix
    from scipy.sparse.csgraph import dijkstra
import numpy as np
    import argparse


    def get_vertex(self,direction):

        return self.__dict__[direction]


    def set_axis(self,x=(1,0,0), y=(0,1,0),z=(0,0,1),axis_type=None):
        """
        axis_type: default None, if not None: 'rotate45_xy': means x,y are 45 degree rotated about z
        """
        if axis_type is None:
            self.x=np.array(x)
            self.y=np.array(y)
            self.z=np.array(z)
        elif axis_type=='rotate45_xy':
            self.x=np.array([1,-1,0])
            self.y=np.array([1,1,0])
            self.z=np.array([0,0,1])
        else:
            raise ValueError('axis_type %s invalid.'%axis_type)

    def get_octahedra(self,atoms,center_atom,atom_symbol,max_distance,var_distance=False,do_force_near_0=True,repeat=True):
        """
        atoms: the atoms
        center_atom: a atom object, or sym_num
        atoms_symbol: symbol of atom in the vertexes
        max_distance: max distance
        var_distance: variable distance.
        """

# Path: pyDFTutils/ase_utils/symbol.py
# def get_symdict(filename='POSCAR',atoms=None):
#     """
#     get a symbol_number: index dict.
#     """
#     if atoms is not None:
#         syms=atoms.get_chemical_symbols()
#     elif filename is not None and atoms is None:
#         syms=read(filename).get_chemical_symbols()
# 
#     symdict=symbol_number(syms)
#     return symdict
# 
# def symbol_number(symbols):
#     """
#     symbols can be also atoms. Thus the chemical symbols will be used.
#     Fe Fe Fe O -> {Fe1:0 Fe2:1 Fe3:2 O1:3}
#     """
#     try:
#         symbs=symbols.copy().get_chemical_symbols()
#     except Exception:
#         symbs=symbols
#     symdict={}
#     result=OrderedDict()
#     for i,sym in enumerate(symbs):
#         if sym in symdict:
#             symdict[sym]=symdict[sym]+1
#         else:
#             symdict[sym]=1
#         result[sym+str(symdict[sym])]=i
#     return result
# 
# def symnum_to_sym(symbol_number):
#     """
#     symnum-> sym. eg: Fe1-> Fe
#     """
#     try:
#         a=re.search('[A-Za-z]+',symbol_number).group()
#         return a
#     except AttributeError:
#         raise AttributeError('%s is not a good symbol_number'%symbol_number)


# Path: pyDFTutils/perovskite/octahedra.py
from ase.utils.geometry import rotate
from ase.io import read
from ase.geometry.cell import cellpar_to_cell, cell_to_cellpar
from pyDFTutils.ase_utils.symbol import get_symdict,symbol_number,symnum_to_sym
from pyDFTutils.ase_utils.ase_utils import scaled_pos_to_pos,force_near_0,pos_to_scaled_pos
from collections import OrderedDict,Iterable
from ase.calculators.neighborlist import NeighborList
    from scipy.sparse import csr_matrix
    from scipy.sparse.csgraph import dijkstra
import numpy as np
    import argparse



    def get_vertex(self,direction):

        return self.__dict__[direction]


    def set_axis(self,x=(1,0,0), y=(0,1,0),z=(0,0,1),axis_type=None):
        """
        axis_type: default None, if not None: 'rotate45_xy': means x,y are 45 degree rotated about z
        """
        if axis_type is None:
            self.x=np.array(x)
            self.y=np.array(y)
            self.z=np.array(z)
        elif axis_type=='rotate45_xy':
            self.x=np.array([1,-1,0])
            self.y=np.array([1,1,0])
            self.z=np.array([0,0,1])
        else:
            raise ValueError('axis_type %s invalid.'%axis_type)

    def get_octahedra(self,atoms,center_atom,atom_symbol,max_distance,var_distance=False,do_force_near_0=True,repeat=True):
        """
        atoms: the atoms
        center_atom: a atom object, or sym_num
        atoms_symbol: symbol of atom in the vertexes
        max_distance: max distance
        var_distance: variable distance.
        """

        self.sym_dict=get_symdict(filename=None,atoms=atoms)