# 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
    rotate atoms so that the x,y,z axis are along the octahedra axises.
    if octahedra is a list of octahedras, use the average value.
    """
    if isinstance(octa,octahedra):
        x,y,z=octa.get_directions()
        rotate(atoms,x,[1,0,0],y,[0,1,0],rotate_cell=True)
    elif isinstance(octa, Iterable):
        xs=[]
        ys=[]
        for oc in octa:
            x,y,z=oc.get_directions()
            xs.append(np.array(x))
            ys.append(np.array(y))
        avg_x=np.average(xs,axis=0)
        avg_y=np.average(ys,axis=0)
        rotate(atoms,avg_x,[1,0,0],avg_y,[0,1,0],rotate_cell=True)
    else:
        raise ValueError("octa should be a octahedra or a iterable object ")
    return atoms

def rotate_atoms_with_octahedra_average(atoms,octahedras):
    """
    rotate atoms so that the x,y,z axis are along the average octahedra axises.
    """


def write_all_octahedra_info(atoms,center_symbol_list,vertex_symbol,max_distance,output='octa_info.csv',axis_type=None,x=(1,0,0),y=(0,1,0),z=(0,0,1), var_distance=False):
    """
    write all octahedra info into a file.
    """

# 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

    rotate atoms so that the x,y,z axis are along the octahedra axises.
    if octahedra is a list of octahedras, use the average value.
    """
    if isinstance(octa,octahedra):
        x,y,z=octa.get_directions()
        rotate(atoms,x,[1,0,0],y,[0,1,0],rotate_cell=True)
    elif isinstance(octa, Iterable):
        xs=[]
        ys=[]
        for oc in octa:
            x,y,z=oc.get_directions()
            xs.append(np.array(x))
            ys.append(np.array(y))
        avg_x=np.average(xs,axis=0)
        avg_y=np.average(ys,axis=0)
        rotate(atoms,avg_x,[1,0,0],avg_y,[0,1,0],rotate_cell=True)
    else:
        raise ValueError("octa should be a octahedra or a iterable object ")
    return atoms

def rotate_atoms_with_octahedra_average(atoms,octahedras):
    """
    rotate atoms so that the x,y,z axis are along the average octahedra axises.
    """


def write_all_octahedra_info(atoms,center_symbol_list,vertex_symbol,max_distance,output='octa_info.csv',axis_type=None,x=(1,0,0),y=(0,1,0),z=(0,0,1), var_distance=False):
    """
    write all octahedra info into a file.
    """

    symdict=symbol_number(atoms)