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

"""
TNG trajectory files --- :mod:`MDAnalysis.coordinates.TNG`
==========================================================


The TNG format :footcite:p:`Lundborg2014` is a format used in `GROMACS`_ for
storage of trajectory and topology information. The TNG format allows a wide
range of compression algorithms and unlike the compressed XTC format can also
store velocities and forces in addition to positions.

The classes in this module are based on the `pytng`_ package for reading TNG
files. The reader is directed to the `pytng documentation`_ for further reading
about how pytng works under the hood. Pytng is an optional dependency and must
be installed to use this Reader. Use of the reader without pytng installed will
raise an `ImportError`. The variable `HAS_PYTNG` indicates whether this module
has pytng availble.

In addition to particle-dependent trajectory information like positions,
forces and velocities, the TNG format can store trajectory metadata and
other arbitrary time dependent data. Additional information can range from
the virial and pressure components to the molecular topology of the system.
This is enabled by a block based system in which binary flags indicate the
presence or absence of various data blocks. The structure of a TNG file is
provided  in the `TNG specification`_. The TNG paper :footcite:p:`Lundborg2014` and
the `pytng documentation`_ are also good resources. The user is encouraged to
read the full list of `TNG blocks`_ to understand the full power of the TNG
format.



Current Limitations
-------------------
Currently there is only a Reader for TNG files and no Writer. This will depend
on upstream changes in pytng. Additionally, it is not currently possible to
read the molecular topology from a TNG file.

There are also some limitations to reading TNG files with pytng.
Currently we do not allow data to be read *off stride*. In essence this
means that all of the critical trajectory data (positions, box, velocities
(if present), forces (if present)) must share the same stride in trajectory
integrator steps. These critical blocks in the TNG file are henceforth called
*special blocks*. Optional blocks (all blocks that are not special blocks)
will not be read if they do not share an integrator step with, or are not
divisible by the shared integrator step of the special blocks.


References
----------

.. footbibliography::


.. Links

.. _GROMACS:
    https://www.gromacs.org/
.. _pytng:
    https://github.com/MDAnalysis/pytng
.. _pytng documentation:
    https://www.mdanalysis.org/pytng/
.. _TNG blocks:
    https://www.mdanalysis.org/pytng/documentation_pages/Blocks.html
.. _TNG specification:
    https://gitlab.com/hugomacdermott/tng/-/blob/master/Trajectoryformatspecification.mk

"""

import warnings
from typing import List, Optional

import numpy as np
from MDAnalysis.coordinates import base
from MDAnalysis.coordinates.timestep import Timestep
from MDAnalysis.lib.mdamath import triclinic_box
from MDAnalysis.lib.util import store_init_arguments

from ..due import Doi, due

try:
    import pytng
except ImportError:
    # Indicates whether pytng is found.
    HAS_PYTNG = False
else:
    # Indicates whether pytng is found.
    HAS_PYTNG = True


