# -*- 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
#
"""\
XDR based trajectory files --- :mod:`MDAnalysis.coordinates.XDR`
================================================================

This module contains helper function and classes to read the XTC and TRR file
formats.

See Also
--------
MDAnalysis.coordinates.XTC: Read and write GROMACS XTC trajectory files.
MDAnalysis.coordinates.TRR: Read and write GROMACS TRR trajectory files.
MDAnalysis.lib.formats.libmdaxdr: Low level xdr format reader
"""

import errno
import numpy as np
from os.path import getctime, getsize, isfile, split, join
import warnings
from filelock import FileLock

from . import base
from ..lib.mdamath import triclinic_box
from ..lib.util import store_init_arguments


def offsets_filename(filename, ending="npz"):
    """Return offset or its lock filename for XDR files.
    For this the filename is appended
    with `_offsets.{ending}`.

    Parameters
    ----------
    filename : str
        filename of trajectory
    ending : str (optional)
        fileending of offsets file

    Returns
    -------
    offset_filename : str

    """
    head, tail = split(filename)
    return join(head, f".{tail}_offsets.{ending}")


def read_numpy_offsets(filename):
    """read offsets into dictionary.

    This assume offsets have been saved using numpy

    Parameters
    ----------
    filename : str
        filename of offsets

    Returns
    -------
    offsets : dict
        dictionary of offsets information

    """
    try:
        return {k: v for k, v in np.load(filename).items()}

    #  `ValueError` is encountered when the offset file is corrupted.
    except (ValueError, IOError):
        warnings.warn(f"Failed to load offsets file {filename}\n")
        return False


