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

"""
XVG auxiliary reader --- :mod:`MDAnalysis.auxiliary.XVG`
========================================================

xvg files are produced by Gromacs during simulation or analysis, formatted
for plotting data with Grace.

Data is column-formatted; time/data selection is enabled by providing column
indices.

Note
----
By default, the time of each step is assumed to be stored in the first column,
in units of ps.


.. autoclass:: XVGStep
   :members:


XVG Readers
-----------
The default :class:`XVGReader` reads and stores the full contents of the .xvg
file on initialisation, while a second reader (:class:`XVGFileReader`) that
reads steps one at a time as required is also provided for when a lower memory
footprint is desired.

Note
----
Data is assumed to be time-ordered.

Multiple datasets, separated in the .xvg file by '&', are currently not
supported (the readers will stop at the first line starting '&').


.. autoclass:: XVGReader
   :members:

.. autoclass:: XVGFileReader
   :members:


.. autofunction:: uncomment

"""
import numbers
import os
import numpy as np
from . import base
from ..lib.util import anyopen


def uncomment(lines):
    """Remove comments from lines in an .xvg file

    Parameters
    ----------
    lines : list of str
        Lines as directly read from .xvg file

    Yields
    ------
    str
        The next non-comment line, with any trailing comments removed
    """
    for line in lines:
        stripped_line = line.strip()
        # ignore blank lines
        if not stripped_line:
            continue
        # '@' must be at the beginning of a line to be a grace instruction
        if stripped_line[0] == "@":
            continue
        # '#' can be anywhere in the line, everything after is a comment
        comment_position = stripped_line.find("#")
        if comment_position > 0 and stripped_line[:comment_position]:
            yield stripped_line[:comment_position]
        elif comment_position < 0 and stripped_line:
            yield stripped_line
        # if comment_position == 0, then the line is empty


class XVGStep(base.AuxStep):
    """AuxStep class for .xvg file format.

    Extends the base AuxStep class to allow selection of time and
    data-of-interest fields (by column index) from the full set of data read
    each step.

    Parameters
    ----------
    time_selector : int | None, optional
        Index of column in .xvg file storing time, assumed to be in ps. Default
        value is 0 (i.e. first column).
    data_selector : list of int | None, optional
        List of indices of columns in .xvg file containing data of interest to
        be stored in ``data``. Default value is ``None``.
    **kwargs
        Other AuxStep options.

    See Also
    --------
    :class:`~MDAnalysis.auxiliary.base.AuxStep`
    """

    def __init__(self, time_selector=0, data_selector=None, **kwargs):
        super(XVGStep, self).__init__(
            time_selector=time_selector, data_selector=data_selector, **kwargs
        )

    def _select_time(self, key):
        if key is None:
            # here so that None is a valid value; just return
            return
        if isinstance(key, numbers.Integral):
            return self._select_data(key)
        else:
            raise ValueError("Time selector must be single index")

    def _select_data(self, key):
        if key is None:
            # here so that None is a valid value; just return
            return
        if isinstance(key, numbers.Integral):
            try:
                return self._data[key]
            except IndexError:
                errmsg = (
                    f"{key} not a valid index for data with "
                    f"{len(self._data)} columns"
                )
                raise ValueError(errmsg) from None
        else:
            return np.array([self._select_data(i) for i in key])


