# -*- 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.
#
# 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
#
"""
Reading trajectories with *chemfiles* --- :mod:`MDAnalysis.coordinates.chemfiles`
=================================================================================

MDAnalysis interoperates with the `chemfiles`_ library. The *chemfiles* C++ library
supports an expanding set of file formats, some of which are not natively supported by
MDAnalysis. Using the *CHEMFILES* reader you can use  `chemfiles`_ for the low-level
file reading. Check the list of `chemfile-supported file formats`_.

.. _chemfiles: https://chemfiles.org
.. _chemfile-supported file formats: https://chemfiles.org/chemfiles/latest/formats.html#list-of-supported-formats

Using the CHEMFILES reader
--------------------------

When reading, set the ``format="CHEMFILES"`` keyword argument and I/O is delegated to
`chemfiles`_. For example::

   >>> import MDAnalysis as mda
   >>> from MDAnalysis.tests import datafiles as data
   >>> u = mda.Universe(data.TPR, data.TRR, format="CHEMFILES")
   >>> print(u.trajectory)
   <ChemfilesReader ~/anaconda3/envs/mda3/lib/python3.8/site-packages/MDAnalysisTests/data/adk_oplsaa.trr with 10 frames of 47681 atoms>

You can then use the :class:`~MDAnalysis.core.universe.Universe` as usual while chemfiles
is handling the I/O transparently in the background.

`chemfiles`_ can also *write* a number of formats for which there are no Writers in
MDAnalysis. For example, to write a mol2 file::

   >>> u = mda.Universe(data.mol2_ligand)
   >>> with mda.Writer("ligand.mol2", format="CHEMFILES") as W:
   ...     W.write(u.atoms)




Classes
-------

Classes to read and write files using the `chemfiles`_ library. This library
provides C++ implementation of multiple formats readers and writers.

.. autoclass:: ChemfilesReader

.. autoclass:: ChemfilesWriter

.. autoclass:: ChemfilesPicklable

Helper functions
----------------

.. autodata:: MIN_CHEMFILES_VERSION
.. autodata:: MAX_CHEMFILES_VERSION
.. autofunction:: check_chemfiles_version

"""
import numpy as np
import warnings
from packaging.version import Version

from . import base
from ..lib.util import store_init_arguments

try:
    import chemfiles
except ImportError:
    HAS_CHEMFILES = False

    # Allow building documentation even if chemfiles is not installed
    import types

    class MockTrajectory:
        pass

    chemfiles = types.ModuleType("chemfiles")
    chemfiles.Trajectory = MockTrajectory
else:
    HAS_CHEMFILES = True


#: Lowest version of chemfiles that is supported by MDAnalysis.
MIN_CHEMFILES_VERSION = Version("0.10")
#: Lowest version of chemfiles that is *not supported* by MDAnalysis.
MAX_CHEMFILES_VERSION = Version("0.11")


def check_chemfiles_version():
    """Check if an appropriate *chemfiles* is available

    Returns ``True`` if a usable chemfiles version is available,
    with :data:`MIN_CHEMFILES_VERSION` <= version <
    :data:`MAX_CHEMFILES_VERSION`

    .. versionadded:: 1.0.0
    """
    if not HAS_CHEMFILES:
        warnings.warn(
            "chemfiles package not found, "
            "try installing it with 'pip install chemfiles'"
        )
        return False
    version = Version(chemfiles.__version__)
    return version >= MIN_CHEMFILES_VERSION and version < MAX_CHEMFILES_VERSION


