from cStringIO import StringIO
from ase.io.xyz import read_xyz
from ase.units import Bohr


def read_gamess_us(filename):
    """Method to read geometry from a GAMESS-US output
    """

    f = filename
    if isinstance(filename, str):
        f = open(filename)

    lines = f.readlines()

    done = False
    i = 0
    xyzstring = ''

    for l, line in enumerate(lines):
        if line.find('CHARGE         X                   Y                   Z') >= 0:
            while not done:
                i += 1

                if not lines[l+i] == '\n':
                    sym = lines[l+i].strip().split()[0]
                    for c in range(2,5):
                        pos = map(float, lines[l+i].strip().split()[2:5])
                    
                    xyzstring += sym + ' ' + str(pos[0]*Bohr)+ ' ' + str(pos[1]*Bohr)+ ' ' + str(pos[2]*Bohr)+'\n'

                else:
                    done = True
        if done:
            break

    xyzstring = str(i-1) + '\n\n' + xyzstring
    atoms = read_xyz(StringIO(xyzstring))

    if type(filename) == str:
        f.close()

    return atoms


def read_gamess_us_input(filename):
    """Method to read geometry from an GAMESS-US input file."""
    f = filename
    if isinstance(filename, str):
        f = open(filename)
    lines = f.readlines()

    # Find geometry region of input file.

    done = False
    i = 0
    xyzstring = ''

    for l, line in enumerate(lines):
        if line.find(' $data') > -1:

            i += 2
            while not done:
                i += 1

                if not (lines[l+i].find('$end') > -1):
                    sym = lines[l+i].strip().split()[0]
                    pos = map(float, lines[l+i].strip().split()[2:5])
                    
                    xyzstring += sym + ' ' + str(pos[0])+ ' ' + str(pos[1])+ ' ' + str(pos[2])+'\n'

                else:
                    done = True
        if done:
            break

    xyzstring = str(i-3) + '\n\n' + xyzstring
    atoms = read_xyz(StringIO(xyzstring))

    if type(filename) == str:
        f.close()

    return atoms


def write_gamess_us(filename, atoms, geometry=None):
    """Method to write GAMESS-US coord file
    """

    if isinstance(filename, str):
        f = open(filename, 'w')
    else:  # Assume it's a 'file-like object'
        f = filename

    f.write(' $data\nInput generated by ASE\n')

    if geometry is not None:
        f.write(str(geometry) + '\n')
    else:
        f.write('C1\n')

#FU| we'll only take integer nuclear charges here and put external charges in $FFDATA
    nums = atoms.get_atomic_numbers()
    for a, atom in enumerate(atoms):
        if atom.tag == -71:  # 71 is ascii G (Ghost)
            symbol = 'bq' + '0'
        else:
            symbol = atom.symbol + ' ' + str(nums[a])
        f.write('  ' + symbol + ' ' +
                str(atom.position[0]) + ' ' +
                str(atom.position[1]) + ' ' +
                str(atom.position[2]) + '\n')
    f.write(' $end\n')
