
def read_aims(filename):
    """Import FHI-aims geometry type files.

    Reads unitcell, atom positions and constraints from
    a geometry.in file.
    """

    from ase import Atoms
    from ase.constraints import FixAtoms, FixCartesian
    import numpy as np

    atoms = Atoms()
    fd = open(filename, 'r')
    lines = fd.readlines()
    fd.close()
    positions = []
    cell = []
    symbols = []
    magmoms = []
    fix = []
    fix_cart = []
    xyz = np.array([0, 0, 0])
    i = -1
    n_periodic = -1
    periodic = np.array([False, False, False])
    for n, line in enumerate(lines):
        inp = line.split()
        if inp == []:
            continue
        if inp[0] == 'atom':
            if xyz.all():
                fix.append(i)
            elif xyz.any():
                fix_cart.append(FixCartesian(i, xyz))
            floatvect = float(inp[1]), float(inp[2]), float(inp[3])
            positions.append(floatvect)
            symbols.append(inp[-1])
            i += 1
            xyz = np.array([0, 0, 0])
        elif inp[0] == 'lattice_vector':
            floatvect = float(inp[1]), float(inp[2]), float(inp[3])
            cell.append(floatvect)
            n_periodic = n_periodic + 1
            periodic[n_periodic] = True
        elif inp[0] == 'initial_moment':
            magmoms.append(float(inp[1]))
        if inp[0] == 'constrain_relaxation':
            if inp[1] == '.true.':
                fix.append(i)
            elif inp[1] == 'x':
                xyz[0] = 1
            elif inp[1] == 'y':
                xyz[1] = 1
            elif inp[1] == 'z':
                xyz[2] = 1
    if xyz.all():
        fix.append(i)
    elif xyz.any():
        fix_cart.append(FixCartesian(i, xyz))
    atoms = Atoms(symbols, positions)
    if len(magmoms) > 0:
        atoms.set_initial_magnetic_moments(magmoms)
    if periodic.any():
        atoms.set_cell(cell)
        atoms.set_pbc(periodic)
    if len(fix):
        atoms.set_constraint([FixAtoms(indices=fix)]+fix_cart)
    else:
        atoms.set_constraint(fix_cart)
    return atoms


def write_aims(filename, atoms, ghosts=None):
    """Method to write FHI-aims geometry files.

    Writes the atoms positions and constraints (only FixAtoms is
    supported at the moment). 
    """

    from ase.constraints import FixAtoms, FixCartesian
    import numpy as np

    if isinstance(atoms, (list, tuple)):
        if len(atoms) > 1:
            raise RuntimeError("Don't know how to save more than "+
                               "one image to FHI-aims input")
        else:
            atoms = atoms[0]

    fd = open(filename, 'w')
    fd.write('#=======================================================\n')
    fd.write('#FHI-aims file: '+filename+'\n')
    fd.write('#Created using the Atomic Simulation Environment (ASE)\n')
    fd.write('#=======================================================\n')
    i = 0
    if atoms.get_pbc().any():
        for n, vector in enumerate(atoms.get_cell()):
            fd.write('lattice_vector ')
            for i in range(3):
                fd.write('%16.16f ' % vector[i])
            fd.write('\n')
    fix_cart = np.zeros([len(atoms),3]) 

    if atoms.constraints:
        for constr in atoms.constraints:
            if isinstance(constr, FixAtoms):
                fix_cart[constr.index] = [1,1,1]
            elif isinstance(constr, FixCartesian):
                fix_cart[constr.a] = -constr.mask+1

    if ghosts is None:
        ghosts = np.zeros(len(atoms))
    else:
        assert len(ghosts) == len(atoms)
    for i, atom in enumerate(atoms):
        if ghosts[i] == 1:
            atomstring = 'empty '
        else:
            atomstring = 'atom '
        fd.write(atomstring)
        for pos in atom.position:
            fd.write('%16.16f ' % pos)
        fd.write(atom.symbol)
        fd.write('\n')
        # (1) all coords are constrained:
        if fix_cart[i].all():
            fd.write('constrain_relaxation .true.\n')
        # (2) some coords are constrained:
        elif fix_cart[i].any():
            xyz = fix_cart[i]
            for n in range(3):
                if xyz[n]:
                    fd.write('constrain_relaxation %s\n' % 'xyz'[n])
        if atom.charge:
            fd.write('initial_charge %16.6f\n' % atom.charge)
        if atom.magmom:
            fd.write('initial_moment %16.6f\n' % atom.magmom)
# except KeyError:
#     continue


def read_energy(filename):
    for line in open(filename, 'r'):
        if line.startswith('  | Total energy corrected'):
            E = float(line.split()[-2])
    return E


