# Source code for MDAnalysis.coordinates.FHIAIMS

    # -*- Mode: python; tab-width: 4; indent-tabs-mode:nil; coding:utf-8 -*-
    # vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4
    #
    # MDAnalysis --- https://www.mdanalysis.org
    # Copyright (c) 2006-2017 The MDAnalysis Development Team and contributors
    # (see the file AUTHORS for the full list of names)
    #
    # Released under the Lesser GNU Public Licence, v2.1 or any higher version
    #
    # Please cite your use of MDAnalysis in published work:
    #
    # R. J. Gowers, M. Linke, J. Barnoud, T. J. E. Reddy, M. N. Melo, S. L. Seyler,
    # D. L. Dotson, J. Domanski, S. Buchoux, I. M. Kenney, and O. Beckstein.
    # MDAnalysis: A Python package for the rapid analysis of molecular dynamics
    # simulations. In S. Benthall and S. Rostrup editors, Proceedings of the 15th
    # Python in Science Conference, pages 102-109, Austin, TX, 2016. SciPy.
    # doi: 10.25080/majora-629e541a-00e
    #
    # N. Michaud-Agrawal, E. J. Denning, T. B. Woolf, and O. Beckstein.
    # MDAnalysis: A Toolkit for the Analysis of Molecular Dynamics Simulations.
    # J. Comput. Chem. 32 (2011), 2319--2327, doi:10.1002/jcc.21787
    #

    """
    FHI-AIMS file format --- :mod:`MDAnalysis.coordinates.FHIAIMS`
    ==============================================================

    Classes to read and write `FHI-AIMS`_ coordinate files.

    The cell vectors are specified by the (optional) lines with the
    ``lattice_vector`` tag::

        lattice_vector x  y  z

    where x, y, and z are expressed in ångström (Å).

    .. Note::

       In the original `FHI-AIMS format`_, up to three lines with
       ``lattice_vector`` are allowed (order matters) where the absent line implies
       no periodicity in that direction.  In MDAnalysis, only the case of no
       ``lattice_vector`` or three ``lattice_vector`` lines are allowed.

    Atomic positions and names are specified either by the ``atom`` or by the
    ``atom_frac`` tags::

        atom           x  y  z  name
        atom_frac      nx ny nz name

    where x, y, and z are expressed in ångström, and nx, ny and nz are real numbers
    in [0, 1] and are used to compute the atomic positions in units of the basic
    cell.

    Atomic velocities can be added on the line right after the corresponding
    ``atom`` in units of Å/ps using the ``velocity`` tag::

        velocity      vx vy vz

    The field name is a string identifying the atomic species.  See also the
    specifications in the official `FHI-AIMS format`_.

    Classes
    -------

    .. autoclass:: FHIAIMSReader
       :members:
       :inherited-members:

    .. autoclass:: FHIAIMSWriter
       :members:
       :inherited-members:

    Developer notes: ``FHIAIMSWriter`` format strings
    -------------------------------------------------

    The :class:`FHIAIMSWriter` class has a :attr:`FHIAIMSWriter.fmt`
    attribute, which is a dictionary of different strings for writing
    lines in ``.in`` files.  These are as follows:

    ``xyz``
      An atom line without velocities.  Requires that the `name` and
      `pos` keys be supplied.  E.g.::

         fmt['xyz'].format(pos=(0.0, 1.0, 2.0), name='O')

    ``vel``
      An line that specifies velocities::

         fmt['xyz'].format(vel=(0.1, 0.2, 0.3))

    ``box_triclinic``
      The (optional) initial lines of the file which gives box dimensions.
      Requires the `box` keyword, as a length 9 vector.  This is a flattened
      version of the (3, 3) triclinic vector representation of the unit
      cell.

    .. Links

    .. _FHI-AIMS: https://aimsclub.fhi-berlin.mpg.de/
    .. _FHI-AIMS format: https://doi.org/10.6084/m9.figshare.12413477.v1

    """
    import re

    import itertools
    import warnings

    import numpy as np

    from . import base
    from .core import triclinic_box, triclinic_vectors
    from ..exceptions import NoDataError
    from ..lib import util
    from ..lib import mdamath

    class FHIAIMSReader(base.SingleFrameReaderBase):
        """Reader for the FHIAIMS geometry format.

        Single frame reader for the `FHI-AIMS`_ input file format.  Reads
        geometry (3D and molecules only), positions (absolut or fractional),
        velocities if given, all according to the `FHI-AIMS format`_
        specifications

        """

        format = ["IN", "FHIAIMS"]
        units = {"time": "ps", "length": "Angstrom", "velocity": "Angstrom/ps"}

        def _read_first_frame(self):
            with util.openany(self.filename, "rt") as fhiaimsfile:
                relative, positions, velocities, lattice_vectors = [], [], [], []
                skip_tags = ["#", "initial_moment"]
                oldline = ""
                for line in fhiaimsfile:
                    line = line.strip()
                    if line.startswith("atom"):
                        positions.append(line.split()[1:-1])
                        relative.append("atom_frac" in line)
                        oldline = line
                        continue
                    if line.startswith("velocity"):
                        if not "atom" in oldline:
                            raise ValueError(
                                "Non-conforming line (velocity must follow atom): ({0})in FHI-AIMS input file {0}".format(
                                    line, self.filename
                                )
                            )
                        velocities.append(line.split()[1:])
                        oldline = line
                        continue
                    if line.startswith("lattice"):
                        lattice_vectors.append(line.split()[1:])
                        oldline = line
                        continue
                    if any([line.startswith(tag) for tag in skip_tags]):
                        oldline = line
                        continue
                    raise ValueError(
                        "Non-conforming line: ({0})in FHI-AIMS input file {0}".format(
                            line, self.filename
                        )
                    )

            # positions and velocities are lists of lists of strings; they will be
            # cast to np.arrays(..., dtype=float32) during assignment to ts.positions/ts.velocities
            lattice_vectors = np.asarray(lattice_vectors, dtype=np.float32)

            if len(velocities) not in (0, len(positions)):
                raise ValueError(
                    "Found incorrect number of velocity tags ({0}) in the FHI-AIMS file, should be {1}.".format(
                        len(velocities), len(positions)
                    )
                )
            if len(lattice_vectors) not in (0, 3):
                raise ValueError(
                    "Found partial periodicity in FHI-AIMS file. This cannot be handled by MDAnalysis."
                )
            if len(lattice_vectors) == 0 and any(relative):
                raise ValueError(
                    "Found relative coordinates in FHI-AIMS file without lattice info."
                )

            # create Timestep

            self.n_atoms = n_atoms = len(positions)
            self.ts = ts = self._Timestep(n_atoms, **self._ts_kwargs)
            ts.positions = positions

            if len(lattice_vectors) > 0:
                ts.dimensions = triclinic_box(*lattice_vectors)
                ts.positions[relative] = np.matmul(
                    ts.positions[relative], lattice_vectors
                )

            if len(velocities) > 0:
                ts.velocities = velocities

            self.ts.frame = 0  # 0-based frame number

        def Writer(self, filename, n_atoms=None, **kwargs):
            """Returns a FHIAIMSWriter for *filename*.

            Parameters
            ----------
            filename: str
                filename of the output FHI-AIMS file

            Returns
            -------
            :class:`FHIAIMSWriter`

            """
            if n_atoms is None:
                n_atoms = self.n_atoms
            return FHIAIMSWriter(filename, n_atoms=n_atoms, **kwargs)

    class FHIAIMSWriter(base.WriterBase):
        """FHI-AIMS Writer.

        Single frame writer for the `FHI-AIMS`_ format.  Writes geometry (3D and
        molecules only), positions (absolut only), velocities if given, all
        according to the `FHI-AIMS format`_ specifications.

        If no atom names are given, it will set each atom name to "X".

        """

        format = ["IN", "FHIAIMS"]
        units = {"time": None, "length": "Angstrom", "velocity": "Angstrom/ps"}

        #: format strings for the FHI-AIMS file (all include newline)
        fmt = {
            # coordinates output format, see  https://doi.org/10.6084/m9.figshare.12413477.v1
            "xyz": "atom {pos[0]:12.8f} {pos[1]:12.8f} {pos[2]:12.8f} {name:<3s}\n",
            "vel": "velocity {vel[0]:12.8f} {vel[1]:12.8f} {vel[2]:12.8f}\n",
            # unitcell
            "box_triclinic": "lattice_vector {box[0]:12.8f} {box[1]:12.8f} {box[2]:12.8f}\nlattice_vector {box[3]:12.8f} {box[4]:12.8f} {box[5]:12.8f}\nlattice_vector {box[6]:12.8f} {box[7]:12.8f} {box[8]:12.8f}\n",
        }

        def __init__(self, filename, convert_units=True, n_atoms=None, **kwargs):
            """Set up the FHI-AIMS Writer

            Parameters
            -----------
            filename : str
                output filename
            n_atoms : int (optional)
                number of atoms

            """
            self.filename = util.filename(filename, ext=".in", keep=True)
            self.n_atoms = n_atoms

        def _write_next_frame(self, obj):
            """Write selection at current trajectory frame to file.

            Parameters
            -----------
            obj : AtomGroup or Universe

            .. versionchanged:: 2.0.0
               Support for deprecated Timestep argument has now been removed.
               Use AtomGroup or Universe as an input instead.
            """
            # write() method that complies with the Trajectory API
            try:
                # make sure to use atoms (Issue 46)
                ag = obj.atoms
                # can write from selection == Universe (Issue 49)
            except AttributeError:
                errmsg = "Input obj is neither an AtomGroup or Universe"
                raise TypeError(errmsg) from None

            # Check for topology information
            missing_topology = []
            try:
                names = ag.names
            except (AttributeError, NoDataError):
                names = itertools.cycle(("X",))
                missing_topology.append("names")

            try:
                atom_indices = ag.ids
            except (AttributeError, NoDataError):
                atom_indices = range(1, ag.n_atoms + 1)
                missing_topology.append("ids")

            if missing_topology:
                warnings.warn(
                    "Supplied AtomGroup was missing the following attributes: "
                    "{miss}. These will be written with default values. "
                    "".format(miss=", ".join(missing_topology))
                )

            positions = ag.positions
            try:
                velocities = ag.velocities
                has_velocities = True
            except (AttributeError, NoDataError):
                has_velocities = False

            with util.openany(self.filename, "wt") as output_fhiaims:
                # Lattice
                try:  # for AtomGroup/Universe
                    tri_dims = obj.universe.trajectory.ts.triclinic_dimensions
                except AttributeError:  # for Timestep
                    tri_dims = obj.triclinic_dimensions
                # full output
                if tri_dims is not None:
                    output_fhiaims.write(
                        self.fmt["box_triclinic"].format(box=tri_dims.flatten())
                    )

                # Atom descriptions and coords
                # Dont use enumerate here,
                # all attributes could be infinite cycles!
                for atom_index, name in zip(range(ag.n_atoms), names):
                    output_fhiaims.write(
                        self.fmt["xyz"].format(
                            pos=positions[atom_index], name=name
                        )
                    )
                    if has_velocities:
                        output_fhiaims.write(
                            self.fmt["vel"].format(vel=velocities[atom_index])
                        )