class ChemfilesReader(base.ReaderBase):
    """Coordinate reader using chemfiles.

    The file format to used is guessed based on the file extension. If no
    matching format is found, a ``ChemfilesError`` is raised. It is also
    possible to manually specify the format to use for a given file.

    .. versionadded:: 1.0.0
    """

    format = "chemfiles"
    units = {"time": "fs", "length": "Angstrom"}

    @store_init_arguments
    def __init__(self, filename, chemfiles_format="", **kwargs):
        """
        Parameters
        ----------
        filename : chemfiles.Trajectory or str
            the chemfiles object to read or filename to read
        chemfiles_format : str (optional)
            if *filename* was a string, use the given format name instead of
            guessing from the extension. The `list of supported formats
            <formats>`_ and the associated names is available in the chemfiles
            documentation.
        **kwargs : dict
            General reader arguments.

        """
        if not check_chemfiles_version():
            raise RuntimeError(
                "Please install Chemfiles > {}"
                "".format(MIN_CHEMFILES_VERSION)
            )
        super(ChemfilesReader, self).__init__(filename, **kwargs)
        self._format = chemfiles_format
        self._cached_n_atoms = None
        self._open()
        self.ts = self._Timestep(self.n_atoms, **self._ts_kwargs)
        self.next()

    @staticmethod
    def _format_hint(thing):
        """Can this Reader read *thing*"""
        # nb, filename strings can still get passed through if
        # format='CHEMFILES' is used
        return HAS_CHEMFILES and isinstance(thing, chemfiles.Trajectory)

    def _open(self):
        self._closed = False
        self._step = 0
        self._frame = -1
        if isinstance(self.filename, chemfiles.Trajectory):
            self._file = self.filename
        else:
            self._file = ChemfilesPicklable(self.filename, "r", self._format)

    def close(self):
        """close reader"""
        if not self._closed:
            self._file.close()
            self._closed = True

    @property
    def n_frames(self):
        """number of frames in trajectory"""
        return self._file.nsteps

    @property
    def n_atoms(self):
        """number of atoms in the first frame of the trajectory"""
        if self._cached_n_atoms is None:
            self._cached_n_atoms = len(self._file.read_step(0).atoms)
        return self._cached_n_atoms

    def _reopen(self):
        """reopen trajectory"""
        self.close()
        self._open()

    def _read_frame(self, i):
        """read frame i"""
        self._step = i
        return self._read_next_timestep()

    def _read_next_timestep(self, ts=None):
        """copy next frame into timestep"""
        if self._step >= self.n_frames:
            raise IOError("trying to go over trajectory limit")
        if ts is None:
            ts = self.ts
        self.ts = ts
        frame = self._file.read_step(self._step)
        self._frame_to_ts(frame, ts)
        ts.frame = frame.step
        self._step += 1
        return ts

    def _frame_to_ts(self, frame, ts):
        """convert a chemfiles frame to a :class:`TimeStep`"""
        if len(frame.atoms) != self.n_atoms:
            raise IOError(
                "The number of atom changed in the trajectory. "
                "This is not supported by MDAnalysis."
            )

        ts.dimensions = np.r_[frame.cell.lengths, frame.cell.angles]
        ts.positions[:] = frame.positions[:]
        if frame.has_velocities():
            ts.has_velocities = True
            ts.velocities[:] = frame.velocities[:]

    def Writer(self, filename, n_atoms=None, **kwargs):
        """Return writer for trajectory format"""
        if n_atoms is None:
            n_atoms = self.n_atoms
        return ChemfilesWriter(filename, n_atoms, **kwargs)