class XDRBaseReader(base.ReaderBase):
    """Base class for libmdaxdr file formats xtc and trr

    This class handles integration of XDR based formats into MDAnalysis. The
    XTC and TRR classes only implement `_write_next_frame` and
    `_frame_to_ts`.

    .. _offsets-label:

    Notes
    -----
    XDR based readers store persistent offsets on disk. The offsets are used to
    enable access to random frames efficiently. These offsets will be generated
    automatically the  first time the  trajectory is opened.  Generally offsets
    are stored  in hidden  `*_offsets.npz` files.  Afterwards opening  the same
    file again is fast. It sometimes can happen that the stored offsets get out
    off sync with the trajectory they refer to. For this the offsets also store
    the number of atoms, size of the file and last modification time. If any of
    them change  the offsets are recalculated.  Writing of the offset  file can
    fail when the  directory where the trajectory file resides  is not writable
    or if the  disk is full. In this  case a warning message will  be shown but
    the offsets will nevertheless be used during the lifetime of the trajectory
    Reader. However, the  next time the trajectory is opened,  the offsets will
    have to be rebuilt again.

    .. versionchanged:: 1.0.0
       XDR offsets read from trajectory if offsets file read-in fails
    .. versionchanged:: 2.0.0
       Add a InterProcessLock when generating offsets
    .. versionchanged:: 2.4.0
       Use a direct read into ts attributes
    .. versionchanged:: 2.9.0
       Changed fasteners.InterProcessLock() to filelock.FileLock
    """

    @store_init_arguments
    def __init__(
        self,
        filename,
        convert_units=True,
        sub=None,
        refresh_offsets=False,
        dt=None,
        **kwargs,
    ):
        """
        Parameters
        ----------
        filename : str
            trajectory filename
        convert_units : bool (optional)
            convert units to MDAnalysis units
        sub : array_like (optional)
            `sub` is an array of indices to pick out the corresponding
            coordinates and load only them; this requires that the topology
            itself is that of the sub system.
        refresh_offsets : bool (optional)
            force refresh of offsets
        dt : float (optional)
            timestep in MDAnalysis units to load trajectory with;
            if `dt` is ``None``, the time is taken from the xdr file;
            else, the time is set to `dt` * frame
        **kwargs : dict
            General reader arguments.

        """
        super(XDRBaseReader, self).__init__(
            filename, convert_units=convert_units, **kwargs
        )
        self._xdr = self._file(self.filename)

        self._sub = sub
        if self._sub is not None:
            self.n_atoms = len(self._sub)
        else:
            self.n_atoms = self._xdr.n_atoms

        if not refresh_offsets:
            self._load_offsets()
        else:
            self._read_offsets(store=True)
        frame = self._xdr.read()
        try:
            xdr_frame = self._xdr.read()
            if dt is None:
                dt = xdr_frame.time - frame.time
            else:
                self._ts_kwargs["dt"] = dt
            self._xdr.seek(1)
        except StopIteration:
            dt = 0

        self.ts = self._Timestep(self.n_atoms, **self._ts_kwargs)
        self._frame = 0
        self._frame_to_ts(frame, self.ts)
        # these should only be initialized once
        self.ts.dt = dt
        self.ts.dimensions = triclinic_box(*frame.box)
        if self.convert_units:
            if self.ts.dimensions is not None:
                self.convert_pos_from_native(self.ts.dimensions[:3])

    @classmethod
    def parse_n_atoms(cls, filename, **kwargs):
        with cls._file(filename) as f:
            n_atoms = f.n_atoms
        return n_atoms

    def close(self):
        """close reader"""
        self._xdr.close()

    def _load_offsets(self):
        """load frame offsets from file, reread them from the trajectory if that
        fails. To prevent the competition of generating the same offset file
        from multiple processes, an `InterProcessLock` is used."""
        fname = offsets_filename(self.filename)
        lock_name = offsets_filename(self.filename, ending="lock")

        #  check if the location of the lock is writable.
        try:
            with FileLock(lock_name) as filelock:
                pass
        except OSError as e:
            if isinstance(e, PermissionError) or e.errno == errno.EROFS:
                warnings.warn(
                    f"Cannot write lock/offset file in same location as "
                    f"{self.filename}. Using slow offset calculation."
                )
                self._read_offsets(store=False)
                return
            else:
                raise

        with FileLock(lock_name) as filelock:
            if not isfile(fname):
                self._read_offsets(store=True)
                return

            # if offsets file read correctly, data will be a dictionary of offsets
            # if not, data will be False
            # if False, offsets should be read from the trajectory
            # this warning can be avoided by loading Universe like:
            # u = mda.Universe(data.TPR, data.XTC, refresh_offsets=True)
            # refer to Issue #1893
            data = read_numpy_offsets(fname)
            if not data:
                warnings.warn(
                    f"Reading offsets from {fname} failed, "
                    "reading offsets from trajectory instead.\n"
                    "Consider setting 'refresh_offsets=True' "
                    "when loading your Universe."
                )
                self._read_offsets(store=True)
                return

            ctime_ok = size_ok = n_atoms_ok = False

            try:
                ctime_ok = getctime(self.filename) == data["ctime"]
                size_ok = getsize(self.filename) == data["size"]
                n_atoms_ok = self._xdr.n_atoms == data["n_atoms"]
            except KeyError:
                # we tripped over some old offset formated file
                pass

            if not (ctime_ok and size_ok and n_atoms_ok):
                warnings.warn(
                    "Reload offsets from trajectory\n "
                    "ctime or size or n_atoms did not match"
                )
                self._read_offsets(store=True)
            else:
                self._xdr.set_offsets(data["offsets"])

    def _read_offsets(self, store=False):
        """read frame offsets from trajectory"""
        fname = offsets_filename(self.filename)
        offsets = self._xdr.offsets
        if store:
            ctime = getctime(self.filename)
            size = getsize(self.filename)
            try:
                np.savez(
                    fname,
                    offsets=offsets,
                    size=size,
                    ctime=ctime,
                    n_atoms=self._xdr.n_atoms,
                )
            except Exception as e:
                warnings.warn(f"Couldn't save offsets because: {e}")

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

    def _reopen(self):
        """reopen trajectory"""
        self.ts.frame = 0
        self._frame = -1
        offsets = self._xdr.offsets.copy()
        self._xdr.close()
        self._xdr.open(self.filename.encode("utf-8"), "r")
        # only restore in case we actually had offsets
        if len(offsets) != 0:
            self._xdr.set_offsets(offsets)

    def _read_frame(self, i):
        """read frame i"""
        self._frame = i - 1
        try:
            self._xdr.seek(i)
            timestep = self._read_next_timestep()
        except IOError:
            warnings.warn("seek failed, recalculating offsets and retrying")
            offsets = self._xdr.calc_offsets()
            self._xdr.set_offsets(offsets)
            self._read_offsets(store=True)
            self._xdr.seek(i)
            timestep = self._read_next_timestep()
        return timestep

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


class XDRBaseWriter(base.WriterBase):
    """Base class for libmdaxdr file formats xtc and trr"""

    def __init__(
        self, filename, n_atoms, convert_units=True, dt=None, **kwargs
    ):
        """
        Parameters
        ----------
        filename : str
            filename of trajectory
        n_atoms : int
            number of atoms to be written
        convert_units : bool (optional)
            convert from MDAnalysis units to format specific units
        dt : float (optional)
            timestep in MDAnalysis units to write trajectory with;
            if `dt` is ``None``, time for a frame is set from the timestep;
            else, the time for a frame is `dt` * frame
        **kwargs : dict
            General writer arguments
        """
        self.filename = filename
        self._convert_units = convert_units
        self.n_atoms = n_atoms
        self._xdr = self._file(self.filename, "w")
        self._dt = dt

    def close(self):
        """close trajectory"""
        self._xdr.close()

    def __del__(self):
        self.close()