# Source code for mdtraj.formats.xyzfile

    ##############################################################################
    # MDTraj: A Python Library for Loading, Saving, and Manipulating
    #         Molecular Dynamics Trajectories.
    # Copyright 2012-2015 Stanford University and the Authors
    #
    # Authors: Christoph Klein
    # Contributors: Robert T. McGibbon
    #
    # MDTraj is free software: you can redistribute it and/or modify
    # it under the terms of the GNU Lesser General Public License as
    # published by the Free Software Foundation, either version 2.1
    # of the License, or (at your option) any later version.
    #
    # This library is distributed in the hope that it will be useful,
    # but WITHOUT ANY WARRANTY; without even the implied warranty of
    # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    # GNU Lesser General Public License for more details.
    #
    # You should have received a copy of the GNU Lesser General Public
    # License along with MDTraj. If not, see <http://www.gnu.org/licenses/>.
    ##############################################################################

    import itertools
    import os
    from datetime import date

    import numpy as np

    import mdtraj
    from mdtraj.formats.registry import FormatRegistry
    from mdtraj.utils import cast_indices, ensure_type, in_units_of, open_maybe_zipped

    __all__ = ["XYZTrajectoryFile", "load_xyz"]

    class _EOF(IOError):
        pass

    @FormatRegistry.register_loader(".xyz")
    @FormatRegistry.register_loader(".xyz.gz")
    def load_xyz(filename, top=None, stride=None, atom_indices=None, frame=None):
        """Load a xyz trajectory file.

        While there is no universal standard for this format, this plugin adheres
        to the same format as the VMD plugin:

        http://www.ks.uiuc.edu/Research/vmd/plugins/molfile/xyzplugin.html

        Most notably, units are in angstroms and anything past the 'z' field is
        ignored.

        Parameters
        ----------
        filename : path-like
            Path of xyz trajectory file.
        top : {str, Trajectory, Topology}
            The xyz format does not contain topology information. Pass in
            either the path to a pdb file, a trajectory, or a topology to supply
            this information.
        stride : int, default=None
            Only read every stride-th frame
        atom_indices : array_like, optional
            If not none, then read only a subset of the atoms coordinates from the
            file.
        frame : int, optional
            Use this option to load only a single frame from a trajectory on disk.
            If frame is None, the default, the entire trajectory will be loaded.
            If supplied, ``stride`` will be ignored.

        Returns
        -------
        trajectory : md.Trajectory
            The resulting trajectory, as an md.Trajectory object.

        See Also
        --------
        mdtraj.XYZTrajectoryFile :  Low level interface to xyz files
        """
        from mdtraj.core.trajectory import _parse_topology

        # We make `top` required. Although this is a little weird, its good because
        # this function is usually called by a dispatch from load(), where top comes
        # from **kwargs. So if its not supplied, we want to give the user an
        # informative error message.
        if top is None:
            raise ValueError('"top" argument is required for load_xyz')

        if not isinstance(filename, (str, os.PathLike)):
            raise TypeError(
                "filename must be of type path-like for load_xyz. you supplied %s",
            )

        topology = _parse_topology(top)
        atom_indices = cast_indices(atom_indices)

        with XYZTrajectoryFile(filename) as f:
            if frame is not None:
                f.seek(frame)
                n_frames = 1
            else:
                n_frames = None
            return f.read_as_traj(
                topology,
                n_frames=n_frames,
                stride=stride,
                atom_indices=atom_indices,
            )

    @FormatRegistry.register_fileobject(".xyz")
    @FormatRegistry.register_fileobject(".xyz.gz")
    class XYZTrajectoryFile:
        """Interface for reading and writing to xyz files.

        This is a file-like object, that both reading or writing depending
        on the `mode` flag. It implements the context manager protocol,
        so you can also use it with the python 'with' statement.

        Parameters
        ----------
        filename : path-like
            The filename to open. A path to a file on disk.
        mode : {'r', 'w'}
            The mode in which to open the file, either 'r' for read or 'w' for
            write.
        force_overwrite : bool
            If opened in write mode, and a file by the name of `filename` already
            exists on disk, should we overwrite it?
        """

        distance_unit = "angstroms"

        def __init__(self, filename, mode="r", force_overwrite=True):
            """Open a xyz file for reading/writing."""
            self._is_open = False
            self._filename = filename
            self._mode = mode
            self._frame_index = 0
            self._n_frames = None
            # track which line we're on. this is not essential, but its useful
            # when reporting errors to the user to say what line it occured on.
            self._line_counter = 0

            if mode == "r":
                self._fh = open_maybe_zipped(filename, "r")
                self._is_open = True
            elif mode == "w":
                self._fh = open_maybe_zipped(filename, "w", force_overwrite)
                self._is_open = True
            else:
                raise ValueError(
                    f'mode must be one of "r" or "w". you supplied "{mode}"',
                )

        def close(self):
            """Close the xyz file."""
            if self._is_open:
                self._fh.close()
                self._is_open = False

        def __del__(self):
            self.close()

        def __enter__(self):
            """Support the context manager protocol."""
            return self

        def __exit__(self, *exc_info):
            """Support the context manager protocol."""
            self.close()

        def read_as_traj(self, topology, n_frames=None, stride=None, atom_indices=None):
            """Read a trajectory from a XYZ file

            Parameters
            ----------
            topology : Topology
                The system topology
            n_frames : int, optional
                If positive, then read only the next `n_frames` frames. Otherwise read all
                of the frames in the file.
            stride : np.ndarray, optional
                Read only every stride-th frame.
            atom_indices : array_like, optional
                If not none, then read only a subset of the atoms coordinates from the
                file. This may be slightly slower than the standard read because it required
                an extra copy, but will save memory.

            Returns
            -------
            trajectory : Trajectory
                A trajectory object containing the loaded portion of the file.
            """
            from mdtraj.core.trajectory import Trajectory

            if atom_indices is not None:
                topology = topology.subset(atom_indices)

            initial = int(self._frame_index)
            xyz = self.read(n_frames=n_frames, stride=stride, atom_indices=atom_indices)
            if len(xyz) == 0:
                return Trajectory(xyz=np.zeros((0, topology.n_atoms, 3)), topology=topology)

            in_units_of(xyz, self.distance_unit, Trajectory._distance_unit, inplace=True)

            if stride is None:
                stride = 1
            time = (stride * np.arange(len(xyz))) + initial
            return Trajectory(xyz=xyz, topology=topology, time=time)

        def read(self, n_frames=None, stride=None, atom_indices=None):
            """Read data from a xyz file.

            Parameters
            ----------
            n_frames : int, None
                The number of frames you would like to read from the file.
                If None, all of the remaining frames will be loaded.
            stride : np.ndarray, optional
                Read only every stride-th frame.
            atom_indices : array_like, optional
                If not none, then read only a subset of the atoms coordinates
                from the file.

            Returns
            -------
            xyz : np.ndarray, shape=(n_frames, n_atoms, 3), dtype=np.float32
            """
            if not self._mode == "r":
                raise ValueError(
                    'read() is only available when file is opened in mode="r"',
                )

            if n_frames is None:
                frame_counter = itertools.count()
            else:
                frame_counter = range(n_frames)

            if stride is None:
                stride = 1

            all_coords = []
            for i in frame_counter:
                try:
                    frame_coords = self._read()
                    if atom_indices is not None:
                        frame_coords = frame_coords[atom_indices, :]
                except _EOF:
                    break

                all_coords.append(frame_coords)

                for j in range(stride - 1):
                    # throw away these frames
                    try:
                        self._read()
                    except _EOF:
                        break

            all_coords = np.array(all_coords)
            return all_coords

        def _read(self):
            """Read a single frame."""

            first = self._fh.readline()  # Number of atoms.
            if first == "":
                raise _EOF()
            else:
                self._n_atoms = int(first)
            self._fh.readline()  # Comment line.
            self._line_counter += 2

            xyz = np.empty(shape=(self._n_atoms, 3))
            types = np.empty(shape=self._n_atoms, dtype=str)

            for i in range(self._n_atoms):
                line = self._fh.readline()
                if line == "":
                    raise _EOF()
                split_line = line.split()
                try:
                    types[i] = split_line[0]
                    xyz[i] = [float(x) for x in split_line[1:4]]
                except Exception:
                    raise OSError(
                        f'xyz parse error on line {self._line_counter:d} of "{self._filename:s}". '
                        "This file does not appear to be a valid "
                        "xyz file.",
                    )
                self._line_counter += 1
            # --- end body ---

            self._frame_index += 1
            return xyz

        def write(self, xyz, types=None):
            """Write one or more frames of data to a xyz file.

            Parameters
            ----------
            xyz : np.ndarray, shape=(n_frames, n_atoms, 3)
                The cartesian coordinates of the atoms to write. By convention for
                this trajectory format, the lengths should be in units of angstroms.
            types : np.ndarray, shape(3, )
                The type of each particle.
            """

            if not self._mode == "w":
                raise ValueError(
                    'write() is only available when file is opened in mode="w"',
                )

            if not types:
                # Make all particles the same type.
                types = ["X" for _ in range(xyz.shape[1])]
            xyz = ensure_type(
                xyz,
                np.float32,
                3,
                "xyz",
                can_be_none=False,
                shape=(None, None, 3),
                warn_on_cast=False,
                add_newaxis_on_deficient_ndim=True,
            )

            for i in range(xyz.shape[0]):
                self._fh.write(f"{xyz.shape[1]}\n")
                self._fh.write(
                    f"Created with MDTraj {mdtraj.__version__}, {str(date.today())}\n",
                )

                for j, coord in enumerate(xyz[i]):
                    self._fh.write(
                        f"{types[j]} {coord[0]:8.3f} {coord[1]:8.3f} {coord[2]:8.3f}\n",
                    )

        def seek(self, offset, whence=0):
            """Move to a new file position.

            Parameters
            ----------
            offset : int
                A number of frames.
            whence : {0, 1, 2}
                0: offset from start of file, offset should be >=0.
                1: move relative to the current position, positive or negative
                2: move relative to the end of file, offset should be <= 0.
                Seeking beyond the end of a file is not supported
            """
            if self._mode == "r":
                advance, absolute = None, None
                if whence == 0 and offset >= 0:
                    if offset >= self._frame_index:
                        advance = offset - self._frame_index
                    else:
                        absolute = offset
                elif whence == 1 and offset >= 0:
                    advance = offset
                elif whence == 1 and offset < 0:
                    absolute = offset + self._frame_index
                elif whence == 2 and offset <= 0:
                    raise NotImplementedError("offsets from the end are not supported yet")
                else:
                    raise OSError("Invalid argument")

                if advance is not None:
                    for i in range(advance):
                        self._read()  # advance and throw away these frames
                elif absolute is not None:
                    self._fh.close()
                    self._fh = open(self._filename)
                    self._frame_index = 0
                    self._line_counter = 0
                    for i in range(absolute):
                        self._read()
                else:
                    raise RuntimeError()

            else:
                raise NotImplementedError("offsets in write mode are not supported yet")

        def tell(self):
            """Current file position.

            Returns
            -------
            offset : int
                The current frame in the file.
            """
            return int(self._frame_index)

        def __len__(self):
            """Number of frames in the file."""
            if str(self._mode) != "r":
                raise NotImplementedError('len() only available in mode="r" currently')
            if not self._is_open:
                raise ValueError("I/O operation on closed file")
            if self._n_frames is None:
                with open(self._filename) as fh:
                    n_atoms = int(fh.readline())
                    self._n_frames = (sum(1 for line in fh) + 1) // (n_atoms + 2)
            return self._n_frames