''' maipulate I/O files for the VASP code '''
import logging
logger = logging.getLogger('piou.config.vasp')
import os
from piou import config
from piou.data import properties

class VASP_input(config.AtomConfigGenerator):
    inpfile = None
    elems = None

    def get_defaultname(self,args):
        return 'POSCAR'

    def __init__(self,coord_file):
        if coord_file is None:
            self.set_is_valid(False)
            return
        self.inpfile = coord_file

    def get_elem_map(self):
        return self.elems

    def _gen_atomic_configuration(self):
        f=open(self.inpfile,'r')
        line=[]
        for lin in f:
            line.append(lin.strip())
        f.close()

        if len(line)<8:
            logger.error('insufficient number of lines in POSCAR')
            return None

        pot = 'POTCAR'
        dir = os.path.dirname(self.inpfile)
        if len(dir)>0:
            pot = dir+os.sep+'POTCAR'
        f=open(pot,'r')
        key = 'VRHFIN'
        elems=[]
        for li in f:
            li = li.strip()
            if li.startswith(key):
                w0 = li.split('=')
                if w0>=2:
                    elems.append(w0[1].strip().split(':')[0])
        msg = 'elements resolved from the POTCAR file :'
        for ele in elems:
            msg += ' '+ele
        print msg

        factor=1.0
        try:
            factor = float(line[1].split()[0])
        except ValueError:
            logger.error("the 1st line is invalid : "+line[1])
            return None

        try:
            av=line[2].split()
            av[0] = float(av[0])*factor
            av[1] = float(av[1])*factor
            av[2] = float(av[2])*factor
            bv=line[3].split()
            bv[0] = float(bv[0])*factor
            bv[1] = float(bv[1])*factor
            bv[2] = float(bv[2])*factor
            cv=line[4].split()
            cv[0] = float(cv[0])*factor
            cv[1] = float(cv[1])*factor
            cv[2] = float(cv[2])*factor
        except ValueError,IndexError:
            logger.error('failed to read the unit cell')
            return None

        unitcell = config.UnitCell(\
        a_vector=[av[0],av[1],av[2]],\
        b_vector=[bv[0],bv[1],bv[2]],\
        c_vector=[cv[0],cv[1],cv[2]])

        nat_per_elem_line = line[5].split()
        ver5=False
        try:
            iii = int(nat_per_elem_line[0])
        except ValueError:
            ver5=True
        if ver5:
            logger.info("detected POSCAR for VASP ver.5 (and over)")
        else:
            logger.info("detected POSCAR for VASP ver.4")

        nat_per_elem = []
        for i in range(len(nat_per_elem_line)):
            try:
                ii = int(nat_per_elem_line[i])
                nat_per_elem.append(ii)
            except ValueError:
                if not ver5:
                    logger.error('invalid specification for the number of atoms : '+str(line[5]))
                    return None
                else:
                    continue
        if len(nat_per_elem)!=len(elems):
            logger.error('inconsistent number of elements defined in the POSCAR and POTCAR files')
            return None

        nextind=6
        if line[6].lower().startswith('s'):
            nextind=7

        atomcount=0
        for i in range(len(nat_per_elem)):
            atomcount += nat_per_elem[i]
        if atomcount+nextind>len(line):
            logger.error('insufficient number of lines in POSCAR')
            return None

        mode=config.FRACTIONAL
        factatom=1.0
        if line[nextind].lower().startswith('c') or \
            line[nextind].lower().startswith('k'):
            mode=config.CARTESIAN
            factatom=factor
        atomconf=config.AtomConfig(coordinate_system=mode,\
        length_unit=config.ANGSTROM,unit_cell=unitcell)
        unitcell.set_length_unit(config.ANGSTROM)

        self.elems=[]
        atomcount = 0
        for ielem in range(len(nat_per_elem)):
            elem = elems[ielem]
            for i in range(nat_per_elem[ielem]):
                self.elems.append(elem)
                atomcount += 1
                index = atomcount+nextind
                pos = line[index].split()
                for i in range(3):
                    try:
                        pos[i] = float(pos[i])*factatom
                    except ValueError,IndexError:
                        logger.error('failed to read the atomic coordinates from the POSCAR file')
                        return None
                atomconf.add_atom(element=elem,rx=pos[0],ry=pos[1],rz=pos[2])

        return [atomconf]

    def export_atomic_configuration(self,atomic_coordinates,to_file=None,frame_no=None,all_frames=False):
        frameno=config.resolve_frame_no(frame_no,len(atomic_coordinates))
        atomic_coordinate = atomic_coordinates[frameno]
        if not atomic_coordinate.is_cell_ready():
            logger.error('the unit cell is not set')
            return

        file=open(to_file,'w')
        file.write('first line is a comment \n')
        file.write('1\n')
        (avec,bvec,cvec) = atomic_coordinate.get_unit_cell().get_lattice_vector(to_unit=config.ANGSTROM)
        file.write(str(avec[0])+" "+str(avec[1])+" "+str(avec[2])+'\n')
        file.write(str(bvec[0])+" "+str(bvec[1])+" "+str(bvec[2])+'\n')
        file.write(str(cvec[0])+" "+str(cvec[1])+" "+str(cvec[2])+'\n')

        elements = atomic_coordinate.get_elem_list()
        for element in elements:
            file.write(' '+element)
        for element in elements:
            file.write(' '+str(self._get_num_elem(element,atomic_coordinate)))
        file.write('\n')
        file.write('cartesian\n')

        for atom in atomic_coordinate:
            pos = atom.get_pos(mode=config.CARTESIAN, to_unit=config.ANGSTROM)
            file.write(str(pos[0])+" "+str(pos[1])+" "+str(pos[2])+'\n')

        file.flush()
        file.close()

        # ppfile
        ppdir = properties.Property().get_property('pp.vasp_pp_dir')
        if ppdir is not None and os.path.exists(ppdir) and os.path.isdir(ppdir):
            potcar = open('POTCAR','w')
            potfound=0
            for element in elements:
                dirs = os.listdir(ppdir)
                for dir in dirs:
                    if str(dir).startswith(element):
                        pot = open(ppdir+os.sep+dir+os.sep+'POTCAR',mode='r')
                        for l in pot:
                            potcar.write(l)
                        pot.close()
                        potfound += 1
                        break

            if potfound<len(elements):
                logger.warn('failed to resolve pp files for some of the elements')

    def _get_num_elem(self,elemname,coord):
        ret = 0
        for atom in coord:
            if elemname==atom.get_element_name():
                ret+=1
        return ret

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