class TNGReader(base.ReaderBase):
    r"""Reader for the TNG format

    The TNG format :footcite:p:`Lundborg2014` is a format used in `GROMACS`_ for
    storage of trajectory and topology information. This reader is are based on
    the `pytng`_ package for reading TNG files. The reader is directed to the
    `pytng documentation`_ and `TNG specification`_ for further reading.

    The TNG format allows a wide range of compression
    algorithms and unlike the compressed XTC format can also store velocities
    and forces in addition to positions.

    The contents of the *special blocks* (positions, box, velocities, forces)
    are read into the timestep if present. Additional blocks are read into the
    `ts.data` dictionary if they are available at the current frame.

    .. versionadded:: 2.4.0
    """

    format = "TNG"
    # NOTE: Time units are in seconds unlike other GROMACS formats
    units = {
        "time": "second",
        "length": "nm",
        "velocity": "nm/ps",
        "force": "kJ/(mol*nm)",
    }

    _box_blockname = "TNG_TRAJ_BOX_SHAPE"
    _positions_blockname = "TNG_TRAJ_POSITIONS"
    _velocities_blockname = "TNG_TRAJ_VELOCITIES"
    _forces_blockname = "TNG_TRAJ_FORCES"
    _special_blocks = [
        _box_blockname,
        _positions_blockname,
        _velocities_blockname,
        _forces_blockname,
    ]

    @due.dcite(
        Doi("10.1002/jcc.23495"), description="The TNG paper", path=__name__
    )
    @store_init_arguments
    def __init__(self, filename: str, convert_units: bool = True, **kwargs):
        """Initialize a TNG trajectory

        Parameters
        ----------
        filename : str
            filename of the trajectory
        convert_units : bool (optional)
            convert into MDAnalysis units

        """
        if not HAS_PYTNG:
            raise ImportError(
                "TNGReader: To read TNG files please install pytng"
            )

        super(TNGReader, self).__init__(filename, **kwargs)

        self.filename = filename
        self.convert_units = convert_units

        self._file_iterator = pytng.TNGFileIterator(self.filename, "r")
        self.n_atoms = self._file_iterator.n_atoms
        self._n_steps = self._file_iterator.n_steps

        # names of the blocks
        self._block_names = list(self._file_iterator.block_ids.keys())
        # block ids, dict of C long long
        self._block_dictionary = self._file_iterator.block_ids
        self._block_strides = self._file_iterator.block_strides
        self._data_frames = self._file_iterator.n_data_frames
        # init all special blocks to not present
        self._special_block_present = {k: False for k in self._special_blocks}

        self.ts = self._Timestep(self.n_atoms, **self._ts_kwargs)

        # check for all the special blocks
        self._has_box = self._box_blockname in self._block_names
        if self._has_box:
            self._special_block_present[self._box_blockname] = True
            self.ts.dimensions = np.zeros(6, dtype=np.float32)

        self._has_positions = self._positions_blockname in self._block_names
        if self._has_positions:
            self._special_block_present[self._positions_blockname] = True
            self.ts.positions = (
                self._file_iterator.make_ndarray_for_block_from_name(
                    self._positions_blockname
                )
            )

        self._has_velocities = self._velocities_blockname in self._block_names
        if self._has_velocities:
            self._special_block_present[self._velocities_blockname] = True
            self.ts.velocities = (
                self._file_iterator.make_ndarray_for_block_from_name(
                    self._velocities_blockname
                )
            )

        self._has_forces = self._forces_blockname in self._block_names
        if self._has_forces:
            self._special_block_present[self._forces_blockname] = True
            self.ts.forces = (
                self._file_iterator.make_ndarray_for_block_from_name(
                    self._forces_blockname
                )
            )

        # check for any additional blocks that will be read into ts.data
        self._additional_blocks = [
            block
            for block in self._block_names
            if block not in self._special_blocks
        ]
        self._check_strides_and_frames()
        self._frame = 0
        # box needs a temporary place to be stored as ts.dimensions is
        # wrong shape initially
        self._box_temp = self._file_iterator.make_ndarray_for_block_from_name(
            self._box_blockname
        )
        self._frame = -1
        self._read_next_timestep()

    def _check_strides_and_frames(self):
        """
        Check that the strides and frame numbers of the blocks in the TNG
        file match up so that the file can be iterated over retrieving data at
        each integrator step
        """

        strides = []
        n_data_frames = []

        for block in self.special_blocks:
            stride = self._block_strides[block]
            strides.append(stride)
            n_data_frame = self._data_frames[block]
            n_data_frames.append(n_data_frame)

        strides_eq = all(v == strides[0] for v in strides)
        frames_eq = all(v == n_data_frames[0] for v in n_data_frames)

        if (not strides_eq) or (not frames_eq):
            raise IOError(
                "Strides of TNG special blocks not equal, file cannot be read"
            )

        self._global_stride = strides[0]
        # NOTE frame number is 0 indexed so increment
        self._n_frames = n_data_frames[0] + 1

        self._additional_blocks_to_read = []
        for block in self._additional_blocks:
            stride_add = self._block_strides[block]
            if stride_add != self._global_stride:
                if stride_add % self._global_stride:  # pragma: no cover
                    warnings.warn(
                        f"TNG additional block {block} does not match"
                        " strides of other blocks and is not"
                        " divisible by the global stride_length."
                        " It will not be read"
                    )
                else:
                    self._additional_blocks_to_read.append(
                        block
                    )  # pragma: no cover
            else:
                self._additional_blocks_to_read.append(block)

    def close(self):
        """close the reader"""
        self._file_iterator._close()

    @staticmethod
    def parse_n_atoms(filename: str) -> int:
        """parse the number of atoms in the TNG trajectory file

        Parameters
        ----------
        filename : str
            The name of the TNG file

        Returns
        -------
        n_atoms : int
            The number of atoms in the TNG file

        """
        if not HAS_PYTNG:
            raise ImportError(
                "TNGReader: To read TNG files please install pytng"
            )
        with pytng.TNGFileIterator(filename, "r") as tng:
            n_atoms = tng.n_atoms
        return n_atoms

    @property
    def n_frames(self) -> int:
        """number of frames in trajectory

        Returns
        -------
        n_frames : int
            The number of data containing frames in the trajectory
        """
        return self._n_frames

    @property
    def blocks(self) -> List[str]:
        """list of the blocks that are in the file

        Returns
        -------

        blocks : list
            The block names present in the TNG trajectory
        """
        return self._block_names

    @property
    def special_blocks(self) -> List[str]:
        """list of the special blocks that are in the file

        Returns
        -------

        special_blocks : list
            The special block names (positions, box, velocities, forces)
            present in the TNG trajectory
        """
        return [k for k, v in self._special_block_present.items() if v]

    @property
    def additional_blocks(self) -> List[str]:
        """list of the additional (non-special) blocks that are being read from
        the trajectory. This may be exclude some blocks present in the file if
        they do not fall on the same trajectory stride as the positions and
        velocities.

        Returns
        -------

        additional_blocks : list
            The additional block names in the TNG trajectory
        """
        return self._additional_blocks_to_read

    def _reopen(self):
        """reopen the trajectory"""
        self.ts.frame = 0
        self._frame = -1
        self._file_iterator._close()
        self._file_iterator._open(self.filename, "r")

    def _frame_to_step(self, frame: int) -> int:
        """Convert a frame number to an integrator step

        Parameters
        ----------
        frame : int
            The frame number

        Returns
        -------
        integrators_step :  int
            The integrator step of the frame in the TNG file

        """
        return frame * self._global_stride

    def _read_frame(self, i: int) -> Timestep:
        """read frame i

        Parameters
        ----------
        i : int
            The trajectory frame to be read

        Returns
        -------
        ts : Timestep
            Data from frame i encapsulated in an MDA `:class:Timestep`
        """
        self._frame = i - 1
        ts = self._read_next_timestep()
        return ts

    def _read_next_timestep(self, ts: Optional[Timestep] = None) -> Timestep:
        """Read next frame into a timestep

        Parameters
        ----------
        ts : Timestep
            The timestep to read the data into

        Returns
        -------
        ts :  Timestep
            The timestep filled with data from the next step
        """
        if self._frame == self.n_frames - 1:
            raise IOError("trying to go over trajectory limit")
        if ts is None:
            ts = self.ts
        # convert from frames to integrator steps
        step = self._frame_to_step(self._frame + 1)
        iterator_step = self._file_iterator.read_step(step)
        self._frame += 1
        ts = self._frame_to_ts(iterator_step, ts)
        return ts

    def _frame_to_ts(
        self, curr_step: "pytng.TNGCurrentIntegratorStep", ts: Timestep
    ) -> Timestep:
        """convert a TNGCurrentIteratorStep to an MDA Timestep

        Parameters
        ----------
        curr_step : pytng.TNGCurrentIntegratorStep
            The current timestep in the TNG trajectory

        ts : Timestep
            The current timestep in the MDA trajectory

        Returns
        -------
        ts: Timestep
            The current timestep in the MDA trajectory with data from TNG
            trajectory integrator step

        Raises
        ------
        IOError
            The reading of one of the attributes from the TNG file failed
        """

        ts.frame = self._frame
        time = curr_step.get_time()
        if self.convert_units:
            time = self.convert_time_from_native(time)
        ts.time = time
        ts.data["step"] = curr_step.step

        if self._has_box:
            curr_step.get_box(self._box_temp)
            ts.dimensions = triclinic_box(*self._box_temp.reshape(3, 3))
            if not curr_step.read_success:
                raise IOError("Failed to read box from TNG file")
            if self.convert_units:
                self.convert_pos_from_native(ts.dimensions[:3])
        if self._has_positions:
            curr_step.get_positions(ts.positions)
            if not curr_step.read_success:
                raise IOError("Failed to read positions from TNG file")
            if self.convert_units:
                self.convert_pos_from_native(ts.positions)
        if self._has_velocities:
            curr_step.get_velocities(ts.velocities)
            if not curr_step.read_success:
                raise IOError("Failed to read velocities from TNG file")
            if self.convert_units:
                self.convert_velocities_from_native(ts.velocities)
        if self._has_forces:
            curr_step.get_forces(ts.forces)
            if not curr_step.read_success:
                raise IOError("Failed to read forces from TNG file")
            if self.convert_units:
                self.convert_forces_from_native(self.ts.forces)

        for block in self._additional_blocks_to_read:
            add_block_stride = self._block_strides[block]
            # check we are on stride for our block
            if not (add_block_stride % self._global_stride):
                block_data = (
                    self._file_iterator.make_ndarray_for_block_from_name(block)
                )
                # additional blocks read into ts.data dictionary
                ts.data[block] = curr_step.get_blockid(
                    self._block_dictionary[block], block_data
                )
                if not curr_step.read_success:
                    raise IOError(
                        f"Failed to read additional block {block} from TNG file"
                    )
        return ts

    def __getstate__(self):
        """Make a dictionary of the class state to pickle Reader instance.

        Must be done manually as pytng uses a non-trivial`__cinit__`.
        """
        state = self.__dict__.copy()
        # cant have PyTNG file iterator in as it is non-pickable
        del state["_file_iterator"]
        return state

    def __setstate__(self, state):
        """Restore class from `state` dictionary in unpickling of Reader
        instance
        """
        self.__dict__ = state
        # reconstruct file iterator
        self._file_iterator = pytng.TNGFileIterator(self.filename, "r")

        # unlike self._read_frame(self._frame),
        # the following lines update the state of the C-level file iterator
        # without updating the ts object.
        # This is necessary to preserve the modification,
        # e.g. changing coordinates, in the ts object.
        # see PR #3722 for more details.
        step = self._frame_to_step(self._frame)
        _ = self._file_iterator.read_step(step)

    def Writer(self):
        """Writer for TNG files

        Raises
        ------
        NotImplementedError
            Currently there is not writer for TNG files pending implementation
            upstream in pytng.
        """
        raise NotImplementedError("There is currently no writer for TNG files")