from ATK.KohnSham import *

'''
NanoLanguage XYZ import filters

Description:

3 utility routines for specifying atom positions in XYZ format in NanoLanguage

- xyzBlockToLists: Converts a sequence of lines with the format "element x y z" to two
                   lists [elements,positions]

- importXYZFile: Reads an XYZ file from disk and uses xyzBlockToLists to convert
                 the atom positions to two lists [elements,positions]
                 
- importXYZFile: Same as above, only the XYZ file should not contain the two initial lines
                 (number of atoms and comment). Very useful for copy/pasted atom lists!

For more details, see each individual routine

Notes:

* At the moment, these routines can only read single-configuration xyz files.
* If the xyz file contains more atoms than specified on line 1, the remaining ones will be ignored.
* If the xyz file contains fewer atoms than specified on line 1, an error will be raised.
* Errors will be raised on all format errors.
* No units are attached to the positions; this is up to the user to handle.
'''


def convert(line,types=()):
    '''
    Utility routine used by xyzBlockToLists to parse a single line in the xyz format
    This is the place where we enforce the xyz format
      element x y z
    where element is an arbitrary string, and x y z are floats
    We allow an arbitrary amount of leading and trailing white spaces, 
    plus any number and any kind of white spaces as separators
    '''
    import re
    a = re.split('\s*',line.strip())
    for i in range(len(types)):
        a[i] = apply(types[i],[a[i]])
    return tuple(a)

    
def xyzBlockToLists (xyz_block,number_of_atoms=None):
    '''
    Converts a block (input argument "xyz_block") of the format
    
      element1 x1 y1 z1
      element2 x2 y2 z2
      ...
      elementN xN yN zN
      
    to two lists, one with the elements (in NanoLanguage format, i.e. as objects)
    and one with positions (without unit!), and returns these as
    
      [elements,positions]
    
    If N is given as input (input argument "number_of_atoms), the xyz block must 
    contain at least this number of lines (additional ones will be ignored). 
    If N is not given, all lines are converted.
    
    Strict checking of the xyz format is performed. This means, that only
    initial and/or trailing blank lines are allowed (internal blanks lines are not 
    allowed in the xyz format), and any line not conforming to the xyz format
    will raise an error.
    
    This can easily be modified, to allow for e.g. comments (start line with #),
    as indicated in a commented-out piece of code below.
    '''  
    import re
    # Handle DOS files as well as Unix
    all_lines = re.split('\n',xyz_block.replace('\r',''))
    # Strip off all leading and trailing blank lines
    while all_lines[-1]=='':
        all_lines = all_lines[0:-1]
    while all_lines[0]=='':
        all_lines = all_lines[1:]
    # If number of atoms is not known, attempt to convert all lines
    if number_of_atoms==None:
        number_of_atoms = len(all_lines)
    # Otherwise, convert as many as requested 
    # (first we check that there are these many lines in the block)
    elif len(all_lines)<number_of_atoms:
        raise Exception('Number of atoms is %i, but only %i lines found in the xyz block' % (number_of_atoms,len(all_lines)))
    # Count lines for reference in error messages 
    # and to exit when the number of wanted lines is read
    line_ix = 0
    # Store elements and positions
    elements = []
    positions = []
    for line in all_lines:
        line_ix = line_ix+1
        # Read until we have the desired number of atoms
        if line_ix>number_of_atoms:
            break
        try: 
            element,x,y,z = convert(line,(str,)+(float,)*3)
            positions.append([x,y,z])
        except:
            raise Exception('Invalid format ("element x y z" expected) on line %i (%s)' % (line_ix,line))
        try: 
            elements.append(eval("PeriodicTable.%s" % element))
        except:
            raise Exception('Invalid element on line %i (%s)' % (line_ix,line))
    return (elements,positions)


def importXYZFile (filename):
    '''
    Read an xyz file from disk (file name given as "filename"),
    and convert it to two lists (returned)
    
      [elements,positions]
      
    by using xyzBlockToLists above.
    
    The file must correspond to the xyz format (right now, this can only handle
    single configuration files!), meaning
    
        N (number of atoms)
        comment
        element1 x1 y1 z1
        element1 x2 y2 z2
        ...
        elementN xN yN zN

    The comment is discarded. If the file contains more atoms than N,
    the additional ones will be ignored. No blank lines are allowed.
    '''
    f = open(filename,'r')
    try:
        s = f.readline()
        number_of_atoms = int(s)
    except:
        raise Exception('Invalid format (expected an integer number) on line 1: %s' % s)
    try:
        comment = f.readline()
    except:
        raise Exception('Invalid format on line 2: File ended abruptly')
    try:
        xyz_block = f.read()
    except:
        raise Exception('Invalid format on line 3: File ended abruptly')
    return xyzBlockToLists(xyz_block,number_of_atoms)

    
def importSimplifiedXYZFile (filename):
    '''
    Read a "simplified" xyz file from disk (file name given as "filename"), 
    and convert it to two lists (returned)
    
      [elements,positions]
      
    by using xyzBlockToLists above.
    
    The file must correspond to the xyz format (right now, this can only handle
    single configuration files!), but should NOT contain the initial
    two lines (number of atoms and comment) meaning
    
        element1 x1 y1 z1
        element1 x2 y2 z2
        ...
        elementN xN yN zN

    All lines are read; no blank lines are allowed.
    '''
    f = open(filename,'r')
    xyz_block = f.read()
    if xyz_block=='':
        raise Exception('Invalid format on line 1: File ended abruptly')
    return xyzBlockToLists(xyz_block)