class VASP_output(config.AtomConfigGenerator):
    def __init__(self,coord_file):
        if coord_file is None:
            self.set_is_valid(False)
            return
        self.inpfile = coord_file

    def supports_frame(self):
        return True

    def supports_export(self):
        return False

    def _gen_atomic_configuration(self):
        dir = os.path.dirname(self.inpfile)
        if len(dir)>0:
            dir = dir+os.sep

        tmp = VASP_input(coord_file=dir+'POSCAR')
        tmplconfig = tmp.get_atomic_configuration()
        cell = tmplconfig.get_unit_cell()
        frame = []
        natm = tmplconfig.get_num_atom()
        f = open(self.inpfile,'r')
        header_count=0
        config_count=0
        elem_map = tmp.get_elem_map()

        for line in f:

            line = line.strip()
            header_count += 1
            if header_count<=5:
                continue

            if config_count==0:
                config_count += 1
                curr_config = config.AtomConfig(\
                coordinate_system=config.FRACTIONAL,length_unit=config.ANGSTROM,\
                unit_cell=cell)
                frame.append(curr_config)
                continue

            pos = line.split()
            if len(pos)<3:
                logger.error('invalid line found in XDATCAR : '+line)
                config_count += 1
                continue

            try:
                curr_config.add_atom(element=elem_map[config_count-1],\
                rx=float(pos[0]),ry=float(pos[1]),rz=float(pos[2]))
            except ValueError, IndexError:
                logger.error('invalid line found in XDATCAR : '+line)
                config_count += 1
                continue

            if config_count==natm:
                config_count = 0
                continue

            config_count += 1

        return frame

    def export_atomic_configuration(self,atomic_coordinates,to_file=None,frame_no=None,all_frames=False):
        pass

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

    def get_defaultname(self,args):
        return 'XDATCAR'