class XVGReader(base.AuxReader):
    """Auxiliary reader to read data from an .xvg file.

    Default reader for .xvg files. All data from the file will be read and
    stored on initialisation.

    Parameters
    ----------
    filename : str
        Location of the file containing the auxiliary data.
    **kwargs
       Other AuxReader options.

    See Also
    --------
    :class:`~MDAnalysis.auxiliary.base.AuxReader`

    Note
    ----
    The file is assumed to be of a size such that reading and storing the full
    contents is practical.
    """

    format = "XVG"
    _Auxstep = XVGStep

    def __init__(self, filename, **kwargs):
        self._auxdata = os.path.abspath(filename)
        with anyopen(filename) as xvg_file:
            lines = xvg_file.readlines()
        auxdata_values = []
        # remove comments before storing
        for i, line in enumerate(uncomment(lines)):
            if line.lstrip()[0] == "&":
                # multiple data sets not supported; stop at end of first set
                break
            auxdata_values.append([float(val) for val in line.split()])
            # check the number of columns is consistent
            if len(auxdata_values[i]) != len(auxdata_values[0]):
                raise ValueError(
                    "Step {0} has {1} columns instead of "
                    "{2}".format(i, auxdata_values[i], auxdata_values[0])
                )
        self._auxdata_values = np.array(auxdata_values)
        self._n_steps = len(self._auxdata_values)
        super(XVGReader, self).__init__(**kwargs)

    def _memory_usage(self):
        return self._auxdata_values.nbytes

    def _read_next_step(self):
        """Read next auxiliary step and update ``auxstep``.

        Returns
        -------
        AuxStep object
            Updated with the data for the new step.

        Raises
        ------
        StopIteration
            When end of auxiliary data set is reached.
        """
        auxstep = self.auxstep
        new_step = self.step + 1
        if new_step < self.n_steps:
            auxstep._data = self._auxdata_values[new_step]
            auxstep.step = new_step
            return auxstep
        else:
            self.rewind()
            raise StopIteration

    def _go_to_step(self, i):
        """Move to and read i-th auxiliary step.

        Parameters
        ----------
        i: int
            Step number (0-indexed) to move to

        Returns
        -------
        :class:`XVGStep`

        Raises
        ------
        ValueError
            If step index not in valid range.
        """
        if i >= self.n_steps or i < 0:
            raise ValueError(
                "Step index {0} is not valid for auxiliary "
                "(num. steps {1})".format(i, self.n_steps)
            )
        self.auxstep.step = i - 1
        self.next()
        return self.auxstep

    def read_all_times(self):
        """Get NumPy array of time at each step.

        Returns
        -------
        NumPy array of float
            Time at each step.
        """
        return self._auxdata_values[:, self.time_selector]


class XVGFileReader(base.AuxFileReader):
    """Auxiliary reader to read (one step at a time) from an .xvg file.

    An alternative XVG reader which reads each step from the .xvg file as
    needed (rather than reading and storing all from the start), for a lower
    memory footprint.

    Parameters
    ----------
    filename : str
       Location of the file containing the auxiliary data.
    **kwargs
       Other AuxReader options.

    See Also
    --------
    :class:`~MDAnalysis.auxiliary.base.AuxFileReader`


    Note
    ----
    The default reader for .xvg files is :class:`XVGReader`.
    """

    format = "XVG-F"
    _Auxstep = XVGStep

    def __init__(self, filename, **kwargs):
        super(XVGFileReader, self).__init__(filename, **kwargs)

    def _read_next_step(self):
        """Read next recorded step in xvg file and update ``auxstep``.

        Returns
        -------
        AuxStep object
            Updated with the data for the new step.

        Raises
        ------
        StopIteration
            When end of file or end of first data set is reached.
        """
        line = next(self.auxfile)
        while True:
            if not line or (line.strip() and line.strip()[0] == "&"):
                # at end of file or end of first set of data (multiple sets
                # currently not supported)
                self.rewind()
                raise StopIteration
            # uncomment the line
            for uncommented in uncomment([line]):
                # line has data in it; add to auxstep + return
                auxstep = self.auxstep
                auxstep.step = self.step + 1
                auxstep._data = [float(i) for i in uncommented.split()]
                # see if we've set n_cols yet...
                try:
                    auxstep._n_cols
                except AttributeError:
                    # haven't set n_cols yet; set now
                    auxstep._n_cols = len(auxstep._data)
                if len(auxstep._data) != auxstep._n_cols:
                    raise ValueError(
                        f"Step {self.step} has "
                        f"{len(auxstep._data)} columns instead "
                        f"of {auxstep._n_cols}"
                    )
                return auxstep
            # line is comment only - move to next
            line = next(self.auxfile)

    def _count_n_steps(self):
        """Iterate through all steps to count total number.

        Returns
        -------
        int
            Total number of steps
        """
        if not self.constant_dt:
            # check if we've already iterated through to build _times list
            try:
                return len(self._times)
            except AttributeError:
                # might as well build _times now, since we'll need to iterate
                # through anyway
                self._times = self.read_all_times()
                return len(self.read_all_times())
        else:
            # don't need _times; iterate here instead
            self._restart()
            count = 0
            for step in self:
                count = count + 1
            return count

    def read_all_times(self):
        """Iterate through all steps to build times list.

        Returns
        -------
        NumPy array of float
            Time of each step
        """
        self._restart()
        times = []
        for step in self:
            times.append(self.time)
        return np.array(times)