##############################################################################
# MDTraj: A Python Library for Loading, Saving, and Manipulating
#         Molecular Dynamics Trajectories.
# Copyright 2012-2014 Stanford University and the Authors
#
# Authors: Jason Swails
# Contributors:
#
# This code for reading Amber restart and inpcrd files was taken from ParmEd,
# which is released under the GNU Lesser General Public License
#
# 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/>.
##############################################################################

"""
This module provides the ability to read Amber inpcrd/restart files as well as
Amber NetCDF restart files. This code was taken from ParmEd and simplified by
removing the functionality that is not needed.
"""

import os
import warnings
from math import ceil

import numpy as np
from packaging.version import Version

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

__all__ = [
    "AmberRestartFile",
    "load_restrt",
    "AmberNetCDFRestartFile",
    "load_ncrestrt",
]


@FormatRegistry.register_loader(".rst7")
@FormatRegistry.register_loader(".restrt")
@FormatRegistry.register_loader(".inpcrd")
def load_restrt(filename, top=None, atom_indices=None):
    """Load an AMBER ASCII restart/inpcrd file. Since this file doesn't contain
    information to specify the topology, you need to supply a topology

    Parameters
    ----------
    filename : path-like
        name of the AMBER restart file
    top : {str, Trajectory, Topology}
        Pass in either the path to a file containing topology information (e.g.,
        a PDB, an AMBER prmtop, or certain types of Trajectory objects) to
        supply the necessary topology information that is not present in these
        files
    atom_indices : array_like, optional
        If not None, then read only a subset of the atoms coordinates from the
        file.

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

    See Also
    --------
    mdtraj.AmberRestartFile : Low level interface to AMBER restart files
    """
    from mdtraj.core.trajectory import _parse_topology

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

    with AmberRestartFile(filename) as f:
        return f.read_as_traj(topology, atom_indices=atom_indices)


