# -*- 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
#


"""XYZ trajectory reader --- :mod:`MDAnalysis.coordinates.XYZ`
==============================================================

The :ref:`XYZ format <xyz-format>` is a loosely defined, simple
coordinate trajectory format. The implemented format definition was
taken from the `VMD xyzplugin`_ and is therefore compatible with VMD.

Note the following:

* Comments are not allowed in the XYZ file (we neither read nor write
  them to remain compatible with VMD).
* The atom name (first column) is ignored during reading.
* The coordinates are assumed to be space-delimited rather than fixed
  width (this may cause issues - see below).
* All fields to the right of the z-coordinate are ignored.
* The unitcell information is all zeros since this is not recorded in
  the XYZ format.

.. rubric:: Units

* Coordinates are in Angstroms.
* The length of a timestep can be set by passing the *dt* argument,
  it's assumed to be in ps (default: 1 ps).

There appears to be no rigid format definition so it is likely users
will need to tweak this class.

.. _xyz-format:

XYZ File format
---------------

Definition used by the :class:`XYZReader` and :class:`XYZWriter` (and
the `VMD xyzplugin`_ from whence the definition was taken)::

    [ comment line            ] !! NOT IMPLEMENTED !! DO NOT INCLUDE
    [ N                       ] # of atoms, required by this xyz reader plugin  line 1
    [ molecule name           ] name of molecule (can be blank)                 line 2
    atom1 x y z [optional data] atom name followed by xyz coords                line 3
    atom2 x y z [ ...         ] and (optionally) other data.
    ...
    atomN x y z [ ...         ]                                                 line N+2


Note
----
* comment lines not implemented (do not include them)
* molecule name: the line is required but the content is ignored
  at the moment
* optional data (after the coordinates) are presently ignored


.. Links
.. _`VMD xyzplugin`: http://www.ks.uiuc.edu/Research/vmd/plugins/molfile/xyzplugin.html

Classes
-------

"""
import itertools
import os
import errno
import numpy as np
import warnings
import logging

logger = logging.getLogger("MDAnalysis.coordinates.XYZ")

from . import base
from .timestep import Timestep
from ..lib import util
from ..lib.util import cached, store_init_arguments
from ..exceptions import NoDataError
from ..version import __version__