class ChemfilesWriter(base.WriterBase):
    """
    Coordinate writer using chemfiles.

    The file format to used is guessed based on the file extension. If no
    matching format is found, a ``ChemfilesError`` is raised. It is also
    possible to manually specify the format to use for a given file.

    Chemfiles support writting to files with varying number of atoms if the
    underlying format support it. This is the case of most of text-based
    formats.

    .. versionadded:: 1.0.0
    """

    format = "chemfiles"
    multiframe = True

    # chemfiles mostly[1] uses these units for the in-memory representation,
    # and convert into the file format units when writing.
    #
    # [1] mostly since some format don't have a specified unit
    # (XYZ for example), so then chemfiles just assume they are in A and fs.
    units = {"time": "fs", "length": "Angstrom"}

    def __init__(
        self,
        filename,
        n_atoms=0,
        mode="w",
        chemfiles_format="",
        topology=None,
        **kwargs,
    ):
        """
        Parameters
        ----------
        filename: str
            filename of trajectory file.
        n_atoms: int
            number of atoms in the trajectory to write. This value is not
            used and can vary during trajectory, if the underlying format
            support it
        mode : str (optional)
            file opening mode: use 'a' to append to an existing file or 'w' to
            create a new file
        chemfiles_format : str (optional)
            use the given format name instead of guessing from the extension.
            The `list of supported formats <formats>`_ and the associated names
            is available in chemfiles documentation.
        topology : Universe or AtomGroup (optional)
            use the topology from this :class:`~MDAnalysis.core.groups.AtomGroup`
            or :class:`~MDAnalysis.core.universe.Universe` to write all the
            timesteps to the file
        **kwargs : dict
            General writer arguments.

        """
        if not check_chemfiles_version():
            raise RuntimeError(
                "Please install Chemfiles > {}"
                "".format(MIN_CHEMFILES_VERSION)
            )
        self.filename = filename
        self.n_atoms = n_atoms
        if mode != "a" and mode != "w":
            raise IOError("Expected 'a' or 'w' as mode in ChemfilesWriter")
        self._file = chemfiles.Trajectory(
            self.filename, mode, chemfiles_format
        )
        self._closed = False
        if topology is not None:
            if isinstance(topology, str):
                self._file.set_topology(topology)
            else:
                topology = self._topology_to_chemfiles(topology, n_atoms)
                self._file.set_topology(topology)

    def close(self):
        """Close the trajectory file and finalize the writing"""
        if hasattr(self, "_closed") and not self._closed:
            self._file.close()
            self._closed = True

    def _write_next_frame(self, obj):
        """Write information associated with ``obj`` at current frame into
        trajectory.

        Topology for the output is taken from the `obj` or default to the value
        of the `topology` keyword supplied to the :class:`ChemfilesWriter`
        constructor.

        If `obj` contains velocities, and the underlying format supports it, the
        velocities are written to the file. Writing forces is unsupported at the
        moment.

        Parameters
        ----------
        obj : AtomGroup or Universe
            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.
        """
        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

        frame = self._timestep_to_chemfiles(ts)
        frame.topology = self._topology_to_chemfiles(obj, len(frame.atoms))

        self._file.write(frame)

    def _timestep_to_chemfiles(self, ts):
        """
        Convert a Timestep to a chemfiles Frame
        """
        # TODO: CONVERTERS?
        frame = chemfiles.Frame()
        frame.resize(ts.n_atoms)
        if ts.has_positions:
            frame.positions[:] = ts.positions[:]
        if ts.has_velocities:
            frame.add_velocities()
            frame.velocities[:] = ts.velocities[:]

        # if there is no cell information in this universe, still pass a valid
        # cell to chemfiles
        if ts.dimensions is None:
            lengths, angles = np.zeros(3), np.array([90, 90, 90.0])
        else:
            lengths = ts.dimensions[:3]
            angles = ts.dimensions[3:]

        frame.cell = chemfiles.UnitCell(lengths, angles)

        return frame

    def _topology_to_chemfiles(self, obj, n_atoms):
        """
        Convert an AtomGroup or an Universe to a chemfiles Topology
        """
        topology = chemfiles.Topology()
        if not hasattr(obj, "atoms"):
            # use an empty topology
            topology.resize(n_atoms)
            return topology

        # (1) add all atoms to the topology
        residues = {}
        for atom in obj.atoms:
            name = getattr(atom, "name", "")
            type = getattr(atom, "type", name)
            chemfiles_atom = chemfiles.Atom(name, type)

            if hasattr(atom, "altLoc"):
                chemfiles_atom["altloc"] = str(atom.altLoc)

            if hasattr(atom, "segid"):
                chemfiles_atom["segid"] = str(atom.segid)

            if hasattr(atom, "segindex"):
                chemfiles_atom["segindex"] = int(atom.segindex)

            if hasattr(atom, "resid"):
                resname = getattr(atom, "resname", "")
                if atom.resid not in residues.keys():
                    residues[atom.resid] = chemfiles.Residue(
                        resname, atom.resid
                    )
                residue = residues[atom.resid]

                atom_idx = len(topology.atoms)
                residue.atoms.append(atom_idx)

                if hasattr(atom, "record_type"):
                    # set corresponding chemfiles residue property
                    if atom.record_type == "ATOM":
                        residue["is_standard_pdb"] = True
                    else:
                        residue["is_standard_pdb"] = False

            topology.atoms.append(chemfiles_atom)

        # (2) add residues to the topology
        for residue in residues.values():
            topology.residues.append(residue)

        # (3) add bonds to the topology
        for bond in getattr(obj, "bonds", []):
            topology.add_bond(bond.atoms[0].ix, bond.atoms[1].ix)

        return topology


class ChemfilesPicklable(chemfiles.Trajectory):
    """Chemfiles file object (read-only) that can be pickled.

    This class provides a file-like object (as returned by
    :class:`chemfiles.Trajectory`) that,
    unlike standard Python file objects,
    can be pickled. Only read mode is supported.

    When the file is pickled, path, mode, and format of the file handle
    are saved. On unpickling, the file is opened by path with mode,
    and saved format.
    This means that for a successful unpickle, the original file still has
    to be accessible with its filename.

    Note
    ----
    Can only be used with reading ('r') mode.
    Upon pickling, the current frame is reset. `universe.trajectory[i]` has
    to be used to return to its original frame.

    Parameters
    ----------
    filename : str
        a filename given a text or byte string.
    mode : 'r' , optional
        only 'r' can be used for pickling.
    format : '', optional
        guessed from the file extension if empty.

    Example
    -------
    ::

        f = ChemfilesPicklable(XYZ, 'r', '')
        print(f.read())
        f.close()

    can also be used as context manager::

        with ChemfilesPicklable(XYZ) as f:
            print(f.read())

    See Also
    ---------
    :class:`MDAnalysis.lib.picklable_file_io.FileIOPicklable`
    :class:`MDAnalysis.lib.picklable_file_io.BufferIOPicklable`
    :class:`MDAnalysis.lib.picklable_file_io.TextIOPicklable`
    :class:`MDAnalysis.lib.picklable_file_io.GzipPicklable`
    :class:`MDAnalysis.lib.picklable_file_io.BZ2Picklable`


    .. versionadded:: 2.0.0
    """

    def __init__(self, path, mode="r", format=""):
        if mode != "r":
            raise ValueError("Only read mode ('r') files can be pickled.")
        super().__init__(path=path, mode=mode, format=format)

    def __getstate__(self):
        return self.path, self._Trajectory__mode, self._Trajectory__format

    def __setstate__(self, args):
        self.__init__(*args)