def read_aims_output(filename, index = -1):
    """  Import FHI-aims output files with all data available, i.e. relaxations, 
    MD information, force information etc etc etc. """
    from ase import Atoms, Atom 
    from ase.calculators.singlepoint import SinglePointCalculator
    from ase.units import Ang, fs
    from ase.constraints import FixAtoms, FixCartesian
    molecular_dynamics = False
    fd = open(filename, 'r')
    cell = []
    images = []
    fix = []
    fix_cart = []    
    n_periodic = -1
    f = None
    pbc = False
    found_aims_calculator = False
    v_unit = Ang/(1000.0*fs)
    while True:
        line = fd.readline()
        if not line:
            break
        if "List of parameters used to initialize the calculator:" in line:
            fd.readline()
            calc = read_aims_calculator(fd)
            calc.out = filename
            found_aims_calculator = True
        if "| Number of atoms                   :" in line:
            inp = line.split()
            n_atoms = int(inp[5])
        if "| Unit cell:" in line:
            if not pbc:
                pbc = True
                for i in range(3):
                    inp = fd.readline().split()
                    cell.append([inp[1],inp[2],inp[3]])
        if "Found relaxation constraint for atom" in line:
            xyz = [0, 0, 0]
            ind = int(line.split()[5][:-1])-1
            if "All coordinates fixed" in line:
                if ind not in fix:
                    fix.append(ind)
            if "coordinate fixed" in line:
                coord = line.split()[6]
                constr_ind = 0
                if coord == 'x':
                    xyz[0] = 1
                elif coord == 'y':
                    xyz[1] = 1
                elif coord == 'z':
                    xyz[2] = 1
                keep = True
                for n,c in enumerate(fix_cart):
                    if ind == c.a:
                        keep = False
                        constr_ind = n
                if keep:
                    fix_cart.append(FixCartesian(ind, xyz))
                else:
                    fix_cart[n].mask[xyz.index(1)] = 0
        if "Atomic structure:" in line and not molecular_dynamics:
            fd.readline()
            atoms = Atoms()
            for i in range(n_atoms):
                inp = fd.readline().split()
                atoms.append(Atom(inp[3],(inp[4],inp[5],inp[6])))
        if "Complete information for previous time-step:" in line:
            molecular_dynamics = True
        if "Updated atomic structure:" in line and not molecular_dynamics:
            fd.readline()
            atoms = Atoms()
            velocities = []
            for i in range(n_atoms):
                inp = fd.readline().split()
                if 'lattice_vector' in inp[0]:
                    cell = []
                    for i in range(3):
                        cell += [[float(inp[1]),float(inp[2]),float(inp[3])]]
                        inp = fd.readline().split()
                    atoms.set_cell(cell)
                    inp = fd.readline().split()
                atoms.append(Atom(inp[4],(inp[1],inp[2],inp[3])))  
                if molecular_dynamics:
                    inp = fd.readline().split()
        if "Atomic structure (and velocities)" in line:
            fd.readline()
            atoms = Atoms()
            velocities = []
            for i in range(n_atoms):
                inp = fd.readline().split()
                atoms.append(Atom(inp[4],(inp[1],inp[2],inp[3])))  
                inp = fd.readline().split()
                velocities += [[float(inp[1])*v_unit,float(inp[2])*v_unit,float(inp[3])*v_unit]]
            atoms.set_velocities(velocities)
            if len(fix):
                atoms.set_constraint([FixAtoms(indices=fix)]+fix_cart)
            else:
                atoms.set_constraint(fix_cart)
            images.append(atoms)
        if "Total atomic forces" in line:
            f = []
            for i in range(n_atoms):
                inp = fd.readline().split()
                f.append([float(inp[2]),float(inp[3]),float(inp[4])])
            if not found_aims_calculator:
                e = images[-1].get_potential_energy()
                images[-1].set_calculator(SinglePointCalculator(atoms,
                                                                energy=e,
                                                                forces=f))
            e = None
            f = None
        if "Total energy corrected" in line:
            e = float(line.split()[5])
            if pbc:
                atoms.set_cell(cell)
                atoms.pbc = True
            if not found_aims_calculator:
                atoms.set_calculator(SinglePointCalculator(atoms, energy=e))
            if not molecular_dynamics: 
                if len(fix):
                    atoms.set_constraint([FixAtoms(indices=fix)]+fix_cart)
                else:
                    atoms.set_constraint(fix_cart)
                images.append(atoms)
            e = None
            if found_aims_calculator:
                calc.set_results(images[-1])
                images[-1].set_calculator(calc)
    fd.close()
    if molecular_dynamics:
        images = images[1:]

    # return requested images, code borrowed from ase/io/trajectory.py
    if isinstance(index, int):
        return images[index]
    else:
        step = index.step or 1
        if step > 0:
            start = index.start or 0
            if start < 0:
                start += len(images)
            stop = index.stop or len(images)
            if stop < 0:
                stop += len(images)
        else:
            if index.start is None:
                start = len(images) - 1
            else:
                start = index.start
                if start < 0:
                    start += len(images)
            if index.stop is None:
                stop = -1
            else:
                stop = index.stop
                if stop < 0:
                    stop += len(images)
        return [images[i] for i in range(start, stop, step)]