class XYZWriter(base.WriterBase):
    """Writes an XYZ file

    The XYZ file format is not formally defined. This writer follows
    the VMD implementation for the molfile `xyzplugin`_.


    Notes
    -----
    By default, the XYZ writer will attempt to use the input
    :class:`~MDAnalysis.core.groups.AtomGroup` or
    :class:`~MDAnalysis.core.universe.Universe` ``elements`` record to assign
    atom names in the XYZ file. If the ``elements`` record is missing, then
    the ``name`` record will be used. In the event that neither of these are
    available, the atoms will all be named ``X``. Please see, the
    `User Guide`_ for more information on how to add topology attributes if
    you wish to add your own elements / atom names to a
    :class:`~MDAnalysis.core.universe.Universe`.


    .. Links

    .. _xyzplugin:
           http://www.ks.uiuc.edu/Research/vmd/plugins/molfile/xyzplugin.html
    .. _User Guide:
           https://userguide.mdanalysis.org/examples/constructing_universe.html#Adding-topology-attributes


    .. versionchanged:: 1.0.0
       Use elements attribute instead of names attribute, if present.
    .. versionchanged:: 2.0.0
       Support for passing timestep to the writer was deprecated in 1.0 and
       has now been removed. As a consequence, custom names can no longer be
       passed to the writer, these should be added to the
       :class:`~MDAnalysis.core.universe.Universe`, or
       :class:`~MDAnalysis.core.groups.AtomGroup` before invoking the writer.
    """

    format = "XYZ"
    multiframe = True
    # these are assumed!
    units = {"time": "ps", "length": "Angstrom"}

    def __init__(
        self,
        filename,
        n_atoms=None,
        convert_units=True,
        remark=None,
        precision=5,
        **kwargs,
    ):
        """Initialize the XYZ trajectory writer

        Parameters
        ----------
        filename: str
            filename of trajectory file. If it ends with "gz" then the file
            will be gzip-compressed; if it ends with "bz2" it will be bzip2
            compressed.
        n_atoms: int (optional)
            Number of atoms in trajectory. By default assume that this is None
            and that this file is used to store several different models
            instead of a single trajectory. If a number is provided each
            written TimeStep has to contain the same number of atoms.
        convert_units : bool (optional)
            convert quantities to default MDAnalysis units of Angstrom upon
            writing  [``True``]
        remark: str (optional)
            single line of text ("molecule name"). By default writes MDAnalysis
            version and frame
        precision: int (optional)
            set precision of saved trjactory to this number of decimal places.

            .. versionadded:: 2.9.0


        .. versionchanged:: 1.0.0
           Removed :code:`default_remark` variable (Issue #2692).
        .. versionchanged:: 2.0.0
           Due to the removal of timestep as an input for writing, the atoms
           parameter is no longer relevant and has been removed. If passing
           an empty universe, please use ``add_TopologyAttr`` to add in the
           required elements or names.
        """
        self.filename = filename
        self.remark = remark
        self.n_atoms = n_atoms
        self.convert_units = convert_units
        self.precision = precision

        # can also be gz, bz2
        self._xyz = util.anyopen(self.filename, "wt")

    def _get_atoms_elements_or_names(self, atoms):
        """Return a list of atom elements (if present) or fallback to atom names"""
        try:
            return atoms.atoms.elements
        except (AttributeError, NoDataError):
            try:
                return atoms.atoms.names
            except (AttributeError, NoDataError):
                wmsg = (
                    "Input AtomGroup or Universe does not have atom "
                    "elements or names attributes, writer will default "
                    "atom names to 'X'"
                )
                warnings.warn(wmsg)
                return itertools.cycle(("X",))

    def close(self):
        """Close the trajectory file and finalize the writing"""
        if self._xyz is not None:
            self._xyz.write("\n")
            self._xyz.close()
        self._xyz = None

    def write(self, obj):
        """Write object `obj` at current trajectory frame to file.

        Atom elements (or names) in the output are taken from the `obj` or
        default to the value of the `atoms` keyword supplied to the
        :class:`XYZWriter` constructor.

        Parameters
        ----------
        obj : Universe or AtomGroup
            The :class:`~MDAnalysis.core.groups.AtomGroup` or
            :class:`~MDAnalysis.core.universe.Universe` to write.


        .. versionchanged:: 2.0.0
           Deprecated support for Timestep argument has now been removed.
           Use AtomGroup or Universe as an input instead.
        """
        # prepare the Timestep and extract atom names if possible
        # (The way it is written it should be possible to write
        # trajectories with frames that differ in atom numbers
        # but this is not tested.)
        try:
            atoms = obj.atoms
        except AttributeError:
            errmsg = "Input obj is neither an AtomGroup or Universe"
            raise TypeError(errmsg) from None
        else:
            if hasattr(obj, "universe"):
                # For AtomGroup and children (Residue, ResidueGroup, Segment)
                ts_full = obj.universe.trajectory.ts
                if ts_full.n_atoms == atoms.n_atoms:
                    ts = ts_full
                else:
                    # Only populate a time step with the selected atoms.
                    ts = ts_full.copy_slice(atoms.indices)
            elif hasattr(obj, "trajectory"):
                # For Universe only --- get everything
                ts = obj.trajectory.ts
            # update atom names
            self.atomnames = self._get_atoms_elements_or_names(atoms)

        self._write_next_frame(ts)

    def _write_next_frame(self, ts=None):
        """
        Write coordinate information in *ts* to the trajectory

        .. versionchanged:: 1.0.0
           Print out :code:`remark` if present, otherwise use generic one
           (Issue #2692).
           Renamed from `write_next_timestep` to `_write_next_frame`.
        """
        if ts is None:
            if not hasattr(self, "ts"):
                raise NoDataError(
                    "XYZWriter: no coordinate data to write to "
                    "trajectory file"
                )
            else:
                ts = self.ts

        if self.n_atoms is not None:
            if self.n_atoms != ts.n_atoms:
                raise ValueError(
                    "n_atoms keyword was specified indicating "
                    "that this should be a trajectory of the "
                    "same model. But the provided TimeStep has a "
                    "different number ({}) then expected ({})"
                    "".format(ts.n_atoms, self.n_atoms)
                )
        else:
            if (
                not isinstance(self.atomnames, itertools.cycle)
                and len(self.atomnames) != ts.n_atoms
            ):
                logger.info(
                    "Trying to write a TimeStep with unknown atoms. "
                    'Expected {} atoms, got {}. Try using "write" if you are '
                    'using "_write_next_frame" directly'.format(
                        len(self.atomnames), ts.n_atoms
                    )
                )
                self.atomnames = np.array([self.atomnames[0]] * ts.n_atoms)

        if self.convert_units:
            coordinates = self.convert_pos_to_native(
                ts.positions, inplace=False
            )
        else:
            coordinates = ts.positions

        # Write number of atoms
        self._xyz.write("{0:d}\n".format(ts.n_atoms))

        # Write remark
        if self.remark is None:
            remark = (
                "frame {} | Written by MDAnalysis {} (release {})\n".format(
                    ts.frame, self.__class__.__name__, __version__
                )
            )

            self._xyz.write(remark)
        else:
            self._xyz.write(self.remark.strip() + "\n")

        # Write content
        for atom, (x, y, z) in zip(self.atomnames, coordinates):
            self._xyz.write(
                "{0!s:>8}  {1:10.{p}f} {2:10.{p}f} {3:10.{p}f}\n"
                "".format(atom, x, y, z, p=self.precision)
            )