@FormatRegistry.register_fileobject(".rst7")
@FormatRegistry.register_fileobject(".restrt")
@FormatRegistry.register_fileobject(".inpcrd")
class AmberRestartFile:
    """Interface for reading and writing AMBER ASCII restart files. This is a
    file-like object, that supports both reading and 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 name of the file to open
    mode : {'r', 'w'}, default='r'
        The mode in which to open the file. Valid options are 'r' or 'w' for
        'read' or 'write'
    force_overwrite : bool, default=False
        In write mode, if a file named `filename` already exists, clobber it and
        overwrite it

    See Also
    --------
    md.AmberNetCDFRestartFile : Low level interface to AMBER NetCDF-format restart files
    """

    distance_unit = "angstroms"

    def __init__(self, filename, mode="r", force_overwrite=True):
        self._closed = True
        self._mode = mode
        self._filename = filename

        if mode not in ("r", "w"):
            raise ValueError("mode must be one of ['r', 'w']")

        if mode == "w" and not force_overwrite and os.path.exists(filename):
            raise OSError('"%s" already exists' % filename)

        if mode == "w":
            self._needs_initialization = True
            self._handle = open(filename, mode)
            self._closed = False
        elif mode == "r":
            with open(filename, mode) as f:
                f.readline()
                words = f.readline().split()
                try:
                    self._n_atoms = int(words[0])
                except (IndexError, ValueError):
                    raise TypeError(
                        '"%s" is not a recognized Amber restart' % filename,
                    )
            self._needs_initialization = False
        else:
            raise RuntimeError()

    @property
    def n_atoms(self):
        self._validate_open()
        if self._needs_initialization:
            raise OSError("The file is uninitialized")
        return self._n_atoms

    @property
    def n_frames(self):
        return 1  # always 1 frame

    def _validate_open(self):
        if self._closed:
            raise OSError("The file is closed.")

    def _parse(self, lines):
        """Parses the file"""
        self._time = None
        try:
            words = lines[1].split()
            self._n_atoms = natom = int(words[0])
        except (IndexError, ValueError):
            raise TypeError("not a recognized Amber restart")

        time = None
        if len(words) >= 2:
            time = float(words[1])

        # variable _ throughout was hasvels, but velocies are not tracked
        lines_per_frame = int(ceil(natom / 2))
        if len(lines) == lines_per_frame + 2:
            hasbox = _ = False
        elif natom in (1, 2) and len(lines) == 4:
            # This is the _only_ case where line counting does not work -- there
            # is either 1 or 2 atoms and there are 4 lines. The 1st 3 lines are
            # the title, natom/time, and coordinates. The 4th are almost always
            # velocities since it's hard to have a periodic system this small.
            # However, velocities (which are scaled down by 20.445) have a ~0%
            # chance of being 60+, so we can pretty easily tell if the last line
            # has box dimensions and angles or velocities. I cannot envision a
            # plausible scenario where the detection here will ever fail
            line = lines[3]
            if natom == 1:
                tmp = [line[i : i + 12] for i in range(0, 72, 12) if line[i : i + 12].strip()]
                if len(tmp) == 3:
                    _ = True
                    hasbox = False
                elif len(tmp) == 6:
                    hasbox = True
                    _ = False
                else:
                    raise TypeError("not a recognized Amber restart")
            else:
                # Ambiguous case
                tmp = [float(line[i : i + 12]) >= 60.0 for i in range(0, 72, 12)]
                if any(tmp):
                    hasbox = True
                    _ = False
                else:
                    _ = True
                    hasbox = False
        elif len(lines) == lines_per_frame + 3:
            hasbox = True
            _ = False
        elif len(lines) == 2 * lines_per_frame + 2:
            hasbox = False
            _ = True
        elif len(lines) == 2 * lines_per_frame + 3:
            hasbox = _ = True
        else:
            raise TypeError(
                "Badly formatted restart file. Has %d lines for %d atoms" % (len(lines), natom),
            )

        coordinates = np.zeros((1, natom, 3))
        if time is None:
            time = np.zeros(1)
        else:
            time = np.asarray((time,))

        # Fill the coordinates
        for i in range(lines_per_frame):
            line = lines[i + 2]  # Skip first two lines
            i2 = i * 2
            coordinates[0, i2, :] = [float(line[j : j + 12]) for j in range(0, 36, 12)]
            i2 += 1
            if i2 < natom:
                coordinates[0, i2, :] = [float(line[j : j + 12]) for j in range(36, 72, 12)]
        if hasbox:
            cell_lengths = np.zeros((1, 3))
            cell_angles = np.zeros((1, 3))
            line = lines[-1]
            cell_lengths[0, :] = [float(line[i : i + 12]) for i in range(0, 36, 12)]
            cell_angles[0, :] = [float(line[i : i + 12]) for i in range(36, 72, 12)]
        else:
            cell_lengths = cell_angles = None

        return coordinates, time, cell_lengths, cell_angles

    def read_as_traj(self, topology, atom_indices=None):
        """Read an AMBER ASCII restart file as a trajectory.

        Parameters
        ----------
        topology : Topology
            The system topology
        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 with 1 frame created from the file.
        """
        from mdtraj.core.trajectory import Trajectory

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

        xyz, time, cell_lengths, cell_angles = self.read(atom_indices=atom_indices)
        xyz = in_units_of(
            xyz,
            self.distance_unit,
            Trajectory._distance_unit,
            inplace=True,
        )
        cell_lengths = in_units_of(
            cell_lengths,
            self.distance_unit,
            Trajectory._distance_unit,
            inplace=True,
        )

        return Trajectory(
            xyz=xyz,
            topology=topology,
            time=time,
            unitcell_lengths=cell_lengths,
            unitcell_angles=cell_angles,
        )

    def read(self, atom_indices=None):
        """Read data from an AMBER ASCII restart file

        Parameters
        ----------
        atom_indices : np.ndarray, dtype=int, optional
            The specific indices of the atoms you'd like to retrieve. If not
            supplied, all of the atoms will be retrieved.

        Returns
        -------
        coordinates : np.ndarray, shape=(1, n_atoms, 3)
            The cartesian coordinates of the atoms, in units of angstroms. These
            files only ever contain 1 frame
        time : np.ndarray, None
            The time corresponding to the frame, in units of picoseconds, or
            None if no time information is present
        cell_lengths : np.ndarray, None
            The lengths (a, b, c) of the unit cell for the frame in angstroms,
            or None if the information is not present in the file
        cell_angles : np.ndarray, None
            The angles (\alpha, \beta, \\gamma) defining the unit cell for each
            frame, or None if the information is not present in the file.
        """
        if self._mode != "r":
            raise OSError(
                "The file was opened in mode=%s. Reading is not allowed." % self._mode,
            )

        with open(self._filename) as f:
            lines = f.readlines()

        coordinates, time, cell_lengths, cell_angles = self._parse(lines)

        if atom_indices is not None:
            atom_slice = ensure_type(
                atom_indices,
                dtype=int,
                ndim=1,
                name="atom_indices",
                warn_on_cast=False,
            )
            if not np.all(atom_slice) >= 0:
                raise ValueError("Entries in atom_slice must be >= 0")
            coordinates = coordinates[:, atom_slice, :]

        return coordinates, time, cell_lengths, cell_angles

    def write(
        self,
        coordinates,
        time=None,
        cell_lengths=None,
        cell_angles=None,
    ):
        """Write one frame of a MD trajectory to disk in the AMBER ASCII restart
        file format.

        Parameters
        ----------
        coordinates : np.ndarray, dtype=np.float32, shape=([1,] n_atoms, 3)
            The cartesian coordinates of each atom, in units of angstroms. Must
            be only a single frame (shape can be (1,N,3) or (N,3) where N is
            the number of atoms)
        time : array-like with 1 element or float, optional
            The time corresponding to this frame. If not specified, a place
            holder of 0 will be written
        cell_lengths : np.ndarray, dtype=np.double, shape=([1,] 3)
            The lengths (a,b,c) of the unit cell for the frame in Angstroms
        cell_angles : np.ndarray, dtype=np.double, shape=([1,] 3)
            The angles between the unit cell vectors for the frame in Degrees
        """
        if self._mode != "w":
            raise OSError(
                "The file was opened in mode=%s. Writing not allowed." % self._mode,
            )
        if not self._needs_initialization:
            # Must have already been written -- can only write once
            raise RuntimeError(
                "restart file has already been written -- can only write one frame to restart files.",
            )
        # These are no-ops.
        # coordinates = in_units_of(coordinates, None, 'angstroms')
        # time = in_units_of(time, None, 'picoseconds')
        # cell_lengths = in_units_of(cell_lengths, None, 'angstroms')
        # cell_angles = in_units_of(cell_angles, None, 'degrees')

        # typecheck all of the input arguments rigorously
        coordinates = ensure_type(
            coordinates,
            np.float32,
            3,
            "coordinates",
            length=None,
            can_be_none=False,
            shape=(1, None, 3),
            warn_on_cast=False,
            add_newaxis_on_deficient_ndim=True,
        )
        n_frames, self._n_atoms = coordinates.shape[0], coordinates.shape[1]
        if n_frames != 1:
            raise ValueError("Can only write 1 frame to a restart file!")
        if time is not None:
            try:
                if isinstance(time, np.ndarray):
                    time = float(time.item())
                else:
                    time = float(time)
            except TypeError:
                raise TypeError("Can only provide a single time")
        else:
            time = 0.0
        cell_lengths = ensure_type(
            cell_lengths,
            np.float64,
            2,
            "cell_lengths",
            length=1,
            can_be_none=True,
            warn_on_cast=False,
            add_newaxis_on_deficient_ndim=True,
        )
        cell_angles = ensure_type(
            cell_angles,
            np.float64,
            2,
            "cell_angles",
            length=1,
            can_be_none=True,
            warn_on_cast=False,
            add_newaxis_on_deficient_ndim=True,
        )
        if (cell_lengths is None and cell_angles is not None) or (cell_lengths is not None and cell_angles is None):
            prov, negl = "cell_lengths", "cell_angles"
            if cell_lengths is None:
                prov, negl = negl, prov
            raise ValueError(
                f'You provided the variable "{prov}" but did not '
                f'provide "{negl}". Either provide both or neither -- '
                "one without the other is meaningless.",
            )

        self._handle.write(
            "Amber restart file (without velocities) written by MDTraj\n",
        )
        self._handle.write("%5d%15.7e\n" % (self._n_atoms, time))
        fmt = "%12.7f%12.7f%12.7f"
        for i in range(self._n_atoms):
            acor = coordinates[0, i, :]
            self._handle.write(fmt % (acor[0], acor[1], acor[2]))
            if i % 2 == 1:
                self._handle.write("\n")
        if self._n_atoms % 2 == 1:
            self._handle.write("\n")
        if cell_lengths is not None:
            self._handle.write(
                fmt
                % (
                    cell_lengths[0, 0],
                    cell_lengths[0, 1],
                    cell_lengths[0, 2],
                ),
            )
            self._handle.write(
                fmt
                % (
                    cell_angles[0, 0],
                    cell_angles[0, 1],
                    cell_angles[0, 2],
                )
                + "\n",
            )
        self._handle.flush()

    def __enter__(self):
        return self

    def __exit__(self, *exc_info):
        self.close()

    def close(self):
        if not self._closed and hasattr(self, "_handle"):
            self._handle.close()
        self._closed = True

    def __del__(self):
        self.close()

    def __len__(self):
        return 1  # All restarts have only 1 frame


