"""Extension to ASE: read and write structures in GEN format

Refer to DFTB+ manual for GEN format description.

Note: GEN format only supports single snapshot.
"""

from ase.atoms import Atoms
from ase.utils import reader, writer


@reader
def read_gen(fileobj):
    """Read structure in GEN format (refer to DFTB+ manual).
       Multiple snapshot are not allowed. """
    image = Atoms()
    lines = fileobj.readlines()
    line = lines[0].split()
    natoms = int(line[0])
    pb_flag = line[1]
    if line[1] not in ['C', 'F', 'S']:
        raise IOError('Error in line #1: only C (Cluster), S (Supercell) '
                      'or F (Fraction) are valid options')

    # Read atomic symbols
    line = lines[1].split()
    # Define a dictionary with symbols-id
    symboldict = dict()
    symbolid = 1
    for symb in line:
        symboldict[symbolid] = symb
        symbolid += 1

    # Read atoms (GEN format supports only single snapshot)
    del lines[:2]
    positions = []
    symbols = []
    for line in lines[:natoms]:
        dummy, symbolid, x, y, z = line.split()[:5]
        symbols.append(symboldict[int(symbolid)])
        positions.append([float(x), float(y), float(z)])
    image = Atoms(symbols=symbols, positions=positions)
    del lines[:natoms]

    # If Supercell, parse periodic vectors.
    # If Fraction, translate into Supercell.
    if pb_flag == 'C':
        return image
    else:
        # Dummy line: line after atom positions is not uniquely defined
        # in gen implementations, and not necessary in DFTB package
        del lines[:1]
        image.set_pbc([True, True, True])
        p = []
        for i in range(3):
            x, y, z = lines[i].split()[:3]
            p.append([float(x), float(y), float(z)])
        image.set_cell([(p[0][0], p[0][1], p[0][2]),
                        (p[1][0], p[1][1], p[1][2]),
                        (p[2][0], p[2][1], p[2][2])])
        if pb_flag == 'F':
            frac_positions = image.get_positions()
            image.set_scaled_positions(frac_positions)
        return image


@writer
def write_gen(fileobj, images):
    """Write structure in GEN format (refer to DFTB+ manual).
       Multiple snapshots are not allowed. """
    if not isinstance(images, (list, tuple)):
        images = [images]

    # Images is kept in a list but a size > 0 is not allowed
    # as GEN format doesn't support multiple snapshots.
    # Images is used as a list for consistency with the other
    # output modules
    if len(images) != 1:
        raise ValueError('images contains more than one structure\n' +
                         'GEN format supports only single snapshot output')

    symbols = images[0].get_chemical_symbols()

    # Define a dictionary with symbols-id
    symboldict = dict()
    for sym in symbols:
        if not (sym in symboldict):
            symboldict[sym] = len(symboldict) + 1
    # An ordered symbol list is needed as ordered dictionary
    # is just available in python 2.7
    orderedsymbols = list(['null'] * len(symboldict.keys()))
    for sym in symboldict.keys():
        orderedsymbols[symboldict[sym] - 1] = sym

    # Check whether the structure is periodic
    # GEN cannot describe periodicity in one or two direction,
    # a periodic structure is considered periodic in all the
    # directions. If your structure is not periodical in all
    # the directions, be sure you have set big periodicity
    # vectors in the non-periodic directions
    if images[0].pbc.any():
        pb_flag = 'S'
    else:
        pb_flag = 'C'

    natoms = len(symbols)
    ind = 0
    for atoms in images:
        fileobj.write('%d  %-5s\n' % (natoms, pb_flag))
        for s in orderedsymbols:
            fileobj.write('%-5s' % s)
        fileobj.write('\n')
        for sym, (x, y, z) in zip(symbols, atoms.get_positions()):
            ind += 1
            symbolid = symboldict[sym]
            fileobj.write('%-6d %d %22.15f %22.15f %22.15f\n' % (ind,
                          symbolid, x, y, z))
    if images[0].pbc.any():
        fileobj.write('%22.15f %22.15f %22.15f \n' % (0.0, 0.0, 0.0))
        fileobj.write('%22.15f %22.15f %22.15f \n' %
                      (images[0].get_cell()[0][0],
                       images[0].get_cell()[0][1],
                       images[0].get_cell()[0][2]))
        fileobj.write('%22.15f %22.15f %22.15f \n' %
                      (images[0].get_cell()[1][0],
                       images[0].get_cell()[1][1],
                       images[0].get_cell()[1][2]))
        fileobj.write('%22.15f %22.15f %22.15f \n' %
                      (images[0].get_cell()[2][0],
                       images[0].get_cell()[2][1],
                       images[0].get_cell()[2][2]))