class XYZReader(base.ReaderBase):
    """Reads from an XYZ file

    :Data:
        :attr:`ts`
          Timestep object containing coordinates of current frame

    :Methods:
        ``len(xyz)``
          return number of frames in xyz
        ``for ts in xyz:``
          iterate through trajectory

    .. Note: this can read both compressed (foo.xyz) and compressed
          (foo.xyz.bz2 or foo.xyz.gz) files; uncompression is handled
          on the fly and also reads streams via
          :class:`~MDAnalysis.lib.util.NamedStream`.

    The XYZ file format follows VMD's xyzplugin_ and is also described
    under :ref:`XYZ format <xyz-format>`.

    .. versionchanged:: 0.11.0
       Frames now 0-based instead of 1-based. Added *dt* and
       *time_offset* keywords (passed to :class:`Timestep`)
    """

    # Phil Fowler:
    # Validation: the geometric centre of 1284 atoms was calculated over
    # 500 frames using both MDAnalysis and a VMD Tcl script. There was
    # exact agreement (measured to three decimal places). bzipped and
    # gzipped versions of the XYZ file were also tested

    format = "XYZ"
    # these are assumed!
    units = {"time": "ps", "length": "Angstrom"}
    _Timestep = Timestep

    @store_init_arguments
    def __init__(self, filename, **kwargs):
        super(XYZReader, self).__init__(filename, **kwargs)

        # the filename has been parsed to be either be foo.xyz or foo.xyz.bz2 by
        # coordinates::core.py so the last file extension will tell us if it is
        # bzipped or not
        root, ext = os.path.splitext(self.filename)
        self.xyzfile = util.anyopen(self.filename)
        self.compression = ext[1:] if ext[1:] != "xyz" else None
        self._cache = dict()

        self.ts = self._Timestep(self.n_atoms, **self._ts_kwargs)
        # Haven't quite figured out where to start with all the self._reopen()
        # etc.
        # (Also cannot just use seek() or reset() because that would break
        # with urllib2.urlopen() streams)
        self._read_next_timestep()

    @property
    @cached("n_atoms")
    def n_atoms(self):
        """number of atoms in a frame"""
        with util.anyopen(self.filename) as f:
            n = f.readline()
        # need to check type of n
        return int(n)

    @property
    @cached("n_frames")
    def n_frames(self):
        try:
            return self._read_xyz_n_frames()
        except IOError:
            return 0

    def _read_xyz_n_frames(self):
        # the number of lines in the XYZ file will be 2 greater than the
        # number of atoms
        linesPerFrame = self.n_atoms + 2
        counter = 0
        offsets = []

        with util.anyopen(self.filename) as f:
            line = True
            while line:
                if not counter % linesPerFrame:
                    offsets.append(f.tell())
                line = f.readline()
                counter += 1

        # need to check this is an integer!
        n_frames = int(counter / linesPerFrame)
        self._offsets = offsets
        return n_frames

    def _read_frame(self, frame):
        self.xyzfile.seek(self._offsets[frame])
        self.ts.frame = frame - 1  # gets +1'd in next
        return self._read_next_timestep()

    def _read_next_timestep(self, ts=None):
        # check that the timestep object exists
        if ts is None:
            ts = self.ts

        f = self.xyzfile

        try:
            # we assume that there are only two header lines per frame
            f.readline()
            f.readline()
            # convert all entries at the end once for optimal speed
            tmp_buf = []
            for i in range(self.n_atoms):
                tmp_buf.append(f.readline().split()[1:4])
            ts.positions = tmp_buf
            ts.frame += 1
            return ts
        except (ValueError, IndexError) as err:
            raise EOFError(err) from None

    def _reopen(self):
        self.close()
        self.open_trajectory()

    def open_trajectory(self):
        if self.xyzfile is not None:
            raise IOError(
                errno.EALREADY, "XYZ file already opened", self.filename
            )

        self.xyzfile = util.anyopen(self.filename)

        # reset ts
        ts = self.ts
        ts.frame = -1

        return self.xyzfile

    def Writer(self, filename, n_atoms=None, **kwargs):
        """Returns a XYZWriter for *filename* with the same parameters as this
        XYZ.

        Parameters
        ----------
        filename: str
            filename of the output trajectory
        n_atoms: int (optional)
            number of atoms. If none is given use the same number of atoms from
            the reader instance is used
        **kwargs:
            See :class:`XYZWriter` for additional kwargs

        Returns
        -------
        :class:`XYZWriter` (see there for more details)

        See Also
        --------
        :class:`XYZWriter`
        """
        if n_atoms is None:
            n_atoms = self.n_atoms
        return XYZWriter(filename, n_atoms=n_atoms, **kwargs)

    def close(self):
        """Close xyz trajectory file if it was open."""
        if self.xyzfile is None:
            return
        self.xyzfile.close()
        self.xyzfile = None