import logging
logger = logging.getLogger('piou.config.lammps')
from piou import config
from piou.config import AtomConfigGenerator

class LAMMPS_input(AtomConfigGenerator):
    coord_file = None
    valid = True
    def __init__(self,coord_file=None):
        if coord_file is None:
            return
        self.coord_file = coord_file

    def supports_frame(self):
        return False

    def _gen_atomic_configuration(self):
        pass

    def supports_import(self):
        return False

    def export_atomic_configuration(self,atomic_coordinates,to_file=None,frame_no=None,all_frames=True):
        import time
        atomic_coordinate = atomic_coordinates[0]
        natm = atomic_coordinate.get_num_atom()
        elist = atomic_coordinate.get_elem_list()
        file = open(to_file,'w')
        file.write('LAMMPS input file generated by the PHASE I/O utility\n');
        file.write(str('# created on : '+time.strftime('%x %X'))+'\n')
        file.write('\n')
        file.write(str(natm)+' atoms\n')
        file.write(str(len(elist))+' atom types\n')
        unit_cell = atomic_coordinate.get_unit_cell()
        if atomic_coordinate.is_cell_ready():
            (a,b,c,alph,bet,gam) = unit_cell.get_lattice_constant() 
            unit_cell.set_lattice_constant(a,b,c,alph,bet,gam)
            (avec,bvec,cvec) = unit_cell.get_lattice_vector(to_unit=config.ANGSTROM)
            file.write('0 '+str(avec[0])+' xlo xhi\n')
            file.write('0 '+str(bvec[1])+' ylo yhi\n')
            file.write('0 '+str(cvec[2])+' zlo zhi\n')
            file.write(str(bvec[0])+' '+str(cvec[0])+' '+str(cvec[1])+' xy xz yz\n')
        file.write('\n')
        file.write('Atoms\n')
        file.write("\n")
        for i in range(natm):
            atomID=i+1
            atom = atomic_coordinate.get_atom_at(i)
            cpos = atom.get_pos(mode=config.CARTESIAN, to_unit=config.ANGSTROM)
            elem = atom.get_element_name()
            eid = -1
            for j in range(len(elist)):
                if elem==elist[j]:
                    eid = j+1
                    break
            file.write(str(atomID)+' '+str(eid)+' '+str(cpos[0])+' '+str(cpos[1])+' '+str(cpos[2])+'\n')
        file.flush()
        file.close()

    def get_name(self):
        return 'LAMMPS input format'

    def get_defaultname(self,args):
        return "coord.lammps"

class LAMMPS_output(AtomConfigGenerator):
    coord_file = None
    valid = True
    def __init__(self,coord_file=None):
        if coord_file is None:
            return
        self.coord_file = coord_file

    def supports_frame(self):
        return True

    def _gen_atomic_configuration(self):
        dumpfile = file(self.coord_file)
        logger.info("parsing LAMMPS dump file : "+self.coord_file)
        natm=0
        currnatm = -1
        in_natm = False
        in_cell = False
        in_coords = False
        cell_count=0
        a0=None
        b0=None
        b1=None
        c0=None
        c1=None
        c2=None
        xlo=0.0
        xhi=0.0
        ylo=0.0
        yhi=0.0
        zlo=0.0
        zhi=0.0
        xy=0.0
        xz=0.0
        yz=0.0
        ret = []
        atom_config=None
        unit_cell=None
        element_index=1
        for line in dumpfile:
            line = line.strip()
            words = line.split()
            if line.startswith('ITEM'):
                if len(words)<2: 
                    continue
                keyw = words[1]
                if keyw=='NUMBER': 
                    in_natm = True
                    continue
                elif keyw=='BOX':
                    in_cell = True
                elif keyw=='ATOMS':
                    if atom_config is None:
                        logger.error('the unit cell is seemingly undefined...')
                        break
                    if 'element' in line:
                        for i in range(len(words)):
                            if words[i]=='element':
                                element_index = i-2
                    in_coords=True
                continue
            if in_natm:
                try:
                    natm = int(line)
                    in_natm = False
                    continue
                except ValueError:
                    logger.error("invalid number of atoms : "+line)
                    in_natm = False
                    break
            if in_cell:
                if len(words)<3:
                    continue
                try:
                    if cell_count==0: 
                        xlo=float(words[0])
                        xhi=float(words[1])
                        xy=float(words[2])
                        cell_count += 1
                        continue
                    elif cell_count==1:
                        ylo=float(words[0])
                        yhi=float(words[1])
                        xz=float(words[2])
                        cell_count += 1
                        continue
                    elif cell_count==2:
                        zlo=float(words[0])
                        zhi=float(words[1])
                        yz=float(words[2])
                        xlo = xlo-min([0,xy,xz,xy+xz])
                        xhi = xhi-max([0,xy,xz,xy+xz])
                        ylo = ylo-min([0,yz])
                        yhi = yhi-max([0,yz])
                        a0=xhi-xlo
                        b1=yhi-ylo
                        c2=zhi-zlo
                        b0=xy
                        c0=xz
                        c1=yz

                        cell_count=0
                        in_cell = False

                        a_vector = [a0, 0.0, 0.0]
                        b_vector = [b0, b1,  0.0]
                        c_vector = [c0, c1,  c2 ]
                        unitcell = config.UnitCell(a_vector=a_vector,b_vector=b_vector,c_vector=c_vector)
                        unitcell.set_length_unit(config.ANGSTROM)
                        atom_config = config.AtomConfig(coordinate_system=config.FRACTIONAL, \
                                      length_unit=config.BOHR,unit_cell=unitcell)
                        ret.append(atom_config)
                        continue
                except ValueError:
                    logger.error("failed to parse the unit cell")
                    in_cell = False
                    continue
            if in_coords:
                if currnatm<0:
                    currnatm=1
                if len(words)<5:
                    continue
                if currnatm>natm:
                    logger.error("too many atoms are defined!")
                    continue
                try:
                    e=words[element_index]
                    x=float(words[2])
                    y=float(words[3])
                    z=float(words[4])
                    atom_config.add_atom(element=e,rx=x,ry=y,rz=z)
                    if currnatm==natm:
                        currnatm=-1
                        in_coords=False
                    else:
                        currnatm += 1
            
                except ValueError:
                    logger.error("failed to parse line : "+line)
                    currnatm=-1
                    in_coords=False
                    continue

        dumpfile.close()
        return ret

    def supports_import(self):
        return True

    def supports_export(self):
        return False

    def get_name(self):
        return 'LAMMPS output format'

    def get_defaultname(self,args):
        return "dump"