@FormatRegistry.register_loader(".ncrst")
def load_ncrestrt(filename, top=None, atom_indices=None):
    """Load an AMBER NetCDF restart/inpcrd file. Since this file doesn't
    contain information to specify the topology, you need to supply a topology

    Parameters
    ----------
    filename : path-like
        name of the AMBER restart file
    top : {path-like, Trajectory, Topology}
        Pass in either the path to a file containing topology information (e.g.,
        a PDB, an AMBER prmtop, or certain types of Trajectory objects) to
        supply the necessary topology information that is not present in these
        files
    atom_indices : array_like, optional
        If not None, then read only a subset of the atoms coordinates from the
        file.

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

    See Also
    --------
    mdtraj.AmberRestartFile : Low level interface to AMBER restart files
    """
    from mdtraj.core.trajectory import _parse_topology

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

    with AmberNetCDFRestartFile(filename) as f:
        return f.read_as_traj(topology, atom_indices=atom_indices)


@FormatRegistry.register_fileobject(".ncrst")
class AmberNetCDFRestartFile:
    """Interface for reading and writing AMBER NetCDF files. This is a file-like
    object, that supports both reading and 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 name of the file to open
    mode : {'r', 'w'}, default='r'
        The mode in which to open the file. Valid options are 'r' or 'w' for
        'read' or 'write'
    force_overwrite : bool, default=False
        In write mode, if a file named `filename` already exists, clobber it and
        overwrite it
    """

    distance_unit = "angstroms"

    def __init__(self, filename, mode="r", force_overwrite=False):
        self._closed = True
        self._mode = mode
        if Version(import_("scipy.version").short_version) < Version("0.12.0"):
            raise ImportError(
                "MDTraj NetCDF support requires scipy>=0.12.0. You have %s" % import_("scipy.version").short_version,
            )
        netcdf = import_("scipy.io").netcdf_file

        if mode not in ("r", "w"):
            raise ValueError("mode must be one of ['r', 'w']")

        if mode == "w" and not force_overwrite and os.path.exists(filename):
            raise OSError('"%s" already exists' % filename)

        # AMBER uses the NetCDF3 format, with 64 bit encodings, which for
        # scipy.io.netcdf_file is "version=2"
        self._handle = netcdf(filename, mode=mode, version=2)
        self._closed = False
        if mode == "w":
            self._needs_initialization = True
        elif mode == "r":
            self._needs_initialization = False
        else:
            raise RuntimeError()

    @property
    def n_atoms(self):
        self._validate_open()
        if self._needs_initialization:
            raise OSError("The file is uninitialized")
        return self._handle.dimensions["atom"]

    @property
    def n_frames(self):
        return 1  # always 1 frame

    def _validate_open(self):
        if self._closed:
            raise OSError("The file is closed.")

    def read_as_traj(self, topology, atom_indices=None):
        """Read an AMBER ASCII restart file as a trajectory.

        Parameters
        ----------
        topology : Topology
            The system topology
        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 with 1 frame created from the file.
        """
        from mdtraj.core.trajectory import Trajectory

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

        xyz, time, cell_lengths, cell_angles = self.read(atom_indices=atom_indices)
        xyz = in_units_of(
            xyz,
            self.distance_unit,
            Trajectory._distance_unit,
            inplace=True,
        )
        cell_lengths = in_units_of(
            cell_lengths,
            self.distance_unit,
            Trajectory._distance_unit,
            inplace=True,
        )

        return Trajectory(
            xyz=xyz,
            topology=topology,
            time=time,
            unitcell_lengths=cell_lengths,
            unitcell_angles=cell_angles,
        )

    def read(self, atom_indices=None):
        """Read data from an AMBER NetCDF restart file

        Parameters
        ----------
        atom_indices : np.ndarray, dtype=int, optional
            The specific indices of the atoms you'd like to retrieve. If not
            supplied, all of the atoms will be retrieved.

        Returns
        -------
        coordinates : np.ndarray, shape=(1, n_atoms, 3)
            The cartesian coordinates of the atoms, in units of angstroms. These
            files only ever contain 1 frame
        time : np.ndarray, None
            The time corresponding to the frame, in units of picoseconds, or
            None if no time information is present
        cell_lengths : np.ndarray, None
            The lengths (a, b, c) of the unit cell for the frame in angstroms,
            or None if the information is not present in the file
        cell_angles : np.ndarray, None
            The angles (\alpha, \beta, \\gamma) defining the unit cell for each
            frame, or None if the information is not present in the file.

        Notes
        -----
        If the file is not a NetCDF file with the appropriate convention, a
        TypeError is raised. If variables that are needed do not exist or if
        illegal values are passed in for parameters, ValueError is raised. If
        I/O errors occur, IOError is raised.
        """
        if self._mode != "r":
            raise OSError(
                "The file was opened in mode=%s. Reading is not allowed." % self._mode,
            )

        if self._closed:
            raise OSError("The file has been closed!")

        if "coordinates" not in self._handle.variables:
            raise ValueError("No coordinates found in the NetCDF file.")

        # Check that conventions are correct
        try:
            conventions = self._handle.Conventions.decode("ascii")
        except UnicodeDecodeError:
            raise TypeError("NetCDF file does not have correct Conventions")
        try:
            convention_version = self._handle.ConventionVersion.decode("ascii")
        except UnicodeDecodeError:
            raise ValueError("NetCDF file does not have correct ConventionVersion")
        except AttributeError:
            raise TypeError("NetCDF file does not have ConventionVersion")
        if not hasattr(self._handle, "Conventions") or conventions != "AMBERRESTART":
            raise TypeError("NetCDF file does not have correct Conventions")
        if convention_version != "1.0":
            raise ValueError(
                "NetCDF restart has ConventionVersion %s. Only Version 1.0 is supported." % convention_version,
            )
        if atom_indices is not None:
            atom_slice = ensure_type(
                atom_indices,
                dtype=int,
                ndim=1,
                name="atom_indices",
                warn_on_cast=False,
            )
            if not np.all(atom_slice) >= 0:
                raise ValueError("Entries in atom_slice must be >= 0")
            coordinates = self._handle.variables["coordinates"][atom_slice, :]
        else:
            coordinates = self._handle.variables["coordinates"][:, :]

        # Get unit cell parameters
        if "cell_lengths" in self._handle.variables:
            cell_lengths = self._handle.variables["cell_lengths"][:]
        else:
            cell_lengths = None
        if "cell_angles" in self._handle.variables:
            cell_angles = self._handle.variables["cell_angles"][:]
        else:
            cell_angles = None

        if cell_lengths is None and cell_angles is not None:
            warnings.warn("cell_lengths were found, but no cell_angles")
        if cell_lengths is not None and cell_angles is None:
            warnings.warn("cell_angles were found, but no cell_lengths")

        if "time" in self._handle.variables:
            time = self._handle.variables["time"].getValue()
        else:
            time = None

        # scipy.io.netcdf variables are mem-mapped, and are only backed by valid
        # memory while the file handle is open. This is _bad_ because we need to
        # support the user opening the file, reading the coordinates, and then
        # closing it, and still having the coordinates be a valid memory
        # segment.
        # https://github.com/mdtraj/mdtraj/issues/440
        if coordinates is not None and not coordinates.flags["WRITEABLE"]:
            coordinates = np.array(coordinates, copy=True)
        if cell_lengths is not None and not cell_lengths.flags["WRITEABLE"]:
            cell_lengths = np.array(cell_lengths, copy=True)
        if cell_angles is not None and not cell_angles.flags["WRITEABLE"]:
            cell_angles = np.array(cell_angles, copy=True)

        # The leading frame dimension is missing on all of these arrays since
        # restart files have only one frame. Reshape them to add this extra
        # dimension
        coordinates = coordinates[np.newaxis, :]
        if cell_lengths is not None:
            cell_lengths = cell_lengths[np.newaxis, :]
        if cell_angles is not None:
            cell_angles = cell_angles[np.newaxis, :]
        if time is not None:
            time = np.asarray([time])

        return coordinates, time, cell_lengths, cell_angles

    def write(
        self,
        coordinates,
        time=None,
        cell_lengths=None,
        cell_angles=None,
    ):
        """Write one frame of a MD trajectory to disk in the AMBER NetCDF
        restart file format.

        Parameters
        ----------
        coordinates : np.ndarray, dtype=np.float32, shape=([1,] n_atoms, 3)
            The cartesian coordinates of each atom, in units of angstroms. Must
            be only a single frame (shape can be (1,N,3) or (N,3) where N is
            the number of atoms)
        time : array-like with 1 element or float, optional
            The time corresponding to this frame. If not specified, a place
            holder of 0 will be written
        cell_lengths : np.ndarray, dtype=np.double, shape=([1,] 3)
            The lengths (a,b,c) of the unit cell for the frame in Angstroms
        cell_angles : np.ndarray, dtype=np.double, shape=([1,] 3)
            The angles between the unit cell vectors for the frame in Degrees

        Notes
        -----
        You must only have one frame to write to this file.
        """
        if self._mode != "w":
            raise OSError(
                "The file was opened in mode=%s. Writing not allowed." % self._mode,
            )
        if not self._needs_initialization:
            # Must have already been written -- can only write once
            raise RuntimeError(
                "NetCDF restart file has already been written -- can only write one frame to restart files.",
            )
        # these are no-ops
        # coordinates = in_units_of(coordinates, None, 'angstroms')
        # time = in_units_of(time, None, 'picoseconds')
        # cell_lengths = in_units_of(cell_lengths, None, 'angstroms')
        # cell_angles = in_units_of(cell_angles, None, 'degrees')

        # typecheck all of the input arguments rigorously
        coordinates = ensure_type(
            coordinates,
            np.float32,
            3,
            "coordinates",
            length=None,
            can_be_none=False,
            shape=(1, None, 3),
            warn_on_cast=False,
            add_newaxis_on_deficient_ndim=True,
        )
        n_frames, n_atoms = coordinates.shape[0], coordinates.shape[1]
        if n_frames != 1:
            raise ValueError("Can only write 1 frame to a restart file!")
        if time is not None:
            try:
                if isinstance(time, np.ndarray):
                    time = float(time.item())
                else:
                    time = float(time)
            except TypeError:
                raise TypeError("Can only provide a single time")
        else:
            time = 0.0
        cell_lengths = ensure_type(
            cell_lengths,
            np.float64,
            2,
            "cell_lengths",
            length=1,
            can_be_none=True,
            warn_on_cast=False,
            add_newaxis_on_deficient_ndim=True,
        )
        cell_angles = ensure_type(
            cell_angles,
            np.float64,
            2,
            "cell_angles",
            length=1,
            can_be_none=True,
            warn_on_cast=False,
            add_newaxis_on_deficient_ndim=True,
        )
        if (cell_lengths is None and cell_angles is not None) or (cell_lengths is not None and cell_angles is None):
            prov, negl = "cell_lengths", "cell_angles"
            if cell_lengths is None:
                prov, negl = negl, prov
            raise ValueError(
                f'You provided the variable "{prov}" but did not '
                f'provide "{negl}". Either provide both or neither -- '
                "one without the other is meaningless.",
            )

        self._initialize_headers(
            n_atoms=n_atoms,
            set_coordinates=True,
            set_time=(time is not None),
            set_cell=(cell_lengths is not None),
        )
        self._needs_initialization = False

        # Write the time, coordinates, and box info
        if time is not None:
            if isinstance(time, np.ndarray):
                self._handle.variables["time"][0] = float(time.item())
            else:
                self._handle.variables["time"][0] = float(time)
        self._handle.variables["coordinates"][:, :] = coordinates[0, :, :]
        if cell_lengths is not None:
            self._handle.variables["cell_angles"][:] = cell_angles[0, :]
            self._handle.variables["cell_lengths"][:] = cell_lengths[0, :]
        self.flush()

    def _initialize_headers(self, n_atoms, set_coordinates, set_time, set_cell):
        """Initialize the headers and convention properties of the NetCDF
        restart file
        """
        ncfile = self._handle
        ncfile.Conventions = "AMBERRESTART"
        ncfile.ConventionVersion = "1.0"
        ncfile.title = "NetCDF Restart file written by MDTraj w/out velocities"
        ncfile.application = "Omnia"
        ncfile.program = "MDTraj"
        ncfile.programVersion = mdtraj.__version__
        # Dimensions
        ncfile.createDimension("spatial", 3)
        ncfile.createDimension("atom", n_atoms)
        if set_cell:
            ncfile.createDimension("cell_spatial", 3)
            ncfile.createDimension("label", 5)
            ncfile.createDimension("cell_angular", 3)
        if set_time:
            ncfile.createDimension("time", 1)
        # Variables
        v = ncfile.createVariable("spatial", "c", ("spatial",))
        v[:] = np.asarray(list("xyz"))
        v = ncfile.createVariable("coordinates", "d", ("atom", "spatial"))
        v.units = "angstrom"
        if set_cell:
            v = ncfile.createVariable(
                "cell_angular",
                "c",
                ("cell_angular", "label"),
            )
            v[0] = np.asarray(list("alpha"))
            v[1] = np.asarray(list("beta "))
            v[2] = np.asarray(list("gamma"))
            v = ncfile.createVariable("cell_spatial", "c", ("cell_spatial",))
            v[:] = np.asarray(list("abc"))
            v = ncfile.createVariable("cell_lengths", "d", ("cell_spatial",))
            v.units = "angstrom"
            v = ncfile.createVariable("cell_angles", "d", ("cell_angular",))
            v.units = "degree"
        if set_time:
            v = ncfile.createVariable("time", "d", ("time",))
            v.units = "picosecond"
        self.flush()

    def __enter__(self):
        return self

    def __exit__(self, *exc_info):
        self.close()

    def close(self):
        if not self._closed and hasattr(self, "_handle"):
            self._handle.close()
        self._closed = True

    def __del__(self):
        self.close()

    def __len__(self):
        return 1  # All restarts have only 1 frame

    def flush(self):
        self._validate_open()
        if self._mode != "w":
            raise OSError("Cannot flush a file opened for reading")
        self._handle.flush()
