##############################################################################
#
# Structure         by DANSE Diffraction group
#                   Simon J. L. Billinge
#                   (c) 2007 trustees of the Michigan State University.
#                   All rights reserved.
#
# File coded by:    Pavol Juhas
#
# See AUTHORS.txt for a list of people who contributed.
# See LICENSE.txt for license information.
#
##############################################################################

"""Parser for XYZ file format, where
first line gives number of atoms
second one has optional title
remaining lines contain symbol, x, y, z
"""

__id__ = "$Id: P_xyz.py 2825 2009-03-09 04:33:12Z juhas $"

import sys

from matter import Structure, Lattice, Atom
from matter import StructureFormatError
from matter.Parsers import StructureParser

class P_xyz(StructureParser):
    """Parser for standard XYZ structure format.
    """

    def __init__(self):
        StructureParser.__init__(self)
        self.format = "xyz"
        return

    def parseLines(self, lines):
        """Parse list of lines in XYZ format.

        Return Structure object or raise StructureFormatError.
        """
        linefields = [l.split() for l in lines]
        # prepare output structure
        stru = Structure()
        # find first valid record
        start = 0
        for field in linefields:
            if len(field) == 0 or field[0] == "#":
                start += 1
            else:
                break
        # first valid line gives number of atoms
        try:
            lfs = linefields[start]
            w1 = linefields[start][0]
            if len(lfs) == 1 and str(int(w1)) == w1:
                p_natoms = int(w1)
                #try to get lattice vectors from description line
                try:
                    latticeVecs = map(float, linefields[start+1])
                    assert len(latticeVecs)==9
                    from matter.Lattice import Lattice
                    reshaped = [latticeVecs[0:3], latticeVecs[3:6], latticeVecs[6:9]]
                    stru.lattice = Lattice(base=reshaped) 
                    needsDescription = True
                except:
                    needsDescription = False
                    stru.description = lines[start+1].strip()
                start += 2
            else:
                emsg = ("%d: invalid XYZ format, missing number of atoms" %
                        (start + 1))
                raise StructureFormatError(emsg)
        except (IndexError, ValueError):
            exc_type, exc_value, exc_traceback = sys.exc_info()
            emsg = ("%d: invalid XYZ format, missing number of atoms" %
                    (start + 1))
            raise StructureFormatError, emsg, exc_traceback
        # find the last valid record
        stop = len(lines)
        while stop > start and len(linefields[stop-1]) == 0:
            stop -= 1
        # get out for empty structure
        if p_natoms == 0 or start >= stop:
            return stru
        # here we have at least one valid record line
        nfields = len(linefields[start])
        if nfields != 4 and nfields != 5:
            emsg = "%d: invalid XYZ format, expected 4 or 5 columns" % (start + 1)
            raise StructureFormatError(emsg)
        # now try to read all record lines
        try:
            p_nl = start
            for fields in linefields[start:] :
                p_nl += 1
                if fields == []:
                    continue
                elif len(fields) != 4 and len(fields) !=5:
                    emsg = ('%d: all lines must have ' +
                            'a symbol, position, and optionally charge') % p_nl
                    raise StructureFormatError(emsg)
                symbol = fields[0]
                symbol = symbol[0].upper() + symbol[1:].lower()
                xyz = [ float(f) for f in fields[1:4] ]                 
                if len(fields)==5:
                    charge = float(fields[4])
                else:
                    charge = 0.0
                stru.addNewAtom(symbol, xyz=xyz, charge=charge)
        except ValueError:
            exc_type, exc_value, exc_traceback = sys.exc_info()
            emsg = "%d: invalid number format" % p_nl
            raise StructureFormatError, emsg, exc_traceback
        # finally check if all the atoms have been read
        if p_natoms is not None and len(stru) != p_natoms:
            emsg = "expected %d atoms, read %d" % (p_natoms, len(stru))
            raise StructureFormatError(emsg)
        if needsDescription:
            stru.generateDescription()
        return stru
    # End of parseLines

    def toLines(self, stru, **kwds):
        """Convert Structure stru to a list of lines in XYZ format.

        Return list of strings.
        """
        use_primitive_unitcell = kwds.get('use_primitive_unitcell')
        if use_primitive_unitcell:
            return self._toLines_primitiveunitcell(stru, **kwds)
        return self._toLines_originalunitcell(stru, **kwds)

    def _toLines_primitiveunitcell(self, stru, **kwds):
        use_fractional_coordinates = kwds.get('use_fractional_coordinates')
        puc = stru.primitive_unitcell
        atoms = puc.atoms
        lines = []
        lines.append( str(len(atoms)) )
        #this next part puts the lattice vectors as the "description"
        if kwds.get('latticeAsDescription'):
            import numpy as np
            lattice = stru.lattice
            base = np.array(lattice.base, copy=1); base.shape = -1
            lines.append(' '.join(map(str, base)))
        elif stru.description:
            lines.append( stru.description )
        else:
            lines.append('\n')
        for a in atoms:
            if use_fractional_coordinates:
                rc = a.xyz
            else:
                rc = a.xyz_cartn
            s = "%-3s %g %g %g" % (a.symbol, rc[0], rc[1], rc[2])
            lines.append(s)
        return lines

    def _toLines_originalunitcell(self, stru, **kwds):
        use_fractional_coordinates = kwds.get('use_fractional_coordinates')
        lines = []
        lines.append( str(len(stru)) )
        #this next part puts the lattice vectors as the "description"
        if kwds.get('latticeAsDescription'):
            import numpy as np
            lattice = stru.lattice
            base = np.array(lattice.base, copy=1); base.shape = -1
            lines.append(' '.join(map(str, base)))
        elif stru.description:
            lines.append( stru.description )
        else:
            lines.append('\n')
        for a in stru:
            if use_fractional_coordinates:
                rc = a.xyz
            else:
                rc = a.xyz_cartn
            s = "%-3s %g %g %g" % (a.symbol, rc[0], rc[1], rc[2])
            lines.append(s)
        return lines
    # End of toLines

# End of class P_xyz

# Routines

def getParser():
    return P_xyz()



import unittest
class TestCase(unittest.TestCase):

    def testWriter(self):
        p = getParser()
        import matter
        a = 1.5
        lattice = matter.Lattice(2*a, 2*a, 2*a, 90,90,90)
        atoms = [matter.Atom('Ni'), matter.Atom('Ni', (0.5,0.5,0.5))]
        struct = Structure(lattice=lattice, atoms=atoms, sgid=229)
        
        print 'original unitcell, cartesian coords'
        print '\n'.join(p.toLines(struct))
        
        print 'original unitcell, fractional coords'
        print '\n'.join(p.toLines(struct, use_fractional_coordinates=1))

        print 'primitive unitcell, cartesian coords'
        print '\n'.join(p.toLines(struct, use_primitive_unitcell=1))

        print 'primitive unitcell, fractional coords'
        print '\n'.join(p.toLines(struct, use_primitive_unitcell=1, use_fractional_coordinates=1))

        return



def main():
    unittest.main()
    return



if __name__ == '__main__': main()


# End of file
