# Source code for MDAnalysis.coordinates.chemfiles

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